Optimize JSON responses (#436)

Optimize the backend that generates the JSON from the Zarr and Parquet
files. All of the list comprehensions in Python that were generating the
GeoJSON responses were really slowing down our responses. By changing
the response format and relying on `pandas.DataFrame.to_json` so that we
can eliminate those loops, we get significantly faster responses (0.3s
instead of 20s).

src/unified_graphics/diag.py

@@ -1,8 +1,7 @@
 import os
 from collections import namedtuple
-from dataclasses import dataclass
 from enum import Enum
-from typing import Generator, List, Union
+from typing import Union
 from urllib.parse import urlparse
 
 import numpy as np
@@ -22,60 +21,13 @@ class MinimLoop(Enum):
     ANALYSIS = "anl"
 
 
-class ValueType(Enum):
-    OBSERVATION = "observation"
-    FORECAST = "forecast"
-
-
 class Variable(Enum):
     MOISTURE = "q"
     PRESSURE = "ps"
     TEMPERATURE = "t"
     WIND = "uv"
 
 
-class VariableType(Enum):
-    SCALAR = "scalar"
-    VECTOR = "vector"
-
-
-Coordinate = namedtuple("Coordinate", "longitude latitude")
-Vector = namedtuple("Vector", "u v")
-
-
-@dataclass
-class Observation:
-    variable: str
-    variable_type: VariableType
-    loop: MinimLoop
-    adjusted: Union[float, Vector]
-    unadjusted: Union[float, Vector]
-    observed: Union[float, Vector]
-    position: Coordinate
-
-    def to_geojson(self):
-        properties = {
-            "type": self.variable_type.value,
-            "variable": self.variable,
-            "loop": self.loop.value,
-        }
-
-        if isinstance(self.adjusted, float):
-            properties["adjusted"] = self.adjusted
-            properties["unadjusted"] = self.unadjusted
-            properties["observed"] = self.observed
-        else:
-            properties["adjusted"] = self.adjusted._asdict()
-            properties["unadjusted"] = self.unadjusted._asdict()
-            properties["observed"] = self.observed._asdict()
-
-        return {
-            "type": "Feature",
-            "properties": properties,
-            "geometry": {"type": "Point", "coordinates": list(self.position)},
-        }
-
-
 ModelMetadata = namedtuple(
     "ModelMetadata",
     (
@@ -226,7 +178,7 @@ def scalar(
     initialization_time: str,
     loop: MinimLoop,
     filters: MultiDict,
-) -> List[Observation]:
+) -> pd.DataFrame:
     data = open_diagnostic(
         diag_zarr,
         model,
@@ -240,21 +192,7 @@ def scalar(
     )
     data = apply_filters(data, filters)
 
-    return [
-        Observation(
-            variable.name.lower(),
-            VariableType.SCALAR,
-            loop,
-            adjusted=float(data["obs_minus_forecast_adjusted"].values[idx]),
-            unadjusted=float(data["obs_minus_forecast_unadjusted"].values[idx]),
-            observed=float(data["observation"].values[idx]),
-            position=Coordinate(
-                float(data["longitude"].values[idx]),
-                float(data["latitude"].values[idx]),
-            ),
-        )
-        for idx in range(data.dims["nobs"])
-    ]
+    return data.to_dataframe()
 
 
 def temperature(
@@ -267,7 +205,7 @@ def temperature(
     initialization_time: str,
     loop: MinimLoop,
     filters: MultiDict,
-) -> List[Observation]:
+) -> pd.DataFrame:
     return scalar(
         diag_zarr,
         model,
@@ -292,7 +230,7 @@ def moisture(
     initialization_time: str,
     loop: MinimLoop,
     filters: MultiDict,
-) -> List[Observation]:
+) -> pd.DataFrame:
     return scalar(
         diag_zarr,
         model,
@@ -317,7 +255,7 @@ def pressure(
     initialization_time: str,
     loop: MinimLoop,
     filters: MultiDict,
-) -> List[Observation]:
+) -> pd.DataFrame:
     return scalar(
         diag_zarr,
         model,
@@ -332,28 +270,6 @@ def pressure(
     )
 
 
-def vector_direction(u, v):
-    direction = (90 - np.degrees(np.arctan2(-v, -u))) % 360
-
-    # Anywhere the magnitude of the vector is 0
-    calm = (np.abs(u) == 0) & (np.abs(v) == 0)
-
-    # numpy.arctan2 treats 0.0 and -0.0 differently. Whenever the second
-    # argument to the function is -0.0, it return pi or -pi depending on the
-    # sign of the first argument. Whenever the second argument is 0.0, it will
-    # return 0.0 or -0.0 depending on the sign of the first argument. We
-    # normalize all calm vectors (magnitude 0) to have a direction of 0.0, per
-    # the NCAR Command Language docs.
-    # http://ncl.ucar.edu/Document/Functions/Contributed/wind_direction.shtml
-    direction[calm] = 0.0
-
-    return direction
-
-
-def vector_magnitude(u, v):
-    return np.sqrt(u**2 + v**2)
-
-
 def wind(
     diag_zarr: str,
     model: str,
@@ -364,7 +280,7 @@ def wind(
     initialization_time: str,
     loop: MinimLoop,
     filters: MultiDict,
-) -> List[Observation]:
+) -> pd.DataFrame | pd.Series:
     data = open_diagnostic(
         diag_zarr,
         model,
@@ -379,59 +295,19 @@ def wind(
 
     data = apply_filters(data, filters)
 
-    omf_adj_u = data["obs_minus_forecast_adjusted"].sel(component="u").values
-    omf_adj_v = data["obs_minus_forecast_adjusted"].sel(component="v").values
-    omf_una_u = data["obs_minus_forecast_unadjusted"].sel(component="u").values
-    omf_una_v = data["obs_minus_forecast_unadjusted"].sel(component="v").values
-    obs_u = data["observation"].sel(component="u").values
-    obs_v = data["observation"].sel(component="v").values
-    lng = data["longitude"].values
-    lat = data["latitude"].values
-
-    return [
-        Observation(
-            "wind",
-            VariableType.VECTOR,
-            loop,
-            adjusted=Vector(
-                round(float(omf_adj_u[idx]), 5), round(float(omf_adj_v[idx]), 5)
-            ),
-            unadjusted=Vector(
-                round(float(omf_una_u[idx]), 5), round(float(omf_una_v[idx]), 5)
-            ),
-            observed=Vector(round(float(obs_u[idx]), 5), round(float(obs_v[idx]), 5)),
-            position=Coordinate(round(float(lng[idx]), 5), round(float(lat[idx]), 5)),
-        )
-        for idx in range(data.dims["nobs"])
-    ]
+    return data.to_dataframe()
 
 
-def magnitude(dataset: List[Observation]) -> Generator[Observation, None, None]:
-    for obs in dataset:
-        if isinstance(obs.adjusted, Vector):
-            adjusted = vector_magnitude(obs.adjusted.u, obs.adjusted.v)
-        else:
-            adjusted = abs(obs.adjusted)
-
-        if isinstance(obs.unadjusted, Vector):
-            unadjusted = vector_magnitude(obs.unadjusted.u, obs.unadjusted.v)
-        else:
-            unadjusted = abs(obs.unadjusted)
-
-        if isinstance(obs.observed, Vector):
-            observed = vector_magnitude(obs.observed.u, obs.observed.v)
-        else:
-            observed = abs(obs.observed)
-
-        yield Observation(
-            obs.variable,
-            obs.variable_type,
-            obs.loop,
-            adjusted=adjusted,
-            unadjusted=unadjusted,
-            observed=observed,
-            position=obs.position,
-        )
+def magnitude(dataset: pd.DataFrame) -> pd.DataFrame:
+    return dataset.groupby(level=0).aggregate(
+        {
+            "obs_minus_forecast_adjusted": np.linalg.norm,
+            "obs_minus_forecast_unadjusted": np.linalg.norm,
+            "observation": np.linalg.norm,
+            "longitude": "first",
+            "latitude": "first",
+        }
+    )
 
 
 def get_model_run_list(

src/unified_graphics/routes.py

@@ -178,13 +178,21 @@ def diagnostics(
         initialization_time,
         diag.MinimLoop(loop),
         request.args,
-    )
+    )[
+        [
+            "obs_minus_forecast_adjusted",
+            "obs_minus_forecast_unadjusted",
+            "observation",
+            "longitude",
+            "latitude",
+        ]
+    ]
 
-    response = jsonify(
-        {"type": "FeatureCollection", "features": [obs.to_geojson() for obs in data]}
-    )
+    if "component" in data.index.names:
+        data = data.unstack()
+        data.columns = ["_".join(col) for col in data.columns]
 
-    return response
+    return data.to_json(orient="records"), {"Content-Type": "application/json"}
 
 
 @bp.route(
@@ -210,11 +218,15 @@ def magnitude(
         initialization_time,
         diag.MinimLoop(loop),
         request.args,
-    )
+    )[
+        [
+            "obs_minus_forecast_adjusted",
+            "obs_minus_forecast_unadjusted",
+            "observation",
+            "longitude",
+            "latitude",
+        ]
+    ]
     data = diag.magnitude(data)
 
-    response = jsonify(
-        {"type": "FeatureCollection", "features": [obs.to_geojson() for obs in data]}
-    )
-
-    return response
+    return data.to_json(orient="records"), {"Content-Type": "application/json"}

src/unified_graphics/static/js/component/Chart2DHistogram.js

@@ -138,11 +138,12 @@ export default class Chart2DHistogram extends ChartElement {
   }
 
   set data(value) {
-    if (Array.isArray(value)) {
-      this.#data = value;
-    } else {
-      this.#data = value.features.map((d) => d.properties.adjusted);
-    }
+    this.#data = value.map((record) => {
+      return {
+        u: record["obs_minus_forecast_adjusted_u"],
+        v: record["obs_minus_forecast_adjusted_v"],
+      };
+    });
 
     // FIXME: This is duplicated across all charts to ensure that they fire
     // this event. This event is used by the ColorRamp component so that it

src/unified_graphics/static/js/component/ChartHistogram.js

@@ -116,7 +116,7 @@ export default class ChartHistogram extends ChartElement {
       case "src":
         fetch(newValue)
           .then((response) => response.json())
-          .then((data) => data.features.map((d) => d.properties.adjusted))
+          .then((data) => data.map((d) => d.obs_minus_forecast_adjusted))
           .then((data) => (this.data = data));
         break;
       default:

src/unified_graphics/static/js/component/ChartMap.js

@@ -149,9 +149,7 @@ export default class ChartMap extends ChartElement {
     if (!observations) return scaleQuantize().range(schemePurples[9]);
 
     /** @type number[] */
-    const [lower, upper] = extent(observations.features, (feature) =>
-      get(feature, this.fill)
-    );
+    const [lower, upper] = extent(observations, (feature) => get(feature, this.fill));
 
     const isDiverging = lower / Math.abs(lower) !== upper / Math.abs(upper);
     const largestBound = Math.max(Math.abs(lower), Math.abs(upper));
@@ -203,7 +201,7 @@ export default class ChartMap extends ChartElement {
 
     if (!(borders || observations)) return;
 
-    this.#projection.fitSize([width, height], observations ?? borders);
+    this.#projection.fitSize([width, height], borders);
 
     const ctx = canvas.getContext("2d");
     if (!ctx) return;
@@ -230,9 +228,10 @@ export default class ChartMap extends ChartElement {
     const fill = this.scale;
     const radius = 2;
 
-    observations.features.forEach((feature) => {
+    console.log(observations);
+    observations.forEach((feature) => {
       const value = get(feature, fillProp);
-      const [x, y] = this.#projection(feature.geometry.coordinates);
+      const [x, y] = this.#projection([feature.longitude, feature.latitude]);
 
       ctx.save();
       ctx.fillStyle = fill(value);

src/unified_graphics/templates/layouts/diag.html

@@ -41,7 +41,7 @@ <h2 class="heading-2 flex-0">{% if minim_loop == "ges" %}Guess{% else %}Analysis
     <chart-container class="padding-2 radius-md bg-white shadow-1">
       <chart-map id="observations-{{ minim_loop }}"
         src="{{ map_url[minim_loop] }}"
-        fill="properties.adjusted"></chart-map>
+        fill="obs_minus_forecast_adjusted"></chart-map>
       <color-ramp slot="legend" for="observations-{{ minim_loop }}" class="font-ui-3xs" format="s"
         >Observation &minus; Forecast</color-ramp>
     </chart-container>

tests/test_diag.py

@@ -234,40 +234,6 @@ def test_open_diagnostic_s3(moto_server, test_dataset, monkeypatch):
     xr.testing.assert_equal(result, expected)
 
 
-# Test cases taken from the examples at
-# http://ncl.ucar.edu/Document/Functions/Contributed/wind_direction.shtml
-@pytest.mark.parametrize(
-    "u,v,expected",
-    (
-        [
-            np.array([10, 0, 0, -10, 10, 10, -10, -10]),
-            np.array([0, 10, -10, 0, 10, -10, 10, -10]),
-            np.array([270, 180, 0, 90, 225, 315, 135, 45]),
-        ],
-        [
-            np.array([0.0, -0.0, 0.0, -0.0]),
-            np.array([0.0, 0.0, -0.0, -0.0]),
-            np.array([0.0, 0.0, 0.0, 0.0]),
-        ],
-    ),
-)
-def test_vector_direction(u, v, expected):
-    result = diag.vector_direction(u, v)
-
-    np.testing.assert_array_almost_equal(result, expected, decimal=5)
-
-
-def test_vector_magnitude():
-    u = np.array([1, 0, 1, 0])
-    v = np.array([0, 1, 1, 0])
-
-    result = diag.vector_magnitude(u, v)
-
-    np.testing.assert_array_almost_equal(
-        result, np.array([1, 1, 1.41421, 0]), decimal=5
-    )
-
-
 @pytest.mark.parametrize(
     "mapping,expected",
     [