ocf-data-sampler 0.0.19__py3-none-any.whl → 0.0.21__py3-none-any.whl

ocf_data_sampler-0.0.21.dist-info/RECORD

@@ -0,0 +1,53 @@
+ocf_data_sampler/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+ocf_data_sampler/config/__init__.py,sha256=YXnAkgHViHB26hSsjiv32b6EbpG-A1kKTkARJf0_RkY,212
+ocf_data_sampler/config/load.py,sha256=4f7vPHAIAmd-55tPxoIzn7F_TI_ue4NxkDcLPoVWl0g,943
+ocf_data_sampler/config/model.py,sha256=lAfryeJnwkFOaeuC74aYju38QrCPugSt5K2XDw1WZYI,8128
+ocf_data_sampler/config/save.py,sha256=wKdctbv0dxIIiQtcRHLRxpWQVhEFQ_FCWg-oNaRLIps,1093
+ocf_data_sampler/data/uk_gsp_locations.csv,sha256=RSh7DRh55E3n8lVAaWXGTaXXHevZZtI58td4d4DhGos,10415772
+ocf_data_sampler/load/__init__.py,sha256=MjgfxilTzyz1RYFoBEeAXmE9hyjknLvdmlHPmlAoiQY,44
+ocf_data_sampler/load/gsp.py,sha256=Gcr1JVUOPKhFRDCSHtfPDjxx0BtyyEhXrZvGEKLPJ5I,759
+ocf_data_sampler/load/satellite.py,sha256=3KlA1fx4SwxdzM-jC1WRaONXO0D6m0WxORnEnwUnZrA,2967
+ocf_data_sampler/load/utils.py,sha256=EQGvVWlGMoSOdbDYuMfVAa0v6wmAOPmHIAemdrTB5v4,1406
+ocf_data_sampler/load/nwp/__init__.py,sha256=SmcrnbygO5xtCKmGR4wtHrj-HI7nOAvnAtfuvRufBGQ,25
+ocf_data_sampler/load/nwp/nwp.py,sha256=O4QnajEZem8BvBgTcYYDBhRhgqPYuJkolHmpMRmrXEA,610
+ocf_data_sampler/load/nwp/providers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ocf_data_sampler/load/nwp/providers/ecmwf.py,sha256=vW-p3vCyQ-CofKo555-gE7VDi5hlpjtjTLfHqWF0HEE,1175
+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_batch/__init__.py,sha256=mrtqwbGik5Zc9MYP5byfCTBm08wMtS2XnTsypC4fPMo,245
+ocf_data_sampler/numpy_batch/gsp.py,sha256=3gwSj0k29JyA8_09zovB8f8Pr-dVhCuMSO1-k4QKAOg,668
+ocf_data_sampler/numpy_batch/nwp.py,sha256=Rv0yfDj902Z2oCwdlRjOs3Kh-F5Fgxjjylh99-lQ9ws,1105
+ocf_data_sampler/numpy_batch/satellite.py,sha256=e6eoNmiiHtzZbDVtBolFzDuE3qwhHN6bL9H86emAUsk,732
+ocf_data_sampler/numpy_batch/sun_position.py,sha256=UW6-WtjrKdCkcguolHUDSLhYFfarknQzzjlCX8YdEOM,1700
+ocf_data_sampler/select/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+ocf_data_sampler/select/dropout.py,sha256=zDpVLMjGb70RRyYKN-WI2Kp3x9SznstT4cMcZ4dsvJg,1066
+ocf_data_sampler/select/fill_time_periods.py,sha256=iTtMjIPFYG5xtUYYedAFBLjTWWUa7t7WQ0-yksWf0-E,440
+ocf_data_sampler/select/find_contiguous_time_periods.py,sha256=6ioB8LeFpFNBMgKDxrgG3zqzNjkBF_jlV9yye2ZYT2E,11925
+ocf_data_sampler/select/geospatial.py,sha256=oHJoKEKubn3v3yKCVeuiPxuGroVA4RyrpNi6ARq5woE,3558
+ocf_data_sampler/select/location.py,sha256=26Y5ZjfFngShBwXieuWSoOA-RLaRzci4TTmcDk3Wg7U,2015
+ocf_data_sampler/select/select_spatial_slice.py,sha256=hWIJe4_VzuQ2iiiQh7V17AXwTILT5kIkUvzG458J_Gw,11220
+ocf_data_sampler/select/select_time_slice.py,sha256=41cch1fQr59fZgv7UHsNGc3OvoynrixT3bmr3_1d7cU,6628
+ocf_data_sampler/torch_datasets/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+ocf_data_sampler/torch_datasets/pvnet_uk_regional.py,sha256=3b44y-cgIa0Nsa4RqFCKUyhs8cIXc04_j0JhMQZFIn8,19127
+tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tests/conftest.py,sha256=O77dmow8mGpGPbZ6Pz7ma7cLaiV1k8mxW1eYg37Avrw,5585
+tests/config/test_config.py,sha256=G_PD_pXib0zdRBPUIn0jjwJ9VyoKaO_TanLN1Mh5Ca4,5055
+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/numpy_batch/test_gsp.py,sha256=uAnDmJKMp8BVm8olTN8d499VGd4dKSZbZbhZq9LSV4M,574
+tests/numpy_batch/test_nwp.py,sha256=2dCmsLztzqPCpUoUv-d3F8sXWzYaHk_9PxqJz99SuzI,1477
+tests/numpy_batch/test_satellite.py,sha256=Ysf0RUAgtxebaJ8F7oiR9ujucZ4ZkvBbGrCSOSPsgHE,1193
+tests/numpy_batch/test_sun_position.py,sha256=oH6aSjsKhHPJ-Hy-G5Z_BMcQclvBuDOEVmPq9c-TVyc,2492
+tests/select/test_dropout.py,sha256=kiycl7RxAQYMCZJlokmx6Da5h_oBpSs8Is8pmSW4gOU,2413
+tests/select/test_fill_time_periods.py,sha256=o59f2YRe5b0vJrG3B0aYZkYeHnpNk4s6EJxdXZluNQg,907
+tests/select/test_find_contiguous_time_periods.py,sha256=G6tJRJd0DMfH9EdfzlKWsmfTbtMwOf3w-2filjJzuIQ,5998
+tests/select/test_location.py,sha256=_WZk2FPYeJ-nIfCJS6Sp_yaVEEo7m31DmMFoZzgyCts,2712
+tests/select/test_select_spatial_slice.py,sha256=m_R0wS5ApoBfE9h28FgdwYqT3gFxle6d3CBtmoqdPmo,5134
+tests/select/test_select_time_slice.py,sha256=rH4h90HdQCoWE7vV7ivMEKhiCStQDEcMBCPamiDuO0k,10147
+tests/torch_datasets/test_pvnet_uk_regional.py,sha256=R_b8wbAYHqW5KQDmLLNIN8Y8uz5vslG2LPYvHHHSvvg,2615
+ocf_data_sampler-0.0.21.dist-info/LICENSE,sha256=F-Q3UFCR-BECSocV55BFDpn4YKxve9PKrm-lTt6o_Tg,1073
+ocf_data_sampler-0.0.21.dist-info/METADATA,sha256=_8fQPhtqm7tXX2_JBrkbDlUKnGY4QhG5q2uy89ag5cs,5167
+ocf_data_sampler-0.0.21.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+ocf_data_sampler-0.0.21.dist-info/top_level.txt,sha256=KaQn5qzkJGJP6hKWqsVAc9t0cMLjVvSTk8-kTrW79SA,23
+ocf_data_sampler-0.0.21.dist-info/RECORD,,

