Mathematica's NDSolve and SciPy's solve_ivp returning different results

Question

This question involves both Python and Mathematica code.

I want to solve a system of complex ODEs numerically. Originally, I used NDSolve, but since I ended up needing to do an extensive parameter search, I decided to switch to using Python, specifically SciPy's solve_ivp function. However, I quickly ran into the following problem:

A, B, C, and D are the four coupled complex ODEs I want to solve for. Notably, abs(A)^2 + abs(B)^2 + abs(C)^2 + abs(D)^2 = 1 always holds (due to the nature of the problem). Indeed, this is always true in NDSolve's solution

.

However, when I transferred the code to Python, the result using solve_ivp is incorrect.

At first, I was sure that there was a typo when I transferred the code to Mathematica. However, I've spent a couple of days looking over the functions, and haven't found anything. I also had a friend familiar with both languages look it over, and they did not find anything either (but that was a slightly different version of the code then the one presented here, which I simplified to make more readable).

Therefore, my question is if there is something in the definition of either NDSolve or solve_ivp that might cause this difference, or if there is a way to reproduce Mathematica's result in using Python.

Mathematica Code:

r[t_, c_, b_] := Sqrt[(c*t)^2 + b^2]
Sol[t0_, tf_, C1_, C2_, a_, c_, b_, d_, \[Phi]1_, \[Phi]2_] := 
 Module[{g}, g = C1*C2 ; 
  NDSolve[{A'[
      t] == -I (A[
          t]*(C1*a/2 + C2*a*UnitStep[d - r[t, c, b]]/2 + 
           g/r[t, c, b]^3) - 
        3*g/r[t, c, b]^5*C11[t]*((c*t)^2 - b^2 - 2*I*c*t*b)), 
    A[t0] == Cos[\[Phi]1]/Sqrt[2],
    
    B'[t] == -I*(D1[
          t]*(2*g/r[t, c, b]^3 - 3*g/r[t, c, b]^5*((c*t)^2 + b^2)) + 
        B[t]*(-g/r[t, c, b]^3 + C1*a/2 - 
           C2*a*UnitStep[d - r[t, c, b]]/2)), 
    B[t0] == Sin[\[Phi]1]/Sqrt[2], 
    C11'[t] ==
     -I*(C11[
          t]*(-C1*a/2 - C2*a*UnitStep[d - r[t, c, b]]/2 + 
           g/r[t, c, b]^3) - 
        3*g/r[t, c, b]^5*A[t]*((c*t)^2 - b^2 + 2*I*c*t*b)), 
    C11[t0] == Sin[\[Phi]1]*E^(-I*\[Phi]2)/Sqrt[2], D1'[t] ==
     
     -I*(B[t]*(2*g/r[t, c, b]^3 - 3*g/r[t, c, b]^5*((c*t)^2 + b^2)) + 
        D1[t]*(-g/r[t, c, b]^3 - C1*a/2 + 
           C2*a*UnitStep[d - r[t, c, b]]/2)), 
    D1[t0] == Cos[\[Phi]1]*E^(-I*\[Phi]2)/Sqrt[2]}, {A, B, C11, 
    D1}, {t, t0, tf}]]

To reproduce result:

S = Sol[-100, 100, 1, 1.25, 100, 1, 1, 10, 0, 0];
Rasterize[
 Plot[{Abs[A[t]]^2 /. S, Abs[B[t]]^2 /. S, Abs[C11[t]]^2 /. S, 
   Abs[D1[t]]^2 /. S, 
   Abs[A[t]]^2 + Abs[B[t]]^2 + Abs[C11[t]]^2 + Abs[D1[t]]^2 /. 
    S}, {t, -100, 100}, 
  PlotLegends -> {"A^2", "B^2", "C^2", "D^2", "Sum=1"}, 
  PlotLabel -> "Working Version with Mathematica NDSolve"]]

Python Code:

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


def r(t, c, b):
    return np.sqrt((c*t)**2 + b**2)


