Repository: Aarhus-Psychiatry-Research/timeseriesflattener

Issue:
fix: remove hardcoded timestamp names

Discussion:

Patch:
diff --git a/src/timeseriesflattener/feature_spec_objects.py b/src/timeseriesflattener/feature_spec_objects.py
index b498393..a1726bf 100644
--- a/src/timeseriesflattener/feature_spec_objects.py
+++ b/src/timeseriesflattener/feature_spec_objects.py
@@ -232,9 +232,6 @@ class TemporalSpec(AnySpec):
     id_col_name: str = "id"
     # Col name for ids in the input dataframe.
 
-    timestamp_col_name: str = "timestamp"
-    # Col name for timestamps in the input dataframe.
-
     loader_kwargs: Optional[dict] = None
 
     # Optional keyword arguments for the data loader
@@ -255,25 +252,6 @@ class TemporalSpec(AnySpec):
 
         super().__init__(**data)
 
-        timestamp_col_type = self.values_df[self.timestamp_col_name].dtype  # type: ignore
-
-        if timestamp_col_type not in ("Timestamp", "datetime64[ns]"):
-            # Convert column dtype to datetime64[ns] if it isn't already
-            log.info(
-                f"{self.feature_name}: Converting timestamp column to datetime64[ns]",
-            )
-
-            self.values_df[self.timestamp_col_name] = pd.to_datetime(
-                self.values_df[self.timestamp_col_name],
-            )
-
-            min_timestamp = min(self.values_df[self.timestamp_col_name])
-
-            if min_timestamp < pd.Timestamp("1971-01-01"):
-                log.warning(
-                    f"{self.feature_name}: Minimum timestamp is {min_timestamp} - perhaps ints were coerced to timestamps?",
-                )
-
         self.resolve_multiple_fn = data["resolve_multiple_fn"]
 
         # override fallback strings with objects
@@ -305,20 +283,6 @@ class PredictorSpec(TemporalSpec):
 
         super().__init__(**data)
 
-    def get_cutoff_date(self) -> pd.Timestamp:
-        """Get the cutoff date from a spec.
-
-        A cutoff date is the earliest date that a prediction time can get data from the values_df.
-        We do not want to include those prediction times, as we might make incorrect inferences.
-        For example, if a spec says to look 5 years into the future, but we only have one year of data,
-        there will necessarily be fewer outcomes - without that reflecting reality. This means our model won't generalise.
-
-        Returns:
-            pd.Timestamp: A cutoff date.
-        """
-        min_val_date = self.values_df[self.timestamp_col_name].min()  # type: ignore
-        return min_val_date + pd.Timedelta(days=self.lookbehind_days)
-
 
 class OutcomeSpec(TemporalSpec):
     """Specification for a single outcome, where the df has been resolved."""
@@ -360,21 +324,6 @@ class OutcomeSpec(TemporalSpec):
 
         return len(self.values_df[col_name].unique()) <= 2  # type: ignore
 
-    def get_cutoff_date(self) -> pd.Timestamp:
-        """Get the cutoff date from a spec.
-
-        A cutoff date is the earliest date that a prediction time can get data from the values_df.
-        We do not want to include those prediction times, as we might make incorrect inferences.
-        For example, if a spec says to look 5 years into the future, but we only have one year of data,
-        there will necessarily be fewer outcomes - without that reflecting reality. This means our model won't generalise.
-
-        Returns:
-            pd.Timestamp: A cutoff date.
-        """
-        max_val_date = self.values_df[self.timestamp_col_name].max()  # type: ignore
-
-        return max_val_date - pd.Timedelta(days=self.lookahead_days)
-
 
 class MinGroupSpec(BaseModel):
     """Minimum specification for a group of features, whether they're looking ahead or behind.
diff --git a/src/timeseriesflattener/flattened_dataset.py b/src/timeseriesflattener/flattened_dataset.py
index fa5dba2..9a1f786 100644
--- a/src/timeseriesflattener/flattened_dataset.py
+++ b/src/timeseriesflattener/flattened_dataset.py
@@ -46,6 +46,7 @@ class SpecCollection(PydanticBaseModel):
     static_specs: list[AnySpec] = []
 
     def __len__(self):
+        """Return number of specs in collection."""
         return (
             len(self.outcome_specs) + len(self.predictor_specs) + len(self.static_specs)
         )
@@ -185,6 +186,7 @@ class TimeseriesFlattener:  # pylint: disable=too-many-instance-attributes
         df: DataFrame,
         pred_times_with_uuid: DataFrame,
         pred_time_uuid_colname: str,
+        pred_timestamp_col_name: str,
     ) -> DataFrame:
         """Ensure all prediction times are represented in the returned
 
@@ -194,6 +196,7 @@ class TimeseriesFlattener:  # pylint: disable=too-many-instance-attributes
             df (DataFrame): Dataframe with prediction times but without uuid.
             pred_times_with_uuid (DataFrame): Dataframe with prediction times and uuid.
             pred_time_uuid_colname (str): Name of uuid column in both df and pred_times_with_uuid.
+            pred_itmestamp_col_name (str): Name of timestamp column in df.
 
         Returns:
             DataFrame: A merged dataframe with all prediction times.
@@ -204,13 +207,14 @@ class TimeseriesFlattener:  # pylint: disable=too-many-instance-attributes
             how="left",
             on=pred_time_uuid_colname,
             suffixes=("", "_temp"),
-        ).drop(["timestamp_pred"], axis=1)
+        ).drop([pred_timestamp_col_name], axis=1)
 
     @staticmethod
     def _resolve_multiple_values_within_interval_days(
         resolve_multiple: Callable,
         df: DataFrame,
         pred_time_uuid_colname: str,
+        val_timestamp_col_name: str,
     ) -> DataFrame:
         """Apply the resolve_multiple function to prediction_times where there
 
@@ -227,14 +231,14 @@ class TimeseriesFlattener:  # pylint: disable=too-many-instance-attributes
         # Convert timestamp val to numeric that can be used for resolve_multiple functions
         # Numeric value amounts to days passed since 1/1/1970
         try:
-            df["timestamp_val"] = (
-                df["timestamp_val"] - dt.datetime(1970, 1, 1)
+            df[val_timestamp_col_name] = (
+                df[val_timestamp_col_name] - dt.datetime(1970, 1, 1)
             ).dt.total_seconds() / 86400
         except TypeError:
             log.info("All values are NaT, returning empty dataframe")
 
         # Sort by timestamp_pred in case resolve_multiple needs dates
-        df = df.sort_values(by="timestamp_val").groupby(pred_time_uuid_colname)
+        df = df.sort_values(by=val_timestamp_col_name).groupby(pred_time_uuid_colname)
 
         if isinstance(resolve_multiple, str):
             resolve_multiple = resolve_multiple_fns.get(resolve_multiple)
@@ -300,6 +304,7 @@ class TimeseriesFlattener:  # pylint: disable=too-many-instance-attributes
         output_spec: AnySpec,
         id_col_name: str,
         pred_time_uuid_col_name: str,
+        timestamp_col_name: str,
         verbose: bool = False,  # noqa
     ) -> DataFrame:
         """Create a dataframe with flattened values (either predictor or
@@ -317,6 +322,7 @@ class TimeseriesFlattener:  # pylint: disable=too-many-instance-attributes
                 static method.
             pred_time_uuid_col_name (str): Name of uuid column in
                 prediction_times_with_uuid_df. Required because this is a static method.
+            timestamp_col_name (str): Name of timestamp column in. Required because this is a static method.
             verbose (bool, optional): Whether to print progress.
 
 
@@ -334,6 +340,9 @@ class TimeseriesFlattener:  # pylint: disable=too-many-instance-attributes
             validate="m:m",
         ).drop(id_col_name, axis=1)
 
+        timestamp_val_col_name = f"{timestamp_col_name}_val"
+        timestamp_pred_col_name = f"{timestamp_col_name}_pred"
+
         # Drop prediction times without event times within interval days
         if isinstance(output_spec, OutcomeSpec):
             direction = "ahead"
