twallema / pySODM

Simulating and Optimising Dynamical Models in Python 3
Other
9 stars 2 forks source link

Rename `state_dimensions` --> `dimensions_per_state` in model declaration #60

Closed twallema closed 1 year ago

twallema commented 1 year ago

Describe your fixes/additions/changes

Previously,

class JumpProcess_SIR_SI(JumpProcess):
    """
    A stochastic, age-stratified SIR model for humans, an unstratified SI model for a disease vector (f.i. mosquito)
    """

    states = ['S', 'I', 'R', 'S_v', 'I_v']
    parameters = ['beta', 'gamma']
    stratified_parameters = ['alpha']
    dimensions = ['age_group']
    state_dimensions = [['age_group'],['age_group'],['age_group'],[],[]]

    @staticmethod
    def compute_rates(t, S, I, R, S_v, I_v, alpha, beta, gamma):

        # Calculate total mosquito population
        N = S + I + R
        N_v = S_v + I_v

        rates = {
            'S': [alpha*(I_v/N_v),],
            'I': [(1/gamma)*np.ones(N.shape),],
            'S_v': [np.sum(alpha*(I/N))*np.ones(N_v.shape), (1/beta)*np.ones(N_v.shape),],
            'I_v': [(1/beta)*np.ones(N_v.shape),]
        }

        return rates

Replace by,

class JumpProcess_SIR_SI(JumpProcess):
    """
    A stochastic, age-stratified SIR model for humans, an unstratified SI model for a disease vector (f.i. mosquito)
    """

    states = ['S', 'I', 'R', 'S_v', 'I_v']
    parameters = ['beta', 'gamma']
    stratified_parameters = ['alpha']
    dimensions = ['age_group']
    dimensions_per_state = [['age_group'],['age_group'],['age_group'],[],[]]

    @staticmethod
    def compute_rates(t, S, I, R, S_v, I_v, alpha, beta, gamma):

        # Calculate total mosquito population
        N = S + I + R
        N_v = S_v + I_v

        rates = {
            'S': [alpha*(I_v/N_v),],
            'I': [(1/gamma)*np.ones(N.shape),],
            'S_v': [np.sum(alpha*(I/N))*np.ones(N_v.shape), (1/beta)*np.ones(N_v.shape),],
            'I_v': [(1/beta)*np.ones(N_v.shape),]
        }

        return rates