Atomizer/docs/06_PROTOCOLS_DETAILED/ZERNIKE_EXTRACTOR.md

# Zernike Coefficient Extractor

## Overview

The Zernike extractor module provides complete wavefront error (WFE) analysis for telescope mirror optimization. It extracts Zernike polynomial coefficients from FEA displacement results and computes RMS metrics used as optimization objectives.

**Location**: `optimization_engine/extractors/extract_zernike.py`

---

## Mathematical Background

### What are Zernike Polynomials?

Zernike polynomials are a set of orthogonal functions defined on the unit disk. They are the standard basis for describing optical aberrations because:

1. **Orthogonality**: Each mode is independent (no cross-talk)
2. **Physical meaning**: Each mode corresponds to a recognizable aberration
3. **RMS property**: Total RMS² = sum of individual coefficient²

### Noll Indexing Convention

We use the Noll indexing scheme (standard in optics):

| Noll j | n | m  | Name                | Physical Meaning |
|--------|---|----|---------------------|------------------|
| 1      | 0 | 0  | Piston              | Constant offset (ignored) |
| 2      | 1 | 1  | Tilt Y              | Pointing error - correctable |
| 3      | 1 | -1 | Tilt X              | Pointing error - correctable |
| 4      | 2 | 0  | Defocus             | Focus error - correctable |
| 5      | 2 | -2 | Astigmatism 45°     | 3rd order aberration |
| 6      | 2 | 2  | Astigmatism 0°      | 3rd order aberration |
| 7      | 3 | -1 | Coma X              | 3rd order aberration |
| 8      | 3 | 1  | Coma Y              | 3rd order aberration |
| 9      | 3 | -3 | Trefoil X           | Triangular aberration |
| 10     | 3 | 3  | Trefoil Y           | Triangular aberration |
| 11     | 4 | 0  | Primary Spherical   | 4th order spherical |
| 12-50  | ...| ...| Higher orders       | Higher-order aberrations |

### Zernike Polynomial Formula

Each Zernike polynomial Z_j(r, θ) is computed as:

```
Z_j(r, θ) = R_n^m(r) × { cos(m·θ)  if m ≥ 0
                       { sin(|m|·θ) if m < 0

where R_n^m(r) = Σ(s=0 to (n-|m|)/2) [(-1)^s × (n-s)! / (s! × ((n+|m|)/2-s)! × ((n-|m|)/2-s)!)] × r^(n-2s)
```

### Wavefront Error Conversion

FEA gives surface displacement in mm. We convert to wavefront error in nm:

```
WFE = 2 × displacement × 10⁶  [nm]
      ↑                    ↑
      optical reflection   mm → nm
```

The factor of 2 accounts for the optical path difference when light reflects off a surface.

---

## Module Structure

### Files

| File | Purpose |
|------|---------|
| `extract_zernike.py` | Core extraction: Zernike fitting, RMS computation, OP2 parsing |
| `zernike_helpers.py` | High-level helpers for optimization integration |
| `extract_zernike_surface.py` | Surface-based extraction (alternative method) |

### Key Classes

#### `ZernikeExtractor`

Main class for Zernike analysis:

```python
from optimization_engine.extractors import ZernikeExtractor

extractor = ZernikeExtractor(
    op2_path="results/model-solution_1.op2",
    bdf_path="results/model.dat",  # Optional, auto-detected
    displacement_unit="mm",        # Unit in OP2 file
    n_modes=50,                    # Number of Zernike modes
    filter_orders=4                # Modes to filter (J1-J4)
)

# Extract single subcase
result = extractor.extract_subcase("20")
print(f"Filtered RMS: {result['filtered_rms_nm']:.2f} nm")

# Extract relative metrics (target vs reference)
relative = extractor.extract_relative(
    target_subcase="40",
    reference_subcase="20"
)
print(f"Relative RMS (40 vs 20): {relative['relative_filtered_rms_nm']:.2f} nm")

# Extract all subcases
all_results = extractor.extract_all_subcases(reference_subcase="20")
```

---

## RMS Metrics Explained

### Global RMS

Raw RMS of the entire wavefront error surface:

```
global_rms = sqrt(mean(WFE²))
```

### Filtered RMS (J1-J4 removed)

RMS after removing correctable aberrations (piston, tip, tilt, defocus):

