add function to sum up gsp values (#203)

* add function to sum up gsp values

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* self PR comment

* add end dateitme into test

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: peterdudfield

nowcasting_datamodel/read/read_gsp.py

@@ -3,10 +3,10 @@
 from datetime import datetime, timezone
 from typing import List, Optional, Union
 
-from sqlalchemy import desc
+from sqlalchemy import desc, func
 from sqlalchemy.orm import Session, contains_eager, joinedload
 
-from nowcasting_datamodel.models import GSPYieldSQL, LocationSQL
+from nowcasting_datamodel.models import GSPYield, GSPYieldSQL, LocationSQL
 
 logger = logging.getLogger(__name__)
 
@@ -217,3 +217,67 @@ def get_gsp_yield_by_location(
             gsp_yield.datetime_utc = gsp_yield.datetime_utc.replace(tzinfo=timezone.utc)
 
     return locations
+
+
+def get_gsp_yield_sum(
+    session: Session,
+    gsp_ids: List[int],
+    start_datetime_utc: datetime,
+    regime: Optional[str] = None,
+    end_datetime_utc: Optional[datetime] = None,
+) -> List[GSPYield]:
+    """
+    Get the sum of gsp yield values.
+
+    :param session: sqlalchemy sessions
+    :param gsp_ids: list of gsp ids that we filter on
+    :param start_datetime_utc: filter values on this start datetime
+    :param regime: filter query on this regim. Can be "in-day" or "day-after"
+    :param end_datetime_utc: optional end datetime filter
+
+    :return: list of GSPYield objects
+    """
+
+    logger.info(f"Getting gsp yield sum for {len(gsp_ids)} gsp systems")
+
+    if regime is None:
+        logger.debug("No regime given, defaulting to 'in-day'")
+        regime = "in-day"
+
+    # start main query
+    query = session.query(
+        GSPYieldSQL.datetime_utc,
+        func.sum(GSPYieldSQL.solar_generation_kw).label("solar_generation_kw"),
+    )
+
+    # join with location table
+    query = query.join(LocationSQL)
+
+    # select only the gsp systems we want
+    query = query.where(LocationSQL.gsp_id.in_(gsp_ids))
+
+    # filter on regime
+    query = query.where(GSPYieldSQL.regime == regime)
+
+    # filter on datetime
+    query = query.where(GSPYieldSQL.datetime_utc >= start_datetime_utc)
+    if end_datetime_utc is not None:
+        query = query.where(GSPYieldSQL.datetime_utc <= end_datetime_utc)
+
+    # group and order by datetime
+    query = query.group_by(GSPYieldSQL.datetime_utc)
+    query = query.order_by(GSPYieldSQL.datetime_utc)
+
+    results = query.all()
+
+    # format results
+    results = [
+        GSPYield(
+            datetime_utc=result.datetime_utc,
+            solar_generation_kw=result.solar_generation_kw,
+            regime=regime,
+        )
+        for result in results
+    ]
+
+    return results

tests/read/test_read_gsp.py

@@ -8,6 +8,7 @@
     get_gsp_yield,
     get_gsp_yield_by_location,
     get_latest_gsp_yield,
+    get_gsp_yield_sum,
 )
 
 logger = logging.getLogger(__name__)
@@ -218,3 +219,25 @@ def test_get_gsp_yield_by_location_extra(db_session):
     )
 
     assert len(locations_with_gsp_yields) == 0
+
+
+def test_get_gsp_yield_sum(db_session):
+    _ = setup_gsp_yields(db_session)
+
+    gsp_yields = get_gsp_yield_sum(
+        session=db_session,
+        gsp_ids=[1, 2],
+        start_datetime_utc=datetime(2022, 1, 1),
+        end_datetime_utc=datetime(2022, 1, 2),
+    )
+
+    # read database
+    assert len(gsp_yields) == 2
+
+    assert gsp_yields[0].datetime_utc == datetime(2022, 1, 1, tzinfo=timezone.utc)
+    assert gsp_yields[1].datetime_utc == datetime(2022, 1, 2, tzinfo=timezone.utc)
+
+    assert (
+        gsp_yields[0].solar_generation_kw == 4
+    )  # 2+2, see hard coded values in 'setup_gsp_yields'
+    assert gsp_yields[1].solar_generation_kw == 1  # 1, see hard coded values in 'setup_gsp_yields'