def _array2mt(self, array, idx):
if isinstance(array, Matrix):
mt = array.astype(self.dtype)
else:
mt = Matrix(array, dtype=self.dtype)
if self.use_dummy_qn:
mt.sigmaqn = np.zeros(mt.pdim_prod, dtype=np.int)
else:
mt.sigmaqn = self._get_sigmaqn(idx)
return mt
def _array2mt(self, array, idx, allow_dump=True):
# convert dtype
if isinstance(array, Matrix):
mt = array.astype(self.dtype)
else:
mt = Matrix(array, dtype=self.dtype)
if mt.pdim[0] != self.pbond_list[idx]:
raise ValueError(
"Matrix physical bond dimension does not match system information"
)
# setup the matrix
mt.sigmaqn = self._get_sigmaqn(idx)
# array too large. Should be stored in disk
# use ``while`` to handle the multiple-exit logic
while allow_dump and self.compress_config.dump_matrix_size < mt.array.nbytes:
dir_with_id = os.path.join(self.compress_config.dump_matrix_dir,
str(id(self)))
if not os.path.exists(dir_with_id):
try:
os.mkdir(dir_with_id)
except:
logger.exception(
"Creating dump dir failed. Working with the matrix in memory."
)
break
dump_name = os.path.join(dir_with_id, f"{idx}.npy")
try:
array = mt.array
if not array.flags.c_contiguous and not array.flags.f_contiguous:
# for faster dump (3x). Costs more memory.
array = np.ascontiguousarray(array)
np.save(dump_name, array)
except:
logger.exception(
"Save matrix to disk failed. Working with the matrix in memory."
)
break
return dump_name
return mt
def __getitem__(self, item):
mt_or_str_or_list = self._mp[item]
if isinstance(mt_or_str_or_list, list):
assert isinstance(item, slice)
for elem in mt_or_str_or_list:
if isinstance(elem, str):
# load all matrices to memory will make
# the dump mechanism pointless
raise IndexError("Can't slice on dump matrices.")
if isinstance(mt_or_str_or_list, str):
try:
mt = Matrix(np.load(mt_or_str_or_list), dtype=self.dtype)
mt.sigmaqn = self._get_sigmaqn(item)
except:
logger.exception(f"Can't load matrix from {mt_or_str_or_list}")
raise RuntimeError("MPS internal structure corrupted.")
else:
if not isinstance(mt_or_str_or_list, (Matrix, type(None))):
raise RuntimeError(
f"Unknown matrix type: {type(mt_or_str_or_list)}")
mt = mt_or_str_or_list
return mt