Source code for netket.hilbert.tensor_hilbert

# Copyright 2021 The NetKet Authors - All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Optional, List, Union

import numpy as np

from .discrete_hilbert import DiscreteHilbert, _is_indexable


# TODO: Make parametric class
class TensorHilbert(DiscreteHilbert):
    r"""Tensor product of several Discrete sub-spaces, representing the space

    In general you should not construct this object directly, but you should
    simply multiply different hilbert spaces together. In this case, Python's
    `*` operator will be interpreted as a tensor product.

    This Hilbert can be used as a replacement anywhere a Uniform Hilbert space
    is not required.

    Examples:
        Couple a bosonic mode with spins

        >>> from netket.hilbert import Spin, Fock
        >>> Fock(3)*Spin(0.5, 5)
        Fock(n_max=3, N=1)*Spin(s=1/2, N=5)
        >>> type(_)
        <class 'netket.hilbert.tensor_hilbert.TensorHilbert'>

    """

[docs] def __init__(self, *hilb_spaces: DiscreteHilbert): r"""Constructs a tensor Hilbert space Args: *hilb: An iterable object containing at least 1 hilbert space. """ self._hilbert_spaces = hilb_spaces self._n_hilbert_spaces = len(hilb_spaces) self._hilbert_i = np.concatenate( [[i for _ in range(hi.size)] for (i, hi) in enumerate(hilb_spaces)] ) self._sizes = tuple([hi.size for hi in hilb_spaces]) shape = np.concatenate([hi.shape for hi in hilb_spaces]) self._cum_sizes = np.cumsum(self._sizes) self._cum_indices = np.concatenate([[0], self._cum_sizes]) self._size = sum(self._sizes) self._delta_indices_i = np.array( [self._cum_indices[i] for i in self._hilbert_i] ) # pre-compute indexing data iff the tensor space is still indexable if all(hi.is_indexable for hi in hilb_spaces) and _is_indexable(shape): self._ns_states = [hi.n_states for hi in self._hilbert_spaces] self._ns_states_r = np.flip(self._ns_states) self._cum_ns_states = np.concatenate([[0], np.cumprod(self._ns_states)]) self._cum_ns_states_r = np.flip( np.cumprod(np.concatenate([[1], np.flip(self._ns_states)]))[:-1] ) self._n_states = np.prod(self._ns_states) super().__init__(shape=shape)
@property def size(self) -> int: return self._size @property def is_finite(self): return all([hi.is_finite for hi in self._hilbert_spaces]) def _sub_index(self, i): for (j, sz) in enumerate(self._cum_sizes): if i < sz: return j # def size_at_index(self, i): # # j = self._sub_index(i) # # return self._hilbert_spaces[j].size_at_index(i-self._cum_indices[j-1]) # return self._hilbert_spaces[self._hilbert_i[i]].size_at_index( # i - self._delta_indices_i[i] # )
[docs] def states_at_index(self, i): # j = self._sub_index(i) # return self._hilbert_spaces[j].states_at_index(i-self._cum_indices[j-1]) return self._hilbert_spaces[self._hilbert_i[i]].states_at_index( i - self._delta_indices_i[i] )
@property def n_states(self) -> int: if not self.is_indexable: raise RuntimeError("The hilbert space is too large to be indexed.") return self._n_states def _numbers_to_states(self, numbers, out): # !!! WARNING # This code assumes that states are stored in a MSB # (Most Significant Bit) format. # We assume that the rightmost-half indexes the LSBs # and the leftmost-half indexes the MSBs # HilbertIndex-generated states respect this, as they are: # 0 -> [0,0,0,0] # 1 -> [0,0,0,1] # 2 -> [0,0,1,0] # etc... rem = numbers for (i, dim) in enumerate(self._ns_states_r): rem, loc_numbers = np.divmod(rem, dim) hi_i = self._n_hilbert_spaces - (i + 1) self._hilbert_spaces[hi_i].numbers_to_states( loc_numbers, out=out[:, self._cum_indices[hi_i] : self._cum_sizes[hi_i]] ) return out def _states_to_numbers(self, states, out): out[:] = 0 temp = out.copy() # !!! WARNING # See note above in numbers_to_states for (i, dim) in enumerate(self._cum_ns_states_r): self._hilbert_spaces[i].states_to_numbers( states[:, self._cum_indices[i] : self._cum_sizes[i]], out=temp ) out += temp * dim return out def __repr__(self): _str = "{}".format(self._hilbert_spaces[0]) for hi in self._hilbert_spaces[1:]: _str += "*{}".format(hi) return _str @property def _attrs(self): return self._hilbert_spaces def __mul__(self, other): spaces_l = self._hilbert_spaces[:-1] space_center_l = self._hilbert_spaces[-1] if isinstance(other, TensorHilbert): space_center_r = other._hilbert_spaces[0] spaces_r = other._hilbert_spaces[1:] else: space_center_r = other spaces_r = tuple() try: spaces_center = space_center_l * space_center_r if isinstance(spaces_center, TensorHilbert): spaces_center = (space_center_l, space_center_r) else: spaces_center = (spaces_center,) except Exception: spaces_center = (space_center_l, space_center_r) return TensorHilbert(*spaces_l, *spaces_center, *spaces_r)
[docs] def ptrace(self, sites: Union[int, List]) -> Optional[DiscreteHilbert]: if isinstance(sites, int): sites = [sites] sites = np.sort(sites) for site in sites: if site < 0 or site >= self.size: raise ValueError( f"Site {site} not in this hilbert space of site {self.size}" ) Nsites = len(sites) if self.size - Nsites == 0: return None else: new_hilberts = [] sz = 0 for hilb in self._hilbert_spaces: sites_this_hilb = ( sites[np.logical_and(sites >= sz, sites < sz + hilb.size)] - sz ) if len(sites_this_hilb) == 0: new_hilberts.append(hilb) else: ptraced_hilb = hilb.ptrace(sites_this_hilb) if ptraced_hilb is not None: new_hilberts.append(ptraced_hilb) sz += hilb.size if len(new_hilberts) == 0: return None elif len(new_hilberts) >= 1: hilb = new_hilberts[0] for h in new_hilberts[1:]: hilb = hilb * h return hilb