Atomizer/optimization_engine/extractors/extract_zernike_figure.py

"""
OPD Zernike Extractor (Most Rigorous Method)
=============================================

This is the RECOMMENDED Zernike extractor for telescope mirror optimization.
It computes surface error using the ACTUAL undeformed mesh geometry as the
reference surface, rather than assuming any analytical shape.

WHY THIS IS THE MOST ROBUST:
----------------------------
1. Works with ANY surface shape (parabola, hyperbola, asphere, freeform)
2. No need to know/estimate focal length or conic constant
3. Uses the actual mesh geometry as ground truth
4. Eliminates errors from prescription/shape approximations
5. Properly accounts for lateral (X, Y) displacement via interpolation
6. Supports ANNULAR APERTURES (mirrors with central holes)

HOW IT WORKS:
-------------
The key insight: The BDF geometry for nodes present in OP2 IS the undeformed
reference surface (i.e., the optical figure before deformation).

1. Load BDF geometry for nodes that have displacements in OP2
2. Build 2D interpolator z_figure(x, y) from undeformed coordinates
3. For each deformed node at (x0+dx, y0+dy, z0+dz):
   - Interpolate z_figure at the deformed (x,y) position
   - Surface error = (z0 + dz) - z_interpolated
4. Fit Zernike polynomials to the surface error map
5. For annular apertures: exclude central hole from fitting & RMS calculations

The interpolation accounts for the fact that when a node moves laterally,
it should be compared against the figure height at its NEW position.

ANNULAR APERTURE SUPPORT:
-------------------------
For mirrors with central holes (e.g., M1 Mirror with 271.5mm diameter hole):
- Set inner_radius to exclude the central obscuration from Zernike fitting
- Standard Zernike polynomials are fit only to the annular region
- RMS calculations are computed only over the annular aperture

USAGE:
------
    from optimization_engine.extractors import ZernikeOPDExtractor

    # Standard usage (full disk)
    extractor = ZernikeOPDExtractor(op2_file)
    result = extractor.extract_subcase('20')

    # With annular aperture (e.g., 271.5mm diameter central hole)
    extractor = ZernikeOPDExtractor(op2_file, inner_radius=135.75)
    result = extractor.extract_subcase('20')

    # Simple convenience function
    from optimization_engine.extractors import extract_zernike_opd
    result = extract_zernike_opd(op2_file, subcase='20', inner_radius=135.75)

Author: Atomizer Framework
Date: 2024
Updated: 2025-01-06 - Added annular aperture support
"""

from pathlib import Path
from typing import Dict, Any, Optional, List, Tuple, Union
import numpy as np
from scipy.interpolate import LinearNDInterpolator, CloughTocher2DInterpolator
import warnings

# Import base Zernike functionality
from .extract_zernike import (
    compute_zernike_coefficients,
    compute_aberration_magnitudes,
    zernike_noll,
    zernike_label,
    read_node_geometry,
    find_geometry_file,
    extract_displacements_by_subcase,
    UNIT_TO_NM,
    DEFAULT_N_MODES,
    DEFAULT_FILTER_ORDERS,
)

try:
    from pyNastran.bdf.bdf import BDF
except ImportError:
    BDF = None


# ============================================================================
# Annular Aperture Support
# ============================================================================

