ocf-data-sampler 0.0.24__tar.gz → 0.0.25__tar.gz

{ocf_data_sampler-0.0.24/ocf_data_sampler.egg-info → ocf_data_sampler-0.0.25}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ocf_data_sampler
-Version: 0.0.24
+Version: 0.0.25
 Summary: Sample from weather data for renewable energy prediction
 Author: James Fulton, Peter Dudfield, and the Open Climate Fix team
 Author-email: info@openclimatefix.org

{ocf_data_sampler-0.0.24 → ocf_data_sampler-0.0.25}/ocf_data_sampler/config/model.py

@@ -102,6 +102,39 @@ class TimeResolutionMixin(Base):
     )
 
 
+class Site(DataSourceMixin, TimeResolutionMixin, DropoutMixin):
+    """Site configuration model"""
+
+    file_path: str = Field(
+        ...,
+        description="The NetCDF files holding the power timeseries.",
+    )
+    metadata_file_path: str = Field(
+        ...,
+        description="The CSV files describing power system",
+    )
+
+    @field_validator("forecast_minutes")
+    def forecast_minutes_divide_by_time_resolution(cls, v: int, info: ValidationInfo) -> int:
+        """Check forecast length requested will give stable number of timesteps"""
+        if v % info.data["time_resolution_minutes"] != 0:
+            message = "Forecast duration must be divisible by time resolution"
+            logger.error(message)
+            raise Exception(message)
+        return v
+
+    @field_validator("history_minutes")
+    def history_minutes_divide_by_time_resolution(cls, v: int, info: ValidationInfo) -> int:
+        """Check history length requested will give stable number of timesteps"""
+        if v % info.data["time_resolution_minutes"] != 0:
+            message = "History duration must be divisible by time resolution"
+            logger.error(message)
+            raise Exception(message)
+        return v
+
+    # TODO validate the netcdf for sites
+    # TODO validate the csv for metadata
+
 class Satellite(DataSourceMixin, TimeResolutionMixin, DropoutMixin):
     """Satellite configuration model"""
 
@@ -240,6 +273,7 @@ class InputData(Base):
     satellite: Optional[Satellite] = None
     nwp: Optional[MultiNWP] = None
     gsp: Optional[GSP] = None
+    site: Optional[Site] = None
 
 
 class Configuration(Base):

ocf_data_sampler-0.0.25/ocf_data_sampler/load/load_dataset.py

@@ -0,0 +1,55 @@
+""" Loads all data sources """
+import xarray as xr
+
+from ocf_data_sampler.config import Configuration
+from ocf_data_sampler.load.gsp import open_gsp
+from ocf_data_sampler.load.nwp import open_nwp
+from ocf_data_sampler.load.satellite import open_sat_data
+from ocf_data_sampler.load.site import open_site
+
+
+def get_dataset_dict(config: Configuration) -> dict[str, dict[xr.DataArray]]:
+    """Construct dictionary of all of the input data sources
+
+    Args:
+        config: Configuration file
+    """
+
+    in_config = config.input_data
+
+    datasets_dict = {}
+
+    # Load GSP data unless the path is None
+    if in_config.gsp and in_config.gsp.gsp_zarr_path:
+        da_gsp = open_gsp(zarr_path=in_config.gsp.gsp_zarr_path).compute()
+
+        # Remove national GSP
+        datasets_dict["gsp"] = da_gsp.sel(gsp_id=slice(1, None))
+
+    # Load NWP data if in config
+    if in_config.nwp:
+
+        datasets_dict["nwp"] = {}
+        for nwp_source, nwp_config in in_config.nwp.items():
+
+            da_nwp = open_nwp(nwp_config.nwp_zarr_path, provider=nwp_config.nwp_provider)
+
+            da_nwp = da_nwp.sel(channel=list(nwp_config.nwp_channels))
+
+            datasets_dict["nwp"][nwp_source] = da_nwp
+
+    # Load satellite data if in config
+    if in_config.satellite:
+        sat_config = config.input_data.satellite
+
+        da_sat = open_sat_data(sat_config.satellite_zarr_path)
+
+        da_sat = da_sat.sel(channel=list(sat_config.satellite_channels))
+
+        datasets_dict["sat"] = da_sat
+
+    if in_config.site:
+        da_sites = open_site(in_config.site)
+        datasets_dict["site"] = da_sites
+
+    return datasets_dict

