claudson/ai/lib/python3.12/site-packages/opt_einsum/testing.py

"""Testing routines for opt_einsum."""

import random
from typing import Any, Dict, List, Literal, Optional, Tuple, Union, overload

import pytest

from opt_einsum.parser import get_symbol
from opt_einsum.typing import ArrayType, PathType, TensorShapeType

_valid_chars = "abcdefghijklmopqABC"
_sizes = [2, 3, 4, 5, 4, 3, 2, 6, 5, 4, 3, 2, 5, 7, 4, 3, 2, 3, 4]
_default_dim_dict = dict(zip(_valid_chars, _sizes))


def build_shapes(string: str, dimension_dict: Optional[Dict[str, int]] = None) -> Tuple[TensorShapeType, ...]:
    """Builds random tensor shapes for testing.

    Parameters:
        string: List of tensor strings to build
        dimension_dict: Dictionary of index sizes, defaults to indices size of 2-7

    Returns:
        The resulting shapes.

    Examples:
        ```python
        >>> shapes = build_shapes('abbc', {'a': 2, 'b':3, 'c':5})
        >>> shapes
        [(2, 3), (3, 3, 5), (5,)]
        ```

    """
    if dimension_dict is None:
        dimension_dict = _default_dim_dict

    shapes = []
    terms = string.split("->")[0].split(",")
    for term in terms:
        dims = [dimension_dict[x] for x in term]
        shapes.append(tuple(dims))
    return tuple(shapes)


def build_views(
    string: str, dimension_dict: Optional[Dict[str, int]] = None, array_function: Optional[Any] = None
) -> Tuple[ArrayType]:
    """Builds random numpy arrays for testing.

    Parameters:
        string: List of tensor strings to build
        dimension_dict: Dictionary of index _sizes
        array_function: Function to build the arrays, defaults to np.random.rand

    Returns:
        The resulting views.

    Examples:
        ```python
        >>> view = build_views('abbc', {'a': 2, 'b':3, 'c':5})
        >>> view[0].shape
        (2, 3, 3, 5)
        ```

    """
    if array_function is None:
        np = pytest.importorskip("numpy")
        array_function = np.random.rand

    views = []
    for shape in build_shapes(string, dimension_dict=dimension_dict):
        if shape:
            views.append(array_function(*shape))
        else:
            views.append(random.random())
    return tuple(views)


@overload
def rand_equation(
    n: int,
    regularity: int,
    n_out: int = ...,
    d_min: int = ...,
    d_max: int = ...,
    seed: Optional[int] = ...,
    global_dim: bool = ...,
    *,
    return_size_dict: Literal[True],
) -> Tuple[str, PathType, Dict[str, int]]: ...


@overload
def rand_equation(
    n: int,
    regularity: int,
    n_out: int = ...,
    d_min: int = ...,
    d_max: int = ...,
    seed: Optional[int] = ...,
    global_dim: bool = ...,
    return_size_dict: Literal[False] = ...,
) -> Tuple[str, PathType]: ...


def rand_equation(
    n: int,
    regularity: int,
    n_out: int = 0,
    d_min: int = 2,
    d_max: int = 9,
    seed: Optional[int] = None,
    global_dim: bool = False,
    return_size_dict: bool = False,
) -> Union[Tuple[str, PathType, Dict[str, int]], Tuple[str, PathType]]:
    """Generate a random contraction and shapes.

    Parameters:
        n: Number of array arguments.
        regularity: 'Regularity' of the contraction graph. This essentially determines how
            many indices each tensor shares with others on average.
        n_out: Number of output indices (i.e. the number of non-contracted indices).
            Defaults to 0, i.e., a contraction resulting in a scalar.
        d_min: Minimum dimension size.
        d_max: Maximum dimension size.
        seed: If not None, seed numpy's random generator with this.
        global_dim: Add a global, 'broadcast', dimension to every operand.
        return_size_dict: Return the mapping of indices to sizes.

    Returns:
        eq: The equation string.
        shapes: The array shapes.
        size_dict: The dict of index sizes, only returned if ``return_size_dict=True``.

    Examples:
        ```python
        >>> eq, shapes = rand_equation(n=10, regularity=4, n_out=5, seed=42)
        >>> eq
        'oyeqn,tmaq,skpo,vg,hxui,n,fwxmr,hitplcj,kudlgfv,rywjsb->cebda'

        >>> shapes
        [(9, 5, 4, 5, 4),
        (4, 4, 8, 5),
        (9, 4, 6, 9),
        (6, 6),
        (6, 9, 7, 8),
        (4,),
        (9, 3, 9, 4, 9),
        (6, 8, 4, 6, 8, 6, 3),
        (4, 7, 8, 8, 6, 9, 6),
        (9, 5, 3, 3, 9, 5)]
        ```
    """
    np = pytest.importorskip("numpy")
    if seed is not None:
        np.random.seed(seed)

    # total number of indices
    num_inds = n * regularity // 2 + n_out
    inputs = ["" for _ in range(n)]
    output = []

    size_dict = {get_symbol(i): np.random.randint(d_min, d_max + 1) for i in range(num_inds)}

    # generate a list of indices to place either once or twice
    def gen():
        for i, ix in enumerate(size_dict):
            # generate an outer index
            if i < n_out:
                output.append(ix)
                yield ix
            # generate a bond
            else:
                yield ix
                yield ix

    # add the indices randomly to the inputs
    for i, ix in enumerate(np.random.permutation(list(gen()))):
        # make sure all inputs have at least one index
        if i < n:
            inputs[i] += ix
        else:
            # don't add any traces on same op
            where = np.random.randint(0, n)
            while ix in inputs[where]:
                where = np.random.randint(0, n)

            inputs[where] += ix

    # possibly add the same global dim to every arg
    if global_dim:
        gdim = get_symbol(num_inds)
        size_dict[gdim] = np.random.randint(d_min, d_max + 1)
        for i in range(n):
            inputs[i] += gdim
        output += gdim

    # randomly transpose the output indices and form equation
    output = "".join(np.random.permutation(output))  # type: ignore
    eq = "{}->{}".format(",".join(inputs), output)

    # make the shapes
    shapes = [tuple(size_dict[ix] for ix in op) for op in inputs]

    ret = (eq, shapes)

    if return_size_dict:
        return ret + (size_dict,)
    else:
        return ret


def build_arrays_from_tuples(path: PathType) -> List[Any]:
    """Build random numpy arrays from a path.

    Parameters:
        path: The path to build arrays from.

    Returns:
    The resulting arrays.
    """
    np = pytest.importorskip("numpy")

    return [np.random.rand(*x) for x in path]