feat: compute "interval ratio" in logical count operators

Also implement the operators via find_rectangles and remove from the
process module the functions count_intervals and
filter_count_intervals.

Author: scymtym

execution_engine/task/process/rectangle.py

@@ -428,110 +428,6 @@ def union_intervals(data: list[PersonIntervals]) -> PersonIntervals:
     return _process_intervals(data, _impl.union_interval_lists)
 
 
-def interval_to_interval_with_count(interval: Interval) -> IntervalWithCount:
-    """
-    Converts an Interval to an IntervalWithCount.
-    """
-    return IntervalWithCount(interval.lower, interval.upper, interval.type, 1)
-
-
-def intervals_to_intervals_with_count(
-    intervals: list[Interval],
-) -> list[IntervalWithCount]:
-    """
-    Converts a list of Intervals to a list of IntervalWithCount.
-    """
-    return [interval_to_interval_with_count(interval) for interval in intervals]
-
-
-def count_intervals(data: list[PersonIntervals]) -> PersonIntervalsWithCount:
-    """
-    Counts the intervals per dict key in the list.
-
-    :param data: A list of dict of intervals.
-    :return: A dict with the unioned intervals.
-    """
-    if not len(data):
-        return dict()
-
-    # assert dfs is a list of dataframes
-    assert isinstance(data, list) and all(
-        isinstance(arr, dict) for arr in data
-    ), "data must be a list of dicts"
-
-    result = {}
-
-    for arr in data:
-        if not len(arr):
-            # if the operation is union, an empty dataframe can be ignored
-            continue
-
-        for group_keys, intervals in arr.items():
-            intervals_with_count = intervals_to_intervals_with_count(intervals)
-            intervals_with_count = _impl.union_rects_with_count(intervals_with_count)
-            if group_keys not in result:
-                result[group_keys] = intervals_with_count
-            else:
-                result[group_keys] = _impl.union_rects_with_count(
-                    result[group_keys] + intervals_with_count
-                )
-
-    return result
-
-
-def filter_count_intervals(
-    data: PersonIntervalsWithCount,
-    min_count: int | None,
-    max_count: int | None,
-    keep_no_data: bool = True,
-    keep_not_applicable: bool = True,
-) -> PersonIntervals:
-    """
-    Filters the intervals per dict key in the list by count.
-
-    :param data: A list of dict of intervals.
-    :param min_count: The minimum count of the intervals.
-    :param max_count: The maximum count of the intervals.
-    :param keep_no_data: Whether to keep NO_DATA intervals (irrespective of the count).
-    :param keep_not_applicable: Whether to keep NOT_APPLICABLE intervals (irrespective of the count).
-    :return: A dict with the unioned intervals.
-    """
-
-    result: PersonIntervals = {}
-
-    interval_filter = []
-
-    if keep_no_data:
-        interval_filter.append(IntervalType.NO_DATA)
-    if keep_not_applicable:
-        interval_filter.append(IntervalType.NOT_APPLICABLE)
-
-    if min_count is None and max_count is None:
-        raise ValueError("min_count and max_count cannot both be None")
-    elif min_count is not None and max_count is not None:
-        for person_id in data:
-            result[person_id] = [
-                Interval(interval.lower, interval.upper, interval.type)
-                for interval in data[person_id]
-                if min_count <= interval.count <= max_count
-                or interval.type in interval_filter
-            ]
-    elif min_count is not None:
-        for person_id in data:
-            result[person_id] = [
-                Interval(interval.lower, interval.upper, interval.type)
-                for interval in data[person_id]
-                if min_count <= interval.count or interval.type in interval_filter
-            ]
-    elif max_count is not None:
-        for person_id in data:
-            result[person_id] = [
-                Interval(interval.lower, interval.upper, interval.type)
-                for interval in data[person_id]
-                if interval.count <= max_count or interval.type in interval_filter
-            ]
-
-    return result
 
 
 def intersect_intervals(data: list[PersonIntervals]) -> PersonIntervals:

execution_engine/task/process/rectangle_python.py

@@ -105,90 +105,6 @@ def union_rects(intervals: list[Interval]) -> list[Interval]:
         return union
 
 
