ode-solvers icon indicating copy to clipboard operation
ode-solvers copied to clipboard

Incorrect divergence in integration?

Open fixgoats opened this issue 9 months ago • 5 comments

Hi, I was attempting to model a system using ode-solvers, but the integration kept exploding on me. With Dopri5 I would always get a step size underflow around x=139, and Rk4 just outputs NaNs at around the same point Edit: I also checked what happens if I go all the way to Dop853 and it still gives a step size underflow at the same point as Dopri5. I assumed it must be some fault in the equations, but I tried rewriting the system in Python using Scipy, and that gave me perfectly reasonable results with no divergence. I'll share the exact code I've been running in Rust and Python so you can see both implementations are being given exactly the same equations and parameters.

The rust version:

use std::{
    f64::consts::PI,
    fs::{create_dir_all, File},
    io::{BufWriter, Write},
    path::Path,
};

use ode_solvers::{Dopri5, System, Vector4};

/// micro eV T^-1
const MU_B: f64 = 57.883818060;
/// ps^-1
const W_XX: f64 = 1e-7;
/// Polariton lifetime: ps
const TAU_P: f64 = 1e1;
/// Reservoir exciton lifetime: ps
const TAU_R: f64 = 3e2;
/// Same spin exciton interaction constant: micro eV
const ALPHA_1: f64 = 6e-1;
/// Opposite spin exciton interaction constant: micro eV
const ALPHA_2: f64 = -6e-2;
/// Spin relaxation probability rate: ps^-1
const W_S: f64 = 4e-2;
// const BETA: f64 = 2.2e5; // micron^2
// const DELTA_SMALL: f64 = 2.2e6; // micron^2
/// micron^2
const S: f64 = 3.2 * 3.2; // micron^2.
const X: f64 = 0.4;
const G: f64 = -0.364;
/// Kelvin
const T: f64 = 1.8;
/// micro eV K^-1
const K_B: f64 = 86.17333262;
const BETA: f64 = 1.0 / (K_B * T);
/// micro eV ps
const HBAR: f64 = 658.2119569;
/// micro eV
const DELTA: f64 = 1e4;
/// micron^2 micro eV^-1  ps^-2
const D: f64 = S * K_B * T / (2.0 * PI * HBAR * HBAR);

/// The system is contained in a 4D vector, the first entry is n_p^+, second is n_p^-, third is
/// n_r^+, fourth is n_r^-
type State = Vector4<f64>;

/// Struct to hold variable simulation parameters
struct ZeemanPolaritons {
    p: f64,
    b: f64,
}

impl System<f64, State> for ZeemanPolaritons {
    /// Here y[0] is n_p^+, y[1] is n_p^-, y[2] is n_r^+, y[3] is n_r^-.
    fn system(&self, _x: f64, y: &State, dy: &mut State) {
        // I set E_{r_0} to 0, it shouldn't make a difference for the simulation
        // E_r^+
        let erp = -0.5 * MU_B * G * self.b
            + ALPHA_1 * (y[2] + y[3] + X * y[0])
            + 2.0 * ALPHA_2 * X * y[1];
        // E_r^-
        let erm =
            0.5 * MU_B * G * self.b + ALPHA_1 * (y[2] + y[3] + X * y[1]) + 2.0 * ALPHA_2 * X * y[0];
        // E^+
        let ep = -DELTA - 0.5 * X * MU_B * G * self.b
            + ALPHA_1 * X * (y[2] + X * y[0])
            + X * ALPHA_2 * (y[3] + X * y[1]);
        // E^-
        let em = -DELTA
            + 0.5 * X * MU_B * G * self.b
            + ALPHA_1 * X * (y[3] + X * y[1])
            + X * ALPHA_2 * (y[2] + X * y[0]);

        let er_diff_plus_exp = (BETA * (erp - erm)).exp();
        let er_diff_minus_exp = (BETA * (erm - erp)).exp();
        dy[0] = W_XX * (y[2].powi(2) * (y[0] + 1.0) - D * y[2] * y[0] * (BETA * (ep - erp)).exp())
            - y[0] / TAU_P;
        dy[1] = W_XX * (y[3].powi(2) * (y[1] + 1.0) - D * y[3] * y[1] * (BETA * (em - erm)).exp())
            - y[1] / TAU_P;
        dy[2] = W_XX * (D * y[2] * y[0] * (BETA * (ep - erp)).exp() - y[2].powi(2) * (y[0] + 1.0))
            + W_S
                * (er_diff_plus_exp / (1.0 + er_diff_plus_exp))
                * (y[3] * er_diff_plus_exp - y[2])
            - y[2] / TAU_R
            + self.p;
        dy[3] = W_XX * (D * y[3] * y[1] * (BETA * (em - erm)).exp() - y[3].powi(2) * (y[1] + 1.0))
            - W_S
                * (er_diff_minus_exp / (1.0 + er_diff_minus_exp))
                * (y[3] * er_diff_minus_exp - y[2])
            - y[3] / TAU_R
            + self.p;
    }
}

