Improve cache coherence

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/core/energy_pdf.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"]

flarestack/core/injector.py

@@ -3,9 +3,10 @@
 import random
 import zipfile
 import zlib
+from typing import TYPE_CHECKING
 
 import numpy as np
-from astropy.table import Table
+from astropy.table import Table, vstack
 from scipy import interpolate, sparse
 
 from flarestack.core.energy_pdf import EnergyPDF, read_e_pdf_dict
@@ -14,6 +15,9 @@
 from flarestack.shared import band_mask_cache_name, k_to_flux
 from flarestack.utils.catalogue_loader import calculate_source_weight
 
+if TYPE_CHECKING:
+    from flarestack.data import SeasonWithMC
+
 logger = logging.getLogger(__name__)
 
 
@@ -163,7 +167,7 @@ def create_dataset(self, scale, angular_error_modifier=None):
             sig_events = []
 
         if len(sig_events) > 0:
-            simulated_data = np.concatenate((bkg_events, sig_events))
+            simulated_data = vstack((bkg_events, sig_events))
         else:
             simulated_data = bkg_events
 
@@ -172,8 +176,8 @@ def create_dataset(self, scale, angular_error_modifier=None):
 
         return simulated_data
 
-    def inject_signal(self, scale):
-        return
+    def inject_signal(self, scale: float) -> Table:
+        raise NotImplementedError
 
     @classmethod
     def register_subclass(cls, inj_name):
@@ -245,7 +249,7 @@ def __init__(self, season, sources, **kwargs):
             logger.warning("No Injection Arguments. Are you unblinding?")
             pass
 
-    def get_mc(self, season):
+    def get_mc(self, season: "SeasonWithMC") -> Table:
         return season.get_mc()
 
     def select_mc_band(self, source):
@@ -337,7 +341,7 @@ def calculate_fluence(self, source, scale, source_mc, band_mask, omega):
 
         return source_mc
 
-    def inject_signal(self, scale):
+    def inject_signal(self, scale: float) -> Table:
         """Randomly select simulated events from the Monte Carlo dataset to
         simulate a signal for each source. The source flux can be scaled by
         the scale parameter.
@@ -370,7 +374,7 @@ def inject_signal(self, scale):
                 n_s = int(n_inj)
 
             try:
-                f_n_inj = float(n_inj[0])
+                f_n_inj = float(n_inj[0])  # type: ignore[index]
             except (TypeError, IndexError):
                 f_n_inj = float(n_inj)
 
@@ -537,14 +541,12 @@ class TableInjector(MCInjector):
     For 1000 sources, calculate_n_exp() is ~60x faster than MCInjector.
     """
 
-    def get_mc(self, season):
-        mc: np.ndarray = season.get_mc()
-        # Sort rows by trueDec, and store as columns in a Table
-        table = Table(mc[np.argsort(mc["trueDec"].copy())])
-        # Prevent in-place modifications
-        for k in table.columns:
-            table[k].setflags(write=False)
-        return table
+    def get_mc(self, season: "SeasonWithMC") -> Table:
+        mc = season.get_mc().copy(copy_data=False)
+        mc.sort("trueDec")
+        for col in mc.columns.values():
+            col.setflags(write=False)
+        return mc
 
     def get_band_mask(self, source, min_dec, max_dec):
         return slice(*np.searchsorted(self._mc["trueDec"], [min_dec, max_dec]))
@@ -570,9 +572,9 @@ def __init__(self, season, sources, **kwargs):
         self.n_exp = self.calculate_n_exp()
         self.conversion_cache = dict()
 
-    def inject_signal(self, scale):
+    def inject_signal(self, scale: float) -> Table:
         # Creates empty signal event array
-        sig_events = np.empty((0,), dtype=self.season.get_background_dtype())
+        sig_events = []
 
         n_tot_exp = 0
 
@@ -598,7 +600,7 @@ def inject_signal(self, scale):
             logger.debug(
                 "Injected {0} events with an expectation of {1:.2f} events for {2}".format(
                     n_s,
-                    n_inj if isinstance(n_inj, float) else float(n_inj[0]),
+                    n_inj if isinstance(n_inj, float) else float(n_inj[0]),  # type: ignore[index]
                     source["source_name"],
                 )
             )
@@ -607,7 +609,20 @@ def inject_signal(self, scale):
             if n_s < 1:
                 continue
 
-            sim_ev = np.empty((n_s,), dtype=self.season.get_background_dtype())
+            sim_ev = Table(
+                np.empty(
+                    (n_s,),
+                    dtype=np.dtype(
+                        self.season.get_background_dtype().descr
+                        + [
+                            ("trueRa", float),
+                            ("trueDec", float),
+                            ("trueE", float),
+                            ("ow", float),
+                        ]
+                    ),
+                )
+            )
 
             # Fills the energy proxy conversion cache
 
@@ -631,16 +646,14 @@ def inject_signal(self, scale):
             sim_ev["raw_sigma"] = sim_ev["sigma"].copy()
 
             sim_ev = self.spatial_pdf.simulate_distribution(source, sim_ev)
+            sim_ev.keep_columns(self.season.get_background_dtype().names)
 
-            sim_ev = sim_ev[list(self.season.get_background_dtype().names)].copy()
-            #
+            sig_events.append(sim_ev)
 
-            # Joins the new events to the signal events
-            sig_events = np.concatenate((sig_events, sim_ev))
-
-        sig_events = np.array(sig_events)
-
-        return sig_events
+        if sig_events:
+            return vstack(sig_events)
+        else:
+            return Table()
 
     def calculate_single_source(self, source, scale):
         # Calculate the effective injection time for simulation. Equal to

flarestack/core/spatial_pdf.py

@@ -4,7 +4,7 @@
 
 import healpy as hp
 import numpy as np