-def union_rects_with_count(
-    intervals: list[IntervalWithCount],
-) -> list[IntervalWithCount]:
-    """
-    Unions the intervals while keeping track of the count of overlapping intervals of the same type.
-    """
-
-    if not len(intervals):
-        return []
-
-    with IntervalType.union_order():
-        events = intervals_to_events(intervals)
-
-        union = []
-
-        last_x_start = -np.inf  # holds the x_min of the currently open output rectangle
-        last_x_end = events[0][
-            0
-        ]  # x variable of the last closed interval (we start with the first x, so we
-        # don't close the first rectangle at the first x)
-        previous_x_visited = -np.inf
-        open_y = SortedDict()
-
-        def get_y_max() -> IntervalType | None:
-            max_key = None
-            for key in reversed(open_y):
-                if open_y[key] > 0:
-                    max_key = key
-                    break
-            return max_key
-
-        for x, start_point, interval in events:
-            y, count_event = interval.type, interval.count
-            if start_point:
-                y_max = get_y_max()
-
-                if x > previous_x_visited and y_max is None:
-                    # no currently open rectangles
-                    last_x_start = x  # start new output rectangle
-                elif y >= y_max:
-                    if x == last_x_end or x == last_x_start:
-                        # we already closed a rectangle at this x, so we don't need to start a new one
-                        open_y[y] = open_y.get(y, 0) + count_event
-                        continue
-
-                    union.append(
-                        IntervalWithCount(
-                            lower=last_x_start,
-                            upper=x - 1,
-                            type=y_max,
-                            count=open_y[y_max],
-                        )
-                    )
-                    last_x_end = x
-                    last_x_start = x
-
-                open_y[y] = open_y.get(y, 0) + count_event
-
-            else:
-                open_y[y] = max(open_y.get(y, 0) - count_event, 0)
-
-                y_max = get_y_max()
-
-                if (y_max is None or (open_y and y_max <= y)) and x > last_x_end:
-                    if y_max is None or y_max < y:
-                        # the closing rectangle has a higher y_max than the currently open ones
-                        count = count_event
-                    else:
-                        # the closing rectangle has the same y_max as the currently open ones
-                        count = open_y[y] + count_event
-
-                    union.append(
-                        IntervalWithCount(
-                            lower=last_x_start, upper=x - 1, type=y, count=count
-                        )
-                    )  # close the previous rectangle at y_max
-                    last_x_end = x
-                    last_x_start = x  # start new output rectangle
-
-            previous_x_visited = x
-
-        return merge_adjacent_intervals(union)
-
-
 def merge_adjacent_intervals(
     intervals: list[IntervalWithCount],
 ) -> list[IntervalWithCount]:

execution_engine/task/task.py

@@ -306,12 +306,50 @@ def handle_binary_logical_operator(
         elif isinstance(self.expr, (logic.Or, logic.NonSimplifiableOr)):
             result = process.union_intervals(data)
         elif isinstance(self.expr, logic.Count):
-            result = process.count_intervals(data)
-            result = process.filter_count_intervals(
-                result,
-                min_count=self.expr.count_min,
-                max_count=self.expr.count_max,
-            )
+            # result = process.count_intervals(data)
+            # result = process.filter_count_intervals(
+            #     result,
+            #     min_count=self.expr.count_min,
+            #     max_count=self.expr.count_max,
+            # )
+            count_min = self.expr.count_min
+            count_max = self.expr.count_max
+            if count_min is None and count_max is None:
+                raise ValueError("count_min and count_max cannot both be None")
+            def interval_counts(
+                    start: int, end: int, intervals: List[AnyInterval]
+            ) -> GeneralizedInterval:
+                positive_count, negative_count, not_applicable_count, no_data_count = 0, 0, 0, 0
+                for interval in intervals:
+                    if interval is None or interval.type is IntervalType.NEGATIVE:
+                        negative_count += 1
+                    elif interval.type is IntervalType.POSITIVE:
+                        positive_count += 1
+                    elif interval.type is IntervalType.NOT_APPLICABLE:
+                        not_applicable_count += 1
+                    elif interval.type is IntervalType.NO_DATA:
+                        no_data_count += 1
+                #
+                if positive_count > 0:
+                    if count_min is None:
+                        interval_type = IntervalType.POSITIVE if (
+                                    positive_count <= count_max) else IntervalType.NEGATIVE
+                        return Interval(start, end, interval_type)
+                    else:
+                        min_good = count_min <= positive_count
+                        max_good = (count_max is None) or (positive_count <= count_max)
+                        interval_type = IntervalType.POSITIVE if (min_good and max_good) else IntervalType.NEGATIVE
+                        ratio = positive_count / count_min
+                        return IntervalWithCount(start, end, interval_type, ratio)
+                if no_data_count > 0:
+                    return Interval(start, end, IntervalType.NO_DATA)
+                if not_applicable_count > 0:
+                    return Interval(start, end, IntervalType.NOT_APPLICABLE)
+                if negative_count > 0:
+                    return Interval(start, end, IntervalType.NEGATIVE)
+
+            result = process.find_rectangles(data, interval_counts)
+
         elif isinstance(self.expr, logic.CappedCount):
 
             def interval_counts(