tests/config/test_config.py

@@ -0,0 +1,152 @@
+import tempfile
+
+import pytest
+from pydantic import ValidationError
+
+from ocf_data_sampler.config import (
+    load_yaml_configuration,
+    Configuration,
+    save_yaml_configuration
+)
+
+
+def test_default():
+    """Test default pydantic class"""
+
+    _ = Configuration()
+
+
+def test_yaml_load_test_config(test_config_filename):
+    """
+    Test that yaml loading works for 'test_config.yaml'
+    and fails for an empty .yaml file
+    """
+
+    # check we get an error if loading a file with no config
+    with tempfile.NamedTemporaryFile(suffix=".yaml") as fp:
+        filename = fp.name
+
+        # check that temp file can't be loaded
+        with pytest.raises(TypeError):
+            _ = load_yaml_configuration(filename)
+
+    # test can load test_config.yaml
+    config = load_yaml_configuration(test_config_filename)
+
+    assert isinstance(config, Configuration)
+
+
+def test_yaml_save(test_config_filename):
+    """
+    Check configuration can be saved to a .yaml file
+    """
+
+    test_config = load_yaml_configuration(test_config_filename)
+
+    with tempfile.NamedTemporaryFile(suffix=".yaml") as fp:
+        filename = fp.name
+
+        # save default config to file
+        save_yaml_configuration(test_config, filename)
+
+        # check the file can be loaded back
+        tmp_config = load_yaml_configuration(filename)
+
+        # check loaded configuration is the same as the one passed to save
+        assert test_config == tmp_config
+
+
+def test_extra_field():
+    """
+    Check an extra parameters in config causes error
+    """
+
+    configuration = Configuration()
+    configuration_dict = configuration.model_dump()
+    configuration_dict["extra_field"] = "extra_value"
+    with pytest.raises(ValidationError, match="Extra inputs are not permitted"):
+        _ = Configuration(**configuration_dict)
+
+
+def test_incorrect_forecast_minutes(test_config_filename):
+    """
+    Check a forecast length not divisible by time resolution causes error
+    """
+
+    configuration = load_yaml_configuration(test_config_filename)
+
+    configuration.input_data.nwp['ukv'].forecast_minutes = 1111
+    with pytest.raises(Exception, match="duration must be divisible by time resolution"):
+        _ = Configuration(**configuration.model_dump())
+
+
+def test_incorrect_history_minutes(test_config_filename):
+    """
+    Check a history length not divisible by time resolution causes error
+    """
+
+    configuration = load_yaml_configuration(test_config_filename)
+
+    configuration.input_data.nwp['ukv'].history_minutes = 1111
+    with pytest.raises(Exception, match="duration must be divisible by time resolution"):
+        _ = Configuration(**configuration.model_dump())
+
+
+def test_incorrect_nwp_provider(test_config_filename):
+    """
+    Check an unexpected nwp provider causes error
+    """
+
+    configuration = load_yaml_configuration(test_config_filename)
+
+    configuration.input_data.nwp['ukv'].nwp_provider = "unexpected_provider"
+    with pytest.raises(Exception, match="NWP provider"):
+        _ = Configuration(**configuration.model_dump())
+
+def test_incorrect_dropout(test_config_filename):
+    """
+    Check a dropout timedelta over 0 causes error and 0 doesn't
+    """
+
+    configuration = load_yaml_configuration(test_config_filename)
+
+    # check a positive number is not allowed
+    configuration.input_data.nwp['ukv'].dropout_timedeltas_minutes = [120]
+    with pytest.raises(Exception, match="Dropout timedeltas must be negative"):
+        _ = Configuration(**configuration.model_dump())
+
+    # check 0 is allowed
+    configuration.input_data.nwp['ukv'].dropout_timedeltas_minutes = [0]
+    _ = Configuration(**configuration.model_dump())
+
+def test_incorrect_dropout_fraction(test_config_filename):
+    """
+    Check dropout fraction outside of range causes error
+    """
+
+    configuration = load_yaml_configuration(test_config_filename)
+
+    configuration.input_data.nwp['ukv'].dropout_fraction= 1.1
+    with pytest.raises(Exception, match="Dropout fraction must be between 0 and 1"):
+        _ = Configuration(**configuration.model_dump())
+
+    configuration.input_data.nwp['ukv'].dropout_fraction= -0.1
+    with pytest.raises(Exception, match="Dropout fraction must be between 0 and 1"):
+        _ = Configuration(**configuration.model_dump())
+
+
+def test_inconsistent_dropout_use(test_config_filename):
+    """
+    Check dropout fraction outside of range causes error
+    """
+
+    configuration = load_yaml_configuration(test_config_filename)
+    configuration.input_data.satellite.dropout_fraction= 1.0
+    configuration.input_data.satellite.dropout_timedeltas_minutes = None
+
+    with pytest.raises(ValueError, match="To dropout fraction > 0 requires a list of dropout timedeltas"):
+        _ = Configuration(**configuration.model_dump())
+    configuration.input_data.satellite.dropout_fraction= 0.0
+    configuration.input_data.satellite.dropout_timedeltas_minutes = [-120, -60]
+    with pytest.raises(ValueError, match="To use dropout timedeltas dropout fraction should be > 0"):
+        _ = Configuration(**configuration.model_dump())