-from numpy.lib.recfunctions import append_fields
+from astropy.table import Table
 from scipy.interpolate import interp1d
 from scipy.stats import norm
 
@@ -44,13 +44,12 @@ def __init__(self, spatial_pdf_dict, season):
 class SignalSpatialPDF:
     """Base Signal Spatial PDF class."""
 
-    subclasses: dict[str, object] = {}
+    subclasses: "dict[str, type[SignalSpatialPDF]]" = {}
 
     def __init__(self, spatial_pdf_dict):
         pass
 
-    @staticmethod
-    def simulate_distribution(source, data):
+    def simulate_distribution(self, source, data: Table) -> Table:
         return data
 
     @staticmethod
@@ -70,7 +69,7 @@ def decorator(subclass):
         return decorator
 
     @classmethod
-    def create(cls, s_pdf_dict):
+    def create(cls, s_pdf_dict) -> "SignalSpatialPDF":
         try:
             s_pdf_name = s_pdf_dict["spatial_pdf_name"]
         except KeyError:
@@ -132,7 +131,7 @@ def rotate(ra1, dec1, ra2, dec2, ra3, dec3):
         ra = phi + np.pi
         return np.atleast_1d(ra), np.atleast_1d(dec)
 
-    def rotate_to_position(self, ev, ra, dec):
+    def rotate_to_position(self, ev: Table, ra: float, dec: float) -> Table:
         """Modifies the events by reassigning the Right Ascension and
         Declination of the events. Rotates the events, so that they are
         distributed as if they originated from the source. Removes the
@@ -162,29 +161,24 @@ def rotate_to_position(self, ev, ra, dec):
         ev["sinDec"] = np.sin(rot_dec)
 
         # Deletes the Monte Carlo information from sampled events
-        non_mc = [
-            name for name in names if name not in ["trueRa", "trueDec", "trueE", "ow"]
-        ]
-        ev = ev[non_mc].copy()
+        ev.remove_columns(["trueRa", "trueDec", "trueE", "ow"])
 
         return ev
 
 
 @SignalSpatialPDF.register_subclass("circular_gaussian")
 class CircularGaussian(SignalSpatialPDF):
-    def simulate_distribution(self, source, data):
-        data["ra"] = np.pi + norm.rvs(size=len(data)) * data["sigma"]
-        data["dec"] = norm.rvs(size=len(data)) * data["sigma"]
-        data["sinDec"] = np.sin(data["dec"])
-        data = append_fields(
-            data,
-            ["trueRa", "trueDec"],
-            [np.ones_like(data["dec"]) * np.pi, np.zeros_like(data["dec"])],
-        ).copy()
+    def simulate_distribution(self, source, data: Table) -> Table:
+        data = data.copy()
+        data["ra"][:] = np.pi + norm.rvs(size=len(data)) * data["sigma"]
+        data["dec"][:] = norm.rvs(size=len(data)) * data["sigma"]
+        data["sinDec"][:] = np.sin(data["dec"])
+        data["trueRa"][:] = np.ones_like(data["dec"]) * np.pi
+        data["trueDec"][:] = np.zeros_like(data["dec"])
 
-        data = self.rotate_to_position(data, source["ra_rad"], source["dec_rad"]).copy()
+        data = self.rotate_to_position(data, source["ra_rad"], source["dec_rad"])
 
-        return data.copy()
+        return data
 
     @staticmethod
     def signal_spatial(source, cut_data):
@@ -316,13 +310,12 @@ def simulate_distribution(self, source, data, gamma=2.0):
         data["ra"] = np.pi + distance * np.cos(phi)
         data["dec"] = distance * np.sin(phi)
         data["sinDec"] = np.sin(data["dec"])
-        data = append_fields(
-            data,
-            ["trueRa", "trueDec"],
+        data.add_columns(
             [np.ones_like(data["dec"]) * np.pi, np.zeros_like(data["dec"])],
-        ).copy()
+            ["trueRa", "trueDec"],
+        )
 
-        data = self.rotate_to_position(data, source["ra_rad"], source["dec_rad"]).copy()
+        data = self.rotate_to_position(data, source["ra_rad"], source["dec_rad"])
 
         return data.copy()
 
@@ -393,7 +386,7 @@ def signal_spatial(source, cut_data, gamma: Optional[float]):  # type: ignore[ov
 
 
 class BackgroundSpatialPDF:
-    subclasses: dict[str, object] = {}
+    subclasses: "dict[str, type[BackgroundSpatialPDF]]" = {}
 
     def __init__(self, spatial_pdf_dict, season):
         pass
@@ -411,7 +404,7 @@ def decorator(subclass):
         return decorator
 
     @classmethod
-    def create(cls, s_pdf_dict, season):
+    def create(cls, s_pdf_dict, season) -> "BackgroundSpatialPDF":
         try:
             s_pdf_name = s_pdf_dict["bkg_spatial_pdf"]
         except KeyError:

flarestack/core/time_pdf.py

@@ -96,7 +96,7 @@ def read_t_pdf_dict(t_pdf_dict):
 
 
 class TimePDF:
-    subclasses: dict[str, object] = {}
+    subclasses: "dict[str, type[TimePDF]]" = {}
 
     def __init__(self, t_pdf_dict, livetime_pdf=None):
         self.t_dict = t_pdf_dict
@@ -128,7 +128,7 @@ def decorator(subclass):
         return decorator
 
     @classmethod
-    def create(cls, t_pdf_dict, livetime_pdf=None):
+    def create(cls, t_pdf_dict, livetime_pdf=None) -> "TimePDF":
         t_pdf_dict = read_t_pdf_dict(t_pdf_dict)
 
         t_pdf_name = t_pdf_dict["time_pdf_name"]