diff --git a/execution_engine/omop/criterion/combination/temporal.py b/execution_engine/omop/criterion/combination/temporal.py index 98c50254..35f762e7 100644 --- a/execution_engine/omop/criterion/combination/temporal.py +++ b/execution_engine/omop/criterion/combination/temporal.py @@ -15,6 +15,8 @@ class TimeIntervalType(StrEnum): MORNING_SHIFT = "morning_shift" AFTERNOON_SHIFT = "afternoon_shift" NIGHT_SHIFT = "night_shift" + NIGHT_SHIFT_BEFORE_MIDNIGHT = "night_shift_before_midnight" + NIGHT_SHIFT_AFTER_MIDNIGHT = "night_shift_after_midnight" DAY = "day" ANY_TIME = "any_time" @@ -259,6 +261,34 @@ def NightShift( return cls.Presence(criterion, category, TimeIntervalType.NIGHT_SHIFT) + @classmethod + def NightShiftBeforeMidnight( + cls, + criterion: Union[Criterion, "CriterionCombination"], + category: CohortCategory, + ) -> "TemporalIndicatorCombination": + """ + Create a NightShiftBeforeMidnight combination of criteria. + """ + + return cls.Presence( + criterion, category, TimeIntervalType.NIGHT_SHIFT_BEFORE_MIDNIGHT + ) + + @classmethod + def NightShiftAfterMidnight( + cls, + criterion: Union[Criterion, "CriterionCombination"], + category: CohortCategory, + ) -> "TemporalIndicatorCombination": + """ + Create a NightShiftAfterMidnight combination of criteria. + """ + + return cls.Presence( + criterion, category, TimeIntervalType.NIGHT_SHIFT_AFTER_MIDNIGHT + ) + @classmethod def Day( cls, diff --git a/execution_engine/settings.py b/execution_engine/settings.py index 3a5f156d..c13cdae5 100644 --- a/execution_engine/settings.py +++ b/execution_engine/settings.py @@ -44,6 +44,12 @@ class TimeIntervalSettings(BaseModel): night_shift: TimeInterval = TimeInterval( start=datetime.time(22, 0, 0), end=datetime.time(5, 59, 59) ) + night_shift_before_midnight: TimeInterval = TimeInterval( + start=datetime.time(22, 0, 0), end=datetime.time(23, 59, 59) + ) + night_shift_after_midnight: TimeInterval = TimeInterval( + start=datetime.time(0, 0, 0), end=datetime.time(5, 59, 59) + ) day: TimeInterval = TimeInterval( start=datetime.time(0, 0, 0), end=datetime.time(23, 59, 59) ) diff --git a/tests/execution_engine/omop/criterion/combination/test_temporal_combination.py b/tests/execution_engine/omop/criterion/combination/test_temporal_combination.py index cae1e71e..9381c4bb 100644 --- a/tests/execution_engine/omop/criterion/combination/test_temporal_combination.py +++ b/tests/execution_engine/omop/criterion/combination/test_temporal_combination.py @@ -683,6 +683,88 @@ def patient_events(self, db_session, person_visit): ), {1: set(), 2: set(), 3: set()}, ), + ####################### + # Partial Night Shifts (before midnight) + ####################### + ( + TemporalIndicatorCombination.NightShiftBeforeMidnight( + c1, category=CohortCategory.POPULATION + ), + { + 1: { + ( + pendulum.parse("2023-03-01 22:00:00+01:00"), + pendulum.parse("2023-03-01 23:59:59+01:00"), + ), + ( + pendulum.parse("2023-03-02 22:00:00+01:00"), + pendulum.parse("2023-03-02 23:59:59+01:00"), + ), + }, + 2: set(), + 3: set(), + }, + ), + ( + TemporalIndicatorCombination.NightShiftBeforeMidnight( + c2, category=CohortCategory.POPULATION + ), + { + 1: { + ( + pendulum.parse("2023-03-02 22:00:00+01:00"), + pendulum.parse("2023-03-02 23:59:59+01:00"), + ), + }, + 2: set(), + 3: set(), + }, + ), + ( + TemporalIndicatorCombination.NightShiftBeforeMidnight( + c3, category=CohortCategory.POPULATION + ), + {1: set(), 2: set(), 3: set()}, + ), + ####################### + # Partial Night Shifts (after midnight) + ####################### + ( + TemporalIndicatorCombination.NightShiftAfterMidnight( + c1, category=CohortCategory.POPULATION + ), + { + 1: { + ( + pendulum.parse("2023-03-02 00:00:00+01:00"), + pendulum.parse("2023-03-02 05:59:59+01:00"), + ), + }, + 2: set(), + 3: set(), + }, + ), + ( + TemporalIndicatorCombination.NightShiftAfterMidnight( + c2, category=CohortCategory.POPULATION + ), + { + 1: { + ( + pendulum.parse("2023-03-03 00:00:00+01:00"), + pendulum.parse("2023-03-03 05:59:59+01:00"), + ), + }, + 2: set(), + 3: set(), + }, + ), + ( + TemporalIndicatorCombination.NightShiftAfterMidnight( + c3, category=CohortCategory.POPULATION + ), + {1: set(), 2: set(), 3: set()}, + ), ], ) def test_combination_on_database(