Abie enhanced consistent rates

artifact_requirements.txt

@@ -22,4 +22,3 @@ matplotlib
 jax
 numpyro
 diffrax
-interpax

src/vivarium_csu_alzheimers/components/alzheimers.py

@@ -20,6 +20,7 @@
     DW_BBBM,
     EMR_BBBM,
     EMR_MCI,
+    MCI_AVG_DURATION,
 )
 from vivarium_csu_alzheimers.constants.models import ALZHEIMERS_DISEASE_MODEL
 from vivarium_csu_alzheimers.data.mci_hazard import hazard
@@ -172,39 +173,70 @@ def _create_disease_model(self) -> AlzheimersModel:
             ),
             excess_mortality_rate=EMR_MCI,
         )
-        alzheimers_state = DiseaseState(
-            ALZHEIMERS_DISEASE_MODEL.ALZHEIMERS_DISEASE_STATE,
+        mild_dementia_state = DiseaseState(
+            ALZHEIMERS_DISEASE_MODEL.MILD_DEMENTIA_STATE,
+            allow_self_transition=True,
+            birth_prevalence=0.0,
+            prevalence=lambda builder: builder.data.load(
+                ALZHEIMERS.MILD_DEMENTIA_CONDITIONAL_PREVALENCE
+            ),
+            disability_weight=lambda builder: builder.data.load(
+                ALZHEIMERS.MILD_DEMENTIA_DISABILITY_WEIGHT
+            ),
+            excess_mortality_rate=lambda builder: builder.data.load(ALZHEIMERS.EMR_MILD),
+        )
+        moderate_dementia_state = DiseaseState(
+            ALZHEIMERS_DISEASE_MODEL.MODERATE_DEMENTIA_STATE,
+            allow_self_transition=True,
+            birth_prevalence=0.0,
+            prevalence=lambda builder: builder.data.load(
+                ALZHEIMERS.MODERATE_DEMENTIA_CONDITIONAL_PREVALENCE
+            ),
+            disability_weight=lambda builder: builder.data.load(
+                ALZHEIMERS.MODERATE_DEMENTIA_DISABILITY_WEIGHT
+            ),
+            excess_mortality_rate=lambda builder: builder.data.load(ALZHEIMERS.EMR_MODERATE),
+        )
+        severe_dementia_state = DiseaseState(
+            ALZHEIMERS_DISEASE_MODEL.SEVERE_DEMENTIA_STATE,
+            allow_self_transition=True,
             birth_prevalence=0.0,
             prevalence=lambda builder: builder.data.load(
-                ALZHEIMERS.DEMENTIA_CONDITIONAL_PREVALENCE
+                ALZHEIMERS.SEVERE_DEMENTIA_CONDITIONAL_PREVALENCE
             ),
-            disability_weight=lambda builder: builder.data.load(ALZHEIMERS.DISABILITY_WEIGHT),
-            excess_mortality_rate=lambda builder: builder.data.load(ALZHEIMERS.EMR),
+            disability_weight=lambda builder: builder.data.load(
+                ALZHEIMERS.SEVERE_DEMENTIA_DISABILITY_WEIGHT
+            ),
+            excess_mortality_rate=lambda builder: builder.data.load(ALZHEIMERS.EMR_SEVERE),
         )
 
-        # Add transitions between states
+        # AD progression transitions
         bbbm_state.add_bbbm_transition(output=mci_state)
         mci_state.add_rate_transition(
-            output=alzheimers_state,
+            output=mild_dementia_state,
+            transition_rate=1 / MCI_AVG_DURATION,
+        )
+        mild_dementia_state.add_rate_transition(
+            output=moderate_dementia_state,
             transition_rate=lambda builder: builder.data.load(
-                ALZHEIMERS.MCI_TO_DEMENTIA_TRANSITION_RATE
+                ALZHEIMERS.MILD_TO_MODERATE_DEMENTIA_TRANSITION_RATE
+            ),
+        )
+        moderate_dementia_state.add_rate_transition(
+            output=severe_dementia_state,
+            transition_rate=lambda builder: builder.data.load(
+                ALZHEIMERS.MODERATE_TO_SEVERE_DEMENTIA_TRANSITION_RATE
             ),
         )
 
         return AlzheimersModel(
             ALZHEIMERS_DISEASE_MODEL.NAME,
             initial_state=bbbm_state,
-            states=[bbbm_state, mci_state, alzheimers_state],
+            states=[
+                bbbm_state,
+                mci_state,
+                mild_dementia_state,
+                moderate_dementia_state,
+                severe_dementia_state,
+            ],
         )
-
-    def _get_alzheimers_disease_state_prevalence(self, builder: Builder) -> pd.DataFrame:
-        """Get the Alzheimer's disease state prevalence table."""
-        bbbm_prevalence = builder.data.load(ALZHEIMERS.BBBM_CONDITIONAL_PREVALENCE)
-        mci_prevalence = builder.data.load(ALZHEIMERS.MCI_CONDITIONAL_PREVALENCE)
-        alz_prevalence = bbbm_prevalence.copy()
-        alz_prevalence["value"] = 1.0
-        alz_prevalence["value"] = alz_prevalence["value"] - (
-            bbbm_prevalence["value"] + mci_prevalence["value"]
-        )
-
-        return alz_prevalence

src/vivarium_csu_alzheimers/components/population.py

@@ -237,7 +237,7 @@ def load_population_structure(self, builder: Builder) -> pd.Series:
         return pop_structure
 
     def load_prevalence(self, builder: Builder) -> pd.Series:
-        prevalence = builder.data.load(data_keys.ALZHEIMERS_CONSISTENT.PREVALENCE_ANY)
+        prevalence = builder.data.load(data_keys.ALZHEIMERS_CONSISTENT.AD_PREVALENCE)
 
         # Updating age_end to match configuration since some simulants are living past 125
         prevalence.loc[prevalence["age_end"] == 125, "age_end"] = self.age_end

src/vivarium_csu_alzheimers/components/testing.py

@@ -118,7 +118,9 @@ def _update_baseline_testing(self, pop: pd.DataFrame) -> None:
         eligible_state = pop[COLUMNS.DISEASE_STATE].isin(
             [
                 ALZHEIMERS_DISEASE_MODEL.MCI_STATE,
-                ALZHEIMERS_DISEASE_MODEL.ALZHEIMERS_DISEASE_STATE,
+                ALZHEIMERS_DISEASE_MODEL.MILD_DEMENTIA_STATE,
+                ALZHEIMERS_DISEASE_MODEL.MODERATE_DEMENTIA_STATE,
+                ALZHEIMERS_DISEASE_MODEL.SEVERE_DEMENTIA_STATE,
             ]
         )
         eligible_test_propensity = pop[COLUMNS.TESTING_PROPENSITY] < (

src/vivarium_csu_alzheimers/constants/data_keys.py

@@ -10,12 +10,12 @@
 class __Population(NamedTuple):
     LOCATION: str = "population.location"
     STRUCTURE: str = "population.structure"
+    SCALING_FACTOR: str = "population.scaling_factor"
     AGE_BINS: str = "population.age_bins"
     DEMOGRAPHY: str = "population.demographic_dimensions"
     TMRLE: str = "population.theoretical_minimum_risk_life_expectancy"
     ACMR: str = "cause.all_causes.cause_specific_mortality_rate"
     LIVE_BIRTH_RATE: str = "covariate.live_births_by_sex.estimate"
-    SCALING_FACTOR: str = "population.scaling_factor"
 
     @property
     def name(self):
@@ -30,19 +30,40 @@ def log_name(self):
 
 
 class __Alzheimers(NamedTuple):
-    PREVALENCE: str = "cause.alzheimers.prevalence"
-    BBBM_CONDITIONAL_PREVALENCE: str = "cause.alzheimers.bbbm_conditional_prevalence"
-    MCI_CONDITIONAL_PREVALENCE: str = "cause.alzheimers.mci_conditional_prevalence"
-    MCI_TO_DEMENTIA_TRANSITION_RATE: str = "cause.alzheimers.mci_to_dementia_transition_rate"
-    SUSCEPTIBLE_TO_BBBM_TRANSITION_COUNT: str = (
-        "cause.alzheimers.susceptible_to_bbbm_transition_count"
+    AD_DEMENTIA_PREVALENCE: str = "cause.alzheimers.ad_dementia_prevalence"
+    MIXED_DEMENTIA_PREVALENCE: str = "cause.alzheimers.mixed_dementia_prevalence"
+
+    AD_DEMENTIA_INCIDENCE_RATE_TOTAL_POPULATION: str = (
+        "cause.alzheimers.ad_dementia_population_incidence_rate"
     )
-    # BBBM to MCI transition rate caluclated during sim using mci_hazard.py and time in state
-    INCIDENCE_RATE_TOTAL_POPULATION: str = "cause.alzheimers.population_incidence_rate"
-    CSMR: str = "cause.alzheimers_disease_and_other_dementias.cause_specific_mortality_rate"
-    EMR: str = "cause.alzheimers.excess_mortality_rate"
-    DISABILITY_WEIGHT: str = "cause.alzheimers_disease_and_other_dementias.disability_weight"
+    MIXED_DEMENTIA_INCIDENCE_RATE_TOTAL_POPULATION: str = (
+        "cause.alzheimers.mixed_dementia_population_incidence_rate"
+    )
+
+    MILD_DEMENTIA_PREVALENCE: str = (
+        "sequela.mild_alzheimers_disease_and_other_dementias.prevalence"
+    )
+    MODERATE_DEMENTIA_PREVALENCE: str = (
+        "sequela.moderate_alzheimers_disease_and_other_dementias.prevalence"
+    )
+    SEVERE_DEMENTIA_PREVALENCE: str = (
+        "sequela.severe_alzheimers_disease_and_other_dementias.prevalence"
+    )
+
+    EMR_COMO: str = "cause.alzheimers_disease_and_other_dementias.excess_mortality_rate"
+    EMR_DISMOD: str = "cause.dementia.excess_mortality_rate"
+
     MCI_DISABILITY_WEIGHT: str = "cause.alzheimers.mci_disability_weight"
+    MILD_DEMENTIA_DISABILITY_WEIGHT: str = (
+        "sequela.mild_alzheimers_disease_and_other_dementias.disability_weight"
+    )
+    MODERATE_DEMENTIA_DISABILITY_WEIGHT: str = (
+        "sequela.moderate_alzheimers_disease_and_other_dementias.disability_weight"
+    )
+    SEVERE_DEMENTIA_DISABILITY_WEIGHT: str = (
+        "sequela.severe_alzheimers_disease_and_other_dementias.disability_weight"
+    )
+
     RESTRICTIONS: str = "cause.alzheimers_disease_and_other_dementias.restrictions"
 
     @property
@@ -58,27 +79,59 @@ def log_name(self):
 
 
 class __AlzheimersConsistent(NamedTuple):
-    PREVALENCE_ANY: str = "cause.alzheimers_consistent.prevalence_any"
-    PREVALENCE: str = "cause.alzheimers.prevalence"
+    AD_PREVALENCE: str = "cause.alzheimers_consistent.alzheimers_prevalence"
+    MIXED_DEMENTIA_PREVALENCE: str = "cause.alzheimers_consistent.mixed_dementia_prevalence"
     BBBM_CONDITIONAL_PREVALENCE: str = (
         "cause.alzheimers_consistent.bbbm_conditional_prevalence"
     )
     MCI_CONDITIONAL_PREVALENCE: str = "cause.alzheimers_consistent.mci_conditional_prevalence"
-    DEMENTIA_CONDITIONAL_PREVALENCE: str = (
-        "cause.alzheimers_consistent.dementia_conditional_prevalence"
+    MILD_DEMENTIA_CONDITIONAL_PREVALENCE: str = (
+        "cause.alzheimers_consistent.mild_dementia_conditional_prevalence"
+    )
+    MODERATE_DEMENTIA_CONDITIONAL_PREVALENCE: str = (
+        "cause.alzheimers_consistent.moderate_dementia_conditional_prevalence"
+    )
+    SEVERE_DEMENTIA_CONDITIONAL_PREVALENCE: str = (
+        "cause.alzheimers_consistent.severe_dementia_conditional_prevalence"
     )
-    MCI_TO_DEMENTIA_TRANSITION_RATE: str = "cause.alzheimers.mci_to_dementia_transition_rate"
+
     SUSCEPTIBLE_TO_BBBM_TRANSITION_COUNT: str = (
         "cause.alzheimers_consistent.susceptible_to_bbbm_transition_count"
     )
     # BBBM to MCI transition rate caluclated during sim using mci_hazard.py and time in state
-    INCIDENCE_RATE_TOTAL_POPULATION: str = (
-        "cause.alzheimers_consistent.population_incidence_dementia"
+    # MCI to MILD transition rate is specified in data_values.py
+    MILD_TO_MODERATE_DEMENTIA_TRANSITION_RATE: str = (
+        "cause.alzheimers_consistent.mild_to_moderate_dementia_transition_rate"
+    )
+    MODERATE_TO_SEVERE_DEMENTIA_TRANSITION_RATE: str = (
+        "cause.alzheimers_consistent.moderate_to_severe_dementia_transition_rate"
+    )
+
+    BBBM_AD_INCIDENCE_RATE: str = (
+        "cause.alzheimers_consistent.susceptible_to_bbbm_ad_transition_rate"
     )
+    MIXED_DEMENTIA_INCIDENCE_RATE_TOTAL_POPULATION: str = (
+        "cause.alzheimers_consistent.mixed_dementia_population_incidence_rate"
+    )
+    MILD_DEMENTIA_INCIDENCE_RATE_TOTAL_POPULATION: str = (
+        "cause.alzheimers_consistent.population_incidence_mild_dementia"
+    )
+
     CSMR: str = "cause.alzheimers_disease_and_other_dementias.cause_specific_mortality_rate"
-    EMR: str = "cause.alzheimers_consistent.excess_mortality_rate"
-    DISABILITY_WEIGHT: str = "cause.alzheimers_disease_and_other_dementias.disability_weight"
+    EMR_MILD: str = "cause.alzheimers_consistent.excess_mortality_rate_mild"
+    EMR_MODERATE: str = "cause.alzheimers_consistent.excess_mortality_rate_moderate"
+    EMR_SEVERE: str = "cause.alzheimers_consistent.excess_mortality_rate_severe"
     MCI_DISABILITY_WEIGHT: str = "cause.alzheimers.mci_disability_weight"
+    MILD_DEMENTIA_DISABILITY_WEIGHT: str = (
+        "sequela.mild_alzheimers_disease_and_other_dementias.disability_weight"
+    )
+    MODERATE_DEMENTIA_DISABILITY_WEIGHT: str = (
+        "sequela.moderate_alzheimers_disease_and_other_dementias.disability_weight"
+    )
+    SEVERE_DEMENTIA_DISABILITY_WEIGHT: str = (
+        "sequela.severe_alzheimers_disease_and_other_dementias.disability_weight"
+    )
+
     RESTRICTIONS: str = "cause.alzheimers_disease_and_other_dementias.restrictions"
 
     @property

src/vivarium_csu_alzheimers/constants/models.py

@@ -11,7 +11,10 @@ class __AlzheimersDiseaseModel(NamedTuple):
     NAME: str = data_keys.ALZHEIMERS.log_name
     BBBM_STATE: str = "alzheimers_blood_based_biomarker_state"
     MCI_STATE: str = "alzheimers_mild_cognitive_impairment_state"
-    ALZHEIMERS_DISEASE_STATE: str = "alzheimers_disease_state"
+    MILD_DEMENTIA_STATE: str = "alzheimers_mild_dementia_state"
+    MODERATE_DEMENTIA_STATE: str = "alzheimers_moderate_dementia_state"
+    SEVERE_DEMENTIA_STATE: str = "alzheimers_severe_dementia_state"
+    MIXED_DEMENTIA_STATE: str = "mixed_dementia_state"
 
     def __iter__(self):
         """Allow iteration over the named tuple field values."""