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.
- {marex-4.0.2 → marex-4.0.4}/.github/workflows/ci.yml +5 -0
- {marex-4.0.2 → marex-4.0.4}/PKG-INFO +1 -1
- marex-4.0.4/examples/batch jobs/run_detect.py +91 -0
- marex-4.0.4/examples/batch jobs/run_track.py +131 -0
- marex-4.0.4/examples/batch jobs/submit_track.sh +79 -0
- {marex-4.0.2 → marex-4.0.4}/marEx/__init__.py +5 -0
- {marex-4.0.2 → marex-4.0.4}/marEx/detect.py +173 -38
- {marex-4.0.2 → marex-4.0.4}/marEx/helper.py +185 -6
- {marex-4.0.2 → marex-4.0.4}/marEx.egg-info/PKG-INFO +1 -1
- {marex-4.0.2 → marex-4.0.4}/marEx.egg-info/SOURCES.txt +3 -0
- {marex-4.0.2 → marex-4.0.4}/tests/conftest.py +6 -0
- {marex-4.0.2 → marex-4.0.4}/.codecov.yml +0 -0
- {marex-4.0.2 → marex-4.0.4}/.github/workflows/README.md +0 -0
- {marex-4.0.2 → marex-4.0.4}/.github/workflows/release.yml +0 -0
- {marex-4.0.2 → marex-4.0.4}/.github/workflows/validate.py +0 -0
- {marex-4.0.2 → marex-4.0.4}/.readthedocs.yaml +0 -0
- {marex-4.0.2 → marex-4.0.4}/CITATION.cff +0 -0
- {marex-4.0.2 → marex-4.0.4}/CONTRIBUTING.md +0 -0
- {marex-4.0.2 → marex-4.0.4}/MANIFEST.in +0 -0
- {marex-4.0.2 → marex-4.0.4}/README.md +0 -0
- {marex-4.0.2 → marex-4.0.4}/examples/gridded data/01_preprocess_extremes.ipynb +0 -0
- {marex-4.0.2 → marex-4.0.4}/examples/gridded data/02_id_track_events.ipynb +0 -0
- {marex-4.0.2 → marex-4.0.4}/examples/gridded data/03_visualise_events.ipynb +0 -0
- {marex-4.0.2 → marex-4.0.4}/examples/regional data/01_preprocess_extremes.ipynb +0 -0
- {marex-4.0.2 → marex-4.0.4}/examples/regional data/02_id_track_events.ipynb +0 -0
- {marex-4.0.2 → marex-4.0.4}/examples/regional data/03_visualise_events.ipynb +0 -0
- {marex-4.0.2 → marex-4.0.4}/examples/unstructured data/01_preprocess_extremes.ipynb +0 -0
- {marex-4.0.2 → marex-4.0.4}/examples/unstructured data/02_id_track_events.ipynb +0 -0
- {marex-4.0.2 → marex-4.0.4}/examples/unstructured data/03_visualise_events.ipynb +0 -0
- {marex-4.0.2 → marex-4.0.4}/marEx/_coverage_init.py +0 -0
- {marex-4.0.2 → marex-4.0.4}/marEx/_dependencies.py +0 -0
- {marex-4.0.2 → marex-4.0.4}/marEx/exceptions.py +0 -0
- {marex-4.0.2 → marex-4.0.4}/marEx/logging_config.py +0 -0
- {marex-4.0.2 → marex-4.0.4}/marEx/plotX/__init__.py +0 -0
- {marex-4.0.2 → marex-4.0.4}/marEx/plotX/base.py +0 -0
- {marex-4.0.2 → marex-4.0.4}/marEx/plotX/gridded.py +0 -0
- {marex-4.0.2 → marex-4.0.4}/marEx/plotX/unstructured.py +0 -0
- {marex-4.0.2 → marex-4.0.4}/marEx/track.py +0 -0
- {marex-4.0.2 → marex-4.0.4}/marEx.egg-info/dependency_links.txt +0 -0
- {marex-4.0.2 → marex-4.0.4}/marEx.egg-info/requires.txt +0 -0
- {marex-4.0.2 → marex-4.0.4}/marEx.egg-info/top_level.txt +0 -0
- {marex-4.0.2 → marex-4.0.4}/pyproject.toml +0 -0
- {marex-4.0.2 → marex-4.0.4}/setup.cfg +0 -0
- {marex-4.0.2 → marex-4.0.4}/tests/__init__.py +0 -0
- {marex-4.0.2 → marex-4.0.4}/tests/coverage_config.py +0 -0
- {marex-4.0.2 → marex-4.0.4}/tests/make_test_data.ipynb +0 -0
- {marex-4.0.2 → marex-4.0.4}/tests/test_detect_detrending.py +0 -0
- {marex-4.0.2 → marex-4.0.4}/tests/test_detect_helpers.py +0 -0
- {marex-4.0.2 → marex-4.0.4}/tests/test_error_handling.py +0 -0
- {marex-4.0.2 → marex-4.0.4}/tests/test_exceptions.py +0 -0
- {marex-4.0.2 → marex-4.0.4}/tests/test_gridded_preprocessing.py +0 -0
- {marex-4.0.2 → marex-4.0.4}/tests/test_gridded_tracking.py +0 -0
- {marex-4.0.2 → marex-4.0.4}/tests/test_integration.py +0 -0
- {marex-4.0.2 → marex-4.0.4}/tests/test_logging_system.py +0 -0
- {marex-4.0.2 → marex-4.0.4}/tests/test_plotx.py +0 -0
- {marex-4.0.2 → marex-4.0.4}/tests/test_track_edge_cases.py +0 -0
- {marex-4.0.2 → marex-4.0.4}/tests/test_track_helpers.py +0 -0
- {marex-4.0.2 → marex-4.0.4}/tests/test_unstructured_preprocessing.py +0 -0
- {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
|
|
|
@@ -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(
|
|
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
|
-
|
|
1509
|
-
|
|
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
|
-
).
|
|
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 =
|
|
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(
|
|
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
|
-
|
|
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
|
-
|
|
1827
|
-
|
|
1828
|
-
|
|
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
|
-
|
|
1832
|
-
|
|
1833
|
-
extremes =
|
|
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
|
-
|
|
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
|
-
|
|
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.
|
|
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
|
-
|
|
2377
|
-
|
|
2378
|
-
|
|
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
|
-
|
|
2381
|
-
|
|
2382
|
-
|
|
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.
|
|
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.
|
|
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
|
-
|
|
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.
|
|
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 =
|
|
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 =
|
|
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)
|
|
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
|
-
"
|
|
42
|
-
|
|
43
|
-
"distributed.
|
|
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=
|
|
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.
|
|
@@ -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
|
|
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
|