Profiling and debugging

Profiling

One of the main reasons people don’t use ABMs is because they can be very slow. While “vanilla Starsim” is quite fast (10,000 agents running for 100 timesteps should take about a second), custom modules, if not properly written, can be quite slow.

The first step of fixing a slow module is to identify the problem. To do this, Starsim includes some built-in profiling tools.

Let’s look at a simple simulation:

import sciris as sc
import starsim as ss
ss.options(jupyter=True)
sc.options(jupyter=True)

pars = dict(
    start = '2000-01-01',
    stop = '2020-01-01',
    diseases = 'sis',
    networks = 'random'
)

# Profile sim
sim = ss.Sim(pars)
prof = sim.profile()

Initializing sim with 10000 agents

Profiling 15 function(s):

 <bound method Sim.run of Sim(n=10000; 2000.01.01—2020.01.01; networks=randomnet; diseases=sis)>

<bound method Sim.start_step of Sim(n=10000; 2000.01.01—2020.01.01; networks=randomnet; diseases=sis [...]

<function Module.start_step at 0x73bc443874c0>

<function Module.start_step at 0x73bc443874c0>

<bound method SIS.step_state of sis(pars=[init_prev, beta, dur_inf, waning, imm_boost]; states=[susc [...]

<bound method DynamicNetwork.step of randomnet(n_edges=50000; pars=[n_contacts, dur, beta]; states=[ [...]

<bound method Infection.step of sis(pars=[init_prev, beta, dur_inf, waning, imm_boost]; states=[susc [...]

<bound method People.step_die of People(n=10000; age=30.2±17.4)>

<bound method People.update_results of People(n=10000; age=30.2±17.4)>

<bound method Network.update_results of randomnet(n_edges=50000; pars=[n_contacts, dur, beta]; state [...]

<function SIS.update_results at 0x73bca11bb4c0>

<function Module.finish_step at 0x73bc443876a0>

<function Module.finish_step at 0x73bc443876a0>

<bound method People.finish_step of People(n=10000; age=30.2±17.4)>

<bound method Sim.finish_step of Sim(n=10000; 2000.01.01—2020.01.01; networks=randomnet; diseases=si [...] 



  Running 2000.01.01 ( 0/21) (0.00 s)  ———————————————————— 5%

  Running 2010.01.01 (10/21) (0.07 s)  ••••••••••—————————— 52%

  Running 2020.01.01 (20/21) (0.11 s)  •••••••••••••••••••• 100%



Elapsed time: 0.113 s





—————————————————————————————————————————————————————————————————————

Profile of networks.Network.update_results: 0.000318099 s (0.281622%)

—————————————————————————————————————————————————————————————————————



Total time: 0.000318099 s

File: /home/cliffk/idm/starsim/starsim/networks.py

Function: Network.update_results at line 278



Line #      Hits         Time  Per Hit   % Time  Line Contents

==============================================================

   278                                               def update_results(self):

   279                                                   """ Store the number of edges in the network """

   280        21      96535.0   4596.9     30.3          super().update_results()

   281        21     215397.0  10257.0     67.7          self.results['n_edges'][self.ti] = len(self)

   282        21       6167.0    293.7      1.9          return







———————————————————————————————————————————————————————————————

Profile of people.People.finish_step: 0.000660023 s (0.584337%)

———————————————————————————————————————————————————————————————



Total time: 0.000660023 s

File: /home/cliffk/idm/starsim/starsim/people.py

Function: People.finish_step at line 524



Line #      Hits         Time  Per Hit   % Time  Line Contents

==============================================================

   524                                               def finish_step(self):

   525                                                   """ Remove dead agents and run post-step updates """

   526        21     515477.0  24546.5     78.1          self.remove_dead()

   527        21     139757.0   6655.1     21.2          self.update_post()

   528        21       4789.0    228.0      0.7          return







——————————————————————————————————————————————————————————

Profile of people.People.step_die: 0.00140846 s (1.24695%)

——————————————————————————————————————————————————————————



Total time: 0.00140846 s

File: /home/cliffk/idm/starsim/starsim/people.py

Function: People.step_die at line 484



Line #      Hits         Time  Per Hit   % Time  Line Contents

==============================================================

   484                                               def step_die(self):

   485                                                   """ Carry out any deaths or removals that took place this timestep """

   486        21    1190674.0  56698.8     84.5          death_uids = ((self.ti_dead <= self.sim.ti) | (self.ti_removed <= self.sim.ti)).uids

   487        21      98337.0   4682.7      7.0          self.alive[death_uids] = False  # Whilst not dead, removed agents should not be included in alive totals

   488                                           

   489                                                   # Execute deaths that took place this timestep (i.e., changing the `alive` state of the agents). This is executed

   490                                                   # before analyzers have run so that analyzers are able to inspect and record outcomes for agents that died this timestep

   491        42      85162.0   2027.7      6.0          for disease in self.sim.diseases():

   492        21      10986.0    523.1      0.8              if isinstance(disease, ss.Disease):

   493        21      17424.0    829.7      1.2                  disease.step_die(death_uids)

   494                                           

   495        21       5874.0    279.7      0.4          return death_uids







———————————————————————————————————————————————————————————————

Profile of diseases.SIS.update_results: 0.00160916 s (1.42463%)

———————————————————————————————————————————————————————————————



Total time: 0.00160916 s

File: /home/cliffk/idm/starsim/starsim/diseases.py

Function: SIS.update_results at line 907



Line #      Hits         Time  Per Hit   % Time  Line Contents

==============================================================

   907                                               @ss.required()

   908                                               def update_results(self):

   909                                                   """ Store the population immunity (susceptibility) """

   910        21    1048050.0  49907.1     65.1          super().update_results()

   911        21     555623.0  26458.2     34.5          self.results['rel_sus'][self.ti] = self.rel_sus.mean()

   912        21       5488.0    261.3      0.3          return







————————————————————————————————————————————————————————————————

Profile of people.People.update_results: 0.00212609 s (1.88229%)

————————————————————————————————————————————————————————————————



Total time: 0.00212609 s

File: /home/cliffk/idm/starsim/starsim/people.py

Function: People.update_results at line 513



Line #      Hits         Time  Per Hit   % Time  Line Contents

==============================================================

   513                                               def update_results(self):

   514                                                   """ Record per-timestep population counts into the simulation results """

   515        21      17019.0    810.4      0.8          ti = self.sim.ti

   516        21       5476.0    260.8      0.3          res = self.sim.results

   517        63     270943.0   4300.7     12.7          for state in self.auto_state_list: # Count each auto-generated BoolState result, e.g. n_alive, n_female

   518        42     513556.0  12227.5     24.2              res[f'n_{state.name}'][ti] = np.count_nonzero(getattr(self, state.name))

   519        21     509041.0  24240.0     23.9          res.new_deaths[ti] = np.count_nonzero(self.ti_dead == ti)

   520        21     429553.0  20454.9     20.2          res.new_emigrants[ti] = np.count_nonzero(self.ti_removed == ti)

   521        21     374395.0  17828.3     17.6          res.cum_deaths[ti] = res.new_deaths[ti] if ti == 0 else res.cum_deaths[ti-1] + res.new_deaths[ti]

   522        21       6108.0    290.9      0.3          return







—————————————————————————————————————————————————————

Profile of sim.Sim.start_step: 0.00326162 s (2.8876%)

—————————————————————————————————————————————————————



Total time: 0.00326162 s

File: /home/cliffk/idm/starsim/starsim/sim.py

Function: Sim.start_step at line 451



Line #      Hits         Time  Per Hit   % Time  Line Contents

==============================================================

   451                                               def start_step(self):

   452                                                   """ Start the step -- only print progress; all actual changes happen in the modules """

   453                                           

   454                                                   # Set the time and if we have reached the end of the simulation, then do nothing

   455        21       6679.0    318.0      0.2          if self.complete:

   456                                                       errormsg = 'Simulation already complete (call sim.init() to re-run)'

   457                                                       raise AlreadyRunError(errormsg)

   458                                           

   459                                                   # Print out progress if needed

   460        21    2322508.0 110595.6     71.2          self.elapsed = self.timer.toc(output=True)

   461        21       6419.0    305.7      0.2          if self.verbose: # Print progress

   462        21       5038.0    239.9      0.2              t = self.t

   463        21     134492.0   6404.4      4.1              simlabel = f'"{self.label}": ' if self.label else ''

   464        21     508197.0  24199.9     15.6              string = f'  Running {simlabel}{t.now("str")} ({t.ti:2.0f}/{t.npts}) ({self.elapsed:0.2f} s) '

   465        21       7307.0    348.0      0.2              if self.verbose >= 1:

   466                                                           sc.heading(string)

   467        21       5984.0    285.0      0.2              elif self.verbose > 0:

   468        21      12339.0    587.6      0.4                  if not (t.ti % int(1.0 / self.verbose)):

   469         3     246683.0  82227.7      7.6                      sc.progressbar(t.ti + 1, t.npts, label=string, length=20, newline=True)

   470        21       5973.0    284.4      0.2          return







—————————————————————————————————————————————————————————————

Profile of modules.Module.start_step: 0.00634658 s (5.61881%)

—————————————————————————————————————————————————————————————



Total time: 0.00634658 s

File: /home/cliffk/idm/starsim/starsim/modules.py

Function: Module.start_step at line 828



Line #      Hits         Time  Per Hit   % Time  Line Contents

==============================================================

   828                                               @required()

   829                                               def start_step(self):

   830                                                   """ Tasks to perform at the beginning of the step """

   831        42      14964.0    356.3      0.2          if self.finalized:

   832                                                       errormsg = f'The module {self._debug_name} has already been run. Did you mean to copy it before running it?'

   833                                                       raise RuntimeError(errormsg)

   834        42     212666.0   5063.5      3.4          if self.dists is not None and ss.options.crn: # Will be None if no distributions are defined; jumping is a no-op under crn=False (see Dist.jump), so skip the per-dist loop entirely

   835        42    6106279.0 145387.6     96.2              self.dists.jump_dt() # Advance random number generators forward for calls on this step

   836        42      12675.0    301.8      0.2          return







———————————————————————————————————————————————————————————

Profile of diseases.SIS.step_state: 0.00930471 s (8.23772%)

———————————————————————————————————————————————————————————



Total time: 0.00930471 s

File: /home/cliffk/idm/starsim/starsim/diseases.py

Function: SIS.step_state at line 868



Line #      Hits         Time  Per Hit   % Time  Line Contents

==============================================================

   868                                               def step_state(self):

   869                                                   """ Progress infectious -> recovered """

   870        21     960030.0  45715.7     10.3          recovered = (self.infected & (self.ti_recovered <= self.ti)).uids

   871        21      62737.0   2987.5      0.7          self.infected[recovered] = False

   872        21      35996.0   1714.1      0.4          self.susceptible[recovered] = True

   873        21    8239939.0 392378.0     88.6          self.update_immunity()

   874        21       6010.0    286.2      0.1          return







———————————————————————————————————————————————————————————————

Profile of networks.DynamicNetwork.step: 0.0239477 s (21.2016%)

———————————————————————————————————————————————————————————————



Total time: 0.0239477 s

File: /home/cliffk/idm/starsim/starsim/networks.py

Function: DynamicNetwork.step at line 437



Line #      Hits         Time  Per Hit   % Time  Line Contents

==============================================================

   437                                               def step(self):

   438                                                   """ Remove expired partnerships and add new ones """

   439        21    9021919.0 429615.2     37.7          self.end_pairs()

   440        21   14919741.0 710463.9     62.3          self.add_pairs()

   441        21       6085.0    289.8      0.0          return







—————————————————————————————————————————————————————————

Profile of diseases.Infection.step: 0.0618573 s (54.764%)

—————————————————————————————————————————————————————————



Total time: 0.0618573 s

File: /home/cliffk/idm/starsim/starsim/diseases.py

Function: Infection.step at line 214



Line #      Hits         Time  Per Hit   % Time  Line Contents

==============================================================

   214                                               def step(self):

   215                                                   """

   216                                                   Perform key infection updates, including infection and setting prognoses

   217                                                   """

   218                                                   # Create new cases

   219        21   47831003.0 2.28e+06     77.3          new_cases, sources, networks = self.infect() # TODO: store outputs in self or use objdict rather than 3 returns

   220                                           

   221                                                   # Set prognoses

   222        21       8084.0    385.0      0.0          if len(new_cases):

   223        21   14011120.0 667196.2     22.7              self.set_outcomes(new_cases, sources)

   224                                           

   225        21       7075.0    336.9      0.0          return new_cases, sources, networks



Figure(672x480)

/home/cliffk/idm/starsim/starsim/loop.py:606: RuntimeWarning:


No CPU timing was recorded: run with sim.run(profile=True) (or loop.run(profile=True)) to populate cpu_time.

This graph (which is a shortcut to sim.loop.plot_cpu()) shows us how much time each step in the integration loop takes. We can get line-by-line detail of where each function is taking time, though:

prof.disp(maxentries=5)

—————————————————————————————————————————————————————

Profile of sim.Sim.start_step: 0.00326162 s (2.8876%)

—————————————————————————————————————————————————————



Total time: 0.00326162 s

File: /home/cliffk/idm/starsim/starsim/sim.py

Function: Sim.start_step at line 451



Line #      Hits         Time  Per Hit   % Time  Line Contents

==============================================================

   451                                               def start_step(self):

   452                                                   """ Start the step -- only print progress; all actual changes happen in the modules """

   453                                           

   454                                                   # Set the time and if we have reached the end of the simulation, then do nothing

   455        21       6679.0    318.0      0.2          if self.complete:

   456                                                       errormsg = 'Simulation already complete (call sim.init() to re-run)'

   457                                                       raise AlreadyRunError(errormsg)

   458                                           

   459                                                   # Print out progress if needed

   460        21    2322508.0 110595.6     71.2          self.elapsed = self.timer.toc(output=True)

   461        21       6419.0    305.7      0.2          if self.verbose: # Print progress

   462        21       5038.0    239.9      0.2              t = self.t

   463        21     134492.0   6404.4      4.1              simlabel = f'"{self.label}": ' if self.label else ''

   464        21     508197.0  24199.9     15.6              string = f'  Running {simlabel}{t.now("str")} ({t.ti:2.0f}/{t.npts}) ({self.elapsed:0.2f} s) '

   465        21       7307.0    348.0      0.2              if self.verbose >= 1:

   466                                                           sc.heading(string)

   467        21       5984.0    285.0      0.2              elif self.verbose > 0:

   468        21      12339.0    587.6      0.4                  if not (t.ti % int(1.0 / self.verbose)):

   469         3     246683.0  82227.7      7.6                      sc.progressbar(t.ti + 1, t.npts, label=string, length=20, newline=True)

   470        21       5973.0    284.4      0.2          return







—————————————————————————————————————————————————————————————

Profile of modules.Module.start_step: 0.00634658 s (5.61881%)

—————————————————————————————————————————————————————————————



Total time: 0.00634658 s

File: /home/cliffk/idm/starsim/starsim/modules.py

Function: Module.start_step at line 828



Line #      Hits         Time  Per Hit   % Time  Line Contents

==============================================================

   828                                               @required()

   829                                               def start_step(self):

   830                                                   """ Tasks to perform at the beginning of the step """

   831        42      14964.0    356.3      0.2          if self.finalized:

   832                                                       errormsg = f'The module {self._debug_name} has already been run. Did you mean to copy it before running it?'

   833                                                       raise RuntimeError(errormsg)

   834        42     212666.0   5063.5      3.4          if self.dists is not None and ss.options.crn: # Will be None if no distributions are defined; jumping is a no-op under crn=False (see Dist.jump), so skip the per-dist loop entirely

   835        42    6106279.0 145387.6     96.2              self.dists.jump_dt() # Advance random number generators forward for calls on this step

   836        42      12675.0    301.8      0.2          return







———————————————————————————————————————————————————————————

Profile of diseases.SIS.step_state: 0.00930471 s (8.23772%)

———————————————————————————————————————————————————————————



Total time: 0.00930471 s

File: /home/cliffk/idm/starsim/starsim/diseases.py

Function: SIS.step_state at line 868



Line #      Hits         Time  Per Hit   % Time  Line Contents

==============================================================

   868                                               def step_state(self):

   869                                                   """ Progress infectious -> recovered """

   870        21     960030.0  45715.7     10.3          recovered = (self.infected & (self.ti_recovered <= self.ti)).uids

   871        21      62737.0   2987.5      0.7          self.infected[recovered] = False

   872        21      35996.0   1714.1      0.4          self.susceptible[recovered] = True

   873        21    8239939.0 392378.0     88.6          self.update_immunity()

   874        21       6010.0    286.2      0.1          return







———————————————————————————————————————————————————————————————

Profile of networks.DynamicNetwork.step: 0.0239477 s (21.2016%)

———————————————————————————————————————————————————————————————



Total time: 0.0239477 s

File: /home/cliffk/idm/starsim/starsim/networks.py

Function: DynamicNetwork.step at line 437



Line #      Hits         Time  Per Hit   % Time  Line Contents

==============================================================

   437                                               def step(self):

   438                                                   """ Remove expired partnerships and add new ones """

   439        21    9021919.0 429615.2     37.7          self.end_pairs()

   440        21   14919741.0 710463.9     62.3          self.add_pairs()

   441        21       6085.0    289.8      0.0          return







—————————————————————————————————————————————————————————

Profile of diseases.Infection.step: 0.0618573 s (54.764%)

—————————————————————————————————————————————————————————



Total time: 0.0618573 s

File: /home/cliffk/idm/starsim/starsim/diseases.py

Function: Infection.step at line 214



Line #      Hits         Time  Per Hit   % Time  Line Contents

==============================================================

   214                                               def step(self):

   215                                                   """

   216                                                   Perform key infection updates, including infection and setting prognoses

   217                                                   """

   218                                                   # Create new cases

   219        21   47831003.0 2.28e+06     77.3          new_cases, sources, networks = self.infect() # TODO: store outputs in self or use objdict rather than 3 returns

   220                                           

   221                                                   # Set prognoses

   222        21       8084.0    385.0      0.0          if len(new_cases):

   223        21   14011120.0 667196.2     22.7              self.set_outcomes(new_cases, sources)

   224                                           

   225        21       7075.0    336.9      0.0          return new_cases, sources, networks

(Note that the names of the functions here refer to the actual functions called, which may not match the graph above. That’s because, for example, ss.SIS does not define its own step() method, but instead inherits step() from Infection. In the graph, this is shown as sis.step(), but is listed in the table as Infection.step(). This is because it’s referring to the actual code being run, so refers to where those lines of code exist in the codebase; there is no code corresponding to SIS.step() since it’s just inherited from Infection.step().)

If you want more detail, you can also define custom functions to follow. For example, we can see that ss.SIS.infect() takes the most time in ss.SIS.step(), so let’s profile that:

prof = sim.profile(follow=ss.SIS.infect, plot=False)
prof.disp()

Initializing sim with 10000 agents

Profiling 1 function(s):

 <function Infection.infect at 0x73bca11b98a0> 



  Running 2000.01.01 ( 0/21) (0.00 s)  ———————————————————— 5%

  Running 2010.01.01 (10/21) (0.04 s)  ••••••••••—————————— 52%

  Running 2020.01.01 (20/21) (0.08 s)  •••••••••••••••••••• 100%



Elapsed time: 88.3 ms





————————————————————————————————————————————————————————————

Profile of diseases.Infection.infect: 0.0349516 s (39.5574%)

————————————————————————————————————————————————————————————



Total time: 0.0349516 s

File: /home/cliffk/idm/starsim/starsim/diseases.py

Function: Infection.infect at line 255



Line #      Hits         Time  Per Hit   % Time  Line Contents

==============================================================

   255                                               def infect(self):

   256                                                   """ Determine who gets infected on this timestep via transmission on the network """

   257        21      11465.0    546.0      0.0          new_cases = []

   258        21       5290.0    251.9      0.0          sources = []

   259        21       4411.0    210.0      0.0          networks = []

   260        21     341341.0  16254.3      1.0          betamap = self.validate_beta()

   261                                           

   262                                                   # Compute effective transmissibility and susceptibility directly on the raw

   263                                                   # (full-length) arrays. This avoids the gather/scatter, full-length astype copy,

   264                                                   # and extra wrapper allocations of the Arr math operators; edges only ever index

   265                                                   # living agents, so stale raw values for inactive agents are never used.

   266        21     314990.0  14999.5      0.9          rel_trans = self.rel_trans.asnew(self.infectious.raw * self.rel_trans.raw, copy=False)

   267        21     173785.0   8275.5      0.5          rel_sus   = self.rel_sus.asnew(self.susceptible.raw * self.rel_sus.raw, copy=False)

   268                                           

   269        42      79997.0   1904.7      0.2          for i, (nkey,route) in enumerate(self.sim.networks.items()):

   270        21      18078.0    860.9      0.1              nk = ss.standardize_netkey(nkey)

   271                                           

   272                                                       # Main use case: networks

   273        21       9397.0    447.5      0.0              if isinstance(route, ss.Network):

   274        21     127366.0   6065.0      0.4                  if len(route): # Skip networks with no edges

   275        21       5824.0    277.3      0.0                      edges = route.edges

   276        21     188440.0   8973.3      0.5                      p1_to_p2 = [edges.p1, edges.p2, betamap[nk][0]]  # p1→p2 direction, beta 0

   277        21     178856.0   8517.0      0.5                      p2_to_p1 = [edges.p2, edges.p1, betamap[nk][1]]  # p2→p1 direction, beta 1

   278        63      28584.0    453.7      0.1                      for src, trg, beta in [p1_to_p2, p2_to_p1]:

   279        42      40517.0    964.7      0.1                          if beta: # Skip networks with no transmission

   280        42     658484.0  15678.2      1.9                              disease_beta = beta.to_prob(self.t.dt) if isinstance(beta, ss.Rate) else beta

   281        42    1437431.0  34224.5      4.1                              beta_per_dt = route.net_beta(disease_beta=disease_beta, disease=self) # Compute beta for this network and timestep

   282        42   25601420.0 609557.6     73.2                              randvals = self.trans_rng.rvs(src, trg) # Generate a new random number based on the two other random numbers

   283        42      24475.0    582.7      0.1                              args = (src, trg, rel_trans, rel_sus, beta_per_dt, randvals) # Set up the arguments to calculate transmission

   284        42    3402477.0  81011.4      9.7                              target_uids, source_uids = self.compute_transmission(*args) # Actually calculate it

   285        42      14000.0    333.3      0.0                              new_cases.append(target_uids)

   286        42       9649.0    229.7      0.0                              sources.append(source_uids)

   287        42     264792.0   6304.6      0.8                              networks.append(np.full(len(target_uids), dtype=ss_int, fill_value=i))

   288                                           

   289                                                       # Handle everything else: mixing pools, environmental transmission, etc.

   290                                                       elif isinstance(route, ss.Route):

   291                                                           # Mixing pools are unidirectional, only use the first beta value

   292                                                           disease_beta = betamap[nk][0].to_prob(self.t.dt) if isinstance(betamap[nk][0], ss.Rate) else betamap[nk][0]

   293                                                           target_uids = route.compute_transmission(rel_sus, rel_trans, disease_beta, disease=self)

   294                                                           new_cases.append(target_uids)

   295                                                           sources.append(np.full(len(target_uids), dtype=ss_int, fill_value=ss.dtypes.int_nan))

   296                                                           networks.append(np.full(len(target_uids), dtype=ss_int, fill_value=i))

   297                                                       else:

   298                                                           errormsg = f'Cannot compute transmission via route {type(route)}; please subclass ss.Route and define a compute_transmission() method'

   299                                                           raise TypeError(errormsg)

   300                                           

   301                                                   # Finalize

   302        21      10435.0    496.9      0.0          if len(new_cases) and len(sources):

   303        21     230669.0  10984.2      0.7              new_cases = ss.uids.concatenate(new_cases)

   304        21    1511146.0  71959.3      4.3              new_cases, inds = new_cases.unique(return_index=True)

   305        21     184583.0   8789.7      0.5              sources = ss.uids.concatenate(sources)[inds]

   306        21      66489.0   3166.1      0.2              networks = np.concatenate(networks)[inds]

   307                                                   else:

   308                                                       new_cases = ss.uids()

   309                                                       sources = ss.uids()

   310                                                       networks = np.empty(0, dtype=ss_int)

   311                                           

   312        21       7230.0    344.3      0.0          return new_cases, sources, networks







————————————————————————————————————————————————————————————

Profile of diseases.Infection.infect: 0.0349516 s (39.5574%)

————————————————————————————————————————————————————————————



Total time: 0.0349516 s

File: /home/cliffk/idm/starsim/starsim/diseases.py

Function: Infection.infect at line 255



Line #      Hits         Time  Per Hit   % Time  Line Contents

==============================================================

   255                                               def infect(self):

   256                                                   """ Determine who gets infected on this timestep via transmission on the network """

   257        21      11465.0    546.0      0.0          new_cases = []

   258        21       5290.0    251.9      0.0          sources = []

   259        21       4411.0    210.0      0.0          networks = []

   260        21     341341.0  16254.3      1.0          betamap = self.validate_beta()

   261                                           

   262                                                   # Compute effective transmissibility and susceptibility directly on the raw

   263                                                   # (full-length) arrays. This avoids the gather/scatter, full-length astype copy,

   264                                                   # and extra wrapper allocations of the Arr math operators; edges only ever index

   265                                                   # living agents, so stale raw values for inactive agents are never used.

   266        21     314990.0  14999.5      0.9          rel_trans = self.rel_trans.asnew(self.infectious.raw * self.rel_trans.raw, copy=False)

   267        21     173785.0   8275.5      0.5          rel_sus   = self.rel_sus.asnew(self.susceptible.raw * self.rel_sus.raw, copy=False)

   268                                           

   269        42      79997.0   1904.7      0.2          for i, (nkey,route) in enumerate(self.sim.networks.items()):

   270        21      18078.0    860.9      0.1              nk = ss.standardize_netkey(nkey)

   271                                           

   272                                                       # Main use case: networks

   273        21       9397.0    447.5      0.0              if isinstance(route, ss.Network):

   274        21     127366.0   6065.0      0.4                  if len(route): # Skip networks with no edges

   275        21       5824.0    277.3      0.0                      edges = route.edges

   276        21     188440.0   8973.3      0.5                      p1_to_p2 = [edges.p1, edges.p2, betamap[nk][0]]  # p1→p2 direction, beta 0

   277        21     178856.0   8517.0      0.5                      p2_to_p1 = [edges.p2, edges.p1, betamap[nk][1]]  # p2→p1 direction, beta 1

   278        63      28584.0    453.7      0.1                      for src, trg, beta in [p1_to_p2, p2_to_p1]:

   279        42      40517.0    964.7      0.1                          if beta: # Skip networks with no transmission

   280        42     658484.0  15678.2      1.9                              disease_beta = beta.to_prob(self.t.dt) if isinstance(beta, ss.Rate) else beta

   281        42    1437431.0  34224.5      4.1                              beta_per_dt = route.net_beta(disease_beta=disease_beta, disease=self) # Compute beta for this network and timestep

   282        42   25601420.0 609557.6     73.2                              randvals = self.trans_rng.rvs(src, trg) # Generate a new random number based on the two other random numbers

   283        42      24475.0    582.7      0.1                              args = (src, trg, rel_trans, rel_sus, beta_per_dt, randvals) # Set up the arguments to calculate transmission

   284        42    3402477.0  81011.4      9.7                              target_uids, source_uids = self.compute_transmission(*args) # Actually calculate it

   285        42      14000.0    333.3      0.0                              new_cases.append(target_uids)

   286        42       9649.0    229.7      0.0                              sources.append(source_uids)

   287        42     264792.0   6304.6      0.8                              networks.append(np.full(len(target_uids), dtype=ss_int, fill_value=i))

   288                                           

   289                                                       # Handle everything else: mixing pools, environmental transmission, etc.

   290                                                       elif isinstance(route, ss.Route):

   291                                                           # Mixing pools are unidirectional, only use the first beta value

   292                                                           disease_beta = betamap[nk][0].to_prob(self.t.dt) if isinstance(betamap[nk][0], ss.Rate) else betamap[nk][0]

   293                                                           target_uids = route.compute_transmission(rel_sus, rel_trans, disease_beta, disease=self)

   294                                                           new_cases.append(target_uids)

   295                                                           sources.append(np.full(len(target_uids), dtype=ss_int, fill_value=ss.dtypes.int_nan))

   296                                                           networks.append(np.full(len(target_uids), dtype=ss_int, fill_value=i))

   297                                                       else:

   298                                                           errormsg = f'Cannot compute transmission via route {type(route)}; please subclass ss.Route and define a compute_transmission() method'

   299                                                           raise TypeError(errormsg)

   300                                           

   301                                                   # Finalize

   302        21      10435.0    496.9      0.0          if len(new_cases) and len(sources):

   303        21     230669.0  10984.2      0.7              new_cases = ss.uids.concatenate(new_cases)

   304        21    1511146.0  71959.3      4.3              new_cases, inds = new_cases.unique(return_index=True)

   305        21     184583.0   8789.7      0.5              sources = ss.uids.concatenate(sources)[inds]

   306        21      66489.0   3166.1      0.2              networks = np.concatenate(networks)[inds]

   307                                                   else:

   308                                                       new_cases = ss.uids()

   309                                                       sources = ss.uids()

   310                                                       networks = np.empty(0, dtype=ss_int)

   311                                           

   312        21       7230.0    344.3      0.0          return new_cases, sources, networks

(Note: you can only follow functions that are called as part of sim.run() this way. To follow other functions, such as those run by sim.init(), you can use sc.profile() directly.)

Debugging

When figuring out what your sim is doing – whether it’s doing something it shouldn’t be, or not doing something it should – sim.loop is your friend. It shows everything that will happen in the sim, and in what order:

import starsim as ss

sim = ss.Sim(
    start = 2000,
    stop = 2002,
    diseases = 'sis',
    networks = 'random',
    verbose = 0,
)
sim.run()
sim.loop.to_df().disp()
# %%

    time  ti  func_order                     label     module       func_name  cpu_time
0   2000   0           0            sim.start_step        sim      start_step       NaN
1   2000   0           1      randomnet.start_step  randomnet      start_step       NaN
2   2000   0           2            sis.start_step        sis      start_step       NaN
3   2000   0           3            sis.step_state        sis      step_state       NaN
4   2000   0           4            randomnet.step  randomnet            step       NaN
5   2000   0           5                  sis.step        sis            step       NaN
6   2000   0           6           people.step_die     people        step_die       NaN
7   2000   0           7     people.update_results     people  update_results       NaN
8   2000   0           8  randomnet.update_results  randomnet  update_results       NaN
9   2000   0           9        sis.update_results        sis  update_results       NaN
10  2000   0          10     randomnet.finish_step  randomnet     finish_step       NaN
11  2000   0          11           sis.finish_step        sis     finish_step       NaN
12  2000   0          12        people.finish_step     people     finish_step       NaN
13  2000   0          13           sim.finish_step        sim     finish_step       NaN
14  2001   1           0            sim.start_step        sim      start_step       NaN
15  2001   1           1      randomnet.start_step  randomnet      start_step       NaN
16  2001   1           2            sis.start_step        sis      start_step       NaN
17  2001   1           3            sis.step_state        sis      step_state       NaN
18  2001   1           4            randomnet.step  randomnet            step       NaN
19  2001   1           5                  sis.step        sis            step       NaN
20  2001   1           6           people.step_die     people        step_die       NaN
21  2001   1           7     people.update_results     people  update_results       NaN
22  2001   1           8  randomnet.update_results  randomnet  update_results       NaN
23  2001   1           9        sis.update_results        sis  update_results       NaN
24  2001   1          10     randomnet.finish_step  randomnet     finish_step       NaN
25  2001   1          11           sis.finish_step        sis     finish_step       NaN
26  2001   1          12        people.finish_step     people     finish_step       NaN
27  2001   1          13           sim.finish_step        sim     finish_step       NaN
28  2002   2           0            sim.start_step        sim      start_step       NaN
29  2002   2           1      randomnet.start_step  randomnet      start_step       NaN
30  2002   2           2            sis.start_step        sis      start_step       NaN
31  2002   2           3            sis.step_state        sis      step_state       NaN
32  2002   2           4            randomnet.step  randomnet            step       NaN
33  2002   2           5                  sis.step        sis            step       NaN
34  2002   2           6           people.step_die     people        step_die       NaN
35  2002   2           7     people.update_results     people  update_results       NaN
36  2002   2           8  randomnet.update_results  randomnet  update_results       NaN
37  2002   2           9        sis.update_results        sis  update_results       NaN
38  2002   2          10     randomnet.finish_step  randomnet     finish_step       NaN
39  2002   2          11           sis.finish_step        sis     finish_step       NaN
40  2002   2          12        people.finish_step     people     finish_step       NaN
41  2002   2          13           sim.finish_step        sim     finish_step       NaN

As you can see, it’s a lot – this is only three timesteps and two modules, and it’s already 41 steps.

The typical way to do debugging is to insert breakpoints or print statements into your modules for custom debugging (e.g., to print a value), or to use analyzers for heavier-lift debugging. Starsim also lets you manually modify the loop by inserting “probes” or other arbitrary functions. For example, if you wanted to check the population size after each time the People object is updated:

def check_pop_size(sim):
    print(f'Population size is {len(sim.people)}')

sim = ss.Sim(diseases='sir', networks='random', demographics=True, dur=10)
sim.init()
sim.loop.insert(check_pop_size, label='people.finish_step')
sim.run()

Initializing sim with 10000 agents
  Running 2000 ( 0/11) (0.00 s)  •——————————————————— 9%
Population size is 10074
Population size is 10191
Population size is 10281
Population size is 10376
Population size is 10454
Population size is 10556
Population size is 10652
Population size is 10742
Population size is 10853
Population size is 10946
  Running 2010 (10/11) (0.04 s)  •••••••••••••••••••• 100%

Population size is 11041

Sim(n=10000; 2000—2010.0; demographics=births, deaths; networks=randomnet; diseases=sir)

In this case, you get the same output as using an analyzer:

def check_pop_size(sim):
    print(f'Population size is {len(sim.people)}')

sim = ss.Sim(diseases='sir', networks='random', demographics=True, dur=10, analyzers=check_pop_size)
sim.run()

Initializing sim with 10000 agents
  Running 2000 ( 0/11) (0.00 s)  •——————————————————— 9%
Population size is 10179
Population size is 10287
Population size is 10379
Population size is 10485
Population size is 10559
Population size is 10665
Population size is 10750
Population size is 10857
Population size is 10951
Population size is 11066
  Running 2010 (10/11) (0.04 s)  •••••••••••••••••••• 100%

Population size is 11163

Sim(n=10000; 2000—2010.0; demographics=births, deaths; networks=randomnet; diseases=sir; analyzers=check_pop_size)

However, inserting functions directly in the loop gives you more control over their exact placement, whereas analyzers are always executed last in the timestep.

The loop also has methods for visualizing itself. You can get a simple representation of the loop with loop.plot():

sim.loop.plot()

Figure(672x480)

Or a slightly more detailed one with loop.plot_step_order():

sim.loop.plot_step_order()

Figure(672x480)

This is especially useful if your simulation has modules with different timesteps, e.g.:

sis = ss.SIS(dt=0.1)
net = ss.RandomNet(dt=0.5)
births = ss.Births(dt=1)
sim = ss.Sim(dt=0.1, dur=5, diseases=sis, networks=net, demographics=births)
sim.init()
sim.loop.plot_step_order()

Initializing sim with 10000 agents
Figure(672x480)

(Note: this is a 3D plot, so it helps if you can plot it in a separate window interactively to be able to move it around, rather than just in a notebook.)