marEx 4.0.4__tar.gz → 4.1__tar.gz

{marex-4.0.4 → marex-4.1}/.github/workflows/ci.yml

@@ -147,7 +147,7 @@ jobs:
     - name: Run preprocessing tests (gridded)
       timeout-minutes: 30  # Windows can be slow for these tests
       run: |
-        pytest tests/test_gridded_preprocessing.py -n 2 -v --tb=short
+        pytest tests/test_gridded_preprocessing.py tests/test_detect_detrending.py -n 2 -v --tb=short
 
     - name: Run preprocessing tests (unstructured)
       timeout-minutes: 30  # Windows can be slow for these tests
@@ -164,6 +164,11 @@ jobs:
       run: |
         pytest tests/test_unstructured_tracking.py -n 2 -v --tb=short
 
+    - name: Run tracking edge case tests
+      timeout-minutes: 30  # Windows can be slow for these tests
+      run: |
+        pytest tests/test_track_edge_cases.py -n 2 -v --tb=short
+
     - name: Run plotting tests
       run: |
         pytest tests/test_plotx.py -n 2 -v --tb=short

{marex-4.0.4 → marex-4.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: marEx
-Version: 4.0.4
+Version: 4.1
 Summary: Marine Extremes Detection and Tracking
 Author-email: Aaron Wienkers <aaron@wienkers.com>
 Project-URL: Homepage, https://github.com/wienkers/marEx
@@ -87,7 +87,7 @@ Requires-Dist: pytest-xdist; extra == "test"
 Requires-Dist: coverage[toml]>=6.0.0; extra == "test"
 Requires-Dist: psutil; extra == "test"
 
-![logo](https://github.com/user-attachments/assets/df99a84b-c88f-447d-a362-ac538d9b924b)
+<img src="https://github.com/user-attachments/assets/55eab465-9aca-4754-a981-e26a36f7d045" alt="marEx_animated" width="100%">
 
 [![CI](https://github.com/wienkers/marEx/actions/workflows/ci.yml/badge.svg)](https://github.com/wienkers/marEx/actions/workflows/ci.yml)
 [![codecov](https://codecov.io/gh/wienkers/marEx/branch/main/graph/badge.svg)](https://codecov.io/gh/wienkers/marEx)

{marex-4.0.4 → marex-4.1}/README.md

@@ -1,4 +1,4 @@
-![logo](https://github.com/user-attachments/assets/df99a84b-c88f-447d-a362-ac538d9b924b)
+<img src="https://github.com/user-attachments/assets/55eab465-9aca-4754-a981-e26a36f7d045" alt="marEx_animated" width="100%">
 
 [![CI](https://github.com/wienkers/marEx/actions/workflows/ci.yml/badge.svg)](https://github.com/wienkers/marEx/actions/workflows/ci.yml)
 [![codecov](https://codecov.io/gh/wienkers/marEx/branch/main/graph/badge.svg)](https://codecov.io/gh/wienkers/marEx)

{marex-4.0.4 → marex-4.1}/marEx/detect.py

@@ -298,6 +298,7 @@ def preprocess_data(
     std_normalise: bool = False,  # for detrend_harmonic
     detrend_orders: Optional[List[int]] = None,  # "
     force_zero_mean: bool = True,  # "
+    reference_period: Optional[Tuple[int, int]] = None,  # for fixed_baseline & detrend_fixed_baseline
     method_percentile: Literal["exact", "approximate"] = "approximate",
     precision: float = 0.01,
     max_anomaly: float = 5.0,
@@ -354,6 +355,12 @@ def preprocess_data(
         Default is 1st order (linear) detrend. `[1,2]` e.g. would use a linear+quadratic detrending.
     force_zero_mean : bool, default=True
         Whether to enforce zero mean in detrended anomalies (detrend_harmonic method only).
+    reference_period : tuple of (int, int), optional
+        Year range (start_year, end_year) inclusive for computing the daily climatology
+        (fixed_baseline and detrend_fixed_baseline only). If None (default), uses all
+        available years. Anomalies are computed for the full time series regardless.
+        Example: reference_period=(1990, 2020) computes the climatology from 1990-2020
+        but outputs anomalies for the entire input time range.
     method_percentile : str, default='approximate'
         Method for percentile calculation ('exact' or 'approximate') for both global_extreme & hobday_extreme methods.
         N.B.: Using the exact percentile calculation requires both careful/thoughtful chunking & sufficient memory,
@@ -560,6 +567,17 @@ def preprocess_data(
             data_info={"data_type": type(da.data).__name__, "shape": da.shape},
         )
 
+    # Validate reference_period before triggering any computation
+    if reference_period is not None and method_anomaly not in ("fixed_baseline", "detrend_fixed_baseline"):
+        raise ConfigurationError(
+            f"reference_period is not supported for method_anomaly='{method_anomaly}'",
+            details="reference_period is only applicable to 'fixed_baseline' and 'detrend_fixed_baseline' methods",
+            suggestions=[
+                "Remove the reference_period parameter, or",
+                "Use method_anomaly='fixed_baseline' or 'detrend_fixed_baseline'",
+            ],
+        )
+
     # Validate that all unmasked data is valid (finite values only)
     logger.debug("Validating data values for NaN/infinite values")
     _validate_data_values(da, dimensions)
@@ -588,6 +606,7 @@ def preprocess_data(
             std_normalise,
             detrend_orders,
             force_zero_mean,
+            reference_period,
             use_temp_checkpoints,
         )
         log_memory_usage(logger, "After anomaly computation", logging.DEBUG)
@@ -723,6 +742,7 @@ def preprocess_data(
                 smooth_days_baseline,
                 window_days_hobday,
                 window_spatial_hobday,
+                reference_period,
             ),
         }
     )
@@ -744,14 +764,18 @@ def preprocess_data(
             }
         )
     elif method_anomaly == "fixed_baseline":
-        ds.attrs.update({})  # No method-specific parameters
+        attrs = {}
+        if reference_period is not None:
+            attrs["reference_period"] = list(reference_period)
+        ds.attrs.update(attrs)
     elif method_anomaly == "detrend_fixed_baseline":
-        ds.attrs.update(
-            {
-                "detrend_orders": detrend_orders,
-                "force_zero_mean": force_zero_mean,
-            }
-        )
+        attrs = {
+            "detrend_orders": detrend_orders,
+            "force_zero_mean": force_zero_mean,
+        }
+        if reference_period is not None:
+            attrs["reference_period"] = list(reference_period)
+        ds.attrs.update(attrs)
 
     if method_extreme == "hobday_extreme":
         ds.attrs.update({"window_days_hobday": window_days_hobday})
@@ -826,6 +850,7 @@ def _get_preprocessing_steps(
     smooth_days_baseline: int,
     window_days_hobday: int,
     window_spatial_hobday: Optional[int],
+    reference_period: Optional[Tuple[int, int]] = None,
 ) -> List[str]:
     """Generate preprocessing steps description based on selected methods."""
     steps = []
@@ -838,10 +863,16 @@ def _get_preprocessing_steps(
         steps.append(f"Rolling climatology using {window_year_baseline} years")
         steps.append(f"Smoothed with {smooth_days_baseline}-day window")
     elif method_anomaly == "fixed_baseline":
-        steps.append("Daily climatology computed from full time series")
+        if reference_period is not None:
+            steps.append(f"Daily climatology computed from {reference_period[0]}-{reference_period[1]}")
+        else:
+            steps.append("Daily climatology computed from full time series")
     elif method_anomaly == "detrend_fixed_baseline":
         steps.append(f"Removed polynomial trend orders={detrend_orders}")
-        steps.append("Daily climatology computed from detrended data")
+        if reference_period is not None:
+            steps.append(f"Daily climatology computed from detrended data ({reference_period[0]}-{reference_period[1]})")
+        else:
+            steps.append("Daily climatology computed from detrended data")
 
     # Extreme method steps
     if method_extreme == "global_extreme":
@@ -869,6 +900,7 @@ def compute_normalised_anomaly(
     std_normalise: bool = False,  # for detrend_harmonic
     detrend_orders: Optional[List[int]] = None,  # "
     force_zero_mean: bool = True,  # "
+    reference_period: Optional[Tuple[int, int]] = None,  # for fixed_baseline & detrend_fixed_baseline
     use_temp_checkpoints: bool = False,
     verbose: Optional[bool] = None,
     quiet: Optional[bool] = None,
@@ -901,7 +933,10 @@ def compute_normalised_anomaly(
         Polynomial orders for trend removal (detrend_harmonic and detrend_fixed_baseline only)
     force_zero_mean : bool, default=True
         Explicitly enforce zero mean in final anomalies (detrend_harmonic and detrend_fixed_baseline only)
-
+    reference_period : tuple of (int, int), optional
+        Year range (start_year, end_year) inclusive for computing the daily climatology
+        (fixed_baseline and detrend_fixed_baseline only). If None (default), uses all
+        available years. Anomalies are computed for the full time series regardless.
 
     Returns
     -------
@@ -997,6 +1032,15 @@ def compute_normalised_anomaly(
     ... )
     >>> # Climatology computed from all available years
 
+    Fixed baseline with a restricted reference period:
+
+    >>> # Compute climatology from 1990-2020 only, but output anomalies for all years
+    >>> result_ref = marEx.compute_normalised_anomaly(
+    ...     sst,
+    ...     method_anomaly="fixed_baseline",
+    ...     reference_period=(1990, 2020)
+    ... )
+
     Fixed detrended baseline:
 
     >>> # Remove long-term trends then compute fixed climatology
@@ -1021,6 +1065,17 @@ def compute_normalised_anomaly(
     # Infer and validate dimensions and coordinates
     dimensions, coordinates = _infer_dims_coords(da, dimensions, coordinates)
 
+    # Validate reference_period is only used with compatible methods
+    if reference_period is not None and method_anomaly not in ("fixed_baseline", "detrend_fixed_baseline"):
+        raise ConfigurationError(
+            f"reference_period is not supported for method_anomaly='{method_anomaly}'",
+            details="reference_period is only applicable to 'fixed_baseline' and 'detrend_fixed_baseline' methods",
+            suggestions=[
+                "Remove the reference_period parameter, or",
+                "Use method_anomaly='fixed_baseline' or 'detrend_fixed_baseline'",
+            ],
+        )
+
     if method_anomaly == "detrend_harmonic":
         logger.debug(
             f"Detrended baseline parameters: std_normalise={std_normalise}, orders={detrend_orders}, zero_mean={force_zero_mean}"
@@ -1032,11 +1087,16 @@ def compute_normalised_anomaly(
             da, window_year_baseline, smooth_days_baseline, dimensions, coordinates, use_temp_checkpoints
         )
     elif method_anomaly == "fixed_baseline":
-        logger.debug("Fixed baseline parameters: using full time series for daily climatology")
-        return _compute_anomaly_fixed_baseline(da, dimensions, coordinates)
+        logger.debug(f"Fixed baseline parameters: reference_period={reference_period}")
+        return _compute_anomaly_fixed_baseline(da, dimensions, coordinates, reference_period)
     elif method_anomaly == "detrend_fixed_baseline":
-        logger.debug(f"Fixed detrended baseline parameters: orders={detrend_orders}, zero_mean={force_zero_mean}")
-        return _compute_anomaly_detrend_fixed_baseline(da, detrend_orders, dimensions, coordinates, force_zero_mean)
+        logger.debug(
+            f"Fixed detrended baseline parameters: orders={detrend_orders}, "
+            f"zero_mean={force_zero_mean}, reference_period={reference_period}"
+        )
+        return _compute_anomaly_detrend_fixed_baseline(
+            da, detrend_orders, dimensions, coordinates, force_zero_mean, reference_period
+        )
     else:
         logger.error(f"Unknown anomaly method: {method_anomaly}")
         raise ConfigurationError(
@@ -1859,16 +1919,41 @@ def _identify_extremes_hobday(
 
     # Group by day-of-year and compute percentile
     if method_percentile == "exact":
-        # Construct rolling window dimension
-        da_windowed = da.rolling({dimensions["time"]: window_days_hobday}, center=True).construct("window")
-
-        # Ensure dayofyear coordinate is computed for groupby (required by newer xarray)
-        if "dayofyear" in da_windowed.coords:
-            da_windowed = da_windowed.assign_coords(dayofyear=da_windowed.dayofyear.compute())
-
-        thresholds = da_windowed.groupby("dayofyear").reduce(
-            np.nanpercentile, q=threshold_percentile, dim=("window", dimensions["time"])
+        # Use apply_ufunc to compute DOY percentiles per spatial chunk in pure numpy.
+        da_ufunc = da.chunk({dimensions["time"]: -1})
+        dayofyear_vals = da_ufunc[coordinates["time"]].dt.dayofyear.values
+        half_w = window_days_hobday // 2
+
+        # Pre-compute boolean masks (which time indices contribute to each DOY)
+        doy_masks = []
+        for doy in range(1, 367):
+            mask = np.zeros(len(dayofyear_vals), dtype=bool)
+            for offset in range(-half_w, half_w + 1):
+                target = ((doy - 1 + offset) % 366) + 1
+                mask |= dayofyear_vals == target
+            doy_masks.append(mask)
+
+        def _doy_percentiles(data, doy_masks, percentile):
+            """Compute per-DOY percentiles. data: (*spatial, time) -> (*spatial, 366)."""
+            result = np.full(data.shape[:-1] + (366,), np.nan, dtype=np.float32)
+            for i, mask in enumerate(doy_masks):
+                if mask.any():
+                    result[..., i] = np.nanpercentile(data[..., mask], percentile, axis=-1)
+            return result
+
+        thresholds = xr.apply_ufunc(
+            _doy_percentiles,
+            da_ufunc,
+            input_core_dims=[[dimensions["time"]]],
+            output_core_dims=[["dayofyear"]],
+            dask="parallelized",
+            kwargs={"doy_masks": doy_masks, "percentile": threshold_percentile},
+            output_dtypes=[np.float32],
+            dask_gufunc_kwargs={"output_sizes": {"dayofyear": 366}},
         )
+
+        # Assign dayofyear coordinate values and move dayofyear to first dimension
+        thresholds = thresholds.assign_coords(dayofyear=np.arange(1, 367)).transpose("dayofyear", ...)
     else:  # Optimised histogram approximation method
         thresholds = _compute_histogram_quantile_2d(
             da,
@@ -2179,6 +2264,7 @@ def _compute_anomaly_detrended(
             func="std",
             isbin=False,
             method="cohorts",
+            dtype=np.float32,
         )
 
         # Calculate 30-day rolling standard deviation with annual wrapped padding
@@ -2214,12 +2300,14 @@ def _compute_anomaly_fixed_baseline(
     da: xr.DataArray,
     dimensions: Optional[Dict[str, str]] = None,
     coordinates: Optional[Dict[str, str]] = None,
+    reference_period: Optional[Tuple[int, int]] = None,
 ) -> xr.Dataset:
     """
     Compute anomalies using fixed baseline method with full time series climatology.
 
-    This method computes a daily climatology using all available years in the dataset,
-    then subtracts this climatology from the original data to obtain anomalies.
+    This method computes a daily climatology using all available years in the dataset
+    (or a specified reference period), then subtracts this climatology from the
+    original data to obtain anomalies.
 
     Parameters
     ----------
@@ -2229,6 +2317,10 @@ def _compute_anomaly_fixed_baseline(
         Mapping of dimensions to names in the data
     coordinates : dict, optional
         Mapping of coordinates to names in the data
+    reference_period : tuple of (int, int), optional
+        Year range (start_year, end_year) inclusive for computing the daily climatology.
+        If None (default), uses all available years. Anomalies are still computed for
+        the full time series.
 
     Returns
     -------
@@ -2238,11 +2330,41 @@ def _compute_anomaly_fixed_baseline(
     # Infer and validate dimensions and coordinates
     dimensions, coordinates = _infer_dims_coords(da, dimensions, coordinates)
 
-    # Compute daily climatology across all years using flox for efficiency
-    logger.debug("Computing daily climatology across all years")
+    # Select data for climatology computation (optionally restricted to reference period)
+    if reference_period is not None:
+        start_year, end_year = reference_period
+        if start_year > end_year:
+            raise ConfigurationError(
+                f"Invalid reference_period: start year ({start_year}) must be <= end year ({end_year})",
+                details="The reference_period tuple must be (start_year, end_year) with start_year <= end_year",
+                suggestions=[f"Swap the order: use reference_period=({end_year}, {start_year})"],
+            )
+        years = da[coordinates["time"]].dt.year
+        year_mask = (years >= start_year) & (years <= end_year)
+        da_for_clim = da.isel({dimensions["time"]: year_mask})
+        if da_for_clim.sizes[dimensions["time"]] == 0:
+            data_min_year = int(years.min().values)
+            data_max_year = int(years.max().values)
+            raise ConfigurationError(
+                f"No data found in reference_period ({start_year}, {end_year})",
+                details=f"Dataset spans {data_min_year}-{data_max_year} but no timesteps fall within the specified period",
+                suggestions=[
+                    f"Adjust reference_period to overlap with data range ({data_min_year}-{data_max_year})",
+                    "Set reference_period=None to use the full time series",
+                ],
+            )
+        logger.debug(
+            f"Using reference_period ({start_year}-{end_year}): "
+            f"{da_for_clim.sizes[dimensions['time']]} of {da.sizes[dimensions['time']]} timesteps"
+        )
+    else:
+        da_for_clim = da
+
+    # Compute daily climatology using flox for efficiency
+    logger.debug("Computing daily climatology across %s", "reference period" if reference_period else "all years")
     daily_climatology = flox.xarray.xarray_reduce(
-        da,
-        da[coordinates["time"]].dt.dayofyear,
+        da_for_clim,
+        da_for_clim[coordinates["time"]].dt.dayofyear,
         dim=dimensions["time"],
         func="nanmean",
         isbin=False,
@@ -2281,13 +2403,14 @@ def _compute_anomaly_detrend_fixed_baseline(
     dimensions: Optional[Dict[str, str]] = None,
     coordinates: Optional[Dict[str, str]] = None,
     force_zero_mean: bool = True,
+    reference_period: Optional[Tuple[int, int]] = None,
 ) -> xr.Dataset:
     """
     Compute anomalies using fixed detrended baseline method.
 
     This method first removes polynomial trends (without harmonics) from the data,
-    then removes a full-time-series daily climatology from the detrended signal.
-    Uses _compute_anomaly_detrended internally to perform the trend removal.
+    then removes a daily climatology from the detrended signal. The trend removal
+    always uses the full time series; only the climatology step respects reference_period.
 
     Parameters
     ----------
@@ -2301,6 +2424,10 @@ def _compute_anomaly_detrend_fixed_baseline(
         Mapping of coordinates to names in the data
     force_zero_mean : bool, default=True
         Whether to enforce zero mean in detrended data
+    reference_period : tuple of (int, int), optional
+        Year range (start_year, end_year) inclusive for computing the daily climatology.
+        If None (default), uses all available years. Only affects the climatology step,
+        not the polynomial detrending.
 
     Returns
     -------
@@ -2329,6 +2456,7 @@ def _compute_anomaly_detrend_fixed_baseline(
         da=detrended_result,
         dimensions=dimensions,
         coordinates=coordinates,
+        reference_period=reference_period,
     )
 
     return final_result

{marex-4.0.4 → marex-4.1}/marEx.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: marEx
-Version: 4.0.4
+Version: 4.1
 Summary: Marine Extremes Detection and Tracking
 Author-email: Aaron Wienkers <aaron@wienkers.com>
 Project-URL: Homepage, https://github.com/wienkers/marEx
@@ -87,7 +87,7 @@ Requires-Dist: pytest-xdist; extra == "test"
 Requires-Dist: coverage[toml]>=6.0.0; extra == "test"
 Requires-Dist: psutil; extra == "test"
 
-![logo](https://github.com/user-attachments/assets/df99a84b-c88f-447d-a362-ac538d9b924b)
+<img src="https://github.com/user-attachments/assets/55eab465-9aca-4754-a981-e26a36f7d045" alt="marEx_animated" width="100%">
 
 [![CI](https://github.com/wienkers/marEx/actions/workflows/ci.yml/badge.svg)](https://github.com/wienkers/marEx/actions/workflows/ci.yml)
 [![codecov](https://codecov.io/gh/wienkers/marEx/branch/main/graph/badge.svg)](https://codecov.io/gh/wienkers/marEx)

{marex-4.0.4 → marex-4.1}/tests/test_gridded_preprocessing.py

@@ -1,9 +1,11 @@
 from pathlib import Path
 
 import numpy as np
+import pytest
 import xarray as xr
 
 import marEx
+from marEx.exceptions import ConfigurationError
 
 from .conftest import assert_percentile_frequency
 
@@ -767,3 +769,80 @@ class TestGriddedPreprocessing:
             assert (
                 0.025 < extreme_frequency < 0.075
             ), f"{method_anomaly}+{method_extreme} produced unreasonable frequency: {extreme_frequency}"
+
+    def test_fixed_baseline_with_reference_period(self):
+        """Test fixed_baseline reference_period restricts climatology to specified years."""
+        import pandas as pd
+
+        # Build a small synthetic dataset: 4 years of daily data at 3x3 grid
+        # Years 2000-2001 have a warm bias (+2), years 2002-2003 are baseline (0)
+        times = pd.date_range("2000-01-01", "2003-12-31", freq="D")
+        rng = np.random.default_rng(42)
+        base = rng.standard_normal((len(times), 3, 3)).astype(np.float32)
+
+        # Add a +2 warm bias to the first two years
+        warm_mask = np.array([t.year <= 2001 for t in times])
+        base[warm_mask] += 2.0
+
+        da = xr.DataArray(
+            base,
+            dims=["time", "lat", "lon"],
+            coords={"time": times, "lat": [0.0, 1.0, 2.0], "lon": [0.0, 1.0, 2.0]},
+        ).chunk({"time": -1})
+
+        dims = {"time": "time", "x": "lon", "y": "lat"}
+
+        # Full-range climatology: mean mixes warm+cool years → anomalies centred on mixed mean
+        result_full = marEx.compute_normalised_anomaly(da, method_anomaly="fixed_baseline", dimensions=dims)
+
+        # Reference period restricted to cool years 2002-2003:
+        # climatology is ~0, so warm years produce large positive anomalies
+        result_ref = marEx.compute_normalised_anomaly(
+            da, method_anomaly="fixed_baseline", reference_period=(2002, 2003), dimensions=dims
+        )
+
+        # Output should preserve all time steps
+        assert result_ref.sizes["time"] == da.sizes["time"]
+
+        # Anomalies during the warm period (2000-2001) should be larger when the
+        # climatology is computed from the cool period only
+        warm_slice = {"time": slice("2000-01-01", "2001-12-31")}
+        mean_anom_full = float(result_full.dat_anomaly.sel(**warm_slice).mean())
+        mean_anom_ref = float(result_ref.dat_anomaly.sel(**warm_slice).mean())
+        assert mean_anom_ref > mean_anom_full + 0.5, (
+            f"Reference-period anomalies during warm years ({mean_anom_ref:.2f}) should be "
+            f"substantially larger than full-range anomalies ({mean_anom_full:.2f})"
+        )
+
+    def test_reference_period_attr_stored_in_preprocess_data(self):
+        """Test that reference_period attribute is stored by preprocess_data."""
+        import pandas as pd
+
+        times = pd.date_range("2000-01-01", "2001-12-31", freq="D")
+        da = xr.DataArray(
+            np.random.default_rng(0).standard_normal((len(times), 3, 3)).astype(np.float32),
+            dims=["time", "lat", "lon"],
+            coords={"time": times, "lat": [0.0, 1.0, 2.0], "lon": [0.0, 1.0, 2.0]},
+        ).chunk({"time": -1})
+
+        result = marEx.preprocess_data(
+            da,
+            method_anomaly="fixed_baseline",
+            method_extreme="global_extreme",
+            threshold_percentile=95,
+            reference_period=(2000, 2000),
+            dimensions={"time": "time", "x": "lon", "y": "lat"},
+            dask_chunks={"time": -1},
+        )
+        assert result.attrs.get("reference_period") == [2000, 2000]
+
+    def test_reference_period_invalid_for_shifting_baseline(self):
+        """Test that reference_period raises ConfigurationError for incompatible methods."""
+        with pytest.raises(ConfigurationError, match="reference_period is not supported"):
+            marEx.preprocess_data(
+                self.sst_data,
+                method_anomaly="shifting_baseline",
+                reference_period=(1990, 2010),
+                dimensions=self.dimensions,
+                dask_chunks=self.dask_chunks,
+            )

{marex-4.0.4 → marex-4.1}/tests/test_track_edge_cases.py

@@ -20,14 +20,14 @@ from marEx.exceptions import ConfigurationError, DataValidationError
 
 
 @pytest.fixture(scope="module")
-def extremes_gridded():
+def extremes_gridded(dask_client):
     """Load gridded extremes data for testing."""
     test_data_path = Path(__file__).parent / "data" / "extremes_gridded.zarr"
     return xr.open_zarr(str(test_data_path), chunks={}).persist()
 
 
 @pytest.fixture(scope="module")
-def extremes_unstructured():
+def extremes_unstructured(dask_client):
     """Load unstructured extremes data for testing."""
     test_data_path = Path(__file__).parent / "data" / "extremes_unstructured.zarr"
     return xr.open_zarr(str(test_data_path), chunks={}).persist()
@@ -45,7 +45,7 @@ def temp_checkpoint_dir():
 class TestCheckpointFunctionality:
     """Test checkpoint save and load functionality."""
 
-    def test_checkpoint_save(self, extremes_unstructured, temp_checkpoint_dir, dask_client_unstructured):
+    def test_checkpoint_save(self, extremes_unstructured, temp_checkpoint_dir, dask_client):
         """Test saving checkpoints during preprocessing (unstructured grids only)."""
         neighbours = extremes_unstructured.neighbours
         cell_areas = extremes_unstructured.cell_areas
@@ -84,7 +84,7 @@ class TestCheckpointFunctionality:
         loaded_zarr = xr.open_zarr(str(checkpoint_zarr))
         assert "data_bin_preproc" in loaded_zarr, "Preprocessed data not in checkpoint"
 
-    def test_checkpoint_load(self, extremes_unstructured, temp_checkpoint_dir, dask_client_unstructured):
+    def test_checkpoint_load(self, extremes_unstructured, temp_checkpoint_dir, dask_client):
         """Test loading from existing checkpoints (unstructured grids only)."""
         neighbours = extremes_unstructured.neighbours
         cell_areas = extremes_unstructured.cell_areas
@@ -144,7 +144,7 @@ class TestCheckpointFunctionality:
 class TestChunkingValidationWarnings:
     """Test chunking validation and rechunking warnings."""
 
-    def test_neighbours_nv_chunking_warning(self, extremes_unstructured, temp_checkpoint_dir, dask_client_unstructured):
+    def test_neighbours_nv_chunking_warning(self, extremes_unstructured, temp_checkpoint_dir, dask_client):
         """Test warning when neighbours has multiple chunks in nv dimension."""
         neighbours = extremes_unstructured.neighbours
         cell_areas = extremes_unstructured.cell_areas
@@ -178,7 +178,7 @@ class TestChunkingValidationWarnings:
 class TestDataValidationEdgeCases:
     """Test edge cases in data validation."""
 
-    def test_empty_data_bin_attrs(self, extremes_gridded, dask_client_gridded):
+    def test_empty_data_bin_attrs(self, extremes_gridded, dask_client):
         """Test handling of DataArray with no attrs."""
         # Remove all attributes
         data_no_attrs = extremes_gridded.extreme_events.copy(deep=True)
@@ -195,7 +195,7 @@ class TestDataValidationEdgeCases:
 
         assert tracker.data_attrs == {}, "Should handle empty attrs gracefully"
 
-    def test_missing_coordinates_unstructured(self, dask_client_unstructured):
+    def test_missing_coordinates_unstructured(self, dask_client):
         """Test error when required coordinates are missing in unstructured data."""
         # Create data with missing coordinates
         test_data = xr.DataArray(
@@ -231,7 +231,7 @@ class TestDataValidationEdgeCases:
 class TestOverlapValidationAndWarnings:
     """Test overlap validation and warning conditions."""
 
-    def test_empty_overlap_list_return(self, extremes_gridded, dask_client_gridded):
+    def test_empty_overlap_list_return(self, extremes_gridded, dask_client):
         """Test handling of empty overlap lists."""
         # Create minimal tracker
         tracker = marEx.tracker(
@@ -254,7 +254,7 @@ class TestOverlapValidationAndWarnings:
 class TestComplexMergeSplitValidation:
     """Test complex merge/split validation and logging."""
 
-    def test_complex_merging_scenario(self, extremes_gridded, dask_client_gridded):
+    def test_complex_merging_scenario(self, extremes_gridded, dask_client):
         """Test complex merging with validation logging."""
         # Create tracker with aggressive merging settings
         tracker = marEx.tracker(
@@ -277,7 +277,7 @@ class TestComplexMergeSplitValidation:
         assert "allow_merging" in tracked_ds.attrs
         assert tracked_ds.attrs["allow_merging"] == 1
 
-    def test_merging_with_temporal_fill(self, extremes_gridded, dask_client_gridded):
+    def test_merging_with_temporal_fill(self, extremes_gridded, dask_client):
         """Test merging with temporal filling."""
         tracker = marEx.tracker(
             extremes_gridded.extreme_events,
@@ -301,10 +301,12 @@ class TestComplexMergeSplitValidation:
 class TestUnstructuredPartitioning:
     """Test unstructured grid partitioning code paths."""
 
-    def test_unstructured_tracking_with_merging(self, extremes_unstructured, temp_checkpoint_dir, dask_client_unstructured):
+    def test_unstructured_tracking_with_merging(self, extremes_unstructured, temp_checkpoint_dir, dask_client):
         """Test unstructured grid tracking with merging to trigger partition code."""
         neighbours = extremes_unstructured.neighbours
-        cell_areas = extremes_unstructured.cell_areas
+        # Compute cell_areas to numpy so it doesn't serialise as a distributed
+        # future into Dask workers
+        cell_areas = extremes_unstructured.cell_areas.compute()
 
         # Create tracker with merging enabled
         tracker = marEx.tracker(