Skip to main content

internals.modules

internals.modules

The base classes of the parts of a sim: ss.modules.Base and ss.Module, which diseases, networks, interventions, and the rest inherit from

They exist so that every part of a model has the same parameters, timeline, results, and place in the integration loop.

Used with: agents and compartments; declared and coded modules

Classes

Name Description
Base The parent class for Sim and Module objects
Module The main base class for all Starsim modules: diseases, networks, interventions, etc.
Preference A reason in agent_reasons() for solver='auto' to run a module as agents that does not stop a compartmental solver that was asked for

Base

internals.modules.Base()

The parent class for Sim and Module objects

It holds what both have in common, such as how they are shown.

Attributes

Name Description
dt Get the current module timestep
now Shortcut to self.t.now()
ti Get the current module time index
timevec Shortcut to self.t.timevec

Methods

Name Description
copy Perform a deep copy of the module/sim, as it is (for a sim, with its results); to run a sim again, use sim.run(), and for a new, uninitialized sim, sim.new()
disp Display the full object
copy
internals.modules.Base.copy(die=True)

Perform a deep copy of the module/sim, as it is (for a sim, with its results); to run a sim again, use sim.run(), and for a new, uninitialized sim, sim.new()

Parameters
Name Type Description Default
die bool whether to raise an exception if copy fails (else, try a shallow copy) True
disp
internals.modules.Base.disp(output=False, **kwargs)

Display the full object

Parameters
Name Type Description Default
output bool whether to return the text rather than print it False
**kwargs dict passed to sc.prepr()

Module

internals.modules.Module(name=None, label=None, **kwargs)

The main base class for all Starsim modules: diseases, networks, interventions, etc.

By convention, keyword arguments to modules are assigned to “_“, which is an alias for None. These are then populated with defaults by mod.define_pars(), which are then updated with mod.update_pars(), which inspects the __init__ function to get these arguments.

Note that there is no functional difference between specifying arguments this way rather than simply via **kwargs, but having the arguments shown in the function signature can make it easier to read, and “_” is a convention to indicate that the default is specified below.

It is also of course OK to specify the actual values rather than “_“; however module arguments are often complex objects (e.g. ss.bernoulli) that can’t be easily specified in the function signature.

Finally, note that you can (and should) call super().__init__() with no arguments: the name, label, and time arguments get correctly updated via self.update_pars().

Parameters

Name Type Description Default
name str a short, key-like name for the module (e.g. “randomnet”) None
label str the full, human-readable name for the module (e.g. “Random network”) None
kwargs dict passed to ss.Timeline() (e.g. start, stop, dt)

Examples

class SIR(ss.Module):
    def __init__(self, pars=_, beta=_, init_prev=_, p_death=_, **kwargs):
        super().__init__() # Call this first with no arguments
        self.define_pars( # Then define the parameters, including their default value
            beta = ss.peryear(0.1),
            init_prev = ss.bernoulli(p=0.01),
            p_death = ss.bernoulli(p=0.3),
        )
        self.update_pars(pars, **kwargs) # Update with any user-supplied parameters, and raise an exception if trying to set a parameter that wasn't defined in define_pars()
        return

Attributes

Name Description
auto_state_list List of “automatic” states with boolean type (ss.BoolState) that were added
derived_state_names Names of the aliases that are defined by a callable (see define_aliases())
state_dict Return a flat dictionary (objdict) of all states
state_list Return a flat list of all states (ss.Arr objects)

Methods

Name Description
advance For a module with a timeline of its own: move its time to the step about to be taken (the sim moves the time of the modules on its own timeline: see ss.Sim.advance())
agent_reasons What needs agents in this module, from what it is rather than by running it, for choosing the solver (see ss.solvers.rules):
brief Show a brief representation of the module, with any extra text after its name (e.g. a network’s number of edges); used by repr
check_bare_units Warn about parameters given as bare numbers in a unit other than the sim’s (see ss.Pars.bare_units())
check_method_calls Check if any required methods were not called.
check_not_given Raise if this module was given to a sim, which runs its own copy, so that action on this one would never be seen
create Create a module instance by name
current_value A parameter (key is its name) with the modifiers active on it (ss.modify('tx.pars.p_treat', ...)), in its own type: one value, or one for each agent if they differ between agents
define_aliases Define aliases for states, e.g. self.define_aliases(infectious='infected')
define_pars Create or merge Pars objects
define_results Add results to the module; in __init__(), they are kept, and added once the sim’s timeline is known
define_states Define states of the module with the same attribute name as the state
finalize Perform any final operations, such as removing unneeded data
finalize_results Finalize results
finish_step Define what should happen at the end of the step, after the results are recorded
from_func Create a module from a function
init_mock Initialize with a mock simulation – for debugging purposes only
init_post Initialize the values of the states, and link its schedules (ss.schedule) to its timeline; the last step of initialization
init_pre Perform initialization steps
init_results Initialize results output; called during init_pre()
link_timepars Find all time parameters in the module and link them to the module’s dt
match_time_inds Find the nearest matching sim time indices for the current module
on_birth What to do when people are born (e.g. to give the children traits from their parents); by default, nothing
on_death What to do when agents die or leave; by default, nothing (see ss.Disease.on_death() for diseases)
plot Plot the module’s results
record_events Record the results about what happened in the step (flows), at its index: e.g. self.results.doses[self.ti] = self.n_doses
record_states Record the results that are the state at the start of the step (stocks), at its index; by default, the count of each state
set_metadata Set metadata for the module
setattribute Method for setting an attribute that does not perform checking against immutable attributes
shrink Shrink the size of the module for saving to disk
start_step Tasks to perform at the beginning of the step
step Define how the module updates over time – the key part of Starsim!!
to_json The module as a reference by its type and name, for sc.jsonify(): e.g. a module among the parameters in sim.to_json(), or in a parameter of another module
update_pars Pull out recognized parameters and update them; raises on unrecognized keys
advance
internals.modules.Module.advance()

For a module with a timeline of its own: move its time to the step about to be taken (the sim moves the time of the modules on its own timeline: see ss.Sim.advance())

agent_reasons
internals.modules.Module.agent_reasons(sim)

What needs agents in this module, from what it is rather than by running it, for choosing the solver (see ss.solvers.rules): each a sentence, which stops a compartmental solver that was asked for, or a ss.modules.Preference, which does not

Here, for a module written as code: compartmental = True or False on it (or its class) says whether it works on groups; without it, keeping a value for each person (define_states()) needs agents. A subclass can say more.

Parameters
Name Type Description Default
sim ss.Sim the sim, with its groups and unit worked out, but no solver yet required
brief
internals.modules.Module.brief(output=False, extra='')

Show a brief representation of the module, with any extra text after its name (e.g. a network’s number of edges); used by repr

Parameters
Name Type Description Default
output bool whether to return the text rather than print it False
extra str text to show after the name ''
check_bare_units
internals.modules.Module.check_bare_units(sim)

Warn about parameters given as bare numbers in a unit other than the sim’s (see ss.Pars.bare_units())

check_method_calls
internals.modules.Module.check_method_calls()

Check if any required methods were not called.

Typically called automatically by sim.run().

check_not_given
internals.modules.Module.check_not_given(action)

Raise if this module was given to a sim, which runs its own copy, so that action on this one would never be seen

create
internals.modules.Module.create(name, *args, **kwargs)

Create a module instance by name

Parameters
Name Type Description Default
name str A string with the name of the module class in lower case, e.g. ‘sir’ required
current_value
internals.modules.Module.current_value(key, uids=None)

A parameter (key is its name) with the modifiers active on it (ss.modify('tx.pars.p_treat', ...)), in its own type: one value, or one for each agent if they differ between agents

A module written as code reads a parameter from self.pars with the modifiers applied, which is one value. To apply a modifier that differs between people (e.g. with who=), its code reads the parameter with current_value(), and says so in its class, allow_modifiers = ('p_treat',).

Parameters
Name Type Description Default
key str the name of the parameter required
uids uids the agents to give the values of (default: one value, or an array of every agent, by UID, if they differ between agents) None
Example
class Treat(ss.Intervention):
    allow_modifiers = ('p_treat',) # ss.modify('treat.pars.p_treat', multiply=0.5, who=...) applies to those people only
    def step(self):
        uids = self.sim.diseases.sis.infected.uids
        treated = uids[np.random.rand(len(uids)) < self.current_value('p_treat', uids)] # ...