def compute_zernike_coefficients_annular(
    x: np.ndarray,
    y: np.ndarray,
    values: np.ndarray,
    n_modes: int,
    inner_radius: float,
    chunk_size: int = 100000
) -> Tuple[np.ndarray, float, float, np.ndarray]:
    """
    Fit Zernike coefficients with ANNULAR APERTURE masking.

    Points inside the central obscuration (r < inner_radius) are EXCLUDED
    from the least-squares fitting. This is the correct approach for mirrors
    with central holes where no actual surface exists.

    Args:
        x, y: Node coordinates (in same units as inner_radius, typically mm)
        values: Surface values at each node (e.g., WFE in nm)
        n_modes: Number of Zernike modes to fit
        inner_radius: Inner radius of annular aperture (same units as x, y)
        chunk_size: Chunk size for memory-efficient processing

    Returns:
        Tuple of:
        - coefficients: Zernike coefficients
        - R_outer: Outer radius (normalization radius)
        - r_inner_normalized: Inner radius normalized to unit disk
        - annular_mask: Boolean mask for points in annular region
    """
    from numpy.linalg import LinAlgError

    # Center coordinates
    x_centered = x - np.mean(x)
    y_centered = y - np.mean(y)

    # Compute radial distance (in original units, before normalization)
    r_physical = np.hypot(x_centered, y_centered)

    # Outer radius for normalization
    R_outer = float(np.max(r_physical))

    # Normalize to unit disk
    r = (r_physical / R_outer).astype(np.float32)
    theta = np.arctan2(y_centered, x_centered).astype(np.float32)

    # Compute normalized inner radius
    r_inner_normalized = inner_radius / R_outer

    # ANNULAR MASK: r must be between inner and outer radius
    # Points inside central hole are EXCLUDED from fitting
    annular_mask = (r >= r_inner_normalized) & (r <= 1.0) & ~np.isnan(values)

    if not np.any(annular_mask):
        raise RuntimeError(
            f"No valid points in annular region for Zernike fitting. "
            f"Inner radius = {inner_radius:.2f}, Outer radius = {R_outer:.2f}"
        )

    idx = np.nonzero(annular_mask)[0]
    m = int(n_modes)

    # Normal equations: (Z^T Z) c = Z^T v
    G = np.zeros((m, m), dtype=np.float64)
    h = np.zeros((m,), dtype=np.float64)
    v = values.astype(np.float64)

    for start in range(0, len(idx), chunk_size):
        chunk_idx = idx[start:start + chunk_size]
        r_chunk = r[chunk_idx]
        theta_chunk = theta[chunk_idx]
        v_chunk = v[chunk_idx]

        # Build Zernike basis for this chunk
        Z_chunk = np.column_stack([
            zernike_noll(j, r_chunk, theta_chunk).astype(np.float32)
            for j in range(1, m + 1)
        ])

        # Accumulate normal equations
        G += (Z_chunk.T @ Z_chunk).astype(np.float64)
        h += (Z_chunk.T @ v_chunk).astype(np.float64)

    # Solve normal equations
    try:
        coeffs = np.linalg.solve(G, h)
    except LinAlgError:
        coeffs = np.linalg.lstsq(G, h, rcond=None)[0]

    return coeffs, R_outer, r_inner_normalized, annular_mask


# ============================================================================
# Figure Geometry Parser
# ============================================================================

def parse_nastran_grid_large(line: str, continuation: str) -> Tuple[int, float, float, float]:
    """
    Parse a GRID* (large field format) card from Nastran BDF.

    Large field format (16-char fields):
    GRID*         ID              CP              X               Y               +
    *             Z               CD

    Args:
        line: First line of GRID* card
        continuation: Continuation line starting with *

    Returns:
        Tuple of (node_id, x, y, z)
    """
    # GRID* line has: GRID* | ID | CP | X | Y | continuation marker
    # Fields are 16 characters each after the 8-char GRID* identifier

    # Parse first line
    node_id = int(line[8:24].strip())
    # CP (coordinate system) at 24:40 - skip
    x = float(line[40:56].strip())
    y = float(line[56:72].strip())

    # Parse continuation line for Z
    # Continuation line: * | Z | CD
    z = float(continuation[8:24].strip())

    return node_id, x, y, z


def parse_nastran_grid_small(line: str) -> Tuple[int, float, float, float]:
    """
    Parse a GRID (small field format) card from Nastran BDF.

    Small field format (8-char fields):
    GRID    ID      CP      X       Y       Z       CD      PS      SEID

    Args:
        line: GRID card line

    Returns:
        Tuple of (node_id, x, y, z)
    """
    parts = line.split()
    node_id = int(parts[1])
    # parts[2] is CP
    x = float(parts[3]) if len(parts) > 3 else 0.0
    y = float(parts[4]) if len(parts) > 4 else 0.0
    z = float(parts[5]) if len(parts) > 5 else 0.0

    return node_id, x, y, z


def load_figure_geometry(figure_path: Union[str, Path]) -> Dict[int, np.ndarray]:
    """
    Load figure node geometry from a Nastran DAT/BDF file.

    Supports both GRID (small field) and GRID* (large field) formats.

    Args:
        figure_path: Path to figure.dat file

    Returns:
        Dict mapping node_id to (x, y, z) coordinates
    """
    figure_path = Path(figure_path)

    if not figure_path.exists():
        raise FileNotFoundError(f"Figure file not found: {figure_path}")

    geometry = {}

    with open(figure_path, 'r', encoding='utf-8', errors='ignore') as f:
        lines = f.readlines()

    i = 0
    while i < len(lines):
        line = lines[i]

        # Skip comments and empty lines
        if line.startswith('$') or line.strip() == '':
            i += 1
            continue

        # Large field format: GRID*
        if line.startswith('GRID*'):
            if i + 1 < len(lines):
                continuation = lines[i + 1]
                try:
                    node_id, x, y, z = parse_nastran_grid_large(line, continuation)
                    geometry[node_id] = np.array([x, y, z])
                except (ValueError, IndexError) as e:
                    warnings.warn(f"Failed to parse GRID* at line {i}: {e}")
                i += 2
                continue

        # Small field format: GRID
        elif line.startswith('GRID') and not line.startswith('GRID*'):
            try:
                node_id, x, y, z = parse_nastran_grid_small(line)
                geometry[node_id] = np.array([x, y, z])
            except (ValueError, IndexError) as e:
                warnings.warn(f"Failed to parse GRID at line {i}: {e}")
            i += 1
            continue

        # Check for END
        if 'ENDDATA' in line:
            break

        i += 1

    if not geometry:
        raise ValueError(f"No GRID cards found in {figure_path}")

    return geometry


