CLT Toolkit API Reference

Docstrings and references for clt_toolkit package.

PROJECT_ROOT = Path(__file__).resolve().parent.parent module-attribute

Compartment

Bases: StateVariable

Class for epidemiological compartments (e.g. Susceptible, Exposed, Infected, etc...).

Attributes:

Name	Type	Description
current_inflow	np.ndarray of shape (A, R	Used to sum up all transition variable realizations incoming to this compartment for age-risk groups.
current_outflow	np.ndarray of shape (A, R	Used to sum up all transition variable realizations outgoing from this compartment for age-risk groups.

See StateVariable docstring for additional attributes and A, R definitions.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

class Compartment(StateVariable):
    """
    Class for epidemiological compartments (e.g. Susceptible,
        Exposed, Infected, etc...).

    Attributes:
        current_inflow (np.ndarray of shape (A, R)):
            Used to sum up all  transition variable realizations
            incoming to this compartment for age-risk groups.
        current_outflow (np.ndarray of shape (A, R)):
            Used to sum up all transition variable realizations
            outgoing from this compartment for age-risk groups.

    See `StateVariable` docstring for additional attributes
        and A, R definitions.
    """

    def __init__(self,
                 init_val):
        super().__init__(np.asarray(init_val, dtype=float))

        self.current_inflow = np.zeros(np.shape(init_val))
        self.current_outflow = np.zeros(np.shape(init_val))

    def update_current_val(self) -> None:
        """
        Updates `current_val` attribute in-place by adding
            `current_inflow` (sum of all incoming transition variables'
            realizations) and subtracting current outflow (sum of all
            outgoing transition variables' realizations).
        """
        self.current_val = self.current_val + self.current_inflow - self.current_outflow

    def reset_inflow(self) -> None:
        """
        Resets `current_inflow` attribute to zero array.
        """
        self.current_inflow = np.zeros(np.shape(self.current_inflow))

    def reset_outflow(self) -> None:
        """
        Resets `current_outflow` attribute to zero array.
        """
        self.current_outflow = np.zeros(np.shape(self.current_outflow))

reset_inflow() -> None

Resets current_inflow attribute to zero array.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def reset_inflow(self) -> None:
    """
    Resets `current_inflow` attribute to zero array.
    """
    self.current_inflow = np.zeros(np.shape(self.current_inflow))

reset_outflow() -> None

Resets current_outflow attribute to zero array.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def reset_outflow(self) -> None:
    """
    Resets `current_outflow` attribute to zero array.
    """
    self.current_outflow = np.zeros(np.shape(self.current_outflow))

update_current_val() -> None

Updates current_val attribute in-place by adding current_inflow (sum of all incoming transition variables' realizations) and subtracting current outflow (sum of all outgoing transition variables' realizations).

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def update_current_val(self) -> None:
    """
    Updates `current_val` attribute in-place by adding
        `current_inflow` (sum of all incoming transition variables'
        realizations) and subtracting current outflow (sum of all
        outgoing transition variables' realizations).
    """
    self.current_val = self.current_val + self.current_inflow - self.current_outflow

DataClassProtocol

Bases: Protocol

Source code in CLT_BaseModel/clt_toolkit/input_parsers.py

class DataClassProtocol(Protocol):
    __dataclass_fields__: dict

DynamicVal

Bases: StateVariable, ABC

Abstract base class for variables that dynamically adjust their values based the current values of other StateVariable instances.

This class should model social distancing (and more broadly, staged-alert policies). For example, if we consider a case where transmission rates decrease when number infected increase above a certain level, we can create a subclass of DynamicVal that models a coefficient that modifies transmission rates, depending on the epi compartments corresponding to infected individuals.

Inherits attributes from StateVariable.

See __init__ docstring for other attributes.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

class DynamicVal(StateVariable, ABC):
    """
    Abstract base class for variables that dynamically adjust
    their values based the current values of other `StateVariable`
    instances.

    This class should model social distancing (and more broadly,
    staged-alert policies). For example, if we consider a
    case where transmission rates decrease when number infected
    increase above a certain level, we can create a subclass of
    DynamicVal that models a coefficient that modifies transmission
    rates, depending on the epi compartments corresponding to
    infected individuals.

    Inherits attributes from `StateVariable`.

    See `__init__` docstring for other attributes.
    """

    def __init__(self,
                 init_val: Optional[np.ndarray | float] = None,
                 is_enabled: Optional[bool] = False):
        """

        Args:
            init_val (Optional[np.ndarray | float]):
                starting value(s) at the beginning of the simulation.
            is_enabled (Optional[bool]):
                if `False`, this dynamic value does not get updated
                during the simulation and defaults to its `init_val`.
                This is designed to allow easy toggling of
                simulations with or without staged alert policies
                and other interventions.
        """

        super().__init__(init_val)
        self.is_enabled = is_enabled

    @abstractmethod
    def update_current_val(self,
                           state: SubpopState,
                           params: SubpopParams) -> None:
        """
        Args:
            state (SubpopState):
                holds subpopulation simulation state (current values of
                `StateVariable` instances).
            params (SubpopParams):
                holds values of epidemiological parameters.
        """

__init__(init_val: Optional[np.ndarray | float] = None, is_enabled: Optional[bool] = False)

Parameters:

Name	Type	Description	Default
init_val	Optional[ndarray | float]	starting value(s) at the beginning of the simulation.	None
is_enabled	Optional[bool]	if False, this dynamic value does not get updated during the simulation and defaults to its init_val. This is designed to allow easy toggling of simulations with or without staged alert policies and other interventions.	False

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def __init__(self,
             init_val: Optional[np.ndarray | float] = None,
             is_enabled: Optional[bool] = False):
    """

    Args:
        init_val (Optional[np.ndarray | float]):
            starting value(s) at the beginning of the simulation.
        is_enabled (Optional[bool]):
            if `False`, this dynamic value does not get updated
            during the simulation and defaults to its `init_val`.
            This is designed to allow easy toggling of
            simulations with or without staged alert policies
            and other interventions.
    """

    super().__init__(init_val)
    self.is_enabled = is_enabled

update_current_val(state: SubpopState, params: SubpopParams) -> None abstractmethod

Parameters:

Name	Type	Description	Default
state	SubpopState	holds subpopulation simulation state (current values of StateVariable instances).	required
params	SubpopParams	holds values of epidemiological parameters.	required

Source code in CLT_BaseModel/clt_toolkit/base_components.py

@abstractmethod
def update_current_val(self,
                       state: SubpopState,
                       params: SubpopParams) -> None:
    """
    Args:
        state (SubpopState):
            holds subpopulation simulation state (current values of
            `StateVariable` instances).
        params (SubpopParams):
            holds values of epidemiological parameters.
    """

EpiMetric

Bases: StateVariable, ABC

Abstract base class for epi metrics in epidemiological model.

This is intended for variables that are aggregate deterministic functions of the SubpopState (including Compartment current_val's, other parameters, and time.)

For example, population-level immunity variables should be modeled as a EpiMetric subclass, with a concrete implementation of the abstract method get_change_in_current_val.

Inherits attributes from StateVariable.

Attributes:

Name	Type	Description
current_val	np.ndarray of shape (A, R	same size as init_val, holds current value of StateVariable for age-risk groups.
change_in_current_val		(np.ndarray of shape (A, R)): initialized to None, but during simulation holds change in current value of EpiMetric for age-risk groups (size A x R, where A is the number of risk groups and R is number of age groups).

See __init__ docstring for other attributes.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

class EpiMetric(StateVariable, ABC):
    """
    Abstract base class for epi metrics in epidemiological model.

    This is intended for variables that are aggregate deterministic functions of
    the `SubpopState` (including `Compartment` `current_val`'s, other parameters,
    and time.)

    For example, population-level immunity variables should be
    modeled as a `EpiMetric` subclass, with a concrete
    implementation of the abstract method `get_change_in_current_val`.

    Inherits attributes from `StateVariable`.

    Attributes:
        current_val (np.ndarray of shape (A, R)):
            same size as init_val, holds current value of `StateVariable`
            for age-risk groups.
        change_in_current_val : (np.ndarray of shape (A, R)):
            initialized to None, but during simulation holds change in
            current value of `EpiMetric` for age-risk groups
            (size A x R, where A is the number of risk groups and R is number
            of age groups).

    See `__init__` docstring for other attributes.
    """

    def __init__(self,
                 init_val):
        """
        Args:
            init_val (np.ndarray of shape (A, R)):
                2D array that contains nonnegative floats,
                corresponding to initial value of dynamic val,
                where i,jth entry corresponds to age group i and
                risk group j.
        """

        super().__init__(init_val)

        self.change_in_current_val = None

    @abstractmethod
    def get_change_in_current_val(self,
                                  state: SubpopState,
                                  params: SubpopParams,
                                  num_timesteps: int) -> np.ndarray:
        """
        Computes and returns change in current value of dynamic val,
        based on current state of the simulation and epidemiological parameters.

        NOTE:
            OUTPUT SHOULD ALREADY BE SCALED BY NUM_TIMESTEPS.

        Output should be a numpy array of size A x R, where A
        is number of age groups and R is number of risk groups.

        Args:
            state (SubpopState):
                holds subpopulation simulation state (current values of
                `StateVariable` instances).
            params (SubpopParams):
                holds values of epidemiological parameters.
            num_timesteps (int):
                number of timesteps per day -- used to determine time interval
                length for discretization.

        Returns:
            (np.ndarray of shape (A, R))
                size A x R, where A is the number of age groups and
                R is number of risk groups.
        """
        pass

    def update_current_val(self) -> None:
        """
        Adds `change_in_current_val` attribute to
        `current_val` attribute in-place.
        """

        self.current_val += self.change_in_current_val

__init__(init_val)

Parameters:

Name	Type	Description	Default
init_val	np.ndarray of shape (A, R	2D array that contains nonnegative floats, corresponding to initial value of dynamic val, where i,jth entry corresponds to age group i and risk group j.	required

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def __init__(self,
             init_val):
    """
    Args:
        init_val (np.ndarray of shape (A, R)):
            2D array that contains nonnegative floats,
            corresponding to initial value of dynamic val,
            where i,jth entry corresponds to age group i and
            risk group j.
    """

    super().__init__(init_val)

    self.change_in_current_val = None

get_change_in_current_val(state: SubpopState, params: SubpopParams, num_timesteps: int) -> np.ndarray abstractmethod

Computes and returns change in current value of dynamic val, based on current state of the simulation and epidemiological parameters.

NOTE

OUTPUT SHOULD ALREADY BE SCALED BY NUM_TIMESTEPS.

Output should be a numpy array of size A x R, where A is number of age groups and R is number of risk groups.

Parameters:

Name	Type	Description	Default
state	SubpopState	holds subpopulation simulation state (current values of StateVariable instances).	required
params	SubpopParams	holds values of epidemiological parameters.	required
num_timesteps	int	number of timesteps per day -- used to determine time interval length for discretization.	required

Returns:

Type	Description
ndarray	(np.ndarray of shape (A, R)) size A x R, where A is the number of age groups and R is number of risk groups.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

@abstractmethod
def get_change_in_current_val(self,
                              state: SubpopState,
                              params: SubpopParams,
                              num_timesteps: int) -> np.ndarray:
    """
    Computes and returns change in current value of dynamic val,
    based on current state of the simulation and epidemiological parameters.

    NOTE:
        OUTPUT SHOULD ALREADY BE SCALED BY NUM_TIMESTEPS.

    Output should be a numpy array of size A x R, where A
    is number of age groups and R is number of risk groups.

    Args:
        state (SubpopState):
            holds subpopulation simulation state (current values of
            `StateVariable` instances).
        params (SubpopParams):
            holds values of epidemiological parameters.
        num_timesteps (int):
            number of timesteps per day -- used to determine time interval
            length for discretization.

    Returns:
        (np.ndarray of shape (A, R))
            size A x R, where A is the number of age groups and
            R is number of risk groups.
    """
    pass

update_current_val() -> None

Adds change_in_current_val attribute to current_val attribute in-place.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def update_current_val(self) -> None:
    """
    Adds `change_in_current_val` attribute to
    `current_val` attribute in-place.
    """

    self.current_val += self.change_in_current_val

Experiment

Class to manage running multiple simulation replications on a SubpopModel or MetapopModel instance and query its results.

Also allows running a batch of simulation replications on a deterministic sequence of values for a given input (for example, to see how output changes as a function of a given input).

Also handles random sampling of inputs from a uniform distribution.

Parameters:

Name	Type	Description	Default
experiment_subpop_models	tuple	tuple of SubpopModel instances associated with the Experiment. If the Experiment is for a MetapopModel, then this tuple contains all the associated SubpopModel instances that comprise that MetapopModel. If the Experiment is for a SubpopModel only, then this tuple contains only that particular SubpopModel.	required
results_df	DataFrame	DataFrame holding simulation results from each simulation replication	required
has_been_run	bool	indicates if self.run_static_inputs has been executed.	required

See __init__ docstring for other attributes.

Source code in CLT_BaseModel/clt_toolkit/experiments.py

class Experiment:
    """
    Class to manage running multiple simulation replications
    on a `SubpopModel` or `MetapopModel` instance and query its results.

    Also allows running a batch of simulation replications on a
    deterministic sequence of values for a given input
    (for example, to see how output changes as a function of
    a given input).

    Also handles random sampling of inputs from a uniform
    distribution.

    Params:
        experiment_subpop_models (tuple):
            tuple of `SubpopModel` instances associated with the `Experiment`.
            If the `Experiment` is for a `MetapopModel`, then this tuple
            contains all the associated `SubpopModel` instances
            that comprise that `MetapopModel.` If the `Experiment` is for
            a `SubpopModel` only, then this tuple contains only that
            particular `SubpopModel`.
        results_df (pd.DataFrame):
            DataFrame holding simulation results from each
            `simulation` replication
        has_been_run (bool):
            indicates if `self.run_static_inputs` has been executed.

    See `__init__` docstring for other attributes.
    """

    def __init__(self,
                 model: SubpopModel | MetapopModel,
                 state_variables_to_record: list,
                 database_filename: str):

        """
        Params:
            model (SubpopModel | MetapopModel):
                SubpopModel or MetapopModel instance on which to
                run multiple replications.
            state_variables_to_record (list[str]):
                list or list-like of strings corresponding to
                state variables to record -- each string must match
                a state variable name on each SubpopModel in
                the MetapopModel.
            database_filename (str):
                must be valid filename with suffix ".db" --
                experiment results are saved to this SQL database
        """

        self.model = model
        self.state_variables_to_record = state_variables_to_record
        self.database_filename = database_filename

        self.has_been_run = False

        # Create experiment_subpop_models tuple
        # If model is MetapopModel instance, then this tuple is a list
        #   of all associated SubpopModel instances
        # If model is a SubpopModel instance, then this tuple
        #   only contains that SubpopModel.
        if isinstance(model, MetapopModel):
            experiment_subpop_models = tuple(model.subpop_models.values())
        elif isinstance(model, SubpopModel):
            experiment_subpop_models = (model,)
        else:
            raise ExperimentError("\"model\" argument must be an instance of SubpopModel "
                                  "or MetapopModel class.")
        self.experiment_subpop_models = experiment_subpop_models

        # Initialize results_df attribute -- this will store
        #   results of experiment run
        self.results_df = None

        # Make sure the state variables to record are valid -- the names
        #   of the state variables to record must match actual state variables
        #   on each SubpopModel
        for subpop_model in self.experiment_subpop_models:
            if not check_is_subset_list(state_variables_to_record,
                                        subpop_model.all_state_variables.keys()):
                raise ExperimentError(
                    f"\"state_variables_to_record\" list is not a subset "
                    "of the state variables on SubpopModel \"{subpop_name}\" -- "
                    "modify \"state_variables_to_record\" and re-initialize experiment.")

    def run_static_inputs(self,
                          num_reps: int,
                          simulation_end_day: int,
                          days_between_save_history: int = 1,
                          results_filename: str = None):
        """
        Runs the associated `SubpopModel` or `MetapopModel` for a
        given number of independent replications until `simulation_end_day`.
        Parameter values and initial values are the same across
        simulation replications. User can specify how often to save the
        history and a CSV file in which to store this history.

        Params:
            num_reps (positive int):
                number of independent simulation replications
                to run in an experiment.
            simulation_end_day (positive int):
                stop simulation at simulation_end_day (i.e. exclusive,
                simulate up to but not including simulation_end_day).
            days_between_save_history (positive int):
                indicates how often to save simulation results.
            results_filename (str):
                if specified, must be valid filename with suffix ".csv" --
                experiment results are saved to this CSV file.
        """

        if self.has_been_run:
            raise ExperimentError("Experiment has already been run. "
                                  "Create a new Experiment instance to simulate "
                                  "more replications.")

        else:
            self.has_been_run = True

            self.create_results_sql_table()

            self.simulate_reps_and_save_results(reps=num_reps,
                                                end_day=simulation_end_day,
                                                days_per_save=days_between_save_history,
                                                inputs_are_static=True,
                                                filename=results_filename)

    def get_state_var_df(self,
                         state_var_name: str,
                         subpop_name: str = None,
                         age_group: int = None,
                         risk_group: int = None,
                         results_filename: str = None) -> pd.DataFrame:
        """
        Get pandas DataFrame of recorded values of `StateVariable` given by
        `state_var_name`, in the `SubpopModel` given by `subpop_name`,
        for the age-risk group given by `age_group` and `risk_group`.
        If `subpop_name` is not specified, then values are summed across all
        associated subpopulations. Similarly, if `age_group` (or `risk_group`)
        is not specified, then values are summed across all age groups
        (or risk groups).

        Args:
            state_var_name (str):
                Name of the `StateVariable` to retrieve.
            subpop_name (Optional[str]):
                The name of the `SubpopModel` for filtering. If None, values are
                summed across all `SubpopModel` instances.
            age_group (Optional[int]):
                The age group to select. If None, values are summed across
                all age groups.
            risk_group (Optional[int]):
                The risk group to select. If None, values are summed across
                all risk groups.
            results_filename (Optional[str]):
                If provided, saves the resulting DataFrame as a CSV.

        Returns:
            A pandas DataFrame where rows represent the replication and columns indicate the
            simulation day (timepoint) of recording. DataFrame values are the `StateVariable`'s
            current_val or the sum of the `StateVariable`'s current_val across subpopulations,
            age groups, or risk groups (the combination of what is summed over is
            specified by the user -- details are in the part of this docstring describing
            this function's parameters).
        """

        if state_var_name not in self.state_variables_to_record:
            raise ExperimentError("\"state_var_name\" is not in \"self.state_variables_to_record\" --"
                                  "function call is invalid.")

        conn = sqlite3.connect(self.database_filename)

        # Query all results table entries where state_var_name matches
        # This will return results across all subpopulations, age groups,
        #   and risk groups
        df = get_sql_table_as_df(conn,
                                 "SELECT * FROM results WHERE state_var_name = ?",
                                 chunk_size=int(1e4),
                                 sql_query_params=(state_var_name,))

        conn.close()

        # Define filter conditions
        filters = {
            "subpop_name": subpop_name,
            "age_group": age_group,
            "risk_group": risk_group
        }

        # Filter DataFrame based on user-specified conditions
        #   (for example, if user specifies subpop_name, return subset of
        #   DataFrame where subpop_name matches)
        conditions = [(df[col] == value) for col, value in filters.items() if value is not None]
        df_filtered = df if not conditions else df[np.logical_and.reduce(conditions)]

        # Group DataFrame based on unique combinations of "rep" and "timepoint" columns
        # Then sum (numeric values only), return the "value" column, and reset the index
        #   so that "rep" and "timepoint" become regular columns and are not the index
        df_aggregated = \
            df_filtered.groupby(["rep",
                                 "timepoint"]).sum(numeric_only=True)["value"].reset_index()

        # Use pivot() function to reshape the DataFrame for its final form
        # The "timepoint" values are spread across new columns
        #   (creating a column for each unique timepoint).
        # The "value" column populates the corresponding cells.
        df_final = df_aggregated.pivot(index="rep",
                                       columns="timepoint",
                                       values="value")

        if results_filename:
            df_final.to_csv(results_filename)

        return df_final

    def log_current_vals_to_sql(self,
                                rep_counter: int,
                                experiment_cursor: sqlite3.Cursor) -> None:
        """
        For each subpopulation and state variable to record
        associated with this `Experiment`, save current values to
        "results" table in SQL database specified by `experiment_cursor`.

        Params:
            rep_counter (int):
                Current replication ID.
            experiment_cursor (sqlite3.Cursor):
                Cursor object connected to the database
                where results should be inserted.
        """

        for subpop_model in self.experiment_subpop_models:
            for state_var_name in self.state_variables_to_record:
                data = format_current_val_for_sql(subpop_model,
                                                  state_var_name,
                                                  rep_counter)
                experiment_cursor.executemany(
                    "INSERT INTO results VALUES (?, ?, ?, ?, ?, ?, ?)", data)

    def log_inputs_to_sql(self,
                          experiment_cursor: sqlite3.Cursor):
        """
        For each subpopulation, add a new table to SQL
        database specified by `experiment_cursor`. Each table
        contains information on inputs that vary across
        replications (either due to random sampling or
        user-specified deterministic sequence). Each table
        contains inputs information from `Experiment` attribute
        `self.inputs_realizations` for a given subpopulation.

        Params:
            experiment_cursor (sqlite3.Cursor):
                Cursor object connected to the database
                where results should be inserted.
        """

        for subpop_model in self.experiment_subpop_models:
            table_name = f'"{subpop_model.name}_inputs"'

            # Get the column names (dynamically, based on table)
            experiment_cursor.execute(f"PRAGMA table_info({table_name})")

            # Extract column names from the table info
            # But exclude the column name "rep"
            columns_info = experiment_cursor.fetchall()
            column_names = [col[1] for col in columns_info if col[1] != "rep"]

            # Create a placeholder string for the dynamic query
            placeholders = ", ".join(["?" for _ in column_names])  # Number of placeholders matches number of columns

            # Create the dynamic INSERT statement
            sql_statement = f"INSERT INTO {table_name} ({', '.join(column_names)}) VALUES ({placeholders})"

            # Create list of lists -- each nested list contains a sequence of values
            #   for that particular input
            subpop_inputs_realizations = self.inputs_realizations[subpop_model.name]
            inputs_vals_over_reps_list = \
                [np.array(subpop_inputs_realizations[input_name]).reshape(-1,1) for input_name in column_names]
            inputs_vals_over_reps_list = np.hstack(inputs_vals_over_reps_list)

            experiment_cursor.executemany(sql_statement, inputs_vals_over_reps_list)

    def simulate_reps_and_save_results(self,
                                       reps: int,
                                       end_day: int,
                                       days_per_save: int,
                                       inputs_are_static: bool,
                                       filename: str = None):
        """
        Helper function that executes main loop over
        replications in `Experiment` and saves results.

        Params:
            reps (int):
                number of independent simulation replications
                to run in an experiment.
            end_day (int):
                stop simulation at end_day (i.e. exclusive,
                simulate up to but not including end_day).
            days_per_save (int):
                indicates how often to save simulation results.
            inputs_are_static (bool):
                indicates if inputs are same across replications.
            filename (str):
                if specified, must be valid filename with suffix ".csv" --
                experiment results are saved to this CSV file.
        """

        # Override each subpop simulation_settings's save_daily_history attribute --
        #   set it to False -- because we will manually save history
        #   to results database according to user-defined
        #   days_between_save_history for all subpops
        for subpop_model in self.experiment_subpop_models:
            subpop_model.simulation_settings = \
                updated_dataclass(subpop_model.simulation_settings, {"save_daily_history": False})

        model = self.model

        # Connect to SQL database
        conn = sqlite3.connect(self.database_filename)
        cursor = conn.cursor()

        # Loop through replications
        for rep in range(reps):

            # Reset model and clear its history
            model.reset_simulation()

            # Simulate model and save results every `days_per_save` days
            while model.current_simulation_day < end_day:
                model.simulate_until_day(min(model.current_simulation_day + days_per_save,
                                             end_day))

                self.log_current_vals_to_sql(rep, cursor)

        self.results_df = get_sql_table_as_df(conn, "SELECT * FROM results", chunk_size=int(1e4))

        if filename:
            self.results_df.to_csv(filename)

        # Commit changes to database and close
        conn.commit()
        conn.close()

    def create_results_sql_table(self):
        """
        Create SQL database and save to `self.database_filename`.
        Create table named `results` with columns `subpop_name`,
        `state_var_name`, `age_group`, `risk_group`, `rep`, `timepoint`,
        and `value` to store results from each replication of experiment.
        """

        # Make sure user is not overwriting database
        if os.path.exists(self.database_filename):
            raise ExperimentError("Database already exists! Overwriting is not allowed. "
                                  "Delete existing .db file or change database_filename "
                                  "attribute.")

        # Connect to the SQLite database and create database
        # Create a cursor object to execute SQL commands
        # Initialize a table with columns given by column_names
        # Commit changes and close the connection
        conn = sqlite3.connect(self.database_filename)
        cursor = conn.cursor()
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS results (
            subpop_name TEXT,
            state_var_name TEXT,
            age_group INT,
            risk_group INT,
            rep INT,
            timepoint INT,
            value FLOAT,
            PRIMARY KEY (subpop_name, state_var_name, age_group, risk_group, rep, timepoint)
        )
        """)
        conn.commit()
        conn.close()

__init__(model: SubpopModel | MetapopModel, state_variables_to_record: list, database_filename: str)

Parameters:

Name	Type	Description	Default
model	SubpopModel | MetapopModel	SubpopModel or MetapopModel instance on which to run multiple replications.	required
state_variables_to_record	list[str]	list or list-like of strings corresponding to state variables to record -- each string must match a state variable name on each SubpopModel in the MetapopModel.	required
database_filename	str	must be valid filename with suffix ".db" -- experiment results are saved to this SQL database	required

Source code in CLT_BaseModel/clt_toolkit/experiments.py

def __init__(self,
             model: SubpopModel | MetapopModel,
             state_variables_to_record: list,
             database_filename: str):

    """
    Params:
        model (SubpopModel | MetapopModel):
            SubpopModel or MetapopModel instance on which to
            run multiple replications.
        state_variables_to_record (list[str]):
            list or list-like of strings corresponding to
            state variables to record -- each string must match
            a state variable name on each SubpopModel in
            the MetapopModel.
        database_filename (str):
            must be valid filename with suffix ".db" --
            experiment results are saved to this SQL database
    """

    self.model = model
    self.state_variables_to_record = state_variables_to_record
    self.database_filename = database_filename

    self.has_been_run = False

    # Create experiment_subpop_models tuple
    # If model is MetapopModel instance, then this tuple is a list
    #   of all associated SubpopModel instances
    # If model is a SubpopModel instance, then this tuple
    #   only contains that SubpopModel.
    if isinstance(model, MetapopModel):
        experiment_subpop_models = tuple(model.subpop_models.values())
    elif isinstance(model, SubpopModel):
        experiment_subpop_models = (model,)
    else:
        raise ExperimentError("\"model\" argument must be an instance of SubpopModel "
                              "or MetapopModel class.")
    self.experiment_subpop_models = experiment_subpop_models

    # Initialize results_df attribute -- this will store
    #   results of experiment run
    self.results_df = None

    # Make sure the state variables to record are valid -- the names
    #   of the state variables to record must match actual state variables
    #   on each SubpopModel
    for subpop_model in self.experiment_subpop_models:
        if not check_is_subset_list(state_variables_to_record,
                                    subpop_model.all_state_variables.keys()):
            raise ExperimentError(
                f"\"state_variables_to_record\" list is not a subset "
                "of the state variables on SubpopModel \"{subpop_name}\" -- "
                "modify \"state_variables_to_record\" and re-initialize experiment.")

create_results_sql_table()

Create SQL database and save to self.database_filename. Create table named results with columns subpop_name, state_var_name, age_group, risk_group, rep, timepoint, and value to store results from each replication of experiment.

Source code in CLT_BaseModel/clt_toolkit/experiments.py

def create_results_sql_table(self):
    """
    Create SQL database and save to `self.database_filename`.
    Create table named `results` with columns `subpop_name`,
    `state_var_name`, `age_group`, `risk_group`, `rep`, `timepoint`,
    and `value` to store results from each replication of experiment.
    """

    # Make sure user is not overwriting database
    if os.path.exists(self.database_filename):
        raise ExperimentError("Database already exists! Overwriting is not allowed. "
                              "Delete existing .db file or change database_filename "
                              "attribute.")

    # Connect to the SQLite database and create database
    # Create a cursor object to execute SQL commands
    # Initialize a table with columns given by column_names
    # Commit changes and close the connection
    conn = sqlite3.connect(self.database_filename)
    cursor = conn.cursor()
    cursor.execute("""
    CREATE TABLE IF NOT EXISTS results (
        subpop_name TEXT,
        state_var_name TEXT,
        age_group INT,
        risk_group INT,
        rep INT,
        timepoint INT,
        value FLOAT,
        PRIMARY KEY (subpop_name, state_var_name, age_group, risk_group, rep, timepoint)
    )
    """)
    conn.commit()
    conn.close()

get_state_var_df(state_var_name: str, subpop_name: str = None, age_group: int = None, risk_group: int = None, results_filename: str = None) -> pd.DataFrame

Get pandas DataFrame of recorded values of StateVariable given by state_var_name, in the SubpopModel given by subpop_name, for the age-risk group given by age_group and risk_group. If subpop_name is not specified, then values are summed across all associated subpopulations. Similarly, if age_group (or risk_group) is not specified, then values are summed across all age groups (or risk groups).

Parameters:

Name	Type	Description	Default
state_var_name	str	Name of the StateVariable to retrieve.	required
subpop_name	Optional[str]	The name of the SubpopModel for filtering. If None, values are summed across all SubpopModel instances.	None
age_group	Optional[int]	The age group to select. If None, values are summed across all age groups.	None
risk_group	Optional[int]	The risk group to select. If None, values are summed across all risk groups.	None
results_filename	Optional[str]	If provided, saves the resulting DataFrame as a CSV.	None

Returns:

Type	Description
DataFrame	A pandas DataFrame where rows represent the replication and columns indicate the
DataFrame	simulation day (timepoint) of recording. DataFrame values are the StateVariable's
DataFrame	current_val or the sum of the StateVariable's current_val across subpopulations,
DataFrame	age groups, or risk groups (the combination of what is summed over is
DataFrame	specified by the user -- details are in the part of this docstring describing
DataFrame	this function's parameters).

Source code in CLT_BaseModel/clt_toolkit/experiments.py

def get_state_var_df(self,
                     state_var_name: str,
                     subpop_name: str = None,
                     age_group: int = None,
                     risk_group: int = None,
                     results_filename: str = None) -> pd.DataFrame:
    """
    Get pandas DataFrame of recorded values of `StateVariable` given by
    `state_var_name`, in the `SubpopModel` given by `subpop_name`,
    for the age-risk group given by `age_group` and `risk_group`.
    If `subpop_name` is not specified, then values are summed across all
    associated subpopulations. Similarly, if `age_group` (or `risk_group`)
    is not specified, then values are summed across all age groups
    (or risk groups).

    Args:
        state_var_name (str):
            Name of the `StateVariable` to retrieve.
        subpop_name (Optional[str]):
            The name of the `SubpopModel` for filtering. If None, values are
            summed across all `SubpopModel` instances.
        age_group (Optional[int]):
            The age group to select. If None, values are summed across
            all age groups.
        risk_group (Optional[int]):
            The risk group to select. If None, values are summed across
            all risk groups.
        results_filename (Optional[str]):
            If provided, saves the resulting DataFrame as a CSV.

    Returns:
        A pandas DataFrame where rows represent the replication and columns indicate the
        simulation day (timepoint) of recording. DataFrame values are the `StateVariable`'s
        current_val or the sum of the `StateVariable`'s current_val across subpopulations,
        age groups, or risk groups (the combination of what is summed over is
        specified by the user -- details are in the part of this docstring describing
        this function's parameters).
    """

    if state_var_name not in self.state_variables_to_record:
        raise ExperimentError("\"state_var_name\" is not in \"self.state_variables_to_record\" --"
                              "function call is invalid.")

    conn = sqlite3.connect(self.database_filename)

    # Query all results table entries where state_var_name matches
    # This will return results across all subpopulations, age groups,
    #   and risk groups
    df = get_sql_table_as_df(conn,
                             "SELECT * FROM results WHERE state_var_name = ?",
                             chunk_size=int(1e4),
                             sql_query_params=(state_var_name,))

    conn.close()

    # Define filter conditions
    filters = {
        "subpop_name": subpop_name,
        "age_group": age_group,
        "risk_group": risk_group
    }

    # Filter DataFrame based on user-specified conditions
    #   (for example, if user specifies subpop_name, return subset of
    #   DataFrame where subpop_name matches)
    conditions = [(df[col] == value) for col, value in filters.items() if value is not None]
    df_filtered = df if not conditions else df[np.logical_and.reduce(conditions)]

    # Group DataFrame based on unique combinations of "rep" and "timepoint" columns
    # Then sum (numeric values only), return the "value" column, and reset the index
    #   so that "rep" and "timepoint" become regular columns and are not the index
    df_aggregated = \
        df_filtered.groupby(["rep",
                             "timepoint"]).sum(numeric_only=True)["value"].reset_index()

    # Use pivot() function to reshape the DataFrame for its final form
    # The "timepoint" values are spread across new columns
    #   (creating a column for each unique timepoint).
    # The "value" column populates the corresponding cells.
    df_final = df_aggregated.pivot(index="rep",
                                   columns="timepoint",
                                   values="value")

    if results_filename:
        df_final.to_csv(results_filename)

    return df_final

log_current_vals_to_sql(rep_counter: int, experiment_cursor: sqlite3.Cursor) -> None

For each subpopulation and state variable to record associated with this Experiment, save current values to "results" table in SQL database specified by experiment_cursor.

Parameters:

Name	Type	Description	Default
rep_counter	int	Current replication ID.	required
experiment_cursor	Cursor	Cursor object connected to the database where results should be inserted.	required

Source code in CLT_BaseModel/clt_toolkit/experiments.py

def log_current_vals_to_sql(self,
                            rep_counter: int,
                            experiment_cursor: sqlite3.Cursor) -> None:
    """
    For each subpopulation and state variable to record
    associated with this `Experiment`, save current values to
    "results" table in SQL database specified by `experiment_cursor`.

    Params:
        rep_counter (int):
            Current replication ID.
        experiment_cursor (sqlite3.Cursor):
            Cursor object connected to the database
            where results should be inserted.
    """

    for subpop_model in self.experiment_subpop_models:
        for state_var_name in self.state_variables_to_record:
            data = format_current_val_for_sql(subpop_model,
                                              state_var_name,
                                              rep_counter)
            experiment_cursor.executemany(
                "INSERT INTO results VALUES (?, ?, ?, ?, ?, ?, ?)", data)

log_inputs_to_sql(experiment_cursor: sqlite3.Cursor)

For each subpopulation, add a new table to SQL database specified by experiment_cursor. Each table contains information on inputs that vary across replications (either due to random sampling or user-specified deterministic sequence). Each table contains inputs information from Experiment attribute self.inputs_realizations for a given subpopulation.

Parameters:

Name	Type	Description	Default
experiment_cursor	Cursor	Cursor object connected to the database where results should be inserted.	required

Source code in CLT_BaseModel/clt_toolkit/experiments.py

def log_inputs_to_sql(self,
                      experiment_cursor: sqlite3.Cursor):
    """
    For each subpopulation, add a new table to SQL
    database specified by `experiment_cursor`. Each table
    contains information on inputs that vary across
    replications (either due to random sampling or
    user-specified deterministic sequence). Each table
    contains inputs information from `Experiment` attribute
    `self.inputs_realizations` for a given subpopulation.

    Params:
        experiment_cursor (sqlite3.Cursor):
            Cursor object connected to the database
            where results should be inserted.
    """

    for subpop_model in self.experiment_subpop_models:
        table_name = f'"{subpop_model.name}_inputs"'

        # Get the column names (dynamically, based on table)
        experiment_cursor.execute(f"PRAGMA table_info({table_name})")

        # Extract column names from the table info
        # But exclude the column name "rep"
        columns_info = experiment_cursor.fetchall()
        column_names = [col[1] for col in columns_info if col[1] != "rep"]

        # Create a placeholder string for the dynamic query
        placeholders = ", ".join(["?" for _ in column_names])  # Number of placeholders matches number of columns

        # Create the dynamic INSERT statement
        sql_statement = f"INSERT INTO {table_name} ({', '.join(column_names)}) VALUES ({placeholders})"

        # Create list of lists -- each nested list contains a sequence of values
        #   for that particular input
        subpop_inputs_realizations = self.inputs_realizations[subpop_model.name]
        inputs_vals_over_reps_list = \
            [np.array(subpop_inputs_realizations[input_name]).reshape(-1,1) for input_name in column_names]
        inputs_vals_over_reps_list = np.hstack(inputs_vals_over_reps_list)

        experiment_cursor.executemany(sql_statement, inputs_vals_over_reps_list)

run_static_inputs(num_reps: int, simulation_end_day: int, days_between_save_history: int = 1, results_filename: str = None)

Runs the associated SubpopModel or MetapopModel for a given number of independent replications until simulation_end_day. Parameter values and initial values are the same across simulation replications. User can specify how often to save the history and a CSV file in which to store this history.

Parameters:

Name	Type	Description	Default
num_reps	positive int	number of independent simulation replications to run in an experiment.	required
simulation_end_day	positive int	stop simulation at simulation_end_day (i.e. exclusive, simulate up to but not including simulation_end_day).	required
days_between_save_history	positive int	indicates how often to save simulation results.	1
results_filename	str	if specified, must be valid filename with suffix ".csv" -- experiment results are saved to this CSV file.	None

Source code in CLT_BaseModel/clt_toolkit/experiments.py

def run_static_inputs(self,
                      num_reps: int,
                      simulation_end_day: int,
                      days_between_save_history: int = 1,
                      results_filename: str = None):
    """
    Runs the associated `SubpopModel` or `MetapopModel` for a
    given number of independent replications until `simulation_end_day`.
    Parameter values and initial values are the same across
    simulation replications. User can specify how often to save the
    history and a CSV file in which to store this history.

    Params:
        num_reps (positive int):
            number of independent simulation replications
            to run in an experiment.
        simulation_end_day (positive int):
            stop simulation at simulation_end_day (i.e. exclusive,
            simulate up to but not including simulation_end_day).
        days_between_save_history (positive int):
            indicates how often to save simulation results.
        results_filename (str):
            if specified, must be valid filename with suffix ".csv" --
            experiment results are saved to this CSV file.
    """

    if self.has_been_run:
        raise ExperimentError("Experiment has already been run. "
                              "Create a new Experiment instance to simulate "
                              "more replications.")

    else:
        self.has_been_run = True

        self.create_results_sql_table()

        self.simulate_reps_and_save_results(reps=num_reps,
                                            end_day=simulation_end_day,
                                            days_per_save=days_between_save_history,
                                            inputs_are_static=True,
                                            filename=results_filename)

simulate_reps_and_save_results(reps: int, end_day: int, days_per_save: int, inputs_are_static: bool, filename: str = None)

Helper function that executes main loop over replications in Experiment and saves results.

Parameters:

Name	Type	Description	Default
reps	int	number of independent simulation replications to run in an experiment.	required
end_day	int	stop simulation at end_day (i.e. exclusive, simulate up to but not including end_day).	required
days_per_save	int	indicates how often to save simulation results.	required
inputs_are_static	bool	indicates if inputs are same across replications.	required
filename	str	if specified, must be valid filename with suffix ".csv" -- experiment results are saved to this CSV file.	None

Source code in CLT_BaseModel/clt_toolkit/experiments.py

def simulate_reps_and_save_results(self,
                                   reps: int,
                                   end_day: int,
                                   days_per_save: int,
                                   inputs_are_static: bool,
                                   filename: str = None):
    """
    Helper function that executes main loop over
    replications in `Experiment` and saves results.

    Params:
        reps (int):
            number of independent simulation replications
            to run in an experiment.
        end_day (int):
            stop simulation at end_day (i.e. exclusive,
            simulate up to but not including end_day).
        days_per_save (int):
            indicates how often to save simulation results.
        inputs_are_static (bool):
            indicates if inputs are same across replications.
        filename (str):
            if specified, must be valid filename with suffix ".csv" --
            experiment results are saved to this CSV file.
    """

    # Override each subpop simulation_settings's save_daily_history attribute --
    #   set it to False -- because we will manually save history
    #   to results database according to user-defined
    #   days_between_save_history for all subpops
    for subpop_model in self.experiment_subpop_models:
        subpop_model.simulation_settings = \
            updated_dataclass(subpop_model.simulation_settings, {"save_daily_history": False})

    model = self.model

    # Connect to SQL database
    conn = sqlite3.connect(self.database_filename)
    cursor = conn.cursor()

    # Loop through replications
    for rep in range(reps):

        # Reset model and clear its history
        model.reset_simulation()

        # Simulate model and save results every `days_per_save` days
        while model.current_simulation_day < end_day:
            model.simulate_until_day(min(model.current_simulation_day + days_per_save,
                                         end_day))

            self.log_current_vals_to_sql(rep, cursor)

    self.results_df = get_sql_table_as_df(conn, "SELECT * FROM results", chunk_size=int(1e4))

    if filename:
        self.results_df.to_csv(filename)

    # Commit changes to database and close
    conn.commit()
    conn.close()

ExperimentError

Bases: Exception

Custom exceptions for experiment errors.

Source code in CLT_BaseModel/clt_toolkit/experiments.py

class ExperimentError(Exception):
    """Custom exceptions for experiment errors."""
    pass

InteractionTerm

Bases: StateVariable, ABC

Abstract base class for variables that depend on the state of more than one SubpopModel (i.e., that depend on more than one SubpopState). These variables are functions of how subpopulations interact.

Inherits attributes from StateVariable.

See __init__ docstring for other attributes.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

class InteractionTerm(StateVariable, ABC):
    """
    Abstract base class for variables that depend on the state of
    more than one `SubpopModel` (i.e., that depend on more than one
    `SubpopState`). These variables are functions of how subpopulations
    interact.

    Inherits attributes from `StateVariable`.

    See `__init__` docstring for other attributes.
    """

    @abstractmethod
    def update_current_val(self,
                           subpop_state: SubpopState,
                           subpop_params: SubpopParams) -> None:
        """
        Subclasses must provide a concrete implementation of
        updating `current_val` in-place.

        Args:
            subpop_params (SubpopParams):
                holds values of subpopulation's epidemiological parameters.
        """

        pass

update_current_val(subpop_state: SubpopState, subpop_params: SubpopParams) -> None abstractmethod

Subclasses must provide a concrete implementation of updating current_val in-place.

Parameters:

Name	Type	Description	Default
subpop_params	SubpopParams	holds values of subpopulation's epidemiological parameters.	required

Source code in CLT_BaseModel/clt_toolkit/base_components.py

@abstractmethod
def update_current_val(self,
                       subpop_state: SubpopState,
                       subpop_params: SubpopParams) -> None:
    """
    Subclasses must provide a concrete implementation of
    updating `current_val` in-place.

    Args:
        subpop_params (SubpopParams):
            holds values of subpopulation's epidemiological parameters.
    """

    pass

JointTransitionTypes

Bases: str, Enum

Defines available options for transition_type in TransitionVariableGroup.

Source code in CLT_BaseModel/clt_toolkit/base_data_structures.py

class JointTransitionTypes(str, Enum):
    """
    Defines available options for `transition_type` in `TransitionVariableGroup`.
    """
    MULTINOM = "multinom"
    MULTINOM_DETERMINISTIC = "multinom_deterministic"
    MULTINOM_TAYLOR_APPROX = "multinom_taylor_approx"
    MULTINOM_TAYLOR_APPROX_DETERMINISTIC = "multinom_taylor_approx_deterministic"
    POISSON = "poisson"
    POISSON_DETERMINISTIC = "poisson_deterministic"

MetapopModel

Bases: ABC

Abstract base class that bundles SubpopModels linked using a travel model.

See __init__ docstring for other attributes.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

class MetapopModel(ABC):
    """
    Abstract base class that bundles `SubpopModel`s linked using
        a travel model.

    See `__init__` docstring for other attributes.
    """

    def __init__(self,
                 subpop_models: list[dict],
                 mixing_params: dict,
                 name: str = ""):
        """
        Params:
            name (str):
                unique identifier for `MetapopModel`.
        """

        # We use both an `objdict` and an `odict` (ordered):
        # - `objdict`: allows convenient dot-access for users (consistent with the rest of the model)
        # - `odict`: preserves the order of subpopulations, which is crucial because
        #   the index in the state and parameter tensors depends on it.
        # The `objdict` is "outwards-facing" for user access, while the `odict`
        # is used internally to ensure tensor indices are consistent.

        subpop_models_dict = sc.objdict()
        for model in subpop_models:
            subpop_models_dict[model.name] = model

        _subpop_models_ordered_dict = sc.odict()
        for model in subpop_models:
            _subpop_models_ordered_dict[model.name] = model

        self.subpop_models = subpop_models_dict
        self._subpop_models_ordered = _subpop_models_ordered_dict

        self.name = name

        # Concrete implementations of `MetapopModel` will generally
        #   do something more with these parameters -- but this is
        #   just default storage here
        self.mixing_params = mixing_params

        for model in self.subpop_models.values():
            model.metapop_model = self

    def __getattr__(self, name):
        """
        Called if normal attribute lookup fails.
        Delegate to `subpop_models` if name matches a key.
        """

        if name in self.subpop_models:
            return self.subpop_models[name]
        else:
            raise AttributeError(f"{type(self).__name__!r} object has no attribute {name!r}")

    def modify_simulation_settings(self,
                                   updates_dict: dict):
        """
        This method applies the changes specified in `updates_dict` to the
        `simulation_settings` attribute of each subpopulation model.
        `SimulationSettings` is a frozen dataclass to prevent users from
        mutating individual subpop settings directly and making subpop
        models have different settings within the same metapop model.
        Instead, a new instance is created with the requested updates.

        Parameters:
            updates_dict (dict):
                Dictionary specifying values to update in a
                `SimulationSettings` instance -- keys must match the
                field names of `SimulationSettings`.
        """

        for subpop_model in self.subpop_models.values():
            subpop_model.modify_simulation_settings(updates_dict)

    def simulate_until_day(self,
                           simulation_end_day: int) -> None:
        """
        Advance simulation model time until `simulation_end_day` in
        `MetapopModel`.

        NOT just the same as looping through each `SubpopModel`'s
        `simulate_until_day` method. On the `MetapopModel`,
        because `SubpopModel` instances are linked with `InteractionTerm`s
        and are not independent of each other, this `MetapopModel`'s
        `simulate_until_day` method has additional functionality.

        Note: the update order at the beginning of each day is very important!

        - First, each `SubpopModel` updates its daily state (computing
        `Schedule` and `DynamicVal` instances).

        - Second, the `MetapopModel` computes quantities that depend
        on more than one subpopulation (i.e. inter-subpop quantities,
        such as the force of infection to each subpopulation in a travel
        model, where these terms depend on the number infected in
        other subpopulations) and then applies the update to each
        `SubpopModel` according to the user-implemented method
        `apply_inter_subpop_updates.`

        - Third, each `SubpopModel` simulates discretized timesteps (sampling
        `TransitionVariable`s, updating `EpiMetric`s, and updating `Compartment`s).

        Note: we only update inter-subpop quantities once a day, not at every timestep
        -- in other words, the travel model state-dependent values are only
        updated daily -- this is to avoid severe computation inefficiency

        Args:
            simulation_end_day (positive int):
                stop simulation at `simulation_end_day` (i.e. exclusive,
                simulate up to but not including `simulation_end_day`).
        """

        if self.current_simulation_day > simulation_end_day:
            raise MetapopModelError(f"Current day counter ({self.current_simulation_day}) "
                                    f"exceeds last simulation day ({simulation_end_day}).")

        # Adding this in case the user manually changes the initial
        #   value or current value of any state variable --
        #   otherwise, the state will not get updated
        # Analogous logic in SubpopModel's `simulate_until_day` method
        for subpop_model in self.subpop_models.values():
            subpop_model.state.sync_to_current_vals(subpop_model.all_state_variables)

        while self.current_simulation_day < simulation_end_day:

            for subpop_model in self.subpop_models.values():
                subpop_model.prepare_daily_state()

            self.apply_inter_subpop_updates()

            for subpop_model in self.subpop_models.values():

                save_daily_history = subpop_model.simulation_settings.save_daily_history
                timesteps_per_day = subpop_model.simulation_settings.timesteps_per_day

                subpop_model._simulate_timesteps(timesteps_per_day)

                if save_daily_history:
                    subpop_model.save_daily_history()

                subpop_model.increment_simulation_day()

    def apply_inter_subpop_updates(self):
        """
        `MetapopModel` subclasses can **optionally** override this method
        with a customized implementation. Otherwise, by default does nothing.

        Called once a day (not for each discretized timestep), after each
        subpop model's daily state is prepared, and before
        discretized transitions are computed.

        This method computes quantities that depend on multiple subpopulations
        (e.g. this is where a travel model should be implemented).

        See `simulate_until_day` method for more details.
        """

        pass

    def reset_simulation(self):
        """
        Resets `MetapopModel` by resetting and clearing
        history on all `SubpopModel` instances in
        `subpop_models`.
        """

        for subpop_model in self.subpop_models.values():
            subpop_model.reset_simulation()

    @property
    def current_simulation_day(self) -> int:
        """
        Returns:
            Current simulation day. The current simulation day of the
            `MetapopModel` should be the same as each individual `SubpopModel`
            in the `MetapopModel`. Otherwise, an error is raised.
        """

        current_simulation_days_list = []

        for subpop_model in self.subpop_models.values():
            current_simulation_days_list.append(subpop_model.current_simulation_day)

        if len(set(current_simulation_days_list)) > 1:
            raise MetapopModelError("Subpopulation models are on different simulation days "
                                    "and are out-of-sync. This may be caused by simulating "
                                    "a subpopulation model independently from the "
                                    "metapopulation model. Fix error and try again.")
        else:
            return current_simulation_days_list[0]

    @property
    def current_real_date(self) -> datetime.date:
        """
        Returns:
            Current real date corresponding to current simulation day.
            The current real date of the `MetapopModel` should be the same as
            each individual `SubpopModel` in the `MetapopModel`.
            Otherwise, an error is raised.
        """

        current_real_dates_list = []

        for subpop_model in self.subpop_models.values():
            current_real_dates_list.append(subpop_model.current_real_date)

        if len(set(current_real_dates_list)) > 1:
            raise MetapopModelError("Subpopulation models are on different real dates \n"
                                    "and are out-of-sync. This may be caused by simulating \n"
                                    "a subpopulation model independently from the \n"
                                    "metapopulation model. Please reset and restart simulation, \n"
                                    "and try again.")
        else:
            return current_real_dates_list[0]

current_real_date: datetime.date property

Returns:

Type	Description
date	Current real date corresponding to current simulation day.
date	The current real date of the MetapopModel should be the same as
date	each individual SubpopModel in the MetapopModel.
date	Otherwise, an error is raised.

current_simulation_day: int property

Returns:

Type	Description
int	Current simulation day. The current simulation day of the
int	MetapopModel should be the same as each individual SubpopModel
int	in the MetapopModel. Otherwise, an error is raised.

__getattr__(name)

Called if normal attribute lookup fails. Delegate to subpop_models if name matches a key.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def __getattr__(self, name):
    """
    Called if normal attribute lookup fails.
    Delegate to `subpop_models` if name matches a key.
    """

    if name in self.subpop_models:
        return self.subpop_models[name]
    else:
        raise AttributeError(f"{type(self).__name__!r} object has no attribute {name!r}")

__init__(subpop_models: list[dict], mixing_params: dict, name: str = '')

Parameters:

Name	Type	Description	Default
name	str	unique identifier for MetapopModel.	''

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def __init__(self,
             subpop_models: list[dict],
             mixing_params: dict,
             name: str = ""):
    """
    Params:
        name (str):
            unique identifier for `MetapopModel`.
    """

    # We use both an `objdict` and an `odict` (ordered):
    # - `objdict`: allows convenient dot-access for users (consistent with the rest of the model)
    # - `odict`: preserves the order of subpopulations, which is crucial because
    #   the index in the state and parameter tensors depends on it.
    # The `objdict` is "outwards-facing" for user access, while the `odict`
    # is used internally to ensure tensor indices are consistent.

    subpop_models_dict = sc.objdict()
    for model in subpop_models:
        subpop_models_dict[model.name] = model

    _subpop_models_ordered_dict = sc.odict()
    for model in subpop_models:
        _subpop_models_ordered_dict[model.name] = model

    self.subpop_models = subpop_models_dict
    self._subpop_models_ordered = _subpop_models_ordered_dict

    self.name = name

    # Concrete implementations of `MetapopModel` will generally
    #   do something more with these parameters -- but this is
    #   just default storage here
    self.mixing_params = mixing_params

    for model in self.subpop_models.values():
        model.metapop_model = self

apply_inter_subpop_updates()

MetapopModel subclasses can optionally override this method with a customized implementation. Otherwise, by default does nothing.

Called once a day (not for each discretized timestep), after each subpop model's daily state is prepared, and before discretized transitions are computed.

This method computes quantities that depend on multiple subpopulations (e.g. this is where a travel model should be implemented).

See simulate_until_day method for more details.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def apply_inter_subpop_updates(self):
    """
    `MetapopModel` subclasses can **optionally** override this method
    with a customized implementation. Otherwise, by default does nothing.

    Called once a day (not for each discretized timestep), after each
    subpop model's daily state is prepared, and before
    discretized transitions are computed.

    This method computes quantities that depend on multiple subpopulations
    (e.g. this is where a travel model should be implemented).

    See `simulate_until_day` method for more details.
    """

    pass

modify_simulation_settings(updates_dict: dict)

This method applies the changes specified in updates_dict to the simulation_settings attribute of each subpopulation model. SimulationSettings is a frozen dataclass to prevent users from mutating individual subpop settings directly and making subpop models have different settings within the same metapop model. Instead, a new instance is created with the requested updates.

Parameters:

Name	Type	Description	Default
updates_dict	dict	Dictionary specifying values to update in a SimulationSettings instance -- keys must match the field names of SimulationSettings.	required

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def modify_simulation_settings(self,
                               updates_dict: dict):
    """
    This method applies the changes specified in `updates_dict` to the
    `simulation_settings` attribute of each subpopulation model.
    `SimulationSettings` is a frozen dataclass to prevent users from
    mutating individual subpop settings directly and making subpop
    models have different settings within the same metapop model.
    Instead, a new instance is created with the requested updates.

    Parameters:
        updates_dict (dict):
            Dictionary specifying values to update in a
            `SimulationSettings` instance -- keys must match the
            field names of `SimulationSettings`.
    """

    for subpop_model in self.subpop_models.values():
        subpop_model.modify_simulation_settings(updates_dict)

reset_simulation()

Resets MetapopModel by resetting and clearing history on all SubpopModel instances in subpop_models.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def reset_simulation(self):
    """
    Resets `MetapopModel` by resetting and clearing
    history on all `SubpopModel` instances in
    `subpop_models`.
    """

    for subpop_model in self.subpop_models.values():
        subpop_model.reset_simulation()

simulate_until_day(simulation_end_day: int) -> None

Advance simulation model time until simulation_end_day in MetapopModel.

NOT just the same as looping through each SubpopModel's simulate_until_day method. On the MetapopModel, because SubpopModel instances are linked with InteractionTerms and are not independent of each other, this MetapopModel's simulate_until_day method has additional functionality.

Note: the update order at the beginning of each day is very important!

• First, each SubpopModel updates its daily state (computing Schedule and DynamicVal instances).

• Second, the MetapopModel computes quantities that depend on more than one subpopulation (i.e. inter-subpop quantities, such as the force of infection to each subpopulation in a travel model, where these terms depend on the number infected in other subpopulations) and then applies the update to each SubpopModel according to the user-implemented method apply_inter_subpop_updates.

• Third, each SubpopModel simulates discretized timesteps (sampling TransitionVariables, updating EpiMetrics, and updating Compartments).

Note: we only update inter-subpop quantities once a day, not at every timestep -- in other words, the travel model state-dependent values are only updated daily -- this is to avoid severe computation inefficiency

Parameters:

Name	Type	Description	Default
simulation_end_day	positive int	stop simulation at simulation_end_day (i.e. exclusive, simulate up to but not including simulation_end_day).	required

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def simulate_until_day(self,
                       simulation_end_day: int) -> None:
    """
    Advance simulation model time until `simulation_end_day` in
    `MetapopModel`.

    NOT just the same as looping through each `SubpopModel`'s
    `simulate_until_day` method. On the `MetapopModel`,
    because `SubpopModel` instances are linked with `InteractionTerm`s
    and are not independent of each other, this `MetapopModel`'s
    `simulate_until_day` method has additional functionality.

    Note: the update order at the beginning of each day is very important!

    - First, each `SubpopModel` updates its daily state (computing
    `Schedule` and `DynamicVal` instances).

    - Second, the `MetapopModel` computes quantities that depend
    on more than one subpopulation (i.e. inter-subpop quantities,
    such as the force of infection to each subpopulation in a travel
    model, where these terms depend on the number infected in
    other subpopulations) and then applies the update to each
    `SubpopModel` according to the user-implemented method
    `apply_inter_subpop_updates.`

    - Third, each `SubpopModel` simulates discretized timesteps (sampling
    `TransitionVariable`s, updating `EpiMetric`s, and updating `Compartment`s).

    Note: we only update inter-subpop quantities once a day, not at every timestep
    -- in other words, the travel model state-dependent values are only
    updated daily -- this is to avoid severe computation inefficiency

    Args:
        simulation_end_day (positive int):
            stop simulation at `simulation_end_day` (i.e. exclusive,
            simulate up to but not including `simulation_end_day`).
    """

    if self.current_simulation_day > simulation_end_day:
        raise MetapopModelError(f"Current day counter ({self.current_simulation_day}) "
                                f"exceeds last simulation day ({simulation_end_day}).")

    # Adding this in case the user manually changes the initial
    #   value or current value of any state variable --
    #   otherwise, the state will not get updated
    # Analogous logic in SubpopModel's `simulate_until_day` method
    for subpop_model in self.subpop_models.values():
        subpop_model.state.sync_to_current_vals(subpop_model.all_state_variables)

    while self.current_simulation_day < simulation_end_day:

        for subpop_model in self.subpop_models.values():
            subpop_model.prepare_daily_state()

        self.apply_inter_subpop_updates()

        for subpop_model in self.subpop_models.values():

            save_daily_history = subpop_model.simulation_settings.save_daily_history
            timesteps_per_day = subpop_model.simulation_settings.timesteps_per_day

            subpop_model._simulate_timesteps(timesteps_per_day)

            if save_daily_history:
                subpop_model.save_daily_history()

            subpop_model.increment_simulation_day()

MetapopModelError

Bases: Exception

Custom exceptions for metapopulation simulation model errors.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

class MetapopModelError(Exception):
    """Custom exceptions for metapopulation simulation model errors."""
    pass

ParamShapes

Bases: str, Enum

Defines allowed structural shapes for parameter sampling.

Specifies which population dimension(s) a parameter varies across. - "age": an array of length A (one value per age group) - "age_risk": a 2D array of shape (A, R) (values per age × risk group) - "scalar": a single value applied to all subpopulations

Used in UniformSamplingSpec.param_shapes to reduce the need for manually expanding arrays.

Source code in CLT_BaseModel/clt_toolkit/sampling.py

class ParamShapes(str, Enum):
    """
    Defines allowed structural shapes for parameter sampling.

    Specifies which population dimension(s) a parameter varies across.
      - "age": an array of length A (one value per age group)
      - "age_risk": a 2D array of shape (A, R) (values per age × risk group)
      - "scalar": a single value applied to all subpopulations

    Used in `UniformSamplingSpec.param_shapes` to reduce the need for
    manually expanding arrays.
    """

    age = "age"
    age_risk = "age_risk"
    scalar = "scalar"

Schedule dataclass

Bases: StateVariable, ABC

Abstract base class for variables that are functions of real-world dates -- for example, contact matrices (which depend on the day of the week and whether the current day is a holiday), historical vaccination data, and seasonality.

Inherits attributes from StateVariable.

See __init__ docstring for other attributes.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

@dataclass
class Schedule(StateVariable, ABC):
    """
    Abstract base class for variables that are functions of real-world
    dates -- for example, contact matrices (which depend on the day of
    the week and whether the current day is a holiday), historical
    vaccination data, and seasonality.

    Inherits attributes from `StateVariable`.

    See `__init__` docstring for other attributes.
    """

    def __init__(self,
                 init_val: Optional[np.ndarray | float] = None,
                 timeseries_df: Optional[dict] = None):
        """
        Args:
            init_val (Optional[np.ndarray | float]):
                starting value(s) at the beginning of the simulation
            timeseries_df (Optional[pd.DataFrame] = None):
                has a "date" column with strings in format `"YYYY-MM-DD"`
                of consecutive calendar days, and other columns
                corresponding to values on those days
        """

        super().__init__(init_val)
        self.timeseries_df = timeseries_df

    @abstractmethod
    def update_current_val(self,
                           params: SubpopParams,
                           current_date: datetime.date) -> None:
        """
        Subpop classes must provide a concrete implementation of
        updating `current_val` in-place.

        Args:
            params (SubpopParams):
                fixed parameters of subpopulation model.
            current_date (date):
                real-world date corresponding to
                model's current simulation day.
        """
        pass

__init__(init_val: Optional[np.ndarray | float] = None, timeseries_df: Optional[dict] = None)

Parameters:

Name	Type	Description	Default
init_val	Optional[ndarray | float]	starting value(s) at the beginning of the simulation	None
timeseries_df	Optional[pd.DataFrame] = None	has a "date" column with strings in format "YYYY-MM-DD" of consecutive calendar days, and other columns corresponding to values on those days	None

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def __init__(self,
             init_val: Optional[np.ndarray | float] = None,
             timeseries_df: Optional[dict] = None):
    """
    Args:
        init_val (Optional[np.ndarray | float]):
            starting value(s) at the beginning of the simulation
        timeseries_df (Optional[pd.DataFrame] = None):
            has a "date" column with strings in format `"YYYY-MM-DD"`
            of consecutive calendar days, and other columns
            corresponding to values on those days
    """

    super().__init__(init_val)
    self.timeseries_df = timeseries_df

update_current_val(params: SubpopParams, current_date: datetime.date) -> None abstractmethod

Subpop classes must provide a concrete implementation of updating current_val in-place.

Parameters:

Name	Type	Description	Default
params	SubpopParams	fixed parameters of subpopulation model.	required
current_date	date	real-world date corresponding to model's current simulation day.	required

Source code in CLT_BaseModel/clt_toolkit/base_components.py

@abstractmethod
def update_current_val(self,
                       params: SubpopParams,
                       current_date: datetime.date) -> None:
    """
    Subpop classes must provide a concrete implementation of
    updating `current_val` in-place.

    Args:
        params (SubpopParams):
            fixed parameters of subpopulation model.
        current_date (date):
            real-world date corresponding to
            model's current simulation day.
    """
    pass

SimulationSettings dataclass

Stores simulation settings.

Attributes:

Name	Type	Description
timesteps_per_day	int	number of discretized timesteps within a simulation day -- more timesteps_per_day mean smaller discretization time intervals, which may cause the model to run slower.
transition_type	str	valid value must be from TransitionTypes, specifying the probability distribution of transitions between compartments.
start_real_date	str	actual date in string format "YYYY-MM-DD" that aligns with the beginning of the simulation.
save_daily_history	bool	set to True to save current_val of StateVariable to history after each simulation day -- set to False if want speedier performance.
transition_variables_to_save	tuple	Names of transition variables whose histories should be saved during the simulation. Saving these can significantly slow execution, so leave this tuple empty for faster performance.

Source code in CLT_BaseModel/clt_toolkit/base_data_structures.py

@dataclass(frozen=True)
class SimulationSettings:
    """
    Stores simulation settings.

    Attributes:
        timesteps_per_day (int):
            number of discretized timesteps within a simulation
            day -- more `timesteps_per_day` mean smaller discretization
            time intervals, which may cause the model to run slower.
        transition_type (str):
            valid value must be from `TransitionTypes`, specifying
            the probability distribution of transitions between
            compartments.
        start_real_date (str):
            actual date in string format "YYYY-MM-DD" that aligns with the
            beginning of the simulation.
        save_daily_history (bool):
            set to `True` to save `current_val` of `StateVariable` to history after each
            simulation day -- set to `False` if want speedier performance.
        transition_variables_to_save (tuple):
            Names of transition variables whose histories should be saved
            during the simulation. Saving these can significantly slow
            execution, so leave this tuple empty for faster performance.
    """

    timesteps_per_day: int = 7
    transition_type: str = TransitionTypes.BINOM
    start_real_date: str = "2024-10-31"
    save_daily_history: bool = True
    transition_variables_to_save: tuple = ()

StateVariable

Parent class of InteractionTerm, Compartment, EpiMetric, DynamicVal, and Schedule classes. All subclasses have the attributes init_val and current_val.

Dimensions

A (int): Number of age groups. R (int): Number of risk groups.

Attributes:

Name	Type	Description
init_val	np.ndarray of shape (A, R	Holds initial value of StateVariable for age-risk groups.
current_val	np.ndarray of shape (A, R	Same size as init_val, holds current value of StateVariable for age-risk groups.
history_vals_list	list[ndarray]	Each element is an A x R array that holds history of compartment states for age-risk groups -- element t corresponds to previous current_val value at end of simulation day t.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

class StateVariable:
    """
    Parent class of `InteractionTerm`, `Compartment`, `EpiMetric`,
    `DynamicVal`, and `Schedule` classes. All subclasses have the
    attributes `init_val` and `current_val`.

    Dimensions:
        A (int):
            Number of age groups.
        R (int):
            Number of risk groups.

    Attributes:
        init_val (np.ndarray of shape (A, R)):
            Holds initial value of `StateVariable` for age-risk groups.
        current_val (np.ndarray of shape (A, R)):
            Same size as `init_val`, holds current value of `StateVariable`
            for age-risk groups.
        history_vals_list (list[np.ndarray]):
            Each element is an A x R array that holds
            history of compartment states for age-risk groups --
            element t corresponds to previous `current_val` value at
            end of simulation day t.
    """

    def __init__(self, init_val=None):
        self._init_val = init_val
        self.current_val = copy.deepcopy(init_val)
        self.history_vals_list = []

    @property
    def init_val(self):
        return self._init_val

    @init_val.setter
    def init_val(self, value):
        """
        We need to use properties/setters because when we change
        `init_val`, we want `current_val` to be updated too!
        """
        self._init_val = value
        self.current_val = copy.deepcopy(value)

    def save_history(self) -> None:
        """
        Saves current value to history by appending `current_val` attribute
        to `history_vals_list` in-place..

        Deep copying is CRUCIAL because `current_val` is a mutable
        `np.ndarray` -- without deep copying, `history_vals_list` would
        have the same value for all elements.
        """
        self.history_vals_list.append(copy.deepcopy(self.current_val))

    def reset(self) -> None:
        """
        Resets `current_val` to `init_val` and resets
        `history_vals_list` attribute to empty list.
        """

        self.current_val = copy.deepcopy(self.init_val)
        self.history_vals_list = []

reset() -> None

Resets current_val to init_val and resets history_vals_list attribute to empty list.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def reset(self) -> None:
    """
    Resets `current_val` to `init_val` and resets
    `history_vals_list` attribute to empty list.
    """

    self.current_val = copy.deepcopy(self.init_val)
    self.history_vals_list = []

save_history() -> None

Saves current value to history by appending current_val attribute to history_vals_list in-place..

Deep copying is CRUCIAL because current_val is a mutable np.ndarray -- without deep copying, history_vals_list would have the same value for all elements.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def save_history(self) -> None:
    """
    Saves current value to history by appending `current_val` attribute
    to `history_vals_list` in-place..

    Deep copying is CRUCIAL because `current_val` is a mutable
    `np.ndarray` -- without deep copying, `history_vals_list` would
    have the same value for all elements.
    """
    self.history_vals_list.append(copy.deepcopy(self.current_val))

SubpopModel

Bases: ABC

Contains and manages all necessary components for simulating a compartmental model for a given subpopulation.

Each SubpopModel instance includes compartments, epi metrics, dynamic vals, a data container for the current simulation state, transition variables and transition variable groups, epidemiological parameters, simulation experiment simulation settings parameters, and a random number generator.

All city-level subpopulation models, regardless of disease type and compartment/transition structure, are instances of this class.

When creating an instance, the order of elements does not matter within compartments, epi_metrics, dynamic_vals, transition_variables, and transition_variable_groups. The "flow" and "physics" information are stored on the objects.

Attributes:

Name	Type	Description
compartments	objdict[str, Compartment]	objdict of all the subpop model's Compartment instances.
transition_variables	objdict[str, TransitionVariable]	objdict of all the subpop model's TransitionVariable instances.
transition_variable_groups	objdict[str, TransitionVariableGroup]	objdict of all the subpop model's TransitionVariableGroup instances.
epi_metrics	objdict[str, EpiMetric]	objdict of all the subpop model's EpiMetric instances.
dynamic_vals	objdict[str, DynamicVal]	objdict of all the subpop model's DynamicVal instances.
schedules	objdict[str, Schedule]	objdict of all the subpop model's Schedule instances.
current_simulation_day	int	tracks current simulation day -- incremented by +1 when simulation_settings.timesteps_per_day discretized timesteps have completed.
current_real_date	date	tracks real-world date -- advanced by +1 day when simulation_settings.timesteps_per_day discretized timesteps have completed.

See __init__ docstring for other attributes.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

class SubpopModel(ABC):
    """
    Contains and manages all necessary components for
    simulating a compartmental model for a given subpopulation.

    Each `SubpopModel` instance includes compartments,
    epi metrics, dynamic vals, a data container for the current simulation
    state, transition variables and transition variable groups,
    epidemiological parameters, simulation experiment simulation settings
    parameters, and a random number generator.

    All city-level subpopulation models, regardless of disease type and
    compartment/transition structure, are instances of this class.

    When creating an instance, the order of elements does not matter
    within `compartments`, `epi_metrics`, `dynamic_vals`,
    `transition_variables`, and `transition_variable_groups`.
    The "flow" and "physics" information are stored on the objects.

    Attributes:
        compartments (sc.objdict[str, Compartment]):
            objdict of all the subpop model's `Compartment` instances.
        transition_variables (sc.objdict[str, TransitionVariable]):
            objdict of all the subpop model's `TransitionVariable` instances.
        transition_variable_groups (sc.objdict[str, TransitionVariableGroup]):
            objdict of all the subpop model's `TransitionVariableGroup` instances.
        epi_metrics (sc.objdict[str, EpiMetric]):
            objdict of all the subpop model's `EpiMetric` instances.
        dynamic_vals (sc.objdict[str, DynamicVal]):
            objdict of all the subpop model's `DynamicVal` instances.
        schedules (sc.objdict[str, Schedule]):
            objdict of all the subpop model's `Schedule` instances.
        current_simulation_day (int):
            tracks current simulation day -- incremented by +1
            when `simulation_settings.timesteps_per_day` discretized timesteps
            have completed.
        current_real_date (datetime.date):
            tracks real-world date -- advanced by +1 day when
            `simulation_settings.timesteps_per_day` discretized timesteps
            have completed.

    See `__init__` docstring for other attributes.
    """

    def __init__(self,
                 state: SubpopState,
                 params: SubpopParams,
                 simulation_settings: SimulationSettings,
                 RNG: np.random.Generator,
                 name: str,
                 metapop_model: MetapopModel = None):

        """
        Params:
            state (SubpopState):
                holds current values of `SubpopModel`'s state variables.
            params (SubpopParams):
                data container for the model's epidemiological parameters,
                such as the "Greek letters" characterizing sojourn times
                in compartments.
            simulation_settings (SimulationSettings):
                data container for the model's simulation settings.
            RNG (np.random.Generator):
                 used to generate stochastic transitions in the model and control
                 reproducibility.
            name (str):
                unique identifier of `SubpopModel`.
            metapop_model (Optional[MetapopModel]):
                if not `None`, is the `MetapopModel` instance
                associated with this `SubpopModel`.
        """

        self.state = copy.deepcopy(state)
        self.params = copy.deepcopy(params)
        self.simulation_settings = copy.deepcopy(simulation_settings)

        self.RNG = RNG

        self.current_simulation_day = 0
        self.start_real_date = self.get_start_real_date()
        self.current_real_date = self.start_real_date

        self.metapop_model = metapop_model
        self.name = name

        self.schedules = self.create_schedules()
        self.compartments = self.create_compartments()
        self.transition_variables = self.create_transition_variables()
        self.transition_variable_groups = self.create_transition_variable_groups()
        self.epi_metrics = self.create_epi_metrics()
        self.dynamic_vals = self.create_dynamic_vals()

        self.all_state_variables = {**self.compartments,
                                    **self.epi_metrics,
                                    **self.dynamic_vals,
                                    **self.schedules}

        # The model's state also has access to the model's
        #   compartments, epi_metrics, dynamic_vals, and schedules --
        #   so that state can easily retrieve each object's
        #   current_val and store it
        self.state.compartments = self.compartments
        self.state.epi_metrics = self.epi_metrics
        self.state.dynamic_vals = self.dynamic_vals
        self.state.schedules = self.schedules

        self.params = updated_dataclass(self.params, {"total_pop_age_risk": self.compute_total_pop_age_risk()})

    def __getattr__(self, name):
        """
        Called if normal attribute lookup fails.
        Delegate to `all_state_variables`, `transition_variables`,
            or `transition_variable_groups` if name matches a key.
        """

        if name in self.all_state_variables:
            return self.all_state_variables[name]
        elif name in self.transition_variables:
            return self.transition_variables[name]
        elif name in self.transition_variable_groups:
            return self.transition_variable_groups[name]
        else:
            raise AttributeError(f"{type(self).__name__!r} object has no attribute {name!r}")

    def modify_simulation_settings(self,
                                   updates_dict: dict):
        """
        This method lets users safely modify simulation settings;
        if this subpop model is associated with a metapop model,
        the same updates are applied to all subpop models on the
        metapop model. See also `modify_simulation_settings` method on
        `MetapopModel`.

        Parameters:
            updates_dict (dict):
                Dictionary specifying values to update in a
                `SimulationSettings` instance -- keys must match the
                field names of `SimulationSettings`.
        """

        self.simulation_settings = \
            updated_dataclass(self.simulation_settings, updates_dict)

    def compute_total_pop_age_risk(self) -> np.ndarray:
        """
        Returns:
            (np.ndarray of shape (A, R))
                A x R array, where A is the number of age groups
                and R is the number of risk groups, corresponding to
                total population for that age-risk group (summed
                over all compartments in the subpop model).
        """

        total_pop_age_risk = np.zeros((self.params.num_age_groups,
                                       self.params.num_risk_groups))

        # At initialization (before simulation is run), each
        #   compartment's current val is equivalent to the initial val
        #   specified in the state variables' init val JSON.
        for compartment in self.compartments.values():
            total_pop_age_risk += compartment.current_val

        return total_pop_age_risk

    def get_start_real_date(self):
        """
        Fetches `start_real_date` from `simulation_settings` -- converts to
            proper datetime.date format if originally given as
            string.

        Returns:
            start_real_date (datetime.date):
                real-world date that corresponds to start of
                simulation.
        """

        start_real_date = self.simulation_settings.start_real_date

        if not isinstance(start_real_date, datetime.date):
            try:
                start_real_date = \
                    datetime.datetime.strptime(start_real_date, "%Y-%m-%d").date()
            except ValueError:
                print("Error: The date format should be YYYY-MM-DD.")

        return start_real_date

    @abstractmethod
    def create_compartments(self) -> sc.objdict[str, Compartment]:
        """
        Create the epidemiological compartments used in the model.
        Subclasses **must override** this method to provide model-specific
        transitions.

        Returns:
            (sc.objdict[str, Compartment]):
                Dictionary mapping compartment names to `Compartment` objects.
        """

        return sc.objdict()

    @abstractmethod
    def create_transition_variables(self) -> sc.objdict[str, TransitionVariable]:
        """
        Create the transition variables specifying how individuals transition
        between epidemiological compartments in the model. Subclasses
        **must override** this method to provide model-specific transitions.

        See `__init__` method -- this method is called after `compartments`
        is assigned via `create_compartments()`, so it can reference the instance's
        compartments.

        Returns:
            (sc.objdict[str, TransitionVariable]):
                Dictionary mapping names to `TransitionVariable` objects.
        """

        return sc.objdict()

    def create_transition_variable_groups(self) -> sc.objdict[str, TransitionVariableGroup]:
        """
        Create the joint transition variables specifying how transitioning
        from compartments with multiple outflows is handled. Subclasses
        can **optionally** override this method to provide model-specific transitions.

        See `__init__` method -- this method is called after `compartments`
        is assigned via `create_compartments()` and `transition_variables` is
        assigned via `create_transition_variables()`, so it can reference the instance's
        compartments and transition variables.

        Returns:
            (sc.objdict[str, TransitionVariableGroup]):
                Dictionary mapping names to `TransitionVariableGroup` objects.
                Default is empty `objdict`.
        """

        return sc.objdict()

    def create_epi_metrics(self) -> sc.objdict[str, EpiMetric]:
        """
        Create the epidemiological metrics that track deterministic functions of
        compartments' current values. Subclasses can **optionally** override this method
        to provide model-specific transitions.

        See `__init__` method -- this method is called after `transition_variables` is
        assigned via `create_transition_variables()`, so it can reference the instance's
        transition variables.

        Returns:
            (sc.objdict[str, EpiMetric]):
                Dictionary mapping names to `EpiMetric` objects. Default is empty `objdict`.
        """

        return sc.objdict()

    def create_dynamic_vals(self) -> sc.objdict[str, DynamicVal]:
        """
        Create dynamic values that change depending on the simulation state.
        Subclasses can **optionally** override this method to provide model-specific transitions.

        Returns:
            (sc.objdict[str, DynamicVal]):
                Dictionary mapping names to `DynamicVal` objects. Default is empty `objdict`.
        """

        return sc.objdict()

    def create_schedules(self) -> sc.objdict[str, Schedule]:
        """
        Create schedules that are deterministic functions of the real-world simulation date.
        Subclasses can **optionally** override this method to provide model-specific transitions.

        Returns:
            (sc.objdict[str, Schedule]):
                Dictionary mapping names to `Schedule` objects. Default is empty `objdict`.
        """

        return sc.objdict()

    def modify_random_seed(self, new_seed_number) -> None:
        """
        Modifies model's `RNG` attribute in-place to new generator
        seeded at `new_seed_number`.

        Args:
            new_seed_number (int):
                used to re-seed model's random number generator.
        """

        self._bit_generator = np.random.MT19937(seed=new_seed_number)
        self.RNG = np.random.Generator(self._bit_generator)

    def simulate_until_day(self,
                           simulation_end_day: int) -> None:
        """
        Advance simulation model time until `simulation_end_day`.

        Advance time by iterating through simulation days,
        which are simulated by iterating through discretized
        timesteps.

        Save daily simulation data as history on each `Compartment`
        instance.

        Args:
            simulation_end_day (positive int):
                stop simulation at `simulation_end_day` (i.e. exclusive,
                simulate up to but not including `simulation_end_day`).
        """

        if self.current_simulation_day > simulation_end_day:
            raise SubpopModelError(f"Current day counter ({self.current_simulation_day}) "
                                   f"exceeds last simulation day ({simulation_end_day}).")

        save_daily_history = self.simulation_settings.save_daily_history
        timesteps_per_day = self.simulation_settings.timesteps_per_day

        # Adding this in case the user manually changes the initial
        #   value or current value of any state variable --
        #   otherwise, the state will not get updated
        self.state.sync_to_current_vals(self.all_state_variables)

        # simulation_end_day is exclusive endpoint
        while self.current_simulation_day < simulation_end_day:

            self.prepare_daily_state()

            self._simulate_timesteps(timesteps_per_day)

            if save_daily_history:
                self.save_daily_history()

            self.increment_simulation_day()

    def _simulate_timesteps(self,
                            num_timesteps: int) -> None:
        """
        Subroutine for `simulate_until_day`.

        Iterates through discretized timesteps to simulate next
        simulation day. Granularity of discretization is given by
        attribute `simulation_settings.timesteps_per_day`.

        Properly scales transition variable realizations and changes
        in dynamic vals by specified timesteps per day.

        Args:
            num_timesteps (int):
                number of timesteps per day -- used to determine time interval
                length for discretization.
        """

        for timestep in range(num_timesteps):

            self.update_transition_rates()

            self.sample_transitions()

            self.update_epi_metrics()

            self.update_compartments()

            self.state.sync_to_current_vals(self.epi_metrics)
            self.state.sync_to_current_vals(self.compartments)

    def prepare_daily_state(self) -> None:
        """
        At beginning of each day, update current value of
        interaction terms, schedules, dynamic values --
        note that these are only updated once a day, not
        for every discretized timestep.
        """

        subpop_state = self.state
        subpop_params = self.params
        current_real_date = self.current_real_date

        # Important note: this order of updating is important,
        #   because schedules do not depend on other state variables,
        #   but dynamic vals may depend on schedules
        # Interaction terms may depend on both schedules
        #   and dynamic vals.

        schedules = self.schedules
        dynamic_vals = self.dynamic_vals

        # Update schedules for current day
        for schedule in schedules.values():
            schedule.update_current_val(subpop_params,
                                        current_real_date)

        self.state.sync_to_current_vals(schedules)

        # Update dynamic values for current day
        for dval in dynamic_vals.values():
            if dval.is_enabled:
                dval.update_current_val(subpop_state, subpop_params)

        self.state.sync_to_current_vals(dynamic_vals)

    def update_epi_metrics(self) -> None:
        """
        Update current value attribute on each associated
            `EpiMetric` instance.
        """

        state = self.state
        params = self.params
        timesteps_per_day = self.simulation_settings.timesteps_per_day

        for metric in self.epi_metrics.values():
            metric.change_in_current_val = \
                metric.get_change_in_current_val(state,
                                                 params,
                                                 timesteps_per_day)
            metric.update_current_val()

    def update_transition_rates(self) -> None:
        """
        Compute current transition rates for each transition variable,
            and store this updated value on each variable's
            current_rate attribute.
        """

        state = self.state
        params = self.params

        for tvar in self.transition_variables.values():
            tvar.current_rate = tvar.get_current_rate(state, params)

    def sample_transitions(self) -> None:
        """
        For each transition variable, sample a random realization
            using its current rate. Handle jointly distributed transition
            variables first (using `TransitionVariableGroup` logic), then
            handle marginally distributed transition variables.
            Use `SubpopModel`'s `RNG` to generate random variables.
        """

        RNG = self.RNG
        timesteps_per_day = self.simulation_settings.timesteps_per_day
        transition_variables_to_save = self.simulation_settings.transition_variables_to_save

        # Obtain transition variable realizations for jointly distributed transition variables
        #   (i.e. when there are multiple transition variable outflows from an epi compartment)
        for tvargroup in self.transition_variable_groups.values():
            tvargroup.current_vals_list = tvargroup.get_joint_realization(RNG,
                                                                          timesteps_per_day)
            tvargroup.update_transition_variable_realizations()

        # Obtain transition variable realizations for marginally distributed transition variables
        #   (i.e. when there is only one transition variable outflow from an epi compartment)
        # If transition variable is jointly distributed, then its realization has already
        #   been computed by its transition variable group container previously,
        #   so skip the marginal computation
        for tvar in self.transition_variables.values():
            if not tvar.is_jointly_distributed:
                tvar.current_val = tvar.get_realization(RNG, timesteps_per_day)

        for name in transition_variables_to_save:
            self.transition_variables[name].save_history()

    def update_compartments(self) -> None:
        """
        Update current value of each `Compartment`, by
            looping through all `TransitionVariable` instances
            and subtracting/adding their current values
            from origin/destination compartments respectively.
        """

        for tvar in self.transition_variables.values():
            tvar.update_origin_outflow()
            tvar.update_destination_inflow()

        for compartment in self.compartments.values():
            compartment.update_current_val()

            # By construction (using binomial/multinomial with or without taylor expansion),
            #   more individuals cannot leave the compartment than are in the compartment
            # However, for Poisson any for ANY deterministic version, it is possible
            #   to have more individuals leaving the compartment than are in the compartment,
            #   and hence negative-valued compartments
            # We use this function to fix this, and also use a differentiable torch
            #   function to be consistent with the torch implementation (this still
            #   allows us to take derivatives in the torch implementation)
            # The syntax is janky here -- we want everything as an array, but
            #   we need to pass a tensor to the torch functional
            if "deterministic" in self.simulation_settings.transition_type:
                compartment.current_val = \
                        np.array(torch.nn.functional.softplus(torch.tensor(compartment.current_val)))

            # After updating the compartment's current value,
            #   reset its inflow and outflow attributes, to
            #   prepare for the next iteration.
            compartment.reset_inflow()
            compartment.reset_outflow()

    def increment_simulation_day(self) -> None:
        """
        Move day counters to next simulation day, both
            for integer simulation day and real date.
        """

        self.current_simulation_day += 1
        self.current_real_date += datetime.timedelta(days=1)

    def save_daily_history(self) -> None:
        """
        Update history at end of each day, not at end of every
           discretization timestep, to be efficient.
        Update history of state variables other than `Schedule`
           instances -- schedules do not have history.
        """
        for svar in self.compartments.values() + \
                    self.epi_metrics.values() + \
                    self.dynamic_vals.values():
            svar.save_history()

    def reset_simulation(self) -> None:
        """
        Reset simulation in-place. Subsequent method calls of
        `simulate_until_day` start from day 0, with original
        day 0 state.

        Returns `current_simulation_day` to 0.
        Restores state values to initial values.
        Clears history on model's state variables.
        Resets transition variables' `current_val` attribute to 0.

        WARNING:
            DOES NOT RESET THE MODEL'S RANDOM NUMBER GENERATOR TO
            ITS INITIAL STARTING SEED. RANDOM NUMBER GENERATOR WILL CONTINUE
            WHERE IT LEFT OFF.

        Use method `modify_random_seed` to reset model's `RNG` to its
        initial starting seed.
        """

        self.current_simulation_day = 0
        self.current_real_date = self.start_real_date

        # AGAIN, MUST BE CAREFUL ABOUT MUTABLE NUMPY ARRAYS -- MUST USE DEEP COPY
        for svar in self.all_state_variables.values():
            setattr(svar, "current_val", copy.deepcopy(svar.init_val))

        self.state.sync_to_current_vals(self.all_state_variables)

        for svar in self.all_state_variables.values():
            svar.reset()

        for tvar in self.transition_variables.values():
            tvar.reset()

        for tvargroup in self.transition_variable_groups.values():
            tvargroup.current_vals_list = []

    def find_name_by_compartment(self,
                                 target_compartment: Compartment) -> str:
        """
        Given `Compartment`, returns name of that `Compartment`.

        Args:
            target_compartment (Compartment):
                Compartment object with a name to look up

        Returns:
            (str):
                Compartment name, given by the key to look
                it up in the `SubpopModel`'s compartments objdict
        """

        for name, compartment in self.compartments.items():
            if compartment == target_compartment:
                return name

__getattr__(name)

Called if normal attribute lookup fails. Delegate to all_state_variables, transition_variables, or transition_variable_groups if name matches a key.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def __getattr__(self, name):
    """
    Called if normal attribute lookup fails.
    Delegate to `all_state_variables`, `transition_variables`,
        or `transition_variable_groups` if name matches a key.
    """

    if name in self.all_state_variables:
        return self.all_state_variables[name]
    elif name in self.transition_variables:
        return self.transition_variables[name]
    elif name in self.transition_variable_groups:
        return self.transition_variable_groups[name]
    else:
        raise AttributeError(f"{type(self).__name__!r} object has no attribute {name!r}")

__init__(state: SubpopState, params: SubpopParams, simulation_settings: SimulationSettings, RNG: np.random.Generator, name: str, metapop_model: MetapopModel = None)

Parameters:

Name	Type	Description	Default
state	SubpopState	holds current values of SubpopModel's state variables.	required
params	SubpopParams	data container for the model's epidemiological parameters, such as the "Greek letters" characterizing sojourn times in compartments.	required
simulation_settings	SimulationSettings	data container for the model's simulation settings.	required
RNG	Generator	used to generate stochastic transitions in the model and control reproducibility.	required
name	str	unique identifier of SubpopModel.	required
metapop_model	Optional[MetapopModel]	if not None, is the MetapopModel instance associated with this SubpopModel.	None

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def __init__(self,
             state: SubpopState,
             params: SubpopParams,
             simulation_settings: SimulationSettings,
             RNG: np.random.Generator,
             name: str,
             metapop_model: MetapopModel = None):

    """
    Params:
        state (SubpopState):
            holds current values of `SubpopModel`'s state variables.
        params (SubpopParams):
            data container for the model's epidemiological parameters,
            such as the "Greek letters" characterizing sojourn times
            in compartments.
        simulation_settings (SimulationSettings):
            data container for the model's simulation settings.
        RNG (np.random.Generator):
             used to generate stochastic transitions in the model and control
             reproducibility.
        name (str):
            unique identifier of `SubpopModel`.
        metapop_model (Optional[MetapopModel]):
            if not `None`, is the `MetapopModel` instance
            associated with this `SubpopModel`.
    """

    self.state = copy.deepcopy(state)
    self.params = copy.deepcopy(params)
    self.simulation_settings = copy.deepcopy(simulation_settings)

    self.RNG = RNG

    self.current_simulation_day = 0
    self.start_real_date = self.get_start_real_date()
    self.current_real_date = self.start_real_date

    self.metapop_model = metapop_model
    self.name = name

    self.schedules = self.create_schedules()
    self.compartments = self.create_compartments()
    self.transition_variables = self.create_transition_variables()
    self.transition_variable_groups = self.create_transition_variable_groups()
    self.epi_metrics = self.create_epi_metrics()
    self.dynamic_vals = self.create_dynamic_vals()

    self.all_state_variables = {**self.compartments,
                                **self.epi_metrics,
                                **self.dynamic_vals,
                                **self.schedules}

    # The model's state also has access to the model's
    #   compartments, epi_metrics, dynamic_vals, and schedules --
    #   so that state can easily retrieve each object's
    #   current_val and store it
    self.state.compartments = self.compartments
    self.state.epi_metrics = self.epi_metrics
    self.state.dynamic_vals = self.dynamic_vals
    self.state.schedules = self.schedules

    self.params = updated_dataclass(self.params, {"total_pop_age_risk": self.compute_total_pop_age_risk()})

compute_total_pop_age_risk() -> np.ndarray

Returns:

Type	Description
ndarray	(np.ndarray of shape (A, R)) A x R array, where A is the number of age groups and R is the number of risk groups, corresponding to total population for that age-risk group (summed over all compartments in the subpop model).

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def compute_total_pop_age_risk(self) -> np.ndarray:
    """
    Returns:
        (np.ndarray of shape (A, R))
            A x R array, where A is the number of age groups
            and R is the number of risk groups, corresponding to
            total population for that age-risk group (summed
            over all compartments in the subpop model).
    """

    total_pop_age_risk = np.zeros((self.params.num_age_groups,
                                   self.params.num_risk_groups))

    # At initialization (before simulation is run), each
    #   compartment's current val is equivalent to the initial val
    #   specified in the state variables' init val JSON.
    for compartment in self.compartments.values():
        total_pop_age_risk += compartment.current_val

    return total_pop_age_risk

create_compartments() -> sc.objdict[str, Compartment] abstractmethod

Create the epidemiological compartments used in the model. Subclasses must override this method to provide model-specific transitions.

Returns:

Type	Description
objdict[str, Compartment]	Dictionary mapping compartment names to Compartment objects.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

@abstractmethod
def create_compartments(self) -> sc.objdict[str, Compartment]:
    """
    Create the epidemiological compartments used in the model.
    Subclasses **must override** this method to provide model-specific
    transitions.

    Returns:
        (sc.objdict[str, Compartment]):
            Dictionary mapping compartment names to `Compartment` objects.
    """

    return sc.objdict()

create_dynamic_vals() -> sc.objdict[str, DynamicVal]

Create dynamic values that change depending on the simulation state. Subclasses can optionally override this method to provide model-specific transitions.

Returns:

Type	Description
objdict[str, DynamicVal]	Dictionary mapping names to DynamicVal objects. Default is empty objdict.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def create_dynamic_vals(self) -> sc.objdict[str, DynamicVal]:
    """
    Create dynamic values that change depending on the simulation state.
    Subclasses can **optionally** override this method to provide model-specific transitions.

    Returns:
        (sc.objdict[str, DynamicVal]):
            Dictionary mapping names to `DynamicVal` objects. Default is empty `objdict`.
    """

    return sc.objdict()

create_epi_metrics() -> sc.objdict[str, EpiMetric]

Create the epidemiological metrics that track deterministic functions of compartments' current values. Subclasses can optionally override this method to provide model-specific transitions.

See __init__ method -- this method is called after transition_variables is assigned via create_transition_variables(), so it can reference the instance's transition variables.

Returns:

Type	Description
objdict[str, EpiMetric]	Dictionary mapping names to EpiMetric objects. Default is empty objdict.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def create_epi_metrics(self) -> sc.objdict[str, EpiMetric]:
    """
    Create the epidemiological metrics that track deterministic functions of
    compartments' current values. Subclasses can **optionally** override this method
    to provide model-specific transitions.

    See `__init__` method -- this method is called after `transition_variables` is
    assigned via `create_transition_variables()`, so it can reference the instance's
    transition variables.

    Returns:
        (sc.objdict[str, EpiMetric]):
            Dictionary mapping names to `EpiMetric` objects. Default is empty `objdict`.
    """

    return sc.objdict()

create_schedules() -> sc.objdict[str, Schedule]

Create schedules that are deterministic functions of the real-world simulation date. Subclasses can optionally override this method to provide model-specific transitions.

Returns:

Type	Description
objdict[str, Schedule]	Dictionary mapping names to Schedule objects. Default is empty objdict.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def create_schedules(self) -> sc.objdict[str, Schedule]:
    """
    Create schedules that are deterministic functions of the real-world simulation date.
    Subclasses can **optionally** override this method to provide model-specific transitions.

    Returns:
        (sc.objdict[str, Schedule]):
            Dictionary mapping names to `Schedule` objects. Default is empty `objdict`.
    """

    return sc.objdict()

create_transition_variable_groups() -> sc.objdict[str, TransitionVariableGroup]

Create the joint transition variables specifying how transitioning from compartments with multiple outflows is handled. Subclasses can optionally override this method to provide model-specific transitions.

See __init__ method -- this method is called after compartments is assigned via create_compartments() and transition_variables is assigned via create_transition_variables(), so it can reference the instance's compartments and transition variables.

Returns:

Type	Description
objdict[str, TransitionVariableGroup]	Dictionary mapping names to TransitionVariableGroup objects. Default is empty objdict.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def create_transition_variable_groups(self) -> sc.objdict[str, TransitionVariableGroup]:
    """
    Create the joint transition variables specifying how transitioning
    from compartments with multiple outflows is handled. Subclasses
    can **optionally** override this method to provide model-specific transitions.

    See `__init__` method -- this method is called after `compartments`
    is assigned via `create_compartments()` and `transition_variables` is
    assigned via `create_transition_variables()`, so it can reference the instance's
    compartments and transition variables.

    Returns:
        (sc.objdict[str, TransitionVariableGroup]):
            Dictionary mapping names to `TransitionVariableGroup` objects.
            Default is empty `objdict`.
    """

    return sc.objdict()

create_transition_variables() -> sc.objdict[str, TransitionVariable] abstractmethod

Create the transition variables specifying how individuals transition between epidemiological compartments in the model. Subclasses must override this method to provide model-specific transitions.

See __init__ method -- this method is called after compartments is assigned via create_compartments(), so it can reference the instance's compartments.

Returns:

Type	Description
objdict[str, TransitionVariable]	Dictionary mapping names to TransitionVariable objects.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

@abstractmethod
def create_transition_variables(self) -> sc.objdict[str, TransitionVariable]:
    """
    Create the transition variables specifying how individuals transition
    between epidemiological compartments in the model. Subclasses
    **must override** this method to provide model-specific transitions.

    See `__init__` method -- this method is called after `compartments`
    is assigned via `create_compartments()`, so it can reference the instance's
    compartments.

    Returns:
        (sc.objdict[str, TransitionVariable]):
            Dictionary mapping names to `TransitionVariable` objects.
    """

    return sc.objdict()

find_name_by_compartment(target_compartment: Compartment) -> str

Given Compartment, returns name of that Compartment.

Parameters:

Name	Type	Description	Default
target_compartment	Compartment	Compartment object with a name to look up	required

Returns:

Type	Description
str	Compartment name, given by the key to look it up in the SubpopModel's compartments objdict

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def find_name_by_compartment(self,
                             target_compartment: Compartment) -> str:
    """
    Given `Compartment`, returns name of that `Compartment`.

    Args:
        target_compartment (Compartment):
            Compartment object with a name to look up

    Returns:
        (str):
            Compartment name, given by the key to look
            it up in the `SubpopModel`'s compartments objdict
    """

    for name, compartment in self.compartments.items():
        if compartment == target_compartment:
            return name

get_start_real_date()

Fetches start_real_date from simulation_settings -- converts to proper datetime.date format if originally given as string.

Returns:

Name	Type	Description
start_real_date	date	real-world date that corresponds to start of simulation.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def get_start_real_date(self):
    """
    Fetches `start_real_date` from `simulation_settings` -- converts to
        proper datetime.date format if originally given as
        string.

    Returns:
        start_real_date (datetime.date):
            real-world date that corresponds to start of
            simulation.
    """

    start_real_date = self.simulation_settings.start_real_date

    if not isinstance(start_real_date, datetime.date):
        try:
            start_real_date = \
                datetime.datetime.strptime(start_real_date, "%Y-%m-%d").date()
        except ValueError:
            print("Error: The date format should be YYYY-MM-DD.")

    return start_real_date

increment_simulation_day() -> None

Move day counters to next simulation day, both for integer simulation day and real date.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def increment_simulation_day(self) -> None:
    """
    Move day counters to next simulation day, both
        for integer simulation day and real date.
    """

    self.current_simulation_day += 1
    self.current_real_date += datetime.timedelta(days=1)

modify_random_seed(new_seed_number) -> None

Modifies model's RNG attribute in-place to new generator seeded at new_seed_number.

Parameters:

Name	Type	Description	Default
new_seed_number	int	used to re-seed model's random number generator.	required

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def modify_random_seed(self, new_seed_number) -> None:
    """
    Modifies model's `RNG` attribute in-place to new generator
    seeded at `new_seed_number`.

    Args:
        new_seed_number (int):
            used to re-seed model's random number generator.
    """

    self._bit_generator = np.random.MT19937(seed=new_seed_number)
    self.RNG = np.random.Generator(self._bit_generator)

modify_simulation_settings(updates_dict: dict)

This method lets users safely modify simulation settings; if this subpop model is associated with a metapop model, the same updates are applied to all subpop models on the metapop model. See also modify_simulation_settings method on MetapopModel.

Parameters:

Name	Type	Description	Default
updates_dict	dict	Dictionary specifying values to update in a SimulationSettings instance -- keys must match the field names of SimulationSettings.	required

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def modify_simulation_settings(self,
                               updates_dict: dict):
    """
    This method lets users safely modify simulation settings;
    if this subpop model is associated with a metapop model,
    the same updates are applied to all subpop models on the
    metapop model. See also `modify_simulation_settings` method on
    `MetapopModel`.

    Parameters:
        updates_dict (dict):
            Dictionary specifying values to update in a
            `SimulationSettings` instance -- keys must match the
            field names of `SimulationSettings`.
    """

    self.simulation_settings = \
        updated_dataclass(self.simulation_settings, updates_dict)

prepare_daily_state() -> None

At beginning of each day, update current value of interaction terms, schedules, dynamic values -- note that these are only updated once a day, not for every discretized timestep.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def prepare_daily_state(self) -> None:
    """
    At beginning of each day, update current value of
    interaction terms, schedules, dynamic values --
    note that these are only updated once a day, not
    for every discretized timestep.
    """

    subpop_state = self.state
    subpop_params = self.params
    current_real_date = self.current_real_date

    # Important note: this order of updating is important,
    #   because schedules do not depend on other state variables,
    #   but dynamic vals may depend on schedules
    # Interaction terms may depend on both schedules
    #   and dynamic vals.

    schedules = self.schedules
    dynamic_vals = self.dynamic_vals

    # Update schedules for current day
    for schedule in schedules.values():
        schedule.update_current_val(subpop_params,
                                    current_real_date)

    self.state.sync_to_current_vals(schedules)

    # Update dynamic values for current day
    for dval in dynamic_vals.values():
        if dval.is_enabled:
            dval.update_current_val(subpop_state, subpop_params)

    self.state.sync_to_current_vals(dynamic_vals)

reset_simulation() -> None

Reset simulation in-place. Subsequent method calls of simulate_until_day start from day 0, with original day 0 state.

Returns current_simulation_day to 0. Restores state values to initial values. Clears history on model's state variables. Resets transition variables' current_val attribute to 0.

WARNING

DOES NOT RESET THE MODEL'S RANDOM NUMBER GENERATOR TO ITS INITIAL STARTING SEED. RANDOM NUMBER GENERATOR WILL CONTINUE WHERE IT LEFT OFF.

Use method modify_random_seed to reset model's RNG to its initial starting seed.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def reset_simulation(self) -> None:
    """
    Reset simulation in-place. Subsequent method calls of
    `simulate_until_day` start from day 0, with original
    day 0 state.

    Returns `current_simulation_day` to 0.
    Restores state values to initial values.
    Clears history on model's state variables.
    Resets transition variables' `current_val` attribute to 0.

    WARNING:
        DOES NOT RESET THE MODEL'S RANDOM NUMBER GENERATOR TO
        ITS INITIAL STARTING SEED. RANDOM NUMBER GENERATOR WILL CONTINUE
        WHERE IT LEFT OFF.

    Use method `modify_random_seed` to reset model's `RNG` to its
    initial starting seed.
    """

    self.current_simulation_day = 0
    self.current_real_date = self.start_real_date

    # AGAIN, MUST BE CAREFUL ABOUT MUTABLE NUMPY ARRAYS -- MUST USE DEEP COPY
    for svar in self.all_state_variables.values():
        setattr(svar, "current_val", copy.deepcopy(svar.init_val))

    self.state.sync_to_current_vals(self.all_state_variables)

    for svar in self.all_state_variables.values():
        svar.reset()

    for tvar in self.transition_variables.values():
        tvar.reset()

    for tvargroup in self.transition_variable_groups.values():
        tvargroup.current_vals_list = []

sample_transitions() -> None

For each transition variable, sample a random realization using its current rate. Handle jointly distributed transition variables first (using TransitionVariableGroup logic), then handle marginally distributed transition variables. Use SubpopModel's RNG to generate random variables.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def sample_transitions(self) -> None:
    """
    For each transition variable, sample a random realization
        using its current rate. Handle jointly distributed transition
        variables first (using `TransitionVariableGroup` logic), then
        handle marginally distributed transition variables.
        Use `SubpopModel`'s `RNG` to generate random variables.
    """

    RNG = self.RNG
    timesteps_per_day = self.simulation_settings.timesteps_per_day
    transition_variables_to_save = self.simulation_settings.transition_variables_to_save

    # Obtain transition variable realizations for jointly distributed transition variables
    #   (i.e. when there are multiple transition variable outflows from an epi compartment)
    for tvargroup in self.transition_variable_groups.values():
        tvargroup.current_vals_list = tvargroup.get_joint_realization(RNG,
                                                                      timesteps_per_day)
        tvargroup.update_transition_variable_realizations()

    # Obtain transition variable realizations for marginally distributed transition variables
    #   (i.e. when there is only one transition variable outflow from an epi compartment)
    # If transition variable is jointly distributed, then its realization has already
    #   been computed by its transition variable group container previously,
    #   so skip the marginal computation
    for tvar in self.transition_variables.values():
        if not tvar.is_jointly_distributed:
            tvar.current_val = tvar.get_realization(RNG, timesteps_per_day)

    for name in transition_variables_to_save:
        self.transition_variables[name].save_history()

save_daily_history() -> None

Update history at end of each day, not at end of every discretization timestep, to be efficient. Update history of state variables other than Schedule instances -- schedules do not have history.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def save_daily_history(self) -> None:
    """
    Update history at end of each day, not at end of every
       discretization timestep, to be efficient.
    Update history of state variables other than `Schedule`
       instances -- schedules do not have history.
    """
    for svar in self.compartments.values() + \
                self.epi_metrics.values() + \
                self.dynamic_vals.values():
        svar.save_history()

simulate_until_day(simulation_end_day: int) -> None

Advance simulation model time until simulation_end_day.

Advance time by iterating through simulation days, which are simulated by iterating through discretized timesteps.

Save daily simulation data as history on each Compartment instance.

Parameters:

Name	Type	Description	Default
simulation_end_day	positive int	stop simulation at simulation_end_day (i.e. exclusive, simulate up to but not including simulation_end_day).	required

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def simulate_until_day(self,
                       simulation_end_day: int) -> None:
    """
    Advance simulation model time until `simulation_end_day`.

    Advance time by iterating through simulation days,
    which are simulated by iterating through discretized
    timesteps.

    Save daily simulation data as history on each `Compartment`
    instance.

    Args:
        simulation_end_day (positive int):
            stop simulation at `simulation_end_day` (i.e. exclusive,
            simulate up to but not including `simulation_end_day`).
    """

    if self.current_simulation_day > simulation_end_day:
        raise SubpopModelError(f"Current day counter ({self.current_simulation_day}) "
                               f"exceeds last simulation day ({simulation_end_day}).")

    save_daily_history = self.simulation_settings.save_daily_history
    timesteps_per_day = self.simulation_settings.timesteps_per_day

    # Adding this in case the user manually changes the initial
    #   value or current value of any state variable --
    #   otherwise, the state will not get updated
    self.state.sync_to_current_vals(self.all_state_variables)

    # simulation_end_day is exclusive endpoint
    while self.current_simulation_day < simulation_end_day:

        self.prepare_daily_state()

        self._simulate_timesteps(timesteps_per_day)

        if save_daily_history:
            self.save_daily_history()

        self.increment_simulation_day()

update_compartments() -> None

Update current value of each Compartment, by looping through all TransitionVariable instances and subtracting/adding their current values from origin/destination compartments respectively.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def update_compartments(self) -> None:
    """
    Update current value of each `Compartment`, by
        looping through all `TransitionVariable` instances
        and subtracting/adding their current values
        from origin/destination compartments respectively.
    """

    for tvar in self.transition_variables.values():
        tvar.update_origin_outflow()
        tvar.update_destination_inflow()

    for compartment in self.compartments.values():
        compartment.update_current_val()

        # By construction (using binomial/multinomial with or without taylor expansion),
        #   more individuals cannot leave the compartment than are in the compartment
        # However, for Poisson any for ANY deterministic version, it is possible
        #   to have more individuals leaving the compartment than are in the compartment,
        #   and hence negative-valued compartments
        # We use this function to fix this, and also use a differentiable torch
        #   function to be consistent with the torch implementation (this still
        #   allows us to take derivatives in the torch implementation)
        # The syntax is janky here -- we want everything as an array, but
        #   we need to pass a tensor to the torch functional
        if "deterministic" in self.simulation_settings.transition_type:
            compartment.current_val = \
                    np.array(torch.nn.functional.softplus(torch.tensor(compartment.current_val)))

        # After updating the compartment's current value,
        #   reset its inflow and outflow attributes, to
        #   prepare for the next iteration.
        compartment.reset_inflow()
        compartment.reset_outflow()

update_epi_metrics() -> None

Update current value attribute on each associated EpiMetric instance.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def update_epi_metrics(self) -> None:
    """
    Update current value attribute on each associated
        `EpiMetric` instance.
    """

    state = self.state
    params = self.params
    timesteps_per_day = self.simulation_settings.timesteps_per_day

    for metric in self.epi_metrics.values():
        metric.change_in_current_val = \
            metric.get_change_in_current_val(state,
                                             params,
                                             timesteps_per_day)
        metric.update_current_val()

update_transition_rates() -> None

Compute current transition rates for each transition variable, and store this updated value on each variable's current_rate attribute.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def update_transition_rates(self) -> None:
    """
    Compute current transition rates for each transition variable,
        and store this updated value on each variable's
        current_rate attribute.
    """

    state = self.state
    params = self.params

    for tvar in self.transition_variables.values():
        tvar.current_rate = tvar.get_current_rate(state, params)

SubpopModelError

Bases: Exception

Custom exceptions for subpopulation simulation model errors.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

class SubpopModelError(Exception):
    """Custom exceptions for subpopulation simulation model errors."""
    pass

SubpopParams dataclass

Bases: ABC

Data container for pre-specified and fixed epidemiological parameters in model.

Assume that SubpopParams fields are constant or piecewise constant throughout the simulation. For variables that are more complicated and time-dependent, use an EpiMetric instead.

Source code in CLT_BaseModel/clt_toolkit/base_data_structures.py

@dataclass(frozen=True)
class SubpopParams(ABC):
    """
    Data container for pre-specified and fixed epidemiological
    parameters in model.

    Assume that `SubpopParams` fields are constant or piecewise
    constant throughout the simulation. For variables that
    are more complicated and time-dependent, use an `EpiMetric`
    instead.
    """

    pass

SubpopState dataclass

Bases: ABC

Holds current values of SubpopModel's simulation state.

Source code in CLT_BaseModel/clt_toolkit/base_data_structures.py

@dataclass
class SubpopState(ABC):
    """
    Holds current values of `SubpopModel`'s simulation state.
    """

    def sync_to_current_vals(self, lookup_dict: dict):
        """
        Updates `SubpopState`'s attributes according to
        data in `lookup_dict.` Keys of `lookup_dict` must match
        names of attributes of `SubpopState` instance.
        """

        for name, item in lookup_dict.items():
            setattr(self, name, item.current_val)

sync_to_current_vals(lookup_dict: dict)

Updates SubpopState's attributes according to data in lookup_dict. Keys of lookup_dict must match names of attributes of SubpopState instance.

Source code in CLT_BaseModel/clt_toolkit/base_data_structures.py

def sync_to_current_vals(self, lookup_dict: dict):
    """
    Updates `SubpopState`'s attributes according to
    data in `lookup_dict.` Keys of `lookup_dict` must match
    names of attributes of `SubpopState` instance.
    """

    for name, item in lookup_dict.items():
        setattr(self, name, item.current_val)

TransitionTypes

Bases: str, Enum

Defines available options for transition_type in TransitionVariable.

Source code in CLT_BaseModel/clt_toolkit/base_data_structures.py

class TransitionTypes(str, Enum):
    """
    Defines available options for `transition_type` in `TransitionVariable`.
    """
    BINOM = "binom"
    BINOM_DETERMINISTIC = "binom_deterministic"
    BINOM_DETERMINISTIC_NO_ROUND = "binom_deterministic_no_round"
    BINOM_TAYLOR_APPROX = "binom_taylor_approx"
    BINOM_TAYLOR_APPROX_DETERMINISTIC = "binom_taylor_approx_deterministic"
    POISSON = "poisson"
    POISSON_DETERMINISTIC = "poisson_deterministic"

TransitionVariable

Bases: ABC

Abstract base class for transition variables in epidemiological model.

For example, in an S-I-R model, the new number infected every iteration (the number going from S to I) in an iteration is modeled as a TransitionVariable subclass, with a concrete implementation of the abstract method get_current_rate.

When an instance is initialized, its get_realization attribute is dynamically assigned, just like in the case of TransitionVariableGroup instantiation.

Dimensions

A (int): Number of age groups. R (int): Number of risk groups.

Attributes:

Name	Type	Description
_transition_type	str	only values defined in TransitionTypes are valid, specifying probability distribution of transitions between compartments.
get_current_rate	function	provides specific implementation for computing current rate as a function of current subpopulation simulation state and epidemiological parameters.
current_rate	np.ndarray of shape (A, R	holds output from get_current_rate method -- used to generate random variable realizations for transitions between compartments.
current_val	np.ndarray of shape (A, R	holds realization of random variable parameterized by current_rate.
history_vals_list	list[ndarray]	each element is the same size of current_val, holds history of transition variable realizations for age-risk groups -- element t corresponds to previous current_val value at end of simulation day t.

See __init__ docstring for other attributes.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

class TransitionVariable(ABC):
    """
    Abstract base class for transition variables in
    epidemiological model.

    For example, in an S-I-R model, the new number infected
    every iteration (the number going from S to I) in an iteration
    is modeled as a `TransitionVariable` subclass, with a concrete
    implementation of the abstract method `get_current_rate`.

    When an instance is initialized, its `get_realization` attribute
    is dynamically assigned, just like in the case of
    `TransitionVariableGroup` instantiation.

    Dimensions:
        A (int):
            Number of age groups.
        R (int):
            Number of risk groups.

    Attributes:
        _transition_type (str):
            only values defined in `TransitionTypes` are valid, specifying
            probability distribution of transitions between compartments.
        get_current_rate (function):
            provides specific implementation for computing current rate
            as a function of current subpopulation simulation state and
            epidemiological parameters.
        current_rate (np.ndarray of shape (A, R)):
            holds output from `get_current_rate` method -- used to generate
            random variable realizations for transitions between compartments.
        current_val (np.ndarray of shape (A, R)):
            holds realization of random variable parameterized by
            `current_rate`.
        history_vals_list (list[np.ndarray]):
            each element is the same size of `current_val`, holds
            history of transition variable realizations for age-risk
            groups -- element t corresponds to previous `current_val`
            value at end of simulation day t.

    See `__init__` docstring for other attributes.
    """

    def __init__(self,
                 origin: Compartment,
                 destination: Compartment,
                 transition_type: TransitionTypes,
                 is_jointly_distributed: str = False):
        """
        Parameters:
            origin (Compartment):
                `Compartment` from which `TransitionVariable` exits.
            destination (Compartment):
                `Compartment` that the `TransitionVariable` enters.
            transition_type (TransitionTypes):
                Specifies probability distribution of transitions between compartments.
            is_jointly_distributed (bool):
                Indicates if transition quantity must be jointly computed
                (i.e. if there are multiple outflows from the origin compartment).
        """

        self.origin = origin
        self.destination = destination

        # Also see __init__ method in TransitionVariableGroup class.
        #   The structure is similar.
        self._transition_type = transition_type
        self._is_jointly_distributed = is_jointly_distributed

        # Assigns appropriate realization method based on transition type.
        # If jointly distributed, no single realization function applies.
        if is_jointly_distributed:
            self.get_realization = None
        else:
            self.get_realization = getattr(self, "get_" + transition_type + "_realization")

        self.current_rate = None
        self.current_val = None

        self.history_vals_list = []

    @property
    def transition_type(self) -> TransitionTypes:
        return self._transition_type

    @property
    def is_jointly_distributed(self) -> bool:
        return self._is_jointly_distributed

    @abstractmethod
    def get_current_rate(self,
                         state: SubpopState,
                         params: SubpopParams) -> np.ndarray:
        """
        Computes and returns current rate of transition variable,
        based on current state of the simulation and epidemiological parameters.

        Args:
            state (SubpopState):
                Holds subpopulation simulation state
                (current values of `StateVariable` instances).
            params (SubpopParams):
                Holds values of epidemiological parameters.

        Returns:
            np.ndarray:
                Holds age-risk transition rate.
        """
        pass

    def update_origin_outflow(self) -> None:
        """
        Adds current realization of `TransitionVariable` to
            its origin `Compartment`'s current_outflow.
            Used to compute total number leaving that
            origin `Compartment`.
        """

        self.origin.current_outflow = self.origin.current_outflow + self.current_val

    def update_destination_inflow(self) -> None:
        """
        Adds current realization of `TransitionVariable` to
            its destination `Compartment`'s `current_inflow`.
            Used to compute total number leaving that
            destination `Compartment`.
        """

        self.destination.current_inflow = self.destination.current_inflow + self.current_val

    def save_history(self) -> None:
        """
        Saves current value to history by appending `current_val`
        attribute to `history_vals_list` in-place..

        Deep copying is CRUCIAL because `current_val` is a mutable
        np.ndarray -- without deep copying, `history_vals_list` would
        have the same value for all elements.
        """
        self.history_vals_list.append(copy.deepcopy(self.current_val))

    def reset(self) -> None:
        """
        Resets `history_vals_list` attribute to empty list.
        """

        self.current_rate = None
        self.current_val = None
        self.history_vals_list = []

    def get_realization(self,
                        RNG: np.random.Generator,
                        num_timesteps: int) -> np.ndarray:
        """
        Generate a realization of the transition process.

        This method is dynamically assigned to the appropriate transition-specific
        function (e.g., `get_binom_realization`) depending on the transition type.
        Provides common interface so realizations can always be obtained via
        ``get_realization``.

        Parameters:
            RNG (np.random.Generator object):
                 Used to generate stochastic transitions in the model and control
                 reproducibility. If deterministic transitions are used, the
                 RNG is passed for a consistent function interface but the RNG
                 is not used.
            num_timesteps (int):
                Number of timesteps per day -- used to determine time interval
                length for discretization.

        Returns:
            (np.ndarray of shape (A, R)):
                Number of transitions for age-risk groups.
        """

        pass

    def get_binom_realization(self,
                              RNG: np.random.Generator,
                              num_timesteps: int) -> np.ndarray:
        """
        Uses `RNG` to generate binomial random variable with
        number of trials equal to population count in the
        origin `Compartment` and probability computed from
        a function of the `TransitionVariable`'s current rate
        -- see `approx_binom_probability_from_rate` function
        for details.

        See `get_realization` for parameters.

        Returns:
            (np.ndarray of shape (A, R))
                Element-wise Binomial distributed transitions for each
                age-risk group, with the probability parameter generated
                using a conversion from rates to probabilities.
        """

        return RNG.binomial(n=np.asarray(self.base_count, dtype=int),
                            p=approx_binom_probability_from_rate(self.current_rate, 1.0 / num_timesteps))

    def get_binom_taylor_approx_realization(self,
                                            RNG: np.random.Generator,
                                            num_timesteps: int) -> np.ndarray:
        """
        Uses `RNG` to generate binomial random variable with
            number of trials equal to population count in the
            origin `Compartment` and probability equal to
            the `TransitionVariable`'s `current_rate` / `num_timesteps`.

        See `get_realization` for parameters.

        Returns:
            (np.ndarray of shape (A, R))
                Element-wise Binomial distributed transitions for each
                age-risk group, with the probability parameter generated
                using a Taylor approximation.
        """
        return RNG.binomial(n=np.asarray(self.base_count, dtype=int),
                            p=self.current_rate * (1.0 / num_timesteps))

    def get_poisson_realization(self,
                                RNG: np.random.Generator,
                                num_timesteps: int) -> np.ndarray:
        """
        Generates realizations from a Poisson distribution.

        The rate is computed element-wise from each age-risk group as:
        (origin compartment population count x `current_rate` / `num_timesteps`)

        See `get_realization` for parameters.

        Returns:
            (np.ndarray of shape (A, R))
                Poisson-distributed integers representing number
                of individuals transitioning in each age-risk group.
        """
        return RNG.poisson(self.base_count * self.current_rate / float(num_timesteps))

    def get_binom_deterministic_realization(self,
                                            RNG: np.random.Generator,
                                            num_timesteps: int) -> np.ndarray:
        """
        Deterministically returns mean of binomial distribution
        (number of trials x probability), where number of trials
        equals population count in the origin `Compartment` and
        probability is computed from a function of the `TransitionVariable`'s
        current rate -- see the `approx_binom_probability_from_rate`
        function for details.

        See `get_realization` for parameters. The `RNG` parameter is not used
        and is only included to maintain a consistent interface.

        Returns:
            (np.ndarray of shape (A, R))
                Number of individuals transitioning compartments in each age-risk group.
        """

        return np.asarray(self.base_count *
                          approx_binom_probability_from_rate(self.current_rate, 1.0 / num_timesteps),
                          dtype=int)

    def get_binom_deterministic_no_round_realization(self,
                                                     RNG: np.random.Generator,
                                                     num_timesteps: int) -> np.ndarray:
        """
        The same as `get_binom_deterministic_realization` except no rounding --
        so the populations can be non-integer. This is used to test the torch
        implementation (because that implementation does not round either).

        See `get_realization` for parameters. The `RNG` parameter is not used
        and is only included to maintain a consistent interface.

        Returns:
            (np.ndarray of shape (A, R))
                (Non-integer) "number of individuals" transitioning compartments in
                each age-risk group.
        """

        return np.asarray(self.base_count *
                          approx_binom_probability_from_rate(self.current_rate, 1.0 / num_timesteps))

    def get_binom_taylor_approx_deterministic_realization(self,
                                                          RNG: np.random.Generator,
                                                          num_timesteps: int) -> np.ndarray:
        """
        Deterministically returns mean of binomial distribution
        (number of trials x probability), where number of trials
        equals population count in the origin `Compartment` and
        probability equals the `TransitionVariable`'s `current_rate` /
        `num_timesteps`.

        See `get_realization` for parameters. The `RNG` parameter is not used
        and is only included to maintain a consistent interface.

        Returns:
            (np.ndarray of shape (A, R))
                Number of individuals transitioning compartments in each age-risk group.
        """

        return np.asarray(self.base_count * self.current_rate / num_timesteps, dtype=int)

    def get_poisson_deterministic_realization(self,
                                              RNG: np.random.Generator,
                                              num_timesteps: int) -> np.ndarray:
        """
        Deterministically returns mean of Poisson distribution,
        given by (population count in the origin `Compartment` x
        `TransitionVariable`'s `current_rate` / `num_timesteps`).

        See `get_realization` for parameters. The `RNG` parameter is not used
        and is only included to maintain a consistent interface.

        Returns:
            (np.ndarray of shape (A, R))
                Number of individuals transitioning compartments in each age-risk group.
        """

        return np.asarray(self.base_count * self.current_rate / num_timesteps, dtype=int)

    @property
    def base_count(self) -> np.ndarray:
        return self.origin.current_val

__init__(origin: Compartment, destination: Compartment, transition_type: TransitionTypes, is_jointly_distributed: str = False)

Parameters:

Name	Type	Description	Default
origin	Compartment	Compartment from which TransitionVariable exits.	required
destination	Compartment	Compartment that the TransitionVariable enters.	required
transition_type	TransitionTypes	Specifies probability distribution of transitions between compartments.	required
is_jointly_distributed	bool	Indicates if transition quantity must be jointly computed (i.e. if there are multiple outflows from the origin compartment).	False

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def __init__(self,
             origin: Compartment,
             destination: Compartment,
             transition_type: TransitionTypes,
             is_jointly_distributed: str = False):
    """
    Parameters:
        origin (Compartment):
            `Compartment` from which `TransitionVariable` exits.
        destination (Compartment):
            `Compartment` that the `TransitionVariable` enters.
        transition_type (TransitionTypes):
            Specifies probability distribution of transitions between compartments.
        is_jointly_distributed (bool):
            Indicates if transition quantity must be jointly computed
            (i.e. if there are multiple outflows from the origin compartment).
    """

    self.origin = origin
    self.destination = destination

    # Also see __init__ method in TransitionVariableGroup class.
    #   The structure is similar.
    self._transition_type = transition_type
    self._is_jointly_distributed = is_jointly_distributed

    # Assigns appropriate realization method based on transition type.
    # If jointly distributed, no single realization function applies.
    if is_jointly_distributed:
        self.get_realization = None
    else:
        self.get_realization = getattr(self, "get_" + transition_type + "_realization")

    self.current_rate = None
    self.current_val = None

    self.history_vals_list = []

get_binom_deterministic_no_round_realization(RNG: np.random.Generator, num_timesteps: int) -> np.ndarray

The same as get_binom_deterministic_realization except no rounding -- so the populations can be non-integer. This is used to test the torch implementation (because that implementation does not round either).

See get_realization for parameters. The RNG parameter is not used and is only included to maintain a consistent interface.

Returns:

Type	Description
ndarray	(np.ndarray of shape (A, R)) (Non-integer) "number of individuals" transitioning compartments in each age-risk group.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def get_binom_deterministic_no_round_realization(self,
                                                 RNG: np.random.Generator,
                                                 num_timesteps: int) -> np.ndarray:
    """
    The same as `get_binom_deterministic_realization` except no rounding --
    so the populations can be non-integer. This is used to test the torch
    implementation (because that implementation does not round either).

    See `get_realization` for parameters. The `RNG` parameter is not used
    and is only included to maintain a consistent interface.

    Returns:
        (np.ndarray of shape (A, R))
            (Non-integer) "number of individuals" transitioning compartments in
            each age-risk group.
    """

    return np.asarray(self.base_count *
                      approx_binom_probability_from_rate(self.current_rate, 1.0 / num_timesteps))

get_binom_deterministic_realization(RNG: np.random.Generator, num_timesteps: int) -> np.ndarray

Deterministically returns mean of binomial distribution (number of trials x probability), where number of trials equals population count in the origin Compartment and probability is computed from a function of the TransitionVariable's current rate -- see the approx_binom_probability_from_rate function for details.

See get_realization for parameters. The RNG parameter is not used and is only included to maintain a consistent interface.

Returns:

Type	Description
ndarray	(np.ndarray of shape (A, R)) Number of individuals transitioning compartments in each age-risk group.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def get_binom_deterministic_realization(self,
                                        RNG: np.random.Generator,
                                        num_timesteps: int) -> np.ndarray:
    """
    Deterministically returns mean of binomial distribution
    (number of trials x probability), where number of trials
    equals population count in the origin `Compartment` and
    probability is computed from a function of the `TransitionVariable`'s
    current rate -- see the `approx_binom_probability_from_rate`
    function for details.

    See `get_realization` for parameters. The `RNG` parameter is not used
    and is only included to maintain a consistent interface.

    Returns:
        (np.ndarray of shape (A, R))
            Number of individuals transitioning compartments in each age-risk group.
    """

    return np.asarray(self.base_count *
                      approx_binom_probability_from_rate(self.current_rate, 1.0 / num_timesteps),
                      dtype=int)

get_binom_realization(RNG: np.random.Generator, num_timesteps: int) -> np.ndarray

Uses RNG to generate binomial random variable with number of trials equal to population count in the origin Compartment and probability computed from a function of the TransitionVariable's current rate -- see approx_binom_probability_from_rate function for details.

See get_realization for parameters.

Returns:

Type	Description
ndarray	(np.ndarray of shape (A, R)) Element-wise Binomial distributed transitions for each age-risk group, with the probability parameter generated using a conversion from rates to probabilities.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def get_binom_realization(self,
                          RNG: np.random.Generator,
                          num_timesteps: int) -> np.ndarray:
    """
    Uses `RNG` to generate binomial random variable with
    number of trials equal to population count in the
    origin `Compartment` and probability computed from
    a function of the `TransitionVariable`'s current rate
    -- see `approx_binom_probability_from_rate` function
    for details.

    See `get_realization` for parameters.

    Returns:
        (np.ndarray of shape (A, R))
            Element-wise Binomial distributed transitions for each
            age-risk group, with the probability parameter generated
            using a conversion from rates to probabilities.
    """

    return RNG.binomial(n=np.asarray(self.base_count, dtype=int),
                        p=approx_binom_probability_from_rate(self.current_rate, 1.0 / num_timesteps))

get_binom_taylor_approx_deterministic_realization(RNG: np.random.Generator, num_timesteps: int) -> np.ndarray

Deterministically returns mean of binomial distribution (number of trials x probability), where number of trials equals population count in the origin Compartment and probability equals the TransitionVariable's current_rate / num_timesteps.

See get_realization for parameters. The RNG parameter is not used and is only included to maintain a consistent interface.

Returns:

Type	Description
ndarray	(np.ndarray of shape (A, R)) Number of individuals transitioning compartments in each age-risk group.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def get_binom_taylor_approx_deterministic_realization(self,
                                                      RNG: np.random.Generator,
                                                      num_timesteps: int) -> np.ndarray:
    """
    Deterministically returns mean of binomial distribution
    (number of trials x probability), where number of trials
    equals population count in the origin `Compartment` and
    probability equals the `TransitionVariable`'s `current_rate` /
    `num_timesteps`.

    See `get_realization` for parameters. The `RNG` parameter is not used
    and is only included to maintain a consistent interface.

    Returns:
        (np.ndarray of shape (A, R))
            Number of individuals transitioning compartments in each age-risk group.
    """

    return np.asarray(self.base_count * self.current_rate / num_timesteps, dtype=int)

get_binom_taylor_approx_realization(RNG: np.random.Generator, num_timesteps: int) -> np.ndarray

Uses RNG to generate binomial random variable with number of trials equal to population count in the origin Compartment and probability equal to the TransitionVariable's current_rate / num_timesteps.

See get_realization for parameters.

Returns:

Type	Description
ndarray	(np.ndarray of shape (A, R)) Element-wise Binomial distributed transitions for each age-risk group, with the probability parameter generated using a Taylor approximation.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def get_binom_taylor_approx_realization(self,
                                        RNG: np.random.Generator,
                                        num_timesteps: int) -> np.ndarray:
    """
    Uses `RNG` to generate binomial random variable with
        number of trials equal to population count in the
        origin `Compartment` and probability equal to
        the `TransitionVariable`'s `current_rate` / `num_timesteps`.

    See `get_realization` for parameters.

    Returns:
        (np.ndarray of shape (A, R))
            Element-wise Binomial distributed transitions for each
            age-risk group, with the probability parameter generated
            using a Taylor approximation.
    """
    return RNG.binomial(n=np.asarray(self.base_count, dtype=int),
                        p=self.current_rate * (1.0 / num_timesteps))

get_current_rate(state: SubpopState, params: SubpopParams) -> np.ndarray abstractmethod

Computes and returns current rate of transition variable, based on current state of the simulation and epidemiological parameters.

Parameters:

Name	Type	Description	Default
state	SubpopState	Holds subpopulation simulation state (current values of StateVariable instances).	required
params	SubpopParams	Holds values of epidemiological parameters.	required

Returns:

Type	Description
ndarray	np.ndarray: Holds age-risk transition rate.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

@abstractmethod
def get_current_rate(self,
                     state: SubpopState,
                     params: SubpopParams) -> np.ndarray:
    """
    Computes and returns current rate of transition variable,
    based on current state of the simulation and epidemiological parameters.

    Args:
        state (SubpopState):
            Holds subpopulation simulation state
            (current values of `StateVariable` instances).
        params (SubpopParams):
            Holds values of epidemiological parameters.

    Returns:
        np.ndarray:
            Holds age-risk transition rate.
    """
    pass

get_poisson_deterministic_realization(RNG: np.random.Generator, num_timesteps: int) -> np.ndarray

Deterministically returns mean of Poisson distribution, given by (population count in the origin Compartment x TransitionVariable's current_rate / num_timesteps).

See get_realization for parameters. The RNG parameter is not used and is only included to maintain a consistent interface.

Returns:

Type	Description
ndarray	(np.ndarray of shape (A, R)) Number of individuals transitioning compartments in each age-risk group.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def get_poisson_deterministic_realization(self,
                                          RNG: np.random.Generator,
                                          num_timesteps: int) -> np.ndarray:
    """
    Deterministically returns mean of Poisson distribution,
    given by (population count in the origin `Compartment` x
    `TransitionVariable`'s `current_rate` / `num_timesteps`).

    See `get_realization` for parameters. The `RNG` parameter is not used
    and is only included to maintain a consistent interface.

    Returns:
        (np.ndarray of shape (A, R))
            Number of individuals transitioning compartments in each age-risk group.
    """

    return np.asarray(self.base_count * self.current_rate / num_timesteps, dtype=int)

get_poisson_realization(RNG: np.random.Generator, num_timesteps: int) -> np.ndarray

Generates realizations from a Poisson distribution.

The rate is computed element-wise from each age-risk group as: (origin compartment population count x current_rate / num_timesteps)

See get_realization for parameters.

Returns:

Type	Description
ndarray	(np.ndarray of shape (A, R)) Poisson-distributed integers representing number of individuals transitioning in each age-risk group.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def get_poisson_realization(self,
                            RNG: np.random.Generator,
                            num_timesteps: int) -> np.ndarray:
    """
    Generates realizations from a Poisson distribution.

    The rate is computed element-wise from each age-risk group as:
    (origin compartment population count x `current_rate` / `num_timesteps`)

    See `get_realization` for parameters.

    Returns:
        (np.ndarray of shape (A, R))
            Poisson-distributed integers representing number
            of individuals transitioning in each age-risk group.
    """
    return RNG.poisson(self.base_count * self.current_rate / float(num_timesteps))

get_realization(RNG: np.random.Generator, num_timesteps: int) -> np.ndarray

Generate a realization of the transition process.

This method is dynamically assigned to the appropriate transition-specific function (e.g., get_binom_realization) depending on the transition type. Provides common interface so realizations can always be obtained via get_realization.

Parameters:

Name	Type	Description	Default
RNG	np.random.Generator object	Used to generate stochastic transitions in the model and control reproducibility. If deterministic transitions are used, the RNG is passed for a consistent function interface but the RNG is not used.	required
num_timesteps	int	Number of timesteps per day -- used to determine time interval length for discretization.	required

Returns:

Type	Description
np.ndarray of shape (A, R)	Number of transitions for age-risk groups.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def get_realization(self,
                    RNG: np.random.Generator,
                    num_timesteps: int) -> np.ndarray:
    """
    Generate a realization of the transition process.

    This method is dynamically assigned to the appropriate transition-specific
    function (e.g., `get_binom_realization`) depending on the transition type.
    Provides common interface so realizations can always be obtained via
    ``get_realization``.

    Parameters:
        RNG (np.random.Generator object):
             Used to generate stochastic transitions in the model and control
             reproducibility. If deterministic transitions are used, the
             RNG is passed for a consistent function interface but the RNG
             is not used.
        num_timesteps (int):
            Number of timesteps per day -- used to determine time interval
            length for discretization.

    Returns:
        (np.ndarray of shape (A, R)):
            Number of transitions for age-risk groups.
    """

    pass

reset() -> None

Resets history_vals_list attribute to empty list.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def reset(self) -> None:
    """
    Resets `history_vals_list` attribute to empty list.
    """

    self.current_rate = None
    self.current_val = None
    self.history_vals_list = []

save_history() -> None

Saves current value to history by appending current_val attribute to history_vals_list in-place..

Deep copying is CRUCIAL because current_val is a mutable np.ndarray -- without deep copying, history_vals_list would have the same value for all elements.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def save_history(self) -> None:
    """
    Saves current value to history by appending `current_val`
    attribute to `history_vals_list` in-place..

    Deep copying is CRUCIAL because `current_val` is a mutable
    np.ndarray -- without deep copying, `history_vals_list` would
    have the same value for all elements.
    """
    self.history_vals_list.append(copy.deepcopy(self.current_val))

update_destination_inflow() -> None

Adds current realization of TransitionVariable to its destination Compartment's current_inflow. Used to compute total number leaving that destination Compartment.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def update_destination_inflow(self) -> None:
    """
    Adds current realization of `TransitionVariable` to
        its destination `Compartment`'s `current_inflow`.
        Used to compute total number leaving that
        destination `Compartment`.
    """

    self.destination.current_inflow = self.destination.current_inflow + self.current_val

update_origin_outflow() -> None

Adds current realization of TransitionVariable to its origin Compartment's current_outflow. Used to compute total number leaving that origin Compartment.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def update_origin_outflow(self) -> None:
    """
    Adds current realization of `TransitionVariable` to
        its origin `Compartment`'s current_outflow.
        Used to compute total number leaving that
        origin `Compartment`.
    """

    self.origin.current_outflow = self.origin.current_outflow + self.current_val

TransitionVariableGroup

Container for TransitionVariable objects to handle joint sampling, when there are multiple outflows from a single compartment.

For example, if all outflows of compartment H are: R and D, i.e. from the hospital, individuals either recover or die, a TransitionVariableGroup that holds both R and D handles the correct correlation structure between R and D.

When an instance is initialized, its get_joint_realization attribute is dynamically assigned to a method according to its transition_type attribute. This enables all instances to use the same method during simulation.

Dimensions

M (int): number of outgoing compartments from the origin compartment A (int): number of age groups R (int): number of risk groups

Attributes:

Name	Type	Description
origin	Compartment	Specifies origin of TransitionVariableGroup -- corresponding populations leave this compartment.
_transition_type	str	Only values defined in JointTransitionTypes are valid, specifies joint probability distribution of all outflows from origin.
transition_variables	list[`TransitionVariable`]	Specifying TransitionVariable instances that outflow from origin -- order does not matter.
get_joint_realization	function	Assigned at initialization, generates realizations according to probability distribution given by transition_type, returns np.ndarray of either shape (M, A, R) or ((M+1), A, R), where M is the length of transition_variables (i.e., number of outflows from origin), A is the number of age groups, R is number of risk groups.
current_vals_list	list	Used to store results from get_joint_realization -- has either M or M+1 arrays of shape (A, R).

See __init__ docstring for other attributes.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

class TransitionVariableGroup:
    """
    Container for `TransitionVariable` objects to handle joint sampling,
    when there are multiple outflows from a single compartment.

    For example, if all outflows of compartment `H` are: `R` and `D`,
    i.e. from the hospital, individuals either recover or die,
    a `TransitionVariableGroup` that holds both `R` and `D` handles
    the correct correlation structure between `R` and `D.`

    When an instance is initialized, its `get_joint_realization` attribute
    is dynamically assigned to a method according to its `transition_type`
    attribute. This enables all instances to use the same method during
    simulation.

    Dimensions:
        M (int):
            number of outgoing compartments from the origin compartment
        A (int):
            number of age groups
        R (int):
            number of risk groups

    Attributes:
        origin (Compartment):
            Specifies origin of `TransitionVariableGroup` --
            corresponding populations leave this compartment.
        _transition_type (str):
            Only values defined in `JointTransitionTypes` are valid,
            specifies joint probability distribution of all outflows
            from origin.
        transition_variables (list[`TransitionVariable`]):
            Specifying `TransitionVariable` instances that outflow from origin --
            order does not matter.
        get_joint_realization (function):
            Assigned at initialization, generates realizations according
            to probability distribution given by `transition_type`,
            returns np.ndarray of either shape (M, A, R) or ((M+1), A, R),
            where M is the length of `transition_variables` (i.e., number of
            outflows from origin), A is the number of age groups, R is number of
            risk groups.
        current_vals_list (list):
            Used to store results from `get_joint_realization` --
            has either M or M+1 arrays of shape (A, R).

    See `__init__` docstring for other attributes.
    """

    def __init__(self,
                 origin: Compartment,
                 transition_type: TransitionTypes,
                 transition_variables: list[TransitionVariable]):
        """
        Args:
            transition_type (TransitionTypes):
                Specifies probability distribution of transitions between compartments.

        See class docstring for other parameters.
        """

        self.origin = origin

        # Using a list is important here because we want to keep the order
        #   of transition variables -- this determines the index in the
        #   current rates array
        self.transition_variables = transition_variables

        # If marginal transition type is any kind of binomial transition,
        #   then its joint transition type is a multinomial counterpart
        # For example, if the marginal transition type is TransitionTypes.BINOM_DETERMINISTIC,
        #   then the joint transition type is JointTransitionTypes.MULTINOM_DETERMINISTIC
        transition_type = transition_type.replace("binom", "multinom")
        self._transition_type = transition_type

        # Dynamically assign a method to get_joint_realization attribute
        #   based on the value of transition_type
        # getattr fetches a method by name
        self.get_joint_realization = getattr(self, "get_" + transition_type + "_realization")

        self.current_vals_list = []

    @property
    def transition_type(self) -> JointTransitionTypes:
        return self._transition_type

    def get_total_rate(self) -> np.ndarray:
        """
        Return the age-risk-specific total transition rate,
        which is the sum of the current rate of each transition variable
        in this transition variable group.

        Used to properly scale multinomial probabilities vector so
        that elements sum to 1.

        Returns:
            (np.ndarray of shape (A, R))
                Array with values corresponding to sum of current rates of
                transition variables in transition variable group, where
                elements correspond to age-risk groups.
        """

        # axis 0: corresponds to outgoing transition variable
        # axis 1: corresponds to age groups
        # axis 2: corresponds to risk groups
        # --> summing over axis 0 gives the total rate for each age-risk group
        return np.sum(self.get_current_rates_array(), axis=0)

    def get_probabilities_array(self,
                                num_timesteps: int) -> list:
        """
        Returns an array of probabilities used for joint binomial
        (multinomial) transitions (`get_multinom_realization` method).

        Returns:
            (np.ndarray of shape (M+1, A, R)
                Contains positive floats <= 1, corresponding to probability
                of transitioning to a compartment for that outgoing compartment
                and age-risk group -- note the "+1" corresponds to the multinomial
                outcome of staying in the same compartment (we can think of as
                transitioning to the same compartment).
        """

        total_rate = self.get_total_rate()

        total_outgoing_probability = approx_binom_probability_from_rate(total_rate,
                                                                        1 / num_timesteps)

        # Create probabilities_list, where element i corresponds to the
        #   transition variable i's current rate divided by the total rate,
        #   multiplized by the total outgoing probability
        # This generates the probabilities array that parameterizes the
        #   multinomial distribution
        probabilities_list = []

        for transition_variable in self.transition_variables:
            probabilities_list.append((transition_variable.current_rate / total_rate) *
                                      total_outgoing_probability)

        # Append the probability that a person stays in the compartment
        probabilities_list.append(1 - total_outgoing_probability)

        return np.asarray(probabilities_list)

    def get_current_rates_array(self) -> np.ndarray:
        """
        Returns an array of current rates of transition variables in
        `transition_variables` -- ith element in array
        corresponds to current rate of ith transition variable.

        Returns:
            (np.ndarray of shape (M, A, R))
                array of positive floats corresponding to current rate
                element-wise for an outgoing compartment and age-risk group
        """

        current_rates_list = []
        for tvar in self.transition_variables:
            current_rates_list.append(tvar.current_rate)

        return np.asarray(current_rates_list)

    def get_joint_realization(self,
                              RNG: np.random.Generator,
                              num_timesteps: int) -> np.ndarray:
        """
        This function is dynamically assigned based on the `TransitionVariableGroup`'s
            `transition_type`. It is set to the appropriate distribution-specific method.

        See `get_realization` for parameters.
        """

        pass

    def get_multinom_realization(self,
                                 RNG: np.random.Generator,
                                 num_timesteps: int) -> np.ndarray:
        """
        Returns an array of transition realizations (number transitioning
        to outgoing compartments) sampled from multinomial distribution.

        See `get_realization` for parameters.

        Returns:
            (np.ndarray of shape (M + 1, A, R))
                contains positive floats with transition realizations
                for individuals going to compartment m in age-risk group (a, r) --
                note the "+1" corresponds to the multinomial outcome of staying
                in the same compartment (not transitioning to any outgoing
                epi compartment).
        """

        probabilities_array = self.get_probabilities_array(num_timesteps)

        num_outflows = len(self.transition_variables)

        num_age_groups, num_risk_groups = np.shape(self.origin.current_val)

        # We use num_outflows + 1 because for the multinomial distribution we explicitly model
        #   the number who stay/remain in the compartment
        realizations_array = np.zeros((num_outflows + 1, num_age_groups, num_risk_groups))

        for age_group in range(num_age_groups):
            for risk_group in range(num_risk_groups):
                realizations_array[:, age_group, risk_group] = RNG.multinomial(
                    np.asarray(self.origin.current_val[age_group, risk_group], dtype=int),
                    probabilities_array[:, age_group, risk_group])

        return realizations_array

    def get_multinom_taylor_approx_realization(self,
                                               RNG: np.random.Generator,
                                               num_timesteps: int) -> np.ndarray:
        """
        Returns an array of transition realizations (number transitioning
        to outgoing compartments) sampled from multinomial distribution
        using Taylor Series approximation for probability parameter.

        See `get_realization` for parameters.

        Returns:
            (np.ndarray of shape (M + 1, A, R))
                contains positive integers with transition realizations
                for individuals going to compartment m in age-risk group (a, r) --
                note the "+1" corresponds to the multinomial outcome of staying
                in the same compartment (not transitioning to any outgoing
                epi compartment).
        """

        num_outflows = len(self.transition_variables)

        current_rates_array = self.get_current_rates_array()

        total_rate = self.get_total_rate()

        # Multiply current rates array by length of time interval (1 / num_timesteps)
        # Also append additional value corresponding to probability of
        #   remaining in current epi compartment (not transitioning at all)
        # Note: "vstack" function here works better than append function because append
        #   automatically flattens the resulting array, resulting in dimension issues
        current_scaled_rates_array = np.vstack((current_rates_array / num_timesteps,
                                                np.expand_dims(1 - total_rate / num_timesteps, axis=0)))

        num_age_groups, num_risk_groups = np.shape(self.origin.current_val)

        # We use num_outflows + 1 because for the multinomial distribution we explicitly model
        #   the number who stay/remain in the compartment
        realizations_array = np.zeros((num_outflows + 1, num_age_groups, num_risk_groups))

        for age_group in range(num_age_groups):
            for risk_group in range(num_risk_groups):
                realizations_array[:, age_group, risk_group] = RNG.multinomial(
                    np.asarray(self.origin.current_val[age_group, risk_group], dtype=int),
                    current_scaled_rates_array[:, age_group, risk_group])

        return realizations_array

    def get_poisson_realization(self,
                                RNG: np.random.Generator,
                                num_timesteps: int) -> np.ndarray:
        """
        Returns an array of transition realizations (number transitioning
        to outgoing compartments) sampled from Poisson distribution.

        See `get_realization` for parameters.

        Returns:
            (np.ndarray of shape (M, A, R))
                contains positive integers with transition realizations
                for individuals going to compartment m in age-risk group (a, r)
        """

        num_outflows = len(self.transition_variables)

        num_age_groups, num_risk_groups = np.shape(self.origin.current_val)

        realizations_array = np.zeros((num_outflows, num_age_groups, num_risk_groups))

        transition_variables = self.transition_variables

        for age_group in range(num_age_groups):
            for risk_group in range(num_risk_groups):
                for outflow_ix in range(num_outflows):
                    realizations_array[outflow_ix, age_group, risk_group] = RNG.poisson(
                        self.origin.current_val[age_group, risk_group] *
                        transition_variables[outflow_ix].current_rate[
                            age_group, risk_group] / num_timesteps)

        return realizations_array

    def get_multinom_deterministic_realization(self,
                                               RNG: np.random.Generator,
                                               num_timesteps: int) -> np.ndarray:
        """
        Deterministic counterpart to `get_multinom_realization` --
        uses mean (n x p, i.e. total counts x probability array) as realization
        rather than randomly sampling.

        See `get_realization` for parameters.

        Returns:
            (np.ndarray of shape (M + 1, A, R))
                contains positive integers with transition realizations
                for individuals going to compartment m in age-risk group (a, r) --
                note the "+1" corresponds to the multinomial outcome of staying
                in the same compartment (not transitioning to any outgoing
                epi compartment).
        """

        probabilities_array = self.get_probabilities_array(num_timesteps)
        return np.asarray(self.origin.current_val * probabilities_array, dtype=int)

    def get_multinom_deterministic_no_round_realization(self,
                                                        RNG: np.random.Generator,
                                                        num_timesteps: int) -> np.ndarray:
        """
        The same as `get_multinom_deterministic_realization` except no rounding --
        so the populations can be non-integer. This is used to test the torch
        implementation (because that implementation does not round either).

        See `get_realization` for parameters.

        Returns:
            (np.ndarray of shape (M + 1, A, R))
                contains positive floats with transition realizations
                for individuals going to compartment m in age-risk group (a, r) --
                note the "+1" corresponds to the multinomial outcome of staying
                in the same compartment (not transitioning to any outgoing
                epi compartment).
        """

        probabilities_array = self.get_probabilities_array(num_timesteps)
        return np.asarray(self.origin.current_val * probabilities_array)

    def get_multinom_taylor_approx_deterministic_realization(self,
                                                             RNG: np.random.Generator,
                                                             num_timesteps: int) -> np.ndarray:
        """
        Deterministic counterpart to `get_multinom_taylor_approx_realization` --
        uses mean (n x p, i.e. total counts x probability array) as realization
        rather than randomly sampling.

        See `get_realization` for parameters.

        Returns:
            (np.ndarray of shape (M + 1, A, R))
                contains positive floats with transition realizations
                for individuals going to compartment m in age-risk group (a, r) --
                note the "+1" corresponds to the multinomial outcome of staying
                in the same compartment (not transitioning to any outgoing
                epi compartment).
        """

        current_rates_array = self.get_current_rates_array()
        return np.asarray(self.origin.current_val * current_rates_array / num_timesteps, dtype=int)

    def get_poisson_deterministic_realization(self,
                                              RNG: np.random.Generator,
                                              num_timesteps: int) -> np.ndarray:
        """
        Deterministic counterpart to `get_poisson_realization` --
        uses mean (rate array) as realization rather than randomly sampling.

        See `get_realization` for parameters.

        Returns:
            (np.ndarray of shape (A, R))
                contains positive integers with transition realizations
                for individuals going to compartment m in age-risk group (a, r) --
        """

        return np.asarray(self.origin.current_val *
                          self.get_current_rates_array() / num_timesteps, dtype=int)

    def reset(self) -> None:
        self.current_vals_list = []

    def update_transition_variable_realizations(self) -> None:
        """
        Updates current_val attribute on all `TransitionVariable`
        instances contained in this `TransitionVariableGroup`.
        """

        # Since the ith element in probabilities_array corresponds to the ith transition variable
        #   in transition_variables, the ith element in multinom_realizations_list
        #   also corresponds to the ith transition variable in transition_variables
        # Update the current realization of the transition variables contained in this group
        for ix in range(len(self.transition_variables)):
            self.transition_variables[ix].current_val = \
                self.current_vals_list[ix, :, :]

__init__(origin: Compartment, transition_type: TransitionTypes, transition_variables: list[TransitionVariable])

Parameters:

Name	Type	Description	Default
transition_type	TransitionTypes	Specifies probability distribution of transitions between compartments.	required

See class docstring for other parameters.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def __init__(self,
             origin: Compartment,
             transition_type: TransitionTypes,
             transition_variables: list[TransitionVariable]):
    """
    Args:
        transition_type (TransitionTypes):
            Specifies probability distribution of transitions between compartments.

    See class docstring for other parameters.
    """

    self.origin = origin

    # Using a list is important here because we want to keep the order
    #   of transition variables -- this determines the index in the
    #   current rates array
    self.transition_variables = transition_variables

    # If marginal transition type is any kind of binomial transition,
    #   then its joint transition type is a multinomial counterpart
    # For example, if the marginal transition type is TransitionTypes.BINOM_DETERMINISTIC,
    #   then the joint transition type is JointTransitionTypes.MULTINOM_DETERMINISTIC
    transition_type = transition_type.replace("binom", "multinom")
    self._transition_type = transition_type

    # Dynamically assign a method to get_joint_realization attribute
    #   based on the value of transition_type
    # getattr fetches a method by name
    self.get_joint_realization = getattr(self, "get_" + transition_type + "_realization")

    self.current_vals_list = []

get_current_rates_array() -> np.ndarray

Returns an array of current rates of transition variables in transition_variables -- ith element in array corresponds to current rate of ith transition variable.

Returns:

Type	Description
ndarray	(np.ndarray of shape (M, A, R)) array of positive floats corresponding to current rate element-wise for an outgoing compartment and age-risk group

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def get_current_rates_array(self) -> np.ndarray:
    """
    Returns an array of current rates of transition variables in
    `transition_variables` -- ith element in array
    corresponds to current rate of ith transition variable.

    Returns:
        (np.ndarray of shape (M, A, R))
            array of positive floats corresponding to current rate
            element-wise for an outgoing compartment and age-risk group
    """

    current_rates_list = []
    for tvar in self.transition_variables:
        current_rates_list.append(tvar.current_rate)

    return np.asarray(current_rates_list)

get_joint_realization(RNG: np.random.Generator, num_timesteps: int) -> np.ndarray

This function is dynamically assigned based on the TransitionVariableGroup's transition_type. It is set to the appropriate distribution-specific method.

See get_realization for parameters.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def get_joint_realization(self,
                          RNG: np.random.Generator,
                          num_timesteps: int) -> np.ndarray:
    """
    This function is dynamically assigned based on the `TransitionVariableGroup`'s
        `transition_type`. It is set to the appropriate distribution-specific method.

    See `get_realization` for parameters.
    """

    pass

get_multinom_deterministic_no_round_realization(RNG: np.random.Generator, num_timesteps: int) -> np.ndarray

The same as get_multinom_deterministic_realization except no rounding -- so the populations can be non-integer. This is used to test the torch implementation (because that implementation does not round either).

See get_realization for parameters.

Returns:

Type	Description
ndarray	(np.ndarray of shape (M + 1, A, R)) contains positive floats with transition realizations for individuals going to compartment m in age-risk group (a, r) -- note the "+1" corresponds to the multinomial outcome of staying in the same compartment (not transitioning to any outgoing epi compartment).

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def get_multinom_deterministic_no_round_realization(self,
                                                    RNG: np.random.Generator,
                                                    num_timesteps: int) -> np.ndarray:
    """
    The same as `get_multinom_deterministic_realization` except no rounding --
    so the populations can be non-integer. This is used to test the torch
    implementation (because that implementation does not round either).

    See `get_realization` for parameters.

    Returns:
        (np.ndarray of shape (M + 1, A, R))
            contains positive floats with transition realizations
            for individuals going to compartment m in age-risk group (a, r) --
            note the "+1" corresponds to the multinomial outcome of staying
            in the same compartment (not transitioning to any outgoing
            epi compartment).
    """

    probabilities_array = self.get_probabilities_array(num_timesteps)
    return np.asarray(self.origin.current_val * probabilities_array)

get_multinom_deterministic_realization(RNG: np.random.Generator, num_timesteps: int) -> np.ndarray

Deterministic counterpart to get_multinom_realization -- uses mean (n x p, i.e. total counts x probability array) as realization rather than randomly sampling.

See get_realization for parameters.

Returns:

Type	Description
ndarray	(np.ndarray of shape (M + 1, A, R)) contains positive integers with transition realizations for individuals going to compartment m in age-risk group (a, r) -- note the "+1" corresponds to the multinomial outcome of staying in the same compartment (not transitioning to any outgoing epi compartment).

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def get_multinom_deterministic_realization(self,
                                           RNG: np.random.Generator,
                                           num_timesteps: int) -> np.ndarray:
    """
    Deterministic counterpart to `get_multinom_realization` --
    uses mean (n x p, i.e. total counts x probability array) as realization
    rather than randomly sampling.

    See `get_realization` for parameters.

    Returns:
        (np.ndarray of shape (M + 1, A, R))
            contains positive integers with transition realizations
            for individuals going to compartment m in age-risk group (a, r) --
            note the "+1" corresponds to the multinomial outcome of staying
            in the same compartment (not transitioning to any outgoing
            epi compartment).
    """

    probabilities_array = self.get_probabilities_array(num_timesteps)
    return np.asarray(self.origin.current_val * probabilities_array, dtype=int)

get_multinom_realization(RNG: np.random.Generator, num_timesteps: int) -> np.ndarray

Returns an array of transition realizations (number transitioning to outgoing compartments) sampled from multinomial distribution.

See get_realization for parameters.

Returns:

Type	Description
ndarray	(np.ndarray of shape (M + 1, A, R)) contains positive floats with transition realizations for individuals going to compartment m in age-risk group (a, r) -- note the "+1" corresponds to the multinomial outcome of staying in the same compartment (not transitioning to any outgoing epi compartment).

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def get_multinom_realization(self,
                             RNG: np.random.Generator,
                             num_timesteps: int) -> np.ndarray:
    """
    Returns an array of transition realizations (number transitioning
    to outgoing compartments) sampled from multinomial distribution.

    See `get_realization` for parameters.

    Returns:
        (np.ndarray of shape (M + 1, A, R))
            contains positive floats with transition realizations
            for individuals going to compartment m in age-risk group (a, r) --
            note the "+1" corresponds to the multinomial outcome of staying
            in the same compartment (not transitioning to any outgoing
            epi compartment).
    """

    probabilities_array = self.get_probabilities_array(num_timesteps)

    num_outflows = len(self.transition_variables)

    num_age_groups, num_risk_groups = np.shape(self.origin.current_val)

    # We use num_outflows + 1 because for the multinomial distribution we explicitly model
    #   the number who stay/remain in the compartment
    realizations_array = np.zeros((num_outflows + 1, num_age_groups, num_risk_groups))

    for age_group in range(num_age_groups):
        for risk_group in range(num_risk_groups):
            realizations_array[:, age_group, risk_group] = RNG.multinomial(
                np.asarray(self.origin.current_val[age_group, risk_group], dtype=int),
                probabilities_array[:, age_group, risk_group])

    return realizations_array

get_multinom_taylor_approx_deterministic_realization(RNG: np.random.Generator, num_timesteps: int) -> np.ndarray

Deterministic counterpart to get_multinom_taylor_approx_realization -- uses mean (n x p, i.e. total counts x probability array) as realization rather than randomly sampling.

See get_realization for parameters.

Returns:

Type	Description
ndarray	(np.ndarray of shape (M + 1, A, R)) contains positive floats with transition realizations for individuals going to compartment m in age-risk group (a, r) -- note the "+1" corresponds to the multinomial outcome of staying in the same compartment (not transitioning to any outgoing epi compartment).

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def get_multinom_taylor_approx_deterministic_realization(self,
                                                         RNG: np.random.Generator,
                                                         num_timesteps: int) -> np.ndarray:
    """
    Deterministic counterpart to `get_multinom_taylor_approx_realization` --
    uses mean (n x p, i.e. total counts x probability array) as realization
    rather than randomly sampling.

    See `get_realization` for parameters.

    Returns:
        (np.ndarray of shape (M + 1, A, R))
            contains positive floats with transition realizations
            for individuals going to compartment m in age-risk group (a, r) --
            note the "+1" corresponds to the multinomial outcome of staying
            in the same compartment (not transitioning to any outgoing
            epi compartment).
    """

    current_rates_array = self.get_current_rates_array()
    return np.asarray(self.origin.current_val * current_rates_array / num_timesteps, dtype=int)

get_multinom_taylor_approx_realization(RNG: np.random.Generator, num_timesteps: int) -> np.ndarray

Returns an array of transition realizations (number transitioning to outgoing compartments) sampled from multinomial distribution using Taylor Series approximation for probability parameter.

See get_realization for parameters.

Returns:

Type	Description
ndarray	(np.ndarray of shape (M + 1, A, R)) contains positive integers with transition realizations for individuals going to compartment m in age-risk group (a, r) -- note the "+1" corresponds to the multinomial outcome of staying in the same compartment (not transitioning to any outgoing epi compartment).

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def get_multinom_taylor_approx_realization(self,
                                           RNG: np.random.Generator,
                                           num_timesteps: int) -> np.ndarray:
    """
    Returns an array of transition realizations (number transitioning
    to outgoing compartments) sampled from multinomial distribution
    using Taylor Series approximation for probability parameter.

    See `get_realization` for parameters.

    Returns:
        (np.ndarray of shape (M + 1, A, R))
            contains positive integers with transition realizations
            for individuals going to compartment m in age-risk group (a, r) --
            note the "+1" corresponds to the multinomial outcome of staying
            in the same compartment (not transitioning to any outgoing
            epi compartment).
    """

    num_outflows = len(self.transition_variables)

    current_rates_array = self.get_current_rates_array()

    total_rate = self.get_total_rate()

    # Multiply current rates array by length of time interval (1 / num_timesteps)
    # Also append additional value corresponding to probability of
    #   remaining in current epi compartment (not transitioning at all)
    # Note: "vstack" function here works better than append function because append
    #   automatically flattens the resulting array, resulting in dimension issues
    current_scaled_rates_array = np.vstack((current_rates_array / num_timesteps,
                                            np.expand_dims(1 - total_rate / num_timesteps, axis=0)))

    num_age_groups, num_risk_groups = np.shape(self.origin.current_val)

    # We use num_outflows + 1 because for the multinomial distribution we explicitly model
    #   the number who stay/remain in the compartment
    realizations_array = np.zeros((num_outflows + 1, num_age_groups, num_risk_groups))

    for age_group in range(num_age_groups):
        for risk_group in range(num_risk_groups):
            realizations_array[:, age_group, risk_group] = RNG.multinomial(
                np.asarray(self.origin.current_val[age_group, risk_group], dtype=int),
                current_scaled_rates_array[:, age_group, risk_group])

    return realizations_array

get_poisson_deterministic_realization(RNG: np.random.Generator, num_timesteps: int) -> np.ndarray

Deterministic counterpart to get_poisson_realization -- uses mean (rate array) as realization rather than randomly sampling.

See get_realization for parameters.

Returns:

Type	Description
ndarray	(np.ndarray of shape (A, R)) contains positive integers with transition realizations for individuals going to compartment m in age-risk group (a, r) --

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def get_poisson_deterministic_realization(self,
                                          RNG: np.random.Generator,
                                          num_timesteps: int) -> np.ndarray:
    """
    Deterministic counterpart to `get_poisson_realization` --
    uses mean (rate array) as realization rather than randomly sampling.

    See `get_realization` for parameters.

    Returns:
        (np.ndarray of shape (A, R))
            contains positive integers with transition realizations
            for individuals going to compartment m in age-risk group (a, r) --
    """

    return np.asarray(self.origin.current_val *
                      self.get_current_rates_array() / num_timesteps, dtype=int)

get_poisson_realization(RNG: np.random.Generator, num_timesteps: int) -> np.ndarray

Returns an array of transition realizations (number transitioning to outgoing compartments) sampled from Poisson distribution.

See get_realization for parameters.

Returns:

Type	Description
ndarray	(np.ndarray of shape (M, A, R)) contains positive integers with transition realizations for individuals going to compartment m in age-risk group (a, r)

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def get_poisson_realization(self,
                            RNG: np.random.Generator,
                            num_timesteps: int) -> np.ndarray:
    """
    Returns an array of transition realizations (number transitioning
    to outgoing compartments) sampled from Poisson distribution.

    See `get_realization` for parameters.

    Returns:
        (np.ndarray of shape (M, A, R))
            contains positive integers with transition realizations
            for individuals going to compartment m in age-risk group (a, r)
    """

    num_outflows = len(self.transition_variables)

    num_age_groups, num_risk_groups = np.shape(self.origin.current_val)

    realizations_array = np.zeros((num_outflows, num_age_groups, num_risk_groups))

    transition_variables = self.transition_variables

    for age_group in range(num_age_groups):
        for risk_group in range(num_risk_groups):
            for outflow_ix in range(num_outflows):
                realizations_array[outflow_ix, age_group, risk_group] = RNG.poisson(
                    self.origin.current_val[age_group, risk_group] *
                    transition_variables[outflow_ix].current_rate[
                        age_group, risk_group] / num_timesteps)

    return realizations_array

get_probabilities_array(num_timesteps: int) -> list

Returns an array of probabilities used for joint binomial (multinomial) transitions (get_multinom_realization method).

Returns:

Type	Description
list	(np.ndarray of shape (M+1, A, R) Contains positive floats <= 1, corresponding to probability of transitioning to a compartment for that outgoing compartment and age-risk group -- note the "+1" corresponds to the multinomial outcome of staying in the same compartment (we can think of as transitioning to the same compartment).

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def get_probabilities_array(self,
                            num_timesteps: int) -> list:
    """
    Returns an array of probabilities used for joint binomial
    (multinomial) transitions (`get_multinom_realization` method).

    Returns:
        (np.ndarray of shape (M+1, A, R)
            Contains positive floats <= 1, corresponding to probability
            of transitioning to a compartment for that outgoing compartment
            and age-risk group -- note the "+1" corresponds to the multinomial
            outcome of staying in the same compartment (we can think of as
            transitioning to the same compartment).
    """

    total_rate = self.get_total_rate()

    total_outgoing_probability = approx_binom_probability_from_rate(total_rate,
                                                                    1 / num_timesteps)

    # Create probabilities_list, where element i corresponds to the
    #   transition variable i's current rate divided by the total rate,
    #   multiplized by the total outgoing probability
    # This generates the probabilities array that parameterizes the
    #   multinomial distribution
    probabilities_list = []

    for transition_variable in self.transition_variables:
        probabilities_list.append((transition_variable.current_rate / total_rate) *
                                  total_outgoing_probability)

    # Append the probability that a person stays in the compartment
    probabilities_list.append(1 - total_outgoing_probability)

    return np.asarray(probabilities_list)

get_total_rate() -> np.ndarray

Return the age-risk-specific total transition rate, which is the sum of the current rate of each transition variable in this transition variable group.

Used to properly scale multinomial probabilities vector so that elements sum to 1.

Returns:

Type	Description
ndarray	(np.ndarray of shape (A, R)) Array with values corresponding to sum of current rates of transition variables in transition variable group, where elements correspond to age-risk groups.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def get_total_rate(self) -> np.ndarray:
    """
    Return the age-risk-specific total transition rate,
    which is the sum of the current rate of each transition variable
    in this transition variable group.

    Used to properly scale multinomial probabilities vector so
    that elements sum to 1.

    Returns:
        (np.ndarray of shape (A, R))
            Array with values corresponding to sum of current rates of
            transition variables in transition variable group, where
            elements correspond to age-risk groups.
    """

    # axis 0: corresponds to outgoing transition variable
    # axis 1: corresponds to age groups
    # axis 2: corresponds to risk groups
    # --> summing over axis 0 gives the total rate for each age-risk group
    return np.sum(self.get_current_rates_array(), axis=0)

update_transition_variable_realizations() -> None

Updates current_val attribute on all TransitionVariable instances contained in this TransitionVariableGroup.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def update_transition_variable_realizations(self) -> None:
    """
    Updates current_val attribute on all `TransitionVariable`
    instances contained in this `TransitionVariableGroup`.
    """

    # Since the ith element in probabilities_array corresponds to the ith transition variable
    #   in transition_variables, the ith element in multinom_realizations_list
    #   also corresponds to the ith transition variable in transition_variables
    # Update the current realization of the transition variables contained in this group
    for ix in range(len(self.transition_variables)):
        self.transition_variables[ix].current_val = \
            self.current_vals_list[ix, :, :]

UniformSamplingSpec dataclass

Holds Uniform distribution info to randomly sample a subpop model's SubpopParams attribute.

Attributes:

Name	Type	Description
lower_bound	[ndarray | float]	Lower bound(s) of the uniform distribution. Can be a scalar, shape (A,) array, or shape (A, R) array depending on param_shape.
upper_bound	[ndarray | float]	Upper bound(s) of the uniform distribution. Must have the same shape as lower_bound.
param_shape	ParamShapes	Describes how the parameter varies across subpopulations (scalar, by age, or by age and risk).
num_decimals	positive int	Optional number of decimals to keep after rounding -- default is 2.

Source code in CLT_BaseModel/clt_toolkit/sampling.py

@dataclass(frozen=True)
class UniformSamplingSpec:
    """
    Holds Uniform distribution info to randomly sample a
    subpop model's `SubpopParams` attribute.

    Attributes:
        lower_bound ([np.ndarray | float]):
            Lower bound(s) of the uniform distribution. Can be a scalar,
            shape (A,) array, or shape (A, R) array depending on `param_shape`.
        upper_bound ([np.ndarray | float]):
            Upper bound(s) of the uniform distribution. Must have the same shape
            as `lower_bound`.
        param_shape (ParamShapes):
            Describes how the parameter varies across subpopulations
            (scalar, by age, or by age and risk).
        num_decimals (positive int):
            Optional number of decimals to keep after rounding -- default is 2.
    """

    lower_bound: Optional[np.ndarray | float] = None
    upper_bound: Optional[np.ndarray | float] = None
    param_shape: Optional[ParamShapes] = None
    num_decimals: Optional[int] = 2

aggregate_daily_tvar_history(metapop_model: MetapopModel, transition_var_name: str) -> np.ndarray

Sum the history values of a given transition variable across all subpopulations and across timesteps per day, so that we have the total number that transitioned compartments in a day.

Parameters:

Name	Type	Description	Default
metapop_model	MetapopModel	The metapopulation model containing subpopulations.	required
transition_var_name	str	Name of the transition variable to sum.	required

Returns:

Name	Type	Description
total	ndarray	Array of shape (num_days, A, R) containing the sum across all subpopulations, where A = number of age groups, and R = number of risk groups. Each element contains the total number of individuals who transitioned that day for the given age and risk group.

Source code in CLT_BaseModel/clt_toolkit/sampling.py

def aggregate_daily_tvar_history(metapop_model: MetapopModel,
                                 transition_var_name: str) -> np.ndarray:
    """
    Sum the history values of a given transition variable
    across all subpopulations and across timesteps per day,
    so that we have the total number that transitioned compartments
    in a day.

    Parameters:
        metapop_model (MetapopModel):
            The metapopulation model containing subpopulations.
        transition_var_name (str):
            Name of the transition variable to sum.

    Returns:
        total (np.ndarray):
            Array of shape (num_days, A, R) containing the sum across all subpopulations,
            where A = number of age groups, and R = number of risk groups.
            Each element contains the total number of individuals who transitioned that
            day for the given age and risk group.
    """
    # Convert each subpop's history list to a NumPy array
    all_arrays = [
        np.asarray(getattr(subpop, transition_var_name).history_vals_list)
        for subpop in metapop_model.subpop_models.values()
    ]

    # Stack along new subpop dimension (axis=0) and sum across subpops
    total = np.sum(np.stack(all_arrays, axis=0), axis=0)

    # Each subpopulation model should have the same simulation settings, so
    #   just grab the first subpop model
    num_timesteps = metapop_model.subpop_models[0].simulation_settings.timesteps_per_day

    # Transition variable history contains values recorded at each TIMESTEP.
    # To get DAILY totals, we sum blocks of `num_timesteps` consecutive timesteps.
    return daily_sum_over_timesteps(total, num_timesteps)

approx_binom_probability_from_rate(rate: np.ndarray, interval_length: int) -> np.ndarray

Converts a rate (events per time) to the probability of any event occurring in the next time interval of length interval_length, assuming the number of events occurring in time interval follows a Poisson distribution with given rate parameter.

The probability of 0 events in interval_length is e^(-rate * interval_length), so the probability of any event in interval_length is 1 - e^(-rate * interval_length).

Rate must be A x R np.ndarray, where A is the number of age groups and R is the number of risk groups. Rate is transformed to A x R np.ndarray corresponding to probabilities.

Parameters:

Name	Type	Description	Default
rate	np.ndarray of shape (A, R	Rate parameters in a Poisson distribution per age-risk group.	required
interval_length	positive int	Length of time interval in simulation days.	required

Returns:

Type	Description
ndarray	np.ndarray of shape (A, R): Array of positive scalars corresponding to probability that any individual in an age-risk group transitions compartments.

Source code in CLT_BaseModel/clt_toolkit/base_components.py

def approx_binom_probability_from_rate(rate: np.ndarray,
                                       interval_length: int) -> np.ndarray:
    """
    Converts a rate (events per time) to the probability of any event
    occurring in the next time interval of length `interval_length`,
    assuming the number of events occurring in time interval
    follows a Poisson distribution with given rate parameter.

    The probability of 0 events in `interval_length` is
    e^(-`rate` * `interval_length`), so the probability of any event
    in `interval_length` is 1 - e^(-`rate` * `interval_length`).

    Rate must be A x R `np.ndarray`, where A is the number of
    age groups and R is the number of risk groups. Rate is transformed to
    A x R `np.ndarray` corresponding to probabilities.

    Parameters:
        rate (np.ndarray of shape (A, R)):
            Rate parameters in a Poisson distribution per age-risk group.
        interval_length (positive int):
            Length of time interval in simulation days.

    Returns:
        np.ndarray of shape (A, R):
            Array of positive scalars corresponding to probability that
            any individual in an age-risk group transitions compartments.
    """

    return 1 - np.exp(-rate * interval_length)

check_is_subset_list(listA: list, listB: list) -> bool

Parameters:

Name	Type	Description	Default
listA	list	list-like of elements to check if subset of listB.	required
listB	list	list-like of elements.	required

Returns:

Type	Description
bool	True if listA is a subset of listB, and False otherwise.

Source code in CLT_BaseModel/clt_toolkit/experiments.py

def check_is_subset_list(listA: list,
                         listB: list) -> bool:
    """
    Params:
        listA (list):
            list-like of elements to check if subset of listB.
        listB (list):
            list-like of elements.

    Returns:
        True if listA is a subset of listB, and False otherwise.
    """

    return all(item in listB for item in listA)

convert_dict_vals_lists_to_arrays(d: dict) -> dict

Converts dictionary of lists to dictionary of arrays to support numpy operations.

Source code in CLT_BaseModel/clt_toolkit/input_parsers.py

def convert_dict_vals_lists_to_arrays(d: dict) -> dict:
    """
    Converts dictionary of lists to dictionary of arrays
    to support `numpy` operations.
    """

    for key, val in d.items():
        if type(val) is list:
            d[key] = np.asarray(val)

    return d

daily_sum_over_timesteps(x: np.ndarray, num_timesteps: int) -> np.ndarray

For example, used for transition variable history, which is saved for every timestep, but we generally would like converted to daily totals.

Parameters:

Name	Type	Description	Default
x	np.ndarray of shape (N, A, R	Array to aggregate -- N is the number of timesteps, A is the number of age groups, R is the number of risk groups.	required
num_timesteps	int	Number of timesteps per day. Must divide N (length of x) evenly.	required

Returns:

Type	Description
np.ndarray of shape (N/n, A, R	Array of daily totals, where each block of num_timesteps consecutive timesteps from x has been summed along the first dimension. The first axis now represents days instead of individual timesteps.

Source code in CLT_BaseModel/clt_toolkit/utils.py

def daily_sum_over_timesteps(x: np.ndarray,
                             num_timesteps: int) -> np.ndarray:
    """
    For example, used for transition variable history, which is
    saved for every timestep, but we generally would like converted to daily totals.

    Params:
        x (np.ndarray of shape (N, A, R)):
            Array to aggregate -- N is the number of
            timesteps, A is the number of age groups,
            R is the number of risk groups.
        num_timesteps (int):
            Number of timesteps per day. Must divide
            N (length of `x`) evenly.

    Returns:
        (np.ndarray of shape (N/n, A, R):
            Array of daily totals, where each block of `num_timesteps`
            consecutive timesteps from `x` has been summed along the
            first dimension. The first axis now represents days instead
            of individual timesteps.
    """

    total_timesteps, A, R = x.shape

    if total_timesteps / num_timesteps != total_timesteps // num_timesteps:
        raise ValueError("x must be shape (N, A, R), where num_timesteps divides N.")

    num_days = int(total_timesteps/num_timesteps)

    # Pretty sweet hack ;)
    x = x.reshape(num_days, num_timesteps, A, R)

    # Sum along the number of timesteps
    return x.sum(axis=1)

format_current_val_for_sql(subpop_model: SubpopModel, state_var_name: str, rep: int) -> list

Processes current_val of given subpop_model's StateVariable specified by state_var_name. Current_val is an A x R numpy array (for age-risk) -- this function "unpacks" it into an (A x R, 1) numpy array (a column vector). Converts metadata (subpop_name, state_var_name, rep, and current_simulation_day) into list of A x R rows, where each row has 7 elements, for consistent row formatting for batch SQL insertion.

Parameters:

Name	Type	Description	Default
subpop_model	SubpopModel	SubpopModel to record.	required
state_var_name	str	StateVariable name to record.	required
rep	int	replication counter to record.	required

Returns:

Name	Type	Description
data	list	list of A x R rows, where each row is a list of 7 elements corresponding to subpop_name, state_var_name, age_group, risk_group, rep, current_simulation_day, and the scalar element of current_val corresponding to that age-risk group.

Source code in CLT_BaseModel/clt_toolkit/experiments.py

def format_current_val_for_sql(subpop_model: SubpopModel,
                               state_var_name: str,
                               rep: int) -> list:
    """
    Processes current_val of given subpop_model's `StateVariable`
    specified by `state_var_name`. Current_val is an A x R
    numpy array (for age-risk) -- this function "unpacks" it into an
    (A x R, 1) numpy array (a column vector). Converts metadata
    (subpop_name, state_var_name, `rep`, and current_simulation_day)
    into list of A x R rows, where each row has 7 elements, for
    consistent row formatting for batch SQL insertion.

    Params:
        subpop_model (SubpopModel):
            SubpopModel to record.
        state_var_name (str):
            StateVariable name to record.
        rep (int):
            replication counter to record.

    Returns:
        data (list):
            list of A x R rows, where each row is a list of 7 elements
            corresponding to subpop_name, state_var_name, age_group, risk_group,
            rep, current_simulation_day, and the scalar element of current_val
            corresponding to that age-risk group.
    """

    current_val = subpop_model.all_state_variables[state_var_name].current_val

    A, R = np.shape(current_val)

    # numpy's default is row-major / C-style order
    # This means the elements are unpacked ROW BY ROW
    current_val_reshaped = current_val.reshape(-1, 1)

    # (AxR, 1) column vector of row indices, indicating the original row in current_val
    #   before reshaping
    # Each integer in np.arange(A) repeated R times
    age_group_indices = np.repeat(np.arange(A), R).reshape(-1, 1)

    # (AxR, 1) column vector of column indices, indicating the original column
    #   each element belonged to in current_val before reshaping
    # Repeat np.arange(R) A times
    risk_group_indices = np.tile(np.arange(R), A).reshape(-1, 1)

    # (subpop_name, state_var_name, age_group, risk_group, rep, timepoint)
    data = np.column_stack(
        (np.full((A * R, 1), subpop_model.name),
         np.full((A * R, 1), state_var_name),
         age_group_indices,
         risk_group_indices,
         np.full((A * R, 1), rep),
         np.full((A * R, 1), subpop_model.current_simulation_day),
         current_val_reshaped)).tolist()

    return data

get_sql_table_as_df(conn: sqlite3.Connection, sql_query: str, sql_query_params: tuple[str] = None, chunk_size: int = int(10000.0)) -> pd.DataFrame

Returns a pandas DataFrame containing data from specified SQL table, retrieved using the provided database connection. Reads in SQL rows in batches of size chunk_size to avoid memory issues for very large tables.

Parameters:

Name	Type	Description	Default
conn	Connection	connection to SQL database.	required
sql_query	str	SQL query/statement to execute on database.	required
sql_query_params	tuple[str]	tuple of strings to pass as parameters to SQL query -- used to avoid SQL injections.	None
chunk_size	positive int	number of rows to read in at a time.	int(10000.0)

Returns:

Type	Description
DataFrame	DataFrame containing data from specified SQL table,
DataFrame	or empty DataFrame if table does not exist.

Source code in CLT_BaseModel/clt_toolkit/experiments.py

def get_sql_table_as_df(conn: sqlite3.Connection,
                        sql_query: str,
                        sql_query_params: tuple[str] = None,
                        chunk_size: int = int(1e4)) -> pd.DataFrame:
    """
    Returns a pandas DataFrame containing data from specified SQL table,
    retrieved using the provided database connection. Reads in SQL rows
    in batches of size `chunk_size` to avoid memory issues for very large
    tables.

    Params:
        conn (sqlite3.Connection):
            connection to SQL database.
        sql_query (str):
            SQL query/statement to execute on database.
        sql_query_params (tuple[str]):
            tuple of strings to pass as parameters to
            SQL query -- used to avoid SQL injections.
        chunk_size (positive int):
            number of rows to read in at a time.

    Returns:
        DataFrame containing data from specified SQL table,
        or empty DataFrame if table does not exist.
    """

    chunks = []

    try:
        for chunk in pd.read_sql_query(sql_query,
                                       conn,
                                       chunksize=chunk_size,
                                       params=sql_query_params):
            chunks.append(chunk)
            df = pd.concat(chunks, ignore_index=True)

    # Handle exception gracefully -- print a warning and
    #   return an empty DataFrame if table given by sql_query
    #   does not exist
    except sqlite3.OperationalError as e:
        if "no such table" in str(e).lower():
            print(f"Warning: table does not exist for query: {sql_query}. "
                  f"Returning empty DataFrame.")
            df = pd.DataFrame()

    return df

load_json_augment_dict(json_filepath: str, d: dict) -> dict

Augments pre-existing dictionary with information from JSON file -- if keys already exist, the previous values are overriden, otherwise the new key-value pairs are added. Lists are automatically converted to numpy arrays for computational compatibility, since JSON does not natively support np.ndarray.

Parameters:

Name	Type	Description	Default
json_filepath	str	Full JSON filepath.	required
d	dict	Dictionary to be augmented with new JSON values.	required

Returns:

Type	Description
dict	Dictionary loaded with JSON information.

Source code in CLT_BaseModel/clt_toolkit/input_parsers.py

def load_json_augment_dict(json_filepath: str,
                           d: dict) -> dict:
    """
    Augments pre-existing dictionary with information
    from `JSON` file -- if keys already exist, the previous values
    are overriden, otherwise the new key-value pairs are added.
    Lists are automatically converted to numpy arrays for
    computational compatibility, since `JSON` does not natively
    support `np.ndarray`.

    Args:
        json_filepath (str):
            Full `JSON` filepath.
        d (dict):
            Dictionary to be augmented with new `JSON` values.

    Returns:
        (dict):
            Dictionary loaded with `JSON` information.
    """

    with open(json_filepath, 'r') as file:
        data = json.load(file)

    data = convert_dict_vals_lists_to_arrays(data)

    for key, val in data.items():
        d[key] = val

    return d

load_json_new_dict(json_filepath: str) -> dict

Loads specified JSON file into new dictionary. Lists are automatically converted to numpy arrays for computational compatibility, since JSON does not natively support np.ndarray.

Parameters:

Name	Type	Description	Default
json_filepath	str	Full JSON filepath.	required

Returns:

Type	Description
dict	Dictionary loaded with JSON information.

Source code in CLT_BaseModel/clt_toolkit/input_parsers.py

def load_json_new_dict(json_filepath: str) -> dict:
    """
    Loads specified `JSON` file into new dictionary.
    Lists are automatically converted to numpy arrays for
    computational compatibility, since `JSON` does not natively
    support `np.ndarray`.

    Args:
        json_filepath (str):
            Full `JSON` filepath.

    Returns:
        (dict):
            Dictionary loaded with `JSON` information.
    """

    # Note: the "with open" is important for file handling
    #   and avoiding resource leaks -- otherwise,
    #   we have to manually close the file, which is a bit
    #   more cumbersome
    with open(json_filepath, 'r') as file:
        data = json.load(file)

    # json does not support numpy, so we must convert
    #   lists to numpy arrays
    return convert_dict_vals_lists_to_arrays(data)

make_dataclass_from_dict(dataclass_ref: Type[DataClassProtocol], d: dict) -> DataClassProtocol

Create instance of class dataclass_ref, based on information in dictionary.

Parameters:

Name	Type	Description	Default
dataclass_ref	Type[DataClassProtocol]	(class, not instance) from which to create instance -- must have dataclass decorator.	required
d	dict	all keys and values respectively must match name and datatype of dataclass_ref instance attributes.	required

Returns:

Name	Type	Description
DataClassProtocol	DataClassProtocol	instance of dataclass_ref with attributes dynamically assigned by json_filepath file contents.

Source code in CLT_BaseModel/clt_toolkit/input_parsers.py

def make_dataclass_from_dict(dataclass_ref: Type[DataClassProtocol],
                             d: dict) -> DataClassProtocol:
    """
    Create instance of class dataclass_ref,
    based on information in dictionary.

    Args:
        dataclass_ref (Type[DataClassProtocol]):
            (class, not instance) from which to create instance --
            must have dataclass decorator.
        d (dict):
            all keys and values respectively must match name and datatype
            of dataclass_ref instance attributes.

    Returns:
        DataClassProtocol:
            instance of dataclass_ref with attributes dynamically
            assigned by json_filepath file contents.
    """

    d = convert_dict_vals_lists_to_arrays(d)

    return dataclass_ref(**d)

make_dataclass_from_json(json_filepath: str, dataclass_ref: Type[DataClassProtocol]) -> DataClassProtocol

Create instance of class dataclass_ref, based on information in json_filepath.

Parameters:

Name	Type	Description	Default
json_filepath	str	path to json file (path includes actual filename with suffix ".json") -- all json fields must match name and datatype of dataclass_ref instance attributes.	required
dataclass_ref	Type[DataClassProtocol]	(class, not instance) from which to create instance -- must have dataclass decorator.	required

Returns:

Name	Type	Description
DataClassProtocol	DataClassProtocol	instance of dataclass_ref with attributes dynamically assigned by json_filepath file contents.

Source code in CLT_BaseModel/clt_toolkit/input_parsers.py

def make_dataclass_from_json(json_filepath: str,
                             dataclass_ref: Type[DataClassProtocol]) -> DataClassProtocol:
    """
    Create instance of class dataclass_ref,
    based on information in json_filepath.

    Args:
        json_filepath (str):
            path to json file (path includes actual filename
            with suffix ".json") -- all json fields must
            match name and datatype of dataclass_ref instance
            attributes.
        dataclass_ref (Type[DataClassProtocol]):
            (class, not instance) from which to create instance --
            must have dataclass decorator.

    Returns:
        DataClassProtocol:
            instance of dataclass_ref with attributes dynamically
            assigned by json_filepath file contents.
    """

    d = load_json_new_dict(json_filepath)

    return make_dataclass_from_dict(dataclass_ref, d)

plot_metapop_basic_compartment_history(metapop_model: MetapopModel, axes: matplotlib.axes.Axes = None)

Plots the compartment data for a metapopulation model.

Parameters:

Name	Type	Description	Default
metapop_model	MetapopModel	Metapopulation model containing compartments.	required
axes	Axes	Matplotlib axes to plot on.	None

Source code in CLT_BaseModel/clt_toolkit/plotting.py

@plot_metapop_decorator
def plot_metapop_basic_compartment_history(metapop_model: MetapopModel,
                                           axes: matplotlib.axes.Axes = None):
    """
    Plots the compartment data for a metapopulation model.

    Args:
        metapop_model (MetapopModel):
            Metapopulation model containing compartments.
        axes (matplotlib.axes.Axes):
            Matplotlib axes to plot on.
    """

    # Iterate over subpop models and plot
    for ix, (subpop_name, subpop_model) in enumerate(metapop_model.subpop_models.items()):
        plot_subpop_basic_compartment_history(subpop_model, axes[ix])

plot_metapop_decorator(plot_func)

Decorator to handle common metapopulation plotting tasks.

Source code in CLT_BaseModel/clt_toolkit/plotting.py

def plot_metapop_decorator(plot_func):
    """
    Decorator to handle common metapopulation plotting tasks.
    """

    @functools.wraps(plot_func)
    def wrapper(metapop_model: MetapopModel,
                savefig_filename = None):

        num_plots = len(metapop_model.subpop_models)
        num_cols = 2
        num_rows = (num_plots + num_cols - 1) // num_cols

        # Create figure and axes
        fig, axes = plt.subplots(num_rows, num_cols, figsize=(5 * num_cols, 4 * num_rows))
        axes = axes.flatten()

        plot_func(metapop_model=metapop_model, axes=axes)

        # Turn off any unused subplots
        for j in range(num_plots, len(axes)):
            fig.delaxes(axes[j])  # Remove empty subplot

        # Adjust layout and save/show the figure
        plt.tight_layout()

        if savefig_filename:
            plt.savefig(savefig_filename, dpi=1200)

        plt.show()

    return wrapper

plot_metapop_epi_metrics(metapop_model: MetapopModel, axes: matplotlib.axes.Axes)

Plots the EpiMetric data for a metapopulation model.

Parameters:

Name	Type	Description	Default
metapop_model	MetapopModel	Metapopulation model containing compartments.	required
axes	Axes	Matplotlib axes to plot on.	required

Source code in CLT_BaseModel/clt_toolkit/plotting.py

@plot_metapop_decorator
def plot_metapop_epi_metrics(metapop_model: MetapopModel,
                             axes: matplotlib.axes.Axes):
    """
    Plots the EpiMetric data for a metapopulation model.

    Args:
        metapop_model (MetapopModel):
            Metapopulation model containing compartments.
        axes (matplotlib.axes.Axes):
            Matplotlib axes to plot on.
    """

    for ix, (subpop_name, subpop_model) in enumerate(metapop_model.subpop_models.items()):
        plot_subpop_epi_metrics(subpop_model, axes[ix])

plot_metapop_total_infected_deaths(metapop_model: MetapopModel, axes: matplotlib.axes.Axes)

Plots the total infected (IP+IS+IA) and deaths data for a metapopulation model.

Parameters:

Name	Type	Description	Default
metapop_model	MetapopModel	Metapopulation model containing compartments.	required
axes	Axes	Matplotlib axes to plot on.	required

Source code in CLT_BaseModel/clt_toolkit/plotting.py

@plot_metapop_decorator
def plot_metapop_total_infected_deaths(metapop_model: MetapopModel,
                                       axes: matplotlib.axes.Axes):
    """
    Plots the total infected (IP+IS+IA) and deaths data for a metapopulation model.

    Args:
        metapop_model (MetapopModel):
            Metapopulation model containing compartments.
        axes (matplotlib.axes.Axes):
            Matplotlib axes to plot on.
    """

    # Iterate over subpop models and plot
    for ix, (subpop_name, subpop_model) in enumerate(metapop_model.subpop_models.items()):
        plot_subpop_total_infected_deaths(subpop_model, axes[ix])

plot_subpop_basic_compartment_history(subpop_model: SubpopModel, ax: matplotlib.axes.Axes = None)

Plots data for a single subpopulation model on the given axis.

Parameters:

Name	Type	Description	Default
subpop_model	SubpopModel	Subpopulation model containing compartments.	required
ax	Axes	Matplotlib axis to plot on.	None

Source code in CLT_BaseModel/clt_toolkit/plotting.py

@plot_subpop_decorator
def plot_subpop_basic_compartment_history(subpop_model: SubpopModel,
                                          ax: matplotlib.axes.Axes = None):
    """
    Plots data for a single subpopulation model on the given axis.

    Args:
        subpop_model (SubpopModel):
            Subpopulation model containing compartments.
        ax (matplotlib.axes.Axes):
            Matplotlib axis to plot on.
    """

    for name, compartment in subpop_model.compartments.items():
        # Compute summed history values for each age-risk group
        history_vals_list = [np.sum(age_risk_group_entry) for
                             age_risk_group_entry in compartment.history_vals_list]

        # Plot data with a label
        ax.plot(history_vals_list, label=name, alpha=0.6)

    # Set axis title and labels
    ax.set_title(f"{subpop_model.name}")
    ax.set_xlabel("Days")
    ax.set_ylabel("Number of individuals")
    ax.legend()

plot_subpop_decorator(plot_func)

Decorator to handle common subpopulation plotting tasks.

Source code in CLT_BaseModel/clt_toolkit/plotting.py

def plot_subpop_decorator(plot_func):
    """
    Decorator to handle common subpopulation plotting tasks.
    """

    @functools.wraps(plot_func)
    def wrapper(subpop_model: SubpopModel,
                ax: matplotlib.axes.Axes = None,
                savefig_filename: str = None):
        """
        Args:
            subpop_model (SubpopModel):
                SubpopModel to plot.
            ax (matplotlib.axes.Axes):
                Matplotlib axis to plot on.
            savefig_filename (str):
                Optional filename to save the figure.
        """

        ax_provided = ax

        # If no axis is provided, create own axis
        if ax is None:
            fig, ax = plt.subplots()

        plot_func(subpop_model=subpop_model, ax=ax)

        if savefig_filename:
            plt.savefig(savefig_filename, dpi=1200)

        if ax_provided is None:
            plt.show()

    return wrapper

plot_subpop_epi_metrics(subpop_model: SubpopModel, ax: matplotlib.axes.Axes = None)

Plots EpiMetric history for a single subpopulation model on the given axis.

Parameters:

Name	Type	Description	Default
subpop_model	SubpopModel	Subpopulation model containing compartments.	required
ax	Axes	Matplotlib axis to plot on.	None

Source code in CLT_BaseModel/clt_toolkit/plotting.py

@plot_subpop_decorator
def plot_subpop_epi_metrics(subpop_model: SubpopModel,
                            ax: matplotlib.axes.Axes = None):
    """
    Plots EpiMetric history for a single subpopulation model on the given axis.

    Args:
        subpop_model (SubpopModel):
            Subpopulation model containing compartments.
        ax (matplotlib.axes.Axes):
            Matplotlib axis to plot on.
    """

    for name, epi_metric in subpop_model.epi_metrics.items():

        # Compute summed history values for each age-risk group
        history_vals_list = [np.average(age_risk_group_entry) for
                             age_risk_group_entry in epi_metric.history_vals_list]

        # Plot data with a label
        ax.plot(history_vals_list, label=name, alpha=0.6)

    # Set axis title and labels
    ax.set_title(f"{subpop_model.name}")
    ax.set_xlabel("Days")
    ax.set_ylabel("Epi Metric Value")
    ax.legend()

plot_subpop_total_infected_deaths(subpop_model: SubpopModel, ax: matplotlib.axes.Axes = None)

Plots data for a single subpopulation model on the given axis.

Parameters:

Name	Type	Description	Default
subpop_model	SubpopModel	Subpopulation model containing compartments.	required
ax	Axes	Matplotlib axis to plot on.	None

Source code in CLT_BaseModel/clt_toolkit/plotting.py

@plot_subpop_decorator
def plot_subpop_total_infected_deaths(subpop_model: SubpopModel,
                                      ax: matplotlib.axes.Axes = None):
    """
    Plots data for a single subpopulation model on the given axis.

    Args:
        subpop_model (SubpopModel):
            Subpopulation model containing compartments.
        ax (matplotlib.axes.Axes):
            Matplotlib axis to plot on.
    """

    infected_compartment_names = [name for name in subpop_model.compartments.keys() if
                                  "I" in name or "H" in name]

    infected_compartments_history = [subpop_model.compartments[compartment_name].history_vals_list
                                     for compartment_name in infected_compartment_names]

    total_infected = np.sum(np.asarray(infected_compartments_history), axis=(0, 2, 3))

    ax.plot(total_infected, label="Total infected", alpha=0.6)

    if "D" in subpop_model.compartments.keys():
        deaths = [np.sum(age_risk_group_entry)
                  for age_risk_group_entry
                  in subpop_model.compartments.D.history_vals_list]

        ax.plot(deaths, label="D", alpha=0.6)

    ax.set_title(f"{subpop_model.name}")
    ax.set_xlabel("Days")
    ax.set_ylabel("Number of individuals")
    ax.legend()

sample_uniform_matrix(lb: [np.ndarray | float], ub: [np.ndarray | float], RNG: np.random.Generator, A: int, R: int, param_shape: str) -> [np.ndarray | float]

Sample a matrix X of shape (A,R) such that X[a,r] ~ (independent) Uniform(low[a,r], high[a,r]). We assume each element is independent, so we do not assume any correlation structure:

Parameters:

Name	Type	Description	Default
lb	np.ndarray of shape (A,) or (A, R) or float	Array or scalar of lower bounds	required
ub	np.ndarray of shape (A,) or (A, R) or float	Array or scalar of upper bounds	required
RNG	Generator	Used to generate Uniform random variables.	required

Returns:

Name	Type	Description
X	np.ndarray of shape (A,) or (A, R) or float	Random matrix or scalar realization where each element is independently sampled from a Uniform distribution with parameters given element-wise by lb and ub -- X is same shape as lb and ub.

Source code in CLT_BaseModel/clt_toolkit/sampling.py

def sample_uniform_matrix(lb: [np.ndarray | float],
                          ub: [np.ndarray | float],
                          RNG: np.random.Generator,
                          A: int,
                          R: int,
                          param_shape: str,
                          ) -> [np.ndarray | float]:
    """
    Sample a matrix X of shape (A,R) such that
    X[a,r] ~ (independent) Uniform(low[a,r], high[a,r]).
    We assume each element is independent, so we do not assume
    any correlation structure:

    Parameters:
        lb (np.ndarray of shape (A,) or (A, R) or float):
            Array or scalar of lower bounds
        ub (np.ndarray of shape (A,) or (A, R) or float):
            Array or scalar of upper bounds
        RNG (np.random.Generator):
            Used to generate Uniform random variables.


    Returns:
        X (np.ndarray of shape (A,) or (A, R) or float):
            Random matrix or scalar realization where each
            element is independently sampled from a Uniform distribution
            with parameters given element-wise by `lb`
            and `ub` -- `X` is same shape as `lb` and `ub`.
    """

    # Use linear transformation of Uniform random variable!
    # Sample standard Uniforms ~[0,1] and apply transformation below
    #   to get Uniforms ~[low, high] element-wise :)

    if param_shape == "age":
        if (np.shape(lb) != (A,) or
                np.shape(ub) != (A,)):
            raise ValueError("With dependence on age, lower bounds and \n"
                             "upper bounds must be arrays of shape (A,). \n"
                             "Fix inputs and try again.")

            U = RNG.uniform(size=lb.shape)
            X = lb + (ub - lb) * U
    elif param_shape == "AR":
        if (np.shape(lb) != (A, R) or
                np.shape(ub) != (A, R)):
            raise ValueError("With dependence on age-risk, lower bounds and \n"
                             "upper bounds must be arrays of shape (A,R). \n"
                             "Fix inputs and try again.")
        U = RNG.uniform(size=lb.shape)
        X = lb + (ub - lb) * U
    elif param_shape == "scalar":
        if not np.isscalar(lb) or not np.isscalar(ub):
            raise ValueError("With dependence type scalar, lower bounds and \n"
                             "upper bounds must be scalars. Fix inputs and try again.")
        U = RNG.uniform()
        X = lb + (ub - lb) * U

    return X

sample_uniform_metapop_params(metapop_model: MetapopModel, sampling_RNG: np.random.Generator, sampling_info: dict[str, dict[str, UniformSamplingSpec]]) -> dict[str, dict[str, np.ndarray]]

Draw parameter realizations from uniform distributions for a metapopulation model.

Parameters:

Name	Type	Description	Default
metapop_model	MetapopModel	The metapop model whose subpopulation parameters are sampled.	required
sampling_RNG	Generator	Random number generator for Uniform sampling.	required
sampling_info	dict[str, dict[str, UniformSamplingSpec]]	Nested dictionary with sampling information. - Outer keys: Either "all_subpop" (apply to all subpopulations) or the name of a subpopulation, matching the name attribute of a SubpopModel in metapop_model. - Inner keys: Parameter names corresponding to attributes of the SubpopParams class associated with the subpop models in metapop_model. - Values: UniformSamplingSpec objects defining lower/upper bounds and shape of the parameter.	required

Returns:

Name	Type	Description
pending_param_updates	dict[str, dict[str, ndarray | float]]	Nested dictionary of sampled parameter values. - Outer keys: subpop names -- similar to description for outer keys of sampling_info argument. But unlike sampling_info, there is no "all_subpop" key here -- if a parameter applies to all subpopulations, the same sampled value appears under each subpopulation key. - Inner keys: parameter names -- same as description for inner keys of sampling_info argument. - Values: sampled parameters (scalar, 1D array, or 2D array) according to the shape specified in UniformSamplingSpec.param_shape.

Source code in CLT_BaseModel/clt_toolkit/sampling.py

def sample_uniform_metapop_params(metapop_model: MetapopModel,
                                  sampling_RNG: np.random.Generator,
                                  sampling_info: dict[str, dict[str, UniformSamplingSpec]]) \
        -> dict[str, dict[str, np.ndarray]]:
    """
    Draw parameter realizations from uniform distributions for a
    metapopulation model.

    Parameters:
        metapop_model (MetapopModel):
            The metapop model whose subpopulation parameters are sampled.
        sampling_RNG (np.random.Generator):
            Random number generator for Uniform sampling.
        sampling_info (dict[str, dict[str, UniformSamplingSpec]]):
            Nested dictionary with sampling information.
            - Outer keys:
                Either "all_subpop" (apply to all subpopulations)
                or the name of a subpopulation, matching the `name`
                attribute of a `SubpopModel` in `metapop_model`.
            - Inner keys:
                Parameter names corresponding to attributes of the
                `SubpopParams` class associated with the subpop models
                in `metapop_model`.
            - Values:
                `UniformSamplingSpec` objects defining lower/upper bounds
                and shape of the parameter.

    Returns:
        pending_param_updates (dict[str, dict[str, np.ndarray | float]]):
            Nested dictionary of sampled parameter values.
            - Outer keys: subpop names -- similar to description for
                outer keys of `sampling_info` argument. But unlike `sampling_info`,
                there is no `"all_subpop"` key here -- if a parameter applies to
                all subpopulations, the same sampled value appears under each
                subpopulation key.
            - Inner keys: parameter names -- same as description for
                inner keys of `sampling_info` argument.
            - Values: sampled parameters (scalar, 1D array, or 2D array) according to
            the shape specified in `UniformSamplingSpec.param_shape`.
    """

    # These dimensions should be the same across subpopulations
    #   -- so just grab the values from the 1st subpop model
    num_age_groups = metapop_model._subpop_models_ordered[0].params.num_age_groups
    num_risk_groups = metapop_model._subpop_models_ordered[0].params.num_risk_groups

    # We do not want to call `updated_dataclass` repeatedly
    #   when we update a single parameter field for each subpopulation,
    #   because this creates a NEW instance (since the dataclass
    #   is frozen and cannot be edited).
    # Instead, we to call `updated_dataclass` once for each
    #   subpop model.
    # So, for each subpop model, we save the parameters that
    #   need to be changed (to reflect the sampling outcomes)
    #   in a dictionary, hence the nested dictionaries.
    pending_param_updates = defaultdict(dict)  # {subpop_id: {param_name: new_value}}

    for subpop_name, params_dict in sampling_info.items():

        for param_name, param_spec in params_dict.items():
            sample = sample_uniform_matrix(param_spec.lower_bound,
                                           param_spec.upper_bound,
                                           sampling_RNG,
                                           num_age_groups,
                                           num_risk_groups,
                                           param_spec.param_shape)

            sample = np.round(sample, param_spec.num_decimals)

        if subpop_name == "all_subpop":  # sample is the same across all subpop models
            for subpop_id in metapop_model.subpop_models.keys():
                pending_param_updates[subpop_id][param_name] = sample
        else:
            pending_param_updates[subpop_name][param_name] = sample

    return pending_param_updates

serialize_dataclass(dc) -> dict

Convert a dataclass or dict to a JSON-serializable dictionary.

Parameters:

Name	Type	Description	Default
dc	obj | dict	The object to serialize. - If a dataclass, it will be converted using asdict(). - All numpy arrays are converted to lists. - Scalars (int, float, str, bool) remain unchanged. - Other objects are converted to strings as a fallback.	required

Returns:

Type	Description
dict	dict Dictionary representation of the object, fully JSON-serializable.

Source code in CLT_BaseModel/clt_toolkit/utils.py

def serialize_dataclass(dc) -> dict:
    """
    Convert a dataclass or dict to a JSON-serializable dictionary.

    Parameters:
        dc (obj | dict):
            The object to serialize.
            - If a dataclass, it will be converted using `asdict()`.
            - All numpy arrays are converted to lists.
            - Scalars (int, float, str, bool) remain unchanged.
            - Other objects are converted to strings as a fallback.

    Returns:
        dict
            Dictionary representation of the object, fully JSON-serializable.
    """

    if is_dataclass(dc):
        dc = asdict(dc)
    elif not isinstance(dc, dict):
        raise TypeError("Object must be a dataclass or dict.")

    return {k: serialize_value(v) for k, v in dc.items()}

serialize_value(value)

Convert a value into a JSON-serializable format.

value (any): The value to serialize. Supported types: - np.ndarray is converted to list - Scalars and None remain unchanged - dict, list, or tuple gets recursively serialized (i.e. in case it's a nested object, etc...) - Any other type is converted to str as a fallback

Returns:

Type	Description
	A version of the input that can be safely serialized to JSON.

Source code in CLT_BaseModel/clt_toolkit/utils.py

def serialize_value(value):
    """
    Convert a value into a JSON-serializable format.

    Parameters:
    value (any):
        The value to serialize. Supported types:
        - `np.ndarray` is converted to `list`
        - Scalars and `None` remain unchanged
        - `dict`, `list`, or `tuple` gets recursively serialized
            (i.e. in case it's a nested object, etc...)
        - Any other type is converted to `str` as a fallback

    Returns:
        A version of the input that can be safely serialized to JSON.
    """

    if isinstance(value, np.ndarray):
        return value.tolist()
    elif isinstance(value, (int, float, str, bool)) or value is None:
        return value
    elif isinstance(value, dict):
        return {k: serialize_value(v) for k, v in value.items()}
    elif isinstance(value, list) or isinstance(value, tuple):
        return [serialize_value(v) for v in value]
    else:
        return str(value)  # fallback for anything else

to_AR_array(x, A, R) -> np.ndarray

Convert scalar, 1D (A,) or 2D (A,R) to a (A,R) array.

Parameters:

Name	Type	Description	Default
x	float | ndarray	Float or array to convert to (A, R) array.	required
A	int	number of age groups.	required
R	int	number of risk groups.	required

Returns:

Type	Description
ndarray	(np.ndarray of shape (A, R))

Source code in CLT_BaseModel/clt_toolkit/utils.py

def to_AR_array(x, A, R) -> np.ndarray:
    """
    Convert scalar, 1D (A,) or 2D (A,R) to a (A,R) array.

    Params:
        x (float | np.ndarray):
            Float or array to convert to (A, R) array.
        A (int):
            number of age groups.
        R (int):
            number of risk groups.

    Returns:
        (np.ndarray of shape (A, R))
    """

    arr = np.asarray(x)

    if arr.ndim == 0:  # scalar
        return np.full((A, R), arr)

    elif arr.ndim == 1:  # shape (A,)
        if arr.shape[0] != A:
            raise ValueError(f"Expected length {A}, got {arr.shape[0]}.")
        return np.tile(arr[:, None], (1, R))  # expand to (A,R)

    elif arr.ndim == 2:  # shape (A,R)
        if arr.shape != (A, R):
            raise ValueError(f"Expected shape ({A},{R}), got {arr.shape}.")
        return arr

    else:
        raise ValueError(f"Unsupported array shape {arr.shape}")

updated_dataclass(original: dataclass, updates: dict) -> object

Return a new dataclass based on original, with fields in updates replaced/added.

Source code in CLT_BaseModel/clt_toolkit/utils.py

def updated_dataclass(original: dataclass,
                      updates: dict) -> object:

    """
    Return a new dataclass based on `original`, with fields in `updates` replaced/added.
    """

    return replace(original, **updates)

updated_dict(original: dict, updates: dict) -> dict

Return a new dictionary based on original, with keys in updates replaced/added.

Parameters:

Name	Type	Description	Default
original	dict	Original dictionary.	required
updates	dict	Dictionary of updates to apply.	required

Returns:

Type	Description
dict	New dictionary with updates applied.

Source code in CLT_BaseModel/clt_toolkit/utils.py

def updated_dict(original: dict,
                 updates: dict) -> dict:
    """
    Return a new dictionary based on `original`, with keys in `updates` replaced/added.

    Parameters:
        original (dict):
            Original dictionary.
        updates (dict):
            Dictionary of updates to apply.

    Returns:
        (dict):
            New dictionary with updates applied.
    """

    return {**original, **updates}