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

tests/select/test_select_time_slice.py

@@ -1,275 +0,0 @@
-from ocf_data_sampler.select.select_time_slice import select_time_slice, select_time_slice_nwp
-
-import numpy as np
-import pandas as pd
-import xarray as xr
-import pytest
-
-
-NWP_FREQ = pd.Timedelta("3h")
-
-@pytest.fixture(scope="module")
-def da_sat_like():
-    """Create dummy data which looks like satellite data"""
-    x = np.arange(-100, 100)
-    y = np.arange(-100, 100)
-    datetimes = pd.date_range("2024-01-02 00:00", "2024-01-03 00:00", freq="5min")
-
-    da_sat = xr.DataArray(
-        np.random.normal(size=(len(datetimes), len(x), len(y))),
-        coords=dict(
-            time_utc=(["time_utc"], datetimes),
-            x_geostationary=(["x_geostationary"], x),
-            y_geostationary=(["y_geostationary"], y),
-        )
-    )
-    return da_sat
-
-
-@pytest.fixture(scope="module")
-def da_nwp_like():
-    """Create dummy data which looks like NWP data"""
-
-    x = np.arange(-100, 100)
-    y = np.arange(-100, 100)
-    datetimes = pd.date_range("2024-01-02 00:00", "2024-01-03 00:00", freq=NWP_FREQ)
-    steps = pd.timedelta_range("0h", "16h", freq="1h")
-    channels = ["t", "dswrf"]
-
-    da_nwp = xr.DataArray(
-        np.random.normal(size=(len(datetimes), len(steps), len(channels), len(x), len(y))),
-        coords=dict(
-            init_time_utc=(["init_time_utc"], datetimes),
-            step=(["step"], steps),
-            channel=(["channel"], channels),
-            x_osgb=(["x_osgb"], x),
-            y_osgb=(["y_osgb"], y),
-        )
-    )
-    return da_nwp
-
-
-@pytest.mark.parametrize("t0_str", ["12:30", "12:40", "12:00"])
-def test_select_time_slice(da_sat_like, t0_str):
-    """Test the basic functionality of select_time_slice"""
-
-    # Slice parameters
-    t0 = pd.Timestamp(f"2024-01-02 {t0_str}")
-    interval_start = pd.Timedelta(-0, "min")
-    interval_end = pd.Timedelta(60, "min")
-    freq = pd.Timedelta("5min")
-
-    # Expect to return these timestamps from the selection
-    expected_datetimes = pd.date_range(t0 +interval_start, t0 + interval_end, freq=freq)
-
-    # Make the selection
-    sat_sample = select_time_slice(
-        da_sat_like,
-        t0=t0,
-        interval_start=interval_start,
-        interval_end=interval_end,
-        sample_period_duration=freq,
-    )
-
-    # Check the returned times are as expected
-    assert (sat_sample.time_utc == expected_datetimes).all()
-
-
-@pytest.mark.parametrize("t0_str", ["00:00", "00:25", "11:00", "11:55"])
-def test_select_time_slice_out_of_bounds(da_sat_like, t0_str):
-    """Test the behaviour of select_time_slice when the selection is out of bounds"""
-
-    # Slice parameters
-    t0 = pd.Timestamp(f"2024-01-02 {t0_str}")
-    interval_start = pd.Timedelta(-30, "min")
-    interval_end = pd.Timedelta(60, "min")
-    freq = pd.Timedelta("5min")
-
-    # The data is available between these times
-    min_time = pd.Timestamp(da_sat_like.time_utc.min().item())
-    max_time = pd.Timestamp(da_sat_like.time_utc.max().item())
-
-    # Expect to return these timestamps within the requested range
-    expected_datetimes = pd.date_range(
-        max(t0 + interval_start, min_time),
-        min(t0 + interval_end, max_time),
-        freq=freq,
-    )
-
-    # Make the partially out of bounds selection
-    sat_sample = select_time_slice(
-        da_sat_like,
-        t0=t0,
-        interval_start=interval_start,
-        interval_end=interval_end,
-        sample_period_duration=freq,
-    )
-
-    # Check the returned times are as expected
-    assert (sat_sample.time_utc == expected_datetimes).all()
-
-    
-    # Check all the values before the first timestamp available in the data are NaN
-    all_nan_space = sat_sample.isnull().all(dim=("x_geostationary", "y_geostationary"))
-    if expected_datetimes[0] < min_time:    
-        assert all_nan_space.sel(time_utc=slice(None, min_time-freq)).all(dim="time_utc") 
-
-    # Check all the values before the first timestamp available in the data are NaN
-    if expected_datetimes[-1] > max_time:    
-        assert all_nan_space.sel(time_utc=slice(max_time+freq, None)).all(dim="time_utc")   
-
-    # Check that none of the values between the first and last available timestamp are NaN
-    any_nan_space = sat_sample.isnull().any(dim=("x_geostationary", "y_geostationary"))
-    assert not any_nan_space.sel(time_utc=slice(min_time, max_time)).any(dim="time_utc")
-
-
-@pytest.mark.parametrize("t0_str", ["10:00", "11:00", "12:00"])
-def test_select_time_slice_nwp_basic(da_nwp_like, t0_str):
-    """Test the basic functionality of select_time_slice_nwp"""
-
-    # Slice parameters
-    t0 = pd.Timestamp(f"2024-01-02 {t0_str}")
-    interval_start = pd.Timedelta(-6, "h")
-    interval_end = pd.Timedelta(3, "h")
-    freq = pd.Timedelta("1h")
-
-    # Make the selection
-    da_slice = select_time_slice_nwp(
-        da_nwp_like,
-        t0,
-        sample_period_duration=freq,
-        interval_start=interval_start,
-        interval_end=interval_end,
-        dropout_timedeltas = None,
-        dropout_frac = 0,
-        accum_channels = [],
-        channel_dim_name = "channel",
-    )
-
-    # Check the target-times are as expected
-    expected_target_times = pd.date_range(t0 + interval_start, t0 + interval_end, freq=freq)
-    assert (da_slice.target_time_utc==expected_target_times).all()
-
-    # Check the init-times are as expected
-    # - Forecast frequency is `NWP_FREQ`, and we can't have selected future init-times
-    expected_init_times = pd.to_datetime(
-        [t if t<t0 else t0 for t in expected_target_times]
-    ).floor(NWP_FREQ)
-    assert (da_slice.init_time_utc==expected_init_times).all()
-
-
-@pytest.mark.parametrize("dropout_hours", [1, 2, 5])
-def test_select_time_slice_nwp_with_dropout(da_nwp_like, dropout_hours):
-    """Test the functionality of select_time_slice_nwp with dropout"""
-
-    t0 = pd.Timestamp("2024-01-02 12:00")
-    interval_start = pd.Timedelta(-6, "h")
-    interval_end = pd.Timedelta(3, "h")
-    freq = pd.Timedelta("1h")
-    dropout_timedelta = pd.Timedelta(f"-{dropout_hours}h")
-
-    da_slice = select_time_slice_nwp(
-        da_nwp_like,
-        t0,
-        sample_period_duration=freq,
-        interval_start=interval_start,
-        interval_end=interval_end,
-        dropout_timedeltas = [dropout_timedelta],
-        dropout_frac = 1,
-        accum_channels = [],
-        channel_dim_name = "channel",
-    )
-
-    # Check the target-times are as expected
-    expected_target_times = pd.date_range(t0 + interval_start, t0 + interval_end, freq=freq)
-    assert (da_slice.target_time_utc==expected_target_times).all()
-
-    # Check the init-times are as expected considering the delay
-    t0_delayed = t0 + dropout_timedelta
-    expected_init_times = pd.to_datetime(
-        [t if t<t0_delayed else t0_delayed for t in expected_target_times]
-    ).floor(NWP_FREQ)
-    assert (da_slice.init_time_utc==expected_init_times).all()
-
-
-@pytest.mark.parametrize("t0_str", ["10:00", "11:00", "12:00"])
-def test_select_time_slice_nwp_with_dropout_and_accum(da_nwp_like, t0_str):
-    """Test the functionality of select_time_slice_nwp with dropout and accumulated variables"""
-
-    # Slice parameters
-    t0 = pd.Timestamp(f"2024-01-02 {t0_str}")
-    interval_start = pd.Timedelta(-6, "h")
-    interval_end = pd.Timedelta(3, "h")
-    freq = pd.Timedelta("1h")
-    dropout_timedelta = pd.Timedelta("-2h")
-
-    t0_delayed = (t0 + dropout_timedelta).floor(NWP_FREQ)
-
-    da_slice = select_time_slice_nwp(
-        da_nwp_like,
-        t0,
-        sample_period_duration=freq,
-        interval_start=interval_start,
-        interval_end=interval_end,
-        dropout_timedeltas=[dropout_timedelta],
-        dropout_frac=1,
-        accum_channels=["dswrf"],
-        channel_dim_name="channel",
-    )
-
-    # Check the target-times are as expected
-    expected_target_times = pd.date_range(t0 + interval_start, t0 + interval_end, freq=freq)
-    assert (da_slice.target_time_utc==expected_target_times).all()
-
-    # Check the init-times are as expected considering the delay
-    expected_init_times = pd.to_datetime(
-        [t if t<t0_delayed else t0_delayed for t in expected_target_times]
-    ).floor(NWP_FREQ)
-    assert (da_slice.init_time_utc==expected_init_times).all()
-
-    # Check channels are as expected
-    assert (da_slice.channel.values == ["t", "diff_dswrf"]).all()
-
-    # Check the accummulated channel has been differenced correctly
-
-    # This part of the data is pulled from the init-time: t0_delayed
-    da_slice_accum = da_slice.sel(
-        target_time_utc=slice(t0_delayed, None), 
-        channel="diff_dswrf"
-    )
-
-    # Get the original data for the t0_delayed init-time, and diff it along steps
-    # then select the steps which are expected to be used in the above slice
-    da_orig_diffed = (
-        da_nwp_like.sel(
-            init_time_utc=t0_delayed, 
-            channel="dswrf", 
-        ).diff(dim="step", label="lower")
-        .sel(step=slice(t0-t0_delayed + interval_start, t0-t0_delayed + interval_end))
-    )
-
-    # Check the values are the same
-    assert (da_slice_accum.values == da_orig_diffed.values).all()
-
-    # Check the non-accummulated channel has not been differenced
-
-    # This part of the data is pulled from the init-time: t0_delayed
-    da_slice_nonaccum = da_slice.sel(
-        target_time_utc=slice(t0_delayed, None), 
-        channel="t"
-    )
-
-    # Get the original data for the t0_delayed init-time, and select the steps which are expected 
-    # to be used in the above slice
-    da_orig = (
-        da_nwp_like.sel(
-            init_time_utc=t0_delayed, 
-            channel="t", 
-        )
-        .sel(step=slice(t0-t0_delayed + interval_start, t0-t0_delayed + interval_end))
-    )
-
-    # Check the values are the same
-    assert (da_slice_nonaccum.values == da_orig.values).all()
-
-