def build_figure_interpolator(
    figure_geometry: Dict[int, np.ndarray],
    method: str = 'linear'
) -> LinearNDInterpolator:
    """
    Build a 2D interpolator for the figure surface z(x, y).

    Args:
        figure_geometry: Dict mapping node_id to (x, y, z)
        method: Interpolation method ('linear' or 'cubic')

    Returns:
        Interpolator function that takes (x, y) and returns z
    """
    # Extract coordinates
    points = np.array(list(figure_geometry.values()))
    x = points[:, 0]
    y = points[:, 1]
    z = points[:, 2]

    # Build interpolator
    xy_points = np.column_stack([x, y])

    if method == 'cubic':
        # Clough-Tocher gives smoother results but is slower
        interpolator = CloughTocher2DInterpolator(xy_points, z)
    else:
        # Linear is faster and usually sufficient
        interpolator = LinearNDInterpolator(xy_points, z)

    return interpolator


# ============================================================================
# Figure-Based OPD Calculation
# ============================================================================

def compute_figure_opd(
    x0: np.ndarray,
    y0: np.ndarray,
    z0: np.ndarray,
    dx: np.ndarray,
    dy: np.ndarray,
    dz: np.ndarray,
    figure_interpolator
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """
    Compute surface error using actual figure geometry as reference.

    The key insight: after lateral displacement, compare the deformed Z
    against what the ACTUAL figure surface Z is at that (x, y) position.

    Args:
        x0, y0, z0: Original node coordinates
        dx, dy, dz: Displacement components from FEA
        figure_interpolator: Interpolator for figure z(x, y)

    Returns:
        Tuple of:
        - x_def: Deformed X coordinates
        - y_def: Deformed Y coordinates
        - surface_error: Difference from ideal figure surface
        - lateral_magnitude: Magnitude of lateral displacement
    """
    # Deformed coordinates
    x_def = x0 + dx
    y_def = y0 + dy
    z_def = z0 + dz

    # Get ideal figure Z at deformed (x, y) positions
    z_figure_at_deformed = figure_interpolator(x_def, y_def)

    # Handle any NaN values (outside interpolation domain)
    nan_mask = np.isnan(z_figure_at_deformed)
    if np.any(nan_mask):
        # For points outside figure, use original Z as fallback
        z_figure_at_deformed[nan_mask] = z0[nan_mask]
        warnings.warn(
            f"{np.sum(nan_mask)} nodes outside figure interpolation domain, "
            "using original Z as fallback"
        )

    # Surface error = deformed Z - ideal figure Z at deformed position
    surface_error = z_def - z_figure_at_deformed

    # Lateral displacement magnitude
    lateral_magnitude = np.sqrt(dx**2 + dy**2)

    return x_def, y_def, surface_error, lateral_magnitude


# ============================================================================
# Main OPD Extractor Class (Most Rigorous Method)
# ============================================================================

class ZernikeOPDExtractor:
    """
    Zernike extractor using actual mesh geometry as the reference surface.

    THIS IS THE RECOMMENDED EXTRACTOR for telescope mirror optimization.

    This is the most rigorous approach for computing WFE because it:
    1. Uses the actual mesh geometry (not an analytical approximation)
    2. Accounts for lateral displacement via interpolation
    3. Works with any surface shape (parabola, hyperbola, asphere, freeform)
    4. No need to know focal length or optical prescription
    5. Supports annular apertures (mirrors with central holes)

    The extractor works in two modes:
    1. Default: Uses BDF geometry for nodes present in OP2 (RECOMMENDED)
    2. With figure_file: Uses explicit figure.dat for reference geometry

    For annular apertures (mirrors with central holes):
    - Set inner_radius to exclude the central obscuration from fitting
    - Example: M1 Mirror with 271.5mm diameter hole -> inner_radius=135.75

    Example:
        # Simple usage - BDF geometry filtered to OP2 nodes (RECOMMENDED)
        extractor = ZernikeOPDExtractor(op2_file)
        results = extractor.extract_all_subcases()

        # With annular aperture (central hole)
        extractor = ZernikeOPDExtractor(op2_file, inner_radius=135.75)
        results = extractor.extract_all_subcases()

        # With explicit figure file (advanced - ensure coordinates match!)
        extractor = ZernikeOPDExtractor(op2_file, figure_file='figure.dat')
    """

    def __init__(
        self,
        op2_path: Union[str, Path],
        figure_path: Optional[Union[str, Path]] = None,
        bdf_path: Optional[Union[str, Path]] = None,
        displacement_unit: str = 'mm',
        n_modes: int = DEFAULT_N_MODES,
        filter_orders: int = DEFAULT_FILTER_ORDERS,
        interpolation_method: str = 'linear',
        inner_radius: Optional[float] = None
    ):
        """
        Initialize figure-based Zernike extractor.

        Args:
            op2_path: Path to OP2 results file
            figure_path: Path to figure.dat with undeformed figure geometry (OPTIONAL)
                        If None, uses BDF geometry for nodes present in OP2
            bdf_path: Path to BDF geometry (auto-detected if None)
            displacement_unit: Unit of displacement in OP2
            n_modes: Number of Zernike modes to fit
            filter_orders: Number of low-order modes to filter for RMS
            interpolation_method: 'linear' or 'cubic' for figure interpolation
            inner_radius: Inner radius of central hole for annular apertures (same
                         units as geometry, typically mm). If None, full disk is used.
                         Example: For M1 Mirror with 271.5mm diameter hole, use 135.75
        """
        self.op2_path = Path(op2_path)
        self.figure_path = Path(figure_path) if figure_path else None
        self.bdf_path = Path(bdf_path) if bdf_path else find_geometry_file(self.op2_path)
        self.displacement_unit = displacement_unit
        self.n_modes = n_modes
        self.filter_orders = filter_orders
        self.interpolation_method = interpolation_method
        self.inner_radius = inner_radius
        self.use_explicit_figure = figure_path is not None
        self.use_annular = inner_radius is not None and inner_radius > 0

        # Unit conversion
        self.nm_scale = UNIT_TO_NM.get(displacement_unit.lower(), 1e6)
        self.um_scale = self.nm_scale / 1000.0
        self.wfe_factor = 2.0 * self.nm_scale  # WFE = 2 * surface error

        # Lazy-loaded data
        self._figure_geo = None
        self._figure_interp = None
        self._node_geo = None
        self._displacements = None

    @property
    def node_geometry(self) -> Dict[int, np.ndarray]:
        """Lazy-load FEM node geometry from BDF."""
        if self._node_geo is None:
            self._node_geo = read_node_geometry(self.bdf_path)
        return self._node_geo

    @property
    def displacements(self) -> Dict[str, Dict[str, np.ndarray]]:
        """Lazy-load displacements from OP2."""
        if self._displacements is None:
            self._displacements = extract_displacements_by_subcase(self.op2_path)
        return self._displacements

    @property
    def figure_geometry(self) -> Dict[int, np.ndarray]:
        """
        Lazy-load figure geometry.

        If explicit figure_path provided, load from that file.
        Otherwise, use BDF geometry filtered to nodes present in OP2.
        """
        if self._figure_geo is None:
            if self.use_explicit_figure and self.figure_path:
                # Load from explicit figure.dat
                self._figure_geo = load_figure_geometry(self.figure_path)
            else:
                # Use BDF geometry filtered to OP2 nodes
                # Get all node IDs from OP2 (any subcase)
                first_subcase = next(iter(self.displacements.values()))
                op2_node_ids = set(int(nid) for nid in first_subcase['node_ids'])

                # Filter BDF geometry to only OP2 nodes
                self._figure_geo = {
                    nid: geo for nid, geo in self.node_geometry.items()
                    if nid in op2_node_ids
                }
        return self._figure_geo

    @property
    def figure_interpolator(self):
        """Lazy-build figure interpolator."""
        if self._figure_interp is None:
            self._figure_interp = build_figure_interpolator(
                self.figure_geometry,
                method=self.interpolation_method
            )
        return self._figure_interp

    def _build_figure_opd_data(self, subcase_label: str) -> Dict[str, np.ndarray]:
        """
        Build OPD data using figure geometry as reference.

        Uses the figure geometry (either from explicit file or BDF filtered to OP2 nodes)
        as the undeformed reference surface.
        """
        if subcase_label not in self.displacements:
            available = list(self.displacements.keys())
            raise ValueError(f"Subcase '{subcase_label}' not found. Available: {available}")

        data = self.displacements[subcase_label]
        node_ids = data['node_ids']
        disp = data['disp']

        # Get figure node IDs
        figure_node_ids = set(self.figure_geometry.keys())

        # Build arrays - only for nodes in figure
        x0, y0, z0 = [], [], []
        dx_arr, dy_arr, dz_arr = [], [], []
        matched_ids = []

        for nid, vec in zip(node_ids, disp):
            nid_int = int(nid)

            # Check if node is in figure
            if nid_int not in figure_node_ids:
                continue

            # Use figure geometry as reference (undeformed surface)
            fig_geo = self.figure_geometry[nid_int]

            x0.append(fig_geo[0])
            y0.append(fig_geo[1])
            z0.append(fig_geo[2])
            dx_arr.append(vec[0])
            dy_arr.append(vec[1])
            dz_arr.append(vec[2])
            matched_ids.append(nid_int)

        if not x0:
            raise ValueError(
                f"No nodes matched between figure ({len(figure_node_ids)} nodes) "
                f"and displacement data ({len(node_ids)} nodes)"
            )

        x0 = np.array(x0)
        y0 = np.array(y0)
        z0 = np.array(z0)
        dx_arr = np.array(dx_arr)
        dy_arr = np.array(dy_arr)
        dz_arr = np.array(dz_arr)

        # Compute figure-based OPD
        x_def, y_def, surface_error, lateral_disp = compute_figure_opd(
            x0, y0, z0, dx_arr, dy_arr, dz_arr,
            self.figure_interpolator
        )

        # Convert to WFE
        wfe_nm = surface_error * self.wfe_factor

        return {
            'node_ids': np.array(matched_ids),
            'x_original': x0,
            'y_original': y0,
            'z_original': z0,
            'dx': dx_arr,
            'dy': dy_arr,
            'dz': dz_arr,
            'x_deformed': x_def,
            'y_deformed': y_def,
            'surface_error': surface_error,
            'wfe_nm': wfe_nm,
            'lateral_disp': lateral_disp,
            'n_figure_nodes': len(self.figure_geometry),
            'n_matched_nodes': len(matched_ids)
        }

    def extract_subcase(
        self,
        subcase_label: str,
        include_coefficients: bool = True,
        include_diagnostics: bool = True
    ) -> Dict[str, Any]:
        """
        Extract Zernike metrics using figure-based OPD method.

        Args:
            subcase_label: Subcase identifier
            include_coefficients: Include all Zernike coefficients
            include_diagnostics: Include lateral displacement diagnostics

        Returns:
            Dict with RMS metrics, coefficients, and diagnostics
        """
        opd_data = self._build_figure_opd_data(subcase_label)

        # Use deformed coordinates for Zernike fitting
        X = opd_data['x_deformed']
        Y = opd_data['y_deformed']
        WFE = opd_data['wfe_nm']

        # Fit Zernike coefficients (with or without annular masking)
        if self.use_annular:
            coeffs, R_outer, r_inner_norm, annular_mask = compute_zernike_coefficients_annular(
                X, Y, WFE, self.n_modes, self.inner_radius
            )
            R_max = R_outer
            n_annular = int(np.sum(annular_mask))
        else:
            coeffs, R_max = compute_zernike_coefficients(X, Y, WFE, self.n_modes)
            annular_mask = np.ones(len(X), dtype=bool)  # All points included
            r_inner_norm = 0.0
            n_annular = len(X)

        # Compute RMS metrics
        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)

        # Low-order Zernike basis
        Z_low = np.column_stack([
            zernike_noll(j, r, theta) for j in range(1, self.filter_orders + 1)
        ])

        # Filtered WFE (J5+)
        wfe_filtered = WFE - Z_low @ coeffs[:self.filter_orders]

        # J1-J3 filtered (for manufacturing/optician)
        Z_j1to3 = np.column_stack([
            zernike_noll(j, r, theta) for j in range(1, 4)
        ])
        wfe_j1to3 = WFE - Z_j1to3 @ coeffs[:3]

        # Compute RMS ONLY over annular region (or full disk if not annular)
        global_rms = float(np.sqrt(np.mean(WFE[annular_mask]**2)))
        filtered_rms = float(np.sqrt(np.mean(wfe_filtered[annular_mask]**2)))
        rms_j1to3 = float(np.sqrt(np.mean(wfe_j1to3[annular_mask]**2)))

        # Aberration magnitudes
        aberrations = compute_aberration_magnitudes(coeffs)

        result = {
            'subcase': subcase_label,
            'method': 'figure_opd_annular' if self.use_annular else 'figure_opd',
            'rms_wfe_nm': filtered_rms,  # Primary metric (J5+)
            'global_rms_nm': global_rms,
            'filtered_rms_nm': filtered_rms,
            'rms_filter_j1to3_nm': rms_j1to3,
            'n_nodes': len(X),
            'n_annular_nodes': n_annular,
            'n_figure_nodes': opd_data['n_figure_nodes'],
            'figure_file': str(self.figure_path.name) if self.figure_path else 'BDF (filtered to OP2)',
            **aberrations
        }

        # Add annular aperture info
        if self.use_annular:
            result.update({
                'inner_radius': self.inner_radius,
                'outer_radius': R_max,
                'obscuration_ratio': r_inner_norm,
            })

        if include_diagnostics:
            lateral = opd_data['lateral_disp']
            # Compute lateral displacement stats over annular region
            result.update({
                'max_lateral_displacement_um': float(np.max(lateral[annular_mask]) * self.um_scale),
                'rms_lateral_displacement_um': float(np.sqrt(np.mean(lateral[annular_mask]**2)) * self.um_scale),
                'mean_lateral_displacement_um': float(np.mean(lateral[annular_mask]) * self.um_scale),
            })

        if include_coefficients:
            result['coefficients'] = coeffs.tolist()
            result['coefficient_labels'] = [
                zernike_label(j) for j in range(1, self.n_modes + 1)
            ]

        return result

    def extract_all_subcases(self) -> Dict[int, Dict[str, Any]]:
        """Extract metrics for all subcases."""
        results = {}
        for label in self.displacements.keys():
            try:
                # Convert label to int if possible for consistent keys
                key = int(label) if label.isdigit() else label
                results[key] = self.extract_subcase(label)
            except Exception as e:
                warnings.warn(f"Failed to extract subcase {label}: {e}")
        return results

    def extract_relative(
        self,
        target_subcase: str,
        reference_subcase: str,
        include_coefficients: bool = False
    ) -> Dict[str, Any]:
        """
        Extract Zernike metrics relative to a reference subcase using OPD method.

        Computes: WFE_relative = WFE_target(node_i) - WFE_reference(node_i)
        for each node, then fits Zernike to the difference field.

        This is the CORRECT way to compute relative WFE for optimization.
        It properly accounts for lateral (X,Y) displacement via OPD interpolation.

        For annular apertures, the central hole is excluded from fitting and RMS.

        Args:
            target_subcase: Subcase to analyze (e.g., "3" for 40 deg)
            reference_subcase: Reference subcase to subtract (e.g., "2" for 20 deg)
            include_coefficients: Whether to include all Zernike coefficients

        Returns:
            Dict with relative metrics: relative_filtered_rms_nm, relative_rms_filter_j1to3, etc.
        """
        # Build OPD data for both subcases
        target_data = self._build_figure_opd_data(target_subcase)
        ref_data = self._build_figure_opd_data(reference_subcase)

        # Build node ID -> WFE map for reference
        ref_node_to_wfe = {
            int(nid): wfe for nid, wfe in zip(ref_data['node_ids'], ref_data['wfe_nm'])
        }

        # Compute node-by-node relative WFE for common nodes
        X_rel, Y_rel, WFE_rel = [], [], []
        lateral_rel = []

        for i, nid in enumerate(target_data['node_ids']):
            nid = int(nid)
            if nid not in ref_node_to_wfe:
                continue

            # Use target's deformed coordinates for Zernike fitting
            X_rel.append(target_data['x_deformed'][i])
            Y_rel.append(target_data['y_deformed'][i])

            # Relative WFE = target WFE - reference WFE
            target_wfe = target_data['wfe_nm'][i]
            ref_wfe = ref_node_to_wfe[nid]
            WFE_rel.append(target_wfe - ref_wfe)

            lateral_rel.append(target_data['lateral_disp'][i])

        X_rel = np.array(X_rel)
        Y_rel = np.array(Y_rel)
        WFE_rel = np.array(WFE_rel)
        lateral_rel = np.array(lateral_rel)

        if len(X_rel) == 0:
            raise ValueError(f"No common nodes between subcases {target_subcase} and {reference_subcase}")

        # Fit Zernike coefficients to the relative WFE (with or without annular masking)
        if self.use_annular:
            coeffs, R_outer, r_inner_norm, annular_mask = compute_zernike_coefficients_annular(
                X_rel, Y_rel, WFE_rel, self.n_modes, self.inner_radius
            )
            R_max = R_outer
            n_annular = int(np.sum(annular_mask))
        else:
            coeffs, R_max = compute_zernike_coefficients(X_rel, Y_rel, WFE_rel, self.n_modes)
            annular_mask = np.ones(len(X_rel), dtype=bool)
            r_inner_norm = 0.0
            n_annular = len(X_rel)

        # Compute RMS metrics
        x_c = X_rel - np.mean(X_rel)
        y_c = Y_rel - np.mean(Y_rel)
        r = np.hypot(x_c / R_max, y_c / R_max)
        theta = np.arctan2(y_c, x_c)

        # Low-order Zernike basis (for filtering)
        Z_low = np.column_stack([
            zernike_noll(j, r, theta) for j in range(1, self.filter_orders + 1)
        ])

        # Filtered WFE (J5+) - this is the primary optimization metric
        wfe_filtered = WFE_rel - Z_low @ coeffs[:self.filter_orders]

        # J1-J3 filtered (for manufacturing/optician workload)
        Z_j1to3 = np.column_stack([
            zernike_noll(j, r, theta) for j in range(1, 4)
        ])
        wfe_j1to3 = WFE_rel - Z_j1to3 @ coeffs[:3]

        # Compute RMS ONLY over annular region (or full disk if not annular)
        global_rms = float(np.sqrt(np.mean(WFE_rel[annular_mask]**2)))
        filtered_rms = float(np.sqrt(np.mean(wfe_filtered[annular_mask]**2)))
        rms_j1to3 = float(np.sqrt(np.mean(wfe_j1to3[annular_mask]**2)))

        # Aberration magnitudes
        aberrations = compute_aberration_magnitudes(coeffs)

        result = {
            'target_subcase': target_subcase,
            'reference_subcase': reference_subcase,
            'method': 'figure_opd_relative_annular' if self.use_annular else 'figure_opd_relative',
            'relative_global_rms_nm': global_rms,
            'relative_filtered_rms_nm': filtered_rms,
            'relative_rms_filter_j1to3': rms_j1to3,
            'n_common_nodes': len(X_rel),
            'n_annular_nodes': n_annular,
            'max_lateral_displacement_um': float(np.max(lateral_rel[annular_mask]) * self.um_scale),
            'rms_lateral_displacement_um': float(np.sqrt(np.mean(lateral_rel[annular_mask]**2)) * self.um_scale),
            **{f'relative_{k}': v for k, v in aberrations.items()}
        }

        # Add annular aperture info
        if self.use_annular:
            result.update({
                'inner_radius': self.inner_radius,
                'outer_radius': R_max,
                'obscuration_ratio': r_inner_norm,
            })

        if include_coefficients:
            result['coefficients'] = coeffs.tolist()
            result['coefficient_labels'] = [
                zernike_label(j) for j in range(1, self.n_modes + 1)
            ]

        return result

    def extract_comparison(self, subcase_label: str) -> Dict[str, Any]:
        """
        Compare figure-based method vs standard Z-only method.
        """
        from .extract_zernike_wfe import ZernikeExtractor

        # Figure-based (this extractor)
        figure_result = self.extract_subcase(subcase_label)

        # Standard Z-only
        standard_extractor = ZernikeExtractor(self.op2_path, self.bdf_path)
        standard_result = standard_extractor.extract_subcase(subcase_label)

        # Compute deltas
        delta_filtered = figure_result['filtered_rms_nm'] - standard_result['filtered_rms_nm']
        pct_diff = 100.0 * delta_filtered / standard_result['filtered_rms_nm'] if standard_result['filtered_rms_nm'] > 0 else 0

        return {
            'subcase': subcase_label,
            'figure_method': {
                'filtered_rms_nm': figure_result['filtered_rms_nm'],
                'global_rms_nm': figure_result['global_rms_nm'],
                'n_nodes': figure_result['n_nodes'],
            },
            'standard_method': {
                'filtered_rms_nm': standard_result['filtered_rms_nm'],
                'global_rms_nm': standard_result['global_rms_nm'],
                'n_nodes': standard_result['n_nodes'],
            },
            'delta': {
                'filtered_rms_nm': delta_filtered,
                'percent_difference': pct_diff,
            },
            'lateral_displacement': {
                'max_um': figure_result.get('max_lateral_displacement_um', 0),
                'rms_um': figure_result.get('rms_lateral_displacement_um', 0),
            },
            'figure_file': str(self.figure_path.name) if self.figure_path else 'BDF (filtered to OP2)',
        }


