Move dynamic chunking function (#35)

* Move dynamic chunking function

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update pyproject.toml

* Update test_cmip_transforms.py

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update pyproject.toml

* Update cmip_transforms.py

* Update pyproject.toml

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update cmip_transforms.py

* Update test_cmip_transforms.py

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update test_cmip_transforms.py

* Update test_cmip_transforms.py

* Update cmip_transforms.py

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update pyproject.toml

* Update cmip_transforms.py

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update test_cmip_transforms.py

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

leap_data_management_utils/cmip_transforms.py

@@ -3,17 +3,95 @@
 """
 
 import datetime
+import logging
+import warnings
 from dataclasses import dataclass
 
 import apache_beam as beam
+import xarray as xr
 import zarr
+from dask.utils import parse_bytes
+from dynamic_chunks.algorithms import (
+ NoMatchingChunks,
+ even_divisor_algo,
+ iterative_ratio_increase_algo,
+)
 from google.cloud import bigquery
 from pangeo_forge_recipes.transforms import Indexed, T
 from tqdm.auto import tqdm
 
 from leap_data_management_utils.cmip_testing import test_all
 from leap_data_management_utils.data_management_transforms import BQInterface
 
+# TODO: I am not sure the chunking function belongs here, but it clutters the recipe and I did not want
+# To open a whole file for this.
+logger = logging.getLogger(__name__)
+
+
+## Dynamic Chunking Wrapper
+def dynamic_chunking_func(ds: xr.Dataset) -> dict[str, int]:
+ logger.info(f'Input Dataset for dynamic chunking {ds =}')
+
+ target_chunk_size = '150MB'
+ target_chunks_aspect_ratio = {
+ 'time': 10,
+ 'x': 1,
+ 'i': 1,
+ 'ni': 1,
+ 'xh': 1,
+ 'nlon': 1,
+ 'lon': 1, # TODO: Maybe import all the known spatial dimensions from xmip?
+ 'y': 1,
+ 'j': 1,
+ 'nj': 1,
+ 'yh': 1,
+ 'nlat': 1,
+ 'lat': 1,
+ }
+ size_tolerance = 0.5
+
+ # Some datasets are smaller than the target chunk size and should not be chunked at all
+ if ds.nbytes < parse_bytes(target_chunk_size):
+ target_chunks = dict(ds.dims)
+
+ else:
+ try:
+ target_chunks = even_divisor_algo(
+ ds,
+ target_chunk_size,
+ target_chunks_aspect_ratio,
+ size_tolerance,
+ allow_extra_dims=True,
+ )
+
+ except NoMatchingChunks:
+ warnings.warn(
+ 'Primary algorithm using even divisors along each dimension failed '
+ 'with. Trying secondary algorithm.'
+ f'Input {ds=}'
+ )
+ try:
+ target_chunks = iterative_ratio_increase_algo(
+ ds,
+ target_chunk_size,
+ target_chunks_aspect_ratio,
+ size_tolerance,
+ allow_extra_dims=True,
+ )
+ except NoMatchingChunks:
+ raise ValueError(
+ 'Could not find any chunk combinations satisfying '
+ 'the size constraint with either algorithm.'
+ f'Input {ds=}'
+ )
+ # If something fails
+ except Exception as e:
+ raise e
+ except Exception as e:
+ raise e
+ logger.info(f'Dynamic Chunking determined {target_chunks =}')
+ return target_chunks
+
 
 @dataclass
 class IIDEntry:

leap_data_management_utils/tests/test_cmip_transforms.py

@@ -1,6 +1,8 @@
+import numpy as np
 import pytest
+import xarray as xr
 
-from leap_data_management_utils.cmip_transforms import IIDEntry
+from leap_data_management_utils.cmip_transforms import IIDEntry, dynamic_chunking_func
 
 
 class TestIIDEntry:
@@ -24,4 +26,11 @@ def test_too_long(self):
  IIDEntry(iid, store, retracted, tests_passed)
 
 
+class TestDynamicChunks:
+ def test_too_small(self):
+ ds = xr.DataArray(np.random.rand(4, 6)).to_dataset(name='data')
+ chunks = dynamic_chunking_func(ds)
+ assert chunks == {'dim_0': 4, 'dim_1': 6}
+
+
 # TODO Its super hard to test anything involving big query, because AFAIK there is no way to mock it.

pyproject.toml

@@ -39,6 +39,8 @@ pangeo-forge=[
  "pangeo-forge-esgf",
  "pangeo-forge-recipes",
  "apache-beam",
+ "dynamic-chunks",
+ "dask"
  ]
 
 catalog = [