Merge branch 'main' into fix/ShapeletPolar_dPIEkappa

Author: Jammy2211

autogalaxy/__init__.py

@@ -6,6 +6,9 @@
 
 from autoconf.dictable import from_dict, from_json, output_to_json, to_dict
 from autoarray.dataset import preprocess  # noqa
+from autoarray.dataset.interferometer.w_tilde import (
+    load_curvature_preload_if_compatible,
+)
 from autoarray.dataset.imaging.dataset import Imaging  # noqa
 from autoarray.dataset.interferometer.dataset import Interferometer  # noqa
 from autoarray.dataset.dataset_model import DatasetModel

autogalaxy/aggregator/interferometer/interferometer.py

@@ -53,7 +53,7 @@ def _interferometer_from(
         )
         uv_wavelengths = fit.value(name="dataset")[3].data
 
-        transformer_class = fit.value(name="dataset.transformer_class")
+        transformer_class = fit.value(name="transformer_class")
 
         dataset = aa.Interferometer(
             data=data,

autogalaxy/interferometer/model/analysis.py

@@ -185,5 +185,6 @@ def save_attributes(self, paths: af.DirectoryPaths):
         super().save_attributes(paths=paths)
 
         paths.save_json(
-            "transformer_class", to_dict(self.dataset.transformer.__class__), "dataset"
+            "transformer_class",
+            to_dict(self.dataset.transformer.__class__),
         )

autogalaxy/profiles/mass/total/isothermal.py

@@ -76,9 +76,6 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Calculate the deflection angles on a grid of (y,x) arc-second coordinates.
 
-        For coordinates (0.0, 0.0) the analytic calculation of the deflection angle gives a NaN. Therefore,
-        coordinates at (0.0, 0.0) are shifted slightly to (1.0e-8, 1.0e-8).
-
         Parameters
         ----------
         grid

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