Changes made to data handlers, sob_terms and time_profile to make cus…

.gitignore

@@ -13,3 +13,9 @@ Pipfile
 **/.venv
 **/.vim
 pyvenv.cfg
+
+build/
+dist/
+
+build/
+dist/

mla/analysis.py

@@ -29,9 +29,12 @@ class SingleSourceLLHAnalysis:
     sob_term_factories: List[SoBTermFactory]
     data_handler_source: Tuple[DataHandler, PointSource]
     _sob_term_factories: List[SoBTermFactory] = field(init=False, repr=False)
-    _data_handler_source: Tuple[DataHandler, PointSource] = field(init=False, repr=False)
-    _trial_generator: SingleSourceTrialGenerator = field(init=False, repr=False)
-    _test_statistic_factory: LLHTestStatisticFactory = field(init=False, repr=False)
+    _data_handler_source: Tuple[DataHandler,
+                                PointSource] = field(init=False, repr=False)
+    _trial_generator: SingleSourceTrialGenerator = field(
+        init=False, repr=False)
+    _test_statistic_factory: LLHTestStatisticFactory = field(
+        init=False, repr=False)
 
     def produce_and_minimize(
         self,
@@ -66,10 +69,12 @@ def sob_term_factories(self) -> List[SoBTermFactory]:
         return self._sob_term_factories
 
     @sob_term_factories.setter
-    def sob_term_factories(self, sob_term_factories: List[SoBTermFactory]) -> None:
+    def sob_term_factories(
+            self,
+            sob_term_factories: List[SoBTermFactory]) -> None:
         """Docstring"""
         self._sob_term_factories = sob_term_factories
-        self._test_statistic_factory = LLHTestStatisticFactory(  # pylint: disable=too-many-function-args
+        self._test_statistic_factory = LLHTestStatisticFactory(
             self.config['LLHTestStatisticFactory'],
             self._sob_term_factories,
         )
@@ -80,8 +85,9 @@ def data_handler_source(self) -> Tuple[DataHandler, PointSource]:
         return self._data_handler_source
 
     @data_handler_source.setter
-    def data_handler_source(
-            self, data_handler_source: Tuple[DataHandler, PointSource]) -> None:
+    def data_handler_source(self,
+                            data_handler_source: Tuple[DataHandler,
+                                                       PointSource]) -> None:
         """Docstring"""
         self._data_handler_source = data_handler_source
         self._trial_generator = SingleSourceTrialGenerator(

mla/core.py

@@ -6,4 +6,5 @@
 def generate_default_config(classes: list) -> dict:
     """Docstring"""
     return {
-        c.__name__: c.generate_config() for c in classes if issubclass(c, Configurable)}
+        c.__name__: c.generate_config()
+        for c in classes if issubclass(c, Configurable)}

mla/data_handlers.py

@@ -31,7 +31,10 @@ class DataHandler(configurable.Configurable):
     __metaclass__ = abc.ABCMeta
 
     @abc.abstractmethod
-    def sample_background(self, n: int, rng: np.random.Generator) -> np.ndarray:
+    def sample_background(
+            self,
+            n: int,
+            rng: np.random.Generator) -> np.ndarray:
         """Docstring"""
 
     @abc.abstractmethod
@@ -47,7 +50,8 @@ def evaluate_background_sindec_pdf(self, events: np.ndarray) -> np.ndarray:
         """Docstring"""
 
     @abc.abstractmethod
-    def build_background_sindec_logenergy_histogram(self, bins: np.ndarray) -> np.ndarray:
+    def build_background_sindec_logenergy_histogram(
+            self, bins: np.ndarray) -> np.ndarray:
         """Docstring"""
 
     @abc.abstractmethod
@@ -77,7 +81,10 @@ class NuSourcesDataHandler(DataHandler):
     _livetime: float = field(init=False, repr=False)
     _sin_dec_bins: np.ndarray = field(init=False, repr=False)
 
-    def sample_background(self, n: int, rng: np.random.Generator) -> np.ndarray:
+    def sample_background(
+            self,
+            n: int,
+            rng: np.random.Generator) -> np.ndarray:
         """Docstring"""
         return rng.choice(self._data, n)
 
@@ -105,7 +112,8 @@ def evaluate_background_sindec_pdf(self, events: np.ndarray) -> np.ndarray:
         """
         return (1 / (2 * np.pi)) * self._dec_spline(events['sindec'])
 
-    def build_background_sindec_logenergy_histogram(self, bins: np.ndarray) -> np.ndarray:
+    def build_background_sindec_logenergy_histogram(
+            self, bins: np.ndarray) -> np.ndarray:
         """Docstring"""
         return np.histogram2d(
             self._data['sindec'],
@@ -188,14 +196,18 @@ def _cut_sim_dec(self) -> None:
             self._sim = self._full_sim[close].copy()
 
             self._sim['ow'] /= 2 * np.pi * (np.min([np.sin(
-                self.config['dec_cut_location'] + self.config['dec_bandwidth (rad)']
+                self.config['dec_cut_location'] +
+                self.config['dec_bandwidth (rad)']
             ), 1]) - np.max([np.sin(
-                self.config['dec_cut_location'] - self.config['dec_bandwidth (rad)']
+                self.config['dec_cut_location'] -
+                self.config['dec_bandwidth (rad)']
             ), -1]))
             self._sim['weight'] /= 2 * np.pi * (np.min([np.sin(
-                self.config['dec_cut_location'] + self.config['dec_bandwidth (rad)']
+                self.config['dec_cut_location'] +
+                self.config['dec_bandwidth (rad)']
             ), 1]) - np.max([np.sin(
-                self.config['dec_cut_location'] - self.config['dec_bandwidth (rad)']
+                self.config['dec_cut_location'] -
+                self.config['dec_bandwidth (rad)']
             ), -1]))
         else:
             self._sim = self._full_sim
@@ -208,7 +220,8 @@ def data_grl(self) -> Tuple[np.ndarray, np.ndarray]:
     @data_grl.setter
     def data_grl(self, data_grl: Tuple[np.ndarray, np.ndarray]) -> None:
         """Docstring"""
-        self._sin_dec_bins = np.linspace(-1, 1, 1 + self.config['sin_dec_bins'])
+        self._sin_dec_bins = np.linspace(-1, 1,
+                                         1 + self.config['sin_dec_bins'])
         self._data = data_grl[0].copy()
         self._grl = data_grl[1].copy()
         if 'sindec' not in self._data.dtype.names:
@@ -232,7 +245,8 @@ def data_grl(self, data_grl: Tuple[np.ndarray, np.ndarray]) -> None:
             self._data['sindec'], bins=self._sin_dec_bins, density=True)
         bin_centers = bins[:-1] + np.diff(bins) / 2
 
-        self._dec_spline = Spline(bin_centers, hist, **self.config['dec_spline_kwargs'])
+        self._dec_spline = Spline(
+            bin_centers, hist, **self.config['dec_spline_kwargs'])
 
     @property
     def livetime(self) -> float:
@@ -301,7 +315,8 @@ def background_time_profile(self, profile: GenericProfile) -> None:
         background_grl = self._grl[background_run_mask]
         self._n_background = background_grl['events'].sum()
         self._n_background /= background_grl['livetime'].sum()
-        self._n_background *= self._contained_livetime(*profile.range, background_grl)
+        self._n_background *= self._contained_livetime(
+            *profile.range, background_grl)
         self._background_time_profile = copy.deepcopy(profile)
 
     @property
@@ -362,31 +377,38 @@ def _contained_livetime(
 
         return contained_livetime
 
-    def sample_background(self, n: int, rng: np.random.Generator) -> np.ndarray:
+    def sample_background(
+            self,
+            n: int,
+            rng: np.random.Generator) -> np.ndarray:
         """Docstring"""
         events = super().sample_background(n, rng)
         return self._randomize_times(events, self._background_time_profile)
 
     def sample_signal(self, n: int, rng: np.random.Generator) -> np.ndarray:
         """Docstring"""
-        events = super().sample_signal(n, rng)
-        return self._randomize_times(events, self._signal_time_profile)
+        self.events = super().sample_signal(n, rng)
+        return self._randomize_times(self.events, self._signal_time_profile)
 
     def _randomize_times(
         self,
         events: np.ndarray,
         time_profile: GenericProfile,
     ) -> np.ndarray:
-        grl_start_cdf = time_profile.cdf(self._grl['start'])
-        grl_stop_cdf = time_profile.cdf(self._grl['stop'])
-        valid = np.logical_and(grl_start_cdf < 1, grl_stop_cdf > 0)
-        rates = grl_stop_cdf[valid] - grl_start_cdf[valid]
-
+        self.grl_start_cdf = time_profile.cdf(self._grl['start'])
+        self.grl_stop_cdf = time_profile.cdf(self._grl['stop'])
+        self.valid = np.logical_and(
+            self.grl_start_cdf <= 1,
+            self.grl_stop_cdf >= 0)
+        self.rates = self.grl_stop_cdf[self.valid] - \
+            self.grl_start_cdf[self.valid]
+        # print(self.rates, self.rates.sum())
+        # print(len(events))
         runs = np.random.choice(
-            self._grl[valid],
+            self._grl[self.valid],
             size=len(events),
             replace=True,
-            p=rates / rates.sum(),
+            p=self.rates / self.rates.sum(),
         )
 
         events['time'] = time_profile.inverse_transform_sample(

mla/minimizers.py

@@ -33,26 +33,31 @@ class Minimizer(configurable.Configurable):
 
     @abc.abstractmethod
     def __call__(
-            self, fitting_params: Optional[List[str]] = None) -> Tuple[float, np.ndarray]:
+        self, fitting_params: Optional[List[str]] = None
+    ) -> Tuple[float, np.ndarray]:
         """Docstring"""
 
 
 @dataclasses.dataclass
 class GridSearchMinimizer(Minimizer):
     """Docstring"""
+
     def __call__(
-            self, fitting_params: Optional[List[str]] = None) -> Tuple[float, np.ndarray]:
+        self, fitting_params: Optional[List[str]] = None
+    ) -> Tuple[float, np.ndarray]:
         """Docstring"""
         if fitting_params is None:
             fitting_key_idx_map = self.test_statistic.params.key_idx_map
             fitting_bounds = self.test_statistic.params.bounds
         else:
             fitting_key_idx_map = {
-                key: val for key, val in self.test_statistic.params.key_idx_map.items()
+                key: val
+                for key, val in self.test_statistic.params.key_idx_map.items()
                 if key in fitting_params
             }
             fitting_bounds = {
-                key: val for key, val in self.test_statistic.params.bounds.items()
+                key: val
+                for key, val in self.test_statistic.params.bounds.items()
                 if key in fitting_params
             }
 
@@ -72,13 +77,22 @@ def __call__(
         ])
 
         return self._minimize(
-            points[grid_ts_values.argmin()], fitting_key_idx_map, fitting_bounds)
+            points[grid_ts_values.argmin()],
+            fitting_key_idx_map, fitting_bounds)
 
-    def _eval_test_statistic(self, point: np.ndarray, fitting_key_idx_map: dict) -> float:
+    def _eval_test_statistic(
+            self,
+            point: np.ndarray,
+            fitting_key_idx_map: dict) -> float:
         """Docstring"""
-        return self.test_statistic(self._param_values(point, fitting_key_idx_map))
-
-    def _param_values(self, point: np.ndarray, fitting_key_idx_map: dict) -> np.ndarray:
+        return self.test_statistic(
+            self._param_values(
+                point, fitting_key_idx_map))
+
+    def _param_values(
+            self,
+            point: np.ndarray,
+            fitting_key_idx_map: dict) -> np.ndarray:
         """Docstring"""
         param_values = self.test_statistic.params.value_array.copy()
 

mla/params.py

@@ -37,7 +37,10 @@ def names(self) -> List[str]:
         return [*self.key_idx_map]
 
     @classmethod
-    def from_dict(cls, value_dict: dict, bounds: Optional[dict] = None) -> 'Params':
+    def from_dict(
+            cls,
+            value_dict: dict,
+            bounds: Optional[dict] = None) -> 'Params':
         """Docstring"""
         value_array = np.array(list(value_dict.values()))
         key_idx_map = {key: i for i, key in enumerate(value_dict)}
@@ -50,8 +53,11 @@ def from_array(
         bounds: Optional[dict] = None,
     ) -> 'Params':
         """Docstring"""
-        value_array = rf.structured_to_unstructured(named_value_array, copy=True)[0]
-        key_idx_map = {name: i for i, name in enumerate(named_value_array.dtype.names)}
+        value_array = rf.structured_to_unstructured(
+            named_value_array, copy=True)[0]
+        key_idx_map = {
+            name: i for i, name in enumerate(
+                named_value_array.dtype.names)}
         return cls._build_params(value_array, key_idx_map, bounds)
 
     @classmethod

mla/sob_terms.py

@@ -173,7 +173,10 @@ def calculate_drop_mask(self, events: np.ndarray) -> np.ndarray:
         return 1 / self.background_time_profile.pdf(events['time']) != 0
 
     def generate_params(self) -> tuple:
-        return self.signal_time_profile.params, self.signal_time_profile.param_bounds
+        return (
+               self.signal_time_profile.params,
+               self.signal_time_profile.param_bounds
+               )
 
 
 @dataclasses.dataclass
@@ -233,8 +236,10 @@ def __post_init__(self) -> None:
 
     def __call__(self, params: Params, events: np.ndarray) -> SoBTerm:
         """Docstring"""
-        sin_dec_idx = np.searchsorted(self._sin_dec_bins[:-1], events['sindec'])
-        log_energy_idx = np.searchsorted(self._log_energy_bins[:-1], events['logE'])
+        sin_dec_idx = np.searchsorted(
+            self._sin_dec_bins[:-1], events['sindec'])
+        log_energy_idx = np.searchsorted(
+            self._log_energy_bins[:-1], events['logE'])
 
         spline_idxs, event_spline_idxs = np.unique(
             [sin_dec_idx - 1, log_energy_idx - 1],
@@ -276,7 +281,8 @@ def _init_sob_map(
             An array of signal-over-background values binned in sin(dec) and
             log(energy) for a given gamma.
         """
-        sig_h = self.data_handler.build_signal_sindec_logenergy_histogram(gamma, bins)
+        sig_h = self.data_handler.build_signal_sindec_logenergy_histogram(
+            gamma, bins)
 
         # Normalize histogram by dec band
         sig_h /= np.sum(sig_h, axis=1)[:, None]
@@ -292,7 +298,10 @@ def _init_sob_map(
             good_bins, good_vals = bin_centers[good], ratio[i][good]
 
             # Do a linear interpolation across the energy range
-            spline = Spline(good_bins, good_vals, **self.config['energy_spline_kwargs'])
+            spline = Spline(
+                good_bins,
+                good_vals,
+                **self.config['energy_spline_kwargs'])
 
             # And store the interpolated values
             ratio[i] = spline(bin_centers)
@@ -306,7 +315,8 @@ def _init_spline_map(self) -> List[List[Spline]]:
         """
         bins = np.array([self._sin_dec_bins, self._log_energy_bins])
         bin_centers = bins[1, :-1] + np.diff(bins[1]) / 2
-        bg_h = self.data_handler.build_background_sindec_logenergy_histogram(bins)
+        bg_h = self.data_handler.build_background_sindec_logenergy_histogram(
+            bins)
 
         # Normalize histogram by dec band
         bg_h /= np.sum(bg_h, axis=1)[:, None]
@@ -323,7 +333,8 @@ def _init_spline_map(self) -> List[List[Spline]]:
         transposed_log_sob_maps = np.log(sob_maps.transpose(1, 2, 0))
 
         splines = [[
-            Spline(self._gamma_bins, log_ratios, **self.config['sob_spline_kwargs'])
+            Spline(self._gamma_bins, log_ratios,
+                   **self.config['sob_spline_kwargs'])
             for log_ratios in dec_bin
         ] for dec_bin in transposed_log_sob_maps]
 

mla/sources.py

@@ -22,13 +22,18 @@
 @dataclasses.dataclass
 class PointSource(configurable.Configurable):
     """Stores a source object name and location"""
+
     def sample(self, size: int = 1) -> tuple:
         """Sample locations.
 
         Args:
             size: number of points to sample
         """
-        return (np.ones(size) * self.config['ra'], np.ones(size) * self.config['dec'])
+        return (
+            np.ones(size) *
+            self.config['ra'],
+            np.ones(size) *
+            self.config['dec'])
 
     def spatial_pdf(self, events: np.ndarray) -> np.ndarray:
         """calculates the signal probability of events.
@@ -72,6 +77,7 @@ def generate_config(cls) -> dict:
 @dataclasses.dataclass
 class GaussianExtendedSource(PointSource):
     """Gaussian Extended Source"""
+
     def sample(self, size: int = 1) -> np.ndarray:
         """Sample locations.