@@ -346,8 +355,8 @@ class TimeseriesFlattener:  # pylint: disable=too-many-instance-attributes
             df,
             direction=direction,
             interval_days=output_spec.interval_days,
-            timestamp_pred_colname="timestamp_pred",
-            timestamp_value_colname="timestamp_val",
+            timestamp_pred_colname=timestamp_pred_col_name,
+            timestamp_value_colname=timestamp_val_col_name,
         )
 
         # Add back prediction times that don't have a value, and fill them with fallback
@@ -355,14 +364,16 @@ class TimeseriesFlattener:  # pylint: disable=too-many-instance-attributes
             df=df,
             pred_times_with_uuid=prediction_times_with_uuid_df,
             pred_time_uuid_colname=pred_time_uuid_col_name,
+            pred_timestamp_col_name=timestamp_pred_col_name,
         ).fillna(output_spec.fallback)
 
-        df["timestamp_val"].replace({output_spec.fallback: pd.NaT}, inplace=True)
+        df[timestamp_val_col_name].replace({output_spec.fallback: pd.NaT}, inplace=True)
 
         df = TimeseriesFlattener._resolve_multiple_values_within_interval_days(
             resolve_multiple=output_spec.resolve_multiple_fn,
             df=df,
             pred_time_uuid_colname=pred_time_uuid_col_name,
+            val_timestamp_col_name=timestamp_val_col_name,
         )
 
         # If resolve_multiple generates empty values,
@@ -415,6 +426,7 @@ class TimeseriesFlattener:  # pylint: disable=too-many-instance-attributes
             id_col_name=self.id_col_name,
             pred_time_uuid_col_name=self.pred_time_uuid_col_name,
             output_spec=feature_spec,
+            timestamp_col_name=self.timestamp_col_name,
         )
 
         # Write df to cache if exists
@@ -614,6 +626,26 @@ class TimeseriesFlattener:  # pylint: disable=too-many-instance-attributes
 
         self._df = df
 