tests/test_sample/test_base.py

@@ -1,164 +0,0 @@
-"""
-Base class testing - SampleBase
-"""
-
-import pytest
-import torch
-import numpy as np
-
-from pathlib import Path
-from ocf_data_sampler.sample.base import (
-    SampleBase,
-    batch_to_tensor,
-    copy_batch_to_device
-)
-
-class TestSample(SampleBase):
-    """
-    SampleBase for testing purposes
-    Minimal implementations - abstract methods
-    """
-
-    def __init__(self):
-        super().__init__()
-        self._data = {}
-
-    def plot(self, **kwargs):
-        """ Minimal plot implementation """
-        return None
-
-    def to_numpy(self) -> None:
-        """ Standard implementation """
-        return {key: np.array(value) for key, value in self._data.items()}
-
-    def save(self, path):
-        """ Minimal save implementation """
-        path = Path(path)
-        with open(path, 'wb') as f:
-            f.write(b'test_data')
-
-    @classmethod
-    def load(cls, path):
-        """ Minimal load implementation """
-        path = Path(path)
-        instance = cls()
-        return instance
-
-
-def test_sample_base_initialisation():
-    """ Initialisation of SampleBase subclass """
-
-    sample = TestSample()
-    assert hasattr(sample, '_data'), "Sample should have _data attribute"
-    assert sample._data == {}, "Sample should start with empty dict"
-
-
-def test_sample_base_save_load(tmp_path):
-    """ Test basic save and load functionality """
-
-    sample = TestSample()
-    sample._data['test_data'] = [1, 2, 3]
-
-    save_path = tmp_path / 'test_sample.dat'
-    sample.save(save_path)
-    assert save_path.exists()
-
-    loaded_sample = TestSample.load(save_path)
-    assert isinstance(loaded_sample, TestSample)
-
-
-def test_sample_base_abstract_methods():
-    """ Test abstract method enforcement """
-    
-    with pytest.raises(TypeError, match="Can't instantiate abstract class"):
-        SampleBase()
-
-
-def test_sample_base_to_numpy():
-    """ Test the to_numpy functionality """
-    import numpy as np
-    
-    sample = TestSample()
-    sample._data = {
-        'int_data': 42,
-        'list_data': [1, 2, 3]
-    }
-    numpy_data = sample.to_numpy()
-
-    assert isinstance(numpy_data, dict)
-    assert all(isinstance(value, np.ndarray) for value in numpy_data.values())
-    assert np.array_equal(numpy_data['list_data'], np.array([1, 2, 3]))
-
-
-def test_batch_to_tensor_nested():
-    """ Test nested dictionary conversion """
-    batch = {
-        'outer': {
-            'inner': np.array([1, 2, 3])
-        }
-    }
-    tensor_batch = batch_to_tensor(batch)
-    
-    assert torch.equal(tensor_batch['outer']['inner'], torch.tensor([1, 2, 3]))
-
-
-def test_batch_to_tensor_mixed_types():
-    """ Test handling of mixed data types """
-    batch = {
-        'tensor_data': np.array([1, 2, 3]),
-        'string_data': 'not_a_tensor',
-        'nested': {
-            'numbers': np.array([4, 5, 6]),
-            'text': 'still_not_a_tensor'
-        }
-    }
-    tensor_batch = batch_to_tensor(batch)
-    
-    assert isinstance(tensor_batch['tensor_data'], torch.Tensor)
-    assert isinstance(tensor_batch['string_data'], str)
-    assert isinstance(tensor_batch['nested']['numbers'], torch.Tensor)
-    assert isinstance(tensor_batch['nested']['text'], str)
-
-
-def test_batch_to_tensor_different_dtypes():
-    """ Test conversion of arrays with different dtypes """
-    batch = {
-        'float_data': np.array([1.0, 2.0, 3.0], dtype=np.float32),
-        'int_data': np.array([1, 2, 3], dtype=np.int64),
-        'bool_data': np.array([True, False, True], dtype=np.bool_)
-    }
-    tensor_batch = batch_to_tensor(batch)
-    
-    assert isinstance(tensor_batch['bool_data'], torch.Tensor)
-    assert tensor_batch['float_data'].dtype == torch.float32
-    assert tensor_batch['int_data'].dtype == torch.int64
-    assert tensor_batch['bool_data'].dtype == torch.bool
-
-
-def test_batch_to_tensor_multidimensional():
-    """ Test conversion of multidimensional arrays """
-    batch = {
-        'matrix': np.array([[1, 2], [3, 4]]),
-        'tensor': np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
-    }
-    tensor_batch = batch_to_tensor(batch)
-    
-    assert tensor_batch['matrix'].shape == (2, 2)
-    assert tensor_batch['tensor'].shape == (2, 2, 2)
-    assert torch.equal(tensor_batch['matrix'], torch.tensor([[1, 2], [3, 4]]))
-
-
-def test_copy_batch_to_device():
-    """ Test moving tensors to a different device """
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    batch = {
-        'tensor_data': torch.tensor([1, 2, 3]),
-        'nested': {
-            'matrix': torch.tensor([[1, 2], [3, 4]])
-        },
-        'non_tensor': 'unchanged'
-    }
-    moved_batch = copy_batch_to_device(batch, device)
-    assert moved_batch['tensor_data'].device == device
-    assert moved_batch['nested']['matrix'].device == device
-    assert moved_batch['non_tensor'] == 'unchanged'  # Non-tensors should remain unchanged

