passed xp everywhere, linked to mge instead mge_numpy

Author: NiekWielders

…galaxy/profiles/mass/abstract/mge_jax.py autogalaxy/profiles/mass/abstract/mge.pyautogalaxy/profiles/mass/abstract/mge_jax.py renamed to autogalaxy/profiles/mass/abstract/mge.py autogalaxy/profiles/mass/abstract/mge_jax.py renamed to autogalaxy/profiles/mass/abstract/mge.py

@@ -1,4 +1,7 @@
 # from .jax_utils import w_f_approx
+import numpy as np
+
+from autogalaxy.profiles.mass.stellar.gaussian import Gaussian
 
 
 class MassProfileMGE:
@@ -13,7 +16,7 @@ def __init__(self):
         pass
 
     @staticmethod
-    def zeta_from(grid, amps, sigmas, axis_ratio):
+    def zeta_from(grid, amps, sigmas, axis_ratio, xp=np):
         """
         The key part to compute the deflection angle of each Gaussian.
         """
@@ -36,8 +39,8 @@ def zeta_from(grid, amps, sigmas, axis_ratio):
 
         # process as one big vectorized calculation
         # could try `jax.lax.scan` instead if this is too much memory
-        w = w_f_approx(inv_sigma_ * z)
-        wq = w_f_approx(inv_sigma_ * zq)
+        w = Gaussian.wofz(inv_sigma_ * z, xp=xp)
+        wq = Gaussian.wofz(inv_sigma_ * zq, xp=xp)
         exp_factor = xp.exp(inv_sigma_**2 * expv)
 
         sigma_func_real = w.imag - exp_factor * wq.imag
@@ -47,15 +50,15 @@ def zeta_from(grid, amps, sigmas, axis_ratio):
         return output_grid.sum(axis=0)
 
     @staticmethod
-    def kesi(p):
+    def kesi(p, xp=np):
         """
         see Eq.(6) of 1906.08263
         """
         n_list = xp.arange(0, 2 * p + 1, 1)
         return (2.0 * p * xp.log(10) / 3.0 + 2.0 * xp.pi * n_list * 1j) ** (0.5)
 
     @staticmethod
-    def eta(p):
+    def eta(p, xp=np):
         """
         see Eq.(6) of 1906.00263
         """
@@ -76,7 +79,7 @@ def decompose_convergence_via_mge(self):
         raise NotImplementedError()
 
     def _decompose_convergence_via_mge(
-        self, func, radii_min, radii_max, func_terms=28, func_gaussians=20
+        self, func, radii_min, radii_max, func_terms=28, func_gaussians=20, xp=np
     ):
         """
 
@@ -117,7 +120,7 @@ def convergence_2d_via_mge_from(self, grid_radii):
         raise NotImplementedError()
 
     def _convergence_2d_via_mge_from(
-        self, grid_radii, func_terms=28, func_gaussians=20
+        self, grid_radii, func_terms=28, func_gaussians=20, xp=np
     ):
         """Calculate the projected convergence at a given set of arc-second gridded coordinates.
 
