Refactor: isolate a standalone forward kernel — 2 repos vs. 1 repo with 2 install paths?

[joglekara]

Context

The package has grown to mix three concerns that have different audiences and dependency footprints:

1. Forward modeling — turning plasma parameters into a Thomson spectrum (physics → synthetic diagnostic).
2. Fitting / inverse — optimization loops, loss functions, data preprocessing, postprocessing.
3. Orchestration — deck/config processing, MLflow runners, plotting.

Two external pulls motivate cleaning this up:

• A 3rd-party request to use the forward solve standalone and let the user decide how the fit is performed (bring-your-own optimizer). This is the strongest argument for isolating a clean forward kernel.
• An internal preference (from the primary user) to decompose into well-defined modules: deck/runners, postprocessor, inverse, data handler, interactive runner, forward module, synthetic diagnostic, and physics.

This issue records an analysis of the current structure and proposes a path. It ends with one decision to make (see bottom).

Findings (from the import graph)

The codebase (~9k LOC) is already close to a clean layered shape. Two facts dominate the design:

1. core/ is already a clean, IO-free forward kernel ✅

Everything under core/ (physics, EDF/params, synthetic diagnostic) imports only jax, scipy, numpy, interpax, equinox, plus one internal helper (utils.vector_tools). It pulls in none of the heavy/dirty deps — no mlflow, optax, opencv, skimage, boto3, pandas, xlrd. Those are confined to the fitting/data/IO/plotting layers.

This is exactly the boundary the 3rd party wants: import tsadar.core... gives a differentiable forward model with light deps and no opinion about how you fit.

2. Three couplings are the actual "muddle" 🔧

These must be resolved for any split (separate packages cannot have circular dependencies) — and they're worth fixing even if we keep one package:

• inverse ↔ utils/process (true cycle): inverse/fitter.py imports process.postprocess; process/postprocess.py imports back into inverse.loss_function and inverse.loops.
• utils/data_handling ↔ utils/process (true cycle): data_handling/load_ts_data.py → process.warpcorr, while process/evaluate_background.py and process/prepare.py → data_handling.
• forward → inverse (wrong-direction edge): forward/calc_vs_data.py (the "interactive runner") imports inverse.fitter and inverse.loss_function. Forward should not depend on inverse.

Recommendation

Modules ≠ repos. The 8-module decomposition is an excellent package decomposition. Splitting it into 8 repos would multiply versioning/CI/release/pinning overhead 8× for a ~9k-LOC codebase whose layers co-evolve — and it isn't even possible today because of the two dependency cycles above.

What we actually want — "use the forward solve standalone, bring your own fit" — is a dependency-isolation problem, solved by a clean layered DAG plus dependency gating.

Target layering (strict acyclic, bottom → top)

The 8 desired modules map cleanly onto this:

utils            (vector_tools, misc)                    ← leaf
physics          (form_factor, generate_spectra, chi)    "Physics"          ─┐
params/edf       (ts_params, distribution_functions)     "EDF"              │ LIGHT-DEP
diagnostic       (thomson_diagnostic, irf)               "synthetic diag"   ─┘ KERNEL
────────────────────────────────────────────────────────────────────────────
data             (load, calibration, lineouts, bg,       "data handler"
                  throughput, warp, feature, visualizer)
plotting         (plotters, lineout_plot, edf_movie)     "postprocessor" primitives
forward          (calc_series = forward + series runner) "forward module"
inverse          (loss, loops, fitter, postprocess)      "inverse"
interactive      (calc_vs_data)                          "interactive runner"
runner/deck      (config, mlflow orchestration)          "tsadar"

The kernel (utils → physics → params/edf → diagnostic) is the standalone forward solver. Everything above it is fitting/data/orchestration.

Prerequisite refactors (needed regardless of repo count)

• Move postprocess.py into inverse/ — it's post-fit processing. Turns the inverse↔process cycle into an intra-package call. (There's already a stray empty tsadar/inverse/process/ dir suggesting this was started.)
• Merge data_handling + process into one data package — they're mutually recursive because the split is artificial; preprocessing is data handling.
• Relocate calc_vs_data (interactive runner) beside inverse (or push shared helpers down: load_data_for_fitting → data, keep LossFunction in inverse) so forward no longer imports inverse.
• Note: "postprocessor = plotting" conflates two things. Plotting primitives are a clean leaf; the post-fit orchestration that calls them belongs in inverse. Keep them separate.

Decision to make 👇

