Error when trying to implement cox proportional hazard model

Hi there,

I am new to NumPyro and I am trying to build a cox proportional hazard model for survival analysis. I followed this example and tried to replicate the model in NumPyro:

def cox_model(df, duration_col='time', event_col='event'):

        lambda0 = numpyro.sample('lambda0', dist.Gamma(0.01, 0.01), sample_shape=(n_intervals,))

        beta = numpyro.sample('beta', dist.Normal(0, 1000))

        lambda_ = numpyro.deterministic('lambda_', jnp.outer(jnp.exp(beta * df.metastasized.values), lambda0))

        mu = numpyro.deterministic('mu', exposure * lambda_)
        
        
        y = numpyro.sample('obs', dist.Poisson(mu), obs=death)
        

        
rng_key = random.PRNGKey(0)
rng_key, alpha_key = random.split(rng_key)
mcmc = MCMC(NUTS(cox_model, target_accept_prob=0.95), num_samples=1000, num_warmup=1000, num_chains=2)
mcmc.run(alpha_key, df)
mcmc_samples = mcmc.get_samples()

but I get this error

ValueError: Poisson distribution got invalid rate parameter

How do I modify my model to prevent this error message. Thanks for your help.

You might want to add some print statements to your model to see why mu is invalid for Poisson likelihood. It seems to get inf or nan. I guess jnp.exp(beta * df.metastasized.values) somehow blow up? You might want to use bounded_exp here.

Thanks for your reply. You are absolutely spot on. The reason mu is invalid is there were some zeros in mu. I tried using the bounded_exp to replace jnp.exp(beta * df.metastasized.values) but I still end up with zeros in the parameter mu. Is there a way to handle these zeros in the mu parameter?

How about adding a very small positive value to it, like mu + jnp.finfo(mu.dtype).tiny? Or clip the value…

Adding a very small positive value works fine. Thank you for your help @fehiepsi