From f5f2741bc7a2357c330293d0c76df64c0e8360d1 Mon Sep 17 00:00:00 2001 From: Johnnie Gray Date: Thu, 26 Sep 2024 13:54:37 -0700 Subject: [PATCH] contract_compressed: add docstring --- docs/changelog.md | 4 +- quimb/tensor/tensor_core.py | 134 ++++++++++++++++++++++++++++++++++-- 2 files changed, 132 insertions(+), 6 deletions(-) diff --git a/docs/changelog.md b/docs/changelog.md index d76e9dd9..8d689321 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -7,8 +7,8 @@ Release notes for `quimb`. **Enhancements:** -- expose [`qtn.edge_coloring`](quimb.tensor.tensor_arbgeom_tebd.edge_coloring) - as top level function and allow layers to be returned grouped. +- expose [`qtn.edge_coloring`](quimb.tensor.tensor_arbgeom_tebd.edge_coloring) as top level function and allow layers to be returned grouped. +- add docstring for [`tn.contract_compressed`](quimb.tensor.tensor_core.TensorNetwork.contract_compressed) (whats-new-1-8-4)= diff --git a/quimb/tensor/tensor_core.py b/quimb/tensor/tensor_core.py index ccb6737b..87faf791 100644 --- a/quimb/tensor/tensor_core.py +++ b/quimb/tensor/tensor_core.py @@ -7326,22 +7326,23 @@ def gauge_simple_temp( def _contract_compressed_tid_sequence( self, seq, + *, + output_inds=None, max_bond=None, cutoff=1e-10, - output_inds=None, tree_gauge_distance=1, canonize_distance=None, canonize_opts=None, canonize_after_distance=None, canonize_after_opts=None, gauge_boundary_only=True, - compress_opts=None, compress_late=True, compress_mode="auto", compress_min_size=None, compress_span=False, compress_matrices=True, compress_exclude=None, + compress_opts=None, equalize_norms=False, gauges=None, gauge_smudge=1e-6, @@ -7354,6 +7355,8 @@ def _contract_compressed_tid_sequence( progbar=False, inplace=False, ): + """Core routine for performing compressed contraction. + """ tn = self if inplace else self.copy() # options relating to the compression itself @@ -7744,6 +7747,7 @@ def contract_around( def contract_compressed( self, optimize, + *, output_inds=None, max_bond=None, cutoff=1e-10, @@ -7753,12 +7757,13 @@ def contract_compressed( canonize_after_distance=None, canonize_after_opts=None, gauge_boundary_only=True, - compress_late=True, + compress_late=False, + compress_mode="auto", compress_min_size=None, - compress_opts=None, compress_span=True, compress_matrices=True, compress_exclude=None, + compress_opts=None, equalize_norms=False, gauges=None, gauge_smudge=1e-6, @@ -7767,9 +7772,124 @@ def contract_compressed( callback_pre_compress=None, callback_post_compress=None, callback=None, + preserve_tensor=False, progbar=False, + inplace=False, **kwargs, ): + """Contract this tensor network using the hyperoptimized approximate + contraction method introduced in https://arxiv.org/abs/2206.07044. + + Only supports non-hyper tensor networks. + + Parameters + ---------- + optimize : str, sequence, PathOptimizer, ContractionTree + The contraction strategy to use. The options are: + + - a string specifying a preset strategy + - an explicit sequence of tuples specifying the contraction path + - a PathOptimizer object from cotengra + - an explicit ContractionTree object from cotengra + + Note that the strategy should be one that specifically targets + compressed contraction, paths for exact contraction will likely + perform badly. See the cotengra documentation for more details. + output_inds : sequence of str, optional + Output indices. Note that hyper indices are not supported and this + is just for specifying the output order. + max_bond : int, optional + The maximum bond dimension to allow during compression. + cutoff : float, optional + The singular value cutoff to use during compression. + tree_gauge_distance : int, optional + The distance to 'tree gauge' around a pair of tensors before + compressing. Depending on if `compress_mode="basic"` this sets + `canonize_distance` and `canonize_after_distance`. + canonize_distance : int, optional + The distance to canonize around a pair of tensors before + compressing. + canonize_opts : dict, optional + Additional keyword arguments to pass to the canonize routine. + canonize_after_distance : int, optional + The distance to canonize around a pair of tensors after + compressing. + canonize_after_opts : dict, optional + Additional keyword arguments to pass to the canonize routine after + compressing. + gauge_boundary_only : bool, optional + Whether to only gauge the 'boundary' tensors, that is, intermediate + tensors. + compress_late : bool, optional + Whether to compress just before contracting the tensors involved or + immediately after. Early compression is cheaper and a better + default especially for contractions beyond planar. Late compression + leaves more information in the tensors for possibly better quality + gauging and compression. Whilst the largest tensor ('contraction + width') is typically unchanged, the total memory and cost can be + quite a lot higher. + compress_mode : {'auto', 'basic', 'virtual-tree', ...}, optional + How to compress a pair of tensors. If 'auto', then 'basic' is used + if `tree_gauge_distance=0` or `gauges` are supplied, otherwise + 'virtual-tree' is used. See `_compress_between_tids` for other + valid options. + compress_min_size : int, optional + Skip compressing a pair of tensors if their contraction would yield + a tensor smaller than this size. + compress_opts : dict, optional + Additional keyword arguments to pass to the core pariwise + compression routine. + compress_span : bool or int, optional + Whether to compress between tensors that are going to be + contracted. If an `int`, this specifies that if two tensors will be + contracted in the next `compress_span` contractions, then their + bonds should be compressed. + compress_matrices : bool, optional + Whether to compress pairs of tensors that are effectively matrices. + compress_exclude : set[int], optional + An explicit set of tensor ids to exclude from compression. + equalize_norms : bool or float, optional + Whether to equalize the norms of the tensors after each operation. + The overall scaling is accumulated, log10, into `tn.exponent`. If + `True`, at the end this exponent is redistributed. If a float, + this is the target norm to equalize tensors to, e.g. `1.0`, and the + exponent is *not* redistributed, which is useful in the case that + the non-log value is beyond standard precision. + gauges : dict[str, array_like], optional + If supplied, use simple update style gauges during the contraction. + The keys should be indices and the values singular value vectors. + Only bonds present in this dictionary will be gauged. + gauge_smudge : float, optional + If using simple update style gauging, add a small value to the + singular values to avoid singularities. + callback_pre_contract : callable, optional + A function to call before contracting a pair of tensors. It should + have signature `fn(tn, (tid1, tid2))`. + callback_post_contract : callable, optional + A function to call after contracting a pair of tensors. It should + have signature `fn(tn, tid)`. + callback_pre_compress : callable, optional + A function to call before compressing a pair of tensors. It should + have signature `fn(tn, (tid1, tid2))`. + callback_post_compress : callable, optional + A function to call after compressing a pair of tensors. It should + have signature `fn(tn, (tid1, tid2))`. + callback : callable, optional + A function to call after each full step of contraction and + compressions. It should have signature `fn(tn, tid)`. + preserve_tensor : bool, optional + If `True`, return a Tensor object even if it represents a scalar. + Ignore if `inplace=True`, in which case a TensorNetwork is always + returned. + progbar : bool, optional + Whether to show a progress bar. + inplace : bool, optional + Whether to perform the contraction inplace. + kwargs : dict, optional + Additional keyword passed to `_contract_compressed_tid_sequence`. + """ + # XXX: pick up max_bond, compress_late from ContractionTree + path = self.contraction_path(optimize, output_inds=output_inds) # generate the list of merges (tid1 -> tid2) @@ -7797,6 +7917,7 @@ def contract_compressed( canonize_after_opts=canonize_after_opts, gauge_boundary_only=gauge_boundary_only, compress_late=compress_late, + compress_mode=compress_mode, compress_min_size=compress_min_size, compress_opts=compress_opts, compress_span=compress_span, @@ -7810,7 +7931,9 @@ def contract_compressed( callback_pre_compress=callback_pre_compress, callback_post_compress=callback_post_compress, callback=callback, + preserve_tensor=preserve_tensor, progbar=progbar, + inplace=inplace, **kwargs, ) @@ -8461,6 +8584,9 @@ def contraction_path(self, optimize=None, **contract_opts): """ if optimize is None: optimize = get_contract_strategy() + + # XXX: short circuit to handle explicit path or tree + return self.contract( all, optimize=optimize, get="path", **contract_opts )