# This code is part of qtealeaves.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.
"""
Dense Matrix Product Operators representing Hamiltonians or observables.
| | | |
-O-O-O-O-
| | | |
There are 3 classes in this module: `MPOSite`, `DenseMPO`, `DenseMPOList`
DenseMPO : is the standard textbook-like MPO that acts on specified sites
(DenseMPO.sites). It is a list of MPOSite-s
MPOSite : each tensor in the DenseMPO is stored as MPOSite. For example,
a 4-body DenseMPO is going to be a list of 4 MPOSites.
DenseMPOList : a list of DenseMPO-s. Here, usually, it's used to represent
Hamiltonians which consist of multiple terms. In this case,
each Hamiltonian term is represented as a DenseMPO. For example,
Ising model contains nearest-neighbour sigma_x sigma_x and
sigma_z terms, and every 2-body term sigma_x sigma_x and
local term sigma_z would be stored as a separate DenseMPO
added to a DenseMPOList
Some useful attributes/functions:
MPOSite : initialize with MPOSite([list of sites], tensor_str_name, pstrength=None,
prefactor=1.0, operators=op_dict, params={})
- self.site = on which site in a system does this MPOSite act
- self.operator = get an actual tensor
DenseMPO : initialize with DenseMPO([list of MPOSite-s], tensor_backend)
- self.num_sites = on how many sites does this MPO act
- self.sites = list of sites on which MPO acts
- self[ind].operator = get tensor on site ind
- self.append(MPOSite) = add MPOSite to Dense MPO
DenseMPOList : initialize with DenseMPOList() and then append DenseMPO-s
- self[ind] = get ind-th DenseMPO in a list
- DenseMPOList.from_model_prepared(...) = creates a DenseMPOList
- object from a given Hamiltonian model
"""
# pylint: disable=too-many-arguments
# pylint: disable=too-many-lines
# pylint: disable=too-many-locals
# pylint: disable=too-many-public-methods
import logging
import warnings
from copy import deepcopy
from typing import Sequence
import numpy as np
from qtealeaves.abstracttns.abstract_matrix_tn import _AbstractMatrixTN
from qtealeaves.convergence_parameters import TNConvergenceParameters
from qtealeaves.mpos.abstracteffop import _AbstractEffectiveMpo
from qtealeaves.operators import TNOperators
from qtealeaves.tensors import TensorBackend
from qtealeaves.tensors.abstracttensor import _AbstractQteaTensor
from qtealeaves.tooling import QTeaLeavesError
from qtealeaves.tooling.operatorstrings import _op_string_mul
from qtealeaves.tooling.parameterized import _ParameterizedClass
from qtealeaves.tooling.restrictedclasses import _RestrictedList
__all__ = ["MPOSite", "DenseMPO", "DenseMPOList"]
logger = logging.getLogger(__name__)
[docs]
class MPOSite(_ParameterizedClass):
"""
One site in a dense MPO term. For example, a 4-body DenseMPO is
going to be a list of 4 MPOSites.
Initialize with MPOSite([list of sites], tensor_str_name, pstrength=None,
prefactor=1.0, operators=op_dict, params={})
**Arguments**
site : integer
Site index.
str_op : str
Key for the operator.
pstrength : pstrength, callable, numeric
Containing the parameterization of the term.
weight : scalar
Scalar constant prefactor.
operators : :class:`TNOperators` or None
If present, operators will be directly extracted.
params : dict or None
If present, parameterization will be directly extracted.
"""
def __init__(self, site, str_op, pstrength, weight, operators=None, params=None):
self.site = site
self.str_op = str_op
self.pstrength = pstrength
self.operator = None if operators is None else operators[(site, str_op)].copy()
self.strength = (
None if params is None else self.eval_numeric_param(self.pstrength, params)
)
if self.pstrength is None:
self.strength = 1.0
self.weight = weight
# --------------------------------------------------------------------------
# Overwritten magic methods
# --------------------------------------------------------------------------
def __repr__(self):
"""
User-friendly representation of object for print().
"""
str_repr = (
f"{self.__class__.__name__} on site {self.site} with ("
f"\noperator={self.operator},"
f"\npstrength={self.pstrength},"
f" weight={self.weight} )"
)
return str_repr
@property
def total_scaling(self):
"""Returns the scaling combining params and weight."""
return self.strength * self.weight
@property
def shape(self):
"""Shape attribute equivalent to tensor's shape: bond dimension of links."""
return self.operator.shape
[docs]
def initialize(self, operators, params):
"""Resolve operators and parameterization for the given input."""
self.set_op(operators)
self.set_param(params)
[docs]
def set_op(self, operators):
"""Resolve operators for the given input."""
if self.str_op is None:
raise QTeaLeavesError("Operator string no longer available.")
self.operator = operators[(self.site, self.str_op)]
[docs]
def set_param(self, params):
"""Resolve parameterization for the given input."""
if self.pstrength is None:
self.strength = 1.0
return
strength = self.eval_numeric_param(self.pstrength, params)
if hasattr(strength, "__len__"):
raise QTeaLeavesError("Strength cannot be a list.")
if strength == 0.0:
warnings.warn("Adding term with zero-coupling.")
self.strength = strength
[docs]
def copy_with_new_op(self, operator):
"""
Create a copy of self, but with replacing the operator with the one passed.
Corresponding string identifier will be set to `None`.
"""
obj = deepcopy(self)
obj.operator = operator
obj.str_op = None
return obj
[docs]
def copy(self):
"""Create a copy of self."""
return deepcopy(self)
[docs]
def convert(self, dtype, device, stream=None):
"""Convert underlying array to the specified data type inplace."""
self.operator.convert(dtype, device, stream=stream)
[docs]
@classmethod
def from_operator(cls, site, operator):
"""
Create MPOSite object from a given operator tensor.
Arguments
---------
site : int
Which site in a system.
operator : :class:`_AbstractQteaBaseTensor`'
Rank-4 tensor which will be the operator for this MPOSite.
Return
------
mpo_site : :class:`MPOSite`'
The resulting MPO site.
"""
if len(operator.shape) != 4:
raise ValueError(
"MPOSite can be built only from rank-4 operators and "
f"not rank-{len(operator.shape)}."
)
str_op = f"MPO_{site}"
mpo_site = cls(
site, str_op, pstrength=None, weight=1, operators=None, params=None
)
# set the operator tensor
mpo_site.operator = operator
return mpo_site
# pylint: disable-next=too-many-instance-attributes
[docs]
class DenseMPO(_AbstractMatrixTN, _RestrictedList, _AbstractEffectiveMpo):
"""
Dense MPO is the standard textbook-like MPO that acts on specified sites
(DenseMPO.sites). It is a list of :class:`MPOSite`'s.
Initialize with DenseMPO([list of MPOSite-s], tensor_backend)
"""
class_allowed = MPOSite
def __init__(
self,
sites: Sequence[MPOSite] | None = None,
convergence_parameters: TNConvergenceParameters | None = None,
is_oqs: bool = False,
tensor_backend: TensorBackend | None = None,
require_singvals: bool = False,
local_dim: int = 2,
):
if tensor_backend is None:
# prevents dangerous default value similar to dict/list
tensor_backend = TensorBackend()
if convergence_parameters is None:
# Not really allowed in the next step
convergence_parameters = TNConvergenceParameters()
if sites is None:
sites = []
_RestrictedList.__init__(self, sites)
_AbstractMatrixTN.__init__(
self,
num_sites=1, # is required, otherwise some local dim fails**
convergence_parameters=convergence_parameters,
local_dim=local_dim,
requires_singvals=require_singvals,
tensor_backend=tensor_backend,
)
# **: Should be fine, although _num_sites is something else
# than the property num_sites defined in this class. Similar
# problem in _AbstractMatrixTN, where _tensors is empty, but
# singular values set according to num_sites. But it looks
# like everything is overwritten consistently.
self.is_oqs = is_oqs
self.eff_ops = {}
self._tensor_network = None
self._phi_for_fit = None
# Need this internal flag for returning AbstractQteaTensor for move_pos.
# Otherwise, we have to copy the whole move_pos logic here.
self._getitem_return_qteatensor = False
# --------------------------------------------------------------------------
# Overwritten magic methods
# --------------------------------------------------------------------------
def __iter__(self):
"""
Iteration needs to be inherited from RestrictedList explicitly.
"""
yield from super(_RestrictedList, self).__iter__()
def __repr__(self):
"""
User-friendly representation of object for print().
"""
return f"{self.__class__.__name__}(sites={self.sites})"
@property
def num_sites(self):
"""Length of the Dense MPO"""
return len(self)
@property
def sites(self):
"""Generate list of site indices."""
sites = [elem.site for elem in self]
return sites
@property
def phi_for_fit(self):
"""Get the TN for variational fitting."""
return self._phi_for_fit
@phi_for_fit.setter
def phi_for_fit(self, value):
"""
Set the TN for variational fitting.
"""
self._phi_for_fit = value
[docs]
def get_tensor_of_site(self, idx):
"""
Return the tensor representing the MPO operator at site `idx`
"""
return self[idx].operator
def _iter_tensors(self):
"""Iterate over all tensors forming the tensor network (for convert etc)."""
for ii in range(self.num_sites):
yield self.get_tensor_of_site(ii)
def __len__(self):
"""
Length needs to be inherited from RestrictedList explicitly.
"""
return super(_RestrictedList, self).__len__()
def __getitem__(self, idxs):
"""Get an entry from the DenseMPO.
Note:
The function resolves tuples via the effective operators
while everything else is resolved via the DenseMPO acting
as a list.
"""
if isinstance(idxs, tuple):
return self.eff_ops[idxs]
value = super(_RestrictedList, self).__getitem__(idxs)
if self._getitem_return_qteatensor and isinstance(value, MPOSite):
# Fix for `move_pos` and CPU / GPU / mixed-device simulations
# where __getitem__ should return _AbstractQteaTensor. It only
# acts on single tensors, so no need to check is value is list
# of `MPOSite`s
return value.operator
return value
def __setitem__(self, index, elem):
"""
New setitem with the possibility of just setting the tensor
"""
if isinstance(index, tuple):
self.eff_ops[index] = elem
else:
if isinstance(elem, _AbstractQteaTensor):
site = self[index]
site.operator = elem
site.str_op = f"MPO_{index}"
else:
site = elem
super(_RestrictedList, self).__setitem__(index, self._check_class(site))
[docs]
def append(self, elem):
"""Overwriting append to extend as well the list of singvals."""
super().append(elem)
# Copy last singvals - assumption: cannot change with local operators
self._singvals.append(self._singvals[-1])
[docs]
def compress_links(self, idx_start, idx_end, trunc=False, conv_params=None):
"""
Compresses links between sites in a dense MPO by performing a QR or SVD,
optionally performs the additional truncation along the way.
Parameters
----------
idx_start : int
MPO site from which to start the compression.
idx_end : int
MPO site on which to end the compression.
trunc : Boolean, optional
If True, the truncation will be done according to the `conv_params`.
Default to False.
conv_params : :py:class:`TNConvergenceParameters`, optional
Convergence parameters for the truncation. Must be specified if
`trunc` is set to True.
Default to `None`.
"""
# pylint: disable=access-member-before-definition
if len(self) > len(self._singvals):
# Note: the linter does not recognize the call to
# super, which initializes _singlvals in the _AbstractMatrixTN
# pylint: disable=attribute-defined-outside-init, access-member-before-definition
self._singvals = [None] * (len(self) * 2)
self.iso_towards(idx_start, True)
self.iso_towards(idx_end, True, trunc, conv_params)
[docs]
def add_identity_on_site(self, idx, link_vertical):
"""
Add identity with the correct links to neighboring terms on site `idx`.
Parameters
----------
idx : int
Site to which add the identity. Goes from 0 to num sites in a system.
link_vertical : link as returned by corresponding QteaTensor
Needed to build the local Hilbert space (in case it is different across
the system).
"""
if len(self) == 0:
raise QTeaLeavesError(
"Cannot use `add_identity_on_site` on empty DenseMPO."
)
sites = np.array(self.sites)
if np.any(sites[1:] - sites[:-1] <= 0):
raise QTeaLeavesError(
"Cannot use `add_identity_on_site` on unsorted DenseMPO."
)
if idx in self.sites:
raise QTeaLeavesError("Site is already in DenseMPO.")
sites = np.array(self.sites + [idx])
# Figure out the index where to insert
# the identity, and the index of the site
# on the left of it.
# This takes care of the periodic boundary
# conditions on the indices of the dense_mpo.
sites_sorted = sorted(sites)
insert_site_index = sites_sorted.index(idx)
left_index = insert_site_index - 1
op = self[left_index].operator
if op is None:
raise QTeaLeavesError(
"Cannot use `add_identity_on_site` on uninitialized DenseMPO."
)
eye_horizontal = op.eye_like(op.links[3])
eye_vertical = op.eye_like(link_vertical)
# Contract together
eye_horizontal.attach_dummy_link(0, False)
eye_vertical.attach_dummy_link(0, True)
eye = eye_horizontal.tensordot(eye_vertical, ([0], [0]))
eye.transpose_update([0, 2, 3, 1])
key = str(id(eye))
op_dict = TNOperators()
op_dict[key] = eye
# add it to the correct site
site = MPOSite(idx, key, None, 1.0, operators=op_dict, params={})
self.insert(insert_site_index, site)
# add the singular values, identity cannot change them, so insert
# the same ones
self._singvals.insert(insert_site_index, self._singvals[insert_site_index])
[docs]
def initialize(self, operators, params):
"""Resolve operators and parameterization for the given input for each site."""
for elem in self:
elem.initialize(operators, params)
[docs]
def move_pos(self, pos, device=None, stream=False):
"""
Move just the tensor in position `pos` with the effective
operators insisting on links of `pos` on another device.
Other objects like effective projectors will be moved as
well. Acts in place.
Note: the implementation of the tensor backend can fallback
to synchronous moves only depending on its implementation.
Parameters
----------
pos : int | Tuple[int]
Integers identifying a tensor in a tensor network.
device : str, optional
Device where you want to send the QteaTensor. If None, no
conversion. Default to None.
stream : bool | stream | None, optional
If True, use a new stream for memory communication given
by the data mover. If False, use synchronous communication
with GPU. If not a boolean (and not None), we assume it
is a stream object compatible with the backend.
None is deprecated and equals to False.
Default to False (Use null stream).
"""
# Fix `move_pos` for CPU / GPU / mixed-device simulations
# where __getitem__ should return _AbstractQteaTensor.
self._getitem_return_qteatensor = True
super().move_pos(pos, device=device, stream=stream)
self._getitem_return_qteatensor = False
[docs]
def sort_sites(self):
"""Sort sites while and install matching link for symmetries."""
sites = [elem.site for elem in self]
# Potentially no fast return possible here, because even if the sites
# are sorted, we need to manage the links for symmetric tensor
# networks
if not self[0].operator.has_symmetry:
if all(sites[ii] <= sites[ii + 1] for ii in range(len(sites) - 1)):
# sites already sorted, return
return self
inds = np.argsort(sites)
dims_l = [elem.operator.shape[0] for elem in self]
dims_r = [elem.operator.shape[3] for elem in self]
max_l = np.max(dims_l)
max_r = np.max(dims_r)
max_chi = max(max_l, max_r)
if max_chi == 1:
return self._sort_sites_chi_one(inds)
raise QTeaLeavesError("For now, we only sort product terms.")
[docs]
def pad_identities(self, num_sites, eye_ops):
"""Pad identities on sites which are not in MPO yet respecting the symmetry."""
sites = np.array([elem.site for elem in self])
if np.any(sites[1:] - sites[:-1] < 1):
sorted_mpo = self.sort_sites()
return sorted_mpo.pad_identities(num_sites, eye_ops)
raise QTeaLeavesError("Not implemtented yet.")
[docs]
def trace(self, local_dim):
"""
Compute the trace of and MPO written in DenseMPO form.
Parameters
----------
local_dim : list[int]
Physical dimension of each of the sites in self.
Returns
-------
operator_trace : float | complex
Corresponds to the trace of the operator.
"""
partial_trace = 1
dims = local_dim.copy()
for mpo_site in self:
# for a dense_mpo the trace is the product of the traces of the
# operators at each site.
if mpo_site.operator.shape[0] != 1 or mpo_site.operator.shape[-1] != 1:
raise NotImplementedError(
"Trace not yet impemented if DenseMPO list contains MPOsite "
"with non-dummy horizontal links."
)
operator = mpo_site.operator.copy()
operator.remove_dummy_link(3)
operator.remove_dummy_link(0)
partial_trace *= operator.trace() * mpo_site.total_scaling
# Cancel contribution from dims
dims[mpo_site.site] = 1
partial_trace *= float(np.prod(np.array(dims)))
return partial_trace
def _sort_sites_chi_one(self, inds):
"""Sorting sites in the case of bond dimension equal to one."""
new_mpo = DenseMPO(is_oqs=self.is_oqs, tensor_backend=self._tensor_backend)
link = self[0].operator.dummy_link(self[0].operator.links[0])
for ii in inds:
# Trivial tensor porting the sector
one = self._tensor_backend(
[link, link],
ctrl="O",
are_links_outgoing=[False, True],
device=self[ii].operator.device,
dtype=self[ii].operator.dtype,
)
one.attach_dummy_link(2, True)
op = self[ii].operator.copy()
op.attach_dummy_link(4, is_outgoing=False)
tens = one.tensordot(op, ([2], [4]))
# We have six links [left, right-1, right-former-left, bra, ket, right-2]
tens.transpose_update([0, 3, 4, 1, 2, 5])
tens.fuse_links_update(3, 5)
mpo_site = self[ii].copy_with_new_op(tens)
mpo_site.str_op = self[ii].str_op
new_mpo.append(mpo_site)
# For the next loop iteration
link = new_mpo[-1].operator.links[-1]
# Check that MPO does conserve symmetry
if not self.is_oqs:
new_mpo[-1].operator.assert_identical_irrep(3)
return new_mpo
[docs]
@classmethod
def from_matrix(
cls,
matrix,
sites,
dim,
conv_params,
tensor_backend=None,
operators=None,
pad_with_identities=False,
):
"""
For a given matrix returns dense MPO form decomposing with SVDs
Parameters
----------
matrix : QteaTensor | ndarray
Matrix to write in (MPO) format
sites : List[int]
Sites to which the MPO is applied
dim : int
Local Hilbert space dimension
conv_params : :py:class:`TNConvergenceParameters`
Convergence parameters. The relevant attribute is the `max_bond_dimension`
tensor_backend : instance of :class:`TensorBackend`
Default for `None` is :class:`QteaTensor` with np.complex128 on CPU.
operators: TNOperators, optional
Operator class. Default for `None` is empty `TNOperator` instance.
pad_with_identities: bool, optional
If True, pad with identities the sites between min(sites) and max(sites)
that have no operator. Default to False.
Return
------
DenseMPO
The MPO decomposition of the matrix
"""
if tensor_backend is None:
# prevents dangerous default value similar to dict/list
tensor_backend = TensorBackend()
if operators is None:
# prevents dangerous default value similar to dict/list
operators = TNOperators()
if not isinstance(matrix, tensor_backend.tensor_cls):
matrix = tensor_backend.tensor_cls.from_elem_array(matrix)
mpo = cls(
convergence_parameters=conv_params,
tensor_backend=tensor_backend,
local_dim=dim,
)
bond_dim = 1
names = []
work = matrix
op_set_name = operators.set_names[0]
if len(operators.set_names) > 1:
raise QTeaLeavesError(
"Can use matrix to MPO decomposition only for one set of operators."
)
site_cnt = 0
for ii in range(sites[0], sites[-1], 1):
if ii in sites:
# dim dim**(n_sites-1)
# | ||
# O --[unfuse]--> O --[fuse upper and lower legs]-->
# | ||
#
# ==O== --[SVD, truncating]--> ==O-o-O==
#
# | |
# --[unfuse]--> O-O ---iterate
# | |
# dim dim**(n_sites-1)
work = np.reshape(
work,
(
bond_dim,
dim,
dim ** (len(sites) - 1 - site_cnt),
dim,
dim ** (len(sites) - 1 - site_cnt),
),
)
tens_left, work, _, _ = work.split_svd(
[0, 1, 3], [2, 4], contract_singvals="R", conv_params=conv_params
)
bond_dim = deepcopy(work.shape[0])
operators[(op_set_name, f"mpo{ii}")] = tens_left
names.append(f"mpo{ii}")
site_cnt += 1
elif pad_with_identities:
operators[(op_set_name, f"id{ii}")] = DenseMPO.generate_mpo_identity(
bond_dim, dim, bond_dim, tensor_backend
)
names.append((op_set_name, f"id{ii}"))
work = work.reshape((work.shape[0], dim, dim, 1))
operators[(op_set_name, f"mpo{sites[-1]}")] = work
names.append(f"mpo{sites[-1]}")
# Note: the linter does not recognize the call to
# super, which initializes _local_dim in the _AbstractMatrixTN
# pylint: disable=attribute-defined-outside-init
mpo._local_dim = np.zeros(len(sites), dtype=int)
cnt = 0
for site, name in zip(sites, names):
mpo.append(MPOSite(site, name, 1, 1, operators=operators))
mpo._local_dim[cnt] = mpo[cnt].operator.shape[1]
mpo._singvals.append(None)
cnt += 1
# pylint: disable=attribute-defined-outside-init
mpo._iso_center = (cnt - 1, cnt)
return mpo
[docs]
@classmethod
def from_tensor_list(
cls,
tensor_list,
conv_params=None,
iso_center=None,
tensor_backend=None,
operators=None,
sites=None,
check_equivalence=True,
):
"""
Initialize the dense MPO from a list of tensors.
Parameters
----------
tensor_list : List[QteaTensor] | List[MPOSite] | list[np.ndarray]
Matrix to write in (MPO) format. `np.ndarray` only allowed
for systems without symmetry.
conv_params : :py:class:`TNConvergenceParameters`, None
Input for handling convergence parameters. Default to None
iso_center : None, int, List[int], str, optional
If None, the center is None.
If str, the iso center is installed
If int, the iso center is that integer.
Default is None
tensor_backend : instance of :class:`TensorBackend`
Default for `None` is :class:`QteaTensor` with np.complex128 on CPU.
operators: TNOperators, optional
Operator class. Default for `None` is empty `TNOperator` instance.
sites : List[int], None
Sites to which the MPO is applied. If None, they are assumed to be
[0, 1, ..., len(tensorlist)-1]. Default to None
check_equivalence : bool, optional
If False, skip equivalence scan for operators and only avoid key-name collisions.
Defaults to True.
Return
------
DenseMPO
The MPO decomposition of the matrix
"""
if tensor_backend is None:
# prevents dangerous default value similar to dict/list
tensor_backend = TensorBackend()
if operators is None:
# prevents dangerous default value similar to dict/list
operators = TNOperators()
if sites is None:
sites = list(range(len(tensor_list)))
mpo = cls(convergence_parameters=conv_params, tensor_backend=tensor_backend)
# Convert to list if _AbstractQteaTensor if not in this format yet
if isinstance(tensor_list[0], MPOSite):
# Convert because they are MPOSites
tensor_list = [ss.operator * ss.weight for ss in tensor_list]
elif isinstance(tensor_list[0], np.ndarray):
# Convert because they are numpy.ndarray
if tensor_backend.base_tensor_cls != tensor_backend.tensor_cls:
# Strong indication that this is a system with symmetric tensors
raise TypeError(
"Numpy ndarray only allowed if tensor_cls equals tensor_base_cls."
)
tensor_list = [tensor_backend.from_elem_array(elem) for elem in tensor_list]
names = []
for ii, tens in enumerate(tensor_list):
names.append(
_duplicate_check_and_set(
operators, f"mpo{ii}", tens, check_equivalence=check_equivalence
)
)
# pylint: disable=attribute-defined-outside-init
mpo._local_dim = np.zeros(len(tensor_list), dtype=int)
cnt = 0
for site, name in zip(sites, names):
mpo.append(
MPOSite(site, name, pstrength=None, weight=1, operators=operators)
)
mpo._local_dim[cnt] = mpo[cnt].operator.shape[1]
mpo._singvals.append(None)
cnt += 1
if isinstance(iso_center, str):
mpo.install_gauge_center()
elif isinstance(iso_center, int):
# Note: the linter does not recognize the call to
# super, which initializes iso_center in the _AbstractMatrixTN
# pylint: disable=attribute-defined-outside-init
mpo._iso_center = (iso_center, iso_center + 2)
elif isinstance(iso_center, (list, tuple)):
# pylint: disable=attribute-defined-outside-init
mpo.iso_center = iso_center
return mpo
[docs]
@staticmethod
def generate_mpo_identity(left_bd, local_dim, right_bd, tensor_backend):
"""
Generate an identity in MPO form with given dimensions.
"""
id_tens = tensor_backend([left_bd, local_dim, local_dim, right_bd])
for ii in range(min(left_bd, right_bd)):
id_tens[ii : ii + 1, :local_dim, :local_dim, ii : ii + 1] = (
id_tens.eye_like(local_dim)
)
return id_tens
############################################################
# Methods for effective operators
############################################################
[docs]
def setup_as_eff_ops(self, tensor_network, measurement_mode=False) -> None:
"""Set this dense MPO as effective ops in TN and initialize."""
self._tensor_network = tensor_network
self.eff_ops[(None, 0)] = self.tensor_backend(ctrl="O", links=[1, 1, 1, 1])
self.eff_ops[(None, (0, 0))] = self.tensor_backend(ctrl="O", links=[1, 1, 1, 1])
self.eff_ops[(None, len(self) - 1)] = self.tensor_backend(
ctrl="O", links=[1, 1, 1, 1]
)
# pylint: disable-next=protected-access
for ii, key in enumerate(tensor_network._iter_physical_links()):
self[key] = self[ii].operator.copy()
self[key].convert(
dtype=tensor_network.dtype,
device=tensor_network.tensor_backend.memory_device,
)
# Last chance to change iso center
if tensor_network.iso_center is None:
tensor_network.isometrize_all()
if tensor_network.iso_center != tensor_network.default_iso_pos:
tensor_network.iso_towards(tensor_network.default_iso_pos)
tensor_network.eff_op = self
tensor_network.build_effective_operators(measurement_mode=measurement_mode)
[docs]
def contr_to_eff_op(
self,
tensor: _AbstractQteaTensor,
pos: int,
pos_links: Sequence[int],
idx_out: int,
) -> None:
"""
Contract operator lists with tensors T and Tdagger to effective operator.
**Arguments**
tensor : :class:`_AbstractQteaTensor`
Tensor to be contracted to effective operator.
pos : int, tuple (depending on TN)
Position of tensor.
pos_links : list of int, tuple (depending on TN)
Position of neighboring tensors where the links
in `tensor` lead to.
idx_out : int
Uncontracted link to be used for effective operator.
"""
ops_list, idx_list, key, ikey = self._helper_contract_to_eff_op(
pos, pos_links, idx_out, keep_none=True
)
phi = tensor if self._phi_for_fit is None else self._phi_for_fit[pos]
## reshuffle to avoid "holes" in the contraction chain
if 0 < idx_out < len(phi.shape) - 1:
idx_list = idx_list[idx_out:] + idx_list[:idx_out]
ops_list = ops_list[idx_out:] + ops_list[:idx_out]
ctens = phi.tensordot(ops_list[0], [(idx_list[0],), (2,)])
## ctens.shape = (rem_t,O1_0,O1_1,O1_3) where rem_t are the remaining indices of tensors
new_idx_list = [
ii - 1 if ii > idx_list[0] else ii for ii in idx_list
] # update legs label
for ii in range(1, len(ops_list)):
op = ops_list[ii]
cidx = new_idx_list[ii]
ctens = ctens.tensordot(op, [(cidx, -1), (2, 0)])
new_idx_list = [idx - 1 if idx > cidx else idx for idx in new_idx_list]
## At the end ctens.shape = (idx_out,O1_0,O2_1,...,Ok_1,Ok_3)
cidx = list(range(2, len(ctens.shape) - 1))
ctens = ctens.tensordot(tensor.conj(), [cidx, idx_list])
## ctens.shape = (idx_out,O1_0,Ok_3,idx_out*)
ctens.transpose_update((1, 3, 0, 2))
self[key] = ctens
if ikey in self.eff_ops:
del self.eff_ops[ikey]
[docs]
def contract_tensor_lists(
self,
tensor: _AbstractQteaTensor,
pos: int,
pos_links: Sequence[int],
custom_ops=None,
pre_return_hook=None,
) -> _AbstractQteaTensor:
"""
Linear operator to contract all the effective operators
around the tensor in position `pos`. Used as a matrix-vector multiplication
function in solvers.
**Arguments**
-------------
tensor : :class:`_AbstractQteaTensor`
Tensor to be contracted to effective operator.
pos : int, tuple (depending on TN)
Position of tensor.
pos_links : list of int, tuple (depending on TN)
Position of neighboring tensors where the links
in `tensor` lead to.
custom_ops : `None` or list of :class:`ITPOTerm`
Ordered list of iTPO terms for tensor, which
should be used instead of information in `pos`
and `pos_links`.
pre_return_hook : `None` - not used for DenseMPO
**Return**
----------
ctens : :class:`_AbstractQteaTensor`
The tensor contracted with effective operators.
"""
if custom_ops is None:
ops_list = []
idx_list = []
for ii, pos_link in enumerate(pos_links):
pos_jj = self.eff_ops[(pos_link, pos)]
ops_list.append(pos_jj)
idx_list.append(ii)
else:
# Required for time evolution backwards step on R-tensor
# and two-tensor update; None entries must be acceptable
ops_list = []
idx_list = []
for ii, elem in enumerate(custom_ops):
ops_list.append(elem)
idx_list.append(ii)
ctens = tensor.tensordot(ops_list[0], [(idx_list[0],), (2,)])
## ctens has shape (rem_t,O1_0,O1_1,O1_3)
new_idx_list = [ii - 1 if ii > idx_list[0] else ii for ii in idx_list]
for ii in range(1, len(ops_list) - 1):
op = ops_list[ii]
cidx = new_idx_list[ii]
ctens = ctens.tensordot(op, [(cidx, -1), (2, 0)])
new_idx_list = [idx - 1 if idx > cidx else idx for idx in new_idx_list]
## after loop ctens has shape (t,O1_0,O1_1,...,Ok-1_1,Ok-1_3)
## final contraction
ctens = ctens.tensordot(ops_list[-1], [(0, 1, -1), (2, 3, 0)])
## ctens.shape = (O1_1,...,Ok_1)
return ctens
############################################################
# Abstract methods that should not work with the DenseMPO
############################################################
[docs]
@classmethod
def from_density_matrix(
cls, rho, n_sites, dim, conv_params, tensor_backend=None, prob=False
):
"""Conversion of density matrix to DenseMPO (not implemented yet)."""
raise NotImplementedError("Feature not yet implemented.")
[docs]
@classmethod
def from_lptn(cls, lptn, conv_params=None, **kwargs):
"""Conversion of LPTN to DenseMPO (not implemented yet)."""
raise NotImplementedError("Feature not yet implemented.")
[docs]
@classmethod
def from_mps(cls, mps, conv_params=None, **kwargs):
"""Conversion of MPS to DenseMPO (not implemented yet)."""
raise NotImplementedError("Feature not yet implemented.")
[docs]
@classmethod
def from_ttn(cls, ttn, conv_params=None, **kwargs):
"""Conversion of TTN to DenseMPO (not implemented yet)."""
raise NotImplementedError("Feature not yet implemented.")
[docs]
@classmethod
def from_tto(cls, tto, conv_params=None, **kwargs):
"""Conversion of TTO to DenseMPO (not implemented yet)."""
raise NotImplementedError("Feature not yet implemented.")
[docs]
@classmethod
def product_state_from_local_states(
cls,
mat,
padding=None,
convergence_parameters=None,
tensor_backend=None,
):
"""
Construct a product (separable) state in a suitable tensor network form, given the local
states of each of the sites.
"""
raise NotImplementedError(
"DenseMPO cannot be initialized from local product state."
)
[docs]
@classmethod
def mpi_bcast(cls, state, comm, tensor_backend, root=0):
"""
Broadcast a whole tensor network.
Arguments
---------
state : :class:`DenseMPO` (for MPI-rank root, otherwise None is acceptable)
State to be broadcasted via MPI.
comm : MPI communicator
Send state to this group of MPI processes.
tensor_backend : :class:`TensorBackend`
Needed to identity data types and tensor classes on receiving
MPI threads (plus checks on sending MPI thread).
root : int, optional
MPI-rank of sending thread with the state.
Default to 0.
"""
raise NotImplementedError("DenseMPO cannot be broadcasted yet.")
[docs]
@classmethod
def from_statevector(
cls, statevector, local_dim=2, conv_params=None, tensor_backend=None
):
"""No statevector for operators"""
raise NotImplementedError("No statevector for operators")
[docs]
def get_rho_i(self, idx):
"""No density matrix for operators"""
raise NotImplementedError("No density matrix for operators")
[docs]
def get_local_kraus_operators(self, dt: float) -> dict[int, _AbstractQteaTensor]:
"""Constructs local Kraus operators from local Lindblad operators."""
raise NotImplementedError("Currently cannot extract Kraus from DenseMPO.")
[docs]
def values(self):
"""Return the values stored in the effective ops dictionary.
Returns:
values (dict_values)
"""
return self.eff_ops.values()
# pylint: disable-next=unused-argument
[docs]
def to_dense(self, true_copy=False):
"""Convert into a TN with dense tensors (without symmetries)."""
return
[docs]
@classmethod
def read(cls, filename, tensor_backend, cmplx=True, order="F"):
"""Read a MPO from a formatted file."""
raise NotImplementedError("No read method for DenseMPO")
[docs]
def apply_projective_operator(self, site, selected_output=None, remove=False):
"""No measurements in MPO"""
raise NotImplementedError("No apply_projective_operator method for DenseMPO")
# pylint: disable-next=unused-argument
[docs]
def build_effective_operators(self, measurement_mode=False):
"""Pass"""
return
# pylint: disable-next=unused-argument
def _convert_singvals(self, dtype, device):
"""Pass"""
return
def _iter_physical_links(self):
"""Pass"""
return
# pylint: disable-next=unused-argument
[docs]
def get_bipartition_link(self, pos_src, pos_dst):
"""Pass"""
return
# pylint: disable-next=unused-argument
[docs]
def get_pos_links(self, pos):
"""Pass"""
return
[docs]
def norm(self):
"""Pass"""
return
# pylint: disable-next=unused-argument
[docs]
def set_singvals_on_link(self, pos_a, pos_b, s_vals):
"""Pass"""
return
# pylint: disable-next=unused-argument
def _update_eff_ops(self, id_step):
"""Pass"""
return
# pylint: disable-next=unused-argument
def _partial_iso_towards_for_timestep(self, pos, next_pos, no_rtens=False):
"""Pass"""
return
# pylint: disable-next=unused-argument
[docs]
def get_pos_partner_link_expansion(self, pos):
"""
Get the position of the partner tensor to use in the link expansion
subroutine
Parameters
----------
pos : int | Tuple[int]
Position w.r.t. which you want to compute the partner
Returns
-------
int | Tuple[int]
Position of the partner
int
Link of pos pointing towards the partner
int
Link of the partner pointing towards pos
"""
[docs]
def to_statevector(self, qiskit_order=False, max_qubit_equivalent=20):
"""No statevector for operators"""
raise NotImplementedError("No statevector for operators")
[docs]
def write(self, filename, cmplx=True):
"""Write the TN in python format into a FORTRAN compatible format."""
raise NotImplementedError("No write method for DenseMPO")
[docs]
class DenseMPOList(_RestrictedList):
"""
A list of dense MPOs, i.e., for building iTPOs or other MPOs.
Initialize with DenseMPOList() and then append DenseMPO-s
Here, usually, it's used to represent Hamiltonians which consist
of multiple terms. In this case, each Hamiltonian term is represented
as a DenseMPO. For example, Ising model contains nearest-neighbour
sigma_x sigma_x and sigma_z terms, and every 2-body term
sigma_x sigma_x and local term sigma_z would be stored as a separate
DenseMPO added to a DenseMPOList
"""
class_allowed = DenseMPO
@property
def has_oqs(self):
"""Return flag if the `DenseMPOList` contains any open system term."""
has_oqs = False
for elem in self:
logger.debug("elem.is_oqs %s %s", has_oqs, elem.is_oqs)
has_oqs = has_oqs or elem.is_oqs
return has_oqs
[docs]
@classmethod
def from_model(cls, model, params, tensor_backend: TensorBackend | None = None):
"""Fill class with :class:`QuantumModel` and its parameters. Not ready
for measurement before initializing with the operators. For full
preparation use `from_model_prepared()`.
"""
if tensor_backend is None:
# prevents dangerous default value similar to dict/list
tensor_backend = TensorBackend()
obj = cls()
lx_ly_lz = model.eval_lvals(params)
for term in model.hterms:
for elem, coords in term.get_interactions(lx_ly_lz, params, dim=model.dim):
weight = term.prefactor
if "weight" in elem:
weight *= elem["weight"]
pstrength = term.strength
# Checked manually, must be a linter-problem with double-inheritance
# pylint: disable-next=abstract-class-instantiated
mpo = DenseMPO(is_oqs=term.is_oqs, tensor_backend=tensor_backend)
for idx, coord in enumerate(coords):
site_term = MPOSite(
coord, elem["operators"][idx], pstrength, weight
)
mpo.append(site_term)
# pylint: disable-next=protected-access
mpo._singvals.append(None)
# Only needed on first site
pstrength = None
weight = 1.0
obj.append(mpo)
return obj
[docs]
@classmethod
def from_model_prepared(
cls,
model,
params,
operators,
tensor_backend: TensorBackend | None = None,
sym=None,
generators=None,
):
"""
Returns fully prepared dense MPOs from a model, ready for the
measurement. (Pay attention: this classmethod has two return
values.)
Arguments
---------
model : :class:`QuantumModel`
The model to build the :class:`DenseMPOList` for.
params : dict
Simulation dictionary containing parameterization and
similar settings for the model and simulation.
operators : :class:`TNOperators`
The operators to be used with the simulation.
tensor_backend : :class:`TensorBackend`, optional
Choose the tensor backend to be used for the operators
and simulation.
Default to `TensorBackend()` (numpy, CPU, complex128, etc.)
sym : list, optional
Information on symmetry in case of symmetric tensors.
Default to `None` (no symmetries).
generators : list, optional
Information on the generators of the symmetry in case of
symmetric tensors.
Default to `None` (no symmetries).
Returns
-------
obj : :class:`DenseMPOListt`
MPO representation of the model for the given
parameterization.
operators : :class:`TNOperators`
Operators converted to the corresponding tensor backend.
"""
if tensor_backend is None:
# prevents dangerous default value similar to dict/list
tensor_backend = TensorBackend()
if sym is None:
sym = []
if generators is None:
generators = []
obj = cls.from_model(model, params, tensor_backend)
# prepare the operators
base_tensor_cls = tensor_backend.base_tensor_cls
tensor_cls = tensor_backend.tensor_cls
operators = base_tensor_cls.convert_operator_dict(
operators,
params=params,
symmetries=[],
generators=[],
base_tensor_cls=base_tensor_cls,
dtype=tensor_backend.dtype,
device=tensor_backend.memory_device,
)
if base_tensor_cls != tensor_cls:
operators = tensor_cls.convert_operator_dict(
operators,
symmetries=sym,
generators=generators,
base_tensor_cls=base_tensor_cls,
)
# initialize mpo with the operators
obj.initialize(operators, params)
return obj, operators
[docs]
def initialize(self, operators, params, do_sort=True):
"""Resolve operators and parameterization for the given input."""
for elem in self:
elem.initialize(operators, params)
if do_sort:
mpos_sorted = self.sort_sites()
for ii, elem in enumerate(mpos_sorted):
self[ii] = elem
[docs]
def sort_sites(self):
"""Sort the sites in each :class:`DenseMPO`."""
mpos_sorted = DenseMPOList()
for elem in self:
elem_sorted = elem.sort_sites()
mpos_sorted.append(elem_sorted)
return mpos_sorted
[docs]
def trace(self, local_dim):
"""
Compute the trace of and MPO written in DenseMPOList form.
Parameters
----------
local_dim : list[int]
Physical dimension of each of the sites in self.
Returns
-------
operator_trace : float | complex
Corresponds to the trace of the operator.
"""
operator_trace = 0
# we compute the trace of each dense_mpo in self and then sum the results
for dense_mpo in self:
# This is potentially dangerous for overflows and problems with
# machine precision if the offsets are very different. 1000
# qubits are about 1e300 so beyond one-thousand qubits we
# soon run into overflows with double precision numbers
# and partial_trace being of the order of one
operator_trace += dense_mpo.trace(local_dim)
return operator_trace
[docs]
def mpo_product(
self,
other,
operators,
self_conj=False,
self_transpose=False,
other_conj=False,
other_transpose=False,
):
"""
Compute the product of two MPOs encoded as DenseMPOList. The order of
product is self*other. If the operators at site j are self_op_j and
other_op_j respectively, `mpo_product` returns the DenseMPOList whose
operator at site j is the operator product self_op_j*other_op_j. The
input arguments `self` and `other` remain unchanged in this method.
Parameters
----------
other : :py:class:`DenseMPOList`
Representing the operator that we are multiplying to self.
operators : :class:`TNOperators`
The operator dictionary for the simulation. The corresponding second
order operators are generated within the function if not set yet.
self_conj : Boolean, optional
Tells if self needs to be complex conjugated.
Default is False.
self_transpose : Boolean, optional
Tells if self needs to be transposed.
Default is False.
other_conj : Boolean, optional
Tells if other needs to be complex conjugated.
Default is False.
other_transpose : Boolean, optional
Tells if other needs to be transposed.
Default is False.
Return
------
mpo_product : :py:class:`DenseMPOList`
The product of the operators represented by self and other.
Details
-------
This function is potentially costly in terms of memory and computation
time. While the computational complexity of this approach cannot be
mitigated within `DenseMPOLists`, the memory requirements can be tuned
by using `_mpo_product_iter`, which is an iterator over smaller
`DenseMPOLists`.
"""
# We will need all combinations of the operators (can be called
# multiple times without generating 4th order etc).
operators.generate_products_2nd_order()
new_dense_mpo_list = []
for elem in self._mpo_product_term_iter(
other,
operators,
self_conj=self_conj,
self_transpose=self_transpose,
other_conj=other_conj,
other_transpose=other_transpose,
):
new_dense_mpo_list.append(elem)
return DenseMPOList(new_dense_mpo_list)
def _mpo_product_iter(
self,
other,
operators,
self_conj=False,
self_transpose=False,
other_conj=False,
other_transpose=False,
batch_size=None,
print_progress=False,
):
"""
Compute the product of two MPOs encoded as DenseMPOList. This version
is the iterator version, i.e., it can return multiple
:class:`DenseMPOList` building the full product via an iterator. The
order of product is self*other. If the operators at site j are
self_op_j and other_op_j respectively, `mpo_product` returns the
DenseMPOList whose operator at site j is the operator product
self_op_j*other_op_j. The input arguments `self` and `other` remain
unchanged in this method.
Parameters
----------
other : :py:class:`DenseMPOList`
Representing the operator that we are multiplying to self.
`other` is the right-hand operator in the multiplcation.
operators : :class:`TNOperators`
The operator dictionary for the simulation. The corresponding second
order operators are generated within the function if not set yet.
self_conj : Boolean, optional
Tells if self needs to be complex conjugated.
Default is False.
self_transpose : Boolean, optional
Tells if self needs to be transposed.
Default is False.
other_conj : Boolean, optional
Tells if other needs to be complex conjugated.
Default is False.
other_transpose : Boolean, optional
Tells if other needs to be transposed.
Default is False.
batch_size : int | None, optional
Build smaller MPOs to keep memory cost low. Efficiency might
suffer to a certain extent at the same time as work has to
be redone.
Default to `None` (one MPO).
print_progress : bool, optional
If True, progress will be printed to logger.
Default to False.
Yields as iterator
------------------
mpo_product : :py:class:`DenseMPOList`
The product of the operators represented by self and other
as one or more :py:class:`DenseMPOList`.
"""
# We will need all combinations of the operators (can be called
# multiple times without generating 4th order etc).
operators.generate_products_2nd_order()
if batch_size is None:
is_next_mpo = lambda ii, nn: False
else:
is_next_mpo = lambda ii, nn: ii % nn == 0
new_dense_mpo_list = []
cntr = 0
nn = len(self) * len(other)
for elem in self._mpo_product_term_iter(
other,
operators,
self_conj=self_conj,
self_transpose=self_transpose,
other_conj=other_conj,
other_transpose=other_transpose,
):
new_dense_mpo_list.append(elem)
cntr += 1
if is_next_mpo(cntr, batch_size):
if print_progress:
progress_str = (
f"_mpo_product_iter: {np.round(cntr / nn, decimals=3)}%"
)
logger.info(progress_str)
yield new_dense_mpo_list
new_dense_mpo_list = []
if len(new_dense_mpo_list) > 0:
yield new_dense_mpo_list
def _mpo_product_term_iter(
self,
other,
operators,
self_conj=False,
self_transpose=False,
other_conj=False,
other_transpose=False,
):
"""
Build the product between two MPO represented as :class:`DenseMPOList`
via an iterator, which returns term by term to build a new
:class:`DenseMPOList`
Arguments
---------
other : :py:class:`DenseMPOList`
Representing the operator that we are multiplying to self.
`other` is the right-hand operator in the multiplcation.
operators : :class:`TNOperators`
The operator dictionary for the simulation. The corresponding second
order operators are generated within the function if not set yet.
self_conj : Boolean, optional
Tells if self needs to be complex conjugated.
Default is False.
self_transpose : Boolean, optional
Tells if self needs to be transposed.
Default is False.
other_conj : Boolean, optional
Tells if other needs to be complex conjugated.
Default is False.
other_transpose : Boolean, optional
Tells if other needs to be transposed.
Default is False.
Yields as iterator
------------------
MPOs : :class:`DenseMPO`
Term-by-term which can build for example
a DenseMPOList representing the product.
"""
# Figure out the name of the operators set once
set_names = operators.set_names
if len(set_names) > 1:
# This would be tricky, but probably another loop would be sufficient
# under the assumption there are no cross-terms between operators sets.
raise QTeaLeavesError("Multiple operators sets not yet supported here.")
set_name = set_names[0]
# Collect id terms in one (added at the end)
weight_id = 0.0
for dense_mpo1 in self:
if dense_mpo1.current_max_bond_dim != 1:
raise ValueError(
"One MPO in `self` has non-dummy links. Product not implemented."
)
for dense_mpo2 in other:
new_mpo_sites_list = []
if dense_mpo2.current_max_bond_dim != 1:
raise ValueError(
"One MPO in `other` has non-dummy links. Product not implemented."
)
# create the new list of sites in the DenseMPO
new_sites = sorted(
list(set().union(dense_mpo1.sites, dense_mpo2.sites))
)
for xx in new_sites:
strength = 1.0
weight = 1.0
op_str = ""
# First MPO
if xx in dense_mpo1.sites:
idx = dense_mpo1.sites.index(xx)
strength *= dense_mpo1[idx].strength
weight *= dense_mpo1[idx].weight
op_str = _op_string_mul(
op_str, dense_mpo1[idx].str_op, self_conj, self_transpose
)
# 2nd MPO
if xx in dense_mpo2.sites:
idx = dense_mpo2.sites.index(xx)
strength *= dense_mpo2[idx].strength
weight *= dense_mpo2[idx].weight
op_str = _op_string_mul(
op_str, dense_mpo2[idx].str_op, other_conj, other_transpose
)
check_entry = operators.check_alternative_op(set_name, op_str)
if isinstance(check_entry, str):
if check_entry == "id":
# We can reduce the number of terms by collecting them
# and add a single offset at the end.
weight_id += weight
continue
op_str = check_entry
name = op_str
# Add actual MPO
pstrength = None
new_mpo_site = MPOSite(
xx, name, pstrength, weight, operators=operators
)
new_mpo_site.strength = strength
new_mpo_sites_list.append(new_mpo_site)
# Checked manually, must be a linter-problem with double-inheritance
# pylint: disable-next=abstract-class-instantiated
yield DenseMPO(new_mpo_sites_list)
if weight_id != 0:
new_mpo_site = MPOSite(0, "id", None, weight_id, operators=operators)
new_mpo_site.strength = 1.0
# Checked manually, must be a linter-problem with double-inheritance
# pylint: disable-next=abstract-class-instantiated
yield DenseMPO([new_mpo_site])
def _duplicate_check_and_set(ops_dict, key, op, check_equivalence=True):
"""
Checks if an equivalent operator is already in the dictionary and returns
the key, otherwise sets operator and returns key passed to the function.
Arguments
---------
ops_dict : :class:`TNOperators`
Operator dictionary. Restricted to single-set of operators (for now).
key : str
Suggested key for operator if not already in the dictionary.
op : :class:`_AbstractQteaTensor`
Operator represented as tensor.
check_equivalence : bool, optional
If False, skip equivalence scan for operators and only avoid key-name collisions.
Defaults to True.
Returns
-------
existing_key : str
The key under which `op` can be found now. Either existing key
or `key` from the function arguments.
"""
if len(ops_dict.set_names) != 1:
# We can resolve this by appending something to the string name
raise NotImplementedError("Not implemented for different sets.")
set_name = ops_dict.set_names[0]
# pylint: disable=protected-access
if key in ops_dict._ops_dicts[set_name]:
original_key = key
idx = 0
while op.are_equal(ops_dict._ops_dicts[set_name][key], tol=10 * op.dtype_eps):
idx += 1
key = f"{original_key}_{idx}"
if key not in ops_dict._ops_dicts[set_name]:
# Condition 1: key not inside dict, exit
break
existing_key = None
if check_equivalence:
for key_ii, op_ii in ops_dict._ops_dicts[set_name].items():
if op.are_equal(op_ii, tol=10 * op.dtype_eps):
existing_key = key_ii
break
if existing_key is None:
ops_dict[(set_name, key)] = op
existing_key = key
# pylint: enable=protected-access
return existing_key