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

tests/load/test_load_sites.py

@@ -0,0 +1,14 @@
+from ocf_data_sampler.load.site import open_site
+import xarray as xr
+
+
+def test_open_site(data_sites):
+    da = open_site(data_sites)
+
+    assert isinstance(da, xr.DataArray)
+    assert da.dims == ("time_utc", "site_id")
+
+    assert "capacity_kwp" in da.coords
+    assert "latitude" in da.coords
+    assert "longitude" in da.coords
+    assert da.shape == (49, 10)

tests/numpy_sample/test_collate.py

@@ -0,0 +1,26 @@
+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 import PVNetUKRegionalDataset
+
+
+def test_pvnet(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]
+
+    batch = stack_np_samples_into_batch([sample1, sample2])
+
+    assert isinstance(batch, dict)
+    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

tests/numpy_sample/test_gsp.py

@@ -0,0 +1,38 @@
+from ocf_data_sampler.load.gsp import open_gsp
+import numpy as np
+
+from ocf_data_sampler.numpy_sample import convert_gsp_to_numpy_sample, GSPSampleKey
+
+def test_convert_gsp_to_numpy_sample(uk_gsp_zarr_path):
+    
+    da = (
+        open_gsp(uk_gsp_zarr_path)
+        .isel(time_utc=slice(0, 10))
+        .sel(gsp_id=1)
+    )
+
+    numpy_sample = convert_gsp_to_numpy_sample(da)
+
+    # Test data structure
+    assert isinstance(numpy_sample, dict), "Should be dict"
+    assert set(numpy_sample.keys()).issubset({
+        GSPSampleKey.gsp,
+        GSPSampleKey.nominal_capacity_mwp,
+        GSPSampleKey.effective_capacity_mwp,
+        GSPSampleKey.time_utc,
+    }), "Unexpected keys"
+
+    # Assert data content and capacity values
+    assert np.array_equal(numpy_sample[GSPSampleKey.gsp], da.values), "GSP values mismatch"
+    assert isinstance(numpy_sample[GSPSampleKey.time_utc], np.ndarray), "Time UTC should be numpy array"
+    assert numpy_sample[GSPSampleKey.time_utc].dtype == float, "Time UTC should be float type"
+    assert numpy_sample[GSPSampleKey.nominal_capacity_mwp] == da.isel(time_utc=0)["nominal_capacity_mwp"].values
+    assert numpy_sample[GSPSampleKey.effective_capacity_mwp] == da.isel(time_utc=0)["effective_capacity_mwp"].values
+
+    # Test with t0_idx
+    t0_idx = 5
+    numpy_sample_with_t0 = convert_gsp_to_numpy_sample(da, t0_idx=t0_idx)
+    assert numpy_sample_with_t0[GSPSampleKey.t0_idx] == t0_idx, "t0_idx not correctly set"
+
+
+

tests/numpy_sample/test_nwp.py

@@ -0,0 +1,52 @@
+import numpy as np
+import pandas as pd
+import xarray as xr
+
+import pytest
+
+from ocf_data_sampler.numpy_sample import convert_nwp_to_numpy_sample, NWPSampleKey
+
+@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_sample(da_nwp_like):
+
+    # Call the function
+    numpy_sample = convert_nwp_to_numpy_sample(da_nwp_like)
+
+    # Assert the output type
+    assert isinstance(numpy_sample, dict)
+
+    # Assert the shape of the numpy sample
+    assert (numpy_sample[NWPSampleKey.nwp] == da_nwp_like.values).all()

tests/numpy_sample/test_satellite.py

@@ -0,0 +1,40 @@
+
+import numpy as np
+import pandas as pd
+import xarray as xr
+
+import pytest
+
+from ocf_data_sampler.numpy_sample import convert_satellite_to_numpy_sample, SatelliteSampleKey
+
+
+@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_sample(da_sat_like):
+
+    # Call the function
+    numpy_sample = convert_satellite_to_numpy_sample(da_sat_like)
+
+    # Assert the output type
+    assert isinstance(numpy_sample, dict)
+
+    # Assert the shape of the numpy sample
+    assert (numpy_sample[SatelliteSampleKey.satellite_actual] == da_sat_like.values).all()

tests/numpy_sample/test_sun_position.py

