Add reordering WRF-Hydro channel-restart reader (get_channel_restart_from_wrf_hydro_reorder) selected via wrf_hydro_channel_restart_input_type

[jameshalgren]

Summary

Add an opt-in alternate reader for WRF-Hydro channel restart files that does not assume the restart file and the crosswalk file share an identical row order. The alternate reader reorders using the segment IDs found inside the restart file itself (link / links).

Selection is controlled by a single new string parameter read from the .yaml configuration file:

• wrf_hydro_channel_restart_input_type
  • absent (or any value other than "reorder") → current default behavior; nhd_io.get_channel_restart_from_wrf_hydro is called exactly as today.
  • "reorder" → the new nhd_io.get_channel_restart_from_wrf_hydro_reorder is called instead.

The new function takes the same argument list as the existing one, so no new file inputs are required. The original function is left untouched.

How we got here (process)

This work started from a re-written version of get_channel_restart_from_wrf_hydro (developed in an adjacent analysis repo, routing-comparison-nwm-nextgen-troute). Rather than edit code directly, we produced two planning documents:

1. A short driver prompt (prompt.md) whose only job was to request a comprehensive, hand-off-ready implementation plan.
2. That plan (insert_new_routing_function.md), suitable for a human or agent executor.

Those planning documents are intentionally not versioned — they capture the thought process, and their content is preserved here in this issue instead. The step-by-step implementation plan below is the actionable artifact.

Background: the difference between the two readers

The current get_channel_restart_from_wrf_hydro (src/troute-network/troute/nhd_io.py:1368) builds the index by reading segment IDs from the crosswalk file and assigning them to the restart rows by position:

with xr.open_dataset(crosswalk_file) as xds:
    xdf = xds[channel_ID_column].to_dataframe()
...
qdf2[channel_ID_column] = xdf          # <-- positional alignment
qdf2 = qdf2.reset_index().set_index([channel_ID_column])

This is correct only when the restart file rows and the crosswalk file rows are in the same order.

The rewritten reader instead reads the IDs from a link variable inside the restart file and indexes on those, so the crosswalk row order is irrelevant. It keeps the same function signature (the crosswalk_file and channel_ID_column arguments are accepted but unused), so it is a drop-in replacement.

Files touched (four)

1. src/troute-network/troute/nhd_io.py — add the new function.
2. src/troute-network/troute/AbstractNetwork.py — dispatch at call site Module reservoir #1.
3. src/troute-network/troute/nhd_network_utilities_v02.py — dispatch at call site Module reservoir #2.
4. src/troute-config/troute/config/compute_parameters.py — declare the parameter.

Config validation in config.py needs no change for the baseline (see companion issue on validation).

Step 1 — Add the new function to nhd_io.py

Insert immediately after the existing get_channel_restart_from_wrf_hydro (after its return, ~line 1430, before def read_lite_restart). Keep the original in place.

def get_channel_restart_from_wrf_hydro_reorder(
    channel_initial_states_file,
    crosswalk_file,
    channel_ID_column,
    us_flow_column="qlink1",
    ds_flow_column="qlink2",
    depth_column="hlink",
    default_us_flow_column="qu0",
    default_ds_flow_column="qd0",
    default_depth_column="h0",
):
    """
    channel_initial_states_file: restart file. The column containing IDs should be named "links" (or change restart_id_column in the function)
    crosswalk_file: doesn't get used
    channel_ID_column: doesn't get used
    us_flow_column: column in the restart file to use for upstream flow initial state
    ds_flow_column: column in the restart file to use for downstream flow initial state
    depth_column: column in the restart file to use for depth initial state
    default_us_flow_column: name used in remainder of program to refer to this column of the dataset
    default_ds_flow_column: name used in remainder of program to refer to this column of the dataset
    default_depth_column: name used in remainder of program to refer to this column of the dataset
    The Restart file gives hlink, qlink1, and qlink2 values for channels
    Unlike the original version, this function does NOT assume that the restart and the crosswalk file share an identical row order
    """

    restart_id_column = "links"  # name of column in restart file containing NWM IDs

    with xr.open_dataset(channel_initial_states_file) as qds:
        variables_to_keep = [us_flow_column, ds_flow_column]
        has_depth = depth_column in qds
        if has_depth:
            variables_to_keep.append(depth_column)

        if "link" in qds.variables and "link" not in variables_to_keep:
            variables_to_keep.append("link")

        qdf = qds[variables_to_keep].to_dataframe().reset_index()

    if "link" in qdf.columns and restart_id_column == "links":
        qdf[restart_id_column] = qdf["link"]

    if restart_id_column not in qdf.columns:
        raise KeyError(
            f"Could not find ID column '{restart_id_column}' in {channel_initial_states_file}. "
            f"Available columns: {qdf.columns.tolist()}"
        )

    if not has_depth:
        qdf[depth_column] = 0.0

    qdf.rename(
        columns={
            us_flow_column: default_us_flow_column,
            ds_flow_column: default_ds_flow_column,
            depth_column: default_depth_column,
        },
        inplace=True,
    )

    qdf[restart_id_column] = qdf[restart_id_column].astype('int64')

    final_columns = [restart_id_column, default_us_flow_column, default_ds_flow_column, default_depth_column]
    q_initial_states = qdf[final_columns].set_index(restart_id_column)

    try:
        q_initial_states.index = q_initial_states.index.astype(q_initial_states.index.dtype)
    except Exception:
        pass

    return q_initial_states

Notes:

• xr, pd, pathlib are already imported in nhd_io.py; do not add imports.
• Debug print statements from the source were intentionally removed.
• Returned DataFrame is indexed by segment ID with columns ["qu0", "qd0", "h0"], matching the original — downstream consumers need no changes.

Step 2 — Dispatch at call site #1 (AbstractNetwork.py, ~line 691)

Before:

            elif restart_parameters.get("wrf_hydro_channel_restart_file", None):
                self._q0 = nhd_io.get_channel_restart_from_wrf_hydro(
                    restart_parameters["wrf_hydro_channel_restart_file"],
                    restart_parameters["wrf_hydro_channel_ID_crosswalk_file"],
                    restart_parameters.get("wrf_hydro_channel_ID_crosswalk_file_field_name", 'link'),
                    restart_parameters.get("wrf_hydro_channel_restart_upstream_flow_field_name", 'qlink1'),
                    restart_parameters.get("wrf_hydro_channel_restart_downstream_flow_field_name", 'qlink2'),
                    restart_parameters.get("wrf_hydro_channel_restart_depth_flow_field_name", 'hlink'),
                    )

After:

            elif restart_parameters.get("wrf_hydro_channel_restart_file", None):
                if restart_parameters.get("wrf_hydro_channel_restart_input_type", None) == "reorder":
                    channel_restart_reader = nhd_io.get_channel_restart_from_wrf_hydro_reorder
                else:
                    channel_restart_reader = nhd_io.get_channel_restart_from_wrf_hydro
                self._q0 = channel_restart_reader(
                    restart_parameters["wrf_hydro_channel_restart_file"],
                    restart_parameters["wrf_hydro_channel_ID_crosswalk_file"],
                    restart_parameters.get("wrf_hydro_channel_ID_crosswalk_file_field_name", 'link'),
                    restart_parameters.get("wrf_hydro_channel_restart_upstream_flow_field_name", 'qlink1'),
                    restart_parameters.get("wrf_hydro_channel_restart_downstream_flow_field_name", 'qlink2'),
                    restart_parameters.get("wrf_hydro_channel_restart_depth_flow_field_name", 'hlink'),
                    )

The subsequent t0_str / self._t0 lines are unchanged.

Step 3 — Dispatch at call site #2 (nhd_network_utilities_v02.py, ~line 220)

Before:

    elif wrf_hydro_channel_restart_file:
        q0 = nhd_io.get_channel_restart_from_wrf_hydro(
            restart_parameters["wrf_hydro_channel_restart_file"],
            restart_parameters["wrf_hydro_channel_ID_crosswalk_file"],
            restart_parameters.get("wrf_hydro_channel_ID_crosswalk_file_field_name", 'link'),
            restart_parameters.get("wrf_hydro_channel_restart_upstream_flow_field_name", 'qlink1'),
            restart_parameters.get("wrf_hydro_channel_restart_downstream_flow_field_name", 'qlink2'),
            restart_parameters.get("wrf_hydro_channel_restart_depth_flow_field_name", 'hlink'),
        )

After:

    elif wrf_hydro_channel_restart_file:
        if restart_parameters.get("wrf_hydro_channel_restart_input_type", None) == "reorder":
            channel_restart_reader = nhd_io.get_channel_restart_from_wrf_hydro_reorder
        else:
            channel_restart_reader = nhd_io.get_channel_restart_from_wrf_hydro
        q0 = channel_restart_reader(
            restart_parameters["wrf_hydro_channel_restart_file"],
            restart_parameters["wrf_hydro_channel_ID_crosswalk_file"],
            restart_parameters.get("wrf_hydro_channel_ID_crosswalk_file_field_name", 'link'),
            restart_parameters.get("wrf_hydro_channel_restart_upstream_flow_field_name", 'qlink1'),
            restart_parameters.get("wrf_hydro_channel_restart_downstream_flow_field_name", 'qlink2'),
            restart_parameters.get("wrf_hydro_channel_restart_depth_flow_field_name", 'hlink'),
        )

Both call sites use .get(..., None) == "reorder", so an absent parameter and any non-"reorder" value fall through to the original reader — guaranteeing unchanged default behavior.

Step 4 — Declare the parameter in the config schema

src/troute-config/troute/config/compute_parameters.py, inside RestartParameters, after wrf_hydro_channel_restart_depth_flow_field_name (~line 123) and before wrf_hydro_waterbody_restart_file:

    wrf_hydro_channel_restart_input_type: Optional[str] = None
    """
    Selects which reader is used for the WRF-Hydro channel restart file.
    If unset (or any value other than 'reorder'), the default reader is used,
    which assumes the restart file and crosswalk file share the same row order.
    If set to 'reorder', an alternate reader is used that reorders restart
    values using the segment IDs stored inside the restart file itself, so the
    crosswalk file row order does not matter.
    """

Optional is already imported. Default None preserves current behavior when omitted.

Step 5 — Config validation (config.py) — no change for baseline

The existing validator check_wrf_hydro_restart_files (lines 218–235) already requires wrf_hydro_channel_ID_crosswalk_file whenever a channel restart file is given. The reorder reader ignores the crosswalk file, but keeping the requirement is harmless and symmetric. Stricter value-checking of the selector is deferred — see the companion validation issue.

Example .yaml snippet

compute_parameters:
  restart_parameters:
    wrf_hydro_channel_restart_file: /path/to/HYDRO_RST.YYYY-MM-DD_HH:MM_DOMAIN1
    wrf_hydro_channel_ID_crosswalk_file: /path/to/Route_Link.nc
    wrf_hydro_channel_ID_crosswalk_file_field_name: link
    wrf_hydro_channel_restart_upstream_flow_field_name: qlink1
    wrf_hydro_channel_restart_downstream_flow_field_name: qlink2
    wrf_hydro_channel_restart_depth_flow_field_name: hlink
    wrf_hydro_channel_restart_input_type: reorder   # <-- new; omit or use any other value for default behavior

Verification

A. Default behavior unchanged (regression): run an existing WRF-Hydro restart config without the new key and confirm identical output; confirm the config still parses; run uv run pytest src/troute-config.

B. New pathway works: add wrf_hydro_channel_restart_input_type: reorder and re-run; confirm q0 is populated with columns ["qu0","qd0","h0"] indexed by segment ID. With a matched-order restart file both readers should agree; with a mismatched-order file only the reorder reader aligns correctly.

C. Isolated smoke test:

from troute import nhd_io
q0 = nhd_io.get_channel_restart_from_wrf_hydro_reorder("<restart_file>", "<crosswalk_file>", "link")
print(q0.columns.tolist())   # ['qu0', 'qd0', 'h0']
print(q0.index.name)         # 'links'

Risks / edge cases

1. Hardcoded ID column. The reader hardcodes restart_id_column = "links", populated from a link variable. Restart files exposing IDs under another name raise a KeyError.
2. Unrecognized selector values fall silently to default (case-sensitive). See companion validation issue.
3. Index dtype. Reorder reader casts IDs to int64; original inherits dtype from the crosswalk — a small behavioral difference.
4. Crosswalk still required by validation even though reorder ignores it (intentional for symmetry).
5. Two call sites must stay in sync — dispatch logic is duplicated in AbstractNetwork.py and nhd_network_utilities_v02.py.

Driver prompt used (prompt.md)

Contents of prompt.md

The driver prompt requested creation of the comprehensive plan above (not source edits). Its key constraints:

• Deliverable = a step-by-step plan document; no source code modified during planning.
• Selection mechanism: one string parameter wrf_hydro_channel_restart_input_type read from the .yaml; absent → unchanged default; value "reorder" → new function.
• Rename the rewritten function to get_channel_restart_from_wrf_hydro_reorder, insert alongside the original at nhd_io.py:1368, same argument list, no new inputs.
• Wire the selector through both call sites (AbstractNetwork.py:692, nhd_network_utilities_v02.py:221).
• Register the parameter in compute_parameters.py (RestartParameters) and note validation in config.py.
• Preserve default behavior when the parameter is absent or non-"reorder".
• Plan must include final function body, config declaration, example .yaml, verification steps, and risks.