AIFS-ENS Initial time-step fix (#615)

* Fix AIFS ENS

Author: NickGeneva

earth2studio/models/px/aifsens.py

@@ -846,13 +846,15 @@ def _fill_input(self, x: torch.Tensor, coords: CoordSystem) -> torch.Tensor:
         out[:, :, 0, indices[valid_mask]] = x[0, 0, 0, ...]
         out[:, :, 1, indices[valid_mask]] = x[0, 0, 1, ...]
 
-        # Drop generated forcing dimension range from output
-        out = torch.cat([out[:, :, :, :92, ...], out[:, :, :, 101:, ...]], dim=3)
+        # Drop generated forcing / invariants from output
+        all_indices = torch.arange(len(VARIABLES))
+        keep = torch.isin(
+            all_indices, self.model.data_indices.data.output.forcing, invert=True
+        )
+        out = out[:, :, :, keep, ...]
 
         # Update coordinates with remaining variable names
-        all_indices = torch.arange(len(VARIABLES))
-        variable_mask = ~torch.isin(all_indices, self.forcing_ids)
-        selected_variables = [VARIABLES[i] for i in all_indices[variable_mask].tolist()]
+        selected_variables = [VARIABLES[i] for i in all_indices[keep].tolist()]
 
         out_coords = coords.copy()
         out_coords["variable"] = np.array(selected_variables)

test/models/px/test_aifsens.py

@@ -235,7 +235,6 @@ def test_aifsens_iter(ensemble, device):
     if not isinstance(time, Iterable):
         time = [time]
 
-    next(p_iter)
     for i, (out, out_coords) in enumerate(p_iter):
         assert len(out.shape) == 6
         assert out.shape == torch.Size(
@@ -245,7 +244,7 @@ def test_aifsens_iter(ensemble, device):
             out_coords["variable"] == p.output_coords(p.input_coords())["variable"]
         ).all()
         assert (out_coords["ensemble"] == np.arange(ensemble)).all()
-        assert out_coords["lead_time"][0] == np.timedelta64(6 * (i + 1), "h")
+        assert out_coords["lead_time"][0] == np.timedelta64(6 * (i), "h")
 
         if i > 5:
             break