- This example only uses a second order ODE, but should get the idea across for how to add a legend.
- Convert the unstructured array into a structured array, by adding
dtype to the sol array.
from scipy.integrate import odeint
import numpy as np
import matplotlib.pyplot as plt
# ode system
def pend(y, t, b, c):
theta, omega = y
dydt = [omega, -b*omega - c*np.sin(theta)]
return dydt
b = np.arange(0.50, 0, -0.25)
c = 5.0
y0 = [np.pi - 0.1, 0.0]
t = np.linspace(0, 10, 101)
ls = ['-', '--'] # list of linestyles
count = 0 # used to index the current linestyle
for i in b:
sol = odeint(pend, y0, t, args=(i, c))
type_to_set = 'float64' # specify the type you need (e.g. 'float64')
dtype = np.dtype([(f'{x}_ode:{i}', type_to_set) for x in range(sol.shape[1])]) # create names and dtypes list
sol_updated = sol.astype(type_to_set).view(dtype) # add names and dtypes to sol
# plot
# plt.figure() # if you want separate plots for each group uncomment this line & put plt.legend in 4 spaces, so it's just after the plotting loop
for name in sol_updated.dtype.names:
plt.plot(t, sol_updated[name], linestyle=ls[count], label=name)
count += 1 # add 1 to the markers indexer
plt.gca().set_prop_cycle(None)
plt.legend(bbox_to_anchor=(1.04, 0.5), loc='center left', borderaxespad=0)
plt.show()
- This is 2 sets of (101, 2) arrays on the same plot with unique line styles for each
ode solution group.
