Array indexing

🧠34% 🤖66%

When a model runs as agents, each person’s properties (age, sex, the state of each disease) are kept in arrays with one value per agent. Code that reads or changes them, in a module written as code or in an analysis after a run, needs to pick out agents correctly, including when some of them have died. This page covers the arrays, how agents are identified, and the rules for indexing them.

A declared model needs none of this: conditions such as who=ss.person.age < 5 or who=measles.I pick out people for you, and work under every solver. In a compartmental run there are no agents, and sim.people has one row for each group instead.

Arrays and UIDs

A state such as alive, female, or a disease’s I is an ss.BoolState, an array of true or false values; a number, such as age, is an ss.FloatArr (or ss.IntArr for whole numbers); the kinds are listed in the appendix. These arrays work like NumPy arrays, with sums, means, comparisons, and arithmetic, but are made for three things nearly every agent-based model needs:

1. Growth: as people are born or arrive, the arrays grow, in chunks rather than one agent at a time, so that adding agents is cheap.
2. Indexing: agents who die or leave stay in the arrays, so that what happened to them can still be looked up, but are left out of counts, sums, and means.
3. Random values: a state’s default can be a distribution (e.g. ss.BoolState('urban', default=ss.bernoulli(p=0.3))), so each new agent draws its own value, with common random numbers so that two runs give each agent the same draw.

Every agent has a UID (a unique identifier): a whole number, given in order from 0 as agents are made. An agent’s UID never changes, and isn’t reused after the agent dies, so a UID can be used to look up any agent, alive or dead.

The active UIDs (auids) are those of the agents who are alive and in the sim now. Each array holds two views of its data:

• raw: a value for every agent ever made, indexed by UID (once agents are born, it also has room at the end for agents not yet made, holding NaN)
• values: a value for each active agent, in the order of auids

To see this, here is a small measles outbreak in 10 people, in which infectious people die at a rate of 5% a day (a much higher death rate than measles has, so that some agents die in a short run):

import starsim as ss

S, I, R = ss.states('S I R')
infection = S >> I @ 1.9*I                  # β = 1.9 per day
recovery  = I >> R @ ss.days(8)             # Infectious for 8 days on average
death     = ss.death(I, rate=ss.perday(0.05)) # Deaths of infectious people
measles = ss.Disease(infection, recovery, death, init_prev=0.3)

sim = ss.Sim(measles, n_agents=10, stop=ss.days(50), rand_seed=3, verbose=0)
sim.run()

ppl = sim.people
print('Agents alive:', len(ppl))
print('Active UIDs:', ppl.auids)
print('Every UID:', ppl.uid.raw)
print('Alive:', ppl.alive.raw)
print('Ages of active agents:', ppl.age)
print('Ages of every agent:', ppl.age.raw)

Agents alive: 7
Active UIDs: [1 3 5 6 7 8 9]
Every UID: [0 1 2 3 4 5 6 7 8 9]
Alive: [False  True False  True False  True  True  True  True  True]
Ages of active agents: <FloatArr "age", len=7, [50.05353   65.253044  52.07521    1.9449059 25.157228  70.84109   39.55616  ]>
Ages of every agent: [10.827496  50.05353   32.99445   65.253044  12.883546  52.07521
  1.9449059 25.157228  70.84109   39.55616  ]

The dead agents are still in raw, but not in auids or in an array’s values.

Rules for indexing

Which agents an operation uses

The rules are:

• Sums, means, and other statistics (.sum(), .mean(), .std(), and so on) use the active agents only.
• Indexing depends on the index:
  • a whole number, or ss.uids(), is a UID, so it can pick any agent, alive or dead (as raw does);
  • [:], or a boolean state such as measles.I, picks among the active agents (as values does);
  • any other slice (e.g. [1:3]) raises an error, since it could mean either; write arr.values[1:3] or arr.raw[1:3].

age = ppl.age
print('Mean age of active agents:', age.mean())
print('The same:', age.values.mean())
print('Ages of the dead:', age[ppl.dead.uids])

Mean age of active agents: 43.55445
The same: 43.55445
Ages of the dead: [10.827496 32.99445  12.883546]

ppl.dead.uids finds the dead, since they aren’t among the active agents.

Indexing that works

print('Age of the agent with UID 0:', ppl.age[0]) # A whole number is a UID

some = ss.uids([0, 1, 2]) # Several UIDs
print('Their ages:', ppl.age[some])

