Discrete dynamics

discrete_dynamic.py

import matplotlib.pyplot as plt
import numpy as np
from scipy.stats import multinomial
import cmdstanpy



def simulate(n_transitions):
    A = 1000
    B = 100
    C = 400
    N = A + B + C
    r1 = 0.02
    r2 = 0.02
    r3 = 0.03
    r4 = 0.01

    transition_matrix = np.array([
        [1-r1, 0, r4],
        [r1, 1-r2, r3],
        [0, r2, 1-r4-r3]
    ])

    X = np.zeros((n_transitions, 3))
    theta = np.zeros((n_transitions, 3))

    X[0] = (A, B, C)
    theta[0] = X[0]/sum(X[0])

    for i in range(1, n_transitions):
        
        theta[i] = np.linalg.matrix_power(transition_matrix, i) @ theta[0]

        Aout = multinomial(n=A, p=[1-r1, r1, 0]).rvs().ravel()
        Bout = multinomial(n=B, p=[0, 1-r2, r2]).rvs().ravel()
        Cout = multinomial(n=C, p=[r4, r3, 1-r3-r4]).rvs().ravel()

        A += - np.delete(Aout, 0).sum() + Cout[0] + Bout[0]
        B += - np.delete(Bout, 1).sum() + Cout[1] + Aout[1]
        C += - np.delete(Cout, 2).sum() + Bout[2] + Aout[2]

        X[i] = (A, B, C)
        
    return X, theta*N


X, t = simulate(365)