vivarium-profiling 0.4.0__py3-none-any.whl

vivarium/profiling/__init__.py

@@ -0,0 +1,5 @@
+"""vivarium.profiling
+
+Profiling and benchmarking tools for Vivarium simulations.
+
+"""

vivarium/profiling/_version.py

@@ -0,0 +1 @@
+__version__ = "0.4.0"\n

vivarium/profiling/components/risks/effect.py

@@ -0,0 +1,45 @@
+import re
+
+import pandas as pd
+from vivarium.framework.engine import Builder
+from vivarium_public_health.risks.effect import (
+    NonLogLinearRiskEffect as NonLogLinearRiskEffect_,
+)
+from vivarium_public_health.risks.effect import RiskEffect as RiskEffect_
+
+"""
+Basic Risk Effect Wrapper for use in the MultiComponentParser
+"""
+
+
+class RiskEffect(RiskEffect_):
+    def get_filtered_data(
+        self, builder: Builder, data_source: str | float | pd.DataFrame
+    ) -> float | pd.DataFrame:
+        data = super().get_data(builder, data_source)
+
+        if isinstance(data, pd.DataFrame):
+            # filter data to only include the target entity and measure
+            correct_target_mask = True
+            columns_to_drop = []
+            if "affected_entity" in data.columns:
+                # THIS IS THE ONLY CHANGE! We need to filter to the non-suffixed name
+                correct_target_mask &= data["affected_entity"] == re.sub(
+                    r"(_\d+)$", "", self.target.name
+                )
+                columns_to_drop.append("affected_entity")
+            if "affected_measure" in data.columns:
+                correct_target_mask &= data["affected_measure"] == self.target.measure
+                columns_to_drop.append("affected_measure")
+            data = data[correct_target_mask].drop(columns=columns_to_drop)
+        return data
+
+
+class NonLogLinearRiskEffect(NonLogLinearRiskEffect_, RiskEffect):
+    pass
+
+
+from vivarium_public_health.risks.effect import (
+    NonLogLinearRiskEffect as NonLogLinearRiskEffect_,
+)
+from vivarium_public_health.risks.effect import RiskEffect as RiskEffect_

vivarium/profiling/constants/data_keys.py

