Merge pull request #79 from bandframework/bugfix-depr

Bugfix depr

requirements.txt

@@ -5,7 +5,7 @@ numpy>=1.21.5
 scipy>=1.10.1
 mpmath>=1.2.1
 tqdm>=4.62.3
-jitr>=1.1
+jitr>=2.1.1
 mkdocstrings>=0.22.0
 mkdocstrings-python>=1.1.2
 mkdocs-material>=9.1.18

src/rose/free_solutions.py

@@ -3,25 +3,21 @@
 """
 
 import numpy as np
-from scipy.special import spherical_jn, spherical_yn
 from mpmath import coulombf, coulombg
-from scipy.interpolate import interp1d
-from scipy.misc import derivative
+from scipy.interpolate import UnivariateSpline
 
 
 def F(rho, ell, eta):
     """
     Bessel function of the first kind.
     """
-    # return rho*spherical_jn(ell, rho)
     return np.complex128(coulombf(ell, eta, rho))
 
 
 def G(rho, ell, eta):
     """
     Bessel function of the second kind.
     """
-    # return -rho*spherical_yn(ell, rho)
     return np.complex128(coulombg(ell, eta, rho))
 
 
@@ -39,18 +35,31 @@ def H_minus(rho, ell, eta):
     return G(rho, ell, eta) - 1j * F(rho, ell, eta)
 
 
-def H_plus_prime(rho, ell, eta, dx=1e-6):
+def coulomb_func_deriv(func, s, l, eta):
     """
-    Derivative of the Hankel function (first kind) with respect to rho.
+    Derivative of Coulomb functions F, G, and Coulomb Hankel functions H+ and H-
     """
-    return derivative(lambda z: H_plus(z, ell, eta), rho, dx=dx)
+    # recurrance relations from https://dlmf.nist.gov/33.4
+    # dlmf Eq. 33.4.4
+    R = np.sqrt(1 + eta**2 / (l + 1) ** 2)
+    S = (l + 1) / s + eta / (l + 1)
+    Xl = func(s, l, eta)
+    Xlp = func(s, l + 1, eta)
+    return S * Xl - R * Xlp
 
 
-def H_minus_prime(rho, ell, eta, dx=1e-6):
+def H_plus_prime(s, l, eta):
     """
