Use King function for Northern Tracks PSF (#444)

* Use astropy Tables for data everywhere

* Load npz file if present

* Improve cache coherence in NTSeason.simulate_background

Cuts time in Injector.create_dataset() from ~1500 ms to ~600 ms

* CircularGaussian: do not attempt to modify data in-place

* Blacken/isorten following #415

* Use Table for EffectiveAreaInjector

* Annotate factory functions

* Avoid modifying columns in place

* Use instance variables in NorthernTracksKDE

* Normalize KDE PDF to region of interest

* Speed up TableInjector

by a factor ~2 when drawing ~175 events from 11.5e6 MC events across 1000 sources

* Speed up StaticFloor

numpy broadcasting is much faster than a list comprehension

* Cache effective injection time for steady sources

There's no point in calculating this tens of thousands of times, since it's independent of the source for a steady time pdf

* Remove upper bound on n_s

* Squeeze the last bits of fluff out of StdMatrixKDEEnabledLLH.get_spatially_coincident_indices

Unexpected bottlenecks:
- Repeatedly slicing into astropy Tables (overhead from copying units)
- Temporary arrays in ra_dist and angular_distance

* Evaluate SoB for different gammas lazily

Keys appear as before, but values are only evaluated when accessed. Speeds up initialization when gamma is not being fitted.

* Use less accurate (but much faster) angular distance calculation in NorthernTracksKDE

* Annotate PDF factory functions

* HACK: allow StdMatrixKDEEnabledLLH to use CircularGaussian pdf

* Revert "Annotate PDF factory functions"

This reverts commit 9afa53e13556e783a73909fa2751ed98bad35f27.

* Normalize gaussian PDF

* WIP: King function PSF

* chore: annotate BaseInjector

* Allow Season to simulate background only in the source box

* Evaluate energy SoB lazily

* Add log of King PDF

* Add an option to disable energy S/B interpolation

The 2nd-order spline interpolation does not necessarily reflect the normalization of the underyling signal or background energy PDFs

* Use splines for King PDF parameters

* Speed up make_season_weight 10x

180 ms -> 11 ms for 22k sources

* Microoptimization: sparse matrix projection 25% faster

* Microoptimization: slice result of background pdf calculation instead of inputs

* mypy cleanup

* Guard annotation-only imports

* Cosmetic: raw string for latex

* Return excluded events from create_dataset in unblinded injectors

* Consider excluded events in flare likelihoods (but require ==0)

* Handle unbounded n_s in 1D llh scan

Author: jvansanten

flarestack/analyses/angular_error_floor/test_dynamic_pull_correction.py

@@ -5,7 +5,6 @@
 
 import matplotlib.pyplot as plt
 import numpy as np
-from numpy.lib.recfunctions import append_fields
 
 from flarestack.core.astro import angular_distance
 from flarestack.data.icecube.ps_tracks.ps_v002_p01 import IC86_1_dict
@@ -93,9 +92,7 @@ def weighted_quantile(values, quantiles, weight):
         cut_mc = mc[mask]
 
         percentile = np.ones_like(cut_mc["ra"]) * np.nan
-        cut_mc = append_fields(
-            cut_mc, "percentile", percentile, usemask=False, dtypes=[np.float]
-        )
+        cut_mc.add_column(percentile, name="percentile")
 
         weights = cut_mc["ow"] * cut_mc["trueE"] ** -gamma
         # weights = np.ones_like(cut_mc["ow"])

flarestack/analyses/ccsn/necker_2019/check_injections-compare_flarestack_injections.ipynb

@@ -2,6 +2,8 @@
 import logging
 import os
 import pickle as Pickle
+from collections.abc import Callable, Mapping
+from typing import Any, Collection, Iterator, KeysView, TypeVar
 
 import numexpr
 import numpy as np
@@ -36,7 +38,7 @@
 
 
 class BaseFloorClass(object):
-    subclasses: dict[str, object] = {}
+    subclasses: dict[str, type["BaseFloorClass"]] = {}
 
     def __init__(self, floor_dict):
         self.floor_dict = floor_dict
@@ -56,7 +58,7 @@ def decorator(subclass):
         return decorator
 
     @classmethod
-    def create(cls, floor_dict):
+    def create(cls, floor_dict) -> "BaseFloorClass":
         floor_name = floor_dict["floor_name"]
 
         if floor_name not in cls.subclasses:
@@ -129,7 +131,7 @@ def __init__(self, floor_dict):
         )
 
     def floor(self, data):
-        return np.array([self.min_error for _ in data])
+        return self.min_error
 
 
 class BaseQuantileFloor(BaseFloorClass):
