Merge pull request #283 from CODEX-CELIDA/interval-ratio-removed-criterion-combination

Interval ratio removed criterion combination

Author: glichtner

execution_engine/execution_engine.py

@@ -434,6 +434,23 @@ def execute_recommendation(
         assert start_datetime.tzinfo, "start_datetime must be timezone-aware"
         assert end_datetime.tzinfo, "end_datetime must be timezone-aware"
 
+        # We cut off microseconds to ensure second-level precision.
+        # In Python, timestamps are usually floored when cast to int,
+        # but in PostgreSQL they might be rounded. To avoid subtle bugs
+        # where a 1-second difference appears due to different rounding
+        # strategies, we explicitly strip microseconds and assert that
+        # the timestamp is exactly an integer.
+        start_datetime = start_datetime.replace(microsecond=0)
+        end_datetime = end_datetime.replace(microsecond=0)
+
+        assert start_datetime.timestamp() == int(
+            start_datetime.timestamp()
+        ), f"start_datetime still contains sub-second precision: {start_datetime}"
+
+        assert end_datetime.timestamp() == int(
+            end_datetime.timestamp()
+        ), f"end_datetime still contains sub-second precision: {end_datetime}"
+
         date_format = "%Y-%m-%d %H:%M:%S %z"
 
         logging.info(

execution_engine/execution_graph/graph.py

@@ -104,6 +104,10 @@ def traverse(
 
         traverse(expr, category=category)
 
+        # make sure base node has still the correct category - it might have changed duing  traversal if the base node
+        # is used in an interval_criterion of a TemporalCount operator
+        graph.nodes[base_node].update({"category": CohortCategory.BASE})
+
         if hash(expr) != expr_hash:
             raise ValueError("Expression has been modified during traversal")
 

execution_engine/omop/criterion/abstract.py

@@ -27,7 +27,8 @@
 from execution_engine.util.interval import IntervalType
 from execution_engine.util.serializable import SerializableDataClassABC
 from execution_engine.util.sql import SelectInto, select_into
-from execution_engine.util.types import PersonIntervals, TimeRange
+from execution_engine.util.types import PersonIntervals
+from execution_engine.util.types.timerange import TimeRange
 
 __all__ = [
     "Criterion",

execution_engine/omop/criterion/point_in_time.py

@@ -10,7 +10,8 @@
 from execution_engine.omop.criterion.concept import ConceptCriterion
 from execution_engine.task.process import get_processing_module
 from execution_engine.util.interval import IntervalType
-from execution_engine.util.types import PersonIntervals, TimeRange, Timing
+from execution_engine.util.types import PersonIntervals, Timing
+from execution_engine.util.types.timerange import TimeRange
 from execution_engine.util.value import Value
 
 process = get_processing_module()

execution_engine/omop/db/celida/tables.py

@@ -27,8 +27,8 @@
 from execution_engine.util.interval import IntervalType
 
 # Use the "public" schema so that tables in different schemas can
-# these enums easily without introducing depedencies between the
-# respective schemas. Note that replicate the enum definitions in each
+# use these enums easily without introducing dependencies between the
+# respective schemas. Note that replicating the enum definitions in each
 # schema would not work when data must be exchanged between the
 # schemas because enum definitions in separate schemas, even if
 # identical in terms of enum values, are considered distinct and
@@ -37,7 +37,6 @@
 CohortCategoryEnum = Enum(CohortCategory, name="cohort_category", schema="public")
 
 
-
 class Recommendation(Base):  # noqa: D101
     __tablename__ = "recommendation"
     __table_args__ = {"schema": SCHEMA_NAME}
@@ -170,7 +169,9 @@ class ResultInterval(Base):  # noqa: D101
     interval_start: Mapped[datetime]
     interval_end: Mapped[datetime]
     interval_type = mapped_column(IntervalTypeEnum)
-
+    interval_ratio: Mapped[float] = mapped_column(
+        nullable=True
+    )
     execution_run: Mapped["ExecutionRun"] = relationship(
         primaryjoin="ResultInterval.run_id == ExecutionRun.run_id",
     )

execution_engine/omop/db/celida/views.py

@@ -199,6 +199,7 @@ def interval_result_view() -> Select:
             rri.c.interval_type,
             rri.c.interval_start,
             rri.c.interval_end,
+            rri.c.interval_ratio,
         )
         .select_from(rri)
         .outerjoin(pip, (rri.c.pi_pair_id == pip.c.pi_pair_id))

execution_engine/omop/sqlclient.py

@@ -58,6 +58,35 @@ def _enable_database_triggers(
     cursor.close()
 
 
+def datetime_cols_to_epoch(stmt: sqlalchemy.Select) -> sqlalchemy.Select:
+    """
+    Given a SQLAlchemy 2.0 Select that has columns labeled 'interval_start'
+    or 'interval_end', replace those column expressions with
+    EXTRACT(EPOCH FROM <expr>)::BIGINT so they become integer timestamps.
+
+    Returns a new Select object with the replaced columns.
+    """
+    new_columns = []
+
+    for col in stmt.selected_columns:
+        label = getattr(col, "name")
+
+        if label in ("interval_start", "interval_end"):
+            # We'll wrap col in EXTRACT(EPOCH FROM col)::BIGINT,
+            new_col = (
+                sqlalchemy.func.extract("epoch", col)
+                .cast(sqlalchemy.BigInteger)
+                .label(label)
+            )
+            new_columns.append(new_col)
+        else:
+            new_columns.append(col)
+
+    new_stmt = stmt.with_only_columns(*new_columns, maintain_column_froms=True)
+
+    return new_stmt
+
+
 class OMOPSQLClient:
     """A client for the OMOP SQL database.
 

execution_engine/task/process/__init__.py

@@ -3,6 +3,10 @@
 import sys
 import types
 from collections import namedtuple
+from typing import TypeVar
+
+from execution_engine.util.interval import IntervalType
+from execution_engine.util.types.timerange import TimeRange
 
 
 def get_processing_module(
@@ -39,6 +43,35 @@ def get_processing_module(
 
 Interval = namedtuple("Interval", ["lower", "upper", "type"])
 IntervalWithCount = namedtuple("IntervalWithCount", ["lower", "upper", "type", "count"])
-IntervalWithTypeCounts = namedtuple(
-    "IntervalWithTypeCounts", ["lower", "upper", "counts"]
-)
+
+AnyInterval = Interval | IntervalWithCount
+GeneralizedInterval = None | AnyInterval
+
+TInterval = TypeVar("TInterval", bound=AnyInterval)
+
+
+def interval_like(interval: TInterval, start: int, end: int) -> TInterval:
+    """
+    Return a copy of the given interval with its lower and upper bounds replaced.
+
+    Args:
+        interval (I): The interval to copy. Must be one of Interval or IntervalWithCount.
+        start (datetime): The new lower bound.
+        end (datetime): The new upper bound.
+
+    Returns:
+        I: A copy of the interval with updated lower and upper bounds.
+    """
+
+    return interval._replace(lower=start, upper=end)  # type: ignore[return-value]
+
+
+def timerange_to_interval(tr: TimeRange, type_: IntervalType) -> Interval:
+    """
+    Converts a timerange to an interval with the supplied type.
+    """
+    return Interval(
+        lower=int(tr.start.timestamp()),
+        upper=int(tr.end.timestamp()),
+        type=type_,
+    )