vivarium-public-health 3.0.10__tar.gz → 3.0.11__tar.gz

{vivarium_public_health-3.0.10 → vivarium_public_health-3.0.11}/.github/workflows/build.yml

@@ -11,9 +11,15 @@ on:
     - cron: "0 8 * * *"
   
 jobs:
-  call_build:
+  build:
+    strategy:
+      matrix:
+        python-version: ["3.9", "3.10", "3.11"]
     uses: 
       ihmeuw/vivarium_build_utils/.github/workflows/build.yml@main
     with:
       dependencies: "layered_config_tree,vivarium"
-      use_mypy: false
+      skip_mypy: true
+    secrets:
+      NOTIFY_EMAIL: ${{ secrets.NOTIFY_EMAIL }}
+      NOTIFY_PASSWORD: ${{ secrets.NOTIFY_PASSWORD }}

{vivarium_public_health-3.0.10 → vivarium_public_health-3.0.11}/CHANGELOG.rst

@@ -1,3 +1,7 @@
+**3.0.11 - 11/07/24**
+
+ - Add ScaledPopulation component to population package
+
 **3.0.10 - 10/29/24**
 
  - Import subpackage classes to top-level namespace

{vivarium_public_health-3.0.10 → vivarium_public_health-3.0.11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vivarium_public_health
-Version: 3.0.10
+Version: 3.0.11
 Summary: Components for modelling diseases, risks, and interventions with ``vivarium``
 Home-page: https://github.com/ihmeuw/vivarium_public_health
 Author: The vivarium developers

{vivarium_public_health-3.0.10 → vivarium_public_health-3.0.11}/src/vivarium_public_health/__init__.py

@@ -30,6 +30,7 @@ from vivarium_public_health.population import (
     FertilityCrudeBirthRate,
     FertilityDeterministic,
     Mortality,
+    ScaledPopulation,
 )
 from vivarium_public_health.results import (
     CategoricalRiskObserver,

vivarium_public_health-3.0.11/src/vivarium_public_health/_version.py

@@ -0,0 +1 @@
+__version__ = "3.0.11"

{vivarium_public_health-3.0.10 → vivarium_public_health-3.0.11}/src/vivarium_public_health/population/__init__.py

@@ -3,5 +3,5 @@ from .add_new_birth_cohorts import (
     FertilityCrudeBirthRate,
     FertilityDeterministic,
 )
-from .base_population import BasePopulation, generate_population
+from .base_population import BasePopulation, ScaledPopulation, generate_population
 from .mortality import Mortality

{vivarium_public_health-3.0.10 → vivarium_public_health-3.0.11}/src/vivarium_public_health/population/base_population.py

@@ -8,6 +8,8 @@ characteristics to simulants.
 
 """
 
+from __future__ import annotations
+
 from typing import Callable, Dict, Iterable, List
 
 import numpy as np
@@ -76,12 +78,11 @@ class BasePopulation(Component):
         # Validate configuration for deprecated keys
         self._validate_config_for_deprecated_keys()
 
-        source_population_structure = load_population_structure(builder)
+        source_population_structure = self._load_population_structure(builder)
         self.demographic_proportions = assign_demographic_proportions(
             source_population_structure,
             include_sex=self.config.include_sex,
         )
-
         self.randomness = self.get_randomness_streams(builder)
         self.register_simulants = builder.randomness.register_simulants
 
@@ -138,7 +139,6 @@ class BasePopulation(Component):
         demographic_proportions = self.get_demographic_proportions_for_creation_time(
             self.demographic_proportions, pop_data.creation_time.year
         )
-
         self.population_view.update(
             generate_population(
                 simulant_ids=pop_data.index,
@@ -202,6 +202,50 @@ class BasePopulation(Component):
                 f"Public Health. Use the new key '{mapper[key]}' instead."
             )
 
+    def _load_population_structure(self, builder: Builder) -> pd.DataFrame:
+        return load_population_structure(builder)
+
+
+class ScaledPopulation(BasePopulation):
+    """This component is to be used in place of BasePopulation when all simulants are
+    a subset of the total population and need to be rescaled. The base population
+    structure is multiplied by a provided scaling factor. This scaling factor
+    can be a dataframe passed in or a string that corresponds to an artifact key.
+    If providing an artifact key, users can specify that in the configuration file.
+    For example:
+
+    .. code-block:: yaml
+
+    components:
+        vivarium_public_health:
+            population:
+                - ScaledPopulation("some.artifact.key")
+
+
+    """
+
+    def __init__(self, scaling_factor: str | pd.DataFrame):
+        super().__init__()
+        self.scaling_factor = scaling_factor
+        """Set a multiplicative scaling factor for the population structure."""
+
+    def _load_population_structure(self, builder: Builder) -> pd.DataFrame:
+        scaling_factor = self.get_data(builder, self.scaling_factor)
+        population_structure = load_population_structure(builder)
+        if not isinstance(scaling_factor, pd.DataFrame):
+            raise ValueError(
+                f"Scaling factor must be a pandas DataFrame. Provided value: {scaling_factor}"
+            )
+        scaling_factor = scaling_factor.set_index(
+            [col for col in scaling_factor.columns if col != "value"]
+        )
+        population_structure = population_structure.set_index(
+            [col for col in population_structure.columns if col != "value"]
+        )
+        scaled_population_structure = (population_structure * scaling_factor).reset_index()
+
+        return scaled_population_structure
+
 
 class AgeOutSimulants(Component):
     """Component for handling aged-out simulants"""

{vivarium_public_health-3.0.10 → vivarium_public_health-3.0.11}/src/vivarium_public_health/population/data_transformations.py

@@ -13,6 +13,7 @@ from typing import Tuple, Union
 
 import numpy as np
 import pandas as pd
+from vivarium.framework.engine import Builder
 from vivarium.framework.randomness import RandomnessStream
 
 _SORT_ORDER = ["location", "year_start", "sex", "age_start"]
@@ -493,11 +494,12 @@ def get_cause_deleted_mortality_rate(all_cause_mortality_rate, list_of_csmrs):
     )
 
 
-def load_population_structure(builder):
+def load_population_structure(builder: Builder) -> pd.DataFrame:
     data = builder.data.load("population.structure")
     # create an age column which is the midpoint of the age group
     data["age"] = data.apply(lambda row: (row["age_start"] + row["age_end"]) / 2, axis=1)
     data["location"] = builder.data.load("population.location")
+
     return data
 
 

{vivarium_public_health-3.0.10 → vivarium_public_health-3.0.11}/src/vivarium_public_health.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vivarium_public_health
-Version: 3.0.10
+Version: 3.0.11
 Summary: Components for modelling diseases, risks, and interventions with ``vivarium``
 Home-page: https://github.com/ihmeuw/vivarium_public_health
 Author: The vivarium developers

{vivarium_public_health-3.0.10 → vivarium_public_health-3.0.11}/tests/population/test_add_new_birth_cohort.py

@@ -130,7 +130,7 @@ def test_FertilityCrudeBirthRate_extrapolate(base_config, base_plugins):
         }
     )
     pop_size = base_config.population.population_size
-    true_pop_size = 50000  # What's available in the mock artifact
+    true_pop_size = 25000  # What's available in the mock artifact
     live_births_by_sex = 500
     components = [TestPopulation(), FertilityCrudeBirthRate()]
 

{vivarium_public_health-3.0.10 → vivarium_public_health-3.0.11}/tests/population/test_base_population.py

@@ -1,15 +1,16 @@
 import math
-from pathlib import Path
 
 import numpy as np
 import pandas as pd
 import pytest
 from vivarium import InteractiveContext
 from vivarium.testing_utilities import get_randomness
+from vivarium_testing_utils import FuzzyChecker
 
 import vivarium_public_health.population.base_population as bp
 import vivarium_public_health.population.data_transformations as dt
-from tests.test_utilities import make_uniform_pop_data
+from tests.mock_artifact import MockArtifact
+from tests.test_utilities import make_uniform_pop_data, simple_pop_structure
 from vivarium_public_health import utilities
 
 
@@ -283,6 +284,71 @@ def test__assign_demography_with_age_bounds_error(base_simulants, include_sex):
         )
 
 
+@pytest.mark.parametrize("constructor_type", ["string", "data"])
+def test_scaled_population(
+    constructor_type, config, base_plugins, mocker, fuzzy_checker: FuzzyChecker
+):
+    config.update(
+        {
+            "population": {
+                "population_size": 1_000_000,
+                "include_sex": "Both",
+            },
+            "time": {"step_size": 1},
+        },
+        layer="override",
+    )
+
+    # Simple pop data
+    pop_structure = simple_pop_structure()
+    # Simple scalar data to pass to ScaledPopulation
+    scalar_data = simple_pop_structure().drop(columns=["location"])
+    scalar_values = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]
+    scalar_data["value"] = scalar_values
+
+    # Add data to artifact and mock return for plugin
+    mock_art = MockArtifact()
+    mock_art.write("population.structure", pop_structure)
+    mocker.patch(
+        "tests.mock_artifact.MockArtifactManager._load_artifact",
+    ).return_value = mock_art
+
+    if constructor_type == "string":
+        mock_art.write("population.scalar", scalar_data)
+        scaling_factor = "population.scalar"
+    else:
+        scaling_factor = scalar_data
+    scaled_pop = bp.ScaledPopulation(scaling_factor)
+    sim = InteractiveContext(
+        components=[scaled_pop], configuration=config, plugin_configuration=base_plugins
+    )
+    pop = sim.get_population()
+    # Use FuzzyChecker to compare population structure to demographic proportion by
+    # iterating through each age_group/sex combination
+    scaled_structure = pop_structure.copy()
+    scaled_structure["value"] = scaled_structure["value"] * scalar_data["value"]
+
+    for row in range(len(scaled_structure)):
+        row_data = scaled_structure.iloc[row]
+        age_start = row_data["age_start"]
+        sex = row_data["sex"]
+        # Get proportion of each age group
+        target_proportion = row_data["value"] / scaled_structure["value"].sum()
+        number_of_sims = len(
+            pop.loc[
+                (pop["age"] >= age_start)
+                & (pop["age"] <= row_data["age_end"])
+                & (pop["sex"] == sex)
+            ]
+        )
+        fuzzy_checker.fuzzy_assert_proportion(
+            observed_numerator=number_of_sims,
+            observed_denominator=len(pop),
+            target_proportion=target_proportion,
+            name=f"scaled_pop_proportion_check_{sex}_{age_start}",
+        )
+
+
 def _check_population(simulants, initial_age, step_size, include_sex):
     assert len(simulants) == len(simulants.age.unique())
     assert simulants.age.min() > initial_age

{vivarium_public_health-3.0.10 → vivarium_public_health-3.0.11}/tests/population/test_data_transformations.py

@@ -127,7 +127,7 @@ def test_smooth_ages(pop_data, include_sex):
             "sex": (
                 ["Male", "Female"] * 10000 if include_sex == "Both" else [include_sex] * 20000
             ),
-            "location": [1, 2] * 10000,
+            "location": [1] * 20000,
         }
     )
     randomness = get_randomness()

{vivarium_public_health-3.0.10 → vivarium_public_health-3.0.11}/tests/test_utilities.py

@@ -64,7 +64,7 @@ def make_uniform_pop_data(age_bin_midpoint=False):
     age_bins = [(b.age_start, b.age_end) for b in make_age_bins().itertuples()]
     sexes = ("Male", "Female")
     years = zip(range(1990, 2018), range(1991, 2019))
-    locations = (1, 2)
+    locations = list(range(1, 2))
 
     age_bins, sexes, years, locations = zip(*product(age_bins, sexes, years, locations))
     mins, maxes = zip(*age_bins)
@@ -116,3 +116,39 @@ def make_age_bins():
         names=["age_start", "age_end", "age_group_name"],
     )
     return pd.DataFrame(index=idx).reset_index()
+
+
+def simple_pop_structure() -> pd.DataFrame:
+    # Create simple population structure
+    age_idx = pd.MultiIndex.from_tuples(
+        [
+            (0.0, 25.0, "Young People"),
+            (25.0, 50.0, "Old People"),
+            (50.0, 75.0, "Ancient People"),
+            (75.0, 100.0, "People Who Beat the Odds"),
+        ],
+        names=["age_start", "age_end", "age_group_name"],
+    )
+    age_df = pd.DataFrame(index=age_idx).reset_index()
+    age_bins = [(group.age_start, group.age_end) for group in age_df.itertuples()]
+    sexes = ("Male", "Female")
+    location = ["Kenya"]
+    years = zip(range(2021, 2022), range(2022, 2023))
+
+    age_bins, sexes, years, location = zip(*product(age_bins, sexes, years, location))
+    mins, maxes = zip(*age_bins)
+    year_starts, year_ends = zip(*years)
+
+    pop_structure = pd.DataFrame(
+        {
+            "age_start": mins,
+            "age_end": maxes,
+            "sex": sexes,
+            "year_start": year_starts,
+            "year_end": year_ends,
+            "location": location,
+            # We have 8 demographic groups (4 age groups * 2 sexes)
+            "value": [1.0, 1.0, 2.0, 2.0, 3.0, 3.0, 4.0, 4.0],
+        }
+    )
+    return pop_structure

vivarium_public_health-3.0.10/src/vivarium_public_health/_version.py

@@ -1 +0,0 @@
-__version__ = "3.0.10"