define_aliases
internals.modules.Module.define_aliases(overwrite=False, **kwargs)

Define aliases for states, e.g. self.define_aliases(infectious='infected')

An alias is only used if the attribute isn’t otherwise defined, so it acts as a default that a subclass overrides simply by defining a state of the same name. E.g. ss.AgentDisease treats everyone infected as infectious this way, while letting a subclass define infectious as a state (or another alias) in its own right.

Each alias is either the name of another state, or a callable taking the module as its only argument. A callable behaves like a property: it is recomputed on each access, so it can be read but not written to.

Note: a callable alias defined as a lambda can’t be saved with plain pickle (sc.save() and ss.MultiSim are both fine, since they use dill). Use a module-level function, or a property, if plain pickling is required.

A callable alias also generates an automatic result, in the same way a BoolState does: an alias named infected produces the result infected. A string alias does not, since the state it points to is already counted under its own name.

Parameters
Name Type Description Default
overwrite bool if True, replace any existing state or alias of the same name, rather than raising False
kwargs dict mapping of alias name to either the name of a state, or a callable
Examples
self.define_aliases(infectious='infected') # Everyone infected is infectious
self.define_aliases(infected=lambda self: self.exposed | self.infectious) # Both E and I are infected
self.define_states(ss.BoolState('exposed'), infectious='infected') # Aliases can also be defined alongside states
define_pars
internals.modules.Module.define_pars(inherit=True, **kwargs)

Create or merge Pars objects

Parameters
Name Type Description Default
inherit bool whether to merge with the parameters already defined, rather than replace them True
**kwargs dict the parameters
define_results
internals.modules.Module.define_results(*args)

Add results to the module; in __init__(), they are kept, and added once the sim’s timeline is known

Parameters
Name Type Description Default
*args ss.Result/list/dict the results; a list or dict is passed to ss.Result()
define_states
internals.modules.Module.define_states(
    *args,
    reset=False,
    lock=True,
    overwrite=False,
    **kwargs,
)

Define states of the module with the same attribute name as the state

In addition to registering the state with the module by attribute, it adds it to mod._auto_states, which is used by mod.state_list and mod.state_dict.

Any keyword arguments other than the ones listed below are passed to define_aliases(), so states and their aliases can be defined together.

Parameters
Name Type Description Default
args states list of states to add
reset bool / str / list if True, remove all existing states and aliases; if a name or list of names, remove only those False
lock bool if True, prevent the state attributes from being overwritten after definition True
overwrite bool if True, replace any existing state or alias of the same name, rather than raising False
kwargs dict passed to define_aliases()
Examples
self.define_states(ss.BoolState('exposed'), reset='ti_infected') # Add a state and remove an inherited one
self.define_states(ss.BoolState('infectious'), overwrite=True) # Replace an inherited state or alias
finalize
internals.modules.Module.finalize()

Perform any final operations, such as removing unneeded data

finalize_results
internals.modules.Module.finalize_results()

Finalize results

finish_step
internals.modules.Module.finish_step()

Define what should happen at the end of the step, after the results are recorded

from_func
internals.modules.Module.from_func(func)

Create a module from a function

Wraps a plain function into a module of the calling class, so that simple, one-off logic can be used anywhere a module is expected without writing a full class. The function must take the sim as its only argument; it becomes the module’s step() method and is therefore called once on every timestep (so it must do its own time gating, e.g. if sim.now == 2015:). The module’s name is taken from the function’s name.

This is the mechanism behind “function-based” interventions and analyzers: passing a function to interventions=, analyzers=, or modules= automatically calls from_func() on the appropriate class (ss.Intervention, ss.Analyzer, or ss.Module respectively), so usually you don’t need to call this method directly. Call it explicitly when you want to construct the module yourself, e.g. to add it to a sim later.

Keyword arguments of the function with defaults are the module’s parameters, passed to it at each call; one given as ss.par with bounds can be fitted (sim.fit()), sampled (sim.sensitivity()), and set by name (ss.Sim(pars=...), ss.parallel(pars=...)) as any other. This is how to fit a value that is not a parameter of any module, e.g. one set in code (see the second example).

