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

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,22 @@
+from ocf_data_sampler.load.gsp import open_gsp
+
+from ocf_data_sampler.numpy_batch import convert_gsp_to_numpy_batch
+from ocf_data_sampler.numpy_batch.gsp import GSPBatchKey
+
+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[GSPBatchKey.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_data_sampler.numpy_batch.nwp 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_data_sampler.numpy_batch.satellite import SatelliteBatchKey
+
+
+@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[SatelliteBatchKey.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_data_sampler.numpy_batch.gsp import GSPBatchKey
+
+
+@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_prefix="gsp")
+
+    assert GSPBatchKey.gsp_solar_elevation in batch
+    assert GSPBatchKey.gsp_solar_azimuth in batch
+
+    # The solar coords are normalised in the function
+    assert (batch[GSPBatchKey.gsp_solar_elevation]>=0).all()
+    assert (batch[GSPBatchKey.gsp_solar_elevation]<=1).all()
+    assert (batch[GSPBatchKey.gsp_solar_azimuth]>=0).all()
+    assert (batch[GSPBatchKey.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)

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")
+    history_duration = pd.Timedelta(60, "min")
+    forecast_duration = 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,
+        history_duration=history_duration,
+        forecast_duration=forecast_duration,
+        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")
+
+    datetimes = (
+        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)):
+        history_duration = 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(
+            datetimes=datetimes,
+            history_duration=history_duration,
+            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"