LCF 批处理

#!/usr/bin/env python3
"""Linear combination fitting of flattened XANES with Larch."""
import argparse
import csv
import numpy as np
from larch.io import read_athena
from larch.xafs import pre_edge, find_e0
from larch.math import interp
from scipy.optimize import nnls

def flattened(group, pre=(-150, -30), norm=(150, 800)):
    e0 = find_e0(group.energy, group.mu)
    pre_edge(group, e0=e0, pre1=pre[0], pre2=pre[1],
             norm1=norm[0], norm2=norm[1], nvict=2)
    return np.asarray(group.energy), np.asarray(group.flat), e0

def window(e, y, e0, emin, emax):
    m = (e >= e0 + emin) & (e <= e0 + emax)
    return e[m], y[m]

def lcf(e_unk, y_unk, stds):
    # interpolate standards onto unknown grid
    A = np.column_stack([interp(e, y, e_unk) for e, y, _ in stds])
    w, _ = nnls(A, y_unk)
    if w.sum() > 0:
        w = w / w.sum()
    fit = A @ w
    r = np.sum((y_unk - fit) ** 2) / np.sum(y_unk ** 2)
    return w, fit, r

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--unknowns", required=True)
    ap.add_argument("--standards", required=True)
    ap.add_argument("--emin", type=float, default=-20)
    ap.add_argument("--emax", type=float, default=50)
    ap.add_argument("--csv", default="lcf_report.csv")
    args = ap.parse_args()

    std_proj = read_athena(args.standards)
    stds = []
    names = []
    for name, g in std_proj._athena_groups.items():
        e, y, e0 = flattened(g)
        stds.append((e, y, e0))
        names.append(name)

    rows = []
    unk_proj = read_athena(args.unknowns)
    for name, g in unk_proj._athena_groups.items():
        e, y, e0 = flattened(g)
        ew, yw = window(e, y, e0, args.emin, args.emax)
        packed = []
        for se, sy, se0 in stds:
            sw, ssy = window(se, sy, se0, args.emin, args.emax)
            packed.append((sw, ssy, se0))
        w, fit, r = lcf(ew, yw, packed)
        row = {"group": name, "E0": f"{e0:.2f}", "R": f"{r:.5f}"}
        for n, wi in zip(names, w):
            row[n] = f"{100 * wi:.2f}"
        rows.append(row)
        print(name, row)

    with open(args.csv, "w", newline="") as f:
        wri = csv.DictWriter(f, fieldnames=rows[0].keys())
        wri.writeheader()
        wri.writerows(rows)

if __name__ == "__main__":
    main()