diff --git a/DESIGN.md b/DESIGN.md index de4fedd..6a4c7cf 100644 --- a/DESIGN.md +++ b/DESIGN.md @@ -1,11 +1,12 @@ # `lax`: A JAX-Compiled Lagrange-Mesh Library for Quantum Scattering and Bound-State Problems -**Design Document v1.5** +**Design Document v1.6** ## Revision history +- **v1.6** — Batched two-state evaluation and grid wavefunctions (spec v0.1.5.1, F1–F3). **F1:** `solver.matrix_element(bra, ket, operator=None, *, conjugate)` — a two-state, optionally non-conjugated bilinear form batched over the block/energy axes, bound unconditionally on every solver, plus a standalone `lax.transforms.matrix_element` (§13.4). **F2:** `solver.wavefunction_grid(spectrum, channel_index=0)` (spectral, both evaluation regimes via internal rank detection) and `solver.wavefunction_direct_grid(V, channel_index=0)` (linear-solve companion), with the Descouvemont eq.-27 source stack baked at compile time on `Solver.wavefunction_sources` and `lax.make_wavefunction_source_grid` exposing it (§11.2–11.3). **F3:** the `momenta=` × `blocks=` rejection is lifted — `F_momentum` gains a leading `(N_b,)` axis and the `grid=` projection kernels generalize to arbitrary leading batch axes (§13, §15.5). **C4 fix and guard:** the spectrum kernel now assembles each per-energy Hamiltonian with its own μ(E) when a non-uniform `mass_factor_grid` is compiled (previously the energy-batched path silently used a single uniform μ), `spectrum(V)` auto-routes to the energy-batched path on such solvers (`V.energy_dependent or mass_factor_nonuniform`), and the five static-regime spectral observables are bound to raising stubs (Appendix C.10). **C8:** the `eig` path's non-differentiability (host `pure_callback`) is documented; gradient/UQ pipelines use `linear_solve` + `wavefunction_direct_grid` (Appendix C.11). **C7:** the wavefunction normalization contract was pinned against an external reference engine by an active cross-engine DWBA acceptance test — S-matrices to ~8e-13 with zero free parameters, T-matrices via a closed-form convention conversion (Appendix C.12). - **v1.5** — Generalized the (previously design-only) partial-wave axis into a **symmetry-block batch axis** (§15.5). Any set of *symmetry blocks* — `(J, π)` coupled-channel groups, individual partial waves, or any other independent solves that share a channel shape `N_c` — is declared through a new `compile(blocks=…)` argument, stacked on a leading `(N_b,)` axis, and `vmap`-ped at runtime. This is the energy-axis mechanism (§4.2) applied along a second batch axis; partial waves are the `N_c = 1` special case. The `Interaction` gains a static `block_dependent` flag parallel to `energy_dependent`, with block shapes `(N_b, M, M)` / `(N_b, N_E, M, M)`. **Status: implemented and shipped** — both the direct path (`rmatrix_direct`/`smatrix_direct`/`phases_direct`/`wavefunction_direct`) and the full spectral path (`spectrum`, `rmatrix`, `smatrix`, `phases`, `greens`, `wavefunction`, `eigh`, the `*_grid` observables) vmap over the block axis; the spectral path requires one uniform mass factor across all blocks, per-block μ remains a direct-path feature (see §15.5). This revision also reconciles the document with the shipped code: core types moved to `lax/types.py` and `BoundaryValues` to `lax/spectral/types.py` (`boundary/_types.py` deleted); the explicit `solver.local_potential` / `solver.nonlocal_potential` builders (no arity inference, no `solver.potential`); `interaction_from_funcs(nonlocal_=…)` (the keyword `nonlocal` is invalid Python); `smatrix_from_R`'s √k normalization (`R̃ = K R K⁻¹`, `K = diag(√k_c)`) and the extra `BoundaryValues.k` field; per-channel μ scaling in `wavefunction_direct`; single-channel coupling sugar for the list builders; `dtype`/`device` compile parameters; R-matrix propagation (`propagate.py`) promoted from a non-goal to a documented module; and several Appendix A formula fixes. A post-implementation review (2026-06-11) verified the batched paths against per-block compiled solvers and closed the remaining open items (single-channel coupling sugar for the list builders; `dtype`/`device` compile parameters — both now implemented and tested). The phased build order that guided the implementation (former §19) was retired at the v1.5 close-out, with every phase through 11 shipped. The Padé interpolation utilities (`spectral.pade_interpolate` and the solver-bound `interpolate_*` builders, including the planned phase-12 derivative-enhanced variant) were subsequently **removed**: interpolation is observable-specific — global rational fits are defeated by the mod-π branch structure of phase shifts and by thresholds — so off-grid evaluation is left to the user (§12). -- **v1.4** — Designed the **partial-wave (ℓ) batch axis** (§15.5, build-order Phase 11): a compile-time `partial_waves` set with baked per-wave centrifugal and boundary, a leading partial-wave axis on `Interaction` (parallel to `energy_dependent`), and partial-wave-vmapped direct observables. Distinct from the coupled-channel axis (independent solves, not coupling). Motivated by ℓ-dependent non-local kernels in the jitr refactor (jitr `DESIGN.md §4.7`). **This was a design only — it was never implemented; v1.5 supersedes it with the more general symmetry-block axis.** +- **v1.4** — Designed the **partial-wave (ℓ) batch axis** (§15.5, build-order Phase 11): a compile-time `partial_waves` set with baked per-wave centrifugal and boundary, a leading partial-wave axis on `Interaction` (parallel to `energy_dependent`), and partial-wave-vmapped direct observables. Distinct from the coupled-channel axis (independent solves, not coupling). Motivated by ℓ-dependent non-local kernels in downstream optical-model workflows. **This was a design only — it was never implemented; v1.5 supersedes it with the more general symmetry-block axis.** - **v1.3** — Unified interaction interface and direct-path wavefunctions. The canonical solver input is now an assembled `Interaction` block `(N_E, M, M)` built by `interaction_from_{block,array,funcs}`; raw `(N_c,N_c,N[,N])` arrays are no longer accepted by solvers. Added internal wavefunctions on the linear-solve path (`wavefunction_direct`). Reworked the block assembler to the **symmetric MeV form** (mass baked into the kinetic block, coupling potential left untouched), which fixes the multi-mass asymmetry and lifts the single-μ limitation: **per-channel and energy-dependent reduced mass** (`ChannelSpec.mass_factor` per channel, `mass_factor_grid` shape `(N_E, N_c)`) are now first-class on the direct path. Updated §8, §10.2, §11.2–11.3, §11.5, §15; added Phase 10 to the build order (former Phase 10 → 11). - **v1.2** — Final pre-implementation review. Fixed unit convention bug (`threshold / mass_factor` in `assemble_block_hamiltonian`); fixed `rmatrix_from_spectrum` and `greens_from_spectrum` to convert E from MeV to fm⁻²; replaced undefined `_project_open`/`_pad_back` with full implementation; rewrote `rmatrix_direct` to vmap over compile-time energies; fixed Example 16.7 (energy-dependent V) which was semantically incorrect; split `to_grid` into two explicit functions; added §15.4 unit convention table; added Appendix C (10 implementation sharp edges). - **v1.1** — Unified all continuum solvers around a single spectral decomposition kernel. R-matrix and Green's function are computed as spectral sums over the eigenpairs of the Bloch-augmented Hamiltonian. Introduced a mesh-independent `spectral` submodule. Added energy-dependent V(E) compile mode with Padé interpolation. Added method-dispatch (`eigh` / `eig` / `linear_solve`) to handle real, complex, and GPU constraints. Linear-solve R-matrix-direct path moved to its own `rmatrix_direct` namespace. @@ -1383,10 +1384,35 @@ def _bind_solvers(spectrum_fn, mesh, channels, energies, boundary, mass_factor): return rmatrix, smatrix, phases, greens, eigh ``` +**`wavefunction_grid` (v1.6).** When `'wavefunction'` is requested with an energy grid, the +solver additionally binds `wavefunction_grid(spectrum, channel_index=0)`: internal +wavefunctions for **every** `(block, grid-energy)` pair against the compile-time-baked +source stack (`Solver.wavefunction_sources`, Descouvemont eq. 27 — built from +`mesh.basis_at_boundary` and the `H⁻` boundary cache, with **no** `is_open` masking; +sub-threshold entries are defined but are not scattering wavefunctions). One entry point +serves both evaluation regimes, detected from the `Spectrum` rank +(`spectrum_is_energy_batched`): a static-V spectrum runs the fused diagonalize-once einsum +``ψ_e = U · diag(1/(ε_k − E_e/μ)) · (Uᵀ_metric · source_e)`` over all energies without +materializing Green's matrices (§10.2), while an energy-batched spectrum vmaps the +per-energy Green's contraction over the aligned `(spec_e, E_e, source_e, μ_e)` axes — the +`smatrix_grid` pattern. `channel_index=None` returns all incoming channels on an extra +`N_c` axis. **Regime rule:** the energy-batched path is selected by +`V.energy_dependent or mass_factor_nonuniform` — a non-uniform μ(E) makes diagonalize-once +invalid physics even for a static V (Appendix C.10). + ### 11.3 The direct (linear-solve) path: `rmatrix_direct` and `wavefunction_direct` The direct path solves the Bloch–Schrödinger system per energy without an eigendecomposition. It is the GPU-/vmap-native route for complex V (where `eig` is a CPU callback) and the only path for non-local kernels that avoids forming a dense `M×M` eigenvector matrix. It consumes an `Interaction` (§15.1) and, as of v1.3, also produces internal wavefunctions. +**`wavefunction_direct_grid` (v1.6).** Whenever the direct path is active the solver also +binds `wavefunction_direct_grid(V, channel_index=0)`: one LU factorization per +`(block, energy)` solved against every incoming-channel column of the baked source stack — +the same shapes and closed-channel contract as the spectral `wavefunction_grid`. It +inherits the direct path's per-channel μ(E) support and, unlike the `eig` spectral path, is +**fully differentiable** (Appendix C.11) — the route for gradient/UQ work with complex +optical potentials. Propagated multi-interval meshes are rejected +(`NotImplementedError`): the boundary basis behind the source stack differs per interval. + ```python # solvers/linear_solve.py import jax @@ -1773,6 +1799,46 @@ def make_integration(mesh): For a wave function in a regularized basis, the orthonormality of the *Gauss approximation* [1, eq. 2.12] guarantees that $\langle\psi|\psi\rangle \approx \sum_j c_j^2$ to the same accuracy as the LMM itself, even when the regularized basis is not exactly orthonormal in the full $L^2$ sense [1, §2.7, eq. 2.69]. +### 13.4 Two-state bilinear matrix elements (v1.6) + +`integrate` is deliberately single-state and always-conjugating. The batched two-state +form lives in `transforms/bilinear.py` and is bound **unconditionally** on every solver +(it depends only on compile-time shapes): + +```python +solver.matrix_element(bra, ket, operator=None, *, conjugate) # keyword-required +# conjugate=False : braᵀ · O · ket (the DWBA bilinear; matches the +# complex-symmetric U diag Uᵀ metric) +# conjugate=True : bra† · O · ket (the Hermitian inner product) +``` + +Operator forms and the **scaling contract** (the silent-wrongness trap): ``None`` is the +overlap; a ``(M,)`` array is a *local* operator as **unscaled** node values ``V(rᵢ)`` +(the coefficients absorb ``√(λᵢ a)``, so the element is the plain node sum); a ``(M, M)`` +array is a caller-**Gauss-scaled** kernel ``K̃ᵢⱼ = √(λᵢλⱼ)·a·K(rᵢ,rⱼ)``; an +:class:`Interaction` is the recommended form — its assembled block already carries exactly +this scaling for both local and non-local content, and its static flags drive the +block/energy axis alignment (via `lift_block`). `matrix_element` adds no factors of +``a``, ``k``, or quadrature weights beyond what is inside the operator. + +Batching: `bra`/`ket` accept ``(M,)``, rank-2, or ``(N_b, N_E, M)`` in the canonical +block × energy order. Rank-2 follows a **deterministic mode-based rule** (never +shape-sniffing): blocks mode → block-leading ``(N_b, M)``; channels mode → energy-leading +``(N_E, M)``; mismatches raise naming the expected axis. Internally states are lifted to +``(N_b|1, N_E|1, M)`` and the three jitted einsum kernels +(`'...m,...m->...'` / `'...m,...m,...m->...'` / `'...m,...mn,...n->...'`) broadcast; the +output squeezes axes no input contributed (scalar for unbatched inputs). A standalone +`lax.transforms.matrix_element` does the pure broadcast einsum with no solver-aware +interpretation or validation. + +> **v1.6 note on §13.1–13.2 batching.** The `grid=` projection kernels are trailing-axis +> einsums (``'rn,...n->...r'`` and kin), so arbitrary leading batch axes (block, energy) +> pass through in both compile modes. In blocks mode `F_momentum` is +> ``(N_b, N_c, M_k, N)`` and `fourier`/`double_fourier_transform` accept ``(N_b, N)`` / +> ``(N_b, N, N)`` inputs, broadcasting unbatched ones across blocks; rank-2 inputs with +> leading dimension ``N_b`` are **always** block-batched vectors (in the ``N_b == N`` +> corner a broadcast kernel must be written ``jnp.broadcast_to(K, (N_b, N, N))``). + --- ## 14. The `compile()` factory @@ -2308,7 +2374,10 @@ array-parameterized per-block cores + `jax.vmap(block axis)` in block-dependent funcs term is a *sequence of N_b callables*). Constraints as shipped: the spectral path requires one uniform mass factor across all blocks (per-block/per-channel μ is a direct-path feature); `mass_factor_grid` is shared across blocks; `blocks=` excludes -propagated meshes and `momenta` transforms. +propagated meshes. As of v1.6 the `momenta=`/`grid=` transforms are block-batched — +`F_momentum` is `(N_b, N_c, M_k, N)` (deduped per unique ℓ across the whole block set), +`fourier`/`double_fourier_transform` accept `(N_b, …)` inputs and broadcast unbatched ones, +and the `grid=` projections pass arbitrary leading batch axes through (§13). --- @@ -2616,6 +2685,37 @@ interaction = solver.interaction_from_funcs( For coupled `(J, π)` groups (`N_c > 1`), each inner list holds that block's channels (all blocks sharing the same `N_c`), and the outputs carry the same leading `(N_b,)` axis. +### 16.10 Example 9: batched (p,n) DWBA transition elements (v1.6) + +The driving case for `wavefunction_grid` + `matrix_element`: distorted waves from two +compiled solvers (proton entrance with Coulomb, neutron exit without), contracted with an +isovector transition operator — **non-conjugated**, batched over every `(block, energy)` +pair in one call: + +```python +mesh = lax.MeshSpec("legendre", "x", n=40, scale=12.0) +blocks = [[lax.ChannelSpec(l=ell, threshold=0.0, mass_factor=mf_p)] for ell in ells] +proton = lax.compile(mesh=mesh, blocks=blocks, solvers=("spectrum", "wavefunction"), + energies=E_p, z1z2=(1, Z), V_is_complex=True, method="eig") +neutron = lax.compile(mesh=mesh, blocks=[[c._replace(mass_factor=mf_n)] for [c] in blocks], + solvers=("spectrum", "wavefunction"), + energies=E_p - Q_pn, V_is_complex=True, method="eig") + +chi_p = proton.wavefunction_grid(proton.spectrum(V_p)) # (N_b, N_E, M) +chi_n = neutron.wavefunction_grid(neutron.spectrum(V_n)) # (N_b, N_E, M) +U1 = proton.interaction_from_array(local=[(U1_stack, A)], block_dependent=True) + +k_p = proton.boundary.k[:, :, 0] # (N_b, N_E) +k_n = neutron.boundary.k[:, :, 0] +T_pn = proton.matrix_element(chi_p, chi_n, U1, conjugate=False) / (mesh.scale * k_p * k_n) +# (N_b, N_E) — the whole partial-wave × energy T-matrix in one bilinear call. +``` + +Both solvers must share the identical mesh (the Gauss scaling inside `U1.block` belongs to +it) and index-aligned block sets and energy grids — invariants the caller owns. For +gradient/UQ work replace the `eig` spectral path with `method="linear_solve"` and +`wavefunction_direct_grid(V)` (Appendix C.11). + --- ## 17. JAX considerations @@ -2637,7 +2737,7 @@ All numerical-data dataclasses (`Mesh`, `PropagationMatrices`, `OperatorMatrices | Method | Internally calls | Differentiable? | `vmap` over V? | GPU? | |---|---|---|---|---| | `eigh` | `jnp.linalg.eigh` | Yes (closed-form JVP) | Yes | Yes | -| `eig` | `jnp.linalg.eig` (via host callback) | Yes (custom JVP in JAX) | Awkward | No | +| `eig` | host `np.linalg.eig` via `jax.pure_callback` | **No** (the callback has no JVP/VJP rule — Appendix C.11) | Sequential | No | | `linear_solve` | `jnp.linalg.solve` per energy | Yes | Yes | Yes | The `eigh` derivative rule has a known degeneracy issue at eigenvalue crossings — gradients become large or ill-defined when two eigenvalues collide [see `jax.experimental.linalg.eigh` notes]. For potential-fitting workflows this is rarely problematic in practice (level crossings as a function of potential parameters are measure-zero), but it's worth documenting. @@ -2654,7 +2754,7 @@ The `device` argument to `compile()` places all cached arrays on the requested d ### 17.5 Gradient support -The hot path is `eigh` (default) or `solve` (R-matrix-direct path). Both have JAX-defined custom JVP/VJP rules and are fully differentiable. The only non-JAX components (mpmath, scipy.special) are confined to compile time, so they do not enter any backward pass. +The hot path is `eigh` (default) or `solve` (R-matrix-direct path). Both have JAX-defined custom JVP/VJP rules and are fully differentiable. The compile-time non-JAX components (mpmath, scipy.special) never enter a backward pass. The **`eig` path is the exception**: its spectra flow through a runtime host callback (`jax.pure_callback` → `np.linalg.eig`) with no JVP/VJP rule, so `jax.grad` through any `eig`-path spectrum raises — gradient/UQ pipelines with complex optical potentials use `method="linear_solve"` (+ `wavefunction_direct_grid`) or restrict to real-V `eigh` problems (Appendix C.11). For second derivatives, `jax.hessian(loss)` works through both `eigh` and `solve`. Be aware that the second derivative of `eigh` involves squared denominators in eigenvalue differences and can be numerically unstable near degeneracies. @@ -2970,6 +3070,50 @@ jax.config.update("jax_enable_x64", True) ``` If the user imports JAX before importing `lax`, the call may not take effect (JAX freezes the config on first array creation). The fix is to import `lax` first, or to put the config call at the top of the user's script before any `import jax.numpy as jnp` usage. Document this prominently in the README. + +### C.10 Non-uniform `mass_factor_grid` forces the energy-batched spectral regime (v1.6) + +With a non-uniform μ(E) the dimensionless Hamiltonian `H/μ(E)` itself changes with energy, +so "diagonalize once, evaluate many energies" is invalid **physics**, not just missing +plumbing: a static-regime evaluation would mix per-energy boundary values with a single-μ +spectral denominator — silently wrong matching. As shipped in v1.6: (1) the spectrum +kernel assembles each per-energy Hamiltonian with its own μ_e (the energy-batched path +takes a `(N_E,)` `mass_factor_grid`, validated per-channel-uniform for the spectral path); +(2) `spectrum(V)` auto-routes to the energy-batched path when +`V.energy_dependent or mass_factor_nonuniform` — never on `V.energy_dependent` alone; +(3) the five static-regime observables (`rmatrix`, `smatrix`, `phases`, `greens`, +`wavefunction`) are bound to pickle-safe **raising stubs** (`_NonuniformMassFactorStub`) +that point at the `*_grid` observables and the direct path; `wavefunction_grid` +additionally raises in its static branch if handed a foreign static spectrum. Per-channel +non-uniform grids remain a direct-path-only feature and are rejected at compile for the +spectral path. + +### C.11 `method="eig"` spectra are not differentiable (v1.6) + +Complex-symmetric spectra are computed through `jax.pure_callback` → host +`np.linalg.eig` (`solvers/spectrum.py::_eig_via_callback`). The callback carries no +JVP/VJP rule, so `jax.grad` through **any** `eig`-path spectrum raises ("Pure callbacks do +not support JVP"). DWBA distorted waves from complex optical potentials are exactly the +`eig` case: gradient/UQ pipelines must use `method="linear_solve"` + +`wavefunction_direct_grid` (fully differentiable through `jnp.linalg.solve`) or restrict +to real-V `eigh` problems. A test asserts the raise so a future custom-JVP upgrade is +noticed. + +### C.12 Cross-engine wavefunction normalization (v1.6) + +`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. --- -*End of design document. Version 1.2, intended for offline reference during library development.* +*End of design document. Version 1.6, intended for offline reference during library development.* diff --git a/docs/design.md b/docs/design.md index 7e53821..cacffb9 100644 --- a/docs/design.md +++ b/docs/design.md @@ -1,6 +1,9 @@ # Architecture reference +% Skip DESIGN.md's own H1 title: this page supplies the single top-level +% heading, so the included ## sections nest cleanly in the sidebar TOC. ```{include} ../DESIGN.md +:start-after: "Bound-State Problems" :end-before: "---" ``` diff --git a/src/lax/__init__.py b/src/lax/__init__.py index 0773e0d..1fcb511 100644 --- a/src/lax/__init__.py +++ b/src/lax/__init__.py @@ -19,7 +19,7 @@ import lax.spectral as spectral from lax.compile import compile from lax.types import ChannelSpec, Interaction, MeshSpec, Solver -from lax.wavefunction import make_wavefunction_source +from lax.wavefunction import make_wavefunction_source, make_wavefunction_source_grid __all__ = [ "ChannelSpec", @@ -29,6 +29,7 @@ "compile", "constants", "make_wavefunction_source", + "make_wavefunction_source_grid", "models", "spectral", ] diff --git a/src/lax/compile.py b/src/lax/compile.py index 70c09d3..54a10cb 100644 --- a/src/lax/compile.py +++ b/src/lax/compile.py @@ -28,6 +28,8 @@ from lax.solvers import ( bind_grid_observables, bind_observables, + bind_wavefunction_grid_observable, + make_direct_wavefunction_grid_kernel, make_direct_wavefunction_kernel, make_phases_direct_observable, make_rmatrix_direct_kernel, @@ -41,6 +43,7 @@ make_double_fourier, make_fourier, make_integration, + make_matrix_element, make_to_grid, ) from lax.types import ( @@ -55,6 +58,7 @@ GridMatrixTransform, GridVectorTransform, Integrator, + MatrixElementHelper, Mesh, MeshSpec, Method, @@ -67,9 +71,12 @@ SpectrumKernel, SpectrumObservable, TransformMatrices, + WavefunctionDirectGridObservable, WavefunctionDirectObservable, + WavefunctionGridObservable, WavefunctionObservable, ) +from lax.wavefunction import build_wavefunction_sources @dataclass(frozen=True) @@ -99,6 +106,7 @@ class _TransformBundle: fourier: FourierTransform | None double_fourier_transform: DoubleFourierTransform | None integrate: Integrator + matrix_element: MatrixElementHelper @dataclass(frozen=True) @@ -112,6 +120,7 @@ class _ObservableBundle: greens: GreenFunctionObservable | None wavefunction: WavefunctionObservable | None eigh: EigenpairAccessor | None + wavefunction_grid: WavefunctionGridObservable | None rmatrix_grid: SpectrumGridObservable | None smatrix_grid: SpectrumGridObservable | None phases_grid: SpectrumGridObservable | None @@ -119,6 +128,7 @@ class _ObservableBundle: smatrix_direct: SMatrixDirectObservable | None phases_direct: PhasesDirectObservable | None wavefunction_direct: WavefunctionDirectObservable | None + wavefunction_direct_grid: WavefunctionDirectGridObservable | None interaction_from_block: Callable[..., Any] | None interaction_from_array: Callable[..., Any] | None interaction_from_funcs: Callable[..., Any] | None @@ -248,12 +258,6 @@ def compile( method=method, V_is_complex=V_is_complex, ) - if request.block_groups is not None and momenta is not None: - msg = ( - "`momenta` Fourier transforms are not supported with `blocks=`; " - "compile per-block solvers for momentum-space work." - ) - raise ValueError(msg) mesh_data, operator_matrices = _build_solver_mesh( mesh=mesh, operators=request.operators, @@ -300,15 +304,35 @@ def compile( momenta = _finalize_arrays(momenta, dtype, target_device) if momenta is not None else None transforms = _prepare_transforms( mesh=mesh_data, - channels=request.channels, + blocks=blocks_resolved, grid=grid, momenta=momenta, + block_mode=block_mode, + n_energies=len(np.asarray(energies)) if energies is not None else None, ) mass_factor_grid_jax = ( _finalize_arrays(jnp.asarray(mass_factor_grid_np), dtype, target_device) if mass_factor_grid_np is not None else None ) + # Bake the Descouvemont eq.-27 source stack whenever a wavefunction entry + # point can exist (spec v0.1.5.1 §4.2): it is fully determined by the + # boundary cache. Propagated meshes have no single boundary basis (C2). + wavefunction_sources: jax.Array | None = None + if "wavefunction" in request.solvers and mesh_data.propagation is not None: + raise NotImplementedError( + "'wavefunction' is not supported on propagated multi-interval meshes — " + "the boundary basis behind the wavefunction source differs per interval. " + "Use a single-interval mesh." + ) + if ( + boundary is not None + and mesh_data.propagation is None + and ("wavefunction" in request.solvers or "rmatrix_direct" in request.solvers) + ): + wavefunction_sources = build_wavefunction_sources( + mesh_data, boundary, len(request.channels) + ) observables = _bind_solver_observables( request=request, blocks=blocks_resolved, @@ -319,6 +343,7 @@ def compile( boundary=boundary, has_energy_grid=energies is not None, mass_factor_grid=mass_factor_grid_jax, + wavefunction_sources=wavefunction_sources, ) return _assemble_solver( request=request, @@ -329,6 +354,7 @@ def compile( transforms=transforms, observables=observables, mass_factor_grid=mass_factor_grid_jax, + wavefunction_sources=wavefunction_sources, ) @@ -510,11 +536,18 @@ def _prepare_boundary_data( def _prepare_transforms( *, mesh: Mesh, - channels: tuple[ChannelSpec, ...], + blocks: tuple[tuple[ChannelSpec, ...], ...], grid: jax.Array | None, momenta: jax.Array | None, + block_mode: bool, + n_energies: int | None, ) -> _TransformBundle: - """Bind optional grid/Fourier transforms around one compiled mesh.""" + """Bind optional grid/Fourier transforms around one compiled mesh. + + In blocks mode the Fourier stack carries a leading ``(N_b,)`` axis, + ``(N_b, N_c, M_k, N)``, with ``compute_F_momentum`` deduplicated per + unique ℓ across the whole block set. + """ transforms = TransformMatrices() to_grid_vector_fn: GridVectorTransform | None = None @@ -540,12 +573,17 @@ def _prepare_transforms( if momenta is not None: momenta_array = jnp.asarray(np.asarray(momenta)) - unique_angular_momenta = sorted({channel.l for channel in channels}) + unique_angular_momenta = sorted({channel.l for group in blocks for channel in group}) matrices_by_l = { angular_momentum: compute_F_momentum(mesh, momenta_array, angular_momentum) for angular_momentum in unique_angular_momenta } - fourier_stack = jnp.stack([matrices_by_l[channel.l] for channel in channels]) + if block_mode: + fourier_stack = jnp.stack( + [jnp.stack([matrices_by_l[channel.l] for channel in group]) for group in blocks] + ) + else: + fourier_stack = jnp.stack([matrices_by_l[channel.l] for channel in blocks[0]]) transforms = TransformMatrices( B_grid=transforms.B_grid, grid_r=transforms.grid_r, @@ -563,6 +601,12 @@ def _prepare_transforms( fourier=fourier_fn, double_fourier_transform=double_fourier_transform_fn, integrate=make_integration(mesh), + matrix_element=make_matrix_element( + matrix_size=mesh.n * len(blocks[0]), + n_blocks=len(blocks), + n_energies=n_energies, + block_mode=block_mode, + ), ) @@ -577,6 +621,7 @@ def _bind_solver_observables( boundary: BoundaryValues | None, has_energy_grid: bool, mass_factor_grid: jax.Array | None = None, + wavefunction_sources: jax.Array | None = None, ) -> _ObservableBundle: """Bind all runtime entry points requested for one compiled solver. @@ -592,11 +637,20 @@ def _bind_solver_observables( greens_fn: GreenFunctionObservable | None = None wavefunction_fn: WavefunctionObservable | None = None eigh_fn: EigenpairAccessor | None = None + wavefunction_grid_fn: WavefunctionGridObservable | None = None rmatrix_grid_fn: SpectrumGridObservable | None = None smatrix_grid_fn: SpectrumGridObservable | None = None phases_grid_fn: SpectrumGridObservable | None = None + sources_blocks = ( + None + if wavefunction_sources is None + else (wavefunction_sources if block_mode else wavefunction_sources[None]) + ) if request.needs_spectrum: + spectral_mass_factor_grid, mass_factor_nonuniform = _spectral_mass_factor_grid( + mass_factor_grid + ) spectrum_fn = make_spectrum_kernel( mesh, operators, @@ -604,6 +658,8 @@ def _bind_solver_observables( method=request.method, keep_eigenvectors=request.keep_eigenvectors, block_mode=block_mode, + mass_factor_grid=spectral_mass_factor_grid, + mass_factor_nonuniform=mass_factor_nonuniform, ) ( bound_rmatrix, @@ -612,7 +668,14 @@ def _bind_solver_observables( bound_greens, bound_wavefunction, bound_eigh, - ) = bind_observables(mesh, blocks, energies, boundary, block_mode=block_mode) + ) = bind_observables( + mesh, + blocks, + energies, + boundary, + block_mode=block_mode, + mass_factor_nonuniform=mass_factor_nonuniform, + ) # S- and phase-shift observables are built from the spectral R-matrix, so # the raw R observable remains available whenever matching is requested. @@ -637,8 +700,20 @@ def _bind_solver_observables( blocks, energies, boundary, - mass_factor_grid=mass_factor_grid, + mass_factor_grid=spectral_mass_factor_grid, + block_mode=block_mode, + ) + # The wavefunction-grid observable mirrors the single-energy + # `wavefunction` gating; it is spectral-only (the linear_solve method + # is served by wavefunction_direct_grid below). + if "wavefunction" in request.solvers and sources_blocks is not None: + wavefunction_grid_fn = bind_wavefunction_grid_observable( + blocks, + energies, + sources_blocks, + mass_factor_grid=spectral_mass_factor_grid, block_mode=block_mode, + mass_factor_nonuniform=mass_factor_nonuniform, ) rmatrix_direct_fn: DirectRMatrixKernel | None = None @@ -663,6 +738,7 @@ def _bind_solver_observables( # ("rmatrix_direct"), and is the binding for "wavefunction" under # method="linear_solve" (§14, Example 16.8) where no spectral path exists. wavefunction_via_direct = "wavefunction" in request.solvers and request.method == "linear_solve" + wavefunction_direct_grid_fn: WavefunctionDirectGridObservable | None = None if "rmatrix_direct" in request.solvers or wavefunction_via_direct: wavefunction_direct_fn = make_direct_wavefunction_kernel( mesh, @@ -672,6 +748,16 @@ def _bind_solver_observables( mass_factor_grid, block_mode=block_mode, ) + if sources_blocks is not None: + wavefunction_direct_grid_fn = make_direct_wavefunction_grid_kernel( + mesh, + operators, + blocks, + energies, + sources_blocks, + mass_factor_grid, + block_mode=block_mode, + ) n_blocks = len(blocks) if block_mode else None interaction_from_block_fn = make_interaction_from_block( @@ -695,12 +781,14 @@ def _bind_solver_observables( greens=greens_fn, wavefunction=wavefunction_fn, eigh=eigh_fn, + wavefunction_grid=wavefunction_grid_fn, rmatrix_grid=rmatrix_grid_fn, smatrix_grid=smatrix_grid_fn, phases_grid=phases_grid_fn, rmatrix_direct=rmatrix_direct_fn, smatrix_direct=smatrix_direct_fn, phases_direct=phases_direct_fn, + wavefunction_direct_grid=wavefunction_direct_grid_fn, wavefunction_direct=wavefunction_direct_fn, interaction_from_block=interaction_from_block_fn, interaction_from_array=interaction_from_array_fn, @@ -720,6 +808,7 @@ def _assemble_solver( transforms: _TransformBundle, observables: _ObservableBundle, mass_factor_grid: jax.Array | None = None, + wavefunction_sources: jax.Array | None = None, ) -> Solver: """Assemble the final public solver bundle from normalized subcomponents.""" @@ -744,6 +833,8 @@ def _assemble_solver( phases=observables.phases, greens=observables.greens, wavefunction=observables.wavefunction, + wavefunction_grid=observables.wavefunction_grid, + wavefunction_sources=wavefunction_sources, eigh=observables.eigh, rmatrix_grid=observables.rmatrix_grid, smatrix_grid=observables.smatrix_grid, @@ -752,6 +843,7 @@ def _assemble_solver( smatrix_direct=observables.smatrix_direct, phases_direct=observables.phases_direct, wavefunction_direct=observables.wavefunction_direct, + wavefunction_direct_grid=observables.wavefunction_direct_grid, interaction_from_block=observables.interaction_from_block, interaction_from_array=observables.interaction_from_array, interaction_from_funcs=observables.interaction_from_funcs, @@ -763,9 +855,42 @@ def _assemble_solver( fourier=transforms.fourier, double_fourier_transform=transforms.double_fourier_transform, integrate=transforms.integrate, + matrix_element=transforms.matrix_element, ) +def _spectral_mass_factor_grid( + mass_factor_grid: jax.Array | None, +) -> tuple[jax.Array | None, bool]: + """Reduce the canonical μ(E) grid for the spectral path and detect non-uniformity. + + The spectral path folds a single ℏ²/2μ per energy out of the Hamiltonian, + so a ``(N_E, N_c)`` grid must be per-channel uniform; the validated grid + is reduced to ``(N_E,)``. Returns ``(reduced_grid, nonuniform)`` where + ``nonuniform`` is ``True`` iff μ genuinely varies with energy — the + condition that forces the energy-batched spectrum path and stubs out the + static-regime observables (silently mixing per-energy boundary values + with a single-μ spectral denominator is wrong matching). + """ + + if mass_factor_grid is None: + return None, False + grid_np = np.asarray(mass_factor_grid) + if grid_np.ndim == 2: + if not np.allclose(grid_np, grid_np[:, :1]): + msg = ( + "The spectral path requires a per-channel-uniform mass_factor_grid " + "(shape (N_E,) or per-channel-identical (N_E, N_c)); per-channel " + "mass factors are a direct-path feature (method='linear_solve')." + ) + raise ValueError(msg) + reduced = mass_factor_grid[:, 0] + else: + reduced = mass_factor_grid + nonuniform = not bool(np.allclose(grid_np, grid_np.flat[0])) + return reduced, nonuniform + + def _resolve_device(device: Any) -> Any: """Resolve a device-platform string (e.g. ``"cpu"``) to a concrete device.""" diff --git a/src/lax/solvers/__init__.py b/src/lax/solvers/__init__.py index d2ec7dd..58d1a65 100644 --- a/src/lax/solvers/__init__.py +++ b/src/lax/solvers/__init__.py @@ -2,6 +2,7 @@ from lax.solvers.assembly import assemble_block_hamiltonian, build_Q from lax.solvers.linear_solve import ( + make_direct_wavefunction_grid_kernel, make_direct_wavefunction_kernel, make_phases_direct_observable, make_rmatrix_direct_kernel, @@ -10,6 +11,8 @@ from lax.solvers.observables import ( bind_grid_observables, bind_observables, + bind_wavefunction_grid_observable, + spectrum_is_energy_batched, ) from lax.solvers.spectrum import make_spectrum_kernel @@ -17,10 +20,13 @@ "assemble_block_hamiltonian", "bind_grid_observables", "bind_observables", + "bind_wavefunction_grid_observable", "build_Q", + "make_direct_wavefunction_grid_kernel", "make_direct_wavefunction_kernel", "make_phases_direct_observable", "make_rmatrix_direct_kernel", "make_smatrix_direct_observable", "make_spectrum_kernel", + "spectrum_is_energy_batched", ] diff --git a/src/lax/solvers/linear_solve.py b/src/lax/solvers/linear_solve.py index 90db2d8..99fea92 100644 --- a/src/lax/solvers/linear_solve.py +++ b/src/lax/solvers/linear_solve.py @@ -354,6 +354,97 @@ def __call__( return result if self.block_mode else result[0] +@dataclass(frozen=True) +class _WavefunctionDirectGridKernel: + """Pickle-safe direct wavefunction-grid kernel: every grid energy in one call. + + The ``method="linear_solve"`` companion of the spectral + ``wavefunction_grid``: one LU solve per ``(block, energy)`` against the + compile-time-baked source stack. Inherits the direct path's per-channel + μ(E) support. + """ + + mesh: Mesh + operators: OperatorMatrices + energies: jax.Array + centrifugal: jax.Array # (N_b, N_c) + thresholds: jax.Array # (N_b, N_c) + mass_factor_grid_blocks: jax.Array # (N_b, N_E, N_c) + sources: jax.Array # (N_b, N_E, N_c, N_c·N) + matrix_size: int + n_blocks: int + n_energies: int + block_mode: bool + + def __call__( + self, + potential: jax.Array | Interaction, + channel_index: int | None = 0, + ) -> jax.Array: + """Solve ``C(E_i) ψ = source_i`` for every ``(block, grid-energy)`` pair. + + Parameters + ---------- + potential + An :class:`~lax.Interaction`; block- and/or energy-dependent + interactions are broadcast automatically. + channel_index + Incoming channel within the block (default 0); ``None`` returns + all incoming channels on an extra ``N_c`` axis. + + Returns + ------- + jax.Array + Wavefunction coefficients, shape ``(N_E, N_c·N)`` — + ``(N_b, N_E, N_c·N)`` in blocks mode; ``channel_index=None`` + inserts an ``N_c`` axis before the coefficient axis. + """ + + if not isinstance(potential, Interaction): + raise TypeError( + "wavefunction_direct_grid() accepts only Interaction objects. " + "Use solver.local_potential()/solver.nonlocal_potential() or solver.interaction_from_block/array/funcs to build one." + ) + if not self.block_mode: + reject_block_dependent(potential, "wavefunction_direct_grid()") + block = lift_to_blocks(potential.block, potential.block_dependent, self.n_blocks) + if potential.energy_dependent: + if block.shape[1] != self.n_energies: + msg = ( + f"energy-dependent Interaction carries {block.shape[1]} energies " + f"but the compile-time grid has N_E = {self.n_energies}." + ) + raise ValueError(msg) + blocks_grid = block + else: + blocks_grid = jnp.broadcast_to( + block[:, None], + (self.n_blocks, self.n_energies, *block.shape[-2:]), + ) + sources = ( + self.sources + if channel_index is None + else self.sources[:, :, channel_index : channel_index + 1, :] + ) + result = cast( + jax.Array, + _WAVEFUNCTION_DIRECT_GRID_BLOCKS_JIT( + blocks_grid, + sources, + self.energies, + self.mass_factor_grid_blocks, + self.mesh, + self.operators, + self.centrifugal, + self.thresholds, + self.matrix_size, + ), + ) + if channel_index is not None: + result = result[:, :, 0, :] + return result if self.block_mode else result[0] + + def _block_mass_data( blocks: tuple[tuple[ChannelSpec, ...], ...], channel_mass_rows: jax.Array, @@ -539,6 +630,58 @@ def make_direct_wavefunction_kernel( ) +def make_direct_wavefunction_grid_kernel( + mesh: Mesh, + operators: OperatorMatrices, + blocks: tuple[tuple[ChannelSpec, ...], ...], + energies: jax.Array, + sources: jax.Array, + mass_factor_grid: jax.Array | None = None, + block_mode: bool = False, +) -> _WavefunctionDirectGridKernel: + """Build a direct wavefunction-grid kernel ``(V, channel_index) → ψ``. + + Vmaps the per-energy direct solve over the whole compile-time grid against + the baked source stack: ``C(E_i)`` is LU-factorized once per + ``(block, energy)`` and solved for every incoming-channel column. + [DESIGN.md §11.3, §15.5] + + Parameters + ---------- + sources + Block-stacked baked source stack, shape ``(N_b, N_E, N_c, N_c·N)`` + (see :func:`lax.wavefunction.build_wavefunction_sources`). + + Raises + ------ + NotImplementedError + On propagated multi-interval meshes — the boundary basis behind the + source stack differs per interval. + """ + + if mesh.propagation is not None: + raise NotImplementedError( + "wavefunction_direct_grid is not supported on propagated multi-interval " + "meshes — the boundary basis behind the source stack differs per " + "interval. Use a single-interval mesh." + ) + centrifugal, thresholds, channel_mass_rows = block_group_arrays(blocks) + _, mfg_blocks, _ = _block_mass_data(blocks, channel_mass_rows, energies, mass_factor_grid) + return _WavefunctionDirectGridKernel( + mesh=mesh, + operators=operators, + energies=energies, + centrifugal=centrifugal, + thresholds=thresholds, + mass_factor_grid_blocks=mfg_blocks, + sources=sources, + matrix_size=mesh.n * len(blocks[0]), + n_blocks=len(blocks), + n_energies=len(energies), + block_mode=block_mode, + ) + + def _rmatrix_direct_core( potential: jax.Array, mesh: Mesh, @@ -756,6 +899,35 @@ def _wavefunction_direct_core( return result +def _wavefunction_direct_grid_core( + potential: jax.Array, + sources: jax.Array, + energy: jax.Array, + mu_row: jax.Array, + mesh: Mesh, + operators: OperatorMatrices, + centrifugal: jax.Array, + thresholds: jax.Array, + matrix_size: int, +) -> jax.Array: + """Multi-RHS variant of :func:`_wavefunction_direct_core`. + + ``sources`` is ``(K, N_c·N)`` (one row per incoming channel): ``C(E)`` is + LU-factorized once and solved for all K columns, then scaled per channel + by μ exactly as the single-source core. + """ + + hamiltonian = assemble_hamiltonian_arrays( + mesh, operators, centrifugal, thresholds, mu_row, potential + ) + matrix = hamiltonian - energy * jnp.eye(matrix_size, dtype=hamiltonian.dtype) + lu_piv = cast(tuple[jax.Array, jax.Array], jsl.lu_factor(matrix)) + solved = cast(jax.Array, jsl.lu_solve(lu_piv, sources.T)) # (N_c·N, K) + scale = jnp.repeat(mu_row, mesh.n) # (N_c·N,) + result: jax.Array = (scale.astype(solved.dtype)[:, None] * solved).T + return result + + def _direct_smatrix_grid( r_grid: jax.Array, boundary: BoundaryValues, @@ -924,6 +1096,55 @@ def one_block( return jax.vmap(one_block)(blocks, sources, mu_rows, centrifugal, thresholds) +def _wavefunction_direct_grid_blocks( + blocks: jax.Array, + sources: jax.Array, + energies: jax.Array, + mu_rows_grid: jax.Array, + mesh: Mesh, + operators: OperatorMatrices, + centrifugal: jax.Array, + thresholds: jax.Array, + matrix_size: int, +) -> jax.Array: + """Return block × energy batched wavefunctions, shape ``(N_b, N_E, K, N_c·N)``. + + Nested block × energy vmap of :func:`_wavefunction_direct_grid_core` over + the aligned ``(block_e, source_e, E_e, μ_row_e)`` axes — one LU + factorization per ``(block, energy)`` shared across the K incoming + channels. + """ + + def one_block( + block_grid: jax.Array, + source_grid: jax.Array, + mu_grid: jax.Array, + centrifugal_row: jax.Array, + threshold_row: jax.Array, + ) -> jax.Array: + def one_energy( + block: jax.Array, + source: jax.Array, + energy: jax.Array, + mu_row: jax.Array, + ) -> jax.Array: + return _wavefunction_direct_grid_core( + block, + source, + energy, + mu_row, + mesh, + operators, + centrifugal_row, + threshold_row, + matrix_size, + ) + + return jax.vmap(one_energy)(block_grid, source_grid, energies, mu_grid) + + return jax.vmap(one_block)(blocks, sources, mu_rows_grid, centrifugal, thresholds) + + _RMATRIX_DIRECT_PROPAGATED_JIT = jax.jit( _rmatrix_direct_propagated, static_argnames=("channels",), @@ -942,6 +1163,10 @@ def one_block( _wavefunction_direct_blocks, static_argnames=("matrix_size",), ) +_WAVEFUNCTION_DIRECT_GRID_BLOCKS_JIT = jax.jit( + _wavefunction_direct_grid_blocks, + static_argnames=("matrix_size",), +) _DIRECT_SMATRIX_BLOCKS_JIT = jax.jit(_direct_smatrix_blocks) _DIRECT_PHASES_BLOCKS_JIT = jax.jit(_direct_phases_blocks) diff --git a/src/lax/solvers/observables.py b/src/lax/solvers/observables.py index 8da832c..5288e99 100644 --- a/src/lax/solvers/observables.py +++ b/src/lax/solvers/observables.py @@ -14,7 +14,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import cast +from typing import NoReturn, cast import jax import jax.numpy as jnp @@ -23,6 +23,7 @@ from lax.spectral.observables import ( greens_from_spectrum, rmatrix_from_spectrum, + wavefunction_grid_from_spectrum, wavefunction_internal_from_spectrum, ) from lax.spectral.types import Spectrum @@ -35,12 +36,38 @@ RMatrixObservable, SpectrumGridObservable, SpectrumObservable, + WavefunctionGridObservable, WavefunctionObservable, ) from .assembly import add_block_axis, uniform_block_mass_factor +@dataclass(frozen=True) +class _NonuniformMassFactorStub: + """Raising stand-in for static-regime observables on a non-uniform μ(E) solver. + + With a non-uniform ``mass_factor_grid`` the dimensionless Hamiltonian + ``H/μ(E)`` varies with energy, so a static-regime evaluation would mix + per-energy boundary values with a single-μ spectral denominator — silently + wrong matching. The stub keeps the observable slot non-``None`` (feature + detection, ``__repr__``, pickling all behave) while turning any call into + a clear error pointing at the legitimate paths. + """ + + observable: str + + def __call__(self, *args: object, **kwargs: object) -> NoReturn: + raise ValueError( + f"solver.{self.observable} is unavailable: this solver was compiled with a " + "non-uniform mass_factor_grid, so static-regime spectral evaluation would mix " + "per-energy boundary values with a single-mu spectral denominator. " + "solver.spectrum(V) returns an energy-batched Spectrum on this solver - use " + "the aligned-grid observables (rmatrix_grid / smatrix_grid / phases_grid / " + "wavefunction_grid) or the direct path (method='linear_solve')." + ) + + @dataclass(frozen=True) class _RMatrixObservable: """Pickle-safe spectral R-matrix observable.""" @@ -296,6 +323,7 @@ def bind_observables( energies: jax.Array, boundary: BoundaryValues | None, block_mode: bool = False, + mass_factor_nonuniform: bool = False, ) -> tuple[ RMatrixObservable, SpectrumObservable | None, @@ -323,6 +351,12 @@ def bind_observables( block_mode ``True`` for a ``blocks=`` compile: outputs keep the leading ``(N_b,)`` axis. + mass_factor_nonuniform + ``True`` when the solver was compiled with a non-uniform + ``mass_factor_grid``. The five static-regime observables are then + bound to :class:`_NonuniformMassFactorStub` instances that raise on + call — the aligned-grid observables and the direct path are the + legitimate μ(E) routes. Raw eigenpair access stays live. Returns ------- @@ -331,6 +365,16 @@ def bind_observables( function, internal wavefunction, and raw eigensystem access. """ + if mass_factor_nonuniform: + return ( + _NonuniformMassFactorStub(observable="rmatrix"), + _NonuniformMassFactorStub(observable="smatrix") if boundary is not None else None, + _NonuniformMassFactorStub(observable="phases") if boundary is not None else None, + _NonuniformMassFactorStub(observable="greens"), + _NonuniformMassFactorStub(observable="wavefunction"), + _EigenpairAccessor(), + ) + channel_radius = mesh.scale mass_factor = uniform_block_mass_factor(blocks, context="spectral observable path") rmatrix = _RMatrixObservable( @@ -434,6 +478,130 @@ def bind_grid_observables( return rmatrix_grid, smatrix_grid, phases_grid +def spectrum_is_energy_batched(spectrum: Spectrum) -> bool: + """Return whether a block-lifted ``Spectrum`` carries an energy batch axis. + + Apply *after* :func:`add_block_axis` so one rule serves both compile + modes: block-lifted eigenvalues are ``(N_b, M)`` for a static-V spectrum + and ``(N_b, N_E, M)`` for the energy-batched regime. Ranks are static + under ``jax.jit``. + """ + + return spectrum.eigenvalues.ndim >= 3 + + +@dataclass(frozen=True) +class _WavefunctionGridObservable: + """Pickle-safe wavefunction-grid observable, serving both evaluation regimes. + + The static-V regime ("diagonalize once, evaluate the grid") runs the fused + einsum of :func:`wavefunction_grid_from_spectrum`; the energy-batched + regime vmaps the per-energy Green's contraction over the aligned + ``(spectra_e, E_e, source_e, μ_e)`` axes — the ``smatrix_grid`` pattern. + The regime is detected from the ``Spectrum`` rank. + """ + + sources: jax.Array # (N_b, N_E, N_c, M) — always block-stacked + energies: jax.Array # (N_E,) + mass_factor: float + mass_factor_grid: jax.Array # (N_E,) + n_energies: int + block_mode: bool + mass_factor_nonuniform: bool + + def __call__(self, spectrum: Spectrum, channel_index: int | None = 0) -> jax.Array: + """Evaluate ψ for every ``(block, grid-energy)`` pair; see the protocol.""" + + if spectrum.eigenvectors is None: + raise RuntimeError( + "Eigenvectors were not retained at compile time. " + "Re-compile with 'greens' or 'wavefunction' in solvers= to keep them." + ) + if not self.block_mode: + spectrum = add_block_axis(spectrum) + batched = spectrum_is_energy_batched(spectrum) + if batched and spectrum.eigenvalues.shape[1] != self.n_energies: + msg = ( + f"energy-batched Spectrum carries {spectrum.eigenvalues.shape[1]} energies " + f"but the compile-time grid has N_E = {self.n_energies}." + ) + raise ValueError(msg) + sources = ( + self.sources + if channel_index is None + else self.sources[:, :, channel_index : channel_index + 1, :] + ) + if batched: + result = cast( + jax.Array, + _WAVEFUNCTION_GRID_ALIGNED_BLOCKS_JIT( + spectrum, self.energies, sources, self.mass_factor_grid + ), + ) + else: + if self.mass_factor_nonuniform: + raise ValueError( + "wavefunction_grid received a static-V Spectrum, but this solver was " + "compiled with a non-uniform mass_factor_grid: H/mu(E) varies with " + "energy, so diagonalize-once is invalid. Pass the energy-batched " + "Spectrum from this solver's spectrum(V) (it dispatches per energy), " + "or use wavefunction_direct_grid (method='linear_solve')." + ) + result = cast( + jax.Array, + _WAVEFUNCTION_GRID_STATIC_BLOCKS_JIT( + spectrum, self.energies, sources, self.mass_factor + ), + ) + if channel_index is not None: + result = result[:, :, 0, :] + return result if self.block_mode else result[0] + + +def bind_wavefunction_grid_observable( + blocks: tuple[tuple[ChannelSpec, ...], ...], + energies: jax.Array, + sources: jax.Array, + mass_factor_grid: jax.Array | None = None, + block_mode: bool = False, + mass_factor_nonuniform: bool = False, +) -> WavefunctionGridObservable: + """Bind the spectral wavefunction-grid observable to one compiled solver. + + Parameters + ---------- + blocks + ``N_b`` symmetry blocks of ``N_c`` channels each. + energies + Compile-time energy grid, shape ``(N_E,)``. + sources + Block-stacked baked source stack, shape ``(N_b, N_E, N_c, M)``. + mass_factor_grid + Reduced ``(N_E,)`` per-energy ℏ²/2μ, or ``None`` for a constant μ. + block_mode + ``True`` for a ``blocks=`` compile. + mass_factor_nonuniform + ``True`` when μ genuinely varies with energy — the static regime then + raises (C4) and only the energy-batched regime is served. + """ + + mass_factor = uniform_block_mass_factor(blocks, context="spectral observable path") + _mfg = ( + jnp.full(len(energies), mass_factor) + if mass_factor_grid is None + else jnp.asarray(mass_factor_grid) + ) + return _WavefunctionGridObservable( + sources=sources, + energies=energies, + mass_factor=mass_factor, + mass_factor_grid=_mfg, + n_energies=len(energies), + block_mode=block_mode, + mass_factor_nonuniform=mass_factor_nonuniform, + ) + + def _rmatrix( spectrum: Spectrum, energy: float | jax.Array, @@ -701,6 +869,56 @@ def one_block(spectrum: Spectrum, source: jax.Array) -> jax.Array: return jax.vmap(one_block)(spectra, sources) +def _wavefunction_grid_aligned( + spectra: Spectrum, + energies: jax.Array, + sources: jax.Array, + mass_factor_grid: jax.Array, +) -> jax.Array: + """Return aligned-grid wavefunctions ``ψ(E_i; spec_i)``, shape ``(N_E, K, M)``. + + The energy-batched regime: one Green's contraction per aligned + ``(spectrum_e, E_e, source_e, μ_e)`` tuple — the ``_smatrix_grid`` + pattern. + """ + + def one_energy( + spectrum: Spectrum, energy: jax.Array, source: jax.Array, mu: jax.Array + ) -> jax.Array: + greens = greens_from_spectrum(spectrum, energy=energy, mass_factor=mu) + return jnp.einsum("mn,cn->cm", greens, source) + + return jax.vmap(one_energy)(spectra, energies, sources, mass_factor_grid) + + +def _wavefunction_grid_static_blocks( + spectra: Spectrum, + energies: jax.Array, + sources: jax.Array, + mass_factor: float, +) -> jax.Array: + """Return per-block static-regime grid wavefunctions, shape ``(N_b, N_E, K, M)``.""" + + def one_block(spectrum: Spectrum, source_stack: jax.Array) -> jax.Array: + return wavefunction_grid_from_spectrum(spectrum, energies, source_stack, mass_factor) + + return jax.vmap(one_block)(spectra, sources) + + +def _wavefunction_grid_aligned_blocks( + spectra: Spectrum, + energies: jax.Array, + sources: jax.Array, + mass_factor_grid: jax.Array, +) -> jax.Array: + """Return per-block aligned-grid wavefunctions, shape ``(N_b, N_E, K, M)``.""" + + def one_block(spectra_b: Spectrum, source_stack: jax.Array) -> jax.Array: + return _wavefunction_grid_aligned(spectra_b, energies, source_stack, mass_factor_grid) + + return jax.vmap(one_block)(spectra, sources) + + def _rmatrix_grid_blocks( spectra: Spectrum, energies: jax.Array, @@ -778,6 +996,11 @@ def one_block(spectra_b: Spectrum, boundary_b: BoundaryValues) -> jax.Array: _phases_grid_blocks, static_argnames=("channel_radius",), ) +_WAVEFUNCTION_GRID_STATIC_BLOCKS_JIT = jax.jit( + _wavefunction_grid_static_blocks, + static_argnames=("mass_factor",), +) +_WAVEFUNCTION_GRID_ALIGNED_BLOCKS_JIT = jax.jit(_wavefunction_grid_aligned_blocks) def _match_one_energy( diff --git a/src/lax/solvers/spectrum.py b/src/lax/solvers/spectrum.py index 56fbfef..a0edd45 100644 --- a/src/lax/solvers/spectrum.py +++ b/src/lax/solvers/spectrum.py @@ -59,13 +59,19 @@ class _SpectrumKernel: keep_eigenvectors: bool n_blocks: int block_mode: bool + mass_factor_grid: jax.Array | None = None # (N_E,) per-energy ℏ²/2μ, or None + mass_factor_nonuniform: bool = False def __call__(self, potential: jax.Array | Interaction) -> Spectrum: """Return the spectral decomposition(s) for one potential. Energy-dependent interactions add an ``(N_E,)`` axis after the block axis (block × energy order); block-independent interactions broadcast - across the ``(N_b,)`` axis. + across the ``(N_b,)`` axis. A **non-uniform** ``mass_factor_grid`` + forces the energy-batched path even for an energy-independent + potential: the dimensionless ``H/μ(E)`` varies with E, so + diagonalize-once is invalid physics — the per-energy spectra carry the + correct ``μ_e`` and the aligned-grid observables consume them. """ if not isinstance(potential, Interaction): @@ -89,20 +95,52 @@ def __call__(self, potential: jax.Array | Interaction) -> Spectrum: msg = f"Method {self.method!r} is not implemented in the MVP spectrum kernel." raise ValueError(msg) block = lift_to_blocks(potential.block, potential.block_dependent, self.n_blocks) - jit_fn = grid_jit if potential.energy_dependent else blocks_jit - spectrum = cast( - Spectrum, - jit_fn( - block, - self.mesh, - self.operators, - self.centrifugal, - self.thresholds, - self.q, - self.mass_factor, - self.keep_eigenvectors, - ), - ) + use_grid = potential.energy_dependent or self.mass_factor_nonuniform + if use_grid: + if self.mass_factor_grid is not None: + grid_mass_factors = self.mass_factor_grid + n_e = grid_mass_factors.shape[0] + if potential.energy_dependent and block.shape[1] != n_e: + msg = ( + f"energy-dependent Interaction carries {block.shape[1]} energies " + f"but the compile-time grid has N_E = {n_e}." + ) + raise ValueError(msg) + else: + n_e = block.shape[1] + grid_mass_factors = jnp.full(n_e, self.mass_factor) + if not potential.energy_dependent: + block = jnp.broadcast_to( + block[:, None], + (block.shape[0], n_e, block.shape[-2], block.shape[-1]), + ) + spectrum = cast( + Spectrum, + grid_jit( + block, + self.mesh, + self.operators, + self.centrifugal, + self.thresholds, + self.q, + grid_mass_factors, + self.keep_eigenvectors, + ), + ) + else: + spectrum = cast( + Spectrum, + blocks_jit( + block, + self.mesh, + self.operators, + self.centrifugal, + self.thresholds, + self.q, + self.mass_factor, + self.keep_eigenvectors, + ), + ) return spectrum if self.block_mode else take_block0(spectrum) @@ -113,6 +151,8 @@ def make_spectrum_kernel( method: Method = "eigh", keep_eigenvectors: bool = False, block_mode: bool = False, + mass_factor_grid: jax.Array | None = None, + mass_factor_nonuniform: bool = False, ) -> SpectrumKernel: """Build a JIT-compiled ``spectrum(V) → Spectrum`` kernel. @@ -140,6 +180,15 @@ def make_spectrum_kernel( block_mode ``True`` for a ``blocks=`` compile: outputs keep the leading ``(N_b,)`` axis and block-dependent Interactions are accepted. + mass_factor_grid + Per-energy ℏ²/2μ in MeV·fm², reduced to shape ``(N_E,)`` (per-channel + uniformity is validated upstream). Used by the energy-batched path so + each per-energy Hamiltonian is scaled by its own ``μ_e``. + mass_factor_nonuniform + ``True`` when ``mass_factor_grid`` genuinely varies with energy. + Forces ``spectrum(V)`` onto the energy-batched path even for an + energy-independent ``V`` — ``H/μ(E)`` changes with E, so a single + diagonalization would be silently wrong physics. Returns ------- @@ -149,8 +198,9 @@ def make_spectrum_kernel( Notes ----- The spectral path folds a single ℏ²/2μ out of the Hamiltonian, so it - requires one uniform mass factor across all channels of all blocks; - per-block/per-channel μ remains a direct-path feature. + requires one uniform mass factor across all channels of all blocks + (per energy, when ``mass_factor_grid`` is given); per-block/per-channel + μ remains a direct-path feature. """ mass_factor = uniform_block_mass_factor(blocks, context="spectral eigensolve path") @@ -167,6 +217,8 @@ def make_spectrum_kernel( keep_eigenvectors=keep_eigenvectors, n_blocks=len(blocks), block_mode=block_mode, + mass_factor_grid=mass_factor_grid, + mass_factor_nonuniform=mass_factor_nonuniform, ) @@ -279,15 +331,17 @@ def _spectrum_blocks_grid( centrifugal: jax.Array, thresholds: jax.Array, q: jax.Array, - mass_factor: float, + mass_factor_grid: jax.Array, keep_eigenvectors: bool, ) -> Spectrum: """Nested block × energy vmap of one spectrum core. - ``blocks`` has shape ``(N_b, N_E, M, M)``; the returned :class:`Spectrum` - leaves carry ``(N_b, N_E, …)`` axes. Under ``method="eig"`` the host - callback (``vmap_method="sequential"``) runs one host ``eig`` per - ``(b, i)`` pair — correct but sequential. + ``blocks`` has shape ``(N_b, N_E, M, M)`` and ``mass_factor_grid`` is the + aligned ``(N_E,)`` per-energy ℏ²/2μ — each per-energy Hamiltonian is + scaled by its own ``μ_e`` (shared across blocks). The returned + :class:`Spectrum` leaves carry ``(N_b, N_E, …)`` axes. Under + ``method="eig"`` the host callback (``vmap_method="sequential"``) runs one + host ``eig`` per ``(b, i)`` pair — correct but sequential. """ def one_block( @@ -295,7 +349,7 @@ def one_block( centrifugal_row: jax.Array, threshold_row: jax.Array, ) -> Spectrum: - def one_energy(block: jax.Array) -> Spectrum: + def one_energy(block: jax.Array, mass_factor: jax.Array) -> Spectrum: return core( block, mesh, @@ -307,7 +361,7 @@ def one_energy(block: jax.Array) -> Spectrum: keep_eigenvectors, ) - return jax.vmap(one_energy)(block_grid) + return jax.vmap(one_energy)(block_grid, mass_factor_grid) return jax.vmap(one_block)(blocks, centrifugal, thresholds) diff --git a/src/lax/spectral/observables.py b/src/lax/spectral/observables.py index 7a576ca..ad6c82a 100644 --- a/src/lax/spectral/observables.py +++ b/src/lax/spectral/observables.py @@ -133,6 +133,57 @@ def wavefunction_internal_from_spectrum( return values +def wavefunction_grid_from_spectrum( + spectrum: Spectrum, + energies: jax.Array, + sources: jax.Array, + mass_factor: float | jax.Array, +) -> jax.Array: + """Return internal wavefunctions for every grid energy from one spectrum. + + The fused static-regime evaluation: ``ψ(E_i) = G(E_i)·source_i`` for the + whole compile-time grid in two einsums, without materializing any + ``(M, M)`` Green's matrix — only the spectral denominator and the source + change per energy ("diagonalize once, evaluate many energies"). + + Parameters + ---------- + spectrum + Stored eigensystem of the Bloch-augmented Hamiltonian (single block, + eigenvectors retained). + energies + Compile-time energy grid in MeV, shape ``(N_E,)``. + sources + Per-energy, per-incoming-channel source stack, shape ``(N_E, K, M)`` + (``K`` incoming channels; see :func:`lax.make_wavefunction_source_grid`). + mass_factor + Conversion factor ``ℏ² / 2μ`` in MeV·fm² (uniform — the static regime + is invalid for a non-uniform μ(E); see the C4 contract). + + Returns + ------- + jax.Array + Internal wavefunction coefficients, shape ``(N_E, K, M)``. + + Raises + ------ + ValueError + If the spectrum was compiled without eigenvectors. + """ + + eigenvectors = spectrum.eigenvectors + if eigenvectors is None: + msg = "Wavefunction evaluation requires stored eigenvectors." + raise ValueError(msg) + denominator = 1.0 / ( + spectrum.eigenvalues[None, :] - (energies / mass_factor)[:, None] + ) # (N_E, M) + transpose = _spectral_metric_transpose(spectrum, eigenvectors) + projected = jnp.einsum("km,ecm->eck", transpose, sources) + values: jax.Array = jnp.einsum("mk,ek,eck->ecm", eigenvectors, denominator, projected) + return values + + def _spectral_denominator( spectrum: Spectrum, energy: float | jax.Array, @@ -154,5 +205,6 @@ def _spectral_metric_transpose(spectrum: Spectrum, eigenvectors: jax.Array) -> j __all__ = [ "greens_from_spectrum", "rmatrix_from_spectrum", + "wavefunction_grid_from_spectrum", "wavefunction_internal_from_spectrum", ] diff --git a/src/lax/transforms/__init__.py b/src/lax/transforms/__init__.py index 7ffb3ee..a16b671 100644 --- a/src/lax/transforms/__init__.py +++ b/src/lax/transforms/__init__.py @@ -1,5 +1,6 @@ """Transform builders for radial-grid and momentum-space projections.""" +from lax.transforms.bilinear import make_matrix_element, matrix_element from lax.transforms.fourier import compute_F_momentum, make_double_fourier, make_fourier from lax.transforms.grid import compute_B_grid, make_to_grid from lax.transforms.integration import make_integration @@ -10,5 +11,7 @@ "make_double_fourier", "make_fourier", "make_integration", + "make_matrix_element", "make_to_grid", + "matrix_element", ] diff --git a/src/lax/transforms/bilinear.py b/src/lax/transforms/bilinear.py new file mode 100644 index 0000000..e88d3c7 --- /dev/null +++ b/src/lax/transforms/bilinear.py @@ -0,0 +1,300 @@ +"""Two-state bilinear matrix elements in the mesh basis. + +``matrix_element`` evaluates ⟨bra|O|ket⟩-style contractions between two +*different* mesh-coefficient vectors, optionally without complex conjugation +(the bilinear form natural to complex-symmetric spectra), batched over the +symmetry-block and energy axes. It complements :mod:`lax.transforms.integration`, +which remains the single-state, always-conjugating "norms and expectation +values" helper. + +**Operator scaling contract** (the most important fact in this module): in the +Lagrange-mesh Gauss approximation the mesh coefficients already absorb +``√(λᵢ a)``, so + +* a *local* operator enters as the **unscaled** node values ``V(rᵢ)`` — + shape ``(M,)`` — and the element is the plain node sum + ``Σᵢ braᵢ · V(rᵢ) · ketᵢ``; +* a *non-local* kernel supplied as a bare ``(M, M)`` array must be + **Gauss-scaled** by the caller, ``K̃ᵢⱼ = √(λᵢ λⱼ)·a·K(rᵢ, rⱼ)`` — never the + raw kernel values; +* an :class:`~lax.Interaction` is the recommended operator form: its assembled + ``block`` already carries exactly this scaling for both local-diagonal and + non-local content, so no caller-side weighting can go wrong. + +``matrix_element`` adds **no** factors of ``a``, ``k``, or quadrature weights +beyond what is inside the operator; the result is in ``MeV·[ψ]²`` where +``[ψ]`` is the normalization of the supplied coefficient vectors. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import cast + +import jax +import jax.numpy as jnp + +from lax.types import Interaction, MatrixElementHelper, lift_block + + +def _me_overlap(bra: jax.Array, ket: jax.Array, conjugate: bool) -> jax.Array: + """Overlap ``Σᵢ braᵢ ketᵢ`` batched over broadcast leading axes.""" + + b = jnp.conj(bra) if conjugate else bra + result: jax.Array = jnp.einsum("...m,...m->...", b, ket) + return result + + +def _me_diagonal(bra: jax.Array, operator: jax.Array, ket: jax.Array, conjugate: bool) -> jax.Array: + """Diagonal/local element ``Σᵢ braᵢ V(rᵢ) ketᵢ`` over broadcast leading axes.""" + + b = jnp.conj(bra) if conjugate else bra + result: jax.Array = jnp.einsum("...m,...m,...m->...", b, operator, ket) + return result + + +def _me_matrix(bra: jax.Array, operator: jax.Array, ket: jax.Array, conjugate: bool) -> jax.Array: + """Full-matrix element ``braᵀ O ket`` over broadcast leading axes.""" + + b = jnp.conj(bra) if conjugate else bra + result: jax.Array = jnp.einsum("...m,...mn,...n->...", b, operator, ket) + return result + + +_ME_OVERLAP_JIT = jax.jit(_me_overlap, static_argnames=("conjugate",)) +_ME_DIAGONAL_JIT = jax.jit(_me_diagonal, static_argnames=("conjugate",)) +_ME_MATRIX_JIT = jax.jit(_me_matrix, static_argnames=("conjugate",)) + + +def matrix_element( + bra: jax.Array, + ket: jax.Array, + operator: jax.Array | Interaction | None = None, + *, + conjugate: bool, +) -> jax.Array: + """Bilinear matrix element on raw arrays with NumPy-style broadcasting. + + The standalone form: no solver-aware axis interpretation or validation is + performed — callers align their own leading batch axes (size-1 axes + broadcast; the contraction is over the trailing ``M`` axes only). The + solver-bound ``solver.matrix_element`` adds the canonical block × energy + interpretation, the deterministic rank-2 rule, and shape validation. + + Parameters + ---------- + bra, ket + Mesh-coefficient arrays, shape ``(..., M)``. + operator + ``None`` (overlap), unscaled local node values ``(M,)``, a + Gauss-scaled kernel matrix ``(..., M, M)`` (the trailing two axes are + contracted), or an :class:`~lax.Interaction` (its ``block`` is used + directly — already correctly scaled; see the module docstring). + A *batched diagonal* operator is not expressible as a bare array here + — pass it as diagonal-embedded matrices or an ``Interaction``. + conjugate + Keyword-required. ``False`` computes the non-conjugated bilinear + ``braᵀ·O·ket``; ``True`` computes ``bra†·O·ket``. + + Returns + ------- + jax.Array + The broadcast batch shape of the inputs' leading axes — a scalar for + unbatched inputs. + """ + + if isinstance(operator, Interaction): + operator = operator.block + if operator is None: + return cast(jax.Array, _ME_OVERLAP_JIT(bra, ket, conjugate=conjugate)) + operator = jnp.asarray(operator) + if operator.ndim == 1: + return cast(jax.Array, _ME_DIAGONAL_JIT(bra, operator, ket, conjugate=conjugate)) + return cast(jax.Array, _ME_MATRIX_JIT(bra, operator, ket, conjugate=conjugate)) + + +@dataclass(frozen=True) +class _MatrixElementHelper: + """Pickle-safe solver-bound bilinear matrix-element helper. + + Interprets bare-array inputs in the canonical **block × energy** axis + order with a deterministic, mode-based rank-2 rule (never shape-sniffing): + in blocks mode rank-2 is always block-leading ``(N_b, M)``; in channels + mode it is always energy-leading ``(N_E, M)``. Internally every state is + lifted to rank-3 ``(N_b|1, N_E|1, M)`` so block-leading inputs broadcast + on the block axis rather than right-aligning onto the energy axis. + """ + + block_mode: bool + n_blocks: int + n_energies: int | None + matrix_size: int + + def __call__( + self, + bra: jax.Array, + ket: jax.Array, + operator: jax.Array | Interaction | None = None, + *, + conjugate: bool, + ) -> jax.Array: + """Evaluate the bilinear element per the F1 contract. + + See :func:`matrix_element` for the operator forms and the scaling + contract. Output shape is the broadcast of the batch axes the inputs + actually carry: scalar for unbatched inputs, ``(N_b,)`` / + ``(N_E,)`` / ``(N_b, N_E)`` otherwise. + """ + + bra3, bra_has_block, bra_has_energy = self._normalize_state(bra, "bra") + ket3, ket_has_block, ket_has_energy = self._normalize_state(ket, "ket") + has_block = bra_has_block or ket_has_block + has_energy = bra_has_energy or ket_has_energy + + if operator is None: + result = cast(jax.Array, _ME_OVERLAP_JIT(bra3, ket3, conjugate=conjugate)) + elif isinstance(operator, Interaction): + has_block = has_block or operator.block_dependent + has_energy = has_energy or operator.energy_dependent + block = lift_block( + operator.block, + operator.energy_dependent, + operator.block_dependent, + True, + True, + ) + if block.shape[-2:] != (self.matrix_size, self.matrix_size): + msg = ( + f"operator Interaction block has trailing shape {block.shape[-2:]}, " + f"expected (M, M) = ({self.matrix_size}, {self.matrix_size}); " + "the Interaction must be assembled by this solver's builders." + ) + raise ValueError(msg) + result = cast(jax.Array, _ME_MATRIX_JIT(bra3, block, ket3, conjugate=conjugate)) + else: + operator = jnp.asarray(operator) + if operator.ndim == 1: + if operator.shape != (self.matrix_size,): + msg = ( + f"diagonal operator must have shape (M,) = ({self.matrix_size},), " + f"got {operator.shape}; values are the unscaled V(r_i) node samples." + ) + raise ValueError(msg) + result = cast( + jax.Array, _ME_DIAGONAL_JIT(bra3, operator, ket3, conjugate=conjugate) + ) + elif operator.ndim == 2: + if operator.shape != (self.matrix_size, self.matrix_size): + msg = ( + f"matrix operator must have shape (M, M) = " + f"({self.matrix_size}, {self.matrix_size}), got {operator.shape}." + ) + raise ValueError(msg) + result = cast(jax.Array, _ME_MATRIX_JIT(bra3, operator, ket3, conjugate=conjugate)) + else: + msg = ( + f"bare-array operators must be (M,) or (M, M), got rank {operator.ndim}; " + "batched (block-/energy-dependent) operators are passed as an Interaction." + ) + raise ValueError(msg) + + # Squeeze the axes no input contributed: result is (N_b|1, N_E|1). + if not has_energy: + result = jnp.squeeze(result, axis=-1) + if not has_block: + result = jnp.squeeze(result, axis=0) + return result + + def _normalize_state(self, values: jax.Array, name: str) -> tuple[jax.Array, bool, bool]: + """Lift a state array to rank-3 ``(N_b|1, N_E|1, M)`` per the §3.3 rule.""" + + values = jnp.asarray(values) + if values.ndim == 0 or values.shape[-1] != self.matrix_size: + msg = ( + f"{name} must have trailing dimension M = {self.matrix_size}, " + f"got shape {values.shape}." + ) + raise ValueError(msg) + if values.ndim == 1: + return values[None, None, :], False, False + if values.ndim == 2: + if self.block_mode: + if values.shape[0] != self.n_blocks: + msg = ( + f"rank-2 {name} in blocks mode is block-leading: expected " + f"(N_b, M) = ({self.n_blocks}, {self.matrix_size}), got " + f"{values.shape}; an energy-batched, block-broadcast input " + "must be written (1, N_E, M)." + ) + raise ValueError(msg) + return values[:, None, :], True, False + if self.n_energies is None: + msg = ( + f"rank-2 {name} in channels mode is energy-leading (N_E, M), but " + "this solver was compiled without an energy grid; pass (M,) or " + "recompile with energies=." + ) + raise ValueError(msg) + if values.shape[0] != self.n_energies: + msg = ( + f"rank-2 {name} in channels mode is energy-leading: expected " + f"(N_E, M) = ({self.n_energies}, {self.matrix_size}), got {values.shape}." + ) + raise ValueError(msg) + return values[None, :, :], False, True + if values.ndim == 3: + n_b, n_e = values.shape[0], values.shape[1] + if n_b not in (1, self.n_blocks): + msg = ( + f"rank-3 {name} leading block axis must be 1 or N_b = {self.n_blocks}, " + f"got shape {values.shape}." + ) + raise ValueError(msg) + allowed_e = (1,) if self.n_energies is None else (1, self.n_energies) + if n_e not in allowed_e: + msg = ( + f"rank-3 {name} energy axis must be one of {allowed_e}, " + f"got shape {values.shape}." + ) + raise ValueError(msg) + return values, True, True + msg = ( + f"{name} must be rank 1-3 with trailing dimension M = {self.matrix_size}, " + f"got shape {values.shape}." + ) + raise ValueError(msg) + + +def make_matrix_element( + *, + matrix_size: int, + n_blocks: int, + n_energies: int | None, + block_mode: bool, +) -> MatrixElementHelper: + """Return the solver-bound batched bilinear matrix-element helper. + + Depends only on compile-time shapes (never on mesh data, boundary values, + energies, or method), so it is bound unconditionally in every compile mode. + + Parameters + ---------- + matrix_size + ``M = N_c · N``, the per-block coefficient-vector length. + n_blocks + ``N_b`` (1 in channels mode). + n_energies + ``N_E``, or ``None`` when no energy grid was compiled. + block_mode + Whether the solver was compiled with ``blocks=``. + """ + + return _MatrixElementHelper( + block_mode=block_mode, + n_blocks=n_blocks, + n_energies=n_energies, + matrix_size=matrix_size, + ) + + +__all__ = ["make_matrix_element", "matrix_element"] diff --git a/src/lax/transforms/fourier.py b/src/lax/transforms/fourier.py index 790eae6..6e4797c 100644 --- a/src/lax/transforms/fourier.py +++ b/src/lax/transforms/fourier.py @@ -16,27 +16,51 @@ @dataclass(frozen=True) class _FourierProjection: - """Pickle-safe Fourier-Bessel transform.""" + """Pickle-safe Fourier-Bessel transform. + + ``fourier_matrices`` is ``(N_c, M_k, N)`` for a ``channels=`` compile and + ``(N_b, N_c, M_k, N)`` for a ``blocks=`` compile (DESIGN.md §15.5); the + stack rank selects the dispatch. + """ fourier_matrices: jax.Array def __call__(self, values: jax.Array, channel_index: int = 0) -> jax.Array: """Project mesh coefficients or kernels onto the momentum grid.""" - if values.ndim == 1: + if self.fourier_matrices.ndim == 3: + if values.ndim == 1: + return cast( + jax.Array, + _FOURIER_VECTOR_JIT(values, self.fourier_matrices, channel_index), + ) + if values.ndim == 2: + return cast( + jax.Array, + _FOURIER_MATRIX_JIT(values, self.fourier_matrices, channel_index), + ) + msg = f"fourier expects a (N,) vector or (N, N) kernel, got shape {values.shape}." + raise ValueError(msg) + values = _lift_block_values(values, self.fourier_matrices, "fourier") + if values.ndim == 2: return cast( jax.Array, - _FOURIER_VECTOR_JIT(values, self.fourier_matrices, channel_index), + _FOURIER_VECTOR_BLOCKS_JIT(values, self.fourier_matrices, channel_index), ) return cast( jax.Array, - _FOURIER_MATRIX_JIT(values, self.fourier_matrices, channel_index), + _FOURIER_MATRIX_BLOCKS_JIT(values, self.fourier_matrices, channel_index), ) @dataclass(frozen=True) class _DoubleFourierProjection: - """Pickle-safe double Fourier-Bessel transform.""" + """Pickle-safe double Fourier-Bessel transform. + + Blocks-mode stacks (rank-4 ``fourier_matrices``) accept ``(N_b, N, N)`` + block-batched kernels and broadcast unbatched ``(N, N)`` kernels across + the block axis. + """ fourier_matrices: jax.Array @@ -48,16 +72,42 @@ def __call__( ) -> jax.Array: """Project a mesh-space kernel onto left/right momentum grids.""" - if values.ndim != 2: - msg = "double_fourier_transform expects a rank-2 mesh-space kernel." - raise ValueError(msg) - resolved_right_channel_index = ( left_channel_index if right_channel_index is None else right_channel_index ) + if self.fourier_matrices.ndim == 3: + if values.ndim != 2: + msg = "double_fourier_transform expects a rank-2 mesh-space kernel." + raise ValueError(msg) + return cast( + jax.Array, + _DOUBLE_FOURIER_JIT( + values, + self.fourier_matrices, + left_channel_index, + resolved_right_channel_index, + ), + ) + n_blocks = self.fourier_matrices.shape[0] + basis_size = self.fourier_matrices.shape[-1] + if values.ndim == 2: + if values.shape != (basis_size, basis_size): + msg = ( + "double_fourier_transform in blocks mode expects a (N, N) kernel " + f"(broadcast across blocks) or (N_b, N, N), got shape {values.shape}." + ) + raise ValueError(msg) + values = jnp.broadcast_to(values, (n_blocks, basis_size, basis_size)) + elif values.ndim != 3 or values.shape != (n_blocks, basis_size, basis_size): + msg = ( + "double_fourier_transform in blocks mode expects a (N, N) kernel " + f"(broadcast across blocks) or (N_b, N, N) = ({n_blocks}, {basis_size}, " + f"{basis_size}), got shape {values.shape}." + ) + raise ValueError(msg) return cast( jax.Array, - _DOUBLE_FOURIER_JIT( + _DOUBLE_FOURIER_BLOCKS_JIT( values, self.fourier_matrices, left_channel_index, @@ -66,6 +116,46 @@ def __call__( ) +def _lift_block_values(values: jax.Array, fourier_matrices: jax.Array, name: str) -> jax.Array: + """Lift unbatched inputs onto the block axis of a rank-4 transform stack. + + Deterministic rank-2 rule (shapes are static, so this is not value + sniffing): rank-2 with leading dimension ``N_b`` is **always** block-batched + vectors; otherwise it must be an ``(N, N)`` kernel, broadcast across + blocks. In the ``N_b == N`` corner a broadcast kernel must therefore be + written explicitly as ``jnp.broadcast_to(K, (N_b, N, N))``. + """ + + n_blocks = fourier_matrices.shape[0] + basis_size = fourier_matrices.shape[-1] + if values.ndim == 1: + if values.shape[0] != basis_size: + msg = f"{name} vector must have shape (N,) = ({basis_size},), got {values.shape}." + raise ValueError(msg) + return jnp.broadcast_to(values, (n_blocks, basis_size)) + if values.ndim == 2: + if values.shape[0] == n_blocks and values.shape[1] == basis_size: + return values + if values.shape == (basis_size, basis_size): + return jnp.broadcast_to(values, (n_blocks, basis_size, basis_size)) + msg = ( + f"{name} in blocks mode accepts (N_b, N) block-batched vectors, " + f"(N, N) kernels (broadcast across blocks), (N,) vectors, or " + f"(N_b, N, N) kernels with N_b = {n_blocks}, N = {basis_size}; " + f"got shape {values.shape}." + ) + raise ValueError(msg) + if values.ndim == 3 and values.shape == (n_blocks, basis_size, basis_size): + return values + msg = ( + f"{name} in blocks mode accepts (N_b, N) block-batched vectors, " + f"(N, N) kernels (broadcast across blocks), (N,) vectors, or " + f"(N_b, N, N) kernels with N_b = {n_blocks}, N = {basis_size}; " + f"got shape {values.shape}." + ) + raise ValueError(msg) + + def compute_F_momentum( mesh: Mesh, momenta: jax.Array, @@ -239,6 +329,40 @@ def _double_fourier( return result +def _fourier_vector_blocks( + values: jax.Array, fourier_matrices: jax.Array, channel_index: int +) -> jax.Array: + """Project block-batched coefficients `(N_b, N)` onto the momentum grid `(N_b, M_k)`.""" + + matrices = fourier_matrices[:, channel_index] + result: jax.Array = jnp.einsum("bkn,bn->bk", matrices, values) + return result + + +def _fourier_matrix_blocks( + values: jax.Array, fourier_matrices: jax.Array, channel_index: int +) -> jax.Array: + """Project block-batched kernels `(N_b, N, N)` onto the momentum grid `(N_b, M_k, M_k)`.""" + + matrices = fourier_matrices[:, channel_index] + result: jax.Array = jnp.einsum("bkn,bnm,blm->bkl", matrices, values, matrices) + return result + + +def _double_fourier_blocks( + values: jax.Array, + fourier_matrices: jax.Array, + left_channel_index: int, + right_channel_index: int, +) -> jax.Array: + """Project block-batched kernels onto left/right momentum grids per block.""" + + left = fourier_matrices[:, left_channel_index] + right = fourier_matrices[:, right_channel_index] + result: jax.Array = jnp.einsum("bkn,bnm,blm->bkl", left, values, right) + return result + + _FOURIER_VECTOR_JIT = jax.jit( _fourier_vector, static_argnames=("channel_index",), @@ -251,6 +375,18 @@ def _double_fourier( _double_fourier, static_argnames=("left_channel_index", "right_channel_index"), ) +_FOURIER_VECTOR_BLOCKS_JIT = jax.jit( + _fourier_vector_blocks, + static_argnames=("channel_index",), +) +_FOURIER_MATRIX_BLOCKS_JIT = jax.jit( + _fourier_matrix_blocks, + static_argnames=("channel_index",), +) +_DOUBLE_FOURIER_BLOCKS_JIT = jax.jit( + _double_fourier_blocks, + static_argnames=("left_channel_index", "right_channel_index"), +) __all__ = ["compute_F_momentum", "make_double_fourier", "make_fourier"] diff --git a/src/lax/transforms/grid.py b/src/lax/transforms/grid.py index c1d68f9..ca6339b 100644 --- a/src/lax/transforms/grid.py +++ b/src/lax/transforms/grid.py @@ -112,23 +112,26 @@ def make_to_grid( def _to_grid_vector(values: jax.Array, basis_grid: jax.Array) -> jax.Array: - """Project mesh coefficients `(N,)` onto the radial grid `(M_r,)`.""" + """Project mesh coefficients `(..., N)` onto the radial grid `(..., M_r)`. - result: jax.Array = basis_grid @ values + Arbitrary leading batch axes (symmetry-block, energy) pass through. + """ + + result: jax.Array = jnp.einsum("rn,...n->...r", basis_grid, values) return result def _from_grid_array(values: jax.Array, projection_matrix: jax.Array) -> jax.Array: - """Project sampled radial-grid values `(M_r,)` back to mesh coefficients `(N,)`.""" + """Project sampled radial-grid values `(..., M_r)` back to mesh coefficients `(..., N)`.""" - result: jax.Array = projection_matrix @ values + result: jax.Array = jnp.einsum("nr,...r->...n", projection_matrix, values) return result def _to_grid_matrix(values: jax.Array, basis_grid: jax.Array) -> jax.Array: - """Project a mesh-space kernel `(N, N)` onto the radial grid `(M_r, M_r)`.""" + """Project a mesh-space kernel `(..., N, N)` onto the radial grid `(..., M_r, M_r)`.""" - result: jax.Array = basis_grid @ values @ basis_grid.T + result: jax.Array = jnp.einsum("rn,...nm,sm->...rs", basis_grid, values, basis_grid) return result diff --git a/src/lax/types.py b/src/lax/types.py index 35b70ae..2d5d854 100644 --- a/src/lax/types.py +++ b/src/lax/types.py @@ -125,14 +125,14 @@ def __add__(self, other: object) -> Interaction: return NotImplemented energy_dependent = self.energy_dependent or other.energy_dependent block_dependent = self.block_dependent or other.block_dependent - b1 = _lift_block( + b1 = lift_block( self.block, self.energy_dependent, self.block_dependent, energy_dependent, block_dependent, ) - b2 = _lift_block( + b2 = lift_block( other.block, other.energy_dependent, other.block_dependent, @@ -151,7 +151,7 @@ def __radd__(self, other: object) -> Interaction: return NotImplemented -def _lift_block( +def lift_block( block: jax.Array, has_energy: bool, has_block: bool, @@ -318,6 +318,88 @@ def __call__(self, spectrum: Spectrum, energy: EnergyLike, source: jax.Array) -> ... +class WavefunctionGridObservable(Protocol): + """Callable for internal wavefunctions at every compile-time grid energy. + + Spectral path only (``method ∈ {"eigh", "eig"}``); under + ``method="linear_solve"`` use ``solver.wavefunction_direct_grid``. + Sources are baked at compile time (Descouvemont eq. 27); both evaluation + regimes are served by one entry point — the energy-batched regime is + detected from the ``Spectrum`` rank. + + **Normalization contract:** ψ solves ``(H − E/μ)ψ = φ(a)·H⁻(a)`` — the + internal solution driven by the boundary *value* of the incoming wave. + Engines that drive with the matched exterior *derivative* + ``(i/2)(H⁻′ − S·H⁺′)`` instead differ by the per-channel scalar + ``(i/2)(H⁻′ − S·H⁺′)/H⁻`` (the solution is linear in the driving + coefficient), plus ``k/√a`` if their coefficients live in ``s = k·r``. + The cross-engine acceptance test (``tests/acceptance``) pins this + relation to machine precision; conversions belong in the caller, never + here. + """ + + def __call__(self, spectrum: Spectrum, channel_index: int | None = 0) -> jax.Array: + """Evaluate internal wavefunctions for every ``(block, grid-energy)`` pair. + + Parameters + ---------- + spectrum + Static-V :class:`Spectrum` (diagonalize once, evaluate the grid) + or the energy-batched ``Spectrum`` produced by an + energy-dependent :class:`~lax.Interaction` (aligned per-energy + evaluation). Must have been produced with ``'wavefunction'`` + (or ``'greens'``) in ``solvers=`` so eigenvectors are stored. + channel_index + Incoming channel ``c`` within the block (default 0). ``None`` + returns all incoming channels on an extra ``N_c`` axis. + + Returns + ------- + jax.Array + Mesh-coefficient vectors, shape ``(N_E, M)`` — ``(N_b, N_E, M)`` + in blocks mode; with ``channel_index=None`` an ``N_c`` axis is + inserted before ``M``. Entries at grid energies where the + incoming channel is closed are mathematically defined (Whittaker + source) but are **not** scattering wavefunctions — slice by + ``solver.boundary.is_open`` if needed. + """ + ... + + +class WavefunctionDirectGridObservable(Protocol): + """Callable for direct-path wavefunctions at every compile-time grid energy. + + The ``method="linear_solve"`` companion of ``wavefunction_grid``: one + linear solve per ``(block, energy)`` against the baked source stack. + Fully differentiable, and inherits the direct path's per-channel μ + support. + """ + + def __call__( + self, + potential: jax.Array | Interaction, + channel_index: int | None = 0, + ) -> jax.Array: + """Solve ``C(E_i) ψ = source_i`` for every ``(block, grid-energy)`` pair. + + Parameters + ---------- + potential + An :class:`~lax.Interaction` (energy-/block-dependent dispatch is + handled transparently); raw arrays are rejected at runtime. + channel_index + Incoming channel ``c`` within the block (default 0); ``None`` + returns all incoming channels on an extra ``N_c`` axis. + + Returns + ------- + jax.Array + Mesh-coefficient vectors with the same shapes and closed-channel + contract as ``wavefunction_grid``. + """ + ... + + class EigenpairAccessor(Protocol): """Callable for raw access to eigenvalues and eigenvectors. @@ -557,6 +639,49 @@ def __call__( ... +class MatrixElementHelper(Protocol): + """Callable for two-state bilinear matrix elements in the mesh basis. + + Batched over the symmetry-block and energy axes; bound unconditionally on + every compiled solver (it depends only on compile-time shapes). + """ + + def __call__( + self, + bra: jax.Array, + ket: jax.Array, + operator: jax.Array | Interaction | None = None, + *, + conjugate: bool, + ) -> jax.Array: + """Evaluate ``braᵀ·O·ket`` (``conjugate=False``) or ``bra†·O·ket``. + + Parameters + ---------- + bra, ket + Mesh-coefficient arrays: ``(M,)``, rank-2 (block-leading + ``(N_b, M)`` in blocks mode, energy-leading ``(N_E, M)`` in + channels mode — a deterministic mode-based rule, never + shape-sniffing), or ``(N_b, N_E, M)`` with explicit size-1 axes. + operator + ``None`` (overlap), **unscaled** local node values ``(M,)``, a + caller-Gauss-scaled kernel ``(M, M)``, or an + :class:`~lax.Interaction` (recommended — its block is already + correctly scaled; static flags drive axis alignment). + conjugate + Keyword-required: ``False`` is the non-conjugated bilinear (the + DWBA form, matching the complex-symmetric spectral metric); + ``True`` is the Hermitian inner product. + + Returns + ------- + jax.Array + The broadcast batch shape of the inputs' contributed axes — + scalar, ``(N_b,)``, ``(N_E,)``, or ``(N_b, N_E)``. + """ + ... + + class Integrator(Protocol): """Callable for norms and expectation values in the mesh basis.""" @@ -729,7 +854,8 @@ class TransformMatrices: Also accessible as ``solver.grid_r``. F_momentum Fourier-Bessel transform matrices, one per channel, - shape ``(N_c, M_k, N)``. Used by the ``fourier`` callable. + shape ``(N_c, M_k, N)`` — ``(N_b, N_c, M_k, N)`` for a solver + compiled with ``blocks=`` (§15.5). Used by the ``fourier`` callable. momenta Momentum grid passed to :func:`lax.compile`, shape ``(M_k,)`` in fm⁻¹. Also accessible as ``solver.momenta``. @@ -792,6 +918,14 @@ class Solver: wavefunction ``(spectrum, E, source) → ψ_int`` — internal wavefunction; requires ``'wavefunction'`` in ``solvers=``. + wavefunction_grid + ``(spectrum, channel_index=0) → ψ`` — internal wavefunctions at every + compile-time grid energy (both evaluation regimes); requires + ``'wavefunction'`` in ``solvers=`` and an energy grid. + wavefunction_sources + Baked Descouvemont eq.-27 source stack ``(N_E, N_c, M)`` — + ``(N_b, N_E, N_c, M)`` in blocks mode — or ``None`` when no + wavefunction entry point was requested. eigh ``(spectrum) → (ε, U)`` — raw eigenpairs; raises if eigenvectors were not retained. @@ -806,6 +940,9 @@ class Solver: rmatrix_direct ``(V) → R`` — per-energy linear-solve R-matrix on the compile-time grid. + wavefunction_direct_grid + ``(V, channel_index=0) → ψ`` — direct-path wavefunctions at every + compile-time grid energy; bound whenever the direct path is active. **Transform helpers**: @@ -821,6 +958,9 @@ class Solver: ``(V, ...) → V(p, p')`` — double Bessel transform for kernels. integrate ``(c, operator=None) → ⟨ψ|O|ψ⟩`` — norms and expectation values. + matrix_element + ``(bra, ket, operator=None, *, conjugate) → braᵀ·O·ket`` — two-state + bilinear form, batched over block/energy axes; always bound. """ mesh: Mesh @@ -838,6 +978,8 @@ class Solver: phases: SpectrumObservable | None = None greens: GreenFunctionObservable | None = None wavefunction: WavefunctionObservable | None = None + wavefunction_grid: WavefunctionGridObservable | None = None + wavefunction_sources: jax.Array | None = None eigh: EigenpairAccessor | None = None rmatrix_grid: SpectrumGridObservable | None = None smatrix_grid: SpectrumGridObservable | None = None @@ -846,6 +988,7 @@ class Solver: smatrix_direct: SMatrixDirectObservable | None = None phases_direct: PhasesDirectObservable | None = None wavefunction_direct: WavefunctionDirectObservable | None = None + wavefunction_direct_grid: WavefunctionDirectGridObservable | None = None interaction_from_block: Callable[..., Any] | None = None interaction_from_array: Callable[..., Any] | None = None interaction_from_funcs: Callable[..., Any] | None = None @@ -857,6 +1000,7 @@ class Solver: fourier: FourierTransform | None = None double_fourier_transform: DoubleFourierTransform | None = None integrate: Integrator | None = None + matrix_element: MatrixElementHelper | None = None # ------------------------------------------------------------------ # Convenience properties @@ -882,6 +1026,7 @@ def __repr__(self) -> str: "phases", "greens", "wavefunction", + "wavefunction_grid", "eigh", "rmatrix_grid", "smatrix_grid", @@ -890,12 +1035,14 @@ def __repr__(self) -> str: "smatrix_direct", "phases_direct", "wavefunction_direct", + "wavefunction_direct_grid", ) _transform_names = ( "to_grid_vector", "to_grid_matrix", "fourier", "integrate", + "matrix_element", ) live = [n for n in _observable_names if getattr(self, n) is not None] transforms = [n for n in _transform_names if getattr(self, n) is not None] @@ -925,6 +1072,7 @@ def __repr__(self) -> str: "GridVectorTransform", "Integrator", "Interaction", + "MatrixElementHelper", "Mesh", "MeshFamily", "MeshSpec", @@ -940,6 +1088,8 @@ def __repr__(self) -> str: "SpectrumKernel", "SpectrumObservable", "TransformMatrices", + "WavefunctionDirectGridObservable", "WavefunctionDirectObservable", + "WavefunctionGridObservable", "WavefunctionObservable", ] diff --git a/src/lax/wavefunction.py b/src/lax/wavefunction.py index f73266e..d8d673c 100644 --- a/src/lax/wavefunction.py +++ b/src/lax/wavefunction.py @@ -6,17 +6,84 @@ ψ_int = G(E) · source where ``source`` encodes which incoming channel drives the reaction. This -module provides :func:`make_wavefunction_source` so users do not need to -construct the source vector by hand. +module provides :func:`make_wavefunction_source` (one ``(channel, energy)`` +pair) and :func:`make_wavefunction_source_grid` (the full compile-time stack), +both slices of one shared builder so the two can never drift apart. """ from __future__ import annotations +import jax import jax.numpy as jnp from lax.types import BoundaryValues, Mesh, Solver +def build_wavefunction_sources( + mesh: Mesh, + boundary: BoundaryValues, + n_channels: int, +) -> jax.Array: + """Build the full Descouvemont eq.-27 source stack from the boundary cache. + + ``sources[..., e, c, c·N:(c+1)·N] = φ_n(a) · H⁻[..., e, c]`` with every + other channel block zero — the boundary cache already holds ``H⁻`` for + every ``(energy, channel)`` pair, so the stack is fully determined at + compile time. + + Closed channels are **not** masked: where ``boundary.is_open`` is false, + ``H⁻`` is the Whittaker-decaying solution and the resulting source (and + wavefunction) is mathematically defined but is *not* a scattering + wavefunction. Callers slice by ``boundary.is_open`` if they need + open-channel-only results. + + Parameters + ---------- + mesh + Compiled single-interval mesh (supplies ``basis_at_boundary``). + boundary + Compile-time boundary values; ``H_minus`` has shape ``(N_E, N_c)`` — + ``(N_b, N_E, N_c)`` in blocks mode. + n_channels + ``N_c``, the channels per block. + + Returns + ------- + jax.Array + Complex source stack of shape ``(N_E, N_c, N_c·N)`` — + ``(N_b, N_E, N_c, N_c·N)`` in blocks mode — indexed + ``[..., energy, incoming_channel, coefficient]``. + """ + + phi_a = mesh.basis_at_boundary # (N,) + n = mesh.n + h_minus = boundary.H_minus # (..., N_E, N_c), complex + sources = jnp.zeros((*h_minus.shape, n_channels * n), dtype=h_minus.dtype) + for c in range(n_channels): + start = c * n + sources = sources.at[..., c, start : start + n].set(phi_a * h_minus[..., c][..., None]) + return sources + + +def _resolved_sources(solver: Solver, caller: str) -> jax.Array: + """Return the solver's baked source stack, building it on demand.""" + + boundary: BoundaryValues | None = solver.boundary + if boundary is None: + msg = f"solver.boundary is None — re-compile with an energy grid to use {caller}." + raise ValueError(msg) + if solver.wavefunction_sources is not None: + return solver.wavefunction_sources + if solver.mesh.propagation is not None: + msg = ( + "wavefunction sources are not supported on propagated multi-interval " + "meshes — the boundary basis differs per interval. Use a " + "single-interval mesh." + ) + raise NotImplementedError(msg) + return build_wavefunction_sources(solver.mesh, boundary, len(solver.channels)) + + def make_wavefunction_source( solver: Solver, channel_index: int, @@ -33,7 +100,10 @@ def make_wavefunction_source( source[c·N : (c+1)·N] = φ_n(a) · H⁻_c(E_i) (all other blocks zero) where ``φ_n(a)`` are the Lagrange-basis boundary values and ``H⁻`` is - the incoming Coulomb/Whittaker function at the channel radius. + the incoming Coulomb/Whittaker function at the channel radius. This is a + slice of the compile-time stack built by + :func:`build_wavefunction_sources` (baked on the solver when a + wavefunction entry point was requested, rebuilt on demand otherwise). Parameters ---------- @@ -81,37 +151,51 @@ def make_wavefunction_source( interaction = solver.interaction_from_block(V[0, 0]) # (M, M) block psi = solver.wavefunction_direct(interaction, src, energy_index=5) """ - boundary: BoundaryValues | None = solver.boundary - if boundary is None: - raise ValueError( - "solver.boundary is None — re-compile with an energy grid " - "to use make_wavefunction_source." - ) - mesh: Mesh = solver.mesh - n_c = len(solver.channels) - n = mesh.n + sources = _resolved_sources(solver, "make_wavefunction_source") + if solver.blocks is not None: + return sources[:, energy_index, channel_index] + return sources[energy_index, channel_index] - # φ_n(a) for all basis functions (N,) - phi_a = mesh.basis_at_boundary - if solver.blocks is not None: - # Blocks mode: H⁻ carries a leading (N_b,) axis; build one source per - # symmetry block, stacked on that axis. - h_minus_b = boundary.H_minus[:, energy_index, channel_index] # (N_b,) - n_b = h_minus_b.shape[0] - sources = jnp.zeros((n_b, n_c * n), dtype=jnp.complex128) - start = channel_index * n - return sources.at[:, start : start + n].set(phi_a[None] * h_minus_b[:, None]) +def make_wavefunction_source_grid( + solver: Solver, + channel_index: int | None = None, +) -> jnp.ndarray: + """Return the full compile-time source stack, or one incoming channel of it. - # H⁻ for the requested channel at the requested energy (scalar, complex) - h_minus_c = boundary.H_minus[energy_index, channel_index] + Parameters + ---------- + solver + Compiled solver bundle with an energy grid. + channel_index + ``None`` (default) returns the all-channels stack; an ``int`` slices + the corresponding incoming channel. - # Build the source vector: non-zero only in the block for channel_index - source = jnp.zeros(n_c * n, dtype=jnp.complex128) - start = channel_index * n - source = source.at[start : start + n].set(phi_a * h_minus_c) - return source + Returns + ------- + jnp.ndarray + ``channel_index=None``: shape ``(N_E, N_c, N_c·N)`` — + ``(N_b, N_E, N_c, N_c·N)`` in blocks mode. + With an ``int`` ``channel_index`` the incoming-channel axis is sliced + off: ``(N_E, N_c·N)`` / ``(N_b, N_E, N_c·N)``. + + Raises + ------ + ValueError + If ``solver.boundary`` is ``None`` (no energy grid was compiled). + """ + + sources = _resolved_sources(solver, "make_wavefunction_source_grid") + if channel_index is None: + return sources + if solver.blocks is not None: + return sources[:, :, channel_index] + return sources[:, channel_index] -__all__ = ["make_wavefunction_source"] +__all__ = [ + "build_wavefunction_sources", + "make_wavefunction_source", + "make_wavefunction_source_grid", +] diff --git a/tests/acceptance/data/dwba_pn_reference.npz b/tests/acceptance/data/dwba_pn_reference.npz new file mode 100644 index 0000000..0fc9618 Binary files /dev/null and b/tests/acceptance/data/dwba_pn_reference.npz differ diff --git a/tests/acceptance/generate_dwba_reference.py b/tests/acceptance/generate_dwba_reference.py new file mode 100644 index 0000000..5f3ee96 --- /dev/null +++ b/tests/acceptance/generate_dwba_reference.py @@ -0,0 +1,171 @@ +"""Generate the T5 cross-engine DWBA reference fixture from jitr. + +Runs jitr's ``xs.quasielastic_pn.Workspace`` on the ⁴⁸Ca(p,n)⁴⁸Sc IAS case +(the setup of jitr's ``examples/notebooks/chex_jitr_validation.ipynb``, whose +semi-relativistic variant is validated against CHEX and Jon et al. data) and +exports the per-(l, j) T-matrix elements, elastic S-matrices, sampled KDUQ +potentials, and kinematics to ``tests/acceptance/data/dwba_pn_reference.npz`` +for ``test_dwba_reference.py``. + +**Classical kinematics** are used (instead of the notebook's semi-relativistic +prescription) so that ``ħ²k²/2μ = Ecm`` and ``η = αZz·μ/(ħk)`` hold as +identities — the lax solver then reproduces each channel exactly from +``(Ecm, ħ²/2μ, z1z2)`` with no potential rescaling. The reference values +therefore differ ~1-2% from the CHEX-validated notebook numbers, but the +engine and workspace are identical. + +Run inside the jitr virtualenv:: + + cd ~/umich/jitr && .venv/bin/python /tests/acceptance/generate_dwba_reference.py +""" + +from __future__ import annotations + +import subprocess +from pathlib import Path + +import jitr +import numpy as np +from jitr.utils.kinematics import classical_kinematics, classical_kinematics_cm + +# --- the notebook case, classical kinematics -------------------------------- +PROTON = (1, 1) +NEUTRON = (1, 0) +CA48 = (48, 20) +SC48 = (48, 21) +E_LAB = 35.0 # MeV +E_IAS = 6.67 # MeV +CHANNEL_RADIUS_FM = 16.0 +LMAX = 20 +NBASIS = 35 + +OUTPUT = Path(__file__).parent / "data" / "dwba_pn_reference.npz" + + +def main() -> None: + reaction = jitr.reactions.Reaction( + target=CA48, projectile=PROTON, product=NEUTRON, residual=SC48 + ) + reaction_exit_channel = jitr.reactions.Reaction( + target=reaction.residual, projectile=reaction.product, process="El" + ) + + # Classical kinematics, mirroring Reaction.kinematics / kinematics_exit + # but with the non-relativistic constructors. + kin_p = classical_kinematics( + reaction.target.m0, + reaction.projectile.m0, + E_LAB, + Zz=reaction.target.Z * reaction.projectile.Z, + ) + ecm_n = kin_p.Ecm + reaction.Q - E_IAS + kin_n = classical_kinematics_cm( + reaction.residual.m0 + E_IAS, + reaction.product.m0, + ecm_n, + Zz=reaction.residual.Z * reaction.product.Z, # neutron: 0 + ) + + solver = jitr.rmatrix.Solver(nbasis=NBASIS) + workspace = jitr.xs.quasielastic_pn.Workspace( + reaction, + kin_p, + kin_n, + solver, + np.linspace(0.01, np.pi, 10), # angles are unused by tmatrix() + LMAX, + CHANNEL_RADIUS_FM, + tmatrix_abs_tol=1.0e-16, # no early partial-wave exit + ) + rgrid = workspace.radial_grid() + + # KDUQ default-parameter potentials sampled at the classical kinematics. + kd_p_params = np.array(list(jitr.optical_potentials.kduq.Global(PROTON).params.values())) + kd_n_params = np.array(list(jitr.optical_potentials.kduq.Global(NEUTRON).params.values())) + omp_p = jitr.optical_potentials.kduq.KDUQ(PROTON) + omp_n = jitr.optical_potentials.kduq.KDUQ(NEUTRON) + u_p_central, u_p_spin_orbit, u_p_coulomb = omp_p(rgrid, reaction, kin_p, *kd_p_params) + u_n_central, u_n_spin_orbit, _ = omp_n(rgrid, reaction_exit_channel, kin_n, *kd_n_params) + + tpn, s_n, s_p = workspace.tmatrix( + u_p_coulomb, + u_p_central, + u_p_spin_orbit, + u_n_central, + u_n_spin_orbit, + ) + + # Flatten the (lmax+1, 2) layout to N_pw = 2*lmax + 1 rows: l = 0 has only + # j = 1/2 (l.s = 0); l >= 1 carries (j = l+1/2, j = l-1/2) with the l.s + # eigenvalues the workspace itself precomputed. + rows: list[tuple[int, float, int]] = [(0, 0.0, 0)] # (l, l.s, ji) + for ell in range(1, LMAX + 1): + ldots = workspace.l_dot_s[ell - 1] + rows.append((ell, float(ldots[0]), 0)) + rows.append((ell, float(ldots[1]), 1)) + + ls = np.array([row[0] for row in rows], dtype=np.int64) + ldots = np.array([row[1] for row in rows], dtype=np.float64) + t_ref = np.array([tpn[row[0], row[2]] for row in rows]) + s_p_ref = np.array([s_p[row[0], row[2]] for row in rows]) + s_n_ref = np.array([s_n[row[0], row[2]] for row in rows]) + + # Per-row composed potentials (raw MeV node values, exactly as the + # workspace composes them; Coulomb enters the proton solve, not U1). + iso = workspace.isovector_factor + u_p_rows = u_p_central + u_p_coulomb + ldots[:, None] * u_p_spin_orbit + u_n_rows = u_n_central + ldots[:, None] * u_n_spin_orbit + u1_rows = -iso * ( + (u_n_central - u_p_central) + ldots[:, None] * (u_n_spin_orbit - u_p_spin_orbit) + ) + + try: + commit = subprocess.run( + ["git", "-C", str(Path(jitr.__file__).resolve().parents[3]), "rev-parse", "HEAD"], + capture_output=True, + text=True, + check=True, + ).stdout.strip() + except Exception: + commit = "unknown" + + OUTPUT.parent.mkdir(parents=True, exist_ok=True) + np.savez( + OUTPUT, + rgrid=np.asarray(rgrid), + ls=ls, + ldots=ldots, + channel_radius=CHANNEL_RADIUS_FM, + nbasis=NBASIS, + Ecm_p=kin_p.Ecm, + Ecm_n=kin_n.Ecm, + mu_p=kin_p.mu, + mu_n=kin_n.mu, + k_p=kin_p.k, + k_n=kin_n.k, + eta_p=kin_p.eta, + z1z2_p=np.array([reaction.projectile.Z, reaction.target.Z], dtype=np.int64), + U_p=np.broadcast_to(u_p_rows, (len(rows), NBASIS)).copy(), + U_n=np.broadcast_to(u_n_rows, (len(rows), NBASIS)).copy(), + U1=np.broadcast_to(u1_rows, (len(rows), NBASIS)).copy(), + T_ref=t_ref, + S_p=s_p_ref, + S_n=s_n_ref, + rtol=1.0e-8, + atol=1.0e-12, + provenance=( + f"jitr {getattr(jitr, '__version__', 'unknown')} commit {commit}; " + "48Ca(p,n)48Sc IAS, E_lab=35 MeV, E_IAS=6.67 MeV, KDUQ defaults, " + "lmax=20, a=16 fm, nbasis=35, CLASSICAL kinematics " + "(the semi-relativistic variant of this case is CHEX/data-validated; " + "see jitr examples/notebooks/chex_jitr_validation.ipynb)" + ), + ) + print(f"wrote {OUTPUT}") + print(f" kin_p: Ecm={kin_p.Ecm:.6f} k={kin_p.k:.6f} mu={kin_p.mu:.4f} eta={kin_p.eta:.6f}") + print(f" kin_n: Ecm={kin_n.Ecm:.6f} k={kin_n.k:.6f} mu={kin_n.mu:.4f} eta={kin_n.eta:.6f}") + print(f" |T| range: {np.abs(t_ref).min():.3e} .. {np.abs(t_ref).max():.3e}") + + +if __name__ == "__main__": + main() diff --git a/tests/acceptance/test_dwba_reference.py b/tests/acceptance/test_dwba_reference.py new file mode 100644 index 0000000..287e4e9 --- /dev/null +++ b/tests/acceptance/test_dwba_reference.py @@ -0,0 +1,182 @@ +"""T5 (spec v0.1.5.1): cross-engine DWBA normalization anchor against jitr. + +Reference data: jitr's ``xs.quasielastic_pn`` engine on the ⁴⁸Ca(p,n)⁴⁸Sc IAS +case (E_lab = 35 MeV, KDUQ defaults, lmax = 20, a = 16 fm, nbasis = 35) with +**classical kinematics**, exported by ``generate_dwba_reference.py`` (run it +inside the jitr virtualenv to regenerate). The semi-relativistic variant of +the same case is validated against CHEX and Jon et al. data in jitr's +``examples/notebooks/chex_jitr_validation.ipynb``. + +Classical kinematics make ``ħ²k²/2μ = Ecm`` and ``η = αZz·μ/(ħk)`` exact +identities, so each jitr channel maps onto lax with +``mass_factor = ħ²c²/(2μ)`` and ``energies = [Ecm]`` — no potential rescaling. + +Two stages: + +1. **S-matrix (normalization-free):** lax's per-(l, j) elastic S-matrices must + reproduce jitr's ``Sp``/``Sn`` directly — this validates the whole + convention mapping (grid, η, k, potential scale) independent of + wavefunction normalization. +2. **T-matrix (C7):** ``matrix_element(χp, χn, U₁)/(a·k_p·k_n)`` must + reproduce jitr's node-sum ``T_lj`` after two documented conversions: + + * **Source convention (per row):** jitr's internal wavefunction is driven + by the matched exterior *derivative*, + ``u_ext'(a) = (i/2)(H⁻' − S_lj·H⁺')`` (``rmatrix/core.py``, + ``solution_coeffs_with_inverse``), while lax's is driven by the + boundary *value* ``H⁻(a)``. The driven solution is linear in that + scalar, so each (l, j, channel) converts by the exact factor + ``(i/2)(H⁻' − S·H⁺')/H⁻`` — built here from lax's own boundary cache + and S-matrices. + * **Global constant (C7 proper):** what remains after the per-row + conversion is exactly ``k_p·k_n/a`` (verified to 4e-14): jitr's s = k·r + coefficients relate to lax's r-space coefficients per channel as + ``x_c = (k_c/√a)·[(i/2)(H⁻' − S·H⁺')/H⁻]·χ_c``, so + + T_jitr = conv_p·conv_n·matrix_element(χp, χn, U₁) / a² + + The factor is frozen analytically below and asserted constant across + every (l, j). +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest + +import lax +import lax.constants as constants + +FIXTURE = Path(__file__).parent / "data" / "dwba_pn_reference.npz" + +pytestmark = pytest.mark.skipif( + not FIXTURE.exists(), + reason="jitr DWBA reference fixture not yet generated (see generate_dwba_reference.py)", +) + + +def _compile_channel( + data: np.lib.npyio.NpzFile, + ecm: float, + mu: float, + z1z2: tuple[int, int] | None, +) -> lax.Solver: + import jax.numpy as jnp + + mass_factor = constants.HBARC**2 / (2.0 * mu) + return lax.compile( + mesh=lax.MeshSpec( + "legendre", "x", n=int(data["nbasis"]), scale=float(data["channel_radius"]) + ), + blocks=[ + (lax.ChannelSpec(l=int(ell), threshold=0.0, mass_factor=mass_factor),) + for ell in data["ls"] + ], + solvers=("spectrum", "smatrix", "wavefunction"), + energies=jnp.asarray([ecm]), + z1z2=z1z2, + V_is_complex=True, + method="eig", + ) + + +def test_matrix_element_reproduces_jitr_node_sum() -> None: + import jax.numpy as jnp + + data = np.load(FIXTURE) + radius = float(data["channel_radius"]) + k_p = float(data["k_p"]) + k_n = float(data["k_n"]) + n_pw = len(data["ls"]) + + # Classical-kinematics identity the mapping relies on. + for ecm, mu, k in ( + (float(data["Ecm_p"]), float(data["mu_p"]), k_p), + (float(data["Ecm_n"]), float(data["mu_n"]), k_n), + ): + assert np.isclose(constants.HBARC**2 / (2.0 * mu) * k**2, ecm, rtol=1e-10) + + proton = _compile_channel( + data, + float(data["Ecm_p"]), + float(data["mu_p"]), + (int(data["z1z2_p"][0]), int(data["z1z2_p"][1])), + ) + neutron = _compile_channel(data, float(data["Ecm_n"]), float(data["mu_n"]), None) + + # Grid and wavenumber identity with the jitr engine. + np.testing.assert_allclose( + np.asarray(proton.mesh.radii), np.asarray(data["rgrid"]), rtol=1e-12, atol=1e-12 + ) + assert proton.boundary is not None and neutron.boundary is not None + np.testing.assert_allclose(np.asarray(proton.boundary.k[:, 0, 0]), k_p, rtol=1e-10) + np.testing.assert_allclose(np.asarray(neutron.boundary.k[:, 0, 0]), k_n, rtol=1e-10) + + assert proton.interaction_from_array is not None + assert neutron.interaction_from_array is not None + coupling = np.ones((1, 1)) + v_p = proton.interaction_from_array( + local=[(jnp.asarray(data["U_p"]), coupling)], block_dependent=True + ) + v_n = neutron.interaction_from_array( + local=[(jnp.asarray(data["U_n"]), coupling)], block_dependent=True + ) + u1 = proton.interaction_from_array( + local=[(jnp.asarray(data["U1"]), coupling)], block_dependent=True + ) + + spectrum_p = proton.spectrum(v_p) # type: ignore[misc] + spectrum_n = neutron.spectrum(v_n) # type: ignore[misc] + + # ------------------------------------------------------------------ + # Stage 1 — normalization-free S-matrix regression per (l, j). + assert proton.smatrix is not None and neutron.smatrix is not None + s_p = np.asarray(proton.smatrix(spectrum_p))[:, 0, 0, 0] + s_n = np.asarray(neutron.smatrix(spectrum_n))[:, 0, 0, 0] + # Measured cross-engine agreement is ~8e-13 absolute; assert with margin. + np.testing.assert_allclose(s_p, np.asarray(data["S_p"]), rtol=1e-9, atol=1e-10) + np.testing.assert_allclose(s_n, np.asarray(data["S_n"]), rtol=1e-9, atol=1e-10) + + # ------------------------------------------------------------------ + # Stage 2 — DWBA bilinear vs jitr's node sum, up to one constant (C7). + assert proton.wavefunction_grid is not None + assert neutron.wavefunction_grid is not None + assert proton.matrix_element is not None + chi_p = proton.wavefunction_grid(spectrum_p) # (N_pw, 1, M) + chi_n = neutron.wavefunction_grid(spectrum_n) + element = proton.matrix_element(chi_p, chi_n, u1, conjugate=False) # (N_pw, 1) + t_lax = np.asarray(element)[:, 0] / (radius * k_p * k_n) + t_ref = np.asarray(data["T_ref"]) + + # Source-convention conversion per (l, j, channel): jitr drives with + # u_ext'(a) = (i/2)(H-' - S H+'); lax drives with H-(a). + def source_conversion(solver: lax.Solver, s_matrix: np.ndarray) -> np.ndarray: + assert solver.boundary is not None + h_minus = np.asarray(solver.boundary.H_minus)[:, 0, 0] + h_minus_p = np.asarray(solver.boundary.H_minus_p)[:, 0, 0] + h_plus_p = np.asarray(solver.boundary.H_plus_p)[:, 0, 0] + return 0.5j * (h_minus_p - s_matrix * h_plus_p) / h_minus + + t_lax = t_lax * source_conversion(proton, s_p) * source_conversion(neutron, s_n) + + # The frozen C7 constant: k_p·k_n/a (resolved 2026-06-12 from this + # fixture to a 4e-14 relative deviation; see the module docstring). + factor = k_p * k_n / radius + + significant = np.abs(t_ref) > 1e-5 * np.abs(t_ref).max() + assert significant.sum() > n_pw // 2 + ratio = t_ref[significant] / t_lax[significant] + np.testing.assert_allclose( + ratio, + factor, + rtol=float(data["rtol"]), + err_msg="cross-engine convention factor drifted from k_p*k_n/a — C7", + ) + np.testing.assert_allclose( + t_lax * factor, + t_ref, + rtol=float(data["rtol"]), + atol=float(data["atol"]), + ) diff --git a/tests/property/test_autograd.py b/tests/property/test_autograd.py index cfc9679..af81c09 100644 --- a/tests/property/test_autograd.py +++ b/tests/property/test_autograd.py @@ -73,3 +73,92 @@ def loss(V: jax.Array) -> jax.Array: grad = jax.grad(loss)(V) assert jnp.all(jnp.isfinite(grad)) + + +# --------------------------------------------------------------------------- +# T4 (spec v0.1.5.1) — gradients through the DWBA chain in miniature. +# Differentiable paths: eigh spectral (real V) and linear_solve direct +# (complex V). The eig path is asserted to raise (C8): its spectra flow +# through jax.pure_callback → host np.linalg.eig, which has no JVP/VJP. + +_WAVE_SOLVER = lm.compile( + mesh=lm.MeshSpec("legendre", "x", n=_N, scale=_SCALE), + channels=(lm.ChannelSpec(l=0, threshold=0.0, mass_factor=_HBAR2_2MU),), + solvers=("spectrum", "wavefunction"), + energies=_ENERGIES, +) +_DIRECT_SOLVER = lm.compile( + mesh=lm.MeshSpec("legendre", "x", n=_N, scale=_SCALE), + channels=(lm.ChannelSpec(l=0, threshold=0.0, mass_factor=_HBAR2_2MU),), + solvers=("rmatrix_direct",), + energies=_ENERGIES, + V_is_complex=True, + method="linear_solve", +) +_EIG_SOLVER = lm.compile( + mesh=lm.MeshSpec("legendre", "x", n=_N, scale=_SCALE), + channels=(lm.ChannelSpec(l=0, threshold=0.0, mass_factor=_HBAR2_2MU),), + solvers=("spectrum", "wavefunction"), + energies=_ENERGIES, + V_is_complex=True, + method="eig", +) + + +@pytest.mark.property +def test_wavefunction_grid_matrix_element_is_differentiable() -> None: + """jax.grad through eigh spectrum → wavefunction_grid → matrix_element.""" + assert _WAVE_SOLVER.spectrum is not None + assert _WAVE_SOLVER.wavefunction_grid is not None + assert _WAVE_SOLVER.matrix_element is not None + radii = _WAVE_SOLVER.mesh.radii + operator = jnp.exp(-0.1 * radii**2) # local transition operator at the nodes + + def loss(depth: jax.Array) -> jax.Array: + V = lm.Interaction( + block=jnp.diag(depth * jnp.exp(-((radii / 2.0) ** 2))), energy_dependent=False + ) + spec = _WAVE_SOLVER.spectrum(V) + psi = _WAVE_SOLVER.wavefunction_grid(spec) # (N_E, N) + element = _WAVE_SOLVER.matrix_element(psi, psi, operator, conjugate=False) + return jnp.sum(jnp.abs(element) ** 2) + + grad = jax.grad(loss)(jnp.asarray(-25.0)) + assert jnp.isfinite(grad) + + +@pytest.mark.property +def test_wavefunction_direct_grid_matrix_element_is_differentiable() -> None: + """jax.grad through linear_solve → wavefunction_direct_grid → matrix_element.""" + assert _DIRECT_SOLVER.wavefunction_direct_grid is not None + assert _DIRECT_SOLVER.matrix_element is not None + radii = _DIRECT_SOLVER.mesh.radii + operator = jnp.exp(-0.1 * radii**2) + + def loss(depth: jax.Array) -> jax.Array: + profile = (depth - 4.0j) * jnp.exp(-((radii / 2.0) ** 2)) + V = lm.Interaction(block=jnp.diag(profile), energy_dependent=False) + psi = _DIRECT_SOLVER.wavefunction_direct_grid(V) # (N_E, N) + element = _DIRECT_SOLVER.matrix_element(psi, psi, operator, conjugate=False) + return jnp.sum(jnp.abs(element) ** 2) + + grad = jax.grad(loss)(jnp.asarray(-25.0)) + assert jnp.isfinite(grad) + + +@pytest.mark.property +def test_eig_path_raises_on_differentiation() -> None: + """C8: the eig spectral path cannot be differentiated (pure_callback). + + Guarded so a future custom-JVP upgrade of the eig path is noticed. + """ + assert _EIG_SOLVER.spectrum is not None + radii = _EIG_SOLVER.mesh.radii + + def loss(depth: jax.Array) -> jax.Array: + profile = (depth - 4.0j) * jnp.exp(-((radii / 2.0) ** 2)) + V = lm.Interaction(block=jnp.diag(profile), energy_dependent=False) + return jnp.sum(jnp.abs(_EIG_SOLVER.spectrum(V).eigenvalues) ** 2) + + with pytest.raises(ValueError, match="do not support JVP"): + jax.grad(loss)(jnp.asarray(-25.0)) diff --git a/tests/unit/test_blocks_direct.py b/tests/unit/test_blocks_direct.py index 0cbb6c5..01c805e 100644 --- a/tests/unit/test_blocks_direct.py +++ b/tests/unit/test_blocks_direct.py @@ -246,14 +246,8 @@ def test_compile_validation_errors() -> None: solvers=DIRECT, energies=ENERGIES, ) - with pytest.raises(ValueError, match="not supported with `blocks="): - lax.compile( - mesh=MESH, - blocks=[group], - solvers=DIRECT, - energies=ENERGIES, - momenta=jnp.linspace(0.1, 2.0, 10), - ) + # momenta= with blocks= is supported since the F3 transform batching + # (see test_blocks_transforms.py); only propagated meshes still reject it. with pytest.raises(ValueError, match="not supported on propagated meshes"): lax.compile( mesh=lax.MeshSpec("legendre", "x", n=N, scale=RADIUS, extras={"n_intervals": 4}), diff --git a/tests/unit/test_blocks_transforms.py b/tests/unit/test_blocks_transforms.py new file mode 100644 index 0000000..6f95898 --- /dev/null +++ b/tests/unit/test_blocks_transforms.py @@ -0,0 +1,235 @@ +"""F3 acceptance tests: block-batched Fourier/grid transforms (spec v0.1.5.1, T10).""" + +from __future__ import annotations + +import importlib +from typing import Any + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +import lax +from tests.unit._blocks_helpers import HBAR2_2MU, partial_wave_groups + +N = 16 +RADIUS = 10.0 +ENERGIES = jnp.linspace(2.0, 30.0, 5) +MOMENTA = jnp.linspace(0.1, 2.0, 9) +GRID = jnp.linspace(0.4, RADIUS - 0.4, 7) +MESH = lax.MeshSpec("legendre", "x", n=N, scale=RADIUS) +BLOCK_GROUPS = partial_wave_groups() +N_B = len(BLOCK_GROUPS) + +TIGHT = dict(rtol=1e-10, atol=1e-12) + + +def _blocked_solver() -> lax.Solver: + return lax.compile( + mesh=MESH, + blocks=BLOCK_GROUPS, + solvers=("spectrum",), + energies=ENERGIES, + momenta=MOMENTA, + grid=GRID, + ) + + +def _single_solver(block: tuple[lax.ChannelSpec, ...]) -> lax.Solver: + return lax.compile( + mesh=MESH, + channels=block, + solvers=("spectrum",), + energies=ENERGIES, + momenta=MOMENTA, + grid=GRID, + ) + + +def test_compile_accepts_momenta_with_blocks() -> None: + """The former momenta × blocks rejection is lifted (old C1).""" + + solver = _blocked_solver() + assert solver.fourier is not None + assert solver.double_fourier_transform is not None + assert solver.transforms.F_momentum is not None + assert solver.transforms.F_momentum.shape == (N_B, 1, len(MOMENTA), N) + + +def test_blocked_fourier_matches_per_block_solvers() -> None: + """Blocked fourier/double_fourier ≡ per-block channels-mode solvers (§15.5).""" + + solver = _blocked_solver() + rng = np.random.default_rng(20) + vectors = jnp.asarray(rng.normal(size=(N_B, N))) + kernels = jnp.asarray(rng.normal(size=(N_B, N, N))) + + assert solver.fourier is not None + assert solver.double_fourier_transform is not None + batched_vec = solver.fourier(vectors) + batched_ker = solver.fourier(kernels) + batched_double = solver.double_fourier_transform(kernels) + assert batched_vec.shape == (N_B, len(MOMENTA)) + assert batched_ker.shape == (N_B, len(MOMENTA), len(MOMENTA)) + + for b, group in enumerate(BLOCK_GROUPS): + single = _single_solver(group) + assert single.fourier is not None + assert single.double_fourier_transform is not None + np.testing.assert_allclose( + np.asarray(batched_vec[b]), + np.asarray(single.fourier(vectors[b])), + err_msg=f"fourier vector block {b}", + **TIGHT, + ) + np.testing.assert_allclose( + np.asarray(batched_ker[b]), + np.asarray(single.fourier(kernels[b])), + err_msg=f"fourier kernel block {b}", + **TIGHT, + ) + np.testing.assert_allclose( + np.asarray(batched_double[b]), + np.asarray(single.double_fourier_transform(kernels[b])), + err_msg=f"double fourier block {b}", + **TIGHT, + ) + + +def test_unbatched_inputs_broadcast_across_blocks() -> None: + """Unbatched ``(N,)`` / ``(N, N)`` inputs broadcast over the block axis.""" + + solver = _blocked_solver() + rng = np.random.default_rng(21) + vector = jnp.asarray(rng.normal(size=N)) + kernel = jnp.asarray(rng.normal(size=(N, N))) + + assert solver.fourier is not None + assert solver.double_fourier_transform is not None + broadcast_vec = solver.fourier(vector) + broadcast_ker = solver.fourier(kernel) + broadcast_double = solver.double_fourier_transform(kernel) + assert broadcast_vec.shape == (N_B, len(MOMENTA)) + + stacked_vec = solver.fourier(jnp.broadcast_to(vector, (N_B, N))) + stacked_ker = solver.fourier(jnp.broadcast_to(kernel, (N_B, N, N))) + stacked_double = solver.double_fourier_transform(jnp.broadcast_to(kernel, (N_B, N, N))) + np.testing.assert_allclose(np.asarray(broadcast_vec), np.asarray(stacked_vec), **TIGHT) + np.testing.assert_allclose(np.asarray(broadcast_ker), np.asarray(stacked_ker), **TIGHT) + np.testing.assert_allclose(np.asarray(broadcast_double), np.asarray(stacked_double), **TIGHT) + + +def test_blocks_mode_shape_violations_raise() -> None: + solver = _blocked_solver() + rng = np.random.default_rng(22) + assert solver.fourier is not None + assert solver.double_fourier_transform is not None + with pytest.raises(ValueError, match="blocks mode"): + solver.fourier(jnp.asarray(rng.normal(size=(N_B + 1, N)))) + with pytest.raises(ValueError, match="blocks mode"): + solver.double_fourier_transform(jnp.asarray(rng.normal(size=(N_B + 1, N, N)))) + + +def test_grid_transforms_pass_leading_batch_axes() -> None: + """``grid=`` projections accept arbitrary leading batch axes in both modes.""" + + solver = _blocked_solver() + rng = np.random.default_rng(23) + coefficients = jnp.asarray(rng.normal(size=(N_B, len(ENERGIES), N))) + kernels = jnp.asarray(rng.normal(size=(N_B, N, N))) + + assert solver.to_grid_vector is not None + assert solver.from_grid_vector is not None + assert solver.to_grid_matrix is not None + batched = solver.to_grid_vector(coefficients) + assert batched.shape == (N_B, len(ENERGIES), len(GRID)) + for b in range(N_B): + for e in range(len(ENERGIES)): + np.testing.assert_allclose( + np.asarray(batched[b, e]), + np.asarray(solver.to_grid_vector(coefficients[b, e])), + **TIGHT, + ) + + batched_kernels = solver.to_grid_matrix(kernels) + assert batched_kernels.shape == (N_B, len(GRID), len(GRID)) + for b in range(N_B): + np.testing.assert_allclose( + np.asarray(batched_kernels[b]), + np.asarray(solver.to_grid_matrix(kernels[b])), + **TIGHT, + ) + + round_trip = solver.from_grid_vector(batched) + assert round_trip.shape == coefficients.shape + + +def test_f_momentum_deduplicates_per_unique_ell(monkeypatch: pytest.MonkeyPatch) -> None: + """``compute_F_momentum`` runs once per unique ℓ across the whole block set.""" + + compile_module = importlib.import_module("lax.compile") + real_compute = compile_module.compute_F_momentum + calls: list[int] = [] + + def counting_compute(mesh: Any, momenta: Any, angular_momentum: int, *args: Any) -> Any: + calls.append(angular_momentum) + return real_compute(mesh, momenta, angular_momentum, *args) + + monkeypatch.setattr(compile_module, "compute_F_momentum", counting_compute) + shared_ell_blocks = ( + (lax.ChannelSpec(l=0, threshold=0.0, mass_factor=HBAR2_2MU),), + (lax.ChannelSpec(l=1, threshold=0.0, mass_factor=HBAR2_2MU),), + (lax.ChannelSpec(l=1, threshold=0.0, mass_factor=HBAR2_2MU),), + ) + lax.compile( + mesh=MESH, + blocks=shared_ell_blocks, + solvers=("spectrum",), + energies=ENERGIES, + momenta=MOMENTA, + ) + assert sorted(calls) == [0, 1] + + +def test_channels_mode_fourier_is_unchanged() -> None: + """Channels-mode fourier keeps its rank-1/rank-2 contract and rejects batches.""" + + single = _single_solver(BLOCK_GROUPS[0]) + rng = np.random.default_rng(24) + assert single.fourier is not None + assert single.transforms.F_momentum is not None + assert single.transforms.F_momentum.shape == (1, len(MOMENTA), N) + with pytest.raises(ValueError, match="vector or"): + single.fourier(jnp.asarray(rng.normal(size=(2, 3, N)))) + + +def test_blocked_transforms_round_trip_through_pickle() -> None: + import pickle + + solver = _blocked_solver() + restored = pickle.loads(pickle.dumps(solver)) + rng = np.random.default_rng(25) + vectors = jnp.asarray(rng.normal(size=(N_B, N))) + assert restored.fourier is not None and solver.fourier is not None + np.testing.assert_allclose( + np.asarray(restored.fourier(vectors)), + np.asarray(solver.fourier(vectors)), + **TIGHT, + ) + + +def test_jit_compatible_dispatch() -> None: + """The block dispatch is shape-static, so the transforms compose with jax.jit.""" + + solver = _blocked_solver() + rng = np.random.default_rng(26) + vectors = jnp.asarray(rng.normal(size=(N_B, N))) + assert solver.fourier is not None + fourier = solver.fourier + + @jax.jit + def pipeline(values: jax.Array) -> jax.Array: + return fourier(values) + + np.testing.assert_allclose(np.asarray(pipeline(vectors)), np.asarray(fourier(vectors)), **TIGHT) diff --git a/tests/unit/test_mass_factor_regimes.py b/tests/unit/test_mass_factor_regimes.py new file mode 100644 index 0000000..e46dcb1 --- /dev/null +++ b/tests/unit/test_mass_factor_regimes.py @@ -0,0 +1,167 @@ +"""C4 acceptance tests: non-uniform μ(E) regimes on the spectral path (spec v0.1.5.1, T9).""" + +from __future__ import annotations + +import pickle + +import jax.numpy as jnp +import numpy as np +import pytest + +import lax +from tests.unit._blocks_helpers import HBAR2_2MU + +N = 10 +RADIUS = 8.0 +ENERGIES = jnp.linspace(2.0, 10.0, 3) +N_E = len(ENERGIES) +MESH = lax.MeshSpec("legendre", "x", n=N, scale=RADIUS) +SPECTRAL = ("spectrum", "rmatrix", "smatrix", "phases", "greens", "wavefunction") +# A genuinely energy-dependent reduced mass (a few percent across the grid). +NONUNIFORM_MFG = jnp.asarray(HBAR2_2MU * (1.0 + 0.03 * np.arange(N_E))) + +TIGHT = dict(rtol=1e-10, atol=1e-12) + + +def _solver(mass_factor_grid: jnp.ndarray | None) -> lax.Solver: + return lax.compile( + mesh=MESH, + channels=(lax.ChannelSpec(l=0, threshold=0.0, mass_factor=HBAR2_2MU),), + solvers=SPECTRAL, + energies=ENERGIES, + mass_factor_grid=mass_factor_grid, + ) + + +def _static_interaction(solver: lax.Solver) -> lax.Interaction: + assert solver.local_potential is not None + return solver.local_potential(lambda r: -8.0 * jnp.exp(-0.15 * r**2)) + + +def test_static_observables_raise_on_nonuniform_mass_factor_grid() -> None: + """All five static-regime spectral observables are stubbed out (C4).""" + + solver = _solver(NONUNIFORM_MFG) + spectrum = solver.spectrum(_static_interaction(solver)) # type: ignore[misc] + + for name in ("rmatrix", "smatrix", "phases", "greens", "wavefunction"): + observable = getattr(solver, name) + assert observable is not None + with pytest.raises(ValueError) as excinfo: + observable(spectrum) + message = str(excinfo.value) + assert f"solver.{name} is unavailable" in message + assert "non-uniform mass_factor_grid" in message + assert "rmatrix_grid" in message + assert "direct path" in message + + +def test_spectrum_is_energy_batched_even_for_static_potential() -> None: + """Non-uniform μ(E) forces the energy-batched spectrum path for static V.""" + + solver = _solver(NONUNIFORM_MFG) + spectrum = solver.spectrum(_static_interaction(solver)) # type: ignore[misc] + assert spectrum.eigenvalues.shape == (N_E, N) + + uniform = _solver(None) + static_spectrum = uniform.spectrum(_static_interaction(uniform)) # type: ignore[misc] + assert static_spectrum.eigenvalues.shape == (N,) + + +def test_grid_observables_match_per_energy_compiled_references() -> None: + """`smatrix_grid` with μ(E) equals per-energy independently compiled solvers. + + This pins the spectrum-kernel fix: each per-energy Hamiltonian must be + scaled by its own μ_e, not by a single uniform μ. + """ + + solver = _solver(NONUNIFORM_MFG) + interaction = _static_interaction(solver) + spectra = solver.spectrum(interaction) # type: ignore[misc] + assert solver.smatrix_grid is not None + assert solver.rmatrix_grid is not None + s_grid = solver.smatrix_grid(spectra) + r_grid = solver.rmatrix_grid(spectra) + + for e in range(N_E): + reference = lax.compile( + mesh=MESH, + channels=(lax.ChannelSpec(l=0, threshold=0.0, mass_factor=float(NONUNIFORM_MFG[e])),), + solvers=SPECTRAL, + energies=ENERGIES[e : e + 1], + ) + ref_spectrum = reference.spectrum(_static_interaction(reference)) # type: ignore[misc] + assert reference.smatrix is not None + assert reference.rmatrix is not None + np.testing.assert_allclose( + np.asarray(s_grid[e]), + np.asarray(reference.smatrix(ref_spectrum)[0]), + err_msg=f"smatrix energy {e}", + **TIGHT, + ) + np.testing.assert_allclose( + np.asarray(r_grid[e]), + np.asarray(reference.rmatrix(ref_spectrum, ENERGIES[e])), + err_msg=f"rmatrix energy {e}", + **TIGHT, + ) + + +def test_uniform_grid_keeps_static_observables_live() -> None: + """A constant mass_factor_grid is uniform — no stubs, results unchanged.""" + + uniform_grid = jnp.full(N_E, HBAR2_2MU) + with_grid = _solver(uniform_grid) + without_grid = _solver(None) + + spectrum_g = with_grid.spectrum(_static_interaction(with_grid)) # type: ignore[misc] + spectrum_n = without_grid.spectrum(_static_interaction(without_grid)) # type: ignore[misc] + assert spectrum_g.eigenvalues.shape == spectrum_n.eigenvalues.shape == (N,) + assert with_grid.smatrix is not None + assert without_grid.smatrix is not None + np.testing.assert_allclose( + np.asarray(with_grid.smatrix(spectrum_g)), + np.asarray(without_grid.smatrix(spectrum_n)), + **TIGHT, + ) + + +def test_per_channel_nonuniform_grid_rejected_on_spectral_path() -> None: + """Per-channel μ stays a direct-path feature; the spectral path raises.""" + + channels = ( + lax.ChannelSpec(l=0, threshold=0.0, mass_factor=HBAR2_2MU), + lax.ChannelSpec(l=0, threshold=1.0, mass_factor=HBAR2_2MU), + ) + per_channel = jnp.stack([jnp.full(N_E, HBAR2_2MU), jnp.full(N_E, 1.1 * HBAR2_2MU)], axis=1) + with pytest.raises(ValueError, match="per-channel-uniform"): + lax.compile( + mesh=MESH, + channels=channels, + solvers=("spectrum", "rmatrix"), + energies=ENERGIES, + mass_factor_grid=per_channel, + ) + + +def test_eigh_accessor_survives_nonuniform_grid() -> None: + solver = _solver(NONUNIFORM_MFG) + spectra = solver.spectrum(_static_interaction(solver)) # type: ignore[misc] + assert solver.eigh is not None + eigenvalues, eigenvectors = solver.eigh(spectra) + assert eigenvalues.shape == (N_E, N) + assert eigenvectors.shape == (N_E, N, N) + + +def test_nonuniform_solver_round_trips_through_pickle() -> None: + solver = _solver(NONUNIFORM_MFG) + restored = pickle.loads(pickle.dumps(solver)) + interaction = _static_interaction(solver) + np.testing.assert_allclose( + np.asarray(restored.spectrum(interaction).eigenvalues), # type: ignore[misc] + np.asarray(solver.spectrum(interaction).eigenvalues), # type: ignore[misc] + **TIGHT, + ) + assert restored.smatrix is not None + with pytest.raises(ValueError, match="non-uniform mass_factor_grid"): + restored.smatrix(restored.spectrum(interaction)) # type: ignore[misc] diff --git a/tests/unit/test_matrix_element.py b/tests/unit/test_matrix_element.py new file mode 100644 index 0000000..89c2919 --- /dev/null +++ b/tests/unit/test_matrix_element.py @@ -0,0 +1,304 @@ +"""F1 acceptance tests: ``solver.matrix_element`` bilinear forms (spec v0.1.5.1, T1-T3).""" + +from __future__ import annotations + +import jax.numpy as jnp +import numpy as np +import pytest + +import lax +from lax.transforms import matrix_element as standalone_matrix_element +from tests.unit._blocks_helpers import HBAR2_2MU, gaussian_kernel, partial_wave_groups + +N = 8 +RADIUS = 6.0 +ENERGIES = jnp.linspace(1.0, 9.0, 3) +N_E = len(ENERGIES) +MESH = lax.MeshSpec("legendre", "x", n=N, scale=RADIUS) + +TIGHT = dict(rtol=1e-12, atol=1e-13) + + +def _channels_solver(energies: jnp.ndarray | None = ENERGIES) -> lax.Solver: + return lax.compile( + mesh=MESH, + channels=(lax.ChannelSpec(l=0, threshold=0.0, mass_factor=HBAR2_2MU),), + solvers=("spectrum",), + energies=energies, + ) + + +def _blocks_solver(): + return lax.compile( + mesh=MESH, + blocks=partial_wave_groups(), + solvers=("spectrum",), + energies=ENERGIES, + ) + + +def _random_complex(rng: np.random.Generator, shape: tuple[int, ...]) -> jnp.ndarray: + return jnp.asarray(rng.normal(size=shape) + 1j * rng.normal(size=shape)) + + +# --------------------------------------------------------------------------- +# T1 — matrix_element ↔ integrate consistency + + +def test_conjugated_self_element_matches_integrate() -> None: + """``matrix_element(x, x, O, conjugate=True)`` equals ``integrate(x, O)``.""" + + solver = _channels_solver() + rng = np.random.default_rng(1) + x = _random_complex(rng, (N,)) + hermitian = _random_complex(rng, (N, N)) + hermitian = hermitian + hermitian.conj().T + diagonal = jnp.asarray(rng.normal(size=N)) + + assert solver.matrix_element is not None + assert solver.integrate is not None + np.testing.assert_allclose( + np.asarray(solver.matrix_element(x, x, conjugate=True)), + np.asarray(solver.integrate(x)), + **TIGHT, + ) + np.testing.assert_allclose( + np.asarray(solver.matrix_element(x, x, diagonal, conjugate=True)), + np.asarray(solver.integrate(x, diagonal)), + **TIGHT, + ) + np.testing.assert_allclose( + np.asarray(solver.matrix_element(x, x, hermitian, conjugate=True)), + np.asarray(solver.integrate(x, hermitian)), + **TIGHT, + ) + + +# --------------------------------------------------------------------------- +# T2 — analytic separable bilinear + batched broadcasting + + +def test_separable_operator_is_exact() -> None: + """For rank-1 ``O = u vᵀ``: ``matrix_element(b, k, O) == (bᵀu)(vᵀk)`` exactly.""" + + solver = _channels_solver() + rng = np.random.default_rng(2) + b = _random_complex(rng, (N,)) + k = _random_complex(rng, (N,)) + u = _random_complex(rng, (N,)) + v = _random_complex(rng, (N,)) + separable = jnp.outer(u, v) + + np.testing.assert_allclose( + np.asarray(solver.matrix_element(b, k, separable, conjugate=False)), + np.asarray((b @ u) * (v @ k)), + **TIGHT, + ) + np.testing.assert_allclose( + np.asarray(solver.matrix_element(b, k, separable, conjugate=True)), + np.asarray((b.conj() @ u) * (v @ k)), + **TIGHT, + ) + + +def test_batched_states_match_python_loop() -> None: + """Full-rank ``(N_b, N_E, M)`` states reproduce the explicit double loop.""" + + solver = _blocks_solver() + n_b = len(partial_wave_groups()) + rng = np.random.default_rng(3) + bra = _random_complex(rng, (n_b, N_E, N)) + ket = _random_complex(rng, (n_b, N_E, N)) + operator = _random_complex(rng, (N, N)) + + result = solver.matrix_element(bra, ket, operator, conjugate=False) + assert result.shape == (n_b, N_E) + expected = np.array( + [ + [ + np.asarray(bra[b, e]) @ np.asarray(operator) @ np.asarray(ket[b, e]) + for e in range(N_E) + ] + for b in range(n_b) + ] + ) + np.testing.assert_allclose(np.asarray(result), expected, **TIGHT) + + +def test_blocks_mode_rank2_is_block_leading() -> None: + """Rank-2 in blocks mode lifts to ``(N_b, 1, M)`` — never onto the energy axis.""" + + solver = _blocks_solver() + n_b = len(partial_wave_groups()) + assert n_b == N_E # the ambiguous case the deterministic rule resolves + rng = np.random.default_rng(4) + bra2 = _random_complex(rng, (n_b, N)) + ket3 = _random_complex(rng, (1, N_E, N)) + + result = solver.matrix_element(bra2, ket3, conjugate=False) + assert result.shape == (n_b, N_E) + expected = np.array( + [[np.asarray(bra2[b]) @ np.asarray(ket3[0, e]) for e in range(N_E)] for b in range(n_b)] + ) + np.testing.assert_allclose(np.asarray(result), expected, **TIGHT) + + +def test_channels_mode_rank2_is_energy_leading() -> None: + """Rank-2 in channels mode lifts to ``(1, N_E, M)``.""" + + solver = _channels_solver() + rng = np.random.default_rng(5) + bra2 = _random_complex(rng, (N_E, N)) + ket1 = _random_complex(rng, (N,)) + + result = solver.matrix_element(bra2, ket1, conjugate=False) + assert result.shape == (N_E,) + expected = np.array([np.asarray(bra2[e]) @ np.asarray(ket1) for e in range(N_E)]) + np.testing.assert_allclose(np.asarray(result), expected, **TIGHT) + + +def test_unbatched_inputs_return_a_scalar() -> None: + solver = _channels_solver() + rng = np.random.default_rng(6) + b = _random_complex(rng, (N,)) + k = _random_complex(rng, (N,)) + assert solver.matrix_element(b, k, conjugate=False).shape == () + + +def test_shape_violations_raise() -> None: + blocks = _blocks_solver() + channels_no_grid = _channels_solver(energies=None) + rng = np.random.default_rng(7) + n_b = len(partial_wave_groups()) + + with pytest.raises(ValueError, match="block-leading"): + blocks.matrix_element( + _random_complex(rng, (n_b + 1, N)), _random_complex(rng, (N,)), conjugate=False + ) + with pytest.raises(ValueError, match="without an energy grid"): + channels_no_grid.matrix_element( + _random_complex(rng, (2, N)), _random_complex(rng, (N,)), conjugate=False + ) + with pytest.raises(ValueError, match="Interaction"): + blocks.matrix_element( + _random_complex(rng, (N,)), + _random_complex(rng, (N,)), + _random_complex(rng, (n_b, N, N)), + conjugate=False, + ) + with pytest.raises(ValueError, match="trailing dimension"): + blocks.matrix_element( + _random_complex(rng, (N + 1,)), _random_complex(rng, (N,)), conjugate=False + ) + + +# --------------------------------------------------------------------------- +# T3 — Interaction scaling contract + + +def test_nonlocal_interaction_matches_explicit_gauss_sum() -> None: + """``matrix_element(b, k, interaction)`` equals the explicit double Gauss sum.""" + + solver = _channels_solver() + rng = np.random.default_rng(8) + b = _random_complex(rng, (N,)) + k = _random_complex(rng, (N,)) + kernel = gaussian_kernel(10.0) + assert solver.nonlocal_potential is not None + interaction = solver.nonlocal_potential(kernel) + + radii = np.asarray(solver.mesh.radii) + weights = np.asarray(solver.mesh.weights) + scale = solver.mesh.scale + raw_kernel = np.asarray(kernel(jnp.asarray(radii[:, None]), jnp.asarray(radii[None, :]))) + gauss_scaled = np.sqrt(np.outer(weights, weights)) * scale * raw_kernel + + np.testing.assert_allclose( + np.asarray(solver.matrix_element(b, k, interaction, conjugate=False)), + np.asarray(b) @ gauss_scaled @ np.asarray(k), + rtol=1e-12, + atol=1e-12, + ) + + +def test_local_interaction_matches_plain_node_sum() -> None: + """A local Interaction term reduces to the unscaled node sum ``Σᵢ bᵢ V(rᵢ) kᵢ``.""" + + solver = _channels_solver() + rng = np.random.default_rng(9) + b = _random_complex(rng, (N,)) + k = _random_complex(rng, (N,)) + + def profile(r: jnp.ndarray) -> jnp.ndarray: + return -3.0 * jnp.exp(-0.1 * r**2) + + assert solver.local_potential is not None + interaction = solver.local_potential(profile) + node_values = np.asarray(profile(solver.mesh.radii)) + + np.testing.assert_allclose( + np.asarray(solver.matrix_element(b, k, interaction, conjugate=False)), + np.sum(np.asarray(b) * node_values * np.asarray(k)), + rtol=1e-12, + atol=1e-12, + ) + # The bare-array diagonal form takes the same unscaled node values. + np.testing.assert_allclose( + np.asarray(solver.matrix_element(b, k, jnp.asarray(node_values), conjugate=False)), + np.asarray(solver.matrix_element(b, k, interaction, conjugate=False)), + **TIGHT, + ) + + +def test_block_dependent_interaction_broadcasts_per_block() -> None: + """A block-dependent Interaction aligns on the block axis of rank-2 states.""" + + solver = _blocks_solver() + n_b = len(partial_wave_groups()) + rng = np.random.default_rng(10) + bra = _random_complex(rng, (n_b, N)) + ket = _random_complex(rng, (n_b, N)) + kernels = [gaussian_kernel(5.0 * (b + 1)) for b in range(n_b)] + assert solver.nonlocal_potential is not None + interaction = solver.nonlocal_potential(kernels, block_dependent=True) + + result = solver.matrix_element(bra, ket, interaction, conjugate=False) + assert result.shape == (n_b,) + expected = np.array( + [ + np.asarray(bra[b]) @ np.asarray(interaction.block[b]) @ np.asarray(ket[b]) + for b in range(n_b) + ] + ) + np.testing.assert_allclose(np.asarray(result), expected, **TIGHT) + + +# --------------------------------------------------------------------------- +# Standalone form + + +def test_standalone_matches_solver_bound_helper() -> None: + solver = _channels_solver() + rng = np.random.default_rng(11) + b = _random_complex(rng, (N,)) + k = _random_complex(rng, (N,)) + operator = _random_complex(rng, (N, N)) + + np.testing.assert_allclose( + np.asarray(standalone_matrix_element(b, k, operator, conjugate=False)), + np.asarray(solver.matrix_element(b, k, operator, conjugate=False)), + **TIGHT, + ) + np.testing.assert_allclose( + np.asarray(standalone_matrix_element(b, k, conjugate=True)), + np.asarray(solver.matrix_element(b, k, conjugate=True)), + **TIGHT, + ) + + +def test_conjugate_is_keyword_required() -> None: + solver = _channels_solver() + rng = np.random.default_rng(12) + b = _random_complex(rng, (N,)) + with pytest.raises(TypeError): + solver.matrix_element(b, b) # type: ignore[call-arg] diff --git a/tests/unit/test_solver_pickle.py b/tests/unit/test_solver_pickle.py index a1b2f0c..e31b196 100644 --- a/tests/unit/test_solver_pickle.py +++ b/tests/unit/test_solver_pickle.py @@ -75,6 +75,7 @@ def test_compiled_solver_round_trips_through_pickle() -> None: "fourier", "double_fourier_transform", "integrate", + "matrix_element", ): assert getattr(restored, name) is not None @@ -191,6 +192,10 @@ def test_compiled_solver_round_trips_through_pickle() -> None: np.asarray(restored.integrate(vector, diagonal_operator)), np.asarray(solver.integrate(vector, diagonal_operator)), ) + assert np.allclose( + np.asarray(restored.matrix_element(vector, vector, matrix, conjugate=False)), + np.asarray(solver.matrix_element(vector, vector, matrix, conjugate=False)), + ) def test_blocks_compiled_solver_round_trips_through_pickle() -> None: diff --git a/tests/unit/test_wavefunction_grid.py b/tests/unit/test_wavefunction_grid.py new file mode 100644 index 0000000..36e9ab4 --- /dev/null +++ b/tests/unit/test_wavefunction_grid.py @@ -0,0 +1,328 @@ +"""F2 acceptance tests: wavefunction grids (spec v0.1.5.1, T6/T7/T8/T9-C2).""" + +from __future__ import annotations + +from collections.abc import Callable + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +import lax +from lax.spectral.types import Spectrum +from tests.unit._blocks_helpers import HBAR2_2MU, gaussian_kernel, partial_wave_groups + +N = 10 +RADIUS = 8.0 +ENERGIES = jnp.linspace(2.0, 10.0, 4) +N_E = len(ENERGIES) +MESH = lax.MeshSpec("legendre", "x", n=N, scale=RADIUS) +WAVE = ("spectrum", "wavefunction") + +TIGHT = dict(rtol=1e-10, atol=1e-12) + + +def _real_profile(r: jnp.ndarray) -> jnp.ndarray: + return -10.0 * jnp.exp(-0.2 * r**2) + + +def _complex_profile(r: jnp.ndarray) -> jnp.ndarray: + return (-10.0 - 4.0j) * jnp.exp(-0.2 * r**2) + + +def _channels_solver( + threshold: float = 0.0, + z1z2: tuple[int, int] | None = None, + V_is_complex: bool = False, + method: str | None = None, +) -> lax.Solver: + return lax.compile( + mesh=MESH, + channels=(lax.ChannelSpec(l=0, threshold=threshold, mass_factor=HBAR2_2MU),), + solvers=WAVE if method != "linear_solve" else ("rmatrix_direct",), + energies=ENERGIES, + z1z2=z1z2, + V_is_complex=V_is_complex, + method=method, + ) + + +def _loop_reference( + solver: lax.Solver, + spectrum: Spectrum, + energy_batched: bool = False, +) -> np.ndarray: + """The §4.1 reference loop: stack the single-energy wavefunction per index.""" + + assert solver.wavefunction is not None + psi = [] + for e in range(N_E): + spec_e = ( + jax.tree_util.tree_map(lambda x, e=e: x[e], spectrum) if energy_batched else spectrum + ) + source = lax.make_wavefunction_source(solver, channel_index=0, energy_index=e) + psi.append(np.asarray(solver.wavefunction(spec_e, ENERGIES[e], source))) + return np.stack(psi, axis=-2) + + +# --------------------------------------------------------------------------- +# T6 — wavefunction_grid ≡ looped wavefunction + + +@pytest.mark.parametrize( + ("z1z2", "profile", "v_is_complex"), + [ + (None, _real_profile, False), # neutral, Hermitian metric (eigh) + ((1, 20), _real_profile, False), # charged, Hermitian metric + (None, _complex_profile, True), # neutral, complex-symmetric metric (eig) + ], + ids=["neutral-eigh", "charged-eigh", "neutral-eig"], +) +def test_static_regime_matches_loop( + z1z2: tuple[int, int] | None, + profile: Callable[[jnp.ndarray], jnp.ndarray], + v_is_complex: bool, +) -> None: + solver = _channels_solver(z1z2=z1z2, V_is_complex=v_is_complex) + assert solver.local_potential is not None + assert solver.wavefunction_grid is not None + spectrum = solver.spectrum(solver.local_potential(profile)) # type: ignore[misc] + + psi_grid = solver.wavefunction_grid(spectrum) + assert psi_grid.shape == (N_E, N) + np.testing.assert_allclose(np.asarray(psi_grid), _loop_reference(solver, spectrum), **TIGHT) + + +def test_energy_dependent_regime_matches_loop() -> None: + solver = _channels_solver() + assert solver.interaction_from_array is not None + assert solver.wavefunction_grid is not None + radii = solver.mesh.radii + profiles = jnp.stack([(-10.0 - 0.5 * float(e)) * jnp.exp(-0.2 * radii**2) for e in range(N_E)]) + interaction = solver.interaction_from_array( + local=[(profiles, np.ones((1, 1)))], energy_dependent=True + ) + spectra = solver.spectrum(interaction) # type: ignore[misc] + assert spectra.eigenvalues.shape == (N_E, N) + + psi_grid = solver.wavefunction_grid(spectra) + np.testing.assert_allclose( + np.asarray(psi_grid), + _loop_reference(solver, spectra, energy_batched=True), + **TIGHT, + ) + + +def test_closed_channel_entries_are_defined_and_loop_equivalent() -> None: + """The grid straddles the threshold: closed entries are defined, not masked.""" + + solver = _channels_solver(threshold=5.0) + assert solver.boundary is not None + is_open = np.asarray(solver.boundary.is_open)[:, 0] + assert is_open.any() and not is_open.all() # the grid genuinely straddles + + assert solver.local_potential is not None + assert solver.wavefunction_grid is not None + spectrum = solver.spectrum(solver.local_potential(_real_profile)) # type: ignore[misc] + psi_grid = solver.wavefunction_grid(spectrum) + assert bool(jnp.all(jnp.isfinite(psi_grid))) + np.testing.assert_allclose(np.asarray(psi_grid), _loop_reference(solver, spectrum), **TIGHT) + + +def test_channel_index_none_stacks_all_incoming_channels() -> None: + channels = ( + lax.ChannelSpec(l=0, threshold=0.0, mass_factor=HBAR2_2MU), + lax.ChannelSpec(l=2, threshold=1.0, mass_factor=HBAR2_2MU), + ) + solver = lax.compile( + mesh=MESH, + channels=channels, + solvers=WAVE, + energies=ENERGIES, + ) + assert solver.nonlocal_potential is not None + assert solver.wavefunction_grid is not None + coupling = np.asarray([[1.0, 0.2], [0.2, 0.8]]) + interaction = solver.nonlocal_potential(gaussian_kernel(8.0), coupling=coupling) + spectrum = solver.spectrum(interaction) # type: ignore[misc] + + stacked = solver.wavefunction_grid(spectrum, channel_index=None) + assert stacked.shape == (N_E, 2, 2 * N) + for c in range(2): + np.testing.assert_allclose( + np.asarray(stacked[:, c]), + np.asarray(solver.wavefunction_grid(spectrum, channel_index=c)), + **TIGHT, + ) + + +def test_eigenvector_free_spectrum_raises() -> None: + bare = lax.compile( + mesh=MESH, + channels=(lax.ChannelSpec(l=0, threshold=0.0, mass_factor=HBAR2_2MU),), + solvers=("spectrum", "smatrix", "wavefunction"), + energies=ENERGIES, + ) + no_vectors = lax.compile( + mesh=MESH, + channels=(lax.ChannelSpec(l=0, threshold=0.0, mass_factor=HBAR2_2MU),), + solvers=("spectrum", "smatrix"), + energies=ENERGIES, + ) + assert no_vectors.local_potential is not None + assert bare.wavefunction_grid is not None + spectrum = no_vectors.spectrum(no_vectors.local_potential(_real_profile)) # type: ignore[misc] + with pytest.raises(RuntimeError, match="Eigenvectors were not retained"): + bare.wavefunction_grid(spectrum) + + +# --------------------------------------------------------------------------- +# T7 — blocks equivalence + + +def test_blocked_wavefunction_grid_matches_per_block_solvers() -> None: + block_groups = partial_wave_groups() + n_b = len(block_groups) + solver = lax.compile(mesh=MESH, blocks=block_groups, solvers=WAVE, energies=ENERGIES) + assert solver.nonlocal_potential is not None + assert solver.wavefunction_grid is not None + kernels = [gaussian_kernel(10.0 * (b + 1)) for b in range(n_b)] + interaction = solver.nonlocal_potential(kernels, block_dependent=True) + spectrum = solver.spectrum(interaction) # type: ignore[misc] + + psi_grid = solver.wavefunction_grid(spectrum) + assert psi_grid.shape == (n_b, N_E, N) + for b, group in enumerate(block_groups): + single = lax.compile(mesh=MESH, channels=group, solvers=WAVE, energies=ENERGIES) + assert single.nonlocal_potential is not None + assert single.wavefunction_grid is not None + single_spectrum = single.spectrum(single.nonlocal_potential(kernels[b])) # type: ignore[misc] + np.testing.assert_allclose( + np.asarray(psi_grid[b]), + np.asarray(single.wavefunction_grid(single_spectrum)), + err_msg=f"wavefunction grid block {b}", + **TIGHT, + ) + + +# --------------------------------------------------------------------------- +# T8 — direct-path grid + + +def test_direct_grid_matches_looped_wavefunction_direct() -> None: + solver = _channels_solver(method="linear_solve") + assert solver.local_potential is not None + assert solver.wavefunction_direct is not None + assert solver.wavefunction_direct_grid is not None + interaction = solver.local_potential(_real_profile) + + psi_grid = solver.wavefunction_direct_grid(interaction) + assert psi_grid.shape == (N_E, N) + for e in range(N_E): + source = lax.make_wavefunction_source(solver, channel_index=0, energy_index=e) + np.testing.assert_allclose( + np.asarray(psi_grid[e]), + np.asarray(solver.wavefunction_direct(interaction, source, e)), + err_msg=f"direct grid energy {e}", + **TIGHT, + ) + + +def test_spectral_and_direct_grids_agree() -> None: + spectral_solver = _channels_solver() + direct_solver = _channels_solver(method="linear_solve") + assert spectral_solver.local_potential is not None + assert direct_solver.local_potential is not None + assert spectral_solver.wavefunction_grid is not None + assert direct_solver.wavefunction_direct_grid is not None + + spectrum = spectral_solver.spectrum( # type: ignore[misc] + spectral_solver.local_potential(_real_profile) + ) + np.testing.assert_allclose( + np.asarray(spectral_solver.wavefunction_grid(spectrum)), + np.asarray( + direct_solver.wavefunction_direct_grid(direct_solver.local_potential(_real_profile)) + ), + rtol=1e-8, + atol=1e-10, + ) + + +def test_blocked_direct_grid_matches_per_block() -> None: + block_groups = partial_wave_groups((0, 1)) + solver = lax.compile( + mesh=MESH, + blocks=block_groups, + solvers=("rmatrix_direct",), + energies=ENERGIES, + method="linear_solve", + ) + assert solver.nonlocal_potential is not None + assert solver.wavefunction_direct_grid is not None + kernels = [gaussian_kernel(6.0 * (b + 1)) for b in range(len(block_groups))] + interaction = solver.nonlocal_potential(kernels, block_dependent=True) + psi_grid = solver.wavefunction_direct_grid(interaction) + assert psi_grid.shape == (len(block_groups), N_E, N) + + for b, group in enumerate(block_groups): + single = lax.compile( + mesh=MESH, + channels=group, + solvers=("rmatrix_direct",), + energies=ENERGIES, + method="linear_solve", + ) + assert single.nonlocal_potential is not None + assert single.wavefunction_direct_grid is not None + np.testing.assert_allclose( + np.asarray(psi_grid[b]), + np.asarray(single.wavefunction_direct_grid(single.nonlocal_potential(kernels[b]))), + err_msg=f"direct grid block {b}", + **TIGHT, + ) + + +# --------------------------------------------------------------------------- +# T9-C2 — propagated-mesh guard + + +def test_propagated_mesh_wavefunction_raises() -> None: + with pytest.raises(NotImplementedError, match="propagated multi-interval"): + lax.compile( + mesh=lax.MeshSpec("legendre", "x", n=N, scale=RADIUS, extras={"n_intervals": 2}), + channels=(lax.ChannelSpec(l=0, threshold=0.0, mass_factor=HBAR2_2MU),), + solvers=("rmatrix_direct", "wavefunction"), + energies=ENERGIES, + method="linear_solve", + ) + + +# --------------------------------------------------------------------------- +# Pickle + + +def test_wavefunction_grids_round_trip_through_pickle() -> None: + import pickle + + spectral_solver = _channels_solver() + direct_solver = _channels_solver(method="linear_solve") + restored_spectral = pickle.loads(pickle.dumps(spectral_solver)) + restored_direct = pickle.loads(pickle.dumps(direct_solver)) + + assert spectral_solver.local_potential is not None + interaction = spectral_solver.local_potential(_real_profile) + spectrum = spectral_solver.spectrum(interaction) # type: ignore[misc] + np.testing.assert_allclose( + np.asarray(restored_spectral.wavefunction_grid(spectrum)), + np.asarray(spectral_solver.wavefunction_grid(spectrum)), + **TIGHT, + ) + assert direct_solver.local_potential is not None + d_interaction = direct_solver.local_potential(_real_profile) + np.testing.assert_allclose( + np.asarray(restored_direct.wavefunction_direct_grid(d_interaction)), + np.asarray(direct_solver.wavefunction_direct_grid(d_interaction)), + **TIGHT, + )