Vectorize the calculation of interval overlap

Author: BSchilperoort

src/xarray_regrid/methods.py

@@ -1,6 +1,5 @@
 from typing import Literal
 
-import numpy as np
 import xarray as xr
 
 from xarray_regrid import utils
@@ -68,16 +67,16 @@ def conservative_regrid(
         source_intervals = utils.to_intervalindex(
             source_coords, resolution=source_coords[1]-source_coords[0]
         )
-        overlap = np.zeros((source_intervals.size, target_intervals.size), dtype=float)
-        for i, source_iv in enumerate(source_intervals):
-            for j, target_iv in enumerate(target_intervals):
-                overlap[i,j] = utils.overlaps(source_iv, target_iv)
+        overlap = utils.overlap(source_intervals, target_intervals)
         weights = utils.normalize_overlap(overlap)
 
-        dot_array = utils.create_dot_dataarray(weights, coord, target_coords, source_coords)
-
+        # TODO: Use `sparse.COO(weights)`. xr.dot does not support this. Much faster!
+        dot_array = utils.create_dot_dataarray(
+            weights, coord, target_coords, source_coords
+        )
+        # TODO: modify weights to correct for latitude.
         dataarrays = [
-            da.dot(dot_array).rename({f"target_{coord}": coord}).rename(da.name)
+            xr.dot(da, dot_array).rename({f"target_{coord}": coord}).rename(da.name)
             for da in dataarrays
         ]
     return xr.merge(dataarrays)  # TODO: add other coordinates/data variables back in.

src/xarray_regrid/utils.py

@@ -108,18 +108,35 @@ def to_intervalindex(coords: np.ndarray, resolution: float) -> pd.IntervalIndex:
     )
 
 
-def overlaps(a: pd.Interval, b: pd.Interval):
-    """Return the overlap (fraction) between two Pandas intervals."""
-    return max(
-        min(a.right, b.right) - max(a.left, b.left),
-        0
+def overlap(a: pd.IntervalIndex, b: pd.IntervalIndex) -> np.ndarray:
+    """Calculate the overlap between two sets of intervals.
+
+    Args:
+        a: Pandas IntervalIndex containing the first set of intervals.
+        b: Pandas IntervalIndex containing the second set of intervals.
+
+    Returns:
+        2D numpy array containing overlap (as a fraction) between the intervals of a
+            and b. If there is no overlap, the value will be 0.
+    """
+    # TODO: newaxis on B and transpose is MUCH faster on benchmark.
+    #  likely due to it being the bigger dimension.
+    #  size(a) > size(b) leads to better perf than size(b) > size(a)
+    mins = np.minimum(
+        a.right.to_numpy(),
+        b.right.to_numpy()[:, np.newaxis]
+    )
+    maxs = np.maximum(
+        a.left.to_numpy(),
+        b.left.to_numpy()[:, np.newaxis]
     )
+    return np.maximum(mins-maxs, 0).T
 
 
 def normalize_overlap(overlap: np.ndarray) -> np.ndarray:
     """Normalize overlap values so they sum up to 1.0 along the first axis."""
     overlap_sum = overlap.sum(axis=0)
-    overlap_sum[overlap_sum==0] = 1e-6  # Avoid dividing by 0
+    overlap_sum[overlap_sum==0] = 1e-12  # Avoid dividing by 0.
     return (overlap / overlap_sum)