Merge pull request #263 from Jammy2211/feature/multipole_jax_fix

Feature/multipole jax fix

Author: Jammy2211

autogalaxy/profiles/mass/total/power_law_multipole.py

@@ -113,8 +113,6 @@ def __init__(
             grid=grid
         )
         """
-        from astropy import units
-
         super().__init__(centre=centre, ell_comps=(0.0, 0.0))
 
         self.m = int(m)
@@ -123,10 +121,40 @@ def __init__(
         self.slope = slope
 
         self.multipole_comps = multipole_comps
-        self.k_m, self.angle_m = convert.multipole_k_m_and_phi_m_from(
-            multipole_comps=multipole_comps, m=m
+
+    def k_m_and_angle_m_from(self, xp=np) -> Tuple[float, float]:
+        """
+        Return the multipole normalization ``k_m`` and orientation angle ``angle_m``.
+
+        The multipole normalization and angle are computed from the multipole component
+        parameters ``(epsilon_1, epsilon_2)`` using
+        :func:`convert.multipole_k_m_and_phi_m_from`. The returned angle is converted
+        from degrees to radians.
+
+        The numerical backend can be selected via the ``xp`` argument, allowing this
+        method to be used with both NumPy and JAX (e.g. inside ``jax.jit``-compiled
+        code).
+
+        Parameters
+        ----------
+        xp
+            Numerical backend module, typically ``numpy`` or ``jax.numpy``.
+
+        Returns
+        -------
+        k_m
+            The multipole normalization.
+        angle_m
+            The multipole orientation angle in radians.
+        """
+        from astropy import units
+
+        k_m, angle_m = convert.multipole_k_m_and_phi_m_from(
+            multipole_comps=self.multipole_comps, m=self.m, xp=xp
         )
-        self.angle_m *= units.deg.to(units.rad)
+        angle_m *= units.deg.to(units.rad)
+
+        return k_m, angle_m
 
     def get_shape_angle(
         self,
@@ -198,15 +226,17 @@ def deflections_yx_2d_from(
         """
         radial_grid, polar_angle_grid = radial_and_angle_grid_from(grid=grid, xp=xp)
 
+        k_m, angle_m = self.k_m_and_angle_m_from(xp=xp)
+
         a_r = (
             -(
                 (3.0 - self.slope)
                 * self.einstein_radius ** (self.slope - 1.0)
                 * radial_grid ** (2.0 - self.slope)
             )
             / (self.m**2.0 - (3.0 - self.slope) ** 2.0)
-            * self.k_m
-            * xp.cos(self.m * (polar_angle_grid - self.angle_m))
+            * k_m
+            * xp.cos(self.m * (polar_angle_grid - angle_m))
         )
 
         a_angle = (
@@ -216,8 +246,8 @@ def deflections_yx_2d_from(
                 * radial_grid ** (2.0 - self.slope)
             )
             / (self.m**2.0 - (3.0 - self.slope) ** 2.0)
-            * self.k_m
-            * xp.sin(self.m * (polar_angle_grid - self.angle_m))
+            * k_m
+            * xp.sin(self.m * (polar_angle_grid - angle_m))
         )
 
         return xp.stack(
@@ -242,13 +272,14 @@ def convergence_2d_from(
             The grid of (y,x) arc-second coordinates the convergence is computed on.
         """
         r, angle = radial_and_angle_grid_from(grid=grid, xp=xp)
+        k_m, angle_m = self.k_m_and_angle_m_from(xp=xp)
 
         return (
             1.0
             / 2.0
             * (self.einstein_radius / r) ** (self.slope - 1)
-            * self.k_m
-            * xp.cos(self.m * (angle - self.angle_m))
+            * k_m
+            * xp.cos(self.m * (angle - angle_m))
         )
 
     @aa.grid_dec.to_array