Minor revision of injector-related code

flarestack/core/data_types.py

@@ -2,8 +2,6 @@
 This module provides the basic data types used by the other modules of flarestack.
 """
 
-import numpy as np
-
 """ Catalogue data type """
 catalogue_dtype = [
     ("ra_rad", float),

flarestack/core/injector.py

@@ -8,7 +8,10 @@
 from flarestack.core.energy_pdf import EnergyPDF, read_e_pdf_dict
 from flarestack.core.time_pdf import TimePDF, read_t_pdf_dict
 from flarestack.core.spatial_pdf import SpatialPDF
-from flarestack.utils.catalogue_loader import calculate_source_weight
+from flarestack.utils.catalogue_loader import (
+    distance_scaled_weight,
+    distance_scaled_weight_sum,
+)
 from scipy import sparse, interpolate
 from flarestack.shared import k_to_flux
 
@@ -78,7 +81,9 @@ def __init__(self, season, sources, **kwargs):
         self.sources = sources
 
         if len(sources) > 0:
-            self.weight_scale = calculate_source_weight(self.sources)
+            self.weight_norm = distance_scaled_weight_sum(sources)
+        else:
+            raise RuntimeError("Catalogue is empty!")
 
         try:
             self.sig_time_pdf = TimePDF.create(
@@ -318,7 +323,7 @@ def calculate_fluence(self, source, scale, source_mc, band_mask, omega):
         # standard candles with flux proportional to 1/d^2 multiplied by the
         # sources weight
 
-        weight = calculate_source_weight(source) / self.weight_scale
+        weight = distance_scaled_weight(source) / self.weight_norm
 
         # Calculate the fluence, using the effective injection time.
         fluence = inj_flux * eff_inj_time * weight
@@ -609,7 +614,7 @@ def inject_signal(self, scale):
 
         return sig_events
 
-    def calculate_single_source(self, source, scale):
+    def calculate_single_source(self, source: np.ndarray, scale: float) -> float:
         # Calculate the effective injection time for simulation. Equal to
         # the overlap between the season and the injection time PDF for
         # the source, scaled if the injection PDF is not uniform in time.
@@ -622,7 +627,7 @@ def calculate_single_source(self, source, scale):
         # standard candles with flux proportional to 1/d^2 multiplied by the
         # sources weight
 
-        weight = calculate_source_weight(source) / self.weight_scale
+        weight = distance_scaled_weight(source) / self.weight_norm
 
         # Calculate the fluence, using the effective injection time.
         fluence = inj_flux * eff_inj_time * weight

flarestack/core/minimisation.py

@@ -22,7 +22,11 @@
 import matplotlib as mpl
 from flarestack.core.time_pdf import TimePDF, Box, Steady
 from flarestack.core.angular_error_modifier import BaseAngularErrorModifier
-from flarestack.utils.catalogue_loader import load_catalogue, calculate_source_weight
+from flarestack.utils.catalogue_loader import (
+    load_catalogue,
+    distance_scaled_weight,
+    distance_scaled_weight_sum,
+)
 from flarestack.utils.asimov_estimator import estimate_discovery_potential
 
 logger = logging.getLogger(__name__)
@@ -619,27 +623,21 @@ def run(self, n_trials, scale=1.0, seed=None):
     def make_season_weight(self, params, season):
         src = self.sources
 
-        weight_scale = calculate_source_weight(src)
-
-        # dist_weight = src["distance_mpc"] ** -2
-        # base_weight = src["base_weight"]
+        weight_norm = distance_scaled_weight_sum(src)
 
         llh = self.get_likelihood(season.season_name)
-        acc = []
 
-        time_weights = []
-        source_weights = []
+        acc, time_weights = [], []
 
         for source in src:
             time_weights.append(llh.sig_time_pdf.effective_injection_time(source))
             acc.append(llh.acceptance(source, params))
-            source_weights.append(calculate_source_weight(source) / weight_scale)
 
         time_weights = np.array(time_weights)
-        source_weights = np.array(source_weights)
-
         acc = np.array(acc).T[0]
 
+        source_weights = distance_scaled_weight(src) / weight_norm
+
         w = acc * time_weights
         w *= source_weights
 

flarestack/utils/asimov_estimator.py

@@ -8,7 +8,11 @@
 from flarestack.core.llh import LLH
 from flarestack.core.astro import angular_distance
 from flarestack.shared import k_to_flux
-from flarestack.utils.catalogue_loader import load_catalogue, calculate_source_weight
+from flarestack.utils.catalogue_loader import (
+    load_catalogue,
+    distance_scaled_weight,
+    distance_scaled_weight_sum,
+)
 from scipy.stats import norm
 import logging
 
@@ -37,7 +41,7 @@ def ts_weight(n_s):
 
     # def weight_ts(ts, n_s)
 
-    weight_scale = calculate_source_weight(sources)
+    weight_norm = distance_scaled_weight_sum(sources)
 
     livetime = 0.0
 
@@ -89,14 +93,16 @@ def signalness(sig_over_background):
         sig_times = np.array(
             [llh.sig_time_pdf.effective_injection_time(x) for x in sources]
         )
-        source_weights = np.array([calculate_source_weight(x) for x in sources])
-        mean_time = np.sum(sig_times * source_weights) / weight_scale
+
+        source_weights = distance_scaled_weight(sources)
+
+        mean_time = np.sum(sig_times * source_weights) / weight_norm
 
         # print(source_weights)
 
         fluences = (
             np.array([x * sig_times[i] for i, x in enumerate(source_weights)])
-            / weight_scale
+            / weight_norm
         )
         # print(sources.dtype.names)
         # print(sources["dec_rad"], np.sin(sources["dec_rad"]))

flarestack/utils/catalogue_loader.py

@@ -2,18 +2,15 @@
 from numpy.lib.recfunctions import append_fields, rename_fields
 
 
-def calculate_source_weight(sources) -> float:
-    """
-    Depending on the dimension of the input, calculate:
-    - the sum of the weights for a vector of sources
-    - the absolute weight of an individual source.
+def distance_scaled_weight(sources: np.ndarray) -> np.ndarray:
+    return sources["base_weight"] * sources["distance_mpc"] ** -2
 
