diff --git a/external/loaders/loaders/batches/_batch.py b/external/loaders/loaders/batches/_batch.py index 8b17468258..8d0daf6899 100644 --- a/external/loaders/loaders/batches/_batch.py +++ b/external/loaders/loaders/batches/_batch.py @@ -256,9 +256,26 @@ def _get_batch(mapper: Mapping[str, xr.Dataset], keys: Iterable[str],) -> xr.Dat return ds +def _add_dQ1_dQ2(ds): + ds["dQ1"] = ( + ds.air_temperature + - ds.air_temperature_before_interval_update_at_next_time_step.shift(time=1) + ) / 900.0 + ds["dQ2"] = ( + ds.specific_humidity + - ds.specific_humidity_before_interval_update_at_next_time_step.shift(time=1) + ) / 900.0 + return ds + + @curry def _open_dataset(fs: fsspec.AbstractFileSystem, variable_names, filename): - return xr.open_dataset(fs.open(filename), engine="h5netcdf")[variable_names] + ds = xr.open_dataset(fs.open(filename), engine="h5netcdf")[variable_names] + try: + ds = _add_dQ1_dQ2(ds) + except (AttributeError): + pass + return ds @batches_functions.register