Merge pull request #5 from NSAPH-Projects:mauriciogtec/issue4

Fix installer #4

Author: mauriciogtec

MANIFEST.in renamed to _deprecated/MANIFEST.in

@@ -1,3 +1,44 @@
 [build-system]
-requires = ["setuptools>=42", "wheel"]
+requires = ["setuptools", "wheel"]
 build-backend = "setuptools.build_meta"
+
+[project]
+name = "weather2alert"
+version = "0.1.0"
+description = "A gym environment for optimizing heat alert issuance during heatwaves"
+readme = "README.md"
+authors = [{ name="Anonymous", email="[email protected]" }]
+license = { text = "MIT" }
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+]
+requires-python = ">=3.10"
+dependencies = [
+    "scipy",
+    "tqdm",
+    "pyarrow",
+    "pandas",
+    "torch",
+    "gymnasium",
+]
+
+# Optional dependencies go here under 'optional-dependencies'
+[project.optional-dependencies]
+dev = ["pytest"]
+
+[tool.setuptools]
+packages = ["src.weather2alert"]
+include-package-data = true
+
+[tool.setuptools.package-data]
+weather2alert = [
+    "weights/nn_full_medicare/*",
+    "weights/nn_debug_medicare/*",
+    "weights/master.yaml",
+    "data/processed/*.parquet",
+    "data/raw/*",
+]

setup.py renamed to _deprecated/setup.py

@@ -0,0 +1 @@
+

pyproject.toml

@@ -1,9 +1,4 @@
 import pandas as pd
