Files
Atomizer/tools/generate_synthetic_wfe.py
Antoine 4146e9d8f1 feat(V&V): Updated to FEA CSV format + real M2 mesh injection
- Output now matches WFE_from_CSV_OPD format: ,X,Y,Z,DX,DY,DZ (meters)
- Suite regenerated using real M2 mesh (357 nodes, 308mm diameter)
- All 14 clean test cases: PASS (0.000 nm error)
- 3 noisy cases: expected FAIL due to low node count amplifying noise
- Added --inject mode to use real FEA mesh geometry
- Added lateral displacement test case
2026-03-09 15:56:23 +00:00

644 lines
21 KiB
Python

#!/usr/bin/env python3
"""
Synthetic WFE Surface Generator for Zernike Pipeline Validation
================================================================
Generates synthetic FEA-style CSV files with known Zernike content for
validating the WFE_from_CSV_OPD tool. Output matches the exact format
used by Atomizer's Zernike analysis pipeline.
Output CSV format (matching WFE_from_CSV_OPD):
,X,Y,Z,DX,DY,DZ
0, x_meters, y_meters, z_meters, dx_meters, dy_meters, dz_meters
Where:
X, Y, Z = undeformed node positions (meters)
DX, DY, DZ = displacement vector (meters)
The tool computes OPD from displacements (rigorous method accounts for
lateral DX/DY via interpolation on the reference surface).
Two operating modes:
1. PURE SYNTHETIC: Generate a flat/spherical mirror mesh + inject known
Zernike displacements as DZ. Ground truth is exact.
2. INJECT INTO REAL: Load a real FEA CSV, zero out DZ, then inject known
Zernike content. Tests the full pipeline including real mesh geometry.
Usage:
# Pure synthetic (flat mirror, M1 params)
python generate_synthetic_wfe.py --zernike "5:100,7:50" -o test.csv
# Pure synthetic with spherical surface
python generate_synthetic_wfe.py --preset realistic --surface sphere --roc 5000 -o test.csv
# Inject into real M2 data
python generate_synthetic_wfe.py --inject m2_data.csv --zernike "5:100" -o test_injected.csv
# Full validation suite
python generate_synthetic_wfe.py --suite -d validation_suite/
# Full suite with injection into real mesh
python generate_synthetic_wfe.py --suite --inject m2_data.csv -d validation_suite/
Author: Mario (Atomizer V&V)
Created: 2026-03-09
Project: P-Zernike-Validation (GigaBIT M1)
"""
import sys
import os
import argparse
import json
import csv
from pathlib import Path
from math import factorial
from typing import Dict, List, Tuple, Optional
import numpy as np
# ============================================================================
# Zernike Polynomial Mathematics (matching extract_zernike.py conventions)
# ============================================================================
def noll_indices(j: int) -> Tuple[int, int]:
"""Convert Noll index j to radial order n and azimuthal frequency m."""
if j < 1:
raise ValueError("Noll index j must be >= 1")
count = 0
n = 0
while True:
if n == 0:
ms = [0]
elif n % 2 == 0:
ms = [0] + [m for k in range(1, n // 2 + 1) for m in (-2 * k, 2 * k)]
else:
ms = [m for k in range(0, (n + 1) // 2) for m in (-(2 * k + 1), (2 * k + 1))]
for m in ms:
count += 1
if count == j:
return n, m
n += 1
def zernike_radial(n: int, m: int, r: np.ndarray) -> np.ndarray:
"""Compute radial component R_n^m(r)."""
R = np.zeros_like(r)
m_abs = abs(m)
for s in range((n - m_abs) // 2 + 1):
coef = ((-1) ** s * factorial(n - s) /
(factorial(s) *
factorial((n + m_abs) // 2 - s) *
factorial((n - m_abs) // 2 - s)))
R += coef * r ** (n - 2 * s)
return R
def zernike_noll(j: int, r: np.ndarray, theta: np.ndarray) -> np.ndarray:
"""Evaluate Noll-indexed Zernike polynomial Z_j(r, theta)."""
n, m = noll_indices(j)
R = zernike_radial(n, m, r)
if m == 0:
return R
elif m > 0:
return R * np.cos(m * theta)
else:
return R * np.sin(-m * theta)
def zernike_name(j: int) -> str:
"""Get common optical name for Zernike mode."""
n, m = noll_indices(j)
names = {
(0, 0): "Piston",
(1, -1): "Tilt X", (1, 1): "Tilt Y",
(2, 0): "Defocus",
(2, -2): "Astigmatism 45°", (2, 2): "Astigmatism 0°",
(3, -1): "Coma X", (3, 1): "Coma Y",
(3, -3): "Trefoil X", (3, 3): "Trefoil Y",
(4, 0): "Primary Spherical",
(4, -2): "2nd Astig X", (4, 2): "2nd Astig Y",
(4, -4): "Quadrafoil X", (4, 4): "Quadrafoil Y",
(5, -1): "2nd Coma X", (5, 1): "2nd Coma Y",
(5, -3): "2nd Trefoil X", (5, 3): "2nd Trefoil Y",
(5, -5): "Pentafoil X", (5, 5): "Pentafoil Y",
(6, 0): "2nd Spherical",
}
if (n, m) in names:
return names[(n, m)]
return f"Z({n},{m:+d})"
# ============================================================================
# Mesh Generation
# ============================================================================
def generate_mesh(
diameter_mm: float = 1200.0,
inner_radius_mm: float = 0.0,
n_rings: int = 20,
surface_type: str = "flat",
roc_mm: float = 0.0,
conic: float = 0.0,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
"""
Generate a mirror mesh with realistic node distribution.
Args:
diameter_mm: Mirror outer diameter in mm
inner_radius_mm: Inner radius (0 for full disk)
n_rings: Number of radial rings
surface_type: "flat", "sphere", "parabola"
roc_mm: Radius of curvature in mm (for sphere/parabola)
conic: Conic constant (-1 = parabola, 0 = sphere)
Returns:
x_m, y_m, z_m: Node positions in meters
"""
outer_r_mm = diameter_mm / 2.0
# Generate rings with increasing node count
all_x = []
all_y = []
if inner_radius_mm > 0:
r_values = np.linspace(inner_radius_mm, outer_r_mm, n_rings)
else:
# Include center point
r_values = np.linspace(0, outer_r_mm, n_rings + 1)
for i, r in enumerate(r_values):
if r < 1e-6:
all_x.append(0.0)
all_y.append(0.0)
else:
# More nodes at larger radii (like real FEA mesh)
n_pts = max(6, int(6 + i * 3))
angles = np.linspace(0, 2 * np.pi, n_pts, endpoint=False)
for a in angles:
all_x.append(r * np.cos(a))
all_y.append(r * np.sin(a))
x_mm = np.array(all_x)
y_mm = np.array(all_y)
# Compute Z (surface sag)
if surface_type == "flat":
z_mm = np.zeros_like(x_mm)
elif surface_type in ("sphere", "parabola"):
if roc_mm <= 0:
raise ValueError("roc_mm must be > 0 for curved surfaces")
r2 = x_mm**2 + y_mm**2
if surface_type == "sphere":
# z = R - sqrt(R^2 - r^2)
z_mm = roc_mm - np.sqrt(roc_mm**2 - r2)
else:
# Parabola: z = r^2 / (2R)
z_mm = r2 / (2 * roc_mm)
else:
raise ValueError(f"Unknown surface_type: {surface_type}")
# Convert to meters
return x_mm / 1000.0, y_mm / 1000.0, z_mm / 1000.0
def load_real_mesh(csv_path: str) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
"""
Load node positions from a real FEA CSV file.
Returns:
x_m, y_m, z_m: Node positions in meters
"""
x_list, y_list, z_list = [], [], []
with open(csv_path) as f:
reader = csv.DictReader(f)
for row in reader:
x_list.append(float(row['X']))
y_list.append(float(row['Y']))
z_list.append(float(row['Z']))
return np.array(x_list), np.array(y_list), np.array(z_list)
# ============================================================================
# Surface Synthesis
# ============================================================================
def synthesize_displacements(
x_m: np.ndarray,
y_m: np.ndarray,
coefficients: Dict[int, float],
diameter_mm: float = None,
noise_rms_nm: float = 0.0,
include_lateral: bool = False,
seed: int = 42,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Dict]:
"""
Generate synthetic displacement vectors from Zernike coefficients.
The Zernike content is injected into DZ (surface-normal displacement).
DX and DY are set to zero unless include_lateral is True (adds small
realistic lateral displacements).
Args:
x_m, y_m: Node positions in meters
coefficients: {Noll_index: amplitude_nm}
diameter_mm: Mirror diameter (auto-detected if None)
noise_rms_nm: Gaussian noise RMS in nm
include_lateral: Add small realistic DX/DY
seed: Random seed
Returns:
dx_m, dy_m, dz_m: Displacement vectors in meters
metadata: Ground truth info
"""
# Auto-detect diameter from mesh
if diameter_mm is None:
r_mm = np.sqrt(x_m**2 + y_m**2) * 1000.0
diameter_mm = 2.0 * np.max(r_mm)
outer_r_m = diameter_mm / 2000.0 # meters
# Normalize to unit disk
r_m = np.sqrt(x_m**2 + y_m**2)
r_norm = r_m / outer_r_m
theta = np.arctan2(y_m, x_m)
# Build DZ from Zernike modes (in nm, then convert to meters)
dz_nm = np.zeros_like(x_m)
for j, amp_nm in coefficients.items():
Z_j = zernike_noll(j, r_norm, theta)
dz_nm += amp_nm * Z_j
rms_clean_nm = np.sqrt(np.mean(dz_nm**2))
# Add noise
rng = np.random.default_rng(seed)
if noise_rms_nm > 0:
noise = rng.normal(0, noise_rms_nm, size=dz_nm.shape)
dz_nm += noise
rms_noisy_nm = np.sqrt(np.mean(dz_nm**2))
# Convert to meters
dz_m = dz_nm * 1e-9
# DX, DY
if include_lateral:
# Small lateral displacements (~1/10 of DZ magnitude)
scale = np.max(np.abs(dz_m)) * 0.1
dx_m = rng.normal(0, scale, size=x_m.shape)
dy_m = rng.normal(0, scale, size=y_m.shape)
else:
dx_m = np.zeros_like(x_m)
dy_m = np.zeros_like(y_m)
metadata = {
"input_coefficients": {str(j): amp for j, amp in coefficients.items()},
"coefficient_names": {str(j): zernike_name(j) for j in coefficients},
"n_points": len(x_m),
"diameter_mm": float(diameter_mm),
"rms_nm_clean": float(rms_clean_nm),
"rms_nm_with_noise": float(rms_noisy_nm),
"noise_rms_nm": noise_rms_nm,
"include_lateral": include_lateral,
"seed": seed,
"units": {
"positions": "meters",
"displacements": "meters",
"coefficients": "nanometers",
}
}
return dx_m, dy_m, dz_m, metadata
# ============================================================================
# Output (matching WFE_from_CSV_OPD format exactly)
# ============================================================================
def write_csv_fea(
filepath: str,
x_m: np.ndarray,
y_m: np.ndarray,
z_m: np.ndarray,
dx_m: np.ndarray,
dy_m: np.ndarray,
dz_m: np.ndarray,
):
"""
Write FEA-style CSV matching the WFE_from_CSV_OPD input format.
Format:
,X,Y,Z,DX,DY,DZ
0,x,y,z,dx,dy,dz
1,...
"""
filepath = Path(filepath)
filepath.parent.mkdir(parents=True, exist_ok=True)
with open(filepath, 'w', newline='') as f:
writer = csv.writer(f)
writer.writerow(['', 'X', 'Y', 'Z', 'DX', 'DY', 'DZ'])
for i in range(len(x_m)):
writer.writerow([
i,
f"{x_m[i]:.16e}",
f"{y_m[i]:.16e}",
f"{z_m[i]:.16e}",
f"{dx_m[i]:.5e}",
f"{dy_m[i]:.16e}",
f"{dz_m[i]:.16e}",
])
print(f" Written: {filepath} ({len(x_m)} nodes)")
def write_metadata(filepath: str, metadata: Dict):
"""Write ground truth metadata as JSON."""
filepath = Path(filepath)
filepath.parent.mkdir(parents=True, exist_ok=True)
with open(filepath, 'w') as f:
json.dump(metadata, f, indent=2)
print(f" Metadata: {filepath}")
# ============================================================================
# Preset Test Cases
# ============================================================================
# Realistic M1 mirror: typical gravity-induced deformation pattern
PRESET_REALISTIC = {
5: 80.0, # Astigmatism 0° — dominant gravity mode
6: 45.0, # Astigmatism 45°
7: 30.0, # Coma X
8: 20.0, # Coma Y
9: 15.0, # Trefoil X
11: 10.0, # Primary Spherical
13: 5.0, # Secondary Astigmatism
16: 3.0, # Secondary Coma
22: 2.0, # Secondary Spherical
}
# Single-mode test cases
SINGLE_MODE_TESTS = {
"Z05_astig_0deg": {5: 100.0},
"Z06_astig_45deg": {6: 100.0},
"Z07_coma_x": {7: 100.0},
"Z08_coma_y": {8: 100.0},
"Z09_trefoil_x": {9: 100.0},
"Z10_trefoil_y": {10: 100.0},
"Z11_spherical": {11: 100.0},
"Z22_2nd_spherical": {22: 50.0},
"Z37_high_order": {37: 30.0},
}
# Edge case tests
EDGE_CASE_TESTS = {
"near_zero": {5: 0.1, 7: 0.05, 11: 0.01},
"large_amplitude": {5: 500.0, 7: 300.0},
"many_modes": {j: 100.0 / j for j in range(5, 51)},
}
def generate_single_case(
name: str,
coeffs: Dict[int, float],
output_dir: Path,
x_m: np.ndarray,
y_m: np.ndarray,
z_m: np.ndarray,
diameter_mm: float,
noise_rms_nm: float = 0.0,
include_lateral: bool = False,
):
"""Generate a single test case."""
print(f"\nGenerating: {name}")
dx, dy, dz, meta = synthesize_displacements(
x_m, y_m, coeffs, diameter_mm,
noise_rms_nm=noise_rms_nm,
include_lateral=include_lateral,
)
write_csv_fea(output_dir / f"{name}.csv", x_m, y_m, z_m, dx, dy, dz)
write_metadata(output_dir / f"{name}_truth.json", meta)
return meta
def generate_test_suite(
output_dir: str,
inject_csv: str = None,
diameter_mm: float = None,
inner_radius_mm: float = 0.0,
surface_type: str = "flat",
roc_mm: float = 0.0,
):
"""Generate the full validation test suite."""
output_dir = Path(output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
# Get mesh
if inject_csv:
print(f"Loading real mesh from: {inject_csv}")
x_m, y_m, z_m = load_real_mesh(inject_csv)
if diameter_mm is None:
r_mm = np.sqrt(x_m**2 + y_m**2) * 1000.0
diameter_mm = 2.0 * np.max(r_mm)
mesh_source = f"real FEA mesh ({inject_csv})"
else:
if diameter_mm is None:
diameter_mm = 1200.0
x_m, y_m, z_m = generate_mesh(
diameter_mm, inner_radius_mm,
n_rings=25, surface_type=surface_type, roc_mm=roc_mm
)
mesh_source = f"synthetic {surface_type} (D={diameter_mm}mm)"
print(f"\n Mesh: {mesh_source}")
print(f" Nodes: {len(x_m)}")
print(f" Diameter: {diameter_mm:.1f} mm")
all_cases = {}
# 1. Single-mode tests
print("\n=== Single-Mode Tests ===")
for name, coeffs in SINGLE_MODE_TESTS.items():
meta = generate_single_case(name, coeffs, output_dir, x_m, y_m, z_m, diameter_mm)
all_cases[name] = meta
# 2. Realistic multi-mode
print("\n=== Realistic Multi-Mode ===")
meta = generate_single_case("realistic_gravity", PRESET_REALISTIC, output_dir,
x_m, y_m, z_m, diameter_mm)
all_cases["realistic_gravity"] = meta
# 3. Noisy versions
print("\n=== Noisy Tests ===")
for noise_level in [1.0, 5.0, 10.0]:
name = f"realistic_noise_{noise_level:.0f}nm"
meta = generate_single_case(name, PRESET_REALISTIC, output_dir,
x_m, y_m, z_m, diameter_mm,
noise_rms_nm=noise_level)
all_cases[name] = meta
# 4. Edge cases
print("\n=== Edge Case Tests ===")
for name, coeffs in EDGE_CASE_TESTS.items():
meta = generate_single_case(name, coeffs, output_dir, x_m, y_m, z_m, diameter_mm)
all_cases[name] = meta
# 5. With lateral displacement (tests rigorous OPD method)
print("\n=== Lateral Displacement Test ===")
meta = generate_single_case("realistic_with_lateral", PRESET_REALISTIC, output_dir,
x_m, y_m, z_m, diameter_mm, include_lateral=True)
all_cases["realistic_with_lateral"] = meta
# Summary
print(f"\n{'='*60}")
print(f"Generated {len(all_cases)} test cases in: {output_dir}")
print(f"Mesh source: {mesh_source}")
print(f"{'='*60}")
# Write suite manifest
manifest = {
"suite": "Zernike Pipeline Validation",
"generated": "2026-03-09",
"mesh_source": mesh_source,
"diameter_mm": diameter_mm,
"n_modes": 50,
"cases": all_cases,
}
manifest_path = output_dir / "suite_manifest.json"
with open(manifest_path, 'w') as f:
json.dump(manifest, f, indent=2)
print(f"Manifest: {manifest_path}")
return all_cases
# ============================================================================
# CLI
# ============================================================================
def parse_zernike_string(s: str) -> Dict[int, float]:
"""Parse 'j1:amp1,j2:amp2,...' format."""
coeffs = {}
for pair in s.split(","):
parts = pair.strip().split(":")
if len(parts) != 2:
raise ValueError(f"Invalid format: '{pair}'. Use 'j:amplitude'")
j = int(parts[0])
amp = float(parts[1])
coeffs[j] = amp
return coeffs
def main():
parser = argparse.ArgumentParser(
description="Generate synthetic WFE surfaces for Zernike validation",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Examples:
# Pure astigmatism (100nm) on flat synthetic mesh
python generate_synthetic_wfe.py --zernike "5:100" -o test_astig.csv
# Realistic multi-mode
python generate_synthetic_wfe.py --preset realistic -o test_realistic.csv
# Inject into real M2 mesh
python generate_synthetic_wfe.py --inject m2_real.csv --zernike "5:100,7:50" -o test.csv
# Full suite using real mesh geometry
python generate_synthetic_wfe.py --suite --inject m2_real.csv -d validation_suite/
# Full suite with synthetic spherical mirror
python generate_synthetic_wfe.py --suite --surface sphere --roc 5000 -d validation_suite/
"""
)
parser.add_argument("--zernike", "-z", type=str,
help="Zernike coefficients as 'j:amp_nm,...'")
parser.add_argument("--preset", choices=["realistic"],
help="Use preset coefficient set")
parser.add_argument("--suite", action="store_true",
help="Generate full validation test suite")
parser.add_argument("--inject", type=str,
help="Real FEA CSV to use as mesh source (keeps geometry, replaces displacements)")
parser.add_argument("--output", "-o", type=str, default="synthetic_wfe.csv",
help="Output CSV file path")
parser.add_argument("--output-dir", "-d", type=str, default="validation_suite",
help="Output directory for --suite mode")
parser.add_argument("--diameter", type=float, default=None,
help="Mirror diameter in mm (auto-detected from mesh if not set)")
parser.add_argument("--inner-radius", type=float, default=0.0,
help="Inner radius in mm for annular aperture")
parser.add_argument("--surface", choices=["flat", "sphere", "parabola"],
default="flat", help="Surface type for synthetic mesh")
parser.add_argument("--roc", type=float, default=5000.0,
help="Radius of curvature in mm (for sphere/parabola)")
parser.add_argument("--noise", type=float, default=0.0,
help="Gaussian noise RMS in nm")
parser.add_argument("--lateral", action="store_true",
help="Include small lateral DX/DY displacements")
args = parser.parse_args()
print("=" * 60)
print("Synthetic WFE Surface Generator")
print("Zernike Pipeline Validation — GigaBIT")
print("=" * 60)
if args.suite:
generate_test_suite(
args.output_dir,
inject_csv=args.inject,
diameter_mm=args.diameter,
inner_radius_mm=args.inner_radius,
surface_type=args.surface,
roc_mm=args.roc,
)
else:
# Determine coefficients
if args.preset == "realistic":
coeffs = PRESET_REALISTIC
elif args.zernike:
coeffs = parse_zernike_string(args.zernike)
else:
print("\nERROR: Specify --zernike, --preset, or --suite")
sys.exit(1)
# Get mesh
if args.inject:
x_m, y_m, z_m = load_real_mesh(args.inject)
diameter_mm = args.diameter
else:
diameter_mm = args.diameter or 1200.0
x_m, y_m, z_m = generate_mesh(
diameter_mm, args.inner_radius,
n_rings=25, surface_type=args.surface, roc_mm=args.roc
)
print(f" Nodes: {len(x_m)}")
if diameter_mm:
print(f" Diameter: {diameter_mm:.1f} mm")
print(f"\n Coefficients:")
for j, amp in sorted(coeffs.items()):
print(f" Z{j:2d} ({zernike_name(j):20s}): {amp:8.2f} nm")
dx, dy, dz, meta = synthesize_displacements(
x_m, y_m, coeffs, diameter_mm,
noise_rms_nm=args.noise,
include_lateral=args.lateral,
)
print(f"\n RMS (clean): {meta['rms_nm_clean']:.3f} nm")
if args.noise > 0:
print(f" RMS (noisy): {meta['rms_nm_with_noise']:.3f} nm")
write_csv_fea(args.output, x_m, y_m, z_m, dx, dy, dz)
meta_path = Path(args.output).with_suffix('.json')
write_metadata(str(meta_path), meta)
print("\nDone!")
if __name__ == "__main__":
main()