feat: handle lookahead-based outcome resolution

Aarhus-Psychiatry-Research · Aug 21, 2023 · 489003f · 489003f
1 parent 2918452
commit 489003f
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 5 deletions.
diff --git a/psycop/common/feature_generation/sequences/timeseries_windower_python/patient.py b/psycop/common/feature_generation/sequences/timeseries_windower_python/patient.py
@@ -46,9 +46,10 @@ def add_static_events(self, features: List[StaticFeature]):
     def get_static_events(self) -> Sequence[StaticFeature]:
         return self._static_features
 
-    def to_prediction_sequences(
+    def to_prediction_times(
         self,
         lookbehind: dt.timedelta,
+        lookahead: dt.timedelta,
         outcome_timestamp: dt.datetime,
         prediction_timestamps: Sequence[dt.datetime],
     ) -> list[PredictionTime]:
@@ -63,13 +64,17 @@ def to_prediction_sequences(
                 end=prediction_timestamp,
             )
 
+            outcome_within_lookahead = outcome_timestamp <= (
+                prediction_timestamp + lookahead
+            )
+
             # 2. Return prediction sequences
             prediction_sequences.append(
                 PredictionTime(
                     patient=self,
                     prediction_timestamp=prediction_timestamp,
                     temporal_events=filtered_events,
-                    outcome=outcome_timestamp,
+                    outcome=outcome_within_lookahead,
                     static_features=self._static_features,
                 ),
             )

diff --git a/psycop/common/feature_generation/sequences/timeseries_windower_python/prediction_time.py b/psycop/common/feature_generation/sequences/timeseries_windower_python/prediction_time.py
@@ -22,9 +22,8 @@
 class PredictionTime:
     """A cut sequence of events for a patient, ready to issue a prediction."""
 
-    # TODO: rename to PredictionTime and add pointer to Patient and drop slices
     patient: Patient
     temporal_events: Sequence[TemporalEvent]
     static_features: Sequence[StaticFeature]
     prediction_timestamp: dt.datetime
-    outcome: dt.datetime | None  # TODO: Maybe this should be a bool for the given lookahead window? Where do we want this parsing to happen?
+    outcome: bool
diff --git a/psycop/common/feature_generation/sequences/timeseries_windower_python/test_patient.py b/psycop/common/feature_generation/sequences/timeseries_windower_python/test_patient.py
@@ -31,7 +31,7 @@ def test_temporal_event_filtering(self):
         )
         patient.add_temporal_events(temporal_events)
 
-        prediction_sequences = patient.to_prediction_sequences(
+        prediction_sequences = patient.to_prediction_times(
             lookbehind=dt.timedelta(days=2),
             prediction_timestamps=[dt.datetime(2021, 1, 2), dt.datetime(2021, 1, 4)],
             outcome_timestamp=dt.datetime(2021, 1, 5),