AIFS-ENS Initial time-step fix

[NickGeneva]

Earth2Studio Pull Request

Description

Discussion here: #578

Basically the initial step of the generator was returning the invariants (104 variables) instead of the expected 100.
The iterator test skipped this step so it wasnt caught and got through the previous PR.

from pathlib import Path

import earth2studio.run as run
from earth2studio.data import CDS
from earth2studio.io import NetCDF4Backend
from earth2studio.models.px import AIFSENS

package = AIFSENS.load_default_package()
model = AIFSENS.load_model(package)
data = CDS()
io = NetCDF4Backend(file_name=Path("/tmp/output.nc"), backend_kwargs={'mode': 'w'})

nsteps = 1
io = run.deterministic(["2024-01-01"], nsteps, model, data, io)

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.
• Assess and address Greptile feedback (AI code review bot for guidance; use discretion, addressing all feedback is not required).

Dependencies

[greptile-apps]

Greptile Overview

Greptile Summary

This PR fixes a critical bug in the AIFS-ENS model where the iterator was incorrectly skipping the initial time-step (0h lead time) and the output was incorrectly including 4 invariant variables.

Key Changes:

• Replaced hardcoded index slicing ([:92] and [101:]) with dynamic filtering using self.model.data_indices.data.output.forcing in _fill_input method
• Ensures all 13 forcing/invariant variables (4 invariants + 9 time-dependent forcings) are properly removed from output instead of only 9
• Fixed test to validate iterator returns initial condition at 0h lead time by removing next(p_iter) call
• Updated lead time assertion from 6 * (i + 1) to 6 * (i) to match corrected iterator behavior

Impact:

• Output now correctly has 100 variables instead of 104 (fixes invariants being incorrectly included)
• Iterator now properly yields initial condition first, matching expected behavior for prognostic models

Confidence Score: 4/5

• This PR is safe to merge with one minor style improvement recommended
• The changes correctly fix a critical bug where invariant variables were incorrectly included in output and the iterator skipped the initial time-step. The logic is sound and internally consistent. One pre-existing type annotation issue was identified (not introduced by this PR) that should be corrected for code clarity.
• No files require special attention - the changes are straightforward bug fixes

Important Files Changed

File Analysis

Filename	Score	Overview
earth2studio/models/px/aifsens.py	4/5	Fixed forcing/invariant filtering logic to use model data indices instead of hardcoded ranges, one pre-existing type annotation issue found but not introduced by this PR
test/models/px/test_aifsens.py	5/5	Corrected iterator test to validate initial time-step (0h) instead of skipping it

[greptile-apps]

Additional Comments (1)

earth2studio/models/px/aifsens.py
return type annotation incorrect - should be tuple[torch.Tensor, CoordSystem] not torch.Tensor

    def _fill_input(self, x: torch.Tensor, coords: CoordSystem) -> tuple[torch.Tensor, CoordSystem]:

[NickGeneva]

/blossom-ci

[loliverhennigh]

Nice! Less chance of indexing issues.

[NickGeneva]

/blossom-ci

[NickGeneva]

/blossom-ci

[NickGeneva]

/blossom-ci

[NickGeneva]

/blossom-ci