DSM2-1228 feat: add GTM restart file creation functionality and enhance related modules

Author: dwr-psandhu

pydsm/__init__.py

@@ -4,6 +4,14 @@
 
 __author__ = """Nicky Sandhu"""
 __email__ = "psandhu@water.ca.gov"
-__all__ = []
+__all__ = [
+    "__version__",
+    # GTM helper exports
+    "build_timewindow_for_time",
+    "get_interpolated_cell_concentrations",
+    "write_gtm_restart",
+]
 
 from ._version import __version__
+from .gtmh5 import build_timewindow_for_time, get_interpolated_cell_concentrations
+from .create_gtm_restart import write_gtm_restart

pydsm/cli.py

@@ -8,6 +8,7 @@
 from pydsm import repeating_timeseries
 from pydsm import create_cd_inp
 from pydsm import extend_dss_ts
+from pydsm import write_gtm_restart
 
 
 import pandas as pd
@@ -317,6 +318,28 @@ def csv_to_dss(
     )
 
 
+@click.command(name="create-gtm-restart")
+@click.argument("tidefile", type=click.Path(exists=True))
+@click.argument("target_time", type=str)
+@click.argument("outfile", type=click.Path())
+@click.option(
+    "-c",
+    "--constituent",
+    default="ec",
+    show_default=True,
+    help="Constituent to export (e.g. ec)",
+)
+def create_gtm_restart_cmd(tidefile, target_time, outfile, constituent):
+    """Create a GTM restart file from a GTM/Qual tide HDF5 file.
+
+    TIDEFILE: Path to HDF5 tide file.
+    TARGET_TIME: Desired time (e.g. '05FEB2020 0300'). Nearest model output time will be used.
+    OUTFILE: Destination restart file name.
+    """
+    path = write_gtm_restart(tidefile, target_time, outfile, constituent=constituent)
+    click.echo(f"Wrote restart file: {path}")
+
+
 # Add the commands to the group repeating
 repeating.add_command(create_repeating)
 repeating.add_command(extend_repeating)