```python
# Subtract low-order contribution
WFE_filtered = WFE - Σ(j=1 to 4) c_j × Z_j(r, θ)
filtered_rms = sqrt(mean(WFE_filtered²))
```

**This is typically the primary optimization objective** because:
- Piston (J1): Doesn't affect imaging
- Tip/Tilt (J2-J3): Corrected by telescope pointing
- Defocus (J4): Corrected by focus mechanism

### Optician Workload (J1-J3 removed)

RMS for manufacturing assessment - keeps defocus because it requires material removal:

```python
# Subtract only piston and tilt
WFE_j1to3 = WFE - Σ(j=1 to 3) c_j × Z_j(r, θ)
rms_filter_j1to3 = sqrt(mean(WFE_j1to3²))
```

### Relative RMS Between Subcases

Measures gravity-induced deformation relative to a reference orientation:

```python
# Compute difference surface
ΔWFE = WFE_target - WFE_reference

# Fit Zernike to difference
Δc = zernike_fit(ΔWFE)

# Filter and compute RMS
relative_filtered_rms = sqrt(Σ(j=5 to 50) Δc_j²)
```

---

## Coefficient-Based vs Surface-Based RMS

Due to Zernike orthogonality, these two methods are mathematically equivalent:

### Method 1: Coefficient-Based (Fast)
```python
# From coefficients directly
filtered_rms = sqrt(Σ(j=5 to 50) c_j²)
```

### Method 2: Surface-Based (More accurate for irregular meshes)
```python
# Reconstruct and subtract low-order surface
WFE_low = Σ(j=1 to 4) c_j × Z_j(r, θ)
WFE_filtered = WFE - WFE_low
filtered_rms = sqrt(mean(WFE_filtered²))
```

The module uses the surface-based method for maximum accuracy with FEA meshes.

---

## Usage in Optimization

### Simple: Single Objective

```python
from optimization_engine.extractors import extract_zernike_filtered_rms

def objective(trial):
    # ... run simulation ...

    rms = extract_zernike_filtered_rms(
        op2_file=sim_dir / "model-solution_1.op2",
        subcase="20"
    )
    return rms
```

### Multi-Subcase: Weighted Sum

```python
from optimization_engine.extractors import ZernikeExtractor

def objective(trial):
    # ... run simulation ...

    extractor = ZernikeExtractor(op2_path)

    # Extract relative metrics
    rel_40_20 = extractor.extract_relative("3", "2")['relative_filtered_rms_nm']
    rel_60_20 = extractor.extract_relative("4", "2")['relative_filtered_rms_nm']
    mfg_90 = extractor.extract_relative("1", "2")['relative_rms_filter_j1to3']

    # Weighted objective
    weighted = (
        5.0 * (rel_40_20 / 4.0) +   # Target: 4 nm
        5.0 * (rel_60_20 / 10.0) +  # Target: 10 nm
        1.0 * (mfg_90 / 20.0)       # Target: 20 nm
    ) / 11.0

    return weighted
```

### Using Helper Classes

```python
from optimization_engine.extractors.zernike_helpers import ZernikeObjectiveBuilder

builder = ZernikeObjectiveBuilder(
    op2_finder=lambda: sim_dir / "model-solution_1.op2"
)

builder.add_relative_objective("3", "2", weight=5.0)  # 40 vs 20
builder.add_relative_objective("4", "2", weight=5.0)  # 60 vs 20
builder.add_relative_objective("1", "2",
                               metric="relative_rms_filter_j1to3",
                               weight=1.0)  # 90 vs 20

objective = builder.build_weighted_sum()
value = objective()  # Returns combined metric
```

---

## Output Dictionary Reference

### `extract_subcase()` Returns:

| Key | Type | Description |
|-----|------|-------------|
| `subcase` | str | Subcase identifier |
| `global_rms_nm` | float | Global RMS WFE (nm) |
| `filtered_rms_nm` | float | Filtered RMS (J1-J4 removed) |
| `rms_filter_j1to3` | float | J1-J3 filtered RMS (keeps defocus) |
| `n_nodes` | int | Number of nodes analyzed |
| `defocus_nm` | float | Defocus magnitude (J4) |
| `astigmatism_rms_nm` | float | Combined astigmatism (J5+J6) |
| `coma_rms_nm` | float | Combined coma (J7+J8) |
| `trefoil_rms_nm` | float | Combined trefoil (J9+J10) |
| `spherical_nm` | float | Primary spherical (J11) |

