optics.py

# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""
Optical lithography simulation using CuPy FFT or cuTile FFT.

Implements the Hopkins/SOCS (Sum of Coherent Systems) model:
    Aerial(x,y) = Σᵢ wᵢ |Mask ⊛ Kernelᵢ|²

where:
    - Mask is the binary input pattern
    - Kernelᵢ are the optical kernels (point spread functions)
    - wᵢ are the kernel weights (scales)
    - ⊛ denotes convolution (implemented via FFT)
"""

import numpy as np
import cupy as cp
import torch
from pathlib import Path

# Optional import of cuTile FFT
try:
    from test_fft import cutile_fft_2d
    CUTILE_AVAILABLE = True
except ImportError:
    CUTILE_AVAILABLE = False


# =============================================================================
# cuTile FFT 2D wrappers
# =============================================================================

# FFT configuration for common sizes
FFT_CONFIG = {
    2048: {"factors": (8, 16, 16), "atom_packing_dim": 64},
    1024: {"factors": (8, 8, 16), "atom_packing_dim": 64},
    512: {"factors": (8, 8, 8), "atom_packing_dim": 32},
    256: {"factors": (4, 8, 8), "atom_packing_dim": 16},
}


def cutile_ifft_2d(
    x: torch.Tensor,
    factors: tuple[int, int, int],
    atom_packing_dim: int,
    chunk_size: int = 8,
) -> torch.Tensor:
    """
    Perform 2D IFFT using the conjugate trick: IFFT(x) = conj(FFT(conj(x))) / N

    Args:
        x: Input tensor of shape (H, W) with complex64 dtype on CUDA.
        factors: FFT factors (F0, F1, F2).
        atom_packing_dim: Packing dimension for cuTile kernel.
        chunk_size: Chunk size for batched processing.

    Returns:
        2D IFFT result of same shape.
    """
    H, W = x.shape
    # resolve_conj() materializes the conjugate so view_as_real() works
    x_conj = torch.conj(x).resolve_conj()
    fft_result = cutile_fft_2d(x_conj, factors, atom_packing_dim, chunk_size)
    # resolve_conj() again for the final conjugate
    return torch.conj(fft_result).resolve_conj() / (H * W)


def load_scales(path: str | Path, device: str = "cpu") -> torch.Tensor:
    """
    Load kernel scale factors from a PyTorch file.

    Args:
        path: Path to the .pt file containing scales.
        device: Device to load tensors to.

    Returns:
        1D tensor of shape (num_kernels,).
    """
    return torch.load(str(path), map_location=device, weights_only=False)


def kernel_multiply_freq(
    mask_fft: cp.ndarray,
    kernels: cp.ndarray,
    num_kernels: int,
) -> cp.ndarray:
    """
    Multiply mask FFT with kernels using OpenILT's quadrant mapping.

    The kernels are 35x35 spatial-domain PSFs centered at (17,17).
    This function implements the frequency-domain multiplication by
    mapping kernel quadrants to the corners of the FFT where low
    frequencies reside.

    Args:
        mask_fft: FFT of mask, shape (H, W) complex.
        kernels: Spatial-domain kernels, shape (num_kernels, kh, kw) complex.
        num_kernels: Number of kernels to use.

    Returns:
        Complex array of shape (num_kernels, H, W) - convolution results in freq domain.
    """
    kh, kw = kernels.shape[1], kernels.shape[2]
    knxh, knyh = kh // 2, kw // 2  # 17 for 35x35 kernels
    H, W = mask_fft.shape

    # Output array
    output = cp.zeros((num_kernels, H, W), dtype=cp.complex64)

    # Map kernel quadrants to FFT corners (OpenILT's approach)
    # The kernel is centered at (knxh, knyh). To do convolution in freq domain,
    # we need to account for the circular shift that FFT implies.
    #
    # FFT corners: top-left = low freq, top-right = neg y freq, etc.
    # Kernel quadrants need to be mapped accordingly:
    #   output[top-left] = mask_fft[top-left] * kernel[bottom-right]
    #   output[top-right] = mask_fft[top-right] * kernel[bottom-left]
    #   output[bottom-left] = mask_fft[bottom-left] * kernel[top-right]
    #   output[bottom-right] = mask_fft[bottom-right] * kernel[top-left]

    k = kernels[:num_kernels]  # (num_kernels, kh, kw)

    # Top-left corner of output: indices [0:knxh+1, 0:knyh+1]
    # corresponds to kernel bottom-right: [-(knxh+1):, -(knyh+1):]
    output[:, :knxh+1, :knyh+1] = mask_fft[None, :knxh+1, :knyh+1] * k[:, -(knxh+1):, -(knyh+1):]

    # Top-right corner of output: indices [0:knxh+1, -knyh:]
    # corresponds to kernel bottom-left: [-(knxh+1):, :knyh]
    output[:, :knxh+1, -knyh:] = mask_fft[None, :knxh+1, -knyh:] * k[:, -(knxh+1):, :knyh]

    # Bottom-left corner of output: indices [-knxh:, 0:knyh+1]
    # corresponds to kernel top-right: [:knxh, -(knyh+1):]
    output[:, -knxh:, :knyh+1] = mask_fft[None, -knxh:, :knyh+1] * k[:, :knxh, -(knyh+1):]

    # Bottom-right corner of output: indices [-knxh:, -knyh:]
    # corresponds to kernel top-left: [:knxh, :knyh]
    output[:, -knxh:, -knyh:] = mask_fft[None, -knxh:, -knyh:] * k[:, :knxh, :knyh]

    return output


def apply_optical_model(
    mask: np.ndarray | cp.ndarray,
    kernels: np.ndarray | cp.ndarray,
    scales: np.ndarray | cp.ndarray,
    num_kernels: int | None = None,
) -> cp.ndarray:
    """
    Apply the Hopkins/SOCS optical model using FFT convolution.

    Computes: Aerial(x,y) = Σᵢ wᵢ |Mask ⊛ Kernelᵢ|²

    Uses OpenILT-compatible normalization and kernel multiplication.

    Args:
        mask: 2D binary mask array of shape (H, W).
        kernels: Complex kernels of shape (num_kernels, kh, kw) in spatial domain.
        scales: Kernel weights of shape (num_kernels,).
        num_kernels: Number of kernels to use (default: all).

    Returns:
        2D aerial image array of shape (H, W).
    """
    # Move to GPU if needed
    if isinstance(mask, np.ndarray):
        mask = cp.asarray(mask)
    if isinstance(kernels, np.ndarray):
        kernels = cp.asarray(kernels)
    if isinstance(scales, np.ndarray):
        scales = cp.asarray(scales)
    if isinstance(scales, torch.Tensor):
        scales = cp.asarray(scales.numpy())
    if isinstance(kernels, torch.Tensor):
        kernels = cp.asarray(kernels.numpy())

    # Determine number of kernels
    if num_kernels is None:
        num_kernels = kernels.shape[0]
    else:
        num_kernels = min(num_kernels, kernels.shape[0])

    n_kernels = num_kernels
    H, W = mask.shape

    # Convert mask to complex and transform to frequency domain
    # Use 'forward' normalization to match OpenILT
    mask_complex = mask.astype(cp.complex64)
    mask_freq = cp.fft.fft2(mask_complex, norm='forward')

    # Compute aerial image: sum of weighted squared magnitudes
    aerial = cp.zeros((H, W), dtype=cp.float32)

    for i in range(n_kernels):
        # Multiply in frequency domain using quadrant mapping
        conv_freq = cp.zeros((H, W), dtype=cp.complex64)
        kh, kw = kernels.shape[1], kernels.shape[2]
        knxh, knyh = kh // 2, kw // 2

        k = kernels[i]
        conv_freq[:knxh+1, :knyh+1] = mask_freq[:knxh+1, :knyh+1] * k[-(knxh+1):, -(knyh+1):]
        conv_freq[:knxh+1, -knyh:] = mask_freq[:knxh+1, -knyh:] * k[-(knxh+1):, :knyh]
        conv_freq[-knxh:, :knyh+1] = mask_freq[-knxh:, :knyh+1] * k[:knxh, -(knyh+1):]
        conv_freq[-knxh:, -knyh:] = mask_freq[-knxh:, -knyh:] * k[:knxh, :knyh]

        # Transform back to spatial domain
        conv_spatial = cp.fft.ifft2(conv_freq, norm='forward')

        # Accumulate weighted squared magnitude
        aerial += scales[i] * cp.abs(conv_spatial) ** 2

    return aerial


def apply_optical_model_cutile(
    mask: np.ndarray | torch.Tensor,
    kernels: np.ndarray | torch.Tensor,
    scales: np.ndarray | torch.Tensor,
    num_kernels: int | None = None,
) -> np.ndarray:
    """
    Apply the Hopkins/SOCS optical model using cuTile FFT.

    Uses the custom cuTile 1D FFT kernel applied separably for 2D FFT.

    Args:
        mask: 2D binary mask array of shape (H, W).
        kernels: Complex kernels of shape (num_kernels, kh, kw) in spatial domain.
        scales: Kernel weights of shape (num_kernels,).
        num_kernels: Number of kernels to use (default: all).

    Returns:
        2D aerial image array of shape (H, W) as numpy array.
    """
    if not CUTILE_AVAILABLE:
        raise RuntimeError("cuTile FFT not available. Check test_fft.py import.")

    # Convert to torch tensors on CUDA
    if isinstance(mask, np.ndarray):
        mask = torch.from_numpy(mask).cuda()
    if isinstance(kernels, np.ndarray):
        kernels = torch.from_numpy(kernels).cuda()
    elif isinstance(kernels, torch.Tensor) and not kernels.is_cuda:
        kernels = kernels.cuda()
    if isinstance(scales, np.ndarray):
        scales = torch.from_numpy(scales).cuda()
    elif isinstance(scales, torch.Tensor) and not scales.is_cuda:
        scales = scales.cuda()

    # Determine number of kernels
    if num_kernels is None:
        num_kernels = kernels.shape[0]
    else:
        num_kernels = min(num_kernels, kernels.shape[0])

    H, W = mask.shape
    kh, kw = kernels.shape[1], kernels.shape[2]
    knxh, knyh = kh // 2, kw // 2

    # Get FFT configuration for this size
    if W not in FFT_CONFIG:
        raise ValueError(f"Unsupported image size {W}. Supported: {list(FFT_CONFIG.keys())}")
    config = FFT_CONFIG[W]
    factors = config["factors"]
    atom_packing_dim = config["atom_packing_dim"]
    chunk_size = 8

    # Convert mask to complex and compute FFT
    mask_complex = mask.to(torch.complex64)
    mask_freq = cutile_fft_2d(mask_complex, factors, atom_packing_dim, chunk_size)

    # Compute aerial image: sum of weighted squared magnitudes
    aerial = torch.zeros((H, W), dtype=torch.float32, device='cuda')

    for i in range(num_kernels):
        # Multiply in frequency domain using quadrant mapping
        conv_freq = torch.zeros((H, W), dtype=torch.complex64, device='cuda')

        k = kernels[i]
        conv_freq[:knxh+1, :knyh+1] = mask_freq[:knxh+1, :knyh+1] * k[-(knxh+1):, -(knyh+1):]
        conv_freq[:knxh+1, -knyh:] = mask_freq[:knxh+1, -knyh:] * k[-(knxh+1):, :knyh]
        conv_freq[-knxh:, :knyh+1] = mask_freq[-knxh:, :knyh+1] * k[:knxh, -(knyh+1):]
        conv_freq[-knxh:, -knyh:] = mask_freq[-knxh:, -knyh:] * k[:knxh, :knyh]

        # Transform back to spatial domain using IFFT (conjugate trick)
        conv_spatial = cutile_ifft_2d(conv_freq, factors, atom_packing_dim, chunk_size)

        # Accumulate weighted squared magnitude
        # Use real/imag manually to avoid torch.abs() JIT compilation on Blackwell GPUs
        magnitude_sq = conv_spatial.real**2 + conv_spatial.imag**2
        aerial += scales[i] * magnitude_sq

    return aerial.cpu().numpy()


def apply_optical_model_batched(
    mask: np.ndarray | cp.ndarray,
    kernels: np.ndarray | cp.ndarray,
    scales: np.ndarray | cp.ndarray,
    num_kernels: int | None = None,
) -> cp.ndarray:
    """
    Apply the Hopkins/SOCS optical model using batched FFT (more efficient).

    Same as apply_optical_model but processes all kernels in parallel.
    Uses OpenILT-compatible normalization and kernel multiplication.

    Args:
        mask: 2D binary mask array of shape (H, W).
        kernels: Complex kernels of shape (num_kernels, kh, kw) in spatial domain.
        scales: Kernel weights of shape (num_kernels,).
        num_kernels: Number of kernels to use (default: all).

    Returns:
        2D aerial image array of shape (H, W).
    """
    # Move to GPU if needed
    if isinstance(mask, np.ndarray):
        mask = cp.asarray(mask)
    if isinstance(kernels, np.ndarray):
        kernels = cp.asarray(kernels)
    if isinstance(scales, np.ndarray):
        scales = cp.asarray(scales)
    if isinstance(scales, torch.Tensor):
        scales = cp.asarray(scales.numpy())
    if isinstance(kernels, torch.Tensor):
        kernels = cp.asarray(kernels.numpy())

    # Determine number of kernels
    if num_kernels is None:
        num_kernels = kernels.shape[0]
    else:
        num_kernels = min(num_kernels, kernels.shape[0])

    n_kernels = num_kernels
    H, W = mask.shape

    # Convert mask to complex and transform to frequency domain
    # Use 'forward' normalization to match OpenILT
    mask_complex = mask.astype(cp.complex64)
    mask_freq = cp.fft.fft2(mask_complex, norm='forward')

    # Use the quadrant-based kernel multiplication
    conv_freq = kernel_multiply_freq(mask_freq, kernels, n_kernels)  # (n_kernels, H, W)

    # Batched inverse FFT with forward normalization
    conv_spatial = cp.fft.ifft2(conv_freq, axes=(1, 2), norm='forward')  # (n_kernels, H, W)

    # Squared magnitude
    magnitude_sq = cp.abs(conv_spatial) ** 2  # (n_kernels, H, W)

    # Weighted sum over kernels
    # scales is (n_kernels,), reshape to (n_kernels, 1, 1) for broadcasting
    scales_subset = scales[:n_kernels]
    aerial = cp.sum(scales_subset[:, None, None] * magnitude_sq, axis=0)

    return aerial


# =============================================================================
# Main - Generate aerial image and display
# =============================================================================

if __name__ == "__main__":
    import argparse
    import matplotlib.pyplot as plt
    from display import (
        load_mask_image, 
        load_kernels, 
        display_lithography_result,
        display_lithography_comparison,
    )

    parser = argparse.ArgumentParser(description="Apply optical model and display result")
    parser.add_argument("--save", type=str, default=None,
                        help="Save figure to file instead of displaying")
    parser.add_argument("--num-kernels", type=int, default=24,
                        help="Number of kernels to use (default: 24)")
    parser.add_argument("--cupy", action="store_true",
                        help="Use CuPy FFT implementation (default if no FFT specified)")
    parser.add_argument("--cutile", action="store_true",
                        help="Use cuTile FFT instead of CuPy FFT")
    parser.add_argument("--compare", action="store_true",
                        help="Run both CuPy and cuTile, display comparison")
    parser.add_argument("--cutline", type=int, default=None,
                        help="Y position for intensity cutline (default: 1088)")
    parser.add_argument("--threshold", type=float, default=0.225,
                        help="Print threshold for cutline (default: 0.225)")
    parser.add_argument("--cutline-xlim", type=str, default="500,1600",
                        help="X-axis limits for intensity profile (default: 500,1600)")
    parser.add_argument("--debug-frames", action="store_true",
                        help="Draw red frames around subplots for debugging layout")
    parser.add_argument("--mock", action="store_true",
                        help="Use mask as placeholder for aerial images (fast layout iteration)")
    args = parser.parse_args()

    if args.cutile and not CUTILE_AVAILABLE:
        print("Warning: cuTile FFT not available, falling back to CuPy")
        args.cutile = False
    
    # Parse cutline x-limits
    cutline_xlim = tuple(int(x) for x in args.cutline_xlim.split(","))

    # Paths to OpenILT data
    script_dir = Path(__file__).parent
    data_dir = script_dir / "openilt_data"

    mask_path = data_dir / "tmp" / "CurvILT_target1.png"
    kernel_path = data_dir / "kernel" / "kernels" / "focus.pt"
    scales_path = data_dir / "kernel" / "scales" / "focus.pt"

    # Load data
    print(f"Loading mask from: {mask_path}")
    mask = load_mask_image(mask_path)
    print(f"  Mask shape: {mask.shape}")

    print(f"Loading kernels from: {kernel_path}")
    kernels = load_kernels(kernel_path)
    print(f"  Kernels shape: {kernels.shape}")

    print(f"Loading scales from: {scales_path}")
    scales = load_scales(scales_path)
    print(f"  Scales shape: {scales.shape}")

    # Determine cutline position (default: 1088)
    if args.cutline is None:
        cutline_y = 1088
    else:
        cutline_y = args.cutline

    # Apply optical model(s) or use mock data
    if args.mock:
        print("\nUsing mock aerial images (mask as placeholder)...")
        aerial_cupy = mask.astype(np.float32) * 0.5  # Scale to similar range
        aerial_cutile = mask.astype(np.float32) * 0.5
    else:
        print(f"\nApplying optical model with {args.num_kernels} kernels...")
    
    if args.compare:
        if not args.mock:
            # Run both implementations for comparison
            print("  Running CuPy FFT implementation...")
            aerial_gpu = apply_optical_model_batched(
                mask, kernels, scales, num_kernels=args.num_kernels
            )
            aerial_cupy = cp.asnumpy(aerial_gpu)
            print(f"    CuPy aerial range: [{aerial_cupy.min():.4f}, {aerial_cupy.max():.4f}]")
            
            print("  Running cuTile FFT implementation...")
            aerial_cutile = apply_optical_model_cutile(
                mask, kernels, scales, num_kernels=args.num_kernels
            )
            print(f"    cuTile aerial range: [{aerial_cutile.min():.4f}, {aerial_cutile.max():.4f}]")
        
        # Compute difference stats
        diff = np.abs(aerial_cupy - aerial_cutile)
        print(f"    Max absolute difference: {diff.max():.6f}")
        print(f"    Mean absolute difference: {diff.mean():.6f}")
        
        print(f"  Cutline at Y = {cutline_y}")
        
        # Display comparison
        print("\nCreating comparison visualization...")
        fig = display_lithography_comparison(
            mask=mask,
            kernels=kernels,
            aerial_cupy=aerial_cupy,
            aerial_cutile=aerial_cutile,
            num_kernels_to_show=24,
            suptitle="Optical Imaging: CuPy vs cuTile FFT",
            cutline_y=cutline_y,
            threshold=args.threshold,
            cutline_xlim=cutline_xlim,
            debug_frames=args.debug_frames,
        )
    else:
        # Single implementation mode
        if args.cutile:
            print("  Using cuTile FFT implementation")
            aerial = apply_optical_model_cutile(
                mask, kernels, scales, num_kernels=args.num_kernels
            )
        else:
            # Default to CuPy (batched implementation)
            print("  Using CuPy FFT implementation")
            aerial_gpu = apply_optical_model_batched(
                mask, kernels, scales, num_kernels=args.num_kernels
            )
            aerial = cp.asnumpy(aerial_gpu)
        
        print(f"  Aerial image shape: {aerial.shape}")
        print(f"  Aerial range: [{aerial.min():.4f}, {aerial.max():.4f}]")
        print(f"  Cutline at Y = {cutline_y}")

        # Display result
        print("\nCreating visualization...")
        fft_impl = "cuTile" if args.cutile else "CuPy"
        fig = display_lithography_result(
            mask=mask,
            kernels=kernels,
            aerial=aerial,
            num_kernels_to_show=24,
            suptitle=f"Optical Imaging: FFT Convolution ({fft_impl})",
            cutline_y=cutline_y,
            threshold=args.threshold,
        )

    if not args.compare:
        plt.tight_layout()

    if args.save:
        plt.savefig(args.save, dpi=150, bbox_inches="tight")
        print(f"Saved figure to: {args.save}")
    else:
        plt.show()