批量归一化 + 展平

#!/usr/bin/env python3
"""Batch-normalize XANES groups from an Athena project."""
from pathlib import Path
import argparse
import numpy as np
from larch.io import read_athena
from larch.xafs import find_e0, pre_edge

def flatten_mu(e, mu, e0):
    """Simple post-edge flatten so μ → 1 after the edge."""
    post = e > e0
    if post.sum() < 10:
        return mu
    coef = np.polyfit(e[post], mu[post] - 1.0, 1)
    out = mu.copy()
    out[post] = mu[post] - np.polyval(coef, e[post])
    return out

def main():
    p = argparse.ArgumentParser()
    p.add_argument("prj")
    p.add_argument("--pre", nargs=2, type=float, default=[-150, -30])
    p.add_argument("--norm", nargs=2, type=float, default=[150, 800])
    p.add_argument("--nvict", type=int, default=2)
    p.add_argument("--out", default="xanes_export")
    args = p.parse_args()

    proj = read_athena(args.prj)
    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)

    for name, g in proj._athena_groups.items():
        e, mu = np.asarray(g.energy), np.asarray(g.mu)
        e0 = find_e0(e, mu)
        pre_edge(
            g,
            e0=e0,
            pre1=args.pre[0],
            pre2=args.pre[1],
            norm1=args.norm[0],
            norm2=args.norm[1],
            nvict=args.nvict,
        )
        flat = flatten_mu(e, np.asarray(g.flat), e0)
        dest = out / f"{name}_flat.csv"
        np.savetxt(
            dest,
            np.column_stack([e, flat]),
            delimiter=",",
            header="energy_eV,flattened_mu",
            comments="",
        )
        print(f"{name:30s}  E0={e0:8.2f}  step={g.edge_step:.4f}  -> {dest}")

if __name__ == "__main__":
    main()