tests/test_sample/test_site_sample.py

@@ -1,195 +0,0 @@
-"""
-Site class testing - SiteSample
-"""
-
-import pytest
-import numpy as np
-import xarray as xr
-import pandas as pd
-
-from pathlib import Path
-from xarray import Dataset
-
-from ocf_data_sampler.sample.site import SiteSample
-
-
-@pytest.fixture
-def sample_data():
-    """ Fixture creation with sample data """
-
-    #  Time periods specified
-    init_time = pd.Timestamp("2023-01-01 00:00")
-    target_times = pd.date_range("2023-01-01 00:00", periods=4, freq="1h")
-    sat_times = pd.date_range("2023-01-01 00:00", periods=7, freq="5min")
-    site_times = pd.date_range("2023-01-01 00:00", periods=4, freq="30min")
-
-    # Defined steps for NWP data
-    steps = [(t - init_time) for t in target_times]
-    
-    # Create the sample dataset
-    return Dataset(
-        data_vars={
-            'nwp-ukv': (
-                ["nwp-ukv__target_time_utc", "nwp-ukv__channel", 
-                 "nwp-ukv__y_osgb", "nwp-ukv__x_osgb"],
-                np.random.rand(4, 1, 2, 2)
-            ),
-            'satellite': (
-                ["satellite__time_utc", "satellite__channel",
-                 "satellite__y_geostationary", "satellite__x_geostationary"],
-                np.random.rand(7, 1, 2, 2)
-            ),
-            'site': (["site__time_utc"], np.random.rand(4))
-        },
-        coords={
-            # NWP coords
-            'nwp-ukv__target_time_utc': target_times,
-            'nwp-ukv__channel': ['dswrf'],
-            'nwp-ukv__y_osgb': [0, 1],
-            'nwp-ukv__x_osgb': [0, 1],
-            'nwp-ukv__init_time_utc': init_time,
-            'nwp-ukv__step': ('nwp-ukv__target_time_utc', steps),
-            
-            # Sat coords
-            'satellite__time_utc': sat_times,
-            'satellite__channel': ['vis'],
-            'satellite__y_geostationary': [0, 1],
-            'satellite__x_geostationary': [0, 1],
-            
-            # Site coords
-            'site__time_utc': site_times,
-            'site__capacity_kwp': 1000.0,
-            'site__site_id': 1,
-            'site__longitude': -3.5,
-            'site__latitude': 51.5,
-            
-            # Site features as coords
-            'site__solar_azimuth': ('site__time_utc', np.random.rand(4)),
-            'site__solar_elevation': ('site__time_utc', np.random.rand(4)),
-            'site__date_cos': ('site__time_utc', np.random.rand(4)),
-            'site__date_sin': ('site__time_utc', np.random.rand(4)),
-            'site__time_cos': ('site__time_utc', np.random.rand(4)),
-            'site__time_sin': ('site__time_utc', np.random.rand(4))
-        }
-    )
-
-
-def test_site_sample_init():
-    """ Test initialisation """
-    sample = SiteSample()
-    assert isinstance(sample._data, dict)
-    assert len(sample._data) == 0
-
-
-def test_site_sample_with_data(sample_data):
-    """ Testing of defined sample with actual data """
-    sample = SiteSample()
-    sample._data = sample_data
-    
-    # Assert data structure
-    assert isinstance(sample._data, Dataset)
-    
-    # Assert dimensions / shapes
-    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",
-    }
-    assert set(sample._data.dims) == expected_dims    
-    assert sample._data["satellite"].values.shape == (7, 1, 2, 2)
-    assert sample._data["nwp-ukv"].values.shape == (4, 1, 2, 2)
-    assert sample._data["site"].values.shape == (4,)
-
-
-def test_save_load(tmp_path, sample_data):
-    """ Save and load functionality """
-    sample = SiteSample()
-    sample._data = sample_data    
-    filepath = tmp_path / "test_sample.nc"
-    sample.save(filepath)
-    
-    # Assert file exists and has content
-    assert filepath.exists()
-    assert filepath.stat().st_size > 0
-    
-    # Load and verify
-    loaded = SiteSample.load(filepath)
-    assert isinstance(loaded, SiteSample)
-    assert isinstance(loaded._data, Dataset)
-    
-    # Compare original / loaded data
-    xr.testing.assert_identical(sample._data, loaded._data)
-
-
-def test_invalid_save_format(sample_data):
-    """ Saving with invalid format """
-    sample = SiteSample()
-    sample._data = sample_data
-    with pytest.raises(ValueError, match="Only .nc format is supported"):
-        sample.save("invalid.txt")
-
-
-def test_invalid_load_format():
-    """ Loading with invalid format """
-    with pytest.raises(ValueError, match="Only .nc format is supported"):
-        SiteSample.load("invalid.txt")
-
-
-def test_invalid_data_type():
-    """ Handling of invalid data types """
-    sample = SiteSample()
-    sample._data = {"invalid": "data"}
-    
-    with pytest.raises(TypeError, match="Data must be xarray Dataset"):
-        sample.to_numpy()
-    
-    with pytest.raises(TypeError, match="Data must be xarray Dataset for saving"):
-        sample.save("test.nc")
-
-
-def test_to_numpy(sample_data):
-    """ To numpy conversion """
-    sample = SiteSample()
-    sample._data = sample_data
-    numpy_data = sample.to_numpy()
-    
-    # Assert structure
-    assert isinstance(numpy_data, dict)
-    assert 'site' in numpy_data
-    assert 'nwp' in numpy_data
-    
-    # Check site - numpy array instead of dict
-    site_data = numpy_data['site']
-    assert isinstance(site_data, np.ndarray)
-    assert site_data.ndim == 1
-    assert len(site_data) == 4
-    assert np.all(site_data >= 0) and np.all(site_data <= 1)
-    
-    # Check NWP
-    assert 'ukv' in numpy_data['nwp']
-    nwp_data = numpy_data['nwp']['ukv']
-    assert 'nwp' in nwp_data
-    assert nwp_data['nwp'].shape == (4, 1, 2, 2)
-
-
-def test_data_consistency(sample_data):
-    """ Consistency of data across operations """
-    sample = SiteSample()
-    sample._data = sample_data    
-    numpy_data = sample.to_numpy()
-    
-    # Assert components remain consistent after conversion above
-    assert numpy_data['nwp']['ukv']['nwp'].shape == (4, 1, 2, 2)
-    assert 'site' in numpy_data
-    
-    # Update site data checks to expect numpy array
-    assert isinstance(numpy_data['site'], np.ndarray)
-    assert numpy_data['site'].shape == (4,)
-    assert np.all(numpy_data['site'] >= 0)
-    assert np.all(numpy_data['site'] <= 1)