{ocf_data_sampler-0.0.24 → ocf_data_sampler-0.0.25}/ocf_data_sampler/load/nwp/providers/ecmwf.py

@@ -9,7 +9,6 @@ from ocf_data_sampler.load.utils import (
 )
 
 
-
 def open_ifs(zarr_path: Path | str | list[Path] | list[str]) -> xr.DataArray:
     """
     Opens the ECMWF IFS NWP data
@@ -27,10 +26,14 @@ def open_ifs(zarr_path: Path | str | list[Path] | list[str]) -> xr.DataArray:
     ds = ds.rename(
         {
             "init_time": "init_time_utc",
-            "variable": "channel",
         }
     )
 
+    # LEGACY SUPPORT
+    # rename variable to channel if it exists
+    if "variable" in ds:
+        ds = ds.rename({"variable": "channel"})
+
     # Check the timestamps are unique and increasing
     check_time_unique_increasing(ds.init_time_utc)
 

ocf_data_sampler-0.0.25/ocf_data_sampler/load/site.py

@@ -0,0 +1,30 @@
+import pandas as pd
+import xarray as xr
+import numpy as np
+
+from ocf_data_sampler.config.model import Site
+
+
+def open_site(sites_config: Site) -> xr.DataArray:
+
+    # Load site generation xr.Dataset
+    data_ds = xr.open_dataset(sites_config.file_path)
+
+    # Load site generation data
+    metadata_df = pd.read_csv(sites_config.metadata_file_path, index_col="site_id")
+
+    # Add coordinates
+    ds = data_ds.assign_coords(
+        latitude=(metadata_df.latitude.to_xarray()),
+        longitude=(metadata_df.longitude.to_xarray()),
+        capacity_kwp=data_ds.capacity_kwp,
+    )
+
+    # Sanity checks
+    assert np.isfinite(data_ds.capacity_kwp.values).all()
+    assert (data_ds.capacity_kwp.values > 0).all()
+    assert metadata_df.index.is_unique
+
+    return ds.generation_kw
+
+

{ocf_data_sampler-0.0.24 → ocf_data_sampler-0.0.25}/ocf_data_sampler/numpy_batch/__init__.py

@@ -4,4 +4,5 @@ from .gsp import convert_gsp_to_numpy_batch, GSPBatchKey
 from .nwp import convert_nwp_to_numpy_batch, NWPBatchKey
 from .satellite import convert_satellite_to_numpy_batch, SatelliteBatchKey
 from .sun_position import make_sun_position_numpy_batch
+from .site import convert_site_to_numpy_batch
 

ocf_data_sampler-0.0.25/ocf_data_sampler/numpy_batch/site.py

@@ -0,0 +1,29 @@
+"""Convert site to Numpy Batch"""
+
+import xarray as xr
+
+
+class SiteBatchKey:
+
+    generation = "site"
+    site_capacity_kwp = "site_capacity_kwp"
+    site_time_utc = "site_time_utc"
+    site_t0_idx = "site_t0_idx"
+    site_solar_azimuth = "site_solar_azimuth"
+    site_solar_elevation = "site_solar_elevation"
+    site_id = "site_id"
+
+
+def convert_site_to_numpy_batch(da: xr.DataArray, t0_idx: int | None = None) -> dict:
+    """Convert from Xarray to NumpyBatch"""
+
+    example = {
+        SiteBatchKey.generation: da.values,
+        SiteBatchKey.site_capacity_kwp: da.isel(time_utc=0)["capacity_kwp"].values,
+        SiteBatchKey.site_time_utc: da["time_utc"].values.astype(float),
+    }
+
+    if t0_idx is not None:
+        example[SiteBatchKey.site_t0_idx] = t0_idx
+
+    return example

ocf_data_sampler-0.0.25/ocf_data_sampler/select/__init__.py

@@ -0,0 +1,8 @@
+from .fill_time_periods import fill_time_periods
+from .find_contiguous_time_periods import (
+    find_contiguous_t0_periods,
+    intersection_of_multiple_dataframes_of_periods,
+)
+from .location import Location
+from .spatial_slice_for_dataset import slice_datasets_by_space
+from .time_slice_for_dataset import slice_datasets_by_time

{ocf_data_sampler-0.0.24 → ocf_data_sampler-0.0.25}/ocf_data_sampler/select/dropout.py

@@ -1,3 +1,4 @@
+""" Functions for simulating dropout in time series data """
 import numpy as np
 import pandas as pd
 import xarray as xr
@@ -5,7 +6,7 @@ import xarray as xr
 
 def draw_dropout_time(
         t0: pd.Timestamp,
-        dropout_timedeltas: list[pd.Timedelta] | None,
+        dropout_timedeltas: list[pd.Timedelta] | pd.Timedelta | None,
         dropout_frac: float = 0,
     ):
 

{ocf_data_sampler-0.0.24 → ocf_data_sampler-0.0.25}/ocf_data_sampler/select/geospatial.py

@@ -55,6 +55,23 @@ def lon_lat_to_osgb(
     return _lon_lat_to_osgb(xx=x, yy=y)
 
 
+def lon_lat_to_geostationary_area_coords(
+    longitude: Union[Number, np.ndarray],
+    latitude: Union[Number, np.ndarray],
+    xr_data: xr.DataArray,
+) -> tuple[Union[Number, np.ndarray], Union[Number, np.ndarray]]:
+    """Loads geostationary area and transformation from lat-lon to geostationary coords
+
+    Args:
+        longitude: longitude
+        latitude: latitude
+        xr_data: xarray object with geostationary area
+
+    Returns:
+        Geostationary coords: x, y
+    """
+    return coordinates_to_geostationary_area_coords(longitude, latitude, xr_data, WGS84)
+
 def osgb_to_geostationary_area_coords(
     x: Union[Number, np.ndarray],
     y: Union[Number, np.ndarray],
@@ -70,6 +87,31 @@ def osgb_to_geostationary_area_coords(
     Returns:
         Geostationary coords: x, y
     """