+    def _get_cutoff_date_from_spec(self, spec: TemporalSpec) -> pd.Timestamp:
+        """Get the cutoff date from a spec.
+
+        A cutoff date is the earliest date that a prediction time can get data from the values_df.
+        We do not want to include those prediction times, as we might make incorrect inferences.
+        For example, if a spec says to look 5 years into the future, but we only have one year of data,
+        there will necessarily be fewer outcomes - without that reflecting reality. This means our model won't generalise.
+
+        Returns:
+            pd.Timestamp: A cutoff date.
+        """
+
+        if isinstance(spec, PredictorSpec):
+            min_val_date = spec.values_df[self.timestamp_col_name].min()  # type: ignore
+            return min_val_date + pd.Timedelta(days=spec.lookbehind_days)
+
+        if isinstance(spec, OutcomeSpec):
+            max_val_date = spec.values_df[self.timestamp_col_name].max()  # type: ignore
+            return max_val_date - pd.Timedelta(days=spec.lookahead_days)
+
     @print_df_dimensions_diff
     def _drop_pred_time_if_insufficient_look_distance(self, df: pd.DataFrame):
         """Drop prediction times if there is insufficient look distance.
@@ -638,7 +670,7 @@ class TimeseriesFlattener:  # pylint: disable=too-many-instance-attributes
         cutoff_date_ahead = pd.Timestamp("2200-01-01")
 
         for spec in spec_batch:
-            spec_cutoff_date = spec.get_cutoff_date()
+            spec_cutoff_date = self._get_cutoff_date_from_spec(spec=spec)
 
             if isinstance(spec, OutcomeSpec):
                 cutoff_date_ahead = min(cutoff_date_ahead, spec_cutoff_date)
@@ -692,6 +724,30 @@ class TimeseriesFlattener:  # pylint: disable=too-many-instance-attributes
             if col not in spec.values_df.columns:  # type: ignore
                 raise ValueError(f"Missing required column: {col}")
 
+    def _check_that_spec_df_timestamp_col_is_correctly_formatted(
+        self,
+        spec: TemporalSpec,
+    ):
+        """Check that timestamp column is correctly formatted. Attempt to coerce if possible."""
+        timestamp_col_type = spec.values_df[self.timestamp_col_name].dtype  # type: ignore
+
+        if timestamp_col_type not in ("Timestamp", "datetime64[ns]"):
+            # Convert column dtype to datetime64[ns] if it isn't already
+            log.info(
+                f"{spec.feature_name}: Converting timestamp column to datetime64[ns]",
+            )
+
+            spec.values_df[self.timestamp_col_name] = pd.to_datetime(  # type: ignore
+                spec.values_df[self.timestamp_col_name],  # type: ignore
+            )
+
+            min_timestamp = min(spec.values_df[self.timestamp_col_name])  # type: ignore
+
+            if min_timestamp < pd.Timestamp("1971-01-01"):
+                log.warning(
+                    f"{spec.feature_name}: Minimum timestamp is {min_timestamp} - perhaps ints were coerced to timestamps?",
+                )
+
     def add_spec(
         self,
         spec: Union[list[AnySpec], AnySpec],
@@ -720,6 +776,11 @@ class TimeseriesFlattener:  # pylint: disable=too-many-instance-attributes
 
             self._check_that_spec_df_has_required_columns(spec=spec_i)
 
+            if isinstance(spec_i, TemporalSpec):
+                self._check_that_spec_df_timestamp_col_is_correctly_formatted(
+                    spec=spec_i,
+                )
+
             if isinstance(spec_i, OutcomeSpec):
                 self.unprocessed_specs.outcome_specs.append(spec_i)
             elif isinstance(spec_i, PredictorSpec):


Test patch:
diff --git a/tests/test_timeseriesflattener/test_flattened_dataset/test_flattened_dataset.py b/tests/test_timeseriesflattener/test_flattened_dataset/test_flattened_dataset.py
index 7eea49e..c4d36b0 100644
--- a/tests/test_timeseriesflattener/test_flattened_dataset/test_flattened_dataset.py
+++ b/tests/test_timeseriesflattener/test_flattened_dataset/test_flattened_dataset.py
@@ -106,22 +106,24 @@ def test_compute_specs(
 
 def test_drop_pred_time_if_insufficient_look_distance():
     # Create a sample DataFrame with some test data
+    # Uses datetime to also test that using another column name works
     pred_time_df = pd.DataFrame(
         {
             "id": [1, 1, 1, 1],
-            "timestamp": ["2022-01-01", "2022-01-02", "2022-01-03", "2022-01-04"],
+            "datetime": ["2022-01-01", "2022-01-02", "2022-01-03", "2022-01-04"],
         },
     )
 
     ts_flattener = TimeseriesFlattener(
         prediction_times_df=pred_time_df,
         drop_pred_times_with_insufficient_look_distance=True,
+        timestamp_col_name="datetime",
     )
 
     pred_val_df = pd.DataFrame(
         {
             "id": [1],
-            "timestamp": ["2022-01-01"],
+            "datetime": ["2022-01-01"],
             "value": [1],
         },
     )
@@ -138,7 +140,7 @@ def test_drop_pred_time_if_insufficient_look_distance():
     out_val_df = pd.DataFrame(
         {
             "id": [1],
-            "timestamp": ["2022-01-05"],
+            "datetime": ["2022-01-05"],
             "value": [4],
         },
     )
@@ -157,7 +159,7 @@ def test_drop_pred_time_if_insufficient_look_distance():
     out_df = ts_flattener.get_df()
 
     # Assert that the correct rows were dropped from the DataFrame
-    expected_df = pd.DataFrame({"timestamp": ["2022-01-02", "2022-01-03"]})
+    expected_df = pd.DataFrame({"datetime": ["2022-01-02", "2022-01-03"]})
     # Convert to datetime to avoid a warning
-    expected_df = expected_df.astype({"timestamp": "datetime64[ns]"})
-    pd.testing.assert_series_equal(out_df["timestamp"], expected_df["timestamp"])
+    expected_df = expected_df.astype({"datetime": "datetime64[ns]"})
+    pd.testing.assert_series_equal(out_df["datetime"], expected_df["datetime"])