Merge pull request #16 from nsidc/ardwa-4

ARDWA-4

Author: scott-lewis-nsidc

pyproject.toml

@@ -25,6 +25,9 @@ dependencies = [
   "loguru",
   "tqdm",
   "fastapi ~=0.111.0",
+  "matplotlib",
+  "numpy",
+  "pandas",
   "pydantic ~=2.0",
   "pydantic-settings",
   "sqlalchemy ~=2.0",
@@ -61,8 +64,6 @@ docs = [
 ui = [
   "jupyterlab",
   "leafmap",
-  "matplotlib",
-  "pandas",
 ]
 
 [project.urls]

src/aross_stations_db/api/v1/output.py

@@ -1,4 +1,5 @@
 import datetime as dt
+import io
 from typing import Annotated
 
 from annotated_types import Ge, Le
@@ -8,8 +9,12 @@
     Point,
 )
 from geojson_pydantic.types import Position2D
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
 from pydantic import BaseModel
 from sqlalchemy import Row
+from sqlalchemy.orm.query import RowReturningQuery
 
 from aross_stations_db.db.tables import Station
 
@@ -64,6 +69,87 @@ def timeseries_query_results_to_json(
     ]
 
 
+def timeseries_query_results_to_bar_plot_buffer(
+    query: RowReturningQuery[tuple[dt.datetime, int]],
+    start: dt.date,
+    end: dt.date,
+    format: str = 'png'
+) -> io.BytesIO:
+    data = pd.read_sql(query.statement, query.session.connection())
+    data.set_index('month', inplace=True)
+
+    start_str = start.strftime("%Y-%m")
+    end_str = end.strftime("%Y-%m")
+
+    title_parts = [
+        f"Monthly Rain-on-Snow Events",
+        f"[{start_str} - {end_str}]"
+    ]
+    title = "\n".join(title_parts)
+
+    if len(data) == 0:
+        data.loc[start_str] = 0
+        data.loc[end_str] = 0
+
+    
+    data.index = pd.to_datetime(data.index)
+    data.index = data.index.strftime("%Y-%m")
+    data = add_missing_plot_months(data)
+
+    plot = data.plot(
+        kind="bar",
+        title=title,
+        ylabel="Event Count",
+        xlabel="Month",
+        rot=45,
+        legend=False,
+    )
+
+    ticks = plot.get_xticklines()
+    labels = plot.get_xticklabels()
+
+    # Evenly space out the tick labels to avoid overcrowding
+    indexes = np.linspace(0, len(labels)-1, num=12, dtype=int)
+    for i, l in enumerate(labels):
+        if i not in indexes:
+            l.set_visible(False)
+        else:
+            l.set_horizontalalignment('right')
+            l.set_rotation_mode('anchor')
+            ticks[i*2].set_markersize(6)
+    
+    plt.tight_layout()
+
+    # Create the buffer and return it so it can be sent to the requester
+    buffer = io.BytesIO()
+    plt.savefig(buffer, format="png")
+    plt.close()
+
+    buffer.seek(0)
+
+    return buffer
+
+
+# If there are any months missing in the dataframe, add a "count 0" entry for them
+def add_missing_plot_months(df: pd.DataFrame) -> pd.DataFrame:
+    start = df.index[0]
+    end = df.index[-1]
+
+    syear, smonth = map(int, start.split('-'))
+    eyear, emonth = map(int, end.split('-'))
+
+    while [syear, smonth] != [eyear, emonth]:
+        smonth += 1
+        if smonth > 12:
+            smonth = 1
+            syear += 1
+        key = f"{syear:04}-{smonth:02}"
+        if key not in df.index:
+            df.loc[key] = [0]
+        
+    return df.sort_index()
+
+
 class ClimatologyJsonElement(BaseModel):
     month: Annotated[int, Ge(1), Le(12)]
     event_count: int

src/aross_stations_db/api/v1/timeseries.py

@@ -1,19 +1,23 @@
 import datetime as dt
 from typing import Annotated
 
-from fastapi import APIRouter, Depends, Query
+from fastapi import APIRouter, Depends, Form, Query
+from fastapi.responses import StreamingResponse
 from geoalchemy2 import WKTElement
+import pandas as pd
 from sqlalchemy.orm import Session
 
 from aross_stations_db.api.dependencies import get_db_session
 from aross_stations_db.api.v1.output import (
     TimeseriesJsonElement,
     timeseries_query_results_to_json,
+    timeseries_query_results_to_bar_plot_buffer,
 )
 from aross_stations_db.db.query import timeseries_query
 
-router = APIRouter()
+from loguru import logger
 
+router = APIRouter()
 
 @router.get("/monthly")
 def get_monthly_timeseries(
@@ -22,8 +26,58 @@ def get_monthly_timeseries(
     start: Annotated[dt.datetime, Query(description="ISO-format timestamp")],
     end: Annotated[dt.datetime, Query(description="ISO-format timestamp")],
     polygon: Annotated[str | None, WKTElement, Query(description="WKT shape")] = None,
+    stations: Annotated[list[str], Query(description="List of station identifiers")] = [],
 ) -> list[TimeseriesJsonElement]:
     """Get a monthly timeseries of events matching query parameters."""
-    query = timeseries_query(db=db, start=start, end=end, polygon=polygon)
+    logger.debug(f"STATIONS: {stations}")
+    query = timeseries_query(db=db, start=start, end=end, polygon=polygon, stations=stations)
 
     return timeseries_query_results_to_json(query.all())
+
+
+@router.post("/monthly")
+def post_monthly_timeseries(
+    db: Annotated[Session, Depends(get_db_session)],
+    *,
+    start: Annotated[dt.datetime, Form(description="ISO-format timestamp")],
+    end: Annotated[dt.datetime, Form(description="ISO-format timestamp")],
+    polygon: Annotated[str | None, WKTElement, Form(description="WKT shape")] = None,
+    stations: Annotated[list[str], Form(description="List of station identifiers")] = [],
+) -> list[TimeseriesJsonElement]:
+    """Get a monthly timeseries of events matching query parameters."""
+    logger.debug(f"STATIONS: {stations}")
+    query = timeseries_query(db=db, start=start, end=end, polygon=polygon, stations=stations)
+
+    return timeseries_query_results_to_json(query.all())
+
+
+@router.get("/monthly/png")
+def get_monthly_timeseries_png(
+    db: Annotated[Session, Depends(get_db_session)],
+    *,
+    start: Annotated[dt.datetime, Query(description="ISO-format timestamp")],
+    end: Annotated[dt.datetime, Query(description="ISO-format timestamp")],
+    polygon: Annotated[str | None, WKTElement, Query(description="WKT shape")] = None,
+    stations: Annotated[list[str], Query(Description="List of station identifiers")] = [],
+) -> StreamingResponse:
+    """Get a monthly timeseries image plot of events matching query parameters."""
+    query = timeseries_query(db=db, start=start, end=end, polygon=polygon, stations=stations)
+    buffer = timeseries_query_results_to_bar_plot_buffer(query, start, end, 'png')
+
+    return StreamingResponse(buffer, media_type="image/png")
+
+
+@router.post("/monthly/png")
+def post_monthly_timeseries_png(
+    db: Annotated[Session, Depends(get_db_session)],
+    *,
+    start: Annotated[dt.datetime, Form(description="ISO-format timestamp")],
+    end: Annotated[dt.datetime, Form(description="ISO-format timestamp")],
+    polygon: Annotated[str | None, WKTElement, Form(description="WKT shape")] = None,
+    stations: Annotated[list[str], Form(Description="List of station identifiers")] = [],
+) -> StreamingResponse:
+    """Get a monthly timeseries image plot of events matching query parameters."""
+    query = timeseries_query(db=db, start=start, end=end, polygon=polygon, stations=stations)
+    buffer = timeseries_query_results_to_bar_plot_buffer(query, start, end, 'png')
+    
+    return StreamingResponse(buffer, media_type="image/png")

src/aross_stations_db/db/query.py

@@ -46,7 +46,6 @@ def station_data_query(
     end: dt.datetime | None,
     stations: list[str] = [],
 ) -> RowReturningQuery[tuple[str, dt.datetime, dt.datetime, bool | None, int, int, int, int]]:
-    print(type(stations))
     query = db.query(
         Event.station_id,
         Event.time_start,
@@ -72,6 +71,7 @@ def timeseries_query(
     start: dt.datetime,
     end: dt.datetime,
     polygon: str | None = None,
+    stations: list[str] | None = None,
 ) -> RowReturningQuery[tuple[dt.datetime, int]]:
     query = db.query(
         func.date_trunc("month", Event.time_start, type_=DateTime).label("month"),
@@ -87,6 +87,9 @@ def timeseries_query(
             )
         )
 
+    if stations:
+        query = query.filter(Event.station_id.in_(stations))
+    
     return query.group_by("month").order_by("month")