Skip to content

zeta_from made jax compatible, zeta_from and wofz in higher precision… #272

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
NiekWielders merged 3 commits into from
Feb 6, 2026
+55 −62
Merged

zeta_from made jax compatible, zeta_from and wofz in higher precision… #272

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
7e2468c
zeta_from made jax compatible, zeta_from and wofz in higher precision…
Feb 6, 2026
b69159f
Merge branch 'main' into feature/gaussian_mass
Feb 6, 2026
0f1fca2
fixed negative y handling
Feb 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
103 changes: 41 additions & 62 deletions autogalaxy/profiles/mass/stellar/gaussian.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -189,45 +189,31 @@ def axis_ratio(self, xp=np):
return xp.where(axis_ratio < 0.9999, axis_ratio, 0.9999)

def zeta_from(self, grid: aa.type.Grid2DLike, xp=np):
q = self.axis_ratio(xp)
q2 = q**2.0
q = xp.asarray(self.axis_ratio(xp), dtype=xp.float64)
q2 = q * q

ind_pos_y = grid.array[:, 0] >= 0
shape_grid = np.shape(grid)
output_grid = np.zeros((shape_grid[0]), dtype=np.complex128)
y = xp.asarray(grid.array[:, 0], dtype=xp.float64)
x = xp.asarray(grid.array[:, 1], dtype=xp.float64)

scale_factor = q / (self.sigma * xp.sqrt(2.0 * (1.0 - q2)))
ind_pos_y = y >= 0

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
scale = q / (xp.asarray(self.sigma, dtype=xp.float64)
* xp.sqrt(xp.asarray(2.0, dtype=xp.float64) * (1.0 - q2)))

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
xs = x * scale
ys = xp.abs(y) * scale

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))
z1 = xs + 1j * ys
z2 = q * xs + 1j * ys / q

if xp == np:
from scipy.special import wofz

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))
)
exp_term = xp.exp(
-(xs * xs) * (1.0 - q2)
- (ys * ys) * (1.0 / q2 - 1.0)
)

else:
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))
)
core = -1j * (self.wofz(z1, xp=xp) - exp_term * self.wofz(z2, xp=xp))

return output_grid
return xp.where(ind_pos_y, core, xp.conj(core))

def wofz(self, z, xp=np):
"""
Expand All @@ -236,71 +222,64 @@ def wofz(self, z, xp=np):
Valid for all complex z. JIT + autodiff safe.
"""

z = xp.asarray(z)
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.sqrt(xp.pi)

# --- Regions 1 to 4 ---
r1_s1 = xp.array([2.5, 2.0, 1.5, 1.0, 0.5])
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 coef in r1_s1:
t = z - coef / t
for c in r1_s1:
t = z - c / t

w_large = 1j / (t * sqrt_pi)
w_large = 1j * inv_sqrt_pi / t

# --- Region 5: special small-imaginary case ---
U5 = xp.array([1.320522, 35.7668, 219.031, 1540.787, 3321.990, 36183.31])
V5 = xp.array(
[1.841439, 61.57037, 364.2191, 2186.181, 9022.228, 24322.84, 32066.6]
)
# ---------- 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 = 1 / sqrt_pi
t = inv_sqrt_pi
for u in U5:
t = u + z2 * t

s = 1.0
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: remaining small-|z| region ---
U6 = xp.array([5.9126262, 30.180142, 93.15558, 181.92853, 214.38239, 122.60793])
V6 = xp.array(
[
10.479857,
53.992907,
170.35400,
348.70392,
457.33448,
352.73063,
122.60793,
]
)
# ---------- 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)

Comment on lines +233 to 266
Copy link

Copilot AI Feb 6, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

wofz() allocates multiple coefficient arrays (r1_s1, U5, V5, U6, V6) on every call. Since zeta_from() calls wofz() twice per evaluation, this can add avoidable overhead for large grids. Consider hoisting these coefficients to module- or class-level constants (e.g. NumPy arrays) and only casting to xp as needed, to reduce per-call allocations.

Copilot uses AI. Check for mistakes.
t = 1 / sqrt_pi
t = inv_sqrt_pi
for u in U6:
t = u - 1j * z * t

s = 1.0
s = xp.asarray(1.0, dtype=xp.float64)
for v in V6:
s = v - 1j * z * s

w6 = t / s

# --- Regions ---
# ---------- 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)
)

# --- Combine regions using pure array logic ---
w = w6
w = xp.where(reg2, w5, w)
w = xp.where(reg1, w_large, w)
Expand Down
14 changes: 14 additions & 0 deletions test_autogalaxy/profiles/mass/stellar/test_gaussian.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,13 @@ def test__deflections_2d_via_analytic_from():
assert deflections[0, 0] == pytest.approx(1.024423, 1.0e-4)
assert deflections[0, 1] == pytest.approx(0.0, abs=1.0e-4)

deflections = mp.deflections_2d_via_analytic_from(
grid=ag.Grid2DIrregular([[-1.0, 0.0]])
)

assert deflections[0, 0] == pytest.approx(-1.024423, 1.0e-4)
assert deflections[0, 1] == pytest.approx(0.0, abs=1.0e-4)

mp = ag.mp.Gaussian(
centre=(0.0, 0.0),
ell_comps=(0.0, 0.111111),
Expand All @@ -37,6 +44,13 @@ def test__deflections_2d_via_analytic_from():
assert deflections[0, 0] == pytest.approx(0.554062, 1.0e-4)
assert deflections[0, 1] == pytest.approx(0.177336, 1.0e-4)

deflections = mp.deflections_2d_via_analytic_from(
grid=ag.Grid2DIrregular([[-0.5, -0.2]])
)

assert deflections[0, 0] == pytest.approx(-0.554062, 1.0e-4)
assert deflections[0, 1] == pytest.approx(-0.177336, 1.0e-4)

mp = ag.mp.Gaussian(
centre=(0.0, 0.0),
ell_comps=(0.0, 0.111111),
Expand Down
Loading