+
+    return coordinates_to_geostationary_area_coords(x, y, xr_data, OSGB36)
+
+
+
+def coordinates_to_geostationary_area_coords(
+    x: Union[Number, np.ndarray],
+    y: Union[Number, np.ndarray],
+    xr_data: xr.DataArray,
+    crs_from: int
+) -> tuple[Union[Number, np.ndarray], Union[Number, np.ndarray]]:
+    """Loads geostationary area and transformation from respective coordiates to geostationary coords
+
+        Args:
+            x: osgb east-west, or latitude
+            y: osgb north-south, or longitude
+            xr_data: xarray object with geostationary area
+            crs_from: the cordiates system of x,y
+
+        Returns:
+            Geostationary coords: x, y
+        """
+
+    assert crs_from in [OSGB36, WGS84], f"Unrecognized coordinate system: {crs_from}"
+
     # Only load these if using geostationary projection
     import pyresample
 
@@ -80,7 +122,7 @@ def osgb_to_geostationary_area_coords(
     )
     geostationary_crs = geostationary_area_definition.crs
     osgb_to_geostationary = pyproj.Transformer.from_crs(
-        crs_from=OSGB36, crs_to=geostationary_crs, always_xy=True
+        crs_from=crs_from, crs_to=geostationary_crs, always_xy=True
     ).transform
     return osgb_to_geostationary(xx=x, yy=y)
 

{ocf_data_sampler-0.0.24 → ocf_data_sampler-0.0.25}/ocf_data_sampler/select/select_spatial_slice.py

