no circular imports

Author: NiekWielders

autogalaxy/profiles/mass/abstract/mge.py

@@ -1,8 +1,92 @@
 # from .jax_utils import w_f_approx
 import numpy as np
 
-from autogalaxy.profiles.mass.stellar.gaussian import Gaussian
 
+def wofz(z, xp=np):
+    """
+    JAX-compatible Faddeeva function w(z) = exp(-z^2) * erfc(-i z)
+    Based on the Poppe–Wijers / Zaghloul–Ali rational approximations.
+    Valid for all complex z. JIT + autodiff safe.
+    """
+
+    z = xp.asarray(z, dtype=xp.complex128)
+    x = xp.real(z)
+    y = xp.imag(z)
+
+    r2 = x * x + y * y
+    y2 = y * y
+    z2 = z * z
+
+    sqrt_pi = xp.asarray(xp.sqrt(xp.pi), dtype=xp.float64)
+    inv_sqrt_pi = xp.asarray(1.0 / sqrt_pi, dtype=xp.float64)
+
+    # ---------- Large-|z| continued fraction ----------
+    r1_s1 = xp.asarray([2.5, 2.0, 1.5, 1.0, 0.5], dtype=xp.float64)
+
+    t = z
+    for c in r1_s1:
+        t = z - c / t
+
+    w_large = 1j * inv_sqrt_pi / t
+
+    # ---------- Region 5 ----------
+    U5 = xp.asarray(
+        [1.320522, 35.7668, 219.031, 1540.787, 3321.990, 36183.31], dtype=xp.float64
+    )
+    V5 = xp.asarray(
+        [1.841439, 61.57037, 364.2191, 2186.181, 9022.228, 24322.84, 32066.6],
+        dtype=xp.float64,
+    )
+
+    t = inv_sqrt_pi
+    for u in U5:
+        t = u + z2 * t
+
+    s = xp.asarray(1.0, dtype=xp.float64)
+    for v in V5:
+        s = v + z2 * s
+
+    w5 = xp.exp(-z2) + 1j * z * t / s
+
+    # ---------- Region 6 ----------
+    U6 = xp.asarray(
+        [5.9126262, 30.180142, 93.15558, 181.92853, 214.38239, 122.60793],
+        dtype=xp.float64,
+    )
+    V6 = xp.asarray(
+        [
+            10.479857,
+            53.992907,
+            170.35400,
+            348.70392,
+            457.33448,
+            352.73063,
+            122.60793,
+        ],
+        dtype=xp.float64,
+    )
+
+    t = inv_sqrt_pi
+    for u in U6:
+        t = u - 1j * z * t
+
+    s = xp.asarray(1.0, dtype=xp.float64)
+    for v in V6:
+        s = v - 1j * z * s
+
+    w6 = t / s
+
+    # ---------- Region logic ----------
+    reg1 = (r2 >= 62.0) | ((r2 >= 30.0) & (r2 < 62.0) & (y2 >= 1e-13))
+    reg2 = ((r2 >= 30) & (r2 < 62) & (y2 < 1e-13)) | (
+            (r2 >= 2.5) & (r2 < 30) & (y2 < 0.072)
+    )
+
+    w = w6
+    w = xp.where(reg2, w5, w)
+    w = xp.where(reg1, w_large, w)
+
+    return w
 
 class MassProfileMGE:
     """
@@ -39,8 +123,8 @@ def zeta_from(grid, amps, sigmas, axis_ratio, xp=np):
 
         # process as one big vectorized calculation
         # could try `jax.lax.scan` instead if this is too much memory
-        w = Gaussian.wofz(inv_sigma_ * z, xp=xp)
-        wq = Gaussian.wofz(inv_sigma_ * zq, xp=xp)
+        w = wofz(inv_sigma_ * z, xp=xp)
+        wq = wofz(inv_sigma_ * zq, xp=xp)
         exp_factor = xp.exp(inv_sigma_**2 * expv)
 
         sigma_func_real = w.imag - exp_factor * wq.imag