diff --git a/CHANGELOG.md b/CHANGELOG.md index 3c71d0c..47afef9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 Highlights: new `pyradtran.data` data layer with bundled libRadtran data, a LEGO-style "blocks" aerosol API, OPAC folding with real Mie phase functions, and a publication-ready visualization/workflow suite. Contains **breaking changes** to the aerosol API (see Removed). +#### Example self-validation + analysis expansion +- `compute_energy_budget` + `assert_energy_conservation` (`pyradtran.core.postprocess`): column energy identity `F_inc = eup_TOA + (1-a)(edir+edn)_surf + F_abs_atm` with hard physical-bound assertions. +- `parse_heating_ascii` (`pyradtran.core.output_parser`): parses libRadtran's wide heating-rate format (K/day per zout per wavelength). `Runner.execute` auto-detects heating mode (scene has `heating_rate` but no `output_user` quantities) and dispatches to it. +- Viz: `plot_size_distributions`, `plot_phase_functions`, `plot_legendre_decay`, `plot_block_spectral_optics`, `plot_drf_spectral` (direct radiative forcing), `plot_spectral_attribution`. +- `examples/multicomponent_viz`: expanded from 10 to 23 plots (per-block spectral optics, scattering phase functions, Legendre decay, heating-rate profile, DRF via no-aerosol baseline, spectral attribution) with a hard energy-conservation assertion and a column-τ-stacking check. Heating rates require a second uvspec invocation (libRadtran's `heating_rate` mode replaces flux output rather than appending a column); the demo runs both and merges. +- `examples/multicomponent_viz/canonical.py`: shared scene config incl. `build_scene_heating` (used by both the demo and the regression test so they cannot drift). +- `tests/fixtures/multicomponent_baseline.json` + `tests/test_multicomponent_regression.py` (slow/uvspec-gated): end-to-end regression baseline for the canonical 3-block scene. +- `scripts/regen_baseline.py` to regenerate the baseline. + ### Added #### Data layer (`pyradtran.data`) @@ -49,7 +58,20 @@ Highlights: new `pyradtran.data` data layer with bundled libRadtran data, a LEGO - Comprehensive viz + workflow demo: `MieSpecies` → DISORT → all plots + attribution. - Sphinx user guide, expanded API reference, and rewritten README. +#### YAML configuration front-end (`pyradtran.config`) +- `config_version: 1` YAML schema (strict: scene / aerosol / analysis; block kinds `mie`/`bulk`/`opac_preset`/`explicit_layer`; placements `od_inversion`/`mass`/`exponential`/`tabulated`), `load_config` / `export_config`, and the `run_config` orchestrator (energy-conservation assertion, heating second run, DRF baseline, leave-one-out attribution, plot registry, NetCDF export). +- CLI `pyradtran run|validate|export-config` (also `python -m pyradtran`): `validate` checks a config without invoking uvspec; `export-config` writes canonical YAML. +- Hard round-trip guarantee, enforced by `tests/test_config_roundtrip.py`: a YAML config and its API-built twin emit byte-identical uvspec input and byte-identical `.master` layer files; `export_config` output re-loads identically. +- `examples/multicomponent_viz/canonical.yaml` — YAML twin of the canonical scene, generated from `canonical.py` by `make_yaml.py`. +#### Benchmarks (`pyradtran.benchmarks`) +- Randles et al. (2013) AeroCom shortwave RT intercomparison replication: `run_randles2013` (3 cases × 2 AFGL atmospheres × 2 SZAs × 2 bands = 24 uvspec runs; Å=1 power law, 0–2 km linear taper, HG phase function as PMOM `beta_l = g**l`, LBL-median normalization constants). +- Bundled reference `reference/randles2013_lbl.json`: 64 tabulated LBL values plus a `_meta` block (thresholds: fluxes ±8% PASS / ±12% WARN; RF ≤15% or ≤1.5 W/m²). +- `compare_benchmark` / `format_report` / `write_report` / `plot_benchmark_overlay`: per-row status classification, Markdown + CSV report, LBL-overlay PNG. +- Ready-to-run example `examples/randles2013_benchmark/run_benchmark.py`: matrix -> report + overlay PNG in one command; `--from-json` regenerates reports offline; non-zero exit on any FAIL (CI gate). +- Regression test `tests/test_benchmark_randles.py` (slow, uvspec-gated): real run vs libRadtran 2.0.6 — 68 comparison rows, 64 PASS / 0 WARN / 0 FAIL (+4 report-only n/a). +- Ratified deviation: benchmark results persist as JSON, not NetCDF (nested ragged dict). ### Changed +- `evaluate_blocks_on_grid` (`pyradtran.core.postprocess`) now also returns per-block `ssa` and `g` (was `tau`/`rho_kg_m3` only). - `OpacPreset` / `OpacCustom` are now `PlacedBlock` factories; OPAC convenience functions expose `output_dir`. - Runner resolves the libRadtran data path through `DataResolver`. - **Default `Output.format` is now `"ascii"`** (was `"netcdf"`): uvspec's NetCDF output is broken in many libRadtran builds (a 0-byte `.nc` from a libnetcdf ABI mismatch), so the previous default crashed every simulation that omitted `format=`. ASCII works everywhere and yields an equivalent `xarray.Dataset`; NetCDF remains available via `format="netcdf"`. The convenience functions no longer hardcode `"netcdf"` — they inherit the ASCII default. @@ -62,6 +84,9 @@ Highlights: new `pyradtran.data` data layer with bundled libRadtran data, a LEGO - **Breaking:** removed `ExternalFile` / `ExternalAerosol` and the `run_with_aerosol` convenience function — the only non-LEGO aerosol path. Pre-computed explicit files now go through `DirectLayerOpticsBlock` mixed into a `CompositeAerosol`. Also removed the unused `AerosolModel.set_tau_at_wvl` / `king_byrne` fields (never set by any code path). ### Fixed +- ASCII output parser names the 8th uvspec column `heating_rate` (was `col_7`) so `HEATING_RATE_COLUMN` resolves when libRadtran emits heating rates alongside the standard 7 flux columns. +- `_trapz` guard added to `tests/test_composite_aerosol_unit.py` normalization tests (numpy <2.0 compat; production `mie.py` already had the guard). +- `run_satellite` test now skips cleanly when libRadtran's `MPS` satellite-geometry file is absent (not bundled). - libRadtran integration issues for explicit aerosol coupling. - ASCII column-name mapping and aerosol-layer file cache collision. - Mixing-rule HG fallback now uses `g_l` (PMOM) Legendre moments consistently across species. @@ -70,7 +95,10 @@ Highlights: new `pyradtran.data` data layer with bundled libRadtran data, a LEGO - CI branch triggers, lint, and test assertions realigned for green CI. - Convenience functions no longer crash on the NetCDF parse path: they now use ASCII output (see **Changed**). `_parse_netcdf` raises a clear error pointing at `format='ascii'` when uvspec produces an empty/missing NetCDF file instead of xarray's cryptic backend-mismatch message. - `run_with_opac_preset` / `run_with_opac_custom` now set `disort_intcor="moments"` on their DISORT solver. The OPAC folding produces Legendre-moment phase functions, which DISORT rejects without this flag (`you need to specify 'disort_intcor moments'`); these functions previously crashed at runtime. +- **i550 wavelength-index bug** in the multicomponent demo and its regression test: `argmin(|wavelength - 0.55|)` compared an nm grid against 0.55 µm, silently selecting 401 nm everywhere a 550 nm scalar was labeled (energy log, heating log, baseline fixture). Also removed a spurious `× 1000` on the DRF wavelength axis (same µm/nm confusion) that made `np.interp(550, ...)` clamp to the 401 nm endpoint. The committed baseline now holds true 550 nm values (edir_surf 798.75, eup_toa 194.79 W/m², F_abs_atm 336.94 W/m²). +- `Runner.execute_many` no longer swallows worker failures: failing scenes were silently returned to callers as `(idx, exception)` result entries; the first failure now raises `RuntimeError(f"scene {i} failed")` chained to the original exception, after cancelling outstanding futures. +- `Runner.execute_many` uses a thread pool instead of a process pool: scenes carrying `CompositeAerosol` can hold unpicklable payloads (aerosol3D bulk size distributions store local closures), which crashed `ProcessPoolExecutor` pickling; uvspec runs as a subprocess and releases the GIL while waiting, so threads parallelize the invocations just as well and nothing needs pickling. ## [0.1.0] - 2026-05-12 ### Added diff --git a/docs/api.rst b/docs/api.rst index 6b1718b..17b30f4 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -253,3 +253,67 @@ Workflow :members: :undoc-members: :show-inheritance: + +YAML Configuration Front-end +---------------------------- + +.. automodule:: pyradtran.config + :members: + :undoc-members: + :show-inheritance: + +Configuration Schema +^^^^^^^^^^^^^^^^^^^^ + +.. automodule:: pyradtran.config.schema + :members: + :undoc-members: + :show-inheritance: + +Configuration Loader +^^^^^^^^^^^^^^^^^^^^ + +.. automodule:: pyradtran.config.loader + :members: + :undoc-members: + :show-inheritance: + +Config Orchestrator +^^^^^^^^^^^^^^^^^^^ + +.. automodule:: pyradtran.config.orchestrator + :members: + :undoc-members: + :show-inheritance: + +Command Line +^^^^^^^^^^^^ + +.. automodule:: pyradtran.config.cli + :members: + :undoc-members: + :show-inheritance: + +Benchmarks +---------- + +.. automodule:: pyradtran.benchmarks + :members: + :undoc-members: + :show-inheritance: + +Randles et al. (2013) +^^^^^^^^^^^^^^^^^^^^^ + +.. automodule:: pyradtran.benchmarks.randles2013 + :members: + :undoc-members: + :show-inheritance: + +Benchmark Comparison and Reporting +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +.. automodule:: pyradtran.benchmarks.compare + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/guide/benchmarks.rst b/docs/guide/benchmarks.rst new file mode 100644 index 0000000..c846e92 --- /dev/null +++ b/docs/guide/benchmarks.rst @@ -0,0 +1,209 @@ +Benchmarks +========== + +``pyradtran.benchmarks`` replicates published radiative-transfer +intercomparisons end-to-end with real uvspec runs and compares the results +against the published line-by-line (LBL) reference values. It currently ships +one benchmark: the shortwave aerosol experiment of Randles et al. (2013). + +Randles et al. (2013) protocol +------------------------------ + +Randles, C. A. et al., "Intercomparison of shortwave radiative transfer +schemes to improve the aerosol and surface albedo treatment in climate +models", Atmos. Chem. Phys. 13, 2347-2362 (2013) — the AeroCom RT experiment, +Tables 3, 5, A3 and A4. Three cases: + +.. list-table:: + :header-rows: 1 + + * - Case + - Aerosol + - AOD550 + - Ångström + - SSA + - g + * - ``case1`` + - none (Rayleigh only) + - — + - — + - — + - — + * - ``case2a`` + - lowest 2 km + - 0.2 + - 1 + - 1.0 + - 0.7 + * - ``case2b`` + - as 2a + - 0.2 + - 1 + - 0.8 + - 0.7 + +Fixed conditions: surface albedo 0.2 (Lambert), SZA 30° and 75°, AFGL subarctic +winter (``saw``) and tropics (``trop``) atmospheres, no clouds, and two bands — +broadband 0.2–4.0 µm (``bb``) and UV-VIS 0.2–0.7 µm (``uvvis``). The full +matrix, 3 cases × 2 atmospheres × 2 SZAs × 2 bands, is 24 uvspec invocations +(~20–60 s). The solver setup is DISORT with 16 streams, ``disort_intcor +moments``, pseudospherical geometry and reptran correlated-k with +band-integrated output; the phase function is supplied as 32 Legendre moments. + +Replication assumptions +----------------------- + +The paper leaves details to each participating model; this replication fixes +them as: + +1. Built-in AFGL profiles are used (the participants used their + individual-model profiles; the comparison tolerance absorbs this). +2. "Aerosol linear in the lowest 2 km" is a linear taper from the surface to + zero at 2 km, ``w(z) = max(0, 1 - z / 2 km)``. +3. ``g = 0.7`` Henyey–Greenstein phase function represented through PMOM + moments ``beta_l = g**l`` (the DISORT/libRadtran coefficient form). +4. Ångström exponent 1 means ``tau(lambda) = 0.2 * (lambda / 0.55 um)**-1``. + +Normalization +------------- + +Each result is divided by the same run's band TOA incident flux and multiplied +by the LBL median constants, which removes the solar-constant choice. +Formally ``F_norm = F_raw / edir_toa_raw(same run) * C[band][sza]`` with: + +.. list-table:: + :header-rows: 1 + + * - Constant (W/m²) + - SZA 30° + - SZA 75° + * - Broadband ``bb`` + - 1189.28 + - 355.43 + * - UV-VIS ``uvvis`` + - 563.38 + - 168.37 + +Reference and thresholds +------------------------ + +``pyradtran/benchmarks/reference/randles2013_lbl.json`` bundles the 64 +tabulated LBL values (per case and quantity, for the four configs +``saw30`` / ``saw75`` / ``trop30`` / ``trop75``) plus a ``_meta`` block holding +the citation and the comparison thresholds: + +.. list-table:: + :header-rows: 1 + + * - Quantity + - PASS + - WARN + - FAIL + * - Fluxes (everything except ``*_rf``) + - ``|rel| <= 8%`` + - ``|rel| <= 12%`` + - beyond + * - Radiative forcing (``toa_rf`` / ``sfc_rf``) + - ``|rel| <= 15%`` **or** ``|abs| <= 1.5 W/m²`` + - ``|rel| <= 25%`` or ``|abs| <= 3 W/m²`` + - beyond both + +Case-2a atmospheric forcing (``toa_rf - sfc_rf``) is below 1 W/m² and +RSD-invalid per the paper; those four rows are classified ``n/a`` and are +report-only. + +Running the benchmark +--------------------- + +.. code-block:: python + + from pyradtran.benchmarks import ( + CASES, + compare_benchmark, + format_report, + load_reference, + plot_benchmark_overlay, + run_randles2013, + write_report, + ) + + reference = load_reference() # bundled LBL values + _meta + print(CASES, sorted(k for k in reference if not k.startswith("_"))) + + def full_benchmark(out_dir, uvspec_exe=None, data_path=None): + """Run the matrix (needs uvspec), compare against LBL, emit report + plot.""" + results = run_randles2013(out_dir, uvspec_exe=uvspec_exe, data_path=data_path) + rows = compare_benchmark(results) # one status row per case/quantity/config + md, csv = write_report(rows, f"{out_dir}/randles2013.md") + png = plot_benchmark_overlay(rows, f"{out_dir}/randles2013_overlay.png") + return format_report(rows), md, csv, png + +:func:`~pyradtran.benchmarks.randles2013.run_randles2013` runs the whole +matrix (pass ``cases=`` to restrict it) and returns a nested results dict. + +The regression suite runs the real matrix. It is marked ``slow`` and skips +automatically when uvspec or the libRadtran data directory are not found: + +.. code-block:: console + + $ python -m pytest tests/test_benchmark_randles.py -v -s # ~25 s, real uvspec runs + $ python -m pytest -q -m "not slow" # the rest of the suite + +A ready-to-run driver for the whole workflow lives at +``examples/randles2013_benchmark/run_benchmark.py``: it runs the matrix, +compares against the reference and writes the Markdown/CSV report and the +overlay PNG to its ``output/`` directory. ``--from-json`` regenerates the +report from an earlier results JSON without re-running uvspec; the exit +code is non-zero if any row FAILs. + +Outputs +------- + +- ``compare_benchmark(results)`` — one row per (case, quantity, config) with + ``lbl``, ``ours``, ``abs_diff``, ``rel_diff`` and ``status`` (``PASS`` / + ``WARN`` / ``FAIL`` / ``n/a``). +- ``format_report(rows)`` — a Markdown report string, grouped by case with a + status-count summary. +- ``write_report(rows, path)`` — writes the Markdown report plus a CSV twin + next to it; returns both paths. +- ``plot_benchmark_overlay(rows, path)`` — a PNG overlaying this run's values + on the LBL reference. + +.. note:: + + ``run_randles2013`` persists its results as **JSON**, not NetCDF: the + matrix is a nested, ragged dict (per-case quantities differ), which NetCDF + cannot represent. The JSON additionally carries a ``results_path`` key + pointing at the artifact. + +Validation result +----------------- + +Real run against libRadtran 2.0.6 (uvspec, 2026-08-16): the regression suite +passes — 28 tests in 22.84 s — and the full-matrix comparison yields +**68 rows: 64 PASS / 0 WARN / 0 FAIL** (plus the 4 report-only ``n/a`` rows). +The three largest deviations, all comfortably inside the 8% PASS band: + +.. list-table:: + :header-rows: 1 + + * - Row (case / quantity / config) + - LBL + - pyRadtran + - rel. + * - case1 / absorptance / trop75 + - 0.307 + - 0.2858 + - -6.91% + * - case1 / nir_sfc_down / trop75 + - 101.1 + - 104.95 + - +3.81% + * - case1 / direct_bb_sfc_down / trop75 + - 179.6 + - 184.26 + - +2.59% + +The paper's central case-2b result reproduces: positive TOA forcing at SZA 30 +(ours +11.4 / +10.3 vs LBL +11.6 / +10.3 W/m² for ``saw`` / ``trop``) and +negative surface forcing in every configuration. diff --git a/docs/guide/config.rst b/docs/guide/config.rst new file mode 100644 index 0000000..e3ce99c --- /dev/null +++ b/docs/guide/config.rst @@ -0,0 +1,264 @@ +YAML Configuration +================== + +A pyRadtran experiment can be written as a single YAML file — scene, aerosol +and analysis intents — and run end-to-end without writing Python. The +:mod:`pyradtran.config` package maps the file onto exactly the builders the +Python API uses, so YAML and API-built scenes are interchangeable: for +equivalent input they produce **byte-identical uvspec input and +byte-identical aerosol layer files** (see `Round-trip guarantee`_). + +.. code-block:: python + + from pyradtran.config import load_config, export_config + + experiment = { + "config_version": 1, + "name": "demo", + "scene": { + "atmosphere": {"profile": "us"}, + "source": {"sza": 30.0}, + "wavelength": {"min_nm": 400.0, "max_nm": 700.0}, + "solver": {"method": "disort", "streams": 16, "disort_intcor": "moments"}, + "surface": {"albedo": 0.1}, + }, + } + loaded = load_config(experiment) # validates + builds; no files, no uvspec + path = export_config(loaded.config, "demo.yaml") # canonical YAML + assert load_config(path).config == loaded.config # re-loads identically + +:func:`~pyradtran.config.load_config` accepts a YAML file path, a raw ``dict`` +or an already-built :class:`~pyradtran.config.PyRadtranConfig` and returns a +``LoadedConfig`` (validated config, built :class:`~pyradtran.scene.Scene`, and +the optional composite aerosol). The schema is strict: unknown keys, unknown +block kinds and typo'd fields raise ``pydantic.ValidationError`` / +``ValueError`` instead of being silently ignored. + +File layout +----------- + +The schema version is pinned by ``config_version: 1``:: + + config_version: 1 + name: my_experiment # free-form label (default: "unnamed") + scene: { ... } # required — radiation scene + aerosol: { ... } # optional — LEGO composite aerosol + analysis: { ... } # optional — post-run intents + +Scene +----- + +Every sub-dict of ``scene`` is forwarded as keyword arguments to the matching +:class:`~pyradtran.scene.Scene` builder method: + +.. list-table:: + :header-rows: 1 + + * - Key + - Builder method + - Notes + * - ``atmosphere`` + - ``set_atmosphere`` + - required; e.g. ``profile: us`` (US standard atmosphere) + * - ``source`` + - ``set_source_solar`` / ``set_source_thermal`` + - inner key ``source: solar`` (default; needs ``sza``) or ``source: thermal`` + * - ``wavelength`` + - ``set_wavelength`` + - ``min_nm`` required, ``max_nm`` optional + * - ``solver`` + - ``set_solver`` + - optional; defaults ``disort``, 16 streams, ``disort_intcor: null``, no pseudospherical + * - ``surface`` + - ``set_surface`` + - optional; e.g. ``albedo: 0.1`` + * - ``output`` + - ``set_output`` + - optional; e.g. ``quantities``, ``zout``, ``format`` + +Aerosol +------- + +The ``aerosol`` section builds one +:class:`~pyradtran.models.aerosol_composite.CompositeAerosol` from LEGO blocks +(see :doc:`aerosols`). Shared keys: + +- ``wavelength_grid_um`` — required, µm; must match the block data +- ``altitude_grid_km`` — required, strictly descending, at least two levels +- ``n_legendre`` — Legendre moments per layer (default 32) +- ``output_dir`` — where the explicit ``.master`` / ``.LAYER`` files are + written at run time. A relative path resolves against the *working* + directory, not the config file; loading and validating never write files. +- ``blocks`` — list of blocks, discriminated by ``kind``: + +.. list-table:: + :header-rows: 1 + + * - ``kind`` + - Builds + - Specific fields + * - ``mie`` + - ``MieSpecies`` + - ``refractive_index`` (``wavelength_um`` / ``n_real`` / ``k_imag``), ``size_distribution``, ``particle_density_kg_m3``, ``phase_function`` (``hg`` default or ``mie``), ``integration`` + * - ``bulk`` + - ``BulkSpecies`` + - ``file`` — path to a NetCDF ``BulkAerosolOpticsData`` (Aerosol3D) + * - ``opac_preset`` + - OPAC preset block + - ``preset`` (e.g. ``maritime_clean``), ``rh_pct`` (50), ``species_names``, ``data_path``, ``n_legendre`` (32) + * - ``explicit_layer`` + - ``DirectLayerOpticsBlock`` + - ``master_path`` — pre-computed explicit file set + +``size_distribution`` carries ``kind`` (``lognormal``, ``modified_gamma``, +``discrete``, ``monodisperse``), a free-form ``params`` dict (for lognormal: +``r_g_um``, ``sigma_g``) and ``number_density_per_m3`` (default 1.0 — the +placement, not the distribution, sets the column loading). + +Placements +---------- + +Each block carries exactly one ``placement``, discriminated by ``kind``: + +.. list-table:: + :header-rows: 1 + + * - ``kind`` + - Meaning + - Fields + * - ``od_inversion`` + - exponential mass profile scaled so the block's optical depth at ``ref_nm`` equals ``tau_ref`` + - ``tau_ref``, ``ref_nm`` (550.0), ``scale_height_km`` + * - ``exponential`` + - exponential profile by density + - ``rho0_kg_m3``, ``scale_height_km`` + * - ``mass`` + - per-layer mass concentrations on the altitude grid + - ``kg_m3_per_layer`` + * - ``tabulated`` + - arbitrary tabulated profile + - ``z_km``, ``kg_m3`` (equal-length lists) + +Analysis intents +---------------- + +The optional ``analysis`` section drives what +:func:`~pyradtran.config.run_config` does beyond the main flux run: + +.. list-table:: + :header-rows: 1 + + * - Key + - Effect + * - ``energy_conservation`` + - compute the column energy budget and assert ``F_inc = eup_TOA + (1-a)(edir+edn)_surf + F_abs_atm`` within ``tol`` (default 0.05); optional surface ``albedo`` override + * - ``heating`` + - second uvspec run in heating mode (libRadtran's heating mode replaces flux output); rates merged into ``RunResult.rt`` + * - ``drf`` + - additional no-aerosol baseline run; ``RunResult.drf`` holds the spectral direct radiative forcing (negative = cooling) + * - ``attribution`` + - N+1 leave-one-out runs; ``RunResult.attribution`` (see :doc:`analysis`) + * - ``plots`` + - plot names resolved against a fixed registry (below) + * - ``save_netcdf`` + - path; the main run is persisted to NetCDF (``RunResult.netcdf_path``) + +Plot registry — valid ``analysis.plots`` names: + +- ``rt_spectral``, ``rt_overview``, ``rt_budget`` — main run +- ``rt_flux_profile_edir`` / ``rt_flux_profile_edn`` / ``rt_flux_profile_eup`` +- ``rt_heating_rate`` (requires ``heating: true``) +- ``drf_spectral`` (requires ``drf: true``) +- ``attribution_edir``, ``attribution_spectral`` (require ``attribution: true``) + +Command line +------------ + +The ``pyradtran`` console script (also runnable as ``python -m pyradtran``) +wraps the front-end: + +.. code-block:: console + + $ pyradtran validate canonical.yaml # schema check; no uvspec, no files + $ pyradtran run canonical.yaml # main run + analysis intents + $ pyradtran run canonical.yaml --plot-dir plots # redirect plots + $ pyradtran export-config canonical.yaml -o copy.yaml + +``run`` additionally accepts ``--uvspec`` (executable path) and ``--data-path`` +(libRadtran data directory), both forwarded to the ``Runner``. + +Canonical example walkthrough +----------------------------- + +``examples/multicomponent_viz/canonical.yaml`` is the YAML twin of the +canonical multicomponent scene used by the example gallery and the regression +tests. It is **generated** from ``canonical.py`` by ``make_yaml.py`` — do not +edit it by hand. Trimmed to its skeleton: + +.. code-block:: yaml + + config_version: 1 + name: multicomponent_viz_canonical + scene: + atmosphere: {profile: us} + source: {sza: 30.0} + wavelength: {min_nm: 401.0, max_nm: 699.0} + solver: {method: disort, streams: 16, disort_intcor: moments, pseudospherical: true} + surface: {albedo: 0.1} + output: + quantities: [lambda, edir, edn, eup] + format: ascii + zout: [0, 1, 2, 4, 6, 8, 10, toa] + heating_rate: local + aerosol: + wavelength_grid_um: [0.4, 0.45, 0.5, 0.55, 0.6, 0.65, 0.7] + altitude_grid_km: [10.0, 8.0, 6.0, 4.0, 2.0, 1.0, 0.0] + n_legendre: 32 + output_dir: output + blocks: + - kind: mie + name: black_carbon + refractive_index: {wavelength_um: [0.3, 0.4, 0.55, 0.7, 1.0], + n_real: [1.95, ...], k_imag: [0.79, ...]} + size_distribution: {kind: lognormal, params: {r_g_um: 0.1, sigma_g: 2.0}} + particle_density_kg_m3: 1800.0 + phase_function: hg + placement: {kind: od_inversion, tau_ref: 0.15, ref_nm: 550.0, scale_height_km: 1.5} + # ... two more mie blocks on the same grids: + # sulfate tau_ref 0.15, scale_height_km 2.0 + # mineral_dust tau_ref 0.20, scale_height_km 3.0 + analysis: null + +Three externally-mixed Mie blocks share the 7-wavelength / 7-altitude grids; +their ``od_inversion`` placements scale the exponential profiles so the column +holds τ550 = 0.15 + 0.15 + 0.20 = 0.5. + +Because ``output_dir: output`` is relative, run the file from the example +directory so the generated ``.master`` / ``.LAYER`` files land next to it: + +.. code-block:: console + + $ cd examples/multicomponent_viz + $ pyradtran run canonical.yaml + +(or copy the file and point ``aerosol.output_dir`` / ``--plot-dir`` elsewhere). + +Round-trip guarantee +-------------------- + +The YAML front-end is not an approximation of the Python API — it is the same +code path: + +1. A YAML config and its API-built twin produce **identical uvspec input + text** (``test_yaml_matches_api_uvspec_text``). +2. The explicit aerosol ``.master`` layer files they write are + **byte-identical** (``test_yaml_master_file_bytes_identical``). +3. :func:`~pyradtran.config.export_config` serializes a config as canonical + YAML that re-loads into an equal config: + ``load_config(export_config(cfg, p)).config == cfg``. +4. ``canonical.yaml`` is regenerated from ``canonical.py`` by ``make_yaml.py`` + and ``test_generated_canonical_yaml_matches_api`` fails if the two drift. + +Guarantees 1, 2 and 4 are enforced by ``tests/test_config_roundtrip.py``; +the export round trip (3) by ``tests/test_config_schema.py`` +(``test_export_config_round_trip``). diff --git a/docs/index.rst b/docs/index.rst index b05c144..fb7d1c2 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -11,12 +11,14 @@ pyRadtran provides a Pythonic, type-safe API for configuring and executing libRa guide/installation guide/quickstart + guide/config guide/scene_builder guide/aerosols guide/blocks guide/data guide/visualization guide/analysis + guide/benchmarks guide/convenience guide/parallel diff --git a/examples/multicomponent_viz/canonical.py b/examples/multicomponent_viz/canonical.py new file mode 100644 index 0000000..16b009b --- /dev/null +++ b/examples/multicomponent_viz/canonical.py @@ -0,0 +1,189 @@ +"""Canonical scene configuration for the multicomponent_viz example. + +Shared by run_demo.py and tests/test_multicomponent_regression.py so the +regression baseline and the demo cannot drift apart. +""" + +import logging +from pathlib import Path + +from pyradtran.models.aerosol_composite import ( + CompositeAerosol, + IntegrationConfig, + MieSpecies, + RefractiveIndex, + SizeDistribution, +) +from pyradtran.models.blocks import PlacedBlock, od_to_mass_profile +from pyradtran.scene import Scene + +logger = logging.getLogger(__name__) + +EXAMPLE_DIR = Path(__file__).parent +OUTPUT_DIR = EXAMPLE_DIR / "output" + +# --- 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, + }, +] + + +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: + """Build the RT scene with the given aerosol composite attached.""" + 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 build_scene_no_aerosol() -> Scene: + """Same scene as build_scene() minus the aerosol — DRF baseline.""" + 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"]) + ) + + +def build_scene_heating(aerosol: CompositeAerosol) -> Scene: + """Scene variant for heating-rate output (libRadtran heating_rate mode). + + Identical to build_scene() but drops ``output_user`` quantities. With + ``heating_rate`` set and no ``output_user``, uvspec emits heating rates + (K/day, wide format) instead of fluxes — a second run is required because + libRadtran cannot emit both in one invocation. + """ + c = SCENE_KW + heating_output = {k: v for k, v in c["output"].items() if k != "quantities"} + 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(**heating_output) + .set_aerosol(aerosol) + ) diff --git a/examples/multicomponent_viz/canonical.yaml b/examples/multicomponent_viz/canonical.yaml new file mode 100644 index 0000000..4ee01eb --- /dev/null +++ b/examples/multicomponent_viz/canonical.yaml @@ -0,0 +1,176 @@ +# canonical.yaml — YAML twin of the canonical multicomponent_viz scene +# (SCENE_KW + three Mie blocks with od_inversion placements on shared grids). +# +# GENERATED by make_yaml.py from canonical.py — do not edit by hand; +# regenerate with: +# cd examples/multicomponent_viz && python make_yaml.py +# +# `aerosol.output_dir: output` is RELATIVE: explicit .master/.LAYER files are +# written relative to the working directory at run time (loading and +# validating write no files). Either run from the example directory: +# cd examples/multicomponent_viz && pyradtran run canonical.yaml +# or point the output/plot dirs elsewhere: set `aerosol.output_dir` in a copy +# of this file and pass `--plot-dir` to `pyradtran run`. +# +# config_version 1. +config_version: 1 +name: multicomponent_viz_canonical +scene: + atmosphere: + profile: us + altitude: 0.0 + source: + sza: 30.0 + wavelength: + min_nm: 401.0 + max_nm: 699.0 + solver: + method: disort + streams: 16 + disort_intcor: moments + pseudospherical: true + surface: + albedo: 0.1 + output: + quantities: + - lambda + - edir + - edn + - eup + format: ascii + zout: + - 0 + - 1 + - 2 + - 4 + - 6 + - 8 + - 10 + - toa + heating_rate: local +aerosol: + wavelength_grid_um: + - 0.4 + - 0.45 + - 0.5 + - 0.55 + - 0.6 + - 0.65 + - 0.7 + altitude_grid_km: + - 10.0 + - 8.0 + - 6.0 + - 4.0 + - 2.0 + - 1.0 + - 0.0 + n_legendre: 32 + output_dir: output + blocks: + - kind: mie + name: black_carbon + refractive_index: + wavelength_um: + - 0.3 + - 0.4 + - 0.55 + - 0.7 + - 1.0 + n_real: + - 1.95 + - 1.95 + - 1.95 + - 1.95 + - 1.95 + k_imag: + - 0.79 + - 0.79 + - 0.79 + - 0.79 + - 0.79 + size_distribution: + kind: lognormal + params: + r_g_um: 0.1 + sigma_g: 2.0 + number_density_per_m3: 1.0 + particle_density_kg_m3: 1800.0 + phase_function: hg + integration: null + placement: + kind: od_inversion + tau_ref: 0.15 + ref_nm: 550.0 + scale_height_km: 1.5 + - kind: mie + name: sulfate + refractive_index: + wavelength_um: + - 0.3 + - 0.4 + - 0.55 + - 0.7 + - 1.0 + n_real: + - 1.53 + - 1.53 + - 1.53 + - 1.53 + - 1.53 + k_imag: + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + size_distribution: + kind: lognormal + params: + r_g_um: 0.15 + sigma_g: 1.7 + number_density_per_m3: 1.0 + particle_density_kg_m3: 1770.0 + phase_function: hg + integration: null + placement: + kind: od_inversion + tau_ref: 0.15 + ref_nm: 550.0 + scale_height_km: 2.0 + - kind: mie + name: mineral_dust + refractive_index: + wavelength_um: + - 0.3 + - 0.4 + - 0.55 + - 0.7 + - 1.0 + n_real: + - 1.53 + - 1.53 + - 1.53 + - 1.53 + - 1.53 + k_imag: + - 0.008 + - 0.008 + - 0.008 + - 0.008 + - 0.008 + size_distribution: + kind: lognormal + params: + r_g_um: 0.5 + sigma_g: 2.2 + number_density_per_m3: 1.0 + particle_density_kg_m3: 2600.0 + phase_function: hg + integration: null + placement: + kind: od_inversion + tau_ref: 0.2 + ref_nm: 550.0 + scale_height_km: 3.0 +analysis: null diff --git a/examples/multicomponent_viz/make_yaml.py b/examples/multicomponent_viz/make_yaml.py new file mode 100644 index 0000000..4251db9 --- /dev/null +++ b/examples/multicomponent_viz/make_yaml.py @@ -0,0 +1,101 @@ +"""Regenerate ``canonical.yaml`` — the YAML twin of the ``canonical.py`` scene. + +The emitted config describes the exact same experiment as +``canonical.build_scene(canonical.build_composite())`` (same SCENE_KW, same +three Mie blocks with od_inversion placements, same shared grids), except +that ``aerosol.output_dir`` is the RELATIVE path ``output``: the explicit +``.master``/``.LAYER`` files are written relative to the working directory +at run time (loading/validating a config writes nothing). + +Run from this directory (any cwd works — ``canonical.py`` is found next to +this script): + + cd examples/multicomponent_viz && python make_yaml.py +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import canonical # script directory is on sys.path when run as a script + +from pyradtran.config import export_config + +TARGET = Path(__file__).resolve().parent / "canonical.yaml" + +HEADER = """\ +# canonical.yaml — YAML twin of the canonical multicomponent_viz scene +# (SCENE_KW + three Mie blocks with od_inversion placements on shared grids). +# +# GENERATED by make_yaml.py from canonical.py — do not edit by hand; +# regenerate with: +# cd examples/multicomponent_viz && python make_yaml.py +# +# `aerosol.output_dir: output` is RELATIVE: explicit .master/.LAYER files are +# written relative to the working directory at run time (loading and +# validating write no files). Either run from the example directory: +# cd examples/multicomponent_viz && pyradtran run canonical.yaml +# or point the output/plot dirs elsewhere: set `aerosol.output_dir` in a copy +# of this file and pass `--plot-dir` to `pyradtran run`. +# +# config_version 1. +""" + + +def config_dict() -> dict: + """Config equivalent to canonical.build_composite() + build_scene(). + + Mirrors ``tests/test_config_roundtrip.py::_config_dict`` (same shape, + same constants) but with the relative ``output_dir: output``; that test + enforces the YAML/API equivalence, and this generator keeps the committed + file in sync with canonical.py. + """ + blocks = [ + { + "kind": "mie", + "name": b["name"], + "refractive_index": { + "wavelength_um": canonical._WL_RI, + "n_real": b["n_real"], + "k_imag": b["k_imag"], + }, + "size_distribution": { + "kind": "lognormal", + "params": {"r_g_um": b["r_g_um"], "sigma_g": b["sigma_g"]}, + }, + "particle_density_kg_m3": b["density"], + "placement": { + "kind": "od_inversion", + "tau_ref": b["tau_550"], + "ref_nm": canonical.REF_NM, + "scale_height_km": b["scale_height_km"], + }, + } + for b in canonical.BLOCKS + ] + return { + "config_version": 1, + "name": "multicomponent_viz_canonical", + "scene": dict(canonical.SCENE_KW), + "aerosol": { + "wavelength_grid_um": canonical.WAVELENGTHS_UM, + "altitude_grid_km": canonical.ALTITUDE_GRID_KM, + "n_legendre": canonical.N_LEGENDRE, + "output_dir": "output", + "blocks": blocks, + }, + } + + +def main() -> int: + """Write canonical.yaml = comment header + canonical YAML config.""" + export_config(config_dict(), TARGET) # pure YAML (no comments) + body = TARGET.read_text(encoding="utf-8") + TARGET.write_text(HEADER + body, encoding="utf-8") + print(f"wrote {TARGET}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/examples/multicomponent_viz/run_demo.py b/examples/multicomponent_viz/run_demo.py index 7043bf8..aa8d0c6 100644 --- a/examples/multicomponent_viz/run_demo.py +++ b/examples/multicomponent_viz/run_demo.py @@ -1,29 +1,25 @@ """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. +Self-contained (pyRadtran APIs only, no aerosol3d dependency). Scene config + +composite builders live in ``canonical.py`` (shared with the regression test). - 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 +Scientific self-validation wired in: + - B1: column tau@550 logged + checked against sum of block targets. + - B2: column energy conservation asserted (F_inc = eup_TOA + (1-a)(edir+edn)_surf + + F_abs_atm); script exits non-zero if the physics is violated. + - T3 DRF: no-aerosol baseline run -> direct radiative forcing spectrum. + - B3: ``--dump-baseline PATH`` writes the scalar regression fixture. Requires libRadtran. Set PYRADTRAN_DATA_PATH to its data/ dir (or rely on the bundled resolver). Usage: python run_demo.py + python run_demo.py --dump-baseline tests/fixtures/multicomponent_baseline.json """ +import argparse +import json import logging import os from pathlib import Path @@ -32,31 +28,44 @@ matplotlib.use("Agg") import matplotlib.pyplot as plt - -from pyradtran import Runner, Scene -from pyradtran.core.output_parser import HEATING_RATE_COLUMN +import numpy as np + +from canonical import ( + ALBEDO, + ALTITUDE_GRID_KM, + BLOCKS, + N_LEGENDRE, + OUTPUT_DIR, + WAVELENGTHS_UM, + _WL_RI, + build_composite, + build_scene, + build_scene_heating, + build_scene_no_aerosol, +) +from pyradtran import Runner from pyradtran.core.postprocess import ( add_budget_vars, + assert_energy_conservation, 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.models.aerosol_composite import RefractiveIndex, SizeDistribution from pyradtran.viz import ( plot_block_profiles, + plot_block_spectral_optics, plot_budget, plot_component_attribution, plot_composite_optics, + plot_drf_spectral, plot_flux_profile, plot_heating_rate, + plot_legendre_decay, + plot_phase_functions, plot_rt_overview, + plot_size_distributions, plot_spectral, + plot_spectral_attribution, save, set_theme, ) @@ -65,77 +74,8 @@ 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 @@ -144,70 +84,16 @@ def _save(fig, name): 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 main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--dump-baseline", + type=str, + default=None, + help="Write B3 scalar baseline JSON to PATH and continue.", ) + args, _ = parser.parse_known_args() - -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: @@ -223,6 +109,14 @@ def main(): grid_ds = evaluate_composite_on_grid( composite, WAVELENGTHS_UM, ALTITUDE_GRID_KM, n_legendre=N_LEGENDRE ) + # B1: column tau@550 must stack to the sum of block targets. + col_tau_550 = float(np.sum(grid_ds["tau"].sel(wavelength=0.55).values)) + target_sum = sum(b["tau_550"] for b in BLOCKS) + logger.info(" column tau@550 = %.4f (target sum %.2f)", col_tau_550, target_sum) + assert np.isclose(col_tau_550, target_sum, rtol=0.01), ( + f"B1 FAIL: column tau@550={col_tau_550:.4f}, expected {target_sum:.4f}" + ) + for q in ("tau", "ssa", "g"): fig, _ = plot_composite_optics(grid_ds, quantity=q) _save(fig, f"composite_{q}.png") @@ -234,32 +128,104 @@ def main(): fig, _ = plot_block_profiles(block_dict, quantity="rho") _save(fig, "block_rho_profiles.png") + # T2: input diagnostics — size distributions, per-block spectral optics, + # phase functions, Legendre decay (the Mie->DISORT scattering physics). + blocks_ri_sd = { + b["name"]: ( + RefractiveIndex(wavelength_um=_WL_RI, n_real=b["n_real"], k_imag=b["k_imag"]), + SizeDistribution( + kind="lognormal", params={"r_g_um": b["r_g_um"], "sigma_g": b["sigma_g"]} + ), + ) + for b in BLOCKS + } + fig, _ = plot_size_distributions({n: sd for n, (_, sd) in blocks_ri_sd.items()}) + _save(fig, "size_distributions.png") + for q in ("tau", "ssa", "g"): + fig, _ = plot_block_spectral_optics(block_dict, quantity=q) + _save(fig, f"block_spectral_{q}.png") + fig, _ = plot_phase_functions(blocks_ri_sd, wavelength_um=0.55) + _save(fig, "phase_functions_550.png") + fig, _ = plot_legendre_decay( + composite, wavelength_um=0.55, n_legendre=N_LEGENDRE, altitude_grid_km=ALTITUDE_GRID_KM + ) + _save(fig, "legendre_decay_550.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)) + # B2: column energy conservation (hard assertion). + budget = assert_energy_conservation(rt, albedo=ALBEDO, tol=0.05) + i550 = int(np.argmin(np.abs(budget.wavelength - 550.0))) + logger.info( + " energy@550: F_inc=%.1f up_TOA=%.1f abs_surf=%.1f abs_atm=%.1f (W/m²)", + budget.f_incident[i550], + budget.f_up_toa[i550], + budget.f_abs_surface[i550], + budget.f_abs_atm[i550], + ) + 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") + # T1: exercise edn + eup flux profiles too (was edir-only). + for var in ("edn", "eup"): + fig, _ = plot_flux_profile(rt, variable=var, wavelength_nm=550.0) + _save(fig, f"rt_flux_profile_{var}.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") + # T1: heating rate needs its own uvspec run — libRadtran's heating_rate + # mode replaces flux output (wide format) rather than appending a column. + logger.info("=== Heating-rate run (K/day, separate invocation) ===") + rt_heat = Runner.execute(build_scene_heating(composite), data_path=data_path) + rt = rt.assign(heating_rate=rt_heat["heating_rate"]) + logger.info( + " heating@550 (K/day): %s", + ", ".join( + f"z={z:.0f}km={v:.3e}" + for z, v in zip( + rt["zout"].values, + rt["heating_rate"].isel(wavelength=i550).values, + ) + ), + ) + fig, _ = plot_heating_rate(rt, wavelength_nm=550.0) + _save(fig, "rt_heating_rate.png") rt.to_netcdf(str(OUTPUT_DIR / "rt_full.nc")) fig, _ = plot_rt_overview(rt, wavelength_nm=550.0) _save(fig, "rt_overview.png") + # --- T3: direct radiative forcing (with aerosol - without aerosol) --- + logger.info("=== DRF baseline (no aerosol) ===") + rt_noaer = Runner.execute(build_scene_no_aerosol(), data_path=data_path) + wl_nm = rt["wavelength"].values + surf_idx = int(np.argmin(rt["zout"].values)) + toa_idx = int(np.argmax(rt["zout"].values)) + net_toa_aer = rt["edir"].isel(zout=toa_idx).values - rt["eup"].isel(zout=toa_idx).values + net_toa_noaer = ( + rt_noaer["edir"].isel(zout=toa_idx).values - rt_noaer["eup"].isel(zout=toa_idx).values + ) + net_sfc_aer = (rt["edir"] + rt["edn"] - rt["eup"]).isel(zout=surf_idx).values + net_sfc_noaer = ( + (rt_noaer["edir"] + rt_noaer["edn"] - rt_noaer["eup"]).isel(zout=surf_idx).values + ) + drf_toa = net_toa_aer - net_toa_noaer # <0 = cooling (IPCC convention) + drf_surf = net_sfc_aer - net_sfc_noaer + fig, _ = plot_drf_spectral(wl_nm, drf_toa, drf_surf) + _save(fig, "drf_spectral.png") + logger.info( + " DRF@550: TOA=%.2f surf=%.2f atm=%.2f W/m²", + float(np.interp(550.0, wl_nm, drf_toa)), + float(np.interp(550.0, wl_nm, drf_surf)), + float(np.interp(550.0, wl_nm, drf_toa - drf_surf)), + ) + # --- Component attribution (leave-one-out) --- logger.info("=== Component attribution (N+1 DISORT runs) ===") @@ -272,6 +238,34 @@ def execute_many(scenes): logger.info(" contributions: %s", list(result.contributions)) fig, _ = plot_component_attribution(result, variable="edir", level="surface") _save(fig, "attribution_edir.png") + # T1: attribution at additional variable/level. plot_component_attribution + # treats non-"surface" level as an integer zout index, so resolve "toa". + toa_zout_idx = int(np.argmax(result.full["zout"].values)) + fig, _ = plot_component_attribution(result, variable="eup", level=toa_zout_idx) + _save(fig, "attribution_eup_toa.png") + fig, _ = plot_component_attribution(result, variable="edn", level="surface") + _save(fig, "attribution_edn_surf.png") + # T3: spectral attribution (per-block contribution across the band). + fig, _ = plot_spectral_attribution(result, variable="edir", level="surface") + _save(fig, "attribution_edir_spectral.png") + + # --- B3: optional scalar baseline dump --- + if args.dump_baseline: + i550_rt = int(np.argmin(np.abs(rt["wavelength"].values - 550.0))) + scalars = { + "_meta": { + "scene": "multicomponent_viz", + "note": "regenerate via run_demo.py --dump-baseline", + }, + "edir_surf_550nm": float(rt["edir"].isel(wavelength=i550_rt, zout=surf_idx).values), + "edn_surf_550nm": float(rt["edn"].isel(wavelength=i550_rt, zout=surf_idx).values), + "eup_surf_550nm": float(rt["eup"].isel(wavelength=i550_rt, zout=surf_idx).values), + "eup_toa_550nm": float(rt["eup"].isel(wavelength=i550_rt, zout=toa_idx).values), + "column_tau_550nm": col_tau_550, + "F_abs_atm_550nm": float(budget.f_abs_atm[i550]), + } + Path(args.dump_baseline).write_text(json.dumps(scalars, indent=2)) + logger.info("wrote baseline -> %s", args.dump_baseline) logger.info("=== Comprehensive demo complete. Figures in %s ===", OUTPUT_DIR) diff --git a/examples/randles2013_benchmark/.gitignore b/examples/randles2013_benchmark/.gitignore new file mode 100644 index 0000000..ea1472e --- /dev/null +++ b/examples/randles2013_benchmark/.gitignore @@ -0,0 +1 @@ +output/ diff --git a/examples/randles2013_benchmark/run_benchmark.py b/examples/randles2013_benchmark/run_benchmark.py new file mode 100644 index 0000000..786cafb --- /dev/null +++ b/examples/randles2013_benchmark/run_benchmark.py @@ -0,0 +1,86 @@ +#!/usr/bin/env python3 +"""Randles et al. (2013) AeroCom shortwave benchmark — end-to-end example. + +Runs the full case matrix (3 cases x 2 AFGL atmospheres x 2 SZAs x 2 bands += 24 uvspec invocations, ~20-60 s), normalizes against the LBL median +constants, compares against the bundled reference +(``pyradtran.benchmarks.reference.randles2013_lbl.json``) and writes + +- ``output/randles2013_results.json`` — raw + normalized results +- ``output/randles2013_report.md`` / ``.csv`` — PASS/WARN/FAIL comparison +- ``output/randles2013_overlay.png`` — LBL-vs-pyRadtran bar overlay + +``--from-json PATH`` regenerates the report and plot from an earlier +results JSON without re-running uvspec. The exit code is non-zero if any +row FAILs, so the driver doubles as a CI gate. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") # headless rendering; mirrors examples/multicomponent_viz + +from pyradtran.benchmarks import ( # noqa: E402 + compare_benchmark, + format_report, + plot_benchmark_overlay, + run_randles2013, + write_report, +) + +HERE = Path(__file__).resolve().parent +DEFAULT_OUT = HERE / "output" + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description="Randles 2013 benchmark driver") + parser.add_argument("--uvspec", default=None, help="uvspec binary (auto-detected if omitted)") + parser.add_argument( + "--data-path", default=None, help="libRadtran data directory (auto-detected)" + ) + parser.add_argument( + "--cases", nargs="*", default=None, help="subset of case1 case2a case2b (default: all)" + ) + parser.add_argument( + "--from-json", default=None, help="reuse an existing results JSON; skips uvspec" + ) + parser.add_argument( + "--out-dir", default=str(DEFAULT_OUT), help=f"output directory (default: {DEFAULT_OUT})" + ) + args = parser.parse_args(argv) + + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + if args.from_json: + results = json.loads(Path(args.from_json).read_text(encoding="utf-8")) + print(f"[randles2013] reusing results from {args.from_json}") + else: + print(f"[randles2013] running the case matrix -> {out_dir}") + results = run_randles2013( + out_dir, uvspec_exe=args.uvspec, data_path=args.data_path, cases=args.cases or None + ) + + rows = compare_benchmark(results) + md_path, csv_path = write_report(rows, out_dir / "randles2013_report.md") + png_path = plot_benchmark_overlay(rows, out_dir / "randles2013_overlay.png") + + print(format_report(rows)) + print(f"[randles2013] results: {results.get('results_path', '(in-memory)')}") + print(f"[randles2013] report: {md_path}") + print(f"[randles2013] csv: {csv_path}") + print(f"[randles2013] plot: {png_path}") + + n_fail = sum(row["status"] == "FAIL" for row in rows) + n_warn = sum(row["status"] == "WARN" for row in rows) + print(f"[randles2013] {len(rows)} rows: {n_fail} FAIL, {n_warn} WARN") + return 1 if n_fail else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/pyproject.toml b/pyproject.toml index 47c923a..0242f06 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,12 +23,16 @@ classifiers = [ requires-python = ">=3.11" dependencies = [ "pydantic>=2.0", + "pyyaml>=6.0", "xarray", "numpy", "netCDF4", ] # viz modules import matplotlib lazily; matplotlib stays in the `plot` extra (not core). +[project.scripts] +pyradtran = "pyradtran.config.cli:main" + [project.optional-dependencies] dev = [ "pytest>=7.0", diff --git a/scripts/regen_baseline.py b/scripts/regen_baseline.py new file mode 100644 index 0000000..866edc2 --- /dev/null +++ b/scripts/regen_baseline.py @@ -0,0 +1,23 @@ +"""Regenerate tests/fixtures/multicomponent_baseline.json. + +Usage: MPLBACKEND=Agg python scripts/regen_baseline.py +""" + +import os +import subprocess +import sys +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] + + +def main(): + target = ROOT / "tests" / "fixtures" / "multicomponent_baseline.json" + demo = ROOT / "examples" / "multicomponent_viz" / "run_demo.py" + env = {**os.environ, "MPLBACKEND": "Agg"} + subprocess.run([sys.executable, str(demo), "--dump-baseline", str(target)], env=env, check=True) + print(f"wrote {target}") + + +if __name__ == "__main__": + main() diff --git a/src/pyradtran/__main__.py b/src/pyradtran/__main__.py new file mode 100644 index 0000000..684defd --- /dev/null +++ b/src/pyradtran/__main__.py @@ -0,0 +1,8 @@ +"""Module entry point: ``python -m pyradtran`` delegates to the config CLI.""" + +import sys + +from pyradtran.config.cli import main + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/pyradtran/benchmarks/__init__.py b/src/pyradtran/benchmarks/__init__.py new file mode 100644 index 0000000..9567a21 --- /dev/null +++ b/src/pyradtran/benchmarks/__init__.py @@ -0,0 +1,33 @@ +"""Benchmarks: published radiative-transfer intercomparison case matrices. + +Currently ships the Randles et al. (2013) shortwave benchmark +(:mod:`pyradtran.benchmarks.randles2013`) with its bundled LBL reference +values, plus comparison/report/plot tooling +(:mod:`pyradtran.benchmarks.compare`). +""" + +from __future__ import annotations + +from pyradtran.benchmarks.compare import ( + compare_benchmark, + format_report, + plot_benchmark_overlay, + write_report, +) +from pyradtran.benchmarks.randles2013 import ( + CASES, + RandlesAerosol, + load_reference, + run_randles2013, +) + +__all__ = [ + "CASES", + "RandlesAerosol", + "compare_benchmark", + "format_report", + "load_reference", + "plot_benchmark_overlay", + "run_randles2013", + "write_report", +] diff --git a/src/pyradtran/benchmarks/compare.py b/src/pyradtran/benchmarks/compare.py new file mode 100644 index 0000000..69eb5b3 --- /dev/null +++ b/src/pyradtran/benchmarks/compare.py @@ -0,0 +1,463 @@ +"""Compare benchmark results against the bundled LBL reference. + +Turns a results dict from :func:`pyradtran.benchmarks.run_randles2013` +(shape ``{case: {band: {atm: {sza: {quantity: value}}}}}``, optionally plus a +stray ``results_path`` string key) into flat comparison rows, a Markdown/CSV +report, and a LBL-vs-pyRadtran overlay figure. + +Status classification +--------------------- +Fluxes (every quantity except ``*_rf``): PASS if ``|rel| <= 8%``, WARN if +``|rel| <= 12%``, else FAIL (reference ``_meta`` thresholds). + +Radiative forcing (``toa_rf``/``sfc_rf``): PASS if ``|rel| <= 15%`` OR +``|abs| <= 1.5 W/m2`` (reference ``_meta``). The paper only fixes the PASS +band; by analogy we set WARN if ``|rel| <= 25%`` OR ``|abs| <= 3 W/m2``, and +FAIL beyond both. ``rel_diff`` is in percent; when the reference value is 0 +the absolute difference (quantity units) is used as the fallback so the +thresholds remain well defined. + +Case-2a atmospheric RF (``toa_rf - sfc_rf``, <1 W/m2 and RSD-invalid per the +paper) is emitted as report-only rows with status ``"n/a"``. + +Robustness +---------- +- Non-case keys in ``results`` (e.g. ``results_path``) are ignored: only the + known case names are read. +- SZA keys may be ints (in-memory) or strings ("30"; results loaded back from + JSON); both are accepted. +- RF reference values are broadband-only, so RF rows are emitted exclusively + from the ``bb`` band; any ``rf_*`` entries under ``uvvis`` are ignored. +""" + +from __future__ import annotations + +import csv +from pathlib import Path + +from pyradtran.benchmarks.randles2013 import CASES, load_reference + +#: Config keys of the reference file, in canonical report order. +CONFIGS: tuple[str, ...] = ("saw30", "saw75", "trop30", "trop75") + +# Thresholds (kept in sync with reference/_meta "thresholds"; the RF WARN band +# is our documented extension of the paper's PASS-only rule). +FLUX_PASS_PCT = 8.0 +FLUX_WARN_PCT = 12.0 +RF_PASS_REL_PCT = 15.0 +RF_PASS_ABS_W_M2 = 1.5 +RF_WARN_REL_PCT = 25.0 +RF_WARN_ABS_W_M2 = 3.0 + +#: Status -> Markdown emoji (plain words are used in the CSV). +STATUS_EMOJI: dict[str, str] = {"PASS": "✅", "WARN": "⚠️", "FAIL": "❌"} + +#: Report-only status for case-2a atmospheric RF rows. +NA_STATUS = "n/a" + +#: Reference quantity -> (results band, results quantity key), per case. +_CASE2_QUANTITIES: dict[str, tuple[str, str]] = { + "total_bb_sfc_down": ("bb", "total_sfc_down"), + "diffuse_bb_toa_up": ("bb", "eup_toa"), + "uvvis_sfc_down": ("uvvis", "total_sfc_down"), + # Broadband only: the reference RF is a bb difference; uvvis rf_* ignored. + "toa_rf": ("bb", "rf_toa"), + "sfc_rf": ("bb", "rf_sfc"), +} +_QUANTITY_MAP: dict[str, dict[str, tuple[str, str]]] = { + "case1": { + "direct_bb_sfc_down": ("bb", "edir_sfc"), + "diffuse_bb_sfc_down": ("bb", "edn_sfc"), + "diffuse_bb_toa_up": ("bb", "eup_toa"), + "uvvis_sfc_down": ("uvvis", "total_sfc_down"), + "nir_sfc_down": ("bb", "nir_sfc_down"), + "absorptance": ("bb", "absorptance"), + }, + "case2a": _CASE2_QUANTITIES, + "case2b": _CASE2_QUANTITIES, +} + +#: Column order of the CSV side of :func:`write_report`. +CSV_COLUMNS: tuple[str, ...] = ( + "benchmark", + "case", + "quantity", + "band", + "atm", + "sza", + "units", + "lbl", + "ours", + "abs_diff", + "rel_diff", + "status", +) + + +def _split_config(config: str) -> tuple[str, int]: + """Split a config key ("saw30"/"trop75") into (atmosphere, sza).""" + return config[:-2], int(config[-2:]) + + +def _lookup_ours(results: dict, case: str, quantity: str, config: str) -> float | None: + """Fetch our value for (case, reference quantity, config). + + Returns None when the results dict does not carry the entry (subset run). + SZA keys are accepted as int or str (JSON round trip). + """ + band, key = _QUANTITY_MAP[case][quantity] + atm, sza = _split_config(config) + entry = results.get(case, {}).get(band, {}).get(atm) + if not isinstance(entry, dict): + return None + for sza_key in (sza, str(sza)): + run = entry.get(sza_key) + if isinstance(run, dict): + value = run.get(key) + if value is not None: + return float(value) + return None + + +def _status_for(quantity: str, abs_diff: float, rel_diff: float) -> str: + """Classify one comparison using the flux or RF rule (see module docstring).""" + if quantity.endswith("_rf"): + if abs(rel_diff) <= RF_PASS_REL_PCT or abs(abs_diff) <= RF_PASS_ABS_W_M2: + return "PASS" + if abs(rel_diff) <= RF_WARN_REL_PCT or abs(abs_diff) <= RF_WARN_ABS_W_M2: + return "WARN" + return "FAIL" + if abs(rel_diff) <= FLUX_PASS_PCT: + return "PASS" + if abs(rel_diff) <= FLUX_WARN_PCT: + return "WARN" + return "FAIL" + + +def compare_benchmark(results: dict, benchmark: str = "randles2013") -> list[dict]: + """Compare run results against the bundled LBL reference. + + Args: + results: Output of :func:`run_randles2013` (possibly loaded back from + its saved JSON, in which case SZA keys are strings and a + ``results_path`` key is present; both are handled). + benchmark: Benchmark whose reference to compare against. Currently + only ``"randles2013"``. + + Returns: + One row per (case, quantity, atm, sza) with fields ``case``, + ``quantity``, ``band``, ``atm``, ``sza``, ``units``, ``lbl``, ``ours``, + ``abs_diff``, ``rel_diff`` (percent; absolute fallback when ``lbl == 0``) + and ``status`` (PASS/WARN/FAIL; case-2a ``atmospheric_rf`` rows are + report-only with status ``"n/a"``). Quantities missing from the results + dict are skipped, so subset runs yield fewer rows. + + Raises: + ValueError: If ``benchmark`` names no known benchmark. + """ + if benchmark != "randles2013": + raise ValueError(f"Unknown benchmark '{benchmark}'. Valid: ['randles2013']") + reference = load_reference() + + rows: list[dict] = [] + # Iterate only known case keys: stray keys (results_path, ...) are skipped. + for case in CASES: + if case not in results: + continue + for quantity, config_values in reference[case].items(): + band = _QUANTITY_MAP[case][quantity][0] + for config in CONFIGS: + atm, sza = _split_config(config) + lbl = float(config_values[config]) + ours = _lookup_ours(results, case, quantity, config) + if ours is None: + continue + abs_diff = ours - lbl + rel_diff = abs_diff if lbl == 0.0 else abs_diff / abs(lbl) * 100.0 + rows.append( + { + "benchmark": benchmark, + "case": case, + "quantity": quantity, + "band": band, + "atm": atm, + "sza": sza, + "units": "1" if quantity == "absorptance" else "W/m2", + "lbl": lbl, + "ours": ours, + "abs_diff": abs_diff, + "rel_diff": rel_diff, + "status": _status_for(quantity, abs_diff, rel_diff), + } + ) + if case == "case2a": + rows.extend(_atmospheric_rf_rows(results, reference)) + return rows + + +def _atmospheric_rf_rows(results: dict, reference: dict) -> list[dict]: + """Report-only rows for the case-2a atmospheric RF (``toa_rf - sfc_rf``). + + The paper flags case-2a atmospheric RF as <1 W/m2 with invalid RSD, so it + is reported (LBL value derived as the difference of the tabulated toa/sfc + reference RFs) but not classified; ``status`` is ``"n/a"``. + """ + rows = [] + for config in CONFIGS: + atm, sza = _split_config(config) + toa = _lookup_ours(results, "case2a", "toa_rf", config) + sfc = _lookup_ours(results, "case2a", "sfc_rf", config) + if toa is None or sfc is None: + continue + ours = toa - sfc + lbl = reference["case2a"]["toa_rf"][config] - reference["case2a"]["sfc_rf"][config] + abs_diff = ours - lbl + rel_diff = abs_diff if lbl == 0.0 else abs_diff / abs(lbl) * 100.0 + rows.append( + { + "benchmark": "randles2013", + "case": "case2a", + "quantity": "atmospheric_rf", + "band": "bb", + "atm": atm, + "sza": sza, + "units": "W/m2", + "lbl": float(lbl), + "ours": ours, + "abs_diff": abs_diff, + "rel_diff": rel_diff, + "status": NA_STATUS, + } + ) + return rows + + +def _cell(row: dict) -> str: + """Render one row as a Markdown table cell: ``lbl / ours (rel) status``.""" + rel = ( + f"{row['rel_diff']:+.1f}%" + if row["lbl"] != 0.0 + else f"{row['rel_diff']:+.2f} (abs)" # percent fallback not defined + ) + status = row["status"] + tag = STATUS_EMOJI.get(status, status) + return f"{row['lbl']:g} / {row['ours']:g} ({rel}) {tag}" + + +def format_report(rows: list[dict]) -> str: + """Render comparison rows as Markdown. + + One table per case with quantities as rows and the four configs + (saw30/saw75/trop30/trop75) as columns; each cell shows + ``LBL / ours (rel diff) status-emoji``. A header line carries the global + counts (``N pass / N warn / N fail``) and each case heading repeats its + own counts. + """ + counts = {"PASS": 0, "WARN": 0, "FAIL": 0, NA_STATUS: 0} + for row in rows: + counts[row["status"]] = counts.get(row["status"], 0) + 1 + + lines: list[str] = [] + benchmark = rows[0]["benchmark"] if rows else "benchmark" + lines.append(f"# {benchmark} — pyRadtran vs LBL reference") + lines.append("") + summary = f"**{counts['PASS']} pass / {counts['WARN']} warn / {counts['FAIL']} fail**" + if counts[NA_STATUS]: + summary += f" (+{counts[NA_STATUS]} report-only n/a)" + lines.append(f"{summary} — {len(rows)} rows") + lines.append("") + lines.append( + f"Thresholds — fluxes: PASS |rel| ≤ {FLUX_PASS_PCT:g}% / WARN ≤ {FLUX_WARN_PCT:g}% / " + f"else FAIL; RF: PASS |rel| ≤ {RF_PASS_REL_PCT:g}% or |abs| ≤ {RF_PASS_ABS_W_M2:g} W/m2, " + f"WARN ≤ {RF_WARN_REL_PCT:g}% or ≤ {RF_WARN_ABS_W_M2:g} W/m2, else FAIL. " + f"Cells: LBL / ours (rel diff)." + ) + + if not rows: + lines.append("") + lines.append("(no rows)") + return "\n".join(lines) + "\n" + + case_order = [c for c in CASES if any(row["case"] == c for row in rows)] + for case in case_order: + case_rows = [row for row in rows if row["case"] == case] + cc = {"PASS": 0, "WARN": 0, "FAIL": 0, NA_STATUS: 0} + for row in case_rows: + cc[row["status"]] = cc.get(row["status"], 0) + 1 + heading = f"## {case} — {cc['PASS']} pass / {cc['WARN']} warn / {cc['FAIL']} fail" + if cc[NA_STATUS]: + heading += f" (+{cc[NA_STATUS]} n/a)" + lines.append("") + lines.append(heading) + lines.append("") + lines.append("| quantity | " + " | ".join(CONFIGS) + " |") + lines.append("| --- | " + " | ".join(["---"] * len(CONFIGS)) + " |") + by_key: dict[tuple[str, str], dict] = { + (row["quantity"], f"{row['atm']}{row['sza']}"): row for row in case_rows + } + quantities = list(dict.fromkeys(row["quantity"] for row in case_rows)) + for quantity in quantities: + units = next(r["units"] for r in case_rows if r["quantity"] == quantity) + cells = [] + for config in CONFIGS: + row = by_key.get((quantity, config)) + cells.append(_cell(row) if row else "—") + lines.append(f"| {quantity} ({units}) | " + " | ".join(cells) + " |") + return "\n".join(lines) + "\n" + + +def write_report(rows: list[dict], path: str | Path) -> tuple[Path, Path]: + """Write the Markdown report to ``path`` and a CSV to the same stem. + + The CSV uses :data:`CSV_COLUMNS` with plain status words (PASS/WARN/FAIL, + ``n/a``) instead of emojis. + + Returns: + ``(markdown_path, csv_path)``. + """ + md_path = Path(path) + md_path.parent.mkdir(parents=True, exist_ok=True) + md_path.write_text(format_report(rows), encoding="utf-8") + + csv_path = md_path.with_suffix(".csv") + with open(csv_path, "w", newline="", encoding="utf-8") as f: + writer = csv.DictWriter(f, fieldnames=list(CSV_COLUMNS)) + writer.writeheader() + for row in rows: + writer.writerow({col: _csv_cell(row.get(col)) for col in CSV_COLUMNS}) + return md_path, csv_path + + +def _csv_cell(value) -> str: + if isinstance(value, float): + return f"{value:.6g}" + return str(value) + + +# Short tick labels for the overlay panels (full names stay in the report). +_QUANTITY_SHORT: dict[str, str] = { + "direct_bb_sfc_down": "direct", + "diffuse_bb_sfc_down": "diffuse", + "diffuse_bb_toa_up": "diffuse", + "uvvis_sfc_down": "uvvis", + "nir_sfc_down": "nir", + "total_bb_sfc_down": "total", + "toa_rf": "toa", + "sfc_rf": "sfc", + "atmospheric_rf": "atm", +} + + +def plot_benchmark_overlay(rows: list[dict], path: str | Path) -> Path: + """Grouped LBL-vs-pyRadtran bar chart per case and save it as PNG. + + One row of panels per case: surface-downwelling fluxes, TOA-upwelling + fluxes, and — for the aerosol cases 2a/2b — radiative forcing. Each group + is one (config, quantity) pair with two adjacent bars (LBL reference in + gray, pyRadtran in blue). The dimensionless ``absorptance`` is not plotted + (it would break the shared W/m2 flux axis); it stays in the report. + + matplotlib is imported lazily (Agg-safe headless rendering). + + Args: + rows: Rows from :func:`compare_benchmark`. + path: Destination PNG path. + + Returns: + The saved path. + + Raises: + ValueError: If ``rows`` is empty. + """ + if not rows: + raise ValueError("plot_benchmark_overlay requires at least one comparison row") + + from pyradtran.viz._style import get_palette, require_mpl, set_theme + + require_mpl() + import matplotlib.pyplot as plt + + set_theme() + ours_color = get_palette(1)[0] + lbl_color = "#7f7f7f" + + by_key = {(row["case"], row["quantity"], f"{row['atm']}{row['sza']}"): row for row in rows} + case_order = [c for c in CASES if any(row["case"] == c for row in rows)] + + panels = ( + ("sfc_down", "SFC downwelling fluxes"), + ("toa_up", "TOA upwelling fluxes"), + ("rf", "Radiative forcing"), + ) + + def panel_quantities(case: str) -> dict[str, list[str]]: + groups: dict[str, list[str]] = {"sfc_down": [], "toa_up": [], "rf": []} + for quantity in dict.fromkeys(r["quantity"] for r in rows if r["case"] == case): + if quantity.endswith("_rf"): + groups["rf"].append(quantity) + elif quantity.endswith("_sfc_down"): + groups["sfc_down"].append(quantity) + elif quantity.endswith("_toa_up"): + groups["toa_up"].append(quantity) + return groups + + fig, axes = plt.subplots( + nrows=len(case_order), + ncols=3, + figsize=(15, 3.6 * len(case_order) + 0.8), + squeeze=False, + ) + fig.suptitle("Randles 2013 benchmark — LBL vs pyRadtran (W/m2)", fontsize=13) + + for i, case in enumerate(case_order): + groups = panel_quantities(case) + for j, (panel, title) in enumerate(panels): + ax = axes[i][j] + quantities = groups[panel] + if not quantities: + ax.set_visible(False) + continue + # Subset runs may cover only part of CONFIGS: skip absent + # (config, quantity) pairs — like format_report's "—" cells — + # so x positions stay contiguous instead of raising KeyError. + positions = [ + (config, quantity) + for config in CONFIGS + for quantity in quantities + if (case, quantity, config) in by_key + ] + left = [by_key[(case, quantity, config)]["lbl"] for config, quantity in positions] + right = [by_key[(case, quantity, config)]["ours"] for config, quantity in positions] + x = range(len(positions)) + width = 0.38 + ax.bar( + [p - width / 2 for p in x], + left, + width=width, + color=lbl_color, + label="LBL reference", + ) + ax.bar( + [p + width / 2 for p in x], right, width=width, color=ours_color, label="pyRadtran" + ) + ax.set_xticks( + list(x), + [ + f"{config}\n{_QUANTITY_SHORT.get(quantity, quantity)}" + for config, quantity in positions + ], + fontsize=7, + ) + if panel == "rf": + ax.axhline(0.0, color="black", linewidth=0.8) + ax.set_title(f"{case} — {title}", fontsize=10) + if j == 0: + ax.set_ylabel(f"{case}\nW/m2") + if i == 0 and j == 0: + ax.legend(loc="best", fontsize=8) + + out = Path(path) + out.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(out, dpi=150, bbox_inches="tight") + plt.close(fig) + return out diff --git a/src/pyradtran/benchmarks/randles2013.py b/src/pyradtran/benchmarks/randles2013.py new file mode 100644 index 0000000..c8fa012 --- /dev/null +++ b/src/pyradtran/benchmarks/randles2013.py @@ -0,0 +1,394 @@ +"""Randles et al. (2013) shortwave benchmark: case matrix + runner. + +Reproduces the shortwave radiative-transfer intercomparison protocol of +Randles, C. A. et al. (2013), "Intercomparison of shortwave radiative +transfer schemes to improve the aerosol and surface albedo treatment in +climate models", Atmos. Chem. Phys. 13, 2347-2362 (Tables 2/3/5/A3/A4). + +Protocol +-------- +Case 1 Rayleigh-only (no aerosol). +Case 2a AOD550 = 0.2, Angstroem exponent 1, SSA = 1.0, g = 0.7. +Case 2b As 2a but SSA = 0.8. +Surface albedo 0.2 Lambert; SZA {30, 75} deg; atmospheres AFGL subarctic +winter ("saw") and tropics ("trop"); no clouds. Aerosol confined to the +lowest 2 km, linearly decreasing to zero at 2 km. Bands: broadband +0.2-4.0 um ("bb") and UV-VIS 0.2-0.7 um ("uvvis"). + +Normalization +------------- +Each model divides its own band TOA incident flux and multiplies by the +LBL median constants (W/m2): BB {SZA30: 1189.28, SZA75: 355.43}, +UVVIS {SZA30: 563.38, SZA75: 168.37}. This removes the solar-constant +choice. Formally ``F_norm = F_raw / edir_toa_raw(same run) * C[band][sza]``. + +Replication assumptions +----------------------- +The paper leaves details to each participating model; we fix them as: + +1. Built-in AFGL profiles are used (the paper's participants used their + individual-model profiles; the comparison tolerance absorbs this). +2. "Aerosol linear in the lowest 2 km" = a linear taper from the surface + to zero at 2 km, ``w(z) = max(0, 1 - z / 2 km)``. +3. g = 0.7 Henyey-Greenstein phase function represented through PMOM + moments ``beta_l = g**l`` (the DISORT/libRadtran coefficient form). +4. Angstroem exponent 1 means ``tau(lambda) = 0.2 * (lambda / 0.55 um)**-1``. + +Comparison thresholds live in the reference file ``_meta`` block +(fluxes: PASS within +/-8%, WARN within +/-12%, else FAIL; RF: PASS if +|diff| <= 15% OR <= 1.5 W/m2). +""" + +from __future__ import annotations + +import json +import os +from collections.abc import Sequence +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import xarray as xr + +from pyradtran.core.runner import Runner +from pyradtran.models.aerosol_composite import CompositeAerosol, LayerOptics +from pyradtran.scene import Scene + +#: Benchmark cases: Rayleigh-only, and the two aerosol loadings. +CASES: tuple[str, ...] = ("case1", "case2a", "case2b") + +#: Atmosphere keys -> pyRadtran AFGL profile aliases (sw/trop). +ATMOSPHERES: dict[str, str] = {"saw": "sw", "trop": "tp"} + +#: Solar zenith angles in degrees. +SZAS: tuple[int, ...] = (30, 75) + +#: Spectral bands: key -> (min_nm, max_nm). +BANDS: dict[str, tuple[float, float]] = { + "bb": (200.0, 4000.0), + "uvvis": (200.0, 700.0), +} + +#: LBL median band TOA incident flux constants (W/m2) per band and SZA. +NORMALIZATION_CONSTANTS: dict[str, dict[int, float]] = { + "bb": {30: 1189.28, 75: 355.43}, + "uvvis": {30: 563.38, 75: 168.37}, +} + +#: Lambertian surface albedo (protocol Table 2). +SURFACE_ALBEDO = 0.2 + +#: Aerosol optical depth at the reference wavelength. +AOD_550 = 0.2 +#: Reference wavelength for the Angstroem-1 power law (um). +WL_REF_UM = 0.55 +#: Aerosol top altitude (km): linear taper reaches zero here. +AEROSOL_TOP_KM = 2.0 + +#: Vertical grid (km, strictly descending): 3 layers, lowest 2 aerosol-bearing. +ALTITUDE_GRID_KM: list[float] = [3.0, 2.0, 1.0, 0.0] + +#: Number of Legendre moments written for the aerosol phase function. +N_LEGENDRE = 32 + +#: Wavelength grid (um) for the explicit aerosol layer files. Geometrically +#: spaced so uvspec's linear-in-wavelength interpolation of the Angstroem-1 +#: power law stays accurate (max relative error ~3e-4) across both bands. +AEROSOL_WAVELENGTH_GRID_UM: list[float] = np.geomspace(0.2, 4.0, 65).tolist() + +#: Path of the bundled LBL reference file. +_REFERENCE_PATH = Path(__file__).resolve().parent / "reference" / "randles2013_lbl.json" + + +def load_reference() -> dict: + """Load the bundled LBL reference values (Tables 3/5/A3/A4). + + Returns: + Dict with keys ``case1``/``case2a``/``case2b``, each mapping a + quantity name to ``{"saw30": v, "saw75": v, "trop30": v, "trop75": v}``, + plus a ``_meta`` block (citation, normalization constants, thresholds). + """ + with open(_REFERENCE_PATH) as f: + return json.load(f) + + +def taper_layer_weights(altitude_km) -> np.ndarray: + """Per-layer column weights of the linear taper ``w(z) = max(0, 1 - z/2)``. + + The weight of layer *i* (between consecutive descending grid levels) is + ``integral_layer w(z) dz / integral_0^2km w(z) dz``. ``w`` is piecewise + linear, so the trapezoid over the layer endpoints is exact. + + Returns: + Array of shape ``(n_layer,)`` summing to 1 on a grid covering 0-2 km. + """ + z = np.asarray(altitude_km, dtype=float) + w = np.clip(1.0 - z / AEROSOL_TOP_KM, 0.0, None) + dz = z[:-1] - z[1:] # positive: grid is strictly descending + integrals = 0.5 * (w[:-1] + w[1:]) * dz + total = float(integrals.sum()) + if total <= 0.0: + raise ValueError("altitude grid does not intersect the 0-2 km aerosol layer") + return integrals / total + + +@dataclass(frozen=True) +class RandlesAerosol: + """Piece implementing the Randles 2013 case-2 aerosol (duck-typed Piece). + + Column optical depth follows the Angstroem-1 law + ``tau_col(lambda) = AOD_550 * (lambda / 0.55 um)**-1`` and is distributed + over layers with the linear-taper weights (assumption 2). SSA is constant + per case (2a: 1.0, 2b: 0.8), g is constant 0.7, and the phase function is + Henyey-Greenstein via PMOM moments ``beta_l = g**l`` (assumption 3). + """ + + ssa: float + g: float = 0.7 + + @property + def name(self) -> str: + """Signature used for explicit-aerosol file naming.""" + return f"RandlesAerosol(ssa={self.ssa:g})" + + def to_layer_optics(self, wl_um, altitude_km, n_legendre: int = 32) -> LayerOptics: + """Return LayerOptics for this piece on the given grids. + + Args: + wl_um: Wavelengths (um), shape (n_wl,). + altitude_km: Descending layer-boundary altitudes (km), len n_layer+1. + n_legendre: Number of Legendre moments. + + Returns: + LayerOptics with tau/ssa/g of shape (n_wl, n_layer) and + legendre_moments of shape (n_wl, n_layer, n_legendre). + """ + wl = np.atleast_1d(np.asarray(wl_um, dtype=float)) + weights = taper_layer_weights(altitude_km) + n_layer = len(weights) + + tau_col = AOD_550 * (wl / WL_REF_UM) ** -1.0 + tau = np.outer(tau_col, weights) + ssa = np.full((len(wl), n_layer), float(self.ssa)) + g = np.full((len(wl), n_layer), float(self.g)) + beta = self.g ** np.arange(n_legendre, dtype=float) + moments = np.broadcast_to(beta, (len(wl), n_layer, n_legendre)).copy() + return LayerOptics(tau=tau, ssa=ssa, g=g, legendre_moments=moments) + + +def build_scene( + case: str, + atm: str, + sza: int, + band: str, + *, + aerosol_output_dir: str | os.PathLike | None = None, +) -> Scene: + """Build one benchmark Scene. + + Args: + case: One of :data:`CASES`. ``case1`` omits the aerosol entirely. + atm: Key of :data:`ATMOSPHERES` ("saw" or "trop"). + sza: Solar zenith angle (30 or 75). + band: Key of :data:`BANDS` ("bb" or "uvvis"). + aerosol_output_dir: Directory for the explicit aerosol layer files + (aerosol cases only). Defaults to ``./aerosol`` (CompositeAerosol). + + Returns: + A Scene with AFGL profile + reptran, DISORT 16 streams + + ``disort_intcor moments`` + pseudospherical, albedo 0.2, and ASCII + band-integrated (``output_process integrate``) output at surface and + TOA with quantities [lambda, edir, edn, eup]. + """ + if case not in CASES: + raise ValueError(f"Unknown case '{case}'. Valid: {list(CASES)}") + if atm not in ATMOSPHERES: + raise ValueError(f"Unknown atmosphere '{atm}'. Valid: {list(ATMOSPHERES)}") + if sza not in SZAS: + raise ValueError(f"Unknown sza {sza}. Valid: {list(SZAS)}") + if band not in BANDS: + raise ValueError(f"Unknown band '{band}'. Valid: {list(BANDS)}") + + scene = ( + Scene() + .set_atmosphere(profile=ATMOSPHERES[atm], mol_abs_param="reptran") + .set_source_solar(sza=float(sza)) + .set_wavelength(*BANDS[band]) + .set_solver(method="disort", streams=16, disort_intcor="moments", pseudospherical=True) + .set_surface(albedo=SURFACE_ALBEDO) + .set_output( + quantities=["lambda", "edir", "edn", "eup"], + format="ascii", + zout=[0, "toa"], + process="integrate", + ) + ) + if case != "case1": + ssa = 1.0 if case == "case2a" else 0.8 + scene = scene.set_aerosol( + CompositeAerosol( + pieces=[RandlesAerosol(ssa=ssa)], + wavelength_grid_um=AEROSOL_WAVELENGTH_GRID_UM, + altitude_grid_km=ALTITUDE_GRID_KM, + n_legendre=N_LEGENDRE, + output_dir=Path(aerosol_output_dir) if aerosol_output_dir else None, + ) + ) + return scene + + +def _extract_fluxes(ds: xr.Dataset) -> dict[str, float]: + """Pull surface/TOA band fluxes from one parsed ``integrate`` dataset. + + The parsed dataset has dims (wavelength: 1, zout: 2); wavelength is the + collapsed band dimension of ``output_process integrate`` output. + """ + zout = np.asarray(ds["zout"].values, dtype=float) + i_sfc = int(np.argmin(zout)) + i_toa = int(np.argmax(zout)) + + def flux(quantity: str, i_level: int) -> float: + arr = np.asarray(ds[quantity].values, dtype=float) + if arr.ndim != 2 or arr.shape[0] != 1: + raise ValueError( + f"Expected band-integrated output (single wavelength) for " + f"{quantity}; got shape {arr.shape}" + ) + return float(arr[0, i_level]) + + return { + "edir_sfc": flux("edir", i_sfc), + "edn_sfc": flux("edn", i_sfc), + "eup_sfc": flux("eup", i_sfc), + "edir_toa": flux("edir", i_toa), + "eup_toa": flux("eup", i_toa), + } + + +def _prewrite_aerosol_files(scenes: list[Scene]) -> None: + """Write explicit aerosol files serially before threaded execution. + + All runs of one aerosol case share the same content hash, hence the same + layer files. Writing them here (in the main thread) means the worker + threads only hit the on-disk cache — no concurrent writers on one file. + """ + seen: set[str] = set() + for scene in scenes: + aerosol = scene.aerosol + if aerosol is None: + continue + sig = str(aerosol) + if sig in seen: + continue + aerosol.to_uvspec_lines() + seen.add(sig) + + +def run_randles2013( + output_dir: str | os.PathLike, + *, + uvspec_exe: str | None = None, + data_path: str | None = None, + cases: Sequence[str] | None = None, +) -> dict: + """Run the Randles 2013 case matrix and return normalized fluxes. + + Executes ``len(cases) * 2 atmospheres * 2 SZAs * 2 bands`` uvspec + invocations (24 with all cases) via :meth:`Runner.execute_many`, then + normalizes every flux ``F_norm = F_raw / edir_toa_raw(same run) * + C[band][sza]`` with the LBL median constants. + + Args: + output_dir: Directory for the results file and aerosol layer files. + uvspec_exe: Path to the uvspec binary (auto-detected if None). + data_path: libRadtran data directory (auto-detected if None). + cases: Subset of :data:`CASES` to run (default: all). RF entries for + the aerosol cases are only produced when ``case1`` is included, + since RF is the case-2 minus case-1 difference. + + Returns: + Nested dict ``{case: {band: {atm: {sza: {quantity: value}}}}}`` with + quantities ``edir_sfc``, ``edn_sfc``, ``eup_sfc``, ``eup_toa``, + ``total_sfc_down`` (all normalized W/m2); case1 broadband additionally + carries ``absorptance`` (dimensionless, same-run ratio) and + ``nir_sfc_down`` (bb total minus uvvis total); the aerosol cases + additionally carry ``rf_toa``/``rf_sfc`` (net case-2 minus case-1; + net_toa = edir - eup at TOA, net_sfc = edir + edn - eup at surface). + The whole dict is saved as ``/randles2013_results.json`` + and its path is returned under key ``results_path``. + + Raises: + ValueError: If ``cases`` contains an unknown case name. + RuntimeError: If any uvspec invocation fails (propagated from + :meth:`Runner.execute_many`). + """ + out_dir = Path(output_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + selected = CASES if cases is None else tuple(dict.fromkeys(cases)) + unknown = set(selected) - set(CASES) + if unknown: + raise ValueError(f"Unknown case(s) {sorted(unknown)}. Valid: {list(CASES)}") + + runs = [ + (case, atm, sza, band) + for case in selected + for atm in ATMOSPHERES + for sza in SZAS + for band in BANDS + ] + scenes = [ + build_scene(case, atm, sza, band, aerosol_output_dir=out_dir / "aerosol") + for case, atm, sza, band in runs + ] + _prewrite_aerosol_files(scenes) + + datasets = Runner.execute_many(scenes, uvspec_exe=uvspec_exe, data_path=data_path) + + results: dict = {} + raw_by_run: dict[tuple[str, str, str, int], dict[str, float]] = {} + for (case, atm, sza, band), ds in zip(runs, datasets, strict=True): + raw = _extract_fluxes(ds) + raw_by_run[(case, band, atm, sza)] = raw + scale = NORMALIZATION_CONSTANTS[band][sza] / raw["edir_toa"] + entry = { + "edir_sfc": raw["edir_sfc"] * scale, + "edn_sfc": raw["edn_sfc"] * scale, + "eup_sfc": raw["eup_sfc"] * scale, + "eup_toa": raw["eup_toa"] * scale, + "total_sfc_down": (raw["edir_sfc"] + raw["edn_sfc"]) * scale, + } + results.setdefault(case, {}).setdefault(band, {}).setdefault(atm, {})[sza] = entry + + if "case1" in results: + for atm in ATMOSPHERES: + for sza in SZAS: + raw_bb = raw_by_run[("case1", "bb", atm, sza)] + # Dimensionless same-run ratio: C cancels, so compute from raw. + absorbed_surf = (1.0 - SURFACE_ALBEDO) * (raw_bb["edir_sfc"] + raw_bb["edn_sfc"]) + absorptance = 1.0 - (raw_bb["eup_toa"] + absorbed_surf) / raw_bb["edir_toa"] + bb = results["case1"]["bb"][atm][sza] + uvvis = results["case1"]["uvvis"][atm][sza] + bb["absorptance"] = absorptance + bb["nir_sfc_down"] = bb["total_sfc_down"] - uvvis["total_sfc_down"] + + for case in ("case2a", "case2b"): + if case not in results: + continue + for band in BANDS: + for atm in ATMOSPHERES: + for sza in SZAS: + base = results["case1"][band][atm][sza] + pert = results[case][band][atm][sza] + # net_toa = edir_toa - eup_toa = C - eup_toa: the C + # constants cancel in the case-2 minus case-1 difference. + pert["rf_toa"] = base["eup_toa"] - pert["eup_toa"] + pert["rf_sfc"] = (pert["total_sfc_down"] - pert["eup_sfc"]) - ( + base["total_sfc_down"] - base["eup_sfc"] + ) + + results_path = out_dir / "randles2013_results.json" + results["results_path"] = str(results_path) + with open(results_path, "w") as f: + json.dump(results, f, indent=2) + return results diff --git a/src/pyradtran/benchmarks/reference/randles2013_lbl.json b/src/pyradtran/benchmarks/reference/randles2013_lbl.json new file mode 100644 index 0000000..d6f6995 --- /dev/null +++ b/src/pyradtran/benchmarks/reference/randles2013_lbl.json @@ -0,0 +1,42 @@ +{ + "_meta": { + "paper": "Randles, C. A. et al. (2013): Intercomparison of shortwave radiative transfer schemes to improve the aerosol and surface albedo treatment in climate models. Atmospheric Chemistry and Physics, 13, 2347-2362.", + "tables": ["Table 3", "Table 5", "Table A3", "Table A4"], + "protocol": "Case 1 Rayleigh-only; Case 2a AOD550 0.2, Angstroem 1, SSA 1.0, g 0.7; Case 2b SSA 0.8 (rest same as 2a). Surface albedo 0.2 Lambert. SZA 30/75 deg. AFGL subarctic winter (saw) + tropics (trop). Aerosol confined to lowest 2 km, linearly decreasing to zero at 2 km. No clouds.", + "units": "W/m2 (absorptance: dimensionless)", + "normalization": "Each model divides its own band TOA incident flux and multiplies by the LBL median constants (removes the solar-constant choice).", + "normalization_constants_w_m2": { + "bb": {"sza30": 1189.28, "sza75": 355.43}, + "uvvis": {"sza30": 563.38, "sza75": 168.37} + }, + "thresholds": { + "flux": {"pass_pct": 8, "warn_pct": 12}, + "rf": {"pass_rel_pct": 15, "pass_abs_w_m2": 1.5}, + "case2a_atmospheric_rf": "report-only (<1 W/m2; RSD invalid per paper)" + }, + "quantity_naming": "Quantities tagged bb refer to the 0.2-4.0 um broadband, uvvis to 0.2-0.7 um. toa_rf/sfc_rf are broadband case-2 minus case-1 differences. Config keys: {saw,trop}{30,75}.", + "value_count": 64 + }, + "case1": { + "direct_bb_sfc_down": {"saw30": 942.4, "saw75": 216.2, "trop30": 844.5, "trop75": 179.6}, + "diffuse_bb_sfc_down": {"saw30": 64.4, "saw75": 37.2, "trop30": 64.0, "trop75": 36.8}, + "diffuse_bb_toa_up": {"saw30": 227.6, "saw75": 82.6, "trop30": 204.7, "trop75": 75.2}, + "uvvis_sfc_down": {"saw30": 489.2, "saw75": 115.8, "trop30": 489.1, "trop75": 115.7}, + "nir_sfc_down": {"saw30": 519.1, "saw75": 138.0, "trop30": 421.1, "trop75": 101.1}, + "absorptance": {"saw30": 0.134, "saw75": 0.201, "trop30": 0.221, "trop75": 0.307} + }, + "case2a": { + "total_bb_sfc_down": {"saw30": 994.6, "saw75": 226.5, "trop30": 896.0, "trop75": 192.8}, + "diffuse_bb_toa_up": {"saw30": 236.3, "saw75": 102.9, "trop30": 212.9, "trop75": 93.2}, + "uvvis_sfc_down": {"saw30": 480.6, "saw75": 101.5, "trop30": 480.2, "trop75": 101.5}, + "toa_rf": {"saw30": -8.6, "saw75": -20.3, "trop30": -8.2, "trop75": -18.0}, + "sfc_rf": {"saw30": -9.7, "saw75": -21.5, "trop30": -10.0, "trop75": -18.9} + }, + "case2b": { + "total_bb_sfc_down": {"saw30": 954.1, "saw75": 206.2, "trop30": 857.7, "trop75": 174.5}, + "diffuse_bb_toa_up": {"saw30": 216.1, "saw75": 89.8, "trop30": 194.4, "trop75": 81.7}, + "uvvis_sfc_down": {"saw30": 452.4, "saw75": 89.0, "trop30": 452.0, "trop75": 89.0}, + "toa_rf": {"saw30": 11.6, "saw75": -7.2, "trop30": 10.3, "trop75": -6.5}, + "sfc_rf": {"saw30": -42.1, "saw75": -37.8, "trop30": -40.6, "trop75": -33.6} + } +} diff --git a/src/pyradtran/config/__init__.py b/src/pyradtran/config/__init__.py new file mode 100644 index 0000000..be9c2f8 --- /dev/null +++ b/src/pyradtran/config/__init__.py @@ -0,0 +1,29 @@ +"""YAML configuration front-end for pyRadtran. + +A config file describes one experiment (scene + aerosol + analysis intents); +:func:`~pyradtran.config.load_config` maps it onto the same builders the +Python API uses, so YAML and API-built scenes are interchangeable. +""" + +from pyradtran.config.loader import ( + LoadedConfig, + build_aerosol, + build_scene, + export_config, + load_config, +) +from pyradtran.config.orchestrator import DrfResult, RunResult, run_config +from pyradtran.config.schema import CONFIG_VERSION, PyRadtranConfig + +__all__ = [ + "CONFIG_VERSION", + "DrfResult", + "LoadedConfig", + "PyRadtranConfig", + "RunResult", + "build_aerosol", + "build_scene", + "export_config", + "load_config", + "run_config", +] diff --git a/src/pyradtran/config/cli.py b/src/pyradtran/config/cli.py new file mode 100644 index 0000000..93030ac --- /dev/null +++ b/src/pyradtran/config/cli.py @@ -0,0 +1,141 @@ +"""Command-line interface: ``pyradtran run | validate | export-config``. + +argparse stdlib only. :func:`main` returns a process exit code so it can be +called directly from tests; config load/validation errors are reported as a +single concise stderr message, while anything else propagates with its +traceback (far more useful during development). +""" + +from __future__ import annotations + +import argparse +import sys + +import numpy as np +from pydantic import ValidationError + +from pyradtran.config.loader import export_config, load_config +from pyradtran.config.orchestrator import RunResult, run_config +from pyradtran.config.schema import PyRadtranConfig + +#: Wavelength (nm) at which budget/DRF scalars are summarized. +_SUMMARY_NM = 550.0 + + +def main(argv: list[str] | None = None) -> int: + """CLI entry point. ``argv`` defaults to ``sys.argv[1:]``.""" + args = _build_parser().parse_args(argv) + try: + return args.func(args) + except (ValidationError, ValueError, TypeError, FileNotFoundError) as e: + print(f"error: {e}", file=sys.stderr) + return 1 + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="pyradtran", + description="Run libRadtran experiments described by YAML configs.", + ) + sub = parser.add_subparsers(dest="command", required=True) + + run = sub.add_parser("run", help="execute a config: main run + analysis intents") + run.add_argument("config", help="path to the YAML config") + run.add_argument("--uvspec", help="path to the uvspec executable") + run.add_argument("--data-path", help="libRadtran data directory") + run.add_argument("--plot-dir", help="directory for requested plots") + run.set_defaults(func=_cmd_run) + + validate = sub.add_parser("validate", help="validate a config (no uvspec invocation)") + validate.add_argument("config", help="path to the YAML config") + validate.set_defaults(func=_cmd_validate) + + export = sub.add_parser("export-config", help="write a config as canonical YAML") + export.add_argument("config", help="path to the YAML config") + export.add_argument("--output", "-o", required=True, help="output YAML path") + export.set_defaults(func=_cmd_export) + + return parser + + +# --- Commands --- + + +def _cmd_run(args: argparse.Namespace) -> int: + loaded = load_config(args.config) + result = run_config( + loaded, + uvspec_exe=args.uvspec, + data_path=args.data_path, + plot_dir=args.plot_dir, + ) + _print_summary(loaded.config, result) + return 0 + + +def _cmd_validate(args: argparse.Namespace) -> int: + cfg = load_config(args.config).config + profile = cfg.scene.atmosphere.get("profile", "?") + n_blocks = len(cfg.aerosol.blocks) if cfg.aerosol else 0 + print(f"ok: {cfg.name} — scene {profile}, {n_blocks} aerosol blocks, intents: {_intents(cfg)}") + return 0 + + +def _cmd_export(args: argparse.Namespace) -> int: + cfg = load_config(args.config).config + print(export_config(cfg, args.output)) + return 0 + + +# --- Summary rendering --- + + +def _print_summary(cfg: PyRadtranConfig, result: RunResult) -> None: + """Human-readable one-block-per-topic digest of a finished run.""" + print(f"config: {cfg.name}") + sizes = " ".join(f"{dim}={n}" for dim, n in result.rt.sizes.items()) + print(f"rt: {sizes} — data_vars: {', '.join(result.rt.data_vars)}") + + if result.budget is not None: + b, i = result.budget, _nearest_index(result.budget.wavelength) + print( + f"energy budget @{b.wavelength[i]:.0f} nm: " + f"F_inc={b.f_incident[i]:.4g} up_TOA={b.f_up_toa[i]:.4g} " + f"abs_surf={b.f_abs_surface[i]:.4g} abs_atm={b.f_abs_atm[i]:.4g} W/m²" + ) + if result.drf is not None: + d, i = result.drf, _nearest_index(result.drf.wavelength_nm) + print( + f"drf @{d.wavelength_nm[i]:.0f} nm: " + f"TOA={d.toa[i]:.4g} surface={d.surface[i]:.4g} " + f"atmosphere={d.atmosphere[i]:.4g} W/m²" + ) + if "heating_rate" in result.rt: + print("heating_rate: merged") + if result.attribution is not None: + print(f"attribution: {', '.join(result.attribution.contributions)}") + for fig in result.figures: + print(f"figure: {fig}") + if result.netcdf_path is not None: + print(f"netcdf: {result.netcdf_path}") + + +def _nearest_index(wavelengths: np.ndarray) -> int: + """Index of the grid point closest to the summary wavelength.""" + return int(np.argmin(np.abs(np.asarray(wavelengths, dtype=float) - _SUMMARY_NM))) + + +def _intents(cfg: PyRadtranConfig) -> list[str]: + """Names of the enabled analysis intents, in schema declaration order.""" + analysis = cfg.analysis + if analysis is None: + return [] + intents: list[str] = [] + if analysis.energy_conservation is not None: + intents.append("energy_conservation") + intents.extend(flag for flag in ("drf", "heating", "attribution") if getattr(analysis, flag)) + if analysis.plots: + intents.append("plots") + if analysis.save_netcdf: + intents.append("save_netcdf") + return intents diff --git a/src/pyradtran/config/loader.py b/src/pyradtran/config/loader.py new file mode 100644 index 0000000..f2626d7 --- /dev/null +++ b/src/pyradtran/config/loader.py @@ -0,0 +1,223 @@ +"""Config loader: YAML document → validated config → Scene + CompositeAerosol. + +The loader calls the same builder chain as an API user would, so a YAML +config and the equivalent Python code produce identical uvspec input text. +This equivalence is enforced by the round-trip regression test +(``tests/test_config_roundtrip.py``). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +import yaml + +from pyradtran.config.schema import ( + CONFIG_VERSION, + AerosolSection, + BlockSpec, + BulkBlockSpec, + ExponentialPlacement, + MassPlacement, + MieBlockSpec, + OdInversionPlacement, + OpacPresetBlockSpec, + PyRadtranConfig, + SceneSection, +) +from pyradtran.models.aerosol import OpacPreset +from pyradtran.models.aerosol_composite import ( + BulkSpecies, + CompositeAerosol, + IntegrationConfig, + MieSpecies, +) +from pyradtran.models.blocks import ( + DirectLayerOpticsBlock, + ExponentialProfile, + MassProfile, + PlacedBlock, + TabulatedProfile, + od_to_mass_profile, +) +from pyradtran.scene import Scene + + +@dataclass +class LoadedConfig: + """A validated config plus its built runtime objects.""" + + config: PyRadtranConfig + scene: Scene + aerosol: CompositeAerosol | None + + +def load_config(source: str | Path | dict | PyRadtranConfig) -> LoadedConfig: + """Load a config from a YAML file path, a raw dict, or a built config. + + Raises :class:`pydantic.ValidationError` / ``ValueError`` on invalid input + (unknown ``config_version``, unknown block ``kind``, typo'd keys, ...). + """ + if isinstance(source, PyRadtranConfig): + cfg = source + else: + if isinstance(source, str | Path): + path = Path(source) + raw = yaml.safe_load(path.read_text(encoding="utf-8")) + if not isinstance(raw, dict): + raise ValueError(f"{path}: config must be a YAML mapping") + elif isinstance(source, dict): + raw = source + else: + raise TypeError(f"unsupported config source: {type(source)!r}") + cfg = PyRadtranConfig.model_validate(raw) + aerosol = build_aerosol(cfg.aerosol) if cfg.aerosol is not None else None + scene = build_scene(cfg.scene, aerosol=aerosol) + return LoadedConfig(config=cfg, scene=scene, aerosol=aerosol) + + +def export_config(source: dict | PyRadtranConfig, path: str | Path) -> Path: + """Validate and serialize a config as canonical YAML. + + Accepts a raw dict (validated on the way out) or a built config; returns + the written path. The output re-loads identically + (``load_config(export_config(cfg, p))`` builds the same scene). + """ + cfg = source if isinstance(source, PyRadtranConfig) else PyRadtranConfig.model_validate(source) + out = Path(path) + out.write_text(yaml.safe_dump(cfg.model_dump(mode="json"), sort_keys=False), encoding="utf-8") + return out + + +# --- Scene assembly (mirrors examples/multicomponent_viz/canonical.py) --- + + +def build_scene(section: SceneSection, aerosol: CompositeAerosol | None = None) -> Scene: + """Map a validated ``scene`` section onto the Scene builder. + + Calls mirror the canonical example's builder chain exactly so the + generated uvspec text is byte-identical for equivalent inputs. + """ + scene = Scene().set_atmosphere(**section.atmosphere) + + src = dict(section.source) + src_kind = src.pop("source", "solar") + if src_kind == "solar": + if "sza" not in src: + raise ValueError("scene.source: 'sza' is required for a solar source") + scene = scene.set_source_solar(sza=src.pop("sza"), **src) + elif src_kind == "thermal": + scene = scene.set_source_thermal(**src) + else: + raise ValueError(f"scene.source: unknown source {src_kind!r}") + + wl = dict(section.wavelength) + if "min_nm" not in wl: + raise ValueError("scene.wavelength: 'min_nm' is required") + scene = scene.set_wavelength(wl.pop("min_nm"), wl.pop("max_nm", None), **wl) + + sol = dict(section.solver) + scene = scene.set_solver( + method=sol.pop("method", "disort"), + streams=sol.pop("streams", 16), + disort_intcor=sol.pop("disort_intcor", None), + pseudospherical=sol.pop("pseudospherical", False), + **sol, + ) + + if section.surface is not None: + scene = scene.set_surface(**section.surface) + if section.output: + scene = scene.set_output(**section.output) + if aerosol is not None: + scene = scene.set_aerosol(aerosol) + return scene + + +# --- Aerosol assembly (mirrors canonical.build_composite) --- + + +def build_aerosol(section: AerosolSection) -> CompositeAerosol: + """Assemble the LEGO blocks described by the ``aerosol`` section.""" + pieces = [piece for spec in section.blocks for piece in _pieces_for(spec, section)] + return CompositeAerosol( + pieces=pieces, + wavelength_grid_um=list(section.wavelength_grid_um), + altitude_grid_km=list(section.altitude_grid_km), + n_legendre=section.n_legendre, + output_dir=Path(section.output_dir) if section.output_dir else None, + ) + + +def _pieces_for(spec: BlockSpec, section: AerosolSection) -> list: + if isinstance(spec, MieBlockSpec): + species = MieSpecies( + refractive_index=spec.refractive_index, + size_distribution=spec.size_distribution, + particle_density_kg_m3=spec.particle_density_kg_m3, + integration_config=( + spec.integration if spec.integration is not None else IntegrationConfig() + ), + phase_function=spec.phase_function, + name=spec.name, + ) + return [PlacedBlock(block=species, profile=_placement(spec.placement, species, section))] + + if isinstance(spec, BulkBlockSpec): + species = BulkSpecies(bulk=_read_bulk(spec.file), name=spec.name) + return [PlacedBlock(block=species, profile=_placement(spec.placement, species, section))] + + if isinstance(spec, OpacPresetBlockSpec): + preset = OpacPreset( + name=spec.preset, + rh_pct=spec.rh_pct, + species_names=spec.species_names, + data_path=spec.data_path, + n_legendre=spec.n_legendre, + ) + return preset.to_placed_blocks() + + # ExplicitLayerBlockSpec: already a Piece (no placement needed) + return [DirectLayerOpticsBlock(master_path=spec.master_path, name=spec.name)] + + +def _read_bulk(file: str): + """Read an aerosol3D bulk-optics NetCDF (lazy import; optional dependency).""" + try: + from Aerosol3D.bulk.datastructs import BulkAerosolOpticsData + except ImportError as e: # pragma: no cover - depends on optional install + raise ImportError( + "config block kind 'bulk' requires the aerosol3D package " + "(Aerosol3D.bulk.datastructs.BulkAerosolOpticsData)" + ) from e + return BulkAerosolOpticsData.from_netcdf(file) + + +def _placement(placement, species, section: AerosolSection): + if isinstance(placement, OdInversionPlacement): + return od_to_mass_profile( + species, + tau_ref=placement.tau_ref, + ref_nm=placement.ref_nm, + altitude_km=list(section.altitude_grid_km), + scale_height_km=placement.scale_height_km, + ) + if isinstance(placement, MassPlacement): + return MassProfile(kg_m3_per_layer=tuple(placement.kg_m3_per_layer)) + if isinstance(placement, ExponentialPlacement): + return ExponentialProfile( + rho0_kg_m3=placement.rho0_kg_m3, scale_height_km=placement.scale_height_km + ) + # TabulatedPlacement + return TabulatedProfile(z_km=tuple(placement.z_km), kg_m3=tuple(placement.kg_m3)) + + +__all__ = [ + "CONFIG_VERSION", + "LoadedConfig", + "build_aerosol", + "build_scene", + "export_config", + "load_config", +] diff --git a/src/pyradtran/config/orchestrator.py b/src/pyradtran/config/orchestrator.py new file mode 100644 index 0000000..b0a327f --- /dev/null +++ b/src/pyradtran/config/orchestrator.py @@ -0,0 +1,244 @@ +"""Orchestrator: execute a config's main run plus its analysis intents. + +:func:`run_config` drives the full pipeline a config describes: + +1. main uvspec run (fluxes) of the loaded scene; +2. optional analyses — energy conservation, heating (second uvspec + invocation, merged), DRF vs a no-aerosol baseline, leave-one-out + attribution (N+1 runs); +3. optional plots (names resolved against a fixed registry) and NetCDF + persistence. + +The multi-run analyses mirror the proven flow of +``examples/multicomponent_viz/run_demo.py``. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass, field +from pathlib import Path + +import numpy as np +import xarray as xr + +from pyradtran.config.loader import LoadedConfig, build_scene, load_config +from pyradtran.config.schema import AnalysisSection, SceneSection +from pyradtran.core.postprocess import EnergyBudget, assert_energy_conservation +from pyradtran.core.runner import Runner +from pyradtran.workflow import AttributionResult, compute_component_attribution + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class DrfResult: + """Spectral direct radiative forcing (aerosol minus no-aerosol), W/m². + + Sign convention: negative TOA/surface = cooling (IPCC). + """ + + wavelength_nm: np.ndarray + toa: np.ndarray + surface: np.ndarray + + @property + def atmosphere(self) -> np.ndarray: + """TOA minus surface forcing (energy deposited in the column).""" + return self.toa - self.surface + + +@dataclass +class RunResult: + """Everything a config run produced.""" + + rt: xr.Dataset # main flux run (+ merged heating_rate if requested) + budget: EnergyBudget | None = None + drf: DrfResult | None = None + attribution: AttributionResult | None = None + figures: list[Path] = field(default_factory=list) + netcdf_path: Path | None = None + + +def run_config( + source: str | Path | dict | LoadedConfig, + *, + uvspec_exe: str | None = None, + data_path: str | None = None, + plot_dir: str | Path | None = None, +) -> RunResult: + """Load (if needed) and run a config end-to-end. + + Args: + source: YAML path, raw dict, or an already-loaded config. + uvspec_exe / data_path: forwarded to :meth:`Runner.execute`. + plot_dir: directory for requested plots (defaults to the config's + aerosol ``output_dir``, else the current directory). + + Returns: + :class:`RunResult`. Raises ``AssertionError`` if an energy-budget + intent is violated, or ``ValueError`` for unsatisfiable intents. + """ + loaded = source if isinstance(source, LoadedConfig) else load_config(source) + section = loaded.config.scene + analysis = loaded.config.analysis or AnalysisSection() + run = dict(uvspec_exe=uvspec_exe, data_path=data_path) + + rt = Runner.execute(loaded.scene, **run) + + budget = None + if analysis.energy_conservation is not None: + spec = analysis.energy_conservation + albedo = spec.albedo if spec.albedo is not None else _scene_albedo(section) + budget = assert_energy_conservation(rt, albedo=albedo, tol=spec.tol) + + if analysis.heating: + # libRadtran's heating_rate mode replaces flux output — second run. + heat_scene = build_scene(_heating_section(section), aerosol=loaded.aerosol) + rt_heat = Runner.execute(heat_scene, **run) + rt = rt.assign(heating_rate=rt_heat["heating_rate"]) + + drf = None + if analysis.drf: + _require_aerosol(loaded, "analysis.drf") + rt_noaer = Runner.execute(build_scene(section), **run) + drf = _compute_drf(rt, rt_noaer) + + attribution = None + if analysis.attribution: + _require_aerosol(loaded, "analysis.attribution") + + def execute_many(scenes): + # Sequential: Runner.execute_many swallows exceptions and cannot + # pickle CompositeAerosol scenes (fixed separately in T7). + return [Runner.execute(s, **run) for s in scenes] + + attribution = compute_component_attribution( + lambda comp: build_scene(section, aerosol=comp), loaded.aerosol, execute_many + ) + + figures = [] + if analysis.plots: + target = Path(plot_dir) if plot_dir is not None else _default_plot_dir(loaded) + target.mkdir(parents=True, exist_ok=True) + figures = _make_plots( + analysis.plots, + RunResult(rt=rt, budget=budget, drf=drf, attribution=attribution), + target, + ) + + netcdf_path = None + if analysis.save_netcdf: + netcdf_path = Path(analysis.save_netcdf) + rt.to_netcdf(netcdf_path) + + return RunResult( + rt=rt, + budget=budget, + drf=drf, + attribution=attribution, + figures=figures, + netcdf_path=netcdf_path, + ) + + +# --- helpers --- + + +def _scene_albedo(section: SceneSection) -> float: + surface = section.surface or {} + if "albedo" not in surface: + raise ValueError( + "analysis.energy_conservation: no 'albedo' given and the scene " + "surface defines none — set analysis.energy_conservation.albedo" + ) + return float(surface["albedo"]) + + +def _heating_section(section: SceneSection) -> SceneSection: + """Scene variant whose output drops quantities (switches uvspec to + heating-rate mode) — mirrors canonical.build_scene_heating.""" + return section.model_copy( + update={"output": {k: v for k, v in section.output.items() if k != "quantities"}} + ) + + +def _require_aerosol(loaded: LoadedConfig, intent: str) -> None: + if loaded.aerosol is None: + raise ValueError(f"{intent}: requires an 'aerosol' section in the config") + + +def _compute_drf(rt: xr.Dataset, rt_noaer: xr.Dataset) -> DrfResult: + surf = int(np.argmin(rt["zout"].values)) + toa = int(np.argmax(rt["zout"].values)) + + def net_toa(ds): + return (ds["edir"] - ds["eup"]).isel(zout=toa).values + + def net_sfc(ds): + return (ds["edir"] + ds["edn"] - ds["eup"]).isel(zout=surf).values + + return DrfResult( + wavelength_nm=np.asarray(rt["wavelength"].values, dtype=float), + toa=net_toa(rt) - net_toa(rt_noaer), + surface=net_sfc(rt) - net_sfc(rt_noaer), + ) + + +def _default_plot_dir(loaded: LoadedConfig) -> Path: + if loaded.config.aerosol is not None and loaded.config.aerosol.output_dir: + return Path(loaded.config.aerosol.output_dir) + return Path.cwd() + + +def _make_plots(names: list[str], result: RunResult, target: Path) -> list[Path]: + """Render the requested plots by name; returns the written PNG paths.""" + import matplotlib.pyplot as plt + + from pyradtran import viz + from pyradtran.core.postprocess import add_budget_vars + + registry = { + "rt_spectral": lambda: viz.plot_spectral(result.rt), + "rt_flux_profile_edir": lambda: viz.plot_flux_profile( + result.rt, variable="edir", wavelength_nm=550.0 + ), + "rt_flux_profile_edn": lambda: viz.plot_flux_profile( + result.rt, variable="edn", wavelength_nm=550.0 + ), + "rt_flux_profile_eup": lambda: viz.plot_flux_profile( + result.rt, variable="eup", wavelength_nm=550.0 + ), + "rt_budget": lambda: viz.plot_budget(add_budget_vars(result.rt)), + "rt_heating_rate": lambda: viz.plot_heating_rate(result.rt, wavelength_nm=550.0), + "rt_overview": lambda: viz.plot_rt_overview(result.rt, wavelength_nm=550.0), + "drf_spectral": lambda: viz.plot_drf_spectral( + result.drf.wavelength_nm, result.drf.toa, result.drf.surface + ), + "attribution_edir": lambda: viz.plot_component_attribution( + result.attribution, variable="edir", level="surface" + ), + "attribution_spectral": lambda: viz.plot_spectral_attribution( + result.attribution, variable="edir", level="surface" + ), + } + figures = [] + for name in names: + if name not in registry: + raise ValueError(f"unknown plot {name!r}; valid: {sorted(registry)}") + if result.drf is None and name.startswith("drf_"): + raise ValueError(f"plot {name!r} requires analysis.drf: true") + if result.attribution is None and name.startswith("attribution_"): + raise ValueError(f"plot {name!r} requires analysis.attribution: true") + if name == "rt_heating_rate" and "heating_rate" not in result.rt: + raise ValueError(f"plot {name!r} requires analysis.heating: true") + fig, _ = registry[name]() + path = target / f"{name}.png" + fig.savefig(path, dpi=150, bbox_inches="tight") + plt.close(fig) + figures.append(path) + logger.info(" plot -> %s", path) + return figures + + +__all__ = ["DrfResult", "RunResult", "run_config"] diff --git a/src/pyradtran/config/schema.py b/src/pyradtran/config/schema.py new file mode 100644 index 0000000..9bdfdf2 --- /dev/null +++ b/src/pyradtran/config/schema.py @@ -0,0 +1,213 @@ +"""YAML configuration schema (``config_version: 1``). + +A config file describes one radiative-transfer experiment: + +- ``scene`` mirrors the :class:`~pyradtran.scene.Scene` builder kwargs; +- ``aerosol`` assembles LEGO blocks (discriminated on ``kind``) with a + column ``placement``; +- ``analysis`` declares post-processing intents (orchestrated by + :mod:`pyradtran.config.orchestrator`). + +The loader (:mod:`pyradtran.config.loader`) maps a validated config onto the +same builder calls the Python API uses, so a YAML config and an equivalent +API-built scene produce identical uvspec input text. +""" + +from __future__ import annotations + +from typing import Annotated, Any, Literal + +from pydantic import BaseModel, ConfigDict, Field + +from pyradtran.models.aerosol import OpacPresetName +from pyradtran.models.aerosol_composite import ( + IntegrationConfig, + RefractiveIndex, + SizeDistribution, +) + +CONFIG_VERSION = 1 + + +class _Strict(BaseModel): + """Base for config models: reject unknown keys (typo protection).""" + + model_config = ConfigDict(extra="forbid") + + +# --- Scene section (builder-kwargs pass-through) --- + + +class SceneSection(_Strict): + """Mirror of the Scene builder kwargs. + + Each sub-dict is passed verbatim to the corresponding ``Scene.set_*()`` + call. Unknown keys are rejected downstream by the frozen, extra=forbid + uvspec option models at load time, so typos fail loudly. + + Keys: + + - ``atmosphere``: ``profile`` (+ ``altitude``, ``mol_modify``, ...) + - ``source``: ``sza`` (+ ``phi0``, ...); ``source: solar|thermal`` optional + - ``wavelength``: ``min_nm``, ``max_nm`` (+ ``unit``, ...) + - ``solver``: ``method``, ``streams`` (+ ``disort_intcor``, ``pseudospherical``, ...) + - ``surface``: ``albedo`` (+ BRDF/BPDF alternatives) + - ``output``: ``quantities``, ``format``, ``zout``, ``heating_rate``, + ``process`` (e.g. ``integrate`` for band-integrated fluxes), ... + """ + + atmosphere: dict[str, Any] + source: dict[str, Any] + wavelength: dict[str, Any] + solver: dict[str, Any] = Field(default_factory=dict) + surface: dict[str, Any] | None = None + output: dict[str, Any] = Field(default_factory=dict) + + +# --- Column placements --- + + +class OdInversionPlacement(_Strict): + """Invert a target column optical depth into an exponential mass profile. + + The placed block's column OD at ``ref_nm`` equals ``tau_ref`` exactly on + the composite's altitude grid (``od_to_mass_profile``). + """ + + kind: Literal["od_inversion"] + tau_ref: float = Field(gt=0) + ref_nm: float = Field(default=550.0, gt=0) + scale_height_km: float = Field(gt=0) + + +class MassPlacement(_Strict): + """Explicit per-layer mass concentration (kg/m^3), descending-altitude order.""" + + kind: Literal["mass"] + kg_m3_per_layer: list[float] = Field(min_length=1) + + +class ExponentialPlacement(_Strict): + """``rho(z) = rho0 * exp(-z / H)``.""" + + kind: Literal["exponential"] + rho0_kg_m3: float = Field(ge=0) + scale_height_km: float = Field(gt=0) + + +class TabulatedPlacement(_Strict): + """Mass concentration tabulated vs altitude (linear interpolation).""" + + kind: Literal["tabulated"] + z_km: list[float] = Field(min_length=1) + kg_m3: list[float] = Field(min_length=1) + + +PlacementSpec = Annotated[ + OdInversionPlacement | MassPlacement | ExponentialPlacement | TabulatedPlacement, + Field(discriminator="kind"), +] + + +# --- Blocks --- + + +class MieBlockSpec(_Strict): + """Mie-computed species: refractive index + size distribution + placement.""" + + kind: Literal["mie"] + name: str = "MieSpecies" + refractive_index: RefractiveIndex + size_distribution: SizeDistribution + particle_density_kg_m3: float = Field(gt=0) + phase_function: Literal["hg", "mie"] = "hg" + integration: IntegrationConfig | None = None + placement: PlacementSpec + + +class BulkBlockSpec(_Strict): + """Bulk optics NetCDF written by aerosol3D ``BulkAerosolOpticsData``. + + Requires the aerosol3D package at load time (only for the reader). + """ + + kind: Literal["bulk"] + name: str = "BulkSpecies" + file: str + placement: PlacementSpec + + +class OpacPresetBlockSpec(_Strict): + """OPAC preset mixture; self-placed via its own preset mass column.""" + + kind: Literal["opac_preset"] + name: str = "OPAC preset" + preset: OpacPresetName + rh_pct: float = Field(default=50.0, ge=0, le=100) + species_names: list[str] | None = None + data_path: str | None = None + n_legendre: int = Field(default=32, ge=1) + + +class ExplicitLayerBlockSpec(_Strict): + """Pre-computed explicit aerosol file set (``.master`` + ``.LAYER``).""" + + kind: Literal["explicit_layer"] + name: str = "explicit_file" + master_path: str + + +BlockSpec = Annotated[ + MieBlockSpec | BulkBlockSpec | OpacPresetBlockSpec | ExplicitLayerBlockSpec, + Field(discriminator="kind"), +] + + +class AerosolSection(_Strict): + """Composite-aerosol assembly: shared grids + LEGO blocks.""" + + wavelength_grid_um: list[float] = Field(min_length=1) + altitude_grid_km: list[float] = Field(min_length=2) + n_legendre: int = Field(default=32, ge=1) + output_dir: str | None = None + blocks: list[BlockSpec] = Field(min_length=1) + + +# --- Analysis section --- + + +class EnergyConservationSpec(_Strict): + """Column energy identity assertion (``assert_energy_conservation``).""" + + tol: float = Field(default=0.05, gt=0) + albedo: float | None = Field(default=None, ge=0, le=1) # None -> scene surface + + +class AnalysisSection(_Strict): + """Post-processing intents executed after the main run. + + - ``energy_conservation``: assert the column energy budget identity. + - ``drf``: direct radiative forcing vs an auto-built no-aerosol baseline. + - ``heating``: second uvspec invocation in heating-rate mode (libRadtran + cannot emit fluxes and heating in one run) and merge into the result. + - ``attribution``: leave-one-out component attribution (N+1 runs). + - ``plots``: names resolved against the viz plot registry (T3). + - ``save_netcdf``: filename for the merged result dataset. + """ + + energy_conservation: EnergyConservationSpec | None = None + drf: bool = False + heating: bool = False + attribution: bool = False + plots: list[str] = Field(default_factory=list) + save_netcdf: str | None = None + + +class PyRadtranConfig(_Strict): + """Top-level config document (``config_version: 1``).""" + + config_version: Literal[1] + name: str = "unnamed" + scene: SceneSection + aerosol: AerosolSection | None = None + analysis: AnalysisSection | None = None diff --git a/src/pyradtran/core/output_parser.py b/src/pyradtran/core/output_parser.py index 7895ab1..da2f62a 100644 --- a/src/pyradtran/core/output_parser.py +++ b/src/pyradtran/core/output_parser.py @@ -105,6 +105,53 @@ def _parse_netcdf(path: Path) -> xr.Dataset: _STANDARD_COLUMNS = ["wavelength", "edir", "edn", "eup", "udir", "udn", "uup"] +def parse_heating_ascii( + path: str | Path, zout_levels_km: Sequence[float] +) -> xr.Dataset: + """Parse uvspec heating-rate ASCII output (wide format). + + libRadtran's ``heating_rate`` mode emits a format distinct from the + standard 7-column flux output: one header row (``0.0`` marker followed by + the zout altitudes), then one row per wavelength (wavelength followed by + the heating rate in K/day at each zout level). With ``output_user`` set, + heating is computed but discarded — so this format only appears when + ``heating_rate`` is requested WITHOUT ``output_user``. + + Args: + path: Path to the uvspec heating-rate output file. + zout_levels_km: Physical zout altitudes in km (one per heating column). + + Returns: + Dataset with dims ``(wavelength, zout)`` and a single ``heating_rate`` + data variable (K/day). + """ + path = Path(path) + rows: list[list[float]] = [] + with open(path) as f: + for line in f: + stripped = line.strip() + if stripped and not stripped.startswith("#"): + rows.append([float(x) for x in stripped.split()]) + if not rows: + raise ValueError(f"Empty heating-rate output file: {path}") + # Drop the header marker row (col0 == 0.0; no real wavelength is 0). + if abs(rows[0][0]) < 1e-9: + rows = rows[1:] + data = np.asarray(rows, dtype=float) # (n_wavelength, 1 + n_zout) + wavelength = data[:, 0] + heating = data[:, 1:] + zout = np.asarray(zout_levels_km, dtype=float) + if heating.shape[1] != zout.size: + raise ValueError( + f"heating-rate output has {heating.shape[1]} per-level columns but " + f"zout_levels_km has {zout.size} entries" + ) + return xr.Dataset( + {"heating_rate": (("wavelength", "zout"), heating)}, + coords={"wavelength": wavelength, "zout": zout}, + ) + + def _parse_ascii( path: Path, n_zout: int = 1, @@ -135,6 +182,11 @@ def _parse_ascii( if column_names is None: if n_cols == 7: column_names = _STANDARD_COLUMNS + elif n_cols == 8: + # libRadtran appends a K/day heating-rate column to the standard + # flux output when heating_rate is enabled. Name it explicitly so + # HEATING_RATE_COLUMN appears in the dataset (was col_7). + column_names = _STANDARD_COLUMNS + [HEATING_RATE_COLUMN] elif n_cols == 2: column_names = ["wavelength", "value"] else: diff --git a/src/pyradtran/core/postprocess.py b/src/pyradtran/core/postprocess.py index a231445..f815471 100644 --- a/src/pyradtran/core/postprocess.py +++ b/src/pyradtran/core/postprocess.py @@ -72,6 +72,83 @@ def compute_budget(ds: xr.Dataset, *, f_incident: str = "edir") -> BudgetResult: ) +@dataclass(frozen=True) +class EnergyBudget: + """Spectral column energy budget (W/m² per wavelength). + + Identity holds by construction: + ``f_incident == f_up_toa + f_abs_surface + f_abs_atm``. ``f_abs_atm`` is the + atmospheric-absorption residual; physically it must be >= 0. + """ + + f_incident: np.ndarray # edir@toa = F0 * mu0 + f_up_toa: np.ndarray # eup@toa (reflected + back-scattered out the top) + f_abs_surface: np.ndarray # (1 - albedo) * (edir + edn)@surface + f_abs_atm: np.ndarray # residual = f_incident - f_up_toa - f_abs_surface + wavelength: np.ndarray + + +def compute_energy_budget( + ds: xr.Dataset, albedo: float, *, f_incident: str = "edir" +) -> EnergyBudget: + """Per-wavelength column energy budget in W/m². + + Convention: ``f_incident = edir @ TOA`` (the incident solar beam, = F0·μ0; + matches :func:`add_budget_vars`). Surface = lowest zout level, TOA = highest. + """ + if "zout" not in ds.dims or ds.sizes["zout"] < 2: + raise ValueError( + "Energy budget 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_up_toa = ds["eup"].isel(zout=toa).values + f_down_sfc = (ds["edir"].isel(zout=surface) + ds["edn"].isel(zout=surface)).values + f_abs_surface = (1.0 - albedo) * f_down_sfc + f_abs_atm = f_inc - f_up_toa - f_abs_surface + return EnergyBudget( + f_incident=f_inc, + f_up_toa=f_up_toa, + f_abs_surface=f_abs_surface, + f_abs_atm=f_abs_atm, + wavelength=np.asarray(ds["wavelength"].values, dtype=float), + ) + + +def assert_energy_conservation( + ds: xr.Dataset, albedo: float, *, tol: float = 0.05, f_incident: str = "edir" +) -> EnergyBudget: + """Compute the budget and assert physical bounds. + + Hard checks (exact physics, within ``tol`` for numerical noise): + 1. ``f_abs_atm >= -tol * f_incident`` (atmosphere cannot create energy) + 2. ``f_up_toa <= f_incident * (1 + tol)`` (cannot reflect more than incident) + + The identity ``f_inc == f_up_toa + f_abs_surface + f_abs_atm`` holds by + construction (f_abs_atm is defined as the residual) and is not re-asserted. + Returns the budget for logging / plotting. Raises AssertionError on violation. + """ + b = compute_energy_budget(ds, albedo, f_incident=f_incident) + eps = np.maximum(tol * b.f_incident, 1e-9) + # Check the more specific physical impossibility first: eup@toa cannot + # exceed f_incident. (When it does, the residual f_abs_atm is also negative, + # so this must precede the absorption check to report the right root cause.) + if np.any(b.f_up_toa > b.f_incident + eps): + worst = float(np.max(b.f_up_toa - b.f_incident)) + raise AssertionError( + f"TOA upwelling exceeds incident ({worst:.3f} W/m² over F_inc): " + "physically impossible — check eup/edir column mapping." + ) + if np.any(b.f_abs_atm < -eps): + worst = float(np.min(b.f_abs_atm)) + raise AssertionError( + f"Atmospheric absorption negative ({worst:.3f} W/m² < -{tol:.0%}·F_inc): " + "column creates energy — check flux parsing / zout ordering." + ) + return b + + def _layer_centers(z_km: np.ndarray) -> np.ndarray: z = np.asarray(z_km, dtype=float) return 0.5 * (z[:-1] + z[1:]) @@ -116,7 +193,11 @@ def evaluate_blocks_on_grid(comp, wl_um, z_km, n_legendre: int = 32) -> dict[str 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)}, + { + "tau": (("wavelength", "layer"), lo.tau), + "ssa": (("wavelength", "layer"), lo.ssa), + "g": (("wavelength", "layer"), lo.g), + }, coords={ "wavelength": wl, "layer": np.arange(centers.size), diff --git a/src/pyradtran/core/runner.py b/src/pyradtran/core/runner.py index ddc3509..a620746 100644 --- a/src/pyradtran/core/runner.py +++ b/src/pyradtran/core/runner.py @@ -10,7 +10,7 @@ import os import shutil import subprocess -from concurrent.futures import ProcessPoolExecutor, as_completed +from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from pathlib import Path from tempfile import TemporaryDirectory @@ -18,7 +18,7 @@ import xarray as xr -from pyradtran.core.output_parser import parse_output, resolve_zout_tokens +from pyradtran.core.output_parser import parse_heating_ascii, parse_output, resolve_zout_tokens from pyradtran.data.resolver import DataResolver if TYPE_CHECKING: @@ -193,13 +193,31 @@ def execute( 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, + # libRadtran heating_rate mode emits a wide format (wavelength + + # per-zout K/day) distinct from the standard long-format flux + # output. When heating is requested WITHOUT output_user quantities, + # use the dedicated heating parser. (With output_user set, heating + # is computed but discarded — the standard parser handles that.) + is_heating_output = ( + output_format != "netcdf" + and scene.output is not None + and scene.output.heating_rate is not None + and not scene.output.quantities ) + if is_heating_output: + if not zout_levels_km: + raise ValueError( + "heating_rate output requires zout levels to parse the wide format" + ) + result = parse_heating_ascii(out_file, zout_levels_km=zout_levels_km) + else: + 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() result.attrs["uvspec_exe"] = uvspec @@ -218,6 +236,19 @@ def execute_many( ) -> list[xr.Dataset]: """Execute multiple uvspec simulations in parallel. + Runs on a *thread* pool (``ThreadPoolExecutor``): uvspec executes as a + subprocess, so the GIL is released while each worker waits on it — the + threads parallelize the actual uvspec invocations, and scenes never + cross a process boundary. This matters because scenes may carry + unpicklable payloads (e.g. ``CompositeAerosol`` pieces backed by + aerosol3D bulk data with local closures), which a process pool cannot + ship to workers. + + Fails fast: futures are consumed in submission order, and the first + failure raises ``RuntimeError(f"scene {i} failed")`` (with the original + exception chained) after cancelling any outstanding futures, instead of + being swallowed into the result list. + Args: scenes: List of configured Scene objects. uvspec_exe: Path to uvspec binary. Auto-detected if None. @@ -228,7 +259,11 @@ def execute_many( config: Optional ``RunnerConfig`` overriding global defaults. Returns: - List of xarray.Dataset results, one per scene. + List of xarray.Dataset results, one per scene, in input order. + + Raises: + RuntimeError: If any scene fails to execute; the original exception + is chained as the cause. """ cfg = config or Runner._config resolved_max_workers = max_workers if max_workers is not None else cfg.max_workers @@ -237,9 +272,8 @@ def execute_many( resolved_keep_temp = keep_temp if keep_temp is not None else cfg.keep_temp resolved_timeout = timeout if timeout is not None else cfg.timeout - results = [] - with ProcessPoolExecutor(max_workers=resolved_max_workers) as executor: - futures = { + with ThreadPoolExecutor(max_workers=resolved_max_workers) as executor: + futures = [ executor.submit( Runner.execute, scene, @@ -247,15 +281,15 @@ def execute_many( data_path=resolved_data_path, keep_temp=resolved_keep_temp, timeout=resolved_timeout, - ): i - for i, scene in enumerate(scenes) - } - for future in as_completed(futures): - idx = futures[future] + ) + for scene in scenes + ] + results = [] + for i, future in enumerate(futures): try: - results.append((idx, future.result())) - except Exception as e: - results.append((idx, e)) - - results.sort(key=lambda x: x[0]) - return [r[1] for r in results] + results.append(future.result()) + except Exception as exc: + for f in futures: + f.cancel() + raise RuntimeError(f"scene {i} failed") from exc + return results diff --git a/src/pyradtran/viz/__init__.py b/src/pyradtran/viz/__init__.py index 2c9b640..62fd953 100644 --- a/src/pyradtran/viz/__init__.py +++ b/src/pyradtran/viz/__init__.py @@ -2,7 +2,10 @@ from pyradtran.viz._style import get_palette, save, set_theme from pyradtran.viz.attribution import AttributionLike, plot_component_attribution +from pyradtran.viz.blocks import plot_block_spectral_optics from pyradtran.viz.composite import plot_block_profiles, plot_composite_optics +from pyradtran.viz.forcing import plot_drf_spectral, plot_spectral_attribution +from pyradtran.viz.inputs import plot_legendre_decay, plot_phase_functions, plot_size_distributions from pyradtran.viz.rt import ( plot_budget, plot_flux_profile, @@ -24,4 +27,10 @@ "plot_block_profiles", "plot_component_attribution", "AttributionLike", + "plot_size_distributions", + "plot_phase_functions", + "plot_legendre_decay", + "plot_block_spectral_optics", + "plot_drf_spectral", + "plot_spectral_attribution", ] diff --git a/src/pyradtran/viz/blocks.py b/src/pyradtran/viz/blocks.py new file mode 100644 index 0000000..6d7e1a1 --- /dev/null +++ b/src/pyradtran/viz/blocks.py @@ -0,0 +1,80 @@ +"""Per-block spectral optics plot: tau / ssa / g vs wavelength, one line per block. + +For tau: column-integrated (sum over layers). For ssa/g: tau-weighted column +average (the physically meaningful column value). +""" + +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 + +_VALID = {"tau", "ssa", "g"} +_LABELS = { + "tau": "Column optical depth τ(λ)", + "ssa": "Single-scattering albedo ⟨ω⟩(λ)", + "g": "Asymmetry parameter ⟨g⟩(λ)", +} + + +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 _column_reduce(ds: xr.Dataset, quantity: str) -> np.ndarray: + """Reduce (wavelength, layer) -> (wavelength,) for plotting.""" + tau = np.asarray(ds["tau"].values, dtype=float) # (wl, layer) + if quantity == "tau": + return tau.sum(axis=1) + q = np.asarray(ds[quantity].values, dtype=float) # (wl, layer) + w = np.abs(tau) + wsum = w.sum(axis=1, keepdims=True) + safe = np.where(wsum > 0, wsum, 1.0) + return (q * w / safe).sum(axis=1) # tau-weighted mean over layers + + +def plot_block_spectral_optics( + per_block_ds_dict: dict[str, xr.Dataset], + *, + quantity: str = "tau", + ax=None, + save_path=None, +): + """One line per block of column tau / ssa / g vs wavelength. + + Args: + per_block_ds_dict: ``{block_name: xr.Dataset}`` from + :func:`pyradtran.core.postprocess.evaluate_blocks_on_grid`. Each + dataset must carry the requested quantity (tau/ssa/g) on a + (wavelength, layer) grid. + quantity: one of ``"tau"`` (column sum), ``"ssa"`` or ``"g"`` + (tau-weighted column mean). + """ + if quantity not in _VALID: + raise ValueError(f"Unknown quantity {quantity!r}; choose from {sorted(_VALID)}") + 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): + if quantity not in bds.data_vars: + continue + wl = np.asarray(bds["wavelength"].values, dtype=float) + y = _column_reduce(bds, quantity) + ax.plot(wl, y, color=color, label=name, linewidth=1.5, marker="o", markersize=3) + ax.set_xlabel("Wavelength (µm)") + ax.set_ylabel(_LABELS[quantity]) + ax.legend(loc="best") + if save_path is not None: + save(fig, Path(save_path)) + return fig, ax diff --git a/src/pyradtran/viz/forcing.py b/src/pyradtran/viz/forcing.py new file mode 100644 index 0000000..74883e2 --- /dev/null +++ b/src/pyradtran/viz/forcing.py @@ -0,0 +1,96 @@ +"""Climate-forcing plots: direct radiative forcing (DRF) + spectral attribution. + +DRF sign convention follows IPCC: negative = cooling (aerosol reflects/absorbs +more than the clean column). DRF_atm (shaded) = atmospheric-absorption forcing. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np + +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 plot_drf_spectral( + wavelength_nm: np.ndarray, + drf_toa: np.ndarray, + drf_surf: np.ndarray, + ax=None, + save_path=None, +): + """Two lines (TOA, surface DRF) + shaded atmospheric-absorption forcing band. + + Args: + wavelength_nm: wavelength axis in nm. + drf_toa, drf_surf: per-wavelength DRF in W/m² (negative = cooling). + """ + fig, ax = _ensure_axes(ax) + ax.axhline(0.0, color="k", linewidth=0.8, linestyle="--") + ax.plot(wavelength_nm, drf_toa, color="#1f77b4", linewidth=1.8, label="TOA") + ax.plot(wavelength_nm, drf_surf, color="#d62728", linewidth=1.8, label="Surface") + ax.fill_between( + wavelength_nm, + drf_toa, + drf_surf, + color="#ff7f0e", + alpha=0.25, + label="Atmosphere", + ) + ax.set_xlabel("Wavelength (nm)") + ax.set_ylabel("Direct radiative forcing (W/m²)") + ax.legend(loc="best") + if save_path is not None: + save(fig, Path(save_path)) + return fig, ax + + +def plot_spectral_attribution( + result, + variable: str = "edir", + level: str = "surface", + ax=None, + save_path=None, +): + """Stacked area of per-block contribution(λ) for one variable/level. + + Args: + result: :class:`pyradtran.workflow.AttributionResult` (contributions are + already spectral xr.Datasets). + variable: flux variable to attribute (edir, edn, eup, ...). + level: "surface" (min zout) or "toa" (max zout). + """ + fig, ax = _ensure_axes(ax) + zout = result.full["zout"].values + level_idx = int(np.argmin(zout)) if level == "surface" else int(np.argmax(zout)) + wl = np.asarray(result.full["wavelength"].values, dtype=float) + + names = list(result.contributions) + colors = get_palette(len(names)) + stack = np.zeros_like(wl, dtype=float) + for color, name in zip(colors, names, strict=True): + c = np.asarray( + result.contributions[name][variable].isel(zout=level_idx).values, dtype=float + ) + ax.fill_between(wl, stack, stack + c, color=color, alpha=0.55, label=name, linewidth=0) + stack = stack + c + ax.plot(wl, stack, color="k", linewidth=1.0, linestyle=":", label="Σ contributions") + ax.set_xlabel("Wavelength (nm)") + ax.set_ylabel(f"Attribution of {variable} @ {level} (W/m²)") + ax.legend(loc="best") + if save_path is not None: + save(fig, Path(save_path)) + return fig, ax diff --git a/src/pyradtran/viz/inputs.py b/src/pyradtran/viz/inputs.py new file mode 100644 index 0000000..44b1667 --- /dev/null +++ b/src/pyradtran/viz/inputs.py @@ -0,0 +1,152 @@ +"""Input-diagnostic plots: size distributions, phase functions, Legendre decay. + +Pure data-in -> fig-out. No Runner, no libRadtran. Phase functions use +monodisperse Mie at the characteristic radius (lognormal r_g / monodisperse +radius_um) — an approximation showing the characteristic angular structure. +For exact polydisperse moments, use plot_legendre_decay. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np + +try: + from numpy import trapezoid as _trapz +except ImportError: # numpy < 2.0 + from numpy import trapz as _trapz # type: ignore[attr-defined] + +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 _characteristic_radius_um(sd) -> float: + params = sd.params + if sd.kind == "lognormal": + return float(params["r_g_um"]) + if sd.kind == "monodisperse": + return float(params["radius_um"]) + if sd.kind == "modified_gamma": + return float(params["r_c_um"]) + raise NotImplementedError( + f"plot_phase_functions does not support kind={sd.kind!r} " + "(no single characteristic radius)." + ) + + +def plot_size_distributions( + blocks: dict, + n_points: int = 200, + r_min_um: float = 1e-3, + r_max_um: float = 50.0, + ax=None, + save_path=None, +): + """dn/dlog10(r) per block on a log-radius axis. + + Args: + blocks: mapping ``{name: SizeDistribution}``. + """ + fig, ax = _ensure_axes(ax) + r = np.logspace(np.log10(r_min_um), np.log10(r_max_um), n_points) + colors = get_palette(len(blocks)) + for color, (name, sd) in zip(colors, blocks.items(), strict=True): + dn = sd.evaluate(r) + dn_dlogr = dn * r * np.log(10) # dn/dr -> dn/dlog10(r) + ax.plot(r, dn_dlogr, color=color, label=name, linewidth=1.5) + ax.set_xscale("log") + ax.set_xlabel("Radius (µm)") + ax.set_ylabel("dn/dlog₁₀(r) (per µm³)") + ax.legend(loc="best") + if save_path is not None: + save(fig, Path(save_path)) + return fig, ax + + +def plot_phase_functions( + blocks: dict, + wavelength_um: float = 0.55, + n_angles: int = 181, + ax=None, + save_path=None, +): + """Normalized phase function P(θ) per block at the characteristic radius. + + Args: + blocks: mapping ``{name: (RefractiveIndex, SizeDistribution)}``. + wavelength_um: wavelength at which to evaluate refractive index + Mie. + """ + from pyradtran.optics.mie import bhmie + + fig, ax = _ensure_axes(ax) + angles = np.linspace(0.0, 180.0, n_angles) + mu = np.cos(np.deg2rad(angles)) + colors = get_palette(len(blocks)) + for color, (name, (ri, sd)) in zip(colors, blocks.items(), strict=True): + m = complex(ri.at(np.array([wavelength_um]))[0]) + r_um = _characteristic_radius_um(sd) + x = 2.0 * np.pi * r_um / wavelength_um + res = bhmie(x, m, n_angles=n_angles) + intensity = np.abs(res["S1"]) ** 2 + np.abs(res["S2"]) ** 2 + # Normalize so integral over the sphere = 1: 2*pi * trapz(P, mu) = 1. + norm = 2.0 * np.pi * _trapz(intensity, mu) + P = intensity / norm if norm > 0 else intensity + ax.semilogy(angles, P, color=color, label=name, linewidth=1.5) + ax.set_xlabel("Scattering angle θ (°)") + ax.set_ylabel("Phase function P(θ) (sr⁻¹)") + ax.set_xlim(0.0, 180.0) + ax.legend(loc="best") + if save_path is not None: + save(fig, Path(save_path)) + return fig, ax + + +def plot_legendre_decay( + comp, + wavelength_um: float = 0.55, + n_legendre: int = 32, + altitude_grid_km=(10.0, 0.0), + ax=None, + save_path=None, +): + """Column-tau-weighted Legendre moments β_l vs l at one wavelength. + + Shows the anisotropy / how many moments carry weight (relates to DISORT + streams needed). Uses ``comp.evaluate`` -> ``LayerOptics.legendre_moments``. + """ + wl = np.asarray([wavelength_um], dtype=float) + z = np.asarray(altitude_grid_km, dtype=float) + lo = comp.evaluate(wl_um=wl, z_km=z, n_legendre=n_legendre) + moments = np.asarray(lo.legendre_moments) # (wl, layer, n_moments) or (wl, n_moments) + tau = np.asarray(lo.tau) # (wl, layer) + if moments.ndim == 3: # (wl, layer, n_moments) + per_layer = moments[0] # (layer, n_moments) + w = tau[0] # (layer,) + wsum = w.sum() + beta = (per_layer * w[:, None]).sum(axis=0) / wsum if wsum > 0 else per_layer.mean(axis=0) + elif moments.ndim == 2: # (wl, n_moments) — already column + beta = moments[0] + else: + beta = np.ravel(moments)[:n_legendre] + ls = np.arange(beta.size) + + fig, ax = _ensure_axes(ax) + ax.stem(ls, beta, basefmt=" ") + ax.set_xlabel("Legendre moment index l") + ax.set_ylabel("βₗ (tau-weighted column)") + ax.set_title(f"Phase-function moments @ {wavelength_um} µm") + if save_path is not None: + save(fig, Path(save_path)) + return fig, ax diff --git a/tests/fixtures/multicomponent_baseline.json b/tests/fixtures/multicomponent_baseline.json new file mode 100644 index 0000000..7e27380 --- /dev/null +++ b/tests/fixtures/multicomponent_baseline.json @@ -0,0 +1,12 @@ +{ + "_meta": { + "scene": "multicomponent_viz", + "note": "regenerate via run_demo.py --dump-baseline" + }, + "edir_surf_550nm": 798.7457, + "edn_surf_550nm": 442.3922, + "eup_surf_550nm": 124.1138, + "eup_toa_550nm": 194.7893, + "column_tau_550nm": 0.5, + "F_abs_atm_550nm": 336.9375899999998 +} \ No newline at end of file diff --git a/tests/test_benchmark_compare.py b/tests/test_benchmark_compare.py new file mode 100644 index 0000000..40e4f40 --- /dev/null +++ b/tests/test_benchmark_compare.py @@ -0,0 +1,323 @@ +"""Offline unit tests for benchmark comparison, reporting and overlay plots. + +Synthetic results dicts are built from the bundled LBL reference itself +(ours = lbl x multiplier, plus fixed RF offsets), so no uvspec is needed. +""" + +import csv +from collections import Counter +from pathlib import Path + +import pytest + +from pyradtran.benchmarks import ( + compare_benchmark, + format_report, + load_reference, + plot_benchmark_overlay, + write_report, +) +from pyradtran.benchmarks.compare import CONFIGS, _status_for +from pyradtran.benchmarks.randles2013 import CASES + +REF = load_reference() + +#: case1 carries 6 reference quantities, the aerosol cases 5 each; case2a +#: additionally gets 4 report-only atmospheric_rf rows. +EXPECTED_ROWS = {"case1": 24, "case2a": 24, "case2b": 20} + + +def build_results(flux_mult=1.0, rf_offset=0.0, *, string_sza_keys=False, with_path=False): + """Synthetic run results whose entries sit at known offsets from the reference. + + Flux-like entries are scaled by ``flux_mult``; RF entries are shifted by + ``rf_offset`` (W/m2). ``uvvis`` bands carry bogus ``rf_*`` entries (99.9) + that the comparison must ignore (reference RF is broadband-only). + """ + results: dict = {} + + def put(case, band, atm, sza, entry): + target = results.setdefault(case, {}).setdefault(band, {}).setdefault(atm, {}) + target[str(sza) if string_sza_keys else sza] = entry + + for config in CONFIGS: + atm, sza = config[:-2], int(config[-2:]) + c1 = REF["case1"] + put( + "case1", + "bb", + atm, + sza, + { + "edir_sfc": c1["direct_bb_sfc_down"][config] * flux_mult, + "edn_sfc": c1["diffuse_bb_sfc_down"][config] * flux_mult, + "eup_sfc": 100.0, + "eup_toa": c1["diffuse_bb_toa_up"][config] * flux_mult, + "total_sfc_down": 0.0, + "nir_sfc_down": c1["nir_sfc_down"][config] * flux_mult, + "absorptance": c1["absorptance"][config] * flux_mult, + }, + ) + put( + "case1", + "uvvis", + atm, + sza, + { + "edir_sfc": 1.0, + "edn_sfc": 1.0, + "eup_sfc": 1.0, + "eup_toa": 1.0, + "total_sfc_down": c1["uvvis_sfc_down"][config] * flux_mult, + }, + ) + for case in ("case2a", "case2b"): + rc = REF[case] + put( + case, + "bb", + atm, + sza, + { + "edir_sfc": 1.0, + "edn_sfc": 1.0, + "eup_sfc": 1.0, + "eup_toa": rc["diffuse_bb_toa_up"][config] * flux_mult, + "total_sfc_down": rc["total_bb_sfc_down"][config] * flux_mult, + "rf_toa": rc["toa_rf"][config] + rf_offset, + "rf_sfc": rc["sfc_rf"][config] + rf_offset, + }, + ) + put( + case, + "uvvis", + atm, + sza, + { + "edir_sfc": 1.0, + "edn_sfc": 1.0, + "eup_sfc": 1.0, + "eup_toa": 1.0, + "total_sfc_down": rc["uvvis_sfc_down"][config] * flux_mult, + "rf_toa": 99.9, + "rf_sfc": 99.9, + }, + ) + if with_path: + results["results_path"] = "/tmp/randles2013_results.json" + return results + + +class TestCompareBenchmark: + @pytest.mark.parametrize( + ("mult", "expected"), + [(1.0, "PASS"), (1.10, "WARN"), (1.20, "FAIL")], + ) + def test_flux_status_at_0_10_20_percent(self, mult, expected): + rows = compare_benchmark(build_results(flux_mult=mult)) + flux_rows = [r for r in rows if not r["quantity"].endswith("_rf")] + assert len(flux_rows) == 48 + statuses = {r["status"] for r in flux_rows} + assert statuses == {expected} + for row in flux_rows: + assert row["rel_diff"] == pytest.approx((mult - 1.0) * 100.0, abs=1e-9) + assert row["abs_diff"] == pytest.approx(row["lbl"] * (mult - 1.0), abs=1e-9) + + def test_exact_results_all_pass_except_report_only(self): + rows = compare_benchmark(build_results()) + assert Counter(r["status"] for r in rows) == Counter({"PASS": 64, "n/a": 4}) + + def test_row_shape_and_counts_per_case(self): + rows = compare_benchmark(build_results()) + required = { + "case", + "quantity", + "atm", + "sza", + "lbl", + "ours", + "abs_diff", + "rel_diff", + "status", + } + assert all(required <= set(row) for row in rows) + assert Counter(r["case"] for r in rows) == Counter(EXPECTED_ROWS) + assert len(rows) == sum(EXPECTED_ROWS.values()) == 68 + for case in CASES: + assert {r["atm"] for r in rows if r["case"] == case} == {"saw", "trop"} + assert {r["sza"] for r in rows if r["case"] == case} == {30, 75} + + def test_rf_rows_broadband_only_and_report_only_atmospheric(self): + rows = compare_benchmark(build_results()) + rf_rows = [r for r in rows if r["quantity"].endswith("_rf")] + assert {r["quantity"] for r in rf_rows} == {"toa_rf", "sfc_rf", "atmospheric_rf"} + # The uvvis rf_* entries (99.9) never leak: RF rows read band "bb" only. + assert all(r["band"] == "bb" for r in rf_rows) + assert all(r["status"] == "PASS" for r in rf_rows if r["quantity"] != "atmospheric_rf") + atm_rows = [r for r in rf_rows if r["quantity"] == "atmospheric_rf"] + assert len(atm_rows) == 4 + assert all(r["status"] == "n/a" for r in atm_rows) + assert all(r["case"] == "case2a" for r in atm_rows) + saw30 = next(r for r in atm_rows if r["atm"] == "saw" and r["sza"] == 30) + assert saw30["lbl"] == pytest.approx(-8.6 - (-9.7)) # toa_rf - sfc_rf + assert saw30["ours"] == pytest.approx(saw30["lbl"]) + + def test_rf_abs_clause_fires_beyond_the_rel_band(self): + results = build_results() + # case2b toa_rf trop75: smallest-magnitude RF reference (-6.5 W/m2), + # so a +1.4 offset is >15% relative but <=1.5 W/m2 absolute. + results["case2b"]["bb"]["trop"][75]["rf_toa"] = -6.5 + 1.4 + row = _find(compare_benchmark(results), "case2b", "toa_rf", "trop", 75) + assert row["abs_diff"] == pytest.approx(1.4) + assert abs(row["rel_diff"]) > 15.0 # rel clause cannot fire + assert row["status"] == "PASS" # via the 1.5 W/m2 clause + # Past 1.5 but within 3 W/m2 (and >25% relative) -> WARN. + results["case2b"]["bb"]["trop"][75]["rf_toa"] = -6.5 + 2.0 + row = _find(compare_benchmark(results), "case2b", "toa_rf", "trop", 75) + assert row["status"] == "WARN" + # Beyond both bands -> FAIL. + results["case2b"]["bb"]["trop"][75]["rf_toa"] = -6.5 + 5.0 + row = _find(compare_benchmark(results), "case2b", "toa_rf", "trop", 75) + assert row["status"] == "FAIL" + + def test_rf_rule_boundary_direct(self): + # 0.5 W/m2 abs at 100% rel -> PASS via the 1.5 W/m2 clause. + assert _status_for("toa_rf", abs_diff=0.5, rel_diff=100.0) == "PASS" + # 2.0 W/m2 at 100% rel -> not PASS (WARN via the 3 W/m2 clause). + assert _status_for("toa_rf", abs_diff=2.0, rel_diff=100.0) == "WARN" + # 3.5 W/m2 at 200% rel -> beyond both bands. + assert _status_for("toa_rf", abs_diff=3.5, rel_diff=200.0) == "FAIL" + # Relative clause alone. + assert _status_for("sfc_rf", abs_diff=10.0, rel_diff=10.0) == "PASS" + + def test_results_path_key_ignored_and_string_sza_keys_accepted(self): + rows = compare_benchmark(build_results(string_sza_keys=True, with_path=True)) + assert len(rows) == 68 + assert Counter(r["case"] for r in rows) == Counter(EXPECTED_ROWS) + assert all(r["status"] in ("PASS", "n/a") for r in rows) + + def test_missing_quantity_rows_skipped(self): + results = build_results() + for atm in ("saw", "trop"): + for sza in (30, 75): + del results["case1"]["bb"][atm][sza]["nir_sfc_down"] + rows = compare_benchmark(results) + assert Counter(r["case"] for r in rows) == Counter( + {"case1": 20, "case2a": 24, "case2b": 20} + ) + assert "nir_sfc_down" not in {r["quantity"] for r in rows} + + def test_empty_results_give_no_rows(self): + assert compare_benchmark({}) == [] + + def test_unknown_benchmark_raises(self): + with pytest.raises(ValueError, match="benchmark"): + compare_benchmark(build_results(), benchmark="ipcc_ar6") + + +class TestReport: + def test_markdown_contains_counts_and_statuses(self): + rows = compare_benchmark(build_results(flux_mult=1.20)) + md = format_report(rows) + assert "**16 pass / 0 warn / 48 fail** (+4 report-only n/a)" in md + assert "❌" in md and "✅" in md + # case1 carries only flux quantities -> all 24 fail at +20%. + assert "## case1 — 0 pass / 0 warn / 24 fail" in md + + rows = compare_benchmark(build_results(flux_mult=1.10)) + md = format_report(rows) + assert "**16 pass / 48 warn / 0 fail**" in md + assert "⚠️" in md + + def test_markdown_one_table_per_case_with_all_configs(self): + md = format_report(compare_benchmark(build_results())) + for case in CASES: + assert f"## {case} —" in md + assert "| quantity | saw30 | saw75 | trop30 | trop75 |" in md + assert "942.4 / 942.4 (+0.0%) ✅" in md # direct_bb_sfc_down saw30 + assert "atmospheric_rf (W/m2)" in md # report-only row present + + def test_empty_rows_report(self): + md = format_report([]) + assert "no rows" in md + + def test_write_report_markdown_and_csv_roundtrip(self, tmp_path): + rows = compare_benchmark(build_results(flux_mult=1.10)) + md_path, csv_path = write_report(rows, tmp_path / "report.md") + assert isinstance(md_path, Path) and isinstance(csv_path, Path) + assert md_path == tmp_path / "report.md" + assert csv_path == tmp_path / "report.csv" + assert md_path.is_file() and csv_path.is_file() + assert "48 warn" in md_path.read_text(encoding="utf-8") + + with open(csv_path, newline="", encoding="utf-8") as f: + reader = csv.DictReader(f) + csv_rows = list(reader) + assert len(csv_rows) == len(rows) == 68 + required = { + "case", + "quantity", + "atm", + "sza", + "lbl", + "ours", + "abs_diff", + "rel_diff", + "status", + } + assert required <= set(reader.fieldnames) + assert len(reader.fieldnames) == 12 + assert {r["status"] for r in csv_rows} <= {"PASS", "WARN", "FAIL", "n/a"} + assert csv_rows[0]["case"] == rows[0]["case"] + assert csv_rows[0]["quantity"] == rows[0]["quantity"] + assert float(csv_rows[0]["lbl"]) == pytest.approx(rows[0]["lbl"]) + assert float(csv_rows[0]["rel_diff"]) == pytest.approx(rows[0]["rel_diff"]) + + +class TestOverlayPlot: + def test_writes_nonempty_png(self, tmp_path): + rows = compare_benchmark(build_results(flux_mult=1.10)) + out = tmp_path / "overlay.png" + returned = plot_benchmark_overlay(rows, out) + assert Path(returned) == out + assert out.is_file() + with open(out, "rb") as f: + assert f.read(8) == b"\x89PNG\r\n\x1a\n" # PNG magic + assert out.stat().st_size > 10_000 + + def test_case1_only_rows_render(self, tmp_path): + results = {"case1": build_results()["case1"]} + out = tmp_path / "case1.png" + plot_benchmark_overlay(compare_benchmark(results), out) + assert out.stat().st_size > 10_000 + + def test_empty_rows_raise(self, tmp_path): + with pytest.raises(ValueError, match="row"): + plot_benchmark_overlay([], tmp_path / "empty.png") + + def test_partial_config_rows_skip_missing_pairs(self, tmp_path): + # One config of one quantity missing (nir_sfc_down @ trop30 only): + # the overlay must skip the absent pair instead of raising KeyError, + # like format_report renders it as "—". + results = build_results() + del results["case1"]["bb"]["trop"][30]["nir_sfc_down"] + rows = compare_benchmark(results) + assert any( + r["quantity"] == "nir_sfc_down" and f"{r['atm']}{r['sza']}" == "trop75" for r in rows + ) + assert not any( + r["quantity"] == "nir_sfc_down" and f"{r['atm']}{r['sza']}" == "trop30" for r in rows + ) + out = tmp_path / "partial.png" + returned = plot_benchmark_overlay(rows, out) + assert Path(returned) == out + assert out.is_file() + assert out.stat().st_size > 10_000 + + +def _find(rows, case, quantity, atm, sza): + return next( + r + for r in rows + if r["case"] == case and r["quantity"] == quantity and r["atm"] == atm and r["sza"] == sza + ) diff --git a/tests/test_benchmark_randles.py b/tests/test_benchmark_randles.py new file mode 100644 index 0000000..c6217d2 --- /dev/null +++ b/tests/test_benchmark_randles.py @@ -0,0 +1,189 @@ +"""Randles et al. (2013) benchmark regression — real uvspec runs (slow + gated). + +Runs the full benchmark matrix — 3 cases x saw/trop x SZA 30/75 x bb/uvvis, +24 uvspec invocations, ~20-60 s — once per session and asserts the paper's +headline results against the bundled LBL reference (loaded from +``load_reference()`` at test time, never hardcoded here): + +1. Case 1 (Rayleigh-only): direct/diffuse broadband SFC-down, diffuse + broadband TOA-up, UV-VIS and NIR SFC-down fluxes within +/-8% (the + reference PASS band, :data:`FLUX_PASS_PCT`) of LBL for all four configs + (saw30/saw75/trop30/trop75) — 20 assertions, supersets the brief's + "8 minimum" (the brief only demands SAW30 + SZA75). +2. Case 2b @ SZA30: TOA radiative forcing POSITIVE for both atmospheres — + the paper's central sign result (SSA-0.8 absorbing aerosol over the + bright Rayleigh atmosphere / 0.2-albedo surface warms at TOA). +3. Case 2a: SFC radiative forcing NEGATIVE for all 4 configs. +4. :func:`compare_benchmark` over the full matrix: zero FAIL rows + (WARN tolerated — see below); row summary printed (visible with ``-s``) + and the Markdown/CSV report + overlay PNG written into the fixture + tmp dir (asserted to exist and render with real data). + +Expected WARN-tolerant quantities +--------------------------------- +- ``case1/absorptance``: reptran-coarse band integration shifts broadband + absorptance ~-4% to -7% relative to the LBL median. On the current + libRadtran 2.0.6 run it stays inside the 8% PASS band, but it sits + closest to the edge of all quantities; version drift is expected to + push it into WARN (<=12%), not FAIL. It is therefore deliberately NOT + part of the +/-8% test above and is only covered by the zero-FAIL gate. +- Small-flux SZA-75 rows (e.g. ``case1/nir_sfc_down`` trop75, ~+3.8%) may + similarly drift between PASS and WARN without indicating a regression. +- ``case2a/atmospheric_rf`` rows are report-only (``"n/a"``: <1 W/m2, RSD + invalid per the paper) and are excluded from the FAIL gate. + +Gating: requests the conftest ``uvspec_exe``/``data_path`` fixtures, so the +whole module skips cleanly without libRadtran. Marked slow — deselect with +``-m "not slow"``. +""" + +import csv +from pathlib import Path + +import pytest + +from pyradtran.benchmarks import ( + compare_benchmark, + load_reference, + plot_benchmark_overlay, + run_randles2013, + write_report, +) +from pyradtran.benchmarks.compare import CONFIGS, FLUX_PASS_PCT, NA_STATUS + +pytestmark = pytest.mark.slow + +#: Case-1 flux quantities asserted against the +/-8% band: reference +#: quantity -> (results band, results key). Covers the brief's "SFC-down BB +#: and TOA-up BB quantities" plus the UV-VIS/NIR split. ``absorptance`` is +#: excluded on purpose (reptran-coarse shift — see module docstring). +_CASE1_FLUX_QUANTITIES: dict[str, tuple[str, str]] = { + "direct_bb_sfc_down": ("bb", "edir_sfc"), + "diffuse_bb_sfc_down": ("bb", "edn_sfc"), + "diffuse_bb_toa_up": ("bb", "eup_toa"), + "uvvis_sfc_down": ("uvvis", "total_sfc_down"), + "nir_sfc_down": ("bb", "nir_sfc_down"), +} + +#: Memoization cache: the conftest gate fixtures are function-scoped and +#: pytest forbids a module-scoped fixture from requesting them (ScopeMismatch), +#: so the expensive full-matrix run is cached here instead — same +#: once-per-session semantics. +_MATRIX_CACHE: dict = {} + + +@pytest.fixture +def randles_results(uvspec_exe, data_path, tmp_path_factory): + """Run the full 24-invocation Randles matrix once per session. + + Skips (via the gate fixtures) when uvspec or the libRadtran data + directory is absent. The results dict (plus its saved + ``randles2013_results.json`` under ``results_path``) is shared by every + test in this module. + """ + if "results" not in _MATRIX_CACHE: + out_dir = tmp_path_factory.mktemp("randles2013") + _MATRIX_CACHE["results"] = run_randles2013( + out_dir, uvspec_exe=uvspec_exe, data_path=data_path + ) + return _MATRIX_CACHE["results"] + + +def _split_config(config: str) -> tuple[str, int]: + """Split a config key ("saw30"/"trop75") into (atmosphere, sza).""" + return config[:-2], int(config[-2:]) + + +@pytest.mark.parametrize("quantity", list(_CASE1_FLUX_QUANTITIES)) +@pytest.mark.parametrize("config", CONFIGS) +def test_case1_fluxes_within_pass_band(randles_results, quantity, config): + """Brief item 1: Case-1 SFC-down/TOA-up fluxes within +/-8% of LBL.""" + reference = load_reference() + band, key = _CASE1_FLUX_QUANTITIES[quantity] + atm, sza = _split_config(config) + lbl = float(reference["case1"][quantity][config]) + ours = randles_results["case1"][band][atm][sza][key] + rel_pct = (ours - lbl) / abs(lbl) * 100.0 + assert abs(rel_pct) <= FLUX_PASS_PCT, ( + f"case1 {quantity} {config}: ours {ours:.2f} vs LBL {lbl:.2f} W/m2 " + f"({rel_pct:+.2f}% exceeds the +/-{FLUX_PASS_PCT:g}% PASS band)" + ) + + +@pytest.mark.parametrize("atm", ["saw", "trop"]) +def test_case2b_toa_rf_positive_at_sza30(randles_results, atm): + """Brief item 2: Case-2b TOA RF is positive at SZA30 (paper's sign result).""" + rf = randles_results["case2b"]["bb"][atm][30]["rf_toa"] + lbl = float(load_reference()["case2b"]["toa_rf"][f"{atm}30"]) + assert rf > 0.0, ( + f"case2b TOA RF @ {atm} SZA30 = {rf:.2f} W/m2, expected positive " + f"(LBL reference {lbl:+.1f} W/m2)" + ) + + +@pytest.mark.parametrize("config", CONFIGS) +def test_case2a_sfc_rf_negative(randles_results, config): + """Brief item 3: Case-2a SFC RF is negative in all 4 configs (less light down).""" + atm, sza = _split_config(config) + rf = randles_results["case2a"]["bb"][atm][sza]["rf_sfc"] + lbl = float(load_reference()["case2a"]["sfc_rf"][config]) + assert rf < 0.0, ( + f"case2a SFC RF @ {config} = {rf:.2f} W/m2, expected negative " + f"(LBL reference {lbl:+.1f} W/m2)" + ) + + +def test_compare_benchmark_no_fail_rows(randles_results): + """Full-matrix comparison: zero FAIL rows (WARN tolerated, n/a report-only). + + Prints the status counts plus every non-PASS row (visible with ``-s``) + so a WARN/FAIL breakdown is directly readable from the test log. + """ + rows = compare_benchmark(randles_results) + assert rows, "compare_benchmark produced no rows from the full matrix" + + fails = [r for r in rows if r["status"] == "FAIL"] + warns = [r for r in rows if r["status"] == "WARN"] + passes = [r for r in rows if r["status"] == "PASS"] + report_only = [r for r in rows if r["status"] == NA_STATUS] + + print( + f"\nRandles benchmark comparison: {len(rows)} rows — " + f"{len(passes)} pass / {len(warns)} warn / {len(fails)} fail " + f"(+{len(report_only)} report-only {NA_STATUS})" + ) + for row in rows: + if row["status"] != "PASS": + print( + f" {row['status']:>4s} {row['case']}/{row['quantity']}/" + f"{row['atm']}{row['sza']}: LBL {row['lbl']:g} vs ours " + f"{row['ours']:g} ({row['rel_diff']:+.2f}%)" + ) + + assert not fails, f"{len(fails)} FAIL row(s) — see printed comparison above" + + +def test_report_and_overlay_render_from_real_data(randles_results): + """Markdown/CSV report + overlay PNG render from the real matrix output. + + The artifacts are written into the fixture tmp dir (next to the saved + results JSON) once per session as a sanity check that the report and + figure tooling works with real uvspec data. + """ + rows = compare_benchmark(randles_results) + out_dir = Path(randles_results["results_path"]).parent + + md_path, csv_path = write_report(rows, out_dir / "randles2013_report.md") + png_path = plot_benchmark_overlay(rows, out_dir / "randles2013_overlay.png") + + assert md_path.is_file() and md_path.stat().st_size > 0 + assert csv_path.is_file() + assert png_path.is_file() + + with open(csv_path, newline="", encoding="utf-8") as f: + n_csv_rows = sum(1 for _ in csv.DictReader(f)) + assert n_csv_rows == len(rows), f"CSV has {n_csv_rows} rows, expected {len(rows)}" + + with open(png_path, "rb") as f: + assert f.read(8) == b"\x89PNG\r\n\x1a\n" # PNG magic + assert png_path.stat().st_size > 10_000 diff --git a/tests/test_benchmark_randles_unit.py b/tests/test_benchmark_randles_unit.py new file mode 100644 index 0000000..d09e69c --- /dev/null +++ b/tests/test_benchmark_randles_unit.py @@ -0,0 +1,271 @@ +"""Offline unit tests for the Randles 2013 benchmark (no uvspec needed).""" + +import json + +import numpy as np +import pytest +import xarray as xr + +from pyradtran.benchmarks import CASES, RandlesAerosol, load_reference, run_randles2013 +from pyradtran.benchmarks.randles2013 import ( + AEROSOL_WAVELENGTH_GRID_UM, + ALTITUDE_GRID_KM, + BANDS, + N_LEGENDRE, + NORMALIZATION_CONSTANTS, + SZAS, + build_scene, + taper_layer_weights, +) + +ALT_GRID = np.array(ALTITUDE_GRID_KM) + + +class TestTaperWeights: + def test_weights_on_benchmark_grid(self): + w = taper_layer_weights(ALT_GRID) + np.testing.assert_allclose(w, [0.0, 0.25, 0.75]) + + def test_weights_sum_to_one(self): + assert taper_layer_weights(ALT_GRID).sum() == pytest.approx(1.0, abs=1e-15) + + def test_taller_grid_same_result(self): + """Aerosol-free layers above 2 km get zero weight.""" + w = taper_layer_weights(np.array([12.0, 3.0, 2.0, 1.0, 0.0])) + np.testing.assert_allclose(w, [0.0, 0.0, 0.25, 0.75]) + + +class TestRandlesAerosol: + def _piece(self, ssa=1.0): + return RandlesAerosol(ssa=ssa) + + def test_tau_column_at_550(self): + lo = self._piece().to_layer_optics(np.array([0.55]), ALT_GRID, n_legendre=N_LEGENDRE) + assert lo.tau.shape == (1, 3) + assert lo.tau.sum() == pytest.approx(0.2, abs=1e-12) + assert lo.tau[0, 0] == 0.0 # layer [3,2] km: above the aerosol top + assert lo.tau[0, 1] == pytest.approx(0.2 * 0.25) # layer [2,1] km + assert lo.tau[0, 2] == pytest.approx(0.2 * 0.75) # layer [1,0] km + + def test_angstrom_exponent_one(self): + piece = self._piece() + tau30 = piece.to_layer_optics(np.array([0.30]), ALT_GRID).tau.sum() + tau55 = piece.to_layer_optics(np.array([0.55]), ALT_GRID).tau.sum() + assert tau30 / tau55 == pytest.approx(0.55 / 0.30, rel=1e-12) + + def test_ssa_constant_per_case(self): + wl = np.array([0.3, 0.55, 2.0]) + for ssa in (1.0, 0.8): + lo = RandlesAerosol(ssa=ssa).to_layer_optics(wl, ALT_GRID) + assert lo.ssa.shape == (3, 3) + np.testing.assert_allclose(lo.ssa, ssa) + + def test_case2b_ssa_is_08(self): + scene = build_scene("case2b", "saw", 30, "bb") + assert scene.aerosol.pieces[0].ssa == 0.8 + + def test_hg_moments(self): + n_leg = 16 + lo = self._piece().to_layer_optics(np.array([0.55]), ALT_GRID, n_legendre=n_leg) + assert lo.legendre_moments.shape == (1, 3, n_leg) + for l in range(8): + np.testing.assert_allclose(lo.legendre_moments[:, :, l], 0.7**l) + + def test_piece_protocol_duck_typing(self): + """RandlesAerosol is accepted as a CompositeAerosol piece.""" + from pyradtran.models.aerosol_composite import CompositeAerosol + + comp = CompositeAerosol( + pieces=[self._piece()], + wavelength_grid_um=AEROSOL_WAVELENGTH_GRID_UM, + altitude_grid_km=ALTITUDE_GRID_KM, + ) + mixed = comp.evaluate(np.array([0.55]), ALT_GRID) + assert mixed.tau.sum() == pytest.approx(0.2, abs=1e-12) + + +class TestBuildScene: + def test_case1_has_no_aerosol(self): + scene = build_scene("case1", "saw", 30, "bb") + assert scene.aerosol is None + text = scene.build_input() + assert "aerosol" not in text + assert "atmosphere_file subarctic_winter" in text + assert "mol_abs_param reptran" in text + assert "sza 30.0" in text + assert "albedo 0.2" in text + assert "rte_solver disort" in text + assert "number_of_streams 16" in text + assert "pseudospherical" in text + assert "disort_intcor moments" in text + assert "output_process integrate" in text + assert "output_user lambda edir edn eup" in text + assert "output_format ascii" in text + + def test_case2a_has_aerosol_piece(self, tmp_path): + scene = build_scene("case2a", "trop", 75, "uvvis", aerosol_output_dir=tmp_path) + assert scene.aerosol is not None + assert isinstance(scene.aerosol.pieces[0], RandlesAerosol) + assert scene.aerosol.pieces[0].ssa == 1.0 + assert scene.aerosol.altitude_grid_km == ALTITUDE_GRID_KM + text = scene.build_input() + assert "wavelength 200.0 700.0" in text + assert "aerosol_file explicit" in text + master = [ln for ln in text.splitlines() if ln.startswith("aerosol_file explicit")][0] + assert tmp_path.name in master + + def test_invalid_args_raise(self): + with pytest.raises(ValueError, match="case"): + build_scene("case9", "saw", 30, "bb") + with pytest.raises(ValueError, match="atmosphere"): + build_scene("case1", "mlw", 30, "bb") + with pytest.raises(ValueError, match="sza"): + build_scene("case1", "saw", 60, "bb") + with pytest.raises(ValueError, match="band"): + build_scene("case1", "saw", 30, "swir") + + +def _synthetic_dataset(edir_toa=1000.0, edir_sfc=800.0, edn_sfc=100.0, eup_sfc=50.0, eup_toa=200.0): + """Mimic parsed ``process=integrate`` output: dims (wavelength: 1, zout: 2).""" + return xr.Dataset( + { + "edir": (("wavelength", "zout"), [[edir_sfc, edir_toa]]), + "edn": (("wavelength", "zout"), [[edn_sfc, 0.0]]), + "eup": (("wavelength", "zout"), [[eup_sfc, eup_toa]]), + }, + coords={"wavelength": [200.0], "zout": [0.0, 120.0]}, + ) + + +class TestRunRandles2013Normalization: + """End-to-end runner logic with uvspec monkeypatched out.""" + + @pytest.fixture() + def run_with_stub(self, tmp_path, monkeypatch): + def _run(cases=None): + def fake_execute_many(scenes, **kwargs): + expected = 24 if cases is None else len(cases) * 8 + assert len(scenes) == expected + return [_synthetic_dataset() for _ in scenes] + + monkeypatch.setattr( + "pyradtran.core.runner.Runner.execute_many", staticmethod(fake_execute_many) + ) + return run_randles2013(tmp_path, cases=cases) + + return _run + + def test_normalized_flux_values(self, run_with_stub): + results = run_with_stub() + c = NORMALIZATION_CONSTANTS["bb"][30] + entry = results["case1"]["bb"]["saw"][30] + assert entry["edir_sfc"] == pytest.approx(800.0 / 1000.0 * c) + assert entry["edn_sfc"] == pytest.approx(100.0 / 1000.0 * c) + assert entry["eup_sfc"] == pytest.approx(50.0 / 1000.0 * c) + assert entry["eup_toa"] == pytest.approx(200.0 / 1000.0 * c) + assert entry["total_sfc_down"] == pytest.approx(900.0 / 1000.0 * c) + + def test_structure_and_derived(self, run_with_stub): + results = run_with_stub() + for case in CASES: + for band in BANDS: + for atm in ("saw", "trop"): + assert set(results[case][band][atm]) == set(SZAS) + # Identical stubbed fluxes -> zero RF and absorptance = 1 - (200+0.8*900)/1000 + entry = results["case2a"]["bb"]["trop"][75] + assert entry["rf_toa"] == pytest.approx(0.0, abs=1e-12) + assert entry["rf_sfc"] == pytest.approx(0.0, abs=1e-12) + c1 = results["case1"]["bb"]["saw"][30] + assert c1["absorptance"] == pytest.approx(1.0 - (200.0 + 0.8 * 900.0) / 1000.0) + # bb and uvvis both stubbed the same -> NIR difference is C_bb - C_uvvis + assert c1["nir_sfc_down"] == pytest.approx( + (NORMALIZATION_CONSTANTS["bb"][30] - NORMALIZATION_CONSTANTS["uvvis"][30]) * 0.9 + ) + + def test_results_json_saved(self, run_with_stub, tmp_path): + results = run_with_stub() + path = tmp_path / "randles2013_results.json" + assert results["results_path"] == str(path) + saved = json.loads(path.read_text()) + assert ( + saved["case1"]["bb"]["saw"]["30"]["edir_sfc"] + == results["case1"]["bb"]["saw"][30]["edir_sfc"] + ) + assert saved["case2b"]["uvvis"]["trop"]["75"]["rf_toa"] == pytest.approx(0.0, abs=1e-12) + + def test_case_subset_without_case1_omits_rf(self, run_with_stub, tmp_path): + results = run_with_stub(cases=["case2a"]) + assert "case1" not in results + assert "rf_toa" not in results["case2a"]["bb"]["saw"][30] + assert "rf_sfc" not in results["case2a"]["uvvis"]["trop"][75] + + def test_unknown_case_raises(self, tmp_path): + with pytest.raises(ValueError, match="case9"): + run_randles2013(tmp_path, cases=["case9"]) + + +class TestReferenceFile: + def test_loads_with_exact_structure(self): + ref = load_reference() + assert set(ref) == {"_meta", "case1", "case2a", "case2b"} + assert set(ref["case1"]) == { + "direct_bb_sfc_down", + "diffuse_bb_sfc_down", + "diffuse_bb_toa_up", + "uvvis_sfc_down", + "nir_sfc_down", + "absorptance", + } + for case in ("case2a", "case2b"): + assert set(ref[case]) == { + "total_bb_sfc_down", + "diffuse_bb_toa_up", + "uvvis_sfc_down", + "toa_rf", + "sfc_rf", + } + + def test_all_values_finite_floats_and_count(self): + ref = load_reference() + n = 0 + for case in ("case1", "case2a", "case2b"): + for quantity, configs in ref[case].items(): + assert set(configs) == {"saw30", "saw75", "trop30", "trop75"}, quantity + for value in configs.values(): + assert isinstance(value, float), f"{case}/{quantity}: {value!r}" + assert np.isfinite(value) + n += 1 + assert n == 64 # 6 quantities x 4 configs (case1) + 2 x (5 x 4) + + def test_anchor_values_match_locked_facts(self): + ref = load_reference() + c1 = ref["case1"] + assert c1["direct_bb_sfc_down"] == { + "saw30": 942.4, + "saw75": 216.2, + "trop30": 844.5, + "trop75": 179.6, + } + assert c1["absorptance"]["trop75"] == 0.307 + assert c1["nir_sfc_down"]["saw30"] == 519.1 + c2a = ref["case2a"] + assert c2a["toa_rf"]["saw30"] == -8.6 + assert c2a["sfc_rf"]["trop75"] == -18.9 + c2b = ref["case2b"] + assert c2b["toa_rf"]["trop30"] == 10.3 + assert c2b["sfc_rf"]["saw30"] == -42.1 + + def test_meta_documentation(self): + ref = load_reference() + meta = ref["_meta"] + assert "Randles" in meta["paper"] + consts = meta["normalization_constants_w_m2"] + assert consts["bb"] == {"sza30": 1189.28, "sza75": 355.43} + assert consts["uvvis"] == {"sza30": 563.38, "sza75": 168.37} + assert meta["thresholds"]["flux"] == {"pass_pct": 8, "warn_pct": 12} + assert meta["thresholds"]["rf"] == {"pass_rel_pct": 15, "pass_abs_w_m2": 1.5} + + def test_reference_file_exists_for_wheel(self): + from pyradtran.benchmarks.randles2013 import _REFERENCE_PATH + + assert _REFERENCE_PATH.is_file() diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..d408774 --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,71 @@ +"""T4: CLI — validate / export-config / run error paths (no uvspec needed). + +A real uvspec `run` is exercised by the canonical YAML example (T6) and the +uvspec-gated suites; here `main` is called directly and via `python -m`. +""" + +import subprocess +import sys +from pathlib import Path + +import yaml + +from pyradtran.config import load_config +from pyradtran.config.cli import main + +#: Minimal valid config: scene only, no aerosol/analysis sections. +MINIMAL = { + "config_version": 1, + "name": "cli_test", + "scene": { + "atmosphere": {"profile": "us", "altitude": 0.0}, + "source": {"sza": 30.0}, + "wavelength": {"min_nm": 400, "max_nm": 700}, + }, +} + + +def _write(path: Path, cfg: dict) -> Path: + path.write_text(yaml.safe_dump(cfg), encoding="utf-8") + return path + + +def test_validate_ok(tmp_path, capsys): + cfg = _write(tmp_path / "config.yaml", MINIMAL) + assert main(["validate", str(cfg)]) == 0 + out = capsys.readouterr().out + assert "ok:" in out + assert "cli_test" in out and "us" in out + + +def test_validate_bad_config_version(tmp_path, capsys): + bad = dict(MINIMAL, config_version=2) + path = _write(tmp_path / "bad.yaml", bad) + assert main(["validate", str(path)]) == 1 + assert capsys.readouterr().err + + +def test_export_config_roundtrip(tmp_path, capsys): + cfg = _write(tmp_path / "config.yaml", MINIMAL) + out_path = tmp_path / "exported.yaml" + assert main(["export-config", str(cfg), "-o", str(out_path)]) == 0 + assert str(out_path) in capsys.readouterr().out + + assert load_config(out_path).scene.build_input() == load_config(cfg).scene.build_input() + + +def test_run_missing_config(tmp_path, capsys): + missing = tmp_path / "does_not_exist.yaml" + assert main(["run", str(missing)]) == 1 + assert str(missing) in capsys.readouterr().err + + +def test_python_dash_m_entry_point(tmp_path): + cfg = _write(tmp_path / "config.yaml", MINIMAL) + proc = subprocess.run( + [sys.executable, "-m", "pyradtran", "validate", str(cfg)], + capture_output=True, + text=True, + ) + assert proc.returncode == 0, proc.stderr + assert "ok:" in proc.stdout diff --git a/tests/test_composite_aerosol_unit.py b/tests/test_composite_aerosol_unit.py index a292471..dbeab3d 100644 --- a/tests/test_composite_aerosol_unit.py +++ b/tests/test_composite_aerosol_unit.py @@ -1,4 +1,11 @@ import numpy as np + +# numpy >=2.0 renamed np.trapz -> np.trapezoid; env has 1.26.x. Mirror mie.py guard. +try: + from numpy import trapezoid as _trapz +except ImportError: # numpy < 2.0 + from numpy import trapz as _trapz # type: ignore[attr-defined] + import pytest from pyradtran.models.aerosol_composite import ( @@ -55,7 +62,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.trapezoid(dn, r) + total = _trapz(dn, r) assert np.isclose(total, 1.0, rtol=0.01) def test_monodisperse_peak_location(self): @@ -81,7 +88,7 @@ def test_modified_gamma_normalization(self): ) r = np.logspace(-2, 2, 10000) dn = sd.evaluate(r) - total = np.trapezoid(dn, r) + total = _trapz(dn, r) assert np.isclose(total, 1.0, rtol=0.01) def test_discrete_normalization(self): @@ -91,7 +98,7 @@ def test_discrete_normalization(self): ) r = np.logspace(-2, 2, 10000) dn = sd.evaluate(r) - total = np.trapezoid(dn, r) + total = _trapz(dn, r) assert np.isclose(total, 1.0, rtol=0.01) def test_modified_gamma_invalid_params_raises(self): diff --git a/tests/test_config_orchestrator.py b/tests/test_config_orchestrator.py new file mode 100644 index 0000000..b578492 --- /dev/null +++ b/tests/test_config_orchestrator.py @@ -0,0 +1,107 @@ +"""T3: orchestrator unit tests (synthetic data — no uvspec needed). + +End-to-end orchestration is exercised by the canonical YAML example (T6) and +the uvspec-gated regression suite. +""" + +import numpy as np +import pytest +import xarray as xr + +from pyradtran.config.orchestrator import ( + RunResult, + _compute_drf, + _heating_section, + _make_plots, + _scene_albedo, +) +from pyradtran.config.schema import SceneSection + + +def _scene_section(output=None, surface=None) -> SceneSection: + return SceneSection( + atmosphere={"profile": "us", "altitude": 0.0}, + source={"sza": 30.0}, + wavelength={"min_nm": 400.0, "max_nm": 700.0}, + surface=surface, + output=output if output is not None else {}, + ) + + +def _synthetic_rt(**extra) -> xr.Dataset: + """2 wl × 2 zout flux dataset with hand-checkable values.""" + wl = [400.0, 550.0] + zout = [0.0, 120.0] # index 0 = surface, 1 = TOA + edir = xr.DataArray([[100.0, 1000.0], [80.0, 800.0]], dims=("wavelength", "zout")) + edn = xr.DataArray([[50.0, 0.0], [40.0, 0.0]], dims=("wavelength", "zout")) + eup = xr.DataArray([[30.0, 200.0], [24.0, 160.0]], dims=("wavelength", "zout")) + return xr.Dataset( + { + "edir": (("wavelength", "zout"), edir.data), + "edn": (("wavelength", "zout"), edn.data), + "eup": (("wavelength", "zout"), eup.data), + }, + coords={"wavelength": wl, "zout": zout}, + **extra, + ) + + +class TestHeatingSection: + def test_drops_quantities(self): + out = {"quantities": ["lambda", "edir"], "format": "ascii", "heating_rate": "local"} + section = _scene_section(output=out) + heat = _heating_section(section) + assert "quantities" not in heat.output + assert heat.output["heating_rate"] == "local" + assert heat.output["format"] == "ascii" + + def test_original_untouched(self): + out = {"quantities": ["lambda"], "heating_rate": "local"} + section = _scene_section(output=out) + _heating_section(section) + assert section.output["quantities"] == ["lambda"] + + +class TestSceneAlbedo: + def test_from_surface(self): + assert _scene_albedo(_scene_section(surface={"albedo": 0.2})) == 0.2 + + def test_missing_raises(self): + with pytest.raises(ValueError, match="albedo"): + _scene_albedo(_scene_section()) + + +class TestComputeDrf: + def test_hand_computed(self): + # With aerosol (wl=550): net_toa = 800-160 = 640, net_sfc = 80+40-24 = 96. + # No aerosol: edir+20/edn+10/eup+6 everywhere, plus edir+100 and eup+50 + # at TOA → net_toa = 920-216 = 704, net_sfc = 100+50-30 = 120. + rt_noaer = _synthetic_rt() + for var, delta in (("edir", 20.0), ("edn", 10.0), ("eup", 6.0)): + rt_noaer[var].data += delta + rt_noaer["edir"].data[:, 1] += 100.0 # extra beam at TOA + rt_noaer["eup"].data[:, 1] += 50.0 + drf = _compute_drf(_synthetic_rt(), rt_noaer) + assert np.allclose(drf.toa, 640.0 - 704.0) # -64 W/m² (both wl) + assert np.allclose(drf.surface, 96.0 - 120.0) # -24 W/m² + assert np.allclose(drf.atmosphere, drf.toa - drf.surface) + assert np.allclose(drf.wavelength_nm, [400.0, 550.0]) + + +class TestMakePlots: + def test_unknown_plot_rejected(self, tmp_path): + with pytest.raises(ValueError, match="unknown plot"): + _make_plots(["not_a_plot"], RunResult(rt=_synthetic_rt()), tmp_path) + + def test_drf_plot_requires_drf(self, tmp_path): + with pytest.raises(ValueError, match="analysis.drf"): + _make_plots(["drf_spectral"], RunResult(rt=_synthetic_rt()), tmp_path) + + def test_heating_plot_requires_heating(self, tmp_path): + with pytest.raises(ValueError, match="analysis.heating"): + _make_plots(["rt_heating_rate"], RunResult(rt=_synthetic_rt()), tmp_path) + + def test_renders_rt_spectral(self, tmp_path): + paths = _make_plots(["rt_spectral"], RunResult(rt=_synthetic_rt()), tmp_path) + assert paths == [tmp_path / "rt_spectral.png"] + assert (tmp_path / "rt_spectral.png").stat().st_size > 0 diff --git a/tests/test_config_roundtrip.py b/tests/test_config_roundtrip.py new file mode 100644 index 0000000..8fdf276 --- /dev/null +++ b/tests/test_config_roundtrip.py @@ -0,0 +1,147 @@ +"""T2 acceptance: YAML config → Scene → uvspec text round-trip. + +A config built from the canonical example's constants must produce the exact +same uvspec input text (modulo the explicit-aerosol file path) as the scene +built via the Python API — and, thanks to the layer writer's content-hash +naming, byte-identical .master/.LAYER files. + +T5 extends this with error-case, discriminated-kind, and export tests. +""" + +import re +import sys +from pathlib import Path + +import pytest + +_EXAMPLES_DIR = Path(__file__).resolve().parents[1] / "examples" / "multicomponent_viz" +sys.path.insert(0, str(_EXAMPLES_DIR)) +try: + import canonical # noqa: E402 +finally: + sys.path.pop(0) + +from pyradtran.config import load_config # noqa: E402 + + +def _config_dict(tmp_path: Path) -> dict: + """Config equivalent to canonical.build_composite() + build_scene().""" + blocks = [] + for b in canonical.BLOCKS: + blocks.append( + { + "kind": "mie", + "name": b["name"], + "refractive_index": { + "wavelength_um": canonical._WL_RI, + "n_real": b["n_real"], + "k_imag": b["k_imag"], + }, + "size_distribution": { + "kind": "lognormal", + "params": {"r_g_um": b["r_g_um"], "sigma_g": b["sigma_g"]}, + }, + "particle_density_kg_m3": b["density"], + "placement": { + "kind": "od_inversion", + "tau_ref": b["tau_550"], + "ref_nm": canonical.REF_NM, + "scale_height_km": b["scale_height_km"], + }, + } + ) + return { + "config_version": 1, + "name": "roundtrip_test", + "scene": dict(canonical.SCENE_KW), + "aerosol": { + "wavelength_grid_um": canonical.WAVELENGTHS_UM, + "altitude_grid_km": canonical.ALTITUDE_GRID_KM, + "n_legendre": canonical.N_LEGENDRE, + "output_dir": str(tmp_path), + "blocks": blocks, + }, + } + + +def _canonical_texts(): + aerosol = canonical.build_composite() + scene = canonical.build_scene(aerosol) + return scene.build_input(), aerosol + + +_EXPLICIT = re.compile(r"^(aerosol_file explicit )\S+$", re.MULTILINE) + + +def _mask_paths(text: str) -> str: + return _EXPLICIT.sub(r"\1", text) + + +def test_yaml_matches_api_uvspec_text(tmp_path): + cfg_text = load_config(_config_dict(tmp_path)).scene.build_input() + api_text, _ = _canonical_texts() + assert _mask_paths(cfg_text) == _mask_paths(api_text) + + +def test_yaml_master_file_bytes_identical(tmp_path): + """Same content → same content-hash filename → byte-identical files.""" + _canonical_texts() # writes the reference explicit files to OUTPUT_DIR + load_config(_config_dict(tmp_path)).scene.build_input() # writes to tmp_path + + cfg_master = next(Path(tmp_path).glob("scene_*.master")) + api_master = canonical.OUTPUT_DIR / cfg_master.name + assert api_master.is_file(), "canonical master missing from OUTPUT_DIR" + + # The .master embeds absolute .LAYER paths — normalize the directory prefix. + def _normalized(path: Path) -> bytes: + return path.read_bytes().replace(str(path.parent).encode(), b"") + + assert _normalized(cfg_master) == _normalized(api_master) + + # Layer files are pure numbers — must be byte-identical as-is. + for cfg_layer in sorted(Path(tmp_path).glob("scene_*_layer_*.LAYER")): + api_layer = canonical.OUTPUT_DIR / cfg_layer.name + assert api_layer.is_file(), f"missing layer file {api_layer.name}" + assert cfg_layer.read_bytes() == api_layer.read_bytes() + + +def test_load_from_yaml_file_roundtrip(tmp_path): + """The dict and the YAML-file entry points build identical scenes.""" + import yaml + + cfg_path = tmp_path / "config.yaml" + cfg_path.write_text(yaml.safe_dump(_config_dict(tmp_path)), encoding="utf-8") + from_dict = load_config(_config_dict(tmp_path / "nested")) + from_file = load_config(cfg_path) + assert _mask_paths(from_dict.scene.build_input()) == _mask_paths(from_file.scene.build_input()) + + +def test_generated_canonical_yaml_matches_api(monkeypatch): + """The committed canonical.yaml (written by make_yaml.py) reproduces the + canonical.py-built scene: identical uvspec input text modulo the + explicit-aerosol file path (relative `output` vs absolute OUTPUT_DIR).""" + monkeypatch.chdir(_EXAMPLES_DIR) + loaded = load_config("canonical.yaml") + api_text = canonical.build_scene(canonical.build_composite()).build_input() + assert _mask_paths(loaded.scene.build_input()) == _mask_paths(api_text) + + +def test_unknown_block_kind_rejected(tmp_path): + cfg = _config_dict(tmp_path) + cfg["aerosol"]["blocks"][0]["kind"] = "miee" # typo + with pytest.raises(Exception, match="kind"): + load_config(cfg) + + +def test_unknown_top_level_key_rejected(tmp_path): + cfg = _config_dict(tmp_path) + cfg["scen"] = {} # typo + with pytest.raises(Exception): + load_config(cfg) + + +def test_bad_config_version_rejected(tmp_path): + cfg = _config_dict(tmp_path) + cfg["config_version"] = 2 + with pytest.raises(Exception, match="config_version"): + load_config(cfg) diff --git a/tests/test_config_schema.py b/tests/test_config_schema.py new file mode 100644 index 0000000..f0cd666 --- /dev/null +++ b/tests/test_config_schema.py @@ -0,0 +1,202 @@ +"""T5: config schema error cases, discriminated kinds, and export_config. + +Characterization tests for :mod:`pyradtran.config`: + +- loader semantic errors in the ``scene`` section (``ValueError``); +- the discriminated ``placement`` union building the matching profile class; +- pydantic constraint violations rejected with ``ValidationError``; +- the ``explicit_layer`` block building a bare direct piece (no file read); +- the ``export_config`` YAML round trip. + +All offline — no uvspec invocation, no libRadtran data. +""" + +import sys +from pathlib import Path + +import pytest +from pydantic import ValidationError + +_EXAMPLES_DIR = Path(__file__).resolve().parents[1] / "examples" / "multicomponent_viz" +sys.path.insert(0, str(_EXAMPLES_DIR)) +try: + import canonical # noqa: E402 +finally: + sys.path.pop(0) + +from pyradtran.config import export_config, load_config # noqa: E402 +from pyradtran.models.blocks import ( # noqa: E402 + DirectLayerOpticsBlock, + ExponentialProfile, + MassProfile, + PlacedBlock, + TabulatedProfile, +) + +_N_LAYER = len(canonical.ALTITUDE_GRID_KM) - 1 + +_OD_PLACEMENT = { + "kind": "od_inversion", + "tau_ref": canonical.BLOCKS[0]["tau_550"], + "ref_nm": canonical.REF_NM, + "scale_height_km": canonical.BLOCKS[0]["scale_height_km"], +} + +_MASS_PLACEMENT = {"kind": "mass", "kg_m3_per_layer": [1e-8] * _N_LAYER} + + +def _mie_block(placement: dict) -> dict: + """A one-block Mie spec derived from the canonical example's first species.""" + b = canonical.BLOCKS[0] + return { + "kind": "mie", + "name": b["name"], + "refractive_index": { + "wavelength_um": canonical._WL_RI, + "n_real": b["n_real"], + "k_imag": b["k_imag"], + }, + "size_distribution": { + "kind": "lognormal", + "params": {"r_g_um": b["r_g_um"], "sigma_g": b["sigma_g"]}, + }, + "particle_density_kg_m3": b["density"], + "placement": placement, + } + + +def _config_dict(tmp_path: Path, block: dict) -> dict: + """A minimal one-block config over the canonical scene and grids.""" + return { + "config_version": 1, + "name": "schema_test", + "scene": dict(canonical.SCENE_KW), + "aerosol": { + "wavelength_grid_um": canonical.WAVELENGTHS_UM, + "altitude_grid_km": canonical.ALTITUDE_GRID_KM, + "n_legendre": canonical.N_LEGENDRE, + "output_dir": str(tmp_path), + "blocks": [block], + }, + } + + +# --- Scene-section semantic errors (loader ValueError) --- + + +def test_solar_source_without_sza_rejected(tmp_path): + cfg = _config_dict(tmp_path, _mie_block(_MASS_PLACEMENT)) + cfg["scene"]["source"] = {"source": "solar"} + with pytest.raises(ValueError, match="sza"): + load_config(cfg) + + +def test_wavelength_without_min_nm_rejected(tmp_path): + cfg = _config_dict(tmp_path, _mie_block(_MASS_PLACEMENT)) + cfg["scene"]["wavelength"] = {"max_nm": 699.0} + with pytest.raises(ValueError, match="min_nm"): + load_config(cfg) + + +def test_unknown_source_kind_rejected(tmp_path): + cfg = _config_dict(tmp_path, _mie_block(_MASS_PLACEMENT)) + cfg["scene"]["source"] = {"source": "flashlight", "sza": 30.0} + with pytest.raises(ValueError, match="flashlight"): + load_config(cfg) + + +# --- Placement discrimination --- + + +_PLACEMENT_CASES = [ + pytest.param(_OD_PLACEMENT, MassProfile, id="od_inversion"), + pytest.param(_MASS_PLACEMENT, MassProfile, id="mass"), + pytest.param( + {"kind": "exponential", "rho0_kg_m3": 1e-8, "scale_height_km": 1.5}, + ExponentialProfile, + id="exponential", + ), + pytest.param( + {"kind": "tabulated", "z_km": [10.0, 5.0, 0.0], "kg_m3": [1e-9, 1e-8, 5e-8]}, + TabulatedProfile, + id="tabulated", + ), +] + + +@pytest.mark.parametrize(("placement", "profile_cls"), _PLACEMENT_CASES) +def test_placement_kind_builds_matching_profile(tmp_path, placement, profile_cls): + loaded = load_config(_config_dict(tmp_path, _mie_block(placement))) + piece = loaded.aerosol.pieces[0] + assert isinstance(piece, PlacedBlock) + assert isinstance(piece.profile, profile_cls) + + kind = placement["kind"] + if kind == "od_inversion": # the inversion yields one mass per composite layer + assert len(piece.profile.kg_m3_per_layer) == _N_LAYER + assert all(m > 0.0 for m in piece.profile.kg_m3_per_layer) + elif kind == "mass": + assert piece.profile.kg_m3_per_layer == tuple(placement["kg_m3_per_layer"]) + elif kind == "exponential": + assert piece.profile.rho0_kg_m3 == placement["rho0_kg_m3"] + assert piece.profile.scale_height_km == placement["scale_height_km"] + else: # tabulated + assert piece.profile.z_km == tuple(placement["z_km"]) + assert piece.profile.kg_m3 == tuple(placement["kg_m3"]) + + +# --- Constraint violations (pydantic ValidationError) --- + + +@pytest.mark.parametrize("tau_ref", [0.0, -0.15], ids=["zero", "negative"]) +def test_od_inversion_nonpositive_tau_rejected(tmp_path, tau_ref): + placement = {"kind": "od_inversion", "tau_ref": tau_ref, "scale_height_km": 1.5} + with pytest.raises(ValidationError, match="tau_ref"): + load_config(_config_dict(tmp_path, _mie_block(placement))) + + +def test_single_level_altitude_grid_rejected(tmp_path): + cfg = _config_dict(tmp_path, _mie_block(_MASS_PLACEMENT)) + cfg["aerosol"]["altitude_grid_km"] = [5.0] + with pytest.raises(ValidationError, match="altitude_grid_km"): + load_config(cfg) + + +def test_invalid_phase_function_rejected(tmp_path): + block = _mie_block(_MASS_PLACEMENT) + block["phase_function"] = "delta" + with pytest.raises(ValidationError, match="phase_function"): + load_config(_config_dict(tmp_path, block)) + + +def test_unknown_opac_preset_rejected(tmp_path): + block = {"kind": "opac_preset", "preset": "not_a_preset"} + with pytest.raises(ValidationError, match="preset"): + load_config(_config_dict(tmp_path, block)) + + +# --- explicit_layer: bare direct piece, no file access --- + + +def test_explicit_layer_block_builds_bare_direct_piece(tmp_path): + master = tmp_path / "no_such_dir" / "fake.master" + block = {"kind": "explicit_layer", "name": "fake_explicit", "master_path": str(master)} + loaded = load_config(_config_dict(tmp_path, block)) + piece = loaded.aerosol.pieces[0] + assert isinstance(piece, DirectLayerOpticsBlock) + assert not isinstance(piece, PlacedBlock) # direct route: no placement wrapper + assert piece.master_path == str(master) + assert not master.exists() # building must not read (or create) the file + + +# --- export_config round trip --- + + +def test_export_config_round_trip(tmp_path): + cfg = _config_dict(tmp_path, _mie_block(_OD_PLACEMENT)) + exported = export_config(cfg, tmp_path / "exported.yaml") + assert exported.is_file() + + via_file = load_config(exported) + via_dict = load_config(cfg) + assert via_file.scene.build_input() == via_dict.scene.build_input() diff --git a/tests/test_convenience.py b/tests/test_convenience.py index 19597fd..1954295 100644 --- a/tests/test_convenience.py +++ b/tests/test_convenience.py @@ -161,8 +161,23 @@ def test_returns_dataset(self, uvspec_exe, data_path): assert isinstance(result, xr.Dataset) +def _need_mps(data_path): + """Skip unless libRadtran's MPS satellite-geometry netCDF is available. + + MPS ships with libRadtran (not bundled here). Gate the satellite test so + CI without the full libRadtran data tree skips cleanly instead of failing. + """ + import os + + for root in (data_path, os.environ.get("PYRADTRAN_DATA_PATH")): + if root and os.path.isfile(os.path.join(root, "MPS")): + return + pytest.skip("libRadtran MPS satellite-geometry file not available") + + class TestRunSatellite: def test_returns_dataset(self, uvspec_exe, data_path): + _need_mps(data_path) from pyradtran.convenience import run_satellite result = run_satellite( diff --git a/tests/test_energy_budget.py b/tests/test_energy_budget.py new file mode 100644 index 0000000..136f64a --- /dev/null +++ b/tests/test_energy_budget.py @@ -0,0 +1,65 @@ +"""B2: column energy budget + conservation assertion. + +Identity (per wavelength): f_inc = f_up_toa + f_abs_surface + f_abs_atm, +where f_inc = edir@toa, f_up_toa = eup@toa, f_abs_surface = (1-albedo)*(edir+edn)@surf, +and f_abs_atm is the residual (atmospheric absorption, must be >= 0). +""" + +import numpy as np +import pytest +import xarray as xr + +from pyradtran.core.postprocess import ( + EnergyBudget, + assert_energy_conservation, + compute_energy_budget, +) + + +def _rt_dataset(edir_toa, eup_toa, edir_surf, edn_surf): + """2-level (surface, toa) synthetic dataset. zout sorted ascending.""" + wl = [0.55] + return xr.Dataset( + { + "edir": (("wavelength", "zout"), [[edir_surf, edir_toa]]), + "edn": (("wavelength", "zout"), [[edn_surf, 0.0]]), + "eup": (("wavelength", "zout"), [[0.0, eup_toa]]), + }, + coords={"wavelength": wl, "zout": [0.0, 120.0]}, + ) + + +def test_energy_budget_identity_holds_by_construction(): + ds = _rt_dataset(edir_toa=800.0, eup_toa=200.0, edir_surf=500.0, edn_surf=100.0) + b = compute_energy_budget(ds, albedo=0.1) + assert isinstance(b, EnergyBudget) + assert np.isclose(b.f_incident[0], 800.0) + assert np.isclose(b.f_up_toa[0], 200.0) + assert np.isclose(b.f_abs_surface[0], 0.9 * (500.0 + 100.0)) + # residual = 800 - 200 - 540 = 60 + assert np.isclose(b.f_abs_atm[0], 60.0) + + +def test_assert_energy_conservation_passes_for_physical_fluxes(): + ds = _rt_dataset(edir_toa=800.0, eup_toa=200.0, edir_surf=500.0, edn_surf=100.0) + b = assert_energy_conservation(ds, albedo=0.1) # should not raise + assert b.f_abs_atm[0] >= 0 + + +def test_assert_energy_conservation_raises_when_atmosphere_creates_energy(): + # Impossible: eup@toa + (1-a)(edir+edn)@surf > f_inc => negative atmospheric absorption. + ds = _rt_dataset(edir_toa=800.0, eup_toa=500.0, edir_surf=500.0, edn_surf=100.0) + with pytest.raises(AssertionError, match="[Aa]tmospheric absorption"): + assert_energy_conservation(ds, albedo=0.1) + + +def test_assert_energy_conservation_raises_when_toa_upwelling_exceeds_incident(): + ds = _rt_dataset(edir_toa=800.0, eup_toa=900.0, edir_surf=0.0, edn_surf=0.0) + with pytest.raises(AssertionError, match="[Uu]pwelling"): + assert_energy_conservation(ds, albedo=0.1) + + +def test_compute_energy_budget_requires_2d_dataset(): + ds = xr.Dataset({"edir": ("wavelength", [800.0])}, coords={"wavelength": [0.55]}) + with pytest.raises(ValueError, match="zout"): + compute_energy_budget(ds, albedo=0.1) diff --git a/tests/test_multiblock_tau_stacking.py b/tests/test_multiblock_tau_stacking.py new file mode 100644 index 0000000..a23cf43 --- /dev/null +++ b/tests/test_multiblock_tau_stacking.py @@ -0,0 +1,77 @@ +"""B1: multi-block external-mixture column tau must equal sum of target taus. + +Validates od_to_mass_profile (per-block tau->mass inversion) + composite +stacking: 3 Mie blocks with target tau@550 = {0.15, 0.15, 0.20} must yield a +composite column tau@550 ~= 0.50. No libRadtran — pure analytic Mie mixing. +""" + +import numpy as np + +from pyradtran.core.postprocess import 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 + +WL = [0.40, 0.55, 0.70] +ALT = [10.0, 8.0, 6.0, 4.0, 2.0, 1.0, 0.0] +_TARGETS = [ + { + "name": "bc", "n": 1.95, "k": 0.79, "r_g": 0.10, + "sigma": 2.0, "rho": 1800.0, "H": 1.5, "tau": 0.15, + }, + { + "name": "sulfate", "n": 1.53, "k": 0.0, "r_g": 0.15, + "sigma": 1.7, "rho": 1770.0, "H": 2.0, "tau": 0.15, + }, + { + "name": "dust", "n": 1.53, "k": 0.008, "r_g": 0.50, + "sigma": 2.2, "rho": 2600.0, "H": 3.0, "tau": 0.20, + }, +] + + +def _build_composite(): + pieces = [] + for b in _TARGETS: + ri = RefractiveIndex(wavelength_um=WL, n_real=[b["n"]] * 3, k_imag=[b["k"]] * 3) + sd = SizeDistribution(kind="lognormal", params={"r_g_um": b["r_g"], "sigma_g": b["sigma"]}) + species = MieSpecies( + refractive_index=ri, + size_distribution=sd, + particle_density_kg_m3=b["rho"], + integration_config=IntegrationConfig(), + name=b["name"], + ) + profile = od_to_mass_profile( + species, tau_ref=b["tau"], ref_nm=550.0, altitude_km=ALT, scale_height_km=b["H"] + ) + pieces.append(PlacedBlock(block=species, profile=profile)) + return CompositeAerosol( + pieces=pieces, wavelength_grid_um=WL, altitude_grid_km=ALT, n_legendre=4, output_dir="." + ) + + +def test_column_tau_550_equals_sum_of_targets(): + comp = _build_composite() + grid = evaluate_composite_on_grid(comp, WL, ALT, n_legendre=4) + col_tau_550 = float(np.sum(grid["tau"].sel(wavelength=0.55).values)) + expected = sum(b["tau"] for b in _TARGETS) + assert np.isclose(col_tau_550, expected, rtol=0.01), ( + f"column tau@550={col_tau_550:.4f}, expected sum={expected:.4f} (1% tol)" + ) + + +def test_per_block_column_tau_550_matches_own_target(): + """Each block alone also inverts to its own target (single-block sanity).""" + comp = _build_composite() + blocks = evaluate_blocks_on_grid(comp, WL, ALT, n_legendre=4) + for b in _TARGETS: + col = float(np.sum(blocks[b["name"]]["tau"].sel(wavelength=0.55).values)) + assert np.isclose(col, b["tau"], rtol=0.01), ( + f"{b['name']}: column tau@550={col:.4f}, target={b['tau']:.4f}" + ) diff --git a/tests/test_multicomponent_regression.py b/tests/test_multicomponent_regression.py new file mode 100644 index 0000000..a8e77b8 --- /dev/null +++ b/tests/test_multicomponent_regression.py @@ -0,0 +1,73 @@ +"""B3: canonical multicomponent scene regression — key scalars vs committed baseline. + +Guards the end-to-end pipeline (LEGO composite -> explicit .LAYER -> DISORT) +against silent drift. Gated: skipped without uvspec + libRadtran data; marked slow. +""" + +import json +import sys +from pathlib import Path + +import numpy as np +import pytest + +pytestmark = pytest.mark.slow + +_FIXTURE = Path(__file__).parent / "fixtures" / "multicomponent_baseline.json" +_TOL = 0.02 # 2% — DISORT is deterministic; absorbs minor libRadtran-version drift +_EXAMPLES_DIR = Path(__file__).resolve().parents[1] / "examples" / "multicomponent_viz" + + +@pytest.fixture +def canonical_rt(uvspec_exe, data_path): + """Run the canonical scene once per session and return the dataset.""" + sys.path.insert(0, str(_EXAMPLES_DIR)) + try: + from canonical import build_composite, build_scene # noqa: E402 + finally: + sys.path.pop(0) + + from pyradtran.core.runner import Runner + + scene = build_scene(build_composite()) + return Runner.execute(scene, uvspec_exe=uvspec_exe, data_path=data_path) + + +def _baseline(): + if not _FIXTURE.is_file(): + pytest.skip("baseline fixture missing — run scripts/regen_baseline.py") + return json.loads(_FIXTURE.read_text()) + + +def test_scalar_baseline_matches(canonical_rt): + rt = canonical_rt + base = _baseline() + i550 = int(np.argmin(np.abs(rt["wavelength"].values - 550.0))) + surf_idx = int(np.argmin(rt["zout"].values)) + toa_idx = int(np.argmax(rt["zout"].values)) + checks = { + "edir_surf_550nm": float(rt["edir"].isel(wavelength=i550, zout=surf_idx).values), + "edn_surf_550nm": float(rt["edn"].isel(wavelength=i550, zout=surf_idx).values), + "eup_surf_550nm": float(rt["eup"].isel(wavelength=i550, zout=surf_idx).values), + "eup_toa_550nm": float(rt["eup"].isel(wavelength=i550, zout=toa_idx).values), + } + for key, got in checks.items(): + expected = float(base[key]) + assert np.isclose(got, expected, rtol=_TOL), ( + f"{key}: got {got:.4f}, baseline {expected:.4f} (>{_TOL:.0%} drift)" + ) + + +def test_atmospheric_absorption_matches(canonical_rt): + """F_abs_atm@550 from a fresh run must match the committed baseline.""" + from pyradtran.core.postprocess import compute_energy_budget + + rt = canonical_rt + base = _baseline() + budget = compute_energy_budget(rt, albedo=0.1) + i550 = int(np.argmin(np.abs(budget.wavelength - 550.0))) + got = float(budget.f_abs_atm[i550]) + expected = float(base["F_abs_atm_550nm"]) + assert np.isclose(got, expected, rtol=_TOL), ( + f"F_abs_atm@550: got {got:.4f}, baseline {expected:.4f} (>{_TOL:.0%} drift)" + ) diff --git a/tests/test_output_parser_heating.py b/tests/test_output_parser_heating.py new file mode 100644 index 0000000..64a7c10 --- /dev/null +++ b/tests/test_output_parser_heating.py @@ -0,0 +1,35 @@ +"""T1: 8-column uvspec ASCII output (fluxes + heating_rate) must name the 8th +column 'heating_rate', not 'col_7'.""" + +from pathlib import Path + +from pyradtran.core.output_parser import HEATING_RATE_COLUMN, parse_output + + +def _write_ascii(tmp_path: Path, n_zout_rows: int, n_cols: int) -> Path: + p = tmp_path / "out.dat" + lines = [] + for w in range(500, 510): # 10 wavelengths + for _ in range(n_zout_rows): + vals = [float(w)] + [0.1 * (c + 1) for c in range(n_cols - 1)] + lines.append(" ".join(f"{v:.4f}" for v in vals)) + p.write_text("\n".join(lines) + "\n") + return p + + +def test_8_column_ascii_names_heating_rate(tmp_path): + p = _write_ascii(tmp_path, n_zout_rows=1, n_cols=8) + ds = parse_output(p, format="ascii", n_zout=1) + assert HEATING_RATE_COLUMN in ds.data_vars, ( + f"expected '{HEATING_RATE_COLUMN}' in {list(ds.data_vars)}" + ) + # Standard 7 flux columns still present. + for col in ("wavelength", "edir", "edn", "eup", "udir", "udn", "uup"): + assert col in ds.data_vars or col in ds.coords + + +def test_7_column_ascii_unchanged(tmp_path): + p = _write_ascii(tmp_path, n_zout_rows=1, n_cols=7) + ds = parse_output(p, format="ascii", n_zout=1) + assert HEATING_RATE_COLUMN not in ds.data_vars + assert "col_0" not in ds.data_vars # standard naming kicked in diff --git a/tests/test_parse_heating_ascii.py b/tests/test_parse_heating_ascii.py new file mode 100644 index 0000000..1cd4fe6 --- /dev/null +++ b/tests/test_parse_heating_ascii.py @@ -0,0 +1,57 @@ +"""T1 follow-up: parse_heating_ascii handles uvspec's wide heating-rate format. + +libRadtran heating_rate mode emits: 1 header row (0.0 marker + zout levels), +then per-wavelength rows (wavelength + K/day at each zout). Distinct from the +standard long-format flux output. +""" + +import numpy as np + +from pyradtran.core.output_parser import parse_heating_ascii + + +def test_parse_heating_ascii_wide_format(tmp_path): + p = tmp_path / "heat.out" + p.write_text( + "0.000e+00 0.000000 2.000000 5.000000\n" + "400.120 5.99e-06 6.63e-06 7.09e-06\n" + "400.360 6.99e-06 7.64e-06 8.10e-06\n" + ) + ds = parse_heating_ascii(p, zout_levels_km=[0.0, 2.0, 5.0]) + assert "heating_rate" in ds.data_vars + assert ds["heating_rate"].dims == ("wavelength", "zout") + assert ds.sizes["wavelength"] == 2 + assert ds.sizes["zout"] == 3 + assert list(ds["zout"].values) == [0.0, 2.0, 5.0] + assert np.isclose(float(ds["heating_rate"].isel(wavelength=0, zout=1).values), 6.63e-06) + assert np.isclose(float(ds["heating_rate"].isel(wavelength=1, zout=2).values), 8.10e-06) + + +def test_parse_heating_ascii_raises_on_column_count_mismatch(tmp_path): + import pytest + + p = tmp_path / "heat.out" + p.write_text( + "0.000e+00 0.000000 2.000000\n" + "400.120 5.99e-06 6.63e-06\n" + ) + with pytest.raises(ValueError, match="per-level columns"): + parse_heating_ascii(p, zout_levels_km=[0.0, 2.0, 5.0]) # 3 zout, file has 2 + + +def test_parse_heating_ascii_empty_file_raises(tmp_path): + import pytest + + p = tmp_path / "heat.out" + p.write_text("") + with pytest.raises(ValueError, match="Empty"): + parse_heating_ascii(p, zout_levels_km=[0.0, 2.0]) + + +def test_parse_heating_ascii_handles_no_header_marker(tmp_path): + """If libRadtran ever omits the marker row, still parse as wide data.""" + p = tmp_path / "heat.out" + p.write_text("400.120 5.99e-06 6.63e-06\n400.360 6.99e-06 7.64e-06\n") + ds = parse_heating_ascii(p, zout_levels_km=[0.0, 2.0]) + assert ds.sizes["wavelength"] == 2 + assert float(ds["heating_rate"].isel(wavelength=0, zout=0).values) == 5.99e-06 diff --git a/tests/test_runner_parallel.py b/tests/test_runner_parallel.py new file mode 100644 index 0000000..35a6938 --- /dev/null +++ b/tests/test_runner_parallel.py @@ -0,0 +1,182 @@ +"""Tests for Runner.execute_many parallel behavior: fail-fast + no pickling. + +``execute_many`` runs scenes on a *thread* pool: uvspec executes as a +subprocess (the GIL is released while waiting on it), so threads parallelize +the actual work while scenes stay in-process and never need to be pickled. +Failures must surface immediately as ``RuntimeError`` naming the failing +scene index instead of being swallowed into the result list. +""" + +import pickle +from types import SimpleNamespace + +import numpy as np +import pytest +import xarray as xr + +from pyradtran import Scene +from pyradtran.core.runner import Runner +from pyradtran.models.aerosol_composite import ( + BulkSpecies, + CompositeAerosol, + MieSpecies, + RefractiveIndex, + SizeDistribution, +) +from pyradtran.models.blocks import MassProfile, PlacedBlock + + +def _simple_scene() -> Scene: + return ( + Scene().set_atmosphere(profile="us").set_source_solar(sza=30.0).set_wavelength(400.0, 700.0) + ) + + +def _composite_scene() -> Scene: + """Scene carrying a CompositeAerosol with one Mie block (the docs case).""" + ri = RefractiveIndex(wavelength_um=[0.40, 0.55], n_real=[1.5, 1.5], k_imag=[0.0, 0.0]) + sp = MieSpecies( + refractive_index=ri, + size_distribution=SizeDistribution( + kind="lognormal", params={"r_g_um": 0.1, "sigma_g": 2.0} + ), + particle_density_kg_m3=1800.0, + name="t", + ) + comp = CompositeAerosol( + pieces=[PlacedBlock(block=sp, profile=MassProfile(kg_m3_per_layer=(1e-8, 1e-9)))], + wavelength_grid_um=[0.4, 0.55], + altitude_grid_km=[1.0, 0.0], + n_legendre=32, + output_dir=None, + ) + return _simple_scene().set_aerosol(comp) + + +def _unpicklable_composite_scene() -> Scene: + """Scene carrying a composite whose payload cannot pickle. + + Reproduces the multicomponent-example failure: an aerosol3D-backed + ``BulkSpecies`` stores a local closure (``SizeDistribution.lognormal``'s + ``_pdf``) on the bulk object, so pickling the scene raises. + """ + + def _pdf(r): # local closure -> unpicklable, like aerosol3D SizeDistribution + return r + + l_vals = np.arange(8) + bulk = SimpleNamespace( + wavelength_nm=np.linspace(400.0, 600.0, 3), + C_ext=np.full(3, 10.0), + C_sca=np.full(3, 9.0), + SSA=np.full(3, 0.9), + g=np.full(3, 0.5), + beta=np.tile((2 * l_vals + 1) * 0.5**l_vals, (3, 1)), + legendre_moments_beta=np.tile(0.5**l_vals, (3, 1)), + size_distribution=SimpleNamespace(moment=lambda p: 1.0, _pdf=_pdf), + effective_density_kg_m3=1800.0, + ) + comp = CompositeAerosol( + pieces=[ + PlacedBlock( + block=BulkSpecies(bulk=bulk, name="b"), + profile=MassProfile(kg_m3_per_layer=(1e-8, 1e-9)), + ) + ], + wavelength_grid_um=[0.4, 0.55], + altitude_grid_km=[1.0, 0.0], + n_legendre=32, + output_dir=None, + ) + return _simple_scene().set_aerosol(comp) + + +class TestExecuteManyFailFast: + def test_first_scene_failure_raises_naming_index(self, monkeypatch): + """A failing scene must raise RuntimeError naming its index, with the + original exception chained — not be swallowed into the results.""" + + def fake_execute(scene, **kwargs): + raise RuntimeError("boom") + + monkeypatch.setattr(Runner, "execute", staticmethod(fake_execute)) + + with pytest.raises(RuntimeError, match="scene 0") as excinfo: + Runner.execute_many([_simple_scene(), _simple_scene()], max_workers=2) + assert isinstance(excinfo.value.__cause__, RuntimeError) + assert "boom" in str(excinfo.value.__cause__) + + def test_later_scene_failure_raises_naming_index(self, monkeypatch): + """Only the second scene fails -> the error names scene 1.""" + + def fake_execute(scene, **kwargs): + if getattr(scene, "_tag", None) == 0: + return xr.Dataset({"x": ("zout", [0])}) + raise ValueError("late boom") + + monkeypatch.setattr(Runner, "execute", staticmethod(fake_execute)) + scenes = [_simple_scene(), _simple_scene()] + for i, s in enumerate(scenes): + s._tag = i + + with pytest.raises(RuntimeError, match="scene 1") as excinfo: + Runner.execute_many(scenes, max_workers=2) + assert isinstance(excinfo.value.__cause__, ValueError) + + def test_both_succeed_returns_datasets_in_order(self, monkeypatch): + """Two succeeding scenes -> both Datasets returned in submission order.""" + + def fake_execute(scene, **kwargs): + return xr.Dataset({"x": ("zout", [scene._tag])}) + + monkeypatch.setattr(Runner, "execute", staticmethod(fake_execute)) + scenes = [_simple_scene(), _simple_scene()] + for i, s in enumerate(scenes): + s._tag = i + + results = Runner.execute_many(scenes, max_workers=2) + + assert len(results) == 2 + assert all(isinstance(r, xr.Dataset) for r in results) + assert results[0]["x"].item() == 0 + assert results[1]["x"].item() == 1 + + +class TestExecuteManyNoPickling: + def test_composite_scene_accepted(self, monkeypatch): + """A Scene carrying a CompositeAerosol runs without pickling errors.""" + + def fake_execute(scene, **kwargs): + return xr.Dataset({"x": ("zout", [42])}) + + monkeypatch.setattr(Runner, "execute", staticmethod(fake_execute)) + scene = _composite_scene() + + results = Runner.execute_many([scene, scene], max_workers=2) + + assert len(results) == 2 + for r in results: + assert isinstance(r, xr.Dataset) + assert r["x"].item() == 42 + + def test_unpicklable_composite_scene_accepted(self, monkeypatch): + """Regression (multicomponent example): a scene whose composite payload + cannot pickle (aerosol3D bulk closures) must still run — execute_many + must not put scenes on a process boundary.""" + + def fake_execute(scene, **kwargs): + return xr.Dataset({"x": ("zout", [7])}) + + monkeypatch.setattr(Runner, "execute", staticmethod(fake_execute)) + scene = _unpicklable_composite_scene() + + # Guard: the fixture genuinely fails to pickle (the original bug). + with pytest.raises(Exception): + pickle.dumps(scene) + + results = Runner.execute_many([scene, scene], max_workers=2) + + assert len(results) == 2 + for r in results: + assert isinstance(r, xr.Dataset) + assert r["x"].item() == 7 diff --git a/tests/test_viz_blocks.py b/tests/test_viz_blocks.py new file mode 100644 index 0000000..efcb564 --- /dev/null +++ b/tests/test_viz_blocks.py @@ -0,0 +1,78 @@ +"""Unit tests for per-block spectral optics plot + evaluate_blocks_on_grid ssa/g.""" + +import matplotlib + +matplotlib.use("Agg") + +import numpy as np + +from pyradtran.core.postprocess import evaluate_blocks_on_grid +from pyradtran.models.aerosol_composite import ( + CompositeAerosol, + IntegrationConfig, + MieSpecies, + RefractiveIndex, + SizeDistribution, +) +from pyradtran.models.blocks import MassProfile, PlacedBlock + +WL = [0.40, 0.55, 0.70] +ALT = [2.0, 0.0] + + +def _composite(): + ri = RefractiveIndex( + wavelength_um=WL, n_real=[1.95, 1.95, 1.95], k_imag=[0.79, 0.79, 0.79] + ) + sd = SizeDistribution(kind="lognormal", params={"r_g_um": 0.1, "sigma_g": 2.0}) + sp = MieSpecies( + refractive_index=ri, + size_distribution=sd, + particle_density_kg_m3=1800.0, + integration_config=IntegrationConfig(), + name="bc", + ) + return CompositeAerosol( + pieces=[PlacedBlock(block=sp, profile=MassProfile(kg_m3_per_layer=(1e-7,)))], + wavelength_grid_um=WL, + altitude_grid_km=ALT, + n_legendre=4, + output_dir=".", + ) + + +def test_evaluate_blocks_on_grid_now_has_ssa_and_g(): + blocks = evaluate_blocks_on_grid(_composite(), WL, ALT, n_legendre=4) + ds = blocks["bc"] + assert "ssa" in ds.data_vars + assert "g" in ds.data_vars + assert ds["ssa"].shape == (3, 1) # (wavelength, layer) + + +def test_plot_block_spectral_optics_tau_one_line_per_block(): + from pyradtran.viz.blocks import plot_block_spectral_optics + + blocks = evaluate_blocks_on_grid(_composite(), WL, ALT, n_legendre=4) + fig, ax = plot_block_spectral_optics(blocks, quantity="tau") + assert len(ax.lines) == 1 + assert "Wavelength" in ax.get_xlabel() + + +def test_plot_block_spectral_optics_ssa_in_unit_interval(): + from pyradtran.viz.blocks import plot_block_spectral_optics + + blocks = evaluate_blocks_on_grid(_composite(), WL, ALT, n_legendre=4) + fig, ax = plot_block_spectral_optics(blocks, quantity="ssa") + line = ax.lines[0] + y = line.get_ydata() + assert np.all(y >= 0.0) and np.all(y <= 1.0) + + +def test_plot_block_spectral_optics_bad_quantity_raises(): + import pytest + + from pyradtran.viz.blocks import plot_block_spectral_optics + + blocks = evaluate_blocks_on_grid(_composite(), WL, ALT, n_legendre=4) + with pytest.raises(ValueError): + plot_block_spectral_optics(blocks, quantity="nope") diff --git a/tests/test_viz_forcing.py b/tests/test_viz_forcing.py new file mode 100644 index 0000000..4a78a8b --- /dev/null +++ b/tests/test_viz_forcing.py @@ -0,0 +1,60 @@ +"""Unit tests for forcing plots: DRF spectral + spectral attribution.""" + +import matplotlib + +matplotlib.use("Agg") + +import numpy as np +import xarray as xr + +from pyradtran.workflow import AttributionResult + + +def test_plot_drf_spectral_draws_two_lines_and_fill(): + from pyradtran.viz.forcing import plot_drf_spectral + + wl = np.array([400.0, 550.0, 700.0]) + drf_toa = np.array([-1.0, -1.5, -0.8]) + drf_surf = np.array([-2.0, -3.0, -1.5]) + fig, ax = plot_drf_spectral(wl, drf_toa, drf_surf) + assert len(ax.lines) >= 2 + assert len(ax.collections) >= 1 # the fill_between band + assert "Wavelength" in ax.get_xlabel() + + +def test_plot_drf_spectral_zero_line_drawn(): + from pyradtran.viz.forcing import plot_drf_spectral + + wl = np.array([400.0, 700.0]) + fig, ax = plot_drf_spectral(wl, np.array([-1.0, -0.5]), np.array([-2.0, -1.0])) + ydata = [list(l.get_ydata()) for l in ax.lines] + assert [0.0, 0.0] in ydata # zero reference line + + +def _attribution_result_spectral(): + wl = np.array([400.0, 550.0, 700.0]) + zout = [0.0, 120.0] + full = xr.Dataset( + {"edir": (("wavelength", "zout"), [[500.0, 800.0], [450.0, 800.0], [480.0, 800.0]])}, + coords={"wavelength": wl, "zout": zout}, + ) + contributions = { + "bc": full - xr.Dataset( + {"edir": (("wavelength", "zout"), [[510.0, 800.0]] * 3)}, + coords={"wavelength": wl, "zout": zout}, + ), + "dust": full - xr.Dataset( + {"edir": (("wavelength", "zout"), [[520.0, 800.0]] * 3)}, + coords={"wavelength": wl, "zout": zout}, + ), + } + return AttributionResult(full=full, contributions=contributions) + + +def test_plot_spectral_attribution_draws_one_stack_per_block(): + from pyradtran.viz.forcing import plot_spectral_attribution + + result = _attribution_result_spectral() + fig, ax = plot_spectral_attribution(result, variable="edir", level="surface") + legend_texts = [t.get_text() for t in ax.get_legend().get_texts()] + assert "bc" in legend_texts and "dust" in legend_texts diff --git a/tests/test_viz_inputs.py b/tests/test_viz_inputs.py new file mode 100644 index 0000000..681c0bf --- /dev/null +++ b/tests/test_viz_inputs.py @@ -0,0 +1,77 @@ +"""Unit tests for input-diagnostic viz (size dist, phase function, Legendre decay).""" + +import matplotlib + +matplotlib.use("Agg") + +from pyradtran.models.aerosol_composite import ( + CompositeAerosol, + IntegrationConfig, + MieSpecies, + RefractiveIndex, + SizeDistribution, +) +from pyradtran.models.blocks import MassProfile, PlacedBlock + + +def _blocks(): + ri = RefractiveIndex(wavelength_um=[0.40, 0.70], n_real=[1.95, 1.95], k_imag=[0.79, 0.79]) + sd = SizeDistribution(kind="lognormal", params={"r_g_um": 0.1, "sigma_g": 2.0}) + return {"bc": sd}, {"bc": (ri, sd)} + + +def test_plot_size_distributions_draws_one_line_per_block(): + from pyradtran.viz.inputs import plot_size_distributions + + blocks, _ = _blocks() + fig, ax = plot_size_distributions(blocks) + assert len(ax.lines) == 1 + assert ax.get_xscale() == "log" + assert "Radius" in ax.get_xlabel() + + +def test_plot_phase_functions_draws_one_line_per_block_on_log_y(): + from pyradtran.viz.inputs import plot_phase_functions + + _, ri_sd = _blocks() + fig, ax = plot_phase_functions(ri_sd, wavelength_um=0.55) + assert len(ax.lines) == 1 + assert ax.get_yscale() == "log" + assert "Scattering angle" in ax.get_xlabel() + + +def test_plot_phase_functions_raises_on_discrete_unsupported(): + """Discrete SD has no single characteristic radius — raise clearly.""" + import pytest + + from pyradtran.viz.inputs import plot_phase_functions + + ri = RefractiveIndex(wavelength_um=[0.40, 0.70], n_real=[1.5, 1.5], k_imag=[0.0, 0.0]) + sd = SizeDistribution( + kind="discrete", params={"radius_um": [0.1, 1.0], "weights": [1.0, 1.0]} + ) + with pytest.raises(NotImplementedError): + plot_phase_functions({"d": (ri, sd)}) + + +def test_plot_legendre_decay_draws_stem(): + from pyradtran.viz.inputs import plot_legendre_decay + + ri = RefractiveIndex(wavelength_um=[0.40, 0.70], n_real=[1.5, 1.5], k_imag=[0.01, 0.01]) + sd = SizeDistribution(kind="lognormal", params={"r_g_um": 0.5, "sigma_g": 2.0}) + sp = MieSpecies( + refractive_index=ri, + size_distribution=sd, + particle_density_kg_m3=1000.0, + integration_config=IntegrationConfig(), + name="x", + ) + comp = CompositeAerosol( + pieces=[PlacedBlock(block=sp, profile=MassProfile(kg_m3_per_layer=(1e-7,)))], + wavelength_grid_um=[0.55], + altitude_grid_km=[2.0, 0.0], + n_legendre=8, + output_dir=".", + ) + fig, ax = plot_legendre_decay(comp, wavelength_um=0.55, n_legendre=8) + assert "moment" in ax.get_xlabel().lower() or "l" in ax.get_xlabel().lower()