Antoine/Atomizer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add PSD figure generation script for CDR reports

Browse Source 
Generates publication-ready PSD figures from OP2 FEA data:
- psd_multi_angle.png: PSD curves for all elevation angles
- psd_band_decomposition.png: LSF/MSF/HSF bar chart
- psd_40deg_composite.png: WFE surface + PSD side-by-side
- psd_trajectory.png: Band evolution vs elevation

Uses Atomizer's full OPD pipeline for consistency with CDR WFE numbers.
This commit is contained in:
Antoine
2026-02-06 17:38:48 +00:00
parent 38d0994d29
commit a5059dd64a
1 changed files with 271 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

271
tools/generate_psd_figures.py Normal file
View File

@@ -0,0 +1,271 @@
#!/usr/bin/env python3
"""
Generate PSD figures for CDR — using Atomizer's full OPD pipeline.
Ensures exact consistency with CDR WFE numbers.
"""
import sys
sys.path.insert(0, '/home/papa/repos/Atomizer')
import numpy as np
from scipy.interpolate import griddata
from numpy.fft import fft2, fftshift
import warnings
warnings.filterwarnings('ignore')
import os
os.environ['MPLBACKEND'] = 'Agg'
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from optimization_engine.extractors.extract_zernike import (
zernike_noll, compute_zernike_coefficients, DEFAULT_FILTER_ORDERS
)
OP2_PATH = '/home/papa/obsidian-vault/2-Projects/P04-GigaBIT-M1/FEA/CDR-Optimization-V7/iter0_BEST_OF_ALL/assy_m1_assyfem1_sim1-solution_1.op2'
OUTPUT_DIR = '/home/papa/obsidian-vault/2-Projects/P04-GigaBIT-M1/CDR-Report/Images'
ANGLES = ['20', '40', '60', '90']
APERTURE_RADIUS_MM = 575.0
GRID_SIZE = 512
BANDS = {
'LSF': (0.05, 2.0, '#3b82f6'),
'MSF': (2.0, 20.0, '#f59e0b'),
'HSF': (20.0, 100.0,'#ef4444'),
}
def extract_filtered_wfe(ext, angle_str):
"""
Use Atomizer's full OPD pipeline + Zernike filtering.
Gets the WFE array (2× surface, figure-corrected) then removes J1-J4.
"""
# Get the full OPD data (this is what extract_subcase uses internally)
opd_data = ext._build_figure_opd_data(angle_str)
X = opd_data['x_original']
Y = opd_data['y_original']
WFE = opd_data['wfe_nm'] # Already 2× surface error, figure-corrected
# Now do the exact same Zernike filtering as extract_subcase
coeffs, R_max = compute_zernike_coefficients(X, Y, WFE, n_modes=50)
x_c = X - np.mean(X)
y_c = Y - np.mean(Y)
r = np.hypot(x_c / R_max, y_c / R_max)
theta = np.arctan2(y_c, x_c)
filter_orders = DEFAULT_FILTER_ORDERS # 4
Z_low = np.column_stack([
zernike_noll(j, r, theta) for j in range(1, filter_orders + 1)
])
wfe_filtered = WFE - Z_low @ coeffs[:filter_orders]
X_c = X - np.mean(X)
Y_c = Y - np.mean(Y)
rms_filtered = np.sqrt(np.mean(wfe_filtered**2))
return X_c, Y_c, wfe_filtered, rms_filtered
def compute_psd(X, Y, Z_nm, n_grid=GRID_SIZE):
extent = APERTURE_RADIUS_MM
x_grid = np.linspace(-extent, extent, n_grid)
y_grid = np.linspace(-extent, extent, n_grid)
Xg, Yg = np.meshgrid(x_grid, y_grid)
Rg = np.sqrt(Xg**2 + Yg**2)
mask = Rg <= APERTURE_RADIUS_MM
Zg = griddata((X, Y), Z_nm, (Xg, Yg), method='linear', fill_value=0.0)
Zg[~mask] = 0.0
Zg = np.nan_to_num(Zg, nan=0.0)
hx = np.hanning(n_grid)
hy = np.hanning(n_grid)
Zw = Zg * np.outer(hy, hx) * mask
F = fftshift(fft2(Zw))
psd_2d = np.abs(F)**2
dx = 2 * extent / n_grid
freqs = fftshift(np.fft.fftfreq(n_grid, d=dx))
FX, FY = np.meshgrid(freqs, freqs)
FR_cpa = np.sqrt(FX**2 + FY**2) * (2 * APERTURE_RADIUS_MM)
bin_edges = np.logspace(np.log10(0.03), np.log10(n_grid/4), 80)
freqs_out, psd_out = [], []
for i in range(len(bin_edges) - 1):
sel = (FR_cpa >= bin_edges[i]) & (FR_cpa < bin_edges[i+1]) & mask
if np.sum(sel) > 2:
psd_out.append(np.mean(psd_2d[sel]) * dx**2)
freqs_out.append(np.sqrt(bin_edges[i] * bin_edges[i+1]))
return np.array(freqs_out), np.array(psd_out), Zg, mask
def compute_band_rms(X, Y, Z_nm, f_lo, f_hi, n_grid=GRID_SIZE):
extent = APERTURE_RADIUS_MM
x_grid = np.linspace(-extent, extent, n_grid)
y_grid = np.linspace(-extent, extent, n_grid)
Xg, Yg = np.meshgrid(x_grid, y_grid)
Rg = np.sqrt(Xg**2 + Yg**2)
mask = Rg <= APERTURE_RADIUS_MM
Zg = griddata((X, Y), Z_nm, (Xg, Yg), method='linear', fill_value=0.0)
Zg[~mask] = 0.0
Zg = np.nan_to_num(Zg, nan=0.0)
F = np.fft.fft2(Zg)
dx = 2 * extent / n_grid
FX, FY = np.meshgrid(np.fft.fftfreq(n_grid, d=dx), np.fft.fftfreq(n_grid, d=dx))
FR_cpa = np.sqrt(FX**2 + FY**2) * (2 * APERTURE_RADIUS_MM)
band = (FR_cpa >= f_lo) & (FR_cpa < f_hi)
F_band = np.zeros_like(F)
F_band[band] = F[band]
Z_band = np.real(np.fft.ifft2(F_band))
return np.sqrt(np.mean(Z_band[mask]**2))
def main():
from optimization_engine.extractors.extract_zernike_figure import ZernikeOPDExtractor
print("Loading FEA data (Atomizer full OPD pipeline)...")
ext = ZernikeOPDExtractor(OP2_PATH)
# Reference values from Atomizer
all_subcases = ext.extract_all_subcases()
print("\n Atomizer reference WFE (filtered_rms_nm):")
for a in ANGLES:
print(f" {a}°: {all_subcases[int(a)]['filtered_rms_nm']:.2f} nm")
os.makedirs(OUTPUT_DIR, exist_ok=True)
all_results = {}
print("\n Computing PSD from Atomizer OPD pipeline:")
for angle in ANGLES:
print(f"\n {angle}°:")
X, Y, Z_filt, rms_filt = extract_filtered_wfe(ext, angle)
ref_rms = all_subcases[int(angle)]['filtered_rms_nm']
print(f" Our filtered RMS: {rms_filt:.2f} nm | Atomizer ref: {ref_rms:.2f} nm | Δ: {abs(rms_filt - ref_rms):.2f} nm")
freqs, psd, Zg, mask = compute_psd(X, Y, Z_filt)
band_rms = {bname: compute_band_rms(X, Y, Z_filt, flo, fhi) for bname, (flo, fhi, _) in BANDS.items()}
all_results[angle] = {
'freqs': freqs, 'psd': psd, 'Zg': Zg, 'mask': mask,
'band_rms': band_rms, 'total_rms': rms_filt, 'ref_rms': ref_rms,
}
print(f" LSF: {band_rms['LSF']:.2f} | MSF: {band_rms['MSF']:.2f} | HSF: {band_rms['HSF']:.2f} nm")
# ── FIGURES ──
print("\n--- Generating figures ---")
ca = {'20': '#22c55e', '40': '#3b82f6', '60': '#f59e0b', '90': '#ef4444'}
# FIG 1: PSD overlay
fig, ax = plt.subplots(figsize=(12, 7))
for angle in ANGLES:
r = all_results[angle]
v = r['psd'] > 0
ax.loglog(r['freqs'][v], r['psd'][v], '-', color=ca[angle], linewidth=2, alpha=0.85,
label=f"{angle}° (WFE = {r['total_rms']:.1f} nm RMS)")
for bname, (flo, fhi, color) in BANDS.items():
ax.axvspan(flo, fhi, alpha=0.08, color=color, zorder=0)
ax.text(0.3, 0.02, 'LSF', transform=ax.transAxes, fontsize=12, fontweight='bold', color='#3b82f6', alpha=0.7, ha='center')
ax.text(0.6, 0.02, 'MSF', transform=ax.transAxes, fontsize=12, fontweight='bold', color='#f59e0b', alpha=0.7, ha='center')
ax.text(0.88, 0.02, 'HSF', transform=ax.transAxes, fontsize=12, fontweight='bold', color='#ef4444', alpha=0.7, ha='center')
ax.set_xlabel('Spatial Frequency [cycles / aperture]', fontsize=13)
ax.set_ylabel('Power Spectral Density [nm²·mm²]', fontsize=13)
ax.set_title('Residual WFE Power Spectral Density — Tony Hull Methodology\nM1 Mirror, CDR V7 (Zernike J1–J4 corrected)', fontsize=14, fontweight='bold')
ax.legend(fontsize=11, loc='upper right', framealpha=0.9)
ax.set_xlim(0.05, 100); ax.grid(True, which='both', alpha=0.3); ax.tick_params(labelsize=11)
plt.tight_layout()
fig.savefig(os.path.join(OUTPUT_DIR, 'psd_multi_angle.png'), dpi=200, bbox_inches='tight')
plt.close(fig); print(" ✓ psd_multi_angle.png")
# FIG 2: Band decomposition
fig, ax = plt.subplots(figsize=(10, 6))
x_pos = np.arange(len(ANGLES)); width = 0.22
for i, (bname, (flo, fhi, color)) in enumerate(BANDS.items()):
vals = [all_results[a]['band_rms'][bname] for a in ANGLES]
bars = ax.bar(x_pos + i*width - width, vals, width, label=f'{bname} ({flo}–{fhi} c/a)',
color=color, alpha=0.8, edgecolor='white', linewidth=0.5)
for bar, val in zip(bars, vals):
if val > 0.3:
ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.3,
f'{val:.1f}', ha='center', va='bottom', fontsize=9, fontweight='bold')
totals = [all_results[a]['total_rms'] for a in ANGLES]
ax.plot(x_pos, totals, 'ks--', markersize=8, linewidth=1.5, label='Total WFE (J5+)',
markeredgecolor='white', markeredgewidth=1, zorder=5)
for xp, t in zip(x_pos, totals):
ax.text(xp, t + 0.8, f'{t:.1f}', ha='center', va='bottom', fontsize=9, fontweight='bold', color='#333')
ax.set_xticks(x_pos); ax.set_xticklabels([f'{a}°' for a in ANGLES], fontsize=12)
ax.set_xlabel('Elevation Angle', fontsize=13); ax.set_ylabel('RMS WFE [nm]', fontsize=13)
ax.set_title('Residual WFE Band Decomposition (LSF / MSF / HSF)\nZernike J1–J4 Corrected (Piston, Tilt, Defocus)', fontsize=14, fontweight='bold')
ax.legend(fontsize=10, loc='upper left'); ax.grid(True, axis='y', alpha=0.3); ax.tick_params(labelsize=11)
plt.tight_layout()
fig.savefig(os.path.join(OUTPUT_DIR, 'psd_band_decomposition.png'), dpi=200, bbox_inches='tight')
plt.close(fig); print(" ✓ psd_band_decomposition.png")
# FIG 3: 40° composite
r40 = all_results['40']
fig, axes = plt.subplots(1, 2, figsize=(16, 7))
ax1 = axes[0]
Zp = np.where(r40['mask'], r40['Zg'], np.nan)
vmax = np.nanpercentile(np.abs(Zp), 99)
im = ax1.imshow(Zp, extent=[-APERTURE_RADIUS_MM, APERTURE_RADIUS_MM, -APERTURE_RADIUS_MM, APERTURE_RADIUS_MM],
cmap='RdBu_r', vmin=-vmax, vmax=vmax, origin='lower')
plt.colorbar(im, ax=ax1, shrink=0.8, label='Residual WFE [nm]')
ax1.set_xlabel('X [mm]', fontsize=12); ax1.set_ylabel('Y [mm]', fontsize=12)
ax1.set_title(f'Residual WFE at 40° Elevation\n(J1–J4 removed, RMS = {r40["total_rms"]:.1f} nm)', fontsize=13, fontweight='bold')
ax1.set_aspect('equal')
ax2 = axes[1]
v = r40['psd'] > 0
ax2.loglog(r40['freqs'][v], r40['psd'][v], '-', color='#3b82f6', linewidth=2.5)
for bname, (flo, fhi, color) in BANDS.items():
ax2.axvspan(flo, fhi, alpha=0.1, color=color)
ax2.annotate(f'{bname}\n{r40["band_rms"][bname]:.1f} nm',
xy=(np.sqrt(flo*fhi), 1), xycoords=('data', 'axes fraction'),
fontsize=10, fontweight='bold', color=color,
ha='center', va='top', xytext=(0, -10), textcoords='offset points')
ax2.set_xlabel('Spatial Frequency [cycles/aperture]', fontsize=12); ax2.set_ylabel('PSD [nm²·mm²]', fontsize=12)
ax2.set_title('Power Spectral Density at 40°\nTony Hull Method', fontsize=13, fontweight='bold')
ax2.set_xlim(0.05, 100); ax2.grid(True, which='both', alpha=0.3)
plt.tight_layout()
fig.savefig(os.path.join(OUTPUT_DIR, 'psd_40deg_composite.png'), dpi=200, bbox_inches='tight')
plt.close(fig); print(" ✓ psd_40deg_composite.png")
# FIG 4: Trajectory
fig, ax = plt.subplots(figsize=(10, 6))
an = [int(a) for a in ANGLES]
for bname, (flo, fhi, color) in BANDS.items():
ax.plot(an, [all_results[a]['band_rms'][bname] for a in ANGLES], 'o-', color=color, linewidth=2.5, markersize=8,
label=f'{bname} ({flo}–{fhi} c/a)', markeredgecolor='white', markeredgewidth=1)
ax.plot(an, [all_results[a]['total_rms'] for a in ANGLES], 's--', color='#1e293b', linewidth=2, markersize=8,
label='Total WFE (J5+)', markeredgecolor='white', markeredgewidth=1)
ax.set_xlabel('Elevation Angle [°]', fontsize=13); ax.set_ylabel('RMS WFE [nm]', fontsize=13)
ax.set_title('Residual WFE PSD Band Trajectory vs Elevation\nZernike J1–J4 Corrected', fontsize=14, fontweight='bold')
ax.legend(fontsize=11); ax.grid(True, alpha=0.3); ax.set_xticks(an); ax.tick_params(labelsize=11)
plt.tight_layout()
fig.savefig(os.path.join(OUTPUT_DIR, 'psd_trajectory.png'), dpi=200, bbox_inches='tight')
plt.close(fig); print(" ✓ psd_trajectory.png")
# Summary
print("\n" + "="*75)
print("RESIDUAL WFE PSD — Tony Hull Method (Atomizer OPD Pipeline)")
print("="*75)
print(f"{'Angle':>6} {'WFE RMS':>8} {'Ref':>8} {'Δ':>6} {'LSF':>8} {'MSF':>8} {'HSF':>8}")
print("-"*75)
for a in ANGLES:
r = all_results[a]
d = abs(r['total_rms'] - r['ref_rms'])
print(f"{a:>5}° {r['total_rms']:>8.2f} {r['ref_rms']:>8.2f} {d:>6.2f} {r['band_rms']['LSF']:>8.2f} {r['band_rms']['MSF']:>8.2f} {r['band_rms']['HSF']:>8.2f}")
print("="*75)
if __name__ == '__main__':
main()