kruparell-data assimilation by kruparell · Pull Request #283 · google-research/flood-forecasting

kruparell · 2026-09-02T11:21:52Z

Variational Data Assimilation for Hydrological Forecasting

Summary

This PR introduces a gradient-based Variational Data Assimilation (DA) framework for hydrological forecasting. This framework optimizes latent model components (such as static and dynamic catchment representations) exposed through a standardized, model-agnostic interface.

Files Changed and Rationale

1. `googlehydrology/modelzoo/basemodel.py` & `mean_embedding_forecast_lstm.py`

BaseModel contract: Added supported_assimilation_components property and get_supported_assimilation_components() helper.
MeanEmbeddingForecastLSTM overrides:
- Checks data['assimilation_overrides'] (with backward-compatible top-level keys) for 'static_embedding', 'hindcast_embedding', and 'forecast_embedding'.
- If supplied, reuses the optimized representations directly instead of re-evaluating embedding submodules.
- Attaches all generated representations to the returned dictionary for downstream inspection.

2. `googlehydrology/utils/cmal_deterministic.py`

Standalone mathematical utilities for CMAL:
- calc_cmal_mean(mu, b, tau): Vectorized numerical quadrature computing the expected value $\mathbb{E}[y]$ from asymmetric Laplace parameters.
- ensure_y_hat(pred, use_median=True/False): Unifies prediction dictionaries and ensures y_hat is populated with robust handling of non-finite outputs.

3. `googlehydrology/utils/assimilationconfig.py`

Config dataclass validating DA hyperparameters: assimilation_window, assimilation_lead_time, assimilation_components, regularization_weight, learning_rate, epochs, and component-specific learning rates and regularization weights.

4. `googlehydrology/evaluation/assimilation.py`

Self-contained Var DA assimilation engine:
- Optimizes exposed latent representations over rolling assimilation windows using observed streamflow.
- Applies intensive scale-invariant regularization (torch.mean((param - base_param) ** 2)).
- Implements local _create_assimilation_optimizer supporting Adam, AdamW, SGD, RMSprop, and Adagrad without altering package-level training contracts.
- Updates evaluation metrics and rollout predictions across sequential time horizons.

5. `googlehydrology/run.py`, `evaluate.py`, & `tester.py`

Integrates run_data_assimilation flag into CLI commands and pipeline workflows (evaluate and infer).
Updates documentation in README.md and docs/source/usage/quickstart.rst.

6. `test/test_assimilation.py` & `test/test_cli.py`

Comprehensive test suite validating:
- Multi-component embedding optimization with model weight freezing.
- Model-agnostic fallback on mock architectures implementing the contract.
- Regularization penalty enforcement and loss minimization.
- CLI and pipeline entrypoints with Data Assimilation enabled.

Verification & Testing

All unit and integration tests pass cleanly:

$ pytest test/test_assimilation.py test/test_cli.py -v
============================= test session starts ==============================
test/test_assimilation.py::AssimilationTest::test_initialization PASSED
test/test_assimilation.py::AssimilationTest::test_embedding_gradient_update PASSED
test/test_assimilation.py::AssimilationTest::test_model_weights_remain_frozen PASSED
test/test_assimilation.py::AssimilationTest::test_regularization_loss PASSED
test/test_assimilation.py::AssimilationTest::test_component_specific_learning_rates PASSED
test/test_assimilation.py::AssimilationTest::test_missing_component_handling PASSED
test/test_assimilation.py::AssimilationTest::test_generic_model_contract PASSED
test/test_assimilation.py::AssimilationTest::test_rollout_prediction_stitching PASSED
...
test/test_cli.py::test_run_main_infer_data_assimilation PASSED
======================== 19 passed in 23.96s ========================

sylvesterkaczmarek

Precipitation DA silently falls back to the first hindcast feature when no precipitation key matches. That can turn assimilation_targets=['precip'] into optimisation of temperature, radiation, or another forcing while still producing a plausible rollout. Could this fail explicitly, or require a configured forcing key, instead of selecting list(hind_dict.keys())[0]? For DA this seems safer than silently changing the physical variable being assimilated.

