From d7c3af09752c3075077415235a737293a9e75bde Mon Sep 17 00:00:00 2001 From: Brian Yu Date: Thu, 15 Oct 2026 22:35:28 +0800 Subject: [PATCH 1/7] Faster Runge-Kutta step (1.36x), bit-for-bit identical results Weighted sums use sum(map(mul, ...)), which keeps the products and their order (and Python 3.12's compensated float summation). A unit test pins the step against the plain loop; all 14 paper runs checked identical. --- src/outflow/integrator.py | 24 ++++++++++---------- tests/unit/test_integrator.py | 41 +++++++++++++++++++++++++++++++++++ 2 files changed, 53 insertions(+), 12 deletions(-) diff --git a/src/outflow/integrator.py b/src/outflow/integrator.py index 052fafc..d8273a0 100644 --- a/src/outflow/integrator.py +++ b/src/outflow/integrator.py @@ -12,6 +12,7 @@ """ from collections.abc import Callable, Sequence +from operator import mul Derivative = Callable[[float, tuple[float, ...]], Sequence[float]] @@ -27,6 +28,7 @@ ) _C5 = (37.0 / 378.0, 0.0, 250.0 / 621.0, 125.0 / 594.0, 0.0, 512.0 / 1771.0) _C4 = (2825.0 / 27648.0, 0.0, 18575.0 / 48384.0, 13525.0 / 55296.0, 277.0 / 14336.0, 1.0 / 4.0) +_ERROR_WEIGHTS = tuple(c5 - c4 for c5, c4 in zip(_C5, _C4, strict=True)) # fifth - fourth order # Growth factor used when the error estimate is exactly zero (the formula would divide by # zero). The paper code never met this case in its runs. @@ -49,20 +51,18 @@ def cash_karp_step( As in the paper code, the new state is the embedded fourth-order solution and the error is (fifth order - fourth order). """ - n = len(y) - stages = [tuple(h * s for s in first_slope)] + # stages[i] holds the stage increments k1, k2, ... of component i. The weighted sums use + # sum(map(mul, weights, increments)), which runs its loop in C: the solver spends most of + # its time in this function. The products and their order are those of the plain form + # sum(c * k for c, k in zip(weights, increments)), so the results are bit-for-bit the same. + stages = [[h * s] for s in first_slope] for stage in range(1, 6): - coefficients = _B[stage] - state = tuple( - y[i] + sum(c * k[i] for c, k in zip(coefficients, stages, strict=False)) - for i in range(n) - ) + state = [yi + sum(map(mul, _B[stage], ki)) for yi, ki in zip(y, stages, strict=True)] slope = derivative(r + _A[stage] * h, state) - stages.append(tuple(h * s for s in slope)) - new_y = tuple(y[i] + sum(c * k[i] for c, k in zip(_C4, stages, strict=True)) for i in range(n)) - error = tuple( - sum((c5 - c4) * k[i] for c5, c4, k in zip(_C5, _C4, stages, strict=True)) for i in range(n) - ) + for ki, s in zip(stages, slope, strict=True): + ki.append(h * s) + new_y = tuple(yi + sum(map(mul, _C4, ki)) for yi, ki in zip(y, stages, strict=True)) + error = tuple(sum(map(mul, _ERROR_WEIGHTS, ki)) for ki in stages) return new_y, error diff --git a/tests/unit/test_integrator.py b/tests/unit/test_integrator.py index 7bf1bf7..754441d 100644 --- a/tests/unit/test_integrator.py +++ b/tests/unit/test_integrator.py @@ -1,6 +1,7 @@ """Cash-Karp Runge-Kutta step and the paper code's error measure and step control.""" import math +import random import pytest @@ -37,6 +38,46 @@ def derivative(_r, state): assert errors[0] / errors[1] == pytest.approx(32, rel=0.2) +def _reference_cash_karp_step(derivative, h, r, y, first_slope): + """The step as first ported (a plain loop over stages), kept to pin down the arithmetic.""" + n = len(y) + stages = [tuple(h * s for s in first_slope)] + for stage in range(1, 6): + state = tuple( + y[i] + sum(c * k[i] for c, k in zip(integrator._B[stage], stages, strict=False)) + for i in range(n) + ) + stages.append(tuple(h * s for s in derivative(r + integrator._A[stage] * h, state))) + new_y = tuple( + y[i] + sum(c * k[i] for c, k in zip(integrator._C4, stages, strict=True)) for i in range(n) + ) + error = tuple( + sum( + (c5 - c4) * k[i] + for c5, c4, k in zip(integrator._C5, integrator._C4, stages, strict=True) + ) + for i in range(n) + ) + return new_y, error + + +def test_cash_karp_step_is_bit_identical_to_the_plain_loop(): + """The faster step gives exactly the same floats as the plain loop (no change in results).""" + rng = random.Random(1) + + def derivative(r, state): + v, p, pc = state + return (math.sin(r * v) + p / (1.0 + pc * pc), v * p - r, math.exp(-abs(pc)) * v) + + for _ in range(500): + y = tuple(rng.uniform(-3.0, 3.0) * 10.0 ** rng.randint(-20, 5) for _ in range(3)) + r = rng.uniform(0.0, 2.0) + h = 10.0 ** rng.uniform(-8.0, 0.0) + slope = derivative(r, y) + expected = _reference_cash_karp_step(derivative, h, r, y, slope) + assert integrator.cash_karp_step(derivative, h, r, y, slope) == expected + + @pytest.mark.parametrize( "values, errors", [ From c84e5173af2598d20ecb63ee698f6bb04a248552 Mon Sep 17 00:00:00 2001 From: Brian Yu Date: Thu, 15 Oct 2026 22:35:28 +0800 Subject: [PATCH 2/7] Central-density search: Brent's method for the last stage (default) Finds r_s = r_sb with about a quarter of the trials, to the same 1e-12 precision; rho_0 agrees with the paper code's bisection to 2e-12. method='bisection' keeps the paper code's search. --- docs/physics.md | 8 +++-- src/outflow/boundary.py | 60 ++++++++++++++++++++++++++++++++++--- tests/unit/test_boundary.py | 18 +++++++++++ 3 files changed, 80 insertions(+), 6 deletions(-) diff --git a/docs/physics.md b/docs/physics.md index 6a49ff7..64dccfc 100644 --- a/docs/physics.md +++ b/docs/physics.md @@ -74,11 +74,15 @@ switch at r_sb (density for r < r_sb, Alfvén speed for r ≤ r_sb). The port fo | `central_sound_speed_squared` | eq. 35 at r = 0 | with the L'Hôpital ratio v_A/v | | `sonic_trial`, `_sonic_trial`, `SonicTrial` | §2.3.2, §2.4 (code: `rsonic()`) | integrate from r = 0 until sonic or turning over | | `_bisect` | §2.3.2 (code: `gss()`) | bisection on ρ₀ | +| `_brent`, `sonic_radius_mismatch` | §2.3.2 (not in the paper code; replaces the last `gss()` stage) | Brent's method on r_s(ρ₀) − r_sb (default) | | `find_central_density` | §2.3.2, §2.4 (code: `CRline()`) | ρ₀ for which the sonic radius is r_sb | The search: start at ρ₀ = 0.1 M☉/pc³, multiply by 10 until a trial turns over or reaches r_sb, -bisect for the critical ("peak") ρ₀, then bisect for the ρ₀ whose sonic radius is r_sb. Trials -use the inner density and Alfvén forms at every radius (the code's `'find0'` tag). +bisect for the critical ("peak") ρ₀, then find the ρ₀ whose sonic radius is r_sb. The paper +code bisects in that last stage; by default the port uses Brent's method (scipy's `brentq`) +there, to the same precision with fewer trials (`method="bisection"` gives the paper code's +search). Trials use the inner density and Alfvén forms at every radius (the code's `'find0'` +tag). ## Solver (`solver.py`, `integrator.py`) diff --git a/src/outflow/boundary.py b/src/outflow/boundary.py index bdc9c8f..065f9b4 100644 --- a/src/outflow/boundary.py +++ b/src/outflow/boundary.py @@ -10,8 +10,10 @@ (dv/dr <= 0 while still subsonic: a "breeze", rho_0 too high) or reaches r_sb. 2. Bisect between 0 and that rho_0 for the critical value that separates winds that turn supersonic from breezes (the "peak"); this gives the largest possible sonic radius. -3. If that sonic radius lies beyond r_sb, bisect between 0 and the peak value for the rho_0 - whose sonic radius is r_sb. +3. If that sonic radius lies beyond r_sb, find the rho_0 between 0 and the peak value whose + sonic radius is r_sb. The paper code bisects; by default this port uses Brent's method, + which reaches the same precision with fewer trials (``method="bisection"`` restores the + paper code's search). Trial solutions use the inner forms of density and Alfven speed at every radius (``searching=True``, the paper code's tag ``'find0'``). @@ -36,6 +38,10 @@ SMALLEST_STEP_FRACTION = 1e-50 # Safety limit (not in the paper code): stop multiplying rho_0 by 10 after this many tries. MAX_DENSITY_INCREASES = 40 +# Brent search (not in the paper code): halve rho_0 at most this many times looking for a sonic +# radius inside r_sb before falling back to the bisection. 40 halvings reach the bisection's +# own precision (2^-40 ~ 1e-12). +MAX_BRACKET_HALVINGS = 40 # Fallback for the first step from r = 0 (not in the paper code). Near r = 0 the equations are # unstable within a step: round-off in the Runge-Kutta stages grows by a factor of ~40 per # stage, whatever the step size, so the error estimate of the first step has a floor of about @@ -181,12 +187,55 @@ def _bisect(equations: WindEquations, params, low: float, high: float, find_peak return trial.radius, low -def find_central_density(params: WindParameters, equations: WindEquations | None = None): +def _brent(equations: WindEquations, params, peak_density: float) -> float: + """rho_0 in (0, peak_density) whose sonic radius is r_sb, by Brent's method. + + Replaces the paper code's second bisection (decision of 2026-10-15, docs/decisions.md): + the sonic radius changes smoothly with rho_0 below the peak, so Brent's method (bisection + combined with secant and inverse quadratic steps) needs about a quarter of the trials. It + stops at the same absolute precision as the bisection, SEARCH_TOLERANCE * peak_density. + Falls back to the bisection if the bracket cannot be set up. + """ + from scipy.optimize import brentq # imported here: only needed for this search + + def sonic_radius_mismatch(density: float) -> float: + trial = sonic_trial(equations, central_state(params, density)) + return trial.radius - params.starburst_radius # > 0: rho_0 too high + + high = peak_density + if not sonic_radius_mismatch(high) > 0: + return _bisect(equations, params, 0.0, peak_density, find_peak=False)[1] + # Lower end of the bracket: halve rho_0 until the sonic radius lies inside r_sb (these are + # the same first trials the bisection makes). + low = peak_density / 2.0 + for _ in range(MAX_BRACKET_HALVINGS): + mismatch = sonic_radius_mismatch(low) + if mismatch == 0: + return low + if mismatch < 0: + return brentq(sonic_radius_mismatch, low, high, xtol=SEARCH_TOLERANCE * peak_density) + high, low = low, low / 2.0 + return _bisect(equations, params, 0.0, peak_density, find_peak=False)[1] + + +def find_central_density( + params: WindParameters, equations: WindEquations | None = None, method: str = "brent" +): """rho_0 for which the wind's sonic radius is r_sb, as in the paper code's ``CRline()``. + Args: + params: model inputs. + equations: the wind equations for ``params`` (built if not given). + method: how the last stage finds the rho_0 with sonic radius r_sb: ``"brent"`` + (default; the whole solve is about 1.6x faster) or ``"bisection"`` (the paper + code's method). + Both stop at the same precision; the first two stages are the paper code's. + Raises: NoSolutionError: if no trial density turns the flow over within the safety limit. """ + if method not in ("brent", "bisection"): + raise ValueError(f"method must be 'brent' or 'bisection', got {method!r}") if equations is None: equations = WindEquations(params) density = FIRST_DENSITY_TRIAL @@ -203,6 +252,9 @@ def find_central_density(params: WindParameters, equations: WindEquations | None trial = sonic_trial(equations, central_state(params, density)) sonic_radius, peak_density = _bisect(equations, params, 0.0, density, find_peak=True) if sonic_radius > params.starburst_radius: - _, density = _bisect(equations, params, 0.0, peak_density, find_peak=False) + if method == "brent": + density = _brent(equations, params, peak_density) + else: + _, density = _bisect(equations, params, 0.0, peak_density, find_peak=False) # Otherwise the paper code keeps the last trial of step 1 (see docs/open-items.md). return density diff --git a/tests/unit/test_boundary.py b/tests/unit/test_boundary.py index a75ce3e..7b80daf 100644 --- a/tests/unit/test_boundary.py +++ b/tests/unit/test_boundary.py @@ -25,6 +25,24 @@ def test_central_state_eqs_44_45(): assert state.velocity_gradient == pytest.approx(q / (3 * rho0), rel=1e-14) +def test_unknown_search_method_is_rejected(): + with pytest.raises(ValueError, match="brent"): + boundary.find_central_density(WindParameters.paper_default(), method="secant") + + +@pytest.mark.slow +@pytest.mark.parametrize("run_name", ["fig6_thermal", "fig7_B0_100uG", "fig8_kappa_1e5"]) +def test_brent_search_finds_the_bisection_density(run_name): + """Brent's method (default) and the paper code's bisection agree on rho_0. + + Both stop when rho_0 is known to 1e-12 of the search range, so they agree to a few 1e-12. + """ + run = next(r for figure in (6, 7, 8) for r in paper_runs(figure) if r.name == run_name) + brent = boundary.find_central_density(run.params, method="brent") + bisection = boundary.find_central_density(run.params, method="bisection") + assert brent == pytest.approx(bisection, rel=1e-11) + + def test_first_step_fallback_rescues_round_off_start(): """A density where the first step from r = 0 sits on its round-off floor (Fig. 7, 100 uG). From 4e82cdd121de5bcde9738f96a4be8a2d29f2f77b Mon Sep 17 00:00:00 2001 From: Brian Yu Date: Thu, 15 Oct 2026 22:35:28 +0800 Subject: [PATCH 3/7] App: figure tabs solve their runs in parallel (Fig. 7: 20-40 s -> 5.5 s) One worker process per run, falling back to one after another if processes cannot start. The page code is now under if __name__ == '__main__' because workers re-import the script. --- app/streamlit_app.py | 68 ++++++++++++++++++++++++++++++++++--------- tests/app/test_app.py | 28 ++++++++++++++++++ 2 files changed, 83 insertions(+), 13 deletions(-) diff --git a/app/streamlit_app.py b/app/streamlit_app.py index 90ff47e..65666b8 100644 --- a/app/streamlit_app.py +++ b/app/streamlit_app.py @@ -11,8 +11,11 @@ """ import math +import os import tempfile import traceback +from concurrent.futures import ProcessPoolExecutor, as_completed +from concurrent.futures.process import BrokenProcessPool from dataclasses import dataclass from pathlib import Path @@ -60,10 +63,18 @@ def solve_safely(build_params, outer_radius_in_sb: float) -> RunOutcome: ``build_params`` is a function returning :class:`WindParameters`, so that invalid inputs are reported the same way as solver failures. """ - try: + + def solve(): params = build_params() - outer_radius = outer_radius_in_sb * params.starburst_radius - return RunOutcome(profile=solve_wind(params, outer_radius=outer_radius)) + return solve_wind(params, outer_radius=outer_radius_in_sb * params.starburst_radius) + + return outcome_of(solve) + + +def outcome_of(compute) -> RunOutcome: + """Call ``compute`` (which returns a profile); turn every failure into a plain message.""" + try: + return RunOutcome(profile=compute()) except NoSolutionError as error: return RunOutcome(error=error.message, details=error.details) except ValueError as error: @@ -469,16 +480,43 @@ class FigureRunResult: def run_figure(spec: FigureSpec): - """Recompute every run of one figure, showing "run k of n" while it works.""" + """Recompute every run of one figure, in parallel, showing "k of n finished".""" runs = paper_runs(spec.number) - progress = st.progress(0.0, text=f"Run 1 of {len(runs)}: {runs[0].label}…") - results = [] - for index, run in enumerate(runs): - progress.progress(index / len(runs), text=f"Run {index + 1} of {len(runs)}: {run.label}…") - outcome = solve_safely(lambda run=run: run.params, run.outer_radius_in_sb) - results.append(FigureRunResult(run, outcome, load_paper_profile(run))) + progress = st.progress(0.0, text=f"Solving the {len(runs)} runs at the same time…") + + def show_progress(finished: int): + progress.progress(finished / len(runs), text=f"{finished} of {len(runs)} runs finished…") + + outcomes = solve_runs_in_parallel(runs, show_progress) progress.empty() - st.session_state[f"fig{spec.number}_results"] = results + st.session_state[f"fig{spec.number}_results"] = [ + FigureRunResult(run, outcome, load_paper_profile(run)) + for run, outcome in zip(runs, outcomes, strict=True) + ] + + +def solve_runs_in_parallel(runs: list[PaperRun], on_finished) -> list[RunOutcome]: + """Solve the runs in separate processes (one per run, up to the number of CPU cores). + + The runs are independent, so this changes only the waiting time, not the results. If + worker processes cannot be started (some hosting platforms forbid them), the runs are + solved one after another instead. ``on_finished(k)`` is called as each run finishes. + """ + workers = min(len(runs), os.cpu_count() or 1) + try: + with ProcessPoolExecutor(max_workers=workers) as pool: + futures = [ + pool.submit(solve_wind, run.params, outer_radius=run.outer_radius) for run in runs + ] + for finished, _ in enumerate(as_completed(futures), start=1): + on_finished(finished) + return [outcome_of(future.result) for future in futures] + except (BrokenProcessPool, OSError, NotImplementedError): + outcomes = [] + for finished, run in enumerate(runs, start=1): + outcomes.append(solve_safely(lambda run=run: run.params, run.outer_radius_in_sb)) + on_finished(finished) + return outcomes def figure_series(results: list[FigureRunResult], show_published: bool) -> list[plots.Series]: @@ -526,7 +564,7 @@ def figure_tab(spec: FigureSpec): f"Run the {len(runs)} cases of Fig. {spec.number}", type="primary", key=f"fig{spec.number}_run", - help="Each case takes a few seconds; the whole figure about half a minute.", + help="The cases are solved at the same time; a few seconds on a multi-core computer.", ) if pressed: run_figure(spec) @@ -636,4 +674,8 @@ def main(): about_tab() -main() +# Streamlit runs this script as "__main__". The worker processes of the figure tabs load it +# under another name (Windows and macOS start workers by re-importing the main script), and +# this guard stops them from drawing the page or starting workers of their own. +if __name__ == "__main__": + main() diff --git a/tests/app/test_app.py b/tests/app/test_app.py index 1319efc..83803e3 100644 --- a/tests/app/test_app.py +++ b/tests/app/test_app.py @@ -149,6 +149,34 @@ def test_fig7_run_plots_four_runs_and_four_published_overlays(): assert "published" not in kinds +def test_figure_runs_one_by_one_when_processes_are_unavailable(monkeypatch): + """Some hosts forbid worker processes; the figure tab then solves the runs in turn.""" + import concurrent.futures + + import outflow + from outflow import load_paper_profile, paper_runs + + def no_processes(*args, **kwargs): + raise OSError("worker processes are not allowed here") + + published = {run.params: load_paper_profile(run) for run in paper_runs(7)} + + def published_profile(params, outer_radius=None, cooling=False): + return published[params] # stands in for the solver: instant and deterministic + + # The app imports both names afresh on every run, so patch them where they come from. + monkeypatch.setattr(concurrent.futures, "ProcessPoolExecutor", no_processes) + monkeypatch.setattr(outflow, "solve_wind", published_profile) + app = start_app() + app.button(key="fig7_run").click() + app.run() + assert not app.exception, app.exception + assert not app.error, [e.value for e in app.error] + traces = plotted_traces(app) + velocity = [t for t in traces if t["meta"]["panel"] == "v" and not t["meta"]["ratio"]] + assert [t["meta"]["kind"] for t in velocity].count("computed") == 4 + + def test_about_tab_has_paper_links_and_energy_note(): app = start_app() about = app.tabs[5] From 3730693fb8640792e87d4fe350a734efb777ce6a Mon Sep 17 00:00:00 2001 From: Brian Yu Date: Thu, 15 Oct 2026 22:35:28 +0800 Subject: [PATCH 4/7] Docs: speed decisions, Brent search note, parallel figure tabs, Numba for later --- CLAUDE.md | 8 +++++++- README.md | 8 +++++++- docs/architecture.md | 7 ++++++- docs/decisions.md | 37 +++++++++++++++++++++++++++++++++++-- docs/known-issues.md | 14 ++++++++++++++ docs/open-items.md | 5 +++++ 6 files changed, 74 insertions(+), 5 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 3ec2eda..cf05c4f 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -36,7 +36,13 @@ Read these before substantial changes: output, now stored as package data in `src/outflow/data/paper/`, not `legacy/outflow.py` (which has a temperature bug and auto-enables cooling). - Defaults follow the paper code exactly (values, rounded constants, tolerances, numerical - method). Changing any of them is a decision for the user. + method). Changing any of them is a decision for the user. One such decision so far: the last + stage of the ρ₀ search uses Brent's method (`docs/decisions.md`, "Speed"). +- A speed-up that changes results is a numerics change (the user's decision). To show a + speed-up changes nothing, use `tests/unit/test_integrator.py` (bit-identical RK step) and + compare all 14 paper runs before and after. +- The app script keeps its `if __name__ == "__main__": main()` guard: the figure tabs' worker + processes re-import it. - Cooling is off by default and is not exposed in the UI. ## Conventions diff --git a/README.md b/README.md index 92ac702..b496d80 100644 --- a/README.md +++ b/README.md @@ -78,7 +78,8 @@ streamlit run app/streamlit_app.py Tabs: **Explore** (change any input, overlay runs, download CSV), **Fig. 6–9** (recompute the paper's run sets and compare them with the published curves) and **About**. Every run is -computed when you press its button; a figure tab takes 20–40 s. +computed when you press its button. A figure tab solves its runs at the same time, one per CPU +core, so it takes about as long as its slowest run (about 5 s on a multi-core computer). ## How the results compare with the paper @@ -87,6 +88,11 @@ published runs are recomputed and compared between 0.2 r_sb and the outer radius relative differences are about 3e-5 in velocity and 2e-4 in pressure (the tests allow 1e-3). The thermal-only case matches the analytic Chevalier & Clegg wind to 3e-6. +One step differs from the paper code, for speed: the last stage of the search for the central +density uses Brent's method instead of bisection, to the same precision (central densities agree +to 2e-12). `outflow.boundary.find_central_density(params, method="bisection")` gives the paper +code's search. + **A note on Ė.** The published Figs. 6–9 were made with an energy injection rate Ė = 4.2e41 erg/s (the value in the paper's code), although the paper's Table 1 lists 1e42 erg/s. The defaults here use 4.2e41 erg/s so that they reproduce the figures. diff --git a/docs/architecture.md b/docs/architecture.md index 51a7e88..4429a5b 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -37,7 +37,7 @@ paper.load_paper_profile (published runs, package data CSV) | `physics.py` | single formulas of the paper (density, Alfvén speed, c\*², luminosity, radiative force and heating, NFW gravity) | | `equations.py` | `WindEquations`: d(v, P, P_C)/dr, inner and outer equations | | `integrator.py` | Cash–Karp step, error measure and step-size rule of the paper code | -| `boundary.py` | central values (eqs. 44–45) and the search for ρ₀ | +| `boundary.py` | central values (eqs. 44–45) and the search for ρ₀ (last stage: Brent's method by default, or the paper code's bisection) | | `solver.py` | `solve_wind` and the outward integration, stall detection | | `profile.py` | `WindProfile`: arrays, derived quantities, interpolation, comparison, CSV read/write | | `paper.py` | the 14 published runs (`paper_runs`) and their saved profiles (`load_paper_profile`) | @@ -61,6 +61,11 @@ sit on top. Runs are computed fresh on each button press (no caching). Explore keeps up to 8 runs of the session in `st.session_state` for overlay. Cooling is not offered in the app. +A figure tab solves its runs in parallel (`solve_runs_in_parallel`: a `ProcessPoolExecutor` +with one worker per run, up to the number of CPU cores), and falls back to solving them one +after another if worker processes cannot be started. Workers re-import the app script, so its +page code runs only under `if __name__ == "__main__"`. + ## Other folders | Path | Role | diff --git a/docs/decisions.md b/docs/decisions.md index 9129c6c..e81f46f 100644 --- a/docs/decisions.md +++ b/docs/decisions.md @@ -32,7 +32,9 @@ minor decisions made without asking; the user can overturn them. thermal-only limit. (user, default accepted) - No command-line tool. The scriptable interface is the Python API; the UI is the only app. This overrides the "CLI/UI parity" app convention. (user) -- Speed improvements are out of scope. (Claude's call; never raised as a goal) +- Speed improvements are out of scope. (Claude's call; never raised as a goal) **Reversed + 2026-10-15:** the user asked for speed-ups 1, 2 and 4 of the proposal (see "Speed"); Numba + (compiling the hot loop) is left for later. ## Physics and numerics (follow the paper code) @@ -60,6 +62,8 @@ minor decisions made without asking; the user can overturn them. (dv/dr = q/3ρ₀, dv_A/dr = B₀/(r_sb √(4πρ₀))). - ρ₀ found as in `CRline.CRline()`: start at 0.1 (natural units), multiply by 10 until the trial solution turns over or reaches r_sb, bisect for the peak, then bisect for r_s = r_sb. + **Changed 2026-10-15 (user):** the last stage uses Brent's method by default (see "Speed"); + the paper code's bisection stays available as `method="bisection"`. - Source-term ("inner") equations are used while v² ≤ c*² (subsonic), as the paper code does; density and Alfvén speed switch at r_sb (ρ: r < r_sb; v_A: r ≤ r_sb). Documented in `docs/physics.md` as the paper code's convention. @@ -90,6 +94,34 @@ minor decisions made without asking; the user can overturn them. - The temperature is computed from each radius's own density (the `outflow.py` bug is not carried over). +## Speed (2026-10-15, user: "implement 1, 2, 4, leave numba for later") + +Baseline: the 14 paper runs took 95 s in total; 96 % of a run is the ρ₀ search (81 trial +integrations for the paper default), and half of each trial was the Runge–Kutta step's Python +overhead. Built on branch `feature/faster-solver`: + +1. **Faster Runge–Kutta step, same results.** `cash_karp_step` computes its weighted sums with + `sum(map(mul, weights, increments))`. The products and their order are unchanged, so every + result is bit-for-bit the same as before (all 14 runs checked; a unit test pins it against + the plain loop). Writing the sums as plain `a + b + c` would be faster still (~1.6× vs + 1.36×) but not identical: since Python 3.12, `sum()` adds floats with compensated + summation. (Claude's call: keep results identical.) 95 s → 70 s. +2. **Figure tabs solve their runs in parallel**, one process per run (up to the number of CPU + cores), with a fallback to one after another if processes cannot be started. Results are + unchanged. Fig. 7 tab: 20–40 s before, 5.5 s now (measured in the real app on a 20-core + machine). The app script is guarded by `if __name__ == "__main__"`, because worker + processes re-import it. +3. **Brent's method for the last stage of the ρ₀ search** (finding r_s = r_sb), stopping at the + same absolute precision as the bisection (1e-12 of the search range). The first two stages + (×10 scan, peak bisection) are unchanged: the peak is where the trial switches from wind to + breeze, a jump that suits bisection only. Measured over the 14 runs: ρ₀ differs from the + bisection result by ≤ 1.8e-12; profiles move by ≤ 2.1e-4 relative (the transonic solution + is that sensitive to ρ₀ at this precision); the match to the published runs stays within + 1.7e-4 (was 2.1e-4). 70 s → 43.5 s. The paper code's search remains available as + `find_central_density(..., method="bisection")`. + +Overall: 95 s → 43.5 s for the 14 runs (2.2×); a default Explore run about 1.5 s. + ## Verification - "Same results" = differences not visible by eye. Test threshold: every reference run matches @@ -158,7 +190,8 @@ From the builder check on these documents. All minor; the user can overturn any. made in this session are kept in the page's memory so they can be overlaid, and can be cleared. - **App, progress:** a progress indicator shows "run k of n" while a figure tab computes - (about 20–40 s locally). + (about 20–40 s locally). Since 2026-10-15 the runs are solved in parallel and it shows + "k of n runs finished" (about 5 s locally). - **App, axes:** as in the published figures. x is r/r_sb on a linear axis (0–5 for Figs. 6–8, 0–30 for Fig. 9); v in km/s, n in cm⁻³ (log), T in K (log), P in erg cm⁻³ (log). The Explore tab has a log-r toggle. diff --git a/docs/known-issues.md b/docs/known-issues.md index c3ba800..c79dbcb 100644 --- a/docs/known-issues.md +++ b/docs/known-issues.md @@ -69,6 +69,20 @@ The central-density search stops after 40 increases by ×10 (`NoSolutionError`). integration stops after 2 000 000 steps: inside 1.1 r_sb this raises `NoSolutionError`, beyond it the profile is returned as stalled. Neither limit is reached by the published runs. +## Profiles are sensitive to ρ₀ at the 1e-4 level + +Near the sonic point the solution is very sensitive to the central density: changing ρ₀ by +about 1e-12 (the search's own precision) moves the profiles by up to 2e-4 relative. This is why +Brent's method and the paper code's bisection, which agree on ρ₀ to 2e-12, give profiles that +differ by up to 2e-4, and why the match to the published runs cannot be much better than +that. Differences of this size are far below what is visible in a plot. + +## Figure tabs on machines with few cores + +The figure tabs solve their runs in parallel, one process per CPU core. On a machine or host +with one or two cores (e.g. free cloud hosting) the gain is correspondingly smaller. Where +worker processes cannot be started at all, the runs are solved one after another. + ## Repeated radii in the published data The published profiles repeat some radii (where the step size fell below round-off near r_sb, diff --git a/docs/open-items.md b/docs/open-items.md index 9b2b09a..f0ae29d 100644 --- a/docs/open-items.md +++ b/docs/open-items.md @@ -61,6 +61,11 @@ here; a limit or trap we live with goes in `known-issues.md`. Each item in one p ## Later +- Numba (compiling the integrator and equations to machine code) for a further speed-up, + likely 20–50× (not measured). Costs: a new dependency, a few seconds of compiling on first + use, and the equations rewritten as plain functions, which reads less like the paper. (user, + 2026-10-15: "leave numba for later") + - Port Figs. 1–5 scripts and the §3.5 line-emission work (Figs. 10–11). - Validate the experimental cooling option against `Silich.py` (Fig. 2). - Possible erratum on Ė in Table 1 (user's decision). From 2d3d8d2fa4c0c55b7a315335514b60c7ccf133d6 Mon Sep 17 00:00:00 2001 From: Brian Yu Date: Thu, 15 Oct 2026 22:55:57 +0800 Subject: [PATCH 5/7] =?UTF-8?q?Optional=20Numba-compiled=20=CF=81=E2=82=80?= =?UTF-8?q?=20search:=20bit-identical=20results,=20~30x=20faster?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit src/outflow/_compiled.py repeats the trial integration line by line. Exponents are read at run time so powers call pow() as Python does, and sums repeat Python 3.12's compensated sum(), so all 14 paper runs are bit-for-bit identical (tested). Installed with the 'fast' extra; without Numba the Python code runs. --- .github/workflows/ci.yml | 2 +- environment.yml | 3 +- pyproject.toml | 2 + src/outflow/_compiled.py | 330 ++++++++++++++++++++++++++++++++++++ src/outflow/boundary.py | 59 ++++++- tests/unit/test_compiled.py | 94 ++++++++++ 6 files changed, 487 insertions(+), 3 deletions(-) create mode 100644 src/outflow/_compiled.py create mode 100644 tests/unit/test_compiled.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 5a8d5d4..de787c8 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -20,7 +20,7 @@ jobs: cache: pip - name: Install - run: pip install -e ".[dev]" + run: pip install -e ".[dev,fast]" - name: Lint run: | diff --git a/environment.yml b/environment.yml index f128a26..4d23c54 100644 --- a/environment.yml +++ b/environment.yml @@ -14,4 +14,5 @@ dependencies: - plotly==7.1.0 - pytest==9.1.1 - ruff==0.17.0 - - -e .[dev] + - numba==0.68.0 + - -e .[dev,fast] diff --git a/pyproject.toml b/pyproject.toml index e5bb568..0a20b4b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,6 +20,8 @@ dependencies = [ [project.optional-dependencies] dev = ["pytest>=8", "ruff>=0.5"] +# Compiles the central-density search with Numba: same results, about 30x faster. +fast = ["numba>=0.60"] [tool.setuptools.dynamic] version = { attr = "outflow.__version__" } diff --git a/src/outflow/_compiled.py b/src/outflow/_compiled.py new file mode 100644 index 0000000..66f8f41 --- /dev/null +++ b/src/outflow/_compiled.py @@ -0,0 +1,330 @@ +"""Compiled copy (Numba) of the trial integration used by the central-density search. + +This module is optional: it is used only when Numba is installed, and only for winds without +cooling. It repeats, line by line, the arithmetic of :mod:`outflow.physics`, +:mod:`outflow.equations`, :mod:`outflow.integrator` and ``boundary._sonic_trial``; those +modules remain the readable reference, and the tests check that this copy gives bit-for-bit +the same floats. + +Two details keep the results identical to the Python code: + +- **Powers.** Python evaluates ``x**2`` by calling the C library's ``pow``, which is not + always the exactly rounded ``x*x``. A compiler that sees a constant exponent replaces + ``x**2`` (even ``x**2.0``) by ``x*x``. So every exponent here is read from the parameter + array at run time (``p[TWO]``, ``p[THREE]``, ...): the compiler cannot see its value and calls + ``pow``, as Python does. +- **Sums.** Since Python 3.12, ``sum()`` adds floats with Neumaier's compensated summation; + ``_compensated_sum`` repeats that algorithm. +""" + +import math + +import numpy as np +from numba import njit + +from outflow.constants import FOUR_PI, GAMMA_CR, GAMMA_GAS, GRAVITATIONAL_CONSTANT, SPEED_OF_LIGHT +from outflow.integrator import _A, _B, _C4, _ERROR_WEIGHTS +from outflow.physics import NFW_SERIES_LIMIT, NFW_SERIES_TERMS + +# Cash-Karp tableau as arrays (row `stage` of _STAGE_WEIGHTS holds _B[stage], zero-padded). +_STAGE_NODES = np.array(_A) +_STAGE_WEIGHTS = np.zeros((6, 5)) +for _stage, _row in enumerate(_B): + _STAGE_WEIGHTS[_stage, : len(_row)] = _row +_FOURTH_ORDER_WEIGHTS = np.array(_C4) +_ERROR_WEIGHTS_ARRAY = np.array(_ERROR_WEIGHTS) + +# Positions of the inputs in the packed parameter array (see pack_parameters). +( + STARBURST_RADIUS, + MASS_RATE, + TERMINAL_VELOCITY, + MASS_INJECTION_RATE, + THERMAL_INJECTION_RATE, + CR_INJECTION_RATE, + OPACITY, + MAGNETIC_FIELD, + NFW_DENSITY, + NFW_SCALE_RADIUS, + # Exponents, passed at run time so that powers call pow() as Python does (see above). + ONE, + TWO, + THREE, + GROWTH_EXPONENT, + SHRINK_EXPONENT, +) = range(15) + + +def pack_parameters(params) -> np.ndarray: + """The WindParameters values the equations use, and the exponents, as a float array.""" + return np.array( + [ + params.starburst_radius, + params.mass_rate, + params.terminal_velocity, + params.mass_injection_rate, + params.thermal_energy_injection_rate, + params.cr_energy_injection_rate, + params.opacity, + params.magnetic_field, + params.nfw_characteristic_density, + params.nfw_scale_radius, + 1.0, + 2.0, + 3.0, + 0.2, # step growth exponent of integrator.next_step_size + 0.25, # step shrink exponent of integrator.next_step_size + ] + ) + + +@njit(cache=True) +def _compensated_sum(values, count): + """sum(values[:count]) exactly as Python 3.12's built-in sum() adds floats.""" + total = 0.0 + compensation = 0.0 + for i in range(count): + x = values[i] + t = total + x + if abs(total) >= abs(x): + compensation += (total - t) + x + else: + compensation += (x - t) + total + total = t + if compensation != 0.0 and math.isfinite(compensation): + total += compensation + return total + + +# --- physics.py and equations.py -------------------------------------------------------------- + + +@njit(cache=True) +def _density(r, v, p, searching): + starburst_radius = p[STARBURST_RADIUS] + if searching or r < starburst_radius: + return p[MASS_RATE] * r / (FOUR_PI * starburst_radius ** p[THREE] * v) + return p[MASS_RATE] / (FOUR_PI * r ** p[TWO] * v) + + +@njit(cache=True) +def _alfven_speed(r, density, p, searching): + starburst_radius = p[STARBURST_RADIUS] + if searching or r <= starburst_radius: + field = p[MAGNETIC_FIELD] * r / starburst_radius + else: + field = p[MAGNETIC_FIELD] * starburst_radius ** p[TWO] / r ** p[TWO] + return field / math.sqrt(FOUR_PI * density) + + +@njit(cache=True) +def _sound_speed_squared_from(P, Pc, rho, v, v_alfven): + cr_factor = ( + (2.0 * v + v_alfven) * (v - (GAMMA_GAS - 1.0) * v_alfven) / (2.0 * v * (v + v_alfven)) + ) + return GAMMA_GAS * P / rho + GAMMA_CR * Pc * cr_factor / rho + + +@njit(cache=True) +def _sound_speed_squared(r, v, P, Pc, p, searching): + rho = _density(r, v, p, searching) + return _sound_speed_squared_from(P, Pc, rho, v, _alfven_speed(r, rho, p, searching)) + + +@njit(cache=True) +def _luminosity(v, P, Pc, rho, cr_velocity_ratio, p): + return p[MASS_RATE] * ( + p[TERMINAL_VELOCITY] ** p[TWO] / 2.0 + - v ** p[TWO] / 2.0 + - GAMMA_GAS * P / ((GAMMA_GAS - 1.0) * rho) + - GAMMA_CR * cr_velocity_ratio * Pc / (rho * (GAMMA_CR - 1.0)) + ) + + +@njit(cache=True) +def _radiative_force(r, luminosity, p, inner): + opacity, starburst_radius = p[OPACITY], p[STARBURST_RADIUS] + if inner: + return opacity * luminosity * r / (FOUR_PI * SPEED_OF_LIGHT * starburst_radius ** p[THREE]) + return opacity * luminosity / (FOUR_PI * SPEED_OF_LIGHT * r ** p[TWO]) + + +@njit(cache=True) +def _radiative_power_factor(r, luminosity, p): + opacity, starburst_radius = p[OPACITY], p[STARBURST_RADIUS] + if r == 0: + return 3.0 * opacity * luminosity / (FOUR_PI * starburst_radius ** p[TWO] * SPEED_OF_LIGHT) + if r == starburst_radius: + return ( + 3.0 + * opacity + * luminosity + / (8.0 * math.pi * starburst_radius ** p[TWO] * SPEED_OF_LIGHT) + ) + shape = 2.0 * starburst_radius * r - (r ** p[TWO] - starburst_radius ** p[TWO]) * math.log( + (r + starburst_radius) / abs(r - starburst_radius) + ) + return ( + 3.0 + * opacity + * luminosity + * shape + / (16.0 * math.pi * starburst_radius ** p[THREE] * SPEED_OF_LIGHT * r) + ) + + +@njit(cache=True) +def _nfw_gravity(r, p): + density, scale_radius = p[NFW_DENSITY], p[NFW_SCALE_RADIUS] + if r > scale_radius * NFW_SERIES_LIMIT: + shape = math.log((scale_radius + r) / scale_radius) - r / (scale_radius + r) + return ( + FOUR_PI + * GRAVITATIONAL_CONSTANT + * density + * scale_radius ** p[THREE] + * shape + / r ** p[TWO] + ) + x = -r / scale_radius + terms = np.empty(NFW_SERIES_TERMS) + for i in range(NFW_SERIES_TERMS): + terms[i] = (i + 1) * x ** (i * p[ONE]) / (i + 2) # x**i, exponent opaque to compiler + series = _compensated_sum(terms, NFW_SERIES_TERMS) + return FOUR_PI * GRAVITATIONAL_CONSTANT * density * scale_radius * series + + +@njit(cache=True) +def _derivatives(r, v, P, Pc, p, searching, out): + """WindEquations.derivatives without cooling (r > 0), written into ``out``.""" + gamma, gamma_cr = GAMMA_GAS, GAMMA_CR + rho = _density(r, v, p, searching) + v_alfven = _alfven_speed(r, rho, p, searching) + cooling = 0.0 + cs2 = _sound_speed_squared_from(P, Pc, rho, v, v_alfven) + luminosity = _luminosity(v, P, Pc, rho, (v + v_alfven) / v, p) + power_factor = _radiative_power_factor(r, luminosity, p) + gravity = _nfw_gravity(r, p) + subsonic = v ** p[TWO] <= cs2 + force = _radiative_force(r, luminosity, p, subsonic) + external = ( + (gamma - 1.0) * (cooling / rho - power_factor * v) + gamma * force * v - gravity * v + ) / (v ** p[TWO] - cs2) + cr_factor = gamma_cr * Pc * (2.0 * v + v_alfven) / (2.0 * (v + v_alfven)) + gas_factor = gamma * P - gamma_cr * Pc * (gamma - 1.0) * v_alfven * (2.0 * v + v_alfven) / ( + 2.0 * v * (v + v_alfven) + ) + if subsonic: + source_cr = -3.0 * gamma_cr * Pc / r + (gamma_cr - 1.0) * p[CR_INJECTION_RATE] / ( + v + v_alfven + ) + source_gas = ( + -(gamma - 1.0) * v_alfven * source_cr / v + + (gamma - 1.0) * p[THERMAL_INJECTION_RATE] / v + - 3.0 * gamma * P / r + + 3.0 * (gamma - 1.0) * rho * v ** p[TWO] / (2.0 * r) + ) + dv_cr = cs2 * v / (r * (cs2 - v ** p[TWO])) + v * ( + source_gas + source_cr + p[MASS_INJECTION_RATE] * v + ) / (rho * (cs2 - v ** p[TWO])) + dv_dr = dv_cr + external + log_density_gradient = 1.0 / r - dv_dr / v + dpc_dr = log_density_gradient * cr_factor + source_cr + dp_dr = log_density_gradient * gas_factor + source_gas + else: + dv_dr = 2.0 * cs2 * v / (r * (v ** p[TWO] - cs2)) + external + log_density_gradient = -2.0 / r - dv_dr / v + dpc_dr = log_density_gradient * cr_factor + dp_dr = log_density_gradient * gas_factor + dp_dr = dp_dr + (gamma - 1.0) * (rho * power_factor - rho * force - cooling / v) + out[0], out[1], out[2] = dv_dr, dp_dr, dpc_dr + + +# --- integrator.py ---------------------------------------------------------------------------- + + +@njit(cache=True) +def _cash_karp_step(h, r, y, first_slope, p, searching, new_y, error, stages, work, slope): + """integrator.cash_karp_step for the 3-component wind state; results in new_y, error.""" + for i in range(3): + stages[i, 0] = h * first_slope[i] + state = np.empty(3) + for stage in range(1, 6): + for i in range(3): + for j in range(stage): + work[j] = _STAGE_WEIGHTS[stage, j] * stages[i, j] + state[i] = y[i] + _compensated_sum(work, stage) + _derivatives(r + _STAGE_NODES[stage] * h, state[0], state[1], state[2], p, searching, slope) + for i in range(3): + stages[i, stage] = h * slope[i] + for i in range(3): + for j in range(6): + work[j] = _FOURTH_ORDER_WEIGHTS[j] * stages[i, j] + new_y[i] = y[i] + _compensated_sum(work, 6) + for j in range(6): + work[j] = _ERROR_WEIGHTS_ARRAY[j] * stages[i, j] + error[i] = _compensated_sum(work, 6) + + +@njit(cache=True) +def _relative_error(values, errors): + largest = 0.0 + for i in range(3): + if values[i] != 0.0: + ratio = abs(errors[i] / values[i]) + if ratio > largest: + largest = ratio + return largest + + +@njit(cache=True) +def _next_step_size(h, delta, tolerance, accepted, growth_when_error_is_zero, p): + if delta == 0.0: + return h * growth_when_error_is_zero + exponent = p[GROWTH_EXPONENT] if accepted else p[SHRINK_EXPONENT] + return 0.9 * h * (tolerance / delta) ** exponent + + +# --- boundary._sonic_trial -------------------------------------------------------------------- + + +@njit(cache=True) +def sonic_trial( + p, + pressure, + cr_pressure, + velocity_gradient, + central_cs2, + initial_step, + step_tolerance, + first_step_tolerance, + smallest_step_fraction, + growth_when_error_is_zero, +): + """boundary._sonic_trial: returns (stopping radius, dv/dr there, or inf if it failed).""" + starburst_radius = p[STARBURST_RADIUS] + r = 0.0 + y = np.array([0.0, pressure, cr_pressure]) + slope = np.array([velocity_gradient, 0.0, 0.0]) + h = initial_step + cs2 = central_cs2 + gradient = velocity_gradient + new_y, error, scratch_slope = np.empty(3), np.empty(3), np.empty(3) + stages, work = np.empty((3, 6)), np.empty(6) + while cs2 >= y[0] ** p[TWO] and gradient > 0: + _cash_karp_step(h, r, y, slope, p, True, new_y, error, stages, work, scratch_slope) + delta = _relative_error(new_y, error) + if h < starburst_radius * smallest_step_fraction: + return r, math.inf + if delta <= (first_step_tolerance if r == 0 else step_tolerance): + r = r + h + y[:] = new_y + cs2 = _sound_speed_squared(r, y[0], y[1], y[2], p, True) + if r + h == r and y[0] ** p[TWO] < 0.9 * cs2: + return r, math.inf + _derivatives(r, y[0], y[1], y[2], p, True, slope) + gradient = slope[0] + h = _next_step_size(h, delta, step_tolerance, True, growth_when_error_is_zero, p) + else: + h = _next_step_size(h, delta, step_tolerance, False, growth_when_error_is_zero, p) + return r, gradient diff --git a/src/outflow/boundary.py b/src/outflow/boundary.py index 065f9b4..1873970 100644 --- a/src/outflow/boundary.py +++ b/src/outflow/boundary.py @@ -19,6 +19,7 @@ (``searching=True``, the paper code's tag ``'find0'``). """ +import functools import math from dataclasses import dataclass @@ -26,7 +27,12 @@ from outflow.constants import GAMMA_CR, GAMMA_GAS, PARSEC, SOLAR_MASS from outflow.equations import WindEquations from outflow.errors import NoSolutionError -from outflow.integrator import cash_karp_step, next_step_size, relative_error +from outflow.integrator import ( + _GROWTH_WHEN_ERROR_IS_ZERO, + cash_karp_step, + next_step_size, + relative_error, +) from outflow.parameters import WindParameters # Numerical settings of the paper code, converted from its natural units (pc, Msun). @@ -50,6 +56,10 @@ # the integration restarted, accepting the first step at this looser tolerance. FIRST_STEP_FALLBACK_TOLERANCE = 100 * STEP_TOLERANCE +# Use the compiled copy of the trial integration (outflow._compiled) when Numba is installed. +# It gives the same results, about 30x faster. Set to False to run the plain Python version. +USE_COMPILED_TRIALS = True + # Marker for a trial whose step size collapsed before it became sonic or turned over (the # paper code returns the string 'RKfail'). It is treated as larger than any gradient, which is # how Python 2 compares a string with a number; see docs/open-items.md. @@ -131,6 +141,53 @@ def sonic_trial(equations: WindEquations, state: CentralState) -> SonicTrial: def _sonic_trial( equations: WindEquations, state: CentralState, first_step_tolerance: float ) -> SonicTrial: + """One trial: the compiled copy if Numba is installed and there is no cooling, else Python. + + Both give bit-for-bit the same result (tests/unit/test_compiled.py). + """ + if USE_COMPILED_TRIALS and equations.cooling_function is None: + kernel = _compiled_kernel() + if kernel is not None: + radius, gradient = kernel.sonic_trial( + kernel.pack_parameters(equations.params), + state.pressure, + state.cr_pressure, + state.velocity_gradient, + central_sound_speed_squared(state), + INITIAL_STEP, + STEP_TOLERANCE, + first_step_tolerance, + SMALLEST_STEP_FRACTION, + _GROWTH_WHEN_ERROR_IS_ZERO, + ) + return SonicTrial(radius, gradient) + return _sonic_trial_python(equations, state, first_step_tolerance) + + +def compiled_trials_available() -> bool: + """True if the central-density search runs compiled (Numba installed and switched on).""" + return USE_COMPILED_TRIALS and _compiled_kernel() is not None + + +@functools.cache +def _compiled_kernel(): + """The Numba-compiled trial integration (outflow._compiled), or None without Numba. + + Imported on first use, not when the package is imported (Numba takes a moment to load). + The first use after installing compiles it (a few seconds); the compiled code is then + cached on disk. + """ + try: + from outflow import _compiled + except ImportError: + return None + return _compiled + + +def _sonic_trial_python( + equations: WindEquations, state: CentralState, first_step_tolerance: float +) -> SonicTrial: + """The trial integration in plain Python: the readable reference for outflow._compiled.""" starburst_radius = equations.starburst_radius def derivative(r, y): diff --git a/tests/unit/test_compiled.py b/tests/unit/test_compiled.py new file mode 100644 index 0000000..a2bebd7 --- /dev/null +++ b/tests/unit/test_compiled.py @@ -0,0 +1,94 @@ +"""The Numba-compiled trial integration gives bit-for-bit the same results as the Python code.""" + +import random + +import numpy as np +import pytest + +pytest.importorskip("numba") + +from outflow import _compiled, boundary, load_paper_profile, paper_runs, solve_wind # noqa: E402 +from outflow.equations import WindEquations # noqa: E402 +from outflow.parameters import WindParameters # noqa: E402 + +ALL_RUNS = [run for figure in (6, 7, 8, 9) for run in paper_runs(figure)] + + +def test_compensated_sum_is_pythons_sum(): + rng = random.Random(3) + for _ in range(2000): + count = rng.randint(1, 6) + values = [rng.uniform(-1.0, 1.0) * 10.0 ** rng.randint(-20, 20) for _ in range(count)] + if rng.random() < 0.3: # cancellation, where compensation matters most + values.append(-sum(values[:-1]) * (1 + 1e-15)) + array = np.array(values) + assert _compiled._compensated_sum(array, len(values)) == sum(values) + + +def trial_cases(): + """(name, params, central density): the published rho_0 of every paper run and a factor of + 3 either side (breezes and winds), plus inputs that switch physics off.""" + cases = [] + for run in ALL_RUNS: + density = float(load_paper_profile(run).rho[0]) + for factor in (1 / 3, 1.0, 3.0): + cases.append((f"{run.name} x{factor:.2f}", run.params, density * factor)) + default = WindParameters.paper_default() + switched_off = { + "no cosmic rays": dict(cr_fraction=0.0), + "no radiation force": dict(opacity=0.0), + "no magnetic field": dict(magnetic_field=0.0), + "no halo": dict(halo_mass=0.0), + } + density = float(load_paper_profile(ALL_RUNS[0]).rho[0]) + for name, change in switched_off.items(): + cases.append((name, default.replace(**change), density)) + return cases + + +@pytest.mark.parametrize("name, params, density", trial_cases(), ids=lambda v: str(v)[:30]) +@pytest.mark.parametrize( + "first_step_tolerance", + [boundary.STEP_TOLERANCE, boundary.FIRST_STEP_FALLBACK_TOLERANCE], + ids=["paper tolerance", "first-step fallback"], +) +def test_compiled_trial_is_bit_identical(name, params, density, first_step_tolerance, monkeypatch): + equations = WindEquations(params) + state = boundary.central_state(params, density) + python = boundary._sonic_trial_python(equations, state, first_step_tolerance) + compiled = boundary._sonic_trial(equations, state, first_step_tolerance) + assert boundary._compiled_kernel() is not None + assert compiled == python, name + + +def test_compiled_trial_reproduces_a_failed_first_step(): + """Fig. 7, 100 uG: at this density the first step from r = 0 fails at the paper tolerance.""" + run = next(r for r in ALL_RUNS if r.name == "fig7_B0_100uG") + equations = WindEquations(run.params) + state = boundary.central_state(run.params, 5.973568084651753e-24) + python = boundary._sonic_trial_python(equations, state, boundary.STEP_TOLERANCE) + compiled = boundary._sonic_trial(equations, state, boundary.STEP_TOLERANCE) + assert compiled == python + + +def test_switch_selects_python(monkeypatch): + calls = [] + monkeypatch.setattr(boundary, "USE_COMPILED_TRIALS", False) + monkeypatch.setattr( + boundary, + "_sonic_trial_python", + lambda *args: calls.append(args) or boundary.SonicTrial(0.0, 1.0), + ) + params = WindParameters.paper_default() + boundary._sonic_trial(WindEquations(params), boundary.central_state(params, 1e-24), 1e-12) + assert len(calls) == 1 + + +@pytest.mark.slow +@pytest.mark.parametrize("run", ALL_RUNS, ids=lambda run: run.name) +def test_every_paper_run_is_bit_identical(run, monkeypatch): + compiled = solve_wind(run.params, outer_radius=run.outer_radius) + monkeypatch.setattr(boundary, "USE_COMPILED_TRIALS", False) + python = solve_wind(run.params, outer_radius=run.outer_radius) + for name in ("r", "v", "rho", "P", "Pc", "T", "vA"): + np.testing.assert_array_equal(getattr(compiled, name), getattr(python, name), name) From be444de6e07737d63d403dcedd3f881afdf82412 Mon Sep 17 00:00:00 2001 From: Brian Yu Date: Thu, 15 Oct 2026 22:55:57 +0800 Subject: [PATCH 6/7] App: with the compiled search, figure tabs solve runs in-process (Fig. 7 about 1.5 s) Starting worker processes took longer than the compiled runs; the parallel pool is kept for installs without Numba. --- app/streamlit_app.py | 33 ++++++++++++++++++++++++--------- tests/app/test_app.py | 30 +++++++++++++++++++++++++++++- 2 files changed, 53 insertions(+), 10 deletions(-) diff --git a/app/streamlit_app.py b/app/streamlit_app.py index 65666b8..fd25791 100644 --- a/app/streamlit_app.py +++ b/app/streamlit_app.py @@ -31,6 +31,7 @@ paper_runs, solve_wind, ) +from outflow.boundary import compiled_trials_available from outflow.solver import DEFAULT_OUTER_RADIUS_IN_SB TAB_NAMES = ( @@ -480,14 +481,22 @@ class FigureRunResult: def run_figure(spec: FigureSpec): - """Recompute every run of one figure, in parallel, showing "k of n finished".""" + """Recompute every run of one figure, showing "k of n finished". + + With the compiled search (Numba installed) a run takes a fraction of a second, less than + starting a worker process, so the runs are solved here one after another. Without it they + are solved in parallel processes. + """ runs = paper_runs(spec.number) - progress = st.progress(0.0, text=f"Solving the {len(runs)} runs at the same time…") + progress = st.progress(0.0, text=f"Solving the {len(runs)} runs…") def show_progress(finished: int): progress.progress(finished / len(runs), text=f"{finished} of {len(runs)} runs finished…") - outcomes = solve_runs_in_parallel(runs, show_progress) + if compiled_trials_available(): + outcomes = solve_runs_one_by_one(runs, show_progress) + else: + outcomes = solve_runs_in_parallel(runs, show_progress) progress.empty() st.session_state[f"fig{spec.number}_results"] = [ FigureRunResult(run, outcome, load_paper_profile(run)) @@ -512,11 +521,16 @@ def solve_runs_in_parallel(runs: list[PaperRun], on_finished) -> list[RunOutcome on_finished(finished) return [outcome_of(future.result) for future in futures] except (BrokenProcessPool, OSError, NotImplementedError): - outcomes = [] - for finished, run in enumerate(runs, start=1): - outcomes.append(solve_safely(lambda run=run: run.params, run.outer_radius_in_sb)) - on_finished(finished) - return outcomes + return solve_runs_one_by_one(runs, on_finished) + + +def solve_runs_one_by_one(runs: list[PaperRun], on_finished) -> list[RunOutcome]: + """Solve the runs in this process, one after another.""" + outcomes = [] + for finished, run in enumerate(runs, start=1): + outcomes.append(solve_safely(lambda run=run: run.params, run.outer_radius_in_sb)) + on_finished(finished) + return outcomes def figure_series(results: list[FigureRunResult], show_published: bool) -> list[plots.Series]: @@ -564,7 +578,8 @@ def figure_tab(spec: FigureSpec): f"Run the {len(runs)} cases of Fig. {spec.number}", type="primary", key=f"fig{spec.number}_run", - help="The cases are solved at the same time; a few seconds on a multi-core computer.", + help="About a second with Numba installed; otherwise the cases are solved in parallel " + "and take a few seconds on a multi-core computer.", ) if pressed: run_figure(spec) diff --git a/tests/app/test_app.py b/tests/app/test_app.py index 83803e3..0b5b90e 100644 --- a/tests/app/test_app.py +++ b/tests/app/test_app.py @@ -149,6 +149,32 @@ def test_fig7_run_plots_four_runs_and_four_published_overlays(): assert "published" not in kinds +@pytest.mark.slow +def test_figure_runs_in_worker_processes_without_the_compiled_search(monkeypatch): + import concurrent.futures + + import outflow + + pools = [] + real_pool = concurrent.futures.ProcessPoolExecutor + + def counting_pool(*args, **kwargs): + pools.append(kwargs.get("max_workers")) + return real_pool(*args, **kwargs) + + monkeypatch.setattr(concurrent.futures, "ProcessPoolExecutor", counting_pool) + monkeypatch.setattr(outflow.boundary, "compiled_trials_available", lambda: False) + app = start_app() + app.button(key="fig7_run").click() + app.run(timeout=300) + assert not app.exception, app.exception + assert not app.error, [e.value for e in app.error] + assert pools == [4] + traces = plotted_traces(app) + velocity = [t for t in traces if t["meta"]["panel"] == "v" and not t["meta"]["ratio"]] + assert [t["meta"]["kind"] for t in velocity].count("computed") == 4 + + def test_figure_runs_one_by_one_when_processes_are_unavailable(monkeypatch): """Some hosts forbid worker processes; the figure tab then solves the runs in turn.""" import concurrent.futures @@ -164,9 +190,11 @@ def no_processes(*args, **kwargs): def published_profile(params, outer_radius=None, cooling=False): return published[params] # stands in for the solver: instant and deterministic - # The app imports both names afresh on every run, so patch them where they come from. + # The app imports these names afresh on every run, so patch them where they come from. monkeypatch.setattr(concurrent.futures, "ProcessPoolExecutor", no_processes) monkeypatch.setattr(outflow, "solve_wind", published_profile) + # Without the compiled search the app would use worker processes. + monkeypatch.setattr(outflow.boundary, "compiled_trials_available", lambda: False) app = start_app() app.button(key="fig7_run").click() app.run() From 8f953c6c5af1074a77a70ca53cabb59e60b646d3 Mon Sep 17 00:00:00 2001 From: Brian Yu Date: Thu, 15 Oct 2026 22:55:58 +0800 Subject: [PATCH 7/7] Docs: Numba extra, compiled search kept in step with the Python reference --- CLAUDE.md | 6 +++++- README.md | 12 +++++++----- docs/architecture.md | 11 +++++++---- docs/decisions.md | 23 ++++++++++++++++++++++- docs/known-issues.md | 15 ++++++++++++--- docs/open-items.md | 6 ++---- docs/physics.md | 7 +++++++ 7 files changed, 62 insertions(+), 18 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index cf05c4f..61d3ec8 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -19,7 +19,7 @@ Read these before substantial changes: ## Environment - Conda env: `outflow` (Python 3.12) -- Install: `conda activate outflow; pip install -e ".[dev]"` +- Install: `conda activate outflow; pip install -e ".[dev,fast]"` (`fast` = Numba, optional) - Run the app: `streamlit run app/streamlit_app.py` - Tests: `pytest` (fast); `pytest -m slow` for the full paper-reproduction set - Lint: `ruff check . ; ruff format --check .` @@ -41,6 +41,10 @@ Read these before substantial changes: - A speed-up that changes results is a numerics change (the user's decision). To show a speed-up changes nothing, use `tests/unit/test_integrator.py` (bit-identical RK step) and compare all 14 paper runs before and after. +- `src/outflow/_compiled.py` is a Numba copy of the trial integration. Any change to + `physics.py`, `equations.py`, `integrator.py` or `boundary._sonic_trial_python` must be made + there too, keeping every exponent opaque (`x ** p[TWO]`, never `x**2`) and sums compensated; + `pytest tests/unit/test_compiled.py` (and `-m slow`) must stay bit-identical. - The app script keeps its `if __name__ == "__main__": main()` guard: the figure tabs' worker processes re-import it. - Cooling is off by default and is not exposed in the UI. diff --git a/README.md b/README.md index b496d80..38f3e78 100644 --- a/README.md +++ b/README.md @@ -37,10 +37,12 @@ Python 3.12. With conda: ```powershell conda create -n outflow python=3.12 conda activate outflow -pip install -e ".[dev]" +pip install -e ".[dev,fast]" ``` -Any other Python 3.12 environment works the same way with `pip install -e ".[dev]"`. +Any other Python 3.12 environment works the same way. The `fast` extra installs Numba, which +compiles the slowest part of the solver: same results to the last digit, about 30 times faster +(a run takes about 0.1 s instead of 1.4 s). Without it, everything works in plain Python. ## Quickstart @@ -57,7 +59,7 @@ profile.to_csv("wind.csv") # r, v, rho, P, Pc, T, vA with a header of the input cr_driven = WindParameters.from_astro_units(thermal_fraction=0.1, cr_fraction=0.8) ``` -The run takes a few seconds. Main entry points: +The run takes about a second (0.1 s with the `fast` extra). Main entry points: - `WindParameters` – the inputs (cgs), with `paper_default()`, `from_astro_units()` and `replace()`. - `solve_wind(params, outer_radius=None)` – returns a `WindProfile`: arrays `r, v, rho, P, Pc, @@ -78,8 +80,8 @@ streamlit run app/streamlit_app.py Tabs: **Explore** (change any input, overlay runs, download CSV), **Fig. 6–9** (recompute the paper's run sets and compare them with the published curves) and **About**. Every run is -computed when you press its button. A figure tab solves its runs at the same time, one per CPU -core, so it takes about as long as its slowest run (about 5 s on a multi-core computer). +computed when you press its button. With the `fast` extra a figure tab takes about a second; +without it, the tab solves its runs in parallel and takes about 5 s on a multi-core computer. ## How the results compare with the paper diff --git a/docs/architecture.md b/docs/architecture.md index 4429a5b..891caf0 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -38,6 +38,7 @@ paper.load_paper_profile (published runs, package data CSV) | `equations.py` | `WindEquations`: d(v, P, P_C)/dr, inner and outer equations | | `integrator.py` | Cash–Karp step, error measure and step-size rule of the paper code | | `boundary.py` | central values (eqs. 44–45) and the search for ρ₀ (last stage: Brent's method by default, or the paper code's bisection) | +| `_compiled.py` | optional Numba copy of the trial integration (physics, equations, Cash–Karp step, `_sonic_trial_python`), bit-identical, ~30× faster; loaded on first use only if Numba is installed | | `solver.py` | `solve_wind` and the outward integration, stall detection | | `profile.py` | `WindProfile`: arrays, derived quantities, interpolation, comparison, CSV read/write | | `paper.py` | the 14 published runs (`paper_runs`) and their saved profiles (`load_paper_profile`) | @@ -61,10 +62,12 @@ sit on top. Runs are computed fresh on each button press (no caching). Explore keeps up to 8 runs of the session in `st.session_state` for overlay. Cooling is not offered in the app. -A figure tab solves its runs in parallel (`solve_runs_in_parallel`: a `ProcessPoolExecutor` -with one worker per run, up to the number of CPU cores), and falls back to solving them one -after another if worker processes cannot be started. Workers re-import the app script, so its -page code runs only under `if __name__ == "__main__"`. +A figure tab solves its runs one after another in the app's own process when the compiled +search is available (a run then takes about 0.1 s, less than starting a worker process). +Without Numba it solves them in parallel (`solve_runs_in_parallel`: a `ProcessPoolExecutor` +with one worker per run, up to the number of CPU cores), and falls back to one after another +if worker processes cannot be started. Workers re-import the app script, so its page code runs +only under `if __name__ == "__main__"`. ## Other folders diff --git a/docs/decisions.md b/docs/decisions.md index e81f46f..68ae2dc 100644 --- a/docs/decisions.md +++ b/docs/decisions.md @@ -120,7 +120,28 @@ overhead. Built on branch `feature/faster-solver`: 1.7e-4 (was 2.1e-4). 70 s → 43.5 s. The paper code's search remains available as `find_central_density(..., method="bisection")`. -Overall: 95 s → 43.5 s for the 14 runs (2.2×); a default Explore run about 1.5 s. +Overall: 95 s → 43.5 s for the 14 runs (2.2×); a default Explore run about 1.4 s. + +4. **Numba-compiled trial integration** (user, 2026-10-15: "try numba now"; branch + `experiment/numba`). `src/outflow/_compiled.py` is a line-by-line compiled copy of the trial + integration of the ρ₀ search (physics, equations without cooling, Cash–Karp step, + `_sonic_trial`). The Python modules stay the readable reference. Decisions (Claude's calls): + - **Bit-identical, not just close.** Python's `x**2` calls the C library's `pow`, which + differs from the exactly rounded `x*x` in ~0.06 % of cases (`x**3`: ~26 %), and a + compiler replaces constant powers by multiplications. So every exponent is read from the + parameter array at run time (the compiler cannot fold it), and sums repeat Python 3.12's + compensated `sum()`. Result: all 14 runs bit-for-bit identical; tests compare compiled and + Python trials on 45 inputs × 2 first-step tolerances, and all 14 full runs (slow). + - **Optional:** extra `fast` (`pip install -e ".[fast]"`), in `environment.yml` and in CI; + without Numba the Python code runs. Loaded on first use, not at import. + - **Only the trials are compiled** (96 % of the time). The outward integration stays in + Python (~0.1 s per run); cooled winds use the Python code. + - **App:** with the compiled search, figure tabs solve their runs one after another in the + app's process (worker start-up would take longer than the runs); the parallel pool is + kept for installs without Numba. + Measured: 14 runs 43 s → 1.4 s; default solve 1.37 s → 0.10 s; Fig. 7 tab 1.5 s (3.6 s on + the first press after the app starts, which loads Numba). Compiling takes ~5 s once after + installing; the compiled code is cached on disk (~0.2 s to load). ## Verification diff --git a/docs/known-issues.md b/docs/known-issues.md index c79dbcb..654ddf5 100644 --- a/docs/known-issues.md +++ b/docs/known-issues.md @@ -77,10 +77,19 @@ Brent's method and the paper code's bisection, which agree on ρ₀ to 2e-12, gi differ by up to 2e-4, and why the match to the published runs cannot be much better than that. Differences of this size are far below what is visible in a plot. -## Figure tabs on machines with few cores +## The compiled search is a second copy of the physics -The figure tabs solve their runs in parallel, one process per CPU core. On a machine or host -with one or two cores (e.g. free cloud hosting) the gain is correspondingly smaller. Where +`_compiled.py` repeats the arithmetic of `physics.py`, `equations.py`, `integrator.py` and the +trial integration. A change to any of those must be made in both places; +`tests/unit/test_compiled.py` fails if the two stop giving identical results for the tested +inputs. A change that never alters a result on those inputs (for example `v*v` instead of +`v**2` in one place) can go unnoticed, so edit both copies the same way. The first run after +installing compiles the code (about 5 s); later runs load it from Numba's on-disk cache. + +## Figure tabs without Numba on machines with few cores + +Without Numba the figure tabs solve their runs in parallel, one process per CPU core; on a +machine or host with one or two cores (e.g. free cloud hosting) the gain is smaller. Where worker processes cannot be started at all, the runs are solved one after another. ## Repeated radii in the published data diff --git a/docs/open-items.md b/docs/open-items.md index f0ae29d..59fe007 100644 --- a/docs/open-items.md +++ b/docs/open-items.md @@ -61,10 +61,8 @@ here; a limit or trap we live with goes in `known-issues.md`. Each item in one p ## Later -- Numba (compiling the integrator and equations to machine code) for a further speed-up, - likely 20–50× (not measured). Costs: a new dependency, a few seconds of compiling on first - use, and the equations rewritten as plain functions, which reads less like the paper. (user, - 2026-10-15: "leave numba for later") +- Compile the outward integration too (now ~0.1 s per run in Python; only worth it if many + runs are needed, e.g. parameter scans). - Port Figs. 1–5 scripts and the §3.5 line-emission work (Figs. 10–11). - Validate the experimental cooling option against `Silich.py` (Fig. 2). diff --git a/docs/physics.md b/docs/physics.md index 64dccfc..fc61419 100644 --- a/docs/physics.md +++ b/docs/physics.md @@ -73,6 +73,8 @@ switch at r_sb (density for r < r_sb, Alfvén speed for r ≤ r_sb). The port fo | `central_state`, `CentralState` | eqs. 44–45 | P₀ and P_C,0 from ρ₀ | | `central_sound_speed_squared` | eq. 35 at r = 0 | with the L'Hôpital ratio v_A/v | | `sonic_trial`, `_sonic_trial`, `SonicTrial` | §2.3.2, §2.4 (code: `rsonic()`) | integrate from r = 0 until sonic or turning over | +| `_sonic_trial_python` | §2.3.2 (code: `rsonic()`) | the trial in plain Python, the readable reference | +| `_compiled_kernel`, `compiled_trials_available` | §2.3.2 (compiled copy of `rsonic()`) | the Numba copy in `_compiled.py`, used when installed; bit-identical | | `_bisect` | §2.3.2 (code: `gss()`) | bisection on ρ₀ | | `_brent`, `sonic_radius_mismatch` | §2.3.2 (not in the paper code; replaces the last `gss()` stage) | Brent's method on r_s(ρ₀) − r_sb (default) | | `find_central_density` | §2.3.2, §2.4 (code: `CRline()`) | ρ₀ for which the sonic radius is r_sb | @@ -84,6 +86,11 @@ there, to the same precision with fewer trials (`method="bisection"` gives the p search). Trials use the inner density and Alfvén forms at every radius (the code's `'find0'` tag). +If Numba is installed (`pip install -e ".[fast]"`), each trial runs in `_compiled.py`, a +compiled line-by-line copy of `physics.py`, `equations.py`, `integrator.py` and +`_sonic_trial_python` (without cooling). It gives bit-for-bit the same floats; read the Python +modules for the physics. `boundary.USE_COMPILED_TRIALS = False` switches it off. + ## Solver (`solver.py`, `integrator.py`) - `solve_wind` finds ρ₀ and calls `integrate_outward`, which integrates from r = 0 to the outer