"""Conversion staging cache functions."""
from __future__ import annotations
from collections.abc import Collection, Mapping, Sequence
from pathlib import PurePath
import numpy as np
import pyproj
import shapely
import structlog
import xarray as xr
import healpix_convert.core.utils as utils
from healpix_convert.core.coarsen import select_source_levels
from healpix_convert.core.conversion_models import (
ConvertStagingCache,
InputGroupSpatialInfo,
InputSpatialInfo,
OutputGroupInfo,
OutputSpatialInfo,
PyramidMultiscaleGroupInfo,
zarr_group_path,
)
from healpix_convert.core.healpix_conventions import Healpix
from healpix_convert.core.multiscales_conventions import HealpixMultiscales
from healpix_convert.settings.common import (
ConvertSettings,
HealpixDenseChunkSettings,
HealpixGroupSettings,
validate_convert_settings,
)
log = structlog.get_logger()
def _check_and_filter_input_groups(
datatrees: list[xr.DataTree],
groups: str | Sequence[str] | None,
) -> list[xr.DataTree]:
"""Return xr.DataTree objects with user group selection.
- check that all input datasets have the same hierarchical structure
- filter input datasets with user-defined groups
"""
dt0 = datatrees[0]
if any(not dt0.isomorphic(dt) for dt in datatrees[1:]):
raise ValueError(
"input Zarr datasets must all have the same hierarchical structure (groups)"
)
if groups is not None:
if isinstance(groups, str):
dt0 = dt0.match(groups)
else:
lgroups = list(groups)
dt0 = dt0.filter(lambda dt: dt.path in lgroups)
return [dt.filter_like(dt0) for dt in datatrees]
def _extract_root_spatial_info(datatrees: list[xr.DataTree]) -> InputSpatialInfo:
"""Extract spatial information from input Zarr datasets (root group)."""
spatial_info = InputSpatialInfo()
for dt in datatrees:
spatial_info.crs.append(utils.get_crs_from_stac(dt))
spatial_info.geometry_latlon.append(utils.get_wgs84_polygon_from_stac(dt))
spatial_info.geometry_xy.append(utils.get_proj_polygon_from_stac(dt))
return spatial_info
def _extract_spatial_info(
ds: xr.Dataset, default_crs: pyproj.CRS
) -> InputGroupSpatialInfo | None:
"""Extract spatial information in a group of a Zarr input dataset.
Return None if no spatial information is found.
"""
info = utils.extract_spatial_info_stac(ds, default_crs)
if info is None:
info = utils.extract_spatial_info_cf(ds)
if info is not None:
info = InputGroupSpatialInfo(**info)
return info
def _find_input_spatial_groups(
datatrees: list[xr.DataTree],
settings: ConvertSettings,
default_crss: list[pyproj.CRS],
) -> dict[PurePath, InputGroupSpatialInfo]:
"""Find spatial groups in input datasets, i.e., groups to convert to HEALPix.
- extract and return spatial information for each group found
- check that spatial information is consistent across all input datasets
- check that HEALPix conversion settings are given for each group found
"""
spatial_groups = {}
for path, dts in xr.group_subtrees(*datatrees):
no_data = [not dt.has_data for dt in dts]
if PurePath(path) in settings.exclude_groups:
continue
if all(no_data):
# a spatial group must have data
continue
elif any(no_data):
raise ValueError(
f"inconsistent group {path!r} across input datasets "
"(non-compatible data vs. no data)"
)
spatial_info = _extract_spatial_info(dts[0].dataset, default_crss[0])
for dt, default_crs in zip(dts[1:], default_crss[1:]):
info = _extract_spatial_info(dt.dataset, default_crs)
if spatial_info is None and info is None:
continue
elif spatial_info is not None and info is not None:
try:
spatial_info += info
except ValueError:
raise ValueError(
f"inconsistent group {path!r} across input datasets "
"(non-compatible spatial properties)"
)
else:
raise ValueError(
f"inconsistent group {path!r} across input datasets "
"(non-compatible spatial vs. non-spatial)"
)
if spatial_info is not None:
path = PurePath(path)
if not isinstance(settings.group_settings.get(path), HealpixGroupSettings):
raise KeyError(
f"HEALPix conversion settings are missing for group {path!r}"
)
spatial_groups[PurePath(path)] = spatial_info
return spatial_groups
#: Smallest number of refinement levels between a derived level and the
#: HEALPix cells its chunks are aligned with, i.e. the chunks of a derived
#: level hold at least ``4 ** MIN_CHUNK_DEPTH`` = 256 cells. Without such a
#: floor, following the chunk level of the converted data would divide the
#: number of cells per chunk by four at every derived level.
MIN_CHUNK_DEPTH = 4
def _chunk_cells(extent: shapely.Polygon, chunk: HealpixDenseChunkSettings):
"""Chunk cells of a converted level, i.e. the HEALPix cells of the chunk
level covering the spatial extent of the output, in ascending (nested)
order.
These are the cells the conversion writes one chunk for, so they say which
cells the level holds.
"""
assert chunk.healpix.refinement_level is not None
cell_ids, _ = utils.compute_output_chunk_info(
extent, chunk.healpix.refinement_level, chunk.healpix.ellipsoid.name
)
return np.sort(cell_ids)
def _coarsest_whole_level(chunk_cells, chunk_level: int) -> int:
"""Coarsest refinement level whose cells the converted cells cover
entirely.
A cell of a coarser level than the chunk level is covered entirely only if
all its children chunk cells are converted, which is not the case at the
edges of a partial spatial extent. Aggregating the cells of the converted
data onto such a cell would mix the cells of several of them, so no
derived level can go below that level.
"""
cells = chunk_cells
for level in range(chunk_level, 0, -1):
parents = cells >> np.uint64(2)
counts = np.bincount(np.unique(parents, return_inverse=True)[1])
if not np.all(counts == 4):
return level
cells = np.unique(parents)
return 0
def _derived_chunk_level(
level: int, native_level: int, native_chunk_level: int, coarsest_level: int = 0
) -> int:
"""Refinement level of the chunk cells of a derived level: a chunk holds
the ``4 ** depth`` children of one HEALPix cell ``depth`` levels above the
derived level.
The depth is the one of the converted chunks, so that the chunks of a
derived level hold as many cells as the ones of the level it is aggregated
from, with :py:data:`MIN_CHUNK_DEPTH` as a floor.
To ensure HEALPix-aligned chunks for all derived levels, the chunk cell
refinement level stops at ``coarsest_level``: the coarsest level that is
entirely covered by converted data (for a full-earth extent this
corresponds to level 0). Beyond that level threshold, some of the chunk
cells would be only partly covered, such that the actual chunks would be
cut along the array (i.e., not HEALPix-aligned).
"""
native_depth = native_level - native_chunk_level
return max(level - max(native_depth, MIN_CHUNK_DEPTH), coarsest_level)
def _derived_level_sources(
info: PyramidMultiscaleGroupInfo,
coarsen_from: str,
native_variables: Mapping[PurePath, Collection[str]],
) -> dict[PurePath, dict[str, PurePath]]:
"""Native level each variable of each derived level of a pyramid is
aggregated from (derived level path -> variable name -> native level path).
``native_variables`` gives the variables converted in each native level
(native level path -> variable names): a variable converted at several
native levels is aggregated from the one ``coarsen_from`` selects, so the
variables of a derived level may come from different native levels (see
:py:func:`~healpix_convert.core.coarsen.select_source_levels`).
"""
native: dict[int, PurePath] = {}
variables: dict[int, Collection[str]] = {}
for level_path, (healpix, _) in info.native_levels.items():
assert healpix.refinement_level is not None
native[healpix.refinement_level] = level_path
variables[healpix.refinement_level] = native_variables.get(level_path, ())
sources: dict[PurePath, dict[str, PurePath]] = {}
for level_path, (healpix, _) in info.derived_levels.items():
assert healpix.refinement_level is not None
selected = select_source_levels(
variables, healpix.refinement_level, coarsen_from=coarsen_from
)
sources[level_path] = {
name: native[source] for name, source in selected.items()
}
return sources
def _multiscales_metadata(
group_path: PurePath,
levels: dict[PurePath, Healpix],
add_dggs: bool,
*,
variables: Mapping[PurePath, Collection[str]] | None = None,
derived_from: dict[PurePath, dict[str, PurePath]] | None = None,
resampling_method: str | None = None,
asset_type: str = "auto",
) -> HealpixMultiscales:
"""Multiscales metadata of a multiscale group: its levels, finest first.
The assets are paths relative to the group holding the metadata, as the
Multiscales Zarr convention asks, and all of them refer to the same kind of
node, which the convention also asks: the level groups with ``asset_type``
``"group"``, one array of a level with ``"array"`` (the arrays of the
native levels are given by ``variables``, those of a derived level by
``derived_from``).
A level in ``derived_from`` is a derived level of a pyramid: it is
described by the native level its variables are aggregated from, the
transformation from that level (how much longer a cell is here) and the
method its cells are aggregated with.
One item per group names a single source, which cannot describe a derived
level whose variables are aggregated from different native levels. Such a
pyramid is rejected with ``"group"`` and described by one item per array
with ``"auto"``, which otherwise keeps the level groups as the assets.
"""
derived_from = derived_from or {}
variables = variables or {}
layout: list[dict] = []
mixed = [
level_path
for level_path, sources in derived_from.items()
if len(set(sources.values())) > 1
]
if asset_type == "group" and mixed:
sources = derived_from[mixed[0]]
raise ValueError(
f"the variables of the derived level "
f"{mixed[0].relative_to(group_path).as_posix()!r} of the multiscale "
f"group {group_path.as_posix()!r} are aggregated from different "
f"native levels ({', '.join(sorted({p.name for p in sources.values()}))}"
"), which one multiscales layout item per group cannot describe: "
"set the 'asset_type' multiscale setting to 'array' or 'auto'"
)
per_array = asset_type == "array" or (asset_type == "auto" and bool(mixed))
def derived_fields(
level_healpix: Healpix, source_path: PurePath, name: str | None = None
) -> dict:
"""What a derived level, or one of its arrays, is aggregated from."""
depth = levels[source_path].refinement_level - level_healpix.refinement_level
source_asset = source_path.relative_to(group_path)
return {
"derived_from": source_asset if name is None else source_asset / name,
# one factor for the single (cell) dimension: a cell of this level
# is 2 ** depth source cells long, hence aggregates 4 ** depth of
# them over the two dimensions of the sphere
"transform": {"scale": [float(2**depth)]},
"resampling_method": resampling_method,
}
for level_path, level_healpix in sorted(
levels.items(), key=lambda item: item[1].refinement_level, reverse=True
):
asset = level_path.relative_to(group_path)
item: dict = {"asset": asset}
if add_dggs:
item["dggs"] = level_healpix
# a native level, converted from the input data, has no source
sources = derived_from.get(level_path, {})
if not per_array:
if sources:
item |= derived_fields(level_healpix, set(sources.values()).pop())
layout.append(item)
continue
names = sorted(sources) if sources else sorted(variables.get(level_path, ()))
if not names:
raise ValueError(
f"no array is known for the level {asset.as_posix()!r} of the "
f"multiscale group {group_path.as_posix()!r}, which one "
"multiscales layout item per array requires"
)
for name in names:
array_item = item | {"asset": asset / name}
if sources:
array_item |= derived_fields(level_healpix, sources[name], name)
layout.append(array_item)
metadata: dict = {"layout": layout}
if resampling_method is not None:
# the same method aggregates every derived level of a pyramid
metadata["resampling_method"] = resampling_method
return HealpixMultiscales.model_validate(metadata)
def _set_pyramid_multiscale_groups(
settings: ConvertSettings,
io_path_map: dict[PurePath, PurePath],
output_spatial: OutputSpatialInfo,
) -> dict[PurePath, PyramidMultiscaleGroupInfo]:
"""Compute the structure of the pyramids given in the multiscale settings:
the levels of each multiscale group holding one, with their HEALPix grid
and their chunks.
The derived levels are the ``derived_levels`` of the pyramid settings. They
are written next to the native levels or in a child group of the multiscale
group, as its ``subgroup_name`` pyramid setting says.
"""
pyramid_multiscale_groups: dict[PurePath, PyramidMultiscaleGroupInfo] = {}
for path, pyramid in settings.multiscale_settings.pyramid_groups.items():
native_paths = settings.native_refinement_levels(path)
finest_level = max(native_paths)
finest_settings = settings.group_settings[native_paths[finest_level]]
# the dense chunk method of every child group of a pyramid
# multiscale group is enforced when the settings are validated
assert isinstance(finest_settings.chunk, HealpixDenseChunkSettings)
assert finest_settings.chunk.healpix.refinement_level is not None
chunk_level = finest_settings.chunk.healpix.refinement_level
coarsest_level = _coarsest_whole_level(
_chunk_cells(output_spatial.geometry_latlon, finest_settings.chunk),
chunk_level,
)
too_coarse = sorted(
level for level in pyramid.derived_levels if level < coarsest_level
)
if too_coarse:
raise ValueError(
f"pyramid multiscale group {path.as_posix()!r}: the converted "
f"cells do not cover whole HEALPix cells below the refinement "
f"level {coarsest_level} (the spatial extent of the output covers "
f"the cells of the chunk level {chunk_level} partially), so the "
f"derived levels {too_coarse} cannot be generated"
)
native_levels = {}
for input_path in native_paths.values():
group_settings = settings.group_settings[input_path]
assert isinstance(group_settings.chunk, HealpixDenseChunkSettings)
assert group_settings.chunk.healpix.refinement_level is not None
native_levels[io_path_map[input_path]] = (
group_settings.healpix,
group_settings.chunk.healpix.refinement_level,
)
# the derived levels are written next to the native levels, or in a
# child group of the multiscale group
group_path = io_path_map[path]
if pyramid.subgroup_name is None:
derived_path = group_path
else:
derived_path = group_path / pyramid.subgroup_name
healpix = finest_settings.healpix
derived_levels = {}
for level in sorted(pyramid.derived_levels, reverse=True):
derived_levels[derived_path / str(level)] = (
healpix.model_copy(update={"refinement_level": level}),
_derived_chunk_level(level, finest_level, chunk_level, coarsest_level),
)
pyramid_multiscale_groups[path] = PyramidMultiscaleGroupInfo(
native_levels=native_levels,
derived_levels=derived_levels,
)
return pyramid_multiscale_groups
def _converted_arrays(
input_spatial_groups: Mapping[PurePath, InputGroupSpatialInfo],
input_path: PurePath,
) -> Collection[str]:
"""Arrays an input group holds once converted: its spatial arrays, the ones
the conversion resamples onto the HEALPix cells.
"""
group_spatial = input_spatial_groups.get(input_path)
return () if group_spatial is None else group_spatial.spatial_arrays
def _set_output_groups(
datatrees: list[xr.DataTree],
input_spatial_groups: dict[PurePath, InputGroupSpatialInfo],
settings: ConvertSettings,
output_spatial: OutputSpatialInfo,
) -> OutputGroupInfo:
multiscale_groups: dict[PurePath, HealpixMultiscales] = {}
excluded_groups: list[PurePath] = []
io_path_map: dict[PurePath, PurePath] = {}
multiscales_settings = settings.multiscale_settings
for path, dts in xr.group_subtrees(*datatrees):
path = PurePath(path)
if path in settings.exclude_groups:
excluded_groups.append(path)
continue
if path not in io_path_map:
io_path_map[path] = path
# datatrees isomorphism has been checked before
dt0 = dts[0]
# configure multiscales
if path in multiscales_settings.groups:
levels: dict[PurePath, Healpix] = {}
level_variables: dict[PurePath, Collection[str]] = {}
for child_dt in dt0.children.values():
child_path = PurePath(child_dt.path.lstrip("/"))
child_settings = settings.group_settings.get(child_path)
if isinstance(child_settings, HealpixGroupSettings):
group_output_path = child_path.parent / str(
child_settings.healpix.refinement_level
)
io_path_map[child_path] = group_output_path
levels[group_output_path] = child_settings.healpix
level_variables[group_output_path] = _converted_arrays(
input_spatial_groups, child_path
)
if not levels:
raise ValueError(
f"multiscale group {path} has no child group to convert to HEALPix"
)
multiscale_groups[path] = _multiscales_metadata(
path,
levels,
multiscales_settings.add_healpix_positioning,
variables=level_variables,
asset_type=multiscales_settings.asset_type,
)
pyramid_multiscale_groups = _set_pyramid_multiscale_groups(
settings, io_path_map, output_spatial
)
# the arrays a native level holds once converted, which tell the levels
# the variables of a derived level are aggregated from
native_variables = {
io_path_map[input_path]: group_spatial.spatial_arrays
for input_path, group_spatial in input_spatial_groups.items()
if input_path in io_path_map
}
for path, info in pyramid_multiscale_groups.items():
# the metadata of a pyramid multiscale group describes all its levels,
# the derived ones included
all_levels = {}
for level_path, (level_healpix, _) in (
info.native_levels | info.derived_levels
).items():
all_levels[level_path] = level_healpix
pyramid = multiscales_settings.pyramid_groups[path]
multiscale_groups[path] = _multiscales_metadata(
io_path_map[path],
all_levels,
multiscales_settings.add_healpix_positioning,
variables=native_variables,
derived_from=_derived_level_sources(
info, pyramid.coarsen_from, native_variables
),
resampling_method=pyramid.reduce,
asset_type=multiscales_settings.asset_type,
)
return OutputGroupInfo(
multiscale_groups=multiscale_groups,
pyramid_multiscale_groups=pyramid_multiscale_groups,
excluded_groups=excluded_groups,
io_path_map=io_path_map,
)
def _check_consolidate_groups(
output_groups: OutputGroupInfo, settings: ConvertSettings
) -> None:
"""Warn about the groups named by the ``groups`` consolidate metadata
setting that the output dataset does not hold.
"""
paths = output_groups.group_paths()
missing = [
path
for path in settings.metadata.consolidate.groups
if zarr_group_path(path) not in paths
]
if missing:
log.warning(
f"the 'consolidate' metadata settings name {len(missing)} group(s) "
"that the output dataset does not hold (will be ignored)\n"
+ "\n".join(f" - {path}" for path in missing)
)
def _set_output_spatial_info(
input_spatial: InputSpatialInfo,
output_extent: dict | shapely.Polygon | None,
) -> OutputSpatialInfo:
if output_extent is None:
geometry = utils.compute_output_extent_latlon(*input_spatial.get_extents())
elif isinstance(output_extent, shapely.Polygon):
geometry = output_extent
else:
geometry = shapely.geometry.shape(output_extent)
if not isinstance(output_extent, shapely.Polygon):
raise ValueError("output extent must represent a single polygon.")
return OutputSpatialInfo(geometry_latlon=geometry)
[docs]
def create_staging_cache(
input_paths: Sequence[str],
settings: dict | ConvertSettings,
*,
groups: str | Sequence[str] | None = None,
output_extent: dict | shapely.Polygon | None = None,
) -> ConvertStagingCache:
"""Create a temporary (staging) conversion data cache by parsing input dataset(s)
and conversion settings.
Useful for testing and debugging or for multi-stage conversion workflows.
Parameters
----------
input_paths : list of paths
Path(s) to the input Zarr dataset(s) to convert to HEALPix.
settings : dict or object
Healpix conversion settings.
groups : str or list, optional
Groups (paths) in the input datasets to process. It can be
either a string (single group or unix-like pattern) or a
list of groups. By default, all groups found in input datasets
are processed.
output_extent : dict or :py:class:`shapely.Polygon`, optional
Spatial extent of the output HEALPix dataset, given as a
(GEOJSON-like or object) polygon with lat-lon coordinates.
If not given (default), the spatial extent will be the union
of the spatial extents of the input datasets (single polygon).
Returns
-------
:py:class:`ConvertStagingCache`
A reusable data cache that contains information extracted from
the input dataset(s) and/or pre-computed for the output dataset.
"""
cache = ConvertStagingCache()
settings = validate_convert_settings(settings)
# --- open and filter input Zarr datasets
log.info("••• reading Zarr datasets...")
input_datatrees = _check_and_filter_input_groups(
utils.open_datatrees(list(input_paths)), groups
)
cache.input_datatrees = {path: dt for path, dt in zip(input_paths, input_datatrees)}
# --- extract (root) spatial info in input Zarr datasets
log.info("••• extracting spatial info from input datasets (root group)...")
cache.input_spatial = _extract_root_spatial_info(input_datatrees)
unique_crss = set(cache.input_spatial.crs)
log.info(
f"found {len(unique_crss)} CRS(s) in input datasets:\n"
+ "\n".join(f" - {crs.name}" for crs in unique_crss)
)
# --- compute or set output spatial extent
cache.output_spatial = _set_output_spatial_info(cache.input_spatial, output_extent)
log.info(
"spatial extent of the output HEALPix dataset (WGS84):\n"
f" - geometry: {cache.output_spatial.geometry_latlon}\n"
f" - bbox: {cache.output_spatial.bbox}"
)
# --- detect input spatial groups (to convert to HEALPix)
log.info(
"••• detecting spatial groups in input datasets (to convert to HEALPix)..."
)
default_crss = cache.input_spatial.crs
cache.input_spatial_groups = _find_input_spatial_groups(
input_datatrees, settings, default_crss=default_crss
)
log.info(
f"found {len(cache.input_spatial_groups)} spatial groups to convert to HEALPix:\n"
+ "\n".join(f" - {path}" for path in cache.input_spatial_groups)
)
# --- configure HEALPix output groups
log.info("••• configuring groups in the output HEALPix dataset...")
cache.output_groups = _set_output_groups(
input_datatrees,
cache.input_spatial_groups,
settings,
cache.output_spatial,
)
log.info(
f"configured {len(cache.output_groups.multiscale_groups)} multiscale groups:\n"
+ "\n".join(f" - {path}" for path in cache.output_groups.multiscale_groups)
)
pyramid_group_lines = []
for path, info in cache.output_groups.pyramid_multiscale_groups.items():
levels = [
healpix.refinement_level for healpix, _ in info.derived_levels.values()
]
pyramid_group_lines.append(
f" - {path} (derived levels: {sorted(levels, reverse=True)})"
)
log.info(
f"configured {len(cache.output_groups.pyramid_multiscale_groups)} pyramid "
"multiscale groups:\n" + "\n".join(pyramid_group_lines)
)
log.info(
f"skip {len(cache.output_groups.excluded_groups)} group(s) in output dataset:\n"
+ "\n".join(f" - {path}" for path in cache.output_groups.excluded_groups)
)
log.info(
"output HEALPix dataset groups:\n"
+ "\n".join(f" - {path}" for path in cache.output_groups.io_path_map.values())
)
_check_consolidate_groups(cache.output_groups, settings)
log.info("••• finished creating conversion staging cache.")
return cache