diff --git a/build_flu_age_training_datasets.ipynb b/build_flu_age_training_datasets.ipynb
new file mode 100644
index 0000000..3ac8cf2
--- /dev/null
+++ b/build_flu_age_training_datasets.ipynb
@@ -0,0 +1,626 @@
+{
+ "cells": [
+ {
+ "cell_type": "code",
+ "execution_count": 2,
+ "id": "46cfd1c6",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import numpy as np\n",
+ "import pandas as pd\n",
+ "import xarray as xr"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 3,
+ "id": "c01bc22f",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from influpaint.datasets import mixer as dataset_mixer\n",
+ "from influpaint.utils import converters"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "1cfd665d",
+ "metadata": {},
+ "source": [
+ "Need for this file:\n",
+ " - the imports above, with modifications in `converters.py` and `mixer.py` found in `adding-new-datasource` branch\n",
+ " - the output from `prep_age_flu_hosp_data.ipynb` (at top level of `influpaint` directory, also in `adding-new-datasource` branch)\n",
+ " - your desired output path for .nc NetCDF files (for final cell to run)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "be286f08",
+ "metadata": {},
+ "source": [
+ "This file is an attempted modification of the logic found in influpaint/2-build_training_flu_datasets_ipynb.py"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 4,
+ "id": "cd01cb57",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "all_datasets_df = pd.read_csv(\n",
+ " \"/Users/emprzy/Documents/work/miscellaneous/influpaint_data/training_data.csv\",\n",
+ " dtype={\n",
+ " 'location_code': str,\n",
+ " 'value': np.float64,\n",
+ " 'fluseason_fraction': np.float64,\n",
+ " 'season_week': np.int64,\n",
+ " 'fluseason': np.int64,\n",
+ " 'datasetH1': str,\n",
+ " 'datasetH2': str,\n",
+ " 'sample': str \n",
+ " },\n",
+ " parse_dates=['week_enddate']\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 13,
+ "id": "95349996",
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/html": [
+ "
\n",
+ "\n",
+ "
\n",
+ " \n",
+ " \n",
+ " | \n",
+ " location_code | \n",
+ " season_week | \n",
+ " week_enddate | \n",
+ " sample | \n",
+ " fluseason | \n",
+ " datasetH1 | \n",
+ " datasetH2 | \n",
+ " fluseason_fraction | \n",
+ " 0-130 | \n",
+ " 0-4 | \n",
+ " 18-49 | \n",
+ " 5-17 | \n",
+ " 50-64 | \n",
+ " 65-130 | \n",
+ "
\n",
+ " \n",
+ " \n",
+ " \n",
+ " | 0 | \n",
+ " 01 | \n",
+ " 1 | \n",
+ " 2024-08-03 | \n",
+ " 1 | \n",
+ " 2024 | \n",
+ " NHSN | \n",
+ " NHSN | \n",
+ " 0.841096 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ "
\n",
+ " \n",
+ " | 1 | \n",
+ " 01 | \n",
+ " 2 | \n",
+ " 2024-08-10 | \n",
+ " 1 | \n",
+ " 2024 | \n",
+ " NHSN | \n",
+ " NHSN | \n",
+ " 0.860274 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ "
\n",
+ " \n",
+ " | 2 | \n",
+ " 01 | \n",
+ " 3 | \n",
+ " 2024-08-17 | \n",
+ " 1 | \n",
+ " 2024 | \n",
+ " NHSN | \n",
+ " NHSN | \n",
+ " 0.879452 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ "
\n",
+ " \n",
+ " | 3 | \n",
+ " 01 | \n",
+ " 4 | \n",
+ " 2024-08-24 | \n",
+ " 1 | \n",
+ " 2024 | \n",
+ " NHSN | \n",
+ " NHSN | \n",
+ " 0.898630 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ "
\n",
+ " \n",
+ " | 4 | \n",
+ " 01 | \n",
+ " 5 | \n",
+ " 2024-08-31 | \n",
+ " 1 | \n",
+ " 2024 | \n",
+ " NHSN | \n",
+ " NHSN | \n",
+ " 0.917808 | \n",
+ " 1.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ "
\n",
+ " \n",
+ " | ... | \n",
+ " ... | \n",
+ " ... | \n",
+ " ... | \n",
+ " ... | \n",
+ " ... | \n",
+ " ... | \n",
+ " ... | \n",
+ " ... | \n",
+ " ... | \n",
+ " ... | \n",
+ " ... | \n",
+ " ... | \n",
+ " ... | \n",
+ " ... | \n",
+ "
\n",
+ " \n",
+ " | 2647 | \n",
+ " 56 | \n",
+ " 48 | \n",
+ " 2025-06-28 | \n",
+ " 1 | \n",
+ " 2024 | \n",
+ " NHSN | \n",
+ " NHSN | \n",
+ " 0.739726 | \n",
+ " 1.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ " 1.0 | \n",
+ "
\n",
+ " \n",
+ " | 2648 | \n",
+ " 56 | \n",
+ " 49 | \n",
+ " 2025-07-05 | \n",
+ " 1 | \n",
+ " 2024 | \n",
+ " NHSN | \n",
+ " NHSN | \n",
+ " 0.758904 | \n",
+ " 1.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ " 1.0 | \n",
+ " 0.0 | \n",
+ "
\n",
+ " \n",
+ " | 2649 | \n",
+ " 56 | \n",
+ " 50 | \n",
+ " 2025-07-12 | \n",
+ " 1 | \n",
+ " 2024 | \n",
+ " NHSN | \n",
+ " NHSN | \n",
+ " 0.778082 | \n",
+ " 1.0 | \n",
+ " 0.0 | \n",
+ " 1.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ "
\n",
+ " \n",
+ " | 2650 | \n",
+ " 56 | \n",
+ " 51 | \n",
+ " 2025-07-19 | \n",
+ " 1 | \n",
+ " 2024 | \n",
+ " NHSN | \n",
+ " NHSN | \n",
+ " 0.797260 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ "
\n",
+ " \n",
+ " | 2651 | \n",
+ " 56 | \n",
+ " 52 | \n",
+ " 2025-07-26 | \n",
+ " 1 | \n",
+ " 2024 | \n",
+ " NHSN | \n",
+ " NHSN | \n",
+ " 0.816438 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ " 0.0 | \n",
+ "
\n",
+ " \n",
+ "
\n",
+ "
2652 rows × 14 columns
\n",
+ "
"
+ ],
+ "text/plain": [
+ " location_code season_week week_enddate sample fluseason datasetH1 \\\n",
+ "0 01 1 2024-08-03 1 2024 NHSN \n",
+ "1 01 2 2024-08-10 1 2024 NHSN \n",
+ "2 01 3 2024-08-17 1 2024 NHSN \n",
+ "3 01 4 2024-08-24 1 2024 NHSN \n",
+ "4 01 5 2024-08-31 1 2024 NHSN \n",
+ "... ... ... ... ... ... ... \n",
+ "2647 56 48 2025-06-28 1 2024 NHSN \n",
+ "2648 56 49 2025-07-05 1 2024 NHSN \n",
+ "2649 56 50 2025-07-12 1 2024 NHSN \n",
+ "2650 56 51 2025-07-19 1 2024 NHSN \n",
+ "2651 56 52 2025-07-26 1 2024 NHSN \n",
+ "\n",
+ " datasetH2 fluseason_fraction 0-130 0-4 18-49 5-17 50-64 65-130 \n",
+ "0 NHSN 0.841096 0.0 0.0 0.0 0.0 0.0 0.0 \n",
+ "1 NHSN 0.860274 0.0 0.0 0.0 0.0 0.0 0.0 \n",
+ "2 NHSN 0.879452 0.0 0.0 0.0 0.0 0.0 0.0 \n",
+ "3 NHSN 0.898630 0.0 0.0 0.0 0.0 0.0 0.0 \n",
+ "4 NHSN 0.917808 1.0 0.0 0.0 0.0 0.0 0.0 \n",
+ "... ... ... ... ... ... ... ... ... \n",
+ "2647 NHSN 0.739726 1.0 0.0 0.0 0.0 0.0 1.0 \n",
+ "2648 NHSN 0.758904 1.0 0.0 0.0 0.0 1.0 0.0 \n",
+ "2649 NHSN 0.778082 1.0 0.0 1.0 0.0 0.0 0.0 \n",
+ "2650 NHSN 0.797260 0.0 0.0 0.0 0.0 0.0 0.0 \n",
+ "2651 NHSN 0.816438 0.0 0.0 0.0 0.0 0.0 0.0 \n",
+ "\n",
+ "[2652 rows x 14 columns]"
+ ]
+ },
+ "execution_count": 13,
+ "metadata": {},
+ "output_type": "execute_result"
+ }
+ ],
+ "source": [
+ "\n",
+ "x = all_datasets_df[all_datasets_df['datasetH1'] == 'NHSN']\n",
+ "x = x[x['location_code'] != 'US']\n",
+ "x = x.pivot_table(\n",
+ " index=['location_code', 'season_week', 'week_enddate', 'sample', 'fluseason', 'datasetH1', 'datasetH2', 'fluseason_fraction'],\n",
+ " columns='age_group',\n",
+ " values='value'\n",
+ ").reset_index()\n",
+ "x.columns.name = None\n",
+ "x"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 14,
+ "id": "2a36a001",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "\n",
+ "RangeIndex: 2652 entries, 0 to 2651\n",
+ "Data columns (total 14 columns):\n",
+ " # Column Non-Null Count Dtype \n",
+ "--- ------ -------------- ----- \n",
+ " 0 location_code 2652 non-null object \n",
+ " 1 season_week 2652 non-null int64 \n",
+ " 2 week_enddate 2652 non-null datetime64[ns]\n",
+ " 3 sample 2652 non-null object \n",
+ " 4 fluseason 2652 non-null int64 \n",
+ " 5 datasetH1 2652 non-null object \n",
+ " 6 datasetH2 2652 non-null object \n",
+ " 7 fluseason_fraction 2652 non-null float64 \n",
+ " 8 0-130 2652 non-null float64 \n",
+ " 9 0-4 2652 non-null float64 \n",
+ " 10 18-49 2652 non-null float64 \n",
+ " 11 5-17 2652 non-null float64 \n",
+ " 12 50-64 2652 non-null float64 \n",
+ " 13 65-130 2652 non-null float64 \n",
+ "dtypes: datetime64[ns](1), float64(7), int64(2), object(4)\n",
+ "memory usage: 290.2+ KB\n"
+ ]
+ }
+ ],
+ "source": [
+ "x.info()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "681bf760",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# age_groups = sorted(all_datasets_df['age_group'].unique())\n",
+ "# pivot the table to have one column per age group\n",
+ "df_wide = all_datasets_df.pivot_table(\n",
+ " index=['location_code', 'season_week', 'week_enddate', 'sample', 'fluseason', 'datasetH1', 'datasetH2', 'fluseason_fraction'],\n",
+ " columns='age_group',\n",
+ " values='value'\n",
+ ").reset_index()\n",
+ "df_wide.columns.name = None"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 26,
+ "id": "eb62f2da",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "AGE_COLUMNS = [\n",
+ " '0-130',\n",
+ " '0-4',\n",
+ " '18-49',\n",
+ " '5-17',\n",
+ " '50-64',\n",
+ " '65-130'\n",
+ " ]\n",
+ "build_frames_config = { \n",
+ " \"NHSN\": {\"multiplier\": 2080, \"to_scale\": False}, # creates 10k frames of training data, ~20% are surveillance\n",
+ " \"SMH_R4\": {\"multiplier\": 1, \"to_scale\": False},\n",
+ " \"SMH_R5\": {\"multiplier\": 1, \"to_scale\": False},\n",
+ "}\n",
+ "season_axis = set(all_datasets_df['location_code']) # my SeasonAxis replacement (PATCH)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 27,
+ "id": "b20c5d28",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Pre-computing intelligent fill lookup table...\n",
+ " Building location data lookup...\n",
+ " Lookup table built for 52 locations\n",
+ "Building frames...\n",
+ "Processing NHSN (multiplier=2080)...\n"
+ ]
+ },
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ " \r"
+ ]
+ },
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Processing SMH_R4 (multiplier=1)...\n"
+ ]
+ },
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ " \r"
+ ]
+ },
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Processing SMH_R5 (multiplier=1)...\n"
+ ]
+ },
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ " \r"
+ ]
+ },
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Created 10000 total frames:\n",
+ " NHSN: 2080 frames (multiplier=2080)\n",
+ " SMH_R4: 600 frames (multiplier=1)\n",
+ " SMH_R5: 7320 frames (multiplier=1)\n"
+ ]
+ }
+ ],
+ "source": [
+ "frame_list = dataset_mixer.build_frames(\n",
+ " df_wide, \n",
+ " build_frames_config,\n",
+ " season_axis=season_axis, \n",
+ " fill_missing_locations=\"random\",\n",
+ ") \n",
+ "\n",
+ "# version below is for the unpivoted data (i don't think this the right way to do this)\n",
+ "#frame_list = dataset_mixer.build_frames(\n",
+ " #all_datasets_df, \n",
+ " #build_frames_config,\n",
+ " #season_axis=season_axis, \n",
+ " #fill_missing_locations=\"random\",\n",
+ "#) "
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 29,
+ "id": "da132fd5",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def build_dataset_from_framelist(frame_list, season_setup: set[str]):\n",
+ " main_origins = []\n",
+ " for i, frame in enumerate(frame_list):\n",
+ " df = frame_list[i]\n",
+ " df[\"fluseason\"] = i # Normalize sample ID\n",
+ " frame_list[i] = df\n",
+ " \n",
+ " # Validation\n",
+ " assert df.season_week.max() == 53 and df.season_week.min() == 1\n",
+ " \n",
+ " # Track Origin\n",
+ " main_origins.append(df[\"origin\"].mode()[0] if not df[\"origin\"].mode().empty else None)\n",
+ "\n",
+ " all_frames_df = pd.concat(frame_list).reset_index(drop=True)\n",
+ " \n",
+ " # PATCHED to accommodate 6 channels\n",
+ " array_list = converters.dataframe_to_arraylist( \n",
+ " df=all_frames_df, \n",
+ " locations=season_setup, \n",
+ " age_columns=AGE_COLUMNS\n",
+ " )\n",
+ "\n",
+ " array = np.array(array_list) # Shape: (Samples, 6, 64, 64)\n",
+ "\n",
+ " flu_payload_array = xr.DataArray(\n",
+ " array, \n",
+ " coords={\n",
+ " 'sample': np.arange(array.shape[0]),\n",
+ " 'feature': AGE_COLUMNS, # Coordinates for the 6 age groups\n",
+ " 'season_week_padded': np.arange(64), # Time axis (0-63)\n",
+ " 'location_padded': np.arange(64) # Spatial axis (0-63)\n",
+ " }, \n",
+ " dims=[\"sample\", \"feature\", \"season_week_padded\", \"location_padded\"]\n",
+ " )\n",
+ " return flu_payload_array, main_origins\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 30,
+ "id": "e2d9b0b9",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "flu_payload_array, main_origins = build_dataset_from_framelist(frame_list, season_setup=season_axis)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "af15d010",
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/plain": [
+ "(10000, 6, 64, 64)"
+ ]
+ },
+ "execution_count": 31,
+ "metadata": {},
+ "output_type": "execute_result"
+ }
+ ],
+ "source": [
+ "flu_payload_array.shape\n",
+ "# AVAILABLE_DATASETS = [\"20.8S79.2M\"]"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "acfa481c",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "flu_payload_array.to_netcdf(\"/Users/emprzy/Documents/work/miscellaneous/influpaint_data/TS_30S70M_2025-07-17.nc\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "c8f387cc",
+ "metadata": {},
+ "source": [
+ "Skipping visualization, since there are 6 channels "
+ ]
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "base",
+ "language": "python",
+ "name": "python3"
+ },
+ "language_info": {
+ "codemirror_mode": {
+ "name": "ipython",
+ "version": 3
+ },
+ "file_extension": ".py",
+ "mimetype": "text/x-python",
+ "name": "python",
+ "nbconvert_exporter": "python",
+ "pygments_lexer": "ipython3",
+ "version": "3.13.7"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/emily_train.ipynb b/emily_train.ipynb
new file mode 100644
index 0000000..171019d
--- /dev/null
+++ b/emily_train.ipynb
@@ -0,0 +1,148 @@
+{
+ "cells": [
+ {
+ "cell_type": "code",
+ "execution_count": 1,
+ "id": "4f2d8f69",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import torch\n",
+ "import sys\n",
+ "import os\n",
+ "from torch.utils.data import DataLoader"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 2,
+ "id": "99c59435",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Environment: Mac\n",
+ "Project root: /Users/emprzy/Documents/work/influpaint\n",
+ "Data path: /Users/emprzy/Documents/work/miscellaneous/influpaint_data/TS_30S70M_2025-07-17.nc\n"
+ ]
+ }
+ ],
+ "source": [
+ "if os.path.exists(\"/nas/longleaf/home/emprzy\"):\n",
+ " project_root = \"/nas/longleaf/home/emprzy/influpaint\"\n",
+ " # This matches the directory you just created\n",
+ " data_path = os.path.join(project_root, \"training_datasets/TS_30S70M_2025-07-17.nc\")\n",
+ "else:\n",
+ " project_root = \"/Users/emprzy/Documents/work/influpaint\"\n",
+ " data_path = \"/Users/emprzy/Documents/work/miscellaneous/influpaint_data/TS_30S70M_2025-07-17.nc\"\n",
+ "\n",
+ "if project_root not in sys.path:\n",
+ " sys.path.insert(0, project_root)\n",
+ "\n",
+ "# Point to the specific CoPaint folder\n",
+ "copaint_path = os.path.join(project_root, \"influpaint/batch/CoPaint4influpaint\")\n",
+ "if copaint_path not in sys.path:\n",
+ " sys.path.insert(0, copaint_path)\n",
+ "\n",
+ "print(f\"Environment: {'Longleaf' if 'nas' in project_root else 'Mac'}\")\n",
+ "print(f\"Project root: {project_root}\")\n",
+ "print(f\"Data path: {data_path}\")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 3,
+ "id": "4e9a67a1",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from influpaint.batch.scenarios import get_training_scenario, create_scenario_objects, print_available_scenarios\n",
+ "from influpaint.batch.config import transform_library\n",
+ "from influpaint.datasets import loaders as training_datasets"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 5,
+ "id": "df346a15",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Using device: cpu\n"
+ ]
+ }
+ ],
+ "source": [
+ "# season_setup = SeasonAxis.for_flusight(remove_us=True, remove_territories=True) \n",
+ "image_size = 64\n",
+ "channels = 6\n",
+ "batch_size=512\n",
+ "epochs=3000\n",
+ "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
+ "print(f\"Using device: {device}\")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "f7472d8a",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "created dataset with max [188028. 109653.8 109653.8 109653.8 109653.8 109653.8], full dataset has shape (10000, 6, 64, 64)\n",
+ "test passed: back and forth transformation are ok ✅\n"
+ ]
+ }
+ ],
+ "source": [
+ "scn_id = 868 # i868::m_U500cRx1224::ds_30S70M::tr_Sqrt::ri_No\n",
+ "experiment_name = \"emily_first_train\" # MLflow experiment name\n",
+ "scenario_spec = get_training_scenario(scn_id)\n",
+ "ddpm, dataset, transform, enrich, scaling_per_channel, data_mean, data_sd = create_scenario_objects(\n",
+ " scenario_spec, image_size, channels, batch_size, epochs, device) # PATCH: removed season_setup param (don't need it b/c it is only used for `datasets`, which i will set explicitly)\n",
+ "# dataset = training_datasets.FluDataset.from_xarray(\"/Users/emprzy/Documents/work/miscellaneous/influpaint_data/TS_30S70M_2025-07-17.nc\",channels=channels,)\n",
+ "# don't need ^this^ line because i modified create_scenario_objects()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "edf53cb8",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, drop_last=True)\n",
+ "losses = ddpm.train(dataloader, mlflow_logging=True)"
+ ]
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "base",
+ "language": "python",
+ "name": "python3"
+ },
+ "language_info": {
+ "codemirror_mode": {
+ "name": "ipython",
+ "version": 3
+ },
+ "file_extension": ".py",
+ "mimetype": "text/x-python",
+ "name": "python",
+ "nbconvert_exporter": "python",
+ "pygments_lexer": "ipython3",
+ "version": "3.13.7"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/influpaint/batch/config.py b/influpaint/batch/config.py
index 1b9fc64..f394445 100644
--- a/influpaint/batch/config.py
+++ b/influpaint/batch/config.py
@@ -40,46 +40,62 @@
}
def unet_library(image_size, channels):
- unet_spec = { "Rx124":
+ # unet_library() logic has been patched
+ base_width = image_size
+
+ unet_spec = {
+ "Rx124":
nn_blocks.Unet(
- dim=image_size,
+ dim=base_width,
+ init_dim=base_width,
channels=channels,
dim_mults=(1, 2, 4,),
+ resnet_block_groups=1,
use_convnext=False
),
"Cx124":
nn_blocks.Unet(
- dim=image_size,
+ dim=base_width,
+ init_dim=base_width,
channels=channels,
dim_mults=(1, 2, 4,),
+ resnet_block_groups=1,
use_convnext=True
),
"Rx1224":
nn_blocks.Unet(
- dim=image_size,
+ dim=base_width,
+ init_dim=base_width,
channels=channels,
dim_mults=(1, 2, 2, 4,),
+ resnet_block_groups=1,
use_convnext=False
),
"Cx1224":
nn_blocks.Unet(
- dim=image_size,
+ dim=base_width,
+ init_dim=base_width,
channels=channels,
dim_mults=(1, 2, 2, 4,),
+ resnet_block_groups=1,
use_convnext=True
),
"Rx12448":
nn_blocks.Unet(
- dim=image_size,
+ dim=base_width,
+ init_dim=base_width,
channels=channels,
dim_mults=(1, 2, 4, 4, 8,),
+ resnet_block_groups=1,
use_convnext=False
),
"Cx12448":
nn_blocks.Unet(
- dim=image_size,
+ dim=base_width,
+ init_dim=base_width,
channels=channels,
dim_mults=(1, 2, 4, 4, 8,),
+ resnet_block_groups=1,
use_convnext=True
),
}
diff --git a/influpaint/batch/scenarios.py b/influpaint/batch/scenarios.py
index 33a3b4f..076d2e4 100644
--- a/influpaint/batch/scenarios.py
+++ b/influpaint/batch/scenarios.py
@@ -8,6 +8,7 @@
import itertools
from .config import AVAILABLE_DDPMS,AVAILABLE_UNETS, AVAILABLE_DATASETS, AVAILABLE_TRANSFORMS, AVAILABLE_ENRICHMENTS, AVAILABLE_COPAINT_CONFIGS
from .config import CONFIG_BASELINE
+from ..datasets import loaders as training_datasets
@dataclass(frozen=True)
@@ -164,8 +165,10 @@ def get_inpainting_scenario(scenario_id: int) -> InpaintingScenario:
return scenarios[scenario_id]
-# Simple helper for research use
-def create_scenario_objects(scenario_spec: TrainingScenario, season_setup, image_size=64, channels=1, batch_size=512, epochs=800, device="cuda"):
+# --- PATCH ---
+# REMOVED season_axis as a required param (don't need it b/c i will overwrite `dataset` manually)
+# PATCH cascades into `config.py::get_dataset()` and `config.py::dataset_library()`
+def create_scenario_objects(scenario_spec: TrainingScenario, season_setup=None, image_size=64, channels=1, batch_size=512, epochs=800, device="cuda"):
"""Create actual objects from scenario spec - one function does everything"""
from .config import ddpm_library, unet_library, get_dataset, transform_library
import numpy as np
@@ -176,7 +179,8 @@ def create_scenario_objects(scenario_spec: TrainingScenario, season_setup, image
ddpm_spec = ddpm_library(image_size, channels, epochs, device, batch_size, unet=unet)
ddpm = ddpm_spec[scenario_spec.ddpm_name]
- dataset = get_dataset(scenario_spec.dataset_name, season_setup, channels)
+ # PATCH to avoid original dataset init
+ dataset = training_datasets.FluDataset.from_xarray("/Users/emprzy/Documents/work/miscellaneous/influpaint_data/TS_30S70M_2025-07-17.nc",channels=channels,)
# Create transforms
# scaling_per_channel = np.array(max(dataset.max_per_feature, gt1.gt_xarr.max(dim=["date", "place"])))
@@ -191,7 +195,7 @@ def create_scenario_objects(scenario_spec: TrainingScenario, season_setup, image
enrich = transform_enrich[scenario_spec.enrich_name]
# Configure dataset
- dataset.add_transform(
+ dataset.add_transform( # had to PATCH transforms.py to make this work
transform=transform["reg"],
transform_inv=transform["inv"],
transform_enrich=enrich,
@@ -199,6 +203,7 @@ def create_scenario_objects(scenario_spec: TrainingScenario, season_setup, image
)
return ddpm, dataset, transform, enrich, scaling_per_channel, data_mean, data_std
+# --- END PATCH ---
def print_available_scenarios():
diff --git a/influpaint/datasets/mixer.py b/influpaint/datasets/mixer.py
index d3cb6fc..c0452f1 100644
--- a/influpaint/datasets/mixer.py
+++ b/influpaint/datasets/mixer.py
@@ -68,9 +68,18 @@
import pandas as pd
import numpy as np
-from ..utils.season_axis import SeasonAxis
+# from ..utils.season_axis import SeasonAxis
from tqdm import tqdm
+# --- PATCHED: most functions in this file have been patched to accommodate my data ---
+AGE_COLUMNS = [
+ '0-130',
+ '0-4',
+ '18-49',
+ '5-17',
+ '50-64',
+ '65-130'
+ ]
def _validate_required_columns(df: pd.DataFrame, required_columns: list,
context: str) -> None:
@@ -100,7 +109,7 @@ def _validate_config_consistency(config: dict) -> tuple[bool, bool]:
return has_proportions, has_multipliers
-def build_frames(all_datasets_df: pd.DataFrame, config: dict, season_axis: SeasonAxis,
+def build_frames(all_datasets_df: pd.DataFrame, config: dict, season_axis: set[str],
fill_missing_locations: str = "error", scaling_distribution: np.ndarray = None) -> list:
"""
Build complete epidemic frames from hierarchical dataset structure.
@@ -121,7 +130,7 @@ def build_frames(all_datasets_df: pd.DataFrame, config: dict, season_axis: Seaso
- Values: Either {"multiplier": int} or {"proportion": float, "total": int}
- Optional: {"to_scale": bool} to enable frame scaling
- season_axis (SeasonAxis): Season axis object providing location definitions
+ season_axis (SeasonAxis): Season axis object providing location definitions (MODIFIED TO BE A set[str] OF LOCATIONS)
fill_missing_locations (str): Strategy for handling missing locations:
- "error": Fail if any expected location is missing (default)
@@ -155,8 +164,8 @@ def build_frames(all_datasets_df: pd.DataFrame, config: dict, season_axis: Seaso
"""
# Validate input dataframe
required_columns = ['datasetH1', 'datasetH2', 'fluseason', 'sample',
- 'location_code', 'season_week', 'value', 'week_enddate']
- _validate_required_columns(all_datasets_df, required_columns, "Input dataframe")
+ 'location_code', 'season_week', 'week_enddate'] + AGE_COLUMNS
+ _validate_required_columns(all_datasets_df, required_columns, "Input WIDE dataframe")
# Validate config references existing H1 datasets
available_h1 = set(all_datasets_df['datasetH1'].unique())
@@ -227,7 +236,7 @@ def build_frames(all_datasets_df: pd.DataFrame, config: dict, season_axis: Seaso
# Clean up frames by removing unnecessary metadata columns
cleaned_frames = []
- essential_columns = ['location_code', 'season_week', 'value', 'week_enddate', 'origin']
+ essential_columns = ['location_code', 'season_week', 'week_enddate', 'origin'] + AGE_COLUMNS
for frame in all_frames:
# Keep only essential columns that have actual data
@@ -294,7 +303,7 @@ def _calculate_explicit_multipliers(config: dict) -> dict:
def _build_h1_frames(h1_data: pd.DataFrame, h1_name: str, multiplier: int,
- season_axis: SeasonAxis, fill_missing_locations: str,
+ season_axis: set[str], fill_missing_locations: str,
all_datasets_df: pd.DataFrame = None, global_lookup: dict = None,
should_scale: bool = False, scaling_distribution: np.ndarray = None) -> list:
"""Build frames for a single H1 dataset with replication."""
@@ -342,11 +351,11 @@ def _build_h1_frames(h1_data: pd.DataFrame, h1_name: str, multiplier: int,
return frames
-def _pad_frame_complete(frame: pd.DataFrame, season_axis: SeasonAxis,
+def _pad_frame_complete(frame: pd.DataFrame, season_axis: set[str],
fill_missing_locations: str, all_datasets_df: pd.DataFrame = None,
global_lookup: dict = None) -> pd.DataFrame:
"""Ensure frame has complete weekly and location coverage."""
- expected_locations = set(season_axis.locations)
+ expected_locations = season_axis # CHANGE: expected locations is just the set i passed in
actual_locations = set(frame['location_code'].unique())
missing_locations = expected_locations - actual_locations
@@ -431,10 +440,11 @@ def _fill_missing_with_zeros(frame: pd.DataFrame, missing_locations: set) -> pd.
missing_row = {
'location_code': location,
'season_week': week,
- 'value': 0.0,
'origin': origin,
**metadata
}
+ for col in AGE_COLUMNS:
+ missing_row[col] = 0.0
# Add week_enddate if it exists in original frame
if ref_enddate is not None:
@@ -541,7 +551,7 @@ def _build_global_lookup_table(all_datasets_df: pd.DataFrame) -> dict:
(year_data['sample'] == first_combo['sample'])
]
lookup[location][h1_name]['by_year'][year] = {
- 'data': specific_data[['season_week', 'value', 'week_enddate']].copy(),
+ 'data': specific_data[['season_week', *AGE_COLUMNS, 'week_enddate']].copy(),
'h2': first_combo['datasetH2'],
'sample': first_combo['sample']
}
@@ -551,7 +561,7 @@ def _build_global_lookup_table(all_datasets_df: pd.DataFrame) -> dict:
for (year, sample), sample_data in sample_groups:
key = f"{year}_{sample}"
lookup[location][h1_name]['by_sample'][key] = {
- 'data': sample_data[['season_week', 'value', 'week_enddate']].copy(),
+ 'data': sample_data[['season_week', *AGE_COLUMNS, 'week_enddate']].copy(),
'h2': sample_data['datasetH2'].iloc[0],
'year': year,
'sample': sample
@@ -683,11 +693,17 @@ def _pad_single_location(frame: pd.DataFrame, location: str) -> pd.DataFrame:
# Handle empty input
if frame.empty:
# Create empty frame for all weeks - metadata will be preserved by _preserve_columns later
- missing_data = [{
- "season_week": week,
- "location_code": location,
- "value": 0
- } for week in range(1, 54)]
+ missing_data = []
+ for week in range(1, 54):
+ row = {
+ "season_week": week,
+ "location_code": location,
+ }
+ # Add a 0.0 value for every age column
+ for col in AGE_COLUMNS:
+ row[col] = 0.0
+
+ missing_data.append(row)
return pd.DataFrame(missing_data)
# Get min/max weeks if data exists
@@ -706,26 +722,22 @@ def _pad_single_location(frame: pd.DataFrame, location: str) -> pd.DataFrame:
if not frame.empty:
first_row = frame.iloc[0]
for col in frame.columns:
- if col not in ["season_week", "location_code", "value"]:
+ if col not in ["season_week", "location_code"] + AGE_COLUMNS:
metadata[col] = first_row[col]
# Batch create missing rows for better performance
missing_rows = []
for week in missing_weeks:
- # Determine fill value based on position
if week < min_week or week > max_week:
- new_value = 0 # External gaps filled with zeros
+ new_values = {col: 0.0 for col in AGE_COLUMNS}
else:
- # Internal gaps filled with previous week's value
previous_week = frame[frame["season_week"] == week-1]
- new_value = previous_week["value"].values[0] if not previous_week.empty else 0
-
- missing_row = {
- "season_week": week,
- "location_code": location,
- "value": new_value,
- **metadata # Include all metadata from original frame
- }
+ if not previous_week.empty:
+ new_values = {col: previous_week[col].values[0] for col in AGE_COLUMNS}
+ else:
+ new_values = {col: 0.0 for col in AGE_COLUMNS}
+
+ missing_row = {**new_values, "season_week": week, "location_code": location, **metadata}
missing_rows.append(missing_row)
# Single concat instead of multiple
@@ -736,7 +748,7 @@ def _pad_single_location(frame: pd.DataFrame, location: str) -> pd.DataFrame:
return frame.sort_values("season_week").reset_index(drop=True)
-def _apply_frame_scaling(frame: pd.DataFrame, scaling_distribution: np.ndarray) -> pd.DataFrame:
+def _apply_frame_scaling(frame: pd.DataFrame, scaling_distribution: np.ndarray) -> pd.DataFrame: # did not need to modify for my wide dataset b/c i'm not scaling
"""
Apply scaling to a frame based on US peak sum scaling distribution.
diff --git a/influpaint/datasets/read_datasources.py b/influpaint/datasets/read_datasources.py
index f60d429..aa904f3 100644
--- a/influpaint/datasets/read_datasources.py
+++ b/influpaint/datasets/read_datasources.py
@@ -1,6 +1,7 @@
import pandas as pd
import numpy as np
-from helpers.delphi_epidata import Epidata
+# from helpers.delphi_epidata import Epidata
+from delphi_epidata import Epidata
from ..utils.season_axis import SeasonAxis
import xarray as xr
diff --git a/influpaint/datasets/transforms.py b/influpaint/datasets/transforms.py
index ffaaadf..6e0d637 100644
--- a/influpaint/datasets/transforms.py
+++ b/influpaint/datasets/transforms.py
@@ -12,13 +12,17 @@ def transform_randomscale(image, max, min):
return image * scale
-def transform_channelwisescale(
- image, scale
-): # TODO write for three channel like it was above
+# PATCHED to deal with 3d frames
+def transform_channelwisescale(image, scale):
+ if isinstance(scale, np.ndarray) and scale.ndim == 1:
+ scale = scale.reshape(-1, 1, 1)
return image * scale
+# PATCHED to deal with 3d frames
def transform_channelwisescale_inv(image, scale):
+ if isinstance(scale, np.ndarray) and scale.ndim == 1:
+ scale = scale.reshape(-1, 1, 1)
return image / scale
diff --git a/influpaint/utils/converters.py b/influpaint/utils/converters.py
index e9d50ec..ae95d19 100644
--- a/influpaint/utils/converters.py
+++ b/influpaint/utils/converters.py
@@ -1,6 +1,6 @@
import pandas as pd
import numpy as np
-from helpers.delphi_epidata import Epidata
+from delphi_epidata import Epidata
from ..utils.season_axis import SeasonAxis
import xarray as xr
@@ -70,28 +70,48 @@ def dataframe_to_xarray(
return df_xarr
-
+# --- PATCHED to accommodate 6 channels ---
def dataframe_to_arraylist(
- df: pd.DataFrame, season_setup: SeasonAxis = None, value_column="value",
-) -> np.ndarray:
-
+ df: pd.DataFrame,
+ locations: set[str],
+ age_columns: list[str],
+) -> list:
+ """
+ Standardizes data into a list of arrays shaped (6, 64, 64).
+ Uses pivot_table to resolve duplicates and matches the original tower pivot logic.
+ """
samples = []
+ sorted_locs = sorted(list(locations))
+ unique_seasons = sorted(df["fluseason"].unique())
- df_piv = df.pivot(
- columns="location_code",
- values=value_column,
- index=["fluseason", "season_week"],
- )
- for season in df_piv.index.unique(level="fluseason"):
- array = df_piv.loc[season][
- season_setup.locations
- ].sort_index().to_numpy() # make sure order is right w.r.t season_setup locations and the time is right
- # TODO: should give an error when dates are missing because it would be missaligned
+ towers = {}
+ for age_group in age_columns:
+ pivoted_tower = df.pivot_table(
+ index=["fluseason", "season_week"],
+ columns="location_code",
+ values=age_group,
+ aggfunc='mean'
+ )
+
+ towers[age_group] = pivoted_tower.reindex(columns=sorted_locs)
- array[np.isnan(array)] = 0 # replace NaNs with 0
+ for season in unique_seasons:
+ channel_slices = []
+
+ for age_group in age_columns:
+ try:
+ season_slice = towers[age_group].loc[season].sort_index().to_numpy()
+ except KeyError:
+ season_slice = np.zeros((53, len(sorted_locs)))
- samples.append(
- np.array([padto64x64(array)])
- ) # pad to 64x64 and add a dimension for channel
+ season_slice = np.nan_to_num(season_slice, nan=0.0)
+
+ # Pad from (53, 52) to (64, 64)
+ padded_array = padto64x64(season_slice)
+
+ channel_slices.append(padded_array)
+ # channel stacking
+ # Stack 6 age groups to get (6, 64, 64) for a sample
+ samples.append(np.stack(channel_slices))
return samples
\ No newline at end of file
diff --git a/influpaint/utils/ground_truth.py b/influpaint/utils/ground_truth.py
index 7b12757..e1ee0c0 100644
--- a/influpaint/utils/ground_truth.py
+++ b/influpaint/utils/ground_truth.py
@@ -7,6 +7,10 @@
from tqdm.auto import tqdm
from . import season_axis
from .season_axis import SeasonAxis
+import requests
+import time
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
import numpy as np
@@ -19,196 +23,624 @@
from ..datasets import read_datasources
from . import converters
-
-def pad_dataframe(df, season_setup):
- # Make sure gt_df and gt_df_final have values (even if NaN) for all dates in the season
- date_range = pd.date_range(start=df.week_enddate.min(), periods=52, freq='W-SAT')
- locations = df.location_code.unique()
- # Create expanded dataframe with all combinations of dates and locations
- expanded_df = pd.DataFrame([(d, l) for d in date_range for l in locations],
- columns=['week_enddate', 'location_code'])
-
- # Calculate the season columns
- expanded_df['fluseason'] = expanded_df.week_enddate.apply(season_setup.get_fluseason_year)
- expanded_df['fluseason_fraction'] = expanded_df.week_enddate.apply(season_setup.get_fluseason_fraction)
- expanded_df['season_week'] = expanded_df.week_enddate.apply(season_setup.get_season_week)
-
- # Merge with original data to get values where they exist
- padded_df = expanded_df.merge(
- df[['week_enddate', 'location_code', 'value']],
- on=['week_enddate', 'location_code'],
- how='left'
- )
- return padded_df
+# map for static method _fetch_nhsn_data()
+LOCATIONS_ABBREV = [
+ 'AL', 'AK', 'AZ', 'AR', 'CA', 'CO', 'CT', 'DE', 'DC', 'FL', 'GA',
+ 'HI', 'ID', 'IL', 'IN', 'IA', 'KS', 'KY', 'LA', 'ME', 'MD', 'MA',
+ 'MI', 'MN', 'MS', 'MO', 'MT', 'NE', 'NV', 'NH', 'NJ', 'NM', 'NY',
+ 'NC', 'ND', 'OH', 'OK', 'OR', 'PA', 'RI', 'SC', 'SD', 'TN', 'TX',
+ 'UT', 'VT', 'VA', 'WA', 'WV', 'WI', 'WY', 'US'
+ ]
+
+
+def pad_dataframe(df, season_setup, nhsn: bool = False, mask_date=None):
+ # --- START EMILY PATCH ---
+ if nhsn:
+ AGE_CHANNELS = ['0-4', '5-17', '18-49', '50-64', '0-130', '65-130']
+ date_range = pd.date_range(start=df.week_enddate.min(), periods=52, freq='W-SAT')
+ locations = season_setup
+ expanded_df = pd.DataFrame([(d, l) for d in date_range for l in locations],
+ columns=['week_enddate', 'location_code'])
+ padded_df = expanded_df.merge(
+ df,
+ on=['week_enddate', 'location_code'],
+ how='left'
+ )
+ # only check for NaNs on or before the mask_date
+ if mask_date is not None:
+ check_df = padded_df[padded_df['week_enddate'] <= pd.to_datetime(mask_date)]
+
+ if check_df[AGE_CHANNELS].isna().any().any():
+ missing = check_df[check_df[AGE_CHANNELS].isna().any(axis=1)].iloc[0]
+ raise ValueError(
+ f"STRICT CHECK FAILED: NHSN data is missing for {missing['location_code']} "
+ f"on {missing['week_enddate'].date()}. This date is before or on your "
+ f"mask_date ({mask_date.date()}), so data must be present."
+ )
+ # basically, i don't have a contingency plan yet for if there are NaNs, i just
+ # wanted to check if they existed or not (seems like they don't). Filling them
+ # in is non-trivial without a SeasonAxis object.
+
+ return padded_df
+ # --- END EMILY PATCH ---
+ else:
+ # Make sure gt_df and gt_df_final have values (even if NaN) for all dates in the season
+ date_range = pd.date_range(start=df.week_enddate.min(), periods=52, freq='W-SAT')
+ locations = df.location_code.unique()
+ # Create expanded dataframe with all combinations of dates and locations
+ expanded_df = pd.DataFrame([(d, l) for d in date_range for l in locations],
+ columns=['week_enddate', 'location_code'])
+
+ # Calculate the season columns
+ expanded_df['fluseason'] = expanded_df.week_enddate.apply(season_setup.get_fluseason_year)
+ expanded_df['fluseason_fraction'] = expanded_df.week_enddate.apply(season_setup.get_fluseason_fraction)
+ expanded_df['season_week'] = expanded_df.week_enddate.apply(season_setup.get_season_week)
+
+ # Merge with original data to get values where they exist
+ padded_df = expanded_df.merge(
+ df[['week_enddate', 'location_code', 'value']],
+ on=['week_enddate', 'location_code'],
+ how='left'
+ )
+ return padded_df
class GroundTruth():
- def __init__(self, season_first_year: str,
- data_date: datetime.datetime,
- mask_date: datetime.datetime,
- from_final_data:bool=False,
- channels=1,
- image_size=64,
- nogit=False,
- payload=None,
- payload_season_first_year=None,
- dataset_coords: xr.core.coordinates.DataArrayCoordinates=None):
- self.season_first_year = season_first_year
- self.data_date = data_date
- self.mask_date = mask_date
- self.channels = channels
- self.image_size=image_size
-
-
- if not nogit: self.git_checkout_data_rev(target_date=None)
-
- self.season_setup = SeasonAxis.for_flusight(remove_territories=True, remove_us=True)
-
- flusight = read_datasources.get_from_epidata(dataset=f"flusight{self.season_first_year}", season_setup=self.season_setup, write=False)
- flusight = self.season_setup.add_season_columns(flusight, do_fluseason_year=True)
- gt_df_final = flusight[flusight["fluseason"] == int(self.season_first_year)]
+ def __init__(
+ self,
+ season_first_year: str,
+ gt_df: pd.DataFrame,
+ gt_df_final: pd.DataFrame,
+ mask_date: datetime.datetime,
+ season_setup: SeasonAxis | list, # changed for Emily's purposes to also be a list
+ channels=1,
+ image_size=64,
+ previous_data=None,
+ dataset_coords: xr.core.coordinates.DataArrayCoordinates = None,
+ nhsn: bool = False
+ ):
+ # --- START EMILY PATCH --- mostly the same pipeline, some small adjustments
+ if nhsn:
+ AGE_CHANNELS = ['0-4', '5-17', '18-49', '50-64', '0-130', '65-130']
+ self.season_first_year = season_first_year
+ self.mask_date = pd.to_datetime(mask_date)
+ self.channels = channels
+ self.image_size = image_size
+ self.season_setup = season_setup # PATCH, option to pass season_setup as a list of locs instead of a SeasonAxis object
+
+ self.gt_df = gt_df
+ self.gt_df_final = gt_df_final
+
+ # note that previous data should be None always for NHSN processing
+ if previous_data is None:
+ previous_data = []
+ if isinstance(previous_data, list):
+ previous_data = pd.concat(previous_data, ignore_index=True) if previous_data else pd.DataFrame()
+ self.previous_data = previous_data
+
+ # note that dataset_coords should be None always for NHSN processing
+ if dataset_coords is not None:
+ pass # dont need anything here
+
+ self.gt_df = self.gt_df[self.gt_df["location_code"].isin(self.season_setup)]
+ self.gt_df_final = self.gt_df_final[self.gt_df_final["location_code"].isin(self.season_setup)]
+
+ # note that previous_data should be empty always for NHSN (don't need this logic)
+ if not self.previous_data.empty:
+ # self.previous_data = self.season_setup.add_season_columns(self.previous_data, do_fluseason_year=True)
+ pass
+
+ self.gt_df = pad_dataframe(self.gt_df, self.season_setup, nhsn=True, mask_date=self.mask_date)
+ self.gt_df_final = pad_dataframe(self.gt_df_final, self.season_setup, nhsn=True, mask_date=self.mask_date)
+
+ # If the last data_point is not in the last week, we need to update the mask to be in the week after the last data point
+ last_non_nan_datadate = self.gt_df.week_enddate[self.gt_df['0-130'].notna()].max().to_pydatetime()
+ if self.mask_date > last_non_nan_datadate + datetime.timedelta(days=7):
+ self.mask_date = last_non_nan_datadate + datetime.timedelta(days=2)
+ print(f" WARNING: mask_date is after last non-NaN data date, setting mask_date to {self.mask_date}")
+
+ self.gt_xarr = converters.dataframe_to_xarray_nhsn(
+ self.gt_df,
+ season_setup=self.season_setup,
+ xarray_name="gt_NHSN_incidHosp",
+ age_channels=AGE_CHANNELS,
+ pad_to=self.image_size
+ )
+
+ self.gt_final_xarr = converters.dataframe_to_xarray_nhsn(
+ self.gt_df_final,
+ season_setup=self.season_setup,
+ xarray_name="gt_NHSN_incidHosp_final",
+ age_channels=AGE_CHANNELS,
+ pad_to=self.image_size
+ )
+ # Find the largest index of the data dates that are before the mask date
+ dates = pd.to_datetime(self.gt_xarr.coords["date"].values)
+ self.inpaintfrom_idx = sum(dates < self.mask_date)
+
+ self.gt_keep_mask = np.ones((channels, image_size, image_size))
+ self.gt_keep_mask[:, self.inpaintfrom_idx :, :] = 0
+
+ print(f"Masking, >> {self.inpaintfrom_idx} weeks already in data, inpainting the next ones")
+ # --- END EMILY PATCH ---
- if from_final_data:
- gt_df = gt_df_final.copy()
else:
- if not nogit: self.git_checkout_data_rev(target_date=data_date)
- flusight = read_datasources.get_from_epidata(dataset=f"flusight{self.season_first_year}", season_setup=self.season_setup, write=False)
- flusight = self.season_setup.add_season_columns(flusight, do_fluseason_year=True)
- gt_df = flusight[flusight["fluseason"] == int(self.season_first_year)]
- if not nogit: self.git_checkout_data_rev(target_date=None)
-
-
- self.gt_df = gt_df[gt_df["location_code"].isin(self.season_setup.locations)]
- self.gt_df_final = gt_df_final[gt_df_final["location_code"].isin(self.season_setup.locations)]
-
- # generates past data
- past_data_1 = read_datasources.get_from_epidata(dataset=f"flusight2024", season_setup=self.season_setup, write=False)
- past_data_2 = read_datasources.get_from_epidata(dataset=f"flusight2024", season_setup=self.season_setup, write=False)
-
- # Add season columns to past data
- past_data_1 = self.season_setup.add_season_columns(past_data_1, do_fluseason_year=True)
- past_data_2 = self.season_setup.add_season_columns(past_data_2, do_fluseason_year=True)
-
- self.previous_data = [past_data_1, past_data_2]
-
-
-
- self.gt_df = pad_dataframe(self.gt_df, self.season_setup)
- self.gt_df_final = pad_dataframe(self.gt_df_final, self.season_setup)
-
- last_non_nan_datadate = self.gt_df.week_enddate[self.gt_df.value.notna()].max().to_pydatetime()
- # If the last data_point is not in the last week, we need to update the mask to be in the week after the last data point
- if self.mask_date > last_non_nan_datadate + datetime.timedelta(days=7):
- self.mask_date = last_non_nan_datadate + datetime.timedelta(days=2)
- print(f" WARNING: mask_date is after last non-NaN data date, setting mask_date to {self.mask_date}")
-
-
- if payload is not None:
- if payload_season_first_year is None:
- payload_season_first_year = season_first_year
- import dataset_mixer
- payload = self.season_setup.add_season_columns(payload, do_fluseason_year=True)
- this_payload = payload[payload["fluseason"] == int(payload_season_first_year)]
- self.gt_df = pd.concat([self.gt_df, this_payload], ignore_index=True)
- self.gt_df_final = pd.concat([self.gt_df_final, this_payload], ignore_index=True)
- self.previous_data.append(payload)
- location_codes = self.gt_df.location_code.unique()
- new_locations = pd.DataFrame({"location_code": sorted(location_codes)})
- # Ensure location_code is of type string
- new_locations['location_code'] = new_locations['location_code'].astype(str)
- # Merge with season_setup.locations_df to get the location names
- new_locations = new_locations.merge(self.season_setup.locations_df,
- on='location_code',
- how='left')
-
- # Fill missing location names with the location code
- new_locations['location_name'] = new_locations['location_name'].fillna(new_locations['location_code'])
- new_locations = new_locations[['location_code', 'location_name']]
- self.season_setup.update_locations(new_locations)
-
- if dataset_coords is not None:
- # change the flusetup locations to be in the same order as flu_payload_array.coords["place"]
- self.season_setup.reorder_locations(list(dataset_coords["place"].values))
-
- # Concatenate all previous data and ensure it has season columns
- self.previous_data = pd.concat(self.previous_data, ignore_index=True).drop_duplicates()
- self.previous_data = self.season_setup.add_season_columns(self.previous_data, do_fluseason_year=True)
-
- self.gt_xarr = converters.dataframe_to_xarray(self.gt_df, season_setup=self.season_setup,
- xarray_name = "gt_flusight_incidHosp",
- xarrax_features = "incidHosp")
-
- self.gt_final_xarr = converters.dataframe_to_xarray(self.gt_df_final, season_setup=self.season_setup,
- xarray_name = "gt_flusight_incidHos_final",
- xarrax_features = "incidHosp")
-
- # Find the largest index of the data dates that are before the mask date
- dates = pd.to_datetime(self.gt_xarr.coords['date'].values)
- self.inpaintfrom_idx = sum(dates < self.mask_date)
-
- self.gt_keep_mask = np.ones((channels,image_size,image_size))
- self.gt_keep_mask[:,self.inpaintfrom_idx:,:] = 0
-
- print(f"Masking, >> {self.inpaintfrom_idx} weeks already in data, inpainting the next ones")
-
-
- def git_checkout_data_rev(self, target_date=None):
+ self.season_first_year = season_first_year
+ self.mask_date = pd.to_datetime(mask_date)
+ self.channels = channels
+ self.image_size = image_size
+ self.season_setup = season_setup
+
+ self.gt_df = gt_df
+ self.gt_df_final = gt_df_final
+
+ if previous_data is None:
+ previous_data = []
+ if isinstance(previous_data, list):
+ previous_data = pd.concat(previous_data, ignore_index=True) if previous_data else pd.DataFrame()
+ self.previous_data = previous_data
+
+ if dataset_coords is not None:
+ # change the flusetup locations to be in the same order as flu_payload_array.coords["place"]
+ self.season_setup.reorder_locations(list(dataset_coords["place"].values))
+
+ self.gt_df = self.gt_df[self.gt_df["location_code"].isin(self.season_setup.locations)]
+ self.gt_df_final = self.gt_df_final[self.gt_df_final["location_code"].isin(self.season_setup.locations)]
+
+ if not self.previous_data.empty:
+ self.previous_data = self.season_setup.add_season_columns(self.previous_data, do_fluseason_year=True)
+
+ self.gt_df = pad_dataframe(self.gt_df, self.season_setup)
+ self.gt_df_final = pad_dataframe(self.gt_df_final, self.season_setup)
+
+ last_non_nan_datadate = self.gt_df.week_enddate[self.gt_df.value.notna()].max().to_pydatetime()
+ # If the last data_point is not in the last week, we need to update the mask to be in the week after the last data point
+ if self.mask_date > last_non_nan_datadate + datetime.timedelta(days=7):
+ self.mask_date = last_non_nan_datadate + datetime.timedelta(days=2)
+ print(f" WARNING: mask_date is after last non-NaN data date, setting mask_date to {self.mask_date}")
+
+ self.gt_xarr = converters.dataframe_to_xarray(
+ self.gt_df,
+ season_setup=self.season_setup,
+ xarray_name="gt_flusight_incidHosp",
+ xarrax_features="incidHosp",
+ pad_to=image_size,
+ )
+
+ self.gt_final_xarr = converters.dataframe_to_xarray(
+ self.gt_df_final,
+ season_setup=self.season_setup,
+ xarray_name="gt_flusight_incidHos_final",
+ xarrax_features="incidHosp",
+ pad_to=image_size,
+ )
+
+ # Find the largest index of the data dates that are before the mask date
+ dates = pd.to_datetime(self.gt_xarr.coords["date"].values)
+ self.inpaintfrom_idx = sum(dates < self.mask_date)
+
+ self.gt_keep_mask = np.ones((channels, image_size, image_size))
+ self.gt_keep_mask[:, self.inpaintfrom_idx :, :] = 0
+
+ print(f"Masking, >> {self.inpaintfrom_idx} weeks already in data, inpainting the next ones")
+
+ @staticmethod
+ def _git_checkout_repo_rev(repo_path, target_date=None, main_branch="main"):
import pygit2
- if self.season_first_year == "2023":
- repo_path = "Flusight/2023-2024/FluSight-forecast-hub-official/"
- main_branch = "main"
- elif self.season_first_year == "2022":
- repo_path = "Flusight/2022-2023/FluSight-forecast-hub-official/"
- main_branch = "master"
- elif self.season_first_year == "2024":
- repo_path = "Flusight/2024-2025/FluSight-forecast-hub-official/"
- main_branch = "main"
- print(repo_path)
-
- # Open the existing repository
+
repo = pygit2.Repository(repo_path)
if target_date is not None:
- # Find the commit closest to the target date
closest_commit = None
for commit in repo.walk(repo.head.target, pygit2.GIT_SORT_TIME):
if commit.commit_time <= target_date.timestamp():
closest_commit = commit
break
- # Check out the commit
if closest_commit:
repo.checkout_tree(closest_commit.tree)
repo.set_head(closest_commit.id)
- print(f"Checked out commit on {target_date} (SHA: {closest_commit.id}, {commit.commit_time}) for repo {repo_path}")
+ print(
+ f"Checked out commit on {target_date} (SHA: {closest_commit.id}, {commit.commit_time}) for repo {repo_path}"
+ )
else:
- print("ERROR: No commit found for the specified date on repo {repo_path}.")
+ print(f"ERROR: No commit found for the specified date on repo {repo_path}.")
else:
- repo.checkout("refs/heads/" + main_branch)
+ repo.checkout("refs/heads/" + main_branch)
print(f"Restored git repo {repo_path}")
+ @staticmethod
+ def _flusight_repo_info(season_first_year: str):
+ if season_first_year == "2023":
+ return "Flusight/2023-2024/FluSight-forecast-hub-official/", "main"
+ if season_first_year == "2022":
+ return "Flusight/2022-2023/FluSight-forecast-hub-official/", "master"
+ if season_first_year == "2024":
+ return "Flusight/2024-2025/FluSight-forecast-hub-official/", "main"
+ if season_first_year == "2025":
+ return "Flusight/2024-2025/FluSight-forecast-hub-official/", "main"
+ raise ValueError(f"Unsupported FluSight season_first_year: {season_first_year}")
+
+ # METHOD ADDITION (PATCH) to help get NHSN data in the right season week arrangement
+ @staticmethod
+ def _get_season_fraction(ts, start_month: int, start_day: int): # modified from season_axis.py
+ if pd.isna(ts):
+ return float('nan')
+ if isinstance(ts, datetime.datetime):
+ ts = ts.date()
+ try:
+ season_start = datetime.date(ts.year, start_month, start_day)
+ except AttributeError:
+ season_start = datetime.date(ts.year, start_month, start_day)
+ if ts < season_start:
+ season_start = datetime.date(ts.year - 1, start_month, start_day)
+
+ days_since_start = (ts - season_start).days
+ return days_since_start / 365
+
+ @staticmethod
+ def _fetch_nhsn_data(resource_id: str = 'ua7e-t2fy') -> pd.DataFrame:
+ """Distilled logic for fetching, cleaning, and mapping NHSN data."""
+ data_url = f"https://data.cdc.gov/resource/{resource_id}.json"
+ metadata_url = f"https://data.cdc.gov/api/views/{resource_id}.json"
+ session = requests.Session()
+ retries = Retry(total=5, backoff_factor=1, status_forcelist=[500, 502, 503, 504])
+ session.mount('https://', HTTPAdapter(max_retries=retries))
+
+ # paginated retrieval
+ all_data, offset, batch_size = [], 0, 1000
+ while True:
+ params = {"$limit": batch_size, "$offset": offset}
+ data_response = session.get(data_url, params=params, timeout=30)
+ data_response.raise_for_status()
+ batch_data = data_response.json()
+ if not batch_data:
+ break
+ all_data.extend(batch_data)
+ offset += batch_size
+ time.sleep(0.1)
+
+ data = pd.DataFrame(all_data).drop(columns=['respseason'], errors='ignore')
+
+ non_numeric = ['jurisdiction', 'weekendingdate']
+ for col in data.columns:
+ if col not in non_numeric:
+ data[col] = pd.to_numeric(data[col], errors='coerce')
+
+ data = data.replace(np.nan, value=None)
+ data.loc[data['jurisdiction'].str.lower() == 'usa', 'jurisdiction'] = 'US'
+ data = data[data['jurisdiction'].isin(LOCATIONS_ABBREV)].copy()
+ data['weekendingdate'] = pd.to_datetime(data['weekendingdate']).dt.strftime('%Y-%m-%d')
+ cdc_metadata = requests.get(metadata_url).json()
+ column_name_map = {col['fieldName']: col['name'] for col in cdc_metadata['columns']}
+ data = data.rename(columns=column_name_map, errors="ignore")
+
+ return data.sort_values(by=['Geographic aggregation', 'Week Ending Date'])
+
+
+ @classmethod
+ def for_flusight(
+ cls,
+ season_first_year: str,
+ data_date: datetime.datetime,
+ mask_date: datetime.datetime,
+ from_final_data: bool = False,
+ channels=1,
+ image_size=64,
+ nogit=False,
+ payload=None,
+ payload_season_first_year=None,
+ dataset_coords: xr.core.coordinates.DataArrayCoordinates = None,
+ ):
+ data_date = pd.to_datetime(data_date)
+ repo_path, main_branch = cls._flusight_repo_info(season_first_year)
+ if not nogit:
+ cls._git_checkout_repo_rev(repo_path, target_date=None, main_branch=main_branch)
+
+ season_setup = SeasonAxis.for_flusight(remove_territories=True, remove_us=True)
+ flusight = read_datasources.get_from_epidata(
+ dataset=f"flusight{season_first_year}", season_setup=season_setup, write=False
+ )
+ flusight = season_setup.add_season_columns(flusight, do_fluseason_year=True)
+ gt_df_final = flusight[flusight["fluseason"] == int(season_first_year)]
+
+ if from_final_data:
+ gt_df = gt_df_final.copy()
+ else:
+ if not nogit:
+ cls._git_checkout_repo_rev(repo_path, target_date=data_date, main_branch=main_branch)
+ flusight = read_datasources.get_from_epidata(
+ dataset=f"flusight{season_first_year}", season_setup=season_setup, write=False
+ )
+ flusight = season_setup.add_season_columns(flusight, do_fluseason_year=True)
+ gt_df = flusight[flusight["fluseason"] == int(season_first_year)]
+ if not nogit:
+ cls._git_checkout_repo_rev(repo_path, target_date=None, main_branch=main_branch)
+
+ previous_data = []
+ for past_year in [int(season_first_year) - 1, int(season_first_year) - 2]:
+ try:
+ past_df = read_datasources.get_from_epidata(
+ dataset=f"flusight{past_year}", season_setup=season_setup, write=False
+ )
+ previous_data.append(past_df)
+ except Exception as exc:
+ print(f" WARNING: could not load flusight{past_year} for historical data: {exc}")
+
+ if payload is not None:
+ if payload_season_first_year is None:
+ payload_season_first_year = season_first_year
+ payload = season_setup.add_season_columns(payload, do_fluseason_year=True)
+ this_payload = payload[payload["fluseason"] == int(payload_season_first_year)]
+ gt_df = pd.concat([gt_df, this_payload], ignore_index=True)
+ gt_df_final = pd.concat([gt_df_final, this_payload], ignore_index=True)
+ previous_data.append(payload)
+
+ location_codes = gt_df.location_code.unique()
+ new_locations = pd.DataFrame({"location_code": sorted(location_codes)})
+ new_locations["location_code"] = new_locations["location_code"].astype(str)
+ new_locations = new_locations.merge(
+ season_setup.locations_df, on="location_code", how="left"
+ )
+ new_locations["location_name"] = new_locations["location_name"].fillna(
+ new_locations["location_code"]
+ )
+ new_locations = new_locations[["location_code", "location_name"]]
+ season_setup.update_locations(new_locations)
+
+ return cls(
+ season_first_year=season_first_year,
+ gt_df=gt_df,
+ gt_df_final=gt_df_final,
+ mask_date=mask_date,
+ season_setup=season_setup,
+ channels=channels,
+ image_size=image_size,
+ previous_data=previous_data,
+ dataset_coords=dataset_coords, # TODO, need to ensure that the week and location coords are the same for this as they are for the training data
+ )
+
+ @classmethod
+ def from_nhsn(
+ cls,
+ season_first_year: int,
+ # data_date: datetime.datetime, # don't need this param for NHSN; no vintaging yet
+ mask_date: datetime.datetime,
+ channels=6,
+ image_size=64
+ # that's all the args we need i believe. these i removed:
+ # from_final_data: bool = False,
+ # nogit = False,
+ # payload = None,
+ # payload_season_first_year = None,
+ # dataset_coords: xr.core.coordinates.DataArrayCoordinates = None,
+ ):
+ pd.set_option('future.no_silent_downcasting', True)
+ # fetcht the data
+ data = cls._fetch_nhsn_data(resource_id='ua7e-t2fy')
+ # only keep the columns we need
+ data = data[[
+ "Week Ending Date",
+ "Geographic aggregation",
+ "Number of Pediatric Influenza Admissions, 0-4 years",
+ "Number of Pediatric Influenza Admissions, 5-17 years",
+ "Total Pediatric Influenza Admissions",
+ "Number of Adult Influenza Admissions, 18-49 years",
+ "Number of Adult Influenza Admissions, 50-64 years",
+ "Number of Adult Influenza Admissions, 65-74 years",
+ "Number of Adult Influenza Admissions, 75+ years",
+ "Total Adult Influenza Admissions",
+ ]].copy()
+ # renaming, removing, and small calculations
+ data = data.fillna(0) # fill all Na values w/ 0
+ data['0-130'] = data["Total Pediatric Influenza Admissions"] + data["Total Adult Influenza Admissions"]
+ data['65+'] = data['Number of Adult Influenza Admissions, 65-74 years'] + data['Number of Adult Influenza Admissions, 75+ years']
+ data = data.rename(columns={
+ 'Week Ending Date': 'week_enddate',
+ 'Number of Pediatric Influenza Admissions, 0-4 years': '0-4',
+ 'Number of Pediatric Influenza Admissions, 5-17 years': '5-17',
+ 'Number of Adult Influenza Admissions, 18-49 years': '18-49',
+ 'Number of Adult Influenza Admissions, 50-64 years': '50-64',
+ })
+ data = data.drop(columns=['Number of Adult Influenza Admissions, 65-74 years', 'Number of Adult Influenza Admissions, 75+ years', 'Total Pediatric Influenza Admissions', 'Total Adult Influenza Admissions'])
+ # cast week_enddate column as date
+ data['week_enddate'] = pd.to_datetime(data['week_enddate'])
+ # add fluseason column, filter for just the year we want
+ data['fluseason'] = data['week_enddate'].dt.year
+ data.loc[data['week_enddate'].dt.month < 8, 'fluseason'] -= 1
+ data = data[data['fluseason'].isin([season_first_year])]
+ # add location_code, sample, datasetH1, datasetH2 columns
+ locations = pd.read_csv('/proj/jlessler/projects/emprzy_influpaint/age_channels/influpaint/locations.csv') #USED ABSOLUTE PATH! Would have to change if others used
+ data = data.merge(
+ locations[['abbreviation', 'location']],
+ left_on='Geographic aggregation',
+ right_on='abbreviation',
+ how='left'
+ )
+ data = data.rename(columns={'location': 'location_code'}).drop(columns=['abbreviation', 'Geographic aggregation'])
+ data['sample'] = "1"
+ data['datasetH1'] = 'NHSN'
+ data['datasetH2'] = 'NHSN'
+ # add season_week column
+ data = data.sort_values(['location_code', 'fluseason', 'week_enddate'])
+ data['season_week'] = data.groupby(['location_code', 'fluseason']).cumcount() + 1
+ # get fluseason_fraction (using static method i created above)
+ data['fluseason_fraction'] = data['week_enddate'].apply(
+ cls._get_season_fraction,
+ start_month=10,
+ start_day=1
+ )
+ # melt NHSN age grouping columns into singular `age_group`
+ data = data.rename(columns={'65+': '65-130'})
+ age_columns = ['0-4', '5-17', '18-49', '50-64', '0-130', '65-130']
+ id_vars = [col for col in data.columns if col not in age_columns]
+ data = data.melt(
+ id_vars=id_vars,
+ value_vars=age_columns,
+ var_name='age_group',
+ value_name='value'
+ )
+ # ensure correct dtypes for all columns
+ data['week_enddate'] = pd.to_datetime(data['week_enddate'], errors='coerce')
+ data = data.astype({
+ 'location_code': str,
+ 'value': 'float64',
+ 'fluseason_fraction': 'float64',
+ 'season_week': 'int64',
+ 'fluseason': 'int64',
+ 'datasetH1': str,
+ 'datasetH2': str,
+ 'sample': str,
+ 'age_group': str
+ })
+ # remove US (because this happens in the from_flusight() method)
+ data = data[data['location_code'] != 'US']
+ # filter by season_first_year (param)
+ data = data[data['fluseason'] == season_first_year]
+ # pivot wide, rename to conform with other class methods
+ gt_df_final = data.pivot_table(
+ index=['location_code', 'season_week', 'week_enddate', 'sample', 'fluseason', 'datasetH1', 'datasetH2', 'fluseason_fraction'],
+ columns='age_group',
+ values='value'
+ ).reset_index()
+ gt_df_final.columns.name = None
+ gt_df = gt_df_final.copy()
+
+ return cls(
+ season_first_year=season_first_year,
+ gt_df=gt_df,
+ gt_df_final=gt_df_final,
+ mask_date=mask_date,
+ season_setup=list(set(gt_df_final["location_code"])), # passing as a list of locs instead
+ channels=channels,
+ image_size=image_size,
+ previous_data=None,
+ dataset_coords=None,
+ nhsn=True, # param added as a PATCH
+ )
+
+
+ @classmethod
+ def from_metrocast(
+ cls,
+ season_first_year: str,
+ data_date: datetime.datetime,
+ mask_date: datetime.datetime,
+ channels=1,
+ image_size=128,
+ nogit=False,
+ dataset_coords: xr.core.coordinates.DataArrayCoordinates = None,
+ repo_path="Flusight/metrocast/flu-metrocast",
+ data_path="Flusight/metrocast/flu-metrocast/target-data/latest-data.csv",
+ main_branch="main",
+ ):
+ data_date = pd.to_datetime(data_date)
+ if not nogit:
+ cls._git_checkout_repo_rev(repo_path, target_date=None, main_branch=main_branch)
+
+ season_setup = SeasonAxis.for_metrocast()
+
+ latest_df = pd.read_csv(data_path, parse_dates=["target_end_date"])
+ latest_df = latest_df.rename(
+ columns={
+ "target_end_date": "week_enddate",
+ "location": "location_code",
+ "observation": "value",
+ }
+ )
+ latest_df["location_code"] = latest_df["location_code"].astype(str).str.strip()
+ latest_df["target"] = latest_df["target"].astype(str).str.strip()
+
+ flu_df = latest_df[latest_df["target"] == "Flu ED visits pct"].copy()
+ ili_df = latest_df[latest_df["target"] == "ILI ED visits pct"].copy()
+
+ # Combine both targets: use Flu ED visits pct when available, fill gaps with ILI ED visits pct.
+ flu_key = flu_df[["week_enddate", "location_code"]].drop_duplicates()
+ ili_fill = ili_df.merge(flu_key, on=["week_enddate", "location_code"], how="left", indicator=True)
+ ili_fill = ili_fill[ili_fill["_merge"] == "left_only"].drop(columns="_merge")
+
+ flu_df["target_source"] = "Flu ED visits pct"
+ ili_fill["target_source"] = "ILI ED visits pct"
+ latest_df = pd.concat([flu_df, ili_fill], ignore_index=True)
+
+ full_df = season_setup.add_season_columns(latest_df, do_fluseason_year=True)
+ gt_df_final = full_df[full_df["fluseason"] == int(season_first_year)]
+
+ if not nogit:
+ cls._git_checkout_repo_rev(repo_path, target_date=data_date, main_branch=main_branch)
+ latest_df = pd.read_csv(data_path, parse_dates=["target_end_date"])
+ latest_df = latest_df.rename(
+ columns={
+ "target_end_date": "week_enddate",
+ "location": "location_code",
+ "observation": "value",
+ }
+ )
+ latest_df["location_code"] = latest_df["location_code"].astype(str).str.strip()
+ latest_df["target"] = latest_df["target"].astype(str).str.strip()
+
+ flu_df = latest_df[latest_df["target"] == "Flu ED visits pct"].copy()
+ ili_df = latest_df[latest_df["target"] == "ILI ED visits pct"].copy()
+
+ flu_key = flu_df[["week_enddate", "location_code"]].drop_duplicates()
+ ili_fill = ili_df.merge(flu_key, on=["week_enddate", "location_code"], how="left", indicator=True)
+ ili_fill = ili_fill[ili_fill["_merge"] == "left_only"].drop(columns="_merge")
+
+ flu_df["target_source"] = "Flu ED visits pct"
+ ili_fill["target_source"] = "ILI ED visits pct"
+ latest_df = pd.concat([flu_df, ili_fill], ignore_index=True)
+ cls._git_checkout_repo_rev(repo_path, target_date=None, main_branch=main_branch)
+ full_df = season_setup.add_season_columns(latest_df, do_fluseason_year=True)
+ gt_df = full_df[full_df["fluseason"] == int(season_first_year)]
+
+ previous_data = full_df.copy()
+
+ return cls(
+ season_first_year=season_first_year,
+ gt_df=gt_df,
+ gt_df_final=gt_df_final,
+ mask_date=mask_date,
+ season_setup=season_setup,
+ channels=channels,
+ image_size=image_size,
+ previous_data=previous_data,
+ dataset_coords=dataset_coords,
+ )
+
def plot(self):
season_start_date = datetime.date(int(self.season_first_year), self.season_setup.season_start_month, self.season_setup.season_start_day)
- fig, axes = plt.subplots(8, 8, sharex=True, figsize=(14,16))
+ n_locations = len(self.season_setup.locations)
+ n_plots = n_locations + 1 # include US aggregate
+ n_cols = math.ceil(math.sqrt(n_plots))
+ n_rows = math.ceil(n_plots / n_cols)
+ fig, axes = plt.subplots(
+ n_rows,
+ n_cols,
+ sharex=True,
+ figsize=(max(8, n_cols * 2.2), max(8, n_rows * 2.0)),
+ )
gt_piv = self.gt_df.pivot(index = "week_enddate", columns='location_code', values='value')
gt_piv_final = self.gt_df_final.pivot(index = "week_enddate", columns='location_code', values='value')
- ax = axes.flat[0]
+ axes_flat = np.atleast_1d(axes).flat
+ ax = axes_flat[0]
ax.plot(gt_piv[self.season_setup.locations].sum(axis=1), color="black", linewidth=2,label="datadate")
ax.plot(gt_piv_final[self.season_setup.locations].sum(axis=1), lw=1, color='r', ls='-.', label="final")
ax.legend()
ax.set_ylim(0)
ax.set_title("US")
for idx, pl in enumerate(gt_piv.columns):
- ax = axes.flat[idx+1]
+ if idx + 1 >= n_plots:
+ break
+ ax = axes_flat[idx + 1]
ax.plot(gt_piv[pl], lw=2, color='k')
ax.plot(gt_piv_final[pl], lw=1, color='r', ls='-.')
na_mask = gt_piv.isna()
ax.plot(gt_piv[na_mask].index,
gt_piv[na_mask],
- marker='o',
+ marker='o',
color="pink",
- fillstyle='full',
- markeredgecolor='red',
+ fillstyle='full',
+ markeredgecolor='red',
markersize=5,
markeredgewidth=1)
ax.set_title(self.season_setup.get_location_name(pl))
@@ -217,9 +649,45 @@ def plot(self):
ax.set_xlim(season_start_date, season_start_date + datetime.timedelta(days=365))
#ax.set_xticks(season_setup.get_dates(52).resample("M"))
#ax.plot(pd.date_range(season_setup.fluseason_startdate, season_setup.fluseason_startdate + datetime.timedelta(days=64*7), freq="W-SAT"), data.flu_dyn[-50:,0,:,idx].T, c='r', lw=.5, alpha=.2)
+ for extra_ax in list(axes_flat)[n_plots:]:
+ extra_ax.set_visible(False)
fig.tight_layout()
fig.autofmt_xdate()
+ def _get_historical_series(self, location_code):
+ if self.previous_data is None or self.previous_data.empty:
+ return []
+
+ if "season_week" not in self.previous_data.columns:
+ self.previous_data = self.season_setup.add_season_columns(self.previous_data, do_fluseason_year=True)
+
+ calendar = self.season_setup.get_season_calendar(int(self.season_first_year))
+ calendar = calendar[["season_week", "saturday"]]
+
+ hist = self.previous_data[self.previous_data["location_code"] == location_code].copy()
+ if hist.empty:
+ return []
+
+ series = []
+ for hist_season in sorted(hist["fluseason"].dropna().unique()):
+ if int(hist_season) == int(self.season_first_year):
+ continue
+ season_data = hist[hist["fluseason"] == hist_season][["season_week", "value"]].dropna()
+ if season_data.empty:
+ continue
+ season_data = season_data.groupby("season_week", as_index=False)["value"].mean()
+ season_data = season_data.merge(calendar, on="season_week", how="inner").sort_values("season_week")
+ if season_data.empty:
+ continue
+ series.append(
+ (
+ hist_season,
+ pd.to_datetime(season_data["saturday"]).to_numpy(),
+ season_data["value"].to_numpy(),
+ )
+ )
+ return series
+
def plot_mask(self):
# check that it stitch
fig, axes = plt.subplots(1, 4, figsize=(8,8), dpi=200, sharex=True, sharey=True)
@@ -233,7 +701,7 @@ def plot_mask(self):
axes[1].imshow(self.gt_keep_mask[0], alpha=.3, cmap = cmap_rainbow)
axes[1].set_title("Inpainting mask", fontsize=8)
-
+
axes[2].imshow(self.gt_xarr.data[0], cmap=cmap_greys)
@@ -255,23 +723,23 @@ def export_forecasts(self, fluforecasts_ti, forecasts_national, directory=".", p
target_dates = pd.date_range(forecast_date, forecast_date + datetime.timedelta(days=4*7), freq="W-SAT")
target_dict= dict(zip(
- target_dates,
+ target_dates,
[f"{n} wk ahead inc flu hosp" for n in range(1,5)]))
print(target_dates)
#pd.DataFrame(colums=["forecast_date","target_end_date","location","type","quantile","value","target"])
df_list=[]
for qt in myutils.flusight_quantiles:
- a = pd.DataFrame(np.quantile(fluforecasts_ti[:,:,:,:len(self.season_setup.locations)], qt, axis=0)[0],
- columns= self.season_setup.locations, index=pd.date_range(season_start_date, season_start_date + datetime.timedelta(days=64*7), freq="W-SAT")).loc[target_dates]
+ a = pd.DataFrame(np.quantile(fluforecasts_ti[:,:,:,:len(self.season_setup.locations)], qt, axis=0)[0],
+ columns= self.season_setup.locations, index=pd.date_range(season_start_date, season_start_date + datetime.timedelta(days=self.image_size*7), freq="W-SAT")).loc[target_dates]
#a["US"] = a.sum(axis=1)
a["US"] = pd.DataFrame(np.quantile(forecasts_national, qt, axis=0)[0],
- columns= ["US"], index=pd.date_range(season_start_date, season_start_date + datetime.timedelta(days=64*7), freq="W-SAT")).loc[target_dates]
+ columns= ["US"], index=pd.date_range(season_start_date, season_start_date + datetime.timedelta(days=self.image_size*7), freq="W-SAT")).loc[target_dates]
a = a.reset_index().rename(columns={'index': 'target_end_date'})
a = pd.melt(a,id_vars="target_end_date",var_name="location")
a["quantile"] = '{:<.3f}'.format(qt)
-
+
df_list.append(a)
df = pd.concat(df_list)
@@ -292,7 +760,7 @@ def export_forecasts(self, fluforecasts_ti, forecasts_national, directory=".", p
# check for Error when validating format: Entries in `value` must be non-decreasing as quantiles increase:
for tg in target_dates:
old_vals = np.zeros(len(self.season_setup.locations)+1)
- for dfd in df_list: # very important to not call this df: it overwrites in namesapce the exported df
+ for dfd in df_list: # avoid naming this df; it would shadow the exported df
new_vals = dfd[dfd["target_end_date"]==tg]["value"].to_numpy()
if not (new_vals-old_vals >= 0).all():
num_negative = sum((new_vals-old_vals) < 0)
@@ -306,11 +774,16 @@ def export_forecasts(self, fluforecasts_ti, forecasts_national, directory=".", p
if save_plot:
self.plot_forecasts(fluforecasts_ti, forecasts_national, directory=directory, prefix=prefix, forecast_date=forecast_date)
-
- def plot_forecasts(self, fluforecasts_ti, forecasts_national, directory=".", prefix="", forecast_date=None):
+
+ def plot_forecasts(self, fluforecasts_ti, forecasts_national, directory=".", prefix="", forecast_date=None, mode="flusight"):
forecast_date_str=str(forecast_date)
if forecast_date == None:
forecast_date = self.mask_date
+ if forecasts_national is None:
+ if mode == "metrocast":
+ forecasts_national = fluforecasts_ti.sum(axis=-1)
+ else:
+ raise ValueError("forecasts_national is required for mode='flusight'")
idx_now = self.inpaintfrom_idx-1
idx_horizon = idx_now+4
@@ -326,6 +799,9 @@ def plot_forecasts(self, fluforecasts_ti, forecasts_national, directory=".", pre
color_gt = "black"
color_past='grey'
+ y_label = "New Hosp. Admissions" if mode != "metrocast" else "ED visits pct"
+ national_title = "National" if mode != "metrocast" else "Aggregate"
+ median_q = 0.5 if 0.5 in myutils.flusight_quantiles else myutils.flusight_quantiles[12]
nplace_toplot = len(self.season_setup.locations)
#nplace_toplot = 3 # less plots for faster iteration
@@ -337,14 +813,14 @@ def plot_forecasts(self, fluforecasts_ti, forecasts_national, directory=".", pre
#if self.season_first_year == "2023" or self.season_first_year == "2024":
- # gt2022 = GroundTruth(season_first_year="2022",
+ # gt2022 = GroundTruth(season_first_year="2022",
# data_date=datetime.datetime.combine(datetime.date(2023,7,15), datetime.datetime.min.time()),
# mask_date=datetime.datetime.today(),
# channels=self.channels,
# image_size=self.image_size,
# payload=pd.read_csv("custom_datasets/nc_payload_gt.csv", parse_dates=["week_enddate"]))
#if self.season_first_year == "2024":
- # gt2023 = GroundTruth(season_first_year="2023",
+ # gt2023 = GroundTruth(season_first_year="2023",
# data_date=datetime.datetime.combine(datetime.date(2023,7,15), datetime.datetime.min.time()),
# mask_date=datetime.datetime.today(),
# channels=self.channels,
@@ -356,15 +832,15 @@ def plot_forecasts(self, fluforecasts_ti, forecasts_national, directory=".", pre
fig, axes = plt.subplots(nplace_toplot+1, 2, figsize=(10,nplace_toplot*3.5), dpi=200)
for iax in range(2):
ax = axes[0][iax]
-
- x = np.arange(64)
+
+ x = np.arange(self.image_size)
if iax == 0:
x_lims_idx = (0, 51)
- x_lims = (pd.to_datetime(self.gt_xarr["date"][x_lims_idx[0]].values),
+ x_lims = (pd.to_datetime(self.gt_xarr["date"][x_lims_idx[0]].values),
pd.to_datetime(self.gt_xarr["date"][x_lims_idx[1]].values))
elif iax == 1:
x_lims_idx = (idx_now-3, idx_horizon)
- x_lims = (pd.to_datetime(self.gt_xarr["date"][x_lims_idx[0]].values),
+ x_lims = (pd.to_datetime(self.gt_xarr["date"][x_lims_idx[0]].values),
pd.to_datetime(self.gt_xarr["date"][x_lims_idx[1]].values))
# US WIDE: quantiles and median, US-wide
for iqt in plot_spec["quantiles_idx"]:
@@ -372,12 +848,12 @@ def plot_forecasts(self, fluforecasts_ti, forecasts_national, directory=".", pre
# TODO: not exactly true that it is the sum of quantiles (sum of quantile is not quantile of sum)
ylo = np.quantile(forecasts_national, myutils.flusight_quantile_pairs[iqt,0], axis=0)[0]
yup = np.quantile(forecasts_national, myutils.flusight_quantile_pairs[iqt,1], axis=0)[0]
- ax.fill_between(self.gt_xarr["date"][plotrange],
- ylo[plotrange],
- yup[plotrange],
- alpha=.1,
+ ax.fill_between(self.gt_xarr["date"][plotrange],
+ ylo[plotrange],
+ yup[plotrange],
+ alpha=.1,
color=plot_spec["color"])
-
+
# widest quantile pair is the first one. We take the up quantile of it + a few % as x_lim
if iqt == plot_spec["quantiles_idx"][0]:
if plot_past_median:
@@ -386,19 +862,24 @@ def plot_forecasts(self, fluforecasts_ti, forecasts_national, directory=".", pre
max_y_value = max(yup[self.inpaintfrom_idx:x_lims_idx[1]])
max_y_value = max(max_y_value, self.gt_xarr.data[0,:self.inpaintfrom_idx].sum(axis=1)[x_lims_idx[0]:x_lims_idx[1]].max())
max_y_value = max_y_value + max_y_value*.05 # 10% more
-
+
# median
- ax.plot(self.gt_xarr["date"][plotrange],
- np.quantile(forecasts_national, myutils.flusight_quantiles[12], axis=0)[0][plotrange], color=plot_spec["color"], marker='.', label='forecast median')
-
+ ax.plot(
+ self.gt_xarr["date"][plotrange],
+ np.quantile(forecasts_national, median_q, axis=0)[0][plotrange],
+ color=plot_spec["color"],
+ marker=".",
+ label="forecast median",
+ )
+
# ground truth
ax.plot(self.gt_xarr["date"][:self.inpaintfrom_idx],
self.gt_xarr.data[0,:self.inpaintfrom_idx].sum(axis=1), color=color_gt, marker = '.', lw=.5, label='ground-truth')
ax.plot(self.gt_xarr["date"][self.inpaintfrom_idx:],
- self.gt_xarr.data[0,self.inpaintfrom_idx:].sum(axis=1),
- color='red',
- marker = '.',
- lw=.1,
+ self.gt_xarr.data[0,self.inpaintfrom_idx:].sum(axis=1),
+ color='red',
+ marker = '.',
+ lw=.1,
label='ground-truth',
markersize=.4)
@@ -409,7 +890,7 @@ def plot_forecasts(self, fluforecasts_ti, forecasts_national, directory=".", pre
if iax==0:
ax.legend(fontsize=8)
-
+
#ax.set_xticks(np.arange(0,53,13))
@@ -418,7 +899,7 @@ def plot_forecasts(self, fluforecasts_ti, forecasts_national, directory=".", pre
ax.axvline(self.gt_xarr["date"][idx_now].values, c='k', lw=1, ls='-.')
if iax == 0:
ax.axvline(self.gt_xarr["date"][idx_horizon].values, c='k', lw=1, ls='-.')
- ax.set_title("National")
+ ax.set_title(national_title)
sns.despine(ax = ax, trim = True, offset=4)
@@ -448,31 +929,38 @@ def plot_forecasts(self, fluforecasts_ti, forecasts_national, directory=".", pre
location_name=self.season_setup.get_location_name(self.season_setup.locations[ipl])
ax = axes[ipl+1][iax]
# median
- ax.plot(self.gt_xarr["date"][plotrange],
- np.quantile(fluforecasts_ti, myutils.flusight_quantiles[12], axis=0)[0,:,ipl][plotrange], color=plot_spec["color"], marker = '.', lw=.5)
+ ax.plot(
+ self.gt_xarr["date"][plotrange],
+ np.quantile(fluforecasts_ti, median_q, axis=0)[0, :, ipl][plotrange],
+ color=plot_spec["color"],
+ marker=".",
+ lw=.5,
+ )
# ground truth
ax.plot(self.gt_xarr["date"][:self.inpaintfrom_idx],
self.gt_xarr.data[0,:self.inpaintfrom_idx, ipl], color=color_gt, marker = '.', lw=.5)
ax.plot(self.gt_xarr["date"][self.inpaintfrom_idx:],
self.gt_xarr.data[0,self.inpaintfrom_idx:, ipl], color='red', marker = '.', lw=.1, markersize=.4)
- # TODO I'm here
-
- this_hist_data = self.previous_data[self.previous_data["location_code"]==self.season_setup.locations[ipl]]
- for hist_season in this_hist_data["fluseason"].unique():
- if int(hist_season) != int(self.season_first_year):
- hist_data = this_hist_data[this_hist_data["fluseason"]==hist_season]
- hist_data = hist_data.pivot(index = "season_week", columns='location_code', values='value').sort_index()
- thisthing = hist_data[self.season_setup.locations[ipl]]
- # TODO MATCH HERE !!!!
- ax.plot(self.gt_xarr["date"][0:len(thisthing)], thisthing, color=color_past, ls='dashed', lw=.5, label=f"{hist_season}")
+ for hist_season, hist_dates, hist_values in self._get_historical_series(
+ self.season_setup.locations[ipl]
+ ):
+ ax.plot(
+ hist_dates,
+ hist_values,
+ color=color_past,
+ ls="dashed",
+ lw=.5,
+ label=f"{hist_season}",
+ )
ax.axvline(self.gt_xarr["date"][idx_now].values, c='k', lw=1, ls='-.')
if iax == 0:
ax.axvline(self.gt_xarr["date"][idx_horizon].values, c='k', lw=1, ls='-.')
ax.set_xlim(x_lims)
ax.set_ylim(bottom=0, top=max_y_value[ipl])
- if iax==0: ax.set_ylabel("New Hosp. Admissions")
+ if iax == 0:
+ ax.set_ylabel(y_label)
ax.set_title(location_name)
# rotate the x axis labels
ax.tick_params(axis='x', rotation=45)
@@ -484,33 +972,96 @@ def plot_forecasts(self, fluforecasts_ti, forecasts_national, directory=".", pre
fig.tight_layout()
plt.savefig(f"{directory}/{prefix}-{forecast_date_str}-plot{plot_title}.pdf")
- def export_forecasts_2023(self, fluforecasts_ti, forecasts_national, directory=".", prefix="", forecast_date=None, save_plot=True, nochecks=False, rate_trend=True):
+ def export_forecasts_2023(self, fluforecasts_ti, forecasts_national=None, directory=".", prefix="", forecast_date=None, save_plot=True, nochecks=False, rate_trend=True, mode="flusight"):
forecast_date_str=str(forecast_date)
if forecast_date == None:
forecast_date = self.mask_date
season_start_date = datetime.date(int(self.season_first_year), self.season_setup.season_start_month, self.season_setup.season_start_day)
- target_dates = pd.date_range(forecast_date, forecast_date + datetime.timedelta(days=3*7), freq="W-SAT")
-
- target_dict= dict(zip(
- target_dates,
- [f"{n}" for n in range(0,4)]))
-
- df_list=[]
+ reference_date = pd.to_datetime(forecast_date).date()
+ reference_date_str = str(reference_date)
+ base_index = pd.date_range(
+ season_start_date,
+ season_start_date + datetime.timedelta(days=self.image_size * 7),
+ freq="W-SAT",
+ )
+ target_dates = [reference_date + datetime.timedelta(days=7 * h) for h in range(4)]
+ target_dates = pd.to_datetime(target_dates)
+ horizon_map = {pd.to_datetime(d): h for h, d in enumerate(target_dates)}
+
+ df_list = []
for qt in myutils.flusight_quantiles:
- a = pd.DataFrame(np.quantile(fluforecasts_ti[:,:,:,:len(self.season_setup.locations)], qt, axis=0)[0],
- columns= self.season_setup.locations, index=pd.date_range(season_start_date, season_start_date + datetime.timedelta(days=64*7), freq="W-SAT")).loc[target_dates]
- #a["US"] = a.sum(axis=1)
- a["US"] = pd.DataFrame(np.quantile(forecasts_national, qt, axis=0)[0],
- columns= ["US"], index=pd.date_range(season_start_date, season_start_date + datetime.timedelta(days=64*7), freq="W-SAT")).loc[target_dates]
-
- a = a.reset_index().rename(columns={'index': 'target_end_date'})
- a = pd.melt(a,id_vars="target_end_date",var_name="location")
- a["output_type_id"] = "{:.3f}".format(qt).rstrip('0').rstrip('.')# " #'{:<.3f}'.format(qt)
-
+ a = pd.DataFrame(
+ np.quantile(
+ fluforecasts_ti[:, :, :, : len(self.season_setup.locations)], qt, axis=0
+ )[0],
+ columns=self.season_setup.locations,
+ index=base_index,
+ ).loc[target_dates]
+
+ a = a.reset_index().rename(columns={"index": "target_end_date"})
+ a = pd.melt(a, id_vars="target_end_date", var_name="location")
+ a["output_type_id"] = "{:.3f}".format(qt).rstrip("0").rstrip(".")
df_list.append(a)
- df = pd.concat(df_list)
+ df = pd.concat(df_list, ignore_index=True)
+
+ if mode == "metrocast":
+ df["reference_date"] = reference_date_str
+ df["output_type"] = "quantile"
+ df["horizon"] = df["target_end_date"].map(horizon_map)
+ df["target"] = np.where(
+ df["location"] == "nyc", "ILI ED visits pct", "Flu ED visits pct"
+ )
+ df = df[
+ [
+ "reference_date",
+ "target",
+ "horizon",
+ "target_end_date",
+ "location",
+ "output_type",
+ "output_type_id",
+ "value",
+ ]
+ ]
+
+ if not nochecks:
+ assert sum(df["value"] < 0) == 0
+ assert sum(df["value"].isna()) == 0
+
+ df.to_csv(f"{directory}/{reference_date_str}-{prefix}.csv", index=False)
+ if save_plot:
+ if forecasts_national is None:
+ forecasts_national = fluforecasts_ti.sum(axis=-1)
+ self.plot_forecasts(
+ fluforecasts_ti,
+ forecasts_national,
+ directory=directory,
+ prefix=prefix,
+ forecast_date=forecast_date,
+ mode=mode,
+ )
+ return
+
+ if forecasts_national is None:
+ raise ValueError("forecasts_national is required for mode='flusight'")
+
+ target_dict = {d: f"{h}" for d, h in horizon_map.items()}
+ updated_df_list = []
+ for qt, dfd in zip(myutils.flusight_quantiles, df_list):
+ us_vals = pd.DataFrame(
+ np.quantile(forecasts_national, qt, axis=0)[0],
+ columns=["US"],
+ index=base_index,
+ ).loc[target_dates]
+ us_vals = us_vals.reset_index().rename(columns={"index": "target_end_date"})
+ us_vals = pd.melt(us_vals, id_vars="target_end_date", var_name="location")
+ us_vals["output_type_id"] = "{:.3f}".format(qt).rstrip("0").rstrip(".")
+ dfd = pd.concat([dfd, us_vals], ignore_index=True)
+ updated_df_list.append(dfd)
+
+ df = pd.concat(updated_df_list, ignore_index=True)
df["reference_date"] = forecast_date_str
df["target"] = "wk inc flu hosp"
df["horizon"] = df["target_end_date"].map(target_dict)
@@ -518,7 +1069,7 @@ def export_forecasts_2023(self, fluforecasts_ti, forecasts_national, directory="
df = df[["reference_date","target","horizon","target_end_date","location","output_type","output_type_id","value"]]
df
- # Suppress verbose output for column information
+ # Suppress verbose output for column information
# for col in df.columns:
# print(col)
# print(df[col].unique())
@@ -530,7 +1081,7 @@ def export_forecasts_2023(self, fluforecasts_ti, forecasts_national, directory="
# check for Error when validating format: Entries in `value` must be non-decreasing as quantiles increase:
for tg in target_dates:
old_vals = np.zeros(len(self.season_setup.locations)+1)
- for dfd in df_list: # very important to not call this df: it overwrites in namesapce the exported df
+ for dfd in updated_df_list: # very important to not call this df: it overwrites in namesapce the exported df
new_vals = dfd[dfd["target_end_date"]==tg]["value"].to_numpy()
if not (new_vals-old_vals >= 0).all():
num_negative = sum((new_vals-old_vals) < 0)
@@ -544,7 +1095,7 @@ def export_forecasts_2023(self, fluforecasts_ti, forecasts_national, directory="
# df_list=[]
# for sim_id in np.arange(fluforecasts_ti.shape[0]):
# #for qt in myutils.flusight_quantiles:
-# a = pd.DataFrame(fluforecasts_ti[:,:,:,:len(self.season_setup.locations)],
+# a = pd.DataFrame(fluforecasts_ti[:,:,:,:len(self.season_setup.locations)],
# columns= self.season_setup.locations, index=pd.date_range(self.season_setup.fluseason_startdate, self.season_setup.fluseason_startdate + datetime.timedelta(days=64*7), freq="W-SAT")).loc[target_dates]
# a["US"] = pd.DataFrame(forecasts_national[sim_id],
# columns= ["US"], index=pd.date_range(self.season_setup.fluseason_startdate, self.season_setup.fluseason_startdate + datetime.timedelta(days=64*7), freq="W-SAT")).loc[target_dates]
@@ -552,7 +1103,7 @@ def export_forecasts_2023(self, fluforecasts_ti, forecasts_national, directory="
# a = a.reset_index().rename(columns={'index': 'target_end_date'})
# a = pd.melt(a,id_vars="target_end_date",var_name="location")
#
-#
+#
# df_list.append(a)
#
# df2 = pd.concat(df_list)
@@ -566,5 +1117,4 @@ def export_forecasts_2023(self, fluforecasts_ti, forecasts_national, directory="
df.to_csv(f"{directory}/{forecast_date_str}-{prefix}.csv", index=False)
if save_plot:
- self.plot_forecasts(fluforecasts_ti, forecasts_national, directory=directory, prefix=prefix, forecast_date=forecast_date)
-
+ self.plot_forecasts(fluforecasts_ti, forecasts_national, directory=directory, prefix=prefix, forecast_date=forecast_date)
\ No newline at end of file
diff --git a/prep_age_flu_hosp_data.ipynb b/prep_age_flu_hosp_data.ipynb
new file mode 100644
index 0000000..21f2f63
--- /dev/null
+++ b/prep_age_flu_hosp_data.ipynb
@@ -0,0 +1,835 @@
+{
+ "cells": [
+ {
+ "cell_type": "code",
+ "execution_count": 1,
+ "id": "07601a26",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import numpy as np\n",
+ "import pandas as pd\n",
+ "import requests\n",
+ "import time\n",
+ "from requests.adapters import HTTPAdapter\n",
+ "from urllib3.util.retry import Retry\n",
+ "import datetime\n",
+ "from pathlib import Path \n",
+ "import re"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "b92f303c",
+ "metadata": {},
+ "source": [
+ "Need for this file:\n",
+ "\n",
+ "- The imports above\n",
+ "- locations data: Path to a locations.csv file (to merge location name to location code). Path will be read in as a pd.DataFrame\n",
+ "- current FluSMH GitHub repository local clone. The file will use the abs path to the `model-output` directory:\n",
+ " e.g. \"your/path/to/flu-scenario-modeling-hub/model-output\"\n",
+ "- archive FluSMH GitHub repository locial clone. The file will use the abs path to the `data-processed` directory\n",
+ " e.g. \"your/path/to/flu-scenario-modeling-hub_archive/data-processed\"\n",
+ "- output path (abs) for processed files\n",
+ " replace in saving for loop"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 2,
+ "id": "1322a17e",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# modified from respi\n",
+ "\n",
+ "LOCATIONS_ABBREV = [\n",
+ " 'AL', 'AK', 'AZ', 'AR', 'CA', 'CO', 'CT', 'DE', 'DC', 'FL', 'GA', \n",
+ " 'HI', 'ID', 'IL', 'IN', 'IA', 'KS', 'KY', 'LA', 'ME', 'MD', 'MA', \n",
+ " 'MI', 'MN', 'MS', 'MO', 'MT', 'NE', 'NV', 'NH', 'NJ', 'NM', 'NY', \n",
+ " 'NC', 'ND', 'OH', 'OK', 'OR', 'PA', 'RI', 'SC', 'SD', 'TN', 'TX', \n",
+ " 'UT', 'VT', 'VA', 'WA', 'WV', 'WI', 'WY', 'US'\n",
+ " ]\n",
+ "\n",
+ "class NHSNDataProcessor:\n",
+ " def __init__(self, resource_id, replace_column_names: bool = True):\n",
+ " self.replace_column_names = replace_column_names\n",
+ " self.data_url = \"https://data.cdc.gov/resource/\" + f\"{resource_id}.json\"\n",
+ " self.metadata_url = \"https://data.cdc.gov/api/views/\" + f\"{resource_id}.json\"\n",
+ " self.output_dict = {}\n",
+ "\n",
+ " self._process_data()\n",
+ "\n",
+ " \n",
+ " def _process_data(self):\n",
+ " \"\"\"Fetches, processes, and structures NHSN data into self.output_dict\"\"\"\n",
+ " # Get data set up \n",
+ " data = pd.DataFrame(self._retrieve_data_from_endpoint_aslist()) # read from endpoint\n",
+ " data = data.drop(columns=['respseason'])\n",
+ " non_numeric_cols = ['jurisdiction', 'weekendingdate'] # make numeric cols not strings\n",
+ " for col in data.columns:\n",
+ " if col not in non_numeric_cols:\n",
+ " data[col] = pd.to_numeric(data[col], errors='raise')\n",
+ " data = data.replace(np.nan, value=None) # cleanse NaN values \n",
+ " data.loc[data['jurisdiction'].str.lower() == 'usa', 'jurisdiction'] = 'US' # change USA jurisdiction to US\n",
+ " data = data[data['jurisdiction'].isin(LOCATIONS_ABBREV)].copy() # filter out unwanted regions\n",
+ " data['weekendingdate'] = pd.to_datetime(data['weekendingdate']).dt.strftime('%Y-%m-%d') # ensure date columns are dates\n",
+ " # Get metadata set up\n",
+ " cdc_metadata = (requests.get(self.metadata_url)).json() \n",
+ "\n",
+ " if self.replace_column_names: \n",
+ " data = self._replace_column_names(data, cdc_metadata) \n",
+ " data = data.sort_values(by=['Geographic aggregation', 'Week Ending Date'])\n",
+ " self.output_dict[\"processed_df\"] = data\n",
+ "\n",
+ " def _retrieve_data_from_endpoint_aslist(self) -> list[dict]:\n",
+ " \"\"\"Downloads NHSN data from the endpoint with pagination and retries.\"\"\"\n",
+ " \n",
+ " session = requests.Session()\n",
+ " retries = Retry(total=5,\n",
+ " backoff_factor=1,\n",
+ " status_forcelist=[500, 502, 503, 504])\n",
+ " session.mount('https://', HTTPAdapter(max_retries=retries))\n",
+ " \n",
+ " all_data = []\n",
+ " offset = 0\n",
+ " batch_size = 1000\n",
+ " while True:\n",
+ " params = {\"$limit\": batch_size, \"$offset\": offset}\n",
+ " try:\n",
+ " # Use the configured session to make the request\n",
+ " data_response = session.get(self.data_url, params=params, timeout=30)\n",
+ " data_response.raise_for_status()\n",
+ " batch_data = data_response.json()\n",
+ " if not batch_data:\n",
+ " break\n",
+ " all_data.extend(batch_data)\n",
+ " offset += batch_size\n",
+ " time.sleep(0.1)\n",
+ " except Exception as e:\n",
+ " raise\n",
+ " return all_data\n",
+ " \n",
+ " def _replace_column_names(self, data: pd.DataFrame, cdc_metadata: dict) -> pd.DataFrame:\n",
+ " \"\"\"Replace short-form column names with long-form column names\"\"\"\n",
+ " column_name_map = {\n",
+ " col_info['fieldName']: col_info['name']\n",
+ " for col_info in cdc_metadata['columns']\n",
+ " }\n",
+ " return data.rename(columns=column_name_map, errors=\"ignore\")\n",
+ "\n",
+ "\n",
+ "NHSNData = NHSNDataProcessor(resource_id='ua7e-t2fy')\n",
+ "\n",
+ " "
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 3,
+ "id": "41c04aac",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "/var/folders/pf/s416pvp93gd610_55fzf_f280000gp/T/ipykernel_53075/2467725948.py:19: FutureWarning: Downcasting object dtype arrays on .fillna, .ffill, .bfill is deprecated and will change in a future version. Call result.infer_objects(copy=False) instead. To opt-in to the future behavior, set `pd.set_option('future.no_silent_downcasting', True)`\n",
+ " data = data.fillna(0) # fill all Na values w/ 0\n"
+ ]
+ }
+ ],
+ "source": [
+ "data = NHSNData.output_dict['processed_df'] # only has 2024-25 seasons for the metrics we need\n",
+ "# has all weeks for all location in 2024 and 2025 so far though! \n",
+ "\n",
+ "# initial column filtering\n",
+ "data = data[[ \n",
+ " \"Week Ending Date\", \n",
+ " \"Geographic aggregation\",\n",
+ " \"Number of Pediatric Influenza Admissions, 0-4 years\",\n",
+ " \"Number of Pediatric Influenza Admissions, 5-17 years\",\n",
+ " \"Total Pediatric Influenza Admissions\",\n",
+ " \"Number of Adult Influenza Admissions, 18-49 years\",\n",
+ " \"Number of Adult Influenza Admissions, 50-64 years\",\n",
+ " \"Number of Adult Influenza Admissions, 65-74 years\",\n",
+ " \"Number of Adult Influenza Admissions, 75+ years\",\n",
+ " \"Total Adult Influenza Admissions\",\n",
+ " ]].copy()\n",
+ "\n",
+ "# renaming, removing, and small calculations\n",
+ "data = data.fillna(0) # fill all Na values w/ 0\n",
+ "data['0-130'] = data[\"Total Pediatric Influenza Admissions\"] + data[\"Total Adult Influenza Admissions\"]\n",
+ "data['65+'] = data['Number of Adult Influenza Admissions, 65-74 years'] + data['Number of Adult Influenza Admissions, 75+ years']\n",
+ "data = data.rename(columns={\n",
+ " 'Week Ending Date': 'week_enddate', \n",
+ " 'Number of Pediatric Influenza Admissions, 0-4 years': '0-4',\n",
+ " 'Number of Pediatric Influenza Admissions, 5-17 years': '5-17',\n",
+ " 'Number of Adult Influenza Admissions, 18-49 years': '18-49',\n",
+ " 'Number of Adult Influenza Admissions, 50-64 years': '50-64',\n",
+ " })\n",
+ "data = data.drop(columns=['Number of Adult Influenza Admissions, 65-74 years', 'Number of Adult Influenza Admissions, 75+ years', 'Total Pediatric Influenza Admissions', 'Total Adult Influenza Admissions'])\n",
+ "\n",
+ "# remove US loc\n",
+ "# data = data[data['Geographic aggregation'] != 'US'], opting to keep US in now\n",
+ "# cast week_enddate column as date\n",
+ "data['week_enddate'] = pd.to_datetime(data['week_enddate'])\n",
+ "# add fluseason column\n",
+ "data['fluseason'] = data['week_enddate'].dt.year\n",
+ "data.loc[data['week_enddate'].dt.month < 8, 'fluseason'] -= 1\n",
+ "# keep only 2024 fluseason (that's where the age group we need start, and we won't use current season)\n",
+ "data = data[data['fluseason'].isin([2024])]\n",
+ "# add location_code, sample, datasetH1, datasetH2 columns\n",
+ "locations = pd.read_csv('/Users/emprzy/Documents/work/miscellaneous/locations.csv')\n",
+ "data = data.merge(\n",
+ " locations[['abbreviation', 'location']],\n",
+ " left_on='Geographic aggregation',\n",
+ " right_on='abbreviation',\n",
+ " how='left'\n",
+ ")\n",
+ "data = data.rename(columns={'location': 'location_code'}).drop(columns=['abbreviation', 'Geographic aggregation'])\n",
+ "data['sample'] = \"1\"\n",
+ "data['datasetH1'] = 'NHSN'\n",
+ "data['datasetH2'] = 'NHSN'\n",
+ "# add season_week column\n",
+ "data = data.sort_values(['location_code', 'fluseason', 'week_enddate'])\n",
+ "data['season_week'] = data.groupby(['location_code', 'fluseason']).cumcount() + 1\n",
+ "# add fluseason_fraction\n",
+ "def get_season_fraction(ts, start_month: int, start_day: int): # modified from influpaint season_axis.py\n",
+ " if pd.isna(ts):\n",
+ " return float('nan')\n",
+ " if isinstance(ts, datetime.datetime):\n",
+ " ts = ts.date()\n",
+ " try:\n",
+ " season_start = datetime.date(ts.year, start_month, start_day)\n",
+ " except AttributeError:\n",
+ " season_start = datetime.date(ts.year, start_month, start_day)\n",
+ " if ts < season_start:\n",
+ " season_start = datetime.date(ts.year - 1, start_month, start_day)\n",
+ " \n",
+ " days_since_start = (ts - season_start).days\n",
+ " return days_since_start / 365\n",
+ "\n",
+ "data['fluseason_fraction'] = data['week_enddate'].apply(\n",
+ " get_season_fraction, \n",
+ " start_month=10, \n",
+ " start_day=1\n",
+ ")\n",
+ "\n",
+ "NHSN = data\n",
+ "\n",
+ "# melt NHSN age grouping columns into singular `age_group`\n",
+ "NHSN = NHSN.rename(columns={'65+': '65-130'})\n",
+ "age_columns = ['0-4', '5-17', '18-49', '50-64', '0-130', '65-130']\n",
+ "id_vars = [col for col in NHSN.columns if col not in age_columns]\n",
+ "NHSN = NHSN.melt(\n",
+ " id_vars=id_vars,\n",
+ " value_vars=age_columns,\n",
+ " var_name='age_group',\n",
+ " value_name='value'\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 4,
+ "id": "2cd01265",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "LOCATIONS = set(data['location_code'])"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 5,
+ "id": "2603072c",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "\n",
+ "\n",
+ "def fetch_SMH_submissions(base_path: str) -> list:\n",
+ " base_path = Path(base_path)\n",
+ " keep_files = []\n",
+ " YEAR_PATTERN = re.compile(r'^(2022-08|2023|2024|2025)')\n",
+ " # YEAR_PATTERN = re.compile(r'^(2023|2024|2025)')\n",
+ "\n",
+ " for folder in base_path.iterdir(): # if it's a folder\n",
+ " if folder.is_dir():\n",
+ " \n",
+ " # iterate over files in that folder\n",
+ " for file in folder.iterdir():\n",
+ " if file.is_file() and not file.name.startswith('.'):\n",
+ " full_file_path = file.resolve()\n",
+ "\n",
+ " if full_file_path.suffix.lower() in ['.csv', '.parquet']:\n",
+ " if YEAR_PATTERN.match(file.name):\n",
+ " keep_files.append(full_file_path)\n",
+ "\n",
+ " else: continue # skip files that aren't csv or parquet\n",
+ "\n",
+ " else: continue # skip folders once inside of each models' dir\n",
+ "\n",
+ " else: # if it's a file (skipping .md files)\n",
+ " continue\n",
+ "\n",
+ " return keep_files\n",
+ "\n",
+ "\n",
+ "def parse_SMH_submissions(file_paths: dict[str: Path]) -> dict[str: Path]:\n",
+ " # return object\n",
+ " good_files = {}\n",
+ " AGE_GROUPS = {'0-130', '0-4', '18-49', '5-17', '50-64', '65-130'}\n",
+ "\n",
+ " for _, list_of_paths in file_paths.items():\n",
+ " for path in list_of_paths:\n",
+ " if path.suffix.lower() == '.csv':\n",
+ " data = pd.read_csv(path)\n",
+ " elif path.suffix.lower() == '.parquet':\n",
+ " data = pd.read_parquet(path)\n",
+ " else:\n",
+ " raise ValueError(f'Unexpected file type received: {path.suffix}')\n",
+ "\n",
+ " if not 'output_type' in data.columns: # 2022 round 1 doesn't have output_type column\n",
+ " if 'type' in data.columns:\n",
+ " data = data[data['type'] == 'point'] # only keep sample output type\n",
+ " else:\n",
+ " data = data[data['output_type'] == 'sample'] # only keep sample output_type\n",
+ "\n",
+ " # check if all required age groups are present\n",
+ " if not (AGE_GROUPS.issubset(set(data['age_group']))):\n",
+ " continue\n",
+ " \n",
+ " # check by age group\n",
+ " is_file_valid = True\n",
+ " age_group_gbo = data.groupby('age_group')\n",
+ " for name, age_group_df in age_group_gbo:\n",
+ " if name not in AGE_GROUPS:\n",
+ " continue\n",
+ " has_locations = len(age_group_df['location'].unique()) >= 20\n",
+ " has_target = age_group_df['target'].str.contains('inc hosp', case=False, na=False).any()\n",
+ " if not (has_locations and has_target):\n",
+ " is_file_valid = False\n",
+ " break \n",
+ "\n",
+ " if is_file_valid:\n",
+ " good_files[path.name] = {\n",
+ " 'file_path': path,\n",
+ " 'scenarios': set(data['scenario_id']),\n",
+ " 'num_locations': len(data['location'].unique())\n",
+ " }\n",
+ "\n",
+ " return good_files"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 6,
+ "id": "9e3d7663",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "/var/folders/pf/s416pvp93gd610_55fzf_f280000gp/T/ipykernel_53075/3944370396.py:37: DtypeWarning: Columns (4) have mixed types. Specify dtype option on import or set low_memory=False.\n",
+ " data = pd.read_csv(path)\n",
+ "/var/folders/pf/s416pvp93gd610_55fzf_f280000gp/T/ipykernel_53075/3944370396.py:37: DtypeWarning: Columns (4) have mixed types. Specify dtype option on import or set low_memory=False.\n",
+ " data = pd.read_csv(path)\n"
+ ]
+ }
+ ],
+ "source": [
+ "base_paths = {\"past\": \"/Users/emprzy/Documents/work/flu-scenario-modeling-hub_archive/data-processed\", \"current\": \"/Users/emprzy/Documents/work/flu-scenario-modeling-hub/model-output\"}\n",
+ "files_to_parse = {}\n",
+ "for origin, path in base_paths.items():\n",
+ " files_to_parse[origin] = fetch_SMH_submissions(base_path=path)\n",
+ "\n",
+ "good_files = parse_SMH_submissions(file_paths=files_to_parse) # note that there are none from 2022! 15 total good files"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 7,
+ "id": "071a45b6",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def impute_season_week(data: pd.DataFrame, fluseason: int, trajectories: int) -> pd.DataFrame: # critical: assumes only one origin_date per df\n",
+ " if len(set(data['origin_date'])) != 1:\n",
+ " raise ValueError(f\"Len of unique origin_dates is {len(set(data['origin_date']))}\") \n",
+ "\n",
+ " # prelim set up\n",
+ " zero_horizon = pd.to_datetime(data['origin_date'].iloc[0])\n",
+ " one_horizon = pd.to_datetime(data['origin_date'].iloc[0])+pd.Timedelta(days=7)\n",
+ " horizon_to_week_enddate = {0: {\"week_enddate\": zero_horizon}, 1: {\"week_enddate\": one_horizon}}\n",
+ "\n",
+ " # add all necessary weeks to dict\n",
+ " desired_start = pd.to_datetime(datetime.datetime(fluseason, 8, 1))\n",
+ " current_date = zero_horizon\n",
+ " week_offset = 0\n",
+ " while current_date > desired_start:\n",
+ " current_date -= pd.Timedelta(days=7) \n",
+ " week_offset += 1\n",
+ " for i in range(1,week_offset):\n",
+ " j = i*-1\n",
+ " horizon_to_week_enddate[j] = {}\n",
+ " horizon_to_week_enddate[j][\"week_enddate\"] = (horizon_to_week_enddate[0][\"week_enddate\"] - pd.Timedelta(days=7*i))\n",
+ " for i in range(2,54):\n",
+ " horizon_to_week_enddate[i] = {}\n",
+ " horizon_to_week_enddate[i][\"week_enddate\"] = (horizon_to_week_enddate[0][\"week_enddate\"] + pd.Timedelta(days=7*i))\n",
+ "\n",
+ " # trim off horizons we don't need\n",
+ " next_season = fluseason + 1\n",
+ " keys_to_delete = []\n",
+ " for key, value in horizon_to_week_enddate.items():\n",
+ " if value[\"week_enddate\"] > pd.to_datetime(datetime.datetime(next_season, 8, 1)):\n",
+ " keys_to_delete.append(key)\n",
+ " for key in keys_to_delete:\n",
+ " del horizon_to_week_enddate[key]\n",
+ "\n",
+ " # add season_week to dict\n",
+ " sorted_keys = sorted(horizon_to_week_enddate.keys())\n",
+ " for i, key in enumerate(sorted_keys, start=1):\n",
+ " horizon_to_week_enddate[key][\"season_week\"] = i \n",
+ "\n",
+ " # add missing horizons\n",
+ " missing_horizons = set(horizon_to_week_enddate.keys()) - set(data['horizon'])\n",
+ " if missing_horizons:\n",
+ " unique_groups = data[['scenario_id', 'location', 'age_group']].drop_duplicates()\n",
+ " new_rows_list = []\n",
+ " for h in missing_horizons:\n",
+ " temp_df = unique_groups.copy()\n",
+ " temp_df['horizon'] = h\n",
+ " temp_df['value'] = 0\n",
+ " temp_df = temp_df.loc[temp_df.index.repeat(trajectories)].reset_index(drop=True) \n",
+ " new_rows_list.append(temp_df)\n",
+ " missing_data_df = pd.concat(new_rows_list, ignore_index=True)\n",
+ " data = pd.concat([data, missing_data_df], ignore_index=True)\n",
+ " \n",
+ " \n",
+ " # match horizon to season_week, week_enddate\n",
+ " sw_map = {k: v['season_week'] for k, v in horizon_to_week_enddate.items()}\n",
+ " we_map = {k: v['week_enddate'] for k, v in horizon_to_week_enddate.items()}\n",
+ " data['season_week'] = data['horizon'].map(sw_map)\n",
+ " data['week_enddate'] = data['horizon'].map(we_map)\n",
+ "\n",
+ " # add fluseason_fraction column (function used is defined above; with NHSN data)\n",
+ " data['fluseason_fraction'] = data['week_enddate'].apply(\n",
+ " get_season_fraction, \n",
+ " start_month=10, \n",
+ " start_day=1\n",
+ " )\n",
+ "\n",
+ " # add sample column (CRITICAL: assumes all trajectory numbers are equal across all combinations)\n",
+ " sort_cols = ['scenario_id', 'location', 'age_group', 'season_week']\n",
+ " data = data.sort_values(by=sort_cols).reset_index(drop=True)\n",
+ " data['sample'] = np.tile(np.arange(1, trajectories + 1), len(data) // trajectories)\n",
+ "\n",
+ " # drop useless columns\n",
+ " cols_to_drop = [c for c in ['origin_date', 'target', 'output_type', 'output_type_id', 'horizon'] if c in data.columns]\n",
+ " data = data.drop(columns=cols_to_drop)\n",
+ " data = data.rename(columns={'location': 'location_code'})\n",
+ "\n",
+ " return data\n",
+ "\n",
+ "\n",
+ "def filter_smh(data: pd.DataFrame, fluseason: int, datasetH1: str, datasetH2: str, trajectories: int) -> pd.DataFrame:\n",
+ " data = data[data['output_type'] == 'sample'] # only keep sample output type\n",
+ " data = data[data['target'] == 'inc hosp'] # only keep inc hosp target\n",
+ " data['location'] = data['location'].astype(str) # ensure location column is type <'str'>\n",
+ " data = data[data['location'].isin(LOCATIONS)] # only keep 50 states + D.C. + U.S. (52 total locations)\n",
+ " data = impute_season_week(data=data, fluseason=fluseason, trajectories=trajectories) # add season_week, week_enddate\n",
+ " data['fluseason'] = fluseason # add fluseason column based on explicit specification\n",
+ " # add identifying columns\n",
+ " data['datasetH1'] = datasetH1\n",
+ " datasetH2_list = []\n",
+ " for row in data.itertuples():\n",
+ " value = datasetH2 + row.scenario_id\n",
+ " datasetH2_list.append(value)\n",
+ " data['datasetH2'] = datasetH2_list\n",
+ " data = data.drop(columns=['scenario_id'])\n",
+ " return data\n",
+ "\n",
+ "\n",
+ "def get_trajectory_count(data: pd.DataFrame) -> int:\n",
+ " \"\"\"\n",
+ " Helper function to count how many trajectories per model, per location, per SMH scenario.\n",
+ "\n",
+ " Arbitrary grouping; just needed a number.\n",
+ " \"\"\"\n",
+ " x = data[\n",
+ " (data['scenario_id'] == data['scenario_id'].iloc[0]) & \n",
+ " (data['horizon'] == 1) & \n",
+ " (data['location'] == \"37\") & \n",
+ " (data['age_group'] == '0-4') & \n",
+ " (data['target'] == 'inc hosp') & \n",
+ " (data['output_type'] == 'sample')\n",
+ " ]\n",
+ " return len(x)\n",
+ " \n",
+ "\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 8,
+ "id": "3e45b237",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# trajectories counted using get_trajectory_count(), but code for this is ommitted from this file\n",
+ "\n",
+ "one = pd.read_parquet(\"/Users/emprzy/Documents/work/flu-scenario-modeling-hub_archive/data-processed/USC-SIkJalpha/2023-09-03-USC-SIkJalpha.parquet\")\n",
+ "one = filter_smh(one, 2023, datasetH1=\"SMH_R4\", datasetH2=\"round4_USC-SIkJalpha_\", trajectories=100)\n",
+ "\n",
+ "two = pd.read_parquet(\"/Users/emprzy/Documents/work/flu-scenario-modeling-hub_archive/data-processed/USC-SIkJalpha/2024-08-11-USC-SIkJalpha.parquet\")\n",
+ "two = filter_smh(two, 2024, datasetH1=\"SMH_R5\", datasetH2=\"round5_USC-SIkJalpha_\", trajectories=100)\n",
+ "two = two.drop(columns=['run_grouping', 'stochastic_run']) \n",
+ "\n",
+ "three = pd.read_parquet(\"/Users/emprzy/Documents/work/flu-scenario-modeling-hub_archive/data-processed/NotreDame-FRED/2024-08-11-NotreDame-FRED.gz.parquet\")\n",
+ "three = filter_smh(three, 2024, datasetH1=\"SMH_R5\", datasetH2=\"round5_NotreDame-FRED_\", trajectories=200)\n",
+ "three = three.drop(columns=['run_grouping', 'stochastic_run']) \n",
+ "\n",
+ "four = pd.read_parquet(\"/Users/emprzy/Documents/work/flu-scenario-modeling-hub_archive/data-processed/SigSci-SWIFT/2024-08-11-SigSci-SWIFT.parquet\")\n",
+ "four = filter_smh(four, 2024, datasetH1=\"SMH_R5\", datasetH2=\"round5_SigSci-SWIFT_\", trajectories=300)\n",
+ "four = four.drop(columns=['run_grouping', 'stochastic_run'])\n",
+ "\n",
+ "five = pd.read_parquet(\"/Users/emprzy/Documents/work/flu-scenario-modeling-hub_archive/data-processed/UVA-FluXSim/2024-08-11-UVA-FluXSim.parquet\")\n",
+ "five = filter_smh(five, 2024, datasetH1=\"SMH_R5\", datasetH2=\"round5_UVA-FluXSim_\", trajectories=100)\n",
+ "five = five.drop(columns=['run_grouping', 'stochastic_run'])\n",
+ "\n",
+ "six = pd.read_parquet(\"/Users/emprzy/Documents/work/flu-scenario-modeling-hub_archive/data-processed/UVA-EscapeFlu/2024-08-11-UVA-EscapeFlu.parquet\")\n",
+ "six = filter_smh(six, 2024, datasetH1=\"SMH_R5\", datasetH2=\"round5_UVA-EscapeFlu_\", trajectories=120)\n",
+ "six = six.drop(columns=['run_grouping', 'stochastic_run'])\n",
+ "\n",
+ "seven = pd.read_parquet(\"/Users/emprzy/Documents/work/flu-scenario-modeling-hub_archive/data-processed/ACCIDDA-FlepiMoP/2024-08-11-ACCIDDA-FlepiMoP.gz.parquet\")\n",
+ "seven = filter_smh(seven, 2024, datasetH1=\"SMH_R5\", datasetH2=\"round5_ACCIDDA-FlepiMoP_\", trajectories=100)\n",
+ "seven = seven.drop(columns=['run_grouping', 'stochastic_run'])\n",
+ "\n",
+ "eight = pd.read_parquet(\"/Users/emprzy/Documents/work/flu-scenario-modeling-hub_archive/data-processed/MOBS_NEU-GLEAM_FLU/2024-08-11-MOBS_NEU-GLEAM_FLU.gz.parquet\")\n",
+ "eight = filter_smh(eight, 2024, datasetH1=\"SMH_R5\", datasetH2=\"round5_MOBS_NEU-GLEAM_FLU_\", trajectories=300)\n",
+ "eight = eight.drop(columns=['run_grouping', 'stochastic_run'])\n",
+ "\n",
+ "# !!! THE FOLLOWING DATASETS REPRESENT THE CURRENT SMH ROUND (2025-26 season) !!! \n",
+ "\n",
+ "nine = pd.read_parquet(\"/Users/emprzy/Documents/work/flu-scenario-modeling-hub/model-output/UT-ImmunoSEIRS/2025-08-10-UT-ImmunoSEIRS.gz.parquet\")\n",
+ "nine = filter_smh(nine, 2025, datasetH1=\"SMH_R6\", datasetH2=\"round6_UT-ImmunoSEIRS_\", trajectories=300)\n",
+ "nine = nine.drop(columns=['run_grouping', 'stochastic_run'])\n",
+ "\n",
+ "ten = pd.read_parquet(\"/Users/emprzy/Documents/work/flu-scenario-modeling-hub/model-output/CEPH-MetaFlu/2025-08-10-CEPH-MetaFlu.gz.parquet\")\n",
+ "ten = filter_smh(ten, 2025, datasetH1=\"SMH_R6\", datasetH2=\"round6_CEPH-MetaFlu_\", trajectories=300)\n",
+ "ten = ten.drop(columns=['run_grouping', 'stochastic_run'])\n",
+ "\n",
+ "eleven = pd.read_parquet(\"/Users/emprzy/Documents/work/flu-scenario-modeling-hub/model-output/UVA-FluXSim/2025-08-10-UVA-FluXSim.parquet\")\n",
+ "eleven = filter_smh(eleven, 2025, datasetH1=\"SMH_R6\", datasetH2=\"round6_UVA-FluXSim_\", trajectories=300)\n",
+ "eleven = eleven.drop(columns=['run_grouping', 'stochastic_run'])\n",
+ "\n",
+ "twelve = pd.read_parquet(\"/Users/emprzy/Documents/work/flu-scenario-modeling-hub/model-output/ACCIDDA-FlepiMoP/2025-08-10-ACCIDDA-FlepiMoP.gz.parquet\")\n",
+ "twelve = filter_smh(twelve, 2025, datasetH1=\"SMH_R6\", datasetH2=\"round6_ACCIDDA-FlepiMoP_\", trajectories=600)\n",
+ "twelve = twelve.drop(columns=['run_grouping', 'stochastic_run'])\n",
+ "\n",
+ "thirteen = pd.read_parquet(\"/Users/emprzy/Documents/work/flu-scenario-modeling-hub/model-output/MOBS_NEU-GLEAM_FLU/2025-08-10-MOBS_NEU-GLEAM_FLU.gz.parquet\")\n",
+ "thirteen = filter_smh(thirteen, 2025, datasetH1=\"SMH_R6\", datasetH2=\"round6_MOBS_NEU-GLEAM_FLU_\", trajectories=600)\n",
+ "thirteen = thirteen.drop(columns=['run_grouping', 'stochastic_run'])\n",
+ "\n",
+ "fourteen = pd.read_parquet(\"/Users/emprzy/Documents/work/flu-scenario-modeling-hub/model-output/UNCC-Hierbin/2025-08-10-UNCC-Hierbin.gz.parquet\")\n",
+ "fourteen = filter_smh(fourteen, 2025, datasetH1=\"SMH_R6\", datasetH2=\"round6_UNCC-Hierbin_\", trajectories=300)\n",
+ "fourteen = fourteen.drop(columns=['run_grouping', 'stochastic_run'])\n",
+ "\n",
+ "fifteen = pd.read_parquet(\"/Users/emprzy/Documents/work/flu-scenario-modeling-hub/model-output/PSI-M3/2025-08-10-PSI-M3.gz.parquet\")\n",
+ "fifteen = filter_smh(fifteen, 2025, datasetH1=\"SMH_R6\", datasetH2=\"round6_PSI-M3_\", trajectories=600)\n",
+ "fifteen = fifteen.drop(columns=['run_grouping', 'stochastic_run'])"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 9,
+ "id": "c48b9acc",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# files not separated by age group. save all data to CSV\n",
+ "finished_files = {\n",
+ " \"one\": {\"df\": one, \"name\": \"round4_USC-SIkJalpha.csv\"}, \n",
+ " \"two\": {\"df\": two, \"name\": \"round5_USC-SIkJalpha.csv\"}, \n",
+ " \"three\": {\"df\": three, \"name\": \"round5_NotreDame-FRED.csv\"}, \n",
+ " \"four\": {\"df\": four, \"name\": \"round5_SigSci-SWIFT.csv\"}, \n",
+ " \"five\": {\"df\": five, \"name\": \"round5_UVA-FluXSim.csv\"},\n",
+ " \"six\": {\"df\": six, \"name\": \"round5_UVA-EscapeFlu.csv\"}, \n",
+ " \"seven\": {\"df\": seven, \"name\": \"round5_ACCIDDA-FlepiMoP.csv\"}, \n",
+ " \"eight\": {\"df\": eight, \"name\": \"round5_MOBS_NEU-GLEAM_FLU.csv\"}, \n",
+ " \"nine\": {\"df\": nine, \"name\": \"round6_UT-ImmunoSEIRS.csv\"}, # start of current season (nine-fiften are un-saved to CSV)\n",
+ " \"ten\": {\"df\": ten, \"name\": \"round6_CEPH-MetaFlu.csv\"}, \n",
+ " \"eleven\": {\"df\": eleven, \"name\": \"round6_UVA-FluXSim.csv\"}, \n",
+ " \"twelve\": {\"df\": twelve, \"name\": \"round6_ACCIDDA-FlepiMoP.csv\"}, \n",
+ " \"thirteen\": {\"df\": thirteen, \"name\": \"round6_MOBS_NEU-GLEAM_FLU.csv\"}, \n",
+ " \"fourteen\": {\"df\": fourteen, \"name\": \"round6_UNCC-Hierbin.csv\"}, \n",
+ " \"fifteen\": {\"df\": fifteen, \"name\": \"round6_PSI-M3.csv\"},\n",
+ " \"NHSN\": {\"df\": NHSN, \"name\": \"2024_NHSN_surveillance.csv\"} # surveillance\n",
+ "}\n",
+ "for file, info in finished_files.items():\n",
+ " full_path = '/Users/emprzy/Documents/work/miscellaneous/influpaint_data' + '/' + info[\"name\"]\n",
+ " info[\"df\"].to_csv(full_path, index=False)\n",
+ "\n",
+ "training_data = pd.concat([one, two, three, four, five, six, seven, eight, NHSN]) # includes non-current SMH forecasts and surveillance data\n",
+ "all_data = pd.concat([\n",
+ " one, two, three, four, five, six, seven,\n",
+ " eight, nine, ten, eleven, twelve,\n",
+ " thirteen, fourteen, fifteen, NHSN\n",
+ "])\n",
+ "training_data.to_csv('/Users/emprzy/Documents/work/miscellaneous/influpaint_data/training_data.csv', index=False)\n",
+ "all_data.to_csv('/Users/emprzy/Documents/work/miscellaneous/influpaint_data/all_data.csv', index=False)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "48abdff6",
+ "metadata": {},
+ "source": [
+ "Datatypes for each column:\n",
+ "- week_enddate: pandas._libs.tslibs.timestamps.Timestamp\n",
+ "- location_code: str\n",
+ "- value: numpy.float64\n",
+ "- fluseason_fraction: numpy.float64\n",
+ "- season_week: numpy.int64\n",
+ "- fluseason: numpy.int64\n",
+ "- datasetH1: str\n",
+ "- datasetH2: str\n",
+ "- sample: str"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "50118cc5",
+ "metadata": {},
+ "source": [
+ "# --- DATA IS ALIGNED WITH THE GOAL AT THIS POINT! ---"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "4fff38b2",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "t = pd.read_csv(\n",
+ " \"/Users/emprzy/Documents/work/miscellaneous/influpaint_data/training_data.csv\",\n",
+ " dtype={\n",
+ " 'location_code': str,\n",
+ " 'value': np.float64,\n",
+ " 'fluseason_fraction': np.float64,\n",
+ " 'season_week': np.int64,\n",
+ " 'fluseason': np.int64,\n",
+ " 'datasetH1': str,\n",
+ " 'datasetH2': str,\n",
+ " 'sample': str \n",
+ " },\n",
+ " parse_dates=['week_enddate']\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "df41adff",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "for dH1 in t['datasetH1'].unique():\n",
+ " h1df= t[t['datasetH1'] == dH1]\n",
+ " print(f\"datasetH1: {dH1}, nH2= {len(h1df['datasetH2'].unique())}\")\n",
+ " for dH2 in h1df['datasetH2'].unique():\n",
+ " h2df = h1df[h1df['datasetH2'] == dH2]\n",
+ " print(f\" - datasetH2: {dH2}, shape: {h2df.shape}, years: {len(h2df['fluseason'].unique())}, samples: {len(h2df['sample'].unique())} ===> n_frames={len(h2df['fluseason'].unique())* len(h2df['sample'].unique())}\")\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 14,
+ "id": "7ccaa3e9",
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/html": [
+ "\n",
+ "\n",
+ "
\n",
+ " \n",
+ " \n",
+ " | \n",
+ " week_enddate | \n",
+ " location_code | \n",
+ " value | \n",
+ " fluseason_fraction | \n",
+ " season_week | \n",
+ " fluseason | \n",
+ " datasetH1 | \n",
+ " datasetH2 | \n",
+ " sample | \n",
+ "
\n",
+ " \n",
+ " \n",
+ " \n",
+ " | 0 | \n",
+ " 2010-10-09 | \n",
+ " 02 | \n",
+ " 0.875146 | \n",
+ " 0.189041 | \n",
+ " 10 | \n",
+ " 2010 | \n",
+ " fluview | \n",
+ " fluview | \n",
+ " 1 | \n",
+ "
\n",
+ " \n",
+ " | 1 | \n",
+ " 2010-10-16 | \n",
+ " 02 | \n",
+ " 1.128270 | \n",
+ " 0.208219 | \n",
+ " 11 | \n",
+ " 2010 | \n",
+ " fluview | \n",
+ " fluview | \n",
+ " 1 | \n",
+ "
\n",
+ " \n",
+ " | 2 | \n",
+ " 2010-10-23 | \n",
+ " 02 | \n",
+ " 0.586042 | \n",
+ " 0.227397 | \n",
+ " 12 | \n",
+ " 2010 | \n",
+ " fluview | \n",
+ " fluview | \n",
+ " 1 | \n",
+ "
\n",
+ " \n",
+ " | 3 | \n",
+ " 2010-10-30 | \n",
+ " 02 | \n",
+ " 0.967742 | \n",
+ " 0.246575 | \n",
+ " 13 | \n",
+ " 2010 | \n",
+ " fluview | \n",
+ " fluview | \n",
+ " 1 | \n",
+ "
\n",
+ " \n",
+ " | 4 | \n",
+ " 2010-11-06 | \n",
+ " 02 | \n",
+ " 0.683851 | \n",
+ " 0.265753 | \n",
+ " 14 | \n",
+ " 2010 | \n",
+ " fluview | \n",
+ " fluview | \n",
+ " 1 | \n",
+ "
\n",
+ " \n",
+ " | 5 | \n",
+ " 2010-11-13 | \n",
+ " 02 | \n",
+ " 0.951904 | \n",
+ " 0.284932 | \n",
+ " 15 | \n",
+ " 2010 | \n",
+ " fluview | \n",
+ " fluview | \n",
+ " 1 | \n",
+ "
\n",
+ " \n",
+ "
\n",
+ "
"
+ ],
+ "text/plain": [
+ " week_enddate location_code value fluseason_fraction season_week \\\n",
+ "0 2010-10-09 02 0.875146 0.189041 10 \n",
+ "1 2010-10-16 02 1.128270 0.208219 11 \n",
+ "2 2010-10-23 02 0.586042 0.227397 12 \n",
+ "3 2010-10-30 02 0.967742 0.246575 13 \n",
+ "4 2010-11-06 02 0.683851 0.265753 14 \n",
+ "5 2010-11-13 02 0.951904 0.284932 15 \n",
+ "\n",
+ " fluseason datasetH1 datasetH2 sample \n",
+ "0 2010 fluview fluview 1 \n",
+ "1 2010 fluview fluview 1 \n",
+ "2 2010 fluview fluview 1 \n",
+ "3 2010 fluview fluview 1 \n",
+ "4 2010 fluview fluview 1 \n",
+ "5 2010 fluview fluview 1 "
+ ]
+ },
+ "execution_count": 14,
+ "metadata": {},
+ "output_type": "execute_result"
+ }
+ ],
+ "source": [
+ "goal = pd.read_parquet(\"/Users/emprzy/Documents/work/miscellaneous/josephs_processed_influpaint_data.parquet\")\n",
+ "goal.head(6)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "309c5fcb",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# helper function i created to check how many data types are in a certain column\n",
+ "def type_count(df: pd.DataFrame, col_name: str) -> dict[str, int]:\n",
+ " types = {}\n",
+ " for value in df[col_name]:\n",
+ " current_type = str(type(value))\n",
+ " if current_type not in types:\n",
+ " types[current_type] = 1\n",
+ " else:\n",
+ " types[current_type] += 1\n",
+ " \n",
+ " return types"
+ ]
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "base",
+ "language": "python",
+ "name": "python3"
+ },
+ "language_info": {
+ "codemirror_mode": {
+ "name": "ipython",
+ "version": 3
+ },
+ "file_extension": ".py",
+ "mimetype": "text/x-python",
+ "name": "python",
+ "nbconvert_exporter": "python",
+ "pygments_lexer": "ipython3",
+ "version": "3.13.7"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}