women = ppl.female.uids # The UIDs for which a boolean state is true
print('Ages of women:', ppl.age[women])

m = sim.diseases.measles # The sim's copy of the disease, which ran
print('Recovered:', m.R.true()) # The UIDs of active agents in R
print('Not recovered:', m.R.false())

Age of the agent with UID 0: 10.827496
Their ages: [10.827496 50.05353  32.99445 ]
Ages of women: [50.05353   65.253044   1.9449059 25.157228  70.84109   39.55616  ]
Recovered: [1 3 5 6 7 8 9]
Not recovered: []

A boolean state, such as female or a disease’s I, has .uids (the UIDs for which it is true), and .true() and .false() (the active agents for which it is true or false). .uids is the same as .true(), so it leaves out the dead, and & ppl.alive isn’t needed; only ppl.dead looks beyond the active agents. Comparisons give boolean arrays with the same methods, e.g. (ppl.age < 15).uids. The sim’s own copy of a module is the one that ran: use sim.diseases.measles, not the measles given to the sim, whose arrays are empty.

Indexing that doesn’t

A list of whole numbers raises an error, since it could be UIDs or positions among the active agents:

import sciris as sc

with sc.tryexcept():
    ppl.age[[0, 1, 2]] # Use ss.uids([0, 1, 2]) instead

<class 'TypeError'> Indexing an Arr (age) by ([0, 1, 2]) is ambiguous or not supported. Use ss.uids() instead, or index Arr.raw or Arr.values.

raw includes the dead, so statistics of it differ from those of the active agents:

print('Mean age (active agents):', age.mean())
print('Mean age (raw, including the dead):', age.raw.mean())

Mean age (active agents): 43.55445
Mean age (raw, including the dead): 36.15867

Combining and changing selections

UIDs (ss.uids) are sorted arrays of whole numbers, with methods for combining them as sets: a.intersect(b) (in both), a.union(b) (in either), a.remove(b) (in a but not b), and a.concatenate(b) (both, without sorting). Comparisons and boolean states combine with & (and), | (or), and ~ (not) before being turned into UIDs:

children = (ppl.age < 15).uids
print('Children:', children)
print('Girls:', children.intersect(ppl.female.uids))
print('Boys:', children.remove(ppl.female.uids))
print('Girls, the same:', ((ppl.age < 15) & ppl.female).uids)
print('Children or women:', children.union(ppl.female.uids))

Children: [6]
Girls: [6]
Boys: []
Girls, the same: [6]
Children or women: [1 3 6 7 8 9]

Values are changed by UID in the same way as they are read: arr[uids] = value sets them, and arr[uids] += value adds to them. A single value is given to everyone selected, and an array gives each their own:

women = ppl.female.uids
print('Ages before:', ppl.age[women])
ppl.age[women] += 1 # A year older
print('Ages after:', ppl.age[women])

Ages before: [50.05353   65.253044   1.9449059 25.157228  70.84109   39.55616  ]
Ages after: [51.05353   66.253044   2.9449058 26.157228  71.84109   40.55616  ]

&, |, and ~ are only for true-or-false arrays; on an array of numbers they raise an error, since they would mean something else (bitwise operations). Compare first: (ppl.age > 15) & (ppl.age < 50), not ppl.age & .... To compare with several values, arr.isin([2, 3]) is shorter than (arr == 2) | (arr == 3).

Indexing in a module

The rules are the same inside a module written as code, where indexing is used most: at each step, a module picks out the agents it acts on, by UID, and reads or changes their values. Here an analyzer records the mean age of the people infectious with measles at each step. It finds them with the disease’s I.uids, reads their ages by UID, and checks for an empty selection, since nobody may be infectious:

import numpy as np

class AgeOfInfectious(ss.Analyzer):
    """ Record the mean age of the people infectious with measles at each step """
    def init_results(self):
        super().init_results()
        self.define_results(ss.Result('mean_age', label='Mean age of the infectious', kind='ratio'))
        return

    def record_states(self): # At the start of each step, as the other stocks are
        super().record_states()
        infectious = self.sim.diseases.measles.I.uids
        if len(infectious):
            self.results.mean_age[self.ti] = self.sim.people.age[infectious].mean()
        else:
            self.results.mean_age[self.ti] = np.nan
        return

