Python scipy solve_ivp / odeint: print & plot time/solution-dependent parameters

Question

I have been using scipy.integrate.solve_ivp to solve a system of 2nd order ODEs.

Assuming the code available here (note this uses odeint, but similar approach with solve_ivp): https://scipy-cookbook.readthedocs.io/items/CoupledSpringMassSystem.html

Now, making some parameters (b1 and b2) time & solution-dependent, say: e.g.

b1 = 0.8 * x1

if (x2 >= 1.0):
    b2 = 20*x2
else:
    b2 = 1.0

The solve_ivp solutions can easily be printed and plotted.

I'm also able to print part of the time/solution-dependent parameters by inserting print(b1) within the defined function. But I'm confused with this output. What I'm after is to print & plot the time/solution-dependent parameters set within the defined function being solved at the solution timestamps: e.g. print(t, b1).

Hoping this makes sense. Any hints would be great.

Cheers

Sample code (with odeint and b1 & b2 as dynamic parameters):

    def vectorfield(w, t, p):
    """
    Defines the differential equations for the coupled spring-mass system.

    Arguments:
        w :  vector of the state variables:
                  w = [x1,y1,x2,y2]
        t :  time
        p :  vector of the parameters:
                  p = [m1,m2,k1,k2,L1,L2,b1,b2]
    """
    x1, y1, x2, y2 = w
    m1, m2, k1, k2, L1, L2, = p
    
    # Friction coefficients
    b1 = 0.8 * x1
    
    if (x2 >= 1.0):
        b2 = 20*x2
    else:
        b2 = 1.0
    
    # Create f = (x1',y1',x2',y2'):
    f = [y1,
         (-b1 * y1 - k1 * (x1 - L1) + k2 * (x2 - x1 - L2)) / m1,
         y2,
         (-b2 * y2 - k2 * (x2 - x1 - L2)) / m2]
    
    print('b1', b1)
    # print('b2', b2)
    
    return f

# Use ODEINT to solve the differential equations defined by the vector field
from scipy.integrate import odeint

# Parameter values
# Masses:
m1 = 1.0
m2 = 1.5
# Spring constants
k1 = 8.0
k2 = 40.0
# Natural lengths
L1 = 0.5
L2 = 1.0

# Initial conditions
# x1 and x2 are the initial displacements; y1 and y2 are the initial velocities
x1 = 0.5
y1 = 0.0
x2 = 2.25
y2 = 0.0

# ODE solver parameters
abserr = 1.0e-8
relerr = 1.0e-6
stoptime = 10.0
numpoints = 101

# Create the time samples for the output of the ODE solver.
t = [stoptime * float(i) / (numpoints - 1) for i in range(numpoints)]

# Pack up the parameters and initial conditions:
p = [m1, m2, k1, k2, L1, L2]
w0 = [x1, y1, x2, y2]

# Call the ODE solver.
wsol = odeint(vectorfield, w0, t, args=(p,),
              atol=abserr, rtol=relerr)

# Print Output
print(t)
# print(wsol)
# print(t, b1)
# print(t, b2)


# Plot the solution that was generated
from numpy import loadtxt
from pylab import figure, plot, xlabel, grid, legend, title, savefig
from matplotlib.font_manager import FontProperties

figure(1, figsize=(6, 4.5))

xlabel('t')
grid(True)
lw = 1

plot(t, wsol[:,0], 'b', linewidth=lw)
plot(t, wsol[:,1], 'g', linewidth=lw)
plot(t, wsol[:,2], 'r', linewidth=lw)
plot(t, wsol[:,3], 'c', linewidth=lw)
# plot(t, b1, 'k', linewidth=lw)

legend((r'$x_1$', r'$x_2$',r'$x_3$', r'$x_4$'), prop=FontProperties(size=16))
title('Mass Displacements for the\nCoupled Spring-Mass System')
# savefig('two_springs.png', dpi=100)

Answer 1

New ODE Solver API

Adapting your code to use solve_ivp (which is the modern scipy API to solve ODE) we can solve the system in a non intrusive way:

import numpy as np
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt

numpoints = 1001
t0 = np.linspace(0, stoptime, numpoints)

def func(t, x0, beta):
    return vectorfield(x0, t, beta)

#sol1 = odeint(vectorfield, w0, t0, args=(p,), atol=abserr, rtol=relerr)
sol2 = solve_ivp(func, [t0[0], t0[-1]], w0, args=(p,), t_eval=t0, method='LSODA')