Function-based modules are otherwise minimal: they have no states or results of their own. As soon as you need define_states() or define_results(), write a proper subclass instead.

Parameters
Name Type Description Default
func callable a function taking the sim, and optionally keyword arguments with defaults (its parameters) required
Returns
Name Type Description
mod Module an instance of cls whose step() calls func(sim)
Example
import starsim as ss

def knockdown(sim): # Halve transmission for one year starting in 2015
    if sim.now == 2015:
        sim.diseases.sis.rel_trans[:] *= 0.5

# Explicit construction ...
intv = ss.Intervention.from_func(knockdown)
sim = ss.Sim(diseases='sis', networks='random', interventions=intv)

# ... or equivalently, let the sim convert it automatically:
sim = ss.Sim(diseases='sis', networks='random', interventions=knockdown)

# A value set in code, fitted: the fraction of transmission left in 2015
def knockdown(sim, frac=ss.par(0.5, bounds=[0.1, 1])):
    if sim.now == 2015:
        sim.diseases.sis.rel_trans[:] *= frac
init_mock
internals.modules.Module.init_mock(n_agents=100, dur=10)

Initialize with a mock simulation – for debugging purposes only

Parameters
Name Type Description Default
n_agents int the number of agents in the mock sim 100
dur float the duration of the mock sim 10
init_post
internals.modules.Module.init_post()

Initialize the values of the states, and link its schedules (ss.schedule) to its timeline; the last step of initialization

init_pre
internals.modules.Module.init_pre(sim, force=False)

Perform initialization steps

This method is called once, as part of initializing a Sim. Note: after initialization, initialized=False until init_post() is called (which is after distributions are initialized).

Parameters
Name Type Description Default
sim ss.Sim the sim required
force bool whether to run again if already done False
init_results
internals.modules.Module.init_results()

Initialize results output; called during init_pre()

By default, modules all report on counts for any explicitly defined “States”, e.g. if a disease contains an ss.BoolState called ‘susceptible’ it will automatically contain a Result for ‘susceptible’ (the number susceptible). For identical behavior that does not automatically generate results, use ss.BoolArr instead of ss.BoolState.

match_time_inds
internals.modules.Module.match_time_inds(inds=None)

Find the nearest matching sim time indices for the current module

Parameters
Name Type Description Default
inds array indices to return if the module and sim have the same number of timesteps; if None, all of them None
on_birth
internals.modules.Module.on_birth(parents, children)

What to do when people are born (e.g. to give the children traits from their parents); by default, nothing

Called for every module that defines it, wherever children are born (by ss.Births, at a delivery of ss.Pregnancy, or by a declared birth transition), once they are people.

Parameters
Name Type Description Default
parents ss.uids the parent of each child (so a parent of several children appears once for each) required
children ss.uids the children required
on_death
internals.modules.Module.on_death(uids)

What to do when agents die or leave; by default, nothing (see ss.Disease.on_death() for diseases)

Called by people.step_die() for every module that defines it, once the deaths of the step are resolved, each to one cause (people.cause_of_death; people.died_of(module) gives those of a module), and before the results are recorded. It isn’t called in a step with no deaths.

Parameters
Name Type Description Default
uids ss.uids the agents who died or left in this step required
plot
internals.modules.Module.plot(*keys, **kwargs)

Plot the module’s results

Parameters
Name Type Description Default
keys str / list what to plot: 'states' (the number in each state, together in one axes; the default, if the module has states), 'flows' (each transition’s flow, in axes of its own), 'all' (every result, each in axes of its own; the default otherwise), or the names of results: a list of them in one axes (['S', 'I']), and each argument in axes of its own ('S', 'I')
kwargs dict passed to ss.Results.plot(), e.g. fig, figsize, annualize
Examples
sir.plot()           # The states
sir.plot('flows')    # Each flow
sir.plot(['S', 'I']) # S and I, together
sir.plot('S', 'I')   # S and I, apart
record_events
internals.modules.Module.record_events()

Record the results about what happened in the step (flows), at its index: e.g. self.results.doses[self.ti] = self.n_doses

