Zero-Inflated Poisson Model: compute_log_likelihood, r2_score

[craig0a]

Hello,

Attempting to compute the model evaluation metrics compute_log_likelihood and r2_score for the fish example zero-inflated poisson resulted in the following errors:

TypeError                                 Traceback (most recent call last)
Cell In[12], line 1
----> 1 zip_model.compute_log_likelihood(idata=zip_idata)

File ~/.pyenv/versions/3.10.14/envs/baseline_3.10.14/lib/python3.10/site-packages/bambi/models.py:935, in Model.compute_log_likelihood(self, idata, data, inplace)
    930 idata = self._compute_likelihood_params(
    931     idata, data, include_group_specific, sample_new_groups
    932 )
    934 required_kwargs = {"model": self, "posterior": idata.posterior, "data": data}
--> 935 log_likelihood = self.family.log_likelihood(**required_kwargs)
    936 log_likelihood = log_likelihood.to_dataset(name=response_aliased_name)
    938 if "log_likelihood" in idata:

File ~/.pyenv/versions/3.10.14/envs/baseline_3.10.14/lib/python3.10/site-packages/bambi/families/family.py:217, in Family.log_likelihood(self, model, posterior, data, **kwargs)
    213     output_array = pm.logp(
    214         pm.Truncated.dist(response_dist.dist(**kwargs), lower=lower, upper=upper), y_values
    215     ).eval()
    216 else:
--> 217     output_array = pm.logp(response_dist.dist(**kwargs), y_values).eval()
    219 return xr.DataArray(output_array, coords=coords)

File ~/.pyenv/versions/3.10.14/envs/baseline_3.10.14/lib/python3.10/site-packages/pymc/logprob/basic.py:207, in logp(rv, value, warn_rvs, **kwargs)
    205 value = pt.as_tensor_variable(value, dtype=rv.dtype)
    206 try:
--> 207     return _logprob_helper(rv, value, **kwargs)
    208 except NotImplementedError:
    209     fgraph = construct_ir_fgraph({rv: value})

File ~/.pyenv/versions/3.10.14/envs/baseline_3.10.14/lib/python3.10/site-packages/pymc/logprob/abstract.py:81, in _logprob_helper(rv, *values, **kwargs)
     79 def _logprob_helper(rv, *values, **kwargs):
     80     """Help call `_logprob` dispatcher."""
---> 81     logprob = _logprob(rv.owner.op, values, *rv.owner.inputs, **kwargs)
     83     name = rv.name
     84     if (not name) and (len(values) == 1):

File ~/.pyenv/versions/3.10.14/lib/python3.10/functools.py:889, in singledispatch.<locals>.wrapper(*args, **kw)
    885 if not args:
    886     raise TypeError(f'{funcname} requires at least '
    887                     '1 positional argument')
--> 889 return dispatch(args[0].__class__)(*args, **kw)

File ~/.pyenv/versions/3.10.14/envs/baseline_3.10.14/lib/python3.10/site-packages/pymc/distributions/mixture.py:349, in marginal_mixture_logprob(op, values, rng, weights, *components, **kwargs)
    347     components_logp = logp(components[0], pt.expand_dims(value, mix_axis))
    348 else:
--> 349     components_logp = pt.stack(
    350         [logp(component, value) for component in components],
    351         axis=-1,
    352     )
    354 mix_logp = pt.logsumexp(pt.log(weights) + components_logp, axis=-1)
    356 mix_logp = check_parameters(
    357     mix_logp,
    358     0 <= weights,
   (...)
    361     msg="0 <= weights <= 1, sum(weights) == 1",
    362 )

File ~/.pyenv/versions/3.10.14/envs/baseline_3.10.14/lib/python3.10/site-packages/pytensor/tensor/basic.py:2955, in stack(tensors, axis)
   2953     dtype = ps.upcast(*[i.dtype for i in tensors])
   2954     return MakeVector(dtype)(*tensors)
-> 2955 return join(axis, *[shape_padaxis(t, axis) for t in tensors])

File ~/.pyenv/versions/3.10.14/envs/baseline_3.10.14/lib/python3.10/site-packages/pytensor/tensor/basic.py:2795, in join(axis, *tensors_list)
   2793     return tensors_list[0]
   2794 else:
-> 2795     return _join(axis, *tensors_list)

File ~/.pyenv/versions/3.10.14/envs/baseline_3.10.14/lib/python3.10/site-packages/pytensor/graph/op.py:293, in Op.__call__(self, name, return_list, *inputs, **kwargs)
    249 def __call__(
    250     self, *inputs: Any, name=None, return_list=False, **kwargs
    251 ) -> Variable | list[Variable]:
    252     r"""Construct an `Apply` node using :meth:`Op.make_node` and return its outputs.
    253 
    254     This method is just a wrapper around :meth:`Op.make_node`.
   (...)
    291 
    292     """
--> 293     node = self.make_node(*inputs, **kwargs)
    294     if name is not None:
    295         if len(node.outputs) == 1:

File ~/.pyenv/versions/3.10.14/envs/baseline_3.10.14/lib/python3.10/site-packages/pytensor/tensor/basic.py:2488, in Join.make_node(self, axis, *tensors)
   2485 ndim = tensors[0].type.ndim
   2487 if not builtins.all(x.ndim == ndim for x in tensors):
