Python Forum

Full Version: Python scipy odeint: solving with solution-dependent functions
You're currently viewing a stripped down version of our content. View the full version with proper formatting.
Hi,

I have been implementing a code to solve a system of coupled ODEs using scipy.odeint, similarly to the following one which is for a coupled-spring mass system (from https://scipy-cookbook.readthedocs.io/it...ystem.html):

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, b1, b2 = p

    # 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]
    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
# Friction coefficients
b1 = 0.8
b2 = 0.5

# 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 = 250

# Create the time samples for the output of the ODE solver.
# I use a large number of points, only because I want to make
# a plot of the solution that looks nice.
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, b1, b2]
w0 = [x1, y1, x2, y2]

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

with open('two_springs.dat', 'w') as f:
    # Print & save the solution.
    for t1, w1 in zip(t, wsol):
        print(f, t1, w1[0], w1[1], w1[2], w1[3])

# Plot the solution that was generated

from numpy import loadtxt
from pylab import figure, plot, xlabel, grid, legend, title, savefig #, hold, legend, title, savefig
from matplotlib.font_manager import FontProperties

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

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

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

#plot(t, x1, 'b', linewidth=lw)
#plot(t, x2, 'g', linewidth=lw)
plot(t, wsol[:,0], 'b', linewidth=lw)
plot(t, wsol[:,2], 'g', linewidth=lw)

legend((r'$x_1$', r'$x_2$'), prop=FontProperties(size=16))
title('Mass Displacements for the\nCoupled Spring-Mass System')
savefig('two_springs.png', dpi=100)
But I've been struggling to set some parameters as functions that are variable-dependent, e.g. say: at each time step, if x2>1.5: m1=10*x2*y2, else: m1=x2*y2

Also, is it possible to define some parameters as "non-derivative" variables to be solved by odeint, e.g. say: making b1 and b2 to have initial conditions, then solving for them as well at each time step

Should they be set as additional variables within the odeint function, e.g. as part of w or f=[] ?

Any help would be great. Cheers