Sunode error: When changing to a smaller dtype, its size must be a divisor of the size of original dtype

Questions
1

Hello,

I am trying to use the sunode ODE solver instead of the DifferentialEquations solver, but I get the error “When changing to a smaller dtype, its size must be a divisor of the size of original dtype” in doing so. Any pointers or suggestions?

import sunode
import sunode.wrappers.as_theano

def SIR_sunode(t, y, p):
    """Right hand side of Lotka-Volterra equation.

    All inputs are dataclasses of sympy variables, or in the case
    of non-scalar variables numpy arrays of sympy variables.
    """
    return {
        'S': -p.lam * y.S * y.I,
        'I': p.lam * y.S * y.I - p.mu * y.I,
    }


with pm.Model() as model:
    sigma = pm.HalfCauchy('sigma', 2, shape=2)
    S_init = yobs[0,0]
    I_init = yobs[0,1]
    lam = pm.Lognormal('lambda', np.log(4), 1.5)
    mu = pm.Lognormal('mu', np.log(1), 1)
    
    #sir_curves = sir_model(y0=[0.99, 0.01], theta=[lam, mu])
    
    sir_curves, _, problem, solver, _, _ = sunode.wrappers.as_theano.solve_ivp(
        y0={
	    # The initial conditions of the ode. Each variable
	    # needs to specify a theano or numpy variable and a shape.
	    # This dict can be nested.
            'S': (S_init, ()),
            'I': (I_init, ()),
        },
        params={
	    # Each parameter of the ode. sunode will only compute derivatives
	    # with respect to theano variables. The shape needs to be specified
	    # as well. It it infered automatically for numpy variables.
	    # This dict can be nested.
            'lam': (lam, ()),
            'mu': (mu, ()),
        },
	# A functions that computes the right-hand-side of the ode using
	# sympy variables.
        rhs=SIR_sunode,
	# The time points where we want to access the solution
        tvals=times,
        t0=times[0],
    )

    S = pm.Lognormal('S', mu=pm.math.log(sir_curves['S']), sigma=sigma, observed=yobs[:,0])
    I = pm.Lognormal('I', mu=pm.math.log(sir_curves['I']), sigma=sigma, observed=yobs[:,1])
    step = pm.Metropolis()
    trace = pm.sample(400, tune=100, cores=2)
    data = az.from_pymc3(trace=trace)
2

paging @aseyboldt

3

Turns out that is a bug. I’ll fix it in the next release, but to work around it you can add any fixed parameter (no need to use it). So

        params={
	    # Each parameter of the ode. sunode will only compute derivatives
	    # with respect to theano variables. The shape needs to be specified
	    # as well. It it infered automatically for numpy variables.
	    # This dict can be nested.
            'lam': (lam, ()),
            'mu': (mu, ()),
            '_dummy': (np.array(1.), ()),
        },

You probably also want to change the likelihood a bit:

    S = pm.Lognormal('S', mu=pm.math.log(sir_curves['S']), sigma=sigma, observed=yobs[:,0])

Taking the log of the sir doesn’t make much sense for a lognormal likelihood I think, and you need sigma[0].

3 Likes
4

Thank you!

That fix worked and you are right about ‘sigma’. As for the log part, that did look a little odd to me but I am using this example from the PyMC3 page https://docs.pymc.io/notebooks/ODE_API_introduction.html.

Regards

5

Looking back at that page, I realized that the reason for the log was

y = odeint(SIR_non_normalized, t=times, y0=[S0, I0], args=((lam, mu),), rtol=1e-8)
yobs = np.random.lognormal(loc=np.log(y[1::]), scale=[0.2, 0.3])

This seems to help numerically as well…