From 0729d0308b3d6aaeca8ad9e14d9e276eb79e274f Mon Sep 17 00:00:00 2001 From: Michael Larson Date: Sun, 26 Jul 2026 21:25:18 -0400 Subject: [PATCH 1/4] First pass at folding the extended source stuff into the normal KingFitter code. --- kingmaker/fitting.py | 208 ++++++++++++++++++++++++++---------------- kingmaker/utils.py | 47 +++++++++- kingmaker/wrapper.py | 140 +++++++++++++++++++++------- tests/test_fitting.py | 24 ++--- tests/test_wrapper.py | 31 ++++--- 5 files changed, 310 insertions(+), 140 deletions(-) diff --git a/kingmaker/fitting.py b/kingmaker/fitting.py index d20008b..6448d8c 100644 --- a/kingmaker/fitting.py +++ b/kingmaker/fitting.py @@ -6,7 +6,7 @@ from .distribution import _cdf_and_gradient from .pdf import KingPDF -from .utils import angular_distance +from .utils import angular_distance, sample_with_extension class KingPSFFitter: @@ -55,13 +55,20 @@ class KingPSFFitter: Spectral indices (gamma) for reweighting. Default is [2.0]. angular_cutoff : float, optional Maximum angular separation for King PDF. Default is pi. + extension_grid : array-like, optional + Source extension radii in radians, non-negative. Each event's true + position is displaced by a Rayleigh(extension)-magnitude offset + before computing dpsi. Default is [0.0], the point-source case. + rng : np.random.Generator, optional + Random number generator for the extension smearing draws. Defaults + to a fixed seed so repeated fits reproduce the same result. Attributes ---------- fit_alpha : ndarray - Fitted alpha parameters for each bin. + Fitted alpha parameters, shape (n_extension, n_gamma, *bins). fit_beta : ndarray - Fitted beta parameters for each bin. + Fitted beta parameters, shape (n_extension, n_gamma, *bins). histograms : ndarray Histogram values for each bin. uncertainties : ndarray @@ -86,6 +93,7 @@ def __init__( true_energy_name: str = "trueE", spectral_indices: Optional[Union[List[float], npt.NDArray[np.floating]]] = None, angular_cutoff: float = np.pi, + extension_grid: Optional[Union[List[float], npt.NDArray[np.floating]]] = None, ) -> None: """Initialize the KingPSFFitter.""" self.signal_events = signal_events @@ -103,6 +111,14 @@ def __init__( ) self.angular_cutoff = angular_cutoff + self.extension_grid = np.sort( + np.atleast_1d( + np.asarray(extension_grid if extension_grid is not None else [0.0], dtype=np.float64) + ) + ) + if np.any(self.extension_grid < 0): + raise ValueError("extension_grid contains negative values.") + # Initialize King PDF self.king_pdf = KingPDF(angular_cutoff=angular_cutoff) @@ -112,12 +128,16 @@ def __init__( self.bin_names = list(self.parametrization_bins.keys()) self.parametrization_shape = [len(bins) - 1 for bins in self.parametrization_bins.values()] - # Calculate angular distances - self.dpsi = angular_distance( - self.signal_events["ra"], - self.signal_events["dec"], - self.signal_events[self.true_ra_name], - self.signal_events[self.true_dec_name], + # Find default alpha value for failing bins. + self._alpha_guess = float( + np.median( + angular_distance( + self.signal_events["ra"], + self.signal_events["dec"], + self.signal_events[self.true_ra_name], + self.signal_events[self.true_dec_name], + ) + ) ) # Bin events @@ -265,30 +285,36 @@ def _bin_events(self) -> Dict[str, npt.NDArray[np.integer]]: def _initialize_storage(self) -> None: """Initialize arrays to store fit results and diagnostics.""" - shape_with_gamma = [len(self.spectral_indices)] + self.parametrization_shape + shape = [len(self.extension_grid), len(self.spectral_indices)] + self.parametrization_shape # Fit parameters - self.fit_alpha = np.full(shape_with_gamma, np.median(self.dpsi)) - self.fit_beta = np.full(shape_with_gamma, 2.25) + self.fit_alpha = np.full(shape, self._alpha_guess) + self.fit_beta = np.full(shape, 2.25) # Diagnostics - self.histograms = np.zeros(shape_with_gamma + [self.dpsi_nbins], dtype=float) - self.uncertainties = np.zeros(shape_with_gamma + [self.dpsi_nbins], dtype=float) - self.dpsi_bins = np.zeros(shape_with_gamma + [self.dpsi_nbins + 1], dtype=float) - self.fit_quality = np.zeros(shape_with_gamma, dtype=float) - self.event_counts = np.zeros(shape_with_gamma, dtype=int) - - def fit_all_bins(self, verbose: bool = True) -> Dict[str, npt.NDArray]: + self.histograms = np.zeros(shape + [self.dpsi_nbins], dtype=float) + self.uncertainties = np.zeros(shape + [self.dpsi_nbins], dtype=float) + self.dpsi_bins = np.zeros(shape + [self.dpsi_nbins + 1], dtype=float) + self.fit_quality = np.zeros(shape, dtype=float) + self.event_counts = np.zeros(shape, dtype=int) + + def fit_all_bins( + self, verbose: bool = True, rng: Optional[np.random.Generator] = None + ) -> Dict[str, npt.NDArray]: """ Fit King PSF parameters in all bins. - Iterates over all bins defined by parametrization_bins and spectral_indices, - fitting King distribution parameters to the angular error distribution. + Iterates over all bins defined by parametrization_bins, spectral_indices, + and extension_grid, fitting King distribution parameters to the angular + error distribution. Parameters ---------- verbose : bool, optional Print progress information. Default is True. + rng : np.random.Generator, optional + Random number generator for the extension smearing draws. Defaults + to a fixed seed so repeated fits reproduce the same result. Returns ------- @@ -301,65 +327,87 @@ def fit_all_bins(self, verbose: bool = True) -> Dict[str, npt.NDArray]: - 'dpsi_bins': angular error bin edges - 'fit_quality': chi-square values - 'event_counts': number of events per bin + - 'parametrization_bins': bin edges + - 'extension_grid': the extension values fit """ + if rng is None: + rng = np.random.default_rng(0) + if verbose: print(f"Fitting King PSF in {np.prod(self.parametrization_shape)} bins...") print(f" Spectral indices: {self.spectral_indices}") + print(f" Extensions: {self.extension_grid}") print(f" Binning dimensions: {self.bin_names}") - # Iterate over spectral indices - for g_idx, gamma in enumerate(self.spectral_indices): - if verbose: - print(f"\n Spectral index γ = {gamma:.2f}") - - # Calculate event weights - if self.weight_field is not None: - weights = self.signal_events[self.weight_field] * self.signal_events[ - self.true_energy_name - ] ** (-gamma) - else: - weights = np.ones(len(self.signal_events)) - - # Iterate over all bin combinations - n_fitted = 0 - n_skipped = 0 - - total_bins = np.prod(self.parametrization_shape) - for bin_indices in tqdm(np.ndindex(*self.parametrization_shape), total=total_bins): - flat_idx = int(np.ravel_multi_index(bin_indices, tuple(self.parametrization_shape))) - event_idx = self._event_sort_order[ - self._bin_boundaries[flat_idx] : self._bin_boundaries[flat_idx + 1] - ] - - if self.remove_weight_outliers and len(event_idx) > 0: - bin_weights = weights[event_idx] - idx_range = [ - int(len(bin_weights) * self.weight_outlier_percentiles[0] / 100), - int(len(bin_weights) * self.weight_outlier_percentiles[1] / 100), - ] - idx = np.digitize(bin_weights, np.unique(bin_weights)) - event_idx = event_idx[(idx_range[0] <= idx) & (idx <= idx_range[1])] - - n_events = len(event_idx) - param_idx = tuple([g_idx] + list(bin_indices)) - self.event_counts[param_idx] = n_events - - # Skip if insufficient events - if n_events < self.minimum_counts: - n_skipped += 1 - continue - - # Fit this bin - success = self._fit_single_bin(event_idx, weights, param_idx) - if success: - n_fitted += 1 + reco_ra = self.signal_events["ra"] + reco_dec = self.signal_events["dec"] + true_ra = self.signal_events[self.true_ra_name] + true_dec = self.signal_events[self.true_dec_name] + trueE = self.signal_events[self.true_energy_name] if self.weight_field is not None else None + ow = self.signal_events[self.weight_field] if self.weight_field is not None else None + + n_fitted = 0 + n_skipped = 0 + total_bins = np.prod(self.parametrization_shape) + for bin_indices in tqdm(np.ndindex(*self.parametrization_shape), total=total_bins): + flat_idx = int(np.ravel_multi_index(bin_indices, tuple(self.parametrization_shape))) + event_idx = self._event_sort_order[ + self._bin_boundaries[flat_idx] : self._bin_boundaries[flat_idx + 1] + ] + if len(event_idx) == 0: + continue + + bin_reco_ra = reco_ra[event_idx] + bin_reco_dec = reco_dec[event_idx] + bin_true_ra = true_ra[event_idx] + bin_true_dec = true_dec[event_idx] + bin_trueE = trueE[event_idx] if trueE is not None else None + bin_ow = ow[event_idx] if ow is not None else None + + for ext_idx, extension in enumerate(self.extension_grid): + if extension == 0.0: + bin_dpsi = angular_distance(bin_reco_ra, bin_reco_dec, bin_true_ra, bin_true_dec) else: - n_skipped += 1 - - if verbose: - print(f" Fitted {n_fitted} bins, skipped {n_skipped} bins") + smeared_ra, smeared_dec = sample_with_extension( + bin_true_ra, bin_true_dec, extension, rng + ) + bin_dpsi = angular_distance(bin_reco_ra, bin_reco_dec, smeared_ra, smeared_dec) + + for g_idx, gamma in enumerate(self.spectral_indices): + if bin_ow is not None: + bin_weights = bin_ow * bin_trueE ** (-gamma) + else: + bin_weights = np.ones(len(event_idx)) + + local_idx = np.arange(len(event_idx)) + if self.remove_weight_outliers and len(local_idx) > 0: + idx_range = [ + int(len(local_idx) * self.weight_outlier_percentiles[0] / 100), + int(len(local_idx) * self.weight_outlier_percentiles[1] / 100), + ] + idx = np.digitize(bin_weights, np.unique(bin_weights)) + local_idx = local_idx[(idx_range[0] <= idx) & (idx <= idx_range[1])] + + n_events = len(local_idx) + param_idx = (ext_idx, g_idx) + tuple(bin_indices) + self.event_counts[param_idx] = n_events + + # Skip if insufficient events + if n_events < self.minimum_counts: + n_skipped += 1 + continue + + # Fit this bin + success = self._fit_single_bin( + bin_dpsi[local_idx], bin_weights[local_idx], param_idx + ) + if success: + n_fitted += 1 + else: + n_skipped += 1 if verbose: + print(f"\nFitted {n_fitted} bins, skipped {n_skipped} bins") print("\nFitting complete!") return { @@ -371,6 +419,7 @@ def fit_all_bins(self, verbose: bool = True) -> Dict[str, npt.NDArray]: "fit_quality": self.fit_quality, "event_counts": self.event_counts, "parametrization_bins": self.parametrization_bins, # type: ignore[dict-item] + "extension_grid": self.extension_grid, } def _cdf_chi2(self, cdf_hist, cdf_variance, bins, alpha, beta): @@ -404,8 +453,8 @@ def _cdf_chi2(self, cdf_hist, cdf_variance, bins, alpha, beta): def _fit_single_bin( self, - event_idx: npt.NDArray[np.intp], - weights: npt.NDArray[np.floating], + masked_dpsi: npt.NDArray[np.floating], + masked_weights: npt.NDArray[np.floating], param_idx: Tuple[int, ...], ) -> bool: """ @@ -413,10 +462,10 @@ def _fit_single_bin( Parameters ---------- - event_idx : ndarray - Integer indices of events in this bin. - weights : ndarray - Event weights. + masked_dpsi : ndarray + Angular errors for events in this bin. + masked_weights : ndarray + Event weights for events in this bin, not yet normalized. param_idx : tuple Index tuple for storing results. @@ -425,10 +474,7 @@ def _fit_single_bin( bool True if fit succeeded, False otherwise. """ - # Extract events in this bin - masked_dpsi = self.dpsi[event_idx] - masked_weights = weights[event_idx] - masked_weights /= masked_weights.sum() # Normalize + masked_weights = masked_weights / masked_weights.sum() # Normalize # Create bins for this subset. Also calculate the # phase space parameter while we're here. We'll need diff --git a/kingmaker/utils.py b/kingmaker/utils.py index 07ea426..ecf38dd 100644 --- a/kingmaker/utils.py +++ b/kingmaker/utils.py @@ -1,4 +1,4 @@ -from typing import Tuple, Union +from typing import Optional, Tuple, Union import numpy as np import numpy.typing as npt from numba import njit, prange @@ -103,6 +103,51 @@ def angular_distance( return np.arccos(np.minimum(np.maximum(cosDist, -1.0), 1.0)) # type: ignore[no-any-return] +def sample_with_extension( + true_ra: Union[float, npt.NDArray[np.floating]], + true_dec: Union[float, npt.NDArray[np.floating]], + extension: Union[float, npt.NDArray[np.floating]], + rng: Optional[np.random.Generator] = None, +) -> Tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]: + """ + Sample a position offset from (true_ra, true_dec) by a Rayleigh(extension) + magnitude at a uniformly random bearing, simulating a source's angular extent. + + Parameters + ---------- + true_ra, true_dec : float or ndarray + True source position(s) in radians. + extension : float or ndarray + Rayleigh scale of the angular offset, in radians. + rng : np.random.Generator, optional + Random number generator. If None, uses np.random.default_rng(). + + Returns + ------- + ra, dec : ndarray + Sampled positions in radians. + """ + if rng is None: + rng = np.random.default_rng() + + true_ra, true_dec, extension = np.broadcast_arrays(true_ra, true_dec, extension) + d = rng.rayleigh(extension) + theta = rng.uniform(0, 2 * np.pi, size=np.shape(d)) + + sin_dec = np.sin(true_dec) + cos_dec = np.cos(true_dec) + sin_d = np.sin(d) + cos_d = np.cos(d) + + sin_dec2 = np.clip(sin_dec * cos_d + cos_dec * sin_d * np.cos(theta), -1.0, 1.0) + dec = np.arcsin(sin_dec2) + ra = np.mod( + true_ra + np.arctan2(np.sin(theta) * sin_d * cos_dec, cos_d - sin_dec * sin_dec2), + 2 * np.pi, + ) + return ra, dec + + @njit(cache=True) def _premask_events( ra_i: float, diff --git a/kingmaker/wrapper.py b/kingmaker/wrapper.py index 2928c70..0f7590d 100644 --- a/kingmaker/wrapper.py +++ b/kingmaker/wrapper.py @@ -42,6 +42,7 @@ class then fits King distribution parameters using the requested parameter binni # Source-level information source_ras: Optional[npt.NDArray[Any]] = None source_decs: Optional[npt.NDArray[Any]] = None + source_extensions: Optional[npt.NDArray[Any]] = None # Have some place to cache the per-event information so we don't need to # recalculate it every time we evaluate the PDF. @@ -81,6 +82,7 @@ def __init__( marginalization_points_beta: Optional[npt.NDArray[np.floating]] = None, marginalization_n_signed_delta_dec: int = 200, marginalization_n_ra_bins: int = 100, + extension_grid: Optional[npt.NDArray[np.floating]] = None, ): # Store some of the configuration parameters for this instance. # Note that we don't need to store the signal events, dpsi_nbins, @@ -117,6 +119,7 @@ def __init__( true_ra_name=true_ra_name, true_dec_name=true_dec_name, true_energy_name=true_energy_name, + extension_grid=extension_grid, ) fitted_parameters = fitter.fit_all_bins(verbose=True) if cache_parameters and (cache_name is not None): @@ -136,9 +139,20 @@ def __init__( self.keys.append(key) self.bin_centers.append((edges[:-1] + edges[1:]) / 2) - # And grab the fitted alpha/beta arrays + # Extension grid (bin-center-style values, nearest-snapped per source). + if "extension_grid" not in fitted_parameters: + raise ValueError(f"Cache {cache_name!r} has no extension_grid. Delete it and refit.") + self.extension_grid = np.sort(np.atleast_1d(fitted_parameters["extension_grid"])) + + # And grab the fitted alpha/beta arrays, shape (n_extension, n_gamma, *bins). self.alpha_values = fitted_parameters["alpha"] self.beta_values = fitted_parameters["beta"] + expected_ndim = 2 + len(self.parametrization_bins) + if self.alpha_values.ndim != expected_ndim or self.beta_values.ndim != expected_ndim: + raise ValueError( + f"Cached alpha/beta have {self.alpha_values.ndim} dimensions, expected " + f"{expected_ndim} (extension, gamma, *bins). Delete {cache_name!r} and refit." + ) # Instantiate the PDF object. self.king_pdf = KingPDF(angular_cutoff=angular_cutoff) @@ -188,7 +202,12 @@ def _events_match(self, events: npt.NDArray[Any]) -> bool: result &= np.array_equal(self.events["dec"][::10], events["dec"][::10]) return result - def _sources_match(self, source_ras: npt.NDArray[Any], source_decs: npt.NDArray[Any]) -> bool: + def _sources_match( + self, + source_ras: npt.NDArray[Any], + source_decs: npt.NDArray[Any], + source_extensions: Optional[npt.NDArray[Any]] = None, + ) -> bool: if self.source_ras is None: return False if self.source_decs is None: @@ -201,6 +220,10 @@ def _sources_match(self, source_ras: npt.NDArray[Any], source_decs: npt.NDArray[ return False if len(self.source_decs) != len(source_decs): return False + if source_extensions is not None and not np.array_equal( + self.source_extensions, source_extensions + ): + return False return np.array_equal(self.source_ras, source_ras) and np.array_equal( self.source_decs, source_decs ) @@ -210,6 +233,7 @@ def set_events( events: npt.NDArray[Any], source_ras: Optional[npt.NDArray[np.floating]], source_decs: Optional[npt.NDArray[np.floating]], + source_extensions: Optional[npt.NDArray[np.floating]] = None, ) -> None: """ Cache per-event King PDF values for each spectral index ahead of a call @@ -220,8 +244,9 @@ def set_events( computed, and the King PDF is evaluated and cached for every spectral index in ``spectral_indices``. This must be called before :meth:`evaluate_pdf`. Calling it again with the same ``events``, - ``source_ras``, and ``source_decs`` as the previous call is a cheap - no-op, so it is safe to call once per trial without checking first. + ``source_ras``, ``source_decs``, and ``source_extensions`` as the + previous call is a cheap no-op, so it is safe to call once per trial + without checking first. Parameters ---------- @@ -234,6 +259,10 @@ def set_events( source_decs : ndarray Source declination(s) in radians. Must have the same length as ``source_ras``. + source_extensions : ndarray, optional + Source extension radii in radians, nearest-snapped to the fitted + ``extension_grid``. Defaults to zero (point source) for every + source. Raises ------ @@ -246,7 +275,9 @@ def set_events( Support for multiple simultaneous sources is experimental and logs a one-time warning; results should be checked carefully in that case. """ - if self._events_match(events) and self._sources_match(source_ras, source_decs): + if self._events_match(events) and self._sources_match( + source_ras, source_decs, source_extensions + ): return self.events = events @@ -271,6 +302,17 @@ def set_events( ) self.multiple_source_warning_logged = True + self.source_extensions = ( + np.zeros(len(source_ras)) + if source_extensions is None + else np.asarray(source_extensions, dtype=np.float64) + ) + if len(self.source_extensions) != len(source_ras): + raise ValueError( + "source_extensions must have the same length as source_ras and source_decs." + ) + ext_idx_per_source = self._nearest_extension_index(self.source_extensions) + # Calculate the (event, source) angular distances via a single compiled # pass that pre-filters on a dec/RA bounding box before the haversine, # returning only the (event, source) pairs within the cutoff. @@ -279,6 +321,7 @@ def set_events( event_rows, event_cols, self.event_distances = _pre_mask_and_distance( events["ra"], events["dec"], source_ras, source_decs, cutoff ) + ext_idx_pairs = ext_idx_per_source[event_cols] # alpha/beta/norm are looked up once per event (they don't depend on # source), so event_mask is the per-event OR across sources, and each @@ -304,9 +347,9 @@ def set_events( # event_distances are already within angular_cutoff by construction. values = self.king_pdf.pdf_from_norm( self.event_distances, - all_alpha[i][pair_position], - all_beta[i][pair_position], - all_norm[i][pair_position], + all_alpha[ext_idx_pairs, i, pair_position], + all_beta[ext_idx_pairs, i, pair_position], + all_norm[ext_idx_pairs, i, pair_position], ) self._pdf_matrices.append( csr_array( @@ -316,39 +359,66 @@ def set_events( ) ) - # Marginalized path: precompute one sparse (n_events, n_sources) matrix per - # spectral index. The first call establishes the sparsity structure; all - # subsequent calls pass that result as mask= so interpn operates on the same - # fixed (event, source) pairs. The mask path in MarginalizedKingPDF.evaluate() - # does not filter values > 0, so every matrix in _marg_matrices is guaranteed to - # share identical .indices and .indptr — a requirement for the scalar lerp - # on .data in evaluate_marginalized_pdf. + # Marginalized path: marg_source_decs is assumed to correspond 1:1 with + # source_ras/source_decs/source_extensions. if self.mkpdf is not None: + if len(self._marg_source_decs) != len(source_ras): + raise ValueError( + "marginalization_source_decs must have the same length as " + "source_ras/source_decs/source_extensions passed to set_events." + ) all_alpha_full, all_beta_full = self._lookup_all_events_grid(events) + + unique_ext = np.unique(ext_idx_per_source) + group_source_idx = [np.flatnonzero(ext_idx_per_source == e) for e in unique_ext] + group_masks: List[Optional[csr_array]] = [None] * len(unique_ext) + self._marg_matrices = [] - mask_marg: Optional[csr_array] = None for i in range(len(self.spectral_indices)): - mat = self.mkpdf.evaluate( - self._marg_source_decs, - events["dec"], - all_alpha_full[i], - all_beta_full[i], - mask=mask_marg, + data_parts, row_parts, col_parts = [], [], [] + for g, (ext_idx, source_idx) in enumerate(zip(unique_ext, group_source_idx)): + mat = self.mkpdf.evaluate( + self._marg_source_decs[source_idx], + events["dec"], + all_alpha_full[ext_idx, i], + all_beta_full[ext_idx, i], + mask=group_masks[g], + ) + if group_masks[g] is None: + group_masks[g] = mat + coo = mat.tocoo() + data_parts.append(coo.data) + row_parts.append(coo.row) + col_parts.append(source_idx[coo.col]) + self._marg_matrices.append( + csr_array( + ( + np.concatenate(data_parts), + (np.concatenate(row_parts), np.concatenate(col_parts)), + ), + shape=(len(events), len(self._marg_source_decs)), + dtype=np.float64, + ) ) - if mask_marg is None: - mask_marg = mat - self._marg_matrices.append(mat) return + def _nearest_extension_index(self, source_extensions): + """Nearest-bin lookup of extension_grid index for each source.""" + centers = self.extension_grid + if len(centers) == 1: + return np.zeros(len(source_extensions), dtype=np.intp) + i = np.searchsorted(centers, source_extensions).clip(1, len(centers) - 1) + return np.where(source_extensions - centers[i - 1] < centers[i] - source_extensions, i - 1, i) + def _lookup_event_grid(self, events): """ - Nearest-bin lookup of alpha, beta, and norm for each (unmasked) event. + Nearest-bin lookup of alpha, beta, and norm for each (unmasked) event, + for every extension in extension_grid. Shared by :meth:`get_alpha_beta` and :meth:`set_events` so the nearest-bin index computation is only ever done once per call. """ - # Nearest-bin lookup. Extracts each field individually after masking. def index(centers, values): i = np.searchsorted(centers, values).clip(1, len(centers) - 1) return np.where(values - centers[i - 1] < centers[i] - values, i - 1, i) @@ -358,13 +428,13 @@ def index(centers, values): for i, key in enumerate(self.keys) ) - idx = (slice(None), *event_indices) + idx = (slice(None), slice(None), *event_indices) return self.alpha_values[idx], self.beta_values[idx], self.norm_values[idx] def _lookup_all_events_grid(self, events): """ Nearest-bin lookup of alpha and beta for every event, without applying - ``event_mask``. + ``event_mask``, for every extension in extension_grid. Used by :meth:`set_events` to supply per-event PSF parameters for :meth:`MarginalizedKingPDF.evaluate`, which performs its own angular @@ -372,8 +442,8 @@ def _lookup_all_events_grid(self, events): Returns ------- - alpha : ndarray, shape (n_gamma, n_events) - beta : ndarray, shape (n_gamma, n_events) + alpha : ndarray, shape (n_extension, n_gamma, n_events) + beta : ndarray, shape (n_extension, n_gamma, n_events) """ def index(centers, values): @@ -383,10 +453,10 @@ def index(centers, values): event_indices = tuple( index(self.bin_centers[i], events[key]) for i, key in enumerate(self.keys) ) - idx = (slice(None), *event_indices) + idx = (slice(None), slice(None), *event_indices) return self.alpha_values[idx], self.beta_values[idx] - def get_alpha_beta(self, events): + def get_alpha_beta(self, events, extension_index: int = 0): """ Look up fitted alpha/beta parameters for each event via nearest-bin lookup. @@ -399,6 +469,8 @@ def get_alpha_beta(self, events): Events to look up. Must contain the fields referenced by ``parametrization_bins``. Only events selected by the mask set in the most recent :meth:`set_events` call are returned. + extension_index : int, optional + Index into ``extension_grid``. Default is 0. Returns ------- alpha : ndarray, shape (n_gamma, n_masked_events) @@ -407,7 +479,7 @@ def get_alpha_beta(self, events): Fitted beta values for each spectral index and event. """ alpha, beta, _ = self._lookup_event_grid(events) - return alpha, beta + return alpha[extension_index], beta[extension_index] def get_alpha_beta_gamma(self, gamma, events=None, alpha=None, beta=None): """ diff --git a/tests/test_fitting.py b/tests/test_fitting.py index 13547d8..23e7e67 100644 --- a/tests/test_fitting.py +++ b/tests/test_fitting.py @@ -81,11 +81,11 @@ def test_bin_names(self, fitter): assert fitter.bin_names == ["aux"] def test_fit_alpha_shape_before_fitting(self, fitter): - # Shape should be (n_spectral_indices=1, n_bins=3). - assert fitter.fit_alpha.shape == (1, 3) + # Shape should be (n_extension=1, n_spectral_indices=1, n_bins=3). + assert fitter.fit_alpha.shape == (1, 1, 3) def test_fit_beta_shape_before_fitting(self, fitter): - assert fitter.fit_beta.shape == (1, 3) + assert fitter.fit_beta.shape == (1, 1, 3) def test_explicit_bin_edges_count(self): """Passing k explicit edges should produce k-1 bins.""" @@ -135,29 +135,29 @@ def result(self): def test_all_bins_have_enough_events(self, result): fitter, _, _ = result - assert np.all(fitter.event_counts[0] >= 100) + assert np.all(fitter.event_counts[0, 0] >= 100) def test_alpha_consistent_across_bins(self, result): """Coefficient of variation of fitted alpha should be small (<30%).""" fitter, _, _ = result - alphas = fitter.fit_alpha[0] + alphas = fitter.fit_alpha[0, 0] assert alphas.std() / alphas.mean() < 0.30 def test_beta_consistent_across_bins(self, result): """Coefficient of variation of fitted beta should be small (<30%).""" fitter, _, _ = result - betas = fitter.fit_beta[0] + betas = fitter.fit_beta[0, 0] assert betas.std() / betas.mean() < 0.30 def test_mean_alpha_roughly_accurate(self, result): """Mean fitted alpha should be within 30% of the true value.""" fitter, alpha_true, _ = result - assert_allclose(fitter.fit_alpha[0].mean(), alpha_true, rtol=0.30) + assert_allclose(fitter.fit_alpha[0, 0].mean(), alpha_true, rtol=0.30) def test_mean_beta_roughly_accurate(self, result): """Mean fitted beta should be within 40% of the true value.""" fitter, _, beta_true = result - assert_allclose(fitter.fit_beta[0].mean(), beta_true, rtol=0.40) + assert_allclose(fitter.fit_beta[0, 0].mean(), beta_true, rtol=0.40) # --------------------------------------------------------------------------- @@ -204,22 +204,22 @@ def test_parametrization_shape(self, result): def test_alpha_decreases_with_energy(self, result): """Higher-energy bins should produce a smaller fitted alpha.""" - alphas = result.fit_alpha[0] + alphas = result.fit_alpha[0, 0] assert alphas[0] > alphas[1] > alphas[2] def test_beta_increases_with_energy(self, result): """Higher-energy bins should produce a larger fitted beta.""" - betas = result.fit_beta[0] + betas = result.fit_beta[0, 0] assert betas[0] < betas[1] < betas[2] def test_alpha_values_per_bin(self, result): """Fitted alpha per bin should be within 40% of the true value.""" - alphas = result.fit_alpha[0] + alphas = result.fit_alpha[0, 0] for i, (alpha_true, _, _) in enumerate(self._GROUP_PARAMS): assert_allclose(alphas[i], alpha_true, rtol=0.40, err_msg=f"bin {i}: alpha mismatch") def test_beta_values_per_bin(self, result): """Fitted beta per bin should be within 50% of the true value.""" - betas = result.fit_beta[0] + betas = result.fit_beta[0, 0] for i, (_, beta_true, _) in enumerate(self._GROUP_PARAMS): assert_allclose(betas[i], beta_true, rtol=0.50, err_msg=f"bin {i}: beta mismatch") diff --git a/tests/test_wrapper.py b/tests/test_wrapper.py index 3757ca6..38db66d 100644 --- a/tests/test_wrapper.py +++ b/tests/test_wrapper.py @@ -18,18 +18,24 @@ SPECTRAL_INDICES = np.array([1.0, 2.0, 3.0]) BIN_EDGES = np.array([0.0, 1.0, 2.0, 3.0]) # bin centers: 0.5, 1.5, 2.5 +EXTENSION_GRID = np.array([0.0]) # single point-source extension +# shape (n_extension=1, n_gamma=3, n_bins=3) ALPHA_VALUES = np.array( [ - [np.radians(0.5), np.radians(1.0), np.radians(1.5)], - [np.radians(0.6), np.radians(1.1), np.radians(1.6)], - [np.radians(0.7), np.radians(1.2), np.radians(1.7)], + [ + [np.radians(0.5), np.radians(1.0), np.radians(1.5)], + [np.radians(0.6), np.radians(1.1), np.radians(1.6)], + [np.radians(0.7), np.radians(1.2), np.radians(1.7)], + ] ] ) BETA_VALUES = np.array( [ - [2.0, 2.5, 3.0], - [2.1, 2.6, 3.1], - [2.2, 2.7, 3.2], + [ + [2.0, 2.5, 3.0], + [2.1, 2.6, 3.1], + [2.2, 2.7, 3.2], + ] ] ) @@ -41,6 +47,7 @@ def _make_likelihood(tmp_path, angular_cutoff=np.pi): parametrization_bins=np.array({"aux": BIN_EDGES}, dtype=object), alpha=ALPHA_VALUES, beta=BETA_VALUES, + extension_grid=EXTENSION_GRID, ) return KingSpatialLikelihood( signal_events=np.empty(0), @@ -95,8 +102,8 @@ def test_matches_direct_pdf_computation(self, likelihood): for gamma_idx, gamma in enumerate(SPECTRAL_INDICES): result = likelihood.evaluate_pdf(events, gamma=gamma).toarray().ravel() - alpha = ALPHA_VALUES[gamma_idx][bin_idx] - beta = BETA_VALUES[gamma_idx][bin_idx] + alpha = ALPHA_VALUES[0][gamma_idx][bin_idx] + beta = BETA_VALUES[0][gamma_idx][bin_idx] expected = king_pdf.pdf(dist, alpha, beta) np.testing.assert_allclose(result, expected, rtol=1e-10) @@ -145,8 +152,8 @@ def test_matches_direct_pdf_computation_per_source(self, likelihood): bin_idx = np.array([_bin_index(a) for a in events["aux"]]) gamma_idx = int(np.searchsorted(SPECTRAL_INDICES, 2.0)) - alpha = ALPHA_VALUES[gamma_idx][bin_idx] - beta = BETA_VALUES[gamma_idx][bin_idx] + alpha = ALPHA_VALUES[0][gamma_idx][bin_idx] + beta = BETA_VALUES[0][gamma_idx][bin_idx] king_pdf = KingPDF(angular_cutoff=likelihood.king_pdf.angular_cutoff) dense = result.toarray() @@ -178,8 +185,8 @@ def test_sparse_structure_with_partial_and_overlapping_coverage(self, tmp_path): bin_idx = np.array([_bin_index(a) for a in events["aux"]]) gamma_idx = int(np.searchsorted(SPECTRAL_INDICES, 2.0)) - alpha = ALPHA_VALUES[gamma_idx][bin_idx] - beta = BETA_VALUES[gamma_idx][bin_idx] + alpha = ALPHA_VALUES[0][gamma_idx][bin_idx] + beta = BETA_VALUES[0][gamma_idx][bin_idx] king_pdf = KingPDF(angular_cutoff=cutoff) for j, (src_ra, src_dec) in enumerate(zip(src_ras, src_decs)): From a4e6ae75105dea47ce7fcbb174a21003f830fd31 Mon Sep 17 00:00:00 2001 From: Michael Larson Date: Sun, 26 Jul 2026 21:46:54 -0400 Subject: [PATCH 2/4] Remove old extended source class since extended sources are folded into the fitter now --- kingmaker/pdf.py | 460 +---------------------------------------------- 1 file changed, 1 insertion(+), 459 deletions(-) diff --git a/kingmaker/pdf.py b/kingmaker/pdf.py index ff0c6f8..4f193d3 100644 --- a/kingmaker/pdf.py +++ b/kingmaker/pdf.py @@ -4,8 +4,7 @@ import healpy as hp from scipy.interpolate import interpn from scipy.sparse import csr_array -from scipy.special import legendre_p_all, sph_harm_y_all, gammaln -from numpy.polynomial.laguerre import laggauss +from scipy.special import legendre_p_all, sph_harm_y_all from .distribution import _log10pi from .distribution import _norm, _unnormalized_pdf, _unnormalized_cdf @@ -1116,460 +1115,3 @@ def sample( colatitude, longitude = hp.pix2ang(self.nside, pixel_indices) return longitude, np.pi / 2 - colatitude # reco_ra, reco_dec - - -class ExtendedSourceKingPDF: - """ - King PSF convolved with a Rayleigh (Gaussian) source extension. - - Precomputes a 4D lookup table of convolved PDF values over - (log10(alpha), log10(beta), log10(extension), psi) using Gauss-Laguerre - quadrature on the inverse-gamma scale mixture representation of the King - distribution. At runtime, evaluates each event via quadrilinear - interpolation into this table. - - Both the King PSF and the Rayleigh extension are axially symmetric, so the - convolution depends only on the scalar angular distance psi between the - event and the source — not on the full sky coordinates of either. - - Parameters - ---------- - angular_cutoff : float, optional - Maximum angular separation in radians. Default is pi. - maximum_sigma : float, optional - Number of source-extension radii beyond the King angular_cutoff - that ``evaluate()`` considers for each source. For a source with - extension r₀, events at angular distance greater than - ``maximum_sigma * r₀ + angular_cutoff`` are skipped. Default is 3. - points_alpha : ndarray, optional - Alpha grid points in radians. Default: 30 log-spaced values from - 0.05 degrees to pi. - points_beta : ndarray, optional - Beta grid points. Default: 20 log-spaced values from ~1.023 to 10. - points_extension : ndarray, optional - Extension radius grid points in radians. Default: 20 log-spaced values - from 0.05 degrees to 5 degrees. - points_psi : ndarray, optional - Angular separation grid points in radians. Default: 0 followed by 500 - log-spaced values from 1e-4 rad to angular_cutoff. - n_quad : int, optional - Number of Gauss-Laguerre quadrature nodes for the scale mixture - integral. Default is 32. - """ - - def __init__( - self, - *, - angular_cutoff: float = np.pi, - maximum_sigma: float = 3.0, - points_alpha: Optional[npt.NDArray[np.floating]] = None, - points_beta: Optional[npt.NDArray[np.floating]] = None, - points_extension: Optional[npt.NDArray[np.floating]] = None, - points_psi: Optional[npt.NDArray[np.floating]] = None, - n_quad: int = 32, - ) -> None: - self.angular_cutoff = float(angular_cutoff) - self.maximum_sigma = float(maximum_sigma) - - self.n_quad = n_quad - - self._points_alpha = np.sort( - np.asarray( - points_alpha - if points_alpha is not None - else np.logspace(np.log10(np.radians(0.05)), _log10pi, 30), - dtype=np.float64, - ) - ) - self._points_beta = np.sort( - np.asarray( - points_beta if points_beta is not None else np.logspace(0.01, 1, 20), - dtype=np.float64, - ) - ) - self._points_extension = np.sort( - np.asarray( - points_extension - if points_extension is not None - else np.logspace(np.log10(np.radians(0.05)), np.log10(np.radians(5.0)), 20), - dtype=np.float64, - ) - ) - # The psi table must reach the widest possible search window so that - # interpn stays in-bounds for all events accepted by evaluate(). - # That window is maximum_sigma * max_ext + angular_cutoff, capped at π. - _psi_max = min( - self.maximum_sigma * self._points_extension[-1] + angular_cutoff, - np.pi, - ) - self._points_psi = np.sort( - np.asarray( - points_psi - if points_psi is not None - else np.concatenate([[0.0], np.logspace(-4, np.log10(_psi_max), 500)]), - dtype=np.float64, - ) - ) - - if np.any(self._points_alpha <= 0): - raise ValueError( - "points_alpha contains values <= 0. The King distribution is not defined here." - ) - if np.any(self._points_beta <= 1): - raise ValueError( - "points_beta contains values <= 1. The King distribution is not defined here." - ) - if np.any(self._points_extension <= 0): - raise ValueError( - "points_extension contains values <= 0. Extension radius must be positive." - ) - if np.any(self._points_extension > np.radians(5.0) + 1e-12): - raise ValueError( - "points_extension contains values > 5 degrees. The flat-sky (Rayleigh) " - "approximation error exceeds ~0.5% at this scale; use a vMF extension instead." - ) - - self._log10_points_alpha = np.log10(self._points_alpha) - self._log10_points_beta = np.log10(self._points_beta) - self._log10_points_extension = np.log10(self._points_extension) - - # Gauss-Laguerre nodes and weights for the scale mixture integral. - # The substitution t = beta*alpha^2 / v^2 maps the InvGamma weight - # exp(-beta*alpha^2/v^2) to the standard Gauss-Laguerre form e^{-t}. - self._quad_nodes, self._quad_weights = laggauss(n_quad) - - self._build_table() - - def _scale_mixture_pdf( - self, - psi: npt.NDArray[np.floating], - alpha: npt.NDArray[np.floating], - beta: npt.NDArray[np.floating], - extension: float, - ) -> npt.NDArray[np.floating]: - """ - Evaluate the King–Rayleigh convolution via Gauss-Laguerre quadrature. - - Derivation - ---------- - Consider integrating a Rayleigh PDF over an unknown scale v^2, weighted - by an InvGamma(kappa, c) density: - - integral_0^inf [psi/v^2 * exp(-psi^2/(2v^2))] <- Rayleigh - * [c^kappa/Gamma(kappa) * (v^2)^{-kappa-1} * exp(-c/v^2)] - dv^2 - - The two exponentials combine: exp(-psi^2/(2v^2)) * exp(-c/v^2) - = exp(-(psi^2/2 + c)/v^2). Collecting powers of v^2 gives an integrand - proportional to (v^2)^{-(kappa+2)} * exp(-A/v^2) with A = c + psi^2/2. - That integral is a standard gamma-function result: Gamma(kappa+1)/A^{kappa+1}. - Substituting back: - - result = psi * kappa * c^kappa / (c + psi^2/2)^{kappa+1} - = psi/c * kappa / (1 + psi^2/(2c))^{kappa+1} - - Matching to the flat-sky King PDF psi/alpha^2 * (1-1/beta) - * (1 + psi^2/(2*beta*alpha^2))^{-beta} requires kappa = beta-1 and - c = beta*alpha^2. The King distribution is therefore exactly equal - to an InvGamma-weighted integral of Rayleigh distributions. - - For the convolution with a Rayleigh source extension of radius r0: the - event position is the vector sum of an independent PSF displacement - (Rayleigh with scale v^2) and a source-extent displacement (Rayleigh - with scale r0^2). Adding two independent 2D Gaussian displacements - produces a 2D Gaussian with combined scale v^2 + r0^2. Replacing v^2 - with v^2 + r0^2 inside the integral above gives the convolution. - - The substitution t = c/v^2 (so v^2 = c/t) maps exp(-c/v^2) to exp(-t), - putting the integral in standard Gauss-Laguerre form - integral_0^inf f(t) * e^{-t} dt, evaluated here with fixed nodes and - weights from laggauss(n_quad). After the substitution the integral - becomes: - - p_conv = psi / Gamma(beta-1) - * integral_0^inf t^(beta-1) / (c + r0^2 * t) - * exp(-psi^2 * t / (2*(c + r0^2*t))) - * e^{-t} dt - - Setting r0 = 0 recovers the flat-sky King PDF exactly. Uses psi^2/2 - throughout rather than 1 - cos(psi); error is below 0.5% for psi and - r0 below 5 degrees. - - Parameters - ---------- - psi : ndarray - Angular separation(s) from the source in radians. - alpha : ndarray - King alpha parameter(s) in radians. - beta : ndarray - King beta parameter(s). Must be > 1. - extension : float - Rayleigh source extension radius in radians. Must be <= 5 degrees. - - Returns - ------- - ndarray - Convolved PDF values, same shape as the broadcast of inputs. - Zero wherever psi = 0. - - Raises - ------ - NotImplementedError - If extension exceeds 5 degrees (np.radians(5)), where the - flat-sky approximation error exceeds ~0.5% and a vMF extension - should be used instead. - """ - if float(extension) > np.radians(5.0) + 1e-12: - raise NotImplementedError( - f"extension={np.degrees(float(extension)):.2f} deg exceeds the 5-degree " - "limit for the flat-sky (Rayleigh) approximation. " - "A von Mises-Fisher extension needs to be implemented " - " for larger angular scales." - ) - - psi, alpha, beta = np.broadcast_arrays( - np.asarray(psi, dtype=np.float64), - np.asarray(alpha, dtype=np.float64), - np.asarray(beta, dtype=np.float64), - ) - - r0_sq = float(extension) ** 2 - c = beta * alpha**2 # InvGamma scale: beta * alpha^2 - - # Broadcast data arrays against the quadrature axis - t = self._quad_nodes # (n_quad,) - w = self._quad_weights # (n_quad,) - c_q = c[..., np.newaxis] # (..., 1) - psi_q = psi[..., np.newaxis] # (..., 1) - beta_q = beta[..., np.newaxis] # (..., 1) - - denom = c_q + r0_sq * t # (..., n_quad); always positive - - # t^(beta-1) via log to stay in float64 range for large beta or t - t_pow = np.exp((beta_q - 1.0) * np.log(t)) # (..., n_quad) - - integrand = psi_q * t_pow / denom * np.exp(-(psi_q**2) * t / (2.0 * denom)) - # (..., n_quad); naturally zero when psi = 0 - - quad_sum = (w * integrand).sum(axis=-1) # (...) - - return quad_sum / np.exp(gammaln(beta - 1.0)) - - def _build_table(self) -> None: - """ - Precompute the convolved PDF on the (alpha, beta, extension, psi) grid. - - Evaluates _scale_mixture_pdf over the full four-dimensional parameter - space by looping over extension values (keeping one (n_alpha, n_beta, - n_psi) slice in memory at a time) and stacking the results. Stores the - result as self._table with shape (n_alpha, n_beta, n_extension, n_psi) - for use by interpn at runtime. - """ - # Broadcast axes over (alpha, beta, psi) for each extension slice. - # Shape annotations assume n_alpha, n_beta, n_psi grid sizes. - alpha_g = self._points_alpha[:, np.newaxis, np.newaxis] # (n_alpha, 1, 1) - beta_g = self._points_beta[np.newaxis, :, np.newaxis] # (1, n_beta, 1) - psi_g = self._points_psi[np.newaxis, np.newaxis, :] # (1, 1, n_psi) - - slices = [ - self._scale_mixture_pdf(psi_g, alpha_g, beta_g, float(ext)) - for ext in self._points_extension - ] # each element: (n_alpha, n_beta, n_psi) - - self._table = np.stack(slices, axis=2) # (n_alpha, n_beta, n_extension, n_psi) - - def pdf( - self, - x: Union[float, npt.NDArray[np.floating]], - alpha: Union[float, npt.NDArray[np.floating]], - beta: Union[float, npt.NDArray[np.floating]], - extension: Union[float, npt.NDArray[np.floating]], - ) -> npt.NDArray[np.floating]: - """ - Evaluate the convolved PDF at angular separation(s) x. - - Parameters - ---------- - x : float or ndarray - Angular separation(s) from the source in radians. - alpha : float or ndarray - Per-event King alpha parameter in radians. - beta : float or ndarray - Per-event King beta parameter. - extension : float or ndarray - Source extension radius in radians. - - Returns - ------- - ndarray - Convolved PDF values in probability/steradian. - """ - x = np.asarray(x, dtype=np.float64) - alpha = np.asarray(alpha, dtype=np.float64) - beta = np.asarray(beta, dtype=np.float64) - extension = np.asarray(extension, dtype=np.float64) - x, alpha, beta, extension = np.broadcast_arrays(x, alpha, beta, extension) - shape = x.shape - x = np.atleast_1d(x).ravel() - alpha = np.atleast_1d(alpha).ravel() - beta = np.atleast_1d(beta).ravel() - extension = np.atleast_1d(extension).ravel() - - result = np.zeros(len(x)) - in_bounds = x <= self._points_psi[-1] - - if np.any(in_bounds): - queries = np.column_stack( - [ - np.log10(alpha[in_bounds]), - np.log10(beta[in_bounds]), - np.log10(extension[in_bounds]), - x[in_bounds], - ] - ) - - # _table stores the 1D radial density f(ψ): ∫₀^∞ f dψ = 1. - # Per-steradian density: p(ψ) = f(ψ) / (2π ψ) [flat-sky: dΩ = 2π ψ dψ]. - f_psi = interpn( - ( - self._log10_points_alpha, - self._log10_points_beta, - self._log10_points_extension, - self._points_psi, - ), - self._table, - queries, - method="linear", - bounds_error=True, - ) - - x_in = x[in_bounds] - with np.errstate(invalid="ignore", divide="ignore"): - result[in_bounds] = np.where(x_in > 0.0, f_psi / (2.0 * np.pi * x_in), 0.0) - - return result.reshape(shape) - - def evaluate( - self, - source_ras: npt.NDArray[np.floating], - source_decs: npt.NDArray[np.floating], - source_extensions: npt.NDArray[np.floating], - event_ras: npt.NDArray[np.floating], - event_decs: npt.NDArray[np.floating], - alpha: npt.NDArray[np.floating], - beta: npt.NDArray[np.floating], - *, - mask: Optional[csr_array] = None, - ) -> csr_array: - """ - Evaluate the extended-source convolved PDF for all (event, source) pairs. - - Iterates over sources, applies a declination pre-filter and a full - great-circle distance check to identify pairs within ``angular_cutoff``, - then evaluates the convolved PDF for those pairs using ``pdf()``. On - repeated calls where the source and event positions are unchanged, pass - the result of a previous call as ``mask`` to skip the masking loop and - go straight to vectorized PDF evaluation. - - Parameters - ---------- - source_ras : ndarray, shape (n_sources,) - Source right ascensions in radians. - source_decs : ndarray, shape (n_sources,) - Source declinations in radians. - source_extensions : ndarray, shape (n_sources,) - Per-source angular extension radii in radians. Must be within the - range covered by ``points_extension`` provided at construction. - event_ras : ndarray, shape (n_events,) - Reconstructed event right ascensions in radians. - event_decs : ndarray, shape (n_events,) - Reconstructed event declinations in radians. - alpha : ndarray, shape (n_events,) - Per-event King alpha parameter in radians. - beta : ndarray, shape (n_events,) - Per-event King beta parameter. - mask : csr_array, optional - Sparse array whose nonzero structure encodes the valid - (event, source) pairs. When provided, the masking loop is skipped - and only the indexed pairs are evaluated. Pass the result of a - previous :meth:`evaluate` call to reuse the geometry. - - Returns - ------- - csr_array, shape (n_events, n_sources) - Sparse array of convolved PDF values in probability/steradian, - indexed ``[event_index, source_index]``. - """ - source_ras = np.atleast_1d(np.asarray(source_ras, dtype=np.float64)) - source_decs = np.atleast_1d(np.asarray(source_decs, dtype=np.float64)) - source_extensions = np.atleast_1d(np.asarray(source_extensions, dtype=np.float64)) - event_ras = np.asarray(event_ras, dtype=np.float64) - event_decs = np.asarray(event_decs, dtype=np.float64) - alpha = np.asarray(alpha, dtype=np.float64) - beta = np.asarray(beta, dtype=np.float64) - - n_events = len(event_ras) - n_sources = len(source_ras) - - if mask is not None: - rows, cols = mask.nonzero() - psi = angular_distance( - source_ras[cols], - source_decs[cols], - event_ras[rows], - event_decs[rows], - ) - vals = self.pdf(psi, alpha[rows], beta[rows], source_extensions[cols]) - nonzero = vals > 0.0 - return csr_array( - (vals[nonzero], (rows[nonzero], cols[nonzero])), - shape=(n_events, n_sources), - dtype=np.float64, - ) - - row_chunks, col_chunks, val_chunks = [], [], [] - for source_index, (src_ra, src_dec, src_ext) in enumerate( - zip(source_ras, source_decs, source_extensions) - ): - radius = min( - self.maximum_sigma * src_ext + self.angular_cutoff, - self._points_psi[-1], - ) - dec_mask = np.abs(event_decs - src_dec) <= radius - candidate_indices = np.flatnonzero(dec_mask) - if len(candidate_indices) == 0: - continue - - psi = angular_distance( - src_ra, - src_dec, - event_ras[candidate_indices], - event_decs[candidate_indices], - ) - within_cutoff = psi <= radius - event_indices = candidate_indices[within_cutoff] - if len(event_indices) == 0: - continue - - vals = self.pdf( - psi[within_cutoff], - alpha[event_indices], - beta[event_indices], - np.full(int(within_cutoff.sum()), src_ext), - ) - nonzero = vals > 0.0 - row_chunks.append(event_indices[nonzero]) - col_chunks.append(np.full(int(nonzero.sum()), source_index, dtype=np.intp)) - val_chunks.append(vals[nonzero]) - - if not row_chunks: - return csr_array((n_events, n_sources), dtype=np.float64) - - return csr_array( - ( - np.concatenate(val_chunks), - (np.concatenate(row_chunks), np.concatenate(col_chunks)), - ), - shape=(n_events, n_sources), - dtype=np.float64, - ) From 186d638e116c076aea5bf1787ddd9abf25e40f78 Mon Sep 17 00:00:00 2001 From: Michael Larson Date: Mon, 27 Jul 2026 00:20:55 -0400 Subject: [PATCH 3/4] Update tests --- tests/test_extended_source_king_pdf.py | 310 ------------------------- tests/test_fitting.py | 82 +++++++ tests/test_utils.py | 71 +++++- tests/test_wrapper.py | 93 ++++++++ 4 files changed, 243 insertions(+), 313 deletions(-) delete mode 100644 tests/test_extended_source_king_pdf.py diff --git a/tests/test_extended_source_king_pdf.py b/tests/test_extended_source_king_pdf.py deleted file mode 100644 index 9b4cf17..0000000 --- a/tests/test_extended_source_king_pdf.py +++ /dev/null @@ -1,310 +0,0 @@ -""" -Unit tests for ExtendedSourceKingPDF. - -Covers initialization (composition, not inheritance), pdf() correctness -(normalization, boundary behaviour, shape), and evaluate() correctness -(sparse output, geometry screening, mask reuse). -""" - -import numpy as np -import pytest -from numpy.testing import assert_allclose -from scipy.sparse import csr_array - -from kingmaker.pdf import ExtendedSourceKingPDF, KingPDF - - -# --------------------------------------------------------------------------- -# Shared fixtures -# --------------------------------------------------------------------------- - - -# Minimal grid shared by most pdf() tests; built once per module. -@pytest.fixture(scope="module") -def ext_pdf(): - return ExtendedSourceKingPDF( - points_alpha=np.radians(np.logspace(-1, 1, 10)), - points_beta=np.logspace(np.log10(1.01), 1, 8), - points_extension=np.radians(np.logspace(-1.5, np.log10(4.9), 8)), - points_psi=np.concatenate([[0.0], np.logspace(-4, np.log10(np.pi), 200)]), - n_quad=16, - ) - - -# Fixture for evaluate() tests: small angular_cutoff so far events are screened. -@pytest.fixture(scope="module") -def ext_eval(): - return ExtendedSourceKingPDF( - angular_cutoff=np.radians(10.0), - points_alpha=np.radians(np.logspace(-1, 1, 10)), - points_beta=np.logspace(np.log10(1.01), 1, 8), - points_extension=np.radians(np.logspace(-1.5, np.log10(4.9), 8)), - points_psi=np.concatenate([[0.0], np.logspace(-4, np.log10(np.pi), 200)]), - n_quad=16, - ) - - -PARAM_CASES = [ - pytest.param(np.radians(0.5), 2.0, np.radians(0.5), id="narrow-moderate-small-ext"), - pytest.param(np.radians(1.0), 2.5, np.radians(1.0), id="moderate-moderate-med-ext"), - pytest.param(np.radians(2.0), 4.0, np.radians(1.5), id="wide-heavy-med-ext"), -] - - -# --------------------------------------------------------------------------- -# Initialization -# --------------------------------------------------------------------------- - - -class TestExtendedSourceKingPDFInit: - def test_not_instance_of_king_pdf(self, ext_pdf): - assert not isinstance(ext_pdf, KingPDF) - - def test_default_angular_cutoff(self, ext_pdf): - assert ext_pdf.angular_cutoff == pytest.approx(np.pi) - - def test_custom_angular_cutoff(self): - cutoff = np.radians(5.0) - ext = ExtendedSourceKingPDF( - angular_cutoff=cutoff, - points_alpha=np.radians([0.5, 1.0]), - points_beta=np.array([1.5, 3.0]), - points_extension=np.radians([0.5, 1.0]), - n_quad=4, - ) - assert ext.angular_cutoff == pytest.approx(cutoff) - - def test_default_maximum_sigma(self, ext_pdf): - assert ext_pdf.maximum_sigma == pytest.approx(3.0) - - def test_custom_maximum_sigma(self): - ext = ExtendedSourceKingPDF( - maximum_sigma=5.0, - points_alpha=np.radians([0.5, 1.0]), - points_beta=np.array([1.5, 3.0]), - points_extension=np.radians([0.5, 1.0]), - n_quad=4, - ) - assert ext.maximum_sigma == pytest.approx(5.0) - - def test_table_shape(self, ext_pdf): - expected = ( - len(ext_pdf._log10_points_alpha), - len(ext_pdf._log10_points_beta), - len(ext_pdf._log10_points_extension), - len(ext_pdf._points_psi), - ) - assert ext_pdf._table.shape == expected - - def test_table_finite(self, ext_pdf): - assert np.all(np.isfinite(ext_pdf._table)) - - def test_table_nonneg(self, ext_pdf): - assert np.all(ext_pdf._table >= 0.0) - - -# --------------------------------------------------------------------------- -# pdf() -# --------------------------------------------------------------------------- - - -class TestExtendedSourceKingPDFPdf: - @pytest.mark.parametrize("alpha, beta, extension", PARAM_CASES) - def test_pdf_valid(self, ext_pdf, alpha, beta, extension): - psi = np.linspace(0, np.radians(5), 50) - vals = ext_pdf.pdf(psi, np.full_like(psi, alpha), np.full_like(psi, beta), extension) - assert np.all(vals >= 0) - assert np.all(np.isfinite(vals)) - - def test_zero_at_psi_zero(self, ext_pdf): - val = ext_pdf.pdf(0.0, np.radians(1.0), 2.5, np.radians(1.0)) - assert val == 0.0 - - def test_zero_beyond_psi_max(self): - """Angles above _points_psi[-1] (but still ≤ π) must return 0. - - Use angular_cutoff=10° without a custom points_psi so the default - upper bound is max_sigma*max_ext + angular_cutoff ≈ 16° < π, leaving - room for test points on the sphere that are out-of-table. - """ - small = ExtendedSourceKingPDF( - angular_cutoff=np.radians(10.0), - points_alpha=np.radians([0.5, 1.0, 2.0]), - points_beta=np.array([1.5, 2.5, 5.0]), - points_extension=np.radians([0.5, 1.0, 2.0]), - n_quad=4, - ) - psi_max = small._points_psi[-1] - assert psi_max < np.pi, "fixture must end before π for this test to be meaningful" - beyond = np.array([psi_max + np.radians(5.0), psi_max + np.radians(20.0)]) - beyond = beyond[beyond <= np.pi] - alpha = np.full(len(beyond), np.radians(1.0)) - beta = np.full(len(beyond), 2.5) - ext = np.full(len(beyond), np.radians(1.0)) - assert np.all(small.pdf(beyond, alpha, beta, ext) == 0.0) - - def test_output_shape_array(self, ext_pdf): - psi = np.linspace(0.01, np.radians(5), 20) - vals = ext_pdf.pdf(psi, np.radians(1.0), 2.5, np.radians(1.0)) - assert vals.shape == psi.shape - - def test_scalar_input_finite(self, ext_pdf): - val = ext_pdf.pdf(np.radians(1.0), np.radians(1.0), 2.5, np.radians(1.0)) - assert np.isfinite(val) - - def test_oob_alpha_raises(self, ext_pdf): - with pytest.raises(ValueError): - ext_pdf.pdf(np.radians(1.0), np.radians(0.001), 2.5, np.radians(1.0)) - - def test_oob_extension_raises(self, ext_pdf): - with pytest.raises(ValueError): - ext_pdf.pdf(np.radians(1.0), np.radians(1.0), 2.5, np.radians(10.0)) - - @pytest.mark.parametrize("alpha, beta, extension", PARAM_CASES) - def test_normalization(self, ext_pdf, alpha, beta, extension): - """∫ pdf(ψ) 2π ψ dψ ≈ 1 (flat-sky).""" - psi = np.linspace(1e-4, ext_pdf._points_psi[-1], 30_000) - dpsi = psi[1] - psi[0] - vals = ext_pdf.pdf( - psi, - np.full_like(psi, alpha), - np.full_like(psi, beta), - np.full_like(psi, extension), - ) - integral = np.sum(vals * 2.0 * np.pi * psi) * dpsi - assert_allclose(integral, 1.0, rtol=0.02) - - @pytest.mark.parametrize("alpha, beta, extension", PARAM_CASES) - def test_small_extension_approaches_king(self, ext_pdf, alpha, beta, extension): - """Convolved PDF with the smallest grid extension should be close to flat-sky King.""" - tiny_ext = ext_pdf._points_extension[0] - psi = np.radians([0.5, 1.0, 2.0]) - psi = psi[psi < alpha * 3] # stay in the PSF core where flat-sky is accurate - if len(psi) == 0: - pytest.skip("no test angles within PSF core for this alpha") - - flat_norm = (beta - 1.0) / (2.0 * np.pi * beta * alpha**2) - flat_king = flat_norm * (1.0 + psi**2 / (2.0 * beta * alpha**2)) ** (-beta) - conv = ext_pdf.pdf( - psi, - np.full_like(psi, alpha), - np.full_like(psi, beta), - np.full_like(psi, tiny_ext), - ) - assert_allclose(conv, flat_king, rtol=0.15) - - def test_larger_extension_broader(self, ext_pdf): - """Larger extension shifts probability outward, reducing the PDF near psi=0.""" - alpha = np.radians(1.0) - beta = 2.5 - psi_near = np.radians(0.1) - val_small = ext_pdf.pdf(psi_near, alpha, beta, ext_pdf._points_extension[0]) - val_large = ext_pdf.pdf(psi_near, alpha, beta, ext_pdf._points_extension[-1]) - assert val_small > val_large - - -# --------------------------------------------------------------------------- -# evaluate() -# --------------------------------------------------------------------------- - - -class TestExtendedSourceKingPDFEvaluate: - def test_returns_csr_array(self, ext_eval): - result = ext_eval.evaluate( - np.array([0.0]), - np.array([0.0]), - np.array([np.radians(1.0)]), - np.array([0.0]), - np.array([0.0]), - np.array([np.radians(1.0)]), - np.array([2.5]), - ) - assert isinstance(result, csr_array) - - def test_output_shape(self, ext_eval): - src_ras = np.radians([0.0, 45.0]) - src_decs = np.radians([0.0, 10.0]) - src_exts = np.radians([1.0, 1.0]) - ev_ras = np.radians(np.linspace(0, 5, 8)) - ev_decs = np.zeros(8) - alpha = np.full(8, np.radians(1.0)) - beta = np.full(8, 2.5) - result = ext_eval.evaluate(src_ras, src_decs, src_exts, ev_ras, ev_decs, alpha, beta) - assert result.shape == (8, 2) - - def test_nonneg(self, ext_eval): - rng = np.random.default_rng(0) - ev_ras = rng.uniform(0, 2 * np.pi, 30) - ev_decs = np.arcsin(rng.uniform(-1, 1, 30)) - alpha = np.full(30, np.radians(1.0)) - beta = np.full(30, 2.5) - result = ext_eval.evaluate( - np.array([0.0]), - np.array([0.0]), - np.array([np.radians(1.0)]), - ev_ras, - ev_decs, - alpha, - beta, - ) - assert np.all(result.toarray() >= 0) - - def test_near_source_positive(self, ext_eval): - """Events close to a source should get a positive PDF value.""" - result = ext_eval.evaluate( - np.array([0.0]), - np.array([0.0]), - np.array([np.radians(1.0)]), - np.array([np.radians(0.1)]), - np.array([0.0]), - np.array([np.radians(1.0)]), - np.array([2.5]), - ) - assert result.toarray()[0, 0] > 0 - - def test_zero_beyond_search_radius(self, ext_eval): - """Events beyond maximum_sigma * ext + angular_cutoff should be zero.""" - src_ext = np.radians(1.0) - radius = ext_eval.maximum_sigma * src_ext + ext_eval.angular_cutoff - # Place one event just inside and one well outside - psi_far = min(radius + np.radians(5.0), np.pi) - result = ext_eval.evaluate( - np.array([0.0]), - np.array([0.0]), - np.array([src_ext]), - np.array([np.radians(0.5), psi_far]), - np.array([0.0, 0.0]), - np.array([np.radians(1.0), np.radians(1.0)]), - np.array([2.5, 2.5]), - ).toarray() - assert result[0, 0] > 0 - assert result[1, 0] == 0.0 - - def test_mask_gives_same_result(self, ext_eval): - rng = np.random.default_rng(42) - src_ras = np.radians([0.0, 45.0]) - src_decs = np.radians([0.0, 10.0]) - src_exts = np.radians([1.0, 2.0]) - ev_ras = rng.uniform(0, 2 * np.pi, 30) - ev_decs = np.arcsin(rng.uniform(-1, 1, 30)) - alpha = np.full(30, np.radians(1.0)) - beta = np.full(30, 2.5) - first = ext_eval.evaluate(src_ras, src_decs, src_exts, ev_ras, ev_decs, alpha, beta) - second = ext_eval.evaluate( - src_ras, src_decs, src_exts, ev_ras, ev_decs, alpha, beta, mask=first - ) - assert_allclose(first.toarray(), second.toarray(), rtol=1e-12) - - def test_two_sources_prefer_nearest(self, ext_eval): - """An event near source 0 should get a higher PDF for source 0 than source 1.""" - src_ras = np.radians([0.0, 90.0]) - src_decs = np.radians([0.0, 0.0]) - src_exts = np.radians([1.0, 1.0]) - ev_ras = np.radians([1.0]) - ev_decs = np.radians([0.0]) - alpha = np.array([np.radians(1.0)]) - beta = np.array([2.5]) - result = ext_eval.evaluate( - src_ras, src_decs, src_exts, ev_ras, ev_decs, alpha, beta - ).toarray() - assert result[0, 0] > result[0, 1] diff --git a/tests/test_fitting.py b/tests/test_fitting.py index 23e7e67..13907f8 100644 --- a/tests/test_fitting.py +++ b/tests/test_fitting.py @@ -223,3 +223,85 @@ def test_beta_values_per_bin(self, result): betas = result.fit_beta[0, 0] for i, (_, beta_true, _) in enumerate(self._GROUP_PARAMS): assert_allclose(betas[i], beta_true, rtol=0.50, err_msg=f"bin {i}: beta mismatch") + + +# --------------------------------------------------------------------------- +# extension_grid +# --------------------------------------------------------------------------- + + +class TestKingPSFFitterExtensionGrid: + def test_default_is_point_source(self): + rng = np.random.default_rng(RNG_SEED) + events = _make_events(500, np.radians(1.0), 2.5, "aux", np.zeros(500), rng) + fitter = KingPSFFitter( + events, parametrization_bins={"aux": [-1.0, 1.0]}, minimum_counts=100, weight_field=None + ) + assert_allclose(fitter.extension_grid, [0.0]) + + def test_negative_extension_raises(self): + rng = np.random.default_rng(RNG_SEED) + events = _make_events(500, np.radians(1.0), 2.5, "aux", np.zeros(500), rng) + with pytest.raises(ValueError): + KingPSFFitter( + events, + parametrization_bins={"aux": [-1.0, 1.0]}, + minimum_counts=100, + weight_field=None, + extension_grid=[-0.1, 0.0], + ) + + def test_extension_grid_is_sorted(self): + rng = np.random.default_rng(RNG_SEED) + events = _make_events(500, np.radians(1.0), 2.5, "aux", np.zeros(500), rng) + fitter = KingPSFFitter( + events, + parametrization_bins={"aux": [-1.0, 1.0]}, + minimum_counts=100, + weight_field=None, + extension_grid=[np.radians(2.0), 0.0, np.radians(1.0)], + ) + assert_allclose(fitter.extension_grid, [0.0, np.radians(1.0), np.radians(2.0)]) + + @pytest.fixture(scope="class") + def multi_ext_result(self): + rng = np.random.default_rng(RNG_SEED) + alpha_true, beta_true = np.radians(1.0), 2.5 + n = 100_000 + events = _make_events(n, alpha_true, beta_true, "aux", np.zeros(n), rng) + extension_grid = np.radians([0.0, 1.0, 2.0]) + fitter = KingPSFFitter( + events, + parametrization_bins={"aux": [-1.0, 1.0]}, + dpsi_nbins=100, + minimum_counts=100, + weight_field=None, + extension_grid=extension_grid, + ) + result = fitter.fit_all_bins(verbose=False) + return result, alpha_true, beta_true, extension_grid + + def test_shape_matches_extension_grid(self, multi_ext_result): + result, _, _, extension_grid = multi_ext_result + assert result["alpha"].shape == (len(extension_grid), 1, 1) + assert result["beta"].shape == (len(extension_grid), 1, 1) + assert_allclose(result["extension_grid"], extension_grid) + + def test_fitted_values_finite_and_valid(self, multi_ext_result): + result, _, _, _ = multi_ext_result + assert np.all(np.isfinite(result["alpha"])) + assert np.all(np.isfinite(result["beta"])) + assert np.all(result["alpha"] > 0) + assert np.all(result["beta"] > 1) + + def test_zero_extension_recovers_point_source_fit(self, multi_ext_result): + """extension=0 should reproduce the un-smeared point-source fit.""" + result, alpha_true, beta_true, _ = multi_ext_result + assert_allclose(result["alpha"][0, 0, 0], alpha_true, rtol=0.1) + assert_allclose(result["beta"][0, 0, 0], beta_true, rtol=0.2) + + def test_alpha_increases_with_extension(self, multi_ext_result): + """A wider source extension should widen the fitted PSF.""" + result, _, _, _ = multi_ext_result + alphas = result["alpha"][:, 0, 0] + assert alphas[0] < alphas[1] < alphas[2] diff --git a/tests/test_utils.py b/tests/test_utils.py index 4f94e6f..bac0a1e 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,14 +1,13 @@ """ Unit tests for kingmaker.utils. -Covers angular_distance (known angles, symmetry, self-distance) and -meshgrid2d (shape, values, dtype preservation). +Covers angular_distance, meshgrid2d, and sample_with_extension. """ import numpy as np from numpy.testing import assert_allclose -from kingmaker.utils import angular_distance, meshgrid2d +from kingmaker.utils import angular_distance, meshgrid2d, sample_with_extension # --------------------------------------------------------------------------- @@ -113,3 +112,69 @@ def test_dtype_preserved(self): ga, gb = meshgrid2d(a, b) assert ga.dtype == np.float32 assert gb.dtype == np.float32 + + +# --------------------------------------------------------------------------- +# sample_with_extension +# --------------------------------------------------------------------------- + + +class TestSampleWithExtension: + def test_zero_extension_is_noop(self): + rng = np.random.default_rng(0) + ra, dec = sample_with_extension(1.0, 0.5, 0.0, rng) + assert_allclose(ra, 1.0) + assert_allclose(dec, 0.5) + + def test_output_shape_matches_input(self): + rng = np.random.default_rng(1) + true_ra = np.linspace(0, 2 * np.pi, 50) + true_dec = np.linspace(-1.0, 1.0, 50) + ra, dec = sample_with_extension(true_ra, true_dec, np.radians(2.0), rng) + assert ra.shape == (50,) + assert dec.shape == (50,) + + def test_finite_and_in_range(self): + rng = np.random.default_rng(2) + n = 5000 + true_ra = rng.uniform(0, 2 * np.pi, n) + true_dec = np.arcsin(rng.uniform(-1, 1, n)) + ra, dec = sample_with_extension(true_ra, true_dec, np.radians(3.0), rng) + assert np.all(np.isfinite(ra)) + assert np.all(np.isfinite(dec)) + assert np.all(dec >= -np.pi / 2 - 1e-9) + assert np.all(dec <= np.pi / 2 + 1e-9) + assert np.all(ra >= 0.0) + assert np.all(ra < 2 * np.pi) + + def test_rayleigh_scale_recovered(self): + """Offset magnitude follows Rayleigh(extension): mean ~ scale*sqrt(pi/2).""" + rng = np.random.default_rng(3) + n = 200_000 + extension = np.radians(2.0) + true_ra = np.full(n, 1.0) + true_dec = np.full(n, 0.3) + ra, dec = sample_with_extension(true_ra, true_dec, extension, rng) + offset = angular_distance(true_ra, true_dec, ra, dec) + expected_mean = extension * np.sqrt(np.pi / 2) + assert_allclose(offset.mean(), expected_mean, rtol=0.02) + + def test_no_pole_crash(self): + rng = np.random.default_rng(4) + for dec0 in [np.pi / 2, -np.pi / 2, np.radians(89.9), np.radians(-89.9)]: + ra, dec = sample_with_extension(0.0, dec0, np.radians(3.0), rng) + assert np.isfinite(ra) + assert np.isfinite(dec) + assert -np.pi / 2 - 1e-9 <= dec <= np.pi / 2 + 1e-9 + + def test_seeded_rng_is_reproducible(self): + true_ra, true_dec, extension = 0.5, 0.2, np.radians(1.5) + ra1, dec1 = sample_with_extension(true_ra, true_dec, extension, np.random.default_rng(7)) + ra2, dec2 = sample_with_extension(true_ra, true_dec, extension, np.random.default_rng(7)) + assert_allclose(ra1, ra2) + assert_allclose(dec1, dec2) + + def test_default_rng_when_none(self): + ra, dec = sample_with_extension(0.0, 0.0, np.radians(1.0)) + assert np.isfinite(ra) + assert np.isfinite(dec) diff --git a/tests/test_wrapper.py b/tests/test_wrapper.py index 38db66d..15b01c9 100644 --- a/tests/test_wrapper.py +++ b/tests/test_wrapper.py @@ -59,6 +59,32 @@ def _make_likelihood(tmp_path, angular_cutoff=np.pi): ) +# Three extensions; alpha grows with extension index so per-source +# differentiation is directly observable. +MULTI_EXTENSION_GRID = np.radians([0.0, 1.0, 3.0]) +MULTI_EXT_ALPHA_VALUES = np.stack([ALPHA_VALUES[0] * scale for scale in (1.0, 2.0, 4.0)]) +MULTI_EXT_BETA_VALUES = np.stack([BETA_VALUES[0] for _ in range(3)]) + + +def _make_multi_ext_likelihood(tmp_path, angular_cutoff=np.pi): + cache_path = tmp_path / "king_cache_multi_ext.npz" + np.savez( + cache_path, + parametrization_bins=np.array({"aux": BIN_EDGES}, dtype=object), + alpha=MULTI_EXT_ALPHA_VALUES, + beta=MULTI_EXT_BETA_VALUES, + extension_grid=MULTI_EXTENSION_GRID, + ) + return KingSpatialLikelihood( + signal_events=np.empty(0), + parametrization_bins={"aux": 3}, + spectral_indices=SPECTRAL_INDICES, + cache_parameters=True, + cache_name=str(cache_path), + angular_cutoff=angular_cutoff, + ) + + def _make_events(n_per_bin, rng, offset_scale=np.radians(2.0)): """n_per_bin events at each of the 3 known bin centers (0.5, 1.5, 2.5), at small random offsets from a source at (ra=0, dec=0).""" @@ -326,3 +352,70 @@ def test_shape_matches_event_count(self, likelihood): result = likelihood.evaluate_pdf(events_10, gamma=2.0) assert result.shape == (len(events_10), 1) + + +class TestNearestExtensionIndex: + def test_snaps_to_nearest(self, tmp_path): + likelihood = _make_multi_ext_likelihood(tmp_path) + idx = likelihood._nearest_extension_index(np.radians([0.4, 1.6, 2.9])) + np.testing.assert_array_equal(idx, [0, 1, 2]) + + +class TestSourceExtensions: + def test_default_is_zero(self, tmp_path): + likelihood = _make_multi_ext_likelihood(tmp_path) + rng = np.random.default_rng(10) + events = _make_events(5, rng) + likelihood.set_events(events, source_ras=np.array([0.0]), source_decs=np.array([0.0])) + np.testing.assert_array_equal(likelihood.source_extensions, [0.0]) + + def test_narrower_extension_source_has_higher_nearby_pdf(self, tmp_path): + likelihood = _make_multi_ext_likelihood(tmp_path) + src_ras = np.array([0.0, 0.0]) + src_decs = np.array([0.0, 0.0]) + src_exts = np.array([MULTI_EXTENSION_GRID[0], MULTI_EXTENSION_GRID[2]]) + + dtype = [("ra", float), ("dec", float), ("aux", float)] + events = np.zeros(1, dtype=dtype) + events["ra"], events["dec"], events["aux"] = np.radians(0.1), 0.0, 0.5 + + likelihood.set_events( + events, source_ras=src_ras, source_decs=src_decs, source_extensions=src_exts + ) + result = likelihood.evaluate_pdf(events, gamma=2.0).toarray() + assert result[0, 0] > result[0, 1] + + def test_missing_extension_grid_key_raises(self, tmp_path): + cache_path = tmp_path / "stale_no_key.npz" + np.savez( + cache_path, + parametrization_bins=np.array({"aux": BIN_EDGES}, dtype=object), + alpha=ALPHA_VALUES[0], + beta=BETA_VALUES[0], + ) + with pytest.raises(ValueError): + KingSpatialLikelihood( + signal_events=np.empty(0), + parametrization_bins={"aux": 3}, + spectral_indices=SPECTRAL_INDICES, + cache_parameters=True, + cache_name=str(cache_path), + ) + + def test_wrong_ndim_raises(self, tmp_path): + cache_path = tmp_path / "stale_wrong_ndim.npz" + np.savez( + cache_path, + parametrization_bins=np.array({"aux": BIN_EDGES}, dtype=object), + alpha=ALPHA_VALUES[0], + beta=BETA_VALUES[0], + extension_grid=np.array([0.0]), + ) + with pytest.raises(ValueError): + KingSpatialLikelihood( + signal_events=np.empty(0), + parametrization_bins={"aux": 3}, + spectral_indices=SPECTRAL_INDICES, + cache_parameters=True, + cache_name=str(cache_path), + ) From 652ac0b6f327ce5e2f3b5677fb71bdf9de9b328b Mon Sep 17 00:00:00 2001 From: Michael Larson Date: Thu, 30 Jul 2026 15:29:01 -0400 Subject: [PATCH 4/4] Fix up docs testing. --- docs/examples.rst | 13 +++++-------- kingmaker/fitting.py | 13 +++++++++---- 2 files changed, 14 insertions(+), 12 deletions(-) diff --git a/docs/examples.rst b/docs/examples.rst index f9cc3b4..25b1bd0 100644 --- a/docs/examples.rst +++ b/docs/examples.rst @@ -141,7 +141,7 @@ bin. ) results = fitter.fit_all_bins(verbose=True) - alpha_fit = results["alpha"] # shape (n_gamma, n_logE, n_dec) + alpha_fit = results["alpha"] # shape (n_extension, n_gamma, n_logE, n_dec) beta_fit = results["beta"] # Continuous evaluation between bin centers: @@ -192,8 +192,9 @@ above: from kingmaker.wrapper import KingSpatialLikelihood import numpy as np - # Source catalog for the signal-subtraction (marginalized) path. - catalog_decs = np.radians(np.linspace(-60, 60, 13)) + # Stand-in "data" events and a point-source position for one trial. + data_events = signal_events[:1000] + source_ra, source_dec = 0.5, 0.2 wrapper = KingSpatialLikelihood( signal_events=signal_events, @@ -202,14 +203,10 @@ above: cache_parameters=False, # Enable the RA-marginalized path for signal-subtraction likelihoods. enable_marginalization=True, - marginalization_source_decs=catalog_decs, + marginalization_source_decs=np.array([source_dec]), marginalization_angular_cutoff=np.radians(10.0), ) - # Stand-in "data" events and a point-source position for one trial. - data_events = signal_events[:1000] - source_ra, source_dec = 0.5, 0.2 - # Per trial: cache per-event parameters once, then evaluate as needed. # set_events precomputes both the standard and marginalized PDF matrices. wrapper.set_events( diff --git a/kingmaker/fitting.py b/kingmaker/fitting.py index 6448d8c..101fb01 100644 --- a/kingmaker/fitting.py +++ b/kingmaker/fitting.py @@ -559,7 +559,7 @@ def _fit_single_bin( return False - def get_interpolator(self, gamma_index: int = 0) -> Tuple[Any, Any]: + def get_interpolator(self, gamma_index: int = 0, extension_index: int = 0) -> Tuple[Any, Any]: """ Get an interpolator for fitted parameters at a given spectral index. @@ -567,6 +567,8 @@ def get_interpolator(self, gamma_index: int = 0) -> Tuple[Any, Any]: ---------- gamma_index : int, optional Index of the spectral index to use. Default is 0. + extension_index : int, optional + Index into extension_grid to use. Default is 0. Returns ------- @@ -586,7 +588,7 @@ def get_interpolator(self, gamma_index: int = 0) -> Tuple[Any, Any]: # Create interpolators alpha_interp = RegularGridInterpolator( tuple(bin_centers), - self.fit_alpha[gamma_index], + self.fit_alpha[extension_index, gamma_index], method="linear", bounds_error=False, fill_value=self.fit_alpha[gamma_index].mean(), @@ -594,7 +596,7 @@ def get_interpolator(self, gamma_index: int = 0) -> Tuple[Any, Any]: beta_interp = RegularGridInterpolator( tuple(bin_centers), - self.fit_beta[gamma_index], + self.fit_beta[extension_index, gamma_index], method="linear", bounds_error=False, fill_value=self.fit_beta[gamma_index].mean(), @@ -606,6 +608,7 @@ def plot_fit( self, bin_indices: Union[Tuple[int, ...], Dict[str, int]], gamma_index: int = 0, + extension_index: int = 0, ax: Optional[Any] = None, ) -> Any: """ @@ -618,6 +621,8 @@ def plot_fit( mapping bin names to indices. gamma_index : int, optional Index of spectral index. Default is 0. + extension_index : int, optional + Index into extension_grid to use. Default is 0. ax : matplotlib.axes.Axes, optional Axes to plot on. If None, creates new figure. @@ -640,7 +645,7 @@ def plot_fit( if isinstance(bin_indices, dict): bin_indices = tuple(bin_indices[key] for key in self.bin_names) - param_idx = tuple([gamma_index] + list(bin_indices)) + param_idx = (extension_index, gamma_index) + tuple(bin_indices) # Get histogram data hist = self.histograms[param_idx]