Source code for xcube.core.mldataset.base

# Copyright (c) 2018-2024 by xcube team and contributors
# Permissions are hereby granted under the terms of the MIT License:

from typing import Any, Dict, Optional

import xarray as xr

from xcube.core.gridmapping import GridMapping
from xcube.core.schema import rechunk_cube
from xcube.core.subsampling import AggMethods
from xcube.core.subsampling import get_dataset_agg_methods
from xcube.core.subsampling import subsample_dataset
from xcube.core.tilingscheme import get_num_levels
from xcube.util.assertions import assert_instance
from xcube.util.assertions import assert_true
from .lazy import LazyMultiLevelDataset

[docs] class BaseMultiLevelDataset(LazyMultiLevelDataset): """A multi-level dataset whose level datasets are created by down-sampling a base dataset. Args: base_dataset: The base dataset for the level at index zero. grid_mapping: Optional grid mapping for *base_dataset*. num_levels: Optional number of levels. ds_id: Optional dataset identifier. agg_methods: Optional aggregation methods. May be given as string or as mapping from variable name pattern to aggregation method. Valid aggregation methods are None, "first", "min", "max", "mean", "median". If None, the default, "first" is used for integer variables and "mean" for floating point variables. """ def __init__( self, base_dataset: xr.Dataset, grid_mapping: Optional[GridMapping] = None, num_levels: Optional[int] = None, agg_methods: Optional[AggMethods] = "first", ds_id: Optional[str] = None, ): assert_instance(base_dataset, xr.Dataset, name="base_dataset") if grid_mapping is not None: assert_instance(grid_mapping, GridMapping, name="grid_mapping") if grid_mapping is None: # TODO (forman): why not computing it lazily? grid_mapping = GridMapping.from_dataset(base_dataset, tolerance=1e-4) self._agg_methods = get_dataset_agg_methods( base_dataset, xy_dim_names=grid_mapping.xy_dim_names, agg_methods=agg_methods, ) self._base_dataset = base_dataset super().__init__(grid_mapping=grid_mapping, num_levels=num_levels, ds_id=ds_id) @property def agg_methods(self): return self._agg_methods def _get_num_levels_lazily(self) -> int: gm = self.grid_mapping return get_num_levels(gm.size, gm.tile_size) def _get_dataset_lazily(self, index: int, parameters: Dict[str, Any]) -> xr.Dataset: """Compute the dataset at level *index*: If *index* is zero, return the base image passed to constructor, otherwise down-sample the dataset for the level at given *index*. Args: index: the level index parameters: currently unused Returns: the dataset for the level at *index*. """ assert_instance(index, int, name="index") if index < 0: index = self.num_levels + index assert_true(0 <= index < self.num_levels, message="index out of range") if index == 0: level_dataset = self._base_dataset else: level_dataset = subsample_dataset( self._base_dataset, 2**index, xy_dim_names=self.grid_mapping.xy_dim_names, agg_methods=self._agg_methods, ) # Tile each level according to grid mapping tile_size = self.grid_mapping.tile_size if tile_size is not None: level_dataset, _ = rechunk_cube( level_dataset, self.grid_mapping, tile_size=tile_size ) return level_dataset