Graphical Differential equations

graph_diff_eq.py

import itertools as it
import matplotlib.pyplot as plt
import numpy as np
from scipy.integrate import odeint

mu = 0.3
L = 10
g = 9.8
DELTA = 0.01
T = 1000
DX = .1
DY = .1


def c(v, min_, max_):
    p = (v - min_) / (max_ - min_)
    return 'y' if p < 0.3 else 'orange' if p < 0.6 else 'r'


def pendulum(y, t, mu, c):
    theta, theta_dot = y
    return [
        theta_dot,
        -mu * theta_dot - c * np.sin(theta)
    ]


def spring(y, t, mu, L):
    x, v = y
    return [
        v,
        -mu * v - L * x
    ]


func = pendulum
MX = 10
MY = 10

y0 = [np.pi - 0.1, 1.0]
t = np.arange(0, T, DELTA)
sol = odeint(func, y0, t, args=(mu, g/L))
plt.plot(sol[:, 0], sol[:, 1],
         label=r'$y(t) = \left[\stackrel{\theta(t)}{\omega(t)}\right]$')
plt.legend()
plt.xlabel(r"$\theta(t)$")
plt.ylabel(r"$\omega(t)$")

mX = np.min(sol[:, 0])
mY = np.min(sol[:, 1])
MX = np.max(sol[:, 0])
MY = np.max(sol[:, 1])

x = np.arange(mX, MX, DX)
y = np.arange(mY, MY, DY)

xx, yy = np.meshgrid(x, y)
uu, vv = func([xx, yy], t, mu, g / L)

norm = np.sqrt(uu*uu + vv*vv)
max_norm = np.max(norm)
min_norm = np.min(norm)
colors = [c(v, min_norm, max_norm) for v in np.array(norm).reshape(1, -1)[0]]
norm[norm == 0] = 1
uu = uu / norm
vv = vv / norm

plt.quiver(xx, yy, uu, vv, width=.001, color=colors)
plt.grid()
plt.show()