plot-timedomain.py

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

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

# import matplotlib as mpl
# mpl.rcParams['path.simplify_threshold'] = 0.1


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']))
        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


if __name__ == "__main__":
    parser = argparse.ArgumentParser(prog='plot-polarilog-data', description='Plot recorded data from Polarilog saved as compressed Numpy arrays (npz/hdf5).', epilog='')
    parser.add_argument('-i', action='store', required=True, help='Filename to read recorded data from.')
    parser.add_argument('-r', action='store', default='0:4', help='Range of colums in numpy slice notation e.g.: 3:5')
    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('-s', action='store_true', help='If set, save rendered image.')
    parser.add_argument('-d', action='store_false', help='If set, supress the displaying of rendered image.')
    parser.add_argument('-df', action='store', default=None, help='Decimation factor [Default = %(default)s]')
    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='Changes of polarization', help='Graph title')
    parser.add_argument('-xlab', action='store', default='Time', help='Graph label for X axis')
    parser.add_argument('-ylab', action='store', default=None, help='Graph label for Y axis')
    parser.add_argument('-m', '--metric', action='store', default=None, choices=['angle', 'distance', 'angular_velocity'], help='Compute and plot derived metric [Default = %(default)s]')
    parser.add_argument('-step', action='store', default=1, type=int, help='Step size for metric computation (number of samples between compared points) [Default = %(default)s]')

    args = parser.parse_args()
    print(args)
    fn = args.i
    bn = os.path.basename(fn)
    fs = int(args.f)

    # range of channels (typically 0-4)
    r = list(map(lambda x: int(x), args.r.split(':')))

    # load data
    attrs = dict()
    a = np.array([])
    try:
        with np.load(fn) as fp:
            first_key = fp.files[0]
            a = fp[first_key]
    except ValueError:
        # not numpy's pickled data
        with h5py.File(fn, "r") as hf:
            first_key = list(hf.keys())[0]
            a = np.array(hf[first_key])
            attrs = dict(hf[first_key].attrs)

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

    print(f'Recorded data has a following shape: {a.shape}')
    a = a[c1:c2, :]
    print(f'Clipped data has a following shape: {a.shape}')
    print(a)
    ypadding = 1.
    if (len(a.shape) > 1):
        ymin = np.min(a[1:, r[0]:r[1]])  # 1:4
        ymax = np.max(a[1:, r[0]:r[1]])  # 1:4
    else:
        ymin = np.min(a[1:])  # 1:4
        ymax = np.max(a[1:])  # 1:4

    print(f'ymin= {ymin} ymax={ymax}')
    ymin = ymin - abs(ymin) * ypadding
    ymax = ymax + abs(ymax) * ypadding
    print(f'ymin= {ymin} ymax={ymax} (padded)')

    # 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}')

    # decimation factor
    decf = 1
    if args.df is not None:
        decf = int(args.df)
        fs = fs / decf

    # Determine if we need to plot a metric
    plot_metric = args.metric is not None

    if plot_metric:
        # Create figure with two subplots sharing X axis
        fig, (ax_main, ax_metric) = plt.subplots(2, 1, figsize=(15, 10), sharex=True)
        
        # Apply decimation to data before plotting
        if (len(a.shape) > 1):
            a_decimated = a[::decf, :]
        else:
            a_decimated = a[::decf]
        
        # Plot main data
        if (len(a_decimated.shape) > 1):
            ax_main.plot(a_decimated[:, r[0]:r[1]])
            # Use "Sx" (Stokes) if calibrated, otherwise "CHx"
            if attrs.get('calibration_status') == 'calibrated':
                channel_labels = [f'S{i}' for i in range(r[0], r[1])]
            else:
                channel_labels = [f'CH{i}' for i in range(r[0], r[1])]
            ax_main.legend(channel_labels, loc='upper right')
        else:
            ax_main.plot(a_decimated)
        
        ax_main.grid()
        ax_main.set_xlim((0, a_decimated.shape[0]))
        ax_main.set_ylim((ymin, ymax))
        
        # Compute the metric
        if args.metric == 'angle':
            metric_data = compute_angle(a, step=args.step)
            metric_label = f'Angle [rad] (step={args.step})'
        elif args.metric == 'distance':
            metric_data = compute_distance(a, step=args.step)
            metric_label = f'Distance (step={args.step})'
        elif args.metric == 'angular_velocity':
            metric_data = compute_angular_velocity(a, fs=fs * decf, step=args.step)
            metric_label = f'Angular velocity [rad/s] (step={args.step})'
        
        # Apply decimation to metric data
        metric_data_decimated = metric_data[::decf]
        
        # Plot metric on the bottom subplot
        ax_metric.plot(metric_data_decimated, color='red', linewidth=1.5)
        ax_metric.grid()
        ax_metric.set_ylabel(metric_label)
        ax_metric.set_xlabel(args.xlab if fromtime is None else args.xlab + f' [tz={fromtime.tzinfo}]')
    else:
        # Original behavior - single plot
        fig = plt.figure(figsize=(15, 7))
        ax = plt.axes(xlim=(0, a.shape[0]), ylim=(ymin, ymax))
        ax.grid()

        # decimation factor
        decf = 1
        if args.df is not None:
            decf = int(args.df)
            fs = fs / decf
            ax.set_xlim((0, a.shape[0] / decf))

        # plot it
        if (len(a.shape) > 1):
            a = a[::decf, :]
            plt.plot(a[:, r[0]:r[1]])
            # Use "Sx" (Stokes) if calibrated, otherwise "CHx"
            if attrs.get('calibration_status') == 'calibrated':
                channel_labels = [f'S{i}' for i in range(r[0], r[1])]
            else:
                channel_labels = [f'CH{i}' for i in range(r[0], r[1])]
            plt.legend(channel_labels, loc='upper right')
        else:
            a = a[::decf]
            plt.plot(a)

    # xtics and its labels
    if plot_metric:
        # Use ax_main for setting xticks when metric is plotted
        timelen = a_decimated.shape[0] / fs
        timestep = timelen / 10
        xtickscnt = int(timelen / timestep) + 1  # +1 is for 0
        ax_main.set_xticks(np.linspace(0, a_decimated.shape[0], num=xtickscnt))
        if fromtime is not None:
            unit = 's'
            if timestep < 10.0:
                unit = 'ms'
                timelen = int(timelen * 1000000)
                timestep = int(timestep * 1000000)
            xlabels = []
            for i in range(0, int(timelen) + int(timestep), int(timestep)):
                if unit == 's':
                    when = fromtime + timedelta(seconds=i)
                    xlabels.append(str(when.strftime("%H:%M:%S")))
                else:
                    when = fromtime + timedelta(microseconds=i)
                    xlabels.append(str(when.strftime("%H:%M:%S")))
            ax_main.set_xticklabels(xlabels[0:xtickscnt], rotation=90)
        else:
            # how many decimal places should be rounded
            exp = int(f'{timelen:e}'.split('e')[-1])
            if exp > 0:
                decimal_places = 0
            else:
                decimal_places = abs(exp) + 1

            labels = [f'{i:.{decimal_places}f}' for i in np.arange(0, timelen + timestep, timestep)]
            ax_main.set_xticklabels(labels[0:xtickscnt], rotation=90)
        
        del a, a_decimated, metric_data, metric_data_decimated  # free mem
        
        # default
        if args.xlab == 'Time' and fromtime is not None:
            args.xlab += f' [tz={fromtime.tzinfo}]'

        # If calibrated and user not provided any specific Y-label, show Stokes [-], Channel voltage [V] otherwise
        if args.ylab is None:
            if attrs.get('calibration_status') == 'calibrated':
                args.ylab = 'Stokes [-]'
            else:
                args.ylab = 'Channel voltage [V]'

        ax_main.set_ylabel(args.ylab)
        ax_main.set_title(args.t)
    else:
        # Original behavior
        timelen = a.shape[0] / fs
        timestep = timelen / 10
        xtickscnt = int(timelen / timestep) + 1  # +1 is for 0
        ax.set_xticks(np.linspace(0, a.shape[0], num=xtickscnt))
        if fromtime is not None:
            unit = 's'
            if timestep < 10.0:
                unit = 'ms'
                timelen = int(timelen * 1000000)
                timestep = int(timestep * 1000000)
            xlabels = []
            for i in range(0, int(timelen) + int(timestep), int(timestep)):
                if unit == 's':
                    when = fromtime + timedelta(seconds=i)
                    xlabels.append(str(when.strftime("%H:%M:%S")))
                else:
                    when = fromtime + timedelta(microseconds=i)
                    xlabels.append(str(when.strftime("%H:%M:%S")))
            ax.set_xticklabels(xlabels[0:xtickscnt], rotation=90)
        else:
            # how many decimal places should be rounded
            exp = int(f'{timelen:e}'.split('e')[-1])
            if exp > 0:
                decimal_places = 0
            else:
                decimal_places = abs(exp) + 1

            labels = [f'{i:.{decimal_places}f}' for i in np.arange(0, timelen + timestep, timestep)]
            ax.set_xticklabels(labels[0:xtickscnt], rotation=90)

        del a  # free mem

        # default
        if args.xlab == 'Time' and fromtime is not None:
            args.xlab += f' [tz={fromtime.tzinfo}]'

        # If calibrated and user not provided any specific Y-label, show Stokes [-], Channel voltage [V] otherwise
        if args.ylab is None:
            if attrs.get('calibration_status') == 'calibrated':
                args.ylab = 'Stokes [-]'
            else:
                args.ylab = 'Channel voltage [V]'

        plt.ylabel(args.ylab)
        plt.xlabel(args.xlab)
        plt.title(args.t)

    if args.d:
        plt.show()

    if args.s:
        bn = os.path.basename(fn)
        bn_no_ext = ''.join(bn.split('.')[:-1])
        path = os.path.dirname(fn)
        if plot_metric:
            if args.step > 1:
                fig.savefig(os.path.join(path, f'{bn_no_ext}_time-domain_{args.metric}_step{args.step}.png'), dpi=300, bbox_inches='tight', pad_inches=0.1)
            else:
                fig.savefig(os.path.join(path, f'{bn_no_ext}_time-domain_{args.metric}.png'), dpi=300, bbox_inches='tight', pad_inches=0.1)
        else:
            fig.savefig(os.path.join(path, f'{bn_no_ext}_time-domain.png'), dpi=300, bbox_inches='tight', pad_inches=0.1)