ocf-data-sampler 0.1.10__py3-none-any.whl → 0.1.16__py3-none-any.whl

tests/test_sample/test_uk_regional_sample.py

@@ -1,163 +0,0 @@
-"""
-UK Regional class testing - UKRegionalSample
-"""
-
-import pytest
-import numpy as np
-import torch
-import tempfile
-
-from ocf_data_sampler.numpy_sample import (
-    GSPSampleKey,
-    SatelliteSampleKey,
-    NWPSampleKey
-)
-
-from ocf_data_sampler.sample.uk_regional import UKRegionalSample
-
-
-# Fixture define
-@pytest.fixture
-def pvnet_config_filename(tmp_path):
-    """ Minimal config file - testing """
-    config_content = """
-    input_data:
-        gsp:
-            zarr_path: ""
-            time_resolution_minutes: 30
-            interval_start_minutes: -180
-            interval_end_minutes: 0
-        nwp:
-            ukv:
-                zarr_path: ""
-                image_size_pixels_height: 64
-                image_size_pixels_width: 64
-                time_resolution_minutes: 60
-                interval_start_minutes: -180
-                interval_end_minutes: 0
-                channels: ["t", "dswrf"]
-                provider: "ukv"
-        satellite:
-            zarr_path: ""
-            image_size_pixels_height: 64
-            image_size_pixels_width: 64
-            time_resolution_minutes: 30
-            interval_start_minutes: -180
-            interval_end_minutes: 0
-            channels: ["HRV"]
-    """
-    config_file = tmp_path / "test_config.yaml"
-    config_file.write_text(config_content)
-    return str(config_file)
-
-
-def create_test_data():
-    """ Synthetic data generation """
-
-    # Field / spatial coordinates
-    nwp_data = {
-        'nwp': np.random.rand(4, 1, 2, 2),
-        'x': np.array([1, 2]),
-        'y': np.array([1, 2]),
-        NWPSampleKey.channel_names: ['test_channel']
-    }
-    
-    return {
-        'nwp': {
-            'ukv': nwp_data
-        },
-        GSPSampleKey.gsp: np.random.rand(7),
-        SatelliteSampleKey.satellite_actual: np.random.rand(7, 1, 2, 2),
-        GSPSampleKey.solar_azimuth: np.random.rand(7),
-        GSPSampleKey.solar_elevation: np.random.rand(7)
-    }
-
-
-# UKRegionalSample testing
-def test_sample_init():
-    """ Initialisation """
-    sample = UKRegionalSample()
-    assert hasattr(sample, '_data'), "Sample should have _data attribute"
-    assert isinstance(sample._data, dict)
-    assert len(sample._data) == 0
-
-
-def test_sample_save_load():
-   sample = UKRegionalSample()
-   sample._data = create_test_data()
-   
-   with tempfile.NamedTemporaryFile(suffix='.pt') as tf:
-       sample.save(tf.name)
-       loaded = UKRegionalSample.load(tf.name)
-       
-       assert set(loaded._data.keys()) == set(sample._data.keys())
-       assert isinstance(loaded._data['nwp'], dict)
-       assert 'ukv' in loaded._data['nwp']
-
-       assert loaded._data[GSPSampleKey.gsp].shape == (7,)
-       assert loaded._data[SatelliteSampleKey.satellite_actual].shape == (7, 1, 2, 2) 
-       assert loaded._data[GSPSampleKey.solar_azimuth].shape == (7,)
-       assert loaded._data[GSPSampleKey.solar_elevation].shape == (7,)
-
-       np.testing.assert_array_almost_equal(
-           loaded._data[GSPSampleKey.gsp], 
-           sample._data[GSPSampleKey.gsp]
-       )
-
-
-def test_save_unsupported_format():
-   """ Test saving - unsupported file format """
-   sample = UKRegionalSample()
-   sample._data = create_test_data()
-  
-   with tempfile.NamedTemporaryFile(suffix='.npz') as tf:
-       with pytest.raises(ValueError, match="Only .pt format is supported"):
-           sample.save(tf.name)
-
-
-def test_load_unsupported_format():
-    """ Test loading - unsupported file format """
-
-    with tempfile.NamedTemporaryFile(suffix='.npz') as tf:
-        with pytest.raises(ValueError, match="Only .pt format is supported"):
-            UKRegionalSample.load(tf.name)
-
-
-def test_load_corrupted_file():
-    """ Test loading - corrupted / empty file """
-
-    with tempfile.NamedTemporaryFile(suffix='.pt') as tf:
-        with open(tf.name, 'wb') as f:
-            f.write(b'corrupted data')
-        
-        with pytest.raises(Exception):
-            UKRegionalSample.load(tf.name)
-
-
-def test_to_numpy():
-    """ To numpy conversion check """
-    sample = UKRegionalSample()
-    sample._data = {
-        'nwp': {
-            'ukv': {
-                'nwp': np.random.rand(4, 1, 2, 2),
-                'x': np.array([1, 2]),
-                'y': np.array([1, 2])
-            }
-        },
-        GSPSampleKey.gsp: np.random.rand(7),
-        SatelliteSampleKey.satellite_actual: np.random.rand(7, 1, 2, 2),
-        GSPSampleKey.solar_azimuth: np.random.rand(7),
-        GSPSampleKey.solar_elevation: np.random.rand(7)
-    }
-    
-    numpy_data = sample.to_numpy()
-    
-    # Check returned data matches
-    assert numpy_data == sample._data
-    assert len(numpy_data) == len(sample._data)
-    
-    # Assert specific keys and types
-    assert 'nwp' in numpy_data
-    assert isinstance(numpy_data['nwp']['ukv']['nwp'], np.ndarray)
-    assert numpy_data[GSPSampleKey.gsp].shape == (7,)

