Atomizer/tools/test_zernike_recentering.py

"""
Test: Is recentering Z per subcase equivalent to removing piston from relative WFE?

Theory:
- Option A: Recenter each subcase, then subtract
  WFE_rel_A = (dZ_20 - mean(dZ_20)) - (dZ_90 - mean(dZ_90))

- Option B: Subtract raw, then remove piston via Zernike J1
  WFE_rel_B = dZ_20 - dZ_90
  Then filter J1 from Zernike fit

Mathematically:
WFE_rel_A = dZ_20 - dZ_90 - mean(dZ_20) + mean(dZ_90)
          = (dZ_20 - dZ_90) - (mean(dZ_20) - mean(dZ_90))

If nodes are identical: mean(dZ_20) - mean(dZ_90) = mean(dZ_20 - dZ_90)
So: WFE_rel_A = WFE_rel_B - mean(WFE_rel_B)

This should equal WFE_rel_B with J1 (piston) removed, BUT only if:
1. Piston Zernike Z1 = 1 (constant)
2. Sampling is uniform (or Z1 coefficient = mean for non-uniform)

For annular apertures and non-uniform mesh sampling, this might not be exact!
"""

import numpy as np
from pathlib import Path
import sys

# Add project root to path
sys.path.insert(0, str(Path(__file__).parent.parent))

from optimization_engine.extractors.extract_zernike import (
    ZernikeExtractor,
    compute_zernike_coefficients,
    compute_rms_metrics,
    zernike_noll,
    DEFAULT_N_MODES,
    DEFAULT_FILTER_ORDERS,
)