fn save(states: &[State], filename: &Path) {
    if let Some(dir) = filename.parent() {
        if let Err(e) = create_dir_all(dir) {
            println!("Could not create directory. Error: {:?}", e);
            return;
        }
    }
    let file = match File::create(filename) {
        Err(e) => {
            println!("Could not open file. Error: {:?}", e);
            return;
        }
        Ok(buf) => buf,
    };
    let mut buf = BufWriter::new(file);

    // Write state vector in a csv format
    for state in states.iter() {
        for val in state.iter() {
            buf.write_fmt(format_args!("{}, ", val)).unwrap();
        }
        buf.write_fmt(format_args!("\n")).unwrap();
    }
    if let Err(e) = buf.flush() {
        println!("Could not write to file. Error: {:?}", e);
    }
}

fn main() {
    let y0 = State::new(0.0, 0.0, 0.0, 0.0);

    let system = ZeemanPolaritons { p: 9.0, b: 3.5 };
    let mut solver = Dopri5::new(system, 0.0, 200., 1e-1, y0, 0.001, 1e-6);
    let res = solver.integrate();
    match res {
        Ok(stats) => {
            println!("{}", stats);
            let path = Path::new("./outputs/zeeman_polaritons.dat");
            save(solver.y_out(), &path);
        }
        Err(e) => println!("Error: {}", e),
    }
}

And the Python version:

from scipy.integrate import RK45
import numpy as np

mub = 57.883818060
wxx = 1e-7
taup = 1e1
taur = 3e2
alpha1 = 6e-1
alpha2 = -6e-2
ws = 4e-2
S = 3.2 * 3.2
x = 0.4
g = -0.364
T = 1.8
kB = 86.17333262
beta = 1 / (kB * T)
Delta = 1e4
hbar = 658.2119
D = S * kB * T / (2 * np.pi * hbar**2)
b = 3.5
p = 9


def f(t, y):
    erp = (
        -0.5 * mub * g * b + alpha1 * (y[2] + y[3] + x * y[0]) + 2.0 * alpha2 * x * y[1]
    )
    erm = (
        0.5 * mub * g * b + alpha1 * (y[2] + y[3] + x * y[1]) + 2.0 * alpha2 * x * y[0]
    )
    ep = (
        -Delta
        - 0.5 * x * mub * g * b
        + alpha1 * x * (y[2] + x * y[0])
        + alpha2 * x * (y[3] + x * y[1])
    )
    em = (
        -Delta
        + 0.5 * x * mub * g * b
        + alpha1 * x * (y[3] + x * y[1])
        + alpha2 * x * (y[2] + x * y[0])
    )
    er_diff_plus_exp = np.exp(beta * (erp - erm))
    er_diff_minus_exp = np.exp(beta * (erm - erp))
    return np.array(
        [
            wxx
            * (y[2] ** 2 * (y[0] + 1.0) - D * y[2] * y[0] * np.exp(beta * (ep - erp)))
            - y[0] / taup,
            wxx
            * (y[3] ** 2 * (y[1] + 1.0) - D * y[3] * y[1] * np.exp(beta * (em - erm)))
            - y[1] / taup,
            wxx
            * (D * y[2] * y[0] * np.exp(beta * (ep - erp)) - y[2] ** 2 * (y[0] + 1.0))
            + ws
            * (er_diff_plus_exp / (1.0 + er_diff_plus_exp))
            * (y[3] * er_diff_plus_exp - y[2])
            - y[2] / taur
            + p,
            wxx
            * (D * y[3] * y[1] * np.exp(beta * (em - erm)) - y[2] ** 2 * (y[0] + 1.0))
            - ws
            * (er_diff_minus_exp / (1.0 + er_diff_minus_exp))
            * (y[3] * er_diff_minus_exp - y[2])
            - y[3] / taur
            + p,
        ]
    )


solver = RK45(f, 0, np.array([0, 0, 0, 0]), 80000)
while solver.status == "running":
    solver.step()

if solver.status == "failed":
    print("integration failed")
else:
    np.save("outputs/zeeman_polaritons.npy", solver.y)

The fact that Scipy handles this integration with no problems while ode-solvers doesn't leads me to believe there is some issue with ode-solvers.

fixgoats avatar Apr 12 '25 14:04 fixgoats

Does version 0.5 integrate the problem correctly? I also seemed to notice a regression somewhere between 0.5 and 0.6.1.

joshburkart avatar Apr 13 '25 20:04 joshburkart

Thank you for reporting this, I need to look into it. If version 0.5 does integrate the problem correctly, please stick with it for now.

srenevey avatar Apr 14 '25 00:04 srenevey

Tried downgrading to 0.5, it also has the same error.

fixgoats avatar Apr 14 '25 02:04 fixgoats

For what it's worth, I tried implementing this with Boost's odeint in C++ as well and its explicit adaptive RK solver also couldn't handle the problem. So the real question might be, how is it that Scipy's RK45 and Matlab's ode integrator can handle the problem when these methods, that should be practically identical, can't.

fixgoats avatar Apr 18 '25 22:04 fixgoats

So it turns out I had some parameters wrong and I didn't inspect the scipy solution closely enough to notice it's actually a bit suspect. So whether the somewhat strange scipy solution that at least doesn't have NaNs, or the Rust and C++ versions that return NaNs are the better solutions, I'm not sure. With the correct parameters the system behaves the same with any method.

fixgoats avatar May 12 '25 16:05 fixgoats