vivarium-public-health 2.0.2__py3-none-any.whl → 2.1.1__py3-none-any.whl

vivarium_public_health/_version.py

@@ -1 +1 @@
-__version__ = "2.0.2"
+__version__ = "2.1.1"

vivarium_public_health/disease/state.py

@@ -11,10 +11,11 @@ from typing import Callable, Dict, List, Optional
 import numpy as np
 import pandas as pd
 from vivarium.framework.engine import Builder
-from vivarium.framework.lookup import LookupTableData
+from vivarium.framework.lookup import LookupTable, LookupTableData
 from vivarium.framework.population import PopulationView, SimulantData
+from vivarium.framework.randomness import RandomnessStream
 from vivarium.framework.state_machine import State, Transient, Transition, Trigger
-from vivarium.framework.values import list_combiner, union_post_processor
+from vivarium.framework.values import Pipeline, list_combiner, union_post_processor
 
 from vivarium_public_health.disease.transition import (
     ProportionTransition,
@@ -323,65 +324,41 @@ class DiseaseState(BaseDiseaseState):
         self.clock = builder.time.clock()
 
         prevalence_data = self.load_prevalence_data(builder)
-        self.prevalence = builder.lookup.build_table(
-            prevalence_data, key_columns=["sex"], parameter_columns=["age", "year"]
-        )
+        self.prevalence = self.get_prevalence(builder, prevalence_data)
 
         birth_prevalence_data = self.load_birth_prevalence_data(builder)
-        self.birth_prevalence = builder.lookup.build_table(
-            birth_prevalence_data, key_columns=["sex"], parameter_columns=["year"]
-        )
+        self.birth_prevalence = self.get_birth_prevalence(builder, birth_prevalence_data)
 
         dwell_time_data = self.load_dwell_time_data(builder)
-        self.dwell_time = builder.value.register_value_producer(
-            f"{self.state_id}.dwell_time",
-            source=builder.lookup.build_table(
-                dwell_time_data, key_columns=["sex"], parameter_columns=["age", "year"]
-            ),
-            requires_columns=["age", "sex"],
-        )
+        self.dwell_time = self.get_dwell_time_pipeline(builder, dwell_time_data)
 
         disability_weight_data = self.load_disability_weight_data(builder)
         self.has_disability = is_non_zero(disability_weight_data)
-        self.base_disability_weight = builder.lookup.build_table(
-            disability_weight_data, key_columns=["sex"], parameter_columns=["age", "year"]
-        )
-        self.disability_weight = builder.value.register_value_producer(
-            f"{self.state_id}.disability_weight",
-            source=self.compute_disability_weight,
-            requires_columns=["age", "sex", "alive", self.model],
+        self.base_disability_weight = self.get_base_disability_weight(
+            builder, disability_weight_data
         )
+
+        self.disability_weight = self.get_disability_weight_pipeline(builder)
+
         builder.value.register_value_modifier(
             "disability_weight", modifier=self.disability_weight
         )
 
         excess_mortality_data = self.load_excess_mortality_rate_data(builder)
         self.has_excess_mortality = is_non_zero(excess_mortality_data)
-        self.base_excess_mortality_rate = builder.lookup.build_table(
-            excess_mortality_data, key_columns=["sex"], parameter_columns=["age", "year"]
-        )
-        self.excess_mortality_rate = builder.value.register_rate_producer(
-            self.excess_mortality_rate_pipeline_name,
-            source=self.compute_excess_mortality_rate,
-            requires_columns=["age", "sex", "alive", self.model],
-            requires_values=[self.excess_mortality_rate_paf_pipeline_name],
-        )
-        paf = builder.lookup.build_table(0)
-        self.joint_paf = builder.value.register_value_producer(
-            self.excess_mortality_rate_paf_pipeline_name,
-            source=lambda idx: [paf(idx)],
-            preferred_combiner=list_combiner,
-            preferred_post_processor=union_post_processor,
+        self.base_excess_mortality_rate = self.get_base_excess_mortality_rate(
+            builder, excess_mortality_data
         )
+        self.excess_mortality_rate = self.get_excess_mortality_rate_pipeline(builder)
+        self.joint_paf = self.get_joint_paf(builder)
+
         builder.value.register_value_modifier(
             "mortality_rate",
             modifier=self.adjust_mortality_rate,
             requires_values=[self.excess_mortality_rate_pipeline_name],
         )
 
-        self.randomness_prevalence = builder.randomness.get_stream(
-            f"{self.state_id}_prevalent_cases"
-        )
+        self.randomness_prevalence = self.get_randomness_prevalence(builder)
 
     #################
     # Setup methods #
@@ -393,12 +370,26 @@ class DiseaseState(BaseDiseaseState):
         else:
             return builder.data.load(f"{self.cause_type}.{self.state_id}.prevalence")
 
+    def get_prevalence(
+        self, builder: Builder, prevalence_data: LookupTableData
+    ) -> LookupTable:
+        return builder.lookup.build_table(
+            prevalence_data, key_columns=["sex"], parameter_columns=["age", "year"]
+        )
+
     def load_birth_prevalence_data(self, builder: Builder) -> LookupTableData:
         if "birth_prevalence" in self._get_data_functions:
             return self._get_data_functions["birth_prevalence"](builder, self.state_id)
         else:
             return 0
 
+    def get_birth_prevalence(
+        self, builder: Builder, birth_prevalence_data: LookupTableData
+    ) -> LookupTable:
+        return builder.lookup.build_table(
+            birth_prevalence_data, key_columns=["sex"], parameter_columns=["year"]
+        )
+
     def load_dwell_time_data(self, builder: Builder) -> LookupTableData:
         if "dwell_time" in self._get_data_functions:
             dwell_time = self._get_data_functions["dwell_time"](builder, self.state_id)
@@ -414,6 +405,17 @@ class DiseaseState(BaseDiseaseState):
 
         return dwell_time
 
+    def get_dwell_time_pipeline(
+        self, builder: Builder, dwell_time_data: LookupTableData
+    ) -> Pipeline:
+        return builder.value.register_value_producer(
+            f"{self.state_id}.dwell_time",
+            source=builder.lookup.build_table(
+                dwell_time_data, key_columns=["sex"], parameter_columns=["age", "year"]
+            ),
+            requires_columns=["age", "sex"],
+        )
+
     def load_disability_weight_data(self, builder: Builder) -> LookupTableData:
         if "disability_weight" in self._get_data_functions:
             disability_weight = self._get_data_functions["disability_weight"](
@@ -429,6 +431,20 @@ class DiseaseState(BaseDiseaseState):
 
         return disability_weight
 
+    def get_base_disability_weight(
+        self, builder: Builder, disability_weight_data: LookupTableData
+    ) -> LookupTable:
+        return builder.lookup.build_table(
+            disability_weight_data, key_columns=["sex"], parameter_columns=["age", "year"]
+        )
+
+    def get_disability_weight_pipeline(self, builder: Builder) -> Pipeline:
+        return builder.value.register_value_producer(
+            f"{self.state_id}.disability_weight",
+            source=self.compute_disability_weight,
+            requires_columns=["age", "sex", "alive", self.model],
+        )
+
     def load_excess_mortality_rate_data(self, builder: Builder) -> LookupTableData:
         if "excess_mortality_rate" in self._get_data_functions:
             return self._get_data_functions["excess_mortality_rate"](builder, self.state_id)
@@ -439,6 +455,33 @@ class DiseaseState(BaseDiseaseState):
                 f"{self.cause_type}.{self.state_id}.excess_mortality_rate"
             )
 
+    def get_base_excess_mortality_rate(
+        self, builder: Builder, excess_mortality_data: LookupTableData
+    ) -> LookupTable:
+        return builder.lookup.build_table(
+            excess_mortality_data, key_columns=["sex"], parameter_columns=["age", "year"]
+        )
+
+    def get_excess_mortality_rate_pipeline(self, builder: Builder) -> Pipeline:
+        return builder.value.register_rate_producer(
+            self.excess_mortality_rate_pipeline_name,
+            source=self.compute_excess_mortality_rate,
+            requires_columns=["age", "sex", "alive", self.model],
+            requires_values=[self.excess_mortality_rate_paf_pipeline_name],
+        )
+
+    def get_joint_paf(self, builder: Builder) -> Pipeline:
+        paf = builder.lookup.build_table(0)
+        return builder.value.register_value_producer(
+            self.excess_mortality_rate_paf_pipeline_name,
+            source=lambda idx: [paf(idx)],
+            preferred_combiner=list_combiner,
+            preferred_post_processor=union_post_processor,
+        )
+
+    def get_randomness_prevalence(self, builder: Builder) -> RandomnessStream:
+        return builder.randomness.get_stream(f"{self.state_id}_prevalent_cases")
+
     ##################
     # Public methods #
     ##################

{vivarium_public_health-2.0.2.dist-info → vivarium_public_health-2.1.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vivarium-public-health
-Version: 2.0.2
+Version: 2.1.1
 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-2.0.2.dist-info → vivarium_public_health-2.1.1.dist-info}/RECORD

@@ -1,12 +1,12 @@
 vivarium_public_health/__about__.py,sha256=RgWycPypKZS80TpSX7o41cREnG8PfguNHDHLuLyl820,487
 vivarium_public_health/__init__.py,sha256=tomMOl3PI7O8GdxDWGBiBjT0Bwd31GpyQTYTzwIv108,361
-vivarium_public_health/_version.py,sha256=tATvJM5shAzfspHYjdVwpV2w3-gDA119NlEYi5X2lFY,22
+vivarium_public_health/_version.py,sha256=zPJIgPGcoSNiD0qme18OnYJYE3A9VVytlhO-V5DaAW0,22
 vivarium_public_health/utilities.py,sha256=_X9sZQ7flsi2sVWQ9zrf8GJw8QwZsPZm3NUjx1gu7bM,2555
 vivarium_public_health/disease/__init__.py,sha256=RuuiRcvAJfX9WQGt_WZZjxN7Cu3E5rMTmuaRS-UaFPM,419
 vivarium_public_health/disease/model.py,sha256=7xFGo6JqPjQjm2VUZ3u3ThXWSzmitTOCZ8N9PTTL6MU,8253
 vivarium_public_health/disease/models.py,sha256=uiB2qUlxBsPPPmHJ8Cgot_T1ItZ8RYSNVOBtxtn93Y0,3478
 vivarium_public_health/disease/special_disease.py,sha256=gl8aK0z6afCxiCZxgJafLe4xmbR91zk3079hsi2pUAw,14751
-vivarium_public_health/disease/state.py,sha256=yZ76I4mRHfa-KG2lhCJZpVWkgxip0j6r3AkQD5SpTzg,21165
+vivarium_public_health/disease/state.py,sha256=V_X3DPioZJ3SlR8Rp4o8fkdyeBPmEDDjLYOBq6JOcqA,22778
 vivarium_public_health/disease/transition.py,sha256=Or5nucRzeGG-UuE_CGkPZ9qE35-Ed9I9LWHj4rjknCc,5334
 vivarium_public_health/metrics/__init__.py,sha256=bWAvvdCm_7RPIazo12qFohA2x5-_EV6ceV8IhKS37sk,209
 vivarium_public_health/metrics/disability.py,sha256=zm0vAG00wj44CHjYGdT2_pebgARa3XXIerrR06t80rc,3984
@@ -40,8 +40,8 @@ vivarium_public_health/treatment/__init__.py,sha256=wONElu9aJbBYwpYIovYPYaN_GYfV
 vivarium_public_health/treatment/magic_wand.py,sha256=iPKFN3VjfiMy_XvN94UqM-FUrGuI0ULwmOdAGdOepYQ,1979
 vivarium_public_health/treatment/scale_up.py,sha256=zTPMfKG38dGBj14QpBeLlmUixx1y_XIMElmTpkklydA,6714
 vivarium_public_health/treatment/therapeutic_inertia.py,sha256=VwZ7t90zzfGoBusduIvcE4lDe5zTvzmHiUNB3u2I52Y,2339
-vivarium_public_health-2.0.2.dist-info/LICENSE.txt,sha256=mN4bNLUQNcN9njYRc_3jCZkfPySVpmM6MRps104FxA4,1548
-vivarium_public_health-2.0.2.dist-info/METADATA,sha256=nbj5JXxhmRd7lWBWvUFvHz8iInnQEFzy823NmkmSn8M,3430
-vivarium_public_health-2.0.2.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
-vivarium_public_health-2.0.2.dist-info/top_level.txt,sha256=VVInlpzCFD0UNNhjOq_j-a29odzjwUwYFTGfvqbi4dY,23
-vivarium_public_health-2.0.2.dist-info/RECORD,,
+vivarium_public_health-2.1.1.dist-info/LICENSE.txt,sha256=mN4bNLUQNcN9njYRc_3jCZkfPySVpmM6MRps104FxA4,1548
+vivarium_public_health-2.1.1.dist-info/METADATA,sha256=QrLx-cPWz8E_pgj_6ZG3qO0y5KiywQsqVKSBdviNZ3A,3430
+vivarium_public_health-2.1.1.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
+vivarium_public_health-2.1.1.dist-info/top_level.txt,sha256=VVInlpzCFD0UNNhjOq_j-a29odzjwUwYFTGfvqbi4dY,23
+vivarium_public_health-2.1.1.dist-info/RECORD,,