include index capability and time vector input in more_positive

seismostats/analysis/avalue/more_positive.py

@@ -73,7 +73,7 @@ def calculate(self,
                                  b_value=b_value,
                                  )
 
-    def _filter_magnitudes(self):
+    def _filter_magnitudes(self) -> np.ndarray:
         '''
         Filter out magnitudes below the completeness magnitude.
         '''
@@ -88,8 +88,9 @@ def _estimate(self) -> float:
         srt = np.argsort(self.times)
         self.magnitudes = self.magnitudes[srt]
         self.times = self.times[srt]
+        self.idx = self.idx[srt]
 
-        # differences
+        # find next larger event (if it exists)
         idx_next_larger = find_next_larger(
             self.magnitudes, self.delta_m, self.dmc)
         time_diffs = self.times[idx_next_larger] - self.times
@@ -108,5 +109,11 @@ def _estimate(self) -> float:
         time_factor = sum(tau / total_time)
         n_more_pos = sum(~idx_no_next) / time_factor
 
+        # make sure that all attributes are consistent
+        idx_next_larger = idx_next_larger[~idx_no_next]
+        self.magnitudes = self.magnitudes[idx_next_larger]
+        self.times = self.times[idx_next_larger]
+        self.idx = self.idx[idx_next_larger]
+
         # estimate a-value
         return np.log10(n_more_pos)

seismostats/analysis/avalue/positive.py

@@ -95,7 +95,6 @@ def _estimate(self) -> float:
 
         # only consider events with magnitude difference >= dmc
         is_larger = mag_diffs >= self.dmc - self.delta_m / 2
-        mag_diffs = mag_diffs[is_larger]
         time_diffs = time_diffs[is_larger]
         self.magnitudes = self.magnitudes[1:][is_larger]
         self.idx = self.idx[1:][is_larger]

seismostats/analysis/avalue/tests/test_avalue.py

@@ -101,6 +101,7 @@ def test_estimate_a_positive():
                       datetime(2000, 1, 5)])
     estimator.calculate(mags, mc=-1, delta_m=0.1, times=times)
     assert (mags[estimator.idx] == np.array([0.9, 0.2, 0.5])).all()
+    assert (mags[estimator.idx] == estimator.magnitudes).all()
 
     # test that time array is correctly used
     mags = np.array([0, 0.9, 0.5, 0.2, -1])
@@ -111,6 +112,7 @@ def test_estimate_a_positive():
                       datetime(2000, 1, 3)])
     estimator.calculate(mags, mc=-1, delta_m=0.1, times=times)
     assert (mags[estimator.idx] == np.array([0.9, 0.2, 0.5])).all()
+    assert (mags[estimator.idx] == estimator.magnitudes).all()
 
 
 @pytest.mark.filterwarnings("ignore")
@@ -142,3 +144,29 @@ def test_estimate_a_more_positive():
     a = estimator.calculate(
         mags, 0, 0.1, times, b_value=1, dmc=0.1)
     assert_almost_equal(10**a, 16.0)
+
+    # test that index works correctly
+    mags = np.array([0, 0.9, -1, 0.2, 0.5, 1])
+    times = np.array([datetime(2000, 1, 1),
+                      datetime(2000, 1, 2),
+                      datetime(2000, 1, 3),
+                      datetime(2000, 1, 4),
+                      datetime(2000, 1, 5),
+                      datetime(2000, 1, 6)])
+    estimator.calculate(mags, mc=-1, delta_m=0.1, times=times, b_value=1)
+    assert (mags[estimator.idx] == np.array([0.9, 1, 0.2, 0.5, 1])).all()
+    assert (estimator.idx == np.array([1, 5, 3, 4, 5])).all()
+    assert (mags[estimator.idx] == estimator.magnitudes).all()
+    assert (estimator.times == times[estimator.idx]).all()
+
+    # test that time array is correctly used
+    mags = np.array([0, 0.9, 0.5, 0.2, -1])
+    times = np.array([datetime(2000, 1, 1),
+                      datetime(2000, 1, 2),
+                      datetime(2000, 1, 5),
+                      datetime(2000, 1, 4),
+                      datetime(2000, 1, 3)])
+    estimator.calculate(mags, mc=-1, delta_m=0.1, times=times, b_value=1)
+    assert (mags[estimator.idx] == np.array([0.9, 0.2, 0.5])).all()
+    assert (mags[estimator.idx] == estimator.magnitudes).all()
+    assert (estimator.times == times[estimator.idx]).all()

seismostats/analysis/bvalue/more_positive.py

@@ -29,6 +29,7 @@ def calculate(self,
                   magnitudes: np.ndarray,
                   mc: float,
                   delta_m: float,
+                  times: np.ndarray | None = None,
                   weights: np.ndarray | None = None,
                   dmc: float | None = None) -> float:
         '''
@@ -39,11 +40,19 @@ def calculate(self,
             magnitudes: Array of magnitudes
             mc:         Completeness magnitude
             delta_m:    Discretization of magnitudes.
+            times:      Vector of times of the events, in any format (datetime,
+                    float, etc.). If None, it is assumed that the events are
+                    ordered in time.
             weights:    Array of weights for the magnitudes.
-            dmc:        Cutoff value for the differences (differences
-                        below this value are not considered). If None,
-                        the cutoff is set to delta_m.
+            dmc:        Cutoff value for the differences (differences below
+                    this value are not considered). If None, the cutoff is set
+                    to delta_m.
         '''
+        if times is None:
+            self.times = None
+        else:
+            self.times: np.ndarray = np.array(times)
+
         self.dmc: float = dmc if dmc is not None else delta_m
 
         if self.dmc < 0:
@@ -57,19 +66,41 @@ def calculate(self,
                                  delta_m=delta_m,
                                  weights=weights)
 
+    def _filter_magnitudes(self) -> np.ndarray:
+        '''
+        Filter out magnitudes below the completeness magnitude.
+        '''
+        idx = super()._filter_magnitudes()
+
+        if self.times is not None:
+            self.times = self.times[idx]
+        return idx
+
     def _estimate(self) -> float:
+        if self.times is not None:
+            srt = np.argsort(self.times)
+            self.magnitudes = self.magnitudes[srt]
+            self.times = self.times[srt]
+            if self.weights is not None:
+                self.weights = self.weights[srt]
+            self.idx = self.idx[srt]
+
+        # calculate mg diffs to next larger magnitude
         mag_diffs = -np.ones(len(self.magnitudes) - 1) * self.delta_m
-
         idx_next_larger = find_next_larger(
             self.magnitudes, self.delta_m, self.dmc)
         mag_diffs = self.magnitudes[idx_next_larger] - self.magnitudes
 
         idx_no_next = idx_next_larger == 0
         self.magnitudes = mag_diffs[~idx_no_next]
 
+        # make sure that all attributes are consistent
+        idx_next_larger = idx_next_larger[~idx_no_next]
+        self.idx = self.idx[idx_next_larger]
         if self.weights is not None:
-            weights = self.weights[idx_next_larger]
-            self.weights = weights[~idx_no_next]
+            self.weights = self.weights[idx_next_larger]
+        if self.times is not None:
+            self.times = self.times[idx_next_larger]
 
         return _mle_estimator(self.magnitudes,
                               mc=self.dmc,

seismostats/analysis/bvalue/tests/test_bvalues.py

@@ -215,3 +215,22 @@ def test_estimate_b_more_positive(
     assert_almost_equal(b_estimate, b_est_correct)
     assert_almost_equal(b_estimate, b_estimate_weighted)
     assert_almost_equal(b_estimate, b_estimate_half_weighted)
+
+    # test that index works correctly
+    mags = np.array([0, 0.9, -1, 0.2, 0.5, 1])
+    estimator.calculate(mags, mc=-1, delta_m=0.1)
+    assert (mags[estimator.idx] == np.array([0.9, 1, 0.2, 0.5, 1])).all()
+    assert (estimator.idx == np.array([1, 5, 3, 4, 5])).all()
+    assert (estimator.times is None)
+
+    # test that time array is correctly used
+    mags = np.array([0, 0.9, 0.5, 0.2, -1])
+    times = np.array([datetime(2000, 1, 1),
+                      datetime(2000, 1, 2),
+                      datetime(2000, 1, 5),
+                      datetime(2000, 1, 4),
+                      datetime(2000, 1, 3)])
+
+    estimator.calculate(mags, mc=-1, delta_m=0.1, times=times)
+    assert (mags[estimator.idx] == np.array([0.9, 0.2, 0.5])).all()
+    assert (estimator.times == times[estimator.idx]).all()

seismostats/analysis/bvalue/utils.py

@@ -93,7 +93,7 @@ def find_next_larger(magnitudes: np.array,
     for ii in range(len(magnitudes) - 1):
         for jj in range(ii + 1, len(magnitudes)):
             mag_diff_loop = magnitudes[jj] - magnitudes[ii]
-            if mag_diff_loop > dmc - delta_m / 2:
+            if mag_diff_loop >= dmc - delta_m / 2:
                 idx_next_larger[ii] = jj
                 break
     return idx_next_larger.astype(int)