Skip to content

Feature/linalg mixed precision#275

Merged
Jammy2211 merged 2 commits intomainfrom
feature/linalg_mixed_precision
Feb 8, 2026
Merged

Feature/linalg mixed precision#275
Jammy2211 merged 2 commits intomainfrom
feature/linalg_mixed_precision

Conversation

@Jammy2211
Copy link
Copy Markdown
Collaborator

This pull request introduces mixed precision support for linear algebra calculations in the inversion pipeline, allowing targeted computations to use single precision for significant speed improvements on GPUs, reduced VRAM usage, and minimal impact on accuracy.

The changes propagate the new use_mixed_precision setting throughout the codebase, update relevant functions and classes to accept and utilize this option, and refactor the mapping matrix and curvature matrix calculations to handle mixed precision correctly.

Enhancements for inversion and light profile configuration:

  • Added settings parameter (defaulting to SettingsInversion) to the LinearLightProfile class and related functions, ensuring that inversion settings can be consistently passed and used throughout the codebase. (autogalaxy/profiles/light/linear/abstract.py, autogalaxy/galaxy/to_inversion.py) [1] [2] [3] [4]
  • Updated the image convolution logic to utilize the use_mixed_precision setting from settings, allowing for performance optimization during PSF convolution. (autogalaxy/profiles/light/linear/abstract.py)

Code readability and cleanup:

  • Refactored and reformatted several code blocks for clarity, including splitting long lines and improving list formatting in mathematical calculations and array initializations. (autogalaxy/profiles/mass/stellar/gaussian.py, autogalaxy/profiles/basis.py, autogalaxy/abstract_fit.py) [1] [2] [3] [4] [5]
  • Removed unnecessary blank lines and added debugging print statements for light profile intensity inspection. (autogalaxy/analysis/chaining_util.py)

Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

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

Pull request overview

Adds mixed-precision plumbing to the inversion pipeline so selected linear-algebra and PSF convolution steps can run in lower precision (e.g. for GPU speed/VRAM wins), and propagates inversion settings through the light-profile linear object and inversion construction codepaths.

Changes:

  • Thread SettingsInversion into LightProfileLinearObjFuncList and GalaxiesToInversion construction paths.
  • Pass use_mixed_precision into PSF convolution for the operated mapping matrix.
  • Apply small readability-only formatting refactors across a few math-heavy/profile utilities.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 7 comments.

Show a summary per file
File Description
autogalaxy/profiles/mass/stellar/gaussian.py Formatting-only refactor of numeric expressions / array literals.
autogalaxy/profiles/light/linear/abstract.py Adds settings propagation to the linear light-profile obj-func list and uses it in PSF convolution.
autogalaxy/profiles/basis.py Formatting-only refactor of list comprehension layout.
autogalaxy/galaxy/to_inversion.py Propagates inversion settings into linear light-profile obj-func list creation and mapper factory call.
autogalaxy/analysis/chaining_util.py Touches code near debug output (prints still present).
autogalaxy/abstract_fit.py Formatting-only refactor.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment on lines 308 to 313
blurred_image_2d = self.psf.convolved_image_from(
image=image_2d, blurring_image=blurring_image_2d, xp=self._xp
image=image_2d,
blurring_image=blurring_image_2d,
use_mixed_precision=self.settings.use_mixed_precision,
xp=self._xp,
)
Copy link

Copilot AI Feb 8, 2026

Choose a reason for hiding this comment

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

Kernel2D.convolved_image_from(...) is now called with use_mixed_precision=.... Because autoarray is an unpinned dependency, this will raise TypeError: got an unexpected keyword argument in environments using an older autoarray that doesn’t support this kwarg. Consider (1) pinning/bumping the minimum autoarray version that introduces this parameter, or (2) adding a backward-compatible call path (e.g. only passing the kwarg when supported).

Copilot uses AI. Check for mistakes.
Comment on lines 308 to 313
blurred_image_2d = self.psf.convolved_image_from(
image=image_2d, blurring_image=blurring_image_2d, xp=self._xp
image=image_2d,
blurring_image=blurring_image_2d,
use_mixed_precision=self.settings.use_mixed_precision,
xp=self._xp,
)
Copy link

