How to plot using arays in an array?

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])