ocf-data-sampler 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

ocf_data_sampler/numpy_sample/collate.py

@@ -1,75 +1,64 @@
-from ocf_data_sampler.numpy_sample import NWPSampleKey
-
 import numpy as np
-import logging
-from typing import Union
-
-logger = logging.getLogger(__name__)
-
-
-
-def stack_np_samples_into_batch(dict_list):
-#     """
-#     Stacks Numpy samples into a batch
 
-#     Args:
-#         dict_list: A list of dict-like Numpy samples to stack
 
-#     Returns:
-#         The stacked NumpySample object, aka a batch
-#     """
+def stack_np_samples_into_batch(dict_list: list[dict]) -> dict:
+    """Stacks list of dict samples into a dict where all samples are joined along a new axis
 
-    if not dict_list:
-        raise ValueError("Input is empty")
+    Args:
+        dict_list: A list of dict-like samples to stack
 
-    # Extract keys from first dict - structure
-    sample = {}
-    sample_keys = list(dict_list[0].keys())
+    Returns:
+        Dict of the samples stacked with new batch dimension on axis 0
+    """
+    
+    batch = {}
+
+    keys = list(dict_list[0].keys())
+
+    for key in keys:
+        # NWP is nested so treat separately
+        if key == "nwp":
+            batch["nwp"] = {}
+
+            # Unpack NWP provider keys
+            nwp_providers = list(dict_list[0]["nwp"].keys())
+
+            for nwp_provider in nwp_providers:
+                # Keys can be different for different NWPs
+                nwp_keys = list(dict_list[0]["nwp"][nwp_provider].keys())
+                
+                # Create dict to store NWP batch for this provider
+                nwp_provider_batch = {}
+                
+                for nwp_key in nwp_keys:
+                    # Stack values under each NWP key for this provider
+                    nwp_provider_batch[nwp_key] = stack_data_list(
+                        [d["nwp"][nwp_provider][nwp_key] for d in dict_list],
+                        nwp_key,
+                    )
+
+                batch["nwp"][nwp_provider] = nwp_provider_batch
 
-    # Process - handle NWP separately due to nested structure
-    for sample_key in sample_keys:
-        if sample_key == "nwp":
-            sample["nwp"] = process_nwp_data(dict_list)
         else:
-            # Stack arrays for the given key across all dicts
-            sample[sample_key] = stack_data_list([d[sample_key] for d in dict_list], sample_key)
-    return sample
+            batch[key] = stack_data_list([d[key] for d in dict_list], key)
 
+    return batch
 
-def process_nwp_data(dict_list):
-    """Stacks data for NWP, handling nested structure"""
-    
-    nwp_sample = {}
-    nwp_sources = dict_list[0]["nwp"].keys()
 
-    # Stack data for each NWP source independently
-    for nwp_source in nwp_sources:
-        nested_keys = dict_list[0]["nwp"][nwp_source].keys()
-        nwp_sample[nwp_source] = {
-            key: stack_data_list([d["nwp"][nwp_source][key] for d in dict_list], key)
-            for key in nested_keys
-        }
-    return nwp_sample
+def _key_is_constant(key: str):
+    return key.endswith("t0_idx") or key.endswith("channel_names")
 
-def _key_is_constant(sample_key):
-    return sample_key.endswith("t0_idx") or sample_key == NWPSampleKey.channel_names
 
-
-def stack_data_list(data_list: list,sample_key: Union[str, NWPSampleKey],):
-    """How to combine data entries for each key
+def stack_data_list(data_list: list, key: str):
+    """Stack a sequence of data elements along a new axis
 
      Args:
