add code to keep dusk and dawn, add extra test

Author: peterdudfield

nowcasting_dataset/time.py

@@ -2,6 +2,7 @@
 import logging
 import random
 import warnings
+from datetime import timedelta
 from typing import Dict, Iterable, List, Optional, Tuple
 
 import numpy as np
@@ -40,6 +41,31 @@ def select_daylight_datetimes(
     across all locations.
 
     """
+
+    if keep_dawn_dusk_hours > 0:
+        # get dawn datetimes by
+        # - adding X hours the original datetimes,
+        # - running through select_daylight_datetimes
+        # - minusing X hours from the results
+        # similar for dusk
+        temp_keep_dawn_dusk_hours = 0
+        dawn = select_daylight_datetimes(
+            datetimes=datetimes + timedelta(hours=keep_dawn_dusk_hours),
+            locations=locations,
+            ghi_threshold=ghi_threshold,
+            keep_dawn_dusk_hours=temp_keep_dawn_dusk_hours,
+        ) - timedelta(hours=keep_dawn_dusk_hours)
+
+        dusk = select_daylight_datetimes(
+            datetimes=datetimes - timedelta(hours=keep_dawn_dusk_hours),
+            locations=locations,
+            ghi_threshold=ghi_threshold,
+            keep_dawn_dusk_hours=temp_keep_dawn_dusk_hours,
+        ) + timedelta(hours=keep_dawn_dusk_hours)
+
+        # take union of dusk and dawn
+        return dusk.union(dawn)
+
     ghi_for_all_locations = []
     for x, y in locations:
         lat, lon = geospatial.osgb_to_lat_lon(x, y)

tests/test_time.py

@@ -27,6 +27,31 @@ def test_select_daylight_datetimes_dusk():
     np.testing.assert_array_equal(daylight_datetimes, correct_daylight_datetimes)
 
 
+def test_select_daylight_datetimes_dusk_non_gridded_date():
+    """Test day light filter works with keeping dusk and dawn extra hours
+
+    This is for non-gridded date.
+    This is to check that no extra datetimes are being returned
+    """
+    datetimes = pd.DatetimeIndex(
+        [
+            "2020-01-01 07:00",
+            "2020-01-01 16:00",
+            "2020-01-01 17:00",
+            "2020-01-01 18:00",
+            "2020-01-01 19:00",
+        ]
+    )
+    correct_daylight_datetimes = pd.DatetimeIndex(
+        ["2020-01-01 07:00", "2020-01-01 16:00", "2020-01-01 17:00", "2020-01-01 18:00"]
+    )
+    locations = [(0, 0), (20_000, 20_000)]
+    daylight_datetimes = nd_time.select_daylight_datetimes(
+        datetimes=datetimes, locations=locations, keep_dawn_dusk_hours=2
+    )
+    np.testing.assert_array_equal(daylight_datetimes, correct_daylight_datetimes)
+
+
 @pytest.mark.parametrize("min_seq_length", [2, 3, 12])
 def test_get_contiguous_time_periods_1_with_1_chunk(min_seq_length):
     """Test getting continuous chunks of data with 1 chunk of data"""