vivarium-public-health 4.1.1__py3-none-any.whl → 4.2.0__py3-none-any.whl

vivarium_public_health/_version.py

@@ -1 +1 @@
-__version__ = "4.1.1"
+__version__ = "4.2.0"

vivarium_public_health/risks/implementations/low_birth_weight_and_short_gestation.py

@@ -14,6 +14,8 @@ from typing import Any
 
 import numpy as np
 import pandas as pd
+from layered_config_tree import ConfigurationError
+from loguru import logger
 from vivarium.framework.engine import Builder
 from vivarium.framework.lifecycle import LifeCycleError
 from vivarium.framework.population import SimulantData
@@ -21,9 +23,12 @@ from vivarium.framework.resource import Resource
 from vivarium.framework.values import Pipeline
 
 from vivarium_public_health.risks import Risk, RiskEffect
-from vivarium_public_health.risks.data_transformations import get_exposure_post_processor
+from vivarium_public_health.risks.data_transformations import (
+    get_exposure_post_processor,
+    pivot_categorical,
+)
 from vivarium_public_health.risks.distributions import PolytomousDistribution
-from vivarium_public_health.utilities import get_lookup_columns, to_snake_case
+from vivarium_public_health.utilities import EntityString, get_lookup_columns, to_snake_case
 
 CATEGORICAL = "categorical"
 BIRTH_WEIGHT = "birth_weight"
@@ -31,16 +36,76 @@ GESTATIONAL_AGE = "gestational_age"
 
 
 class LBWSGDistribution(PolytomousDistribution):
+    @property
+    def categories(self) -> list[str]:
+        # These need to be sorted so the cumulative sum is in the correct order of categories
+        # and results are therefore reproducible and correct
+        return sorted(self.lookup_tables[self.exposure_key].value_columns)
 
     #################
     # Setup methods #
     #################
 
+    def __init__(
+        self,
+        risk: EntityString,
+        distribution_type: str,
+        exposure_data: int | float | pd.DataFrame | None = None,
+    ) -> None:
+        super().__init__(risk, distribution_type, exposure_data)
+        self.exposure_key = "birth_exposure"
+
     # noinspection PyAttributeOutsideInit
     def setup(self, builder: Builder) -> None:
         super().setup(builder)
         self.category_intervals = self.get_category_intervals(builder)
 
