Fix up docs testing.

Author: mjlarson

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(

kingmaker/fitting.py

@@ -559,14 +559,16 @@ 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.
 
         Parameters
         ----------
         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,15 +588,15 @@ 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(),
         )
 
         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]