first version of write_levels() function (uses xcube)

Author: forman

tests/test_levels.py

@@ -0,0 +1,73 @@
+# Copyright © 2024 Norman Fomferra and contributors
+# Permissions are hereby granted under the terms of the MIT License:
+# https://opensource.org/licenses/MIT.
+import json
+import unittest
+
+
+from zappend.fsutil import FileObj
+from zappend.levels import write_levels
+from zappend.levels import get_variables_config
+from .helpers import clear_memory_fs
+from .helpers import make_test_dataset
+
+try:
+    import xcube
+except:
+    xcube = None
+
+
+class GetVariablesConfigTest(unittest.TestCase):
+    def test_it(self):
+        dataset = make_test_dataset()
+        variables = get_variables_config(dataset, dict(x=512, y=256, time=1))
+        self.assertEqual(
+            {
+                "x": {"encoding": {"chunks": None}},
+                "y": {"encoding": {"chunks": None}},
+                "time": {"encoding": {"chunks": None}},
+                "chl": {"encoding": {"chunks": [1, 256, 512]}},
+                "tsm": {"encoding": {"chunks": [1, 256, 512]}},
+            },
+            variables,
+        )
+
+
+@unittest.skipIf(xcube is None, reason="xcube is not installed")
+class WriteLevelsTest(unittest.TestCase):
+    def setUp(self):
+        clear_memory_fs()
+
+    def test_it(self):
+        source_path = "memory://source.zarr"
+        make_test_dataset(
+            uri=source_path,
+            dims=("time", "y", "x"),
+            shape=(3, 1024, 2048),
+            chunks=(1, 128, 256),
+            crs="EPSG:4326",
+        )
+
+        target_dir = FileObj("memory://target.levels")
+        self.assertFalse(target_dir.exists())
+
+        write_levels(source_path=source_path, target_path=target_dir.uri)
+
+        self.assertTrue(target_dir.exists())
+
+        levels_file = target_dir.for_path(".zlevels")
+        self.assertTrue(levels_file.exists())
+        levels_info = json.loads(levels_file.read())
+        self.assertEqual(
+            {
+                "version": "1.0",
+                "num_levels": 4,
+                "agg_methods": {"chl": "mean", "tsm": "mean"},
+                "use_saved_levels": False,
+            },
+            levels_info,
+        )
+
+        # ds = xr.open_zarr(target_dir.uri + "/0.zarr")
+        # self.assertEqual({"time": 1, "y": 10, "x": 20}, ds.sizes)
+        # self.assertEqual({"x", "y", "time", "chl", "tsm"}, set(ds.variables))

zappend/levels.py