-        data_list: List of data entries to combine
-        sample_key: Key identifying the data type
+        data_list: List of data elements to combine
+        key: string identifying the data type
     """
-    if _key_is_constant(sample_key):
+    if _key_is_constant(key):
         # These are always the same for all examples.
         return data_list[0]
-    try:
+    else:
         return np.stack(data_list)
-    except Exception as e:
-        logger.debug(f"Could not stack the following shapes together, ({sample_key})")
-        shapes = [example.shape for example in data_list]
-        logger.debug(shapes)
-        logger.error(e)
-        raise e
+

ocf_data_sampler/sample/uk_regional.py

@@ -65,7 +65,9 @@ class UKRegionalSample(SampleBase):
             raise ValueError(f"Only .pt format is supported: {path.suffix}")
         
         instance = cls()
-        instance._data = torch.load(path)
+        # TODO: We should move away from using torch.load(..., weights_only=False)
+        # This is not recommended
+        instance._data = torch.load(path, weights_only=False)
         logger.debug(f"Successfully loaded UKRegionalSample from {path}")
         return instance
 

ocf_data_sampler/select/fill_time_periods.py

@@ -4,7 +4,7 @@ import pandas as pd
 import numpy as np
 
 
-def fill_time_periods(time_periods: pd.DataFrame, freq: pd.Timedelta):
+def fill_time_periods(time_periods: pd.DataFrame, freq: pd.Timedelta) -> pd.DatetimeIndex:
     start_dts = pd.to_datetime(time_periods["start_dt"].values).ceil(freq)
     end_dts = pd.to_datetime(time_periods["end_dt"].values)
     date_ranges = [pd.date_range(start_dt, end_dt, freq=freq) for start_dt, end_dt in zip(start_dts, end_dts)]

ocf_data_sampler/select/time_slice_for_dataset.py

@@ -1,5 +1,6 @@
 """ Slice datasets by time"""
 import pandas as pd
+import xarray as xr
 
 from ocf_data_sampler.config import Configuration
 from ocf_data_sampler.select.dropout import draw_dropout_time, apply_dropout_time
@@ -64,16 +65,8 @@ def slice_datasets_by_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(gsp_config.time_resolution_minutes),
-            interval_end=minutes(gsp_config.interval_end_minutes),
-        )
-
-        sliced_datasets_dict["gsp"] = select_time_slice(
+        
+        da_gsp_past = select_time_slice(
             datasets_dict["gsp"],
             t0,
             sample_period_duration=minutes(gsp_config.time_resolution_minutes),
@@ -81,17 +74,27 @@ def slice_datasets_by_time(
             interval_end=minutes(0),
         )
 
-        # Dropout on the GSP, but not the future GSP
+        # Dropout on the past 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"], 
+        da_gsp_past = apply_dropout_time(
+            da_gsp_past, 
             gsp_dropout_time
         )
+        
+        da_gsp_future = select_time_slice(
+            datasets_dict["gsp"],
+            t0,
+            sample_period_duration=minutes(gsp_config.time_resolution_minutes),
+            interval_start=minutes(gsp_config.time_resolution_minutes),
+            interval_end=minutes(gsp_config.interval_end_minutes),
+        )
+        
+        sliced_datasets_dict["gsp"] = xr.concat([da_gsp_past, da_gsp_future], dim="time_utc")
 
     if "site" in datasets_dict:
         site_config = config.input_data.site

ocf_data_sampler/torch_datasets/datasets/__init__.py

@@ -1,11 +1,6 @@
-from .pvnet_uk_regional import PVNetUKRegionalDataset
+from .pvnet_uk import PVNetUKRegionalDataset, PVNetUKConcurrentDataset
 
 from .site import (
     convert_netcdf_to_numpy_sample,
     SitesDataset
-)
-
-__all__ = [
-    'convert_netcdf_to_numpy_sample',
-    'SitesDataset'
-]
+)

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

@@ -1,15 +1,20 @@
-"""Torch dataset for PVNet"""
+"""Torch dataset for UK PVNet"""
+
+import pkg_resources
 
 import numpy as np
 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.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.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,
@@ -17,13 +22,16 @@ from ocf_data_sampler.numpy_sample import (
     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.numpy_sample.collate import stack_np_samples_into_batch
+from ocf_data_sampler.select.geospatial import osgb_to_lon_lat
+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.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)
 
@@ -65,9 +73,10 @@ def process_and_combine_datasets(
     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 = dataset_dict["gsp"]
         da_gsp = da_gsp / da_gsp.effective_capacity_mwp
-
+        
+        # Convert to NumpyBatch
         numpy_modalities.append(
             convert_gsp_to_numpy_sample(
                 da_gsp, 
@@ -105,6 +114,7 @@ def process_and_combine_datasets(
 
     return combined_sample
 
+
 def compute(xarray_dict: dict) -> dict:
     """Eagerly load a nested dictionary of xarray DataArrays"""
     for k, v in xarray_dict.items():
@@ -114,10 +124,8 @@ def compute(xarray_dict: dict) -> dict:
             xarray_dict[k] = v.compute(scheduler="single-threaded")
     return xarray_dict
 
-def find_valid_t0_times(
-    datasets_dict: dict,
-    config: Configuration,
-):
+
+def find_valid_t0_times(datasets_dict: dict, config: Configuration) -> pd.DatetimeIndex:
     """Find the t0 times where all of the requested input data is available
 
     Args:
@@ -167,7 +175,7 @@ class PVNetUKRegionalDataset(Dataset):
         self, 
         config_filename: str, 
         start_time: str | None = None,
-        end_time: str| None = None,
+        end_time: str | None = None,
         gsp_ids: list[int] | None = None,
     ):
         """A torch Dataset for creating PVNet UK GSP samples
@@ -253,7 +261,7 @@ class PVNetUKRegionalDataset(Dataset):
     def get_sample(self, t0: pd.Timestamp, gsp_id: int) -> dict:
         """Generate a sample for the given coordinates. 
         
-        Useful for users to generate samples by GSP ID.
+        Useful for users to generate specific samples.
         
         Args:
             t0: init-time for sample
@@ -265,4 +273,94 @@ class PVNetUKRegionalDataset(Dataset):
 
         location = self.location_lookup[gsp_id]
         
-        return self._get_sample(t0, location)
+        return self._get_sample(t0, location)
+    
+    
+class PVNetUKConcurrentDataset(Dataset):
+    def __init__(
+        self, 
+        config_filename: str, 
+        start_time: str | None = None,
+        end_time: str | None = None,
+        gsp_ids: list[int] | None = None,
+    ):
+        """A torch Dataset for creating concurrent samples of PVNet UK regional data
+        
+        Each concurrent sample includes the data from all GSPs for a single t0 time
+
+        Args:
+            config_filename: Path to the configuration file
+            start_time: Limit the init-times to be after this
+            end_time: Limit the init-times to be before this
+            gsp_ids: List of all GSP IDs included in each sample. Defaults to all
+        """
+        
+        config = load_yaml_configuration(config_filename)
+        
+        datasets_dict = get_dataset_dict(config)
+        
+        # Get t0 times where all input data is available
+        valid_t0_times = find_valid_t0_times(datasets_dict, config)
+
+        # Filter t0 times to given range
+        if start_time is not None:
+            valid_t0_times = valid_t0_times[valid_t0_times>=pd.Timestamp(start_time)]
+            
+        if end_time is not None:
+            valid_t0_times = valid_t0_times[valid_t0_times<=pd.Timestamp(end_time)]
+
+        # Construct list of locations to sample from
+        locations = get_gsp_locations(gsp_ids)
+                        
+        # Assign coords and indices to self
+        self.valid_t0_times = valid_t0_times
+        self.locations = locations
+
+        # Assign config and input data to self
+        self.datasets_dict = datasets_dict
+        self.config = config
+        
+        
+    def __len__(self):
+        return len(self.valid_t0_times)
+    
+    
+    def _get_sample(self, t0: pd.Timestamp) -> dict:
+        """Generate a concurrent PVNet sample for given init-time
+        
+        Args:
+            t0: init-time for sample
+        """
+        # Slice by time then load to avoid loading the data multiple times from disk
+        sample_dict = slice_datasets_by_time(self.datasets_dict, t0, self.config)
+        sample_dict = compute(sample_dict)
+        
+        gsp_samples = []
+        
+        # Prepare sample for each GSP
+        for location in self.locations:
+            gsp_sample_dict = slice_datasets_by_space(sample_dict, location, self.config)
+            gsp_numpy_sample = process_and_combine_datasets(
+                gsp_sample_dict, self.config, t0, location
+            )
+            gsp_samples.append(gsp_numpy_sample)
+        
+        # Stack GSP samples
+        return stack_np_samples_into_batch(gsp_samples)
+    
+        
+    def __getitem__(self, idx):
+        return self._get_sample(self.valid_t0_times[idx])
+    
+
+    def get_sample(self, t0: pd.Timestamp) -> dict:
+        """Generate a sample for the given init-time. 
+        
+        Useful for users to generate specific samples.
+        
+        Args:
+            t0: init-time for sample
+        """
+        # Check data is availablle for init-time t0
+        assert t0 in self.valid_t0_times        
+        return self._get_sample(t0)

{ocf_data_sampler-0.1.0.dist-info → ocf_data_sampler-0.1.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: ocf_data_sampler
-Version: 0.1.0
+Version: 0.1.2
 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.1.0.dist-info → ocf_data_sampler-0.1.2.dist-info}/RECORD

@@ -19,7 +19,7 @@ ocf_data_sampler/load/nwp/providers/ecmwf.py,sha256=2iR1Iy542lo51rC6XFLV-3pbUE68
 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=nY5C6CcuxiWZ_jrXRzWtN7WyKXhJImSiVTIG6Rz4B_4,401
-ocf_data_sampler/numpy_sample/collate.py,sha256=y8QFUhskaAfOMP22aVkexwyGAwLHbNE-q1pOZ6txWKA,2227
+ocf_data_sampler/numpy_sample/collate.py,sha256=Onl_aKhsZ4pbFJsh70orjsHk523GHxrpRirH2vJq_GA,1911
 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
@@ -29,32 +29,32 @@ ocf_data_sampler/numpy_sample/sun_position.py,sha256=UklhucCxCT6GMlAhCWL6c4cfWrd
 ocf_data_sampler/sample/__init__.py,sha256=02CM7E5nKkGiYbVW-kvzjNd4RaqGuHCkDChtmDBDUoA,248
 ocf_data_sampler/sample/base.py,sha256=4U78tczCRsKMDwU4HkD20nyGyYjIBSZV5neF2mT--2M,1197
 ocf_data_sampler/sample/site.py,sha256=0BvDXs0kxTjUq7kWpeoITK_uN4uE0w1IvEFXZUoKOb0,2507
-ocf_data_sampler/sample/uk_regional.py,sha256=FPaFi6qaTsi1ag42pfVKDZhopt3cDjQsF4rVI8k2qWo,4244
+ocf_data_sampler/sample/uk_regional.py,sha256=D1A6nQB1PYCmxb3FzU9gqbNufQfx__wcprcDm50jCJw,4381
 ocf_data_sampler/select/__init__.py,sha256=E4AJulEbO2K-o0UlG1fgaEteuf_1ZFjHTvrotXSb4YU,332
 ocf_data_sampler/select/dropout.py,sha256=HCx5Wzk8Oh2Z9vV94Jy-ALJsHtGduwvMaQOleQXp5z0,1142
-ocf_data_sampler/select/fill_time_periods.py,sha256=iTtMjIPFYG5xtUYYedAFBLjTWWUa7t7WQ0-yksWf0-E,440
+ocf_data_sampler/select/fill_time_periods.py,sha256=h0XD1Ds_wUUoy-7bILxmN8AIbjlQ6YdXRKuCk_Is5jo,460
 ocf_data_sampler/select/find_contiguous_time_periods.py,sha256=q7IaNfX95A3z9XHqbhgtkZ4Js1gn5K9Qyp6DVLbsL-Q,11093
 ocf_data_sampler/select/geospatial.py,sha256=4xL-9y674jjoaXeqE52NHCHVfknciE4OEGsZtn9DvP4,4911
 ocf_data_sampler/select/location.py,sha256=26Y5ZjfFngShBwXieuWSoOA-RLaRzci4TTmcDk3Wg7U,2015
 ocf_data_sampler/select/select_spatial_slice.py,sha256=WNxwur9Q5oetvogATw8-hNejDuEwrXHzuZIovFDjNJA,11488
 ocf_data_sampler/select/select_time_slice.py,sha256=9M-yvDv9K77XfEys_OIR31_aVB56sNWk3BnCnkCgcPI,4725
 ocf_data_sampler/select/spatial_slice_for_dataset.py,sha256=3tRrMBXr7s4CnClbVSIq7hpls3H4Y3qYTDwswcxCCCE,1763
