Search code examples
pythonnumpydifferential-equations

Numpy array not updating during a for loop


I am trying to code something to solve differential equations, here the case is rather simple, it's the famous harmonic oscillator, however I am trying to make a general code that would work for an ODE of order n.

During my for loop, the values in X are not updating and I checked that the member on the right side of the equation should change it. The initial value of X, meaning X[0] is [1, 0] and the right side of the X[k+1] is [1, -0.1] but when I print the value of X[k+1] it's still [0, 0], the value that was supposed to be replaced.

from matplotlib import pyplot as plt
import numpy as np


def deriv(X_k, omega):
    functions = [lambda _: X_k[1], lambda _: -omega**2*X_k[0]]
    X_dot_k = np.array([f(1) for f in functions])
    return X_dot_k


step = 0.1 # Time step
omega = 1 # Angular velocity
dimension = 2 # Important when there is n equations.
t0, tf, x0, v0 = 0, 10, 1, 0
t = np.linspace(t0, tf, int((tf-t0)/step) + 1)
X = np.asarray([[0 for _ in range(dimension)] for _ in t])
X[0] = [x0, v0] # Inital conditions.

for k in range(len(t)-1):
    X[k+1] = X[k] + deriv(X[k], omega)*step


plt.plot(t, X[:, 0], label="position")
plt.xlabel("time (s)")
plt.ylabel("position (AU)")
plt.title("Position in function of time.")
plt.show()

Solution

  • Your array consists of all integers. deriv always returns [0,-1] and the entire array rounds to [1,0].

    So, just make it:

    X = np.asarray([[0 for _ in range(dimension)] for _ in t]).astype(float)
    

    Or even better:

    X = np.zeros((len(t),dimension))