zarr in forcingprocessor

[JordanLaserGit]

Forcingprocessor was originally designed to serve as the forcings generation for the NRDS system. This motivated the design choice to optimize for large domain (>10,000 catchments), short time (<1000 time steps) processing of NWM forcing files. Forcingprocessor's internal algorithm performs a multi processed read and processing of each netcdf required. This is performant at scale, but would be considered a network heavy solution in the context of small domain (<100 catchments) simulations. Also, forcingprocessor is not currently capable of processing anything other than the NWM forcing files provided by nwmurl.

From here, we can implement algorithms (that are already in related tooling) that will generalize forcingprocessor to use the most efficient internal algorithm based on the domain size and time step count. For example, if forcingprocessor starts up with a large geopackage and a few number of nwm files listed in filenamelist.txt, then the original algorithm is used. However if a small geopackage and a large number of nwm files, forcingprocessor could then switch to a zarr based algorithm that is more suitable in that case.

Action items:

1. Implement zarr algorithm in forcingprocessor with switch based on geopackage and filenamelist size.
2. Implement additional forcing sources like AORC, this may motivate creating another tool similar to nwmurl, but perhaps aorc_url or something more generally named so it could be used for any forcing source (i.e. forcing_url).

If these two are implemented, forcingprocessor could serve as a general purpose forcings tool that can be integrated in any workflow requiring forcings generation.

[quinnylee]

I was going to open an issue about this but I'm glad I looked in discussions first!

I'm running into writing the commands for the dHBV2 datastream right now. It's a small domain (44 catchments) requiring many timesteps (>8500 time steps), meaning that forcingprocessor is very slow and borderline unusable for the dHBV2. Given that we still want to run the dHBV2 datastream on an operational cadence, real-time, I wonder if switching to our kerchunked operational zarrs https://ciroh-nwm-zarr-copy.s3.amazonaws.com/ instead of the operational netCDFs would help speed things up? @jameshalgren suggested we could pull the standard analysis-assimilation 3-hour lookbacks for every hour and use the oldest hour from each file in the final preprocessed forcing input file. That is currently not how forcingprocessor is set up to work, but because dHBV2 requires inputs that are a non-standard forecast length, we would have to do some data wrangling anyway.

Open to thoughts from anyone @leoglonz @JordanLaserGit @harshavemula-ua @JoshCu

[JordanLaserGit]

Happy to talk with said student on how to implement the zarr feature in fp/nrds/aws.

[quinnylee]

@Racta-1 @pratikshac15

[JoshCu]

Unfortunately all kerchunk does is make it much faster to index the data and help reduce the dask task overhead. e.g. when you call xarray.open_dataset(). kerchunk saves us having a for loop that lazily opens every single file and then merging them before we have a continuous xarray dataset, but when downloading the data we're still stuck with whatever chunk format the underlying netcdf files use. To speed it up we'd need to reformat and keep our own copy of the data.

[toakinrele]

That makes sense,. It explains why the speedup plateaus where it does. The chunks are sized for the full grid, so even with kerchunk pointing us straight at the right chunk, we're still pulling and decompressing way more than we actually need. Given that, is maintaining a reformatted copy worth exploring as a follow-up?

[JoshCu]

Maintaining a copy of the raster data is probably too much for us to keep, but once it's been processed against the hydrofabric it's much smaller. we sort of have a copy here in the datastream bucket.
Depending on how much we want to store we could do some processing like below that would make it roughly 20x smaller than the raster data.

process.py

# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "dask>=2026.7.1",
#     "duckdb>=1.5.2",
#     "h5netcdf>=1.8.1",
#     "h5py>=3.16.0",
#     "pandas>=2.0",
#     "xarray>=2026.7.0",
# ]
# ///
import duckdb
import pandas as pd
import xarray as xr

rg_size = 10_000


def write_parquet(con, select_sql, path, rg_size, kv_metadata=None, compression_level=3):
  rg_clause = f", ROW_GROUP_SIZE {rg_size}" if rg_size else ""
  kv_clause = ""
  if kv_metadata:
      pairs = ", ".join(f"'{k}': '{v}'" for k, v in kv_metadata.items())
      kv_clause = f", KV_METADATA {{{pairs}}}"
  con.execute(f"""
      COPY ({select_sql})
      TO '{path}' (FORMAT PARQUET{rg_clause}, COMPRESSION zstd, COMPRESSION_LEVEL {compression_level}{kv_clause})
  """)


def preprocess(ds):
  ds = ds.set_coords("ids")
  ds = ds.swap_dims({"catchment-id": "ids"})
  ds = ds.rename(
      {"ids": "catchment_id"}
  )  # optional, avoids the hyphen (bad for dot-access/query)
  time_var = ds["Time"].values
  time_1d = time_var[0, :]
  time_index = pd.to_datetime(time_1d, unit="s")
  ds = ds.assign_coords(time=("time", time_index))
  ds = ds.drop_vars("Time")
  return ds


def main():
  ds = xr.open_mfdataset(
      "./forcings/*.nc", preprocess=preprocess, concat_dim="catchment_id", combine="nested"
  )
  ds.to_netcdf("conus_18h.nc")

  df = ds.to_dataframe().reset_index()

  con = duckdb.connect()
  con.register("raw", df)

  cols = con.execute("DESCRIBE SELECT * FROM raw").fetchall()
  select_list = ", ".join(
      f'"{name}"::REAL AS "{name}"' if dtype == "DOUBLE" else f'"{name}"'
      for name, dtype, *_ in cols
  )
  write_parquet(
      con,
      f'SELECT {select_list} FROM raw ORDER BY catchment_id, "time"',
      "catchment_timeseries.parquet",
      rg_size,
  )

  # (scale factor, value expression in native units -> converted units)
  scales = {
      "UGRD_10maboveground": (10, '"UGRD_10maboveground"'),
      "VGRD_10maboveground": (10, '"VGRD_10maboveground"'),
      "TMP_2maboveground": (10, '("TMP_2maboveground" - 273.15)'),  # K -> degC
      "PRES_surface": (10, '("PRES_surface" / 100)'),  # Pa -> hPa
      "SPFH_2maboveground": (10000, '"SPFH_2maboveground"'),
      "DLWRF_surface": (1, '"DLWRF_surface"'),
      "DSWRF_surface": (1, '"DSWRF_surface"'),
      "APCP_surface": (10, '"APCP_surface"'),
      "precip_rate": (100, '"precip_rate"'),
  }

  scaled_cols = ", ".join(
      f'CAST(ROUND({expr} * {scale}) AS SMALLINT) AS "{col}"'
      for col, (scale, expr) in scales.items()
  )

  write_parquet(
      con,
      f'SELECT catchment_id, "time", {scaled_cols} FROM raw ORDER BY catchment_id, "time"',
      "compressed_timeseries.parquet",
      rg_size,
      kv_metadata={f"scale_{k}": scale for k, (scale, _) in scales.items()},
  )


if __name__ == "__main__":
  main()

The forcings we have saved are formatted to work well in ngen so if were were going to backfill data then we should probably save them in a format best for storage space and querying