Skip to content

Infer array namespace (xp) from input in ModelCore/CoreSolver dense paths #66

Description

@jc-macdonald

Summary

Make ModelCore and CoreSolver's dense/elementwise code paths array-namespace-polymorphic, parallel to op_system's namespace-from-input pattern (op_system PR #101/#102: compiled eval_fn infers its array module from the input via y.__array_namespace__() instead of a compile-time xp/backend selection).

Why this is in scope now, narrowly

An earlier plan (#26 and its since-closed children #19/#21/#22/#23/#25) tried to solve backend-swappability with an injected LinearSolverBackend Protocol object covering all of CoreSolver, including the SciPy-sparse implicit-solve machinery. That's been dropped as premature (see #26). This issue is the piece of that plan that's actually cheap and immediately useful: most of op_engine's code has no SciPy dependency at all today and is already Array-API-shaped.

  • src/op_engine/model_core.py: zero scipy import; every array op (np.asarray, np.zeros, np.diff, .reshape, .mean(), etc.) is elementwise/broadcast/reduction. It already has a dead stub anticipating this: ModelCoreOptions.xp: object = np (docstring: "forward-compatibility hook for GPU backends") and an unused class ArrayBackend(Protocol) — stored on self.xp in __init__ but never read again.
  • src/op_engine/core_solver.py: the stepping/error-control logic (_error_norm, _propose_dt, _step_euler_doubling, _step_heun, _advance_adaptive_to_time, run()) uses only np.abs/add/allclose/asarray/copyto/divide/eye/isfinite/maximum/mean/multiply/sqrt/subtract/zeros/zeros_like — no np.linalg, no np.kron, no SciPy. The SciPy-specific boundary is already isolated to a single chokepoint, CoreSolver._as_scipy_operator() (its own docstring already calls it "the single boundary where backend-specific operator requirements are enforced"), used only by the implicit/linearized methods.

Scope

  • Replace the dead ModelCoreOptions.xp = np / self.xp pattern with real inference: derive xp from the input state array (y.__array_namespace__()) at the point each dense operation runs, rather than storing a module reference at construction time.
  • Apply the same inference to CoreSolver's explicit-method stepping/error-control code (euler, heun, and the doubling/error-norm helpers) so a JAX array passed as initial state flows through those methods without modification.
  • Remove (or repurpose) the unused ArrayBackend Protocol stub in model_core.py if inference makes it redundant — don't leave dead abstraction next to live code.
  • Explicit non-goal: matrix_ops.py's SciPy-sparse implicit-solve/factorization-caching machinery, and CoreSolver's implicit/linearized methods (implicit-euler, trapezoidal, bdf2, ros2, imex-*) that route through _as_scipy_operator. Those stay SciPy-first; see the new JAX-engine tracking issue for how that gap actually gets closed (not via namespace polymorphism — SciPy's LU/factorized isn't JAX-traceable regardless of how the input array is inferred).

Acceptance criteria

  • CoreSolver with NumPy arrays and the euler/heun methods behaves identically to today (bit-for-bit where currently deterministic).
  • Passing JAX arrays as initial state through ModelCore and CoreSolver.run() with an explicit (non-implicit) method works without a compile-time backend selection, mirroring how op_system.compile_spec no longer needs xp=/backend=.
  • New tests assert namespace-polymorphism: same input (as NumPy and as JAX arrays) produces matching results through the affected code paths.
  • just quality (ruff + mypy --strict) and just test pass for both op_engine and flepimop2-op_engine.

Relationship to other issues

Supersedes the "inject backend into CoreSolver"-shaped work previously tracked in #21/#22 for the dense paths only — see #26 for the full revised plan. Does not touch anything #19/#20/#23/#25 covered for the SciPy-sparse boundary.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    architectureStructural or design-level changesenhancementNew feature or request

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions