Quaternions

animate.py

from mpl_toolkits.mplot3d import axes3d
import matplotlib.pyplot as plt
from IPython.display import HTML
import matplotlib.animation as animation

import numpy as np
import quaternion


def q(phi, v):
    v1 = np.sin(phi/2) * (v / np.linalg.norm(v))
    return np.quaternion(np.cos(phi/2), v1[0], v1[1], v1[2])

def rotate(p, phi, v, f=1):
    _phi = phi * f
    _p = np.quaternion(0, p[0], p[1], p[2])
    r = q(_phi, v) * _p * q(-_phi, v)
    return quaternion.as_float_array(r)[1:]

def animate_multiple(p, phi, v, filename=None, bound=2, time=15, az_start=60, fig_size=(8,8), fps=30, start=3, end=3, bitrate=2000, axis=True):
    fig = plt.figure(figsize=fig_size)
    ax = plt.subplot(111, projection='3d')
    
    phi = np.array(phi)
    v = np.array(v)
    p = np.array(p)
    
    if(phi.size * 3 != v.size):
        raise TypeError("p, phi, v, arrays have to be the same size")
    
    n = phi.size

    f_count = int(time / (start + end + n) * fps)
    
    frames = []
    last_p = p
    for i in range(n):
        frames.append([rotate(last_p, phi[i], v[i] , f / f_count) for f in range(f_count)])
        last_p = rotate(last_p, phi[i], v[i] , 1)

    def init():
        ax.set_zlim(-bound, bound)
        ax.set_xlim(-bound, bound)
        ax.set_ylim(-bound, bound)
        if(axis):
            plot_line(ax, [-bound, 0, 0], [bound,0,0], c="gray")
            plot_line(ax, [0, -bound, 0], [0,bound,0], c="gray")
            plot_line(ax, [0,  0,-bound], [0,0,bound], c="gray")
            for m in range(n):
                _v = (bound / np.max(v[m])) * v[m]
                plot_line(ax, -_v, _v, c="red")
        ax.scatter3D(p[0], p[1], p[2], c="blue")
        return fig,

    def update(i):
        if i < f_count * start:
            ax.view_init(elev=30., azim= az_start + 360 / (f_count * start) * i)
        elif i < (n + start) * f_count :
            m = int((i / f_count)) - start
            q = frames[m][i % f_count]
            ax.view_init(elev=30., azim=az_start)
            ax.scatter3D(q[0], q[1], q[2], c="blue")
        else:
            ax.view_init(elev=30., azim= az_start + 360 / (f_count * end) * (i - (start + n) * f_count))
        return fig,

    ani = FuncAnimation(fig, update, frames=(start + end + n) * f_count, interval= 1000 / fps, 
                        init_func=init, blit=True)
    
    if(filename is not None):
        Writer = animation.writers['ffmpeg']
        writer = Writer(fps=fps, metadata=dict(artist='Krzysztof Zdulski & Bartosz Walusiak'), bitrate=bitrate)
        ani.save(filename, writer=writer)
        return
    return HTML(ani.to_html5_video())