Skip to content

Feature/jax in image dict #387

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
Jammy2211 merged 2 commits into from
Feb 11, 2026
Merged

Feature/jax in image dict #387

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
30ed351
xp added to galaxy_image_dict
Feb 11, 2026
fd83da4
black
Feb 11, 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
2 changes: 2 additions & 0 deletions autolens/imaging/fit_imaging.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -172,6 +172,7 @@ def galaxy_image_dict(self) -> Dict[ag.Galaxy, np.ndarray]:

galaxy_image_2d_dict = self.tracer.galaxy_image_2d_dict_from(
grid=self.grids.lp,
xp=self._xp
)

galaxy_linear_obj_image_dict = self.galaxy_linear_obj_data_dict_from(
Expand Down Expand Up @@ -199,6 +200,7 @@ def galaxy_model_image_dict(self) -> Dict[ag.Galaxy, np.ndarray]:
grid=self.grids.lp,
psf=self.dataset.psf,
blurring_grid=self.grids.blurring,
xp=self._xp
)

galaxy_linear_obj_image_dict = self.galaxy_linear_obj_data_dict_from(
Expand Down
6 changes: 4 additions & 2 deletions autolens/interferometer/fit_interferometer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -173,7 +173,9 @@ def galaxy_image_dict(self) -> Dict[ag.Galaxy, np.ndarray]:
For modeling, this dictionary is used to set up the `adapt_images` that adapt certain pixelizations to the
data being fitted.
"""
galaxy_image_dict = self.tracer.galaxy_image_2d_dict_from(grid=self.grids.lp)
galaxy_image_dict = self.tracer.galaxy_image_2d_dict_from(
grid=self.grids.lp, xp=self._xp
)

galaxy_linear_obj_image_dict = self.galaxy_linear_obj_data_dict_from(
use_operated=False
Expand All @@ -194,7 +196,7 @@ def galaxy_model_visibilities_dict(self) -> Dict[ag.Galaxy, np.ndarray]:
are solved for first via the inversion.
"""
galaxy_model_visibilities_dict = self.tracer.galaxy_visibilities_dict_from(
grid=self.grids.lp, transformer=self.dataset.transformer
grid=self.grids.lp, transformer=self.dataset.transformer, xp=self._xp
)

galaxy_linear_obj_visibilities_dict = self.galaxy_linear_obj_data_dict_from(
Expand Down
8 changes: 5 additions & 3 deletions autolens/lens/tracer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -579,7 +579,7 @@ def image_2d_via_input_plane_image_from(
return aa.Array2D(values=image, mask=grid.mask, xp=xp)

def galaxy_image_2d_dict_from(
self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None
self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None
Copy link

Copilot AI Feb 11, 2026

Choose a reason for hiding this comment

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

Adding xp as the second positional parameter changes the positional call signature (e.g. existing user code like galaxy_image_2d_dict_from(grid, True) would now pass True as xp). To preserve backward compatibility, keep operated_only as the second parameter and make xp keyword-only (e.g. by moving xp after operated_only or introducing a * before xp).

Suggested change
self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None
self,
grid: aa.type.Grid2DLike,
operated_only: Optional[bool] = None,
*,
xp=np,

Copilot uses AI. Check for mistakes.
) -> Dict[ag.Galaxy, np.ndarray]:
"""
Returns a dictionary associating every `Galaxy` object in the `Tracer` with its corresponding 2D image, using
Expand All @@ -605,12 +605,14 @@ def galaxy_image_2d_dict_from(

galaxy_image_2d_dict = dict()

traced_grid_list = self.traced_grid_2d_list_from(grid=grid)
traced_grid_list = self.traced_grid_2d_list_from(grid=grid, xp=xp)

for plane_index, galaxies in enumerate(self.planes):
image_2d_list = [
galaxy.image_2d_from(
grid=traced_grid_list[plane_index], operated_only=operated_only
grid=traced_grid_list[plane_index],
operated_only=operated_only,
xp=xp,
)
for galaxy in galaxies
]
Expand Down
8 changes: 2 additions & 6 deletions test_autolens/interferometer/test_fit_interferometer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -180,12 +180,8 @@ def test___galaxy_image_dict(interferometer_7, interferometer_7_grid):
tracer=tracer,
)

assert fit.galaxy_image_dict[g0_linear][4] == pytest.approx(
1.00018622848, 1.0e-2
)
assert fit.galaxy_image_dict[g1_linear][3] == pytest.approx(
-0.017435532289, 1.0e-2
)
assert fit.galaxy_image_dict[g0_linear][4] == pytest.approx(1.00018622848, 1.0e-2)
assert fit.galaxy_image_dict[g1_linear][3] == pytest.approx(-0.017435532289, 1.0e-2)

pixelization = al.Pixelization(
mesh=al.mesh.RectangularUniform(shape=(3, 3)),
Expand Down
Loading