def test_recentering_equivalence_synthetic():
    """Test with synthetic data where we control everything."""
    print("=" * 70)
    print("TEST 1: Synthetic Data")
    print("=" * 70)

    np.random.seed(42)

    # Create annular aperture (like a telescope mirror)
    n_points = 5000
    r_inner = 0.3  # 30% central obscuration
    r_outer = 1.0

    # Generate random points in annulus
    r = np.sqrt(np.random.uniform(r_inner**2, r_outer**2, n_points))
    theta = np.random.uniform(0, 2*np.pi, n_points)
    X = r * np.cos(theta) * 500  # Scale to 500mm radius
    Y = r * np.sin(theta) * 500

    # Simulate Z-displacements for two subcases (in mm, will convert to nm)
    # Subcase 90: Some aberration pattern + piston
    piston_90 = 0.001  # 1 micron mean displacement
    dZ_90 = piston_90 + 0.0005 * (X**2 + Y**2) / 500**2  # Add some power
    dZ_90 += 0.0002 * np.random.randn(n_points)  # Add noise

    # Subcase 20: Different aberration + different piston
    piston_20 = 0.003  # 3 micron mean displacement
    dZ_20 = piston_20 + 0.0008 * (X**2 + Y**2) / 500**2  # More power
    dZ_20 += 0.0003 * X / 500  # Add some tilt
    dZ_20 += 0.0002 * np.random.randn(n_points)  # Add noise

    # Convert to WFE in nm (2x for reflection, 1e6 for mm->nm)
    wfe_factor = 2.0 * 1e6
    WFE_90 = dZ_90 * wfe_factor
    WFE_20 = dZ_20 * wfe_factor

    print(f"\nInput data:")
    print(f"  Points: {n_points} (annular, r_inner={r_inner})")
    print(f"  Mean WFE_90: {np.mean(WFE_90):.2f} nm")
    print(f"  Mean WFE_20: {np.mean(WFE_20):.2f} nm")
    print(f"  Mean difference: {np.mean(WFE_20) - np.mean(WFE_90):.2f} nm")

    # =========================================================================
    # Option A: Recenter each subcase BEFORE subtraction
    # =========================================================================
    WFE_90_centered = WFE_90 - np.mean(WFE_90)
    WFE_20_centered = WFE_20 - np.mean(WFE_20)
    WFE_rel_A = WFE_20_centered - WFE_90_centered

    # Fit Zernike to option A
    coeffs_A, R_max_A = compute_zernike_coefficients(X, Y, WFE_rel_A, DEFAULT_N_MODES)

    # Compute RMS metrics for option A
    rms_A = compute_rms_metrics(X, Y, WFE_rel_A, DEFAULT_N_MODES, DEFAULT_FILTER_ORDERS)

    print(f"\nOption A (recenter before subtraction):")
    print(f"  Mean WFE_rel_A: {np.mean(WFE_rel_A):.6f} nm (should be ~0)")
    print(f"  J1 (piston) coefficient: {coeffs_A[0]:.6f} nm")
    print(f"  Global RMS: {rms_A['global_rms_nm']:.2f} nm")
    print(f"  Filtered RMS (J1-J4 removed): {rms_A['filtered_rms_nm']:.2f} nm")

    # =========================================================================
    # Option B: Subtract raw, then analyze (current implementation)
    # =========================================================================
    WFE_rel_B = WFE_20 - WFE_90  # No recentering

    # Fit Zernike to option B
    coeffs_B, R_max_B = compute_zernike_coefficients(X, Y, WFE_rel_B, DEFAULT_N_MODES)

    # Compute RMS metrics for option B
    rms_B = compute_rms_metrics(X, Y, WFE_rel_B, DEFAULT_N_MODES, DEFAULT_FILTER_ORDERS)

    print(f"\nOption B (current: subtract raw, filter J1-J4 after):")
    print(f"  Mean WFE_rel_B: {np.mean(WFE_rel_B):.2f} nm")
    print(f"  J1 (piston) coefficient: {coeffs_B[0]:.2f} nm")
    print(f"  Global RMS: {rms_B['global_rms_nm']:.2f} nm")
    print(f"  Filtered RMS (J1-J4 removed): {rms_B['filtered_rms_nm']:.2f} nm")

    # =========================================================================
    # Compare
    # =========================================================================
    print(f"\n" + "=" * 70)
    print("COMPARISON")
    print("=" * 70)

    # Check if J1 coefficient in B equals the mean
    j1_vs_mean_diff = abs(coeffs_B[0] - np.mean(WFE_rel_B))
    print(f"\n1. Does J1 = mean(WFE)?")
    print(f"   J1 coefficient: {coeffs_B[0]:.6f} nm")
    print(f"   Mean of WFE:    {np.mean(WFE_rel_B):.6f} nm")
    print(f"   Difference:     {j1_vs_mean_diff:.6f} nm")
    print(f"   --> {'YES (negligible diff)' if j1_vs_mean_diff < 0.01 else 'NO (significant diff!)'}")

    # Check if filtered RMS is the same
    filtered_rms_diff = abs(rms_A['filtered_rms_nm'] - rms_B['filtered_rms_nm'])
    print(f"\n2. Is filtered RMS the same?")
    print(f"   Option A filtered RMS: {rms_A['filtered_rms_nm']:.6f} nm")
    print(f"   Option B filtered RMS: {rms_B['filtered_rms_nm']:.6f} nm")
    print(f"   Difference:            {filtered_rms_diff:.6f} nm")
    print(f"   --> {'EQUIVALENT' if filtered_rms_diff < 0.01 else 'NOT EQUIVALENT!'}")

    # Check all coefficients (J2 onwards should be identical)
    coeff_diffs = np.abs(coeffs_A[1:] - coeffs_B[1:])  # Skip J1
    max_coeff_diff = np.max(coeff_diffs)
    print(f"\n3. Are non-piston coefficients (J2+) identical?")
    print(f"   Max difference in J2-J{DEFAULT_N_MODES}: {max_coeff_diff:.6f} nm")
    print(f"   --> {'YES' if max_coeff_diff < 0.01 else 'NO!'}")

    # The key insight: for non-uniform sampling, J1 might not equal mean exactly
    # Let's check how Z1 (piston polynomial) behaves
    x_c = X - np.mean(X)
    y_c = Y - np.mean(Y)
    R_max = np.max(np.hypot(x_c, y_c))
    r_norm = np.hypot(x_c / R_max, y_c / R_max)
    theta_norm = np.arctan2(y_c, x_c)

    Z1 = zernike_noll(1, r_norm, theta_norm)  # Piston polynomial
    print(f"\n4. Piston polynomial Z1 statistics:")
    print(f"   Z1 should be constant = 1 for all points")
    print(f"   Mean(Z1):   {np.mean(Z1):.6f}")
    print(f"   Std(Z1):    {np.std(Z1):.6f}")
    print(f"   Min(Z1):    {np.min(Z1):.6f}")
    print(f"   Max(Z1):    {np.max(Z1):.6f}")

    return filtered_rms_diff < 0.01


def test_recentering_with_real_data():
    """Test with real OP2 data if available."""
    print("\n" + "=" * 70)
    print("TEST 2: Real Data (if available)")
    print("=" * 70)

    # Look for a real study with OP2 data
    studies_dir = Path(__file__).parent.parent / "studies"
    op2_files = list(studies_dir.rglob("*.op2"))

    if not op2_files:
        print("  No OP2 files found in studies directory. Skipping real data test.")
        return None

    # Use the first one found
    op2_file = op2_files[0]
    print(f"\n  Using: {op2_file.relative_to(studies_dir.parent)}")

    try:
        extractor = ZernikeExtractor(op2_file)
        subcases = list(extractor.displacements.keys())
        print(f"  Available subcases: {subcases}")

        if len(subcases) < 2:
            print("  Need at least 2 subcases for relative comparison. Skipping.")
            return None

        # Pick two subcases
        ref_subcase = subcases[0]
        target_subcase = subcases[1]
        print(f"  Comparing: {target_subcase} vs {ref_subcase}")

        # Get raw data
        X_t, Y_t, WFE_t = extractor._build_dataframe(target_subcase)
        X_r, Y_r, WFE_r = extractor._build_dataframe(ref_subcase)

        # Build node matching (same logic as extract_relative)
        target_data = extractor.displacements[target_subcase]
        ref_data = extractor.displacements[ref_subcase]

        ref_node_to_idx = {
            int(nid): i for i, nid in enumerate(ref_data['node_ids'])
        }

        X_common, Y_common = [], []
        WFE_target_common, WFE_ref_common = [], []

        for i, nid in enumerate(target_data['node_ids']):
            nid = int(nid)
            if nid not in ref_node_to_idx:
                continue
            ref_idx = ref_node_to_idx[nid]
            geo = extractor.node_geometry.get(nid)
            if geo is None:
                continue

            X_common.append(geo[0])
            Y_common.append(geo[1])
            WFE_target_common.append(target_data['disp'][i, 2] * extractor.wfe_factor)
            WFE_ref_common.append(ref_data['disp'][ref_idx, 2] * extractor.wfe_factor)

        X = np.array(X_common)
        Y = np.array(Y_common)
        WFE_target = np.array(WFE_target_common)
        WFE_ref = np.array(WFE_ref_common)

        print(f"  Common nodes: {len(X)}")
        print(f"  Mean WFE target ({target_subcase}): {np.mean(WFE_target):.2f} nm")
        print(f"  Mean WFE ref ({ref_subcase}): {np.mean(WFE_ref):.2f} nm")

        # Option A: Recenter before
        WFE_target_centered = WFE_target - np.mean(WFE_target)
        WFE_ref_centered = WFE_ref - np.mean(WFE_ref)
        WFE_rel_A = WFE_target_centered - WFE_ref_centered

        rms_A = compute_rms_metrics(X, Y, WFE_rel_A, DEFAULT_N_MODES, DEFAULT_FILTER_ORDERS)

        # Option B: Current (no recentering)
        WFE_rel_B = WFE_target - WFE_ref
        rms_B = compute_rms_metrics(X, Y, WFE_rel_B, DEFAULT_N_MODES, DEFAULT_FILTER_ORDERS)

        print(f"\n  Option A (recenter before):")
        print(f"    Filtered RMS: {rms_A['filtered_rms_nm']:.4f} nm")

        print(f"\n  Option B (current, filter after):")
        print(f"    Filtered RMS: {rms_B['filtered_rms_nm']:.4f} nm")

        diff = abs(rms_A['filtered_rms_nm'] - rms_B['filtered_rms_nm'])
        pct_diff = 100 * diff / rms_B['filtered_rms_nm'] if rms_B['filtered_rms_nm'] > 0 else 0

        print(f"\n  Difference: {diff:.4f} nm ({pct_diff:.4f}%)")
        print(f"  --> {'EQUIVALENT' if pct_diff < 0.1 else 'NOT EQUIVALENT!'}")

        return pct_diff < 0.1

    except Exception as e:
        print(f"  Error: {e}")
        import traceback
        traceback.print_exc()
        return None


def test_annular_zernike_piston():
    """
    Specifically test whether J1 = mean for annular apertures.

    For a filled circular aperture with uniform sampling, J1 = mean exactly.
    For annular apertures or non-uniform sampling, this may not hold!
    """
    print("\n" + "=" * 70)
    print("TEST 3: Annular Aperture - Does J1 = mean?")
    print("=" * 70)

    np.random.seed(123)

    # Test different inner radii
    for r_inner in [0.0, 0.2, 0.4, 0.6]:
        n_points = 10000
        r_outer = 1.0

        # Generate points in annulus
        if r_inner > 0:
            r = np.sqrt(np.random.uniform(r_inner**2, r_outer**2, n_points))
        else:
            r = np.sqrt(np.random.uniform(0, r_outer**2, n_points))
        theta = np.random.uniform(0, 2*np.pi, n_points)
        X = r * np.cos(theta) * 500
        Y = r * np.sin(theta) * 500

        # Create WFE with known piston
        true_piston = 1000.0  # nm
        WFE = true_piston + 100 * np.random.randn(n_points)

        # Fit Zernike
        coeffs, _ = compute_zernike_coefficients(X, Y, WFE, DEFAULT_N_MODES)

        j1_coeff = coeffs[0]
        wfe_mean = np.mean(WFE)
        diff = abs(j1_coeff - wfe_mean)

        print(f"\n  r_inner = {r_inner}:")
        print(f"    True piston:   {true_piston:.2f} nm")
        print(f"    Mean(WFE):     {wfe_mean:.2f} nm")
        print(f"    J1 coefficient: {j1_coeff:.2f} nm")
        print(f"    |J1 - mean|:   {diff:.4f} nm")
        print(f"    --> {'J1 ≈ mean' if diff < 1.0 else 'J1 ≠ mean!'}")


if __name__ == "__main__":
    print("\n" + "=" * 70)
    print("ZERNIKE RECENTERING EQUIVALENCE TEST")
    print("=" * 70)
    print("\nQuestion: Is recentering Z per subcase equivalent to")
    print("          removing piston (J1) from relative WFE?")
    print("=" * 70)

    # Run tests
    test1_passed = test_recentering_equivalence_synthetic()
    test2_passed = test_recentering_with_real_data()
    test_annular_zernike_piston()

    # Summary
    print("\n" + "=" * 70)
    print("SUMMARY")
    print("=" * 70)
    print(f"\nSynthetic data test: {'PASSED' if test1_passed else 'FAILED'}")
    if test2_passed is not None:
        print(f"Real data test:      {'PASSED' if test2_passed else 'FAILED'}")
    else:
        print("Real data test:      SKIPPED")

    print("\nConclusion:")
    if test1_passed:
        print("  For standard Zernike on circular/annular apertures,")
        print("  recentering Z before subtraction IS equivalent to")
        print("  filtering J1 (piston) after Zernike fit.")
        print("\n  The current implementation is correct!")
    else:
        print("  WARNING: Recentering and J1 filtering are NOT equivalent!")
        print("  Consider updating the extractor to recenter before subtraction.")