Extend plotter

Author: nocturnalastro

src/vse_sync_pp/plot.py

@@ -6,51 +6,62 @@
 
 import numpy as np
 import matplotlib.pyplot as plt
+from collections import namedtuple
 
 from .common import open_input
 
 from .parsers import PARSERS
 
 
+Axis = namedtuple("Axis", ["desc", "attr", "scale", "scale_kwargs"], defaults=[None, None, None, None])
+TIMESERIES = Axis("Time", "timestamp")
+
+
 class Plotter():
     """Rudimentary plotter of data values against timestamp"""
-    def __init__(self, y_name, y_desc=None):
-        self._x_name = 'timestamp'
-        self._y_name = y_name
-        if y_desc is None:
-            self._y_desc = y_name
-        else:
-            self._y_desc = y_desc
+    def __init__(self, x, y):
+        self._x = x
+        self._y = y
         self._x_data = []
         self._y_data = []
 
+    @staticmethod
+    def _extract_attr(axis, data):
+        return getattr(data, axis.attr)
+
+    def _set_yscale(self, ax):
+        if self._y.scale is not None:
+            ax.set_yscale(self._y.scale, **(self._y.scale_kwargs or {}))
+        elif any((abs(v) > 10 for v in self._y_data)):
+            ax.set_yscale("symlog", linthresh=10)
+
     def append(self, data):
         """Append x and y data points extracted from `data`"""
-        x_val = getattr(data, self._x_name)
-        self._x_data.append(x_val)
-        y_val = getattr(data, self._y_name)
-        self._y_data.append(y_val)
+        self._x_data.append(self._extract_attr(self._x, data))
+        self._y_data.append(self._extract_attr(self._y, data))
 
     def plot(self, filename):
         """Plot data to `filename`"""
         fig, (ax1, ax2) = plt.subplots(2, constrained_layout=True)
         fig.set_size_inches(10, 8)
         ax1.axhline(0, color='black')
-        if any((abs(v) > 10 for v in self._y_data)):
-            ax1.set_yscale('symlog', linthresh=10)
+        self._set_yscale(ax1)
+        if self._x.scale is not None:
+            ax1.set_xscale(self._x.scale, **(self._x.scale_kwargs or {}))
         ax1.plot(self._x_data, self._y_data, '.')
         ax1.grid()
-        ax1.set_title(f'{self._x_name} vs {self._y_desc}')
+        ax1.set_title(f'{self._x.desc} vs {self._y.desc}')
         counts, bins = np.histogram(
             np.array(self._y_data, dtype=float),
             bins='scott',
         )
         ax2.hist(bins[:-1], bins, weights=counts)
-        ax2.set_yscale('symlog', linthresh=10)
-        ax2.set_title(f'Histogram of {self._y_desc}')
+        self._set_yscale(ax2)
+        if self._x.scale is not None:
+            ax2.set_xscale(self._x.scale, **(self._x.scale_kwargs or {}))
+        ax2.set_title(f'Histogram of {self._y.desc}')
         plt.savefig(filename)
 
-
 def main():
     """Plot data parsed from log messages from a single source.
 
@@ -75,7 +86,7 @@ def main():
     )
     args = aparser.parse_args()
     parser = PARSERS[args.parser]()
-    plotter = Plotter(parser.y_name)
+    plotter = Plotter(TIMESERIES, Axis(parser.y_name, parser.y_name))
     with open_input(args.input) as fid:
         method = parser.canonical if args.canonical else parser.parse
         for parsed in method(fid):