Untitled

@nb.njit
def lv_drift_ode(x: np.ndarray, theta: np.ndarray):
    k1, k2, k3 = theta
    x1, x2 = x
    mu = np.zeros(2)
    mu[0] = k1 * x1 - k1 / 4 - k2 * x1 * x2 + k2 * x1 / 4 - k2 * x2 / 4
    mu[1] = k2 * x1 * x2 - k2 * x1 / 4 + k2 * x2 / 4 - k3 * x2 - k3 / 4
    return mu


@nb.njit
def lv_diffusion_ode(x: np.ndarray, theta: np.ndarray, noise: np.ndarray):
    k1, k2, k3 = theta
    x1, x2 = x
    sig = np.zeros(2)
    sig[0] = sqrt(k1 * x1) * noise[0] - sqrt(k2 * x1 * x2) * noise[1]
    sig[1] = sqrt(k2 * x1 * x2) * noise[1] - sqrt(k3 * x2) * noise[2]
    return sig
    
@nb.njit
def lv_splitting(
    x0: np.ndarray,
    n: int,
    A: int,
    dt: float,
    drift_ode: Callable,
    diffusion_ode: Callable,
    theta: np.ndarray,
):
    # Timestep.
    h = dt / A
    sq3 = sqrt(3)

    # Allocate.
    x = np.zeros((1, 2, n * A + 1), dtype=np.float32)
    xaux = np.zeros(2, dtype=np.float32)
    W = np.zeros(3, dtype=np.float32)
    H = np.zeros(3, dtype=np.float32)
    noise = np.zeros(3, dtype=np.float32)

    # Solve.
    x[0, :, 0] = x0
    for i in range(n * A):

        # Generate noise.
        W[:] = np.random.normal(loc=0, scale=h, size=3)
        H[:] = np.random.normal(loc=0, scale=h / 12, size=3)

        # First ODE.
        mu = drift_ode(x[0, :, i], theta)
        xaux[:] = x[0, :, i] + mu * (3 - sq3) * h / 6

        # Second ODE.
        noise[:] = (1 / 2) * W + sq3 * H
        sig = diffusion_ode(xaux, theta, noise)
        xaux[:] = xaux[:] + sig

        # Third ODE.
        mu = drift_ode(xaux, theta)
        xaux[:] = xaux + mu * sq3 * h / 3

        # Fourth ODE.
        noise[:] = (1 / 2) * W - sq3 * H
        sig = diffusion_ode(xaux, theta, noise)
        xaux[:] = xaux[:] + sig

        # Fifth ODE.
        mu = drift_ode(xaux, theta)
        x[0, :, i + 1] = xaux + mu * (3 - sq3) * h / 6
    return x