refactor sun position calculation

ocf_data_sampler/numpy_batch/__init__.py

@@ -3,5 +3,5 @@
 from .gsp import convert_gsp_to_numpy_batch
 from .nwp import convert_nwp_to_numpy_batch
 from .satellite import convert_satellite_to_numpy_batch
-from .add_sun_position import add_sun_position_to_numpy_batch
+from .sun_position import make_sun_position_numpy_batch
 

ocf_data_sampler/numpy_batch/sun_position.py

@@ -0,0 +1,66 @@
+
+import pvlib
+import numpy as np
+import pandas as pd
+from ocf_datapipes.batch import BatchKey, NumpyBatch
+
+
+def calculate_azimuth_and_elevation(
+    datetimes: pd.DatetimeIndex, 
+    lon: float, 
+    lat: float
+) -> tuple[np.ndarray, np.ndarray]:
+    """Calculate the solar coordinates for multiple datetimes at a single location
+    
+    Args:
+        datetimes: The datetimes to calculate for
+        lon: The longitude
+        lat: The latitude
+
+    Returns:
+        np.ndarray: The azimuth of the datetimes in degrees
+        np.ndarray: The elevation of the datetimes in degrees
+    """
+
+    solpos = pvlib.solarposition.get_solarposition(
+        time=datetimes,
+        longitude=lon,
+        latitude=lat,
+        method='nrel_numpy'
+    )
+    azimuth = solpos["azimuth"].values
+    elevation = solpos["elevation"].values
+    return azimuth, elevation
+
+
+def make_sun_position_numpy_batch(
+        datetimes: pd.DatetimeIndex, 
+        lon: float, 
+        lat: float, 
+        key_preffix: str = "gsp"
+) -> NumpyBatch:
+    """Creates NumpyBatch with standardized solar coordinates
+
+    Args:
+        datetimes: The datetimes to calculate solar angles for
+        lon: The longitude
+        lat: The latitude
+    """
+
+    azimuth, elevation = calculate_azimuth_and_elevation(datetimes, lon, lat)
+
+    # Normalise
+
+    # Azimuth is in range [0, 360] degrees
+    azimuth = azimuth / 360
+
+    # Elevation is in range [-90, 90] degrees
+    elevation = elevation / 180 + 0.5
+
+    # Make NumpyBatch
+    sun_numpy_batch: NumpyBatch = {
+        BatchKey[key_preffix + "_solar_azimuth"]: azimuth,
+        BatchKey[key_preffix + "_solar_elevation"]: elevation,
+    }
+
+    return sun_numpy_batch

ocf_data_sampler/torch_datasets/pvnet_uk_regional.py

@@ -23,15 +23,16 @@
     convert_gsp_to_numpy_batch,
     convert_nwp_to_numpy_batch,
     convert_satellite_to_numpy_batch,
-    add_sun_position_to_numpy_batch,
+    make_sun_position_numpy_batch,
 )
 
 
 from ocf_datapipes.config.model import Configuration
 from ocf_datapipes.config.load import load_yaml_configuration
+from ocf_datapipes.batch import BatchKey, NumpyBatch
 
 from ocf_datapipes.utils.location import Location
-from ocf_datapipes.batch import BatchKey, NumpyBatch
+from ocf_datapipes.utils.geospatial import osgb_to_lon_lat
 
 from ocf_datapipes.utils.consts import (
     NWP_MEANS,
@@ -40,9 +41,8 @@
     RSS_STD,
 )
 
-from ocf_datapipes.training.common import (
-    is_config_and_path_valid, concat_xr_time_utc, normalize_gsp,
-)
+from ocf_datapipes.training.common import concat_xr_time_utc, normalize_gsp
+
 
 
 xr.set_options(keep_attrs=True)
@@ -347,7 +347,12 @@ def merge_dicts(list_of_dicts: list[dict]) -> dict:
     return combined_dict
 
 
-def process_and_combine_datasets(dataset_dict: dict, config: Configuration) -> NumpyBatch:
+def process_and_combine_datasets(
+        dataset_dict: dict, 
+        config: Configuration,
+        t0: pd.Timedelta,
+        location: Location,
+    ) -> NumpyBatch:
     """Normalize and convert data to numpy arrays"""    
 
     numpy_modalities = []
