plot.py

#!/usr/bin/env python3

# Extremely simple signal plotting tool

import os
import sys
import binascii
import random
from scipy.signal import butter,lfilter,freqz
import scipy.signal
from numpy import *
import time
import getopt
import matplotlib as mpl
import sparkgap.filemanager

TRIGGERS = 0

lastTime = 0.0
lastX = 0

def onclick(event):
  global lastTime, lastX, OFFSET
  t = time.time()
  if t - lastTime < 0.200:
    print("debounce - nope")
    return
  elif event.xdata is None:
    print("skip - event.xdata (click on graph) is none")
    return
  else:
    lastTime = t
    if lastX == 0:
      lastX = int(event.xdata)
      lastX += OFFSET
      print("MARK: %d" % lastX)
    else:
      localX = int(event.xdata)
      fromX = min(lastX,localX)
      toX = max(lastX,localX)
      dist = toX - fromX
      fromX += OFFSET
      toX += OFFSET
      print("FROM %d TO %d DIST %d" % (fromX,toX,dist))
      lastX = localX


def butter_bandpass(lowcut,highcut,fs,order=5):
  nyq = 0.5 * fs
  low = lowcut / nyq
  high = highcut / nyq
  b,a = butter(order, [low, high], btype = 'band')
  return b,a

def butter_bandpass_filter(data,lowcut,highcut,fs,order=5):
  b,a = butter_bandpass(lowcut,highcut,fs,order=order)
  y = lfilter(b,a,data)
  return y

def butter_lowpass(cutoff, fs, order=5):
  nyq = 0.5 * fs
  normal_cutoff = cutoff / nyq
  b, a = butter(order, normal_cutoff, btype='low', analog=False)
  return b, a

def butter_lowpass_filter(data, cutoff, fs, order=5):
  b, a = butter_lowpass(cutoff, fs, order=order)
  y = lfilter(b, a, data)
  return y

def getTraceConfig(r_str):
  r = []
  if "," in r_str:
    tokens = r_str.split(",")
    print(tokens)
  else:
    tokens = [r_str]
  for t in tokens:
    if "-" in t:
      (t1,t2) = t.split("-")
      r += list(range(int(t1),int(t2)))
    else:
      r += [int(t)]
  return r

OFFSET = 0
COUNT = 0
RULER = []
# NUM_TRACES = 1
TRACES = []
GAIN_FACTOR =  31622.0
ADDITIONAL_FILES = []

BANDPASS_LOWCUT=1000000
BANDPASS_HIGHCUT=5000000
BANDPASS_SR = 249999999
BANDPASS_ORDER=1
BANDPASS_EN = False

LOWPASS_CUTOFF = 10000
LOWPASS_SR = 40000000
LOWPASS_ORDER = 5
LOWPASS_EN = False

FFT_BASEFREQ = 40000000
FFT_EN = False

SPECGRAM_EN = False
SPECGRAM_SR = 0

PLOT_SHOWN = False # dirty hack

TITLE = "Single Trace Plot"
XAXIS = "Sample Count"
YAXIS = "Power"

def configure_fft(arg):
  global FFT_BASEFREQ,FFT_EN,TITLE,XAXIS
  FFT_BASEFREQ = float(arg)
  TITLE = "FFT Plot (%d Hz Sample Rate)" % FFT_BASEFREQ
  XAXIS = "Frequency"
  FFT_EN = True

def configure_specgram(arg):
  global SPECGRAM_EN, TITLE, XAXIS, YAXIS, SPECGRAM_SR
  TITLE = "Spectogram View"
  print("SPECGRAM_EN = True")
  SPECGRAM_EN = True
  SPECGRAM_SR = float(arg)

def configure_lowpass(in_str):
  global LOWPASS_CUTOFF, LOWPASS_SR, LOWPASS_ORDER, LOWPASS_EN, TITLE
  try:
    (cutoff,samplerate,order) = in_str.split(",")
  except:
    print("syntax: -l 10000,40000000,5 (cutoff, samplerate, order)")
    sys.exit(0)
  LOWPASS_CUTOFF = float(cutoff)
  LOWPASS_SR = float(samplerate)
  LOWPASS_ORDER = int(order)
  LOWPASS_EN = True
  TITLE = "Low Pass (%d Hz SR, %d Hz Cutoff)" % (LOWPASS_SR,LOWPASS_CUTOFF)

def configure_bandpass(in_str):
  global BANDPASS_LOWCUT, BANDPASS_HIGHCUT,BANDPASS_ORDER, BANDPASS_SR, BANDPASS_EN, TITLE
  try:
    (lowcut,highcut,samplerate,order) = in_str.split(",")
  except:
    print("syntax -b (lowcut,highcut,sr,order)")
    sys.exit(0)
  BANDPASS_LOWCUT = float(lowcut)
  BANDPASS_HIGHCUT = float(highcut)
  BANDPASS_SR = float(samplerate)
  BANDPASS_ORDER = int(order)
  BANDPASS_EN = True
  TITLE = "Band pass (%d Hz to %d Hz, %d SR)" % (BANDPASS_LOWCUT,BANDPASS_HIGHCUT,BANDPASS_SR)

def usage():
  print(" plot.py : part of the fuckshitfuck toolkit")
  print("----------------------------------------------")
  print(" -h : prints this message")
  print(" -o : offset to start plotting samples from")
  print(" -n : number of samples from offset to plot")
  print(" -c : select traces")
  print(" -f : input npz file (can be multiple)")
  print(" -r : print vertical ruler at point (NOT IMPLEMENTED)")
  print(" -l [cutoff,samplerate,order] : lowpass mode - units in hz")
  print(" -b [lowcut,highcut,samplerate,order] : bandpass mode - units in hz")
  print(" -F [samplerate] : plot fft, base freq in hz")
  print(" -s [samplerate] : plot spectrogram")
  print(" -w [filename] : write output to file. suppresses window.")

mpl.rcParams['agg.path.chunksize'] = 10000 
CONFIG_WRITEFILE = None
SPECIAL_TEST = False

