Accelerate vizapp

libs/ledger/ledger/events.py

@@ -5,10 +5,10 @@
 import ligo.segments
 import numpy as np
 
-from ledger.injections import InterferometerResponseSet
+from ledger.injections import InterferometerResponseSet, shift_mask
 from ledger.ledger import Ledger, metadata, parameter
 
-SECONDS_IN_YEAR = 31556952
+SECONDS_PER_YEAR = 31556952  # 60 * 60 * 24 * 365.2425
 F = TypeVar("F", np.ndarray, float)
 
 
@@ -77,10 +77,7 @@ def get_shift(self, shift: np.ndarray) -> "EventSet":
             EventSet containing only events from the specified shift.
         """
         # downselect to all events from a given shift
-        mask = self.shift == shift
-        if self.shift.ndim == 2:
-            mask = mask.all(axis=-1)
-        return self[mask]
+        return self[shift_mask(self.shift, shift)]
 
     def nb(self, threshold: F) -> F:
         """Calculate number of events above detection threshold.
@@ -121,7 +118,7 @@ def min_far(self):
         Returns:
             Minimum FAR in yr^-1.
         """
-        return (1 / self.Tb) * SECONDS_IN_YEAR
+        return (1 / self.Tb) * SECONDS_PER_YEAR
 
     def far(self, threshold: F) -> F:
         """Calculate false alarm rate (FAR) for a given detection threshold.
@@ -137,7 +134,7 @@ def far(self, threshold: F) -> F:
             FAR in yr^-1. Returns min_far if threshold exceeds all events.
         """
         nb = self.nb(threshold)
-        far = SECONDS_IN_YEAR * nb / self.Tb
+        far = SECONDS_PER_YEAR * nb / self.Tb
         return np.maximum(far, self.min_far)
 
     def significance(self, threshold: F, T: float) -> F:

libs/ledger/ledger/injections.py

@@ -18,6 +18,14 @@
 MSUN = 1.988409902147041637325262574352366540e30
 
 
+def shift_mask(shifts: np.ndarray, target) -> np.ndarray:
+    """Boolean mask of rows in `shifts` equal to `target`."""
+    mask = shifts == np.asarray(target)
+    if shifts.ndim == 2:
+        mask = mask.all(axis=-1)
+    return mask
+
+
 def chirp_mass(
     m1: float | np.ndarray, m2: float | np.ndarray
 ) -> float | np.ndarray:
@@ -764,10 +772,7 @@ def get_shift(self, shift):
         Returns:
             InjectionParameterSet with injections with the specified shift.
         """
-        mask = self.shift == shift
-        if self.shift.ndim == 2:
-            mask = mask.all(axis=-1)
-        return self[mask]
+        return self[shift_mask(self.shift, shift)]
 
     def get_times(self, start: float | None = None, end: float | None = None):
         """

libs/ledger/ledger/ledger.py

@@ -258,6 +258,20 @@ def read(cls, fname: PATH):
         with h5py.File(fname, "r") as f:
             return cls._load_with_idx(f, None)
 
+    @classmethod
+    def read_idx(cls, fname: PATH, idx: np.ndarray):
+        """Read only the rows at `idx` from `fname`.
+
+        Args:
+            fname: Path to HDF5 file to read.
+            idx: Row indices to load.
+
+        Returns:
+            Instance of the ledger class populated with just those rows.
+        """
+        with h5py.File(fname, "r") as f:
+            return cls._load_with_idx(f, idx)
+
     @classmethod
     def sample_from_file(cls, fname: PATH, N: int, replace: bool = False):
         """Sample data from HDF5 file for out-of-memory operations.

libs/ledger/tests/test_events.py

@@ -69,7 +69,7 @@ def test_far(self):
         det_stats = np.arange(10)
         times = np.arange(10)
         shifts = np.array([0] * 4 + [1] * 5 + [2])
-        Tb = 2 * events.SECONDS_IN_YEAR
+        Tb = 2 * events.SECONDS_PER_YEAR
         obj = events.EventSet(det_stats, times, shifts, Tb)
 
         assert obj.far(5) == 2.5

libs/p_astro/p_astro/background.py

@@ -1,5 +1,5 @@
 import numpy as np
-from ledger.events import SECONDS_IN_YEAR, EventSet
+from ledger.events import SECONDS_PER_YEAR, EventSet
 from numpy.polynomial import Polynomial
 from scipy.stats import gaussian_kde
 
@@ -56,7 +56,7 @@ def __init__(self, *args, split: float | None = None, **kwargs):
     @property
     def scale_factor(self):
         """Scale factor to convert from rate density to count rate"""
