add united states again

CDCgov · Feb 5, 2025 · 2f34ea9 · 2f34ea9
1 parent a37b3b8
commit 2f34ea9
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 21 deletions.
diff --git a/README.md b/README.md
@@ -131,24 +131,24 @@ loc_table = forecasttools.location_table
 print(loc_table)
 ```
 
-    shape: (58, 4)
-    ┌───────────────┬────────────┬─────────────────────────────┬────────────┐
-    │ location_code ┆ short_name ┆ long_name                   ┆ population │
-    │ ---           ┆ ---        ┆ ---                         ┆ ---        │
-    │ str           ┆ str        ┆ str                         ┆ i64        │
-    ╞═══════════════╪════════════╪═════════════════════════════╪════════════╡
-    │ US            ┆ US         ┆ United States               ┆ 334735155  │
-    │ 01            ┆ AL         ┆ Alabama                     ┆ 5024279    │
-    │ 02            ┆ AK         ┆ Alaska                      ┆ 733391     │
-    │ 04            ┆ AZ         ┆ Arizona                     ┆ 7151502    │
-    │ 05            ┆ AR         ┆ Arkansas                    ┆ 3011524    │
-    │ …             ┆ …          ┆ …                           ┆ …          │
-    │ 66            ┆ GU         ┆ Guam                        ┆ null       │
-    │ 69            ┆ MP         ┆ Northern Mariana Islands    ┆ null       │
-    │ 72            ┆ PR         ┆ Puerto Rico                 ┆ 3285874    │
-    │ 74            ┆ UM         ┆ U.S. Minor Outlying Islands ┆ null       │
-    │ 78            ┆ VI         ┆ U.S. Virgin Islands         ┆ null       │
-    └───────────────┴────────────┴─────────────────────────────┴────────────┘
+    shape: (58, 5)
+    ┌───────────────┬────────────┬─────────────────────────────┬────────────┬──────────┐
+    │ location_code ┆ short_name ┆ long_name                   ┆ population ┆ is_state │
+    │ ---           ┆ ---        ┆ ---                         ┆ ---        ┆ ---      │
+    │ str           ┆ str        ┆ str                         ┆ i64        ┆ bool     │
+    ╞═══════════════╪════════════╪═════════════════════════════╪════════════╪══════════╡
+    │ US            ┆ US         ┆ United States               ┆ 334735155  ┆ false    │
+    │ 01            ┆ AL         ┆ Alabama                     ┆ 5024279    ┆ true     │
+    │ 02            ┆ AK         ┆ Alaska                      ┆ 733391     ┆ true     │
+    │ 04            ┆ AZ         ┆ Arizona                     ┆ 7151502    ┆ true     │
+    │ 05            ┆ AR         ┆ Arkansas                    ┆ 3011524    ┆ true     │
+    │ …             ┆ …          ┆ …                           ┆ …          ┆ …        │
+    │ 66            ┆ GU         ┆ Guam                        ┆ null       ┆ false    │
+    │ 69            ┆ MP         ┆ Northern Mariana Islands    ┆ null       ┆ false    │
+    │ 72            ┆ PR         ┆ Puerto Rico                 ┆ 3285874    ┆ false    │
+    │ 74            ┆ UM         ┆ U.S. Minor Outlying Islands ┆ null       ┆ false    │
+    │ 78            ┆ VI         ┆ U.S. Virgin Islands         ┆ null       ┆ false    │
+    └───────────────┴────────────┴─────────────────────────────┴────────────┴──────────┘
 
 Using `./forecasttools/data.py`, the location table was created by
 running the following:
@@ -174,7 +174,7 @@ united_states = forecasttools.united_states
 print(united_states)
 ```
 
