import numpy as np import matplotlib.pyplot as plt from math import isclose from numpy import linalg as LA import matplotlib.cm as cm

def plot_circle(v_1, v_2, ax, **kwargs):

   angles = np.linspace(0, 2*np.pi, 100)
   points = v_1[:,np.newaxis] * np.cos(angles) + v_2[:,np.newaxis] * np.sin(angles)
   ax.plot(points[0,:], points[1,:], **kwargs)

def plot_vector_field(A, xmin=-5, xmax=5, ymin=-5, ymax=5, title=""):

   axx, axy = A[0,0], A[0,1]
   ayx, ayy = A[1,0], A[1,1]
   det = axx * ayy - axy * ayx
   tr = axx + ayy
   eigen_vals, eigen_vects = LA.eig(A)
   is_critical = abs(eigen_vals[0] - eigen_vals[1]) / abs(eigen_vals[0]) < 1e-2
   delta = tr**2 - 4*det
   is_rotational = delta <= 0 and not is_critical

   # Initialize plotting object
   fig, axes = plt.subplot_mosaic("133;233", figsize=(18,12))
   colormap=cm.viridis

   # pole-zero plot
   ax = axes['1']
   ax.scatter(eigen_vals[0].real, eigen_vals[0].imag, color=colormap(eigen_vals[0].real))
   ax.scatter(eigen_vals[1].real, eigen_vals[1].imag, color=colormap(eigen_vals[1].real))
   r = np.sqrt(abs(eigen_vals[0] * eigen_vals[1]))
   plot_circle(np.array([r, 0]), np.array([0, r]), ax, color='k', alpha=0.3)

   ax.plot([xmin, xmax], np.zeros(2), color='k', alpha=0.2)
   ax.plot(np.zeros(2), [ymin, ymax], color='k', alpha=0.2)
   ax.set_aspect('equal')
   ax.set_xlim([-2, 2])
   ax.set_ylim([-2, 2])
   ax.set_xlabel('Real')
   ax.set_ylabel('Imag')
   ax.set_title('pole-zero plot')

   # stability plot
   ax = axes['2']
   xs = np.linspace(xmin, xmax, 100)
   ys = xs**2 / 4
   ax.plot(xs, ys)
   ax.scatter(tr, det, color='red')
   
   ax.plot([xmin, xmax], np.zeros(2), color='k', alpha=0.2)
   ax.plot(np.zeros(2), [ymin, ymax], color='k', alpha=0.2)
   ax.set_aspect('equal')
   ax.set_xlim([-4,2])
   ax.set_ylim([-1, 5])
   ax.set_xlabel('Tr(A)')
   ax.set_ylabel('Det(A)')
   ax.set_title('stability plot')

   # vector field plot
   ax = axes['3']
   x, y = np.meshgrid(np.linspace(xmin, xmax, 10), np.linspace(ymin, ymax, 10))
   vx = axx * x + axy * y
   vy = ayx * x + ayy * y

   ax.quiver(x,y, vx, vy, units='xy', scale=6, color='g', headwidth=3, width=0.04)    
   # Plot the circle, or fast and slow manifolds
   if is_rotational:
       v_1 = np.array(eigen_vects[:,0].real)
       v_2 = np.array(eigen_vects[:,0].imag)
       # normalize
       radius = (xmax - xmin) / 4
       norm = max(np.linalg.norm(v_1), np.linalg.norm(v_2)) / radius
       v_1 /= norm
       v_2 /= norm
       plot_circle(v_1, v_2, ax, color=colormap(eigen_vals[0].real))
       
   elif is_critical:
       v_1 = eigen_vects[:,0]
       length = (xmax - xmin) * 2
       lengths = np.linspace(-length, length, 100)
       points = v_1[:,np.newaxis] * lengths
       ax.plot(points[0,:], points[1,:], color=colormap(eigen_vals[0]))

   else:
       v_1 = eigen_vects[:,0]
       v_2 = eigen_vects[:,1]

       length = (xmax - xmin) * 2
       lengths = np.linspace(-length, length, 100)
       points = v_1[:,np.newaxis] * lengths
       ax.plot(points[0,:], points[1,:], color=colormap(eigen_vals[0]))
       
       points = v_2[:,np.newaxis] * lengths
       ax.plot(points[0,:], points[1,:], color=colormap(eigen_vals[1]))

   ax.plot([xmin, xmax], np.zeros(2), color='k', alpha=0.2)
   ax.plot(np.zeros(2), [ymin, ymax], color='k', alpha=0.2)
   ax.set_aspect('equal')
   ax.set_xlim([xmin, xmax])
   ax.set_ylim([ymin, ymax])
   ax.set_xlabel('$x$')
   ax.set_ylabel('$\\dot{x}$')
   
   ax.set_title('phase portrait')
   fig.suptitle(title)
   fig.tight_layout()
   return fig

import tempfile import os import imageio

plt.rc('figure', titlesize=16)

with tempfile.TemporaryDirectory() as temp_dir:

   n_frames = 201
   
   omegas = [1.0] * n_frames
   gammas = (1-np.cos(np.linspace(0, np.pi, n_frames//2)))/2
   gammas = list(gammas) + [gammas[-1]] + list(gammas + 1)
   for i in range(n_frames):
       omega = omegas[i]
       gamma = gammas[i]
       operator_A = np.array([[0, 1], [-omega**2, -2*gamma]])
       fig = plot_vector_field(operator_A, title=f"$\\ddot x + 2\\gamma \\dot x + \\omega^2x = 0$,\n$\\omega = {omega:1.1f}$, $\\gamma = {gamma:0.3f}$")
       filename = os.path.join(temp_dir, f"plot_{i:03d}.png")
       fig.savefig(filename)
       plt.close(fig)

   # Compile images into GIF
   fps = 12
   images = []
   for i in range(n_frames):
       filename = os.path.join(temp_dir, f"plot_{i:03d}.png")
       images.append(imageio.v2.imread(filename))
   imageio.mimsave(f"phase_portrait_omega_{omega:1.1f}.gif", images, duration=1/fps)