@@ -0,0 +1,240 @@
+from typing import NamedTuple
+
+from vivarium_public_health.utilities import TargetString
+
+#############
+# Data Keys #
+#############
+
+METADATA_LOCATIONS = "metadata.locations"
+
+
+class __Population(NamedTuple):
+    LOCATION: str = "population.location"
+    STRUCTURE: str = "population.structure"
+    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"
+
+    @property
+    def name(self):
+        return "population"
+
+    @property
+    def log_name(self):
+        return "population"
+
+
+POPULATION = __Population()
+
+
+class __LRI(NamedTuple):
+
+    # Keys that will be loaded into the artifact. must have a colon type declaration
+    PREVALENCE: TargetString = TargetString("cause.lower_respiratory_infections.prevalence")
+    INCIDENCE_RATE: TargetString = TargetString(
+        "cause.lower_respiratory_infections.incidence_rate"
+    )
+    REMISSION_RATE: TargetString = TargetString(
+        "cause.lower_respiratory_infections.remission_rate"
+    )
+    DISABILITY_WEIGHT: TargetString = TargetString(
+        "cause.lower_respiratory_infections.disability_weight"
+    )
+    EMR: TargetString = TargetString(
+        "cause.lower_respiratory_infections.excess_mortality_rate"
+    )
+    CSMR: TargetString = TargetString(
+        "cause.lower_respiratory_infections.cause_specific_mortality_rate"
+    )
+    RESTRICTIONS: TargetString = TargetString(
+        "cause.lower_respiratory_infections.restrictions"
+    )
+
+    @property
+    def name(self):
+        return "lower_respiratory_infections"
+
+    @property
+    def log_name(self):
+        return self.name.replace("_", " ")
+
+
+LRI = __LRI()
+
+
+class __LRI2(NamedTuple):
+
+    # Keys that will be loaded into the artifact. must have a colon type declaration
+    PREVALENCE: TargetString = TargetString("cause.lower_respiratory_infections_2.prevalence")
+    INCIDENCE_RATE: TargetString = TargetString(
+        "cause.lower_respiratory_infections_2.incidence_rate"
+    )
+    REMISSION_RATE: TargetString = TargetString(
+        "cause.lower_respiratory_infections_2.remission_rate"
+    )
+    DISABILITY_WEIGHT: TargetString = TargetString(
+        "cause.lower_respiratory_infections_2.disability_weight"
+    )
+    EMR: TargetString = TargetString(
+        "cause.lower_respiratory_infections_2.excess_mortality_rate"
+    )
+    CSMR: TargetString = TargetString(
+        "cause.lower_respiratory_infections_2.cause_specific_mortality_rate"
+    )
+    RESTRICTIONS: TargetString = TargetString(
+        "cause.lower_respiratory_infections_2.restrictions"
+    )
+
+    @property
+    def name(self):
+        return "lower_respiratory_infections_2"
+
+    @property
+    def log_name(self):
+        return self.name.replace("_", " ")
+
+
+LRI2 = __LRI2()
+
+
+class __LRI3(NamedTuple):
+
+    # Keys that will be loaded into the artifact. must have a colon type declaration
+    PREVALENCE: TargetString = TargetString("cause.lower_respiratory_infections_3.prevalence")
+    INCIDENCE_RATE: TargetString = TargetString(
+        "cause.lower_respiratory_infections_3.incidence_rate"
+    )
+    REMISSION_RATE: TargetString = TargetString(
+        "cause.lower_respiratory_infections_3.remission_rate"
+    )
+    DISABILITY_WEIGHT: TargetString = TargetString(
+        "cause.lower_respiratory_infections_3.disability_weight"
+    )
+    EMR: TargetString = TargetString(
+        "cause.lower_respiratory_infections_3.excess_mortality_rate"
+    )
+    CSMR: TargetString = TargetString(
+        "cause.lower_respiratory_infections_3.cause_specific_mortality_rate"
+    )
+    RESTRICTIONS: TargetString = TargetString(
+        "cause.lower_respiratory_infections_3.restrictions"
+    )
+
+    @property
+    def name(self):
+        return "lower_respiratory_infections_3"
+
+    @property
+    def log_name(self):
+        return self.name.replace("_", " ")
+
+
+LRI3 = __LRI3()
+
+
+class __LRI4(NamedTuple):
+
+    # Keys that will be loaded into the artifact. must have a colon type declaration
+    PREVALENCE: TargetString = TargetString("cause.lower_respiratory_infections_4.prevalence")
+    INCIDENCE_RATE: TargetString = TargetString(
+        "cause.lower_respiratory_infections_4.incidence_rate"
+    )
+    REMISSION_RATE: TargetString = TargetString(
+        "cause.lower_respiratory_infections_4.remission_rate"
+    )
+    DISABILITY_WEIGHT: TargetString = TargetString(
+        "cause.lower_respiratory_infections_4.disability_weight"
+    )
+    EMR: TargetString = TargetString(
+        "cause.lower_respiratory_infections_4.excess_mortality_rate"
+    )
+    CSMR: TargetString = TargetString(
+        "cause.lower_respiratory_infections_4.cause_specific_mortality_rate"
+    )
+    RESTRICTIONS: TargetString = TargetString(
+        "cause.lower_respiratory_infections_4.restrictions"
+    )
+
+    @property
+    def name(self):
+        return "lower_respiratory_infections_4"
+
+    @property
+    def log_name(self):
+        return self.name.replace("_", " ")
+
+
+LRI4 = __LRI4()
+
+
+class __SBP(NamedTuple):
+    DISTRIBUTION: TargetString = TargetString(
+        "risk_factor.high_systolic_blood_pressure.distribution"
+    )
+    EXPOSURE_MEAN: TargetString = TargetString(
+        "risk_factor.high_systolic_blood_pressure.exposure"
+    )
+    EXPOSURE_SD: TargetString = TargetString(
+        "risk_factor.high_systolic_blood_pressure.exposure_standard_deviation"
+    )
+    EXPOSURE_WEIGHTS: TargetString = TargetString(
+        "risk_factor.high_systolic_blood_pressure.exposure_distribution_weights"
+    )
+    RELATIVE_RISK: TargetString = TargetString(
+        "risk_factor.high_systolic_blood_pressure.relative_risk"
+    )
+    PAF: TargetString = TargetString(
+        "risk_factor.high_systolic_blood_pressure.population_attributable_fraction"
+    )
+    TMRED: TargetString = TargetString("risk_factor.high_systolic_blood_pressure.tmred")
+    RELATIVE_RISK_SCALAR: TargetString = TargetString(
+        "risk_factor.high_systolic_blood_pressure.relative_risk_scalar"
+    )
+
+    @property
+    def name(self):
+        return "high_systolic_blood_pressure"
+
+    @property
+    def log_name(self):
+        return self.name.replace("_", " ")
+
+
+SBP = __SBP()
+
+
+class __Water(NamedTuple):
+    DISTRIBUTION: TargetString = TargetString("risk_factor.unsafe_water_source.distribution")
+    EXPOSURE: TargetString = TargetString("risk_factor.unsafe_water_source.exposure")
+    CATEGORIES: TargetString = TargetString("risk_factor.unsafe_water_source.categories")
+    RELATIVE_RISK: TargetString = TargetString(
+        "risk_factor.unsafe_water_source.relative_risk"
+    )
+    PAF: TargetString = TargetString(
+        "risk_factor.unsafe_water_source.population_attributable_fraction"
+    )
+
+    @property
+    def name(self):
+        return "unsafe_water_source"
+
+    @property
+    def log_name(self):
+        return self.name.replace("_", " ")
+
+
+WATER = __Water()
+
+
+MAKE_ARTIFACT_KEY_GROUPS = [
+    POPULATION,
+    LRI,
+    LRI2,
+    LRI3,
+    LRI4,
+    SBP,
+    WATER,
+]

vivarium/profiling/constants/data_values.py

@@ -0,0 +1,34 @@
+from datetime import datetime
+
+############################
+# Disease Model Parameters #
+############################
+
+REMISSION_RATE = 0.1
+MEAN_SOJOURN_TIME = 10
+
+
+##############################
+# Screening Model Parameters #
+##############################
+
+PROBABILITY_ATTENDING_SCREENING_KEY = "probability_attending_screening"
+PROBABILITY_ATTENDING_SCREENING_START_MEAN = 0.25
+PROBABILITY_ATTENDING_SCREENING_START_STDDEV = 0.0025
+PROBABILITY_ATTENDING_SCREENING_END_MEAN = 0.5
+PROBABILITY_ATTENDING_SCREENING_END_STDDEV = 0.005
+
+FIRST_SCREENING_AGE = 21
+MID_SCREENING_AGE = 30
+LAST_SCREENING_AGE = 65
+
+
+###################################
+# Scale-up Intervention Constants #
+###################################
+SCALE_UP_START_DT = datetime(2021, 1, 1)
+SCALE_UP_END_DT = datetime(2030, 1, 1)
+SCREENING_SCALE_UP_GOAL_COVERAGE = 0.50
+SCREENING_SCALE_UP_DIFFERENCE = (
+    SCREENING_SCALE_UP_GOAL_COVERAGE - PROBABILITY_ATTENDING_SCREENING_START_MEAN
+)

vivarium/profiling/constants/metadata.py

@@ -0,0 +1,42 @@
+from typing import NamedTuple
+
+import pandas as pd
+
+####################
+# Project metadata #
+####################
+
+# Underscore form is intentional: this string is used to build paths on the
+# shared cluster filesystem (e.g. /mnt/team/.../{PROJECT_NAME}/artifacts/).
+# Changing it would break access to existing artifacts.
+PROJECT_NAME = "vivarium_profiling"
+CLUSTER_PROJECT = "proj_simscience_prod"
+
+CLUSTER_QUEUE = "all.q"
+MAKE_ARTIFACT_MEM = 10  # GB
+MAKE_ARTIFACT_CPU = 1
+MAKE_ARTIFACT_RUNTIME = "3:00:00"
+MAKE_ARTIFACT_SLEEP = 10
+
+LOCATIONS = [
+    "Pakistan",
+]
+
+ARTIFACT_INDEX_COLUMNS = [
+    "sex",
+    "age_start",
+    "age_end",
+    "year_start",
+    "year_end",
+]
+
+DRAW_COUNT = 1000
+ARTIFACT_COLUMNS = pd.Index([f"draw_{i}" for i in range(DRAW_COUNT)])
+
+
+class __Scenarios(NamedTuple):
+    baseline: str = "baseline"
+    # TODO - add scenarios here
+
+
+SCENARIOS = __Scenarios()

vivarium/profiling/constants/models.py

@@ -0,0 +1,11 @@
+from vivarium.profiling.constants import data_keys
+
+###########################
+# Disease Model variables #
+###########################
+
+# TODO input details of model states
+SOME_MODEL_NAME = data_keys.SOME_DISEASE.name
+SUSCEPTIBLE_STATE_NAME = f"susceptible_to_{SOME_MODEL_NAME}"
+FIRST_STATE_NAME = "first_state"
+SECOND_STATE_NAME = "second_state"

vivarium/profiling/constants/paths.py

@@ -0,0 +1,10 @@
+from pathlib import Path
+
+import vivarium.profiling
+from vivarium.profiling.constants import metadata
+
+BASE_DIR = Path(vivarium.profiling.__file__).resolve().parent
+
+ARTIFACT_ROOT = Path(
+    f"/mnt/team/simulation_science/pub/models/{metadata.PROJECT_NAME}/artifacts/"
+)

vivarium/profiling/constants/scenarios.py

