线站列映射

#!/usr/bin/env python3
"""Convert beamline columns to energy, mu for Athena."""
import argparse
import numpy as np

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("src")
    ap.add_argument("--energy", type=int, default=0)
    ap.add_argument("--i0", type=int, default=1)
    ap.add_argument("--it", type=int, default=2)
    ap.add_argument("--ifluo", type=int, default=None)
    ap.add_argument("--kev", action="store_true", help="energy column is keV")
    ap.add_argument("--out", default="mu.dat")
    args = ap.parse_args()

    raw = np.loadtxt(args.src, comments=("#", "%", ";"))
    e = raw[:, args.energy]
    if args.kev:
        e = e * 1000.0
    i0 = np.clip(raw[:, args.i0], 1e-12, None)
    if args.ifluo is not None:
        mu = raw[:, args.ifluo] / i0
        kind = "fluorescence"
    else:
        it = np.clip(raw[:, args.it], 1e-12, None)
        mu = np.log(i0 / it)
        kind = "transmission"
    np.savetxt(
        args.out,
        np.column_stack([e, mu]),
        header=f"energy_eV  mu_{kind}",
        comments="# ",
    )
    print(f"wrote {args.out}: {len(e)} points, {kind}")

if __name__ == "__main__":
    main()