Mlarson/extended sources rewrite

[mjlarson]

Bit of a rewrite to handle extended sources in the normal KingSpatialLikelihood class

[codecov-commenter]

Codecov Report

❌ Patch coverage is 60.62992% with 50 lines in your changes missing coverage. Please review.
✅ Project coverage is 58.45%. Comparing base (3a78cb4) to head (061aee2).

Files with missing lines	Patch %	Lines
kingmaker/wrapper.py	8.69%	42 Missing ⚠️
kingmaker/fitting.py	87.69%	8 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main       #7      +/-   ##
==========================================
+ Coverage   53.46%   58.45%   +4.98%     
==========================================
  Files           6        6              
  Lines         937      905      -32     
==========================================
+ Hits          501      529      +28     
+ Misses        436      376      -60     

Flag	Coverage Δ
test_basic	14.80% <7.87%> (+0.07%)	⬆️
test_fitting	35.46% <59.84%> (+4.09%)	⬆️
test_king_pdf	28.61% <7.87%> (+0.55%)	⬆️
test_template_smeared_king_pdf	26.29% <7.87%> (+0.47%)	⬆️
test_utils	16.24% <18.11%> (+1.51%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

[Copilot]

🟡 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 KingPSFFitter outputs to shape (n_extension, n_gamma, *bins) by fitting across an extension_grid, using Rayleigh smearing of true positions.
• Update KingSpatialLikelihood to load/store extension_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; remove ExtendedSourceKingPDF implementation 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.

[Copilot]

🟡 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

[Copilot]

Copilot review overview

🟡 Changes recommended

Unresolved correctness, validation, compatibility, and example-integrity issues remain.

Review effort: Lite
Findings: 1

Open (1)

• Preserve alpha and beta ordering when sorting extension grid · New

Resolved since last review (4)

• _sources_match ignores extension changes when the caller omits source_extensions (None). If a…
• fill_value indexes fit_alpha with gamma_index only, but fit_alpha is now shaped (n_extension,…
• When loading cached fit results, extension_grid is not validated against alpha/beta. If…
• The new per-source source_extensions behavior interacts with the marginalization path (which…

Previously missed (5)

In code that hasn't changed since last review

Validate extension grid shape, values, and non-emptiness

kingmaker/​fitting.py:120

This only rejects negative values; NaN, infinity, an empty grid, or a multidimensional array can pass and then produce invalid or empty fitted parameter axes. Validate that extension_grid is non-empty, one-dimensional, finite, and non-negative before fitting.

Use selected extension when indexing beta fallback values

kingmaker/​fitting.py:602

When extension_index is nonzero, the beta interpolator's out-of-bounds fill still uses self.fit_beta[gamma_index], indexing the extension axis instead of the selected extension. This can return the wrong fallback mean (or raise an IndexError); use the same (extension_index, gamma_index) slice as the interpolator values.

Validate source extensions before nearest-bin lookup

kingmaker/​wrapper.py:330

source_extensions is passed to searchsorted without validating it as a radius. Negative or NaN values are silently snapped to an unrelated extension bin, while a scalar produces a TypeError at len; reject non-1-D, non-finite, or negative inputs before calling _nearest_extension_index.

Preserve or update marginalized declination catalog semantics

kingmaker/​wrapper.py:385

This new check changes the existing marginalized API from an independent source-declination catalog to a 1:1 mapping with the trial sources. The repository's README still constructs 13 marginalization_source_decs values and then calls set_events with one source, which now raises here; either preserve the catalog semantics or update all supported callers and the output contract as an intentional breaking change.

Document rng under the correct API

kingmaker/​fitting.py:64

This rng entry is documented under KingPSFFitter's constructor parameters, but __init__ has no such argument; the generator is accepted by fit_all_bins instead. Users following this section will pass rng to construction and get a TypeError; remove this entry here or move it to the constructor's actual API description.