Solving ODE numerically with Python

Question

I am solving an ODE for an harmonic oscillator numerically with Python. When I add a driving force it makes no difference, so I'm guessing something is wrong with the code. Can anyone see the problem? The (h/m)*f0*np.cos(wd*i) part is the driving force.

import numpy as np
import matplotlib.pyplot as plt

# This code solves the ODE mx'' + bx' + kx = F0*cos(Wd*t)
# m is the mass of the object in kg, b is the damping constant in Ns/m
# k is the spring constant in N/m, F0 is the driving force in N,
# Wd is the frequency of the driving force and x is the position 

# Setting up

timeFinal= 16.0   # This is how far the graph will go in seconds
steps = 10000     # Number of steps
dT = timeFinal/steps      # Step length 
time = np.linspace(0, timeFinal, steps+1)   
# Creates an array with steps+1 values from 0 to timeFinal

# Allocating arrays for velocity and position
vel = np.zeros(steps+1)
pos = np.zeros(steps+1)

# Setting constants and initial values for vel. and pos.
k = 0.1
m = 0.01
vel0 = 0.05
pos0 = 0.01
freqNatural = 10.0**0.5
b = 0.0
F0 = 0.01
Wd = 7.0
vel[0] = vel0    #Sets the initial velocity
pos[0] = pos0    #Sets the initial position



# Numerical solution using Euler's
# Splitting the ODE into two first order ones
# v'(t) = -(k/m)*x(t) - (b/m)*v(t) + (F0/m)*cos(Wd*t)
# x'(t) = v(t)
# Using the definition of the derivative we get
# (v(t+dT) - v(t))/dT on the left side of the first equation
# (x(t+dT) - x(t))/dT on the left side of the second 
# In the for loop t and dT will be replaced by i and 1

for i in range(0, steps):
    vel[i+1] = (-k/m)*dT*pos[i] + vel[i]*(1-dT*b/m) + (dT/m)*F0*np.cos(Wd*i)
    pos[i+1] = dT*vel[i] + pos[i]

# Ploting
#----------------
# With no damping
plt.plot(time, pos, 'g-', label='Undampened')

# Damping set to 10% of critical damping
b = (freqNatural/50)*0.1

# Using Euler's again to compute new values for new damping
for i in range(0, steps):
    vel[i+1] = (-k/m)*dT*pos[i] + vel[i]*(1-(dT*(b/m))) + (F0*dT/m)*np.cos(Wd*i)
    pos[i+1] = dT*vel[i] + pos[i]

plt.plot(time, pos, 'b-', label = '10% of crit. damping')
plt.plot(time, 0*time, 'k-')      # This plots the x-axis
plt.legend(loc = 'upper right')

#---------------
plt.show()

Answer 1

The problem here is with the term np.cos(Wd*i). It should be np.cos(Wd*i*dT), that is note that dT has been added into the correct equation, since t = i*dT.

If this correction is made, the simulation looks reasonable. Here's a version with F0=0.001. Note that the driving force is clear in the continued oscillations in the damped condition.

The problem with the original equation is that np.cos(Wd*i) just jumps randomly around the circle, rather than smoothly moving around the circle, causing no net effect in the end. This can be best seen by plotting it directly, but the easiest thing to do is run the original form with F0 very large. Below is F0 = 10 (ie, 10000x the value used in the correct equation), but using the incorrect form of the equation, and it's clear that the driving force here just adds noise as it randomly moves around the circle.

Answer 2

Note that your ODE is well behaved and has an analytical solution. So you could utilize sympy for an alternate approach:

import sympy as sy    
sy.init_printing()  # Pretty printer for IPython

t,k,m,b,F0,Wd = sy.symbols('t,k,m,b,F0,Wd', real=True)  # constants

consts = {k:  0.1, # values
          m:  0.01,
          b:  0.0,
          F0: 0.01,
          Wd: 7.0}

x = sy.Function('x')(t)  # declare variables
dx = sy.Derivative(x, t)
d2x = sy.Derivative(x, t, 2)

# the ODE:
ode1 = sy.Eq(m*d2x + b*dx + k*x, F0*sy.cos(Wd*t))

sl1 = sy.dsolve(ode1, x) # solve ODE
xs1 = sy.simplify(sl1.subs(consts)).rhs # substitute constants


# Examining the solution, we note C3 and C4 are superfluous
xs2 = xs1.subs({'C3':0, 'C4':0})
dxs2 = xs2.diff(t)

print("Solution x(t) = ")
print(xs2)
print("Solution x'(t) = ")
print(dxs2)

gives

Solution x(t) = 
C1*sin(3.16227766016838*t) + C2*cos(3.16227766016838*t) - 0.0256410256410256*cos(7.0*t)
Solution x'(t) = 
3.16227766016838*C1*cos(3.16227766016838*t) - 3.16227766016838*C2*sin(3.16227766016838*t) + 0.179487179487179*sin(7.0*t)

The constants C1,C2 can be determined by evaluating x(0),x'(0) for the initial conditions.