-
Notifications
You must be signed in to change notification settings - Fork 16
Fix/ShapeletPolar_dPIEkappa #259
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.
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
Changes from 12 commits
cea025a
007ebdb
182eb57
630d736
cd299a0
f589484
846ca38
c711bce
cdea5ae
7e606d3
90604bb
fb7463f
802c0e5
9fcde1c
676940c
898142b
9656fa0
4c7a22f
ddc41e6
4121217
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change | ||||
|---|---|---|---|---|---|---|
|
|
@@ -3,20 +3,73 @@ | |||||
|
|
||||||
| import autoarray as aa | ||||||
|
|
||||||
|
|
||||||
| from autogalaxy.profiles.light.decorators import ( | ||||||
| check_operated_only, | ||||||
| ) | ||||||
| from autogalaxy.profiles.light.standard.shapelets.abstract import AbstractShapelet | ||||||
|
|
||||||
| import jax.numpy as jnp | ||||||
| from jax.scipy.special import gammaln | ||||||
|
|
||||||
| def genlaguerre_jax(n, alpha, x): | ||||||
| """ | ||||||
| Generalized (associated) Laguerre polynomial L_n^alpha(x) | ||||||
| calculated using the explicit summation formula, optimized for JAX vectorization. | ||||||
|
|
||||||
| Parameters: | ||||||
| n (int): Degree of the polynomial (static Python integer). | ||||||
| alpha (Numeric): Parameter alpha > -1. | ||||||
| x (Array): Input array (evaluation points). | ||||||
| """ | ||||||
| # 0. Input Validation (Requires static Python int n) | ||||||
| if not isinstance(n, int) or n < 0: | ||||||
| # Use Python's math.isnan/isinf check if n is float, otherwise type error | ||||||
| raise ValueError(f"Degree n must be a non-negative Python integer (static), got {n}.") | ||||||
|
|
||||||
| # Base Case L0 | ||||||
| if n == 0: | ||||||
| return jnp.ones_like(x) | ||||||
|
|
||||||
| # 1. Generate k values for summation range [0, 1, 2, ..., n] | ||||||
| k_values = jnp.arange(n + 1) # (n+1,) | ||||||
|
|
||||||
| # 2. Reshape inputs for broadcasting (x: (M, 1), k: (1, n+1)) | ||||||
| x_expanded = jnp.expand_dims(x, axis=-1) | ||||||
| k_values_expanded = jnp.expand_dims(k_values, axis=0) | ||||||
|
|
||||||
| # --- A. Binomial Factor (BF) Calculation --- | ||||||
| # BF = exp( log( (n+alpha)! / ((n-k)! * (alpha+k)!) ) ) | ||||||
|
|
||||||
| log_N_plus_alpha_fact = gammaln(n + alpha + 1) | ||||||
|
|
||||||
| log_BF_k = ( | ||||||
| log_N_plus_alpha_fact | ||||||
| - gammaln(n - k_values + 1) # log( (n-k)! ) | ||||||
| - gammaln(alpha + k_values + 1) # log( (alpha+k)! ) | ||||||
| ) | ||||||
|
|
||||||
| BF_k = jnp.exp(log_BF_k) # Shape: (n+1,) | ||||||
|
|
||||||
| # --- B. Term Factor (TF) Calculation --- | ||||||
| # TF = (-x)^k / k! | ||||||
|
|
||||||
| # Note: jnp.math.gamma(k_values + 1) is equivalent to k! in log-gamma space | ||||||
| TF_k = jnp.power(-x_expanded, k_values_expanded) / jnp.exp(gammaln(k_values_expanded + 1)) | ||||||
| # TF_k Shape: (M, n+1) | ||||||
|
|
||||||
| # --- C. Final Summation --- | ||||||
| # Sum over the last axis (axis=1), which corresponds to k | ||||||
| # BF_k broadcasts over the M dimension of TF_k | ||||||
| return jnp.sum(BF_k * TF_k, axis=1) | ||||||
|
Comment on lines
+13
to
+69
|
||||||
|
|
||||||
| class ShapeletPolar(AbstractShapelet): | ||||||
| def __init__( | ||||||
| self, | ||||||
| n: int, | ||||||
| m: int, | ||||||
| centre: Tuple[float, float] = (0.0, 0.0), | ||||||
| ell_comps: Tuple[float, float] = (0.0, 0.0), | ||||||
| q: float = 1.0, | ||||||
| phi: float = 0.0, | ||||||
| intensity: float = 1.0, | ||||||
| beta: float = 1.0, | ||||||
| ): | ||||||
|
|
@@ -39,20 +92,25 @@ def __init__( | |||||
| The m order of the shapelets basis function in the x-direction. | ||||||
| centre | ||||||
| The (y,x) arc-second coordinates of the profile (shapelet) centre. | ||||||
| ell_comps | ||||||
| The first and second ellipticity components of the elliptical coordinate system. | ||||||
| q | ||||||
| The axis-ratio of the elliptical coordinate system, where a perfect circle has q=1.0. | ||||||
| phi | ||||||
| The position angle (in degrees) of the elliptical coordinate system, measured counter-clockwise from the | ||||||
| positive x-axis. | ||||||
| intensity | ||||||
| Overall intensity normalisation of the light profile (units are dimensionless and derived from the data | ||||||
| the light profile's image is compared too, which is expected to be electrons per second). | ||||||
| beta | ||||||
| The characteristic length scale of the shapelet basis function, defined in arc-seconds. | ||||||
| """ | ||||||
|
|
||||||
| self.n = n | ||||||
| self.m = m | ||||||
| self.n = int(n) | ||||||
| self.m = int(m) | ||||||
| self.phi = float(phi) | ||||||
| self.q = float(q) | ||||||
|
||||||
|
|
||||||
| super().__init__( | ||||||
| centre=centre, ell_comps=ell_comps, beta=beta, intensity=intensity | ||||||
| centre=centre, beta=beta, intensity=intensity | ||||||
| ) | ||||||
|
Comment on lines
116
to
118
|
||||||
|
|
||||||
| @property | ||||||
|
|
@@ -62,7 +120,6 @@ def coefficient_tag(self) -> str: | |||||
| @aa.over_sample | ||||||
| @aa.grid_dec.to_array | ||||||
| @check_operated_only | ||||||
| @aa.grid_dec.transform | ||||||
| def image_2d_from( | ||||||
| self, | ||||||
| grid: aa.type.Grid2DLike, | ||||||
|
|
@@ -86,11 +143,11 @@ def image_2d_from( | |||||
| image | ||||||
| The image of the Polar Shapelet evaluated at every (y,x) coordinate on the transformed grid. | ||||||
| """ | ||||||
| from scipy.special import genlaguerre | ||||||
| from jax.scipy.special import factorial | ||||||
|
|
||||||
| laguerre = genlaguerre(n=(self.n - xp.abs(self.m)) / 2.0, alpha=xp.abs(self.m)) | ||||||
|
|
||||||
| grid = aa.util.geometry.transform_grid_2d_to_reference_frame( | ||||||
| grid_2d=grid.array, centre=self.centre, angle=self.phi, xp=xp | ||||||
|
||||||
| grid_2d=grid.array, centre=self.centre, angle=self.phi, xp=xp | |
| grid_2d=grid, centre=self.centre, angle=self.phi, xp=xp |
Copilot
AI
Jan 17, 2026
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The manual transformation uses self.phi directly as the rotation angle, but this doesn't account for the elliptical coordinate system transformations that the parent class handles. The removed @aa.grid_dec.transform decorator would have applied the proper transformations including centre translation, rotation, and elliptical scaling based on ell_comps. The manual approach bypasses this and may produce incorrect results.
Copilot
AI
Jan 17, 2026
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The axis ratio self.q is applied inconsistently in the calculation. The code divides grid[:, 1] by self.q when computing rsq, but this doesn't account for the full elliptical transformation. In an elliptical coordinate system, both coordinates should be scaled properly, and the transformation should be applied after rotation to the major/minor axis frame. The current implementation may not correctly represent an elliptical shapelet.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The hardcoded import of
jax.numpy as jnpat the module level forces a dependency on JAX for all users of this module, even if they are not using JAX. This breaks the existing pattern in the codebase wherexpis used as a parameter to switch between numpy and JAX. Consider making JAX an optional dependency and importing it conditionally, or implementing a fallback to scipy.special.genlaguerre when JAX is not available.