From 291f211bb7ebd0792e0c546cfbfb16eda7f612c5 Mon Sep 17 00:00:00 2001 From: benmsanderson Date: Thu, 2 Jul 2026 12:47:04 +0200 Subject: [PATCH 1/2] perf(fair2): vectorise forcing-aggregation output extraction Profiling a native-calibration ensemble run (py-spy, ssp245, 400-842 members) showed ~68% of wall time in the output extractor's forcing aggregations -- specifically repeated xarray `.sel(specie=...).isel(...)` orthogonal indexing, run once per (scenario, member) for every aggregation. FaIR's actual integration was ~6%, and the conc-driven `unstep_concentration` path ~0.3% (so conc-vs-emissions is not the cost). Materialise `f.forcing` to a numpy array once (species->index map) and do the per-member sums / single-species lookups with integer indexing. Same species order and summation order, so results are numerically identical. Benchmarks (ssp245, native AR7 1.6.0 calibration, single core): members before after 200 23.9 s 14.4 s 400 63.9 s 21.5 s 842 ~370 s 39.6 s (~9x; superlinear scaling removed) GSAT, ERF|Greenhouse Gases and ERF|Aerosols (ensemble mean AND std) match the previous xarray path to 6 d.p. Adds a unit equivalence test (numpy aggregation == direct xarray reduction) and the existing modern-adapter snapshot is unchanged. No FaIR internals touched. Co-Authored-By: Claude Opus 4.8 --- .../fair2_adapter/_output_extractor.py | 71 +++++----- .../adapters/test_fair2_output_extractor.py | 126 ++++++++++++++++++ 2 files changed, 167 insertions(+), 30 deletions(-) create mode 100644 tests/unit/adapters/test_fair2_output_extractor.py diff --git a/src/openscm_runner/adapters/fair2_adapter/_output_extractor.py b/src/openscm_runner/adapters/fair2_adapter/_output_extractor.py index 5868fe00..6db32d48 100644 --- a/src/openscm_runner/adapters/fair2_adapter/_output_extractor.py +++ b/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,7 +197,8 @@ 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. @@ -187,12 +206,12 @@ def _build_forcing_aggregations( # noqa: PLR0912, PLR0915 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 {} diff --git a/tests/unit/adapters/test_fair2_output_extractor.py b/tests/unit/adapters/test_fair2_output_extractor.py new file mode 100644 index 00000000..626e60ee --- /dev/null +++ b/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, + ) From 2f3abcb436dd3b332a129c74f1127b3e963bf2d6 Mon Sep 17 00:00:00 2001 From: benmsanderson Date: Thu, 2 Jul 2026 12:47:42 +0200 Subject: [PATCH 2/2] docs(changelog): add fragment for #109 Co-Authored-By: Claude Opus 4.8 --- changelog/109.improvement.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changelog/109.improvement.md diff --git a/changelog/109.improvement.md b/changelog/109.improvement.md new file mode 100644 index 00000000..5b3257ab --- /dev/null +++ b/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.