Format code, improve typing. Add pandas-stubs.

Author: BSchilperoort

benchmarks/benchmarking_conservative.ipynb

@@ -1095,7 +1095,8 @@
    ],
    "source": [
     "import numpy as np\n",
-    "rmse = np.sqrt(np.mean((data_regrid - data_cdo)**2))[\"tp\"].to_numpy() * 1000\n",
+    "\n",
+    "rmse = np.sqrt(np.mean((data_regrid - data_cdo) ** 2))[\"tp\"].to_numpy() * 1000\n",
     "print(f\"RMSE: {rmse:.4f} mm\")"
    ]
   },

pyproject.toml

@@ -51,6 +51,7 @@ dev = [
   "mypy",
   "pytest",
   "pytest-cov",
+  "pandas-stubs", # Adds typing for pandas.
 ]
 
 [tool.hatch.version]

src/xarray_regrid/methods.py

@@ -61,18 +61,18 @@ def conservative_regrid(
         target_coords = coords[coord].to_numpy()
         # TODO: better resolution/IntervalIndex inference
         target_intervals = utils.to_intervalindex(
-            target_coords, resolution=target_coords[1]-target_coords[0]
+            target_coords, resolution=target_coords[1] - target_coords[0]
         )
         source_coords = data[coord].to_numpy()
         source_intervals = utils.to_intervalindex(
-            source_coords, resolution=source_coords[1]-source_coords[0]
+            source_coords, resolution=source_coords[1] - source_coords[0]
         )
         overlap = utils.overlap(source_intervals, target_intervals)
         weights = utils.normalize_overlap(overlap)
 
         # 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
+            weights, str(coord), target_coords, source_coords
         )
         # TODO: modify weights to correct for latitude.
         dataarrays = [

src/xarray_regrid/utils.py

@@ -102,7 +102,7 @@ def to_intervalindex(coords: np.ndarray, resolution: float) -> pd.IntervalIndex:
     """
     return pd.IntervalIndex(
         [
-            pd.Interval(left=coord - resolution/2, right=coord + resolution/2)
+            pd.Interval(left=coord - resolution / 2, right=coord + resolution / 2)
             for coord in coords
         ]
     )
@@ -122,25 +122,25 @@ def overlap(a: pd.IntervalIndex, b: pd.IntervalIndex) -> np.ndarray:
     # 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
+    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])
+    overlap: np.ndarray = np.maximum(mins - maxs, 0).T
+    return overlap
 
 
 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-12  # Avoid dividing by 0.
-    return (overlap / overlap_sum)
+    overlap_sum: np.ndarray = overlap.sum(axis=0)
+    overlap_sum[overlap_sum == 0] = 1e-12  # Avoid dividing by 0.
+    return overlap / overlap_sum  # type: ignore
 
 
-def create_dot_dataarray(weights, coord, target_coords, source_coords):
+def create_dot_dataarray(
+    weights: np.ndarray,
+    coord: str,
+    target_coords: np.ndarray,
+    source_coords: np.ndarray,
+) -> xr.DataArray:
     """Create a DataArray to be used at dot product compatible with xr.dot."""
     return xr.DataArray(
         data=weights,