@@ -0,0 +1,29 @@
+from typing import NamedTuple
+
+#############
+# Scenarios #
+#############
+
+
+class InterventionScenario:
+    def __init__(
+        self,
+        name: str,
+        # todo add additional interventions
+        # has_treatment_one: bool = False,
+        # has_treatment_two: bool = False,
+    ):
+        self.name = name
+        # self.has_treatment_one = has_treatment_one
+        # self.has_treatment_two = has_treatment_two
+
+
+class __InterventionScenarios(NamedTuple):
+    BASELINE: InterventionScenario = InterventionScenario("baseline")
+    # todo add additional intervention scenarios
+
+    def __get_item__(self, item):
+        return self._asdict()[item]
+
+
+INTERVENTION_SCENARIOS = __InterventionScenarios()

vivarium/profiling/data/builder.py

@@ -0,0 +1,128 @@
+"""Modularized functions for building project data artifacts.
+
+This module is an abstraction around the load portion of our artifact building ETL pipeline.
+The intent is to be declarative so it's easy to see what is put into the artifact and how.
+Some degree of verbosity/boilerplate is fine in the interest of transparency.
+
+.. admonition::
+
+   Logging in this module should be done at the ``debug`` level.
+
+"""
+from pathlib import Path
+
+import pandas as pd
+from loguru import logger
+from vivarium.framework.artifact import Artifact, EntityKey
+
+from vivarium.profiling.constants import data_keys
+from vivarium.profiling.data import loader
+
+
+def open_artifact(output_path: Path, location: str) -> Artifact:
+    """Creates or opens an artifact at the output path.
+
+    Parameters
+    ----------
+    output_path
+        Fully resolved path to the artifact file.
+    location
+        Proper GBD location name represented by the artifact.
+
+    Returns
+    -------
+        A new artifact.
+
+    """
+    if not output_path.exists():
+        logger.debug(f"Creating artifact at {str(output_path)}.")
+    else:
+        logger.debug(f"Opening artifact at {str(output_path)} for appending.")
+
+    artifact = Artifact(output_path)
+
+    key = data_keys.METADATA_LOCATIONS
+    if key not in artifact:
+        artifact.write(key, [location])
+
+    return artifact
+
+
+def load_and_write_data(
+    artifact: Artifact, key: str, location: str, years: str | None, replace: bool
+):
+    """Loads data and writes it to the artifact if not already present.
+
+    Parameters
+    ----------
+    artifact
+        The artifact to write to.
+    key
+        The entity key associated with the data to write.
+    location
+        The location associated with the data to load and the artifact to
+        write to.
+    replace
+        Flag which determines whether to overwrite existing data
+
+    """
+    if key in artifact and not replace:
+        logger.debug(f"Data for {key} already in artifact.  Skipping...")
+    else:
+        logger.debug(f"Loading data for {key} for location {location}.")
+        # years is either a string we want to convert to an int, 'all', or None
+        years = int(years) if years and years != "all" else years
+        data = loader.get_data(key, location, years)
+        if key not in artifact:
+            logger.debug(f"Writing data for {key} to artifact.")
+            artifact.write(key, data)
+        else:  # key is in artifact, but should be replaced
+            logger.debug(f"Replacing data for {key} in artifact.")
+            artifact.replace(key, data)
+    return artifact.load(key)
+
+
+def write_data(artifact: Artifact, key: str, data: pd.DataFrame):
+    """Writes data to the artifact if not already present.
+
+    Parameters
+    ----------
+    artifact
+        The artifact to write to.
+    key
+        The entity key associated with the data to write.
+    data
+        The data to write.
+
+    """
+    if key in artifact:
+        logger.debug(f"Data for {key} already in artifact.  Skipping...")
+    else:
+        logger.debug(f"Writing data for {key} to artifact.")
+        artifact.write(key, data)
+    return artifact.load(key)
+
+
+# TODO - writing and reading by draw is necessary if you are using
+#        LBWSG data. Find the read function in utilities.py
+def write_data_by_draw(artifact: Artifact, key: str, data: pd.DataFrame):
+    """Writes data to the artifact on a per-draw basis. This is useful
+    for large datasets like Low Birthweight Short Gestation (LBWSG).
+
+    Parameters
+    ----------
+    artifact
+        The artifact to write to.
+    key
+        The entity key associated with the data to write.
+    data
+        The data to write.
+
+    """
+    with pd.HDFStore(artifact.path, complevel=9, mode="a") as store:
+        key = EntityKey(key)
+        artifact._keys.append(key)
+        store.put(f"{key.path}/index", data.index.to_frame(index=False))
+        data = data.reset_index(drop=True)
+        for c in data.columns:
+            store.put(f"{key.path}/{c}", data[c])