@@ -137,7 +140,7 @@ def _convergence_2d_via_mge_from(
         return convergence.sum(axis=0)
 
     def _deflections_2d_via_mge_from(
-        self, grid, sigmas_factor=1.0, func_terms=28, func_gaussians=20
+        self, grid, sigmas_factor=1.0, func_terms=28, func_gaussians=20, xp=np
     ):
         axis_ratio = xp.min(xp.array([self.axis_ratio(xp), 0.9999]))
 
@@ -147,7 +150,7 @@ def _deflections_2d_via_mge_from(
         sigmas *= sigmas_factor
 
         angle = self.zeta_from(
-            grid=grid, amps=amps, sigmas=sigmas, axis_ratio=axis_ratio
+            grid=grid, amps=amps, sigmas=sigmas, axis_ratio=axis_ratio, xp=xp
         )
 
         angle *= xp.sqrt((2.0 * xp.pi) / (1.0 - axis_ratio**2.0))

autogalaxy/profiles/mass/dark/abstract.py

@@ -5,7 +5,7 @@
 
 from autogalaxy.profiles.mass.abstract.abstract import MassProfile
 from autogalaxy.cosmology.model import LensingCosmology
-from autogalaxy.profiles.mass.abstract.mge_numpy import (
+from autogalaxy.profiles.mass.abstract.mge import (
     MassProfileMGE,
 )
 
@@ -85,9 +85,9 @@ def convergence_2d_via_mge_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs)
 
         elliptical_radii = self.elliptical_radii_grid_from(grid=grid, xp=xp, **kwargs)
 
-        return self._convergence_2d_via_mge_from(grid_radii=elliptical_radii)
+        return self._convergence_2d_via_mge_from(grid_radii=elliptical_radii, xp=xp)
 
-    def tabulate_integral(self, grid, tabulate_bins, **kwargs):
+    def tabulate_integral(self, grid, tabulate_bins, xp=np, **kwargs):
         """Tabulate an integral over the convergence of deflection potential of a mass profile. This is used in \
         the gNFW profile classes to speed up the integration procedure.
 
@@ -99,15 +99,15 @@ def tabulate_integral(self, grid, tabulate_bins, **kwargs):
             The number of bins to tabulate the inner integral of this profile.
         """
         eta_min = 1.0e-4
-        eta_max = 1.05 * np.max(self.elliptical_radii_grid_from(grid=grid, **kwargs))
+        eta_max = 1.05 * xp.max(self.elliptical_radii_grid_from(grid=grid, xp=xp, **kwargs))
 
-        minimum_log_eta = np.log10(eta_min)
-        maximum_log_eta = np.log10(eta_max)
+        minimum_log_eta = xp.log10(eta_min)
+        maximum_log_eta = xp.log10(eta_max)
         bin_size = (maximum_log_eta - minimum_log_eta) / (tabulate_bins - 1)
 
         return eta_min, eta_max, minimum_log_eta, maximum_log_eta, bin_size
 
-    def decompose_convergence_via_mge(self, **kwargs):
+    def decompose_convergence_via_mge(self, xp=np, **kwargs):
         rho_at_scale_radius = (
             self.kappa_s / self.scale_radius
         )  # density parameter of 3D gNFW
@@ -124,9 +124,9 @@ def gnfw_3d(r):
             )
 
         amplitude_list, sigma_list = self._decompose_convergence_via_mge(
-            func=gnfw_3d, radii_min=radii_min, radii_max=radii_max
+            func=gnfw_3d, radii_min=radii_min, radii_max=radii_max, xp=xp
         )
-        amplitude_list *= np.sqrt(2.0 * np.pi) * sigma_list
+        amplitude_list *= xp.sqrt(2.0 * xp.pi) * sigma_list
         return amplitude_list, sigma_list
 
     def coord_func_f(self, grid_radius: np.ndarray, xp=np) -> np.ndarray:

autogalaxy/profiles/mass/stellar/sersic.py

@@ -5,7 +5,7 @@
 import autoarray as aa
 
 from autogalaxy.profiles.mass.abstract.abstract import MassProfile
-from autogalaxy.profiles.mass.abstract.mge_numpy import (
+from autogalaxy.profiles.mass.abstract.mge import (
     MassProfileMGE,
 )
 from autogalaxy.profiles.mass.abstract.cse import (
@@ -126,7 +126,7 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
     def deflections_2d_via_mge_from(
-        self, grid: aa.type.Grid2DLike, func_terms=28, func_gaussians=20, **kwargs
+        self, grid: aa.type.Grid2DLike, func_terms=28, func_gaussians=20, xp=np, **kwargs
     ):
         """
         Calculate the projected 2D deflection angles from a grid of (y,x) arc second coordinates, by computing and
@@ -146,6 +146,7 @@ def deflections_2d_via_mge_from(
             sigmas_factor=np.sqrt(self.axis_ratio()),
             func_terms=func_terms,
             func_gaussians=func_gaussians,
+            xp=xp
         )
 
     @aa.grid_dec.to_vector_yx
@@ -209,6 +210,7 @@ def convergence_2d_via_mge_from(
             grid_radii=eccentric_radii,
             func_terms=func_terms,
             func_gaussians=func_gaussians,
+            xp=xp
         )
 
     @aa.over_sample
@@ -238,7 +240,7 @@ def convergence_func(self, grid_radius: float) -> float:
 
     @aa.grid_dec.to_array
     def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
-        return np.zeros(shape=grid.shape[0])
+        return xp.zeros(shape=grid.shape[0])
 
     def image_2d_via_radii_from(self, radius: np.ndarray):
         """
@@ -258,7 +260,7 @@ def image_2d_via_radii_from(self, radius: np.ndarray):
         )
 
     def decompose_convergence_via_mge(
-        self, func_terms=28, func_gaussians=20
+        self, func_terms=28, func_gaussians=20, xp=np
     ) -> Tuple[List, List]:
         radii_min = self.effective_radius / 100.0
         radii_max = self.effective_radius * 20.0
@@ -267,7 +269,7 @@ def sersic_2d(r):
             return (
                 self.mass_to_light_ratio
                 * self.intensity
-                * np.exp(
+                * xp.exp(
                     -self.sersic_constant
                     * (((r / self.effective_radius) ** (1.0 / self.sersic_index)) - 1.0)
                 )
@@ -279,6 +281,7 @@ def sersic_2d(r):
             radii_max=radii_max,
             func_terms=func_terms,
             func_gaussians=func_gaussians,
+            xp=xp
         )
 
     def decompose_convergence_via_cse(

autogalaxy/profiles/mass/stellar/sersic_core.py

@@ -103,7 +103,7 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np):
             ),
         )
 
-    def decompose_convergence_via_mge(self):
+    def decompose_convergence_via_mge(self, xp=np):
         radii_min = self.effective_radius / 50.0
         radii_max = self.effective_radius * 20.0
 
@@ -113,7 +113,7 @@ def core_sersic_2D(r):
                 * self.intensity_prime()
                 * (1.0 + (self.radius_break / r) ** self.alpha)
                 ** (self.gamma / self.alpha)
-                * np.exp(
+                * xp.exp(
                     -self.sersic_constant
                     * (
                         (r**self.alpha + self.radius_break**self.alpha)
@@ -124,7 +124,7 @@ def core_sersic_2D(r):
             )
 
         return self._decompose_convergence_via_mge(
-            func=core_sersic_2D, radii_min=radii_min, radii_max=radii_max
+            func=core_sersic_2D, radii_min=radii_min, radii_max=radii_max, xp=xp
         )
 
     def intensity_prime(self, xp=np):

autogalaxy/profiles/mass/stellar/sersic_gradient.py

@@ -150,7 +150,7 @@ def convergence_func(self, grid_radius: float) -> float:
             * self.image_2d_via_radii_from(grid_radius)
         )
 
-    def decompose_convergence_via_mge(self, **kwargs):
+    def decompose_convergence_via_mge(self, xp=np, **kwargs):
         radii_min = self.effective_radius / 100.0
         radii_max = self.effective_radius * 20.0
 
@@ -162,14 +162,14 @@ def sersic_gradient_2D(r):
                     ((self.axis_ratio() * r) / self.effective_radius)
                     ** -self.mass_to_light_gradient
                 )
-                * np.exp(
+                * xp.exp(
                     -self.sersic_constant
                     * (((r / self.effective_radius) ** (1.0 / self.sersic_index)) - 1.0)
                 )
             )
 
         return self._decompose_convergence_via_mge(
-            func=sersic_gradient_2D, radii_min=radii_min, radii_max=radii_max
+            func=sersic_gradient_2D, radii_min=radii_min, radii_max=radii_max, xp=xp
         )
 
     def decompose_convergence_via_cse(