diff --git a/pyblock2/driver/core.py b/pyblock2/driver/core.py index 9e7c6fea..b837a4a8 100644 --- a/pyblock2/driver/core.py +++ b/pyblock2/driver/core.py @@ -557,6 +557,7 @@ def __init__( scratch="./nodex", clean_scratch=True, restart_dir=None, + restart_dir_per_sweep=None, n_threads=None, n_mkl_threads=1, symm_type=SymmetryTypes.SU2, @@ -585,6 +586,10 @@ def __init__( restart_dir : None or str If not None, MPS will be copied to the given directory after each DMRG sweep. Default is None (MPS will not be copied). + restart_dir_per_sweep : None or str + If not None, MPS will be copied to the given directory after each DMRG sweep, + and the MPSs from different sweeps will be kept in separate directories. + Default is None (MPS will not be copied). n_threads : None or int Number of threads. When MPI is used, this is the number of threads for each MPI processor. Default is None, and the max number of threads available on this node will be used. @@ -623,6 +628,7 @@ def __init__( self._scratch = scratch self._restart_dir = restart_dir + self._restart_dir_per_sweep = restart_dir_per_sweep self.stack_mem = stack_mem self.stack_mem_ratio = stack_mem_ratio self.fp_codec_cutoff = fp_codec_cutoff @@ -685,6 +691,19 @@ def restart_dir(self, restart_dir): self._restart_dir = restart_dir self.frame.restart_dir = restart_dir + @property + def restart_dir_per_sweep(self): + """ + If not None, MPS will be copied to the given directory after each DMRG sweep, + and the MPSs from different sweeps will be kept in separate directories. + """ + return self._restart_dir_per_sweep + + @restart_dir_per_sweep.setter + def restart_dir_per_sweep(self, restart_dir_per_sweep): + self._restart_dir_per_sweep = restart_dir_per_sweep + self.frame.restart_dir_per_sweep = restart_dir_per_sweep + def set_symm_type(self, symm_type, reset_frame=True): """ Change the symmetry type of this :class:`DMRGDriver`. @@ -750,6 +769,9 @@ def set_symm_type(self, symm_type, reset_frame=True): self.mpi.barrier() self.frame.restart_dir = self.restart_dir + if self.restart_dir_per_sweep is not None: + self.frame.restart_dir_per_sweep = self.restart_dir_per_sweep + def set_symmetry_groups(self, *args, hint=None): """ Set the combination of symmetry sub-groups for ``symm_type = SAny``.