Merge pull request #79 from hubverse-org/fix/location-dtype-inference

Fix dtype inference for models with numeric-only location codes

Author: nickreich

src/hub_predtimechart/__init__.py

@@ -1 +1 @@
-__version__ = "2.2.4"
+__version__ = "2.2.5"

src/hub_predtimechart/app/generate_json_files.py

@@ -5,6 +5,7 @@
 
 import click
 import pandas as pd
+import pyarrow.compute as pc
 import structlog
 
 from hub_predtimechart.generate_data import forecast_data_for_model_df
@@ -86,13 +87,14 @@ def _generate_forecast_json_files(hub_config: HubConfigPtc, output_dir: Path, is
             for model_id in hub_config.model_id_to_metadata:  # ex: ['Flusight-baseline', 'MOBS-GLEAM_FLUH', ...]
                 model_output_file = hub_config.model_output_file_for_ref_date(model_id, reference_date)
                 if model_output_file:
-                    if model_output_file.suffix == '.csv':
-                        model_id_to_df[model_id] = pd.read_csv(model_output_file, usecols=df_cols_to_use)
-                    elif model_output_file.suffix in ['.parquet', '.pqt']:
-                        model_id_to_df[model_id] = pd.read_parquet(model_output_file, columns=df_cols_to_use)
-                    else:
-                        raise RuntimeError(f"unsupported model output file type: {model_output_file!r}. "
-                                           f"Only .csv and .parquet are supported")
+                    # Use hubdata's to_table() method with filtering to load only this model's data
+                    # for this reference_date. This applies the schema from tasks.json, ensuring
+                    # task_id columns (like location) are properly typed as strings, preventing
+                    # dtype inference issues with numeric-only values like "01", "02"
+                    filter_expr = (pc.field('model_id') == model_id) & \
+                                  (pc.field(hub_config.reference_date_col_name) == date.fromisoformat(reference_date))
+                    pa_table = hub_config.to_table(columns=df_cols_to_use, filter=filter_expr)
+                    model_id_to_df[model_id] = pa_table.to_pandas()
 
             if not model_id_to_df:  # no model outputs for reference_date
                 continue

tests/expected/example-complex-forecast-hub/forecasts/wk-inc-flu-hosp_01_2022-10-22.json

@@ -22,5 +22,13 @@
     "q0.5": [118, 175],
     "q0.75": [133, 193],
     "q0.975": [165, 233]
+  },
+  "Test-NumericOnly": {
+    "target_end_date": ["2022-10-29", "2022-11-05"],
+    "q0.025": [1542, 1443],
+    "q0.25": [1542, 1443],
+    "q0.5": [1724, 1724],
+    "q0.75": [1906, 2006],
+    "q0.975": [2334, 2562]
   }
 }

tests/expected/example-complex-forecast-hub/forecasts/wk-inc-flu-hosp_02_2022-10-22.json

@@ -0,0 +1,28 @@
+{
+    "Test-NumericOnly": {
+        "target_end_date": [
+            "2022-10-29",
+            "2022-11-05"
+        ],
+        "q0.025": [
+            100.0,
+            110.0
+        ],
+        "q0.25": [
+            150.0,
+            160.0
+        ],
+        "q0.5": [
+            200.0,
+            210.0
+        ],
+        "q0.75": [
+            250.0,
+            260.0
+        ],
+        "q0.975": [
+            300.0,
+            310.0
+        ]
+    }
+}

tests/hub_predtimechart/test_generate_json_files.py

