cotengra.utils

Various utilities for cotengra.

Attributes

_einsum_symbols_base

Exceptions

BadTrial

Use this to indicate that a trial contraction tree was bad.

Classes

oset	An ordered set which stores elements as the keys of dict (ordered as of
MaxCounter	Simple class to keep track of the maximum in a likely changing
BitSet
BitMembers
DiskDict	A simple persistent dict. The keys should be filesystem compatible
GumbelBatchedGenerator	Non numpy version of gumbel number generator.
Contraction	Built-in immutable sequence.

Functions

getter(index)	Adapted from toolz.
deprecated(fn, old_name, new_name)
prod(it)	Compute the product of sequence of numbers it.
get_rng([seed])	Get a source of random numbers.
compute_size_by_dict(indices, size_dict)	Computes the product of sizes of indices based on size_dict.
get_symbol(i)	Get the symbol corresponding to int i - runs through the usual 52
get_symbol_map(inputs)	Get a mapping of arbitrary hashable 'indices' to single unicode symbols,
rand_equation(n, reg[, n_out, n_hyper_in, ...])	A more advanced version of opt_einsum.testing.rand_equation that
tree_equation(n[, d_min, d_max, n_outer, seed])	Create a random contraction equation that corresponds to a tree.
networkx_graph_to_equation(G[, d_min, d_max, seed])	Turn a networkx graph into a cotengra style contraction, randomly
randreg_equation(n, reg[, d_min, d_max, seed])	Create a random contraction equation that corresponds to a random
perverse_equation(n[, num_indices, min_rank, ...])	Create a weird but valid einsum equation with lots of hyper-edges,
rand_tree(n, reg[, n_out, n_hyper_in, n_hyper_out, ...])	Get a random contraction tree (note, not a tree like equation).
lattice_equation(dims[, cyclic, d_min, d_max, seed])	Create a random contraction equation that corresponds to a lattice.
find_output_str(lhs)	Compute the output string from the left-hand side only of an equation.
eq_to_inputs_output(eq)	Convert a einsum equation into an explicit list of list of characters
inputs_output_to_eq(inputs, output[, canonicalize])	Convert an explicit list of inputs and output to a str einsum equation.
shapes_inputs_to_size_dict(shapes, inputs)	Convert a list of shapes and inputs to a size dictionary.
make_rand_size_dict_from_inputs(inputs[, d_min, ...])	Get a random size dictionary for a given set of inputs.
make_shapes_from_inputs(inputs, size_dict)	Make example shapes to match inputs and index sizes.
make_arrays_from_inputs(inputs, size_dict[, seed, dtype])	Make example arrays to match inputs and index sizes.
make_arrays_from_eq(eq[, d_min, d_max, seed, ...])	Create a set of example arrays to match an einsum equation directly.
find_output_from_inputs(inputs)	Find the output indices for a given set of inputs. The outputs are
is_edge_path(optimize)	Check if the optimize path is a list of indices or a single string.
canonicalize_inputs(inputs[, output, shapes, ...])	Return a canonicalized version of the inputs and output, with the
convert_from_interleaved(args)	Convert from interleaved format array0, input0, array1, input1, ...
check_ellipsis(term)	Check if a einsum term has exactly one ellipsis ('...') or else no
parse_equation_ellipses(eq, shapes[, tuples])
parse_einsum_input(args[, shapes, tuples, constants])	Reproduce einsum input parsing, which handles both interleaved input,
save_to_json(inputs, output, size_dict, filename)	Save a contraction to a json file.
load_from_json(filename)	Load a contraction from a json file.

Module Contents

cotengra.utils.getter(index)

Adapted from toolz.

cotengra.utils.deprecated(fn, old_name, new_name)

cotengra.utils.prod(it)

Compute the product of sequence of numbers it.

class cotengra.utils.oset(it=())

An ordered set which stores elements as the keys of dict (ordered as of python 3.6). ‘A few times’ slower than using a set directly for small sizes, but makes everything deterministic.

__slots__ = ('_d',)

_d

classmethod _from_dict(d)

classmethod from_dict(d)

Public method makes sure to copy incoming dictionary.

copy()

add(k)

discard(k)

remove(k)

clear()

update(*others)

union(*others)

intersection_update(*others)

intersection(*others)

difference_update(*others)

difference(*others)

symmetric_difference(other)

__eq__(other)

__or__(other)

__ior__(other)

__and__(other)

__iand__(other)

__sub__(other)

__isub__(other)

__len__()

__iter__()

__contains__(x)

__repr__()

class cotengra.utils.MaxCounter(it=None)

Simple class to keep track of the maximum in a likely changing sequence of elements.

Parameters:

it (None or sequence of hashable, optional) – The initial items to add.

Examples

>>> mc = MaxCounter([1, 2, 3, 3])
>>> mc.max()
3

>>> mc.discard(3)
>>> mc.max()
3

>>> mc.discard(3)
>>> mc.max()
2

>>> mc.add(10)
>>> mc.max()
10

__slots__ = ('_c', '_max_element')

_c

copy()

discard(x)

Discard element x and possibly update the maximum.

add(x)

Add element x and possibly update the maximum.

max()

The maximum element in this list.

class cotengra.utils.BitSet(it)

__slots__ = ('members', 'map', 'size', 'infimum', 'supremum', 'hashkey')

members

map

size

supremum

infimum

hashkey

asint(elem)

fromint(n)

frommembers(it=())

__call__

class cotengra.utils.BitMembers

__slots__ = ('i', 'bitset')

classmethod fromint(bitset, n)

classmethod frommembers(bitset, it=())

__int__()

__hash__

__eq__(other)

__len__()

__iter__()

add(elem)

clear()

copy()

__bool__()

__contains__(elem)

discard(elem)

remove(elem)

difference_update(*others)

__isub__

difference(*others)

__sub__

intersection_update(*others)

__iand__

intersection(*others)

__and__

isdisjoint(other)

issubset(other)

issuperset(other)

symmetric_difference_update(other)

__ixor__

symmetric_difference(other)

__xor__

update(*others)

__ior__

union(*others)

__or__

__repr__()

class cotengra.utils.DiskDict(directory=None, max_retries=3, retry_delay=0.01)

A simple persistent dict. The keys should be filesystem compatible strings, or tuples of strings, in which case it will be used as a sub-directory structure. The values should be picklable. The directory will be created if it does not exist. Values are loaded into memory once they are accessed.

Parameters:

• directory (str or pathlib.Path, optional) – The directory to store the files in. If None, the files will not be stored on disk and only kept in memory.

• max_retries (int, optional) – The maximum number of retries to read a file if it is not completely written yet. Default is 3.

• retry_delay (float, optional) – The delay between retries in seconds. Default is 0.01.

__slots__ = ('_directory', '_mem_cache', '_path', 'max_retries', 'retry_delay')

_mem_cache

_directory = None

max_retries = 3

retry_delay

clear()

Clear the memory cache and delete all value files, but leave directory structure.

cleanup(delete_dir=False)

Delete all files and subdirectories and optionally delete the root directory.

__contains__(k)

__setitem__(k, v)

__delitem__(k)

__getitem__(k)

get(k, default=None)

keys()

values()

items()

cotengra.utils.get_rng(seed=None)

Get a source of random numbers.

Parameters:

seed (None or int or random.Random, optional) – The seed for the random number generator. If None, use the default random number generator. If an integer, use a new random number generator with the given seed. If a random.Random instance, use that instance.

class cotengra.utils.GumbelBatchedGenerator(seed=None)

Non numpy version of gumbel number generator.

rng

__call__()

exception cotengra.utils.BadTrial

Bases: Exception

Use this to indicate that a trial contraction tree was bad.

cotengra.utils.compute_size_by_dict(indices, size_dict)

Computes the product of sizes of indices based on size_dict.

Parameters:

• indices (iterable[str] or iterable[int]) – The indices of the term.

• size_dict (dict or list) – Mapping (or list/tuple if the indices are indexing integers, which can be slightly faster) of indices to sizes.

Returns:

d – The resulting product.

Return type:

int

Examples

>>> compute_size_by_dict('abbc', {'a': 2, 'b':3, 'c':5})
90

cotengra.utils._einsum_symbols_base = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'

cotengra.utils.get_symbol(i)

Get the symbol corresponding to int i - runs through the usual 52 letters before resorting to unicode characters, starting at chr(192) and skipping surrogates.

Examples

get_symbol(2) #> ‘c’

get_symbol(200) #> ‘Ŕ’

get_symbol(20000) #> ‘京’

cotengra.utils.get_symbol_map(inputs)

Get a mapping of arbitrary hashable ‘indices’ to single unicode symbols, matching the canonicalization of the expression.

class cotengra.utils.Contraction

Bases: collections.namedtuple('Contraction', ('inputs', 'output', 'shapes', 'size_dict'))

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable’s items.

If the argument is a tuple, the return value is the same object.

property eq

Get the einsum equation string corresponding to this contraction.

make_arrays(seed=None, dtype='float64')

Make example arrays to match the inputs, shapes and size_dict of this contraction.

Parameters:

• seed (None or int, optional) – Seed for random generator for repeatibility, by default None.

• dtype (str or np.dtype, optional) – The dtype of the arrays to create, by default ‘float64’.

Returns:

arrays

Return type:

list[np.ndarray]

get_hypergraph()

Get the hypergraph corresponding to this contraction.

plot(**kwargs)

Draw the hypergraph corresponding to this contraction.

draw

cotengra.utils.rand_equation(n, reg, n_out=0, n_hyper_in=0, n_hyper_out=0, d_min=2, d_max=3, seed=None)

A more advanced version of opt_einsum.testing.rand_equation that can also generate both inner and outer hyper-edges. Mostly useful for generating test instances covering all edge cases.

Parameters:

• n (int) – The number of tensors.

• reg (int) – The average number of indices per tensor if no hyper-edges, i.e. total number of inds = n * reg // 2.

• n_out (int, optional) – The number of output indices.

• n_hyper_in (int, optional) – The number of inner hyper-indices.

• n_hyper_out (int, optional) – The number of outer hyper-indices.

• d_min (int, optional) – The minimum dimension size.

• d_max (int, optional) – The maximum dimension size.

• seed (None or int, optional) – Seed for np.random for repeatibility.

Returns:

• inputs (list[list[str]])

• output (list[str])

• shapes (list[tuple[int]])

• size_dict (dict[str, int])

cotengra.utils.tree_equation(n, d_min=2, d_max=3, n_outer=0, seed=None)

Create a random contraction equation that corresponds to a tree.

Parameters:

• n (int) – The number of tensors.

• d_min (int, optional) – The minimum size of an index.

• d_max (int, optional) – The maximum size of an index.

• n_outer (int, optional) – The number of outer indices.

cotengra.utils.networkx_graph_to_equation(G, d_min=2, d_max=3, seed=None)

Turn a networkx graph into a cotengra style contraction, randomly sampling index sizes for each edge.

Parameters:

• G (nx.Graph) – The graph to convert.

• d_min (int, optional) – The minimum size of an index.

• d_max (int, optional) – The maximum size of an index.

• seed (None, int or np.random.Generator, optional) – Seed for repeatibility.

Returns:

• inputs (list[list[str]])

• output (list[str])

• shapes (list[tuple[int]])

• size_dict (dict[str, int])

cotengra.utils.randreg_equation(n, reg, d_min=2, d_max=3, seed=None)

Create a random contraction equation that corresponds to a random regular graph.

Parameters:

• n (int) – The number of terms.

• reg (int) – The degree of the graph.

• d_min (int, optional) – The minimum size of an index.

• d_max (int, optional) – The maximum size of an index.

• seed (None or int, optional) – Seed for networkx and np.random.default_rng for repeatibility.

Returns:

• inputs (list[list[str]])

• output (list[str])

• shapes (list[tuple[int]])

• size_dict (dict[str, int])

cotengra.utils.perverse_equation(n, num_indices=3, min_rank=0, max_rank=6, d_min=2, d_max=3, n_outer=2, seed=None)

Create a weird but valid einsum equation with lots of hyper-edges, repeated indices, scalars and repeated terms.

Parameters:

• n (int) – The number of tensors.

• num_indices (int, optional) – The number of indices to use.

• min_rank (int, optional) – The minimum rank of a tensor.

• max_rank (int, optional) – The maximum rank of a tensor.

• d_min (int, optional) – The minimum size of an index.

• d_max (int, optional) – The maximum size of an index.

• n_outer (int, optional) – The number of outer indices.

• seed (None or int, optional) – Seed for random.Random for repeatibility.

cotengra.utils.rand_tree(n, reg, n_out=0, n_hyper_in=0, n_hyper_out=0, d_min=2, d_max=3, seed=None, optimize='greedy')

Get a random contraction tree (note, not a tree like equation).

cotengra.utils.lattice_equation(dims, cyclic=False, d_min=2, d_max=None, seed=None)

Create a random contraction equation that corresponds to a lattice.

Parameters:

• dims (sequence of int) – The size of each dimension, with the dimensionality being the length of the sequence.

• cyclic (bool or sequence of bool, optional) – Whether each dimension is cyclic or not. If a sequence, must be the same length as dims.

• d_min (int, optional) – The minimum size of an index.

• d_max (int, optional) – The maximum size of an index. If None, defaults to d_min, i.e. all indices are the same size.

• seed (None or int, optional) – Seed for random.Random for repeatibility.

cotengra.utils.find_output_str(lhs)

Compute the output string from the left-hand side only of an equation. This is any indices that appear only once in the left-hand side, in sorted order.

Parameters:

lhs (str) – The comma separated list of indices on the left-hand side of an einsum equation.

Returns:

rhs – The output string of the einsum equation.

Return type:

str

Examples

>>> find_output_str('cb,ba')
'ac'

cotengra.utils.eq_to_inputs_output(eq)

Convert a einsum equation into an explicit list of list of characters and an output list of characters.

Parameters:

eq (str) – The einsum equation, with or without output.

Returns:

• inputs (list[list[str]]) – The input terms.

• output (list[str]) – The output term.

cotengra.utils.inputs_output_to_eq(inputs, output, canonicalize=False)

Convert an explicit list of inputs and output to a str einsum equation.

Parameters:

• inputs (list[list[str]]) – The input terms.

• output (list[str]) – The output term.

• canonicalize (bool, optional) – Whether to canonicalize (map into [a-zA-Z]) equation, by default False.

Returns:

eq – The einsum equation.

Return type:

str

cotengra.utils.shapes_inputs_to_size_dict(shapes, inputs)

Convert a list of shapes and inputs to a size dictionary.

Parameters:

• shapes (list[tuple[int]]) – The shapes of each input.

• inputs (list[list[str]]) – The input terms.

Returns:

size_dict – The index size dictionary.

Return type:

dict[str, int]

cotengra.utils.make_rand_size_dict_from_inputs(inputs, d_min=2, d_max=3, seed=None)

Get a random size dictionary for a given set of inputs.

Parameters:

• inputs (list[list[str]]) – The input terms.

• d_min (int, optional) – The minimum dimension, by default 2.

• d_max (int, optional) – The maximum dimension, by default 3.

• seed (int, optional) – The random seed, by default None.

Returns:

size_dict – The index size dictionary.

Return type:

dict[str, int]

cotengra.utils.make_shapes_from_inputs(inputs, size_dict)

Make example shapes to match inputs and index sizes.

Parameters:

• inputs (list[list[str]]) – The input terms.

• size_dict (dict[str, int]) – The index size dictionary.

Returns:

shapes – The example shapes.

Return type:

list[tuple[int]]

cotengra.utils.make_arrays_from_inputs(inputs, size_dict, seed=None, dtype='float64')

Make example arrays to match inputs and index sizes.

Parameters:

• inputs (list[list[str]]) – The input terms.

• size_dict (dict[str, int]) – The index size dictionary.

• seed (int, optional) – The random seed, by default None.

• dtype ({'float32', 'float64', 'complex64', 'complex128'}, optional) – The dtype of the arrays, by default ‘float64’.

Returns:

arrays – The example arrays.

Return type:

list[numpy.ndarray]

cotengra.utils.make_arrays_from_eq(eq, d_min=2, d_max=3, seed=None, size_dict=None, dtype='float64')

Create a set of example arrays to match an einsum equation directly.

Parameters:

• eq (str) – The einsum equation.

• d_min (int, optional) – The minimum dimension, by default 2.

• d_max (int, optional) – The maximum dimension, by default 3.

• seed (int, optional) – The random seed, by default None.

• size_dict (dict[str, int], optional) – The index size dictionary. If supplied this is used instead of generating a random one.

• dtype ({'float32', 'float64', 'complex64', 'complex128'}, optional) – The dtype of the arrays, by default ‘float64’.

Returns:

arrays – The example arrays.

Return type:

list[numpy.ndarray]

cotengra.utils.find_output_from_inputs(inputs)

Find the output indices for a given set of inputs. The outputs are calculated as the set of indices that appear only once, in the order they appear in the inputs. This is different to einsum where they in sorted order since we only require the indices to be hashable.

Parameters:

inputs (Sequence[Sequence[Hashable]]) – The input terms.

Returns:

output

Return type:

tuple[Hashable]

cotengra.utils.is_edge_path(optimize)

Check if the optimize path is a list of indices or a single string.

cotengra.utils.canonicalize_inputs(inputs, output=None, shapes=None, size_dict=None, optimize=None)

Return a canonicalized version of the inputs and output, with the indices labelled from ‘a’, ‘b’, ‘c’, … according to the order they appear in the equation.

If output is not provided, it will be computed as the set of indices that appear once, in the order they appear on the inputs (note that this is different to einsum where they in sorted order since we only require the indices to be hashable and not comparable).

If either shapes or size_dict is provided, then also compute the corresponding canonicalized size_dict for new index labels.

Parameters:

• inputs (Sequence[Sequence[Hashable]]) – The input terms.

• output (Sequence[Hashable], optional) – The output term. If not supplied it will be computed.

• shapes (None or Sequence[tuple[int]], optional) – The shapes of each input.

• size_dict (None or dict[Hashable, int], optional) – The index size dictionary.

Returns:

• new_inputs (tuple[tuple[str]]) – The canonicalized input terms.

• new_output (tuple[str]) – The canonicalized output term.

• new_size_dict (dict[str, int] or None) – The canonicalized index size dictionary, None if size_dict or shapes was not provided.

• new_optimize (None or str or Sequence[str] or Sequence[int], optional) – The canonicalized optimize path, None if optimize was not provided.

cotengra.utils.convert_from_interleaved(args)

Convert from interleaved format array0, input0, array1, input1, ... to a single equation and list of arrays. The arrays can just be shapes.

cotengra.utils.check_ellipsis(term)

Check if a einsum term has exactly one ellipsis (’…’) or else no dots at all.

cotengra.utils.parse_equation_ellipses(eq, shapes, tuples=False)

cotengra.utils.parse_einsum_input(args, shapes=False, tuples=False, constants=None)

Reproduce einsum input parsing, which handles both interleaved input, ellipsis expansions, and output calculation. The main processing is cached.

Parameters:

• args (tuple) – The arguments to einsum.

• shapes (bool, optional) – Whether arrays are being supplied (False) or shapes (True).

• tuples (bool, optional) – Whether to return parsed indices as strings or tuples.

cotengra.utils.save_to_json(inputs, output, size_dict, filename)

Save a contraction to a json file.

Parameters:

• inputs (list[list[str]]) – The input terms.

• output (list[str]) – The output term.

• size_dict (dict[str, int]) – The index size dictionary.

• filename (str) – The filename to save to.

cotengra.utils.load_from_json(filename)

Load a contraction from a json file.

Parameters:

filename (str) – The filename to load from.

Returns:

• inputs (list[list[str]]) – The input terms.

• output (list[str]) – The output term.

• size_dict (dict[str, int]) – The index size dictionary.