MOSFiT 2.0

[guillochon]

Summary

This is MOSFiT 2.0: a packaging-and-runtime overhaul on top of current master. Development installs move to uv and pyproject.toml; the default sampler is dynesty 3.1; torch is optional; likelihoods can use a local process pool; several hot-path modules are vectorized without changing the math. The TDE viscous delay is now a jitted closed-form recurrence instead of a 1000-point trapezoid (~500× on that kernel, ~3.3× on a full PS1-10jh ln_likelihood).

Packaging and Python

• Replace setup.py / setup.cfg / MANIFEST.in as the source of truth with hatchling + pyproject.toml + uv.lock. Version is 2.0.0.
• requires-python = ">=3.11" (classifiers 3.11–3.14). Local .python-version pins 3.14.7 for uv; any 3.11+ interpreter is allowed.
• Runtime floor is driven by astrocats 0.5.0 (NumPy 2.3 / Astropy 7.1). Drop the vendored astrocats 0.3.37 wheel. astrocats>=0.5.0 resolves from PyPI.
• Optional extras: mpi (mpi4py), sedona (PyTorch, SESN SEDONA emulator only), docs (Sphinx).
• import mosfit does not import torch. SED modules lazy-load; sesn_sedona tells the user to uv sync --extra sedona.
• NumPy 2: generative dummy times (max_time default) no longer pass an empty string into linspace.
• Docker: python:3.14.7-slim, uv sync --frozen --extra mpi, no torch. Smoke-check that import mosfit does not pull in torch. .dockerignore added.
• conda-forge: hatchling sdist (no setup.py, no required PyTorch). Template in recipe/. astrocats >= 0.5.0 must land on conda-forge before the mosfit feedstock can solve (channel still has 0.3.37).

Sampling

• Default -D / Python method is dynesty (was ensembler). ensembler and ultranest remain selectable.
• Require dynesty >= 3.1 (evid_thresh / iterator attributes). Nested live points are nlive = 20 * ndim; --num-walkers does not set nlive.
• -R is remaining evidence (dlogz, default 0.02) for dynesty; PSRF still applies to ensembler.
• Nested sampling: do not rebuild sampler.results every step; disable unused bound history (save_bounds=False); snapshot Results once after baseline and after each combine_runs. CPU lulls between batches (ellipsoid rebuilds on the parent) are expected.

Parallelism and likelihood speed

• --max-cores N: local schwimmbad process pool (Windows spawn). Default 1. mpirun still takes precedence.
• Vectorize SED/photometry by unique band (blackbody, losextinction, photometry). NumPy residuals in diagonal/kernel.
• Photometry: interpolate filter curves onto sample_wavelengths once per sample grid (same integrals).
• Keep SEDs as rectangular (n_obs, n_wav) float64 arrays instead of object-dtype row lists (no np.stack round trip through photometry). Isolated blackbody + photometry-like path is 1.85×.
• TDE Fallback: cache Tout coefficients, np.interp instead of interp1d, NumPy NaN/array assembly. Golden draws match the previous engine.
• Viscous delay: replace the 1000-point log-spaced trapezoid with a Numba piecewise-linear exponential recurrence (exact integral of the same interpolant). Isolated kernel ~526×; PS1-10jh ln_likelihood ~3.3×.

Viscous exponential recurrence

The old viscous transform spent most of a TDE likelihood on a (n_times × ~1000) interp + exp + trapezoid. L(t) was already piecewise linear on the dense grid; the convolution (1/τ) ∫ L(s) exp((s−t)/τ) ds has a closed form on each segment. That walk is JIT-compiled.

PS1-10jh, x = 0.5, same machine. Kernel times are 40 repeats after warmup; TDE times are 20 repeats after 3 warmups (post-rectangular-SED baseline 10.035 ms).

---
config:
  xyChart:
    width: 700
    height: 320
    showLegend: true
  themeVariables:
    xyChart:
      plotColorPalette: "#8b949e, #3fb950"
---
xychart-beta
    title "PS1-10jh TDE ln_likelihood (ms)"
    x-axis ["1000-pt trapz", "Jitted recurrence"]
    y-axis "Mean wall time (ms)" 0 --> 12
    bar "Wall time" [10.04, 3.02]

Loading

---
config:
  xyChart:
    width: 700
    height: 260
    showLegend: false
  themeVariables:
    xyChart:
      plotColorPalette: "#4493f8"
---
xychart-beta horizontal
    title "Speedup vs 1000-point viscous quadrature"
    x-axis ["Viscous kernel", "TDE ln_likelihood"]
    y-axis "Speedup" 0 --> 550
    bar [526, 3.32]

Loading

Path	Before	After	Speedup
Viscous kernel (PS1-10jh arrays)	7.22 ms	0.014 ms	526×
PS1-10jh TDE ln_likelihood	10.04 ms	3.02 ms	3.32×

Physics vs the old trapezoid (same piecewise-linear L):

• Isolated recurrence vs 1000-pt trapz: max relative 6.5×10⁻⁵ (τ = 3.5); ≤ 3.7×10⁻⁴ for τ from 10⁻³ to 100.
• Constant L matches L(1 − e^{−t/τ}) to 1e-12.
• Restoring the trapz on PS1-10jh recovers lnL −1032.9597765404137 exactly. The recurrence is −1032.7869483449351 (Δ +0.173). Model magnitudes: max 0.83 mmag, median 0.28 mmag.

The millimag shift is the old rule not being fully converged, not a different delay kernel.

Rectangular SED arrays

Wall-clock speedup from keeping SEDs as float64 (n_obs, n_wav) arrays. Same machine, consecutive runs with the patch stashed then restored. Synthetic path: 400 observations × 25 wavelengths, 6 bands; 5 warmup + 40 timed calls. TDE ln_likelihood in this table is before the viscous recurrence (viscous still dominated).

---
config:
  xyChart:
    width: 700
    height: 360
    showLegend: true
  themeVariables:
    xyChart:
      plotColorPalette: "#8b949e, #3fb950"
---
xychart-beta
    title "Mean wall time, synthetic SED path (us)"
    x-axis ["Blackbody process", "BB + photometry-like", "Photometry consume"]
    y-axis "Mean wall time (us)" 0 --> 700
    bar "Object-dtype lists (before)" [329, 599, 270]
    bar "Rectangular float64 (after)" [237, 323, 86]

Loading

---
config:
  xyChart:
    width: 700
    height: 280
    showLegend: false
  themeVariables:
    xyChart:
      plotColorPalette: "#4493f8"
---
xychart-beta horizontal
    title "SED-path speedup vs object-dtype lists (before / after)"
    x-axis ["Blackbody process", "BB + photometry-like", "Photometry consume", "TDE lnL (pre-viscous)"]
    y-axis "Speedup" 0 --> 3.5
    bar [1.39, 1.85, 3.12, 0.99]

Loading

Path	Before	After	Speedup
Blackbody process	329.3 µs	236.9 µs	1.39×
Blackbody + photometry-like	599.5 µs	323.4 µs	1.85×
Photometry-like consume only (np.stack vs 2-D slice)	270.2 µs	86.5 µs	3.12×
PS1-10jh TDE ln_likelihood (before viscous recurrence)	9.958 ms	10.035 ms	~1.00×

The 3.12× gain is the restack this change removes. That was a few hundred microseconds inside a ~10 ms TDE likelihood, so it did not move the full-model number until viscous was replaced.

Tests, CI, and coverage

GitHub Actions (.github/workflows/tests.yml) runs uv sync --frozen then uv run python mosfit/tests/run_all.py on Ubuntu 3.11 / 3.14 and Windows 3.14, plus a job that pip-installs the hatchling sdist (conda-forge shape, no uv.lock). That runner discovers every mosfit/tests/_test_*.py plus the repo-root generative smoke test.py, collects line coverage, and fails if coverage drops below 70%.

Local line coverage of the package (omitting tests and the optional torch SEDONA extra):