@@ -0,0 +1,81 @@
+import numpy as np
+import pandas as pd
+import pytest
+
+from ocf_data_sampler.numpy_sample.sun_position import (
+    calculate_azimuth_and_elevation, make_sun_position_numpy_sample
+)
+
+from ocf_data_sampler.numpy_sample import GSPSampleKey
+
+
+@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_sample():
+
+    datetimes = pd.date_range("2024-06-20 12:00", "2024-06-20 16:00", freq="30min")
+    lon, lat = 0, 51.5
+
+    sample = make_sun_position_numpy_sample(datetimes, lon, lat, key_prefix="gsp")
+
+    assert GSPSampleKey.solar_elevation in sample
+    assert GSPSampleKey.solar_azimuth in sample
+
+    # The solar coords are normalised in the function
+    assert (sample[GSPSampleKey.solar_elevation]>=0).all()
+    assert (sample[GSPSampleKey.solar_elevation]<=1).all()
+    assert (sample[GSPSampleKey.solar_azimuth]>=0).all()
+    assert (sample[GSPSampleKey.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)

tests/select/test_find_contiguous_time_periods.py

@@ -0,0 +1,202 @@
+import pandas as pd
+
+from ocf_data_sampler.select.find_contiguous_time_periods import (
+    find_contiguous_t0_periods, find_contiguous_t0_periods_nwp, 
+    intersection_of_multiple_dataframes_of_periods,
+)
+
+
+
+def test_find_contiguous_t0_periods():
+
+    # Create 5-minutely data timestamps
+    freq = pd.Timedelta(5, "min")
+    interval_start = pd.Timedelta(-60, "min")
+    interval_end = pd.Timedelta(15, "min")
+
+    datetimes = (
+        pd.date_range("2023-01-01 12:00", "2023-01-01 17:00", freq=freq)
+        .delete([5, 6, 30])
+    )
+
+    periods = find_contiguous_t0_periods(
+        datetimes=datetimes,
+        interval_start=interval_start,
+        interval_end=interval_end,
+        sample_period_duration=freq,
+    )
+
+    expected_results = pd.DataFrame(
+        {
+            "start_dt": pd.to_datetime(
+                [
+                    "2023-01-01 13:35",
+                    "2023-01-01 15:35",
+                ]
+            ),
+            "end_dt": pd.to_datetime(
+                [
+                    "2023-01-01 14:10",
+                    "2023-01-01 16:45",
+                ]
+            ),
+        },
+    )
+
+    assert periods.equals(expected_results)
+
+
+def test_find_contiguous_t0_periods_nwp():
+
+    # These are the expected results of the test
+    expected_results = [
+        pd.DataFrame(
+            {
+                "start_dt": pd.to_datetime(["2023-01-01 03:00", "2023-01-02 03:00"]),
+                "end_dt": pd.to_datetime(["2023-01-01 21:00", "2023-01-03 06:00"]),
+            },
+        ),
+        pd.DataFrame(
+            {
+                "start_dt": pd.to_datetime(
+                    [
+                        "2023-01-01 05:00",
+                        "2023-01-02 05:00",
+                    ]
+                ),
+                "end_dt": pd.to_datetime(
+                    [
+                        "2023-01-01 21:00",
+                        "2023-01-03 06:00",
+                    ]
+                ),
+            },
+        ),
+        pd.DataFrame(
+            {
+                "start_dt": pd.to_datetime(
+                    [
+                        "2023-01-01 05:00",
+                        "2023-01-02 05:00",
+                        "2023-01-02 14:00",
+                    ]
+                ),
+                "end_dt": pd.to_datetime(
+                    [
+                        "2023-01-01 18:00",
+                        "2023-01-02 09:00",
+                        "2023-01-03 03:00",
+                    ]
+                ),
+            },
+        ),
+        pd.DataFrame(
+            {
+                "start_dt": pd.to_datetime(
+                    [
+                        "2023-01-01 05:00",
+                        "2023-01-01 11:00",
+                        "2023-01-02 05:00",
+                        "2023-01-02 14:00",
+                    ]
+                ),
+                "end_dt": pd.to_datetime(
+                    [
+                        "2023-01-01 06:00",
+                        "2023-01-01 15:00",
+                        "2023-01-02 06:00",
+                        "2023-01-03 00:00",
+                    ]
+                ),
+            },
+        ),
+        pd.DataFrame(
+            {
+                "start_dt": pd.to_datetime(
+                    [
+                        "2023-01-01 06:00",
+                        "2023-01-01 12:00",
+                        "2023-01-02 06:00",
+                        "2023-01-02 15:00",
+                    ]
+                ),
+                "end_dt": pd.to_datetime(
+                    [
+                        "2023-01-01 09:00",
+                        "2023-01-01 18:00",
+                        "2023-01-02 09:00",
+                        "2023-01-03 03:00",
+                    ]
+                ),
+            },
+        ),
+    ]
+
+    # Create 3-hourly init times with a few time stamps missing
+    freq = pd.Timedelta(3, "h")
+
+    init_times = (
+        pd.date_range("2023-01-01 03:00", "2023-01-02 21:00", freq=freq)
+        .delete([1, 4, 5, 6, 7, 9, 10])
+    )
+
+    # Choose some history durations and max stalenesses
+    history_durations_hr = [0, 2, 2, 2, 2]
+    max_stalenesses_hr = [9, 9, 6, 3, 6]
+    max_dropouts_hr = [0, 0, 0, 0, 3]
+
+    for i in range(len(expected_results)):
+        interval_start = pd.Timedelta(-history_durations_hr[i], "h")
+        max_staleness = pd.Timedelta(max_stalenesses_hr[i], "h")
+        max_dropout = pd.Timedelta(max_dropouts_hr[i], "h")
+
+        time_periods = find_contiguous_t0_periods_nwp(
+            init_times=init_times,
+            interval_start=interval_start,
+            max_staleness=max_staleness,
+            max_dropout=max_dropout,
+        )
+
+        # Check if results are as expected
+        assert time_periods.equals(expected_results[i])
+
+
+def test_intersection_of_multiple_dataframes_of_periods():
+    periods_1 = pd.DataFrame(
+        {
+            "start_dt": pd.to_datetime(["2023-01-01 05:00", "2023-01-01 14:10"]),
+            "end_dt": pd.to_datetime(["2023-01-01 13:35", "2023-01-01 18:00"]),
+        },
+    )
+
+    periods_2 = pd.DataFrame(
+        {
+            "start_dt": pd.to_datetime(["2023-01-01 12:00"]),
+            "end_dt": pd.to_datetime(["2023-01-02 00:00"]),
+        },
+    )
+
+    periods_3 = pd.DataFrame(
+        {
+            "start_dt": pd.to_datetime(["2023-01-01 00:00", "2023-01-01 13:00"]),
+            "end_dt": pd.to_datetime(["2023-01-01 12:30", "2023-01-01 23:00"]),
+        },
+    )
+
+    expected_result = pd.DataFrame(
+        {
+            "start_dt": pd.to_datetime(
+                ["2023-01-01 12:00", "2023-01-01 13:00", "2023-01-01 14:10"]
+            ),
+            "end_dt": pd.to_datetime([
+                "2023-01-01 12:30", "2023-01-01 13:35", "2023-01-01 18:00"]
+            ),
+        },
+    )
+
+    overlaping_periods = intersection_of_multiple_dataframes_of_periods(
+        [periods_1, periods_2, periods_3]
+    )   
+
+    # Check if results are as expected
+    assert overlaping_periods.equals(expected_result)

tests/select/test_location.py

@@ -0,0 +1,67 @@
+from ocf_data_sampler.select.location import Location
+import pytest
+
+
+def test_make_valid_location_object_with_default_coordinate_system():
+    x, y = -1000.5, 50000
+    location = Location(x=x, y=y)
+    assert location.x == x, "location.x value not set correctly"
+    assert location.y == y, "location.x value not set correctly"
+    assert (
+        location.coordinate_system == "osgb"
+    ), "location.coordinate_system value not set correctly"
+
+
+def test_make_valid_location_object_with_osgb_coordinate_system():
+    x, y, coordinate_system = 1.2, 22.9, "osgb"
+    location = Location(x=x, y=y, coordinate_system=coordinate_system)
+    assert location.x == x, "location.x value not set correctly"
+    assert location.y == y, "location.x value not set correctly"
+    assert (
+        location.coordinate_system == coordinate_system
+    ), "location.coordinate_system value not set correctly"
+
+
+def test_make_valid_location_object_with_lon_lat_coordinate_system():
+    x, y, coordinate_system = 1.2, 1.2, "lon_lat"
+    location = Location(x=x, y=y, coordinate_system=coordinate_system)
+    assert location.x == x, "location.x value not set correctly"
+    assert location.y == y, "location.x value not set correctly"
+    assert (
+        location.coordinate_system == coordinate_system
+    ), "location.coordinate_system value not set correctly"
+
+
+def test_make_invalid_location_object_with_invalid_osgb_x():
+    x, y, coordinate_system = 10000000, 1.2, "osgb"
+    with pytest.raises(ValueError) as err:
+        _ = Location(x=x, y=y, coordinate_system=coordinate_system)
+    assert err.typename == "ValidationError"
+
+
+def test_make_invalid_location_object_with_invalid_osgb_y():
+    x, y, coordinate_system = 2.5, 10000000, "osgb"
+    with pytest.raises(ValueError) as err:
+        _ = Location(x=x, y=y, coordinate_system=coordinate_system)
+    assert err.typename == "ValidationError"
+
+
+def test_make_invalid_location_object_with_invalid_lon_lat_x():
+    x, y, coordinate_system = 200, 1.2, "lon_lat"
+    with pytest.raises(ValueError) as err:
+        _ = Location(x=x, y=y, coordinate_system=coordinate_system)
+    assert err.typename == "ValidationError"
+
+
+def test_make_invalid_location_object_with_invalid_lon_lat_y():
+    x, y, coordinate_system = 2.5, -200, "lon_lat"
+    with pytest.raises(ValueError) as err:
+        _ = Location(x=x, y=y, coordinate_system=coordinate_system)
+    assert err.typename == "ValidationError"
+
+
+def test_make_invalid_location_object_with_invalid_coordinate_system():
+    x, y, coordinate_system = 2.5, 1000, "abcd"
+    with pytest.raises(ValueError) as err:
+        _ = Location(x=x, y=y, coordinate_system=coordinate_system)
+    assert err.typename == "ValidationError"