flepimop2-op_engine: options-surface parity with flepimop2-op_system + jax extra

[jc-macdonald]

Summary

Bring flepimop2-op_engine's options surface up to parity with flepimop2-op_system's, and add a jax optional extra, so op_engine's new diffrax/lineax-backed CoreSolver methods (#67) and sparse/dense implicit-solve dispatch (#69) can actually be wired up through the flepimop2 provider the same way COVID19_USA's diffrax_engine.py is wired up today.

Current state (confirmed by direct comparison)

flepimop2-op_system's OpSystemSystem (flepimop2-op_system/src/flepimop2/system/op_system/__init__.py) compiles the spec once and publishes a rich self.options surface — 23 keys, including pytree_stepper_fn, block_pytree_stepper_fn, history_stepper_fn, body_eval_fn, block_history_stepper_fn, block_body_eval_fn, factorize_axes, block_axes, template_shapes — read via stepper.option(name, default). This is exactly the contract COVID19_USA's diffrax_engine.py depends on.

flepimop2-op_engine's OpEngineFlepimop2Engine (flepimop2-op_engine/src/flepimop2/engine/op_engine/__init__.py) is an EngineABC (consumer role, not a SystemABC) and today only ever reads 4 option keys off whatever system it's given — operators, jacobian, operator_axis, mixing_kernels — and only ever calls the plain bound stepper() (flat/scalar eval_fn path). It never touches pytree_stepper_fn, history_stepper_fn, body_eval_fn, or any block-axis variant. If op_engine were pointed at an op_system-compiled spec today, it would silently use only the scalar path and none of the vectorized/history machinery.

Structural drift beyond scope, noted for awareness (not necessarily to fix in this issue): the two adapters declare their module identity differently (OpSystemSystem's module="..." class-declaration kwarg vs OpEngineFlepimop2Engine's module: Literal[...] Pydantic field), and their pyproject.tomls have diverged in dependency style (op-system>=0.2.0 registry pin vs op-engine @ git+...@main direct git pin) and version (0.2.0 vs 0.1.0).

Scope

• Extend OpEngineFlepimop2Engine to read and use the op_system-style vectorized/history option keys when present and when the configured CoreSolver method supports them (this issue should land after or alongside Add diffrax/lineax-backed adaptive methods to op_engine's CoreSolver, ported from COVID19_USA's diffrax_engine.py #67/Multi-backend implicit solve: Array-API dense fallback + SciPy/CuPy sparse registry #69, since the SciPy explicit/dense-implicit CoreSolver path doesn't need most of these — the point is enabling the new diffrax-backed methods to consume them the way diffrax_engine.py does).
• Add [project.optional-dependencies] jax = [...] to flepimop2-op_engine/pyproject.toml, mirroring flepimop2-op_system's jax/jax-inference extras, once op_engine itself has a jax extra (see Add diffrax/lineax-backed adaptive methods to op_engine's CoreSolver, ported from COVID19_USA's diffrax_engine.py #67).
• Evaluate (don't necessarily fix wholesale) the op-engine dependency pin style — if/when op_engine starts publishing releases the way op_system now does (see op_system's recent 0.2.0 release), switch from the git-ref pin to a registry version pin for consistency.

Acceptance criteria

• flepimop2-op_engine can drive a diffrax-backed CoreSolver solve using pytree_stepper_fn/history_stepper_fn/body_eval_fn/block-axis options from an op_system-compiled spec, exercised by a test analogous to COVID19_USA's own coverage.
• Existing SciPy-CoreSolver-backed behavior is unchanged (this is additive, not a replacement of the existing 4-key read path).
• just quality and just test pass for both provider packages.

Relationship to other issues

Depends on #67 (needs the diffrax-backed methods to exist before the provider can wire them up) and #69 (needs the sparse/dense implicit-solve dispatch for full method coverage) — can be scoped/designed in parallel but the "use them" acceptance criterion is blocked on both landing. See #26 for the overall tracking issue.

[jc-macdonald]

Design revision (2026-08-18)

Light correction to framing only — this issue's actual scope (bringing flepimop2-op_engine's options surface up to parity with flepimop2-op_system's) is unchanged. What's corrected: references to "op_engine's forthcoming JAX engine" now read as "op_engine's new diffrax/lineax-backed CoreSolver methods" — per #67's revision, these land as additional methods on the existing CoreSolver, not a separate engine class. Also now depends on #69 (sparse/dense implicit-solve dispatch) in addition to #67, since full method coverage includes the dense-fallback and sparse-registry-backed implicit paths, not only the diffrax methods. See #26 for the full revised plan.