Using pytensor random number generator with scan for Simulator class

Questions v5
11

No in fact it is the quite opposite of sparse. Jumps from one state to another are modelled with Multinomials so anything can go anywhere. Ofcourse you could try to cut it down beyond a threshold. But with such a large state space, would sparse matrices even help?

It is actually the Wright-Fisher model with fitness and mutations (a very commonly used model but generally with 2 variants). The description of the state space I used in the code is as follows: There are six variants each variant can have numbers up to some fixed N (say on the order of thousands) and the sum of variants is also equal to N. So state space is roughly N^6.

Transition is calculated here:

def transition(*args):

    state,  f, m = args

    p = _calculate_P_pyt(state, f, m, nvariants)
    next_rng, next_state = pm.Multinomial.dist(N, p, rng=rng).owner.outputs

    return next_state, {rng: next_rng}

The _calculate_P_pyt just calculates some probability but transitioning to any other state is a multinomial.