From 7573f4ab8a510fe5f1df0ec85893074cb9f3261a Mon Sep 17 00:00:00 2001 From: Fan Zhang Date: Fri, 26 Jun 2026 17:44:03 +0800 Subject: [PATCH 01/15] feat(output_parser): retain physical zout altitudes; add heating-rate column constant --- src/pyradtran/core/output_parser.py | 51 +++++++++++++++++- src/pyradtran/core/runner.py | 8 ++- tests/test_output_parser_zout.py | 83 +++++++++++++++++++++++++++++ 3 files changed, 139 insertions(+), 3 deletions(-) create mode 100644 tests/test_output_parser_zout.py diff --git a/src/pyradtran/core/output_parser.py b/src/pyradtran/core/output_parser.py index abe1579..f087e8c 100644 --- a/src/pyradtran/core/output_parser.py +++ b/src/pyradtran/core/output_parser.py @@ -6,17 +6,53 @@ from __future__ import annotations +from collections.abc import Sequence from pathlib import Path import numpy as np import xarray as xr +#: Name of the heating-rate data variable when libRadtran's heat output is +#: requested (e.g. via ``solver.dynamic_heat_unit``). The ASCII parser handles +#: arbitrary column names; this constant pins the convention used across +#: pyRadtran so callers and the viz layer agree on the variable name. +HEATING_RATE_COLUMN = "heating_rate" + + +def resolve_zout_tokens( + zout: Sequence[float | str], atmosphere_top_km: float = 120.0 +) -> list[float]: + """Resolve uvspec ``zout`` tokens (which may contain ``"toa"``/``"surface"``) + into numeric altitudes in km above ground level. + + Args: + zout: Output levels as given to uvspec (floats and/or keyword strings). + atmosphere_top_km: Altitude (km) to which ``"toa"``/``"top"`` resolve. + + Returns: + List of float altitudes in km, same length and order as ``zout``. + """ + resolved: list[float] = [] + for z in zout: + if isinstance(z, str): + token = z.strip().lower() + if token in ("toa", "top"): + resolved.append(float(atmosphere_top_km)) + elif token in ("surface", "ground", "sfc", "0"): + resolved.append(0.0) + else: + resolved.append(float(z)) + else: + resolved.append(float(z)) + return resolved + def parse_output( output_path: str | Path, format: str = "netcdf", n_zout: int = 1, column_names: list[str] | None = None, + zout_levels_km: Sequence[float] | None = None, ) -> xr.Dataset: """Parse uvspec output file into xarray.Dataset. @@ -25,6 +61,8 @@ def parse_output( format: Output format — "netcdf" or "ascii". n_zout: Number of zout levels in output (for ASCII multi-level). column_names: Custom column names for output_user ASCII output. + zout_levels_km: Physical zout altitudes in km. When provided (ASCII path), + used as the ``zout`` coordinate instead of an integer index. Returns: xarray.Dataset with wavelength (and optionally zout) as coordinates. @@ -34,7 +72,12 @@ def parse_output( if format == "netcdf": return _parse_netcdf(output_path) else: - return _parse_ascii(output_path, n_zout=n_zout, column_names=column_names) + return _parse_ascii( + output_path, + n_zout=n_zout, + column_names=column_names, + zout_levels_km=zout_levels_km, + ) def _parse_netcdf(path: Path) -> xr.Dataset: @@ -52,6 +95,7 @@ def _parse_ascii( path: Path, n_zout: int = 1, column_names: list[str] | None = None, + zout_levels_km: Sequence[float] | None = None, ) -> xr.Dataset: """Read ASCII uvspec output. @@ -108,5 +152,8 @@ def _parse_ascii( coords["wavelength"] = reshaped[:, 0, i] else: data_vars[name] = (("wavelength", "zout"), reshaped[:, :, i]) - coords["zout"] = np.arange(n_zout) + if zout_levels_km is not None: + coords["zout"] = np.asarray(zout_levels_km, dtype=float) + else: + coords["zout"] = np.arange(n_zout) return xr.Dataset(data_vars, coords=coords) diff --git a/src/pyradtran/core/runner.py b/src/pyradtran/core/runner.py index b201191..ddc3509 100644 --- a/src/pyradtran/core/runner.py +++ b/src/pyradtran/core/runner.py @@ -18,7 +18,7 @@ import xarray as xr -from pyradtran.core.output_parser import parse_output +from pyradtran.core.output_parser import parse_output, resolve_zout_tokens from pyradtran.data.resolver import DataResolver if TYPE_CHECKING: @@ -188,11 +188,17 @@ def execute( if output_format != "netcdf" and scene.output and scene.output.quantities: col_names = ["wavelength" if q == "lambda" else q for q in scene.output.quantities] + # Preserve physical zout altitudes (ASCII path); "toa" -> top-of-atmosphere km. + zout_levels_km = None + if output_format != "netcdf" and scene.output and scene.output.zout: + zout_levels_km = resolve_zout_tokens(scene.output.zout) + result = parse_output( out_file, format=output_format, n_zout=n_zout, column_names=col_names, + zout_levels_km=zout_levels_km, ) result.attrs["input_config"] = input_text.strip() diff --git a/tests/test_output_parser_zout.py b/tests/test_output_parser_zout.py new file mode 100644 index 0000000..0a55093 --- /dev/null +++ b/tests/test_output_parser_zout.py @@ -0,0 +1,83 @@ +"""Tests for physical zout-altitude retention and the heating-rate column.""" +from __future__ import annotations + +from pathlib import Path + +from pyradtran.core.output_parser import ( + HEATING_RATE_COLUMN, + parse_output, + resolve_zout_tokens, +) + + +def _write_ascii(tmp_path: Path, text: str) -> Path: + p = tmp_path / "out.dat" + p.write_text(text) + return p + + +def test_resolve_zout_tokens_handles_strings_and_floats(): + assert resolve_zout_tokens([0, "toa"], atmosphere_top_km=120.0) == [0.0, 120.0] + assert resolve_zout_tokens(["surface", "top", 5.0], atmosphere_top_km=100.0) == [ + 0.0, + 100.0, + 5.0, + ] + + +def test_parse_ascii_uses_physical_zout_levels(tmp_path): + # 2 wavelengths x 2 zout levels, standard 7 columns. + # Each wavelength's zout levels are consecutive; zout index 0 = 0.0 km + # (surface, attenuated), index 1 = 120 km (TOA, unattenuated). + text = "\n".join( + [ + "300 0.8 0.4 0.1 0.0 0.0 0.0", + "300 1.0 0.5 0.1 0.0 0.0 0.0", + "500 0.9 0.5 0.2 0.0 0.0 0.0", + "500 1.0 0.6 0.2 0.0 0.0 0.0", + ] + ) + path = _write_ascii(tmp_path, text) + ds = parse_output( + path, + format="ascii", + n_zout=2, + zout_levels_km=[0.0, 120.0], + ) + assert list(ds["zout"].values) == [0.0, 120.0] + assert ds["edir"].sel(zout=0.0).values.tolist() == [0.8, 0.9] # surface = low z + + +def test_parse_ascii_falls_back_to_integer_index_when_no_levels(tmp_path): + text = "\n".join( + [ + "300 1.0 0.5 0.1 0.0 0.0 0.0", + "300 0.8 0.4 0.1 0.0 0.0 0.0", + ] + ) + path = _write_ascii(tmp_path, text) + ds = parse_output(path, format="ascii", n_zout=2) + assert list(ds["zout"].values) == [0, 1] # legacy integer index preserved + + +def test_parse_ascii_heating_rate_column(tmp_path): + # 1 wavelength x 1 zout, 8 columns (standard 7 + heating_rate). + text = "300 1.0 0.5 0.1 0.0 0.0 0.0 0.234\n" + path = _write_ascii(tmp_path, text) + ds = parse_output( + path, + format="ascii", + n_zout=1, + column_names=[ + "wavelength", + "edir", + "edn", + "eup", + "udir", + "udn", + "uup", + HEATING_RATE_COLUMN, + ], + ) + assert HEATING_RATE_COLUMN in ds.data_vars + assert float(ds[HEATING_RATE_COLUMN].values.item()) == 0.234 From 989d3c5108ca91446194d49a45f9b067dd2dc9ca Mon Sep 17 00:00:00 2001 From: Fan Zhang Date: Fri, 26 Jun 2026 17:45:44 +0800 Subject: [PATCH 02/15] feat(postprocess): add T/R/A budget derivation (add_budget_vars, compute_budget) --- src/pyradtran/core/postprocess.py | 72 +++++++++++++++++++++++++++++++ tests/test_postprocess_budget.py | 57 ++++++++++++++++++++++++ 2 files changed, 129 insertions(+) create mode 100644 src/pyradtran/core/postprocess.py create mode 100644 tests/test_postprocess_budget.py diff --git a/src/pyradtran/core/postprocess.py b/src/pyradtran/core/postprocess.py new file mode 100644 index 0000000..a258592 --- /dev/null +++ b/src/pyradtran/core/postprocess.py @@ -0,0 +1,72 @@ +"""Pure dataset post-processing for RT outputs (no Runner / no libRadtran). + +Transforms raw parsed flux datasets into physically meaningful quantities. +Kept separate from :mod:`pyradtran.viz` so the same derivations serve numeric +use (logging, NetCDF export, comparisons) without importing matplotlib. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np +import xarray as xr + + +@dataclass(frozen=True) +class BudgetResult: + """Spectral transmittance / reflectance / absorptance of the column.""" + + transmittance: np.ndarray + reflectance: np.ndarray + absorptance: np.ndarray + wavelength: np.ndarray + + +def _boundary_indices(zout: np.ndarray) -> tuple[int, int]: + """Return ``(surface_idx, toa_idx)`` — the min- and max-altitude zout levels.""" + z = np.asarray(zout, dtype=float) + return int(np.argmin(z)), int(np.argmax(z)) + + +def _budget_arrays(ds: xr.Dataset, *, f_incident: str) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + if "zout" not in ds.dims or ds.sizes["zout"] < 2: + raise ValueError( + "Budget derivation requires a 2-D dataset (wavelength, zout) with at " + "least surface + TOA levels." + ) + surface, toa = _boundary_indices(ds["zout"].values) + f_inc = ds[f_incident].isel(zout=toa).values + f_down_sfc = (ds["edir"].isel(zout=surface) + ds["edn"].isel(zout=surface)).values + f_up_toa = ds["eup"].isel(zout=toa).values + with np.errstate(divide="ignore", invalid="ignore"): + trans = np.where(f_inc != 0, f_down_sfc / f_inc, np.nan) + refl = np.where(f_inc != 0, f_up_toa / f_inc, np.nan) + absorp = 1.0 - trans - refl + return trans, refl, absorp + + +def add_budget_vars(ds: xr.Dataset, *, f_incident: str = "edir") -> xr.Dataset: + """Return a copy of ``ds`` with ``transmittance``/``reflectance``/``absorptance``. + + Conventions: ``T = (edir+edn)@surface / f_incident@TOA``, + ``R = eup@TOA / f_incident@TOA``, ``A = 1 - T - R``. Surface = lowest zout + level, TOA = highest zout level. + """ + trans, refl, absorp = _budget_arrays(ds, f_incident=f_incident) + out = ds.copy() + out["transmittance"] = ("wavelength", trans) + out["reflectance"] = ("wavelength", refl) + out["absorptance"] = ("wavelength", absorp) + return out + + +def compute_budget(ds: xr.Dataset, *, f_incident: str = "edir") -> BudgetResult: + """Typed companion to :func:`add_budget_vars` returning a :class:`BudgetResult`.""" + trans, refl, absorp = _budget_arrays(ds, f_incident=f_incident) + return BudgetResult( + transmittance=trans, + reflectance=refl, + absorptance=absorp, + wavelength=np.asarray(ds["wavelength"].values, dtype=float), + ) diff --git a/tests/test_postprocess_budget.py b/tests/test_postprocess_budget.py new file mode 100644 index 0000000..64e229f --- /dev/null +++ b/tests/test_postprocess_budget.py @@ -0,0 +1,57 @@ +"""Tests for T/R/A budget derivation in pyradtran.core.postprocess.""" +from __future__ import annotations + +import numpy as np +import xarray as xr + +from pyradtran.core.postprocess import BudgetResult, add_budget_vars, compute_budget + + +def _flux_ds() -> xr.Dataset: + wl = np.array([300.0, 500.0, 800.0]) + zout = np.array([0.0, 120.0]) # surface (min) , TOA (max) + # (n_wl, n_zout); surface attenuated, TOA = incident direct = 1.0 + edir = np.array([[0.5, 1.0], [0.6, 1.0], [0.7, 1.0]]) + edn = np.array([[0.1, 0.0], [0.1, 0.0], [0.1, 0.0]]) + eup = np.array([[0.05, 0.2], [0.06, 0.2], [0.07, 0.2]]) + return xr.Dataset( + { + "edir": (("wavelength", "zout"), edir), + "edn": (("wavelength", "zout"), edn), + "eup": (("wavelength", "zout"), eup), + }, + coords={"wavelength": wl, "zout": zout}, + ) + + +def test_add_budget_vars_adds_three_vars_and_closes(): + ds = add_budget_vars(_flux_ds()) + for name in ("transmittance", "reflectance", "absorptance"): + assert name in ds.data_vars + T, R, A = ds["transmittance"].values, ds["reflectance"].values, ds["absorptance"].values + np.testing.assert_allclose(T + R + A, np.ones(3), atol=1e-9) + # T = (edir+edn)@surface / edir@TOA = (0.6, 0.7, 0.8) + np.testing.assert_allclose(T, [0.6, 0.7, 0.8]) + # R = eup@TOA / edir@TOA = 0.2 + np.testing.assert_allclose(R, [0.2, 0.2, 0.2]) + + +def test_budget_quantities_in_unit_interval(): + ds = add_budget_vars(_flux_ds()) + for name in ("transmittance", "reflectance", "absorptance"): + v = ds[name].values + assert np.all(v >= 0.0) and np.all(v <= 1.0) + + +def test_compute_budget_returns_dataclass(): + result = compute_budget(_flux_ds()) + assert isinstance(result, BudgetResult) + np.testing.assert_allclose(result.transmittance, [0.6, 0.7, 0.8]) + assert result.wavelength.tolist() == [300.0, 500.0, 800.0] + + +def test_add_budget_vars_does_not_mutate_input(): + ds = _flux_ds() + before = set(ds.data_vars) + add_budget_vars(ds) + assert set(ds.data_vars) == before From 0a1d783b21d710705113bb0dae1274ddd2e3f834 Mon Sep 17 00:00:00 2001 From: Fan Zhang Date: Fri, 26 Jun 2026 17:47:32 +0800 Subject: [PATCH 03/15] feat(postprocess): add evaluate_composite_on_grid / evaluate_blocks_on_grid --- src/pyradtran/core/postprocess.py | 62 ++++++++++++++++++++++++++++++ tests/test_postprocess_grid.py | 64 +++++++++++++++++++++++++++++++ 2 files changed, 126 insertions(+) create mode 100644 tests/test_postprocess_grid.py diff --git a/src/pyradtran/core/postprocess.py b/src/pyradtran/core/postprocess.py index a258592..1dfc006 100644 --- a/src/pyradtran/core/postprocess.py +++ b/src/pyradtran/core/postprocess.py @@ -70,3 +70,65 @@ def compute_budget(ds: xr.Dataset, *, f_incident: str = "edir") -> BudgetResult: absorptance=absorp, wavelength=np.asarray(ds["wavelength"].values, dtype=float), ) + + +def _layer_centers(z_km: np.ndarray) -> np.ndarray: + z = np.asarray(z_km, dtype=float) + return 0.5 * (z[:-1] + z[1:]) + + +def evaluate_composite_on_grid( + comp, wl_um, z_km, n_legendre: int = 32 +) -> xr.Dataset: + """Evaluate the mixed composite optics on a (wavelength, layer) grid. + + Calls ``comp.evaluate(...)`` (pure analytic mixing — no RT) and wraps the + resulting :class:`LayerOptics` as an :class:`xarray.Dataset` with a + ``layer`` dimension indexed by layer-center altitude. + """ + wl = np.asarray(wl_um, dtype=float) + z = np.asarray(z_km, dtype=float) + lo = comp.evaluate(wl_um=wl, z_km=z, n_legendre=n_legendre) + centers = _layer_centers(z) + return xr.Dataset( + { + "tau": (("wavelength", "layer"), lo.tau), + "ssa": (("wavelength", "layer"), lo.ssa), + "g": (("wavelength", "layer"), lo.g), + }, + coords={ + "wavelength": wl, + "layer": np.arange(centers.size), + "altitude_km": ("layer", centers), + }, + ) + + +def evaluate_blocks_on_grid( + comp, wl_um, z_km, n_legendre: int = 32 +) -> dict[str, xr.Dataset]: + """Per-piece ``(tau, rho)`` on the grid; one dataset per block name. + + ``tau`` comes from each piece's ``to_layer_optics``; ``rho_kg_m3`` is added + only for pieces that carry a ``profile`` (e.g. :class:`PlacedBlock`). + """ + wl = np.asarray(wl_um, dtype=float) + z = np.asarray(z_km, dtype=float) + centers = _layer_centers(z) + out: dict[str, xr.Dataset] = {} + for piece in comp.pieces: + lo = piece.to_layer_optics(wl, z, n_legendre=n_legendre) + name = getattr(piece, "name", type(piece).__name__) + ds = xr.Dataset( + {"tau": (("wavelength", "layer"), lo.tau)}, + coords={ + "wavelength": wl, + "layer": np.arange(centers.size), + "altitude_km": ("layer", centers), + }, + ) + profile = getattr(piece, "profile", None) + if profile is not None: + ds["rho_kg_m3"] = ("layer", np.asarray(profile.evaluate(centers), dtype=float)) + out[name] = ds + return out diff --git a/tests/test_postprocess_grid.py b/tests/test_postprocess_grid.py new file mode 100644 index 0000000..a628103 --- /dev/null +++ b/tests/test_postprocess_grid.py @@ -0,0 +1,64 @@ +"""Tests for composite/block grid evaluation helpers.""" +from __future__ import annotations + +import numpy as np +import xarray as xr + +from pyradtran.core.postprocess import evaluate_blocks_on_grid, evaluate_composite_on_grid +from pyradtran.models.aerosol_composite import LayerOptics + + +def _layer_optics(n_wl: int, n_layer: int, n_leg: int, tau: float) -> LayerOptics: + return LayerOptics( + tau=np.full((n_wl, n_layer), tau), + ssa=np.full((n_wl, n_layer), 0.5), + g=np.full((n_wl, n_layer), 0.7), + legendre_moments=np.zeros((n_wl, n_layer, n_leg)), + ) + + +class _Profile: + def evaluate(self, centers): + return np.full_like(np.asarray(centers, dtype=float), 1e-6) + + +class _FakePiece: + def __init__(self, name: str, tau: float): + self.name = name + self._tau = tau + self.profile = _Profile() + + def to_layer_optics(self, wl_um, altitude_km, n_legendre: int = 32) -> LayerOptics: + return _layer_optics(len(wl_um), len(altitude_km) - 1, n_legendre, self._tau) + + +class _FakeComposite: + def __init__(self, pieces): + self.pieces = pieces + + def evaluate(self, wl_um=None, z_km=None, n_legendre=None) -> LayerOptics: + n_leg = n_legendre or 32 + return _layer_optics(len(wl_um), len(z_km) - 1, n_leg, 0.9) + + +def test_evaluate_composite_on_grid_shape_and_coords(): + wl = np.array([0.3, 0.5, 0.8]) + z = np.array([5.0, 3.0, 1.0, 0.0]) # 3 layers + comp = _FakeComposite([_FakePiece("a", 0.5), _FakePiece("b", 0.4)]) + ds = evaluate_composite_on_grid(comp, wl, z, n_legendre=16) + assert set(["tau", "ssa", "g"]).issubset(ds.data_vars) + assert ds["tau"].shape == (3, 3) # (n_wl, n_layer) + assert ds.sizes["layer"] == 3 + # layer altitude coord = layer centers + np.testing.assert_allclose(ds["altitude_km"].values, [4.0, 2.0, 0.5]) + + +def test_evaluate_blocks_on_grid_returns_one_ds_per_piece_with_rho(): + wl = np.array([0.5]) + z = np.array([2.0, 0.0]) # 1 layer + comp = _FakeComposite([_FakePiece("soot", 0.5), _FakePiece("sulfate", 0.2)]) + out = evaluate_blocks_on_grid(comp, wl, z, n_legendre=8) + assert set(out.keys()) == {"soot", "sulfate"} + assert out["soot"]["tau"].shape == (1, 1) + assert "rho_kg_m3" in out["soot"].data_vars + np.testing.assert_allclose(out["soot"]["rho_kg_m3"].values, [1e-6]) From a6c885a6771b0b17b1dfe412c1e8cbfa715b3bff Mon Sep 17 00:00:00 2001 From: Fan Zhang Date: Fri, 26 Jun 2026 17:49:13 +0800 Subject: [PATCH 04/15] feat(workflow): add component-attribution orchestration (N+1 leave-one-out) --- src/pyradtran/workflow/__init__.py | 5 ++ src/pyradtran/workflow/attribution.py | 60 +++++++++++++++++++++ tests/test_workflow_attribution.py | 78 +++++++++++++++++++++++++++ 3 files changed, 143 insertions(+) create mode 100644 src/pyradtran/workflow/__init__.py create mode 100644 src/pyradtran/workflow/attribution.py create mode 100644 tests/test_workflow_attribution.py diff --git a/src/pyradtran/workflow/__init__.py b/src/pyradtran/workflow/__init__.py new file mode 100644 index 0000000..12798c5 --- /dev/null +++ b/src/pyradtran/workflow/__init__.py @@ -0,0 +1,5 @@ +"""RT orchestration workflows (drive the Runner to produce structured results).""" + +from pyradtran.workflow.attribution import AttributionResult, compute_component_attribution + +__all__ = ["AttributionResult", "compute_component_attribution"] diff --git a/src/pyradtran/workflow/attribution.py b/src/pyradtran/workflow/attribution.py new file mode 100644 index 0000000..c6a57dd --- /dev/null +++ b/src/pyradtran/workflow/attribution.py @@ -0,0 +1,60 @@ +"""Component-attribution workflow. + +Drives N+1 RT runs (one full composite + one leave-one-out per piece) via an +injected ``execute_many`` and returns per-block contributions +``full - leave_one_out``. The plot layer consumes :class:`AttributionResult` +by duck-typing; this module owns the data contract and the orchestration. +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass + +import xarray as xr + + +@dataclass(frozen=True) +class AttributionResult: + """Full-field run plus per-block contribution datasets.""" + + full: xr.Dataset + contributions: dict[str, xr.Dataset] + + +def compute_component_attribution( + build_scene: Callable, + composite, + execute_many: Callable[[list], list[xr.Dataset]], +) -> AttributionResult: + """Compute per-block RT contributions by leave-one-out subtraction. + + Args: + build_scene: Callable mapping a (possibly leave-one-out) composite to a + runnable :class:`~pyradtran.scene.Scene`. + composite: A composite with ``.pieces`` and ``.model_copy(update=...)``. + execute_many: Callable running a list of scenes in parallel and returning + the parsed datasets in the same order (e.g. + ``lambda scenes: Runner.execute_many(scenes, uvspec_exe=..., data_path=...)``). + + Returns: + :class:`AttributionResult` whose ``contributions[piece.name]`` is + ``full - leave_piece_out``. + """ + pieces = list(composite.pieces) + names = [getattr(p, "name", f"piece_{i}") for i, p in enumerate(pieces)] + + scenes = [build_scene(composite)] + for i in range(len(pieces)): + remaining = [p for j, p in enumerate(pieces) if j != i] + sub = composite.model_copy(update={"pieces": remaining}) + scenes.append(build_scene(sub)) + + datasets = execute_many(scenes) + full = datasets[0] + + contributions: dict[str, xr.Dataset] = {} + for i, name in enumerate(names): + removed = datasets[i + 1] + contributions[name] = full - removed + return AttributionResult(full=full, contributions=contributions) diff --git a/tests/test_workflow_attribution.py b/tests/test_workflow_attribution.py new file mode 100644 index 0000000..6191cc8 --- /dev/null +++ b/tests/test_workflow_attribution.py @@ -0,0 +1,78 @@ +"""Tests for component-attribution orchestration (no libRadtran).""" +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np +import xarray as xr + +from pyradtran.workflow.attribution import ( + AttributionResult, + compute_component_attribution, +) + + +@dataclass(frozen=True) +class _Piece: + name: str + + +@dataclass(frozen=True) +class _Composite: + pieces: tuple + + def model_copy(self, *, update): + pieces = update.get("pieces", self.pieces) + return _Composite(pieces=tuple(pieces)) + + +def _ds(edir: float) -> xr.Dataset: + return xr.Dataset( + {"edir": ("wavelength", np.array([edir]))}, + coords={"wavelength": np.array([500.0])}, + ) + + +def test_attribution_contribution_is_full_minus_removed(): + comp = _Composite(pieces=(_Piece("soot"), _Piece("dust"))) + # full=1.0; removing soot -> 0.7; removing dust -> 0.9 + results = { + ("soot", "dust"): _ds(1.0), + ("dust",): _ds(0.7), # soot removed + ("soot",): _ds(0.9), # dust removed + } + + def build_scene(c): + return c # the "scene" is just the composite + + def execute_many(scenes): + return [results[tuple(p.name for p in s.pieces)] for s in scenes] + + result = compute_component_attribution(build_scene, comp, execute_many) + assert isinstance(result, AttributionResult) + assert set(result.contributions.keys()) == {"soot", "dust"} + # contribution(soot) = full - removed(soot) = 1.0 - 0.7 = 0.3 + np.testing.assert_allclose(result.contributions["soot"]["edir"].values, [0.3]) + np.testing.assert_allclose(result.contributions["dust"]["edir"].values, [0.1]) + np.testing.assert_allclose(result.full["edir"].values, [1.0]) + + +def test_attribution_contributions_sum_approximates_full(): + comp = _Composite(pieces=(_Piece("a"), _Piece("b"), _Piece("c"))) + # Additive blocks: full = a+b+c, so contributions sum back to full. + results = { + ("a", "b", "c"): _ds(0.9), + ("b", "c"): _ds(0.5), + ("a", "c"): _ds(0.6), + ("a", "b"): _ds(0.7), + } + + def build_scene(c): + return c + + def execute_many(scenes): + return [results[tuple(p.name for p in s.pieces)] for s in scenes] + + result = compute_component_attribution(build_scene, comp, execute_many) + total = sum(result.contributions[name]["edir"].values[0] for name in ("a", "b", "c")) + np.testing.assert_allclose(total, 0.9, atol=1e-9) From 3285fb52d9603f90e22e757c3eee29b1c04c5b67 Mon Sep 17 00:00:00 2001 From: Fan Zhang Date: Fri, 26 Jun 2026 17:51:20 +0800 Subject: [PATCH 05/15] feat(viz): add publication theme, colorblind palette, lazy-matplotlib guard --- src/pyradtran/viz/__init__.py | 1 + src/pyradtran/viz/_style.py | 84 +++++++++++++++++++++++++++++++++++ tests/test_viz_style.py | 41 +++++++++++++++++ 3 files changed, 126 insertions(+) create mode 100644 src/pyradtran/viz/__init__.py create mode 100644 src/pyradtran/viz/_style.py create mode 100644 tests/test_viz_style.py diff --git a/src/pyradtran/viz/__init__.py b/src/pyradtran/viz/__init__.py new file mode 100644 index 0000000..db2f30d --- /dev/null +++ b/src/pyradtran/viz/__init__.py @@ -0,0 +1 @@ +"""pyRadtran visualization layer (pure data-in -> fig-out; lazy matplotlib).""" diff --git a/src/pyradtran/viz/_style.py b/src/pyradtran/viz/_style.py new file mode 100644 index 0000000..f77a50a --- /dev/null +++ b/src/pyradtran/viz/_style.py @@ -0,0 +1,84 @@ +"""Publication-grade theme, colorblind-safe palette, and save helper. + +matplotlib is imported lazily: :func:`require_mpl` is the single chokepoint so +``import pyradtran`` never requires matplotlib. +""" + +from __future__ import annotations + +from pathlib import Path + +# Okabe-Ito colorblind-safe palette. +_OKABE_ITO = [ + "#0072B2", # blue + "#D55E00", # vermillion + "#009E73", # bluish green + "#CC79A7", # reddish purple + "#F0E442", # yellow + "#56B4E9", # sky blue + "#E69F00", # orange + "#000000", # black +] + + +def require_mpl(): + """Import and return matplotlib, or raise an actionable ImportError.""" + try: + import matplotlib + except ImportError as e: # pragma: no cover - exercised only without matplotlib + raise ImportError( + "pyradtran.viz requires matplotlib; install it with `pip install pyradtran[plot]`" + ) from e + return matplotlib + + +def set_theme(style: str = "publication") -> None: + """Apply rcParams for the requested style. + + ``publication`` (default): higher DPI, larger fonts, tighter layout. + ``diagnostic``: lighter settings for quick iteration. + """ + mpl = require_mpl() + if style == "publication": + mpl.rcParams.update( + { + "figure.dpi": 150, + "savefig.dpi": 300, + "font.size": 11, + "axes.labelsize": 11, + "axes.titlesize": 12, + "legend.fontsize": 9, + "xtick.labelsize": 9, + "ytick.labelsize": 9, + "axes.grid": True, + "grid.alpha": 0.3, + "figure.autolayout": True, + } + ) + elif style == "diagnostic": + mpl.rcParams.update( + { + "figure.dpi": 100, + "savefig.dpi": 100, + "font.size": 9, + "axes.grid": True, + "grid.alpha": 0.3, + "figure.autolayout": True, + } + ) + else: + raise ValueError(f"Unknown style: {style!r} (use 'publication' or 'diagnostic')") + + +def get_palette(n: int) -> list[str]: + """Return ``n`` colorblind-safe colors (cycling the Okabe-Ito set).""" + if n <= 0: + return [] + return [_OKABE_ITO[i % len(_OKABE_ITO)] for i in range(n)] + + +def save(fig, path, *, formats=("pdf", "png")) -> None: + """Save ``fig`` to ``path`` with each extension in ``formats`` appended.""" + base = Path(path) + for ext in formats: + fig.savefig(base.with_suffix(f".{ext}")) diff --git a/tests/test_viz_style.py b/tests/test_viz_style.py new file mode 100644 index 0000000..89f6147 --- /dev/null +++ b/tests/test_viz_style.py @@ -0,0 +1,41 @@ +"""Tests for viz theme/palette/save (headless).""" +from __future__ import annotations + +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +from pyradtran.viz._style import get_palette, require_mpl, save, set_theme + + +def test_require_mpl_returns_matplotlib_module(): + assert require_mpl() is matplotlib + + +def test_set_theme_changes_a_distinguishable_param(): + set_theme("publication") + pub_dpi = plt.rcParams["figure.dpi"] + set_theme("diagnostic") + diag_dpi = plt.rcParams["figure.dpi"] + assert pub_dpi != diag_dpi + + +def test_palette_returns_requested_count_and_consistent(): + a = get_palette(3) + b = get_palette(5) + assert len(a) == 3 and len(b) == 5 + # Same palette family: first 3 of the 5-color request equal the 3-color request. + assert a == b[:3] + + +def test_save_writes_requested_formats(tmp_path: Path): + set_theme("publication") + fig, ax = plt.subplots() + ax.plot([0, 1], [0, 1]) + out = tmp_path / "fig" + save(fig, out, formats=("png",)) + assert (tmp_path / "fig.png").exists() + plt.close(fig) From b3fe36d61c96f153eadd66d2e3d159f0df10e4dd Mon Sep 17 00:00:00 2001 From: Fan Zhang Date: Fri, 26 Jun 2026 17:53:01 +0800 Subject: [PATCH 06/15] feat(viz/rt): add plot_spectral, plot_flux_profile, plot_heating_rate --- src/pyradtran/viz/rt.py | 119 +++++++++++++++++++++++++++++++++++++ tests/test_viz_rt_lines.py | 62 +++++++++++++++++++ 2 files changed, 181 insertions(+) create mode 100644 src/pyradtran/viz/rt.py create mode 100644 tests/test_viz_rt_lines.py diff --git a/src/pyradtran/viz/rt.py b/src/pyradtran/viz/rt.py new file mode 100644 index 0000000..12d0655 --- /dev/null +++ b/src/pyradtran/viz/rt.py @@ -0,0 +1,119 @@ +"""RT-result plots consuming xarray.Dataset (pure data-in -> fig-out).""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import xarray as xr + +from pyradtran.core.output_parser import HEATING_RATE_COLUMN +from pyradtran.viz._style import get_palette, require_mpl, save, set_theme + + +def _ensure_axes(ax=None): + require_mpl() + import matplotlib.pyplot as plt + + set_theme() + if ax is None: + fig, ax = plt.subplots() + else: + fig = ax.figure + return fig, ax + + +def _maybe_save(fig, save_path): + if save_path is not None: + save(fig, Path(save_path)) + + +def _surface_index(ds: xr.Dataset) -> int: + """Index of the surface zout level (lowest altitude), or 0 for 1-D datasets.""" + if "zout" not in ds.dims: + return 0 + z = np.asarray(ds["zout"].values, dtype=float) + return int(np.argmin(z)) + + +def plot_spectral( + ds: xr.Dataset, + *, + variables=("edir", "edn", "eup"), + level="surface", + ax=None, + save_path=None, +): + """Plot selected flux variables vs wavelength at a single zout level. + + For 2-D datasets, ``level="surface"`` selects the lowest-altitude zout level. + """ + fig, ax = _ensure_axes(ax) + wl = np.asarray(ds["wavelength"].values, dtype=float) + idx = _surface_index(ds) if level == "surface" else int(level) + colors = get_palette(len(variables)) + for color, var in zip(colors, variables, strict=True): + y = ds[var].values + y = y[:, idx] if y.ndim == 2 else y + ax.plot(wl, y, color=color, label=var, linewidth=1.5) + ax.set_xlabel("Wavelength (nm)") + ax.set_ylabel("Irradiance") + ax.legend(loc="best") + _maybe_save(fig, save_path) + return fig, ax + + +def plot_flux_profile( + ds: xr.Dataset, + *, + variable: str = "edir", + wavelength_nm: float = 550.0, + ax=None, + save_path=None, +): + """Plot a flux variable vs altitude (km) at the wavelength nearest ``wavelength_nm``.""" + if "zout" not in ds.dims: + raise ValueError("plot_flux_profile requires a 2-D dataset with a zout dimension.") + fig, ax = _ensure_axes(ax) + wl = np.asarray(ds["wavelength"].values, dtype=float) + z = np.asarray(ds["zout"].values, dtype=float) + i = int(np.argmin(np.abs(wl - wavelength_nm))) + ax.plot(ds[variable].isel(wavelength=i).values, z, linewidth=1.5) + ax.set_xlabel(f"{variable} @ {wl[i]:.0f} nm") + ax.set_ylabel("Altitude (km)") + _maybe_save(fig, save_path) + return fig, ax + + +def plot_heating_rate( + ds: xr.Dataset, + *, + wavelength_nm: float | None = None, + ax=None, + save_path=None, +): + """Plot heating rate vs altitude (km). If ``wavelength_nm`` is None, plot all + wavelengths as faint background lines plus the nearest-wavelength bold line.""" + if HEATING_RATE_COLUMN not in ds.data_vars: + raise ValueError( + f"Dataset has no '{HEATING_RATE_COLUMN}' variable; request heating-rate " + "output from libRadtran (set dynamic_heat_unit)." + ) + if "zout" not in ds.dims: + raise ValueError("plot_heating_rate requires a 2-D dataset with a zout dimension.") + fig, ax = _ensure_axes(ax) + wl = np.asarray(ds["wavelength"].values, dtype=float) + z = np.asarray(ds["zout"].values, dtype=float) + heat = ds[HEATING_RATE_COLUMN].values + if wavelength_nm is None: + for i in range(wl.size): + ax.plot(heat[i, :], z, linewidth=0.5, alpha=0.4) + else: + i = int(np.argmin(np.abs(wl - wavelength_nm))) + ax.plot(heat[i, :], z, linewidth=1.5) + ax.set_xlabel(f"Heating rate @ {wl[i]:.0f} nm") + ax.set_ylabel("Altitude (km)") + if wavelength_nm is None: + ax.set_xlabel("Heating rate (all wavelengths)") + _maybe_save(fig, save_path) + return fig, ax diff --git a/tests/test_viz_rt_lines.py b/tests/test_viz_rt_lines.py new file mode 100644 index 0000000..97e7c2e --- /dev/null +++ b/tests/test_viz_rt_lines.py @@ -0,0 +1,62 @@ +"""Tests for RT line/profile plots (headless, synthetic data).""" +from __future__ import annotations + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import xarray as xr + +from pyradtran.core.output_parser import HEATING_RATE_COLUMN +from pyradtran.viz.rt import plot_flux_profile, plot_heating_rate, plot_spectral + + +def _spectral_ds() -> xr.Dataset: + wl = np.array([300.0, 500.0, 800.0]) + return xr.Dataset( + { + "edir": ("wavelength", np.array([1.0, 0.9, 0.8])), + "edn": ("wavelength", np.array([0.1, 0.15, 0.2])), + "eup": ("wavelength", np.array([0.05, 0.08, 0.1])), + }, + coords={"wavelength": wl}, + ) + + +def _profile_ds() -> xr.Dataset: + wl = np.array([300.0, 550.0, 800.0]) + zout = np.array([0.0, 4.0, 120.0]) # surface, mid, TOA + edir = np.array([[0.5, 0.7, 1.0], [0.6, 0.8, 1.0], [0.7, 0.9, 1.0]]) + heat = np.array([[0.1, 0.05, 0.0], [0.12, 0.06, 0.0], [0.14, 0.07, 0.0]]) + return xr.Dataset( + { + "edir": (("wavelength", "zout"), edir), + HEATING_RATE_COLUMN: (("wavelength", "zout"), heat), + }, + coords={"wavelength": wl, "zout": zout}, + ) + + +def test_plot_spectral_draws_one_line_per_variable(): + fig, ax = plot_spectral(_spectral_ds(), variables=("edir", "edn", "eup")) + assert len(ax.lines) == 3 + assert "wavelength" in ax.get_xlabel().lower() + plt.close(fig) + + +def test_plot_flux_profile_uses_physical_altitude_axis(): + ds = _profile_ds() + fig, ax = plot_flux_profile(ds, variable="edir", wavelength_nm=550.0) + assert len(ax.lines) == 1 + # y-axis is altitude in km -> max ~120 + ylim = ax.get_ylim() + assert max(abs(v) for v in ylim) >= 100.0 + plt.close(fig) + + +def test_plot_heating_rate_profiles_against_altitude(): + ds = _profile_ds() + fig, ax = plot_heating_rate(ds, wavelength_nm=550.0) + assert len(ax.lines) == 1 + plt.close(fig) From 17af936940ff119fe12ec47032c2ccc1dc5dede9 Mon Sep 17 00:00:00 2001 From: Fan Zhang Date: Fri, 26 Jun 2026 17:55:04 +0800 Subject: [PATCH 07/15] feat(viz/rt): add plot_budget (stacked T/R/A) and plot_rt_overview --- src/pyradtran/viz/rt.py | 44 ++++++++++++++++++++++++++++++++ tests/test_viz_rt_budget.py | 50 +++++++++++++++++++++++++++++++++++++ 2 files changed, 94 insertions(+) create mode 100644 tests/test_viz_rt_budget.py diff --git a/src/pyradtran/viz/rt.py b/src/pyradtran/viz/rt.py index 12d0655..276bf82 100644 --- a/src/pyradtran/viz/rt.py +++ b/src/pyradtran/viz/rt.py @@ -8,6 +8,7 @@ import xarray as xr from pyradtran.core.output_parser import HEATING_RATE_COLUMN +from pyradtran.core.postprocess import add_budget_vars from pyradtran.viz._style import get_palette, require_mpl, save, set_theme @@ -117,3 +118,46 @@ def plot_heating_rate( ax.set_xlabel("Heating rate (all wavelengths)") _maybe_save(fig, save_path) return fig, ax + + +def plot_budget( + ds: xr.Dataset, + *, + components=("transmittance", "reflectance", "absorptance"), + ax=None, + save_path=None, +): + """Stacked-area plot of T/R/A vs wavelength. Requires a budget-enriched dataset.""" + missing = [c for c in components if c not in ds.data_vars] + if missing: + raise ValueError( + f"Dataset missing budget variables {missing}; run add_budget_vars(ds) first." + ) + fig, ax = _ensure_axes(ax) + wl = np.asarray(ds["wavelength"].values, dtype=float) + colors = get_palette(len(components)) + values = np.vstack([np.asarray(ds[c].values, dtype=float) for c in components]) + ax.stackplot(wl, values, labels=list(components), colors=colors, alpha=0.85) + ax.set_xlabel("Wavelength (nm)") + ax.set_ylabel("Fraction of incident flux") + ax.set_ylim(0.0, 1.0) + ax.legend(loc="best") + _maybe_save(fig, save_path) + return fig, ax + + +def plot_rt_overview(ds: xr.Dataset, *, wavelength_nm: float = 550.0): + """Convenience multi-panel: spectral (surface) + flux profile + budget.""" + require_mpl() + import matplotlib.pyplot as plt + + set_theme() + fig, axes = plt.subplots(1, 3, figsize=(13, 4)) + plot_spectral(ds, ax=axes[0]) + if "zout" in ds.dims: + plot_flux_profile(ds, wavelength_nm=wavelength_nm, ax=axes[1]) + else: + axes[1].text(0.5, 0.5, "no zout dim", ha="center", va="center") + plot_budget(add_budget_vars(ds), ax=axes[2]) + fig.tight_layout() + return fig, axes diff --git a/tests/test_viz_rt_budget.py b/tests/test_viz_rt_budget.py new file mode 100644 index 0000000..7af187e --- /dev/null +++ b/tests/test_viz_rt_budget.py @@ -0,0 +1,50 @@ +"""Tests for budget and overview plots (headless, synthetic data).""" +from __future__ import annotations + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import pytest +import xarray as xr + +from pyradtran.core.postprocess import add_budget_vars +from pyradtran.viz.rt import plot_budget, plot_rt_overview + + +def _flux_ds() -> xr.Dataset: + wl = np.array([300.0, 500.0, 800.0]) + zout = np.array([0.0, 120.0]) + edir = np.array([[0.5, 1.0], [0.6, 1.0], [0.7, 1.0]]) + edn = np.array([[0.1, 0.0], [0.1, 0.0], [0.1, 0.0]]) + eup = np.array([[0.05, 0.2], [0.06, 0.2], [0.07, 0.2]]) + return xr.Dataset( + { + "edir": (("wavelength", "zout"), edir), + "edn": (("wavelength", "zout"), edn), + "eup": (("wavelength", "zout"), eup), + }, + coords={"wavelength": wl, "zout": zout}, + ) + + +def test_plot_budget_requires_enriched_dataset_and_draws_components(): + ds = add_budget_vars(_flux_ds()) + fig, ax = plot_budget(ds) + # stackplot + explicit legend -> one legend entry per component (3). + assert len(ax.legend_.get_texts()) == 3 + plt.close(fig) + + +def test_plot_budget_raises_on_unenriched_dataset(): + with pytest.raises(ValueError): + plot_budget(_flux_ds()) + + +def test_plot_rt_overview_has_three_panels(): + fig, ax = plot_rt_overview(_flux_ds(), wavelength_nm=500.0) + # plot_rt_overview returns the figure and a numpy array (or tuple) of 3 axes. + axes = fig.get_axes() + assert len(axes) >= 3 + plt.close(fig) From 79923f9531b558e24e8c2e6d5e44ee614fe3d75c Mon Sep 17 00:00:00 2001 From: Fan Zhang Date: Fri, 26 Jun 2026 17:56:53 +0800 Subject: [PATCH 08/15] feat(viz/composite): add plot_composite_optics and plot_block_profiles --- src/pyradtran/viz/composite.py | 78 ++++++++++++++++++++++++++++++++++ tests/test_viz_composite.py | 56 ++++++++++++++++++++++++ 2 files changed, 134 insertions(+) create mode 100644 src/pyradtran/viz/composite.py create mode 100644 tests/test_viz_composite.py diff --git a/src/pyradtran/viz/composite.py b/src/pyradtran/viz/composite.py new file mode 100644 index 0000000..739061f --- /dev/null +++ b/src/pyradtran/viz/composite.py @@ -0,0 +1,78 @@ +"""Composite and per-block diagnostic plots (pure data-in -> fig-out).""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import xarray as xr + +from pyradtran.viz._style import get_palette, require_mpl, save, set_theme + +_QUANTITY_LABELS = { + "tau": "Optical depth τ", + "ssa": "Single-scattering albedo", + "g": "Asymmetry parameter g", +} +_BLOCK_QUANTITY_LABELS = { + "tau": "Optical depth τ (column sum)", + "rho": "Mass concentration (kg/m³)", +} + + +def _ensure_axes(ax=None): + require_mpl() + import matplotlib.pyplot as plt + + set_theme() + if ax is None: + fig, ax = plt.subplots() + else: + fig = ax.figure + return fig, ax + + +def plot_composite_optics(ds: xr.Dataset, *, quantity: str = "tau", ax=None, save_path=None): + """pcolormesh of τ / ssa / g over (wavelength, altitude) for the mixed composite.""" + if quantity not in ds.data_vars: + raise ValueError(f"Dataset has no '{quantity}' variable.") + fig, ax = _ensure_axes(ax) + wl = np.asarray(ds["wavelength"].values, dtype=float) + alt = np.asarray(ds["altitude_km"].values, dtype=float) + mesh = ax.pcolormesh(wl, alt, ds[quantity].values.T, shading="auto") + ax.set_xlabel("Wavelength (µm)") + ax.set_ylabel("Altitude (km)") + ax.set_title(_QUANTITY_LABELS.get(quantity, quantity)) + fig.colorbar(mesh, ax=ax, label=_QUANTITY_LABELS.get(quantity, quantity)) + if save_path is not None: + save(fig, Path(save_path)) + return fig, ax + + +def plot_block_profiles( + per_block_ds_dict: dict[str, xr.Dataset], + *, + quantity: str = "tau", + ax=None, + save_path=None, +): + """Per-block τ(z) (column-summed over wavelength) or ρ(z) vs altitude.""" + fig, ax = _ensure_axes(ax) + colors = get_palette(len(per_block_ds_dict)) + for color, (name, bds) in zip(colors, per_block_ds_dict.items(), strict=True): + alt = np.asarray(bds["altitude_km"].values, dtype=float) + if quantity == "tau": + y = np.asarray(bds["tau"].values, dtype=float).sum(axis=0) # sum over wavelength + elif quantity == "rho": + if "rho_kg_m3" not in bds.data_vars: + continue # block has no mass profile (e.g. DirectLayerOpticsBlock) + y = np.asarray(bds["rho_kg_m3"].values, dtype=float) + else: + raise ValueError(f"Unknown quantity: {quantity!r} (use 'tau' or 'rho')") + ax.plot(y, alt, color=color, label=name, linewidth=1.5, marker="o", markersize=3) + ax.set_xlabel(_BLOCK_QUANTITY_LABELS.get(quantity, quantity)) + ax.set_ylabel("Altitude (km)") + ax.legend(loc="best") + if save_path is not None: + save(fig, Path(save_path)) + return fig, ax diff --git a/tests/test_viz_composite.py b/tests/test_viz_composite.py new file mode 100644 index 0000000..dd68ef3 --- /dev/null +++ b/tests/test_viz_composite.py @@ -0,0 +1,56 @@ +"""Tests for composite/block diagnostic plots (headless, synthetic data).""" +from __future__ import annotations + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import xarray as xr + +from pyradtran.viz.composite import plot_block_profiles, plot_composite_optics + + +def _composite_ds() -> xr.Dataset: + wl = np.array([0.3, 0.5, 0.8]) + alt = np.array([4.0, 2.0, 0.5]) + tau = np.array([[0.9, 0.5, 0.1], [0.8, 0.4, 0.08], [0.7, 0.3, 0.06]]) + return xr.Dataset( + { + "tau": (("wavelength", "layer"), tau), + "ssa": (("wavelength", "layer"), np.full((3, 3), 0.9)), + "g": (("wavelength", "layer"), np.full((3, 3), 0.6)), + }, + coords={"wavelength": wl, "layer": np.arange(3), "altitude_km": ("layer", alt)}, + ) + + +def _block_dict() -> dict[str, xr.Dataset]: + alt = np.array([4.0, 2.0, 0.5]) + return { + "soot": xr.Dataset( + {"tau": (("wavelength", "layer"), np.full((2, 3), 0.5)), + "rho_kg_m3": ("layer", np.array([1e-6, 5e-7, 1e-7]))}, + coords={"wavelength": np.array([0.3, 0.5]), "layer": np.arange(3), + "altitude_km": ("layer", alt)}, + ), + "dust": xr.Dataset( + {"tau": (("wavelength", "layer"), np.full((2, 3), 0.2)), + "rho_kg_m3": ("layer", np.array([2e-6, 1e-6, 2e-7]))}, + coords={"wavelength": np.array([0.3, 0.5]), "layer": np.arange(3), + "altitude_km": ("layer", alt)}, + ), + } + + +def test_plot_composite_optics_draws_pcolormesh(): + fig, ax = plot_composite_optics(_composite_ds(), quantity="tau") + assert len(ax.collections) >= 1 # a QuadMesh was added + plt.close(fig) + + +def test_plot_block_profiles_draws_one_line_per_block(): + fig, ax = plot_block_profiles(_block_dict(), quantity="tau") + # 2 blocks, summed over wavelength -> 2 lines (one per block). + assert len(ax.lines) == 2 + plt.close(fig) From 891f66a3e197fd94aa54acb9a5d6f8558b5822e5 Mon Sep 17 00:00:00 2001 From: Fan Zhang Date: Fri, 26 Jun 2026 17:59:00 +0800 Subject: [PATCH 09/15] feat(viz/attribution): add plot_component_attribution (LEGO killer plot) --- src/pyradtran/viz/attribution.py | 79 ++++++++++++++++++++++++++++++++ tests/test_viz_attribution.py | 44 ++++++++++++++++++ 2 files changed, 123 insertions(+) create mode 100644 src/pyradtran/viz/attribution.py create mode 100644 tests/test_viz_attribution.py diff --git a/src/pyradtran/viz/attribution.py b/src/pyradtran/viz/attribution.py new file mode 100644 index 0000000..674c188 --- /dev/null +++ b/src/pyradtran/viz/attribution.py @@ -0,0 +1,79 @@ +"""Component-attribution plot (the LEGO killer plot). + +Consumes an AttributionResult by duck-typing (AttributionLike). This module +intentionally does NOT import pyradtran.workflow — the workflow produces the +result, viz renders it. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Protocol, runtime_checkable + +import numpy as np +import xarray as xr + +from pyradtran.viz._style import get_palette, require_mpl, save, set_theme + + +@runtime_checkable +class AttributionLike(Protocol): + full: xr.Dataset + contributions: dict[str, xr.Dataset] + + +def _get(obj, key): + """Read ``key`` from an attribute-style (``.full``) or mapping (``["full"]``) object.""" + if hasattr(obj, key): + return getattr(obj, key) + if isinstance(obj, dict): + return obj[key] + raise TypeError( + f"Expected an object with .{key} or ['{key}']; got {type(obj).__name__}" + ) + + +def plot_component_attribution( + result: AttributionLike, + *, + variable: str = "edir", + level="surface", + ax=None, + save_path=None, +): + """Plot per-block contributions (full - leave_one_out) plus the full curve.""" + require_mpl() + import matplotlib.pyplot as plt + + set_theme() + if ax is None: + fig, ax = plt.subplots() + else: + fig = ax.figure + + full = _get(result, "full") + contributions = _get(result, "contributions") + wl = np.asarray(full["wavelength"].values, dtype=float) + + idx = 0 + if "zout" in full[variable].dims: + z = np.asarray(full["zout"].values, dtype=float) + idx = int(np.argmin(z)) if level == "surface" else int(level) + full_y = full[variable].values[:, idx] + else: + full_y = np.asarray(full[variable].values, dtype=float) + + n = 1 + len(contributions) + colors = get_palette(n) + ax.plot(wl, full_y, color=colors[0], linewidth=2.0, label="full") + for color, (name, ds_c) in zip(colors[1:], contributions.items(), strict=True): + y = np.asarray(ds_c[variable].values, dtype=float) + if "zout" in ds_c[variable].dims: + y = y[:, idx] + ax.plot(wl, y, color=color, linewidth=1.5, label=name) + ax.set_xlabel("Wavelength (nm)") + ax.set_ylabel(f"{variable} contribution") + ax.legend(loc="best") + if save_path is not None: + save(fig, Path(save_path)) + return fig, ax diff --git a/tests/test_viz_attribution.py b/tests/test_viz_attribution.py new file mode 100644 index 0000000..56b160e --- /dev/null +++ b/tests/test_viz_attribution.py @@ -0,0 +1,44 @@ +"""Tests for the component-attribution plot (headless, synthetic data).""" +from __future__ import annotations + +from dataclasses import dataclass + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import xarray as xr + +from pyradtran.viz.attribution import plot_component_attribution + + +@dataclass(frozen=True) +class _Result: + full: xr.Dataset + contributions: dict + + +def _ds(value: float) -> xr.Dataset: + return xr.Dataset( + {"edir": ("wavelength", np.array([value]))}, + coords={"wavelength": np.array([500.0])}, + ) + + +def test_plot_component_attribution_draws_full_plus_each_block(): + result = _Result( + full=_ds(1.0), + contributions={"soot": _ds(0.6), "dust": _ds(0.4)}, + ) + fig, ax = plot_component_attribution(result, variable="edir") + # 1 full line + 2 contribution lines = 3 lines. + assert len(ax.lines) == 3 + plt.close(fig) + + +def test_plot_component_attribution_works_with_plain_dict_duck_type(): + result = {"full": _ds(1.0), "contributions": {"a": _ds(1.0)}} + fig, ax = plot_component_attribution(result, variable="edir") + assert len(ax.lines) == 2 + plt.close(fig) From b844ccddff23f4d38840e06a9f4c79d43a67150a Mon Sep 17 00:00:00 2001 From: Fan Zhang Date: Fri, 26 Jun 2026 18:03:08 +0800 Subject: [PATCH 10/15] feat(viz): wire public exports for viz/postprocess/workflow (matplotlib stays optional) --- pyproject.toml | 1 + src/pyradtran/__init__.py | 16 ++++++++++++ src/pyradtran/core/postprocess.py | 8 ++---- src/pyradtran/viz/__init__.py | 26 ++++++++++++++++++++ src/pyradtran/viz/attribution.py | 4 +-- tests/test_public_api.py | 41 +++++++++++++++++++++++++++++++ 6 files changed, 87 insertions(+), 9 deletions(-) create mode 100644 tests/test_public_api.py diff --git a/pyproject.toml b/pyproject.toml index ccb9f55..6d7ad44 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,6 +27,7 @@ dependencies = [ "numpy", "netCDF4", ] +# viz modules import matplotlib lazily; matplotlib stays in the `plot` extra (not core). [project.optional-dependencies] dev = [ diff --git a/src/pyradtran/__init__.py b/src/pyradtran/__init__.py index 79dd230..0f1df93 100644 --- a/src/pyradtran/__init__.py +++ b/src/pyradtran/__init__.py @@ -17,6 +17,13 @@ run_with_opac_custom, run_with_opac_preset, ) +from pyradtran.core.postprocess import ( + BudgetResult, + add_budget_vars, + compute_budget, + evaluate_blocks_on_grid, + evaluate_composite_on_grid, +) from pyradtran.core.runner import Runner, RunnerConfig from pyradtran.data import DataResolver from pyradtran.models import ThreeDConfig @@ -51,6 +58,7 @@ od_to_mass_profile, ) from pyradtran.scene import Scene +from pyradtran.workflow import AttributionResult, compute_component_attribution __all__ = [ "Scene", @@ -85,6 +93,14 @@ "ExponentialProfile", "TabulatedProfile", "od_to_mass_profile", + # Post-processing + workflow entry points + "BudgetResult", + "add_budget_vars", + "compute_budget", + "evaluate_composite_on_grid", + "evaluate_blocks_on_grid", + "compute_component_attribution", + "AttributionResult", # Convenience functions "run_3d", "run_solar_transmittance", diff --git a/src/pyradtran/core/postprocess.py b/src/pyradtran/core/postprocess.py index 1dfc006..a231445 100644 --- a/src/pyradtran/core/postprocess.py +++ b/src/pyradtran/core/postprocess.py @@ -77,9 +77,7 @@ def _layer_centers(z_km: np.ndarray) -> np.ndarray: return 0.5 * (z[:-1] + z[1:]) -def evaluate_composite_on_grid( - comp, wl_um, z_km, n_legendre: int = 32 -) -> xr.Dataset: +def evaluate_composite_on_grid(comp, wl_um, z_km, n_legendre: int = 32) -> xr.Dataset: """Evaluate the mixed composite optics on a (wavelength, layer) grid. Calls ``comp.evaluate(...)`` (pure analytic mixing — no RT) and wraps the @@ -104,9 +102,7 @@ def evaluate_composite_on_grid( ) -def evaluate_blocks_on_grid( - comp, wl_um, z_km, n_legendre: int = 32 -) -> dict[str, xr.Dataset]: +def evaluate_blocks_on_grid(comp, wl_um, z_km, n_legendre: int = 32) -> dict[str, xr.Dataset]: """Per-piece ``(tau, rho)`` on the grid; one dataset per block name. ``tau`` comes from each piece's ``to_layer_optics``; ``rho_kg_m3`` is added diff --git a/src/pyradtran/viz/__init__.py b/src/pyradtran/viz/__init__.py index db2f30d..2c9b640 100644 --- a/src/pyradtran/viz/__init__.py +++ b/src/pyradtran/viz/__init__.py @@ -1 +1,27 @@ """pyRadtran visualization layer (pure data-in -> fig-out; lazy matplotlib).""" + +from pyradtran.viz._style import get_palette, save, set_theme +from pyradtran.viz.attribution import AttributionLike, plot_component_attribution +from pyradtran.viz.composite import plot_block_profiles, plot_composite_optics +from pyradtran.viz.rt import ( + plot_budget, + plot_flux_profile, + plot_heating_rate, + plot_rt_overview, + plot_spectral, +) + +__all__ = [ + "set_theme", + "get_palette", + "save", + "plot_spectral", + "plot_flux_profile", + "plot_heating_rate", + "plot_budget", + "plot_rt_overview", + "plot_composite_optics", + "plot_block_profiles", + "plot_component_attribution", + "AttributionLike", +] diff --git a/src/pyradtran/viz/attribution.py b/src/pyradtran/viz/attribution.py index 674c188..6cf6e90 100644 --- a/src/pyradtran/viz/attribution.py +++ b/src/pyradtran/viz/attribution.py @@ -28,9 +28,7 @@ def _get(obj, key): return getattr(obj, key) if isinstance(obj, dict): return obj[key] - raise TypeError( - f"Expected an object with .{key} or ['{key}']; got {type(obj).__name__}" - ) + raise TypeError(f"Expected an object with .{key} or ['{key}']; got {type(obj).__name__}") def plot_component_attribution( diff --git a/tests/test_public_api.py b/tests/test_public_api.py new file mode 100644 index 0000000..b7713fc --- /dev/null +++ b/tests/test_public_api.py @@ -0,0 +1,41 @@ +"""Smoke test: public API is importable without matplotlib installed at import time. + +(`import pyradtran` must not pull matplotlib; viz symbols import lazily.) +""" +from __future__ import annotations + + +def test_pyradtran_imports_without_matplotlib_dependency(): + import importlib + + import pyradtran + + importlib.reload(pyradtran) # no exception -> top-level import is matplotlib-free + + +def test_public_symbols_exist(): + import pyradtran + + for name in ( + "add_budget_vars", + "compute_budget", + "evaluate_composite_on_grid", + "evaluate_blocks_on_grid", + "compute_component_attribution", + "AttributionResult", + ): + assert hasattr(pyradtran, name), f"pyradtran missing public symbol: {name}" + + +def test_viz_namespace_exposes_plot_functions(): + from pyradtran.viz import ( # noqa: F401 + plot_block_profiles, + plot_budget, + plot_component_attribution, + plot_composite_optics, + plot_flux_profile, + plot_heating_rate, + plot_rt_overview, + plot_spectral, + set_theme, + ) From 173ce6c02c06f2af477f2e8bc0d88b6413a2d379 Mon Sep 17 00:00:00 2001 From: Fan Zhang Date: Fri, 26 Jun 2026 20:25:08 +0800 Subject: [PATCH 11/15] docs(viz): fix misleading 'column sum' label (spectral sum) and heating-rate docstring (code review) --- src/pyradtran/viz/composite.py | 10 ++++++++-- src/pyradtran/viz/rt.py | 7 +++++-- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/src/pyradtran/viz/composite.py b/src/pyradtran/viz/composite.py index 739061f..14cbf35 100644 --- a/src/pyradtran/viz/composite.py +++ b/src/pyradtran/viz/composite.py @@ -15,7 +15,7 @@ "g": "Asymmetry parameter g", } _BLOCK_QUANTITY_LABELS = { - "tau": "Optical depth τ (column sum)", + "tau": "Optical depth τ (spectral sum)", "rho": "Mass concentration (kg/m³)", } @@ -56,7 +56,13 @@ def plot_block_profiles( ax=None, save_path=None, ): - """Per-block τ(z) (column-summed over wavelength) or ρ(z) vs altitude.""" + """Per-block τ(z) (spectrally summed over wavelength) or ρ(z) vs altitude. + + The τ curve is ``tau.sum(axis=0)`` — the per-layer sum over wavelength, a + rough "spectral optical burden per layer" overview. For a physically + meaningful τ(z) at a specific band, slice the block dataset to that + wavelength before passing it in. + """ fig, ax = _ensure_axes(ax) colors = get_palette(len(per_block_ds_dict)) for color, (name, bds) in zip(colors, per_block_ds_dict.items(), strict=True): diff --git a/src/pyradtran/viz/rt.py b/src/pyradtran/viz/rt.py index 276bf82..9a82b9a 100644 --- a/src/pyradtran/viz/rt.py +++ b/src/pyradtran/viz/rt.py @@ -93,8 +93,11 @@ def plot_heating_rate( ax=None, save_path=None, ): - """Plot heating rate vs altitude (km). If ``wavelength_nm`` is None, plot all - wavelengths as faint background lines plus the nearest-wavelength bold line.""" + """Plot heating rate vs altitude (km). + + If ``wavelength_nm`` is None, plot all wavelengths as faint background + lines; otherwise plot the nearest-wavelength line bold. + """ if HEATING_RATE_COLUMN not in ds.data_vars: raise ValueError( f"Dataset has no '{HEATING_RATE_COLUMN}' variable; request heating-rate " From 0be4f4dce94b847dde247cb56e7b44209ab5f91b Mon Sep 17 00:00:00 2001 From: Fan Zhang Date: Fri, 26 Jun 2026 22:07:43 +0800 Subject: [PATCH 12/15] fix(optics): numpy 2.x compat for _trapz shim (np.trapz removed in 2.0) --- src/pyradtran/optics/mie.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/pyradtran/optics/mie.py b/src/pyradtran/optics/mie.py index 0e53bcd..3849b01 100644 --- a/src/pyradtran/optics/mie.py +++ b/src/pyradtran/optics/mie.py @@ -7,7 +7,12 @@ import numpy as np -_trapz = getattr(np, "trapezoid", np.trapz) +# numpy >=2.0 removed np.trapz in favor of np.trapezoid; import whichever exists. +# The eager getattr(np, "trapezoid", np.trapz) default would crash on numpy 2.x. +try: + from numpy import trapezoid as _trapz +except ImportError: # numpy < 2.0 + from numpy import trapz as _trapz # type: ignore[attr-defined] # Import here to avoid circular import: mie.py uses SizeDistribution # from this module, and aerosol_composite.py imports integrate_size_distribution. From 2b7dfee8160525cf019dba81649143a8fcedcc99 Mon Sep 17 00:00:00 2001 From: Fan Zhang Date: Sat, 27 Jun 2026 13:23:26 +0800 Subject: [PATCH 13/15] fix(viz/attribution): make AttributionLike read-only so frozen AttributionResult satisfies it --- src/pyradtran/viz/attribution.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/pyradtran/viz/attribution.py b/src/pyradtran/viz/attribution.py index 6cf6e90..9ce3ca2 100644 --- a/src/pyradtran/viz/attribution.py +++ b/src/pyradtran/viz/attribution.py @@ -18,8 +18,13 @@ @runtime_checkable class AttributionLike(Protocol): - full: xr.Dataset - contributions: dict[str, xr.Dataset] + # Read-only: consumers only read these. Declaring as properties lets a + # frozen dataclass (AttributionResult) satisfy the protocol. + @property + def full(self) -> xr.Dataset: ... + + @property + def contributions(self) -> dict[str, xr.Dataset]: ... def _get(obj, key): From 19bfc802b0eb080fa1c2d1cd6445584a11cb24ce Mon Sep 17 00:00:00 2001 From: Fan Zhang Date: Sat, 27 Jun 2026 13:23:26 +0800 Subject: [PATCH 14/15] feat(examples): add comprehensive pyRadtran viz+workflow demo (MieSpecies -> DISORT -> all plots + attribution) --- examples/multicomponent_viz/.gitignore | 3 + examples/multicomponent_viz/run_demo.py | 280 ++++++++++++++++++++++++ 2 files changed, 283 insertions(+) create mode 100644 examples/multicomponent_viz/.gitignore create mode 100644 examples/multicomponent_viz/run_demo.py diff --git a/examples/multicomponent_viz/.gitignore b/examples/multicomponent_viz/.gitignore new file mode 100644 index 0000000..179b486 --- /dev/null +++ b/examples/multicomponent_viz/.gitignore @@ -0,0 +1,3 @@ +output/ +__pycache__/ +*.pyc diff --git a/examples/multicomponent_viz/run_demo.py b/examples/multicomponent_viz/run_demo.py new file mode 100644 index 0000000..7043bf8 --- /dev/null +++ b/examples/multicomponent_viz/run_demo.py @@ -0,0 +1,280 @@ +"""Comprehensive pyRadtran demo: LEGO blocks -> DISORT -> full viz + workflow. + +Self-contained (pyRadtran APIs only, no aerosol3d dependency). Three aerosol +blocks with contrasting optics are built from refractive indices + lognormal +size distributions via ``MieSpecies``, externally mixed into one +``CompositeAerosol``, and run through DISORT. The script then exercises the +FULL ``pyradtran.viz`` plot surface and the component-attribution workflow. + + Composite diagnostics (analytic mixing, no RT): + - evaluate_composite_on_grid -> plot_composite_optics (tau / ssa / g) + - evaluate_blocks_on_grid -> plot_block_profiles (per-block tau(z), rho(z)) + + RT result plots (full composite, real DISORT): + - plot_spectral, plot_flux_profile, plot_budget (via add_budget_vars), + plot_heating_rate (if libRadtran emits it), plot_rt_overview + + Workflow (component attribution, leave-one-out): + - compute_component_attribution -> plot_component_attribution + +Requires libRadtran. Set PYRADTRAN_DATA_PATH to its data/ dir (or rely on the +bundled resolver). + +Usage: + python run_demo.py +""" + +import logging +import os +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +from pyradtran import Runner, Scene +from pyradtran.core.output_parser import HEATING_RATE_COLUMN +from pyradtran.core.postprocess import ( + add_budget_vars, + evaluate_blocks_on_grid, + evaluate_composite_on_grid, +) +from pyradtran.models.aerosol_composite import ( + CompositeAerosol, + IntegrationConfig, + MieSpecies, + RefractiveIndex, + SizeDistribution, +) +from pyradtran.models.blocks import PlacedBlock, od_to_mass_profile +from pyradtran.viz import ( + plot_block_profiles, + plot_budget, + plot_component_attribution, + plot_composite_optics, + plot_flux_profile, + plot_heating_rate, + plot_rt_overview, + plot_spectral, + save, + set_theme, +) +from pyradtran.workflow import compute_component_attribution + +logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") +logger = logging.getLogger(__name__) + +EXAMPLE_DIR = Path(__file__).parent +OUTPUT_DIR = EXAMPLE_DIR / "output" +OUTPUT_DIR.mkdir(exist_ok=True) + +# --- Spectral / grid (composite optics on a coarse grid; DISORT resamples) --- +WAVELENGTHS_UM = [0.40, 0.45, 0.50, 0.55, 0.60, 0.65, 0.70] +N_LEGENDRE = 32 +ALTITUDE_GRID_KM = [10.0, 8.0, 6.0, 4.0, 2.0, 1.0, 0.0] # descending (TOA -> surface) +REF_NM = 550.0 +ZOUT_LEVELS = [0, 1, 2, 4, 6, 8, 10, "toa"] + +# --- RT scene --- +SZA_DEG = 30.0 +ALBEDO = 0.1 +N_STREAMS = 16 +SCENE_KW = { + "atmosphere": {"profile": "us", "altitude": 0.0}, + "source": {"sza": SZA_DEG}, + "wavelength": {"min_nm": 401.0, "max_nm": 699.0}, + "solver": { + "method": "disort", + "streams": N_STREAMS, + "disort_intcor": "moments", + "pseudospherical": True, + }, + "surface": {"albedo": ALBEDO}, + "output": { + "quantities": ["lambda", "edir", "edn", "eup"], + "format": "ascii", + "zout": ZOUT_LEVELS, + "heating_rate": "local", + }, +} + +# --- Three LEGO blocks: refractive index, lognormal size dist, placement --- +_WL_RI = [0.30, 0.40, 0.55, 0.70, 1.00] # µm; constant n/k across the band +BLOCKS = [ + { + "name": "black_carbon", + "n_real": [1.95] * 5, + "k_imag": [0.79] * 5, + "density": 1800.0, + "r_g_um": 0.10, + "sigma_g": 2.0, + "tau_550": 0.15, + "scale_height_km": 1.5, + }, + { + "name": "sulfate", + "n_real": [1.53] * 5, + "k_imag": [0.0] * 5, + "density": 1770.0, + "r_g_um": 0.15, + "sigma_g": 1.7, + "tau_550": 0.15, + "scale_height_km": 2.0, + }, + { + "name": "mineral_dust", + "n_real": [1.53] * 5, + "k_imag": [0.008] * 5, + "density": 2600.0, + "r_g_um": 0.50, + "sigma_g": 2.2, + "tau_550": 0.20, + "scale_height_km": 3.0, + }, +] + +_FLUX_VARS = {"edir", "edn", "eup", "udir", "udn", "uup"} + + +def _save(fig, name): + path = OUTPUT_DIR / name + save(fig, str(path), formats=("png",)) + plt.close(fig) + logger.info(" saved %s", path.name) + + +def build_composite() -> CompositeAerosol: + """Build MieSpecies for each block, invert its target OD@550 into a mass + profile via the API (od_to_mass_profile), and assemble the composite.""" + pieces = [] + for b in BLOCKS: + ri = RefractiveIndex(wavelength_um=_WL_RI, n_real=b["n_real"], k_imag=b["k_imag"]) + sd = SizeDistribution(kind="lognormal", params={"r_g_um": b["r_g_um"], "sigma_g": b["sigma_g"]}) + species = MieSpecies( + refractive_index=ri, + size_distribution=sd, + particle_density_kg_m3=b["density"], + integration_config=IntegrationConfig(), + name=b["name"], + ) + profile = od_to_mass_profile( + species, + tau_ref=b["tau_550"], + ref_nm=REF_NM, + altitude_km=ALTITUDE_GRID_KM, + scale_height_km=b["scale_height_km"], + ) + pieces.append(PlacedBlock(block=species, profile=profile)) + logger.info( + " block %-14s m=%.2f%+.3fi r_g=%.2fµm tau@550=%.2f H=%.1fkm", + b["name"], b["n_real"][0], b["k_imag"][0], b["r_g_um"], b["tau_550"], b["scale_height_km"], + ) + return CompositeAerosol( + pieces=pieces, + wavelength_grid_um=list(WAVELENGTHS_UM), + altitude_grid_km=list(ALTITUDE_GRID_KM), + n_legendre=N_LEGENDRE, + output_dir=OUTPUT_DIR, + ) + + +def build_scene(aerosol: CompositeAerosol) -> Scene: + c = SCENE_KW + return ( + Scene() + .set_atmosphere(profile=c["atmosphere"]["profile"], altitude=c["atmosphere"]["altitude"]) + .set_source_solar(sza=c["source"]["sza"]) + .set_wavelength(c["wavelength"]["min_nm"], c["wavelength"]["max_nm"]) + .set_solver( + method=c["solver"]["method"], + streams=c["solver"]["streams"], + disort_intcor=c["solver"].get("disort_intcor"), + pseudospherical=c["solver"].get("pseudospherical", False), + ) + .set_surface(albedo=c["surface"]["albedo"]) + .set_output(**c["output"]) + .set_aerosol(aerosol) + ) + + +def _resolve_heating_var(ds): + if HEATING_RATE_COLUMN in ds.data_vars: + return HEATING_RATE_COLUMN + extras = [v for v in ds.data_vars if v not in _FLUX_VARS and v != "wavelength"] + if len(extras) == 1: + return extras[0] + return None + + +def main(): + set_theme("publication") + data_path = os.environ.get("PYRADTRAN_DATA_PATH") # None -> bundled resolver + if data_path is None: + logger.info("PYRADTRAN_DATA_PATH unset; using bundled DataResolver") + else: + logger.info("using libRadtran data: %s", data_path) + + logger.info("=== Building 3-block composite (pyRadtran MieSpecies) ===") + composite = build_composite() + + # --- Composite diagnostics (analytic) --- + logger.info("=== Composite diagnostics (analytic) ===") + grid_ds = evaluate_composite_on_grid( + composite, WAVELENGTHS_UM, ALTITUDE_GRID_KM, n_legendre=N_LEGENDRE + ) + for q in ("tau", "ssa", "g"): + fig, _ = plot_composite_optics(grid_ds, quantity=q) + _save(fig, f"composite_{q}.png") + block_dict = evaluate_blocks_on_grid( + composite, WAVELENGTHS_UM, ALTITUDE_GRID_KM, n_legendre=N_LEGENDRE + ) + fig, _ = plot_block_profiles(block_dict, quantity="tau") + _save(fig, "block_tau_profiles.png") + fig, _ = plot_block_profiles(block_dict, quantity="rho") + _save(fig, "block_rho_profiles.png") + + # --- RT (full composite) --- + logger.info("=== RT run (full composite, DISORT) ===") + rt = Runner.execute(build_scene(composite), data_path=data_path) + logger.info(" RT data_vars=%s dims=%s", list(rt.data_vars), dict(rt.sizes)) + + fig, _ = plot_spectral(rt) + _save(fig, "rt_spectral.png") + fig, _ = plot_flux_profile(rt, variable="edir", wavelength_nm=550.0) + _save(fig, "rt_flux_profile_edir.png") + rt_budget = add_budget_vars(rt) + fig, _ = plot_budget(rt_budget) + _save(fig, "rt_budget.png") + + hvar = _resolve_heating_var(rt) + if hvar is not None: + if hvar != HEATING_RATE_COLUMN: + rt = rt.rename({hvar: HEATING_RATE_COLUMN}) + fig, _ = plot_heating_rate(rt, wavelength_nm=550.0) + _save(fig, "rt_heating_rate.png") + else: + logger.warning(" no heating-rate column in RT output; skipping plot_heating_rate") + + rt.to_netcdf(str(OUTPUT_DIR / "rt_full.nc")) + fig, _ = plot_rt_overview(rt, wavelength_nm=550.0) + _save(fig, "rt_overview.png") + + # --- Component attribution (leave-one-out) --- + logger.info("=== Component attribution (N+1 DISORT runs) ===") + + def execute_many(scenes): + # Sequential: Runner.execute_many (ProcessPoolExecutor) swallows exceptions + # and fails on CompositeAerosol scenes (pickling). Sequential surfaces errors. + return [Runner.execute(s, data_path=data_path) for s in scenes] + + result = compute_component_attribution(build_scene, composite, execute_many) + logger.info(" contributions: %s", list(result.contributions)) + fig, _ = plot_component_attribution(result, variable="edir", level="surface") + _save(fig, "attribution_edir.png") + + logger.info("=== Comprehensive demo complete. Figures in %s ===", OUTPUT_DIR) + + +if __name__ == "__main__": + main() From 16ede57208838ee20c08a7c2a835b28c9ffdf8f2 Mon Sep 17 00:00:00 2001 From: Fan Zhang Date: Sat, 27 Jun 2026 14:08:53 +0800 Subject: [PATCH 15/15] fix(ci): green CI - ci.yml: install .[dev,plot] so matplotlib-backed viz tests run - test_bulk_species: importorskip('Aerosol3D') (aerosol3d absent in CI) - aerosol.py: OpacPresetName(str, Enum) -> StrEnum (ruff UP042) - scene.py: set_satellite passed 'source' twice -> build defaults dict (fixes 'SourceConfig() got multiple values for source') - test_composite_aerosol_unit: np.trapz -> np.trapezoid (numpy 2.x) - test_postprocess_grid: drop unused xarray import (F401) - ruff format (12 files) --- .github/workflows/ci.yml | 2 +- src/pyradtran/models/aerosol.py | 4 ++-- src/pyradtran/models/blocks.py | 14 ++++-------- src/pyradtran/scene.py | 7 +++++- tests/test_blocks.py | 32 ++++++++++++++++++++-------- tests/test_bulk_species.py | 3 +++ tests/test_composite_aerosol_unit.py | 6 +++--- tests/test_output_parser_zout.py | 1 + tests/test_postprocess_budget.py | 1 + tests/test_postprocess_grid.py | 2 +- tests/test_public_api.py | 1 + tests/test_viz_attribution.py | 1 + tests/test_viz_composite.py | 27 ++++++++++++++++------- tests/test_viz_rt_budget.py | 1 + tests/test_viz_rt_lines.py | 1 + tests/test_viz_style.py | 1 + tests/test_workflow_attribution.py | 5 +++-- 17 files changed, 72 insertions(+), 37 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a86ec09..490af62 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -18,7 +18,7 @@ jobs: - name: Install dependencies run: | python -m pip install --upgrade pip - pip install -e ".[dev]" + pip install -e ".[dev,plot]" - name: Run tests run: pytest tests/ -v --tb=short --cov=pyradtran --cov-report=xml - name: Upload coverage diff --git a/src/pyradtran/models/aerosol.py b/src/pyradtran/models/aerosol.py index d9f2269..edb90cf 100644 --- a/src/pyradtran/models/aerosol.py +++ b/src/pyradtran/models/aerosol.py @@ -13,7 +13,7 @@ from __future__ import annotations from abc import abstractmethod -from enum import Enum +from enum import StrEnum from pydantic import BaseModel, Field, model_validator @@ -47,7 +47,7 @@ def _validate_opac_species_names(names: list[str]) -> None: ) -class OpacPresetName(str, Enum): +class OpacPresetName(StrEnum): """OPAC preset mixture profile names. These correspond to files in data/aerosol/OPAC/standard_aerosol_files/. diff --git a/src/pyradtran/models/blocks.py b/src/pyradtran/models/blocks.py index a11fb65..4a5b209 100644 --- a/src/pyradtran/models/blocks.py +++ b/src/pyradtran/models/blocks.py @@ -168,9 +168,7 @@ class PlacedBlock: def name(self) -> str: return self.block.name - def to_layer_optics( - self, wl_um: np.ndarray, altitude_km, n_legendre: int = 32 - ) -> LayerOptics: + def to_layer_optics(self, wl_um: np.ndarray, altitude_km, n_legendre: int = 32) -> LayerOptics: wl = np.asarray(wl_um, dtype=float) z = np.asarray(altitude_km, dtype=float) n_wl = wl.shape[0] @@ -201,7 +199,7 @@ def to_layer_optics( # Henyey-Greenstein fill: g_l = g^l (matches LoadedSpecies / _fill_hg_moments) moments_layer = np.zeros((n_wl, n_layer, n_legendre)) for l in range(n_legendre): - moments_layer[:, :, l] = g_layer ** l + moments_layer[:, :, l] = g_layer**l # Apply per-block modify (tau / ssa / gg scale or set). for entry in self.modify: @@ -221,9 +219,7 @@ def to_layer_optics( else: g_layer = np.full_like(g_layer, entry.value) - return LayerOptics( - tau=tau, ssa=ssa_layer, g=g_layer, legendre_moments=moments_layer - ) + return LayerOptics(tau=tau, ssa=ssa_layer, g=g_layer, legendre_moments=moments_layer) @dataclass(frozen=True) @@ -238,9 +234,7 @@ class DirectLayerOpticsBlock: master_path: str name: str = "explicit_file" - def to_layer_optics( - self, wl_um: np.ndarray, altitude_km, n_legendre: int = 32 - ) -> LayerOptics: + def to_layer_optics(self, wl_um: np.ndarray, altitude_km, n_legendre: int = 32) -> LayerOptics: from pyradtran.optics.layer_writer import read_explicit_aerosol tau, ssa, g, moments, _wl_file, _alt_file = read_explicit_aerosol(self.master_path) diff --git a/src/pyradtran/scene.py b/src/pyradtran/scene.py index 01e0ab9..f8f9955 100644 --- a/src/pyradtran/scene.py +++ b/src/pyradtran/scene.py @@ -234,7 +234,12 @@ def set_satellite( if new.source is not None: new.source = new.source.model_copy(update={**sat_updates, **source_kwargs}) else: - new.source = SourceConfig(source="solar", sza=0.0, **sat_updates, **source_kwargs) + # Build defaults then let caller overrides win (avoids passing + # `source` twice when source_kwargs already supplies it). + defaults: dict = {"source": "solar", "sza": 0.0} + defaults.update(sat_updates) + defaults.update(source_kwargs) + new.source = SourceConfig(**defaults) return new def set_dynamic( diff --git a/tests/test_blocks.py b/tests/test_blocks.py index 76e3561..f2b875e 100644 --- a/tests/test_blocks.py +++ b/tests/test_blocks.py @@ -55,8 +55,11 @@ def test_inverts_target_od_discrete(self): block = _mie_block() alt = [4.0, 2.0, 0.0] # 2 layers mp = od_to_mass_profile( - block, tau_ref=0.5, ref_nm=550.0, - altitude_km=alt, scale_height_km=2.0, + block, + tau_ref=0.5, + ref_nm=550.0, + altitude_km=alt, + scale_height_km=2.0, ) assert isinstance(mp, MassProfile) dz_m = -np.diff(alt) * 1000.0 # layer thicknesses in m @@ -70,17 +73,22 @@ def test_raises_on_nonpositive_beta(self): # from MieSpecies, so check the guard via a tiny stub block. class _ZeroBlock: name = "zero" + def intensive(self, wl_um, n_legendre=32): from pyradtran.models.aerosol_composite import SpeciesOptics - return SpeciesOptics(beta_ext_per_mass=np.array([0.0]), - ssa=np.array([0.0]), g=np.array([0.0])) + + return SpeciesOptics( + beta_ext_per_mass=np.array([0.0]), ssa=np.array([0.0]), g=np.array([0.0]) + ) + @property def mass_per_particle_kg(self) -> float: return 1.0 with pytest.raises(ValueError): - od_to_mass_profile(_ZeroBlock(), tau_ref=0.5, ref_nm=550.0, - altitude_km=[4.0, 0.0], scale_height_km=2.0) + od_to_mass_profile( + _ZeroBlock(), tau_ref=0.5, ref_nm=550.0, altitude_km=[4.0, 0.0], scale_height_km=2.0 + ) class TestAerosolBlockProtocol: @@ -144,7 +152,9 @@ def test_roundtrip_write_read(self): ) sd = SizeDistribution(kind="monodisperse", params={"radius_um": 0.5}) block = MieSpecies( - refractive_index=ri, size_distribution=sd, particle_density_kg_m3=1000.0, + refractive_index=ri, + size_distribution=sd, + particle_density_kg_m3=1000.0, integration_config=IntegrationConfig(n_radius_grid=30), ) alt = [6.0, 4.0, 2.0, 0.0] # 3 layers @@ -155,9 +165,13 @@ def test_roundtrip_write_read(self): with tempfile.TemporaryDirectory() as d: master = write_explicit_aerosol( - tau=original.tau, ssa=original.ssa, g=original.g, + tau=original.tau, + ssa=original.ssa, + g=original.g, legendre_moments=original.legendre_moments, - wavelength_um=wl, altitude_km=np.asarray(alt), output_dir=Path(d), + wavelength_um=wl, + altitude_km=np.asarray(alt), + output_dir=Path(d), source_signatures=["test"], ) recovered = DirectLayerOpticsBlock( diff --git a/tests/test_bulk_species.py b/tests/test_bulk_species.py index 6989201..97b0300 100644 --- a/tests/test_bulk_species.py +++ b/tests/test_bulk_species.py @@ -3,6 +3,9 @@ from types import SimpleNamespace import numpy as np +import pytest + +pytest.importorskip("Aerosol3D") # needs the aerosol3d sibling repo (absent in CI) from Aerosol3D.bulk.datastructs import SizeDistribution from pyradtran.models.aerosol_composite import BulkSpecies diff --git a/tests/test_composite_aerosol_unit.py b/tests/test_composite_aerosol_unit.py index 26b71b2..a292471 100644 --- a/tests/test_composite_aerosol_unit.py +++ b/tests/test_composite_aerosol_unit.py @@ -55,7 +55,7 @@ def test_lognormal_normalization(self): sd = SizeDistribution(kind="lognormal", params={"r_g_um": 0.5, "sigma_g": 2.0}) r = np.logspace(-2, 2, 10000) dn = sd.evaluate(r) - total = np.trapz(dn, r) + total = np.trapezoid(dn, r) assert np.isclose(total, 1.0, rtol=0.01) def test_monodisperse_peak_location(self): @@ -81,7 +81,7 @@ def test_modified_gamma_normalization(self): ) r = np.logspace(-2, 2, 10000) dn = sd.evaluate(r) - total = np.trapz(dn, r) + total = np.trapezoid(dn, r) assert np.isclose(total, 1.0, rtol=0.01) def test_discrete_normalization(self): @@ -91,7 +91,7 @@ def test_discrete_normalization(self): ) r = np.logspace(-2, 2, 10000) dn = sd.evaluate(r) - total = np.trapz(dn, r) + total = np.trapezoid(dn, r) assert np.isclose(total, 1.0, rtol=0.01) def test_modified_gamma_invalid_params_raises(self): diff --git a/tests/test_output_parser_zout.py b/tests/test_output_parser_zout.py index 0a55093..80eafcb 100644 --- a/tests/test_output_parser_zout.py +++ b/tests/test_output_parser_zout.py @@ -1,4 +1,5 @@ """Tests for physical zout-altitude retention and the heating-rate column.""" + from __future__ import annotations from pathlib import Path diff --git a/tests/test_postprocess_budget.py b/tests/test_postprocess_budget.py index 64e229f..181b17e 100644 --- a/tests/test_postprocess_budget.py +++ b/tests/test_postprocess_budget.py @@ -1,4 +1,5 @@ """Tests for T/R/A budget derivation in pyradtran.core.postprocess.""" + from __future__ import annotations import numpy as np diff --git a/tests/test_postprocess_grid.py b/tests/test_postprocess_grid.py index a628103..d99101c 100644 --- a/tests/test_postprocess_grid.py +++ b/tests/test_postprocess_grid.py @@ -1,8 +1,8 @@ """Tests for composite/block grid evaluation helpers.""" + from __future__ import annotations import numpy as np -import xarray as xr from pyradtran.core.postprocess import evaluate_blocks_on_grid, evaluate_composite_on_grid from pyradtran.models.aerosol_composite import LayerOptics diff --git a/tests/test_public_api.py b/tests/test_public_api.py index b7713fc..8b7276e 100644 --- a/tests/test_public_api.py +++ b/tests/test_public_api.py @@ -2,6 +2,7 @@ (`import pyradtran` must not pull matplotlib; viz symbols import lazily.) """ + from __future__ import annotations diff --git a/tests/test_viz_attribution.py b/tests/test_viz_attribution.py index 56b160e..e3e45e3 100644 --- a/tests/test_viz_attribution.py +++ b/tests/test_viz_attribution.py @@ -1,4 +1,5 @@ """Tests for the component-attribution plot (headless, synthetic data).""" + from __future__ import annotations from dataclasses import dataclass diff --git a/tests/test_viz_composite.py b/tests/test_viz_composite.py index dd68ef3..8b0450c 100644 --- a/tests/test_viz_composite.py +++ b/tests/test_viz_composite.py @@ -1,4 +1,5 @@ """Tests for composite/block diagnostic plots (headless, synthetic data).""" + from __future__ import annotations import matplotlib @@ -29,16 +30,26 @@ def _block_dict() -> dict[str, xr.Dataset]: alt = np.array([4.0, 2.0, 0.5]) return { "soot": xr.Dataset( - {"tau": (("wavelength", "layer"), np.full((2, 3), 0.5)), - "rho_kg_m3": ("layer", np.array([1e-6, 5e-7, 1e-7]))}, - coords={"wavelength": np.array([0.3, 0.5]), "layer": np.arange(3), - "altitude_km": ("layer", alt)}, + { + "tau": (("wavelength", "layer"), np.full((2, 3), 0.5)), + "rho_kg_m3": ("layer", np.array([1e-6, 5e-7, 1e-7])), + }, + coords={ + "wavelength": np.array([0.3, 0.5]), + "layer": np.arange(3), + "altitude_km": ("layer", alt), + }, ), "dust": xr.Dataset( - {"tau": (("wavelength", "layer"), np.full((2, 3), 0.2)), - "rho_kg_m3": ("layer", np.array([2e-6, 1e-6, 2e-7]))}, - coords={"wavelength": np.array([0.3, 0.5]), "layer": np.arange(3), - "altitude_km": ("layer", alt)}, + { + "tau": (("wavelength", "layer"), np.full((2, 3), 0.2)), + "rho_kg_m3": ("layer", np.array([2e-6, 1e-6, 2e-7])), + }, + coords={ + "wavelength": np.array([0.3, 0.5]), + "layer": np.arange(3), + "altitude_km": ("layer", alt), + }, ), } diff --git a/tests/test_viz_rt_budget.py b/tests/test_viz_rt_budget.py index 7af187e..d4f347d 100644 --- a/tests/test_viz_rt_budget.py +++ b/tests/test_viz_rt_budget.py @@ -1,4 +1,5 @@ """Tests for budget and overview plots (headless, synthetic data).""" + from __future__ import annotations import matplotlib diff --git a/tests/test_viz_rt_lines.py b/tests/test_viz_rt_lines.py index 97e7c2e..1f36dcd 100644 --- a/tests/test_viz_rt_lines.py +++ b/tests/test_viz_rt_lines.py @@ -1,4 +1,5 @@ """Tests for RT line/profile plots (headless, synthetic data).""" + from __future__ import annotations import matplotlib diff --git a/tests/test_viz_style.py b/tests/test_viz_style.py index 89f6147..adf76fd 100644 --- a/tests/test_viz_style.py +++ b/tests/test_viz_style.py @@ -1,4 +1,5 @@ """Tests for viz theme/palette/save (headless).""" + from __future__ import annotations from pathlib import Path diff --git a/tests/test_workflow_attribution.py b/tests/test_workflow_attribution.py index 6191cc8..d643590 100644 --- a/tests/test_workflow_attribution.py +++ b/tests/test_workflow_attribution.py @@ -1,4 +1,5 @@ """Tests for component-attribution orchestration (no libRadtran).""" + from __future__ import annotations from dataclasses import dataclass @@ -38,8 +39,8 @@ def test_attribution_contribution_is_full_minus_removed(): # full=1.0; removing soot -> 0.7; removing dust -> 0.9 results = { ("soot", "dust"): _ds(1.0), - ("dust",): _ds(0.7), # soot removed - ("soot",): _ds(0.9), # dust removed + ("dust",): _ds(0.7), # soot removed + ("soot",): _ds(0.9), # dust removed } def build_scene(c):