# ============================================================================
# Convenience Functions
# ============================================================================

def extract_zernike_opd(
    op2_file: Union[str, Path],
    figure_file: Optional[Union[str, Path]] = None,
    subcase: Union[int, str] = 1,
    inner_radius: Optional[float] = None,
    **kwargs
) -> Dict[str, Any]:
    """
    Convenience function for OPD-based Zernike extraction (most rigorous method).

    THIS IS THE RECOMMENDED FUNCTION for telescope mirror optimization.

    Args:
        op2_file: Path to OP2 results
        figure_file: Path to explicit figure.dat (uses BDF geometry if None - RECOMMENDED)
        subcase: Subcase identifier
        inner_radius: Inner radius of central hole for annular apertures (mm).
                     If None, full disk is used.
                     Example: For M1 Mirror with 271.5mm diameter hole, use 135.75
        **kwargs: Additional args for ZernikeOPDExtractor

    Returns:
        Dict with Zernike metrics, aberrations, and lateral displacement info
    """
    extractor = ZernikeOPDExtractor(
        op2_file,
        figure_path=figure_file,
        inner_radius=inner_radius,
        **kwargs
    )
    return extractor.extract_subcase(str(subcase))


def extract_zernike_opd_filtered_rms(
    op2_file: Union[str, Path],
    subcase: Union[int, str] = 1,
    inner_radius: Optional[float] = None,
    **kwargs
) -> float:
    """
    Extract filtered RMS WFE using OPD method (most rigorous).

    Primary metric for mirror optimization.

    Args:
        op2_file: Path to OP2 results
        subcase: Subcase identifier
        inner_radius: Inner radius of central hole for annular apertures (mm).
                     If None, full disk is used.
        **kwargs: Additional args for ZernikeOPDExtractor

    Returns:
        Filtered RMS WFE in nanometers
    """
    result = extract_zernike_opd(op2_file, subcase=subcase, inner_radius=inner_radius, **kwargs)
    return result['filtered_rms_nm']


