Understanding coords, indexation, Data, ..., for multilevel models

[OriolAbril]

@covertg

Is there “broadcasting” of shape/dimension when performing operations on variables? If so, how does it work?

Yes, random variables are aesara tensors, so the broadcasting rules explained in https://aesara.readthedocs.io/en/latest/tutorial/broadcasting.html apply. TL; DR they are similar to numpy ones but tensors also have a true/false indicator to “turn off” broadcasting on some dimensions if desired.

@covertg

When performing indexation on vector or array variables (variables that have dimensionality more complex than just a scalar), most of the multilevel model examples that I’ve found end up doing operations in the dimension of the total number of observations. That is, they perform fancy indexation based on various other dimensions, but that indexation yields a vector that is as long as the number of observations in the data. Is this the recommended way to proceed?

Yes, I’d recommend factorizing into 2d vectors with length as the number of observations in a style similar to long format dataframes. There are two reasons for that. The first is that it works also with ragged data; keeping every hierarchy level as its own dimension only works if there are the same number of observations per group which is generally not the case. The second is I haven’t had time to do proper benchmarks but I am quite sure it is faster to implement the models this way unless no indexing were required if working with multidimensional arrays of parameters.

@covertg

My confusion here is that the important dimensionality is sometimes (often?) not as large as the number of observations, and so I wonder if there is a more efficient and understandable way to proceed.

Using your example, if you have M observations per combination of city and time, then you can strucutre you observations as a 3d array with dimensions time: 2, city: 3, observation: 5 → n_obs = 30. You can then use broadcasting only to combine the 1d (length 2) time parameters vector with the 1d (length 3) citiy parameters, no indexing would be needed here and I think it would be quite intuitive, but that is generally not the case and you have 5 observations for Berlin and 7 for Paris.

I think it is important to note that the number of parameters and the number of observations are independent and you always need to evaluate the model on all your observations. In a simple linear regression we have 3 parameters a, b, \sigma but the log likelihood is the sum of the linear model evaluated at every observations; PyMC does that automatically (and aesara can optimize the operations done sometimes) when given the observed kwarg with all the observations.

Second note, the pm.Deterministics and the city_temperature[cities_idx] do not add more parameters to the model, they only create transformations and arrays with the original parameters restructured to compute the right quantity when working with arrays.

@covertg

Is there a recommended way to represent multifactor/multilevel data for convenient use in pymc? I have been experimenting with various implementation details: dataframes that are “long”, versus dataframes that are “wide”, using pandas MultiIndex , using one pm.Data object to store multiple columns’ worth of data (columns are considered in the “long” format), using multiple pm.Data objects to store each column of data, and so on. Different combinations of these options can give rise to different ways to specify models in pymc, and it’s been difficult to know if I’m making my life easier or harder depending on how I represent my data.

In my opinion, going back to what I said above you need 1d vectors and all the factors as indexes (you can use pandas.factorize for this like you have done) of if same # of observations per category then multidimensional array (not compatible with pandas anymore unless you have only one level). I can’t really speak about having the levels as “raw” columns or as multiindex though.

---

---

@covertg

I’ve deviated from the original notebook by expressing the data in “long” format. It makes the most sense to me as it’s very easy to handle arbitrary combinations of factor levels, as well as possibly imbalanced observations, in this format. I hope this is a wise enough first step.

I think it is

@covertg

This samples just as well as the original example. However, it strikes me that the likelihood node only needs to have a size of 3 (instead, it has size of 174)

Here is where the parameters vs observation difference comes into play. The likelihood can’t be of size 3 because the likelihood combines both data and parameters. There are only 3 possible values for the mean yes, but we need to evaluate the normal (whatever one of the 3 means we need to use) at each of the observations. In the exponent of the pdf (x - \mu) you need to take the x into account too, not only the \mu.

Moreover, to do posterior predictive checks you need to preserve that structure.

However, this can be quite inconvenient if you want to analyze the posterior predictive samples/predictions at each city (i.e. to get the probability the temperature will be higher than t in each city).

To get around this I think you should consider two options:

• You can update the cities_idx_pm with pm.set_data and replace it with an arange of length the number of cities. This is easy and convenient if you only have one category, but it quickly becomes cumbersome. To add the time component you’d need to use 0, 0, 1, 1, 2, 2 as city_idx and 0, 1, 0, 1, 0, 1 as time_idx…
• Repeat the mu computation with xarray and use xarray-einstats to sample from the posterior predictive. Broadcasting and alignment will be automatic and you’ll get a multidimensional array with hierarchies as different dimensions.

@covertg

Again to question 2: what is the difference between these two models? In this case, it doesn’t make a difference to our results, since sampling still works fine and the stochastic and deterministic variables have the shapes that we want.

The likelihood with size 3 looks very wrong to me. I still don’t grok the new size argument, but you might be generating a 3, 174 array to evaluate the log likelihood on. What is the shape of the data in the log likelihood group of the result? Do you get the same results?

@covertg

city_time_temperature = city_temperature[cities_idx_pm] + time_offset[times_idx_pm] # is there any way to make this dims (cities, times) ?

Yes, the example does have the same number of observations per category, so you can store the observations in a time, city, date 3d array. Then this line can become city_time_temperature = city_temperature[None, :] + time_offset[:, None] which might work as input directly if the observed is the 3d size, or you might need an extra length 1 dim for broadcasting to work so use mu=city_time_temerature[:, :, None].

However, big however. That is the mean of the temperatures, it is not what you observed. I’d recommend working with the posterior predictive samples to get values you can interpret as what you’d expect to see. Already with the current model, you can get posterior predictive samples as:

from xarray_einstats.stats import XrContinuousRV
from scipy.stats import norm

post = idata.posterior
temperature = XrContinuousRV(norm, post["city_temperature"] + post["time_offset"], 0.5).rvs()

I recently wrote a Discourse topic to introduce xarray-einstats.

@covertg

To question #1: I wondered if any type of intelligent broadcasting is possible, something like:

Magic dimension/label based broadcasting is (at least for now) outside the scope of pymc and aesara. IIRC there is an open issue on aesara about this though you might want to search for it. All of this and more however already works after sampling once you get the results as InferenceData.