cotengra.utils ============== .. py:module:: cotengra.utils .. autoapi-nested-parse:: Various utilities for cotengra. Attributes ---------- .. autoapisummary:: cotengra.utils.NODE_TYPE cotengra.utils._einsum_symbols_base Exceptions ---------- .. autoapisummary:: cotengra.utils.BadTrial Classes ------- .. autoapisummary:: cotengra.utils.oset cotengra.utils.MaxCounter cotengra.utils.BitSet cotengra.utils.BitMembers cotengra.utils.BitSetInt cotengra.utils.DiskDict cotengra.utils.GumbelBatchedGenerator cotengra.utils.Contraction Functions --------- .. autoapisummary:: cotengra.utils.getter cotengra.utils.deprecated cotengra.utils.prod cotengra.utils.fast_bitmember_iter cotengra.utils.fast_bitmember_len cotengra.utils.node_from_seq cotengra.utils.get_rng cotengra.utils.compute_size_by_dict cotengra.utils.get_symbol cotengra.utils.get_symbol_map cotengra.utils.rand_equation cotengra.utils.tree_equation cotengra.utils.networkx_graph_to_equation cotengra.utils.randreg_equation cotengra.utils.perverse_equation cotengra.utils.rand_tree cotengra.utils.lattice_equation cotengra.utils.find_output_str cotengra.utils.eq_to_inputs_output cotengra.utils.inputs_output_to_eq cotengra.utils.shapes_inputs_to_size_dict cotengra.utils.make_rand_size_dict_from_inputs cotengra.utils.make_shapes_from_inputs cotengra.utils.make_arrays_from_inputs cotengra.utils.make_arrays_from_eq cotengra.utils.find_output_from_inputs cotengra.utils.is_edge_path cotengra.utils.canonicalize_inputs cotengra.utils.convert_from_interleaved cotengra.utils.check_ellipsis cotengra.utils.parse_equation_ellipses cotengra.utils.parse_einsum_input cotengra.utils.save_to_json cotengra.utils.load_from_json Module Contents --------------- .. py:function:: getter(index) Adapted from `toolz`. .. py:function:: deprecated(fn, old_name, new_name) .. py:function:: prod(it) Compute the product of sequence of numbers ``it``. .. py:class:: 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. .. py:attribute:: __slots__ :value: ('_d',) .. py:attribute:: _d .. py:method:: _from_dict(d) :classmethod: .. py:method:: from_dict(d) :classmethod: Public method makes sure to copy incoming dictionary. .. py:method:: copy() .. py:method:: add(k) .. py:method:: discard(k) .. py:method:: remove(k) .. py:method:: clear() .. py:method:: update(*others) .. py:method:: union(*others) .. py:method:: intersection_update(*others) .. py:method:: intersection(*others) .. py:method:: difference_update(*others) .. py:method:: difference(*others) .. py:method:: symmetric_difference(other) .. py:method:: __eq__(other) .. py:method:: __or__(other) .. py:method:: __ior__(other) .. py:method:: __and__(other) .. py:method:: __iand__(other) .. py:method:: __sub__(other) .. py:method:: __isub__(other) .. py:method:: __len__() .. py:method:: __iter__() .. py:method:: __contains__(x) .. py:method:: __repr__() .. py:class:: MaxCounter(it=None) Simple class to keep track of the maximum in a likely changing sequence of elements. :param it: The initial items to add. :type it: None or sequence of hashable, optional .. rubric:: 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 .. py:attribute:: __slots__ :value: ('_c', '_max_element') .. py:attribute:: _c .. py:method:: copy() .. py:method:: discard(x) Discard element ``x`` and possibly update the maximum. .. py:method:: add(x) Add element ``x`` and possibly update the maximum. .. py:method:: max() The maximum element in this list. .. py:class:: BitSet(it) .. py:attribute:: __slots__ :value: ('members', 'map', 'size', 'infimum', 'supremum', 'hashkey') .. py:attribute:: members .. py:attribute:: map .. py:attribute:: size .. py:attribute:: supremum .. py:attribute:: infimum .. py:attribute:: hashkey .. py:method:: asint(elem) .. py:method:: fromint(n) .. py:method:: frommembers(it=()) .. py:attribute:: __call__ .. py:class:: BitMembers .. py:attribute:: __slots__ :value: ('i', 'bitset') .. py:method:: fromint(bitset, n) :classmethod: .. py:method:: frommembers(bitset, it=()) :classmethod: .. py:method:: __int__() .. py:attribute:: __hash__ .. py:method:: __eq__(other) .. py:method:: __len__() .. py:method:: __iter__() .. py:method:: add(elem) .. py:method:: clear() .. py:method:: copy() .. py:method:: __bool__() .. py:method:: __contains__(elem) .. py:method:: discard(elem) .. py:method:: remove(elem) .. py:method:: difference_update(*others) .. py:attribute:: __isub__ .. py:method:: difference(*others) .. py:attribute:: __sub__ .. py:method:: intersection_update(*others) .. py:attribute:: __iand__ .. py:method:: intersection(*others) .. py:attribute:: __and__ .. py:method:: isdisjoint(other) .. py:method:: issubset(other) .. py:method:: issuperset(other) .. py:method:: symmetric_difference_update(other) .. py:attribute:: __ixor__ .. py:method:: symmetric_difference(other) .. py:attribute:: __xor__ .. py:method:: update(*others) .. py:attribute:: __ior__ .. py:method:: union(*others) .. py:attribute:: __or__ .. py:method:: __repr__() .. py:class:: BitSetInt Bases: :py:obj:`int` A bitset for keeping track of dense sets of integers. Subclass of ``int`` that overrides ``-``, ``&`` and ``|`` and provides ``__len__`` and ``__iter__``. As of python 3.8 doesn't seem much better than using frozenset[int] (slower and limited memory savings). :param it: If sequence of ``int``, treat these as the set members. If raw ``int`` interpret as bitset directly. :type it: Iterable[int] or int .. py:method:: __hash__() Return hash(self). .. py:method:: infimum() :classmethod: .. py:method:: supremum(size) :classmethod: .. py:method:: __len__() .. py:method:: iter_sparse() .. py:method:: __iter__() .. py:method:: __contains__(i) .. py:method:: __sub__(i) Return self-value. .. py:method:: difference(i) .. py:method:: __and__(i) Return self&value. .. py:method:: intersection(i) .. py:method:: __or__(i) Return self|value. .. py:method:: union() .. py:method:: __repr__() Return repr(self). .. py:function:: fast_bitmember_iter(self) .. py:function:: fast_bitmember_len(self) .. py:data:: NODE_TYPE :value: 'frozenset[int]' .. py:function:: node_from_seq(it) .. py:class:: 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. :param directory: The directory to store the files in. If None, the files will not be stored on disk and only kept in memory. :type directory: str or pathlib.Path, optional :param max_retries: The maximum number of retries to read a file if it is not completely written yet. Default is 3. :type max_retries: int, optional :param retry_delay: The delay between retries in seconds. Default is 0.01. :type retry_delay: float, optional .. py:attribute:: __slots__ :value: ('_directory', '_mem_cache', '_path', 'max_retries', 'retry_delay') .. py:attribute:: _mem_cache .. py:attribute:: _directory :value: None .. py:attribute:: max_retries :value: 3 .. py:attribute:: retry_delay .. py:method:: clear() .. py:method:: cleanup(delete_dir=False) .. py:method:: __contains__(k) .. py:method:: __setitem__(k, v) .. py:method:: __getitem__(k) .. py:function:: get_rng(seed=None) Get a source of random numbers. :param seed: 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. :type seed: None or int or random.Random, optional .. py:class:: GumbelBatchedGenerator(seed=None) Non numpy version of gumbel number generator. .. py:attribute:: rng .. py:method:: __call__() .. py:exception:: BadTrial Bases: :py:obj:`Exception` Use this to indicate that a trial contraction tree was bad. .. py:function:: compute_size_by_dict(indices, size_dict) Computes the product of sizes of ``indices`` based on ``size_dict``. :param indices: The indices of the term. :type indices: iterable[str] or iterable[int] :param size_dict: Mapping (or list/tuple if the indices are indexing integers, which can be slightly faster) of indices to sizes. :type size_dict: dict or list :returns: **d** -- The resulting product. :rtype: int .. rubric:: Examples >>> compute_size_by_dict('abbc', {'a': 2, 'b':3, 'c':5}) 90 .. py:data:: _einsum_symbols_base :value: 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ' .. py:function:: 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. .. rubric:: Examples get_symbol(2) #> 'c' get_symbol(200) #> 'Ŕ' get_symbol(20000) #> '京' .. py:function:: get_symbol_map(inputs) Get a mapping of arbitrary hashable 'indices' to single unicode symbols, matching the canonicalization of the expression. .. py:class:: Contraction Bases: :py:obj:`tuple` .. py:attribute:: inputs .. py:attribute:: output .. py:attribute:: shapes .. py:attribute:: size_dict .. py:function:: 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. :param n: The number of tensors. :type n: int :param reg: The average number of indices per tensor if no hyper-edges, i.e. total number of inds ``= n * reg // 2``. :type reg: int :param n_out: The number of output indices. :type n_out: int, optional :param n_hyper_in: The number of inner hyper-indices. :type n_hyper_in: int, optional :param n_hyper_out: The number of outer hyper-indices. :type n_hyper_out: int, optional :param d_min: The minimum dimension size. :type d_min: int, optional :param d_max: The maximum dimension size. :type d_max: int, optional :param seed: Seed for ``np.random`` for repeatibility. :type seed: None or int, optional :returns: * **inputs** (*list[list[str]]*) * **output** (*list[str]*) * **shapes** (*list[tuple[int]]*) * **size_dict** (*dict[str, int]*) .. py:function:: tree_equation(n, d_min=2, d_max=3, n_outer=0, seed=None) Create a random contraction equation that corresponds to a tree. :param n: The number of tensors. :type n: int :param d_min: The minimum size of an index. :type d_min: int, optional :param d_max: The maximum size of an index. :type d_max: int, optional :param n_outer: The number of outer indices. :type n_outer: int, optional .. py:function:: 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. :param G: The graph to convert. :type G: nx.Graph :param d_min: The minimum size of an index. :type d_min: int, optional :param d_max: The maximum size of an index. :type d_max: int, optional :param seed: Seed for repeatibility. :type seed: None, int or np.random.Generator, optional :returns: * **inputs** (*list[list[str]]*) * **output** (*list[str]*) * **shapes** (*list[tuple[int]]*) * **size_dict** (*dict[str, int]*) .. py:function:: randreg_equation(n, reg, d_min=2, d_max=3, seed=None) Create a random contraction equation that corresponds to a random regular graph. :param n: The number of terms. :type n: int :param reg: The degree of the graph. :type reg: int :param d_min: The minimum size of an index. :type d_min: int, optional :param d_max: The maximum size of an index. :type d_max: int, optional :param seed: Seed for ``networkx`` and ``np.random.default_rng`` for repeatibility. :type seed: None or int, optional :returns: * **inputs** (*list[list[str]]*) * **output** (*list[str]*) * **shapes** (*list[tuple[int]]*) * **size_dict** (*dict[str, int]*) .. py:function:: 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. :param n: The number of tensors. :type n: int :param num_indices: The number of indices to use. :type num_indices: int, optional :param min_rank: The minimum rank of a tensor. :type min_rank: int, optional :param max_rank: The maximum rank of a tensor. :type max_rank: int, optional :param d_min: The minimum size of an index. :type d_min: int, optional :param d_max: The maximum size of an index. :type d_max: int, optional :param n_outer: The number of outer indices. :type n_outer: int, optional :param seed: Seed for ``random.Random`` for repeatibility. :type seed: None or int, optional .. py:function:: 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). .. py:function:: lattice_equation(dims, cyclic=False, d_min=2, d_max=None, seed=None) Create a random contraction equation that corresponds to a lattice. :param dims: The size of each dimension, with the dimensionality being the length of the sequence. :type dims: sequence of int :param cyclic: Whether each dimension is cyclic or not. If a sequence, must be the same length as ``dims``. :type cyclic: bool or sequence of bool, optional :param d_min: The minimum size of an index. :type d_min: int, optional :param d_max: The maximum size of an index. If ``None``, defaults to ``d_min``, i.e. all indices are the same size. :type d_max: int, optional :param seed: Seed for ``random.Random`` for repeatibility. :type seed: None or int, optional .. py:function:: 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. :param lhs: The comma separated list of indices on the left-hand side of an einsum equation. :type lhs: str :returns: **rhs** -- The output string of the einsum equation. :rtype: str .. rubric:: Examples >>> find_output_str('cb,ba') 'ac' .. py:function:: eq_to_inputs_output(eq) Convert a einsum equation into an explicit list of list of characters and an output list of characters. :param eq: The einsum equation, with or without output. :type eq: str :returns: * **inputs** (*list[list[str]]*) -- The input terms. * **output** (*list[str]*) -- The output term. .. py:function:: inputs_output_to_eq(inputs, output, canonicalize=False) Convert an explicit list of inputs and output to a str einsum equation. :param inputs: The input terms. :type inputs: list[list[str]] :param output: The output term. :type output: list[str] :param canonicalize: Whether to canonicalize (map into [a-zA-Z]) equation, by default False. :type canonicalize: bool, optional :returns: **eq** -- The einsum equation. :rtype: str .. py:function:: shapes_inputs_to_size_dict(shapes, inputs) Convert a list of shapes and inputs to a size dictionary. :param shapes: The shapes of each input. :type shapes: list[tuple[int]] :param inputs: The input terms. :type inputs: list[list[str]] :returns: **size_dict** -- The index size dictionary. :rtype: dict[str, int] .. py:function:: 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. :param inputs: The input terms. :type inputs: list[list[str]] :param d_min: The minimum dimension, by default 2. :type d_min: int, optional :param d_max: The maximum dimension, by default 3. :type d_max: int, optional :param seed: The random seed, by default None. :type seed: int, optional :returns: **size_dict** -- The index size dictionary. :rtype: dict[str, int] .. py:function:: make_shapes_from_inputs(inputs, size_dict) Make example shapes to match inputs and index sizes. :param inputs: The input terms. :type inputs: list[list[str]] :param size_dict: The index size dictionary. :type size_dict: dict[str, int] :returns: **shapes** -- The example shapes. :rtype: list[tuple[int]] .. py:function:: make_arrays_from_inputs(inputs, size_dict, seed=None, dtype='float64') Make example arrays to match inputs and index sizes. :param inputs: The input terms. :type inputs: list[list[str]] :param size_dict: The index size dictionary. :type size_dict: dict[str, int] :param seed: The random seed, by default None. :type seed: int, optional :param dtype: The dtype of the arrays, by default 'float64'. :type dtype: {'float32', 'float64', 'complex64', 'complex128'}, optional :returns: **arrays** -- The example arrays. :rtype: list[numpy.ndarray] .. py:function:: make_arrays_from_eq(eq, d_min=2, d_max=3, seed=None, dtype='float64') Create a set of example arrays to match an einsum equation directly. :param eq: The einsum equation. :type eq: str :param d_min: The minimum dimension, by default 2. :type d_min: int, optional :param d_max: The maximum dimension, by default 3. :type d_max: int, optional :param seed: The random seed, by default None. :type seed: int, optional :param dtype: The dtype of the arrays, by default 'float64'. :type dtype: {'float32', 'float64', 'complex64', 'complex128'}, optional :returns: **arrays** -- The example arrays. :rtype: list[numpy.ndarray] .. py:function:: 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. :param inputs: The input terms. :type inputs: Sequence[Sequence[Hashable]] :returns: **output** :rtype: tuple[Hashable] .. py:function:: is_edge_path(optimize) Check if the optimize path is a list of indices or a single string. .. py:function:: 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. :param inputs: The input terms. :type inputs: Sequence[Sequence[Hashable]] :param output: The output term. If not supplied it will be computed. :type output: Sequence[Hashable], optional :param shapes: The shapes of each input. :type shapes: None or Sequence[tuple[int]], optional :param size_dict: The index size dictionary. :type size_dict: None or dict[Hashable, int], optional :returns: * **inputs** (*tuple[tuple[str]]*) -- The canonicalized input terms. * **output** (*tuple[str]*) -- The canonicalized output term. * **size_dict** (*dict[str, int] or None*) -- The canonicalized index size dictionary, ``None`` if ``size_dict`` or ``shapes`` was not provided. .. py:function:: 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. .. py:function:: check_ellipsis(term) Check if a einsum term has exactly one ellipsis ('...') or else no dots at all. .. py:function:: parse_equation_ellipses(eq, shapes, tuples=False) .. py:function:: 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. :param args: The arguments to ``einsum``. :type args: tuple :param shapes: Whether arrays are being supplied (``False``) or shapes (``True``). :type shapes: bool, optional :param tuples: Whether to return parsed indices as strings or tuples. :type tuples: bool, optional .. py:function:: save_to_json(inputs, output, size_dict, filename) Save a contraction to a json file. :param inputs: The input terms. :type inputs: list[list[str]] :param output: The output term. :type output: list[str] :param size_dict: The index size dictionary. :type size_dict: dict[str, int] :param filename: The filename to save to. :type filename: str .. py:function:: load_from_json(filename) Load a contraction from a json file. :param filename: The filename to load from. :type filename: str :returns: * **inputs** (*list[list[str]]*) -- The input terms. * **output** (*list[str]*) -- The output term. * **size_dict** (*dict[str, int]*) -- The index size dictionary.