tests/conftest.py

@@ -6,11 +6,16 @@ import pytest
 import xarray as xr
 import tempfile
 
+_top_test_directory = os.path.dirname(os.path.realpath(__file__))
+
+@pytest.fixture()
+def test_config_filename():
+    return f"{_top_test_directory}/test_data/configs/test_config.yaml"
 
 
 @pytest.fixture(scope="session")
 def config_filename():
-    return f"{os.path.dirname(os.path.abspath(__file__))}/test_data/pvnet_test_config.yaml"
+    return f"{os.path.dirname(os.path.abspath(__file__))}/test_data/configs/pvnet_test_config.yaml"
 
 
 @pytest.fixture(scope="session")

tests/load/test_load_gsp.py

@@ -0,0 +1,15 @@
+from ocf_data_sampler.load.gsp import open_gsp
+import xarray as xr
+
+
+def test_open_gsp(uk_gsp_zarr_path):
+    da = open_gsp(uk_gsp_zarr_path)
+
+    assert isinstance(da, xr.DataArray)
+    assert da.dims == ("time_utc", "gsp_id")
+
+    assert "nominal_capacity_mwp" in da.coords
+    assert "effective_capacity_mwp" in da.coords
+    assert "x_osgb" in da.coords
+    assert "y_osgb" in da.coords
+    assert da.shape == (49, 318)

