"""Module for fetching forecast and observation data based on the provided configuration.
Functions:
- get_fcst_info: Get the forecast window size, time step, and reference time for a given NWM configuration.
- check_existing_obs_data: Check existing parquet files of USGS observations and get the dates for previously downloaded data.
- check_missing_obs_data: Check for missing observation data in the specified directory.
- safe_fetch_usgs: Fetch USGS streamflow data for a list of gage IDs and a date range, and save to parquet files.
- get_txdot_gage_list: Get the list of gage IDs for retrieving TxDOT streamflow observations.
- process_txdot_data: Process raw TxDOT streamflow data for a specific gage ID.
- fetch_txdot_gage_data: Fetch streamflow data for a list of TxDOT gage IDs and a date range, and resample temporally as needed.
- retrieve_usgs_obs: Retrieve USGS streamflow observations given configuration and a list of gage IDs.
- get_fcst_files: Get the list of forecast files for a given dataset.
- retrieve_fcsts_ngencerf: Retrieve NWM forecasts given the configurations and list of locations from the ngenCERF data source.
- retrieve_fcsts_gcs: Retrieve NWM forecasts from Google Cloud Storage based on the provided configuration and list of locations.
- retrieve_fcsts: Retrieve NWM forecast data for the specified locations and configuration.
- get_obs_files: Get the list of existing observation files.
- read_obs_data: Read observation data from a list of file paths (in csv or parquet format) and concatenate them into a single DataFrame.
"""
import gc
import glob
import os
import sys
import warnings
from pathlib import Path
import aiohttp
import dask
import dataretrieval.nwis as nwis
import pandas as pd
import teehr.loading.nwm.nwm_points as tlp
import xarray as xr
from aiohttp.client_exceptions import ClientOSError, ServerDisconnectedError
from dask.distributed import ( # install with 'pip install dask[complete]'
Client,
LocalCluster,
)
from pandas.api.types import is_numeric_dtype, is_object_dtype, is_string_dtype
from teehr.loading.usgs.usgs import usgs_to_parquet
from .identify_location_ids import find_locations_in_crosswalk
from .nwm_configs import ForecastConfig
from .settings import default_txdot_gage_list
from .utils import create_time_sequence, get_n_workers, read_data, save_data
warnings.filterwarnings("ignore", message="Compute Engine Metadata server unavailable")
import logging
logger = logging.getLogger(__name__)
# Disable worker profiling
dask.config.set({"distributed.worker.profile.enabled": False})
[docs]
def get_fcst_info(conf: dict) -> tuple[int, int, list[pd.Timestamp]]:
"""Get the forecast window size, time step, and reference time for a given NWM configuration.
Args:
conf: Dictionary containing NWM configuration.
Returns:
Tuple containing forecast window size, time step, and list of reference times.
"""
# validate forecast cycle
nwm_config = conf["general"]["nwm_configuration"]
fc = ForecastConfig(conf["file_paths"]["fcst_config_file"])
fc.validate_cycle_info(nwm_config)
# get time step and forecast window for the configuration and cycle
win_size, time_step = fc.get_fcst_window_timestep(nwm_config)
return win_size, time_step
[docs]
def check_existing_obs_data(obs_dir: str | Path) -> list:
"""Check existing parquet files of usgs obs and get the dates for previously downloaded data.
Args:
obs_dir (str | Path): Directory containing the observation parquet files.
Returns:
list: List of dates for which observation data is available.
"""
dates0 = list()
parquet_files = glob.glob(str(obs_dir) + "/*.parquet")
if len(parquet_files):
periods = sorted(
[os.path.basename(x).split(".")[0].split("_") for x in parquet_files]
)
for p1 in periods:
if len(p1) == 1:
p1 = [p1[0], p1[0]]
dates0 = dates0 + create_time_sequence(
p1[0], p1[1], freq_hour=24, start_hour=0, end_hour=23
)
dates0 = sorted(list(set(dates0)))
return dates0
[docs]
def check_missing_obs_data(obs_dir: str | Path, conf: dict, gages: list):
"""Check for missing observation data in the specified directory.
Args:
obs_dir (str | Path): Directory containing the observation parquet files.
conf (dict): Dictionary containing NWM configuration.
gages (list): List of gage IDs to check for missing data.
Returns:
None
"""
# Get existing observation dates
df = pd.DataFrame()
parquet_files = glob.glob(str(obs_dir) + "/*.parquet")
for p in parquet_files:
df = pd.concat([df, pd.read_parquet(p)], ignore_index=True)
# If no observation data is found, log warning and exit
if df.empty:
max_show = 5
gages_list = list(gages)
if len(gages_list) > max_show:
shown = gages_list[:max_show]
msg = (
f"No observation data is available for gages {shown} "
f"(showing first {max_show} of {len(gages_list)})."
" Verification cannot proceed. Exit."
)
else:
msg = (
f"No observation data is available for gages {gages_list}."
" Verification cannot proceed. Exit."
)
logger.warning(msg)
sys.exit(0)
if "value_time" not in df.columns:
msg = f"'value_time' column not found in observation data in {obs_dir}."
logger.error(msg)
raise ValueError(msg)
df["value_time"] = pd.to_datetime(df["value_time"])
existing_dates = df["value_time"].unique().tolist()
# Get required observation dates
time_step = 1
if conf["general"]["nwm_configuration"] != "ngen_simulation":
fcst_win, time_step = get_fcst_info(conf)
conf1 = conf["general"]
for i1, dataset in enumerate(conf1["dataset_name"]):
if conf1["nwm_configuration"] == "ngen_simulation":
start_date = pd.to_datetime(conf1["eval_start_date"][i1])
end_date = pd.to_datetime(conf1["eval_end_date"][i1])
else:
start_date = pd.to_datetime(
conf1["forecast_start_date"][i1]
) + pd.Timedelta(hours=time_step)
end_date = pd.to_datetime(conf1["forecast_end_date"][i1]) + pd.Timedelta(
hours=fcst_win
)
required_dates = create_time_sequence(start_date, end_date, freq_hour=time_step)
# Check for missing data. For ngenCERF (single gage), check by dates; otherwise, check by gages
if conf["nwm_forecast"]["data_source"].upper() == "NGENCERF":
# Check for missing dates
max_show = 5 # number of dates to show in logging
missing_dates = [d for d in required_dates if d not in existing_dates]
if missing_dates:
total_missing = len(missing_dates)
shown = missing_dates[:max_show]
shown_str = ", ".join(d.strftime("%Y-%m-%d %H:%M:%S") for d in shown)
if total_missing > max_show:
msg = (
f"{dataset} - Missing observation data for "
f"{total_missing} timestamps (showing first {max_show}): {shown_str} ..."
)
else:
msg = (
f"{dataset} - Missing observation data for "
f"{total_missing} timestamps: {shown_str}"
)
logger.warning(msg)
else:
# first filter data with dates within required dates range
df_required = df[
(df["value_time"] >= min(required_dates))
& (df["value_time"] <= max(required_dates))
]
# then check for missing gages
existing_gages = df_required["location_id"].unique().tolist()
logger.info(
f"{dataset} - Number of gages with observation data available: {len(existing_gages)}"
)
existing_gages_no_prefix = [g.split("-", 1)[1] for g in existing_gages]
missing_gages = [g for g in gages if g not in existing_gages_no_prefix]
if missing_gages:
logger.warning(
f"{dataset} - Missing observation data for {len(missing_gages)} gages."
)
logger.debug(f"{dataset} - Missing gages: {missing_gages}")
[docs]
def safe_fetch_usgs(
site_codes: list, dates: list, conf: dict, out_dir: str, hourly: bool = True
):
"""Fetch USGS streamflow data for a list of gage IDs and a date range, and save to parquet files.
Args:
site_codes (list): List of USGS gage IDs.
dates (list): List of dates for which to fetch data.
conf (dict): Dictionary containing configuration options.
out_dir (str | Path): Directory to save the parquet files.
hourly (bool): Whether to filter data to hourly values.
Returns:
None
"""
try:
usgs_to_parquet(
sites=site_codes,
start_date=min(dates),
end_date=max(dates) + pd.Timedelta(hours=23),
output_parquet_dir=out_dir,
chunk_by=conf["chunk_by"],
filter_to_hourly=hourly,
overwrite_output=conf["overwrite_output"],
)
except aiohttp.client_exceptions.ContentTypeError as e:
raise RuntimeError(
"USGS request returned non-JSON response.\n"
f"Sites: {site_codes}\n"
f"Date range: {min(dates)} → {max(dates)}\n"
f"Hourly: {hourly}\n"
"This often means the USGS service returned an HTML error page "
"(bad parameters, rate limit, or service outage)."
) from e
[docs]
def get_txdot_gage_list(config: dict) -> list:
"""Get the list of gage IDs for retrieving TxDOT streamflow observations.
Args:
config (dict): Dictionary containing configuration options.
Returns:
list: List of TxDOT gage IDs.
"""
txdot_gage_file = config["file_paths"].get("txdot_gage_file", None)
if not txdot_gage_file:
return default_txdot_gage_list
if txdot_gage_file and not txdot_gage_file.exists():
msg = f"TxDOT gage file not found at {txdot_gage_file}. Cannot retrieve TxDOT streamflow observations."
logger.error(msg)
raise FileNotFoundError(msg)
try:
with open(txdot_gage_file, "r") as f:
gage_list = [line.strip() for line in f if line.strip()]
except Exception as e:
msg = f"Error reading TxDOT gage file at {txdot_gage_file}: {e}"
logger.error(msg)
raise e
return gage_list
[docs]
def process_txdot_data(df: pd.DataFrame, gage_id: str, freq: str = "H") -> pd.DataFrame:
"""Process raw TxDOT streamflow data for a specific gage ID.
Args:
df (pd.DataFrame): DataFrame containing raw TxDOT streamflow data.
gage_id (str): Gage ID for which to process data.
freq (str): Frequency for resampling the data ("H" for hourly, "D" for daily).
Returns:
pd.DataFrame: DataFrame containing processed TxDOT streamflow data.
"""
CFS_TO_CMS = 0.3048**3
flow_cols = [
"00060",
"00060_3",
"00060_300060_rq-30",
"00060_downstream, [rq-30",
"00060_learned discharge",
]
# select the first available column for discharge
selected_col = None
for col in flow_cols:
if col in df.columns:
selected_col = col
break
if selected_col is None:
raise ValueError(f"No expected discharge columns found for gage {gage_id}")
# Convert units
df[selected_col] = df[selected_col] * CFS_TO_CMS
# Resample with specified frequency, keeping the first value in each period
df = df.resample(freq).first()
# Keep only selected column for discharge
df = df[[selected_col]]
# Rename selected column to "value"
df = df.rename(columns={selected_col: "value"})
return df
[docs]
def fetch_txdot_gage_data(
site_codes: list, dates: list, conf: dict, out_dir: str, hourly: bool = True
):
"""Fetch streamflow data for a list of TxDOT gage IDs and a date range, and resample temporally as needed.
Args:
site_codes (list): List of TxDOT gage IDs.
dates (list): List of dates for which to fetch data.
conf (dict): Dictionary containing configuration options.
out_dir (str | Path): Directory to save the parquet files.
hourly (bool): Whether to filter data to hourly values.
Returns:
None
"""
dates_str = list(dict.fromkeys(d.strftime("%Y-%m-%d") for d in dates))
date_start = min(dates_str)
date_end = max(dates_str)
df_all = pd.DataFrame()
for gage_id in site_codes:
try:
# convert d1 in UTC to local time for the gage location (for TX: America/Chicago timezone)
tz = "America/Chicago"
local_start = pd.Timestamp(f"{date_start} 00:00:00", tz="UTC").tz_convert(
tz
)
local_end = pd.Timestamp(f"{date_end} 23:59:59", tz="UTC").tz_convert(tz)
df = nwis.get_record(
sites=gage_id,
service="iv",
start=local_start.strftime("%Y-%m-%dT%H:%M:%S"),
end=local_end.strftime("%Y-%m-%dT%H:%M:%S"),
access="3",
)
except Exception as e:
logger.warning(
f"Failed to fetch TxDOT data for gage {gage_id} for date range {date_start} to {date_end}: {e}"
)
if df.empty:
logger.warning(
f"No TxDOT data returned for gage {gage_id} for date range {date_start} to {date_end}"
)
continue
# logger.info(f"df before processing for gage {gage_id} on {date_start} to {date_end}:\n{df.head()}")
df_processed = process_txdot_data(df, gage_id, freq="H" if hourly else "D")
df_processed["location_id"] = f"usgs-{gage_id}"
df_processed["value_time"] = df_processed.index
df_processed = df_processed[["location_id", "value_time", "value"]]
# remove tz info in value_time (which is in UTC)
if df_processed["value_time"].dt.tz is not None:
df_processed["value_time"] = df_processed["value_time"].dt.tz_localize(None)
# add a few extra columns for compatibility with teehr
df_processed["reference_time"] = df_processed["value_time"]
df_processed["variable_name"] = conf["general"]["variable_name"]
df_processed["measurement_unit"] = "m3/s"
df_all = pd.concat([df_all, df_processed], ignore_index=True)
if not df_all.empty:
start1 = df_all["value_time"].min().strftime("%Y-%m-%dT%H:%M:%S")
end1 = df_all["value_time"].max().strftime("%Y-%m-%dT%H:%M:%S")
output_file = Path(out_dir) / f"{start1}_{end1}_txdot.parquet"
df_all.to_parquet(output_file, index=False)
[docs]
def get_obs_files(conf: dict) -> list[Path]:
"""Get the list of existing observation files."""
obs_path = conf.get("obs_data_dir", None)
obs_file = conf.get("obs_data_file", None)
# gather observation files from both obs_data_dir and obs_data_file (if specified)
obs_files = []
if obs_path:
obs_path = Path(obs_path)
if not obs_path.exists():
msg = f"Observation data path {obs_path} does not exist. Ignore this path."
logger.warning(msg)
else:
# identify all observation files in csv or parquet format in obs_path
obs_files = list(obs_path.glob("*.parquet")) + list(obs_path.glob("*.csv"))
if obs_file:
obs_file = Path(obs_file)
if not obs_file.exists():
msg = f"Observation data file {obs_file} does not exist. Ignore this file."
logger.warning(msg)
else:
obs_files = obs_files + [Path(obs_file)]
return obs_files
[docs]
def read_obs_data(locations: list, conf: dict, output_dir: Path) -> pd.DataFrame:
"""Read observation data from a list of file paths (in csv or parquet format) and concatenate them into a single DataFrame."""
# get gage agency
agency_lookup = get_gage_agency(
locations, next(iter(conf["file_paths"]["crosswalk_file"].values()))
)
df_all = pd.DataFrame()
obs_files = get_obs_files(conf.get("file_paths", {}))
for file_path in obs_files:
file_path = Path(file_path)
logger.info(f"Reading observation data from file: {file_path}")
df = read_data(file_path)
# function to detect time column
def detect_time_column(df: pd.DataFrame, time_col: str = "time"):
for col in df.columns:
if is_object_dtype(df[col]) or is_string_dtype(df[col]):
parsed = pd.to_datetime(df[col], errors="coerce")
if parsed.notna().all():
if col.lower() != time_col:
logger.info(
f"Using '{col}' as 'time' column for the observation file"
)
return col
msg = "No column in the observation file contains fully valid datetimes"
logger.error(msg)
raise ValueError(msg)
# function to detect flow column
def detect_flow_column(df: pd.DataFrame, flow_col: str = "obs_flow"):
for col in df.columns:
if is_numeric_dtype(df[col]):
if col.lower() != flow_col:
logger.info(
f"Using '{col}' as 'obs_flow' column for the observation file"
)
return col
msg = "No column in the observation file contains numeric flow values"
logger.error(msg)
raise ValueError(msg)
# if 'time' column not present, try to detect it
time_col = "time" if "time" in df.columns else detect_time_column(df)
df = df.rename(columns={time_col: "value_time"})
df["value_time"] = pd.to_datetime(df["value_time"], errors="raise")
# if 'obs_flow' column not present, try to detect it
flow_col = "obs_flow" if "obs_flow" in df.columns else detect_flow_column(df)
df = df.rename(columns={flow_col: "value"})
# add a few extra columns for compatibility with teehr if not already present
df["reference_time"] = df["value_time"]
df["variable_name"] = conf.get("general", {}).get("variable_name", "streamflow")
df["measurement_unit"] = "m3/s"
df["configuration"] = "observed"
# get location_id from file name if not present in the dataframe
if "location_id" not in df.columns:
if "_" not in file_path.stem:
msg = (
f"Cannot determine location_id from filename '{file_path.name}'. "
"Expected filename stem to contain '_' "
"(e.g., '01123000_hourly_discharge.csv')."
)
logger.warning(msg)
continue
location_id = file_path.stem.split("_")[0]
if location_id not in locations:
msg = (
f"Location ID '{location_id}' from filename '{file_path.name}' "
"is not in the list of specified locations. Skipping this file."
)
logger.info(msg)
continue
# add gage agency as a prefix to location_id for consistency with teehr (e.g., 'usgs-01123000')
df["location_id"] = (
f"{agency_lookup.get(location_id, 'unknown')}-{location_id}"
)
else:
# filter to only the specified locations
df = df[df["location_id"].isin(locations)]
# add gage agency as a prefix to location_id
agency = df["location_id"].map(agency_lookup)
df["location_id"] = agency + "-" + df["location_id"].astype(str)
# reorder columns
df = df[
[
"location_id",
"reference_time",
"value_time",
"value",
"variable_name",
"measurement_unit",
"configuration",
]
]
df_all = pd.concat([df_all, df], ignore_index=True)
# remove duplicates if any
df_all = df_all.drop_duplicates(subset=["location_id", "value_time"])
# save data to parquet files
if not df_all.empty:
start1 = df_all["value_time"].min().strftime("%Y-%m-%dT%H:%M:%S")
end1 = df_all["value_time"].max().strftime("%Y-%m-%dT%H:%M:%S")
output_file = Path(output_dir) / f"{start1}_{end1}_local.parquet"
output_file.parent.mkdir(parents=True, exist_ok=True)
df_all.to_parquet(output_file, index=False)
return
[docs]
def get_gage_agency(gages: list, cwt_file: str | Path) -> dict:
"""Get the agency for each gage ID based on the crosswalk file."""
cwt = read_data(cwt_file)
lookup, miss_ids = find_locations_in_crosswalk(
gages,
cwt,
)
gage_agency = {gage: lookup[gage].split("-")[0] for gage in gages if gage in lookup}
return gage_agency
[docs]
def retrieve_usgs_obs(locations: dict, conf: dict, output_dir: Path):
"""Retrieve USGS streamflow observations given configuration and a list of gage IDs.
Args:
locations (dict): Dictionary containing USGS gage IDs for which observations are to be retrieved.
conf (dict): Dictionary defining the configurations (e.g., config.yaml).
output_dir (Path): Path to store the observation data.
Data retrieved will be saved in parquet files by chunk (e.g., month) in the data directory defined in conf
"""
# get the list of unique USGS gage IDs
list_all = list(
{item for subdict in locations.values() for item in subdict["primary"]}
)
# gage category by agency
cwt_file = next(iter(conf["file_paths"]["crosswalk_file"].values()))
gage_agency = get_gage_agency(list_all, cwt_file)
list_usgs = [k for k, v in gage_agency.items() if v.lower() == "usgs"]
list_txdot = get_txdot_gage_list(conf)
list_txdot = [gage for gage in list_txdot if gage in list_usgs]
list_usgs = [gage for gage in list_usgs if gage not in list_txdot]
# read obs data in existing files
read_obs_data(list_all, conf, output_dir)
# if flow_observation.usgs section not present or empty in config, skip retrieval
conf2 = conf["flow_observation"]["usgs"]
if not conf2:
# Check for missing observation data after retrieval
check_missing_obs_data(output_dir, conf, list_all)
return
# check existing parquet files of usgs obs and get the dates for previously downloaded data
dates0 = list()
overwrite = (conf2 or {}).get("overwrite_output", False)
if not overwrite:
dates0 = check_existing_obs_data(str(output_dir))
dates0 = [x.strftime("%Y-%m-%d") for x in dates0]
if len(dates0) > 0:
logger.info(
f" Existing observed parquet files for {min(dates0)} to {max(dates0)} will be used"
)
# identify start and end dates of observations required by all NWM forecasts datasets
dates = list()
conf1 = conf["general"]
for i1 in range(len(conf1["forecast_start_date"])):
start_date = conf1["forecast_start_date"][i1]
end_date = conf1["forecast_end_date"][i1]
if conf1["nwm_configuration"] == "ngen_simulation":
fcst_win1 = 0
timestep1 = 1
else:
fcst_win1, timestep1 = get_fcst_info(conf)
# adjust start/end date based on forecast window
if fcst_win1 < 0:
start_date = pd.Timestamp(start_date) + pd.Timedelta(
fcst_win1, unit="hours"
)
start_date = start_date.strftime("%Y-%m-%d %H:%M:%S")
else:
end_date = pd.Timestamp(end_date) + pd.Timedelta(
fcst_win1 + 24, unit="hours"
)
end_date = end_date.strftime("%Y-%m-%d %H:%M:%S")
# create the list of dates required
dates = dates + create_time_sequence(
start_date, end_date, freq_hour=24, start_hour=0, end_hour=23
)
dates = sorted(list(set(dates)))
dates = [x.strftime("%Y-%m-%d") for x in dates]
# dates that require data downloading
dates1 = [d1 for d1 in dates if d1 not in dates0]
if len(dates1) == 0:
logger.info(
" Observed data for all required dates and locations already exist"
)
else:
# create data path
output_dir.mkdir(parents=True, exist_ok=True)
# break the list of dates into consecutive chunks
dates2 = sorted([pd.Timestamp(d1) for d1 in dates1])
dates2 = [pd.Timestamp(d1) for d1 in dates2]
date_list = list()
for i1 in range(len(dates2)):
if i1 == 0:
list1 = [dates2[i1]]
else:
if (dates2[i1] - dates2[i1 - 1]).days == 1:
list1 = list1 + [dates2[i1]]
else:
date_list.append(list1)
list1 = [dates2[i1]]
date_list.append(list1)
# determine n_workers dynamically
mem_limit = conf2["memory_per_worker_gb"]
n_workers = get_n_workers(mem_limit)
# loop through the date chunks to download USGS data
hourly = False if timestep1 < 1 else True
for d1 in date_list:
# fetch data for standard usgs (non-TxDOT) gages
if list_usgs:
use_dask = (
conf["nwm_forecast"]["data_source"] != "ngenCERF" and n_workers > 1
)
if use_dask:
logger.info(
f" Using Dask with {n_workers} workers to download USGS data ..."
)
# use a local dask cluster to fetch the data
with (
LocalCluster(
n_workers=n_workers,
processes=True,
memory_limit=f"{mem_limit}GB",
dashboard_address=None,
) as cluster,
Client(cluster) as client,
):
try:
safe_fetch_usgs(
list_usgs, d1, conf2, str(output_dir), hourly
)
except (ServerDisconnectedError, ClientOSError) as e:
logger.warning(
f"Failed to fetch USGS data after retries: {e}"
)
else:
logger.info(
" Running USGS fetch without Dask (single-process mode)"
)
try:
safe_fetch_usgs(list_usgs, d1, conf2, str(output_dir), hourly)
except (ServerDisconnectedError, ClientOSError) as e:
logger.warning(f"Failed to fetch USGS data after retries: {e}")
# fetch data for TxDOT gages
if list_txdot:
logger.info(f" Fetching TxDOT gages data for gages: {list_txdot} ...")
fetch_txdot_gage_data(list_txdot, d1, conf, str(output_dir), hourly)
# clean up memory
gc.collect()
# if output_dir is empty after retrieval, give warning
if not any(output_dir.iterdir()):
msg = "No observation data retrieved for the specified gage IDs and date range. Verification cannot proceed. Exit."
logger.error(msg)
raise Exception(msg)
logger.info(f" Observation data are saved in parquet files at: {output_dir}")
# Check for missing observation data after retrieval
check_missing_obs_data(output_dir, conf, list_all)
[docs]
def get_fcst_files(conf: dict, dataset: str) -> list[Path]:
"""Get the list of forecast files for a given dataset.
Args:
conf (dict): Dictionary containing configuration options.
dataset (str): Name of the dataset for which to get forecast files.
Returns:
list[Path]: List of Path objects representing the forecast files.
"""
fcst_path = conf["file_paths"].get("fcst_data_dir", None)
if fcst_path:
fcst_path = Path(fcst_path[dataset])
if not fcst_path.exists():
msg = f"Forecast data path {fcst_path} does not exist for dataset {dataset}. Verification cannot proceed. Exit."
logger.error(msg)
raise FileNotFoundError(msg)
fcst_file = conf["file_paths"].get("fcst_data_file", None)
if fcst_file:
if fcst_path:
# make sure fcst_file is a string (not a dictionary)
if isinstance(fcst_file, dict):
msg = (
"fcst_data_file should be a string pattern when fcst_data_dir is specified, "
f"but got a dictionary for dataset {dataset}. Verification cannot proceed. Exit."
)
logger.error(msg)
raise ValueError(msg)
# identify all forecast files recursively in fcst_path that match the fcst_file pattern
fcst_files = list(fcst_path.rglob(fcst_file))
else:
if isinstance(fcst_file, dict):
fcst_files = [Path(fcst_file[dataset])]
else:
fcst_files = [Path(fcst_file)]
else:
if conf["nwm_forecast"]["data_source"].upper() != "GCS":
msg = f"fcst_data_file not specified in configuration for dataset {dataset}. Verification cannot proceed. Exit."
logger.error(msg)
raise ValueError(msg)
if not fcst_files:
msg = f"No forecast files found for dataset {dataset} with the specified configuration. Verification cannot proceed. Exit."
logger.error(msg)
raise FileNotFoundError(msg)
return fcst_files
[docs]
def retrieve_fcsts_ngencerf(locations: dict, conf: dict, data_paths: dict):
"""Retrieve NWM forecasts given the configurations and list of locations from the ngenCERF data source.
The forecast files are generated by ngenCERF on the server.
Args:
locations (dict): Dictionary containing secondary ID (NWM link IDs) for which forecasts are to be retrieved.
conf (dict): Dictionary defining the configurations (e.g., config.yaml).
data_paths (dict): Dictionary containing paths to store the data.
Returns:
None
"""
# get time step and forecast window for the configuration and cycle
win_size, time_step = get_fcst_info(conf)
# read forecast data from file for each dataset
for dataset in conf["general"]["dataset_name"]:
logger.info(f" Retrieving forecast data for {dataset} ...")
fcst_files = get_fcst_files(conf, dataset)
for fcst_file in fcst_files:
df_fcst = read_data(
fcst_file, dtype={"sim_flow": float}, parse_dates=["Time"]
)
# make sure the required columns are present
if not {"Time", "sim_flow"}.issubset(df_fcst.columns):
logger.warning(
f"Required columns 'Time' and 'sim_flow' not found in forecast file {fcst_file}. Skipping this file."
)
continue
if df_fcst.empty:
logger.warning(f"No data found in forecast file {fcst_file}")
continue
# order the dataframe by time
df_fcst = df_fcst.sort_values("Time")
# assume reference time is one timestep before the first datetime in the forecast data when win_size is positive,
# or one timestep after the last datetime when win_size is negative.
if win_size > 0:
start_time = df_fcst["Time"].min()
end_time = start_time + pd.Timedelta(hours=win_size - time_step)
reference_time = start_time - pd.Timedelta(hours=time_step)
elif win_size < 0:
end_time = df_fcst["Time"].max()
start_time = end_time + pd.Timedelta(hours=win_size + time_step)
reference_time = end_time + pd.Timedelta(hours=time_step)
else:
msg = "Forecast window size cannot be zero."
logger.error(msg)
raise ValueError(msg)
# make sure the time period covers the forecast window
if df_fcst["Time"].max() < end_time:
msg = (
f"Forecast time period {start_time} to {end_time} is not covered by data. "
f"Available data time period is {start_time} to {df_fcst['Time'].max()}"
)
logger.warning(msg)
# if extra data is found, give warning too
if df_fcst["Time"].max() > end_time:
msg = (
f"Extra forecast data found beyond the expected time period: "
f"Available data time period is {start_time} to {df_fcst['Time'].max()}. "
f"Expected time period is {start_time} to {end_time}."
)
logger.warning(msg)
# filter data by the overlapping period
df_fcst = df_fcst[
(df_fcst["Time"] >= start_time) & (df_fcst["Time"] <= end_time)
]
# make sure every time step is covered
all_times = pd.date_range(
start=start_time, end=end_time, freq=f"{time_step}H"
)
missing_times = all_times[~all_times.isin(df_fcst["Time"])]
if not missing_times.empty:
msg = f"Missing forecast data for time steps: {missing_times.min()} to {missing_times.max()}"
logger.warning(msg)
# rename columns to match the format expected by teehr
df_fcst = df_fcst.rename(
columns={"sim_flow": "value", "Time": "value_time"}
)
# add additional columns to match the format expected by teehr
df_fcst["reference_time"] = reference_time
df_fcst["configuration"] = conf["general"]["nwm_configuration"]
df_fcst["variable_name"] = conf["general"]["variable_name"]
df_fcst["measurement_unit"] = "m3/s"
# get the location_id for the forecast data
loc_fcst_list = list(
{
item
for subdict in locations.values()
for item in subdict["secondary"]
}
)
if len(loc_fcst_list) == 0:
msg = f"No secondary_location_id ({locations[dataset]['secondary']}) found in locations for forecast data. Verification cannot proceed."
logger.error(msg)
raise ValueError(msg)
if len(loc_fcst_list) > 1:
logger.warning(
"Multiple secondary_location_id found in locations. Using the first one for location_id in forecast data."
)
df_fcst["location_id"] = (
conf["general"]["nwm_version"][0] + "-" + str(loc_fcst_list[0])
)
# save forecast data to parquet files
output_file = Path(
data_paths.get("fcst_link").get(dataset)
) / reference_time.strftime("%Y%m%dT%H.parquet")
save_data(df_fcst, output_file)
logger.info(f" Forecast data for {dataset} saved to {output_file.parent}")
[docs]
def retrieve_fcsts_gcs(locations: dict, conf: dict, data_paths: dict):
"""Retrieve NWM forecasts given the configurations and list of locations from Google Cloud Storage.
Args:
locations (dict): Dictionary containing secondary ID (NWM link IDs) for which forecasts are to be retrieved.
conf (dict): Dictionary defining the configurations (e.g., config.yaml).
data_paths (dict): Dictionary containing paths to store the data.
Returns:
None. Data retrieved will be saved in parquet files by forecast cycle in the data directory defined in conf.
"""
output_dir = data_paths.get("fcst")
json_dir = data_paths.get("fcst_json")
data_link_dir = data_paths.get("fcst_link")
# get some general information
conf1 = conf["general"]
conf2 = conf["nwm_forecast"]
config = conf1["nwm_configuration"]
# determine n_workers dynamically
mem_limit = conf2["memory_per_worker_gb"]
n_workers = get_n_workers(mem_limit)
# loop through datasets
for i1, dataset in enumerate(conf1["dataset_name"]):
locations_nwm = locations[dataset]["secondary"]
fetch = conf2["fetch_fcst"][i1]
if fetch:
version = conf1["nwm_version"][i1]
start_date = conf1["forecast_start_date"][i1]
end_date = conf1["forecast_end_date"][i1]
# get forecast configuration and cycle frequency
fc = ForecastConfig(conf["file_paths"]["fcst_config_file"])
fc.validate_cycle_info(config)
cycle_config = fc.get_cycle_info(config)
logger.info(
f" ======== Fetch data for NWM dataset {dataset}: {version} {config} {start_date} to {end_date} =========="
)
# check existing parquet files for NWM forecasts
parquet_files = glob.glob(str(output_dir[dataset]) + "/*.parquet")
hours = sorted([os.path.basename(x).split(".")[0] for x in parquet_files])
cycles0 = [pd.Timestamp(h1) for h1 in hours]
# determine all cycles needed
cycles = []
for c1 in cycle_config:
cycle_start, cycle_end, cycle_freq, fcst_win, fcst_timestep = c1
cycles.extend(
create_time_sequence(
start_date,
end_date,
cycle_freq,
)
)
# cycles not in existing parquet files
cycles1 = [c1 for c1 in cycles if c1 not in cycles0]
if len(cycles1) == 0:
logger.info(
f" All parquet files already exist at: {output_dir[dataset]}"
)
else:
if len(cycles1) < len(cycles):
logger.info(
f" Some parquet files already exist at: {output_dir[dataset]}"
)
# create the data paths
output_dir[dataset].mkdir(parents=True, exist_ok=True)
json_dir[dataset].mkdir(parents=True, exist_ok=True)
# determine the dates
dates1 = sorted(list(set([c1.strftime("%Y-%m-%d") for c1 in cycles1])))
for d1 in dates1:
logger.info(f" Retriving NWM forecast data for {d1} ...")
# use a local dask cluster to fetch the data
with (
LocalCluster(
n_workers=n_workers,
processes=True,
memory_limit=f"{mem_limit}GB",
dashboard_address=None,
) as cluster,
Client(cluster) as client,
):
# fectch NWM forecasts data (1-month short-range took around 1.5 hours)
tlp.nwm_to_parquet(
configuration=config,
output_type=conf2["output_type"],
variable_name=conf1["variable_name"],
start_date=d1,
ingest_days=1,
location_ids=locations_nwm,
json_dir=str(json_dir[dataset]),
output_parquet_dir=str(output_dir[dataset]),
nwm_version=version,
data_source=conf2["data_source"],
kerchunk_method=conf2["kerchunk_method"],
t_minus_hours=conf2["t_minus"],
process_by_z_hour=conf2["process_by_z_hour"],
stepsize=conf2["stepsize"],
ignore_missing_file=conf2["ignore_missing_file"],
overwrite_output=conf2["overwrite_output"],
)
# clean up memory
gc.collect()
logger.info(
f" NWM forecast data are saved in parquet files at: {output_dir[dataset]}"
)
for c1 in cycles:
c1_str = c1.strftime("%Y%m%dT%H")
link1 = Path(data_link_dir[dataset], c1_str + ".parquet")
if not link1.parent.exists():
link1.parent.mkdir(parents=True)
if link1.is_symlink():
link1.unlink()
target1 = Path(output_dir[dataset], c1_str + ".parquet")
link1.symlink_to(target1)
[docs]
def retrieve_ngen_simulation(locations: dict, conf: dict, data_paths: dict):
"""Retrieve NGEN simulation data for the specified configuration.
Based on the data source specified in the configuration, the appropriate retrieval function is called.
Args:
locations (dict): Dictionary containing secondary ID (NWM feature IDs) and primary ID (USGS gage IDs).
conf (dict): Dictionary defining the configurations (e.g., config.yaml).
data_paths (dict): Dictionary containing paths to store the data.
Returns:
None. Data retrieved will be saved in parquet files by forecast cycle in the data directory defined in conf.
"""
# read forecast data from file for each dataset
for idx, (dataset, nwm_ver) in enumerate(
zip(conf["general"]["dataset_name"], conf["general"]["nwm_version"])
):
logger.info(f" Retrieving forecast data for dataset {dataset} ...")
# for ngen simulation, there should be only one file for each dataset
fcst_files = get_fcst_files(conf, dataset)
fcst_file = fcst_files[0]
# make sure fcst_file is a valid file
if not Path(fcst_file).is_file():
msg = f"Forecast file {fcst_file} does not exist. Verification cannot proceed. Exit."
logger.error(msg)
raise ValueError(msg)
df_sim = extract_flow_for_gages(
nc_file=Path(fcst_file),
locations=locations[dataset],
start_time=pd.to_datetime(conf["general"]["eval_start_date"][idx]),
end_time=pd.to_datetime(conf["general"]["eval_end_date"][idx]),
flow_var="flow",
feature_id_var="feature_id",
time_var="time",
)
if df_sim.empty:
msg = f"No simulation data found for {fcst_file} after extraction. Verification cannot proceed. Exit."
logger.error(msg)
raise ValueError(msg)
# add additional columns to match the format expected by teehr
df_sim["reference_time"] = df_sim["value_time"]
df_sim["configuration"] = conf["general"]["nwm_configuration"]
df_sim["variable_name"] = conf["general"]["variable_name"]
df_sim["measurement_unit"] = "m3/s"
# save forecast data to parquet files
start_time = df_sim["value_time"].min()
end_time = df_sim["value_time"].max()
output_file = Path(
data_paths.get("fcst_link").get(dataset)
) / start_time.strftime(
"%Y%m%dT%H" + "-" + end_time.strftime("%Y%m%dT%H") + ".parquet"
)
save_data(df_sim, output_file)
logger.info(f" NGEN simulation data saved to {output_file}")
[docs]
def retrieve_fcsts(locations: dict, conf: dict, data_paths: dict):
"""Retrieve NWM forecast data for the specified locations and configuration.
Based on the data source specified in the configuration, the appropriate retrieval function is called.
Args:
locations (dict): Dictionary containing secondary ID (NWM link IDs) for which forecasts are to be retrieved.
conf (dict): Dictionary defining the configurations (e.g., config.yaml).
data_paths (dict): Dictionary containing paths to store the data.
Returns:
None. Data retrieved will be saved in parquet files by forecast cycle in the data directory defined in conf.
"""
if conf["nwm_forecast"]["data_source"].upper() == "GCS":
retrieve_fcsts_gcs(locations, conf, data_paths)
elif conf["nwm_forecast"]["data_source"].upper() in ["NGENCERF", "HINDCAST"]:
retrieve_fcsts_ngencerf(locations, conf, data_paths)
elif conf["nwm_forecast"]["data_source"].upper() == "NGENSIM":
retrieve_ngen_simulation(locations, conf, data_paths)
else:
msg = (
f"Data source {conf['nwm_forecast']['data_source']} not recognized. "
f"Supported data sources are (case insensitive) 'GCS', 'NGENSIM', 'NGENCERF', and 'HINDCAST'."
)
logger.error(msg)
raise ValueError(msg)