-ocf_data_sampler/select/time_slice_for_dataset.py,sha256=P7cAARfDzjttGDvpKt2zuA4WkLoTmSXy_lBpI8RiA6k,4249
-ocf_data_sampler/torch_datasets/datasets/__init__.py,sha256=JMfMQ6DxCWQiwm-Xwdy_b0gnGqZnORSR_SGrLM1QEe4,201
-ocf_data_sampler/torch_datasets/datasets/pvnet_uk_regional.py,sha256=xxeX4Js9LQpydehi3BS7k9psqkYGzgJuM17uTYux40M,8742
+ocf_data_sampler/select/time_slice_for_dataset.py,sha256=Z7pOiilSHScxmBKZNG18K5J-S4ifdXXAYGZoHRHD3AY,4324
+ocf_data_sampler/torch_datasets/datasets/__init__.py,sha256=jfJSFcR0eO1AqeH7S3KnGjsBqVZT5w3oyi784PUR6Q0,146
+ocf_data_sampler/torch_datasets/datasets/pvnet_uk.py,sha256=rodvVSR4sh8qZ2hLdI8qAc3lyxq5U7cVGfS4rRKCzbs,11944
 ocf_data_sampler/torch_datasets/datasets/site.py,sha256=5T8nkTMUHHFidZRuFOunYeKAqNuyZ8V7sikBoBOBwwA,16033
 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=DfrH0Pm552Tnl35eZn2UHCfOn2lHRiEQCcUCJIhycSU,8021
+tests/conftest.py,sha256=RlC7YYtBLipUzFS1tQxela1SgHCxSpReUKEJ4429PwQ,7689
 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_satellite.py,sha256=IQ8ISRZKCEoi8IsJoPpXZJTolD0mwjnl2E7762RM_PM,524
 tests/load/test_load_sites.py,sha256=T9lSEnGPI8FQISudVYHHNTHeplNS62Vrx48jaZ6J_Jo,364
-tests/numpy_sample/test_collate.py,sha256=ngbJ8vIewnAvkXx-PpfuSMVNM82_SYaZPLhJkZZw7s0,867
+tests/numpy_sample/test_collate.py,sha256=RqHCD5_LTRpe4r6kqC_2TKhmhM_IHYM0ZtFUvSjDqcM,654
 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
@@ -69,12 +69,11 @@ tests/select/test_select_time_slice.py,sha256=nYrdlmZlGEygJKiE26bADiluNPN1qt5kD4
 tests/test_sample/test_base.py,sha256=ljtB38MmscTGN6OvUgclBceNnfx6m7AN8iHYDml9XW4,2189
 tests/test_sample/test_site_sample.py,sha256=Gln-Or060cUWvA7Q7c1vsthgCttOAM2z9yBI9zUIrDw,6238
 tests/test_sample/test_uk_regional_sample.py,sha256=gkeQWC2wC757jKJz_QBmDMFQjn3R54q_tEo948yyxCY,4840
-tests/torch_datasets/conftest.py,sha256=eRCzHE7cxS4AoskExkCGFDBeqItktAYNAdkfpMoFCeE,629
-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=J1ZDE5V5MRlq7EuZ1zUu-aFRGTDJIiO-ZZzkOXvDdWA,6757
-ocf_data_sampler-0.1.0.dist-info/LICENSE,sha256=F-Q3UFCR-BECSocV55BFDpn4YKxve9PKrm-lTt6o_Tg,1073
-ocf_data_sampler-0.1.0.dist-info/METADATA,sha256=9hgn-WJMx51JmLEiZmd5oiEwjZtcutfZObTtnwxUT2k,12173
-ocf_data_sampler-0.1.0.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-ocf_data_sampler-0.1.0.dist-info/top_level.txt,sha256=Faob6N6cFdPc5eUpCTYcXgCaNhi4XLLteUL5W5ayYmg,31
-ocf_data_sampler-0.1.0.dist-info/RECORD,,
+tests/torch_datasets/test_merge_and_fill_utils.py,sha256=GtuQg82BM1eHQjT7Ik1x1zaVcuc7KJO4_NC9stXsd4s,1123
+tests/torch_datasets/test_pvnet_uk.py,sha256=OzT9ArdnWPa3iJKggxc2-7npkDqWmZyS5pzM4M08NZU,5566
+tests/torch_datasets/test_site.py,sha256=5MH5zkHFJXekwpnV6nHuSxt_sRNu9_mxiUjfWqmEhr0,6966
+ocf_data_sampler-0.1.2.dist-info/LICENSE,sha256=F-Q3UFCR-BECSocV55BFDpn4YKxve9PKrm-lTt6o_Tg,1073
+ocf_data_sampler-0.1.2.dist-info/METADATA,sha256=tWyfIpvmOufUWc7LquOXZ6g5Le_WhwihhZQBSZ0WKhA,12173
+ocf_data_sampler-0.1.2.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+ocf_data_sampler-0.1.2.dist-info/top_level.txt,sha256=Faob6N6cFdPc5eUpCTYcXgCaNhi4XLLteUL5W5ayYmg,31
+ocf_data_sampler-0.1.2.dist-info/RECORD,,

tests/conftest.py

@@ -1,14 +1,15 @@
+import pytest
+
 import os
 import numpy as np
 import pandas as pd
-import pytest
 import xarray as xr
-import tempfile
-from typing import Generator
+import dask.array
 
 from ocf_data_sampler.config.model import Site
 from ocf_data_sampler.config import load_yaml_configuration, save_yaml_configuration
 
+
 _top_test_directory = os.path.dirname(os.path.realpath(__file__))
 
 @pytest.fixture()
@@ -18,40 +19,27 @@ def test_config_filename():
 
 @pytest.fixture(scope="session")
 def config_filename():
-    return f"{os.path.dirname(os.path.abspath(__file__))}/test_data/configs/pvnet_test_config.yaml"
+    return f"{_top_test_directory}/test_data/configs/pvnet_test_config.yaml"
 
 
 @pytest.fixture(scope="session")
