ocf-data-sampler 0.1.11__py3-none-any.whl → 0.1.17__py3-none-any.whl

tests/select/test_select_spatial_slice.py

@@ -1,154 +0,0 @@
-import numpy as np
-import xarray as xr
-from ocf_data_sampler.select.location import Location
-import pytest
-
-from ocf_data_sampler.select.select_spatial_slice import (
-    select_spatial_slice_pixels, _get_idx_of_pixel_closest_to_poi
-)
-
-@pytest.fixture(scope="module")
-def da():
-    # Create dummy data
-    x = np.arange(-100, 100)
-    y = np.arange(-100, 100)
-
-    da = xr.DataArray(
-        np.random.normal(size=(len(x), len(y))),
-        coords=dict(
-            x_osgb=(["x_osgb"], x),
-            y_osgb=(["y_osgb"], y),
-        )
-    )
-    return da
-
-
-def test_get_idx_of_pixel_closest_to_poi(da):
-    
-    idx_location  = _get_idx_of_pixel_closest_to_poi(
-        da,
-        location=Location(x=10, y=10, coordinate_system="osgb"),
-    )
-
-    assert idx_location.coordinate_system == "idx"
-    assert idx_location.x == 110
-    assert idx_location.y == 110
-
-
-
-
-def test_select_spatial_slice_pixels(da):
-
-    # Select window which lies within x-y bounds of the data
-    da_sliced = select_spatial_slice_pixels(
-        da,
-        location=Location(x=-90, y=-80, coordinate_system="osgb"),
-        width_pixels=10,
-        height_pixels=10,
-        allow_partial_slice=True,
-    )
-
-
-    assert isinstance(da_sliced, xr.DataArray)
-    assert (da_sliced.x_osgb.values == np.arange(-95, -85)).all()
-    assert (da_sliced.y_osgb.values == np.arange(-85, -75)).all()
-    # No padding in this case so no NaNs
-    assert not da_sliced.isnull().any()
-
-
-    # Select window where the edge of the window lies right on the edge of the data
-    da_sliced = select_spatial_slice_pixels(
-        da,
-        location=Location(x=-90, y=-80, coordinate_system="osgb"),
-        width_pixels=20,
-        height_pixels=20,
-        allow_partial_slice=True,
-    )
-
-    assert isinstance(da_sliced, xr.DataArray)
-    assert (da_sliced.x_osgb.values == np.arange(-100, -80)).all()
-    assert (da_sliced.y_osgb.values == np.arange(-90, -70)).all()
-    # No padding in this case so no NaNs
-    assert not da_sliced.isnull().any()
-
-    # Select window which is partially outside the boundary of the data - padded on left
-    da_sliced = select_spatial_slice_pixels(
-        da,
-        location=Location(x=-90, y=-80, coordinate_system="osgb"),
-        width_pixels=30,
-        height_pixels=30,
-        allow_partial_slice=True,
-    )
-
-    assert isinstance(da_sliced, xr.DataArray)
-    assert (da_sliced.x_osgb.values == np.arange(-105, -75)).all()
-    assert (da_sliced.y_osgb.values == np.arange(-95, -65)).all()
-    # Data has been padded on left by 5 NaN pixels
-    assert da_sliced.isnull().sum() == 5*len(da_sliced.y_osgb)
-
-
-    # Select window which is partially outside the boundary of the data - padded on right
-    da_sliced = select_spatial_slice_pixels(
-        da,
-        location=Location(x=90, y=-80, coordinate_system="osgb"),
-        width_pixels=30,
-        height_pixels=30,
-        allow_partial_slice=True,
-    )
-
-    assert isinstance(da_sliced, xr.DataArray)
-    assert (da_sliced.x_osgb.values == np.arange(75, 105)).all()
-    assert (da_sliced.y_osgb.values == np.arange(-95, -65)).all()
-    # Data has been padded on right by 5 NaN pixels
-    assert da_sliced.isnull().sum() == 5*len(da_sliced.y_osgb)
-
-
-    location = Location(x=-90, y=-0, coordinate_system="osgb")
-
-    # Select window which is partially outside the boundary of the data - padded on top
-    da_sliced = select_spatial_slice_pixels(
-        da,
-        location=Location(x=-90, y=95, coordinate_system="osgb"),
-        width_pixels=20,
-        height_pixels=20,
-        allow_partial_slice=True,
-    )
-
-    assert isinstance(da_sliced, xr.DataArray)
-    assert (da_sliced.x_osgb.values == np.arange(-100, -80)).all()
-    assert (da_sliced.y_osgb.values == np.arange(85, 105)).all()
-    # Data has been padded on top by 5 NaN pixels
-    assert da_sliced.isnull().sum() == 5*len(da_sliced.x_osgb)
-
-    # Select window which is partially outside the boundary of the data - padded on bottom
-    da_sliced = select_spatial_slice_pixels(
-        da,
-        location=Location(x=-90, y=-95, coordinate_system="osgb"),
-        width_pixels=20,
-        height_pixels=20,
-        allow_partial_slice=True,
-    )
-
-    assert isinstance(da_sliced, xr.DataArray)
-    assert (da_sliced.x_osgb.values == np.arange(-100, -80)).all()
-    assert (da_sliced.y_osgb.values == np.arange(-105, -85)).all()
-    # Data has been padded on bottom by 5 NaN pixels
-    assert da_sliced.isnull().sum() == 5*len(da_sliced.x_osgb)
-
-    # Select window which is partially outside the boundary of the data - padded right and bottom
-    da_sliced = select_spatial_slice_pixels(
-        da,
-        location=Location(x=90, y=-80, coordinate_system="osgb"),
-        width_pixels=50,
-        height_pixels=50,
-        allow_partial_slice=True,
-    )
-
-    assert isinstance(da_sliced, xr.DataArray)
-    assert (da_sliced.x_osgb.values == np.arange(65, 115)).all()
-    assert (da_sliced.y_osgb.values == np.arange(-105, -55)).all()
-    # Data has been padded on right by 15 pixels and bottom by 5 NaN pixels
-    assert da_sliced.isnull().sum() == 15*len(da_sliced.y_osgb) + 5*len(da_sliced.x_osgb) - 15*5
-
-
-

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