Antoine/Atomizer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
f9373bee9941022ff9f27a73f3fd582c5709430e
Atomizer/tools/generate_synthetic_wfe.py
Antoine f9373bee99 feat(V&V): Zernike pipeline validation - synthetic WFE generator + round-trip validator 
- generate_synthetic_wfe.py: Creates synthetic OPD surfaces from known Zernike coefficients
- validate_zernike_roundtrip.py: Round-trip validation (generate → fit → compare)
- validation_suite/: 18 test cases (single mode, multi-mode, noisy, edge cases)
- All 18 test cases pass with 0.000 nm error (clean) and <0.3 nm (10nm noise)
- M1 params: 1200mm dia, 135.75mm inner radius, 50 Noll modes

Project: P-Zernike-Validation (GigaBIT M1)
Requested by: Adyn Miles (StarSpec) for risk reduction before M2/M3 ordering
2026-03-09 15:49:06 +00:00

527 lines
18 KiB
Python
Raw Blame History

#!/usr/bin/env python3
"""
Synthetic WFE Surface Generator for Zernike Pipeline Validation
================================================================
Generates synthetic Optical Path Difference (OPD) maps from user-defined
Zernike coefficients. Used to validate the Atomizer Zernike fitting pipeline
by creating "known truth" surfaces that can be round-tripped through the
WFE_from_CSV_OPD tool.
Features:
- Noll-indexed Zernike polynomials (standard optical convention)
- Full-disk or annular aperture support
- Configurable grid density, mirror diameter, noise level
- Multiple output formats: CSV (for WFE_from_CSV_OPD), NumPy, plot
- Preset test cases for common validation scenarios
Usage:
# Single mode test (pure astigmatism)
python generate_synthetic_wfe.py --mode single --zernike "5:100" --output test_astig.csv
# Multi-mode realistic mirror
python generate_synthetic_wfe.py --mode multi --output test_multi.csv
# Custom coefficients
python generate_synthetic_wfe.py --zernike "5:80,7:45,9:25,11:15" --output test_custom.csv
# With noise
python generate_synthetic_wfe.py --mode multi --noise 2.0 --output test_noisy.csv
# Full test suite (generates all validation cases)
python generate_synthetic_wfe.py --suite --output-dir validation_suite/
Output CSV format (matching WFE_from_CSV_OPD):
x(mm), y(mm), dz(mm)
where dz is surface displacement (OPD = 2 * dz for reflective surfaces)
Author: Mario (Atomizer V&V)
Created: 2026-03-09
Project: P-Zernike-Validation (GigaBIT M1)
"""
import sys
import os
import argparse
import json
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})"
# ============================================================================
# Surface Generation
# ============================================================================
def generate_grid(
diameter_mm: float = 1200.0,
inner_radius_mm: float = 135.75,
n_points_radial: int = 200,
grid_type: str = "cartesian",
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
"""
Generate a 2D grid of points within the mirror aperture.
Returns:
x_mm, y_mm: Physical coordinates in mm
mask: Boolean array (True = inside aperture)
"""
outer_radius = diameter_mm / 2.0
if grid_type == "cartesian":
# Regular Cartesian grid
n = n_points_radial * 2
x_1d = np.linspace(-outer_radius, outer_radius, n)
y_1d = np.linspace(-outer_radius, outer_radius, n)
x_mm, y_mm = np.meshgrid(x_1d, y_1d)
x_mm = x_mm.ravel()
y_mm = y_mm.ravel()
elif grid_type == "scattered":
# Random scattered points (simulates real mesh)
n_total = (n_points_radial * 2) ** 2
rng = np.random.default_rng(42)
x_mm = rng.uniform(-outer_radius, outer_radius, n_total)
y_mm = rng.uniform(-outer_radius, outer_radius, n_total)
else:
raise ValueError(f"Unknown grid_type: {grid_type}")
# Apply aperture mask
r = np.sqrt(x_mm**2 + y_mm**2)
if inner_radius_mm > 0:
mask = (r <= outer_radius) & (r >= inner_radius_mm)
else:
mask = r <= outer_radius
return x_mm[mask], y_mm[mask], mask
def synthesize_surface(
x_mm: np.ndarray,
y_mm: np.ndarray,
coefficients: Dict[int, float],
diameter_mm: float = 1200.0,
noise_rms_nm: float = 0.0,
seed: int = 42,
) -> Tuple[np.ndarray, Dict]:
"""
Generate a synthetic OPD surface from Zernike coefficients.
Args:
x_mm, y_mm: Physical coordinates in mm
coefficients: Dict of {Noll_index: amplitude_nm}
diameter_mm: Mirror diameter in mm
noise_rms_nm: RMS of Gaussian noise to add (nm)
seed: Random seed for reproducibility
Returns:
opd_nm: OPD values in nanometers at each point
metadata: Dict with ground truth info
"""
outer_radius = diameter_mm / 2.0
# Normalize to unit disk
r_norm = np.sqrt(x_mm**2 + y_mm**2) / outer_radius
theta = np.arctan2(y_mm, x_mm)
# Build surface from Zernike modes
opd_nm = np.zeros_like(x_mm)
for j, amp_nm in coefficients.items():
Z_j = zernike_noll(j, r_norm, theta)
opd_nm += amp_nm * Z_j
# Compute ground truth RMS (before noise)
rms_total = np.sqrt(np.mean(opd_nm**2))
# Add noise if requested
if noise_rms_nm > 0:
rng = np.random.default_rng(seed)
noise = rng.normal(0, noise_rms_nm, size=opd_nm.shape)
opd_nm += noise
rms_with_noise = np.sqrt(np.mean(opd_nm**2))
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_mm),
"diameter_mm": diameter_mm,
"rms_nm_clean": float(rms_total),
"rms_nm_with_noise": float(rms_with_noise),
"noise_rms_nm": noise_rms_nm,
"seed": seed,
}
return opd_nm, metadata
# ============================================================================
# Output Formats
# ============================================================================
def write_csv_opd(
filepath: str,
x_mm: np.ndarray,
y_mm: np.ndarray,
opd_nm: np.ndarray,
opd_unit: str = "nm",
):
"""
Write OPD surface to CSV format compatible with WFE_from_CSV_OPD.
Default format: x(mm), y(mm), opd(nm)
Can also output in mm for displacement: x(mm), y(mm), dz(mm)
NOTE: Update this function once the exact WFE_from_CSV_OPD format is confirmed.
"""
filepath = Path(filepath)
filepath.parent.mkdir(parents=True, exist_ok=True)
if opd_unit == "nm":
header = "x(mm),y(mm),opd(nm)"
data = np.column_stack([x_mm, y_mm, opd_nm])
elif opd_unit == "mm":
# Convert nm to mm for displacement
dz_mm = opd_nm / 1e6
header = "x(mm),y(mm),dz(mm)"
data = np.column_stack([x_mm, y_mm, dz_mm])
else:
raise ValueError(f"Unknown opd_unit: {opd_unit}")
np.savetxt(filepath, data, delimiter=",", header=header, comments="",
fmt="%.10e")
print(f" Written: {filepath} ({len(x_mm)} points)")
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}, # Sub-nm coefficients
"large_amplitude": {5: 500.0, 7: 300.0}, # Large deformation
"many_modes": {j: 100.0 / j for j in range(5, 51)}, # All modes, decreasing
}
def generate_test_suite(
output_dir: str,
diameter_mm: float = 1200.0,
inner_radius_mm: float = 135.75,
n_points_radial: int = 200,
):
"""Generate the full validation test suite."""
output_dir = Path(output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
all_cases = {}
# 1. Single-mode tests
print("\n=== Single-Mode Tests ===")
for name, coeffs in SINGLE_MODE_TESTS.items():
print(f"\nGenerating: {name}")
x, y, _ = generate_grid(diameter_mm, inner_radius_mm, n_points_radial)
opd, meta = synthesize_surface(x, y, coeffs, diameter_mm)
write_csv_opd(output_dir / f"{name}.csv", x, y, opd)
write_metadata(output_dir / f"{name}_truth.json", meta)
all_cases[name] = meta
# 2. Realistic multi-mode
print("\n=== Realistic Multi-Mode ===")
name = "realistic_gravity"
print(f"\nGenerating: {name}")
x, y, _ = generate_grid(diameter_mm, inner_radius_mm, n_points_radial)
opd, meta = synthesize_surface(x, y, PRESET_REALISTIC, diameter_mm)
write_csv_opd(output_dir / f"{name}.csv", x, y, opd)
write_metadata(output_dir / f"{name}_truth.json", meta)
all_cases[name] = 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"
print(f"\nGenerating: {name}")
x, y, _ = generate_grid(diameter_mm, inner_radius_mm, n_points_radial)
opd, meta = synthesize_surface(x, y, PRESET_REALISTIC, diameter_mm, noise_rms_nm=noise_level)
write_csv_opd(output_dir / f"{name}.csv", x, y, opd)
write_metadata(output_dir / f"{name}_truth.json", meta)
all_cases[name] = meta
# 4. Edge cases
print("\n=== Edge Case Tests ===")
for name, coeffs in EDGE_CASE_TESTS.items():
print(f"\nGenerating: {name}")
x, y, _ = generate_grid(diameter_mm, inner_radius_mm, n_points_radial)
opd, meta = synthesize_surface(x, y, coeffs, diameter_mm)
write_csv_opd(output_dir / f"{name}.csv", x, y, opd)
write_metadata(output_dir / f"{name}_truth.json", meta)
all_cases[name] = meta
# 5. Full-disk (no central hole) for comparison
print("\n=== Full-Disk (No Hole) ===")
name = "realistic_full_disk"
print(f"\nGenerating: {name}")
x, y, _ = generate_grid(diameter_mm, 0.0, n_points_radial)
opd, meta = synthesize_surface(x, y, PRESET_REALISTIC, diameter_mm)
meta["inner_radius_mm"] = 0.0
write_csv_opd(output_dir / f"{name}.csv", x, y, opd)
write_metadata(output_dir / f"{name}_truth.json", meta)
all_cases[name] = meta
# 6. Scattered points (simulates real FEA mesh)
print("\n=== Scattered Grid (FEA-like) ===")
name = "realistic_scattered"
print(f"\nGenerating: {name}")
x, y, _ = generate_grid(diameter_mm, inner_radius_mm, n_points_radial, grid_type="scattered")
opd, meta = synthesize_surface(x, y, PRESET_REALISTIC, diameter_mm)
meta["grid_type"] = "scattered"
write_csv_opd(output_dir / f"{name}.csv", x, y, opd)
write_metadata(output_dir / f"{name}_truth.json", meta)
all_cases[name] = meta
# Summary
print(f"\n{'='*60}")
print(f"Generated {len(all_cases)} test cases in: {output_dir}")
print(f"{'='*60}")
# Write suite manifest
manifest = {
"suite": "Zernike Pipeline Validation",
"generated": "2026-03-09",
"mirror_diameter_mm": diameter_mm,
"inner_radius_mm": inner_radius_mm,
"n_zernike_modes": 50,
"n_points_radial": n_points_radial,
"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)
python generate_synthetic_wfe.py --zernike "5:100" -o test_astig.csv
# Realistic multi-mode mirror
python generate_synthetic_wfe.py --preset realistic -o test_realistic.csv
# Full validation suite
python generate_synthetic_wfe.py --suite -d validation_suite/
# Custom with noise
python generate_synthetic_wfe.py --zernike "5:80,7:45,11:15" --noise 2.0 -o test.csv
"""
)
parser.add_argument("--zernike", "-z", type=str,
help="Zernike coefficients as 'j:amp_nm,...' (e.g. '5:100,7:50')")
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("--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=1200.0,
help="Mirror diameter in mm (default: 1200)")
parser.add_argument("--inner-radius", type=float, default=135.75,
help="Inner radius in mm for annular aperture (default: 135.75, 0=full disk)")
parser.add_argument("--grid-points", type=int, default=200,
help="Number of radial grid points (default: 200, total ~ 4*N^2)")
parser.add_argument("--noise", type=float, default=0.0,
help="Gaussian noise RMS in nm (default: 0)")
parser.add_argument("--grid-type", choices=["cartesian", "scattered"],
default="cartesian",
help="Grid type (default: cartesian)")
parser.add_argument("--unit", choices=["nm", "mm"], default="nm",
help="OPD output unit (default: nm)")
args = parser.parse_args()
print("=" * 60)
print("Synthetic WFE Surface Generator")
print("Zernike Pipeline Validation — GigaBIT M1")
print("=" * 60)
print(f" Mirror diameter: {args.diameter} mm")
print(f" Inner radius: {args.inner_radius} mm")
print(f" Grid points: ~{(args.grid_points*2)**2}")
if args.suite:
generate_test_suite(
args.output_dir,
args.diameter,
args.inner_radius,
args.grid_points,
)
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)
print(f"\n Coefficients:")
for j, amp in sorted(coeffs.items()):
print(f" Z{j:2d} ({zernike_name(j):20s}): {amp:8.2f} nm")
# Generate
x, y, _ = generate_grid(args.diameter, args.inner_radius,
args.grid_points, args.grid_type)
opd, meta = synthesize_surface(x, y, coeffs, args.diameter,
noise_rms_nm=args.noise)
print(f"\n Points in aperture: {len(x)}")
print(f" RMS (clean): {meta['rms_nm_clean']:.3f} nm")
if args.noise > 0:
print(f" RMS (noisy): {meta['rms_nm_with_noise']:.3f} nm")
# Write outputs
print(f"\n Output unit: {args.unit}")
write_csv_opd(args.output, x, y, opd, opd_unit=args.unit)
meta_path = Path(args.output).with_suffix('.json')
write_metadata(str(meta_path), meta)
print("\nDone!")
if __name__ == "__main__":
main()
Reference in New Issue View Git Blame Copy Permalink