Feature/gaussian mass

autogalaxy/profiles/mass/stellar/gaussian.py

@@ -193,29 +193,36 @@ def zeta_from(self, grid: aa.type.Grid2DLike, xp=np):
         q = self.axis_ratio(xp)
         q2 = q ** 2.0
 
-        y = grid.array[:, 0]
-        x = grid.array[:, 1]
+        ind_pos_y = grid.array[:, 0] >=0
+        shape_grid = np.shape(grid)
+        output_grid = np.zeros((shape_grid[0]), dtype=np.complex128)
 
-        scale = q / (self.sigma * xp.sqrt(2.0 * (1.0 - q2)))
+        scale_factor = q / (self.sigma * xp.sqrt(2.0 * (1.0 - q2)))
 
-        xs = x * scale
-        ys = y * scale
+        xs_0 = grid.array[:, 1][ind_pos_y] * scale_factor
+        ys_0 = grid.array[:, 0][ind_pos_y] * scale_factor
+        xs_1 = grid.array[:, 1][~ind_pos_y] * scale_factor
+        ys_1 = -grid.array[:, 0][~ind_pos_y] * scale_factor
 
-        z1 = xs + 1j * ys
-        z2 = q * xs + 1j * ys / q
+        z1_0 = xs_0 + 1j * ys_0
+        z2_0 = q * xs_0 + 1j * ys_0 / q
+        z1_1 = xs_1 + 1j * ys_1
+        z2_1 = q * xs_1 + 1j * ys_1 / q
 
-        exp_term = xp.exp(-(xs ** 2) * (1.0 - q2) - ys ** 2 * (1.0 / q2 - 1.0))
+        exp_term_0 = xp.exp(-(xs_0 ** 2) * (1.0 - q2) - ys_0 ** 2 * (1.0 / q2 - 1.0))
+        exp_term_1 = xp.exp(-(xs_1 ** 2) * (1.0 - q2) - ys_1 ** 2 * (1.0 / q2 - 1.0))
 
-        if xp is np:
+        if xp == np:
             from scipy.special import wofz
 
-            core = -1j * (wofz(z1) - exp_term * wofz(z2))
+            output_grid[ind_pos_y] = -1j * (wofz(z1_0) - exp_term_0 * wofz(z2_0))
+            output_grid[~ind_pos_y] = xp.conj(-1j * (wofz(z1_1) - exp_term_1 * wofz(z2_1)))
 
         else:
-            core = -1j * (self.wofz(z1, xp=xp) - exp_term * self.wofz(z2, xp=xp))
+            output_grid[ind_pos_y] = -1j * (self.wofz(z1_0, xp=xp) - exp_term_0 * self.wofz(z2_0, xp=xp))
+            output_grid[~ind_pos_y] = xp.conj(-1j * (self.wofz(z1_1, xp=xp) - exp_term_1 * self.wofz(z2_1, xp=xp)))
 
-        # symmetry: zeta(x, -y) = conj(zeta(x, y))
-        return xp.where(y >= 0, core, xp.conj(core))
+        return output_grid
 
     def wofz(self, z, xp=np):
         """

test_autogalaxy/profiles/mass/stellar/test_gaussian.py

@@ -240,3 +240,27 @@ def test__image_2d_via_radii_from__correct_value():
     intensity = mp.image_2d_via_radii_from(grid_radii=ag.ArrayIrregular(3.0))
 
     assert intensity == pytest.approx(0.32465, 1e-2)
+
+def test__wofz():
+    from scipy.special import wofz
+
+    mp = ag.mp.Gaussian(
+        centre=(0.0, 0.0), ell_comps=(0.0, 0.0), intensity=1.0, sigma=2.0
+    )
+
+    wofz_approx_reg_1 = mp.wofz(20.0 + 1j * 0.001)
+    wofz_approx_reg_2 = mp.wofz(2.0 + 1j * 0.001)
+    wofz_approx_reg_3 = mp.wofz(1.0 + 1j * 0.001)
+
+    assert wofz_approx_reg_1 == pytest.approx(wofz(20.0 + 1j * 0.001), 1e-4)
+    assert wofz_approx_reg_2 == pytest.approx(wofz(2.0 + 1j * 0.001), 1e-4)
+    assert wofz_approx_reg_3 == pytest.approx(wofz(1.0 + 1j * 0.001), 1e-4)
+
+    wofz_approx_reg_1 = mp.wofz(7.0 + 1j * 0.1)
+    wofz_approx_reg_2 = mp.wofz(7.0 + 1j * 1e-11)
+    wofz_approx_reg_3 = mp.wofz(2.0 + 1j * 1.0)
+
+    assert wofz_approx_reg_1 == pytest.approx(wofz(7.0 + 1j * 0.1), 1e-4)
+    assert wofz_approx_reg_2 == pytest.approx(wofz(7.0 + 1j * 1e-11), 1e-4)
+    assert wofz_approx_reg_3 == pytest.approx(wofz(2.0 + 1j * 1.0), 1e-4)
+