diff --git a/.circleci/config.yml b/.circleci/config.yml index 79d2572c..16dd2fc5 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -38,7 +38,9 @@ jobs: name: Run tests # This assumes pytest is installed via the install-package step above command: | - pytest --durations=0 + pip install pytest-xdist + export RENO_NUM_THREADS=1 + pytest -n 4 --durations=0 pip install primme==3.2.* pytest --durations=0 renormalizer/mps/tests/test_gs.py::test_multistate - run: diff --git a/renormalizer/__init__.py b/renormalizer/__init__.py index 6b4e5f9b..74e00c3d 100644 --- a/renormalizer/__init__.py +++ b/renormalizer/__init__.py @@ -2,20 +2,22 @@ # Author: Jiajun Ren import os import sys +import warnings -if "numpy" in sys.modules: - raise ImportError("renormalizer should be imported before numpy is imported") -# set environment variables to limit NumPy cpu usage -# Note that this should be done before NumPy is imported -for env in ["MKL_NUM_THREADS", "NUMEXPR_NUM_THREADS", "OMP_NUM_THREADS"]: - os.environ[env] = "1" - del env +reno_num_threads = os.environ.get("RENO_NUM_THREADS") +if reno_num_threads is not None: + # set environment variables to limit NumPy cpu usage + # Note that this should be done before NumPy is imported -# NEP-18 not working. For compatibility of newer NumPy version. See gh-cupy/cupy#2130 -os.environ["NUMPY_EXPERIMENTAL_ARRAY_FUNCTION"] = "0" + if "numpy" in sys.modules: + warnings.warn("renormalizer should be imported before numpy for `RENO_NUM_THREADS` to take effect") + else: + for env in ["MKL_NUM_THREADS", "NUMEXPR_NUM_THREADS", "OMP_NUM_THREADS"]: + os.environ[env] = reno_num_threads + del env -del os, sys +del os, sys, warnings from renormalizer.utils.log import init_log diff --git a/renormalizer/mps/backend.py b/renormalizer/mps/backend.py index e6978d64..ce3deaf2 100644 --- a/renormalizer/mps/backend.py +++ b/renormalizer/mps/backend.py @@ -67,8 +67,10 @@ def __init__(self): self.first_mp = False self._real_dtype = None self._complex_dtype = None - #self.use_32bits() - self.use_64bits() + if os.environ.get("RENO_FP32") is None: + self.use_64bits() + else: + self.use_32bits() def free_all_blocks(self): if not USE_GPU: @@ -130,5 +132,12 @@ def dtypes(self): def dtypes(self, target): self.real_dtype, self.complex_dtype = target + @property + def canonical_atol(self): + if self.is_32bits: + return 1e-4 + else: + return 1e-5 + backend = Backend() diff --git a/renormalizer/mps/matrix.py b/renormalizer/mps/matrix.py index abefbf5a..876a2b30 100644 --- a/renormalizer/mps/matrix.py +++ b/renormalizer/mps/matrix.py @@ -90,21 +90,25 @@ def r_combine(self): def l_combine(self): return self.reshape(self.l_combine_shape) - def check_lortho(self, rtol=1e-5, atol=1e-8): + def check_lortho(self, atol=None): """ check L-orthogonal """ + if atol is None: + atol = backend.canonical_atol tensm = asxp(self.array.reshape([np.prod(self.shape[:-1]), self.shape[-1]])) s = tensm.T.conj() @ tensm - return xp.allclose(s, xp.eye(s.shape[0]), rtol=rtol, atol=atol) + return xp.allclose(s, xp.eye(s.shape[0]), atol=atol) - def check_rortho(self, rtol=1e-5, atol=1e-8): + def check_rortho(self, atol=None): """ check R-orthogonal """ + if atol is None: + atol = backend.canonical_atol tensm = asxp(self.array.reshape([self.shape[0], np.prod(self.shape[1:])])) s = tensm @ tensm.T.conj() - return xp.allclose(s, xp.eye(s.shape[0]), rtol=rtol, atol=atol) + return xp.allclose(s, xp.eye(s.shape[0]), atol=atol) def to_complex(self): # `xp.array` always creates new array, so to_complex means copy, which is diff --git a/renormalizer/mps/mp.py b/renormalizer/mps/mp.py index f305914a..289d73e0 100644 --- a/renormalizer/mps/mp.py +++ b/renormalizer/mps/mp.py @@ -157,21 +157,21 @@ def move_qnidx(self, dstidx: int): self.qn[idx] = [self.qntot - i for i in self.qn[idx]] self.qnidx = dstidx - def check_left_canonical(self, rtol=1e-5, atol=1e-8): + def check_left_canonical(self, atol=None): """ check L-canonical """ for i in range(len(self)-1): - if not self[i].check_lortho(rtol, atol): + if not self[i].check_lortho(atol): return False return True - def check_right_canonical(self, rtol=1e-5, atol=1e-8): + def check_right_canonical(self, atol=None): """ check R-canonical """ for i in range(1, len(self)): - if not self[i].check_rortho(rtol, atol): + if not self[i].check_rortho(atol): return False return True @@ -189,18 +189,18 @@ def is_right_canonical(self): """ return self.qnidx == 0 - def ensure_left_canonical(self, rtol=1e-5, atol=1e-8): + def ensure_left_canonical(self, atol=None): if self.to_right or self.qnidx != self.site_num-1 or \ - (not self.check_left_canonical(rtol, atol)): + (not self.check_left_canonical(atol)): self.move_qnidx(0) self.to_right = True return self.canonicalise() else: return self - - def ensure_right_canonical(self, rtol=1e-5, atol=1e-8): + + def ensure_right_canonical(self, atol=None): if (not self.to_right) or self.qnidx != 0 or \ - (not self.check_right_canonical(rtol, atol)): + (not self.check_right_canonical(atol)): self.move_qnidx(self.site_num - 1) self.to_right = False return self.canonicalise() @@ -226,9 +226,9 @@ def _update_ms( self, idx: int, u: np.ndarray, vt: np.ndarray, sigma=None, qnlset=None, qnrset=None, m_trunc=None ): r""" update mps directly after svd - + """ - + if m_trunc is None: m_trunc = u.shape[1] u = u[:, :m_trunc] @@ -286,8 +286,8 @@ def _switch_direction(self): # assert self.check_right_canonical() def _get_big_qn(self, cidx: List[int], swap=False): - r""" get the quantum number of L-block and R-block renormalized basis - + r""" get the quantum number of L-block and R-block renormalized basis + Parameters ---------- cidx : list @@ -302,7 +302,7 @@ def _get_big_qn(self, cidx: List[int], swap=False): super-R-block (active site + R-block if necessary) quantum number qnmat : np.ndarray L-block + active site + R-block quantum number - + """ if len(cidx) == 2: @@ -311,7 +311,7 @@ def _get_big_qn(self, cidx: List[int], swap=False): elif len(cidx) > 2: assert False assert self.qnidx in cidx - + sigmaqn = [np.array(self._get_sigmaqn(idx)) for idx in cidx] if swap: assert len(sigmaqn) == 2 @@ -354,7 +354,7 @@ def mp_norm(self) -> float: def add(self, other: "MatrixProduct"): assert self.qntot == other.qntot assert self.site_num == other.site_num - + new_mps = self.metacopy() if other.dtype == backend.complex_dtype: new_mps.dtype = backend.complex_dtype @@ -404,7 +404,7 @@ def add(self, other: "MatrixProduct"): new_mps[-1] = concatenate((self[-1], other[-1]), axis=0) else: assert False - + #assert self.qnidx == other.qnidx new_mps.move_qnidx(other.qnidx) new_mps.to_right = other.to_right @@ -434,7 +434,7 @@ def compress(self, temp_m_trunc=None, ret_s=False): assert self.qnidx == 0 else: assert self.qnidx == self.site_num-1 - + if self.compress_config.bonddim_should_set: self.compress_config.set_bonddim(len(self)+1) # used for logging at exit @@ -481,18 +481,18 @@ def compress(self, temp_m_trunc=None, ret_s=False): else: # return singular value list return self, s_list - + def variational_compress(self, mpo=None, guess=None): r"""Variational compress an mps/mpdm/mpo - + Parameters ---------- - mpo : renormalizer.mps.Mpo, optional + mpo : renormalizer.mps.Mpo, optional Default is ``None``. if mpo is not ``None``, the returned mps is an approximation of ``mpo @ self`` guess : renormalizer.mps.MatrixProduct, optional - Initial guess of compressed mps/mpdm/mpo. Default is ``None``. - + Initial guess of compressed mps/mpdm/mpo. Default is ``None``. + Note ---- the variational compress related configurations is defined in @@ -503,12 +503,12 @@ def variational_compress(self, mpo=None, guess=None): mp : renormalizer.mps.MatrixProduct a new compressed mps/mpdm/mpo, ``self`` is not overwritten. ``guess`` is overwritten. - + """ if mpo is None: raise NotImplementedError("Recommend to use svd to compress a single mps/mpo/mpdm.") - + if guess is None: # a minimal representation of self and mpo compressed_mpo = mpo.copy().canonicalise().compress( @@ -530,13 +530,13 @@ def variational_compress(self, mpo=None, guess=None): logger.debug(f"isweep: {isweep}") logger.debug(f"mmax, percent: {mmax}, {percent}") logger.debug(f"mps bond dims: {mps.bond_dims}") - + for imps in mps.iter_idx_list(full=True): if method == "2site" and \ ((mps.to_right and imps == mps.site_num-1) or ((not mps.to_right) and imps == 0)): break - + if mps.to_right: lmethod, rmethod = "System", "Enviro" else: @@ -549,7 +549,7 @@ def variational_compress(self, mpo=None, guess=None): elif method == "2site": if mps.to_right: lidx = imps - 1 - cidx = [imps, imps+1] + cidx = [imps, imps+1] ridx = imps + 2 else: lidx = imps - 2 @@ -571,7 +571,7 @@ def variational_compress(self, mpo=None, guess=None): # get the quantum number pattern qnbigl, qnbigr, qnmat = mps._get_big_qn(cidx) - + # center mo cmo = [asxp(mpo[idx]) for idx in cidx] if method == "1site": @@ -587,9 +587,9 @@ def variational_compress(self, mpo=None, guess=None): if mps.compress_config.ofs is not None: # need to swap the original MPS. Tedious to implement and probably not useful. raise NotImplementedError("OFS for variational compress not implemented") - + mps._switch_direction() - + # check convergence if isweep > 0 and percent == 0: error = mps.distance(mps_old) / np.sqrt(mps.dot(mps.conj()).real) @@ -597,21 +597,21 @@ def variational_compress(self, mpo=None, guess=None): if error < mps.compress_config.vrtol: logger.info("Variational compress is converged!") break - + mps_old = mps.copy() else: logger.warning("Variational compress is not converged! Please increase the procedure!") - + # remove the redundant bond dimension near the boundary of the MPS mps.canonicalise() logger.info(f"{mps}") - + return mps - + def _update_mps(self, cstruct, cidx, qnbigl, qnbigr, Mmax, percent=0): r"""update mps with basis selection algorithm of J. Chem. Phys. 120, 3172 (2004). - + Parameters --------- cstruct : ndarray, List[ndarray] @@ -627,21 +627,21 @@ def _update_mps(self, cstruct, cidx, qnbigl, qnbigr, Mmax, percent=0): percent : float, int The percentage of renormalized basis which is equally selected from each quantum number section rather than according to singular - values. ``percent`` is defined in ``procedure`` of + values. ``percent`` is defined in ``procedure`` of `renormalizer.utils.configs.OptimizeConfig` and ``vprocedure`` of `renormalizer.utils.configs.CompressConfig`. Returns ------- - averaged_ms : + averaged_ms : if ``cstruct`` is a list, ``averaged_ms`` is a list of rotated ms of each element in ``cstruct`` as a single site calculation. It is used for better initial guess in SA-DMRG algorithm. Otherwise, ``None`` is returned. - ``self`` is overwritten inplace. - + ``self`` is overwritten inplace. + """ - + system = "L" if self.to_right else "R" # step 1: get the selected U, S, V @@ -720,9 +720,9 @@ def _update_mps(self, cstruct, cidx, qnbigl, qnbigr, Mmax, percent=0): ms, msdim, msqn, compms = select_basis( Vset, SVset, qnrnew, Uset, Mmax, percent=percent ) - ms = xp.moveaxis(ms.reshape(list(qnbigr.shape) + [msdim]), -1, 0) + ms = xp.moveaxis(ms.reshape(list(qnbigr.shape) + [msdim]), -1, 0) compms = compms.reshape(list(qnbigl.shape) + [msdim]) - + else: # state-averaged method ddm = 0.0 @@ -796,14 +796,14 @@ def _update_mps(self, cstruct, cidx, qnbigl, qnbigr, Mmax, percent=0): if cidx[0] != 0: if type(cstruct) is list: for c in rotated_c: - averaged_ms.append(tensordot(self[cidx[0] - 1], c, axes=1)) + averaged_ms.append(tensordot(self[cidx[0] - 1], c, axes=1)) self[cidx[0] - 1] = tensordot(self[cidx[0] - 1], compms, axes=1) self.qn[cidx[0]] = msqn self.qnidx = cidx[0] - 1 else: if type(cstruct) is list: for c in rotated_c: - averaged_ms.append(tensordot(c, self[cidx[0]], axes=1)) + averaged_ms.append(tensordot(c, self[cidx[0]], axes=1)) self[cidx[0]] = tensordot(compms, self[cidx[0]], axes=1) self.qnidx = 0 else: @@ -993,12 +993,12 @@ def _array2mt(self, array, idx, allow_dump=True): return dump_name return mt - + def build_empty_mp(self, num): self._mp = [[None]] * num - + def dump(self, fname, other_attrs=None): - + if other_attrs is None: other_attrs = [] elif isinstance(other_attrs, str): @@ -1096,7 +1096,7 @@ def __sub__(self, other: "MatrixProduct"): def append(self, array): new_mt = self._array2mt(array, len(self)) self._mp.append(new_mt) - + def __str__(self): if self.is_mps: string = "mps" @@ -1107,7 +1107,7 @@ def __str__(self): else: assert False template_str = "{} current size: {}, Matrix product bond dim:{}" - + return template_str.format(string, sizeof_fmt(self.total_bytes), self.bond_dims,) def __del__(self): diff --git a/renormalizer/mps/svd_qn.py b/renormalizer/mps/svd_qn.py index 2213fd48..fef2b5fb 100644 --- a/renormalizer/mps/svd_qn.py +++ b/renormalizer/mps/svd_qn.py @@ -4,7 +4,7 @@ import scipy.linalg -from renormalizer.mps.backend import np +from renormalizer.mps.backend import np, backend logger = logging.getLogger(__name__) @@ -53,13 +53,13 @@ def add_orthonormal_basis(u): # add `n` basis. `n` is empirical m, n = u.shape assert 2 * n < m - assert np.allclose(u.T.conj() @ u, np.eye(n)) + assert np.allclose(u.T.conj() @ u, np.eye(n), atol=backend.canonical_atol) a = np.random.rand(m,n) a = a - u @ (u.T.conj() @ a) q, _ = scipy.linalg.qr(a, mode='economic') res = np.concatenate([u, q], axis=1) - assert np.allclose(res.T.conj() @ res, np.eye(2 * n)) + assert np.allclose(res.T.conj() @ res, np.eye(2 * n), atol=backend.canonical_atol) return res diff --git a/renormalizer/mps/tests/test_evolve.py b/renormalizer/mps/tests/test_evolve.py index e199685e..acd36be5 100644 --- a/renormalizer/mps/tests/test_evolve.py +++ b/renormalizer/mps/tests/test_evolve.py @@ -8,6 +8,7 @@ import numpy as np from renormalizer.model import Model +from renormalizer.mps.backend import backend from renormalizer.mps import Mps, Mpo, MpDm from renormalizer.utils import EvolveMethod, EvolveConfig, CompressConfig, CompressCriteria, Quantity, OFS from renormalizer.tests.parameter_exact import qutip_clist, qutip_h, model @@ -23,7 +24,7 @@ def f(model, run_qutip=True): init_mpdm = MpDm.from_mps(init_mps).expand_bond_dimension(hint_mpo=tentative_mpo) e = init_mps.expectation(tentative_mpo) mpo = Mpo(model, offset=Quantity(e)) - + if run_qutip: # calculate result in ZT. FT result is exactly the same TIME_LIMIT = 10 @@ -98,6 +99,8 @@ def test_pc_tdrk(init_state, rk_solver): @pytest.mark.parametrize("with_mu", (True, False)) @pytest.mark.parametrize("force_ovlp", (True, False)) def test_tdvp_vmf(init_state, with_mu, force_ovlp, atol): + if backend.is_32bits: + pytest.skip("VMF too stiff for 32 bit float point operation") mps = init_state.copy() method = EvolveMethod.tdvp_mu_vmf if with_mu else EvolveMethod.tdvp_vmf mps.evolve_config = EvolveConfig(method, ivp_rtol=1e-4, ivp_atol=1e-7, force_ovlp=force_ovlp) diff --git a/renormalizer/transport/tests/test_kubo.py b/renormalizer/transport/tests/test_kubo.py index d911cd84..d5900ba0 100644 --- a/renormalizer/transport/tests/test_kubo.py +++ b/renormalizer/transport/tests/test_kubo.py @@ -7,6 +7,7 @@ from renormalizer.model import Phonon, Mol, HolsteinModel, Model from renormalizer.model.basis import BasisSimpleElectron, BasisSHO from renormalizer.model.op import Op +from renormalizer.mps.backend import backend from renormalizer.transport.kubo import TransportKubo from renormalizer.utils import Quantity, CompressConfig, EvolveConfig, EvolveMethod, CompressCriteria from renormalizer.utils.qutip_utils import get_clist, get_blist, get_holstein_hamiltonian, get_qnidx, \ @@ -58,6 +59,8 @@ def get_qutip_holstein_kubo(model, temperature, time_series): def test_peierls_kubo(): + if backend.is_32bits: + pytest.skip("VMF too stiff for 32 bit float point operation") # number of mol n = 4 # electronic coupling diff --git a/renormalizer/utils/log.py b/renormalizer/utils/log.py index 750544d6..ace982d9 100644 --- a/renormalizer/utils/log.py +++ b/renormalizer/utils/log.py @@ -7,21 +7,19 @@ import numpy as np -root_logger = logging.root +package_logger = logging.getLogger("renormalizer") default_stream_handler = logging.StreamHandler() default_formatter = logging.Formatter("%(asctime)s[%(levelname)s] %(message)s") def init_log(level=logging.DEBUG): - root_logger.setLevel(level) + package_logger.setLevel(level) default_stream_handler.setLevel(logging.DEBUG) default_stream_handler.setFormatter(default_formatter) - default_stream_handler.addFilter(logging.Filter("renormalizer")) - - root_logger.addHandler(default_stream_handler) + package_logger.addHandler(default_stream_handler) def set_stream_level(level): @@ -29,8 +27,8 @@ def set_stream_level(level): def disable_stream_output(): - if default_stream_handler in root_logger.handlers: - root_logger.removeHandler(default_stream_handler) + if default_stream_handler in package_logger.handlers: + package_logger.removeHandler(default_stream_handler) def register_file_output(file_path, mode="w", level=DEBUG): @@ -38,7 +36,7 @@ def register_file_output(file_path, mode="w", level=DEBUG): file_handler.setLevel(level) file_handler.setFormatter(default_formatter) file_handler.addFilter(logging.Filter("renormalizer")) - root_logger.addHandler(file_handler) + package_logger.addHandler(file_handler) NP_ERRCONFIG = {"divide": "raise", "over": "raise", "under": "warn", "invalid": "raise"} diff --git a/setup.py b/setup.py index ab25031f..69691149 100644 --- a/setup.py +++ b/setup.py @@ -11,7 +11,7 @@ setuptools.setup( name="renormalizer", - version="0.0.3", + version="0.0.4", packages=setuptools.find_packages(), long_description=long_description, long_description_content_type="text/markdown", @@ -24,4 +24,4 @@ # How to publish the library to pypi # python setup.py sdist -# twine upload -s dist/* \ No newline at end of file +# twine upload -s dist/*