if __name__ == "__main__":
  opts, remainder = getopt.getopt(sys.argv[1:],"tb:s:hl:n:o:c:r:f:F:w:",["spectrogram=","help","lowpass=","samples=","offset=","count=","ruler=","file=","fft=","highlight=","bandpass=","test"])
  for opt,arg in opts:
    if opt in ("-h","--help"):
      usage()
      sys.exit(0)
    elif opt in ("-s","--spectrogram"):
      configure_specgram(arg)
    elif opt in ("-o","--offset"):
      OFFSET = int(float(arg))
    elif opt in ("-n","--samples"):
      COUNT = int(float(arg))
    elif opt in ("-c","--count"):
      TRACES = getTraceConfig(arg)
    elif opt in ("-f","--file"):
      ADDITIONAL_FILES.append(arg)
    elif opt in ("-l","--lowpass"):
      configure_lowpass(arg)
    elif opt in ("-b","--bandpass"):
      configure_bandpass(arg)
    elif opt in ("-t","--test"):
      print("SPECIAL TEST MODE")
      SPECIAL_TEST = True
    elif opt in ("-F","--fft"):
      configure_fft(arg)
    elif opt in ("-w"):
      CONFIG_WRITEFILE = arg
    elif opt in ("-r","--ruler"):
      RULER.append(int(float(arg)))
    else:
      print("Unknown argument: %s" % opt)
      sys.exit(0) 
  if CONFIG_WRITEFILE is not None:
    mpl.use("Agg")
  if LOWPASS_EN and BANDPASS_EN:
    print("You can't have both lowpass and bandpass filters (yet!)")
    sys.exit(0)
  import matplotlib.pyplot as plt
  if [FFT_EN, LOWPASS_EN, SPECGRAM_EN, BANDPASS_EN].count(True) > 1:
    print("You can only select one of -F (FFT), -l (LOWPASS) or -b (BANDPASS)")
    sys.exit(0)
  if SPECGRAM_EN == False:
    fig, ax1 = plt.subplots()
  if len(ADDITIONAL_FILES) != 1:
    print("TraceManager no longer supports multiple files by design. Try something else")
    sys.exit(0)
  for f in ADDITIONAL_FILES:
    if not os.path.isfile(f):
      print("Fatal: could not open %s" % f)
      sys.exit(0)
    tm = sparkgap.filemanager.TraceManager(f)
    if len(TRACES) == 0:
      print("You must specify at least one trace with -c")
      sys.exit(0)
    if len(TRACES) != 1:
      TITLE = "TRACE PLOT"
    if ((OFFSET != 0) and (COUNT == 0)):
      print("Fix: Assume you want to start at %d, finish at trace end" % OFFSET)
      d = tm.getSingleTrace(0)
      COUNT = len(d) - OFFSET
    for i in TRACES:
      if OFFSET == 0 and COUNT == 0:
        d = tm.getSingleTrace(i)
        di = tm.getSingleData(i)
        do = tm.getSingleDataOut(i)
        print(di)
        print(do)
      else:
        d = tm.getSingleTrace(i)[OFFSET:OFFSET+COUNT]
      # print("Continuing...")
      if LOWPASS_EN: # this code is disgusting but fuck you
        print("LOWPASS")
        d = tm.getSingleTrace(i)
        # d = df['traces'][i]
        if OFFSET == 0 and COUNT == 0:
          if SPECIAL_TEST:
            lowpassed_d = butter_lowpass_filter(d,LOWPASS_CUTOFF,LOWPASS_SR,LOWPASS_ORDER)
            std_dev = std(lowpassed_d)
            avg_dev = average(lowpassed_d)
            # peaks,_ = scipy.signal.find_peaks(lowpassed_d,prominence=[0,0.5],rel_height=0.9)
            plt.plot(lowpassed_d)
            # plt.plot(peaks,lowpassed_d[peaks],"x")
          else:
            plt.plot(butter_lowpass_filter(d,LOWPASS_CUTOFF,LOWPASS_SR,LOWPASS_ORDER))
        else:
          plt.plot(butter_lowpass_filter(d,LOWPASS_CUTOFF,LOWPASS_SR,LOWPASS_ORDER)[OFFSET:OFFSET+COUNT])
      elif BANDPASS_EN:
        print("BANDPASS")
        d = tm.getSingleTrace(i)
        # d = df['traces'][i]
        if OFFSET == 0 and COUNT == 0:
          plt.plot(butter_bandpass_filter(d,BANDPASS_LOWCUT,BANDPASS_HIGHCUT,BANDPASS_SR,BANDPASS_ORDER))
        else:
          plt.plot(butter_bandpass_filter(d,BANDPASS_LOWCUT,BANDPASS_HIGHCUT,BANDPASS_SR,BANDPASS_ORDER)[OFFSET:OFFSET+COUNT])
      elif FFT_EN:
        print("FFT")
        n = len(d)
        k = arange(n)
        T = n / FFT_BASEFREQ
        frq = k / T
        frq = frq[list(range(n // 2))]
        Y = fft.fft(d)/n
        Y = Y[list(range(n //  2))]
        plt.plot(frq,abs(Y),'r') 
      elif SPECGRAM_EN:
        print("Specgram enable")
        fig, (ax1, ax2) = plt.subplots(nrows=2)
        ax1.set_title("Power Trace")
        ax1.set_ylabel("Power")
        ax1.set_xlabel("Sample Count")
        ax1.plot(d)
        ax2.set_title("Spectogram")
        ax2.set_ylabel("Frequency Component")
        ax2.set_xlabel("Time")
        ax2.specgram(d * 100,NFFT=1024,Fs=SPECGRAM_SR,noverlap=900)
        ax1.margins(0)
        fig.canvas.manager.set_window_title("plot.py")
        if CONFIG_WRITEFILE is not None:
          print("Saving to %s..." % CONFIG_WRITEFILE)
          plt.savefig(CONFIG_WRITEFILE)
        else:
          plt.show()
          PLOT_SHOWN = True
      else:
        plt.plot(d)
  if PLOT_SHOWN is False:
    plt.title(TITLE)
    plt.ylabel(YAXIS)
    plt.xlabel(XAXIS)
    plt.grid()
    fig.canvas.manager.set_window_title("plot.py")
    if CONFIG_WRITEFILE is not None:
      print("Saving to %s..." % CONFIG_WRITEFILE)
      plt.savefig(CONFIG_WRITEFILE)
    else:
      print("Connecting Event - Focus the window to activate")
      fig.canvas.mpl_connect("button_press_event",onclick)
      plt.show()