Merge pull request #123 from washingtonpost/extrapolation-rule

small changes to extrapolation

Author: lennybronner

src/elexmodel/client.py

@@ -328,10 +328,16 @@ def get_estimates(
                 start_date=model_parameters.get("versioned_start_date", None),
                 end_date=model_parameters.get("versioned_end_date", None),
             )
-            print(
-                "Fetching versioned data between ", versioned_data_handler.start_date, versioned_data_handler.end_date
+            LOG.info(
+                "Fetching versioned data between %s and %s",
+                versioned_data_handler.start_date,
+                versioned_data_handler.end_date,
             )
-            versioned_data_handler.get_versioned_results(model_settings.get("versioned_filepath", None))
+            versioned_results = versioned_data_handler.get_versioned_results(
+                model_settings.get("versioned_filepath", None)
+            )
+            if versioned_results is None:
+                versioned_data_handler = None
         else:
             versioned_data_handler = None
 

src/elexmodel/handlers/data/VersionedData.py

@@ -1,3 +1,5 @@
+from datetime import datetime
+
 import numpy as np
 import pandas as pd
 from dateutil import tz
@@ -30,9 +32,8 @@ def __init__(
             target_bucket = "elex-models-prod"
         else:
             target_bucket = TARGET_BUCKET
-
-        start_date = start_date.astimezone(tz=tz.gettz("UTC")) if start_date else None
-        end_date = end_date.astimezone(tz=tz.gettz("UTC")) if start_date else None
+        start_date = datetime.fromisoformat(start_date).astimezone(tz=tz.gettz("UTC")) if start_date else None
+        end_date = datetime.fromisoformat(end_date).astimezone(tz=tz.gettz("UTC")) if start_date else None
         # versioned results natively are in UTC but we'll convert it back to timezone in tzinfo
         self.s3_client = s3.S3VersionUtil(target_bucket, start_date, end_date, tzinfo)
 
@@ -66,6 +67,9 @@ def get_versioned_results(self, filepath=None):
             path = f"{S3_FILE_PATH}/{self.election_id}/results/{self.office_id}/{self.geographic_unit_type}/current.csv"
 
         data = self.s3_client.get(path, self.sample)
+        if data is None:
+            self.data = data
+            return data
         estimandizer = Estimandizer()
         data, _ = estimandizer.add_estimand_results(data, self.estimands, False)
         self.data = data.sort_values("last_modified")

src/elexmodel/handlers/s3.py

@@ -105,7 +105,11 @@ def list_versions(self, path, **kwargs):
         if "Versions" in response:
             versions = response["Versions"]
 
-        if response["IsTruncated"] and len(versions) > 0 and versions[-1]["LastModified"] >= self.start_date:
+        if (
+            response["IsTruncated"]
+            and len(versions) > 0
+            and (self.start_date is None or versions[-1]["LastModified"] >= self.start_date)
+        ):
             versions += self.list_versions(
                 path,
                 KeyMarker=response["NextKeyMarker"],
@@ -145,7 +149,8 @@ def make_request(self, path, *, version=None, **kwargs):
     def get(self, path, sample=2):
         versions = self.list_versions(path)
         if len(versions) == 0:
-            raise ValueError(f"No versions found for {path}")
+            LOG.info(f"No versions found for {path}")
+            return None
 
         # Instead of asking for the results of downloads synchronously, we're
         # queuing the futures and then waiting for them to complete.

src/elexmodel/models/BootstrapElectionModel.py

@@ -1,6 +1,7 @@
 from __future__ import annotations  # pylint: disable=too-many-lines
 
 import logging
+from datetime import timedelta
 from itertools import combinations
 
 import numpy as np
@@ -108,7 +109,7 @@ def __init__(self, model_settings={}, versioned_data_handler=None):
         self.contest_correlations = model_settings.get("contest_correlations", [])
 
         # impose perfect correlation in the national summary aggregation
-        self.national_summary_correlation = model_settings.get("national_summary_correlation", False)
+        self.national_summary_correlation = model_settings.get("national_summary_correlation", True)
         self.stop_model_call = None
         # Assume that we have a baseline normalized margin
         # (D^{Y'} - R^{Y'}) / (D^{Y'} + R^{Y'}) is one of the covariates
@@ -795,6 +796,7 @@ def _extrapolate_unit_margin(self, reporting_units: pd.DataFrame, nonreporting_u
         all_units = pd.concat([reporting_units, nonreporting_units], axis=0).copy()
         missing_columns = list(set(self.versioned_data_handler.data.columns) - set(all_units.columns))
         all_units[missing_columns] = self.versioned_data_handler.data[missing_columns].max()
+        all_units["last_modified"] = self.versioned_data_handler.data["last_modified"].max() + timedelta(seconds=1)
 
         self.versioned_data_handler.data = pd.concat(
             [self.versioned_data_handler.data, all_units[self.versioned_data_handler.data.columns]], axis=0