Batch Generation¶

def build_wfe_budget(radial_budget: list = [0, 100, 50, 36, 18, 9, 5]) -> list:
    """Build a flat WFE budget list from a per-radial-order specification.

    Each entry ``radial_budget[i]`` is repeated ``i + 1`` times to account for
    the increasing number of Zernike modes per radial order (order 0 has 1
    mode, order 1 has 2 modes, etc.).

    Args:
        radial_budget (list[int]): WFE amplitude in nm RMS for each radial
            order, starting from order 0 (piston).  The default
            ``[0, 100, 50, 36, 18, 9, 5]`` covers orders 0–6.

    Returns:
        list[int]: Flat budget list whose total length equals the number of
            Zernike modes covered by the specification.

    Example:
        >>> build_wfe_budget([0, 100, 50])
        [0, 100, 100, 50, 50, 50]
    """
    WFE_BUDGET = []
    for i, budget in enumerate(radial_budget):
        WFE_BUDGET.extend([budget] * (i + 1))
    return WFE_BUDGET

def generate_coefficients(wfe_budget, sec_factor: float = 1.0) -> np.ndarray:
    """Draw a random realisation of Zernike coefficients from a uniform distribution.

    Each coefficient is drawn independently from a uniform distribution
    over ``[-budget_i * sec_factor, +budget_i * sec_factor]``
    (converted from nm to metres internally).

    Args:
        wfe_budget (list[float]): WFE amplitude in nm RMS for each Zernike
            coefficient.  Typically the output of :func:`build_wfe_budget`.
        sec_factor (float): Global scaling factor applied to every budget
            entry before sampling.  Useful to reduce the amplitude of
            secondary-mirror-sensitive modes.  Defaults to ``1.0``.

    Returns:
        numpy.ndarray: Array of random Zernike coefficients in **metres**,
            with the same length as ``wfe_budget``.

    Example:
        >>> import numpy as np
        >>> np.random.seed(0)
        >>> budget = build_wfe_budget([0, 50, 30])   # 6 coefficients
        >>> coefs = generate_coefficients(budget)
        >>> coefs.shape
        (6,)
    """
    return np.random.uniform(
        low=-1e-9 * np.array(wfe_budget) * sec_factor,
        high=1e-9 * np.array(wfe_budget) * sec_factor
    )