State arrays¶

This tutorial covers `StateArray``, a key data structure in laser-measles that provides convenient access to epidemiological state compartments.

Section 1: StateArray fundamentals¶

What is a StateArray?¶

StateArray is a numpy array wrapper that extends np.ndarray to provide attribute-based access to epidemiological state compartments. Instead of remembering that states[0] is Susceptible and states[1] is Infectious, you can use intuitive names like states.S and states.I.

In [ ]:

import numpy as np

from laser.measles.utils import StateArray

import numpy as np from laser.measles.utils import StateArray

Construction¶

StateArray is constructed with two parameters:

• input_array: A numpy array (typically 2D with shape (num_states, num_patches))
• state_names: A list of state compartment names

In [ ]:

# Example: Create a StateArray for a 3-patch SIR model
num_patches = 3
num_states = 3

# Create underlying numpy array
data = np.array(
    [
        [1000, 800, 1200],  # Susceptible population in each patch
        [10, 20, 5],  # Infectious population in each patch
        [0, 0, 0],  # Recovered population in each patch
    ]
)

# Wrap with StateArray
states = StateArray(data, state_names=["S", "I", "R"])
print("StateArray shape:", states.shape)
print("State names:", states._state_names)

# Example: Create a StateArray for a 3-patch SIR model num_patches = 3 num_states = 3 # Create underlying numpy array data = np.array( [ [1000, 800, 1200], # Susceptible population in each patch [10, 20, 5], # Infectious population in each patch [0, 0, 0], # Recovered population in each patch ] ) # Wrap with StateArray states = StateArray(data, state_names=["S", "I", "R"]) print("StateArray shape:", states.shape) print("State names:", states._state_names)

Data storage¶

StateArray uses standard numpy array storage with additional metadata:

• The underlying data is stored as a regular numpy array
• _state_names stores the list of state compartment names
• _state_indices provides a mapping from names to array indices

In [ ]:

print("Underlying data type:", type(states.view(np.ndarray)))
print("State indices mapping:", states._state_indices)

print("Underlying data type:", type(states.view(np.ndarray))) print("State indices mapping:", states._state_indices)

Access patterns¶

StateArray supports both traditional numeric indexing and intuitive attribute access:

In [ ]:

# Numeric access (backward compatible)
print("Susceptible (numeric):", states[0])
print("Infectious (numeric):", states[1])

# Attribute access (intuitive)
print("Susceptible (attribute):", states.S)
print("Infectious (attribute):", states.I)

# Both approaches access the same data
print("Same data?", np.array_equal(states[0], states.S))

# Numeric access (backward compatible) print("Susceptible (numeric):", states[0]) print("Infectious (numeric):", states[1]) # Attribute access (intuitive) print("Susceptible (attribute):", states.S) print("Infectious (attribute):", states.I) # Both approaches access the same data print("Same data?", np.array_equal(states[0], states.S))

Section 2: StateArray in practice¶

Usage in patches LaserFrame¶

In laser-measles models, StateArray is used as the states property of the patches LaserFrame. This provides a convenient interface for accessing and modifying epidemiological compartments across spatial patches.

In [ ]:

# Example showing how models initialize `StateArray`
# (This mimics what happens in actual model initialization)

# Simulate patch populations
patch_pops = np.array([1000, 800, 1200])
num_patches = len(patch_pops)

# Create states array for SEIR model
seir_states = np.zeros((4, num_patches))
seir_states[0] = patch_pops  # Initialize all as Susceptible

# Wrap with StateArray
patch_states = StateArray(seir_states, state_names=["S", "E", "I", "R"])

print("Initial populations:")
print(f"Susceptible: {patch_states.S}")
print(f"Exposed: {patch_states.E}")
print(f"Infectious: {patch_states.I}")
print(f"Recovered: {patch_states.R}")

# Example showing how models initialize `StateArray` # (This mimics what happens in actual model initialization) # Simulate patch populations patch_pops = np.array([1000, 800, 1200]) num_patches = len(patch_pops) # Create states array for SEIR model seir_states = np.zeros((4, num_patches)) seir_states[0] = patch_pops # Initialize all as Susceptible # Wrap with StateArray patch_states = StateArray(seir_states, state_names=["S", "E", "I", "R"]) print("Initial populations:") print(f"Susceptible: {patch_states.S}") print(f"Exposed: {patch_states.E}") print(f"Infectious: {patch_states.I}") print(f"Recovered: {patch_states.R}")

Practical examples¶

StateArray supports all numpy operations while maintaining readable code:

In [ ]:

# Example 1: Calculate prevalence
total_pop = patch_states.sum(axis=0)
prevalence = patch_states.I / total_pop
print("Prevalence per patch:", prevalence)

# Example 2: Simulate some infections
new_infections = np.array([5, 3, 8])
patch_states.S -= new_infections
patch_states.E += new_infections

print("After infections:")
print(f"Susceptible: {patch_states.S}")
print(f"Exposed: {patch_states.E}")

# Example 3: Slicing operations work as expected
print("Total infectious across all patches:", patch_states.I.sum())
print("States in first patch:", patch_states[:, 0])

# Example 1: Calculate prevalence total_pop = patch_states.sum(axis=0) prevalence = patch_states.I / total_pop print("Prevalence per patch:", prevalence) # Example 2: Simulate some infections new_infections = np.array([5, 3, 8]) patch_states.S -= new_infections patch_states.E += new_infections print("After infections:") print(f"Susceptible: {patch_states.S}") print(f"Exposed: {patch_states.E}") # Example 3: Slicing operations work as expected print("Total infectious across all patches:", patch_states.I.sum()) print("States in first patch:", patch_states[:, 0])

Benefits¶

StateArray provides several key advantages:

1. Readable code: states.S is more intuitive than states[0].
2. Maintainability: Adding/removing states doesn't break numeric indices.
3. Backward compatibility: Existing code using numeric indexing still works
4. Full NumPy support: All numpy operations work seamlessly.
5. Error prevention: Typos in state names raise AttributeError immediately.
6. Flexibility: Works with different model types (SIR, SEIR, etc.).

In [ ]:

# Example: Error handling for invalid state names
try:
    invalid_state = patch_states.X  # X is not a valid state
except AttributeError as e:
    print(f"Error caught: {e}")

# Example: Different state configurations
sir_states = StateArray(np.zeros((3, 2)), state_names=["S", "I", "R"])
seirs_states = StateArray(np.zeros((5, 2)), state_names=["S", "E", "I", "R", "S2"])

print("SIR state names:", sir_states._state_names)
print("SEIRS state names:", seirs_states._state_names)

# Example: Error handling for invalid state names try: invalid_state = patch_states.X # X is not a valid state except AttributeError as e: print(f"Error caught: {e}") # Example: Different state configurations sir_states = StateArray(np.zeros((3, 2)), state_names=["S", "I", "R"]) seirs_states = StateArray(np.zeros((5, 2)), state_names=["S", "E", "I", "R", "S2"]) print("SIR state names:", sir_states._state_names) print("SEIRS state names:", seirs_states._state_names)

Summary¶

StateArray is a wrapper that makes epidemiological modeling code more readable and maintainable. It provides intuitive access to disease compartments while preserving all the performance and functionality of numpy arrays. In laser-measles, it is used for the patches LaserFrame.