Skip to content

final v1.4 updates #5

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
beykyle merged 2 commits into from
Jun 11, 2026
Merged

final v1.4 updates #5

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3ca8a2b
final v1.4 updates
beykyle Jun 11, 2026
703fc84
update readme and docs
beykyle Jun 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 0 additions & 7 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
8 changes: 5 additions & 3 deletions docs/api.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
74 changes: 55 additions & 19 deletions examples/alpha_pb_demo.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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",
Expand All @@ -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",
Expand All @@ -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",
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand All @@ -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",
")"
Expand All @@ -254,17 +276,25 @@
"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",
" \":\",\n",
" 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",
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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",
" )"
]
},
{
Expand Down
12 changes: 9 additions & 3 deletions examples/descouvemont_closed_channels_demo.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down
15 changes: 11 additions & 4 deletions examples/descouvemont_np_demo.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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",
Expand Down Expand Up @@ -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",
Expand All @@ -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",
"]"
]
},
Expand Down
12 changes: 9 additions & 3 deletions examples/descouvemont_o16_ca44_demo.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down
Loading
Loading