-    Dividing the absolute weight of a single source by the sum of the weights of all sources in the catalogue gives the relative weight of the source, i.e. the fraction of the fitted flux that is produced by it. The fraction of the injected flux may be different, since the "injection weight modifier" is not accounted for.
-    """
-    return np.sum(sources["base_weight"] * sources["distance_mpc"] ** -2)
 
+def distance_scaled_weight_sum(sources: np.ndarray) -> float:
+    return np.sum(distance_scaled_weight(sources))
 
-def load_catalogue(path):
+
+def load_catalogue(path) -> np.ndarray:
     sources = np.load(path)
 
     # Maintain backwards-compatibility
@@ -57,36 +54,15 @@ def load_catalogue(path):
         )
 
     # Check that all sources have a unique name
-
     if len(set(sources["source_name"])) < len(sources["source_name"]):
         raise Exception(
             "Some sources in catalogue do not have unique "
             "names. Please assign unique names to each source."
         )
 
-    # Rescale 'base_weight'
-    # sources["base_weight"] /= np.mean(sources["base_weight"])
+    # TODO: add a check on `injection_weight_modifier`
 
-    # Order sources
+    # Order sources by distance
     sources = np.sort(sources, order="distance_mpc")
 
     return sources
-
-
-# def convert_catalogue(path):
-#     print "Converting", path
-#     cat = load_catalogue(path)
-#     print cat
-#     # np.save(path, cat)
-
-# if __name__ == "__main__":
-#     import os
-#     from flarestack.analyses.agn_cores.shared_agncores import agncores_cat_dir
-#
-#     # for path in os.listdir(agncores_cat_dir):
-#     for path in ["radioloud_radioselected_100brightest_srcs.npy"]:
-#         filename = agncores_cat_dir + path
-#         cat = load_catalogue(filename)
-#         cat["base_weight"] = cat['injection_weight_modifier']
-#         cat['injection_weight_modifier'] = np.ones_like(cat["base_weight"])
-#         np.save(filename, cat)