Adding basic virtual zarr usage for harmonic parameters

setup.cfg

@@ -72,12 +72,21 @@ app =
     hvplot
     geoviews
     datashader
+vzarr =
+    tqdm
+    imagecodecs
+    zarr<3
+    fastparquet
+    aiohttp
+    tifffile
+    kerchunk
 all =
     %(test)s
     %(remote)s
     %(docs)s
     %(dev)s
     %(app)s
+    %(vzarr)s
 
 [options.entry_points]
 console_scripts =

src/dask_flood_mapper/harmonic_params.py

@@ -4,9 +4,41 @@
 
 import numpy as np
 import xarray as xr
+from dask_flood_mapper.processing import order_orbits
 from numba import njit, prange
 
-from dask_flood_mapper.processing import order_orbits
+
+def create_harmonic_parameters_zarr(
+    sig0_dc: xr.Dataset,
+    min_nobs: int = 32,
+    k: int = 3,
+):
+    param_names = model_coords(k)
+    template = (
+        sig0_dc.isel(obs=slice(len(param_names)))
+        .rename({"obs": "param"})
+        .drop_vars("time")
+    )
+    template["param"] = param_names
+    hpar_dc = xr.map_blocks(
+        reduce_ds_to_harmonic_parameters,
+        obj=sig0_dc,
+        kwargs={
+            "fit_var_name": "sig0",
+            "k": k,
+            "x_var_name": "X",
+            "y_var_name": "Y",
+            "min_nobs": min_nobs,
+        },
+        template=template,
+    )
+    hpar_dc = hpar_dc.rename({"sig0": "harmonic_parameters"})
+    hpar_dc = hpar_dc.where(hpar_dc.sel(param="NOBS") >= min_nobs).drop_sel(
+        param="NOBS"
+    )
+    hpar_dc = hpar_dc.harmonic_parameters.to_dataset(dim="param")
+
+    return hpar_dc
 
 
 def create_harmonic_parameters(
@@ -65,24 +97,52 @@ def process_harmonic_parameters_datacube(
     return sig0_dc, hpar_dc, orbit_sig0
 
 
+def reduce_ds_to_harmonic_parameters(
+    ts_ds: xr.Dataset, fit_var_name: str, min_nobs: int = 0, **kwargs
+):
+    extra_dims = [dim for dim in ts_ds.dims if dim not in ts_ds.squeeze().dims]
+    ts_xr = ts_ds[fit_var_name]
+
+    # if all pixels have too few observations, skip the regression and return all NaNs
+    too_few_obs_short_circuit = ts_xr.count(dim="obs").max().values < min_nobs
+    ts_dtimes = ts_ds["time.dayofyear"].squeeze(drop=True).values
+    if too_few_obs_short_circuit:
+        ts_xr = ts_xr * np.nan
+    out_dataarray = reduce_to_harmonic_parameters(
+        ts_xr.squeeze(drop=True), dtimes=ts_dtimes, **kwargs
+    )
+    out_dataset = xr.Dataset(
+        {
+            fit_var_name: out_dataarray.expand_dims(dim=extra_dims).transpose(
+                *ts_xr.rename({"obs": "param"}).dims
+            )
+        },
+        coords={
+            dim: ts_ds[dim]
+            for dim in ts_ds.dims
+            if (dim in extra_dims or dim in out_dataarray.dims) and dim in ts_ds.coords
+        },
+    )
+    return out_dataset
+
+
 def reduce_to_harmonic_parameters(
     ts_xr: xr.DataArray,
-    dtimes: np.ndarray,
     x_var_name: str = "x",
     y_var_name: str = "y",
     **kwargs,  # noqa: ANN003
-) -> xr.DataArray:
-    """Reduce a time series to harmonic parameters."""
-    params_arr = harmonic_regression(ts_xr.data, dtimes=dtimes, **kwargs)
+):
+    params_arr = harmonic_regression(ts_xr.values, **kwargs)
     k: int = kwargs.get("k", 3)
     out_dims: list[str] = ["param", y_var_name, x_var_name]
+    coords_dict = {"param": model_coords(k)}
+    if x_var_name in ts_xr.coords:
+        coords_dict[x_var_name] = ts_xr[x_var_name]
+    if y_var_name in ts_xr.coords:
+        coords_dict[y_var_name] = ts_xr[y_var_name]
     return xr.DataArray(
         data=params_arr,
-        coords={
-            "param": model_coords(k),
-            x_var_name: ts_xr[x_var_name],
-            y_var_name: ts_xr[y_var_name],
-        },
+        coords=coords_dict,
         dims=out_dims,
     )
 
@@ -100,6 +160,11 @@ def harmonic_regression(
     # should be in dayofyear format
     t = dtimes
 
+    # drop t and arr where t is nan for efficiency in regression
+    valid_time = ~np.isnan(t)
+    t = t[valid_time]
+    arr = arr[valid_time, ...]  # type: ignore
+
     # prepare A-matrix
     num_dims: int = 3
     if len(arr.shape) != num_dims:
@@ -115,13 +180,11 @@ def harmonic_regression(
 
     # run regression
     param = np.full((nx + 2, rows, cols), np.nan, dtype=np.float32)
-    _fast_harmonic_regression(
-        arr=arr,
-        a_matrix=a,
-        k=k,
-        red=redundancy,
-        param=param,
-    )
+    arr = arr.astype(np.float32)
+    if np.all(np.isnan(arr)):
+        # All NaN array, return NaN params
+        return param
+    _fast_harmonic_regression(arr=arr, a_matrix=a, k=k, red=redundancy, param=param)
     return param
 
 

src/dask_flood_mapper/vzarr/__init__.py

@@ -0,0 +1 @@
+from .read import open_s1_datacube  # noqa