# Backwards compatibility aliases
ZernikeFigureExtractor = ZernikeOPDExtractor  # Deprecated: use ZernikeOPDExtractor
extract_zernike_figure = extract_zernike_opd  # Deprecated: use extract_zernike_opd
extract_zernike_figure_rms = extract_zernike_opd_filtered_rms  # Deprecated


# ============================================================================
# Module Exports
# ============================================================================

__all__ = [
    # Primary exports (new names)
    'ZernikeOPDExtractor',
    'extract_zernike_opd',
    'extract_zernike_opd_filtered_rms',

    # Annular aperture support
    'compute_zernike_coefficients_annular',

    # Utility functions
    'load_figure_geometry',
    'build_figure_interpolator',
    'compute_figure_opd',

    # Backwards compatibility (deprecated)
    'ZernikeFigureExtractor',
    'extract_zernike_figure',
    'extract_zernike_figure_rms',
]


if __name__ == '__main__':
    import sys

    if len(sys.argv) > 1:
        op2_file = Path(sys.argv[1])
        figure_file = Path(sys.argv[2]) if len(sys.argv) > 2 else None

        print(f"Analyzing: {op2_file}")
        print("=" * 60)

        try:
            extractor = ZernikeFigureExtractor(op2_file, figure_path=figure_file)
            print(f"Figure file: {extractor.figure_path}")
            print(f"Figure nodes: {len(extractor.figure_geometry)}")
            print()

            for label in extractor.displacements.keys():
                print(f"\n{'=' * 60}")
                print(f"Subcase {label}")
                print('=' * 60)

                comparison = extractor.extract_comparison(label)

                print(f"\nStandard (Z-only) method:")
                print(f"  Filtered RMS: {comparison['standard_method']['filtered_rms_nm']:.2f} nm")
                print(f"  Nodes: {comparison['standard_method']['n_nodes']}")

                print(f"\nFigure-based OPD method:")
                print(f"  Filtered RMS: {comparison['figure_method']['filtered_rms_nm']:.2f} nm")
                print(f"  Nodes: {comparison['figure_method']['n_nodes']}")

                print(f"\nDifference:")
                print(f"  Delta: {comparison['delta']['filtered_rms_nm']:+.2f} nm ({comparison['delta']['percent_difference']:+.1f}%)")

                print(f"\nLateral Displacement:")
                print(f"  Max: {comparison['lateral_displacement']['max_um']:.3f} um")
                print(f"  RMS: {comparison['lateral_displacement']['rms_um']:.3f} um")

        except Exception as e:
            print(f"Error: {e}")
            import traceback
            traceback.print_exc()
            sys.exit(1)
    else:
        print("Figure-Based Zernike Extractor")
        print("=" * 40)
        print("\nUsage: python extract_zernike_figure.py <op2_file> [figure.dat]")
        print("\nThis extractor uses actual figure geometry instead of parabola approximation.")