astrophotography/pipeline/run.py

"""The push-button entry point: process a session, or a whole archive.

    python run.py                          # process $ASTRO_SESSION
    python run.py PATH                     # process one session
    python run.py --all PATH               # find and process every session
    python run.py --stage register PATH    # run one stage
    python run.py --force PATH             # ignore existing outputs

Stages run in order and each is skipped when its output already exists, so an
interrupted run continues rather than starting again. A failure in one stage of
one session does not stop the others: the failure is recorded and the batch
carries on, because in a batch of twenty the useful outcome is nineteen results
and one clear error, not nothing.
"""
import argparse
import os
import sys
import time
import traceback

import layout
import session as session_mod

STAGES = ("ingest", "measure", "register", "solve", "colour",
          "science", "report")


def _exists(path):
    return os.path.exists(path) and os.path.getsize(path) > 0


def process(root, stages=STAGES, force=False, verbose=True):
    """Run the requested stages for one session. Returns a result summary."""
    s = session_mod.Session(root)
    result = {"root": root, "stages": {}, "ok": True}
    t0 = time.time()

    if "ingest" in stages:
        s.ingest(verbose=verbose)
    s.scan(verbose=verbose)
    if not s.frames:
        print("    no light frames found - skipping")
        result["ok"] = False
        result["error"] = "no frames"
        return result
    s.save_manifest()
    result["target"] = s.target
    result["telescope"] = s.telescope
    result["palette"] = s.palette
    result["frames"] = len(s.frames)

    results = {}
    try:
        if "measure" in stages:
            import measure
            cache = os.path.join(root, layout.INTERMEDIATES, "_measure.npz")
            if force or not _exists(cache):
                print("    measure")
                measure.run(s, force=force, verbose=False)
            else:
                print("    measure (cached)")
            _, results["stats"] = measure.load(s)

        if "register" in stages:
            import measure
            import register
            masters = os.path.join(root, layout.MASTERS)
            have = _exists(os.path.join(masters, f"master-{s.filters[0]}.fit"))
            if force or not have:
                print("    register")
                _, ref = register.run(s, verbose=verbose)
            else:
                print("    register (cached)")
                _, st = measure.load(s)
                ref = measure.choose_reference(s, st)
            results["reference"] = ref

        if "solve" in stages:
            import astrometry
            print("    solve")
            try:
                results["solve"] = astrometry.run(s, verbose=verbose)
            except Exception as exc:                          # noqa: BLE001
                # An unsolved session still produces images; record and go on.
                print(f"      solve failed: {type(exc).__name__}: {exc}")
                results["solve"] = None

        if "colour" in stages:
            import colour
            print("    colour")
            results["final"] = colour.run(s, verbose=verbose)

        if "science" in stages:
            import science
            print("    science")
            try:
                results["science"] = science.run(s, verbose=verbose)
            except Exception as exc:                          # noqa: BLE001
                # Measurements are a bonus on top of the images; losing them
                # must not lose the pictures too.
                print(f"      science failed: {type(exc).__name__}: {exc}")
                results["science"] = None

        if "report" in stages:
            import report
            print("    report")
            report.write(s, results, verbose=verbose)

    except Exception as exc:                                  # noqa: BLE001
        result["ok"] = False
        result["error"] = f"{type(exc).__name__}: {exc}"
        print(f"    FAILED: {result['error']}")
        if verbose:
            traceback.print_exc(limit=3)

    result["solved"] = bool(results.get("solve"))
    result["seconds"] = time.time() - t0
    return result


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("path", nargs="?", default=None,
                    help="session directory, or a tree with --all")
    ap.add_argument("--all", action="store_true",
                    help="find and process every session beneath PATH")
    ap.add_argument("--stage", action="append", choices=STAGES,
                    help="run only this stage (repeatable)")
    ap.add_argument("--force", action="store_true",
                    help="re-run stages even if their output exists")
    args = ap.parse_args(argv)

    root = args.path or layout.SESSION
    stages = tuple(args.stage) if args.stage else STAGES

    targets = session_mod.discover(root) if args.all else [root]
    if not targets:
        print(f"no sessions found under {root}")
        return 1
    print(f"{len(targets)} session(s), stages: {', '.join(stages)}\n")

    results = []
    for i, t in enumerate(targets, 1):
        print(f"[{i}/{len(targets)}] {os.path.relpath(t, root)}")
        results.append(process(t, stages=stages, force=args.force))
        print()

    print("=" * 72)
    print(f"{'target':16s} {'scope':16s} {'palette':8s} "
          f"{'frames':>6s} {'solved':>7s} {'time':>7s}")
    for r in results:
        if not r.get("ok") and "target" not in r:
            print(f"{os.path.basename(r['root']):16s} "
                  f"FAILED: {r.get('error', 'unknown')}")
            continue
        print(f"{r.get('target', '?')[:16]:16s} "
              f"{r.get('telescope', '?')[:16]:16s} "
              f"{r.get('palette', '?'):8s} {r.get('frames', 0):6d} "
              f"{'yes' if r.get('solved') else 'no':>7s} "
              f"{r.get('seconds', 0):6.0f}s")
    failed = [r for r in results if not r.get("ok")]
    unsolved = [r for r in results if r.get("ok") and not r.get("solved")]
    print(f"\n{len(results) - len(failed)}/{len(results)} processed"
          + (f", {len(unsolved)} without an astrometric solution" if unsolved
             else ""))
    return 1 if failed else 0


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