Source code for src.hydro_utils
"""Hydrological and unit-conversion helpers for the forecast pipeline.
Small, mostly stateless functions used throughout data processing and
forecasting:
- time helpers (:func:`seconds_in_month`)
- grid geometry (:func:`calculate_grid_cell_areas`)
- physical conversions (:func:`calculate_evaporation_rate`,
:func:`convert_mm_to_cms`)
- model loading (:func:`load_model`)
Lake surface areas and other constants are defined inline where used.
"""
import calendar
import numpy as np
import pandas as pd
import joblib
[docs]
def seconds_in_month(year, month):
"""
Calculate the number of seconds in a given month of a specific year.
Parameters
----------
year (int): The year for which the month is being calculated.
month (int): The month for which the number of seconds is being calculated (1-12).
Returns
----------
int: The number of seconds in the specified month.
Raises
----------
ValueError: If the month is not between 1 and 12.
"""
# Check if the month is between 1 and 12
if month < 1 or month > 12:
raise ValueError("ERROR: Month must be between 1 and 12.")
# Get the number of days in the specified month and year
num_days = calendar.monthrange(year, month)[1]
# Convert the number of days to seconds (days * 24 hours * 60 minutes * 60 seconds)
return num_days * 24 * 60 * 60
[docs]
def calculate_grid_cell_areas(lon, lat):
"""
Calculate the area of each grid cell given latitude and longitude arrays.
Parameters
----------
lon (array-like): 1D array of longitudes in degrees.
lat (array-like): 1D array of latitudes in degrees.
Returns
----------
numpy.ndarray: A 2D array of grid cell areas in square meters.
Raises
----------
ValueError: If lat or lon are not 1D arrays.
"""
# Check that lon and lat are 1D arrays
if len(lon.shape) != 1 or len(lat.shape) != 1:
raise ValueError("ERROR: Both 'lon' and 'lat' must be 1D arrays.")
# Radius of Earth in meters
R = 6371000.0
# Convert latitude to radians
lat_rad = np.radians(lat)
# Calculate grid cell width in radians
dlat = np.radians(lat[1] - lat[0])
dlon = np.radians(lon[1] - lon[0])
# Initialize the area array
area = np.zeros((len(lat), len(lon)))
# Calculate area of each grid cell in square meters
for i in range(len(lat)):
for j in range(len(lon)):
area[i, j] = R**2 * dlat * dlon * np.cos(lat_rad[i])
return area
[docs]
def calculate_evaporation_rate(temperature_K, latent_heat_flux):
"""
Convert latent heat flux to evaporation rate.
Parameters
----------
temperature_K : float or array
Air temperature in Kelvin.
latent_heat_flux_W_m2 : float or array
Latent heat flux in W/m².
Notes:
No input validation is performed. Negative ``temperature_K`` or
``latent_heat`` values will produce numerically valid but physically
meaningless results — callers are responsible for sanity-checking inputs.
Sub-optimal: a future revision should raise ``ValueError`` for inputs
outside physical ranges.
Returns
-------
evaporation_rate_mm_s : float or array
Evaporation rate in mm/s.
"""
# Convert temperature to Celsius
temperature_C = temperature_K - 273.15
# Latent heat of vaporization (MJ/kg)
lambda_MJ_kg = 2.501 - 0.002361 * temperature_C
# Convert W/m² -> MJ/(s·m²)
latent_heat_flux_MJ = latent_heat_flux * 1e-6
# MJ/(s·m²) / MJ/kg = kg/(m²·s) = mm/s
evaporation_rate_mm_s = latent_heat_flux_MJ / lambda_MJ_kg
return evaporation_rate_mm_s
[docs]
def convert_mm_to_cms(df):
"""
Converts the 'value [mm]' in the dataframe to 'value [cms]' (cubic meters per second) based on lake surface area and the number of seconds in the month.
Parameters
----------
- df (pd.DataFrame): DataFrame containing the columns 'value [mm]', 'lake', and a multi-index with 'month' and 'year'.
Returns
----------
- pd.DataFrame: DataFrame with a new column 'value [cms]' representing the value in cubic meters per second.
Notes:
Recognized lake names are 'superior', 'michigan-huron', 'erie', and
'ontario'. Any other value in the 'lake' column silently yields a
'value [cms]' of 0 (surface area defaults to 0 via ``dict.get``).
Sub-optimal: a future revision should raise on unknown lake names
rather than masking the issue with zeros.
"""
# Dictionary storing the surface area (in square meters) for each lake
lake_sa = {
'superior': 82097 * 1000000, # Lake Superior area in square meters
'michigan-huron': (57753 + 59560) * 1000000, # Lake Michigan-Huron combined area in square meters
'erie': 25655 * 1000000, # Lake Erie area in square meters
'ontario': 19009 * 1000000 # Lake Ontario area in square meters
}
df_units = df.copy() # Create a copy of the input DataFrame to avoid modifying the original
# Apply the conversion formula for each row in the dataframe
df_units['value [cms]'] = df_units.apply(
lambda row: (
# Convert mm to meters, multiply by the lake surface area, and divide by seconds in the given month
(row['value [mm]'] / 1000) * lake_sa.get(row['lake'], 0) /
seconds_in_month(row['year'], row['month'])
), axis=1
)
# Return the modified DataFrame with the new 'value [cms]' column
return df_units
[docs]
def load_model(model_name: str, models_info: list):
"""
Load a serialized model by its name.
Parameters
----------
model_name : str
The short identifier of the model (e.g., "GP", "RF").
models_info : list of dict
Each dict must contain:
- "model": the short identifier
- "path" : the path to the saved model file
Returns
-------
object
The deserialized model object.
Raises
------
ValueError
If `model_name` isn’t found in `models_info`.
FileNotFoundError
If the `.joblib` file cannot be located.
"""
try:
path = next(m["path"] for m in models_info if m["model"] == model_name)
except StopIteration:
raise ValueError(f"Model '{model_name}' not found in models_info")
return joblib.load(path)