ocf-data-sampler 0.0.44__py3-none-any.whl → 0.0.46__py3-none-any.whl

ocf_data_sampler/config/save.py

@@ -9,7 +9,6 @@ Example:
 """
 
 import json
-
 from pathlib import Path
 from typing import Union
 
@@ -18,7 +17,6 @@ import yaml
 
 from ocf_data_sampler.config import Configuration
 
-
 def save_yaml_configuration(
     configuration: Configuration,
     filename: Union[str, Path],
@@ -35,7 +33,7 @@ def save_yaml_configuration(
         Path: The path where the configuration was saved
 
     Raises:
-        ValueError: If filename is None or if writing to the specified path fails
+        ValueError: If filename is None, directory doesn't exist, or if writing to the specified path fails
         TypeError: If the configuration cannot be serialized
     """
     if filename is None:
@@ -50,24 +48,37 @@ def save_yaml_configuration(
 
         filepath = Path(filename)
 
-        # For local files, check if directory exists before proceeding
+        # For local paths, check if parent directory exists before attempting to create
         if filepath.is_absolute():
-            directory = filepath.parent
-            if not directory.exists():
+            if not filepath.parent.exists():
                 raise ValueError("Directory does not exist")
+            
+            # Only try to create directory if it's in a writable location
+            try:
+                filepath.parent.mkdir(parents=True, exist_ok=True)
+            except PermissionError:
+                raise ValueError(f"Permission denied when accessing directory {filepath.parent}")
 
         # Serialize configuration to JSON-compatible dictionary
         config_dict = json.loads(configuration.model_dump_json())
 
-        # Save to YAML file using fsspec
-        with fsspec.open(str(filepath), mode='w') as yaml_file:
-            yaml.safe_dump(config_dict, yaml_file, default_flow_style=False)
+        # Write to file directly for local paths
+        if filepath.is_absolute():
+            try:
+                with open(filepath, 'w') as f:
+                    yaml.safe_dump(config_dict, f, default_flow_style=False)
+            except PermissionError:
+                raise ValueError(f"Permission denied when writing to {filename}")
+        else:
+            # Use fsspec for cloud storage
+            with fsspec.open(str(filepath), mode='w') as yaml_file:
+                yaml.safe_dump(config_dict, yaml_file, default_flow_style=False)
 
         return filepath
 
     except json.JSONDecodeError as e:
         raise TypeError(f"Failed to serialize configuration: {str(e)}") from e
-    except PermissionError as e:
-        raise ValueError(f"Permission denied when writing to {filename}") from e
     except (IOError, OSError) as e:
+        if "Permission denied" in str(e):
+            raise ValueError(f"Permission denied when writing to {filename}") from e
         raise ValueError(f"Failed to write configuration to {filename}: {str(e)}") from e

ocf_data_sampler/numpy_sample/__init__.py

@@ -1,5 +1,6 @@
 """Conversion from Xarray to NumpySample"""
 
+from .datetime_features import make_datetime_numpy_dict
 from .gsp import convert_gsp_to_numpy_sample, GSPSampleKey
 from .nwp import convert_nwp_to_numpy_sample, NWPSampleKey
 from .satellite import convert_satellite_to_numpy_sample, SatelliteSampleKey

ocf_data_sampler/numpy_sample/datetime_features.py

@@ -0,0 +1,46 @@
+"""Functions to create trigonometric date and time inputs"""
+
+import numpy as np
+import pandas as pd
+from numpy.typing import NDArray
+
+
+def _get_date_time_in_pi(
+    dt: pd.DatetimeIndex,
+) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
+    """
+    Change the datetimes, into time and date scaled in radians
+    """
+
+    day_of_year = dt.dayofyear
+    minute_of_day = dt.minute + dt.hour * 60
+
+    # converting into positions on sin-cos circle
+    time_in_pi = (2 * np.pi) * (minute_of_day / (24 * 60))
+    date_in_pi = (2 * np.pi) * (day_of_year / 365)
+
+    return date_in_pi, time_in_pi
+
+
+def make_datetime_numpy_dict(datetimes: pd.DatetimeIndex, key_prefix: str = "wind") -> dict:
+    """ Make dictionary of datetime features"""
+
+    if datetimes.empty:
+        raise ValueError("Input datetimes is empty for 'make_datetime_numpy_dict' function")
+
+    time_numpy_sample = {}
+
+    date_in_pi, time_in_pi = _get_date_time_in_pi(datetimes)
+
+    # Store
+    date_sin_batch_key = key_prefix + "_date_sin"
+    date_cos_batch_key = key_prefix + "_date_cos"
+    time_sin_batch_key = key_prefix + "_time_sin"
+    time_cos_batch_key = key_prefix + "_time_cos"
+
+    time_numpy_sample[date_sin_batch_key] = np.sin(date_in_pi)
+    time_numpy_sample[date_cos_batch_key] = np.cos(date_in_pi)
+    time_numpy_sample[time_sin_batch_key] = np.sin(time_in_pi)
+    time_numpy_sample[time_cos_batch_key] = np.cos(time_in_pi)
+
+    return time_numpy_sample

ocf_data_sampler/torch_datasets/{pvnet_uk_regional.py → datasets/pvnet_uk_regional.py}

@@ -5,16 +5,114 @@ import pandas as pd
 import pkg_resources
 import xarray as xr
 from torch.utils.data import Dataset
-
 from ocf_data_sampler.config import Configuration, load_yaml_configuration
 from ocf_data_sampler.load.load_dataset import get_dataset_dict
 from ocf_data_sampler.select import fill_time_periods, Location, slice_datasets_by_space, slice_datasets_by_time
 from ocf_data_sampler.utils import minutes
-from ocf_data_sampler.torch_datasets.process_and_combine import process_and_combine_datasets, compute
-from ocf_data_sampler.torch_datasets.valid_time_periods import find_valid_time_periods
+from ocf_data_sampler.torch_datasets.utils.valid_time_periods import find_valid_time_periods
+from ocf_data_sampler.constants import NWP_MEANS, NWP_STDS, RSS_MEAN, RSS_STD
+from ocf_data_sampler.numpy_sample import (
+    convert_nwp_to_numpy_sample,
+    convert_satellite_to_numpy_sample,
+    convert_gsp_to_numpy_sample,
+    make_sun_position_numpy_sample,
+)
+from ocf_data_sampler.torch_datasets.utils.merge_and_fill_utils import (
+    merge_dicts,
+    fill_nans_in_arrays,
+)
+from ocf_data_sampler.numpy_sample.gsp import GSPSampleKey
+from ocf_data_sampler.numpy_sample.nwp import NWPSampleKey
+from ocf_data_sampler.select.geospatial import osgb_to_lon_lat
 
 xr.set_options(keep_attrs=True)
 
+def process_and_combine_datasets(
+    dataset_dict: dict,
+    config: Configuration,
+    t0: pd.Timestamp,
+    location: Location,
+    target_key: str = 'gsp'
+) -> dict:
+
+    """Normalise 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].provider
+            da_nwp = (da_nwp - NWP_MEANS[provider]) / NWP_STDS[provider]
+
+            # Convert to NumpyBatch
+            nwp_numpy_modalities[nwp_key] = convert_nwp_to_numpy_sample(da_nwp)
+
+        # Combine the NWPs into NumpyBatch
+        numpy_modalities.append({NWPSampleKey.nwp: nwp_numpy_modalities})
+
+
+    if "sat" in dataset_dict:
+        # Standardise
+        da_sat = dataset_dict["sat"]
+        da_sat = (da_sat - RSS_MEAN) / RSS_STD
+
+        # Convert to NumpyBatch
+        numpy_modalities.append(convert_satellite_to_numpy_sample(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_sample(
+                da_gsp, 
+                t0_idx=-gsp_config.interval_start_minutes / gsp_config.time_resolution_minutes
+            )
+        )
+
+        # Add coordinate data
+        # TODO: Do we need all of these?
+        numpy_modalities.append(
+            {
+                GSPSampleKey.gsp_id: location.id,
+                GSPSampleKey.x_osgb: location.x,
+                GSPSampleKey.y_osgb: location.y,
+            }
+        )
+
+    if target_key == 'gsp':
+        # Make sun coords NumpySample
+        datetimes = pd.date_range(
+            t0+minutes(gsp_config.interval_start_minutes),
+            t0+minutes(gsp_config.interval_end_minutes),
+            freq=minutes(gsp_config.time_resolution_minutes),
+        )
+
+        lon, lat = osgb_to_lon_lat(location.x, location.y)
+
+    numpy_modalities.append(
+        make_sun_position_numpy_sample(datetimes, lon, lat, key_prefix=target_key)
+    )
+
+    # 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 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
 
 def find_valid_t0_times(
     datasets_dict: dict,
@@ -48,7 +146,7 @@ def get_gsp_locations(gsp_ids: list[int] | None = None) -> list[Location]:
 
     # Load UK GSP locations
     df_gsp_loc = pd.read_csv(
-        pkg_resources.resource_filename(__name__, "../data/uk_gsp_locations.csv"),
+        pkg_resources.resource_filename(__name__, "../../data/uk_gsp_locations.csv"),
         index_col="gsp_id",
     )
 

ocf_data_sampler/torch_datasets/{site.py → datasets/site.py}

@@ -17,12 +17,14 @@ from ocf_data_sampler.select import (
     slice_datasets_by_time, slice_datasets_by_space
 )
 from ocf_data_sampler.utils import minutes
-from ocf_data_sampler.torch_datasets.valid_time_periods import find_valid_time_periods
-from ocf_data_sampler.torch_datasets.process_and_combine import merge_dicts, fill_nans_in_arrays
+from ocf_data_sampler.torch_datasets.utils.valid_time_periods import find_valid_time_periods
+from ocf_data_sampler.torch_datasets.utils.merge_and_fill_utils import merge_dicts, fill_nans_in_arrays
 from ocf_data_sampler.numpy_sample import (
     convert_site_to_numpy_sample, 
     convert_satellite_to_numpy_sample, 
-    convert_nwp_to_numpy_sample
+    convert_nwp_to_numpy_sample,
+    make_datetime_numpy_dict,
+    make_sun_position_numpy_sample,
 )
 from ocf_data_sampler.numpy_sample import NWPSampleKey
 from ocf_data_sampler.constants import NWP_MEANS, NWP_STDS
@@ -234,10 +236,26 @@ class SitesDataset(Dataset):
             da_sites = dataset_dict["site"]
             da_sites = da_sites / da_sites.capacity_kwp
             data_arrays.append(("site", da_sites))
-        
+
         combined_sample_dataset = self.merge_data_arrays(data_arrays)
 
-        # TODO add solar + time features for sites
+        # add datetime features
+        datetimes = pd.DatetimeIndex(combined_sample_dataset.site__time_utc.values)
+        datetime_features = make_datetime_numpy_dict(datetimes=datetimes, key_prefix="site")
+        datetime_features_xr = xr.Dataset(datetime_features, coords={"site__time_utc": datetimes})
+        combined_sample_dataset = xr.merge([combined_sample_dataset, datetime_features_xr])
+
+        # add sun features
+        sun_position_features = make_sun_position_numpy_sample(
+            datetimes=datetimes,
+            lon=combined_sample_dataset.site__longitude.values,
+            lat=combined_sample_dataset.site__latitude.values,
+            key_prefix="site",
+        )
+        sun_position_features_xr = xr.Dataset(
+            sun_position_features, coords={"site__time_utc": datetimes}
+        )
+        combined_sample_dataset = xr.merge([combined_sample_dataset, sun_position_features_xr])
 
         # Fill any nan values
         return combined_sample_dataset.fillna(0.0)

ocf_data_sampler/torch_datasets/utils/merge_and_fill_utils.py

@@ -0,0 +1,25 @@
+import numpy as np
+
+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(sample: dict) -> dict:
+    """Fills all NaN values in each np.ndarray in the sample dictionary with zeros.
+
+    Operation is performed in-place on the sample.
+    """
+    for k, v in sample.items():
+        if isinstance(v, np.ndarray) and np.issubdtype(v.dtype, np.number):
+            if np.isnan(v).any():
+                sample[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 sample

{ocf_data_sampler-0.0.44.dist-info → ocf_data_sampler-0.0.46.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: ocf_data_sampler
-Version: 0.0.44
+Version: 0.0.46
 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
@@ -56,7 +56,7 @@ Requires-Dist: mkdocs-material>=8.0; extra == "docs"
 # ocf-data-sampler
  
 <!-- ALL-CONTRIBUTORS-BADGE:START - Do not remove or modify this section -->
-[![All Contributors](https://img.shields.io/badge/all_contributors-10-orange.svg?style=flat-square)](#contributors-)
+[![All Contributors](https://img.shields.io/badge/all_contributors-11-orange.svg?style=flat-square)](#contributors-)
 <!-- ALL-CONTRIBUTORS-BADGE:END -->
 
 [![tags badge](https://img.shields.io/github/v/tag/openclimatefix/ocf-data-sampler?include_prereleases&sort=semver&color=FFAC5F)](https://github.com/openclimatefix/ocf-data-sampler/tags)
@@ -135,6 +135,7 @@ Thanks goes to these wonderful people ([emoji key](https://allcontributors.org/d
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/felix-e-h-p"><img src="https://avatars.githubusercontent.com/u/137530077?v=4?s=100" width="100px;" alt="Felix"/><br /><sub><b>Felix</b></sub></a><br /><a href="https://github.com/openclimatefix/ocf-data-sampler/commits?author=felix-e-h-p" title="Code">💻</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://timothyajaniportfolio-b6v3zq29k-timthegreat.vercel.app/"><img src="https://avatars.githubusercontent.com/u/60073728?v=4?s=100" width="100px;" alt="Ajani Timothy"/><br /><sub><b>Ajani Timothy</b></sub></a><br /><a href="https://github.com/openclimatefix/ocf-data-sampler/commits?author=Tim1119" title="Code">💻</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://rupeshmangalam.vercel.app/"><img src="https://avatars.githubusercontent.com/u/91172425?v=4?s=100" width="100px;" alt="Rupesh Mangalam"/><br /><sub><b>Rupesh Mangalam</b></sub></a><br /><a href="https://github.com/openclimatefix/ocf-data-sampler/commits?author=RupeshMangalam21" title="Code">💻</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="http://siddharth7113.github.io"><img src="https://avatars.githubusercontent.com/u/114160268?v=4?s=100" width="100px;" alt="Siddharth"/><br /><sub><b>Siddharth</b></sub></a><br /><a href="https://github.com/openclimatefix/ocf-data-sampler/commits?author=siddharth7113" title="Code">💻</a></td>
     </tr>
   </tbody>
 </table>

{ocf_data_sampler-0.0.44.dist-info → ocf_data_sampler-0.0.46.dist-info}/RECORD

@@ -4,7 +4,7 @@ ocf_data_sampler/utils.py,sha256=rKA0BHAyAG4f90zEcgxp25EEYrXS-aOVNzttZ6Mzv2k,250
 ocf_data_sampler/config/__init__.py,sha256=YXnAkgHViHB26hSsjiv32b6EbpG-A1kKTkARJf0_RkY,212
 ocf_data_sampler/config/load.py,sha256=4f7vPHAIAmd-55tPxoIzn7F_TI_ue4NxkDcLPoVWl0g,943
 ocf_data_sampler/config/model.py,sha256=sXmh7IadwXDT-7lxEl5_b3vjovZgZYR77EXy4GHaf4w,7276
-ocf_data_sampler/config/save.py,sha256=JhPyX1DJAW-mi3GfkDF4nlrEXNt-yfX9ibBy8CZ6ENI,2464
+ocf_data_sampler/config/save.py,sha256=gB44isAZWUlCe3L6VBkLkngWC9GFpcCfAM57gy-0dkg,3156
 ocf_data_sampler/data/uk_gsp_locations.csv,sha256=RSh7DRh55E3n8lVAaWXGTaXXHevZZtI58td4d4DhGos,10415772
 ocf_data_sampler/load/__init__.py,sha256=MjgfxilTzyz1RYFoBEeAXmE9hyjknLvdmlHPmlAoiQY,44
 ocf_data_sampler/load/gsp.py,sha256=Gcr1JVUOPKhFRDCSHtfPDjxx0BtyyEhXrZvGEKLPJ5I,759
@@ -18,8 +18,9 @@ ocf_data_sampler/load/nwp/providers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQ
 ocf_data_sampler/load/nwp/providers/ecmwf.py,sha256=2iR1Iy542lo51rC6XFLV-3pbUE68dWjlHa6TVJzx3ac,1280
 ocf_data_sampler/load/nwp/providers/ukv.py,sha256=79Bm7q-K_GJPYMy62SUIZbRWRF4-tIaB1dYPEgLD9vo,1207
 ocf_data_sampler/load/nwp/providers/utils.py,sha256=Sy2exG1wpXLLhMXYdsfR-DZMR3txG1_bBmBdchlc-yA,848
-ocf_data_sampler/numpy_sample/__init__.py,sha256=jWvN4BDYAgoFkdGBxjNrYioZdZNVbIcP4RsNGJCEGm0,345
+ocf_data_sampler/numpy_sample/__init__.py,sha256=nY5C6CcuxiWZ_jrXRzWtN7WyKXhJImSiVTIG6Rz4B_4,401
 ocf_data_sampler/numpy_sample/collate.py,sha256=y8QFUhskaAfOMP22aVkexwyGAwLHbNE-q1pOZ6txWKA,2227
+ocf_data_sampler/numpy_sample/datetime_features.py,sha256=U-9uRplfZ7VYFA4qBduI8OkG2x_65RYIP8wrLG4i-Nw,1441
 ocf_data_sampler/numpy_sample/gsp.py,sha256=5UaWO_aGRRVQo82wnDaT4zBKHihOnIsXiwgPjM8vGFM,1005
 ocf_data_sampler/numpy_sample/nwp.py,sha256=_seQNWsut3IzPsrpipqImjnaM3XNHZCy5_5be6syivk,1297
 ocf_data_sampler/numpy_sample/satellite.py,sha256=8OaTvkPjzSjotcdKsa6BKmmlBKDBunbhDN4Pjo0Grxs,910
@@ -35,21 +36,22 @@ ocf_data_sampler/select/select_spatial_slice.py,sha256=WNxwur9Q5oetvogATw8-hNejD
 ocf_data_sampler/select/select_time_slice.py,sha256=D5P_cSvnv8Qs49K5au7lPxDr9U_VmDn42s5leMzHt0k,6122
 ocf_data_sampler/select/spatial_slice_for_dataset.py,sha256=3tRrMBXr7s4CnClbVSIq7hpls3H4Y3qYTDwswcxCCCE,1763
 ocf_data_sampler/select/time_slice_for_dataset.py,sha256=LMw8KnOCKnPjD0m4UubAWERpaiQtzRKkI2cSh5a0A-M,4335
-ocf_data_sampler/torch_datasets/__init__.py,sha256=nJUa2KzVa84ZoM0PT2AbDz26ennmAYc7M7WJVfypPMs,85
-ocf_data_sampler/torch_datasets/process_and_combine.py,sha256=ygRBf70Jd6turmwLrv0Fsxo_Ge-54N5S2NUntyRxtJk,4305
-ocf_data_sampler/torch_datasets/pvnet_uk_regional.py,sha256=QRFqbdfNchVWj4y70n-rJdFvFGvQj-WpZLdFqWjnOTw,5543
-ocf_data_sampler/torch_datasets/site.py,sha256=2fN2tw5tCXoOkh8vzRWWm_5FNGglGwJInTanJgV4Xuo,14455
-ocf_data_sampler/torch_datasets/valid_time_periods.py,sha256=Qo65qUHtle_bW5tLTYr7empHTRv-lpjvfx_6GNJj3Xg,4371
+ocf_data_sampler/torch_datasets/datasets/__init__.py,sha256=nJUa2KzVa84ZoM0PT2AbDz26ennmAYc7M7WJVfypPMs,85
+ocf_data_sampler/torch_datasets/datasets/pvnet_uk_regional.py,sha256=xxeX4Js9LQpydehi3BS7k9psqkYGzgJuM17uTYux40M,8742
+ocf_data_sampler/torch_datasets/datasets/site.py,sha256=7gTtXG3DFzs_0XlYK0oleFPT-Gena_NSngcG_FAnY54,15394
+ocf_data_sampler/torch_datasets/utils/merge_and_fill_utils.py,sha256=hIbekql64eXsNDFIoEc--GWxwdVWrh2qKegdOi70Bow,874
+ocf_data_sampler/torch_datasets/utils/valid_time_periods.py,sha256=Qo65qUHtle_bW5tLTYr7empHTRv-lpjvfx_6GNJj3Xg,4371
 scripts/refactor_site.py,sha256=asZ27hQ4IyXgCCUaFJqcz1ObBNcV2W3ywqHBpSXA_fc,1728
 tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tests/conftest.py,sha256=68hH-HPdHPLvLrtYJU8bjfkdGKbhPfNveLKvUs6_Lr0,7970
-tests/config/test_config.py,sha256=eaye_F7-el4tTP4n2vRME8qlV0b2jaKUX4HhgOUpa7E,5203
+tests/conftest.py,sha256=DfrH0Pm552Tnl35eZn2UHCfOn2lHRiEQCcUCJIhycSU,8021
+tests/config/test_config.py,sha256=Vq_kTL5tJcwEP-hXD_Nah5O6cgafo99iX6Fw1AN5NDY,5288
 tests/config/test_save.py,sha256=rA_XVxP1pOxB--5Ebujz4T5o-VbcrCbg2VSlSq2iI0o,1318
 tests/load/test_load_gsp.py,sha256=aT_nqaSXmUTcdHzuTT7AmXJr3R31k4OEN-Fv3eLxlQE,424
 tests/load/test_load_nwp.py,sha256=3qyyDkB1q9t3tyAwogfotNrxqUOpXXimco1CImoEWGg,753
 tests/load/test_load_satellite.py,sha256=STX5AqqmOAgUgE9R1xyq_sM3P1b8NKdGjO-hDhayfxM,524
 tests/load/test_load_sites.py,sha256=T9lSEnGPI8FQISudVYHHNTHeplNS62Vrx48jaZ6J_Jo,364
-tests/numpy_sample/test_collate.py,sha256=oWVNRR0LyF0iqaxjsq6hPMmkSGIyvGVlAM6q2rqvXqQ,840
+tests/numpy_sample/test_collate.py,sha256=ngbJ8vIewnAvkXx-PpfuSMVNM82_SYaZPLhJkZZw7s0,867
+tests/numpy_sample/test_datetime_features.py,sha256=o4t3KeKFdGrOBQ77rNFcDuDMQSD23ileCS5T5AP3wG4,1769
 tests/numpy_sample/test_gsp.py,sha256=FLlq4SlJ-9cSRAepf4_ksA6PsUVKegnKEAc5pUojCJ0,1458
 tests/numpy_sample/test_nwp.py,sha256=yf4u7mAU0E3FQ4xAH6YjuHuHBzzFoXjHSFNkOVJUdSM,1455
 tests/numpy_sample/test_satellite.py,sha256=cCqtn5See-uSNfh89COGTUQNuFm6sIZ8QmBVHsuUeRI,1189
@@ -61,11 +63,11 @@ tests/select/test_location.py,sha256=_WZk2FPYeJ-nIfCJS6Sp_yaVEEo7m31DmMFoZzgyCts
 tests/select/test_select_spatial_slice.py,sha256=7EX9b6g-pMdACQx3yefjs5do2s-Rho2UmKevV4oglsU,5147
 tests/select/test_select_time_slice.py,sha256=K1EJR5TwZa9dJf_YTEHxGtvs398iy1xS2lr1BgJZkoo,9603
 tests/torch_datasets/conftest.py,sha256=eRCzHE7cxS4AoskExkCGFDBeqItktAYNAdkfpMoFCeE,629
-tests/torch_datasets/test_process_and_combine.py,sha256=mbjQdqzLhox-U2sc1Ec68xLz95b3XOyPa7WchgxUM88,4256
-tests/torch_datasets/test_pvnet_uk_regional.py,sha256=dLY861PMyQ_buTksP8d0UXzfKsZ_CFNgceSYVGXRIRs,2134
-tests/torch_datasets/test_site.py,sha256=TSmdpD0OFid6e8eRgeB5FGxglUN1L2t4cIuMTfjA7jE,4539
-ocf_data_sampler-0.0.44.dist-info/LICENSE,sha256=F-Q3UFCR-BECSocV55BFDpn4YKxve9PKrm-lTt6o_Tg,1073
-ocf_data_sampler-0.0.44.dist-info/METADATA,sha256=DXdwieKe_o3sJ5BSZ8-zN8xawG3_1YDLQRHqgWV89xM,11431
-ocf_data_sampler-0.0.44.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-ocf_data_sampler-0.0.44.dist-info/top_level.txt,sha256=Faob6N6cFdPc5eUpCTYcXgCaNhi4XLLteUL5W5ayYmg,31
-ocf_data_sampler-0.0.44.dist-info/RECORD,,
+tests/torch_datasets/test_merge_and_fill_utils.py,sha256=ueA0A7gZaWEgNdsU8p3CnKuvSnlleTUjEhSw2HUUROM,1229
+tests/torch_datasets/test_pvnet_uk_regional.py,sha256=FCiFueeFqrsXe7gWguSjBz5ZeUrvyhGbGw81gaVvkHM,5087
+tests/torch_datasets/test_site.py,sha256=0tnjgx6z4VlzjoF_V2p3Y2t2Z1d0o_07Vwb-FH_c3tU,4640
+ocf_data_sampler-0.0.46.dist-info/LICENSE,sha256=F-Q3UFCR-BECSocV55BFDpn4YKxve9PKrm-lTt6o_Tg,1073
+ocf_data_sampler-0.0.46.dist-info/METADATA,sha256=S8ScJ8z3O0O5qhgGZmdI0Ugan2Yz4dH0nGj9R8N1sgs,11788
+ocf_data_sampler-0.0.46.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+ocf_data_sampler-0.0.46.dist-info/top_level.txt,sha256=Faob6N6cFdPc5eUpCTYcXgCaNhi4XLLteUL5W5ayYmg,31
+ocf_data_sampler-0.0.46.dist-info/RECORD,,

tests/config/test_config.py

@@ -2,7 +2,7 @@ import tempfile
 
 import pytest
 from pydantic import ValidationError
-
+from pathlib import Path
 from ocf_data_sampler.config import (
     load_yaml_configuration,
     Configuration,
@@ -21,39 +21,37 @@ def test_load_yaml_configuration(test_config_filename):
     Test that yaml loading works for 'test_config.yaml'
     and fails for an empty .yaml file
     """
-
-    # check we get an error if loading a file with no config
-    with tempfile.NamedTemporaryFile(suffix=".yaml") as fp:
-        filename = fp.name
-
-        # check that temp file can't be loaded
+    # Create temporary directory instead of file
+    with tempfile.TemporaryDirectory() as temp_dir:
+        # Create path for empty file
+        empty_file = Path(temp_dir) / "empty.yaml"
+        
+        # Create an empty file
+        empty_file.touch()
+        
+        # Test loading empty file
         with pytest.raises(TypeError):
-            _ = load_yaml_configuration(filename)
-
-    # test can load test_config.yaml
-    config = load_yaml_configuration(test_config_filename)
-
-    assert isinstance(config, Configuration)
-
+            _ = load_yaml_configuration(str(empty_file))
 
 def test_yaml_save(test_config_filename):
     """
     Check configuration can be saved to a .yaml file
     """
-
     test_config = load_yaml_configuration(test_config_filename)
-
-    with tempfile.NamedTemporaryFile(suffix=".yaml") as fp:
-        filename = fp.name
-
-        # save default config to file
-        save_yaml_configuration(test_config, filename)
-
-        # check the file can be loaded back
-        tmp_config = load_yaml_configuration(filename)
-
-        # check loaded configuration is the same as the one passed to save
-        assert test_config == tmp_config
+    
+    with tempfile.TemporaryDirectory() as temp_dir:
+        # Create path for config file
+        config_path = Path(temp_dir) / "test_config.yaml"
+        
+        # Save configuration
+        saved_path = save_yaml_configuration(test_config, config_path)
+        
+        # Verify file exists
+        assert saved_path.exists()
+        
+        # Test loading saved configuration
+        loaded_config = load_yaml_configuration(str(saved_path))
+        assert loaded_config == test_config
 
 
 def test_extra_field_error():

tests/conftest.py

@@ -1,10 +1,10 @@
 import os
-
 import numpy as np
 import pandas as pd
 import pytest
 import xarray as xr
 import tempfile
+from typing import Generator
 
 from ocf_data_sampler.config.model import Site
 from ocf_data_sampler.config import load_yaml_configuration, save_yaml_configuration
@@ -201,7 +201,7 @@ def ds_uk_gsp():
 
 
 @pytest.fixture(scope="session")
-def data_sites() -> Site:
+def data_sites() -> Generator[Site, None, None]:
     """
     Make fake data for sites
     Returns: filename for netcdf file, and csv metadata

tests/numpy_sample/test_collate.py

@@ -1,6 +1,6 @@
 from ocf_data_sampler.numpy_sample import GSPSampleKey, SatelliteSampleKey
 from ocf_data_sampler.numpy_sample.collate import stack_np_samples_into_batch
-from ocf_data_sampler.torch_datasets import PVNetUKRegionalDataset
+from ocf_data_sampler.torch_datasets.datasets.pvnet_uk_regional import PVNetUKRegionalDataset
 
 
 def test_pvnet(pvnet_config_filename):

tests/numpy_sample/test_datetime_features.py

@@ -0,0 +1,47 @@
+import numpy as np
+import pandas as pd
+import pytest
+
+from ocf_data_sampler.numpy_sample.datetime_features import make_datetime_numpy_dict
+
+
+def test_calculate_azimuth_and_elevation():
+
+    # Pick the day of the summer solstice
+    datetimes = pd.to_datetime(["2024-06-20 12:00", "2024-06-20 12:30", "2024-06-20 13:00"])
+
+    # Calculate sun angles
+    datetime_features = make_datetime_numpy_dict(datetimes)
+
+    assert len(datetime_features) == 4
+
+    assert len(datetime_features["wind_date_sin"]) == len(datetimes)
+    assert (datetime_features["wind_date_cos"] != datetime_features["wind_date_sin"]).all()
+
+    # assert all values are between -1 and 1
+    assert all(np.abs(datetime_features["wind_date_sin"]) <= 1)
+    assert all(np.abs(datetime_features["wind_date_cos"]) <= 1)
+    assert all(np.abs(datetime_features["wind_time_sin"]) <= 1)
+    assert all(np.abs(datetime_features["wind_time_cos"]) <= 1)
+
+
+def test_make_datetime_numpy_batch_custom_key_prefix():
+    # Test function correctly applies custom prefix to dict keys
+    datetimes = pd.to_datetime(["2024-06-20 12:00", "2024-06-20 12:30", "2024-06-20 13:00"])
+    key_prefix = "solar"
+
+    datetime_features = make_datetime_numpy_dict(datetimes, key_prefix=key_prefix)
+
+    # Assert dict contains expected quantity of keys and verify starting with custom prefix
+    assert len(datetime_features) == 4
+    assert all(key.startswith(key_prefix) for key in datetime_features.keys())
+
+
+def test_make_datetime_numpy_batch_empty_input():
+    # Verification that function raises error for empty input
+    datetimes = pd.DatetimeIndex([])
+
+    with pytest.raises(
+        ValueError, match="Input datetimes is empty for 'make_datetime_numpy_dict' function"
+    ):
+        make_datetime_numpy_dict(datetimes)

tests/torch_datasets/test_merge_and_fill_utils.py

@@ -0,0 +1,42 @@
+import numpy as np
+
+from ocf_data_sampler.torch_datasets.utils.merge_and_fill_utils import (
+    merge_dicts,
+    fill_nans_in_arrays,
+)
+
+def test_merge_dicts():
+    """Test merge_dicts function"""
+    dict1 = {"a": 1, "b": 2}
+    dict2 = {"c": 3, "d": 4}
+    dict3 = {"e": 5}
+    
+    result = merge_dicts([dict1, dict2, dict3])
+    assert result == {"a": 1, "b": 2, "c": 3, "d": 4, "e": 5}
+    
+    # Test key overwriting
+    dict4 = {"a": 10, "f": 6}
+    result = merge_dicts([dict1, dict4])
+    assert result["a"] == 10
+
+
+def test_fill_nans_in_arrays():
+    """Test the fill_nans_in_arrays function"""
+    array_with_nans = np.array([1.0, np.nan, 3.0, np.nan])
+    nested_dict = {
+        "array1": array_with_nans,
+        "nested": {
+            "array2": np.array([np.nan, 2.0, np.nan, 4.0])
+        },
+        "string_key": "not_an_array"
+    }
+    
+    result = fill_nans_in_arrays(nested_dict)
+    
+    assert not np.isnan(result["array1"]).any()
+    assert np.array_equal(result["array1"], np.array([1.0, 0.0, 3.0, 0.0]))
+    assert not np.isnan(result["nested"]["array2"]).any()
+    assert np.array_equal(result["nested"]["array2"], np.array([0.0, 2.0, 0.0, 4.0]))
+    assert result["string_key"] == "not_an_array"
+
+

tests/torch_datasets/test_pvnet_uk_regional.py

@@ -1,11 +1,88 @@
-import pytest
+import numpy as np
+import pandas as pd
+import xarray as xr
+import dask.array as da
 import tempfile
 
-from ocf_data_sampler.torch_datasets import PVNetUKRegionalDataset
-from ocf_data_sampler.config import load_yaml_configuration, save_yaml_configuration
+from ocf_data_sampler.torch_datasets.datasets.pvnet_uk_regional import PVNetUKRegionalDataset
+from ocf_data_sampler.config.save import save_yaml_configuration
+from ocf_data_sampler.config.load import load_yaml_configuration
 from ocf_data_sampler.numpy_sample import NWPSampleKey, GSPSampleKey, SatelliteSampleKey
+from ocf_data_sampler.torch_datasets.datasets.pvnet_uk_regional import process_and_combine_datasets, compute
+from ocf_data_sampler.select.location import Location
 
+def test_process_and_combine_datasets(pvnet_config_filename):
 
+    # Load in config for function and define location 
+    config = load_yaml_configuration(pvnet_config_filename)
+    t0 = pd.Timestamp("2024-01-01 00:00")
+    location = Location(coordinate_system="osgb", x=1234, y=5678, id=1)
+
+    nwp_data = xr.DataArray(
+        np.random.rand(4, 2, 2, 2),
+        dims=["time_utc", "channel", "y", "x"],
+        coords={
+            "time_utc": pd.date_range("2024-01-01 00:00", periods=4, freq="h"),
+            "channel": ["t2m", "dswrf"],
+            "step": ("time_utc", pd.timedelta_range(start='0h', periods=4, freq='h')),
+            "init_time_utc": pd.Timestamp("2024-01-01 00:00")
+        }
+    )
+
+    sat_data = xr.DataArray(
+        np.random.rand(7, 1, 2, 2),
+        dims=["time_utc", "channel", "y", "x"],
+        coords={
+            "time_utc": pd.date_range("2024-01-01 00:00", periods=7, freq="5min"),
+            "channel": ["HRV"],
+            "x_geostationary": (["y", "x"], np.array([[1, 2], [1, 2]])),
+            "y_geostationary": (["y", "x"], np.array([[1, 1], [2, 2]]))
+        }
+    )
+
+    # Combine as dict
+    dataset_dict = {
+        "nwp": {"ukv": nwp_data},
+        "sat": sat_data
+    }
+
+    # Call relevant function
+    result = process_and_combine_datasets(dataset_dict, config, t0, location)
+
+    # Assert result is dict - check and validate
+    assert isinstance(result, dict)
+    assert NWPSampleKey.nwp in result
+    assert result[SatelliteSampleKey.satellite_actual].shape == (7, 1, 2, 2)
+    assert result[NWPSampleKey.nwp]["ukv"][NWPSampleKey.nwp].shape == (4, 1, 2, 2)
+
+def test_compute():
+    """Test compute function with dask array"""
+    da_dask = xr.DataArray(da.random.random((5, 5)))
+
+    # Create a nested dictionary with dask array
+    nested_dict = {
+        "array1": da_dask,
+        "nested": {
+            "array2": da_dask
+        }
+    }
+
+    # Ensure initial data is lazy - i.e. not yet computed
+    assert not isinstance(nested_dict["array1"].data, np.ndarray)
+    assert not isinstance(nested_dict["nested"]["array2"].data, np.ndarray)
+
+    # Call the compute function
+    result = compute(nested_dict)
+
+    # Assert that the result is an xarray DataArray and no longer lazy
+    assert isinstance(result["array1"], xr.DataArray)
+    assert isinstance(result["nested"]["array2"], xr.DataArray)
+    assert isinstance(result["array1"].data, np.ndarray)
+    assert isinstance(result["nested"]["array2"].data, np.ndarray)
+
+    # Ensure there no NaN values in computed data
+    assert not np.isnan(result["array1"].data).any()
+    assert not np.isnan(result["nested"]["array2"].data).any()
 
 def test_pvnet(pvnet_config_filename):
 

tests/torch_datasets/test_site.py

@@ -1,8 +1,6 @@
 import pandas as pd
-
-from ocf_data_sampler.torch_datasets import SitesDataset
-from ocf_data_sampler.torch_datasets.site import convert_from_dataset_to_dict_datasets
 import numpy as np
+from ocf_data_sampler.torch_datasets.datasets.site import SitesDataset, convert_from_dataset_to_dict_datasets
 from xarray import Dataset, DataArray
 
 
@@ -22,7 +20,9 @@ def test_site(site_config_filename):
     # Expected dimensions and data variables
     expected_dims = {'satellite__x_geostationary', 'site__time_utc', 'nwp-ukv__target_time_utc',
                      'nwp-ukv__x_osgb', 'satellite__channel', 'satellite__y_geostationary',
-                     'satellite__time_utc', 'nwp-ukv__channel', 'nwp-ukv__y_osgb'}
+                     'satellite__time_utc', 'nwp-ukv__channel', 'nwp-ukv__y_osgb', 'site_solar_azimuth',
+                     'site_solar_elevation', 'site_date_cos', 'site_time_cos', 'site_time_sin', 'site_date_sin'}
+
     expected_data_vars = {"nwp-ukv", "satellite", "site"}
 
     # Check dimensions

ocf_data_sampler/torch_datasets/process_and_combine.py

@@ -1,131 +0,0 @@
-import numpy as np
-import pandas as pd
-import xarray as xr
-from typing import Optional
-
-from ocf_data_sampler.config import Configuration
-from ocf_data_sampler.constants import NWP_MEANS, NWP_STDS,RSS_MEAN,RSS_STD
-from ocf_data_sampler.numpy_sample import (
-    convert_nwp_to_numpy_sample,
-    convert_satellite_to_numpy_sample,
-    convert_gsp_to_numpy_sample,
-    make_sun_position_numpy_sample,
-)
-from ocf_data_sampler.numpy_sample.gsp import GSPSampleKey
-from ocf_data_sampler.numpy_sample.nwp import NWPSampleKey
-from ocf_data_sampler.select.geospatial import osgb_to_lon_lat
-from ocf_data_sampler.select.location import Location
-from ocf_data_sampler.utils import minutes
-
-
-def process_and_combine_datasets(
-    dataset_dict: dict,
-    config: Configuration,
-    t0: Optional[pd.Timestamp] = None,
-    location: Optional[Location] = None,
-    target_key: str = 'gsp'
-) -> dict:
-
-    """Normalise 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].provider
-            da_nwp = (da_nwp - NWP_MEANS[provider]) / NWP_STDS[provider]
-            # Convert to NumpySample
-            nwp_numpy_modalities[nwp_key] = convert_nwp_to_numpy_sample(da_nwp)
-
-        # Combine the NWPs into NumpySample
-        numpy_modalities.append({NWPSampleKey.nwp: nwp_numpy_modalities})
-
-
-    if "sat" in dataset_dict:
-        # Standardise
-        da_sat = dataset_dict["sat"]
-        da_sat = (da_sat - RSS_MEAN) / RSS_STD
-
-        # Convert to NumpySample
-        numpy_modalities.append(convert_satellite_to_numpy_sample(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_sample(
-                da_gsp, 
-                t0_idx=-gsp_config.interval_start_minutes / gsp_config.time_resolution_minutes
-            )
-        )
-
-        # Add coordinate data
-        # TODO: Do we need all of these?
-        numpy_modalities.append(
-            {
-                GSPSampleKey.gsp_id: location.id,
-                GSPSampleKey.x_osgb: location.x,
-                GSPSampleKey.y_osgb: location.y,
-            }
-        )
-
-    if target_key == 'gsp':
-        # Make sun coords NumpySample
-        datetimes = pd.date_range(
-            t0+minutes(gsp_config.interval_start_minutes),
-            t0+minutes(gsp_config.interval_end_minutes),
-            freq=minutes(gsp_config.time_resolution_minutes),
-        )
-
-        lon, lat = osgb_to_lon_lat(location.x, location.y)
-
-    numpy_modalities.append(
-        make_sun_position_numpy_sample(datetimes, lon, lat, key_prefix=target_key)
-    )
-
-    # 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(sample: dict) -> dict:
-    """Fills all NaN values in each np.ndarray in the sample dictionary with zeros.
-
-    Operation is performed in-place on the sample.
-    """
-    for k, v in sample.items():
-        if isinstance(v, np.ndarray) and np.issubdtype(v.dtype, np.number):
-            if np.isnan(v).any():
-                sample[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 sample
-
-
-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

tests/torch_datasets/test_process_and_combine.py

@@ -1,126 +0,0 @@
-import numpy as np
-import pandas as pd
-import xarray as xr
-import dask.array as da
-
-from ocf_data_sampler.config import load_yaml_configuration
-from ocf_data_sampler.select.location import Location
-from ocf_data_sampler.numpy_sample import NWPSampleKey, GSPSampleKey, SatelliteSampleKey
-from ocf_data_sampler.torch_datasets import PVNetUKRegionalDataset
-
-from ocf_data_sampler.torch_datasets.process_and_combine import (
-    process_and_combine_datasets,
-    merge_dicts,
-    fill_nans_in_arrays,
-    compute,
-)
-
-
-def test_process_and_combine_datasets(pvnet_config_filename):
-
-    # Load in config for function and define location 
-    config = load_yaml_configuration(pvnet_config_filename)
-    t0 = pd.Timestamp("2024-01-01 00:00")
-    location = Location(coordinate_system="osgb", x=1234, y=5678, id=1)
-
-    nwp_data = xr.DataArray(
-        np.random.rand(4, 2, 2, 2),
-        dims=["time_utc", "channel", "y", "x"],
-        coords={
-            "time_utc": pd.date_range("2024-01-01 00:00", periods=4, freq="h"),
-            "channel": ["t2m", "dswrf"],
-            "step": ("time_utc", pd.timedelta_range(start='0h', periods=4, freq='h')),
-            "init_time_utc": pd.Timestamp("2024-01-01 00:00")
-        }
-    )
-
-    sat_data = xr.DataArray(
-        np.random.rand(7, 1, 2, 2),
-        dims=["time_utc", "channel", "y", "x"],
-        coords={
-            "time_utc": pd.date_range("2024-01-01 00:00", periods=7, freq="5min"),
-            "channel": ["HRV"],
-            "x_geostationary": (["y", "x"], np.array([[1, 2], [1, 2]])),
-            "y_geostationary": (["y", "x"], np.array([[1, 1], [2, 2]]))
-        }
-    )
-
-    # Combine as dict
-    dataset_dict = {
-        "nwp": {"ukv": nwp_data},
-        "sat": sat_data
-    }
-
-    # Call relevant function
-    result = process_and_combine_datasets(dataset_dict, config, t0, location)
-
-    # Assert result is dict - check and validate
-    assert isinstance(result, dict)
-    assert NWPSampleKey.nwp in result
-    assert result[SatelliteSampleKey.satellite_actual].shape == (7, 1, 2, 2)
-    assert result[NWPSampleKey.nwp]["ukv"][NWPSampleKey.nwp].shape == (4, 1, 2, 2)
-
-
-def test_merge_dicts():
-    """Test merge_dicts function"""
-    dict1 = {"a": 1, "b": 2}
-    dict2 = {"c": 3, "d": 4}
-    dict3 = {"e": 5}
-    
-    result = merge_dicts([dict1, dict2, dict3])
-    assert result == {"a": 1, "b": 2, "c": 3, "d": 4, "e": 5}
-    
-    # Test key overwriting
-    dict4 = {"a": 10, "f": 6}
-    result = merge_dicts([dict1, dict4])
-    assert result["a"] == 10
-
-
-def test_fill_nans_in_arrays():
-    """Test the fill_nans_in_arrays function"""
-    array_with_nans = np.array([1.0, np.nan, 3.0, np.nan])
-    nested_dict = {
-        "array1": array_with_nans,
-        "nested": {
-            "array2": np.array([np.nan, 2.0, np.nan, 4.0])
-        },
-        "string_key": "not_an_array"
-    }
-    
-    result = fill_nans_in_arrays(nested_dict)
-    
-    assert not np.isnan(result["array1"]).any()
-    assert np.array_equal(result["array1"], np.array([1.0, 0.0, 3.0, 0.0]))
-    assert not np.isnan(result["nested"]["array2"]).any()
-    assert np.array_equal(result["nested"]["array2"], np.array([0.0, 2.0, 0.0, 4.0]))
-    assert result["string_key"] == "not_an_array"
-
-
-def test_compute():
-    """Test compute function with dask array"""
-    da_dask = xr.DataArray(da.random.random((5, 5)))
-
-    # Create a nested dictionary with dask array
-    nested_dict = {
-        "array1": da_dask,
-        "nested": {
-            "array2": da_dask
-        }
-    }
-
-    # Ensure initial data is lazy - i.e. not yet computed
-    assert not isinstance(nested_dict["array1"].data, np.ndarray)
-    assert not isinstance(nested_dict["nested"]["array2"].data, np.ndarray)
-
-    # Call the compute function
-    result = compute(nested_dict)
-
-    # Assert that the result is an xarray DataArray and no longer lazy
-    assert isinstance(result["array1"], xr.DataArray)
-    assert isinstance(result["nested"]["array2"], xr.DataArray)
-    assert isinstance(result["array1"].data, np.ndarray)
-    assert isinstance(result["nested"]["array2"].data, np.ndarray)
-
-    # Ensure there no NaN values in computed data
-    assert not np.isnan(result["array1"].data).any()
-    assert not np.isnan(result["nested"]["array2"].data).any()