@@ -373,7 +378,6 @@ def process_and_combine_datasets(dataset_dict: dict, config: Configuration) -> N
         # Convert to NumpyBatch
         numpy_modalities.append(convert_satellite_to_numpy_batch(da_sat))
 
-
     # GSP always assumed to be in data
     gsp_config = config.input_data.gsp
     da_gsp = concat_xr_time_utc([dataset_dict["gsp"], dataset_dict["gsp_future"]])
@@ -386,12 +390,20 @@ def process_and_combine_datasets(dataset_dict: dict, config: Configuration) -> N
         )
     )
 
+    # Make sun coords NumpyBatch
+    datetimes = pd.date_range(
+        t0-minutes(gsp_config.history_minutes),
+        t0+minutes(gsp_config.forecast_minutes),
+        freq=minutes(gsp_config.time_resolution_minutes),
+    )
+
+    lon, lat = osgb_to_lon_lat(location.x, location.y)
+
+    numpy_modalities.append(make_sun_position_numpy_batch(datetimes, lon, lat))
+
     # Combine all the modalities
     combined_sample = merge_dicts(numpy_modalities)
 
-    # Add sun coords
-    combined_sample = add_sun_position_to_numpy_batch(combined_sample, modality_name="gsp")
-
     return combined_sample
 
 
@@ -491,7 +503,7 @@ def _get_sample(self, t0: pd.Timestamp, location: Location) -> NumpyBatch:
         sample_dict = slice_datasets_by_time(sample_dict, t0, self.config)
         sample_dict = compute(sample_dict)
 
-        sample = process_and_combine_datasets(sample_dict, self.config)
+        sample = process_and_combine_datasets(sample_dict, self.config, t0, location)
 
         return sample
 

tests/numpy_batch/test_add_sun_position.py → tests/numpy_batch/test_sun_position.py

@@ -2,21 +2,21 @@
 import pandas as pd
 import pytest
 
-from ocf_data_sampler.numpy_batch.add_sun_position import (
-    get_azimuth_and_elevation, add_sun_position_to_numpy_batch
+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_get_azimuth_and_elevation_solstice(lat):
+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 = get_azimuth_and_elevation(datetimes, lon=0, lat=lat)
+    azimuth, elevation = calculate_azimuth_and_elevation(datetimes, lon=0, lat=lat)
 
     assert len(azimuth)==len(datetimes)
     assert len(elevation)==len(datetimes)
@@ -25,7 +25,7 @@ def test_get_azimuth_and_elevation_solstice(lat):
     assert np.abs(elevation - (90-23.5+lat)) < 1
 
 
-def test_get_azimuth_and_elevation_random():
+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
@@ -44,7 +44,7 @@ def test_get_azimuth_and_elevation_random():
         lat = np.random.uniform(low=-90, high=90)
 
         # Calculate sun angles
-        azimuth, elevation = get_azimuth_and_elevation(datetimes, lon=lon, lat=lat)
+        azimuth, elevation = calculate_azimuth_and_elevation(datetimes, lon=lon, lat=lat)
 
         azimuths.append(azimuth.item())
         elevations.append(elevation.item())
@@ -64,21 +64,16 @@ def test_get_azimuth_and_elevation_random():
     assert elevations.max() > 70
 
 
-def test_add_sun_position_to_numpy_batch():
-    datetimes = pd.to_datetime(["2024-06-20 12:00"]).values.astype(float) / 1000
+def test_make_sun_position_numpy_batch():
 
-    batch: NumpyBatch = {
-        BatchKey.gsp_time_utc: datetimes,
-        BatchKey.gsp_x_osgb: 0,
-        BatchKey.gsp_y_osgb: 0,
-    }
+    datetimes = pd.date_range("2024-06-20 12:00", "2024-06-20 16:00", freq="30min")
+    lon, lat = 0, 51.5
 
-    batch = add_sun_position_to_numpy_batch(batch, modality_name="gsp")
+    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()