import copy
from typing import Any
import autoray as a
import jax.numpy as jnp
import numpy as np
import quimb.tensor as qtn
from jax.nn.initializers import Initializer
from quimb import *
from ..initializers import *
from .model import Model
[docs]
class TensorNetwork(Model, qtn.tensor_1d.TensorNetwork1DFlat):
"""A Trainable TensorNetwork class.
See :class:`quimb.tensor.tensor_core.TensorNetwork` for explanation of other attributes and methods.
"""
_EXTRA_PROPS = ("_L", "_site_tag_id", "cyclic")
[docs]
def __init__(
self, tensors, site_tag_id: str = "I{}", cyclic: bool = False, **kwargs
):
"""Initialize :class:`tn4ml.models.tn.ParametrizedTensorNetwork`.
Parameters
----------
tensors : list or TensorNetwork
List of tensors of :class:`quimb.tensor.tensor_core.Tensor` or :class:quimb.tensor.tensor_core.TensorNetwork.
kwargs : dict
Additional arguments.
"""
if isinstance(tensors, TensorNetwork):
Model.__init__(self)
return
Model.__init__(self)
qtn.tensor_1d.TensorNetwork1DFlat.__init__(self, tensors, **kwargs)
self._L = len(self.tensors)
self.cyclic = cyclic
self._site_tag_id = site_tag_id
[docs]
def canonize(self, where, cur_orthog="calc", info=None, bra=None, inplace=False):
"""Canonizes the tensor network."""
self.canonicalize(
where, cur_orthog=cur_orthog, info=info, bra=bra, inplace=inplace
)
[docs]
def copy(self, virtual: bool = False, deep: bool = False):
"""Copy the model.
Returns
-------
Model of the same type.
"""
if deep:
return copy.deepcopy(self)
model = self.__class__(self, virtual=virtual)
for key in self.__dict__:
model.__dict__[key] = self.__dict__[key]
return model
[docs]
def norm(self, **contract_opts) -> float:
"""Calculate norm of :class:`tn4ml.models.tn.TensorNetwork`.
Parameters
----------
contract_opts : Optional
Arguments passed to ``contract()``.
Returns
-------
float
Norm of :class:`tn4ml.models.smpo.SpacedMatrixProductOperator`
"""
norm = self.conj() & self
return norm.contract(**contract_opts) ** 0.5
[docs]
def normalize(self, insert=None) -> None:
"""Normalize tensors of :class:`tn4ml.models.tn.TensorNetwork`.
Parameters
----------
insert : int
Index of tensor divided by norm. *Default = None*. When `None` the norm division is distributed across all tensors.
"""
if not self.tensors:
raise ValueError("The tensor network is empty.")
if self.L > 200: # for large systems
for i, tensor in enumerate(self.tensors):
if i == 0:
self.left_canonize_site(i)
elif i == self.L - 1:
tensor.modify(data=tensor.data / jnp.linalg.norm(tensor.data))
else:
tensor.modify(data=tensor.data / jnp.linalg.norm(tensor.data))
self.left_canonize_site(i)
else:
norm = self.norm()
if insert is None:
for tensor in self.tensors:
tensor.modify(data=tensor.data / a.do("power", norm, 1 / self.L))
else:
if not (0 <= insert < len(self.tensors)):
raise IndexError(
f"Insert index {insert} is out of bounds for the tensor list."
)
self.tensors[insert].modify(data=self.tensors[insert].data / norm)
def trainable_wrapper(
tn: qtn.tensor_1d.TensorNetwork1DFlat, **kwargs
) -> qtn.tensor_1d.TensorNetwork1DFlat:
"""Create a trainable wrapper around qtn.tensor_1d.TensorNetwork1DFlat.
Parameters
----------
tn : :class:`quimb.tensor.TensorNetwork`
Tensor Network to be trained.
Returns
-------
:class:`tn4ml.models.tn.TensorNetwork`
"""
tensors = tn.tensors
return TensorNetwork(tensors, **kwargs)
def TN_initialize( # noqa: N802
arrays: list | None = None,
shapes: list | None = None,
key: Any = None,
initializer: Initializer = None,
inds: list | None = None,
tags_id: str = "I{}",
cyclic: bool = False,
dtype: Any = jnp.float_,
**kwargs,
) -> TensorNetwork:
"""Initialize a TensorNetwork.
Parameters
----------
arrays : list
List of arrays to be used as tensors. *Default = None*.
If None, shapes must be provided.
shapes : list
List of shapes of tensors. Each shape should be in LRP(P) format : (left, right, physical)
*Default = None*. If None, arrays must be provided.
key : Any
Random key for initialization. *Default = None*.
initializer : from `tn4ml.initializers` or `jax.nn.initializers`
Initializer for tensors. *Default = None*.
If None, tensors are initialized with random values. Only provided if arrays is None.
inds : sequence of arrays of str
List of indices for tensors. *Default = None*.
Neeeds to be provided because its showing connectivity between tensors.
Example for TN with 3 tensors:
>>> inds = [['bond0', 'k0'], ['bond0', 'bond1', 'k2'], ['bond1', 'k3']]
tags_id : str
Tag identifier for tensors. *Default = 'I{}'*.
The tag identifier should have a single placeholder for tag number.
dtype : Any
Data type for tensors. *Default = jnp.float_*.
kwargs : dict
Additional arguments.
Returns
-------
:class:`tn4ml.models.tn.TensorNetwork`
"""
if arrays is None and shapes is None:
raise ValueError("Provide either arrays or shapes to create Tensor Network.")
L = len(arrays) if arrays is not None else len(shapes)
if inds is None:
raise ValueError(
"Provide indices for tensors - connectivity map between tensors."
)
tensors = []
if arrays is not None:
if len(arrays) != len(inds):
raise ValueError("Number of tensors and indices should be same.")
for i, array in enumerate(arrays):
tensors.append(qtn.Tensor(array, inds=inds[i], tags=tags_id.format(i)))
elif shapes is not None:
if len(shapes) != len(inds):
raise ValueError("Number of tensors and indices should be same.")
for i, shape in zip(range(1, L + 1), shapes, strict=False):
if initializer is not None:
array = initializer(key, shape, dtype)
else:
rng = np.random.default_rng()
array = np.asarray(rng.normal(0.0, 1.0, shape), dtype)
tensors.append(
qtn.Tensor(array, inds=inds[i - 1], tags=tags_id.format(i - 1))
)
tn = TensorNetwork(tensors, cyclic=cyclic, site_tag_id=tags_id, **kwargs)
# normalize
tn.normalize()
return tn
