plot.py

#!/usr/bin/env python

import argparse
import json

import datetime as dt
import matplotlib.dates as mdates
import matplotlib.pyplot as plt
import pandas as pd

from matplotlib.ticker import EngFormatter, PercentFormatter


def plotter(ax, data, params):
    defaults = {
            'linewidth': 0.5,
        }
    params = defaults | params
    out, = ax.plot(data, **params)
    return out


def stepper(ax, data, params):
    defaults = {
            'linewidth': 0.5,
        }
    params = defaults | params
    out, = ax.step(data.index, data.values, where='post', **params)
    return out


def scatter(ax, data, params):
    defaults = {
           's': 0.1,
           'linewidths': 0.5,
        }
    params = defaults | params
    out = ax.scatter(data.index, data.values, **params)
    return out


def read_rtcp(file, basetime):
    df = pd.read_csv(
            file,
            index_col=0,
            names=['time', 'rate'],
            header=None,
            usecols=[0, 1],
        )
    if not basetime:
        basetime = df.index[0]

    df.index = pd.to_datetime(df.index - basetime, unit='ms')
    df['rate'] = df['rate'].apply(lambda x: x * 8)
    df = df.resample('1s').sum()
    return df


def read_rtp(file, basetime):
    df = pd.read_csv(
            file,
            index_col=0,
            names=['time', 'rate'],
            header=None,
            usecols=[0, 6]
        )

    if not basetime:
        basetime = df.index[0]

    df.index = pd.to_datetime(df.index - basetime, unit='ms')
    df['rate'] = df['rate'].apply(lambda x: x * 8)
    df = df.resample('1s').sum()
    return df


def read_capacity(file, basetime):
    df = pd.read_csv(
            file,
            index_col=0,
            names=['time', 'bandwidth'],
            header=None,
            usecols=[0, 1],
        )
    if not basetime:
        basetime = df.index[0]

    df.index = pd.to_datetime(df.index - basetime, unit='ms')
    return df


def read_cc_qdelay(file, basetime):
    df = pd.read_csv(
            file,
            index_col=0,
            names=['time', 'queue delay'],
            header=None,
            usecols=[0, 2]
        )

    if not basetime:
        basetime = df.index[0]

    df.index = pd.to_datetime(df.index - basetime, unit='ms')
    return df


def read_cc_target_rate(file, basetime):
    df = pd.read_csv(
            file,
            index_col=0,
            names=['time', 'target bitrate'],
            header=None,
            usecols=[0, 1]
        )

    if not basetime:
        basetime = df.index[0]

    df.index = pd.to_datetime(df.index - basetime, unit='ms')
    return df


def read_rtp_loss(send_file, receive_file, basetime):
    df_send = pd.read_csv(
            send_file,
            index_col=1,
            names=['time_send', 'nr'],
            header=None,
            usecols=[0, 8],
        )
    df_receive = pd.read_csv(
            receive_file,
            index_col=1,
            names=['time_receive', 'nr'],
            header=None,
            usecols=[0, 8],
        )

    if not basetime:
        basetime = df_send.index[0]

    df_all = df_send.merge(df_receive, on=['nr'], how='left', indicator=True)
    df_all.index = pd.to_datetime(df_all['time_send'] - basetime, unit='ms')
    df_all['lost'] = df_all['_merge'] == 'left_only'
    df_all = df_all.resample('1s').agg({'time_send': 'count', 'lost': 'sum'})
    df_all['loss_rate'] = df_all['lost'] / df_all['time_send']

    df = df_all.drop('time_send', axis=1)
    df = df.drop('lost', axis=1)

    return df


def read_rtp_latency(send_file, receive_file, basetime):
    df_send = pd.read_csv(
            send_file,
            index_col=1,
            names=['time_send', 'nr'],
            header=None,
            usecols=[0, 8],
        )
    df_receive = pd.read_csv(
            receive_file,
            index_col=1,
            names=['time_receive', 'nr'],
            header=None,
            usecols=[0, 8],
        )

    if not basetime:
        basetime = df_send.index[0]

    df = df_send.merge(df_receive, on='nr')
    df['diff'] = (df['time_receive'] - df['time_send']) / 1000.0
    df['time'] = pd.to_datetime(df['time_send'] - basetime, unit='ms')
    df = df.drop(['time_send', 'time_receive', 'time'], axis=1)

    return df


