From 3ca8a2bcda79e60fa1534038acb7e3a9883491db Mon Sep 17 00:00:00 2001 From: beykyle Date: Thu, 11 Jun 2026 14:12:59 -0400 Subject: [PATCH 1/2] final v1.4 updates --- examples/alpha_pb_demo.ipynb | 74 +- .../descouvemont_closed_channels_demo.ipynb | 12 +- examples/descouvemont_np_demo.ipynb | 15 +- examples/descouvemont_o16_ca44_demo.ipynb | 12 +- examples/energy_dependent_demo.ipynb | 61 +- examples/fourier_demo.ipynb | 47 +- examples/hydrogen_demo.ipynb | 25 +- examples/yamaguchi_demo.ipynb | 18 +- src/lax/__init__.py | 3 +- src/lax/boundary/__init__.py | 8 +- src/lax/boundary/_types.py | 841 ------------------ src/lax/boundary/coulomb.py | 17 +- src/lax/compile.py | 114 ++- src/lax/constants.py | 11 +- src/lax/meshes/_basis_eval.py | 2 +- src/lax/meshes/_quadrature.py | 0 src/lax/meshes/_registry.py | 3 +- src/lax/meshes/_utils.py | 24 + src/lax/meshes/laguerre.py | 24 +- src/lax/meshes/legendre.py | 35 +- src/lax/models/optical.py | 8 +- src/lax/models/reid.py | 3 +- src/lax/operators/interaction.py | 204 ++--- src/lax/propagate.py | 2 +- src/lax/solvers/__init__.py | 2 - src/lax/solvers/assembly.py | 44 +- src/lax/solvers/linear_solve.py | 545 ++++++------ src/lax/solvers/observables.py | 161 +--- src/lax/solvers/spectrum.py | 51 +- src/lax/spectral/__init__.py | 3 +- src/lax/spectral/matching.py | 2 +- src/lax/spectral/types.py | 48 +- src/lax/transforms/fourier.py | 2 +- src/lax/transforms/grid.py | 4 +- src/lax/transforms/integration.py | 2 +- src/lax/types.py | 794 ++++++++++++++++- src/lax/wavefunction.py | 4 +- tests/benchmarks/test_alpha_pb_optical.py | 23 +- .../benchmarks/test_coupled_closed_channel.py | 12 +- .../test_descouvemont_closed_channels.py | 4 + .../benchmarks/test_descouvemont_o16_ca44.py | 4 + tests/benchmarks/test_hydrogen.py | 22 +- tests/benchmarks/test_phase8_meshes.py | 8 +- tests/benchmarks/test_yamaguchi.py | 14 +- tests/benchmarks/test_yamaguchi_fourier.py | 4 +- tests/conftest.py | 22 + tests/property/test_autograd.py | 15 +- tests/property/test_hermiticity.py | 8 +- tests/property/test_unitarity.py | 12 +- tests/unit/test_energy_dependent_flow.py | 56 +- tests/unit/test_interaction_builders.py | 122 ++- tests/unit/test_solver_direct.py | 72 +- tests/unit/test_solver_pickle.py | 10 +- tests/unit/test_solver_spectrum.py | 66 +- tests/unit/test_spectral.py | 2 +- tests/unit/test_transforms_fourier.py | 2 +- tests/unit/test_transforms_grid.py | 3 +- 57 files changed, 1931 insertions(+), 1775 deletions(-) delete mode 100644 src/lax/boundary/_types.py delete mode 100644 src/lax/meshes/_quadrature.py create mode 100644 src/lax/meshes/_utils.py diff --git a/examples/alpha_pb_demo.ipynb b/examples/alpha_pb_demo.ipynb index b414145..c4b2dee 100644 --- a/examples/alpha_pb_demo.ipynb +++ b/examples/alpha_pb_demo.ipynb @@ -52,7 +52,9 @@ " ],\n", " dtype=np.complex128,\n", ")\n", - "ALPHA_PB_MASS_FACTOR = lm.constants.hbar2_over_2mu(4.001506, 207.9767) # α + ²⁰⁸Pb MeV·fm²\n", + "ALPHA_PB_MASS_FACTOR = lm.constants.hbar2_over_2mu(\n", + " 4.001506, 207.9767\n", + ") # α + ²⁰⁸Pb MeV·fm²\n", "BENCHMARK_L = 20\n", "CHANNEL_RADIUS = 14.0\n", "\n", @@ -79,7 +81,11 @@ "def complex_solver(method: str, solvers: tuple[str, ...]) -> lm.Solver:\n", " return lm.compile(\n", " mesh=lm.MeshSpec(\"legendre\", \"x\", n=60, scale=CHANNEL_RADIUS),\n", - " channels=(lm.ChannelSpec(l=BENCHMARK_L, threshold=0.0, mass_factor=ALPHA_PB_MASS_FACTOR),),\n", + " channels=(\n", + " lm.ChannelSpec(\n", + " l=BENCHMARK_L, threshold=0.0, mass_factor=ALPHA_PB_MASS_FACTOR\n", + " ),\n", + " ),\n", " operators=(\"T+L\",),\n", " solvers=solvers,\n", " energies=OPTICAL_ENERGIES,\n", @@ -92,7 +98,11 @@ "def real_solver(method: str, solvers: tuple[str, ...]) -> lm.Solver:\n", " return lm.compile(\n", " mesh=lm.MeshSpec(\"legendre\", \"x\", n=60, scale=CHANNEL_RADIUS),\n", - " channels=(lm.ChannelSpec(l=BENCHMARK_L, threshold=0.0, mass_factor=ALPHA_PB_MASS_FACTOR),),\n", + " channels=(\n", + " lm.ChannelSpec(\n", + " l=BENCHMARK_L, threshold=0.0, mass_factor=ALPHA_PB_MASS_FACTOR\n", + " ),\n", + " ),\n", " operators=(\"T+L\",),\n", " solvers=solvers,\n", " energies=OPTICAL_ENERGIES,\n", @@ -101,7 +111,9 @@ " )\n", "\n", "\n", - "def smatrix_from_direct_rmatrix(solver: lm.Solver, potential: jnp.ndarray) -> np.ndarray:\n", + "def smatrix_from_direct_rmatrix(\n", + " solver: lm.Solver, potential: jnp.ndarray\n", + ") -> np.ndarray:\n", " assert solver.rmatrix_direct is not None\n", " assert solver.boundary is not None\n", " r_values = solver.rmatrix_direct(potential)\n", @@ -156,7 +168,9 @@ "fig, axes = plt.subplots(1, 2, figsize=(13, 4.6))\n", "axes[0].plot(r_plot, np.asarray(nuclear_real), label=\"real nuclear\", linewidth=2.2)\n", "axes[0].plot(r_plot, np.asarray(coulomb), label=\"Coulomb\", linewidth=2.2)\n", - "axes[0].plot(r_plot, np.asarray(total.real), \"--\", label=\"total real part\", linewidth=2.0)\n", + "axes[0].plot(\n", + " r_plot, np.asarray(total.real), \"--\", label=\"total real part\", linewidth=2.0\n", + ")\n", "axes[0].set_title(r\"$\\alpha + {}^{208}\\mathrm{Pb}$ real potential pieces\")\n", "axes[0].set_xlabel(\"r [fm]\")\n", "axes[0].set_ylabel(\"MeV\")\n", @@ -197,15 +211,19 @@ "solver_complex_eig = complex_solver(\"eig\", (\"spectrum\", \"smatrix\"))\n", "solver_complex_direct = complex_solver(\"linear_solve\", (\"rmatrix_direct\",))\n", "\n", - "potential_real = solver_real.potential(lambda r: jnp.real(optical_potential(r, imag_depth=0.0)))\n", - "potential_complex_eig = solver_complex_eig.potential(\n", + "potential_real = solver_real.local_potential(\n", + " lambda r: jnp.real(optical_potential(r, imag_depth=0.0))\n", + ")\n", + "potential_complex_eig = solver_complex_eig.local_potential(\n", " lambda r: optical_potential(r, imag_depth=10.0)\n", ")\n", - "potential_complex_direct = solver_complex_direct.potential(\n", + "potential_complex_direct = solver_complex_direct.local_potential(\n", " lambda r: optical_potential(r, imag_depth=10.0)\n", ")\n", "\n", - "smatrix_real = np.asarray(solver_real.smatrix(solver_real.spectrum(potential_real)))[:, 0, 0]\n", + "smatrix_real = np.asarray(solver_real.smatrix(solver_real.spectrum(potential_real)))[\n", + " :, 0, 0\n", + "]\n", "smatrix_complex_eig = np.asarray(\n", " solver_complex_eig.smatrix(solver_complex_eig.spectrum(potential_complex_eig))\n", ")[:, 0, 0]\n", @@ -227,8 +245,12 @@ " )\n", "\n", "print()\n", - "print(f\"max |eig - Appendix A| = {np.max(np.abs(smatrix_complex_eig - APPENDIX_A_S)):.3e}\")\n", - "print(f\"max |direct - Appendix A| = {np.max(np.abs(smatrix_complex_direct - APPENDIX_A_S)):.3e}\")\n", + "print(\n", + " f\"max |eig - Appendix A| = {np.max(np.abs(smatrix_complex_eig - APPENDIX_A_S)):.3e}\"\n", + ")\n", + "print(\n", + " f\"max |direct - Appendix A| = {np.max(np.abs(smatrix_complex_direct - APPENDIX_A_S)):.3e}\"\n", + ")\n", "print(\n", " f\"max |eig - direct| = {np.max(np.abs(smatrix_complex_eig - smatrix_complex_direct)):.3e}\"\n", ")" @@ -254,8 +276,12 @@ "source": [ "fig, axes = plt.subplots(1, 2, figsize=(13, 4.8))\n", "\n", - "axes[0].plot(OPTICAL_ENERGIES, APPENDIX_A_S.real, \"o-\", label=\"Appendix A real\", linewidth=2.2)\n", - "axes[0].plot(OPTICAL_ENERGIES, smatrix_complex_eig.real, \"--\", label=\"eig real\", linewidth=2.0)\n", + "axes[0].plot(\n", + " OPTICAL_ENERGIES, APPENDIX_A_S.real, \"o-\", label=\"Appendix A real\", linewidth=2.2\n", + ")\n", + "axes[0].plot(\n", + " OPTICAL_ENERGIES, smatrix_complex_eig.real, \"--\", label=\"eig real\", linewidth=2.0\n", + ")\n", "axes[0].plot(\n", " OPTICAL_ENERGIES,\n", " smatrix_complex_direct.real,\n", @@ -263,8 +289,12 @@ " label=\"direct real\",\n", " linewidth=2.0,\n", ")\n", - "axes[0].plot(OPTICAL_ENERGIES, APPENDIX_A_S.imag, \"o-\", label=\"Appendix A imag\", linewidth=2.2)\n", - "axes[0].plot(OPTICAL_ENERGIES, smatrix_complex_eig.imag, \"--\", label=\"eig imag\", linewidth=2.0)\n", + "axes[0].plot(\n", + " OPTICAL_ENERGIES, APPENDIX_A_S.imag, \"o-\", label=\"Appendix A imag\", linewidth=2.2\n", + ")\n", + "axes[0].plot(\n", + " OPTICAL_ENERGIES, smatrix_complex_eig.imag, \"--\", label=\"eig imag\", linewidth=2.0\n", + ")\n", "axes[0].plot(\n", " OPTICAL_ENERGIES,\n", " smatrix_complex_direct.imag,\n", @@ -355,7 +385,9 @@ " lm.compile(\n", " mesh=lm.MeshSpec(\"legendre\", \"x\", n=60, scale=CHANNEL_RADIUS),\n", " channels=(\n", - " lm.ChannelSpec(l=angular_momentum, threshold=0.0, mass_factor=ALPHA_PB_MASS_FACTOR),\n", + " lm.ChannelSpec(\n", + " l=angular_momentum, threshold=0.0, mass_factor=ALPHA_PB_MASS_FACTOR\n", + " ),\n", " ),\n", " operators=(\"T+L\",),\n", " solvers=(\"spectrum\", \"smatrix\", \"phases\"),\n", @@ -383,10 +415,14 @@ "abs_s_curves = {}\n", "\n", "for solver, angular_momentum in zip(solvers, partial_waves):\n", - " potential = solver.potential(lambda r: optical_potential(r, imag_depth=10.0))\n", + " potential = solver.local_potential(lambda r: optical_potential(r, imag_depth=10.0))\n", " spectrum = solver.spectrum(potential)\n", - " phase_curves[angular_momentum] = np.asarray(solver.phases(spectrum)[:, 0]) * (180.0 / np.pi)\n", - " abs_s_curves[angular_momentum] = np.abs(np.asarray(solver.smatrix(spectrum)[:, 0, 0]))" + " phase_curves[angular_momentum] = np.asarray(solver.phases(spectrum)[:, 0]) * (\n", + " 180.0 / np.pi\n", + " )\n", + " abs_s_curves[angular_momentum] = np.abs(\n", + " np.asarray(solver.smatrix(spectrum)[:, 0, 0])\n", + " )" ] }, { diff --git a/examples/descouvemont_closed_channels_demo.ipynb b/examples/descouvemont_closed_channels_demo.ipynb index 182a878..ba320f0 100644 --- a/examples/descouvemont_closed_channels_demo.ipynb +++ b/examples/descouvemont_closed_channels_demo.ipynb @@ -35,7 +35,9 @@ " candidate = root / \"tests\" / \"benchmarks\" / \"data\"\n", " if candidate.is_dir():\n", " return candidate\n", - " msg = \"Could not locate tests/benchmarks/data from the current notebook environment.\"\n", + " msg = (\n", + " \"Could not locate tests/benchmarks/data from the current notebook environment.\"\n", + " )\n", " raise FileNotFoundError(msg)\n", "\n", "\n", @@ -114,9 +116,13 @@ "rows = []\n", "for energy_index, energy in enumerate(energies):\n", " boundary = boundary_at_energy(solver.boundary, energy_index)\n", - " smatrix = np.asarray(lm.spectral.open_channel_smatrix_from_R(r_values[energy_index], boundary))\n", + " smatrix = np.asarray(\n", + " lm.spectral.open_channel_smatrix_from_R(r_values[energy_index], boundary)\n", + " )\n", " open_count = lm.models.open_channel_count(model, float(energy))\n", - " amplitudes, phases = lm.models.first_column_amplitudes_and_phases(smatrix, open_count)\n", + " amplitudes, phases = lm.models.first_column_amplitudes_and_phases(\n", + " smatrix, open_count\n", + " )\n", " rows.append(\n", " {\n", " \"energy_mev\": float(energy),\n", diff --git a/examples/descouvemont_np_demo.ipynb b/examples/descouvemont_np_demo.ipynb index 5221914..f54d5fd 100644 --- a/examples/descouvemont_np_demo.ipynb +++ b/examples/descouvemont_np_demo.ipynb @@ -56,7 +56,9 @@ " candidate = root / \"tests\" / \"benchmarks\" / \"data\"\n", " if candidate.is_dir():\n", " return candidate\n", - " msg = \"Could not locate tests/benchmarks/data from the current notebook environment.\"\n", + " msg = (\n", + " \"Could not locate tests/benchmarks/data from the current notebook environment.\"\n", + " )\n", " raise FileNotFoundError(msg)\n", "\n", "\n", @@ -65,7 +67,9 @@ "\n", "energies = np.asarray(reference[\"energies\"], dtype=np.float64)\n", "channels = lm.models.reid_np_j1_channels()\n", - "mesh = lm.MeshSpec(\"legendre\", \"x\", n=int(reference[\"n_basis\"]), scale=float(reference[\"scale\"]))\n", + "mesh = lm.MeshSpec(\n", + " \"legendre\", \"x\", n=int(reference[\"n_basis\"]), scale=float(reference[\"scale\"])\n", + ")\n", "\n", "{\n", " \"energies_mev\": energies.tolist(),\n", @@ -131,7 +135,8 @@ "source": [ "sample_radii = np.asarray([0.5, 1.0, 2.0, 4.0, 6.0], dtype=np.float64)\n", "central, tensor, spin_orbit = [\n", - " np.asarray(values) for values in lm.models.reid_soft_core_triplet_components(sample_radii)\n", + " np.asarray(values)\n", + " for values in lm.models.reid_soft_core_triplet_components(sample_radii)\n", "]\n", "\n", "[\n", @@ -141,7 +146,9 @@ " \"tensor_mev\": float(v_t),\n", " \"spin_orbit_mev\": float(v_ls),\n", " }\n", - " for radius, v_c, v_t, v_ls in zip(sample_radii, central, tensor, spin_orbit, strict=True)\n", + " for radius, v_c, v_t, v_ls in zip(\n", + " sample_radii, central, tensor, spin_orbit, strict=True\n", + " )\n", "]" ] }, diff --git a/examples/descouvemont_o16_ca44_demo.ipynb b/examples/descouvemont_o16_ca44_demo.ipynb index 177d262..e137dbc 100644 --- a/examples/descouvemont_o16_ca44_demo.ipynb +++ b/examples/descouvemont_o16_ca44_demo.ipynb @@ -35,7 +35,9 @@ " candidate = root / \"tests\" / \"benchmarks\" / \"data\"\n", " if candidate.is_dir():\n", " return candidate\n", - " msg = \"Could not locate tests/benchmarks/data from the current notebook environment.\"\n", + " msg = (\n", + " \"Could not locate tests/benchmarks/data from the current notebook environment.\"\n", + " )\n", " raise FileNotFoundError(msg)\n", "\n", "\n", @@ -114,9 +116,13 @@ "rows = []\n", "for energy_index, energy in enumerate(energies):\n", " boundary = boundary_at_energy(solver.boundary, energy_index)\n", - " smatrix = np.asarray(lm.spectral.open_channel_smatrix_from_R(r_values[energy_index], boundary))\n", + " smatrix = np.asarray(\n", + " lm.spectral.open_channel_smatrix_from_R(r_values[energy_index], boundary)\n", + " )\n", " open_count = lm.models.open_channel_count(model, float(energy))\n", - " amplitudes, phases = lm.models.first_column_amplitudes_and_phases(smatrix, open_count)\n", + " amplitudes, phases = lm.models.first_column_amplitudes_and_phases(\n", + " smatrix, open_count\n", + " )\n", " rows.append(\n", " {\n", " \"energy_mev\": float(energy),\n", diff --git a/examples/energy_dependent_demo.ipynb b/examples/energy_dependent_demo.ipynb index 9df1e9f..37f57ba 100644 --- a/examples/energy_dependent_demo.ipynb +++ b/examples/energy_dependent_demo.ipynb @@ -4,21 +4,7 @@ "cell_type": "markdown", "id": "4b552bb9", "metadata": {}, - "source": [ - "# Energy-dependent potentials\n", - "\n", - "This notebook demonstrates the `energy_dependent=True` workflow:\n", - "\n", - "1. Define a potential whose depth varies with energy.\n", - "2. Compile a solver with `energy_dependent=True`.\n", - "3. Compute a batched `Spectrum` via `jax.vmap`.\n", - "4. Use the aligned-grid helpers (`phases_grid`) to get phase shifts.\n", - "5. Build a Padé interpolant and compare it against the full reference grid.\n", - "\n", - "The toy system is an s-wave Gaussian optical potential\n", - "$V(r; E) = -V_0(E)\\,e^{-r^2/a^2}$ with a linear energy dependence\n", - "$V_0(E) = V_1 + V_2 E$." - ] + "source": "# Energy-dependent potentials\n\nThis notebook demonstrates the `energy_dependent=True` workflow:\n\n1. Define a potential whose depth varies with energy.\n2. Compile a solver with `energy_dependent=True`.\n3. Compute a batched `Spectrum` via `solver.spectrum`'s internal dispatch.\n4. Use the aligned-grid helpers (`phases_grid`) to get phase shifts.\n5. Build a Padé interpolant and compare it against the full reference grid.\n\nThe toy system is an s-wave Gaussian optical potential\n$V(r; E) = -V_0(E)\\,e^{-r^2/a^2}$ with a linear energy dependence\n$V_0(E) = V_1 + V_2 E$." }, { "cell_type": "code", @@ -29,7 +15,6 @@ "source": [ "from __future__ import annotations\n", "\n", - "import jax\n", "import jax.numpy as jnp\n", "import matplotlib.pyplot as plt\n", "import numpy as np\n", @@ -108,7 +93,7 @@ "cell_type": "markdown", "id": "2bd8fd80", "metadata": {}, - "source": "## Build the energy-dependent potential grid\n\nUse `solver.potential(fn, energy_dependent=True)` to build an\nenergy-dependent `Interaction` whose `.block` has shape `(N_E, M, M)`.\nPass `.block` to `jax.vmap(solver.spectrum)` for the diagonal\n`(spec_i, E_i)` pairing." + "source": "## Build the energy-dependent potential grid\n\nUse `solver.local_potential(fn, energy_dependent=True)` to build an\nenergy-dependent `Interaction` whose `.block` has shape `(N_E, M, M)`.\nPass the `Interaction` straight to `solver.spectrum` — it dispatches over\nthe leading energy axis internally for the diagonal `(spec_i, E_i)` pairing." }, { "cell_type": "code", @@ -118,8 +103,8 @@ "outputs": [], "source": [ "# Build the energy-dependent interaction; .block has shape (N_E, M, M)\n", - "assert solver.potential is not None\n", - "interaction = solver.potential(gaussian_potential, energy_dependent=True)\n", + "assert solver.local_potential is not None\n", + "interaction = solver.local_potential(gaussian_potential, energy_dependent=True)\n", "V_grid = interaction.block # (N_E, M, M) = (N_E, N, N) for single channel\n", "print(\"V_grid shape:\", V_grid.shape)" ] @@ -128,32 +113,18 @@ "cell_type": "markdown", "id": "ef4cd0d4", "metadata": {}, - "source": [ - "## Compute the batched spectrum and aligned-grid phase shifts\n", - "\n", - "`jax.vmap(solver.spectrum)` vectorises over the leading energy axis of\n", - "`V_grid`, producing one `Spectrum` per energy. `solver.phases_grid` then\n", - "evaluates each spectrum at its own compile-time energy — the diagonal\n", - "pairing `(spec_i, E_i)` that is physically correct for energy-dependent V." - ] + "source": "## Compute the batched spectrum and aligned-grid phase shifts\n\n`solver.spectrum` dispatches over the leading energy axis of the\nenergy-dependent `Interaction` internally, producing one `Spectrum` per\nenergy. `solver.phases_grid` then evaluates each spectrum at its own\ncompile-time energy — the diagonal pairing `(spec_i, E_i)` that is\nphysically correct for energy-dependent V." }, { "cell_type": "code", - "execution_count": 5, + "execution_count": null, "id": "9941c8f3", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "phase shifts at coarse grid (deg): [-32.7 -53.67 -66.63 -76.23 -83.9 89.71 84.24 79.45 75.19 71.37\n", - " 67.89 64.7 61.76 59.03 56.48 54.09]\n" - ] - } - ], + "outputs": [], "source": [ - "spectra = jax.vmap(solver.spectrum)(V_grid) # batched Spectrum (N_E, ...)\n", + "spectra = solver.spectrum(\n", + " interaction\n", + ") # batched Spectrum (N_E, ...) via internal dispatch\n", "phases_coarse = np.asarray(solver.phases_grid(spectra))[:, 0] # (N_E,) rad\n", "print(\"phase shifts at coarse grid (deg):\", np.degrees(phases_coarse).round(2))" ] @@ -176,7 +147,9 @@ "metadata": {}, "outputs": [], "source": [ - "interp_phases = solver.interpolate_phases(jnp.asarray(phases_coarse[:, None])) # (N_E, 1)\n", + "interp_phases = solver.interpolate_phases(\n", + " jnp.asarray(phases_coarse[:, None])\n", + ") # (N_E, 1)\n", "\n", "# Evaluate interpolant on a fine grid\n", "phases_interp = np.asarray(interp_phases(energies_fine))[:, 0] # (N_FINE,)" @@ -208,11 +181,11 @@ " energies=energies_fine,\n", ")\n", "\n", - "assert solver_ref.potential is not None\n", + "assert solver_ref.local_potential is not None\n", "phases_ref = np.empty(N_FINE)\n", "for idx, e in enumerate(np.asarray(energies_fine)):\n", " e_val = float(e)\n", - " interaction_e = solver_ref.potential(lambda r: gaussian_potential(r, e_val))\n", + " interaction_e = solver_ref.local_potential(lambda r: gaussian_potential(r, e_val))\n", " spec_e = solver_ref.spectrum(interaction_e)\n", " phases_ref[idx] = float(solver_ref.phases(spec_e)[idx, 0])\n", "\n", @@ -253,7 +226,9 @@ "\n", "axes[0].plot(E_fine, np.degrees(phases_ref), label=\"reference (fine grid)\", lw=2)\n", "axes[0].plot(E_fine, np.degrees(phases_interp), \"--\", label=\"Padé interpolant\", lw=1.8)\n", - "axes[0].scatter(E_coarse, np.degrees(phases_coarse), zorder=5, label=\"coarse knots\", s=30)\n", + "axes[0].scatter(\n", + " E_coarse, np.degrees(phases_coarse), zorder=5, label=\"coarse knots\", s=30\n", + ")\n", "axes[0].set_xlabel(\"Energy [MeV]\")\n", "axes[0].set_ylabel(\"Phase shift [deg]\")\n", "axes[0].set_title(r\"Energy-dependent Gaussian: $\\ell=0$ phase shift\")\n", diff --git a/examples/fourier_demo.ipynb b/examples/fourier_demo.ipynb index 51ed3de..8a06c84 100644 --- a/examples/fourier_demo.ipynb +++ b/examples/fourier_demo.ipynb @@ -130,7 +130,9 @@ " \"Gaussian x polynomial l=0\",\n", " 0,\n", " lambda r: r**2 * (3.0 / (2.0 * beta) - r**2) * np.exp(-beta * r**2),\n", - " lambda k: ((k**2) / (4.0 * beta**2)) * np.exp(-(k**2) / (4.0 * beta)) / (2.0 * beta) ** 1.5,\n", + " lambda k: ((k**2) / (4.0 * beta**2))\n", + " * np.exp(-(k**2) / (4.0 * beta))\n", + " / (2.0 * beta) ** 1.5,\n", " ),\n", " (\n", " \"Gaussian l=1\",\n", @@ -220,14 +222,18 @@ "momenta = np.asarray(numerical_solver.transforms.momenta)\n", "\n", "for row, (name, profile) in enumerate(numerical_profiles):\n", - " coeffs = numerical_solver.from_grid_vector(lambda r: jnp.asarray(profile(np.asarray(r))))\n", + " coeffs = numerical_solver.from_grid_vector(\n", + " lambda r: jnp.asarray(profile(np.asarray(r)))\n", + " )\n", " reconstructed = np.asarray(numerical_solver.to_grid_vector(coeffs))\n", " transformed = np.asarray(numerical_solver.fourier(coeffs))\n", " expected = profile(grid_r)\n", " print(name)\n", " print(\" grid relative error =\", relative_error(reconstructed, expected))\n", " axes[row, 0].plot(grid_r, expected, label=\"input profile\", linewidth=2.0)\n", - " axes[row, 0].plot(grid_r, reconstructed, \"--\", label=\"mesh reconstruction\", linewidth=1.8)\n", + " axes[row, 0].plot(\n", + " grid_r, reconstructed, \"--\", label=\"mesh reconstruction\", linewidth=1.8\n", + " )\n", " axes[row, 0].set_title(f\"{name} in r-space\")\n", " axes[row, 0].set_xlabel(\"r\")\n", " axes[row, 0].set_ylabel(\"u(r)\")\n", @@ -318,33 +324,10 @@ }, { "cell_type": "code", - "execution_count": 6, + "execution_count": null, "id": "3468e66a", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Hydrogen 1s r-space norm (dr) ≈ 0.99999958\n", - "Hydrogen 1s k-space amplitude norm (dk) ≈ 1.98766850\n", - "Analytic 1s k-space amplitude norm on this grid ≈ 1.98766922\n", - "The current solver.fourier(...) convention computes the partial-wave amplitude\n", - "sqrt(2/pi) ∫ j_l(k r) u_l(r) dr, so its plain-dk norm is not expected\n", - "to equal the r-space dr norm.\n" - ] - }, - { - "data": { - "image/png": "iVBORw0KGgoAAAANSUhEUgAABKUAAAG4CAYAAACUzVSpAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjksIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvJkbTWQAAAAlwSFlzAAAPYQAAD2EBqD+naQAAl5VJREFUeJzs3XlcVPX6B/DPLDDDOuyLiOKuiALihktqYmhl2W2xLBcsS5M2bhs/b5rdCiszWyzLm2lZaVmZlWlKkpqYK+a+I4isIjvMwMz5/TFyYBIUYWbODHzer9e5zpz5nplnRm98eeb5Pl+ZIAgCiIiIiIiIiIiIrEgudQBERERERERERNT2MClFRERERERERERWx6QUERERERERERFZHZNSRERERERERERkdUxKERERERERERGR1TEpRUREREREREREVsekFBERERERERERWR2TUkREREREREREZHVMShERERERERERkdUxKUXUxqWnp0Mmk2HFihVSh2L3Ro4ciZEjR0odBhEREVkY50/mY+7508svvwyZTIaCggKzPScRWQ6TUkQ2ZsWKFZDJZNi7d2+Dj48cORJhYWFWjsr+lZWVYd68eRg7diy8vLw4kSQiImpFOH+yDM6fiMjSlFIHQERkDQUFBXjllVfQoUMHhIeHIyUlxeyv8dtvv5n9OYmIiIikwvkTEVkak1JE1CTl5eVwcXGROoxmCwwMRHZ2NgICArB3714MGDDA7K/h6Oho9uf8J3v/eyAiImpL7P3ndmuZPxGR7eLyPSI7N2LECISHhzf4WI8ePRAbGyveLyoqwrRp06DRaODh4YGpU6eiqKjoquumTZsGV1dXnDlzBrfeeivc3Nzw4IMPAjBOrv79738jODgYKpUKPXr0wMKFCyEIgslzVFZW4sknn4SPjw/c3Nxwxx13ICsrCzKZDC+//LLJ2KysLEyfPh3+/v5QqVTo3bs3li9fbjImJSUFMpkM33zzDV577TW0b98earUao0ePxunTp6/7OalUKgQEBFx3HADs3bsXsbGx8PHxgZOTEzp16oTp06df97p/9kRoacy1PRGOHj2KSZMmwdPTE8OGDWt0/KlTp3D33XcjICAAarUa7du3x/3334/i4mJxjEwmQ3x8PL788kv06NEDarUaUVFR2LZtm8lznT9/Ho8//jh69OgBJycneHt7495770V6evpVr1tUVIRnnnkGISEhUKlUaN++PaZMmWLSy0Gr1WLevHno2rUrVCoVgoOD8fzzz0Or1V73cyAiIjI3zp9a7/ypIefPn0fXrl0RFhaG3NzcFr2P2n5iCxcuxDvvvIOOHTvCyckJI0aMwOHDh02e6++//8a0adPQuXNnqNVqBAQEYPr06bh06dJVr5uVlYWHH34Y7dq1g0qlQqdOnTBr1izodDpxTFFREZ5++mnx31HXrl3xxhtvwGAwNOtzIbIFrJQislHFxcUNNmisrq42uT958mTMmDEDhw8fNumVsGfPHpw8eRL/+c9/AACCIODOO+/Ejh07MHPmTPTq1Qs//PADpk6d2uDr19TUIDY2FsOGDcPChQvh7OwMQRBwxx13YOvWrXj44YcRERGBTZs24bnnnkNWVhbeeecd8fpp06bhm2++weTJkzF48GD88ccfuO222656ndzcXAwePFhMlvj6+uLXX3/Fww8/jJKSEjz99NMm4xcsWAC5XI5nn30WxcXFePPNN/Hggw/ir7/+avJney15eXm45ZZb4OvrixdffBEeHh5IT0/H999/3+znbGnM9957L7p164bXX3/9qslrLZ1Oh9jYWGi1WjzxxBMICAhAVlYWfv75ZxQVFUGj0Yhj//jjD6xZswZPPvkkVCoVPvzwQ4wdOxa7d+8W/w3t2bMHO3fuxP3334/27dsjPT0dH330EUaOHImjR4/C2dkZgLHXxPDhw3Hs2DFMnz4d/fr1Q0FBAdavX48LFy7Ax8cHBoMBd9xxB3bs2IFHH30UvXr1wqFDh/DOO+/g5MmTWLduXbM/WyIiovo4f+L86Z/OnDmDm2++GV5eXti8eTN8fHzM8j4+//xzlJaWYvbs2aiqqsK7776Lm2++GYcOHYK/vz8AYPPmzTh79izi4uIQEBCAI0eO4JNPPsGRI0ewa9cuyGQyAMDFixcxcOBAFBUV4dFHH0XPnj2RlZWFtWvXoqKiAo6OjqioqMCIESOQlZWFxx57DB06dMDOnTuRmJiI7OxsLF68+IY+FyKbIRCRTfnss88EANc8evfuLY4vKioS1Gq18MILL5g8z5NPPim4uLgIZWVlgiAIwrp16wQAwptvvimOqampEYYPHy4AED777DPx/NSpUwUAwosvvmjynLXP8eqrr5qcv+eeewSZTCacPn1aEARB2LdvnwBAePrpp03GTZs2TQAgzJs3Tzz38MMPC4GBgUJBQYHJ2Pvvv1/QaDRCRUWFIAiCsHXrVgGA0KtXL0Gr1Yrj3n33XQGAcOjQoWt+rvXt2bPnqvdc64cffhAACHv27Gny89UaMWKEMGLECPF+S2OeN2+eAEB44IEHrvvaBw4cEAAI33777TXH1f4b2rt3r3ju/PnzglqtFu666y7xXO3nXl9qaqoAQPj888/Fc3PnzhUACN9///1V4w0GgyAIgvDFF18Icrlc2L59u8njS5cuFQAIf/7553XfHxER0bVw/mTE+VPd/Ck/P184duyY0K5dO2HAgAFCYWHhdWNpyvs4d+6cAEBwcnISLly4IJ7/66+/BADCM888I55raD719ddfCwCEbdu2ieemTJkiyOXyBl+3dj713//+V3BxcRFOnjxp8viLL74oKBQKISMj47rvj8gWcfkekY1asmQJNm/efNXRt29fk3EajQZ33nknvv76a7GKRq/XY82aNZgwYYLYx2DDhg1QKpWYNWuWeK1CocATTzzRaAz1x9Y+h0KhwJNPPmly/t///jcEQcCvv/4KANi4cSMA4PHHHzcZ98/XEgQB3333HcaPHw9BEFBQUCAesbGxKC4uxv79+02uiYuLM+k9MHz4cADA2bNnG30fN8LDwwMA8PPPP1/1rWpztTTmmTNnXndMbSXUpk2bUFFRcc2x0dHRiIqKEu936NABd955JzZt2gS9Xg8AcHJyEh+vrq7GpUuX0LVrV3h4eJj8nXz33XcIDw/HXXfdddXr1H779+2336JXr17o2bOnyd/xzTffDADYunXrdd8fERFRU3D+xPlTrcOHD2PEiBEICQnBli1b4Onped1rbuR9TJgwAUFBQeL9gQMHYtCgQdiwYYN4rv58qqqqCgUFBRg8eDAAiH9HBoMB69atw/jx49G/f/+rXqf+fGr48OHw9PQ0+TuPiYmBXq+/qhUDkb3g8j0iGzVw4MAGfzDV/iCqb8qUKVizZg22b9+Om266CVu2bEFubi4mT54sjjl//jwCAwPh6upqcm2PHj0afH2lUon27dubnDt//jzatWsHNzc3k/O9evUSH6/9Uy6Xo1OnTibjunbtanI/Pz8fRUVF+OSTT/DJJ580GEdeXp7J/Q4dOpjcr51gXL58ucHrb9SIESNw9913Y/78+XjnnXcwcuRITJgwAZMmTYJKpWrWc7Y05vqfY2VlpUmPKAAICAhAp06dkJCQgEWLFuHLL7/E8OHDcccdd+Chhx4yWboHAN26dbvqNbp3746Kigrk5+cjICAAlZWVSEpKwmeffYasrCyTZYP1X//MmTO4++67rxn/qVOncOzYMfj6+jb4+D//jomIiJqL8ycjzp+A8ePHw9/fH5s2bbrq76+srAxlZWXifYVCAV9f3xt6H43Np7755hvxfmFhIebPn4/Vq1df9XdSO5/Kz89HSUmJyTLShpw6dQp///0351PU6jApRdQKxMbGwt/fH6tWrcJNN92EVatWISAgADExMc1+TpVKBbncssWUtU0ZH3rooUZ7M/zzm02FQtHgOKGRXks3SiaTYe3atdi1axd++uknbNq0CdOnT8fbb7+NXbt2XTWpaYqWxlz/W7Y1a9YgLi6uwed5++23MW3aNPz444/47bff8OSTTyIpKQm7du26aoJ8PU888QQ+++wzPP3004iOjoZGo4FMJsP9999/w800DQYD+vTpg0WLFjX4eHBw8A09HxERkTlw/tS650933303Vq5ciS+//BKPPfaYyWMLFy7E/PnzxfsdO3YUG5ib833cd9992LlzJ5577jlERETA1dUVBoMBY8eObdZ8asyYMXj++ecbfLx79+439HxEtoJJKaJWQKFQYNKkSVixYgXeeOMNrFu3DjNmzDD5Yd6xY0ckJyejrKzM5AfqiRMnmvw6HTt2xJYtW1BaWmrybd/x48fFx2v/NBgMOHfunMm3SP/cMcXX1xdubm7Q6/UtmgBawuDBgzF48GC89tpr+Oqrr/Dggw9i9erVeOSRRySNKzY2Fps3b2708T59+qBPnz74z3/+g507d2Lo0KFYunQpXn31VXHMqVOnrrru5MmTcHZ2Fr99W7t2LaZOnYq3335bHFNVVXXVbkNdunS5aqeZf+rSpQsOHjyI0aNHiyXoREREUuP8yfxsaf701ltvQalU4vHHH4ebmxsmTZokPjZlyhSTHY3rfwEINO19NDafCgkJAWCs6EpOTsb8+fMxd+7cRq/z9fWFu7t7k+ZTZWVlNvd3TtRS7ClF1EpMnjwZly9fxmOPPYaysjI89NBDJo/feuutqKmpwUcffSSe0+v1eP/995v8Grfeeiv0ej0++OADk/PvvPMOZDIZxo0bBwDiNsoffvihybh/vpZCocDdd9+N7777rsEfxPn5+U2OzVwuX7581TdwERERAACtVmv1eP4pMDAQMTExJgcAlJSUoKamxmRsnz59IJfLr4o7NTXVpNdEZmYmfvzxR9xyyy3iRFyhUFz1Obz//vtiz6lad999Nw4ePIgffvjhqlhrr7/vvvuQlZWFZcuWXTWmsrIS5eXlTX37REREZsX5k3nY4vxJJpPhk08+wT333IOpU6di/fr14mOdO3c2mUsNHToUwI29j3Xr1iErK0u8v3v3bvz111/i32ftnOqfz/fPXfLkcjkmTJiAn376CXv37r3qfdSfT6WmpmLTpk1XjSkqKrpqHkhkL1gpRdRKREZGIiwsTGwq3a9fP5PHx48fj6FDh+LFF19Eeno6QkND8f3331/Vn+haxo8fj1GjRmHOnDlIT09HeHg4fvvtN/z44494+umn0aVLFwBAVFQU7r77bixevBiXLl0StzQ+efIkAJhUyyxYsABbt27FoEGDMGPGDISGhqKwsBD79+/Hli1bUFhYaIZPx+iDDz5AUVERLl68CAD46aefcOHCBQDG5WoajQYrV67Ehx9+iLvuugtdunRBaWkpli1bBnd3d9x6661mi8Xcfv/9d8THx+Pee+9F9+7dUVNTgy+++EKcuNYXFhaG2NhYPPnkk1CpVOLkt34Z++23344vvvgCGo0GoaGhSE1NxZYtW+Dt7W3yXM899xzWrl2Le++9F9OnT0dUVBQKCwuxfv16LF26FOHh4Zg8eTK++eYbzJw5E1u3bsXQoUOh1+tx/PhxfPPNN9i0aVOD/T+IiIgsjfOn67Pn+ZNcLseqVaswYcIE3HfffdiwYYO40UpDbuR9dO3aFcOGDcOsWbOg1WqxePFieHt7i8vr3N3dcdNNN+HNN99EdXU1goKC8Ntvv+HcuXNXve7rr7+O3377DSNGjMCjjz6KXr16ITs7G99++y127NgBDw8PPPfcc1i/fj1uv/12TJs2DVFRUSgvL8ehQ4ewdu1apKenw8fHx7wfIJE1WHWvPyK6rtotjRvbinbEiBEmWxrX9+abbwoAhNdff73Bxy9duiRMnjxZcHd3FzQajTB58mThwIEDDW5p7OLi0uBzlJaWCs8884zQrl07wcHBQejWrZvw1ltvidvV1iovLxdmz54teHl5Ca6ursKECROEEydOCACEBQsWmIzNzc0VZs+eLQQHBwsODg5CQECAMHr0aOGTTz4Rx9RuD/ztt9+aXFu7LW9D2xP/U8eOHRvdJvrcuXOCIAjC/v37hQceeEDo0KGDoFKpBD8/P+H2228X9u7de93nb2xL4+bGXH9L4+s5e/asMH36dKFLly6CWq0WvLy8hFGjRglbtmwxGQdAmD17trBq1SqhW7dugkqlEiIjI4WtW7eajLt8+bIQFxcn+Pj4CK6urkJsbKxw/PhxoWPHjsLUqVNNxl66dEmIj48XgoKCBEdHR6F9+/bC1KlTTbap1ul0whtvvCH07t1bUKlUgqenpxAVFSXMnz9fKC4uvu77IyIiuhbOnzh/qtXQ/KmiokIYMWKE4OrqKuzatavRa5vyPmrjeOutt4S3335bCA4OFlQqlTB8+HDh4MGDJs934cIF4a677hI8PDwEjUYj3HvvvcLFixcFAMK8efNMxp4/f16YMmWK4OvrK6hUKqFz587C7NmzBa1WK44pLS0VEhMTha5duwqOjo6Cj4+PMGTIEGHhwoWCTqe75udCZKtkgmCm7nZEJLl3330XzzzzDNLT06/ascQWpKWlITIyEqtWrcKDDz4odThtkkwmw+zZs69aQkBERNRWcf5ENyI9PR2dOnXCW2+9hWeffVbqcIjsHntKEbUSgiDg008/xYgRI2xiQlVZWXnVucWLF0Mul+Omm26SICIiIiIiU5w/ERFJiz2liOxceXk51q9fj61bt+LQoUP48ccfpQ4JAPDmm29i3759GDVqFJRKJX799Vf8+uuvePTRRxEcHCx1eERERNSGcf5ERGQbmJQisnP5+fmYNGkSPDw88H//93+44447pA4JADBkyBBs3rwZ//3vf1FWVoYOHTrg5Zdfxpw5c6QOjYiIiNo4zp+IiGwDe0oREREREREREZHVsacUERERERERERFZHZNSRERERERERERkdW2up5TBYMDFixfh5uYGmUwmdThERERkowRBQGlpKdq1awe5nN/jAZxHERERUdM0dR7V5pJSFy9e5M4VRERE1GSZmZlo37691GHYBM6jiIiI6EZcbx7V5pJSbm5uAIwfjLu7u8TREBERka0qKSlBcHCwOHcgzqOIiIioaZo6j2pzSanaUnN3d3dOpoiIiOi6uEytDudRREREdCOuN49igwQiIiIiIiIiIrI6JqWIiIiIiIiIiMjqmJQiIiIiIiIiIiKrY1KKiIiIiIiIiIisjkkpIiIiIiIiIiKyOialiIiIiIiIiIjI6piUIiIiIiIiIiIiq2NSioiIiIiIiIiIrI5JKSIiIiI7tG3bNowfPx7t2rWDTCbDunXrrnuNVqvFnDlz0LFjR6hUKoSEhGD58uWWD5aIiIioAUqpAyAiIiKiG1deXo7w8HBMnz4d//rXv5p0zX333Yfc3Fx8+umn6Nq1K7Kzs2EwGCwcKREREVHDmJQiIiIiskPjxo3DuHHjmjx+48aN+OOPP3D27Fl4eXkBAEJCQiwUHREREdH1cfkeERERURuwfv169O/fH2+++SaCgoLQvXt3PPvss6isrJQ6NACAIAgQBEHqMIiIiMiKWCkloVO5pZjzw2GUamswb3woBnf2ljokIiIiaqXOnj2LHTt2QK1W44cffkBBQQEef/xxXLp0CZ999lmD12i1Wmi1WvF+SUmJxeLbcCgHT60+AFe1Eq4qJTRODujs64oe/q6I7OCJASFecFTy+1QiIqLWhEkpiVTrDXh45V5kFFYAAKav2IOtz46Ev7ta4siIiIioNTIYDJDJZPjyyy+h0WgAAIsWLcI999yDDz/8EE5OTlddk5SUhPnz51slvnJtDWoMAooqqlFUUY0Llytx5GIJfrryuLOjAsO7+eCeqGCM6uELpYIJKiIiInvHn+YS2XAoW0xIAUCFTo9Ptp2VMCIiIiJqzQIDAxEUFCQmpACgV69eEAQBFy5caPCaxMREFBcXi0dmZqbF4nNRKdG7nTs6ejvDx9URSrnM5PEKnR6bjuRixud7MfSN3/FRyhmUa2ssFg8RERFZHiulJLL5aO5V5346eBFzbu0F+T8mYUREREQtNXToUHz77bcoKyuDq6srAODkyZOQy+Vo3759g9eoVCqoVCqrxHdb30Dc1jdQvK+rMSD9UjmOXCzG9pMF+ONkPi6V6wAAuSVavLHxOJZtP4vHR3bB1CEhcGDlFBERkd3hT28JGAwCUs9cAgC4qZQY1cMXAJBXqsWBzCIJIyMiIiJ7UVZWhrS0NKSlpQEAzp07h7S0NGRkZAAwVjlNmTJFHD9p0iR4e3sjLi4OR48exbZt2/Dcc89h+vTpDS7dk5qjUo7u/m64K7I9Fk2MwO45Mfhs2gDE9PJH7fd3heU6vPrLMdz23nbsSS+UNmAiIiK6YUxKSeB8YYX4Td/ATl4YGxYgPpZ6pkCqsIiIiMiO7N27F5GRkYiMjAQAJCQkIDIyEnPnzgUAZGdniwkqAHB1dcXmzZtRVFSE/v3748EHH8T48ePx3nvvSRL/jVLIZRjV0w//m9ofvz0zAuPD20F2JTl1MrcM9y5Nxcvrj6CqWi9toERERNRkXL4ngePZdTvXhAVpEN3ZR7yfevYS4m/uJkVYREREZEdGjhwJQRAafXzFihVXnevZsyc2b95swaiso6ufK95/IBKPDu+MOesO4e8LxQCAFTvTsevsJbz/QCS6+btJHCURERFdDyulJHAsp1S83SvQDcFeTgi4suve35nFMBgan2ASERERkVGf9hr88PhQzBsfCpXSOK09nlOKO5f8iY2HcySOjoiIiK6HSSkJnMipq5TqEeAOmUyGsCB3AECptgYXLldKFRoRERGRXVHIZYgb2gnr44ehu7+xgXuFTo+Zq/bh/eRT16wmIyIiImkxKSWB85cqAABKuQwdvJwBAKGB7uLjRy4WSxIXERERkb3qEeCG9fHDMCGinXju7c0n8eJ3h1CjN0gYGRERETWGSSkrEwQBmYXGpFR7TycormwfE9pOI445Wq/nFBERERE1jdpBgXcmRuCFsT3FJuhr9mZi9lf72QCdiIjIBjEpZWWXK6pRrjNOioKvVEkBQO929SulmJQiIiIiag6ZTIZZI7vggwf6wUFhzExtOpKLR1buZWKKiIjIxjApZWUZV6qkANOkVHtPJ7irjZshHmVSioiIiKhFbusbiOXTBsDZUQEA2HG6ALNW7YOuhkv5iIiIbIVNJKWWLFmCkJAQqNVqDBo0CLt372507MiRIyGTya46brvtNitG3HxZ9ZqYt/d0Em/LZDL0utJXKqekCpfLdVaPjYiIiKg1Gd7NF188PFBMTG09kY8nvt6PavaYIiIisgmSJ6XWrFmDhIQEzJs3D/v370d4eDhiY2ORl5fX4Pjvv/8e2dnZ4nH48GEoFArce++9Vo68eXJLqsTbgRq1yWNd/VzF2+culVstJiIiIqLWKqqjF5ZPGwC1g3Hau+lILhK+OQiDgbvyERERSU3ypNSiRYswY8YMxMXFITQ0FEuXLoWzszOWL1/e4HgvLy8EBASIx+bNm+Hs7GyXSSl/N9OkVGffuqTU2XwmpYiIiIjMYXBnb/xvygA4Ko1T358OXsSCjccljoqIiIgkTUrpdDrs27cPMTEx4jm5XI6YmBikpqY26Tk+/fRT3H///XBxcWnwca1Wi5KSEpNDSvWTUn7u/0hK+dS9h3MFZVaLiYiIiKi1G9bNBx892E/c+fiTbWfxRWq6tEERERG1cZImpQoKCqDX6+Hv729y3t/fHzk5Ode9fvfu3Th8+DAeeeSRRsckJSVBo9GIR3BwcIvjboncEq14O+Afy/c61UtKsVKKiIiIyLxG9/LHK3f2Fu/PW38EycdyJYyIiIiobZN8+V5LfPrpp+jTpw8GDhzY6JjExEQUFxeLR2ZmphUjvFptpZSLowKuKqXJY+09ncSti88VMClFREREZG4PDuqIx0Z0BgAYBCD+qwM4ls2dj4mIiKQgaVLKx8cHCoUCubmm31Dl5uYiICDgmteWl5dj9erVePjhh685TqVSwd3d3eSQUkGZsVLKx0111WNKhRwdvJwBGJNSbMBJREREZH4vxPbEbX0DAQCV1Xo89sU+FFVw52MiIiJrkzQp5ejoiKioKCQnJ4vnDAYDkpOTER0dfc1rv/32W2i1Wjz00EOWDtNsavQGlFTVAAA8nR0bHFPb7FxbY8DF4kqrxUZERETUVsjlMrx9bzj6BGkAABmFFXhqdRr0/EKQiIjIqiRfvpeQkIBly5Zh5cqVOHbsGGbNmoXy8nLExcUBAKZMmYLExMSrrvv0008xYcIEeHt7WzvkZiuurBZvezo7NDimfl+p9IIKi8dERERE1BapHRRYOjkKXi7GLwr/OJmPdzaflDgqIiKitkV5/SGWNXHiROTn52Pu3LnIyclBREQENm7cKDY/z8jIgFxumjs7ceIEduzYgd9++02KkJvtcr2y8MYqpYKvLN8DgKwiJqWIiIiILCXIwwkfTIrE5E93Q28Q8MHW0wgP9sCYUP/rX0xEREQtJnlSCgDi4+MRHx/f4GMpKSlXnevRowcEwf7Kqy9X1FVKeTSSlGrv6STevnCZy/eIiIiILGlIFx8kjuuJV385BgB4bu1BbHhyONp5OF3nSiIiImopyZfvtSWXy+tXSjW8fC+4XlIqs5CVUkRERESW9vCwThjb27jJTlFFNZ5enYYavUHiqIiIiFo/JqWsqKh+pZRLw5VSQR51y/dYKUVERERkeTKZDG/c3RdBV6qjdqcX4r3fT0scFRERUevHpJQVmfaUarhSyslRAR9XY8KKSSkiIiIi69A4O+C9ByKgkMsAAO//fgo7zxRIHBUREVHrxqSUFdXvKdVYo3MACPI0VkvlllZBW6O3eFxEREREBER19ELCmO4AAEEAnv3mIEqqqq9zFRERETUXk1JWVFSvUsqjkUopoK7ZuSAA2UVVFo+LiIiIiIxmjeiC6M7eAICLxVV45aejEkdERETUejEpZUX1l+95NdJTCuAOfERERERSkctleOvevnBVGTepXrvvAn47kiNxVERERK0Tk1JW1NTle+0965qdZ17mDnxERERE1tTe0xlzx4eK9//vh0O4VKaVMCIiIqLWiUkpK6pdvqd2kEPtoGh0XLBJpRSTUkRERETWdm9Ue8T08gcAFJTp8H8/HIIgCBJHRURE1LowKWVFtZVS16qSAkwrpbh8j4iIiMj6ZDIZkv7VR2y5sOlILtalZUkcFRERUevCpJSVCIIgVkp5XDcpxZ5SRERERFLzdVPhtQlh4v35Px1FAZfxERERmQ2TUlZSrtOjWm8s+fa8xs57AKB2UMDHVQWAy/eIiIiIpDSuTyBu7xsIACiqqMZ/f+ZufERERObCpJSVlFTWNTnXOF07KQXUVUvllmihrdFbLC4iIiIiurZ543uL87cf0y5i6/E8iSMiIiJqHZiUspIybY14u3aL4Wtp56EWb+eVsEyciIiISCq+bir857Ze4v05PxwymdsRERFR8zApZSX1Jy4uTUhKBbjX9ZXKLq6ySExERERE1DT3RLXH0K7eAICLxVVYuOmExBERERHZPyalrKSsqi4p5aa+flIqUFNXKZVdzGbnRERERFKSyWRIuqsv1A7G6fPK1HTsO39Z4qiIiIjsG5NSVlJ+o5VS9ZJSOayUIiIion/Ytm0bxo8fj3bt2kEmk2HdunVNvvbPP/+EUqlERESExeJrjTp4O+PfY3oAAATBuIyvRm+QOCoiIiL7xaSUlZTeYE8p00opJqWIiIjIVHl5OcLDw7FkyZIbuq6oqAhTpkzB6NGjLRRZ6xY3NAShge4AgOM5pfg89bzEEREREdkvJqWspPwGk1KslCIiIqJrGTduHF599VXcddddN3TdzJkzMWnSJERHR1sostZNqZDjvxPCxPuLNp9EXgnnakRERM3BpJSV3GhSys9NDZnMeDubEx0iIiIyg88++wxnz57FvHnzmjReq9WipKTE5CAgqqMn7h8QDMC4mc1rG45JHBEREZF9YlLKSkpvsKeUo1IOH1cVACCHjc6JiIiohU6dOoUXX3wRq1atglJ5/bkIACQlJUGj0YhHcHCwhaO0H8+P7QkPZwcAwI9pF7HzTIHEEREREdkfJqWspH6lVFN23wPq+krllWpRzSaaRERE1Ex6vR6TJk3C/Pnz0b179yZfl5iYiOLiYvHIzMy0YJT2xcvFES+M7Snen/vjEehqOF8jIiK6EUxKWUlZ1Y1VSgFAgLsxKSUIQH6p1iJxERERUetXWlqKvXv3Ij4+HkqlEkqlEq+88goOHjwIpVKJ33//vcHrVCoV3N3dTQ6qM7F/MCKCPQAAp/PKsPzPc9IGREREZGeYlLKSMq1evO2iUjTpGu7AR0RERObg7u6OQ4cOIS0tTTxmzpyJHj16IC0tDYMGDZI6RLskl8vw6oQwyK/0AX13yylcLGLbBSIioqZqWskOtViZtlq87aZyaNI1ARon8TZ34CMiIqL6ysrKcPr0afH+uXPnkJaWBi8vL3To0AGJiYnIysrC559/DrlcjrCwMJPr/fz8oFarrzpPNyYsSIOHBnfE56nnUVmtx4Jfj+O9ByKlDouIiMgusFLKSsqvVErJZYDaoWkfu2mlFL91IyIiojp79+5FZGQkIiONCZCEhARERkZi7ty5AIDs7GxkZGRIGWKb8e9besDLxREAsP7gRexNL5Q4IiIiIvvApJSV1DY6d1UpIZPJmnRNQL2kFCuliIiIqL6RI0dCEISrjhUrVgAAVqxYgZSUlEavf/nll5GWlmaVWFs7jZMDEsbUNZCf/9NRGAyChBERERHZByalrKS0XlKqqUwqpUqYlCIiIiKyVfcPCEbPADcAwKGsYny3/4LEEREREdk+JqWsRKyUUjc9KeXvzkopIiIiInugVMgx9/ZQ8f6bm06gTFtzjSuIiIiISSkr0BsEVOiMPaVcbqBSSu2gEPsTMClFREREZNuGdPVBbG9/AEB+qRZLtp6+zhVERERtG5NSVlCuq/uW7EaW7wFAwJVqqdySKujZm4CIiIjIps25NRSOCuMU+9Pt53D+UrnEEREREdkuJqWsoKyq+Ump2r5SNQYBl8q0Zo2LiIiIiMyrg7czHh7eCQCg0xvw+oZjEkdERERku5iUsoJybQsqpeo3O+cSPiIiIiKbN3tUV/i6qQAAm47kYufpAokjIiIisk2SJ6WWLFmCkJAQqNVqDBo0CLt3777m+KKiIsyePRuBgYFQqVTo3r07NmzYYKVom6d+k8sb6SkF/GMHPialiIiIiGyeq0qJ52J7iPdf+fko2zAQERE1QNKk1Jo1a5CQkIB58+Zh//79CA8PR2xsLPLy8hocr9PpMGbMGKSnp2Pt2rU4ceIEli1bhqCgICtHfmPKWlQp5STezi6uNFtMRERERGQ59/Rrjz5BGgDA8ZxSfLfvgsQRERER2R5Jk1KLFi3CjBkzEBcXh9DQUCxduhTOzs5Yvnx5g+OXL1+OwsJCrFu3DkOHDkVISAhGjBiB8PBwK0d+Y2p33gMAZ5Xihq6tbXQOAHml7ClFREREZA/kchnm3NZLvL/wtxOoqLf5DREREUmYlNLpdNi3bx9iYmLqgpHLERMTg9TU1AavWb9+PaKjozF79mz4+/sjLCwMr7/+OvR6fYPjbUVVdV18Tg43lpTyd1eJt3NLuHyPiIiIyF4M7uyNmF7+AIxfLv5v+zmJIyIiIrItkiWlCgoKoNfr4e/vb3Le398fOTk5DV5z9uxZrF27Fnq9Hhs2bMBLL72Et99+G6+++mqjr6PValFSUmJyWFtLklJ+9SulSlgpRURERGRPXhzXEwq5DACw9I8zyCvll4xERES1JG90fiMMBgP8/PzwySefICoqChMnTsScOXOwdOnSRq9JSkqCRqMRj+DgYCtGbFRZb/mek+ONJaXc1UqoHYx/TayUIiIiIrIvXf1c8cBA4/yzQqfH4i2nJI6IiIjIdkiWlPLx8YFCoUBubq7J+dzcXAQEBDR4TWBgILp37w6Foi6x06tXL+Tk5ECn0zV4TWJiIoqLi8UjMzPTfG+iiSqrDeJt9Q1WSslkMvhfqZZiUoqIiIjI/jwd013c7Gb17gycyi2VOCIiIiLbIFlSytHREVFRUUhOThbPGQwGJCcnIzo6usFrhg4ditOnT8NgqEvynDx5EoGBgXB0dGzwGpVKBXd3d5PD2ipbsHwPAPzdjEmpkqoak6orIiIiIrJ9Pq4qzBrZBQBgEICkX49LHBEREZFtkHT5XkJCApYtW4aVK1fi2LFjmDVrFsrLyxEXFwcAmDJlChITE8Xxs2bNQmFhIZ566imcPHkSv/zyC15//XXMnj1bqrfQJCY9pW5w+R4A+NVrds4+BERERET2Z/rQTuKuyr8fz8PO0wUSR0RERCQ9SZNSEydOxMKFCzF37lxEREQgLS0NGzduFJufZ2RkIDs7WxwfHByMTZs2Yc+ePejbty+efPJJPPXUU3jxxRelegtNYtJTqjmVUvWaneey2TkRERGR3XFyVODZ2B7i/dc2HIPBIEgYERERkfSUUgcQHx+P+Pj4Bh9LSUm56lx0dDR27dpl4ajMq/7yvRvtKQUA/vUqpdhXioiIiMg+3RUZhE93nMOx7BIcuViCHw9m4a7I9lKHRUREJBm72n3PXpkmpW78IzetlGJSioiIiMgeKeQyzLm1l3j/rY0nTNo8EBERtTVMSllBVQuX7/m51SWl8kq5fI+IiIjIXg3r5oMR3X0BABeLq/DZn+nSBkRERCQhJqWsoLKFjc65fI+IiIio9fi/W3tBLjPe/nDraRSW66QNiIiISCJMSlmByfI9ZXN23+PyPSIiIqLWokeAG+6NCgYAlGprsGTraYkjIiIikgaTUlZQu/ueSimHvPZrsRvgqlLCVWXsSZ/H3feIiIiI7N4zY7pDpTROxb9IPY/MwgqJIyIiIrI+JqWsoLaBZXOW7tXyu7KEj5VSRERERPYvQKPG9GGdAAA6vQHvbD4pcURERETWx6SUFdQu32tOk/Na/leanZfr9CjT1pglLiIiIiKSzswRXaBxcgAA/JCWhaMXSySOiIiIyLqYlLKC2uV7LUpKsdk5ERERUauicXJA/KiuAABBAN7cdFziiIiIiKyLSSkrqKo2AADULUpKsdk5ERERUWszOboj2mmM87yUE/lIPXNJ4oiIiIish0kpC9MbBOj0xqRUy3pK1SWl2OyciIiIqHVQOyjwzJju4v0FG49DEAQJIyIiIrIeJqUsrLbJOcDle0RERER0tX/1a4/u/q4AgIOZRdh4OEfiiIiIiKyDSSkLq6yXlDLf8j1WShERERG1Fgq5DC+M7Snef2vTCVRfqbQnIiJqzZiUsrDaJucAoHZo/sddu/seAOSWslKKiIiIqDW5uacfBoZ4AQDOFpTjm72ZEkdERERkeUxKWZi5lu/51Vu+l8fle0REREStikwmwwvj6qql3t1yChW6GgkjIiIisjwmpSys/vK9ljQ6VzsooHFyAMDle0REREStUVRHT8T29gcA5JVq8dmf6dIGREREZGFMSllY/eV7LamUAuqaneeWVHFXFiIiojZu27ZtGD9+PNq1aweZTIZ169Zdc/z333+PMWPGwNfXF+7u7oiOjsamTZusEyw12XOxPSCXGW8vTTmDwnKdtAERERFZEJNSFmauRudAXbNzbY0BJZUs5yYiImrLysvLER4ejiVLljRp/LZt2zBmzBhs2LAB+/btw6hRozB+/HgcOHDAwpHSjejq54b7+gcDAEq1NViy9bTEEREREVmOUuoAWrsqMy3fAwC/fzQ71zg7tOj5iIiIyH6NGzcO48aNa/L4xYsXm9x//fXX8eOPP+Knn35CZGSkmaOjlng6pjt+OJAFbY0BX6Sex7QhIQj2cpY6LCIiIrNjpZSFVZqp0TlQt3wPMC7hIyIiImoug8GA0tJSeHl5NTpGq9WipKTE5CDLC9CoMX1YJwCATm/AO5tPShwRERGRZTApZWGVOoN4u+VJqXqVUmx2TkRERC2wcOFClJWV4b777mt0TFJSEjQajXgEBwdbMcK2beaILuImNz+kZeFYNhOCRETU+jApZWEmPaVauHyPlVJERERkDl999RXmz5+Pb775Bn5+fo2OS0xMRHFxsXhkZmZaMcq2TePkgNmjugAABAF4c+NxiSMiIiIyPyalLKx+Tym1smUft1+9Sqk8JqWIiIioGVavXo1HHnkE33zzDWJiYq45VqVSwd3d3eQg65kSHYJ2GuP8b+uJfKSeuSRxRERERObFpJSFaS2w+x7A5XtERER0477++mvExcXh66+/xm233SZ1OHQdagcFnhnTXby/YONxCIIgYURERETmxaSUhWn1dT2lHFtYKeXrWm/5XikrpYiIiNqysrIypKWlIS0tDQBw7tw5pKWlISMjA4Bx6d2UKVPE8V999RWmTJmCt99+G4MGDUJOTg5ycnJQXFwsRfjURP/q1x7d/V0BAAczi7DxcI7EEREREZkPk1IWpq2uS0qpWpiUclTK4e3iCADIY6UUERFRm7Z3715ERkYiMjISAJCQkIDIyEjMnTsXAJCdnS0mqADgk08+QU1NDWbPno3AwEDxeOqppySJn5pGIZfhhbE9xftvbTqBmnpfehIREdkzpdQBtHbaGvNVSgHGvlKXynXIK62CwSBALpe1+DmJiIjI/owcOfKaS7lWrFhhcj8lJcWyAZHF3NzTDwNDvLA7vRBnC8rxzd4LmDSog9RhERERtRgrpSxMV1O/UqplPaWAuh34qvUCLlfoWvx8RERERGTbZDIZXhhXVy31zpaTqNDVSBgRERGReTApZWHamrpG5y1dvgcA/m5sdk5ERETU1kR19MQtof4AgPxSLT7dfk7iiIiIiFqOSSkLM62UMkNSyp3NzomIiIjaoufH9oTiSuuGj7edRUEZv6AkIiL7xqSUhWnNvHzPz72uUiqvhEkpIiIioraiq58rJg4IBgCUaWvwfvIpiSMiIiJqGSalLMxk+Z6DOSqluHyPiIiIqK16enQ3ODkYv+j88q8MpBeUSxwRERFR8zEpZWH1l+85Ksy8fI+VUkRERERtip+7GjOGdwIA1BgEvPXbCYkjIiIiaj4mpSysdvmeg0IG+ZUeAC3BSikiIiKitu3REV3g7eIIAPjl72wczCySNiAiIqJmsomk1JIlSxASEgK1Wo1BgwZh9+7djY5dsWIFZDKZyaFWqxsdL7XaSilzVEkBgLeLI2pzW3lsdE5ERETU5riqlHhydDfxftKvxyAIgoQRERERNY/kSak1a9YgISEB8+bNw/79+xEeHo7Y2Fjk5eU1eo27uzuys7PF4/z581aM+MbUVkqpHFre5BwAlAo5fFyNS/i4fI+IiIiobXpgYAeEeDsDAHadLUTKiXyJIyIiIrpxkielFi1ahBkzZiAuLg6hoaFYunQpnJ2dsXz58kavkclkCAgIEA9/f38rRnxjahudq5Tm+6hrl/Dll2qhN/BbMSIiIqK2xlEpx3OxPcX7C349znkhERHZHUmTUjqdDvv27UNMTIx4Ti6XIyYmBqmpqY1eV1ZWho4dOyI4OBh33nknjhw5Yo1wm0VcvmfWpJSxUsogAJfK2FeKiIiIqC26tU8AwoM9AAAnckvx3f4L0gZERER0gyRNShUUFECv119V6eTv74+cnJwGr+nRoweWL1+OH3/8EatWrYLBYMCQIUNw4ULDP4S1Wi1KSkpMDmsSl++ZMSnlx2bnRERERG2eTCZD4ri6aql3Np9EVbVewoiIiIhujOTL925UdHQ0pkyZgoiICIwYMQLff/89fH198fHHHzc4PikpCRqNRjyCg4OtGq/WEpVSbvWTUuwrRURERNRWDe7sjdE9/QAA2cVV+OzPdGkDIiIiugGSJqV8fHygUCiQm5trcj43NxcBAQFNeg4HBwdERkbi9OnTDT6emJiI4uJi8cjMzGxx3E1VozeIa/tVSvM0Ogfqlu8BQC534CMiIiJq014Y11PcnfnDlNO4XK6TNiAiIqImkjQp5ejoiKioKCQnJ4vnDAYDkpOTER0d3aTn0Ov1OHToEAIDAxt8XKVSwd3d3eSwFp3eUBeHBRqdA1y+R0RERNTWdfd3wz1R7QEApVU1WLK14S9riYiIbI3ky/cSEhKwbNkyrFy5EseOHcOsWbNQXl6OuLg4AMCUKVOQmJgojn/llVfw22+/4ezZs9i/fz8eeughnD9/Ho888ohUb6FRtU3OAfMu3/OrVymVx+V7RERERG3eM2O6i1+Cfp56HpmFFRJHREREdH1KqQOYOHEi8vPzMXfuXOTk5CAiIgIbN24Um59nZGRALq9L6Fy+fBkzZsxATk4OPD09ERUVhZ07dyI0NFSqt9AobY01KqWYlCIiIiJq6wI1Tnh4WCd8mHIGOr0Bb/92Aovvj5Q6LCIiomuSPCkFAPHx8YiPj2/wsZSUFJP777zzDt555x0rRNVy2ur6SSnz9ZTycnaEUi5DjUFADpfvERERERGAmSO74OvdGbhcUY11aRfxyPDOCAvSSB0WERFRoyRfvtea6fR1W/Kac/meXC4Tq6VYKUVEREREAOCudkD8zd3E+29sPC5hNERERNfHpJQFVVVbZvkeULcDX2G5Dtoa/XVGExEREVFb8NDgDmjv6QQA2H6qANtP5UscERERUeOYlLKg+rvvmbNSCgACNHV9pfK4hI+IiIiIYGwZ8VxsD/F+0obj0BsECSMiIiJqHJNSFmSpnlKAabPz7GIu4SMiIiIio/F926HPlV5SR7NL8P3+CxJHRERE1DAmpSyo/rI6cy/fC6iXlMphXykiIiIiukIulyHx1p7i/bc2nUCFrkbCiIiIiBrGpJQF6Wqss3wvl5VSRERERFTPkC4+GBPqDwDIK9Xi4z/OShwRERHR1ZiUsiBtjeUanbNSioiIiIiuJXFcTyjlMgDAx9vOIIdfZBIRkY1hUsqCdJZMSmmYlCIiIiKixnX2dcXk6I4AjLtCv7XphMQRERERmWJSyoJMK6Us1+icy/eIiIiIqCFPje4GjZMDAOC7/Rdw6EKxxBERERHVYVLKgkwanTuY96NWOyjg4WycYLBSioiIiIga4uHsiCdHdxPvv/rLUQiCIGFEREREdZiUsiCTRucK83/UtX2l8kq0MBg4uSAiIiKiq00e3BEh3s4AgL/OFeK3o7kSR0RERGTEpJQFmSzfM3OlFFC3hE+nN6CwQmf25yciIiIi++eolCPx1l7i/aQNx0y+PCUiIpIKk1IWZLJ8z8w9pQAgsH6zc/aVIiIiIqJG3BLqj0GdvAAA6Zcq8MWu8xJHRERExKSURZks3zPz7nvAP5qds68UERERETVCJpPhpdtDIZMZ77+XfApFrLQnIiKJMSllQaa771mgp1T9SikmpYiIiNqUbdu2Yfz48WjXrh1kMhnWrVt33WtSUlLQr18/qFQqdO3aFStWrLB4nGQ7woI0+FdkewBAcWU13k0+JXFERETU1jEpZUGWrpQKqF8pxeV7REREbUp5eTnCw8OxZMmSJo0/d+4cbrvtNowaNQppaWl4+umn8cgjj2DTpk0WjpRsyXOxPaC+0uv0i9TzOJtfJnFERETUlimlDqA1M62UMn9PqfrL91gpRURE1LaMGzcO48aNa/L4pUuXolOnTnj77bcBAL169cKOHTvwzjvvIDY21lJhko0J0Kjx2E1d8G7yKdQYBLy+4Tj+N7W/1GEREVEbxUopC6rf6NwilVImy/e0Zn9+IiIiaj1SU1MRExNjci42NhapqamNXqPValFSUmJykP17bERn+LmpAABbjuVi+6l8iSMiIqK2ikkpC9LVCOJtR4X5P2pPZwcx2cXle0RERHQtOTk58Pf3Nznn7++PkpISVFZWNnhNUlISNBqNeAQHB1sjVLIwZ0clXhjbU7w//6ejqNYbrnEFERGRZTApZUH1f7hbIiklk8nEvlLZxQ1PJomIiIiaKzExEcXFxeKRmZkpdUhkJndFBiEi2AMAcDqvDJ+nnpc2ICIiapOYlLKg+kkpB6XMIq9Rm5QqqapBpU5/ndFERETUVgUEBCA3N9fkXG5uLtzd3eHk5NTgNSqVCu7u7iYHtQ5yuQzz7+gt3l+85SQulbEdBBERWReTUhZkkpSyQKUUAPhr2OyciIiIri86OhrJyckm5zZv3ozo6GiJIiKphQd74N6o9gCA0qoaLPzthMQRERFRW8OklAXp9HU9pZRyS1VKqcTbOewrRURE1GaUlZUhLS0NaWlpAIBz584hLS0NGRkZAIxL76ZMmSKOnzlzJs6ePYvnn38ex48fx4cffohvvvkGzzzzjBThk414fmxPuKmMG3Kv3pOJQxeKJY6IiIjaEialLKi6xlgp5aiQQyazTFLK372uUiqXlVJERERtxt69exEZGYnIyEgAQEJCAiIjIzF37lwAQHZ2tpigAoBOnTrhl19+webNmxEeHo63334b//vf/xAbGytJ/GQbfN1UeHJ0NwCAIADzfzoCQRCucxUREZF5KKUOoDWrXb7noLBMQgoAAjV1PSC4fI+IiKjtGDly5DWTBytWrGjwmgMHDlgwKrJHU4eE4Os9GTibX4695y9j/cGLuDMiSOqwiIioDWCllAWJSSml5T7mAA2X7xERERFR8zkq5Zh7e6h4//UNx1CurZEwIiIiaiuYlLKg6is9pSzV5Bzg8j0iIiIiarmRPfwwuqcfACC3RIsPU05LHBEREbUFTEpZkE5f11PKUvzc6pJS2ayUIiIiIqJmeun2UHHeumzbOZy/VC5xRERE1NoxKWVB1ugp5aiUw8fVEQArpYiIiIio+UJ8XDB9WCcAxi9XX17PpudERGRZTEpZUO3ue5ZcvgfULeHLK9VCb+DEgYiIiIia54mbuyLgytxy64l8bDqSK3FERETUmjEpZUHVBsv3lAKAQI1x4qA3CLhUprXoaxERERFR6+WiUmLu+Lqm56/8dAQVOjY9JyIiy2BSykIEQbDK7nuAabPzHC7hIyIiIqIWGBcWgOHdfAAAF4ur8F4ym54TEZFlMCllIXqDgNol+I4W7CkFQCyxBoAcNjsnIiIiohaQyWR45c4wsen5/7afxancUomjIiKi1sgmklJLlixBSEgI1Go1Bg0ahN27dzfputWrV0Mmk2HChAmWDbAZqvV1vZ0svXwvQMMd+IiIiIjIfDr5uGDmiM4AgBqDgJd+PMym50REZHaSJ6XWrFmDhIQEzJs3D/v370d4eDhiY2ORl5d3zevS09Px7LPPYvjw4VaK9MborizdAyyflGrn4STevlhcadHXIiIiIqK24fFRXRHsZZxn7jpbiPUHL0ocERERtTaSJ6UWLVqEGTNmIC4uDqGhoVi6dCmcnZ2xfPnyRq/R6/V48MEHMX/+fHTu3NmK0TZdtURJqewiVkoRERERUcupHRSYf0dv8f5/fz6GkqpqCSMiIqLWRtKklE6nw759+xATEyOek8vliImJQWpqaqPXvfLKK/Dz88PDDz983dfQarUoKSkxOayhflLKUWnZnlKB9ZbvXSxipRQRERERmcfNPf0xJtQfAFBQpsWi305KHBEREbUmkialCgoKoNfr4e/vb3Le398fOTk5DV6zY8cOfPrpp1i2bFmTXiMpKQkajUY8goODWxx3U1TXWK+nlNpBAW8XRwBMShERERGRec0bHwq1g3E++3lqOg5nFUscERERtRaSL9+7EaWlpZg8eTKWLVsGHx+fJl2TmJiI4uJi8cjMzLRwlEbW7CkF1C3hyy3VoqbeaxMRERERtUR7T2c8cXM3AIBBAOasOwy9gU3PiYio5ZRSvriPjw8UCgVyc3NNzufm5iIgIOCq8WfOnEF6ejrGjx8vnjMYjAkYpVKJEydOoEuXLibXqFQqqFQqC0R/bdbsKQUA7TzUOJRVDL1BQF6p1qTPFBEREdmG6upq5OTkoKKiAr6+vvDy8pI6JKImmTG8M77ffwFn8stxMLMIn6emI25oJ6nDIiIiOydppZSjoyOioqKQnJwsnjMYDEhOTkZ0dPRV43v27IlDhw4hLS1NPO644w6MGjUKaWlpVlua1xQmPaUUlu0pBQCBmno78HEJHxERkc0oLS3FRx99hBEjRsDd3R0hISHo1asXfH190bFjR8yYMQN79uyROkyia3JUyrHg7r7i/bc2nUAW55xERNRCki/fS0hIwLJly7By5UocO3YMs2bNQnl5OeLi4gAAU6ZMQWJiIgBArVYjLCzM5PDw8ICbmxvCwsLg6Ogo5VsxYe1KqaB6lVEXi7kDHxERkS1YtGgRQkJC8NlnnyEmJgbr1q1DWloaTp48idTUVMybNw81NTW45ZZbMHbsWJw6dUrqkIkaNSDEC5MGdQAAVOj0mLvuMASBy/iIiKj5WrR8zxwl6BMnTkR+fj7mzp2LnJwcREREYOPGjWLz84yMDMjlkufObpiufqNzpfV6SgGslCIiIrIVe/bswbZt29C7d+8GHx84cCCmT5+OpUuX4rPPPsP27dvRrVs3K0dJ1HQvjO2JLUdzkVeqRfLxPGw4lIPb+gZKHRYREdmpG05KlZaWYtWqVVi9ejV2794NnU4HQRAgk8nQvn173HLLLXj00UcxYMCAJj9nfHw84uPjG3wsJSXlmteuWLHiBqK3Hil6StViUoqIiMg2fP311+LtvLw8+Pn5NThOpVJh5syZ1gqLqNk0Tg6Yf0dvzPpyPwBg3vojGNbVBxpnB4kjIyIie3RD2RKWoDedtXtKsVKKiIjItt1zzz3Q6/UNPlZTU2PlaIiab2xYAGJ6GVc1FJRpsWDjMYkjIiIie3VDSanaEvTdu3fjpZdeQmxsLPr06YOuXbuK5eefffYZcnJyMGHCBGzfvt1Scds8a1dK+bqq4HAl+XWxiD2liIiIbI2HhweefPLJq85funQJMTExEkRE1DwymQz/ndAbrirjoouvd2di19lLEkdFRET26IayJV9//XWjPRHqqy1Bnz59erMDs3c6fb2eUlZISsnlMgRojEv4LhazUoqIiMjWfP7559i8eTOWL18unjt27BgGDhwIFxcXCSMjunGBGic8P7aHeP//fjiEquqGKwGJiIga06xsSXV1NZRKJQ4fPmzueFqN6pp6lVJWaHQOGCcHAFBUUY0KHZcBEBER2RIPDw989913eO6557B7925s2rQJ0dHRmDBhAn766SepwyO6YQ8O6ojIDh4AgLP55fhw62lpAyIiIrvTrN33HBwc0KFDh0b7IpD1e0oBQJBJX6kqdPVztcrrEhERUcP+9a9/ISIiQjz69OmDDz74ALfeeiuqqqrw/vvvIy4uTuowiZpFIZdhwb/64vb3t6NaL+CjP85gXJ9A9Ap0lzo0IiKyE80u4ZkzZw7+7//+D4WFheaMp9Won5RSyq1TKcUd+IiIiGxLly5dsH37djzyyCMICQmBt7c3li1bBkEQMGnSJPTr1w/V1dVSh0nUbD0C3DBzRBcAQLVewHNrD5rMg4mIiK6lWZVSAPDBBx/g9OnTaNeuHTp27HhVL4T9+/e3ODh7ZtJTysrL9wAmpYiIiGzBW2+9Jd7OyspCWloa0tLS4O3tja1bt+LTTz+FUqlEz549cfDgQQkjJWq++Ju7YtORHJzMLcPhrBJ8/McZxN/cTeqwiIjIDjQ7KTVhwgQzhtH6SL58r5g78BEREdmSoKAgBAUF4bbbbhPPlZWVIS0tjQkpsmsqpQIL7w3HXR/uhN4g4N3kUxgTGoAeAW5Sh0ZERDau2UmpefPmmTOOVsek0bkVdt8DgHYerJQiIiKyJ66urhg2bBiGDRsmdShELdK3vQceu6kzPkw5g2q9gGe/PYjvHx9itXkwERHZpxv6KSEIwvUHEQDTSilr/TAOZE8pIiIim5KRkXFD47OysiwUCZHlPRXTDd2ubLRzKKsYn2w7K3FERERk624oW9K7d2+sXr0aOp3umuNOnTqFWbNmYcGCBS0Kzp6Z9JSyUlLKXe0AN5Wx+C2by/eIiIgkN2DAADz22GPYs2dPo2OKi4uxbNkyhIWF4bvvvrNidETmpVIq8Na94ZBf6VyxeMtJnMgplTYoIiKyaTe0fO/999/HCy+8gMcffxxjxoxB//790a5dO6jValy+fBlHjx7Fjh07cOTIEcTHx2PWrFmWitvmmfSUUlqnpxRgXMJ3IrcUWUWVEAQBMpn1XpuIiIhMHT16FK+99hrGjBkDtVqNqKioq+ZOR44cQb9+/fDmm2/i1ltvlTpkohaJCPbAozd1wdI/zoi78X0/awiUXMZHREQNuKGk1OjRo7F3717s2LEDa9aswZdffonz58+jsrISPj4+iIyMxJQpU/Dggw/C09PTUjHbBSmW7wHGJXwnckuhqzHgUrkOPq4qq702ERERmfL29saiRYvw2muv4ZdffsGOHTtM5k4PPvggYmNjERYWJnWoRGbzdEw3bD6agzP55fj7QjE+3nYWs0d1lTosIiKyQc1qdM6GnNcnVVKqfrPz7KIqJqWIiIhsgJOTE26++Wbcc889Zn3eJUuW4K233kJOTg7Cw8Px/vvvY+DAgY2OX7x4MT766CNkZGTAx8cH99xzD5KSkqBWqxu9huhGqR2Mu/Hd/dFOGATjMr4R3X0RFqSROjQiIrIxzd5975VXXrnm43Pnzm3uU7cKuhrr95QCgKB6Samsogr0ac8f/kRERLbAx8cHQUFBCA8PNzm6d+/erOX2a9asQUJCApYuXYpBgwZh8eLFiI2NxYkTJ+Dn53fV+K+++govvvgili9fjiFDhuDkyZOYNm0aZDIZFi1aZI63SCSK7OCJmSO6iLvxPbMmDT89MQxqB4XUoRERkQ1pdlLqhx9+MLlfXV2Nc+fOQalUokuXLm0+KWXSU8qqlVJ133ReuMwd+IiIiKQyc+ZMPPjggxg+fDgA4NChQ0hLS8PBgwexZ88efPLJJygsLIRarUZYWBj++uuvG3r+RYsWYcaMGYiLiwMALF26FL/88guWL1+OF1988arxO3fuxNChQzFp0iQAQEhICB544IEbfl2ipno6pjtSTuTjaHYJTuWV4c2NJzB3fKjUYRERkQ1pdlLqwIEDV50rKSnBtGnTcNddd7UoqNbAZPmeFRudB3s6i7eZlCIiIpJO//79MWXKFJw7dw6AcRfj3r1748EHHwQACIKAjRs34oknnsDo0aNv6Ll1Oh327duHxMRE8ZxcLkdMTAxSU1MbvGbIkCFYtWoVdu/ejYEDB+Ls2bPYsGEDJk+e3OjraLVaaLVa8X5JSckNxUltm6NSjsX3R+D293dAV2PA8j/PYXQvPwzt6iN1aEREZCPMWsLj7u6O+fPn46WXXjLn09olqXpKBXvVT0pVWO11iYiIyNRff/0lJqAaIpPJMG7cOKxatQo5OTk39NwFBQXQ6/Xw9/c3Oe/v79/oc02aNAmvvPIKhg0bBgcHB3Tp0gUjR47E//3f/zX6OklJSdBoNOIRHBx8Q3ESdfd3w4tje4r3n/32IIorqiWMiIiIbInZsyXFxcUoLi4299PanWq9ND2lfF1VcFQaXy+zkJVSREREUtm5cyceeeSR644bPHgwtm7davF4UlJS8Prrr+PDDz/E/v378f333+OXX37Bf//730avSUxMFOd2xcXFyMzMtHic1PpMGxKCoV29AQDZxVWYu/6wxBEREZGtaPbyvffee8/kviAIyM7OxhdffIFx48a1ODB7J1VPKblchvYeTjhbUI4LlysgCEKzmqcSERFRyyxYsABz5szBl19+CQBwdXVFnz59EB4ejr59+yI8PBw9e/bEnj17UFpaekPP7ePjA4VCgdzcXJPzubm5CAgIaPCal156CZMnTxYTZX369EF5eTkeffRRzJkzB3L51fMVlUoFlYo7+VLLyOUyLLw3HLHvbENJVQ1+TLuI0b38cUd4O6lDIyIiiTU7KfXOO++Y3JfL5fD19cXUqVNN+hu0VabL96ybFAryNCalynV6XK6ohpeLo1Vfn4iIiIDx48dj/Pjx4v21a9ciLS0NaWlpePfdd3HmzBnxy6NrVSs1xNHREVFRUUhOTsaECRMAAAaDAcnJyYiPj2/wmoqKiqsSTwqFcSc0QRAauoTIbAI1TvjvhDA8tToNAPCfHw6hf0dPtKu3czQREbU9zU5K1TbtpIbprizfk8kAhdy6San6faUyCyuYlCIiIrIBY8eOxdixY8X7FRUVOHfuHLy9vRutbrqWhIQETJ06Ff3798fAgQOxePFilJeXi7vxTZkyBUFBQUhKSgJgTJItWrQIkZGRGDRoEE6fPo2XXnoJ48ePF5NTRJZ0Z0QQthzLw08HL6KkqgZPr07DVzMGQWnFVQVERGRbmp2UomurrjFWSjko5FZfPvfPHfjCgz2s+vpERER0fc7Ozujdu3ezr584cSLy8/Mxd+5c5OTkICIiAhs3bhSbn2dkZJhURv3nP/+BTCbDf/7zH2RlZcHX1xfjx4/Ha6+91uL3QtRUr94Zhv3nLyOrqBK70wvx/u+n8cyY7lKHRUREEmFSykJql+9Zs59UrfaedWXQmdyBj4iIqNWKj49vdLleSkqKyX2lUol58+Zh3rx5VoiMqGEaZwe890AE7vt4F/QGAe//fgrRXbwxuLO31KEREZEEWCtrIbVJKWv3kwKuXr5HRERERGQrojp64ZmYbgAAgwA8vToNl8t1EkdFRERSYFLKQqqv9JSSYo18cL1KqQuXK63++kRERERE1zJrZFdEX6mOyimpwnNr/2bDfSKiNohJKQsRK6Ws3OQcALxcHOHkYGxYyuV7RERERGRrFHIZFt8fIW7Is+VYLj5PPS9xVEREZG1MSllIjUG6SimZTIZgL2O11IXLlTAY+K0TEREREdkWf3c1Ft7bV7z/2oZjOHqxRMKIiIjI2piUspCaK5VSSgkqpYC6Hfh0NQYUlGkliYGIiIiI6Fpu7umPuKEhAIzz1viv9qNMWyNtUEREZDVMSllIXaWUNEkp7sBHRERERPbgxXE9ERbkDgA4W1COF75jfykioraCSSkLqaltdC6X5iM23YGPzc6JiIiIyDaplAosmdQPbmolAOCXv7OxYme6tEEREZFVMCllIdWGK43OJauUqktKXWClFBERERHZsI7eLnj73nDx/mu/HMO+85cljIiIiKyBSSkLMBgE1FYcKyTqKWWyfI+VUkRERERk427pHYDHbuoMwNgKI/6r/Sgs10kcFRERWZJNJKWWLFmCkJAQqNVqDBo0CLt372507Pfff4/+/fvDw8MDLi4uiIiIwBdffGHFaK+vtkoKkGb3PeAfy/dYKUVEREREduDZ2B4YGOIFAMgursJTqw9Az52kiYhaLcmTUmvWrEFCQgLmzZuH/fv3Izw8HLGxscjLy2twvJeXF+bMmYPU1FT8/fffiIuLQ1xcHDZt2mTlyBtX208KkG75nsbJAR7ODgCA85eYlCIiIiIi2+egkOP9SZHwcVUBALafKsD7v5+SOCoiIrIUyZNSixYtwowZMxAXF4fQ0FAsXboUzs7OWL58eYPjR44cibvuugu9evVCly5d8NRTT6Fv377YsWOHlSNvXP2klFSNzgHj2nwAuFhcCW2NXrI4iIiIiIiayt9djfceiEBtF4x3k08h5UTDX1gTEZF9kzQppdPpsG/fPsTExIjn5HI5YmJikJqaet3rBUFAcnIyTpw4gZtuuqnBMVqtFiUlJSaHpdXUX74nUU8pAAjxNi7hEwT2lSIiIiIi+zGkiw/+fUsPAMa57BNfH8C5gnKJoyIiInOTNClVUFAAvV4Pf39/k/P+/v7Iyclp9Lri4mK4urrC0dERt912G95//32MGTOmwbFJSUnQaDTiERwcbNb30JCaeuvelRIt3wPqKqUA4Pwl/hAnIiIiIvsxa0QX3BJq/D2htKoGMz7fizJtjcRRERGROUm+fK853NzckJaWhj179uC1115DQkICUlJSGhybmJiI4uJi8cjMzLR4fNV66RudA3WVUgCQzr5SRERERGRH5HIZFk2MQDc/VwDA6bwyPLMmDQY2PiciajUkTUr5+PhAoVAgNzfX5Hxubi4CAgIavU4ul6Nr166IiIjAv//9b9xzzz1ISkpqcKxKpYK7u7vJYWkmjc4lXL7HSikiIiIismeuKiWWTekPd7USALD5aC7eY+NzIqJWQ9KklKOjI6KiopCcnCyeMxgMSE5ORnR0dJOfx2AwQKvVWiLEZqm/fE8hYaNzVkoRERERkb0L8XHBew9Eio3PF285hU1HGm/1QURE9kPy5XsJCQlYtmwZVq5ciWPHjmHWrFkoLy9HXFwcAGDKlClITEwUxyclJWHz5s04e/Ysjh07hrfffhtffPEFHnroIanewlXqNzp3kLCnlJeLI9xUxm+VWClFRERERPZqZA8/PD+2p3g/YU0aTuaWShgRERGZg1LqACZOnIj8/HzMnTsXOTk5iIiIwMaNG8Xm5xkZGZDXqzYqLy/H448/jgsXLsDJyQk9e/bEqlWrMHHiRKnewlXqL9+TstG5TCZDRx9nHM4qwYXLlajWG+AgYY8rIiIiIqLmeuymzjhysQQ/HbyIcp0eMz7fix8eHwovF0epQyMiomaSPCkFAPHx8YiPj2/wsX82MH/11Vfx6quvWiGq5jNpdC7h8j0A6OjlgsNZJdAbBGRdrkSIj8v1LyIiIiIisjEymQxv3t0XZ/LKcDS7BOcvVWDmF/vwxSMDoVIqpA6PiIiagWUzFqCv11NKKWGjcwDoaNJXikv4iIiIiMh+OTkqsGxqf/i4qgAAu9MLkfj9IQgCd+QjIrJHTEpZQLXJ8j1pP+IQkx342OyciIiIiOxbkIcTPp3aH2oH4zz7+/1ZWLL1tMRRERFRczApZQG20ugcYKUUEREREbU+4cEeeOe+CPH+wt9O4ue/L0oXEBERNQuTUhZg0uhc4p5S9XtIsVKKiIiIiFqLcX0C8UL9Hfm+OYj9GZcljIiIiG4Uk1IWUGOwjd33AMDPTSWWNrNSioiIiIhak5kjOuO+/u0BALoaAx79fC8yC/lFLBGRvWBSygJqTHbfkzYpJZPJxL5SmYUVJk3YiYiIiIjsmUwmw6sT+mBwZy8AQEGZDlM/243L5TqJIyMioqZgUsoCqg220+gcqOsrVa0XcLGoUuJoiIiIiIjMx1Epx9KHotDZ1/hF7Nn8ckxfuQeVOr3EkRER0fVInzFphepXSknd6BwAOvu6irfP5JdJGAkRERERkfl5ODtiZdxA+LqpAAAHMooQ/9V+k3k5ERHZHialLKB+TymFxMv3AKCLSVKKfaWIiIiIqPUJ9nLGyriBcFMpAQDJx/Mw54fDEAS2ryAislVMSllA/d33HCTefQ8AuvjW7cDHSikiIiIiaq1C27nj48lR4mqFNXsz8c6WUxJHRUREjZE+Y9IK1RjqNTq3teV7eUxKEREREVHrNaSrDxbdFyHefy/5FFbtOi9dQERE1CgmpSygWm9bjc41Tg7i+nou3yMiIiKi1m58eDu8dHuoeP+lHw/jp4MXJYyIiIgaIn3GpBXS16+UsoGeUkDdEr6CMi2KK6oljoaIiIjMYcmSJQgJCYFarcagQYOwe/fua44vKirC7NmzERgYCJVKhe7du2PDhg1WipbIuh4e1gmP3dQZACAIwDNr0pB8LFfiqIiIqD4mpSzApFLKZpJSdUv4TrOvFBERkd1bs2YNEhISMG/ePOzfvx/h4eGIjY1FXl5eg+N1Oh3GjBmD9PR0rF27FidOnMCyZcsQFBRk5ciJrOfFcT1x/4BgAMbNiGZ9uR87zxRIHBUREdViUsoCTBqd28DyPeCfO/AxKUVERGTvFi1ahBkzZiAuLg6hoaFYunQpnJ2dsXz58gbHL1++HIWFhVi3bh2GDh2KkJAQjBgxAuHh4VaOnMh6ZDIZXrurD27vGwgA0NUY8MjKvdifcVniyIiICGBSyiJsrdE5AHTxY1KKiIiotdDpdNi3bx9iYmLEc3K5HDExMUhNTW3wmvXr1yM6OhqzZ8+Gv78/wsLC8Prrr0Ov1zf6OlqtFiUlJSYHkb1RyGV4Z2IERvf0AwBU6PSYtnw3jmXz3zMRkdSYlLKAGkNdpZTCZpbvuYi3z+Sx2TkREZE9KygogF6vh7+/v8l5f39/5OTkNHjN2bNnsXbtWuj1emzYsAEvvfQS3n77bbz66quNvk5SUhI0Go14BAcHm/V9EFmLg0KOJQ/2Q3RnbwBASVUNJn/6F87yy1oiIkkxKWUBNfq6SilbWb7XTuMEtYMxFv7wJSIiansMBgP8/PzwySefICoqChMnTsScOXOwdOnSRq9JTExEcXGxeGRmZloxYiLzUjsosGxqf0QEewAACsp0eGDZLpwr4Be2RERSsY2MSStji43O5XIZOvsYl/CdL6yArsZwnSuIiIjIVvn4+EChUCA313QnsdzcXAQEBDR4TWBgILp37w6FQiGe69WrF3JycqDT6Rq8RqVSwd3d3eQgsmeuKiVWxg1Er0Djv+XcEi3u/ySViSkiIokwKWUB9XtK2UqlFFDXV0pvEJBRyB+8RERE9srR0RFRUVFITk4WzxkMBiQnJyM6OrrBa4YOHYrTp0/DUG+ecvLkSQQGBsLR0dHiMRPZCo2zA758ZBB6BrgBYGKKiEhKtpMxaUX0NthTCjDtK3WafaWIiIjsWkJCApYtW4aVK1fi2LFjmDVrFsrLyxEXFwcAmDJlChITE8Xxs2bNQmFhIZ566imcPHkSv/zyC15//XXMnj1bqrdAJBkvF0d8NWMwE1NERBJTSh1Aa1R/+Z6Djey+BwBdfLkDHxERUWsxceJE5OfnY+7cucjJyUFERAQ2btwoNj/PyMiAXF73/WNwcDA2bdqEZ555Bn379kVQUBCeeuopvPDCC1K9BSJJ1SamJi3bheM5pWJiavWj0ejk43L9JyAiohZjUsoC6jc6V8ptpxitq19dUupkbqmEkRAREZE5xMfHIz4+vsHHUlJSrjoXHR2NXbt2WTgqIvvRWGLqy0cGm8ydiYjIMmwnY9KKVNdbvqe0sUqp2sbrJ3KYlCIiIiIiamgp38SPU3HkYrHEkRERtX5MSlmA3mT3Pdv5iB2VcnS+0lfqTH4ZqvXcgY+IiIiIqDYxFXplV75L5To88Mku7Dt/WeLIiIhaN9vJmLQi9Xffs6VKKQDoEWD8QVutF3A2n40ciYiIiIgAY2Lq60cHo18HDwBASVUNJn/6F3aeLpA2MCKiVoxJKQswaXRuQ5VSAMSyZAA4nlMiYSRERERERLZF4+SALx4ehKFdvQEAFTo9pq3Yg+RjuRJHRkTUOtlWxqSVsOlKKf+6pBT7ShERERERmXJRKfHp1AGI6eUHANDVGPDYF/uw/uBFiSMjImp9mJSygJp6lVIKuY0lpQKYlCIiIiIiuha1gwIfPRSF8eHtAAA1BgFPfn0Ay3eckzgyIqLWhUkpC6ipt/ueg8K2PuL2nk5wcVQAAI4zKUVERERE1CAHhRyLJ0bggYEdxHOv/HwUSRuOwVBvvk9ERM1nWxmTVqLmyq52MpntVUrJZDJ0v1ItlVVUidKqaokjIiIiIiKyTQq5DK/fFYYnR3cTz3287Sz+/e1B6Gq4kzURUUsxKWUBtY3Oba3Jea36zc5P5rJaioiIiIioMTKZDAljuuO1u8JQ+33zDwey8PDKPSjT1kgbHBGRnbPNrImd018p57W1Kqlaps3OyySMhIiIiIjIPjw4qCM+eigKKqXxV6jtpwrwwCe7kFdaJXFkRET2yyaSUkuWLEFISAjUajUGDRqE3bt3Nzp22bJlGD58ODw9PeHp6YmYmJhrjpdC9ZXd92xt571aPQLcxdsnckokjISIiIiIyH7E9g7Al48MgsbJAQBwKKsYEz74E8eyOacmImoOyZNSa9asQUJCAubNm4f9+/cjPDwcsbGxyMvLa3B8SkoKHnjgAWzduhWpqakIDg7GLbfcgqysLCtH3rja3fdsrcl5rfrL99jsnIiIiIio6fqHeGHtzGgEeTgBAC4WV+Gej3Zi6/GGf38hIqLGSZ41WbRoEWbMmIG4uDiEhoZi6dKlcHZ2xvLlyxsc/+WXX+Lxxx9HREQEevbsif/9738wGAxITk62cuSNq210rrTR5XueLo7wc1MBMCalBIG7hxARERERNVU3fzf8MHsIwoM9AADlOj0eXrkHn/15jnNrIqIbIGlSSqfTYd++fYiJiRHPyeVyxMTEIDU1tUnPUVFRgerqanh5eTX4uFarRUlJiclhaTVXekrZalIKAHq3My7hK66sxoXLlRJHQ0RERERkX/zc1Fjz6GDc1icQAGAQgPk/HcXcH4+IX1ITEdG1SZqUKigogF6vh7+/v8l5f39/5OTkNOk5XnjhBbRr184ksVVfUlISNBqNeAQHB7c47usRk1I2unwPAPoEacTbh7KKJYyEiIiIiMg+qR0UeP+BSMSP6iqe+2LXecSt2IOiCp2EkRER2QfbzZo0wYIFC7B69Wr88MMPUKvVDY5JTExEcXGxeGRmZlo8rmq9bTc6B4Cwekmpw0xKERERERE1i1wuw7OxPbDw3nA4XJn/bz9VgDs++BPHuakQEdE1SZqU8vHxgUKhQG5ursn53NxcBAQEXPPahQsXYsGCBfjtt9/Qt2/fRsepVCq4u7ubHJYmNjqX227OL4yVUkREREREZnNPVHusengQvF0cAQAZhRW4a8lO/Pz3RYkjIyKyXZJmTRwdHREVFWXSpLy2aXl0dHSj17355pv473//i40bN6J///7WCPWG6K8s31PYcE+pQI1a/IF5OKuYDRmJiIiIiFpoUGdvrH9imNgqo7Jaj/ivDmDBr8fF3xGIiKiO5KU8CQkJWLZsGVauXIljx45h1qxZKC8vR1xcHABgypQpSExMFMe/8cYbeOmll7B8+XKEhIQgJycHOTk5KCsrk+otXKXaYFy+52DDy/dkMplYLXW5ohpZRWx2TkRERETUUkEeTvh2ZjT+1S9IPLf0jzPsM0VE1ADJk1ITJ07EwoULMXfuXERERCAtLQ0bN24Um59nZGQgOztbHP/RRx9Bp9PhnnvuQWBgoHgsXLhQqrdgQm8QUFt0ZMuNzgEgLKhuKePhLK53JyIiIiIyB7WDAm/fG45540PF1RPbTubjjg/+ZD9XIqJ6lFIHAADx8fGIj49v8LGUlBST++np6ZYPqAWq623/qrTh5XuA6Q58h7OKMTbs2n28iIiIiIioaWQyGeKGdkKvQHfM/nI/LpXrkFFYgX99uBMv3d4LDw3uCJnMtn9fICKyNNsu5bFD9deK2/LuewCbnRMRERERWdrgK32mwtsb5946vQEv/XgE8V8fQGlVtcTRERFJi0kpM6upl5RS2PDue4BxvbunswMANjsnIiIiIrIUY5+pIYgbGiKe++XvbIx/fweOXOSXw0TUdtl21sQO1a+UcrDx5Xv1m51fKtchp6RK4oiIiIiIiFonR6Uc88b3xtKHouCmNnZRSb9Ugbs+3Ikv/zrPL4iJqE1iUsrMagx1PaUUNp6UAkyX8B3M5Lc0RERERESWNDYsAL88MRx9a5fz1Rgw54fDmP3Vfu7OR0RtDpNSZmZPPaUAICLYQ7x9IOOydIEQEREREbURHbyd8e3MaEwbEiKe23AoB7GLt2HHqQLpAiMisjImpcysRm8/PaUAoF8HT/H2fialiIiIiIisQqVU4OU7emPpQ/2gcTL2ec0t0eKhT//Cqz8fRVW1XuIIiYgsz/azJnamfqWUHRRKwddNhQ5ezgCAgxeKoasxXOcKIiIiIiIyl7Fhgdj09E0Y1tVHPPe/HecwYcmfOJFTKmFkRESWx6SUmdnT7nu1ojoaq6V0NQYczS6ROBoiIiIiorYlQKPG59MH4j+39YKjwvg7xPGcUoz/YAf+t/0sDAY2QSei1sk+siZ2xFBv1wylHTQ6B4B+HeuW8O07zyV8RERERETWJpfL8Mjwzvgxfih6+LsBMH5p/OovxzDxk1ScKyiXOEIiIvNjUsrMTHpK2cP6PQBR9ftKMSlFRERERCSZXoHu+DF+KKYP7SSe25N+GWMXb8P/tp81aRdCRGTvmJQyM5Pd9+ykUqpHgBtcHBUA2OyciIjInixZsgQhISFQq9UYNGgQdu/e3aTrVq9eDZlMhgkTJlg2QCJqFrWDAnPHh+LrGYPF/q/aK1VT9y7diTP5ZRJHSERkHkxKmVmNoa5RuMJOklIKuQwRHTwAANnFVbhYVCltQERERHRda9asQUJCAubNm4f9+/cjPDwcsbGxyMvLu+Z16enpePbZZzF8+HArRUpEzRXdxRsbnx6OaUNCxHP7M4pw67vb8fEfZ1Cj5yZFRGTfmJQyM3uslAL+sYSP1VJEREQ2b9GiRZgxYwbi4uIQGhqKpUuXwtnZGcuXL2/0Gr1ejwcffBDz589H586drRgtETWXs6MSL9/RG988Fo0Q77qqqaRfj+POJX/i7wtF0gZIRNQCTEqZWf3d9+R2lJSKrNfsfG86k1JERES2TKfTYd++fYiJiRHPyeVyxMTEIDU1tdHrXnnlFfj5+eHhhx+2RphEZEYDO3nh16duwsPDOkF25deMIxdLMGHJn3h5/RGUVlVLGyARUTMwKWVm9lop1S/YU/zhtie9UNpgiIiI6JoKCgqg1+vh7+9vct7f3x85OTkNXrNjxw58+umnWLZsWZNfR6vVoqSkxOQgIuk4OSrw0u2hWDtzCHoGGHfoMwjAip3piFn0BzYezoYgsBE6EdkPJqXMrH5SSiG3n49X4+yAXgHuAICj2SUoqtBJHBERERGZS2lpKSZPnoxly5bBx8enydclJSVBo9GIR3BwsAWjJKKmiuroiZ+eGIbEcT2hdjD+zpFbosXMVfvxyMq9yCyskDhCIqKmsZ+siZ2w10opwNhIEQAEAfjrHKuliIiIbJWPjw8UCgVyc3NNzufm5iIgIOCq8WfOnEF6ejrGjx8PpVIJpVKJzz//HOvXr4dSqcSZM2cafJ3ExEQUFxeLR2ZmpkXeDxHdOAeFHI+N6ILNz4zAqB6+4vnk43mIWfQHFm0+iUqdXsIIiYiuj0kpM6sxqZSyr6TU4M7e4u1dZy9JGAkRERFdi6OjI6KiopCcnCyeMxgMSE5ORnR09FXje/bsiUOHDiEtLU087rjjDowaNQppaWmNVkCpVCq4u7ubHERkW4K9nLF82gB8+GA/+LmpABgbob+XfAqj307Bz39f5JI+IrJZSqkDaG30hrptWe2tUmpgJy/IZMZKqdQzTEoRERHZsoSEBEydOhX9+/fHwIEDsXjxYpSXlyMuLg4AMGXKFAQFBSEpKQlqtRphYWEm13t4eADAVeeJyP7IZDLc2icQw7v54L3kU/jsz3TUGARcLK5C/FcH8EWn83j5jt7oFcjEMhHZFlZKmZk9V0ppnBzQu53xB9XxnFIUlrOvFBERka2aOHEiFi5ciLlz5yIiIgJpaWnYuHGj2Pw8IyMD2dnZEkdJRNbkpnbAnNtCsfHpmzCie92Svr/OFeK297bjP+sO4TLn+ERkQ1gpZWZ6O05KAUB0Z28czjLurPPX2UsY1ydQ4oiIiIioMfHx8YiPj2/wsZSUlGteu2LFCvMHREQ2oaufK1bEDcDvx/Pwys9Hcf5SBQwCsGpXBtanXcTjo7pi2pAQqB0UUodKRG0cK6XMrEZvv43Ogbpm5wCw43SBhJEQEREREVFzyWQyjO7lj9+euQnPj+0BZ0djAqqkqgYLfj2OmxemYO2+CyZfqhMRWRuTUmamF+pXStnfxzuokzccFMZk2h8n89kUkYiIiIjIjqmUCjw+siu2PjsS9/VvD9mV780vFlfh2W8P4rb3tiPlRB7n/UQkCfvLmti4+t802GOllItKiQEhXgCAC5crca6gXOKIiIiIiIiopfzd1XjznnD8+tRwjOpR12/qeE4ppn22Bw99+hcOXSiWMEIiaouYlDIze250Xqt+U8Q/TuZLGAkREREREZlTzwB3fBY3EF/PGIy+7TXi+T9PX8L4D3bg0c/34lh2iYQRElFbwqSUmen1BvG2UmGfSambmJQiIiIiImrVort4Y93jQ/H+A5Ho4OUsnv/taC7Gvbsds7/cj1O5pRJGSERtAZNSZtYaKqV6BrjBz00FANh19hKqqvUSR0REREREROYml8swPrwdtiSMwCt39hZ/BwCAXw5l45bF2/DU6gM4m18mYZRE1JoxKWVm9XtKKWT2mZSSyWTiEr6qagP+OlcocURERERERGQpjko5pkSHYNvzo/DS7aHwcXUEAAgC8GPaRcQs+gMJ36ThdB6TU0RkXkxKmVlrqJQCgFE9/cTbm4/mSBgJERERERFZg9pBgYeHdcK250chcVxPeDo7AAAMAvD9/iyMeecPPP7lPhzOYkN0IjIPJqXMzFB/9z077SkFGPtKOSqN/zy2HOUWsUREREREbYWzoxKPjeiC7S/cjOdie0DjZExOCQKw4VAObn9/B6Yu343dXFFBRC3EpJSZmVZK2e/H66pSYmgXbwBATkkVDvHbECIiIiKiNsVVpcTsUV2x44VReHFcT/i41vWc+uNkPu77OBX3Lt2JrSf4JTYRNY/9Zk1sVP2eUko7Xr4HALf0DhBv/3YkV8JIiIiIiIhIKm5qB8wc0QU7XhiF/97ZG0EeTuJje9IvI+6zPYhdvA1r9mRwkyQiuiGSJ6WWLFmCkJAQqNVqDBo0CLt372507JEjR3D33XcjJCQEMpkMixcvtl6gTdRaekoBwOhefqjt1b75KJNSRERERERtmdpBgcnRIUh5biTevjccXXxdxMdO5pbhhe8OYeiC37F4y0kUlGkljJSI7IWkSak1a9YgISEB8+bNw/79+xEeHo7Y2Fjk5eU1OL6iogKdO3fGggULEBAQ0OAYqekNBvG2vSel/NzUiAz2AACcyC3FGW4FS0RERETU5jko5Lg7qj02PzMCSx/qh/4dPcXHLpXrsHjLKQxZ8DteWPs3TuaWShgpEdk6SZNSixYtwowZMxAXF4fQ0FAsXboUzs7OWL58eYPjBwwYgLfeegv3338/VCpVg2Ok1poqpQDg1j6B4u2fDl6UMBIiIiIiIrIlcrkMY8MCsXbWEPzw+BDc3jdQ/B1IV2PAmr2ZuOWdbZj86V/YcjTXpNUJEREgYVJKp9Nh3759iImJqQtGLkdMTAxSU1OlCqvFDK2opxQA3N63nbiE76eDF9nAkIiIiIiIrhLZwRMfTOqHP54biRnDO8FNpRQf236qAI98vhc3vbkVS7ae5tI+IhJJlpQqKCiAXq+Hv7+/yXl/f3/k5OSY7XW0Wi1KSkpMDktqbZVSARo1BoR4AQDO5JfjaLZlPz8iIiIiIrJf7T2dMee2UKT+32jMvT0UwV51TdGziirx1qYTiE5KxpNfH8Ce9EJ+6U3Uxkne6NzSkpKSoNFoxCM4ONiir2e6+17r+HjvCG8n3v7pYLaEkRARERERkT1wVSkxfVgnpDw7Csun9ceoHr7iCoxqvYD1By/i3qWpGPfudnyx6zxKqqqlDZiIJCFZ1sTHxwcKhQK5uaa7uuXm5pq1iXliYiKKi4vFIzMz02zP3ZDWVikFAOPCAsT3sj4ti2vBiYiIiIioSRRyGW7u6Y/P4gbij2dH4bERneHp7CA+fjynFC+tO4wBr27BM2vSkHrmkklLFCJq3SRLSjk6OiIqKgrJycniOYPBgOTkZERHR5vtdVQqFdzd3U0OS9K3sp5SAODtqsKI7r4AgIvFVfjzdIHEERERERERkb3p4O2MxHG9kJo4Gu9MDEe/Dh7iY9oaA344kIUHlu3CyIUp+OD3U8gurpQuWCKyCknXlyUkJGDZsmVYuXIljh07hlmzZqG8vBxxcXEAgClTpiAxMVEcr9PpkJaWhrS0NOh0OmRlZSEtLQ2nT5+W6i1cpTVWSgHAff3bi7fX7LVstRkREREREbVeagcF7opsj+8fH4pfnhyGqdEdoXGqq57KKKzAwt9OYuiC3zHts93YcCgbVdV6CSMmIktRXn+I5UycOBH5+fmYO3cucnJyEBERgY0bN4rNzzMyMiCv15fp4sWLiIyMFO8vXLgQCxcuxIgRI5CSkmLt8BukNxjE260pKXVzT394uzjiUrkOm4/k4nK5Dp4ujlKHRUREREREdqx3Ow3m36lB4q298NvRXHyzJxM7rqzMMAhAyol8pJzIh5taiVvDAjEhMgiDOnlB3op+1yJqy2RCG9vuoKSkBBqNBsXFxRZZyvfYF3ux6YixT9bu/xsNP3e12V9DKq/+fBT/23EOADBvfCjihnaSOCIiIiLLsfScwR7xMyEia8gsrMDafRewdt8FZBVdvYQvUKPGHRHtMCEiCL0C+d8iIlvU1DlD69gezoboW+nyPQCYOKBu58JVu86zASEREREREZldsJcznhnTHdufH4UvHh6If/ULgoujQnw8u7gKH/9xFuPe3Y6xi7fho5QzDSaviMj2Sbp8rzWqMWl03rpyft383TCwkxd2nyvEmfxy7DhdgJuuNEAnIiIiIiIyJ7lchuHdfDG8my8qJ+ix+VgufjyQhT9O5ou/dx3PKcXxjcfxxsbjiAj2wG19AjGuTwDaezpLHD0RNQWTUmZmUimlaF2VUgAQNyQEu88VAgBW7kxnUoqIiIiIiCzOyVGBO8Lb4Y7wdigs1+GXvy/ihwNZ2J9RJI5JyyxCWmYRXttwDOHBHritTwDGhQUi2IsJKiJbxaSUmdXo61dKtb6k1JhQf7TTqHGxuAq/n8hDekE5QnxcpA6LiIiIiIjaCC8XR0yODsHk6BBkXKrA+oNZ+PnvbBzPKRXHHMwswsHMIry+4Tj6ttfg1j6BGNs7gL+7ENmY1rW+zAbUr5SSy1pfUkqpkGNydAgAQBCAZdvPShsQERERERG1WR28nRF/czdsfPom/P7vEXgutgdC/9H8/O8LxVjw63GMXJiCmEV/4I2Nx7Hv/GX2yCWyAayUMrMag0G83RorpQDggYHB+OD3UyjX6fHt3gt4anS3VrXLIBERERER2Z/Ovq6YPaorZo/qinMF5dhwKBsbDmXjyMUScczpvDKczivDRyln4OPqiNE9/RET6o9hXX3gVK+ZOhFZB5NSZla7ek8mMzbma408nB3x0OCO+HjbWej0BizbfhZzbguVOiwiIiIiIiIAQCcfFzFBlV5Qjo1HcrDlaC72ZVyGcOV3toIyHdbszcSavZlQO8gxrKsvxoT6YUR3PwRo+KU7kTUwKWVm+iuVUq21SqrWw8M74bOd6dDVGPDlXxmYOaILvF1VUodFRERERERkIsTHBTNHdMHMEV1wqUyL34/nYfPRXGw/VYDKaj0AoKragC3HcrHlWC4AoGeAG0b08MWI7r7o39ELjkp2viGyBCalzKy20bmilSel/NzUeGBAMFamnkeFTo8lW89g7nhWSxERERERke3ydlXh3v7BuLd/MKqq9fjzdAG2HMvF5qN5KCjTiuOO55TieE4pPv7jLFwcFRjS1QcjrySp2ntyNz8ic2FSysxqG50r5a0/k/74qK5YszcTVdUGrNp1HnFDQ7jdKhERERER2QW1gwKje/ljdC9/vDZBQNqFIqScyMcfJ/Lwd1axuMyvXKfH5qO52HzUWEXVxdcFw7v5YmhXHwzq7AV3tYOE74LIvrX+zImV1SalWnulFAD4u6sxfWgnAIBOb8CizScljoiIiKhtWbJkCUJCQqBWqzFo0CDs3r270bHLli3D8OHD4enpCU9PT8TExFxzPBFRWyKXy9CvgycSxnTHj/HDsHdODBZPjMBdkUHwdnE0GXsmvxwrdqZjxud7ETH/N9y55E+8ufE4/jxdgKorywGJqGlYKWVmNW0oKQUAj43ogi//ykBxZTV+OJCFydEd0a+Dp9RhERERtXpr1qxBQkICli5dikGDBmHx4sWIjY3FiRMn4Ofnd9X4lJQUPPDAAxgyZAjUajXeeOMN3HLLLThy5AiCgoIkeAdERLbL21WFCZFBmBAZBINBwOGLxcYqqpP5OJBxGVd+7YNBAA5mFuFgZhE+TDkDR6Uc/Tt6YkgXbwzp6oO+QRooFawFIWqMTBBqixLbhpKSEmg0GhQXF8Pd3d3szz90we/IKqqEr5sKe+bEmP35bdHyHefwys9HAQBhQe74cfawNpOUIyKi1svSc4aWGjRoEAYMGIAPPvgAAGAwGBAcHIwnnngCL7744nWv1+v18PT0xAcffIApU6Y06TVt/TMhIrKG4spq7Dp7CalnLuHP0wU4lVfW6FhXlRL9OnpiYIgnBoR4ITzYA2oHhRWjJZJGU+cMrJQys7qeUm0nKTMluiO+2ZuJ4zmlOJxVgq/+Oo/J0SFSh0VERNRq6XQ67Nu3D4mJieI5uVyOmJgYpKamNuk5KioqUF1dDS8vL0uFSUTUKmmcHBDbOwCxvQMAAHklVdh55hJ2ninAn6cvIauoUhxbpq3BtpP52HYyHwDgqJAjPFiDASFeGNDJC1EdPdmTito0JqXMrK0t3wMApUKOV+4Mw30fGyfBC349jpE9/Nj0nIiIyEIKCgqg1+vh7+9vct7f3x/Hjx9v0nO88MILaNeuHWJiGq/s1mq10GrrdqMqKSlpXsBERK2Yn7taXOonCAIyCivw5+lL+PNMAXafK0R+ad1/R3V6A/akX8ae9MtAyhnIZUCvQHcMCPHCwE5e6NfBEwEatYTvhsi6mJQyM73BAKBtVUoBwMBOXpjYPxhr9maiXKfHv789iNUzBkPexj4HIiIie7BgwQKsXr0aKSkpUKsb/+UnKSkJ8+fPt2JkRET2TSaToaO3Czp6u2DSoA4QBAHnL1Vgd3oh9pwrxO70Qpy/VCGONwjAkYslOHKxBCt2pgMAAjVqRHbwQGSwJyI7eCAsSMMlf9RqMSllZm2xUqrWf27vhR2nC5BVVInd5wqx/M9zeGR4Z6nDIiIianV8fHygUCiQm5trcj43NxcBAQHXvHbhwoVYsGABtmzZgr59+15zbGJiIhISEsT7JSUlCA4Obn7gRERtjEwmQ4iPC0J8XHBff+N/P/NKquolqS7jeE4J6nd6zi6uQvahHGw4lAPAWPAQ2s4dkcEeiOxgTFR18HKGTNb2fuek1odJKTPTt+GklJvaAW/fF44Hlu2CIABvbjwhNvMjIiIi83F0dERUVBSSk5MxYcIEAMZG58nJyYiPj2/0ujfffBOvvfYaNm3ahP79+1/3dVQqFVQqlbnCJiIiGJf73d63HW7v2w6AsXH6vvOF2H++CPszLuNgZhHKdXpxfI1BwN8XivH3hWKsTD0PAPB2cURYkAZ922sQFqRBnyANAjVqJqrI7jApZWZ1lVJtc9vPwZ298fDQTvjfjnPQ6Q2YuWof1scPg68bJ7RERETmlJCQgKlTp6J///4YOHAgFi9ejPLycsTFxQEApkyZgqCgICQlJQEA3njjDcydOxdfffUVQkJCkJNj/Abe1dUVrq6ukr0PIqK2TuPkgJt7+uPmnsY+gXqDgFN5pTiQUYQDGZdxIKPoqh3+LpXr8MfJfPxxpYE6APi4OooJqj5BGvRpr0GAOxNVZNuYlDKztrj73j89P7Yn0jKLsPf8ZWQXV2H2l/ux6pFBcFS2zUQdERGRJUycOBH5+fmYO3cucnJyEBERgY0bN4rNzzMyMiCv9yXZRx99BJ1Oh3vuucfkeebNm4eXX37ZmqETEdE1KOQy9AxwR88AdzwwsAMAYzXV3xeKxERVWmYRLldUm1xXUKZDyol8pJxoOFEVGuiOXoHu6ODlzN6/ZDNkglB/9WrrV1JSAo1Gg+LiYri7u5v1uQVBQKfEDQCAiGAPrJs91KzPb0/ySqsw/v0dyC0x7jTxr8ggLLw3nP/xIyIiu2HJOYO94mdCRGQbBEFAVlElDl0oxqGsuqPoH4mqhrg4KtAjwA09rySpQgPd0CPAHa4q1qyQ+TR1zsB/dWZkqJfea8uVUgDg56bGx5P7476PU6GrMeD7A1lwVSsx/47eLB8lIiIiIiJqAZlMhvaezmjv6YxxfQIBGBNVFy5X4nBWMf7OKjb+eaEYxZWmiapynR77M4qwP6PI5HxHb2f0CjAmqnoGuqFXgDvaezqxsIAsikkpM6oxGMTbbbHR+T9FBHvg/Qci8fiX+6E3CPg89TycHBR4cVxPJqaIiIiIiIjMSCaTIdjLGcFeDSeqjuWU4lh2CY5ll+DC5cqrrj9/qQLnL1Vg45Ec8ZyTgwJd/VzRzc8V3fzd0M3PFd393ZisIrNhUsqM9PVKpZQK/h8UAGJ7B+Cte/oi4ZuDAICPt51FcWU1Xp0QBqWCPaaIiIiIiIgspaFEFQCUVFXjRL0k1dHsUpzIKUFVtcHk+spqvbg0sD61g/xKssoN3fyNf3b3d0V7T2cWaNANYVLKjGrqJaXkrAQS/atfe1To9Hjpx8MQBGD1nkwUlOnw7v0RcOG6ZSIiIiIiIqtyVztgQIgXBoR4ief0BgHnL5XjWLYxWXUytxSn8spw/lK5SasaAKiqNuBwVgkOZ5WYnFc7yBHi7YIuvq7o5OOCzr4uV/50hcbJwRpvjewMMwJmpNfXq5RidtjEQ4M7QuPkgIRv0lCtF7DlWC7uXPInlj7UD1393KQOj4iIiIiIqE1TyGXo7OuKzr6uuK1vXVVVVbUeZ/PLcSqvFKdyy3AytxSn88qQ3kiy6nhOKY7nlF71/D6ujsYElY8rOvm6oPOVpFUHLxfu1N6GMSllRvp6Gxkq5Pw/1T+ND28HLxdHzPxiH0q1NTidV4Y7PvgTL9/RG/dGtWefKSIiIiIiIhujdlAgtJ07QtuZ7qBWVa3HuYJyMUlVW1mVcanCZBVRrYIyHQrKdNiTftnkvFwGBHs5I8TbBR29ndHByxkd691WOygs+v5IWkxKmZFJTylWSjVoaFcfrH9iGGat2ofjOaWo0Onx/Nq/8fPf2Xj9rjC093SWOkQiIiIiIiK6DrWDAr0Cjbv11VetN+DC5UqczS/DuYJynMkvF2/nlWqveh6DUNdkvSF+bqorCSpjoqqjt7FHVkcvZ3i5OLK4wc4xKWVG9bPBCjY6b1QnHxf88PhQvPTjYazddwEAsO1kPmIW/YEZwzvjsRFd4MpeU0RERERERHbHQSFHJx9jL6l/Kq2qRnpBBc4WlOFsfjnOFtQlrCp0+gafL69Ui7xS7VUVVgDgqlKig5exoirI0wlBHk5o7+mEIE8ntPdwhruTkkkrG8ff/M2IPaWazslRgYX3hmNcWADm/HAYOSVVqKo24P3fT+Pr3RmYPqwTHhzUkc3wiIiIiIiIWgk3tQP6tNegT3uNyXlBEJBfpkXGlYqpjELjcf5SOTIKK1BQpmvw+cq0NTiaXYKj2SUNPu6mUorJqiDPKwkrD2fxtjcrrSTHpJQZ1Rjqts/kNphNM7qXPwZ08sLizafwxa50VOsFFJTp8ObGE/hw6xncPyAY9w/sgK5+rlKHSkRERERERBYgk8ng56aGn5sa/evtCFirXFtzJUlVgYzCcpPEVdblygZ7WAFAqbam0cbrgHG3wHYeV5JWHk4I0KgRqFEjUOOEQI0aARo13NQslLAkJqXMqH5PKQWzrU3mrnbA3PGhmDqkI97adAK/HMqGIBiz3v/bcQ7/23EOkR088K9+7TGmlz8CNGqpQyYiIiIiIiIrcVEpG+xfBQA1egOyi6uQVVSJrMuVuHC5EllFFcgqMt6+WFSJan3DSauqaoNxGWF+eaOv7apSismqAPcrf9ZLWgVq1NA4ObDiqpmYlDKj+tlZJXtK3bCO3i74YFI//LugHMu2n8XafRegqzFWnx3IKMKBjCK8tO4wwoLccXNPfwzv5oM+QRruxkBERERERNRGKRVyBHsZm583xGAwLg28cLniSsLqSuLqciUuXDYmr6qqDQ1eCxiLJU7nleF0XlmjY9QOcgRqnBDgroafuwp+bipj5Ze7Cr5uxvu+bmq4q9nj6p9sIim1ZMkSvPXWW8jJyUF4eDjef/99DBw4sNHx3377LV566SWkp6ejW7dueOONN3DrrbdaMeKGmVRKcfles3XyccHrd/VBwpju+DHtIr7dm2lSbnk4qwSHs0rwXvIpOCrkCAtyR1RHT/QKdEePADd08XVlooqIiIiIiIggl8vg766Gv7saUR2vflwQBBSW65BdXIWc4ipkl1Qhp7hSvJ9TXIWLxddOXFVVG3CuoBznChqvuAKMySvf2oTVlWSVn7taTFzVJrK8nB0hbyM5BcmTUmvWrEFCQgKWLl2KQYMGYfHixYiNjcWJEyfg5+d31fidO3figQceQFJSEm6//XZ89dVXmDBhAvbv34+wsDAJ3kGd+kkppVwuYSStg4+rCg8P64SHh3XCkYvF+O1ILn4/nodDWcXiGJ3egP0ZRdifUSSeU8hlCPF2RkdvFwR7OqG9pzOCvYwN7XzcHOHl4giVkkkrIiIiIiKitk4mk8HbVQVvVxXCgjQNjhEEASWVNcguqUtWGf80TV6Vamuu+VpV1QZkFlYis7DymuOUchl8XFXwcXOEt4sK3q6O8HFVwdvFEd6uKvjU3ne1/99vZYIgNLy40koGDRqEAQMG4IMPPgAAGAwGBAcH44knnsCLL7541fiJEyeivLwcP//8s3hu8ODBiIiIwNKlS6/7eiUlJdBoNCguLoa7+9XrUVti3/nLuPujnQCAh4d1wku3h5r1+ckot6QKf5zIx570Quw7fxlnr5ONboibWin+n9rD2RHuaiVcVEq4qpVwVSnhplbCxdF4TuUgh0ppPBwVCqgc5HBUyOv9qYBSLoNCLoNCJmszGW0iotbOknMGe8XPhIiIqHFl2hrklVQhr1RrPEqqkF97u7QKeSXG28WV1WZ93fq/34oJrCvJq7qklvG2xsnBKr+zNnXOIGmllE6nw759+5CYmCiek8vliImJQWpqaoPXpKamIiEhweRcbGws1q1bZ8lQm8S0UoqJCUvxd1fjvgHBuG9AMACgoEyLQ1nFOJFTihNXdlY4k18m9qNqSGlVDUqraq5bXtlcdQkqiImq+kkrpVwG+ZXHZZChdlmxDBDXGMuu/E/tvySZTFbvtvG62tv4x3V150zHyepfQC3GT9J8+M+SWmJYVx/8+5YeUodBREREbZyrSglXX1d09r327vFV1XoxWZVfWnXlT+2VpFVdUquwXGeSZ2jMjfx+K5cBns6O8HQxVlndGhaAaUM7Nfk9mpukSamCggLo9Xr4+/ubnPf398fx48cbvCYnJ6fB8Tk5OQ2O12q10Gq14v2SkpIWRt24GkNdEoQ9pazHx1WFUT38MKpH3XLP+s3sMguNDewuFlehsEyHgjLt/7d398FR3WUbx6+8bUIpCWAgLwUir8ECgRJKZvGxiQYbKjpgnZFqp0atrWBwiNUqrVpE/wjF1tY6jK3TaXGcUiwdganV1hRIHGOKJZAhIGYg5il1SojlaZMQSAK7v+ePuGs277tszuGc/X5mtmx2T8J99e4pN3d2T3Shs/fXjq7hX14ZKZ/fyCcj+cbkywMA+rhp4ji7SwAAABi1lKSEYS/OHuD3G7VdvqILnd1672KPLlzs6XO/O+Tj0f791m+kC509utDZI0laPG3wtyxaxfZrSo21iooKbd261ZLfq89OiqWUzUa6mF1A91Wf2i5d0cXuq723rqvB+53dV3Wx26fuqz71XPWr+6r/P7/2/9ivKz6//Mb0LqNM7/88fH7T5zHT+5gx8vt7F5g+f+97kwN778D9wBtq+z4nowHH6T/HGpn/fk7wHwMfD/l6uGb2vvEZAAAAgNvFx8dp0vjeVzXNGXjJ7QG6r/r0f529y6v3+iytej/ufez9S73Pv3+pR5d6fJo8PnnsgwzD1qVUenq6EhISdP78+ZDHz58/r8zMzEE/JzMzM6zjH3rooZC3+7W3t2v69OnXWPng/mduuv5322r5/fzl3ymSExM0NTVBozi/AQAAAAC4biUnJigrbZyy0kb3KvKuKz7bv9lu64+I83g8ys/P14EDB4KP+f1+HThwQF6vd9DP8Xq9IcdLUmVl5ZDHJycnKzU1NeQ21gLXDwIAAAAAALgepSQlaJzH3p/cZ/vb9x544AGVlpZq2bJlWr58uZ588kl1dnbqK1/5iiTpS1/6km666SZVVFRIkjZt2qTCwkI9/vjjWr16tXbv3q0jR47oV7/6lZ0xAAAAAAAAEAbbl1Lr1q3Tv//9bz3yyCNqaWnRkiVL9NprrwUvZn727FnFx//3BV0rVqzQrl279IMf/EAPP/yw5s6dq3379mnhwoV2RQAAAAAAAECY4oyx+x2E1mpvb1daWpra2toseSsfAABwJmaGgfh3AgAARmO0M4Ot15QCAAAAAABAbGIpBQAAAAAAAMuxlAIAAAAAAIDlWEoBAAAAAADAciylAAAAAAAAYDmWUgAAAAAAALAcSykAAAAAAABYjqUUAAAAAAAALMdSCgAAAAAAAJZLtLsAqxljJEnt7e02VwIAAK5ngVkhMDuAOQoAAIzOaOeomFtKdXR0SJKmT59ucyUAAMAJOjo6lJaWZncZ1wXmKAAAEI6R5qg4E2Pf/vP7/Xr33Xc1YcIExcXFRf3rt7e3a/r06XrnnXeUmpoa9a9/vYmlvLGUVSKvm8VSVom8bjbWWY0x6ujoUHZ2tuLjueKBxBwVbbGUN5aySuR1s1jKKpHXza6XOSrmXikVHx+vadOmjfnvk5qa6vr/iPuKpbyxlFUir5vFUlaJvG42lll5hVQo5qixEUt5YymrRF43i6WsEnndzO45im/7AQAAAAAAwHIspQAAAAAAAGA5llJRlpycrC1btig5OdnuUiwRS3ljKatEXjeLpawSed0slrLGiljraSzljaWsEnndLJaySuR1s+sla8xd6BwAAAAAAAD245VSAAAAAAAAsBxLKQAAAAAAAFiOpRQAAAAAAAAsx1Iqynbs2KEPf/jDSklJUUFBgf72t7/ZXVLU/ehHP1JcXFzIbf78+XaXFTV//vOf9ZnPfEbZ2dmKi4vTvn37Qp43xuiRRx5RVlaWxo0bp5UrV+r06dP2FBsFI+X98pe/PKDfq1atsqfYa1RRUaFbb71VEyZM0NSpU7V27Vo1NjaGHNPV1aWysjJ96EMf0o033qjPfe5zOn/+vE0VX5vR5C0qKhrQ3/Xr19tUceR++ctfKi8vT6mpqUpNTZXX69Uf//jH4PNu6qs0cl639HUw27ZtU1xcnMrLy4OPua2/bhfurLRnzx7Nnz9fKSkpWrRokf7whz9YVGl0hJN3586dA87dlJQUC6uN3EjzxGCqqqq0dOlSJScna86cOdq5c+eY1xkN4Watqqoa0Ne4uDi1tLRYU/A1Gs08MRgnnruRZHXyeTvSPDEYJ/Y1INy8Tu5tf4PNT4Oxo78spaLot7/9rR544AFt2bJFR48e1eLFi1VSUqLW1la7S4u6BQsW6Ny5c8HbX/7yF7tLiprOzk4tXrxYO3bsGPT57du366mnntLTTz+tw4cPa/z48SopKVFXV5fFlUbHSHkladWqVSH9fvHFFy2sMHqqq6tVVlamN998U5WVlbpy5Ypuv/12dXZ2Bo/51re+pVdeeUV79uxRdXW13n33Xd155502Vh250eSVpPvuuy+kv9u3b7ep4shNmzZN27ZtU11dnY4cOaJPfOITWrNmjU6ePCnJXX2VRs4ruaOv/b311lt65plnlJeXF/K42/rrZuHOSn/961/1hS98Qffee6+OHTumtWvXau3atTpx4oTFlUcmktkwNTU15Nx9++23Law4cqOZJ/pqbm7W6tWr9fGPf1z19fUqLy/X1772Nb3++utjXOm1CzdrQGNjY0hvp06dOkYVRtdo54m+nHruRpJVcu55O5p5oi+n9jUg3LySc3vb11DzU3+29dcgapYvX27KysqCH/t8PpOdnW0qKipsrCr6tmzZYhYvXmx3GZaQZPbu3Rv82O/3m8zMTPPTn/40+NgHH3xgkpOTzYsvvmhDhdHVP68xxpSWlpo1a9bYUs9Ya21tNZJMdXW1Maa3l0lJSWbPnj3BY06dOmUkmdraWrvKjJr+eY0xprCw0GzatMm+osbQpEmTzLPPPuv6vgYE8hrjzr52dHSYuXPnmsrKypB8sdJftwh3Vvr85z9vVq9eHfJYQUGB+frXvz6mdUZLuHmff/55k5aWZlF1Y2eweaK/7373u2bBggUhj61bt86UlJSMYWXRN5qshw4dMpLM+++/b0lNY22weaI/p5+7AaPJ6pbzNqDvPNGfW/ra13B53dDboeanwdjVX14pFSU9PT2qq6vTypUrg4/Fx8dr5cqVqq2ttbGysXH69GllZ2dr1qxZuvvuu3X27Fm7S7JEc3OzWlpaQvqclpamgoICV/Y5oKqqSlOnTlVubq42bNigCxcu2F1SVLS1tUmSJk+eLEmqq6vTlStXQvo7f/58zZgxwxX97Z834IUXXlB6eroWLlyohx56SJcuXbKjvKjx+XzavXu3Ojs75fV6Xd/X/nkD3NbXsrIyrV69OqSPkvvPWzeJZFaqra0d0POSkhJH9DbS2fDixYvKycnR9OnTR/wOvpM5ubeRWrJkibKysvTJT35SNTU1dpcTsaHmib7c0t/RZJXccd4ONU/05Za+SqPLKzm/t0PNT4Oxq7+JY/rVY8h7770nn8+njIyMkMczMjL0j3/8w6aqxkZBQYF27typ3NxcnTt3Tlu3btXHPvYxnThxQhMmTLC7vDEVeO//YH12ynUBwrVq1SrdeeedmjlzppqamvTwww/rjjvuUG1trRISEuwuL2J+v1/l5eX66Ec/qoULF0rq7a/H49HEiRNDjnVDfwfLK0lf/OIXlZOTo+zsbB0/flzf+9731NjYqN/97nc2VhuZhoYGeb1edXV16cYbb9TevXt18803q76+3pV9HSqv5K6+StLu3bt19OhRvfXWWwOec/N56zaRzEotLS2O/TM3kry5ubl67rnnlJeXp7a2Nj322GNasWKFTp48qWnTpllRtmWG6m17e7suX76scePG2VRZ9GVlZenpp5/WsmXL1N3drWeffVZFRUU6fPiwli5dand5YRlqnujPyeduwGizOv28HW6e6M8NfQ0nr9N7O9z8NBi7+stSCmG74447gvfz8vJUUFCgnJwcvfTSS7r33nttrAxj4a677greX7RokfLy8jR79mxVVVWpuLjYxsquTVlZmU6cOOGq66ENZ6i8999/f/D+okWLlJWVpeLiYjU1NWn27NlWl3lNcnNzVV9fr7a2Nr388ssqLS1VdXW13WWNmaHy3nzzza7q6zvvvKNNmzapsrLSsRcXBUbL6/WGfMd+xYoV+shHPqJnnnlGP/nJT2ysDNciNzdXubm5wY9XrFihpqYmPfHEE/rNb35jY2Xhi6X5abRZnX7eDjdPuFE4eZ3cWyfNT7x9L0rS09OVkJAw4Kf9nD9/XpmZmTZVZY2JEydq3rx5OnPmjN2ljLlAL2OxzwGzZs1Senq6o/u9ceNG/f73v9ehQ4dCvsuRmZmpnp4effDBByHHO72/Q+UdTEFBgSQ5sr8ej0dz5sxRfn6+KioqtHjxYv385z93bV+HyjsYJ/e1rq5Ora2tWrp0qRITE5WYmKjq6mo99dRTSkxMVEZGhiv760aRzEqZmZmO/TM3GrNhUlKSbrnlFkeeuyMZqrepqamuepXUUJYvX+64voYzTzj53JXCy9qf087bcOYJp/dVCi9vf07q7Ujzk8/nG/A5dvWXpVSUeDwe5efn68CBA8HH/H6/Dhw4MOx7VN3g4sWLampqUlZWlt2ljLmZM2cqMzMzpM/t7e06fPiw6/sc8K9//UsXLlxwZL+NMdq4caP27t2rgwcPaubMmSHP5+fnKykpKaS/jY2NOnv2rCP7O1LewdTX10uSI/vbn9/vV3d3t+v6OpRA3sE4ua/FxcVqaGhQfX198LZs2TLdfffdwfux0F83iGRW8nq9IcdLUmVlpSN6G43Z0OfzqaGhwZHn7kic3NtoqK+vd0xfI5knnNrfSLL25/Tzdrh5wql9Hc5weftzUm9Hmp8GuwyLbf0d08uox5jdu3eb5ORks3PnTvP3v//d3H///WbixImmpaXF7tKi6tvf/rapqqoyzc3NpqamxqxcudKkp6eb1tZWu0uLio6ODnPs2DFz7NgxI8n87Gc/M8eOHTNvv/22McaYbdu2mYkTJ5r9+/eb48ePmzVr1piZM2eay5cv21x5ZIbL29HRYb7zne+Y2tpa09zcbN544w2zdOlSM3fuXNPV1WV36WHbsGGDSUtLM1VVVebcuXPB26VLl4LHrF+/3syYMcMcPHjQHDlyxHi9XuP1em2sOnIj5T1z5oz58Y9/bI4cOWKam5vN/v37zaxZs8xtt91mc+Xh27x5s6murjbNzc3m+PHjZvPmzSYuLs786U9/Msa4q6/GDJ/XTX0dSv+fHuO2/rrZSLPSPffcYzZv3hw8vqamxiQmJprHHnvMnDp1ymzZssUkJSWZhoYGuyKEJdy8W7duNa+//rppamoydXV15q677jIpKSnm5MmTdkUYtZHmp82bN5t77rknePw///lPc8MNN5gHH3zQnDp1yuzYscMkJCSY1157za4IoxZu1ieeeMLs27fPnD592jQ0NJhNmzaZ+Ph488Ybb9gVISyjmZ/ccu5GktXJ5+1I85Nb+hoQbl4n93Yw/een66W/LKWi7Be/+IWZMWOG8Xg8Zvny5ebNN9+0u6SoW7duncnKyjIej8fcdNNNZt26debMmTN2lxU1gR/b2/9WWlpqjDHG7/ebH/7whyYjI8MkJyeb4uJi09jYaG/R12C4vJcuXTK33367mTJliklKSjI5OTnmvvvuc+yidbCckszzzz8fPOby5cvmG9/4hpk0aZK54YYbzGc/+1lz7tw5+4q+BiPlPXv2rLntttvM5MmTTXJyspkzZ4558MEHTVtbm72FR+CrX/2qycnJMR6Px0yZMsUUFxcHBwxj3NVXY4bP66a+DqX/UOW2/rrdcLNSYWFh8M/bgJdeesnMmzfPeDwes2DBAvPqq69aXPG1CSdveXl58NiMjAzzqU99yhw9etSGqsM30vxUWlpqCgsLB3zOkiVLjMfjMbNmzQr58/h6Fm7WRx991MyePdukpKSYyZMnm6KiInPw4EF7io/AaOYnt5y7kWR18nk70vzklr4GhJvXyb0dTP/56Xrpb5wxxkT/9VcAAAAAAADA0LimFAAAAAAAACzHUgoAAAAAAACWYykFAAAAAAAAy7GUAgAAAAAAgOVYSgEAAAAAAMByLKUAAAAAAABgOZZSAAAAAAAAsBxLKQAAAAAAAFiOpRQAAAAAIGxFRUUqLy+3uwwADsZSCgAAAAAAAJZjKQUA/fT09NhdAgAAAAC4HkspADGvqKhIGzduVHl5udLT01VSUmJ3SQAAAI7z6quvKi0tTS+88ILdpQBwiES7CwCA68Gvf/1rbdiwQTU1NXaXAgAA4Di7du3S+vXrtWvXLn3605+2uxwADsFSCgAkzZ07V9u3b7e7DAAAAMfZsWOHvv/97+uVV15RYWGh3eUAcBCWUgAgKT8/3+4SAAAAHOfll19Wa2urampqdOutt9pdDgCH4ZpSACBp/PjxdpcAAADgOLfccoumTJmi5557TsYYu8sB4DAspQAAAAAAEZk9e7YOHTqk/fv365vf/Kbd5QBwGN6+BwAAAACI2Lx583To0CEVFRUpMTFRTz75pN0lAXAIllIAAAAAgGuSm5urgwcPqqioSAkJCXr88cftLgmAA8QZ3vgLAAAAAAAAi3FNKQAAAAAAAFiOpRQAAAAAAAAsx1IKAAAAAAAAlmMpBQAAAAAAAMuxlAIAAAAAAIDlWEoBAAAAAADAciylAAAAAAAAYDmWUgAAAAAAALAcSykAAAAAAABYjqUUAAAAAAAALMdSCgAAAAAAAJZjKQUAAAAAAADL/T8CIN5pEMUzhwAAAABJRU5ErkJggg==", - "text/plain": [ - "
" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], + "outputs": [], "source": [ "def hydrogen_solver(l: int) -> lm.Solver:\n", " return lm.compile(\n", @@ -357,8 +340,10 @@ " )\n", "\n", "\n", - "def hydrogen_potential(solver: lm.Solver) -> jnp.ndarray:\n", - " return jnp.asarray((-1.0 / solver.mesh.radii)[None, None, :])\n", + "def hydrogen_potential(solver: lm.Solver) -> lm.Interaction:\n", + " # Build the Coulomb -1/r term as an Interaction (spectrum accepts only Interactions).\n", + " assert solver.local_potential is not None\n", + " return solver.local_potential(lambda r: -1.0 / r)\n", "\n", "\n", "solver = hydrogen_solver(0)\n", @@ -421,4 +406,4 @@ }, "nbformat": 4, "nbformat_minor": 5 -} +} \ No newline at end of file diff --git a/examples/hydrogen_demo.ipynb b/examples/hydrogen_demo.ipynb index b9e5818..2695ece 100644 --- a/examples/hydrogen_demo.ipynb +++ b/examples/hydrogen_demo.ipynb @@ -48,7 +48,9 @@ "def hydrogen_solver(angular_momentum: int) -> lm.Solver:\n", " return lm.compile(\n", " mesh=lm.MeshSpec(\"laguerre\", \"x\", n=30, scale=2.0),\n", - " channels=(lm.ChannelSpec(l=angular_momentum, threshold=0.0, mass_factor=HBAR2_2MU),),\n", + " channels=(\n", + " lm.ChannelSpec(l=angular_momentum, threshold=0.0, mass_factor=HBAR2_2MU),\n", + " ),\n", " operators=(\"T\", \"1/r\"),\n", " solvers=(\"spectrum\", \"wavefunction\"),\n", " grid=jnp.linspace(0.0, 40.0, 3000),\n", @@ -63,7 +65,10 @@ " prefactor = (\n", " 2.0\n", " / (n**2)\n", - " * math.sqrt(math.factorial(n - angular_momentum - 1) / math.factorial(n + angular_momentum))\n", + " * math.sqrt(\n", + " math.factorial(n - angular_momentum - 1)\n", + " / math.factorial(n + angular_momentum)\n", + " )\n", " )\n", " radial = (\n", " prefactor\n", @@ -74,7 +79,9 @@ " return radii * radial\n", "\n", "\n", - "def momentum_u_analytic(n: int, angular_momentum: int, momenta: np.ndarray) -> np.ndarray:\n", + "def momentum_u_analytic(\n", + " n: int, angular_momentum: int, momenta: np.ndarray\n", + ") -> np.ndarray:\n", " if n == 1 and angular_momentum == 0:\n", " return np.sqrt(2.0 / np.pi) * 2.0 / (1.0 + momenta**2)\n", " if n == 2 and angular_momentum == 0:\n", @@ -83,7 +90,9 @@ " if n == 2 and angular_momentum == 1:\n", " denominator = momenta**2 + 0.25\n", " return np.sqrt(2.0 / (6.0 * np.pi)) * momenta / (denominator**2)\n", - " raise ValueError(f\"No analytic momentum-space form for (n, l)=({n}, {angular_momentum}).\")\n", + " raise ValueError(\n", + " f\"No analytic momentum-space form for (n, l)=({n}, {angular_momentum}).\"\n", + " )\n", "\n", "\n", "def normalized_and_aligned(\n", @@ -121,8 +130,8 @@ "solver_s = hydrogen_solver(0)\n", "solver_p = hydrogen_solver(1)\n", "\n", - "spectrum_s = solver_s.spectrum(solver_s.potential(lambda r: -1.0 / r))\n", - "spectrum_p = solver_p.spectrum(solver_p.potential(lambda r: -1.0 / r))\n", + "spectrum_s = solver_s.spectrum(solver_s.local_potential(lambda r: -1.0 / r))\n", + "spectrum_p = solver_p.spectrum(solver_p.local_potential(lambda r: -1.0 / r))\n", "\n", "energy_rows = [\n", " (\"1s\", float(np.asarray(spectrum_s.eigenvalues)[0]) * HBAR2_2MU, -0.5),\n", @@ -134,7 +143,9 @@ "\n", "print(\"State numerical analytic abs error\")\n", "for label, numerical, analytic in energy_rows:\n", - " print(f\"{label:>3} {numerical: .10f} {analytic: .10f} {abs(numerical - analytic):.3e}\")" + " print(\n", + " f\"{label:>3} {numerical: .10f} {analytic: .10f} {abs(numerical - analytic):.3e}\"\n", + " )" ] }, { diff --git a/examples/yamaguchi_demo.ipynb b/examples/yamaguchi_demo.ipynb index 68a7ea0..592b623 100644 --- a/examples/yamaguchi_demo.ipynb +++ b/examples/yamaguchi_demo.ipynb @@ -62,7 +62,9 @@ "def yamaguchi_solver(angular_momentum: int, energies: jax.Array) -> lm.Solver:\n", " return lm.compile(\n", " mesh=lm.MeshSpec(\"legendre\", \"x\", n=20, scale=15.0),\n", - " channels=(lm.ChannelSpec(l=angular_momentum, threshold=0.0, mass_factor=HBAR2_2MU),),\n", + " channels=(\n", + " lm.ChannelSpec(l=angular_momentum, threshold=0.0, mass_factor=HBAR2_2MU),\n", + " ),\n", " operators=(\"T+L\",),\n", " solvers=(\"spectrum\", \"phases\"),\n", " energies=energies,\n", @@ -134,14 +136,16 @@ "\n", "for angular_momentum in partial_waves:\n", " solver = yamaguchi_solver(angular_momentum, energies)\n", - " potential = solver.potential(yamaguchi_kernel)\n", + " potential = solver.nonlocal_potential(yamaguchi_kernel)\n", " spectrum = solver.spectrum(potential)\n", " principal_branch = np.asarray(solver.phases(spectrum)[:, 0]) * (180.0 / np.pi)\n", " phase_curves[angular_momentum] = unwrap_phase_shift_deg(principal_branch)\n", "\n", "analytic_s_wave = yamaguchi_s_wave_analytic_phase_deg(np.asarray(energies))\n", "max_s_wave_error = np.max(np.abs(phase_curves[0] - analytic_s_wave))\n", - "print(f\"Maximum |δ_mesh - δ_analytic| for l=0 on this grid: {max_s_wave_error:.3e} degrees\")" + "print(\n", + " f\"Maximum |δ_mesh - δ_analytic| for l=0 on this grid: {max_s_wave_error:.3e} degrees\"\n", + ")" ] }, { @@ -165,8 +169,12 @@ "fig, axes = plt.subplots(1, 2, figsize=(13, 4.8))\n", "\n", "phase_curves[0] += 180\n", - "axes[0].plot(np.asarray(energies), analytic_s_wave, label=\"analytic s-wave\", linewidth=2.5)\n", - "axes[0].plot(np.asarray(energies), phase_curves[0], \"--\", label=\"spectral s-wave\", linewidth=2.0)\n", + "axes[0].plot(\n", + " np.asarray(energies), analytic_s_wave, label=\"analytic s-wave\", linewidth=2.5\n", + ")\n", + "axes[0].plot(\n", + " np.asarray(energies), phase_curves[0], \"--\", label=\"spectral s-wave\", linewidth=2.0\n", + ")\n", "axes[0].set_title(r\"Yamaguchi $\\ell=0$ phase shift\")\n", "axes[0].set_xlabel(\"Energy [MeV]\")\n", "axes[0].set_ylabel(\"Phase shift [deg]\")\n", diff --git a/src/lax/__init__.py b/src/lax/__init__.py index b9bc791..46be737 100644 --- a/src/lax/__init__.py +++ b/src/lax/__init__.py @@ -17,9 +17,8 @@ import lax.constants as constants import lax.models as models import lax.spectral as spectral -from lax.boundary._types import Solver from lax.compile import compile -from lax.types import ChannelSpec, Interaction, MeshSpec +from lax.types import ChannelSpec, Interaction, MeshSpec, Solver from lax.wavefunction import make_wavefunction_source __all__ = [ diff --git a/src/lax/boundary/__init__.py b/src/lax/boundary/__init__.py index c55c68e..638cd9b 100644 --- a/src/lax/boundary/__init__.py +++ b/src/lax/boundary/__init__.py @@ -1,13 +1,9 @@ -"""Boundary-value helpers and internal solver pytrees.""" +"""Boundary-value helpers (Coulomb and Whittaker functions at the channel radius).""" -from lax.boundary._types import BoundaryValues, Mesh, OperatorMatrices, Solver, TransformMatrices from lax.boundary.coulomb import compute_boundary_values +from lax.spectral.types import BoundaryValues __all__ = [ "BoundaryValues", - "Mesh", - "OperatorMatrices", - "Solver", - "TransformMatrices", "compute_boundary_values", ] diff --git a/src/lax/boundary/_types.py b/src/lax/boundary/_types.py deleted file mode 100644 index 1e029d4..0000000 --- a/src/lax/boundary/_types.py +++ /dev/null @@ -1,841 +0,0 @@ -"""Internal pytrees and callable protocols. See DESIGN.md §6.""" - -from __future__ import annotations - -from collections.abc import Callable -from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, Protocol - -import jax - -from lax.types import ChannelSpec, Interaction, MeshFamily, Method, Regularization - -if TYPE_CHECKING: - from lax.spectral.types import Spectrum - -type EnergyLike = float | jax.Array - - -class SpectrumKernel(Protocol): - """Callable that maps a potential to its spectral decomposition.""" - - def __call__( - self, - potential: jax.Array | Interaction, - ) -> Spectrum: - """Return the spectral decomposition for one assembled potential. - - Parameters - ---------- - potential - Assembled potential: either a raw ``jax.Array`` of shape - ``(N_c, N_c, N)`` for local or ``(N_c, N_c, N, N)`` for non-local, - or an :class:`~lax.Interaction` (energy-independent only). - - Returns - ------- - Spectrum - Eigendecomposition of the Bloch-augmented Hamiltonian. - """ - ... - - -class RMatrixObservable(Protocol): - """Callable that evaluates the spectral R-matrix at an arbitrary energy.""" - - def __call__(self, spectrum: Spectrum, energy: EnergyLike) -> jax.Array: - """Evaluate the R-matrix at one energy. - - Parameters - ---------- - spectrum - Spectral decomposition produced by ``solver.spectrum(V)``. - energy - Energy in MeV (scalar). - - Returns - ------- - jax.Array - R-matrix, shape ``(N_c, N_c)``. - """ - ... - - -class SpectrumObservable(Protocol): - """Callable that evaluates one observable on the compile-time energy grid.""" - - def __call__(self, spectrum: Spectrum) -> jax.Array: - """Evaluate an observable on the compile-time energy grid. - - Parameters - ---------- - spectrum - Spectral decomposition produced by ``solver.spectrum(V)``. - - Returns - ------- - jax.Array - Observable values on the compile-time grid, shape ``(N_E, ...)``. - """ - ... - - -class SpectrumGridObservable(Protocol): - """Callable for aligned-grid observables from a batched Spectrum.""" - - def __call__(self, spectra: Spectrum) -> jax.Array: - """Evaluate one observable per compile-time energy / Spectrum pair. - - Parameters - ---------- - spectra - Batched ``Spectrum`` produced by ``jax.vmap(solver.spectrum)(V_grid)``, - with a leading batch axis of size ``N_E``. - - Returns - ------- - jax.Array - Observable values, shape ``(N_E, ...)``. - """ - ... - - -class GreenFunctionObservable(Protocol): - """Callable that evaluates the Green's function at an arbitrary energy.""" - - def __call__(self, spectrum: Spectrum, energy: EnergyLike) -> jax.Array: - """Evaluate the Green's function at one energy. - - Parameters - ---------- - spectrum - Spectral decomposition; must have been produced with - ``'greens'`` in ``solvers=``. - energy - Energy in MeV (scalar). - - Returns - ------- - jax.Array - Resolvent ``(H - E/μ)⁻¹``, shape ``(M, M)``. - """ - ... - - -class WavefunctionObservable(Protocol): - """Callable that reconstructs the internal scattering wavefunction.""" - - def __call__(self, spectrum: Spectrum, energy: EnergyLike, source: jax.Array) -> jax.Array: - """Evaluate the internal wavefunction at one energy. - - Parameters - ---------- - spectrum - Spectral decomposition; must have been produced with - ``'wavefunction'`` in ``solvers=``. - energy - Energy in MeV (scalar). - source - Mesh-space driving term, shape ``(N_c · N,)``. Use - :func:`lax.make_wavefunction_source` to build it. - - Returns - ------- - jax.Array - Internal wavefunction coefficients in the mesh basis, shape ``(M,)``. - """ - ... - - -class EigenpairAccessor(Protocol): - """Callable for raw access to eigenvalues and eigenvectors. - - Raises ``RuntimeError`` if eigenvectors were not retained at compile time - (i.e. neither ``'greens'`` nor ``'wavefunction'`` was in ``solvers=``). - """ - - def __call__(self, spectrum: Spectrum) -> tuple[jax.Array, jax.Array]: - """Return the stored eigenvalues and eigenvectors. - - Parameters - ---------- - spectrum - Spectral decomposition with retained eigenvectors. - - Returns - ------- - tuple[jax.Array, jax.Array] - ``(eigenvalues, eigenvectors)`` — shapes ``(M,)`` and ``(M, M)``. - - Raises - ------ - RuntimeError - If eigenvectors were not retained at compile time. - """ - ... - - -class DirectRMatrixKernel(Protocol): - """Callable that computes the direct R-matrix via per-energy linear solves.""" - - def __call__(self, potential: jax.Array | Interaction) -> jax.Array: - """Evaluate the direct R-matrix on the compile-time energy grid. - - Parameters - ---------- - potential - Assembled potential array, shape ``(N_c, N_c, N)`` / ``(N_c, N_c, N, N)`` - for energy-independent V, or an :class:`~lax.Interaction` object - (dispatch is handled transparently at the Python level). - - Returns - ------- - jax.Array - R-matrix on the compile-time grid, shape ``(N_E, N_c, N_c)``. - """ - ... - - -class SMatrixDirectObservable(Protocol): - """Callable that computes the direct S-matrix via per-energy linear solves.""" - - def __call__(self, potential: jax.Array | Interaction) -> jax.Array: - """Evaluate the S-matrix on the compile-time energy grid. - - Parameters - ---------- - potential - Assembled potential or :class:`~lax.Interaction`. - - Returns - ------- - jax.Array - S-matrix on the compile-time grid, shape ``(N_E, N_c, N_c)``, complex. - """ - ... - - -class PhasesDirectObservable(Protocol): - """Callable that computes direct phase shifts via per-energy linear solves.""" - - def __call__(self, potential: jax.Array | Interaction) -> jax.Array: - """Evaluate phase shifts on the compile-time energy grid. - - Parameters - ---------- - potential - Assembled potential or :class:`~lax.Interaction`. - - Returns - ------- - jax.Array - Phase shifts, shape ``(N_E, N_c)``, in radians. - """ - ... - - -class WavefunctionDirectObservable(Protocol): - """Callable that reconstructs the wavefunction via a direct linear solve.""" - - def __call__( - self, - potential: jax.Array | Interaction, - source: jax.Array, - energy_index: int, - ) -> jax.Array: - """Solve ``C(E_i) ψ = source`` for the internal wavefunction. - - Parameters - ---------- - potential - Assembled potential or :class:`~lax.Interaction`. - source - Mesh-space driving term, shape ``(N_c·N,)``. - energy_index - Index into the compile-time energy grid (Python int). - - Returns - ------- - jax.Array - Internal wavefunction coefficient vector, shape ``(N_c·N,)``. - """ - ... - - -class DirectGridObservable(Protocol): - """Callable for aligned-grid observables from per-energy potential batches.""" - - def __call__(self, potentials: jax.Array) -> jax.Array: - """Evaluate one observable per compile-time energy / potential pair. - - Parameters - ---------- - potentials - Per-energy potentials, shape ``(N_E, N_c, N_c, N)``. - - Returns - ------- - jax.Array - Observable values, shape ``(N_E, ...)``. - """ - ... - - -class InterpolatorBuilder(Protocol): - """Callable that builds a Padé interpolator over the compile-time grid.""" - - def __call__( - self, - values: jax.Array, - order: tuple[int, int] | None = None, - ) -> Callable[[EnergyLike], jax.Array]: - """Build a Padé interpolator over the solver's compile-time energy grid. - - Parameters - ---------- - values - Observable samples, shape ``(N_E, ...)``. - order - Padé numerator/denominator degrees ``(p, q)`` with ``p + q + 1 == N_E``. - Defaults to the diagonal approximant. - - Returns - ------- - Callable[[EnergyLike], jax.Array] - JIT-compiled interpolant; call it at any energy to evaluate. - """ - ... - - -class GridVectorTransform(Protocol): - """Callable that projects mesh coefficients onto a fine radial grid.""" - - def __call__(self, values: jax.Array) -> jax.Array: - """Project mesh coefficients onto a radial grid. - - Parameters - ---------- - values - Mesh coefficient vector, shape ``(N,)``. - - Returns - ------- - jax.Array - Radial function values, shape ``(M_r,)``. - """ - ... - - -class FromGridVectorTransform(Protocol): - """Callable that projects fine-grid values back onto the mesh basis.""" - - def __call__( - self, - values: jax.Array | Callable[[jax.Array], jax.Array], - ) -> jax.Array: - """Project sampled radial-grid values or a callable profile onto the mesh basis. - - Parameters - ---------- - values - Either a ``(M_r,)`` array of grid samples or a callable - ``f(r_grid) → (M_r,)`` that is evaluated on the compile-time grid. - - Returns - ------- - jax.Array - Mesh coefficient vector, shape ``(N,)``. - """ - ... - - -class GridMatrixTransform(Protocol): - """Callable that projects a mesh-space kernel onto a fine radial grid.""" - - def __call__(self, values: jax.Array) -> jax.Array: - """Project a mesh-space kernel onto a radial grid. - - Parameters - ---------- - values - Mesh kernel matrix, shape ``(N, N)``. - - Returns - ------- - jax.Array - Kernel on the fine grid, shape ``(M_r, M_r)``. - """ - ... - - -class FourierTransform(Protocol): - """Callable that maps mesh coefficients or kernels to momentum space.""" - - def __call__(self, values: jax.Array, channel_index: int = 0) -> jax.Array: - """Project mesh coefficients or a kernel onto a momentum grid. - - Parameters - ---------- - values - Mesh vector ``(N,)`` or kernel ``(N, N)``. - channel_index - Which channel's angular momentum to use for the spherical Bessel - transform. - - Returns - ------- - jax.Array - Momentum-space array, shape ``(M_k,)`` or ``(M_k, M_k)``. - """ - ... - - -class DoubleFourierTransform(Protocol): - """Callable for the double Bessel transform of a mesh-space kernel.""" - - def __call__( - self, - values: jax.Array, - left_channel_index: int = 0, - right_channel_index: int | None = None, - ) -> jax.Array: - """Project a mesh-space kernel onto left/right momentum grids. - - Parameters - ---------- - values - Mesh kernel, shape ``(N, N)``. - left_channel_index - Channel angular momentum for the left (row) transform. - right_channel_index - Channel angular momentum for the right (column) transform. - Defaults to ``left_channel_index``. - - Returns - ------- - jax.Array - Double-transformed kernel, shape ``(M_k, M_k)``. - """ - ... - - -class Integrator(Protocol): - """Callable for norms and expectation values in the mesh basis.""" - - def __call__(self, values: jax.Array, operator: jax.Array | None = None) -> jax.Array: - """Integrate mesh coefficients with an optional operator insertion. - - Parameters - ---------- - values - Mesh coefficient vector, shape ``(N,)``. - operator - Optional ``(N, N)`` operator matrix. When ``None``, computes - the norm ``⟨ψ|ψ⟩``. - - Returns - ------- - jax.Array - Scalar expectation value ``⟨ψ|O|ψ⟩``. - """ - ... - - -@jax.tree_util.register_dataclass -@dataclass(frozen=True) -class PropagationMatrices: - """Precomputed subinterval-propagation matrices for Legendre-x meshes. - - Produced once by :func:`lax.propagate.build_legendre_x_propagation` and - stored inside the :class:`Mesh`. All matrices are in fm⁻² units. - - Attributes - ---------- - n_intervals - Number of subintervals (static). - basis_size_per_interval - Number of Legendre basis functions per subinterval (static). - interval_width - Width of each subinterval in fm (static). - local_nodes - Legendre quadrature nodes on ``(0, 1)``, shape ``(basis_size,)``. - local_weights - Legendre quadrature weights, shape ``(basis_size,)``. - kinetic - Per-interval kinetic matrices, shape ``(n_intervals, basis_size, basis_size)``. - blo0 - Bloch surface-overlap matrix for the left boundary of the first interval. - blo1 - Bloch surface-overlap matrix at the left boundary of subsequent intervals. - blo2 - Bloch surface-overlap matrix at the right boundary of interior intervals. - q1 - Left surface-projector vectors, shape ``(n_intervals, basis_size)``. - q2 - Right surface-projector vectors, shape ``(n_intervals, basis_size)``. - """ - - n_intervals: int = field(metadata={"static": True}) - basis_size_per_interval: int = field(metadata={"static": True}) - interval_width: float = field(metadata={"static": True}) - local_nodes: jax.Array - local_weights: jax.Array - kinetic: jax.Array - blo0: jax.Array - blo1: jax.Array - blo2: jax.Array - q1: jax.Array - q2: jax.Array - - -@jax.tree_util.register_dataclass -@dataclass(frozen=True) -class Mesh: - """Concrete mesh data cached inside a compiled solver. - - Produced by the mesh registry and embedded in the :class:`Solver` at - compile time. Static fields are baked into the JAX JIT cache key; - changing them requires recompilation. - - Attributes - ---------- - family - Mesh family, e.g. ``"legendre"`` or ``"laguerre"`` (static). - regularization - Endpoint regularization, e.g. ``"x"`` or ``"x(1-x)"`` (static). - n - Number of basis functions per channel (static). - scale - Physical scale: channel radius ``a`` in fm for finite-interval - meshes, or the Laguerre scaling factor ``h`` in fm (static). - n_intervals - Number of subintervals for propagated meshes; ``1`` otherwise (static). - basis_size_per_interval - Basis functions per subinterval; equals ``n`` when ``n_intervals == 1`` - (static). - nodes - Canonical mesh nodes on ``(0, 1)`` or ``(0, ∞)``, shape ``(n,)``. - weights - Gauss quadrature weights λ_i, shape ``(n,)``. - radii - Physical radial mesh points r_i = scale · x_i, shape ``(n,)``. - basis_at_boundary - Lagrange basis values φ_j(a) at the channel surface, shape ``(n,)``. - All zeros for semi-infinite (Laguerre) meshes. - propagation - Subinterval propagation matrices, or ``None`` for single-interval meshes. - """ - - family: MeshFamily = field(metadata={"static": True}) - regularization: Regularization = field(metadata={"static": True}) - n: int = field(metadata={"static": True}) - scale: float = field(metadata={"static": True}) - n_intervals: int = field(metadata={"static": True}) - basis_size_per_interval: int = field(metadata={"static": True}) - nodes: jax.Array - weights: jax.Array - radii: jax.Array - basis_at_boundary: jax.Array - propagation: PropagationMatrices | None = None - - -@jax.tree_util.register_dataclass -@dataclass(frozen=True) -class OperatorMatrices: - """Precomputed single-channel operator matrices in fm⁻² units. - - All populated fields are ``(N, N)`` symmetric real matrices in the - Lagrange-mesh basis. Unrequested operators are ``None``. - - Attributes - ---------- - T - Kinetic-energy matrix ``-d²/dr²`` (Laguerre meshes, no Bloch term). - TpL - Bloch-augmented kinetic matrix ``T + L(B=0)``. This is the standard - operator for R-matrix calculations on finite-interval (Legendre) meshes. - T_alpha - Hyperradial kinetic matrix for α-type coordinates (Laguerre meshes with - three-body regularization). - D - First-derivative matrix ``d/dr``. - inv_r - Diagonal ``1/r`` matrix. - inv_r2 - Diagonal ``1/r²`` matrix. - """ - - T: jax.Array | None = None - TpL: jax.Array | None = None - T_alpha: jax.Array | None = None - D: jax.Array | None = None - inv_r: jax.Array | None = None - inv_r2: jax.Array | None = None - - -@jax.tree_util.register_dataclass -@dataclass(frozen=True) -class BoundaryValues: - """Coulomb and Whittaker boundary values at the channel radius. - - Precomputed at compile time by ``mpmath`` for every ``(energy, channel)`` - pair. Open channels use Coulomb Hankel functions; closed channels use - Whittaker functions that decay exponentially into the barrier. - - Attributes - ---------- - H_plus - Outgoing Coulomb Hankel function ``H⁺ = G + iF`` at ``r = a``, - shape ``(N_E, N_c)``, complex. - H_minus - Incoming Coulomb Hankel function ``H⁻ = G - iF`` at ``r = a``, - shape ``(N_E, N_c)``, complex. - H_plus_p - ``ρ · d/dρ H⁺`` evaluated at ``ρ = ka``, - shape ``(N_E, N_c)``, complex. - H_minus_p - ``ρ · d/dρ H⁻`` evaluated at ``ρ = ka``, - shape ``(N_E, N_c)``, complex. - is_open - Boolean mask: ``True`` for open channels (``E > E_threshold``), - shape ``(N_E, N_c)``. - k - Channel wave numbers ``k_c(E)`` in fm⁻¹, shape ``(N_E, N_c)``. - """ - - H_plus: jax.Array - H_minus: jax.Array - H_plus_p: jax.Array - H_minus_p: jax.Array - is_open: jax.Array - k: jax.Array - - -@jax.tree_util.register_dataclass -@dataclass(frozen=True) -class TransformMatrices: - """Precomputed matrices for radial-grid and momentum-space transforms. - - All fields default to ``None``; only those corresponding to the ``grid`` - and ``momenta`` arguments supplied to :func:`lax.compile` are populated. - - Attributes - ---------- - B_grid - Basis-evaluation matrix ``B[k, j] = f_j(r_k)``, - shape ``(M_r, N)``. Used by ``to_grid_vector`` and ``to_grid_matrix``. - grid_r - Fine radial grid passed to :func:`lax.compile`, shape ``(M_r,)`` in fm. - 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. - momenta - Momentum grid passed to :func:`lax.compile`, shape ``(M_k,)`` in fm⁻¹. - Also accessible as ``solver.momenta``. - """ - - B_grid: jax.Array | None = None - grid_r: jax.Array | None = None - F_momentum: jax.Array | None = None - momenta: jax.Array | None = None - - -@dataclass(frozen=True) -class Solver: - """Compiled solver bundle produced by :func:`lax.compile`. - - Holds all compile-time caches (mesh, operators, boundary values, transform - matrices) alongside JIT-compiled, pickle-safe runtime callables. Call - ``print(solver)`` to see which observables were compiled. - - Attributes - ---------- - mesh - Compiled mesh data (nodes, weights, radii, boundary values). - operators - Precomputed single-channel operator matrices in fm⁻². - channels - Channel definitions baked into the solver at compile time. - energies - Compile-time energy grid in MeV, shape ``(N_E,)``. - boundary - Coulomb/Whittaker boundary values at ``r = a``, or ``None`` if no - energy grid was supplied. - transforms - Precomputed radial-grid and momentum-space transform matrices. - method - Linear-algebra backend: ``"eigh"``, ``"eig"``, or ``"linear_solve"``. - mass_factor_grid - Per-energy ℏ²/2μ values in MeV·fm², shape ``(N_E,)``, or ``None`` - when a constant mass factor is used. Stored here so the aligned-grid - observables can use the correct μ(E) at each energy point. - - **Spectral-path observables** (present when ``method`` is ``"eigh"``/``"eig"``): - - spectrum - ``(V) → Spectrum`` — one eigendecomposition per potential. - rmatrix - ``(spectrum, E) → R(E)`` — R-matrix at any scalar energy. - smatrix - ``(spectrum) → S`` — S-matrix on the compile-time energy grid. - phases - ``(spectrum) → δ`` — phase shifts ``(N_E, N_c)`` in radians. - greens - ``(spectrum, E) → G(E)`` — Green's function; requires ``'greens'`` - in ``solvers=``. - wavefunction - ``(spectrum, E, source) → ψ_int`` — internal wavefunction; requires - ``'wavefunction'`` in ``solvers=``. - eigh - ``(spectrum) → (ε, U)`` — raw eigenpairs; raises if eigenvectors - were not retained. - rmatrix_grid - ``(spectra) → R`` — aligned-grid R for energy-dependent workflows. - smatrix_grid - ``(spectra) → S`` — aligned-grid S. - phases_grid - ``(spectra) → δ`` — aligned-grid phases. - - **Direct-path observables** (present when ``"rmatrix_direct"`` in ``solvers=``): - - rmatrix_direct - ``(V) → R`` — per-energy linear-solve R-matrix on the compile-time grid. - - **Padé interpolation builders** (present whenever ``energies`` was supplied): - - interpolate_rmatrix, interpolate_smatrix, interpolate_phases - ``(samples, order=None) → callable`` — build a Padé interpolant over - the compile-time grid. - - **Transform helpers**: - - to_grid_vector - ``(c) → ψ(r)`` — mesh coefficients to fine radial grid. - from_grid_vector - ``(ψ_or_fn) → c`` — fine grid values back to mesh coefficients. - to_grid_matrix - ``(V) → V(r, r')`` — mesh kernel to fine radial grid. - fourier - ``(c, channel_index=0) → ũ(k)`` — momentum-space transform. - double_fourier_transform - ``(V, ...) → V(p, p')`` — double Bessel transform for kernels. - integrate - ``(c, operator=None) → ⟨ψ|O|ψ⟩`` — norms and expectation values. - """ - - mesh: Mesh - operators: OperatorMatrices - channels: tuple[ChannelSpec, ...] - energies: jax.Array - boundary: BoundaryValues | None - transforms: TransformMatrices - method: Method - mass_factor_grid: jax.Array | None = None - spectrum: SpectrumKernel | None = None - rmatrix: RMatrixObservable | None = None - smatrix: SpectrumObservable | None = None - phases: SpectrumObservable | None = None - greens: GreenFunctionObservable | None = None - wavefunction: WavefunctionObservable | None = None - eigh: EigenpairAccessor | None = None - rmatrix_grid: SpectrumGridObservable | None = None - smatrix_grid: SpectrumGridObservable | None = None - phases_grid: SpectrumGridObservable | None = None - rmatrix_direct: DirectRMatrixKernel | None = None - smatrix_direct: SMatrixDirectObservable | None = None - phases_direct: PhasesDirectObservable | None = None - wavefunction_direct: WavefunctionDirectObservable | None = None - interaction_from_block: Callable[..., Any] | None = None - interaction_from_array: Callable[..., Any] | None = None - interaction_from_funcs: Callable[..., Any] | None = None - potential: Callable[..., Any] | None = None - interpolate_rmatrix: InterpolatorBuilder | None = None - interpolate_smatrix: InterpolatorBuilder | None = None - interpolate_phases: InterpolatorBuilder | None = None - to_grid_vector: GridVectorTransform | None = None - from_grid_vector: FromGridVectorTransform | None = None - to_grid_matrix: GridMatrixTransform | None = None - fourier: FourierTransform | None = None - double_fourier_transform: DoubleFourierTransform | None = None - integrate: Integrator | None = None - - # ------------------------------------------------------------------ - # Convenience properties - - @property - def grid_r(self) -> jax.Array | None: - """Radial grid passed to :func:`lax.compile`, or ``None``.""" - return self.transforms.grid_r - - @property - def momenta(self) -> jax.Array | None: - """Momentum grid passed to :func:`lax.compile`, or ``None``.""" - return self.transforms.momenta - - # ------------------------------------------------------------------ - # Human-readable repr - - def __repr__(self) -> str: - _observable_names = ( - "spectrum", - "rmatrix", - "smatrix", - "phases", - "greens", - "wavefunction", - "eigh", - "rmatrix_grid", - "smatrix_grid", - "phases_grid", - "rmatrix_direct", - "smatrix_direct", - "phases_direct", - "wavefunction_direct", - ) - _transform_names = ( - "to_grid_vector", - "to_grid_matrix", - "fourier", - "integrate", - ) - 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] - n_e = len(self.energies) - return ( - f"Solver({self.mesh.family}/{self.mesh.regularization} " - f"n={self.mesh.n} scale={self.mesh.scale}fm, " - f"method={self.method}, {n_e} {'energy' if n_e == 1 else 'energies'})\n" - f" observables: {' '.join(live) or 'none'}\n" - f" transforms: {' '.join(transforms) or 'none'}" - ) - - -__all__ = [ - "BoundaryValues", - "DoubleFourierTransform", - "DirectGridObservable", - "DirectRMatrixKernel", - "EigenpairAccessor", - "EnergyLike", - "FourierTransform", - "FromGridVectorTransform", - "GreenFunctionObservable", - "GridMatrixTransform", - "GridVectorTransform", - "InterpolatorBuilder", - "Integrator", - "Mesh", - "OperatorMatrices", - "PropagationMatrices", - "RMatrixObservable", - "Solver", - "SpectrumGridObservable", - "SpectrumKernel", - "SpectrumObservable", - "TransformMatrices", - "PhasesDirectObservable", - "SMatrixDirectObservable", - "WavefunctionDirectObservable", - "WavefunctionObservable", -] diff --git a/src/lax/boundary/coulomb.py b/src/lax/boundary/coulomb.py index ea56fce..1121b91 100644 --- a/src/lax/boundary/coulomb.py +++ b/src/lax/boundary/coulomb.py @@ -10,7 +10,8 @@ import numpy as np import scipy.special as sc -from lax.boundary._types import BoundaryValues +from lax import constants +from lax.spectral.types import BoundaryValues from lax.types import ChannelSpec # reportUnknownMemberType is suppressed globally in pyproject.toml (JAX/mpmath stubs). @@ -158,17 +159,11 @@ def _fill_open_channel( rho = k * channel_radius eta = _sommerfeld(z1z2, k, mf) if z1z2 is not None else 0.0 - def coulombf_at_rho(rho_value: float) -> object: - return _mp_coulombf(channel.l, eta, rho_value) - - def coulombg_at_rho(rho_value: float) -> object: - return _mp_coulombg(channel.l, eta, rho_value) - if eta == 0.0: f_value, g_value, d_f, d_g = _neutral_open_channel_values(channel.l, rho) else: - f_value = complex(coulombf_at_rho(rho)) # pyright: ignore[reportArgumentType] -- mpmath values are complex-convertible numeric objects. - g_value = complex(coulombg_at_rho(rho)) # pyright: ignore[reportArgumentType] -- mpmath values are complex-convertible numeric objects. + f_value = complex(_mp_coulombf(channel.l, eta, rho)) # pyright: ignore[reportArgumentType] -- mpmath values are complex-convertible numeric objects. + g_value = complex(_mp_coulombg(channel.l, eta, rho)) # pyright: ignore[reportArgumentType] -- mpmath values are complex-convertible numeric objects. d_f, d_g = _open_channel_derivatives(channel.l, eta, rho, f_value, g_value) H_plus[energy_index, channel_index] = g_value + 1.0j * f_value @@ -218,7 +213,9 @@ def _sommerfeld(z1z2: tuple[int, int], k: float, mass_factor: float) -> float: """Return the Sommerfeld parameter in the fm^-2 convention.""" z1, z2 = z1z2 - return z1 * z2 * 1.44 / (2.0 * mass_factor * k) + # constants.E2 is read at call time so a test can override the Coulomb + # constant (e.g. to the rounded 1.44 used by published benchmark references). + return z1 * z2 * constants.E2 / (2.0 * mass_factor * k) def _mp_coulombf(l: int, eta: float, rho: float) -> object: diff --git a/src/lax/compile.py b/src/lax/compile.py index 4a752df..ae9d916 100644 --- a/src/lax/compile.py +++ b/src/lax/compile.py @@ -17,37 +17,12 @@ import numpy as np from lax.boundary import compute_boundary_values -from lax.boundary._types import ( - BoundaryValues, - DirectRMatrixKernel, - DoubleFourierTransform, - EigenpairAccessor, - FourierTransform, - FromGridVectorTransform, - GreenFunctionObservable, - GridMatrixTransform, - GridVectorTransform, - Integrator, - InterpolatorBuilder, - Mesh, - OperatorMatrices, - PhasesDirectObservable, - RMatrixObservable, - SMatrixDirectObservable, - Solver, - SpectrumGridObservable, - SpectrumKernel, - SpectrumObservable, - TransformMatrices, - WavefunctionDirectObservable, - WavefunctionObservable, -) from lax.meshes import build_mesh from lax.operators.interaction import ( make_interaction_from_array, make_interaction_from_block, make_interaction_from_funcs, - make_potential_builder, + make_potential_builders, ) from lax.solvers import ( bind_grid_observables, @@ -67,7 +42,34 @@ make_integration, make_to_grid, ) -from lax.types import ChannelSpec, MeshSpec, Method +from lax.types import ( + BoundaryValues, + ChannelSpec, + DirectRMatrixKernel, + DoubleFourierTransform, + EigenpairAccessor, + FourierTransform, + FromGridVectorTransform, + GreenFunctionObservable, + GridMatrixTransform, + GridVectorTransform, + Integrator, + InterpolatorBuilder, + Mesh, + MeshSpec, + Method, + OperatorMatrices, + PhasesDirectObservable, + RMatrixObservable, + SMatrixDirectObservable, + Solver, + SpectrumGridObservable, + SpectrumKernel, + SpectrumObservable, + TransformMatrices, + WavefunctionDirectObservable, + WavefunctionObservable, +) @dataclass(frozen=True) @@ -117,7 +119,8 @@ class _ObservableBundle: interaction_from_block: Callable[..., Any] | None interaction_from_array: Callable[..., Any] | None interaction_from_funcs: Callable[..., Any] | None - potential: Callable[..., Any] | None + local_potential: Callable[..., Any] | None + nonlocal_potential: Callable[..., Any] | None interpolate_rmatrix: InterpolatorBuilder | None interpolate_smatrix: InterpolatorBuilder | None interpolate_phases: InterpolatorBuilder | None @@ -158,7 +161,7 @@ def compile( solvers Runtime entry points to expose on the returned :class:`~lax.Solver`. The potential passed to ``solver.spectrum(V)`` or ``solver.rmatrix_direct(V)`` - must be an :class:`~lax.Interaction`. Build one with ``solver.potential(fn)`` + must be an :class:`~lax.Interaction`. Build one with ``solver.local_potential(fn)``/``solver.nonlocal_potential(fn)`` or the ``solver.interaction_from_{block,array,funcs}`` builders. energies Compile-time energy grid used for boundary-value-dependent observables and @@ -295,20 +298,30 @@ def _resolve_compile_request( channels_tuple = tuple(channels) operators_set = set(operators) solvers_set = frozenset(solvers) - needs_spectrum = bool( - solvers_set & {"spectrum", "rmatrix", "smatrix", "phases", "greens", "wavefunction"} + + selected_method = method or _default_method(V_is_complex) + if selected_method not in {"eigh", "eig", "linear_solve"}: + msg = f"Method {selected_method!r} is not implemented in the MVP compile() path." + raise ValueError(msg) + + # "wavefunction" is served by the spectral path under eigh/eig, but by the + # direct wavefunction_direct kernel under linear_solve (§14, Example 16.8). + wants_wavefunction = "wavefunction" in solvers_set + wavefunction_via_spectrum = wants_wavefunction and selected_method in {"eigh", "eig"} + wavefunction_via_direct = wants_wavefunction and selected_method == "linear_solve" + + needs_spectrum = ( + bool(solvers_set & {"spectrum", "rmatrix", "smatrix", "phases", "greens"}) + or wavefunction_via_spectrum ) needs_boundary = bool(solvers_set & {"smatrix", "phases", "rmatrix_direct"}) - if (needs_boundary or energy_dependent) and energies is None: + # wavefunction_direct indexes the compile-time energy grid, so it needs one. + if (needs_boundary or energy_dependent or wavefunction_via_direct) and energies is None: msg = "`energies` is required for continuum solvers or energy-dependent potentials." raise ValueError(msg) - if needs_spectrum or "rmatrix_direct" in solvers_set: + if needs_spectrum or "rmatrix_direct" in solvers_set or wavefunction_via_direct: operators_set.add("T+L") - selected_method = method or _default_method(V_is_complex) - if selected_method not in {"eigh", "eig", "linear_solve"}: - msg = f"Method {selected_method!r} is not implemented in the MVP compile() path." - raise ValueError(msg) if needs_spectrum and selected_method not in {"eigh", "eig"}: msg = f"Method {selected_method!r} is not implemented in the MVP spectrum path." raise ValueError(msg) @@ -320,7 +333,7 @@ def _resolve_compile_request( method=selected_method, needs_spectrum=needs_spectrum, needs_boundary=needs_boundary, - keep_eigenvectors=bool(solvers_set & {"greens", "wavefunction"}), + keep_eigenvectors=bool(solvers_set & {"greens"}) or wavefunction_via_spectrum, ) @@ -540,18 +553,19 @@ def _bind_solver_observables( boundary, mass_factor_grid, ) - from lax.solvers.linear_solve import ( - _DirectRMatrixKernel, # noqa: PLC0415 # pyright: ignore[reportPrivateUsage] - ) - - if isinstance(rmatrix_direct_fn, _DirectRMatrixKernel): - smatrix_direct_fn = make_smatrix_direct_observable(rmatrix_direct_fn, boundary) - phases_direct_fn = make_phases_direct_observable(smatrix_direct_fn) + smatrix_direct_fn = make_smatrix_direct_observable(rmatrix_direct_fn, boundary) + phases_direct_fn = make_phases_direct_observable(smatrix_direct_fn) + # wavefunction_direct is available whenever the direct path is active + # ("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" + if "rmatrix_direct" in request.solvers or wavefunction_via_direct: wavefunction_direct_fn = make_direct_wavefunction_kernel( mesh, operators, request.channels, energies, + mass_factor_grid, ) interpolate_rmatrix_fn: InterpolatorBuilder | None = None @@ -567,7 +581,9 @@ def _bind_solver_observables( interaction_from_block_fn = make_interaction_from_block(mesh, request.channels, energies) interaction_from_array_fn = make_interaction_from_array(mesh, request.channels, energies) interaction_from_funcs_fn = make_interaction_from_funcs(mesh, request.channels, energies) - potential_fn = make_potential_builder(mesh, request.channels, energies) + local_potential_fn, nonlocal_potential_fn = make_potential_builders( + mesh, request.channels, energies + ) return _ObservableBundle( spectrum=spectrum_fn, @@ -587,7 +603,8 @@ def _bind_solver_observables( interaction_from_block=interaction_from_block_fn, interaction_from_array=interaction_from_array_fn, interaction_from_funcs=interaction_from_funcs_fn, - potential=potential_fn, + local_potential=local_potential_fn, + nonlocal_potential=nonlocal_potential_fn, interpolate_rmatrix=interpolate_rmatrix_fn, interpolate_smatrix=interpolate_smatrix_fn, interpolate_phases=interpolate_phases_fn, @@ -638,7 +655,8 @@ def _assemble_solver( interaction_from_block=observables.interaction_from_block, interaction_from_array=observables.interaction_from_array, interaction_from_funcs=observables.interaction_from_funcs, - potential=observables.potential, + local_potential=observables.local_potential, + nonlocal_potential=observables.nonlocal_potential, interpolate_rmatrix=observables.interpolate_rmatrix, interpolate_smatrix=observables.interpolate_smatrix, interpolate_phases=observables.interpolate_phases, diff --git a/src/lax/constants.py b/src/lax/constants.py index 004af94..dc8c622 100644 --- a/src/lax/constants.py +++ b/src/lax/constants.py @@ -40,7 +40,16 @@ """Electron mass in MeV/c² (PDG).""" E2: float = ALPHA * HBARC -"""e² ≈ 1.44 MeV·fm (Coulomb coupling constant).""" +"""e² in MeV·fm (Coulomb coupling constant), the exact ``α·ℏc`` ≈ 1.43996. + +This is the physically correct value and the single source of truth for the +Coulomb constant — reference it (e.g. from ``boundary.coulomb`` and +``models.optical``) rather than hard-coding a literal. Some published +benchmark references were generated with the conventional rounded value +``1.44``; those specific tests override this constant locally (see the +``legacy_coulomb_constant`` fixture in ``tests/conftest.py``) rather than +forcing the rounded value on user applications. +""" def hbar2_over_2mu(m1_amu: float, m2_amu: float) -> float: diff --git a/src/lax/meshes/_basis_eval.py b/src/lax/meshes/_basis_eval.py index 0d01928..faf4cdc 100644 --- a/src/lax/meshes/_basis_eval.py +++ b/src/lax/meshes/_basis_eval.py @@ -10,7 +10,7 @@ import scipy.special as sc from numpy.polynomial import Legendre -from lax.boundary._types import Mesh +from lax.types import Mesh type BasisEvaluator = Callable[[Mesh, np.ndarray], np.ndarray] diff --git a/src/lax/meshes/_quadrature.py b/src/lax/meshes/_quadrature.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/lax/meshes/_registry.py b/src/lax/meshes/_registry.py index 42c8ed1..8fccec7 100644 --- a/src/lax/meshes/_registry.py +++ b/src/lax/meshes/_registry.py @@ -4,8 +4,7 @@ from collections.abc import Callable -from lax.boundary._types import Mesh, OperatorMatrices -from lax.types import MeshFamily, Regularization +from lax.types import Mesh, MeshFamily, OperatorMatrices, Regularization type Builder = Callable[..., tuple[Mesh, OperatorMatrices]] diff --git a/src/lax/meshes/_utils.py b/src/lax/meshes/_utils.py new file mode 100644 index 0000000..104c022 --- /dev/null +++ b/src/lax/meshes/_utils.py @@ -0,0 +1,24 @@ +"""Shared helpers for compile-time mesh construction.""" + +from __future__ import annotations + +import jax +import jax.numpy as jnp +import numpy as np + + +def to_jax_array(values: np.ndarray) -> jax.Array: + """Convert a NumPy array to a runtime JAX array with an explicit type.""" + + array: jax.Array = jnp.asarray(values) + return array + + +def diagonal_operator(values: np.ndarray) -> jax.Array: + """Construct a diagonal JAX operator from compile-time diagonal values.""" + + matrix: jax.Array = jnp.diag(to_jax_array(values)) + return matrix + + +__all__ = ["diagonal_operator", "to_jax_array"] diff --git a/src/lax/meshes/laguerre.py b/src/lax/meshes/laguerre.py index 94dfbe5..c1498ef 100644 --- a/src/lax/meshes/laguerre.py +++ b/src/lax/meshes/laguerre.py @@ -4,14 +4,13 @@ import math -import jax -import jax.numpy as jnp import numpy as np import scipy.special as sc -from lax.boundary._types import Mesh, OperatorMatrices - -from ._registry import register +from lax.meshes._registry import register +from lax.meshes._utils import diagonal_operator as _diagonal_operator +from lax.meshes._utils import to_jax_array as _to_jax_array +from lax.types import Mesh, OperatorMatrices @register("laguerre", "x") @@ -228,19 +227,4 @@ def _modified_laguerre_x2_kinetic( return matrix -def _to_jax_array(values: np.ndarray) -> jax.Array: - """Convert a NumPy array to a runtime JAX array with an explicit type.""" - - array: jax.Array = jnp.asarray(values) - return array - - -def _diagonal_operator(values: np.ndarray) -> jax.Array: - """Construct a diagonal JAX operator from compile-time diagonal values.""" - - diagonal = _to_jax_array(values) - matrix: jax.Array = jnp.diag(diagonal) - return matrix - - __all__ = ["build_laguerre_modified_x2", "build_laguerre_x"] diff --git a/src/lax/meshes/legendre.py b/src/lax/meshes/legendre.py index 49e98cd..3c7ffd4 100644 --- a/src/lax/meshes/legendre.py +++ b/src/lax/meshes/legendre.py @@ -2,14 +2,14 @@ from __future__ import annotations -import jax import jax.numpy as jnp import numpy as np -from lax.boundary._types import Mesh, OperatorMatrices +from lax.meshes._registry import register +from lax.meshes._utils import diagonal_operator as _diagonal_operator +from lax.meshes._utils import to_jax_array as _to_jax_array from lax.propagate import build_legendre_x_propagation - -from ._registry import register +from lax.types import Mesh, OperatorMatrices @register("legendre", "x") @@ -64,16 +64,8 @@ def build_legendre_x( n=n, scale=a, n_intervals=n_intervals, operators=operators ) - x_raw: np.ndarray - w_raw: np.ndarray - x_raw, w_raw = np.polynomial.legendre.leggauss(n) - nodes: np.ndarray = 0.5 * (x_raw + 1.0) - weights: np.ndarray = 0.5 * w_raw - radii: np.ndarray = a * nodes - - parity: np.ndarray = np.where(np.arange(n) % 2 == 0, 1.0, -1.0) - boundary_sign: np.ndarray = -parity if n % 2 == 1 else parity - basis_at_boundary: np.ndarray = boundary_sign / np.sqrt(a * nodes * (1.0 - nodes)) + nodes, weights, radii = _shifted_legendre_quadrature(n, a) + basis_at_boundary = _legendre_boundary_values(nodes, a) TpL: np.ndarray = _legendre_x_t_plus_l(nodes, a) D: np.ndarray = _legendre_x_derivative(nodes, a) @@ -463,21 +455,6 @@ def _legendre_x_three_halves_t_plus_l(nodes: np.ndarray, scale: float) -> np.nda return matrix -def _to_jax_array(values: np.ndarray) -> jax.Array: - """Convert a NumPy array to a runtime JAX array with an explicit type.""" - - array: jax.Array = jnp.asarray(values) - return array - - -def _diagonal_operator(values: np.ndarray) -> jax.Array: - """Construct a diagonal JAX operator from compile-time diagonal values.""" - - diagonal = _to_jax_array(values) - matrix: jax.Array = jnp.diag(diagonal) - return matrix - - __all__ = [ "build_legendre_x", "build_legendre_x_one_minus_x", diff --git a/src/lax/models/optical.py b/src/lax/models/optical.py index d8482c4..23a0851 100644 --- a/src/lax/models/optical.py +++ b/src/lax/models/optical.py @@ -15,9 +15,9 @@ import jax.numpy as jnp import numpy as np +from lax import constants from lax._angular import wigner_3j, wigner_6j -from lax.boundary._types import Solver -from lax.types import ChannelSpec, Interaction +from lax.types import ChannelSpec, Interaction, Solver @dataclass(frozen=True) @@ -333,7 +333,9 @@ def uniform_sphere_coulomb_potential( Coulomb potential in MeV. """ - prefactor = projectile_charge * target_charge * 1.44 + # constants.E2 is read at call time so tests can override the Coulomb + # constant (e.g. to the rounded 1.44 used by published benchmark references). + prefactor = projectile_charge * target_charge * constants.E2 inside = prefactor * (3.0 - (radii / radius) ** 2) / (2.0 * radius) outside = prefactor / radii return jnp.where(radii <= radius, inside, outside) diff --git a/src/lax/models/reid.py b/src/lax/models/reid.py index e28ae2e..8e19081 100644 --- a/src/lax/models/reid.py +++ b/src/lax/models/reid.py @@ -8,9 +8,8 @@ import jax.numpy as jnp import numpy as np -from lax.boundary._types import Solver from lax.constants import hbar2_over_2mu -from lax.types import ChannelSpec, Interaction +from lax.types import ChannelSpec, Interaction, Solver NN_MASS_FACTOR: Final[float] = hbar2_over_2mu(1.008665, 1.008665) diff --git a/src/lax/operators/interaction.py b/src/lax/operators/interaction.py index 7ea429b..7a3420f 100644 --- a/src/lax/operators/interaction.py +++ b/src/lax/operators/interaction.py @@ -2,17 +2,22 @@ from __future__ import annotations -import inspect from collections.abc import Sequence from dataclasses import dataclass -from typing import Any +from typing import Any, Literal import jax +import jax.core import jax.numpy as jnp import numpy as np -from lax.boundary._types import Mesh -from lax.types import ChannelSpec, Interaction +from lax.types import ChannelSpec, Interaction, Mesh + + +def _is_concrete(value: jax.Array) -> bool: + """Return False when ``value`` is a JAX tracer (inside jit/vmap/grad).""" + + return not isinstance(value, jax.core.Tracer) @dataclass(frozen=True) @@ -52,7 +57,7 @@ def _validate_A(self, A: jax.Array, label: str) -> None: raise ValueError( f"{label}: coupling matrix A must be ({self.N_c},{self.N_c}), got {A.shape}." ) - if not jnp.allclose(A, A.T, atol=1e-12): + if _is_concrete(A) and not jnp.allclose(A, A.T, atol=1e-12): raise ValueError(f"{label}: coupling matrix A must be symmetric.") def __call__( @@ -72,100 +77,68 @@ def __call__( to the full (c,cp) block. """ N = self.N - N_c = self.N_c N_E = self.N_E - M = N_c * N - dtype = jnp.float64 - all_terms = list(local) + list(nonlocal_) - if all_terms: - g0, _ = all_terms[0] - dtype = jnp.asarray(g0).dtype - - if energy_dependent: - block = jnp.zeros((N_E, M, M), dtype=dtype) - else: - block = jnp.zeros((M, M), dtype=dtype) + M = self.N_c * N + terms: list[jax.Array] = [] for term_idx, (g_raw, a_raw) in enumerate(local): g = jnp.asarray(g_raw) a_mat = jnp.asarray(a_raw) self._validate_A(a_mat, f"local term {term_idx}") - if energy_dependent: - if g.ndim != 2 or g.shape != (N_E, N): - raise ValueError( - f"local term {term_idx}: energy_dependent=True requires g shape " - f"({N_E}, {N}), got {g.shape}." - ) - for c in range(N_c): - for cp in range(N_c): - if a_mat[c, cp] == 0: - continue - row_start = c * N - col_start = cp * N - diag_blocks = jax.vmap(jnp.diag)(a_mat[c, cp] * g) # (N_E, N, N) - block = block.at[ - :, row_start : row_start + N, col_start : col_start + N - ].add(diag_blocks) - else: - if g.ndim != 1 or g.shape != (N,): - raise ValueError( - f"local term {term_idx}: energy_dependent=False requires g shape " - f"({N},), got {g.shape}." - ) - for c in range(N_c): - for cp in range(N_c): - if a_mat[c, cp] == 0: - continue - row_start = c * N - col_start = cp * N - block = block.at[row_start : row_start + N, col_start : col_start + N].add( - jnp.diag(a_mat[c, cp] * g) - ) + expected = (N_E, N) if energy_dependent else (N,) + if g.shape != expected: + raise ValueError( + f"local term {term_idx}: energy_dependent={energy_dependent} requires " + f"g shape {expected}, got {g.shape}." + ) + kernel = jax.vmap(jnp.diag)(g) if energy_dependent else jnp.diag(g) + terms.append(_coupling_kron(a_mat, kernel)) for term_idx, (g_raw, a_raw) in enumerate(nonlocal_): g = jnp.asarray(g_raw) a_mat = jnp.asarray(a_raw) self._validate_A(a_mat, f"nonlocal term {term_idx}") - if energy_dependent: - if g.ndim != 3 or g.shape != (N_E, N, N): - raise ValueError( - f"nonlocal term {term_idx}: energy_dependent=True requires g shape " - f"({N_E}, {N}, {N}), got {g.shape}." - ) - scaled = g * self.gauss_scale[None, :, :] # (N_E, N, N) - for c in range(N_c): - for cp in range(N_c): - if a_mat[c, cp] == 0: - continue - row_start = c * N - col_start = cp * N - block = block.at[ - :, row_start : row_start + N, col_start : col_start + N - ].add(a_mat[c, cp] * scaled) - else: - if g.ndim != 2 or g.shape != (N, N): - raise ValueError( - f"nonlocal term {term_idx}: energy_dependent=False requires g shape " - f"({N}, {N}), got {g.shape}." - ) - scaled = g * self.gauss_scale # (N, N) - for c in range(N_c): - for cp in range(N_c): - if a_mat[c, cp] == 0: - continue - row_start = c * N - col_start = cp * N - block = block.at[row_start : row_start + N, col_start : col_start + N].add( - a_mat[c, cp] * scaled - ) + expected = (N_E, N, N) if energy_dependent else (N, N) + if g.shape != expected: + raise ValueError( + f"nonlocal term {term_idx}: energy_dependent={energy_dependent} requires " + f"g shape {expected}, got {g.shape}." + ) + terms.append(_coupling_kron(a_mat, g * self.gauss_scale)) + + shape = (N_E, M, M) if energy_dependent else (M, M) + dtype = jnp.result_type(*terms) if terms else jnp.float64 + block: jax.Array = jnp.zeros(shape, dtype=dtype) + for term in terms: + block = block + term + # Value checks need concrete arrays: inside jit/vmap the block is a tracer + # and boolean conversion would crash, so symmetry is validated only when the + # Interaction is built outside jax transformations. first_block = block[0] if energy_dependent else block - if not jnp.allclose(first_block, first_block.T, atol=1e-10): + if _is_concrete(first_block) and not jnp.allclose(first_block, first_block.T, atol=1e-10): raise ValueError("Assembled Interaction block is not symmetric.") return Interaction(block=block, energy_dependent=energy_dependent) +def _coupling_kron(a_mat: jax.Array, kernel: jax.Array) -> jax.Array: + """Return the coupled-channel block ``A ⊗ kernel``. + + ``kernel`` is ``(N, N)`` or, with a leading energy axis, ``(N_E, N, N)``; + the result is ``(M, M)`` or ``(N_E, M, M)`` with ``M = N_c·N`` and + block ``(c, cp)`` equal to ``A[c, cp] · kernel``. + """ + + n_c = a_mat.shape[0] + n = kernel.shape[-1] + if kernel.ndim == 3: + batched: jax.Array = jnp.einsum("cd,eij->ecidj", a_mat, kernel) + return batched.reshape(kernel.shape[0], n_c * n, n_c * n) + flat: jax.Array = jnp.einsum("cd,ij->cidj", a_mat, kernel) + return flat.reshape(n_c * n, n_c * n) + + @dataclass(frozen=True) class _InteractionFromFuncs: """Pickle-safe callable that evaluates potential functions then delegates to array builder.""" @@ -192,15 +165,14 @@ def __call__( g(r, r', E) — energy-dependent; r,r' shape (N,N), E scalar, returns (N,N) """ r = self.radii - N_E = self.N_E ri, rj = jnp.meshgrid(r, r, indexing="ij") # (N, N) + # One device→host transfer up front instead of one per energy sample. + energy_values = [float(energy) for energy in np.asarray(self.energies)] local_arrays: list[Any] = [] for g_fn, a_mat in local: if energy_dependent: - g_arr = jnp.stack( - [g_fn(r, float(self.energies[ie])) for ie in range(N_E)] - ) # (N_E, N) + g_arr = jnp.stack([g_fn(r, energy) for energy in energy_values]) # (N_E, N) else: g_arr = g_fn(r) # (N,) local_arrays.append((g_arr, a_mat)) @@ -208,9 +180,7 @@ def __call__( nonlocal_arrays: list[Any] = [] for g_fn, a_mat in nonlocal_: if energy_dependent: - g_arr = jnp.stack( - [g_fn(ri, rj, float(self.energies[ie])) for ie in range(N_E)] - ) # (N_E, N, N) + g_arr = jnp.stack([g_fn(ri, rj, energy) for energy in energy_values]) # (N_E, N, N) else: g_arr = g_fn(ri, rj) # (N, N) nonlocal_arrays.append((g_arr, a_mat)) @@ -271,14 +241,17 @@ def make_interaction_from_funcs( @dataclass(frozen=True) class _PotentialBuilder: - """Pickle-safe ``solver.potential(fn, coupling, energy_dependent)`` callable. + """Pickle-safe single-kind ``solver.{local,nonlocal}_potential`` callable. - Wraps ``_InteractionFromFuncs`` with arity-based local/nonlocal dispatch and - a sensible ``coupling`` default for single-channel problems. + Wraps ``_InteractionFromFuncs`` for one fixed ``kind`` (``"local"`` or + ``"nonlocal"``), with a sensible ``coupling`` default for single-channel + problems. The local/nonlocal choice is made explicit by the entry point the + caller picks, so there is no arity inference. """ n_c: int funcs_builder: _InteractionFromFuncs + kind: Literal["local", "nonlocal"] def __call__( self, @@ -292,10 +265,10 @@ def __call__( Parameters ---------- fn - Potential function. Arity determines local vs nonlocal: + Potential function for this builder's fixed ``kind``: - * ``energy_dependent=False``: ``fn(r)`` → local, ``fn(r, r')`` → nonlocal - * ``energy_dependent=True``: ``fn(r, E)`` → local, ``fn(r, r', E)`` → nonlocal + * local — ``fn(r)`` or, with ``energy_dependent=True``, ``fn(r, E)`` + * nonlocal — ``fn(r, r')`` or ``fn(r, r', E)`` coupling ``(N_c, N_c)`` symmetric coupling matrix. Defaults to ``[[1.0]]`` @@ -318,47 +291,32 @@ def __call__( "re-compile with energies=... to use energy-dependent potentials." ) - n_args = len(inspect.signature(fn).parameters) # type: ignore[arg-type] - if energy_dependent: - if n_args == 2: - return self.funcs_builder(local=[(fn, coupling)], energy_dependent=True) # type: ignore[list-item] - elif n_args == 3: - return self.funcs_builder(nonlocal_=[(fn, coupling)], energy_dependent=True) # type: ignore[list-item] - else: - raise ValueError( - f"For energy_dependent=True, fn must take 2 args fn(r, E) " - f"(local) or 3 args fn(r, r', E) (nonlocal); got {n_args}." - ) - else: - if n_args == 1: - return self.funcs_builder(local=[(fn, coupling)], energy_dependent=False) # type: ignore[list-item] - elif n_args == 2: - return self.funcs_builder(nonlocal_=[(fn, coupling)], energy_dependent=False) # type: ignore[list-item] - else: - raise ValueError( - f"fn must take 1 arg fn(r) (local) or 2 args fn(r, r') " - f"(nonlocal); got {n_args}." - ) + if self.kind == "local": + return self.funcs_builder(local=[(fn, coupling)], energy_dependent=energy_dependent) # type: ignore[list-item] + return self.funcs_builder(nonlocal_=[(fn, coupling)], energy_dependent=energy_dependent) # type: ignore[list-item] -def make_potential_builder( +def make_potential_builders( mesh: Mesh, channels: tuple[ChannelSpec, ...], energies: jax.Array, -) -> _PotentialBuilder: - """Return ``solver.potential`` — a simple callable that builds an Interaction from a function. +) -> tuple[_PotentialBuilder, _PotentialBuilder]: + """Return ``(local_potential, nonlocal_potential)`` builders for one solver. - The returned callable auto-detects local vs nonlocal from function arity and - defaults ``coupling`` to ``[[1.0]]`` for single-channel problems. See - :class:`_PotentialBuilder` for the full signature. + Each builds an :class:`~lax.Interaction` from a function of the corresponding + kind, defaulting ``coupling`` to ``[[1.0]]`` for single-channel problems. Both + share one underlying ``_InteractionFromFuncs``. See :class:`_PotentialBuilder`. """ funcs_builder = make_interaction_from_funcs(mesh, channels, energies) - return _PotentialBuilder(n_c=len(channels), funcs_builder=funcs_builder) + n_c = len(channels) + local = _PotentialBuilder(n_c=n_c, funcs_builder=funcs_builder, kind="local") + nonlocal_ = _PotentialBuilder(n_c=n_c, funcs_builder=funcs_builder, kind="nonlocal") + return local, nonlocal_ __all__ = [ "make_interaction_from_array", "make_interaction_from_block", "make_interaction_from_funcs", - "make_potential_builder", + "make_potential_builders", ] diff --git a/src/lax/propagate.py b/src/lax/propagate.py index d153d7b..66110e5 100644 --- a/src/lax/propagate.py +++ b/src/lax/propagate.py @@ -6,7 +6,7 @@ import jax.numpy as jnp import numpy as np -from lax.boundary._types import PropagationMatrices +from lax.types import PropagationMatrices def build_legendre_x_propagation( diff --git a/src/lax/solvers/__init__.py b/src/lax/solvers/__init__.py index 0a917ee..2051e38 100644 --- a/src/lax/solvers/__init__.py +++ b/src/lax/solvers/__init__.py @@ -4,7 +4,6 @@ from lax.solvers.linear_solve import ( make_direct_wavefunction_kernel, make_phases_direct_observable, - make_rmatrix_direct_grid_observable, make_rmatrix_direct_kernel, make_smatrix_direct_observable, ) @@ -23,7 +22,6 @@ "build_Q", "make_direct_wavefunction_kernel", "make_phases_direct_observable", - "make_rmatrix_direct_grid_observable", "make_rmatrix_direct_kernel", "make_smatrix_direct_observable", "make_spectrum_kernel", diff --git a/src/lax/solvers/assembly.py b/src/lax/solvers/assembly.py index 13ba092..d27c090 100644 --- a/src/lax/solvers/assembly.py +++ b/src/lax/solvers/assembly.py @@ -2,11 +2,12 @@ from __future__ import annotations +from typing import cast + import jax import jax.numpy as jnp -from lax.boundary._types import Mesh, OperatorMatrices -from lax.types import ChannelSpec +from lax.types import ChannelSpec, Mesh, OperatorMatrices def assemble_block_hamiltonian( @@ -140,4 +141,41 @@ def _require_operator(operator: jax.Array | None, name: str) -> jax.Array: return operator -__all__ = ["assemble_block_hamiltonian", "build_Q"] +def uniform_mass_factor(channels: tuple[ChannelSpec, ...], context: str = "solver path") -> float: + """Return the single mass factor shared by all channels, or raise. + + Several kernels (the spectral eigensolve and the propagated direct solve) + fold ``ℏ²/2μ`` out of the Hamiltonian using one scalar. They are only + correct when every channel shares that mass factor, so this validates the + assumption at compile time rather than silently producing wrong physics. + + Parameters + ---------- + channels + Channel definitions. + context + Short description of the caller, interpolated into the error message. + + Returns + ------- + float + The shared ``channels[0].mass_factor``. + + Raises + ------ + ValueError + If the channels do not all share one mass factor. + """ + + mass_factor = channels[0].mass_factor + for channel in channels[1:]: + if channel.mass_factor != mass_factor: + msg = ( + f"The {context} requires a uniform mass_factor across channels; " + "use the per-energy/per-channel direct grid path for multi-μ systems." + ) + raise ValueError(msg) + return cast(float, mass_factor) + + +__all__ = ["assemble_block_hamiltonian", "build_Q", "uniform_mass_factor"] diff --git a/src/lax/solvers/linear_solve.py b/src/lax/solvers/linear_solve.py index 4ceee54..cdd47e4 100644 --- a/src/lax/solvers/linear_solve.py +++ b/src/lax/solvers/linear_solve.py @@ -6,21 +6,23 @@ from typing import TYPE_CHECKING, cast import jax +import jax.core import jax.numpy as jnp +import jax.scipy.linalg as jsl import numpy as np -from lax.boundary._types import ( +from lax.spectral.matching import open_channel_smatrix_from_R, phases_from_S +from lax.types import ( BoundaryValues, - DirectGridObservable, + ChannelSpec, DirectRMatrixKernel, + Interaction, Mesh, OperatorMatrices, PropagationMatrices, ) -from lax.spectral.matching import phases_from_S, smatrix_from_R -from lax.types import ChannelSpec, Interaction -from .assembly import assemble_block_hamiltonian, build_Q +from .assembly import assemble_block_hamiltonian, build_Q, uniform_mass_factor if TYPE_CHECKING: pass @@ -50,14 +52,14 @@ class _DirectRMatrixKernel: mesh: Mesh operators: OperatorMatrices channels: tuple[ChannelSpec, ...] + static_channels: tuple[ChannelSpec, ...] energies: jax.Array q: jax.Array q_prime: jax.Array channel_radius: float matrix_size: int - mass_factor: float | jax.Array - boundary: BoundaryValues | None mass_factor_grid: jax.Array + energy_uniform_mu: bool def __call__(self, potential: jax.Array | Interaction) -> jax.Array: """Evaluate the direct R-matrix on the compile-time energy grid. @@ -65,137 +67,143 @@ def __call__(self, potential: jax.Array | Interaction) -> jax.Array: Parameters ---------- potential - :class:`~lax.Interaction` object built by ``solver.potential()`` or + :class:`~lax.Interaction` object built by ``solver.local_potential()``/``solver.nonlocal_potential()`` or ``solver.interaction_from_{block,array,funcs}()``. Energy-dependent interactions (``energy_dependent=True``) use the per-energy block path. - For propagated meshes, local energy-independent Interactions are supported: - the per-interval ``(N_c, N_c, N)`` array is extracted from the block diagonals. + Notes + ----- + ``mass_factor_grid`` is honoured on every path. When μ is uniform across + the energy grid (the common case) the fast static path is used with a + single Hamiltonian assembly and the per-channel-μ surface projector + ``Q'``. When μ varies with energy, an energy-independent block is + broadcast over the grid and solved per energy with ``mass_factor_grid`` + in both ``C(E_i)`` and ``Q'`` — matching §11.3. """ - # Propagated meshes use per-interval raw (N_c, N_c, N) arrays. - # Extract from Interaction.block by taking the diagonal of each sub-block. - if self.mesh.propagation is not None: - if not isinstance(potential, Interaction): - raise TypeError( - "rmatrix_direct() accepts only Interaction objects. " - "Use solver.potential(fn) or solver.interaction_from_block(block)." - ) - if potential.energy_dependent: - raise TypeError( - "rmatrix_direct() does not support energy-dependent Interactions " - "on propagated meshes." - ) - N_c = len(self.channels) - N = self.mesh.n - # Propagated path supports only local Interactions: sub-blocks must be diagonal. - for c in range(N_c): - for cp in range(N_c): - sub = np.asarray(potential.block[c * N : (c + 1) * N, cp * N : (cp + 1) * N]) - if np.any(sub != np.diag(np.diag(sub))): - raise ValueError( - "rmatrix_direct() on propagated meshes supports only local " - "Interactions. Non-local propagated direct solves are not supported." - ) - potential = jnp.stack( - [ - jnp.stack( - [ - jnp.diag(potential.block[c * N : (c + 1) * N, cp * N : (cp + 1) * N]) - for cp in range(N_c) - ] - ) - for c in range(N_c) - ] + if not isinstance(potential, Interaction): + raise TypeError( + "rmatrix_direct() accepts only Interaction objects. " + "Use solver.local_potential()/solver.nonlocal_potential() or solver.interaction_from_block/array/funcs to build one." ) + + if potential.energy_dependent: return cast( jax.Array, - _RMATRIX_DIRECT_JIT( - potential, + _RMATRIX_DIRECT_GRID_JIT( + potential.block, self.mesh, self.operators, self.channels, self.energies, - self.q_prime, + self.q, self.channel_radius, self.matrix_size, - self.mass_factor, - self.boundary, + self.mass_factor_grid, ), ) - - if not isinstance(potential, Interaction): - raise TypeError( - "rmatrix_direct() accepts only Interaction objects. " - "Use solver.potential() or solver.interaction_from_block/array/funcs to build one." - ) - - if potential.energy_dependent: + if self.energy_uniform_mu: + # μ is constant across energies: one assembly, per-channel-μ Q'. return cast( jax.Array, - _RMATRIX_DIRECT_GRID_JIT( + _RMATRIX_DIRECT_JIT( potential.block, self.mesh, self.operators, - self.channels, + self.static_channels, self.energies, - self.q, + self.q_prime, self.channel_radius, self.matrix_size, - self.mass_factor, - self.boundary, - self.mass_factor_grid, ), ) + # Energy-independent block but energy-dependent μ(E): broadcast the block + # over the grid and reuse the per-energy grid solver so each C(E_i)/Q' + # carries the correct mass_factor_grid row. + n_e = self.energies.shape[0] + block_grid = jnp.broadcast_to(potential.block, (n_e, *potential.block.shape)) return cast( jax.Array, - _RMATRIX_DIRECT_JIT( - potential.block, + _RMATRIX_DIRECT_GRID_JIT( + block_grid, self.mesh, self.operators, self.channels, self.energies, - self.q_prime, + self.q, self.channel_radius, self.matrix_size, - self.mass_factor, - self.boundary, + self.mass_factor_grid, ), ) @dataclass(frozen=True) -class _DirectRMatrixGridObservable: - """Pickle-safe aligned-grid direct R-matrix observable.""" +class _PropagatedDirectRMatrixKernel: + """Pickle-safe direct R-matrix kernel for subinterval-propagated meshes. + + Selected at compile time by :func:`make_rmatrix_direct_kernel` when + ``mesh.propagation`` is set, so the non-propagated kernel carries no + propagation special cases. Supports local energy-independent + Interactions only: the per-interval ``(N_c, N_c, N)`` array is extracted + from the block diagonals. + """ mesh: Mesh - operators: OperatorMatrices channels: tuple[ChannelSpec, ...] energies: jax.Array - q: jax.Array - channel_radius: float - matrix_size: int mass_factor: float | jax.Array boundary: BoundaryValues | None - mass_factor_grid: jax.Array - def __call__(self, potentials: jax.Array | Interaction) -> jax.Array: - """Evaluate `R(E_i; V_i)` across the compile-time energy grid.""" + def __call__(self, potential: jax.Array | Interaction) -> jax.Array: + """Evaluate the propagated direct R-matrix on the compile-time energy grid.""" + if not isinstance(potential, Interaction): + raise TypeError( + "rmatrix_direct() accepts only Interaction objects. " + "Use solver.local_potential(fn)/solver.nonlocal_potential(fn) or solver.interaction_from_block(block)." + ) + if potential.energy_dependent: + raise TypeError( + "rmatrix_direct() does not support energy-dependent Interactions " + "on propagated meshes." + ) + n_c = len(self.channels) + n = self.mesh.n + # Propagated path supports only local Interactions: sub-blocks must be + # diagonal. Validate only outside jax transformations — inside vmap/jit + # the block is a tracer and host inspection would crash; the diagonality + # contract is a property of the user's input, checked once when concrete. + if not isinstance(potential.block, jax.core.Tracer): + host_block = np.asarray(potential.block) # one device→host transfer + for c in range(n_c): + for cp in range(n_c): + sub = host_block[c * n : (c + 1) * n, cp * n : (cp + 1) * n] + if not np.allclose(sub, np.diag(np.diag(sub))): + raise ValueError( + "rmatrix_direct() on propagated meshes supports only local " + "Interactions. Non-local propagated direct solves are not supported." + ) + local_potential = jnp.stack( + [ + jnp.stack( + [ + jnp.diag(potential.block[c * n : (c + 1) * n, cp * n : (cp + 1) * n]) + for cp in range(n_c) + ] + ) + for c in range(n_c) + ] + ) return cast( jax.Array, - _RMATRIX_DIRECT_GRID_JIT( - potentials, + _RMATRIX_DIRECT_PROPAGATED_JIT( + local_potential, self.mesh, - self.operators, self.channels, self.energies, - self.q, - self.channel_radius, - self.matrix_size, self.mass_factor, self.boundary, - self.mass_factor_grid, ), ) @@ -204,7 +212,7 @@ def __call__(self, potentials: jax.Array | Interaction) -> jax.Array: class _SMatrixDirectObservable: """Pickle-safe direct S-matrix observable derived from rmatrix_direct.""" - rmatrix_direct: _DirectRMatrixKernel + rmatrix_direct: DirectRMatrixKernel boundary: BoundaryValues def __call__(self, potential: jax.Array | Interaction) -> jax.Array: @@ -236,6 +244,7 @@ class _WavefunctionDirectKernel: channels: tuple[ChannelSpec, ...] energies: jax.Array matrix_size: int + mass_factor_grid: jax.Array def __call__( self, @@ -263,7 +272,7 @@ def __call__( if not isinstance(potential, Interaction): raise TypeError( "wavefunction_direct() accepts only Interaction objects. " - "Use solver.potential() or solver.interaction_from_block/array/funcs to build one." + "Use solver.local_potential()/solver.nonlocal_potential() or solver.interaction_from_block/array/funcs to build one." ) block = potential.block[energy_index] if potential.energy_dependent else potential.block @@ -273,6 +282,7 @@ def __call__( block, source, self.energies[energy_index], + self.mass_factor_grid[energy_index], self.mesh, self.operators, self.channels, @@ -322,110 +332,78 @@ def make_rmatrix_direct_kernel( JIT-compiled callable: ``kernel(V) → R`` with shape ``(N_E, N_c, N_c)``. """ + if mesh.propagation is not None: + # The propagation recursion converts the potential to fm⁻² with a single + # scalar mass factor, so it is only valid for a uniform-μ channel set. + propagated_mass_factor = uniform_mass_factor( + channels, context="propagated direct R-matrix path" + ) + return cast( + DirectRMatrixKernel, + _PropagatedDirectRMatrixKernel( + mesh=mesh, + channels=channels, + energies=energies, + mass_factor=propagated_mass_factor, + boundary=boundary, + ), + ) + q = build_Q(mesh, channels) - q_prime = _build_q_prime(q, channels, mesh.n) - channel_radius = mesh.scale - matrix_size = mesh.n * len(channels) - mass_factor = channels[0].mass_factor # used by propagated path only n_e = len(energies) n_c = len(channels) + # compile() pre-broadcasts every supported input to the canonical (N_E, N_c) + # form via compile._broadcast_mass_factor_grid, so the only cases left here + # are "no grid" (use each channel's static mass_factor) and a ready (N_E, N_c) + # array. Validate rather than re-implement the broadcast. if mass_factor_grid is None: _mfg: jax.Array = jnp.broadcast_to( jnp.array([c.mass_factor for c in channels], dtype=float), (n_e, n_c) ) - elif jnp.asarray(mass_factor_grid).ndim == 1: - _mfg = jnp.broadcast_to(jnp.asarray(mass_factor_grid)[:, None], (n_e, n_c)) else: _mfg = jnp.asarray(mass_factor_grid) + if _mfg.shape != (n_e, n_c): + msg = ( + f"mass_factor_grid must be pre-broadcast to (N_E, N_c)=({n_e}, {n_c}); " + f"got {_mfg.shape}. Build the solver via lax.compile()." + ) + raise ValueError(msg) + # When μ is the same at every energy, the energy-independent path can assemble + # the Hamiltonian once and use a single per-channel-μ Q'. The "static" channels + # carry that uniform μ so both H and Q' honour mass_factor_grid even when it + # differs from each ChannelSpec.mass_factor. When μ varies with energy the + # static path is bypassed (see _DirectRMatrixKernel.__call__). + mfg_np = np.asarray(_mfg) + energy_uniform_mu = bool(np.allclose(mfg_np, mfg_np[0:1, :])) + if energy_uniform_mu: + uniform_mu = mfg_np[0] + static_channels = tuple( + ChannelSpec(l=ch.l, threshold=ch.threshold, mass_factor=float(uniform_mu[i])) + for i, ch in enumerate(channels) + ) + else: + static_channels = channels + q_prime = _build_q_prime(q, static_channels, mesh.n) return cast( DirectRMatrixKernel, _DirectRMatrixKernel( mesh=mesh, operators=operators, channels=channels, + static_channels=static_channels, energies=energies, q=q, q_prime=q_prime, - channel_radius=channel_radius, - matrix_size=matrix_size, - mass_factor=mass_factor, - boundary=boundary, - mass_factor_grid=_mfg, - ), - ) - - -def make_rmatrix_direct_grid_observable( - mesh: Mesh, - operators: OperatorMatrices, - channels: tuple[ChannelSpec, ...], - energies: jax.Array, - boundary: BoundaryValues | None, - mass_factor_grid: jax.Array | None = None, -) -> DirectGridObservable: - """Build a JIT-compiled aligned-grid ``R(E_i; V_i)`` kernel. - - For energy-dependent potentials: given a stack of potentials - ``V_grid[i]`` at compile-time energy ``E_i``, returns ``R[i]`` evaluated - with that exact ``(V_i, E_i)`` pairing — the physically correct aligned-grid - evaluation. - - Parameters - ---------- - mesh - Compiled mesh. - operators - Precomputed operator matrices. - channels - Channel definitions. - energies - Compile-time energy grid in MeV, shape ``(N_E,)``. - boundary - Compile-time boundary values, or ``None``. - mass_factor_grid - Per-energy ℏ²/2μ values in MeV·fm², shape ``(N_E,)``, or ``None`` - for constant mass factor. - - Returns - ------- - DirectGridObservable - JIT-compiled callable: ``kernel(V_grid) → R`` where ``V_grid`` has a - leading ``N_E`` axis and ``R`` has shape ``(N_E, N_c, N_c)``. - """ - - q = build_Q(mesh, channels) - channel_radius = mesh.scale - matrix_size = mesh.n * len(channels) - mass_factor = channels[0].mass_factor # used by propagated path only - n_e = len(energies) - n_c = len(channels) - if mass_factor_grid is None: - _mfg: jax.Array = jnp.broadcast_to( - jnp.array([c.mass_factor for c in channels], dtype=float), (n_e, n_c) - ) - elif jnp.asarray(mass_factor_grid).ndim == 1: - _mfg = jnp.broadcast_to(jnp.asarray(mass_factor_grid)[:, None], (n_e, n_c)) - else: - _mfg = jnp.asarray(mass_factor_grid) - return cast( - DirectGridObservable, - _DirectRMatrixGridObservable( - mesh=mesh, - operators=operators, - channels=channels, - energies=energies, - q=q, - channel_radius=channel_radius, - matrix_size=matrix_size, - mass_factor=mass_factor, - boundary=boundary, + channel_radius=mesh.scale, + matrix_size=mesh.n * len(channels), mass_factor_grid=_mfg, + energy_uniform_mu=energy_uniform_mu, ), ) def make_smatrix_direct_observable( - rmatrix_kernel: _DirectRMatrixKernel, + rmatrix_kernel: DirectRMatrixKernel, boundary: BoundaryValues | None, ) -> _SMatrixDirectObservable | None: """Build a direct S-matrix observable from a direct R-matrix kernel.""" @@ -450,20 +428,39 @@ def make_direct_wavefunction_kernel( operators: OperatorMatrices, channels: tuple[ChannelSpec, ...], energies: jax.Array, + mass_factor_grid: jax.Array | None = None, ) -> _WavefunctionDirectKernel: """Build a direct wavefunction kernel ``(V, source, i) → ψ``. - Solves ``C(E_i) ψ = source`` where ``C = H_MeV − E_i · I`` using - ``jnp.linalg.solve``. [DESIGN.md §11.3] + Solves ``C(E_i) ψ = source`` where ``C = H_MeV − E_i · I`` is assembled with + the per-energy per-channel mass factor ``mass_factor_grid[i]`` and the solve + is scaled per channel by μ to match the fm⁻² spectral Green's convention. + ``mass_factor_grid`` follows the canonical ``(N_E, N_c)`` form; when ``None`` + each channel's static ``mass_factor`` is used. [DESIGN.md §11.3] """ matrix_size = mesh.n * len(channels) + n_e = len(energies) + n_c = len(channels) + if mass_factor_grid is None: + _mfg: jax.Array = jnp.broadcast_to( + jnp.array([c.mass_factor for c in channels], dtype=float), (n_e, n_c) + ) + else: + _mfg = jnp.asarray(mass_factor_grid) + if _mfg.shape != (n_e, n_c): + msg = ( + f"mass_factor_grid must be pre-broadcast to (N_E, N_c)=({n_e}, {n_c}); " + f"got {_mfg.shape}. Build the solver via lax.compile()." + ) + raise ValueError(msg) return _WavefunctionDirectKernel( mesh=mesh, operators=operators, channels=channels, energies=energies, matrix_size=matrix_size, + mass_factor_grid=_mfg, ) @@ -476,15 +473,9 @@ def _rmatrix_direct( q_prime: jax.Array, channel_radius: float, matrix_size: int, - mass_factor: float | jax.Array, - boundary: BoundaryValues | None, ) -> jax.Array: """Return the direct R-matrix across the compile-time energy grid. - Dispatches to the propagated or non-propagated path depending on - ``mesh.propagation``, and to the local or non-local potential path - depending on ``potential.ndim``. - The Hamiltonian is assembled in MeV (symmetric form), and the C matrix is ``H_MeV − E·I``. The surface projector ``Q'`` carries the per-channel sqrt(m_c) factor so that ``R = Q'^T C^{-1} Q' / a`` equals the fm⁻² result @@ -494,10 +485,9 @@ def _rmatrix_direct( ---------- potential Assembled potential in MeV. Local: ``(N_c, N_c, N)``; non-local: - ``(N_c, N_c, N, N)``. - mesh, operators, channels, energies, q_prime, channel_radius, matrix_size, mass_factor, boundary + ``(N_c, N_c, N, N)``; pre-assembled block: ``(M, M)``. + mesh, operators, channels, energies, q_prime, channel_radius, matrix_size Compile-time cached data forwarded from the kernel dataclass. - ``mass_factor`` is used only on the propagated path (fm⁻² units). Returns ------- @@ -505,45 +495,6 @@ def _rmatrix_direct( R-matrix on the compile-time energy grid, shape ``(N_E, N_c, N_c)``. """ - if mesh.propagation is not None and potential.ndim == 3: - if boundary is None: - msg = "Boundary values are required for propagated direct R-matrix solves." - raise ValueError(msg) - propagation = mesh.propagation - - def propagated_one_energy( - energy: jax.Array, - h_plus: jax.Array, - h_plus_p: jax.Array, - is_open: jax.Array, - ) -> jax.Array: - return _propagated_rmatrix_at_energy( - potential, - propagation, - channels, - energy, - h_plus, - h_plus_p, - is_open, - mass_factor, - ) - - result: jax.Array = jax.vmap(propagated_one_energy)( - energies, - boundary.H_plus, - boundary.H_plus_p, - boundary.is_open, - ) - return result - - if mesh.propagation is not None and potential.ndim == 4: - msg = ( - "Subinterval propagation is defined only for local potentials in the direct " - "linear-solve formulation. Non-local propagated solves are not mathematically " - "supported." - ) - raise ValueError(msg) - hamiltonian = assemble_block_hamiltonian( mesh, operators, @@ -552,18 +503,16 @@ def propagated_one_energy( ) def one_energy(energy: jax.Array) -> jax.Array: - # Hamiltonian is in MeV; C = H_MeV − E·I. + # Hamiltonian is in MeV; C = H_MeV − E·I. Factor C once (§11.3: one + # factorization of C(E_i) serves the surface solve here and the + # wavefunction solve in _wavefunction_direct) and reuse it for every + # column of Q'. matrix = hamiltonian - energy * jnp.eye( matrix_size, dtype=hamiltonian.dtype, ) - solved = cast( - jax.Array, - jnp.linalg.solve( - matrix, - q_prime, - ), - ) + lu_piv = cast(tuple[jax.Array, jax.Array], jsl.lu_factor(matrix)) + solved = cast(jax.Array, jsl.lu_solve(lu_piv, q_prime)) values: jax.Array = (q_prime.T @ solved) / channel_radius return values @@ -571,6 +520,64 @@ def one_energy(energy: jax.Array) -> jax.Array: return result +def _rmatrix_direct_propagated( + potential: jax.Array, + mesh: Mesh, + channels: tuple[ChannelSpec, ...], + energies: jax.Array, + mass_factor: float | jax.Array, + boundary: BoundaryValues | None, +) -> jax.Array: + """Return the propagated direct R-matrix across the compile-time energy grid. + + Parameters + ---------- + potential + Local potential in MeV, shape ``(N_c, N_c, N)``. + mesh, channels, energies, mass_factor, boundary + Compile-time cached data forwarded from the propagated kernel dataclass. + ``mass_factor`` converts the potential to fm⁻² inside the recursion. + + Returns + ------- + jax.Array + R-matrix on the compile-time energy grid, shape ``(N_E, N_c, N_c)``. + """ + + if boundary is None: + msg = "Boundary values are required for propagated direct R-matrix solves." + raise ValueError(msg) + propagation = mesh.propagation + if propagation is None: + msg = "Propagated direct solves require a mesh with propagation data." + raise ValueError(msg) + + def propagated_one_energy( + energy: jax.Array, + h_plus: jax.Array, + h_plus_p: jax.Array, + is_open: jax.Array, + ) -> jax.Array: + return _propagated_rmatrix_at_energy( + potential, + propagation, + channels, + energy, + h_plus, + h_plus_p, + is_open, + mass_factor, + ) + + result: jax.Array = jax.vmap(propagated_one_energy)( + energies, + boundary.H_plus, + boundary.H_plus_p, + boundary.is_open, + ) + return result + + def _rmatrix_direct_grid( potentials: jax.Array, mesh: Mesh, @@ -580,8 +587,6 @@ def _rmatrix_direct_grid( q: jax.Array, channel_radius: float, matrix_size: int, - mass_factor: float | jax.Array, - boundary: BoundaryValues | None, mass_factor_grid: jax.Array, ) -> jax.Array: """Return aligned-grid ``R(E_i; V_i)`` samples for energy-dependent potentials. @@ -594,10 +599,9 @@ def _rmatrix_direct_grid( ---------- potentials Per-energy potentials in MeV. Local: ``(N_E, N_c, N_c, N)``; non-local: - ``(N_E, N_c, N_c, N, N)``. - mesh, operators, channels, energies, q, channel_radius, matrix_size, mass_factor, boundary - Compile-time cached data. ``q`` is the unscaled surface projector; - ``mass_factor`` is used only on the propagated path. + ``(N_E, N_c, N_c, N, N)``; pre-assembled blocks: ``(N_E, M, M)``. + mesh, operators, channels, energies, q, channel_radius, matrix_size + Compile-time cached data. ``q`` is the unscaled surface projector. mass_factor_grid Dense ``(N_E, N_c)`` mass-factor array, always present (promoted at compile time from scalar / ``(N_E,)`` / per-channel inputs). @@ -608,47 +612,6 @@ def _rmatrix_direct_grid( R-matrix samples, shape ``(N_E, N_c, N_c)``. """ - if mesh.propagation is not None and potentials.ndim == 4: - if boundary is None: - msg = "Boundary values are required for propagated direct R-matrix solves." - raise ValueError(msg) - propagation = mesh.propagation - - def propagated_one_energy( - potential: jax.Array, - energy: jax.Array, - h_plus: jax.Array, - h_plus_p: jax.Array, - is_open: jax.Array, - ) -> jax.Array: - return _propagated_rmatrix_at_energy( - potential, - propagation, - channels, - energy, - h_plus, - h_plus_p, - is_open, - mass_factor, - ) - - result: jax.Array = jax.vmap(propagated_one_energy)( - potentials, - energies, - boundary.H_plus, - boundary.H_plus_p, - boundary.is_open, - ) - return result - - if mesh.propagation is not None and potentials.ndim == 5: - msg = ( - "Subinterval propagation is defined only for local potentials in the direct " - "linear-solve formulation. Non-local propagated solves are not mathematically " - "supported." - ) - raise ValueError(msg) - def one_energy( potential: jax.Array, energy: jax.Array, @@ -665,7 +628,8 @@ def one_energy( n = mesh.n scale = jnp.repeat(jnp.sqrt(mu_row), n) # (N_c·N,) q_prime_mu: jax.Array = scale[:, None] * q - solved = cast(jax.Array, jnp.linalg.solve(matrix, q_prime_mu)) + lu_piv = cast(tuple[jax.Array, jax.Array], jsl.lu_factor(matrix)) + solved = cast(jax.Array, jsl.lu_solve(lu_piv, q_prime_mu)) return (q_prime_mu.T @ solved) / channel_radius # mass_factor_grid is (N_E, N_c); vmap slices to (N_c,) per energy step. @@ -676,6 +640,7 @@ def _wavefunction_direct( potential: jax.Array, source: jax.Array, energy: jax.Array, + mu_row: jax.Array, mesh: Mesh, operators: OperatorMatrices, channels: tuple[ChannelSpec, ...], @@ -683,15 +648,24 @@ def _wavefunction_direct( ) -> jax.Array: """Solve the internal wavefunction on the MeV direct path. - Computes ``m₀ · (H_MeV − E·I)⁻¹ source`` where ``m₀ = channels[0].mass_factor``. - The ``m₀`` factor makes the result equal to the fm⁻² spectral Green's function: - ``G_spectral = (H_fm2 − E/m)⁻¹ = m · (H_MeV − E·I)⁻¹``. + Assembles ``C = H_MeV(μ) − E·I`` using the per-channel mass factors in + ``mu_row`` (shape ``(N_c,)``), solves ``C ψ̃ = source``, then scales channel + block ``c`` by ``μ_c`` so the result equals the fm⁻² spectral Green's + function ``G_spectral = μ · (H_MeV − E·I)⁻¹`` channel-by-channel. For a + uniform μ this reduces to the previous ``m₀ · solve`` behaviour. """ - hamiltonian = assemble_block_hamiltonian(mesh, operators, channels, potential) + updated = tuple( + ChannelSpec(l=ch.l, threshold=ch.threshold, mass_factor=mu_row[i]) + for i, ch in enumerate(channels) + ) + hamiltonian = assemble_block_hamiltonian(mesh, operators, updated, potential) matrix = hamiltonian - energy * jnp.eye(matrix_size, dtype=hamiltonian.dtype) - m0 = channels[0].mass_factor # evaluated at JIT-trace time (static) - result: jax.Array = cast(jax.Array, m0 * jnp.linalg.solve(matrix, source)) + # Same C(E_i) factorization formulation as the direct R-matrix solve (§11.3). + lu_piv = cast(tuple[jax.Array, jax.Array], jsl.lu_factor(matrix)) + solved = cast(jax.Array, jsl.lu_solve(lu_piv, source)) + scale = jnp.repeat(mu_row, mesh.n) # (N_c·N,) — per-channel μ on the diagonal + result: jax.Array = scale.astype(solved.dtype) * solved return result @@ -699,9 +673,15 @@ def _direct_smatrix_grid( r_grid: jax.Array, boundary: BoundaryValues, ) -> jax.Array: - """Match an (N_E, N_c, N_c) R-matrix grid to the S-matrix grid.""" + """Match an (N_E, N_c, N_c) R-matrix grid to the S-matrix grid. - return jax.vmap(smatrix_from_R)(r_grid, boundary) + Uses :func:`open_channel_smatrix_from_R` so closed channels are decoupled + and projected out exactly as on the spectral path — keeping + ``smatrix_direct``/``phases_direct`` consistent with ``smatrix``/``phases``. + For an all-open channel set this is identical to bare ``smatrix_from_R``. + """ + + return jax.vmap(open_channel_smatrix_from_R)(r_grid, boundary) def _direct_phases_grid(s_grid: jax.Array) -> jax.Array: @@ -714,6 +694,10 @@ def _direct_phases_grid(s_grid: jax.Array) -> jax.Array: _rmatrix_direct, static_argnames=("channels", "matrix_size"), ) +_RMATRIX_DIRECT_PROPAGATED_JIT = jax.jit( + _rmatrix_direct_propagated, + static_argnames=("channels",), +) _RMATRIX_DIRECT_GRID_JIT = jax.jit( _rmatrix_direct_grid, static_argnames=("channels", "matrix_size"), @@ -883,7 +867,6 @@ def _surface_projector( __all__ = [ "make_direct_wavefunction_kernel", "make_phases_direct_observable", - "make_rmatrix_direct_grid_observable", "make_rmatrix_direct_kernel", "make_smatrix_direct_observable", ] diff --git a/src/lax/solvers/observables.py b/src/lax/solvers/observables.py index 30e1625..ff9b964 100644 --- a/src/lax/solvers/observables.py +++ b/src/lax/solvers/observables.py @@ -15,8 +15,17 @@ import jax import jax.numpy as jnp -from lax.boundary._types import ( +from lax.spectral.interpolation import pade_interpolate +from lax.spectral.matching import open_channel_smatrix_from_R, phases_from_S +from lax.spectral.observables import ( + greens_from_spectrum, + rmatrix_from_spectrum, + wavefunction_internal_from_spectrum, +) +from lax.spectral.types import Spectrum +from lax.types import ( BoundaryValues, + ChannelSpec, EigenpairAccessor, GreenFunctionObservable, InterpolatorBuilder, @@ -26,15 +35,8 @@ SpectrumObservable, WavefunctionObservable, ) -from lax.spectral.interpolation import pade_interpolate -from lax.spectral.matching import phases_from_S, smatrix_from_R -from lax.spectral.observables import ( - greens_from_spectrum, - rmatrix_from_spectrum, - wavefunction_internal_from_spectrum, -) -from lax.spectral.types import Spectrum -from lax.types import ChannelSpec + +from .assembly import uniform_mass_factor @dataclass(frozen=True) @@ -637,53 +639,6 @@ def _phases_grid( ) -def _decouple_closed_channels( - rmatrix: jax.Array, - h_plus: jax.Array, - h_plus_p: jax.Array, - is_open: jax.Array, -) -> jax.Array: - """Fold closed-channel Whittaker boundary conditions into an effective R-matrix.""" - - bloch = _closed_channel_bloch(h_plus, h_plus_p, is_open) - identity: jax.Array = jnp.eye( - rmatrix.shape[0], - dtype=rmatrix.dtype, - ) - correction = identity - rmatrix @ jnp.diag(bloch) - return cast( - jax.Array, - jnp.linalg.solve( - correction.T, - rmatrix.T, - ).T, - ) - - -def _match_rmatrix( - rmatrix: jax.Array, - h_plus: jax.Array, - h_minus: jax.Array, - h_plus_p: jax.Array, - h_minus_p: jax.Array, - is_open: jax.Array, - k: jax.Array, -) -> jax.Array: - """Convert one channel-space R-matrix into the physical S-matrix.""" - - decoupled_r = _decouple_closed_channels(rmatrix, h_plus, h_plus_p, is_open) - projected_r, boundary_slice = _project_open_channels( - decoupled_r, - h_plus, - h_minus, - h_plus_p, - h_minus_p, - is_open, - k, - ) - return smatrix_from_R(projected_r, boundary_slice) - - def _match_one_energy( rmatrix: jax.Array, h_plus: jax.Array, @@ -693,83 +648,22 @@ def _match_one_energy( is_open: jax.Array, k: jax.Array, ) -> jax.Array: - """Match one energy sample after the caller has supplied a concrete ``k`` array.""" - - return _match_rmatrix(rmatrix, h_plus, h_minus, h_plus_p, h_minus_p, is_open, k) - - -def _project_open_channels( - rmatrix: jax.Array, - h_plus: jax.Array, - h_minus: jax.Array, - h_plus_p: jax.Array, - h_minus_p: jax.Array, - is_open: jax.Array, - k: jax.Array, -) -> tuple[jax.Array, BoundaryValues]: - """Project the decoupled R-matrix and boundary values onto the open-channel subspace. - - Closed-channel rows and columns of R are zeroed via an ``is_open`` mask, - and the corresponding Hankel function entries are replaced by 1 in - ``H_plus`` (to avoid divide-by-zero in the matching formula) and by 0 in - ``H_minus``, ``H_plus_p``, and ``H_minus_p``. The shapes remain - ``(N_c, N_c)`` / ``(N_c,)`` so JAX sees static shapes inside JIT. + """Match one channel-space R-matrix sample to the physical S-matrix. - Parameters - ---------- - rmatrix - Full channel-space R-matrix, shape ``(N_c, N_c)``. - h_plus, h_minus, h_plus_p, h_minus_p - Boundary value arrays, shape ``(N_c,)``, complex. - is_open - Boolean mask for open channels, shape ``(N_c,)``. - k - Wave numbers in fm⁻¹, shape ``(N_c,)``. - - Returns - ------- - tuple[jax.Array, BoundaryValues] - Masked R-matrix and masked boundary slice for use in - :func:`smatrix_from_R`. + Re-wraps the per-energy boundary arrays as a :class:`BoundaryValues` + slice and delegates the closed-channel decoupling, open-channel + projection, and matching to :func:`lax.spectral.matching.open_channel_smatrix_from_R`. """ - mask = is_open.astype(rmatrix.dtype) - projected_r = rmatrix * mask[:, None] * mask[None, :] - closed_dtype = h_plus.dtype - ones: jax.Array = jnp.ones_like( - h_plus, - dtype=closed_dtype, - ) - mask_complex = is_open.astype(closed_dtype) - ones_k: jax.Array = jnp.ones_like(k, dtype=k.dtype) - k_values = k * is_open.astype(k.dtype) + ones_k * (1 - is_open.astype(k.dtype)) - boundary_slice = BoundaryValues( - H_plus=h_plus * mask_complex + ones * (1.0 - mask_complex), - H_minus=h_minus * mask_complex, - H_plus_p=h_plus_p * mask_complex, - H_minus_p=h_minus_p * mask_complex, + H_plus=h_plus, + H_minus=h_minus, + H_plus_p=h_plus_p, + H_minus_p=h_minus_p, is_open=is_open, - k=k_values, - ) - return projected_r, boundary_slice - - -def _closed_channel_bloch( - h_plus: jax.Array, - h_plus_p: jax.Array, - is_open: jax.Array, -) -> jax.Array: - """Return the closed-channel Bloch boundary parameter `B_c = H'_c / H_c`.""" - - ratio = h_plus_p / h_plus - closed_mask = jnp.logical_not(is_open) - zeros: jax.Array = jnp.zeros_like(ratio) - return jnp.where( - closed_mask, - ratio, - zeros, + k=k, ) + return open_channel_smatrix_from_R(rmatrix, boundary_slice) def _boundary_wave_numbers(boundary: BoundaryValues) -> jax.Array: @@ -779,14 +673,9 @@ def _boundary_wave_numbers(boundary: BoundaryValues) -> jax.Array: def _uniform_mass_factor(channels: tuple[ChannelSpec, ...]) -> float: - """Return the shared mass factor expected by the MVP observables.""" - - mass_factor = channels[0].mass_factor - for channel in channels[1:]: - if channel.mass_factor != mass_factor: - msg = "The MVP solver path requires a uniform mass_factor across channels." - raise ValueError(msg) - return cast(float, mass_factor) + """Return the shared mass factor expected by the spectral observables.""" + + return uniform_mass_factor(channels, context="spectral observable path") __all__ = [ diff --git a/src/lax/solvers/spectrum.py b/src/lax/solvers/spectrum.py index abb73b9..0dfd1d2 100644 --- a/src/lax/solvers/spectrum.py +++ b/src/lax/solvers/spectrum.py @@ -9,11 +9,10 @@ import jax.numpy as jnp import numpy as np -from lax.boundary._types import Mesh, OperatorMatrices, SpectrumKernel from lax.spectral.types import Spectrum -from lax.types import ChannelSpec, Interaction, Method +from lax.types import ChannelSpec, Interaction, Mesh, Method, OperatorMatrices, SpectrumKernel -from .assembly import assemble_block_hamiltonian, build_Q +from .assembly import assemble_block_hamiltonian, build_Q, uniform_mass_factor @dataclass(frozen=True) @@ -31,36 +30,42 @@ def __call__( self, potential: jax.Array | Interaction, ) -> Spectrum: - """Return the spectral decomposition for one assembled potential. + """Return the spectral decomposition for one potential. Parameters ---------- potential - :class:`~lax.Interaction` object built by ``solver.potential()`` or - ``solver.interaction_from_{block,array,funcs}()``. Must have - ``energy_dependent=False`` (one eigendecomposition per potential). - For energy-dependent workflows, vmap over per-energy blocks:: - - jax.vmap(solver.spectrum)(interaction_list) + :class:`~lax.Interaction` object built by ``solver.local_potential()``/``solver.nonlocal_potential()`` or + ``solver.interaction_from_{block,array,funcs}()``. An + energy-independent interaction yields one :class:`Spectrum`; an + energy-dependent interaction (``energy_dependent=True``) is dispatched + internally over its leading ``(N_E,)`` block axis and yields a batched + :class:`Spectrum` — the caller need not ``jax.vmap`` by hand + (§4.3/§11.1). Returns ------- Spectrum - Eigendecomposition of the Bloch-augmented Hamiltonian. + Eigendecomposition of the Bloch-augmented Hamiltonian (batched over + energy when ``potential.energy_dependent``). """ - if isinstance(potential, Interaction): - if potential.energy_dependent: - raise TypeError( - "spectrum() does not accept energy-dependent Interactions directly. " - "Vmap over per-energy blocks: jax.vmap(solver.spectrum)(interaction.block)." - ) - potential = potential.block + if not isinstance(potential, Interaction): + raise TypeError( + "spectrum() accepts only Interaction objects. " + "Use solver.local_potential()/solver.nonlocal_potential() or solver.interaction_from_block/array/funcs to build one." + ) + if potential.energy_dependent: + return jax.vmap(self._spectrum_one)(potential.block) + return self._spectrum_one(potential.block) + + def _spectrum_one(self, block: jax.Array) -> Spectrum: + """Eigendecompose one ``(M, M)`` assembled block via the chosen backend.""" if self.method == "eigh": return cast( Spectrum, _SPECTRUM_EIGH_JIT( - potential, + block, self.mesh, self.operators, self.channels, @@ -72,7 +77,7 @@ def __call__( return cast( Spectrum, _SPECTRUM_EIG_JIT( - potential, + block, self.mesh, self.operators, self.channels, @@ -121,6 +126,11 @@ def make_spectrum_kernel( JIT-compiled callable: ``kernel(V) → Spectrum``. """ + # The eigh/eig kernels fold a single ℏ²/2μ out of the Hamiltonian + # (H_MeV / m0); validate that assumption here so the kernel cannot be + # built for a multi-μ channel set and silently produce wrong physics, + # even when bound without the observable layer. + uniform_mass_factor(channels, context="spectral eigensolve path") q = build_Q(mesh, channels) return _SpectrumKernel( mesh=mesh, @@ -212,6 +222,7 @@ def numpy_eig(matrix: np.ndarray) -> tuple[np.ndarray, np.ndarray]: numpy_eig, result_shape, complex_hamiltonian, + vmap_method="sequential", ) return cast(tuple[jax.Array, jax.Array], eigensystem) diff --git a/src/lax/spectral/__init__.py b/src/lax/spectral/__init__.py index 2323afd..7fb46ea 100644 --- a/src/lax/spectral/__init__.py +++ b/src/lax/spectral/__init__.py @@ -13,9 +13,10 @@ rmatrix_from_spectrum, wavefunction_internal_from_spectrum, ) -from lax.spectral.types import Spectrum +from lax.spectral.types import BoundaryValues, Spectrum __all__ = [ + "BoundaryValues", "Spectrum", "CoupledChannelParameters", "coupled_channel_parameters_from_S", diff --git a/src/lax/spectral/matching.py b/src/lax/spectral/matching.py index 3aea817..79b4d3b 100644 --- a/src/lax/spectral/matching.py +++ b/src/lax/spectral/matching.py @@ -8,7 +8,7 @@ import jax import jax.numpy as jnp -from lax.boundary._types import BoundaryValues +from lax.spectral.types import BoundaryValues @dataclass(frozen=True) diff --git a/src/lax/spectral/types.py b/src/lax/spectral/types.py index c55f526..a09b518 100644 --- a/src/lax/spectral/types.py +++ b/src/lax/spectral/types.py @@ -1,4 +1,10 @@ -"""Mesh-independent spectral decomposition type. See DESIGN.md §10.1.""" +"""Mesh-independent spectral types. See DESIGN.md §10.1. + +Holds the two pure-data pytrees the spectral submodule operates on: +:class:`Spectrum` (eigenpairs + surface amplitudes) and +:class:`BoundaryValues` (Coulomb/Whittaker matching data). Both depend only +on JAX, keeping ``lax.spectral`` independent of the rest of the package. +""" from __future__ import annotations @@ -7,6 +13,44 @@ import jax +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class BoundaryValues: + """Coulomb and Whittaker boundary values at the channel radius. + + Precomputed at compile time by ``mpmath`` for every ``(energy, channel)`` + pair. Open channels use Coulomb Hankel functions; closed channels use + Whittaker functions that decay exponentially into the barrier. + + Attributes + ---------- + H_plus + Outgoing Coulomb Hankel function ``H⁺ = G + iF`` at ``r = a``, + shape ``(N_E, N_c)``, complex. + H_minus + Incoming Coulomb Hankel function ``H⁻ = G - iF`` at ``r = a``, + shape ``(N_E, N_c)``, complex. + H_plus_p + ``ρ · d/dρ H⁺`` evaluated at ``ρ = ka``, + shape ``(N_E, N_c)``, complex. + H_minus_p + ``ρ · d/dρ H⁻`` evaluated at ``ρ = ka``, + shape ``(N_E, N_c)``, complex. + is_open + Boolean mask: ``True`` for open channels (``E > E_threshold``), + shape ``(N_E, N_c)``. + k + Channel wave numbers ``k_c(E)`` in fm⁻¹, shape ``(N_E, N_c)``. + """ + + H_plus: jax.Array + H_minus: jax.Array + H_plus_p: jax.Array + H_minus_p: jax.Array + is_open: jax.Array + k: jax.Array + + @jax.tree_util.register_dataclass @dataclass(frozen=True) class Spectrum: @@ -41,4 +85,4 @@ class Spectrum: is_hermitian: bool = field(metadata={"static": True}) -__all__ = ["Spectrum"] +__all__ = ["BoundaryValues", "Spectrum"] diff --git a/src/lax/transforms/fourier.py b/src/lax/transforms/fourier.py index 3355494..d2e8c5f 100644 --- a/src/lax/transforms/fourier.py +++ b/src/lax/transforms/fourier.py @@ -10,8 +10,8 @@ import numpy as np import scipy.special as sc -from lax.boundary._types import DoubleFourierTransform, FourierTransform, Mesh, TransformMatrices from lax.meshes._basis_eval import basis_at +from lax.types import DoubleFourierTransform, FourierTransform, Mesh, TransformMatrices @dataclass(frozen=True) diff --git a/src/lax/transforms/grid.py b/src/lax/transforms/grid.py index 1636221..c1d68f9 100644 --- a/src/lax/transforms/grid.py +++ b/src/lax/transforms/grid.py @@ -10,13 +10,13 @@ import jax.numpy as jnp import numpy as np -from lax.boundary._types import ( +from lax.meshes._basis_eval import basis_at +from lax.types import ( FromGridVectorTransform, GridMatrixTransform, GridVectorTransform, Mesh, ) -from lax.meshes._basis_eval import basis_at @dataclass(frozen=True) diff --git a/src/lax/transforms/integration.py b/src/lax/transforms/integration.py index 3707fe7..c79d0a1 100644 --- a/src/lax/transforms/integration.py +++ b/src/lax/transforms/integration.py @@ -8,7 +8,7 @@ import jax import jax.numpy as jnp -from lax.boundary._types import Integrator, Mesh +from lax.types import Integrator, Mesh @dataclass(frozen=True) diff --git a/src/lax/types.py b/src/lax/types.py index ed126a8..f7072fb 100644 --- a/src/lax/types.py +++ b/src/lax/types.py @@ -1,12 +1,18 @@ -"""Public type aliases and user-facing specifications. See ``DESIGN.md §6``.""" +"""Core types: user-facing specs, solver-bundle pytrees, and runtime protocols. See ``DESIGN.md §6``.""" from __future__ import annotations +from collections.abc import Callable from dataclasses import dataclass, field -from typing import Literal +from typing import TYPE_CHECKING, Any, Literal, Protocol import jax +from lax.spectral.types import BoundaryValues + +if TYPE_CHECKING: + from lax.spectral.types import Spectrum + type MeshFamily = Literal["legendre", "laguerre"] type Regularization = Literal[ "x", @@ -68,10 +74,10 @@ class ChannelSpec: Channel threshold in MeV. Assembly code converts it to fm^-2 using ``mass_factor``. mass_factor - Conversion factor ``ℏ² / 2μ`` in MeV·fm². Defaults to ``1.0``, - which is physically meaningless for any real nucleus — always set - this explicitly. Use :func:`lax.constants.hbar2_over_2mu` to - compute the correct value from particle masses in AMU, e.g. + Conversion factor ``ℏ² / 2μ`` in MeV·fm². Required — there is no + default, since any fixed value would be physically meaningless for + an arbitrary nucleus. Use :func:`lax.constants.hbar2_over_2mu` to + compute it from particle masses in AMU, e.g. ``lax.constants.hbar2_over_2mu(1.008665, 1.008665)`` ≈ 41.47 MeV·fm² for nucleon–nucleon systems. """ @@ -120,11 +126,787 @@ def __radd__(self, other: object) -> Interaction: return NotImplemented +type EnergyLike = float | jax.Array + + +class SpectrumKernel(Protocol): + """Callable that maps a potential to its spectral decomposition.""" + + def __call__( + self, + potential: jax.Array | Interaction, + ) -> Spectrum: + """Return the spectral decomposition for one potential. + + Parameters + ---------- + potential + An :class:`~lax.Interaction`. Energy-independent interactions yield + one :class:`Spectrum`; energy-dependent interactions are dispatched + internally over the energy axis and yield a batched :class:`Spectrum`. + Raw potential arrays are rejected at runtime — build an Interaction + first via ``solver.local_potential()`` / ``solver.nonlocal_potential()`` + or ``solver.interaction_from_*``. + + Returns + ------- + Spectrum + Eigendecomposition of the Bloch-augmented Hamiltonian. + """ + ... + + +class RMatrixObservable(Protocol): + """Callable that evaluates the spectral R-matrix at an arbitrary energy.""" + + def __call__(self, spectrum: Spectrum, energy: EnergyLike) -> jax.Array: + """Evaluate the R-matrix at one energy. + + Parameters + ---------- + spectrum + Spectral decomposition produced by ``solver.spectrum(V)``. + energy + Energy in MeV (scalar). + + Returns + ------- + jax.Array + R-matrix, shape ``(N_c, N_c)``. + """ + ... + + +class SpectrumObservable(Protocol): + """Callable that evaluates one observable on the compile-time energy grid.""" + + def __call__(self, spectrum: Spectrum) -> jax.Array: + """Evaluate an observable on the compile-time energy grid. + + Parameters + ---------- + spectrum + Spectral decomposition produced by ``solver.spectrum(V)``. + + Returns + ------- + jax.Array + Observable values on the compile-time grid, shape ``(N_E, ...)``. + """ + ... + + +class SpectrumGridObservable(Protocol): + """Callable for aligned-grid observables from a batched Spectrum.""" + + def __call__(self, spectra: Spectrum) -> jax.Array: + """Evaluate one observable per compile-time energy / Spectrum pair. + + Parameters + ---------- + spectra + Batched ``Spectrum`` produced by calling ``solver.spectrum`` on an + energy-dependent :class:`~lax.Interaction` (internal dispatch), with a + leading batch axis of size ``N_E``. + + Returns + ------- + jax.Array + Observable values, shape ``(N_E, ...)``. + """ + ... + + +class GreenFunctionObservable(Protocol): + """Callable that evaluates the Green's function at an arbitrary energy.""" + + def __call__(self, spectrum: Spectrum, energy: EnergyLike) -> jax.Array: + """Evaluate the Green's function at one energy. + + Parameters + ---------- + spectrum + Spectral decomposition; must have been produced with + ``'greens'`` in ``solvers=``. + energy + Energy in MeV (scalar). + + Returns + ------- + jax.Array + Resolvent ``(H - E/μ)⁻¹``, shape ``(M, M)``. + """ + ... + + +class WavefunctionObservable(Protocol): + """Callable that reconstructs the internal scattering wavefunction.""" + + def __call__(self, spectrum: Spectrum, energy: EnergyLike, source: jax.Array) -> jax.Array: + """Evaluate the internal wavefunction at one energy. + + Parameters + ---------- + spectrum + Spectral decomposition; must have been produced with + ``'wavefunction'`` in ``solvers=``. + energy + Energy in MeV (scalar). + source + Mesh-space driving term, shape ``(N_c · N,)``. Use + :func:`lax.make_wavefunction_source` to build it. + + Returns + ------- + jax.Array + Internal wavefunction coefficients in the mesh basis, shape ``(M,)``. + """ + ... + + +class EigenpairAccessor(Protocol): + """Callable for raw access to eigenvalues and eigenvectors. + + Raises ``RuntimeError`` if eigenvectors were not retained at compile time + (i.e. neither ``'greens'`` nor ``'wavefunction'`` was in ``solvers=``). + """ + + def __call__(self, spectrum: Spectrum) -> tuple[jax.Array, jax.Array]: + """Return the stored eigenvalues and eigenvectors. + + Parameters + ---------- + spectrum + Spectral decomposition with retained eigenvectors. + + Returns + ------- + tuple[jax.Array, jax.Array] + ``(eigenvalues, eigenvectors)`` — shapes ``(M,)`` and ``(M, M)``. + + Raises + ------ + RuntimeError + If eigenvectors were not retained at compile time. + """ + ... + + +class DirectRMatrixKernel(Protocol): + """Callable that computes the direct R-matrix via per-energy linear solves.""" + + def __call__(self, potential: jax.Array | Interaction) -> jax.Array: + """Evaluate the direct R-matrix on the compile-time energy grid. + + Parameters + ---------- + potential + An :class:`~lax.Interaction` object (energy-independent or + energy-dependent; dispatch is handled transparently). Raw potential + arrays are rejected at runtime. + + Returns + ------- + jax.Array + R-matrix on the compile-time grid, shape ``(N_E, N_c, N_c)``. + """ + ... + + +class SMatrixDirectObservable(Protocol): + """Callable that computes the direct S-matrix via per-energy linear solves.""" + + def __call__(self, potential: jax.Array | Interaction) -> jax.Array: + """Evaluate the S-matrix on the compile-time energy grid. + + Parameters + ---------- + potential + An :class:`~lax.Interaction`; raw arrays are rejected at runtime. + + Returns + ------- + jax.Array + S-matrix on the compile-time grid, shape ``(N_E, N_c, N_c)``, complex. + + Notes + ----- + ``potential`` must be an :class:`~lax.Interaction`; raw arrays are + rejected at runtime. + """ + ... + + +class PhasesDirectObservable(Protocol): + """Callable that computes direct phase shifts via per-energy linear solves.""" + + def __call__(self, potential: jax.Array | Interaction) -> jax.Array: + """Evaluate phase shifts on the compile-time energy grid. + + Parameters + ---------- + potential + An :class:`~lax.Interaction`; raw arrays are rejected at runtime. + + Returns + ------- + jax.Array + Phase shifts, shape ``(N_E, N_c)``, in radians. + """ + ... + + +class WavefunctionDirectObservable(Protocol): + """Callable that reconstructs the wavefunction via a direct linear solve.""" + + def __call__( + self, + potential: jax.Array | Interaction, + source: jax.Array, + energy_index: int, + ) -> jax.Array: + """Solve ``C(E_i) ψ = source`` for the internal wavefunction. + + Parameters + ---------- + potential + An :class:`~lax.Interaction`; raw arrays are rejected at runtime. + source + Mesh-space driving term, shape ``(N_c·N,)``. + energy_index + Index into the compile-time energy grid (Python int). + + Returns + ------- + jax.Array + Internal wavefunction coefficient vector, shape ``(N_c·N,)``. + """ + ... + + +class InterpolatorBuilder(Protocol): + """Callable that builds a Padé interpolator over the compile-time grid.""" + + def __call__( + self, + values: jax.Array, + order: tuple[int, int] | None = None, + ) -> Callable[[EnergyLike], jax.Array]: + """Build a Padé interpolator over the solver's compile-time energy grid. + + Parameters + ---------- + values + Observable samples, shape ``(N_E, ...)``. + order + Padé numerator/denominator degrees ``(p, q)`` with ``p + q + 1 == N_E``. + Defaults to the diagonal approximant. + + Returns + ------- + Callable[[EnergyLike], jax.Array] + JIT-compiled interpolant; call it at any energy to evaluate. + """ + ... + + +class GridVectorTransform(Protocol): + """Callable that projects mesh coefficients onto a fine radial grid.""" + + def __call__(self, values: jax.Array) -> jax.Array: + """Project mesh coefficients onto a radial grid. + + Parameters + ---------- + values + Mesh coefficient vector, shape ``(N,)``. + + Returns + ------- + jax.Array + Radial function values, shape ``(M_r,)``. + """ + ... + + +class FromGridVectorTransform(Protocol): + """Callable that projects fine-grid values back onto the mesh basis.""" + + def __call__( + self, + values: jax.Array | Callable[[jax.Array], jax.Array], + ) -> jax.Array: + """Project sampled radial-grid values or a callable profile onto the mesh basis. + + Parameters + ---------- + values + Either a ``(M_r,)`` array of grid samples or a callable + ``f(r_grid) → (M_r,)`` that is evaluated on the compile-time grid. + + Returns + ------- + jax.Array + Mesh coefficient vector, shape ``(N,)``. + """ + ... + + +class GridMatrixTransform(Protocol): + """Callable that projects a mesh-space kernel onto a fine radial grid.""" + + def __call__(self, values: jax.Array) -> jax.Array: + """Project a mesh-space kernel onto a radial grid. + + Parameters + ---------- + values + Mesh kernel matrix, shape ``(N, N)``. + + Returns + ------- + jax.Array + Kernel on the fine grid, shape ``(M_r, M_r)``. + """ + ... + + +class FourierTransform(Protocol): + """Callable that maps mesh coefficients or kernels to momentum space.""" + + def __call__(self, values: jax.Array, channel_index: int = 0) -> jax.Array: + """Project mesh coefficients or a kernel onto a momentum grid. + + Parameters + ---------- + values + Mesh vector ``(N,)`` or kernel ``(N, N)``. + channel_index + Which channel's angular momentum to use for the spherical Bessel + transform. + + Returns + ------- + jax.Array + Momentum-space array, shape ``(M_k,)`` or ``(M_k, M_k)``. + """ + ... + + +class DoubleFourierTransform(Protocol): + """Callable for the double Bessel transform of a mesh-space kernel.""" + + def __call__( + self, + values: jax.Array, + left_channel_index: int = 0, + right_channel_index: int | None = None, + ) -> jax.Array: + """Project a mesh-space kernel onto left/right momentum grids. + + Parameters + ---------- + values + Mesh kernel, shape ``(N, N)``. + left_channel_index + Channel angular momentum for the left (row) transform. + right_channel_index + Channel angular momentum for the right (column) transform. + Defaults to ``left_channel_index``. + + Returns + ------- + jax.Array + Double-transformed kernel, shape ``(M_k, M_k)``. + """ + ... + + +class Integrator(Protocol): + """Callable for norms and expectation values in the mesh basis.""" + + def __call__(self, values: jax.Array, operator: jax.Array | None = None) -> jax.Array: + """Integrate mesh coefficients with an optional operator insertion. + + Parameters + ---------- + values + Mesh coefficient vector, shape ``(N,)``. + operator + Optional ``(N, N)`` operator matrix. When ``None``, computes + the norm ``⟨ψ|ψ⟩``. + + Returns + ------- + jax.Array + Scalar expectation value ``⟨ψ|O|ψ⟩``. + """ + ... + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class PropagationMatrices: + """Precomputed subinterval-propagation matrices for Legendre-x meshes. + + Produced once by :func:`lax.propagate.build_legendre_x_propagation` and + stored inside the :class:`Mesh`. All matrices are in fm⁻² units. + + Attributes + ---------- + n_intervals + Number of subintervals (static). + basis_size_per_interval + Number of Legendre basis functions per subinterval (static). + interval_width + Width of each subinterval in fm (static). + local_nodes + Legendre quadrature nodes on ``(0, 1)``, shape ``(basis_size,)``. + local_weights + Legendre quadrature weights, shape ``(basis_size,)``. + kinetic + Per-interval kinetic matrices, shape ``(n_intervals, basis_size, basis_size)``. + blo0 + Bloch surface-overlap matrix for the left boundary of the first interval. + blo1 + Bloch surface-overlap matrix at the left boundary of subsequent intervals. + blo2 + Bloch surface-overlap matrix at the right boundary of interior intervals. + q1 + Left surface-projector vectors, shape ``(n_intervals, basis_size)``. + q2 + Right surface-projector vectors, shape ``(n_intervals, basis_size)``. + """ + + n_intervals: int = field(metadata={"static": True}) + basis_size_per_interval: int = field(metadata={"static": True}) + interval_width: float = field(metadata={"static": True}) + local_nodes: jax.Array + local_weights: jax.Array + kinetic: jax.Array + blo0: jax.Array + blo1: jax.Array + blo2: jax.Array + q1: jax.Array + q2: jax.Array + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class Mesh: + """Concrete mesh data cached inside a compiled solver. + + Produced by the mesh registry and embedded in the :class:`Solver` at + compile time. Static fields are baked into the JAX JIT cache key; + changing them requires recompilation. + + Attributes + ---------- + family + Mesh family, e.g. ``"legendre"`` or ``"laguerre"`` (static). + regularization + Endpoint regularization, e.g. ``"x"`` or ``"x(1-x)"`` (static). + n + Number of basis functions per channel (static). + scale + Physical scale: channel radius ``a`` in fm for finite-interval + meshes, or the Laguerre scaling factor ``h`` in fm (static). + n_intervals + Number of subintervals for propagated meshes; ``1`` otherwise (static). + basis_size_per_interval + Basis functions per subinterval; equals ``n`` when ``n_intervals == 1`` + (static). + nodes + Canonical mesh nodes on ``(0, 1)`` or ``(0, ∞)``, shape ``(n,)``. + weights + Gauss quadrature weights λ_i, shape ``(n,)``. + radii + Physical radial mesh points r_i = scale · x_i, shape ``(n,)``. + basis_at_boundary + Lagrange basis values φ_j(a) at the channel surface, shape ``(n,)``. + All zeros for semi-infinite (Laguerre) meshes. + propagation + Subinterval propagation matrices, or ``None`` for single-interval meshes. + """ + + family: MeshFamily = field(metadata={"static": True}) + regularization: Regularization = field(metadata={"static": True}) + n: int = field(metadata={"static": True}) + scale: float = field(metadata={"static": True}) + n_intervals: int = field(metadata={"static": True}) + basis_size_per_interval: int = field(metadata={"static": True}) + nodes: jax.Array + weights: jax.Array + radii: jax.Array + basis_at_boundary: jax.Array + propagation: PropagationMatrices | None = None + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class OperatorMatrices: + """Precomputed single-channel operator matrices in fm⁻² units. + + All populated fields are ``(N, N)`` symmetric real matrices in the + Lagrange-mesh basis. Unrequested operators are ``None``. + + Attributes + ---------- + T + Kinetic-energy matrix ``-d²/dr²`` (Laguerre meshes, no Bloch term). + TpL + Bloch-augmented kinetic matrix ``T + L(B=0)``. This is the standard + operator for R-matrix calculations on finite-interval (Legendre) meshes. + T_alpha + Hyperradial kinetic matrix for α-type coordinates (Laguerre meshes with + three-body regularization). + D + First-derivative matrix ``d/dr``. + inv_r + Diagonal ``1/r`` matrix. + inv_r2 + Diagonal ``1/r²`` matrix. + """ + + T: jax.Array | None = None + TpL: jax.Array | None = None + T_alpha: jax.Array | None = None + D: jax.Array | None = None + inv_r: jax.Array | None = None + inv_r2: jax.Array | None = None + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class TransformMatrices: + """Precomputed matrices for radial-grid and momentum-space transforms. + + All fields default to ``None``; only those corresponding to the ``grid`` + and ``momenta`` arguments supplied to :func:`lax.compile` are populated. + + Attributes + ---------- + B_grid + Basis-evaluation matrix ``B[k, j] = f_j(r_k)``, + shape ``(M_r, N)``. Used by ``to_grid_vector`` and ``to_grid_matrix``. + grid_r + Fine radial grid passed to :func:`lax.compile`, shape ``(M_r,)`` in fm. + 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. + momenta + Momentum grid passed to :func:`lax.compile`, shape ``(M_k,)`` in fm⁻¹. + Also accessible as ``solver.momenta``. + """ + + B_grid: jax.Array | None = None + grid_r: jax.Array | None = None + F_momentum: jax.Array | None = None + momenta: jax.Array | None = None + + +@dataclass(frozen=True) +class Solver: + """Compiled solver bundle produced by :func:`lax.compile`. + + Holds all compile-time caches (mesh, operators, boundary values, transform + matrices) alongside JIT-compiled, pickle-safe runtime callables. Call + ``print(solver)`` to see which observables were compiled. + + Attributes + ---------- + mesh + Compiled mesh data (nodes, weights, radii, boundary values). + operators + Precomputed single-channel operator matrices in fm⁻². + channels + Channel definitions baked into the solver at compile time. + energies + Compile-time energy grid in MeV, shape ``(N_E,)``. + boundary + Coulomb/Whittaker boundary values at ``r = a``, or ``None`` if no + energy grid was supplied. + transforms + Precomputed radial-grid and momentum-space transform matrices. + method + Linear-algebra backend: ``"eigh"``, ``"eig"``, or ``"linear_solve"``. + mass_factor_grid + Per-energy ℏ²/2μ values in MeV·fm², shape ``(N_E,)``, or ``None`` + when a constant mass factor is used. Stored here so the aligned-grid + observables can use the correct μ(E) at each energy point. + + **Spectral-path observables** (present when ``method`` is ``"eigh"``/``"eig"``): + + spectrum + ``(V) → Spectrum`` — one eigendecomposition per potential. + rmatrix + ``(spectrum, E) → R(E)`` — R-matrix at any scalar energy. + smatrix + ``(spectrum) → S`` — S-matrix on the compile-time energy grid. + phases + ``(spectrum) → δ`` — phase shifts ``(N_E, N_c)`` in radians. + greens + ``(spectrum, E) → G(E)`` — Green's function; requires ``'greens'`` + in ``solvers=``. + wavefunction + ``(spectrum, E, source) → ψ_int`` — internal wavefunction; requires + ``'wavefunction'`` in ``solvers=``. + eigh + ``(spectrum) → (ε, U)`` — raw eigenpairs; raises if eigenvectors + were not retained. + rmatrix_grid + ``(spectra) → R`` — aligned-grid R for energy-dependent workflows. + smatrix_grid + ``(spectra) → S`` — aligned-grid S. + phases_grid + ``(spectra) → δ`` — aligned-grid phases. + + **Direct-path observables** (present when ``"rmatrix_direct"`` in ``solvers=``): + + rmatrix_direct + ``(V) → R`` — per-energy linear-solve R-matrix on the compile-time grid. + + **Padé interpolation builders** (present whenever ``energies`` was supplied): + + interpolate_rmatrix, interpolate_smatrix, interpolate_phases + ``(samples, order=None) → callable`` — build a Padé interpolant over + the compile-time grid. + + **Transform helpers**: + + to_grid_vector + ``(c) → ψ(r)`` — mesh coefficients to fine radial grid. + from_grid_vector + ``(ψ_or_fn) → c`` — fine grid values back to mesh coefficients. + to_grid_matrix + ``(V) → V(r, r')`` — mesh kernel to fine radial grid. + fourier + ``(c, channel_index=0) → ũ(k)`` — momentum-space transform. + double_fourier_transform + ``(V, ...) → V(p, p')`` — double Bessel transform for kernels. + integrate + ``(c, operator=None) → ⟨ψ|O|ψ⟩`` — norms and expectation values. + """ + + mesh: Mesh + operators: OperatorMatrices + channels: tuple[ChannelSpec, ...] + energies: jax.Array + boundary: BoundaryValues | None + transforms: TransformMatrices + method: Method + mass_factor_grid: jax.Array | None = None + spectrum: SpectrumKernel | None = None + rmatrix: RMatrixObservable | None = None + smatrix: SpectrumObservable | None = None + phases: SpectrumObservable | None = None + greens: GreenFunctionObservable | None = None + wavefunction: WavefunctionObservable | None = None + eigh: EigenpairAccessor | None = None + rmatrix_grid: SpectrumGridObservable | None = None + smatrix_grid: SpectrumGridObservable | None = None + phases_grid: SpectrumGridObservable | None = None + rmatrix_direct: DirectRMatrixKernel | None = None + smatrix_direct: SMatrixDirectObservable | None = None + phases_direct: PhasesDirectObservable | None = None + wavefunction_direct: WavefunctionDirectObservable | None = None + interaction_from_block: Callable[..., Any] | None = None + interaction_from_array: Callable[..., Any] | None = None + interaction_from_funcs: Callable[..., Any] | None = None + local_potential: Callable[..., Any] | None = None + nonlocal_potential: Callable[..., Any] | None = None + interpolate_rmatrix: InterpolatorBuilder | None = None + interpolate_smatrix: InterpolatorBuilder | None = None + interpolate_phases: InterpolatorBuilder | None = None + to_grid_vector: GridVectorTransform | None = None + from_grid_vector: FromGridVectorTransform | None = None + to_grid_matrix: GridMatrixTransform | None = None + fourier: FourierTransform | None = None + double_fourier_transform: DoubleFourierTransform | None = None + integrate: Integrator | None = None + + # ------------------------------------------------------------------ + # Convenience properties + + @property + def grid_r(self) -> jax.Array | None: + """Radial grid passed to :func:`lax.compile`, or ``None``.""" + return self.transforms.grid_r + + @property + def momenta(self) -> jax.Array | None: + """Momentum grid passed to :func:`lax.compile`, or ``None``.""" + return self.transforms.momenta + + # ------------------------------------------------------------------ + # Human-readable repr + + def __repr__(self) -> str: + _observable_names = ( + "spectrum", + "rmatrix", + "smatrix", + "phases", + "greens", + "wavefunction", + "eigh", + "rmatrix_grid", + "smatrix_grid", + "phases_grid", + "rmatrix_direct", + "smatrix_direct", + "phases_direct", + "wavefunction_direct", + ) + _transform_names = ( + "to_grid_vector", + "to_grid_matrix", + "fourier", + "integrate", + ) + 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] + n_e = len(self.energies) + return ( + f"Solver({self.mesh.family}/{self.mesh.regularization} " + f"n={self.mesh.n} scale={self.mesh.scale}fm, " + f"method={self.method}, {n_e} {'energy' if n_e == 1 else 'energies'})\n" + f" observables: {' '.join(live) or 'none'}\n" + f" transforms: {' '.join(transforms) or 'none'}" + ) + + __all__ = [ + "BoundaryValues", "ChannelSpec", + "DirectRMatrixKernel", + "DoubleFourierTransform", + "EigenpairAccessor", + "EnergyLike", + "FourierTransform", + "FromGridVectorTransform", + "GreenFunctionObservable", + "GridMatrixTransform", + "GridVectorTransform", + "Integrator", "Interaction", + "InterpolatorBuilder", + "Mesh", "MeshFamily", "MeshSpec", "Method", + "OperatorMatrices", + "PhasesDirectObservable", + "PropagationMatrices", "Regularization", + "RMatrixObservable", + "SMatrixDirectObservable", + "Solver", + "SpectrumGridObservable", + "SpectrumKernel", + "SpectrumObservable", + "TransformMatrices", + "WavefunctionDirectObservable", + "WavefunctionObservable", ] diff --git a/src/lax/wavefunction.py b/src/lax/wavefunction.py index fdcb0ff..a8e6b98 100644 --- a/src/lax/wavefunction.py +++ b/src/lax/wavefunction.py @@ -14,7 +14,7 @@ import jax.numpy as jnp -from lax.boundary._types import BoundaryValues, Mesh, Solver +from lax.types import BoundaryValues, Mesh, Solver def make_wavefunction_source( @@ -67,7 +67,7 @@ def make_wavefunction_source( ... solvers=("spectrum", "wavefunction"), ... energies=energies, ... ) - >>> V = solver.potential(lambda r1, r2: jnp.zeros_like(r1)) + >>> V = solver.nonlocal_potential(lambda r1, r2: jnp.zeros_like(r1)) >>> spec = solver.spectrum(V) >>> src = lax.make_wavefunction_source(solver, channel_index=0, energy_index=5) >>> psi = solver.wavefunction(spec, energies[5], src) diff --git a/tests/benchmarks/test_alpha_pb_optical.py b/tests/benchmarks/test_alpha_pb_optical.py index 26dc3e1..4b2e028 100644 --- a/tests/benchmarks/test_alpha_pb_optical.py +++ b/tests/benchmarks/test_alpha_pb_optical.py @@ -12,9 +12,14 @@ ALPHA_PB_REFERENCES, SingleChannelCollisionReference, ) +from tests.conftest import LEGACY_COULOMB_E2 pytest.importorskip("jax") +# Reference S-matrices were prepared with the rounded e² = 1.44; use it library-wide +# for this module (boundary Sommerfeld parameter) and in the reference potential below. +pytestmark = pytest.mark.usefixtures("legacy_coulomb_constant") + ALPHA_PB_MASS_FACTOR = 20.736 / (4.0 * 208.0 / (4.0 + 208.0)) @@ -25,7 +30,7 @@ def _optical_potential(r: jax.Array, imag_depth: float) -> jax.Array: radius = 1.1132 * (208.0 ** (1.0 / 3.0) + 4.0 ** (1.0 / 3.0)) diffuseness = 0.5803 woods_saxon = 1.0 / (1.0 + jnp.exp((r - radius) / diffuseness)) - coulomb = 2.0 * 82.0 * 1.44 / r + coulomb = 2.0 * 82.0 * LEGACY_COULOMB_E2 / r return -v0 * woods_saxon - 1.0j * imag_depth * woods_saxon + coulomb @@ -110,7 +115,7 @@ def test_alpha_pb_optical_matches_published_appendix_a( """Published Descouvemont Example 1 collision-matrix values stay visible in the suite.""" solver = _complex_solver(reference, "linear_solve", ("rmatrix_direct",)) - V = solver.potential(lambda r: _optical_potential(r, imag_depth=10.0)) + V = solver.local_potential(lambda r: _optical_potential(r, imag_depth=10.0)) smatrix = _smatrix_from_direct_rmatrix(solver, V)[:, 0, 0] assert np.allclose(smatrix, reference.collision_matrix, atol=1.0e-4, rtol=1.0e-4) @@ -122,7 +127,7 @@ def test_alpha_pb_optical_eig_matches_appendix_a() -> None: reference = ALPHA_PB_REFERENCE_A14_N60_NS1 solver = _complex_solver(reference, "eig", ("spectrum", "smatrix")) - potential = solver.potential(lambda r: _optical_potential(r, imag_depth=10.0)) + potential = solver.local_potential(lambda r: _optical_potential(r, imag_depth=10.0)) assert solver.spectrum is not None assert solver.smatrix is not None @@ -138,7 +143,7 @@ def test_alpha_pb_optical_direct_matches_appendix_a() -> None: reference = ALPHA_PB_REFERENCE_A14_N60_NS1 solver = _complex_solver(reference, "linear_solve", ("rmatrix_direct",)) - V = solver.potential(lambda r: _optical_potential(r, imag_depth=10.0)) + V = solver.local_potential(lambda r: _optical_potential(r, imag_depth=10.0)) smatrix = _smatrix_from_direct_rmatrix(solver, V)[:, 0, 0] assert np.allclose(smatrix, reference.collision_matrix, atol=1.0e-4, rtol=1.0e-4) @@ -153,14 +158,16 @@ def test_alpha_pb_optical_method_paths_agree() -> None: real_direct_solver = _real_solver(reference, "linear_solve", ("rmatrix_direct",)) complex_solver = _complex_solver(reference, "eig", ("spectrum", "smatrix")) complex_direct_solver = _complex_solver(reference, "linear_solve", ("rmatrix_direct",)) - real_V_spectral = real_spectrum_solver.potential( + real_V_spectral = real_spectrum_solver.local_potential( lambda r: jnp.real(_optical_potential(r, imag_depth=0.0)) ) - real_V_direct = real_direct_solver.potential( + real_V_direct = real_direct_solver.local_potential( lambda r: jnp.real(_optical_potential(r, imag_depth=0.0)) ) - complex_V_spectral = complex_solver.potential(lambda r: _optical_potential(r, imag_depth=10.0)) - complex_V_direct = complex_direct_solver.potential( + complex_V_spectral = complex_solver.local_potential( + lambda r: _optical_potential(r, imag_depth=10.0) + ) + complex_V_direct = complex_direct_solver.local_potential( lambda r: _optical_potential(r, imag_depth=10.0) ) diff --git a/tests/benchmarks/test_coupled_closed_channel.py b/tests/benchmarks/test_coupled_closed_channel.py index 21a030b..fa10b7b 100644 --- a/tests/benchmarks/test_coupled_closed_channel.py +++ b/tests/benchmarks/test_coupled_closed_channel.py @@ -7,7 +7,7 @@ import lax as lm from lax.boundary import BoundaryValues -from lax.solvers.observables import _decouple_closed_channels, _project_open_channels +from lax.spectral.matching import _decouple_closed_channels, _project_open_channels pytest.importorskip("jax") @@ -69,12 +69,12 @@ def _toy_interaction(solver: lm.Solver, *, coupled: bool) -> object: A00 = np.array([[1.0, 0.0], [0.0, 0.0]]) A01 = np.array([[0.0, 1.0], [1.0, 0.0]]) A11 = np.array([[0.0, 0.0], [0.0, 1.0]]) - assert solver.potential is not None - interaction = solver.potential(_diagonal_open, coupling=A00) + solver.potential( + assert solver.local_potential is not None + interaction = solver.local_potential(_diagonal_open, coupling=A00) + solver.local_potential( _diagonal_closed, coupling=A11 ) if coupled: - interaction = interaction + solver.potential(_channel_coupling, coupling=A01) + interaction = interaction + solver.local_potential(_channel_coupling, coupling=A01) return interaction @@ -149,8 +149,8 @@ def test_coupled_closed_channel_decoupled_limit_matches_single_channel() -> None coupled_solver = _coupled_solver("eigh", ("spectrum", "smatrix")) single_channel_solver = _single_channel_solver() coupled_V = _toy_interaction(coupled_solver, coupled=False) - assert single_channel_solver.potential is not None - single_channel_V = single_channel_solver.potential(_diagonal_open) + assert single_channel_solver.local_potential is not None + single_channel_V = single_channel_solver.local_potential(_diagonal_open) assert coupled_solver.spectrum is not None assert coupled_solver.smatrix is not None diff --git a/tests/benchmarks/test_descouvemont_closed_channels.py b/tests/benchmarks/test_descouvemont_closed_channels.py index fd8113e..d4867b5 100644 --- a/tests/benchmarks/test_descouvemont_closed_channels.py +++ b/tests/benchmarks/test_descouvemont_closed_channels.py @@ -20,6 +20,10 @@ pytest.importorskip("jax") +# Descouvemont reference data was prepared with the rounded e² = 1.44; use it library-wide +# for this module (boundary Sommerfeld parameter and the rotor-model Coulomb potential). +pytestmark = pytest.mark.usefixtures("legacy_coulomb_constant") + def _solver(reference: CoupledColumnReference, method: str, solvers: tuple[str, ...]) -> lm.Solver: """Compile the Descouvemont Example 4 α + 12C benchmark.""" diff --git a/tests/benchmarks/test_descouvemont_o16_ca44.py b/tests/benchmarks/test_descouvemont_o16_ca44.py index d388250..216cce7 100644 --- a/tests/benchmarks/test_descouvemont_o16_ca44.py +++ b/tests/benchmarks/test_descouvemont_o16_ca44.py @@ -19,6 +19,10 @@ pytest.importorskip("jax") +# Descouvemont reference data was prepared with the rounded e² = 1.44; use it library-wide +# for this module (boundary Sommerfeld parameter and the rotor-model Coulomb potential). +pytestmark = pytest.mark.usefixtures("legacy_coulomb_constant") + def _solver(reference: CoupledColumnReference, method: str, solvers: tuple[str, ...]) -> lm.Solver: """Compile the Descouvemont Example 3 16O + 44Ca benchmark.""" diff --git a/tests/benchmarks/test_hydrogen.py b/tests/benchmarks/test_hydrogen.py index 48ca944..4adf59f 100644 --- a/tests/benchmarks/test_hydrogen.py +++ b/tests/benchmarks/test_hydrogen.py @@ -9,7 +9,7 @@ import scipy.special as sc import lax as lm -from lax.boundary._types import Solver +from lax.types import Solver pytest.importorskip("jax") pytest.importorskip("scipy") @@ -118,8 +118,8 @@ def test_hydrogen_ground_state_laguerre_x() -> None: solver = _hydrogen_solver(0) assert solver.spectrum is not None - assert solver.potential is not None - spectrum = solver.spectrum(solver.potential(lambda r: -1.0 / r)) + assert solver.local_potential is not None + spectrum = solver.spectrum(solver.local_potential(lambda r: -1.0 / r)) ground_state = float(np.asarray(spectrum.eigenvalues)[0]) * solver.channels[0].mass_factor assert abs(ground_state + 0.5) < 1.0e-10 @@ -142,8 +142,8 @@ def test_hydrogen_bound_state_energies( solver = _hydrogen_solver(angular_momentum) assert solver.spectrum is not None - assert solver.potential is not None - spectrum = solver.spectrum(solver.potential(lambda r: -1.0 / r)) + assert solver.local_potential is not None + spectrum = solver.spectrum(solver.local_potential(lambda r: -1.0 / r)) physical_energies = np.asarray(spectrum.eigenvalues) * solver.channels[0].mass_factor expected = np.asarray([-0.5 / (n**2) for n in principal_quantum_numbers], dtype=np.float64) @@ -177,8 +177,8 @@ def test_hydrogen_wavefunctions_match_analytic_radial_forms( assert solver.to_grid_vector is not None assert solver.transforms.grid_r is not None - assert solver.potential is not None - spectrum = solver.spectrum(solver.potential(lambda r: -1.0 / r)) + assert solver.local_potential is not None + spectrum = solver.spectrum(solver.local_potential(lambda r: -1.0 / r)) assert spectrum.eigenvectors is not None eigenvector = np.asarray(spectrum.eigenvectors)[:, state_index] @@ -217,8 +217,8 @@ def test_hydrogen_wavefunctions_match_analytic_momentum_forms( assert solver.fourier is not None assert solver.transforms.momenta is not None - assert solver.potential is not None - spectrum = solver.spectrum(solver.potential(lambda r: -1.0 / r)) + assert solver.local_potential is not None + spectrum = solver.spectrum(solver.local_potential(lambda r: -1.0 / r)) assert spectrum.eigenvectors is not None eigenvector = np.asarray(spectrum.eigenvectors)[:, state_index] @@ -255,8 +255,8 @@ def test_hydrogen_momentum_norm_matches_current_fourier_convention() -> None: assert solver.transforms.grid_r is not None assert solver.transforms.momenta is not None - assert solver.potential is not None - spectrum = solver.spectrum(solver.potential(lambda r: -1.0 / r)) + assert solver.local_potential is not None + spectrum = solver.spectrum(solver.local_potential(lambda r: -1.0 / r)) assert spectrum.eigenvectors is not None eigenvector = jnp.asarray(np.asarray(spectrum.eigenvectors)[:, 0]) diff --git a/tests/benchmarks/test_phase8_meshes.py b/tests/benchmarks/test_phase8_meshes.py index 5870f8a..43adfe2 100644 --- a/tests/benchmarks/test_phase8_meshes.py +++ b/tests/benchmarks/test_phase8_meshes.py @@ -23,8 +23,8 @@ def test_confined_hydrogen_ground_state_legendre_x_one_minus_x() -> None: ) assert solver.spectrum is not None - assert solver.potential is not None - potential = solver.potential(lambda r: -1.0 / r) + assert solver.local_potential is not None + potential = solver.local_potential(lambda r: -1.0 / r) spectrum = solver.spectrum(potential) ground_state = float(np.asarray(spectrum.eigenvalues)[0]) * HBAR2_2MU @@ -45,8 +45,8 @@ def test_harmonic_oscillator_ground_state_modified_laguerre_x2() -> None: ) assert solver.spectrum is not None - assert solver.potential is not None - potential = solver.potential(lambda r: 0.25 * r**2) + assert solver.local_potential is not None + potential = solver.local_potential(lambda r: 0.25 * r**2) spectrum = solver.spectrum(potential) ground_state = float(np.asarray(spectrum.eigenvalues)[0]) diff --git a/tests/benchmarks/test_yamaguchi.py b/tests/benchmarks/test_yamaguchi.py index 950dec4..b38366f 100644 --- a/tests/benchmarks/test_yamaguchi.py +++ b/tests/benchmarks/test_yamaguchi.py @@ -17,7 +17,7 @@ the prototype and the library both recover it at `a=15, N=20`. This test is the keystone benchmark: it exercises the full chain - solver.potential → spectrum → rmatrix → smatrix → phases. + solver.local_potential → spectrum → rmatrix → smatrix → phases. """ import numpy as np @@ -99,7 +99,7 @@ def test_yamaguchi_phase_shifts(reference: YamaguchiReference) -> None: solvers=("spectrum", "phases"), energies=jnp.asarray(reference.energies), ) - potential = solver.potential(_yamaguchi_kernel) + potential = solver.nonlocal_potential(_yamaguchi_kernel) spectrum = solver.spectrum(potential) phases_deg = np.asarray(solver.phases(spectrum))[:, 0] * (180.0 / np.pi) @@ -125,7 +125,7 @@ def test_yamaguchi_phase_shifts_direct_rmatrix(reference: YamaguchiReference) -> energies=reference.energies, method="linear_solve", ) - V = solver.potential(_yamaguchi_kernel) + V = solver.nonlocal_potential(_yamaguchi_kernel) phases_deg = np.asarray(_phase_from_direct_rmatrix(solver, V))[:, 0] * (180.0 / np.pi) assert np.allclose(phases_deg, reference.phases_deg, atol=1.0e-2, rtol=0.0) @@ -149,7 +149,7 @@ def test_yamaguchi_direct_matches_spectral(reference: YamaguchiReference) -> Non solvers=("spectrum", "rmatrix", "phases", "rmatrix_direct"), energies=reference.energies, ) - V = solver.potential(_yamaguchi_kernel) + V = solver.nonlocal_potential(_yamaguchi_kernel) spectrum = solver.spectrum(V) spectral_delta = np.asarray(solver.phases(spectrum))[:, 0] direct_delta = np.asarray(_phase_from_direct_rmatrix(solver, V))[:, 0] @@ -172,7 +172,7 @@ def test_yamaguchi_large_radius_matches_baye_reference(a, n, E, ref_deg, tol): solvers=("spectrum", "phases"), energies=jnp.array([E]), ) - potential = solver.potential(_yamaguchi_kernel) + potential = solver.nonlocal_potential(_yamaguchi_kernel) delta = float(solver.phases(solver.spectrum(potential))[0, 0]) * (180.0 / np.pi) assert abs(delta - ref_deg) < tol, ( @@ -205,8 +205,8 @@ def test_complex_yamaguchi_eig_matches_real_limit(a, n, E): V_is_complex=True, method="eig", ) - real_potential = real_solver.potential(_yamaguchi_kernel) - complex_potential = complex_solver.potential( + real_potential = real_solver.nonlocal_potential(_yamaguchi_kernel) + complex_potential = complex_solver.nonlocal_potential( lambda r1, r2: _complex_yamaguchi_kernel(r1, r2, imag_strength) ) real_phase = float(real_solver.phases(real_solver.spectrum(real_potential))[0, 0]) diff --git a/tests/benchmarks/test_yamaguchi_fourier.py b/tests/benchmarks/test_yamaguchi_fourier.py index 7369c80..fa530b0 100644 --- a/tests/benchmarks/test_yamaguchi_fourier.py +++ b/tests/benchmarks/test_yamaguchi_fourier.py @@ -32,8 +32,8 @@ def yamaguchi_kernel(r1: jax.Array, r2: jax.Array) -> jax.Array: solvers=("spectrum", "wavefunction"), momenta=jnp.linspace(0.1, 2.0, 20), ) - assert solver.potential is not None - potential = solver.potential(yamaguchi_kernel) + assert solver.nonlocal_potential is not None + potential = solver.nonlocal_potential(yamaguchi_kernel) assert solver.spectrum is not None assert solver.fourier is not None diff --git a/tests/conftest.py b/tests/conftest.py index c8a149f..2900be6 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,8 +2,11 @@ from __future__ import annotations +import pytest from hypothesis import HealthCheck, settings +import lax.constants + # Register the CI hypothesis profile. GitHub Actions sets CI=true which makes # hypothesis look for this profile; defining it explicitly here ensures the # settings are applied consistently across all CI runs. @@ -15,3 +18,22 @@ derandomize=True, suppress_health_check=(HealthCheck.too_slow,), ) + +# Conventional rounded e² (MeV·fm) used by the published Descouvemont / optical-model +# benchmark references. The library default (lax.constants.E2 = ALPHA*HBARC ≈ 1.43996) +# is the exact physical value; only the benchmarks that compare against 1.44-prepared +# data override it via the legacy_coulomb_constant fixture below. +LEGACY_COULOMB_E2 = 1.44 + + +@pytest.fixture +def legacy_coulomb_constant(monkeypatch: pytest.MonkeyPatch) -> None: + """Override the Coulomb constant to the rounded 1.44 of the benchmark references. + + Both library Coulomb sites (``boundary.coulomb._sommerfeld`` and + ``models.optical.uniform_sphere_coulomb_potential``) read ``lax.constants.E2`` at + call time, so this single patch covers the boundary Sommerfeld parameter and the + rotor-model Coulomb potential for the duration of a test. + """ + + monkeypatch.setattr(lax.constants, "E2", LEGACY_COULOMB_E2) diff --git a/tests/property/test_autograd.py b/tests/property/test_autograd.py index 9398588..cfc9679 100644 --- a/tests/property/test_autograd.py +++ b/tests/property/test_autograd.py @@ -34,6 +34,17 @@ def _local_gaussian_potential(draw: st.DrawFn) -> jax.Array: return jnp.asarray(V[None, None, :]) # shape (1, 1, N) +def _interaction(V: jax.Array) -> lm.Interaction: + """Wrap a raw ``(1, 1, N)`` local potential as an Interaction for spectrum(). + + The assembled single-channel block is ``diag(V[0, 0])`` — identical to the + assembly the kernel performed internally — so the result (and its gradient + w.r.t. ``V``) is unchanged. + """ + + return lm.Interaction(block=jnp.diag(V[0, 0]), energy_dependent=False) + + @pytest.mark.property @settings(deadline=None) @given(V=_local_gaussian_potential()) @@ -42,7 +53,7 @@ def test_spectrum_eigenvalues_are_differentiable(V: jax.Array) -> None: assert _SOLVER.spectrum is not None def loss(V: jax.Array) -> jax.Array: - return jnp.sum(_SOLVER.spectrum(V).eigenvalues) + return jnp.sum(_SOLVER.spectrum(_interaction(V)).eigenvalues) grad = jax.grad(loss)(V) assert jnp.all(jnp.isfinite(grad)) @@ -57,7 +68,7 @@ def test_smatrix_is_differentiable(V: jax.Array) -> None: assert _SOLVER.smatrix is not None def loss(V: jax.Array) -> jax.Array: - spec = _SOLVER.spectrum(V) + spec = _SOLVER.spectrum(_interaction(V)) return jnp.sum(_SOLVER.smatrix(spec).real) grad = jax.grad(loss)(V) diff --git a/tests/property/test_hermiticity.py b/tests/property/test_hermiticity.py index f8ebd2f..f3754f7 100644 --- a/tests/property/test_hermiticity.py +++ b/tests/property/test_hermiticity.py @@ -35,6 +35,12 @@ def _local_gaussian_potential(draw: st.DrawFn) -> jax.Array: return jnp.asarray(V[None, None, :]) # shape (1, 1, N) +def _interaction(V: jax.Array) -> lm.Interaction: + """Wrap a raw (1, 1, N) local potential as an Interaction (block = diag(V[0, 0])).""" + + return lm.Interaction(block=jnp.diag(V[0, 0]), energy_dependent=False) + + @pytest.mark.property @settings(deadline=None) @given(V=_local_gaussian_potential()) @@ -50,5 +56,5 @@ def test_block_hamiltonian_is_real_symmetric(V: jax.Array) -> None: def test_eigenvalues_are_real_for_real_symmetric_potential(V: jax.Array) -> None: """eigh path returns real eigenvalues for a real symmetric Hamiltonian.""" assert _SOLVER.spectrum is not None - spectrum = _SOLVER.spectrum(V) + spectrum = _SOLVER.spectrum(_interaction(V)) assert jnp.issubdtype(spectrum.eigenvalues.dtype, jnp.floating) diff --git a/tests/property/test_unitarity.py b/tests/property/test_unitarity.py index f31dc28..cb538af 100644 --- a/tests/property/test_unitarity.py +++ b/tests/property/test_unitarity.py @@ -8,8 +8,8 @@ from hypothesis import given, settings import lax as lm -from lax.boundary._types import BoundaryValues from lax.spectral import smatrix_from_R +from lax.types import BoundaryValues pytest.importorskip("jax") @@ -36,6 +36,12 @@ def _local_gaussian_potential(draw: st.DrawFn) -> jax.Array: return jnp.asarray(V[None, None, :]) # shape (1, 1, N) +def _interaction(V: jax.Array) -> lm.Interaction: + """Wrap a raw (1, 1, N) local potential as an Interaction (block = diag(V[0, 0])).""" + + return lm.Interaction(block=jnp.diag(V[0, 0]), energy_dependent=False) + + @pytest.mark.property @settings(deadline=None) @given(V=_local_gaussian_potential()) @@ -43,7 +49,7 @@ def test_smatrix_is_unitary(V: jax.Array) -> None: """S-matrix diagonal elements lie on the unit circle for open channels.""" assert _SOLVER.smatrix is not None assert _SOLVER.spectrum is not None - spectrum = _SOLVER.spectrum(V) + spectrum = _SOLVER.spectrum(_interaction(V)) S = np.asarray(_SOLVER.smatrix(spectrum)) # shape (N_E, 1, 1) assert _SOLVER.boundary is not None for i in range(len(_ENERGIES)): @@ -64,7 +70,7 @@ def test_smatrix_agrees_with_per_energy_rmatrix_path(V: jax.Array) -> None: assert _SOLVER.smatrix is not None assert _SOLVER.boundary is not None - spectrum = _SOLVER.spectrum(V) + spectrum = _SOLVER.spectrum(_interaction(V)) S_grid = np.asarray(_SOLVER.smatrix(spectrum)) # shape (N_E, 1, 1) for i, energy in enumerate(np.asarray(_ENERGIES)): diff --git a/tests/unit/test_energy_dependent_flow.py b/tests/unit/test_energy_dependent_flow.py index 3a96b36..79824e8 100644 --- a/tests/unit/test_energy_dependent_flow.py +++ b/tests/unit/test_energy_dependent_flow.py @@ -13,16 +13,30 @@ HBAR2_2MU = 41.472 -def _make_potential(solver: lm.Solver, energy: jax.Array) -> jax.Array: - """Return a smooth energy-dependent non-local potential.""" +def _spectra(solver: lm.Solver, energies: jax.Array) -> object: + """Build the energy-dependent non-local Interaction and decompose it. + + The bare form factor ``g(r, r'; E)`` (shape ``(N_E, N, N)``) is passed to + ``interaction_from_array``, which owns the Gauss scaling ``√(λ_i λ_j)·a`` and the + coupling kron; ``spectrum`` then dispatches over the energy axis internally. + """ radii = solver.mesh.radii radius_i, radius_j = jnp.meshgrid(radii, radii, indexing="ij") - weights_i, weights_j = jnp.meshgrid(solver.mesh.weights, solver.mesh.weights, indexing="ij") - kernel = ( - -2.0 * 1.25 * (0.23 + 1.25) ** 2 * jnp.exp(-1.25 * (radius_i + radius_j)) + 0.01 * energy - ) * HBAR2_2MU - return (kernel * jnp.sqrt(weights_i * weights_j) * solver.mesh.scale)[None, None, :, :] + + def kernel(energy: jax.Array) -> jax.Array: # bare g(r, r'; E), shape (N, N) + return ( + -2.0 * 1.25 * (0.23 + 1.25) ** 2 * jnp.exp(-1.25 * (radius_i + radius_j)) + + 0.01 * energy + ) * HBAR2_2MU + + g = jax.vmap(kernel)(energies) # (N_E, N, N) + assert solver.interaction_from_array is not None + assert solver.spectrum is not None + interaction = solver.interaction_from_array( + nonlocal_=[(g, np.ones((1, 1)))], energy_dependent=True + ) + return solver.spectrum(interaction) def _manual_smatrix_grid(solver: lm.Solver, spectra: object, energies: jax.Array) -> jax.Array: @@ -72,13 +86,12 @@ def test_energy_dependent_smatrix_grid_matches_manual_flow() -> None: energies=energies, energy_dependent=True, ) - potentials = jax.vmap(lambda energy: _make_potential(solver, energy))(energies) assert solver.spectrum is not None assert solver.smatrix_grid is not None assert solver.phases_grid is not None - spectra = jax.vmap(solver.spectrum)(potentials) + spectra = _spectra(solver, energies) manual_smatrix = _manual_smatrix_grid(solver, spectra, energies) bound_smatrix = solver.smatrix_grid(spectra) bound_phases = solver.phases_grid(spectra) @@ -105,13 +118,12 @@ def test_energy_dependent_fixed_spectrum_semantics_are_unchanged() -> None: energies=energies, energy_dependent=True, ) - potentials = jax.vmap(lambda energy: _make_potential(solver, energy))(energies) assert solver.spectrum is not None assert solver.smatrix is not None assert solver.smatrix_grid is not None - spectra = jax.vmap(solver.spectrum)(potentials) + spectra = _spectra(solver, energies) fixed_spectrum = jax.tree.map(lambda leaf: leaf[0], spectra) fixed_grid = solver.smatrix(fixed_spectrum) manual_fixed_grid = _manual_smatrix_from_fixed_spectrum(solver, fixed_spectrum, energies) @@ -144,13 +156,6 @@ def test_energy_dependent_pade_flow_matches_dense_grid() -> None: energy_dependent=True, ) - sparse_potentials = jax.vmap(lambda energy: _make_potential(sparse_solver, energy))( - sparse_energies - ) - dense_potentials = jax.vmap(lambda energy: _make_potential(dense_solver, energy))( - dense_energies - ) - assert sparse_solver.spectrum is not None assert sparse_solver.smatrix_grid is not None assert sparse_solver.interpolate_smatrix is not None @@ -158,8 +163,8 @@ def test_energy_dependent_pade_flow_matches_dense_grid() -> None: assert dense_solver.spectrum is not None assert dense_solver.smatrix_grid is not None - sparse_spectra = jax.vmap(sparse_solver.spectrum)(sparse_potentials) - dense_spectra = jax.vmap(dense_solver.spectrum)(dense_potentials) + sparse_spectra = _spectra(sparse_solver, sparse_energies) + dense_spectra = _spectra(dense_solver, dense_energies) sparse_s = sparse_solver.smatrix_grid(sparse_spectra) dense_s = dense_solver.smatrix_grid(dense_spectra) @@ -202,20 +207,18 @@ def test_constant_mass_factor_grid_reproduces_scalar_result() -> None: mass_factor_grid=mu_grid, ) - potentials = jax.vmap(lambda e: _make_potential(scalar_solver, e))(energies) - assert scalar_solver.spectrum is not None assert grid_solver.spectrum is not None assert scalar_solver.smatrix_grid is not None assert grid_solver.smatrix_grid is not None # Scalar path: spectrum uses ChannelSpec.mass_factor - scalar_spectra = jax.vmap(scalar_solver.spectrum)(potentials) + scalar_spectra = _spectra(scalar_solver, energies) scalar_phases = scalar_solver.phases_grid(scalar_spectra) # Grid path: mass factor is baked in at compile time; spectrum uses ChannelSpec.mass_factor # (which equals mu_scalar), and phases_grid uses the constant mass_factor_grid boundary. - grid_spectra = jax.vmap(grid_solver.spectrum)(potentials) + grid_spectra = _spectra(grid_solver, energies) grid_phases = grid_solver.phases_grid(grid_spectra) assert np.allclose( @@ -250,7 +253,6 @@ def test_varying_mass_factor_grid_changes_phases() -> None: mass_factor_grid=mu_grid, ) - potentials = jax.vmap(lambda e: _make_potential(const_solver, e))(energies) mu_grid_np = np.asarray(mu_grid) assert const_solver.spectrum is not None @@ -258,13 +260,13 @@ def test_varying_mass_factor_grid_changes_phases() -> None: assert const_solver.phases_grid is not None assert mu_solver.phases_grid is not None - const_spectra = jax.vmap(const_solver.spectrum)(potentials) + const_spectra = _spectra(const_solver, energies) const_phases = const_solver.phases_grid(const_spectra) # Mass factors are baked at compile time; spectrum always uses ChannelSpec.mass_factor. # The energy-dependent mu enters through the compiled boundary values (k_c, η_c) in # phases_grid, so mu_phases differ from const_phases even though the spectra agree. - mu_spectra = jax.vmap(mu_solver.spectrum)(potentials) + mu_spectra = _spectra(mu_solver, energies) mu_phases = mu_solver.phases_grid(mu_spectra) # Phases must differ — the energy-dependent mu changes the boundary matching (k_c, η_c). diff --git a/tests/unit/test_interaction_builders.py b/tests/unit/test_interaction_builders.py index d46e173..8c2d5e5 100644 --- a/tests/unit/test_interaction_builders.py +++ b/tests/unit/test_interaction_builders.py @@ -2,6 +2,7 @@ from __future__ import annotations +import jax import jax.numpy as jnp import numpy as np import pytest @@ -26,6 +27,13 @@ def _make_mesh_channels(): return solver +def _is_block_diagonal(block: jnp.ndarray) -> bool: + """Return True when ``block`` is (numerically) diagonal.""" + + arr = np.asarray(block) + return bool(np.allclose(arr, np.diag(np.diag(arr)))) + + # --------------------------------------------------------------------------- # make_interaction_from_block # --------------------------------------------------------------------------- @@ -239,10 +247,10 @@ def yamaguchi_kernel(r1: jnp.ndarray, r2: jnp.ndarray) -> jnp.ndarray: ) assert solver.rmatrix_direct is not None - assert solver.potential is not None + assert solver.nonlocal_potential is not None - # Build via solver.potential (canonical API) - V = solver.potential(yamaguchi_kernel) + # Build via solver.nonlocal_potential (canonical API) + V = solver.nonlocal_potential(yamaguchi_kernel) r_from_potential = np.asarray(solver.rmatrix_direct(V)) # Build via interaction_from_block (pre-assembled with Gauss scaling applied manually) @@ -254,3 +262,111 @@ def yamaguchi_kernel(r1: jnp.ndarray, r2: jnp.ndarray) -> jnp.ndarray: r_from_block = np.asarray(solver.rmatrix_direct(interaction)) assert np.allclose(r_from_potential, r_from_block, atol=1.0e-12, rtol=1.0e-12) + + +# --------------------------------------------------------------------------- +# Builders inside jax transformations (DESIGN Examples 16.3 / 16.4 patterns) +# --------------------------------------------------------------------------- + + +def _yamaguchi_solver() -> lm.Solver: + return lm.compile( + mesh=lm.MeshSpec("legendre", "x", n=8, scale=8.0), + channels=(lm.ChannelSpec(l=0, threshold=0.0, mass_factor=41.472),), + operators=("T+L",), + solvers=("rmatrix_direct",), + energies=jnp.asarray([0.1, 5.0]), + ) + + +def test_interaction_builder_under_vmap() -> None: + """vmap over potential parameters builds a batched Interaction (Example 16.4).""" + + solver = _yamaguchi_solver() + assert solver.interaction_from_funcs is not None + + def make_interaction(alpha: jax.Array, beta: jax.Array) -> lm.Interaction: + def kernel(r1: jax.Array, r2: jax.Array) -> jax.Array: + return -2.0 * beta * (alpha + beta) ** 2 * jnp.exp(-beta * (r1 + r2)) * 41.472 + + assert solver.interaction_from_funcs is not None + return solver.interaction_from_funcs(nonlocal_=[(kernel, jnp.ones((1, 1)))]) + + alphas = jnp.asarray([0.2, 0.3]) + betas = jnp.asarray([1.4, 1.5]) + batched = jax.vmap(make_interaction)(alphas, betas) + + M = solver.mesh.n + assert batched.block.shape == (2, M, M) + # Each slice equals the eagerly built Interaction for those parameters. + eager = make_interaction(alphas[1], betas[1]) + assert np.allclose(np.asarray(batched.block[1]), np.asarray(eager.block), atol=1e-13) + + +def test_interaction_builder_under_jit() -> None: + """A jitted function may build the Interaction inside the trace (Example 16.3).""" + + solver = _yamaguchi_solver() + + @jax.jit + def build_block(depth: jax.Array) -> jax.Array: + assert solver.interaction_from_funcs is not None + interaction = solver.interaction_from_funcs( + nonlocal_=[(lambda r1, r2: -depth * jnp.exp(-(r1 + r2)), jnp.ones((1, 1)))] + ) + return interaction.block + + block = build_block(jnp.asarray(5.0)) + assert block.shape == (solver.mesh.n, solver.mesh.n) + + +def test_interaction_builder_still_validates_eagerly() -> None: + """Outside transformations, an asymmetric assembled block still raises.""" + + solver = _yamaguchi_solver() + N = solver.mesh.n + asymmetric = jnp.triu(jnp.ones((N, N))) # K(r, r') != K(r', r) + + assert solver.interaction_from_array is not None + with pytest.raises(ValueError, match="not symmetric"): + solver.interaction_from_array(nonlocal_=[(asymmetric, np.array([[1.0]]))]) + + +# --------------------------------------------------------------------------- +# solver.local_potential / solver.nonlocal_potential entry points +# --------------------------------------------------------------------------- + + +def test_local_potential_builds_diagonal_block() -> None: + """`solver.local_potential(fn)` assembles a local term (diagonal sub-block).""" + + solver = _make_mesh_channels() + assert solver.local_potential is not None + + interaction = solver.local_potential(lambda r: -50.0 * jnp.exp(-((r / 2.0) ** 2))) + assert not interaction.energy_dependent + assert _is_block_diagonal(interaction.block), "local_potential produced a non-diagonal block" + + +def test_nonlocal_potential_builds_full_block() -> None: + """`solver.nonlocal_potential(fn)` assembles a nonlocal term (full sub-block).""" + + solver = _make_mesh_channels() + assert solver.nonlocal_potential is not None + + interaction = solver.nonlocal_potential(lambda r, rp: jnp.exp(-((r - rp) ** 2))) + assert not interaction.energy_dependent + assert not _is_block_diagonal(interaction.block), "nonlocal_potential produced a diagonal block" + + +def test_local_potential_energy_dependent_carries_energy_axis() -> None: + """`solver.local_potential(fn(r, E), energy_dependent=True)` carries an (N_E,) axis.""" + + solver = _make_mesh_channels() + assert solver.local_potential is not None + + interaction = solver.local_potential( + lambda r, e: -3.0 * jnp.exp(-((r / 2.0) ** 2)) + 0.02 * e, energy_dependent=True + ) + assert interaction.energy_dependent + assert interaction.block.shape[0] == len(solver.energies) diff --git a/tests/unit/test_solver_direct.py b/tests/unit/test_solver_direct.py index e3c6203..a8f7b69 100644 --- a/tests/unit/test_solver_direct.py +++ b/tests/unit/test_solver_direct.py @@ -76,7 +76,7 @@ def test_compile_exposes_direct_rmatrix_kernel() -> None: assert solver.rmatrix_direct is not None assert solver.smatrix_direct is not None assert solver.phases_direct is not None - assert solver.potential is not None + assert solver.local_potential is not None assert solver.interpolate_rmatrix is not None assert solver.interpolate_smatrix is not None assert solver.interpolate_phases is not None @@ -164,7 +164,7 @@ def yamaguchi_kernel(r1: jax.Array, r2: jax.Array) -> jax.Array: energies=energies, ) - V = solver.potential(yamaguchi_kernel) + V = solver.nonlocal_potential(yamaguchi_kernel) spectrum = solver.spectrum(V) assert solver.rmatrix is not None @@ -192,7 +192,7 @@ def test_direct_rmatrix_grid_matches_manual_per_energy_solve() -> None: energy_dependent=True, ) - interaction = solver.potential(_energy_dep_V, energy_dependent=True) + interaction = solver.local_potential(_energy_dep_V, energy_dependent=True) result = np.asarray(solver.rmatrix_direct(interaction)) expected = [] @@ -230,10 +230,10 @@ def test_rmatrix_direct_propagated_rejects_nonlocal_interaction() -> None: def nonlocal_kernel(r1: jax.Array, r2: jax.Array) -> jax.Array: return -2.0 * jnp.exp(-0.5 * (r1 + r2)) - assert propagated_solver.potential is not None + assert propagated_solver.nonlocal_potential is not None assert propagated_solver.rmatrix_direct is not None - nonlocal_interaction = propagated_solver.potential(nonlocal_kernel) + nonlocal_interaction = propagated_solver.nonlocal_potential(nonlocal_kernel) with pytest.raises(ValueError, match="Non-local propagated"): propagated_solver.rmatrix_direct(nonlocal_interaction) @@ -287,12 +287,12 @@ def test_direct_grid_observables_match_spectral_grid_for_real_energy_dependent_p assert direct_solver.smatrix_direct is not None assert direct_solver.phases_direct is not None # deprecated aligned-grid observables are no longer wired - spectral_interaction = spectral_solver.potential(_energy_dep_V, energy_dependent=True) - direct_interaction = direct_solver.potential(_energy_dep_V, energy_dependent=True) + spectral_interaction = spectral_solver.local_potential(_energy_dep_V, energy_dependent=True) + direct_interaction = direct_solver.local_potential(_energy_dep_V, energy_dependent=True) - # Spectral path: vmap spectrum over the per-energy (M, M) block slices. - # Raw 2D blocks pass through the Interaction check in _SpectrumKernel.__call__. - spectra = jax.vmap(spectral_solver.spectrum)(spectral_interaction.block) + # Spectral path: spectrum() dispatches the energy-dependent Interaction + # internally over the per-energy block axis (no manual vmap needed). + spectra = spectral_solver.spectrum(spectral_interaction) spectral_r = np.asarray(spectral_solver.rmatrix_grid(spectra)) spectral_s = np.asarray(spectral_solver.smatrix_grid(spectra)) @@ -337,8 +337,8 @@ def test_mass_factor_grid_broadcast_scalar_reproduces_uniform() -> None: mass_factor_grid=jnp.full((2,), m), # (N_E,) — broadcasts to (N_E, N_c) ) - interaction_uniform = solver_uniform.potential(_energy_dep_V, energy_dependent=True) - interaction_grid = solver_grid.potential(_energy_dep_V, energy_dependent=True) + interaction_uniform = solver_uniform.local_potential(_energy_dep_V, energy_dependent=True) + interaction_grid = solver_grid.local_potential(_energy_dep_V, energy_dependent=True) r_uniform = np.asarray(solver_uniform.rmatrix_direct(interaction_uniform)) r_grid = np.asarray(solver_grid.rmatrix_direct(interaction_grid)) @@ -371,8 +371,8 @@ def test_mass_factor_grid_2d_reproduces_uniform() -> None: mass_factor_grid=jnp.full((2, 1), m), # explicit (N_E, N_c) shape ) - interaction_uniform = solver_uniform.potential(_energy_dep_V, energy_dependent=True) - interaction_grid = solver_grid.potential(_energy_dep_V, energy_dependent=True) + interaction_uniform = solver_uniform.local_potential(_energy_dep_V, energy_dependent=True) + interaction_grid = solver_grid.local_potential(_energy_dep_V, energy_dependent=True) r_uniform = np.asarray(solver_uniform.rmatrix_direct(interaction_uniform)) r_grid = np.asarray(solver_grid.rmatrix_direct(interaction_grid)) @@ -454,11 +454,11 @@ def V_ch1_fn(r: jax.Array, E: float) -> jax.Array: A0 = np.array([[1.0, 0.0], [0.0, 0.0]]) A1 = np.array([[0.0, 0.0], [0.0, 1.0]]) - V_two = two_ch.potential(V_ch0_fn, coupling=A0, energy_dependent=True) + two_ch.potential( - V_ch1_fn, coupling=A1, energy_dependent=True - ) - V_ch0 = ch0_solver.potential(V_ch0_fn, energy_dependent=True) - V_ch1 = ch1_solver.potential(V_ch1_fn, energy_dependent=True) + V_two = two_ch.local_potential( + V_ch0_fn, coupling=A0, energy_dependent=True + ) + two_ch.local_potential(V_ch1_fn, coupling=A1, energy_dependent=True) + V_ch0 = ch0_solver.local_potential(V_ch0_fn, energy_dependent=True) + V_ch1 = ch1_solver.local_potential(V_ch1_fn, energy_dependent=True) r_two = np.asarray(two_ch.rmatrix_direct(V_two)) # (N_E, 2, 2) r_ch0 = np.asarray(ch0_solver.rmatrix_direct(V_ch0)) # (N_E, 1, 1) @@ -497,9 +497,9 @@ def yamaguchi_kernel(r1: jax.Array, r2: jax.Array) -> jax.Array: assert solver.wavefunction is not None assert solver.wavefunction_direct is not None - assert solver.potential is not None + assert solver.nonlocal_potential is not None - V = solver.potential(yamaguchi_kernel) + V = solver.nonlocal_potential(yamaguchi_kernel) spec = solver.spectrum(V) for energy_index in range(len(energies)): @@ -512,3 +512,33 @@ def yamaguchi_kernel(r1: jax.Array, r2: jax.Array) -> jax.Array: assert np.allclose(psi_spec, psi_dir, atol=1.0e-10, rtol=1.0e-10), ( f"wavefunction_direct mismatch at energy_index={energy_index}" ) + + +def test_wavefunction_under_linear_solve_binds_direct_kernel() -> None: + """Under `method="linear_solve"`, `"wavefunction"` binds the direct kernel, not a spectral path.""" + + solver = lm.compile( + mesh=lm.MeshSpec("legendre", "x", n=10, scale=8.0), + channels=(lm.ChannelSpec(l=0, threshold=0.0, mass_factor=HBAR2_2MU),), + solvers=("rmatrix_direct", "wavefunction"), + method="linear_solve", + energies=jnp.asarray([0.5, 1.0, 2.0]), + ) + assert solver.wavefunction_direct is not None + # No spectral path exists under linear_solve. + assert solver.wavefunction is None + assert solver.spectrum is None + + +def test_wavefunction_only_under_linear_solve_binds_direct_kernel() -> None: + """`solvers=("wavefunction",), method="linear_solve"` binds `wavefunction_direct`, no spectrum.""" + + solver = lm.compile( + mesh=lm.MeshSpec("legendre", "x", n=10, scale=8.0), + channels=(lm.ChannelSpec(l=0, threshold=0.0, mass_factor=HBAR2_2MU),), + solvers=("wavefunction",), + method="linear_solve", + energies=jnp.asarray([0.5, 1.0, 2.0]), + ) + assert solver.wavefunction_direct is not None + assert solver.wavefunction is None diff --git a/tests/unit/test_solver_pickle.py b/tests/unit/test_solver_pickle.py index db5ea98..957e12d 100644 --- a/tests/unit/test_solver_pickle.py +++ b/tests/unit/test_solver_pickle.py @@ -2,7 +2,6 @@ import pickle -import jax import jax.numpy as jnp import numpy as np import pytest @@ -68,7 +67,8 @@ def test_compiled_solver_round_trips_through_pickle() -> None: "rmatrix_direct", "smatrix_direct", "phases_direct", - "potential", + "local_potential", + "nonlocal_potential", "interpolate_rmatrix", "interpolate_smatrix", "interpolate_phases", @@ -105,9 +105,9 @@ def test_compiled_solver_round_trips_through_pickle() -> None: fresh_spectrum = solver.spectrum(interaction) restored_spectrum = restored.spectrum(interaction) - # Vmap over the per-energy block slices (raw 2D arrays pass through spectrum's Interaction check). - fresh_spectra_grid = jax.vmap(solver.spectrum)(interaction_grid.block) - restored_spectra_grid = jax.vmap(restored.spectrum)(interaction_grid.block) + # Energy-dependent Interaction → spectrum() dispatches over the energy axis. + fresh_spectra_grid = solver.spectrum(interaction_grid) + restored_spectra_grid = restored.spectrum(interaction_grid) assert np.allclose(np.asarray(restored.mesh.nodes), np.asarray(solver.mesh.nodes)) assert np.allclose(np.asarray(restored.mesh.weights), np.asarray(solver.mesh.weights)) diff --git a/tests/unit/test_solver_spectrum.py b/tests/unit/test_solver_spectrum.py index 1cd92f7..e4b8799 100644 --- a/tests/unit/test_solver_spectrum.py +++ b/tests/unit/test_solver_spectrum.py @@ -6,7 +6,6 @@ import lax as lm from lax.boundary import BoundaryValues -from lax.boundary._types import OperatorMatrices from lax.meshes.legendre import build_legendre_x from lax.operators import make_interaction_from_array from lax.solvers import ( @@ -16,17 +15,37 @@ make_rmatrix_direct_kernel, make_spectrum_kernel, ) -from lax.solvers.observables import ( +from lax.spectral.matching import ( _closed_channel_bloch, _decouple_closed_channels, _project_open_channels, + smatrix_from_R, ) -from lax.spectral.matching import smatrix_from_R -from lax.types import ChannelSpec, MeshSpec +from lax.types import ChannelSpec, MeshSpec, OperatorMatrices pytest.importorskip("jax") +def _interaction(potential: object) -> lm.Interaction: + """Assemble a raw ``(N_c, N_c, N[, N])`` potential into an Interaction block. + + ``spectrum()`` accepts only :class:`~lax.Interaction` objects; these unit + tests build raw potential arrays directly, so this mirrors the assembly the + kernel previously did internally (diagonal sub-blocks for local terms, full + sub-blocks for non-local terms) and wraps the result. + """ + + pot = jnp.asarray(potential) + n_c = pot.shape[0] + n = pot.shape[2] + block = jnp.zeros((n_c * n, n_c * n), dtype=pot.dtype) + for c in range(n_c): + for cp in range(n_c): + sub = jnp.diag(pot[c, cp]) if pot.ndim == 3 else pot[c, cp] + block = block.at[c * n : (c + 1) * n, cp * n : (cp + 1) * n].set(sub) + return lm.Interaction(block=block, energy_dependent=False) + + def test_build_q_places_boundary_values_in_channel_blocks() -> None: """`build_Q` places the boundary vector in each channel block.""" @@ -71,7 +90,7 @@ def test_make_spectrum_kernel_matches_direct_eigh() -> None: potential = jnp.asarray([[[0.5, 1.0, 1.5, 2.0]]]) kernel = make_spectrum_kernel(mesh, operators, channels, keep_eigenvectors=True) - spectrum = kernel(potential) + spectrum = kernel(_interaction(potential)) # Assembler returns H_MeV; spectrum path divides by m0 before eigh → fm⁻² eigenvalues. m0 = channels[0].mass_factor hamiltonian = np.asarray(assemble_block_hamiltonian(mesh, operators, channels, potential)) / m0 @@ -98,7 +117,7 @@ def test_make_spectrum_kernel_matches_complex_symmetric_eig() -> None: channels, method="eig", keep_eigenvectors=True, - )(potential) + )(_interaction(potential)) # Assembler returns H_MeV; spectrum path divides by m0 before eig → fm⁻² eigenvalues. m0 = channels[0].mass_factor hamiltonian = np.asarray(assemble_block_hamiltonian(mesh, operators, channels, potential)) / m0 @@ -140,7 +159,9 @@ def test_bind_observables_matches_direct_spectral_helpers() -> None: is_open=jnp.asarray([[True]]), k=jnp.ones((1, 1)), ) - spectrum = make_spectrum_kernel(mesh, operators, channels, keep_eigenvectors=True)(potential) + spectrum = make_spectrum_kernel(mesh, operators, channels, keep_eigenvectors=True)( + _interaction(potential) + ) rmatrix, smatrix, phases, greens, wavefunction, eigh = bind_observables( mesh, @@ -181,7 +202,7 @@ def test_compile_binds_wavefunction_observable() -> None: assert solver.spectrum is not None assert solver.wavefunction is not None - spectrum = solver.spectrum(potential) + spectrum = solver.spectrum(_interaction(potential)) values = np.asarray(solver.wavefunction(spectrum, 0.5, jnp.asarray([1.0, 0.0, 0.0, 0.0]))) assert values.shape == (4,) @@ -208,7 +229,7 @@ def test_coupled_channel_rmatrix_matches_direct_solver() -> None: ] ) spectrum_kernel = make_spectrum_kernel(mesh, operators, channels, keep_eigenvectors=True) - spectrum = spectrum_kernel(potential) + spectrum = spectrum_kernel(_interaction(potential)) rmatrix, smatrix, phases, _, _, _ = bind_observables( mesh, channels, @@ -272,7 +293,7 @@ def test_closed_channel_decoupling_matches_direct_bloch_updated_rmatrix() -> Non k=jnp.asarray([[np.sqrt(energy / channels[0].mass_factor), 1.0]]), ) spectrum_kernel = make_spectrum_kernel(mesh, operators, channels, keep_eigenvectors=True) - spectrum = spectrum_kernel(potential) + spectrum = spectrum_kernel(_interaction(potential)) rmatrix, smatrix, _, _, _, _ = bind_observables( mesh, channels, @@ -359,3 +380,28 @@ def test_assemble_block_hamiltonian_requires_t_plus_l() -> None: channels, jnp.asarray([[[0.0, 0.0, 0.0]]]), ) + + +def test_spectrum_rejects_raw_array() -> None: + """`solver.spectrum` accepts only `Interaction` objects, not raw potential arrays.""" + + solver = lm.compile( + mesh=MeshSpec("legendre", "x", n=8, scale=8.0), + channels=(ChannelSpec(l=0, threshold=0.0, mass_factor=41.472),), + solvers=("spectrum",), + ) + assert solver.spectrum is not None + with pytest.raises(TypeError, match="only Interaction"): + solver.spectrum(jnp.zeros((1, 1, 8))) + + +def test_make_spectrum_kernel_rejects_multi_mu() -> None: + """`make_spectrum_kernel` validates the uniform-μ assumption it folds out of H.""" + + mesh, operators = build_legendre_x(n=6, scale=8.0, operators={"T+L"}) + channels = ( + ChannelSpec(l=0, threshold=0.0, mass_factor=2.0), + ChannelSpec(l=0, threshold=0.0, mass_factor=3.0), + ) + with pytest.raises(ValueError, match="uniform mass_factor"): + make_spectrum_kernel(mesh, operators, channels) diff --git a/tests/unit/test_spectral.py b/tests/unit/test_spectral.py index 75288a9..2ed0afb 100644 --- a/tests/unit/test_spectral.py +++ b/tests/unit/test_spectral.py @@ -5,7 +5,6 @@ import pytest import lax as lm -from lax.boundary._types import BoundaryValues from lax.models import interaction_from_reid_np_j1, reid_np_j1_channels from lax.spectral import ( Spectrum, @@ -15,6 +14,7 @@ rmatrix_from_spectrum, smatrix_from_R, ) +from lax.types import BoundaryValues pytest.importorskip("jax") diff --git a/tests/unit/test_transforms_fourier.py b/tests/unit/test_transforms_fourier.py index eff2c24..c59ab26 100644 --- a/tests/unit/test_transforms_fourier.py +++ b/tests/unit/test_transforms_fourier.py @@ -10,9 +10,9 @@ import scipy.special as sc import lax as lm -from lax.boundary._types import Solver from lax.meshes._basis_eval import basis_at from lax.transforms.fourier import compute_F_momentum +from lax.types import Solver pytest.importorskip("jax") pytest.importorskip("scipy") diff --git a/tests/unit/test_transforms_grid.py b/tests/unit/test_transforms_grid.py index d6c5592..71fcc56 100644 --- a/tests/unit/test_transforms_grid.py +++ b/tests/unit/test_transforms_grid.py @@ -143,12 +143,13 @@ def test_legendre_to_grid_preserves_norm_for_bound_state() -> None: potential = jnp.asarray( [[[-2.5, -2.2, -1.8, -1.3, -0.9, -0.6, -0.35, -0.2, -0.1, -0.05, -0.02, -0.01]]] ) + interaction = lm.Interaction(block=jnp.diag(potential[0, 0]), energy_dependent=False) assert solver.spectrum is not None assert solver.to_grid_vector is not None assert solver.transforms.grid_r is not None - spectrum = solver.spectrum(potential) + spectrum = solver.spectrum(interaction) assert spectrum.eigenvectors is not None eigenvector = np.asarray(spectrum.eigenvectors)[:, 0] grid_values = np.asarray(solver.to_grid_vector(jnp.asarray(eigenvector))) From 703fc84967a4cc4ca5be16be18a20970714d8a2e Mon Sep 17 00:00:00 2001 From: beykyle Date: Thu, 11 Jun 2026 14:36:11 -0400 Subject: [PATCH 2/2] update readme and docs --- README.md | 7 ------- docs/api.rst | 8 +++++--- 2 files changed, 5 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index 5ff458d..06a5d86 100644 --- a/README.md +++ b/README.md @@ -106,13 +106,6 @@ solver = lax.compile( - **Laguerre regularizations:** `x`, `modified_x^2` - **Methods:** `eigh`, `eig`, `linear_solve` -Propagation is intentionally narrower than the general API surface: - -- propagated meshes are supported **only** for **local potentials** on the direct - `linear_solve` path; -- propagated meshes are not supported for spectrum-derived observables or - grid/momentum transforms. - ## Running tests ```bash diff --git a/docs/api.rst b/docs/api.rst index 0ca262e..3493716 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -99,11 +99,13 @@ directly — they are accessed through :func:`compile` and the :class:`Solver` b .. autoclass:: lax.boundary.BoundaryValues -.. autoclass:: lax.boundary.Mesh +.. autofunction:: lax.boundary.compute_boundary_values -.. autoclass:: lax.boundary.OperatorMatrices +.. rubric:: lax.types -.. autofunction:: lax.boundary.compute_boundary_values +.. autoclass:: lax.types.Mesh + +.. autoclass:: lax.types.OperatorMatrices .. rubric:: lax.meshes