-> 2488     raise TypeError(
   2489         "Only tensors with the same number of dimensions can be joined. "
   2490         f"Input ndims were: {[x.ndim for x in tensors]}"
   2491     )
   2493 try:
   2494     static_axis = int(get_scalar_constant_value(axis))

TypeError: Only tensors with the same number of dimensions can be joined. Input ndims were: [2, 4]

AttributeError                            Traceback (most recent call last)
Cell In[26], line 1
----> 1 zip_model.r2_score(idata=zip_idata)

AttributeError: 'Model' object has no attribute 'r2_score'

Would it be possible to have these functionalities implemented?

Thank you!

Best regards,
Æ

[GStechschulte]

Hey @craig0a thanks for opening the issue.

I get the same error as you when calling the compute_log_likelihood method. An alternative could be to pass the following idata_kwargs to populate the InferenceData object with a log-likelihood group where the associated values are in the count field.

zip_idata = zip_model.fit(
        draws=1000,
        target_accept=0.95,
        random_seed=1234,
        chains=4,
        idata_kwargs={"log_likelihood": True}
    )
print(zip_idata["log_likelihood"])

Dimensions:  (chain: 4, draw: 1000, __obs__: 248)
Coordinates:
  * chain    (chain) int64 32B 0 1 2 3
  * draw     (draw) int64 8kB 0 1 2 3 4 5 6 7 ... 993 994 995 996 997 998 999
  * __obs__  (__obs__) int64 2kB 0 1 2 3 4 5 6 7 ... 241 242 243 244 245 246 247
Data variables:
    count    (chain, draw, __obs__) float64 8MB -0.2121 -0.5855 ... -0.5203
Attributes:
    created_at:                  2025-06-24T18:45:29.181597+00:00
    arviz_version:               0.21.0
    inference_library:           pymc
    inference_library_version:   5.23.0
    modeling_interface:          bambi
    modeling_interface_version:  0.15.1.dev14+gcdea205d.d20250618

Calling r2_score works for me.

zip_idata = zip_model.fit(
      draws=1000,
      target_accept=0.95,
      random_seed=1234,
      chains=4
  )

print(zip_model.r2_score(idata=zip_idata, summary=True))

r2        0.274613
r2_std    0.053568
dtype: float64

What version of Bambi are you using?

[craig0a]

Hello @GStechschulte,

Thank you so much for your help!

• I am using bambi 0.15.0 with Python 3.10.14
• The log-likelihood solution works for the example test case:

zip_model_Camper = bmb.Model(
    "count ~ livebait + camper + persons + child", 
    fish_data, 
    family='zero_inflated_poisson'
)

zip_idata_Camper = zip_model_Camper.fit(
    draws=1000, 
    target_accept=0.95, 
    random_seed=1234, 
    chains=4,
    idata_kwargs={"log_likelihood": True}
)


zip_model_noCamper = bmb.Model(
    "count ~ livebait + persons + child", 
    fish_data, 
    family='zero_inflated_poisson'
)

zip_idata_noCamper = zip_model_noCamper.fit(
    draws=1000, 
    target_accept=0.95, 
    random_seed=1234, 
    chains=4,
    idata_kwargs={"log_likelihood": True}
)



compare_dict = {
    'CamperModel':zip_idata_Camper,
    'NoCamperModel':zip_idata_noCamper
}

az.compare(compare_dict)

/Users/acraig/.pyenv/versions/3.10.14/envs/baseline_3.10.14/lib/python3.10/site-packages/arviz/stats/stats.py:797: UserWarning: Estimated shape parameter of Pareto distribution is greater than 0.70 for one or more samples. You should consider using a more robust model, this is because importance sampling is less likely to work well if the marginal posterior and LOO posterior are very different. This is more likely to happen with a non-robust model and highly influential observations.
  warnings.warn(
/Users/acraig/.pyenv/versions/3.10.14/envs/baseline_3.10.14/lib/python3.10/site-packages/arviz/stats/stats.py:797: UserWarning: Estimated shape parameter of Pareto distribution is greater than 0.70 for one or more samples. You should consider using a more robust model, this is because importance sampling is less likely to work well if the marginal posterior and LOO posterior are very different. This is more likely to happen with a non-robust model and highly influential observations.
  warnings.warn(

rank	elpd_loo	p_loo	elpd_diff	weight	se	dse	warning	scale
CamperModel	0	-569.058570	36.640069	0.000000	0.669443	74.938651	0.000000	True	log
NoCamperModel	1	-569.539038	26.941039	0.480468	0.330557	68.653508	13.906014	True	log

Thank you again!

Best,
Æ

[GStechschulte]

@tomicapretto I noticed the docstring of compute_log_likelihood says

NOTE: This is a new feature and it may not work in all cases.

Do we want to keep this issue open as an example of the method not working?

[GStechschulte]

@craig0a regarding the r2_score error. This is still weird. Did you also manage to get this working?

[tomicapretto]

@tomicapretto I noticed the docstring of compute_log_likelihood says

NOTE: This is a new feature and it may not work in all cases.

Do we want to keep this issue open as an example of the method not working?

I think it's good to keep it until we know it works with all families. It's been some time now since I worked on that, and I don't remember all the details, but I do remember thinking it would not work in all cases, thus the message.