kruparell · 2026-09-03T07:47:01Z

Precipitation DA silently falls back to the first hindcast feature when no precipitation key matches. That can turn assimilation_targets=['precip'] into optimisation of temperature, radiation, or another forcing while still producing a plausible rollout. Could this fail explicitly, or require a configured forcing key, instead of selecting list(hind_dict.keys())[0]? For DA this seems safer than silently changing the physical variable being assimilated.

Thanks for calling this out! You're completely right. Falling back to list(hind_dict.keys())[0] was a silent bug that could optimize radiation or temperature when precipitation keys didn't match.

I have pushed an update that:

Optimizes all matching precipitation features (hres_total_precipitation, graphcast_total_precipitation, imerg_precipitation, cpc_precipitation) together.
Raises an explicit ValueError listing available keys if no precipitation key is found, eliminating the silent fallback.
Supports optional explicit configuration via precip_forcing_keys or precip_forcing_key in AssimilationConfig.

grey-nearing

Temp & partial review while you reconfigure the PRs.

grey-nearing · 2026-09-02T15:27:27Z

+
+                            mask = ~torch.isnan(t_sub) & ~torch.isnan(p_sub)
+                            if mask.any():
+                                loss = torch.mean((p_sub[mask] - t_sub[mask]) ** 2)


Perhaps we should add the loss function to training/loss.py as a child of baseloss. And then either: (1) add an assimilation_loss arguemnt to config.py, or (2) have a specific base_loss inheritance class that is an assimilation_loss and hard-code assimilation.py to use that. The former is better unless you want to make sure the user is unable to change the assimilation loss. I would do that only if you think there is a potential for a user to make a mistake in the sense that assimilation can theoretically or practically only use one specific loss. If there are specific loss functions you want to allow for assimilation or ban for assimilation, you can add that check in config.py (there is a special section for checks on the config file).

grey-nearing · 2026-09-02T15:33:00Z


-        with torch.inference_mode():
+        if data_assimilation:
+            assimilation = Assimilation(self.cfg.assimilation_config)


Is there a check whether the assimilation_config argument is present in the config file if the cli flag for assimilation is used?

I've added this in now

grey-nearing · 2026-09-04T07:43:52Z

+
+                            mask = ~torch.isnan(t_sub) & ~torch.isnan(p_sub)
+                            if mask.any():
+                                loss = torch.mean((p_sub[mask] - t_sub[mask]) ** 2)


Same as above.

grey-nearing · 2026-09-04T07:44:01Z

+
+                            mask = ~torch.isnan(t_sub) & ~torch.isnan(p_sub)
+                            if mask.any():
+                                loss = torch.mean((p_sub[mask] - t_sub[mask]) ** 2)


Same as above.

grey-nearing · 2026-09-04T07:46:09Z

+
+def _get_var_lr(lr_cfg: Any, var_name: str) -> float:
+    """Retrieves target-specific learning rate from lr_cfg."""
+    if isinstance(lr_cfg, dict):


Outside of DA, this dict is typically meant to be {epoch_int: rl}.

grey-nearing · 2026-09-04T08:01:41Z

This notebook also has paths to your personal machine.

I have not reviewed either of the notebooks yet.

grey-nearing · 2026-09-04T08:14:23Z

+                                loss = torch.mean((p_sub[mask] - t_sub[mask]) ** 2)
+                                reg_loss = 0.0
+                                if mask_e_stat and e_stat_opt.requires_grad:
+                                    reg_loss = reg_loss + bg_stat_w * torch.sum((e_stat_opt - e_stat_base) ** 2)


Why do you sum here instead of mean?

grey-nearing · 2026-09-04T08:15:35Z

+
+    def _parse_target_flags(self) -> Tuple[bool, bool, bool, bool]:
+        """Parses self.targets into boolean flags for (c_hc, h_hc, c_fc, h_fc)."""
+        opt_c_hc = any(k in self.targets for k in ['c_n_hindcast', 'c_0_hindcast', 'c_hc', 'c_n', 'c_0'])


