musrfit/tests/musrview_check/musrview_check.py

#!/usr/bin/env python3
"""
Runs musrview --png on a given msr-file, then compares the generated PNGs
against reference images using pixel-level comparison (Pillow).

Usage:
  musrview_check.py <musrview> <msr-file> <ref-dir> <test-name>
                    [--tol T] [--generate] [musrview-opts...]

Modes:
  default     Compare generated PNGs against references in <ref-dir>/<test-name>/
  --generate  Generate reference PNGs into <ref-dir>/<test-name>/ (no comparison)

Tolerance metric: mean absolute pixel difference normalised to [0, 1].
  0.0 = identical, 1.0 = maximally different.  Default tolerance: 0.03 (~3%).
"""

import argparse
import glob
import os
import shutil
import subprocess
import sys
import time


def check_python_deps():
    """Make sure Pillow and numpy are importable; print a clear, actionable
    error message (rather than a bare traceback) if they are not."""
    missing = []
    for module, pip_name in (("PIL", "Pillow"), ("numpy", "numpy")):
        try:
            __import__(module)
        except ImportError:
            missing.append(pip_name)
    if missing:
        print(f"**ERROR** missing required python package(s): {', '.join(missing)}")
        print(f"  install with: {sys.executable} -m pip install {' '.join(missing)}")
        return False
    return True


def pixel_diff(img_a_path, img_b_path):
    """Return the mean absolute pixel difference normalised to [0, 1]."""
    from PIL import Image
    import numpy as np

    a = np.asarray(Image.open(img_a_path).convert("RGBA"), dtype=np.float64)
    b = np.asarray(Image.open(img_b_path).convert("RGBA"), dtype=np.float64)

    if a.shape != b.shape:
        return 1.0  # completely different dimensions

    return np.mean(np.abs(a - b)) / 255.0


def main():
    # ---- argument parsing ----------------------------------------------------
    parser = argparse.ArgumentParser(description="musrview PNG integration test")
    parser.add_argument("musrview", help="path to musrview executable")
    parser.add_argument("msr_file", help="path to msr input file")
    parser.add_argument("ref_dir", help="root reference directory")
    parser.add_argument("test_name", help="test name (subdirectory in ref_dir)")
    parser.add_argument("--tol", type=float, default=0.03,
                        help="tolerance for pixel comparison (default 0.03)")
    parser.add_argument("--generate", action="store_true",
                        help="generate reference PNGs instead of comparing")

    # everything after the known args is forwarded to musrview
    args, musrview_opts = parser.parse_known_args()

    if not args.generate and not check_python_deps():
        return 1

    msr_basename = os.path.splitext(os.path.basename(args.msr_file))[0]
    ref_subdir = os.path.join(args.ref_dir, args.test_name)

    # ---- run musrview --------------------------------------------------------
    # musrview always writes its PNGs next to the msr-file, so the test has to
    # run there and clean up after itself.
    work_dir = os.path.dirname(os.path.abspath(args.msr_file))
    png_glob = os.path.join(work_dir, f"{msr_basename}_*.png")

    # Snapshot the PNGs that are already there together with their mtimes. A PNG
    # counts as produced by this run if it is either new, or pre-existing but
    # rewritten (a leftover from an earlier aborted run -- musrview simply
    # overwrites it). Matching on mtime rather than mere existence keeps stale
    # leftovers from masking the real output, while still ignoring the PNGs of a
    # sibling test running concurrently on the same msr-file.
    pre_existing = {p: os.stat(p).st_mtime_ns for p in glob.glob(png_glob)}

    generated = []

    def collect():
        """Return the PNGs in work_dir that this musrview run created/rewrote."""
        found = []
        for p in glob.glob(png_glob):
            if pre_existing.get(p) != os.stat(p).st_mtime_ns:
                found.append(p)
        return sorted(found)

    def cleanup():
        """Remove everything this run produced -- must happen on every exit path,
        otherwise the leftovers pile up in doc/examples."""
        for png in generated:
            try:
                os.remove(png)
            except OSError:
                pass

    try:
        cmd = [args.musrview, args.msr_file, "--png"] + musrview_opts
        print(f"running: {' '.join(cmd)}")
        result = subprocess.run(cmd, capture_output=True, text=True, cwd=work_dir)
        if result.returncode != 0:
            generated = collect()
            print(f"**ERROR** musrview returned exit code {result.returncode}")
            print(result.stdout + result.stderr)
            return 1

        # musrview has already exited by the time subprocess.run() returns, but on
        # some systems the PNG file's directory entry becomes visible to this
        # process only after a further delay (observed running under ctest, up to
        # several seconds) -- poll for a while instead of failing on the first
        # empty glob.
        for _ in range(150):
            generated = collect()
            if generated:
                break
            time.sleep(0.1)
        if not generated:
            print(f"**ERROR** no PNGs matching '{msr_basename}_*.png' found")
            print(f"  checked: {work_dir}")
            if result.stdout:
                print(result.stdout)
            return 1

        # ---- generate mode ---------------------------------------------------
        if args.generate:
            os.makedirs(ref_subdir, exist_ok=True)
            for png in generated:
                dst = os.path.join(ref_subdir, os.path.basename(png))
                shutil.copy2(png, dst)
                print(f"  saved reference: {dst}")
            print(f"GENERATE: {len(generated)} reference PNG(s) written to {ref_subdir}")
            return 0

        # ---- compare mode ----------------------------------------------------
        if not os.path.isdir(ref_subdir):
            print(f"**ERROR** reference directory not found: {ref_subdir}")
            return 1

        failures = 0
        compared = 0
        for png_path in generated:
            name = os.path.basename(png_path)
            ref_path = os.path.join(ref_subdir, name)
            if not os.path.isfile(ref_path):
                print(f"FAIL: no reference PNG for {name}")
                failures += 1
                continue

            diff = pixel_diff(png_path, ref_path)
            compared += 1
            if diff > args.tol:
                print(f"FAIL: {name}  diff={diff:.6f} > tol={args.tol:.6f}")
                failures += 1
            else:
                print(f"PASS: {name}  diff={diff:.6f} <= tol={args.tol:.6f}")

        if compared == 0:
            print("**ERROR** no PNGs were compared")
            return 1

        if failures:
            print(f"\n{failures} of {len(generated)} PNG(s) FAILED")
            return 1

        print(f"\nAll {compared} PNG(s) PASSED (tol={args.tol})")
        return 0
    finally:
        cleanup()


if __name__ == "__main__":
    sys.exit(main())