Curvature¤

def r2_coords(spec_in: Array | tuple) -> Array:
    x, y = unpack_coord_spec(spec_in)
    return (y**2)[:, None] + (x**2)[None, :]

def quad_phase(coords: Array | tuple, lam: float, curv: float | Array) -> Array:
    return np.exp(1j * np.pi * curv * r2_coords(coords) / lam)

def apply_curv(
    u: Array, coords: Array | tuple, lam: float, curv: float | Array
) -> Array:
    return u * quad_phase(coords, lam, curv)

def remove_curv(
    u: Array, coords: Array | tuple, lam: float, curv: float | Array
) -> Array:
    return u * quad_phase(coords, lam, curv).conj()

def propagate_curv(ABCD: Array, curv_in: float | Array) -> Array:
    a, b, c, d = ABCD.flatten()
    return (c + d * curv_in) / (a + b * curv_in)

def residual_abcd(ABCD: Array, curv_in: float, curv_out: float) -> Array:
    a, b, c, d = ABCD.flatten()
    a_res = a + b * curv_in
    d_res = d - b * curv_out
    c_res = (a_res * d_res - 1.0) / b
    return np.array([[a_res, b], [c_res, d_res]])

def residual_curv_cancel(ABCD: Array):
    a, b, c, d = ABCD.flatten()
    curv_in = -a / b
    curv_out = d / b
    return curv_in, curv_out

def factorise_curv(
    ABCD: Array,
    curv_in: float | Array | None,
    curv_out: float | Array | None,
    mode: str = "physical",
) -> tuple[Array, float | Array, float | Array]:
    if mode == "physical":
        if curv_in is None:
            curv_in = 0.0
        if curv_out is None:
            curv_out = propagate_curv(ABCD, curv_in)
    elif mode == "cancel":
        curv_in, curv_out = residual_curv_cancel(ABCD)
    elif mode == "manual":
        if curv_in is None or curv_out is None:
            raise ValueError("mode='manual' requires curv_in and curv_out.")
    else:
        raise ValueError(f"Unknown mode '{mode}'")

    ABCD_res = residual_abcd(ABCD, curv_in, curv_out)
    return ABCD_res, curv_in, curv_out