Skip to content

Comments

feat: multi-component hydrograph loss function#136

Merged
taddyb merged 5 commits intoDeepGroundwater:masterfrom
taddyb:nnse
Feb 20, 2026
Merged

feat: multi-component hydrograph loss function#136
taddyb merged 5 commits intoDeepGroundwater:masterfrom
taddyb:nnse

Conversation

@taddyb
Copy link
Collaborator

@taddyb taddyb commented Feb 20, 2026

Summary

  • Replaces the single MSE loss with a configurable multi-component hydrograph loss (hydrograph_loss) targeting peak amplitude, baseflow, timing, and overall fit — each per-gage normalized to prevent large basins from dominating
  • Adds LossConfig (Pydantic) to ExperimentConfig with sensible defaults so existing configs work unchanged
  • Adds per-component loss logging in train.py for diagnostics during training
  • Adds comprehensive unit tests for the loss function and an integration test for the training loop

Changes

  • src/ddr/validation/losses.py (new): hydrograph_loss with four components — overall MSE, peak regime (P98), baseflow regime (P30), and temporal gradient loss. Per-gage loss capped at 10 to prevent variance-collapse blowup.
  • src/ddr/validation/configs.py: LossConfig with weights, percentiles, and eps. Added as loss field on ExperimentConfig with defaults.
  • src/ddr/validation/__init__.py: Re-exports hydrograph_loss.
  • scripts/train.py: Uses hydrograph_loss instead of mse_loss; logs individual component values.
  • benchmarks/: Minor config/import alignment.
  • tests/validation/test_losses.py (new): Unit tests for all loss components.
  • tests/integration/ (new): End-to-end training loop integration test.

Test plan

  • uv run pytest tests/validation/test_losses.py — unit tests for loss components
  • uv run pytest tests/integration/test_training.py — integration test for training loop
  • Verify existing configs still work (LossConfig defaults applied automatically)
  • Training run on HPC to validate gradient flow through all components

🤖 Generated with Claude Code

@taddyb taddyb merged commit 9c48986 into DeepGroundwater:master Feb 20, 2026
4 checks passed
@taddyb taddyb deleted the nnse branch February 20, 2026 02:02
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant