#!/usr/bin/env python3 # -*- coding: utf-8 -*- import sys import os import numpy as np import pandas as pd import matplotlib.pyplot as plt from spectrum import pburg, Periodogram # ------------------------ # Default parameters # ------------------------ infile = "" ndata = 300 yscale = "log" order = 8 # AR model order # If True, subtract mean before MEM/FFT. # EXAFS-like data often benefits from removing DC component. remove_mean = False argv = sys.argv nargs = len(argv) if nargs > 1: infile = argv[1] if nargs > 2: order = int(argv[2]) if nargs > 3: yscale = argv[3] if nargs > 4: # Example: # python mem_fft_Raxis.py data.xlsx 8 log 1 # remove_mean = True if 4th argument is 1, true, yes, y remove_mean = argv[4].lower() in ["1", "true", "yes", "y"] def generate_input(nmax: int): """ Generate test signal. In this synthetic case, t is dimensionless from 0 to 1. The converted R-axis is still calculated from dt, but for real EXAFS data, t should be k [A^-1]. """ if nmax < 100: print("Warning: nmax < 100. Setting nmax = 100.") nmax = 100 dt = 1.0 / (nmax - 1) t = np.linspace(0, 1.0, nmax) w1 = 124.5 w2 = 62.0 w3 = 67.0 w4 = 3.0 x = ( np.sin(2 * np.pi * t * w1) + np.sin(2 * np.pi * t * w2) + np.sin(2 * np.pi * t * w3) + np.sin(2 * np.pi * t * w4) ) return t, x def load_xy(infile): """ Load two-column data. Column 0: x-axis, e.g. k [A^-1] Column 1: signal """ ext = os.path.splitext(infile)[1].lower() if ext == ".xlsx": try: df = pd.read_excel(infile, sheet_name=0) t = df.iloc[:, 0].to_numpy(dtype=float) x = df.iloc[:, 1].to_numpy(dtype=float) except Exception as e: print(f"Error reading Excel file: {e}") sys.exit(1) elif ext == ".csv": try: arr = np.loadtxt(infile, delimiter=",", skiprows=1, usecols=[0, 1]) t, x = arr.T except Exception as e: print(f"Error reading CSV file: {e}") sys.exit(1) elif ext == ".txt": try: arr = np.loadtxt(infile, skiprows=1, usecols=[0, 1]) t, x = arr.T except Exception as e: print(f"Error reading TXT file: {e}") sys.exit(1) else: print("Unsupported file format. Please use .xlsx, .csv, or .txt.") sys.exit(1) mask = np.isfinite(t) & np.isfinite(x) t = t[mask] x = x[mask] if len(t) < 4: print("Error: not enough valid data points.") sys.exit(1) # Sort by x-axis if needed if np.any(np.diff(t) < 0): print("WARNING: x-axis is not monotonic. Sorting by x-axis.") idx = np.argsort(t) t = t[idx] x = x[idx] return t, x def estimate_dx(t, rtol=1e-3): """ Estimate x-axis step. For EXAFS: t = k [A^-1] dx = dk [A^-1] pburg/Periodogram frequencies are treated as cycles/sample. Therefore: R = 2*pi*frequency_cycles_per_sample/dx """ diffs = np.diff(t) dx = np.mean(diffs) if dx == 0: print("Error: x-axis step is zero.") sys.exit(1) 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/FFT assume uniformly sampled data.") print(" Consider interpolation onto a uniform grid if needed.") return dx, rel_std def find_main_peak(R, psd, rmin=1e-12): """ Find main peak excluding R=0. """ R = np.asarray(R) psd = np.asarray(psd) mask = np.isfinite(R) & np.isfinite(psd) & (R > rmin) if not np.any(mask): return np.nan, np.nan R2 = R[mask] psd2 = psd[mask] idx = np.argmax(psd2) return R2[idx], psd2[idx] print() print(f"{infile=}") print(f"{order=}") print(f"{yscale=}") print(f"{remove_mean=}") # ------------------------ # Load or generate data # ------------------------ if infile == "": t, x_raw = generate_input(ndata) else: t, x_raw = load_xy(infile) dx, rel_std_dx = estimate_dx(t) if remove_mean: x = x_raw - np.mean(x_raw) else: x = x_raw.copy() nfft = len(x) print() print("Input data summary") print(f" n_data = {len(x)}") print(f" nfft = {nfft}") print(f" x_min = {np.min(t):.10g}") print(f" x_max = {np.max(t):.10g}") print(f" dx = {dx:.10g}") print(f" rel_std_dx = {rel_std_dx:.3e}") print(f" mean(raw) = {np.mean(x_raw):.10g}") print(f" mean(used) = {np.mean(x):.10g}") print() # ------------------------ # MEM spectrum, Burg method # ------------------------ burg_spec = pburg(x, order=order, NFFT=nfft) mem_psd = np.asarray(burg_spec.psd) freqs_mem_sample = np.asarray(burg_spec.frequencies()) # ------------------------ # FFT-based PSD # ------------------------ fft_spec = Periodogram(x, NFFT=nfft) fft_psd = np.asarray(fft_spec.psd) freqs_fft_sample = np.asarray(fft_spec.frequencies()) # ------------------------ # Convert horizontal axis # ------------------------ # pburg/Periodogram frequencies are cycles/sample. # # If input t is k [A^-1], dx is dk [A^-1]. # Physical frequency in cycles per A^-1 is: # # f_phys = f_sample / dk # # For Fourier kernel exp(i k R), # # R = 2*pi*f_phys # # Therefore: # # R = 2*pi*f_sample/dk # R_mem = 2.0 * np.pi * freqs_mem_sample / dx R_fft = 2.0 * np.pi * freqs_fft_sample / dx peak_R_mem, peak_psd_mem = find_main_peak(R_mem, mem_psd) peak_R_fft, peak_psd_fft = find_main_peak(R_fft, fft_psd) print("Main peak excluding R=0") print(f" MEM: R = {peak_R_mem:.10g}, PSD = {peak_psd_mem:.10g}") print(f" FFT: R = {peak_R_fft:.10g}, PSD = {peak_psd_fft:.10g}") print() # ------------------------ # Save results to Excel # ------------------------ if infile != "": filebody = infile.rsplit(".", 1)[0] output_filename = f"{filebody}-mem_fft_R_order={order}.xlsx" n_out = max(len(R_mem), len(R_fft)) def pad_array(a, n): a = np.asarray(a) out = np.full(n, np.nan) out[: len(a)] = a return out df_input = pd.DataFrame({ "x_axis": t, "signal_raw": x_raw, "signal_used": x, }) df_output = pd.DataFrame({ "Frequency_MEM_cycles_per_sample": pad_array(freqs_mem_sample, n_out), "R_MEM": pad_array(R_mem, n_out), "PSD_MEM": pad_array(mem_psd, n_out), "Frequency_FFT_cycles_per_sample": pad_array(freqs_fft_sample, n_out), "R_FFT": pad_array(R_fft, n_out), "PSD_FFT": pad_array(fft_psd, n_out), }) df_summary = pd.DataFrame({ "item": [ "input_file", "n_data", "nfft", "order", "yscale", "remove_mean", "x_min", "x_max", "dx", "relative_std_dx", "mean_raw", "mean_used", "main_peak_R_MEM", "main_peak_PSD_MEM", "main_peak_R_FFT", "main_peak_PSD_FFT", "conversion", ], "value": [ infile, len(x), nfft, order, yscale, remove_mean, np.min(t), np.max(t), dx, rel_std_dx, np.mean(x_raw), np.mean(x), peak_R_mem, peak_psd_mem, peak_R_fft, peak_psd_fft, "R = 2*pi*frequency_cycles_per_sample/dx", ], }) try: with pd.ExcelWriter(output_filename, engine="openpyxl") as writer: df_input.to_excel(writer, sheet_name="Input", index=False) df_output.to_excel(writer, sheet_name="MEM_FFT", index=False) df_summary.to_excel(writer, sheet_name="Summary", index=False) print(f"MEM and FFT results saved to {output_filename}") except Exception as e: print(f"Error saving results to Excel file: {e}") # ------------------------ # Plotting # ------------------------ fig, axes = plt.subplots(nrows=2, ncols=1, figsize=(8, 8)) # Input signal axes[0].plot(t, x_raw, label="Raw signal", linewidth=2) if remove_mean: axes[0].plot(t, x, label="Mean-removed signal", linewidth=1) axes[0].set_xlabel("x") axes[0].set_ylabel("signal") axes[0].set_title("Input signal") axes[0].legend() axes[0].grid(True) # MEM/FFT spectra with converted R-axis axes[1].plot(R_mem, mem_psd, label=f"MEM (pburg), order={order}", linewidth=2) axes[1].plot(R_fft, fft_psd, "--", label="FFT PSD", linewidth=1) if yscale == "log": axes[1].set_yscale("log") axes[1].set_xlabel("R") axes[1].set_ylabel("PSD") axes[1].set_title(f"Spectral Estimate: MEM (order={order}) vs FFT") axes[1].legend() axes[1].grid(True) plt.tight_layout() plt.pause(0.1) input("\nPress ENTER to terminate>>\n")