Add time varying transition probabilities for DiscreteMarkovChain

[ddgpalmer]

Hi. Jesse Grabowski recommended I raise an issue here based on a question I asked in PyMC discourse. Would it be possible to add capacity for time-varying transition probabilities in the DiscreteMarkovChain class? For my own purposes, I would like to use this to use a different transition matrix on a few designated transitions while using the 'main' on on others, but it's possible to imagine a number of possible use cases.

[ricardoV94]

While conceptually it is not hard the API will be tricky because we have 4+ dimensions with different meanings all in the same transition array: *batch dimensions, lags, time, probability.

We also don't allow time varying parameters in any PyMC time series and ask users to implement a CustomDist when they need it to keep our API simpler.

We can offer it but need to decide how to do it. Add a time_varying flag to know dimension x will be time and not lags or batch?

[ricardoV94]

Also time varying parameters tend to be impractical with large time series. Usually you don't really want to define a unique transition per time step but something that changes after x steps or cycles back and forth or is parametrized at runtime by another variable.

Again for this it would be better to define your custom implementation. The pre built distribution couldn't offer this functionality without basically reinventing Scan

[jessegrabowski]

I'm a llittle bit more bullish on this. I do something similar in the kalman filters, where I inspect the incoming matrices and check if there's a batch dimension. If there is, the variable is assigned to sequence in the scan. Otherwise, it goes to non-sequence.

Right now, P has (n_lags - 1) + 2 dimensions. We get all that information up front, so it would be easy enough to check if there's a batch dim. If you think that's too black magic, we could also add a "time_varying_P" flag that would do the same thing in a more transparent way (we would still do the dim check to make sure the user is doing the right thing)

[ricardoV94]

batch dim could just be a batch dimension, which we should definitely support as that's the case in all our distributions. When you write GRW(... sigma=vector) it means you have a vector of independent GRW each with scalar sigma (a matrix of values).

If you always assume the dim before lags is time, users will need to add a dummy one to skip it and distinguish from the first batch dimension. That's not a big problem, the point is you need to disambiguate when you don't want time varying effects, either with a dummy dim or a flag

I'm still suspecting that in most cases you probably want something that's parametrized inside the scan step to avoid allocating huge arrays outside of the scan. If we think I'm wrong I don't have any other objection.

The example in the discourse you wrote fits my suspicion though. Only two matrices that cycle back and forth