• Previous three scripts only (_test_speedup_numerics.py, _test_pool_likelihood.py, test.py): 47%
• Current suite: 75%

The obsolete CLI script test.sh (catalog-by-name fetch, --offline, -F covariance, upload) is removed.

Scripts in the suite:

• mosfit/tests/_test_speedup_numerics.py: blackbody, photometry, diagonal, mm83, fallback goldens, viscous recurrence vs old quadrature; import mosfit must not import torch
• mosfit/tests/_test_pool_likelihood.py: serial vs 10-core TDE ln_likelihood on a fixed vector
• mosfit/tests/_test_lsq12dlf.py: load mosfit/tests/LSQ12dlf.json and evaluate one SLSN likelihood
• mosfit/tests/_test_models_likelihood.py: load every packaged model except sesn_sedona and evaluate one likelihood
• mosfit/tests/_test_ensembler.py: short -D ensembler smoke on dummy exppow data
• mosfit/tests/_test_modules.py: direct process() coverage for engines, parameters, SEDs, transforms, outputs
• mosfit/tests/_test_unit.py: utils, plotting helpers, printer, fetcher (local JSON only), ASCII converter, CLI/get_parser
• mosfit/tests/_test_ptsampler.py: vendored PTSampler / MOSSampler
• test.py: generative fit_events for magni and slsn (iterations=1)

Fixtures: mosfit/tests/PS1-10jh.json (Open TDE Catalog) and mosfit/tests/LSQ12dlf.json.

Documentation

README and Sphinx updated for Python 3.11+, uv, extras, dynesty default, --max-cores, -R/dlogz, SEDONA extra, Docker, the Python Fitter API, and MOSFiT 2.0. Autodoc skips torch-backed SEDONA classes so RTD can build without the sedona extra. Release notes: CHANGELOG.md (stage under Unreleased, move under a version heading when tagging). conda-forge notes: recipe/.

Breaking / behavioral notes

• This is MOSFiT 2.0.0. Python < 3.11 is unsupported.
• Default sampler is dynesty, not ensemble MCMC. Old scripts that assumed emcee defaults should pass -D ensembler.
• Torch is not a core dependency. sesn_sedona needs uv sync --extra sedona or pip install 'mosfit[sedona]'.
• Event data remains local files only (no catalog download by name); that predates most of this branch but docs now match it.
• Published conda-forge / PyPI wheels may lag this branch until a release. conda-forge also needs astrocats 0.5.0 on that channel first.

Test plan

• uv run python mosfit/tests/run_all.py (full suite + coverage; local result 75%, fail-under 70%)
• GitHub Actions Tests workflow green on Ubuntu 3.11 / 3.14 and Windows 3.14
• uv sync on Python 3.11 and 3.14; import mosfit must succeed and must not import torch
• Short dynesty smoke: uv run mosfit -e mosfit/tests/PS1-10jh.json -m tde --max-cores 4 --no-fracking -i 1
• uv sync --extra sedona then load/import the SESN SEDONA emulator
• uv sync --extra docs and sphinx-build docs docs/_build/html
• Docker build from repo root; container import mosfit without torch
• Windows: --max-cores 2 on a small local JSON fit (spawn)

[guillochon]

This should knock out #246 and supercede #247 .

[Copilot]

🟡 Changes recommended

There are a few confirmed import/runtime hazards (package-dir writes in photometry, contradictory torch mocking in Sphinx config, and brittle astrocats version parsing) that can break installs/docs or import mosfit in real environments.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

MOSFiT 2.0 packaging/runtime overhaul moving the project to pyproject.toml/hatchling + uv, switching the default sampler to dynesty, adding a local process pool for likelihood evaluation, and landing multiple performance/vectorization improvements (including a jitted viscous-delay recurrence) alongside a substantially expanded test/CI setup.

Changes:

• Replace legacy packaging/test scripts with hatchling + pyproject.toml, uv.lock, updated Docker/RTD/conda templates, and a new changelog.
• Switch default sampling to dynesty (>=3.1) and add --max-cores local process-pool parallelism.
• Vectorize/optimize hot-path modules (SED/photometry, fallback, viscous transform) and add a coverage-enforced test runner + expanded test suite.

File summaries

File	Description
test.sh	Removed legacy CLI test script.
setup.py	Removed setuptools build script.
setup.cfg	Removed legacy setuptools metadata.
MANIFEST.in	Removed legacy sdist manifest.
.coveragerc	Removed legacy coverage config (moved to pyproject).
README.md	Updated install/use docs for uv/extras/dynesty/cores.
pyproject.toml	Added hatchling project config, deps, coverage settings.
recipe/README.md	Added conda-forge feedstock update notes/template docs.
recipe/meta.yaml	Added conda recipe template for 2.0+.
mosfit/utils.py	NumPy 2 cleanups + safer open_atomic dir handling.
mosfit/tests/run_all.py	Added coverage-based test runner used by CI.
mosfit/tests/PTF10hgi.txt	Removed obsolete fixture used by old test.sh.
mosfit/tests/event_list.txt	Removed obsolete fixture used by old test.sh.
mosfit/tests/dummies.py	Added shared dummy objects for tests.
mosfit/tests/_test_unit.py	Added broad unit/regression coverage for utilities/CLI.
mosfit/tests/_test_speedup_numerics.py	Added numeric goldens + “no torch import” checks.
mosfit/tests/_test_ptsampler.py	Added PTSampler/MOSSampler coverage.
mosfit/tests/_test_pool_likelihood.py	Added serial vs pool likelihood consistency test.
mosfit/tests/_test_modules.py	Added direct process() coverage for many modules.
mosfit/tests/_test_models_likelihood.py	Added “load all models and eval one lnL” smoke.
mosfit/tests/_test_lsq12dlf.py	Added LSQ12dlf fixture likelihood smoke.
mosfit/tests/_test_ensembler.py	Added short emcee/ensembler smoke.
mosfit/tests/init.py	Marked tests as a package for helpers.
mosfit/strings.json	Updated CLI help strings (dynesty default, max-cores).
mosfit/samplers/nester.py	Dynesty 3.1+ integration + results snapshot optimization.
mosfit/samplers/ensembler.py	Removed deprecated np.int usage.
mosfit/requirements.txt	Deprecated file; points users to pyproject.
mosfit/modules/transforms/viscous.py	Replaced trapz quadrature with numba-jitted recurrence.
mosfit/modules/transforms/diffusion.py	Switched trapz → trapezoid (NumPy 2).
mosfit/modules/transforms/diffusion_csm.py	Switched trapz → trapezoid (NumPy 2).
mosfit/modules/transforms/diffusion_aspherical.py	Switched trapz → trapezoid (NumPy 2).
mosfit/modules/seds/synchrotron.py	Updated to rectangular SED allocation/fill.
mosfit/modules/seds/sesn_sedona.py	Made torch optional with clear extra-install error.
mosfit/modules/seds/sed.py	Added rectangular SED helpers + sum logic.
mosfit/modules/seds/losextinction.py	Vectorized extinction and rectangular SED handling.
mosfit/modules/seds/line.py	Vectorized rectangular SED line profile application.
mosfit/modules/seds/blackbody.py	Vectorized unique-band blackbody evaluation.
mosfit/modules/seds/blackbody_supressed.py	Converted to rectangular SEDs + NumPy 2 updates.
mosfit/modules/seds/blackbody_cutoff.py	Converted to rectangular SEDs.
mosfit/modules/seds/init.py	Implemented lazy-loading SED exports to avoid torch import.
mosfit/modules/observables/photometry.py	Filter-cache interpolation + rectangular SED consumption.
mosfit/modules/engines/fallback.py	NumPy-ized fallback + removed SciPy interp1d dependency.
mosfit/modules/energetics/bns_ejecta.py	Switched trapz → trapezoid.
mosfit/modules/energetics/bns_ejecta_generative.py	Switched trapz → trapezoid.
mosfit/modules/arrays/kernel.py	Vectorized observation array assembly.
mosfit/modules/arrays/diagonal.py	Vectorized residuals/errors; robust None→NaN handling.
mosfit/modules/init.py	Avoid eager SED star-import; add lazy SED getattr/dir.
mosfit/model.py	Fixed dynamic loader module qualname (Numba cache importability).
mosfit/main.py	Added --max-cores CLI flag; dynesty default.
mosfit/fitter.py	Added local process-pool attach + dynesty defaults + NumPy2 fix.
mosfit/emulators/sesn_sedona/profiles.py	Made torch optional with clear extra-install error.
mosfit/init.py	Version bump to 2.0.0 + astrocats floor check update.
LICENSE	Updated copyright year.
docs/requirements.txt	Deprecated; points to pyproject docs extra.
docs/package.rst	Added Fitter API docs including max_cores.
docs/models.rst	Documented sedona extra + torch import behavior.
docs/installation.rst	Updated install guide for uv/extras/conda/Docker.
docs/index.rst	Updated docs landing page for MOSFiT 2.0.
docs/help.rst	Updated FAQ for Python version, torch, max-cores.
docs/getting_started.rst	Updated getting started for dynesty + max-cores.
docs/fitting.rst	Updated sampler docs for dynesty default and semantics.
docs/assessing.rst	Clarified -R semantics by sampler.
docs/conf.py	Adjusted autodoc behavior to skip sesn_sedona docs import.
Dockerfile	Rebased on Python 3.14.7 + uv sync; assert no torch installed/imported.
CHANGELOG.md	Added Keep-a-Changelog formatted changelog.
.readthedocs.yaml	Updated RTD build to use pip install with docs extra.
.python-version	Added uv-local CPython pin.
.gitignore	Added uv/coverage ignores + root-file negations.
.github/workflows/tests.yml	Added uv-based CI tests + sdist install job.
.github/workflows/docker-publish.yml	Updated docker workflow paths/notes for new packaging.
.dockerignore	Added docker context exclusions.