-    ['United States', 'Alabama', 'Alaska', 'Arizona', 'Arkansas', 'California', 'Colorado', 'Connecticut', 'Delaware', 'District of Columbia', 'Florida', 'Georgia', 'Hawaii', 'Idaho', 'Illinois', 'Indiana', 'Iowa', 'Kansas', 'Kentucky', 'Louisiana', 'Maine', 'Maryland', 'Massachusetts', 'Michigan', 'Minnesota', 'Mississippi', 'Missouri', 'Montana', 'Nebraska', 'Nevada', 'New Hampshire', 'New Jersey', 'New Mexico', 'New York', 'North Carolina', 'North Dakota', 'Ohio', 'Oklahoma', 'Oregon', 'Pennsylvania', 'Rhode Island', 'South Carolina', 'South Dakota', 'Tennessee', 'Texas', 'Utah', 'Vermont', 'Virginia', 'Washington', 'West Virginia', 'Wisconsin', 'Wyoming', 'American Samoa', 'Guam', 'Northern Mariana Islands', 'Puerto Rico', 'U.S. Minor Outlying Islands', 'U.S. Virgin Islands']
+    ['Alabama', 'Alaska', 'Arizona', 'Arkansas', 'California', 'Colorado', 'Connecticut', 'Delaware', 'Florida', 'Georgia', 'Hawaii', 'Idaho', 'Illinois', 'Indiana', 'Iowa', 'Kansas', 'Kentucky', 'Louisiana', 'Maine', 'Maryland', 'Massachusetts', 'Michigan', 'Minnesota', 'Mississippi', 'Missouri', 'Montana', 'Nebraska', 'Nevada', 'New Hampshire', 'New Jersey', 'New Mexico', 'New York', 'North Carolina', 'North Dakota', 'Ohio', 'Oklahoma', 'Oregon', 'Pennsylvania', 'Rhode Island', 'South Carolina', 'South Dakota', 'Tennessee', 'Texas', 'Utah', 'Vermont', 'Virginia', 'Washington', 'West Virginia', 'Wisconsin', 'Wyoming']
 
 ## Example FluSight Hub Submission
 

diff --git a/forecasttools/__init__.py b/forecasttools/__init__.py
@@ -27,12 +27,12 @@
     validate_iter_has_expected_types,
 )
 
-# location table (from Census data)
+# location table (from Census data; contains territory data)
 location_table_path = importlib.resources.files(__package__).joinpath(
     "location_table.parquet"
 )
 location_table = pl.read_parquet(location_table_path)
-united_states = location_table["long_name"].to_list()
+united_states = location_table.filter(pl.col("is_state")).select("long_name").to_series().to_list()
 
 # load example flusight submission
 example_flusight_submission_path = importlib.resources.files(

diff --git a/forecasttools/data.py b/forecasttools/data.py
@@ -7,6 +7,8 @@
 an example FluSight submission.
 """
 
+
+
 import os
 import pathlib
 from urllib import error, request
@@ -106,6 +108,18 @@ def merge_pop_data_and_loc_data(
     if save_path.exists() and not overwrite:
         print(f"File already exists at {save_path}. Skipping writing.")
         return
+    us_states = [
+        "Alabama", "Alaska", "Arizona", "Arkansas", "California", "Colorado",
+        "Connecticut", "Delaware", "Florida", "Georgia", "Hawaii", "Idaho",
+        "Illinois", "Indiana", "Iowa", "Kansas", "Kentucky", "Louisiana",
+        "Maine", "Maryland", "Massachusetts", "Michigan", "Minnesota",
+        "Mississippi", "Missouri", "Montana", "Nebraska", "Nevada",
+        "New Hampshire", "New Jersey", "New Mexico", "New York",
+        "North Carolina", "North Dakota", "Ohio", "Oklahoma", "Oregon",
+        "Pennsylvania", "Rhode Island", "South Carolina", "South Dakota",
+        "Tennessee", "Texas", "Utah", "Vermont", "Virginia", "Washington",
+        "West Virginia", "Wisconsin", "Wyoming"
+    ]
     pop_df = pl.read_parquet(population_path).select(
         [
             pl.col("STNAME").alias("long_name"),
@@ -123,10 +137,14 @@ def merge_pop_data_and_loc_data(
         .otherwise(pl.col("population"))
         .alias("population")
     )
+    merged_df = merged_df.with_columns(
+        pl.col("long_name").is_in(us_states).alias("is_state")
+    )
     merged_df.write_parquet(save_path)
     print(f"File successfully written to {save_path}")
 
 
+
 def make_census_dataset(
     file_save_path: str,
 ) -> None:

diff --git a/forecasttools/location_table.parquet b/forecasttools/location_table.parquet