-from sklearn.preprocessing import StandardScaler
-import torch
-
-
-# import matplotlib.pyplot as plt
 
 WESTERN_STATES = [
     "AZ",

requirements.txt

@@ -1,75 +1,84 @@
-from importlib import resources
+import os
+from importlib.util import find_spec
 from typing import Literal
 
-import gymnasium as gym
 import numpy as np
-from gymnasium import spaces
-from scipy.special import expit as sigmoid
 import pandas as pd
 import torch
 import yaml
+from gymnasium import Env, spaces
+from scipy.special import expit as sigmoid
 
-from weather2alert.datautils import get_similar_counties
-
-
-with open(resources.path("weather2alert.weights", "master.yaml"), "r") as f:
-    VALID_WEIGHTS = yaml.safe_load(f)
+from .datautils import get_similar_counties
 
 
-class HeatAlertEnv(gym.Env):
+class HeatAlertEnv(Env):
     """Class to simulate the environment for the online RL agent."""
 
     def __init__(
-        self, weights: str = "nn_full_medicare", valid_years: list | None = None
+        self,
+        weights: str = "nn_full_medicare",
+        years: list | None = None,
+        fips_list: list | None = None,
+        similar_climate_counties: bool = False,
+        budget: int | None = None,
     ):
         """Initialize the environment."""
         super().__init__()
-        self.valid_years = valid_years
-        # assert (
-        #     weights in VALID_WEIGHTS
-        # ), f"Invalid weights: {weights}, valid weights are {VALID_WEIGHTS}"
+        self.valid_years = years
+        self.similar_climate_counties = similar_climate_counties
+        self.budget = budget
+        if years is None:
+            years = list(range(2006, 2017))
 
         # load state and confounders data
-        path = resources.path(
-            "weather2alert.data.processed", "exogenous_states.parquet"
-        )
-        exogenous_states = pd.read_parquet(path)
-        path = resources.path(
-            "weather2alert.data.processed", "endogenous_states_actions.parquet"
+
+        # check if path data/processed exists, then we are working with local data
+        if False: # os.path.exists("data/processed"):
+            root = "./"
+        else:
+            root = find_spec("weather2alert").submodule_search_locations[0]
+
+        processed_path = os.path.join(root, "data/processed")
+        weights_path = os.path.join(root, "weights")
+
+        exogenous_states = pd.read_parquet(processed_path + "/exogenous_states.parquet")
+        endogenous_states_actions = pd.read_parquet(
+            processed_path + "/endogenous_states_actions.parquet"
         )
-        endogenous_states_actions = pd.read_parquet(path)
         merged = pd.merge(
             exogenous_states, endogenous_states_actions, on=["fips", "date"]
         )
         merged["year"] = merged.date.str[:4].astype(int)
 
         # make sure merged is order by fips date and remove dates outside of the range
         # 152 days of the summer starting on May 1st to Sep 30th
-        month =  merged.date.str[5:7]
+        month = merged.date.str[5:7]
         merged = merged[(month >= "05") & (month <= "09")].copy()
         merged = merged.drop_duplicates(["fips", "date"])
 
         # merged.set_index(["fips", "date"], inplace=True)
-        confounders = pd.read_parquet(
-            resources.path("weather2alert.data.processed", "confounders.parquet")
-        )
+        confounders = pd.read_parquet(processed_path + "/confounders.parquet")
 
         self.merged = merged.set_index(["fips", "year"])
         self.confounders = confounders
 
         # load posterior parameters and config
-        weights_dir = "weather2alert.weights." + weights
-        path = resources.path(weights_dir, "posterior_samples.pt")
-        posterior_samples = torch.load(path, weights_only=True)
-        self.fips_list = posterior_samples["fips_list"]
+        posterior_samples = torch.load(
+            f"{weights_path}/{weights}/posterior_samples.pt", weights_only=True
+        )
+
+        self.fips_list = fips_list
+        if fips_list is None:
+            self.fips_list = posterior_samples["fips_list"]
 
         self.baseline_coefs = {
             k: v for k, v in posterior_samples.items() if k.startswith("baseline")
         }
         self.effectiveness_coefs = {
             k: v for k, v in posterior_samples.items() if k.startswith("effectiveness")
         }
-        with open(resources.path(weights_dir, "config.yaml"), "r") as f:
+        with open(rf"{weights_path}/{weights}/config.yaml", "r") as f:
             self.config = yaml.safe_load(f)
 
         # get num posterior samples
@@ -91,12 +100,15 @@ def __init__(
         for k in self.merged.columns:
             if k.startswith("bspline_"):
                 self.merged[k.replace("bspline_", "bsplines_")] = self.merged[k]
-        # ----
+
+        if self.valid_years is None:
+            self.valid_years = list(self.merged.index.get_level_values("year").unique())
 
     def _get_episode(
         self,
         location: str,
         augment: bool = False,
+        year: int | None = None,
     ):
         if augment:
             # get similar counties
@@ -109,11 +121,9 @@ def _get_episode(
             self.location_index = self.fips_list.index(location)
 
         # split by year and index by dos, drop data
-        valid_years = self.valid_years
-        if self.valid_years is None:
-            valid_years = self.merged.loc[self.location].index.unique()
+        if year is None:
+            year = self.rng.choice(self.valid_years)
 
-        year = self.rng.choice(valid_years)
         year_data = self.merged.loc[(location, year)]
         year_data = (
             year_data.reset_index().drop(columns=["fips", "year"]).set_index("date")
@@ -123,8 +133,9 @@ def _get_episode(
     def reset(
         self,
         location: str | None = None,
-        similar_climate_counties: bool = False,
+        similar_climate_counties: bool | None = None,
         seed: int | None = None,
+        budget: int | None = None,
         sample_budget: bool = False,
         sample_budget_type: Literal["less_than", "centered"] = "less_than",
     ):
@@ -133,6 +144,9 @@ def reset(
             seed = np.random.randint(0, 10000)
         self.rng = np.random.default_rng(seed)
 
+        if similar_climate_counties is None:
+            similar_climate_counties = self.similar_climate_counties
+
         # if location is None, pick a random location
         if location is None:
             location = self.rng.choice(self.fips_list)
@@ -150,18 +164,24 @@ def reset(
         self.alert_streak = 0
         self.t = 0  # day of summer indicator
 
-        b = self.ep["remaining_budget"].iloc[0]
+        if self.budget is None:
+            self.budget = (
+                self.ep["remaining_budget"].iloc[0] if budget is None else budget
+            )
+
         if sample_budget:
+            b = self.budget
             if sample_budget_type == "less_than":
                 self.budget = self.rng.integers(0, b + 1)
             elif sample_budget_type == "centered":
                 self.budget = self.rng.integers(0.5 * b, 1.5 * b + 1)
-        else:
-            self.budget = b
+        self.remaining_budget = self.budget
 
         self.at_budget = False
         self.observation = self._get_obs()
-        return self.observation, self._get_info()
+        if not hasattr(self, "feat_names"):
+            self.feat_names = self.observation.index.tolist()
+        return self.observation.values, self._get_info()
 
     def _get_obs(self):
         row = self.ep.iloc[self.t].copy()
@@ -195,10 +215,10 @@ def _get_reward(self, action):
             x = row[k.replace("effectiveness_", "")]
             v = v[self.coef_index, 0, li].item()
             effectiveness_contribs.append(x * v)
-        effectiveness = sigmoid(sum(effectiveness_contribs))
+        effectiveness = sigmoid(sum(effectiveness_contribs)) * (row["heat_qi"] > 0.5)
 
-        # reward is 1 - normalized hospitalization rate
-        reward = float(1 - baseline * (1 - effectiveness * action))
+        # reward is - normalized hospitalization rate / 10_000
+        reward = float(-10_000 * baseline * (1 - effectiveness * action))
 
         if action == 1 and self.at_budget:
             reward = -1
@@ -208,8 +228,9 @@ def _get_reward(self, action):
     def _get_info(self) -> dict:
         return {
             "episode_index": self.ep_index,
-            "budget": self.budget,
-            "feature_names": self.ep.columns.tolist(),
+            "remaining_budget": self.remaining_budget,
+            "at_budget": self.at_budget,
+            "feature_names": self.feat_names,
             "location": self.location,
             "location_index": self.location_index,
         }
@@ -225,17 +246,20 @@ def step(self, action: int):
             actual_action = action
 
         self.actual_alert_buffer.append(actual_action)
+        if actual_action == 1:
+            self.remaining_budget -= 1
 
         # compute reward for the new state
         reward = self._get_reward(actual_action)
 
         # advance state
-        self.t += 1
-        observation = self._get_obs().values
-        done = self.t == self.n_days - 1
-        self.alert_streak = self.alert_streak + 1 if actual_action else 0
+        done = self.t >= self.n_days - 1
+        if not done:
+            self.observation = self._get_obs()
+            self.t += 1
+            self.alert_streak = self.alert_streak + 1 if actual_action else 0
 
-        return observation, reward, done, False, self._get_info()
+        return self.observation.values, reward, done, False, self._get_info()
 
 
 if __name__ == "__main__":