Changed PressureLevelModel.data_from_xarray to add sim_time to data…

…set if not included. PiperOrigin-RevId: 612952757
google-research · Mar 5, 2024 · 2f0a834 · 2f0a834
1 parent 29c6bd8
commit 2f0a834
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/neuralgcm/api.py b/neuralgcm/api.py
@@ -146,6 +146,12 @@ def data_from_xarray(
       self, dataset: xarray.Dataset
   ) -> tuple[Inputs, Forcings]:
     """Extracts data and forcings from xarray.Dataset."""
+    ref_datetime = self._structure.specs.aux_features['reference_datetime']
+    dataset = xarray_utils.ds_with_sim_time(
+        dataset,
+        self._structure.specs.physics_specs,
+        reference_datetime=ref_datetime,
+    )
     dataset_coords = model_builder.coordinate_system_from_dataset(dataset)
     if not np.allclose(
         dataset_coords.horizontal.longitudes,