@@ -197,6 +199,30 @@ def create_function(self, pickled_array):
         return lambda data, params: func(data["logE"])
 """
 
+_K = TypeVar("_K")
+_V = TypeVar("_V")
+
+
+class LazyDict(Mapping[_K, _V]):
+    def __init__(self, keys: Collection[_K], func: Callable[[_K], _V]):
+        self._keys = frozenset(keys)
+        self._func = func
+        self._store: dict[_K, _V] = dict()
+
+    def keys(self) -> KeysView[_K]:
+        return self._keys  # type: ignore[return-value]
+
+    def __iter__(self) -> Iterator[_K]:
+        return iter(self._keys)
+
+    def __len__(self):
+        return len(self._keys)
+
+    def __getitem__(self, key: _K) -> _V:
+        if key not in self._store:
+            self._store[key] = self._func(key)
+        return self._store[key]
+
 
 @BaseFloorClass.register_subclass("quantile_floor_1d_e")
 class QuantileFloor1D(BaseQuantileFloor, BaseDynamicFloorClass):
@@ -218,9 +244,9 @@ def create_function(self, pickled_array):
 
 
 class BaseAngularErrorModifier(object):
-    subclasses: dict[str, object] = {}
+    subclasses: dict[str, type["BaseAngularErrorModifier"]] = {}
 
-    def __init__(self, pull_dict):
+    def __init__(self, pull_dict: dict[str, Any]) -> None:
         self.season = pull_dict["season"]
         self.floor = BaseFloorClass.create(pull_dict)
         self.pull_dict = pull_dict
@@ -252,7 +278,7 @@ def create(
         floor_name="static_floor",
         aem_name="no_modifier",
         **kwargs,
-    ):
+    ) -> "BaseAngularErrorModifier":
         pull_dict = dict()
         pull_dict["season"] = season
         pull_dict["e_pdf_dict"] = e_pdf_dict
@@ -278,15 +304,16 @@ def pull_correct_dynamic(self, data, params):
 
     def create_spatial_cache(self, cut_data, SoB_pdf):
         if len(inspect.getfullargspec(SoB_pdf)[0]) == 2:
-            SoB = dict()
-            for gamma in get_gamma_support_points(precision=self.precision):
-                SoB[gamma] = np.log(SoB_pdf(cut_data, gamma))
+            SoB = LazyDict(
+                get_gamma_support_points(precision=self.precision),
+                lambda gamma: np.log(SoB_pdf(cut_data, gamma)),
+            )
         else:
             SoB = SoB_pdf(cut_data)
         return SoB
 
     def estimate_spatial(self, gamma, spatial_cache):
-        if isinstance(spatial_cache, dict):
+        if isinstance(spatial_cache, Mapping):
             return self.estimate_spatial_dynamic(gamma, spatial_cache)
         else:
             return spatial_cache

flarestack/core/angular_error_modifier.py

@@ -2,6 +2,7 @@
 Function taken from IceCube astro package.
 """
 
+import numexpr
 import numpy as np
 
 
@@ -33,3 +34,27 @@ def angular_distance(lon1, lat1, lon2, lat2):
     cd = np.cos(lon2 - lon1)
 
     return np.arctan2(np.hypot(c2 * sd, c1 * s2 - s1 * c2 * cd), s1 * s2 + c1 * c2 * cd)
+
+
+# Fast angular distance using numexpr to avoid creating large temporary arrays
+fast_angular_distance = numexpr.NumExpr(
+    "arccos(sin(lat1) * sin(lat2) + cos(lat1) * cos(lat2) * cos(fmod(lon2 - lon1 + pi, 2 * pi) - pi))",
+    signature=[
+        ("lon1", np.float64),
+        ("lat1", np.float64),
+        ("lon2", np.float64),
+        ("lat2", np.float64),
+        ("pi", np.float64),
+    ],
+)
+
+# NB: numexpr implements fmod but not mod, and mod is equivalent to abs(fmod(...))
+in_ra_window = numexpr.NumExpr(
+    "abs(abs(fmod(lon1 - lon2 + pi, 2 * pi)) - pi) < dPhi",
+    signature=[
+        ("lon1", np.float64),
+        ("lon2", np.float64),
+        ("pi", np.float64),
+        ("dPhi", np.float64),
+    ],
+)

flarestack/core/astro.py

@@ -65,7 +65,7 @@ class EnergyPDF(object):
     A base class for Energy PDFs.
     """
 
-    subclasses: dict[str, object] = {}
+    subclasses: "dict[str, type[EnergyPDF]]" = {}
 
     def __init__(self, e_pdf_dict):
         """
@@ -103,7 +103,7 @@ def decorator(subclass):
         return decorator
 
     @classmethod
-    def create(cls, e_pdf_dict):
+    def create(cls, e_pdf_dict) -> "EnergyPDF":
         e_pdf_dict = read_e_pdf_dict(e_pdf_dict)
 
         e_pdf_name = e_pdf_dict["energy_pdf_name"]