Skip to content

[Draft] Data assimilation tutorial and fine-tune template updates - #351

Draft
kruparell wants to merge 7 commits into
google-research:mainfrom
kruparell:fix/finetune-statics-path
Draft

kruparell wants to merge 7 commits into
google-research:mainfrom
kruparell:fix/finetune-statics-path

Conversation

@kruparell

Copy link
Copy Markdown
Contributor

'run finetune' inherits every setting not present in the fine-tune config from the base run's config.yml. The shipped 5-basin-example run was trained elsewhere, so its statics_data_dir (/home/gsnearing/...) does not exist for other users and fine-tuning from the tutorial failed with 'Could not find subdataset directories'. The template already overrides targets_data_dir; override statics_data_dir the same way.

- backend.load_data_and_metrics: warn when basins are missing from a results
  zarr (run infer appends basins one at a time, so a file that is still being
  written is incomplete) and recompute cached metrics when they do not match
  the loaded basins.
- backend.plot_hydrograph_comparison: clear ValueError naming the missing
  basin and the available ones instead of xarray's KeyError.
- Notebook: DA label no longer depends on model_name (re-assigned by the
  fine-tuning section); DA basin dropdown is built from basins present in both
  result sets; stripped stale cell outputs.
- Regenerate precalculated_metrics_data_assimilation.csv for 5-basin-example
  from the shipped zarr (the committed cache came from an earlier run).
- DA hydrograph cell also draws the fine-tuned model when the selected basin
  is the fine-tuning basin.
- Discussion cell: per-lead NSE of base / base+DA / fine-tuned for four test
  basins (fine-tuned with the tutorial settings) and what it implies.
… DA snippet with the tutorial template

- Ship a second fine-tuned run (finetune-camels_12115000, trained with the
  tutorial's fine-tuning settings) with regular and DA results.
- New notebook section plotting NSE-vs-lead and hydrographs for base,
  base+DA, fine-tuned and fine-tuned+DA on camels_13235000 (fine-tuning
  wins beyond lead 1) and camels_12115000 (DA wins at most leads), with a
  four-basin comparison table; Discussion bullet shortened accordingly.
- README: DA snippet uses the tutorial template values (window 7, reg 0.1)
  and points to the global-model example (window 90, reg 0.5).
@grey-nearing
grey-nearing force-pushed the fix/finetune-statics-path branch from 7a4f952 to 7867086 Compare October 8, 2026 22:50
@grey-nearing
grey-nearing marked this pull request as draft October 8, 2026 22:50
@grey-nearing grey-nearing changed the title Tutorial: set statics_data_dir in the fine-tune template [Draft] Data assimilation tutorial and fine-tune template updates Oct 8, 2026
@grey-nearing

Copy link
Copy Markdown
Collaborator

Thanks @kruparell! We are planning to retire the tutorial/ directory in favor of the upcoming UI, which will replace the basic need for standalone notebook tutorials and avoid committing additional Zarr/checkpoint data to the repository.

To make sure none of your tutorial work is lost while merging your documentation, validation fixes, and unit/integration tests from #352 right away:

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.

2 participants