TheRock/rocm-libraries/shared/rocroller/scripts/show_matrix

#!/usr/bin/env python3

# Copyright Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT

"""
Render a matrix stored in a msgpack file as an image.
"""

import argparse
import os
import sys
from pathlib import Path

repo_dir = Path(__file__).resolve().parent.parent
sys.path.append(str(repo_dir / "scripts" / "lib"))

import matrix_utils  # noqa: E402

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)

    parser.add_argument("--verbose", "-v", action="count", default=0)
    parser.add_argument("input", nargs="+", help="Input file")
    parser.add_argument("--output", default=[], nargs="+", help="Output file")

    args = parser.parse_args()

    if len(args.output) == 0:

        def makeOutput(inFile):
            base, _ = os.path.splitext(inFile)
            return base + ".png"

        args.output = list(map(makeOutput, args.input))

    if len(args.output) != len(args.input):
        print(
            "If specified, number of output arguments must equal "
            "number of input arguments"
        )
        exit(-1)

    for inFile, outFile in zip(args.input, args.output):
        if args.verbose > 0:
            print(args)

        fig = matrix_utils.showMatrix(inFile)
        fig.savefig(outFile)
        print(f"Wrote {outFile}")