NA = 0.6

dx_nm = 20
dz_nm = 30
wav_nm = 759

nx = 512
nr = nx//2 - 1
nz = 256

base_kwargs_3d = {
    "n_pix_pupil": nr,
    "na": NA,
    "wavelength": wav_nm,
    "pix_size": dx_nm,
    "defocus_step": dz_nm,
}

def _():
    waists = np.linspace(0.1, 1, 100)
    #radii = np.arange(nx // 2) * dx_nm / 2 / wav_nm # units of wavelength
    rayleigh_limit = 0.61 * wav_nm / NA
    radii = np.arange(nx // 2) * dx_nm / rayleigh_limit # units of Rayleigh diffraction limit

    argmins = []
    valid_waists = []

    for waist in waists:
        prop = ScalarSphericalPropagator(
            n_pix_psf=nx, n_defocus=1, envelope=waist, **base_kwargs_3d
        )
        ff = prop.compute_focus_field().squeeze().numpy()
        ip = np.abs(ff[nx//2, nx//2:])**2

        imin, _ = find_first_local_minimum(ip)
        if imin is not None:
            valid_waists.append(waist)
            argmins.append(radii[imin])
        else:
            print(f"waist {waist:.3f} has no minimum")

    fig, ax = plt.subplots(figsize=(6, 4))
    ax.plot(valid_waists, argmins, marker='o')
    ax.set_xlabel(r"$w_0/R$")
    ax.set_ylabel(r"$r_{\min}/\zeta_0$")
    ax.set_title("First minimum position vs beam waist")

    return mo.mpl.interactive(fig)