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 93.17073% with 14 lines in your changes missing coverage. Please review.
✅ Project coverage is 79.37%. Comparing base (3a78cb4) to head (ac94135).

Files with missing lines	Patch %	Lines
kingmaker/fitting.py	91.01%	8 Missing ⚠️
kingmaker/pdf.py	75.00%	3 Missing ⚠️
kingmaker/wrapper.py	97.61%	2 Missing ⚠️
kingmaker/utils.py	94.44%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##             main       #7       +/-   ##
===========================================
+ Coverage   53.46%   79.37%   +25.90%     
===========================================
  Files           6        6               
  Lines         937      921       -16     
===========================================
+ Hits          501      731      +230     
+ Misses        436      190      -246     

Flag	Coverage Δ
test_basic	14.54% <19.02%> (-0.18%)	⬇️
test_fitting	38.21% <60.00%> (+6.84%)	⬆️
test_king_pdf	27.90% <20.00%> (-0.17%)	⬇️
test_template_smeared_king_pdf	26.05% <20.00%> (+0.23%)	⬆️
test_utils	16.06% <25.85%> (+1.34%)	⬆️
test_wrapper	38.65% <48.29%> (?)

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.

[Copilot]

Copilot review overview

🟡 Changes recommended

Several moderate validation, edge-case, and backward-compatibility issues remain unresolved.

Review effort: Lite
Findings: 2

Open (2)

• Reject nonfinite, empty, or non-1D grid values · New
• Reject NaN source extensions · New

Resolved since last review (1)

• Preserve alpha and beta ordering when sorting extension grid

[Copilot]

Copilot review overview

🟡 Changes recommended

kingmaker/wrapper.py has a critical import-blocking indentation error and a moderate malformed-cache validation issue.

Review effort: Lite
Findings: 1

Open (1)

• Missing indented suite causes wrapper.py import failure · New

Resolved since last review (2)

• Reject NaN source extensions
• Reject nonfinite, empty, or non-1D grid values

[Copilot]

Copilot review overview

🔵 Needs a closer look

Unresolved validation, compatibility, empty-input, regression-test, and notebook-output issues remain.

Review effort: Lite
Findings: None

Resolved since last review (1)

• Missing indented suite causes wrapper.py import failure

Previously missed (2)

In code that hasn't changed since last review

Reject empty cached extension grids during validation

kingmaker/​wrapper.py:158

This validation accepts an empty cached extension_grid because it checks dimensionality and values but not size. The constructor then succeeds with zero-length alpha/beta axes, and the first set_events call fails later at self.extension_grid[[0, -1]] with an indexing error instead of reporting the invalid cache here. Reject empty grids as the fitter already does.

Add regression coverage for grouped multi-extension marginalization

kingmaker/​wrapper.py:431

This new grouped path is not covered by the added marginalization tests, which only exercise the single-extension/single-source case. Please add a regression test with sources mapped to different extension indices and verify both the source columns and the shared CSR structure across spectral indices; this remapping is distinct from the existing point-source path.

[Copilot]

Copilot review overview

🟡 Changes recommended

Resolve the public API compatibility and extension-grid validation issues before approval.

Review effort: Lite
Findings: 2

Open (2)

• Preserve positional API when adding extension_index · New
• Reject empty extension grids before indexing · New

Previously missed (2)

In code that hasn't changed since last review

Regenerate output for added point-source extension axis

examples/​fitting_demo.ipynb:573

The notebook's stored output still says alpha/beta have shape (3, 10, 5, 10), but the revised fitter adds the default point-source extension axis, so this run now produces (1, 3, 10, 5, 10). Please regenerate or update the output to keep the example consistent with the changed indexing.

Update output for new point-source extension dimension

examples/​likelihood_demo.ipynb:279

The stored output for this notebook still reports the pre-rewrite shape (2, 3), but the updated code now produces (1, 2, 3) for the default point-source extension grid. Please regenerate or update the cell output so the rendered example does not contradict the new indexing and shape documentation.

[Copilot]

Copilot review overview

🔵 Needs a closer look

Unresolved moderate issues affect API compatibility, validation, cache safety, and marginalized evaluation.

Review effort: Lite
Findings: 2

Open (2)

• Reject empty extension grids before indexing
• Preserve positional API when adding extension_index

[Copilot]

Copilot review overview

🔵 Needs a closer look

Moderate issues remain in pole sampling, empty-catalog marginalization, cache validation, and default extension handling; documentation and example updates are also needed.

Review effort: Lite
Findings: None

Resolved since last review (2)

• Reject empty extension grids before indexing
• Preserve positional API when adding extension_index

Previously missed (1)

In code that hasn't changed since last review

Add marginalized mixed-extension regression coverage

kingmaker/​wrapper.py:420

The new group_source_idx/group_masks path is only exercised by the existing point-source marginalization tests; no test passes multiple source extensions with enable_marginalization. Because this code changes both source-column remapping and mask reuse, add a mixed-extension regression test that evaluates the marginalized matrix.