Mlarson/extended sources rewrite - #7
Open
mjlarson wants to merge 4 commits into
Open
Conversation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #7 +/- ##
==========================================
+ Coverage 53.46% 58.77% +5.30%
==========================================
Files 6 6
Lines 937 900 -37
==========================================
+ Hits 501 529 +28
+ Misses 436 371 -65
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. |
There was a problem hiding this comment.
🟡 Changes recommended
There are correctness and robustness issues in the wrapper’s source-caching/validation logic (notably extension-aware cache invalidation and input validation) plus a breaking API removal that leaves in-repo examples outdated.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
This PR refactors Kingmaker’s spatial likelihood pipeline to support extended sources by incorporating an explicit extension_grid axis into fitted PSF parameter grids and enabling per-source extension selection (nearest-snapped) during PDF evaluation.
Changes:
- Extend
KingPSFFitteroutputs to shape(n_extension, n_gamma, *bins)by fitting across anextension_grid, using Rayleigh smearing of true positions. - Update
KingSpatialLikelihoodto load/storeextension_grid, select per-source extension indices, and evaluate both standard and RA-marginalized PDFs with extension-aware parameter lookup. - Add
sample_with_extension()utility + tests; expand wrapper/fitter tests for extension behavior; removeExtendedSourceKingPDFimplementation and its dedicated test module.
File summaries
| File | Description |
|---|---|
| tests/test_wrapper.py | Updates wrapper tests for extension-aware alpha/beta grids; adds new tests for nearest extension selection and per-source extensions. |
| tests/test_utils.py | Adds unit tests for new sample_with_extension() utility. |
| tests/test_fitting.py | Updates fitter tests for new (n_extension, ...) parameter shapes; adds extension grid behavioral tests. |
| tests/test_extended_source_king_pdf.py | Removes tests for ExtendedSourceKingPDF (class removed). |
| kingmaker/wrapper.py | Adds source_extensions support, loads/validates extension_grid, and makes PDF caching/evaluation extension-aware. |
| kingmaker/utils.py | Introduces sample_with_extension() for Rayleigh source-extension sampling. |
| kingmaker/pdf.py | Removes ExtendedSourceKingPDF implementation and related imports. |
| kingmaker/fitting.py | Extends fitting loop to include extension_grid and uses sample_with_extension() to smear truth positions during fitting. |
| docs/examples.rst | Updates documentation to reflect the new extension axis in fitted results and tweaks likelihood example. |
Review details
Suppressed comments (2)
docs/examples.rst:153
- The example now documents alpha_fit as having an extension axis, but the subsequent get_interpolator()/plot_fit calls omit extension_index. This can confuse readers who provide an extension_grid with multiple entries and wonder why only the point-source slice is used.
# Continuous evaluation between bin centers:
alpha_interp, beta_interp = fitter.get_interpolator(gamma_index=0)
point = np.array([[3.5, np.arcsin(0.0)]]) # [logE, dec]
alpha_value = alpha_interp(point)
# Inspect a single bin's fit against its histogram:
ax = fitter.plot_fit(bin_indices=(2, 2), gamma_index=0)
kingmaker/wrapper.py:309
- set_events accepts source_extensions but does not validate that the radii are finite and non-negative. Negative/NaN values will silently snap to the nearest extension_grid entry, masking upstream data issues and producing incorrect likelihood values.
self.source_extensions = (
np.zeros(len(source_ras))
if source_extensions is None
else np.asarray(source_extensions, dtype=np.float64)
)
- Files reviewed: 9/9 changed files
- Comments generated: 4
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+223
to
+226
| if source_extensions is not None and not np.array_equal( | ||
| self.source_extensions, source_extensions | ||
| ): | ||
| return False |
Comment on lines
+142
to
+155
| # 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." | ||
| ) |
Comment on lines
+262
to
+265
| source_extensions : ndarray, optional | ||
| Source extension radii in radians, nearest-snapped to the fitted | ||
| ``extension_grid``. Defaults to zero (point source) for every | ||
| source. |
mjlarson
force-pushed
the
mlarson/extended_sources_rewrite
branch
from
September 3, 2026 19:06
effbd81 to
652ac0b
Compare
There was a problem hiding this comment.
🟡 Changes recommended
There are confirmed correctness issues in the new extension-aware caching and interpolator defaults, plus an example notebook still references the removed ExtendedSourceKingPDF API.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (3)
kingmaker/pdf.py:7
- ExtendedSourceKingPDF has been removed from kingmaker.pdf, but the repository still contains examples/extended_source_demo.ipynb importing and documenting it. As-is, that notebook will break for users and in any documentation build/test that executes it.
from scipy.special import legendre_p_all, sph_harm_y_all
kingmaker/fitting.py:602
- fill_value indexes fit_beta with gamma_index only, but fit_beta is now shaped (n_extension, n_gamma, *bins). This can raise IndexError when n_extension == 1 and gamma_index > 0, and it also uses the wrong slice for non-default extension_index.
fill_value=self.fit_beta[gamma_index].mean(),
kingmaker/wrapper.py:226
- _sources_match only compares source_extensions when the caller passes a non-None array. If a previous set_events call used explicit non-zero extensions and a later call passes source_extensions=None (meaning “default to zeros”), this method can incorrectly treat the sources as unchanged and reuse stale cached matrices.
if source_extensions is not None and not np.array_equal(
self.source_extensions, source_extensions
):
return False
- Files reviewed: 9/9 changed files
- Comments generated: 1
- Review effort level: Lite
| self.fit_alpha[extension_index, gamma_index], | ||
| method="linear", | ||
| bounds_error=False, | ||
| fill_value=self.fit_alpha[gamma_index].mean(), |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Bit of a rewrite to handle extended sources in the normal KingSpatialLikelihood class