From e398aaf5551021eb6b7605ba5eaae88e4deb9659 Mon Sep 17 00:00:00 2001 From: beykyle Date: Fri, 12 Jun 2026 10:29:19 -0400 Subject: [PATCH 1/3] first crack at adding wavefunctions on the direct path --- DESIGN.md | 140 +++++++++- src/lax/__init__.py | 3 +- src/lax/compile.py | 151 ++++++++++- src/lax/solvers/__init__.py | 6 + src/lax/solvers/linear_solve.py | 225 ++++++++++++++++ src/lax/solvers/observables.py | 225 +++++++++++++++- src/lax/solvers/spectrum.py | 102 ++++++-- src/lax/spectral/observables.py | 52 ++++ src/lax/transforms/__init__.py | 3 + src/lax/transforms/bilinear.py | 300 ++++++++++++++++++++++ src/lax/transforms/fourier.py | 156 ++++++++++- src/lax/transforms/grid.py | 15 +- src/lax/types.py | 148 ++++++++++- src/lax/wavefunction.py | 144 ++++++++--- tests/acceptance/test_dwba_reference.py | 132 ++++++++++ tests/property/test_autograd.py | 89 +++++++ tests/unit/test_blocks_direct.py | 10 +- tests/unit/test_blocks_transforms.py | 235 +++++++++++++++++ tests/unit/test_mass_factor_regimes.py | 167 ++++++++++++ tests/unit/test_matrix_element.py | 304 ++++++++++++++++++++++ tests/unit/test_solver_pickle.py | 5 + tests/unit/test_wavefunction_grid.py | 328 ++++++++++++++++++++++++ 22 files changed, 2837 insertions(+), 103 deletions(-) create mode 100644 src/lax/transforms/bilinear.py create mode 100644 tests/acceptance/test_dwba_reference.py create mode 100644 tests/unit/test_blocks_transforms.py create mode 100644 tests/unit/test_mass_factor_regimes.py create mode 100644 tests/unit/test_matrix_element.py create mode 100644 tests/unit/test_wavefunction_grid.py diff --git a/DESIGN.md b/DESIGN.md index de4fedd..384e620 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). - **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,34 @@ 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. --- -*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/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..d9c500f 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,78 @@ 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. + """ + + 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 +629,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 +844,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 +908,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 +930,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 +948,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 +968,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 +978,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 +990,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 +1016,7 @@ def __repr__(self) -> str: "phases", "greens", "wavefunction", + "wavefunction_grid", "eigh", "rmatrix_grid", "smatrix_grid", @@ -890,12 +1025,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 +1062,7 @@ def __repr__(self) -> str: "GridVectorTransform", "Integrator", "Interaction", + "MatrixElementHelper", "Mesh", "MeshFamily", "MeshSpec", @@ -940,6 +1078,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/test_dwba_reference.py b/tests/acceptance/test_dwba_reference.py new file mode 100644 index 0000000..96bde60 --- /dev/null +++ b/tests/acceptance/test_dwba_reference.py @@ -0,0 +1,132 @@ +"""T5 (spec v0.1.5.1): cross-engine DWBA normalization anchor — fixture-gated. + +Compares ``matrix_element(χp, χn, U₁) / (a·k_p·k_n)`` built from lax +wavefunctions against reference-engine node-sum ``T_lj`` golden data per +``(partial wave, energy)``. The fixture is exported from an independent +reference R-matrix engine; this test activates automatically once the file is +dropped into ``tests/acceptance/data/dwba_pn_reference.npz``. + +Fixture schema (``np.savez``): + +================== ========================= ===================================== +key shape meaning +================== ========================= ===================================== +``energies`` ``(N_E,)`` entrance CM energies, MeV +``exit_energies`` ``(N_E,)`` exit CM energies, MeV +``ls`` ``(N_pw,)`` int orbital ℓ per (ℓ, j) partial wave +``channel_radius`` scalar channel radius a, fm +``nbasis`` scalar int Lagrange-Legendre-x mesh size N +``mass_factor_p`` scalar entrance ℏ²/2μ, MeV·fm² +``mass_factor_n`` scalar exit ℏ²/2μ, MeV·fm² +``z1z2_p`` ``(2,)`` int entrance charges (exit is neutral) +``U_p``, ``U_n`` ``(N_pw, N)`` complex distorting potentials sampled on the + lax mesh radii, MeV +``U1`` ``(N_pw, N)`` complex isovector transition operator node + samples, MeV (local) +``T_ref`` ``(N_pw, N_E)`` complex reference node-sum T_lj +``rtol``/``atol`` scalars comparison tolerances +================== ========================= ===================================== + +The C7 convention factor: the test first resolves the (single, constant) +relative normalization between the two engines from the data, asserts it is +constant across every ``(partial wave, energy)``, and then compares. Once the +factor is pinned it must be documented on ``wavefunction_grid`` and frozen +here as an explicit constant. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest + +import lax + +FIXTURE = Path(__file__).parent / "data" / "dwba_pn_reference.npz" + +pytestmark = pytest.mark.skipif( + not FIXTURE.exists(), + reason="reference-engine DWBA fixture not yet exported (see module docstring)", +) + + +def test_matrix_element_reproduces_reference_node_sum() -> None: + import jax.numpy as jnp + + data = np.load(FIXTURE) + energies = jnp.asarray(data["energies"]) + exit_energies = jnp.asarray(data["exit_energies"]) + ls = [int(ell) for ell in data["ls"]] + radius = float(data["channel_radius"]) + nbasis = int(data["nbasis"]) + mesh = lax.MeshSpec("legendre", "x", n=nbasis, scale=radius) + + proton = lax.compile( + mesh=mesh, + blocks=[ + (lax.ChannelSpec(l=ell, threshold=0.0, mass_factor=float(data["mass_factor_p"])),) + for ell in ls + ], + solvers=("spectrum", "wavefunction"), + energies=energies, + z1z2=(int(data["z1z2_p"][0]), int(data["z1z2_p"][1])), + V_is_complex=True, + method="eig", + ) + neutron = lax.compile( + mesh=mesh, + blocks=[ + (lax.ChannelSpec(l=ell, threshold=0.0, mass_factor=float(data["mass_factor_n"])),) + for ell in ls + ], + solvers=("spectrum", "wavefunction"), + energies=exit_energies, + V_is_complex=True, + method="eig", + ) + assert proton.interaction_from_array is not None + assert neutron.interaction_from_array is not None + assert proton.matrix_element is not None + assert proton.wavefunction_grid is not None + assert neutron.wavefunction_grid is not None + assert proton.boundary is not None + assert neutron.boundary 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 + ) + + chi_p = proton.wavefunction_grid(proton.spectrum(v_p)) # type: ignore[misc] + chi_n = neutron.wavefunction_grid(neutron.spectrum(v_n)) # type: ignore[misc] + k_p = proton.boundary.k[:, :, 0] # (N_pw, N_E) + k_n = neutron.boundary.k[:, :, 0] + + element = proton.matrix_element(chi_p, chi_n, u1, conjugate=False) + t_lax = np.asarray(element / (radius * k_p * k_n)) + t_ref = np.asarray(data["T_ref"]) + + # C7: resolve the constant cross-engine convention factor from the data, + # then require it to be the SAME for every (partial wave, energy). + ratio = t_ref / t_lax + factor = ratio.flat[np.argmax(np.abs(t_ref))] + np.testing.assert_allclose( + ratio, + factor, + rtol=float(data["rtol"]), + atol=float(data["atol"]), + err_msg="cross-engine convention factor is not constant (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, + ) From 54df40cb8b01b085fb1029bb33e54d9cdea89291 Mon Sep 17 00:00:00 2001 From: beykyle Date: Fri, 12 Jun 2026 11:33:12 -0400 Subject: [PATCH 2/3] add DWBA reference test --- src/lax/types.py | 10 + tests/acceptance/data/dwba_pn_reference.npz | Bin 0 -> 78378 bytes tests/acceptance/generate_dwba_reference.py | 171 ++++++++++++++++ tests/acceptance/test_dwba_reference.py | 208 ++++++++++++-------- 4 files changed, 310 insertions(+), 79 deletions(-) create mode 100644 tests/acceptance/data/dwba_pn_reference.npz create mode 100644 tests/acceptance/generate_dwba_reference.py diff --git a/src/lax/types.py b/src/lax/types.py index d9c500f..2d5d854 100644 --- a/src/lax/types.py +++ b/src/lax/types.py @@ -326,6 +326,16 @@ class WavefunctionGridObservable(Protocol): 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: diff --git a/tests/acceptance/data/dwba_pn_reference.npz b/tests/acceptance/data/dwba_pn_reference.npz new file mode 100644 index 0000000000000000000000000000000000000000..0fc9618c3786e0f22d5391b546ed2b1426e895b7 GIT binary patch literal 78378 zcmd43c{o+!{`YU5_TFY0GbECEp2A-CHY74c2oX&}Whz6+kd#yuGBg=dQerETTI)ll zq9n>(NrlXX3YBLqN6&MebDclV^ZP#6^*j5Y`+9Hc>fZNjzwXaHtQ>7w+0>Z+@seXY zV%F}pxcuT~5@rhB9lB?S`o7>OW+r~79Oi%g+TZr9{pO91TbP*+G96m8V`td*&^5-g zYfN?-u2Gd;vul57c&P6_pZ%dbb`tMf`38mU#NQ9|_YK~OXC)my)wQa!v9kZ)Up!0| zFI{guJ40ovcCuM@=Ps2gOP@KYIl^gE0ejJFP#06LPE*`|uZhWhM4j&RxbLlX{P7tY} z`ScbI3nQhUG?-iRL=lPVb;?a!B@laSw?01AmB>Fn6I5PZMTnUM(bsza+vldD1cwOqWTN$}+Rqpa( zzY5YVInjI6R~5-tS}WmgtcJYHWxaU&ry9cFIV?cQS4X59CfBW!)j)12ES&ptR0GNM z+FO^^qJb#oyk49CT?29ary&i0&~xbTkaFNd8Wi?-JfAL)=im0Uw%p74Z_Bfs|F)#% z{I?}9=f5pwIsa{Wm-FA2Z#n;M`IqzGR$w{*Z3UO}-&SZj|80er^WXNL$Nzp!3CG{@ zC%)1_JNAeFUn1}yex?5(BJhtMO;MM72Foq-kCtmWb1i57<;=gFg_pDNau!?8V#`@- zIZG{P+2t&|oL4R9Rm*wxa$dch6_>N(a#mW-O3PV!IV&$`)#a?ZoYj}J`f}D>&YH_v zYdLEzXYJ*zy_|KHv(9qXUCz48Sr7ghfd*@SrN0OW;s|W__uaQ|XOK^*?~XkO{x=K9 znEvtl+x~kU{D&<358aN>3d`R%&41C^>DSK%{`Q3j@5??v->^OZmwosTmG=Lj5Aw^O z=|7K0e)I+8Zy$)Sq~-R#KEZ#j{P*Mk+y4JO4Cq{m>q~j_e;I~1y7&IcDnu9d{kL{u z%0KjsisZU~=>l!brHQ|zz=e;(-UI)=gTT%5w14RVlb6|-f9b)0?HIIYwpLuL z|JNQ^4g{(E?Exn~27&**L(rl=Ibq~6L;dHka@25v$Uk@Bzjg#3{>rgjC-v=LIxyf~ zd*N>f2p8=P_xqatsPm)xA)lGx74iVqU_GEba*6gWD7Y*()UtUv*=M|!=*LXT; zaJDgNv%Mo3t{&IjMsBA8sfuE-!-WhF`9^ZhU(n#ESX#@XHyIqi_FHzNG!PakQ#pKy z3=0!`pLFKX;ChI|z8_8a-!}(50R}XfG_HBTT0wzrD&bm*w&q}@U0K_qzz;V4pUN^* z%%H((>U&6?5Ij#c?#K8H!1LFo7qy5ad>b;byQ82CuA@e`3X)_&{9WQyue|_lUwiD6 zq`4Th#*H`yZ5yI;2!;7osIWtNoTLCBzD&GFO32>&WfY5T>poQ=L_%b(L-MwrZpGkC zaQGd45@N%_=J!~r3VS8%o}w_!jf9^rjpb#_#?GwsjuN!xLc~1A+PXeOVf!39*PTQ; z5W9ER(4#&+SmHzJq;M~GWV-2d$;uWw_SS3n=ZQ6J$Sy72B+N<~J1P5iMf6HmB(HgY z>ahcY*zZ$aD|HQ75Y^!cWu}q^^p;-A{m3w8#Kba4@`TPHI;(K^=SBuf`l%|EiT$CLzc4ZCv@J##*2s#+ zNd2Le8&Hqj16l#c%@JvTF_yu1WuKP_@gZ()j11YF7mwvOy+%Kw< zrN_Z>%_P*V-K0wB%`d7)=hTYiU-2mQMq|?DH*}~e!kF)sk)TZD&?$X>I`q3etp21- zhMO6MCe|}FU_Qv7bJTBB~JVt4d?oi2kj7A0rLA&cULo{d=aP#1@CPSgTl--uk zGBl@MId4t=OS{ zHkJaxe7g>1J~9Vht}e5^zxhD^v10DzeJYH&x0N0^E(CLb-0wt)8Ny2bM<1s}Btg}6 zk0+0r3XF{~Kd_I)*9=t|zIHN#Fr*r`tHMwWCJIa&XFNYsleMm>Mk%sGCGQCdVN(vo z=)=s&_K%~OyuY-o2R8}vOpCPRO>4#O9t%4AUWbG{5eO5Jgu z7qhV7fx3$B4P3~@Lg_uUw^3NIoY1QnoR{0zGcyf6e6SK*_PqT#FIxMBl)pC7F_Jvj z)^VJdBP4s7XEbF@!+$X=3g@NW*kRp7s30~WF*Buu^Kz)-Sfx+l0@^3sT1N1anNoAB zQ*#h48tjfGc+oo%{U%GN9laMb9Zv94n9+a7Y@r^V+o5-t;N>M}Sf80K1FhYL3J|<> zld?R1+Z3a)$9@ID%OT^)QP;d2l=f48iI=D1q8nvB($RwFx1STdq@h2~{s&j=eOvZ~Feg z)=V-i9*B59bBPA3;+Pdj7#aM&R+j{n(ID@((p*m<8Qde;pWi8^LHum=l;nOgun#Be z30|Oqqf@m1>Nql#i>qDu5=(liQOIfDCz^em->eKZOSDDx30$9ajhd?hF!;DZ@eIc<)^c{!MvdimgMIu_rm z@ck>!i%Kk$?G1Zn%qAm8>=4e2(s+f^gJ?laGNK@Y@XJZJx|QZP7f{!T*|!8QC*jSv zuZDwYZt-*m!ApzJ+k3*s?PyFR8qogRU}WY#{t{gk5Y{!8;wbJ@trpctB^g0tA(z@PN*6ZT)SKjs+zh5-S4!#_F!-xzC2KIk~ zf79S{{{`J%DjB{#&763;NQ3hWz4J$%$S@z09gsLr118-UXy*wstP2~b9meOg-QV~x zt+J&+!nE7t5>tG>D%;4Tc8d=dU3BRRQ3wcbyX_U_CIoB3B2qQvjNrZe@SNJH1nhd0 z7vSTn3JD|6nihg&VL0V-+x#io=5p-Nc5z-_T#N&` zTkTH`csq(U{q|caw8(>GzdHPO6D+itildG`O`qy{P400pB8=M%U z3t5=`aq;2HxL*`T9T&^GqOj0u`&~geFOm`Q0V6IxSabVmMF7rAu#ox9M^EUO@t13_ zM{!;XKJTt;GgHQ1-I5a_{6fC2KIR-Oh`lf2dqnu9IdYXDc5wmani0z-641}eOXX&(1wBfKXc$GcwyTx zen@_OF?zHtMugyHZ8+Oa!<-zn(BRT~f|pd;An|AJ>8Q|gw_bKQu?cycB zWRs-=7SxhZ-s%TSetFw@uw!~I9@S6|ELnI>1wZn=qk_7;&=LRZ*JXVgu6d5We#}jz6Xi%4}^&%pM42rLvlje`m zAZBmKo$8Zhc&pa_G-N9c*4KRq54%T(80E?~Hx9f$YB@7ovX%nx&ExbpY&3_us}+Hp zt$AVVACs%uK!un0jV z^Fk>aEHB)ojFs3N%q9FX8T%^2Emjay9!)ML{9<-m?%aop1$1rBjR}Gmy|`5GE9(Z) zk&h)g1TQyAoH?l`?P$i2bMXW(CzNa{k69Yf*lCl81TVrPHBWptF;FHnQk39DO-T36 z4yR(Y%r8rv;HA;o`t!@e9JFCx-Q~Oq5+Y`>!E$&1F3lZ6J4cz-%W)MDk9_|>mV|0tYCeAT#gR%Ef4q@ zVid^ewBf6*rh%2i!#88w_#nvoO>LWyIT-KG=8IJj0+*pIpP{wJP+&mj_Xi2Mto=Od z_C7Tb9m(jr<|_*?y>8LfKnSLEYy!I1h=KC($8vV(57fkc-1R(C>@ck!pe8-dju@o% zoel0B#e&aX{nR|egS=XiQn4qw6)O$j@k&RXggjDmifQDk!bAjxHa`2zjdZTLV%PH{ z3scdIknzX;Ql^uXCGt873xAxt;ED5+L=QsFJNRIeo^4Nd;=GjC?9A$1E?{ zSQVqTVKMB{A=dc@ju$uAqH zer-=wO+xpqiCglEOqY><+0S@1|Kck9fcK`5R`Tj)?~Ls@AhNd&r>2 zTk-wDZW>(De7Do{9A1w!n_p>)r$IsCi}Slv$iTI`!@vj6-wlRK+mDmM8suIL1ks@M zxIc@{1u~FmQ8RNUG+CH-t`7$+_2!P9p3;HrqNnKg)rjcPcF` zn8&b>hn|F>c)g;l?b23yz7>0aFa9T)MnW9BtTXqms>0;ndn?Old*_hXfmVc|a5RN;{|$8cWej5P9I_Rz7_7N0m~ab6~$ zJ{Qb&RmRv@TqI&}UaEP&HGVxTh#C6?ydwNE*_gtgczXf;eRumJ!OPoq)FyYcLDWRv z;}XHk@$qNQd(7HV>lFz}1TU5?3PmMs4XAYIpT`6*S5_ToQF39RPMe3s2wqs;^j4WU z7o(qLcS{hwh)dMTPZZ~%esy(AeqsK(*WvZHbac0L%939qq+cHP8A(AatF|opWn*>6 zsaN_*=m=%qZKB?F?)Y@>DsuvQK33=0{9`LHe8YM%;xRWo*UqckRYeE3(5pWJ>q+2w z^F?m%dOEym^ZmlXL54WL6=fT6KO~PB@#=|_!Rp4DW3>u&kX6@tG$Tz0st?2OvKk$% zuieV4mL)^h<@M+8B6L_iKyEs?nL_vb`Zw8_2paGk!oMZ z;e)S6G10YSWk*JM5VF@PKheZitRP_9)K(P|Lem8eDfTKX!O!W_;wNroBm3Oo)9EaX zyYscH2JRQFm^&(mUq)dE&o!po;=GLVxv#L_;DZT&y8CV$&da#x`KrKrI+lH4CgJg*b3?{A;K?l$Hq6x1PWq+S_Ke(*~9H>8Fy|0^<`X-B6!KS&kcH@G>Dcf z!5)H_3)P}iE7aQ2Tv@{)f|sz%jx)Kl_2`OwleY+7B)u0^1}z!rnzSob(S?TA5qj-FeF^l(I3ngmi+S3m%nmjZwgxECAj1l*4J;h za;qevk8;eG{Nmg7E6eXkJjyB5Gu!oL9n>HO&&(qEAaJdzZpSlo$dO;r?f8vs0QvI{~MFi%`W(Q`HGK2 z?+r2-yUNHWO3?sbS6m7FM26U_$24vhn?oB(xg!i;EDNiZK54<3+)?jyE2Q8hAWY;Nt6|T$ zQvh;b7L)@+4mwBizzfl@9(sXdaACIshi}~!^<;G9{zsPVFnV+R!>j}+LhU!5GUXh@ zOa~LMvs;spxiLH0Uzb|3!AD2B2P{ZP+w-n|F~utE&MDcm!eiWsf(b|TPO)svbnf^L zbvrI3YI5r1)1Oh;4!s*r?{Qu%`zaZeXdmpk#6|v-I4^Zs%%Z>E(y{Nc^wZyQUfL!Y z2Tyw_V}AFv@(I7hUJ_i9mMDnXwA=q7{Br5f`0-VD7tq#j3H~UYm!osjdSvrK^pqc^ zgy3bjE^_%S(vJSt4LC~hvKYo>bc~|`#qPmVf){;H57Vj5476x$vjo9QXLiyVvCYM( z?0Z8=f|tO4bwBb|d_Ian>`PCVc@CAi}4S<9$g3=wIY! zSZrnC0>=vlUSAQIGnU!MA}R*Yy@WN?GJB|p%l2+QFTf6Z4mOBVDLbNP@W>>#c@%qq zYKEEi^B~p*S;CLwTCu!sUVO*!^)*m6K87$?VdIR`hxBp3JRGn~sF}>dhGMl(eYfF4 zY90;M4>v_&k$YGa*WtW4_tD1|EqpM;pi8r!I4>$pP5M-njy+m;V%-qV%Y4|K)05iD z*xjSY&P3q6)ckn<=ZL=`R=L5?l<-UKtI6iKXBW^@8PooRU%sV(V&YL8M7_R{w-dZz z`wNQ3l-tpd%no}9UTWtLc;B6_M-OzU-6VKPRZ#!B(t?48v9gf~UXt0qzt5r-qaDGV zOMcOkEB*R51Fv^y`Ir3Cq8k|Z&?Oz+gHo6L;+Sd}5!juA)`vWMM%26aYEZZIwMnQf z&tfGp9~E_Bv*O?5(eRXv3)VA65HfCjK$*f1)v<412A?&D4-&V+|Hx7xsAQOuqCo?J z^}ml))Zz2VmFGput~9us9*WEqlA(6<)gmr88U#jf-hQH(499aG1XY>f`+E(p=l9+x z!=8Z0wmiSgA=Xz6-93lbFFR-6c7~b5RppKt*JBj8mb%VQc{>6r!q>&GzT<_3D;c9o zBPOuZ+R^PCzW}HvZIKGIHh?z8;>Vst!f^jp(ud$)O=wMivVN<%6l87U`&xZg9$t?h z_d5GT7Ch^nGef($V7)L$_vJ7#;O$swKL6k=mCwuS%EAVAXj&Q9HAUq_fTtc^>w`uAZrx%fNYw{n~z3=>r{mQd21U6X)f&aBfes zw=xEyd{VJEFBY0QI~9`!F?VYn6+N66M|~I8hjj~R@%$@U!Y}GI*9tKDAlhF5Hwa!V z4^@vnptYk_&d$dPUhawaw(Q|*Kz%MIHWR!&J*yIV*NuVBBt4QOc)97W67tQp81=Zo zBSrWHEh)ULT#B!!O&%}#cs8oXwnt)mqfjL{zcrh{6q?R;P~j0U*1m$=RY(` zLUmN`E%~MOjdzzJTLP-d9n*iJZasXweALr6feZ9gw;a)uwSbT$YNX^L608=z<>VKG&dXn{ke#ticZ`oH}5|6$)S+p$Zo5uQj!I@`9AX9mVmFPl4Qmw zpNZxQy2`5+_VCbfNN$I?cJW5&^7#p zxq>PSM|yj+^iV5y#?@V&p+G{+Y%gX%`OUz#u(k;9{m6|N!)3{m(Jbuo>h#AC zaKD&QB-}SOMq$zG%<7bJUhZAfXsD$6U|VOB$!<6=dy0S53f9rFX=c@`&p0n)_AgcD z)Ri%hmg@q9U#525xHYm<5G(A6ZD<>#1J^=29aT@{qYL-2BoX3u;dDMrsT z^aKfB-g)R%VW;u^I=wV2Vm>;(4mJcir=z{m15E@kC}W%H@y-@;5jo;yEiwz`gH3wzE?Y9*zD3J2)NN%7I z4Vv21gnM6*K}JZhs0YFK?>Q~c+^Zmiz>Cvcw$f-&B1uMuD#>vD$ccOTt7)M2ZsV85 zW-{c&1xB3d#m|E#4tH$lr9j;3ZD}lYbBM7y=Q4Pe0$Cds8f8QfXg2(PO;msndfwi2 zWBa)d&K<}KsWA`$r;DN*VuJc`U8X_MU0ejD>YBttuV?^oXpy#jkQ8`+di%lm&?+#S zsJ9<%lZ7-Cp_YpbR)-z!C{NGj4tdOHQiLWSKE;tN>B2>9Noc*nCrP(YYB~E zQ}^a1Gd)R2lAqcA%O$PY_pY@3QyWM~$DR48y=qn1+wr0uSI4;#!vt^3W3t&;Vp+@V zCEPDT-f1_Y^MqeW_0G2;MG)IvC+bf4g?Zn-nC;jC8XK|&7M(`4j-PvKw(||H1TC@&x)-fZsM730@w2+&9v_BOMj^a&gHo8dqG^IHpoiw{YDTOXo)<-2#(! zNoZH?+>&2Ty~&SEU{65j?JuXfRNKH^j`hZR?>S*iuXZ@w)dG~#P8d$Rk%0Tz&+nt> z>2O`6$~fgE2`G|QH;n`7a45M!Y+#ZErdHj%>bBD%cmtpMhwmil%Ud5cbbt;Iotvr_ zzmkA!hpVt`A?}-X9r1a?B#5j%#By_j4qdsN$G9d)aIV{~@b?z>Vptej_HkDbZaiw3xniVHJ$GKaFT!shW` zM{vJ9+H5ybS09DNP@)|raKEfNAepn>*au6936kH0^Kxyyg6Uul9g}7jeA)s!=V|Jv8AXPk}!J&Bkhi=HTS-u6pnZ1vuoSU6fBzfp4p$;W0WN)QU7c ze9yiPo)z_+==2tVh}t*i5BhW=QLNv`#a0CB+fMFd$y5V7@554GvZX-w;Egu&b~(5< zF^3qvlLd6ISZhT#J5VW_Zz7I}!4%_~n{dN4RX}EdKVUOE9Q%Fo@!)YzZP#u1HqV?Q(aalh~$+Q9eoxDV!?IMSSt^TJdbc`{*$js>dtDlOo=L~Ywy z7`jUtJMQw6bu8o!PZeyO}7`8N`WY^py zcwyo$N*1(kN2N0qP7}PS?@XUpB{!f2J2_qvyo|85jIZ}(pywC`(gZJVebMSA9>wTO z@$!`fFFSNnSMI)vpHDw%S@O%TK5=QSUFqm0)n7|~X^2tZJ^VEVHS&pF@=M4X=c{*2 zlF*p-Mt6vM*X~W1SQAGA`Y7{};mXPla4LvfT*j0W3O1HI_Q%2)vFA*+hW77~Tx)CCT<+uxblA!qTNmw_=mQduKooL)PM zRUH2Fbay=u@-w(hxh1?6OQbmlambPoW51Z+q2C#pR`CY&o)6rJ;-U49NguPY_)uG2 zJKQfDF2)^Eco2ofy}$FB8~019Ip;nRJ^Xx{&DnS(&dcD~CaAql#~e=WEgZmkab{uN zL0+SbUD`UyNBG6hGEks;n;>>%;&wIRm#fMTe108UK!r3~^a;PrxivnRTsesDR++XZ zcu~wZL-($3M;or3-A3?|$~|~qXrdlHt1?_n@X~SU9`e|XfgU!!%SrIUj{Tu?nHHmK z^w01Uyco7@)K5Br^HP;)Nz6xrHUS?WIHsf6^qM8V)xhT{@rk^R-k^ zNJ6!}ktM(QcXNp#-{R4`pK3DC3h6?{Yq6`gOahP{jTRzP2s~hF)`>`^0Gj%sLKm+; zOh(IeG{2D{aQu9M@d~{Doc*F_{hSO9b=w|uh|@sQJ8JGCp2xh_xCYIbL-Gofq}8MN zd1OgVOJSZlRB}3SdKu#NOM$@`w+{#?riim~3{l`i$yZ5{Z)V_}Q1s;I0Y1ptIH*x4 zWDEVSv9pBs*`BCu=3@zI$?6-eFS$#JV*3I;h3J*?Rv3xbnc zI%!{IVd~-h_{&*VsN=SD{+K2PzH*0GIzF1Ao@J`AD)3;3-Fr3ee<JfTrD$( z6`c+~xhj~1P`WhSEpN7B?BkhRjU7oy_eCk82Hh%*W8)fq<4JDBwy~POOeq^{+$xkJ zu#pSd_F~P6RitQa>-ha&vbbML%Ds)OPWxcS^Q?C+!mnC>ny|{`J?^TTQ zbRCu^c&VQI?U{TFpO504?Fn9tJ03Y5+ntW81#f93>RlHXj)JCXoRfL;oXcl!AU7_KOW_Pjl`v5n)F(Po6!MT{Vii ziO|z!Yj_ZyYDV9okXG!Fq1_41l_aEg?`Xx-Zw&1G%&{(h+%Fg0&yR$>&%(ULO9O>) zzqAfkNSwYKg&nT^8a?-b1Nlaora#f}!J5*HoYv#KC}hObPgc^gSQlQ24>&I&eNs#- zdu(C$Pyxau5|6qpl}- zv2l4CLSNO6l2fj4C3tDwPU6%at4F_CD;5&GSkk|+|2AQuY9%%t1TUJuU3aRlD@Mm6 zfREs1khABA!ktqqNGcWYdo1fajg+pM=10gj(>VV7=Fpr>Pti8zMW z9|zYx^kBo!13d;VkKpG^cY^MxQ1J652_2h5o@_L5UDW89?OVu6OHfZo>3|}SJ8@5e&m?(JLC7K&ffG1GOWV#ehpt_o#IB`4Er6tubPbwSKo2i zfcvHU$n4BrzG%!-GGVtC?icDn==RJ^A8bYF%8OTUUXr%0Oz9b?W4xDSjQ-%fnBLF| z@Y$z~8MqvkIgInNE$mLg>C=K(rs(5$gkOY2#)mDM7SM08<=Y9ryu5kH@udAAy5`c! zCj>8b*{7v<+qR>rEk7~|UWzI#6{q+b(2uIA9Rx3wgTdRoy&0(Shd4Qc7v^nEqg%X- zQHyLR8G;vX2+1p|#Mjf*t|h-jz1*Mt-aj3!tEyh|%aF>uo*T1xy{pUDLd-{O+cK)I zQj^feop+b~^6d4dEuP#7sJkpBqrAip3=PvdzWv0{=U(9TNL0tW)9NT4sEePT zuYX8_m^C{xQdljZF6)v(Nht|_#2WrkO|SqD(??8OPLrU;^OF2)x+Pe(My$SALjpG8 zAz>p9D@c7hzf;AO0{+A6ex}>w^-S)g>Fg&0V3T@BJfaGLiF&TGKbhhnB`NeRDNP$R zhPaJ(2+G1L4N*@dOdQhfs(%cRi9-LR{TGglwbY+<|MkayutASZ#D4QBHsqEc=g0Fm zN3q}6Q{1|5@F4pZYEP}**NTOyufOggK|&0oJa$M=GB7U7C;P4Yxe+1#sUzEZvoI6j zDC$|yg}7;atoT_Kg(W1Z@qfYnvZd#pc$BgaHX$tQPse%Lu2-C3rb&zN^+M1>djJj#~0dmaA|KpM5%76ts8AFJUgsu1u}?db-zZ$uD-w@id)5 zh2~Js(!42o1qEuJtht;wY7Vs@Pj$Bs;rsd5wX`ZG%^~CRAj1{U>Us~qn&9^--`Oj# zw88JsIlp98tlDA@YvrU3UHm9;Seno3=3xY0KFY6BP~?RiMo!sCtts?9vl;B};RDwP z(p;e$hTx-qy~p6S01Wf+N8d2f0v=XlYnu-uP$AyAzshF~aLNny7>lliPp+C9Bz2`= zXtF?5n@J8f1!HS@Z!-adWxvC{3t~Y3qf+|h>33@OLt6&lPIge(=vdnMlM}Iee{@z< zWejuOG4!Y62njj2QEui-RVyaKd_iV#GYNT{Vs>MNNfoAk<2vUH+%G;PzP=Zl*%*o6 zT`URri?hi?ZwKLMY>SO$n>p^6PLY(81Gzrf*Dv}h*KuAf?@zLFex+mmU2*k`I4`no zqAi<4m9Z@Hj?08!*06OEpaSv~%|R{9*xpd~s(u;g`rj&HIScAR6@! zX(4#Yeajx0v9TScj`yD7TPeKD`ek}PV<@M>gBAx_v)h;d5h=Ps4dvW8PK38_A|6~xYs%Z(Vx9W^GHS<7z zPRi3Q6boo%I`FISCl9wd_+a6i4wHuhjtv9Omh5_PW)K(T% zE~W#I2VYqUaLNKVo8t2dx&#P27DUB<5{1hGt)f{oh17$7d4f*g*kIw-V5C(u8=`J_ z$e`%zC{~O-mp)U-gOnomNM}GR7P#8t*%46^a_7i1(eP0Q7M{h<8TXzWIWA-&oPtY4e^dUASNBOPVH?6?`zE3mk{da9#?!du?>f=vYs`1NS?e zmmBGhhHE92v1DD>li@fol(XEchuj4**E5`&gkQ3fQl};27SP=1*VqWZFv(ntej_r7 zhV1P&CU`m7*xNO+vK=k3>2f1@@f36I-8Wp1R>T*cCwN)aW|baf#6VLcf>;S&jz)1_ zi#8}mBNn8V{DK_%7Oj+wpErii&v7WEH4%@_C*2ptcB;Zu-S@7bX#qGGap?g6Mg(#P2fs~DQ-FW} z8Rvnc=HP$IW8jt&1>WlJy|kgj9PWGgs^rb$`z>-3Lu~l^soDI(X~qnG|64ADs&K;` zcmm}Ddc`S___5S!osl_wzM1Uk7(oH1x6w15_6QVz^g40Jk{7}ZgTmgdGzI>M7UL@% z{GdN;9KO-i05ZHLXWSM9fLF04&52VJ1{T~hErnKqt5)7)y1fEe@=ZQuqOFAB=ocy> ztPs#~6a|Y?zhU z5SE|6wd@WmW7Ch!<>GN(CNjTMA7u(+Rf{JL3BQCNsqIaEwSeYeQptp0WY>OPGvhpn zHm#g_P4IGXbxcj2LpwT4YRV&c5&bpqR3+4aD$G#&2wpUX1iSb6G0?YsUaJXS%07yh zQv8b1TEW>Rzm(f;5$(T&-)DdiOMdxY|9I?VU^?1q9J1sWa_>p|khv5TcnX&MlAZGs z%cUiud`Sk?#C&A=`g!gcDFGcT+x&%k-X1=(@?|#OW`l$2=iX)PvV__i&iS?1cp%=UZo%AL ztOM20bxLD?$iP&WuBn``B;-1$EeQ1C`!7)omQSwaVYBhi`^3lDKnih9v@KynZc-*s zlopI)8#?$0H{|djpOT+`JF>eK+aZArmkW~+opi(IXG09E|1Du`K2%d--G`DD8~Iy!m*Z7lmb9f9-W?mNx9N@x(3 zRaev{c&RwNGwqRNJ36X7=|b>Q5PHb+C3)JCU-tJ1XRaKNN85r_cAt+{0uJu%FA)}k0L@f^vK0snyv@@V6XgZf zI`#LB0p|F+Vom2CClVX7`)8I|sVyi{2e8|AhlRgA*vkatHVXh?gZl7=5NT}oPt5b&A zm_x()(iz+@)@9z>F)O37%~xJr^~3$b6?^MLT(J*!t$WXw+c+;z!!m#S{Gek=cBXAC zxL-c4{FFtDR>mA-WeN$u+))rc?R`!VYX}S7N%%$KzRYIDH+a3<-=0PI<>_GV*{iOD zsMOZ0y#y~Vf!|mMo!Zed<*Ew=FRu#T3}uQmppBg`-x0ik%qJo39Sl_CYLf!NOZ=7k zl%eg#=uB&)9Kj2-(2LD~YVi3eVV5H@AL-Qv%YF_@M|Z}xEcxXkW9KpFg%q@K_RErA z-rt^hENGF0-WG^h@(b&0Zeu-40$TkwDvRr!1Ki!y8d@R92C61MXg|(dLNGYrV@~FQ z9ZU|D>6RAY{AKgo6h9KIO}kPlF-nKsb?s^<7f7HIy-OZ@PY3>_C(QRsNU%ciK*<3P z3%GZE@pA@VuS|}F@J{ZqfRsu9iOK-{J}_-|^oEcn^!CR`m1>fJU1;FPw3-z}Z;v}@ zCrbi-xc_Z`h&6auc|Hx;M21_Kv{$aQH8hn6kPEd0;H2U9?o|A|Q|EX1_zhuk2t(c; zW7X1y56urBb@a)k3Xxczo^E-Fa0SQ7-9p{*2gCt6W9>d zkAD5P@<*|$z|A*VPw^nK%`c0!ceG+&{3e{;0wkm||7x7fCkAGxpQtzbmK! z{DWxAM5r3UOPk>&HBY=9)gyCkB6t~ke)(3}r+Rc>t@Js97bf3gi9S6BI<)>Y6TwU0 zY|TicPBF@IIEqa0l0U)oXEYwKcc(^Z1TWnILLay{q@!~Mo0t4@|3%cbElnvXLn3j> zFClLpD?VMBgwAKOEcvCfr*r!CXgqpjb6$1S>D3_VsyCmQB?!1$ww%E4=c~N24@Hc5 z!J4xMJdagt_w zhw%H8MSN^bMil5LeZOY=6M?NyToTq7P=NV%aiK>K73}5S8^1Zi3(6KktYx3qflIzz zlZX*N%J^1HO7zBcDC6TGBZSHvC^Z9sR|Wqcrb+2nr5etstdO)xWB zL-4YB_nuw#JBv~E(>bdMUR=77OsRYLdfI$+$uD8IC0$JRrK8jJDs9Akq?Kygp!XZU z@7=Hcl9-Q_UQwQ(uuMYTFIO-5rJ>!bCzdw>1<99}hu1s8sh1f}dBLpkVaWBWSGy&w zd3Gtf+m{DyZZi~a_*y{T*2axn@p^+zW*Z}b$pW@pF!vq6_fOhc;`Rp2;P-n!@GE%Z z`!h2NSKUuaTfnJPT)B%!NB|wp1G|zfpxKYPiF~Q zpF>ulPgNP}9p?e7;&PM35o?%teRWnDzfU8@z}l>8t>N1pJM79LKNN?4dpP)r2DALR ztV3V$`?)_3@T))21xmAXzh;vRpoyDgSnfzcp@-k1akD6Bx++>7hNU+ll?n?zx*zsnmomX+c{3M+UK}3A8c)fcMHE>?GHtuLI zFQ8*Pb+m)Kalf$3TqIlZDPuvT2Rwvd?)G2R7jqKC*6Mqd5`IbQb4xJ_TR>$3lG+Hr za1I_fsG|&`F}E;9f|uE^k<n+EF97jL{W310RowZ9kAWuRG= zvWsVNUO096$>UnZ=u@~#B6wM^bPjtFhws;E?ppFoeu)1~#`<(L(EH~Tf|s@53-(1k z#h)k9Fk149Qf-N>gH#gg%JJ|PQSW~1$#YH{iARr^adWs8%YtUnYo1esf>1WhlmX6E z2!GDP@ZZV{e$KnUpVl^qvs&C48a@=Tk1@UAk!=o*SeNX$GzDI#v~#guFo#15^$Lwr z6kvRwS9-VI98O^zPG{CrfKT@P>-J6r#9i*^T2xbjMJs0URwNbVk;{TRi+O>)X|O5F zW*r<0cB2W7Xt9^GpU!#CHa}Jkwz9kBI^2-DIeKHBZgANjg3j#SVq! zGf5BA@bkjsHAR1DW7q@An-4E;;?Fl4vZ?CSw_-YLZx)&Ql8{eZ65nAPs<60#%In{7 zzZ}iNLWRt;vCQxf#=CI87)#8vmnueMPRy>cNw{AOKXvL!-txhGo}E`gab9dyjDFnx zO~=;2!T>w&m(o2ITQw4uG4e!D5#blQwY%Ec^MY8v=Q@n=%Tbjv>)n0$dE+Wt4dEAY zKh{8AuR#=gDPIU)_~zc7H+F4DMg5bD2wtqlQ-$2c8&IxG&YuWgKBgZ%-msg2w!OKk zMEFH{t$=0z?qYP$#sGPOmy%1j+l=qy=K(CiPQ-k~xY?L;c7HmWzcea|;z-;*l|HI%MuK;h5>@v^EMVIU_0(GlB#1tdRU^V}0r`l>dX_X2 zyf*6QxuS0YexU+(dk>Hxqd}qO-4zSq>&Sgji{Se=k5`^{J75V~MSq?h|H=c(CwVow zo2=lQM3-A&9S^u3yL>cr`+CTpxAGaeLxR-eZ8kH`>p|jydT8kjeyF>VP)@VKpZ8w< zKwJ*zp}d?o$K6H`44XbKq*lp5iJNbFX!1%>^xl~8vOyH?KI;y5PHo0Q`ux^~zh#5C zsp5+s25gA(sSU^PpB}|b&%IlBAPRrJaq~N_Q(IfHjlXm()yO2IoRR#wsgHrJ?)U_meCvou#)1#r^U<_uVC~3sKmK3-0sTxL>w)MOlC2^TDn>xZI=3*^RO0Ejl_KP{(Oo=heQ%u8IiQ)7v0{iRrD|T`LzAib^lOtu$l^!%ry`KB*Xma z?4TK-539eX$MZtmCDz3773N^~AV=0CgaYAT63_D_m_w6(S!k6i1ymR7P`6{|P~8~8 zvPv2E&8w#e8}a{N;a>0W3#;8JV5UoB@~T1LO}{(?X~z9h?(D@bN`;@^z9Taad7LU!2;sRYDU8^m#vFaB|te0@V+`z7skcD=$;Wh|iN?&?IG7oM)j7`01+*yKx@F~Tnf+c(v7_Aj7= zw7DmQU+QJ<%B z>K+EVdhZHl!Y@ptcjXW4DMpL_7jb9)P4ySHe>ul+%=1|0nTRrF$Z#m7WR?<{hZ0hv ziOfTVj1fghBq5pFdy|=uBx9kB5tSijif3!B=eySP7kqzey_d6AKb-5l-uJ$*`@S~4 zrr3^N)?B1mPXRp7(D3UYUl#IxYrpx#17@S?e|)i0|E|1F5e1G^HvZ#FEy;NB#*t8v z^G5mg&ihFF1L9+54FNIdmW^NPnkX^?NXHvU9!iZ_=Z(E2nP+ae-bvYo9{)Ik^cq>jpX2@5o@ti!tpD^;pxT(i=_DT1?ouEa+NX@1 z+b(fWWeTD@az48zECo=}IW>ZF4esw|HNWTSdJliDCsAn)649ezj;Ho)`2FAnMjE2$ z;G~==LH;7%Z+tLSqhZ?(4VDy~t#Mw19Qc(cG*kxTCABv0ej@1%xqcN{ZAgaRg=#;; zF<-8av#9u`1;RTQlkP=ezO+7Z+4F?X9<~`oE{bEl3^Ac0%M5*34=>sGW4`Eocuzh> zBLhnx>>};(MFzU2O&nu`^d}NbcK9NEAvS5%YXe-B;(WQom)D*`4i>c2U})bd-W|QD zJIC4bvwr{_rALi+^iuVSl$$-{Bz7P*Dt(P@M0tVZHe|$*}cIU9;4FwFd z`v3TH<&zNm#@ArrF4Sctn#qbzYq85ciDE~6>bJLLg0vA$t*2sjE)zau7jwI zq)Nj=8PNl)){Gi2UG(4c8;yau9}{-viGMMk_g{3A|9nFg`%Q0B8cX);qUk@szE$G+ zyz@t9WG7Q}P{M`ScFixie0wZ2%BI+wS9MNy;u*T8dHJC*bG4l46X4g zx0O(=zR}fHz=l#zllnN#`+Z*jP&W^nO;R$J@83L>#UQw18%Bj3=F5Dl`)P4r%e(c(6Z1Hq zsUb)P=((BCVJqh42D~I`Q2d;;&kaW^Na%-y?rNBW6iOvDc zm%Y)+LwyQC@VMQm*%Qnck7E=uudD4LclEo!Em$vI<0A{KL^5pOU$ab$`QrWRy7=xG z8Q6L@C2xl>@7>coS{}3EsU{+?D%Q)_2UQ23PT+ZOr_#Y4zIcwD*Smaf8aUbd|K8C{ zdqjuwJ=+i9cAje=iSKd ztyg4pv4%T)gqI#A4&1A-)ghx(pQEQmTG}M)p z+CuTZeS=&3v@_ms{4#%j)~Xw7&B%s^P~-Zy&B^#mUm4sCxR%)Tk)&fn)_5}wlA-a* z6p7=QF9ccr&QJFP;XU1Mu`8G_3cjYxmQ?oe`+pZ`g|S|E%DT$o9J8nlWwKwe*YURtozIoYQ{ymS@UyFUhWcJ%T&+>PxyIUdyC|I)mp7vGC~_1z6PFX-age|#}^ zHe#CO4h8PE#D9FL`ZVEJHXRK1O^5L+r?Nyg|i*Q5&5s6Yw2tWKJ+4)5uvw?8JzOlKD!(pb zr3NmD~yQvl#;8|TpK;koJi*&Fk^n46S2Cij)o^%W_~|p zK|bWkP2YG$^v3sYA`d>VjVFY!+us#NOWpKN3WvCn`^{7F$2wS1O&ZssCwcie-8=$n3qY6rb60h^xY~pu>V7Q zCV$Ri9=03P#N~G4{YLj6oH8BV@aW-cL+C-TCw`N3ME}0cJ$yU` zuHR}+VlkoBDbR7!KYt(y9*fO+Q-S%i>lQ}^6|{#-`(~#*v0mb}OD|~CkYVNp-s^Oj zFDE}aMzh_OffC&Q5}{ZxX=~Mdj!)PiW7Uko4qwuG1S*bv-2ivO+-7(9!j>UiBIP^{ zUZqR^-OcMq&{tIQGC1Y7`hcC~( zX~(24V7)|r{>PW>CFZ-y4fs4w-ulNE@AMXt^h@zTw)@dPzFeEhECj?Tzz}Ekk1xmd zVkR#h3k7o@EB^7Nw&&pKN9-ZMqJVoXP}>Ymx~6vO&QPG&DSc{7#fFG-{@pdTJ`x%Z z)&0ZUgXi_-qnOKS=+Sx4qpzM>l984B>!dR_^vG9yIGOt}8O@zhDf{6>k9;m)6rR6H zMx%S=Tq+H5|2AU(0Ye%C^ijddziAil&qRL=9iTBp<%YF-9~0@&IkG8Hy5caBC>br} z6{AB@2d+4GbYY%M71T>(f4uve{_{=IqbRYbH>N_B1#R0Zcw6<8(WbrhnRnVeh}b(Q zKloh*@73Nvk$zDSo$0?b@%FbMs&Hqbxl_oEq}YNe{(K&S^Qvp}h@{gbe9yYC zCK<~7S+?WAe2Knv+0o!$AWUFqsByu3N!8M7qoJ^eY*rS6yRlxZtogreCh0=~<5wp> zV7^?Bq`Gl_TN?h7Hl*3%OXfrI^%6ri_}!&-WgpfH`U@VvcHRK*r}OXa@Wu6KWz=`d zX`o!TghvG_bZ+)%*EchL08jm{>hI`9KmK%s{-=6y^sr_0j$Wh<@BUhnD+4ALu4_Ad zDUXdedZ7Nbahd0hKM`06(dr1x%|cQs4jASmXW4x#)wo? zlpJmmbkU)G6Jehs7*V^=fzV%;c%3lPcH;mY?zf14^rKqRLRV^^_&yO~Mt)-lzUd69 zp{>il`xIWYAd~A99>=rxp`Y0wiq#m|kxe{7$KtR6YHyP+mG#AZISm5_ao=z?Y)Q)V zf&dDhRw_xdqeAM-=L*|was9%S(5&V@0`E?9>2Qtlpbp>919=~}4n2NX&(kGvv`~;hh0)t$9evc zbczLiNT*Cw;GQ8*J0r{&;njW5Db<6ZQMId76Xr{LnbfFxqdn|;O4&5KB?QfoR zWawRRmB@hkvVWJ+ciKc582Zk%V23aEF&{lP^V#4Yx&Qoj_+msCzIXN;-nY*aSl{7G zjZHuAu#iuDrCQPnkZ_5=7aF8yLhFA8BRo})qyAa?KRpF4Vyy42x6;avuJ z_B+b&@a5LNQb>6b*Sm%g zKN1HI5{dP#qzqEUez>UTvW+0hu{03pO?<@qVM{MmPCnHvT%!FuUqpeo`Zv%!(YiPjyyFgM#sai7C^9!)=Q?(hXD zQtmlOF%2gC4pHs!h*a^%e&dC%LpyqLyw6&+r5g|O7$pAj#Ta=WXMmBw zExohppFb^2Pr1Yq3Iw#m|M5j**h5JeZx9JRJLg*d;Lo84VvBZ7*Emq}esW&>el3Le zV;2+|nbF)k|5ECF9i%wlf0kIoh#0cc(%WTq(E!^@b*B>}x-)E-@=O`~&&SLzC*c0C z&25Sjw>cf;W&H8|g%n1#7WT2^y0Q*xW^Z6u6l6lwIWOPY#%rNtA6<+N=rN)BR2bHa1ph`=N+BDRd-O7Uhc9cobNSIfHo$@AKDr>R z7bmXR@OvK9z=r28-40)zPRsK%o%;ZCcd?i4=w(*An*xb606Eq@OFMe0F3NmS~ z>kaQ=bn&m-=GVI1hFl?)5QA%D)chmK4f&6 zuWePuh#tv5{1DFn0?*G)X>}w>(xU>qufxhmaDEk6lKy`KbZD7@AbH!v5Rpjg_iN1Q zPz1Q4seSSY%5qVQRf8l{xOr;zCpwC9!XVl~A3jV~xemyEliM7!XZ?Z0$7V)Lc00qtc_ z$3w-&W00g%Hmnj?R+bE>$S*o`Fkh01AMf%+1j6Lkp_63H7q@2cV^#ueGC>cjWw(;DLHT8~i?WM=vye zb2lBO%RsSs^zR+M&=v~!zL&s$z_B;fJ9@FX{pARcAJ$8;(LcV3K5z{v(1{1!ZLLi^ zda0gCWS4mp3BpT4|MA6_d4A(6dngDhG5E)q8^dymgX6*Ak>cgclY#U)rM~gr-JKjr zW;kqkyIB*-r5)U7yN?-B?`GAUi`GHk0$hK{cQPU+`q8u?-0%Gq-9mpK_j{+!4UM;j zbkXoUt=aeUj405%Ys2J&4oY=@FF=ux{p|GMk3_s5efyEL++LiIM=cmFS1PZCgw`a4 z8_zK#sqdcxxgM(`y$6jj8|S_39`wCU>%SM7Z`zGrJ42FN!(fq#1gRCT*U*5e) zq4Olss#7(c3Jn+iI-2K6(m5GLQ8?&25Bp=)RRZVmeq;LJwoqR;baZK&)$ya(`N=U; z_{O0e{*KfTELtV$Am`TJYo}7+x!4y!YA|0o+&Y65^n#!P=g+@Wm@n%Mi6%h_`kvB;n!02|8LOr z{|hev-+<}={rCTO&~#7SlhtGEWK{HER%X%`3wkEYTa}!Z-QD60R6d4^Bme6@!T)ce61m(5P$wp%bO!8lFR#|k#o&VQpbbc1A?LLOP7CREj(yOon}3PV0q>{9U( zfwNEd-!Z7rfcrl-upKKW!TvI_3oOGDFhY&?Q=`N(V2S({lsn4~)!u&lWl%5*oa@`r z8E;CM#OF9P70?cjZkCyrYA%5ijbmh?hALooy}%LM3K`8uh)w-BT+YrT~->IeR! zPci>h3re}3T*&vu zh1CrWWu?+e@JcRBNXk11bnGF=9#E_R(Ee`pTbBS}@!{Iofpr6vcZ6@xYdnY%^j7Wm z8$76fXFOVfuaAIe=sm6%;lV8aw2quT69n}5=R!g%wH3t;JV-7yeNySnJ$;l<@syarM1xoySXRUfb&*wJlI#Q?6nbA2GgdpN zfr{-9{nlsTL<)KnlrA3>5nHUwEy;6yC`rB8@7yn8wEJo5V_$5zt3>W6KiJKLp4OL{ zPQRu=>@}t({I7n)J(@*C_UtK07Oi-uss9CDrTO9Q((B!r?IbZ{Gg9zo@nm3TTp2~VmZ;gnThCa1dyyrjfTTBvgh z0RL^TNCe>^-zu}8C!80Es4pSZS2c!=PJ~NVm`M@Q%A{bfRTenfQg7`F9qzf zZXqLX^9wWuv_vFH|3Q@v^W!q%fprQ!5lPg49{f^7M!$>5-by?~9{I~=Tu8|<1w<&z6SZ%%4;*}5z zw$7Tf4`)aFrRG~K6L3H+XUmLsBLxnK)txHsS%FrYx!+WuPr<1*CvttvXZV#TTbD(< z4W^zuLk`z%hK6|>de(`>u#4VatU%=j-2YvEYf&{G)=bxj93Z5??Ew6qq#Ok)wqNGlNZG^@4TV6<-mUb%!g?<-P{o zG9%L8t|dWxaR8;7B;a!F1a%YgQ#$3)T~J3KA-TE6qeD0pDJC|c@33A5hH7}wrx z2TLXPb6xwEz?q|u)mU1qfB@}V;ZwE~VEsvA&zMpncp1gYeO;>`SbFUbs9ekdT>kyO z`ZZWDVH#|-OYtCJ^Hs2RPXn-K*50=k83DX=m&CkhYC%Kn%V%{h=)5^8>;`h&6Z*EHg12pAc=eu`=fI|GU5fCC2_G6)Po$_lD}qk)=!|xHvLjVoo}GS#2Lw*;YHuy4Kt11gWhPPlfexlN zjAtH9LArf0&-Y#U0?p5`jPpph!Ns>vSHo8KCBmXs>o&z2h}M5a^E}A07%08owsjAfFqGJwCWUf z;A!p)pXc7eKw@E3?^bmsU_Z~TaN8pYgf%WNny6QR05)%%%5wqWShaw5!n!U36k`=f ze5laF{zb1>W@I$Uaz=xNOhifxtqzNi$OwJzzj0|d=1zF6GdraR5aS_q$Tg#ls?vv3Qo`J?|V)*+HLZu6KWVH64LHLpp5v6|OIN-jlk7zHy z7dFo$BCqFBx+^03$gj>>bZ;UR+MJr)l6#DUPQI{Musjp<~_LcJF3*y#sNQ+KsJbiM^Aie(J8qkSZJn|tV!Q4^0*@|ij zSidIp)K*~`gstVKGIq1Wd54F`m@7uX0k-hIp;MGll;<9m#O-!qNUNeYDzgO2O*@rn zJE{Po>*TE}iwO|7>DKmAxe#cYZ=hV&eqhLahMe;&1N1L&AGlxM0Tfr3=gODkfvIdEFu+Bucr>G1+O;mbjwGFQ`=Uc3Sn zeb6=^@CpE@9AwI=*AJt#c@@(iK9p!y{mh;EYM^v(4bYKa2Cw{@ zaXphZpUH8+0G*OQ$zak$K-O{ae58y4N(iPtJ>E}1=Cn5o{*y33+;Q%2-hL#YmkPa< zwPFT{cZAB9Z<>HuucsVx!F-DNFX-%7W+Jkd52bIx^$*t-Ep^9Stj}rc&Px$w1XJ^> zIPO!UcnQw*50m)f4a6)ZG7gPHhTRD)B0Gjh(R6_TzZf<#K@gaqS zSMPGPiJ;J^89pgzcBJN8rMpLr8CiTb(JOjEff9tvq8W|1qlP!awxn7hI+9=Ob-5UNiQ8AT@R!1~8)`SEB;#RG^^Iwkq!g&*Wt{k| z(hur9xHi+n6%0>Zxlr`n$qEYbU-FZ??FK&ul%fXJ z`EK%*1|;M;WUc0tASqIN(Q8Wr9t-%%s35QmGC2&+suMV1n7CTj^#`LM$%m=<^KD8< zwMY4R-i3B>VnM;{_K77>DBfss_vKh zCo+Hy@m_w#rw*{+dTo-+%E@?QdPm3OWrmg%csfAwraHu$%$%H5Z#@4v zwiDqMCn`|85*WlHk4~$f*9ocNL!v%sZ1w8}5pTe8zrsXzq^RxI-V;rauKkf-4eX>q zHm|>spvDS(|K4aQuW}0h*G#XzI6ev$16&tlkF-J0eU)}UCYxZ%K625_!(zB?lif4% zt^`_D{BgF^j)%WwN1$F?GTeU|N$u+PgKY&ie!KN=K@<8b(i_26@MnU$%IJzK3~FAi z-_N57Pno-|B)gfyrCl!)UI&OkGrzT%a6b*$)?q-XeoKO1Q?1!_OC(_L1-9O`aXH45Gh@I(^KDdE%e%OF0y9XvS45dK|k3DBMTBS+Cy1!nJ$4s)Ae zy;x|W5VbL$M#Il!vk2{gVMwE}Z&3UX1zTTHIv6f{oEHDGV)b4d7;kiQ|E| z2ypx5pR%aoT2Q|k=av=_3^-GN@uoej1Z&Uj2>t;|C6g7HZZ^g@5I7beBQcN z-_LdMqyfr!smo6B2iHUTw<$gGGfnbXUGg;|(wrVk+umh>jyd60k0mvFAIu6c>5~xDKJQezgRdc6^0-dyU42Vem1E;Q)3H*qkg2_fI2UdWj9N&9HCu<$jRB)_woJWH9{B zElP9K-U=p9WoN6VxIv~P4i`lhG$5zKMYj zbU*Zrmw-AEUHtvmhX=neYAMT9^`wdD$(zy1ju70BF@;gevP3kv$rmr^NJd6W8@^0h zM3k?^tn(c6DqWQMWxOX5HJ(fg2l%`^TC86sJU~Rgw`0s-;EOooZ|>gRJVu3Pr{rBY zpXnf5jTbXiIB1lYm&kT-UIYDd6jhO%V?sIEaBC$+8R`2P8@}!5M3raK<2{$;(8GpD z6b_wyDC%OocKo~mnxf8AK9SCj_O(4Z#!5|(^4ve&O2xhZ^w52~4HtgFdv%%JXKJUQ z-gEo?q#O5$i++wnHo%kA_nr?Lq;U1_y{nOh7~Z$atXnI-Tm_Zgo1y74f+>)WVN zaWb@;jO@<-;0JZ9UhroK--5Ls=s%foD|oGif%IX_6(0GmKd-z;6OLOq_pqHdg?D2< zt_%c;z~V{I7*0nGC?Dl|tG9&&iod~qY3b`i0!qhMj^16>Ec+H z?FW+9&qo8-Gl1o{2USY>m@fm)+4O(n0nsedEViiuu<6-snBI*5v1;pU$^Es!K!s+S zBQO{wj)>e}OsfRPm^e0@1B1Zd%{aP+-4$S2$UbcHVgNY%eK2fr%?QQDn}>`0?n2h3 zx6+!J43TQZ>s?~51T-v8N^^-YK$lOVmX<;Sx}Y;taNphl=~f5mCe{%Uu5eFOS{NV| zz(&8Rovg$rCQN_PlMdf-%P{IlKj0zje?v37s7k;oMTI`U8S}+W$9&|i$+6t0NJ3IBB zxj`KsCIg{44HzfGvR7%)6#5^KWJo$K0wpeN-n&e#3Dx8T^Xs!nFqk7R%bZgZian10 zc$sGzjZ5bbcDuHNe67&;PPR*6gUkQik@70g zmo=XFCS(Gv9eTs2Bv=SWw~BpTUHd`(Pp(h;BNNXy$kC2LAYFTdKG2{7#Ady@ zn{zq}RD?h2x>J!b=3dQ;;uY9fhNGcJr53_Vm9VMe1mV)D#4ik}v zNbsBuHoA^z2|N@!Ktx@+= z@-Suz6vVu;rMD=V#lgaDt=9=8-0%5yKK2a50elgM{nfNY5;giVWu6u-uY>m^-`dZCc*OmJ^2NV3x2k+qbXQ;a=a>zbEe{!>i^dSG5QH zU`;D^6EDRrDA`KqXB%w=ssDcd)bgJz>xoOx<%(xsSM~gKGOxg^v+&WDNzNb|1yoKt->eKMVHszg zqHA0`(Br5UVBlQ>uicIdFMOy1uLRdq`}J|Zjxh6FSi2B>DXRi}68(UCqp!&2ZwBCW z8xFdijrnq9Pau9rgVDDEX(sgzApfsJ(#7})JnFcqliOVj+AfATCjb+0*IWU07GKV1HjfIbeD!~bx7LsA-()_+YvIi0!- zcYO>{lfldCudfK`GTo@GP?P~Gbfo0 z;spbSeZi4aP}pg_ih^?t!t9XS!W?bzn8eEHm+}^v8^D{z>Rb#X=1P=5x|hOdcrcMH z5D)v_TSRYs!vn~Ben03c@`JIv0>5z>1w+o$cY%Pl6}%a=wssu2!Amo>d;-%NkSOue zxMU*&JJw_3r?@m>&&4h2{0AhcdD`Z3i;yH_t?buH=UfJoImx>3c5}e) z$neAEcSgZmLGx6{R7!Y7kT%4{sU4V27Kd3lEP*^`Be#>KRp95Kh4R_R3E(#KmEUo9 zAuv1ewnh7VKX`xa2cyl$48VK0-M;B3=F4UPxo|QbJiDM|y8BlHs8L#VV+e==!@F+gS;^zG)e&I0S(eR&y9(Q~|yQ?ERf_DgZCqg(`nq(LjS$ z1;!sQ(xCd;XO_?X_0hpY@;21vxL@BVWuPigMy`_MHf@%;zk7nxHQknsdIMr8x3!7r zPxLh9Lu)*U9I@Bps16aW1?tiL+DArx9YTVS@#jT%oqgvqq>t1uv+9<_6462v`v3>N z2&#PY)rkdK9OP>q=i)JmgOqf7IqVgw@kKw2DH|?_knWS02X1fD;&b_W+o-=fx_RuW z`p+n4G_!DxwjyXh4nWDL*XP5FF;6~&qm?p9xQ`UFXMqoyd#_T<=QWGvagxz7#Glk6=inw+9CtmTen6YOb=8z26ZB2&(dtR-0Q;)@%$(Na zL6x%mDEpfRpkD2<_fBF2$Tr-LEN!m^#|t;c8LtO}IL}+*PwrKM?_I&)&)*CJWUG9ekddnJ8n6H9%72>??iwxWBUeV&BwFy!g*8(^yrE`zuWHJ}Q{3w3sxe$FMJs)W?csww<=yA;Kv`18lYZhbX7WQaa4{&#pE5%*)J{kC)b z4bfuLBca`yR4B1rU`+mu0jl~*CHz2*9#IU-jO$<{pYw!Lexoi2qU;ZRleAYEMViZO z{aE2e&KfMuL4IO*@g{~b*@zuwiJZ$yNM}Jsd@-WR*%au`f}hS};3_PxtyXKhJq5pB ztg8K`JqFVk9zL63YJ=n2N6h5LTi{Q{ye4esr?zMcHATDi z%`HWsZxnBWCDze}@{#8;DI~c1nBG`JP7?lY7cHn`TLwiEr|(#ab3kqLY#E>EQ6Oat zr_{12;bMDk%es9#*ox^_HE>-5CujvO*gUTSbnGnyTX9$~@oeIWyoI3IV`vYzUqAR# z(4=#IAOoCYjedCfcLz8S3qAji$Aj(-U3#0<1|S*mv+IX{1mK?hl=pj29bgYLUe$IB z1{L@IOv3I;a3dzte8MgWm`y$NOFCWw^i3byC0hr8;FrBlq~Gf3vuk=8_(_AJ<#|ua zYU!iycE`_eyNIZ;O9w=;lF`YvC+`HDaep^3?8^`iqH!qu?KXnz6Jeik3E_HV6!B=X zy7wp%MXOmIKFdQ!nQG;izn#N_zK&GfN2>G@DT|Gyl}SXA!5_8{VZSPi;%EARFcn(Y zz2n`CgG+$$z69A*)F`t<DZ;aBsuQf>9`}flmwB%6 zW%xlaO*u_?&1J$fTG+s!G2pf5%Wdnt{n4bJV+?85gV++104Q- z&Cik|K;eJ8-o0t81wG@J&&Axp7m>bvzV$bz5@;)^PQ&0JFmpS6p_!!ud`@CMv3fZG zOro=g?yni69}Iu<7sj{Y-?0OQ$6gttE3ACy=hz5nL%m+kcNqV^KL$9p+{FFeH|NxP z@OdVaw4e14u4f_{iyv>^H$WwS{m#ngU_Z%(?@G;M13btX$6$&3!>fk(GM(CSeZxu- zeeym52@echy2EXV7If{V9r1H(`pwamyN0Np#jnF@n1CFnk9W#17@`{q+3R*^sn7>5 z;c_bMpZj0a&#s!KLnSS>_NE28=>6sJw?rWhw4p!qGP6z@8MLoxu}t!!GL7Hny~AQC zqFU+bf*v~(v}>^Q`^17m3I{W|GboU;n5evG(<=Pg?W)Uma|%AceizKRk3oOknw~sT z8@y_r(GR4_VH8P!bC5e8X2z)>`!1FW6A!r^OL*c3*%yU% zQ#^yAX4cwec)|+0F`HXBPrJc^DDGz7VGX#vdpW80uPJ==;K^qtGZDzm^)#VLRTI+u zsZ+d|NP^`vK;(t0B>eW-hw%&ZGEgWI8XDQl0iO$0Gwux^1$UEV4>0FZ!t+-SF~y&0 z2hS8K=9JGbfj<<>eh$y7fY@>MpzOp6pr@`S-@sJ}4rc$n^fRCzED#=E3hvDS50Val z690?&!bg)}I)*RyiZ*S!zu5rp*G?t|8xu zXt~8k_2VyGe{{GR`upMjZXWe%J_RzGpLhJP)`EzRTKy36-b+RmH5r2TxL)D=;}O=n zjs504gU#J{oFff8S9|M0+&uOL?{Gq3GRGQ{W&Kx+!-zK-2&)?oq}u zt_D-1n2vC#&HybW#aI!%6+l9q+Y2F~e^n6iL5jdBb`~VBr^e=YRuP%^Y z<#8e&5>5%9c2!G;V~_GK@{RjJk#+j|sD^x-Ls`~+8c^Hp53wxT8^ zdg>jamNtdk#ZeKbB1Ir)$-XE{Ar06VID1iVkOZ|<2EP3Dmw+nQ&GaI4mcfNl%J*ZL z?9eika_#YdqoA*RW9!2{N~p^|YnPwW4zz_dsSeUCfw~)A%!9*Kz|}Xm^p_I$r%CRK z35JCrjO|KY0!KfXoi<^Yrp^QwBaiBRV=-Sc45@4X;vflLBhSo=2H>#ontf|Z1Sqp0 za&>@OV6}VUGxe=t5b@_zfOSMAxOeaSC|hU{hubl+ zcq#8E$%!p^;9+ygfAfauOZe2}-apu{)+sC(#OIZFw6ErU@qEm|gwJmhDh&|JyXGRE z1OhrM`IL9}Dz0C*y36?deIL&_abl?>9i*~LbP8@?KiT&$t#5Ty}^F!^e zb}Nj?j|+ZnRt$4qZk5fyDusXYMvvrj#KY}bS-EN5RCuvFJ9p`k9~^MMGH-P)7yuKE9Bz38>Bncl)*iq0X=z=UTf2t;XZKSH8v9w_$F5+QTd1_%sH6(%KR<~ zUR3Z~6VaE1{lDH^A}}t4a)oIz7Znc3|E*_XHh2^)cctE*Dx!p6g+Ez{pK1qIgVI?N zeoNrdTFX`Uf-0aLS+rXx4g1q$ZmLV{g`lLDs=eh#KiIbS&rf=v0X}k5^Y#-u!GXdP zuk${~gFQUwZoQ*;3n(Z?Gs}G<0I|QHdsw;-(6QQV7&r%m8p@0%hL4rtQyZV7)~O&s z7iT;4!m+cFdT3*; z7tcPFdD9k*`@8R>`fPczza;ri=>hha{56+eSmAnQyw@yz7C%3Gb)#PCt3F;>8Wgz^ zL`3!y7Wu?0`l#?HN4o7BBC1$1Gfhd+L;h;#_T`&Xp0##HjRgmHspgY z@Fzn8%LVT0DL0!*Aie4Un6yYt^&rYicA)Ju^(_wEBc&KAz&V{cdcgX z2SYJezi^N;0f+5)Xnh#w3&Zwj*X?*<5UXD5P}%@~&4sQ`rA2_cPtBz$no?Q1yxS4?cU|oCO*;)xXp*hr!3f;H)qN5w3vdv0 z8}U9bK0i71ENx8n7@*=rH~WS#>=%XVFcvTvA`kd?RmqQlex0)sQN_VVd~|arN#_Wt zzVUjehWueve4?hOJr4Wf(*#$;uZK~J)WedoaU!x95w|W^K8(I40(_w?9b(8HvGseU zhm4wJx59?m(Y#>OgW``WsNr?o8#BaWL{!>sbE+Fn&Ip3Xh)~+%SKL_41`Xl{+pNI*&Q;buhGorca7Xf|%T) zrt~FW&K?ao&aV8ip34lnMIYy0Iwk@g6`KpvF(fE^MZ7!Zgd}`RqW>sC zw+vJ(ihSm^Ibh7I`49QmM}ZBA_`bG`5<32OC8Nf&9S5VCn{r=W0!B-jH+`N|0h{_g z@XJH&H^wCDh_Ds{QM%oD)wv%q*Bu*r+LZzLG8(NS=rCW@QbyKB;(>m?k2*coTX3G- zM@aFGz!x#~c6TY%frx>T>urv~K&$S2VEJezi1Bdt_+}LZ-0aKZ-khoc)+6)xLrepJ zA6?sz`@i?2B_HoeuZtu!x#@H$5C8rXF5skKe13_piIToitB)8ie%~*65Bq_C=}uW*sZiLi_~?;59W>XdPo040qg*<8_dSf(K+IPvJA;-UpgWW;pp(LqNR^i^V3)geqCeY$-rMei{eVpcxSPp>V3q}D%$q^a_w8L!qsegy^4GB6 zx48)o9^EWBl4S)UEyemh4_C;`srXG#Llbgdq99PPm_k+X`NV@eBCtJ?J@WGQL3nRT z#yt>Q<;_oh_H&(;fOck8mk${(1DC!4^$)S^Fo5cvl2rdFpg$z)C@D$_kJY{GzV@&k zO!v%W*=#HV&p0t+>v$DV{0qFwB`1JjHCwQvaUsYkFk%^_?Fa4hkDpU8WP%UmSlf%azqLt+K!z^yQb>jA(kX>#rOnh9!MHPcU-*n~Ok->ToAK8#*@U+*u2l!*UC zVEuI^L$v9vu-$AWMQrkM`p{Y<f-IHk?&sj00$tP-fM8Z)Zx!U=94a!(tEvDTjVEYA2!>7p+ppjzM{kzr(3eS`F zo2NAbKW`YPsAfJeL=Rc$O%+1t(B9G0d*Yz}CCg0D-N$fvfyMZJss~t^?pN(i4~ErH zo(CsXOyIL?A30pe0o>LEHQB$&gBf%8rrJS6XdwxeeA5zwU0{^B*HHm_YRs$oBIseV zs`nLxxfr|?Fjdl_TSS$z?p2Ro{g zXyLbenD6s)`=Pp1 z3`#9mN>NFf<%ZuyALJ6i!(Y2{68{sprlUVgkMor~&K$g|f<@P)HY*fqaA!%@i?11~ zctxo&&+8dArWcQCCQqy4`1IFXCf?MzVb}T>u^CmYNMEDB?M01SPW>lOSEhJqT-{?i}jo~@8aY7(W6COb`D+kL3y5mXX^uSAueviL^P8q+j0Fuwm=``GkwK9&0il^rn~uk8QXyV zuXl<608K17rg?EWn-aGhU)MTvQ3JCjI^TC!r^4Tv%=Fq<)v?q5(7ORdzBJ;MZfpLF z(3_-`25Lf2&RuE=VWrl-RC_izbeFNZEfV&h!$=iC!I9`xihr% z@iGqZyy+_+Yok20Gx+RpJZ=cbeOL5~)Pi#UNes zzgI5Qi|Aan%&v`#9AJN5@Ts8J5E2R;aB;wtU?yB}XhN?UnYX_B@bBINN*mLN3eG4; zG18ROLb*hLy8owPC}S>S^ww(CkL^asyw=ZtYfM9F&y_nSIa`smiOAs>{c$L+Huh~O zQyo(OQuodLW(4}F)TUA@Uxfl~rh?k91|yZ@ERm#NW$0JVs*jqXAKFiAz1Vxb6fwuO zT67dtHK=b`Iyn*{cp>81qr{>4>r#O&2Ol2 zw=uIiACWKd_pslPCi12EnJr5eqCa%=pl{jzAZlzmU+(MvOcnc?#)_I{QsXpT9b%_A zF;8@pYo2zN8vBSOv~W_WVCjdu*(@_?a8!4E+GE=P1qp0yiwn|W<(9XbE@@}*%1}Ai zsw)F#ofw>N-E#^*3gs2zD`LT(P1kh?bPr=|$r6j&@qJh#pji74g)q+FtxjoT!HY|7 zp18KE!G$%$uj_k#CG_P${3V)Sw0LFSuKFI)PtBSrC5K?N^`Ruos5{zyA zR1ID>qetP%vOcr(h%3$S--lo2XzPBzs0ZH|Qgf(w`DT!dN)-5QPHc6dH@6owds)+w zVf(A*!dryCTs`D+lRBRGz8ON_X4j&30vUB@Gb2#ntv#nmuPc!~WnbUFuwZnTHQBP- zvkZM_X{fP!=!bZ$zq*BPmLPVMVJ+@pA5Rys1fZ)xB)we_qBB`SRK6nD;h7ENqccFuKZ9~G9R@hcMM z*TR8sUR|->ONAdzCbdqPX=8=xRE5`sA8Q}o(J!Z&wK1=Z&pv~z^qA8+?t3=juLx8S zF;Vwn$K8%BzMmpx@h0!Ae_0eamQwWMPMQSpmv3|qvYUXYs;bs}xlyR>(oDNG+XxOFaO#IE zDCWFyx0Nmq9$sC3*Yy4|jI)PQ`zCn6q-JBP848B|_2Ek2&X~YEDoL4i9S7)4^hLAYh{5wuX2#ou0B`rTtfZf{MpC=05Q`4GjpBSRM-y`ZD3HztS6r%Jehs5o;YPVy%=_bcjHM z6|Ebul&jF%VK)k6t6=o`*9g<=zh&sLmU7ViML*=OSsS0_4zAaBTmaGB#w%V378v!H%!mliY0yzsk7ull>p5advP#Sk@AYZdrpH z*@=ELPZ!1Cf5iLnMDaKudlWVP`>yppIZPERe?AcKg3u%LxAUI@=2h@i0(V974mD;P zHds<=Q^tZh>Q_`h&|uzB5$!3LbC}!V^4%3pI()2U)!O8xEOuo7c=mlh1CA5uyRiTJ z30zpUIr^}d1()*(TzY!^5PrpY?QAG5Csuggb&WhGglCfP)JnPVVy>GuYHA*wn0;>r zJ9ig57R~?Kt!hJy1#fVvKVK%{k)mw^3bz2VTM?lx6BA(CU`QHCBg0EQ-D@lUP4Ga@ zXJBLRC-^w!S`Iz=FyB9XhW%?kP&^s4yKNN*Y!S9A4_+t1qt>W0$t7Z*$@8%LeMv5A1?A$4zO6vK+9ix zNbEFu_fcLH9C#(7{hDGXiXQABksb1Y@D(F+?9lC(dnlo=VraASbu$_nate=_ znn!Lg{#$tWryPyxIfQ-NH-@?)!REPfE;>g?x5sp~3#IY!?)}7(j@T~zdPurS=u7*s z#y;A3lyA7v;GRh=Anbo-Y56U4ddPfAKf0ygJ_PLcIVfTQKMUVUt zdUU(W>t6}_;D2jmf2a>ylaw|O^EbfH$8&zjkFSB2dS>%w0t~X$GNk0^Oo`>KPnz7n zuYu3aPCl(aM1{|k8Jk@^sDTahatjIS6n>PGr1^kN1E*g*v^h)EX9Rk)=@g|kaN-Ys znrB2l6}jPkd{=t36YrpmXcK zGx%e|$EqG*ZXBv;{(9P63|E_rl(F-&W1%Yg{xUl@97b8_Vi!un4TKltT+})Q>U(=P zUY`K{ri8XEn^8y+wW&MxyAgJN$J|j6CFV+WR@Yjj^Fh?$Sl>07A~;a`t09Ft4(6(U z#?^K_hT4}OE&Co3{YKzZ=V=Ru5_|dd=h7x1{PpUM*H#Wdd-HJPZX^#tQDO5{%Me!W zG9zxC6M{ASog~|*3h?UTKdt6F^uQgzaM;dY47{(L^DNz7KpOQmne^8=KyA$E+Eb?? zG`&#Du=gt^T)aE+o>{9Im2j`~_s1@v>CU<2xF_YPCU;WnM(G&pJfdhjN|TEge|Y&) zq;;bU61gweYSR#34!WQzLg1~~!I5aR+aA19J9TG4p&fISofxe$=E=|>}Lf-RB zp*@ztNb5yO(uMglr12{%BIAM|>dj|0Uvex(4J!F3?X`Un&ga)poe{&k#ijkXTp2MP z3%~fYa>AdV))BuDM}tct}z1uW& zmO;BGO4@HQ0B-E)jjsFq<#*e+m}2JeWjxWt39e^qHmg!w5&=;^4L78&u#zv zdZrvTajcA-U?uv~oDLT5mvd3ex492N^Ihm6QXPHAnT|S|#4qaKAoN8!Rj`&W9)*AU z!z7bdi@KBFE#JwGKz{?-F;7}0a`XPuJs%m2yffzYZrvy-lZ_Ddzm zBAetV8s>xKda}OyZW`dE`=8l8$NxbP>12D~87+KODX*eRpAyeGa1YA-r-4oQg@xUD zsIYTX$IzZL8o1s~QFrtr6}G93E6)&4%%@$W#9QPCR_|>TqGc{w=`k4}EP3qWD80g?j zLXYPT7f|7pFDKn6baZh1g$1@pHM_9anc}aXOx5v9bW7Z4Np{?zK7Fn8zccs_eZ=NJ zJ8s;js3fiVSq#rNe_3$lV8<_3Zl*rYVZ-L}+Y?%WBpmT*>1}z}I$&F&K-Mc0u&>*p zBk9p7xcdbESsHJI6snQGk;X*bj@nb5>O?;DD`vc9)h~iL{V=$ZtrO=@TP~Zp8qKNBgNpl$HTs!jRn-K6Zt33kpoH!TvRXH4u!Cf2ckZ{pn77S2n6#WWP81YhPA3vb-PM zLtqb)<&l8s)gy$yP;I7Xbj2Ym7etoksY6jub7?7UB2e?22A78ys?hugKWYn$VANs} z&bzc+hL$2bV{~=>kaVyk6Qye@I>^VDdQH;@l_lW(<(Y$+tw(fpWQLK*pQ8`*5+LaS zF*@$Td>UfThZOn0ez^4emEq0Z#A1@(EDi#;(ul8Jbgxc_Q~`KgEcA zi0SCk@j@e2{3S!lk%>Op+0 zs)CAAhZCDk@y1rZ7R2s3qbdt=ym-u^b@x!?Ui?PN;i>f~I~K0h?jwf#v6uQ?o7R291Dsu4tald$fVJ*_- z(Xs37ptoDE&gah%0h*@mUnNn(nDuM7_=0BiAf9Xge}nUgzi4AJbG{rE^xr!d$2f*I zTonzK&2y1SdjIv;f4Y#rwe!0)o^%wi;l!F{N$5+l!MY25Ji00~V6>i6iyRHgA3gpr z0aMXP@g z!2_Y}KPvT#n5VQ(dU&Xn5*MiZq%iH)#`?Pwf9*V@#JQ6qk}<71cyv0s*_o*C^qi1V zt@@>d&-WdV&-d7c4~nRs@A{#Ri%SFwXxP{>F0r^YASH)S2OB8xTp|{io)xXjBI02Y z`P;{v32@)RldYz{zuE8!!3$PL{76{6_GQ}10~e_eWU3!9M;^YqJ!=uA{LTeq1V7u>Jt;uFAl=-w8mhOH6{`E(g(jy5!=DW{n zX>-BQpZCg0<){fT-KeOpj&J}ch6vAH1h=E$-+FH6bwfaxw9xT0LZE(+FGaml0m$qh z^b_6bAzr|&_d=o==tcIO+xSQL=GxAm=JMbGA-_mbjvGUSuRG>-4Ve@HA;R)?}&#;->pEc{tC$fydr1d<|#%z}~L-penZu9u-~F@fZN zT7D?rs0K58l%k;JXp*_Q4>FyvXL~tB%mapVY1O;##!ph$_*1nBkZXy2mP;e??k8}O z?3jtglIDS%`zNR|12xxife=-!SM%v#0g?Yaj$}$LCh{S{hsGC*bBTV^A37Hvc~zXq z*cTkwPK}d0ex`9fQo&|S{H!V`X>dDpyz3`ovEazY+}Sk>T6~P+NDwChR?}p3nzH~p z{I9tynHi+!B?Ic16b&|}3>dyG6f`ww7`B)iEPCr;eX6L|hI z`+i(dQlRUaIX$> zm|G;~1N`)64`uZOyXDHycMqGOel8%CJGu^fK5lezPvt{}p7OaF$$TiSWeMSQh=X4) zkBl<)lZbd_!KOKXLuwhIsHi= zhUdp$TV@HtFQuw`UFmW#XF01cutpEpNQy_Ua*IMo=M=Nk8Q-H$G()OipVmN~lF!yKBVWzm7np7%n~O*h=L4 ztLFRfs9<#Kh@SGn>t%>5A?sycvLDLjaC@Qqw*+O+3BHlM?}MuNeLq~^G{k=G+jmFC zS77G0F?ldY3rmSBzftC<#OAlBekT(3*n4*}Wa^20rsLGtsTdm#T(E8vEFwqvbyw~` zZ8Olo9BfoA`l>{HlkT9u9Z2X+svSkmF)FOylh2vBq=74*F9+*xQ{rAt9vb>yP5f5F zLRqpwaQV%!z-clsNlRx0wlKoQ|<@VzlOYCj z>7uU$Ru)i)5!<~?KMsgfe@_3+Y6$VaSX+KPN(nfi{6U*yGy1^fb2|U|0-8c@uXAF{ z5sRtf3H8roh+m{S*kdaP=?8!Ci^}gtca@GS%6v>igD2N%P0zFiESDdU{8x5SQ&)8ORj6L&r=s9;wKWl!%# zqP~0lX|9!%Dkg87u!CV*)jpTkeVa$kU_g#tEk}NKnd^EaqCYM0>)goIT;zJP;Qh*vE_5ZYTH8}F z9kD9SB%ixT=u1-V+t0h>5gSkSkI+Z8D88`?cjQJOdJf)Xl?Rn*^MHb({lj21?e724 z-?9ujE|@&qdhCbJD|5~&{whJ5oJFc~Q9fuJj~y}bH^MGnyoFp7%OG-Ua`@^n0Zx3j zTR3cs0y77Eb645Wz)xS3;nplArlmXYcEL{rAHe^3G>Ll3Pv&sOpTz(Bg zIcUtDMEK$EN_6J=2vOk(Z<_Yfy_$H+cd5LJsIRERm=6rnYvJ0%LMvuZDRIs%?=Yh` z+W3?PdDc#k5?4-=TVK}c;*GGxz_}$#Y`QA@(^2?5uAR$^nd+n?>Nz%=E7KbIFxBT@ z;a}NsVXW-gzlQR-#N}6|9k5k%ouUMJapQIyxn^{uH0yWz>jjkXLG5kf!*X;u;a$f>Gtr;^ z^7Fxq^&Ir+{g-y@;%*e3PJ2t|LmJ|LU#_-GiO`qA4(7D>ICPtHntn>O4z=0G%I98* zKyUZqCpK5BkPYMR2NNd2$mst1yPs4a5QTkFwxOCI(pe{T!><&b&T(oxuk3?vvHizz z1T$mNkq!1(S0;S;U^iQGzY;#%{IcNJ4h_jPqz{-lf2Q@AW4hm-yF-1Hc-H}UxcQX_Uyvnn;or(+a+-xH{Ns_Ar@Kshf3Zy z5W$u|tY712tXRVOG2?SJUc62q_mtmpVum*ixtwV{__SK^y+`%DnBCr?;*$JZrdFEO%So< zp}8CBz9%02jeaYl8C8pf!#d37-b5hwH2#5$_bO4aMOph;Ofb4xo+sITxeR?YWJ=M0 z>WAiBK>oyd3Hqa`JfQT@2c5OKn5DC6gs*YQj$fQu0?TQi$q`O%yo;?jy?%rOf4W!k zZ;n$F7j|)J^NvvBGvubF-EkT?IsD;!2Vp9V#6$f5cM)LuVA75}5wGa(SMC-s)WD^n zYhn`-S3vs#N7Ieb4Ld$@tcmPCB^buxOdM~o<|5J7LM?`yRxp0584KKJ{6_J zpC-#vPn}0ckaqgXIT+E4x5PKn2Lqr z*zDzu8&81aP8YG=#RCHN)jb;&3W1c^2D=hL6L6|Mr~aeX0VwarE>ZkP@EQILvRXwN z!tct@@n4S%fs91s@YXt@TVFXwn3vTYR_ z{ua;RcqtflE4Ru#+Vuey(r)DKDEpyTUmE*m14~h22w6ts95H`ne0^Q^-VR8_PT%pI zVZw{{YE10eO8EPwKIvDywD@q~H7D~56|CG|e@&36&)$m6$oO%NsJGmr7uRW}#vF3o z8qVs3Kezu*pg|Kg))!M1HY55)lLm2N=F8Oh+v=c4$zc`T6e;L9V@-n-&$a(jf2o9@ zUrZN#cb*oj1fO`XUM!Ej3Ic6wWOm`2=90$7zEfC0uSV$b{@s{k*7pxjzYq@MR&zM{ znH85)v+U>J&50HHT^><i4W#5evP(l&ZS)@?u}R^k=NM+3WJKYZ@J+;3t zG+csWZN7PFJ|gZ5T_wjke`9P^Wy}>hu?UGd^nY($)W)6>*?Z~hDX?GE^Po&AP250b z(($>25<9=s^lK&Rz4W?*`cd3eI7u~NdN0u*imn?>EjmbrKVCbU3XK|g#_-+NI6W1n zFD_I1cUco_i!B_=BG0t92o=svO`L7$bdWZh4Zm>OE4Fz{0iUo4 zlFww~#)jw6I&GggUfSRkI=#q-MR*NNHs;xJq($Vp_l_j29B*s5_u?inwcpDPRG)yw zy^n`YuZ_V6#yzFbOw0uY1|OwRXoiN@gNTzoAFd33`5xY01YXsAV!5NSpf;V{{W$mu z%+v32J#dHMc|79O&X5cN@sk0y_xMa;_@T4M?`{W>)>UlY$t3s;lYvqriH4vj`SfUk zgb)l=m^`7`QUG-!scRzF=^^Nhs3d=l7?8+JefOsqkkyE_opcfh47{ExrZXHu=MBz| z@T^k8$>8WKv69WGl440euw(&^=ggGjh;nq5Hi=WLo9IsqD5g^`=b(JSb$@o$jq;4& zcX<@0p^Qpt>~Mk5m#^0?RlmfcTFKpk*N)eroVfscSF;F|x~q~w$EgZgKH+@$*C-g7 zZ9J=%V)}rtzEUgBKj()ip9?PoJ8beeEVh3n#%dRS*SvMCX!InG53SYuWw{%(S-WL}wIF_OBqILeFDq7heg9PI zcQzcd%;G-q-Jg9IkN9*Uy{4`O(Suf4>7dIoTxaBGkKBLadU?{5!U zFQT(KST#z99Yol_7&VZ7AkNia28}}$;7|W-&ueH#NA>zC5kc-&jol0E3cA?`br;B97(~*8z@xU?p7IgIIAGs#hc$De(sHiQZ7U^z` zMIR}QKw`NqWfR-z6`_#BO7%((Fb~E$dcVJawWqLoo-<6;ut7FlF`?h^x}N} zAJFvp+3EKG15MMCI7pXH9&^2Qhl0dH`aiz(|NZ>`&6%!A(2l1HCSm>nju|IGR4uMENvr8uXgJ8IFkald) z0e&83O4#3F4aI(*s!IU|Fd|ns!$q$Nr`ASN)JbPx%qVd4#fStPv!b=t(<9<5@d+#b zYBpd?;c)!^ju9-ib@J2 zR$ikv)Ok0*Xjqnrmz|qBeYUF*Ir8Rb7K3cGWpX3Ys4yRWpW#c}Bb0!QW#{DW?`0tE zDVB}nKA~vIfi>4LG9IZOYFBOc^Fpr8a^Y|K!qFyg-Kt@S13G@%2q%ZG$tKKT>xHQx^vzFCpQcSzWMpDmr(RzDnf$}=-2KKCa{ z)r3yN=Js_;(5l&wCyG6E@QC088!xWj=dLNQvU&}2j*toQWNlT+u5w@E*Y*Jj@DHg z9|FClJc}@y2AKFMKECy;6*kqFW{*5Bg1oJ!4>QW8V4%Et)Je2`Gu=n0-&oyG#l?>Med`gFPx!oyqx(z<2V2W&DjbaWBomF*uz z!{zBRK5HXrlD2qd)~f{>RXCeX>bD_ZlaLeir_0cPW`g0xn^lNP-Jw3^Vm7Kju<)^| zARl=Ju$`(GOhAs^Gy4}pGEm@5;J`ocQ1pyBZ2WaZJYtckH&;{R?%=Xs-C`C9q!Qt*T^xu_dr*{+gh3Qem8yO~3Bpd%Wwx`D(Q8bydoz^!( zKOD=KvuoW@hj5J++Gl`@tHmmHYDu{5;L;QSNkT6uL;_=CNEpQ4nqEIceBWXfo84hV zKJ-ohn1K)(cAXkK+d$~gnU<~2uY6=MFFeQmA90_YjSts%^OIpa>A2K2H{!nhxFC`* zLWcEQA99*4NZ3=?zJ*(!47y(~_TQ3>V=(g=*M9c=14bvj z%a7cc0M7BAnk!$)V963!T|)Z@`BL_tOI58PXLg}IptJ~%CXNXn(j)XG z$#ZhXAr<&70#b&br$UDD{1xf$5U^2%&zHo5!2FkZ;mJP^uzHU9=bIjDP$=~|oGyc{OG7^HWGlN0yHEh6l)+t`Ej1dOzw5H`d zts!HQ$fcsM%S1id*@j+Z6fs2oIR1NU1gXn;xLf$Qpfg6zB@vg~P*go-X{XF zqNJ=w(M)c2)W+G!xkg>hz9b*TeLLxrDV%`BHy1n&A~Fd7my~SAolx|gSK`$4{dn|7 za6MK(zzbb{EU;!r4o8X2~o2eeS?^dHZg0QB7SVXj`aO{ZpaxKA?NHU+2+q~Xk>FH#fsYiwc9({S1FS4&2A}G3yFTn58~cw zy+y)m=+e4hUq56ePQTyqChkw@zW?gx`ic3l$JY0V_vdHT+^qa+KWN#VOyG5qh?LQ)54xnr#AZGO=ACr=Q<=jF}BDxKtYDlsl|7{xd}d1>0T{0;=aY~ z+4WM`aR&8QB+}{{0he9Cl!{WHf-XO4xjlYuA z?f|{KjEae^)^JGK&C2t(0noglKHfvA30uA$B2QLjA^dw^rd_`ToGB8z!>=s_DO}Tv zP33GrKU1w_noHaluEzAiYisDXZ%zKu`ekIwp!#E$e-w4_pM12jI)Zq08@}B3XhB6M z{!J9>wIP9%Chty4m!UWfAL;4!DkMQ^*TAQjjbi@>kDhy*kM{1~)fBKl0R?+aQw0QP zprXaU)iqwBNV9BZ|73VP8aV!HRn^xE`EQx9&h>;NTbI0>Zfy?8sX}S-!>a(4;4mE- z=ywfSx|!_K-0(&j&tG>;Gn%5^#eZxR34NK2IkqZ&(FoDzH%L9MaYJ7QijEEKH9$wc z@6&f3BjLlVH)g0W4M4@n&u7(1L_9;OI%jS`22LjDuQB&Y_>^sCTkh{?5;hZ|mUYx7 zgNTN(<1snn@7-Qp9El`Dz6dH&Yu|wmsEDW#8GxOpSf@(sZQyNky;l5w5O&kfy^Gh` z0Lphp1*s*&Fd9ntK9X1nvteJbuy1V)3}2T|wLkd-oDSbZpLq~I;uEjNe@&5r`ijD% z<7_`*sq?~vI==?+Yir^2sB49_t4#H(s0d7X{;J9wmV&-TnrEwLD%|~dLD%6`D)cjGJYv}x_;<=9YWfb(L>~ovw zDB=>)?X}$*LChplre~Kg^cB`e>`p69Qwy)l#%P$nIc~z_wJcvg>_?XKQ1HDjh085G5KsXvY z;Gmn+?SMKxR}3Wc1JHjQQq?R$*HH7miY37vZ}j;bC9^r3Df-3qUQVgs9o>9S5hQPH zgfy<-UF)uQLz8OTg=cvTkcpvx#}hsxex>?6-BH~Ks+lG~(usI=aty3PYWm?pf!w$n z@%=N;HmF=1?+5MpFsEW7UNO4!R;FjVACv;inYWyX{Q9W0unwU==Q4N>d0!>&%k8Rz zNyPcG<-t`G4HDj=znuSYu^)E1^P}7BB*L#TWr}FYa6fx2E6jKYysMsFFuhEM!(Td^ z;KdefTUqKwcn?5ivO2~2OY5L+<*a?^>kptk8MFOg`4XrO$qI1jjKQYPubr5Wzu^vd z#;1k76Y!nb%JNE$3{3H-^0dAT!TP4GxQH$f2t|9gi=Dv`1Z*=l~?=2c-0jki<;4XB2irjY%&!wmA&Y%U?7T@zfMJtogfYOx*w+sLiXM zCipIQ2YMro9m#NPK_E-*AqkKEn_l1~;*S|d`lF8FBz)a4O!LWcGT40b_U0wx8UN-5 zT4q8|0zOfnLB#tttx&n}OO6aLeu$V{wIN}JovfGT#C_OC>*g<%NSLO5|Cs$VGL*!P zmW_|^0F%>Mmi<}-aB%5Z(Z|qj@CxVK!%VQSee15j@bKJ#&6(-hcU{8}8CRbZ`gR%A zjXC$J>>G!&vO<=D>OVx>bMgW+!MWP@_oDNn8GyIBzpUQy{{a4&GUM;j4RBR^>2Pdw zE0o)Psibc#f_u-G?DR~C_~k$yDK9V;$jJhvxx7?hnWg!BYBU5?Ng^_rB!l2P$8EFV ze-5zAkX1o;&>HHe?=r1CG=NlRCD-)5nvk|ex|79v2C|GB-t*2$0B`vhg+hX_=cUt? z_4zX!2v62jG*vOel{`IO4)-;r_Q^ziv11ucZuQy$@gwGv$s&L5&3GR3PIM(2X{yb!A5#TjIm3MgV z{$AOE_s1_i(Ye?MWRwWbMEvo9{%}iVZa;W5MY262;*+ZucZN562)%ge?lnT3LHcHZ z2cb9eUe5-U2|YTq&`x#iXFmkd7K&(ClJMgOn(F0=exMd{&f2X)!h9z_o>L^w$Ny=x zQZfKnCL3ZVWc^2t}z>11Z*F@yPqM|C?%WFBEmc^O;n4IQoj>2cugZ zP>%eOrx7m$(1W>W9(lgRdz8?aTeaqmY8$5nE$K~B$)Y_oJ)ti}oc&nxf)O%2FEN!ZWgIdhUI;O z`MC}xEZC>Wy339XeW$;j-=|6FmH7FN>0&Z`+@Sp!wzdNr`|47|ZVrH!z}J&zB)B9A)(81PyHw@{j z5nH!6_At#xEZ0-NP?zTuiw+HPd=Dj{sROGGg;5zuz3q3RYCtG@CDs{qDmor%?ma`v z6XHeqn3qIJL*dA~J-9cz&jF>*GrqAY2tWl7?{$m^Uqe?GHH_YnZlf1m82{lgMM|x* zab9G1lzGJDy`qT`S|6~yJJ;xj<~sQ%wh4V{)>AOZF5Ll3%L$d*;a+ebYR^<4@?Y`e zNs^D^`@wv?Q>o`BiQtzV4oq$72b&L80W8i$yt*LziHY1#$V_N%gdGW={NOysI@Ax~ zs(z5ujD%;nE}C2>J};Us)vhJ-FIf}*@6+A=U>tO-{}C$FaU)-)_hS^>tJC;rP`+c11k42a~66m!8;oJo@(YX zkdpXU-GuYOetqP{Ng{p`yTUNcmkJR~O)n3JrUJ99xcAxo5U^smsJcuW1V!gFuVz#^0Fdri zD}Azt{@=f7PdORD$11y6uQd%Qkd@y|{4EQsY-1b`+a#dj=mkz=Wg*Df@!hRk%m$nD z@?n{oME#{b>`JZa8j9<!N)|98wELv@7 zcTMM!H^<6_dZlv*jxE-1rcj~+m^7UArzlGdTHFemqBD(=fW*^_vr$Iq1 zxGkA^mNf`Vp(nO3y1B*yu<&xJL8f?}r26q~)SeQM(O?wtXBMFAc%Pv+9vjpEqFMRmgLI zdj!6{QH%3kT!uc0Da$AN;}Fo5Wwk)}m*Cyaj4Wgkd%O0sUWif{0FtUQDO>6X+}eu# zv+%Y7jwq^YHubiGjMs(N$>bvFe@FAy%c>OG3fOpUi8|fmc5vJ&Nrlh{JQ62nLO`%C z>Xor#5S&-8rEH^d1TGu)_)nwOa8WFH;dYV%c+>109pcr5_B$-UFYug!olYk<(^Uyz zwJka{azzN(Exy!nH?zUfM_i+7b&N3B-~ROO?KLEM|NfR#&oXNIFL_ExV)TDVd(UvJ z|L}eMCW(w9p-Ga+NW(0dk&)cllo?TG*%?trWQJ^_LLU(!^L5?I%q&!H`?fMOLiX?V z{U68o$^YT+=ZV8{A00jLcD~Q+9M^^9SHi#j+C@PQ0>=)@1vMj*Yl{CZvv#y%Tu?Tv zSc)VEaV3<#8clL~z>HxQQidskHS26A$9j@MVdsP+_gw;M`$FFlaUg^9SM?xS%va*`NA~7^Bf~kfyKc>v z1k!15x-ZLJWZ3yZ=KigVpRaoff66+@aJzJC(Orr_()E69oz_SOFWRKvhY1Ai-Q5Wj zswabsh1<*5>^smc1L=Em#!LzUtK73re@=1#-!2U_HQiv$!4s?0iBa%)Z1o?_A{kiz6PE0LHvqoI=@w$c^>8o9 zG}Z1&3&@|m&ZB8v2;U6T=^}+Oe=!sDp=M14ey4wbGJO)^$5dL4>|1}BJI7CBK;;K6 z;%j?XOKjlND(wre8VfksW*zH{@1$rr1s7K7jH*LMKiDZKxiDacad8+e8Xr9q zW7CYzF!OQtUTsH3$+_=3g-g-p_3tY`m#R?f%B|*$*RoJCN7ZA_>^zjDJivSFa18n} zGw~wFD-AuNd!sn!7=Wg=YM2HCq7jqqO;d@7E-2A5Qoyq-5J~R$)XHhHK_k@%<#jWC zP=JT#?H_JtNTO6){XE8(u|I;2Uug}{Eit^GwaXcKzX}XJr>>1MifPW>sdhwVRW9|n zS&4{QpfhXpE)^+C>7)|hz5%#pxpdUG5bLw|lQiWEd%^XU#6jDa1k%#So{sBUxx&)*Zz z-KX9LV0})&S9DvOKuYO<{b8*i>vz9Tn^~|BNR@Uic{d9BAmRaIKpABl{3}jzoC+Jj z{HE0N;M@kxW#$$!BTY|YG*aGtxSLqyu^h8*(k2)S) zmI#$UUJr`=^@qsU)h0R$e(+4+M&FIm7Eaj5R)&sQ0GkrEl|%v&Mu-DUJI7UEN-81q z^>HZ>J&5F1)`j5trK~hhGra%GQuVuH3oD3H6x{;qvHr3YZFJUi4V}67YC^YX8A+@S zs!a-0(2eLdPo3QqBssR{N_j{#ve@3g9Bt8#Zu*m@>aLa|%|#pP{=LnmC$rflm(cf`Sz|tF={vDyWsP{hr<9RgjS8^V0`I%^=?|$ zKpW*5T-l;;c0}uUE**knL_D|qZLx{&Ku*c-gu-3FAmf<2S)T^h_q}XrlFi8=c-S%3 z7UPFzMtJ%6A~N{)OGY|izCwLJgQ$e@LQqdbalbi%#3lbavj_7VSy`v!&;Q79=)1O<(kj!iF z@iWgBB-p=sc_0Jx<8)t%#=2>pel!f{b|yD4DLM0dxWZwjoW zPbCZf8U>n*FGD;rPqH=fziXI20J>S?FHd#ULm;C>zw$x8Rc~+3xSFB4hNUC+tFFjn*Y$Y+ z-3CZhk0(+ZQSg<|lHAXzNYflI67ziqpkA@` za()?}f3Av-OE>kx^Ps75F09AAQQ`M~V~6J%(4l`HMIaTLMGFb)_JZfkusj{!2a{U* zB*agr7rajvD0pH0H-0rP+`_6C_oS%{L$DsxYR#s<64eXVnL0cVj0hx8iErtCOZeyG z73M`p2&BIpMDdbueef~mzlKlv`RR4*to0`i0AD%3-3t3AjL+$<(3uRwk%rIi4pJ*{ ze|KfpT<9oV{Tr1>cXI~jec!qLz-jNc{P!^#5BuO|U*fAI^+EVzZy3PZ(E!V$@s_(L zTA?Pc- z0bAgAJ2|u>#W3}XgqWg6e z`!-e(6UWP4ni#-@LHqE?qcy~%QSr^KcNx9WIFna!mV#~zS@CqxQgA+I(fT~z*Gdx{ zXwz_~9p$~U_x z0~vpQKMp|eE_k=Phesm{y@$)lQx|k`yGbW;C=luYI3+aPZ-cIg*SL}C2fBL#;oaa zv)gck5~m;#_Y=frLgezU5=hJ$-|b$hl7UUEs<{a3wW*A<`i^hO5H^1*Ar<3|!|B@z zgNO`#8t3-wVt<>1RxY6WGuB_ScOMPCf#)MeK_2f)GL))BUL6r7ke)XxhwuDH#_uCC zz4iDGz`=T0WW-RHv5ckutOeTO#f7XI9Az0Cd#^%h(yT%9bv*9-O(#;X${Yv4U# zRj%qe0If6ME&=^t=#&$(zl{5F^xOMn^r)sFgSw%kym1twCK|2gFpiN}y5sv_41nNt zYx6;zuE~gD_CN310-q1IihsUY2#?ewKHHoshVPwY7Al(wa8FD>=!HunSZReCDJ1y= zi{s}P2CKfXPXFFKC&vcfj?rA1tF(ZySc`&rQzAJ1Z2ZGAuMB4Yvc5<5OM=bKnP-Ky zLZJ7~y(AR(=_N)IYL6BHs#Bx_F&*k)dLjd_vpXm zgMWsR$*A(LbllG#-@JUzNxmIzXldv0o+(AFU61vz&sU-2&IYY=S25px-_uU>IuG?* zHr;IIh(Xi>W5&om4NbH9)4JOSpet!_pBno|qc_U@2MV2BP)_`RCDq>p(JOPaFSYeH zh)a&}`(%a>@*?c0P-A>~(~~!oz2b`KIUK4bX$%mJyV;$LPG_Vncuwc6sy3RDJXxSt z>4@@$^kP!?6VZ-#VXNXhDw5JP^&jz>0XWO+d20ahJ#y}PpNqqIKga#lvL9)fulBkP z(0TU4`@FjZUpz0VyzUx{G{Je7l}(o~&k3aLg^8Tx+gN|8Jajb|>%a1EHU1KudO`7T z`JPVfUw@D*{BkwD7r&1L*~YP-q!+yvv75dRCbcyzsrd+`vtPR5@ZUc8S+Bd?3({9xD9VnW>^TM+e&-_J3P^*SF~wX$>~kfu^v zJAx_^M7vL9Ms?BmWC;#`jXzdpE7y5nYSuWf#>_MnCN zavt`rGBUyQw3x|RA(V$sPkU)33CEzA+)atVs5JDj;yg2VKmei+r;m|}h(_$R;YoZ! zF6jMeCA0f*Aj%RhTD&=kpO;zj+YLoNsM?!n>-kgcPv_y~m%H4M@z+bd`wknR`D=2Q z`Z2!HbMBl_Fw#cmU0i-ATOAS2vf*_erLrIM6N~-7<`RSmq@1~2?hmPC_$49C=i9Rb-(&}?lLE7QM;srUEL~qig=zF1u1xhkYAqSschi1j)Tn^mvOJcL&@P6n{YAohk4Syi_*#Krk%Le?$- z`+|o^iI1a^v)X}Kdq)>^_Cxu`;I}}eySHoeW*vSX`MX|XNcTZ*O8PRKF4&)*wyLHq zyQ1VXk)wa84bZQbX;QQpUxemqr9`f2BbGgtZXl3~*NqQ~tYO_tODi-y6Z<`m_)bS)FQg}Wt!m=^-n|nX>3{F`!V>~# zjxW}K3xqV9K4JZK3;RmHu>L!(v&%T$3+L$solTDUVtx3qrfA&zUN{m_^)Y)ttj|y0ELGTns6Qgymjnl4>jERO%w`jM9l|`tGKb*~ zMb7k&-wJH^Kc7;n8wJJ#{_Dq6X5i7BHABS!?&Yee{VK&dh&Qf_3NaRgaB$$vq4SFk z0K5}^GRrM+KJvQnd+H(xX*n%bugJopcJ6);^vOs!mx!r2#Pj1 zFX~;r>+^S3xH3?8{!#}63|#6?RP$d$`}fbhDH~iy0ZYi?@&yX=u(1>FiPp@_sFZ-#B{ z_2@JtTT&~e|0Dnzj>qe5V0=-d&Lr6eyP#9e?;CX}foSgOmGasl%y;u@=!ZV~pwE(D z_qYd|p_Zgv*+3dM^tJTxmi`d~bd;ak+I_$o@oEdmC>m=cshIt8qHT_7uZ-NvVe)Y1McBS+{@^-F0rNnPIXG#CK9pnwMGtm9fcJM*||nUfCE=i#O zt|IJOEd*v;FPaQ5^1$IAj1?Z)tT1uiGW~ue)?f4>O;&vkRoJW>^OY>435WTvs(lok z^X*bQJ~ND(?3{z%n>3@WX>GOLGVSOp?baJtzEU*tUZm{vOcipR)^68Oz<#5`>hXn) zJjCT+{8;*63=+1Yoa}N*LpsN9H__b>KvC0S8$bP`k&Nbc{TT-rRNB-aPmS?q??)B! z;V*ce{yY`2nCgR^LkiA|JTybBTgM}gVSEv~@_`|n$^hk+^)cIhcSbBFeaD1V@bgkF zFMO`t5xq{4FzII|q6Oc`)zBF#(rPp%QKe}RsK&OFue1?JtlbNFr&;=-Q*75QrMEag zNzInk9^DHuLE&>;iJ0#S8XFio_rkNi7kZ`ezVAf$KR=x(aK)<%lK=x)gJTdt9tf0XaawblXWJ}!(<&A{5~ex566x?ODM6x z`ON?NT3?zCg6)EjVq)|r{K!95sa`t_^z%D6O*2>E+Sg?!C(0-|@pzs%Q!@i{e0AM| zb7LS7%ef(6-v`$}ryXE+7zEk)7o+mK8sT-?z6hhu7Fdj9a(T#D1Ru3UZ*@H=2F=Kx zzOt{0;5_nu;SK(t#?ba&fH446&peaW*7F0O4}^CIPTPWfbb?jZJkCpf+IvFnGk>bL+jw0~t@H*k-y|LGk=Ty+qEqDv;$pQXLtMe78%- z^FmxudSQCN`H?^*)%Vp;oPzbbdm5$xihWSsUuRbHATxCE*}Q|=ZZ{;6wbq?+)BvgO zsf}V8bVgSqv<;MQX`=&LR8%tUj%Z(?hrjFzBI3TjNADlsHYnIAD@3XOfQ@9?;UoC{ zu(+S)Z7~%YS_8B6BCy^YM${I)>Q9Cuf!U%nn6DJxdSO?XiFty^klRT;oM+KFaKbhl z>n}rfY+HD~nj`F)%fRzb#VcER8eY8bPTEO@1^;_zolcW{*nvaOg9amv$e^oE3VvF! z4X-UcsDzZrz~J4RQkS_2qTX{u8DGh8=G`ZICFV6qHEBFso7WG6Ut-PLGL~TF(E@d) z83hK~Pd8qSnFOM7p7WigQMgELvZY}|hFZZ6-h0jiQ1JPM?RZ!{EbzwZ^_sT8Q>t)X zC$U2KSK?xL_7LVTj~y!WdlNwO_SoG`Yn-pMnU+-!@rSyv0{0Dn`+{Oag=JB!4McqP zN68;?9`j0{et`}V_!iG-6^ttbOBrKsScfE#qnEdTlncQaU#51=Ieag$H7Zy04J%OA zBp5`VGeFXn?SZgsYv^;(u=DPZ%gE!%+lMPm6r}skCNpzt80~xrx_uGvm!Yy8{LwDm zj;4mb@uu;XVm&_4W^1Yn#pe(^zRTgf6jPA$-Sj+Ev+o_p;(-`^Z*q8#>BBS>%+vht zn{5Doe{74hV|<~%IhiqG=YnKLx`TN;0udNjIy!#Fe!wm>N>7Rpa%kDZuk^qSH8KA( zG5G6>fcWja3BdsQ^&B4k_RSfsdKOU&DQhFs9W6WFGDnoI!>_c(L`3&ueV^wu;r{jc z+|`LALvXLIrqZFCKyoYyRyrZj2UCYjg*S2uBuxXC> zhv#o0Wuv+kh$%HWm%K9yo#&jdiH*;|Zj+N4skoP$Iy8Q)b+`}yF{_V0^Be>l&CLQL zTO&N?PX zqEh#A#5JR^Tz&CW`*y^6xd_&AAElI78ujTwHCoLLEN`^HdCcy@vZl{@NKrN^_~E%2 zL~JKKAB{~zKi~bftqTl5w26g-zERPraJ`vv{+SEnOo+a`e>4yUUeH(08Nq&|z)uU7 zQXkax;m|SlU^8UuE|Ip6)(yRl{FOg>%mBHL2Uh-v@x`#o=i$}c+Gv((u|~DS5wW!f zCv))-(ODh=p~BHEn2h}~v!}NkVk*6BrZL}$6@SEBIoSgx8S_tnYGHm`P182)g6E&h z_Tm;gF zStq&X0U4?+_#V|4?ZD;pM)P!6$)Guu{B$&F8}5Bta<4x@h6&{sFlQN^$uA zo4NYPa_k@E47|vDY10qVrQLR&K1+}n88R5DNC7X>*;fMglMr@NGkX#1cchX?a*YNV z4iBExnX?*zpuK`=KfLOJ@jtt3$e;!4*IP{A2^K=Dh@JS+gP6aRD~o*VNPu}@`{-hk z2#=4LKUWFx2ffLaBEu11*i*>1k0HVaf@%n{-5)H#{e;ojoE8zv-l#i8j40#108i&H zwn)N%#H8w6_HJmm2@&s~kUV*;M2A%q&V;~`K>nO)C z3t!{>@9ks4yZln5fnU5raQalY;?h3zzF0PxSw72GuQgOxrFZHtSxK+H<2IJ;s2h4aj8 z75PN)IC}Vc?Ku?)si^&K+U7Lq8KfIo=JRgHDpQFw(j88+EZ@fY_g96a)#VfDn zqNCB+%Fy`>VJ;}j_sPmE#+Q!J5yrz~I4||veBYT*KB(2M=CoLd8OpKpZSBJN!sh)c z_&k>ZdR3<~es#zhxzOZ9D4S@bhGt5D-Zw{d=-|te9gHtmR@JEQWo|)ajrB`EyKXqs zTmGu@1c4-WYs{_YTMwjG;wI*G0!i>p>3v_!3tDb8)P`XHIi#C$eFD!rjECRte~h!us%qjTeqb@beHX?(pCa_LH1FUk<$40n=Z?R8P*5 zLB-YGUnyuC&N4dAneFtz=hLYy4Hlc=v+G6Rgee)c&vu0bRQ-dGCZ+=SdHW$Qhu!m+ z(Gpl>aaG;nqX2uMgxB$FlTbhXr|ZNRt8&vswF1$?CW^+4o05!4LrqU?v1p-*Bf~gG#Lue{+7?|A#p8f~%a%SmedJ|CYKJrWyvw9TKuH_@_r!5(Z>b}StEJyB!$?HN zaI19r1~rMA_Rw{w&>^@Sq_WXALLdp>`S_5{un!K89r|(P@?|V?M(hP(%4X zj`bRM<5mHG+=si_E1l)i2cD*<@5kVL74e9@BRkIDJWI5;{rGtY9?d_P)*}qTe*&F{ zFEei8dnJGUBF<3YAE%56J`4z_B1JtGq{bjw%RF9PbQU682QT^`9fzl$Pb^0``$07) zE6_h`5a0J~RdyC{1Rcg_WrZxQAoA|IZk<38h)t^2G5KNsLeY9koKA%2l}h_%e&KxG zl3nAe;{g!YcEaYSl^>j6TFj=uWDACQN5+=_;r9_!bdXmO5yoxSgvKtaK&F>w{w^6Q zc=TiIsVbW=OmfQodi{V0B2(XAz1PbMS|Sghoa@K>OG;0|q33Jp&&wN{mwzuKY5Ctv zy3!Q%p>2w-6wlS9Jr8fuCpDw?3b|u+586?C5Z&k8o26*jolaDNqZ(zO4gFMQh5L=N zv9|qnd5DUqz>wx*3_A51hn4B)-`RFgD zyzUXkn+-*oQyc`+#rKI+C-v|=w)yG zt>y-WO*pAC_A~Sh8Gbuk#J_s+575Aww^wlAaXv{uuwP;c4t*+3kfNu+$5*o22To1` zb#XN7KD$xq!FNIbJ4lActIb6#ngeh?$#BmngL;tvbVTT>S_@p8)o?k^SqN_2#`DXJ z#n36VdVsw$0nP_LRJAh1`8t1%Q)+Jh;K>)QVASIa5l8)T1Mf6}q=dG+hH2fV^Ch z?j*N{dWYgaX67v;krNF!vak;z)N?BL?D#Os@TsRX>NKOnP5$0zFSg@*1X`~+xJps1 zd8p_6u`0BGNUihEW&FI@pRaC8%0q7|L;E=Q$DpVEeIIum(h%?YgExe%1JK1Fy~z&m zXykDFoyc)p7u0!{U=-FGi28Kdj~+(2Pf`zkk%>OYB4IyWkfRxzRowNB4&%$i=NEIH zZ0e(pOWl^cFuvq8UMS>O#CfSYUxDpnM-=6~zI=;;h`PGe)ekvSlM4QHxEI$9!E5Wi zM|P(0{Cq`~f!Vna&V6mm_L6iDIQJfSYK z3J#AFlaK0;!H%h6wVv)Q=xrwv-t2BaCP+@%= zXDiHoU{VjePy}X9H-B?K!TcrHgXQc>B1{_`tgG!ygeP*Vk9c_lVD|?d?wj}gz>z+2 z^MkxCSd^sftJ<)DOJ@$ONtY5qr{otSt)vQ6N&m1ml9z&a>{=fuIE3NOpL3crt~@aC zMnEsNpA~ch(l{K4u>Mk#WY8M1h9*31ByUeGqnz0Dx<@Zl5VF;)K8JmQurId*15%oi zlVVAfj7vK@lXpSUU#AqUgckbkIaH1O9y~77e7XK-uU~n%3kp!bdoOo95Z#s9Ox&8leKl)$8{={x)IGXi zhCb8`Nh>f4nbNtTc^y}&+7kw7?Md?`uVH7D>$O>X%}g74snpUqbvh#6Hzkw6i{D4D zBJz8WZ9(th;2DR=F5Kg~XkmN+KOfoUDs8WNz;WT{zK>W>`AU5#`_m0FjJv-2F^>7E zsb)~{aUU{J)4IG9S0|AAYy#5%1>ioC{}BdxoTt00ETnzfnGANd&&SEESTF9}jO|px zeaUdH4(8`O@P@aW$DIcECDTP7(cIpK$`Y&IxQZU|?lcK%C}l%lPNw`7F}C zf6!#I!o!-<2O8lC(G<2N==gfwm3iiW?u)7j(QW?*8Se|%Qg4ldRR`|_?)4tvUOP_l zQXBw!%KW8lm3p`(vvlmPVhePR)*|Pv-P4N3CE)I zuA2NO1ZyqiN2aIoe&D=}yKhri;ahHKg@hm0U;5Vut7O*D$hJ?D|GQ;GYW1G$!#RfI zS1Jg_qr<4XrnmgY&1U2!$3(w=t{wTRpW>D~R*L*IhXsC*R3XXKYh6>)nD0u~jLIhF zA!W744?CD+P|~B`m!bA)$Ux$W=<=Na)GyS;yAR`ww0Rs|?_C#Uq{^J}z9kSvZ$Hz` zslt8OJz}En2|mcA&t@yx!3@3Dt$L-n;EHH7blasj^ilh*15&o_&L}*bKK`_VHu|?G z!e-&4BMQAZ-Rr%Vh^|jb4_%I-Cf(`@$^7?k2$)n_rOz%BNK>4HnR}k~fog%j%a1Yw ziGM65g4EuN^U*Uk7qhW{?epJ?7S1zeU9KPxrxQp+5pTlGQgQxec$eNH)@P>d)W7If z;5?msm31$kpQgn&DxG)rfrKqxAxi{-WI}uKP(gYh_Q zRW`$SNd?kH(7GP&W?51iqKy|FeH9vj^LgV7fiI$w|1{T| z%?KB?e$J3Lcp?xLMTM9-PU61(De8kR6+WmZqfBK9Oc}{-;u8m`Q#x!{@gHOf2pSkA zjTfcsTyHdB0;Pp{hc%pYBZlDKT^B%D5NDRP`K#0>z@p^Fk@^zouDFkj-wRi0_n7^cFuNl2b09Uo|gSuLYkUB1bH17FBlBM8xqjq0lGw`mme`o_% zshg8kZ*adcpLk?e9^b)!*fziSr!s`ES}~I=Bwm_FPYH3g zXWP+9zV$SlBY3~&B}p<5r3#rDDR(hJSHwV4=f9|KbNK?J)1M}F6%HdTwzu+-Ff3gqfui19WcH{oQZ_->r zNp3IXD0efZV*h%oi9XN!E!JN=9SlxjyoozPedcd1?uX<*tUi@OAbs1Hz+Z#?C>F{< zW+2vsrOeFe#7h20-x+WW@F0)~$2J&L_6)!V15c^}&ht`-t^P6B8iF}<{>6cyEhrKB zqVz3_0u%CH^ztsNK&5k^@&NaZ>U~ui?9*pqz-9f#54CX^HkoUFdcPm8M&zviEgA%g zi9NX)wvDjci+Y6abSv1QaB&WmBKYKMDWVcq4A;dQ1*aI2AooS7NY_{*{3SDtoDdCw zkp~q~L>E7p6dQZ9qiG8d8q)4uqq2lLOOAocY9fegzHCvsq5|_3V?Xz(NdeuEIT#M0O?~#e_f1SL)|h-jOlaBDDYAPN3JXdk<*2l z1=uKv>(YgvF0YzVu}+$Mjz>EhT6uO=)SwidDBQ&rbF>pC`50dY=CwA|EVPlxmp4Jr zx*d^-Sl6+~d^pdb{@v($;wJFBWXqcrbwZF~KFNR4I>sE2$Ak5)C_Sy7$QNZ$UCJ)wKsM_;GJEQ8wTp zYhRm7Xb-%RVuEeXe=ukgezg2RA5c9w{BEXg5o8u0iIf!%L-0fWfrXF1p@`A_W4X*I zcuRx`?fudNt)}&IPlN{GlQY+w;8XPwt(r+CEY<>hdl+A|tbc&zQ2F300_HD0es(@d z3Bd85O0Z2m5$+_u->Z4Y9|Qyb9Wwsv3qR$5+^x5>f!ncSJi-Aq?<4J}J+ zEX%%KMx%Vc1+{RGgqm)oplWy+$*a-Y+`it7#Lx2bL<+Z~Mne0s=tHGQhd(m(_D~fH z;ZXXqB7yzJwt)*m@p&lZ?)_J?j4`P8(=XGEdugaf-gGI*A^?ecOy{r#BJSRAmNVp%%D6Jj*yB8#=SU>#8BanP)&f6aTi1T(@^Ku4nus=!T(2{`X zp^g1LMp_teWWSb0#bCa3ar_LMQwGk{OkRxJk9&vHMkTRpSg&EPc)nWkwGS`Nb#nfW z{cGt)zqmYs0bqQ!AZu^}=k09Ul=t5s0z)orA{J~xfy3>fi6RPw-H_V4^l}x{<+^sK z<>Nik5%T9fYG)xU&%orG(KzV-S2edUxF2T8I*m>U_h-|cTt0X;!qn24g;OG}AXs-@ zs9&oHMt0i~jUtPob@NYI1A7weITt))JcaXh{PFSJ;sHR^nY~8z^uu}6u=TgPw&2;p z70FCv389Y@wY+PHke1|@XCtoy43F%iOf;oH@~85Z&pg7={b0`rS|1*0MXAB8qpV== z%iVQ)iU9_<`d|KzUqh{2A6IZl6D>uscsj~c(1dhi>1&+({=IUHFqhSg^!iz+`@P!H zrK2VFdyGnvcx|48{_$!=d+Ei;*Y~qf51U~#O-ml~p7^zON-74mAEv*Um70dq!gNmT z4-Y{53aK8^$3>&b5!Ho&XcsiIeK|p4G7zcBzM8!`gWpG491d@)e9$l_-R*|wW{90> zxIBNa8@ha$wKnaP0rEM)Qgv;_8T}J{#I9k9_j&lotHl3sL>GKUB+O3Zyf^n>QCj^? zXv>h8=~nCn>BS$r<2QETHs`;`!-hS;uK-CA(gaeuityubtcNfc*epC$A&@ekR2k2k zkzsLyhh?`S0pB^-z4H1l)?X@mcNL$-^N}B4`Du(bOCLwM{k@ZK@~G#>FAX1h2iRYXRCe!86AFS`eL4;*&db7BC( z%S(IQ+3MlHcF(1KA}yf)T_crt;R7`7lDd<-`4Ku<4F=OA6X4Unb8gBiIA5o6QvZgD zKU8PbnH$&nLXx290WNDB&<$IU70s}KUYG0QcG5&R*}LytQ-?Aby6dG16-h#5M@hw# z3?T?n-F-No3*Xy$AtcTp&kD8g>9&1c8K5V&G~<=n8rm|rz`*r-8M$?)etkwoLCbdT zd*p|Pk$c1u*-^D-6#KRy-%zj}+4p9plyj6K*@NTK2?JHAuuJ}@w>ZvACE6Qj#^#|j zD`t|bdt(q;_ot$uZ5mouZtc2Z9)NndPyB51j7Dj8dhslGT+p95$p>s0U-JKrQfZfA zzS|soMmWv~-OO?B@4s(`8bT)y{GD?}734A{vsHbh4y|QLtvJun7plZ7i}wY!Ri0}8 z;D{dGst)hlLqwm>(4}WDQp6^S+{!}f1SA4HlYO|E~zd?O>I zWa=;8w=u_?on1g6v0QWI@vG{EA|nQu133iJy6pFv;<8?-%u&4Uhw*00!*NxVjQ4dM zqui^&dQHNYZ&YR?eNfjK7G(bl_n8iqkIP~GjqX16L!NK~$@L)(7pukq92c*l=$|8y zZmGOUMS(+5*T?WPY{+??SJ#)6{`>{BqPw&JO&Y|Qie|^XTfH)l)u7u z92!?cS&}pRA>lMt;tbC7gz5LllENCH-Clb~L97*~o0)jBjEaCc+{MNI1?DebIF#R< zNCN9N{e+RZM6h<{R^z%H03G~;&Yp9m0 zWoT$|8JSaz$zqE5dFiGx<>sKE;vzbMyKkG(KdWGq13v92UVcbr(6|(-{?tWBxvSBD zPh8P)dz{Bq` zIdJ#Qx$Aed(a@D(x~!j$NV0B`D94Zcu*(Pc0NW6H@BR)p=ZN)UvpcXYV=*}) zkN@6pGdBIi^A3ISz_ByO$#CSe?s}#Y&O6C@W`^nGTLJxh22|w;qzdnaV<9HEU&P}~ z6^Qe?NnymY_sV2AHc=2(i05O{{b!UoW-@ft-^fV<%U&QGT~)6!+t}#Udrc{kVSZR=-q_5JI;UW z+>PqB{S5-C6~})ZABDKI>G7qo9`LHSPY*jd0G3pm%r?8~p`H=cfr^_`s|NDI~(;c0!hxfj?+MOs+A>7j7R6Mee3Ha6Q2D-n~@Z zZ!~EB@}D>nN+LvwIjzcY+;3rpRnOYf#30gX$goMr*& zNy7_PI*c#(o#dU`EL@OVfpVF6Qy{wY%HYy)Dc0*4_xiHM`XJ4jrteGlus@xWF)D)b zMPsh#KF5ka;u>OzAYgns`N~Y-^Uau@{z<6=tCx<3g9~|qINb@MdeWb9m-jTXq_;iOa zyA97vyB>)%eXHq(M{@b5`515FPVQI_bD)Rq%=bD z#3LrN%dPP4dnBWIF?dlA!KB|MmMziID#J-PSRM01)NSF4ubE z2N@nevz9GvA(V1rG>P644pgLyZ8Z|%(FHNPcPc6%yWvg6s4oTa`qN3<0>aSDHJvOG zzyo&eaXHq%S>Ya^`)1}m1DJ}KacHHip{kM5uT;y+s9o&srA;LY+U4xD=YM+lz79*o zp`2!vJh$=fnqNCwux4l~y1APp1MZXfUGnB^&qGpQXWL6J#h|PK z2D|W#G$iKlVd5JZfT}Hz(uBuHBdxxNGW4-7XrP&!aU0{y&}_Mv_#Ye8$t~RQ1NorI zDE$Y?5oRcN(t%To(G77dwY!S)8=&sNSQejAXC$22uAynAjVd?Yn+kq8A{i!XJ8l6Y z%K9qN@BV!Qj;!m;*I9pumrl(=-2*$2ZrMVd6X=0|?uM(}7YHPM?x@Oo7BcX6EA0A? zf1i&w#U}YQvA@$f^DYAKUnrCh-@30$2LDjn9bExDAN3^>_gyB#3)2PLw)Gw8kDhT= zCg8j;KKT)>z61J}t8P75Z#f#`aes7q3leX)2_`Z3KoddnOiKXXZ`l96&SdHrSWuNs z7^khm_Q0Q*3wWQ%0oDGt4XZ^k6PB+PHyXyh5J%e%rQcv_bxY^YzEL=F^H{BoQxAOe ztZCb~cL1bS#}o_}>tMK-|7SN>3$)a>6ngc4fCtN`k3aeQ5xN5D>P9^iV7)7nQQ}G> zlovA36L0#%V87{taiuSC?~K@X7}&s7&wsM7lPrLENtHN#o(L)mMCahI%D}s-EaaCf z3Fov}OjHtuASUI6h}B^p*a^6Juq&DsDt8xk?{Z`S!#fvGNQHv_0bYR&C^00ks~qIX&)ySA45dlJvp7F@=@+L(|9^Ns`q5Ww z|ECwuKG)*Bfj>8OXXxEl!+vy~5&i-H?tM;Ncu%x$l|Xjlg*jN% zepsCvG7g1Nr}Yba`eBk}C_Oqf2>lK3ehj{Eghcb#M|>1p!KNpi>8wK$SS1Hmyhz6U zMeDAanM4wtq+g?2TE+RgePV@2uLVHrzd&cL5I^XP@cEZwYYQ>cnat*lmeBKqE2p}J z2$Ozn8qI1dU{py{VPPx<4KHPyQiO#;Wu=d$<0%hR{bVVRnqmceugs{SC9J zL|Z$T!l?&RgWW5&}9iCoFV@I_DmZdod5p6d#35}o@pJ^2Ugbq zyJ~v>F2eu!^}pZ$pI1!_2yy#x|6hOYeJj7*bmYP?H2bqFR?RO1hXT`nr;Y_6weHYP zhtI*_ZiPsH&t)9ti~bsEI|8uc0hw zY~c_2!8x;eS)DP~V6!8gvOTfP%PfL_WJ_ggsmJgAk$rgKZqpGr1k>T%kzbxp$(Cp5 z4b2xe!*aHUz^;>BvL^F1{&dU~h`He$%1)mm`!Y$FpAy~$gHG?Kwsqd%LshSY9@uw) z;mQ&1FSM^=LBFc(t8OuH)fFZmXMF?t_q(~*PZdF}Sx?*ar6@Rg^xg=2(hE2x#+|9* z9tAb8pOZSLVj%4lu|u)k735QgSjW{|AXbNOu!+|V7P>?)KfU0BDLrXjZ{IEWt4F=} z=cq0`eY5hwBM@At^2NoX?PAjrnGM!4C_Ipf)q zLg5;-NA2p#HA25RZSh0n3We2s&W>MPk|$ihUt`?3@>Jo^PZT4K;8w&*(~5;tMxGQ# zwH@7>IVM@Srt6kx>WEmO=qHt}nzqTpxa7-D$2qY=&i0YL@@AWashLw>9o+0MY@4(# z@eOB_FeSwMsbz(~@S6OiG`~SU@F&w_M)z#87ka+;Db1bkBRm+OIbv6tz0_99&^o(` z*-C?8E4g?Cf2$j*b1wBdBm7S~+brfvhH(vX^nV)7o?S||G%3J$YFCrx9x*qn97Lq( zpSvRSCzg{HH+tSiPpcyj-3w-{b1H(3BIxQitRt?`MdLdSljznZ76+bJRulI|*TGNr z=g_`R<_{AU3rUiJdPDCN5$&e3r6sGlm~2Wrq-TGvmX;sXrJzw(N;j%nWti~p(YKe| zOc!|kL7(TxA6p6EX*b;MyeMMX4Z1w<@y9X6+vtzYhZ9KdALM7=#U+VlS)}A7=k>lJ zRU|@J@J^psM)IPJ((@;_kjr{Js<5h@Xl0u}Kcm!2mg;yJ?7VrKm~xzK4BeiSN?V;# z(?f-1LG|SHjLXl-b@!6{0YO>B?E=qJ;3_5;=PIvyc|MgSsM!l?j+j`wugj`RN+z)? z_0~%&#UwbzW`X00UF6D5`+t=Jcjt|(@|E9C8r)K>T)V{Nj`DD>nr0Z;C%Sr{0C(3} zr99tj71@41{29*arM+ zU&Qaz4hL~%wr9Sk(l~c} zKLOusz9aGQo_flLv!EM!{yWXLG1(8O(7J*ZeZcRw2E}Ev3n|+@eLl$NxOe~1tyVy_ z^1DRHS<5^wwp=J8FHdeQf}w&c451|$oojx@f}oSM|?kD*Q_3W)cIG+ z*+D_kHP{j>$lIoPDdoQkIhKSCw9+z94ENkAc1ii zu`ZQmq@zdU-k66CBOozqFBq+KWcld<0>%V&a>T&10-~Q; z*r|}(7~RJofZOTZ+Lkx!C6zS7ybQTX)#$Nrc8aO;_3`3f;A8dSOE$ViRHTJ@JM#ED zPWw_1RZ_eF^;}%&s~xH%9QR$JhWK8{Mjjo}pmN_LlRCV9S10oMly#gIl|XG-ZPm|D z+rLh@(h^7M`9Ju8?M}5%w2Bw4rn1d$_H#~Q!dg|eNa`1&2|Ghf?Bi^kO*Sr}e2Ucj z`3;5Y<;2TXY9r+>jf-P8(qQ;R3JDA1Gdn4EV7oD;?6aM(wSh6)z&9u)EQBlY2nzAx zg&53a+j9A#Jg5)#;|BA98~$(SIy?}n%-m_#ynD?&DG%0$o_hPW2i$Mi1C4f*jfwHA zew2JO)7Beyo$){oau6I66u}d41>QU+6N*ZoQ^Az7GvR7sJ>&n93zYyc#~%;|2m!DG z764O#)j;V3B@Yk*um+d`rT|s|rU9^R2EZ5a!vMF0659|s7LI}Aq83YR$7?Lj2f&8X z8xRBt8~|e*8+bTC09ZLtV!z0N>wiGW222N_jdKB9fFU3lU<^R55rFCA0@pYn8(hwrqeby#v(S09*l! z0Biu01B@}+g0VomFmIs%W^Sy**kPa50GyYJJJSc-upZ-%wy*&>F6x**7%TJ_`{MYR z2Q0t#pXrOSb_C1@xC{V&$9k;8HZ}nBhU4Ja7&FFB^a1NJ7dQtSfO*4uKLGY+1F((J zVSf`q$Uq$*N-lsID*~>u4{De?CfC@8eNl^j@EUExery1)D~tzTGd3_|VI7Wx8pa;1 z`??+dwgb!qECoyjVE#~p<6()mGj^cusQEQ)yY)wYsQV&;_9^ArnCjKU9_` zFYBZ-qmi6SqyQS9pfH)Tg89P?V_^!!aEYnzDZ%>i;|c^kzDEex$3I-Aeos}Y2BTMz z(5Haq0xxc;zf9flYmdxgbdu)!4ubXO9r*dTQ0d52PtR4BEi)SYS=N8%mX^&~cC3?; zLBB7r4RReo1H)xc#&W0HWUJD_WKhYRYhE*d3Znv7zPvWb1w(*V%ARUXYw}3hBJ=N1 zBy*CfruBzl|5Hd(Jhy~U=~6$Wtplh)>`wm7P58c9iQI8r?(*_9osk* z-r{A8wnW9xl`U$A4_#N;j^R_}7?I@qcPRi-|1;&JQ;c^6@rEBlG8k9K%Khj5ptKC{ zelRjgr}/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 index 96bde60..287e4e9 100644 --- a/tests/acceptance/test_dwba_reference.py +++ b/tests/acceptance/test_dwba_reference.py @@ -1,37 +1,42 @@ -"""T5 (spec v0.1.5.1): cross-engine DWBA normalization anchor — fixture-gated. - -Compares ``matrix_element(χp, χn, U₁) / (a·k_p·k_n)`` built from lax -wavefunctions against reference-engine node-sum ``T_lj`` golden data per -``(partial wave, energy)``. The fixture is exported from an independent -reference R-matrix engine; this test activates automatically once the file is -dropped into ``tests/acceptance/data/dwba_pn_reference.npz``. - -Fixture schema (``np.savez``): - -================== ========================= ===================================== -key shape meaning -================== ========================= ===================================== -``energies`` ``(N_E,)`` entrance CM energies, MeV -``exit_energies`` ``(N_E,)`` exit CM energies, MeV -``ls`` ``(N_pw,)`` int orbital ℓ per (ℓ, j) partial wave -``channel_radius`` scalar channel radius a, fm -``nbasis`` scalar int Lagrange-Legendre-x mesh size N -``mass_factor_p`` scalar entrance ℏ²/2μ, MeV·fm² -``mass_factor_n`` scalar exit ℏ²/2μ, MeV·fm² -``z1z2_p`` ``(2,)`` int entrance charges (exit is neutral) -``U_p``, ``U_n`` ``(N_pw, N)`` complex distorting potentials sampled on the - lax mesh radii, MeV -``U1`` ``(N_pw, N)`` complex isovector transition operator node - samples, MeV (local) -``T_ref`` ``(N_pw, N_E)`` complex reference node-sum T_lj -``rtol``/``atol`` scalars comparison tolerances -================== ========================= ===================================== - -The C7 convention factor: the test first resolves the (single, constant) -relative normalization between the two engines from the data, asserts it is -constant across every ``(partial wave, energy)``, and then compares. Once the -factor is pinned it must be documented on ``wavefunction_grid`` and frozen -here as an explicit constant. +"""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 @@ -42,57 +47,75 @@ 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="reference-engine DWBA fixture not yet exported (see module docstring)", + reason="jitr DWBA reference fixture not yet generated (see generate_dwba_reference.py)", ) -def test_matrix_element_reproduces_reference_node_sum() -> None: +def _compile_channel( + data: np.lib.npyio.NpzFile, + ecm: float, + mu: float, + z1z2: tuple[int, int] | None, +) -> lax.Solver: import jax.numpy as jnp - data = np.load(FIXTURE) - energies = jnp.asarray(data["energies"]) - exit_energies = jnp.asarray(data["exit_energies"]) - ls = [int(ell) for ell in data["ls"]] - radius = float(data["channel_radius"]) - nbasis = int(data["nbasis"]) - mesh = lax.MeshSpec("legendre", "x", n=nbasis, scale=radius) - - proton = lax.compile( - mesh=mesh, + 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=ell, threshold=0.0, mass_factor=float(data["mass_factor_p"])),) - for ell in ls + (lax.ChannelSpec(l=int(ell), threshold=0.0, mass_factor=mass_factor),) + for ell in data["ls"] ], - solvers=("spectrum", "wavefunction"), - energies=energies, - z1z2=(int(data["z1z2_p"][0]), int(data["z1z2_p"][1])), + solvers=("spectrum", "smatrix", "wavefunction"), + energies=jnp.asarray([ecm]), + z1z2=z1z2, V_is_complex=True, method="eig", ) - neutron = lax.compile( - mesh=mesh, - blocks=[ - (lax.ChannelSpec(l=ell, threshold=0.0, mass_factor=float(data["mass_factor_n"])),) - for ell in ls - ], - solvers=("spectrum", "wavefunction"), - energies=exit_energies, - 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 - assert proton.matrix_element is not None - assert proton.wavefunction_grid is not None - assert neutron.wavefunction_grid is not None - assert proton.boundary is not None - assert neutron.boundary is not None - coupling = np.ones((1, 1)) v_p = proton.interaction_from_array( local=[(jnp.asarray(data["U_p"]), coupling)], block_dependent=True @@ -104,25 +127,52 @@ def test_matrix_element_reproduces_reference_node_sum() -> None: local=[(jnp.asarray(data["U1"]), coupling)], block_dependent=True ) - chi_p = proton.wavefunction_grid(proton.spectrum(v_p)) # type: ignore[misc] - chi_n = neutron.wavefunction_grid(neutron.spectrum(v_n)) # type: ignore[misc] - k_p = proton.boundary.k[:, :, 0] # (N_pw, N_E) - k_n = neutron.boundary.k[:, :, 0] + 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) - element = proton.matrix_element(chi_p, chi_n, u1, conjugate=False) - t_lax = np.asarray(element / (radius * k_p * k_n)) + # ------------------------------------------------------------------ + # 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"]) - # C7: resolve the constant cross-engine convention factor from the data, - # then require it to be the SAME for every (partial wave, energy). - ratio = t_ref / t_lax - factor = ratio.flat[np.argmax(np.abs(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"]), - atol=float(data["atol"]), - err_msg="cross-engine convention factor is not constant (C7)", + err_msg="cross-engine convention factor drifted from k_p*k_n/a — C7", ) np.testing.assert_allclose( t_lax * factor, From 2c4672038ebd3ab586d636bd10f7ab349d90d81d Mon Sep 17 00:00:00 2001 From: beykyle Date: Fri, 12 Jun 2026 11:49:10 -0400 Subject: [PATCH 3/3] update design to include ref case and fix toc rendering --- DESIGN.md | 18 +++++++++++++++++- docs/design.md | 3 +++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/DESIGN.md b/DESIGN.md index 384e620..6a4c7cf 100644 --- a/DESIGN.md +++ b/DESIGN.md @@ -4,7 +4,7 @@ ## 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). +- **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 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). @@ -3098,6 +3098,22 @@ not support JVP"). DWBA distorted waves from complex optical potentials are exa `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.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: "---" ```