tests/torch_datasets/test_merge_and_fill_utils.py

@@ -1,40 +0,0 @@
-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 np.array_equal(result["array1"], np.array([1.0, 0.0, 3.0, 0.0]))
-    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

@@ -1,167 +0,0 @@
-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": ["t", "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, 2, 2, 2)
-    assert "gsp_id" in sample
-
-
-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,226 +0,0 @@
-import pytest
-
-import numpy as np
-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(tmp_path, site_config_filename):
-
-    # Create dataset object
-    dataset = SitesDataset(site_config_filename)
-
-    assert len(dataset) == 10 * 41
-    # TODO check 41
-
-    # Generate a sample
-    sample = dataset[0]
-
-    assert isinstance(sample, xr.Dataset)
-
-    # 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",
-    }
-
-    expected_coords_subset = {
-        "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"}
-
-
-    sample.to_netcdf(f"{tmp_path}/sample.nc")
-    sample = xr.open_dataset(f"{tmp_path}/sample.nc")
-
-    # Check dimensions
-    assert (
-        set(sample.dims) == expected_dims
-    ), f"Missing or extra dimensions: {set(sample.dims) ^ expected_dims}"
-    # Check data variables
-    assert (
-        set(sample.data_vars) == expected_data_vars
-    ), f"Missing or extra data variables: {set(sample.data_vars) ^ expected_data_vars}"
-
-    for coords in expected_coords_subset:
-        assert coords in sample.coords
-
-    # check the shape of the data is correct
-    # 30 minutes of 5 minute data (inclusive), one channel, 2x2 pixels
-    assert sample["satellite"].values.shape == (7, 1, 2, 2)
-    # 3 hours of 60 minute data (inclusive), one channel, 2x2 pixels
-    assert sample["nwp-ukv"].values.shape == (4, 1, 2, 2)
-    # 1.5 hours of 30 minute data (inclusive)
-    assert sample["site"].values.shape == (4,)
-
-
-def test_site_time_filter_start(site_config_filename):
-
-    # Create dataset object
-    dataset = SitesDataset(site_config_filename, start_time="2024-01-01")
-
-    assert len(dataset) == 0
-
-
-def test_site_time_filter_end(site_config_filename):
-
-    # Create dataset object
-    dataset = SitesDataset(site_config_filename, end_time="2000-01-01")
-
-    assert len(dataset) == 0
-
-
-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(sites_dataset):
-
-    # Generate sample
-    sample_xr = sites_dataset[0]
-
-    sample = convert_from_dataset_to_dict_datasets(sample_xr)
-
-    assert isinstance(sample, dict)
-
-    for key in ["nwp", "satellite", "site"]:
-        assert key in sample
-
-
-def test_site_dataset_with_dataloader(sites_dataset):
-
-    expected_coods = {
-        "site__solar_azimuth",
-        "site__solar_elevation",
-        "site__date_cos",
-        "site__time_cos",
-        "site__time_sin",
-        "site__date_sin",
-    }
-
-    dataloader_kwargs = dict(
-        shuffle=False,
-        batch_size=None,
-        sampler=None,
-        batch_sampler=None,
-        num_workers=1,
-        collate_fn=None,
-        pin_memory=False,  # Only using CPU to prepare samples so pinning is not beneficial
-        drop_last=False,
-        timeout=0,
-        worker_init_fn=None,
-        prefetch_factor=1,
-        persistent_workers=False,  # Not needed since we only enter the dataloader loop once
-    )
-
-    dataloader = DataLoader(sites_dataset, collate_fn=None, batch_size=None)
-
-    sample = next(iter(dataloader))
-
-    # 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):
-
-    # 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": xr.DataArray(
-                raw_nwp_values,
-                dims=["time_utc", "channel", "y", "x"],
-                coords={
-                    "time_utc": pd.date_range("2024-01-01 00:00", periods=4, freq="h"),
-                    "channel": ["dswrf"],  # Single channel
-                },
-            )
-        },
-        "site": xr.DataArray(
-            fake_site_values,
-            dims=["time_utc"],
-            coords={
-                    "time_utc": pd.date_range("2024-01-01 00:00", periods=197, freq="15min"),
-                    "capacity_kwp": 1000,
-                    "site_id": 1,
-                    "longitude": -3.5,
-                    "latitude": 51.5
-                }
-        )
-    }
-    print(f"Input site_dict: {site_dict}")
-
-    # Call function
-    result = sites_dataset.process_and_combine_site_sample_dict(site_dict)
-
-    # Assert to validate output structure
-    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
-    expected_variables = ["nwp-ukv", "site"]
-    for expected_variable in expected_variables:
-        assert expected_variable in result.data_vars, f"Expected variable '{expected_variable}' not found"
-    
-    nwp_result = result["nwp-ukv"]
-    assert nwp_result.shape == (4, 1, 2, 2), f"Unexpected shape for nwp-ukv : {nwp_result.shape}"
-    site_result = result["site"]
-    assert site_result.shape == (197,), f"Unexpected shape for site: {site_result.shape}"
-
-
-def test_potentially_coarsen(ds_nwp_ecmwf):
-    """Test potentially_coarsen function with ECMWF_UK data."""
-    nwp_data = ds_nwp_ecmwf
-    assert nwp_data.ECMWF_UK.shape[3:] == (15, 12)  # Check initial shape (lon, lat)
-
-    data = coarsen_data(xr_data=nwp_data, coarsen_to_deg=2)
-    assert data.ECMWF_UK.shape[3:] == (8, 6)  # Coarsen to every 2 degrees
-
-    data = coarsen_data(xr_data=nwp_data, coarsen_to_deg=3)
-    assert data.ECMWF_UK.shape[3:] == (5, 4)  # Coarsen to every 3 degrees
-
-    data = coarsen_data(xr_data=nwp_data, coarsen_to_deg=1)
-    assert data.ECMWF_UK.shape[3:] == (15, 12)  # No coarsening (same shape)

