#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ mem_gemini_xlsx.py Burg MEM spectrum program for x-y data. Input: .xlsx / .csv / .txt / .dat Default data format: Column 0 : x-axis, e.g. k [A^-1] Column 1 : signal For EXAFS-like data where the Fourier kernel is exp(i k R), the MEM spectrum is evaluated on the R-axis using omega = dx * R where dx is the sampling interval of the input x-axis. Example: python mem_gemini_xlsx.py a-GeO2_glass_0Pa.xlsx --order 6 --xlim 0 8 Author: revised version based on the user's mem_gemini.py """ import argparse import math import os import sys import numpy as np import pandas as pd import matplotlib.pyplot as plt def load_xy_from_file(infile, sheet=0, xcol=0, ycol=1): """ Load x-y data from xlsx/csv/txt/dat. Parameters ---------- infile : str Input filename. sheet : int or str Excel sheet index or name. xcol : int x-axis column index, 0-based. ycol : int signal column index, 0-based. Returns ------- x_axis : ndarray x-axis data. signal : ndarray signal data. """ ext = os.path.splitext(infile)[1].lower() if ext == ".xlsx": df = pd.read_excel(infile, sheet_name=sheet) x_axis = df.iloc[:, xcol].to_numpy(dtype=float) signal = df.iloc[:, ycol].to_numpy(dtype=float) elif ext == ".csv": df = pd.read_csv(infile) x_axis = df.iloc[:, xcol].to_numpy(dtype=float) signal = df.iloc[:, ycol].to_numpy(dtype=float) elif ext in [".txt", ".dat"]: arr = np.loadtxt(infile) if arr.ndim != 2 or arr.shape[1] <= max(xcol, ycol): raise ValueError( f"Input text file must have at least {max(xcol, ycol) + 1} columns." ) x_axis = arr[:, xcol].astype(float) signal = arr[:, ycol].astype(float) else: raise ValueError( "Unsupported file format. Use .xlsx, .csv, .txt, or .dat." ) mask = np.isfinite(x_axis) & np.isfinite(signal) x_axis = x_axis[mask] signal = signal[mask] if len(x_axis) < 4: raise ValueError("Not enough valid data points.") # Sort by x-axis if necessary if np.any(np.diff(x_axis) < 0): print("WARNING: x-axis is not monotonic. Sorting by x-axis.") idx = np.argsort(x_axis) x_axis = x_axis[idx] signal = signal[idx] return x_axis, signal def estimate_dx(x_axis, rtol=1e-3): """ Estimate sampling interval from x-axis. Parameters ---------- x_axis : ndarray x-axis data. rtol : float Relative tolerance for uniformity check. Returns ------- dx : float Mean sampling interval. rel_std : float Relative standard deviation of dx. """ diffs = np.diff(x_axis) dx = np.mean(diffs) if dx == 0: raise ValueError("x-axis step dx is zero.") rel_std = np.std(diffs) / abs(dx) if rel_std > rtol: print("WARNING: x-axis is not uniformly spaced.") print(f" mean dx = {dx:.10g}") print(f" std(diff) = {np.std(diffs):.10g}") print(f" relative std = {rel_std:.3e}") print(" MEM assumes uniformly sampled data.") print(" Consider interpolating onto a uniform grid if needed.") return dx, rel_std def burg_ar_coefficients(signal, order): """ Burg method for AR coefficients. The AR model is represented as x[n] + a[1] x[n-1] + ... + a[p] x[n-p] = e[n] The MEM spectrum denominator is |1 + sum_j a[j] exp(-i omega j)|^2 Parameters ---------- signal : array-like Input signal. order : int AR order. Returns ------- a : ndarray AR coefficients. a[0] = 1. error_var : float Prediction error variance. reflection : ndarray Reflection coefficients. """ x = np.asarray(signal, dtype=float) x = x[np.isfinite(x)] if len(x) < 4: raise ValueError("Not enough finite signal data.") x = x - np.mean(x) n = len(x) if order <= 0: raise ValueError("order must be positive.") if order >= n: raise ValueError("order must be smaller than number of data points.") a = np.zeros(order + 1, dtype=float) a[0] = 1.0 reflection = np.zeros(order + 1, dtype=float) error_var = np.dot(x, x) / n # Forward and backward prediction errors. # Important: # They must be shifted by one point. ef = x[1:].copy() eb = x[:-1].copy() for m in range(1, order + 1): den = np.dot(ef, ef) + np.dot(eb, eb) if den <= 0: k = 0.0 else: k = -2.0 * np.dot(eb, ef) / den reflection[m] = k a_old = a.copy() if m > 1: a[1:m] = a_old[1:m] + k * a_old[m - 1:0:-1] a[m] = k error_var *= max(0.0, 1.0 - k * k) if m < order: ef_old = ef.copy() eb_old = eb.copy() # Update forward/backward errors. # These arrays shrink by one point at each order. ef = ef_old[1:] + k * eb_old[1:] eb = eb_old[:-1] + k * ef_old[:-1] return a, error_var, reflection def burg_order_scan(signal, max_order): """ Scan AR order using FPE, AIC, and BIC-like criteria. Parameters ---------- signal : array-like Input signal. max_order : int Maximum AR order. Returns ------- df_order : DataFrame Order scan table. best : dict Best orders by FPE, AIC, BIC. """ x = np.asarray(signal, dtype=float) x = x[np.isfinite(x)] x = x - np.mean(x) n = len(x) max_order = min(max_order, n - 2) rows = [] for p in range(1, max_order + 1): try: a, err, refl = burg_ar_coefficients(x, p) if err <= 0 or not np.isfinite(err): fpe = np.inf aic = np.inf bic = np.inf else: # Akaike FPE for AR(p) if n - p - 1 > 0: fpe = err * (n + p + 1) / (n - p - 1) else: fpe = np.inf # Simple Gaussian likelihood criteria up to additive constants aic = n * np.log(err) + 2 * p bic = n * np.log(err) + p * np.log(n) rows.append({ "order": p, "error_var": err, "FPE": fpe, "AIC": aic, "BIC": bic, "last_reflection_coeff": refl[p], }) except Exception as e: rows.append({ "order": p, "error_var": np.nan, "FPE": np.inf, "AIC": np.inf, "BIC": np.inf, "last_reflection_coeff": np.nan, }) df_order = pd.DataFrame(rows) def best_order(col): finite = np.isfinite(df_order[col].to_numpy(dtype=float)) if not np.any(finite): return None idx = df_order.loc[finite, col].idxmin() return int(df_order.loc[idx, "order"]) best = { "fpe": best_order("FPE"), "aic": best_order("AIC"), "bic": best_order("BIC"), } return df_order, best def select_order(signal, fixed_order=None, max_order=40, order_method="fpe"): """ Select AR order. Parameters ---------- signal : ndarray Input signal. fixed_order : int or None If given, use this order. max_order : int Maximum order for automatic scan. order_method : str fpe, aic, bic. Returns ------- selected_order : int Selected AR order. df_order : DataFrame Order scan table. order_info : dict Order selection summary. """ n = len(signal) if fixed_order is not None: if fixed_order <= 0: raise ValueError("--order must be positive.") if fixed_order >= n: raise ValueError("--order must be smaller than number of data points.") df_order, best = burg_order_scan(signal, max_order=min(max_order, n - 2)) order_info = { "order_method": "fixed", "fixed_order": fixed_order, "best_fpe": best.get("fpe"), "best_aic": best.get("aic"), "best_bic": best.get("bic"), } return fixed_order, df_order, order_info order_method = order_method.lower() if order_method not in ["fpe", "aic", "bic"]: raise ValueError("order_method must be one of: fpe, aic, bic") df_order, best = burg_order_scan(signal, max_order=max_order) selected_order = best.get(order_method) if selected_order is None: print("WARNING: automatic order selection failed. Using order=8.") selected_order = min(8, n - 2) order_info = { "order_method": order_method, "fixed_order": "", "best_fpe": best.get("fpe"), "best_aic": best.get("aic"), "best_bic": best.get("bic"), } return selected_order, df_order, order_info def mem_spectrum_r_axis_from_ar(a, error_var, dx, rmin, rmax, ns): """ Calculate MEM spectrum on R-axis. For input x = k [A^-1], sample index j corresponds to k_j = j * dx The Fourier-like phase is k_j * R = dx * R * j Therefore, omega = dx * R and PSD(R) = error_var * dx / |1 + sum_j a[j] exp(-i omega j)|^2 Parameters ---------- a : ndarray AR coefficients with a[0] = 1. error_var : float Prediction error variance. dx : float Sampling interval of input x-axis. rmin, rmax : float R-axis range. ns : int Number of spectrum points. Returns ------- R_axis : ndarray R-axis. psd : ndarray MEM PSD. """ order = len(a) - 1 R_axis = np.linspace(rmin, rmax, ns) psd = np.zeros(ns, dtype=float) for i, Rval in enumerate(R_axis): omega = dx * Rval real = 1.0 imag = 0.0 for j in range(1, order + 1): real += a[j] * math.cos(-omega * j) imag += a[j] * math.sin(-omega * j) denom = real * real + imag * imag if denom <= 0: psd[i] = np.inf else: psd[i] = error_var * dx / denom return R_axis, psd def fft_psd_r_axis(signal, dx, nfft=None): """ Calculate FFT-based PSD on R-axis for comparison. numpy.fft.rfftfreq returns cycles per unit of x. For Fourier kernel exp(i k R), R = 2*pi*f Parameters ---------- signal : ndarray Input signal. dx : float x-axis step. nfft : int or None FFT length. Returns ------- R_fft : ndarray R-axis. psd_fft : ndarray FFT PSD. """ x = np.asarray(signal, dtype=float) x = x[np.isfinite(x)] x = x - np.mean(x) if nfft is None: nfft = len(x) fft_vals = np.fft.rfft(x, n=nfft) freq = np.fft.rfftfreq(nfft, d=dx) R_fft = 2.0 * np.pi * freq psd_fft = (np.abs(fft_vals) ** 2) / len(x) return R_fft, psd_fft def find_main_peak(R_axis, psd, r_exclude=1e-12): """ Find main peak excluding R=0. """ mask = np.isfinite(R_axis) & np.isfinite(psd) & (R_axis > r_exclude) if not np.any(mask): return np.nan, np.nan R2 = R_axis[mask] psd2 = psd[mask] idx = np.argmax(psd2) return R2[idx], psd2[idx] def save_excel( outfile, x_axis, signal_raw, signal_used, R_axis, psd, R_fft, psd_fft, df_order, a_ar, reflection, selected_order, pmm, dx, rel_std, peak_R, peak_psd, peak_R_fft, peak_psd_fft, order_info, args, ): """ Save results to Excel. """ df_input = pd.DataFrame({ f"x_axis_{args.x_unit}": x_axis, "signal_raw": signal_raw, "signal_used_mean_removed": signal_used, }) df_spec = pd.DataFrame({ f"R_MEM_{args.r_unit}": R_axis, "PSD_MEM": psd, }) df_fft = pd.DataFrame({ f"R_FFT_{args.r_unit}": R_fft, "PSD_FFT": psd_fft, }) df_coeff = pd.DataFrame({ "j": np.arange(len(a_ar)), "a_ar": a_ar, "reflection_coeff": reflection, }) df_summary = pd.DataFrame({ "item": [ "input_file", "sheet", "xcol", "ycol", "n_data", f"dx [{args.x_unit}]", "relative_std_dx", "selected_order", "order_method", "fixed_order", "best_order_fpe", "best_order_aic", "best_order_bic", "PMM_error_var", f"rmin [{args.r_unit}]", f"rmax [{args.r_unit}]", "ns", f"main_peak_R_MEM [{args.r_unit}]", "main_peak_PSD_MEM", f"main_peak_R_FFT [{args.r_unit}]", "main_peak_PSD_FFT", "MEM_axis_definition", "FFT_axis_definition", ], "value": [ args.infile, args.sheet, args.xcol, args.ycol, len(x_axis), dx, rel_std, selected_order, order_info.get("order_method", ""), order_info.get("fixed_order", ""), order_info.get("best_fpe", ""), order_info.get("best_aic", ""), order_info.get("best_bic", ""), pmm, args.rmin, args.rmax, args.ns, peak_R, peak_psd, peak_R_fft, peak_psd_fft, "MEM: omega = dx * R", "FFT: R = 2*pi*freq, freq in cycles per x-unit", ], }) with pd.ExcelWriter(outfile, engine="openpyxl") as writer: df_input.to_excel(writer, sheet_name="Input", index=False) df_spec.to_excel(writer, sheet_name="MEM_Spectrum", index=False) df_fft.to_excel(writer, sheet_name="FFT_Spectrum", index=False) df_order.to_excel(writer, sheet_name="OrderScan", index=False) df_coeff.to_excel(writer, sheet_name="AR_Coefficients", index=False) df_summary.to_excel(writer, sheet_name="Summary", index=False) def plot_result( x_axis, signal_raw, signal_used, R_axis, psd, R_fft, psd_fft, selected_order, peak_R, args, ): """ Plot input and MEM/FFT spectrum. """ fig, axes = plt.subplots(nrows=2, ncols=1, figsize=(9, 8)) axes[0].plot(x_axis, signal_raw, label="Raw signal", linewidth=1.5) if args.show_mean_removed: axes[0].plot( x_axis, signal_used, label="Mean-removed signal", linewidth=1.0, ) axes[0].set_xlabel(f"x [{args.x_unit}]") axes[0].set_ylabel("Amplitude") axes[0].set_title("Input data") axes[0].grid(True) axes[0].legend() axes[1].plot( R_axis, psd, label=f"MEM / Burg, order={selected_order}", linewidth=2, ) if args.plot_fft: axes[1].plot( R_fft, psd_fft, "--", label="FFT PSD", linewidth=1, ) if np.isfinite(peak_R): axes[1].axvline(peak_R, linestyle="--", linewidth=1) ymax = np.nanmax(psd[np.isfinite(psd)]) if np.any(np.isfinite(psd)) else 1.0 axes[1].text( peak_R, ymax * 0.8, f"R={peak_R:.3g} {args.r_unit}", rotation=90, va="top", ha="right", ) axes[1].set_xlabel(f"R [{args.r_unit}]") axes[1].set_ylabel("PSD") axes[1].set_title("MEM spectrum") axes[1].grid(True) axes[1].legend() if args.yscale == "log": axes[1].set_yscale("log") if args.xlim is not None: axes[1].set_xlim(args.xlim[0], args.xlim[1]) plt.tight_layout() if args.savefig: plt.savefig(args.savefig, dpi=300) print(f"Figure saved to {args.savefig}") if not args.no_show: plt.show() def main(): parser = argparse.ArgumentParser( description=( "Burg MEM spectrum program reading x-y data from .xlsx/.csv/.txt. " "Column 0 is x-axis, e.g. k [A^-1], column 1 is signal by default." ) ) parser.add_argument( "infile", help="Input file: .xlsx, .csv, .txt, or .dat", ) parser.add_argument( "--sheet", default=0, help="Excel sheet name or index. Default: 0", ) parser.add_argument( "--xcol", type=int, default=0, help="x-axis column index, 0-based. Default: 0", ) parser.add_argument( "--ycol", type=int, default=1, help="signal column index, 0-based. Default: 1", ) parser.add_argument( "--order", type=int, default=None, help="Fixed AR order. If specified, this overrides automatic order selection.", ) parser.add_argument( "--order-method", choices=["fpe", "aic", "bic"], default="fpe", help="Automatic AR order selection method. Default: fpe.", ) parser.add_argument( "--max-order", type=int, default=40, help="Maximum AR order for automatic scan. Default: 40.", ) parser.add_argument( "--rmin", type=float, default=0.0, help="Minimum R value for MEM spectrum. Default: 0.", ) parser.add_argument( "--rmax", type=float, default=None, help=( "Maximum R value for MEM spectrum. " "Default: Nyquist-like value pi/dx." ), ) parser.add_argument( "--ns", type=int, default=0, help="Number of spectrum points. Default: number of input data points.", ) parser.add_argument( "--nfft", type=int, default=0, help="FFT length for comparison. Default: number of input data points.", ) parser.add_argument( "--x-unit", default="A^-1", help="Unit label of input x-axis. Default: A^-1", ) parser.add_argument( "--r-unit", default="A", help="Unit label of output R-axis. Default: A", ) parser.add_argument( "--yscale", choices=["linear", "log"], default="linear", help="Y-scale for spectrum plot. Default: linear.", ) parser.add_argument( "--xlim", type=float, nargs=2, default=None, help="Plot xlim for R-axis, e.g. --xlim 0 8", ) parser.add_argument( "--outfile", default="", help="Output Excel file name. Default: _mem_burg_order=.xlsx", ) parser.add_argument( "--savefig", default="", help="Save figure to this file, e.g. result.png", ) parser.add_argument( "--no-show", action="store_true", help="Do not show matplotlib window.", ) parser.add_argument( "--plot-fft", action="store_true", help="Overlay FFT PSD for comparison.", ) parser.add_argument( "--show-mean-removed", action="store_true", help="Show mean-removed signal in input-data plot.", ) args = parser.parse_args() # Excel sheet can be integer-like string try: args.sheet = int(args.sheet) except Exception: pass try: x_axis, signal_raw = load_xy_from_file( args.infile, sheet=args.sheet, xcol=args.xcol, ycol=args.ycol, ) except Exception as e: print(f"Error loading input file: {e}") sys.exit(1) try: dx, rel_std = estimate_dx(x_axis) except Exception as e: print(f"Error estimating dx: {e}") sys.exit(1) n_data = len(signal_raw) if args.ns <= 0: args.ns = n_data if args.nfft <= 0: args.nfft = n_data if args.rmax is None: # Nyquist-like upper limit: # FFT frequency maximum = 1/(2 dx) # Rmax = 2*pi * 1/(2 dx) = pi/dx args.rmax = math.pi / dx signal_used = signal_raw - np.mean(signal_raw) print("Input data summary") print(f" file = {args.infile}") print(f" n_data = {n_data}") print(f" x_min = {np.min(x_axis):.10g} [{args.x_unit}]") print(f" x_max = {np.max(x_axis):.10g} [{args.x_unit}]") print(f" dx = {dx:.10g} [{args.x_unit}]") print(f" rel std dx = {rel_std:.3e}") print(f" mean(signal) = {np.mean(signal_raw):.10g}") print(f" rmin = {args.rmin:.10g} [{args.r_unit}]") print(f" rmax = {args.rmax:.10g} [{args.r_unit}]") print(f" ns = {args.ns}") print() try: selected_order, df_order, order_info = select_order( signal_raw, fixed_order=args.order, max_order=args.max_order, order_method=args.order_method, ) except Exception as e: print(f"Error selecting AR order: {e}") sys.exit(1) if args.order is not None: print(f"Using fixed AR order = {selected_order}") else: print(f"Order selection method = {args.order_method}") print(f"Best order by FPE = {order_info.get('best_fpe')}") print(f"Best order by AIC = {order_info.get('best_aic')}") print(f"Best order by BIC = {order_info.get('best_bic')}") print(f"Using AR order = {selected_order}") try: a_ar, pmm, reflection = burg_ar_coefficients(signal_raw, selected_order) except Exception as e: print(f"Error calculating Burg AR coefficients: {e}") sys.exit(1) print() print("AR/MEM result") print(f" selected_order = {selected_order}") print(f" PMM/error_var = {pmm:.10g}") print(" AR coefficients:") for i, ai in enumerate(a_ar): print(f" a[{i:2d}] = {ai:.10g}") print() R_axis, psd = mem_spectrum_r_axis_from_ar( a_ar, pmm, dx, args.rmin, args.rmax, args.ns, ) R_fft, psd_fft = fft_psd_r_axis( signal_raw, dx, nfft=args.nfft, ) peak_R, peak_psd = find_main_peak(R_axis, psd) peak_R_fft, peak_psd_fft = find_main_peak(R_fft, psd_fft) print("Main peak excluding R=0") print(f" MEM: R = {peak_R:.10g} [{args.r_unit}], PSD = {peak_psd:.10g}") print(f" FFT: R = {peak_R_fft:.10g} [{args.r_unit}], PSD = {peak_psd_fft:.10g}") print() if args.outfile: outfile = args.outfile else: stem, _ = os.path.splitext(args.infile) outfile = f"{stem}_mem_burg_order={selected_order}.xlsx" try: save_excel( outfile, x_axis, signal_raw, signal_used, R_axis, psd, R_fft, psd_fft, df_order, a_ar, reflection, selected_order, pmm, dx, rel_std, peak_R, peak_psd, peak_R_fft, peak_psd_fft, order_info, args, ) print(f"Results saved to {outfile}") except Exception as e: print(f"Error saving Excel file: {e}") plot_result( x_axis, signal_raw, signal_used, R_axis, psd, R_fft, psd_fft, selected_order, peak_R, args, ) if __name__ == "__main__": main()