test_equations.py

import equations
import numpy as np
import matplotlib.pyplot as plt
import math
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import animation
from PIL import Image

# flake8: noqa


def animation_CW(CW, r0, v0, n, dt, nframes):

    xs = []
    ys = []
    zs = []

    for j in range(nframes):
        xs.append(CW(r0, v0, n, j*dt)[0])
        ys.append(CW(r0, v0, n, j*dt)[1])
        zs.append(CW(r0, v0, n, j*dt)[2])
    xmin = min(xs)
    xmax = max(xs)
    ymin = min(ys)
    ymax = max(ys)
    zmin = min(zs)
    zmax = max(zs)

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d', xlim=(xmin, xmax), ylim=(ymin, ymax), zlim=(zmin, zmax))
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z')
    if CW.__name__ == 'CW_solution':
        ax.set_title('Real trajectory')
    if CW.__name__ == 'CW_finite_diff':
        ax.set_title('Finite diff. trajectory')
    def init():
        return(ax)

    def animate(j):
        x = CW(r0, v0, n, j*dt)[0]
        y = CW(r0, v0, n, j*dt)[1]
        z = CW(r0, v0, n, j*dt)[2]
        ax.scatter(x, y, z, marker='.', c='r', alpha=0.5)
        return(ax)
        
    anim = animation.FuncAnimation(fig, animate, init_func=init, frames = nframes, interval = 0.000001)
    filename = f'{CW.__name__}_animation.gif'
    anim.save(filename, writer='pillow')

def CW_solution(r0, v0, n, t):
    
    x0, y0, z0 = r0
    vx0, vy0, vz0 = v0

    a1 = 4*x0 + 2*vy0
    a2 = y0 - 2*vx0
    a3 = -3*x0 - 2*vy0
    a4 = vx0
    a5 = z0
    a6 = vz0

    x = a1 + a3*np.cos(n*t) + a4*np.sin(n*t)
    y = a2 - 3/2*a1*n*t + 2*a4*np.cos(n*t) - 2*a3*np.sin(n*t)
    z = a5*np.cos(n*t) + a6*np.sin(n*t)

    vx = -a3*n*np.sin(n*t) + a4*n*np.cos(n*t)
    vy = -3/2*a1*n - 2*a3*n*np.cos(n*t) - 2*a4*n*np.sin(n*t)
    vz = -a5*n*np.sin(n*t) + a6*n*np.cos(n*t)
    
    r = np.array([x, y, z])
    v = np.array([vx, vy, vz])

    return r


r0 = [5, 0.0, 0]
v0 = [0.1, 0.3, -0.023]
n = 0.0010854
dt = 60
nframes = 300

animation_CW(CW_solution, r0, v0, n, dt, nframes)