Yes, automatic logp inference for the DiscreteMarkovChain variable isn’t supported yet, it’s on my to-do list. Notice that in the example, it uses the HMM as a mean with gaussian noise, so it’s not 100% correct.
You could probably just ignore the DiscreteMarkovChain random variable for now and handle the state transitions inside the scan. Just pass the P matrix as a non_sequence, and use it to parameterize a Bernoulli at each step.
Your solution for additional jumps was the same as what I had in mind. The only glitch is that it doesn’t have memory of what the last shocked level was – when the state transitions back to “repaired”, it will jump all the way to zero, rather than a deque of shocks.
Having more states in the HMM might be a way to fix it; you’ll have to carry the information about the levels of all the states “under” the current state. But that should be possible, you’ll just have a mu1 and mu2, and when the state transitions from 0->1 you update mu1 and keep carrying it along. When the state transitions from 1->2 you update mu2 = mu1 + shock and carry it along, then when you transitions from 2->1 you bring back mu1 and set mu2 back to 0, when you transition from 1->0 set mu1 back to 0. Define mu = pt.stack([0, mu1, mu2]), then the current value is just mu[state] (after you apply all the update logic at each time step)
No reason, I was just being cute.
If you use pm.CustomDist you won’t need to do this. But to get access to custom dist you can’t use the DiscreteMarkovChain, as noted above. I recommend trying to get it up and running just with Bernoulli RVs inside the scan.