marEx 4.0.2__tar.gz → 4.0.3__tar.gz

{marex-4.0.2 → marex-4.0.3}/.github/workflows/ci.yml

@@ -64,6 +64,7 @@ jobs:
     name: Test Suite
     runs-on: ${{ matrix.os }}
     needs: code-quality
+    timeout-minutes: 90  # Windows can be significantly slower than Linux/macOS
     strategy:
       fail-fast: false
       matrix:
@@ -144,18 +145,22 @@ jobs:
         pytest tests/test_logging_system.py -n 2 -v --tb=short
 
     - 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
 
     - name: Run preprocessing tests (unstructured)
+      timeout-minutes: 30  # Windows can be slow for these tests
       run: |
         pytest tests/test_unstructured_preprocessing.py -n 2 -v --tb=short
 
     - name: Run tracking tests (gridded)
+      timeout-minutes: 30  # Windows can be slow for these tests
       run: |
         pytest tests/test_gridded_tracking.py -n 2 -v --tb=short
 
     - name: Run tracking tests (unstructured)
+      timeout-minutes: 30  # Windows can be slow for these tests
       run: |
         pytest tests/test_unstructured_tracking.py -n 2 -v --tb=short
 

{marex-4.0.2 → marex-4.0.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: marEx
-Version: 4.0.2
+Version: 4.0.3
 Summary: Marine Extremes Detection and Tracking
 Author-email: Aaron Wienkers <aaron@wienkers.com>
 Project-URL: Homepage, https://github.com/wienkers/marEx

marex-4.0.3/examples/batch jobs/run_detect.py

@@ -0,0 +1,91 @@
+#!/usr/bin/env python
+"""
+Global Daily SST (OSTIA) Analysis: Identifying Marine Extremes with MarEx-Detect
+
+Pre-Process with Shifting-Baseline & Hobday et al. (2016) Definition for Extremes
+
+Run from login node - this script will launch SLURM jobs via Dask distributed cluster
+
+NOTE: This script uses a distributed cluster because detect.py does not contain
+Numba JIT-compiled functions that execute locally. Unlike run_track.py, this
+approach works correctly for detection/preprocessing tasks. If you prefer to
+follow HPC best practices and run everything in SLURM jobs, consider creating
+a submit_detect.sh script similar to submit_track.sh.
+"""
+
+import os
+from getpass import getuser
+from pathlib import Path
+
+import xarray as xr
+
+import marEx
+import marEx.helper as hpc
+
+
+def main():
+    """Run MarEx preprocessing and extreme detection pipeline on distributed SLURM cluster."""
+    # Configuration from environment variables
+    n_workers = int(os.getenv("DASK_N_WORKERS", "128"))
+    workers_per_node = int(os.getenv("DASK_WORKERS_PER_NODE", "64"))
+    runtime = int(os.getenv("DASK_RUNTIME", "39"))
+    slurm_account = os.getenv("SLURM_ACCOUNT", "bk1377")
+
+    # Lustre Scratch Directory
+    scratch_dir = Path("/scratch") / getuser()[0] / getuser()
+
+    # Start Distributed Dask Cluster on SLURM
+    print("Starting distributed Dask cluster on SLURM...")
+    print(f"  Workers: {n_workers}")
+    print(f"  Workers per node: {workers_per_node}")
+    print(f"  Runtime: {runtime} minutes")
+    print(f"  SLURM account: {slurm_account}")
+
+    client = hpc.start_distributed_cluster(
+        n_workers=n_workers,
+        workers_per_node=workers_per_node,
+        runtime=runtime,
+        scratch_dir=scratch_dir / "clients",
+        account=slurm_account,
+    )
+
+    # Choose optimal chunk size & load data
+    print("Loading data...")
+    file_name = scratch_dir / "mhws" / "ostia.zarr"
+    chunk_size = {"time": 30, "lat": 360, "lon": 720}
+    sst = xr.open_zarr(str(file_name), chunks=chunk_size).sst
+    print(f"Data loaded: {sst}")
+
+    # Process Data using MarEx-Detect helper functions:
+    print("Processing extremes data...")
+    extremes_ds = marEx.preprocess_data(
+        sst,
+        method_anomaly="shifting_baseline",  # Anomalies from a rolling climatology using previous window_year years
+        method_extreme="hobday_extreme",  # Local day-of-year specific thresholds with windowing
+        threshold_percentile=95,  # Use the 95th percentile as the extremes threshold
+        window_year_baseline=15,
+        smooth_days_baseline=21,  # Defines the rolling climatology window (15 years) and smoothing window (21 days)
+        window_days_hobday=11,  # Defines the window (11 days) of compiled samples collected for extremes detection
+        dimensions={
+            "time": "time",
+            "x": "lon",
+            "y": "lat",
+        },  # Define the dimensions of the data -- if 'y' exists, MarEx-Detect knows this is a gridded dataset
+        dask_chunks={"time": 25},  # Dask chunks for *output* data
+        use_temp_checkpoints=True,  # Enable checkpointing to prevent expensive recomputations
+    )
+    print(f"Processing complete: {extremes_ds}")
+
+    # Save Extremes Data to zarr for more efficient parallel I/O
+    print("Saving extremes data...")
+    output_file = scratch_dir / "mhws" / "extremes_binary_gridded_shifting_hobday_batch.zarr"
+    extremes_ds.to_zarr(output_file, mode="w")
+    print(f"Data saved to: {output_file}")
+
+    # Close the Dask cluster
+    client.close()
+    print("Processing complete!")
+
+
+if __name__ == "__main__":
+    main()

marex-4.0.3/examples/batch jobs/run_track.py

@@ -0,0 +1,131 @@
+#!/usr/bin/env python
+"""
+Global Daily Event Analysis: Marine Heatwave ID & Tracking using MarEx
+
+Performs morphological preprocessing, blob identification, temporal tracking,
+and graph reduction to create globally unique IDs for each tracked extreme event.
+
+This script should be run INSIDE a SLURM job (via sbatch submit_track.sh).
+It uses a local Dask cluster that utilizes the compute resources allocated to the job.
+This ensures that all computation, including Numba JIT-compiled functions, runs on
+compute nodes rather than the login node.
+"""
+
+import os
+from getpass import getuser
+from pathlib import Path
+
+import xarray as xr
+
+import marEx
+import marEx.helper as hpc
+
+
+def main():
+    """Run MarEx event tracking pipeline on local Dask cluster within SLURM job allocation."""
+    # Configuration from environment variables
+    n_workers = int(os.getenv("DASK_N_WORKERS", "32"))
+    threads_per_worker = int(os.getenv("DASK_THREADS_PER_WORKER", "1"))
+
+    # Tracking configuration
+    run_basic = os.getenv("RUN_BASIC_TRACKER", "false").lower() == "true"
+    grid_resolution = float(os.getenv("GRID_RESOLUTION", "0.25"))
+    area_filter = int(os.getenv("AREA_FILTER", "600"))
+    r_fill = int(os.getenv("R_FILL", "12"))
+    t_fill = int(os.getenv("T_FILL", "4"))
+    overlap_threshold = float(os.getenv("OVERLAP_THRESHOLD", "0.25"))
+
+    # Lustre Scratch Directory
+    scratch_dir = Path("/scratch") / getuser()[0] / getuser()
+
+    # Start Local Dask Cluster
+    # This cluster runs within the SLURM job allocation, ensuring all computation
+    # (including Numba JIT-compiled functions) executes on compute nodes
+    print("Starting local Dask cluster within SLURM job allocation...")
+    print(f"  Workers: {n_workers}")
+    print(f"  Threads per worker: {threads_per_worker}")
+
+    client = hpc.start_local_cluster(
+        n_workers=n_workers,
+        threads_per_worker=threads_per_worker,
+        scratch_dir=scratch_dir / "clients",
+    )
+
+    # Choose optimal chunk size & load preprocessed data
+    print("\nLoading preprocessed extremes data...")
+    input_file = scratch_dir / "mhws" / "extremes_binary_gridded_shifting_hobday_batch.zarr"
+    chunk_size = {"time": 25, "lat": -1, "lon": -1}
+    ds = xr.open_zarr(str(input_file), chunks=chunk_size)
+    print(f"Data loaded: {ds}")
+
+    # Run ID, Tracking, & Merging
+    print("\n" + "=" * 80)
+    print("Running Full Tracking with Merge/Split Detection...")
+    print("=" * 80)
+
+    tracker = marEx.tracker(
+        ds.extreme_events,
+        ds.mask.where(
+            (ds.lat < 85) & (ds.lat > -90), other=False
+        ),  # Modify Mask: Anisotropy of lat/lon grid near poles biases ID & Tracking
+        grid_resolution=grid_resolution,  # Grid resolution in degrees, used to calculate object areas on the globe
+        area_filter_absolute=area_filter,  # Remove objects smaller than specified cells
+        R_fill=r_fill,  # Fill small holes with radius < R_fill _cells_
+        T_fill=t_fill,  # Allow gaps of T_fill days and still continue event tracking with same ID
+        allow_merging=True,  # Allow extreme events to split/merge. Keeps track of merge events & unique IDs
+        overlap_threshold=overlap_threshold,  # Overlap threshold for merging events
+        nn_partitioning=True,  # Use new NN method to partition merged children areas
+        verbose=True,
+    )
+
+    extreme_events_ds, merges_ds = tracker.run(return_merges=True)
+    print(f"\nTracking complete: {extreme_events_ds}")
+    print(f"Merges dataset: {merges_ds}")
+
+    # Save IDed/Tracked/Merged Events to zarr for efficient parallel I/O
+    print("\nSaving tracked events data...")
+    output_file = scratch_dir / "mhws" / "extreme_events_merged_gridded_shifting_batch.zarr"
+    extreme_events_ds.to_zarr(output_file, mode="w")
+    print(f"Events data saved to: {output_file}")
+
+    # Save Merges Dataset to netcdf
+    print("Saving merges data...")
+    merges_file = scratch_dir / "mhws" / "extreme_events_merged_gridded_shifting_merges_batch.nc"
+    merges_ds.to_netcdf(merges_file, mode="w")
+    print(f"Merges data saved to: {merges_file}")
+
+    # Optionally run basic tracking for comparison
+    if run_basic:
+        print("\n" + "=" * 80)
+        print("Running Basic Tracking (No Merge/Split, No Temporal Gap Filling)...")
+        print("=" * 80)
+
+        tracker_basic = marEx.tracker(
+            ds.extreme_events,
+            ds.mask.where(
+                (ds.lat < 85) & (ds.lat > -90), other=False
+            ),  # Modify Mask: Anisotropy of lat/lon grid near poles biases ID & Tracking
+            area_filter_absolute=area_filter,  # Remove objects smaller than specified cells
+            R_fill=r_fill,  # Fill small holes with radius < R_fill _cells_
+            T_fill=0,  # No temporal hole filling
+            allow_merging=False,  # Do not allow extreme events to split/merge
+        )
+
+        extreme_events_basic_ds = tracker_basic.run()
+        print(f"\nBasic tracking complete: {extreme_events_basic_ds}")
+
+        # Save Basic IDed Events to zarr
+        print("\nSaving basic tracked events data...")
+        basic_output_file = scratch_dir / "mhws" / "extreme_events_basic_gridded_shifting_batch.zarr"
+        extreme_events_basic_ds.to_zarr(basic_output_file, mode="w")
+        print(f"Basic events data saved to: {basic_output_file}")
+
+    # Close the Dask cluster
+    client.close()
+    print("\n" + "=" * 80)
+    print("Tracking complete!")
+    print("=" * 80)
+
+
+if __name__ == "__main__":
+    main()

marex-4.0.3/examples/batch jobs/submit_track.sh

@@ -0,0 +1,79 @@
+#!/bin/bash
+#SBATCH --job-name=marEx_track
+#SBATCH --partition=compute
+#SBATCH --cpus-per-task=128
+#SBATCH --account=bk1377
+#SBATCH --output=slurm_%j.out
+#SBATCH --time=02:59:00
+#SBATCH --mem=0
+#SBATCH --exclusive
+
+# SLURM Batch Script for MarEx Event Tracking
+#
+# This script submits run_track.py as a SLURM job, ensuring all computation
+# (including Numba JIT-compiled functions) runs on allocated compute nodes.
+#
+# Usage:
+#   sbatch submit_track.sh
+#
+# Configuration via environment variables (optional):
+#   export DASK_N_WORKERS=32         # Number of Dask workers
+#   export DASK_THREADS_PER_WORKER=1 # Threads per worker
+#   export RUN_BASIC_TRACKER=true    # Run basic tracker comparison
+#   export GRID_RESOLUTION=0.25      # Grid resolution in degrees
+#   export AREA_FILTER=600           # Minimum object area (cells)
+#   export R_FILL=12                 # Spatial hole filling radius
+#   export T_FILL=4                  # Temporal gap filling (days)
+#   export OVERLAP_THRESHOLD=0.25    # Merge overlap threshold
+#   sbatch submit_track.sh
+
+# Print job information
+echo "========================================="
+echo "SLURM Job Information"
+echo "========================================="
+echo "Job ID: $SLURM_JOB_ID"
+echo "Job Name: $SLURM_JOB_NAME"
+echo "Node: $SLURM_NODELIST"
+echo "Start Time: $(date)"
+echo "Working Directory: $(pwd)"
+echo "========================================="
+echo ""
+
+# Set ulimits for stack size
+ulimit -S -s unlimited
+ulimit -S -c unlimited
+
+# Initialise conda for bash shell
+source $(conda info --base)/etc/profile.d/conda.sh
+conda activate super
+
+# Set default environment variables if not already set
+export DASK_N_WORKERS=${DASK_N_WORKERS:-32}
+export DASK_THREADS_PER_WORKER=${DASK_THREADS_PER_WORKER:-1}
+
+# Print configuration
+echo "Configuration:"
+echo "  DASK_N_WORKERS: $DASK_N_WORKERS"
+echo "  DASK_THREADS_PER_WORKER: $DASK_THREADS_PER_WORKER"
+echo "  RUN_BASIC_TRACKER: ${RUN_BASIC_TRACKER:-false}"
+echo "  GRID_RESOLUTION: ${GRID_RESOLUTION:-0.25}"
+echo "  AREA_FILTER: ${AREA_FILTER:-600}"
+echo "  R_FILL: ${R_FILL:-12}"
+echo "  T_FILL: ${T_FILL:-4}"
+echo "  OVERLAP_THRESHOLD: ${OVERLAP_THRESHOLD:-0.25}"
+echo ""
+
+# Run the tracking script
+echo "Starting MarEx tracking..."
+python run_track.py
+
+# Capture exit status
+EXIT_STATUS=$?
+
+echo ""
+echo "========================================="
+echo "Job completed at: $(date)"
+echo "Exit status: $EXIT_STATUS"
+echo "========================================="
+
+exit $EXIT_STATUS

{marex-4.0.2 → marex-4.0.3}/marEx/__init__.py

@@ -57,6 +57,9 @@ from .exceptions import (  # Main exception hierarchy; Convenience constructors
     wrap_exception,
 )
 
+# Import HPC helper utilities
+from .helper import configure_dask
+
 # Import logging configuration functions
 from .logging_config import (
     configure_logging,
@@ -116,6 +119,8 @@ __all__ = [
     "is_verbose_mode",
     "is_quiet_mode",
     "get_logger",
+    # HPC helper utilities
+    "configure_dask",
 ]
 
 # Version information

{marex-4.0.2 → marex-4.0.3}/marEx/detect.py

@@ -27,6 +27,7 @@ import flox.xarray
 import numpy as np
 import pandas as pd
 import xarray as xr
+from dask import persist
 from dask.base import is_dask_collection
 from numpy.lib.stride_tricks import sliding_window_view
 from numpy.typing import NDArray
@@ -34,7 +35,7 @@ from xhistogram.xarray import histogram
 
 # Coordinate validation imports removed
 from .exceptions import ConfigurationError, create_data_validation_error
-from .helper import fix_dask_tuple_array
+from .helper import checkpoint_to_zarr, fix_dask_tuple_array
 from .logging_config import configure_logging, get_logger, log_dask_info, log_memory_usage, log_timing
 
 # Get module logger
@@ -305,6 +306,7 @@ def preprocess_data(
     coordinates: Optional[Dict[str, str]] = None,
     neighbours: Optional[xr.DataArray] = None,
     cell_areas: Optional[xr.DataArray] = None,
+    use_temp_checkpoints: bool = False,
     verbose: Optional[bool] = None,
     quiet: Optional[bool] = None,
 ) -> xr.Dataset:
@@ -370,6 +372,12 @@ def preprocess_data(
         Neighbour connectivity for spatial clustering.
     cell_areas : xarray.DataArray, optional
         Cell areas for weighted spatial statistics.
+    use_temp_checkpoints : bool, default=False
+        Enable checkpointing to temporary zarr stores to break Dask graph dependencies.
+        When True, intermediate results (anomalies, thresholds, extremes) are saved to
+        temporary zarr files and immediately reloaded, preventing expensive recomputations.
+        Recommended for large datasets on HPC systems where the 2D histogram computation
+        is expensive. Temporary files are automatically cleaned up after reloading.
     verbose : bool, default=None
         Enable verbose logging with detailed progress information.
         If None, uses current global logging configuration.
@@ -580,6 +588,7 @@ def preprocess_data(
             std_normalise,
             detrend_orders,
             force_zero_mean,
+            use_temp_checkpoints,
         )
         log_memory_usage(logger, "After anomaly computation", logging.DEBUG)
 
@@ -612,6 +621,11 @@ def preprocess_data(
         time_sel = (ds[coordinates["time"]].dt.year >= start_year).compute()
         ds = ds.isel({dimensions["time"]: time_sel})
 
+    # Break graph after expensive anomaly computation
+    if use_temp_checkpoints:
+        logger.debug("Checkpointing anomaly dataset to break graph dependencies")
+        ds = checkpoint_to_zarr(ds, name="anomalies", timedim=dimensions["time"])
+
     anomalies = ds.dat_anomaly
 
     # Step 2: Identify extreme events (both methods now return consistent tuple structures)
@@ -636,10 +650,14 @@ def preprocess_data(
             method_percentile,
             precision,
             max_anomaly,
+            use_temp_checkpoints,
         )
         log_memory_usage(logger, "After extreme identification", logging.DEBUG)
 
     # Add extreme events and thresholds to dataset
+    ds_temp = persist(extremes, thresholds)
+    extremes, thresholds = ds_temp
+
     ds["extreme_events"] = extremes
     ds["thresholds"] = thresholds
 
@@ -663,7 +681,15 @@ def preprocess_data(
                 method_percentile,
                 precision,
                 max_anomaly,
+                use_temp_checkpoints,
             )
+
+            # Break graph after standardised extremes computation
+            if use_temp_checkpoints:
+                logger.debug("Checkpointing standardised extremes and thresholds to break graph dependencies")
+                extremes_stn = checkpoint_to_zarr(extremes_stn, name="extremes_stn", timedim=dimensions["time"])
+                thresholds_stn = checkpoint_to_zarr(thresholds_stn, name="thresholds_stn", timedim="dayofyear")
+
             ds["extreme_events_stn"] = extremes_stn
             ds["thresholds_stn"] = thresholds_stn
 
@@ -739,6 +765,14 @@ def preprocess_data(
         chunk_dict["dayofyear"] = time_chunks
     ds = ds.chunk(chunk_dict)
 
+    # Clear encoding metadata that may conflict with actual Dask chunks
+    # This encoding carries over from checkpointing and can cause chunk misalignment errors
+    logger.debug("Clearing encoding metadata for Dask-backed variables")
+    for var in ds.data_vars:
+        if hasattr(ds[var].data, "chunks"):  # Only for Dask-backed variables
+            if hasattr(ds[var], "encoding") and "chunks" in ds[var].encoding:
+                del ds[var].encoding["chunks"]
+
     # Fix encoding issue with saving when calendar & units attribute is present
     if "calendar" in ds[coordinates["time"]].attrs:  # pragma: no cover
         logger.debug("Removing calendar attribute for Zarr compatibility")
@@ -755,7 +789,6 @@ def preprocess_data(
         show_progress=True,
     ):
         ds = ds.persist(optimize_graph=True)
-        ds["thresholds"] = ds.thresholds.compute()  # Patch for a dask-Zarr bug that has problems saving this data array...
         ds["mask"] = ds.mask.compute()
         ds["dat_anomaly"] = fix_dask_tuple_array(ds.dat_anomaly)
 
@@ -833,6 +866,7 @@ def compute_normalised_anomaly(
     std_normalise: bool = False,  # for detrend_harmonic
     detrend_orders: Optional[List[int]] = None,  # "
     force_zero_mean: bool = True,  # "
+    use_temp_checkpoints: bool = False,
     verbose: Optional[bool] = None,
     quiet: Optional[bool] = None,
 ) -> xr.Dataset:
@@ -991,7 +1025,9 @@ def compute_normalised_anomaly(
         return _compute_anomaly_detrended(da, std_normalise, detrend_orders, dimensions, coordinates, force_zero_mean)
     elif method_anomaly == "shifting_baseline":
         logger.debug(f"Shifting baseline parameters: window_years={window_year_baseline}, smooth_days={smooth_days_baseline}")
-        return _compute_anomaly_shifting_baseline(da, window_year_baseline, smooth_days_baseline, dimensions, coordinates)
+        return _compute_anomaly_shifting_baseline(
+            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)
@@ -1028,6 +1064,7 @@ def identify_extremes(
     method_percentile: Literal["exact", "approximate"] = "approximate",
     precision: float = 0.01,
     max_anomaly: float = 5.0,
+    use_temp_checkpoints: bool = False,
     verbose: Optional[bool] = None,
     quiet: Optional[bool] = None,
 ) -> Tuple[xr.DataArray, xr.DataArray]:
@@ -1385,6 +1422,7 @@ def identify_extremes(
             coordinates,
             precision,
             max_anomaly,
+            use_temp_checkpoints,
         )
     else:
         logger.error(f"Unknown extreme method: {method_extreme}")
@@ -1412,6 +1450,7 @@ def rolling_climatology(
     window_year_baseline: int = 15,
     dimensions: Optional[Dict[str, str]] = None,
     coordinates: Optional[Dict[str, str]] = None,
+    use_temp_checkpoints: bool = False,
 ) -> xr.DataArray:
     """
     Compute rolling climatology efficiently using flox cohorts.
@@ -1505,8 +1544,9 @@ def rolling_climatology(
     da = da.assign_coords({"year": years, "dayofyear": doys})
 
     # Get temporal bounds
-    year_vals = years.compute().values
-    doy_vals = doys.compute().values
+    years, doys = persist(years, doys)
+    year_vals = years.values
+    doy_vals = doys.values
     unique_years = np.unique(year_vals)
     min_year = int(unique_years.min().item())
 
@@ -1563,7 +1603,11 @@ def rolling_climatology(
         isbin=(False, False),
         dtype=np.float32,
         fill_value=np.nan,
-    ).persist()
+    ).chunk({"dayofyear": -1})
+
+    if use_temp_checkpoints:
+        logger.debug("Checkpointing climatologies to break graph dependencies")
+        climatologies = checkpoint_to_zarr(climatologies, name="climatologies", timedim=timedim)
 
     # Create index arrays for final mapping
     year_to_idx = pd.Series(range(len(unique_years)), index=unique_years)
@@ -1587,6 +1631,7 @@ def smoothed_rolling_climatology(
     smooth_days_baseline: int = 21,
     dimensions: Optional[Dict[str, str]] = None,
     coordinates: Optional[Dict[str, str]] = None,
+    use_temp_checkpoints: bool = False,
 ) -> xr.DataArray:
     """
     Compute a smoothed rolling climatology using the previous `window_year_baseline` years of data
@@ -1699,9 +1744,11 @@ def smoothed_rolling_climatology(
     timedim = dimensions["time"]
 
     # N.B.: It is more efficient (chunking-wise) to smooth the raw data rather than the climatology
-    da_smoothed = da.rolling({timedim: smooth_days_baseline}, center=True).mean().chunk(dict(zip(da.dims, da.chunks)))
+    da_smoothed = (
+        da.rolling({timedim: smooth_days_baseline}, center=True).mean().chunk(dict(zip(da.dims, da.chunks))).astype(np.float32)
+    )
 
-    clim = rolling_climatology(da_smoothed, window_year_baseline, dimensions, coordinates)
+    clim = rolling_climatology(da_smoothed, window_year_baseline, dimensions, coordinates, use_temp_checkpoints)
 
     return clim
 
@@ -1712,6 +1759,7 @@ def _compute_anomaly_shifting_baseline(
     smooth_days_baseline: int = 21,
     dimensions: Optional[Dict[str, str]] = None,
     coordinates: Optional[Dict[str, str]] = None,
+    use_temp_checkpoints: bool = False,
 ) -> xr.Dataset:
     """
     Compute anomalies using shifting baseline method with smoothed rolling climatology.
@@ -1725,14 +1773,15 @@ def _compute_anomaly_shifting_baseline(
     dimensions, coordinates = _infer_dims_coords(da, dimensions, coordinates)
 
     # Compute smoothed rolling climatology
-    climatology_smoothed = smoothed_rolling_climatology(da, window_year_baseline, smooth_days_baseline, dimensions, coordinates)
+    climatology_smoothed = smoothed_rolling_climatology(
+        da, window_year_baseline, smooth_days_baseline, dimensions, coordinates, use_temp_checkpoints
+    )
 
     # Compute anomaly as difference from climatology
     anomalies = da - climatology_smoothed
 
     # Create ocean/land mask from first time step
-    chunk_dict_mask = {dimensions[dim]: -1 for dim in ["x", "y"] if dim in dimensions}
-    mask = np.isfinite(da.isel({dimensions["time"]: 0})).drop_vars({coordinates["time"]}).chunk(chunk_dict_mask)
+    mask = np.isfinite(da.isel({dimensions["time"]: 0})).drop_vars({coordinates["time"]})
 
     # Build output dataset
     return xr.Dataset({"dat_anomaly": anomalies, "mask": mask})
@@ -1753,6 +1802,7 @@ def _identify_extremes_hobday(
     coordinates: Optional[Dict[str, str]] = None,
     precision: float = 0.01,
     max_anomaly: float = 5.0,
+    use_temp_checkpoints: bool = False,
 ) -> Tuple[xr.DataArray, xr.DataArray]:
     """
     Identify extreme events using day-of-year (i.e. climatological percentile threshold).
@@ -1801,14 +1851,18 @@ def _identify_extremes_hobday(
             "If your time-series is very short, consider using method_percentile='exact'."
         )
 
-    # Add day-of-year coordinate
-    da = da.assign_coords(dayofyear=da[coordinates["time"]].dt.dayofyear)
+    # Add day-of-year coordinate (compute it to avoid chunked groupby issues)
+    da = da.assign_coords(dayofyear=da[coordinates["time"]].dt.dayofyear.compute()).chunk(dict(zip(da.dims, da.chunks))).persist()
 
     # 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"])
         )
@@ -1821,16 +1875,62 @@ def _identify_extremes_hobday(
             dimensions=dimensions,
             precision=precision,
             max_anomaly=max_anomaly,
+            use_temp_checkpoints=use_temp_checkpoints,
         )
 
-    # Ensure spatial dimensions are fully loaded for efficient comparison
-    spatial_chunks = {dimensions[dim]: -1 for dim in ["x", "y"] if dim in dimensions}
-    thresholds = thresholds.chunk(spatial_chunks).persist()
+    if use_temp_checkpoints:
+        logger.debug("Checkpointing thresholds to break graph dependencies")
+        thresholds = checkpoint_to_zarr(thresholds, name="thresholds", timedim="dayofyear")
+
+    # Extract spatial chunk sizes from input data for alignment
+    # Use most common chunk size to handle irregular chunks robustly
+    spatial_chunks = {}
+    for dim_key in ["x", "y"]:
+        if dim_key in dimensions:
+            dim_name = dimensions[dim_key]
+            if dim_name in da.dims:
+                dim_index = da.dims.index(dim_name)
+                chunks_tuple = da.chunks[dim_index]
+                # Get the most common chunk size (handles irregular chunks better)
+                spatial_chunks[dim_name] = max(set(chunks_tuple), key=chunks_tuple.count)
+
+    # Drop time coordinate/dimension to avoid conflicts when comparing with data grouped by dayofyear
+    coords_to_drop = []
+    if coordinates["time"] in thresholds.coords:
+        coords_to_drop.append(coordinates["time"])
+    if dimensions["time"] in thresholds.coords and dimensions["time"] not in thresholds.dims:
+        coords_to_drop.append(dimensions["time"])
+    if "time" in thresholds.coords and "time" not in thresholds.dims:
+        coords_to_drop.append("time")
+    if coords_to_drop:
+        thresholds = thresholds.drop_vars(coords_to_drop)
+
+    # Rechunk thresholds BEFORE comparison to align with input data
+    # This eliminates expensive implicit rechunking during the groupby operation
+    logger.debug(f"Aligning threshold chunks to match input data spatial chunks: {spatial_chunks}")
+    thresholds = thresholds.chunk(spatial_chunks)
 
     # Compare anomalies to day-of-year specific thresholds
-    dayofyear_labels = da[coordinates["time"]].dt.dayofyear.compute()
-    extremes = da.groupby(dayofyear_labels) >= thresholds
-    extremes = extremes.astype(bool).chunk(spatial_chunks).persist()
+    # Assign dayofyear coordinate and use UniqueGrouper for chunked arrays
+    da = da.assign_coords(dayofyear=da[coordinates["time"]].dt.dayofyear)
+    extremes = da.groupby(dayofyear=xr.groupers.UniqueGrouper(labels=np.arange(1, 367))) >= thresholds
+
+    # Drop unnecessary dayofyear coordinate
+    if "dayofyear" in extremes.coords:
+        extremes = extremes.drop_vars("dayofyear")
+
+    # Rechunk to fix irregular time chunks created by groupby operation
+    # Zarr requires uniform chunks, so we rechunk to match input data's time chunks
+    time_dim_index = da.dims.index(dimensions["time"])
+    time_chunk_size = max(set(da.chunks[time_dim_index]), key=da.chunks[time_dim_index].count)
+    rechunk_dict = {dimensions["time"]: time_chunk_size}
+    rechunk_dict.update(spatial_chunks)
+    logger.debug(f"Rechunking extremes to fix irregular chunks from groupby: {rechunk_dict}")
+    extremes = extremes.chunk(rechunk_dict)
+
+    if use_temp_checkpoints:
+        logger.debug("Checkpointing extremes to break graph dependencies")
+        extremes = checkpoint_to_zarr(extremes, name="extremes", timedim=dimensions["time"])
 
     return extremes, thresholds
 
@@ -2085,7 +2185,12 @@ def _compute_anomaly_detrended(
         # Divide anomalies by rolling standard deviation
         # Replace any zeros or extremely small values with NaN to avoid division warnings
         std_rolling_safe = std_rolling.where(std_rolling > 1e-10, np.nan)
-        da_stn = da_detrend.groupby(da_detrend[coordinates["time"]].dt.dayofyear) / std_rolling_safe
+        da_detrend = da_detrend.assign_coords(dayofyear=da_detrend[coordinates["time"]].dt.dayofyear)
+        da_stn = da_detrend.groupby(dayofyear=xr.groupers.UniqueGrouper(labels=np.arange(1, 367))) / std_rolling_safe
+
+        # Drop dayofyear coordinate to avoid merge conflicts
+        if "dayofyear" in da_stn.coords:
+            da_stn = da_stn.drop_vars("dayofyear")
 
         # Rechunk data for efficient processing
         chunk_dict_std = chunk_dict_mask.copy()
@@ -2144,9 +2249,14 @@ def _compute_anomaly_fixed_baseline(
 
     # Compute anomalies by subtracting daily climatology from original data
     logger.debug("Computing anomalies by subtracting daily climatology")
-    anomalies = da.groupby(da[coordinates["time"]].dt.dayofyear) - daily_climatology
+    da = da.assign_coords(dayofyear=da[coordinates["time"]].dt.dayofyear)
+    anomalies = da.groupby(dayofyear=xr.groupers.UniqueGrouper(labels=np.arange(1, 367))) - daily_climatology
     anomalies = anomalies.astype(np.float32)
 
+    # Drop dayofyear coordinate to avoid merge conflicts
+    if "dayofyear" in anomalies.coords:
+        anomalies = anomalies.drop_vars("dayofyear")
+
     # Create ocean/land mask from first time step
     # Handle both spatial (3D) and time-series (1D) data
     spatial_dims = [dim for dim in ["x", "y"] if dim in dimensions]
@@ -2327,6 +2437,7 @@ def _compute_histogram_quantile_2d(
     dimensions: Optional[Dict[str, str]] = None,
     precision: float = 0.01,
     max_anomaly: float = 5.0,
+    use_temp_checkpoints: bool = False,
 ) -> xr.DataArray:
     """
     Efficiently compute quantiles using binned histograms optimised for extreme values.
@@ -2366,23 +2477,31 @@ def _compute_histogram_quantile_2d(
     bin_centers_array[0] = 0.0
 
     bin_centers = xr.DataArray(
-        bin_centers_array,
+        bin_centers_array.astype(np.float32),
         dims=["da_bin"],
-        coords={"da_bin": np.arange(len(bin_centers_array), dtype=np.int32)},
+        coords={"da_bin": np.arange(len(bin_centers_array), dtype=np.uint16)},
         name="bin_centers",
     )
 
     chunk_dict = {dimensions["time"]: -1}
-    for d in ["x", "y"]:
-        if d in dimensions:
-            chunk_dict[dimensions[d]] = 8
+    chunk_dict[dimensions["x"]] = 16
+    if "y" in dimensions:
+        chunk_dict[dimensions["y"]] = 16
+
+    da_bin = (
+        xr.DataArray(
+            np.digitize(da.data, bin_edges) - 1,  # -1 so first bin is 0
+            dims=da.dims,
+            coords=da.coords,
+            name="da_bin",
+        )
+        .chunk(chunk_dict)
+        .astype(np.uint16)
+    )
 
-    da_bin = xr.DataArray(
-        np.digitize(da.data, bin_edges) - 1,  # -1 so first bin is 0
-        dims=da.dims,
-        coords=da.coords,
-        name="da_bin",
-    ).chunk(chunk_dict)
+    if use_temp_checkpoints:
+        logger.debug("Checkpointing binned data to break graph dependencies")
+        da_bin = checkpoint_to_zarr(da_bin, name="da_bin", timedim=dimensions["time"]).chunk(chunk_dict)
 
     # Construct 2D histogram using flox (in doy & anomaly)
     hist_raw = flox.xarray.xarray_reduce(
@@ -2391,9 +2510,9 @@ def _compute_histogram_quantile_2d(
         da_bin,
         dim=[dimensions["time"]],
         func="count",
-        expected_groups=(np.arange(1, 367, dtype=np.int32), np.arange(len(bin_edges) - 1, dtype=np.int32)),
+        expected_groups=(np.arange(1, 367, dtype=np.uint16), np.arange(len(bin_edges) - 1, dtype=np.uint16)),
         isbin=(False, False),
-        dtype=np.int32,
+        dtype=np.uint16,
         fill_value=0,
     )
     hist_raw.name = None
@@ -2421,7 +2540,12 @@ def _compute_histogram_quantile_2d(
         return _rolling_histogram_quantile(hist_chunk, window_days_hobday, q, bin_centers_chunk)
 
     # Rechunk histogram so core dimensions are unchunked for apply_ufunc
-    hist_raw = hist_raw.chunk({"dayofyear": -1, "da_bin": -1})
+    # Create chunk dict for hist_raw that preserves spatial chunks but drops time
+    hist_chunk_dict = {dimensions["x"]: chunk_dict.get(dimensions["x"], 16), "dayofyear": -1, "da_bin": -1}
+    if "y" in dimensions:
+        hist_chunk_dict[dimensions["y"]] = chunk_dict.get(dimensions["y"], 16)
+
+    hist_raw = hist_raw.chunk(hist_chunk_dict)
 
     # Apply the optimised computation using apply_ufunc
     threshold = xr.apply_ufunc(
@@ -2437,8 +2561,16 @@ def _compute_histogram_quantile_2d(
         keep_attrs=True,
     )
 
+    if use_temp_checkpoints:
+        logger.debug("Checkpointing threshold to break graph dependencies")
+        threshold = checkpoint_to_zarr(threshold, name="threshold")
+
+    # Drop time coordinate to avoid conflicts when comparing with data grouped by dayofyear
+    if dimensions["time"] in threshold.coords:
+        threshold = threshold.drop_vars(dimensions["time"])
+
     # Set threshold to NaN for spatial points that contain NaN values
-    nan_mask = da.isnull().any(dim=dimensions["time"])
+    nan_mask = da.isel({dimensions["time"]: 0}).isnull().compute()
     threshold = threshold.where(~nan_mask).persist()
 
     # Validate threshold values against bounds
@@ -2446,7 +2578,7 @@ def _compute_histogram_quantile_2d(
     lower_bound = bin_edges[3]  # We want this to be positive so that constant=0 anomalies will not be "extreme"
 
     # Check if any values are too high (ignore NaN values)
-    too_high = (threshold > upper_bound) & threshold.notnull()
+    too_high = threshold > upper_bound
     if too_high.any():
         warnings.warn(
             f"Quantile values exceed expected range: max={threshold.max().compute():.4f} > {upper_bound:.4f}. "
@@ -2456,7 +2588,7 @@ def _compute_histogram_quantile_2d(
         )
 
     # Check if any values are too low (ignore NaN values)
-    too_low = (threshold < lower_bound) & threshold.notnull()
+    too_low = threshold < lower_bound
     if too_low.any():
         warnings.warn(
             f"Quantile values below expected range in some locations: min={threshold.min().compute():.4f} < {lower_bound:.4f}. "
@@ -2466,7 +2598,7 @@ def _compute_histogram_quantile_2d(
             stacklevel=2,
         )
         # Set too low values to lower bound -- This is to ensure that constant=0 anomalies will not be "extreme"
-        threshold = threshold.where(~too_low, lower_bound).persist()
+        threshold = threshold.where(~too_low, lower_bound)
 
     return threshold
 

{marex-4.0.2 → marex-4.0.3}/marEx/helper.py

@@ -6,11 +6,14 @@ This module provides utilities for setting up and managing Dask clusters
 in HPC environments, with specific support for the DKRZ Levante Supercomputer.
 """
 
+import logging
 import re
+import shutil
 import subprocess
+import uuid
 from getpass import getuser
 from pathlib import Path
-from tempfile import TemporaryDirectory
+from tempfile import TemporaryDirectory, gettempdir
 from typing import Any, Dict, Optional, Union
 
 import dask
@@ -35,12 +38,32 @@ except (ImportError, ValueError):
     SLURMCluster = None
 
 
-# Default configuration values
+# Default configuration values optimised for HPC environments
+# These settings are tuned for large-scale climate data processing
+# and have been tested on DKRZ Levante with marEx workloads
 DEFAULT_DASK_CONFIG = {
+    # Array processing
     "array.slicing.split_large_chunks": False,
-    "distributed.comm.timeouts.connect": "240s",  # Increased from default
-    "distributed.comm.timeouts.tcp": "480s",  # Double the connection timeout
-    "distributed.comm.retry.count": 10,  # More retries before giving up
+    "array.chunk-size": "24MiB",  # Optimal chunk size for oceanographic data
+    # Worker memory management
+    "distributed.worker.memory.target": 0.4,  # Target memory threshold for spilling to disk
+    "distributed.worker.memory.spill": 0.5,  # Spill to disk threshold
+    "distributed.worker.memory.pause": 0.6,  # Pause worker threshold
+    "distributed.worker.memory.terminate": 0.8,  # Terminate worker threshold
+    "distributed.worker.memory.recent-to-old-time": "10s",  # Time to consider data old
+    "distributed.worker.daemon": False,  # Workers are not daemons
+    # Scheduler stability settings
+    "distributed.scheduler.allowed-failures": 50,  # Allow many retries (common on HPC)
+    "distributed.scheduler.work-stealing": False,  # Disable for deterministic execution
+    "distributed.scheduler.worker-ttl": "600s",  # Keep workers alive for 10 minutes
+    # Communication timeouts - increased for HPC network latency
+    "distributed.comm.timeouts.connect": "300s",  # Connection timeout
+    "distributed.comm.timeouts.tcp": "300s",  # TCP timeout
+    "distributed.comm.retry.count": 15,  # More retries before giving up
+    "distributed.comm.retry.delay.min": "3s",  # Min delay between retries
+    "distributed.comm.retry.delay.max": "30s",  # Max delay between retries
+    # Admin and logging
+    "distributed.admin.log-format": "%(name)s - %(levelname)s - %(message)s",  # Log format
 }
 
 # DKRZ-specific paths and configuration
@@ -358,6 +381,12 @@ def start_local_cluster(
 
     # Create cluster and client
     logger.debug("Creating local cluster and client")
+
+    # Configure Bokeh session token expiration via dask config
+    # Set to 24 hours to prevent Bokeh 3.8+ token expiration during long notebooks
+    dask.config.set({"distributed.scheduler.dashboard.bokeh-application.session_token_expiration": 86400000})
+    logger.debug("Set Bokeh session token expiration to 24 hours (86400000ms)")
+
     with log_timing(logger, "Local cluster startup", log_memory=True, show_progress=True):
         cluster = LocalCluster(n_workers=n_workers, threads_per_worker=threads_per_worker, **kwargs)
         client = Client(cluster)
@@ -562,6 +591,18 @@ def start_distributed_cluster(
 
     # Create SLURM cluster
     logger.info("Creating SLURM cluster")
+
+    # Configure Bokeh session token expiration via dask config
+    # Set to 24 hours to prevent Bokeh 3.8+ token expiration during long notebooks
+    dask.config.set({"distributed.scheduler.dashboard.bokeh-application.session_token_expiration": 86400000})
+    logger.debug("Set Bokeh session token expiration to 24 hours (86400000ms)")
+
+    # Merge user-provided scheduler_options from kwargs if present
+    user_scheduler_options = kwargs.pop("scheduler_options", {})
+    scheduler_options = {"dashboard_address": f":{dashboard_address}", **user_scheduler_options}
+
+    logger.debug(f"Scheduler options: {scheduler_options}")
+
     with log_timing(logger, "SLURM cluster creation"):
         cluster = SLURMCluster(
             name="dask-cluster",
@@ -576,7 +617,7 @@ def start_distributed_cluster(
             job_extra_directives=config["job_extra"],
             log_directory=DKRZ_LOG_PATH,
             local_directory=temp_dir.name,
-            scheduler_options={"dashboard_address": f":{dashboard_address}"},
+            scheduler_options=scheduler_options,
             **kwargs,
         )
 
@@ -598,6 +639,144 @@ def start_distributed_cluster(
     return client
 
 
+def checkpoint_to_zarr(
+    data: Union[xr.DataArray, xr.Dataset],
+    name: str = "checkpoint",
+    cleanup: bool = False,
+    timedim: str = "time",
+) -> Union[xr.DataArray, xr.Dataset]:  # pragma: no cover
+    """
+    Save and reload a Dask-backed xarray object to break graph dependencies.
+
+    This function materialises a Dask array/dataset to a temporary file
+    and immediately reloads it, thereby breaking the computational graph.
+    This prevents expensive recomputations when the same data is used multiple
+    times downstream.
+
+    Parameters
+    ----------
+    data : xarray.DataArray or xarray.Dataset
+        Dask-backed xarray object to checkpoint
+    name : str, default='checkpoint'
+        Name prefix for the temporary file (for logging/debugging)
+    cleanup : bool, default=False
+        Whether to delete the temporary file after reloading.
+        By default (False), temp files are kept for the session and cleaned up
+        by the OS temp directory manager. Set to True to immediately delete after reload.
+    timedim : str, default='time'
+        Name of the time dimension for chunking adjustments
+
+    Returns
+    -------
+    xarray.DataArray or xarray.Dataset
+        Reloaded data with broken graph dependencies
+
+    Examples
+    --------
+    >>> import marEx
+    >>> anomalies = marEx.compute_normalised_anomaly(sst)
+    >>> anomalies_checkpointed = marEx.helper.checkpoint_to_zarr(
+    ...     anomalies, name="anomalies"
+    ... )
+    """
+    logger.debug(f"Checkpointing '{name}' to break graph dependencies")
+
+    # Get dask temporary directory, fallback to system temp
+    try:
+        temp_base = dask.config.get("temporary-directory", None)
+        if temp_base is None or not Path(temp_base).exists():
+            temp_base = gettempdir()
+    except Exception:
+        temp_base = gettempdir()
+
+    unique_id = uuid.uuid4().hex[:8]
+    file_path = None
+
+    try:
+        try:
+            zarr_path = Path(temp_base) / f"marEx_checkpoint_{name}_{unique_id}.zarr"
+            file_path = zarr_path
+
+            logger.debug(f"Attempting Zarr checkpoint: {zarr_path}")
+            # Check if time dimension has irregular chunks that need fixing
+            if timedim in data.dims and timedim in data.chunks:
+                time_chunks = data.chunks[timedim]
+                if len(time_chunks) > 1 and len(set(time_chunks)) > 1:
+                    # Chunks are irregular - need to fix for Zarr
+                    first_chunk = time_chunks[0]
+                    total_size = data.sizes[timedim]
+
+                    logger.debug(f"Irregular {timedim} chunks detected: {time_chunks}")
+                    logger.debug(f"Total {timedim} dimension size: {total_size}")
+
+                    # Calculate how many full chunks we can have
+                    n_full_chunks = total_size // first_chunk
+                    remainder = total_size % first_chunk
+
+                    if remainder > 0:
+                        # Need full chunks + one smaller final chunk
+                        new_time_chunks = (first_chunk,) * n_full_chunks + (remainder,)
+                    else:
+                        # All chunks are equal size
+                        new_time_chunks = (first_chunk,) * n_full_chunks
+
+                    data = data.chunk({timedim: new_time_chunks})
+                    logger.debug(f"Adjusted {timedim} chunks for Zarr: {new_time_chunks}")
+                    logger.debug(f"Verification - sum of chunks: {sum(new_time_chunks)}, dimension size: {total_size}")
+
+            with log_timing(logger, f"Saving '{name}' to Zarr", logging.DEBUG, log_memory=False):
+                data.to_zarr(zarr_path, mode="w")
+
+            logger.debug("Zarr save successful, reloading...")
+            with log_timing(logger, f"Reloading '{name}' from Zarr", logging.DEBUG, log_memory=False):
+                if isinstance(data, xr.Dataset):
+                    reloaded = xr.open_zarr(zarr_path, chunks={})
+                else:
+                    ds_temp = xr.open_zarr(zarr_path, chunks={})
+                    reloaded = ds_temp[list(ds_temp.data_vars)[0]]
+
+            logger.info(f"Checkpoint '{name}' saved via Zarr: {zarr_path}")
+            return reloaded
+
+        except (ValueError, OSError) as e:
+            if "incompatible" in str(e) or "chunk" in str(e).lower():
+                logger.warning(f"Zarr failed due to irregular chunks: {str(e)[:200]}")
+                # Clean up failed zarr attempt
+                if zarr_path.exists():
+                    shutil.rmtree(zarr_path)
+            else:
+                raise
+
+    except Exception as e:
+        logger.error(f"Failed to checkpoint '{name}' to disk: {e}")
+        logger.warning(f"Falling back to in-memory persist() only for '{name}'")
+
+        # Fallback to in-memory persist (no disk I/O)
+        try:
+            reloaded = data.persist()
+            from distributed import wait
+
+            wait(reloaded)
+            logger.info(f"Checkpoint '{name}' persisted to distributed memory (no disk)")
+            return reloaded
+        except Exception as e2:
+            logger.error(f"Even persist() failed for '{name}': {e2}")
+            logger.warning("Returning original data without checkpointing")
+            return data
+
+    finally:
+        # Cleanup if requested
+        if cleanup and file_path and file_path.exists():
+            try:
+                if file_path.suffix == ".zarr":
+                    shutil.rmtree(file_path)
+                else:
+                    file_path.unlink()
+                logger.debug(f"Cleaned up checkpoint file: {file_path}")
+            except Exception as e:
+                logger.warning(f"Failed to cleanup {file_path}: {e}")
+
+
 def fix_dask_tuple_array(da: xr.DataArray) -> xr.DataArray:
     """
     Fix a dask array that has tuple (i.e. task) references in its chunks.

{marex-4.0.2 → marex-4.0.3}/marEx.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: marEx
-Version: 4.0.2
+Version: 4.0.3
 Summary: Marine Extremes Detection and Tracking
 Author-email: Aaron Wienkers <aaron@wienkers.com>
 Project-URL: Homepage, https://github.com/wienkers/marEx

{marex-4.0.2 → marex-4.0.3}/marEx.egg-info/SOURCES.txt

@@ -9,6 +9,9 @@ pyproject.toml
 .github/workflows/ci.yml
 .github/workflows/release.yml
 .github/workflows/validate.py
+examples/batch jobs/run_detect.py
+examples/batch jobs/run_track.py
+examples/batch jobs/submit_track.sh
 examples/gridded data/01_preprocess_extremes.ipynb
 examples/gridded data/02_id_track_events.ipynb
 examples/gridded data/03_visualise_events.ipynb

{marex-4.0.2 → marex-4.0.3}/tests/conftest.py

@@ -18,6 +18,12 @@ from dask.distributed import Client, LocalCluster
 # Configure logging to suppress noisy distributed messages
 def pytest_configure(config):
     """Configure pytest with custom logging settings."""
+    # Configure matplotlib to use non-interactive backend for CI/testing
+    # This must be done before importing pyplot to avoid TkAgg errors on Windows
+    import matplotlib
+
+    matplotlib.use("Agg")
+
     # Suppress all distributed logging to ERROR level or higher
     # This removes all the noisy INFO messages from distributed components
     distributed_loggers = [