@@ -16,6 +16,7 @@ def test_generate_forecast_json_files_ecfh(tmp_path):
     json_files = _generate_forecast_json_files(hub_config, output_dir)
     assert set(json_files) == {output_dir / 'wk-inc-flu-hosp_US_2022-10-22.json',
                                output_dir / 'wk-inc-flu-hosp_01_2022-10-22.json',
+                               output_dir / 'wk-inc-flu-hosp_02_2022-10-22.json',
                                output_dir / 'wk-inc-flu-hosp_US_2022-11-19.json',
                                output_dir / 'wk-inc-flu-hosp_01_2022-11-19.json',
                                output_dir / 'wk-inc-flu-hosp_US_2022-12-17.json',
@@ -28,6 +29,7 @@ def test_generate_forecast_json_files_ecfh(tmp_path):
             assert act_data == exp_data
 
 
+
 def test_generate_forecast_json_files_flu_metrocast(tmp_path):
     """
     An integration test of `generate_json_files.py`'s `_generate_json_files()` for flu-metrocast.
@@ -90,6 +92,7 @@ def test_generate_forecast_json_files_skip_files(tmp_path):
     json_files = Path(output_dir).glob("*")
     assert set(json_files) == {output_dir / 'wk-inc-flu-hosp_US_2022-10-22.json',
                                output_dir / 'wk-inc-flu-hosp_01_2022-10-22.json',
+                               output_dir / 'wk-inc-flu-hosp_02_2022-10-22.json',
                                output_dir / 'wk-inc-flu-hosp_US_2022-11-19.json',
                                output_dir / 'wk-inc-flu-hosp_01_2022-11-19.json',
                                output_dir / 'wk-inc-flu-hosp_01_2022-12-17.json'}
@@ -122,6 +125,7 @@ def test_generate_forecast_json_files_regenerate(tmp_path):
     json_files = _generate_forecast_json_files(hub_config, output_dir, True)
     assert set(json_files) == {output_dir / 'wk-inc-flu-hosp_US_2022-10-22.json',
                                output_dir / 'wk-inc-flu-hosp_01_2022-10-22.json',
+                               output_dir / 'wk-inc-flu-hosp_02_2022-10-22.json',
                                output_dir / 'wk-inc-flu-hosp_US_2022-11-19.json',
                                output_dir / 'wk-inc-flu-hosp_01_2022-11-19.json',
                                output_dir / 'wk-inc-flu-hosp_US_2022-12-17.json',

tests/hub_predtimechart/test_hub_config_ptc.py

@@ -21,7 +21,7 @@ def test_hub_config_complex_forecast_hub():
     assert hub_config.initial_checked_models == ['Flusight-baseline']
     assert hub_config.disclaimer == "Most forecasts have failed to reliably predict rapid changes in the trends of reported cases and hospitalizations. Due to this limitation, they should not be relied upon for decisions about the possibility or timing of rapid changes in trends."
     assert (sorted(list(hub_config.model_id_to_metadata.keys())) ==
-            sorted(['Flusight-baseline', 'MOBS-GLEAM_FLUH', 'PSI-DICE']))
+            sorted(['Flusight-baseline', 'MOBS-GLEAM_FLUH', 'PSI-DICE', 'Test-NumericOnly']))
     assert hub_config.target_data_file_name == 'covid-hospital-admissions.csv'
 
     model_task_0 = hub_config.model_tasks[0]  # only one

tests/hubs/example-complex-forecast-hub/hub-config/admin.json

@@ -11,7 +11,7 @@
         "owner": "Infectious-Disease-Modeling-Hubs",
         "repository": "example-complex-forecast-hub"
     },
-    "file_format": ["csv"],
+    "file_format": ["csv", "parquet"],
     "timezone": "US/Eastern",
     "model_output_dir": "model-output",
     "cloud": {

tests/hubs/example-complex-forecast-hub/model-metadata/Test-NumericOnly.yml

@@ -0,0 +1,21 @@
+team_name: "Test"
+team_abbr: "Test"
+model_name: "Test model for numeric-only location codes"
+model_abbr: "NumericOnly"
+model_version: "1.0"
+model_contributors: [
+    {
+        "name": "Test User",
+        "email": "test@example.com"
+    }
+]
+website_url: "https://github.com/hubverse-org/hub-dashboard-predtimechart"
+license: "MIT"
+team_funding: "N/A"
+designated_model: false
+methods: "Test model for regression testing dtype inference with numeric-only location codes like '01', '02'."
+data_inputs: "None - test data"
+methods_long: "This is a test model used to verify that models with exclusively numeric location codes (e.g., '01', '02') are properly handled and not excluded from dashboard visualizations due to dtype inference issues."
+ensemble_of_models: false
+ensemble_of_hub_models: false
+source_notes: "Test data for issue #78"

tests/hubs/example-complex-forecast-hub/model-output/Test-NumericOnly/2022-10-22-Test-NumericOnly.csv

@@ -0,0 +1,21 @@
+location,horizon,output_type_id,value,target_end_date,reference_date,output_type,target
+"01",1,0.025,1542,2022-10-29,2022-10-22,quantile,wk inc flu hosp
+"01",1,0.25,1542,2022-10-29,2022-10-22,quantile,wk inc flu hosp
+"01",1,0.5,1724,2022-10-29,2022-10-22,quantile,wk inc flu hosp
+"01",1,0.75,1906,2022-10-29,2022-10-22,quantile,wk inc flu hosp
+"01",1,0.975,2334,2022-10-29,2022-10-22,quantile,wk inc flu hosp
+"01",2,0.025,1443,2022-11-05,2022-10-22,quantile,wk inc flu hosp
+"01",2,0.25,1443,2022-11-05,2022-10-22,quantile,wk inc flu hosp
+"01",2,0.5,1724,2022-11-05,2022-10-22,quantile,wk inc flu hosp
+"01",2,0.75,2006,2022-11-05,2022-10-22,quantile,wk inc flu hosp
+"01",2,0.975,2562,2022-11-05,2022-10-22,quantile,wk inc flu hosp
+"02",1,0.025,100,2022-10-29,2022-10-22,quantile,wk inc flu hosp
+"02",1,0.25,150,2022-10-29,2022-10-22,quantile,wk inc flu hosp
+"02",1,0.5,200,2022-10-29,2022-10-22,quantile,wk inc flu hosp
+"02",1,0.75,250,2022-10-29,2022-10-22,quantile,wk inc flu hosp
+"02",1,0.975,300,2022-10-29,2022-10-22,quantile,wk inc flu hosp
+"02",2,0.025,110,2022-11-05,2022-10-22,quantile,wk inc flu hosp
+"02",2,0.25,160,2022-11-05,2022-10-22,quantile,wk inc flu hosp
+"02",2,0.5,210,2022-11-05,2022-10-22,quantile,wk inc flu hosp
+"02",2,0.75,260,2022-11-05,2022-10-22,quantile,wk inc flu hosp
+"02",2,0.975,310,2022-11-05,2022-10-22,quantile,wk inc flu hosp