The new API returns pretty much similar results than old API odeint. The solution object returned by solve_ivp provide multiple extra information about the job execution:

  message: 'The solver successfully reached the end of the integration interval.'
     nfev: 553
     njev: 29
      nlu: 29
      sol: None
   status: 0
  success: True
        t: array([ 0.  ,  0.01,  0.02, ...,  9.98,  9.99, 10.  ])
 t_events: None
        y: array([[ 0.5       ,  0.50149705,  0.50596959, ...,  0.60389634,
         0.60370222,  0.6035087 ],
       [ 0.        ,  0.29903815,  0.59476327, ..., -0.01944383,
        -0.01937904, -0.01932493],
       [ 2.25      ,  2.24909287,  2.24669965, ...,  1.62461438,
         1.62435635,  1.62409906],
       [ 0.        , -0.17257692, -0.29938513, ..., -0.02583719,
        -0.02576507, -0.02569245]])
 y_events: None

Catching dynamic parameters

If I correctly understand your question, you want to have your friction parameters as time series as well.

Follow the solver execution

This could be achieved using a callback function or passing an extra data structure and store coefficient value when evaluating the ODE function.

This will be helpful if you want to track the solver execution, but the result will be dependent on the solver integration grid instead of the desired solution time scale.

Anyway, if you wish to do so, just adapt your function:

def vectorfield2(w, t, p, store):
   # ...
   store.append({'time': t, 'coefs': [b1, b2]})
   # ...

c = []
sol1 = odeint(vectorfield2, w0, t0, args=(p, c), atol=abserr, rtol=relerr)

As you can see, the result does not follow the solution time scale but the solver integration grid instead:

[{'time': 0.0, 'coefs': [0.4, 45.0]},
 {'time': 3.331483023263848e-07, 'coefs': [0.4, 45.0]},
 {'time': 3.331483023263848e-07,
  'coefs': [0.4000000000026638, 44.999999999955605]},
 {'time': 6.662966046527696e-07,
  'coefs': [0.4000000000053275, 44.99999999991122]},
 {'time': 6.662966046527696e-07,
  'coefs': [0.40000000000799124, 44.99999999986683]},
 {'time': 0.0001385450759485236,
  'coefs': [0.4000002303394594, 44.99999616108455]}, ...]

Coefficient as a state function

On the other hand, your coefficients only depend on the system states, it means we can compute it inside the ODE function (each time the solver calls it) and afterward when the time dependent solutions are known.

I would suggest to rewrite your system as follow:

def coefs(x):
    # Ensure signature:
    if isinstance(x, (list, tuple)):
        x = np.array(x)
    if len(x.shape) == 1:
        x = x.reshape(1, -1)
    # Compute coefficients
    b1 = 0.8 * x[:,0]
    b2 = 20. * x[:,2]
    q2 = (x[:,2] < 1.0)
    b2[q2] = 1.0
    return np.stack([b1, b2]).T

def system(t, w, p):
    x1, y1, x2, y2 = w
    m1, m2, k1, k2, L1, L2, = p
    b = coefs(w)
    return [
        y1, (-b[:,0] * y1 - k1 * (x1 - L1) + k2 * (x2 - x1 - L2)) / m1,
        y2, (-b[:,1] * y2 - k2 * (x2 - x1 - L2)) / m2,
    ]

sol3 = solve_ivp(system, [t0[0], t0[-1]], w0, args=(p,), t_eval=t0, method='LSODA')

Notice the interface of coefs method is designed to operate on simple state (list or vector) and on time solutions as well (matrix).

Then, it is easy to display time solutions:

fig, axe = plt.subplots()
axe.plot(t0, sol3.y.T)
axe.set_title("Dynamic System: Coupled Spring Mass")
axe.set_xlabel("Time, $t$")
axe.set_ylabel("System Coordinates, $x_i(t)$")
axe.legend([r'$x_%d$' % i for i in range(sol3.y.T.shape[1])])
axe.grid()

And coefficients:

fig, axes = plt.subplots(2,1, sharex=True, sharey=False)
axes[0].plot(t0, coefs(sol3.y.T)[:,0])
axes[1].plot(t0, coefs(sol3.y.T)[:,1])
axes[0].set_title("Dynamic System: Coupled Spring Mass")
axes[1].set_xlabel("Time, $t$")
axes[0].set_ylabel("$b_0(t)$")
axes[1].set_ylabel("$b_1(t)$")
for i in range(2):
    axes[i].grid()

Which seems more likely what you wanted to achieve.