-        return len(self.background) * SECONDS_IN_YEAR / self.background.Tb
+        return len(self.background) * SECONDS_PER_YEAR / self.background.Tb
 
     def fit(self):
         """

projects/plots/plots/core/compute.py

@@ -1,43 +1,38 @@
-from concurrent.futures import ProcessPoolExecutor
-
 import numpy as np
-from tqdm import trange
-
-
-def init_fn(det_stats, w):
-    global detection_statistics, weights
-    detection_statistics = det_stats
-    weights = w
-
-
-def compute_sv(threshold):
-    mask = detection_statistics >= threshold
-    mus = (weights * mask).sum(-1, keepdims=True)
-    var_summands = weights * (mask - mus)
-    stds = (var_summands**2).sum(-1) ** 0.5
-    return mus[:, 0], stds
 
 
 def sensitive_volume(detection_statistics, weights, thresholds):
-    y = np.empty((len(weights), len(thresholds)))
-    err = np.empty((len(weights), len(thresholds)))
-    ex = ProcessPoolExecutor(
-        8, initializer=init_fn, initargs=(detection_statistics, weights)
-    )
-    with ex:
-        fs = {ex.submit(compute_sv, t): i for i, t in enumerate(thresholds)}
-        for _ in trange(len(thresholds)):
-            while True:
-                for future in fs:
-                    if future.done():
-                        future.exception()
-                        break
-                else:
-                    continue
-                break
-
-            i = fs.pop(future)
-            mu, std = future.result()
-            y[:, i] = mu
-            err[:, i] = std
-    return y, err
+    """Mean and standard error of sensitive volume at each threshold.
+
+    Args:
+        detection_statistics: `(N,)` array of injection detection stats.
+        weights: `(C, N)` array of per-injection weights, one row per
+            mass combo.
+        thresholds: `(T,)` array of detection statistic thresholds.
+
+    Returns:
+        `(y, err)`, each `(C, T)`.
+    """
+    order = np.argsort(detection_statistics)
+    ds_sorted = detection_statistics[order]
+    weights_sorted = weights[:, order]
+
+    # Compute the reverse cumulative sums of the weights and the
+    # weights squared. Because the weights are sorted by the
+    # detection statistics, this avoids needing to compute a mask
+    # for each threshold.
+    n = len(ds_sorted)
+    rev_sum_w = np.zeros((weights.shape[0], n + 1))
+    rev_sum_w2 = np.zeros_like(rev_sum_w)
+    rev_sum_w[:, :-1] = np.cumsum(weights_sorted[:, ::-1], axis=-1)[:, ::-1]
+    rev_sum_w2[:, :-1] = np.cumsum(weights_sorted[:, ::-1] ** 2, axis=-1)[
+        :, ::-1
+    ]
+
+    # idxs is the indices of the first detection statistic greater than
+    # each threshold.
+    idxs = np.searchsorted(ds_sorted, thresholds)
+    mu = rev_sum_w[:, idxs]
+    var = (1 - 2 * mu) * rev_sum_w2[:, idxs] + mu**2 * rev_sum_w2[:, :1]
+    err = np.sqrt(np.maximum(var, 0))
+    return mu, err

projects/plots/plots/core/constants.py

@@ -1,2 +1,5 @@
+from ledger.events import SECONDS_PER_YEAR
+
 SECONDS_PER_MONTH = 60 * 60 * 24 * 30
-SECONDS_PER_YEAR = 60 * 60 * 24 * 365.25
+DEFAULT_NUM_FAR_POINTS = 500
+DEFAULT_VIZAPP_MAX_FAR = 100 * SECONDS_PER_YEAR / SECONDS_PER_MONTH

projects/plots/plots/core/data.py

@@ -44,7 +44,8 @@ def load(
         """Read the ledgers and drop unphysical foreground events.
 
         Args:
-            background: HDF5 file readable by `EventSet.read`
+            background: HDF5 file readable by `EventSet.read`. Gets sorted by
+                detection statistic if it isn't already.
             foreground: HDF5 file readable by `RecoveredInjectionSet.read`
             rejected: HDF5 file readable by `InjectionParameterSet.read`
         """
@@ -54,6 +55,11 @@ def load(
             foreground=drop_unphysical(RecoveredInjectionSet.read(foreground)),
             rejected=InjectionParameterSet.read(rejected),
         )
+
+        # Background should be sorted by infer already, but just in case
+        if not data.background.is_sorted_by("detection_statistic"):
+            data.background = data.background.sort_by("detection_statistic")
+
         logging.info("Read in:")
         logging.info(f"\t{len(data.background)} background events")
         logging.info(f"\t{len(data.foreground)} foreground events")