From 44f8db0b75c07983493cc122790826c62a280e6a Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Sun, 12 Apr 2026 12:36:43 +0100 Subject: [PATCH] fix: give each basis independent ell_comps in mge_model_from Previously, centre and ell_comps priors were constructed once outside the per-basis loop, so gaussian_per_basis > 1 added redundant Gaussians with zero extra flexibility (always 4 free params regardless of basis count). Now ell_comps are always independent per basis, and a new centre_per_basis flag optionally gives each basis its own centre priors. Existing calls with gaussian_per_basis=1 produce identical output. Calls with gaussian_per_basis>1 now get the intended extra flexibility (e.g. 6 free params for 2 bases). Closes #341 Co-Authored-By: Claude Opus 4.6 --- autogalaxy/analysis/model_util.py | 161 ++++++++++++-------- test_autogalaxy/analysis/test_model_util.py | 100 ++++++++++++ 2 files changed, 194 insertions(+), 67 deletions(-) diff --git a/autogalaxy/analysis/model_util.py b/autogalaxy/analysis/model_util.py index 57ce2e27..552a38af 100644 --- a/autogalaxy/analysis/model_util.py +++ b/autogalaxy/analysis/model_util.py @@ -11,6 +11,7 @@ def mge_model_from( centre_prior_is_uniform: bool = True, centre: Tuple[float, float] = (0.0, 0.0), centre_fixed: Optional[Tuple[float, float]] = None, + centre_per_basis: bool = False, centre_sigma: float = 0.3, ell_comps_prior_is_uniform: bool = False, ell_comps_uniform_width: float = 0.2, @@ -18,7 +19,7 @@ def mge_model_from( use_spherical: bool = False, ) -> af.Collection: """ - Construct a Multi-Gaussian Expansion (MGE) for the lens or source galaxy light + Construct a Multi-Gaussian Expansion (MGE) for the lens or source galaxy light. This model is designed as a "start here" configuration for lens modeling: @@ -26,52 +27,65 @@ def mge_model_from( Gaussian light profiles with fixed logarithmically spaced widths (`sigma`). - All Gaussians within each basis share common centres and ellipticity components, reducing degeneracy while retaining flexibility. + - Users can combine with a lens mass model of their choice. - - Users can combine with a lens mass model of their choiuce. + When ``gaussian_per_basis > 1``, each basis receives **independent** ellipticity + components (``ell_comps``), allowing the model to represent twisting or varying + isophotes across different radial scales. Centres are **shared** across bases by + default (the common case: one luminosity centre, complex isophotal shape). Set + ``centre_per_basis=True`` to give each basis its own centre priors. - The resulting model provides a good balance of speed, flexibility, and accuracy - for fitting most galaxy-scale strong lenses. + Expected free-parameter counts (elliptical, ``use_spherical=False``): - This code is mostly to make the API simple for new users, hiding the technical - details of setting up an MGE. More advanced users may wish to customize the - model further. + - ``gaussian_per_basis=1`` : 2 centre + 2 ell_comps = 4 + - ``gaussian_per_basis=K`` (shared centre) : 2 + 2K + - ``gaussian_per_basis=K, centre_per_basis=True`` : 2K + 2K = 4K + + Spherical (``use_spherical=True``): no ell_comps, only centres. + + - Shared centre: 2. Per-basis centre: 2K. Parameters ---------- mask_radius The outer radius (in arcseconds) of the circular mask applied to the data. - This determines the maximum Gaussian width (`sigma`) used in the lens MGE. - lens_total_gaussians - Total number of Gaussian light profiles used in the lens MGE basis. - source_total_gaussians - Total number of Gaussian light profiles used in the source MGE basis. - lens_gaussian_per_basis - Number of separate Gaussian bases to include for the lens light profile. - Each basis has `lens_total_gaussians` components. - source_gaussian_per_basis - Number of separate Gaussian bases to include for the source light profile. - Each basis has `source_total_gaussians` components. + This determines the maximum Gaussian width (`sigma`) used in the MGE. + total_gaussians + Total number of Gaussian light profiles used in each basis. + gaussian_per_basis + Number of separate Gaussian bases. Each basis has ``total_gaussians`` + components sharing the same centre and ellipticity. Multiple bases allow + independent ellipticity (and optionally centre) per radial scale group. + centre_prior_is_uniform + If True (default), centre priors are ``UniformPrior(±0.1)`` around + ``centre``. If False, ``GaussianPrior`` with ``centre_sigma``. + centre + (y, x) centre in arcseconds used as the mean/midpoint for centre priors. + centre_fixed + If not None, fix all Gaussian centres to this (y, x) value instead of + making them free parameters. Overrides ``centre_per_basis``. + centre_per_basis + If True, each basis gets independently drawn centre priors. If False + (default), all bases share the same centre. Ignored when ``centre_fixed`` + is set. + centre_sigma + Sigma for ``GaussianPrior`` centre priors (used when + ``centre_prior_is_uniform=False``). + ell_comps_prior_is_uniform + If True, ell_comps priors are ``UniformPrior``. If False (default), + ``TruncatedGaussianPrior``. + ell_comps_uniform_width + Half-width for uniform ell_comps priors. + ell_comps_sigma + Sigma for truncated-Gaussian ell_comps priors. + use_spherical + If True, use ``GaussianSph`` (no ell_comps). If False (default), use + ``Gaussian`` with ellipticity. Returns ------- - model : af.Collection - An `autofit.Collection` containing: - - A lens galaxy at redshift 0.5, with: - * bulge light profile: MGE basis of Gaussians - * mass profile: Isothermal ellipsoid - * external shear - - A source galaxy at redshift 1.0, with: - * bulge light profile: MGE basis of Gaussians - - Notes - ----- - - Lens light Gaussians have widths (sigma) logarithmically spaced between 0.01" - and the mask radius. - - Source light Gaussians have widths logarithmically spaced between 0.01" and 1.0". - - Gaussian centres are free parameters but tied across all components in each - basis to reduce dimensionality. - - This function is a convenience utility: it hides the technical setup of MGE - composition and provides a ready-to-use lens model for quick experimentation. + af.Model + An ``autofit.Model`` wrapping a ``Basis`` of linear Gaussians. """ import os @@ -83,60 +97,73 @@ def mge_model_from( from autogalaxy.profiles.light.linear import Gaussian, GaussianSph from autogalaxy.profiles.basis import Basis - # The sigma values of the Gaussians will be fixed to values spanning 0.01 to the mask radius, 3.0". + # The sigma values of the Gaussians will be fixed to values spanning 0.01 to the mask radius. log10_sigma_list = np.linspace(-4, np.log10(mask_radius), total_gaussians) - # By defining the centre here, it creates two free parameters that are assigned below to all Gaussians. - - if centre_fixed is not None: - centre_0 = centre[0] - centre_1 = centre[1] - elif centre_prior_is_uniform: - centre_0 = af.UniformPrior( - lower_limit=centre[0] - 0.1, upper_limit=centre[0] + 0.1 - ) - centre_1 = af.UniformPrior( - lower_limit=centre[1] - 0.1, upper_limit=centre[1] + 0.1 - ) - else: - centre_0 = af.GaussianPrior(mean=centre[0], sigma=centre_sigma) - centre_1 = af.GaussianPrior(mean=centre[1], sigma=centre_sigma) - if use_spherical: model_cls = GaussianSph else: model_cls = Gaussian - if ell_comps_prior_is_uniform: - - ell_comps_0 = af.UniformPrior(lower_limit=-ell_comps_uniform_width, upper_limit=ell_comps_uniform_width) - ell_comps_1 = af.UniformPrior(lower_limit=-ell_comps_uniform_width, upper_limit=ell_comps_uniform_width) + def _make_centre_priors(): + if centre_fixed is not None: + return centre[0], centre[1] + elif centre_prior_is_uniform: + return ( + af.UniformPrior( + lower_limit=centre[0] - 0.1, upper_limit=centre[0] + 0.1 + ), + af.UniformPrior( + lower_limit=centre[1] - 0.1, upper_limit=centre[1] + 0.1 + ), + ) + else: + return ( + af.GaussianPrior(mean=centre[0], sigma=centre_sigma), + af.GaussianPrior(mean=centre[1], sigma=centre_sigma), + ) + def _make_ell_comps_priors(): + if ell_comps_prior_is_uniform: + return ( + af.UniformPrior(lower_limit=-ell_comps_uniform_width, upper_limit=ell_comps_uniform_width), + af.UniformPrior(lower_limit=-ell_comps_uniform_width, upper_limit=ell_comps_uniform_width), + ) else: + return ( + af.TruncatedGaussianPrior(mean=0.0, sigma=ell_comps_sigma, lower_limit=-1.0, upper_limit=1.0), + af.TruncatedGaussianPrior(mean=0.0, sigma=ell_comps_sigma, lower_limit=-1.0, upper_limit=1.0), + ) - ell_comps_0 = af.TruncatedGaussianPrior(mean=0.0, sigma=ell_comps_sigma, lower_limit=-1.0, upper_limit=1.0) - ell_comps_1 = af.TruncatedGaussianPrior(mean=0.0, sigma=ell_comps_sigma, lower_limit=-1.0, upper_limit=1.0) + # Shared centre priors (used when centre_per_basis=False). + if not centre_per_basis or centre_fixed is not None: + shared_centre_0, shared_centre_1 = _make_centre_priors() bulge_gaussian_list = [] for j in range(gaussian_per_basis): - # A list of Gaussian model components whose parameters are customized belows. + + # Per-basis centre priors when requested. + if centre_per_basis and centre_fixed is None: + centre_0, centre_1 = _make_centre_priors() + else: + centre_0, centre_1 = shared_centre_0, shared_centre_1 + + # Per-basis ell_comps priors (always independent across bases). + if not use_spherical: + ell_comps_0, ell_comps_1 = _make_ell_comps_priors() gaussian_list = af.Collection( af.Model(model_cls) for _ in range(total_gaussians) ) - # Iterate over every Gaussian and customize its parameters. - for i, gaussian in enumerate(gaussian_list): - gaussian.centre.centre_0 = centre_0 # All Gaussians have same y centre. - gaussian.centre.centre_1 = centre_1 # All Gaussians have same x centre. + gaussian.centre.centre_0 = centre_0 + gaussian.centre.centre_1 = centre_1 if not use_spherical: gaussian.ell_comps.ell_comps_0 = ell_comps_0 gaussian.ell_comps.ell_comps_1 = ell_comps_1 - gaussian.sigma = ( - 10 ** log10_sigma_list[i] - ) # All Gaussian sigmas are fixed to values above. + gaussian.sigma = 10 ** log10_sigma_list[i] bulge_gaussian_list += gaussian_list diff --git a/test_autogalaxy/analysis/test_model_util.py b/test_autogalaxy/analysis/test_model_util.py index 71de80ef..7f873d07 100644 --- a/test_autogalaxy/analysis/test_model_util.py +++ b/test_autogalaxy/analysis/test_model_util.py @@ -5,6 +5,106 @@ import autogalaxy as ag +def test__mge_model_from__single_basis_elliptical(): + model = ag.model_util.mge_model_from( + mask_radius=1.0, total_gaussians=5, gaussian_per_basis=1 + ) + assert model.prior_count == 4 + + +def test__mge_model_from__single_basis_spherical(): + model = ag.model_util.mge_model_from( + mask_radius=1.0, total_gaussians=5, gaussian_per_basis=1, use_spherical=True + ) + assert model.prior_count == 2 + + +def test__mge_model_from__two_bases_shared_centre(): + model = ag.model_util.mge_model_from( + mask_radius=1.0, total_gaussians=5, gaussian_per_basis=2 + ) + # 2 shared centre + 2 ell_comps per basis * 2 = 6 + assert model.prior_count == 6 + + +def test__mge_model_from__two_bases_centre_per_basis(): + model = ag.model_util.mge_model_from( + mask_radius=1.0, total_gaussians=5, gaussian_per_basis=2, centre_per_basis=True + ) + # 2 centre per basis * 2 + 2 ell_comps per basis * 2 = 8 + assert model.prior_count == 8 + + +def test__mge_model_from__three_bases_shared_centre(): + model = ag.model_util.mge_model_from( + mask_radius=1.0, total_gaussians=5, gaussian_per_basis=3 + ) + # 2 shared centre + 2 ell_comps * 3 = 8 + assert model.prior_count == 8 + + +def test__mge_model_from__two_bases_spherical_shared_centre(): + model = ag.model_util.mge_model_from( + mask_radius=1.0, total_gaussians=5, gaussian_per_basis=2, use_spherical=True + ) + # Spherical: no ell_comps, shared centre = 2 + assert model.prior_count == 2 + + +def test__mge_model_from__two_bases_spherical_centre_per_basis(): + model = ag.model_util.mge_model_from( + mask_radius=1.0, + total_gaussians=5, + gaussian_per_basis=2, + use_spherical=True, + centre_per_basis=True, + ) + # Spherical + centre_per_basis: 2 centre * 2 = 4 + assert model.prior_count == 4 + + +def test__mge_model_from__centre_fixed_with_two_bases(): + model = ag.model_util.mge_model_from( + mask_radius=1.0, + total_gaussians=5, + gaussian_per_basis=2, + centre_fixed=(0.0, 0.0), + ) + # Centres fixed (0 free), 2 ell_comps per basis * 2 = 4 + assert model.prior_count == 4 + + +def test__mge_model_from__centre_fixed_overrides_centre_per_basis(): + model = ag.model_util.mge_model_from( + mask_radius=1.0, + total_gaussians=5, + gaussian_per_basis=2, + centre_fixed=(0.0, 0.0), + centre_per_basis=True, + ) + # centre_fixed takes precedence: 0 centre + 4 ell_comps = 4 + assert model.prior_count == 4 + + +def test__mge_model_from__backward_compat_single_basis(): + """gaussian_per_basis=1 must produce identical output to the old code.""" + model = ag.model_util.mge_model_from(mask_radius=3.0, total_gaussians=10) + assert model.prior_count == 4 + + instance = model.instance_from_prior_medians() + assert isinstance(instance, ag.lp_basis.Basis) + assert len(instance.profile_list) == 10 + + +def test__mge_model_from__total_gaussians_per_basis(): + """With gaussian_per_basis=2 and total_gaussians=5, should produce 10 Gaussians.""" + model = ag.model_util.mge_model_from( + mask_radius=1.0, total_gaussians=5, gaussian_per_basis=2 + ) + instance = model.instance_from_prior_medians() + assert len(instance.profile_list) == 10 + + def test__mge_point_model_from__returns_basis_model_with_correct_gaussians(): """ mge_point_model_from should return an af.Model wrapping a Basis whose