plot-spectrum.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Plot recorded data from Polbox/Polarilog saved as compressed Numpy arrays (npz/hdf5) as spectrograms.
@author: Martin Šlapák
"""

import argparse
from datetime import timedelta, datetime
from pathlib import Path
import h5py
import matplotlib.pyplot as plt
import numpy as np
import zoneinfo
from polarilog.sopmetrics import compute_angle, compute_distance, compute_angular_velocity

import senzotools as st


def get_from_time(fn, attrs, desired_tz=None):
    '''Try to estimate the time of start of recording'''
    fromtime = None
    tz = None
    try:
        tz = zoneinfo.ZoneInfo(attrs.get('timezone_name', None))
    except (zoneinfo.ZoneInfoNotFoundError, TypeError):
        print(f'Unknown or ambiguous timezone id ({attrs.get("timezone_name", None)}).')

    # at first use file metadata
    if 'epoch_start' in attrs:
        fromtime = datetime.fromtimestamp(int(attrs['epoch_start']), tz=tz)
        # fromtime = fromtime.replace(tzinfo=tz)
        print(f'Starting time = {fromtime} (from metadata, tz={tz})')
    else:
        # if no suitable attributes, try to guess from filename
        try:
            fn = Path(fn)
            fromtime = datetime.strptime(fn.stem[3:], '%Y-%m-%d_%H-%M-%S')  # take fn from 3rd place skipping common "plg" prefix
            print(f'Starting time = {fromtime} (parsed from filename)')
        except Exception:
            print(f'Cannot parse time from: "{fn.stem[3:]}"')

    # offset the time
    if fromtime is not None and desired_tz is not None:
        try:
            dtz = zoneinfo.ZoneInfo(desired_tz)
            fromtime = fromtime.astimezone(dtz)
            print(f'Starting time = {fromtime} (localized to {dtz})')
        except zoneinfo._common.ZoneInfoNotFoundError:
            print(f'Cannot handle desired target timezone ({desired_tz}), no time shifting."')
            pass

    if fromtime is not None:
        print(f'tzinfo={fromtime.tzinfo}')
    return fromtime


parser = argparse.ArgumentParser(description='Plots spectrum from recorded data (npz/hdf5).', epilog='')
parser.add_argument('-i', action='store', help='A filename to read from.', required=True)
parser.add_argument('-k', action='store', help='A key in HDF5 file. If not provided, the 1st will be used.')
parser.add_argument('-r', action='store', default='0:1', help='Range of colums in numpy slice notation e.g.: 1:2, 3:4 [Default = %(default)s]')
parser.add_argument('-f', action='store', default=20000, help='Frequency in Hz [Default = %(default)s Hz]')
parser.add_argument('-clip', action='store', default='0:', help='Take only this range (in samples) [Default = %(default)s]')
parser.add_argument('-transpose', action='store_true', help='If set, the data channels will be transposed.')
parser.add_argument('-s', action='store_true', help='If set, save rendered image.')
parser.add_argument('-sp', action='store_true', help='If set, save "pure" image - without axes, scales, titles only content (but with metadata).')
parser.add_argument('-dpi', action='store', default=300, help='DPI of saved figure [Default = %(default)s]')
parser.add_argument('-d', action='store_false', help='If set, supress the displaying of rendered image.')
parser.add_argument('-uf', action='store', default=10000, help='Upper frequency limit [Default = %(default)s Hz]')
parser.add_argument('-nperseg', action='store', default=2048, help='FFT window size [Default = %(default)s samples]')
parser.add_argument('-vmin', action='store', default=None, help='Min value in dB for colorbar scale. If ommited it will be automatically computated.')
parser.add_argument('-vmax', action='store', default=None, help='Max value in dB for colorbar scale. If ommited it will be automatically computated.')
parser.add_argument('-tz', action='store', default=None, help='Apply this TZ to time labels [Default = %(default)s]')
parser.add_argument('-t', action='store', default='Frequency domain', help='Graph title')
parser.add_argument('-xlab', action='store', default='Time', help='Graph label for X axis')
parser.add_argument('-ylab', action='store', default='Frequency [Hz]', help='Graph label for Y axis')
parser.add_argument('-sum', action='store_true', help='If set, all channels are summed as their 2nd powers and squered.')
parser.add_argument('-m', '--metric', action='store', default=None, choices=['angle', 'distance', 'angular_velocity'], help='Compute spectrum from derived metric instead of raw channels [Default = %(default)s]')
parser.add_argument('-step', action='store', default=1, type=int, help='Step size for metric computation [Default = %(default)s]')

args = parser.parse_args()
print(f'Options: {args}')

uf = int(args.uf)
fs = int(args.f)
nperseg = int(args.nperseg)
vmin = None if args.vmin is None else float(args.vmin)
vmax = None if args.vmax is None else float(args.vmax)
fn = Path(args.i)
bn = fn.name

attrs = dict()
try:
    data = st.read_binary(fn)
except ValueError:
    # not numpy's pickled data
    with h5py.File(args.i, "r") as hf:
        print(f"Keys: {hf.keys()}")
        key = None
        if args.k is not None:
            print(f"HDF5 key: {args.k}")
            if args.k in hf.keys():
                key = args.k
            else:
                key = None
        if key is None:
            key = list(hf.keys())[0]
        data = np.array(hf[key])
        attrs = dict(hf[key].attrs)

if args.transpose:
    data = data.T

try:
    c1, c2 = args.clip.split(':')
    c1 = int(c1)
    if c2 == '':
        c2 = data.shape[0] - 1
    c2 = int(c2)
except ValueError:
    c1 = 0
    c2 = data.shape[0] - 1
    print(f'Invalid clip range: {args.clip}')
print(f'Using clip range: {c1}:{c2}')

print(f'Original file = {fn}')
print(f'Original data.shape = {data.shape}')
data = data[c1:c2, :]
print(f'Clipped data.shape = {data.shape}')

# get start time based on metadata or filename
fromtime = get_from_time(bn, attrs, args.tz)
print(f'fromtime={fromtime}')

if c1 > 0:
    fromtime = fromtime + timedelta(seconds=c1 / fs)
    print(f'Clipped range time start = {fromtime}')

# Process sum or metric BEFORE column slicing (they use all channels)
sum_computed = False
metric_computed = False
if args.sum:
    # vse
    #data = np.sqrt(data[:, 0]**2 + data[:, 1]**2 + data[:, 2]**2 + data[:, 3]**2)
    # bez U0
    data = np.sqrt(data[:, 1]**2 + data[:, 2]**2 + data[:, 3]**2)
    data = data[:, np.newaxis]
    sum_computed = True
    print(f'Summed channels 1-3, data.shape = {data.shape}')

# Compute metric if selected (uses channels 0-4, ignores -r argument)
if args.metric is not None:
    if args.metric == 'angle':
        data = compute_angle(data, step=args.step)
        metric_computed = True
    elif args.metric == 'distance':
        data = compute_distance(data, step=args.step)
        metric_computed = True
    elif args.metric == 'angular_velocity':
        data = compute_angular_velocity(data, fs=fs, step=args.step)
        metric_computed = True
    # Sanitize metric data: replace zeros with small epsilon (problematic for log10 in spectrum)
    data = np.maximum(data, 1e-10)
    data = data[:, np.newaxis]
    print(f'Computed metric "{args.metric}" (step={args.step}), data.shape = {data.shape}')

# Column slicing only if neither sum nor metric is selected
if not sum_computed and not metric_computed:
    print(f'Selecting column(s): {args.r}')
    col_start, col_end = list(map(lambda x: int(x), args.r.split(':')))
    if len(data.shape) > 1:
        # data = np.ndarray.flatten(data[:, col_start:col_end])
        data = data[:, col_start:col_end]
    print(f'Column-sliced data.shape = {data.shape}')
else:
    col_start = 0  # Default for sum/metric

tl = data.shape[0] / fs
print(f'Time length = {tl} s')
print(f'Vmin: {vmin}, Vmax: {vmax}')

for channel in range(0, data.shape[1]):
    if metric_computed:
        print(f'Plotting metric "{args.metric}" (step={args.step})...')
        channel_label = args.metric
    else:
        print(f'Plotting channel {col_start + channel}...')
        channel_label = str(col_start + channel)
    fig = plt.figure(figsize=(10, 5))
    ax = plt.axes()
    ax, f, t, Sxx = st.plot_spectrum(data[:, channel], fs=fs, nperseg=nperseg, cmap='viridis', force_simple=True,
                                     fromtime=fromtime, timelen=tl, timestep=(tl / 20.0),
                                     upperbound=uf, vmin=vmin, vmax=vmax,
                                     xlabel=args.xlab, ylabel=args.ylab, title=args.t, ax=ax)

    ax.grid(linewidth=0.2)
    print(f'Sxx.min: {10 * np.log10(np.min(Sxx)):.2f}, Sxx.max: {10 * np.log10(np.max(Sxx)):.2f}')

    if args.d:
        try:
            mng = plt.get_current_fig_manager()
            mng.resize(*mng.window.maxsize())
        except AttributeError:
            pass
        plt.show()

    meta = {
        'spect-src': Path(fn).name,
        'spect-clip': f'{c1}:{c2}',
        'spect-channel': channel_label,
        'spect-nperseg': str(nperseg),
        'spect-sampling-freq': str(fs),
        'spect-upper-freq': str(uf),
        'spect-vmin': str(vmin),
        'spect-vmax': str(vmax),
        'spect-title': args.t,
        'Software': 'plot-spectrum.py',
    }
    
    appendix = '_sum' if args.sum else f'_ch{channel}'
    if args.metric is not None:
        appendix = f'_{args.metric}'
        if args.step > 1:
            appendix = f'_{args.metric}_step{args.step}'
    
    if args.sp:
        outfn = fn.parent / f"{fn.stem}{appendix}_pure.png"
        plt.imsave(fname=outfn, arr=10 * np.log10(Sxx[::-1, :]), cmap='viridis', format='png', metadata=meta)

    if args.s:
        outfn = fn.parent / f"{fn.stem}{appendix}.png"
        fig.savefig(outfn, dpi=int(args.dpi), bbox_inches='tight', pad_inches=0.1, metadata=meta)

    del Sxx, fig, ax, t, f