lax

JAX-compiled Lagrange-mesh solvers for quantum scattering and bound-state problems.

lax is a low-level numerical engine implementing the Lagrange-mesh method (LMM) for solving the radial Schrödinger equation. It supports bound-state (eigenvalue) and continuum (R-matrix, S-matrix, Green’s function) calculations through a unified spectral-decomposition kernel, and is designed for use inside reaction codes, potential-fitting pipelines, and uncertainty-quantification workflows.

Key features:

• One eigendecomposition supports R-matrix, S-matrix, phase shifts, and Green’s functions at arbitrary energies — no per-energy linear solves.

• Full jit, vmap, and grad compatibility: push batches of potentials through a precompiled solver on CPU or GPU.

• Coupled-channel, non-local potentials, optical (absorptive) potentials, and R-matrix subinterval propagation.

• Legendre (finite-interval) and Laguerre (semi-infinite) mesh families, each with multiple endpoint regularizations.

Reference

• API Reference
  • Top-level
  • lax.constants
  • lax.spectral
  • lax.models
  • Internal modules
• Examples
  • Yamaguchi scattering example
  • Hydrogen bound states on a Laguerre mesh
  • α + \(^{208}\)Pb optical-model example
  • Coupled n–p scattering: the ³S₁–³D₁ deuteron channel
  • Rotor-coupled optical model: ¹⁶O + ⁴⁴Ca inelastic scattering
  • Closed channels and thresholds: α + ¹²C
  • Fourier transform demo
• Architecture reference
  • Revision history
  • 1. Overview and scope
  • 2. Background
  • 3. Design goals and constraints
  • 4. Architecture
  • 5. Module layout
  • 6. Core types
  • 7. Mesh builders and the registry pattern
  • 8. Operators
  • 9. Boundary values: Coulomb, Hankel, and Whittaker functions
  • 10. The spectral submodule
  • 11. Solvers and method dispatch
  • 12. Off-grid energies: interpolation is out of scope
  • 13. Transforms: grid, Fourier, integration
  • 14. The compile() factory
  • 15. Coupled-channel structure
  • 16. Public API and usage examples
  • 17. JAX considerations
  • 18. Testing strategy and benchmarks
  • 19. References
  • [6] P. Descouvemont, D. Baye, The R-matrix theory, Reports on Progress in Physics 73, 036301 (2010). General review of R-matrix theory in nuclear physics; discusses phenomenological R-matrix fitting in terms of poles and reduced widths, which are directly accessible from the Spectrum object.
  • Appendix A: Mesh formula tables
  • Appendix B: Glossary of symbols
  • Appendix C: Implementation sharp edges
  • wavefunction_grid / wavefunction_direct_grid return the internal solution of (H − E/μ)ψ = φ(a)·H⁻(a) — driven by the boundary value of the incoming wave. R-matrix engines that instead drive the internal solution with the matched exterior derivative u_ext'(a) = (i/2)(H⁻′ − S·H⁺′) differ per (l, j, channel) by exactly the scalar (i/2)(H⁻′ − S·H⁺′)/H⁻ — the driven solution is linear in the driving coefficient — times k/√a per channel if their coefficients live in the dimensionless s = k·r coordinate. The cross-engine acceptance test (tests/acceptance/; fixture generated from an external reference engine on a ⁴⁸Ca(p,n) IAS case with classical kinematics, so ħ²k²/2μ = Ecm maps each channel onto lax exactly) pins this relation to machine precision: elastic S-matrices agree to ~8×10⁻¹³ with zero free parameters, and the DWBA T-matrix to rtol 10⁻⁸ after the closed-form conversion T_ref = conv_p·conv_n·matrix_element(χp, χn, U₁)/a². Convention conversions belong in the caller, never inside lax.

Links

• Source code

• Architecture reference (DESIGN.md)

• Issue tracker