vivarium-public-health 3.0.10__py3-none-any.whl → 3.1.0__py3-none-any.whl

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/_version.py

@@ -1 +1 @@
-__version__ = "3.0.10"
+__version__ = "3.1.0"

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/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/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.dist-info → vivarium_public_health-3.1.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vivarium_public_health
-Version: 3.0.10
+Version: 3.1.0
 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.dist-info → vivarium_public_health-3.1.0.dist-info}/RECORD

@@ -1,6 +1,6 @@
 vivarium_public_health/__about__.py,sha256=RgWycPypKZS80TpSX7o41cREnG8PfguNHDHLuLyl820,487
-vivarium_public_health/__init__.py,sha256=jYSQDGc86QMEW4iZ7wO1JCSP189n1zScIn6QKG470Kw,1321
-vivarium_public_health/_version.py,sha256=xS-MdV9u-8rYsWMV1VgxiWeHHmEWb63Js1BsRGmz0xo,23
+vivarium_public_health/__init__.py,sha256=GDeeP-7OlCBwPuv_xQoB1wNmvCaFsqfTB7qnnYApm0w,1343
+vivarium_public_health/_version.py,sha256=YVoF76lT0p3dIsqphNnDWuqSia3gZP1S1eQYXZ9FbSE,22
 vivarium_public_health/utilities.py,sha256=5cl9jjVkOQ1UeXT4DjDMAhaBNNjAsDo-SVJwpv6FDw0,3071
 vivarium_public_health/disease/__init__.py,sha256=VUJHDLlE6ngo2qHNQUtZ8OWH5H_T7_ao-xsYKDkRmHw,443
 vivarium_public_health/disease/model.py,sha256=0WIYDEx-hwlUJp6Zl8m8bUMoWxuVkOWsJvh_YlZiOPs,8234
@@ -17,10 +17,10 @@ vivarium_public_health/mslt/observer.py,sha256=O4rysQzAGE5oDkdXb0E-qjD9TPFphQHTn
 vivarium_public_health/mslt/population.py,sha256=v_p5VkjndAVJMuXaJQc3lAdzUWHlWCEQWH4A-c4phPA,7255
 vivarium_public_health/plugins/__init__.py,sha256=oBW_zfgG_LbwfgTDjUe0btfy9FaDvAbtXho1zQFnz0Y,76
 vivarium_public_health/plugins/parser.py,sha256=v78mj8awpdrB-oqK8udPI_7MZBChoKJOQN_e17fNEj8,31841
-vivarium_public_health/population/__init__.py,sha256=17rtbcNVK5LtCCxAex7P7Q_vYpwbeTepyf3nazS90Yc,225
+vivarium_public_health/population/__init__.py,sha256=x_9rB93q64Krw-kbBDI1-_U_JsPO1_QrL03AwiFlwXI,243
 vivarium_public_health/population/add_new_birth_cohorts.py,sha256=k65Li0LYWl-JFHBUvLjJxkRv12EJw_FVxrOYgbd44q8,9078
-vivarium_public_health/population/base_population.py,sha256=4lUc8EZwzj5Ba36lSmW9yyxcRuBSMLqi_8Fy69ssq5E,17026
-vivarium_public_health/population/data_transformations.py,sha256=QVh_63Wwg9BUkaQm1pMSvBb-wsYrsgyADKIERAiEOVg,22188
+vivarium_public_health/population/base_population.py,sha256=jJMVWv_EO0ckUZHhh01dRbsRgL491Eg3thRxPwNiAeg,18866
+vivarium_public_health/population/data_transformations.py,sha256=YmqyrlrIBtHGp1nFyhesqlNryvB7Vr33eVu4fU4HWsA,22260
 vivarium_public_health/population/mortality.py,sha256=w7b_TUssHjRcnULdXu7MXKfZBjCrlYWbB94oO3JWogI,10264
 vivarium_public_health/results/__init__.py,sha256=rKUZGlRXJgEyFY4a_WJeg3XnC0l34S5guYZ0N9JJS4E,319
 vivarium_public_health/results/columns.py,sha256=V-L3JgTcsk51Zx9PcUwSgaE1iZjuGyfZ8aShPjynadU,495
@@ -42,8 +42,8 @@ vivarium_public_health/treatment/__init__.py,sha256=wONElu9aJbBYwpYIovYPYaN_GYfV
 vivarium_public_health/treatment/magic_wand.py,sha256=i9N57-MEuQv5B6dQ5iVMTAdOPghYcgiRRz-dTzigf1s,1980
 vivarium_public_health/treatment/scale_up.py,sha256=9Cl3_MaCNPtUPPKXf3hWYepIwMCFlFY24jt2jeyFQO4,7006
 vivarium_public_health/treatment/therapeutic_inertia.py,sha256=8Z97s7GfcpfLu1U1ESJSqeEk4L__a3M0GbBV21MFg2s,2346
-vivarium_public_health-3.0.10.dist-info/LICENSE.txt,sha256=mN4bNLUQNcN9njYRc_3jCZkfPySVpmM6MRps104FxA4,1548
-vivarium_public_health-3.0.10.dist-info/METADATA,sha256=pPeIbY7zeBd-vNNaZMJQL51VeEZ1GanLuxYpDs24ZtU,4062
-vivarium_public_health-3.0.10.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-vivarium_public_health-3.0.10.dist-info/top_level.txt,sha256=VVInlpzCFD0UNNhjOq_j-a29odzjwUwYFTGfvqbi4dY,23
-vivarium_public_health-3.0.10.dist-info/RECORD,,
+vivarium_public_health-3.1.0.dist-info/LICENSE.txt,sha256=mN4bNLUQNcN9njYRc_3jCZkfPySVpmM6MRps104FxA4,1548
+vivarium_public_health-3.1.0.dist-info/METADATA,sha256=YGem2GjFv7jkEz7_3VkJMsJjMtHolbwiFgBgL0LF4Mk,4061
+vivarium_public_health-3.1.0.dist-info/WHEEL,sha256=a7TGlA-5DaHMRrarXjVbQagU3Man_dCnGIWMJr5kRWo,91
+vivarium_public_health-3.1.0.dist-info/top_level.txt,sha256=VVInlpzCFD0UNNhjOq_j-a29odzjwUwYFTGfvqbi4dY,23
+vivarium_public_health-3.1.0.dist-info/RECORD,,

{vivarium_public_health-3.0.10.dist-info → vivarium_public_health-3.1.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.3.0)
+Generator: setuptools (75.4.0)
 Root-Is-Purelib: true
 Tag: py3-none-any