Mlarson/extended sources rewrite

.github/workflows/ci.yml

@@ -54,6 +54,7 @@ jobs:
           - test_template_smeared_king_pdf
           - test_utils
           - test_fitting
+          - test_wrapper
 
     steps:
       - name: Checkout code

.pre-commit-config.yaml

@@ -1,6 +1,6 @@
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.11.8
+    rev: v0.16.10
     hooks:
       - id: ruff
       - id: ruff-format

docs/conf.py

@@ -1,5 +1,5 @@
-import sys
 import os
+import sys
 
 sys.path.insert(0, os.path.abspath(".."))
 

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(

examples/fitting_demo.ipynb

@@ -569,8 +569,8 @@
     "bin_idx = (6, 7, 2)  # Mid energy, mid dec, mid sigma\n",
     "gamma_idx = 0\n",
     "\n",
-    "alpha = results['alpha'][gamma_idx][bin_idx]\n",
-    "beta = results['beta'][gamma_idx][bin_idx]\n",
+    "alpha = results['alpha'][0, gamma_idx][bin_idx]\n",
+    "beta = results['beta'][0, gamma_idx][bin_idx]\n",
     "\n",
     "print(f\"Fitted parameters for bin {bin_idx}:\")\n",
     "print(f\"  alpha = {np.degrees(alpha):.4f} degrees\")\n",

examples/likelihood_demo.ipynb

@@ -272,10 +272,11 @@
     "    spectral_indices=[2.0, 3.0],\n",
     "    angular_cutoff=np.pi,\n",
     "    cache_parameters=False,  # disable caching for this demo\n",
+    "    true_energy_name=\"true_energy\",\n",
     ")\n",
     "\n",
     "print(\"KingSpatialLikelihood initialised successfully.\")\n",
-    "print(f\"  alpha_values shape : {wrapper.alpha_values.shape}  (n_spectral_indices × n_energy_bins)\")\n",
+    "print(f\"  alpha_values shape : {wrapper.alpha_values.shape}  (n_extensions × n_spectral_indices × n_energy_bins)\")\n",
     "print(f\"  beta_values  shape : {wrapper.beta_values.shape}\")\n",
     "print(f\"  spectral_indices   : {wrapper.spectral_indices}\")"
    ]
@@ -339,7 +340,7 @@
     "for g_idx, (gamma, color) in enumerate(zip(wrapper.spectral_indices, colors)):\n",
     "    ax1.plot(\n",
     "        bin_centers,\n",
-    "        np.degrees(wrapper.alpha_values[g_idx]),\n",
+    "        np.degrees(wrapper.alpha_values[0, g_idx]),\n",
     "        \"o-\",\n",
     "        markersize=9,\n",
     "        linewidth=2,\n",
@@ -348,7 +349,7 @@
     "    )\n",
     "    ax2.plot(\n",
     "        bin_centers,\n",
-    "        wrapper.beta_values[g_idx],\n",
+    "        wrapper.beta_values[0, g_idx],\n",
     "        \"o-\",\n",
     "        markersize=9,\n",
     "        linewidth=2,\n",
@@ -380,11 +381,11 @@
     "\n",
     "print(\"Fitted α values (degrees):\")\n",
     "for g_idx, gamma in enumerate(wrapper.spectral_indices):\n",
-    "    fitted = np.degrees(wrapper.alpha_values[g_idx])\n",
+    "    fitted = np.degrees(wrapper.alpha_values[0, g_idx])\n",
     "    print(f\"  γ={gamma:.1f}: {fitted}  (true: {true_alphas_deg})\")\n",
     "print(\"\\nFitted β values:\")\n",
     "for g_idx, gamma in enumerate(wrapper.spectral_indices):\n",
-    "    fitted = wrapper.beta_values[g_idx]\n",
+    "    fitted = wrapper.beta_values[0, g_idx]\n",
     "    print(f\"  γ={gamma:.1f}: {fitted}  (true: {true_betas})\")"
    ]
   },
@@ -533,8 +534,8 @@
     "\n",
     "# Fitted parameters at the test energy (interpolated between bin centers)\n",
     "def fitted_params_at(loge, g_idx):\n",
-    "    alpha = np.interp(loge, bin_centers, wrapper.alpha_values[g_idx])\n",
-    "    beta = np.interp(loge, bin_centers, wrapper.beta_values[g_idx])\n",
+    "    alpha = np.interp(loge, bin_centers, wrapper.alpha_values[0, g_idx])\n",
+    "    beta = np.interp(loge, bin_centers, wrapper.beta_values[0, g_idx])\n",
     "    return alpha, beta\n",
     "\n",
     "\n",

kingmaker/__init__.py

@@ -1,12 +1,12 @@
 # Import main classes for convenience
-from .pdf import KingPDF, MarginalizedKingPDF, TemplateSmearedKingPDF
 from .fitting import KingPSFFitter
+from .pdf import KingPDF, MarginalizedKingPDF, TemplateSmearedKingPDF
 from .wrapper import KingSpatialLikelihood
 
 __all__ = [
     "KingPDF",
-    "MarginalizedKingPDF",
-    "TemplateSmearedKingPDF",
     "KingPSFFitter",
     "KingSpatialLikelihood",
+    "MarginalizedKingPDF",
+    "TemplateSmearedKingPDF",
 ]

kingmaker/distribution.py

@@ -1,7 +1,6 @@
-from typing import Union
 import numpy as np
 import numpy.typing as npt
-from numba import njit, vectorize, float32, float64
+from numba import float32, float64, njit, vectorize
 
 _log10pi: float = np.log10(np.pi)
 
@@ -12,10 +11,10 @@
     cache=True,
 )
 def _unnormalized_pdf(
-    x: Union[float, npt.NDArray[np.floating]],
-    alpha: Union[float, npt.NDArray[np.floating]],
-    beta: Union[float, npt.NDArray[np.floating]],
-) -> Union[float, npt.NDArray[np.floating]]:
+    x: float | npt.NDArray[np.floating],
+    alpha: float | npt.NDArray[np.floating],
+    beta: float | npt.NDArray[np.floating],
+) -> float | npt.NDArray[np.floating]:
     """
     Evaluate the unnormalized spherical King function (without solid angle Jacobian):
         f(x) = [1 + (1 - cos x) / (alpha² * beta)]^(-beta)
@@ -43,8 +42,8 @@ def _unnormalized_pdf(
     cache=True,
 )
 def _unnormalized_cdf(
-    x: Union[float, npt.NDArray[np.floating]], alpha: float, beta: float
-) -> Union[float, npt.NDArray[np.floating]]:
+    x: float | npt.NDArray[np.floating], alpha: float, beta: float
+) -> float | npt.NDArray[np.floating]:
     """
     Evaluate the CDF of the radial King function (without solid angle Jacobian).
 
@@ -78,10 +77,10 @@ def _unnormalized_cdf(
     cache=True,
 )
 def _norm(
-    alpha: Union[float, npt.NDArray[np.floating]],
-    beta: Union[float, npt.NDArray[np.floating]],
+    alpha: float | npt.NDArray[np.floating],
+    beta: float | npt.NDArray[np.floating],
     maximum: float,
-) -> Union[float, npt.NDArray[np.floating]]:
+) -> float | npt.NDArray[np.floating]:
     """
     Compute the normalization constant for the King PDF over the sphere.