Ok so numpyfing the relevant parts of the code I was able to get the code to run, more or less (sorry for using the forums as a journal, I am just updating so that people don’t spend time on something that may already be solved!). The estimates are not great so I am trying to figure out why that is (any insight into this is welcome, perhaps what I am doing is not very feasible?) but in case anyone is interested the full code is below.
One question I have is:
In this case my observables are multinomial counts with roughly totalling to 10000. If I use the counts in simulator for distance comparison, the sampler gets stuck (posteriors are fixed numbers). So I set the epsilon to sum of initial_states which makes it a proportion comparison. However it finishes quite quickly. The estimates aren’t great which I wonder is something that happens with sample_smc? I also tried epsilon=1000 instead of sum of initial_state, or playing with thresholds, sometimes it takes more time but does not improve the estimates much. Should I instead try something with potentials?
As an example, with parameters:
f = np.array([0.67786154, 1.89172801, 0.84634297, 0.63536372, 1.19762663])
m = 1e-4
I get (using epsilon=1000 and default for anything else, also don’t mind the log10_m since I actually use -2 - log10_m inside the model, could have named that better!)
mean sd hdi_3% hdi_97% ...
f[0] 0.970 0.087 0.804 1.127
f[1] 1.546 0.094 1.376 1.727
f[2] 0.724 0.046 0.638 0.808
f[3] 0.968 0.085 0.817 1.137
f[4] 0.966 0.082 0.817 1.123
log10_m[0] 0.936 0.592 0.017 1.976
It seems to get it roughly right but not quite. Note that if you check the model, there is an invariance with respect to f, that is f and af for some a>1 would give the same results (because of the division inside _calculate_P. I however don’t think it is the issue here because sd are low and for instance fitted f[0] is larger than real value but f[1] is lower, f[2] is lower, f[3] is larger, f[4] is lower etc. Fixing value of m, or using more observables does not change the result.
import pymc as pm
import numpy as np
import pytensor as pt
import arviz as az
def _WFM_sim(rng, f, m, initial_state, nsims, size=None):
nvariants = initial_state.size
N = initial_state.sum()
# for some reason simultor class expands dimensions and turns constants
# into arrays so need to deal with that
initial_state = initial_state[0]
nsims = nsims.flatten()[0]
state_history = [initial_state]
f = f.flatten()
for i in range(nsims):
P = _calculate_P(state_history[-1], f, m, nvariants)
state_history.append(rng.multinomial(N, P))
out = np.array(state_history)[1:,:]
return out
def _calculate_P(state, f, m, nvariants):
'''
m is assumed to be of size 1 x nvariants
f and state of size nvariants
'''
denom = state*f
M = m.copy()
M = np.tile(M, (nvariants,1))/(nvariants-1)
mask = np.eye(nvariants, dtype=bool)
M[mask] = 1 - m[0,:]
return np.sum(M*denom,axis=1)/np.sum(denom)
def WFM_ABC(data):
initial_state = data[0,:]
obs = data[1:,:]
nsims = obs.shape[0]
with pm.Model():
log10_m = -2-pm.HalfNormal("log10_m", 1, size=1)
f = pm.LogNormal("f", 0, 0.1, size=initial_state.size)
m = pt.tensor.tile(10**log10_m[:,None], (1,initial_state.size))
pm.Simulator("obs", _WFM_sim, params=(f, m, initial_state, nsims),
epsilon=1000, observed=obs)
idata = pm.sample_smc()
return idata
seed = 0
f = np.array([0.67786154, 1.89172801, 0.84634297, 0.63536372, 1.19762663])
m = 1e-4
rng = np.random.default_rng(0)
initial_state = np.array([[100, 100, 10000, 100, 100]])
nobs = 20
obs_data=\
_WFM_sim(rng, f, np.array([[m for _ in range(f.size)]]),
initial_state, np.array([nobs]))
idata = WFM_ABC(obs_data)
print(az.summary(idata))