def main():
    parser = argparse.ArgumentParser(
            formatter_class=argparse.ArgumentDefaultsHelpFormatter
        )
    parser.add_argument('--name', default='', help='Plot name')
    parser.add_argument('--config', default='', help='Use a config file to'
                        ' read metadata such as the basetime')
    parser.add_argument('--capacity', default='', help='Link capacity log')
    parser.add_argument('--rtp-sent', help='Senderside RTP logfile to include'
                        ' in plot')
    parser.add_argument('--rtp-received', help='Receiverside RTP logfile to'
                        ' include in plot')
    parser.add_argument('--rtcp-sent', help='Senderside RTCP logfile to'
                        ' include in plot')
    parser.add_argument('--rtcp-received', help='Receiverside RTCP logfile to'
                        ' include in plot')
    parser.add_argument('--cc', help='CC file to include in plot')
    parser.add_argument('--loss', nargs=2, help='plot loss between an RTP sent'
                        ' log file and an RTP received log file',
                        metavar=('sent_rtp.log', 'received_rtp.log'))
    parser.add_argument('--latency', nargs=2, help='RTP latency plot between'
                        ' an RTP sent log file and an RTP received log file',
                        metavar=('sent_rtp.log', 'received_rtp.log'))
    parser.add_argument('--qdelay', help='SCReAM queue delay')
    parser.add_argument('-o', '--output', required=True, help='output file')
    parser.add_argument('-b', '--basetime', type=int, help='basetime to use in'
                        ' plots, if not given, will be inferred from the input'
                        ' data using the first row')

    args = parser.parse_args()

    if args.config and not args.basetime:
        with open(args.config) as f:
            d = json.load(f)
            args.basetime = d['basetime']

    print(args)

    fig, ax = plt.subplots(figsize=(8, 2), dpi=400)

    labels = []
    if args.capacity:
        data = read_capacity(
                args.capacity,
                args.basetime,
            )
        labels.append(stepper(ax, data, {
            'label': 'Link Capacity',
            }))
    if args.rtp_sent:
        data = read_rtp(
                args.rtp_sent,
                args.basetime,
            )
        labels.append(plotter(ax, data, {
            'label': 'Sent RTP',
        }))

    if args.rtp_received:
        data = read_rtp(
                args.rtp_received,
                args.basetime,
            )
        labels.append(plotter(ax, data, {
            'label': 'Received RTP',
        }))

    if args.rtcp_sent:
        data = read_rtcp(
                args.rtcp_sent,
                args.basetime,
            )
        labels.append(plotter(ax, data, {
            'label': 'Sent RTCP',
        }))

    if args.rtcp_received:
        data = read_rtcp(
                args.rtcp_received,
                args.basetime,
            )
        labels.append(plotter(ax, data, {
            'label': 'Received RTCP',
        }))

    if args.cc:
        data = read_cc_target_rate(
                args.cc,
                args.basetime,
            )
        labels.append(plotter(ax, data, {
            'label': 'CC Target Bitrate',
        }))

    if args.loss:
        data = read_rtp_loss(
                args.loss[0],
                args.loss[1],
                args.basetime,
            )
        labels.append(plotter(ax, data, {
            'label': 'RTP loss',
        }))

    if args.latency:
        data = read_rtp_latency(
                args.latency[0],
                args.latency[1],
                args.basetime,
            )
        labels.append(scatter(ax, data, {
            'label': 'RTP latency',
            }))

    if args.qdelay:
        data = read_cc_qdelay(
                args.qdelay,
                args.basetime,
            )
        labels.append(plotter(ax, data, {
            'label': 'SCReAM Queue Delay',
        }))

    if args.cc or args.rtp_sent or args.rtp_received:
        ax.set_xlabel('Time')
        ax.set_ylabel('Rate')
        ax.set_title(args.name)
        ax.xaxis.set_major_formatter(mdates.DateFormatter("%M:%S"))
        ax.yaxis.set_major_formatter(EngFormatter(unit='bit/s'))
        ax.set_xlim([dt.datetime(1970, 1, 1), dt.datetime(1970, 1, 1,
                    minute=2)])

    if args.loss:
        ax.set_xlabel('Time')
        ax.set_ylabel('Packet Loss')
        ax.set_title(args.name)
        ax.xaxis.set_major_formatter(mdates.DateFormatter("%M:%S"))
        ax.yaxis.set_major_formatter(PercentFormatter(xmax=1.0))
        ax.set_xlim([dt.datetime(1970, 1, 1), dt.datetime(1970, 1, 1,
                    minute=2)])

    # lgd = ax.legend(handles=labels, loc='upper right', bbox_to_anchor=(1,
    #                 1), ncol=2)
    lgd = ax.legend(handles=labels)
    fig.tight_layout()
    fig.savefig(args.output, bbox_extra_artists=(lgd,), bbox_inches='tight')
    # fig.savefig(args.output)


if __name__ == "__main__":
    main()