How to set a numpy array in loop and plot

Question

I'm trying to modify a code from the scipy cookbook

The scipy cookbook solves some state equations then prints them:

wsol = odeint(two_springs.vectorfield, w0, t, args=(p,),
              atol=abserr, rtol=relerr)

for t1, w1 in zip(t, wsol):
    print t1, w1[0], w1[1], w1[2], w1[3]

The scipy cookbook requires you to ru the program and save printed stuff to a file:

program.py > file.dat

then a separate plotting function requires you to load the data from that file

t, x1, xy, x2, y2 = loadtxt('two_springs.dat', unpack=True)

before plotting

plot(t, x1, 'b', linewidth=lw)
plot(t, x2, 'g', linewidth=lw)

I just want to set the t,x1,x2 and plot

Can anyone show me how to do that?

This is what I'm trying to modify:

http://wiki.scipy.org/Cookbook/CoupledSpringMassSystem

Answer 1

I assume that wsol is Nx4 ndarray, thus you just need to un-pack the values:

x1, xy, x2, y2 = wsol.T