def coeff(t, z, k):
    C1, C2, a, c, b, d = k
    g = C1*C2

    A, B, C, D = z

    dAdt = -1j * (A * (C1*a/2 + C2*a*np.heaviside(d-r(t,c,b), 1)/2 +
    g/r(t,c,b)**3) - 3*g/r(t,c,b)**5 * C * ((c * t)**2 - b**2 - 2*1j*c*t*b))

    dBdt = -1j*(D*(2*g/r(t,c,b)**3 - 3*g/r(t,c,b)**5*((c*t)**2 + b**2)) + B *(-g/r(t,c,b)**3 + C1*a/2 - C2*a*np.heaviside(d-r(t,c,b), 1)/2))

    dCdt = -1j * (C*(-C1*a/2 - C2*a*np.heaviside(d-r(t,c,b), 1)/2 +
    g/r(t,c,b)**3) - 3*g/r(t,c,b)**5*A*((c*t)**2 - b**2 + 2*1j*c*t*b))

    dDdt = -1j*(B*(2*g/r(t,c,b)**3- 3*g/r(t,c,b)**5*((c*t)**2 + b**2)) + D *(-g/r(t,c,b)**3 - C1*a/2 + C2*a*np.heaviside(d-r(t,c,b), 1)/2))

    return [dAdt, dBdt, dCdt, dDdt]

def solver(t0, tf, k, phi1, phi2):

    a0 = np.cos(phi1)/np.sqrt(2)

    b0 = np.sin(phi1)/np.sqrt(2)


    c0 = (np.sin(phi1)/np.sqrt(2))*np.exp(-1j*phi2)

    d0 = (np.cos(phi1)/np.sqrt(2))*np.exp(-1j*phi2)

    z0 = (a0, b0, c0, d0)

    return solve_ivp(coeff, [t0, tf], z0, args=(k,))

To reproduce result:

k = [1, 1.25, 100, 1, 1, 10]
sol=solver(-100, 100, k, 0, 0)

plt.plot(sol.t, np.abs(sol.y[0])**2, label="A^2")
plt.plot(sol.t, np.abs(sol.y[1])**2, label="B^2")
plt.plot(sol.t, np.abs(sol.y[2])**2, label = "C^2")
plt.plot(sol.t, np.abs(sol.y[3])**2, label="D^2")
plt.plot(sol.t, np.abs(sol.y[0])**2+np.abs(sol.y[1])**2+np.abs(sol.y[2])**2+np.abs(sol.y[3])**2, label="Sum != 1???")
plt.legend()
plt.title("Not Working solve_ivp Version")
plt.show()

Edit:

I tried splitting the ODE so that I could remove the discontinuity due to the Heaviside Function. It definitely made things better, but not by much:What I think is most interesting is the shift at t=0, as that is not a discontinuity. Even before that point, the function is shifting away from Sum=1, as it did before.

New Code:

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


def r(t, c, b):
    return np.sqrt((c*t)**2 + b**2)


def coeff(t, z, k):
    C1, C2, a, c, b, d = k
    g = C1*C2

    A, B, C, D = z

    dAdt = -1j * (A * (C1*a/2 + C2*a/2 +
    g/r(t,c,b)**3) - 3*g/r(t,c,b)**5 * C * ((c * t)**2 - b**2 - 2*1j*c*t*b))

    dBdt = -1j*(D*(2*g/r(t,c,b)**3 - 3*g/r(t,c,b)**5*((c*t)**2 + b**2)) + B *(-g/r(t,c,b)**3 + C1*a/2 - C2*a/2))

    dCdt = -1j * (C*(-C1*a/2 - C2*a/2 +
    g/r(t,c,b)**3) - 3*g/r(t,c,b)**5*A*((c*t)**2 - b**2 + 2*1j*c*t*b))

    dDdt = -1j*(B*(2*g/r(t,c,b)**3- 3*g/r(t,c,b)**5*((c*t)**2 + b**2)) + D *(-g/r(t,c,b)**3 - C1*a/2 + C2*a))

    return [dAdt, dBdt, dCdt, dDdt]

def coeff2(t, z, k):
    C1, C2, a, c, b, d = k
    g = C1*C2

    A, B, C, D = z

    dAdt = -1j * (A * (C1*a/2 +
    g/r(t,c,b)**3) - 3*g/r(t,c,b)**5 * C * ((c * t)**2 - b**2 - 2*1j*c*t*b))

    dBdt = -1j*(D*(2*g/r(t,c,b)**3 - 3*g/r(t,c,b)**5*((c*t)**2 + b**2)) + B *(-g/r(t,c,b)**3 + C1*a/2))

    dCdt = -1j * (C*(-C1*a/2 +
    g/r(t,c,b)**3) - 3*g/r(t,c,b)**5*A*((c*t)**2 - b**2 + 2*1j*c*t*b))

    dDdt = -1j*(B*(2*g/r(t,c,b)**3- 3*g/r(t,c,b)**5*((c*t)**2 + b**2)) + D *(-g/r(t,c,b)**3 - C1*a/2))

    return [dAdt, dBdt, dCdt, dDdt]

