Ordinary differential equations ODE codes working in version 5.6.1 but no longer works in version 5.16

Questions v5
1

Hi all!

First of all, thank you very much for all of your help. It was a pleasure to use PyMC. This is my first post so please excuse me if I might not ask an interesting question.

I am using PyMC version 5.6.1 and I have a project using Ordinary Differential Equations and PyMC.

Essentially, my codes is this model:

mcmc_ode = DifferentialEquation(
    func=ode_model,
    times=tspan,
    n_states=26,
    n_theta=37,
    t0=t0,
)

I am using
from pymc.ode import DifferentialEquation
and the solver is
from scipy.integrate import odeint

  • which was specified inside ode_model

I have successfully run this project for maybe thousands times already (it is fairly complicated) and did not have any problems on my Macbook M1 on PyMC version 5.6.1

Today, I decided that it is time to upgrade to version 5.16.2 but I have this weird error here:

File ~/opt/anaconda3/envs/pymc5_16_env/lib/python3.11/site-packages/pytensor/link/c/cmodule.py:2654 in compile_str
    raise CompileError(

CompileError: Compilation failed (return status=1):
/Users/ducquangnguyen/opt/anaconda3/envs/pymc5_16_env/bin/clang++ -dynamiclib -g -O3 -fno-math-errno -Wno-unused-label -Wno-unused-variable -Wno-write-strings -Wno-c++11-narrowing -fno-exceptions -fno-unwind-tables -fno-asynchronous-unwind-tables -DNPY_NO_DEPRECATED_API=NPY_1_7_API_VERSION -fPIC -undefined dynamic_lookup -I/Users/ducquangnguyen/opt/anaconda3/envs/pymc5_16_env/lib/python3.11/site-packages/numpy/core/include -I/Users/ducquangnguyen/opt/anaconda3/envs/pymc5_16_env/include/python3.11 -I/Users/ducquangnguyen/opt/anaconda3/envs/pymc5_16_env/lib/python3.11/site-packages/pytensor/link/c/c_code -L/Users/ducquangnguyen/opt/anaconda3/envs/pymc5_16_env/lib -fvisibility=hidden -o /Users/ducquangnguyen/.pytensor/compiledir_macOS-14.6.1-arm64-arm-64bit-arm-3.11.8-64/tmp2bepk3qv/m847b81085fa430c45d8c59a1335e47f53e5f3c808c86bb220d8048e364d86d74.so /Users/ducquangnguyen/.pytensor/compiledir_macOS-14.6.1-arm64-arm-64bit-arm-3.11.8-64/tmp2bepk3qv/mod.cpp
/Users/ducquangnguyen/.pytensor/compiledir_macOS-14.6.1-arm64-arm-64bit-arm-3.11.8-64/tmp2bepk3qv/mod.cpp:26069:32: fatal error: bracket nesting level exceeded maximum of 256
 26069 |         if (!PyErr_Occurred()) {
       |                                ^
/Users/ducquangnguyen/.pytensor/compiledir_macOS-14.6.1-arm64-arm-64bit-arm-3.11.8-64/tmp2bepk3qv/mod.cpp:26069:32: note: use -fbracket-depth=N to increase maximum nesting level
1 error generated.


    └─ ···
 ← add [id BJ] <float64> 'o34'
    └─ ···
 ← add [id BK] <float64> 'o35'
    └─ ···
 ← add [id BL] <float64> 'o36'
    └─ ···
 ← add [id BM] <float64> 'o37'
    └─ ···

HINT: Use a linker other than the C linker to print the inputs' shapes and strides.
HINT: Re-running with most PyTensor optimizations disabled could provide a back-trace showing when this node was created. This can be done by setting the PyTensor flag 'optimizer=fast_compile'. If that does not work, PyTensor optimizations can be disabled with 'optimizer=None'.
HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.

I suppose this might be something related to an upgrade that was done for PyMC - which somehow broke my codes.
I have tried to run ‘exception_verbosity=high’ but to be honest, a lot of those outputs do not make sense to me…

I have never experienced this issue with 5.6.1 so I am not too sure where to start from.

If you could have any ideas, please let me know

2

can you share your model? do you have perhaps a for loop in it?

3

Thanks.
I do not have a ‘for’ loop in my model. But I have time-varying parameters that I am using PyMC to fit (it is quite important for my application).

The ode_model is as followed:

def ode_model(state, t, params): 
    # Define the compartments
    S = state[0]
    I_W = state[1]
    I_DR = state[2]
    D_W = state[3]
    D_DR = state[4]
    T_W1 = state[5]
    T_W2 = state[6]
    T_DR1 = state[7]
    T_DR2 = state[8]
    F_W1 = state[9]
    F_W2 = state[10]
    F_TDR1 = state[11]
    F_ADR1 = state[12]
    F_DR2 = state[13]
    incidence = state[14]
    deaths = state[15]

    # Define parameters
    beta_psi = params[0]  # Infection rate
    delta_I = params[1]    # Infected mortality rate
    delta_T = params[2]    # Treated mortality rate
    delta_F = params[3]    # Failure mortality rate
    rho = params[4]        # Population growth rate
    k = params[5]          # Treatment rate
    h1 = params[6]         # Resistance impact factor 1
    h2 = params[7]         # Resistance impact factor 2

    # Total population
    N = (S + I_W + I_DR + D_W + D_DR + T_W1 + T_W2 + T_DR1 + T_DR2 + 
         F_W1 + F_W2 + F_TDR1 + F_ADR1 + F_DR2)

    # Time-varying forces of infection
    I_Wa = beta_psi * (I_W + D_W + T_W1 + T_W2) / N * S
    I_Ra = beta_psi * (I_DR + D_DR + T_DR1 + T_DR2) / N * S

    # Simplified differential equations with time-varying parameters
    S_prime = rho * S - (I_Wa + I_Ra)
    I_W_prime = I_Wa - b(t, b_asterisk) * I_W - delta_I * I_W
    I_DR_prime = I_Ra - b(t, b_asterisk) * I_DR - delta_I * I_DR
    D_W_prime = b(t, b_k) * I_W - c1(t, c_b) * D_W - delta_I * D_W
    D_DR_prime = b(t, b_k) * I_DR - c2(t, c_k) * D_DR - delta_I * D_DR
    T_W1_prime = c_1(t, c_b) * D_W - f1(t, f1_rate) * T_W1 - delta_T * T_W1
    T_W2_prime = c_1prime(t, c_b) * T_W1 - f2(t, f2_rate) * T_W2 - delta_T * T_W2
    T_DR1_prime = c_2(t, c_k) * D_DR - f3(t, f3_rate) * T_DR1 - delta_T * T_DR1
    T_DR2_prime = c_2prime(t, c_k) * T_DR1 - f4(t, f4_rate) * T_DR2 - delta_T * T_DR2
    F_W1_prime = f1(t, f1_rate) * T_W1 - delta_F * F_W1 - h1 * F_W1
    F_W2_prime = f2(t, f2_rate) * T_W2 - delta_F * F_W2 - h2 * F_W2
    F_TDR1_prime = f3(t, f3_rate) * T_DR1 - delta_F * F_TDR1
    F_ADR1_prime = h1 * F_W1 + h2 * F_W2 - delta_F * F_ADR1
    F_DR2_prime = f4(t, f4_rate) * T_DR2 - delta_F * F_DR2
    incidence_prime = I_Wa + I_Ra
    deaths_prime = (delta_I * (I_W + I_DR + D_W + D_DR) + 
                    delta_T * (T_W1 + T_DR1 + T_W2 + T_DR2) + 
                    delta_F * (F_W1 + F_TDR1 + F_ADR1 + F_W2 + F_DR2))

    return [S_prime, I_W_prime, I_DR_prime, D_W_prime, D_DR_prime, T_W1_prime, T_W2_prime, T_DR1_prime, T_DR2_prime, F_W1_prime, F_W2_prime, F_TDR1_prime, F_ADR1_prime, F_DR2_prime, incidence_prime, deaths_prime]

Some of the parameters, as you can see above are time-varying such as b(t,b_k). Not too sure whether that is what caused the issue.

After declaring my model, my model setup codes are fairly straightforward, except for the part where I need to specify weights on some of my data for fitting.

import pymc as pm
import pytensor.tensor as pt

with pm.Model() as model:
    # Observed data placeholder
    y = pm.MutableData("y", observed_data)
   
    # Simplified parameter initialization
    beta_psi = pm.Gamma('beta_psi', mu=0.1, sigma=0.1)
    delta_I = pm.Gamma('delta_I', mu=0.1, sigma=0.1)
    delta_T = pm.Gamma('delta_T', mu=0.1, sigma=0.1)
    delta_F = pm.Gamma('delta_F', mu=0.1, sigma=0.1)
    k = pm.Uniform('k', lower=0, upper=1)

    # Other key parameters
    phi_inv = pm.Exponential("phi_inv", lam=10)
    phi = pm.Deterministic('phi', 1. / phi_inv)

    # Simplified ODE model
    mcmc_curves = mcmc_ode(
        y0=initial_conditions, 
        theta=[beta_psi, delta_I, delta_T, delta_F, k]
    )

    # Simplified tensor extraction
    tensor_a = mcmc_curves[:, 14]  # Assuming index 14 corresponds to the incidence compartment
    tensor_b = mcmc_curves[:, 2]  # Assuming index 2 corresponds to I_DR
    tensor_c = mcmc_curves[:, 5]  # Assuming index 5 corresponds to T_W1

    # Flatten and concatenate tensors
    tensor_a_flattened = pt.flatten(tensor_a)
    tensor_b_flattened = pt.flatten(tensor_b)
    tensor_c_flattened = pt.flatten(tensor_c)
    concatenated_tensor = pt.concatenate([tensor_a_flattened, tensor_b_flattened, tensor_c_flattened])

    # Likelihood definition
    rv = pm.NegativeBinomial.dist(mu=concatenated_tensor, alpha=phi)
    logp = pm.logp(rv, y)
    potential = pm.Potential("logp_weighted", logp.sum())

with model:   
    # Sample from the posterior using NUTS
    trace = pm.sample(draws=250, tune=250, step=pm.NUTS(), chains=4, cores=4)

Let me know what you think that might cause the error messages above

4

Can you provide a running example I can copy and paste and execute locally? Also if you can pair it down to the absolute minimum to trigger the issue, that would be great

1 Like
5

I narrowed down to the cause of this - which are pm.Potential and complex modelling.

The following codes can trigger the bugs on my laptop:

### Totally redacted model for reviewing
import numpy as np
import pymc as pm
from pymc.ode import DifferentialEquation
from functools import partial
from scipy.integrate import odeint
import random
import pytensor
import pytensor.tensor as pt
pytensor.config.exception_verbosity = 'high'


np.random.seed(334)
random.seed(123)




def expit(x,k,a):
    return a + (k-a)/(1 + np.exp(-10*x)) 

b = partial(expit)
c = partial(expit)
f = partial(expit)


I_W = 1000
S = 3_000_000 - I_W
I_DR = D_W = D_DR = T_W1 = T_DR1 = T_W2 = T_DR2 = F_W1 = F_TDR1 = F_ADR1 = F_W2 = F_DR2 = incidence = deaths = dummy_a = dummy_b = dummy_c = dummy_d = dummy_e = dummy_f = dummy_g = dummy_h = dummy_i = 0
total_PLHIV = I_W

initial_conditions = [S, I_W, I_DR, D_W, D_DR, T_W1, T_W2, T_DR1, T_DR2, F_W1, F_W2, F_TDR1, F_ADR1, F_DR2, incidence, deaths, total_PLHIV, dummy_a, dummy_b, dummy_c, dummy_d, dummy_e, dummy_f, dummy_g, dummy_h, dummy_i]
initial_conditions = [x/sum(initial_conditions) for x in initial_conditions]

def ode_model(state,t,params): 
    S = state[0]
    I_W= state[1]
    I_DR= state[2]
    D_W= state[3]
    D_DR= state[4]
    T_W1= state[5]
    T_W2= state[6]
    T_DR1= state[7]
    T_DR2= state[8]
    F_W1= state[9]
    F_W2= state[10]
    F_TDR1= state[11]
    F_ADR1= state[12]
    F_DR2= state[13]
    incidence= state[14]
    deaths= state[15]
    total_PLHIV = state[16]
    dummy_a = state[17]
    dummy_b = state[18]
    dummy_c = state[19]
    dummy_d = state[20]
    dummy_e = state[21]
    dummy_f = state[22]
    dummy_g = state[23]
    dummy_h = state[24]
    dummy_i = state[25]




    theta = params[0]
    beta_psi = params[1]
    delta_I = params[2]  
    b_k = params[3]
    rho_asterisk = params[4]


    N = S + I_W + I_DR + D_W + D_DR + T_W1 + T_W2 + T_DR1 + T_DR2 + F_W1 + F_W2 + F_TDR1 +F_ADR1 + F_DR2
    I_Wa =   beta_psi*(  I_W+ D_W + T_W1+T_W2 +   F_W1+F_W2)/N*S
    I_Ra =   theta*  beta_psi*(  I_DR+D_DR + T_DR1+T_DR2+   F_TDR1+F_ADR1+F_DR2)/N*S
    
    S_prime = rho_asterisk*S - (I_Wa+I_Ra) 
    I_W_prime = I_Wa -   b(t,k=b_k,a=b_k)*I_W -   delta_I*I_W
    I_DR_prime= I_Ra -   b(t,k=b_k,a=b_k)*I_DR -   delta_I*I_DR
    
    D_W_prime =   b(t,k=b_k,a=b_k)*I_W  -   c(t,k=b_k,a=b_k)*D_W -   delta_I*D_W
    D_DR_prime =   b(t,k=b_k,a=b_k)*I_DR  -   c(t,k=b_k,a=b_k)*D_DR -   delta_I*D_DR
    


    T_W1_prime =   c(t,k=b_k,a=b_k)*D_W  - (  f(t,k=b_k,a=b_k)+  delta_I)*T_W1 
    T_W2_prime =   c(t,k=b_k,a=b_k)*D_W  - (  f(t,k=b_k,a=b_k)+  delta_I)*T_W2 
    T_DR1_prime =   c(t,k=b_k,a=b_k)*D_DR - (  f(t,k=b_k,a=b_k)+  delta_I)*T_DR1
    T_DR2_prime =   c(t,k=b_k,a=b_k)*D_DR  - (  f(t,k=b_k,a=b_k)+  delta_I)*T_DR2

    F_W1_prime =   f(t,k=b_k,a=b_k)*T_W1  - delta_I*F_W1
    F_W2_prime =   f(t,k=b_k,a=b_k)*T_W2  - delta_I*F_W2
    F_TDR1_prime =   f(t,k=b_k,a=b_k)*T_DR1  -   delta_I*F_TDR1
    F_ADR1_prime =     -   delta_I*F_ADR1
    F_DR2_prime =   f(t,k=b_k,a=b_k)*T_DR2  -   delta_I*F_DR2




    incidence =  I_Wa + I_Ra
    deaths =  delta_I*(I_W+I_DR+D_W+D_DR+T_W1+T_DR1+T_W2+T_DR2+F_W1+F_TDR1+F_ADR1+F_W2+F_DR2)
    total_PLHIV_prime = I_Wa + I_Ra - deaths -  c(t,k=b_k,a=b_k)*D_W - c(t,k=b_k,a=b_k)*D_DR + c(t,k=b_k,a=b_k)*D_W +c(t,k=b_k,a=b_k)*D_W +c(t,k=b_k,a=b_k)*D_DR +c(t,k=b_k,a=b_k)*D_DR

    dummy_a = I_Ra
    dummy_b = 0
    dummy_c = 0
    dummy_d = 0
    dummy_e = 0
    dummy_f = 0
    dummy_g = 0
    dummy_h = 0
    dummy_i = 0


    return [S_prime, I_W_prime, I_DR_prime, D_W_prime, D_DR_prime, T_W1_prime, T_W2_prime, T_DR1_prime, T_DR2_prime, F_W1_prime, F_W2_prime, F_TDR1_prime, F_ADR1_prime, F_DR2_prime, incidence, deaths, total_PLHIV_prime, dummy_a, dummy_b, dummy_c, dummy_d, dummy_e, dummy_f, dummy_g, dummy_h, dummy_i]


###### Initiate variables for Bayesian MCMC analysis
t0=0
##### Year 2025
years = 31
tspan = np.arange(0, years, 1)

mcmc_ode = DifferentialEquation(
    func=ode_model,
    times=tspan,
    n_states=26,
    n_theta=5,
    t0=t0,
)

### synthetic, random number generator
observed_data = [3500, 4800, 6200, 7600, 8900, 10500, 12300, 14200, 15700, 17200, 18500, 20300, 21800, 23600, 25500, 27400, 29000, 31200, 33000, 35400, 37600, 40100, 42500, 45300, 48700, 51000, 54900, 59000, 63500, 68400]
observed_data = [i/3000000.0 for i in observed_data]

## new weights for July 13th 2023

weights = [1 if i < len(observed_data) - 6 else  2**(i - len(observed_data) + 7) for i in range(len(observed_data))] 
weights[8] = 5000
weights[-1] = 12500


with pm.Model() as model:
    y = pm.Data("y", observed_data)
   
    theta =  pm.Lognormal('theta', mu=-0.1, sigma=0.25)
    beta_psi = pm.Gamma('beta_psi', alpha=3, beta=23)
    delta_I =  pm.Gamma('delta_I', alpha=5, beta=100)
    b_k =  pm.Gamma('b_k', alpha=4.0, beta=20)
    rho_asterisk = pm.Uniform('rho_asterisk', lower=0.03, upper=0.04)



    phi_inv = pm.Exponential("phi_inv", lam=10)
    phi = pm.Deterministic('phi', 1. / phi_inv)

    mcmc_curves = mcmc_ode(y0=initial_conditions, theta=[theta,beta_psi,delta_I,b_k,rho_asterisk])

    tensor_a_flattened = pt.flatten( (mcmc_curves[:,16])[0:30])


    rv = pm.NegativeBinomial.dist(mu=tensor_a_flattened, alpha=phi)
    logp = pm.logp(rv, y)
    weighted_logp = logp * weights
    potential = pm.Potential("weighted_logp", weighted_logp.sum())


with model:   
    #### multicore processing  
    trace = pm.sample(draws=250, tune=250, step=pm.NUTS(), chains=4, cores=4)

The error from this simplified code is different from my original one which is:


  File ~/opt/anaconda3/envs/pymc5_16_env/lib/python3.11/multiprocessing/connection.py:430 in _recv_bytes
    buf = self._recv(4)

  File ~/opt/anaconda3/envs/pymc5_16_env/lib/python3.11/multiprocessing/connection.py:399 in _recv
    raise EOFError

EOFError

There are 3 ways that this error will disappear, based on this code.

  1. Change cores=1 to turn off the parallel processing
  2. Comment out the pm.Potential potential = pm.Potential("weighted_logp", weighted_logp.sum()). I believe the bug is mainly cause by pm.Potential. In my original code, I also did not have the original error, if I comment out this line.
  3. remove the dummy_a, dummy_b, dummy_c, dummy_d, dummy_e,.... from my ODE. If i remove these, the multiprocessing also worked.

I did not have this error in version 5.6.1

6

If you want me to reproduce the original error, I believe I will need to make this simplified code more complex.