marEx 4.0.2__tar.gz → 4.0.4__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (59) hide show
  1. {marex-4.0.2 → marex-4.0.4}/.github/workflows/ci.yml +5 -0
  2. {marex-4.0.2 → marex-4.0.4}/PKG-INFO +1 -1
  3. marex-4.0.4/examples/batch jobs/run_detect.py +91 -0
  4. marex-4.0.4/examples/batch jobs/run_track.py +131 -0
  5. marex-4.0.4/examples/batch jobs/submit_track.sh +79 -0
  6. {marex-4.0.2 → marex-4.0.4}/marEx/__init__.py +5 -0
  7. {marex-4.0.2 → marex-4.0.4}/marEx/detect.py +173 -38
  8. {marex-4.0.2 → marex-4.0.4}/marEx/helper.py +185 -6
  9. {marex-4.0.2 → marex-4.0.4}/marEx.egg-info/PKG-INFO +1 -1
  10. {marex-4.0.2 → marex-4.0.4}/marEx.egg-info/SOURCES.txt +3 -0
  11. {marex-4.0.2 → marex-4.0.4}/tests/conftest.py +6 -0
  12. {marex-4.0.2 → marex-4.0.4}/.codecov.yml +0 -0
  13. {marex-4.0.2 → marex-4.0.4}/.github/workflows/README.md +0 -0
  14. {marex-4.0.2 → marex-4.0.4}/.github/workflows/release.yml +0 -0
  15. {marex-4.0.2 → marex-4.0.4}/.github/workflows/validate.py +0 -0
  16. {marex-4.0.2 → marex-4.0.4}/.readthedocs.yaml +0 -0
  17. {marex-4.0.2 → marex-4.0.4}/CITATION.cff +0 -0
  18. {marex-4.0.2 → marex-4.0.4}/CONTRIBUTING.md +0 -0
  19. {marex-4.0.2 → marex-4.0.4}/MANIFEST.in +0 -0
  20. {marex-4.0.2 → marex-4.0.4}/README.md +0 -0
  21. {marex-4.0.2 → marex-4.0.4}/examples/gridded data/01_preprocess_extremes.ipynb +0 -0
  22. {marex-4.0.2 → marex-4.0.4}/examples/gridded data/02_id_track_events.ipynb +0 -0
  23. {marex-4.0.2 → marex-4.0.4}/examples/gridded data/03_visualise_events.ipynb +0 -0
  24. {marex-4.0.2 → marex-4.0.4}/examples/regional data/01_preprocess_extremes.ipynb +0 -0
  25. {marex-4.0.2 → marex-4.0.4}/examples/regional data/02_id_track_events.ipynb +0 -0
  26. {marex-4.0.2 → marex-4.0.4}/examples/regional data/03_visualise_events.ipynb +0 -0
  27. {marex-4.0.2 → marex-4.0.4}/examples/unstructured data/01_preprocess_extremes.ipynb +0 -0
  28. {marex-4.0.2 → marex-4.0.4}/examples/unstructured data/02_id_track_events.ipynb +0 -0
  29. {marex-4.0.2 → marex-4.0.4}/examples/unstructured data/03_visualise_events.ipynb +0 -0
  30. {marex-4.0.2 → marex-4.0.4}/marEx/_coverage_init.py +0 -0
  31. {marex-4.0.2 → marex-4.0.4}/marEx/_dependencies.py +0 -0
  32. {marex-4.0.2 → marex-4.0.4}/marEx/exceptions.py +0 -0
  33. {marex-4.0.2 → marex-4.0.4}/marEx/logging_config.py +0 -0
  34. {marex-4.0.2 → marex-4.0.4}/marEx/plotX/__init__.py +0 -0
  35. {marex-4.0.2 → marex-4.0.4}/marEx/plotX/base.py +0 -0
  36. {marex-4.0.2 → marex-4.0.4}/marEx/plotX/gridded.py +0 -0
  37. {marex-4.0.2 → marex-4.0.4}/marEx/plotX/unstructured.py +0 -0
  38. {marex-4.0.2 → marex-4.0.4}/marEx/track.py +0 -0
  39. {marex-4.0.2 → marex-4.0.4}/marEx.egg-info/dependency_links.txt +0 -0
  40. {marex-4.0.2 → marex-4.0.4}/marEx.egg-info/requires.txt +0 -0
  41. {marex-4.0.2 → marex-4.0.4}/marEx.egg-info/top_level.txt +0 -0
  42. {marex-4.0.2 → marex-4.0.4}/pyproject.toml +0 -0
  43. {marex-4.0.2 → marex-4.0.4}/setup.cfg +0 -0
  44. {marex-4.0.2 → marex-4.0.4}/tests/__init__.py +0 -0
  45. {marex-4.0.2 → marex-4.0.4}/tests/coverage_config.py +0 -0
  46. {marex-4.0.2 → marex-4.0.4}/tests/make_test_data.ipynb +0 -0
  47. {marex-4.0.2 → marex-4.0.4}/tests/test_detect_detrending.py +0 -0
  48. {marex-4.0.2 → marex-4.0.4}/tests/test_detect_helpers.py +0 -0
  49. {marex-4.0.2 → marex-4.0.4}/tests/test_error_handling.py +0 -0
  50. {marex-4.0.2 → marex-4.0.4}/tests/test_exceptions.py +0 -0
  51. {marex-4.0.2 → marex-4.0.4}/tests/test_gridded_preprocessing.py +0 -0
  52. {marex-4.0.2 → marex-4.0.4}/tests/test_gridded_tracking.py +0 -0
  53. {marex-4.0.2 → marex-4.0.4}/tests/test_integration.py +0 -0
  54. {marex-4.0.2 → marex-4.0.4}/tests/test_logging_system.py +0 -0
  55. {marex-4.0.2 → marex-4.0.4}/tests/test_plotx.py +0 -0
  56. {marex-4.0.2 → marex-4.0.4}/tests/test_track_edge_cases.py +0 -0
  57. {marex-4.0.2 → marex-4.0.4}/tests/test_track_helpers.py +0 -0
  58. {marex-4.0.2 → marex-4.0.4}/tests/test_unstructured_preprocessing.py +0 -0
  59. {marex-4.0.2 → marex-4.0.4}/tests/test_unstructured_tracking.py +0 -0
@@ -64,6 +64,7 @@ jobs:
64
64
  name: Test Suite
65
65
  runs-on: ${{ matrix.os }}
66
66
  needs: code-quality
67
+ timeout-minutes: 90 # Windows can be significantly slower than Linux/macOS
67
68
  strategy:
68
69
  fail-fast: false
69
70
  matrix:
@@ -144,18 +145,22 @@ jobs:
144
145
  pytest tests/test_logging_system.py -n 2 -v --tb=short
145
146
 
146
147
  - name: Run preprocessing tests (gridded)
148
+ timeout-minutes: 30 # Windows can be slow for these tests
147
149
  run: |
148
150
  pytest tests/test_gridded_preprocessing.py -n 2 -v --tb=short
149
151
 
150
152
  - name: Run preprocessing tests (unstructured)
153
+ timeout-minutes: 30 # Windows can be slow for these tests
151
154
  run: |
152
155
  pytest tests/test_unstructured_preprocessing.py -n 2 -v --tb=short
153
156
 
154
157
  - name: Run tracking tests (gridded)
158
+ timeout-minutes: 30 # Windows can be slow for these tests
155
159
  run: |
156
160
  pytest tests/test_gridded_tracking.py -n 2 -v --tb=short
157
161
 
158
162
  - name: Run tracking tests (unstructured)
163
+ timeout-minutes: 30 # Windows can be slow for these tests
159
164
  run: |
160
165
  pytest tests/test_unstructured_tracking.py -n 2 -v --tb=short
161
166
 
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: marEx
3
- Version: 4.0.2
3
+ Version: 4.0.4
4
4
  Summary: Marine Extremes Detection and Tracking
5
5
  Author-email: Aaron Wienkers <aaron@wienkers.com>
6
6
  Project-URL: Homepage, https://github.com/wienkers/marEx
