Feat/gribjump

hat/compute_hydrostats/stat_calc.py

@@ -1,17 +1,9 @@
-import earthkit.data as ekd
-from earthkit.hydro._readers import find_main_var
+from hat.core import load_da
 import numpy as np
 import xarray as xr
 from hat.compute_hydrostats import stats
 
 
-def load_da(ds_config):
-    ds = ekd.from_source(*ds_config["source"]).to_xarray()
-    var_name = find_main_var(ds, 2)
-    da = ds[var_name]
-    return da
-
-
 def find_valid_subset(sim_da, obs_da, sim_coords, obs_coords, new_coords):
     sim_station_colname = sim_coords.get("s", "station")
     obs_station_colname = obs_coords.get("s", "station")
@@ -35,9 +27,9 @@ def find_valid_subset(sim_da, obs_da, sim_coords, obs_coords, new_coords):
 
 def stat_calc(config):
     sim_config = config["sim"]
-    sim_da = load_da(config["sim"])
+    sim_da, _ = load_da(sim_config, 2)
     obs_config = config["obs"]
-    obs_da = load_da(obs_config)
+    obs_da, _ = load_da(obs_config, 2)
     new_coords = config["output"]["coords"]
     sim_da, obs_da = find_valid_subset(sim_da, obs_da, sim_config["coords"], obs_config["coords"], new_coords)
     stat_dict = {}

hat/core.py

@@ -0,0 +1,10 @@
+import earthkit.data as ekd
+from earthkit.hydro._readers import find_main_var
+
+
+def load_da(ds_config, n_dims):
+    src_name = list(ds_config["source"].keys())[0]
+    ds = ekd.from_source(src_name, **ds_config["source"][src_name]).to_xarray(**ds_config.get("to_xarray_options", {}))
+    var_name = find_main_var(ds, n_dims)
+    da = ds[var_name]
+    return da, var_name

hat/extract_timeseries/extractor.py

@@ -2,24 +2,17 @@
 import pandas as pd
 import xarray as xr
 import numpy as np
-import earthkit.data as ekd
-from earthkit.hydro._readers import find_main_var
+from hat.core import load_da
 
 from hat import _LOGGER as logger
 
 
 def process_grid_inputs(grid_config):
-    src_name = list(grid_config["source"].keys())[0]
-    logger.info(f"Processing grid inputs from source: {src_name}")
-    logger.debug(f"Grid config: {grid_config['source'][src_name]}")
-    ds = ekd.from_source(src_name, **grid_config["source"][src_name]).to_xarray(
-        **grid_config.get("to_xarray_options", {})
-    )
-    var_name = find_main_var(ds, 3)
-    da = ds[var_name]
+    da, var_name = load_da(grid_config, 3)
     logger.info(f"Xarray created from source:\n{da}\n")
-    gridx_colname = grid_config.get("coord_x", "lat")
-    gridy_colname = grid_config.get("coord_y", "lon")
+    coord_config = grid_config.get("coords", {})
+    gridx_colname = coord_config.get("x", "lat")
+    gridy_colname = coord_config.get("y", "lon")
     da = da.sortby([gridx_colname, gridy_colname])
     shape = da[gridx_colname].shape[0], da[gridy_colname].shape[0]
     return da, var_name, gridx_colname, gridy_colname, shape
@@ -61,7 +54,7 @@ def create_mask_from_coords(coords_config, df, gridx, gridy, shape):
     return mask, duplication_indexes
 
 
-def process_inputs(station_config, grid_config):
+def parse_stations(station_config):
     logger.debug(f"Reading station file, {station_config}")
     df = pd.read_csv(station_config["file"])
     filters = station_config.get("filter")
@@ -72,23 +65,39 @@ def process_inputs(station_config, grid_config):
 
     index_config = station_config.get("index", None)
     coords_config = station_config.get("coords", None)
+    index_1d_config = station_config.get("index_1d", None)
+    return index_config, coords_config, index_1d_config, station_names, df
+
+
+def process_inputs(station_config, grid_config):
+    index_config, coords_config, index_1d_config, station_names, df = parse_stations(station_config)
 
+    # TODO: better malformed config handling
     if index_config is not None and coords_config is not None:
         raise ValueError("Use either index or coords, not both.")
 
-    da, da_varname, gridx_colname, gridy_colname, shape = process_grid_inputs(grid_config)
-
-    if index_config is not None:
-        mask, duplication_indexes = create_mask_from_index(index_config, df, shape)
-    elif coords_config is not None:
-        mask, duplication_indexes = create_mask_from_coords(
-            coords_config, df, da[gridx_colname].values, da[gridy_colname].values, shape
-        )
+    if list(grid_config["source"].keys())[0] == "gribjump":
+        assert index_1d_config is not None
+        unique_indices, duplication_indexes = np.unique(df[index_1d_config].values, return_inverse=True)
+        grid_config["source"]["gribjump"]["indices"] = unique_indices
+        masked_da, da_varname = load_da(grid_config, 2)
     else:
-        # default to index approach
-        mask, duplication_indexes = create_mask_from_index(index_config, df, shape)
+        da, da_varname, gridx_colname, gridy_colname, shape = process_grid_inputs(grid_config)
+
+        if index_config is not None:
+            mask, duplication_indexes = create_mask_from_index(index_config, df, shape)
+        elif coords_config is not None:
+            mask, duplication_indexes = create_mask_from_coords(
+                coords_config, df, da[gridx_colname].values, da[gridy_colname].values, shape
+            )
+        else:
+            # default to index approach
+            mask, duplication_indexes = create_mask_from_index(index_config, df, shape)
 
-    return da, da_varname, gridx_colname, gridy_colname, mask, station_names, duplication_indexes
+        logger.info("Extracting timeseries at selected stations")
+        masked_da = apply_mask(da, mask, gridx_colname, gridy_colname)
+
+    return da_varname, station_names, duplication_indexes, masked_da
 
 
 def mask_array_np(arr, mask):
@@ -101,12 +110,12 @@ def apply_mask(da, mask, coordx, coordy):
         da,
         mask,
         input_core_dims=[(coordx, coordy), (coordx, coordy)],
-        output_core_dims=[["station"]],
+        output_core_dims=[["index"]],
         output_dtypes=[da.dtype],
         exclude_dims={coordx, coordy},
         dask="parallelized",
         dask_gufunc_kwargs={
-            "output_sizes": {"station": int(mask.sum())},
+            "output_sizes": {"index": int(mask.sum())},
             "allow_rechunk": True,
         },
     )
@@ -115,13 +124,10 @@ def apply_mask(da, mask, coordx, coordy):
 
 
 def extractor(config):
-    da, da_varname, gridx_colname, gridy_colname, mask, station_names, duplication_indexes = process_inputs(
-        config["station"], config["grid"]
-    )
-    logger.info("Extracting timeseries at selected stations")
-    masked_da = apply_mask(da, mask, gridx_colname, gridy_colname)
+    da_varname, station_names, duplication_indexes, masked_da = process_inputs(config["station"], config["grid"])
     ds = xr.Dataset({da_varname: masked_da})
-    ds = ds.isel(station=duplication_indexes)
+    ds = ds.isel(index=duplication_indexes)
+    ds = ds.rename({"index": "station"})
     ds["station"] = station_names
     if config.get("output", None) is not None:
         logger.info(f"Saving output to {config['output']['file']}")

hat/station_mapping/mapper.py

@@ -47,14 +47,15 @@ def apply_blacklist(blacklist_config, metric_grid, grid_area_coords1, grid_area_
     return metric_grid, grid_area_coords1, grid_area_coords2
 
 
-def outputs_to_df(df, indx, indy, cindx, cindy, errors, grid_area_coords1, grid_area_coords2, filename):
+def outputs_to_df(df, indx, indy, cindx, cindy, errors, grid_area_coords1, grid_area_coords2, shape, filename):
     df["opt_x_index"] = indx
     df["opt_y_index"] = indy
     df["near_x_index"] = cindx
     df["near_y_index"] = cindy
     df["opt_error"] = errors
     df["opt_x_coord"] = grid_area_coords1[indx, 0]
     df["opt_y_coord"] = grid_area_coords2[0, indy]
+    df["opt_1d_index"] = indy + shape[1] * indx
     if filename is not None:
         df.to_csv(filename, index=False)
     return df
@@ -109,6 +110,7 @@ def mapper(config):
         *mapping_outputs,
         grid_area_coords1,
         grid_area_coords2,
+        shape=grid_area_coords1.shape,
         filename=config["output"]["file"] if config.get("output", None) is not None else None,
     )
     generate_summary_plots(df, config.get("plot", None))

notebooks/workflow/hydrostats_computation.ipynb

@@ -19,14 +19,14 @@
    "source": [
     "config = {\n",
     "    \"sim\": {\n",
-    "        \"source\": [\"file\", \"extracted_timeseries.nc\"],\n",
+    "        \"source\": {\"file\": \"extracted_timeseries.nc\"},\n",
     "        \"coords\": {\n",
     "            \"s\": \"station\",\n",
     "            \"t\": \"time\"\n",
     "        }\n",
     "    },\n",
     "    \"obs\": {\n",
-    "        \"source\": [\"file\", \"observations.nc\"],\n",
+    "        \"source\": {\"file\": \"observations.nc\"},\n",
     "        \"coords\": {\n",
     "            \"s\": \"station\",\n",
     "            \"t\": \"time\"\n",
@@ -49,7 +49,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "hat",
    "language": "python",
    "name": "python3"
   },

notebooks/workflow/timeseries_extraction.ipynb

@@ -35,7 +35,7 @@
     "        \"name\": \"station_id\"\n",
     "    },\n",
     "    \"grid\": {\n",
-    "        \"source\": [\"file\", \"./sim.nc\"],\n",
+    "        \"source\": {\"file\": \"./sim.nc\"},\n",
     "        \"coords\": {\n",
     "            \"x\": \"lat\",\n",
     "            \"y\": \"lon\",\n",

pyproject.toml

@@ -42,7 +42,7 @@ dependencies = [
     "ipyleaflet",
     "ipywidgets",
     "earthkit-data>=0.13.8",
-    "earthkit-hydro",
+    "earthkit-hydro>=1.0.0",
     "earthkit-meteo",
     "cfgrib", # check if necessary
     "netCDF4", # check if necessary
@@ -68,6 +68,10 @@ dependencies = [
         "ruff",
         "pre-commit"
     ]
+    gribjump = [
+	"earthkit-data[gribjump]",
+	"gribjumplib==0.10.3.dev20250908"
+    ]
 
 [project.scripts]
     hat-extract-timeseries = "hat.cli:extractor_cli"