Abcd Matrices¤

def abcd_surface_power(power: float | Array) -> Array:
    p = np.asarray(power).reshape(())
    M = np.eye(2, dtype=p.dtype)
    M = M.at[1, 0].set(-p)
    return M

def abcd_lens(focal_length: float | Array) -> Array:
    f = np.asarray(focal_length).reshape(())
    return abcd_surface_power(1.0 / f)

def abcd_mirror(radius: float) -> Array:
    return abcd_surface_power(2.0 / radius)

def abcd_free_space(z: float | Array) -> Array:
    z = np.asarray(z).reshape(())  # force scalar (0-d)
    M = np.eye(2, dtype=z.dtype)
    M = M.at[0, 1].set(z)
    return M

def abcd_fraunhofer(focal_length: float) -> Array:
    return np.array([[0.0, focal_length], [-1.0 / focal_length, 0.0]])

def compose_abcd(matrices: list | tuple) -> Array:
    M = np.eye(2)
    for Mi in matrices:
        M = Mi @ M
    return M