diff --git a/slsim/Deflectors/DeflectorPopulation/all_lens_galaxies.py b/slsim/Deflectors/DeflectorPopulation/all_lens_galaxies.py index 2bab92bc7..dc687f25a 100644 --- a/slsim/Deflectors/DeflectorPopulation/all_lens_galaxies.py +++ b/slsim/Deflectors/DeflectorPopulation/all_lens_galaxies.py @@ -11,6 +11,7 @@ from astropy.table import vstack from slsim.Util.param_util import catalog_with_angular_size_in_arcsec from slsim.Deflectors.deflector import Deflector +from slsim.Util.color_gradient import attach_foreground_deflector_color_gradient class AllLensGalaxies(DeflectorsBase): @@ -26,6 +27,8 @@ def __init__( sky_area, gamma_pl=None, catalog_type="skypy", + foreground_color_gradient=None, + foreground_component_weights=(0.4, 0.6), ): """ :param red_galaxy_list: list of dictionary with elliptical galaxy @@ -52,6 +55,14 @@ def __init__( default, this class considers deflector catalog is generated using skypy pipeline. :type catalog_type: str. "skypy" or None. + :param foreground_color_gradient: Optional dictionary defining a + band-dependent foreground deflector light colour gradient. When provided, + the deflector light model receives the corresponding ``color_gradient`` + plus two reference-band component weights. + :type foreground_color_gradient: dict or None + :param foreground_component_weights: Two reference-band flux weights for + the foreground light components. + :type foreground_component_weights: tuple or list """ red_galaxy_list = catalog_with_angular_size_in_arcsec( galaxy_catalog=red_galaxy_list, input_catalog_type=catalog_type @@ -59,6 +70,19 @@ def __init__( blue_galaxy_list = catalog_with_angular_size_in_arcsec( galaxy_catalog=blue_galaxy_list, input_catalog_type=catalog_type ) + if foreground_color_gradient is not None: + red_galaxy_list = red_galaxy_list.copy() + blue_galaxy_list = blue_galaxy_list.copy() + attach_foreground_deflector_color_gradient( + red_galaxy_list, + foreground_color_gradient, + component_weights=foreground_component_weights, + ) + attach_foreground_deflector_color_gradient( + blue_galaxy_list, + foreground_color_gradient, + component_weights=foreground_component_weights, + ) red_column_names = red_galaxy_list.colnames if "galaxy_type" not in red_column_names: red_galaxy_list["galaxy_type"] = "red" diff --git a/slsim/Deflectors/DeflectorPopulation/elliptical_lens_galaxies.py b/slsim/Deflectors/DeflectorPopulation/elliptical_lens_galaxies.py index b1448ff44..d1910beac 100644 --- a/slsim/Deflectors/DeflectorPopulation/elliptical_lens_galaxies.py +++ b/slsim/Deflectors/DeflectorPopulation/elliptical_lens_galaxies.py @@ -7,6 +7,7 @@ vel_disp_abundance_matching, ) from slsim.Deflectors.deflector import Deflector +from slsim.Util.color_gradient import attach_foreground_deflector_color_gradient class EllipticalLensGalaxies(DeflectorsBase): @@ -21,6 +22,8 @@ def __init__( sky_area, gamma_pl=None, catalog_type="skypy", + foreground_color_gradient=None, + foreground_component_weights=(0.4, 0.6), ): """ @@ -44,10 +47,25 @@ def __init__( gamma for uniform distribution. eg: gamma_pl=2.1, gamma_pl={"mean": a, "std_dev": b}, gamma_pl={"gamma_min": c, "gamma_max": d} :type catalog_type: str. "skypy" or None. + :param foreground_color_gradient: Optional dictionary defining a + band-dependent foreground deflector light colour gradient. When provided, + the deflector light model receives the corresponding ``color_gradient`` + plus two reference-band component weights. + :type foreground_color_gradient: dict or None + :param foreground_component_weights: Two reference-band flux weights for + the foreground light components. + :type foreground_component_weights: tuple or list """ galaxy_list = param_util.catalog_with_angular_size_in_arcsec( galaxy_catalog=galaxy_list, input_catalog_type=catalog_type ) + if foreground_color_gradient is not None: + galaxy_list = galaxy_list.copy() + attach_foreground_deflector_color_gradient( + galaxy_list, + foreground_color_gradient, + component_weights=foreground_component_weights, + ) super().__init__( deflector_table=galaxy_list, kwargs_cut=kwargs_cut, diff --git a/slsim/Deflectors/DeflectorTypes/epl_sersic.py b/slsim/Deflectors/DeflectorTypes/epl_sersic.py index ebccd7fdb..6e730d8f0 100644 --- a/slsim/Deflectors/DeflectorTypes/epl_sersic.py +++ b/slsim/Deflectors/DeflectorTypes/epl_sersic.py @@ -1,5 +1,7 @@ from slsim.Deflectors.DeflectorTypes.epl import EPL from slsim.Util.param_util import ellipticity_slsim_to_lenstronomy +from slsim.Util.color_gradient import component_weights_for_band +import numpy as np class EPLSersic(EPL): @@ -58,6 +60,15 @@ def light_model_lenstronomy(self, band=None): ) ) size_lens_arcsec = self.angular_size_light + if self._has_band_dependent_color_gradient(): + return self._double_sersic_light_model_lenstronomy( + band=band, + mag_lens=mag_lens, + size_lens_arcsec=size_lens_arcsec, + e1_light_lens_lenstronomy=e1_light_lens_lenstronomy, + e2_light_lens_lenstronomy=e2_light_lens_lenstronomy, + center_lens=center_lens, + ) lens_light_model_list = ["SERSIC_ELLIPSE"] kwargs_lens_light = [ { @@ -71,3 +82,104 @@ def light_model_lenstronomy(self, band=None): } ] return lens_light_model_list, kwargs_lens_light + + def _has_band_dependent_color_gradient(self): + """Return whether the deflector light should use chromatic + components.""" + color_gradient = self._deflector_dict.get("color_gradient") + return ( + isinstance(color_gradient, dict) + and color_gradient.get("component_spectral_slopes") is not None + ) + + def _double_sersic_light_model_lenstronomy( + self, + band, + mag_lens, + size_lens_arcsec, + e1_light_lens_lenstronomy, + e2_light_lens_lenstronomy, + center_lens, + ): + """Return two Sersic components with band-dependent component + weights.""" + flux = 10 ** (-mag_lens / 2.5) + w0, w1 = self._weights_for_band(band) + mag_lens0 = -2.5 * np.log10(w0 * flux) + mag_lens1 = -2.5 * np.log10(w1 * flux) + angular_size_0, angular_size_1 = self._component_angular_sizes(size_lens_arcsec) + n_sersic_0, n_sersic_1 = self._component_sersic_indices() + + lens_light_model_list = ["SERSIC_ELLIPSE", "SERSIC_ELLIPSE"] + kwargs_lens_light = [ + { + "magnitude": mag_lens0, + "R_sersic": angular_size_0, + "n_sersic": n_sersic_0, + "e1": e1_light_lens_lenstronomy, + "e2": e2_light_lens_lenstronomy, + "center_x": center_lens[0], + "center_y": center_lens[1], + }, + { + "magnitude": mag_lens1, + "R_sersic": angular_size_1, + "n_sersic": n_sersic_1, + "e1": e1_light_lens_lenstronomy, + "e2": e2_light_lens_lenstronomy, + "center_x": center_lens[0], + "center_y": center_lens[1], + }, + ] + return lens_light_model_list, kwargs_lens_light + + def _weights_for_band(self, band): + """Return Sersic component weights for an imaging band.""" + w0 = float(self._deflector_dict.get("w0", 0.5)) + w1 = self._deflector_dict.get("w1") + if w1 is None: + w1 = 1 - w0 + return component_weights_for_band( + base_weights=(w0, float(w1)), + band=band, + color_gradient=self._deflector_dict.get("color_gradient"), + source_dict=self._deflector_dict, + default_reference="i", + ) + + def _component_angular_sizes(self, size_lens_arcsec): + """Return component half-light radii, using conservative opt-in + defaults.""" + angular_size_0 = self._deflector_dict.get("angular_size_0") + angular_size_1 = self._deflector_dict.get("angular_size_1") + if angular_size_0 is not None and angular_size_1 is not None: + return float(angular_size_0), float(angular_size_1) + + radius_factors = self._deflector_dict.get("color_gradient", {}).get( + "component_radius_factors", (0.5, 1.5) + ) + if len(radius_factors) != 2: + raise ValueError( + "color_gradient['component_radius_factors'] must contain two values." + ) + return ( + float(size_lens_arcsec) * float(radius_factors[0]), + float(size_lens_arcsec) * float(radius_factors[1]), + ) + + def _component_sersic_indices(self): + """Return component Sersic indices for the chromatic light model.""" + n_sersic_0 = self._deflector_dict.get("n_sersic_0") + n_sersic_1 = self._deflector_dict.get("n_sersic_1") + if n_sersic_0 is not None and n_sersic_1 is not None: + return float(n_sersic_0), float(n_sersic_1) + + n_sersic = float(self._deflector_dict["n_sersic"]) + indices = self._deflector_dict.get("color_gradient", {}).get( + "component_sersic_indices", (n_sersic, n_sersic) + ) + if len(indices) != 2: + raise ValueError( + "color_gradient['component_sersic_indices'] must contain two values." + ) + return float(indices[0]), float(indices[1]) diff --git a/slsim/ImageSimulation/image_quality_lenstronomy.py b/slsim/ImageSimulation/image_quality_lenstronomy.py index 35bd8546c..39fc23aae 100644 --- a/slsim/ImageSimulation/image_quality_lenstronomy.py +++ b/slsim/ImageSimulation/image_quality_lenstronomy.py @@ -10,6 +10,28 @@ LSST_BAND_LIST = ["u", "g", "r", "i", "z", "y"] EUCLID_BAND_LIST = ["VIS", "Y", "J", "H"] +_DEFAULT_BAND_CENTRAL_WAVELENGTH_MICRON = { + "u": 0.367, + "g": 0.482, + "r": 0.622, + "i": 0.755, + "z": 0.869, + "y": 0.971, + "VIS": 0.715, + "Y": 1.063, + "J": 1.285, + "H": 1.577, + "F062": 0.620, + "F087": 0.870, + "F106": 1.060, + "F129": 1.290, + "F146": 1.460, + "F158": 1.580, + "F184": 1.840, + "F213": 2.130, + "F814W": 0.805, +} + def check_speclite_name(band): """Checks if the raw band name is a valid speclite filter. @@ -256,3 +278,59 @@ def get_all_supported_bands(): for info in _OBSERVATORY_REGISTRY.values(): all_bands.extend(info["bands"]) return all_bands + + +def get_band_central_wavelength(band): + """Return an approximate central wavelength for a registered band. + + The built-in LSST, Roman, and Euclid bands use fixed wavelength + values in microns. For custom registered observatories, this + function falls back to the band's order within the observatory + registry so callers can still build monotonic band-dependent + behavior. + + :param band: Imaging band name. + :type band: str + :return: Approximate central wavelength in microns. + :rtype: float + :raises ValueError: if the band is not registered. + """ + if band in _DEFAULT_BAND_CENTRAL_WAVELENGTH_MICRON: + return _DEFAULT_BAND_CENTRAL_WAVELENGTH_MICRON[band] + + obs_name = get_observatory(band) + bands = _OBSERVATORY_REGISTRY[obs_name]["bands"] + if len(bands) == 1: + return 0.0 + return float(bands.index(band)) / float(len(bands) - 1) + + +def get_band_normalized_position(band, reference_band=None): + """Return a normalized wavelength position for a registered band. + + Built-in LSST, Roman, and Euclid bands are normalized over the full set of + built-in bands, so mixed-observatory simulations are ordered consistently. + Custom bands fall back to their observatory-local registry order. + + :param band: Imaging band name. + :type band: str + :param reference_band: Optional reference band. If provided, return + ``position(band) - position(reference_band)``. + :type reference_band: str or None + :return: Normalized position or position offset. + :rtype: float + :raises ValueError: if any requested band is not registered. + """ + wavelength = get_band_central_wavelength(band) + + if band in _DEFAULT_BAND_CENTRAL_WAVELENGTH_MICRON: + wavelength_values = list(_DEFAULT_BAND_CENTRAL_WAVELENGTH_MICRON.values()) + min_wavelength = min(wavelength_values) + max_wavelength = max(wavelength_values) + position = (wavelength - min_wavelength) / (max_wavelength - min_wavelength) + else: + position = wavelength + + if reference_band is None: + return position + return position - get_band_normalized_position(reference_band) diff --git a/slsim/Sources/SourceTypes/catalog_source.py b/slsim/Sources/SourceTypes/catalog_source.py index c56ab8292..6a74f7ce5 100644 --- a/slsim/Sources/SourceTypes/catalog_source.py +++ b/slsim/Sources/SourceTypes/catalog_source.py @@ -1,7 +1,12 @@ +from slsim.Sources.SourceTypes.double_sersic import DoubleSersic from slsim.Sources.SourceTypes.single_sersic import SingleSersic from slsim.Sources.SourceTypes.source_base import SourceBase from slsim.Sources.SourceCatalogues.CosmosWebCatalog import galaxy_match as CosmosWeb from slsim.Sources.SourceCatalogues.HSTCosmosCatalog import galaxy_match as HSTCosmos +from slsim.Util.color_gradient import ( + edge_apodized_image, + radial_color_gradient_image, +) from lenstronomy.Util.param_util import ellipticity2phi_q CATALOG_TYPES = ["HST_COSMOS, COSMOS_WEB"] @@ -26,6 +31,9 @@ def __init__( max_scale=1, match_n_sersic=False, sersic_fallback=False, + band_dependent_color_gradient=False, + color_gradient=None, + fallback_double_sersic_kwargs=None, **source_dict, ): """ @@ -54,6 +62,18 @@ def __init__( :type match_n_sersic: bool :param sersic_fallback: If the matching process returns no matches, then fall back on a single sersic profile. :type sersic_fallback: bool + :param band_dependent_color_gradient: If True, apply an opt-in radial + colour-gradient transfer to matched HST_COSMOS images. Failed matches + fall back to a DoubleSersic model with the same ``color_gradient``. + :type band_dependent_color_gradient: bool + :param color_gradient: Dictionary containing colour-gradient settings. + Matched HST_COSMOS images use ``grad_color`` (mag/dex) with + ``reference_band`` defaulting to ``F814W``. DoubleSersic fallback uses + ``component_spectral_slopes``. + :type color_gradient: dict or None + :param fallback_double_sersic_kwargs: Optional overrides for the + DoubleSersic parameters used after a failed HST_COSMOS match. + :type fallback_double_sersic_kwargs: dict or None """ super().__init__(extended_source=True, point_source=False, **source_dict) self.name = "GAL" @@ -65,6 +85,9 @@ def __init__( self._max_scale = max_scale self._match_n_sersic = match_n_sersic self._sersic_fallback = sersic_fallback + self._band_dependent_color_gradient = band_dependent_color_gradient + self._color_gradient = color_gradient + self._fallback_double_sersic_kwargs = fallback_double_sersic_kwargs self.source_dict = source_dict # Process catalog and store as class attribute @@ -93,6 +116,24 @@ def __init__( f"Catalog_type {catalog_type} not supported. Currently only {CATALOG_TYPES} are supported." ) + if self._band_dependent_color_gradient: + if catalog_type != "HST_COSMOS": + raise ValueError( + "band_dependent_color_gradient is currently supported only " + "for HST_COSMOS." + ) + if not isinstance(self._color_gradient, dict): + raise ValueError( + "color_gradient must be a dictionary when " + "band_dependent_color_gradient is enabled." + ) + if self._fallback_double_sersic_kwargs is not None and not isinstance( + self._fallback_double_sersic_kwargs, dict + ): + raise ValueError( + "fallback_double_sersic_kwargs must be a dictionary or None." + ) + self._catalog_type = catalog_type self._catalog_path = catalog_path @@ -146,8 +187,10 @@ def kwargs_extended_light(self, band=None): match_n_sersic=self._match_n_sersic, ) ) - # If the matching failed, fall back on a regular sersic profile + # If matching fails, the optional chromatic mode uses DoubleSersic. if self._image_list is None: + if self._band_dependent_color_gradient: + return self._double_sersic_fallback().kwargs_extended_light(band=band) if self._sersic_fallback: if not hasattr(self, "single_sersic"): self.single_sersic = SingleSersic( @@ -158,10 +201,16 @@ def kwargs_extended_light(self, band=None): **self.source_dict, ) return self.single_sersic.kwargs_extended_light(band=band) - else: - raise ValueError( - "No valid matches found! Try reducing the desired angular size or increasing max_scale." - "Alternatively, enable sersic_fallback to use a single sersic whenever the matching fails." + raise ValueError( + "No valid matches found! Try reducing the desired angular size or increasing max_scale." + "Alternatively, enable sersic_fallback to use a single sersic whenever the matching fails." + ) + + if self._band_dependent_color_gradient and self._image_list is not None: + if not hasattr(self, "_chromatic_template"): + edge_width = self._color_gradient.get("edge_apodization_pixels") + self._chromatic_template = edge_apodized_image( + self._image_list[0], edge_width=edge_width ) if band is None: @@ -170,7 +219,7 @@ def kwargs_extended_light(self, band=None): mag_source = self.extended_source_magnitude(band=band) center_source = self.extended_source_position - image = self._select_image_from_band(band) + image = self._image_for_band(band) light_model_list = ["INTERPOL"] kwargs_extended_source = [ @@ -185,6 +234,51 @@ def kwargs_extended_light(self, band=None): ] return light_model_list, kwargs_extended_source + def _image_for_band(self, band): + """Return the catalog image, optionally with HST chromatic + morphology.""" + if self._band_dependent_color_gradient: + image = self._chromatic_template + else: + image = self._select_image_from_band(band) + + if not self._band_dependent_color_gradient or band is None: + return image + + return radial_color_gradient_image( + image=image, + band=band, + color_gradient=self._color_gradient, + angular_size=self.angular_size, + pixel_scale=self._scale, + default_reference="F814W", + ) + + def _double_sersic_fallback(self): + """Build the chromatic fallback model after a failed HST match.""" + if hasattr(self, "double_sersic"): + return self.double_sersic + + fallback_kwargs = { + "angular_size_0": 0.5 * self.angular_size, + "angular_size_1": self.angular_size, + "n_sersic_0": 4.0, + "n_sersic_1": 1.0, + "w0": 0.4, + "w1": 0.6, + "e1_1": self._e1, + "e2_1": self._e2, + "e1_2": self._e1, + "e2_2": self._e2, + } + fallback_kwargs.update(self._fallback_double_sersic_kwargs or {}) + fallback_kwargs["color_gradient"] = self._color_gradient + self.double_sersic = DoubleSersic( + **fallback_kwargs, + **self.source_dict, + ) + return self.double_sersic + def _select_image_from_band(self, band): """Selects an image based off of the input band. Only relevant for source catalogs that provide images for multiple bands. diff --git a/slsim/Sources/SourceTypes/double_sersic.py b/slsim/Sources/SourceTypes/double_sersic.py index 52eba7c76..799668ca0 100644 --- a/slsim/Sources/SourceTypes/double_sersic.py +++ b/slsim/Sources/SourceTypes/double_sersic.py @@ -2,6 +2,7 @@ from slsim.Sources.SourceTypes.source_base import SourceBase from slsim.Util.param_util import ellipticity_slsim_to_lenstronomy from slsim.Util.param_util import surface_brightness_reff +from slsim.Util.color_gradient import component_weights_for_band class DoubleSersic(SourceBase): @@ -15,6 +16,7 @@ def __init__( n_sersic_1, w0, w1=None, + color_gradient=None, e1_1=0, e2_1=0, e1_2=0, @@ -33,6 +35,10 @@ def __init__( :param e2_2: eccentricity component of second Sersic :param w0: flux weight of first Sersic component :param w1: flux weight of second Sersic component, if =None, will be set w1 = 1 - w0, otherwise it has to match. + :param color_gradient: Optional dictionary defining a band-dependent + two-component colour gradient with lightweight SED slopes. Supported + keys are ``component_spectral_slopes``, ``reference_band``, and + ``min_weight``. Components with larger spectral slopes are redder. :param source_dict: dictionary for SourceBase() option (see documentation) :type source_dict: dict or astropy.table.Table @@ -55,6 +61,7 @@ def __init__( w1 = 1 - w0 assert np.isclose(w0 + w1, 1, rtol=1e-3) self._w1 = w1 + self._color_gradient = color_gradient self._light_model_list = [ "SERSIC_ELLIPSE", @@ -123,8 +130,9 @@ def kwargs_extended_light(self, band=None): center_source = self.extended_source_position # compute magnitude for each sersic component based on weight flux = 10 ** (-mag_source / 2.5) - mag_source0 = -2.5 * np.log10(self._w0 * flux) - mag_source1 = -2.5 * np.log10(self._w1 * flux) + w0, w1 = self._weights_for_band(band) + mag_source0 = -2.5 * np.log10(w0 * flux) + mag_source1 = -2.5 * np.log10(w1 * flux) # convert from slsim to lenstronomy convention. e1_light_source_1_lenstronomy, e2_light_source_1_lenstronomy = ( ellipticity_slsim_to_lenstronomy( @@ -161,6 +169,16 @@ def kwargs_extended_light(self, band=None): ] return self._light_model_list, kwargs_extended_source + def _weights_for_band(self, band): + """Return Sersic component weights for an imaging band.""" + return component_weights_for_band( + base_weights=(self._w0, self._w1), + band=band, + color_gradient=self._color_gradient, + source_dict=self.source_dict, + default_reference="i", + ) + def _shape_light_model(self): """ diff --git a/slsim/Util/color_gradient.py b/slsim/Util/color_gradient.py new file mode 100644 index 000000000..4e5905455 --- /dev/null +++ b/slsim/Util/color_gradient.py @@ -0,0 +1,219 @@ +import numpy as np + + +def default_reference_band(source_dict, default="i"): + """Choose the available source band closest to a default reference band.""" + from slsim.ImageSimulation.image_quality_lenstronomy import ( + get_band_normalized_position, + ) + + available_bands = [ + key.replace("mag_", "", 1) + for key in source_dict + if isinstance(key, str) and key.startswith("mag_") + ] + if not available_bands: + return default + + positions = [ + abs(get_band_normalized_position(band=band, reference_band=default)) + for band in available_bands + ] + return available_bands[int(np.argmin(positions))] + + +def component_weights_for_band( + base_weights, + band, + color_gradient=None, + source_dict=None, + default_reference="i", +): + """Return band-dependent component weights from local SED slopes. + + This is a lightweight, effective-wavelength approximation to chromatic + light components. Instead of integrating a full stellar-population SED + through each bandpass, each component is assigned a local power-law SED, + ``S_k(lambda) proportional lambda**alpha_k``, evaluated at the central + wavelength of the requested band. The reference-band component weights are + then reweighted as + + ``w_k(b) = w_k(ref) * (lambda_b/lambda_ref)**alpha_k / normalization``. + + The approximation is intended to introduce controlled colour gradients in + analytic multi-component light profiles; it is not a replacement for a + stellar population synthesis model. + + Ref: + Hogg et al. 2002, "The K correction", astro-ph/0210394: + broadband fluxes are formally filter-response weighted SED integrals; + this function uses the corresponding effective-wavelength limit. + Conroy 2013, ARA&A, 51, 393: + review of full stellar-population SED modelling, useful context for + what is intentionally omitted by this lightweight approximation. + La Barbera et al. 2005, MNRAS, 358, 1116; La Barbera & de Carvalho 2009, + ApJ, 699, L76: + observational motivation for radial colour gradients in galaxies. + """ + weights = np.asarray(base_weights, dtype=float) + weights = weights / np.sum(weights) + + if band is None or color_gradient is None: + return tuple(float(weight) for weight in weights) + if not isinstance(color_gradient, dict): + raise ValueError("color_gradient must be a dictionary or None.") + + slopes = color_gradient.get( + "component_spectral_slopes", color_gradient.get("sed_slopes") + ) + if slopes is None: + return tuple(float(weight) for weight in weights) + + slopes = np.asarray(slopes, dtype=float) + if slopes.shape != weights.shape: + raise ValueError( + "color_gradient['component_spectral_slopes'] must match the " + "number of components." + ) + + reference_band = color_gradient.get("reference_band") + if reference_band is None: + reference_band = default_reference_band( + source_dict or {}, default=default_reference + ) + + min_weight = float(color_gradient.get("min_weight", 1e-4)) + if not 0 <= min_weight < 1 / len(weights): + raise ValueError( + "color_gradient['min_weight'] must be in [0, 1 / n_components)." + ) + + from slsim.ImageSimulation.image_quality_lenstronomy import ( + get_band_central_wavelength, + ) + + wavelength = get_band_central_wavelength(band) + reference_wavelength = get_band_central_wavelength(reference_band) + if reference_wavelength <= 0: + raise ValueError("The reference band wavelength must be positive.") + + sed_factors = (wavelength / reference_wavelength) ** slopes + sed_weights = weights * sed_factors + sed_weights = sed_weights / np.sum(sed_weights) + sed_weights = np.clip(sed_weights, min_weight, 1 - min_weight) + sed_weights = sed_weights / np.sum(sed_weights) + return tuple(float(weight) for weight in sed_weights) + + +def attach_foreground_deflector_color_gradient( + galaxy_table, + color_gradient, + component_weights=(0.4, 0.6), +): + """Attach opt-in foreground colour-gradient columns to a deflector table. + + The resulting columns are consumed by ``EPLSersic.light_model_lenstronomy`` + to split the foreground light into two chromatic Sersic components. + The operation is in-place and returns ``galaxy_table`` for convenience. + + :param galaxy_table: galaxy/deflector table to annotate + :param color_gradient: dictionary with ``component_spectral_slopes`` and + optional foreground component settings + :param component_weights: two reference-band flux weights for the Sersic + components + :return: annotated galaxy table + """ + if color_gradient is None: + return galaxy_table + if not isinstance(color_gradient, dict): + raise ValueError("color_gradient must be a dictionary or None.") + + weights = np.asarray(component_weights, dtype=float) + if weights.shape != (2,): + raise ValueError("component_weights must contain two values.") + if np.any(weights < 0) or np.sum(weights) <= 0: + raise ValueError("component_weights must be non-negative with positive sum.") + weights = weights / np.sum(weights) + + galaxy_table["color_gradient"] = [ + dict(color_gradient) for _ in range(len(galaxy_table)) + ] + galaxy_table["w0"] = np.full(len(galaxy_table), weights[0]) + galaxy_table["w1"] = np.full(len(galaxy_table), weights[1]) + return galaxy_table + + +def edge_apodized_image(image, edge_width=None): + """Return image multiplied by a cosine taper at the cutout edges.""" + image = np.asarray(image, dtype=float) + if edge_width is None: + edge_width = max(1, min(image.shape) // 20) + edge_width = int(edge_width) + if edge_width <= 0: + return image + + y_grid, x_grid = np.indices(image.shape, dtype=float) + distance_to_edge = np.minimum.reduce( + [ + x_grid, + y_grid, + image.shape[1] - 1 - x_grid, + image.shape[0] - 1 - y_grid, + ] + ) + mask = np.ones_like(image, dtype=float) + edge_region = distance_to_edge < edge_width + mask[edge_region] = 0.5 * ( + 1 - np.cos(np.pi * distance_to_edge[edge_region] / edge_width) + ) + return image * mask + + +def radial_color_gradient_image( + image, + band, + color_gradient, + angular_size, + pixel_scale, + default_reference="F814W", +): + """Apply a d(color)/dlog10(r) gradient to an image and preserve flux. + + See https://arxiv.org/pdf/1006.4056 for details. + """ + if band is None or color_gradient is None: + return image + if not isinstance(color_gradient, dict): + raise ValueError("color_gradient must be a dictionary or None.") + + grad_color = float( + color_gradient.get("grad_color", color_gradient.get("gradient", 0.0)) + ) + if grad_color == 0: + return image + + reference_band = color_gradient.get("reference_band") or default_reference + from slsim.ImageSimulation.image_quality_lenstronomy import ( + get_band_normalized_position, + ) + + band_offset = get_band_normalized_position(band=band, reference_band=reference_band) + if band_offset == 0: + return image + + image = np.asarray(image, dtype=float) + y_grid, x_grid = np.indices(image.shape, dtype=float) + center_x = (image.shape[1] - 1) / 2 + center_y = (image.shape[0] - 1) / 2 + radius = np.hypot(x_grid - center_x, y_grid - center_y) + half_light_radius_pixels = angular_size / pixel_scale + radius = np.maximum(radius, 0.5) + radius_ratio = radius / max(half_light_radius_pixels, 0.5) + + delta_mag = band_offset * grad_color * np.log10(radius_ratio) + chromatic_image = image * 10 ** (-0.4 * delta_mag) + original_flux = np.sum(image) + chromatic_flux = np.sum(chromatic_image) + if chromatic_flux != 0: + chromatic_image *= original_flux / chromatic_flux + return chromatic_image diff --git a/tests/test_Deflectors/test_DeflectorPopulation/test_all_lens_galaxies.py b/tests/test_Deflectors/test_DeflectorPopulation/test_all_lens_galaxies.py index 087467148..98817b606 100644 --- a/tests/test_Deflectors/test_DeflectorPopulation/test_all_lens_galaxies.py +++ b/tests/test_Deflectors/test_DeflectorPopulation/test_all_lens_galaxies.py @@ -139,5 +139,57 @@ def test_all_lens_galaxies_2(): ) +def test_all_lens_galaxies_foreground_color_gradient(): + red_galaxies = foreground_test_galaxy_table() + blue_galaxies = foreground_test_galaxy_table() + kwargs_deflector_cut = {} + kwargs_mass2light = {} + cosmo = FlatLambdaCDM(H0=70, Om0=0.3) + sky_area = Quantity(value=0.05, unit="deg2") + foreground_color_gradient = { + "component_spectral_slopes": [2.0, -1.0], + "reference_band": "i", + } + + galaxy_class = AllLensGalaxies( + red_galaxies, + blue_galaxies, + kwargs_cut=kwargs_deflector_cut, + kwargs_mass2light=kwargs_mass2light, + cosmo=cosmo, + sky_area=sky_area, + catalog_type=None, + foreground_color_gradient=foreground_color_gradient, + foreground_component_weights=(0.4, 0.6), + ) + assert "color_gradient" not in red_galaxies.colnames + assert "color_gradient" not in blue_galaxies.colnames + deflector = galaxy_class.draw_deflector() + model_list, kwargs_light = deflector.light_model_lenstronomy(band="i") + flux0 = 10 ** (-kwargs_light[0]["magnitude"] / 2.5) + flux1 = 10 ** (-kwargs_light[1]["magnitude"] / 2.5) + + assert model_list == ["SERSIC_ELLIPSE", "SERSIC_ELLIPSE"] + assert flux0 / (flux0 + flux1) == pytest.approx(0.4) + + +def foreground_test_galaxy_table(): + return Table( + { + "z": [0.2, 0.3], + "stellar_mass": [10**11, 2 * 10**11], + "angular_size": [0.7, 0.8], + "ellipticity": [0.2, 0.25], + "mag_i": [19.0, 20.0], + "e1_light": [0.1, 0.1], + "e2_light": [0.0, 0.0], + "e1_mass": [0.1, 0.1], + "e2_mass": [0.0, 0.0], + "n_sersic": [4.0, 4.0], + "vel_disp": [200.0, 210.0], + } + ) + + if __name__ == "__main__": pytest.main() diff --git a/tests/test_Deflectors/test_DeflectorPopulation/test_elliptical_lens_galaxies.py b/tests/test_Deflectors/test_DeflectorPopulation/test_elliptical_lens_galaxies.py index bea6d168e..b1f88f898 100644 --- a/tests/test_Deflectors/test_DeflectorPopulation/test_elliptical_lens_galaxies.py +++ b/tests/test_Deflectors/test_DeflectorPopulation/test_elliptical_lens_galaxies.py @@ -5,6 +5,7 @@ from slsim.Util.param_util import vel_disp_from_m_star from slsim.Pipelines.skypy_pipeline import SkyPyPipeline from astropy.units import Quantity +from astropy.table import Table import copy import pytest @@ -104,5 +105,54 @@ def test_elliptical_lens_galaxies_2(): ) +def test_elliptical_lens_galaxies_foreground_color_gradient(): + red_galaxies = foreground_test_galaxy_table() + kwargs_deflector_cut = {} + kwargs_mass2light = {} + cosmo = FlatLambdaCDM(H0=70, Om0=0.3) + sky_area = Quantity(value=0.001, unit="deg2") + foreground_color_gradient = { + "component_spectral_slopes": [2.0, -1.0], + "reference_band": "i", + } + + galaxy_class = EllipticalLensGalaxies( + red_galaxies, + kwargs_cut=kwargs_deflector_cut, + kwargs_mass2light=kwargs_mass2light, + cosmo=cosmo, + sky_area=sky_area, + catalog_type=None, + foreground_color_gradient=foreground_color_gradient, + foreground_component_weights=(0.4, 0.6), + ) + assert "color_gradient" not in red_galaxies.colnames + deflector = galaxy_class.draw_deflector() + model_list, kwargs_light = deflector.light_model_lenstronomy(band="i") + flux0 = 10 ** (-kwargs_light[0]["magnitude"] / 2.5) + flux1 = 10 ** (-kwargs_light[1]["magnitude"] / 2.5) + + assert model_list == ["SERSIC_ELLIPSE", "SERSIC_ELLIPSE"] + assert flux0 / (flux0 + flux1) == pytest.approx(0.4) + + +def foreground_test_galaxy_table(): + return Table( + { + "z": [0.2, 0.3], + "stellar_mass": [10**11, 2 * 10**11], + "angular_size": [0.7, 0.8], + "ellipticity": [0.2, 0.25], + "mag_i": [19.0, 20.0], + "e1_light": [0.1, 0.1], + "e2_light": [0.0, 0.0], + "e1_mass": [0.1, 0.1], + "e2_mass": [0.0, 0.0], + "n_sersic": [4.0, 4.0], + "vel_disp": [200.0, 210.0], + } + ) + + if __name__ == "__main__": pytest.main() diff --git a/tests/test_Deflectors/test_DeflectorTypes/test_epl_sersic.py b/tests/test_Deflectors/test_DeflectorTypes/test_epl_sersic.py index 8677c6c21..8163fd1ba 100644 --- a/tests/test_Deflectors/test_DeflectorTypes/test_epl_sersic.py +++ b/tests/test_Deflectors/test_DeflectorTypes/test_epl_sersic.py @@ -1,4 +1,5 @@ import pytest +import numpy as np from slsim.Deflectors.DeflectorTypes.epl_sersic import EPLSersic from astropy.cosmology import FlatLambdaCDM @@ -102,6 +103,79 @@ def test_halo_porperties(self): gamma = self.sie_sersic.halo_properties["gamma_pl"] assert gamma == 2.0 + def test_light_model_lenstronomy_keeps_single_component_without_gradient(self): + lens_light_model_list, kwargs_lens_light = ( + self.epl_sersic.light_model_lenstronomy(band=None) + ) + assert lens_light_model_list == ["SERSIC_ELLIPSE"] + assert len(kwargs_lens_light) == 1 + + def test_band_dependent_color_gradient_light_model(self): + deflector_dict = self.deflector_dict.copy() + deflector_dict.update( + { + "mag_g": 20, + "mag_i": 19, + "mag_y": 18, + "w0": 0.4, + "w1": 0.6, + "color_gradient": { + "component_spectral_slopes": [2.0, -1.0], + "reference_band": "i", + }, + } + ) + epl_sersic = EPLSersic(**deflector_dict) + + model_g, kwargs_g = epl_sersic.light_model_lenstronomy(band="g") + model_i, kwargs_i = epl_sersic.light_model_lenstronomy(band="i") + model_y, kwargs_y = epl_sersic.light_model_lenstronomy(band="y") + + assert model_g == ["SERSIC_ELLIPSE", "SERSIC_ELLIPSE"] + assert model_i == ["SERSIC_ELLIPSE", "SERSIC_ELLIPSE"] + assert kwargs_i[0]["R_sersic"] < kwargs_i[1]["R_sersic"] + assert np.isclose(component_flux_fraction(kwargs_i), 0.4, rtol=1e-12) + assert component_flux_fraction(kwargs_y) > component_flux_fraction(kwargs_g) + + def test_color_gradient_component_overrides(self): + deflector_dict = self.deflector_dict.copy() + deflector_dict.update( + { + "mag_i": 19, + "angular_size_0": 0.01, + "angular_size_1": 0.08, + "n_sersic_0": 1, + "n_sersic_1": 4, + "color_gradient": { + "component_spectral_slopes": [1.0, 0.0], + "component_radius_factors": [0.25, 2.0], + "component_sersic_indices": [2, 3], + }, + } + ) + _, kwargs_lens_light = EPLSersic(**deflector_dict).light_model_lenstronomy( + band="i" + ) + + assert kwargs_lens_light[0]["R_sersic"] == 0.01 + assert kwargs_lens_light[1]["R_sersic"] == 0.08 + assert kwargs_lens_light[0]["n_sersic"] == 1 + assert kwargs_lens_light[1]["n_sersic"] == 4 + + def test_invalid_color_gradient_component_configuration(self): + deflector_dict = self.deflector_dict.copy() + deflector_dict.update( + { + "mag_i": 19, + "color_gradient": { + "component_spectral_slopes": [1.0, 0.0], + "component_radius_factors": [1.0], + }, + } + ) + with pytest.raises(ValueError, match="component_radius_factors"): + EPLSersic(**deflector_dict).light_model_lenstronomy(band="i") + @pytest.fixture def gamma_epl_sersic_instance(): @@ -131,5 +205,11 @@ def test_mass_model_lenstronomy_gamma(gamma_epl_sersic_instance): assert lens_mass_model_list[0] == "EPL" +def component_flux_fraction(kwargs_lens_light): + flux0 = 10 ** (-kwargs_lens_light[0]["magnitude"] / 2.5) + flux1 = 10 ** (-kwargs_lens_light[1]["magnitude"] / 2.5) + return flux0 / (flux0 + flux1) + + if __name__ == "__main__": pytest.main() diff --git a/tests/test_ImageSimulation/test_image_quality_lenstronomy.py b/tests/test_ImageSimulation/test_image_quality_lenstronomy.py index cb161679d..d676f01d1 100644 --- a/tests/test_ImageSimulation/test_image_quality_lenstronomy.py +++ b/tests/test_ImageSimulation/test_image_quality_lenstronomy.py @@ -8,6 +8,8 @@ get_observatory, register_observatory, get_all_supported_bands, + get_band_central_wavelength, + get_band_normalized_position, ) @@ -248,5 +250,37 @@ def test_get_all_supported_bands_contains_defaults(): assert band in all_bands +def test_default_band_wavelength_ordering(): + assert get_band_central_wavelength("g") < get_band_central_wavelength("i") + assert get_band_central_wavelength("F106") < get_band_central_wavelength("F184") + assert get_band_central_wavelength("VIS") < get_band_central_wavelength("H") + assert get_band_normalized_position("F184", reference_band="g") > 0 + + +def test_band_wavelength_helpers_cover_hst_and_custom_registry_fallback(): + assert get_band_central_wavelength("F814W") == pytest.approx(0.805) + assert get_band_normalized_position("F814W", reference_band="F814W") == 0 + + register_observatory( + name="WavelengthTestObs", + observatory_class=DummyObservatory, + bands=["W1", "W2", "W3"], + ) + assert get_band_central_wavelength("W1") == 0.0 + assert get_band_central_wavelength("W2") == 0.5 + assert get_band_normalized_position("W3", reference_band="W1") == 1.0 + + register_observatory( + name="SingleBandTestObs", + observatory_class=DummyObservatory, + bands=["OnlyBand"], + ) + assert get_band_central_wavelength("OnlyBand") == 0.0 + assert get_band_normalized_position("OnlyBand") == 0.0 + + with pytest.raises(ValueError, match="not recognised"): + get_band_central_wavelength("UnknownBand") + + if __name__ == "__main__": pytest.main() diff --git a/tests/test_Sources/test_SourceTypes/test_catalog_source.py b/tests/test_Sources/test_SourceTypes/test_catalog_source.py index f91c0c0c3..d90d60205 100644 --- a/tests/test_Sources/test_SourceTypes/test_catalog_source.py +++ b/tests/test_Sources/test_SourceTypes/test_catalog_source.py @@ -10,6 +10,7 @@ from slsim.Pipelines import SkyPyPipeline from slsim.Sources.SourcePopulation.galaxies import Galaxies from slsim.Sources.SourceTypes.single_sersic import SingleSersic +from slsim.Sources.SourceTypes.double_sersic import DoubleSersic from slsim.Sources.SourceTypes.catalog_source import CatalogSource from slsim.Sources.source import Source from slsim.Deflectors.deflector import Deflector @@ -134,6 +135,110 @@ def test_select_image_from_band(self): band="wrong", ) + def test_hst_band_dependent_color_gradient(self): + source_dict = dict(self.source1.source_dict) + source = CatalogSource( + angular_size=self.source1.angular_size, + e1=self.source1.ellipticity[0], + e2=self.source1.ellipticity[1], + n_sersic=0.8, + cosmo=self.source1._cosmo, + catalog_path=hst_cosmos_path, + catalog_type="HST_COSMOS", + band_dependent_color_gradient=True, + color_gradient={ + "grad_color": -0.3, + "reference_band": "F814W", + "edge_apodization_pixels": 2, + }, + **source_dict, + ) + _, reference_kwargs = source.kwargs_extended_light(band="i") + reference_image = source._image_for_band(band=None) + + assert not np.allclose(reference_kwargs[0]["image"], reference_image) + np.testing.assert_allclose( + np.sum(reference_kwargs[0]["image"]), np.sum(reference_image) + ) + assert np.all(reference_image[0, :] == 0) + assert np.all(reference_image[:, 0] == 0) + + source._color_gradient["grad_color"] = 0.0 + np.testing.assert_allclose( + source._image_for_band(band="i"), source._image_for_band(band=None) + ) + + def test_hst_chromatic_double_sersic_fallback(self): + source_dict = { + "z": 0.5, + "mag_g": 20.3, + "mag_i": 20.3, + "mag_y": 20.3, + "n_sersic": 0.8, + "angular_size": 1.3, + "e1": 0.09697001616620306, + "e2": 0.040998265256000574, + "center_x": 0.0, + "center_y": 0.0, + } + source = CatalogSource( + cosmo=FlatLambdaCDM(H0=70, Om0=0.3), + catalog_path=hst_cosmos_path, + catalog_type="HST_COSMOS", + max_scale=0.1, + band_dependent_color_gradient=True, + color_gradient={ + "component_spectral_slopes": [2.0, -1.0], + "reference_band": "i", + }, + **source_dict, + ) + source_model, kwargs_light = source.kwargs_extended_light(band="y") + _, kwargs_light_blue = source.kwargs_extended_light(band="g") + + assert source_model == ["SERSIC_ELLIPSE", "SERSIC_ELLIPSE"] + assert isinstance(source.double_sersic, DoubleSersic) + assert len(kwargs_light) == 2 + + flux_y = 10 ** (-np.array([item["magnitude"] for item in kwargs_light]) / 2.5) + flux_g = 10 ** ( + -np.array([item["magnitude"] for item in kwargs_light_blue]) / 2.5 + ) + assert flux_y[0] / np.sum(flux_y) > flux_g[0] / np.sum(flux_g) + + def test_chromatic_catalog_source_validation(self): + source_dict = dict(self.source1.source_dict) + common_kwargs = { + "angular_size": self.source1.angular_size, + "e1": self.source1.ellipticity[0], + "e2": self.source1.ellipticity[1], + "n_sersic": 0.8, + "cosmo": self.source1._cosmo, + "catalog_path": hst_cosmos_path, + "band_dependent_color_gradient": True, + } + + cosmos_web_kwargs = dict(common_kwargs) + cosmos_web_kwargs["catalog_path"] = cosmos_web_path + with pytest.raises(ValueError, match="only for HST_COSMOS"): + CatalogSource( + catalog_type="COSMOS_WEB", + **cosmos_web_kwargs, + **source_dict, + ) + + with pytest.raises(ValueError, match="color_gradient must be a dictionary"): + CatalogSource(catalog_type="HST_COSMOS", **common_kwargs, **source_dict) + + with pytest.raises(ValueError, match="fallback_double_sersic_kwargs"): + CatalogSource( + catalog_type="HST_COSMOS", + color_gradient={"grad_color": -0.1}, + fallback_double_sersic_kwargs="invalid", + **common_kwargs, + **source_dict, + ) + def test_redshift(self): assert self.source1.redshift == 3.5 diff --git a/tests/test_Sources/test_SourceTypes/test_double_sersic.py b/tests/test_Sources/test_SourceTypes/test_double_sersic.py index 930f9ebd3..2303204a0 100644 --- a/tests/test_Sources/test_SourceTypes/test_double_sersic.py +++ b/tests/test_Sources/test_SourceTypes/test_double_sersic.py @@ -1,5 +1,7 @@ from slsim.Sources.SourceTypes.double_sersic import DoubleSersic +from slsim.Util.color_gradient import default_reference_band from slsim.Util.param_util import ellipticity_slsim_to_lenstronomy +import numpy as np import pytest from numpy import testing as npt @@ -82,6 +84,90 @@ def test_surface_brightness_reff(self): result = self.source.surface_brightness_reff(band="i") npt.assert_almost_equal(result, 21.313, decimal=3) + def test_band_dependent_color_gradient(self): + source_dict = dict(self.source_dict) + source_dict.update( + { + "mag_g": 23, + "mag_y": 23, + "color_gradient": { + "component_spectral_slopes": [2.0, -1.0], + "reference_band": "i", + }, + } + ) + source = DoubleSersic(**source_dict) + + _, kwargs_g = source.kwargs_extended_light(band="g") + _, kwargs_y = source.kwargs_extended_light(band="y") + + flux_g0 = 10 ** (-kwargs_g[0]["magnitude"] / 2.5) + flux_g1 = 10 ** (-kwargs_g[1]["magnitude"] / 2.5) + flux_y0 = 10 ** (-kwargs_y[0]["magnitude"] / 2.5) + flux_y1 = 10 ** (-kwargs_y[1]["magnitude"] / 2.5) + + assert flux_y0 / (flux_y0 + flux_y1) > flux_g0 / (flux_g0 + flux_g1) + + def test_color_gradient_disabled_or_missing_slopes_uses_base_weights(self): + assert self.source._weights_for_band("i") == (0.4, 0.6) + + source_dict = dict(self.source_dict) + source_dict["color_gradient"] = {} + no_gradient_source = DoubleSersic(**source_dict) + assert no_gradient_source._weights_for_band("i") == (0.4, 0.6) + + def test_default_reference_band_and_weight_validation(self): + source_dict = dict(self.source_dict) + source_dict.update( + { + "mag_g": 23, + "mag_y": 23, + "color_gradient": {"component_spectral_slopes": [1.0, 0.0]}, + } + ) + source = DoubleSersic(**source_dict) + + assert default_reference_band(source.source_dict) == "i" + assert source._weights_for_band("i") == (0.4, 0.6) + + source_without_magnitudes = dict(self.source_dict) + source_without_magnitudes.pop("mag_i") + source_without_magnitudes["color_gradient"] = { + "component_spectral_slopes": [1.0, 0.0] + } + no_magnitude_source = DoubleSersic(**source_without_magnitudes) + assert default_reference_band(no_magnitude_source.source_dict) == "i" + + source_dict["color_gradient"] = "invalid" + with pytest.raises(ValueError, match="must be a dictionary"): + DoubleSersic(**source_dict)._weights_for_band("i") + + source_dict["color_gradient"] = { + "component_spectral_slopes": [1.0], + } + with pytest.raises(ValueError, match="must match the number of components"): + DoubleSersic(**source_dict)._weights_for_band("i") + + source_dict["color_gradient"] = { + "component_spectral_slopes": [1.0, 0.0], + "min_weight": 0.5, + } + with pytest.raises(ValueError, match="must be in"): + DoubleSersic(**source_dict)._weights_for_band("i") + + def test_color_gradient_clips_component_weight(self): + source_dict = dict(self.source_dict) + source_dict["color_gradient"] = { + "component_spectral_slopes": [100.0, 0.0], + "reference_band": "i", + "min_weight": 0.2, + } + source = DoubleSersic(**source_dict) + + w0, w1 = source._weights_for_band("F213") + assert np.isclose(w0, 0.8) + assert np.isclose(w0 + w1, 1.0) + if __name__ == "__main__": pytest.main() diff --git a/tests/test_Util/test_color_gradient.py b/tests/test_Util/test_color_gradient.py new file mode 100644 index 000000000..c4f963f75 --- /dev/null +++ b/tests/test_Util/test_color_gradient.py @@ -0,0 +1,159 @@ +import numpy as np +import pytest + +from slsim.Util.color_gradient import ( + attach_foreground_deflector_color_gradient, + component_weights_for_band, + default_reference_band, + edge_apodized_image, + radial_color_gradient_image, +) +from slsim.ImageSimulation.image_quality_lenstronomy import register_observatory +from astropy.table import Table + + +class DummyObservatory: + def __init__(self, band, **kwargs): + self.band = band + + +def test_component_weights_for_band_uses_power_law_sed_slopes(): + color_gradient = { + "component_spectral_slopes": [2.0, -1.0], + "reference_band": "i", + } + + w_g = component_weights_for_band((0.4, 0.6), "g", color_gradient) + w_i = component_weights_for_band((0.4, 0.6), "i", color_gradient) + w_y = component_weights_for_band((0.4, 0.6), "y", color_gradient) + + assert w_i == pytest.approx((0.4, 0.6)) + assert w_y[0] > w_g[0] + assert np.sum(w_y) == pytest.approx(1.0) + + +def test_attach_foreground_deflector_color_gradient_adds_table_columns(): + table = Table({"mag_i": [20.0, 21.0]}) + color_gradient = { + "component_spectral_slopes": [2.0, -1.0], + "reference_band": "i", + } + + result = attach_foreground_deflector_color_gradient( + table, color_gradient, component_weights=(2, 3) + ) + + assert result is table + assert table["color_gradient"][0] == color_gradient + assert table["w0"][0] == pytest.approx(0.4) + assert table["w1"][0] == pytest.approx(0.6) + + with pytest.raises(ValueError, match="must be a dictionary"): + attach_foreground_deflector_color_gradient(table, "bad") + with pytest.raises(ValueError, match="two values"): + attach_foreground_deflector_color_gradient(table, color_gradient, (1, 2, 3)) + with pytest.raises(ValueError, match="positive sum"): + attach_foreground_deflector_color_gradient(table, color_gradient, (0, 0)) + + +def test_component_weights_validation_and_default_reference(): + source_dict = {"mag_g": 22, "mag_i": 22, "mag_y": 22} + assert default_reference_band(source_dict) == "i" + + assert component_weights_for_band((2, 3), "i") == pytest.approx((0.4, 0.6)) + assert component_weights_for_band((0.4, 0.6), "i", {}) == pytest.approx((0.4, 0.6)) + + with pytest.raises(ValueError, match="must be a dictionary"): + component_weights_for_band((0.4, 0.6), "i", "bad") + + with pytest.raises(ValueError, match="must match the number of components"): + component_weights_for_band( + (0.4, 0.6), + "i", + {"component_spectral_slopes": [1.0]}, + ) + + with pytest.raises(ValueError, match="must be in"): + component_weights_for_band( + (0.4, 0.6), + "i", + {"component_spectral_slopes": [1.0, 0.0], "min_weight": 0.5}, + ) + + register_observatory( + name="ZeroWavelengthReferenceObs", + observatory_class=DummyObservatory, + bands=["ZW1", "ZW2"], + ) + with pytest.raises(ValueError, match="reference band wavelength"): + component_weights_for_band( + (0.4, 0.6), + "i", + {"component_spectral_slopes": [1.0, 0.0], "reference_band": "ZW1"}, + ) + + +def test_edge_apodized_image_only_tapers_edges(): + image = np.ones((7, 7)) + np.testing.assert_allclose(edge_apodized_image(image, edge_width=0), image) + default_tapered = edge_apodized_image(image) + assert default_tapered[0, 0] == 0 + + tapered = edge_apodized_image(image, edge_width=2) + + assert np.all(tapered[0, :] == 0) + assert np.all(tapered[:, 0] == 0) + assert tapered[3, 3] == pytest.approx(1.0) + assert tapered[1, 3] > tapered[0, 3] + + +def test_radial_color_gradient_image_preserves_flux_and_reference_band(): + image = np.ones((9, 9)) + color_gradient = {"grad_color": -0.4, "reference_band": "F814W"} + + assert ( + radial_color_gradient_image( + image=image, + band=None, + color_gradient=color_gradient, + angular_size=0.3, + pixel_scale=0.03, + ) + is image + ) + with pytest.raises(ValueError, match="must be a dictionary"): + radial_color_gradient_image( + image=image, + band="i", + color_gradient="bad", + angular_size=0.3, + pixel_scale=0.03, + ) + + reference = radial_color_gradient_image( + image=image, + band="F814W", + color_gradient=color_gradient, + angular_size=0.3, + pixel_scale=0.03, + ) + np.testing.assert_allclose(reference, image) + + chromatic = radial_color_gradient_image( + image=image, + band="i", + color_gradient=color_gradient, + angular_size=0.3, + pixel_scale=0.03, + ) + assert not np.allclose(chromatic, image) + assert np.sum(chromatic) == pytest.approx(np.sum(image)) + + no_gradient = radial_color_gradient_image( + image=image, + band="i", + color_gradient={"grad_color": 0.0}, + angular_size=0.3, + pixel_scale=0.03, + ) + np.testing.assert_allclose(no_gradient, image)