### `extract_relative()` Returns:

| Key | Type | Description |
|-----|------|-------------|
| `target_subcase` | str | Target subcase |
| `reference_subcase` | str | Reference subcase |
| `relative_global_rms_nm` | float | Global RMS of difference |
| `relative_filtered_rms_nm` | float | Filtered RMS of difference |
| `relative_rms_filter_j1to3` | float | J1-J3 filtered RMS of difference |
| `relative_defocus_nm` | float | Defocus change |
| `relative_astigmatism_rms_nm` | float | Astigmatism change |
| ... | ... | (all aberrations with `relative_` prefix) |

---

## Subcase Mapping

NX Nastran subcases map to gravity orientations:

| Subcase ID | Elevation | Purpose |
|------------|-----------|---------|
| 1 | 90° (zenith) | Polishing/manufacturing orientation |
| 2 | 20° | Reference (low elevation) |
| 3 | 40° | Mid-range tracking |
| 4 | 60° | High-range tracking |

The **20° orientation is typically used as reference** because:
- It represents typical low-elevation observing
- Polishing is done at 90°, so we measure change from a tracking position

---

## Saving Zernike Coefficients for Surrogate Training

For neural network training, save all 200 coefficients (50 modes × 4 subcases):

```python
import pandas as pd
from optimization_engine.extractors import ZernikeExtractor

extractor = ZernikeExtractor(op2_path)

# Get coefficients for all subcases
rows = []
for j in range(1, 51):
    row = {'noll_index': j}
    for subcase, label in [('1', '90deg'), ('2', '20deg'),
                           ('3', '40deg'), ('4', '60deg')]:
        result = extractor.extract_subcase(subcase, include_coefficients=True)
        row[f'{label}_nm'] = result['coefficients'][j-1]
    rows.append(row)

df = pd.DataFrame(rows)
df.to_csv(f"zernike_coefficients_trial_{trial_num}.csv", index=False)
```

### CSV Format

| noll_index | 90deg_nm | 20deg_nm | 40deg_nm | 60deg_nm |
|------------|----------|----------|----------|----------|
| 1 | 0.05 | 0.03 | 0.04 | 0.04 |
| 2 | -1.23 | -0.98 | -1.05 | -1.12 |
| ... | ... | ... | ... | ... |
| 50 | 0.02 | 0.01 | 0.02 | 0.02 |

**Note**: These are ABSOLUTE coefficients in nm, not relative RMS values.

---

## Error Handling

### Common Issues

1. **"No displacement data found in OP2"**
   - Check that solve completed successfully
   - Verify OP2 file isn't corrupted or incomplete

2. **"Subcase 'X' not found"**
   - List available subcases: `print(extractor.displacements.keys())`
   - Check subcase numbering in NX simulation setup

3. **"No valid points inside unit disk"**
   - Mirror surface nodes may not be properly identified
   - Check BDF node coordinates

4. **pyNastran version warning**
   - `nx version='2506.5' is not supported` - This is just a warning, extraction still works

---

## Dependencies

```python
# Required
pyNastran >= 1.3.4  # OP2/BDF parsing
numpy >= 1.20       # Numerical computations

# Optional (for visualization)
matplotlib          # Plotting Zernike surfaces
```

---

## References

1. **Noll, R. J. (1976)**. "Zernike polynomials and atmospheric turbulence."
   *Journal of the Optical Society of America*, 66(3), 207-211.

2. **Born, M. & Wolf, E. (1999)**. *Principles of Optics* (7th ed.).
   Cambridge University Press. Chapter 9: Aberrations.

3. **Wyant, J. C. & Creath, K. (1992)**. "Basic Wavefront Aberration Theory
   for Optical Metrology." *Applied Optics and Optical Engineering*, Vol. XI.

---

## Module Exports

```python
from optimization_engine.extractors import (
    # Main class
    ZernikeExtractor,

    # Convenience functions
    extract_zernike_from_op2,
    extract_zernike_filtered_rms,
    extract_zernike_relative_rms,

    # Helpers for optimization
    create_zernike_objective,
    create_relative_zernike_objective,
    ZernikeObjectiveBuilder,

    # Low-level utilities
    compute_zernike_coefficients,
    compute_rms_metrics,
    noll_indices,
    zernike_noll,
    zernike_name,
)
```