@@ -0,0 +1,91 @@
1
+ #!/usr/bin/env python
2
+ """
3
+ Global Daily SST (OSTIA) Analysis: Identifying Marine Extremes with MarEx-Detect
4
+
5
+ Pre-Process with Shifting-Baseline & Hobday et al. (2016) Definition for Extremes
6
+
7
+ Run from login node - this script will launch SLURM jobs via Dask distributed cluster
8
+
9
+ NOTE: This script uses a distributed cluster because detect.py does not contain
10
+ Numba JIT-compiled functions that execute locally. Unlike run_track.py, this
11
+ approach works correctly for detection/preprocessing tasks. If you prefer to
12
+ follow HPC best practices and run everything in SLURM jobs, consider creating
13
+ a submit_detect.sh script similar to submit_track.sh.
14
+ """
15
+
16
+ import os
17
+ from getpass import getuser
18
+ from pathlib import Path
19
+
20
+ import xarray as xr
21
+
22
+ import marEx
23
+ import marEx.helper as hpc
24
+
25
+
26
+ def main():
27
+ """Run MarEx preprocessing and extreme detection pipeline on distributed SLURM cluster."""
28
+ # Configuration from environment variables
29
+ n_workers = int(os.getenv("DASK_N_WORKERS", "128"))
30
+ workers_per_node = int(os.getenv("DASK_WORKERS_PER_NODE", "64"))
31
+ runtime = int(os.getenv("DASK_RUNTIME", "39"))
32
+ slurm_account = os.getenv("SLURM_ACCOUNT", "bk1377")
33
+
34
+ # Lustre Scratch Directory
35
+ scratch_dir = Path("/scratch") / getuser()[0] / getuser()
36
+
37
+ # Start Distributed Dask Cluster on SLURM
38
+ print("Starting distributed Dask cluster on SLURM...")
39
+ print(f" Workers: {n_workers}")
40
+ print(f" Workers per node: {workers_per_node}")
41
+ print(f" Runtime: {runtime} minutes")
42
+ print(f" SLURM account: {slurm_account}")
43
+
44
+ client = hpc.start_distributed_cluster(
45
+ n_workers=n_workers,
46
+ workers_per_node=workers_per_node,
47
+ runtime=runtime,
48
+ scratch_dir=scratch_dir / "clients",
49
+ account=slurm_account,
50
+ )
51
+
52
+ # Choose optimal chunk size & load data
53
+ print("Loading data...")
54
+ file_name = scratch_dir / "mhws" / "ostia.zarr"
55
+ chunk_size = {"time": 30, "lat": 360, "lon": 720}
56
+ sst = xr.open_zarr(str(file_name), chunks=chunk_size).sst
57
+ print(f"Data loaded: {sst}")
58
+
59
+ # Process Data using MarEx-Detect helper functions:
60
+ print("Processing extremes data...")
61
+ extremes_ds = marEx.preprocess_data(
62
+ sst,
63
+ method_anomaly="shifting_baseline", # Anomalies from a rolling climatology using previous window_year years
64
+ method_extreme="hobday_extreme", # Local day-of-year specific thresholds with windowing
65
+ threshold_percentile=95, # Use the 95th percentile as the extremes threshold
66
+ window_year_baseline=15,
67
+ smooth_days_baseline=21, # Defines the rolling climatology window (15 years) and smoothing window (21 days)
68
+ window_days_hobday=11, # Defines the window (11 days) of compiled samples collected for extremes detection
69
+ dimensions={
70
+ "time": "time",
71
+ "x": "lon",
72
+ "y": "lat",
73
+ }, # Define the dimensions of the data -- if 'y' exists, MarEx-Detect knows this is a gridded dataset
74
+ dask_chunks={"time": 25}, # Dask chunks for *output* data
75
+ use_temp_checkpoints=True, # Enable checkpointing to prevent expensive recomputations
76
+ )
77
+ print(f"Processing complete: {extremes_ds}")
78
+
79
+ # Save Extremes Data to zarr for more efficient parallel I/O
80
+ print("Saving extremes data...")
81
+ output_file = scratch_dir / "mhws" / "extremes_binary_gridded_shifting_hobday_batch.zarr"
82
+ extremes_ds.to_zarr(output_file, mode="w")
83
+ print(f"Data saved to: {output_file}")
84
+
85
+ # Close the Dask cluster
86
+ client.close()
87
+ print("Processing complete!")
88
+
89
+
90
+ if __name__ == "__main__":
91
+ main()
@@ -0,0 +1,131 @@
1
+ #!/usr/bin/env python
2
+ """
3
+ Global Daily Event Analysis: Marine Heatwave ID & Tracking using MarEx
4
+
5
+ Performs morphological preprocessing, blob identification, temporal tracking,
6
+ and graph reduction to create globally unique IDs for each tracked extreme event.
7
+
8
+ This script should be run INSIDE a SLURM job (via sbatch submit_track.sh).
9
+ It uses a local Dask cluster that utilizes the compute resources allocated to the job.
10
+ This ensures that all computation, including Numba JIT-compiled functions, runs on
11
+ compute nodes rather than the login node.
12
+ """
13
+
14
+ import os
15
+ from getpass import getuser
16
+ from pathlib import Path
17
+
18
+ import xarray as xr
19
+
20
+ import marEx
21
+ import marEx.helper as hpc
22
+
23
+
24
+ def main():
25
+ """Run MarEx event tracking pipeline on local Dask cluster within SLURM job allocation."""
26
+ # Configuration from environment variables
27
+ n_workers = int(os.getenv("DASK_N_WORKERS", "32"))
28
+ threads_per_worker = int(os.getenv("DASK_THREADS_PER_WORKER", "1"))
29
+
30
+ # Tracking configuration
31
+ run_basic = os.getenv("RUN_BASIC_TRACKER", "false").lower() == "true"
32
+ grid_resolution = float(os.getenv("GRID_RESOLUTION", "0.25"))
33
+ area_filter = int(os.getenv("AREA_FILTER", "600"))
34
+ r_fill = int(os.getenv("R_FILL", "12"))
35
+ t_fill = int(os.getenv("T_FILL", "4"))
36
+ overlap_threshold = float(os.getenv("OVERLAP_THRESHOLD", "0.25"))
37
+
38
+ # Lustre Scratch Directory
39
+ scratch_dir = Path("/scratch") / getuser()[0] / getuser()
40
+
41
+ # Start Local Dask Cluster
42
+ # This cluster runs within the SLURM job allocation, ensuring all computation
43
+ # (including Numba JIT-compiled functions) executes on compute nodes
44
+ print("Starting local Dask cluster within SLURM job allocation...")
45
+ print(f" Workers: {n_workers}")
46
+ print(f" Threads per worker: {threads_per_worker}")
47
+
48
+ client = hpc.start_local_cluster(
49
+ n_workers=n_workers,
50
+ threads_per_worker=threads_per_worker,
51
+ scratch_dir=scratch_dir / "clients",
52
+ )
53
+
54
+ # Choose optimal chunk size & load preprocessed data
55
+ print("\nLoading preprocessed extremes data...")
56
+ input_file = scratch_dir / "mhws" / "extremes_binary_gridded_shifting_hobday_batch.zarr"
57
+ chunk_size = {"time": 25, "lat": -1, "lon": -1}
58
+ ds = xr.open_zarr(str(input_file), chunks=chunk_size)
59
+ print(f"Data loaded: {ds}")
60
+
61
+ # Run ID, Tracking, & Merging
62
+ print("\n" + "=" * 80)
63
+ print("Running Full Tracking with Merge/Split Detection...")
64
+ print("=" * 80)
65
+
66
+ tracker = marEx.tracker(
67
+ ds.extreme_events,
68
+ ds.mask.where(
69
+ (ds.lat < 85) & (ds.lat > -90), other=False
70
+ ), # Modify Mask: Anisotropy of lat/lon grid near poles biases ID & Tracking
71
+ grid_resolution=grid_resolution, # Grid resolution in degrees, used to calculate object areas on the globe
72
+ area_filter_absolute=area_filter, # Remove objects smaller than specified cells
73
+ R_fill=r_fill, # Fill small holes with radius < R_fill _cells_
74
+ T_fill=t_fill, # Allow gaps of T_fill days and still continue event tracking with same ID
75
+ allow_merging=True, # Allow extreme events to split/merge. Keeps track of merge events & unique IDs
76
+ overlap_threshold=overlap_threshold, # Overlap threshold for merging events
77
+ nn_partitioning=True, # Use new NN method to partition merged children areas
78
+ verbose=True,
79
+ )
80
+
81
+ extreme_events_ds, merges_ds = tracker.run(return_merges=True)
82
+ print(f"\nTracking complete: {extreme_events_ds}")
83
+ print(f"Merges dataset: {merges_ds}")
84
+
85
+ # Save IDed/Tracked/Merged Events to zarr for efficient parallel I/O
86
+ print("\nSaving tracked events data...")
87
+ output_file = scratch_dir / "mhws" / "extreme_events_merged_gridded_shifting_batch.zarr"
88
+ extreme_events_ds.to_zarr(output_file, mode="w")
89
+ print(f"Events data saved to: {output_file}")
90
+
91
+ # Save Merges Dataset to netcdf
92
+ print("Saving merges data...")
93
+ merges_file = scratch_dir / "mhws" / "extreme_events_merged_gridded_shifting_merges_batch.nc"
94
+ merges_ds.to_netcdf(merges_file, mode="w")
95
+ print(f"Merges data saved to: {merges_file}")
96
+
97
+ # Optionally run basic tracking for comparison
98
+ if run_basic:
99
+ print("\n" + "=" * 80)
100
+ print("Running Basic Tracking (No Merge/Split, No Temporal Gap Filling)...")
101
+ print("=" * 80)
102
+
103
+ tracker_basic = marEx.tracker(
104
+ ds.extreme_events,
105
+ ds.mask.where(
106
+ (ds.lat < 85) & (ds.lat > -90), other=False
107
+ ), # Modify Mask: Anisotropy of lat/lon grid near poles biases ID & Tracking
108
+ area_filter_absolute=area_filter, # Remove objects smaller than specified cells
109
+ R_fill=r_fill, # Fill small holes with radius < R_fill _cells_
110
+ T_fill=0, # No temporal hole filling
111
+ allow_merging=False, # Do not allow extreme events to split/merge
112
+ )
113
+
114
+ extreme_events_basic_ds = tracker_basic.run()
115
+ print(f"\nBasic tracking complete: {extreme_events_basic_ds}")
116
+
117
+ # Save Basic IDed Events to zarr
118
+ print("\nSaving basic tracked events data...")
119
+ basic_output_file = scratch_dir / "mhws" / "extreme_events_basic_gridded_shifting_batch.zarr"
120
+ extreme_events_basic_ds.to_zarr(basic_output_file, mode="w")
121
+ print(f"Basic events data saved to: {basic_output_file}")
122
+
123
+ # Close the Dask cluster
124
+ client.close()
125
+ print("\n" + "=" * 80)
126
+ print("Tracking complete!")
127
+ print("=" * 80)
128
+
129
+
130
+ if __name__ == "__main__":
131
+ main()
@@ -0,0 +1,79 @@
1
+ #!/bin/bash
2
+ #SBATCH --job-name=marEx_track
3
+ #SBATCH --partition=compute
4
+ #SBATCH --cpus-per-task=128
5
+ #SBATCH --account=bk1377
6
+ #SBATCH --output=slurm_%j.out
7
+ #SBATCH --time=02:59:00
8
+ #SBATCH --mem=0
9
+ #SBATCH --exclusive
10
+
11
+ # SLURM Batch Script for MarEx Event Tracking
12
+ #
13
+ # This script submits run_track.py as a SLURM job, ensuring all computation
14
+ # (including Numba JIT-compiled functions) runs on allocated compute nodes.
15
+ #
16
+ # Usage:
17
+ # sbatch submit_track.sh
18
+ #
19
+ # Configuration via environment variables (optional):
20
+ # export DASK_N_WORKERS=32 # Number of Dask workers
21
+ # export DASK_THREADS_PER_WORKER=1 # Threads per worker
22
+ # export RUN_BASIC_TRACKER=true # Run basic tracker comparison
23
+ # export GRID_RESOLUTION=0.25 # Grid resolution in degrees
24
+ # export AREA_FILTER=600 # Minimum object area (cells)
25
+ # export R_FILL=12 # Spatial hole filling radius
26
+ # export T_FILL=4 # Temporal gap filling (days)
27
+ # export OVERLAP_THRESHOLD=0.25 # Merge overlap threshold
28
+ # sbatch submit_track.sh
29
+
30
+ # Print job information
31
+ echo "========================================="
32
+ echo "SLURM Job Information"
33
+ echo "========================================="
34
+ echo "Job ID: $SLURM_JOB_ID"
35
+ echo "Job Name: $SLURM_JOB_NAME"
36
+ echo "Node: $SLURM_NODELIST"
37
+ echo "Start Time: $(date)"
38
+ echo "Working Directory: $(pwd)"
39
+ echo "========================================="
40
+ echo ""
41
+
42
+ # Set ulimits for stack size
43
+ ulimit -S -s unlimited
44
+ ulimit -S -c unlimited
45
+
46
+ # Initialise conda for bash shell
47
+ source $(conda info --base)/etc/profile.d/conda.sh
48
+ conda activate super
49
+
50
+ # Set default environment variables if not already set
51
+ export DASK_N_WORKERS=${DASK_N_WORKERS:-32}
52
+ export DASK_THREADS_PER_WORKER=${DASK_THREADS_PER_WORKER:-1}
53
+
54
+ # Print configuration
55
+ echo "Configuration:"
56
+ echo " DASK_N_WORKERS: $DASK_N_WORKERS"
57
+ echo " DASK_THREADS_PER_WORKER: $DASK_THREADS_PER_WORKER"
58
+ echo " RUN_BASIC_TRACKER: ${RUN_BASIC_TRACKER:-false}"
59
+ echo " GRID_RESOLUTION: ${GRID_RESOLUTION:-0.25}"
60
+ echo " AREA_FILTER: ${AREA_FILTER:-600}"
61
+ echo " R_FILL: ${R_FILL:-12}"
62
+ echo " T_FILL: ${T_FILL:-4}"
63
+ echo " OVERLAP_THRESHOLD: ${OVERLAP_THRESHOLD:-0.25}"
64
+ echo ""
65
+
66
+ # Run the tracking script
67
+ echo "Starting MarEx tracking..."
68
+ python run_track.py
69
+
70
+ # Capture exit status
71
+ EXIT_STATUS=$?
72
+
73
+ echo ""
74
+ echo "========================================="
75
+ echo "Job completed at: $(date)"
76
+ echo "Exit status: $EXIT_STATUS"
77
+ echo "========================================="
78
+
79
+ exit $EXIT_STATUS
@@ -57,6 +57,9 @@ from .exceptions import ( # Main exception hierarchy; Convenience constructors
57
57
  wrap_exception,
58
58
  )