@@ -0,0 +1,176 @@
+# Copyright © 2024 Norman Fomferra and contributors
+# Permissions are hereby granted under the terms of the MIT License:
+# https://opensource.org/licenses/MIT.
+
+import json
+import logging
+from typing import Any, Hashable
+
+import xarray as xr
+import fsspec
+
+from zappend.api import zappend
+
+
+def write_levels(
+    source_path: str,
+    source_storage_options: dict[str, Any] | None = None,
+    target_path: str | None = None,
+    num_levels: int | None = None,
+    agg_methods: dict[str, Any] | None = None,
+    use_saved_levels: bool = False,
+    link_level_zero: bool = False,
+    xy_dim_names: tuple[str, str] | None = None,
+    tile_size: tuple[int, int] | None = None,
+    **zappend_config,
+):
+    from xcube.core.tilingscheme import get_num_levels
+    from xcube.core.gridmapping import GridMapping
+    from xcube.core.subsampling import get_dataset_agg_methods
+    from xcube.core.subsampling import subsample_dataset
+
+    target_dir = zappend_config.pop("target_dir", None)
+    if not target_dir and not target_path:
+        raise ValueError("either 'target_dir' or 'target_path' can be given, not both")
+    if target_dir and target_path and target_dir != target_path:
+        raise ValueError("either 'target_dir' or 'target_path' can be given, not both")
+    target_path = target_path or target_dir
+    target_storage_options = zappend_config.pop(
+        "target_storage_options", source_storage_options or {}
+    )
+    target_fs, target_root = fsspec.core.url_to_fs(
+        target_path, **target_storage_options
+    )
+
+    source_fs, source_root = fsspec.core.url_to_fs(
+        source_path,
+        **(
+            source_storage_options
+            if source_storage_options is not None
+            else target_storage_options
+        ),
+    )
+    source_store = source_fs.get_mapper(root=source_root)
+    source_ds = xr.open_zarr(source_store)
+
+    logger = logging.getLogger("zappend")
+
+    grid_mapping: GridMapping | None = None
+
+    if xy_dim_names is None:
+        grid_mapping = grid_mapping or GridMapping.from_dataset(source_ds)
+        xy_dim_names = grid_mapping.xy_dim_names
+
+    if tile_size is None:
+        grid_mapping = grid_mapping or GridMapping.from_dataset(source_ds)
+        tile_size = grid_mapping.tile_size
+
+    if num_levels is None:
+        grid_mapping = grid_mapping or GridMapping.from_dataset(source_ds)
+        num_levels = get_num_levels(grid_mapping.size, tile_size)
+
+    agg_methods = get_dataset_agg_methods(
+        source_ds,
+        xy_dim_names=xy_dim_names,
+        agg_methods=agg_methods,
+    )
+
+    force_new = zappend_config.pop("force_new", None)
+
+    append_dim = zappend_config.pop("append_dim", "time")
+    append_coord = source_ds.coords[append_dim]
+
+    variables = get_variables_config(
+        source_ds,
+        {
+            xy_dim_names[0]: tile_size[0],
+            xy_dim_names[0]: tile_size[1],
+            append_dim: 1,
+        },
+        variables=zappend_config.pop("variables", None),
+    )
+
+    with target_fs.open(f"{target_root}/.zlevels", "wt") as fp:
+        levels_data: dict[str, Any] = dict(
+            version="1.0",
+            num_levels=num_levels,
+            agg_methods=dict(agg_methods),
+            use_saved_levels=use_saved_levels,
+        )
+        json.dump(levels_data, fp, indent=2)
+
+    if link_level_zero:
+        with target_fs.open(f"{target_root}/0.link", "wt") as fp:
+            fp.write(source_root)
+
+    subsample_dataset_kwargs = dict(xy_dim_names=xy_dim_names, agg_methods=agg_methods)
+
+    for slice_index in range(append_coord.size):
+        slice_ds_indexer = {append_dim: slice(slice_index, slice_index + 1)}
+        slice_ds = source_ds.isel(slice_ds_indexer)
+
+        for level_index in range(num_levels):
+            if level_index == 0:
+                level_slice_ds = slice_ds
+            elif use_saved_levels:
+                prev_level_path = f"{target_root}/{level_index - 1}.zarr"
+                prev_level_store = target_fs.get_mapper(root=prev_level_path)
+                prev_level_ds = xr.open_zarr(prev_level_store)
+                level_slice_ds = subsample_dataset(
+                    prev_level_ds.isel(slice_ds_indexer),
+                    step=2,
+                    **subsample_dataset_kwargs,
+                )
+            else:
+                level_slice_ds = subsample_dataset(
+                    slice_ds,
+                    step=2**level_index,
+                    **subsample_dataset_kwargs,
+                )
+
+            if not link_level_zero or level_index > 0:
+                level_slice_path = f"{target_path}/{level_index}.zarr"
+                zappend(
+                    [level_slice_ds],
+                    target_dir=level_slice_path,
+                    target_storage_options=target_storage_options,
+                    append_dim=append_dim,
+                    force_new=force_new if slice_index == 0 else False,
+                    variables=variables,
+                    **zappend_config,
+                )
+
+        logger.info(f"done writing {target_path}")
+
+
+def get_variables_config(
+    dataset: xr.Dataset,
+    chunk_sizes: dict[Hashable, int],
+    variables: dict[str, dict[str, Any]] | None = None,
+):
+    """Define the chunk sizes for the variables in *dataset*.
+
+    Args:
+        dataset: The dataset
+        chunk_sizes: The chunk sizes
+        variables: Value of the zappend ``variables``
+            configuration parameter
+    Return:
+        A zappend compatible with the zappend ``variables``
+        configuration parameter.
+    """
+    var_configs = dict(variables or {})
+    for var_name, var in dataset.variables.items():
+        var_name = str(var_name)
+        var_config = dict(var_configs.get(var_name, {}))
+        var_encoding = dict(var_config.get("encoding", {}))
+        var_chunks = var_encoding.get("chunks")
+        if not var_chunks and var.dims:
+            if var_name in dataset.coords:
+                var_chunks = None
+            else:
+                var_chunks = [chunk_sizes.get(dim) for dim in var.dims]
+        var_encoding["chunks"] = var_chunks
+        var_config["encoding"] = var_encoding
+        var_configs[var_name] = var_config
+    return var_configs