tests/load/test_load_nwp.py

@@ -0,0 +1,21 @@
+import pandas as pd
+from xarray import DataArray
+import numpy as np
+
+from ocf_data_sampler.load.nwp import open_nwp
+
+
+def test_load_ukv(nwp_ukv_zarr_path):
+    da = open_nwp(zarr_path=nwp_ukv_zarr_path, provider="ukv")
+    assert isinstance(da, DataArray)
+    assert da.dims == ("init_time_utc", "step", "channel", "x_osgb", "y_osgb")
+    assert da.shape == (24 * 7, 11, 4, 50, 100)
+    assert np.issubdtype(da.dtype, np.number)
+
+
+def test_load_ecmwf(nwp_ecmwf_zarr_path):
+    da = open_nwp(zarr_path=nwp_ecmwf_zarr_path, provider="ecmwf")
+    assert isinstance(da, DataArray)
+    assert da.dims == ("init_time_utc", "step", "channel", "longitude", "latitude")
+    assert da.shape == (24 * 7, 15, 3, 15, 12)
+    assert np.issubdtype(da.dtype, np.number)

tests/load/test_load_satellite.py

@@ -0,0 +1,17 @@
+from ocf_data_sampler.load.satellite import open_sat_data
+import xarray as xr
+import numpy as np
+
+
+def test_open_satellite(sat_zarr_path):
+    da = open_sat_data(zarr_path=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
+    # There are 11 channels
+    # There are 49 x 20 pixels
+    assert da.shape == (576, 11, 49, 20)
+    assert np.issubdtype(da.dtype, np.number)
+
+

tests/numpy_batch/test_gsp.py

@@ -0,0 +1,23 @@
+from ocf_datapipes.batch import BatchKey
+from ocf_data_sampler.load.gsp import open_gsp
+
+from ocf_data_sampler.numpy_batch import convert_gsp_to_numpy_batch
+
+
+def test_convert_gsp_to_numpy_batch(uk_gsp_zarr_path):
+
+    da = (
+        open_gsp(uk_gsp_zarr_path)
+        .isel(time_utc=slice(0, 10))
+        .sel(gsp_id=1)
+    )
+
+    # Call the function
+    numpy_batch = convert_gsp_to_numpy_batch(da)
+
+    # Assert the output type
+    assert isinstance(numpy_batch, dict)
+
+    # Assert the shape of the numpy batch
+    assert (numpy_batch[BatchKey.gsp] == da.values).all()
+

tests/numpy_batch/test_nwp.py

@@ -0,0 +1,54 @@
+import numpy as np
+import pandas as pd
+import xarray as xr
+
+import pytest
+
+from ocf_data_sampler.numpy_batch import convert_nwp_to_numpy_batch
+
+from ocf_datapipes.batch import NWPBatchKey
+
+@pytest.fixture(scope="module")
+def da_nwp_like():
+    """Create dummy data which looks like time-sliced NWP data"""
+
+    t0 = pd.to_datetime("2024-01-02 00:00")
+
+    x = np.arange(-100, 100, 10)
+    y = np.arange(-100, 100, 10)
+    steps = pd.timedelta_range("0H", "8H", freq="1H")
+    target_times = t0 + steps
+
+    channels = ["t", "dswrf"]
+    init_times = pd.to_datetime([t0]*len(steps))
+    
+    # Create dummy time-sliced NWP data
+    da_nwp = xr.DataArray(
+        np.random.normal(size=(len(target_times), len(channels), len(x), len(y))),
+        coords=dict(
+            target_times_utc=(["target_times_utc"], target_times),
+            channel=(["channel"], channels),
+            x_osgb=(["x_osgb"], x),
+            y_osgb=(["y_osgb"], y),
+        )
+    )
+
+    # Add extra non-coordinate dimensions
+    da_nwp = da_nwp.assign_coords(
+        init_time_utc=("target_times_utc", init_times),
+        step=("target_times_utc", steps),
+    )
+
+    return da_nwp
+
+
+def test_convert_nwp_to_numpy_batch(da_nwp_like):
+
+    # Call the function
+    numpy_batch = convert_nwp_to_numpy_batch(da_nwp_like)
+
+    # Assert the output type
+    assert isinstance(numpy_batch, dict)
+
+    # Assert the shape of the numpy batch
+    assert (numpy_batch[NWPBatchKey.nwp] == da_nwp_like.values).all()

tests/numpy_batch/test_satellite.py

@@ -0,0 +1,42 @@
+
+import numpy as np
+import pandas as pd
+import xarray as xr
+
+import pytest
+
+from ocf_data_sampler.numpy_batch import convert_satellite_to_numpy_batch
+
+from ocf_datapipes.batch import BatchKey
+
+
+@pytest.fixture(scope="module")
+def da_sat_like():
+    """Create dummy data which looks like satellite data"""
+    x = np.arange(-100, 100, 10)
+    y = np.arange(-100, 100, 10)
+    datetimes = pd.date_range("2024-01-01 12:00", "2024-01-01 12:30", freq="5min")
+    channels = ["VIS008", "IR016"]
+
+    da_sat = xr.DataArray(
+        np.random.normal(size=(len(datetimes), len(channels), len(x), len(y))),
+        coords=dict(
+            time_utc=(["time_utc"], datetimes),
+            channel=(["channel"], channels),
+            x_geostationary=(["x_geostationary"], x),
+            y_geostationary=(["y_geostationary"], y),
+        )
+    )
+    return da_sat
+
+
+def test_convert_satellite_to_numpy_batch(da_sat_like):
+
+    # Call the function
+    numpy_batch = convert_satellite_to_numpy_batch(da_sat_like)
+
+    # Assert the output type
+    assert isinstance(numpy_batch, dict)
+
+    # Assert the shape of the numpy batch
+    assert (numpy_batch[BatchKey.satellite_actual] == da_sat_like.values).all()

tests/numpy_batch/test_sun_position.py

@@ -0,0 +1,81 @@
+import numpy as np
+import pandas as pd
+import pytest
+
+from ocf_data_sampler.numpy_batch.sun_position import (
+    calculate_azimuth_and_elevation, make_sun_position_numpy_batch
+)
+
+from ocf_datapipes.batch import NumpyBatch, BatchKey
+
+
+@pytest.mark.parametrize("lat", [0, 5, 10, 23.5])
+def test_calculate_azimuth_and_elevation(lat):
+
+    # Pick the day of the summer solstice
+    datetimes = pd.to_datetime(["2024-06-20 12:00"])
+
+    # Calculate sun angles
+    azimuth, elevation = calculate_azimuth_and_elevation(datetimes, lon=0, lat=lat)
+
+    assert len(azimuth)==len(datetimes)
+    assert len(elevation)==len(datetimes)
+
+    # elevation should be close to (90 - (23.5-lat) degrees
+    assert np.abs(elevation - (90-23.5+lat)) < 1
+
+
+def test_calculate_azimuth_and_elevation_random():
+    """Test that the function produces the expected range of azimuths and elevations"""
+
+    # Set seed so we know the test should pass
+    np.random.seed(0)
+
+    # Pick the day of the summer solstice
+    datetimes = pd.to_datetime(["2024-06-20 12:00"])
+
+    # Pick 100 random locations and measure their azimuth and elevations
+    azimuths = []
+    elevations = []
+
+    for _ in range(100):
+
+        lon = np.random.uniform(low=0, high=360)
+        lat = np.random.uniform(low=-90, high=90)
+
+        # Calculate sun angles
+        azimuth, elevation = calculate_azimuth_and_elevation(datetimes, lon=lon, lat=lat)
+
+        azimuths.append(azimuth.item())
+        elevations.append(elevation.item())
+
+    azimuths = np.array(azimuths)
+    elevations = np.array(elevations)
+
+    assert (0<=azimuths).all() and (azimuths<=360).all()
+    assert (-90<=elevations).all() and (elevations<=90).all()
+
+    # Azimuth range is [0, 360]
+    assert azimuths.min() < 30
+    assert azimuths.max() > 330
+
+    # Elevation range is [-90, 90]
+    assert elevations.min() < -70
+    assert elevations.max() > 70
+
+
+def test_make_sun_position_numpy_batch():
+
+    datetimes = pd.date_range("2024-06-20 12:00", "2024-06-20 16:00", freq="30min")
+    lon, lat = 0, 51.5
+
+    batch = make_sun_position_numpy_batch(datetimes, lon, lat, key_preffix="gsp")
+
+    assert BatchKey.gsp_solar_elevation in batch
+    assert BatchKey.gsp_solar_azimuth in batch
+
+    # The solar coords are normalised in the function
+    assert (batch[BatchKey.gsp_solar_elevation]>=0).all()
+    assert (batch[BatchKey.gsp_solar_elevation]<=1).all()
+    assert (batch[BatchKey.gsp_solar_azimuth]>=0).all()
+    assert (batch[BatchKey.gsp_solar_azimuth]<=1).all()

tests/select/test_dropout.py

@@ -0,0 +1,75 @@
+from ocf_data_sampler.select.dropout import draw_dropout_time, apply_dropout_time
+
+import numpy as np
+import pandas as pd
+import xarray as xr
+
+import pytest
+
+
+@pytest.fixture(scope="module")
+def da_sample():
+    """Create dummy data which looks like satellite data"""
+
+    datetimes = pd.date_range("2024-01-01 12:00", "2024-01-01 13:00", freq="5min")
+
+    da_sat = xr.DataArray(
+        np.random.normal(size=(len(datetimes),)),
+        coords=dict(
+            time_utc=(["time_utc"], datetimes),
+        )
+    )
+    return da_sat
+
+
+def test_draw_dropout_time():
+    t0 = pd.Timestamp("2021-01-01 04:00:00")
+
+    dropout_timedeltas = pd.to_timedelta([-30, -60], unit="min")
+    dropout_time = draw_dropout_time(t0, dropout_timedeltas, dropout_frac=1)
+
+    assert isinstance(dropout_time, pd.Timestamp)
+    assert dropout_time-t0 in dropout_timedeltas
+
+
+def test_draw_dropout_time_partial():
+    t0 = pd.Timestamp("2021-01-01 04:00:00")
+
+    dropout_timedeltas = pd.to_timedelta([-30, -60], unit="min")
+
+    dropouts = set()
+
+    # Loop over 1000 to have very high probability of seeing all dropouts
+    # The chances of this failing by chance are approx ((2/3)^100)*3 = 7e-18
+    for _ in range(100):
+        dropouts.add(draw_dropout_time(t0, dropout_timedeltas, dropout_frac=2/3))
+
+    # Check all expected dropouts are present
+    dropouts == {None} | set(t0 + dt for dt in dropout_timedeltas)
+
+
+def test_draw_dropout_time_none():
+    t0 = pd.Timestamp("2021-01-01 04:00:00")
+
+    # No dropout timedeltas
+    dropout_time = draw_dropout_time(t0, dropout_timedeltas=None, dropout_frac=1)
+    assert dropout_time is None
+
+    # Dropout fraction is 0
+    dropout_timedeltas = [pd.Timedelta(-30, "min")]
+    dropout_time = draw_dropout_time(t0, dropout_timedeltas=dropout_timedeltas, dropout_frac=0)
+    assert dropout_time is None
+
+    # No dropout timedeltas and dropout fraction is 0
+    dropout_time = draw_dropout_time(t0, dropout_timedeltas=None, dropout_frac=0)
+    assert dropout_time is None
+
+
+@pytest.mark.parametrize("t0_str", ["12:00", "12:30", "13:00"])
+def test_apply_dropout_time(da_sample, t0_str):
+    dropout_time = pd.Timestamp(f"2024-01-01 {t0_str}")
+
+    da_dropout = apply_dropout_time(da_sample, dropout_time)
+
+    assert da_dropout.sel(time_utc=slice(None, dropout_time)).notnull().all()
+    assert da_dropout.sel(time_utc=slice(dropout_time+pd.Timedelta(5, "min"), None)).isnull().all()

tests/select/test_fill_time_periods.py

@@ -0,0 +1,28 @@
+import pandas as pd
+
+from ocf_data_sampler.select.fill_time_periods import fill_time_periods
+
+def test_fill_time_periods():
+    time_periods = pd.DataFrame(
+        {
+            "start_dt": [
+                "2021-01-01 04:10:00", "2021-01-01 09:00:00",
+                "2021-01-01 09:15:00", "2021-01-01 12:00:00"
+            ],
+            "end_dt": [
+                "2021-01-01 06:00:00", "2021-01-01 09:00:00", 
+                "2021-01-01 09:20:00", "2021-01-01 14:45:00"
+            ],
+        }
+    )
+    freq = pd.Timedelta("30min")
+    filled_time_periods = fill_time_periods(time_periods, freq)
+
+    expected_times = [
+        "04:30", "05:00", "05:30", "06:00", "09:00", "12:00", 
+        "12:30", "13:00", "13:30", "14:00", "14:30"
+    ]
+
+    expected_times = pd.DatetimeIndex([f"2021-01-01 {t}" for t in expected_times])
+
+    pd.testing.assert_index_equal(filled_time_periods, expected_times)