Add `DiscreteMarkovChain` distribution #100

jessegrabowski · 2022-12-16T18:13:10Z

Add a distribution for discrete-state Markov chains. Support for deterministic and random initial conditions. Motivated by two recent threads on the discourse asking about this type of model: here and here. This work is derivative of the work presented by @junpenglao and @ricardoV94 in those threads.

Still a lot of work to do so I'm marking this as a draft. In no particular order:

I use .eval() methods in several places to validate inputs, this strikes me as bad but I didn't know what else to do.
Using dims currently breaks the model
The time series dimension is currently specified by the steps argument, but it seems more natural if it were instead set by the size or dims argument (this is related to why dims breaks the model i think)
pm.sample automatically assigns the markov chain RV to pm.Metropolis, don't know how to direct it to BinaryMetropolis or CategoricalMetropolis, depending on the size of the state space.
There's some janky stuff with the steps argument. I internally subtract 1 to account for the fact that x0 will be appended to the scan -- otherwise the resulting chain is length steps + 1, which I don't think will match the user expectation. It gives the right answer but it doesn't feel very clean.

review-notebook-app · 2022-12-16T18:13:14Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

ricardoV94

This looks amazing 🤩 Left some comments below that hopefully you find helpful.

pymc_experimental/distributions/timeseries.py

ricardoV94 · 2022-12-16T21:51:13Z

For the sampler, I think there's a way to make it so your distribution goes with Categorical sampler by default.

pymc_experimental/distributions/timeseries.py

validate `P` in `logp` via `check_parameters` Create `test_discrete_markov_chain.py`

remove shape checks on x0 add `init_dist` argument add support for `dtype` kwarg

`dims` argument now works Add tests associated with `dims`

ricardoV94

Great progress. Some comments/requests below.

Would be so cool to implement logp marginalization of these variables in a future PR! Should be as simple as putting the right DiscreteMarginalizedRV inside Scan? Should come out the same as the forward algorithm.

https://github.com/pymc-devs/pymc-experimental/blob/main/pymc_experimental/marginal_model.py

Another future improvement would be multiple lag dependencies?

Anyway I am getting carried away, this is already very exciting!

pymc_experimental/distributions/timeseries.py

pymc_experimental/tests/distributions/test_discrete_markov_chain.py

pymc_experimental/distributions/timeseries.py

pymc_experimental/tests/distributions/test_discrete_markov_chain.py

pymc_experimental/distributions/timeseries.py

ricardoV94 · 2022-12-18T05:09:32Z

For the sampler, I think there's a way to make it so your distribution goes with Categorical sampler by default.

Speaking of bad APIs...

Here is a solution? Like the test added here you can override the default list of step methods: pymc-devs/pymc@2dd4c8c

You could add a subclass of the categorical metropolis sampler that assigns ideal competence to this variable.

We should definitely refactor that approach, perhaps with dispatching, but for now that is the step assignment API for external libraries...

jessegrabowski · 2022-12-18T13:24:59Z

For the sampler, I think there's a way to make it so your distribution goes with Categorical sampler by default.

Speaking of bad APIs...

Here is a solution? Like the test added here you can override the default list of step methods: pymc-devs/pymc@2dd4c8c

You could add a subclass of the categorical metropolis sampler that assigns ideal competence to this variable.

We should definitely refactor that approach, perhaps with dispatching, but for now that is the step assignment API for external libraries...

I was thinking if I could subclass from Categorical then it would be automatically assigned correctly, but I got some errors doing this.

jessegrabowski · 2022-12-18T16:41:53Z

re: multiple lag dependencies, I think that would be nice as well and should be easy enough to add.
re: Marginalization, I was really hoping this would integrate with that PR as well. The HMM model in particular is quite difficult to sample (see the example notebook as well as discussion here), and I wonder if automatic marginalization could help the sampler along.

ricardoV94 · 2022-12-18T17:06:43Z

Marginalizing HMMs usually sample much better in my experience. We would need to increase the marginalization to support it, though. Right now it only works for vanilla discrete variables, not scans or this specific wrapped scan op either.

Remove unnecessary type handling on distributions.

Add `initval='prior'` as a default argument to `__new__`

jessegrabowski · 2022-12-18T18:42:16Z

