feat: refactor a and b classes to allow for filtering of additional args

seismostats/analysis/avalue/base.py

@@ -18,19 +18,16 @@ def __init__(self) -> Self:
         self.scaling_factor: float | None = None
         self.m_ref: float | None = None
         self.b_value: float | None = None
-        self.idx: np.ndarray | None = None
 
         self.__a_value: float | None = None
 
     def calculate(self,
                   magnitudes: np.ndarray,
                   mc: float,
                   delta_m: float,
-                  *args,
                   scaling_factor: float | None = None,
                   m_ref: float | None = None,
-                  b_value: float | None = None,
-                  **kwargs) -> float:
+                  b_value: float | None = None) -> float:
         '''
         Return the a-value estimate.
 
@@ -59,27 +56,29 @@ def calculate(self,
         self.m_ref = m_ref
         self.b_value = b_value
 
-        self._filtering()
+        self._filter_magnitudes()
         self._sanity_checks()
 
-        self.__a_value = self._estimate(*args, **kwargs)
+        self.__a_value = self._estimate()
         self.__a_value = self._reference_scaling(self.__a_value)
 
         return self.__a_value
 
     @abstractmethod
-    def _estimate(self, *args, **kwargs) -> float:
+    def _estimate(self) -> float:
         '''
         Specific implementation of the a-value estimator.
         '''
         pass
 
-    def _filtering(self) -> np.ndarray:
+    def _filter_magnitudes(self) -> np.ndarray:
         '''
         Filter out magnitudes below the completeness magnitude.
         '''
-        self.idx = self.magnitudes >= self.mc - self.delta_m / 2
-        self.magnitudes = self.magnitudes[self.idx]
+        idx = self.magnitudes >= self.mc - self.delta_m / 2
+        self.magnitudes = self.magnitudes[idx]
+
+        return idx
 
     def _sanity_checks(self):
         '''

seismostats/analysis/avalue/more_positive.py

@@ -56,46 +56,54 @@ def calculate(self,
             # not passed to the function
             raise ValueError("b_value must be given")
 
+        self.times: np.ndarray = np.array(times)
+        self.dmc: float = dmc if dmc is not None else delta_m
+
+        if self.dmc < 0:
+            raise ValueError("dmc must be larger or equal to 0.")
+
+        if self.dmc < delta_m and get_option("warnings") is True:
+            warnings.warn("dmc is smaller than delta_m, not recommended.")
+
         return super().calculate(magnitudes,
                                  mc=mc,
                                  delta_m=delta_m,
-                                 times=times,
                                  scaling_factor=scaling_factor,
                                  m_ref=m_ref,
                                  b_value=b_value,
-                                 dmc=dmc,
                                  )
 
-    def _estimate(self, times: np.ndarray, dmc: float | None = None) -> float:
-        if dmc is None:
-            dmc = self.delta_m
-        elif dmc < 0:
-            raise ValueError("dmc must be larger or equal to 0")
-        elif dmc < self.delta_m and get_option("warnings") is True:
-            warnings.warn("dmc is smaller than delta_m, not recommended")
+    def _filter_magnitudes(self):
+        '''
+        Filter out magnitudes below the completeness magnitude.
+        '''
+        idx = super()._filter_magnitudes()
+
+        self.times = self.times[idx]
 
-        times = np.array(times)
-        times = times[self.idx]
+        return idx
 
+    def _estimate(self) -> float:
         # order the magnitudes and times
-        srt = np.argsort(times)
+        srt = np.argsort(self.times)
         self.magnitudes = self.magnitudes[srt]
-        times = times[srt]
+        self.times = self.times[srt]
 
         # differences
-        idx_next_larger = find_next_larger(self.magnitudes, self.delta_m, dmc)
-        time_diffs = times[idx_next_larger] - times
+        idx_next_larger = find_next_larger(
+            self.magnitudes, self.delta_m, self.dmc)
+        time_diffs = self.times[idx_next_larger] - self.times
 
         # deal with events which do not have a next larger event
         idx_no_next = idx_next_larger == 0
-        time_diffs[idx_no_next] = times[-1] - times[idx_no_next]
+        time_diffs[idx_no_next] = self.times[-1] - self.times[idx_no_next]
 
         # estimate the number of events within the time interval
-        total_time = times[-1] - times[0]
+        total_time = self.times[-1] - self.times[0]
 
         # scale the time
         tau = time_diffs * 10**(-self.b_value
-                                * (self.magnitudes + dmc - self.mc))
+                                * (self.magnitudes + self.dmc - self.mc))
 
         time_factor = sum(tau / total_time)
         n_more_pos = sum(~idx_no_next) / time_factor

seismostats/analysis/avalue/positive.py

@@ -54,43 +54,52 @@ def calculate(self,
         Returns:
             a_pos: a-value of the Gutenberg-Richter distribution
         '''
+
+        self.times: np.ndarray = np.array(times)
+        self.dmc: float = dmc if dmc is not None else delta_m
+
+        if self.dmc < 0:
+            raise ValueError('dmc must be larger or equal to 0.')
+
+        if self.dmc < delta_m and get_option('warnings') is True:
+            warnings.warn('dmc is smaller than delta_m, not recommended.')
+
         return super().calculate(magnitudes,
                                  mc=mc,
                                  delta_m=delta_m,
-                                 times=times,
                                  scaling_factor=scaling_factor,
                                  m_ref=m_ref,
-                                 b_value=b_value,
-                                 dmc=dmc,
+                                 b_value=b_value
                                  )
 
-    def _estimate(self, times: np.ndarray, dmc: float | None = None) -> float:
-        if dmc is None:
-            dmc = self.delta_m
-        elif dmc < 0:
-            raise ValueError('dmc must be larger or equal to 0.')
-        elif dmc < self.delta_m and get_option('warnings') is True:
-            warnings.warn('dmc is smaller than delta_m, not recommended.')
+    def _filter_magnitudes(self) -> np.ndarray:
+        '''
+        Filter out magnitudes below the completeness magnitude.
+        '''
+        idx = super()._filter_magnitudes()
+
+        self.times = self.times[idx]
+
+        return idx
 
-        times = np.array(times)
-        times = times[self.idx]
+    def _estimate(self) -> float:
 
         # order the magnitudes and times
-        idx = np.argsort(times)
+        idx = np.argsort(self.times)
         self.magnitudes = self.magnitudes[idx]
-        times = times[idx]
+        self.times = self.times[idx]
 
         # differences
         mag_diffs = np.diff(self.magnitudes)
-        time_diffs = np.diff(times)
+        time_diffs = np.diff(self.times)
 
         # only consider events with magnitude difference >= dmc
-        idx = mag_diffs > dmc - self.delta_m / 2
+        idx = mag_diffs > self.dmc - self.delta_m / 2
         mag_diffs = mag_diffs[idx]
         time_diffs = time_diffs[idx]
 
         # estimate the number of events within the time interval
-        total_time = times[-1] - times[0]
+        total_time = self.times[-1] - self.times[0]
 
         time_factor = sum(time_diffs / total_time)
         n_pos = sum(idx) / time_factor

seismostats/analysis/avalue/tests/test_avalue.py

@@ -112,3 +112,12 @@ def test_estimate_a_more_positive():
     with pytest.raises(ValueError):
         estimate_a(mags, mc=0, delta_m=0.1, times=times,
                    m_ref=-1, method=AMorePositiveAValueEstimator)
+
+    # check that cutting at mc is handled correctly
+    mags = np.array([-0.5, 0.9, 0.9, 0.9, 0.9, 10.9])
+    times = np.arange(datetime(2000, 1, 1), datetime(
+        2000, 1, 7), timedelta(days=1)).astype(datetime)
+
+    a = estimator.calculate(
+        mags, 0, 0.1, times, b_value=1, dmc=0.1)
+    assert_almost_equal(10**a, 16.0)

seismostats/analysis/bvalue/base.py

@@ -25,9 +25,7 @@ def calculate(self,
                   magnitudes: np.ndarray | list,
                   mc: float,
                   delta_m: float,
-                  *args,
-                  weights: np.ndarray | list | None = None,
-                  **kwargs) -> float:
+                  weights: np.ndarray | list | None = None) -> float:
         '''
         Return the b-value estimate.
 
@@ -46,28 +44,30 @@ def calculate(self,
         self.delta_m = delta_m
         self.weights = None if weights is None else np.array(weights)
 
-        self._filtering()
+        self._filter_magnitudes()
         self._sanity_checks()
 
-        self.__b_value = self._estimate(*args, **kwargs)
+        self.__b_value = self._estimate()
         return self.__b_value
 
     @abstractmethod
-    def _estimate(self, *args, **kwargs) -> float:
+    def _estimate(self) -> float:
         '''
         Specific implementation of the b-value estimator.
         '''
         pass
 
-    def _filtering(self):
+    def _filter_magnitudes(self):
         '''
         Filter out magnitudes below the completeness magnitude.
         '''
-        self.idx = self.magnitudes >= self.mc - self.delta_m / 2
-        self.magnitudes = self.magnitudes[self.idx]
+        idx = self.magnitudes >= self.mc - self.delta_m / 2
+        self.magnitudes = self.magnitudes[idx]
 
         if self.weights is not None:
-            self.weights = self.weights[self.idx]
+            self.weights = self.weights[idx]
+
+        return idx
 
     def _sanity_checks(self):
         '''

seismostats/analysis/bvalue/more_positive.py

@@ -44,22 +44,20 @@ def calculate(self,
                         below this value are not considered). If None,
                         the cutoff is set to delta_m.
         '''
-
-        return super().calculate(magnitudes,
-                                 mc=mc,
-                                 delta_m=delta_m,
-                                 weights=weights,
-                                 dmc=dmc)
-
-    def _estimate(self, dmc: float) -> float:
-
-        self.dmc = dmc or self.delta_m
+        self.dmc: float = dmc if dmc is not None else delta_m
 
         if self.dmc < 0:
             raise ValueError('dmc must be larger or equal to 0')
-        elif self.dmc < self.delta_m and get_option('warnings') is True:
+
+        if self.dmc < delta_m and get_option('warnings') is True:
             warnings.warn('dmc is smaller than delta_m, not recommended')
 
+        return super().calculate(magnitudes,
+                                 mc=mc,
+                                 delta_m=delta_m,
+                                 weights=weights)
+
+    def _estimate(self) -> float:
         mag_diffs = -np.ones(len(self.magnitudes) - 1) * self.delta_m
 
         idx_next_larger = find_next_larger(
@@ -74,6 +72,6 @@ def _estimate(self, dmc: float) -> float:
             self.weights = weights[~idx_no_next]
 
         return _mle_estimator(self.magnitudes,
-                              mc=self.mc,
+                              mc=self.dmc,
                               delta_m=self.delta_m,
                               weights=self.weights)

seismostats/analysis/bvalue/positive.py

@@ -42,22 +42,21 @@ def calculate(self,
                         the cutoff is set to delta_m.
         '''
 
-        return super().calculate(magnitudes,
-                                 mc=mc,
-                                 delta_m=delta_m,
-                                 weights=weights,
-                                 dmc=dmc)
-
-    def _estimate(self, dmc: float | None = None) -> float:
-
-        self.dmc = dmc or self.delta_m
+        self.dmc: float = dmc if dmc is not None else delta_m
 
         if self.dmc < 0:
             raise ValueError('dmc must be larger or equal to 0')
 
-        elif self.dmc < self.delta_m and get_option('warnings') is True:
+        if self.dmc < delta_m and get_option('warnings') is True:
             warnings.warn('dmc is smaller than delta_m, not recommended')
 
+        return super().calculate(magnitudes,
+                                 mc=mc,
+                                 delta_m=delta_m,
+                                 weights=weights)
+
+    def _estimate(self, dmc: float | None = None) -> float:
+
         # only take the values where the next earthquake is d_mc larger than the
         # previous one. delta_m is added to avoid numerical errors
         self.magnitudes = np.diff(self.magnitudes)