59
59
 
60
+ # Import HPC helper utilities
61
+ from .helper import configure_dask
62
+
60
63
  # Import logging configuration functions
61
64
  from .logging_config import (
62
65
  configure_logging,
@@ -116,6 +119,8 @@ __all__ = [
116
119
  "is_verbose_mode",
117
120
  "is_quiet_mode",
118
121
  "get_logger",
122
+ # HPC helper utilities
123
+ "configure_dask",
119
124
  ]
120
125
 
121
126
  # Version information
@@ -27,6 +27,7 @@ import flox.xarray
27
27
  import numpy as np
28
28
  import pandas as pd
29
29
  import xarray as xr
30
+ from dask import persist
30
31
  from dask.base import is_dask_collection
31
32
  from numpy.lib.stride_tricks import sliding_window_view
32
33
  from numpy.typing import NDArray
@@ -34,7 +35,7 @@ from xhistogram.xarray import histogram
34
35
 
35
36
  # Coordinate validation imports removed
36
37
  from .exceptions import ConfigurationError, create_data_validation_error
37
- from .helper import fix_dask_tuple_array
38
+ from .helper import checkpoint_to_zarr, fix_dask_tuple_array
38
39
  from .logging_config import configure_logging, get_logger, log_dask_info, log_memory_usage, log_timing
39
40
 
40
41
  # Get module logger
@@ -305,6 +306,7 @@ def preprocess_data(
305
306
  coordinates: Optional[Dict[str, str]] = None,
306
307
  neighbours: Optional[xr.DataArray] = None,
307
308
  cell_areas: Optional[xr.DataArray] = None,
309
+ use_temp_checkpoints: bool = False,
308
310
  verbose: Optional[bool] = None,
309
311
  quiet: Optional[bool] = None,
310
312
  ) -> xr.Dataset:
@@ -370,6 +372,12 @@ def preprocess_data(
370
372
  Neighbour connectivity for spatial clustering.
371
373
  cell_areas : xarray.DataArray, optional
372
374
  Cell areas for weighted spatial statistics.
375
+ use_temp_checkpoints : bool, default=False
376
+ Enable checkpointing to temporary zarr stores to break Dask graph dependencies.
377
+ When True, intermediate results (anomalies, thresholds, extremes) are saved to
378
+ temporary zarr files and immediately reloaded, preventing expensive recomputations.
379
+ Recommended for large datasets on HPC systems where the 2D histogram computation
380
+ is expensive. Temporary files are automatically cleaned up after reloading.
373
381
  verbose : bool, default=None
374
382
  Enable verbose logging with detailed progress information.
375
383
  If None, uses current global logging configuration.
@@ -580,6 +588,7 @@ def preprocess_data(
580
588
  std_normalise,
581
589
  detrend_orders,
582
590
  force_zero_mean,
591
+ use_temp_checkpoints,
583
592
  )
584
593
  log_memory_usage(logger, "After anomaly computation", logging.DEBUG)
585
594
 
@@ -612,6 +621,13 @@ def preprocess_data(
612
621
  time_sel = (ds[coordinates["time"]].dt.year >= start_year).compute()
613
622
  ds = ds.isel({dimensions["time"]: time_sel})
614
623
 
624
+ # Break graph after expensive anomaly computation
625
+ # Only needed for shifting baseline, because this is the most expensive method
626
+ # Other methods result in an odd chunking structure that cannot be checkpointed easily
627
+ if use_temp_checkpoints and method_anomaly == "shifting_baseline":
628
+ logger.debug("Checkpointing anomaly dataset to break graph dependencies")
629
+ ds = checkpoint_to_zarr(ds, name="anomalies", timedim=dimensions["time"])
630
+
615
631
  anomalies = ds.dat_anomaly
616
632
 
617
633
  # Step 2: Identify extreme events (both methods now return consistent tuple structures)
@@ -636,10 +652,14 @@ def preprocess_data(
636
652
  method_percentile,
637
653
  precision,
638
654
  max_anomaly,
655
+ use_temp_checkpoints,
639
656
  )
640
657
  log_memory_usage(logger, "After extreme identification", logging.DEBUG)
641
658
 
642
659
  # Add extreme events and thresholds to dataset
660
+ ds_temp = persist(extremes, thresholds)
661
+ extremes, thresholds = ds_temp
662
+
643
663
  ds["extreme_events"] = extremes
644
664
  ds["thresholds"] = thresholds
645
665
 
@@ -663,7 +683,15 @@ def preprocess_data(
663
683
  method_percentile,
664
684
  precision,
665
685
  max_anomaly,
686
+ use_temp_checkpoints,
666
687
  )
688
+
689
+ # Break graph after standardised extremes computation
690
+ if use_temp_checkpoints:
691
+ logger.debug("Checkpointing standardised extremes and thresholds to break graph dependencies")
692
+ extremes_stn = checkpoint_to_zarr(extremes_stn, name="extremes_stn", timedim=dimensions["time"])
693
+ thresholds_stn = checkpoint_to_zarr(thresholds_stn, name="thresholds_stn", timedim="dayofyear")
694
+
667
695
  ds["extreme_events_stn"] = extremes_stn
668
696
  ds["thresholds_stn"] = thresholds_stn
669
697
 
@@ -739,6 +767,14 @@ def preprocess_data(
739
767
  chunk_dict["dayofyear"] = time_chunks
740
768
  ds = ds.chunk(chunk_dict)
741
769
 
770
+ # Clear encoding metadata that may conflict with actual Dask chunks
771
+ # This encoding carries over from checkpointing and can cause chunk misalignment errors
772
+ logger.debug("Clearing encoding metadata for Dask-backed variables")
773
+ for var in ds.data_vars:
774
+ if hasattr(ds[var].data, "chunks"): # Only for Dask-backed variables
775
+ if hasattr(ds[var], "encoding") and "chunks" in ds[var].encoding:
776
+ del ds[var].encoding["chunks"]
777
+
742
778
  # Fix encoding issue with saving when calendar & units attribute is present
743
779
  if "calendar" in ds[coordinates["time"]].attrs: # pragma: no cover
744
780
  logger.debug("Removing calendar attribute for Zarr compatibility")
@@ -833,6 +869,7 @@ def compute_normalised_anomaly(
833
869
  std_normalise: bool = False, # for detrend_harmonic
834
870
  detrend_orders: Optional[List[int]] = None, # "
835
871
  force_zero_mean: bool = True, # "
872
+ use_temp_checkpoints: bool = False,
836
873
  verbose: Optional[bool] = None,
837
874
  quiet: Optional[bool] = None,
838
875
  ) -> xr.Dataset:
@@ -991,7 +1028,9 @@ def compute_normalised_anomaly(
991
1028
  return _compute_anomaly_detrended(da, std_normalise, detrend_orders, dimensions, coordinates, force_zero_mean)
992
1029
  elif method_anomaly == "shifting_baseline":
993
1030
  logger.debug(f"Shifting baseline parameters: window_years={window_year_baseline}, smooth_days={smooth_days_baseline}")
994
- return _compute_anomaly_shifting_baseline(da, window_year_baseline, smooth_days_baseline, dimensions, coordinates)
1031
+ return _compute_anomaly_shifting_baseline(
1032
+ da, window_year_baseline, smooth_days_baseline, dimensions, coordinates, use_temp_checkpoints
1033
+ )
995
1034
  elif method_anomaly == "fixed_baseline":
996
1035
  logger.debug("Fixed baseline parameters: using full time series for daily climatology")
997
1036
  return _compute_anomaly_fixed_baseline(da, dimensions, coordinates)
@@ -1028,6 +1067,7 @@ def identify_extremes(
1028
1067
  method_percentile: Literal["exact", "approximate"] = "approximate",
1029
1068
  precision: float = 0.01,
1030
1069
  max_anomaly: float = 5.0,
1070
+ use_temp_checkpoints: bool = False,
1031
1071
  verbose: Optional[bool] = None,
1032
1072
  quiet: Optional[bool] = None,
1033
1073
  ) -> Tuple[xr.DataArray, xr.DataArray]:
@@ -1385,6 +1425,7 @@ def identify_extremes(
1385
1425
  coordinates,
1386
1426
  precision,
1387
1427
  max_anomaly,
1428
+ use_temp_checkpoints,
1388
1429
  )
1389
1430
  else:
1390
1431
  logger.error(f"Unknown extreme method: {method_extreme}")
@@ -1412,6 +1453,7 @@ def rolling_climatology(
1412
1453
  window_year_baseline: int = 15,
1413
1454
  dimensions: Optional[Dict[str, str]] = None,
1414
1455
  coordinates: Optional[Dict[str, str]] = None,
1456
+ use_temp_checkpoints: bool = False,
1415
1457
  ) -> xr.DataArray:
1416
1458
  """
1417
1459
  Compute rolling climatology efficiently using flox cohorts.
@@ -1505,8 +1547,9 @@ def rolling_climatology(
1505
1547
  da = da.assign_coords({"year": years, "dayofyear": doys})
1506
1548
 
1507
1549
  # Get temporal bounds
1508
- year_vals = years.compute().values
1509
- doy_vals = doys.compute().values
1550
+ years, doys = persist(years, doys)
1551
+ year_vals = years.values
1552
+ doy_vals = doys.values
1510
1553
  unique_years = np.unique(year_vals)
1511
1554
  min_year = int(unique_years.min().item())
1512
1555
 
@@ -1563,7 +1606,11 @@ def rolling_climatology(
1563
1606
  isbin=(False, False),
1564
1607
  dtype=np.float32,
1565
1608
  fill_value=np.nan,
1566
- ).persist()
1609
+ ).chunk({"dayofyear": -1})
1610
+
1611
+ if use_temp_checkpoints:
1612
+ logger.debug("Checkpointing climatologies to break graph dependencies")
1613
+ climatologies = checkpoint_to_zarr(climatologies, name="climatologies", timedim=timedim)
1567
1614
 
1568
1615
  # Create index arrays for final mapping
1569
1616
  year_to_idx = pd.Series(range(len(unique_years)), index=unique_years)
@@ -1587,6 +1634,7 @@ def smoothed_rolling_climatology(
1587
1634
  smooth_days_baseline: int = 21,
1588
1635
  dimensions: Optional[Dict[str, str]] = None,
1589
1636
  coordinates: Optional[Dict[str, str]] = None,
1637
+ use_temp_checkpoints: bool = False,
1590
1638
  ) -> xr.DataArray:
1591
1639
  """
1592
1640
  Compute a smoothed rolling climatology using the previous `window_year_baseline` years of data
@@ -1699,9 +1747,11 @@ def smoothed_rolling_climatology(
1699
1747
  timedim = dimensions["time"]
1700
1748
 
1701
1749
  # N.B.: It is more efficient (chunking-wise) to smooth the raw data rather than the climatology
1702
- da_smoothed = da.rolling({timedim: smooth_days_baseline}, center=True).mean().chunk(dict(zip(da.dims, da.chunks)))
1750
+ da_smoothed = (
1751
+ da.rolling({timedim: smooth_days_baseline}, center=True).mean().chunk(dict(zip(da.dims, da.chunks))).astype(np.float32)
1752
+ )
1703
1753
 
1704
- clim = rolling_climatology(da_smoothed, window_year_baseline, dimensions, coordinates)
1754
+ clim = rolling_climatology(da_smoothed, window_year_baseline, dimensions, coordinates, use_temp_checkpoints)
1705
1755
 
1706
1756
  return clim
1707
1757
 
@@ -1712,6 +1762,7 @@ def _compute_anomaly_shifting_baseline(
1712
1762
  smooth_days_baseline: int = 21,
1713
1763
  dimensions: Optional[Dict[str, str]] = None,
1714
1764
  coordinates: Optional[Dict[str, str]] = None,
1765
+ use_temp_checkpoints: bool = False,
1715
1766
  ) -> xr.Dataset:
1716
1767
  """
1717
1768
  Compute anomalies using shifting baseline method with smoothed rolling climatology.
@@ -1725,14 +1776,15 @@ def _compute_anomaly_shifting_baseline(
1725
1776
  dimensions, coordinates = _infer_dims_coords(da, dimensions, coordinates)
1726
1777
 
1727
1778
  # Compute smoothed rolling climatology
1728
- climatology_smoothed = smoothed_rolling_climatology(da, window_year_baseline, smooth_days_baseline, dimensions, coordinates)
1779
+ climatology_smoothed = smoothed_rolling_climatology(
1780
+ da, window_year_baseline, smooth_days_baseline, dimensions, coordinates, use_temp_checkpoints
1781
+ )
1729
1782
 
1730
1783
  # Compute anomaly as difference from climatology
1731
1784
  anomalies = da - climatology_smoothed
1732
1785
 
1733
1786
  # Create ocean/land mask from first time step
1734
- chunk_dict_mask = {dimensions[dim]: -1 for dim in ["x", "y"] if dim in dimensions}
1735
- mask = np.isfinite(da.isel({dimensions["time"]: 0})).drop_vars({coordinates["time"]}).chunk(chunk_dict_mask)
1787
+ mask = np.isfinite(da.isel({dimensions["time"]: 0})).drop_vars({coordinates["time"]})
1736
1788
 
1737
1789
  # Build output dataset
1738
1790
  return xr.Dataset({"dat_anomaly": anomalies, "mask": mask})
@@ -1753,6 +1805,7 @@ def _identify_extremes_hobday(
1753
1805
  coordinates: Optional[Dict[str, str]] = None,
1754
1806
  precision: float = 0.01,
1755
1807
  max_anomaly: float = 5.0,
1808
+ use_temp_checkpoints: bool = False,
1756
1809
  ) -> Tuple[xr.DataArray, xr.DataArray]:
1757
1810
  """
1758
1811
  Identify extreme events using day-of-year (i.e. climatological percentile threshold).
@@ -1801,14 +1854,18 @@ def _identify_extremes_hobday(
1801
1854
  "If your time-series is very short, consider using method_percentile='exact'."
1802
1855
  )
1803
1856
 
1804
- # Add day-of-year coordinate
1805
- da = da.assign_coords(dayofyear=da[coordinates["time"]].dt.dayofyear)
1857
+ # Add day-of-year coordinate (compute it to avoid chunked groupby issues)
1858
+ da = da.assign_coords(dayofyear=da[coordinates["time"]].dt.dayofyear.compute()).chunk(dict(zip(da.dims, da.chunks))).persist()
1806
1859
 
1807
1860
  # Group by day-of-year and compute percentile
1808
1861
  if method_percentile == "exact":
1809
1862
  # Construct rolling window dimension
1810
1863
  da_windowed = da.rolling({dimensions["time"]: window_days_hobday}, center=True).construct("window")
1811
1864
 
1865
+ # Ensure dayofyear coordinate is computed for groupby (required by newer xarray)
1866
+ if "dayofyear" in da_windowed.coords:
1867
+ da_windowed = da_windowed.assign_coords(dayofyear=da_windowed.dayofyear.compute())
1868
+
1812
1869
  thresholds = da_windowed.groupby("dayofyear").reduce(
1813
1870
  np.nanpercentile, q=threshold_percentile, dim=("window", dimensions["time"])
1814
1871
  )
@@ -1821,16 +1878,62 @@ def _identify_extremes_hobday(
1821
1878
  dimensions=dimensions,
1822
1879
  precision=precision,
1823
1880
  max_anomaly=max_anomaly,
1881
+ use_temp_checkpoints=use_temp_checkpoints,
1824
1882
  )
1825
1883
 
1826
- # Ensure spatial dimensions are fully loaded for efficient comparison
1827
- spatial_chunks = {dimensions[dim]: -1 for dim in ["x", "y"] if dim in dimensions}
1828
- thresholds = thresholds.chunk(spatial_chunks).persist()
1884
+ if use_temp_checkpoints:
1885
+ logger.debug("Checkpointing thresholds to break graph dependencies")
1886
+ thresholds = checkpoint_to_zarr(thresholds, name="thresholds", timedim="dayofyear")
1887
+
1888
+ # Extract spatial chunk sizes from input data for alignment
1889
+ # Use most common chunk size to handle irregular chunks robustly
1890
+ spatial_chunks = {}
1891
+ for dim_key in ["x", "y"]:
1892
+ if dim_key in dimensions:
1893
+ dim_name = dimensions[dim_key]
1894
+ if dim_name in da.dims:
1895
+ dim_index = da.dims.index(dim_name)
1896
+ chunks_tuple = da.chunks[dim_index]
1897
+ # Get the most common chunk size (handles irregular chunks better)
1898
+ spatial_chunks[dim_name] = max(set(chunks_tuple), key=chunks_tuple.count)
1899
+
1900
+ # Drop time coordinate/dimension to avoid conflicts when comparing with data grouped by dayofyear
1901
+ coords_to_drop = []
1902
+ if coordinates["time"] in thresholds.coords:
1903
+ coords_to_drop.append(coordinates["time"])
1904
+ if dimensions["time"] in thresholds.coords and dimensions["time"] not in thresholds.dims:
1905
+ coords_to_drop.append(dimensions["time"])
1906
+ if "time" in thresholds.coords and "time" not in thresholds.dims:
1907
+ coords_to_drop.append("time")
1908
+ if coords_to_drop:
1909
+ thresholds = thresholds.drop_vars(coords_to_drop)
1910
+
1911
+ # Rechunk thresholds BEFORE comparison to align with input data
1912
+ # This eliminates expensive implicit rechunking during the groupby operation
1913
+ logger.debug(f"Aligning threshold chunks to match input data spatial chunks: {spatial_chunks}")
1914
+ thresholds = thresholds.chunk(spatial_chunks)
1829
1915
 
1830
1916
  # Compare anomalies to day-of-year specific thresholds
1831
- dayofyear_labels = da[coordinates["time"]].dt.dayofyear.compute()
1832
- extremes = da.groupby(dayofyear_labels) >= thresholds
1833
- extremes = extremes.astype(bool).chunk(spatial_chunks).persist()
1917
+ # Assign dayofyear coordinate and use UniqueGrouper for chunked arrays
1918
+ da = da.assign_coords(dayofyear=da[coordinates["time"]].dt.dayofyear)
1919
+ extremes = da.groupby(dayofyear=xr.groupers.UniqueGrouper(labels=np.arange(1, 367))) >= thresholds
1920
+
1921
+ # Drop unnecessary dayofyear coordinate
1922
+ if "dayofyear" in extremes.coords:
1923
+ extremes = extremes.drop_vars("dayofyear")
1924
+
1925
+ # Rechunk to fix irregular time chunks created by groupby operation
1926
+ # Zarr requires uniform chunks, so we rechunk to match input data's time chunks
1927
+ time_dim_index = da.dims.index(dimensions["time"])
1928
+ time_chunk_size = max(set(da.chunks[time_dim_index]), key=da.chunks[time_dim_index].count)
1929
+ rechunk_dict = {dimensions["time"]: time_chunk_size}
1930
+ rechunk_dict.update(spatial_chunks)
1931
+ logger.debug(f"Rechunking extremes to fix irregular chunks from groupby: {rechunk_dict}")
1932
+ extremes = extremes.chunk(rechunk_dict)
1933
+
1934
+ if use_temp_checkpoints:
1935
+ logger.debug("Checkpointing extremes to break graph dependencies")
1936
+ extremes = checkpoint_to_zarr(extremes, name="extremes", timedim=dimensions["time"])
1834
1937
 
1835
1938
  return extremes, thresholds
1836
1939
 
@@ -2085,7 +2188,12 @@ def _compute_anomaly_detrended(
2085
2188
  # Divide anomalies by rolling standard deviation
2086
2189
  # Replace any zeros or extremely small values with NaN to avoid division warnings
2087
2190
  std_rolling_safe = std_rolling.where(std_rolling > 1e-10, np.nan)
2088
- da_stn = da_detrend.groupby(da_detrend[coordinates["time"]].dt.dayofyear) / std_rolling_safe
2191
+ da_detrend = da_detrend.assign_coords(dayofyear=da_detrend[coordinates["time"]].dt.dayofyear)
2192
+ da_stn = da_detrend.groupby(dayofyear=xr.groupers.UniqueGrouper(labels=np.arange(1, 367))) / std_rolling_safe
2193
+
2194
+ # Drop dayofyear coordinate to avoid merge conflicts
2195
+ if "dayofyear" in da_stn.coords:
2196
+ da_stn = da_stn.drop_vars("dayofyear")
2089
2197
 
2090
2198
  # Rechunk data for efficient processing
2091
2199
  chunk_dict_std = chunk_dict_mask.copy()
@@ -2144,9 +2252,14 @@ def _compute_anomaly_fixed_baseline(
2144
2252
 
2145
2253
  # Compute anomalies by subtracting daily climatology from original data
2146
2254
  logger.debug("Computing anomalies by subtracting daily climatology")
2147
- anomalies = da.groupby(da[coordinates["time"]].dt.dayofyear) - daily_climatology
2255
+ da = da.assign_coords(dayofyear=da[coordinates["time"]].dt.dayofyear)
2256
+ anomalies = da.groupby(dayofyear=xr.groupers.UniqueGrouper(labels=np.arange(1, 367))) - daily_climatology
2148
2257
  anomalies = anomalies.astype(np.float32)
2149
2258
 
2259
+ # Drop dayofyear coordinate to avoid merge conflicts
2260
+ if "dayofyear" in anomalies.coords:
2261
+ anomalies = anomalies.drop_vars("dayofyear")
2262
+
2150
2263
  # Create ocean/land mask from first time step
2151
2264
  # Handle both spatial (3D) and time-series (1D) data
2152
2265
  spatial_dims = [dim for dim in ["x", "y"] if dim in dimensions]
@@ -2327,6 +2440,7 @@ def _compute_histogram_quantile_2d(
2327
2440
  dimensions: Optional[Dict[str, str]] = None,
2328
2441
  precision: float = 0.01,
2329
2442
  max_anomaly: float = 5.0,
2443
+ use_temp_checkpoints: bool = False,
2330
2444
  ) -> xr.DataArray:
2331
2445
  """
2332
2446
  Efficiently compute quantiles using binned histograms optimised for extreme values.
@@ -2366,23 +2480,31 @@ def _compute_histogram_quantile_2d(
2366
2480
  bin_centers_array[0] = 0.0
2367
2481
 
2368
2482
  bin_centers = xr.DataArray(
2369
- bin_centers_array,
2483
+ bin_centers_array.astype(np.float32),
2370
2484
  dims=["da_bin"],
2371
- coords={"da_bin": np.arange(len(bin_centers_array), dtype=np.int32)},
2485
+ coords={"da_bin": np.arange(len(bin_centers_array), dtype=np.uint16)},
2372
2486
  name="bin_centers",
2373
2487
  )
2374
2488
 
2375
2489
  chunk_dict = {dimensions["time"]: -1}
2376
- for d in ["x", "y"]:
2377
- if d in dimensions:
2378
- chunk_dict[dimensions[d]] = 8
2490
+ chunk_dict[dimensions["x"]] = 16
2491
+ if "y" in dimensions:
2492
+ chunk_dict[dimensions["y"]] = 16
2493
+
2494
+ da_bin = (
2495
+ xr.DataArray(
2496
+ np.digitize(da.data, bin_edges) - 1, # -1 so first bin is 0
2497
+ dims=da.dims,
2498
+ coords=da.coords,
2499
+ name="da_bin",
2500
+ )
2501
+ .chunk(chunk_dict)
2502
+ .astype(np.uint16)
2503
+ )
2379
2504
 
2380
- da_bin = xr.DataArray(
2381
- np.digitize(da.data, bin_edges) - 1, # -1 so first bin is 0
2382
- dims=da.dims,
2383
- coords=da.coords,
2384
- name="da_bin",
2385
- ).chunk(chunk_dict)
2505
+ if use_temp_checkpoints:
2506
+ logger.debug("Checkpointing binned data to break graph dependencies")
2507
+ da_bin = checkpoint_to_zarr(da_bin, name="da_bin", timedim=dimensions["time"]).chunk(chunk_dict)
2386
2508
 
2387
2509
  # Construct 2D histogram using flox (in doy & anomaly)
2388
2510
  hist_raw = flox.xarray.xarray_reduce(
@@ -2391,9 +2513,9 @@ def _compute_histogram_quantile_2d(
2391
2513
  da_bin,
2392
2514
  dim=[dimensions["time"]],
2393
2515
  func="count",
2394
- expected_groups=(np.arange(1, 367, dtype=np.int32), np.arange(len(bin_edges) - 1, dtype=np.int32)),
2516
+ expected_groups=(np.arange(1, 367, dtype=np.uint16), np.arange(len(bin_edges) - 1, dtype=np.uint16)),
2395
2517
  isbin=(False, False),
2396
- dtype=np.int32,
2518
+ dtype=np.uint16,
2397
2519
  fill_value=0,
2398
2520
  )
2399
2521
  hist_raw.name = None
@@ -2421,7 +2543,12 @@ def _compute_histogram_quantile_2d(
2421
2543
  return _rolling_histogram_quantile(hist_chunk, window_days_hobday, q, bin_centers_chunk)
2422
2544
 
2423
2545
  # Rechunk histogram so core dimensions are unchunked for apply_ufunc
2424
- hist_raw = hist_raw.chunk({"dayofyear": -1, "da_bin": -1})
2546
+ # Create chunk dict for hist_raw that preserves spatial chunks but drops time
2547
+ hist_chunk_dict = {dimensions["x"]: chunk_dict.get(dimensions["x"], 16), "dayofyear": -1, "da_bin": -1}
2548
+ if "y" in dimensions:
2549
+ hist_chunk_dict[dimensions["y"]] = chunk_dict.get(dimensions["y"], 16)
2550
+
2551
+ hist_raw = hist_raw.chunk(hist_chunk_dict)
2425
2552
 
2426
2553
  # Apply the optimised computation using apply_ufunc
2427
2554
  threshold = xr.apply_ufunc(
@@ -2437,8 +2564,16 @@ def _compute_histogram_quantile_2d(
2437
2564
  keep_attrs=True,
2438
2565
  )
2439
2566
 
2567
+ if use_temp_checkpoints:
2568
+ logger.debug("Checkpointing threshold to break graph dependencies")
2569
+ threshold = checkpoint_to_zarr(threshold, name="threshold")
2570
+
2571
+ # Drop time coordinate to avoid conflicts when comparing with data grouped by dayofyear
2572
+ if dimensions["time"] in threshold.coords:
2573
+ threshold = threshold.drop_vars(dimensions["time"])
2574
+
2440
2575
  # Set threshold to NaN for spatial points that contain NaN values
2441
- nan_mask = da.isnull().any(dim=dimensions["time"])
2576
+ nan_mask = da.isel({dimensions["time"]: 0}).isnull().compute()
2442
2577
  threshold = threshold.where(~nan_mask).persist()
2443
2578
 
2444
2579
  # Validate threshold values against bounds
@@ -2446,7 +2581,7 @@ def _compute_histogram_quantile_2d(
2446
2581
  lower_bound = bin_edges[3] # We want this to be positive so that constant=0 anomalies will not be "extreme"
2447
2582
 
2448
2583
  # Check if any values are too high (ignore NaN values)
2449
- too_high = (threshold > upper_bound) & threshold.notnull()
2584
+ too_high = threshold > upper_bound
2450
2585
  if too_high.any():
2451
2586
  warnings.warn(
2452
2587
  f"Quantile values exceed expected range: max={threshold.max().compute():.4f} > {upper_bound:.4f}. "
@@ -2456,7 +2591,7 @@ def _compute_histogram_quantile_2d(
2456
2591
  )
2457
2592
 
2458
2593
  # Check if any values are too low (ignore NaN values)
2459
- too_low = (threshold < lower_bound) & threshold.notnull()
2594
+ too_low = threshold < lower_bound
2460
2595
  if too_low.any():
2461
2596
  warnings.warn(
2462
2597
  f"Quantile values below expected range in some locations: min={threshold.min().compute():.4f} < {lower_bound:.4f}. "
@@ -2466,7 +2601,7 @@ def _compute_histogram_quantile_2d(
2466
2601
  stacklevel=2,
2467
2602
  )
2468
2603
  # Set too low values to lower bound -- This is to ensure that constant=0 anomalies will not be "extreme"
2469
- threshold = threshold.where(~too_low, lower_bound).persist()
2604
+ threshold = threshold.where(~too_low, lower_bound)
2470
2605
 
2471
2606
  return threshold
2472
2607
 
@@ -6,11 +6,14 @@ This module provides utilities for setting up and managing Dask clusters
6
6
  in HPC environments, with specific support for the DKRZ Levante Supercomputer.
7
7
  """
8
8
 
9
+ import logging
9
10
  import re
11
+ import shutil
10
12
  import subprocess
13
+ import uuid
11
14
  from getpass import getuser
12
15
  from pathlib import Path
13
- from tempfile import TemporaryDirectory
16
+ from tempfile import TemporaryDirectory, gettempdir
14
17
  from typing import Any, Dict, Optional, Union
15
18
 
16
19
  import dask
@@ -35,12 +38,32 @@ except (ImportError, ValueError):
35
38
  SLURMCluster = None
36
39
 
37
40
 
38
- # Default configuration values
41
+ # Default configuration values optimised for HPC environments
42
+ # These settings are tuned for large-scale climate data processing
43
+ # and have been tested on DKRZ Levante with marEx workloads
39
44
  DEFAULT_DASK_CONFIG = {
45
+ # Array processing
40
46
  "array.slicing.split_large_chunks": False,
41
- "distributed.comm.timeouts.connect": "240s", # Increased from default
42
- "distributed.comm.timeouts.tcp": "480s", # Double the connection timeout
43
- "distributed.comm.retry.count": 10, # More retries before giving up
47
+ "array.chunk-size": "24MiB", # Optimal chunk size for oceanographic data
48
+ # Worker memory management
49
+ "distributed.worker.memory.target": 0.4, # Target memory threshold for spilling to disk
50
+ "distributed.worker.memory.spill": 0.5, # Spill to disk threshold
51
+ "distributed.worker.memory.pause": 0.6, # Pause worker threshold
52
+ "distributed.worker.memory.terminate": 0.8, # Terminate worker threshold
53
+ "distributed.worker.memory.recent-to-old-time": "10s", # Time to consider data old
54
+ "distributed.worker.daemon": False, # Workers are not daemons
55
+ # Scheduler stability settings
56
+ "distributed.scheduler.allowed-failures": 50, # Allow many retries (common on HPC)
57
+ "distributed.scheduler.work-stealing": False, # Disable for deterministic execution
58
+ "distributed.scheduler.worker-ttl": "600s", # Keep workers alive for 10 minutes
59
+ # Communication timeouts - increased for HPC network latency
60
+ "distributed.comm.timeouts.connect": "300s", # Connection timeout
61
+ "distributed.comm.timeouts.tcp": "300s", # TCP timeout
62
+ "distributed.comm.retry.count": 15, # More retries before giving up
63
+ "distributed.comm.retry.delay.min": "3s", # Min delay between retries
64
+ "distributed.comm.retry.delay.max": "30s", # Max delay between retries
65
+ # Admin and logging
66
+ "distributed.admin.log-format": "%(name)s - %(levelname)s - %(message)s", # Log format
44
67
  }
45
68
 
46
69
  # DKRZ-specific paths and configuration
@@ -358,6 +381,12 @@ def start_local_cluster(
358
381
 
359
382
  # Create cluster and client
360
383
  logger.debug("Creating local cluster and client")
384
+
385
+ # Configure Bokeh session token expiration via dask config
386
+ # Set to 24 hours to prevent Bokeh 3.8+ token expiration during long notebooks
387
+ dask.config.set({"distributed.scheduler.dashboard.bokeh-application.session_token_expiration": 86400000})
388
+ logger.debug("Set Bokeh session token expiration to 24 hours (86400000ms)")
389
+
361
390
  with log_timing(logger, "Local cluster startup", log_memory=True, show_progress=True):
362
391
  cluster = LocalCluster(n_workers=n_workers, threads_per_worker=threads_per_worker, **kwargs)
363
392
  client = Client(cluster)
@@ -562,6 +591,18 @@ def start_distributed_cluster(
562
591
 
563
592
  # Create SLURM cluster
564
593
  logger.info("Creating SLURM cluster")
594
+
595
+ # Configure Bokeh session token expiration via dask config
596
+ # Set to 24 hours to prevent Bokeh 3.8+ token expiration during long notebooks
597
+ dask.config.set({"distributed.scheduler.dashboard.bokeh-application.session_token_expiration": 86400000})
598
+ logger.debug("Set Bokeh session token expiration to 24 hours (86400000ms)")
599
+
600
+ # Merge user-provided scheduler_options from kwargs if present
601
+ user_scheduler_options = kwargs.pop("scheduler_options", {})
602
+ scheduler_options = {"dashboard_address": f":{dashboard_address}", **user_scheduler_options}
603
+
604
+ logger.debug(f"Scheduler options: {scheduler_options}")
605
+
565
606
  with log_timing(logger, "SLURM cluster creation"):
566
607
  cluster = SLURMCluster(
567
608
  name="dask-cluster",
@@ -576,7 +617,7 @@ def start_distributed_cluster(
576
617
  job_extra_directives=config["job_extra"],
577
618
  log_directory=DKRZ_LOG_PATH,
578
619
  local_directory=temp_dir.name,
579
- scheduler_options={"dashboard_address": f":{dashboard_address}"},
620
+ scheduler_options=scheduler_options,
580
621
  **kwargs,
581
622
  )
582
623
 
@@ -598,6 +639,144 @@ def start_distributed_cluster(
598
639
  return client
599
640
 
600
641
 
642
+ def checkpoint_to_zarr(
643
+ data: Union[xr.DataArray, xr.Dataset],
644
+ name: str = "checkpoint",
645
+ cleanup: bool = False,
646
+ timedim: str = "time",
647
+ ) -> Union[xr.DataArray, xr.Dataset]: # pragma: no cover
648
+ """
649
+ Save and reload a Dask-backed xarray object to break graph dependencies.
650
+
651
+ This function materialises a Dask array/dataset to a temporary file
652
+ and immediately reloads it, thereby breaking the computational graph.
653
+ This prevents expensive recomputations when the same data is used multiple
654
+ times downstream.
655
+
656
+ Parameters
657
+ ----------
658
+ data : xarray.DataArray or xarray.Dataset
659
+ Dask-backed xarray object to checkpoint
660
+ name : str, default='checkpoint'
661
+ Name prefix for the temporary file (for logging/debugging)
662
+ cleanup : bool, default=False
663
+ Whether to delete the temporary file after reloading.
664
+ By default (False), temp files are kept for the session and cleaned up
665
+ by the OS temp directory manager. Set to True to immediately delete after reload.
666
+ timedim : str, default='time'
667
+ Name of the time dimension for chunking adjustments
668
+
669
+ Returns
670
+ -------
671
+ xarray.DataArray or xarray.Dataset
672
+ Reloaded data with broken graph dependencies
673
+
674
+ Examples
675
+ --------
676
+ >>> import marEx
677
+ >>> anomalies = marEx.compute_normalised_anomaly(sst)
678
+ >>> anomalies_checkpointed = marEx.helper.checkpoint_to_zarr(
679
+ ... anomalies, name="anomalies"
680
+ ... )
681
+ """
682
+ logger.debug(f"Checkpointing '{name}' to break graph dependencies")
683
+
684
+ # Get dask temporary directory, fallback to system temp
685
+ try:
686
+ temp_base = dask.config.get("temporary-directory", None)
687
+ if temp_base is None or not Path(temp_base).exists():
688
+ temp_base = gettempdir()
689
+ except Exception:
690
+ temp_base = gettempdir()
691
+
692
+ unique_id = uuid.uuid4().hex[:8]
693
+ file_path = None
694
+
695
+ try:
696
+ try:
697
+ zarr_path = Path(temp_base) / f"marEx_checkpoint_{name}_{unique_id}.zarr"
698
+ file_path = zarr_path
699
+
700
+ logger.debug(f"Attempting Zarr checkpoint: {zarr_path}")
701
+ # Check if time dimension has irregular chunks that need fixing
702
+ if timedim in data.dims and timedim in data.chunks:
703
+ time_chunks = data.chunks[timedim]
704
+ if len(time_chunks) > 1 and len(set(time_chunks)) > 1:
705
+ # Chunks are irregular - need to fix for Zarr
706
+ first_chunk = time_chunks[0]
707
+ total_size = data.sizes[timedim]
708
+
709
+ logger.debug(f"Irregular {timedim} chunks detected: {time_chunks}")
710
+ logger.debug(f"Total {timedim} dimension size: {total_size}")
711
+
712
+ # Calculate how many full chunks we can have
713
+ n_full_chunks = total_size // first_chunk
714
+ remainder = total_size % first_chunk
715
+
716
+ if remainder > 0:
717
+ # Need full chunks + one smaller final chunk
718
+ new_time_chunks = (first_chunk,) * n_full_chunks + (remainder,)
719
+ else:
720
+ # All chunks are equal size
721
+ new_time_chunks = (first_chunk,) * n_full_chunks
722
+
723
+ data = data.chunk({timedim: new_time_chunks})
724
+ logger.debug(f"Adjusted {timedim} chunks for Zarr: {new_time_chunks}")
725
+ logger.debug(f"Verification - sum of chunks: {sum(new_time_chunks)}, dimension size: {total_size}")
726
+
727
+ with log_timing(logger, f"Saving '{name}' to Zarr", logging.DEBUG, log_memory=False):
728
+ data.to_zarr(zarr_path, mode="w")
729
+
730
+ logger.debug("Zarr save successful, reloading...")
731
+ with log_timing(logger, f"Reloading '{name}' from Zarr", logging.DEBUG, log_memory=False):
732
+ if isinstance(data, xr.Dataset):
733
+ reloaded = xr.open_zarr(zarr_path, chunks={})
734
+ else:
735
+ ds_temp = xr.open_zarr(zarr_path, chunks={})
736
+ reloaded = ds_temp[list(ds_temp.data_vars)[0]]
737
+
738
+ logger.info(f"Checkpoint '{name}' saved via Zarr: {zarr_path}")
739
+ return reloaded
740
+
741
+ except (ValueError, OSError) as e:
742
+ if "incompatible" in str(e) or "chunk" in str(e).lower():
743
+ logger.warning(f"Zarr failed due to irregular chunks: {str(e)[:200]}")
744
+ # Clean up failed zarr attempt
745
+ if zarr_path.exists():
746
+ shutil.rmtree(zarr_path)
747
+ else:
748
+ raise
749
+
750
+ except Exception as e:
751
+ logger.error(f"Failed to checkpoint '{name}' to disk: {e}")
752
+ logger.warning(f"Falling back to in-memory persist() only for '{name}'")
753
+
754
+ # Fallback to in-memory persist (no disk I/O)
755
+ try:
756
+ reloaded = data.persist()
757
+ from distributed import wait
758
+
759
+ wait(reloaded)
760
+ logger.info(f"Checkpoint '{name}' persisted to distributed memory (no disk)")
761
+ return reloaded
762
+ except Exception as e2:
763
+ logger.error(f"Even persist() failed for '{name}': {e2}")
764
+ logger.warning("Returning original data without checkpointing")
765
+ return data
766
+
767
+ finally:
768
+ # Cleanup if requested
769
+ if cleanup and file_path and file_path.exists():
770
+ try:
771
+ if file_path.suffix == ".zarr":
772
+ shutil.rmtree(file_path)
773
+ else:
774
+ file_path.unlink()
775
+ logger.debug(f"Cleaned up checkpoint file: {file_path}")
776
+ except Exception as e:
777
+ logger.warning(f"Failed to cleanup {file_path}: {e}")
778
+
779
+
601
780
  def fix_dask_tuple_array(da: xr.DataArray) -> xr.DataArray:
602
781
  """
603
782
  Fix a dask array that has tuple (i.e. task) references in its chunks.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: marEx
3
- Version: 4.0.2
3
+ Version: 4.0.4
4
4
  Summary: Marine Extremes Detection and Tracking
5
5
  Author-email: Aaron Wienkers <aaron@wienkers.com>
6
6
  Project-URL: Homepage, https://github.com/wienkers/marEx
@@ -9,6 +9,9 @@ pyproject.toml
9
9
  .github/workflows/ci.yml
10
10
  .github/workflows/release.yml
11
11
  .github/workflows/validate.py
12
+ examples/batch jobs/run_detect.py
13
+ examples/batch jobs/run_track.py
14
+ examples/batch jobs/submit_track.sh
12
15
  examples/gridded data/01_preprocess_extremes.ipynb
13
16
  examples/gridded data/02_id_track_events.ipynb
14
17
  examples/gridded data/03_visualise_events.ipynb
@@ -18,6 +18,12 @@ from dask.distributed import Client, LocalCluster
18
18
  # Configure logging to suppress noisy distributed messages
19
19
  def pytest_configure(config):
20
20
  """Configure pytest with custom logging settings."""
21
+ # Configure matplotlib to use non-interactive backend for CI/testing
22
+ # This must be done before importing pyplot to avoid TkAgg errors on Windows
23
+ import matplotlib
24
+
25
+ matplotlib.use("Agg")
26
+
21
27
  # Suppress all distributed logging to ERROR level or higher
22
28
  # This removes all the noisy INFO messages from distributed components
23
29
  distributed_loggers = [
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes