can force the direction of the step function

Author: millen1m

eqsig/functions.py

@@ -604,7 +604,7 @@ def generate_smooth_fa_spectrum(smooth_fa_frequencies, fa_frequencies, fa_spectr
     return smooth_fa_spectrum
 
 
-def calc_step_fn_vals_error(values, pow=1):
+def calc_step_fn_vals_error(values, pow=1, dir=None):
     """
     Calculates the error function generated by fitting a step function
     to the values
@@ -618,6 +618,11 @@ def calc_step_fn_vals_error(values, pow=1):
     values: array_like
     pow: int
         The power that the error should be raised to
+    dir: str
+        Desired direction of the step function
+        if 'down', then all upward steps are set to 10x maximum error
+        if 'up', then all downward steps are set to 10x maximum error
+        else, no modification to error
 
     Returns
     -------
@@ -638,6 +643,12 @@ def calc_step_fn_vals_error(values, pow=1):
     err = np.ones_like(values)
     err[:-1] = err_post[1:] + err_pre[:-1]
     err[-1] = np.sum(np.abs(values - np.mean(values)) ** pow)
+    if dir == 'down':  # if step has to be downward, then increase error for upward steps
+        max_err = np.max(err)
+        err = np.where(pre_mean < post_mean, max_err * 10, err)
+    if dir == 'up':
+        max_err = np.max(err)
+        err = np.where(pre_mean > post_mean, max_err * 10, err)
     return err