Copilot AI Feb 8, 2026

Choose a reason for hiding this comment

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

The mixed-precision behavior isn’t covered by tests: existing LightProfileLinearObjFuncList tests construct the class without settings, so they don’t validate that settings.use_mixed_precision is forwarded into the PSF convolution. Add a unit test that sets SettingsInversion(use_mixed_precision=True) (or equivalent) and asserts the mixed-precision path is exercised (e.g. via a spy/mocked convolved_image_from).

Copilot uses AI. Check for mistakes.
Comment on lines 320 to 328
lp_linear_func = LightProfileLinearObjFuncList(
grid=self.dataset.grids.lp,
blurring_grid=self.dataset.grids.blurring,
psf=self.dataset.psf,
light_profile_list=light_profile_list,
regularization=light_profile.regularization,
settings=self.settings_inversion,
xp=self._xp,
)
Copy link

Copilot AI Feb 8, 2026

Choose a reason for hiding this comment

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

LightProfileLinearObjFuncList(...) is now instantiated with settings=self.settings_inversion. If this depends on a newer autoarray API (and autoarray is unpinned), older environments may fail at runtime. Ensure the required autoarray minimum version is declared (or provide a compatibility fallback).

Copilot uses AI. Check for mistakes.
Comment on lines 487 to 493
return mapper_from(
mapper_grids=mapper_grids,
regularization=regularization,
settings=self.settings_inversion,
preloads=self.preloads,
xp=self._xp,
)
Copy link

Copilot AI Feb 8, 2026

Choose a reason for hiding this comment

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

The factory call mapper_from(..., settings=self.settings_inversion, ...) introduces a new keyword argument. With autoarray unpinned, this can break for users on versions where mapper_from doesn’t accept settings. Please align this change with an explicit autoarray minimum version (or guard the kwarg for backwards compatibility).

Copilot uses AI. Check for mistakes.
Comment on lines 503 to 505
print(light_profile.intensity)
print(type(light_profile.intensity))

Copy link

Copilot AI Feb 8, 2026

Choose a reason for hiding this comment

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

The print(...) statements will emit output during normal library usage and can be very noisy in pipeline runs. Please remove these debug prints or replace them with optional logging behind a logger / verbosity flag.

Suggested change
print(light_profile.intensity)
print(type(light_profile.intensity))

Copilot uses AI. Check for mistakes.
blurring_grid: aa.type.Grid1D2DLike,
psf: Optional[aa.Kernel2D],
light_profile_list: List[LightProfileLinear],
regularization=Optional[aa.reg.Regularization],
Copy link

Copilot AI Feb 8, 2026

Choose a reason for hiding this comment

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

regularization=Optional[aa.reg.Regularization] sets the default value to the typing.Optional[...] object rather than None, which can lead to unexpected truthiness / attribute errors if callers rely on the default. This should be a type annotation (e.g. regularization: Optional[...] = None) instead of an assigned default.

Suggested change
regularization=Optional[aa.reg.Regularization],
regularization: Optional[aa.reg.Regularization] = None,

Copilot uses AI. Check for mistakes.
Comment on lines 141 to 150
def __init__(
self,
grid: aa.type.Grid1D2DLike,
blurring_grid: aa.type.Grid1D2DLike,
psf: Optional[aa.Kernel2D],
light_profile_list: List[LightProfileLinear],
regularization=Optional[aa.reg.Regularization],
settings=aa.SettingsInversion(),
xp=np,
):
Copy link

Copilot AI Feb 8, 2026

Choose a reason for hiding this comment

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

The new settings parameter is added to the constructor, but the docstring’s Parameters section does not describe it. Please document what settings controls (e.g. mixed precision, sparse operator flags) to keep the public API self-describing.

Copilot uses AI. Check for mistakes.
@Jammy2211 Jammy2211 merged commit 32dc1ed into main Feb 8, 2026
10 of 14 checks passed
@Jammy2211 Jammy2211 deleted the feature/linalg_mixed_precision branch February 13, 2026 13:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants