Source code for pyvbmc.priors.tile_inputs

import numpy as np


def _squeezed_shape(shape):
    """The shape with its axes of length one dropped."""
    return tuple(length for length in shape if length != 1)




def tile_inputs(*args, size=None, squeeze=False):
    """Tile scalar inputs to have the same dimension as array inputs.

    If all inputs are given as scalars, returned arrays will have shape `size`
    if `size` is a tuple, or shape `(size,)` if `size` is an integer.

    An argument agrees with `size` when the two hold the same lengths once
    their axes of length one are dropped, so that a row, a column and a flat
    array of `D` elements all agree with `size=D` while an array of another
    layout, such as `(2, 2)` against `size=4`, does not.

    Parameters
    ----------
    *args : [Union[float, np.ndarray]]
        The inputs to tile.
    size : Union[int, tuple], optional
        The desired size/shape of the output. By default the shape of the
        non-scalar inputs, or `(1,)` when every input is a scalar.
    squeeze : bool
        If `True`, then drop 1-d axes from inputs. Default `False`.

    Returns
    -------
    args : list of np.ndarray
        The inputs, in their order, each of shape `size`.

    Raises
    ------
    ValueError
        If the non-scalar arguments do not have the same shape, or if they do
        not agree with `size`.
    """
    if type(size) == int:
        size = (size,)
    shape = None

    # Check that all non-scalar inputs have the same shape
    args = list(args)
    for i, arg in enumerate(args):
        if not (np.isscalar(arg)):
            if squeeze:
                arg = args[i] = np.atleast_1d(np.squeeze(np.array(arg)))
            else:
                arg = args[i] = np.array(arg)
            if shape is None:
                shape = arg.shape
            elif arg.shape != shape:
                raise ValueError(
                    f"All inputs should have the same shape, but found inputs with shapes {shape} and {arg.shape}."
                )

    if size is None:
        if shape is None:
            # Default to shape (1,)
            size = (1,)
        else:
            # Or use inferred shape
            size = shape
    elif shape is not None and _squeezed_shape(shape) != _squeezed_shape(size):
        # `reshape` below accepts any array with the right number of
        # elements, which would map the parameters onto the wrong
        # coordinates.
        raise ValueError(
            f"All inputs should agree with size={size}, but found an input "
            f"with shape {shape}."
        )

    for i, arg in enumerate(args):
        if np.isscalar(arg):
            args[i] = np.full(size, arg)
        else:
            args[i] = args[i].reshape(size)

    return args