-def sat_zarr_path():
-
-    # Load dataset which only contains coordinates, but no data
-    ds = xr.open_zarr(
-        f"{os.path.dirname(os.path.abspath(__file__))}/test_data/non_hrv_shell.zarr.zip"
-    ).compute()
-
-    # Add time coord
-    ds = ds.assign_coords(time=pd.date_range("2023-01-01 00:00", "2023-01-02 23:55", freq="5min"))
-
-    # Add data to dataset
-    ds["data"] = xr.DataArray(
-        np.zeros([len(ds[c]) for c in ds.coords], dtype=np.float32),
-        coords=ds.coords,
-    )
-
-    # Transpose to variables, time, y, x (just in case)
-    ds = ds.transpose("variable", "time", "y_geostationary", "x_geostationary")
+def session_tmp_path(tmp_path_factory):
+    return tmp_path_factory.mktemp("data")
 
-    # add 100,000 to x_geostationary, this to make sure the fix index is within the satellite image
-    ds["x_geostationary"] = ds["x_geostationary"] - 200_000
 
-    # Add some NaNs
-    ds["data"].values[:, :, 0, 0] = np.nan
-
-    # make sure channel values are strings
-    ds["variable"] = ds["variable"].astype(str)
-
-    # add data attrs area
-    ds["data"].attrs["area"] = (
+@pytest.fixture(scope="session")
+def sat_zarr_path(session_tmp_path):
+
+    # Define coords for satellite-like dataset
+    variables = [
+        'IR_016', 'IR_039', 'IR_087', 'IR_097', 'IR_108', 'IR_120', 
+        'IR_134', 'VIS006', 'VIS008', 'WV_062', 'WV_073',
+    ]
+    x = np.linspace(start=15002, stop=-1824245, num=100)
+    y = np.linspace(start=4191563, stop=5304712, num=100)
+    times = pd.date_range("2023-01-01 00:00", "2023-01-01 23:55", freq="5min")
+    
+    area_string = (
         """msg_seviri_rss_3km:
         description: MSG SEVIRI Rapid Scanning Service area definition with 3 km resolution
         projection:
@@ -73,16 +61,31 @@ def sat_zarr_path():
             units: m
         """
     )
-
-    # Specifiy chunking
-    ds = ds.chunk({"time": 10, "variable": -1, "y_geostationary": -1, "x_geostationary": -1})
+    
+    # Create satellite-like data with some NaNs
+    data = dask.array.zeros(
+        shape=(len(variables), len(times), len(y), len(x)), 
+        chunks=(-1, 10, -1, -1),
+        dtype=np.float32
+    )
+    data [:, 10, :, :] = np.nan
+    
+    ds = xr.DataArray(
+        data=data,
+        coords=dict(
+            variable=variables,
+            time=times,
+            y_geostationary=y,
+            x_geostationary=x,
+        ),
+        attrs=dict(area=area_string),
+    ).to_dataset(name="data")
 
     # Save temporarily as a zarr
-    with tempfile.TemporaryDirectory() as tmpdir:
-        zarr_path = f"{tmpdir}/test_sat.zarr"
-        ds.to_zarr(zarr_path)
+    zarr_path = session_tmp_path / "test_sat.zarr"
+    ds.to_zarr(zarr_path)
 
-        yield zarr_path
+    yield zarr_path
 
 
 @pytest.fixture(scope="session")
@@ -112,7 +115,7 @@ def ds_nwp_ukv():
 
 
 @pytest.fixture(scope="session")
-def nwp_ukv_zarr_path(ds_nwp_ukv):
+def nwp_ukv_zarr_path(session_tmp_path, ds_nwp_ukv):
     ds = ds_nwp_ukv.chunk(
         {
             "init_time": 1, 
@@ -122,10 +125,9 @@ def nwp_ukv_zarr_path(ds_nwp_ukv):
             "y": 50,
         }
     )
-    with tempfile.TemporaryDirectory() as tmpdir:
-        filename = tmpdir + "/ukv_nwp.zarr"
-        ds.to_zarr(filename)
-        yield filename
+    zarr_path = session_tmp_path / "ukv_nwp.zarr"
+    ds.to_zarr(zarr_path)
+    yield zarr_path
 
 
 @pytest.fixture(scope="session")
@@ -155,7 +157,7 @@ def ds_nwp_ecmwf():
 
 
 @pytest.fixture(scope="session")
-def nwp_ecmwf_zarr_path(ds_nwp_ecmwf):
+def nwp_ecmwf_zarr_path(session_tmp_path, ds_nwp_ecmwf):
     ds = ds_nwp_ecmwf.chunk(
         {
             "init_time": 1, 
@@ -165,10 +167,10 @@ def nwp_ecmwf_zarr_path(ds_nwp_ecmwf):
             "latitude": 50,
         }
     )
-    with tempfile.TemporaryDirectory() as tmpdir:
-        filename = tmpdir + "/ukv_ecmwf.zarr"
-        ds.to_zarr(filename)
-        yield filename
+
+    zarr_path = session_tmp_path / "ukv_ecmwf.zarr"
+    ds.to_zarr(zarr_path)
+    yield zarr_path
 
 
 @pytest.fixture(scope="session")
@@ -201,7 +203,7 @@ def ds_uk_gsp():
 
 
 @pytest.fixture(scope="session")
-def data_sites() -> Generator[Site, None, None]:
+def data_sites(session_tmp_path) -> Site:
     """
     Make fake data for sites
     Returns: filename for netcdf file, and csv metadata
@@ -245,30 +247,27 @@ def data_sites() -> Generator[Site, None, None]:
         "generation_kw": da_gen,
     })
 
-    with tempfile.TemporaryDirectory() as tmpdir:
-        filename = tmpdir + "/sites.netcdf"
-        filename_csv = tmpdir + "/sites_metadata.csv"
-        generation.to_netcdf(filename)
-        meta_df.to_csv(filename_csv)
-
-        site = Site(
-            file_path=filename,
-            metadata_file_path=filename_csv,
-            interval_start_minutes=-30,
-            interval_end_minutes=60,
-            time_resolution_minutes=30,
-        )
+    filename = f"{session_tmp_path}/sites.netcdf"
+    filename_csv = f"{session_tmp_path}/sites_metadata.csv"
+    generation.to_netcdf(filename)
+    meta_df.to_csv(filename_csv)
+
+    site = Site(
+        file_path=filename,
+        metadata_file_path=filename_csv,
+        interval_start_minutes=-30,
+        interval_end_minutes=60,
+        time_resolution_minutes=30,
+    )
 
-        yield site
+    yield site
 
 
 @pytest.fixture(scope="session")
-def uk_gsp_zarr_path(ds_uk_gsp):
-
-    with tempfile.TemporaryDirectory() as tmpdir:
-        filename = tmpdir + "/uk_gsp.zarr"
-        ds_uk_gsp.to_zarr(filename)
-        yield filename
+def uk_gsp_zarr_path(session_tmp_path, ds_uk_gsp):
+    zarr_path = session_tmp_path / "uk_gsp.zarr"
+    ds_uk_gsp.to_zarr(zarr_path)
+    yield zarr_path
 
 
 @pytest.fixture()

tests/load/test_load_satellite.py

@@ -8,10 +8,10 @@ def test_open_satellite(sat_zarr_path):
 
     assert isinstance(da, xr.DataArray)
     assert da.dims == ("time_utc", "channel", "x_geostationary", "y_geostationary")
-    # 576 is 2 days of data at 5 minutes intervals, 12 * 24 * 2
+    # 288 is 1 days of data at 5 minutes intervals, 12 * 24
     # There are 11 channels
-    # There are 49 x 20 pixels
-    assert da.shape == (576, 11, 49, 20)
+    # There are 100 x 100 pixels
+    assert da.shape == (288, 11, 100, 100)
     assert np.issubdtype(da.dtype, np.number)
 
 

tests/numpy_sample/test_collate.py

@@ -1,17 +1,12 @@
-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.datasets.pvnet_uk_regional import PVNetUKRegionalDataset
+from ocf_data_sampler.torch_datasets.datasets.pvnet_uk import PVNetUKRegionalDataset
 
 
-def test_pvnet(pvnet_config_filename):
+def test_stack_np_samples_into_batch(pvnet_config_filename):
 
     # Create dataset object
     dataset = PVNetUKRegionalDataset(pvnet_config_filename)
 
-    assert len(dataset.locations) == 317
-    assert len(dataset.valid_t0_times) == 39
-    assert len(dataset) == 317 * 39
-
     # Generate 2 samples
     sample1 = dataset[0]
     sample2 = dataset[1]
@@ -22,5 +17,5 @@ def test_pvnet(pvnet_config_filename):
     assert "nwp" in batch
     assert isinstance(batch["nwp"], dict)
     assert "ukv" in batch["nwp"]
-    assert GSPSampleKey.gsp in batch
-    assert SatelliteSampleKey.satellite_actual in batch
+    assert "gsp" in batch
+    assert "satellite_actual" in batch

tests/torch_datasets/test_merge_and_fill_utils.py

@@ -33,9 +33,7 @@ def test_fill_nans_in_arrays():
     
     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.py

@@ -0,0 +1,166 @@
+import numpy as np
+import pandas as pd
+import xarray as xr
+import dask.array
+
+from ocf_data_sampler.config import load_yaml_configuration, save_yaml_configuration
+from ocf_data_sampler.torch_datasets.datasets.pvnet_uk import (
+    PVNetUKRegionalDataset,
+    PVNetUKConcurrentDataset,
+    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
+    sample = process_and_combine_datasets(dataset_dict, config, t0, location)
+
+    # Assert result is dict - check and validate
+    assert isinstance(sample, dict)
+    assert "nwp" in sample
+    assert sample["satellite_actual"].shape == (7, 1, 2, 2)
+    assert sample["nwp"]["ukv"]["nwp"].shape == (4, 1, 2, 2)
+
+
+def test_compute():
+    """Test compute function with dask array"""
+    da_dask = xr.DataArray(dask.array.random.random((5, 5)))
+
+    # Create a nested dictionary with dask array
+    lazy_data_dict = {
+        "array1": da_dask,
+        "nested": {
+            "array2": da_dask
+        }
+    }
+
+    computed_data_dict = compute(lazy_data_dict)
+
+    # Assert that the result is no longer lazy
+    assert isinstance(computed_data_dict["array1"].data, np.ndarray)
+    assert isinstance(computed_data_dict["nested"]["array2"].data, np.ndarray)
+
+    
+def test_pvnet_uk_regional_dataset(pvnet_config_filename):
+
+    # Create dataset object
+    dataset = PVNetUKRegionalDataset(pvnet_config_filename)
+
+    assert len(dataset.locations) == 317 # Number of regional GSPs
+    # NB. I have not checked the value (39 below) is in fact correct
+    assert len(dataset.valid_t0_times) == 39
+    assert len(dataset) == 317*39
+
+    # Generate a sample
+    sample = dataset[0]
+
+    assert isinstance(sample, dict)
+
+    for key in [
+        "nwp", "satellite_actual", "gsp",
+        "gsp_solar_azimuth", "gsp_solar_elevation",
+    ]:
+        assert key in sample
+    
+    for nwp_source in ["ukv"]:
+        assert nwp_source in sample["nwp"]
+
+    # Check the shape of the data is correct
+    # 30 minutes of 5 minute data (inclusive), one channel, 2x2 pixels
+    assert sample["satellite_actual"].shape == (7, 1, 2, 2)
+    # 3 hours of 60 minute data (inclusive), one channel, 2x2 pixels
+    assert sample["nwp"]["ukv"]["nwp"].shape == (4, 1, 2, 2)
+    # 3 hours of 30 minute data (inclusive)
+    assert sample["gsp"].shape == (7,)
+    # Solar angles have same shape as GSP data
+    assert sample["gsp_solar_azimuth"].shape == (7,)
+    assert sample["gsp_solar_elevation"].shape == (7,)
+
+
+def test_pvnet_no_gsp(tmp_path, pvnet_config_filename):
+
+    # Create new config without GSP inputs
+    config = load_yaml_configuration(pvnet_config_filename)
+    config.input_data.gsp.zarr_path = ''
+    new_config_path = tmp_path / "pvnet_config_no_gsp.yaml"
+    save_yaml_configuration(config, new_config_path)
+    
+    # Create dataset object
+    dataset = PVNetUKRegionalDataset(new_config_path)
+
+    # Generate a sample
+    _ = dataset[0]
+
+    
+def test_pvnet_uk_concurrent_dataset(pvnet_config_filename):
+
+    # Create dataset object using a limited set of GSPs for test
+    gsp_ids = [1,2,3]
+    num_gsps = len(gsp_ids)
+    
+    dataset = PVNetUKConcurrentDataset(pvnet_config_filename, gsp_ids=gsp_ids)
+
+    assert len(dataset.locations) == num_gsps # Number of regional GSPs
+    # NB. I have not checked the value (39 below) is in fact correct
+    assert len(dataset.valid_t0_times) == 39
+    assert len(dataset) == 39
+
+    # Generate a sample
+    sample = dataset[0]
+
+    assert isinstance(sample, dict)
+
+    for key in [
+        "nwp", "satellite_actual", "gsp",
+        "gsp_solar_azimuth", "gsp_solar_elevation",
+    ]:
+        assert key in sample
+    
+    for nwp_source in ["ukv"]:
+        assert nwp_source in sample["nwp"]
+
+    # Check the shape of the data is correct
+    # 30 minutes of 5 minute data (inclusive), one channel, 2x2 pixels
+    assert sample["satellite_actual"].shape == (num_gsps, 7, 1, 2, 2)
+    # 3 hours of 60 minute data (inclusive), one channel, 2x2 pixels
+    assert sample["nwp"]["ukv"]["nwp"].shape == (num_gsps, 4, 1, 2, 2)
+    # 3 hours of 30 minute data (inclusive)
+    assert sample["gsp"].shape == (num_gsps, 7,)
+    # Solar angles have same shape as GSP data
+    assert sample["gsp_solar_azimuth"].shape == (num_gsps, 7,)
+    assert sample["gsp_solar_elevation"].shape == (num_gsps, 7,)

tests/torch_datasets/test_site.py

@@ -1,11 +1,37 @@
-import pandas as pd
+import pytest
+
 import numpy as np
-from ocf_data_sampler.torch_datasets.datasets.site import SitesDataset, convert_from_dataset_to_dict_datasets, coarsen_data
-from xarray import Dataset, DataArray
+import pandas as pd
 import xarray as xr
 
 from torch.utils.data import DataLoader
 
+from ocf_data_sampler.config import load_yaml_configuration, save_yaml_configuration
+from ocf_data_sampler.torch_datasets.datasets.site import (
+    SitesDataset, convert_from_dataset_to_dict_datasets, coarsen_data
+)
+
+
+
+@pytest.fixture()
+def site_config_filename(tmp_path, config_filename, nwp_ukv_zarr_path, sat_zarr_path, data_sites):
+
+    # adjust config to point to the zarr file
+    config = load_yaml_configuration(config_filename)
+    config.input_data.nwp["ukv"].zarr_path = nwp_ukv_zarr_path
+    config.input_data.satellite.zarr_path = sat_zarr_path
+    config.input_data.site = data_sites
+    config.input_data.gsp = None
+
+    filename = f"{tmp_path}/configuration.yaml"
+    save_yaml_configuration(config, filename)
+    yield filename
+
+
+@pytest.fixture()
+def sites_dataset(site_config_filename):
+    return SitesDataset(site_config_filename)
+
 
 def test_site(site_config_filename):
 
@@ -18,7 +44,7 @@ def test_site(site_config_filename):
     # Generate a sample
     sample = dataset[0]
 
-    assert isinstance(sample, Dataset)
+    assert isinstance(sample, xr.Dataset)
 
     # Expected dimensions and data variables
     expected_dims = {
@@ -85,21 +111,14 @@ def test_site_time_filter_end(site_config_filename):
     assert len(dataset) == 0
 
 
-def test_site_get_sample(site_config_filename):
-
-    # Create dataset object
-    dataset = SitesDataset(site_config_filename)
-
-    assert len(dataset) == 410
-    sample = dataset.get_sample(t0=pd.Timestamp("2023-01-01 12:00"), site_id=1)
+def test_site_get_sample(sites_dataset):
+    sample = sites_dataset.get_sample(t0=pd.Timestamp("2023-01-01 12:00"), site_id=1)
 
 
-def test_convert_from_dataset_to_dict_datasets(site_config_filename):
-    # Create dataset object
-    dataset = SitesDataset(site_config_filename)
+def test_convert_from_dataset_to_dict_datasets(sites_dataset):
 
-    # Generate two samples
-    sample_xr = dataset[0]
+    # Generate sample
+    sample_xr = sites_dataset[0]
 
     sample = convert_from_dataset_to_dict_datasets(sample_xr)
 
@@ -109,9 +128,7 @@ def test_convert_from_dataset_to_dict_datasets(site_config_filename):
         assert key in sample
 
 
-def test_site_dataset_with_dataloader(site_config_filename):
-    # Create dataset object
-    dataset = SitesDataset(site_config_filename)
+def test_site_dataset_with_dataloader(sites_dataset):
 
     expected_coods = {
         "site__solar_azimuth",
@@ -122,10 +139,6 @@ def test_site_dataset_with_dataloader(site_config_filename):
         "site__date_sin",
     }
 
-    sample = dataset[0]
-    for key in expected_coods:
-        assert key in sample
-
     dataloader_kwargs = dict(
         shuffle=False,
         batch_size=None,
@@ -141,25 +154,23 @@ def test_site_dataset_with_dataloader(site_config_filename):
         persistent_workers=False,  # Not needed since we only enter the dataloader loop once
     )
 
-    dataloader = DataLoader(dataset, collate_fn=None, batch_size=None)
+    dataloader = DataLoader(sites_dataset, collate_fn=None, batch_size=None)
 
-    for i, sample in zip(range(1), dataloader):
+    sample = next(iter(dataloader))
+
+    # check that expected_dims is in the sample
+    for key in expected_coods:
+        assert key in sample
 
-        # check that expected_dims is in the sample
-        for key in expected_coods:
-            assert key in sample
 
+def test_process_and_combine_site_sample_dict(sites_dataset):
 
-def test_process_and_combine_site_sample_dict(site_config_filename):
-    # Load config
-    # config = load_yaml_configuration(pvnet_config_filename)
-    site_ds = SitesDataset(site_config_filename)
     # Specify minimal structure for testing
     raw_nwp_values = np.random.rand(4, 1, 2, 2)  # Single channel
     fake_site_values = np.random.rand(197)
     site_dict = {
         "nwp": {
-            "ukv": DataArray(
+            "ukv": xr.DataArray(
                 raw_nwp_values,
                 dims=["time_utc", "channel", "y", "x"],
                 coords={
@@ -168,7 +179,7 @@ def test_process_and_combine_site_sample_dict(site_config_filename):
                 },
             )
         },
-        "site": DataArray(
+        "site": xr.DataArray(
             fake_site_values,
             dims=["time_utc"],
             coords={
@@ -183,10 +194,10 @@ def test_process_and_combine_site_sample_dict(site_config_filename):
     print(f"Input site_dict: {site_dict}")
 
     # Call function
-    result = site_ds.process_and_combine_site_sample_dict(site_dict)
+    result = sites_dataset.process_and_combine_site_sample_dict(site_dict)
 
     # Assert to validate output structure
-    assert isinstance(result, Dataset), "Result should be an xarray.Dataset"
+    assert isinstance(result, xr.Dataset), "Result should be an xarray.Dataset"
     assert len(result.data_vars) > 0, "Dataset should contain data variables"
 
     # Validate variable via assertion and shape of such

tests/torch_datasets/conftest.py

@@ -1,18 +0,0 @@
-import pytest
-
-from ocf_data_sampler.config import load_yaml_configuration, save_yaml_configuration
-
-
-@pytest.fixture()
-def site_config_filename(tmp_path, config_filename, nwp_ukv_zarr_path, sat_zarr_path, data_sites):
-
-    # adjust config to point to the zarr file
-    config = load_yaml_configuration(config_filename)
-    config.input_data.nwp["ukv"].zarr_path = nwp_ukv_zarr_path
-    config.input_data.satellite.zarr_path = sat_zarr_path
-    config.input_data.site = data_sites
-    config.input_data.gsp = None
-
-    filename = f"{tmp_path}/configuration.yaml"
-    save_yaml_configuration(config, filename)
-    return filename

tests/torch_datasets/test_pvnet_uk_regional.py

@@ -1,136 +0,0 @@
-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.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):
-
-    # Create dataset object
-    dataset = PVNetUKRegionalDataset(pvnet_config_filename)
-
-    assert len(dataset.locations) == 317 # no of GSPs not including the National level
-    # NB. I have not checked this value is in fact correct, but it does seem to stay constant
-    assert len(dataset.valid_t0_times) == 39
-    assert len(dataset) == 317*39
-
-    # Generate a sample
-    sample = dataset[0]
-
-    assert isinstance(sample, dict)
-
-    for key in [
-        NWPSampleKey.nwp, SatelliteSampleKey.satellite_actual, GSPSampleKey.gsp,
-        GSPSampleKey.solar_azimuth, GSPSampleKey.solar_elevation,
-    ]:
-        assert key in sample
-    
-    for nwp_source in ["ukv"]:
-        assert nwp_source in sample[NWPSampleKey.nwp]
-
-    # check the shape of the data is correct
-    # 30 minutes of 5 minute data (inclusive), one channel, 2x2 pixels
-    assert sample[SatelliteSampleKey.satellite_actual].shape == (7, 1, 2, 2)
-    # 3 hours of 60 minute data (inclusive), one channel, 2x2 pixels
-    assert sample[NWPSampleKey.nwp]["ukv"][NWPSampleKey.nwp].shape == (4, 1, 2, 2)
-    # 3 hours of 30 minute data (inclusive)
-    assert sample[GSPSampleKey.gsp].shape == (7,)
-    # Solar angles have same shape as GSP data
-    assert sample[GSPSampleKey.solar_azimuth].shape == (7,)
-    assert sample[GSPSampleKey.solar_elevation].shape == (7,)
-
-def test_pvnet_no_gsp(pvnet_config_filename):
-
-    # load config
-    config = load_yaml_configuration(pvnet_config_filename)
-    # remove gsp
-    config.input_data.gsp.zarr_path = ''
-
-    # save temp config file
-    with tempfile.NamedTemporaryFile() as temp_config_file:
-        save_yaml_configuration(config, temp_config_file.name)
-        # Create dataset object
-        dataset = PVNetUKRegionalDataset(temp_config_file.name)
-
-        # Generate a sample
-        _ = dataset[0]