@@ -8,6 +8,7 @@ import xarray as xr
 from ocf_data_sampler.select.location import Location
 from ocf_data_sampler.select.geospatial import (
     lon_lat_to_osgb,
+    lon_lat_to_geostationary_area_coords,
     osgb_to_geostationary_area_coords,
     osgb_to_lon_lat,
     spatial_coord_type,
@@ -101,7 +102,7 @@ def _get_idx_of_pixel_closest_to_poi(
 
 def _get_idx_of_pixel_closest_to_poi_geostationary(
     da: xr.DataArray,
-    center_osgb: Location,
+    center: Location,
 ) -> Location:
     """
     Return x and y index location of pixel at center of region of interest.
@@ -116,7 +117,12 @@ def _get_idx_of_pixel_closest_to_poi_geostationary(
 
     _, x_dim, y_dim = spatial_coord_type(da)
 
-    x, y = osgb_to_geostationary_area_coords(x=center_osgb.x, y=center_osgb.y, xr_data=da)
+    if center.coordinate_system == 'osgb':
+        x, y = osgb_to_geostationary_area_coords(x=center.x, y=center.y, xr_data=da)
+    elif center.coordinate_system == 'lon_lat':
+        x, y = lon_lat_to_geostationary_area_coords(longitude=center.x, latitude=center.y, xr_data=da)
+    else:
+        x,y = center.x, center.y
     center_geostationary = Location(x=x, y=y, coordinate_system="geostationary")
 
     # Check that the requested point lies within the data

ocf_data_sampler-0.0.25/ocf_data_sampler/select/spatial_slice_for_dataset.py

@@ -0,0 +1,53 @@
+""" Functions for selecting data around a given location """
+from ocf_data_sampler.config import Configuration
+from ocf_data_sampler.select.location import Location
+from ocf_data_sampler.select.select_spatial_slice import select_spatial_slice_pixels
+
+
+def slice_datasets_by_space(
+    datasets_dict: dict,
+    location: Location,
+    config: Configuration,
+) -> dict:
+    """Slice the dictionary of input data sources around a given location
+
+    Args:
+        datasets_dict: Dictionary of the input data sources
+        location: The location to sample around
+        config: Configuration object.
+    """
+
+    assert set(datasets_dict.keys()).issubset({"nwp", "sat", "gsp", "site"})
+
+    sliced_datasets_dict = {}
+
+    if "nwp" in datasets_dict:
+
+        sliced_datasets_dict["nwp"] = {}
+
+        for nwp_key, nwp_config in config.input_data.nwp.items():
+
+            sliced_datasets_dict["nwp"][nwp_key] = select_spatial_slice_pixels(
+                datasets_dict["nwp"][nwp_key],
+                location,
+                height_pixels=nwp_config.nwp_image_size_pixels_height,
+                width_pixels=nwp_config.nwp_image_size_pixels_width,
+            )
+
+    if "sat" in datasets_dict:
+        sat_config = config.input_data.satellite
+
+        sliced_datasets_dict["sat"] = select_spatial_slice_pixels(
+            datasets_dict["sat"],
+            location,
+            height_pixels=sat_config.satellite_image_size_pixels_height,
+            width_pixels=sat_config.satellite_image_size_pixels_width,
+        )
+
+    if "gsp" in datasets_dict:
+        sliced_datasets_dict["gsp"] = datasets_dict["gsp"].sel(gsp_id=location.id)
+
+    if "site" in datasets_dict:
+        sliced_datasets_dict["site"] = datasets_dict["site"].sel(site_id=location.id)
+
+    return sliced_datasets_dict

ocf_data_sampler-0.0.25/ocf_data_sampler/select/time_slice_for_dataset.py

@@ -0,0 +1,124 @@
+""" Slice datasets by time"""
+import pandas as pd
+
+from ocf_data_sampler.config import Configuration
+from ocf_data_sampler.select.dropout import draw_dropout_time, apply_dropout_time
+from ocf_data_sampler.select.select_time_slice import select_time_slice_nwp, select_time_slice
+from ocf_data_sampler.time_functions import minutes
+
+
+def slice_datasets_by_time(
+    datasets_dict: dict,
+    t0: pd.Timestamp,
+    config: Configuration,
+) -> dict:
+    """Slice the dictionary of input data sources around a given t0 time
+
+    Args:
+        datasets_dict: Dictionary of the input data sources
+        t0: The init-time
+        config: Configuration object.
+    """
+
+    sliced_datasets_dict = {}
+
+    if "nwp" in datasets_dict:
+
+        sliced_datasets_dict["nwp"] = {}
+
+        for nwp_key, da_nwp in datasets_dict["nwp"].items():
+
+            nwp_config = config.input_data.nwp[nwp_key]
+
+            sliced_datasets_dict["nwp"][nwp_key] = select_time_slice_nwp(
+                da_nwp,
+                t0,
+                sample_period_duration=minutes(nwp_config.time_resolution_minutes),
+                history_duration=minutes(nwp_config.history_minutes),
+                forecast_duration=minutes(nwp_config.forecast_minutes),
+                dropout_timedeltas=minutes(nwp_config.dropout_timedeltas_minutes),
+                dropout_frac=nwp_config.dropout_fraction,
+                accum_channels=nwp_config.nwp_accum_channels,
+            )
+
+    if "sat" in datasets_dict:
+
+        sat_config = config.input_data.satellite
+
+        sliced_datasets_dict["sat"] = select_time_slice(
+            datasets_dict["sat"],
+            t0,
+            sample_period_duration=minutes(sat_config.time_resolution_minutes),
+            interval_start=minutes(-sat_config.history_minutes),
+            interval_end=minutes(-sat_config.live_delay_minutes),
+            max_steps_gap=2,
+        )
+
+        # Randomly sample dropout
+        sat_dropout_time = draw_dropout_time(
+            t0,
+            dropout_timedeltas=minutes(sat_config.dropout_timedeltas_minutes),
+            dropout_frac=sat_config.dropout_fraction,
+        )
+
+        # Apply the dropout
+        sliced_datasets_dict["sat"] = apply_dropout_time(
+            sliced_datasets_dict["sat"],
+            sat_dropout_time,
+        )
+
+    if "gsp" in datasets_dict:
+        gsp_config = config.input_data.gsp
+
+        sliced_datasets_dict["gsp_future"] = select_time_slice(
+            datasets_dict["gsp"],
+            t0,
+            sample_period_duration=minutes(gsp_config.time_resolution_minutes),
+            interval_start=minutes(30),
+            interval_end=minutes(gsp_config.forecast_minutes),
+        )
+
+        sliced_datasets_dict["gsp"] = select_time_slice(
+            datasets_dict["gsp"],
+            t0,
+            sample_period_duration=minutes(gsp_config.time_resolution_minutes),
+            interval_start=-minutes(gsp_config.history_minutes),
+            interval_end=minutes(0),
+        )
+
+        # Dropout on the GSP, but not the future GSP
+        gsp_dropout_time = draw_dropout_time(
+            t0,
+            dropout_timedeltas=minutes(gsp_config.dropout_timedeltas_minutes),
+            dropout_frac=gsp_config.dropout_fraction,
+        )
+
+        sliced_datasets_dict["gsp"] = apply_dropout_time(
+            sliced_datasets_dict["gsp"], gsp_dropout_time
+        )
+
+    if "site" in datasets_dict:
+        site_config = config.input_data.site
+
+        sliced_datasets_dict["site"] = select_time_slice(
+            datasets_dict["site"],
+            t0,
+            sample_period_duration=minutes(site_config.time_resolution_minutes),
+            interval_start=-minutes(site_config.history_minutes),
+            interval_end=minutes(site_config.forecast_minutes),
+        )
+
+        # Randomly sample dropout
+        site_dropout_time = draw_dropout_time(
+            t0,
+            dropout_timedeltas=minutes(site_config.dropout_timedeltas_minutes),
+            dropout_frac=site_config.dropout_fraction,
+        )
+
+        # Apply the dropout
+        sliced_datasets_dict["site"] = apply_dropout_time(
+            sliced_datasets_dict["site"],
+            site_dropout_time,
+        )
+
+    return sliced_datasets_dict

ocf_data_sampler-0.0.25/ocf_data_sampler/time_functions.py

@@ -0,0 +1,11 @@
+import pandas as pd
+
+
+def minutes(minutes: int | list[float]) -> pd.Timedelta | pd.TimedeltaIndex:
+    """Timedelta minutes
+
+    Args:
+        minutes: the number of minutes, single value or list
+    """
+    minutes_delta = pd.to_timedelta(minutes, unit="m")
+    return minutes_delta

ocf_data_sampler-0.0.25/ocf_data_sampler/torch_datasets/process_and_combine.py

@@ -0,0 +1,153 @@
+import numpy as np
+import pandas as pd
+import xarray as xr
+
+from ocf_data_sampler.config import Configuration
+from ocf_data_sampler.constants import NWP_MEANS, NWP_STDS
+from ocf_data_sampler.numpy_batch import (
+    convert_nwp_to_numpy_batch,
+    convert_satellite_to_numpy_batch,
+    convert_gsp_to_numpy_batch,
+    make_sun_position_numpy_batch,
+    convert_site_to_numpy_batch,
+)
+from ocf_data_sampler.numpy_batch.gsp import GSPBatchKey
+from ocf_data_sampler.numpy_batch.nwp import NWPBatchKey
+from ocf_data_sampler.select.geospatial import osgb_to_lon_lat
+from ocf_data_sampler.select.location import Location
+from ocf_data_sampler.time_functions import minutes
+
+
+def process_and_combine_datasets(
+    dataset_dict: dict,
+    config: Configuration,
+    t0: pd.Timestamp,
+    location: Location,
+    sun_position_key: str = 'gsp'
+) -> dict:
+    """Normalize and convert data to numpy arrays"""
+
+    numpy_modalities = []
+
+    if "nwp" in dataset_dict:
+
+        nwp_numpy_modalities = dict()
+
+        for nwp_key, da_nwp in dataset_dict["nwp"].items():
+            # Standardise
+            provider = config.input_data.nwp[nwp_key].nwp_provider
+            da_nwp = (da_nwp - NWP_MEANS[provider]) / NWP_STDS[provider]
+            # Convert to NumpyBatch
+            nwp_numpy_modalities[nwp_key] = convert_nwp_to_numpy_batch(da_nwp)
+
+        # Combine the NWPs into NumpyBatch
+        numpy_modalities.append({NWPBatchKey.nwp: nwp_numpy_modalities})
+
+    if "sat" in dataset_dict:
+        # Satellite is already in the range [0-1] so no need to standardise
+        da_sat = dataset_dict["sat"]
+
+        # Convert to NumpyBatch
+        numpy_modalities.append(convert_satellite_to_numpy_batch(da_sat))
+
+    gsp_config = config.input_data.gsp
+
+    if "gsp" in dataset_dict:
+        da_gsp = xr.concat([dataset_dict["gsp"], dataset_dict["gsp_future"]], dim="time_utc")
+        da_gsp = da_gsp / da_gsp.effective_capacity_mwp
+
+        numpy_modalities.append(
+            convert_gsp_to_numpy_batch(
+                da_gsp, t0_idx=gsp_config.history_minutes // gsp_config.time_resolution_minutes
+            )
+        )
+
+        # Add coordinate data
+        # TODO: Do we need all of these?
+        numpy_modalities.append(
+            {
+                GSPBatchKey.gsp_id: location.id,
+                GSPBatchKey.x_osgb: location.x,
+                GSPBatchKey.y_osgb: location.y,
+            }
+        )
+
+
+    if "site" in dataset_dict:
+        site_config = config.input_data.site
+        da_sites = dataset_dict["site"]
+        da_sites = da_sites / da_sites.capacity_kwp
+
+        numpy_modalities.append(
+            convert_site_to_numpy_batch(
+                da_sites, t0_idx=site_config.history_minutes / site_config.time_resolution_minutes
+            )
+        )
+
+    if sun_position_key == 'gsp':
+        # Make sun coords NumpyBatch
+        datetimes = pd.date_range(
+            t0 - minutes(gsp_config.history_minutes),
+            t0 + minutes(gsp_config.forecast_minutes),
+            freq=minutes(gsp_config.time_resolution_minutes),
+        )
+
+        lon, lat = osgb_to_lon_lat(location.x, location.y)
+        key_prefix = "gsp"
+
+    elif sun_position_key == 'site':
+        # Make sun coords NumpyBatch
+        datetimes = pd.date_range(
+            t0 - minutes(site_config.history_minutes),
+            t0 + minutes(site_config.forecast_minutes),
+            freq=minutes(site_config.time_resolution_minutes),
+        )
+
+        lon, lat = location.x, location.y
+        key_prefix = "site"
+
+    numpy_modalities.append(
+        make_sun_position_numpy_batch(datetimes, lon, lat, key_prefix=key_prefix)
+    )
+
+    # Combine all the modalities and fill NaNs
+    combined_sample = merge_dicts(numpy_modalities)
+    combined_sample = fill_nans_in_arrays(combined_sample)
+
+    return combined_sample
+
+
+def merge_dicts(list_of_dicts: list[dict]) -> dict:
+    """Merge a list of dictionaries into a single dictionary"""
+    # TODO: This doesn't account for duplicate keys, which will be overwritten
+    combined_dict = {}
+    for d in list_of_dicts:
+        combined_dict.update(d)
+    return combined_dict
+
+
+def fill_nans_in_arrays(batch: dict) -> dict:
+    """Fills all NaN values in each np.ndarray in the batch dictionary with zeros.
+
+    Operation is performed in-place on the batch.
+    """
+    for k, v in batch.items():
+        if isinstance(v, np.ndarray) and np.issubdtype(v.dtype, np.number):
+            if np.isnan(v).any():
+                batch[k] = np.nan_to_num(v, copy=False, nan=0.0)
+
+        # Recursion is included to reach NWP arrays in subdict
+        elif isinstance(v, dict):
+            fill_nans_in_arrays(v)
+
+    return batch
+
+
+def compute(xarray_dict: dict) -> dict:
+    """Eagerly load a nested dictionary of xarray DataArrays"""
+    for k, v in xarray_dict.items():
+        if isinstance(v, dict):
+            xarray_dict[k] = compute(v)
+        else:
+            xarray_dict[k] = v.compute(scheduler="single-threaded")
+    return xarray_dict