From 538570d82fd3d8918ccb97568b68851dc4dd9339 Mon Sep 17 00:00:00 2001
From: Josh Bailey <josh@vandervecken.com>
Date: Thu, 6 Jun 2024 22:31:26 +0000
Subject: [PATCH] Add WaterfallPlot class.

---
 .../gamutrfwaterfall/waterfall.py             | 21 ++++---------
 .../gamutrfwaterfall/waterfall_plot.py        | 31 +++++++++++++++----
 2 files changed, 31 insertions(+), 21 deletions(-)

diff --git a/gamutrfwaterfall/gamutrfwaterfall/waterfall.py b/gamutrfwaterfall/gamutrfwaterfall/waterfall.py
index ea33bb9b..1e57eeef 100644
--- a/gamutrfwaterfall/gamutrfwaterfall/waterfall.py
+++ b/gamutrfwaterfall/gamutrfwaterfall/waterfall.py
@@ -4,8 +4,6 @@
 import tempfile
 import time
 import warnings
-import matplotlib
-from matplotlib import style as matplotlibstyle
 
 from gamutrflib.peak_finder import get_peak_finder
 from gamutrflib.zmqbucket import ZmqReceiver, parse_scanners, frame_resample
@@ -16,11 +14,8 @@
     write_scanner_args,
 )
 from gamutrfwaterfall.waterfall_plot import (
-    reset_fig,
-    init_fig,
-    update_fig,
     make_config,
-    WaterfallState,
+    WaterfallPlot,
 )
 
 warnings.filterwarnings(action="ignore", message="Mean of empty slice")
@@ -61,8 +56,6 @@ def serve_waterfall(
     config_vars,
     config_vars_path,
 ):
-    matplotlibstyle.use("fast")
-
     global need_reset_fig
     need_reset_fig = True
     global running
@@ -124,9 +117,7 @@ def sig_handler(_sig=None, _frame=None):
                 config.fft_len,
                 config.freq_resolution,
             )
-
-            state = WaterfallState(config, base_save_path, peak_finder)
-            matplotlib.use(config.engine)
+            plot = WaterfallPlot(config, base_save_path, peak_finder)
             results = [
                 (scan_configs, frame_resample(scan_df, config.freq_resolution * 1e6))
             ]
@@ -138,11 +129,11 @@ def sig_handler(_sig=None, _frame=None):
             need_reconfig = False
             need_init = True
         if need_init:
-            init_fig(config, state, onresize)
+            plot.init_fig(onresize)
             need_init = False
             need_reset_fig = True
         if need_reset_fig:
-            reset_fig(config, state)
+            plot.reset_fig()
             need_reset_fig = False
         last_gap = time.time()
         while True:
@@ -191,8 +182,8 @@ def sig_handler(_sig=None, _frame=None):
         if need_reconfig:
             continue
         if results:
-            update_fig(config, state, results)
-            if config.batch and state.counter % config.reclose_interval == 0:
+            plot.update_fig(results)
+            if plot.need_init():
                 need_init = True
             results = []
         else:
diff --git a/gamutrfwaterfall/gamutrfwaterfall/waterfall_plot.py b/gamutrfwaterfall/gamutrfwaterfall/waterfall_plot.py
index 80572078..d26489a2 100644
--- a/gamutrfwaterfall/gamutrfwaterfall/waterfall_plot.py
+++ b/gamutrfwaterfall/gamutrfwaterfall/waterfall_plot.py
@@ -5,18 +5,15 @@
 import os
 import shutil
 import time
-import warnings
 from pathlib import Path
-import matplotlib.pyplot as plt
 import numpy as np
+import matplotlib
+from matplotlib import style
 from matplotlib.collections import LineCollection
+import matplotlib.pyplot as plt
 from matplotlib.ticker import MultipleLocator, AutoMinorLocator
 from scipy.ndimage import gaussian_filter
 
-warnings.filterwarnings(action="ignore", message="Mean of empty slice")
-warnings.filterwarnings(action="ignore", message="All-NaN slice encountered")
-warnings.filterwarnings(action="ignore", message="Degrees of freedom <= 0 for slice.")
-
 
 def safe_savefig(path):
     basename = os.path.basename(path)
@@ -798,3 +795,25 @@ def make_config(
         save_time,
     )
     return config
+
+
+class WaterfallPlot:
+    def __init__(self, config, base_save_path, peak_finder):
+        self.config = config
+        self.state = WaterfallState(config, base_save_path, peak_finder)
+        matplotlib.use(self.config.engine)
+        style.use("fast")
+
+    def update_fig(self, results):
+        update_fig(self.config, self.state, results)
+
+    def reset_fig(self):
+        reset_fig(self.config, self.state)
+
+    def init_fig(self, onresize):
+        init_fig(self.config, self.state, onresize)
+
+    def need_init(self):
+        return (
+            self.config.batch and self.state.counter % self.config.reclose_interval == 0
+        )