@@ -330,6 +353,8 @@ def csv_to_dss(
 main.add_command(update_hydro_tidefile_with_inp)
 main.add_command(create_dsm2_input_for_cd)
 main.add_command(pretty_print_input)
+# GTM restart creation
+main.add_command(create_gtm_restart_cmd)
 #
 main.add_command(extend_dss_ts.extend_dss_ts)
 if __name__ == "__main__":

pydsm/create_gtm_restart.py

@@ -1,54 +1,173 @@
-#modify an existing gtm restart file to create a new restart file.
+"""Create a GTM restart file from a tide (Qual/GTM HDF5) file.
+
+Generates a restart file containing:
+ - Cell concentrations (interpolated per cell from channel end concentrations)
+ - Reservoir concentrations
+for a requested model time.
+
+Example
+-------
+>>> write_gtm_restart(
+...    tidefile="tests/data/hist_v82_mss2_extran_gtm.h5",
+...    target_time="05FEB2020 0300",
+...    outfile="restart_created.qrf",
+...    constituent="ec",
+... )
+
+Output file format (illustrative):
+30DEC2024 2400/time
+   65744640 /julmin
+         1 /n_column
+      4279 /n_cell
+                   cell_no                ec
+                        1            685.4853218050435544
+... (cells) ...
+         9 /n_resv
+              reservoir_name                ec
+clifton_court                               416.6467054602688336
+... (reservoirs) ...
+"""
+
+from __future__ import annotations
+
+from datetime import datetime, timedelta
+from pathlib import Path
+from typing import Iterable
 
 import numpy as np
 
-base_restart = "../restart_gtm1000_nx2.qrf"
-out_restart = "../restart_gtm1000_ones_nx2.qrt"
-
-def read_restart(restart_fn,extract_details=False):
-   with open(restart_fn,'r') as file:
-      filedata = file.read()
-   filetxt = filedata.split('\n')
-   ncells = int(filetxt[3].split('/')[0])
-   cell_header_str = '\n'.join(filetxt[0:5])
-   cell_str = filetxt[5:5+ncells]
-   cell_array = np.array([l.split() for l in cell_str])
-   cell_name = cell_array[:,0]
-   cell_value = cell_array[:,1]
-   nreser =  int(filetxt[5+ncells].split('/')[0])
-   res_header_str =  '\n'.join(filetxt[5+ncells:5+ncells+2])
-   res_str = filetxt[5+ncells+2:5+ncells+2+nreser]
-   res_array = np.array([l.split() for l in res_str])
-   res_name =  res_array[:,0]
-   res_value = res_array[:,1]
-   if extract_details: 
-      return cell_name, res_name, ncells, nreser,cell_header_str,res_header_str
-   else:
-      field = np.append(cell_value,res_value)
-      return field
-
-
-def write_restart(field,restart_fn,cell_name, res_name, 
-     cell_header_str, res_header_str):
-   """
-   read from a pdaf field_fn and write into a gtm restart_fn
-   """   
-   #field = read_pdaf_field(field_fn)
-   ncells = len(cell_name)
-   nres = len(res_name)
-   cell_list = ['%32s%32s'%(cid,cv) for cid, cv in zip(cell_name, field[:ncells])]
-   cell_str = '\n'.join(cell_list)
-   res_list = [cid.ljust(32) + '%32s'%cv for cid, cv in zip(res_name,field[ncells:ncells+nres])]
-   res_list = '\n'.join(res_list)
-   field_str = '\n'.join([cell_header_str,cell_str,res_header_str,res_list])
-   with open(restart_fn,'w') as file:
-      file.write(field_str)
-
-# Get basic structure from an example restart file. 
-cell_name, res_name, ncells, nreser,cell_header_str,res_header_str = read_restart(base_restart,extract_details=True) 
-field = read_restart(base_restart)
-field_new = np.ones_like(field)   # all one initial condition 
-#field_zero = (field.astype(float)*0.0).astype(str)
-write_restart(field_new,out_restart,cell_name,res_name,
-   cell_header_str, res_header_str)
+from .gtmh5 import (
+    build_timewindow_for_time,
+    get_interpolated_cell_concentrations,
+    _format_time,
+)
+from .qualh5 import QualH5
+
+_EPOCH_JULMIN = datetime(1899, 12, 31, 0, 0)  # Excel-like epoch (matches sample)
+
+
+def _time_label_with_2400(dt: datetime) -> str:
+    """Format time similar to DSM2/expected restart label.
+
+    If the time is exactly on midnight (00:00) we represent it as previous day 2400.
+    """
+    if dt.hour == 0 and dt.minute == 0:
+        prev = dt - timedelta(days=1)
+        return prev.strftime("%d%b%Y").upper() + " 2400"
+    return dt.strftime("%d%b%Y %H%M").upper()
+
+
+def _compute_julmin(dt: datetime) -> int:
+    """Compute julian minutes relative to the _EPOCH_JULMIN baseline.
+
+    This matches the magnitude of values seen in existing restart samples.
+    """
+    return int((dt - _EPOCH_JULMIN).total_seconds() // 60)
+
+
+def _format_cell_header(constituent: str) -> str:
+    return f"{'cell_no':>29}{constituent:>20}              "
+
+
+def _format_res_header(constituent: str) -> str:
+    return f"{'reservoir_name':>29}{constituent:>20}              "
+
+
+def write_gtm_restart(
+    tidefile: str | Path,
+    target_time: str | datetime,
+    outfile: str | Path,
+    constituent: str = "ec",
+) -> Path:
+    """Write a GTM restart file at the requested time.
+
+    Parameters
+    ----------
+    tidefile : str | Path
+       Path to GTM/Qual HDF5 file (tide file).
+    target_time : str | datetime
+       Desired time (will snap to nearest output time).
+    outfile : str | Path
+       Destination restart file path.
+    constituent : str, default 'ec'
+       Constituent to export.
+
+    Returns
+    -------
+    Path to the written restart file.
+    """
+    tidefile = Path(tidefile)
+    outfile = Path(outfile)
+    qualt = QualH5(str(tidefile))
+
+    # Determine timewindow and chosen model time
+    timewindow, model_time = build_timewindow_for_time(str(tidefile), target_time)
+
+    # End of interval (label uses interval end so that 2400 formatting matches sample)
+    output_freq = qualt.get_output_freq()
+    interval_end = model_time + output_freq
+
+    # Gather cell concentrations
+    cell_conc = get_interpolated_cell_concentrations(
+        str(tidefile), timewindow, constituent=constituent
+    ).ravel()
+    n_cells = cell_conc.size
+
+    # Gather reservoir concentrations
+    res_df = qualt.get_reservoirs()
+    res_names = list(res_df["name"].values)
+    if res_names:
+        res_conc_df = qualt.get_reservoir_concentration(
+            constituent, res_names, timewindow=timewindow
+        )
+        # Single row expected
+        res_conc = res_conc_df.iloc[0].values.astype(float)
+    else:  # pragma: no cover - unlikely empty dataset
+        res_conc = np.array([])
+    n_res = len(res_conc)
+
+    # Compose header
+    time_label = _time_label_with_2400(interval_end)
+    julmin = _compute_julmin(interval_end)
+
+    lines: list[str] = []
+    lines.append(f"{time_label}/time")
+    lines.append(f"{julmin:12d} /julmin")
+    lines.append(f"{1:12d} /n_column")
+    lines.append(f"{n_cells:12d} /n_cell")
+    lines.append(_format_cell_header(constituent))
+
+    # Cell data lines: id (1-based) and value
+    for i, val in enumerate(cell_conc, start=1):
+        # Align cell number under 'cell_no' and value under constituent header
+        lines.append(f"{i:32d}{val:32.16f}")
+
+    lines.append(f"{n_res:12d} /n_resv")
+    lines.append(_format_res_header(constituent))
+
+    for name, val in zip(res_names, res_conc):
+        lines.append(f"{name:<32}{val:>32.16f}")
+
+    outfile.write_text("\n".join(lines) + "\n")
+    return outfile
+
+
+__all__ = ["write_gtm_restart"]
+
+
+if __name__ == "__main__":  # rudimentary CLI usage
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Create GTM restart file")
+    parser.add_argument("tidefile", help="Path to GTM/Qual tide HDF5 file")
+    parser.add_argument("target_time", help="Target time (e.g. '05FEB2020 0300')")
+    parser.add_argument("outfile", help="Output restart file path")
+    parser.add_argument(
+        "-c", "--constituent", default="ec", help="Constituent name (default ec)"
+    )
+    args = parser.parse_args()
 
+    path = write_gtm_restart(
+        args.tidefile, args.target_time, args.outfile, constituent=args.constituent
+    )
+    print(f"Wrote restart file: {path}")

pydsm/dsm2h5.py

@@ -153,7 +153,7 @@ def read_table_attrs_as_df(filename, table_path):
     """
     with h5py.File(filename, "r") as f:
         bf = f[table_path]
-        a = pd.DataFrame(bf.attrs.items(), columns=["Name", "Value"], dtype=np.str)
+        a = pd.DataFrame(bf.attrs.items(), columns=["Name", "Value"], dtype=np.str_)
         a = a.append(
             pd.DataFrame([("shape", str(bf.shape))], columns=["Name", "Value"]),
             ignore_index=True,