Antoine
a5059dd64a
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.
272 lines
12 KiB
Python
272 lines
12 KiB
Python
#!/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()
|