+    def build_all_lookup_tables(self, builder: Builder) -> None:
+        try:
+            birth_exposure_data = self.get_data(
+                builder, self.configuration["data_sources"]["birth_exposure"]
+            )
+            birth_exposure_value_columns = self.get_exposure_value_columns(
+                birth_exposure_data
+            )
+
+            if isinstance(birth_exposure_data, pd.DataFrame):
+                birth_exposure_data = pivot_categorical(
+                    builder, self.risk, birth_exposure_data, "parameter"
+                )
+
+            self.lookup_tables["birth_exposure"] = self.build_lookup_table(
+                builder, birth_exposure_data, birth_exposure_value_columns
+            )
+        except ConfigurationError:
+            logger.warning("Birth exposure data for LBWSG is missing from the simulation")
+        try:
+            super().build_all_lookup_tables(builder)
+        except ConfigurationError:
+            logger.warning("The data for LBWSG exposure is missing from the simulation.")
+
+        if (
+            "birth_exposure" not in self.lookup_tables
+            and "exposure" not in self.lookup_tables
+        ):
+            raise ConfigurationError(
+                "The LBWSG distribution requires either 'birth_exposure' or 'exposure' data to be "
+                "available in the simulation."
+            )
+
+    def get_exposure_parameter_pipeline(self, builder: Builder) -> Pipeline:
+        lookup_columns = []
+        if "exposure" in self.lookup_tables:
+            lookup_columns.extend(get_lookup_columns([self.lookup_tables["exposure"]]))
+        if "birth_exposure" in self.lookup_tables:
+            lookup_columns.extend(get_lookup_columns([self.lookup_tables["birth_exposure"]]))
+        return builder.value.register_value_producer(
+            self.parameters_pipeline_name,
+            source=lambda index: self.lookup_tables[self.exposure_key](index),
+            component=self,
+            required_resources=list(set(lookup_columns)),
+        )
+
     def get_category_intervals(self, builder: Builder) -> dict[str, dict[str, pd.Interval]]:
         """Gets the intervals for each category.
 
@@ -203,8 +268,9 @@ class LBWSGRisk(Risk):
     @property
     def configuration_defaults(self) -> dict[str, Any]:
         configuration_defaults = super().configuration_defaults
+        # Add birth exposure data source
         configuration_defaults[self.name]["data_sources"][
-            "exposure"
+            "birth_exposure"
         ] = f"{self.risk}.birth_exposure"
         configuration_defaults[self.name]["distribution_type"] = "lbwsg"
         return configuration_defaults
@@ -224,6 +290,7 @@ class LBWSGRisk(Risk):
     def setup(self, builder: Builder) -> None:
         super().setup(builder)
         self.birth_exposures = self.get_birth_exposure_pipelines(builder)
+        self.configuration_age_end = builder.configuration.population.initialization_age_max
 
     #################
     # Setup methods #
@@ -242,7 +309,7 @@ class LBWSGRisk(Risk):
             self.exposure_distribution.lookup_tables.values()
         )
 
-        def get_pipeline(axis_: str):
+        def get_pipeline(axis_: str) -> Pipeline:
             return builder.value.register_value_producer(
                 self.birth_exposure_pipeline_name(axis_),
                 source=lambda index: self.get_birth_exposure(axis_, index),
@@ -260,12 +327,23 @@ class LBWSGRisk(Risk):
     ########################
 
     def on_initialize_simulants(self, pop_data: SimulantData) -> None:
-        birth_exposures = {
-            self.get_exposure_column_name(axis): self.birth_exposures[
-                self.birth_exposure_pipeline_name(axis)
-            ](pop_data.index)
-            for axis in self.AXES
-        }
+        if pop_data.user_data.get("age_end", self.configuration_age_end) == 0:
+            self.exposure_distribution.exposure_key = "birth_exposure"
+        else:
+            self.exposure_distribution.exposure_key = "exposure"
+
+        try:
+            birth_exposures = {
+                self.get_exposure_column_name(axis): self.birth_exposures[
+                    self.birth_exposure_pipeline_name(axis)
+                ](pop_data.index)
+                for axis in self.AXES
+            }
+        except KeyError:
+            raise ConfigurationError(
+                f"{self.exposure_distribution.exposure_key} data for {self.name} is missing from the "
+                "simulation. Simulants cannot be initialized."
+            )
         self.population_view.update(pd.DataFrame(birth_exposures))
 
     ##################################

{vivarium_public_health-4.1.1.dist-info → vivarium_public_health-4.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: vivarium_public_health
-Version: 4.1.1
+Version: 4.2.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-4.1.1.dist-info → vivarium_public_health-4.2.0.dist-info}/RECORD

@@ -1,6 +1,6 @@
 vivarium_public_health/__about__.py,sha256=RgWycPypKZS80TpSX7o41cREnG8PfguNHDHLuLyl820,487
 vivarium_public_health/__init__.py,sha256=GDeeP-7OlCBwPuv_xQoB1wNmvCaFsqfTB7qnnYApm0w,1343
-vivarium_public_health/_version.py,sha256=NtWqZUv_8plp1gndoYX_eAqe7K65oPMh-hIvB4Evdc8,22
+vivarium_public_health/_version.py,sha256=pU56ZDjCDnUbINe39hYuKXHy9UI9Tf4xjWgnxBlxEDk,22
 vivarium_public_health/utilities.py,sha256=QNXQ6fhAr1HcV-GwKw7wQLz6QyuNxqNvMA-XujKjTgs,3035
 vivarium_public_health/disease/__init__.py,sha256=VUJHDLlE6ngo2qHNQUtZ8OWH5H_T7_ao-xsYKDkRmHw,443
 vivarium_public_health/disease/exceptions.py,sha256=vb30IIV82OiDf2cNZCs_E2rF6mdDDHbnZSND60no5CU,97
@@ -38,13 +38,13 @@ vivarium_public_health/risks/data_transformations.py,sha256=SgdPKc95BBqgMNUdlAQM
 vivarium_public_health/risks/distributions.py,sha256=5gkurGTPqZczIKum5NgJ-DbJ2bxA1tRbLe7jKgCqkQg,18203
 vivarium_public_health/risks/effect.py,sha256=Cw-V1rMPaAPZsirbUElPzksvNuXLQc8TIIk0hgaFzvA,21259
 vivarium_public_health/risks/implementations/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vivarium_public_health/risks/implementations/low_birth_weight_and_short_gestation.py,sha256=gNQGmw-J_wZJm2r3P-4ix4Lj2BM0_YVzY9CC7hezmY8,17780
+vivarium_public_health/risks/implementations/low_birth_weight_and_short_gestation.py,sha256=wKVGHhdIas9WOrgA5wpwkSSDrZwBtaCqkZXfQH6jrkY,21159
 vivarium_public_health/treatment/__init__.py,sha256=wONElu9aJbBYwpYIovYPYaN_GYfVhPXtTeFWSdQMgA0,222
 vivarium_public_health/treatment/magic_wand.py,sha256=zg4I48G-l9fC6-qjvApbM1zNACJimZlX9ZZ9fdHKwvc,1985
 vivarium_public_health/treatment/scale_up.py,sha256=_bsf9c_yVc7-q_-yBcXimISTUfYzPps1auH0uEf7sfQ,7048
 vivarium_public_health/treatment/therapeutic_inertia.py,sha256=ZIHnpuszZwA_BkS53JbSLvpcnX_bqG2knwCUyUgkA9M,2362
-vivarium_public_health-4.1.1.dist-info/licenses/LICENSE.txt,sha256=mN4bNLUQNcN9njYRc_3jCZkfPySVpmM6MRps104FxA4,1548
-vivarium_public_health-4.1.1.dist-info/METADATA,sha256=FH3HU37jR-MEZ04qtBFZms8kDqdSop1Wc6wSFNoQVSg,4237
-vivarium_public_health-4.1.1.dist-info/WHEEL,sha256=zaaOINJESkSfm_4HQVc5ssNzHCPXhJm0kEUakpsEHaU,91
-vivarium_public_health-4.1.1.dist-info/top_level.txt,sha256=VVInlpzCFD0UNNhjOq_j-a29odzjwUwYFTGfvqbi4dY,23
-vivarium_public_health-4.1.1.dist-info/RECORD,,
+vivarium_public_health-4.2.0.dist-info/licenses/LICENSE.txt,sha256=mN4bNLUQNcN9njYRc_3jCZkfPySVpmM6MRps104FxA4,1548
+vivarium_public_health-4.2.0.dist-info/METADATA,sha256=o3j3So0A5hNLBNevTQG955g_nKQE5AaH3dXh67tzkUo,4237
+vivarium_public_health-4.2.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+vivarium_public_health-4.2.0.dist-info/top_level.txt,sha256=VVInlpzCFD0UNNhjOq_j-a29odzjwUwYFTGfvqbi4dY,23
+vivarium_public_health-4.2.0.dist-info/RECORD,,

{vivarium_public_health-4.1.1.dist-info → vivarium_public_health-4.2.0.dist-info}/WHEEL

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