How to plot using arays in an array? - python-3.x

Question

I have this coupled mass system code that runs good and prints results. But I have trouble plotting the graphs for positions and velocities since I am unable to extract values from arrays. I would appreciate some help!
import numpy as np
%matplotlib inline
import matplotlib.pyplot as pl
from scipy.integrate import odeint
def vectorfield(w, t, p):
x1, y1, x2, y2 = w
m1, m2, k1, k2, kc = p
# Create f = (x1',y1',x2',y2'):
f = [y1, (-x1*(k1+kc) + x2*kc)/m1, y2, (x1*kc - x2*(k2+kc)) / m2]
return f
# Parameter values
# Masses:
m1 = 1.0
m2 = 1.0
# Spring constants
k1 = 4.0
k2 = 1.0
kc = 0.1
# Initial conditions
# x1 and x2 are the initial displacements; y1 and y2 are the initial velocities
x1 = -2.0
y1 = 5.0
x2 = 2.0
y2 = 10.0
# ODE solver parameters
abserr = 1.0e-8
relerr = 1.0e-6
stoptime = 100.0
numpoints = 250
t = [stoptime * float(i) / (numpoints - 1) for i in range(numpoints)]
# Pack up the parameters and initial conditions:
p = [m1, m2, k1, k2, kc]
w0 = [x1, y1, x2, y2]
# Call the ODE solver.
wsol = odeint(vectorfield, w0, t, args=(p,), atol=abserr, rtol=relerr)
# Print solution
for t1, w1 in zip(t, wsol):
AZ = [t1, w1[0], w1[1], w1[2], w1[3]]
print(AZ)
I have tried searching the web but wasnt unable to find a fitting solution to plot this. I tried
with open('coupled_masses.dat', 'w') as f:
for t1, w1 in zip(t, wsol):
print(f, t1, w1[0], w1[1], w1[2], w1[3])
import matplotlib.pyplot as plt;
from matplotlib.font_manager import FontProperties;
# get saved values from saved file
t, x1, y1, x2, y2 = np.loadtxt('coupled_masses.dat', unpack=True);
but it doesnt work

Answer 1

Is this what you want? Using list comprehension here and then convert to numpy array.
from scipy.integrate import odeint
def vectorfield(w, t, p):
x1, y1, x2, y2 = w
m1, m2, k1, k2, kc = p
# Create f = (x1',y1',x2',y2'):
f = [y1, (-x1*(k1+kc) + x2*kc)/m1, y2, (x1*kc - x2*(k2+kc)) / m2]
return f
# Parameter values
# Masses:
m1 = 1.0
m2 = 1.0
# Spring constants
k1 = 4.0
k2 = 1.0
kc = 0.1
# Initial conditions
# x1 and x2 are the initial displacements; y1 and y2 are the initial velocities
x1 = -2.0
y1 = 5.0
x2 = 2.0
y2 = 10.0
# ODE solver parameters
abserr = 1.0e-8
relerr = 1.0e-6
stoptime = 100.0
numpoints = 250
t = [stoptime * float(i) / (numpoints - 1) for i in range(numpoints)]
# Pack up the parameters and initial conditions:
p = [m1, m2, k1, k2, kc]
w0 = [x1, y1, x2, y2]
# Call the ODE solver.
wsol = odeint(vectorfield, w0, t, args=(p,), atol=abserr, rtol=relerr)
# Print solution
data = np.array([[t1, w1[0], w1[1], w1[2], w1[3]] for t1, w1 in zip(t, wsol)])
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(7.2, 7.2/2))
ax1.plot(data[:, 0], data[:, 1])
ax2.plot(data[:, 0], data[:, 3])