updated stockwell plot to be properly aligned with axes.

eqsig/__about__.py

@@ -1,4 +1,4 @@
 __project__ = "eqsig"
 __author__ = "Maxim Millen"
-__version__ = "0.5.35"
+__version__ = "0.5.36"
 __license__ = "MIT"

eqsig/stockwell.py

@@ -32,55 +32,18 @@ def plot_stock(splot, asig, norm_x=False, norm_all=False, times=(None, None), fr
     if norm_x:
         b = b / np.max(b, axis=0)
 
-    max_x = len(b[0])
-    max_y = len(b)
-    max_time = asig.time[-1]
     max_freq = 1 / asig.dt / 2
-    if times[0] is not None:
-        start_t = times[0] / max_time * max_x
-    else:
-        start_t = 0
-    if times[1] is not None:
-        end_t = times[1] / max_time * max_x
-    else:
-        end_t = max_x
-    if freqs[0] is not None:
-        start_f = freqs[0] / max_freq * max_y
-    else:
-        start_f = 0
-    if freqs[1] is not None:
-        end_f = freqs[1] / max_freq * max_y
-    else:
-        end_f = max_y
-    #
-    # print(max_x, max_y)
-    # max_x = 2000
-    extent = (-0.5 + start_t, end_t - 0.5, end_f - 0.5, start_f - 0.5)
-    # extent = (-0.5, max_x - 0.5, max_y - 0.5,  - 0.5)
+
+    extent = (0, asig.time[-1], 0, max_freq)
     splot.imshow(b, aspect='auto', extent=extent)
-    xticks = ticker.FuncFormatter(lambda x, pos: '{0:g}'.format(x * asig.dt / 2))
-    s10 = np.ceil(asig.time[-1] / 10)
-    if s10 < 4:
-        step = 5
-    else:
-        step = 10
-    x_majors = np.arange(0, 10 * s10, step) * 2 / asig.dt
-    splot.xaxis.set_major_formatter(xticks)
-    splot.set_xticks(x_majors)
-
-    # y ticks
-    yticks = ticker.FuncFormatter(lambda y, pos: '{0:g}'.format((len(asig.swtf) / 2 - y) / (len(asig.values) * asig.dt)))
-
-    splot.yaxis.set_major_formatter(yticks)
-
-    s10 = np.ceil(max_freq / 10)
-    if s10 < 4:
-        step = 5
-    else:
-        step = 10
-    y_majors = np.arange(0, 10 * s10, step)
-    print(y_majors)
-    splot.set_yticks(np.flipud(y_majors) * len(asig.values) * asig.dt)
+
+
+def get_stockwell_freqs(asig):
+    return (len(asig.swtf) / 2 - np.arange(0, int(len(asig.swtf) / 2))) / (len(asig.values) * asig.dt)
+
+
+def get_stockwell_times(asig):
+    return np.arange(0, len(asig.swtf[0])) * asig.dt / 2
 
 
 def plot_fas_at_time(splot, asig, time):
@@ -160,5 +123,6 @@ def itransform(stock):
     asig = load_signal(conftest.TEST_DATA_DIR + "test_motion_dt0p01.txt", astype="acc_sig")
     asig = eqsig.interp_to_approx_dt(asig, 0.05)
     bf, sps = plt.subplots()
-    plot_stock(sps, asig, norm_x=True, times=(10, 30), freqs=(0, 7))
+    plot_stock(sps, asig, norm_x=True)
+    # , times = (10, 30), freqs = (0, 7)
     plt.show()