i have a final project in physics in which i should numerically solve some 4D ODE system and show its chaotic properties by visualaizing how all the initial conditions evolve over time. i wanted to plot it like the video - https://youtu.be/n7JK4Ht8k8M?t=106 but when trying to run the code with grid of 100x100 it is just taking very long time, and i can't give up on high resulution beacuse i have to show the fractal behvior of the system. i was looking for hours of how to solve this using my gpu (rtx 3070) but found almost nothing related to ODE'S using python. is there any way to approach this? thanks!
#%% Import Packages
import numpy as np
from numpy import cos , sin
from scipy.integrate import odeint
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.colors import to_rgb
import math
#%% Define Simulation parameters
TIME_FINAL = 10
FRAME_NUM = 300
MAP_RES = 1000
#%% Define Control Parameters
g = 9.8
l = 1
m = 1
#%% Define the ODE System
# Define a time axis (number of frames)
t = np.linspace(0, TIME_FINAL, FRAME_NUM)
# Define our model (Contains the equations of motion for our system)
def model(cords , t):
th1 , th2 , p1 , p2 = cords
th1_dot = 6/(m*l**2) * ((2*p1 - 3*cos(th1-th2)*p2) / (16 - 9*(cos(th1 -th2))**2))
th2_dot = 6/(m*l**2) * ((8*p2 - 3*cos(th1-th2)*p1) / (16 - 9*(cos(th1 -th2))**2))
p1_dot = -0.5*m*l**2 * (th1_dot*th2_dot*sin(th1-th2) + 3*(g/l)*sin(th1))
p2_dot = -0.5*m*l**2 * (-th1_dot*th2_dot*sin(th1-th2) + (g/l)*sin(th1))
return [th1_dot , th2_dot , p1_dot , p2_dot]
#%% First Impression
init0 = np.array([1.3 , 0.7 , 0 , 0]) # Adjust the values to observe chaotic behivior
th1 , th2 , p1 , p2 = odeint(model,init0 ,t).T
plt.figure(dpi = 300)
plt.plot(t, th1 , label = r"$\theta1$", linewidth = 1)
plt.plot(t, th2 , label = r"$\theta2$", linewidth = 1)
plt.grid("on")
plt.legend()
plt.title(fr"$\theta1$ and $\theta2$ ({init0})")
#%% Constructing an initial sub phase-space
# Define the ranges for theta and p_theta coordinates
theta1_range = np.linspace(0, 2*np.pi, MAP_RES)
theta2_range = np.linspace(0, 2*np.pi, MAP_RES)
# Create a mesh grid of theta1 and theta2 values
theta1_grid, theta2_grid = np.meshgrid(theta1_range, theta2_range)
# %% Calculating the Sub Phase Space Progression Over Time
p1_i = 0
p2_i = 0
i = 0
phase_space = []
for (theta , phi) in zip(np.ravel(theta1_grid), np.ravel(theta2_grid)) :
initial_condition = [theta , phi ,p1_i , p2_i]
th1 , th2 , _ , __ = odeint(model,initial_condition ,t).T
phase_space.append((th1,th2))
i = i + 1
print(f"Completed {i}/{theta1_grid.size}")
#%% Generating the Color map from the sub phase space data
def rad2rgb(theta, phi, radius = 127):
return [math.floor(radius + radius * np.cos(phi) * np.sin(theta)),
math.floor(radius + radius * np.sin(phi) * np.sin(theta)),
math.floor(radius + radius * np.cos(theta))]
color_maps = []
for i in range(t.size) :
rgb_colors = np.array([rad2rgb(phase[0][i], phase[1][i]) for phase in phase_space])
rgb_colors = rgb_colors.reshape(theta1_grid.shape[0], theta1_grid.shape[1], 3)
color_maps.append(rgb_colors)
#%% Animating the Phase Space and saving into .GIF file
fig, ax = plt.subplots()
def update(frame):
ax.clear()
ax.imshow(color_maps[frame], origin='lower', extent=(0, 2*np.pi, 0, 2*np.pi), aspect='auto')
ax.set_xlabel(r'$\theta1$')
ax.set_ylabel(r'$\theta2$')
ax.set_title(r'Color Map in $(\theta1 , \theta2)$ Space')
anim = FuncAnimation(fig, update, frames= len(color_maps), interval=50)
#anim.save('Fractal.gif', writer='pillow', fps=60, dpi=150)
plt.show()