bambinos / bambi

BAyesian Model-Building Interface (Bambi) in Python.
https://bambinos.github.io/bambi/
MIT License
1.08k stars 124 forks source link

Support PyMC 5.13 and fix bayeux related issues #803

Closed tomicapretto closed 7 months ago

tomicapretto commented 7 months ago

This PR closes #800


Why I modified how we rename dims with bayeux. See the following example:

from io import StringIO

import requests

import arviz as az
import bambi as bmb
import pandas as pd

url = "https://raw.githubusercontent.com/crnolan/pyrba/main/data.txt"  # replace with your url
response = requests.get(url)
df = pd.read_table(StringIO(response.text), delimiter=r"\s+")

results = model.fit(
    num_warmup=250, num_samples=200, num_chains=2, inference_method="numpyro_nuts"
)

With the old implementation, it would have failed. Look at

https://github.com/bambinos/bambi/blob/5c00663ef290744f20f7f9de700d9051a2ea55b7/bambi/backend/pymc.py#L275-L281

what happened in that case was that pymc_model_dims and bayeux_dims where of different lengths and thus the dictionary created was mixing things.

This new approach

        dims_original = list(self.model.coords)

        # Identify bayeux idata and rename dims and coordinates to match PyMC model
        if idata_from == "bayeux":
            cleaned_dims = {
                f"{dim}_0": dim
                for dim in dims_original
                if not dim.endswith("_obs") and f"{dim}_0" in idata.posterior.dims
            }

exploits the fact that we know we can only rename dims passed to the model, they can't finish with _obs and they have to be in the posterior. And most importantly, the "bad" bayeux names are like the original names with an extra _0 (e.g. party_id_dim original and party_id_dim_0 as result in bayeux).