Refactor/consolidate free solutions - #93
Merged
Merged
Conversation
…inetic matrix Solver.free_matrix(coupled=False) assembled the full coupled free Hamiltonian over every channel in `l` -- an ((nbasis * n_channels)**2) complex matrix -- and returned its diagonal blocks as views, so the whole matrix stayed alive behind every block. For a partial-wave sum (elastic and quasielastic (p,n) workspaces call it with all lmax+1 partial waves as channels) that is an lmax-fold memory overhead: ~200 MB per elastic workspace and ~300 MB per (p,n) workspace at lmax=40, nbasis~80, and ~30 GB for a global-optical-potential calibration over ~230 datasets. The blocks are independent by construction (block-diagonal kinetic matrix, overlap * E_i/E_0 energy matrix), so build each nbasis x nbasis block on its own. LagrangeLegendreQuadrature.kinetic_matrix assembled the matrix element by element in Python (nbasis**2 calls per partial wave). The radial part does not depend on l and the centrifugal part is diagonal, so compute the radial part once per quadrature and add l(l+1)/(a x_n)**2 on the diagonal. Identical to the element-wise assembly to round-off (rtol 1e-12 in the new test); workspace build time drops accordingly. tests/test_free_matrix_blocks.py checks both against the original formulas and that a block owns only nbasis**2 elements.
…channels ProjectileTargetSystem.get_partial_wave_channels evaluated H+, H-, H+' and H-' with mpmath's coulombf/coulombg separately for every partial wave, and the derivatives by the l -> l+1 recurrence cost two more evaluations each: ~6 arbitrary-precision calls per l, ~500 per workspace, 15-18 s per elastic or (p,n) workspace at lmax=40, and the same cost for neutrons (eta = 0) where spherical Bessel functions would do. free_solutions.coulomb_hankel_table(rho, eta, lmax) builds all four arrays at once: for eta = 0 from scipy's spherical Bessel functions; otherwise G by upward and F by downward (Miller) three-term recurrence (Abramowitz & Stegun 14.2.3), anchored on two mpmath values each, with the derivatives from the exact l -> l+1 relation. Every l is checked against the Wronskian F'G - FG' = 1 and any failing l falls back to the per-l mpmath evaluation. get_partial_wave_channels calls it once per distinct Sommerfeld parameter. Workspace construction drops to ~0.01-1 s; the per-l functions H_plus, H_minus, H_plus_prime, H_minus_prime are unchanged. tests/test_coulomb_hankel_table.py compares the table with the per-l mpmath functions (rtol 1e-9) for eta in [0, 6], rho in [4, 75], l <= 60, checks the Wronskian up to l = 80, and that the channel asymptotics are unchanged.
…nventions - free_solutions: import the shared array aliases from jitr._types (fixes the ruff I001 that failed CI), return a CoulombHankelTable NamedTuple so callers can name the components, rewrite docstrings in the module's imperative style, drop a dead branch in the recurrence anchor - system: use the named table fields instead of positional indices - quadrature: cache the l-independent Legendre kinetic matrix with functools.cached_property instead of a getattr-guarded attribute - rmatrix: share the per-channel mu/E defaulting between kinetic_matrix, energy_matrix and the uncoupled free_matrix branch - tests: wrap long lines, zip(strict=True), check block ownership with np.shares_memory; black-format everything
One implementation per quantity, replacing the four kinetic-matrix and three asymptotics code paths that had accumulated: Kinetic matrix - LagrangeQuadrature base class: subclasses provide only the l-independent radial matrix at a = 1 (cached) and the default overlap; the shared kinetic_matrix(a, l) scales by 1/a**2 and adds the centrifugal diagonal - vectorise the Laguerre radial matrix (Baye Eqs. 3.75-3.77) like the Legendre one; delete both per-element kinetic_operator_element methods - Solver: one _free_blocks primitive; free_matrix(coupled=True) is its block_diag; remove Solver.kinetic_matrix / energy_matrix (no callers) Asymptotics - _coulomb_FG is the only place mpmath's coulombf/coulombg are called; the recurrence table anchors on it and the Wronskian fallback uses it - H_plus, H_minus, H_plus_prime, H_minus_prime are one-line lookups into coulomb_hankel_table; smatrix builds one table for all four (4x faster) - remove FreeAsymptotics, CoulombAsymptotics, coulomb_func_deriv and the asym= keyword on H_*, smatrix and Wavefunctions: eta == 0 already selects the Riccati-Bessel branch - Asymptotics.from_table gathers a partial wave from the per-eta tables Behaviour changes (bugs found while consolidating, verified against exact integrals of the Baye Eq. 3.70 basis functions) - the Laguerre kinetic matrix omitted the 1/a**2 scaling on its Gauss correction term, so it was wrong for any channel radius other than 1 - the Laguerre default overlap (Eq. 3.71) was missing its identity term - the laguerre basis function called the nonexistent sc.special.eval_laguerre and evaluated L_n instead of L_N Tests now check both bases against closed-form elements in the test file, the Laguerre elements against numerical integrals, the 1/a**2 scaling, and the Coulomb-Hankel table and wrappers against mpmath (with mpmath.diff for the derivatives) rather than against each other.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Solver.free_matrix(... coupled=False)assembled a full coupled matrix then returned the diag part as views, which is slow and takes up memory in the work space. This fixes this by building each block independently at a time.mpmathevery time