Drawing samples with null probability

Misc.
1

I am having a hard time understanding the behavior of the sample function. For example, with the following code:

@partial(jax.jit, static_argnums=(2,))
def get_states_without_transient(sample, y0, start_ind):
    states = odeint(jr_erp_forward_model, y0, t, sample, rtol=1e-6, atol=1e-5, mxstep=1000)[start_ind:, :]
    return states[:, 1] - states[:, 2]

def model_mcmc(priors_par, start_ind=None, y0=None, obs=None):

    # drawing a sample from the prior distribution
    sample = {key: numpyro.sample(key, dist.TruncatedNormal(loc=vals["mean"], scale=vals["sd"])) 
              for key, vals in priors_par.items()}
    jax.debug.print("sample = {sample}; {m}; {sd}", sample=sample["A"], 
                    m=priors_par["A"]["mean"], sd=priors_par["A"]["sd"])

    output = get_states_without_transient(sample, y0, start_ind)
    obs_sd = jnp.std(jnp.array(output))*0.1
    
    numpyro.sample('obs', dist.Normal(output, obs_sd), obs=obs)


nuts_kernel = NUTS(model_mcmc, dense_mass=True)
mcmc_kwparams = dict(num_chains = 1,
                     num_samples = 100,
                     num_warmup = 0)

mcmc = MCMC(nuts_kernel, **mcmc_kwparams)

start_ind = np.where(t>=0)[0][0]
mcmc.run(PRNGKey(1), priors_par=priors_par, start_ind=start_ind, y0=y0, obs=output[t>=0])

I get an output that looks like this:

sample = -0.7482905387878418; 3.25; 0.032499998807907104
sample = -0.7482905387878418; 3.25; 0.032499998807907104
sample: 100%|██████████| 100/100 [00:08<00:00, 12.32it/s, 1 steps of size 1.00e+00. acc. prob=0.00]sample = -31059187712.0; 3.25; 0.032499998807907104
sample = -31059187712.0; 3.25; 0.032499998807907104
sample = -31059187712.0; 3.25; 0.032499998807907104
sample = -31059187712.0; 3.25; 0.032499998807907104
sample = -31059187712.0; 3.25; 0.032499998807907104
sample = -31059187712.0; 3.25; 0.032499998807907104
sample = -31059187712.0; 3.25; 0.032499998807907104

I am not sure how a value drawn from a normal truncated at 0 would be negative. I understand that the sampling can be influenced by the MCMC context in which it is performed, but still, these are just aberrant values. Since these parameters need to have positive values, the ODE just returns a bunch of NaNs and the MCMC sampler sits there doing nothing but always returning the same completely off samples. I am pretty sure that I am missing something very basic, but somehow I could not find something enlightening on that in the documentation.

2

afaik TruncatedNormal defaults to being non-truncated unless you specify high or low:

TruncatedNormal(loc=0.0, scale=1.0, *, low=None, high=None, ...)

checkout this tutorial

also probably use 64 bit precision

also be aware that doing inference with ode’s in the loop can be very numerically unstable depending on the size of the latent space, the number of time steps, etc etc etc