twallema / pySODM

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

Order of non-stratified vs. stratified parameters in the `integrate` function shouldn't matter when initializing a model #9

Open twallema opened 1 year ago

twallema commented 1 year ago

Suggestion

Consider an SIR model stratified into two age groups (0-20 yo, 20-120 yo) with two parameters: beta, the per-contact chance of infection, and gamma, the recovery rate of the infected individuals. In the example beta is stratified and gamma is not.

class SIRstratified(ODEModel):

    # state variables and parameters
    state_names = ['S', 'I', 'R']
    parameter_names = ['gamma']
    parameters_stratified_names = ['beta']
    stratification_names = ['age_groups']

    @staticmethod
    def integrate(t, S, I, R, beta, gamma):
        """Basic SIR model"""
        # Model equations
        N = S + I + R
        dS = -beta*S*I/N
        dI = beta*S*I/N - gamma*I
        dR = gamma*I
        return dS, dI, dR

parameters = {"gamma": 0.2, "beta": np.array([0.8, 0.9])}
initial_states = {"S": [600_000 - 20, 400_000 - 10], "I": [20, 10], "R": [0, 0]}
coordinates = {'age_groups': ['0-20','20-120']}
model = SIRstratified(initial_states, parameters, coordinates=coordinates)

Providing the parameter beta before gamma yields the following error:

ValueError: The parameters in the 'integrate' function definition do not match the parameter_names + parameters_stratified_names + stratification: ['beta', 'gamma'] vs ['gamma', 'beta']

Ideally, the order in which the parameters are given should not matter when initializing the model.

twallema commented 1 year ago

This issue is partly resolved by PR #19. The order of the states, parameters, stratified parameters supplied to the integrate function does not matter as they are now keyworded. However, the user should still return the differentials in the same order as the state_names variable. Perhaps this can be resolved by using a dictionary, however, this may complicate matters when applying numba.njit to the integrate function. To be reconsidered at some later point.