path: root/bin/dlog-viewer

#!/usr/bin/env python3
# vim:tabstop=4:softtabstop=4:shiftwidth=4:textwidth=160:smarttab:expandtab

"""dlog-viewer - View and Convert Keysight .dlog Files

DESCRIPTION

dlog-viewer loads voltage, current, and/or power measurements from .dlog files
produced by devices such as the Keysight N6705B DC Power Analyzer.
Measurements can be exported to CSV or plotted on-screen.

This program is not affiliated with Keysight and has not been thoroughly
tested yet. Use at your own risk.

OPTIONS
"""

import argparse
import csv
import json
import matplotlib.pyplot as plt
from multiprocessing import Pool
import numpy as np
import os
import struct
import sys
import xml.etree.ElementTree as ET


def running_mean(x: np.ndarray, N: int) -> np.ndarray:
    """
    Compute `N` elements wide running average over `x`.

    :param x: 1-Dimensional NumPy array
    :param N: how many items to average. Should be even for optimal results.
    """

    # to ensure that output.shape == input.shape, we need to insert data
    # at the boundaries
    boundary_array = np.insert(x, 0, np.full((N // 2), x[0]))
    boundary_array = np.append(boundary_array, np.full((N // 2 + N % 2 - 1), x[-1]))

    return np.convolve(boundary_array, np.ones((N,)) / N, mode="valid")


def downsample(x: np.ndarray, N: int) -> np.ndarray:
    """Downsample `x` by factor `N`."""
    fill = x.shape[0] % N
    if fill:
        x = np.array(list(x) + [x[-1] for i in range(N - fill)])
    return x.reshape(-1, N).mean(axis=1)


class NpEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.ndarray):
            return obj.tolist()
        else:
            return super(NpEncoder, self).default(obj)


def PELT_get_changepoints(algo, penalty):
    res = (penalty, algo.predict(pen=penalty))
    return res


class PELT:
    def __init__(self, signal, num_samples=None):
        self.signal = signal
        self.model = "l1"
        self.jump = 1
        self.min_dist = 500

        if num_samples is not None:
            self.ds_factor = len(signal) // num_samples
        else:
            self.ds_factor = 1

        self.jump = self.ds_factor
        # if self.ds_factor > 1:
        #    self.signal = downsample(self.signal, self.ds_factor)
        # print(f"ds from {len(signal)} to {len(self.signal)}")

    def norm_signal(self, signal, scaler=25):
        max_val = max(np.abs(signal))
        normed_signal = np.zeros(shape=len(signal))
        for i, signal_i in enumerate(signal):
            normed_signal[i] = signal_i / max_val
            normed_signal[i] = normed_signal[i] * scaler
        return normed_signal

    def get_changepoints(self):
        # imported here as ruptures is only used for changepoint detection
        import ruptures

        algo = ruptures.Pelt(
            model=self.model, jump=self.jump, min_size=self.min_dist
        ).fit(self.norm_signal(self.signal))
        queue = list()
        for i in range(0, 100):
            queue.append((algo, i))
        with Pool() as pool:
            changepoints = pool.starmap(PELT_get_changepoints, queue)
        changepoints_by_penalty = dict()
        for res in changepoints:
            changepoints_by_penalty[res[0]] = res[1]
        num_changepoints = list()
        for i in range(0, 100):
            num_changepoints.append(len(changepoints_by_penalty[i]))

        # Find plateau
        start_index = -1
        end_index = -1
        longest_start = -1
        longest_end = -1
        prev_val = -1
        for i, num_bkpts in enumerate(num_changepoints):
            if num_bkpts != prev_val:
                end_index = i - 1
                if end_index - start_index > longest_end - longest_start:
                    # currently found sequence is the longest found yet
                    longest_start = start_index
                    longest_end = end_index
                start_index = i
            if i == len(num_changepoints) - 1:
                # end sequence with last value
                end_index = i
                # # since it is not guaranteed that this is the end of the plateau, assume the mid
                # # of the plateau was hit.
                # size = end_index - start_index
                # end_index = end_index + size
                # However this is not the clean solution. Better if search interval is widened
                # with range_min and range_max
                if end_index - start_index > longest_end - longest_start:
                    # last found sequence is the longest found yet
                    longest_start = start_index
                    longest_end = end_index
                start_index = i
            prev_val = num_bkpts
        middle_of_plateau = longest_start + (longest_start - longest_start) // 2
        changepoints = np.array(changepoints_by_penalty[middle_of_plateau])
        return changepoints


class DLogChannel:
    def __init__(self, desc_tuple):
        self.slot = desc_tuple[0]
        self.smu = desc_tuple[1]
        self.unit = desc_tuple[2]
        self.data = None

    def __repr__(self):
        return f"""<DLogChannel(slot={self.slot}, smu="{self.smu}", unit="{self.unit}", data={self.data})>"""


class DLog:
    def __init__(self, filename, skip=None, limit=None):
        self.skip_duration = skip
        self.limit_duration = limit
        self.load_dlog(filename)

    def load_dlog(self, filename):
        lines = []
        line = ""

        with open(filename, "rb") as f:
            if ".xz" in filename:
                import lzma

                f = lzma.open(f)

            while line != "</dlog>\n":
                line = f.readline().decode()
                lines.append(line)
            xml_header = "".join(lines)
            raw_header = f.read(8)
            data_offset = f.tell()
            raw_data = f.read()

        xml_header = xml_header.replace("1ua>", "X1ua>")
        xml_header = xml_header.replace("2ua>", "X2ua>")
        dlog = ET.fromstring(xml_header)
        channels = []
        for channel in dlog.findall("channel"):
            channel_id = int(channel.get("id"))
            sense_curr = channel.find("sense_curr").text
            sense_volt = channel.find("sense_volt").text
            model = channel.find("ident").find("model").text
            if sense_volt == "1":
                channels.append((channel_id, model, "V"))
            if sense_curr == "1":
                channels.append((channel_id, model, "A"))

        num_channels = len(channels)

        self.channels = list(map(DLogChannel, channels))
        self.interval = float(dlog.find("frame").find("tint").text)
        self.planned_duration = int(dlog.find("frame").find("time").text)
        self.observed_duration = self.interval * int(len(raw_data) / (4 * num_channels))

        self.timestamps = np.linspace(
            0, self.observed_duration, num=int(len(raw_data) / (4 * num_channels))
        )

        if (
            self.skip_duration is not None
            and self.observed_duration >= self.skip_duration
        ):
            start_offset = 0
            for i, ts in enumerate(self.timestamps):
                if ts >= self.skip_duration:
                    start_offset = i
                    break
            self.timestamps = self.timestamps[start_offset:]
            raw_data = raw_data[start_offset * 4 * num_channels :]

        if (
            self.limit_duration is not None
            and self.observed_duration > self.limit_duration
        ):
            stop_offset = len(self.timestamps) - 1
            for i, ts in enumerate(self.timestamps):
                if ts > self.limit_duration:
                    stop_offset = i
                    break
            self.timestamps = self.timestamps[:stop_offset]
            self.observed_duration = self.timestamps[-1]
            raw_data = raw_data[: stop_offset * 4 * num_channels]

        self.data = np.ndarray(
            shape=(num_channels, int(len(raw_data) / (4 * num_channels))),
            dtype=np.float32,
        )

        iterator = struct.iter_unpack(">f", raw_data)
        channel_offset = 0
        measurement_offset = 0
        for value in iterator:
            if value[0] < -1e6 or value[0] > 1e6:
                print(
                    f"Invalid data value {value[0]} at channel {channel_offset}, measurement {measurement_offset}. Replacing with 0."
                )
                self.data[channel_offset, measurement_offset] = 0
            else:
                self.data[channel_offset, measurement_offset] = value[0]
            if channel_offset + 1 == num_channels:
                channel_offset = 0
                measurement_offset += 1
            else:
                channel_offset += 1

        # An SMU has four slots
        self.slots = [dict(), dict(), dict(), dict()]

        for i, channel in enumerate(self.channels):
            channel.data = self.data[i]
            self.slots[channel.slot - 1][channel.unit] = channel

    def slot_has_data(self, slot):
        return len(self.slots[slot - 1]) > 0

    def slot_has_power(self, slot):
        slot_data = self.slots[slot - 1]
        if "W" in slot_data:
            return True
        if "V" in slot_data and "A" in slot_data:
            return True
        return False

    def observed_duration_equals_expectation(self):
        return int(self.observed_duration) == self.planned_duration

    def all_data_slots_have_power(self):
        for slot in range(4):
            if self.slot_has_data(slot) and not self.slot_has_power(slot):
                return False
        return True


def detect_changepoints(dlog, num_samples):
    ret_slots = list()
    for slot in dlog.slots:
        ret_slot = dict()
        for unit in slot:
            pelt = PELT(running_mean(slot[unit].data, 10), num_samples=num_samples)
            changepoints = pelt.get_changepoints()
            prev = 0
            ret_slot[unit] = list()
            for cp in changepoints:
                cp = cp - 1
                ret_slot[unit].append(
                    {
                        "interval": [dlog.timestamps[prev], dlog.timestamps[cp]],
                        "mean": np.mean(slot[unit].data[prev:cp]),
                    }
                )
                prev = cp
        ret_slots.append(ret_slot)
    return ret_slots


def print_stats(dlog):
    if dlog.observed_duration_equals_expectation():
        print(
            "Measurement duration: {:d} seconds at {:f} µs per sample".format(
                dlog.planned_duration, dlog.interval * 1_000_000
            )
        )
    else:
        print(
            "Measurement duration: {:f} of {:d} seconds at {:f} µs per sample".format(
                dlog.observed_duration, dlog.planned_duration, dlog.interval * 1_000_000
            )
        )

    for channel in dlog.channels:
        min_data = np.min(channel.data)
        max_data = np.max(channel.data)
        mean_data = np.mean(channel.data)
        if channel.unit == "V":
            precision = 3
        else:
            precision = 6
        print(f"Slot {channel.slot} ({channel.smu}):")
        print(f"    Min  {min_data:.{precision}f} {channel.unit}")
        print(f"    Mean {mean_data:.{precision}f} {channel.unit}")
        print(f"    Max  {max_data:.{precision}f} {channel.unit}")
        print()


def show_power_plot(dlog):

    handles = list()

    for slot in dlog.slots:
        if "W" in slot:
            (handle,) = plt.plot(
                dlog.timestamps, slot["W"].data, "b-", label="P", markersize=1
            )
            handles.append(handle)
            (handle,) = plt.plot(
                dlog.timestamps,
                running_mean(slot["W"].data, 10),
                "r-",
                label="mean(P, 10)",
                markersize=1,
            )
            handles.append(handle)
        elif "V" in slot and "A" in slot:
            (handle,) = plt.plot(
                dlog.timestamps,
                slot["V"].data * slot["A"].data,
                "b-",
                label="P = U * I",
                markersize=1,
            )
            handles.append(handle)
            (handle,) = plt.plot(
                dlog.timestamps,
                running_mean(slot["V"].data * slot["A"].data, 10),
                "r-",
                label="mean(P, 10)",
                markersize=1,
            )
            handles.append(handle)

    plt.legend(handles=handles)
    plt.xlabel("Time [s]")
    plt.ylabel("Power [W]")
    plt.grid(True)
    plt.show()


def show_unit_plot(dlog, metric):

    handles = list()

    if metric == "U":
        unit = "V"
    elif metric == "I":
        unit = "A"
    elif metric == "P":
        unit = "W"

    for slot in dlog.slots:
        if unit in slot:
            channel = slot[unit]
            (handle,) = plt.plot(
                dlog.timestamps,
                slot[unit].data,
                "b-",
                label=f"slot {channel.slot} ({channel.smu})",
                markersize=1,
            )
            handles.append(handle)
            (handle,) = plt.plot(
                dlog.timestamps,
                running_mean(slot[unit].data, 10),
                "r-",
                label=f"slot {channel.slot} mean",
                markersize=1,
            )
            handles.append(handle)

    plt.legend(handles=handles)
    plt.xlabel("Time [s]")
    if unit == "V":
        plt.ylabel("Voltage [V]")
    elif unit == "A":
        plt.ylabel("Current [A]")
    elif unit == "W":
        plt.ylabel("Power [W]")
    else:
        plt.ylabel(f"??? [{unit}]")
    plt.grid(True)
    plt.show()


def show_raw_plot(dlog):
    handles = list()

    for channel in dlog.channels:
        label = f"{channel.slot} / {channel.smu} {channel.unit}"
        (handle,) = plt.plot(
            dlog.timestamps, channel.data, "b-", label=label, markersize=1
        )
        handles.append(handle)
        (handle,) = plt.plot(
            dlog.timestamps,
            running_mean(channel.data, 10),
            "r-",
            label=f"mean({label}, 10)",
            markersize=1,
        )
        handles.append(handle)

    plt.legend(handles=handles)
    plt.xlabel("Time [s]")
    plt.ylabel("Voltage [V] / Current [A] / Power [W]")
    plt.grid(True)
    plt.show()


def export_csv(dlog, filename):
    cols, rows = dlog.data.shape
    with open(filename, "w", newline="") as f:
        writer = csv.writer(f)
        channel_header = list(
            map(lambda x: f"{x.smu} in slot {x.slot} [{x.unit}]", dlog.channels)
        )
        writer.writerow(["Timestamp [s]"] + channel_header)
        for row in range(rows):
            writer.writerow([dlog.timestamps[row]] + list(dlog.data[:, row]))


def export_json(dlog, filename, extra_data=dict()):
    json_channels = list()
    for channel in dlog.channels:
        json_channels.append({"smu": channel.smu, "unit": channel.unit})
    json_data = {"channels": json_channels}
    json_data.update(extra_data)
    with open(filename, "w") as f:
        json.dump(json_data, f, cls=NpEncoder)


def main():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter, description=__doc__
    )
    parser.add_argument(
        "--csv-export",
        metavar="FILENAME",
        type=str,
        help="Export measurements to CSV file",
    )
    parser.add_argument(
        "--json-export",
        metavar="FILENAME",
        type=str,
        help="Export analysis results (e.g. changepoints) to JSON file",
    )
    parser.add_argument(
        "--skip",
        metavar="N",
        type=int,
        default=0,
        help="Skip the first N seconds of data. This is useful to avoid startup code influencing the results of a long-running measurement",
    )
    parser.add_argument(
        "--limit",
        type=float,
        metavar="N",
        help="Limit analysis to the first N seconds of data",
    )
    parser.add_argument(
        "--pelt", metavar="JUMP", type=int, help="Perform changepoint detection"
    )
    parser.add_argument(
        "--plot",
        choices=["U", "I", "P", "all"],
        help="Plot voltage/current/power over time",
    )
    parser.add_argument(
        "--stat", help="Print mean voltage, current, and power", action="store_true"
    )
    parser.add_argument(
        "dlog_file", type=str, help="Input filename in Keysight dlog format"
    )

    args = parser.parse_args()

    dlog = DLog(args.dlog_file, args.skip, args.limit)

    if args.stat:
        print_stats(dlog)

    if args.csv_export:
        export_csv(dlog, args.csv_export)

    if args.pelt:
        changepoints = detect_changepoints(dlog, num_samples=args.pelt)
        for i, slot in enumerate(changepoints):
            for unit in slot:
                num_changepoints = len(slot[unit])
                print(
                    f"Found {num_changepoints} changepoints for {dlog.slots[i][unit].smu} {unit}"
                )

    if args.json_export:
        extra_data = dict()
        if args.pelt:
            extra_data["changepoints"] = changepoints
        export_json(dlog, args.json_export, extra_data)

    if args.plot:
        if args.plot == "P" and dlog.all_data_slots_have_power():
            show_power_plot(dlog)
        elif args.plot == "all":
            show_raw_plot(dlog)
        else:
            show_unit_plot(dlog, args.plot)


if __name__ == "__main__":
    main()