Merge pull request #109 from openscm/perf/fair2-vectorise-output-extraction

perf(fair2): vectorise forcing-aggregation output extraction (~9x at 842 members)

Author: benmsanderson

changelog/109.improvement.md

@@ -0,0 +1 @@
+Sped up the FaIR2 adapter's output extraction by materialising FaIR's forcing array to numpy once and aggregating per ensemble member with integer indexing, instead of repeated per-member xarray orthogonal indexing. For an 842-member native-calibration ensemble this cut a pathway from ~370 s to ~40 s (~9x) and removed the superlinear scaling with member count. Results are numerically identical.

src/openscm_runner/adapters/fair2_adapter/_output_extractor.py

@@ -158,17 +158,35 @@ def _concentration_unit(species_name: str) -> str:
 # properties_df, available_species) and returns a 1-D numpy array
 # along the time axis (or None if the recipe can't be computed against
 # the species set FaIR actually ran).
+def _forcing_to_numpy(forcing_da):
+    """
+    Materialise FaIR's forcing DataArray as a plain numpy array once.
+
+    Returns ``(forcing_np, species_index)`` where ``forcing_np`` has axis
+    order ``(timebounds, scenario, config, specie)`` and ``species_index``
+    maps species name -> position on the last axis. Extracting/aggregating
+    per (scenario, member) against this numpy array with integer indexing
+    is orders of magnitude faster than repeated xarray ``.sel().isel()``
+    orthogonal indexing (the dominant cost for large ensembles), and is
+    numerically identical (same species order, same summation order).
+    """
+    da = forcing_da.transpose("timebounds", "scenario", "config", "specie")
+    species = list(da["specie"].values)
+    return np.asarray(da.values), {name: i for i, name in enumerate(species)}
+
+
 def _sum_forcing_over(
-    forcing_da, sc_idx: int, member_offset: int, species_to_sum: list[str]
+    forcing_np,
+    species_index: "dict[str, int]",
+    sc_idx: int,
+    member_offset: int,
+    species_to_sum: list[str],
 ) -> np.ndarray:
-    """Helper: sum FaIR's forcing array over a list of species names."""
-    species_present = [s for s in species_to_sum if s in forcing_da["specie"].values]
-    if not species_present:
+    """Sum the forcing array over a list of species names (numpy path)."""
+    idx = [species_index[s] for s in species_to_sum if s in species_index]
+    if not idx:
         return None
-    sliced = forcing_da.sel(specie=species_present).isel(
-        scenario=sc_idx, config=member_offset
-    )
-    return sliced.sum(dim="specie").values
+    return forcing_np[:, sc_idx, member_offset, idx].sum(axis=-1)
 
 
 def _species_by_property(properties_df, predicate) -> list[str]:
@@ -179,20 +197,21 @@ def _species_by_property(properties_df, predicate) -> list[str]:
 
 
 def _build_forcing_aggregations(  # noqa: PLR0912, PLR0915
-    forcing_da, sc_idx: int, member_offset: int, properties_df
+    forcing_np, species_index, species_in_run, sc_idx: int,
+    member_offset: int, properties_df,
 ) -> "dict[str, tuple[np.ndarray, str]]":
     """
     Build the value arrays for the supported forcing aggregations.
 
     Returns a dict keyed by openscm-runner variable name; values are
     ``(values, unit)`` tuples. Aggregations whose species list does
     not intersect FaIR's actual species are omitted, not zero-filled.
+    Operates on the pre-materialised numpy forcing array (see
+    :func:`_forcing_to_numpy`) for speed.
     """
