Build target mask for streams with empty target channels

[Jubeku]

This PR contains 2 fixes related to multi-stream inputs:

1. Empty target channels

For multi-stream inputs (i.e. ERA5+Obs) we need to make sure that Forcing: False in the input streams (otherwise are not encoded as latent target). Additionally, we might want to set targets: [] to save memory (especially for ERA5 N320; only source channels in the target batch are used for the latent target) - however, empty targets lead to an empty target mask which leads to the stream being ignored in the latent target.

This PR generates the target mask in latent_loss and student_teacher mode despite empty targets: [].

2. Consistent random number (from @kctezcan)

Problem: the random number used to make the noise signal for the diffusion forward process is coming from the source random draw. the random number for calculating the loss weight for diffusion is coming from the target random draw.

Fix: take the arget random number also for the source.

Remark: for this to work, the target_input in the config has to be set:

  target_input: {
    "forecasting" : {
      ...
      masking_strategy_config: {
        diffusion_rn: True,
        noise_distribution: *noise_distribution,
        sigma_min: *sigma_min,
        sigma_max: *sigma_max
        },
      ...

Description

Issue Number

Is this PR a draft? Mark it as draft.

Checklist before asking for review

• I have performed a self-review of my code
• My changes comply with basic sanity checks:
  • I have fixed formatting issues with ./scripts/actions.sh lint
  • I have run unit tests with ./scripts/actions.sh unit-test
  • I have documented my code and I have updated the docstrings.
  • I have added unit tests, if relevant
• I have tried my changes with data and code:
  • I have run the integration tests with ./scripts/actions.sh integration-test
  • (bigger changes) I have run a full training and I have written in the comment the run_id(s): launch-slurm.py --time 60
  • (bigger changes and experiments) I have shared a hegdedoc in the github issue with all the configurations and runs for this experiments
• I have informed and aligned with people impacted by my change:
  • for config changes: the MatterMost channels and/or a design doc
  • for changes of dependencies: the MatterMost software development channel

[Jubeku]

@shmh40, @clessig, do you have any concerns with these 2 dataloading-related fixes?

[Jubeku]

cc @sophie-xhonneux

[sophie-xhonneux]

I don't think this breaks anything. JEPA would run without stream forcing at the moment and that code path is the same. When we set stream forcing it should also now have the same code path (here locally, unsure about elsewhere), as without it, which is sensible. Overall this feels like quite a bit of code duplication and at some point we should revisit this logic. I wouldn't block the PR because of it, but create an issue to remind ourselves that this needs clean-up.

[Jubeku]

I opened an issue: #2741