Transition. Erf fitting. Return fit params in a more convenient way.

cfatools/tsanalyser/transition.py

@@ -224,7 +224,7 @@ def get_transition_duration_erf(df: pd.DataFrame,
                                 first_datetime: datetime.datetime = None,
                                 last_datetime: datetime.datetime = None,
                                 col_value_name: str = "value",
-                                debug_plot: bool = False) -> float:
+                                debug_plot: bool = False) -> tuple:
     df = df.copy()
 
     # Restrict the data to the transition area
@@ -250,7 +250,10 @@ def get_transition_duration_erf(df: pd.DataFrame,
     sigma_0 = 10
 
     params, extras = curve_fit(erfunc, xdata, ydata, p0=[C1_0, C2_0, z0_0, sigma_0])
-    print(params)
+    amplitude = params[0]
+    baseline = params[1]
+    center = params[2]
+    sigma = params[3]
 
     if debug_plot:
         plt.title('Erf fit, sigma=' + "{:.2f}".format(params[3]) + "s")
@@ -258,15 +261,17 @@ def get_transition_duration_erf(df: pd.DataFrame,
         plt.ylabel(col_value_name)
         plt.plot(xdata, ydata, 'k', label="Raw")
         plt.plot(xdata, erfunc(xdata, *params), label="Erf fit")
-        plt.axhline(y=params[1], color="yellow", linestyle="--", label="Step's plateau")
-        plt.axhline(y=params[1] + params[0], color="yellow", linestyle="--")
-        plt.axvline(x=params[2], color="red", linestyle="--", label="Step's 'center'")
-        plt.axvline(x=params[2]-params[3], color="green", linestyle="--", label="Step's width")
-        plt.axvline(x=params[2]+params[3], color="green", linestyle="--")
+        plt.axhline(y=baseline, color="yellow", linestyle="--", label="Step's plateau")
+        plt.axhline(y=baseline + amplitude, color="yellow", linestyle="--")
+        plt.axvline(x=center, color="red", linestyle="--", label="Step's 'center'")
+        plt.axvline(x=center-sigma, color="green", linestyle="--", label="Step's width")
+        plt.axvline(x=center+sigma, color="green", linestyle="--")
         plt.legend()
         plt.show()
 
-    return params
+    center_timestamp = df.index[0] + pd.Timedelta(center, unit='s')
+
+    return baseline, amplitude, center_timestamp, sigma
 
 
 # def __detect_transition__(df: pd.DataFrame, col_value_name: str = "value") -> pd.DataFrame: