Developer tutorial: ABM examples#

This tutorial describes how Starsim can be used to reproduce examples from other Python ABMs.

Mesa: Wealth model#

Although Starsim is intended primarily for modeling disease, it can also be used as a general-purpose agent-based model. This example illustrates a simple “weath model”, in which each agent starts with a single unit of wealth, and on each timestep, every agent with more than zero wealth gives one unit of wealth to another agent.

This tutorial is adapted from the following example:

https://mesa.readthedocs.io/en/stable/tutorials/intro_tutorial.html

Setting up the model#

We could define the wealth model as any type of module, since they all can store states and update them. Here we will define wealth as a subclass of ss.Intervention (though it could equally well be a subclass of ss.Demographics or even ss.Disease, if you are so inclined). All we need to do is update the wealth state (which we can store inside the “intervention”), and we can also use this class to track the wealth distribution over time and plot it. The full model looks like this:

[1]:
# Imports
import numpy as np
import starsim as ss
import matplotlib.pyplot as plt

# Define the model
class WealthModel(ss.Intervention):
    """ A simple wealth transfer model"""

    def init_post(self, bins=10):
        """ Define custom model attributes """
        super().init_post()
        self.npts = self.sim.npts # Number of timepoints
        self.n_agents = len(sim.people) # Number of agents
        self.wealth = np.ones(self.n_agents) # Initial wealth of each agent
        self.bins = np.arange(bins+1) # Bins used for plotting
        self.wealth_dist = np.zeros((self.npts, len(self.bins)-1)) # Wealth distribution over time
        return

    def apply(self, sim):
        """ Transfer wealth between agents -- core model logic """
        self.wealth_hist() # Store the wealth at this time point
        givers = self.wealth > 0 # People need wealth to be givers
        receivers = np.random.choice(sim.people.uid, size=givers.sum()) # Anyone can be a receiver
        self.wealth[givers] -= 1 # Givers are unique, so can use vectorized version
        for receive in receivers: # Vectorized version is: np.add.at(sim.people.wealth.raw, receivers, 1)
            self.wealth[receive] += 1
        return

    def wealth_hist(self):
        """ Calculate the wealth histogram """
        ti = self.sim.ti # Current timestep
        self.wealth_dist[ti,:], _ = np.histogram(self.wealth, bins=self.bins)
        return

    def plot(self):
        """ Plot a 2D histogram of the final wealth distribution """
        plt.figure()
        plt.bar(self.bins[:-1], self.wealth_dist[-1,:])
        plt.title('Wealth distribution at final time point')
        plt.xlabel('Wealth')
        plt.ylabel('Number of agents')
        plt.show()
        return

    def plot3d(self):
        """ Plot a 3D heatmap of the wealth distribution over time """
        plt.figure()
        plt.pcolor(self.wealth_dist.T, cmap='turbo')
        plt.title('Wealth distribution over time')
        plt.xlabel('Time')
        plt.ylabel('Wealth')
        plt.colorbar().set_label('Number of agents', rotation=270)
        plt.show()
        return

# Create sim inputs, including the wealth model
wealth = WealthModel()
pars = dict(
    n_agents = 100, # Number of agents
    start = 0,
    end = 100,
    interventions = wealth,
)

# Run the model
sim = ss.Sim(pars, copy_inputs=False) # copy_inputs=False lets us reuse the "wealth" object from above
sim.run()

# Plot the results
wealth.plot()
wealth.plot3d()
Starsim 1.0.0 (2024-07-10) — © 2023-2024 by IDM
Initializing sim with 100 agents
  Running 0.0 ( 0/101) (0.36 s)  ———————————————————— 1%
  Running 10.0 (10/101) (0.36 s)  ••—————————————————— 11%
  Running 20.0 (20/101) (0.37 s)  ••••———————————————— 21%
  Running 30.0 (30/101) (0.38 s)  ••••••—————————————— 31%
  Running 40.0 (40/101) (0.38 s)  ••••••••———————————— 41%
  Running 50.0 (50/101) (0.39 s)  ••••••••••—————————— 50%
  Running 60.0 (60/101) (0.39 s)  ••••••••••••———————— 60%
  Running 70.0 (70/101) (0.40 s)  ••••••••••••••—————— 70%
  Running 80.0 (80/101) (0.40 s)  ••••••••••••••••———— 80%
  Running 90.0 (90/101) (0.41 s)  ••••••••••••••••••—— 90%
  Running 100.0 (100/101) (0.41 s)  •••••••••••••••••••• 100%

../_images/tutorials_dev_tut_abm_1_1.png
../_images/tutorials_dev_tut_abm_1_2.png

Comparison with Mesa#

While the implementation in Starsim is similar to Mesa, there are a couple key differences:

  • Because Starsim’s people object is vectorized, the wealth definition and update is vectorized as well.

  • Both Mesa and Starsim versions of the model are quite simple, but there is a little less boilerplate in the Starsim version.