DiscreteMarkovChain combined with Normal

1

Hi!

This is from the documentation (slightly changed) and works perfectly:

with pm.Model() as model:
    x0 = pm.Categorical.dist(np.ones(2) / 2, size=(size,))
    P = pm.Dirichlet("P", a=[1, 1], size=(2,))
    discrete_mc = DiscreteMarkovChain("MarkovChain", P=P, init_dist=x0, observed=chains)
    idata = pm.sample()

However, this does not work:

with pm.Model() as model:
    x0 = pm.Categorical.dist(np.ones(2) / 2, size=(size,))
    P = pm.Dirichlet("P", a=[1, 1], size=(2,))
    discrete_mc = DiscreteMarkovChain("MarkovChain", P=P, init_dist=x0, steps=steps-1)
    obs = pm.Normal('obs', mu=pm.math.switch(pm.math.eq(discrete_mc, 0), 0, 1), sigma=0.1, observed=chains)
    idata = pm.sample()

Can anybody help me explain why this does not work? What can I do to fix it? Thank you!

Hopefully this is the correct place to ask this, I’m sorry if that is not the case.

The following are less important details:

The important part of the error is

ParallelSamplingError                     Traceback (most recent call last)
Cell In[31], line 7
      5 obs = pm.Normal('obs', mu=pm.math.switch(pm.math.eq(discrete_mc, 0), 0, 1), sigma=0.1, observed=chains)
      6 # discrete_mc = DiscreteMarkovChain("MarkovChain", P=P, init_dist=x0, observed=chains)
----> 7 idata = pm.sample()

File c:\Users\pasca\anaconda3\envs\tno\Lib\site-packages\pymc\sampling\mcmc.py:846, in sample(draws, tune, chains, cores, random_seed, progressbar, progressbar_theme, step, var_names, nuts_sampler, initvals, init, jitter_max_retries, n_init, trace, discard_tuned_samples, compute_convergence_checks, keep_warning_stat, return_inferencedata, idata_kwargs, nuts_sampler_kwargs, callback, mp_ctx, blas_cores, model, **kwargs)
    844 _print_step_hierarchy(step)
    845 try:
--> 846     _mp_sample(**sample_args, **parallel_args)
    847 except pickle.PickleError:
    848     _log.warning("Could not pickle model, sampling singlethreaded.")

File c:\Users\pasca\anaconda3\envs\tno\Lib\site-packages\pymc\sampling\mcmc.py:1259, in _mp_sample(draws, tune, step, chains, cores, random_seed, start, progressbar, progressbar_theme, traces, model, callback, blas_cores, mp_ctx, **kwargs)
   1257 try:
   1258     with sampler:
-> 1259         for draw in sampler:
   1260             strace = traces[draw.chain]
   1261             strace.record(draw.point, draw.stats)

File c:\Users\pasca\anaconda3\envs\tno\Lib\site-packages\pymc\sampling\parallel.py:471, in ParallelSampler.__iter__(self)
    464 task = progress.add_task(
    465     self._desc.format(self),
    466     completed=self._completed_draws,
    467     total=self._total_draws,
    468 )
    470 while self._active:
--> 471     draw = ProcessAdapter.recv_draw(self._active)
    472     proc, is_last, draw, tuning, stats = draw
    473     self._completed_draws += 1

File c:\Users\pasca\anaconda3\envs\tno\Lib\site-packages\pymc\sampling\parallel.py:338, in ProcessAdapter.recv_draw(processes, timeout)
    336     else:
    337         error = RuntimeError(f"Chain {proc.chain} failed.")
--> 338     raise error from old_error
    339 elif msg[0] == "writing_done":
    340     proc._readable = True

IndexError: index 2 is out of bounds for axis 0 with size 2
Apply node that caused the error: AdvancedSubtensor(LogSoftmax{axis=1}.0, Subtensor{start:stop}.0, Subtensor{start:}.0)
Toposort index: 34
Inputs types: [TensorType(float64, shape=(2, 2)), TensorType(int64, shape=(99,)), TensorType(int64, shape=(99,))]
Inputs shapes: [(2, 2), (99,), (99,)]
Inputs strides: [(16, 8), (8,), (8,)]
Inputs values: [array([[-0.01639124, -4.11919274],
       [-0.77331141, -0.61893506]]), 'not shown', 'not shown']
Outputs clients: [[Sum{axes=None}(AdvancedSubtensor.0)]]

The code for initialization is

import arviz as az
import numpy as np
import pymc as pm
import pytensor
import pandas as pd
import statsmodels.api as sm
import matplotlib.pyplot as plt

from matplotlib import ticker as mtick

from pymc_experimental.distributions.timeseries import DiscreteMarkovChain

true_P = np.array([[0.6, 0.4], [0.4, 0.6]])

def generate_chains(P, steps, n_chains=1):
    output = np.empty((n_chains, steps), dtype="int64")

    x0 = pm.draw(pm.Categorical.dist(p=[0.5, 0.5], shape=(n_chains,)))
    output[:, 0] = x0

    for t in range(1, steps):
        output[:, t] = [
            np.random.choice(range(P.shape[0]), p=P[output[i, t - 1]].ravel()).astype(int)
            for i in range(n_chains)
        ]

    return output.squeeze()

steps = 100
chains = generate_chains(true_P, steps, n_chains=100)
1 Like
2

I had to specify size in x0, which I put as (2,). I imagine you had defined that elsewhere just not in the snippets you shared.

I think the problem you’re seeing is that it’s not using a specialized sampler for the DiscreteMarkovChain. That was fixed in: Specialized DiscreteMarkovChain step sampler by ricardoV94 · Pull Request #359 · pymc-devs/pymc-experimental · GitHub and should be available in a next release of pymc-experimental that I’ll try to get out today.

1 Like