Stratifications

class vivarium.framework.results.stratification.Stratification(name, requires_attributes, categories, excluded_categories, mapper, is_vectorized=False)

Class for stratifying observed quantities by specified characteristics.

Each Stratification represents a set of mutually exclusive and collectively exhaustive categories into which simulants can be assigned. This class includes a stratify method that produces an output column by calling the mapper on the source columns.

Parameters:

• name (str)

• requires_attributes (list[str])

• categories (list[str])

• excluded_categories (list[str])

• mapper (Callable[[DataFrame], Series] | Callable[[Series], str] | None)

• is_vectorized (bool)

name: str

Name of the stratification.

requires_attributes: list[str]

The population attributes needed as input for the mapper.

categories: list[str]

Exhaustive list of all possible stratification values.

excluded_categories: list[str]

List of possible stratification values to exclude from results processing. If None (the default), will use exclusions as defined in the configuration.

mapper: Callable[[DataFrame], Series] | Callable[[Series], str] | None

A callable that maps the population attributes specified by the requires_attributes argument to the stratification categories. It can either map the entire population or an individual simulant. A simulation will fail if the mapper ever produces an invalid value.

is_vectorized: bool = False

True if the mapper function will map the entire population, and False if it will only map a single simulant.

stratify(population)

Applies the mapper to the population sources columns.

This creates a new Series to be added to the population. Any excluded_categories (which have already been removed from categories) will be converted to NaNs in the new column and dropped later at the observation level.

Parameters:

population (pd.DataFrame) – A DataFrame containing the data to be stratified.

Returns:

A Series containing the mapped values to be used for stratifying.

Raises:

ValueError – If the mapper returns any values not in categories or excluded_categories.

Return type:

pd.Series[CategoricalDtype]

vivarium.framework.results.stratification.get_mapped_col_name(col_name)

Returns a new column name to be used for mapped values

Return type:

str

Parameters:

col_name (str)

vivarium.framework.results.stratification.get_original_col_name(col_name)

Returns the original column name given a modified mapped column name.

Return type:

str

Parameters:

col_name (str)