Hello,
[Note: see below for updates, some of the problems here may be resolved]
I want to use the Simulator to setup a Bayesian ABC for a large Markov Chain Model. By large lets say it is a state space of couple thousand raised to power of 10 (and not sparsely connected). So the DiscreteMarkovChain model is a no go. I don’t even know if this ABC approach is statistically feasible but I wanted to try it on some simulated data. But I ran into an issue, possibly due to my misunderstanding of how pytensor rng should be used. The full code is below. When I run that code I got an “setting an array element with a sequence.” error with a not so very helpful error stack. So I added a line where the output of the simulator class is evalled (see commented out part) and that still led to the same error but the error stack now showing that the problem occurs at
next_state = rng.multinomial(N, p)
After playing around with it more and trying different things I think I was also able to get to the bottom of where this is actually happening:
File ~/miniconda3/envs/ACORG/lib/python3.10/site-packages/pytensor/tensor/random/op.py:401, in RandomVariable.perform(self, node, inputs, outputs)
398 rng = copy(rng)
400 outputs[0][0] = rng
--> 401 outputs[1][0] = np.asarray(
402 self.rng_fn(rng, *args, None if size is None else tuple(size)),
403 dtype=self.dtype,
404 )
Update: After digging into this a little bit more, I feel like I may not be using scan correctly in conjunction with Simulator somehow. In the first step self.rng_fn has as args = [[0], [1]] but then at the point this produces an error, args is
[array([ 2.25634373, 1.18186176, 88.82739814, 2.10452112, 1.81074076]), array([0.02283933]), array([ 0, 100, 10000, 100, 10]), array([11], dtype=int8)]
So this is starting to feel more like not using scan correctly in the context of simulate. Note also that _WFM_sim works fine on its own, it returns a ncycles x ndim array when called and evalled on its own.
Full code (the code does not seem to have any syntax errors in it and runs fine but I don’t know why I could not get it highlighted properly):
import pymc as pm
import numpy as np
import pytensor as pt
def _calculate_P_pyt(state, f, m, nvariants):
'''
calculate the transition from state to the next state
'''
denom = state*f
M = pt.tensor.stack([m for _ in range(nvariants)])
M = pt.tensor.tile(M, (1,nvariants))/(nvariants-1)
mask = np.eye(nvariants, dtype=bool)
M = M[mask].set(1-m)
return pt.tensor.sum(M*denom, axis=1)/pt.tensor.sum(denom)
def _WFM_sim(rng, f, m, initial_state, nsims, size=None):
'''
given an initial state, generate a markov chain of length nsims
'''
nvariants = initial_state.size
N = initial_state.sum()
#rng = pt.shared(rng)
def transition(*args):
state, f, m = args
p = _calculate_P_pyt(initial_state, f, m, nvariants)
next_state = rng.multinomial(N, p)
# my other attempt, in this case uncomment pt.shared(rng) above
#next_state = pm.Multinomial.dist(N, p, rng=rng)
return next_state
result, updates = pt.scan(transition,
outputs_info=initial_state,
non_sequences=[f, m],
n_steps=nsims)
return pt.tensor.flatten(result)
def WFM_ABC(initial_state, data):
nsims = data.shape[0]
with pm.Model():
log10_m = -1-pm.HalfNormal("log10_m", 3)
log_f = pm.LogNormal("log_f", 0, 1, size=initial_state.size)
f = pm.math.exp(log_f)
m = 10**log10_m
s = pm.Simulator("s", _WFM_sim, params=(f, m, initial_state, nsims),
epsilon=1, observed=data.flatten())
# commenting this out gives a more informative error
# s.eval()
idata = pm.sample_smc()
idata.extend(pm.sample_posterior_predictive(idata))
# rows are the observed states in a Markov chain starting with i=1
obs_data = np.array([[ 0, 123, 9990, 85, 12],
[ 0, 99, 10012, 93, 6],
[ 1, 127, 9958, 117, 7],
[ 0, 109, 9969, 120, 12],
[ 0, 125, 9973, 99, 13],
[ 1, 116, 9982, 103, 8],
[ 0, 88, 10022, 89, 11],
[ 1, 107, 9995, 101, 6],
[ 0, 123, 9982, 99, 6],
[ 0, 99, 10006, 92, 13]])
initial_state = np.array([ 0, 100, 10000, 100, 10])
WFM_ABC(initial_state, obs_data)