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 randn
from .model import Model
from .tn import TensorNetwork
[docs]
class MatrixProductState(Model, qtn.MatrixProductState):
"""A Trainable MatrixProductState class.
See :class:`quimb.tensor.tensor_1d.MatrixProductState` for explanation of other attributes and methods.
"""
[docs]
def __init__(self, arrays, **kwargs):
"""Initialize the MatrixProductState.
Parameters
----------
arrays : list of array_like
The list of tensors, each of shape ``(D, D, d)``, where ``D`` is the bond dimension and ``d`` is the physical dimension.
**kwargs : dict
Additional arguments to be passed to the parent class.
"""
Model.__init__(self)
qtn.MatrixProductState.__init__(self, arrays, **kwargs)
[docs]
def normalize(self, insert=None):
"""Normalize tensors of :class:`tn4ml.models.mps.MatrixProductState`."""
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:
self.tensors[insert].modify(data=self.tensors[insert].data / norm)
def trainable_wrapper(mps: qtn.MatrixProductState, **kwargs) -> MatrixProductState:
"""Create a trainable wrapper around qtn.MatrixProductState.
Parameters
----------
mps : :class:`quimb.tensor.MatrixProductState`
Matrix Product State to be trained.
Returns
-------
:class:`tn4ml.models.mps.MatrixProductState`
"""
tensors = mps.arrays
return MatrixProductState(tensors, **kwargs)
def generate_shape(
method: str,
L: int,
bond_dim: int = 2,
phys_dim: int = 2,
cyclic: bool = False,
position: int | None = None,
class_index: int | None = None,
class_dim: int | None = None,
) -> tuple:
"""Return a tensor shape.
Parameters
----------
method : str
Method on how to create shapes of tensors.
'even' = exact dimensions as given by parameters, anything else = truncated dimensions.
L : int
Number of tensors.
bond_dim : int
Dimension of virtual indices between tensors. *Default = 4*.
phys_dim : int
Dimension of physical index for individual tensor.
cyclic : bool
Flag for indicating if MatrixProductState this tensor is part of is cyclic. *Default=False*.
position : int
Position of tensor in MatrixProductState.
class_index : int
Index of tensor that is the output node. For classification tasks only.
class_dim : int
Dimension of output node, or number of classes for classification.
Returns
-------
tuple
"""
if method == "even":
shape = (
(bond_dim, bond_dim, phys_dim, class_dim)
if class_index is not None and position == class_index
else (bond_dim, bond_dim, phys_dim)
)
if position == 1:
shape = (
(1, bond_dim, phys_dim, class_dim)
if class_index is not None and position == class_index
else (1, bond_dim, phys_dim)
)
if position == L:
shape = (
(bond_dim, 1, phys_dim, class_dim)
if class_index is not None and position == class_index
else (bond_dim, 1, phys_dim)
)
else:
assert not cyclic
j = (L + 1 - abs(2 * position - L - 1)) // 2 if position > L // 2 else position
chir = min(bond_dim, phys_dim**j)
chil = min(bond_dim, phys_dim ** (j - 1))
if position > L // 2:
(chil, chir) = (chir, chil)
if position == 1:
shape = (
(1, chir, phys_dim, class_dim)
if class_index is not None and position == class_index
else (1, chir, phys_dim)
)
elif position == L:
shape = (
(chil, 1, phys_dim, class_dim)
if class_index is not None and position == class_index
else (chil, 1, phys_dim)
)
else:
shape = (
(chil, chir, phys_dim, class_dim)
if class_index is not None and position == class_index
else (chil, chir, phys_dim)
)
return shape
def generate_ind(
L: int,
shape: tuple,
position: int,
cyclic: bool = False,
class_index: int | None = None,
) -> tuple:
"""Return the names of the tensor indices.
Parameters
----------
shape : tuple
Shape of tensor.
position : int
Position of tensor in MatrixProductState. Goes from 1 to L included.
cyclic : bool
Flag for indicating if MatrixProductState this tensor is part of is cyclic. *Default=False*.
class_index : int
Index of tensor that is the output node (that is having index for number of classes). For classification tasks only.
Returns
-------
tuple
String names of indices.
"""
ind: tuple
if len(shape) == 3:
if position == 1:
if class_index == position:
ind = (f"bond_{position - 1}", f"k{position - 1}", f"b{position - 1}")
else:
ind = (
f"bond_{position - 2}",
f"bond_{position - 1}",
f"k{position - 1}",
)
elif position == L:
if cyclic and class_index != position:
raise ValueError("Cyclic MPS cannot have class_dim")
ind = (
(f"bond_{position - 2}", f"k{position - 1}", f"b{position - 1}")
if class_index == position
else (
f"bond_{position - 2}",
f"bond_{position - 1}",
f"k{position - 1}",
)
)
else:
ind = (f"bond_{position - 2}", f"bond_{position - 1}", f"k{position - 1}")
else:
ind = (
f"bond_{position - 2}",
f"bond_{position - 1}",
f"k{position - 1}",
f"b{position - 1}",
)
return ind
def MPS_initialize( # noqa: N802
L: int,
arrays: list | None = None,
initializer: Initializer = None,
key: Any = None,
dtype: Any = jnp.float_,
shape_method: str = "even",
bond_dim: int = 4,
phys_dim: int = 2,
cyclic: bool = False,
add_identity: bool = False,
add_to_output: bool = False,
boundary: str = "obc",
class_index: int | None = None,
class_dim: int | None = None,
tags_id: str = "I{}",
compress: bool = False,
insert: int | None = None,
canonical_center: int | None = None,
**kwargs,
):
"""Initialize :class:`tn4ml.models.mps.MatrixProductState`.
Parameters
----------
L : int
Number of tensors.
initializer : :class:`jax.nn.initializers.Initializer``
Type of tensor initialization function.
key : Array
Argument key is a PRNG key (e.g. from `jax.random.key()`), used to generate random numbers to initialize the array.
dtype : Any
Type of tensor data (from `jax.numpy.float_`)
shape_method : str
Method to generate shapes for tensors.
bond_dim : int
Dimension of virtual indices between tensors. *Default = 4*.
phys_dim : int
Dimension of physical index for individual tensor.
cyclic : bool
Flag for indicating if MatrixProductState is cyclic. *Default=False*.
add_identity : bool
Flag to add identity to tensors diagonal elements.
add_to_output : bool
Flag for adding identity to diagonal elements of tensors with output indices. *Default=False*.
boundary : str
Boundary condition of MatrixProductState. *Default = 'obc'*. obc = open boundary condition. pbc = periodic boundary condition.
class_index : int
Index of tensor that is the output node for class. For classification tasks only.
class_dim : int
Dimension of output node, or number of classes for classification.
compress : bool
Flag to truncate bond dimensions.
insert : int
Index of tensor divided by norm. When `None` the norm division is distributed across all tensors
canonical_center : int
If not `None` then create canonical form around canonical center index.
kwargs : dict
Additional arguments.
Returns
-------
:class:`tn4ml.models.mps.MatrixProductState`
"""
if cyclic and shape_method != "even":
raise NotImplementedError("Change shape_method to 'even'.")
if (
initializer is not None
and callable(initializer)
and "rand_unitary" in getattr(initializer, "__qualname__", "")
):
if add_identity:
raise ValueError("rand_unitary initializer does not support add_identity.")
if compress:
raise ValueError("rand_unitary initializer does not support compress.")
if insert:
raise ValueError("rand_unitary initializer does not support insert.")
if boundary == "obc":
boundary = None
if arrays is not None:
# This means MPS for classification needs to be created with qtn.tensor_1d.TensorNetwork1DFlat class
assert (
class_index is not None
) # class_index is required when arrays or shapes are provided
if initializer is None:
initializer = randn()
if class_index is not None:
# MPS for classification
if class_index > L:
raise ValueError("class_index should be less than L.")
tensors = []
if arrays is not None:
for i, array in enumerate(arrays):
ind = generate_ind(L, array.shape, i + 1, cyclic, class_index)
tensors.append(qtn.Tensor(array, inds=ind, tags=tags_id.format(i)))
else:
for i in range(1, L + 1):
shape = generate_shape(
shape_method,
L,
bond_dim,
phys_dim,
cyclic,
i,
class_index,
class_dim,
)
ind = generate_ind(L, shape, i, cyclic, class_index)
if callable(initializer) and "rand_unitary" in getattr(
initializer, "__qualname__", ""
):
if i < class_index or i > class_index:
array = initializer(key, shape, dtype)
elif i == class_index:
# Output node
rng = np.random.default_rng()
array = jnp.asarray(rng.normal(0.0, 1.0, shape), dtype)
else:
raise ValueError(
"Check value of class_index. It should be less than L."
)
else:
array = initializer(key, shape, dtype)
if add_identity:
if len(array.shape) == 3:
copy_array = jnp.copy(array)
copy_array = copy_array.at[:, :, 0].add(
jnp.eye(array.shape[0], array.shape[1], dtype=dtype)
)
array = copy_array
elif len(array.shape) == 4: # output node
if add_to_output:
copy_array = jnp.copy(array)
identity = jnp.eye(
array.shape[0], array.shape[1], dtype=dtype
)
identity = jnp.expand_dims(identity, axis=2)
identity = jnp.broadcast_to(
identity,
(
copy_array.shape[0],
copy_array.shape[1],
copy_array.shape[3],
),
)
copy_array = copy_array.at[:, :, 0, :].add(identity)
array = copy_array
else:
raise ValueError(
"Tensors need to always be 3D or 4D in MPS for classification."
)
if boundary == "obc":
aux_array = jnp.zeros(array.shape, dtype=dtype)
if i == 1:
# Left node
aux_array = aux_array.at[:, 0, :].set(array[:, 0, :])
array = aux_array
elif i == L:
# Right node
aux_array = aux_array.at[0, :, :].set(array[0, :, :])
array = aux_array
tensors.append(qtn.Tensor(array, inds=ind, tags=tags_id.format(i - 1)))
mps = TensorNetwork(tensors, cyclic=cyclic, site_tag_id=tags_id, **kwargs)
if L > 200: # for large systems
for i, tensor in enumerate(mps.tensors):
if i == 0:
mps.left_canonize_site(i)
elif i == L - 1:
tensor.modify(data=tensor.data / jnp.linalg.norm(tensor.data))
else:
tensor.modify(data=tensor.data / jnp.linalg.norm(tensor.data))
mps.left_canonize_site(i)
if canonical_center is not None:
mps.canonicalize(canonical_center, inplace=True)
mps.normalize(insert=canonical_center)
else:
# normalize
if canonical_center is None:
mps.normalize()
else:
mps.canonize(canonical_center, inplace=True)
mps.normalize(insert=canonical_center)
else:
# MPS for regression
if arrays is not None:
tensors = []
for array in arrays:
tensors.append(jnp.squeeze(array))
else:
tensors = []
for i in range(1, L + 1):
shape = generate_shape(shape_method, L, bond_dim, phys_dim, cyclic, i)
tensor = initializer(key, shape, dtype)
if callable(initializer) and "rand_unitary" not in getattr(
initializer, "__qualname__", ""
):
if add_identity:
if len(tensor.shape) == 3:
copy_tensor = jnp.copy(tensor)
copy_tensor.at[:, :, 0].add(
jnp.eye(tensor.shape[0], tensor.shape[1], dtype=dtype)
)
tensor = copy_tensor
else:
raise ValueError(
"There was an error in generating shape. They should be 3D"
)
if boundary == "obc":
aux_tensor = jnp.zeros(tensor.shape, dtype=dtype)
if i == 1:
# Left node
aux_tensor = aux_tensor.at[:, 0, :].set(tensor[:, 0, :])
tensor = aux_tensor
elif i == L:
# Right node
aux_tensor = aux_tensor.at[0, :, :].set(tensor[0, :, :])
tensor = aux_tensor
tensors.append(jnp.squeeze(tensor))
if (
not (
callable(initializer)
and "rand_unitary" in getattr(initializer, "__qualname__", "")
)
and insert
and insert < L
and shape_method == "even"
):
tensors[insert] /= jnp.sqrt(phys_dim)
mps = MatrixProductState(tensors, **kwargs)
if compress:
if shape_method == "even":
mps.compress(form="flat", max_bond=bond_dim) # limit bond_dim
else:
raise ValueError('Compress only works with shape_method = "even".')
if L > 200: # for large systems
for i, tensor in enumerate(mps.tensors):
if i == 0:
mps.left_canonize_site(i)
elif i == L - 1:
tensor.modify(data=tensor.data / jnp.linalg.norm(tensor.data))
else:
tensor.modify(data=tensor.data / jnp.linalg.norm(tensor.data))
mps.left_canonize_site(i)
if canonical_center is not None:
mps.canonicalize(canonical_center, inplace=True)
mps.normalize(insert=canonical_center)
else:
if canonical_center is None:
norm = mps.norm()
for tensor in mps.tensors:
tensor.modify(data=tensor.data / a.do("power", norm, 1 / L))
else:
mps.canonicalize(canonical_center, inplace=True)
mps.normalize(insert=canonical_center)
return mps