-    Derivative of the Hankel function (second kind) with respect to rho.
+    Derivative of the Hankel function (first kind) with respect to s
     """
-    return derivative(lambda z: H_minus(z, ell, eta), rho, dx=dx)
+    return coulomb_func_deriv(H_plus, s, l, eta)
+
+
+def H_minus_prime(s, l, eta, dx=1e-6):
+    """
+    Derivative of the Hankel function (second kind) with respect to s.
+    """
+    return coulomb_func_deriv(H_minus, s, l, eta)
 
 
 def phi_free(rho, ell, eta):
@@ -60,14 +69,6 @@ def phi_free(rho, ell, eta):
     return -0.5j * (H_plus(rho, ell, eta) - H_minus(rho, ell, eta))
 
 
-# def phase_shift(u, up, ell, x0):
-#     rl = 1/x0 * (u/up)
-#     return np.log(
-#         (H_minus(x0, ell) - x0*rl*H_minus_prime(x0, ell)) /
-#         (H_plus(x0, ell) - x0*rl*H_plus_prime(x0, ell))
-#     ) / 2j
-
-
 def phase_shift(phi, phi_prime, ell, eta, x0):
     rl = 1 / x0 * (phi / phi_prime)
     return (
@@ -83,8 +84,8 @@ def phase_shift_interp(u, s, ell, eta, x0, dx=1e-6):
     """
     Given the solution, u, on the s grid, return the phase shift (with respect to the free solution).
     """
-    u_func = interp1d(s, u, kind="cubic")
-    rl = 1 / x0 * (u_func(x0) / derivative(u_func, x0, dx=dx))
+    spl = UnivariateSpline(s, u, k=3)
+    rl = 1 / x0 * (spl(x0) / spl.derivative()(x0))
     return (
         np.log(
             (H_minus(x0, ell, eta) - x0 * rl * H_minus_prime(x0, ell, eta))

src/rose/lagrangelegendersolver.py

@@ -15,36 +15,38 @@ def __init__(
         self,
         interaction: Interaction,
         s_0,
-        solver: jitr.RMatrixSolver,
+        solver: jitr.rmatrix.Solver,
     ):
-        l = np.array([interaction.ell])
-        a = np.array([s_0])
         self.s_0 = s_0
         self.domain = [0, s_0]
         self.interaction = interaction
         self.solver = solver
 
         if self.interaction is not None:
             if self.interaction.k_c == 0:
-                self.eta = 0
+                eta = 0
             else:
                 # There is Coulomb, but we haven't (yet) worked out how to emulate
                 # across energies, so we can precompute H+ and H- stuff.
-                self.eta = self.interaction.eta(None)
-
-        self.Hm = H_minus(self.s_0, self.interaction.ell, self.eta)
-        self.Hp = H_plus(self.s_0, self.interaction.ell, self.eta)
-        self.Hmp = H_minus_prime(self.s_0, self.interaction.ell, self.eta)
-        self.Hpp = H_plus_prime(self.s_0, self.interaction.ell, self.eta)
-        self.channels = np.zeros(1, dtype=jitr.channel_dtype)
-        self.channels["weight"] = np.ones(1)
-        self.channels["l"] = l
-        self.channels["a"] = a
-        self.channels["eta"] = self.eta
-        self.channels["Hp"] = np.array([self.Hp])
-        self.channels["Hm"] = np.array([self.Hm])
-        self.channels["Hpp"] = np.array([self.Hpp])
-        self.channels["Hmp"] = np.array([self.Hmp])
+                eta = self.interaction.eta(None)
+
+        self.Hm = H_minus(self.s_0, self.interaction.ell, eta)
+        self.Hp = H_plus(self.s_0, self.interaction.ell, eta)
+        self.Hmp = H_minus_prime(self.s_0, self.interaction.ell, eta)
+        self.Hpp = H_plus_prime(self.s_0, self.interaction.ell, eta)
+
+        self.weight = np.ones(1)
+        self.l = np.array([interaction.ell])
+        self.a = s_0
+        self.eta = np.array([eta])
+        self.couplings = np.array([[0.0]])
+
+        self.asym = jitr.reactions.Asymptotics(
+            np.array([self.Hp]),
+            np.array([self.Hm]),
+            np.array([self.Hpp]),
+            np.array([self.Hmp]),
+        )
 
         # for Energized EIM interactions, E, mu k take up the first 3 spots
         # in the parameter vector, so we offset to account for that
@@ -59,12 +61,9 @@ def __init__(
             self.potential = potential
             self.get_args = self.get_args_neutral
 
-        self.im = jitr.InteractionMatrix(1)
-        self.im.set_local_interaction(self.potential)
-
         # these are always parameter independent - we can precompute them
-        self.basis_boundary = self.solver.precompute_boundaries(a)
-        self.free_matrix = self.solver.free_matrix(a, l)
+        self.basis_boundary = self.solver.precompute_boundaries(self.a)
+        self.free_matrix = self.solver.free_matrix(self.a, self.l)
 
     def get_args_neutral(self, alpha):
         return (
@@ -88,12 +87,17 @@ def clone_for_new_interaction(self, interaction: Interaction):
         return LagrangeRmatrix(interaction, self.s_0, self.solver)
 
     def get_channel_info(self, alpha):
-        self.im.local_args[0, 0] = self.get_args(alpha)
-        self.channels["E"] = self.interaction.E(alpha)
-        self.channels["mu"] = self.interaction.reduced_mass(alpha)
-        self.channels["k"] = self.interaction.momentum(alpha)
+        ch = jitr.reactions.Channels(
+            np.array([self.interaction.E(alpha)]),
+            np.array([self.interaction.momentum(alpha)]),
+            np.array([self.interaction.reduced_mass(alpha)]),
+            self.eta,
+            self.a,
+            self.l,
+            self.couplings,
+        )
 
-        return self.channels, self.im
+        return ch
 
     def phi(
         self,
@@ -103,34 +107,39 @@ def phi(
     ):
         assert s_mesh[-1] <= self.s_0
 
-        ch, im = self.get_channel_info(alpha)
+        ch = self.get_channel_info(alpha)
         R, S, x, uext_prime_boundary = self.solver.solve(
-            im,
             ch,
+            self.asym,
+            local_interaction=self.potential,
+            local_args=self.get_args(alpha),
             free_matrix=self.free_matrix,
             basis_boundary=self.basis_boundary,
             wavefunction=True,
         )
-        return jitr.Wavefunctions(
+        return jitr.reactions.wavefunction.Wavefunctions(
             self.solver,
             x,
             S,
             uext_prime_boundary,
-            self.channels["weight"],
-            jitr.make_channel_data(ch),
+            ch,
+            incoming_weights=self.weight,
         ).uint()[0](s_mesh)
 
     def smatrix(
         self,
         alpha: np.array,
         **kwargs,
     ):
-        ch, im = self.get_channel_info(alpha)
-        R, S, uext_prime_boundary = self.solver.solve(
-            im,
+        ch = self.get_channel_info(alpha)
+        R, S, x, uext_prime_boundary = self.solver.solve(
             ch,
+            self.asym,
+            local_interaction=self.potential,
+            local_args=self.get_args(alpha),
             free_matrix=self.free_matrix,
             basis_boundary=self.basis_boundary,
+            wavefunction=False,
         )
         return S[0, 0]
 
@@ -139,11 +148,15 @@ def rmatrix(
         alpha: np.array,
         **kwargs,
     ):
-        ch, im = self.get_channel_info(alpha)
-        R, S, uext_prime_boundary = self.solver.solve(
-            im,
+        ch = self.get_channel_info(alpha)
+        ch = self.get_channel_info(alpha)
+        R, S, x, uext_prime_boundary = self.solver.solve(
             ch,
+            self.asym,
+            local_interaction=self.potential,
+            local_args=self.get_args(alpha),
             free_matrix=self.free_matrix,
             basis_boundary=self.basis_boundary,
+            wavefunction=False,
         )
         return R[0, 0]

src/rose/scattering_amplitude_emulator.py

@@ -1,6 +1,6 @@
 import pickle
 import numpy as np
-from scipy.special import eval_legendre, gamma
+from scipy.special import eval_legendre, lpmv, gamma
 from tqdm import tqdm
 from numba import njit
 
@@ -10,7 +10,6 @@
 from .schroedinger import SchroedingerEquation
 from .basis import RelativeBasis, CustomBasis, Basis
 from .utility import (
-    eval_assoc_legendre,
     xs_calc_neutral,
     xs_calc_coulomb,
     NucleonNucleusXS,
@@ -243,9 +242,7 @@ def __init__(
         self.angles = angles.copy()
         self.ls = np.arange(self.l_max + 1)[:, np.newaxis]
         self.P_l_costheta = eval_legendre(self.ls, np.cos(self.angles))
-        self.P_1_l_costheta = np.array(
-            [eval_assoc_legendre(l, np.cos(self.angles)) for l in self.ls]
-        )
+        self.P_1_l_costheta = lpmv(1, self.ls, np.cos(self.angles))
 
         # Coulomb scattering amplitude
         if (
@@ -459,13 +456,12 @@ def exact_smatrix_elements(self, alpha):
         Splus = np.zeros(self.l_max, dtype=np.complex128)
         Sminus = np.zeros(self.l_max, dtype=np.complex128)
         Splus[0] = self.rbes[0][0].basis.solver.smatrix(alpha)
-        Sminus[0] = Splus[0]
         for l in range(1, self.l_max):
             Splus[l] = self.rbes[l][0].basis.solver.smatrix(alpha)
             Sminus[l] = self.rbes[l][1].basis.solver.smatrix(alpha)
             if (
-                np.absolute(Splus[l] - 1) < self.Smatrix_abs_tol
-                and np.absolute(Sminus[l] - 1) < self.Smatrix_abs_tol
+                np.absolute(1 - Splus[l]) < self.Smatrix_abs_tol
+                and np.absolute(1 - Sminus[l]) < self.Smatrix_abs_tol
             ):
                 break
 
@@ -566,12 +562,9 @@ def calculate_xs(
             f_c = self.f_c
         else:
             assert np.max(angles) <= np.pi and np.min(angles) >= 0
-            P_l_costheta = np.array(
-                [eval_legendre(l, np.cos(angles)) for l in range(self.l_max)]
-            )
-            P_1_l_costheta = np.array(
-                [eval_assoc_legendre(l, np.cos(angles)) for l in range(self.l_max)]
-            )
+            P_l_costheta = eval_legendre(self.ls, np.cos(angles))
+            P_1_l_costheta = lpmv(1, self.ls, np.cos(angles))
+
             sin2 = np.sin(angles / 2) ** 2
             rutherford = 10 * self.eta**2 / (4 * k**2 * sin2**2)
             f_c = (