Both options give the 3rd party a standalone, light-dep forward solver and let the user own the fit. They differ only in packaging. The prerequisite refactors above are required either way.

Option A — One repo, two install profiles (extras):

• pip install tsadar → kernel only (jax/equinox/interpax/scipy). Drops into a differentiable solver; nothing heavy loads.
• pip install tsadar[fit] → adds optax/jaxopt/mlflow/pandas/opencv for the full pipeline.
• ✅ Lowest overhead, single CI/version, boundaries still enforced as packages. Cross-cutting changes stay in one PR.
• ⚠️ Forward kernel can't release on an independent cadence; one issue tracker for both audiences.

Option B — Two repos (tsadar-core forward + tsadar fitting app):

• tsadar-core = the kernel; tsadar depends on it.
• ✅ Hard wall, independent release cadence, clean story for forward-only consumers.
• ⚠️ 2× CI/versioning; kernel changes that the app needs require a release + bump + coordinated PRs; harder local dev loop.

Recommendation: start with Option A (do the refactors, ship extras), and promote the kernel to its own repo (Option B) only if/when it needs an independent release cadence. A is reversible into B cheaply once the package boundaries are clean; the reverse is not.

Which do we want — A or B?

---

🤖 Generated with Claude Code

[joglekara]

cc @almilder — this is the refactor we discussed. Your call on the A vs B decision at the bottom (1 repo + install extras vs. 2 repos). I'm opening a PR now for the prerequisite refactors (break the 3 couplings), which are required either way and don't depend on the A/B choice.

[almilder]

I think the single repo+ extras is a good way to go. But i think the most important goal of this decomposition is to isolate the diagnostic specifics (calibrations, data loading, k-smearing) that are unique to OMEGA data, and the post-processor from the rest of the code base. The idea is to have the interface reduced to the point that other users can add on their own diagnostic specifics (e.g. NIF OTS) without having to navigate all the jax and AD logic. This is often requested directly or in the context of supplying preprocessed data instead of raw OMEGA files. Breaking the post-processor out will allow users to write their own visualization if desired or plug into multi-diagnostic visualizers.

[joglekara]

Decision recorded: Option A — single repo + install extras. Thanks @almilder.

Reframing the work around your stated #1 goal (isolate the OMEGA diagnostic specifics + post-processor behind a minimal interface, so others can add a diagnostic like NIF OTS / supply preprocessed data / swap visualization, without touching the jax/AD core). The packaging (extras) is secondary to that. Here's the sequenced roadmap:

1. Break the 3 couplings — PR Refactor: break the 3 module couplings blocking a clean forward/inverse split #106 (open, prerequisite for everything below)
   Merge data_handling+process → tsadar/data/; move postprocess → inverse/; move calc_vs_data → inverse/. Pure behavior-preserving moves.

2. Diagnostic data-provider interface — PR Add a diagnostic data-provider interface (isolate OMEGA specifics) #107 (open, stacked on Refactor: break the 3 module couplings blocking a clean forward/inverse split #106) ← your ported some code to jax #1 goal
   The pipeline already consumes only a plain (all_data, sa, all_axes) tuple; all OMEGA-specificity is upstream of that seam. PR Add a diagnostic data-provider interface (isolate OMEGA specifics) #107 turns it into a registry: register_data_provider("nif_ots") + config["data"]["provider"] (default "omega"). New diagnostics need zero jax/AD knowledge.

3. Post-processor extraction (next)
   Break the post-processor into its own pluggable component (config-selected, decoupled from inverse internals) so users can write their own visualization or feed a multi-diagnostic visualizer.

4. k-smearing into the diagnostic layer (next)
   Currently a dead inline ksmear flag in calc_vs_data; make it a first-class instrument option in the synthetic-diagnostic model.

5. Option A install extras (last)
   pip install tsadar (light jax/AD kernel) + tsadar[fit], tsadar[omega], tsadar[viz]. ⚠️ Not just a pyproject change: tsadar/__init__.py eagerly imports runner → fitter → optax/mlflow, so a core-only install needs __init__ import-hygiene first.

One thing I need from you (@almilder): for supplying preprocessed data, #107 already lets you register a trivial provider returning an in-memory (all_data, sa, all_axes). Do you also want a built-in "preprocessed" provider that reads from a file — and if so, what format (.npz / xarray+netCDF / in-memory only)? I left it out rather than guess.

[almilder]

Recent PRs have broken the couplings between inverse, forward, and postprocess. Data interfacing will be done in the future when the need arises.