Point Spread Function v1

Made with Chatgpt 5, Python 3, and Google Colab notebook. Source:

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Physics-accurate telescope diffraction image (4096×4096) using all 24-bit RGB colors exactly once.

Features
- Pupil model: annulus (central obstruction) + spider vanes → Fraunhofer PSF via FFT
- Aberrations: fringe Zernike terms (defocus, astigmatism 0/45°, coma x/y, spherical)
- Intensity shaping: monotone gamma (does not change rank order)
- Exact-once color usage: rank-match intensity to Rec.709 luma; checksum + XOR proof
- Reserved overlays: center crosshair (darkest colors), neat legend box & ticks (lightest colors)

Requires: numpy, pillow
"""

import time
import numpy as np
from PIL import Image

# -------------------------- Zernike (fringe polynomials) --------------------------
# Fringe (unnormalized) variants often used in engineering optics. Coefficients are in waves.
# theta in radians, rho∈[0,1]
def z_defocus(rho, th):      # Z4 fringe ~ 2ρ^2 - 1
    return 2.0 * rho**2 - 1.0

def z_astig0(rho, th):       # Z2 fringe ~ ρ^2 cos 2θ
    return rho**2 * np.cos(2.0 * th, dtype=np.float32)

def z_astig45(rho, th):      # Z3 fringe ~ ρ^2 sin 2θ
    return rho**2 * np.sin(2.0 * th, dtype=np.float32)

def z_coma_x(rho, th):       # Z7 fringe ~ (3ρ^3 - 2ρ) cos θ
    return (3.0 * rho**3 - 2.0 * rho) * np.cos(th, dtype=np.float32)

def z_coma_y(rho, th):       # Z8 fringe ~ (3ρ^3 - 2ρ) sin θ
    return (3.0 * rho**3 - 2.0 * rho) * np.sin(th, dtype=np.float32)

def z_spherical(rho, th):    # Z11 fringe ~ 6ρ^4 - 6ρ^2 + 1
    return 6.0 * rho**4 - 6.0 * rho**2 + 1.0

# -------------------------- Pupil → PSF (FFT) --------------------------
def aperture_fft_intensity(h, w,
                           epsilon=0.33,        # central obstruction ratio D_inner / D_outer
                           vanes=4,             # number of spider vanes
                           vane_thickness=0.008,# fraction of pupil diameter
                           rotate_deg=15.0,     # orientation of vanes [deg]
                           gamma=0.80,          # intensity^gamma for contrast (monotone)
                           # Aberration coefficients in waves (peak-to-valley-ish)
                           c_defocus = 0.00,
                           c_astig0  = 0.00,
                           c_astig45 = 0.00,
                           c_coma_x  = 0.00,
                           c_coma_y  = 0.00,
                           c_spherical=0.00):
    """
    Build Fraunhofer diffraction intensity from a physical pupil.
    Steps:
      1) Pupil amplitude: annulus minus spider vanes
      2) Pupil phase: exp(i*2π*sum c_k * Z_k(rho, theta))
      3) PSF ∝ |FFT(pupil)|^2  (fftshifted)
    """
    H, W = int(h), int(w)

    # Normalized coordinates in [-1,1]; inscribed pupil radius ≈ 1.0
    yy = np.linspace(-1.0, 1.0, H, dtype=np.float32)
    xx = np.linspace(-1.0, 1.0, W, dtype=np.float32)
    X, Y = np.meshgrid(xx, yy, indexing="xy")
    R = np.hypot(X, Y, dtype=np.float32)
    TH = np.arctan2(Y, X, dtype=np.float32)

    # Annulus mask (slightly underfilled to avoid hard edge touching frame)
    outer_R = 0.98
    inner_R = outer_R * np.float32(epsilon)
    outer = (R <= outer_R).astype(np.float32)
    inner = (R <= inner_R).astype(np.float32)
    annulus = outer - inner

    # Rotate vane coordinates
    th = np.deg2rad(np.float32(rotate_deg))
    c, s = np.cos(th), np.sin(th)
    Xr =  c * X + s * Y
    Yr = -s * X + c * Y

    pupil = annulus.copy()
    if vanes > 0 and vane_thickness > 0.0:
        t = np.float32(vane_thickness)
        for k in range(vanes):
            phi = np.float32(k) * (np.pi / np.float32(vanes))
            cp, sp = np.cos(phi), np.sin(phi)
            u =  cp * Xr + sp * Yr
            v = -sp * Xr + cp * Yr
            vane = (np.abs(v) <= (t / 2.0)).astype(np.float32)
            # remove vane only where within outer pupil
            pupil *= (1.0 - vane) + (1.0 - outer)

    # Phase from Zernike aberrations (fringe set). Coeffs are in waves → 2π factor to radians.
    if any(abs(x) > 1e-12 for x in [c_defocus, c_astig0, c_astig45, c_coma_x, c_coma_y, c_spherical]):
        phase_waves = (c_defocus   * z_defocus(R/outer_R, TH) +
                       c_astig0    * z_astig0 (R/outer_R, TH) +
                       c_astig45   * z_astig45(R/outer_R, TH) +
                       c_coma_x    * z_coma_x (R/outer_R, TH) +
                       c_coma_y    * z_coma_y (R/outer_R, TH) +
                       c_spherical * z_spherical(R/outer_R, TH)).astype(np.float32)
        phase = (2.0 * np.pi * phase_waves).astype(np.float32)
        pupil = pupil * np.exp(1j * phase, dtype=np.complex64)
    else:
        pupil = pupil.astype(np.complex64)

    # FFT → PSF
    F = np.fft.fft2(pupil)
    I = np.fft.fftshift(np.abs(F)**2).astype(np.float32)

    # Normalize + gamma (monotone)
    I -= I.min()
    m = I.max()
    if m > 0:
        I /= m
    if gamma != 1.0:
        I = np.power(I, np.float32(gamma), dtype=np.float32)
    return I

# -------------------------- Exact-once mapping with reserved overlays --------------------------
def stable_argsort(a):
    try:
        return np.argsort(a, kind='stable')
    except TypeError:
        return np.argsort(a, kind='mergesort')

def map_allrgb_with_overlays(intensity,
                             use_crosshair=True,
                             use_legend=True,
                             invert=False,
                             legend_box=(32, 4096-32-120, 480, 120),  # (x,y,w,h) bottom-left by default
                             legend_thick=2,
                             ticks=9, tick_len=16, tick_thick=2):
    """
    Map all 24-bit colors exactly once with optional reserved overlays:
      - Crosshair (center horizontal+vertical lines) uses the darkest colors.
      - Legend box (corner rectangle with ticks) uses the lightest colors.
    The remaining pixels are filled by rank-matching intensity with color luma.
    """
    H, W = intensity.shape
    N = H * W

    # 1) Build overlay mask: 0=free, 1=dark-reserved (crosshair), 2=light-reserved (legend)
    mask = np.zeros((H, W), dtype=np.uint8)

    # Crosshair pixels (avoid double-count at center)
    cross_idx = []
    if use_crosshair:
        cy, cx = H // 2, W // 2
        # horizontal
        mask[cy, :] = 1
        # vertical
        mask[:, cx] = np.maximum(mask[:, cx], 1)
        # collect linear indices
        y = np.full(W, cy, dtype=np.int64)
        x = np.arange(W, dtype=np.int64)
        cross_idx.extend((y * W + x).tolist())
        y = np.arange(H, dtype=np.int64)
        x = np.full(H, cx, dtype=np.int64)
        cross_idx.extend((y * W + x).tolist())
        # remove the center double-count
        cross_idx = sorted(set(cross_idx))

    # Legend rectangle (border only) with ticks
    legend_idx = []
    if use_legend:
        lx, ly, lw, lh = legend_box
        lx = int(np.clip(lx, 0, W-1)); ly = int(np.clip(ly, 0, H-1))
        lw = int(np.clip(lw, 1, W - lx)); lh = int(np.clip(lh, 1, H - ly))
        t  = int(max(1, legend_thick))
        # borders
        # top/bottom
        mask[ly:ly+t, lx:lx+lw] = 2
        mask[ly+lh-t:ly+lh, lx:lx+lw] = 2
        # left/right
        mask[ly:ly+lh, lx:lx+t] = 2
        mask[ly:ly+lh, lx+lw-t:lx+lw] = 2

        # ticks along bottom border inside the box (evenly spaced)
        if ticks > 0:
            for i in range(ticks+1):
                tx = lx + int(round(i * (lw-1) / ticks))
                ty0 = ly + lh - t - tick_len
                ty1 = ly + lh - t
                ty0 = max(ly + t, ty0)
                ty1 = min(ly + lh - 1, ty1)
                mask[ty0:ty1, tx:tx+tick_thick] = 2

        legend_idx = np.flatnonzero(mask == 2).tolist()

    # 2) Compute pixel order among FREE pixels only
    flat_I = intensity.ravel()
    free_lin = np.flatnonzero(mask.ravel() == 0)
    free_order = stable_argsort(flat_I[free_lin])
    if invert:
        free_order = free_order[::-1]
    free_lin_sorted = free_lin[free_order]
    del free_order, free_lin

    # 3) Build global color ordering by perceived brightness (Rec.709 luma, fixed-point)
    c = np.arange(N, dtype=np.uint32)
    R = (c >> 16) & 0xFF
    G = (c >> 8)  & 0xFF
    B = (c      ) & 0xFF
    Y256 = (54 * R + 183 * G + 19 * B).astype(np.uint32)
    color_order = stable_argsort(Y256)  # dark→bright
    del Y256, R, G, B

    # Reserve darkest for crosshair, lightest for legend
    D_dark  = len(cross_idx)
    D_light = len(legend_idx)
    if D_dark + D_light > N:
        raise RuntimeError("Overlay area exceeds total pixels.")

    darkest_ids  = color_order[:D_dark].astype(np.uint32, copy=False)
    lightest_ids = color_order[-D_light:].astype(np.uint32, copy=False) if D_light > 0 else np.array([], dtype=np.uint32)
    remain_ids   = color_order[D_dark: N - D_light].astype(np.uint32, copy=False)
    del color_order

    # 4) Assemble output
    img_flat = np.empty((N, 3), dtype=np.uint8)

    # Fill FREE pixels with remaining colors by intensity rank
    co = remain_ids
    img_flat[free_lin_sorted, 0] = (co >> 16).astype(np.uint8)
    img_flat[free_lin_sorted, 1] = ((co >> 8) & 0xFF).astype(np.uint8)
    img_flat[free_lin_sorted, 2] = (co & 0xFF).astype(np.uint8)
    del co

    # Paint crosshair with darkest colors (center-out radial gradient along both axes)
    if D_dark > 0:
        # Sort crosshair indices by distance from center for a nice gradient
        cross_idx_arr = np.array(cross_idx, dtype=np.int64)
        ys = (cross_idx_arr // W).astype(np.int32)
        xs = (cross_idx_arr %  W).astype(np.int32)
        cy, cx = H // 2, W // 2
        d = np.hypot(xs - cx, ys - cy).astype(np.float32)
        order = np.argsort(d, kind='mergesort')  # stable
        ci = cross_idx_arr[order]
        cd = darkest_ids  # dark→less dark
        img_flat[ci, 0] = ((cd >> 16) & 0xFF).astype(np.uint8)
        img_flat[ci, 1] = ((cd >> 8)  & 0xFF).astype(np.uint8)
        img_flat[ci, 2] = ( cd        & 0xFF).astype(np.uint8)
        del cross_idx_arr, ys, xs, d, order, ci, cd

    # Paint legend with lightest colors (simple sequential assignment)
    if D_light > 0:
        li = np.array(legend_idx, dtype=np.int64)
        ld = lightest_ids
        img_flat[li, 0] = ((ld >> 16) & 0xFF).astype(np.uint8)
        img_flat[li, 1] = ((ld >> 8 ) & 0xFF).astype(np.uint8)
        img_flat[li, 2] = ( ld        & 0xFF).astype(np.uint8)
        del li, ld

    img = img_flat.reshape(H, W, 3)
    del img_flat

    # 5) Integrity checks — prove every color used exactly once
    packed = (img[...,0].astype(np.uint32) << 16) | \
             (img[...,1].astype(np.uint32) << 8)  | \
             (img[...,2].astype(np.uint32))
    sum1 = int(packed.sum(dtype=np.uint64))
    xors = np.bitwise_xor.reduce(packed.ravel().astype(np.uint64)).item()
    del packed

    exp_sum1 = N * (N - 1) // 2
    assert sum1 == exp_sum1 and xors == 0, "Color coverage mismatch."

    return img

# -------------------------- Driver --------------------------
def build_image(size=4096,
                invert=False,
                # Pupil + FFT parameters
                epsilon=0.33, vanes=4, vane_thickness=0.008, rotate_deg=15.0, gamma=0.80,
                # Aberration waves (try small values like 0.05–0.20 for visible but tasteful effects)
                c_defocus=0.00, c_astig0=0.00, c_astig45=0.00, c_coma_x=0.00, c_coma_y=0.00, c_spherical=0.00,
                # Overlays
                use_crosshair=True, use_legend=True,
                legend_box=(32, 4096-32-120, 480, 120), legend_thick=2, ticks=9, tick_len=16, tick_thick=2,
                out_path=None):
    H = W = int(size)
    print(f"Image: {W}×{H}  (pixels: {H*W:,})")

    print("Computing aperture FFT intensity (annulus + spider vanes + optional aberrations)...")
    I = aperture_fft_intensity(H, W,
                               epsilon=epsilon, vanes=vanes, vane_thickness=vane_thickness,
                               rotate_deg=rotate_deg, gamma=gamma,
                               c_defocus=c_defocus, c_astig0=c_astig0, c_astig45=c_astig45,
                               c_coma_x=c_coma_x, c_coma_y=c_coma_y, c_spherical=c_spherical)

    print("Mapping all 24-bit colors with reserved overlays (crosshair & legend)...")
    img = map_allrgb_with_overlays(I,
                                   use_crosshair=use_crosshair,
                                   use_legend=use_legend,
                                   invert=invert,
                                   legend_box=legend_box,
                                   legend_thick=legend_thick,
                                   ticks=ticks, tick_len=tick_len, tick_thick=tick_thick)

    ts = time.strftime("%Y%m%d-%H%M%S")
    if out_path is None:
        suffix = f"eps{epsilon:.2f}_v{vanes}_t{vane_thickness:.3f}_rot{rotate_deg:.0f}_g{gamma:.2f}"
        ab = f"d{c_defocus:.2f}_a0{c_astig0:.2f}_a45{c_astig45:.2f}_cx{c_coma_x:.2f}_cy{c_coma_y:.2f}_s{c_spherical:.2f}"
        out_path = f"telescope_psf_allRGB_{W}x{H}_{suffix}_{ab}_{'inv' if invert else 'norm'}_{ts}.png"

    Image.fromarray(img, mode="RGB").save(out_path, compress_level=0, optimize=False)
    print(f"Done. Saved: {out_path}")
    return img

if __name__ == "__main__":
    # Tasteful defaults; tweak freely (all monotone/gamma and optical changes preserve exact-once mapping).
    build_image(
        size=4096,
        invert=False,          # True puts vivid colors onto the dark spikes/rings
        # Pupil
        epsilon=0.33,          # 0.10..0.40 typical; larger → bigger central hole → stronger rings
        vanes=4,               # 3,4,6 common; controls number of spikes
        vane_thickness=0.008,  # 0.006..0.012 looks good
        rotate_deg=25.0,
        gamma=0.80,            # 0.75..0.95: lower → more visible rings/spikes

        # Aberrations (in waves). Try small nonzero for personality:
        c_defocus=0.08,
        c_astig0=0.05,
        c_astig45=0.00,
        c_coma_x=0.06,
        c_coma_y=-0.04,
        c_spherical=0.03,

        # Overlays
        use_crosshair=True,
        use_legend=True,
        legend_box=(32, 4096-32-120, 480, 120),  # x,y,w,h (place near bottom-left)
        legend_thick=2,
        ticks=9, tick_len=16, tick_thick=2,
    )