Loop fusion failed

General
8

The way you want to handle this is to first reshape your data into “long form”. Here’s a dummy dataset with 5 runners and 3 times:

                       times
runner  year                
A      2018-01-01   1.492973
       2019-01-01   1.259906
       2020-01-01   5.536423
B      2019-01-01   7.734591
       2020-01-01  11.157608
C      2018-01-01   2.014740
       2020-01-01   7.937118
D      2018-01-01   2.950629
E      2019-01-01   3.261694
       2020-01-01   6.551939

So you can see that it’s off-balance, we only have 1 year for runner D, but all 3 for runner A.

Next, we need to get indexes for every row, telling us which year and which runner it is. We will use these to slice into random variables. To get these, use pd.factorize:

runner_idx, runners = pd.factorize(index.get_level_values(0))
year_idx, years = pd.factorize(index.get_level_values(1))

These are just arrays of integers:

runner_idx
>>> Out: array([0, 0, 1, 1, 1, 2, 3, 3, 4, 4], dtype=int64)

time_idx
>>> Out: array([0, 1, 2, 0, 1, 1, 2, 1, 2, 1], dtype=int64)

So you can see that the first two rows are runner 0, on years 0, 1 then runner 1 on years 2, 0, 1, and so on. You also get back the list of labels to use in your PyMC coords. Here’s the relevant portion of a PyMC model:

coords = {'runner':runners, 'year':[x.year for x in years]}
with pm.Model(coords=coords) as mod:
    y = pm.MutableData('y', df.values)
    runner_mean = pm.Normal('runner_mean', dims=['runner'])
    year_mean = pm.Normal('year_mean', dims=['year'])
    mu = runner_mean[runner_idx] + year_mean[year_idx]

So you make two effects – one for runners, one for years. Then you can use the indexes to add them together. The indices will ensure that the right effects end up with the right rows. This is equivalent to broadcasting them together, like runner_mean[:, None] + year_mean[None], which would give you the full (n_runners, n_years) matrix, except that the indexes mean only the data you have is computed.

1 Like
Hierarchical model varying indexing of distributions