Add a latent space diagnostic notebook

[aranas]

Problem / Goal

To aid model development we want visual diagnostics of intermediate layer activation in the sea ice prediction pipeline.

As a first target, we want to visualize the multimodal argo-float pipeline to provide evidence for the following research question:

Does incorporating Argo float observations shift the model's internal representation in a physically meaningful way and if so, at which layer does that shift first become detectable?

Initial demonstration of this diagnostic will be a standalone diagnostic notebook on a separate branch, integration into pipeline might follow at a later stage if diagnostic proves valuable.

What does good look like?

• A latent_space_viz feature branch on the pipeline repo that adds an optional --save-activations flag to the inference CLI, which when set triggers hook registration before the forward pass and saves captured activations to disk as xarray datasets or tensors keyed by layer name and timestep.
• A viz/hook_manager.py module that registers forward_hooks on named encoder/processor/decoder modules at runtime, captures activations as the forward pass executes, and writes them to the configured output path.
• A notebooks/layer_diagnostics.ipynb notebook that orchestrates the full diagnostic by invoking the pipeline inference CLI via shell cells, once with argo float data, once without, both with --save-activations set and loads the resulting activation files for visualisation.

Not in scope

• No changes to model code, configs, or data modules
• No production integration, CLI tooling
• No geographic projection or statistical analysis logic

Related context

#98

[aranas]

@IFenton @jemrobinson

1. Should this be best done for UNet or viT?
2. which layers to hook? once after encoder, processor (3 layers throughout model, or attention maps for viT), just before & after decoder?
3. Does the diagnostic approach make sense given the research question. Is a with/without Argo activation comparison a reasonable way to probe float data influence?
4. Is attaching external forward_hooks to named modules feasible with the current architecture, or are there known gotchas?
5. What documentation or interface conventions would you need from the notebook to make future pipeline integration straightforward?

[aranas]

Regarding which model to use:

• I have contacted Erin for a calibrated model checkpoint. If I get this, we could use this to see if diagnostics reflect something meaningful in a reasonable model
• for the argo float data we do not have a fully trained model yet, but we can still inspect encoder changes

[aranas]

Discussion with @louisavz @jemrobinson @IFenton

• concern is that argo float data is not good enough to showcase benefit at this point
• for initial visuals, focus on how comparing with/without era5 to show model leverages it over SIC
• for initial visuals, use naive linear encoder/decoder and Unet