NaN is included in the result of executing the scan function in the pymc model

[Yamaguchi]

Hi all,

I got the output described in the title when I executed the following code.
I do not know why the ode function outputs NaN this time.
Also, if you know of a solution to prevent this, could you please let me know?

I would be thankful if you could help me with this question.

import os
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
import arviz as az

import pymc as pm
import pytensor
import pytensor.tensor as pt

# Initialize random number generator
RANDOM_SEED = 8927
rng = np.random.default_rng(RANDOM_SEED)
az.style.use("arviz-darkgrid")

time = 600
steps_per_second = 1000
n_steps = time * steps_per_second
dt =1 / steps_per_second

def ode_update_func(A, B, C, k1, k2, am, bm):
    A_new = A + (-k1*pt.power(A, am))*dt
    B_new = B + (k1*pt.power(A, am) - k2*pt.power(B, bm))*dt
    C_new = C + (k2*pt.power(B, bm))*dt
    
    return A_new, B_new, C_new

k1_true = 0.1
k2_true = 0.01
am_true = 1.0
bm_true = 1.0

A_init = np.array(1.0, dtype=np.float64)
B_init = np.array(0.0, dtype=np.float64)
C_init = np.array(0.0, dtype=np.float64)

results, updates = pytensor.scan(
    fn=ode_update_func,
    outputs_info=[A_init, B_init, C_init],
    non_sequences=[k1_true, k2_true, am_true, bm_true],
    n_steps=n_steps,
    strict=True
    )

results_true = np.concatenate((np.array([[A_init, B_init, C_init]]), np.stack([result.eval() for result in results]).T))

plt.plot(np.arange(n_steps+1)/steps_per_second, results_true[:, 0], label="A_true")
plt.plot(np.arange(n_steps+1)/steps_per_second, results_true[:, 1], label="B_true")
plt.plot(np.arange(n_steps+1)/steps_per_second, results_true[:, 2], label="C_true")
plt.legend()
plt.show()

with pm.Model() as model:
    k1 = pm.TruncatedNormal("k1", mu=0.3, sigma=0.2, lower=0)
    k2 = pm.TruncatedNormal("k2", mu=0.03, sigma=0.02, lower=0)
    
    am = pm.TruncatedNormal("am", mu=1.0, sigma=0.2, lower=0)
    bm = pm.TruncatedNormal("bm", mu=1.0, sigma=0.2, lower=0)
    
    sigma = pm.InverseGamma("sigma", alpha=5, beta=0.2)
    
    A_init = np.array(1.0, dtype=np.float64)
    B_init = np.array(0.0, dtype=np.float64)
    C_init = np.array(0.0, dtype=np.float64)
    
    results, updates = pytensor.scan(
        fn=ode_update_func,
        outputs_info=[A_init, B_init, C_init],
        non_sequences=[k1, k2, am, bm],
        n_steps=n_steps,
        strict=True
        )
    
    pm.Deterministic("result", pm.math.stack([results[0], results[1], results[2]]))

with model:
    idata_prior = pm.sample_prior_predictive(5, random_seed=rng)

print(idata_prior.prior.result)

for i in range(5):
    plt.plot(np.arange(1, n_steps+1)/steps_per_second, idata_prior.prior.result[0, i, 0, :], label="A_prior")
    plt.plot(np.arange(1, n_steps+1)/steps_per_second, idata_prior.prior.result[0, i, 1, :], label="B_prior")
    plt.plot(np.arange(1, n_steps+1)/steps_per_second, idata_prior.prior.result[0, i, 2, :], label="C_prior")

[Yamaguchi]

I found that during the loop when A or B takes a very small value (-1.5e-33) it is NaN.
Therefore, I modified the code as follows.

def ode_update_func(A, B, C, k1, k2, am, bm):
    A_new = A + (-k1*pt.power(A, am))*dt
    B_new = B + (k1*pt.power(A, am) - k2*pt.power(B, bm))*dt
    C_new = C + (k2*pt.power(B, bm))*dt
    
    A_new = pt.where(pt.le(A_new, 1e-10), np.array(0, dtype=np.float64), A_new)
    B_new = pt.where(pt.le(B_new, 1e-10), np.array(0, dtype=np.float64), B_new)

    return A_new, B_new, C_new

But I am not sure if this is the best way to fix it. If you know any other way, could you please let me know?

[jessegrabowski]

Small values shouldn’t underflow to NaN unless you do an invalid operation on zero, like 1/x. For example, pt.as_tensor(np.array([1e-500])).eval() just gives back zero. I’d make sure am or bm are always strictly positive. One thing I use for debugging scan loops is the pytensor.printing.Print op, which will print the values of a variable to the terminal every time it changes. Docs here.

Keep in mind that just because a prior rules a value out (i.e. it has zero probability) a sampler could still propose such a value, so you could end up with NaNs in your computation. Clipping might be necessary if this is happening.

[Yamaguchi]

Thank you for your response.
After detailed debugging using Print, I found that the problem was caused by the fact that the power output is imaginary in the case where A has a negative value.
Therefore, I think I solved the bug by clipping the value of A to 0 as follows.

A_new = pt.clip(A_new, 0, np.inf)