tests/torch_datasets/test_validate_channels_utils.py

@@ -1,78 +0,0 @@
-"""Tests for channel validation utility functions"""
-
-import pytest
-from ocf_data_sampler.torch_datasets.utils.validate_channels import (
-    validate_channels,
-    validate_nwp_channels,
-    validate_satellite_channels,
-)
-
-
-class TestChannelValidation:
-    """Tests for channel validation functions"""
-    
-    @pytest.mark.parametrize("test_case", [
-        # Base validation - success case
-        {
-            "data_channels": ["channel1", "channel2"],
-            "norm_channels": ["channel1", "channel2", "extra"],
-            "source_name": "test_source",
-            "expect_error": False
-        },
-        # Base validation - error case
-        {
-            "data_channels": ["channel1", "missing_channel"],
-            "norm_channels": ["channel1"],
-            "source_name": "test_source",
-            "expect_error": True,
-            "error_match": "following channels for test_source are missing in normalisation means"
-        },
-        # NWP case - success
-        {
-            "data_channels": ["t2m", "dswrf"],
-            "norm_channels": ["t2m", "dswrf", "extra"],
-            "source_name": "ecmwf",
-            "expect_error": False
-        },
-        # NWP case - error
-        {
-            "data_channels": ["t2m", "missing_channel"],
-            "norm_channels": ["t2m"],
-            "source_name": "ecmwf",
-            "expect_error": True,
-            "error_match": "following channels for ecmwf are missing in normalisation means"
-        },
-        # Satellite case - success
-        {
-            "data_channels": ["IR_016", "VIS006"],
-            "norm_channels": ["IR_016", "VIS006", "extra"],
-            "source_name": "satellite",
-            "expect_error": False
-        },
-        # Satellite case - error
-        {
-            "data_channels": ["IR_016", "missing_channel"],
-            "norm_channels": ["IR_016"],
-            "source_name": "satellite",
-            "expect_error": True,
-            "error_match": "following channels for satellite are missing in normalisation means"
-        }
-    ])
-    def test_channel_validation(self, test_case):
-        """Test channel validation for both base, NWP and satellite data"""
-        if test_case["expect_error"]:
-            with pytest.raises(ValueError, match=test_case["error_match"]):
-                validate_channels(
-                    data_channels=test_case["data_channels"],
-                    means_channels=test_case["norm_channels"],
-                    stds_channels=test_case["norm_channels"],
-                    source_name=test_case["source_name"]
-                )
-        else:
-            # Should not raise any exceptions
-            validate_channels(
-                data_channels=test_case["data_channels"],
-                means_channels=test_case["norm_channels"],
-                stds_channels=test_case["norm_channels"],
-                source_name=test_case["source_name"]
-            )