fix: use forcing data_mask for ensemble training forward pass and loss

[yyexela]

Fixes a bug where masked forcing variables containing NaN could produce NaN loss during ensemble training.

When train_on_batch is called with n_ensemble > 1, the data_mask controlling which forcing values get zeroed before the forward pass was read from input_ensemble_data (the initial-condition slice, which carries no mask) instead of forcing_ensemble_data. The same incorrect source was used for loss accumulation. This meant any sample with a masked (NaN) forcing variable would propagate NaN through the model and produce NaN loss.

Changes:

• TrainStepper._run_steps: pass data_mask=forcing_ensemble_data.data_mask to predict_generator (was input_ensemble_data.data_mask)

• TrainStepper._run_steps: pass data_mask=data.data_mask to _accumulate_step_loss (was input_batch_data.data_mask)

• Tests added

• If dependencies changed, "deps only" image rebuilt and "latest_deps_only_image.txt" file updated

[yyexela]

Claude Description

Root cause

get_start calls subset_names(prognostic_names), which strips data_mask down to only prognostic variable keys. So:

• input_batch_data.data_mask / input_ensemble_data.data_mask — prognostic vars only
• data.data_mask / forcing_ensemble_data.data_mask — all vars (prognostic + forcing + output-only)

Both bugs are the same mistake: using the IC-derived mask where the full mask is needed.

---

Change 1 — predict_generator: zero out NaN forcing before the forward pass

predict_generator uses data_mask to zero out masked forcing values before they enter the model. With input_ensemble_data.data_mask, forcing-only variables like b have no mask entry (stripped by subset_names), so their NaN values are never zeroed and propagate through to produce NaN loss.

Fix: use forcing_ensemble_data.data_mask, which comes from data directly with no subsetting.

---

Change 2 — _accumulate_step_loss: exclude masked samples from loss count

_accumulate_step_loss uses data_mask to exclude masked samples from the loss denominator. With input_batch_data.data_mask, output-only variables have no mask entry, so all their samples are counted — even ones that should be excluded.

Fix: use data.data_mask, which has mask entries for all variables.