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 + )