Lct¤

def lct_sampling_quick(x_in, x_out, lam, ABCD, eps=1e-30):
    """
    Quick sampling diagnostics for Collins LCT on separable grids.

    Returns Nyquist ratios (>=1 is safe-ish):
      - p_kernel: sampling of exp(-i 2π x x' /(λ b)) kernel
      - p_pre:    sampling of input chirp exp(i π a x^2 /(λ b))
      - p_post:   sampling of output chirp exp(i π d x'^2 /(λ b))
    """
    a, b, c, d = [ABCD.reshape(-1)[i] for i in range(4)]

    dx_in = x_in[1] - x_in[0]
    dx_out = x_out[1] - x_out[0]
    X_in = 0.5 * (x_in[-1] - x_in[0])
    X_out = 0.5 * (x_out[-1] - x_out[0])

    # Kernel phase: exp(-i 2π x x' /(λ b))
    # worst phase slope in x is at max |x'|
    dphi_in_max = (2 * jnp.pi / (lam * jnp.abs(b) + eps)) * X_out * dx_in
    dphi_out_max = (2 * jnp.pi / (lam * jnp.abs(b) + eps)) * X_in * dx_out
    p_kernel = jnp.pi / (jnp.maximum(dphi_in_max, dphi_out_max) + eps)

    # Pre/post chirps: phase ~ π a x^2 /(λ b), slope ~ 2π a x /(λ b)
    dphi_pre_max = (2 * jnp.pi * jnp.abs(a) / (lam * jnp.abs(b) + eps)) * X_in * dx_in
    dphi_post_max = (
        (2 * jnp.pi * jnp.abs(d) / (lam * jnp.abs(b) + eps)) * X_out * dx_out
    )
    p_pre = jnp.pi / (dphi_pre_max + eps)
    p_post = jnp.pi / (dphi_post_max + eps)

    return {"p_kernel": p_kernel, "p_pre": p_pre, "p_post": p_post}

def lct_kernels(
    spec_in: Array | tuple, spec_out: Array | tuple, lam: float, ABCD: Array
) -> tuple:
    x_in, y_in = unpack_coord_spec(spec_in)
    x_out, y_out = unpack_coord_spec(spec_out)

    a, b, c, d = ABCD.flatten()

    r2_in = r2_coords((x_in, y_in))
    r2_out = r2_coords((x_out, y_out))

    pre = np.exp(1j * np.pi * a * r2_in / (lam * b))
    post = np.exp(1j * np.pi * d * r2_out / (lam * b))

    dx_in = x_in[1] - x_in[0]
    dy_in = y_in[1] - y_in[0]
    dx_out = x_out[1] - x_out[0]
    dy_out = y_out[1] - y_out[0]

    alpha = -2.0 * np.pi / (lam * b)

    Kx, Ky = mft_kernels(
        spec_in=spec_in, spec_out=spec_out, alpha=alpha, weight=(dx_in, dy_in)
    )

    pref = 1.0 / (1j * lam * b)
    scale = np.sqrt((dx_out * dy_out) / (dx_in * dy_in))
    return pre, Kx, Ky, post, pref, scale

def lct_kernel_prop(
    u_in: Array,
    pre: Array,
    Kx: Array,
    Ky: Array,
    post: Array,
    pref: complex,
    scale: float,
) -> Array:
    u_tmp = pre * u_in
    u_mft = mft(u_tmp, Kx, Ky)
    return pref * scale * u_mft * post

def lct_prop_basic(
    u_in: Array,
    spec_in: Array | tuple,
    spec_out: Array | tuple,
    lam: float,
    ABCD: Array,
) -> Array:
    pre, Kx, Ky, post, pref, scale = lct_kernels(spec_in, spec_out, lam, ABCD)
    return lct_kernel_prop(u_in, pre, Kx, Ky, post, pref, scale)

def lct_prop(
    u_in: Array,
    spec_in: Array | tuple,
    spec_out: Array | tuple,
    lam: float,
    ABCD: Array,
    curv_in: float | None = None,
    curv_out: float | None = None,
    mode: str = "physical",
    strip_input: bool = True,
    return_residual: bool = False,
) -> Array:
    ABCD_res, curv_in, curv_out = factorise_curv(ABCD, curv_in, curv_out, mode)

    u_res_in = remove_curv(u_in, spec_in, lam, curv_in)

    u_res_out = lct_prop_basic(u_res_in, spec_in, spec_out, lam, ABCD_res)

    if return_residual:
        return u_res_out

    return apply_curv(u_res_out, spec_out, lam, curv_out)

def propagate_mono_abcd(self, wavelength, offset=np.zeros(2), return_wf=False):
    import dLux as dl  # keep local

    wf = dl.Wavefront(self.wf_npixels, self.diameter, wavelength).tilt(offset)

    for layer in list(self.layers.values()):
        wf *= layer

    u_in = wf.phasor

    fl = self.fnumber * self.diameter  # meters

    # If your existing code uses self.defocus in nm, convert:
    z_defocus = self.defocus * 1e-9

    abcd = compose_abcd([abcd_lens(fl), abcd_free_space(fl + z_defocus)])

    # Input sampling (pupil plane)
    N_in = self.wf_npixels
    dx_in = self.diameter / self.wf_npixels
    x_in = dlu.nd_coords(N_in, dx_in)

    # Output sampling: match your existing psf_pixel_scale (angular) via x = f * theta
    true_pixel_scale = self.psf_pixel_scale / self.oversample  # arcsec/pix
    theta_pix = dlu.arcsec2rad(true_pixel_scale)  # rad/pix

    N_out = self.psf_npixels * self.oversample
    dx_out = fl * theta_pix  # meters/pix at focal plane

    x_out = dlu.nd_coords(N_out, dx_out)

    u_out = lct_prop_basic(u_in, x_in, x_out, wavelength, abcd)

    wf_out = dl.Wavefront(N_out, N_out * dx_out, wavelength).set(
        ["amplitude", "phase"], [np.abs(u_out), np.angle(u_out)]
    )

    if return_wf:
        return wf_out
    return wf_out.psf