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/CLAUDE.md b/CLAUDE.md index 3ec2eda..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 .` @@ -36,7 +36,17 @@ 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. +- `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. ## Conventions diff --git a/README.md b/README.md index 92ac702..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,7 +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 takes 20–40 s. +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 @@ -87,6 +90,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/app/streamlit_app.py b/app/streamlit_app.py index 90ff47e..fd25791 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 @@ -28,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 = ( @@ -60,10 +64,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 +481,56 @@ 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, 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"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…") + + def show_progress(finished: int): + progress.progress(finished / len(runs), text=f"{finished} of {len(runs)} runs finished…") + + 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"] = 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): + 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]: @@ -526,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="Each case takes a few seconds; the whole figure about half a minute.", + 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) @@ -636,4 +689,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/docs/architecture.md b/docs/architecture.md index 51a7e88..891caf0 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -37,7 +37,8 @@ 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) | +| `_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,6 +62,13 @@ 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 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 | Path | Role | diff --git a/docs/decisions.md b/docs/decisions.md index 9129c6c..68ae2dc 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,55 @@ 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.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 - "Same results" = differences not visible by eye. Test threshold: every reference run matches @@ -158,7 +211,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..654ddf5 100644 --- a/docs/known-issues.md +++ b/docs/known-issues.md @@ -69,6 +69,29 @@ 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. + +## The compiled search is a second copy of the physics + +`_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 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..59fe007 100644 --- a/docs/open-items.md +++ b/docs/open-items.md @@ -61,6 +61,9 @@ here; a limit or trap we live with goes in `known-issues.md`. Each item in one p ## 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). - Possible erratum on Ė in Table 1 (user's decision). diff --git a/docs/physics.md b/docs/physics.md index 6a49ff7..fc61419 100644 --- a/docs/physics.md +++ b/docs/physics.md @@ -73,12 +73,23 @@ 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 | 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). + +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`) 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 bdc9c8f..1873970 100644 --- a/src/outflow/boundary.py +++ b/src/outflow/boundary.py @@ -10,13 +10,16 @@ (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'``). """ +import functools import math from dataclasses import dataclass @@ -24,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). @@ -36,6 +44,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 @@ -44,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. @@ -125,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): @@ -181,12 +244,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 +309,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/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/app/test_app.py b/tests/app/test_app.py index 1319efc..0b5b90e 100644 --- a/tests/app/test_app.py +++ b/tests/app/test_app.py @@ -149,6 +149,62 @@ 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 + + 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 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() + 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] 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). 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) 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", [