Can we use more descriptive variable names? Or at least have a comment that describes the variabel naming convention.

grey-nearing · 2026-09-04T08:16:57Z

+To run evaluation with test-time 4D-Var Data Assimilation:
+
+   ```
+   run evaluate --run-dir /path/to/your/model_run/ --data-assimilation


Can we use inference mode with DA as well? The difference between evaluate vs. infer is that infer saves the timeseries output whereas evaluate only saves the performance metrics.

Yes we can, I've added a description of this to the README

grey-nearing · 2026-09-04T08:18:17Z

+
 def get_optimizer(
-    model: torch.nn.Module, cfg: Config, *, is_gpu: bool = False
+    model: Union[torch.nn.Module, Iterable[torch.Tensor]], cfg: Config, *, is_gpu: bool = False


Flagging this for myself. Why do we need to allow a different model type?

grey-nearing

more temp comments.

grey-nearing · 2026-09-04T09:00:34Z

+                w_end = min(curr_idx + self.window, a_end)
+                win_len = w_end - curr_idx
+
+                chunk_data = _slice_hydrology_batch(data, curr_idx, w_end)


As discussed in person, the process of choosing the assimilation update targets should be model agnostic. I recommend having each model include a required attribute that lists the model components available for assimilation updating, and then require the user to list the components they want assimilation to target int he config file (with appropriate existence checks). Then direct the optimizer to the selected components generically, in a model-agnostic way.

…milationConfig

…orkflows

kruparell requested a review from amitmarkel as a code owner September 2, 2026 11:21

sylvesterkaczmarek reviewed Sep 2, 2026

View reviewed changes

kruparell force-pushed the kruparell-data-assimilation branch from 2a03eeb to 12b9c39 Compare September 3, 2026 11:53

kruparell changed the title ~~Kruparell-data assimilation~~ kruparell-data assimilation Sep 3, 2026

grey-nearing reviewed Sep 4, 2026

View reviewed changes

kruparell force-pushed the kruparell-data-assimilation branch 2 times, most recently from 36cfee5 to 6cebe9e Compare September 7, 2026 13:20

kruparell added 3 commits September 7, 2026 13:30

feat(modelzoo): Add model-agnostic assimilation components and overrides

f59751c

feat(da): Add model-agnostic 4D-Var Data Assimilation engine and Assi…

00f8870

…milationConfig

feat(pipeline): Integrate data assimilation into evaluate and infer w…

34d8143

…orkflows

kruparell force-pushed the kruparell-data-assimilation branch from 6cebe9e to 34d8143 Compare September 7, 2026 13:32

Uh oh!

Conversation

kruparell commented Sep 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Variational Data Assimilation for Hydrological Forecasting

Summary

Files Changed and Rationale

1. googlehydrology/modelzoo/basemodel.py & mean_embedding_forecast_lstm.py

2. googlehydrology/utils/cmal_deterministic.py

3. googlehydrology/utils/assimilationconfig.py

4. googlehydrology/evaluation/assimilation.py

5. googlehydrology/run.py, evaluate.py, & tester.py

6. test/test_assimilation.py & test/test_cli.py

Verification & Testing

Uh oh!

sylvesterkaczmarek left a comment

Choose a reason for hiding this comment

Uh oh!

kruparell commented Sep 3, 2026

Uh oh!

grey-nearing left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

grey-nearing left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

kruparell commented Sep 2, 2026 •

edited

Loading

1. `googlehydrology/modelzoo/basemodel.py` & `mean_embedding_forecast_lstm.py`

2. `googlehydrology/utils/cmal_deterministic.py`

3. `googlehydrology/utils/assimilationconfig.py`

4. `googlehydrology/evaluation/assimilation.py`

5. `googlehydrology/run.py`, `evaluate.py`, & `tester.py`

6. `test/test_assimilation.py` & `test/test_cli.py`