Codim-3 canonical model of Saggio et al 2020 https://elifesciences.org/articles/55632

codim3.py

from scipy.integrate import ode
from numba import njit
import numba as nb
from NumbaLSODA import lsoda_sig, lsoda
from numba import njit, cfunc

@njit
def Parametrization_VectorsGreatCircle(p1,p2,R):
    E = p1 / R
    F = np.cross(np.cross(p1,p2),p1)
    F = F / np.linalg.norm(F)
    return E, F

@njit
def Parametrization_GreatCircle(E,F,R,z):
    mu2=R*(E[0]*np.cos(z)+F[0]*np.sin(z))
    mu1=-R*(E[1]*np.cos(z)+F[1]*np.sin(z))
    nu=R*(E[2]*np.cos(z)+F[2]*np.sin(z))
    return np.complex128(mu2), np.complex128(mu1), np.complex128(nu)

"""
    else:
        raise ValueError(f'N must be 1, 2, or 3, not {N}')
    if not np.isfinite(x_rs):
        # print(eval_resting_state_cartesian(*))
        raise ValueError(f'x_rs not finite: {x_rs}, {mu2,mu1,nu,N}')
"""

@njit
def eval_resting_state_cartesian(mu2,mu1,nu,N):
    if N == 1: # resting state
        x_rs=mu2 / (3*(mu1 / 2 + (mu1 ** 2 / 4 - mu2 ** 3 / 27) ** (1/2)) ** (1/3)) \
                 + (mu1 / 2 + (mu1 ** 2 / 4 - mu2 ** 3 / 27) ** (1/2)) ** (1/3)
    elif N == 2:
        x_rs=- mu2/(6*(mu1/2 + (mu1**2/4 - mu2**3/27)**(1/2))**(1/3)) - (mu1/2 + (mu1**2/4 - mu2**3/27)**(1/2))**(1/3)/2 - (3**(1/2)*1j*(mu2/(3*(mu1/2 + (mu1**2/4 - mu2**3/27)**(1/2))**(1/3)) - (mu1/2 + (mu1**2/4 - mu2**3/27)**(1/2))**(1/3)))/2 
    elif N == 3:
        x_rs= (3**(1/2)*1j*(mu2/(3*(mu1/2 + (mu1**2/4 - mu2**3/27)**(1/2))**(1/3)) - (mu1/2 + (mu1**2/4 - mu2**3/27)**(1/2))**(1/3)))/2 - (mu1/2 + (mu1**2/4 - mu2**3/27)**(1/2))**(1/3)/2 - mu2/(6*(mu1/2 + (mu1**2/4 - mu2**3/27)**(1/2))**(1/3))
    else:
        x_rs = 0.0
    return x_rs

@njit
def codim3(t,x,c,R,dstar,E,F,N,out):
    # for a given value of z (x(3)), gives position along the path
    mu2, mu1, nu = Parametrization_GreatCircle(E,F,R,x[2]);
    # for a given position on the map, gives values of resting (or silent) state
    x_rs=np.real(eval_resting_state_cartesian(mu2,mu1,nu,N))
    # System
    out[0] = x1dot = np.real(- x[1])
    out[1] = x2dot = np.real(-( -x[0]**3 +mu2*x[0] +mu1 + x[1]*( nu + x[0] + x[0]**2)))
    out[2] = zdot =  np.real(-c*(np.sqrt((x[0] - x_rs)**2 + x[1]**2)-dstar))
    # return np.r_[x1dot,x2dot,zdot]

out = np.zeros((3,))
def codim3_(t,x,c,R,dstar,E,F,N):
    codim3(t,x,c,R,dstar,E,F,N,out)
    return out
  
b = 1.0;                            # focus
R = 0.4;                            # radius in unfolding
c = 0.003;                           # velocity of slow variable
dstar =0.3;                        # threshold for slow variable inversion
N=1;
ode_maxstep = 0.1;
tspan = np.r_[:3000:0.5]
x0 = np.r_[0., 0., 0.]

# SN/SH  
A=np.r_[0.3448,0.02285,0.2014]
B=np.r_[0.3351,0.07465,0.2053]
c=0.0001
tspan = np.r_[:1e4:0.5]

E, F = Parametrization_VectorsGreatCircle(A,B,R) 

def make_lsoda_func(R,dstar,N):
    @cfunc(lsoda_sig)
    def rhs(t, x, du, p):
        cef = nb.carray(p, (7,))
        codim3(t,x,cef[0],R,dstar,cef[1:4],cef[4:],N,du)
    return rhs
  
lf = make_lsoda_func(R,dstar,N)

t = np.r_[:1e4:1.0]
E, F = Parametrization_VectorsGreatCircle(A,B,R)
p = np.r_[c, E, F]
usol, success = lsoda(lf.address, x0, t, data=p)
assert success

codim3.stan

functions {

    vector cross(vector a, vector b) {
        return [-a[3]*b[2]+a[2]*b[3], a[3]*b[1]-a[1]*b[3], -a[2]*b[1]+a[1]*b[2]]';
    }
    
    real vnorm(vector x) {
        return sqrt(sum(square(x)));
    }

    vector vectors_great_circle(vector p1, vector p2, real r) {
        vector[3] e = p1 / r;
        vector[3] f = cross(cross(p1, p2), p1);
        f /= vnorm(f);
        return append_row(e,f);
    }
    
    vector great_circle(vector ef, real r, real z) {
        real mu2 = r*(ef[1]*cos(z)+ef[4]*sin(z));
        real mu1 = -r*(ef[2]*cos(z)+ef[5]*sin(z));
        real nu = r*(ef[3]*cos(z)+ef[6]*sin(z));
        return [mu2, mu1, nu]';
    }
    
    real cube(real x) { return x * x * x; }
    
    complex ccube(complex x) { return x * x * x; }
    
    complex csquare(complex x) { return x * x; }
    
    real eval_rs_N1(real rmu2, real rmu1) {
        complex mu2 = rmu2;
        complex mu1 = rmu1;
        complex sr1 = sqrt(csquare(mu1) / 4.0 - ccube(mu2) / 27.0); // , 1.0/2.0);
        complex cr1 = pow(mu1 / 2 + sr1, 1.0/3.0);
        complex xrs = mu2 / (3*cr1) + cr1;
        return get_real(xrs);
    }
    
    vector codim3(real t, vector x, real c, vector ef) {
        real r = 0.4;
        real b = 1.0;
        real dstar = 0.3;
        // for a given value of z (x(3)), gives position along the path
        vector[3] mmn = great_circle(ef,r,x[3]);
        real mu2 = mmn[1];
        real mu1 = mmn[2];
        real nu = mmn[3];
        // for a given position on the map, gives values of resting (or silent) state
        real x_rs = eval_rs_N1(mu2,mu1);
        // system
        real x1dot = -x[2];
        real x2dot = -( -cube(x[1]) +mu2*x[1] +mu1 + x[2]*( nu + x[1] + square(x[1])));
        real zdot =  -c*(sqrt( square(x[1] - x_rs) + square(x[2]) ) - dstar);
        return [x1dot, x2dot, zdot]';
    }
}

data {
    vector[3] x0;
    int nt;
    real ts[nt];
    real c;
    vector[3] a;
    vector[3] b;
    int integrator;
}

transformed data {
    real t0 = 0;
    real r = 0.4;
    vector[6] ef = vectors_great_circle(a, b, r);
    real rel_tol_forward = 1e-6;
    vector[3] abs_tol_forward = rep_vector(1e-6, 3);
    real rel_tol_backward = 1e-6;
    vector[3] abs_tol_backward = rep_vector(1e-6, 3);
    real rel_tol_quad = 1e-5;
    real abs_tol_quad = 1e-5;
    int max_num_steps = 100000;
    int num_steps_between_checkpoints = 100;
    int interpolation_polynomial = 1;
    int solver_forward = 1;
    int solver_backward = 1;
    
    vector[3] ys[nt] = ode_rk45(codim3, x0, t0, ts, c, ef);
    print("done integrating true solution");
}

parameters {
    vector[3] ah;
}

transformed parameters {
    print("ah", ah);
    vector[6] efh = vectors_great_circle(ah, b, r);
    vector[3] yh[nt];
    if (integrator == 1)
        yh = ode_rk45(codim3, x0, t0, ts, c, efh);
    if (integrator == 2)
        yh = ode_ckrk(codim3, x0, t0, ts, c, efh);
    if (integrator == 3)
        yh = ode_adams(codim3, x0, t0, ts, c, efh);
    if (integrator == 4)
        yh = ode_bdf(codim3, x0, t0, ts, c, efh);
    if (integrator == 5)
        yh = ode_adjoint_tol_ctl(codim3, x0, t0, ts,
            rel_tol_forward, abs_tol_forward, rel_tol_backward,
            abs_tol_backward, rel_tol_quad, abs_tol_quad,
            max_num_steps, num_steps_between_checkpoints,
            interpolation_polynomial, solver_forward, solver_backward,
            c, efh);
}

model {
    ah ~ normal(0.1,2);
    for (t in 1:nt)
        ys[t][1] ~ normal(yh[t][1], 0.02);
}

sbi_codim3.py

import torch
from sbi import utils as utils
from sbi import analysis as analysis
from sbi.inference.base import infer
from codim3 import *

nd = 11
prior = utils.BoxUniform(low=-torch.ones(nd), high=torch.ones(nd))

def simulator(p):
    p = np.asarray(p)
    x0, A, B, c, T = p[:3], p[3:6], p[6:9], p[9], p[10]
    t = np.r_[:T:1]
    E, F = Parametrization_VectorsGreatCircle(A,B,R)
    usol, success = lsoda(lf.address, x0, t, data=np.r_[c, E, F])
    assert success
    return torch.as_tensor(usol[4000:7000:10,0])

# x0, A, B, c, T
# sn_sh = np.r_[0,0,0, 0.3448,0.02285,0.2014, 0.3351,0.07465,0.2053, 0.0001, 1e4]
# sn_sn = np.r_[0,0,0, 0.2649,-0.05246,0.2951, 0.2688,0.05363,0.2914, 0.002, 1e2]

posterior = infer(simulator, prior, method='SNPE', num_simulations=100)
# RuntimeError: There are 0 sized dimensions, and they aren't selected, so unique cannot be applied