def solver(t0, tf, k, phi1, phi2, order=1):

    a0 = np.cos(phi1)/np.sqrt(2)

    b0 = np.sin(phi1)/np.sqrt(2)


    c0 = (np.sin(phi1)/np.sqrt(2))*np.exp(-1j*phi2)

    d0 = (np.cos(phi1)/np.sqrt(2))*np.exp(-1j*phi2)

    z0 = (a0, b0, c0, d0)
    if order == 2 :

        return solve_ivp(coeff, [t0, tf], z0, args=(k,))
    else:
        return solve_ivp(coeff2, [t0, tf], z0, args=(k,))


k = [1, 1.25, 100, 1, 1, 10]

t0=-100
tf=100

t1 = -np.sqrt(k[-1]**2 - k[-2]**2)/k[-4] - 0.00000001
t2 = np.sqrt(k[-1]**2 - k[-2]**2)/k[-4] + 0.000000001

sol1 = solver(t0, t1, k, 0, 0)
sol2 = solve_ivp(coeff, [t1,t2], [sol1.y[0][-1], sol1.y[1][-1], sol1.y[2][-1], sol1.y[3][-1]], args=(k,))
sol3 = solve_ivp(coeff2, [t2,tf], [sol2.y[0][-1], sol2.y[1][-1], sol2.y[2][-1], sol2.y[3][-1]], args=(k,))

t = np.concatenate((sol1.t,sol2.t, sol3.t), axis=None)
A = np.concatenate((sol1.y[0],sol2.y[0],sol3.y[0]), axis=None)
B = np.concatenate((sol1.y[1],sol2.y[1],sol3.y[1]), axis=None)
C = np.concatenate((sol1.y[2],sol2.y[2],sol3.y[2]), axis=None)
D = np.concatenate((sol1.y[3],sol2.y[3],sol3.y[3]), axis=None)


plt.plot(t, np.abs(A)**2, label="A^2")
plt.plot(t, np.abs(B)**2, label="B^2")
plt.plot(t, np.abs(C)**2, label = "C^2")
plt.plot(t, np.abs(D)**2, label="D^2")
plt.plot(t, np.abs(A)**2+np.abs(B)**2+np.abs(C)**2+np.abs(D)**2, label="Sum != 1???")
plt.legend()
plt.title("Not Working solve_ivp Version")
print(np.max(np.abs(A)**2+np.abs(B)**2+np.abs(C)**2+np.abs(D)**2))
plt.show()

Answer 1

Mathematica is smart enough to figure out what type of solver is most accurate in most cases; with scipy.integrate.solve_ivp you have to specify the method parameter if the default is not good enough. Unfortunately, not all solvers for solve_ivp work well with complex-valued problems and the other solvers that support complex values behave much worse for your ODE than the default (RK45).

You can find out what solver Mathematica is using under the hood.

Your derivatives include np.heaviside functions, which are discontinuous; this is something that ODE solvers generally don't like; for example RK45 is designed for functions that can be approximated locally by a functions with continuous 4th-order derivatives. Consider figuring out where the jumps in the derivatives are and integrating in segments between them. For example, if you want to solve over the range t=0..1 and the discontinuities are at t=a and t=b; then integrate from 0 to a-1e-9, then from a+1e-9 to b-1e-9, then from b+1e-9 to 1. The plot of the bad solution suggests that there are problems in the derivatives at a few discrete points.

EDIT: upon closer inspection, the factors 1/r(t,c,b)**3 and **5 are sharply peaked at t=0 (increasing by several orders of magnitude); it could be that the ODE becomes 'stiff' there. (Innocuous ODEs like d(A, B)/dt = (-aB, -bA) with positive a, b tend to be stiff). Splitting the variables into real and imaginary parts so that you can use a solver suitable for real-valued stiff equations may be worth a try.