Developer tutorial: Flexible parameters#

Defining default parameters#

When you create a module using Starsim, you have the opportunity to define the default format for parameters. Let’s look at an example from the SIR disease model:

[1]:
import starsim as ss

class SIR(ss.SIR):
    def __init__(self, pars=None, **kwargs):
        super().__init__()
        self.label = 'SIRdemo'
        self.default_pars(
            beta = 0.1,
            init_prev = ss.bernoulli(p=0.01),
            dur_inf = ss.lognorm_ex(mean=6),
            p_death = ss.bernoulli(p=0.01),
        )
        self.update_pars(pars, **kwargs)

        self.add_states(
            ss.BoolArr('recovered', label='Recovered'),
            ss.FloatArr('ti_recovered', label='Time of recovery'),
            ss.FloatArr('ti_dead', label='Time of death'),
        )
        return
Starsim 2.0.0 (2024-10-01) — © 2023-2024 by IDM

The point of self.default_pars() is to set the ground truth for the format that the parameters should take. When users enter their own parameters for defining an instance of this module, the parameter values they enter will be processed within self.update_pars() and will be checked for consistency with the format provided in the default_pars. For example, the parameter p_death in the SIR example above is specified within the default_pars as a Bernoulli distribution. It would be perfectly legitimate to create an instance of the SIR model using any of the following formats:

[2]:
sir1 = ss.SIR(p_death=0.02)
sir2 = ss.SIR(p_death=ss.bernoulli(p=0.2))
sir3 = ss.SIR(pars=dict(p_death=0.03))

However, it would NOT be ok to create an SIR model with e.g. ss.SIR(p_death=ss.lognorm_ex(4)), because if a distribution is defined as a Bernoulli in the default_pars, it can’t be changed. This is only the case for Bernoulli distributions; other distributions can be changed, e.g. ss.SIR(dur_inf=ss.normal(4)) would be okay.

Using callable parameters#

One of the most flexible aspects of how Starsim’s distributions are defined is that they can take callable functions as parameter values. For example, in reality the duration of infection of a disease might vary by age. We could model this as follows:

[3]:
sir = ss.SIR(dur_inf=ss.normal(loc=10))  # Define an SIR model with a default duration of 10 days
sir.pars.dur_inf.set(loc = lambda self, sim, uids: sim.people.age[uids] / 10)  # Change the mean duration so it scales with age
sim = ss.Sim(n_agents=1e3, diseases=sir, networks='random')
sim.run()
Initializing sim with 1000 agents
  Running 2000.0 ( 0/51) (0.00 s)  ———————————————————— 2%
  Running 2010.0 (10/51) (0.05 s)  ••••———————————————— 22%
  Running 2020.0 (20/51) (0.09 s)  ••••••••———————————— 41%
  Running 2030.0 (30/51) (0.12 s)  ••••••••••••———————— 61%
  Running 2040.0 (40/51) (0.15 s)  ••••••••••••••••———— 80%
  Running 2050.0 (50/51) (0.18 s)  •••••••••••••••••••• 100%

[3]:
Sim(n=1000; 2000—2050; networks=randomnet; diseases=sir)

Using similar logic, any other parameter could be set to depend on anything that the sim is aware of, including time or agent properties like age, sex, or health attributes.