-    species_in_run = list(forcing_da["specie"].values)
-
     def sum_over(species_list: list[str]) -> np.ndarray:
         return _sum_forcing_over(
-            forcing_da, sc_idx, member_offset, species_list
+            forcing_np, species_index, sc_idx, member_offset, species_list
         )
 
     ghgs = _species_by_property(
@@ -230,11 +249,7 @@ def maybe(name: str, values):
         out[name] = (values, forcing_unit)
 
     # Anthropogenic = total forcing minus Solar minus Volcanic
-    total = (
-        forcing_da.isel(scenario=sc_idx, config=member_offset)
-        .sum(dim="specie")
-        .values
-    )
+    total = forcing_np[:, sc_idx, member_offset, :].sum(axis=-1)
     natural_species = [s for s in ("Solar", "Volcanic") if s in species_in_run]
     if natural_species:
         natural = sum_over(natural_species)
@@ -283,13 +298,9 @@ def maybe(name: str, values):
         "Effective Radiative Forcing|Solar": "Solar",
     }
     for openscm_name, fair_specie in single_specie_aliases.items():
-        if fair_specie not in species_in_run:
+        if fair_specie not in species_index:
             continue
-        values = (
-            forcing_da.sel(specie=fair_specie)
-            .isel(scenario=sc_idx, config=member_offset)
-            .values
-        )
+        values = forcing_np[:, sc_idx, member_offset, species_index[fair_specie]]
         maybe(openscm_name, values)
 
     # RCMIP-aligned hierarchical aggregations. Most are aliases for
@@ -367,12 +378,8 @@ def maybe(name: str, values):
     for spec_leaf, spec in (
         ("CO2", "CO2"), ("CH4", "CH4"), ("N2O", "N2O"),
     ):
-        if spec in species_in_run:
-            v = (
-                forcing_da.sel(specie=spec)
-                .isel(scenario=sc_idx, config=member_offset)
-                .values
-            )
+        if spec in species_index:
+            v = forcing_np[:, sc_idx, member_offset, species_index[spec]]
             maybe(f"Effective Radiative Forcing|Anthropogenic|{spec_leaf}", v)
 
     return out
