visualize.py

#!/usr/bin/env python

import matplotlib.pyplot as plt
import argparse
import numpy as np
import sys
from pathlib import Path
print("external:api=5")

WORKDIR = Path(__file__).parent

TYPES = {
        "int8": np.int8,
        "int16": np.int16,
        "int32": np.int32,
        "int64": np.int64,
        "uint8": np.uint8,
        "uint16": np.uint16,
        "uint32": np.uint32,
        "uint64": np.uint64,
        "float": np.float32,
        "double": np.float64
}


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--api", type=float)
    parser.add_argument("--config_name")
    parser.add_argument("--input", type=Path)
    parser.add_argument("--decompressed", type=Path)
    parser.add_argument("--dim", type=int, action="append", default=[])
    parser.add_argument("--type", type=lambda name: TYPES[name], choices=TYPES.values())
    return parser.parse_known_args()[0]


def select_slice(arr: np.ndarray) -> np.ndarray:
    if len(arr.shape) == 2:
        return arr
    elif len(arr.shape) == 3:
        x, y, z = arr.shape
        return arr[x//2, :, :]


def plot_data(args):
    print(args, file=sys.stderr)
    # load the data
    uncompressed = np.fromfile(args.input, dtype=args.type)
    decompressed = np.fromfile(args.decompressed, dtype=args.type)
    error_max = uncompressed.max() - uncompressed.min()
    vmin = uncompressed.min() - .1*error_max
    vmax = uncompressed.max() + .1*error_max

    uncompressed = select_slice(uncompressed)
    decompressed = select_slice(decompressed)
    errors = uncompressed - decompressed

    fig, ax = plt.subplots()
    ax.hist(errors.ravel(), range=(-.008*error_max, .008*error_max), bins=100)
    ax.set_title(args.config_name)
    fig.savefig(WORKDIR / "figures" / (args.config_name + "-hist.png"))

    # plot the error in place
    fig, ax = plt.subplots()
    ax.set_title(args.config_name)
    ax.imshow(np.abs(errors), vmin=0, vmax=error_max*.01)
    fig.savefig(WORKDIR / "figures" / (args.config_name + "-err.png"))

    # plot a side-by-side visualization

    fig, ax = plt.subplots(1, 2)
    ax[0].set_title(args.config_name + " uncompressed")
    ax[0].imshow(uncompressed, vmax=vmax, vmin=vmin)
    ax[1].set_title(args.config_name + " decompressed")
    ax[1].imshow(decompressed, vmax=vmax, vmin=vmin)
    fig.savefig(WORKDIR / "figures" / (args.config_name + "-compare.png"))


if __name__ == "__main__":
    args = parse_args()
    if args.input:
        plot_data(args)