Source code for quimb.tensor.tensor_2d_tebd

from itertools import starmap

import numpy as np
import scipy.sparse.linalg as spla
from opt_einsum import shared_intermediates
from autoray import do, dag, conj, reshape

from ..utils import pairwise
from .drawing import get_colors
from .tensor_core import Tensor, contract_strategy
from .optimize import TNOptimizer
from .tensor_2d import (
    gen_2d_bonds,
    calc_plaquette_sizes,
    calc_plaquette_map,
    plaquette_to_sites,
    gen_long_range_path,
    gen_long_range_swap_path,
    swap_path_to_long_range_path,
    nearest_neighbors,
)
from .tensor_arbgeom_tebd import LocalHamGen, TEBDGen


[docs]class LocalHam2D(LocalHamGen): """A 2D Hamiltonian represented as local terms. This combines all two site and one site terms into a single interaction per lattice pair, and caches operations on the terms such as getting their exponential. Parameters ---------- Lx : int The number of rows. Ly : int The number of columns. H2 : array_like or dict[tuple[tuple[int]], array_like] The two site term(s). If a single array is given, assume to be the default interaction for all nearest neighbours. If a dict is supplied, the keys should represent specific pairs of coordinates like ``((ia, ja), (ib, jb))`` with the values the array representing the interaction for that pair. A default term for all remaining nearest neighbours interactions can still be supplied with the key ``None``. H1 : array_like or dict[tuple[int], array_like], optional The one site term(s). If a single array is given, assume to be the default onsite term for all terms. If a dict is supplied, the keys should represent specific coordinates like ``(i, j)`` with the values the array representing the local term for that site. A default term for all remaining sites can still be supplied with the key ``None``. Attributes ---------- terms : dict[tuple[tuple[int]], array_like] The total effective local term for each interaction (with single site terms appropriately absorbed). Each key is a pair of coordinates ``ija, ijb`` with ``ija < ijb``. """ def __init__(self, Lx, Ly, H2, H1=None): self.Lx = int(Lx) self.Ly = int(Ly) # parse two site terms if hasattr(H2, 'shape'): # use as default nearest neighbour term H2 = {None: H2} else: H2 = dict(H2) # possibly fill in default gates default_H2 = H2.pop(None, None) if default_H2 is not None: for coo_a, coo_b in gen_2d_bonds(Lx, Ly, steppers=[ lambda i, j: (i, j + 1), lambda i, j: (i + 1, j), ]): if (coo_a, coo_b) not in H2 and (coo_b, coo_a) not in H2: H2[coo_a, coo_b] = default_H2 super().__init__(H2=H2, H1=H1) @property def nsites(self): """The number of sites in the system. """ return self.Lx * self.Ly def __repr__(self): s = "<LocalHam2D(Lx={}, Ly={}, num_terms={})>" return s.format(self.Lx, self.Ly, len(self.terms))
[docs] def draw( self, ordering='sort', show_norm=True, figsize=None, fontsize=8, legend=True, ax=None, return_fig=False, **kwargs, ): """Plot this Hamiltonian as a network. Parameters ---------- ordering : {'sort', None, 'random'}, optional An ordering of the termns, or an argument to be supplied to :meth:`quimb.tensor.tensor_2d_tebd.LocalHam2D.get_auto_ordering` to generate this automatically. show_norm : bool, optional Show the norm of each term as edge labels. figsize : None or tuple[int], optional Size of the figure, defaults to size of Hamiltonian. fontsize : int, optional Font size for norm labels. legend : bool, optional Whether to show the legend of which terms are in which group. ax : None or matplotlib.Axes, optional Add to a existing set of axes. return_fig : bool, optional Whether to return any newly created figure. """ import matplotlib.pyplot as plt if figsize is None: figsize = (self.Ly, self.Lx) ax_supplied = (ax is not None) if not ax_supplied: fig, ax = plt.subplots(figsize=figsize, constrained_layout=True) ax.axis('off') ax.set_aspect('equal') if ordering is None or isinstance(ordering, str): ordering = self.get_auto_ordering(ordering, **kwargs) data = [] seen = set() n = 0 for ij1, ij2 in ordering: if (ij1 in seen) or (ij2 in seen): # start a new group seen = {ij1, ij2} n += 1 else: seen.add(ij1) seen.add(ij2) ys, xs = zip(ij1, ij2) d = ((xs[1] - xs[0])**2 + (ys[1] - ys[0])**2)**0.5 # offset by the length of bond to distinguish NNN etc. # choose offset direction by parity of first site if d > 2**0.5: xs = [xi + (-1)**int(ys[0]) * 0.02 * d for xi in xs] ys = [yi + (-1)**int(xs[0]) * 0.02 * d for yi in ys] # set coordinates for label with some offset towards left if ij1[1] < ij2[1]: lbl_x0 = (3 * xs[0] + 2 * xs[1]) / 5 lbl_y0 = (3 * ys[0] + 2 * ys[1]) / 5 else: lbl_x0 = (2 * xs[0] + 3 * xs[1]) / 5 lbl_y0 = (2 * ys[0] + 3 * ys[1]) / 5 nrm = do('linalg.norm', self.terms[ij1, ij2]) data.append((xs, ys, n, lbl_x0, lbl_y0, nrm)) num_groups = n + 1 colors = get_colors(range(num_groups)) # do the plotting for xs, ys, n, lbl_x0, lbl_y0, nrm in data: ax.plot(xs, ys, c=colors[n], linewidth=2 * nrm**0.5) if show_norm: label = "{:.3f}".format(nrm) ax.text(lbl_x0, lbl_y0, label, c=colors[n], fontsize=fontsize) # create legend if legend: handles = [] for color in colors.values(): handles += [plt.Line2D([0], [0], marker='o', color=color, linestyle='', markersize=10)] lbls = [f"Group {i + 1}" for i in range(num_groups)] ax.legend(handles, lbls, ncol=max(round(len(handles) / 20), 1), loc='center left', bbox_to_anchor=(1, 0.5)) if ax_supplied: return if return_fig: return fig plt.show()
graph = draw
[docs]class TEBD2D(TEBDGen): """Generic class for performing two dimensional time evolving block decimation, i.e. applying the exponential of a Hamiltonian using a product formula that involves applying local exponentiated gates only. Parameters ---------- psi0 : TensorNetwork2DVector The initial state. ham : LocalHam2D The Hamtiltonian consisting of local terms. tau : float, optional The default local exponent, if considered as time real values here imply imaginary time. max_bond : {'psi0', int, None}, optional The maximum bond dimension to keep when applying each gate. gate_opts : dict, optional Supplied to :meth:`quimb.tensor.tensor_2d.TensorNetwork2DVector.gate`, in addition to ``max_bond``. By default ``contract`` is set to 'reduce-split' and ``cutoff`` is set to ``0.0``. ordering : str, tuple[tuple[int]], callable, optional How to order the terms, if a string is given then use this as the strategy given to :meth:`~quimb.tensor.tensor_2d_tebd.LocalHam2D.get_auto_ordering`. An explicit list of coordinate pairs can also be given. The default is to greedily form an 'edge coloring' based on the sorted list of Hamiltonian pair coordinates. If a callable is supplied it will be used to generate the ordering before each sweep. second_order_reflect : bool, optional If ``True``, then apply each layer of gates in ``ordering`` forward with half the time step, then the same with reverse order. compute_energy_every : None or int, optional How often to compute and record the energy. If a positive integer 'n', the energy is computed *before* every nth sweep (i.e. including before the zeroth). compute_energy_final : bool, optional Whether to compute and record the energy at the end of the sweeps regardless of the value of ``compute_energy_every``. If you start sweeping again then this final energy is the same as the zeroth of the next set of sweeps and won't be recomputed. compute_energy_opts : dict, optional Supplied to :meth:`~quimb.tensor.tensor_2d.PEPS.compute_local_expectation`. By default ``max_bond`` is set to ``max(8, D**2)`` where ``D`` is the maximum bond to use for applying the gate, ``cutoff`` is set to ``0.0`` and ``normalized`` is set to ``True``. compute_energy_fn : callable, optional Supply your own function to compute the energy, it should take the ``TEBD2D`` object as its only argument. callback : callable, optional A custom callback to run after every sweep, it should take the ``TEBD2D`` object as its only argument. If it returns any value that boolean evaluates to ``True`` then terminal the evolution. progbar : boolean, optional Whether to show a live progress bar during the evolution. kwargs Extra options for the specific ``TEBD2D`` subclass. Attributes ---------- state : TensorNetwork2DVector The current state. ham : LocalHam2D The Hamiltonian being used to evolve. energy : float The current of the current state, this will trigger a computation if the energy at this iteration hasn't been computed yet. energies : list[float] The energies that have been computed, if any. its : list[int] The corresponding sequence of iteration numbers that energies have been computed at. taus : list[float] The corresponding sequence of time steps that energies have been computed at. best : dict If ``keep_best`` was set then the best recorded energy and the corresponding state that was computed - keys ``'energy'`` and ``'state'`` respectively. """ def __init__( self, psi0, ham, tau=0.01, D=None, chi=None, imag=True, gate_opts=None, ordering=None, second_order_reflect=False, compute_energy_every=None, compute_energy_final=True, compute_energy_opts=None, compute_energy_fn=None, compute_energy_per_site=False, callback=None, keep_best=False, progbar=True, ): super().__init__( psi0=psi0, ham=ham, tau=tau, D=D, imag=imag, gate_opts=gate_opts, ordering=ordering, second_order_reflect=second_order_reflect, compute_energy_every=compute_energy_every, compute_energy_final=compute_energy_final, compute_energy_opts=compute_energy_opts, compute_energy_fn=compute_energy_fn, compute_energy_per_site=compute_energy_per_site, callback=callback, keep_best=keep_best, progbar=progbar, ) # parse energy computation options if chi is None: chi = max(8, self.D**2) self.compute_energy_opts['max_bond'] = chi self.compute_energy_opts.setdefault('cutoff', 0.0) self.compute_energy_opts.setdefault('normalized', True)
[docs] def compute_energy(self): """Compute and return the energy of the current state. """ return self.state.compute_local_expectation( self.ham.terms, **self.compute_energy_opts )
@property def chi(self): return self.compute_energy_opts['max_bond'] @chi.setter def chi(self, value): self.compute_energy_opts['max_bond'] = round(value) def __repr__(self): s = "<{}(n={}, tau={}, D={}, chi={})>" return s.format( self.__class__.__name__, self.n, self.tau, self.D, self.chi)
[docs]def conditioner(tn, value=None, sweeps=2, balance_bonds=True): """ """ if balance_bonds: for _ in range(sweeps - 1): tn.balance_bonds_() tn.equalize_norms_() tn.balance_bonds_() tn.equalize_norms_(value=value)
[docs]class SimpleUpdate(TEBD2D): """A simple subclass of ``TEBD2D`` that overrides two key methods in order to keep 'diagonal gauges' living on the bonds of a PEPS. The gauges are stored separately from the main PEPS in the ``gauges`` attribute. Before and after a gate is applied they are absorbed and then extracted. When accessing the ``state`` attribute they are automatically inserted or you can call ``get_state(absorb_gauges=False)`` to lazily add them as hyperedge weights only. Reference: https://arxiv.org/abs/0806.3719. Parameters ---------- psi0 : TensorNetwork2DVector The initial state. ham : LocalHam2D The Hamtiltonian consisting of local terms. tau : float, optional The default local exponent, if considered as time real values here imply imaginary time. max_bond : {'psi0', int, None}, optional The maximum bond dimension to keep when applying each gate. gate_opts : dict, optional Supplied to :meth:`quimb.tensor.tensor_2d.TensorNetwork2DVector.gate`, in addition to ``max_bond``. By default ``contract`` is set to 'reduce-split' and ``cutoff`` is set to ``0.0``. ordering : str, tuple[tuple[int]], callable, optional How to order the terms, if a string is given then use this as the strategy given to :meth:`~quimb.tensor.tensor_2d_tebd.LocalHam2D.get_auto_ordering`. An explicit list of coordinate pairs can also be given. The default is to greedily form an 'edge coloring' based on the sorted list of Hamiltonian pair coordinates. If a callable is supplied it will be used to generate the ordering before each sweep. second_order_reflect : bool, optional If ``True``, then apply each layer of gates in ``ordering`` forward with half the time step, then the same with reverse order. compute_energy_every : None or int, optional How often to compute and record the energy. If a positive integer 'n', the energy is computed *before* every nth sweep (i.e. including before the zeroth). compute_energy_final : bool, optional Whether to compute and record the energy at the end of the sweeps regardless of the value of ``compute_energy_every``. If you start sweeping again then this final energy is the same as the zeroth of the next set of sweeps and won't be recomputed. compute_energy_opts : dict, optional Supplied to :meth:`~quimb.tensor.tensor_2d.PEPS.compute_local_expectation`. By default ``max_bond`` is set to ``max(8, D**2)`` where ``D`` is the maximum bond to use for applying the gate, ``cutoff`` is set to ``0.0`` and ``normalized`` is set to ``True``. compute_energy_fn : callable, optional Supply your own function to compute the energy, it should take the ``TEBD2D`` object as its only argument. callback : callable, optional A custom callback to run after every sweep, it should take the ``TEBD2D`` object as its only argument. If it returns any value that boolean evaluates to ``True`` then terminal the evolution. progbar : boolean, optional Whether to show a live progress bar during the evolution. gauge_renorm : bool, optional Whether to actively renormalize the singular value gauges. gauge_smudge : float, optional A small offset to use when applying the guage and its inverse to avoid numerical problems. condition_tensors : bool, optional Whether to actively equalize tensor norms for numerical stability. condition_balance_bonds : bool, optional If and when equalizing tensor norms, whether to also balance bonds as an additional conditioning. long_range_use_swaps : bool, optional If there are long range terms, whether to use swap gates to apply the terms. If ``False``, a long range blob tensor (which won't scale well for long distances) is formed instead. long_range_path_sequence : str or callable, optional If there are long range terms how to generate the path between the two coordinates. If callable, should take the two coordinates and return a sequence of coordinates that links them, else passed to ``gen_long_range_swap_path``. Attributes ---------- state : TensorNetwork2DVector The current state. ham : LocalHam2D The Hamiltonian being used to evolve. energy : float The current of the current state, this will trigger a computation if the energy at this iteration hasn't been computed yet. energies : list[float] The energies that have been computed, if any. its : list[int] The corresponding sequence of iteration numbers that energies have been computed at. taus : list[float] The corresponding sequence of time steps that energies have been computed at. best : dict If ``keep_best`` was set then the best recorded energy and the corresponding state that was computed - keys ``'energy'`` and ``'state'`` respectively. """ def __init__( self, psi0, ham, tau=0.01, D=None, chi=None, gauge_renorm=True, gauge_smudge=1e-6, condition_tensors=True, condition_balance_bonds=True, long_range_use_swaps=False, long_range_path_sequence='random', imag=True, gate_opts=None, ordering=None, second_order_reflect=False, compute_energy_every=None, compute_energy_final=True, compute_energy_opts=None, compute_energy_fn=None, compute_energy_per_site=False, callback=None, keep_best=False, progbar=True, ): super().__init__( psi0=psi0, ham=ham, tau=tau, D=D, chi=chi, imag=imag, gate_opts=gate_opts, ordering=ordering, second_order_reflect=second_order_reflect, compute_energy_every=compute_energy_every, compute_energy_final=compute_energy_final, compute_energy_opts=compute_energy_opts, compute_energy_fn=compute_energy_fn, compute_energy_per_site=compute_energy_per_site, callback=callback, keep_best=keep_best, progbar=progbar, ) self.gauge_renorm = gauge_renorm self.gauge_smudge = gauge_smudge self.condition_tensors = condition_tensors self.condition_balance_bonds = condition_balance_bonds self.gate_opts['long_range_use_swaps'] = long_range_use_swaps self.long_range_path_sequence = long_range_path_sequence def _initialize_gauges(self): """Create unit singular values, stored as tensors. """ # create the gauges like whatever data array is in the first site. data00 = next(iter(self._psi.tensor_map.values())).data self._gauges = dict() for ija, ijb in self._psi.gen_bond_coos(): bnd = self._psi.bond(ija, ijb) d = self._psi.ind_size(bnd) Tsval = Tensor( do('ones', (d,), dtype=data00.dtype, like=data00), inds=[bnd], tags=[ self._psi.site_tag(*ija), self._psi.site_tag(*ijb), 'SU_gauge', ] ) self._gauges[tuple(sorted((ija, ijb)))] = Tsval @property def gauges(self): """The dictionary of bond pair coordinates to Tensors describing the weights (``t = gauges[pair]; t.data``) and index (``t = gauges[pair]; t.inds[0]``) of all the gauges. """ return self._gauges @property def long_range_use_swaps(self): return self.gate_opts['long_range_use_swaps'] @long_range_use_swaps.setter def long_range_use_swaps(self, b): self.gate_opts['long_range_use_swaps'] = bool(b)
[docs] def gate(self, U, where): """Like ``TEBD2D.gate`` but absorb and extract the relevant gauges before and after each gate application. """ ija, ijb = where if callable(self.long_range_path_sequence): long_range_path_sequence = self.long_range_path_sequence(ija, ijb) else: long_range_path_sequence = self.long_range_path_sequence if self.long_range_use_swaps: path = tuple(gen_long_range_swap_path( ija, ijb, sequence=long_range_path_sequence)) string = swap_path_to_long_range_path(path, ija) else: # get the string linking the two sites string = path = tuple(gen_long_range_path( ija, ijb, sequence=long_range_path_sequence)) def env_neighbours(i, j): return tuple(filter( lambda coo: self._psi.valid_coo((coo)) and coo not in string, nearest_neighbors((i, j)) )) # get the relevant neighbours for string of sites neighbours = {site: env_neighbours(*site) for site in string} # absorb the 'outer' gauges from these neighbours for site in string: Tij = self._psi[site] for neighbour in neighbours[site]: Tsval = self.gauges[tuple(sorted((site, neighbour)))] Tij.multiply_index_diagonal_( ind=Tsval.inds[0], x=(Tsval.data + self.gauge_smudge)) # absorb the inner bond gauges equally into both sites along string for site_a, site_b in pairwise(string): Ta, Tb = self._psi[site_a], self._psi[site_b] Tsval = self.gauges[tuple(sorted((site_a, site_b)))] bnd, = Tsval.inds Ta.multiply_index_diagonal_(ind=bnd, x=Tsval.data**0.5) Tb.multiply_index_diagonal_(ind=bnd, x=Tsval.data**0.5) # perform the gate, retrieving new bond singular values info = dict() self._psi.gate_(U, where, absorb=None, info=info, long_range_path_sequence=path, **self.gate_opts) # set the new singualar values all along the chain for site_a, site_b in pairwise(string): bond_pair = tuple(sorted((site_a, site_b))) s = info['singular_values', bond_pair] if self.gauge_renorm: # keep the singular values from blowing up s = s / s[0] Tsval = self.gauges[bond_pair] Tsval.modify(data=s) # absorb the 'outer' gauges from these neighbours for site in string: Tij = self._psi[site] for neighbour in neighbours[site]: Tsval = self.gauges[tuple(sorted((site, neighbour)))] Tij.multiply_index_diagonal_( ind=Tsval.inds[0], x=(Tsval.data + self.gauge_smudge)**-1)
[docs] def get_state(self, absorb_gauges=True): """Return the state, with the diagonal bond gauges either absorbed equally into the tensors on either side of them (``absorb_gauges=True``, the default), or left lazily represented in the tensor network with hyperedges (``absorb_gauges=False``). """ psi = self._psi.copy() if not absorb_gauges: for Tsval in self.gauges.values(): psi &= Tsval else: for (ija, ijb), Tsval in self.gauges.items(): bnd, = Tsval.inds Ta = psi[ija] Tb = psi[ijb] Ta.multiply_index_diagonal_(bnd, Tsval.data**0.5) Tb.multiply_index_diagonal_(bnd, Tsval.data**0.5) if self.condition_tensors: conditioner(psi, balance_bonds=self.condition_balance_bonds) return psi
[docs] def set_state(self, psi): """Set the wavefunction state, this resets the environment gauges to unity. """ self._psi = psi.copy() self._initialize_gauges()
def gate_full_update_als( ket, env, bra, G, where, tags_plq, steps, tol, max_bond, optimize='auto-hq', solver='solve', dense=True, enforce_pos=False, pos_smudge=1e-6, init_simple_guess=True, condition_tensors=True, condition_maintain_norms=True, condition_balance_bonds=True, ): ket_plq = ket.select_any(tags_plq).view_like_(ket) bra_plq = bra.select_any(tags_plq).view_like_(bra) # this is the full target (copy - not virtual) target = ket_plq.gate(G, where, contract=False) | env if init_simple_guess: ket_plq.gate_(G, where, contract='reduce-split', max_bond=max_bond) for site in tags_plq: bra_plq[site].modify(data=conj(ket_plq[site].data)) if condition_tensors: conditioner(ket_plq, balance_bonds=condition_balance_bonds) for site in tags_plq: bra_plq[site].modify(data=conj(ket_plq[site].data)) if condition_maintain_norms: pre_norm = ket_plq[site].norm() overlap = bra_plq | target norm_plq = bra_plq | env | ket_plq xs = dict() x_previous = dict() previous_cost = None with contract_strategy(optimize), shared_intermediates(): for i in range(steps): for site in tags_plq: lix = norm_plq[site, 'BRA'].inds[:-1] rix = norm_plq[site, 'KET'].inds[:-1] # remove site tensors and group their indices if dense: N = (norm_plq.select(site, which='!any') .to_dense(lix, rix)) if enforce_pos: el, ev = do('linalg.eigh', (N + dag(N)) / 2) el = do('clip', el, pos_smudge, None) N = ev @ do('diag', el) @ dag(ev) else: N = (norm_plq.select(site, which='!any') .aslinearoperator(lix, rix)) # target vector (remove lower site tensor and contract to vec) b = (overlap .select((site, 'BRA'), which='!all') .to_dense(overlap[site, 'BRA'].inds[:-1], overlap[site, 'BRA'].inds[-1:])) if solver == 'solve': x = do('linalg.solve', N, b) elif solver == 'lstsq': x = do('linalg.lstsq', N, b, rcond=tol * 1e-3)[0] else: # use scipy sparse linalg solvers if solver in ('lsqr', 'lsmr'): solver_opts = dict(atol=tol, btol=tol) else: solver_opts = dict(tol=tol) # use current site as initial guess (iterate over site ind) x0 = x_previous.get(site, b) x = np.stack([ getattr(spla, solver) (N, b[..., k], x0=x0[..., k], **solver_opts)[0] for k in range(x0.shape[-1]) ], axis=-1) # update the tensors (all 'virtual' TNs above also updated) Tk, Tb = ket[site], bra[site] Tk.modify(data=reshape(x, Tk.shape)) Tb.modify(data=reshape(conj(x), Tb.shape)) # store solution to check convergence xs[site] = x # after updating both sites check for convergence of tensor entries cost_fid = do('trace', do('real', dag(x) @ b)) cost_norm = do('abs', do('trace', dag(x) @ (N @ x))) cost = - 2 * cost_fid + cost_norm converged = ( (previous_cost is not None) and (abs(cost - previous_cost) < tol) ) if converged: break previous_cost = cost for site in tags_plq: x_previous[site] = xs[site] if condition_tensors: if condition_maintain_norms: conditioner( ket_plq, value=pre_norm, balance_bonds=condition_balance_bonds) else: conditioner( ket_plq, balance_bonds=condition_balance_bonds) for site in tags_plq: bra_plq[site].modify(data=conj(ket_plq[site].data)) def gate_full_update_autodiff_fidelity( ket, env, bra, G, where, tags_plq, steps, tol, max_bond, optimize='auto-hq', autodiff_backend='autograd', autodiff_optimizer='L-BFGS-B', init_simple_guess=True, condition_tensors=True, condition_maintain_norms=True, condition_balance_bonds=True, **kwargs, ): ket_plq = ket.select_any(tags_plq).view_like_(ket) bra_plq = bra.select_any(tags_plq).view_like_(bra) # the target sites + gate and also norm (copy - not virtual) target = ket_plq.gate(G, where, contract=False) | env # make initial guess the simple gate tensors if init_simple_guess: ket_plq.gate_(G, where, contract='reduce-split', max_bond=max_bond) for site in tags_plq: bra_plq[site].modify(data=conj(ket_plq[site].data)) if condition_tensors: conditioner(ket_plq, balance_bonds=condition_balance_bonds) for site in tags_plq: bra_plq[site].modify(data=conj(ket_plq[site].data)) if condition_maintain_norms: pre_norm = ket_plq[site].norm() def fidelity(bra_plq): for site in tags_plq: ket_plq[site].modify(data=conj(bra_plq[site].data)) fid = (bra_plq | target).contract(all, optimize=optimize) norm = (bra_plq | env | ket_plq).contract(all, optimize=optimize) return - 2 * do('abs', fid) + do('abs', norm) tnopt = TNOptimizer( bra_plq, loss_fn=fidelity, tags=tags_plq, progbar=False, optimizer=autodiff_optimizer, autodiff_backend=autodiff_backend, **kwargs, ) bra_plq_opt = tnopt.optimize(steps, tol=tol) for site in tags_plq: new_data = bra_plq_opt[site].data ket[site].modify(data=conj(new_data)) bra[site].modify(data=new_data) if condition_tensors: if condition_maintain_norms: conditioner( ket_plq, value=pre_norm, balance_bonds=condition_balance_bonds) else: conditioner( ket_plq, balance_bonds=condition_balance_bonds) for site in tags_plq: bra_plq[site].modify(data=conj(ket_plq[site].data))
[docs]def get_default_full_update_fit_opts(): """The default options for the full update gate fitting procedure. """ return { # general 'tol': 1e-10, 'steps': 20, 'init_simple_guess': True, 'condition_tensors': True, 'condition_maintain_norms': True, # alternative least squares 'als_dense': True, 'als_solver': 'solve', 'als_enforce_pos': False, 'als_enforce_pos_smudge': 1e-6, # automatic differentation optimizing 'autodiff_backend': 'autograd', 'autodiff_optimizer': 'L-BFGS-B', }
[docs]def parse_specific_gate_opts(strategy, fit_opts): """Parse the options from ``fit_opts`` which are relevant for ``strategy``. """ gate_opts = { 'tol': fit_opts['tol'], 'steps': fit_opts['steps'], 'init_simple_guess': fit_opts['init_simple_guess'], 'condition_tensors': fit_opts['condition_tensors'], 'condition_maintain_norms': fit_opts['condition_maintain_norms'], } if 'als' in strategy: gate_opts['solver'] = fit_opts['als_solver'] gate_opts['dense'] = fit_opts['als_dense'] gate_opts['enforce_pos'] = fit_opts['als_enforce_pos'] gate_opts['pos_smudge'] = fit_opts['als_enforce_pos_smudge'] elif 'autodiff' in strategy: gate_opts['autodiff_backend'] = fit_opts['autodiff_backend'] gate_opts['autodiff_optimizer'] = fit_opts['autodiff_optimizer'] return gate_opts
[docs]class FullUpdate(TEBD2D): """Implements the 'Full Update' version of 2D imaginary time evolution, where each application of a gate is fitted to the current tensors using a boundary contracted environment. Parameters ---------- psi0 : TensorNetwork2DVector The initial state. ham : LocalHam2D The Hamtiltonian consisting of local terms. tau : float, optional The default local exponent, if considered as time real values here imply imaginary time. max_bond : {'psi0', int, None}, optional The maximum bond dimension to keep when applying each gate. gate_opts : dict, optional Supplied to :meth:`quimb.tensor.tensor_2d.TensorNetwork2DVector.gate`, in addition to ``max_bond``. By default ``contract`` is set to 'reduce-split' and ``cutoff`` is set to ``0.0``. ordering : str, tuple[tuple[int]], callable, optional How to order the terms, if a string is given then use this as the strategy given to :meth:`~quimb.tensor.tensor_2d_tebd.LocalHam2D.get_auto_ordering`. An explicit list of coordinate pairs can also be given. The default is to greedily form an 'edge coloring' based on the sorted list of Hamiltonian pair coordinates. If a callable is supplied it will be used to generate the ordering before each sweep. second_order_reflect : bool, optional If ``True``, then apply each layer of gates in ``ordering`` forward with half the time step, then the same with reverse order. compute_energy_every : None or int, optional How often to compute and record the energy. If a positive integer 'n', the energy is computed *before* every nth sweep (i.e. including before the zeroth). compute_energy_final : bool, optional Whether to compute and record the energy at the end of the sweeps regardless of the value of ``compute_energy_every``. If you start sweeping again then this final energy is the same as the zeroth of the next set of sweeps and won't be recomputed. compute_energy_opts : dict, optional Supplied to :meth:`~quimb.tensor.tensor_2d.PEPS.compute_local_expectation`. By default ``max_bond`` is set to ``max(8, D**2)`` where ``D`` is the maximum bond to use for applying the gate, ``cutoff`` is set to ``0.0`` and ``normalized`` is set to ``True``. compute_energy_fn : callable, optional Supply your own function to compute the energy, it should take the ``TEBD2D`` object as its only argument. callback : callable, optional A custom callback to run after every sweep, it should take the ``TEBD2D`` object as its only argument. If it returns any value that boolean evaluates to ``True`` then terminal the evolution. progbar : boolean, optional Whether to show a live progress bar during the evolution. fit_strategy : {'als', 'autodiff-fidelity'}, optional Core method used to fit the gate application. * ``'als'``: alternating least squares * ``'autodiff-fidelity'``: local fidelity using autodiff fit_opts : dict, optional Advanced options for the gate application fitting functions. Defaults are inserted and can be accessed via the ``.fit_opts`` attribute. compute_envs_every : {'term', 'group', 'sweep', int}, optional How often to recompute the environments used to the fit the gate application: * ``'term'``: every gate * ``'group'``: every set of commuting gates (the default) * ``'sweep'``: every total sweep * int: every ``x`` number of total sweeps pre_normalize : bool, optional Actively renormalize the state using the computed environments. condition_tensors : bool, optional Whether to actively equalize tensor norms for numerical stability. condition_balance_bonds : bool, optional If and when equalizing tensor norms, whether to also balance bonds as an additional conditioning. contract_optimize : str, optional Contraction path optimizer to use for gate + env + sites contractions. Attributes ---------- state : TensorNetwork2DVector The current state. ham : LocalHam2D The Hamiltonian being used to evolve. energy : float The current of the current state, this will trigger a computation if the energy at this iteration hasn't been computed yet. energies : list[float] The energies that have been computed, if any. its : list[int] The corresponding sequence of iteration numbers that energies have been computed at. taus : list[float] The corresponding sequence of time steps that energies have been computed at. best : dict If ``keep_best`` was set then the best recorded energy and the corresponding state that was computed - keys ``'energy'`` and ``'state'`` respectively. fit_opts : dict Detailed options for fitting the applied gate. """ def __init__( self, psi0, ham, tau=0.01, D=None, chi=None, fit_strategy='als', fit_opts=None, compute_envs_every=1, pre_normalize=True, condition_tensors=True, condition_balance_bonds=True, contract_optimize='auto-hq', imag=True, gate_opts=None, ordering=None, second_order_reflect=False, compute_energy_every=None, compute_energy_final=True, compute_energy_opts=None, compute_energy_fn=None, compute_energy_per_site=False, callback=None, keep_best=False, progbar=True, ): super().__init__( psi0=psi0, ham=ham, tau=tau, D=D, chi=chi, imag=imag, gate_opts=gate_opts, ordering=ordering, second_order_reflect=second_order_reflect, compute_energy_every=compute_energy_every, compute_energy_final=compute_energy_final, compute_energy_opts=compute_energy_opts, compute_energy_fn=compute_energy_fn, compute_energy_per_site=compute_energy_per_site, callback=callback, keep_best=keep_best, progbar=progbar, ) self.fit_strategy = str(fit_strategy) self.fit_opts = get_default_full_update_fit_opts() if fit_opts is not None: bad_opts = set(fit_opts) - set(self.fit_opts) if bad_opts: raise ValueError("Invalid fit option(s): {}".format(bad_opts)) self.fit_opts.update(fit_opts) self.pre_normalize = bool(pre_normalize) self.contract_optimize = str(contract_optimize) self.condition_tensors = bool(condition_tensors) self.condition_balance_bonds = bool(condition_balance_bonds) self.compute_envs_every = compute_envs_every self._env_n = self._env_term_count = self._env_group_count = -1 self._psi.add_tag('KET') @property def fit_strategy(self): return self._fit_strategy @fit_strategy.setter def fit_strategy(self, fit_strategy): self._gate_fit_fn = { 'als': gate_full_update_als, 'autodiff-fidelity': gate_full_update_autodiff_fidelity, }[fit_strategy] self._fit_strategy = fit_strategy
[docs] def set_state(self, psi): self._psi = psi.copy() # ensure the final dimension of each tensor is the physical dim for tag, ind in zip(self._psi.site_tags, self._psi.site_inds): t = self._psi[tag] if t.inds[-1] != ind: new_inds = [i for i in t.inds if i != ind] + [ind] t.transpose_(*new_inds)
@property def compute_envs_every(self): return self._compute_envs_every @compute_envs_every.setter def compute_envs_every(self, x): if x == 'sweep': self._need_to_recompute_envs = lambda: ( (self._n != self._env_n) ) elif x == 'group': self._need_to_recompute_envs = lambda: ( (self._n != self._env_n) or (self._group_count != self._env_group_count) ) elif x == 'term': self._need_to_recompute_envs = lambda: ( (self._n != self._env_n) or (self._group_count != self._env_group_count) or (self._term_count != self._env_term_count) ) else: x = max(1, int(x)) self._need_to_recompute_envs = lambda: (self._n >= self._env_n + x) self._compute_envs_every = x def _maybe_compute_plaquette_envs(self, force=False): """Compute and store the plaquette environments for all local terms. """ # first check if we need to compute the envs if not self._need_to_recompute_envs() and not force: return if self.condition_tensors: conditioner(self._psi, balance_bonds=self.condition_balance_bonds) # useful to store the bra that went into making the norm norm, _, self._bra = self._psi.make_norm(return_all=True) envs = dict() for x_bsz, y_bsz in calc_plaquette_sizes(self.ham.terms): envs.update(norm.compute_plaquette_environments( x_bsz=x_bsz, y_bsz=y_bsz, max_bond=self.chi, cutoff=0.0)) if self.pre_normalize: # get the first plaquette env and use it to compute current norm p0, env0 = next(iter(envs.items())) sites = plaquette_to_sites(p0) tags_plq = tuple(starmap(norm.site_tag, sites)) norm_plq = norm.select_any(tags_plq) | env0 # contract the local plaquette norm nfactor = do( 'abs', norm_plq.contract(all, optimize=self.contract_optimize)) # scale the bra and ket and each of the plaquette environments self._psi.multiply_(nfactor**(-1 / 2), spread_over='all') self._bra.multiply_(nfactor**(-1 / 2), spread_over='all') # scale the envs, taking into account the number of sites missing n = self._psi.num_tensors for ((_, _), (di, dj)), env in envs.items(): n_missing = di * dj env.multiply_(nfactor ** (n_missing / n - 1), spread_over='all') self.plaquette_envs = envs self.plaquette_mapping = calc_plaquette_map(envs) self._env_n = self._n self._env_group_count = self._group_count self._env_term_count = self._term_count
[docs] def presweep(self, i): """Full update presweep - compute envs and inject gate options. """ # inject the specific gate options required (do # here so user can change options between sweeps) self._gate_opts = parse_specific_gate_opts( self.fit_strategy, self.fit_opts) # keep track of number of gates applied, and commutative groups self._term_count = 0 self._group_count = 0 self._current_group = set()
[docs] def compute_energy(self): """Full update compute energy - use the (likely) already calculated plaquette environments. """ self._maybe_compute_plaquette_envs(force=self._n != self._env_n) return self.state.compute_local_expectation( self.ham.terms, plaquette_envs=self.plaquette_envs, plaquette_mapping=self.plaquette_mapping, **self.compute_energy_opts )
[docs] def gate(self, G, where): """Apply the gate ``G`` at sites where, using a fitting method that takes into account the current environment. """ # check if the new term commutes with those applied so far, this is to # decide if we need to recompute the environments swhere = set(where) if self._current_group.isdisjoint(swhere): # if so add it to the grouping self._current_group |= swhere else: # else increment and reset the grouping self._current_group = swhere self._group_count += 1 # get the plaquette containing ``where`` and the sites it contains - # these will all be fitted self._maybe_compute_plaquette_envs() plq = self.plaquette_mapping[tuple(sorted(where))] env = self.plaquette_envs[plq] tags_plq = tuple(starmap(self._psi.site_tag, plaquette_to_sites(plq))) # perform the gate, inplace self._gate_fit_fn( ket=self._psi, env=env, bra=self._bra, G=G, where=where, tags_plq=tags_plq, max_bond=self.D, optimize=self.contract_optimize, condition_balance_bonds=self.condition_balance_bonds, **self._gate_opts ) # increments every gate call regardless self._term_count += 1