Called at the end of every step, after every module’s step and the deaths (see ss.loop.Loop). The state then is that of the next time point, so a result that also needs a stock reads it from its result, recorded at the start of the step (e.g. self.sim.results.alive[self.ti]), or is worked out over the whole run in finalize_results().

record_states
internals.modules.Module.record_states(names=None)

Record the results that are the state at the start of the step (stocks), at its index; by default, the count of each state

Called at the start of every step, before any module’s step (see ss.loop.Loop), so the results at a time point are the state there: e.g. self.results.n_on_art[self.ti] = self.on_art.sum(). A result about what happens in the step (a flow, or one that needs a flow) goes in record_events() instead.

Parameters
Name Type Description Default
names list the states to count (default: every ss.BoolState and derived state) None
set_metadata
internals.modules.Module.set_metadata(name=None, label=None)

Set metadata for the module

Parameters
Name Type Description Default
name str the name of the module; if None, the lowercase name of its class None
label str the label of the module; if None, its name None
setattribute
internals.modules.Module.setattribute(attr, value)

Method for setting an attribute that does not perform checking against immutable attributes

shrink
internals.modules.Module.shrink()

Shrink the size of the module for saving to disk

start_step
internals.modules.Module.start_step()

Tasks to perform at the beginning of the step

step
internals.modules.Module.step()

Define how the module updates over time – the key part of Starsim!!

to_json
internals.modules.Module.to_json()

The module as a reference by its type and name, for sc.jsonify(): e.g. a module among the parameters in sim.to_json(), or in a parameter of another module

update_pars
internals.modules.Module.update_pars(pars=None, **kwargs)

Pull out recognized parameters and update them; raises on unrecognized keys

Parameters
Name Type Description Default
pars dict the parameters to update None
**kwargs dict more parameters, merged into pars

Preference

internals.modules.Preference()

A reason in agent_reasons() for solver='auto' to run a module as agents that does not stop a compartmental solver that was asked for

It exists so that a reason says this in one word.

Functions

Name Description
find_modules Find all subclasses of Module present in Starsim, divided by type
is_module Check whether something is a Starsim module
module_map Map modules to standard types
module_types Return a list of known module types; based on module_map()
register_modules Register custom modules with Starsim so they can be referred to by string.
required Decorator to mark module methods as required.

find_modules

internals.modules.find_modules(key=None, flat=False, verbose=False)

Find all subclasses of Module present in Starsim, divided by type

Parameters

Name Type Description Default
key str the type of module to return, e.g. ‘diseases’; if None, return all None
flat bool whether to return one flat dict, not one for each type False
verbose bool whether to print each custom module as it is loaded False

is_module

internals.modules.is_module(obj)

Check whether something is a Starsim module

module_map

internals.modules.module_map(key=None)

Map modules to standard types

This is the source of truth about module types. Modules with types that do not map to anything in the module_map are treated as custom modules.

module_types

internals.modules.module_types()

Return a list of known module types; based on module_map()

register_modules

internals.modules.register_modules(*args)

Register custom modules with Starsim so they can be referred to by string.

Note: “modules” here refers to Starsim modules. But this function registers Starsim “modules” (e.g. ss.SIR) that are found within Python “modules” (e.g. starsim).

Parameters

Name Type Description Default
args list the additional modules to register; can be either a module or a list of objects

Examples

# Standard use case, register modules automatically
import my_custom_disease_model as mcdm
ss.register_modules(mcdm)
ss.Sim(diseases='mydisease', networks='random').run() # This will work if mcdm.MyDisease() is defined

# Manual usage
my_modules = [mcdm.MyDisease, mcdm.MyNetwork]
ss.register_modules(my_modules)
ss.Sim(diseases='mydisease', networks='mynetwork').run()

required

internals.modules.required()

Decorator to mark module methods as required.

A common gotcha in Starsim is to forget to call super(), or to mistype a method name so it’s never called. This decorator lets you mark methods (of Modules only) to be sure that they are called either on sim initialization or on sim run.

Examples

import starsim.library as ssl

class CustomSIS(ssl.AgentSIS):

    def step(self):
        super().step()
        self.custom_step() # Will raise an exception if this line is not here
        return

    @ss.required() # Mark this method as required on run
    def custom_step(self):
        pass