One of the scan re-writes is throw an error now, but only if I compile the model outside of a model context (e.g. with pm.draw). Curious if you know what might be causing it?

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: save_mem_new_scan
ERROR (pytensor.graph.rewriting.basic): node: for{cpu,scan_fn}(TensorConstant{(1,) of 9}, IncSubtensor{Set;:int64:}.0, RandomGeneratorSharedVariable(<Generator(PCG64) at 0x151E8EE40>), Softmax{axis=1}.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/Users/jessegrabowski/mambaforge/envs/pymc_dev/lib/python3.9/site-packages/pytensor/graph/rewriting/basic.py", line 1933, in process_node
    replacements = node_rewriter.transform(fgraph, node)
  File "/Users/jessegrabowski/mambaforge/envs/pymc_dev/lib/python3.9/site-packages/pytensor/graph/rewriting/basic.py", line 1092, in transform
    return self.fn(fgraph, node)
  File "/Users/jessegrabowski/mambaforge/envs/pymc_dev/lib/python3.9/site-packages/pytensor/scan/rewriting.py", line 1628, in save_mem_new_scan
    subtens = Subtensor(nw_slice)
  File "/Users/jessegrabowski/mambaforge/envs/pymc_dev/lib/python3.9/site-packages/pytensor/tensor/subtensor.py", line 692, in __init__
    self.idx_list = tuple(map(index_vars_to_types, idx_list))
  File "/Users/jessegrabowski/mambaforge/envs/pymc_dev/lib/python3.9/site-packages/pytensor/tensor/subtensor.py", line 592, in index_vars_to_types
    slice_a = index_vars_to_types(a, False)
  File "/Users/jessegrabowski/mambaforge/envs/pymc_dev/lib/python3.9/site-packages/pytensor/tensor/subtensor.py", line 613, in index_vars_to_types
    raise AdvancedIndexingError("Invalid index type or slice for Subtensor")
pytensor.tensor.exceptions.AdvancedIndexingError: Invalid index type or slice for Subtensor

EDIT:: This was caused by ndims = 1 in steps = at.as_tensor_variable(intX(steps), ndims=1). I copied that from the AR distribution, not sure why it is causing a problem here.

Add test for random draws Re-run notebook

jessegrabowski · 2022-12-18T21:01:05Z

I marked it as ready for review, although the problem of sampler assignment still isn't solved. If it was OK'd for the main PyMC code base, I'd just add the distribution to the list of competencies for CategoricalGibbsMetropolis I guess.

Another small issue is that pm.model_to_graphviz doesn't seem to recognize dependency between the hidden state chain and a sequence of state means when I did the HMM example.

Also still need to add support for multiple lags.

pymc_experimental/distributions/timeseries.py

ricardoV94 · 2022-12-18T21:40:07Z

Also still need to add support for multiple lags.

That need not be a blocker unless you want it to

pymc_experimental/distributions/timeseries.py

Add tests for `n_lags` > 1

jessegrabowski · 2023-04-15T23:23:30Z

I think I addressed pretty much everything. I changed an import in marginal_model.py because I did git rebase on my fork's main but not the branch i set up for this PR; I guess that was pretty stupid. I needed to change the imports for logp. I hope it's not too much of a hassle to fix it.

ricardoV94 · 2023-04-17T13:34:50Z

@jessegrabowski I am happy code and test-wise. Do you need help fixing the conflicts and rebasing or are you up to it? We can merge aftewards

twiecki · 2023-04-17T14:05:05Z

The title of the NB needs to be updated.

jessegrabowski · 2023-04-17T14:09:33Z

The NB was a bit of an afterthought, just to show everyone the distribution works as intended. I could doll it up a bit and make a separate PR into pymc-examples?

twiecki · 2023-04-17T14:13:30Z

@jessegrabowski I updated my comment after I saw that this was pymc-experimental, not pymc. So I think a NB here is a good idea, and you can later doll it up for pymc-examples. But the title needs to be fixed.

ricardoV94 · 2023-04-17T19:37:07Z

The math from the docstrings doesn't seem to render correctly: https://pymcio--100.org.readthedocs.build/projects/experimental/en/100/generated/pymc_experimental.distributions.timeseries.DiscreteMarkovChain.html#pymc_experimental.distributions.timeseries.DiscreteMarkovChain

docs/api_reference.rst

…ults to statsmodels

Use short import name in `api_reference.rst` Co-authored-by: Ricardo Vieira <[email protected]>

…ymc-experimental into discrete-markov

add DiscreteMarkovChainRV

c42c94a

ricardoV94 reviewed Dec 16, 2022

View reviewed changes

ricardoV94 reviewed Dec 17, 2022

View reviewed changes

pymc_experimental/distributions/timeseries.py Outdated Show resolved Hide resolved

ricardoV94 added the enhancements New feature or request label Dec 17, 2022

jessegrabowski added 5 commits December 17, 2022 20:53

remove validate_transition_matrix

fd472bd

validate `P` in `logp` via `check_parameters` Create `test_discrete_markov_chain.py`

remove x0 argument

37ef8e0

remove shape checks on x0 add `init_dist` argument add support for `dtype` kwarg

Add reshape logic to rv_op based on size and init_dist

b4e15db

`dims` argument now works Add tests associated with `dims`

Remove moot TODO comments

a3408ba

Update and re-run example notebook

22bfe17

ricardoV94 reviewed Dec 18, 2022

View reviewed changes

jessegrabowski added 2 commits December 18, 2022 19:33

Update pytensor alias to pt

3aef66d

Remove unnecessary type handling on distributions.

Remove moment method

c2d5fc6

Add `initval='prior'` as a default argument to `__new__`

jessegrabowski added 6 commits December 18, 2022 19:47

Wrap tests into test class

d77337c

Add test for default initial distribution warning

82978f0

Replace .dimshuffle with pt.moveaxis in rv_op

b797af3

Fix scan error

d827586

Add test for random draws Re-run notebook

Add test for change_dist_size

0267801

Add code example to DiscreteMarkovChain docstring

b7794ed

jessegrabowski marked this pull request as ready for review December 18, 2022 20:59

ricardoV94 reviewed Dec 18, 2022

View reviewed changes

pymc_experimental/distributions/timeseries.py Outdated Show resolved Hide resolved

jessegrabowski added 5 commits April 15, 2023 15:16

Updates imports following pymc-devs/pymc#6441

ed2983f

Add a moment function to DiscreteMarkovRV

bbaa81e

Raise NotImplementedError if init_dist is not pm.Categorical

9ac136b

Update example notebook with some new plots

2673b12

Fix a bug that broke n_lags > 1

b2df5a2

Add tests for `n_lags` > 1

Rename test function to correctly match test

daffdc8

jessegrabowski added 3 commits April 17, 2023 21:03

rebase from main

c35bc31

Add timeseries.DiscreteMarkovChain to api_reference.rst

170e20b

Remove check on init_dist

eb23686

ricardoV94 reviewed Apr 17, 2023

View reviewed changes

docs/api_reference.rst Outdated Show resolved Hide resolved

jessegrabowski and others added 10 commits April 17, 2023 22:49

Add DiscreteMarkovChain to distribtuions.__all__

c24a86d

Change example notebook title, add subtitles, add plots comparing res…

ee141ef

…ults to statsmodels

Pass init_dist to all tests to avoid UserWarning

f138cac

Fix flakey test_moment_function test

66a3198

Fix latex in docstring

10e2817

Apply suggestions from code review

7ef7ae6

Use short import name in `api_reference.rst` Co-authored-by: Ricardo Vieira <[email protected]>

Fix latex in docstring

0624d2d

Merge branch 'discrete-markov' of https://github.com/jessegrabowski/p…

f07cdba

…ymc-experimental into discrete-markov

Fix latex in docstring

994e3fb

Fix warning in docstring

3c96dc8

ricardoV94 approved these changes Apr 18, 2023

View reviewed changes

ricardoV94 merged commit 666ac8c into pymc-devs:main Apr 20, 2023

ricardoV94 changed the title ~~add DiscreteMarkovChainRV~~ Add DiscreteMarkovChain distribution Apr 20, 2023

jessegrabowski deleted the discrete-markov branch September 17, 2023 16:32

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add `DiscreteMarkovChain` distribution #100

Add `DiscreteMarkovChain` distribution #100

jessegrabowski commented Dec 16, 2022

review-notebook-app bot commented Dec 16, 2022

ricardoV94 left a comment

ricardoV94 commented Dec 16, 2022

ricardoV94 left a comment •

edited

Loading

ricardoV94 commented Dec 18, 2022 •

edited

Loading

jessegrabowski commented Dec 18, 2022

jessegrabowski commented Dec 18, 2022

ricardoV94 commented Dec 18, 2022

jessegrabowski commented Dec 18, 2022 •

edited

Loading

jessegrabowski commented Dec 18, 2022 •

edited

Loading

ricardoV94 commented Dec 18, 2022

jessegrabowski commented Apr 15, 2023

ricardoV94 commented Apr 17, 2023 •

edited

Loading

twiecki commented Apr 17, 2023 •

edited

Loading

jessegrabowski commented Apr 17, 2023

twiecki commented Apr 17, 2023

ricardoV94 commented Apr 17, 2023

Add DiscreteMarkovChain distribution #100

Add DiscreteMarkovChain distribution #100

Conversation

jessegrabowski commented Dec 16, 2022

review-notebook-app bot commented Dec 16, 2022

ricardoV94 left a comment

Choose a reason for hiding this comment

ricardoV94 commented Dec 16, 2022

ricardoV94 left a comment • edited Loading

Choose a reason for hiding this comment

ricardoV94 commented Dec 18, 2022 • edited Loading

jessegrabowski commented Dec 18, 2022

jessegrabowski commented Dec 18, 2022

ricardoV94 commented Dec 18, 2022

jessegrabowski commented Dec 18, 2022 • edited Loading

jessegrabowski commented Dec 18, 2022 • edited Loading

ricardoV94 commented Dec 18, 2022

jessegrabowski commented Apr 15, 2023

ricardoV94 commented Apr 17, 2023 • edited Loading

twiecki commented Apr 17, 2023 • edited Loading

jessegrabowski commented Apr 17, 2023

twiecki commented Apr 17, 2023

ricardoV94 commented Apr 17, 2023

Add `DiscreteMarkovChain` distribution #100

Add `DiscreteMarkovChain` distribution #100

ricardoV94 left a comment •

edited

Loading

ricardoV94 commented Dec 18, 2022 •

edited

Loading

jessegrabowski commented Dec 18, 2022 •

edited

Loading

jessegrabowski commented Dec 18, 2022 •

edited

Loading

ricardoV94 commented Apr 17, 2023 •

edited

Loading

twiecki commented Apr 17, 2023 •

edited

Loading