Merge pull request #91 from openclimatefix/issue/forecast-horizon

Issue/forecast horizon

Author: peterdudfield

nowcasting_datamodel/read/read.py

@@ -8,7 +8,7 @@
 from datetime import datetime, timedelta, timezone
 from typing import List, Optional
 
-from sqlalchemy import desc
+from sqlalchemy import desc, text
 from sqlalchemy.orm import contains_eager, joinedload
 from sqlalchemy.orm.session import Session
 
@@ -121,7 +121,6 @@ def get_latest_forecast(
     gsp_id: Optional[int] = None,
     historic: bool = False,
     start_target_time: Optional[datetime] = None,
-    forecast_horizon_minutes: Optional[int] = None,
 ) -> ForecastSQL:
     """
     Read forecasts
@@ -149,7 +148,6 @@ def get_latest_forecast(
         start_target_time=start_target_time,
         historic=historic,
         gsp_ids=gsp_ids,
-        forecast_horizon_minutes=forecast_horizon_minutes,
     )
 
     if forecasts is None:
@@ -217,7 +215,6 @@ def get_all_gsp_ids_latest_forecast(
     start_target_time: Optional[datetime] = None,
     preload_children: Optional[bool] = False,
     historic: bool = False,
-    forecast_horizon_minutes: Optional[int] = None,
 ) -> List[ForecastSQL]:
     """
     Read forecasts
@@ -228,9 +225,6 @@ def get_all_gsp_ids_latest_forecast(
         Filter: forecast values target time should be larger than this datetime
     :param preload_children: Option to preload children. This is a speed up, if we need them.
     :param historic: Option to load historic values or not
-    :param forecast_horizon_minutes: Optional filter on forecast horizon. For example
-        forecast_horizon_minutes=120, means load the forecast than was made 2 hours before the
-        target time. Note this only works for non-historic data.
 
     return: List of forecasts objects from database
     """
@@ -244,7 +238,6 @@ def get_all_gsp_ids_latest_forecast(
         preload_children=preload_children,
         historic=historic,
         gsp_ids=list(range(0, N_GSP + 1)),
-        forecast_horizon_minutes=forecast_horizon_minutes,
     )
 
 
@@ -254,7 +247,6 @@ def get_latest_forecast_for_gsps(
     start_target_time: Optional[datetime] = None,
     preload_children: Optional[bool] = False,
     historic: bool = False,
-    forecast_horizon_minutes: Optional[int] = None,
     gsp_ids: List[int] = None,
 ):
     """
@@ -267,19 +259,9 @@ def get_latest_forecast_for_gsps(
     :param preload_children: Option to preload children. This is a speed up, if we need them.
     :param historic: Option to load historic values or not
     :param gsp_ids: Option to filter on gsps. If None, then only the lastest forecast is loaded.
-    :param forecast_horizon_minutes: Optional filter on forecast horizon. For example
-        forecast_horizon_minutes=120, means load the forecast than was made 2 hours before the
-        target time. Note this only works for non-historic data.
 
-    return: List of forecasts objects from database
+    :return: List of forecasts objects from database
 
-    :param session:
-    :param start_created_utc:
-    :param start_target_time:
-    :param preload_children:
-    :param historic:
-    :param gsp_ids:
-    :return:
     """
     order_by_cols = []
 
@@ -307,17 +289,6 @@ def get_latest_forecast_for_gsps(
             query=query, start_target_time=start_target_time, historic=historic
         )
 
-    from sqlalchemy import text
-
-    if forecast_horizon_minutes is not None:
-        assert historic is False, Exception(
-            "Loading a forecast horizon only works on non latest data."
-        )
-        query = query.join(ForecastValueSQL).filter(
-            ForecastValueSQL.target_time - ForecastValueSQL.created_utc
-            >= text(f"interval '{forecast_horizon_minutes} minute'")
-        )
-
     query = query.join(LocationSQL)
 
     # option to preload values, makes querying quicker
@@ -374,6 +345,7 @@ def get_forecast_values(
     session: Session,
     gsp_id: Optional[int] = None,
     start_datetime: Optional[datetime] = None,
+    forecast_horizon_minutes: Optional[int] = None,
     only_return_latest: Optional[bool] = False,
 ) -> List[ForecastValueSQL]:
     """
@@ -386,6 +358,9 @@ def get_forecast_values(
         If None is given then all are returned.
     :param only_return_latest: Optional to only return the latest forecast, not all of them.
         Default is False
+    :param forecast_horizon_minutes: Optional filter on forecast horizon. For example
+        forecast_horizon_minutes=120, means load the forecast than was made 2 hours before the
+        target time. Note this only works for non-historic data.
 
     return: List of forecasts values objects from database
 
@@ -405,6 +380,14 @@ def get_forecast_values(
         query = query.filter(ForecastValueSQL.created_utc >= created_utc_filter)
         query = query.filter(ForecastSQL.created_utc >= created_utc_filter)
 
+    if forecast_horizon_minutes is not None:
+
+        # this seems to only work for postgres
+        query = query.filter(
+            ForecastValueSQL.target_time - ForecastValueSQL.created_utc
+            >= text(f"interval '{forecast_horizon_minutes} minute'")
+        )
+
     # filter on gsp_id
     if gsp_id is not None:
         query = query.join(ForecastSQL)

tests/read/test_read_forecast.py

@@ -0,0 +1,30 @@
+import logging
+from datetime import datetime, timezone
+
+from nowcasting_datamodel.fake import make_fake_forecast
+from nowcasting_datamodel.read.read import get_forecast_values
+
+logger = logging.getLogger(__name__)
+
+
+def test_get_latest_forecast_created_utc_gsp(db_session):
+    t0_datetime_utc = datetime(2022, 1, 1, 12, tzinfo=timezone.utc)
+
+    f1 = make_fake_forecast(gsp_id=1, session=db_session, t0_datetime_utc=t0_datetime_utc)
+    f2 = make_fake_forecast(gsp_id=1, session=db_session, t0_datetime_utc=t0_datetime_utc)
+
+    f1.forecast_values[0].created_utc = datetime(2022, 1, 1, 10, tzinfo=timezone.utc)
+    f1.forecast_values[1].created_utc = datetime(2022, 1, 1, 10, tzinfo=timezone.utc)
+
+    f2.forecast_values[0].created_utc = datetime(2022, 1, 1, 12, tzinfo=timezone.utc)
+    f2.forecast_values[1].created_utc = datetime(2022, 1, 1, 12, tzinfo=timezone.utc)
+
+    db_session.add_all([f1, f2])
+    db_session.commit()
+
+    forecast_values = get_forecast_values(
+        session=db_session, forecast_horizon_minutes=120, gsp_id=1
+    )
+    assert len(forecast_values) == 2
+    assert forecast_values[0].created_utc == datetime(2022, 1, 1, 10, tzinfo=timezone.utc)
+    assert forecast_values[0].target_time == datetime(2022, 1, 1, 12, tzinfo=timezone.utc)