diff --git a/CHANGELOG.md b/CHANGELOG.md index feee011..d647265 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,9 @@ ## 0.10.0 (unreleased) +### Features + +- Add support for showing stereo balance as bars (#475) + ## 0.9.1 ### Major Changes diff --git a/corrscope/channel.py b/corrscope/channel.py index 0eab572..5a068c4 100644 --- a/corrscope/channel.py +++ b/corrscope/channel.py @@ -15,7 +15,7 @@ from corrscope.config import ( ) from corrscope.triggers import MainTriggerConfig from corrscope.util import coalesce -from corrscope.wave import Wave, FlattenOrStr +from corrscope.wave import Wave, Flatten, FlattenOrStr if TYPE_CHECKING: from corrscope.corrscope import Config @@ -42,6 +42,7 @@ class ChannelConfig(DumpableAttrs): line_color: Optional[str] = None color_by_pitch: Optional[bool] = None + stereo_bars: Optional[bool] = None # region Legacy Fields trigger_width_ratio = Alias("trigger_width") @@ -87,6 +88,7 @@ class Channel: self.trigger_wave = wave.with_flatten(tflat, return_channels=False) self.render_wave = wave.with_flatten(rflat, return_channels=True) + self.stereo_wave = wave.with_flatten(Flatten.Stereo, return_channels=True) # `subsampling` increases `stride` and decreases `nsamp`. # `width` increases `stride` without changing `nsamp`. @@ -138,3 +140,8 @@ class Channel: return self.render_wave.get_around( trigger_sample, self._render_samp, self.render_stride ) + + def get_render_stereo(self, trigger_sample: int): + return self.stereo_wave.get_around( + trigger_sample, self._render_samp, self.render_stride + ) diff --git a/corrscope/corrscope.py b/corrscope/corrscope.py index d7ced20..a5bd8f2 100644 --- a/corrscope/corrscope.py +++ b/corrscope/corrscope.py @@ -13,6 +13,7 @@ from threading import Thread from typing import Iterator, Optional, List, Callable, Dict, Union, Any import attr +import numpy as np from corrscope import outputs as outputs_ from corrscope.channel import Channel, ChannelConfig, DefaultLabel @@ -24,6 +25,7 @@ from corrscope.renderer import ( RendererConfig, RendererParams, RenderInput, + StereoLevels, ) from corrscope.settings.global_prefs import Parallelism from corrscope.triggers import ( @@ -232,6 +234,16 @@ def worker_render_frame( prev = t2 +def calc_stereo_levels(data: np.ndarray) -> StereoLevels: + def amplitude(chan_data: np.ndarray) -> float: + sq = chan_data * chan_data + mean = np.add.reduce(sq) / len(sq) + root = np.sqrt(mean) + return root + + return (amplitude(data.T[0]), amplitude(data.T[1])) + + class CorrScope: def __init__(self, cfg: Config, arg: Arguments): """cfg is mutated! @@ -389,7 +401,9 @@ class CorrScope: render_inputs = [] trigger_samples = [] # Get render-data from each wave. - for render_wave, channel in zip(self.render_waves, self.channels): + for wave_idx, (render_wave, channel) in enumerate( + zip(self.render_waves, self.channels) + ): sample = round(render_wave.smp_s * time_seconds) # Get trigger. @@ -408,7 +422,25 @@ class CorrScope: if should_render: trigger_samples.append(trigger_sample) data = channel.get_render_around(trigger_sample) - render_inputs.append(RenderInput(data, freq_estimate)) + + stereo_data = None + if ( + renderer.is_stereo_bars(wave_idx) + and not channel.stereo_wave.is_mono + ): + stereo_data = data + # If stereo track is flattened to mono for rendering, + # get raw stereo data. + if stereo_data.shape[1] == 1: + stereo_data = channel.get_render_stereo(trigger_sample) + + stereo_levels = None + if stereo_data is not None: + stereo_levels = calc_stereo_levels(stereo_data) + + render_inputs.append( + RenderInput(data, stereo_levels, freq_estimate) + ) if not should_render: continue @@ -489,7 +521,9 @@ class CorrScope: render_inputs = [] trigger_samples = [] # Get render-data from each wave. - for render_wave, channel in zip(self.render_waves, self.channels): + for wave_idx, (render_wave, channel) in enumerate( + zip(self.render_waves, self.channels) + ): sample = round(render_wave.smp_s * time_seconds) # Get trigger. @@ -508,7 +542,27 @@ class CorrScope: if should_render: trigger_samples.append(trigger_sample) data = channel.get_render_around(trigger_sample) - render_inputs.append(RenderInput(data, freq_estimate)) + + stereo_data = None + if ( + renderer.is_stereo_bars(wave_idx) + and not channel.stereo_wave.is_mono + ): + stereo_data = data + # If stereo track is flattened to mono for rendering, + # get raw stereo data. + if stereo_data.shape[1] == 1: + stereo_data = channel.get_render_stereo( + trigger_sample + ) + + stereo_levels = None + if stereo_data is not None: + stereo_levels = calc_stereo_levels(stereo_data) + + render_inputs.append( + RenderInput(data, stereo_levels, freq_estimate) + ) if not should_render: continue diff --git a/corrscope/renderer.py b/corrscope/renderer.py index 350d5eb..9cea559 100644 --- a/corrscope/renderer.py +++ b/corrscope/renderer.py @@ -73,6 +73,7 @@ import matplotlib.image import matplotlib.patheffects from matplotlib.backends.backend_agg import FigureCanvasAgg from matplotlib.figure import Figure +from matplotlib.patches import Rectangle if TYPE_CHECKING: from matplotlib.artist import Artist @@ -185,6 +186,9 @@ class RendererConfig( v_midline: bool = False h_midline: bool = False + global_stereo_bars: bool = False + stereo_bar_color: str = "#88ffff" + # Label settings label_font: Font = attr.ib(factory=Font) @@ -251,6 +255,10 @@ def freq_to_color(cmap, freq: Optional[float], fallback_color: str) -> str: class LineParam: color: str color_by_pitch: bool + stereo_bars: bool + + +StereoLevels = Tuple[float, float] @attr.dataclass @@ -258,6 +266,7 @@ class RenderInput: # Should Renderer store a Wave and take an int? # Or take an array on each frame? data: np.ndarray + stereo_levels: Optional[StereoLevels] freq_estimate: Optional[float] @staticmethod @@ -266,7 +275,7 @@ class RenderInput: Stable function to construct a RenderInput given only a data array. Used mainly for tests. """ - return RenderInput(data, None) + return RenderInput(data, None, None) @staticmethod def wrap_datas(datas: List[np.ndarray]) -> List["RenderInput"]: @@ -419,6 +428,7 @@ class _RendererBase(ABC): LineParam( color=coalesce(ccfg.line_color, cfg.global_line_color), color_by_pitch=coalesce(ccfg.color_by_pitch, cfg.global_color_by_pitch), + stereo_bars=coalesce(ccfg.stereo_bars, cfg.global_stereo_bars), ) for ccfg in channel_cfgs ] @@ -434,6 +444,9 @@ class _RendererBase(ABC): else: self.render_strides = [1] * self.nplots + def is_stereo_bars(self, wave_idx: int): + return self._line_params[wave_idx].stereo_bars + # Instance functionality @abstractmethod @@ -501,6 +514,22 @@ def px_from_points(pt: Point) -> Pixel: return pt * PIXELS_PER_PT +@attr.dataclass(cmp=False) +class StereoBar: + rect: Rectangle + x_center: float + x_range: float + + def set_range(self, left: float, right: float): + left = -left + + x = self.x_center + left * self.x_range + width = (right - left) * self.x_range + + self.rect.set_x(x) + self.rect.set_width(width) + + class AbstractMatplotlibRenderer(_RendererBase, ABC): """Matplotlib renderer which can use any backend (agg, mplcairo). To pick a backend, subclass and set _canvas_type at the class level. @@ -543,6 +572,9 @@ class AbstractMatplotlibRenderer(_RendererBase, ABC): # [wave][chan] Line2D _wave_chan_to_line: "Optional[List[List[Line2D]]]" = None + # Only for stereo channels, if stereo bars are enabled. + _wave_to_stereo_bar: "List[Optional[StereoBar]]" + def _setup_axes(self, wave_nchans: List[int]) -> None: """ Creates a flat array of Matplotlib Axes, with the new layout. @@ -766,7 +798,7 @@ class AbstractMatplotlibRenderer(_RendererBase, ABC): return ax # Protected API - def __add_lines_stereo(self, dummy_datas: List[np.ndarray]): + def __add_lines_stereo(self, inputs: List[RenderInput]): cfg = self.cfg strides = self.render_strides @@ -775,7 +807,11 @@ class AbstractMatplotlibRenderer(_RendererBase, ABC): # Foreach wave, plot dummy data. lines2d = [] - for wave_idx, wave_data in enumerate(dummy_datas): + wave_to_stereo_bar = [] + for wave_idx, input in enumerate(inputs): + wave_data = input.data + line_params = self._line_params[wave_idx] + # [nsamp][nchan] Amplitude wave_zeros = np.zeros_like(wave_data) @@ -783,11 +819,11 @@ class AbstractMatplotlibRenderer(_RendererBase, ABC): wave_lines = [] xs = calc_xs(len(wave_zeros), strides[wave_idx]) + line_color = line_params.color # Foreach chan for chan_idx, chan_zeros in enumerate(wave_zeros.T): ax = chan_to_axes[chan_idx] - line_color = self._line_params[wave_idx].color chan_line: Line2D = ax.plot( xs, chan_zeros, color=line_color, linewidth=line_width @@ -809,7 +845,34 @@ class AbstractMatplotlibRenderer(_RendererBase, ABC): lines2d.append(wave_lines) self._artists.extend(wave_lines) + # Add stereo bars if enabled and track is stereo. + if input.stereo_levels: + assert self._line_params[wave_idx].stereo_bars + ax = self._wave_to_mono_axes[wave_idx] + + viewport_stride = self.render_strides[wave_idx] * cfg.viewport_width + x_center = calc_center(viewport_stride) + + xlim = ax.get_xlim() + x_range = (xlim[1] - xlim[0]) / 2 + + y_bottom = ax.get_ylim()[0] + + h = abs(y_bottom) / 16 + stereo_rect = Rectangle((x_center, y_bottom - h), 0, 2 * h) + stereo_rect.set_color(cfg.stereo_bar_color) + stereo_rect.set_linewidth(0) + ax.add_patch(stereo_rect) + + stereo_bar = StereoBar(stereo_rect, x_center, x_range) + + wave_to_stereo_bar.append(stereo_bar) + self._artists.append(stereo_rect) + else: + wave_to_stereo_bar.append(None) + self._wave_chan_to_line = lines2d + self._wave_to_stereo_bar = wave_to_stereo_bar def _update_lines_stereo(self, inputs: List[RenderInput]) -> None: """ @@ -817,8 +880,7 @@ class AbstractMatplotlibRenderer(_RendererBase, ABC): - inputs[wave] = ndarray, [samp][chan] = f32 """ if self._wave_chan_to_line is None: - datas = [input.data for input in inputs] - self.__add_lines_stereo(datas) + self.__add_lines_stereo(inputs) lines2d = self._wave_chan_to_line nplots = len(lines2d) @@ -854,6 +916,14 @@ class AbstractMatplotlibRenderer(_RendererBase, ABC): if color_by_pitch: chan_line.set_color(color) + stereo_bar = self._wave_to_stereo_bar[wave_idx] + stereo_levels = inputs[wave_idx].stereo_levels + assert bool(stereo_bar) == bool( + stereo_levels + ), f"wave {wave_idx}: plot={stereo_bar} != values={stereo_levels}" + if stereo_bar: + stereo_bar.set_range(*stereo_levels) + def _add_xy_line_mono( self, name: str,