@@ -406,6 +413,9 @@ def extract_outputs(  # noqa: PLR0913, PLR0912, PLR0915
     rows: list = []
 
     species_in_run = list(f.forcing["specie"].values)
+    # Materialise the forcing array once (dominant cost otherwise is
+    # repeated xarray orthogonal indexing per member); aggregate in numpy.
+    forcing_np, forcing_species_index = _forcing_to_numpy(f.forcing)
 
     # Pre-compute aggregations once per (scenario, member) since the
     # recipes share intermediate sums.
@@ -415,7 +425,8 @@ def extract_outputs(  # noqa: PLR0913, PLR0912, PLR0915
 
             aggregations = (
                 _build_forcing_aggregations(
-                    f.forcing, sc_idx, member_offset, properties_df
+                    forcing_np, forcing_species_index, species_in_run,
+                    sc_idx, member_offset, properties_df,
                 )
                 if properties_df is not None
                 else {}

tests/unit/adapters/test_fair2_output_extractor.py

@@ -0,0 +1,126 @@
+"""Equivalence tests for the FaIR2 forcing-aggregation fast path.
+
+The extractor materialises FaIR's forcing DataArray to numpy once and
+aggregates per (scenario, member) with integer indexing, instead of
+repeated xarray ``.sel().isel()`` orthogonal indexing (the dominant cost
+for large ensembles). These tests pin that the numpy path is numerically
+identical to a direct xarray reduction over the same species.
+"""
+from __future__ import annotations
+
+import numpy as np
+import pandas as pd
+import pytest
+
+xr = pytest.importorskip("xarray")
+
+from openscm_runner.adapters.fair2_adapter._output_extractor import (  # noqa: E402
+    _build_forcing_aggregations,
+    _forcing_to_numpy,
+    _sum_forcing_over,
+)
+
+SPECIES = [
+    "CO2", "CH4", "N2O", "HFC-134a", "CF4", "CFC-11",
+    "Aerosol-radiation interactions", "Aerosol-cloud interactions",
+    "Ozone", "Solar", "Volcanic",
+]
+
+
+def _synthetic_forcing(n_t=10, n_scen=2, n_cfg=5, seed=0):
+    rng = np.random.default_rng(seed)
+    data = rng.standard_normal((n_t, n_scen, n_cfg, len(SPECIES)))
+    return xr.DataArray(
+        data,
+        dims=("timebounds", "scenario", "config", "specie"),
+        coords={
+            "timebounds": np.arange(1750, 1750 + n_t),
+            "scenario": [f"s{i}" for i in range(n_scen)],
+            "config": np.arange(n_cfg),
+            "specie": SPECIES,
+        },
+    )
+
+
+def _properties_df():
+    ghg = {"CO2", "CH4", "N2O", "HFC-134a", "CF4", "CFC-11"}
+    types = {
+        "HFC-134a": "f-gas", "CF4": "f-gas", "CFC-11": "cfc-11",
+    }
+    return pd.DataFrame(
+        {
+            "greenhouse_gas": [s in ghg for s in SPECIES],
+            "type": [types.get(s, "other") for s in SPECIES],
+        },
+        index=SPECIES,
+    )
+
+
+def _xarray_sum(da, sc_idx, member, species_list):
+    present = [s for s in species_list if s in da["specie"].values]
+    if not present:
+        return None
+    return (
+        da.sel(specie=present)
+        .isel(scenario=sc_idx, config=member)
+        .sum(dim="specie")
+        .values
+    )
+
+
+def test_sum_forcing_over_matches_xarray():
+    da = _synthetic_forcing()
+    forcing_np, species_index = _forcing_to_numpy(da)
+    for sc_idx in range(da.sizes["scenario"]):
+        for member in range(da.sizes["config"]):
+            for species_list in (
+                ["CO2", "CH4", "N2O"],
+                ["HFC-134a", "CF4"],
+                ["Solar", "Volcanic"],
+                ["not-a-species"],  # empty intersection -> None
+            ):
+                got = _sum_forcing_over(
+                    forcing_np, species_index, sc_idx, member, species_list
+                )
+                ref = _xarray_sum(da, sc_idx, member, species_list)
+                if ref is None:
+                    assert got is None
+                else:
+                    np.testing.assert_array_equal(got, ref)
+
+
+def test_build_forcing_aggregations_matches_xarray():
+    da = _synthetic_forcing(seed=1)
+    props = _properties_df()
+    forcing_np, species_index = _forcing_to_numpy(da)
+    species_in_run = list(da["specie"].values)
+
+    for sc_idx in range(da.sizes["scenario"]):
+        for member in range(da.sizes["config"]):
+            out = _build_forcing_aggregations(
+                forcing_np, species_index, species_in_run, sc_idx, member, props
+            )
+            erf = "Effective Radiative Forcing"
+
+            # Total anthropogenic = all species minus Solar/Volcanic.
+            total = _xarray_sum(da, sc_idx, member, species_in_run)
+            natural = _xarray_sum(da, sc_idx, member, ["Solar", "Volcanic"])
+            np.testing.assert_allclose(
+                out[f"{erf}|Anthropogenic"][0], total - natural, rtol=0, atol=1e-12
+            )
+            # Greenhouse gases.
+            np.testing.assert_array_equal(
+                out[f"{erf}|Greenhouse Gases"][0],
+                _xarray_sum(da, sc_idx, member,
+                            ["CO2", "CH4", "N2O", "HFC-134a", "CF4", "CFC-11"]),
+            )
+            # F-gases.
+            np.testing.assert_array_equal(
+                out[f"{erf}|F-Gases"][0],
+                _xarray_sum(da, sc_idx, member, ["HFC-134a", "CF4"]),
+            )
+            # Single-species pass-through (Ozone).
+            np.testing.assert_array_equal(
+                out[f"{erf}|Ozone"][0],
+                da.sel(specie="Ozone").isel(scenario=sc_idx, config=member).values,
+            )