Utilities updates and JRC flood conventions

src/hazard/onboard/jrc_river_flood.py

@@ -124,7 +124,10 @@ def onboard(self, target):
                 height,
                 transform,
                 str(crs),
+                index_name="return_period",
                 index_values=self.return_periods,
+                index_units="years",
+                units=self._resource.units,
             )
 
             # Process each return period

src/hazard/sources/osc_zarr.py

@@ -102,8 +102,14 @@ def create_empty(
         index_name: Optional[str] = "index",
         index_values: Any = None,
         chunks=None,
+        index_units: str = "",
+        units: Optional[str] = None,
     ):
-        """Create an empty Zarr array with specified dimensions and geospatial properties."""
+        """Create an empty Zarr array with specified dimensions and geospatial properties.
+
+        index_units are the units of the index (non-spatial) dimension values, e.g. "years" for
+        return periods; units are the units of the array values themselves, e.g. "metres".
+        """
         if index_values is None:
             index_values = [0]
         return self._zarr_create(
@@ -115,6 +121,8 @@ def create_empty(
             index_name=index_name,
             index_values=index_values,
             chunks=chunks,
+            index_units=index_units,
+            units=units,
         )
 
     def remove(self, path: str):
@@ -384,6 +392,8 @@ def _zarr_create(
         index_name: Optional[str] = None,
         index_values: Any = None,
         chunks=None,
+        index_units: str = "",
+        units: Optional[str] = None,
     ):
         """Create Zarr array with given shape and affine transform."""
         if path in self.root:
@@ -419,10 +429,43 @@ def _zarr_create(
         ]:
             index_values = [int(i) for i in index_values]
         self._add_attributes(
-            z.attrs, transform, crs, index_name=index_name, index_values=index_values
+            z.attrs,
+            transform,
+            crs,
+            index_name=index_name,
+            index_values=index_values,
+            index_units=index_units,
+            units=units,
         )
         return z
 
+    def rename_index_dimension(
+        self,
+        path: str,
+        index_name: str,
+        index_units: str = "",
+        units: Optional[str] = None,
+    ):
+        """Rename the index (non-spatial) dimension of an existing array, in place.
+
+        Only the attributes are changed: the index values are moved to "{index_name}_values" and
+        "{index_name}_units" is set; if given, units (of the array values) is also set. The array
+        data is not rewritten.
+        """
+        z = self.root[path]
+        attrs = z.attrs.asdict()
+        dimensions = attrs["dimensions"]
+        old_name = dimensions[0]
+        index_values = attrs.pop(f"{old_name}_values", None)
+        attrs.pop(f"{old_name}_units", None)
+        attrs["dimensions"] = [index_name] + list(dimensions[1:])
+        if index_values is not None:
+            attrs[f"{index_name}_values"] = index_values
+        attrs[f"{index_name}_units"] = index_units
+        if units is not None:
+            attrs["units"] = units
+        z.attrs.put(attrs)
+
     def _add_attributes(
         self,
         attrs: Dict[str, Any],
@@ -431,6 +474,7 @@ def _add_attributes(
         index_name: Optional[str] = None,
         index_values: Any = None,
         index_units: str = "",
+        units: Optional[str] = None,
     ):
         trans_members = [
             transform.a,
@@ -450,6 +494,8 @@ def _add_attributes(
         if index_values is not None:
             attrs[f"{index_name}_values"] = list(index_values)
             attrs[f"{index_name}_units"] = index_units
+        if units is not None:
+            attrs["units"] = units
 
     def _chunks(self, index_values: List[Any]):
         chunk_dim = 1000 if len(index_values) < 10 else 500

src/hazard/utilities/s3_utilities.py

@@ -13,7 +13,7 @@
 import os
 import pathlib
 import sys
-from typing import Callable, List, Optional, Sequence, Set, Tuple
+from typing import Callable, Iterable, List, Optional, Sequence, Set, Tuple, Union
 
 try:
     ## python >3.11 includes tomllib
@@ -378,23 +378,25 @@ def copy_local_to_prod(zarr_dir: str, array_path: str, dry_run=False):
 
 def _sync_local_to_bucket(
     zarr_dir: str,
-    array_path: str,
+    array_paths: Union[str, Iterable[str]],
     s3_target_client,
     target_bucket_name: str,
     zarr_group_prefix: str,
     dry_run: bool = True,
 ) -> Tuple[Set[str], Set[str], Set[str]]:
-    """Synchronize a local Zarr array to the given bucket based on ETag differences.
+    """Synchronize local Zarr arrays to the given bucket based on ETag differences.
 
     Like sync_buckets, but the source is a local directory rather than another S3
     bucket: only files that are missing from the bucket, or whose ETag differs from
-    the local file's computed ETag (see _local_etag), are uploaded.
+    the local file's computed ETag (see _local_etag), are uploaded. The total size of
+    the files to copy, across all array paths, is logged.
 
     Shared implementation for sync_local_to_dev and sync_local_to_prod.
 
     Args:
         zarr_dir (str): Directory of the Zarr group, i.e. /<path>/hazard/hazard.zarr.
-        array_path (str): The path of the array or data set within the group.
+        array_paths (Union[str, Iterable[str]]): The path(s) of the arrays or data sets
+        within the group.
         s3_target_client: boto3 S3 client to sync to.
         target_bucket_name (str): Name of the target bucket.
         zarr_group_prefix (str): Bucket-side location of the zarr group (see
@@ -408,54 +410,72 @@ def _sync_local_to_bucket(
         missing is present locally but not in the bucket; different is present in
         both but with differing content.
     """
+    if isinstance(array_paths, str):
+        array_paths = [array_paths]
     root_path = pathlib.Path(zarr_dir)
-    relative_file_paths = _local_zarr_relative_file_paths(root_path, array_path)
 
     def key(path: str) -> str:
         return _hazard_zarr_key(path, zarr_group_prefix)
 
-    prefix = key(array_path)
-    target_etags = list_object_etags(s3_target_client, target_bucket_name, prefix)
-
-    checked = set(relative_file_paths)
-    missing = set(path for path in relative_file_paths if key(path) not in target_etags)
-    different = set(
-        path
-        for path in relative_file_paths
-        if path not in missing
-        and target_etags[key(path)] != _local_etag(root_path / path)
-    )
+    all_checked: Set[str] = set()
+    all_missing: Set[str] = set()
+    all_different: Set[str] = set()
+    for array_path in array_paths:
+        relative_file_paths = _local_zarr_relative_file_paths(root_path, array_path)
+        prefix = key(array_path)
+        target_etags = list_object_etags(s3_target_client, target_bucket_name, prefix)
 
-    logger.info(
-        f"Checked {len(checked)} local files from {zarr_dir}/{array_path} against bucket {target_bucket_name}"
-    )
-    logger.info(
-        f"Copying {len(missing)} missing files to {target_bucket_name}: "
-        + _first_5_last_5(list(missing))
-    )
-    logger.info(
-        f"Copying {len(different)} different files to {target_bucket_name}: "
-        + _first_5_last_5(list(different))
-    )
+        checked = set(relative_file_paths)
+        missing = set(
+            path for path in relative_file_paths if key(path) not in target_etags
+        )
+        different = set(
+            path
+            for path in relative_file_paths
+            if path not in missing
+            and target_etags[key(path)] != _local_etag(root_path / path)
+        )
 
-    if not dry_run:
-        _copy_local_to_bucket(
-            zarr_dir,
-            array_path,
-            s3_target_client,
-            target_bucket_name,
-            zarr_group_prefix,
-            dry_run=False,
-            relative_file_paths=list(missing | different),
+        logger.info(
+            f"Checked {len(checked)} local files from {zarr_dir}/{array_path} against bucket {target_bucket_name}"
+        )
+        logger.info(
+            f"Copying {len(missing)} missing files to {target_bucket_name}: "
+            + _first_5_last_5(list(missing))
         )
+        logger.info(
+            f"Copying {len(different)} different files to {target_bucket_name}: "
+            + _first_5_last_5(list(different))
+        )
+        logger.info(
+            f"Size of files to copy: {format_size(local_files_size(zarr_dir, missing | different))}"
+        )
+
+        if not dry_run:
+            _copy_local_to_bucket(
+                zarr_dir,
+                array_path,
+                s3_target_client,
+                target_bucket_name,
+                zarr_group_prefix,
+                dry_run=False,
+                relative_file_paths=list(missing | different),
+            )
+
+        all_checked |= checked
+        all_missing |= missing
+        all_different |= different
 
-    return checked, missing, different
+    logger.info(
+        f"Total size of files to copy: {format_size(local_files_size(zarr_dir, all_missing | all_different))}"
+    )
+    return all_checked, all_missing, all_different
 
 
 def sync_local_to_dev(
-    zarr_dir: str, array_path: str, dry_run: bool = True
+    zarr_dir: str, array_paths: Union[str, Iterable[str]], dry_run: bool = True
 ) -> Tuple[Set[str], Set[str], Set[str]]:
-    """Synchronize a local Zarr array to the development bucket based on ETag
+    """Synchronize local Zarr arrays to the development bucket based on ETag
     differences, only uploading files that are missing or different.
     Requires environment variables:
     OSC_S3_BUCKET_DEV=physrisk-hazard-indicators-dev01
@@ -464,7 +484,8 @@ def sync_local_to_dev(
 
     Args:
         zarr_dir (str): Directory of the Zarr group, i.e. /<path>/hazard/hazard.zarr.
-        array_path (str): The path of the array or data set within the group.
+        array_paths (Union[str, Iterable[str]]): The path(s) of the arrays or data sets
+        within the group.
         dry_run (bool, optional): If True, log actions without executing them. Defaults to True.
 
     Returns:
@@ -476,7 +497,7 @@ def sync_local_to_dev(
     )
     return _sync_local_to_bucket(
         zarr_dir,
-        array_path,
+        array_paths,
         s3_target_client,
         target_bucket_name,
         zarr_group_prefix,
@@ -485,9 +506,9 @@ def sync_local_to_dev(
 
 
 def sync_local_to_prod(
-    zarr_dir: str, array_path: str, dry_run: bool = True
+    zarr_dir: str, array_paths: Union[str, Iterable[str]], dry_run: bool = True
 ) -> Tuple[Set[str], Set[str], Set[str]]:
-    """Synchronize a local Zarr array to the production (public) bucket based on
+    """Synchronize local Zarr arrays to the production (public) bucket based on
     ETag differences, only uploading files that are missing or different.
 
     Note the bucket-side zarr group is at "hazard-indicators/hazard.zarr", not
@@ -501,7 +522,8 @@ def sync_local_to_prod(
 
     Args:
         zarr_dir (str): Directory of the Zarr group, i.e. /<path>/hazard/hazard.zarr.
-        array_path (str): The path of the array or data set within the group.
+        array_paths (Union[str, Iterable[str]]): The path(s) of the arrays or data sets
+        within the group.
         dry_run (bool, optional): If True, log actions without executing them. Defaults to True.
 
     Returns:
@@ -513,7 +535,7 @@ def sync_local_to_prod(
     )
     return _sync_local_to_bucket(
         zarr_dir,
-        array_path,
+        array_paths,
         s3_target_client,
         target_bucket_name,
         zarr_group_prefix,
@@ -836,6 +858,21 @@ def _first_5_last_5(items):
     return ", ".join(items[0:5]) + ", ..., " + ", ".join(items[-5:-1])
 
 
+def local_files_size(zarr_dir: str, relative_file_paths: Iterable[str]) -> int:
+    """Total size in bytes of the given files, as paths relative to zarr_dir."""
+    root_path = pathlib.Path(zarr_dir)
+    return sum((root_path / path).stat().st_size for path in relative_file_paths)
+
+
+def format_size(num_bytes: float) -> str:
+    """Format a size in bytes for display, e.g. 1536 -> '1.5 KiB'."""
+    for unit in ["B", "KiB", "MiB", "GiB"]:
+        if abs(num_bytes) < 1024:
+            return f"{num_bytes:.1f} {unit}" if unit != "B" else f"{int(num_bytes)} B"
+        num_bytes /= 1024
+    return f"{num_bytes:.1f} TiB"
+
+
 def get_input_onboarding_data(download_dir=None):
     """Download inputs for IRISIndicator and Jupiter hazards for the onboarding from the production bucket.
 

src/hazard/utilities/tiles.py

@@ -276,6 +276,8 @@ def create_tile_set(
         chunk_size,
         index_name=str(da.dims[0]),
         index_values=index_values,
+        index_units=da[da.dims[0]].attrs.get("units", ""),
+        units=da.attrs.get("units"),
     )
 
     # here we reproject and write the maximum zoom level
@@ -311,6 +313,8 @@ def _create_empty_tile_pyramid(
     chunk_size: int,
     index_name: str = "index",
     index_values: Sequence[Any] = [0],
+    index_units: str = "",
+    units: Optional[str] = None,
 ):
     pixels_per_tile: int = 256
     ulx, uly = mercantile.xy(*mercantile.ul(mercantile.Tile(x=0, y=0, z=0)))
@@ -335,6 +339,8 @@ def _create_empty_tile_pyramid(
             index_name=index_name,
             index_values=index_values,
             chunks=(1, chunk_size, chunk_size),
+            index_units=index_units,
+            units=units,
         )
 
 

src/hazard/utilities/xarray_utilities.py

@@ -177,6 +177,10 @@ def data_array_from_zarr(z: zarr.array) -> xr.DataArray:
     else:
         da.rio.write_crs(crs, inplace=True)
         da = da.rename({"dim_0": index_name, "dim_1": "y", "dim_2": "x"})
+    if index_name + "_units" in z.attrs:
+        da[index_name].attrs["units"] = z.attrs[index_name + "_units"]
+    if "units" in z.attrs:
+        da.attrs["units"] = z.attrs["units"]
     return da
 
 

tests/test_onboarding.py

@@ -68,7 +68,7 @@ def onboarding_output_dir():
 def test_onboard_named_hazard_local(test_output_dir):
     """Used to onboard into a local directory a named hazard."""
     hazard_to_onboard = "JRCRiverFlood"  # "DavydzenkaEtAlLandSubsidence" # "JRCRiverFlood" # "WRIAqueductWaterRisk"
-    maps_only = True  # True to skip prepare/onboard and just regenerate maps
+    maps_only = False  # True to skip prepare/onboard and just regenerate maps
     # from array data already in test_output_dir (e.g. while iterating on map code).
 
     store = zarr.DirectoryStore(
@@ -101,9 +101,11 @@ def test_sync_named_hazard(test_output_dir, s3_credentials):
     Run test_onboard_named_hazard first so that the data this looks for is actually
     present in test_output_dir.
     """
-    hazard_to_onboard = "DavydzenkaEtAlLandSubsidence"
-    dry_run = True
-    maps_only = True  # True to sync only the map arrays (e.g. after regenerating maps).
+    hazard_to_onboard = "JRCRiverFlood"
+    dry_run = False
+    maps_only = (
+        False  # True to sync only the map arrays (e.g. after regenerating maps).
+    )
 
     zarr_dir = os.path.join(test_output_dir, "hazard-indicators", "hazard.zarr")
 
@@ -122,8 +124,7 @@ def test_sync_named_hazard(test_output_dir, s3_credentials):
         if resource.map is not None:
             array_paths.add(str(PurePosixPath(resource.map.path).parent))
 
-    for array_path in sorted(array_paths):
-        s3_utilities.sync_local_to_prod(zarr_dir, array_path, dry_run=dry_run)
+    s3_utilities.sync_local_to_prod(zarr_dir, sorted(array_paths), dry_run=dry_run)
 
 
 @pytest.mark.manual