Review details

Suppressed comments (1)

docs/conf.py:92

• autodoc_mock_imports still mocks torch, but the comment above says not to stub torch in sys.modules because it breaks Astropy. Mocking it here defeats the SESN SEDONA __all__ pruning and can reintroduce the autodoc import-time failure.

autodoc_mock_imports = [
    'torch',
    'mpi4py',
    'schwimmbad.mpi',
]

• Files reviewed: 72/74 changed files
• Comments generated: 3
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

🔵 Needs a closer look

Two moderate process-pool issues must be fixed before approval.

Review details

Suppressed comments (2)

Previously missed (2) — in code that hasn't changed since the last review.

mosfit/fitter.py:245

• fit_events() resets _max_cores to 1 whenever this argument is omitted, so the documented Fitter(max_cores=10).fit_events(...) API always runs serially. Preserve the constructor value unless the caller explicitly overrides it here.
  mosfit/model.py:362
• Giving every CWD override the installed-package name breaks the new Windows/spawn pool path. During worker startup, model_obj is unpickled before _init_likelihood_worker runs, so Python imports this name from mosfit/modules/...: a new custom leaf raises ModuleNotFoundError, while an override of a built-in leaf can silently resolve to the packaged implementation instead. Local custom modules therefore cannot safely use --max-cores; make their source importable in children or rebuild/load the model inside the worker before resolving its classes.

• Files reviewed: 72/74 changed files
• Comments generated: 0 new
• Review effort level: Balanced

[guillochon]

Addressed the suppressed Copilot notes from the latest review in 68c8b1c:

• fit_events max_cores: default is now None, so Fitter(max_cores=10).fit_events(...) keeps the constructor pool size unless the caller passes an explicit override.
• CWD module overrides: custom leaves load as modules.. with the process CWD on sys.path (and the same path is inserted before spawning workers). Packaged modules stay mosfit.modules.. so Numba cache and Windows spawn still import them.