Skip to content

Docs #4

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 10 commits into from
Jun 11, 2026
Merged

Docs #4

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
e2ad3b5
Switch assembler to symmetric MeV form; divide by m0 in spectrum path
beykyle Jun 10, 2026
f8c780a
Add Interaction type, builders, and unified direct-path API (Phase 10…
beykyle Jun 10, 2026
987b917
Extend mass_factor_grid to (N_E, N_c) for per-channel per-energy μ (P…
beykyle Jun 10, 2026
bda162b
Add Interaction dispatch fix, wavefunction_direct round-trip test, bu…
beykyle Jun 10, 2026
cd16535
change potential API
beykyle Jun 10, 2026
7e0e618
static analysis
beykyle Jun 10, 2026
e626457
remove depr functions from api
beykyle Jun 10, 2026
3149a5a
remove depr code paths
beykyle Jun 11, 2026
253fb84
update reid interaction to use updated solver.potential api cleanly
beykyle Jun 11, 2026
d68baa7
update 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
10 changes: 4 additions & 6 deletions docs/api.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,10 +11,6 @@ containing all JIT-compiled observables.

.. autofunction:: compile

.. autofunction:: assemble_local

.. autofunction:: assemble_nonlocal

.. autofunction:: make_wavefunction_source

.. autoclass:: MeshSpec
Expand Down Expand Up @@ -71,7 +67,9 @@ Reusable interaction models and preset system parameters.

.. autofunction:: lax.models.channels_from_rotor_model

.. autofunction:: lax.models.make_rotor_coupled_optical_potential
.. autofunction:: lax.models.rotor_coupled_optical_potential

.. autofunction:: lax.models.interaction_from_rotor_model

.. autofunction:: lax.models.open_channel_count

Expand All @@ -87,7 +85,7 @@ Reusable interaction models and preset system parameters.

.. autofunction:: lax.models.reid_np_j1_channels

.. autofunction:: lax.models.reid_np_j1_potential
.. autofunction:: lax.models.interaction_from_reid_np_j1

.. autofunction:: lax.models.reid_soft_core_triplet_components

Expand Down
134 changes: 82 additions & 52 deletions examples/alpha_pb_demo.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,9 @@
"CHANNEL_RADIUS = 14.0\n",
"\n",
"\n",
"def optical_potential_parts(r: jnp.ndarray, imag_depth: float) -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]:\n",
"def optical_potential_parts(\n",
" r: jnp.ndarray, imag_depth: float\n",
") -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]:\n",
" v0 = 100.0\n",
" radius = 1.1132 * (208.0 ** (1.0 / 3.0) + 4.0 ** (1.0 / 3.0))\n",
" diffuseness = 0.5803\n",
Expand Down Expand Up @@ -111,10 +113,11 @@
" H_plus_p=solver.boundary.H_plus_p[energy_index],\n",
" H_minus_p=solver.boundary.H_minus_p[energy_index],\n",
" is_open=solver.boundary.is_open[energy_index],\n",
" k=solver.boundary.k[energy_index],\n",
" )\n",
" smatrix = lm.spectral.smatrix_from_R(r_values[energy_index], boundary)\n",
" smatrices.append(np.asarray(smatrix))\n",
" return np.stack(smatrices)\n"
" return np.stack(smatrices)"
]
},
{
Expand Down Expand Up @@ -146,7 +149,9 @@
],
"source": [
"r_plot = np.linspace(0.25, 18.0, 500)\n",
"total, nuclear_real, nuclear_imag, coulomb = optical_potential_parts(jnp.asarray(r_plot), imag_depth=10.0)\n",
"total, nuclear_real, nuclear_imag, coulomb = optical_potential_parts(\n",
" jnp.asarray(r_plot), imag_depth=10.0\n",
")\n",
"\n",
"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",
Expand All @@ -162,7 +167,7 @@
"axes[1].set_xlabel(\"r [fm]\")\n",
"axes[1].set_ylabel(\"MeV\")\n",
"\n",
"fig.tight_layout()\n"
"fig.tight_layout()"
]
},
{
Expand All @@ -183,38 +188,30 @@
},
{
"cell_type": "code",
"execution_count": 3,
"execution_count": null,
"id": "f2d9fbd0",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Energy Appendix A S(E) eig S(E) direct S(E)\n",
" 10.0 1.000000+0.000000j 1.000000+0.000000j 1.000000+0.000000j\n",
" 20.0 1.000000+0.000001j 1.000000+0.000001j 1.000000+0.000001j\n",
" 30.0 0.998930+0.009050j 0.998996+0.008555j 0.998996+0.008555j\n",
" 40.0 0.650810+0.295600j 0.666935+0.297765j 0.666935+0.297765j\n",
" 50.0 0.064367+0.041130j 0.080228+0.044648j 0.080228+0.044648j\n",
"\n",
"max |eig - Appendix A| = 1.627e-02\n",
"max |direct - Appendix A| = 1.627e-02\n",
"max |eig - direct| = 4.148e-12\n"
]
}
],
"outputs": [],
"source": [
"solver_real = real_solver(\"eigh\", (\"spectrum\", \"smatrix\"))\n",
"solver_complex_eig = complex_solver(\"eig\", (\"spectrum\", \"smatrix\"))\n",
"solver_complex_direct = complex_solver(\"linear_solve\", (\"rmatrix_direct\",))\n",
"\n",
"potential_real = lm.assemble_local(solver_real.mesh, lambda r: jnp.real(optical_potential(r, imag_depth=0.0)))\n",
"potential_complex = lm.assemble_local(solver_complex_eig.mesh, lambda r: optical_potential(r, imag_depth=10.0))\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",
" lambda r: optical_potential(r, imag_depth=10.0)\n",
")\n",
"potential_complex_direct = solver_complex_direct.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_complex_eig = np.asarray(solver_complex_eig.smatrix(solver_complex_eig.spectrum(potential_complex)))[:, 0, 0]\n",
"smatrix_complex_direct = smatrix_from_direct_rmatrix(solver_complex_direct, potential_complex)[:, 0, 0]\n",
"smatrix_complex_eig = np.asarray(\n",
" solver_complex_eig.smatrix(solver_complex_eig.spectrum(potential_complex_eig))\n",
")[:, 0, 0]\n",
"smatrix_complex_direct = smatrix_from_direct_rmatrix(\n",
" solver_complex_direct, potential_complex_direct\n",
")[:, 0, 0]\n",
"\n",
"print(\"Energy Appendix A S(E) eig S(E) direct S(E)\")\n",
"for energy, appendix_value, eig_value, direct_value in zip(\n",
Expand All @@ -232,7 +229,9 @@
"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(f\"max |eig - direct| = {np.max(np.abs(smatrix_complex_eig - smatrix_complex_direct)):.3e}\")\n"
"print(\n",
" f\"max |eig - direct| = {np.max(np.abs(smatrix_complex_eig - smatrix_complex_direct)):.3e}\"\n",
")"
]
},
{
Expand All @@ -257,24 +256,54 @@
"\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(OPTICAL_ENERGIES, smatrix_complex_direct.real, \":\", label=\"direct real\", linewidth=2.0)\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(OPTICAL_ENERGIES, smatrix_complex_direct.imag, \":\", label=\"direct imag\", linewidth=2.0)\n",
"axes[0].plot(\n",
" OPTICAL_ENERGIES,\n",
" smatrix_complex_direct.imag,\n",
" \":\",\n",
" label=\"direct imag\",\n",
" linewidth=2.0,\n",
")\n",
"axes[0].set_title(\"Complex optical benchmark vs Appendix A\")\n",
"axes[0].set_xlabel(\"Energy [MeV]\")\n",
"axes[0].set_ylabel(r\"$S_{20}(E)$\")\n",
"axes[0].legend(ncol=2, fontsize=9)\n",
"\n",
"axes[1].plot(OPTICAL_ENERGIES, np.abs(smatrix_real), marker=\"o\", label=\"real spectrum |S|\", linewidth=2.2)\n",
"axes[1].plot(OPTICAL_ENERGIES, np.abs(smatrix_complex_eig), marker=\"s\", label=\"complex eig |S|\", linewidth=2.0)\n",
"axes[1].plot(OPTICAL_ENERGIES, np.abs(smatrix_complex_direct), marker=\"^\", label=\"complex direct |S|\", linewidth=2.0)\n",
"axes[1].plot(\n",
" OPTICAL_ENERGIES,\n",
" np.abs(smatrix_real),\n",
" marker=\"o\",\n",
" label=\"real spectrum |S|\",\n",
" linewidth=2.2,\n",
")\n",
"axes[1].plot(\n",
" OPTICAL_ENERGIES,\n",
" np.abs(smatrix_complex_eig),\n",
" marker=\"s\",\n",
" label=\"complex eig |S|\",\n",
" linewidth=2.0,\n",
")\n",
"axes[1].plot(\n",
" OPTICAL_ENERGIES,\n",
" np.abs(smatrix_complex_direct),\n",
" marker=\"^\",\n",
" label=\"complex direct |S|\",\n",
" linewidth=2.0,\n",
")\n",
"axes[1].set_title(\"Magnitude across the three validated paths\")\n",
"axes[1].set_xlabel(\"Energy [MeV]\")\n",
"axes[1].set_ylabel(r\"$|S_{20}(E)|$\")\n",
"axes[1].legend()\n",
"\n",
"fig.tight_layout()\n"
"fig.tight_layout()"
]
},
{
Expand Down Expand Up @@ -325,43 +354,36 @@
" solvers.append(\n",
" lm.compile(\n",
" mesh=lm.MeshSpec(\"legendre\", \"x\", n=60, scale=CHANNEL_RADIUS),\n",
" channels=(lm.ChannelSpec(l=angular_momentum, threshold=0.0, mass_factor=ALPHA_PB_MASS_FACTOR),),\n",
" channels=(\n",
" lm.ChannelSpec(l=angular_momentum, threshold=0.0, mass_factor=ALPHA_PB_MASS_FACTOR),\n",
" ),\n",
" operators=(\"T+L\",),\n",
" solvers=(\"spectrum\", \"smatrix\", \"phases\"),\n",
" energies=fine_energies,\n",
" V_is_complex=True,\n",
" method=\"eig\",\n",
" z1z2=(2, 82),\n",
" )\n",
" )\n"
" )"
]
},
{
"cell_type": "code",
"execution_count": 6,
"execution_count": null,
"id": "phase-scan-solve",
"metadata": {
"tags": [
"skip-benchmark"
]
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 6.16 s, sys: 60.8 ms, total: 6.22 s\n",
"Wall time: 1.47 s\n"
]
}
],
"outputs": [],
"source": [
"%%time\n",
"phase_curves = {}\n",
"abs_s_curves = {}\n",
"\n",
"for solver, angular_momentum in zip(solvers, partial_waves):\n",
" potential = lm.assemble_local(solver.mesh, lambda r: optical_potential(r, imag_depth=10.0))\n",
" potential = solver.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]))"
Expand Down Expand Up @@ -391,8 +413,16 @@
"source": [
"fig, axes = plt.subplots(1, 2, figsize=(13, 4.8))\n",
"for angular_momentum in partial_waves:\n",
" axes[0].plot(np.asarray(fine_energies), phase_curves[angular_momentum], label=fr\"$\\ell={angular_momentum}$\")\n",
" axes[1].plot(np.asarray(fine_energies), abs_s_curves[angular_momentum], label=fr\"$\\ell={angular_momentum}$\")\n",
" axes[0].plot(\n",
" np.asarray(fine_energies),\n",
" phase_curves[angular_momentum],\n",
" label=rf\"$\\ell={angular_momentum}$\",\n",
" )\n",
" axes[1].plot(\n",
" np.asarray(fine_energies),\n",
" abs_s_curves[angular_momentum],\n",
" label=rf\"$\\ell={angular_momentum}$\",\n",
" )\n",
"\n",
"axes[0].set_title(\"Optical-model phase shifts from the spectrum path\")\n",
"axes[0].set_xlabel(\"Energy [MeV]\")\n",
Expand All @@ -404,7 +434,7 @@
"axes[1].set_ylabel(r\"$|S_\\ell(E)|$\")\n",
"axes[1].legend(ncol=2, fontsize=9)\n",
"\n",
"fig.tight_layout()\n"
"fig.tight_layout()"
]
},
{
Expand Down Expand Up @@ -437,4 +467,4 @@
},
"nbformat": 4,
"nbformat_minor": 5
}
}
Loading
Loading