S, I, R = ss.states('S I R')
measles = ss.Disease(S >> I @ 0.4*I, I >> R @ ss.days(8), init_prev=0.01, name='measles')
sim = ss.Sim(measles, analyzers=AgeOfInfectious(), n_agents=5000, stop=ss.days(100), verbose=0)
sim.run()
print('Mean age of everyone:', sim.people.age.mean().round(1))
print('Mean age of the infectious, every 10 days:', sim.results.ageofinfectious.mean_age[::10].round(1))

Mean age of everyone: 33.9
Mean age of the infectious, every 10 days: [34.4 34.3 33.2 34.7 34.3 34.8 31.4 31.3 44.9 56.8 28.8]

Analyzers and Custom modules have more examples of modules that index agents.

Arrays with several columns

Sometimes each person needs several values of the same kind, such as their immunity to each of several strains (strains keep immunity like this). An array with columns= has one row per agent and one column per value, and grows as agents are born like any other; in a module, it is defined like any state, e.g. self.define_states(ss.FloatArr('sus', columns=['wild', 'mutant'], default=1.0)). The first index picks agents, by the rules above, and the second picks columns, by position or by name:

imm = ss.FloatArr('imm', columns=['wild', 'mutant'], default=0, mock=5) # A stand-alone array of 5 agents
uids = ss.uids([0, 1])
imm[uids, 'mutant'] = 0.8 # The same as imm[uids, 1] = 0.8
print(imm[uids])          # A row for each agent
print(imm[:, 'mutant'])   # One column, for every active agent
print((imm.mutant > 0.5).uids) # UIDs from one column

[[0.  0.8]
 [0.  0.8]]
[0.8 0.8 0.  0.  0. ]
[0 1]

A column can be read by name as an attribute (imm.mutant), as a key (imm['mutant']), or with imm.col('mutant'); each gives an ordinary array with one value per agent, and writing to it (imm.mutant = 0, or imm.mutant *= 0.9 in a module’s step()) changes the state. Since such an array has several values per agent, it can’t be turned into UIDs directly; use one column, as above. Also:

• columns=3 gives three unnamed columns, read by position only, and a tuple such as columns=(2, 3) gives more dimensions.
• Column names must be unique, and can’t be the name of something an array already has (such as raw or mean).
• imm[:, [0, 1]] gives columns 0 and 1 for every active agent.
• people.to_df() gives one column of the table per value, e.g. immunity.sus_wild and immunity.sus_mutant (or _0, _1, and so on for unnamed columns).

Good practice

The rules above, in short:

• Index by UIDs (ss.uids([...]), a state’s .uids, or .true()), never by a plain list of numbers or by np.where(mask)[0], which gives positions among the active agents, not UIDs, once anyone has died.
• A boolean state is not a list of agents: measles.I[:] is true and false values; measles.I.uids is the agents.
• Check for no agents: a selection can be empty, e.g. when nobody is infectious.
• Find who just changed state from the transition (measles.infection.uids, measles.infection.since), rather than keeping the previous set by hand.
• Read and write by UID, not through raw: statistics of the array leave out the dead, and a write by UID to a declared disease’s state is tracked, so its counts see it.
• Work on many agents at once, with arrays of UIDs and NumPy, rather than a Python loop over agents.

Further reading

• NumPy: indexing on arrays: the indexing that Starsim’s arrays build on
• LASER: documentation for laser-generic: another framework that keeps agents as arrays, one value per person

Appendix: the kinds of array

The kinds of array that hold what agents are and have (see Arrays), for the states of ss.People(extra_states=...) and of modules:

Array	Holds	Default
ss.BoolState	True or false, for a state people are in that is worth counting (such as alive, female, or a disease’s I): a module records the number in each of its states as a result	False
ss.BoolArr	True or false, for a flag used inside a module (such as preterm), with no result of its own; otherwise the same, with .uids for the agents for whom it is true	False
ss.FloatArr	A number, such as an age or a relative risk	NaN
ss.IntArr	A whole number, such as a count	-999999 (for none, since whole numbers have no NaN)
ss.TimeArr	The step on which something happened (such as ti_dead)	NaN
ss.arrays.IndexArr	The UID of another agent (such as parent), made by Starsim rather than by models	-1
ss.uids	A set of UIDs, as from (people.age < 5).uids, for indexing any array

Each takes default=, a value or a distribution to draw each agent’s value from, and label=, the name used in plots; and columns= (above).