Merge pull request #263 from langmore:fix-chunk-sizes

Weatherbench2 authors · Weatherbench2 authors · commit 80749789e864 · 2025-05-29T01:57:41.000-07:00
PiperOrigin-RevId: 764628744
diff --git a/.gitignore b/.gitignore
@@ -9,3 +9,4 @@ docs/source/.ipynb_checkpoints
 docs/*.zarr
 __pycache__
 .pytest_cache
+*swp
diff --git a/scripts/compute_probabilistic_climatological_forecasts.py b/scripts/compute_probabilistic_climatological_forecasts.py
@@ -939,6 +939,9 @@ def make_index_info(timedelta: np.ndarray, timedelta_offset: int):
       }
   )  # fmt: skip
   output_chunks.update(OUTPUT_CHUNKS.value)
+  output_chunks = {
+      k: min(output_chunks[k], template.sizes[k]) for k in output_chunks
+  }
 
   itemsize = max(var.dtype.itemsize for var in template.values())
 

Original file line number	Diff line number	Diff line change
`@@ -939,6 +939,9 @@ def make_index_info(timedelta: np.ndarray, timedelta_offset: int):`
`939`	`939`	`}`
`940`	`940`	`) # fmt: skip`
`941`	`941`	`output_chunks.update(OUTPUT_CHUNKS.value)`
	`942`	`+ output_chunks = {`
	`943`	`+ k: min(output_chunks[k], template.sizes[k]) for k in output_chunks`
	`944`	`+ }`
`942`	`945`
`943`	`946`	`itemsize = max(var.dtype.itemsize for var in template.values())`
`944`	`947`