bakaano.neuralnet.train

Training pipeline for regional streamflow models.

Role: Build training datasets and train the TCN-based streamflow model.

bakaano.neuralnet.train.model_config_paths(model_path)

Return supported sidecar config paths for a saved Keras model.

bakaano.neuralnet.train.validate_keras_archive(model_path)

Raise a clear error if a .keras archive is missing model weights.

bakaano.neuralnet.train._load_pysheds_grid()

Import pysheds lazily to avoid import-time backend failures.

bakaano.neuralnet.train._normalize_station_ids(station_ids)

Return station ids as a comparable set across scalar/array-like inputs.

bakaano.neuralnet.train._predictor_cache_path(working_dir)

Return the canonical predictor-response cache path.

bakaano.neuralnet.train._preview_items(items, limit=8)

Return a compact preview string for notebook progress messages.

bakaano.neuralnet.train._build_predictor_cache_payload(data, train_start, train_end)

Wrap predictor-response data with period metadata for safe cache reuse.

bakaano.neuralnet.train._load_predictor_cache_if_compatible(cache_path, train_start, train_end)

Load cached predictors only when the stored period matches exactly.

bakaano.neuralnet.train.filter_training_data_by_station_ids(rawdata, station_ids)

Keep only training entries whose embedded station id matches the target set.

bakaano.neuralnet.train.validate_training_runoff_window(rawdata, train_start, train_end)

Require exact routed-runoff coverage for the requested training dates.

bakaano.neuralnet.train.train_streamflow_model(working_dir, study_area, train_start, train_end, grdc_netcdf=None, batch_size=32, num_epochs=300, learning_rate=0.0005, loss_function='asym_laplace_nll', seed=100, routing_method='mfd', catchment_size_threshold=1, area_normalize=True, log_transform=True, lr_schedule='cosine', warmup_epochs=1, min_learning_rate=5e-05, csv_dir=None, lookup_csv=None, id_col='id', lat_col='latitude', lon_col='longitude', date_col='date', discharge_col='discharge', file_pattern='{id}.csv', model_overwrite=True)

Train the Bakaano streamflow model directly from the neuralnet module.

class bakaano.neuralnet.train.DataPreprocessor(working_dir, study_area, grdc_streamflow_nc_file, train_start, train_end, routing_method, catchment_size_threshold)

Bases: object

_load_runoff_entries_for_period()

Load routed runoff entries and require exact daily coverage for training.

_load_optional_rainfall_entries_for_period()

Load routed rainfall entries if available and fully aligned.

_extract_station_rowcol(lat, lon)

Extract the row and column indices for a given latitude and longitude from given raster file.

Parameters:

• lat (float) – The latitude of the station.

• lon (float) – The longitude of the station.

Returns:

• row (int) – The row index corresponding to the given latitude and longitude.

• col (int) – The column index corresponding to the given latitude and longitude.

_snap_coordinates(lat, lon)

Snap the given latitude and longitude to the nearest river segment based on a river grid.

Parameters:

• lat (float) – The latitude to be snapped.

• lon (float) – The longitude to be snapped.

Returns:

• snapped_lat (float) – The latitude of the nearest river segment.

• snapped_lon (float) – The longitude of the nearest river segment.

load_observed_streamflow(grdc_streamflow_nc_file)

Load and filter observed GRDC streamflow data in a schema-robust way. Works for single- and multi-station NetCDFs.

Parameters:

grdc_streamflow_nc_file (str) – Path to GRDC NetCDF file.

Returns:

Filtered GRDC subset for the study area.

Return type:

xarray.Dataset

_open_grdc_dataset(grdc_streamflow_nc_file)

Open GRDC NetCDF with backend fallback for Colab/Drive compatibility.

load_observed_streamflow_from_csv_dir(csv_dir, lookup_csv, id_col='id', lat_col='latitude', lon_col='longitude', date_col='date', discharge_col='discharge', file_pattern='{id}.csv')

Load observed streamflow from per-station CSV files using a lookup table.

The lookup table must include station identifiers and coordinates. The method filters stations to the study area, then loads per-station CSVs by ID.

Parameters:

• csv_dir (str) – Directory containing per-station CSV files.

• lookup_csv (str) – CSV file with station ids and coordinates.

• id_col (str) – Station id column in lookup CSV.

• lat_col (str) – Latitude column in lookup CSV.

• lon_col (str) – Longitude column in lookup CSV.

• date_col (str) – Date column in station CSVs.

• discharge_col (str) – Discharge column in station CSVs.

• file_pattern (str) – Pattern for station CSV filenames (e.g., "{id}.csv").

Returns:

Mapping of station_id to observed discharge DataFrame.

Return type:

dict

get_data()

Extract and preprocess predictor and response variables for each station based on its coordinates.

Returns:

A list containing two elements: - self.data_list: A list of tuples, each containing predictors (DataFrame) and response (DataFrame). - self.catchment: A list of tuples, each containing catchment data (accumulation and slope values).

Return type:

list

class bakaano.neuralnet.train.StreamflowModel(working_dir, batch_size, num_epochs, loss_function, learning_rate=0.0001, train_start=None, train_end=None, seed=100, area_normalize=True, log_transform=True, lr_schedule=None, warmup_epochs=3, min_learning_rate=1e-05)

Bases: object

Role: Define and train the multi-scale TCN streamflow model.

Full-materialization training variant of the regional streamflow model.

Key characteristics (actual behavior): - Prepares per-station scaled series using area normalization (optional). - Materializes all valid 365-day sliding windows in memory. - Trains directly with in-memory NumPy arrays. - Enables XLA globally via tf.config.optimizer.set_jit(True).

_model_config()

Return inference-critical options that must match this checkpoint.

save_model_config(model_path)

Save inference-critical checkpoint settings next to the model.

_build_lr_callback()

Create a learning-rate schedule callback with optional warmup.

prepare_data(data_list)

Prepare the data for training the streamflow prediction model.

This materializes all sliding windows (365), filters NaNs once, and concatenates across stations.

print_training_summary(station_count=None)

Print a compact summary of prepared training tensors.

build_model()

Multi-scale TCN with sequence FiLM, adaptive timescale weighting, attention pooling, rainfall-runoff gating, and persistence correction.

train_model()

load_regional_model(path)

Load a previously saved regional model from disk.

Parameters:

• path (str) – Path to the saved model file.

• Returns – tensorflow.keras.Model: Loaded model instance.

regional_summary()

Print the Keras model summary.