From ec39c5cce95d6717b74f52e55643c8b0da672c25 Mon Sep 17 00:00:00 2001 From: alexiahartzell Date: Tue, 1 Sep 2026 11:19:45 -0500 Subject: [PATCH] MesoHOPS 1.8.0 Highlights of v1.8.0 (vs. mesohops 1.7.0): Tensor-network HOPS (TadHOPS): - src/mesohops/tensor/: new subpackage representing the hierarchy wavefunction as an MPS rather than a flat auxiliary-enumerated vector. hops_tensor_wavefunction.py holds MPS storage, normalization, operator application, and bond-dimension control; hops_tensor_eom.py and tensor_eom_functions.py evaluate the EOM directly on the MPS; mpo_constructors.py builds the Hamiltonian, dipole, and state-number MPOs. Hierarchy depth lives in the MPS core dimensions (k_max + 1 per mode), so n_hier, n_hmodes, and ADAPTIVE_H do not apply on this path. Representation selected by method: 'number' (ground state implicit as the all-zero MPS configuration) or 'fullstate'. - src/mesohops/tensor/tdvp.py: one- and two-site TDVP following Paeckel et al. (2019), in mixed canonical form over bare lists of np.ndarray. Local update solvers: Arnoldi, Lanczos, solve_ivp. - src/mesohops/trajectory/hops_tensor_trajectory.py: user-facing driver. INTEGRATOR accepts 'TDVP1' and 'TDVP2' alongside 'RUNGE_KUTTA'; TDVP_UPDATE_TYPE selects the local update scheme (default 'krylov'). - src/mesohops/integrator/tensor_integrator.py: MPS-aware RK4 path. - src/mesohops/util/tensor_operations.py: core MPS arithmetic (add, compress via Oseledets rounding, extraction, contraction). - Adaptivity is NOT supported for tensor HOPS in this release. The adaptive paths (hops_tensor_basis.update_basis, tensor_functions_adaptive.py, HopsTensorTrajectory.make_adaptive) are present but unfinished, and are not blocked at runtime. Non-adaptive tensor HOPS is unaffected. list_permanent_sites raises NotImplementedError, since HopsTensorBasis never reads it and would otherwise ignore it silently. Hierarchy: - src/mesohops/basis/hops_hierarchy.py: add TRUNCATION_METHOD ('triangular', the default and prior behavior, or 'rectangular'). define_rectangular_hierarchy() admits every per-mode depth combination in [0, MAXHIER]^n_hmodes, matching the MPS core structure. Rectangular truncation raises NotImplementedError for adaptive calculations: the vector adaptive flux filters enforce the MAXHIER boundary through total auxiliary depth, which is inherently triangular. - src/mesohops/basis/hops_hierarchy.py: add_connections no longer short-circuits the k+1 loop on total depth, required for rectangular hierarchies whose per-mode depths may sum above MAXHIER. - src/mesohops/basis/hops_hierarchy.py: a negative _count_by_modes entry now raises ValueError rather than printing and continuing; the MAXHIER > 255 notice moves from print to warnings.warn; remove the private static method _const_aux_edge(). Spectroscopy: - src/mesohops/util/nondyadic_spectroscopy.py: new dispatch on trajectory type, tensor method, and Hilbert-space convention (embedded, vacuum, excited_only). System and storage: - src/mesohops/basis/hops_system.py: add flag_nearest_neighbor_ham, True when all non-zero Hamiltonian elements satisfy |row - col| <= 1. Explicit stored zeros are eliminated first so padded sparse input is not misclassified. - src/mesohops/storage/hops_storage.py: add psi_g_traj storage option capturing <0,...,0|psi> per step, needed by the vacuum-convention spectroscopy path because extract_psi sees only excited-state slots. - src/mesohops/storage/storage_functions.py: add tensor-aware save_phi_traj_tensor, save_phi_norm_tensor, save_max_tensor_complexity. Integration: - src/mesohops/integrator/integrator_rk.py -> integrator.py: module rename. Imports of runge_kutta_step and runge_kutta_variables must be updated. - src/mesohops/trajectory/hops_trajectory.py: move integrator setup into an overridable _setup_integrator() hook so subclasses can register additional integrators. - src/mesohops/trajectory/hops_trajectory.py: add STORE_STEP_TIMING integration parameter (default False). When enabled, storage.metadata["LIST_PROPAGATION_TIME"] holds (t, elapsed) tuples per step rather than one total-elapsed float per propagate() call. Coverage: - .github/workflows/coverage.yml: also run on pushes to master, so Codecov has a base report to diff against; upgrade codecov-action from the sunset v3 uploader to v5; rename the deprecated file: input to files:; set fail_ci_if_error to true. Uploads had been failing silently since the workflow was introduced, leaving every job green and Codecov empty. - .github/workflows/coverage.yml: run the coverage suite at --level 3, overriding the --level=1 in pytest.ini addopts. At level 1 the level-2 and level-3 tests are deselected and contribute nothing, understating measured coverage. - README.md: add coverage badge. Cleanup: - tests/test_hierarchy_class.py -> test_hops_hierarchy.py: rename to match its subject. - tests/test_eom_hops_ksuper.py: correct a stale library name in a docstring. - .gitignore: expand from a single *.pyc line to a full Python ignore set, so coverage and build artifacts stop appearing as untracked noise. - style_guide.md: document when to annotate docstring units, including the [units: dimensionless] convention and bracket contents for heterogeneous container types. - .github/CODEOWNERS: add, gating master merges on approval from a listed owner. Takes effect only where branch protection enables require_code_owner_reviews. Test results: level-1 suite passes 780/780 (2 xfailed, 93 deselected). --- .github/CODEOWNERS | 3 + .github/workflows/coverage.yml | 22 +- .gitignore | 201 +- README.md | 2 + pyproject.toml | 2 +- src/mesohops/basis/hops_hierarchy.py | 205 +- src/mesohops/basis/hops_modes.py | 1 - src/mesohops/basis/hops_system.py | 27 + src/mesohops/eom/hops_eom.py | 19 +- .../{integrator_rk.py => integrator.py} | 61 +- src/mesohops/integrator/tensor_integrator.py | 414 +++ src/mesohops/storage/hops_storage.py | 3 + src/mesohops/storage/storage_functions.py | 104 +- .../tensor/__init__.py} | 0 src/mesohops/tensor/hops_tensor_basis.py | 234 ++ src/mesohops/tensor/hops_tensor_eom.py | 564 ++++ .../tensor/hops_tensor_wavefunction.py | 776 +++++ src/mesohops/tensor/mpo_constructors.py | 1556 ++++++++++ src/mesohops/tensor/tdvp.py | 1364 +++++++++ src/mesohops/tensor/tensor_eom_functions.py | 406 +++ .../tensor/tensor_functions_adaptive.py | 227 ++ src/mesohops/trajectory/hops_dyadic.py | 2 +- .../trajectory/hops_tensor_trajectory.py | 1105 ++++++++ src/mesohops/trajectory/hops_trajectory.py | 34 +- src/mesohops/util/nondyadic_spectroscopy.py | 608 ++++ src/mesohops/util/tensor_operations.py | 909 ++++++ style_guide.md | 14 + tests/integrated_tests/test_tensor_LTC_eom.py | 366 +++ tests/test_dimer_of_dimers_tensor.py | 729 +++++ tests/test_eom_hops_ksuper.py | 2 +- ...rarchy_class.py => test_hops_hierarchy.py} | 304 +- tests/test_hops_storage.py | 120 +- tests/test_hops_system.py | 35 + tests/test_hops_tensor_basis.py | 250 ++ tests/test_hops_tensor_eom.py | 1137 ++++++++ tests/test_hops_tensor_trajectory.py | 1966 +++++++++++++ tests/test_hops_tensor_wavefunction.py | 1567 ++++++++++ tests/test_hops_trajectory.py | 121 + tests/test_integrator_rk.py | 9 +- tests/test_mpo_constructors.py | 2241 +++++++++++++++ tests/test_nondyadic_spectroscopy.py | 1696 +++++++++++ tests/test_tdvp.py | 2522 +++++++++++++++++ tests/test_tensor_basis_shared_refs.py | 135 + tests/test_tensor_eom_functions.py | 1849 ++++++++++++ tests/test_tensor_functions_adaptive.py | 393 +++ tests/test_tensor_integrator.py | 1310 +++++++++ tests/test_tensor_nonuniform_modes.py | 138 + tests/test_tensor_operations.py | 1385 +++++++++ tests/test_timing_tests.py | 17 +- 49 files changed, 27003 insertions(+), 152 deletions(-) create mode 100644 .github/CODEOWNERS rename src/mesohops/integrator/{integrator_rk.py => integrator.py} (74%) create mode 100644 src/mesohops/integrator/tensor_integrator.py rename src/{conftest.py => mesohops/tensor/__init__.py} (100%) create mode 100644 src/mesohops/tensor/hops_tensor_basis.py create mode 100644 src/mesohops/tensor/hops_tensor_eom.py create mode 100644 src/mesohops/tensor/hops_tensor_wavefunction.py create mode 100644 src/mesohops/tensor/mpo_constructors.py create mode 100644 src/mesohops/tensor/tdvp.py create mode 100644 src/mesohops/tensor/tensor_eom_functions.py create mode 100644 src/mesohops/tensor/tensor_functions_adaptive.py create mode 100644 src/mesohops/trajectory/hops_tensor_trajectory.py create mode 100644 src/mesohops/util/nondyadic_spectroscopy.py create mode 100644 src/mesohops/util/tensor_operations.py create mode 100644 tests/integrated_tests/test_tensor_LTC_eom.py create mode 100644 tests/test_dimer_of_dimers_tensor.py rename tests/{test_hierarchy_class.py => test_hops_hierarchy.py} (58%) create mode 100644 tests/test_hops_tensor_basis.py create mode 100644 tests/test_hops_tensor_eom.py create mode 100644 tests/test_hops_tensor_trajectory.py create mode 100644 tests/test_hops_tensor_wavefunction.py create mode 100644 tests/test_mpo_constructors.py create mode 100644 tests/test_nondyadic_spectroscopy.py create mode 100644 tests/test_tdvp.py create mode 100644 tests/test_tensor_basis_shared_refs.py create mode 100644 tests/test_tensor_eom_functions.py create mode 100644 tests/test_tensor_functions_adaptive.py create mode 100644 tests/test_tensor_integrator.py create mode 100644 tests/test_tensor_nonuniform_modes.py create mode 100644 tests/test_tensor_operations.py diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS new file mode 100644 index 0000000..4c495bc --- /dev/null +++ b/.github/CODEOWNERS @@ -0,0 +1,3 @@ +# ABOUTME: Gates master merges - required approval must come from a listed owner. +# ABOUTME: Works with branch protection require_code_owner_reviews; admins may bypass. +* @digbennett @alexiahartzell diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index ea99a34..8ee6c8c 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -3,6 +3,11 @@ name: Code Coverage on: pull_request: branches: [ main, master ] + # Codecov needs a coverage report on the base branch to compute a diff and to + # render the badge. Without push builds there is no baseline, so PR comments + # show "unknown" and the badge never populates. + push: + branches: [ main, master ] permissions: contents: read @@ -33,8 +38,12 @@ jobs: - name: Run tests with coverage continue-on-error: true run: | - # Run tests and generate both coverage and test report - pytest --cov=mesohops --cov-report=xml --cov-report=term --junitxml=test-results.xml -v || true + # Run tests and generate both coverage and test report. + # --level 3 overrides the --level=1 in pytest.ini's addopts so that + # coverage reflects the whole suite. At level 1 the level-2 and + # level-3 tests are deselected and contribute nothing, which + # understates coverage. + pytest --level 3 --cov=mesohops --cov-report=xml --cov-report=term --junitxml=test-results.xml -v || true # Show test results summary echo "=== Test Results Summary ===" @@ -53,13 +62,16 @@ jobs: - name: Upload coverage to Codecov if: matrix.python-version == '3.12' - uses: codecov/codecov-action@v3 + uses: codecov/codecov-action@v5 with: token: ${{ secrets.CODECOV_TOKEN }} - file: ./coverage.xml + files: ./coverage.xml flags: unittests name: Python Coverage - fail_ci_if_error: false + # Fail loudly. Previously this was false, so a rejected upload (missing + # token, sunset uploader) left the job green while Codecov received + # nothing -- the failure was invisible for several releases. + fail_ci_if_error: true verbose: true - name: Publish Test Results diff --git a/.gitignore b/.gitignore index 03b85fd..9b291d8 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,198 @@ -*.pyc +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# UV +# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +#uv.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock +#poetry.toml + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +#pdm.lock +# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it +# in version control. +# https://pdm.fming.dev/latest/usage/project/#working-with-version-control +.pdm.toml +.pdm-python +.pdm-build/ + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm / JetBrains +.idea/ + +# Abstra +# Abstra is an AI-powered process automation framework. +# Ignore directories containing user credentials, local state, and settings. +# Learn more at https://abstra.io/docs +.abstra/ + +# Visual Studio Code +# Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore +# that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore +# and can be added to the global gitignore or merged into this file. However, if you prefer, +# you could uncomment the following to ignore the entire vscode folder +# .vscode/ + +# Ruff stuff: +.ruff_cache/ + +# PyPI configuration file +.pypirc + +# Cursor +# Cursor is an AI-powered code editor. `.cursorignore` specifies files/directories to +# exclude from AI features like autocomplete and code analysis. Recommended for sensitive data +# refer to https://docs.cursor.com/context/ignore-files +.cursorignore +.cursorindexingignore .DS_Store -mesohops.egg-info -.idea -.idea + +# AI coding tools +.claude/ +.private-journal/ +.zencoder/ +.tldr/ diff --git a/README.md b/README.md index b3eb159..13aef97 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,5 @@ +[![codecov](https://codecov.io/gh/MesoscienceLab/mesohops/graph/badge.svg?token=EZAW5BW4P4)](https://codecov.io/gh/MesoscienceLab/mesohops) + # What is MesoHOPS? MesoHOPS is a Python library for running simulations with the Hierarchy of Pure States (HOPS), a formally exact trajectory-based approach for solving the time-evolution of open quantum systems coupled to non-Markovian thermal environments. The main feature of MesoHOPS is the implementation of adaptive HOPS (adHOPS), an extension of the HOPS formalism that leverages the dynamic localization of excitations to construct an adaptive basis. The moving adHOPS basis significantly reduces the computational cost of simulations and exhibits a size-invariant scaling in large systems. diff --git a/pyproject.toml b/pyproject.toml index 6acd42a..58c1303 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "mesohops" -version = "1.7.0" +version = "1.8.0" # numba and numpy are pinned to a single minor version. numba minor # releases have introduced stochastic test failures and memory issues, # and numba binds against a specific numpy ABI. Upper-bound bumps for diff --git a/src/mesohops/basis/hops_hierarchy.py b/src/mesohops/basis/hops_hierarchy.py index b901717..e4d90c9 100644 --- a/src/mesohops/basis/hops_hierarchy.py +++ b/src/mesohops/basis/hops_hierarchy.py @@ -1,4 +1,5 @@ import itertools as it +import warnings from collections import Counter import numpy as np @@ -13,12 +14,18 @@ __author__ = "D. I. G. Bennett, L. Varvelo, J. K. Lynd, B. Z. Citty" __version__ = "1.6" -HIERARCHY_DICT_DEFAULT = {"MAXHIER": int(3), "TERMINATOR": False, "STATIC_FILTERS": []} +HIERARCHY_DICT_DEFAULT = { + "MAXHIER": int(3), + "TERMINATOR": False, + "STATIC_FILTERS": [], + "TRUNCATION_METHOD": 'triangular', +} HIERARCHY_DICT_TYPES = dict( MAXHIER=[type(int())], TERMINATOR=[type(False), type(str())], STATIC_FILTERS=[type([])], + TRUNCATION_METHOD=[type(str())], ) @@ -67,23 +74,26 @@ def __init__(self, hierarchy_param, system_param): Inputs ------ - 1. hierarchy_param : + 1. hierarchy_param: [see hops_basis.py] - a. MAXHIER : int + a. MAXHIER: int Maximum depth in the hierarchy that will be kept in the calculation (options: >= 0). - b. TERMINATOR : bool + b. TERMINATOR: bool True indicates the terminator condition is used while False indicates otherwise (options: False). - c. STATIC_FILTERS : str + c. STATIC_FILTERS: str Total set of nodes defined by MAXHIER can be further filtered. This is a list of filters [(filter_name1, [filter_param_1]), ...] (options: Triangular, Markovian, LongEdge, Domain). + d. TRUNCATION_METHOD: str + Hierarchy truncation scheme + (options: 'triangular', 'rectangular'). - 2. system_param : + 2. system_param: [see hops_system.py] - a. N_HMODES : int + a. N_HMODES: int Number of modes that appear in hierarchy. Returns @@ -100,9 +110,9 @@ def __init__(self, hierarchy_param, system_param): self.param = hierarchy_param if self.param["MAXHIER"] > 255: - print("Warning: using a hierarchy depth greater than 255 can cause " - "integer overflow issues when calculating derivative error. " - "Resetting hierarchy depth to 255.") + warnings.warn("Using a hierarchy depth greater than 255 can cause " + "integer overflow issues when calculating derivative error. " + "Resetting hierarchy depth to 255.") self.param["MAXHIER"] = 255 self.n_hmodes = system_param["N_HMODES"] self._auxiliary_list = [] @@ -149,7 +159,7 @@ def initialize(self, flag_adaptive): Parameters ---------- - 1. flag_adaptive : bool + 1. flag_adaptive: bool True indicates an adaptive calculation while False indicates otherwise. @@ -159,32 +169,48 @@ def initialize(self, flag_adaptive): """ # Prepare the hierarchy # --------------------- - # The hierarchy is only predefined if the basis is not adaptive + trunc = self.param["TRUNCATION_METHOD"] + # Vector adaptive HOPS uses flux filters (hops_fluxfilters.py) that + # enforce the MAXHIER boundary via aux._sum (total depth), which is + # inherently triangular. Rectangular support would require per-mode + # depth checks in the filter logic. + # NOTE: Tensor HOPS bypasses hierarchy.initialize() entirely + # (hierarchy depth is encoded in MPS bond dimensions), so this + # guard only affects the vector path. + if flag_adaptive and trunc == 'rectangular': + raise NotImplementedError( + 'Rectangular truncation is not yet supported for adaptive ' + 'calculations.' + ) if not flag_adaptive: - # If there are no static filters, use the standard triangular hierarchy - # generator - if len(self.param["STATIC_FILTERS"]) == 0: - self.auxiliary_list = self.filter_aux_list( - self.define_triangular_hierarchy(self.n_hmodes, - self.param["MAXHIER"] - ) - ) - # If the first static filter is not Markovian, use the standard - # triangular hierarchy generator and then apply filters - elif not "Markovian" in self.param["STATIC_FILTERS"][0]: + if trunc == 'rectangular': + if len(self.param["STATIC_FILTERS"]) > 0: + warnings.warn( + 'STATIC_FILTERS with rectangular truncation: filters ' + 'will be applied, but note that no tensor-specific ' + 'filters exist yet.' + ) self.auxiliary_list = self.filter_aux_list( - self.define_triangular_hierarchy(self.n_hmodes, - self.param["MAXHIER"] - ) + self.define_rectangular_hierarchy( + self.n_hmodes, self.param["MAXHIER"] + ) ) - # If the first static filter is Markovian, then use the Markovian - # triangular hierarchy generator + elif trunc == 'triangular': + # Triangular truncation with optional static filters + if "Markovian" in (self.param["STATIC_FILTERS"] or [[]])[0]: + list_mark = self.param["STATIC_FILTERS"][0][1] + self.auxiliary_list = self.filter_aux_list( + self.define_markovian_filtered_triangular_hierarchy( + self.n_hmodes, self.param["MAXHIER"], list_mark) + ) + else: + self.auxiliary_list = self.filter_aux_list( + self.define_triangular_hierarchy( + self.n_hmodes, self.param["MAXHIER"] + ) + ) else: - list_mark = self.param["STATIC_FILTERS"][0][1] - self.auxiliary_list = self.filter_aux_list( - self.define_markovian_filtered_triangular_hierarchy( - self.n_hmodes, self.param["MAXHIER"], list_mark) - ) + raise UnsupportedRequest(trunc, 'TRUNCATION_METHOD') else: # Initialize Guess for the hierarchy @@ -200,6 +226,8 @@ def initialize(self, flag_adaptive): self.only_markovian_filter = False else: self.only_markovian_filter = True + + def filter_aux_list(self, list_aux): """ @@ -208,12 +236,12 @@ def filter_aux_list(self, list_aux): Parameters ---------- - 1. list_aux : list(instance(AuxVec)) + 1. list_aux: list(instance(AuxVec)) List of auxiliaries to be filtered. Returns ------- - 1. list_aux : list(instance(AuxVec)) + 1. list_aux: list(instance(AuxVec)) Filtered list of auxiliaries. """ @@ -230,18 +258,18 @@ def apply_filter(self, list_aux, filter_name, params): Parameters ---------- - 1. list_aux : list(instance(AuxVec)) + 1. list_aux: list(instance(AuxVec)) List of auxiliaries that needs to be filtered. - 2. filter_name : str + 2. filter_name: str Name of filter. - 3. params : list + 3. params: list List of parameters for the filter. Returns ------- - 1. list_aux : list(instance(AuxVec)) + 1. list_aux: list(instance(AuxVec)) List of filtered auxiliaries. @@ -302,11 +330,11 @@ def _aux_index(self, aux): Parameters ---------- - 1. aux : instance(AuxVec) + 1. aux: instance(AuxVec) Returns ------- - 1. aux_index : int + 1. aux_index: int Relative or absolute index of a single auxiliary. """ if aux._index is None: @@ -315,27 +343,41 @@ def _aux_index(self, aux): return aux._index @staticmethod - def _const_aux_edge(absindex_mode, depth, n_hmodes): + def define_rectangular_hierarchy(n_hmodes, maxhier): """ - Creates an auxiliary object for an edge node at - a particular depth along a given mode. + Creates a rectangular hierarchy for a given number of modes at a + given depth. Parameters ---------- - 1. absindex_mode : int - Absolute index of the edge mode. + 1. n_hmodes: int + Number of modes that appear in the hierarchy. - 2. depth : int - Depth of the edge auxiliary. - - 3. n_hmodes : int - Number of modes that appear in the hierarchy. + 2. maxhier: int + Maximum hierarchy depth per mode. Returns ------- - 1. aux : instance(AuxVec) + 1. list_aux: list(AuxVec) + List of auxiliaries in the new rectangular hierarchy. """ - return AuxVec([(absindex_mode, depth)], n_hmodes) + list_aux = [] + # Iterate over all Cartesian products of hierarchy depths + # [0, 1, ..., maxhier] across n_hmodes modes. Each tuple + # hier_depths is a vector k = (k_0, k_1, ..., k_{M-1}) where + # each k_i ∈ {0, ..., maxhier}, giving (maxhier+1)^n_hmodes + # auxiliary states in total. + iter_hier_depth = it.product(range(maxhier + 1), repeat=n_hmodes) + for hier_depths in iter_hier_depth: + # Build the sparse (mode_index, depth) representation for + # this auxiliary state. AuxVec stores only nonzero entries, + # so modes with depth 0 are omitted. + list_aux_input = [(idx_mode, depth) for idx_mode, depth + in enumerate(hier_depths) if depth > 0] + list_aux.append( + AuxVec(list_aux_input, n_hmodes) + ) + return list_aux @staticmethod def define_triangular_hierarchy(n_hmodes, maxhier): @@ -345,17 +387,19 @@ def define_triangular_hierarchy(n_hmodes, maxhier): Parameters ---------- - 1. n_hmodes : int + 1. n_hmodes: int Number of modes that appear in the hierarchy. - 2. maxhier : int + 2. maxhier: int Max single value of the hierarchy. Returns ------- - 1. list_aux : list(instance(AuxVec)) + 1. list_aux: list(instance(AuxVec)) List of auxiliaries in the new triangular hierarchy. """ + + list_aux = [] # first loop over hierarchy depth for k in range(maxhier + 1): @@ -377,19 +421,19 @@ def define_markovian_filtered_triangular_hierarchy(n_hmodes, maxhier, Parameters ---------- - 1. n_hmodes : int + 1. n_hmodes: int Number of modes that appear in the hierarchy. - 2. maxhier : int + 2. maxhier: int Max single value of the hierarchy. - 3. list_boolean_mark : list(bool) + 3. list_boolean_mark: list(bool) List by modes. True indicates that the Markovian filter will be applied while False indicates otherwise. Returns ------- - 1. list_aux : list(instance(AuxVec)) + 1. list_aux: list(instance(AuxVec)) List of auxiliaries in the new triangular hierarchy. """ list_not_boolean_mark = [not bool_mark for bool_mark in list_boolean_mark] @@ -423,9 +467,9 @@ def __update_count(self, aux, type): Parameters ---------- - 1. aux : instance(AuxVec) + 1. aux: instance(AuxVec) - 2. type : str + 2. type: str Determines whether to add or remove (options: add, remove). Returns @@ -462,7 +506,7 @@ def __update_modes_in_use(self): elif value > 0: list_modes_in_use.append(mode) else: - print(f'ERROR: _count_by_modes is negative for mode {mode}') + raise ValueError(f'_count_by_modes is negative for mode {mode}') [self._count_by_modes.pop(mode) for mode in list_to_remove] self._list_modes_in_use = sorted(list_modes_in_use) @@ -488,24 +532,23 @@ def add_connections(self): for aux in self.list_aux_add: sum_aux = np.sum(aux) # add connections to k+1 - if sum_aux < self.param['MAXHIER']: - list_id_p1, list_value_connects_p1, list_mode_connects_p1 = aux.get_list_id_up( - self._list_modes_in_use) - for (rel_ind, my_id) in enumerate(list_id_p1): - try: - aux_p1 = self._dict_aux_by_id[my_id] - aux.add_aux_connect(list_mode_connects_p1[rel_ind], aux_p1, 1) - - # We simply keep track of the index connections of new auxiliaries - # Note: It does not matter that the indices will change because we only use this dictionary - # once, immediately after it is created in eom.ksuper - self._new_aux_id_conn_by_mode[list_mode_connects_p1[rel_ind]][ - aux.id] = [aux_p1.id, list_value_connects_p1[rel_ind] + 1] - self._new_aux_index_conn_by_mode[ - list_mode_connects_p1[rel_ind]][aux._index] = [ - aux_p1._index, list_value_connects_p1[rel_ind] + 1] - except: - pass + list_id_p1, list_value_connects_p1, list_mode_connects_p1 = aux.get_list_id_up( + self._list_modes_in_use) + for (rel_ind, my_id) in enumerate(list_id_p1): + try: + aux_p1 = self._dict_aux_by_id[my_id] + aux.add_aux_connect(list_mode_connects_p1[rel_ind], aux_p1, 1) + + # We simply keep track of the index connections of new auxiliaries + # Note: It does not matter that the indices will change because we only use this dictionary + # once, immediately after it is created in eom.ksuper + self._new_aux_id_conn_by_mode[list_mode_connects_p1[rel_ind]][ + aux.id] = [aux_p1.id, list_value_connects_p1[rel_ind] + 1] + self._new_aux_index_conn_by_mode[ + list_mode_connects_p1[rel_ind]][aux._index] = [ + aux_p1._index, list_value_connects_p1[rel_ind] + 1] + except: + pass # Add connections to k-1 if sum_aux > 0: diff --git a/src/mesohops/basis/hops_modes.py b/src/mesohops/basis/hops_modes.py index 98c2fd9..d2292d0 100644 --- a/src/mesohops/basis/hops_modes.py +++ b/src/mesohops/basis/hops_modes.py @@ -48,7 +48,6 @@ def __init__(self, system, hierarchy): self._list_modeidx_abs = [] self._list_l2idx_abs = [] - @property def list_index_L2_by_hmode(self): return self._list_index_L2_by_hmode diff --git a/src/mesohops/basis/hops_system.py b/src/mesohops/basis/hops_system.py index 8badd64..797138f 100644 --- a/src/mesohops/basis/hops_system.py +++ b/src/mesohops/basis/hops_system.py @@ -48,6 +48,7 @@ class HopsSystem: '_list_activel2idx_abs', # Active L2 indices (absolute) '__list_destination_state', # Destination states for each state '__dict_relindex_states', # Relative state indices + 'flag_nearest_neighbor_ham', # True if Hamiltonian has only nearest-neighbor coupling ) def __init__(self, system_param: dict[str, Any] | str | os.PathLike[str] | Path) -> None: @@ -151,6 +152,9 @@ def __init__(self, system_param: dict[str, Any] | str | os.PathLike[str] | Path) self._dict_nzhamiltonian_abs[key] += data else: self._dict_nzhamiltonian_abs[key] = data + self.flag_nearest_neighbor_ham = self._is_nearest_neighbor( + self.param['SPARSE_HAMILTONIAN'] + ) def initialize(self, flag_adaptive: bool, psi_0: np.ndarray) -> None: """ @@ -456,3 +460,26 @@ def reduce_sparse_matrix( return sp.sparse.coo_matrix( (data, (row, col)), shape=(len(iter_states), len(iter_states)) ) + + @staticmethod + def _is_nearest_neighbor(H2_ham_sparse: sp.sparse.spmatrix) -> bool: + """ + Check if a Hamiltonian has only nearest-neighbor coupling. + + Returns True if all non-zero elements satisfy |row - col| <= 1. + + Parameters + ---------- + 1. H2_ham_sparse: sparse matrix + System Hamiltonian in any scipy sparse format. + + Returns + ------- + 1. flag_nearest_neighbor_ham: bool + True if nearest-neighbor, False otherwise. + """ + H2_ham_coo = H2_ham_sparse.tocoo() + # Eliminate explicit zeros so that a user-supplied sparse matrix + # with stored zeros on far off-diagonals is not misclassified. + H2_ham_coo.eliminate_zeros() + return bool(np.all(np.abs(H2_ham_coo.row - H2_ham_coo.col) <= 1)) diff --git a/src/mesohops/eom/hops_eom.py b/src/mesohops/eom/hops_eom.py index 6bdb59b..bcb3d03 100644 --- a/src/mesohops/eom/hops_eom.py +++ b/src/mesohops/eom/hops_eom.py @@ -104,12 +104,12 @@ def __init__(self, eom_params): self.normalized = False elif self.param["EQUATION_OF_MOTION"] == "LINEAR": self.normalized = False + else: raise UnsupportedRequest( "EQUATION_OF_MOTION =" + self.param["EQUATION_OF_MOTION"], type(self).__name__, ) - # Checks adaptive definition # ------------------------- if self.param["ADAPTIVE_H"] or self.param["ADAPTIVE_S"]: @@ -375,7 +375,6 @@ def dsystem_dt( 2. z_mem1_deriv : np.array(complex) Derivative of z_mem with respect to time. """ - # Construct noise terms # --------------------- z_hat1_tmp = (np.conj(z_rnd1_tmp) + compress_zmem( @@ -400,7 +399,6 @@ def dsystem_dt( list_g, list_w, ) - # Check for a low-temperature correction stemming from flux from # Markovian auxiliaries C2_gamma_LT_corr_to_norm_corr = 0 @@ -427,19 +425,19 @@ def dsystem_dt( np.array(list_lt_corr_param), list_avg_L2, list_avg_L2_sq - ) - + ) # Calculates dphi/dt # ----------------- Φ_view_F = np.asarray(Φ).reshape([system.size,hierarchy.size],order="F") Φ_view_C = np.asarray(Φ).reshape([hierarchy.size,system.size],order="C") - Φ_deriv = K2_stable @ Φ + Φ_deriv += (self.K2_k @ Φ_view_C).reshape([hierarchy.size * system.size],order="C") Φ_deriv += ((-1j * system.hamiltonian) @ Φ_view_F).reshape([system.size * hierarchy.size],order="F") - + + Φ_deriv_view_F = np.asarray(Φ_deriv).reshape([system.size,hierarchy.size],order="F") Φ_deriv_view_C = np.asarray(Φ_deriv).reshape([hierarchy.size,system.size],order="C") @@ -449,7 +447,7 @@ def dsystem_dt( Φ_deriv[:system.size] += C2_LT_corr_physical @ np.asarray( Φ[:system.size]) norm_corr += C2_gamma_LT_corr_to_norm_corr - + if self.normalized: Φ_deriv -= norm_corr * Φ @@ -464,13 +462,12 @@ def dsystem_dt( (z_hat1_tmp[j] - 1.0j * z_tmp2[j]) * (list_l2_nz_csr[rel_index] @ Φ_view_red) ) - + Φ_view_red = Φ_view_C[list_hier_mask_Zp1[rel_index][1],:] Z2_kp1_red = Z2_kp1[rel_index][list_hier_mask_Zp1[rel_index][2]] Φ_deriv_view_C[list_hier_mask_Zp1[rel_index][0],:] += np.conj(list_avg_L2[j]) * (Z2_kp1_red @ Φ_view_red) - # Calculates dz/dt # --------------- @@ -485,7 +482,7 @@ def dsystem_dt( list_l2idx_abs, system.list_activel2idx_abs ) - + return Φ_deriv, z_mem1_deriv elif self.param["EQUATION_OF_MOTION"] == "LINEAR": diff --git a/src/mesohops/integrator/integrator_rk.py b/src/mesohops/integrator/integrator.py similarity index 74% rename from src/mesohops/integrator/integrator_rk.py rename to src/mesohops/integrator/integrator.py index 9ae76c3..dc3ba6d 100644 --- a/src/mesohops/integrator/integrator_rk.py +++ b/src/mesohops/integrator/integrator.py @@ -1,13 +1,42 @@ +""" +Time-integration routines for vector-based HOPS. + +Each integrator operates on the flat hierarchy vector (``phi``) used by +``HopsTrajectory``. + +Functions +--------- +runge_kutta_step(dsystem_dt, phi, z_mem, z_rnd, z_rnd2, tau) + Classic RK4 step for the flat hierarchy vector. + +runge_kutta_variables(phi, z_mem, t, noise, noise2, tau, storage, ...) + Gathers noise samples at the three RK4 time points and returns a dict + ready to unpack into runge_kutta_step. +""" +from __future__ import annotations + import copy +from collections.abc import Callable + import numpy as np + +from mesohops.noise.hops_noise import HopsNoise +from mesohops.storage.hops_storage import HopsStorage from mesohops.util.physical_constants import hbar -__title__ = "Integrators, Runge-Kutta" -__author__ = "D. I. G. Bennett" -__version__ = "1.2" +__title__ = 'Integrators' +__author__ = 'D. I. G. Bennett, B. Z. Citty' +__version__ = '1.6' -def runge_kutta_step(dsystem_dt, phi, z_mem, z_rnd, z_rnd2, tau): +def runge_kutta_step( + dsystem_dt: Callable, + phi: np.ndarray, + z_mem: np.ndarray, + z_rnd: np.ndarray, + z_rnd2: np.ndarray, + tau: float, +) -> tuple[np.ndarray, np.ndarray]: """ Performs a single Runge-Kutta step from the current time to a time tau forward. Parameters @@ -53,7 +82,6 @@ def runge_kutta_step(dsystem_dt, phi, z_mem, z_rnd, z_rnd2, tau): else: z_mem_tmp = z_mem + c_rk[i] * kz[i - 1] * tau / hbar phi_tmp = phi + c_rk[i] * k[i - 1] * tau / hbar - # Calculate system derivatives k[i], kz[i] = dsystem_dt( phi_tmp, z_mem_tmp, z_rnd[:, i_zrnd[i]], z_rnd2[:, i_zrnd[i]] @@ -66,8 +94,17 @@ def runge_kutta_step(dsystem_dt, phi, z_mem, z_rnd, z_rnd2, tau): return phi, z_mem -def runge_kutta_variables(phi,z_mem, t, noise, noise2, tau, storage, - list_l2idx_abs,effective_noise_integration=False): +def runge_kutta_variables( + phi: np.ndarray, + z_mem: np.ndarray, + t: float, + noise: HopsNoise, + noise2: HopsNoise, + tau: float, + storage: HopsStorage, + list_l2idx_abs: list[int], + effective_noise_integration: bool = False, +) -> dict: """ Accepts a storage and noise objects and returns the pre-requisite variables for a runge-kutta integration step in a list that can be unraveled to correctly feed @@ -84,7 +121,7 @@ def runge_kutta_variables(phi,z_mem, t, noise, noise2, tau, storage, 5. noise2 : instance(HopsNoise) 6. tau : float Integration time step [units: fs]. - + 7. storage : instance(HopsStorage) 8. effective_noise_integration: bool True indicates that the effective noise @@ -96,12 +133,18 @@ def runge_kutta_variables(phi,z_mem, t, noise, noise2, tau, storage, Dictionary of variables needed for Runge Kutta. """ if effective_noise_integration: + # Number of fine noise sub-steps per integration step tau_ratio = round(tau/noise.param["TAU"]) tau_ratio2 = round(tau / noise2.param["TAU"]) + # Sample noise at fine resolution over [t, t + 1.5*tau) z_rnd_raw = noise.get_noise([t + (i/tau_ratio)*tau for i in range(round(tau_ratio*1.5))],list_l2idx_abs) z_rnd2_raw = noise2.get_noise([t + (i / tau_ratio2) * tau for i in range(round(tau_ratio2 * 1.5))],list_l2idx_abs) + # Average fine noise into 3 bins matching RK4 time-points: + # bin 0: [t, t+tau/2) -> noise at t + # bin 1: [t+tau/2, t+tau) -> noise at t+tau/2 + # bin 2: [t+tau, t+1.5*tau) -> noise at t+tau z_rnd = np.array([np.mean(z_rnd_raw[:,:round(tau_ratio/2)], axis=1), np.mean(z_rnd_raw[:,round(tau_ratio/2):tau_ratio], axis=1), np.mean(z_rnd_raw[:, tau_ratio:], axis=1)]).T @@ -113,5 +156,5 @@ def runge_kutta_variables(phi,z_mem, t, noise, noise2, tau, storage, else: z_rnd = noise.get_noise([t, t + tau * 0.5, t + tau],list_l2idx_abs) z_rnd2 = noise2.get_noise([t, t + tau * 0.5, t + tau],list_l2idx_abs) - + return {"phi": phi, "z_mem": z_mem, "z_rnd": z_rnd, "z_rnd2": z_rnd2, "tau": tau} diff --git a/src/mesohops/integrator/tensor_integrator.py b/src/mesohops/integrator/tensor_integrator.py new file mode 100644 index 0000000..d420cc7 --- /dev/null +++ b/src/mesohops/integrator/tensor_integrator.py @@ -0,0 +1,414 @@ +""" +Time-integration routines for tensor-based HOPS. + +Each integrator operates on a ``HopsTensorEOM`` instance (which wraps the MPS) +used by ``HopsTensorTrajectory``. + +Functions +--------- +runge_kutta_step_tensor(eom, z_mem, z_rnd, z_rnd2, tau) + RK4 step for the MPS wavefunction, using tensor_add at each stage. + +runge_kutta_variables(z_mem, t, noise, noise2, tau, ...) + Gathers noise samples at the three RK4 time points and returns a dict + ready to unpack into runge_kutta_step_tensor. + +single_point_variables(z_mem, t, noise, noise2, tau, ...) + Gathers noise at time t and returns a dict ready to unpack into + a single-point-noise integrator (used by TDVP steps). + +tdvp_step_tensor(eom, z_mem, z_rnd, z_rnd2, tau, method, **kwargs) + TDVP step (1-site or 2-site) for the MPS wavefunction. +""" +from __future__ import annotations + +import numpy as np + +from mesohops.noise.hops_noise import HopsNoise +from mesohops.tensor import tdvp +from mesohops.tensor.hops_tensor_eom import HopsTensorEOM +from mesohops.util.physical_constants import hbar +from mesohops.util.tensor_operations import scale_mps, tensor_add + +__title__ = 'Tensor Integrators' +__author__ = 'D. I. G. Bennett, B. Z. Citty' +__maintainer__ = 'B. Z. Citty' + + +def runge_kutta_step_tensor( + eom: HopsTensorEOM, + z_mem: np.ndarray, + z_rnd: np.ndarray, + z_rnd2: np.ndarray, + tau: float, +) -> tuple[np.ndarray, int]: + """ + Performs a single Runge-Kutta step for the MPS wavefunction. + + Parameters + ---------- + 1. eom: HopsTensorEOM + Equation of motion object wrapping the MPS and MPO. + 2. z_mem: np.ndarray(complex) + Current memory term values. + 3. z_rnd: np.ndarray(complex) + Primary noise at three time points. + 4. z_rnd2: np.ndarray(complex) + Secondary noise at three time points. + 5. tau: float + Integration time step [units: fs]. + + Returns + ------- + 1. z_mem: np.ndarray(complex) + Updated memory terms. + + Side effects + ------------ + Sets `eom.max_complexity_step` to the largest per-call tensor + complexity observed across the four RK4 matvec-then-compress + evaluations of this step. Read via the side-channel by propagate(); + avoids plumbing the scalar through every return value. + """ + wavefunction = eom.wavefunction + list_k_phi = [[] for _ in range(4)] # RK4 stages for dphi/dt + list_k_zmem = [[] for _ in range(4)] # RK4 stages for dz_mem/dt + # RK4 stage time-offsets: evaluate at t, t+tau/2, t+tau/2, t+tau + list_rk_coeff = [0.0, 0.5, 0.5, 1.0] + # Map each RK4 stage to a noise time-point index: 0=t, 1=t+tau/2, 2=t+tau + list_noise_idx = [0, 1, 1, 2] + # Peak uncompressed-core complexity across the 4 RK4 matvec calls; + # tracked locally and published on eom.max_complexity_step at the end. + max_complexity = 0 + + # Deep-copy the initial MPS state as the RK4 checkpoint. + # Each group is copied element-wise for statenumber (nested list-of-lists); + # a single list comprehension handles both nested and flat structures. + list_cores_phi_checkpoint = [ + [c.copy() for c in g] if isinstance(g, list) else g.copy() + for g in wavefunction.list_cores_phi + ] + + for i in range(4): + if i == 0: + z_mem_tmp = z_mem + else: + z_mem_tmp = z_mem + list_rk_coeff[i] * list_k_zmem[i - 1] * tau / hbar + # Scale k[i-1] by the RK4 prefactor, then unscale after. + # Scaling only the first core is equivalent to scaling the whole + # MPS; tensor_add distributes the weight through the boundary core. + scale_mps(list_k_phi[i - 1], list_rk_coeff[i] * tau / hbar) + wavefunction.list_cores_phi = tensor_add( + list_cores_phi_checkpoint, + list_k_phi[i - 1], + wavefunction.mps_epsilon, + wavefunction.bond_dim_max, + ) + # Restore k[i-1] to its unscaled derivative form so the + # final RK4 weighted sum (below) starts from raw k values. + scale_mps(list_k_phi[i - 1], hbar / (list_rk_coeff[i] * tau)) + # Build MPO and get dz; then get d(phi)/dt cores + list_k_zmem[i] = eom.build_generator( + z_mem_tmp, z_rnd[:, list_noise_idx[i]], z_rnd2[:, list_noise_idx[i]] + ) + list_k_phi[i] = eom.derivative() + # eom.last_matvec_complexity was set by the derivative() call. + if eom.last_matvec_complexity > max_complexity: + max_complexity = eom.last_matvec_complexity + + wavefunction.list_cores_phi = list_cores_phi_checkpoint + + # RK4 weighted sum: scale each derivative by its RK4 weight and + # the tau/hbar factor. Scaling only the first core is equivalent + # to scaling the whole MPS; tensor_add absorbs the weight into the + # boundary core. + scale_mps(list_k_phi[0], tau / (6 * hbar)) + scale_mps(list_k_phi[1], tau / (3 * hbar)) + scale_mps(list_k_phi[2], tau / (3 * hbar)) + scale_mps(list_k_phi[3], tau / (6 * hbar)) + + for k in list_k_phi: + wavefunction.list_cores_phi = tensor_add( + wavefunction.list_cores_phi, + k, + wavefunction.mps_epsilon, + wavefunction.bond_dim_max, + ) + z_mem = ( + z_mem + + tau + / hbar + * ( + list_k_zmem[0] + + 2.0 * list_k_zmem[1] + + 2.0 * list_k_zmem[2] + + list_k_zmem[3] + ) + / 6.0 + ) + + eom.max_complexity_step = max_complexity + return z_mem + + +def runge_kutta_variables( + z_mem: np.ndarray, + t: float, + noise: HopsNoise, + noise2: HopsNoise, + tau: float, + list_l2idx_abs: list[int], + effective_noise_integration: bool = False, +) -> dict: + """ + Accepts noise objects and returns the variables needed for a + Runge-Kutta integration step. + + Parameters + ---------- + 1. z_mem: np.ndarray(complex) + Memory terms [units: cm^-1]. + 2. t: float + Integration time point [units: fs]. + 3. noise: instance(HopsNoise) + Primary noise generator. + 4. noise2: instance(HopsNoise) + Secondary noise generator. + 5. tau: float + Integration time step [units: fs]. + 6. list_l2idx_abs: list(int) + Absolute L2 operator indices for noise + sampling. + 7. effective_noise_integration: bool + True uses moving average over + noise; False uses point samples. + + Returns + ------- + 1. variables: dict + Dictionary of variables needed for Runge-Kutta. + """ + if effective_noise_integration: + # Number of fine noise sub-steps per integration step + tau_ratio = round(tau / noise.param['TAU']) + tau_ratio2 = round(tau / noise2.param['TAU']) + # Sample noise at fine resolution over [t, t + 1.5*tau) + z_rnd_raw = noise.get_noise( + [t + (i / tau_ratio) * tau for i in range(round(tau_ratio * 1.5))], + list_l2idx_abs, + ) + z_rnd2_raw = noise2.get_noise( + [t + (i / tau_ratio2) * tau for i in range(round(tau_ratio2 * 1.5))], + list_l2idx_abs, + ) + # Average fine noise into 3 bins matching RK4 time-points: + # bin 0: [t, t+tau/2) -> noise at t + # bin 1: [t+tau/2, t+tau) -> noise at t+tau/2 + # bin 2: [t+tau, t+1.5*tau) -> noise at t+tau + z_rnd = np.array( + [ + np.mean(z_rnd_raw[:, : round(tau_ratio / 2)], axis=1), + np.mean(z_rnd_raw[:, round(tau_ratio / 2) : tau_ratio], axis=1), + np.mean(z_rnd_raw[:, tau_ratio:], axis=1), + ] + ).T + z_rnd2 = np.array( + [ + np.mean(z_rnd2_raw[:, : round(tau_ratio2 / 2)], axis=1), + np.mean(z_rnd2_raw[:, round(tau_ratio2 / 2) : tau_ratio2], axis=1), + np.mean(z_rnd2_raw[:, tau_ratio2:], axis=1), + ] + ).T + + else: + z_rnd = noise.get_noise([t, t + tau * 0.5, t + tau], list_l2idx_abs) + z_rnd2 = noise2.get_noise([t, t + tau * 0.5, t + tau], list_l2idx_abs) + + return {'z_mem': z_mem, 'z_rnd': z_rnd, 'z_rnd2': z_rnd2, 'tau': tau} + + +def single_point_variables( + z_mem: np.ndarray, + t: float, + noise: HopsNoise, + noise2: HopsNoise, + tau: float, + list_l2idx_abs: list[int] | None = None, + effective_noise_integration: bool = False, +) -> dict: + """ + Accepts noise objects and returns the variables needed for a + single-point-noise integration step (used by TDVP). + + Parameters + ---------- + 1. z_mem: np.ndarray(complex) + Memory terms [units: cm^-1]. + 2. t: float + Integration time point [units: fs]. + 3. noise: instance(HopsNoise) + Primary noise generator. + 4. noise2: instance(HopsNoise) + Secondary noise generator. + 5. tau: float + Integration time step [units: fs]. + 6. list_l2idx_abs: list(int) + Absolute L2 operator indices passed to + noise.get_noise for adaptive filtering. + 7. effective_noise_integration: bool + True uses moving average over + noise; False uses point samples. + **Not yet implemented; raises + NotImplementedError.** + + Returns + ------- + 1. variables: dict + Dictionary of variables needed for a + single-point-noise integration step. + """ + if effective_noise_integration: + raise NotImplementedError( + 'effective_noise_integration is not implemented for single_point_variables.' + ) + z_rnd = noise.get_noise([t], list_l2idx_abs) + z_rnd2 = noise2.get_noise([t], list_l2idx_abs) + return {'z_mem': z_mem, 'z_rnd': z_rnd, 'z_rnd2': z_rnd2, 'tau': tau} + + +def _build_tdvp_solver_kwargs(update_type, krylov_conv_tol, **kwargs): + """ + Map trajectory-level TDVP parameters to tdvp.timestep kwargs. + + Parameters + ---------- + 1. update_type: str + TDVP solver type ('arnoldi', 'lanczos', 'krylov', or 'ivp'). + 2. krylov_conv_tol: float + Relative convergence tolerance for arnoldi/lanczos solvers. + 3. **kwargs: dict + Additional parameters (ivp_method, ivp_rtol, ivp_atol, + ivp_max_step) forwarded when update_type is 'ivp'. + + Returns + ------- + 1. solver_kwargs: dict + Keyword arguments for tdvp.timestep. + """ + # 'krylov' is a convenience alias for 'arnoldi' + solver = 'arnoldi' if update_type == 'krylov' else update_type + solver_kwargs = {'solver': solver} + if solver in ('arnoldi', 'lanczos'): + solver_kwargs['conv_tol'] = krylov_conv_tol + elif solver == 'ivp': + solver_kwargs['method'] = kwargs.get('ivp_method', 'BDF') + solver_kwargs['rtol'] = kwargs.get('ivp_rtol', 1e-7) + solver_kwargs['atol'] = kwargs.get('ivp_atol', 1e-9) + max_step = kwargs.get('ivp_max_step', None) + if max_step is not None: + solver_kwargs['max_step'] = max_step + return solver_kwargs + + +def tdvp_step_tensor( + eom, z_mem, z_rnd, z_rnd2, tau, method='1tdvp', + krylov_conv_tol=1e-6, update_type='krylov', **kwargs +): + """ + Performs a single TDVP step (1-site or 2-site) for the MPS wavefunction. + + Builds the MPO via eom.build_generator, then evolves the MPS using the + Strang-split TDVP sweep from tdvp.timestep. The memory terms z_mem are + updated with a first-order step. + + For 2TDVP, bond dimension is controlled by wavefunction.bond_dim_max and + wavefunction.mps_epsilon. + + Parameters + ---------- + 1. eom: HopsTensorEOM + Equation of motion object wrapping the MPS and MPO. + 2. z_mem: np.ndarray(complex) + Current memory term values. + 3. z_rnd: np.ndarray(complex) + Primary noise at this time point. + 4. z_rnd2: np.ndarray(complex) + Secondary noise at this time point. + 5. tau: float + Integration time step [units: fs]. + 6. method: str + TDVP variant ('1tdvp' or '2tdvp'). + 7. krylov_conv_tol: float + Relative convergence tolerance for the local exponential + solver. Iteration stops when the Hochbruck-Lubich + residual estimate drops below this value. + 8. update_type: str + Solver type ('arnoldi', 'lanczos', 'krylov', or 'ivp'). + 9. **kwargs: dict + Forwarded to the solver (e.g. ivp_method, ivp_rtol, ivp_atol). + + Returns + ------- + 1. z_mem: np.ndarray(complex) + Updated memory terms. + + Side effects + ------------ + Sets `eom.max_complexity_step = 0`. TDVP does not evaluate + tensor_matvec_prod, so the peak-size metric tracked by the RK4 + path does not apply; zero is published on the side-channel only + to keep the attribute populated for propagate's storage write. + """ + wavefunction = eom.wavefunction + + # (1) Build MPO and compute dz/dt. + dz_dt = eom.build_generator(z_mem, z_rnd[:, 0], z_rnd2[:, 0]) + + # (2) Initialize TDVP: right-normalize MPS and build environments. + # (initialize recenters the MPS to site 0 internally for stability) + core_M, list_cores_B, L0, list_envs_R = tdvp.initialize( + wavefunction.flat_cores, eom.mpo_cores + ) + + # (3) Strang-split TDVP sweep. + solver_kwargs = _build_tdvp_solver_kwargs(update_type, krylov_conv_tol, **kwargs) + if method == '2tdvp': + solver_kwargs['chi_max'] = wavefunction.bond_dim_max + solver_kwargs['eps'] = wavefunction.mps_epsilon + core_M, list_cores_B, L0, list_envs_R = tdvp.timestep( + tau / hbar, + L0, + list_envs_R, + eom.mpo_cores, + core_M, + list_cores_B, + method=method, + **solver_kwargs, + ) + + # (4) Write updated cores back into wavefunction. + wavefunction.update_phi_from_flat([core_M] + list(list_cores_B)) + + # (5) First-order update for memory terms. + z_mem = z_mem + tau / hbar * dz_dt + + # TDVP does not evaluate tensor_matvec_prod, so the peak-size metric + # tracked by runge_kutta_step_tensor does not apply. + eom.max_complexity_step = 0 + return z_mem + + +# Convenience aliases for backwards compatibility with tests +tdvp1_step_tensor = tdvp_step_tensor + + +def tdvp2_step_tensor( + eom, z_mem, z_rnd, z_rnd2, tau, krylov_conv_tol=1e-6, update_type='krylov', **kwargs +): + """ + 2TDVP step. Convenience wrapper around tdvp_step_tensor with method='2tdvp'. + """ + return tdvp_step_tensor( + eom, z_mem, z_rnd, z_rnd2, tau, method='2tdvp', + krylov_conv_tol=krylov_conv_tol, update_type=update_type, **kwargs + ) diff --git a/src/mesohops/storage/hops_storage.py b/src/mesohops/storage/hops_storage.py index 823a623..c8fbd7d 100644 --- a/src/mesohops/storage/hops_storage.py +++ b/src/mesohops/storage/hops_storage.py @@ -114,6 +114,9 @@ def adaptive(self, new): self.storage_dic.setdefault('psi_traj', True) self.storage_dic.setdefault('t_axis', True) self.storage_dic.setdefault('z_mem', False) + # Tensor-only: per-step <0,...,0|psi> for vacuum-convention + # spectroscopy (mirrors what gs_core gets via psi_traj[0]). + self.storage_dic.setdefault('psi_g_traj', False) if self._adaptive: self.storage_dic.setdefault('aux_list', True) diff --git a/src/mesohops/storage/storage_functions.py b/src/mesohops/storage/storage_functions.py index 2e88b36..1bb4082 100644 --- a/src/mesohops/storage/storage_functions.py +++ b/src/mesohops/storage/storage_functions.py @@ -1,5 +1,9 @@ +import copy + import numpy as np +from mesohops.util.tensor_operations import contract_down_exact, extract_gs_amp + __title__ = "storage functions" __author__ = "L. Varvelo, D. I. G. Bennett, J. K. Lynd" __version__ = "1.2" @@ -183,9 +187,107 @@ def save_list_zmemmodeidx_abs(list_zmemmodeidx_abs, **kwargs): return list_zmemmodeidx_abs +# ============================================================ +# Tensor-specific storage functions +# ============================================================ +# These are registered by HopsTensorTrajectory.__init__ to replace +# the default functions for keys that need MPS-aware handling. +# They receive a 'wavefunction' kwarg (HopsTensorWavefunction instance) +# passed by the tensor trajectory's store_step calls. + + +def save_phi_traj_tensor(wavefunction, **kwargs): + """ + Returns the full hierarchy state as MPS cores (deep copy). + + Parameters + ---------- + 1. wavefunction: HopsTensorWavefunction + Tensor wavefunction instance. + + Returns + ------- + 1. list_cores: list(np.ndarray) + Deep copy of the MPS cores. + """ + return copy.deepcopy(wavefunction.list_cores_phi) + + +def save_max_tensor_complexity(max_tensor_complexity, **kwargs): + """ + Returns the peak uncompressed-MPS complexity observed during this timestep. + + The scalar is the max across all tensor_matvec_prod calls in one + integration step (4 evaluations for RK4, 0 for TDVP). See + util.tensor_operations.calc_mps_complexity for the per-core formula. + + Parameters + ---------- + 1. max_tensor_complexity: int + Peak complexity for this step. + + Returns + ------- + 1. max_tensor_complexity: int + Same value, appended to the storage trace. + """ + return max_tensor_complexity + + +def save_phi_norm_tensor(wavefunction, **kwargs): + """ + Returns the L2-norm of the full hierarchy state via MPS contraction. + + Parameters + ---------- + 1. wavefunction: HopsTensorWavefunction + Tensor wavefunction instance. + + Returns + ------- + 1. phi_norm: float + L2-norm of the full hierarchy state. + """ + n_state = len(wavefunction.M1_modes_per_site) + V1_norm_sq = contract_down_exact( + wavefunction.list_cores_phi, + wavefunction.method, + n_state, + ) + return np.sqrt(np.sum(V1_norm_sq).real) + + +def save_psi_g_traj_tensor(wavefunction, **kwargs): + """ + Returns the amplitude of the all-zeros configuration of the MPS. + + Counterpart to save_psi_traj for the vacuum-convention spectroscopy + code path: under the vacuum convention extract_psi sees only the + excited-state slots, so the physical |g> amplitude (the all-zeros + configuration of the MPS) is not in psi_traj. This hook captures + it per storage step so detection-phase propagation can be a single + traj.propagate(t3_max, t_step) call rather than a per-step Python + loop calling extract_gs_amp. + + Parameters + ---------- + 1. wavefunction: HopsTensorWavefunction + Tensor wavefunction instance. + + Returns + ------- + 1. gs_amp: complex + <0,...,0|psi> at the current timestep. + """ + return extract_gs_amp( + wavefunction.list_cores_phi, wavefunction.method, + ) + + storage_default_func = {'psi_traj':save_psi_traj, 'phi_traj':save_phi_traj, 'phi_norm':save_phi_norm, 't_axis':save_t_axis, 'aux_list':save_aux_list, 'state_list':save_state_list, 'list_nstate':save_list_nstate, 'list_nhier':save_list_nhier, 'list_aux_norm':save_list_aux_norm, 'z_mem':save_z_mem, - 'list_zmemmodeidx_abs': save_list_zmemmodeidx_abs} + 'list_zmemmodeidx_abs': save_list_zmemmodeidx_abs, + 'psi_g_traj': save_psi_g_traj_tensor} diff --git a/src/conftest.py b/src/mesohops/tensor/__init__.py similarity index 100% rename from src/conftest.py rename to src/mesohops/tensor/__init__.py diff --git a/src/mesohops/tensor/hops_tensor_basis.py b/src/mesohops/tensor/hops_tensor_basis.py new file mode 100644 index 0000000..0869408 --- /dev/null +++ b/src/mesohops/tensor/hops_tensor_basis.py @@ -0,0 +1,234 @@ +from __future__ import annotations + +import numpy as np + +from mesohops.basis.hops_modes import HopsModes # noqa: F401 (type hint) +from mesohops.basis.hops_noise_memory import HopsNoiseMemory # noqa: F401 (type hint) +from mesohops.basis.hops_system import HopsSystem # noqa: F401 (type hint) +from mesohops.tensor.hops_tensor_eom import HopsTensorEOM # noqa: F401 (type hint) +from mesohops.tensor.hops_tensor_wavefunction import ( + HopsTensorWavefunction, # noqa: F401 (type hint) +) +from mesohops.tensor.tensor_functions_adaptive import ( + tensor_state_adaptive_check_add_state, + tensor_state_adaptive_check_remove_state, +) + +__title__ = 'TadHOPS Basis' +__author__ = 'B. Z. Citty' +__maintainer__ = 'B. Z. Citty' + + +class HopsTensorBasis: + """ + Holds shared references to the physical basis objects (system, mode, + noise_memory) owned by HopsBasis, plus the tensor-specific EOM, and + manages adaptive-basis logic for a tensor-network HOPS calculation. + + Hierarchy concepts (n_hier, n_hmodes, adaptive_h) do not apply — + tensor HOPS encodes hierarchy depth in MPS core dimensions (k_max + 1 + per mode), not through explicit auxiliary vector enumeration. + + HopsTensorWavefunction manages its own MPS core restructuring via + add_state_cores/remove_state_cores. HopsTensorBasis decides which + states to add/remove and updates the basis bookkeeping (system, mode). + """ + + __slots__ = ( + 'system', # HopsSystem instance + 'mode', # HopsModes instance + 'noise_memory', # HopsNoiseMemory instance + 'eom', # HopsTensorEOM instance (set by trajectory) + 'adaptive', # bool: True when state adaptivity is active + 'delta_s', # float: adaptive threshold (stored in initialize) + ) + + def __init__( + self, + system: HopsSystem, + mode: HopsModes, + noise_memory: HopsNoiseMemory, + ) -> None: + """ + Stores references to the shared basis objects owned by HopsBasis. + Does not construct HopsSystem, HopsModes, or HopsNoiseMemory — + those are owned by HopsBasis and passed in by HopsTensorTrajectory. + + Parameters + ---------- + 1. system : HopsSystem + Shared system instance (owned by HopsBasis). + + 2. mode : HopsModes + Shared mode instance (owned by HopsBasis). + + 3. noise_memory : HopsNoiseMemory + Shared noise memory instance (owned by HopsBasis). + + Returns + ------- + None + """ + self.system = system + self.mode = mode + self.noise_memory = noise_memory + self.eom = None + self.adaptive = False + self.delta_s = 0 + + def initialize(self, delta_s: float) -> None: + """ + Stores tensor-specific adaptive configuration. The shared objects + (system, mode, noise_memory) are already initialized by + HopsTensorTrajectory.initialize before this call. + + Parameters + ---------- + 1. delta_s : float + Adaptive threshold. delta_s > 0 enables state adaptivity. + + Returns + ------- + None + """ + self.delta_s = delta_s + self.adaptive = delta_s > 0 + + def define_basis( + self, + wavefunction: HopsTensorWavefunction, + z_step: np.ndarray, + ) -> tuple[list[int], list[int]]: + """ + Performs state adaptivity: computes errors for removing boundary states + and adding new ones. + + Receives wavefunction explicitly rather than reaching through self.eom. + The parent HopsBasis.define_basis takes (phi, tau, z_step) — tensor + omits tau because tensor HOPS does not do hierarchy adaptivity. + + Parameters + ---------- + 1. wavefunction : HopsTensorWavefunction + Tensor wavefunction to evaluate for adaptivity. + + 2. z_step : np.array(complex) + Noise values for the current time step. + + Returns + ------- + 1. list_states_old : list(int) + Absolute state indices to remove. + + 2. list_states_new : list(int) + Absolute state indices to add. + """ + if not self.adaptive: + return [], [] + + tensor_eom = self.eom + + def calc_deriv_cores(z_mem, z_rnd, z_rnd2): + # build_generator stores MPO in eom.mpo_cores (side effect). + # derivative() now stores its peak-size proxy on + # eom.last_matvec_complexity; the adaptive-basis path + # doesn't track max-across-substeps, so we just leave that + # attribute set and return the cores. + tensor_eom.build_generator(z_mem, z_rnd, z_rnd2) + return tensor_eom.derivative() + + old_state_indices = tensor_state_adaptive_check_remove_state( + wavefunction.list_cores_phi, + self.system.param['HAMILTONIAN'], + z_step, + self.system.param['NSTATES'], + len(self.system.state_list), + self.delta_s, + self.system.state_list, + wavefunction.method, + wavefunction.M1_modes_per_state, + calc_deriv_cores, + ) + list_states_old = [self.system.state_list[i] for i in old_state_indices] + list_states_new = tensor_state_adaptive_check_add_state( + wavefunction.list_cores_phi, + self.system.param['HAMILTONIAN'], + list_states_old, + self.system.param['NSTATES'], + len(self.system.state_list), + self.delta_s, + self.system.state_list, + wavefunction.method, + wavefunction.M1_modes_per_state, + ) + return list_states_old, list_states_new + + def update_basis( + self, + wavefunction: HopsTensorWavefunction, + z_mem: np.ndarray, + list_states_old: list[int], + list_states_new: list[int], + ) -> tuple[HopsTensorWavefunction, np.ndarray]: + """ + Removes and adds states in both the tensor cores and the basis objects, + then inflates bonds if TDVP is in use. + + wavefunction is mutated in place and returned for call-site visibility — + the return value is the same object as the argument. + + z_mem passes through unchanged today. + + TODO: z_mem remapping bug — when states are added/removed, + mode.list_modeidx_abs changes but z_mem is not remapped, causing + noise memory values to be read at wrong array positions on the next + step. The parent calls noise_memory.update_zmem_indexing(z_mem) and + constructs a new z_mem array — the tensor version must do the same. + + The parent returns (phi, z_mem, dsystem_dt) — tensor omits dsystem_dt + because tensor EOM rebuilds the MPO on the fly. After restructuring, + refreshes the persistent MpoBuilder's state-dependent attributes via + eom.refresh_builder. + + Parameters + ---------- + 1. wavefunction : HopsTensorWavefunction + Tensor wavefunction to update in-place. + + 2. z_mem : np.ndarray(complex) + Noise memory drift terms. Passed through unchanged (see + TODO above for the known remapping bug). + + 3. list_states_old : list(int) + Absolute state indices to remove. + + 4. list_states_new : list(int) + Absolute state indices to add. + + Returns + ------- + 1. wavefunction : HopsTensorWavefunction + The same object passed in, mutated in place. + + 2. z_mem : np.ndarray(complex) + Noise memory drift terms, passed through unchanged. + """ + # Remove old states + wavefunction.remove_state_cores(list_states_old, self.system.state_list) + self.system.state_list = sorted( + set(self.system.state_list) - set(list_states_old) + ) + self.mode.list_modeidx_abs = sorted(self.system.list_statemodeidx_abs) + + # Add new states + wavefunction.add_state_cores(list_states_new, self.system.state_list) + self.system.state_list = sorted( + set(self.system.state_list) | set(list_states_new) + ) + self.mode.list_modeidx_abs = sorted(self.system.list_statemodeidx_abs) + + # Refresh the persistent MpoBuilder with the new state data, + # matching the parent pattern where update_basis rebuilds dsystem_dt. + self.eom.refresh_builder() + + return wavefunction, z_mem diff --git a/src/mesohops/tensor/hops_tensor_eom.py b/src/mesohops/tensor/hops_tensor_eom.py new file mode 100644 index 0000000..3cb2bee --- /dev/null +++ b/src/mesohops/tensor/hops_tensor_eom.py @@ -0,0 +1,564 @@ +from __future__ import annotations + +import numpy as np +import scipy.sparse as sp + +from mesohops.basis.hops_modes import HopsModes +from mesohops.basis.hops_noise_memory import HopsNoiseMemory +from mesohops.basis.hops_system import HopsSystem +from mesohops.eom.eom_functions import ( + calc_delta_zmem, + calc_LT_corr, + calc_LT_corr_linear, + calc_LT_corr_to_norm_corr, + compress_zmem, + operator_expectation, +) +from mesohops.tensor.hops_tensor_wavefunction import HopsTensorWavefunction +from mesohops.tensor.mpo_constructors import ( + MpoBuilder, + build_statenumber_operator_mpo, +) +from mesohops.tensor.tensor_eom_functions import ( + build_physical_correction_mpo, + calc_norm_corr_tensor, + tensor_matvec_prod, +) +from mesohops.util.exceptions import UnsupportedRequest +from mesohops.util.tensor_operations import ( + scale_mps, + tensor_add, +) + +__title__ = 'Tensor HOPS EOM' +__author__ = 'B. Z. Citty' +__maintainer__ = 'B. Z. Citty' + + +class HopsTensorEOM: + """ + Builds and applies the time-evolution MPO for tensor HOPS. + + Sits between the MPS algebra (HopsTensorWavefunction) and the time-stepping + algorithms (tensor_integrator.py). The class owns the MPO storage and the logic for + computing the time-evolution operator and dz/dt at each integrator sub-step. + This two-phase design (build_generator then derivative) differs from HopsEOM, + which produces a single dsystem_dt derivative closure. The split is needed + because multi-step integrators (RK4) rebuild the MPO at each sub-step noise + point while reusing the derivative application. + """ + + __slots__ = ( + 'wavefunction', # MPS wavefunction container (HopsTensorWavefunction) + 'system', # System parameters and Hamiltonian (HopsSystem) + 'mode', # Bath mode indexing and coupling strengths (HopsModes) + 'noise_memory', # Noise memory drift terms (HopsNoiseMemory) + 'adaptive', # True when state adaptivity is active (bool) + 'mpo_cores', # Full operator MPO (list[np.ndarray]) + # Hierarchy MPO on the general number path, the combined Hamiltonian + + # hierarchy generator on a recognized topology, and the single MPO of + # the fullstate method. + '_list_cores_op', # Operator MPO cores (list[np.ndarray]) + # Separate Hamiltonian MPO, built only on the general number path. + # Empty everywhere else, which is how _construct_MPO knows whether + # there is a second MPO to add. + '_list_cores_ham', # Hamiltonian MPO cores (list[np.ndarray]) + 'normalization', # Auxiliary scaling convention, 'homps' or 'adhops' (str) + 'mpo_builder', # Persistent MPO builder (MpoBuilder) + 'flag_linear', # True when EQUATION_OF_MOTION is LINEAR (bool) + '_has_lt_corr', # True when LTC corrections should be applied (bool) + '_C2_LT_corr_phys', # LTC physical correction matrix (np.ndarray|None) + '_C2_LT_corr_hier', # LTC hierarchy correction matrix (np.ndarray|None) + # Diagnostic side-channel: complexity observed by the most recent + # derivative() call; max across an integrator step (reset by the + # step function, read by propagate). Avoids plumbing complexity + # through every return value in the derivative → step → _step + # → propagate chain. + 'last_matvec_complexity', # Complexity of most recent derivative() (int) + 'max_complexity_step', # Max complexity during the current step (int) + ) + + def __init__( + self, + wavefunction: HopsTensorWavefunction, + system: HopsSystem, + mode: HopsModes, + noise_memory: HopsNoiseMemory, + adaptive: bool, + eom_param: dict, + normalization: str = 'homps', + ) -> None: + """ + Inputs + ------ + 1. wavefunction: HopsTensorWavefunction + MPS wavefunction container. + 2. system: HopsSystem + System parameters and Hamiltonian. + 3. mode: HopsModes + Bath mode indexing and coupling strengths. + 4. noise_memory: HopsNoiseMemory + Noise memory drift terms. + 5. adaptive: bool + True when state adaptivity is active. + 6. eom_param: dict + Equation-of-motion parameter dictionary; must contain + 'EQUATION_OF_MOTION'. + 7. normalization: str + Scaling convention for the auxiliary vectors + (default 'homps'). 'homps' splits the bath + coupling evenly between the raising and lowering + operators, g_m/sqrt(|g_m|) up and sqrt(|g_m|) + down, so neither direction dominates and the + auxiliary amplitudes stay comparable with + hierarchy depth, which keeps the MPS well + conditioned under truncation. 'adhops' is the + unbalanced scaling of the vector HOPS code, w_m up + and g_m/w_m down, kept so tensor results can be + compared against it term by term. The two differ + by a rescaling of each auxiliary, not by physics. + _build_mpo_parts accepts 'homps' only. + + Returns + ------- + None + """ + self.wavefunction = wavefunction + self.system = system + self.mode = mode + self.noise_memory = noise_memory + self.adaptive = adaptive + self.normalization = normalization + self.flag_linear = eom_param['EQUATION_OF_MOTION'] == 'LINEAR' + self.mpo_cores = [] + self._list_cores_op = [] + self._list_cores_ham = [] + + self._has_lt_corr = False + self._C2_LT_corr_phys = None + self._C2_LT_corr_hier = None + + # Side-channel diagnostics: set by derivative() / step functions. + self.last_matvec_complexity = 0 + self.max_complexity_step = 0 + + self.mpo_builder = MpoBuilder( + k_max=self.wavefunction.k_max, + n_state=self.system.size, + modes_per_state=self.wavefunction.M1_modes_per_state, + n_lop_full=self.mode.n_l2, + ham=self.system.param['HAMILTONIAN'], + state_list=self.system.state_list, + mode=self.mode, + normalization=self.normalization, + n_states_full=self.system.param['NSTATES'], + flag_nearest_neighbor_ham=self.system.flag_nearest_neighbor_ham, + flag_gs_vacuum=self.wavefunction.flag_gs_vacuum, + mpo_epsilon=self.wavefunction.mpo_epsilon, + ) + + def refresh_builder(self) -> None: + """ + Refreshes the persistent MpoBuilder's state-dependent attributes + after an adaptive basis change. + + Clears the MPO cores and the LTC corrections, both of which are + rebuilt from the new active basis on the next build_generator call. + refresh_state_data drops the builder's cached bond factors. + + Parameters + ---------- + None + + Returns + ------- + None + """ + state_list = self.system.state_list + self.mpo_builder.refresh_state_data(len(state_list), state_list) + self.mpo_cores = [] + self._list_cores_op = [] + self._list_cores_ham = [] + self._has_lt_corr = False + self._C2_LT_corr_phys = None + self._C2_LT_corr_hier = None + + def build_generator( + self, + z_mem: np.ndarray, + z_rnd: np.ndarray, + z_rnd2: np.ndarray, + ) -> np.ndarray: + """ + Build the time-evolution MPO into self.mpo_cores and return dz/dt. + + Reads config flags from self.wavefunction. When flag_norm is True, + a nonlinear norm correction is included in the MPO. + + Parameters + ---------- + 1. z_mem: np.ndarray(complex) + Current memory term values. + 2. z_rnd: np.ndarray(complex) + Primary noise at this sub-step (absolute indices). + 3. z_rnd2: np.ndarray(complex) + Secondary noise at this sub-step (absolute indices). + + Returns + ------- + 1. dz_dt: np.ndarray(complex) + Time derivative of the memory terms. + """ + # Local aliases for mode indexing arrays + list_idx_L2_by_hmode = self.mode.list_index_L2_by_hmode + list_L2_coo = self.mode.list_L2_coo + list_absidx_mode = self.mode.list_modeidx_abs + list_absidx_L2 = self.mode.list_l2idx_abs + + # Linear EOM: no zmem, no feedback, no norm correction, dz/dt = 0 + if self.flag_linear: + # LTC for LINEAR: -sum_n c_n L_n^2 applied to physical wf only + list_lt_corr_param = self.system.list_lt_corr_param + if any(list_lt_corr_param): + psi = self.wavefunction.psi + C2_phys = calc_LT_corr_linear( + list_lt_corr_param, self.mode.list_L2_sq_csr, + ) + self._C2_LT_corr_phys = ( + C2_phys.toarray() if sp.issparse(C2_phys) + else np.asarray(C2_phys) + ) + self._C2_LT_corr_hier = None + self._has_lt_corr = True + else: + self._has_lt_corr = False + + list_z_hat = ( + np.conj(z_rnd[list_absidx_L2]) + - 1j * z_rnd2[list_absidx_L2] + ) + self._construct_MPO(list_z_hat, [0.0] * len(list_L2_coo), 0.0) + return np.zeros_like(z_mem) + + # Compute expectation values for each L2 operator. + # + # In the ground-state-as-vacuum convention extract_psi only + # returns the single-excitation amplitudes; the physical GS lives + # at the |0,...,0> configuration and is invisible to traj.psi. + # When wavefunction.flag_gs_vacuum is True we add |<0,...,0|psi>|^2 + # to so the denominator matches the GS-as-state-core + # trajectory's denominator under matching noise, keeping the two + # conventions trajectory-equivalent under NL EOM. + psi = self.wavefunction.psi + if self.wavefunction.flag_gs_vacuum: + norm_sq = self.wavefunction.manifold_norm_sq + list_expect_L2 = [ + (np.conj(psi) @ (list_L2_coo[idx] @ psi)) / norm_sq + for idx in range(len(list_L2_coo)) + ] + else: + list_expect_L2 = [ + operator_expectation(list_L2_coo[idx], psi) + for idx in range(len(list_L2_coo)) + ] + + # Low-temperature correction + list_lt_corr_param = self.system.list_lt_corr_param + norm_corr_lt = 0.0 + if any(list_lt_corr_param): + list_L2_sq_csr = self.mode.list_L2_sq_csr + list_avg_L2_sq = [ + operator_expectation(list_L2_sq_csr[idx], psi) + for idx in range(len(list_L2_sq_csr)) + ] + C2_phys, C2_hier = calc_LT_corr( + list_lt_corr_param, + self.mode.list_L2_csr, + list_expect_L2, + list_L2_sq_csr, + ) + self._C2_LT_corr_phys = ( + C2_phys.toarray() if sp.issparse(C2_phys) + else np.asarray(C2_phys) + ) + self._C2_LT_corr_hier = ( + C2_hier.toarray() if sp.issparse(C2_hier) + else np.asarray(C2_hier) + ) + self._has_lt_corr = True + + if self.wavefunction.flag_norm: + norm_corr_lt = calc_LT_corr_to_norm_corr( + list_lt_corr_param, list_expect_L2, list_avg_L2_sq, + ) + else: + self._has_lt_corr = False + + # TODO: compress_zmem uses list_zmemactivemodeidx_rel which may + # diverge from the mode basis in adaptive runs, producing a wrong + # list_z_hat_rel for the norm correction. + + # First z_hat: uses active z_mem indices (relative) for norm correction. + # This includes only the z_mem modes currently tracked in the basis. + # The full noise coupling is z_hat = conj(z_rnd) + z_mem_compressed + # minus the secondary noise contribution: -1j * z_rnd2 + # (see hops_eom.py line 580: z_hat[j] - 1j * z_rnd2[j]) + list_z_hat_rel = ( + np.conj(z_rnd[list_absidx_L2]) + - 1j * z_rnd2[list_absidx_L2] + + compress_zmem( + z_mem, + list_idx_L2_by_hmode, + self.noise_memory.list_zmemactivemodeidx_rel, + ) + ) + + # Compute the nonlinear norm correction (zero if unnormalized EOM) + if self.wavefunction.flag_norm: + norm_corr = calc_norm_corr_tensor( + self.wavefunction, + psi, + list_z_hat_rel, + list_expect_L2, + self.mode, + self.mode.list_index_L2_by_hmode, + ) + norm_corr_lt + else: + norm_corr = norm_corr_lt + + # Second z_hat: uses absolute mode indices for the MPO construction. + # This differs from the first z_hat because the MPO needs the full + # mode indexing, while the norm correction uses the active subset. + list_z_hat_abs = ( + np.conj(z_rnd[list_absidx_L2]) + - 1j * z_rnd2[list_absidx_L2] + + compress_zmem( + z_mem, list_idx_L2_by_hmode, list_absidx_mode, + ) + ) + + # Build the time-evolution MPO (method-dependent) + if self.wavefunction.method not in ( + 'fullstate', 'number', + ): + raise UnsupportedRequest(self.wavefunction.method, 'build_generator') + self._construct_MPO(list_z_hat_abs, list_expect_L2, norm_corr) + + # Compute dz/dt for the memory term integration + dz_dt = calc_delta_zmem( + z_mem, + list_expect_L2, + self.noise_memory.list_zmemg_abs, + self.noise_memory.list_zmemw_abs, + list_idx_L2_by_hmode, + list_absidx_mode, + self.noise_memory.list_zmemmodeidx_abs, + list_absidx_L2, + self.system.list_activel2idx_abs, + ) + + return dz_dt + + def derivative(self) -> tuple[list[np.ndarray], int]: + """ + Apply the last-built MPO to the wavefunction via tensor_matvec_prod. + + Uses wavefunction.list_cores_phi as the input MPS. + Pure computation — does not mutate the wavefunction. The calling + integrator is responsible for all wavefunction mutations. + + Returns + ------- + 1. cores: list(np.ndarray) + Derivative of the wavefunction in MPS form (-i H|phi>). + + Side effects + ------------ + Sets `self.last_matvec_complexity` to the scalar complexity proxy + of the uncompressed contracted MPS (i.e. its peak size before SVD + truncation) from this call. The step function reads it across + sub-steps rather than the derivative threading it through every + return value to propagate(). + """ + dphi_cores, complexity = tensor_matvec_prod( + self.wavefunction.list_cores_phi, + self.mpo_cores, + self.wavefunction.mps_epsilon, + self.wavefunction.bond_dim_max, + ) + self.last_matvec_complexity = complexity + + # Apply the -1j Schrödinger factor to the MPO-MPS product, + # matching the vector convention where -1j is in the derivative. + scale_mps(dphi_cores, -1j) + + return dphi_cores + + # ------------------------------------------------------------------ + # Private helpers + # ------------------------------------------------------------------ + + def _construct_MPO( + self, + list_z_hat: np.ndarray, + list_expect_L2: list[complex], + norm_corr: float | complex, + ) -> None: + """ + Assembles mpo_cores from hierarchy and Hamiltonian MPO parts. + + Calls _build_mpo_parts to populate _list_cores_op and _list_cores_ham, then + combines them into self.mpo_cores. For number, + the two parts are added via tensor_add; for fullstate, + _list_cores_op is used directly. LTC corrections stored on self + are folded into the final MPO when _has_lt_corr is True. + + Parameters + ---------- + 1. list_z_hat: np.ndarray(complex) + Conjugate noise, indexed by L2. + 2. list_expect_L2: list(complex) + expectation values, one per L2 + operator. + 3. norm_corr: float | complex + Normalization correction term. + + Returns + ------- + None + """ + self._build_mpo_parts(list_z_hat, list_expect_L2, norm_corr) + + method = self.wavefunction.method + mpo_epsilon = self.wavefunction.mpo_epsilon + if self._list_cores_ham: + # Separate Hamiltonian part: add it to the hierarchy MPO. MPO + # compression uses mpo_epsilon (defaults to 0 — exact). The MPS + # epsilon/bond_dim_max control wavefunction approximation and are + # decoupled from operator truncation; setting mpo_epsilon > 0 lets + # the user opt into operator compression independently. + mpo_bond_dim_max = ( + max(c.shape[-1] for c in self._list_cores_op) + + max(c.shape[-1] for c in self._list_cores_ham) + ) + self.mpo_cores = tensor_add( + self._list_cores_op, + self._list_cores_ham, + epsilon=mpo_epsilon, + bond_dim_max=mpo_bond_dim_max, + ) + else: + self.mpo_cores = self._list_cores_op + + # Fold LTC corrections into the assembled MPO + if self._has_lt_corr: + k_max = self.wavefunction.k_max + M1_modes_per_state = self.wavefunction.M1_modes_per_state + n_state = self.mpo_builder.n_state + + # Hierarchy correction: system-space operator on all k levels. + # For fullstate this is in the Hamiltonian channel; + # for statenumber it is a separate operator MPO. + if (self._C2_LT_corr_hier is not None + and method == 'number'): + ltc_hier_mpo = build_statenumber_operator_mpo( + 1j * self._C2_LT_corr_hier, + n_state, k_max, M1_modes_per_state, + ) + max_bond = ( + max(c.shape[-1] for c in self.mpo_cores) + + max(c.shape[-1] for c in ltc_hier_mpo) + ) + self.mpo_cores = tensor_add( + self.mpo_cores, ltc_hier_mpo, + epsilon=mpo_epsilon, bond_dim_max=max_bond, + ) + + # Physical correction: system-space operator at k=0 only + if self._C2_LT_corr_phys is not None: + ltc_phys_mpo = build_physical_correction_mpo( + 1j * self._C2_LT_corr_phys, method, k_max, + M1_modes_per_state, + ) + max_bond = ( + max(c.shape[-1] for c in self.mpo_cores) + + max(c.shape[-1] for c in ltc_phys_mpo) + ) + self.mpo_cores = tensor_add( + self.mpo_cores, ltc_phys_mpo, + epsilon=mpo_epsilon, bond_dim_max=max_bond, + ) + + def _build_mpo_parts( + self, + list_z_hat: np.ndarray, + list_expect_L2: list[complex], + norm_corr: float | complex, + ) -> None: + """ + Build _list_cores_op and _list_cores_ham from noise and system inputs. + + Delegates to MpoBuilder and the appropriate MPO builder method for the + current representation. The number method builds a single combined + generator for any Hamiltonian, each bond carrying as many channels as + the rank of the coupling block that crosses it; fullstate always + builds one MPO, with the Hamiltonian embedded in it. _list_cores_ham + is left empty in every case except the FLAG_MPO_OPTIMIZE=False + reference path, which is how _construct_MPO knows whether there is + anything to add. + + Parameters + ---------- + 1. list_z_hat: np.ndarray(complex) + Conjugate noise (z* + compressed z_mem), + indexed by L2. + 2. list_expect_L2: list(complex) + expectation values, one per L2 + operator. + 3. norm_corr: float | complex + Nonlinear norm correction (zero for linear or + unnormalized EOM). Includes LTC when active. + + Returns + ------- + None + """ + if self.normalization != 'homps': + raise UnsupportedRequest( + self.normalization, + '_build_mpo_parts normalization', + ) + + builder = self.mpo_builder + + self._list_cores_op = [] + self._list_cores_ham = [] + if self.wavefunction.method == 'number': + if not self.wavefunction.flag_mpo_optimize: + # Reference path: the hierarchy and the Hamiltonian are built + # apart and summed by _construct_MPO. Kept as the independent + # check the combined generators are tested against. + self._list_cores_op = builder.build_statenumber_hierarchy_mpo( + list_z_hat, + list_expect_L2, + norm_corr, + ) + self._list_cores_ham = builder.build_statenumber_ham_mpo() + return + # One combined Hamiltonian + hierarchy MPO, each bond built at the + # rank of the coupling block that crosses it, leaving nothing to + # add or compress. _list_cores_ham stays empty to signal that. + self._list_cores_op = builder.build_general_generator_mpo( + list_z_hat, list_expect_L2, norm_corr, + ) + elif self.wavefunction.method == 'fullstate': + # One MPO carries the hierarchy, the Hamiltonian and the LTC + # hierarchy correction together. + self._list_cores_op = builder.build_fullstate_mpo( + list_z_hat, + list_expect_L2, + norm_corr, + C2_lt_corr_hier=self._C2_LT_corr_hier, + ) + else: + raise UnsupportedRequest(self.wavefunction.method, '_build_mpo_parts') diff --git a/src/mesohops/tensor/hops_tensor_wavefunction.py b/src/mesohops/tensor/hops_tensor_wavefunction.py new file mode 100644 index 0000000..e9f83ce --- /dev/null +++ b/src/mesohops/tensor/hops_tensor_wavefunction.py @@ -0,0 +1,776 @@ +""" +MPS wavefunction container for tensor HOPS. + +Defines HopsTensorWavefunction, which stores the hierarchy wavefunction as a Matrix +Product State (MPS) and provides initialization, normalization, operator +application, and bond-dimension management. +""" +from __future__ import annotations + +import bisect + +import numpy as np + +from mesohops.basis.hops_system import HopsSystem +from mesohops.util.exceptions import LockedException, UnsupportedRequest +from mesohops.util.tensor_operations import ( + extract_gs_amp, + extract_psi, + flatten_cores, + unflatten_cores, +) + +__title__ = 'HopsTensorWavefunction Class' +__author__ = 'B. Z. Citty' +__maintainer__ = 'B. Z. Citty' + + +class HopsTensorWavefunction: + """ + Matrix Product State (MPS) representation of the HOPS wavefunction. + + In the standard HOPS formalism, the full state vector is a flat array + indexed by (physical state, auxiliary index). HopsTensorWavefunction re-encodes this + as an MPS whose sites correspond to system states and bath modes. The + exact local structure depends on the encoding method: in the statenumber + encoding each system state is represented by a binary (dimension-2) core, + while mode cores carry a local dimension of ``k_max + 1``, reflecting the + hierarchy depth along that mode axis. The bond dimensions between adjacent + sites are determined by SVD compression and are bounded by ``bond_dim_max``. + + Time evolution is driven either by applying a Matrix Product Operator + (MPO) and then truncating via SVD (Runge-Kutta path), or by a + single-site TDVP sweep that naturally respects the MPS geometry. + + Responsibilities + ---------------- + - Store and update the MPS cores (``list_cores_phi``). + - Provide norm and bond-dimension utilities. + - Apply system-space operators to the MPS. + + MPO construction and time-derivative computation are owned by + ``HopsTensorEOM``, which reads config flags from this class. + + Key attributes + -------------- + list_cores_phi: list(np.ndarray) + MPS cores for the current wavefunction. + k_max: int + Maximum hierarchy depth; local dimension of each mode core = k_max + 1. + mps_epsilon: float + SVD truncation threshold used during MPS compression + (mirrors tensor_param['MPS_EPSILON']). + mpo_epsilon: float + SVD truncation threshold for MPO assembly, applied where MPO parts + are summed and on the coupling-block factorization behind the + general generator (mirrors tensor_param['MPO_EPSILON']). + flag_mpo_optimize: bool + True (default) builds the number-method generator as one MPO at the + rank of each bond's coupling block; False takes the two-MPO + reference path (mirrors tensor_param['FLAG_MPO_OPTIMIZE']). + method: str + MPS representation ('number' or 'fullstate'). + bond_dim_max: int + Hard cap on MPS bond dimension. + flag_tdvp: bool + True when the TDVP integrator is active; False for Runge-Kutta. + flag_norm: bool + True when the normalized nonlinear EOM is used. + _flag_gs_vacuum: bool + True when the physical ground state is the all-zeros MPS + configuration; its amplitude counts toward the nonlinear norm. + Notes + ----- + This class is owned by ``HopsTensorTrajectory`` and is populated during + ``HopsTensorTrajectory.initialize``. It does not hold references to + ``HopsSystem`` or ``HopsModes``; those are owned by ``HopsTensorEOM``, + which holds a reference to this object and reads its config flags. + """ + + __slots__ = ( + # --- Configuration --- + 'k_max', # Maximum hierarchy depth (int) + 'mps_epsilon', # MPS SVD truncation threshold (float) + 'mpo_epsilon', # MPO SVD truncation threshold (float) + 'flag_mpo_optimize', # Number path: combined generator (bool) + 'method', # MPS representation, 'number' or 'fullstate' (str) + 'bond_dim_max', # Maximum MPS bond dimension (int) + 'flag_norm', # True for normalized nonlinear EOM (bool) + 'flag_tdvp', # True when TDVP integrator is active (bool) + '_flag_gs_vacuum', # Ground state as all-zeros vacuum (bool) + # --- MPS data (populated by initialize) --- + 'list_cores_phi', # MPS cores of the wavefunction (list[np.ndarray]) + 'M1_modes_per_site', # Active modes per MPS site, sorted (np.ndarray[int]) + # --- System-derived data (set by initialize) --- + 'M1_modes_per_state', # Bath modes per state (np.ndarray[int]) + 'M1_mode_offset', # Cumulative mode offset per state (np.ndarray[int]) + '__initialized__', # Initialization status flag (bool) + ) + + def __init__(self, k_max: int, tensor_param: dict, integrator_param: dict, + eom_param: dict) -> None: + """ + Initializes the HopsTensorWavefunction configuration. Does not + create system, mode, or noise_memory objects — those live in + HopsTensorBasis. + + Parameters + ---------- + 1. k_max: int + Maximum hierarchy depth. + + 2. tensor_param: dict + Tensor configuration parameters. + a. 'METHOD': str — tensor topology type + b. 'MPS_EPSILON': float — SVD truncation threshold for MPS + c. 'MPO_EPSILON': float — SVD truncation threshold for MPO + d. 'FLAG_MPO_OPTIMIZE': bool — use the combined generator MPO + e. 'BOND_DIM_MAX': int — maximum MPS bond dimension + + 3. integrator_param: dict + Integration parameters. + a. 'INTEGRATOR': str — 'RUNGE_KUTTA', 'TDVP1', or 'TDVP2' + + 4. eom_param: dict + Equation of motion parameters. + a. 'EQUATION_OF_MOTION': str — 'NORMALIZED NONLINEAR', + 'NONLINEAR', or 'LINEAR' + """ + + self.k_max = k_max + self.mps_epsilon = tensor_param.get('MPS_EPSILON', 0.0) + self.mpo_epsilon = tensor_param.get('MPO_EPSILON', 0.0) + self.flag_mpo_optimize = tensor_param.get('FLAG_MPO_OPTIMIZE', True) + self.method = tensor_param['METHOD'] + self.bond_dim_max = tensor_param['BOND_DIM_MAX'] + # True when the normalized nonlinear equation of motion is active + self.flag_norm = ( + eom_param['EQUATION_OF_MOTION'] == 'NORMALIZED NONLINEAR' + ) + # True when a TDVP integrator is active (TDVP1 or TDVP2) + self.flag_tdvp = integrator_param['INTEGRATOR'] in ('TDVP1', 'TDVP2') + # Set by the spectroscopy layer; not a user tensor_param. + self._flag_gs_vacuum = False + # MPS data — populated by initialize + self.list_cores_phi = [] + self.M1_modes_per_site = None + # System-derived data — set by initialize + self.M1_modes_per_state = None + self.M1_mode_offset = None + self.__initialized__ = False + + def initialize(self, phi_0: np.ndarray, system: HopsSystem) -> None: + """ + Sets system-derived constants, initializes dimension scalars, and builds + the initial MPS. + + Parameters + ---------- + 1. phi_0: np.array(complex) + Initial physical wavefunction. + + 2. system: HopsSystem + Initialized system object from HopsTensorBasis. + + Returns + ------- + None + """ + if self.__initialized__: + raise LockedException('initialize', 'HopsTensorWavefunction') + # Per-state mode counts from the L-operator → state mapping built + # during system initialization. + self.M1_modes_per_state = np.array([ + len(system.param['LIST_HMODE_INDICES_BY_STATE'][s]) + for s in range(system.param['NSTATES']) + ], dtype=int) + # The number representation gives every state group the same mode cores, + # so a state carrying no hierarchy modes has no representation here. The + # ground state is the all-zeros configuration, not a state of its own. + if self.method == 'number' and np.any(self.M1_modes_per_state == 0): + raise UnsupportedRequest( + 'a state with no hierarchy modes', + 'HopsTensorWavefunction.initialize with METHOD=number', + ) + self.M1_mode_offset = np.concatenate( + [[0], np.cumsum(self.M1_modes_per_state)] + ) + self.M1_modes_per_site = self.M1_modes_per_state[sorted(system.state_list)] + # Restrict phi_0 to the active state_list before building MPS + self.build_list_cores_phi(phi_0[system.state_list], len(system.state_list)) + # TDVP sweeps require all bond dimensions to match; pad with + # zeros up to bond_dim_max so the first sweep can proceed + if self.flag_tdvp: + self.inflate_bonds_to(self.bond_dim_max, eps=0.0) + self.__initialized__ = True + + def build_list_cores_phi(self, psi_0: np.ndarray, n_state: int) -> None: + """ + Builds the initial MPS form of the wave function given the initial physical + wave function psi_0 in array form. + + Precondition: self.M1_modes_per_state must be set before calling. + + Parameters + ---------- + 1. psi_0: np.array(complex) + Physical wavefunction restricted to the current state_list. + + 2. n_state: int + Number of active states. + + Returns + ------- + None + """ + # Construct MPS form of Phi + if self.method == 'number': + # Each physical site s gets a group [state_core_s, mode_core_s0, ...] + # State core: binary (dim-2), index 1 = occupied, 0 = unoccupied. + # Mode cores: initialized to hierarchy ground state |0> (index 0 = 1). + # States with no bath coupling get a group with only the state core. + for site in range(n_state): + T3_core_state = np.zeros(shape=(1, 2, 1), dtype=np.complex128) + T3_core_state[0, 1, 0] = psi_0[site] + T3_core_state[0, 0, 0] = 1.0 + T3_core_mode = np.zeros( + shape=(1, self.k_max + 1, 1), dtype=np.complex128 + ) + T3_core_mode[0, 0, 0] = 1.0 + # Build group: state core followed by one mode core per mode. + # .copy() prevents all mode cores sharing the same array. + group = [T3_core_state] + [ + T3_core_mode.copy() for _ in range(self.M1_modes_per_state[site]) + ] + self.list_cores_phi.append(group) + + elif self.method == 'fullstate': + # Single system core with physical dim = n_state, then + # one mode core per hierarchy mode across all states, each + # initialized to hierarchy ground state |0>. States with no + # bath coupling contribute 0 mode cores. + T3_core_state = np.zeros(shape=(1, n_state, 1), dtype=np.complex128) + T3_core_state[0, :, 0] = psi_0 + T3_core_mode = np.zeros(shape=(1, self.k_max + 1, 1), dtype=np.complex128) + T3_core_mode[0, 0, 0] = 1.0 + n_total_modes = int(self.M1_mode_offset[-1]) + self.list_cores_phi = [T3_core_state] + [ + T3_core_mode.copy() for _ in range(n_total_modes) + ] + + else: + raise UnsupportedRequest(self.method, 'build_list_cores_phi') + + def normalize(self) -> float: + """ + Normalizes the main wave function phi_0. Always normalizes — the + decision of whether to normalize (based on EOM type) belongs to the + calling trajectory, not the wavefunction. + + Parameters + ---------- + None + + Returns + ------- + 1. norm_psi: float + Norm of psi before normalization. + """ + norm_psi = np.linalg.norm(self.psi) + # Dividing any single core by a scalar divides the full MPS + # contraction result by that scalar, so normalizing via the + # first core is exact regardless of canonical form. + # For statenumber, list_cores_phi[0] is a group (list), so + # index [0][0] to reach the first state core array. + if self.method == 'number': + self.list_cores_phi[0][0] = self.list_cores_phi[0][0] / norm_psi + elif self.method == 'fullstate': + self.list_cores_phi[0] = self.list_cores_phi[0] / norm_psi + else: + raise UnsupportedRequest(self.method, 'normalize') + return norm_psi + + def add_state_cores( + self, list_states_new: list[int], state_list: list[int], + ) -> None: + """ + Adds new state and mode cores to the MPS for newly activated states. + + For fullstate representation, expands the state core to include new + states and appends mode cores at the end of the flat core list. For + statenumber representation, inserts new state groups (state core + + mode cores) at the sorted position within the list-of-lists structure. + + Parameters + ---------- + 1. list_states_new : list(int) + Absolute state indices to add. + + 2. state_list : list(int) + Current list of active state indices before adding. + + Returns + ------- + None + """ + n_state = len(state_list) + list_states_new = sorted(list(list_states_new)) + if self.method == 'fullstate': + # Expand the state core to accommodate new states. The old + # state data is scattered into a larger core at the positions + # where those states appear in the combined sorted list. + list_total_states = sorted(list(state_list) + list(list_states_new)) + new_state_list_len = len(list_total_states) + # Where old states land in the expanded physical dimension + list_old_idx = sorted( + [list_total_states.index(state) for state in state_list] + ) + state_core_shape = self.list_cores_phi[0].shape + new_state_core = np.zeros( + shape=(1, new_state_list_len, state_core_shape[2]), + dtype=np.complex128, + ) + # Source indices: all entries from the current (smaller) state core + old_tensor_indices = np.ix_( + np.array([0]), + np.arange(n_state), + np.arange(state_core_shape[2]), + ) + # Destination indices: scatter into the expanded core at the + # positions corresponding to the old states + new_tensor_indices = np.ix_( + np.array([0]), + np.array(list_old_idx), + np.arange(state_core_shape[2]), + ) + new_state_core[new_tensor_indices] = self.list_cores_phi[0][ + old_tensor_indices + ] + self.list_cores_phi[0] = new_state_core + + # Insert mode cores at sorted positions (not appended at end). + # This ensures list_cores_phi maintains sorted state ordering. + current_sorted = sorted(state_list) + for state in list_states_new: + insert_site = bisect.bisect_left(current_sorted, state) + # Flat core index: 1 (state core) + modes for all sites before insertion + core_pos = 1 + for s in current_sorted[:insert_site]: + core_pos += self.M1_modes_per_state[s] + for m in range(self.M1_modes_per_state[state]): + new_mode_core = np.zeros( + shape=(1, self.k_max + 1, 1), dtype=np.complex128 + ) + new_mode_core[0, 0, 0] = 1 + 0j + self.list_cores_phi.insert(core_pos + m, new_mode_core) + current_sorted.insert(insert_site, state) + + elif self.method == 'number': + current_state_order = sorted(list(state_list)) + for state in list_states_new: + # Insert the new group at the sorted position so the + # MPS site ordering matches the state index ordering + group_idx = bisect.bisect_left(current_state_order, state) + # Match the bond dimension of the neighboring group so + # the MPS bonds remain compatible + bond_dim = ( + self.list_cores_phi[group_idx - 1][-1].shape[2] + if group_idx > 0 + else 1 + ) + # State core: identity along bond dimension at the + # |unoccupied> index (0), so the new state acts as a + # pass-through and does not alter the existing MPS + new_state_core = np.zeros( + (bond_dim, 2, bond_dim), dtype=np.complex128 + ) + for bond in range(bond_dim): + new_state_core[bond, 0, bond] = 1 + 0j + new_group = [new_state_core] + # Mode cores: identity at hierarchy ground state |0> + for mode in range(self.M1_modes_per_state[state]): + new_mode_core = np.zeros( + (bond_dim, self.k_max + 1, bond_dim), dtype=np.complex128 + ) + for bond in range(bond_dim): + new_mode_core[bond, 0, bond] = 1 + 0j + new_group.append(new_mode_core) + self.list_cores_phi.insert(group_idx, new_group) + current_state_order.insert(group_idx, state) + + else: + raise UnsupportedRequest(self.method, 'add_state_cores') + + new_state_list = sorted(set(state_list) | set(list_states_new)) + self.M1_modes_per_site = self.M1_modes_per_state[new_state_list] + + if self.flag_tdvp: + self.inflate_bonds_to(self.bond_dim_max, eps=0.0) + + def remove_state_cores( + self, list_states_old: list[int], state_list: list[int], + ) -> None: + """ + Removes state and mode cores from the MPS for deactivated states. + + For fullstate representation, shrinks the state core to exclude removed + states and absorbs their mode cores into adjacent cores. For + statenumber representation, absorbs each removed group's |0> slices + into the previous group and pops it from the list-of-lists structure. + + Parameters + ---------- + 1. list_states_old : list(int) + Absolute state indices to remove. + + 2. state_list : list(int) + Current list of active state indices before removing. + + Returns + ------- + None + """ + if self.method == 'fullstate': + old_state_indices = [list(state_list).index(s) + for s in list_states_old] + list_remaining_state_idx = list( + set(np.arange(len(state_list))) - set(old_state_indices) + ) + new_length = len(state_list) - len(old_state_indices) + + # Modify state core for remaining states + state_core_shape = self.list_cores_phi[0].shape + new_state_core = np.zeros( + shape=(1, new_length, state_core_shape[2]), dtype=np.complex128 + ) + old_tensor_indices = np.ix_( + np.array([0]), + list_remaining_state_idx, + np.arange(state_core_shape[2]), + ) + new_tensor_indices = np.ix_( + np.array([0]), + np.arange(new_length), + np.arange(state_core_shape[2]), + ) + new_state_core[new_tensor_indices] = self.list_cores_phi[0][ + old_tensor_indices + ] + self.list_cores_phi[0] = new_state_core + + # Build flat indices of mode cores to remove (1-indexed since + # core 0 is the state core). Compute before any mutation. + state_order = sorted(state_list) + list_flat_mode_indices = [] + for old_state in list_states_old: + order_idx = state_order.index(old_state) + base = 1 + for j in range(order_idx): + base += self.M1_modes_per_state[state_order[j]] + n_modes = self.M1_modes_per_state[old_state] + for m in range(n_modes): + list_flat_mode_indices.append(base + m) + + # Remove in reverse order so earlier indices stay valid. + # Each removed mode core is contracted into its left neighbor + # at the hierarchy ground state slice [:, 0, :]. This is valid + # because a removed state has no hierarchy excitations, so its + # |0> slice carries all the weight and higher slices are zero. + for flat_idx in sorted(list_flat_mode_indices, reverse=True): + mode_core = self.list_cores_phi[flat_idx] + mode_shape = np.shape(mode_core) + # Extract the |0> slice as a 2D transfer matrix (bond x bond) + new_mode_core = np.zeros( + shape=(mode_shape[0], mode_shape[2]), + dtype=np.complex128, + ) + new_indices = np.ix_( + np.arange(mode_shape[0]), np.arange(mode_shape[2]) + ) + old_indices = np.ix_( + np.arange(mode_shape[0]), + [0], + np.arange(mode_shape[2]), + ) + new_mode_core[new_indices] = mode_core[old_indices][:, 0, :] + # Absorb the transfer matrix into the adjacent core + self.list_cores_phi[flat_idx - 1] = ( + self.list_cores_phi[flat_idx - 1] @ new_mode_core + ) + self.list_cores_phi.pop(flat_idx) + + elif self.method == 'number': + list_group_indices = sorted( + [sorted(state_list).index(s) + for s in list_states_old], + reverse=True, + ) + for group_idx in list_group_indices: + group = self.list_cores_phi[group_idx] + if group_idx > 0: + prev_group = self.list_cores_phi[group_idx - 1] + # Absorb state core's |0> slice into previous group's last core + prev_group[-1] = prev_group[-1] @ group[0][:, 0, :] + # Chain-absorb each mode core's |0> slice + for mode_core in group[1:]: + prev_group[-1] = prev_group[-1] @ mode_core[:, 0, :] + else: + # TODO: structural bug — when removing the leftmost group + # (group_idx == 0), the |0> slices are discarded instead + # of absorbed into the next group. The |0> slice is a + # (1, Dr0) row vector carrying normalization information. + # Discarding it leaves the next group with wrong left + # boundary bond dimension (Dr0 instead of 1). Fix: absorb + # into next group via + # np.einsum('ij,jkl->ikl', group[0][:, 0, :], + # next_group[0]) + # and chain-absorb mode core |0> slices similarly. + pass + self.list_cores_phi.pop(group_idx) + + else: + raise UnsupportedRequest(self.method, 'remove_state_cores') + + new_state_list = sorted(set(state_list) - set(list_states_old)) + self.M1_modes_per_site = self.M1_modes_per_state[new_state_list] + + def inflate_bonds_to(self, chi_target: int, eps: float = 0.0) -> None: + """ + Pads every internal bond of list_cores_phi up to chi_target by appending + zero (or small random) columns/rows at each bond interface. + + Parameters + ---------- + 1. chi_target: int + Target bond dimension. Bonds already at or above this + value are left unchanged. + + 2. eps: float + If greater than zero, the padding entries are filled with + complex Gaussian noise of standard deviation eps rather than + zeros. Default: 0.0. + + Returns + ------- + None + """ + # Flatten statenumber groups into a single core list so the + # padding loop can treat all representations uniformly. + if self.method == 'number': + list_modes_per_site = [len(g) - 1 for g in self.list_cores_phi] + list_cores = flatten_cores(self.list_cores_phi) + elif self.method == 'fullstate': + list_modes_per_site = None + list_cores = self.list_cores_phi + else: + raise UnsupportedRequest(self.method, 'inflate_bonds_to') + # Walk each internal bond (shared between adjacent cores) and pad + # any bond whose dimension is below chi_target. Each MPS core has + # shape (bond_left, phys, bond_right), so the right bond of core i + # must equal the left bond of core i+1. Padding appends zero + # slices to both sides of the interface to keep them consistent: + # core[i]: (bond_left, phys, bond_right) → (bond_left, phys, chi_target) + # core[i+1]: (bond_right, phys', bond_right') + # → (chi_target, phys', bond_right') + n_cores = len(list_cores) + for i in range(n_cores - 1): + T3_core_cur = list_cores[i] + T3_core_next = list_cores[i + 1] + bond_dim_left, phys_dim, bond_dim_right = T3_core_cur.shape + _, phys_dim_next, bond_dim_right_next = T3_core_next.shape + if bond_dim_right < chi_target: + n_pad_right = chi_target - bond_dim_right + # Pad right bond of current core with zeros (or noise) + T3_pad_cur = np.zeros( + (bond_dim_left, phys_dim, n_pad_right), dtype=T3_core_cur.dtype + ) + if eps > 0: + T3_pad_cur += eps * ( + np.random.randn(*T3_pad_cur.shape) + + 1j * np.random.randn(*T3_pad_cur.shape) + ) + list_cores[i] = np.concatenate([T3_core_cur, T3_pad_cur], axis=2) + # Pad left bond of next core with zeros to match + T3_pad_next = np.zeros( + (n_pad_right, phys_dim_next, bond_dim_right_next), + dtype=T3_core_next.dtype, + ) + list_cores[i + 1] = np.concatenate([T3_core_next, T3_pad_next], axis=0) + # Restore statenumber group structure from the padded flat list + if self.method == 'number': + self.list_cores_phi = unflatten_cores(list_cores, list_modes_per_site) + + def update_phi_from_flat(self, list_cores_flat: list[np.ndarray]) -> None: + """Set list_cores_phi from a flat core list. + + For number representation, restores the list-of-lists structure + using the current group sizes. For fullstate representation, assigns + the flat list directly. + + Parameters + ---------- + 1. list_cores_flat: list(np.ndarray) + Cores in statenumber MPS representation order. + + Returns + ------- + None + """ + if self.method == 'number': + list_modes_per_site = [len(g) - 1 for g in self.list_cores_phi] + self.list_cores_phi = unflatten_cores( + [c.copy() for c in list_cores_flat], list_modes_per_site + ) + else: + self.list_cores_phi = [c.copy() for c in list_cores_flat] + + def restore_phi( + self, tensor_source: list[np.ndarray] | HopsTensorWavefunction, + ) -> None: + """ + Sets list_cores_phi by copying cores from a list or another + HopsTensorWavefunction. + + Always makes an independent copy of the input cores so that subsequent + mutations to tensor_source do not affect this instance. + + Parameters + ---------- + 1. tensor_source: list(np.ndarray) or HopsTensorWavefunction + Source to copy cores from. If a list, each element + must be a 3D ndarray with axes [bond_left, phys_dim, + bond_right]. If a HopsTensorWavefunction, + cores are copied from its list_cores_phi. + + Returns + ------- + None + """ + if isinstance(tensor_source, HopsTensorWavefunction): + source = tensor_source.list_cores_phi + elif isinstance(tensor_source, list): + source = tensor_source + else: + raise TypeError( + 'Expected list or HopsTensorWavefunction, ' + f'got {type(tensor_source).__name__}' + ) + # Deep copy on restore prevents the integrator from mutating + # the checkpoint arrays through aliased references + if self.method == 'number': + # source is a list of groups; deep-copy each ndarray within each group + self.list_cores_phi = [[arr.copy() for arr in g] for g in source] + else: + self.list_cores_phi = [c.copy() for c in source] + + def get_core_shapes(self) -> list[tuple]: + """ + Returns the shape of every core in the wavefunction MPS. + + Parameters + ---------- + None + + Returns + ------- + 1. list_shapes: list(tuple) + List of (left_bond, phys_dim, right_bond) for each core. + """ + if self.method == 'number': + list_shapes = [core.shape for core in flatten_cores(self.list_cores_phi)] + else: + list_shapes = [core.shape for core in self.list_cores_phi] + return list_shapes + + def check_bondsize(self) -> bool: + """ + Checks that adjacent MPS cores have compatible bond dimensions. + + Parameters + ---------- + None + + Returns + ------- + 1. valid: bool + True if all internal bond dimensions are consistent. + """ + # Adjacent cores must have matching bond dimensions: right bond of + # core i must equal left bond of core i+1. + if self.method == 'number': + list_cores_flat = flatten_cores(self.list_cores_phi) + else: + list_cores_flat = self.list_cores_phi + list_bond_dim_left = [core.shape[0] for core in list_cores_flat][1:] + list_bond_dim_right = [core.shape[2] for core in list_cores_flat][:-1] + return np.array_equal(list_bond_dim_left, list_bond_dim_right) + + @property + def manifold_norm_sq(self) -> np.float64: + """ + Physical-wavefunction norm-squared + |<0,...,0|psi>|^2 on + the GS + single-excitation manifold of a vacuum-convention MPS. + + Under the vacuum convention extract_psi only returns the single- + excitation amplitudes; the physical ground-state amplitude lives at + the all-zeros MPS configuration and is invisible to traj.psi. + Adding |<0,...,0|psi>|^2 to recovers the same physical- + wavefunction norm the gs_core-convention trajectory has by + construction. + + Returns + ------- + 1. norm_sq : np.float64 + + |<0,...,0|psi>|^2. + """ + psi = self.psi + norm_sq_excited = np.linalg.norm(psi) ** 2 + gs_amp = extract_gs_amp(self.list_cores_phi, self.method) + return norm_sq_excited + np.abs(gs_amp) ** 2 + + @property + def flat_cores(self) -> list[np.ndarray]: + """All MPS cores as a 1D flat list suitable for tensor algebra routines. + + For number representation this flattens the list-of-lists; + for fullstate representation the list is already flat. + + Returns + ------- + 1. list_cores_flat: list(np.ndarray) + Cores in statenumber MPS representation order. + """ + if self.method == 'number': + return flatten_cores(self.list_cores_phi) + return self.list_cores_phi + + @property + def psi(self) -> np.ndarray: + """ + Physical (system) wavefunction extracted from the MPS. + + Returns + ------- + 1. psi: np.array(complex) + Physical wavefunction of length n_state. + """ + return extract_psi( + self.list_cores_phi, self.method, + self.M1_modes_per_site, + ) + + @property + def flag_gs_vacuum(self) -> bool: + """True when the ground state is tracked as the all-zeros vacuum.""" + return self._flag_gs_vacuum + + @flag_gs_vacuum.setter + def flag_gs_vacuum(self, value: bool) -> None: + # The vacuum (all-zeros) config only exists in the number + # representation, so the flag is meaningless under fullstate. + if value and self.method != 'number': + raise UnsupportedRequest( + f"flag_gs_vacuum=True with method={self.method!r}", + 'HopsTensorWavefunction.flag_gs_vacuum', + ) + self._flag_gs_vacuum = value diff --git a/src/mesohops/tensor/mpo_constructors.py b/src/mesohops/tensor/mpo_constructors.py new file mode 100644 index 0000000..4ea3d69 --- /dev/null +++ b/src/mesohops/tensor/mpo_constructors.py @@ -0,0 +1,1556 @@ +""" +MPO construction routines and shared operator building blocks for +HopsTensorWavefunction. + +Classes +------- +MpoBuilder + Precomputed ladder operators, projectors, and Hamiltonian slices shared + by all MPO builder methods. All elementary operators are computed during + __init__; builder methods use self.X to access them. + +Functions +--------- +build_statenumber_operator_mpo(H2_op, n_state, k_max, M1_modes_per_state) + Build an MPO that applies an arbitrary (n_state x n_state) system operator + to a statenumberrepresentation MPS, acting as identity on all mode cores. + +build_statenumber_dipole_mpo(list_mu, n_state, k_max, M1_modes_per_state, + raise_or_lower) + Build the bond-dim-2 MPO for a sum-of-single-site dipole raise or + lower operator in the statenumberrepresentation, under the + ground-state-as-vacuum convention. + +build_statenumber_dipole_lower_plus_ident_mpo(list_mu, n_state, k_max, + M1_modes_per_state) + Build the bond-dim-3 MPO for (I + mu^-) = identity-everywhere plus + sum-of-single-site dipole lower operator, in the + statenumberrepresentation under the ground-state-as-vacuum + convention. Used by the fluorescence pathway to preserve + single-excitation content while injecting ground-state amplitude. + +Methods on MpoBuilder +--------------------- +build_statenumber_hierarchy_mpo(list_z_hat, list_expect_L2, norm_corr) + Build the hierarchy-interaction MPO (noise, bath frequency, and + coupling terms) for the statenumberrepresentation. Bond dimension 5. + +build_statenumber_ham_mpo() + Build the Hamiltonian MPO for the statenumberrepresentation. + Dispatches to nearest-neighbor or general builder based on + flag_nearest_neighbor_ham. + +build_general_generator_mpo(list_z_hat, list_expect_L2, norm_corr) + Build the combined Hamiltonian + hierarchy MPO for a Hamiltonian of any + coupling pattern, in the number representation on the one-excitation + manifold, so that no MPO addition or compression is needed. Each bond + carries as many channels as the rank of the coupling block that crosses + it, giving peak bond dimension 5 on a chain or star and 7 on a ring. The + two-MPO path assembles 9 (nearest neighbor) or 5 + 2*n_state (general) + before compressing, and is kept as the FLAG_MPO_OPTIMIZE=False reference. + +build_fullstate_mpo(list_z_hat, list_expect_L2, norm_corr) + Build the combined hierarchy+Hamiltonian MPO for the + fullstaterepresentation. The first core covers the full system + Hilbert space; subsequent cores handle each bath mode. + Bond dimension tapers n_lop_full + 2, n_lop_full + 1, ..., 3 as each + L-operator channel closes at its own last mode core. +""" + +from __future__ import annotations + +import numpy as np + +# --- Shared 2x2 elementary operators for statenumber MPO construction --- +# Site projector |1><1| (occupied state) +_P2_SITE = np.array([[0, 0], [0, 1]], dtype=np.complex128) +# Unoccupied-site projector Q = 1-P = |0><0| +_Q2_SITE = np.array([[1, 0], [0, 0]], dtype=np.complex128) +# Raise operator sigma^+ = |1><0| (carries coupling rightward) +_T2_PLUS = np.array([[0, 0], [1, 0]], dtype=np.complex128) +# Lower operator sigma^- = |0><1| (carries coupling leftward) +_T2_MINUS = np.array([[0, 1], [0, 0]], dtype=np.complex128) + +__title__ = 'MPO Constructors' +__author__ = 'B. Z. Citty' +__maintainer__ = 'B. Z. Citty' + + +def _identity_mode_cores(bond_dim, k_max, n_copies): + """Build identity mode cores that pass all bond channels through. + + Parameters + ---------- + 1. bond_dim: int + MPO bond dimension (left = right). + 2. k_max: int + Maximum hierarchy depth (physical dim = k_max + 1). + 3. n_copies: int + Number of identical cores to return. + + Returns + ------- + 1. list_cores: list(np.ndarray) + n_copies cores of shape (bond_dim, k_max+1, k_max+1, bond_dim). + """ + T4_core = np.zeros( + (bond_dim, k_max + 1, k_max + 1, bond_dim), dtype=np.complex128, + ) + idx_bond = np.arange(bond_dim) + T4_core[idx_bond, :, :, idx_bond] = np.eye(k_max + 1, dtype=np.complex128) + return [T4_core.copy() for _ in range(n_copies)] + + +class MpoBuilder: + """ + Precomputed building blocks shared by all MPO builder methods. + + All elementary operators are computed during __init__. A new instance + is created on each call to HopsTensorWavefunction.build_op; there is + no persistent state between timesteps. + + The number-representation builders carry one L-operator per site, ordered + by site, so a site index is also its L2 index into list_z_hat and + list_expect_L2. + + That ordering is a precondition, not a check: L_HIER listed out of site + order gives wrong results in the number representation. + """ + + __slots__ = ( + # --- Scalars --- + 'k_max', # Maximum hierarchy depth (int) + 'n_state', # Number of active states (int) + 'n_lop_full', # Number of L-operators (int) + 'M1_modes_per_state', # Bath modes per state, padded (np.ndarray[int]) + 'M1_mode_offset', # Cumulative mode offset (np.ndarray[int]) + 'list_state_list', # Active state indices (list[int]) + 'n_hmodes', # Total hierarchy mode count (int) + 'n_states_full', # Full system state count (int) + 'flag_nearest_neighbor_ham', # bool + 'flag_gs_vacuum', # Ground state is the all-zeros MPS config (bool) + 'mpo_epsilon', # Relative SVD threshold for the bond factors (float) + # --- Mode-dimension operators --- + 'I2_mode', # Mode-space identity (k_max+1, k_max+1) + 'B2_raise', # Raising (creation) operator (k_max+1, k_max+1) + 'B2_lower', # Lowering (annihilation) operator (k_max+1, k_max+1) + 'N2_occ', # Occupation number operator (k_max+1, k_max+1) + 'C1_coupling_raise', # Raise (b^dagger) coupling prefactors (n_modes,) + 'C1_coupling_lower', # Lower (b) coupling prefactors (n_modes,) + 'list_w', # Mode frequencies (np.ndarray) + 'list_L2_coo', # L-operator matrices (list[sparse]) + 'list_L2_masks', # Per-L2 [rows, cols, ix_]; [0][0] is the L2's site + 'list_index_L2_by_hmode', # L2 index of each hierarchy mode (list[int]) + # --- State-dimension operators (4-D core shape) --- + 'T4_plus', # Raise sigma^+ = |1><0| (1, 2, 2, 1) + 'T4_minus', # Lower sigma^- = |0><1| (1, 2, 2, 1) + 'Q4_site', # Unoccupied projector |0><0| (1, 2, 2, 1) + 'P4_site', # Occupied projector |1><1| (1, 2, 2, 1) + 'I2_state', # State-space identity (n_state, n_state) + # --- Hamiltonian --- + 'H2_ham', # Full Hamiltonian (n_state_full, n_state_full) + # Per-bond SVD factors of the coupling graph, built on first use by + # _bond_factors and dropped when the active basis changes. + '_list_bond_factors', # (Y2_left, Z2_coeff, n_bond) per bond (list|None) + ) + + def __init__( + self, + k_max, + n_state, + modes_per_state, + n_lop_full, + ham, + state_list, + mode, + normalization, + n_states_full=None, + flag_nearest_neighbor_ham=False, + flag_gs_vacuum=False, + mpo_epsilon=0.0, + ): + """ + Builds and stores the elementary operators used by all MPO constructors. + + Inputs + ------ + 1. k_max: int + Maximum hierarchy depth. + + 2. n_state: int + Number of active states (system.size). + + 3. modes_per_state: np.ndarray(int) + Number of bath modes per state. + + 4. n_lop_full: int + Number of L-operators (mode.n_l2). + + 5. ham: np.ndarray + Full Hamiltonian matrix. + + 6. state_list: np.ndarray + Active state indices (system.state_list). + + 7. mode: HopsModes + Mode object providing list_g and list_w. + + 8. normalization: str + Scaling convention for the auxiliary vectors. + 'homps' balances the coupling across the raising + and lowering operators for truncation stability; + 'adhops' reproduces the vector HOPS scaling. Sets + B2_raise, B2_lower, C1_coupling_raise and + C1_coupling_lower below. + + 9. n_states_full: int or None + Full system state count (self.n_states_full). + + 10. flag_nearest_neighbor_ham: bool + Whether the Hamiltonian is nearest-neighbor + (default False). + + 11. flag_gs_vacuum: bool + Whether the physical ground state is the + all-zeros MPS configuration (vacuum + convention). Gates the site-0 damp1 widening + in build_statenumber_hierarchy_mpo + (default False). + + 12. mpo_epsilon: float + Relative threshold on the singular values of each + bond's coupling block in _bond_factors, below which + no channel is opened. Zero (default) keeps every + coupling and gives the exact rank; raising it drops + weak couplings and narrows the general generator. + + Returns + ------- + None + """ + if normalization != 'homps' and normalization != 'adhops': + raise ValueError( + f"Unknown normalization '{normalization}'. Use 'homps' or 'adhops'." + ) + + # --- scalars --- + self.k_max = k_max + self.n_state = n_state + self.n_lop_full = n_lop_full + self.M1_modes_per_state = np.asarray(modes_per_state, dtype=int) + self.M1_mode_offset = np.concatenate( + [[0], np.cumsum(self.M1_modes_per_state)] + ) + + # --- Static data from system/mode (used by builder methods) --- + self.list_state_list = list(state_list) + self.list_w = np.asarray(mode.list_w) + self.list_L2_coo = mode.list_L2_coo + self.list_L2_masks = mode.list_L2_masks + self.list_index_L2_by_hmode = mode.list_index_L2_by_hmode + self.n_hmodes = mode.n_hmodes + self.n_states_full = n_states_full + self.flag_nearest_neighbor_ham = flag_nearest_neighbor_ham + self.flag_gs_vacuum = flag_gs_vacuum + self.mpo_epsilon = mpo_epsilon + + # --- mode-dimension operators --- + # Complex dtype so these combine with complex coupling prefactors + # and ladder operators without repeated upcasting. + self.I2_mode = np.eye(k_max + 1, dtype=np.complex128) + self.B2_raise = np.zeros((k_max + 1, k_max + 1), dtype=np.complex128) + self.B2_lower = np.zeros((k_max + 1, k_max + 1), dtype=np.complex128) + self.N2_occ = np.zeros((k_max + 1, k_max + 1), dtype=np.complex128) + + # Ladder operator values for occupation levels 1..k_max + V1_levels = np.arange(1, k_max + 1, dtype=np.complex128) + + if normalization == 'homps': + # sqrt(n+1) ladder operators; coupling split as g/sqrt|g| and sqrt|g| + V1_sqrt_levels = np.sqrt(V1_levels) + self.B2_raise += np.diag(V1_sqrt_levels, k=-1) + self.B2_lower += np.diag(V1_sqrt_levels, k=1) + self.N2_occ = self.B2_raise @ self.B2_lower + # Prefactor convention of Gao et al., Phys. Rev. A 105, L030202 + # (2022), Eq. (10). + # C1_coupling_raise pairs with B2_raise (b^dagger): + # V_m^+ = g_m / sqrt(|g_m|) + # C1_coupling_lower pairs with B2_lower (b): + # V_m^- = sqrt(|g_m|) + # Physical coupling split: + # C1_coupling_raise * C1_coupling_lower = g_m. + list_g = np.asarray(mode.list_g) + abs_g = np.abs(list_g) + sqrt_abs_g = np.sqrt(abs_g) + # Guard against g=0: both prefactors vanish, so all + # bath coupling terms for that mode are zero. + self.C1_coupling_lower = sqrt_abs_g + with np.errstate(invalid='ignore', divide='ignore'): + self.C1_coupling_raise = np.where(abs_g > 0, list_g / sqrt_abs_g, 0.0) + elif normalization == 'adhops': + # unit ladder operators; prefactors w and g/w + self.C1_coupling_raise = mode.list_w + self.C1_coupling_lower = mode.list_g / mode.list_w + self.B2_raise += np.diag(V1_levels, k=-1) + self.B2_lower += np.diag(np.ones(k_max, dtype=np.complex128), k=1) + # N_occ diagonal: levels 1..k_max at positions 1..k_max + self.N2_occ += np.diag( + np.concatenate([[0], V1_levels]), + ) + + # --- State-dimension operators (pre-reshaped to 4-D core shape) --- + # In the binary occupation basis {|0>, |1>}: + # P4_site = |1><1| occupied-site projector + # Q4_site = |0><0| unoccupied-site projector + # T4_plus = |1><0| carries coupling rightward + # T4_minus = |0><1| carries coupling leftward + # .copy() ensures these are independent of the module-level + # constants; without it, reshape returns a view that shares + # memory, and any accidental mutation would silently corrupt + # the constants for all future MpoBuilder instances. + self.T4_plus = _T2_PLUS.reshape(1, 2, 2, 1).copy() + self.T4_minus = _T2_MINUS.reshape(1, 2, 2, 1).copy() + self.Q4_site = _Q2_SITE.reshape(1, 2, 2, 1).copy() + self.P4_site = _P2_SITE.reshape(1, 2, 2, 1).copy() + self.I2_state = np.eye(n_state, dtype=np.complex128) + + # --- Hamiltonian (full, for slicing inside constructors) --- + self.H2_ham = ham + self._list_bond_factors = None + + def refresh_state_data(self, n_state, state_list): + """ + Updates only the state-dependent attributes after an adaptive basis + change. Constant attributes (ladder operators, coupling prefactors, + mode operators) are unchanged. + + Parameters + ---------- + 1. n_state: int + New number of active states. + 2. state_list: list(int) + New active state indices. + + Returns + ------- + None + """ + self.n_state = n_state + self.list_state_list = list(state_list) + self.I2_state = np.eye(n_state, dtype=np.complex128) + self._list_bond_factors = None + + def build_statenumber_hierarchy_mpo( + self, list_z_hat, list_expect_L2, norm_corr, + ): + """ + Builds the hierarchy interaction MPO cores for statenumberrepresentation. + + # TODO: replace with preprint reference and equation numbers + # before merging to MesoHOPS. + + Parameters + ---------- + 1. list_z_hat: np.ndarray + Conjugate noise values at current timestep, shape (n_state,). + + 2. list_expect_L2: np.ndarray + L-operator expectation values, shape (n_state,). + + 3. norm_corr: float + Normalization correction term. + + Returns + ------- + 1. list_cores_op: list(np.ndarray) + MPO cores: one (l_bond, 2, 2, r_bond) per state site, + then one (l_bond, k_max+1, k_max+1, r_bond) per mode + site. + """ + # Guard: statenumber hierarchy MPO supports one L2 per state (bond dim 5). + list_states_seen = set() + for list_mask in self.list_L2_masks: + if list_mask[0][0] in list_states_seen: + raise NotImplementedError( + 'Multiple L-operators per state is not supported ' + 'for number representation.' + ) + list_states_seen.add(list_mask[0][0]) + + k_max = self.k_max + n_state = self.n_state + M1_modes_per_state = self.M1_modes_per_state + M1_mode_offset = self.M1_mode_offset + I2_mode = self.I2_mode + B2_raise = self.B2_raise + B2_lower = self.B2_lower + C1_coupling_raise = self.C1_coupling_raise + C1_coupling_lower = self.C1_coupling_lower + + Q4_site = self.Q4_site + P4_site = self.P4_site + + list_cores_op = [] + shape_mode_core = (1, k_max + 1, k_max + 1, 1) + # Pre-reshape mode-space operators to 4-D core shape to avoid + # repeated inline reshaping throughout the mode core loop. + I4_mode = I2_mode.reshape(shape_mode_core) + B4_raise = B2_raise.reshape(shape_mode_core) + B4_lower = B2_lower.reshape(shape_mode_core) + N4_occ = self.N2_occ.reshape(shape_mode_core) + + # MPO bond index convention (w = 5): + # 0 = identity channel (pass-through) + # 1 = L-operator / noise channel + # 2 = damping channel 1 (-w * N_occ + L^dagger_avg * B_lower) + # 3 = damping channel 2 (multi-site coupling carry) + # 4 = accumulator / output channel + channel_ident = slice(0, 1) + channel_lop = slice(1, 2) + channel_damp1 = slice(2, 3) + channel_damp2 = slice(3, 4) + channel_accum = slice(4, 5) + for site in range(n_state): + if site == 0: + T4_core_site = np.zeros((1, 2, 2, 5), dtype=np.complex128) + else: + T4_core_site = np.zeros((5, 2, 2, 5), dtype=np.complex128) + T4_core_site[channel_ident, :, :, channel_ident] = Q4_site + # Channels lop and damp1 are gated by P_site (occupied + # projector): the L-op, noise, and damping terms are specific to + # the occupied branch. Identity gating would let SVD-compressed + # unoccupied-branch content contaminate the damping channel. + I4_site = P4_site + Q4_site + # The 1j prefactor on hierarchy state cores is required by + # the TDVP integrator: the TDVP sweep applies + # exp(-1j * delta * MPO), so hierarchy terms need 1j so + # that (-1j)(1j * hier) = hier (real damping), while the + # Hamiltonian MPO (no prefactor) gets (-1j)(H) = -iH + # (Schrodinger rotation). For the RK4 path, derivative() + # applies scale_mps(dphi, -1j) to achieve the same result. + T4_core_site[channel_ident, :, :, channel_lop] = 1j * P4_site + # In the vacuum convention the site-0 damp1 opener is widened + # from P_site to I_site = P + Q. The Q-component opens the NL + # feedback channel on the all-zeros MPS configuration so the + # vacuum amplitude correctly tracks the physical + # $\\overline{\\langle L\\rangle}\\,(g_n/w_n)\\,\\Phi[\\text{vac}, k=e_n]$ + # drift. The widening is a no-op algebraically on + # single-excitation configurations (the P-gated downstream + # Q-relay carries no contribution) but perturbs the MPS bond + # singular-value spectrum, so it is gated to the vacuum + # convention to keep non-vacuum truncation accuracy intact. + if site == 0 and self.flag_gs_vacuum: + T4_core_site[channel_ident, :, :, channel_damp1] = ( + 1j * I4_site + ) + else: + T4_core_site[channel_ident, :, :, channel_damp1] = ( + 1j * P4_site + ) + if site != 0: + T4_core_site[channel_damp1, :, :, channel_damp1] = Q4_site + T4_core_site[channel_damp2, :, :, channel_damp2] = Q4_site + T4_core_site[channel_accum, :, :, channel_accum] = Q4_site + T4_core_site[channel_damp2, :, :, channel_accum] = 1j * P4_site + list_cores_op.append(T4_core_site) + + # Mode cores for this state: hierarchy cores indexed via + # M1_mode_offset into list_w / C1_coupling_raise / + # C1_coupling_lower. + n_modes_site = M1_modes_per_state[site] + n_total_modes = int(M1_mode_offset[-1]) + l2_idx = site + for i in range(n_modes_site): + global_idx = M1_mode_offset[site] + i + # Coupling channel (ch 1): + # C1_coupling_raise * b^dagger + # - C1_coupling_lower * b + (z_hat - norm_corr)/M * I + # where C1_coupling_raise = g/sqrt(|g|) and + # C1_coupling_lower = sqrt(|g|) for homps, so that + # b^dagger pairs with V^+ and b pairs with V^-; + # for adhops the pair is w and g/w instead. + M4_coupling = ( + C1_coupling_raise[global_idx] * B4_raise + - C1_coupling_lower[global_idx] * B4_lower + + (list_z_hat[l2_idx] - norm_corr) + * I4_mode / n_modes_site + ) + # Damping channel (ch 2): + # -w_m * N_occ + C1_coupling_lower * * b + # where -w_m * N_occ is the bath frequency decay and + # C1_coupling_lower * * b is the nonlinear feedback. + M4_damping = ( + -self.list_w[global_idx] * N4_occ + + C1_coupling_lower[global_idx] + * np.conj(list_expect_L2[l2_idx]) * B4_lower + ) + if global_idx == n_total_modes - 1: + # Last mode core in the MPS: right bond collapses + # to dim 1, closing all channels to the output. + T4_core_mode = np.zeros( + (5, k_max + 1, k_max + 1, 1), dtype=np.complex128 + ) + # lop→output: coupling closes to output + T4_core_mode[channel_lop, :, :, 0:0+1] = M4_coupling + # accum→output: accumulator closes to output + T4_core_mode[channel_accum, :, :, 0:0+1] = I4_mode + # damp1→output: damping closes to output + T4_core_mode[channel_damp1, :, :, 0:0+1] = M4_damping + else: + # Interior mode core: relay identity, coupling, + # and damping channels through the bond. + T4_core_mode = np.zeros( + (5, k_max + 1, k_max + 1, 5), + dtype=np.complex128, + ) + # ident→ident: identity pass-through + T4_core_mode[channel_ident, :, :, channel_ident] = I4_mode + # accum→accum: accumulator pass-through + T4_core_mode[channel_accum, :, :, channel_accum] = I4_mode + # lop→accum: coupling into accumulator + T4_core_mode[channel_lop, :, :, channel_accum] = M4_coupling + # damp1→accum: damping into accumulator + T4_core_mode[channel_damp1, :, :, channel_accum] = M4_damping + # ident→damp2: damping relay (multi-site carry) + T4_core_mode[channel_ident, :, :, channel_damp2] = M4_damping + # damp1→damp1: damping channel relay + T4_core_mode[channel_damp1, :, :, channel_damp1] = I4_mode + # damp2→damp2: multi-site damping relay + T4_core_mode[channel_damp2, :, :, channel_damp2] = I4_mode + if i != n_modes_site - 1: + # lop→lop: L-op channel stays open for + # remaining modes of this state + T4_core_mode[channel_lop, :, :, channel_lop] = I4_mode + list_cores_op.append(T4_core_mode) + + return list_cores_op + + def build_statenumber_ham_mpo(self): + """Build the Hamiltonian MPO for number representation. + + Dispatches to nearest-neighbor or general builder based on + self.flag_nearest_neighbor_ham. + + Returns + ------- + 1. list_cores_ham: list(np.ndarray) + Hamiltonian MPO cores. + """ + if self.flag_nearest_neighbor_ham: + return self._build_statenumber_ham_nn_mpo() + return self._build_statenumber_ham_general_mpo() + + def _build_statenumber_ham_nn_mpo(self): + """ + Builds Hamiltonian MPO cores for nearest-neighbor Hamiltonians + (number representation). Bond dimension is 4: one channel each + for left transfer, right transfer, identity pass-through, and + accumulated diagonal energy. + + # TODO: replace with preprint reference and equation numbers + # before merging to MesoHOPS. + + Returns + ------- + 1. list_cores_ham_full: list(np.ndarray) + MPO cores interleaved: one state core + (w_l, 2, 2, w_r) followed by modes_per_state + identity mode cores (w, k_max+1, k_max+1, w), + repeating for each physical site. + """ + k_max = self.k_max + n_state = self.n_state + M1_modes_per_state = self.M1_modes_per_state + H2_ham = self.H2_ham + + T4_plus = self.T4_plus + T4_minus = self.T4_minus + Q4_site = self.Q4_site + P4_site = self.P4_site + + # Convention: H2_ham[i, j] couples state j into state i. + # T4_plus = |1><0| raises occupation (carries coupling rightward). + # T4_minus = |0><1| lowers occupation (carries coupling leftward). + + # 4-channel MPO structure for nearest-neighbor Hamiltonians: + # The bond dimension is 4 because a nearest-neighbor Hamiltonian + # decomposes into exactly four channels: + # index 0 -- left transfer channel (T4_plus, carries coupling rightward) + # index 1 -- right transfer channel (T4_minus, carries coupling leftward) + # index 2 -- identity pass-through (Q4_site, site unoccupied) + # index 3 -- accumulated diagonal Hamiltonian (on-site energy * P4_site) + # The first-site core dispatches into these 4 channels (row vector), + # interior cores relay and close channels, and the last-site core + # receives and contracts them to a single output (column vector). + list_cores_ham_full = [] + + # NOTE: `site` is a relative index into the active basis + # (0..n_state-1), not an absolute site label. The absolute site + # index is obtained via self.list_state_list[site] (was system.state_list). + for site in range(n_state): + state = self.list_state_list[site] + if site == 0: + # First site: row vector core (1, 2, 2, 4). + T4_core_site = np.zeros((1, 2, 2, 4), dtype=np.complex128) + # Slice notation i:i+1 preserves the bond axis (returns a + # (1,m,m,1) view); plain integer indexing i would collapse + # it to (m,m), mismatching the 4-D operator shapes. + # Channel 0: left transfer + T4_core_site[0:0+1, :, :, 0:0+1] = T4_plus + # Channel 1: right transfer + T4_core_site[0:0+1, :, :, 1:1+1] = T4_minus + # Channel 2: identity pass-through + T4_core_site[0:0+1, :, :, 2:2+1] = Q4_site + # Channel 3: diagonal on-site energy + T4_core_site[0:0+1, :, :, 3:3+1] = H2_ham[state, state] * P4_site + elif site == n_state - 1: + # Last site: column vector core (4, 2, 2, 1). + T4_core_site = np.zeros((4, 2, 2, 1), dtype=np.complex128) + state_prev = self.list_state_list[site - 1] + # Channel 0: close left-transfer with nn coupling + T4_core_site[0:0+1, :, :, 0:0+1] = H2_ham[state_prev, state] * T4_minus + # Channel 1: close right-transfer with nn coupling + T4_core_site[1:1+1, :, :, 0:0+1] = H2_ham[state, state_prev] * T4_plus + # Channel 2: diagonal on-site energy + T4_core_site[2:2+1, :, :, 0:0+1] = H2_ham[state, state] * P4_site + # Channel 3: close identity pass-through + T4_core_site[3:3+1, :, :, 0:0+1] = Q4_site + else: + # Interior site: full (4, 2, 2, 4) core. + T4_core_site = np.zeros((4, 2, 2, 4), dtype=np.complex128) + state_prev = self.list_state_list[site - 1] + # Channel 0 -> 3: close left-transfer with nn coupling + T4_core_site[0:0+1, :, :, 3:3+1] = H2_ham[state_prev, state] * T4_minus + # Channel 1 -> 3: close right-transfer with nn coupling + T4_core_site[1:1+1, :, :, 3:3+1] = H2_ham[state, state_prev] * T4_plus + # Channel 2 -> 0: reopen left-transfer for next nn pair + T4_core_site[2:2+1, :, :, 0:0+1] = T4_plus + # Channel 2 -> 1: reopen right-transfer for next nn pair + T4_core_site[2:2+1, :, :, 1:1+1] = T4_minus + # Channel 2 -> 2: identity pass-through + T4_core_site[2:2+1, :, :, 2:2+1] = Q4_site + # Channel 2 -> 3: diagonal on-site energy + T4_core_site[2:2+1, :, :, 3:3+1] = H2_ham[state, state] * P4_site + # Channel 3 -> 3: identity relay for accumulated diagonal + T4_core_site[3:3+1, :, :, 3:3+1] = Q4_site + list_cores_ham_full.append(T4_core_site) + + # Mode cores: identity pass-through for all channels. + list_cores_ham_full.extend( + _identity_mode_cores( + T4_core_site.shape[3], + k_max, + M1_modes_per_state[site], + ) + ) + + return list_cores_ham_full + + def _build_statenumber_ham_general_mpo(self): + """ + Builds Hamiltonian MPO cores for general (non-nearest-neighbor) Hamiltonians + (number representation). Delegates to build_statenumber_operator_mpo. + + # TODO: replace with preprint reference and equation numbers + # before merging to MesoHOPS. + + Returns + ------- + 1. list_cores_ham_full: list(np.ndarray) + MPO cores interleaved: one state core then + modes_per_state mode cores, for each state. + """ + return build_statenumber_operator_mpo( + self.H2_ham, self.n_state, self.k_max, self.M1_modes_per_state, + ) + + def _bond_factors(self): + """ + Factorizes the coupling graph at each bond between state sites. + + The bond after state site l carries the hoppings that start at a site + at or left of l and end at a site right of l. Their amplitudes are the + block H[0:l+1, l+1:n] of the off-diagonal Hamiltonian, and the number + of channels the bond needs is that block's rank. A thin SVD splits it + into Y2_left, whose row l says with what weight site l's sigma^+ opens + each channel, and Z2_coeff, whose first column says with what amplitude + each channel closes on the next site. The charge-lowering family reuses + the conjugates of both, since H is Hermitian. + + Taking the left factor with orthonormal columns is what lets the + site-to-site relay be a projection rather than a solve; the width + itself is set by the rank, and any rank factorization would attain it. + + The active basis and the Hamiltonian are fixed between basis changes, + so the factors are built once and reused until refresh_state_data + drops them. + + Returns + ------- + 1. list_bond_factors: list(tuple) + One (Y2_left, Z2_coeff, n_bond) per bond, in site + order, with n_bond = 0 on the last. + """ + if self._list_bond_factors is not None: + return self._list_bond_factors + + # Off-diagonal couplings in the active basis: the same view of the + # Hamiltonian the generator's site cores are built from. + list_state = self.list_state_list + M2_coupling = np.array( + self.H2_ham[np.ix_(list_state, list_state)], dtype=np.complex128, + ) + np.fill_diagonal(M2_coupling, 0.0) + + list_bond_factors = [] + for site in range(self.n_state - 1): + # Rows are the sites at or left of this bond, columns those right. + M2_cross = M2_coupling[:site + 1, site + 1:] + U2_left, V1_sv, V2_right = np.linalg.svd( + M2_cross, full_matrices=False, + ) + # A channel exists only where the singular value clears the + # threshold. mpo_epsilon raises it above the floor, dropping + # couplings the caller accepts losing; the floor is there because + # a structurally zero block still returns values at roundoff. + tol = max(self.mpo_epsilon, 1e-12) + n_bond = int(np.sum( + V1_sv > tol * (V1_sv[0] if V1_sv.size else 0.0) + )) + list_bond_factors.append(( + U2_left[:, :n_bond], + V1_sv[:n_bond, None] * V2_right[:n_bond, :], + n_bond, + )) + # The last state site has nothing to its right, so no channel is open. + list_bond_factors.append((None, None, 0)) + + self._list_bond_factors = list_bond_factors + return list_bond_factors + + def build_general_generator_mpo(self, list_z_hat, list_expect_L2, norm_corr): + """ + Builds the combined Hamiltonian + hierarchy MPO for a Hamiltonian of + any coupling pattern, in the number representation on the + one-excitation manifold. + + Peak bond dimension 3 + 2 * max(n_bond) over the bonds, which is + minimal. A chain or star has coupling blocks of rank one and a ring of + rank two, giving peak width 5 and 7 respectively. + + Parameters + ---------- + 1. list_z_hat: np.ndarray(complex) + Conjugate noise values at the current timestep, indexed + by L-operator. + + 2. list_expect_L2: np.ndarray(complex) + L-operator expectation values, indexed by + L-operator. + + 3. norm_corr: float | complex + Normalization correction term. + + Returns + ------- + 1. list_cores_op: list(np.ndarray) + MPO cores: one (w_l, 2, 2, w_r) per site, each + followed by that site's (w_l, k_max+1, k_max+1, w_r) + mode cores. + """ + n_site = self.n_state + list_state = self.list_state_list + I2_site = _P2_SITE + _Q2_SITE + n_mode_core = int(self.M1_mode_offset[-1]) + list_bond_factors = self._bond_factors() + list_cores_op = [] + + for site in range(n_site): + l2_idx = site + Y2_left, _, n_bond_out = list_bond_factors[site] + n_bond_in = list_bond_factors[site - 1][2] if site > 0 else 0 + # Outgoing cut is identity, the sigma^+/sigma^- pairs still in + # flight, this site's bath coupling, and the accumulator. + w_in = 1 if site == 0 else 2 + 2 * n_bond_in + w_out = 3 + 2 * n_bond_out + # The accumulator is always the last channel, so its index is + # whatever the width happens to be on that side. + idx_accum_in = w_in - 1 + idx_accum_out = w_out - 1 + idx_coupling = 1 + 2 * n_bond_out + + # Indexed [bond_in, occupation', occupation, bond_out]; each + # (bond_in, bond_out) block is a 2x2 operator on the site. + T4_core_site = np.zeros((w_in, 2, 2, w_out), dtype=np.complex128) + T4_core_site[0, :, :, 0] = I2_site + # Site energy plus this site's noise, both diagonal in occupation. + energy = ( + self.H2_ham[list_state[site], list_state[site]] + + 1j * list_z_hat[l2_idx] + ) + T4_core_site[0, :, :, idx_accum_out] = energy * _P2_SITE + if site == 0: + # The norm correction multiplies the identity on the whole + # chain, so it is emitted once here. + T4_core_site[0, :, :, idx_accum_out] += -1j * norm_corr * I2_site + T4_core_site[0, :, :, idx_coupling] = 1j * _P2_SITE + + # Open: this site's sigma^+ enters each outgoing channel with the + # weight carried by its own row of the left factor. + for idx_chan_out in range(n_bond_out): + amp_open = Y2_left[site, idx_chan_out] + T4_core_site[0, :, :, 1 + idx_chan_out] = amp_open * _T2_PLUS + T4_core_site[0, :, :, 1 + n_bond_out + idx_chan_out] = ( + np.conj(amp_open) * _T2_MINUS + ) + + if site > 0: + T4_core_site[idx_accum_in, :, :, idx_accum_out] = I2_site + Y2_left_prev, Z2_coeff_prev, _ = list_bond_factors[site - 1] + # Close: the hoppings ending on this site leave the incoming + # channels, weighted by the first column of the previous + # bond's coefficients. + for idx_chan_in in range(n_bond_in): + amp_close = Z2_coeff_prev[idx_chan_in, 0] + T4_core_site[1 + idx_chan_in, :, :, idx_accum_out] = ( + amp_close * _T2_MINUS + ) + T4_core_site[ + 1 + n_bond_in + idx_chan_in, :, :, idx_accum_out + ] = np.conj(amp_close) * _T2_PLUS + # Relay: channels that outlive this site are re-expressed in + # the outgoing bond's basis. The incoming left factor has + # orthonormal columns, so the change of basis is a projection. + if n_bond_out: + M2_relay = Y2_left_prev.conj().T @ Y2_left[:site, :] + for idx_chan_in in range(n_bond_in): + for idx_chan_out in range(n_bond_out): + amp_relay = M2_relay[idx_chan_in, idx_chan_out] + T4_core_site[ + 1 + idx_chan_in, :, :, 1 + idx_chan_out + ] = amp_relay * I2_site + T4_core_site[ + 1 + n_bond_in + idx_chan_in, :, :, + 1 + n_bond_out + idx_chan_out, + ] = np.conj(amp_relay) * I2_site + list_cores_op.append(T4_core_site) + + # The identity channel emits each mode's damping, the coupling channel + # closes on each mode in turn, and every other channel passes through. + n_mode_site = self.M1_modes_per_state[site] + for p in range(n_mode_site): + idx_mode = self.M1_mode_offset[site] + p + # The bath channel stays open until this site's last mode, so the + # right bond loses it there while the left bond still matches the + # site core's outgoing width. + coupling_out = 1 if p < n_mode_site - 1 else 0 + w_mode_in = 3 + 2 * n_bond_out + w_mode_out = 2 + 2 * n_bond_out + coupling_out + if idx_mode == n_mode_core - 1: + # Final core of the chain contracts to the scalar boundary. + w_mode_out = 1 + T4_core_mode = np.zeros( + (w_mode_in, self.k_max + 1, self.k_max + 1, w_mode_out), + dtype=np.complex128, + ) + if w_mode_out > 1: + T4_core_mode[0, :, :, 0] = self.I2_mode + for idx_bond in range(1, 1 + 2 * n_bond_out): + T4_core_mode[idx_bond, :, :, idx_bond] = self.I2_mode + T4_core_mode[w_mode_in - 1, :, :, w_mode_out - 1] = self.I2_mode + # Bath frequency decay plus the feedback. + T4_core_mode[0, :, :, w_mode_out - 1] += 1j * ( + -self.list_w[idx_mode] * self.N2_occ + + self.C1_coupling_lower[idx_mode] + * np.conj(list_expect_L2[site]) * self.B2_lower + ) + # Hierarchy raise/lower coupling closes the channel here. + T4_core_mode[idx_coupling, :, :, w_mode_out - 1] += ( + self.C1_coupling_raise[idx_mode] * self.B2_raise + - self.C1_coupling_lower[idx_mode] * self.B2_lower + ) + if coupling_out: + T4_core_mode[idx_coupling, :, :, idx_coupling] = self.I2_mode + list_cores_op.append(T4_core_mode) + + return list_cores_op + + def build_fullstate_mpo(self, list_z_hat, list_expect_L2, norm_corr, + C2_lt_corr_hier=None): + """ + Builds combined hierarchy+Hamiltonian MPO cores for fullstate representation. + + # TODO: replace with preprint reference and equation numbers + # before merging to MesoHOPS. + + Parameters + ---------- + 1. list_z_hat: np.ndarray + Conjugate noise values at current timestep, shape (n_lop_full,). + + 2. list_expect_L2: np.ndarray + L-operator expectation values, shape (n_lop_full,). + + 3. norm_corr: float + Normalization correction term. + + 4. C2_lt_corr_hier: np.ndarray or None + LTC hierarchy correction matrix, shape + (n_state, n_state). Added to the Hamiltonian + channel in the state core. + + Returns + ------- + 1. list_cores_op: list(np.ndarray) + MPO cores: one (1, n_state, n_state, bond_dim) state + core followed by n_lop_full*modes_per_state mode cores. + """ + k_max = self.k_max + n_state = self.n_state + n_lop_full = self.n_lop_full + M1_modes_per_state = self.M1_modes_per_state + M1_mode_offset = self.M1_mode_offset + I2_mode = self.I2_mode + I2_state = self.I2_state + + bond_dim = n_lop_full + 2 + state_shape = (1, n_state, n_state, 1) + + # Named MPO channel slices for the state core. + # Using slices keeps the 4-D indexing axis intact (returns + # a (1, d, d, 1) view) and replaces verbose n_lop_full + # arithmetic throughout the method. + channel_input = slice(0, 1) + channel_damp = slice(n_lop_full, n_lop_full + 1) + channel_ham = slice(n_lop_full + 1, n_lop_full + 2) + + # State core block structure (bond_dim = n_lop_full + 2): + # indices 0..n_lop_full-1 -- L-op channels (one per l-op) + # index n_lop_full -- damping/noise channel + # index n_lop_full+1 -- Hamiltonian + norm correction + list_cores_op = [] + T4_core_site = np.zeros( + (1, n_state, n_state, bond_dim), + dtype=np.complex128, + ) + + # Each l-op opens a bond channel carrying 1j * L_state into the mode + # cores, where it will be combined with raise/lower/noise. Channels are + # ordered by chain position so that the mode cores of a state find its + # channel at the head of their left bond. + channel = 0 + for abs_state in self.list_state_list: + if M1_modes_per_state[abs_state] == 0: + continue + l2_idx = self.list_index_L2_by_hmode[M1_mode_offset[abs_state]] + L2_lop_site = self.list_L2_coo[l2_idx] + T4_core_site[0, L2_lop_site.row, L2_lop_site.col, channel] += ( + 1j * L2_lop_site.data + ) + channel += 1 + # Damping channel: identity into mode damping terms + T4_core_site[channel_input, :, :, channel_damp] += ( + (1j) * I2_state.reshape(state_shape) + ) + # Hamiltonian channel: on-site H minus norm correction + H2_block = self.H2_ham - 1j * norm_corr * I2_state + # LTC hierarchy correction: 1j factor matches + # L-op/damping channels so that derivative()'s -1j + # scaling produces the correct sign. + if C2_lt_corr_hier is not None: + H2_block = H2_block + 1j * C2_lt_corr_hier + T4_core_site[channel_input, :, :, channel_ham] += ( + H2_block.reshape(state_shape) + ) + list_cores_op.append(T4_core_site) + + # Mode cores carry the hierarchy coupling, the damping terms, and the + # completed terms. Left bond of a core belonging to a state with + # n_channel_open channels still live: + # index 0 -- this state's own L-op channel + # indices 1..n_channel_open-1 -- later states' L-op channels + # index idx_damp -- damping/noise channel + # index idx_ham -- sink: Hamiltonian, norm correction, + # and every fully applied L-op term + # A channel closes at its own state's last mode core, so the survivors + # shift down one index there and the bond tapers n_lop_full + 2, + # n_lop_full + 1, ..., 3 — minimal at every cut. + # HopsModes sorts list_state_list and list_index_L2_by_hmode into chain + # order, so this loop emits cores left to right along the MPS. + n_channel_open = n_lop_full + for abs_state in self.list_state_list: + n_modes_l2 = M1_modes_per_state[abs_state] + # A state with no hierarchy modes, such as the spectroscopy ground + # state, occupies a slot in the state core and emits no mode cores. + if n_modes_l2 == 0: + continue + mode_base = M1_mode_offset[abs_state] + # Every hierarchy mode of a state carries the same L-operator, so + # any of this state's modes names it. + l2_idx = self.list_index_L2_by_hmode[mode_base] + idx_damp = n_channel_open + idx_ham = n_channel_open + 1 + for i in range(n_modes_l2): + idx_mode = mode_base + i + # This state's channel closes at its last mode core, dropping + # one channel and shifting the survivors down one index. + shift = 1 if i == n_modes_l2 - 1 else 0 + # Right-bond positions of the surviving channels. + idx_damp_out = idx_damp - shift + idx_out = idx_ham - shift + + # L-op channel: raise/lower + noise + M2_coupling = ( + -self.C1_coupling_lower[idx_mode] * self.B2_lower + + self.C1_coupling_raise[idx_mode] * self.B2_raise + + list_z_hat[l2_idx] * self.I2_mode / n_modes_l2 + ) + # Damping: drift + frequency damping + M2_damping = ( + np.conj(list_expect_L2[l2_idx]) + * self.C1_coupling_lower[idx_mode] * self.B2_lower + - self.list_w[idx_mode] * self.N2_occ + ) + if idx_mode == self.n_hmodes - 1: + # Last mode core: contracts all bond channels + # to scalar output (idx_ham + 1, phys, phys, 1) + T4_core_mode = np.zeros( + (idx_ham + 1, k_max + 1, k_max + 1, 1), + dtype=np.complex128, + ) + T4_core_mode[0, :, :, 0] += M2_coupling + T4_core_mode[idx_damp, :, :, 0] += M2_damping + # Hamiltonian: identity pass-through to output + T4_core_mode[idx_ham, :, :, 0] += I2_mode + else: + T4_core_mode = np.zeros( + (idx_ham + 1, k_max + 1, k_max + 1, idx_out + 1), + dtype=np.complex128, + ) + T4_core_mode[0, :, :, idx_out] += M2_coupling + T4_core_mode[idx_damp, :, :, idx_out] += M2_damping + # Damping identity: keep channel open + T4_core_mode[idx_damp, :, :, idx_damp_out] += I2_mode + # Hamiltonian identity: pass through + T4_core_mode[idx_ham, :, :, idx_out] += I2_mode + # L-op identities: pass the channels still in transit + for t in range(shift, n_channel_open): + T4_core_mode[t, :, :, t - shift] += I2_mode + list_cores_op.append(T4_core_mode) + # This state's channel closed at its last mode core. + n_channel_open -= 1 + + return list_cores_op + +# Standalone function (not an MpoBuilder method) because callers +# (apply_system_operator) don't have an MpoBuilder instance and +# only need n_state, k_max, M1_modes_per_state — not the full +# builder configuration. +def build_statenumber_operator_mpo(H2_op, n_state, k_max, M1_modes_per_state): + """ + Build an MPO that applies an arbitrary (n_state x n_state) system operator + to a number representation MPS, acting as identity on all mode cores. + + Uses the same transfer-matrix structure as _build_statenumber_ham_general_mpo: + diagonal elements via site projectors, off-diagonal elements via daisy-chained + transfer matrices through intermediate sites. Bond dimension is + 4 + 2*(n_state - 2) for n_state >= 3, 4 for n_state == 2, and 2 for + n_state == 1. + + Parameters + ---------- + 1. H2_op: np.ndarray(complex) + System-space operator, shape (n_state, n_state). Already trimmed + to the active basis by the caller. + + 2. n_state: int + Number of active system states. + + 3. k_max: int + Maximum hierarchy depth (mode core physical dim = k_max + 1). + + 4. M1_modes_per_state: np.ndarray(int) + Number of bath-mode cores per system-state core. + + Returns + ------- + 1. list_cores_op: list(np.ndarray) + MPO cores interleaved: one state core (w_l, 2, 2, w_r) + then modes_per_state identity mode cores + (w, k_max+1, k_max+1, w), for each state. + """ + # Pre-reshape 2x2 state operators to 4-D core shape once + T4_plus = _T2_PLUS.reshape(1, 2, 2, 1) + T4_minus = _T2_MINUS.reshape(1, 2, 2, 1) + Q4_site = _Q2_SITE.reshape(1, 2, 2, 1) + P4_site = _P2_SITE.reshape(1, 2, 2, 1) + + # Single-site special case: no off-diagonal terms, bond dim 1. + if n_state == 1: + T4_core_site = np.zeros((1, 2, 2, 1), dtype=np.complex128) + T4_core_site[0, :, :, 0] = H2_op[0, 0] * _P2_SITE + _Q2_SITE + list_cores_op = [T4_core_site] + list_cores_op.extend(_identity_mode_cores(1, k_max, M1_modes_per_state[0])) + return list_cores_op + + # Bond dim = 4 base channels (left-transfer, right-transfer, + # identity, diagonal) + 2 long-range relay channels per + # non-nearest-neighbor state pair. + if n_state == 2: + bond_dim = 4 + else: + bond_dim = int(4 + 2 * (n_state - 2)) + + # Base offset for the right-transfer relay block. + # Left-transfer relays: indices 4 .. 4+(n_state-3) + # Right-transfer relays: indices relay_base .. end + relay_base = 4 + (n_state - 2) + + list_cores_op = [] + + for site in range(n_state): + if site == 0: + T4_core_site = np.zeros( + (1, 2, 2, bond_dim), dtype=np.complex128, + ) + if n_state > 1: + T4_core_site[0:1, :, :, 0:1] = T4_plus + T4_core_site[0:1, :, :, 1:2] = T4_minus + T4_core_site[0:1, :, :, 2:3] = Q4_site + T4_core_site[0:1, :, :, 3:4] = ( + H2_op[site, site] * P4_site + ) + + elif site == n_state - 1: + T4_core_site = np.zeros( + (bond_dim, 2, 2, 1), dtype=np.complex128, + ) + T4_core_site[0:1, :, :, 0:1] = ( + H2_op[site - 1, site] * T4_minus + ) + T4_core_site[1:2, :, :, 0:1] = ( + H2_op[site, site - 1] * T4_plus + ) + T4_core_site[2:3, :, :, 0:1] = ( + H2_op[site, site] * P4_site + ) + T4_core_site[3:4, :, :, 0:1] = Q4_site + for i in range(n_state - 2): + site_coupled = site - 2 - i + idx_left_relay = 4 + i + idx_right_relay = relay_base + i + T4_core_site[ + idx_left_relay:idx_left_relay+1, + :, :, 0:1, + ] = H2_op[site_coupled, site] * T4_minus + T4_core_site[ + idx_right_relay:idx_right_relay+1, + :, :, 0:1, + ] = H2_op[site, site_coupled] * T4_plus + + else: + T4_core_site = np.zeros( + (bond_dim, 2, 2, bond_dim), + dtype=np.complex128, + ) + T4_core_site[0:1, :, :, 3:4] = ( + H2_op[site - 1, site] * T4_minus + ) + T4_core_site[1:2, :, :, 3:4] = ( + H2_op[site, site - 1] * T4_plus + ) + T4_core_site[2:3, :, :, 0:1] = T4_plus + T4_core_site[2:3, :, :, 1:2] = T4_minus + T4_core_site[2:3, :, :, 2:3] = Q4_site + T4_core_site[2:3, :, :, 3:4] = ( + H2_op[site, site] * P4_site + ) + T4_core_site[3:4, :, :, 3:4] = Q4_site + T4_core_site[0:1, :, :, 4:5] = Q4_site + T4_core_site[ + 1:2, :, :, relay_base:relay_base+1, + ] = Q4_site + for i in range(n_state - 3): + idx_left_relay = 4 + i + idx_right_relay = relay_base + i + T4_core_site[ + idx_left_relay:idx_left_relay+1, + :, :, + idx_left_relay+1:idx_left_relay+2, + ] = Q4_site + T4_core_site[ + idx_right_relay:idx_right_relay+1, + :, :, + idx_right_relay+1:idx_right_relay+2, + ] = Q4_site + for i in range(site - 1): + site_coupled = site - 2 - i + idx_left_relay = 4 + i + idx_right_relay = relay_base + i + T4_core_site[ + idx_left_relay:idx_left_relay+1, + :, :, 3:4, + ] = H2_op[site_coupled, site] * T4_minus + T4_core_site[ + idx_right_relay:idx_right_relay+1, + :, :, 3:4, + ] = H2_op[site, site_coupled] * T4_plus + + list_cores_op.append(T4_core_site) + + # Mode cores: identity pass-through for all bond channels. + list_cores_op.extend( + _identity_mode_cores(T4_core_site.shape[3], k_max, M1_modes_per_state[site]) + ) + + return list_cores_op + + +def build_statenumber_dipole_mpo( + list_mu, n_state, k_max, M1_modes_per_state, raise_or_lower, +): + """ + Build the MPO for a sum-of-single-site dipole operator (raise or lower) + in the statenumber representation under the ground-state-as-vacuum + convention. Standalone counterpart to + MpoBuilder.build_statenumber_dipole_{raise,lower}_mpo for callers + (e.g.\\ nondyadic_spectroscopy) that do not have a full MpoBuilder + instance. + + Spectroscopically this is the bare transition-dipole operator that + moves amplitude between the ground state and the single-excitation + manifold: raise (a^dagger) is excitation by a photon (absorption), + lower (a) is de-excitation (emission), with no ground-state or + excited-manifold preservation. + + The operator is + sum_k list_mu[k] * sigma_k, + where sigma = a^dagger (T4_plus = |1><0|) for raise_or_lower='raise', + sigma = a (T4_minus = |0><1|) for raise_or_lower='lower'. Sites with + list_mu[k] == 0 contribute identity on their state core and drop from + the sum. Acts as identity on all mode cores. Bond dimension 2 + between state cores; 1 at the chain boundaries. + + Parameters + ---------- + 1. list_mu: np.ndarray(complex) + Per-site dipole amplitudes, shape (n_state,). Entries + set to 0 mark sites excluded from the dipole's site + selection. + + 2. n_state: int + Number of active system states (== number of excited + states in the vacuum convention). + + 3. k_max: int + Maximum hierarchy depth (mode core physical dim = k_max + 1). + + 4. M1_modes_per_state: np.ndarray(int) + Number of bath-mode cores per system-state core. + + 5. raise_or_lower: str + 'raise' for the raise operator, 'lower' for the + lower operator. + + Returns + ------- + 1. list_cores_op: list(np.ndarray) + MPO cores interleaved: one state core + (w_l, 2, 2, w_r) followed by M1_modes_per_state[site] + identity mode cores, for each state. + """ + if raise_or_lower == 'raise': + sigma = _T2_PLUS.reshape(1, 2, 2, 1) + elif raise_or_lower == 'lower': + sigma = _T2_MINUS.reshape(1, 2, 2, 1) + else: + raise ValueError( + f"raise_or_lower must be 'raise' or 'lower', " + f"got {raise_or_lower!r}." + ) + + if len(list_mu) != n_state: + raise ValueError( + f'list_mu must have length n_state = {n_state}, ' + f'got {len(list_mu)}.' + ) + + # State-core identity I_2 = P + Q reshaped as a (1, 2, 2, 1) core. + I4_site = (_P2_SITE + _Q2_SITE).reshape(1, 2, 2, 1) + + list_cores_op = [] + + for site in range(n_state): + mu_site = list_mu[site] + + if n_state == 1: + T4_core_site = np.zeros((1, 2, 2, 1), dtype=np.complex128) + T4_core_site[0:1, :, :, 0:1] = mu_site * sigma + elif site == 0: + # First state core: row vector (1, 2, 2, 2). + # Right-bond channel 0 means the local sigma operator has + # already been placed on this site; channel 1 means it has + # not been placed yet and the MPO should keep propagating. + T4_core_site = np.zeros((1, 2, 2, 2), dtype=np.complex128) + T4_core_site[0:1, :, :, 0:1] = mu_site * sigma + T4_core_site[0:1, :, :, 1:2] = I4_site + elif site == n_state - 1: + # Last state core: column vector (2, 2, 2, 1). + T4_core_site = np.zeros((2, 2, 2, 1), dtype=np.complex128) + T4_core_site[0:1, :, :, 0:1] = I4_site + T4_core_site[1:2, :, :, 0:1] = mu_site * sigma + else: + # Interior state core: (2, 2, 2, 2). + T4_core_site = np.zeros((2, 2, 2, 2), dtype=np.complex128) + T4_core_site[0:1, :, :, 0:1] = I4_site + T4_core_site[1:2, :, :, 0:1] = mu_site * sigma + T4_core_site[1:2, :, :, 1:2] = I4_site + + list_cores_op.append(T4_core_site) + + list_cores_op.extend( + _identity_mode_cores( + T4_core_site.shape[3], k_max, M1_modes_per_state[site], + ) + ) + + return list_cores_op + + +def build_statenumber_dipole_lower_plus_ident_mpo( + list_mu, n_state, k_max, M1_modes_per_state, +): + """ + Build the MPO for the operator + sum_k list_mu[k] * a_k + I_excited, + where I_excited = sum_k |e_k> -> 0 (killed by both terms) + |e_j> -> list_mu[j] |g> + |e_j> + |e_j, e_k> -> 0 (killed by I_excited) + |e_j, e_k, ...> -> 0 (any multiple-excited state is + annihilated by the excited-block + projector) + + Bond dimension 4 = parallel bond-dim-2 mu^- path (sum-of-local + sigma^- terms) and bond-dim-2 I_excited path (sum-of-local P_site + terms with Q_site elsewhere). W-matrix channels: + 0 = mu^- path, sigma^- has not yet fired -> apply I_site, + 1 = mu^- path, sigma^- already fired -> apply I_site, + 2 = I_excited path, P has not yet fired -> apply Q_site, + 3 = I_excited path, P already fired -> apply Q_site to + kill any later multiple-excited occupation. + + Here, "fired" means the corresponding branch has already been + selected on an earlier site in the MPO walk. + + Parameters + ---------- + 1. list_mu: np.ndarray(complex) + Per-site dipole amplitudes, shape (n_state,). Entries + set to 0 mark sites excluded from the lower's site + selection. + + 2. n_state: int + Number of active system states. + + 3. k_max: int + Maximum hierarchy depth (mode core physical dim = k_max + 1). + + 4. M1_modes_per_state: np.ndarray(int) + Number of bath-mode cores per system-state core. + + Returns + ------- + 1. list_cores_op: list(np.ndarray) + MPO cores interleaved: one state core then mode + cores, for each state. + """ + if len(list_mu) != n_state: + raise ValueError( + f'list_mu must have length n_state = {n_state}, ' + f'got {len(list_mu)}.' + ) + + sigma_minus = _T2_MINUS.reshape(1, 2, 2, 1) + P4_site = _P2_SITE.reshape(1, 2, 2, 1) + Q4_site = _Q2_SITE.reshape(1, 2, 2, 1) + I4_site = P4_site + Q4_site + + list_cores_op = [] + + for site in range(n_state): + mu_site = list_mu[site] + + if n_state == 1: + # Single state core: mu_1 * sigma^- + P (=I_excited on a + # single site reduces to projecting onto |1>). + T4_core_site = np.zeros((1, 2, 2, 1), dtype=np.complex128) + T4_core_site[0:1, :, :, 0:1] = ( + mu_site * sigma_minus + P4_site + ) + elif site == 0: + # First state core: row vector (1, 2, 2, 4). + T4_core_site = np.zeros((1, 2, 2, 4), dtype=np.complex128) + # mu^- path + T4_core_site[0:1, :, :, 0:1] = I4_site # I, not yet fired + T4_core_site[0:1, :, :, 1:2] = mu_site * sigma_minus # fire at site 1 + # I_excited path + T4_core_site[0:1, :, :, 2:3] = Q4_site # Q, not yet fired + T4_core_site[0:1, :, :, 3:4] = P4_site # fire P at site 1 + elif site == n_state - 1: + # Last state core: column vector (4, 2, 2, 1). + T4_core_site = np.zeros((4, 2, 2, 1), dtype=np.complex128) + # mu^- path + T4_core_site[0:1, :, :, 0:1] = mu_site * sigma_minus # fire at last + T4_core_site[1:2, :, :, 0:1] = I4_site # already fired -> I + # I_excited path + T4_core_site[2:3, :, :, 0:1] = P4_site # fire P at last + T4_core_site[3:4, :, :, 0:1] = Q4_site # already fired -> Q + else: + # Interior state core: (4, 2, 2, 4). + T4_core_site = np.zeros((4, 2, 2, 4), dtype=np.complex128) + # mu^- path + T4_core_site[0:1, :, :, 0:1] = I4_site # continue not-fired + T4_core_site[0:1, :, :, 1:2] = mu_site * sigma_minus # fire here + T4_core_site[1:2, :, :, 1:2] = I4_site # continue fired + # I_excited path + T4_core_site[2:3, :, :, 2:3] = Q4_site # continue not-fired + T4_core_site[2:3, :, :, 3:4] = P4_site # fire P here + T4_core_site[3:4, :, :, 3:4] = Q4_site # continue fired + + list_cores_op.append(T4_core_site) + + list_cores_op.extend( + _identity_mode_cores( + T4_core_site.shape[3], k_max, M1_modes_per_state[site], + ) + ) + + return list_cores_op + + +def build_statenumber_dipole_raise_plus_ground_ident_mpo( + list_mu, n_state, k_max, M1_modes_per_state, +): + """ + Build the MPO for the operator + sum_k list_mu[k] * a_k^dagger + I_g, + where I_g is the ground-state projector |g><0,...,0|, the projector onto the all-zeros configuration + of the system cores, in the number representation under the + ground-state-as-vacuum convention. + + Spectroscopically this is the absorption excitation-pulse operator: + the raise term promotes the ground state into the single-excitation + manifold while +I_g preserves the ground-state amplitude the NL + denominator needs. + + Action on the natural basis: + |g> = |0,...,0> -> sum_k list_mu[k] |e_k> + |g> + |e_j> -> 0 (killed by both terms) + + Counterpart to build_statenumber_dipole_lower_plus_ident_mpo: the + raise term excites the all-zeros configuration into the single- + excitation manifold; the +I_g term preserves the all-zeros amplitude + after the raise, matching the gs_core dipole-raise's + "(0,0)=1 keep |g>" behavior. Without this preservation, the + NL mean-field denominator (which the EOM augments with + |gs_amp|^2 under the vacuum convention) collapses to zero after + the raise, breaking trajectory-by-trajectory equivalence with + the gs_core path under NL absorption. + + Bond dimension 3 = parallel bond-dim-2 mu^+ path (sum-of-local + sigma^+ terms) and bond-dim-1 I_g path (Q_site at every site, no + before/after firing distinction). W-matrix channels: + 0 = mu^+ path, sigma^+ has not yet fired -> apply I_site, + 1 = mu^+ path, sigma^+ already fired -> apply I_site, + 2 = I_g path -> apply Q_site. + + Parameters + ---------- + 1. list_mu: np.ndarray(complex) + Per-site dipole amplitudes, shape (n_state,). Entries + set to 0 mark sites excluded from the raise's site + selection. + + 2. n_state: int + Number of active system states (excited states). + + 3. k_max: int + Maximum hierarchy depth (mode core physical dim = k_max + 1). + + 4. M1_modes_per_state: np.ndarray(int) + Number of bath-mode cores per system-state core. + + Returns + ------- + 1. list_cores_op: list(np.ndarray) + MPO cores interleaved: one state core then mode + cores, for each state. + """ + if len(list_mu) != n_state: + raise ValueError( + f'list_mu must have length n_state = {n_state}, ' + f'got {len(list_mu)}.' + ) + + sigma_plus = _T2_PLUS.reshape(1, 2, 2, 1) + P4_site = _P2_SITE.reshape(1, 2, 2, 1) + Q4_site = _Q2_SITE.reshape(1, 2, 2, 1) + I4_site = P4_site + Q4_site + + list_cores_op = [] + + for site in range(n_state): + mu_site = list_mu[site] + + if n_state == 1: + # Single state core: mu_1 * sigma^+ + Q (= I_g on one site). + T4_core_site = np.zeros((1, 2, 2, 1), dtype=np.complex128) + T4_core_site[0:1, :, :, 0:1] = ( + mu_site * sigma_plus + Q4_site + ) + elif site == 0: + # First state core: row vector (1, 2, 2, 3). + T4_core_site = np.zeros((1, 2, 2, 3), dtype=np.complex128) + # mu^+ path + T4_core_site[0:1, :, :, 0:1] = I4_site # not yet fired + T4_core_site[0:1, :, :, 1:2] = mu_site * sigma_plus # fire at site 0 + # I_g path + T4_core_site[0:1, :, :, 2:3] = Q4_site # I_g opener + elif site == n_state - 1: + # Last state core: column vector (3, 2, 2, 1). + T4_core_site = np.zeros((3, 2, 2, 1), dtype=np.complex128) + # mu^+ path + T4_core_site[0:1, :, :, 0:1] = mu_site * sigma_plus # fire at last + T4_core_site[1:2, :, :, 0:1] = I4_site # already fired + # I_g path + T4_core_site[2:3, :, :, 0:1] = Q4_site # close I_g + else: + # Interior state core: (3, 2, 2, 3). + T4_core_site = np.zeros((3, 2, 2, 3), dtype=np.complex128) + # mu^+ path + T4_core_site[0:1, :, :, 0:1] = I4_site # continue not-fired + T4_core_site[0:1, :, :, 1:2] = mu_site * sigma_plus # fire here + T4_core_site[1:2, :, :, 1:2] = I4_site # continue fired + # I_g path + T4_core_site[2:3, :, :, 2:3] = Q4_site # continue I_g + + list_cores_op.append(T4_core_site) + + list_cores_op.extend( + _identity_mode_cores( + T4_core_site.shape[3], k_max, M1_modes_per_state[site], + ) + ) + + return list_cores_op diff --git a/src/mesohops/tensor/tdvp.py b/src/mesohops/tensor/tdvp.py new file mode 100644 index 0000000..55f5de8 --- /dev/null +++ b/src/mesohops/tensor/tdvp.py @@ -0,0 +1,1364 @@ +"""TDVP time-stepping for MPS wavefunctions following Paeckel et al. (2019). + +All functions operate on bare lists of np.ndarray — no MPS container. +MPS cores have shape (Dl, d, Dr); MPO cores have shape (wl, d_out, d_in, wr); +environment tensors have shape (chi, w, chi). + +State convention +---------------- +The MPS is held in mixed canonical form with the non-canonical center +at site 1: + core_M : ndarray(Dl, d, Dr) — center core (site 1) + list_cores_B : list(ndarray(Dl, d, Dr)) — right-canonical cores, + sites 2..L + L0 : ndarray(1, 1, 1) — left boundary environment + list_envs_R : list(ndarray(chi, w, chi)) — right environments + R_2..R_{L+1}; list_envs_R[0]=R_2, + list_envs_R[j-2]=R_j, list_envs_R[L-1]=R_{L+1} + +Public API +---------- +initialize(list_cores_mps, list_cores_mpo) + → core_M, list_cores_B, L0, list_envs_R + +timestep(delta, L0, list_envs_R, list_cores_mpo, core_M, list_cores_B, + method='1tdvp', solver='arnoldi', **kwargs) + → core_M, list_cores_B, L0, list_envs_R +""" + +from __future__ import annotations + +import math +from collections.abc import Callable + +import numpy as np +from scipy.integrate import solve_ivp +from scipy.linalg import expm +from scipy.linalg import svd as scipy_svd + +__title__ = 'TDVP Integrator' +__author__ = 'B. Z. Citty' +__maintainer__ = 'B. Z. Citty' + +# Hard cap on Krylov subspace dimension; convergence (conv_tol) is expected to +# terminate the iteration long before this limit is reached. +_KRYLOV_DIM_MAX = 60 + +# Below this threshold h_{m+1,m} is treated as zero (prevents division by +# zero when extending the Krylov basis; not exposed to callers). +_KRYLOV_BREAKDOWN_TOL = 1e-14 + + +# ============================================================ +# Section 1: Splitting utilities +# ============================================================ + + +def _split_qr(core: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Left-orthogonalize a rank-3 MPS core via QR decomposition. + + Reshapes the core from (Dl, d, Dr) into a matrix (Dl*d, Dr), performs + QR, then reshapes Q back to (Dl, d, chi). The result satisfies + sum_s Q[:,s,:]^dag Q[:,s,:] = I (left-orthogonality condition). + + Parameters + ---------- + 1. core : np.ndarray(Dl, d, Dr) + + Returns + ------- + 1. Q : np.ndarray(Dl, d, chi) — left-orthogonal + 2. R : np.ndarray(chi, Dr) + """ + Dl, d, Dr = core.shape + # merge physical and left-bond indices: (Dl*d, Dr) + H2_core = core.reshape(Dl * d, Dr) + Q, R = np.linalg.qr(H2_core, mode='reduced') + chi = Q.shape[1] + # restore physical index: (Dl, d, chi) + Q = Q.reshape(Dl, d, chi) + return Q, R + + +def _split_rq(core: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Right-orthogonalize a rank-3 MPS core via RQ decomposition. + + Reshapes the core from (Dl, d, Dr) into a matrix (Dl, d*Dr), performs + QR on the transpose, then reshapes Q back to (chi, d, Dr). The result + satisfies sum_s Q[:,s,:] Q[:,s,:]^dag = I (right-orthogonality condition). + + Parameters + ---------- + 1. core : np.ndarray(Dl, d, Dr) + + Returns + ------- + 1. R : np.ndarray(Dl, chi) + 2. Q : np.ndarray(chi, d, Dr) — right-orthogonal + """ + Dl, d, Dr = core.shape + # merge physical and right-bond indices: (Dl, d*Dr) + H2_core = core.reshape(Dl, d * Dr) + # RQ via QR of transpose: H2_core = R @ Q ↔ H2_core^T = Q^T @ R^T + Qt, Rt = np.linalg.qr(H2_core.T.conj(), mode='reduced') + R = Rt.T.conj() # shape (Dl, chi) + Q = Qt.T.conj() # shape (chi, d*Dr) + chi = Q.shape[0] + # restore physical index: (chi, d, Dr) + Q = Q.reshape(chi, d, Dr) + return R, Q + + +def _split_svd( + theta: np.ndarray, + chi_max: int, + eps: float, + d1: int | None = None, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Truncated SVD of a two-site tensor. Returns raw factors; callers absorb S. + + Reshapes theta from (Dl, d1*d2, Dr) into (Dl*d1, d2*Dr), performs full SVD, + applies truncation by eps and chi_max, then reshapes U and Vt back to rank-3 + cores. + + Parameters + ---------- + 1. theta : np.ndarray(Dl, d1*d2, Dr) + 2. chi_max : int — maximum bond dimension (0 = unlimited) + 3. eps : float — singular value truncation threshold + 4. d1 : int | None + Physical dimension of the left site. If None, assumes + d1 == d2 and infers d1 = sqrt(d1*d2). + + Returns + ------- + 1. U : np.ndarray(Dl, d1, chi_new) — left-orthogonal + 2. S : np.ndarray(chi_new,) — singular values (not absorbed here) + 3. Vt : np.ndarray(chi_new, d2, Dr) — right-orthogonal + 4. chi_new : int + + Notes + ----- + Callers are responsible for absorbing S into whichever side becomes the new + center tensor: + - sweep_right_2tdvp: C_j = S[:,None] * Vt (S absorbed right) + - sweep_left_2tdvp: C_{j-1} = U * S[None,:] (S absorbed left) + """ + Dl, d1d2, Dr = theta.shape + if d1 is None: + d1 = round(math.sqrt(d1d2)) + if d1 * d1 != d1d2: + raise ValueError( + '_split_svd: d1 not provided and middle axis is not a perfect ' + f'square. Got d1*d2={d1d2}.' + ) + d2 = d1d2 // d1 + # bipartition: (Dl, d1, d2, Dr) treated as (Dl*d1, d2*Dr) + H2_theta = theta.reshape(Dl * d1, d2 * Dr) + U, S, Vt = scipy_svd(H2_theta, full_matrices=False, lapack_driver='gesvd') + # truncate by threshold + mask = S > eps + if chi_max > 0: + # keep at most chi_max values + mask[chi_max:] = False + chi_new = int(mask.sum()) + chi_new = max(chi_new, 1) # always keep at least one singular value + U = U[:, :chi_new] # (Dl*d1, chi_new) + S = S[:chi_new] # (chi_new,) + Vt = Vt[:chi_new, :] # (chi_new, d2*Dr) + # reshape to rank-3 cores + U = U.reshape(Dl, d1, chi_new) # (Dl, d1, chi_new) — left-orthogonal + Vt = Vt.reshape(chi_new, d2, Dr) # (chi_new, d2, Dr) — right-orthogonal + return U, S, Vt, chi_new + + +# ============================================================ +# Section 2: Environment contractions (Paeckel Alg. 3) +# ============================================================ + + +def contract_left(L_prev: np.ndarray, W: np.ndarray, A: np.ndarray) -> np.ndarray: + """Paeckel Alg. 3 CONTRACT-LEFT. Extend left environment one site to the right. + + Computes the partial sandwich L_j = contracted with L_{j-1}: + + L[b', w', b] = sum_{a', w, a, s', s} + L_prev[a', w, a] * A*.conj()[a', s', b'] + * W[w, s', s, w'] * A[a, s, b] + + Einsum index map: + L_prev[i, j, k] — i=a'(chi_l_bra), j=w_l, k=a(chi_l_ket) + A.conj()[i, l, m] — i=a'(chi_l_bra), l=s'(d_bra), m=b'(chi_r_bra) + W[j, l, n, o] — j=w_l, l=s'(d_out), n=s(d_in), o=w'(w_r) + A[k, n, p] — k=a(chi_l_ket), n=s(d_in), p=b(chi_r_ket) + → L[m, o, p] — m=b'(chi_r_bra), o=w'(w_r), p=b(chi_r_ket) + + Parameters + ---------- + 1. L_prev : np.ndarray(chi_l, w_l, chi_l) + 2. W : np.ndarray(w_l, d_out, d_in, w_r) — MPO core at site j + 3. A : np.ndarray(Dl, d, Dr) — left-canonical MPS core at site j + + Returns + ------- + 1. L : np.ndarray(chi_r, w_r, chi_r) + """ + return np.einsum('ijk,ilm,jlno,knp->mop', L_prev, A.conj(), W, A, + optimize=True) + + +def contract_right(R_next: np.ndarray, W: np.ndarray, B: np.ndarray) -> np.ndarray: + """Paeckel Alg. 3 CONTRACT-RIGHT. Extend right environment one site to the left. + + Computes the partial sandwich R_j = contracted with R_{j+1}: + + R[l, n, p] = sum_{i, j, k, m, o} + R_next[i, j, k] * B*.conj()[l, m, i] + * W[n, m, o, j] * B[p, o, k] + + Einsum index map: + R_next[i, j, k] — i=b'(chi_r_bra), j=w_r, k=b(chi_r_ket) + B.conj()[l, m, i] — l=a'(chi_l_bra), m=s'(d_bra), i=b'(chi_r_bra) + W[n, m, o, j] — n=w'(w_l), m=s'(d_out), o=s(d_in), j=w_r + B[p, o, k] — p=a(chi_l_ket), o=s(d_in), k=b(chi_r_ket) + → R[l, n, p] — l=a'(chi_l_bra), n=w'(w_l), p=a(chi_l_ket) + + Parameters + ---------- + 1. R_next : np.ndarray(chi_r, w_r, chi_r) + 2. W : np.ndarray(w_l, d_out, d_in, w_r) — MPO core at site j + 3. B : np.ndarray(Dl, d, Dr) — right-canonical MPS core at site j + + Returns + ------- + 1. R : np.ndarray(chi_l, w_l, chi_l) + """ + return np.einsum('ijk,lmi,nmoj,pok->lnp', R_next, B.conj(), W, B, + optimize=True) + + +# ============================================================ +# Section 3: Effective Hamiltonian actions +# ============================================================ + + +def _apply_heff_site( + M: np.ndarray, L: np.ndarray, W: np.ndarray, R: np.ndarray, +) -> np.ndarray: + """One-site effective Hamiltonian action: H_j^eff |M⟩ = L · W · R |M⟩. + + Implements the tensor contraction that appears in Paeckel Alg. 5 TIMESTEP + (single-site forward/backward updates). The effective Hamiltonian acts on + the center core M as + + HM[i, l, o] = sum_{j,k,m,n,p} + L[i, j, k] * W[j, l, m, n] * R[o, n, p] * M[k, m, p] + + Einsum index map: + L[i, j, k] — i=a'(chi_l_bra), j=w_l, k=a(chi_l_ket) + W[j, l, m, n] — j=w_l, l=s'(d_out), m=s(d_in), n=w_r + R[o, n, p] — o=b'(chi_r_bra), n=w_r, p=b(chi_r_ket) + M[k, m, p] — k=a(chi_l_ket), m=s(d_in), p=b(chi_r_ket) + → HM[i, l, o] — i=a', l=s', o=b' + + Parameters + ---------- + 1. M : np.ndarray(Dl, d, Dr) + 2. L : np.ndarray(chi_l, w_l, chi_l) + 3. W : np.ndarray(w_l, d_out, d_in, w_r) + 4. R : np.ndarray(chi_r, w_r, chi_r) + + Returns + ------- + 1. HM : np.ndarray(Dl, d, Dr) + """ + return np.einsum('ijk,jlmn,onp,kmp->ilo', L, W, R, M, optimize=True) + + +def _apply_heff_bond(C: np.ndarray, L: np.ndarray, R: np.ndarray) -> np.ndarray: + """Zero-site (bond) effective Hamiltonian action: H_{j+}^eff |C⟩ = L · R |C⟩. + + Used for the backward bond step in 1TDVP (Paeckel Alg. 5, Step 6). The + zero-site effective Hamiltonian has no MPO core — the environments are + contracted directly over the bond matrix C: + + HC[i, l] = sum_{j, k, m} L[i, j, k] * R[l, j, m] * C[k, m] + + The sum over the MPO bond j collapses because both environments already + encode the full MPO up to their respective boundaries. + + Einsum index map: + L[i, j, k] — i=a'(chi_l_bra), j=w, k=a(chi_l_ket) + R[l, j, m] — l=b'(chi_r_bra), j=w, m=b(chi_r_ket) + C[k, m] — k=a(chi_l_ket), m=b(chi_r_ket) + → HC[i, l] — i=a', l=b' + + Parameters + ---------- + 1. C : np.ndarray(chi_l, chi_r) + 2. L : np.ndarray(chi_l, w, chi_l) + 3. R : np.ndarray(chi_r, w, chi_r) + + Returns + ------- + 1. HC : np.ndarray(chi_l, chi_r) + """ + return np.einsum('ijk,ljm,km->il', L, R, C, optimize=True) + + +def _apply_heff_twosite( + theta: np.ndarray, + L: np.ndarray, + W1: np.ndarray, + W2: np.ndarray, + R: np.ndarray, +) -> np.ndarray: + """Two-site effective Hamiltonian action: H_{j,j+1}^eff |θ⟩ = L · W1 · W2 · R |θ⟩. + + Used for the forward two-site step in 2TDVP (Paeckel Alg. 6). theta is + passed as (Dl, d*d, Dr) and reshaped to (Dl, d, d, Dr) internally: + + Htheta[i, l, o, r] = sum_{j,k,m,n,p,q,s} + L[i,j,k] * W1[j,l,m,n] * W2[n,o,p,q] + * R[r,q,s] * th[k,m,p,s] + + Einsum index map: + L[i, j, k] — i=a', j=w_l, k=a + W1[j, l, m, n] — j=w_l, l=s1', m=s1, n=w_m + W2[n, o, p, q] — n=w_m, o=s2', p=s2, q=w_r + R[r, q, s] — r=b', q=w_r, s=b + th[k, m, p, s] — k=a, m=s1, p=s2, s=b + → Htheta[i,l,o,r] + + Parameters + ---------- + 1. theta : np.ndarray(Dl, d1*d2, Dr) + 2. L : np.ndarray(chi_l, w_l, chi_l) + 3. W1 : np.ndarray(w_l, d_out, d_in, w_m) — left MPO core + 4. W2 : np.ndarray(w_m, d_out, d_in, w_r) — right MPO core + 5. R : np.ndarray(chi_r, w_r, chi_r) + + Returns + ------- + 1. Htheta : np.ndarray(Dl, d1*d2, Dr) + """ + Dl, d1d2, Dr = theta.shape + d1 = W1.shape[1] # physical dim of left site (from MPO core) + d2 = W2.shape[1] # physical dim of right site (from MPO core) + theta_4d = theta.reshape(Dl, d1, d2, Dr) + result_4d = np.einsum( + 'ijk,jlmn,nopq,rqs,kmps->ilor', + L, + W1, + W2, + R, + theta_4d, + optimize=True, + ) + return result_4d.reshape(Dl, d1 * d2, Dr) + + +# ============================================================ +# Section 4: Local solvers +# ============================================================ + + +def _arnoldi_expm(v, matvec, dt, conv_tol): + """Krylov-Arnoldi matrix exponential for general (non-Hermitian) H. + + Computes exp(dt * H) |v⟩ without forming H explicitly. The algorithm: + + 1. Build an orthonormal Krylov basis V_m = [v_1, ..., v_m] via Arnoldi + iteration (modified Gram-Schmidt), recording the upper Hessenberg + matrix H_m such that H V_m ≈ V_m H_m + h_{m+1,m} v_{m+1} e_m^T. + 2. After each step, check the a posteriori relative error estimate + (Hochbruck & Lubich 1997, §2.3): + err_rel ≈ h_{m+1,m} · |[exp(dt H_m) e_1]_m| + and terminate as soon as err_rel < conv_tol. expm on the m × m + Hessenberg is negligible compared to the cost of a matvec. + 3. Reconstruct the result: + exp(dt * H) |v⟩ ≈ ‖v‖ · V_m · exp(dt * H_m) · e_1 + + Reference: Hochbruck & Lubich (1997) SIAM J. Numer. Anal. 34(5). + + Parameters + ---------- + 1. v : np.ndarray — state to evolve (any shape; flattened internally) + 2. matvec : callable(v) → same shape as v + 3. dt : complex — time step (may be negative for backward evolution) + 4. conv_tol : float — relative convergence tolerance on the output vector + + Returns + ------- + 1. v_out : np.ndarray, same shape as v + """ + shape = v.shape + v_flat = v.ravel().astype(complex) + n = len(v_flat) + m = min(_KRYLOV_DIM_MAX, n) + + beta = np.linalg.norm(v_flat) + if beta < _KRYLOV_BREAKDOWN_TOL: + return np.zeros_like(v) + + # H2_krylov: columns are Krylov basis vectors; H2_hess: upper Hessenberg + H2_krylov = np.zeros((n, m + 1), dtype=complex) + H2_hess = np.zeros((m + 1, m), dtype=complex) + H2_krylov[:, 0] = v_flat / beta + + m_actual = m + for j in range(m): + w = matvec(H2_krylov[:, j].reshape(shape)).ravel() + # modified Gram-Schmidt orthogonalization + for i in range(j + 1): + H2_hess[i, j] = np.dot(H2_krylov[:, i].conj(), w) + w = w - H2_hess[i, j] * H2_krylov[:, i] + H2_hess[j + 1, j] = np.linalg.norm(w) + # Convergence check: relative error ≈ h_{j+2,j+1} · |φ_j| + # where φ = expm(dt H_{j+1}) e_1. Also catches happy breakdown + # (h ≈ 0) without a separate code path. + V1_unit = np.zeros(j + 1, dtype=complex) + V1_unit[0] = 1.0 + phi = expm(dt * H2_hess[:j + 1, :j + 1]) @ V1_unit + if H2_hess[j + 1, j] * abs(phi[j]) < conv_tol: + m_actual = j + 1 + return (beta * H2_krylov[:, :m_actual] @ phi).reshape(shape) + if H2_hess[j + 1, j] < _KRYLOV_BREAKDOWN_TOL: + # exact happy breakdown not caught by conv_tol (e.g. conv_tol=0); + # guard against division by zero before storing next basis vector + m_actual = j + 1 + break + H2_krylov[:, j + 1] = w / H2_hess[j + 1, j] + + # exponentiate the (m_actual × m_actual) Hessenberg matrix + H2_hessenberg = H2_hess[:m_actual, :m_actual] + V1_unit = np.zeros(m_actual, dtype=complex) + V1_unit[0] = 1.0 + # exp(dt * H2_hessenberg) e_1 → first column of matrix exponential + V1_expm = expm(dt * H2_hessenberg) @ V1_unit + v_out = beta * H2_krylov[:, :m_actual] @ V1_expm + return v_out.reshape(shape) + + +def _lanczos_expm(v, matvec, dt, conv_tol): + """Krylov-Lanczos matrix exponential for Hermitian H. + + Computes exp(dt * H) |v⟩ for Hermitian H without forming H explicitly. + Uses the three-term Lanczos recurrence to build a tridiagonal T_m and + an orthonormal basis Q_m, then reconstructs the result as: + + exp(dt * H) |v⟩ ≈ ‖v‖ · Q_m · exp(dt · T_m) · e_1 + + For Hermitian H the Lanczos recurrence is cheaper than Arnoldi (two inner + products per step instead of j+1), and T_m is real symmetric when H is + Hermitian, so `expm` operates on a real matrix. + + Convergence is checked after each step via the same relative estimate as + _arnoldi_expm: β_{j+1} · |[exp(dt T_{j+1}) e_1]_j| < conv_tol. + + Reference: Park & Light (1986) J. Chem. Phys. 85(10), 5870. + + Parameters + ---------- + 1. v : np.ndarray + 2. matvec : callable(v) → same shape as v + 3. dt : complex + 4. conv_tol : float — relative convergence tolerance on the output vector + + Returns + ------- + 1. v_out : np.ndarray, same shape as v + """ + shape = v.shape + v_flat = v.ravel().astype(complex) + n = len(v_flat) + m = min(_KRYLOV_DIM_MAX, n) + + beta = np.linalg.norm(v_flat) + if beta < _KRYLOV_BREAKDOWN_TOL: + return np.zeros_like(v) + + # alpha: diagonal of T_m; beta_vec: off-diagonal (beta_vec[j] = T[j+1,j]) + alpha = np.zeros(m, dtype=complex) + beta_vec = np.zeros(m, dtype=complex) + # Q: Lanczos basis vectors stored as columns + Q = np.zeros((n, m + 1), dtype=complex) + Q[:, 0] = v_flat / beta + + m_actual = m + v_prev = np.zeros(n, dtype=complex) + for j in range(m): + w = matvec(Q[:, j].reshape(shape)).ravel() + alpha[j] = np.dot(Q[:, j].conj(), w).real + w = w - alpha[j] * Q[:, j] - (beta_vec[j - 1] if j > 0 else 0.0) * v_prev + beta_j = np.linalg.norm(w) + beta_vec[j] = beta_j + # Convergence check: relative error ≈ β_{j+1} · |φ_j| + # where φ = expm(dt T_{j+1}) e_1. + H2_tridiag_cur = ( + np.diag(alpha[:j + 1]) + + np.diag(beta_vec[:j].real, 1) + + np.diag(beta_vec[:j].real, -1) + ) + V1_unit = np.zeros(j + 1, dtype=complex) + V1_unit[0] = 1.0 + phi = expm(dt * H2_tridiag_cur) @ V1_unit + if beta_j * abs(phi[j]) < conv_tol: + m_actual = j + 1 + return (beta * Q[:, :m_actual] @ phi).reshape(shape) + if beta_j < _KRYLOV_BREAKDOWN_TOL: + m_actual = j + 1 + break + v_prev = Q[:, j].copy() + Q[:, j + 1] = w / beta_j + + # build tridiagonal matrix T_m (m_actual × m_actual) + H2_tridiag = ( + np.diag(alpha[:m_actual]) + + np.diag(beta_vec[: m_actual - 1].real, 1) + + np.diag(beta_vec[: m_actual - 1].real, -1) + ) + V1_unit = np.zeros(m_actual, dtype=complex) + V1_unit[0] = 1.0 + V1_expm = expm(dt * H2_tridiag) @ V1_unit + v_out = beta * Q[:, :m_actual] @ V1_expm + return v_out.reshape(shape) + + +def _ivp_solve( + v: np.ndarray, + matvec: Callable, + dt: float | complex, + **ivp_kwargs, +) -> np.ndarray: + """Evolve v under y' = matvec(y) via scipy.integrate.solve_ivp. + + Wraps scipy ODE integration for the local exponential problem dv/dt = H v. + dt may be negative (backward TDVP evolution); in that case the rhs is + negated and integrated forward over |dt|. + + Default solver is 'BDF' (stiff), suitable for the large-norm effective + Hamiltonians that arise in HOPS. Pass method='RK45' for non-stiff problems. + + Parameters + ---------- + 1. v : np.ndarray + 2. matvec : callable(v) → same shape as v + 3. dt : float + 4. **ivp_kwargs passed to solve_ivp (method, rtol, atol, max_step) + + Returns + ------- + 1. v_out : np.ndarray, same shape as v + """ + shape = v.shape + V1_init = v.ravel().astype(complex) + sign = 1.0 if dt >= 0 else -1.0 + t_span = (0.0, abs(dt)) + + def rhs(_t, y): + # sign accounts for backward integration (negate rhs when dt < 0) + return sign * matvec(y.reshape(shape)).ravel() + + method = ivp_kwargs.pop('method', 'BDF') + rtol = ivp_kwargs.pop('rtol', 1e-7) + atol = ivp_kwargs.pop('atol', 1e-9) + max_step = ivp_kwargs.pop('max_step', np.inf) + + sol = solve_ivp( + rhs, + t_span, + V1_init, + method=method, + rtol=rtol, + atol=atol, + max_step=max_step, + dense_output=False, + **ivp_kwargs, + ) + return sol.y[:, -1].reshape(shape) + + +def _solve_local( + v: np.ndarray, + matvec: Callable, + dt: float | complex, + solver: str = 'arnoldi', + **kwargs, +) -> np.ndarray: + """Dispatch local exponential solve: compute exp(dt * H_eff) |v⟩. + + Selects among three backends: + - 'arnoldi' (default): Krylov-Arnoldi, suitable for non-Hermitian H. + kwargs: conv_tol (float, default 1e-6) — relative convergence tolerance. + - 'lanczos': Krylov-Lanczos, more efficient for Hermitian H. + kwargs: conv_tol (float, default 1e-6) — relative convergence tolerance. + - 'ivp': scipy.integrate.solve_ivp, fully adaptive step control. + kwargs: method (str), rtol (float), atol (float), max_step (float). + + The matvec callable encodes H_eff (not -i/hbar * H_eff); the factor -i/hbar + is absorbed into dt by the caller (i.e. dt = -i*delta/hbar for forward + evolution and dt = +i*delta/hbar for the backward bond step). + + Parameters + ---------- + 1. v : np.ndarray + 2. matvec : callable(v) → same shape as v + 3. dt : float or complex + 4. solver : {'arnoldi', 'lanczos', 'ivp'} + 5. **kwargs forwarded to the selected solver + + Returns + ------- + 1. v_out : np.ndarray, same shape as v + """ + if solver == 'arnoldi': + conv_tol = kwargs.pop('conv_tol', 1e-6) + return _arnoldi_expm(v, matvec, dt, conv_tol) + elif solver == 'lanczos': + conv_tol = kwargs.pop('conv_tol', 1e-6) + return _lanczos_expm(v, matvec, dt, conv_tol) + elif solver == 'ivp': + return _ivp_solve(v, matvec, dt, **kwargs) + else: + raise ValueError( + f"Unknown solver '{solver}'. Choose 'arnoldi', 'lanczos', or 'ivp'." + ) + + +# ============================================================ +# Section 5: Initialization and timestep (Paeckel Alg. 3) +# ============================================================ + + +def recenter_to_zero( + tensors: list[np.ndarray], + normalize: bool = False, + copy: bool = True, +) -> list[np.ndarray]: + """ + Put an MPS into mixed canonical form with the orthogonality center at site 0. + + Sweeps left-to-right through sites 1..L-1, QR-factorizing each core to make + it left-canonical (Q^dagger Q = I on the reshaped (Dl*d, Dr) matrix). The + residual R factor is absorbed into the next site, with the final scalar + folded into A[0]. Assumes open boundary conditions (Dr of the last site == 1). + + Parameters + ---------- + 1. tensors: list(np.ndarray) + MPS cores, each shaped (Dl, d, Dr). + 2. normalize: bool + If True, rescale so that ||state|| = 1. + 3. copy: bool + If True (default), copy the cores before modifying. If False, + reuse the input list and replace its elements with the + canonicalized cores (original array objects are not mutated). + + Returns + ------- + 1. list_cores_centered: list(np.ndarray) + The canonicalized MPS cores with orthogonality center at site 0. + """ + if len(tensors) == 0: + return [] if copy else tensors + + list_cores = [t.copy() for t in tensors] if copy else tensors + n_cores = len(list_cores) + + # Validate open boundary condition on the last core + _, _, bond_dim_right_last = list_cores[-1].shape + if bond_dim_right_last != 1: + raise ValueError( + f'Expected open boundaries with last right bond = 1, ' + f'got Dr[{n_cores - 1}] = {bond_dim_right_last}.' + ) + + # Validate all internal bond dimensions before sweeping + for n in range(n_cores - 1): + bond_dim_right_n = list_cores[n].shape[2] + bond_dim_left_next = list_cores[n + 1].shape[0] + if bond_dim_right_n != bond_dim_left_next: + raise ValueError( + f'Bond mismatch at link {n}: ' + f'right dim of site {n} is {bond_dim_right_n}, ' + f'but left dim of site {n + 1} is {bond_dim_left_next}.' + ) + + # Sweep left-to-right through sites 1..L-1, QR-factorizing each core + # so that Q replaces the core (making it left-canonical) and R captures + # the non-orthogonal content. R is then absorbed into the neighboring + # core to preserve the overall state. After the full sweep, all content + # that couldn't be made orthogonal has been pushed into site 0, which + # becomes the orthogonality center. + for n in range(1, n_cores): + bond_dim_left, phys_dim, bond_dim_right = list_cores[n].shape + H2_core_mat = list_cores[n].reshape( + bond_dim_left * phys_dim, bond_dim_right, + ) + Q, R = np.linalg.qr(H2_core_mat, mode='reduced') + rank_qr = Q.shape[1] + list_cores[n] = Q.reshape(bond_dim_left, phys_dim, rank_qr) + + if n < n_cores - 1: + # R holds the part of this core that couldn't be made orthogonal. + # Absorb it into the next core so the product Q @ R @ next is + # unchanged, preserving the overall state. + (bond_dim_left_next, phys_dim_next, bond_dim_right_next) = ( + list_cores[n + 1].shape + ) + H2_core_next_mat = list_cores[n + 1].reshape( + bond_dim_left_next, + phys_dim_next * bond_dim_right_next, + ) + H2_core_next_mat = R @ H2_core_next_mat + list_cores[n + 1] = H2_core_next_mat.reshape( + rank_qr, + phys_dim_next, + bond_dim_right_next, + ) + else: + # At the boundary under OBC the right bond is 1, so reduced QR + # yields R as a (1,1) scalar. All remaining content is folded + # into site 0 — the orthogonality center. + list_cores[0] *= R[0, 0] + + # Optional normalization + # (with n>=1 left-canonical, ||state|| = ||list_cores[0]||_F) + if normalize: + norm_phi = np.linalg.norm(list_cores[0].ravel()) + if norm_phi > 0: + list_cores[0] /= norm_phi + + return list_cores + + +def initialize( + list_cores_mps: list[np.ndarray], + list_cores_mpo: list[np.ndarray], +) -> tuple[ + np.ndarray, list[np.ndarray], np.ndarray, list[np.ndarray] +]: + """Paeckel Alg. 3 INITIALIZE. + + Right-normalizes the MPS and builds all right-environment tensors needed + before the first timestep. The input MPS is first recentered to site 0 + (mixed canonical form) for numerical stability before right-normalization. + + Procedure: + 1. Recenter the MPS to site 0 via a left-to-right QR sweep. + 2. Right-sweep from site L down to site 2: apply _split_rq to each core, + absorb R into the core to the left, producing right-canonical B cores. + 3. Site 1 is left as the non-canonical center core_M (absorbs the last R). + 4. Build right environments by sweeping from right to left using contract_right. + The boundary environment R_{L+1} = [[1]] (shape (1,1,1)). + list_envs_R[j-2] = R_j for j in 2..L+1: + list_envs_R[0] = R_2 (at bond between sites 1 and 2) + ... + list_envs_R[L-1] = R_{L+1} (right boundary, all ones) + + Parameters + ---------- + 1. list_cores_mps : list(np.ndarray(Dl, d, Dr)) — L MPS cores + 2. list_cores_mpo : list(np.ndarray(wl, d', d, wr)) — L MPO cores + + Returns + ------- + 1. core_M : np.ndarray — non-canonical center core at site 1 + 2. list_cores_B : list(np.ndarray) — right-canonical cores, sites 2..L + 3. L0 : np.ndarray(1, 1, 1) — left boundary environment + 4. list_envs_R : list(np.ndarray) — right environments R_2..R_{L+1} + (L elements; list_envs_R[0]=R_2, list_envs_R[L-1]=R_{L+1}) + """ + n_sites = len(list_cores_mps) + # Recenter to site 0 for numerical stability (copy=True preserves input) + list_cores = recenter_to_zero(list_cores_mps, copy=True) + + # Right-normalize: sweep from last site down to site 1 + for j in range(n_sites - 1, 0, -1): + R_mat, Q = _split_rq(list_cores[j]) + # Q is right-canonical; R_mat absorbed into core to the left + list_cores[j] = Q + list_cores[j - 1] = np.einsum( + 'ijk,kl->ijl', + list_cores[j - 1], + R_mat, + ) + + core_M = list_cores[0] + list_cores_B = list_cores[1:] + + # Build right environments: R_{L+1}, R_L, ..., R_2 + # list_envs_R[j] corresponds to R_{j+2} (0-indexed j → site j+2) + R_boundary = np.ones((1, 1, 1), dtype=complex) # R_{L+1} + list_envs_R = [None] * n_sites + list_envs_R[n_sites - 1] = R_boundary + + # sweep right-to-left building R environments + R_cur = R_boundary + for j in range(n_sites - 1, 0, -1): + # site j (0-indexed) → site j+1 (1-indexed) + W_j = list_cores_mpo[j] + B_j = list_cores_B[j - 1] + R_cur = contract_right(R_cur, W_j, B_j) + list_envs_R[j - 1] = R_cur + + # left boundary + L0 = np.ones((1, 1, 1), dtype=complex) + + return core_M, list_cores_B, L0, list_envs_R + + +def timestep( + delta: float | complex, + L0: np.ndarray, + list_envs_R: list[np.ndarray], + list_cores_mpo: list[np.ndarray], + core_M: np.ndarray, + list_cores_B: list[np.ndarray], + method: str = '1tdvp', + solver: str = 'arnoldi', + **kwargs, +) -> tuple[np.ndarray, list[np.ndarray], np.ndarray, list[np.ndarray]]: + """Paeckel Alg. 3 TIMESTEP. Strang splitting: sweep_right(δ/2) + sweep_left(δ/2). + + The Strang splitting ensures second-order accuracy in the time step δ: + + exp(δ H) ≈ sweep_right(δ/2) ∘ sweep_left(δ/2) + + Each half-sweep is itself a sequence of local exponentials (the TDVP + integrator); the composition cancels the leading-order splitting error. + + After the right half-sweep the state is in right-canonical form with center + at site L. The left half-sweep restores it to the initial canonical form + (center at site 1) with updated environments. + + Parameters + ---------- + 1. delta : float + 2. L0 : np.ndarray(1, 1, 1) + 3. list_envs_R : list(np.ndarray) — R_2..R_{L+1} + 4. list_cores_mpo : list(np.ndarray) + 5. core_M : np.ndarray — center MPS core at site 1 + 6. list_cores_B : list(np.ndarray) — right-canonical cores, sites 2..L + 7. method : {'1tdvp', '2tdvp'} + 8. solver : {'arnoldi', 'lanczos', 'ivp'} + + Returns + ------- + 1. core_M : np.ndarray — updated center core at site 1 + 2. list_cores_B : list(np.ndarray) — updated right-canonical cores + 3. L0 : np.ndarray — left boundary (unchanged) + 4. list_envs_R : list(np.ndarray) — updated right environments + """ + half = delta / 2.0 + + if method == '1tdvp': + # right half-sweep: returns (list_envs_L, list_cores_A, core_M_last) + list_envs_L, list_cores_A, core_M_last = sweep_right_1tdvp( + half, + L0, + list_envs_R, + list_cores_mpo, + core_M, + list_cores_B, + solver=solver, + **kwargs, + ) + # boundary R_{L+1} is the last element of list_envs_R (unchanged by right sweep) + R_boundary = list_envs_R[-1] + # left half-sweep: returns (core_M_first, list_cores_B, list_envs_R) + core_M, list_cores_B, list_envs_R = sweep_left_1tdvp( + half, + list_envs_L, + R_boundary, + list_cores_mpo, + list_cores_A, + core_M_last, + solver=solver, + **kwargs, + ) + + elif method == '2tdvp': + chi_max = kwargs.pop('chi_max', 0) + eps = kwargs.pop('eps', 0.0) + list_envs_L, list_cores_A, core_M_last = sweep_right_2tdvp( + half, + L0, + list_envs_R, + list_cores_mpo, + core_M, + list_cores_B, + chi_max, + eps, + solver=solver, + **kwargs, + ) + R_boundary = list_envs_R[-1] + core_M, list_cores_B, list_envs_R = sweep_left_2tdvp( + half, + list_envs_L, + R_boundary, + list_cores_mpo, + list_cores_A, + core_M_last, + chi_max, + eps, + solver=solver, + **kwargs, + ) + + else: + raise ValueError(f"Unknown method '{method}'. Choose '1tdvp' or '2tdvp'.") + + return core_M, list_cores_B, L0, list_envs_R + + +# ============================================================ +# Section 6: 1TDVP sweeps (Paeckel Alg. 5) +# ============================================================ + + +def sweep_right_1tdvp( + delta: float | complex, + L0: np.ndarray, + list_envs_R: list[np.ndarray], + list_cores_mpo: list[np.ndarray], + core_M: np.ndarray, + list_cores_B: list[np.ndarray], + solver: str = 'arnoldi', + **kwargs, +) -> tuple[list[np.ndarray], list[np.ndarray], np.ndarray, np.ndarray]: + """Paeckel Alg. 5 SWEEP-RIGHT for 1TDVP. + + Implements the right-to-left half of the Strang-split 1TDVP integrator. + For each site j = 1..L the algorithm: + + 1. Forward-evolves the center core M_j under H_j^eff for time delta: + M_j ← exp(-i * delta * H_j^eff) |M_j⟩ + (The factor -i is absorbed into dt = -1j * delta passed to _solve_local; + H_j^eff is encoded via the matvec closure using _apply_heff_site.) + + 2. QR-decomposes M_j → A_j (left-canonical) + C_j (non-unitary remainder). + + 3. Updates L_j ← contract_left(L_{j-1}, W_j, A_j). + + 4. For j < L: backward-evolves the bond matrix C_j for time -delta: + C_j ← exp(+i * delta * H_{j+}^eff) |C_j⟩ + This backward step compensates for the gauge choice (Paeckel eq. 46). + Then absorbs M_{j+1} = C_j · B_{j+1} for the next site. + + At the end the state is in mixed canonical form with center at site L. + The left environments L_0..L_{L-1} and left-canonical cores A_1..A_{L-1} + are returned for use by sweep_left_1tdvp. + + Parameters + ---------- + 1. delta : float + 2. L0 : np.ndarray(1, 1, 1) + 3. list_envs_R : list(np.ndarray) — R_2..R_{L+1} + 4. list_cores_mpo : list(np.ndarray) + 5. core_M : np.ndarray — center core at site 1 + 6. list_cores_B : list(np.ndarray) — right-canonical cores, sites 2..L + 7. solver : {'arnoldi', 'lanczos', 'ivp'} + + Returns + ------- + 1. list_envs_L : list(np.ndarray) — L_0..L_{L-1} (L elements) + 2. list_cores_A : list(np.ndarray) — left-canonical cores A_1..A_{L-1} + 3. core_M_last : np.ndarray — center core at site L + """ + n_sites = len(list_cores_mpo) + # list_envs_R[j] = R_{j+2}: + # list_envs_R[0] = R_2 (right of site 1) + # list_envs_R[j-1] = R_{j+1} (right of site j, 1-indexed) + # We build list_envs_L[j] = L_{j+1} (left of site j+1, 1-indexed): + # list_envs_L[0] = L_0 (left boundary) + # list_envs_L[j] = L_j (left of site j+1) + + list_envs_L = [None] * n_sites + list_cores_A = [None] * (n_sites - 1) + + list_envs_L[0] = L0 + M_cur = core_M # current center core + + for j in range(n_sites): + # site index j (0-indexed) → site j+1 (1-indexed, Paeckel notation) + L_cur = list_envs_L[j] # L_{j} (left env to the left of site j+1) + # R_{j+2} = R_{(j+1)+1} (right env to right of site j+1) + R_cur = list_envs_R[j] + W_cur = list_cores_mpo[j] + + # Step 1: forward-evolve M_j under H_j^eff + # dt = -1j * delta (forward TDVP: dM/dt = -i H_eff M) + def matvec_site(v, _L=L_cur, _W=W_cur, _R=R_cur): + return _apply_heff_site(v, _L, _W, _R) + + M_cur = _solve_local(M_cur, matvec_site, -1j * delta, solver=solver, **kwargs) + + # Step 2: QR decompose M_j → A_j (left-canonical) + C_j + A_j, C_j = _split_qr(M_cur) # A_j: (Dl, d, chi), C_j: (chi, Dr) + + # Step 3: update left environment L_j ← contract_left(L_{j-1}, W_j, A_j) + L_new = contract_left(L_cur, W_cur, A_j) + + if j < n_sites - 1: + list_cores_A[j] = A_j + + # Step 4: backward-evolve C_j under H_{j+}^eff (bond Hamiltonian) + # dt = +1j * delta (backward step cancels gauge artifact) + def matvec_bond(v, _L=L_new, _R=R_cur): + return _apply_heff_bond(v, _L, _R) + + C_j = _solve_local(C_j, matvec_bond, +1j * delta, solver=solver, **kwargs) + + # Absorb C_j into next site: M_{j+1} = C_j · B_{j+1} + B_next = list_cores_B[j] # shape (chi_mid, d, Dr) + M_cur = np.einsum('ij,jkl->ikl', C_j, B_next) + + list_envs_L[j + 1] = L_new + else: + # last site; M_cur is the final center core + core_M_last = M_cur + + return list_envs_L, list_cores_A, core_M_last + + +def sweep_left_1tdvp( + delta: float | complex, + list_envs_L: list[np.ndarray], + R_boundary: np.ndarray, + list_cores_mpo: list[np.ndarray], + list_cores_A: list[np.ndarray], + core_M_last: np.ndarray, + solver: str = 'arnoldi', + **kwargs, +) -> tuple[np.ndarray, list[np.ndarray], np.ndarray, list[np.ndarray]]: + """Paeckel Alg. 5 SWEEP-LEFT for 1TDVP. + + Implements the left-to-right half of the Strang-split 1TDVP integrator. + Traverses sites j = L..1, restoring the canonical form with center at site 1 + and rebuilding the right environments list_envs_R. + + For each site j = L..1: + + 1. Forward-evolves M_j under H_j^eff for time delta: + M_j ← exp(-i * delta * H_j^eff) |M_j⟩ + + 2. RQ-decomposes M_j → C_{j-1} (non-unitary remainder) + B_j (right-canonical). + + 3. Updates R_j ← contract_right(R_{j+1}, W_j, B_j). + + 4. For j > 1: backward-evolves C_{j-1} for time -delta: + C_{j-1} ← exp(+i * delta * H_{j-}^eff) |C_{j-1}⟩ + Then absorbs M_{j-1} = A_{j-1} · C_{j-1}. + + After the sweep the state has center at site 1 with core_M_first, and + list_envs_R is ready for the next timestep. + + Parameters + ---------- + 1. delta : float + 2. list_envs_L : list(np.ndarray) — L_0..L_{L-1} (L elements) + 3. R_boundary : np.ndarray — R_{L+1} + 4. list_cores_mpo : list(np.ndarray) + 5. list_cores_A : list(np.ndarray) — left-canonical cores A_1..A_{L-1} + 6. core_M_last : np.ndarray — center core at site L + 7. solver : {'arnoldi', 'lanczos', 'ivp'} + + Returns + ------- + 1. core_M_first : np.ndarray — center core at site 1 + 2. list_cores_B : list(np.ndarray) — right-canonical cores B_2..B_L + 3. list_envs_R : list(np.ndarray) — R_2..R_{L+1} + """ + n_sites = len(list_cores_mpo) + list_envs_R = [None] * n_sites + list_cores_B = [None] * (n_sites - 1) + + list_envs_R[n_sites - 1] = R_boundary # R_{L+1} + M_cur = core_M_last # start at last site + + for j in range(n_sites - 1, -1, -1): + # site index j (0-indexed) → site j+1 (1-indexed) + L_cur = list_envs_L[j] + R_cur = list_envs_R[j] # R_{j+2} (right env to right of site j+1) + W_cur = list_cores_mpo[j] + + # Step 1: forward-evolve M_j under H_j^eff + def matvec_site(v, _L=L_cur, _W=W_cur, _R=R_cur): + return _apply_heff_site(v, _L, _W, _R) + + M_cur = _solve_local(M_cur, matvec_site, -1j * delta, solver=solver, **kwargs) + + # Step 2: RQ decompose M_j → C_{j-1} + B_j (right-canonical) + C_j, B_j = _split_rq(M_cur) # C_j: (Dl, chi), B_j: (chi, d, Dr) + + # Step 3: update right environment R_j ← contract_right(R_{j+1}, W_j, B_j) + R_new = contract_right(R_cur, W_cur, B_j) + + if j > 0: + # store B_j in the right-canonical list (0-indexed: B at site j+1) + list_cores_B[j - 1] = B_j # B at site j+1 (1-indexed); slot j-1 + + # Step 4: backward-evolve C_j under H_{j-}^eff + def matvec_bond(v, _L=L_cur, _R=R_new): + return _apply_heff_bond(v, _L, _R) + + C_j = _solve_local(C_j, matvec_bond, +1j * delta, solver=solver, **kwargs) + + # Absorb: M_{j-1} = A_{j-1} · C_j + # list_cores_A[j-1] = A at site j (0-indexed) = A_{j} (1-indexed) + A_prev = list_cores_A[j - 1] # shape (Dl, d, chi_l) + M_cur = np.einsum('ijk,kl->ijl', A_prev, C_j) + # A_prev: (Dl, d, chi_mid), C_j: (chi_mid, Dr) → M_cur: (Dl, d, Dr) + + list_envs_R[j - 1] = R_new # R_{j+1} goes in slot j-1 + else: + # j == 0: site 1 is the new center; M_cur (forward-evolved) becomes + # core_M_first directly. No RQ decomposition is needed or correct here: + # the decomposed B_j would be right-canonical but list_cores_B[0] is + # already right-canonical from the j==1 iteration above. + # list_envs_R[0] = R_2 was already stored when j==1; do NOT overwrite. + core_M_first = M_cur + + return core_M_first, list_cores_B, list_envs_R + + +# ============================================================ +# Section 7: 2TDVP sweeps (Paeckel Alg. 6) +# ============================================================ + + +def sweep_right_2tdvp( + delta: float | complex, + L0: np.ndarray, + list_envs_R: list[np.ndarray], + list_cores_mpo: list[np.ndarray], + core_M: np.ndarray, + list_cores_B: list[np.ndarray], + chi_max: int, + eps: float, + solver: str = 'arnoldi', + **kwargs, +) -> tuple[list[np.ndarray], list[np.ndarray], np.ndarray, np.ndarray]: + """Paeckel Alg. 6 SWEEP-RIGHT for 2TDVP. + + Implements the right half of the Strang-split 2TDVP integrator. Unlike + 1TDVP, the two-site algorithm operates on merged two-site tensors theta, + which allows bond dimension to grow before being truncated by SVD. This + enables dynamical adaptation of the bond dimension during time evolution. + + For each pair of sites j, j+1 with j = 1..L-1: + + 1. Contract the two-site tensor: + theta_{j,j+1}[a, s1, s2, b] = M_j[a, s1, c] * B_{j+1}[c, s2, b] + Then reshape to (Dl, d*d, Dr) for the two-site effective Hamiltonian. + + 2. Forward-evolve theta under H_{j,j+1}^eff for time delta: + theta ← exp(-i * delta * H_{j,j+1}^eff) |theta⟩ + + 3. Truncated SVD → A_j (left-canonical), S, Vt: + C_j = diag(S) · Vt (center tensor with bond weights absorbed right) + + 4. For j < L-1: + - Update L_j ← contract_left(L_{j-1}, W_j, A_j) + - Backward single-site evolution of the center C_j under H_{j+1}^eff: + C_j ← exp(+i * delta * H_{j+1}^eff) |C_j⟩ + This backward step exactly inverts the two-site forward step's + contribution from site j+1, leaving only the net effect on site j. + - M_{j+1} = C_j for the next iteration (center shifts right). + + Parameters + ---------- + 1. delta : float + 2. L0 : np.ndarray(1, 1, 1) + 3. list_envs_R : list(np.ndarray) — R_2..R_{L+1} + 4. list_cores_mpo : list(np.ndarray) + 5. core_M : np.ndarray — center core at site 1 + 6. list_cores_B : list(np.ndarray) — right-canonical cores, sites 2..L + 7. chi_max : int — maximum bond dimension (0 = unlimited) + 8. eps : float — SVD truncation threshold + 9. solver : {'arnoldi', 'lanczos', 'ivp'} + + Returns + ------- + 1. list_envs_L : list(np.ndarray) — L_0..L_{L-2} (L-1 elements) + 2. list_cores_A : list(np.ndarray) — left-canonical cores A_1..A_{L-1} + 3. core_M_last : np.ndarray — center core at site L + """ + n_sites = len(list_cores_mpo) + # list_envs_L[j] = L_j (left env for site j+1, 1-indexed) + list_envs_L = [None] * (n_sites - 1) + list_cores_A = [None] * (n_sites - 1) + + list_envs_L[0] = L0 + M_cur = core_M + + for j in range(n_sites - 1): + # site j (0-indexed) = site j+1 (1-indexed) + # site j+1 (0-indexed) = site j+2 (1-indexed) + L_cur = list_envs_L[j] # L_{j} (left of site j+1) + R_next = list_envs_R[j + 1] # R_{j+3} = R_{(j+2)+1} — right of site j+2 + W_j = list_cores_mpo[j] + W_jp1 = list_cores_mpo[j + 1] + + # Step 1: contract two-site tensor + # M_cur: (Dl, d, chi_mid), B_{j+1}: (chi_mid, d, Dr) + B_next = list_cores_B[j] # B at site j+2 (0-indexed j → site j+2) + # theta shape (Dl, d, d, Dr) → reshape to (Dl, d*d, Dr) + theta = np.einsum('ijk,klm->ijlm', M_cur, B_next) + Dl, d1, d2, Dr = theta.shape + theta = theta.reshape(Dl, d1 * d2, Dr) + + # Step 2: forward-evolve theta under H_{j,j+1}^eff + def matvec_two(v, _L=L_cur, _W1=W_j, _W2=W_jp1, _R=R_next): + return _apply_heff_twosite(v, _L, _W1, _W2, _R) + + theta = _solve_local(theta, matvec_two, -1j * delta, solver=solver, **kwargs) + + # Step 3: truncated SVD + A_j, S, Vt, _chi = _split_svd(theta, chi_max, eps, d1=d1) + # bond weight absorbed right: C_j = S[:,None,None] * Vt + # shape (chi_new, d2, Dr) + C_j = S[:, None, None] * Vt + + list_cores_A[j] = A_j + + if j < n_sites - 2: + # Step 4: update left environment + L_new = contract_left(L_cur, W_j, A_j) + list_envs_L[j + 1] = L_new + + # Backward single-site evolution of C_j (the new center at site j+2) + # Uses L_j (just built) and R_{j+2} = list_envs_R[j+1] + def matvec_back(v, _L=L_new, _W=W_jp1, _R=R_next): + return _apply_heff_site(v, _L, _W, _R) + + C_j = _solve_local(C_j, matvec_back, +1j * delta, solver=solver, **kwargs) + + M_cur = C_j # center core shifts to site j+2 + + core_M_last = M_cur + return list_envs_L, list_cores_A, core_M_last + + +def sweep_left_2tdvp( + delta: float | complex, + list_envs_L: list[np.ndarray], + R_boundary: np.ndarray, + list_cores_mpo: list[np.ndarray], + list_cores_A: list[np.ndarray], + core_M_last: np.ndarray, + chi_max: int, + eps: float, + solver: str = 'arnoldi', + **kwargs, +) -> tuple[np.ndarray, list[np.ndarray], np.ndarray, list[np.ndarray]]: + """Paeckel Alg. 6 SWEEP-LEFT for 2TDVP. + + Implements the left half of the Strang-split 2TDVP integrator. Traverses + pairs of sites j-1, j for j = L..2, restoring canonical form with center + at site 1 and rebuilding right environments. + + For each pair j-1, j with j = L..2: + + 1. Contract the two-site tensor: + theta_{j-1,j}[a, s1, s2, b] = A_{j-1}[a, s1, c] * M_j[c, s2, b] + Then reshape to (Dl, d*d, Dr). + + 2. Forward-evolve theta under H_{j-1,j}^eff for time delta. + + 3. Truncated SVD → U, S, Vt: + C_{j-1} = U * S[None,None,:] (bond weight absorbed left) + B_j = Vt (right-canonical) + + 4. For j > 2: + - Update R_j ← contract_right(R_{j+1}, W_j, B_j) + - Backward single-site evolution of C_{j-1} under H_{j-1}^eff. + - M_{j-1} = C_{j-1} (center shifts left). + + At the end the center is at site 1. + + Parameters + ---------- + 1. delta : float + 2. list_envs_L : list(np.ndarray) — L_0..L_{L-2} (L-1 elements) + 3. R_boundary : np.ndarray — R_{L+1} + 4. list_cores_mpo : list(np.ndarray) + 5. list_cores_A : list(np.ndarray) — left-canonical cores A_1..A_{L-1} + 6. core_M_last : np.ndarray — center core at site L + 7. chi_max : int + 8. eps : float + 9. solver : {'arnoldi', 'lanczos', 'ivp'} + + Returns + ------- + 1. core_M_first : np.ndarray — center core at site 1 + 2. list_cores_B : list(np.ndarray) — right-canonical cores B_2..B_L + 3. list_envs_R : list(np.ndarray) — R_2..R_{L+1} + """ + n_sites = len(list_cores_mpo) + list_envs_R = [None] * n_sites + list_cores_B = [None] * (n_sites - 1) + + list_envs_R[n_sites - 1] = R_boundary # R_{L+1} + M_cur = core_M_last + + for j in range(n_sites - 1, 0, -1): + # site j (0-indexed) = site j+1 (1-indexed) + # site j-1 (0-indexed) = site j (1-indexed) + L_prev = list_envs_L[j - 1] # L_{j-1} (left of site j, 1-indexed) + R_cur = list_envs_R[j] # R_{j+2} (right of site j+1, 1-indexed) + W_j = list_cores_mpo[j] + W_jm1 = list_cores_mpo[j - 1] + + # Step 1: contract two-site tensor + # A_{j-1}: (Dl, d, chi_mid), M_cur: (chi_mid, d, Dr) + A_prev = list_cores_A[j - 1] # A at site j (1-indexed) + theta = np.einsum('ijk,klm->ijlm', A_prev, M_cur) + Dl, d1, d2, Dr = theta.shape + theta = theta.reshape(Dl, d1 * d2, Dr) + + # Step 2: forward-evolve theta under H_{j-1,j}^eff + def matvec_two(v, _L=L_prev, _W1=W_jm1, _W2=W_j, _R=R_cur): + return _apply_heff_twosite(v, _L, _W1, _W2, _R) + + theta = _solve_local(theta, matvec_two, -1j * delta, solver=solver, **kwargs) + + # Step 3: truncated SVD + U, S, B_j, _chi = _split_svd(theta, chi_max, eps, d1=d1) + # bond weight absorbed left: U is (Dl, d1, chi_new), + # S is (chi_new,); U * S[None, None, :] → (Dl, d, chi_new) + C_jm1 = U * S[None, None, :] # shape (Dl, d, chi_new) + + list_cores_B[j - 1] = B_j # B at site j+1 (1-indexed) in slot j-1 + + if j > 1: + # Step 4: update right environment + R_new = contract_right(R_cur, W_j, B_j) + list_envs_R[j - 1] = R_new # R_{j+1} in slot j-1 + + # Backward single-site evolution of C_{j-1} under H_{j-1}^eff + def matvec_back(v, _L=L_prev, _W=W_jm1, _R=R_new): + return _apply_heff_site(v, _L, _W, _R) + + C_jm1 = _solve_local( + C_jm1, + matvec_back, + +1j * delta, + solver=solver, + **kwargs, + ) + + M_cur = C_jm1 # center shifts left to site j (1-indexed) + + core_M_first = M_cur + # R_2 = list_envs_R[0]; must have been set when j==1 in the loop above. + # When j==1 we skip the "if j > 1" block, so list_envs_R[0] is still None. + # Build R_2 from B at site 2 (list_cores_B[0]) and R_3 = list_envs_R[1]: + if list_envs_R[0] is None and n_sites > 1: + R2 = contract_right(list_envs_R[1], list_cores_mpo[1], list_cores_B[0]) + list_envs_R[0] = R2 + + return core_M_first, list_cores_B, list_envs_R diff --git a/src/mesohops/tensor/tensor_eom_functions.py b/src/mesohops/tensor/tensor_eom_functions.py new file mode 100644 index 0000000..e9171b3 --- /dev/null +++ b/src/mesohops/tensor/tensor_eom_functions.py @@ -0,0 +1,406 @@ +""" +EOM helper functions for HopsTensorEOM. + +These routines sit between the pure MPS algebra (tensor_operations.py) and the +physics layer (hops_tensor_eom.py). They handle MPO-MPS contraction and the +normalization correction factor. + +Functions +--------- +tensor_matvec_prod(list_cores_vec, list_cores_mpo, epsilon, bond_dim_max) + Apply an MPO to an MPS and compress the result via SVD. + +calc_norm_corr_tensor(wavefunction, psi, z_hat, list_avg_L2, mode, + list_index_L2_by_mode) + Compute the normalization correction factor needed for propagating + the normalized nonlinear wave function. + +apply_system_operator(list_cores_phi, O2_op_trimmed, method, k_max, + M1_modes_per_state, epsilon, bond_dim_max) + Apply a pre-trimmed system-space operator to an MPS. +""" + +from __future__ import annotations + +import numpy as np + +from mesohops.basis.hops_modes import HopsModes +from mesohops.tensor.hops_tensor_wavefunction import HopsTensorWavefunction +from mesohops.tensor.mpo_constructors import build_statenumber_operator_mpo +from mesohops.util.exceptions import UnsupportedRequest +from mesohops.util.tensor_operations import ( + calc_mps_complexity, + flatten_cores, + phi_aux, + tensor_compress, + unflatten_cores, +) + +__title__ = 'Tensor EOM Functions' +__author__ = 'B. Z. Citty' +__maintainer__ = 'B. Z. Citty' + + +def tensor_matvec_prod( + list_cores_vec: list[np.ndarray], + list_cores_mpo: list[np.ndarray], + epsilon: float, + bond_dim_max: int, +) -> tuple[list[np.ndarray], int]: + """ + Performs the matrix-vector product in MPS form, then compresses the result. + + Parameters + ---------- + 1. list_cores_vec: list(np.ndarray) + MPS cores of the vector, each shaped + (dim_left, dim_phys, dim_right). + 2. list_cores_mpo: list(np.ndarray) + MPO cores of the operator, each shaped + (dim_left, dim_out, dim_in, dim_right). + 3. epsilon: float + SVD truncation threshold for compression. + 4. bond_dim_max: int + Maximum bond dimension after compression. + + Returns + ------- + 1. list_cores_compressed: list(np.ndarray) + Compressed MPS cores of the result. + 2. complexity: int + Scalar complexity proxy (see calc_mps_complexity) of the + uncompressed contracted MPS, i.e. its peak size before + SVD truncation. + """ + if len(list_cores_vec) == 0: + raise ValueError( + 'tensor_matvec_prod requires at least one core; ' + 'got empty list_cores_vec.' + ) + # Detect nested statenumber MPS (list/tuple-of-list/tuple) and flatten + # before contraction. Both list and tuple are accepted for the outer + # and inner containers so callers aren't forced to wrap tuples; the + # non-nested branch keeps identifying flat input by the ndarray type + # of the first element. + is_nested = len(list_cores_vec) > 0 and isinstance( + list_cores_vec[0], (list, tuple) + ) + if is_nested: + # Each group has 1 state core + N mode cores; subtract 1 to get mode count. + list_modes_per_site = [len(g) - 1 for g in list_cores_vec] + list_cores_vec = flatten_cores(list_cores_vec) + + if len(list_cores_mpo) != len(list_cores_vec): + raise ValueError('MPO and MPS must have the same number of cores.') + + list_cores_compressed = [] + for core_mpo, core_vec in zip(list_cores_mpo, list_cores_vec): + # Contract over the shared physical index (operator applied to state), + # interleaving MPO and MPS bond indices. Bond dimensions multiply, + # requiring compression afterward. + # L,R = MPO bond left/right; l,r = MPS bond left/right; + # o = phys out; i = phys in (contracted) + mpo_dim_left, dim_out, _, mpo_dim_right = core_mpo.shape + vec_dim_left, _, vec_dim_right = core_vec.shape + #TODO: Cite an equation + core_contracted = np.einsum('LoiR,lir->LloRr', core_mpo, core_vec).reshape( + mpo_dim_left * vec_dim_left, dim_out, mpo_dim_right * vec_dim_right + ) + list_cores_compressed.append(core_contracted) + + # Peak-size proxy: sum of per-core complexity on the uncompressed MPS, + # where MPO and MPS bond dimensions have multiplied. This is the largest + # the tensor becomes in one matvec-then-compress cycle. + complexity = calc_mps_complexity(list_cores_compressed) + + # Compress contracted MPS back to tractable bond dimension with SVD + list_cores_result = tensor_compress(list_cores_compressed, epsilon, bond_dim_max) + if is_nested: + list_cores_result = unflatten_cores(list_cores_result, list_modes_per_site) + return list_cores_result, complexity + + +def calc_norm_corr_tensor( + wavefunction: HopsTensorWavefunction, + psi: np.ndarray, + z_hat: np.ndarray, + list_avg_L2: list[complex], + mode: HopsModes, + list_index_L2_by_mode: list[int], +) -> float: + """ + Computes the correction factor for propagating the normalized wave function. + + Parameters + ---------- + 1. wavefunction: HopsTensorWavefunction + MPS wavefunction container (provides list_cores_phi, + method, M1_modes_per_state). + 2. psi: np.ndarray(complex) + Physical (system) wavefunction already extracted from the + MPS. Passed in rather than re-extracted via extract_psi so + a single RK4 step doesn't redo the full contraction on an + unchanged wavefunction. + 3. z_hat: np.ndarray(complex) + Combined noise + memory term, indexed by L2 operator. + 4. list_avg_L2: list(complex) + Expectation values for each L2 operator. + 5. mode: HopsModes + Mode object providing list_g, list_L2_coo. + 6. list_index_L2_by_mode: list(int) + L2 operator index for each hierarchy mode, ordered + by mode position. + + Returns + ------- + 1. delta: float + Norm correction factor. + """ + list_L2 = mode.list_L2_coo + # z-component: sum_m z_hat_m * + delta = np.dot(z_hat, list_avg_L2) + V1_psi = psi + list_g = mode.list_g + n_modes = len(mode.list_modeidx_abs) + + V1_psi_conj = np.conj(V1_psi) + + # Per-mode correction: for each hierarchy mode m, subtract + # and add * , where phi_1_m is the first-order auxiliary + # along mode m scaled by V_m^- = sqrt(|g_m|) (Gao rescaling; see + # X. Gao, J. Ren, A. Eisfeld, Z. Shuai, "Non-Markovian stochastic + # Schrodinger equation: Matrix-product-state approach to the hierarchy of + # pure states," Phys. Rev. A 105, L030202 (2022), + # DOI: 10.1103/PhysRevA.105.L030202). + for (mode_idx, l2_idx) in enumerate(list_index_L2_by_mode): + # Unit vector in auxiliary space selecting mode_idx + list_aux_idx = [0] * n_modes + list_aux_idx[mode_idx] = 1 + # First-order auxiliary wavefunction scaled by V_m^- = sqrt(|g_m|), + # the same prefactor the b rail carries in MpoBuilder. + V1_phi_aux1 = ( + np.sqrt(np.abs(list_g[mode_idx])) + * phi_aux( + wavefunction.list_cores_phi, + wavefunction.method, + wavefunction.M1_modes_per_state, + list_aux_idx, + ) + ) + H2_lop = list_L2[l2_idx] + avg_lop = list_avg_L2[l2_idx] + # - + delta -= V1_psi_conj @ (H2_lop @ V1_phi_aux1) + # + * + delta += (V1_psi_conj @ V1_phi_aux1) * avg_lop + return np.real(delta) + + +def apply_system_operator( + list_cores_phi: list[np.ndarray], + O2_op_trimmed: np.ndarray, + method: str, + k_max: int, + M1_modes_per_state: np.ndarray, + epsilon: float, + bond_dim_max: int, +) -> list[np.ndarray]: + """ + Apply a pre-trimmed system-space operator to an MPS. + + The operator must already be trimmed to the active state_list. + + Parameters + ---------- + 1. list_cores_phi: list(np.ndarray) + MPS cores of the wavefunction. + 2. O2_op_trimmed: np.ndarray(complex) + Operator in active basis, shape (n_active, n_active). + 3. method: str + Tensor encoding type ('fullstate' or + 'number'). + 4. k_max: int + Maximum hierarchy depth. + 5. M1_modes_per_state: np.ndarray(int) + Number of bath modes per physical state. + 6. epsilon: float + SVD truncation threshold. + 7. bond_dim_max: int + Maximum MPS bond dimension. + + Returns + ------- + 1. list_cores_phi: list(np.ndarray) + Updated MPS cores. + """ + if hasattr(O2_op_trimmed, 'toarray'): + O2_op_trimmed = O2_op_trimmed.toarray() + if O2_op_trimmed.ndim != 2 or O2_op_trimmed.shape[0] != O2_op_trimmed.shape[1]: + raise ValueError( + f'O2_op_trimmed must be a square 2-D array, got shape ' + f'{O2_op_trimmed.shape}' + ) + n_op = O2_op_trimmed.shape[0] + + if method == 'fullstate': + # The first core has the full system Hilbert space as its physical + # dimension (1, n_state, bond_right), so the operator can be applied + # directly via matrix multiplication on the physical index. + n_phys = list_cores_phi[0].shape[1] + if n_op != n_phys: + raise ValueError( + f'Operator dimension ({n_op}) does not match physical ' + f'dimension of core 0 ({n_phys})' + ) + # a = bond left; b = bond right; + # i = phys out; j = phys in (contracted) + T3_core_new = np.einsum( + 'ij,ajb->aib', O2_op_trimmed, list_cores_phi[0], + ) + return [T3_core_new] + list_cores_phi[1:] + elif method == 'number': + # Each state occupies a separate binary core, so the operator + # cannot be applied to a single core. Instead, build a full + # operator MPO (with transfer matrices for off-diagonal elements) + # and apply it to the MPS via MPO-MPS contraction + compression. + list_cores_op = build_statenumber_operator_mpo( + O2_op_trimmed, n_op, k_max, M1_modes_per_state, + ) + # Observable/operator application path — peak-size tracking is + # only meaningful inside the RK4 derivative, so discard the scalar. + list_cores_phi_new, _ = tensor_matvec_prod( + list_cores_phi, list_cores_op, epsilon, bond_dim_max, + ) + return list_cores_phi_new + else: + raise NotImplementedError( + f'apply_system_operator not implemented for method={method!r}' + ) + + +def build_physical_correction_mps( + H1_psi_corr: np.ndarray, + method: str, + k_max: int, + M1_modes_per_state: np.ndarray, +) -> list: + """ + Build an MPS representing a correction to only the physical wavefunction + (zero-auxiliary component) of the hierarchy. + + The result is a low-bond-dimension MPS where all mode cores project onto + the m=0 occupation state. For fullstate this is rank-1 (bond dim 1). For + statenumber, each state group gets a state core with the correction + amplitude on the occupied index and zero on unoccupied, with m=0 projector + mode cores. + + Parameters + ---------- + 1. H1_psi_corr: np.ndarray(complex) + The correction vector, shape (n_state,). Typically + C2_LT_corr_physical @ psi or C2_LT_corr_linear @ psi. + + 2. method: str + 'fullstate' or 'number'. + + 3. k_max: int + Maximum hierarchy depth (mode core physical dim = k_max + 1). + + 4. M1_modes_per_state: np.ndarray(int) + Number of bath modes per system state. + + Returns + ------- + 1. list_cores: list + MPS cores representing the correction. Structure matches + the wavefunction format (nested lists for statenumber, + flat list for fullstate). + """ + n_state = len(H1_psi_corr) + + # Mode core projecting onto m=0: shape (1, k_max+1, 1) + T3_mode_proj = np.zeros((1, k_max + 1, 1), dtype=np.complex128) + T3_mode_proj[0, 0, 0] = 1.0 + + if method == 'fullstate': + T3_core_state = H1_psi_corr.reshape(1, n_state, 1) + n_total_modes = int(M1_modes_per_state.sum()) + return [T3_core_state] + [T3_mode_proj.copy() for _ in range(n_total_modes)] + + elif method == 'number': + list_cores = [] + for site in range(n_state): + T3_core_state = np.zeros((1, 2, 1), dtype=np.complex128) + T3_core_state[0, 1, 0] = H1_psi_corr[site] # occupied + # unoccupied index stays 0 — correction only contributes + # when this site is occupied + group = [T3_core_state] + [ + T3_mode_proj.copy() for _ in range(M1_modes_per_state[site]) + ] + list_cores.append(group) + return list_cores + + else: + raise UnsupportedRequest(method, 'build_physical_correction_mps') + + +def build_physical_correction_mpo( + C2_op: np.ndarray, + method: str, + k_max: int, + M1_modes_per_state: np.ndarray, +) -> list[np.ndarray]: + """ + Build a rank-1 MPO for a system-space operator that acts only on + the physical wavefunction (k=0 hierarchy level). + + Mode cores carry |0><0| projectors so the operator is zero on all + auxiliary (k>0) components. For fullstate the result is a flat list; + for statenumber it mirrors the interleaved state/mode core layout + of the main MPO. + + Parameters + ---------- + 1. C2_op: np.ndarray(complex) + System-space operator, shape (n_state, n_state). + 2. method: str + 'fullstate' or 'number'. + 3. k_max: int + Maximum hierarchy depth. + 4. M1_modes_per_state: np.ndarray(int) + Number of bath modes per system state. + + Returns + ------- + 1. list_cores_mpo: list(np.ndarray) + Rank-1 MPO cores, each shaped (1, d, d, 1). + """ + n_state = C2_op.shape[0] + d_mode = k_max + 1 + + # |0><0| projector on mode space + T4_mode_proj = np.zeros((1, d_mode, d_mode, 1), dtype=np.complex128) + T4_mode_proj[0, 0, 0, 0] = 1.0 + + if method == 'fullstate': + T4_state = C2_op.reshape(1, n_state, n_state, 1) + n_total_modes = int(M1_modes_per_state.sum()) + return [T4_state] + [T4_mode_proj.copy() for _ in range(n_total_modes)] + + elif method == 'number': + # Build via the general operator MPO, then replace identity + # mode cores with |0><0| projectors. + list_cores = build_statenumber_operator_mpo( + C2_op, n_state, k_max, M1_modes_per_state, + ) + for i, core in enumerate(list_cores): + # Mode cores have physical dim k_max+1; state cores have dim 2 + if core.shape[1] == d_mode: + T4_proj = np.zeros_like(core) + # Keep bond structure but zero out k>0 + T4_proj[:, 0, 0, :] = core[:, 0, 0, :] + list_cores[i] = T4_proj + return list_cores + + else: + raise UnsupportedRequest(method, 'build_physical_correction_mpo') diff --git a/src/mesohops/tensor/tensor_functions_adaptive.py b/src/mesohops/tensor/tensor_functions_adaptive.py new file mode 100644 index 0000000..3cd3b13 --- /dev/null +++ b/src/mesohops/tensor/tensor_functions_adaptive.py @@ -0,0 +1,227 @@ +from __future__ import annotations + +from collections.abc import Callable + +__title__ = 'Tensor Functions Adaptive' +__author__ = 'B. Z. Citty' +__maintainer__ = 'B. Z. Citty' + +import numpy as np +import scipy as sp + +from mesohops.basis.basis_functions import determine_error_thresh +from mesohops.util.exceptions import UnsupportedRequest +from mesohops.util.physical_constants import hbar +from mesohops.util.tensor_operations import contract_down_exact + + +def tensor_state_adaptive_check_add_state( + list_cores_phi: list, + ham: np.ndarray, + old_states: list[int], + n_state_full: int, + n_state: int, + delta_s: float, + state_list: list[int] | np.ndarray, + method: str, + M1_modes_per_state: np.ndarray, +) -> list[int]: + """ + Adaptive algorithm which checks to see if new states and their corresponding modes + should be added to the tensor train. + + Parameters + ---------- + 1. list_cores_phi: list + MPS cores of the wavefunction. + + 2. ham: np.ndarray + Full Hamiltonian matrix, shape (n_state_full, n_state_full). + + 3. old_states: list(int) + Absolute state indices already marked for removal + (excluded from the new-state candidates). + + 4. n_state_full: int + Total number of states in the full Hilbert space. + + 5. n_state: int + Number of currently active states. + + 6. delta_s: float + Adaptive threshold for state inclusion. + + 7. state_list: list(int) | np.ndarray + Currently active absolute state indices. + + 8. method: str + Tensor encoding type ('fullstaterepresentation' or + 'statenumberrepresentation'). + + 9. M1_modes_per_state: np.ndarray(int) + Number of bath modes per state. + + Returns + ------- + 1. list_new_states: list(int) + Absolute state indices to add. + """ + + if method == 'fullstate': + # Embed the adaptive-basis state core into the full Hilbert space, + # apply H to estimate flux into states outside the current basis. + state_core = list_cores_phi[0] + extended_state_core = np.zeros( + shape=(1, n_state_full, state_core.shape[2]), dtype=np.complex128 + ) + # Scatter adaptive-basis entries into full-space positions + for i in range(state_core.shape[2]): + extended_state_core[np.ix_(np.array([0]), state_list, np.array([i]))] = ( + state_core[np.ix_(np.array([0]), np.arange(n_state), np.array([i]))] + ) + # H @ psi gives the boundary flux into all states + boundary_state_core = np.tensordot(ham, extended_state_core, axes=([1], [1])) + boundary_state_core = np.swapaxes(boundary_state_core, 0, 1) + + # Contract mode indices to get per-state flux magnitude. + # Division by hbar^2 converts to the dimensionless error metric. + list_cores_tmp = list_cores_phi.copy() + list_cores_tmp[0] = boundary_state_core + V1_flux = ( + contract_down_exact(list_cores_tmp, method, n_state) + / hbar**2 + ) + # Zero flux for states already in basis (only check new states) + V1_flux[state_list] = 0 + adaptive_thresh = determine_error_thresh(np.sort(V1_flux), delta_s * delta_s) + list_new_states = np.where(V1_flux > adaptive_thresh)[0] + list_new_states = list(set(list_new_states) - set(old_states)) + return list_new_states + elif method == 'number': + # Contract mode indices to get per-state population sum_k |phi[k,s]|^2 + V1_phi_0 = contract_down_exact(list_cores_phi, method, n_state) + # Embed into full Hilbert space, then compute flux: + # |H|^2 @ |psi|^2 / hbar^2 estimates outgoing flux per state + V1_phi_0_full = np.zeros(n_state_full, dtype=np.complex128) + V1_phi_0_full[state_list] = V1_phi_0[:] + V1_flux = np.abs(ham**2) @ V1_phi_0_full / hbar**2 + # set the flux to states in the basis to 0. + V1_flux[state_list] = 0.0 + adaptive_thresh = determine_error_thresh(np.sort(V1_flux), delta_s * delta_s) + list_new_states = np.where(V1_flux > adaptive_thresh)[0] + list_new_states = list(set(list_new_states) - set(old_states)) + return list_new_states + else: + raise UnsupportedRequest(method, 'tensor_state_adaptive_check_add_state') + + +def tensor_state_adaptive_check_remove_state( + list_cores_phi: list, + ham: np.ndarray, + z_step: np.ndarray, + n_state_full: int, + n_state: int, + delta_s: float, + state_list: list[int] | np.ndarray, + method: str, + M1_modes_per_state: np.ndarray, + dsystem_dt: Callable, +) -> np.ndarray: + """ + Adaptive algorithm which checks to see if states and their corresponding modes + should be removed from the tensor train. + + Parameters + ---------- + 1. list_cores_phi: list + MPS cores of the wavefunction. + + 2. ham: np.ndarray + Full Hamiltonian matrix, shape (n_state_full, n_state_full). + + 3. z_step: np.ndarray + Noise values for the current time step, passed to + dsystem_dt as (z_mem, z_rnd, z_rnd2). + + 4. n_state_full: int + Total number of states in the full Hilbert space. + + 5. n_state: int + Number of currently active states. + + 6. delta_s: float + Adaptive threshold for state removal. + + 7. state_list: list(int) | np.ndarray + Currently active absolute state indices. + + 8. method: str + Tensor encoding type ('fullstaterepresentation' or + 'statenumberrepresentation'). + + 9. M1_modes_per_state: np.ndarray(int) + Number of bath modes per state. + + 10. dsystem_dt: Callable + Derivative closure that takes (z_mem, z_rnd, z_rnd2) + and returns the MPS derivative cores. + + Returns + ------- + 1. old_state_indices: np.ndarray(int) + Relative indices (into state_list) of states + to remove. + """ + list_states = state_list + list_cores_phi = list_cores_phi.copy() + + if method == 'fullstate': + # --- Flux-in: time derivative contribution --- + # Compute d(phi)/dt, divide by hbar, then contract mode indices + # to get per-state derivative magnitude + list_cores_d_phi = dsystem_dt(z_step[2], z_step[0], z_step[1]) + list_cores_d_phi[0] = list_cores_d_phi[0] / hbar + V1_error = contract_down_exact( + list_cores_d_phi, method, n_state + ) + + # --- Flux-out: norm contribution --- + # Contract phi to get per-state norm squared sum_k |phi[k,s]|^2 + V1_norm_sq_by_state = contract_down_exact( + list_cores_phi, method, n_state + ) + + elif method == 'number': + # --- Flux-in: time derivative contribution --- + list_cores_d_phi = dsystem_dt(z_step[2], z_step[0], z_step[1]) + list_cores_d_phi[0] = list_cores_d_phi[0] / hbar + V1_error = contract_down_exact( + list_cores_d_phi, method, n_state + ) + + # --- Flux-out: norm contribution --- + V1_norm_sq_by_state = contract_down_exact( + list_cores_phi, method, n_state + ) + else: + raise UnsupportedRequest(method, 'tensor_state_adaptive_check_remove_state') + + # Combine flux-in and flux-out into total error per state. + # Extract off-diagonal couplings (remove on-site energies), + # then compute sum_j |H_js|^2 * |psi_s|^2 / hbar^2 for each + # basis state s — this estimates outgoing coupling flux. + H2_sparse_hamiltonian = sp.sparse.coo_array(ham) + H2_sparse_couplings = H2_sparse_hamiltonian - sp.sparse.diags( + H2_sparse_hamiltonian.diagonal(0), + format='csc', + shape=H2_sparse_hamiltonian.shape, + ) + # Keep only columns for states in the current basis + H2_sparse_hamiltonian = H2_sparse_couplings[:, list_states] + # sum_j |H_js|^2 for each state s (column-wise squared sum) + V1_norm_sq = np.array(np.sum(np.abs(H2_sparse_hamiltonian).power(2), axis=0)) + V1_error += V1_norm_sq * V1_norm_sq_by_state / hbar**2 + # States with total error below threshold are candidates for removal + adaptive_thresh = determine_error_thresh(np.sort(V1_error), delta_s * delta_s) + old_state_indices = np.where(V1_error <= adaptive_thresh)[0] + return old_state_indices diff --git a/src/mesohops/trajectory/hops_dyadic.py b/src/mesohops/trajectory/hops_dyadic.py index d574dce..bb07cc5 100644 --- a/src/mesohops/trajectory/hops_dyadic.py +++ b/src/mesohops/trajectory/hops_dyadic.py @@ -49,7 +49,7 @@ def __init__(self, system_param, eom_param=None, noise_param=None, 6. integration_param: dict Dictionary of user-defined integration parameters. - [see integrator_rk.py and hops_trajectory.py] + [see integrator.py and hops_trajectory.py] """ diff --git a/src/mesohops/trajectory/hops_tensor_trajectory.py b/src/mesohops/trajectory/hops_tensor_trajectory.py new file mode 100644 index 0000000..4668aa2 --- /dev/null +++ b/src/mesohops/trajectory/hops_tensor_trajectory.py @@ -0,0 +1,1105 @@ +from __future__ import annotations + +import time as timer +import warnings +from collections.abc import Sequence + +import numpy as np +import scipy.sparse as sparse + +from mesohops.integrator.tensor_integrator import ( + runge_kutta_step_tensor, + runge_kutta_variables, + single_point_variables, + tdvp_step_tensor, +) +from mesohops.storage.storage_functions import ( + save_max_tensor_complexity, + save_phi_traj_tensor, + save_phi_norm_tensor, +) +from mesohops.tensor.hops_tensor_wavefunction import HopsTensorWavefunction +from mesohops.tensor.hops_tensor_basis import HopsTensorBasis +from mesohops.tensor.hops_tensor_eom import HopsTensorEOM +from mesohops.tensor.mpo_constructors import ( + build_statenumber_dipole_lower_plus_ident_mpo, + build_statenumber_dipole_mpo, + build_statenumber_dipole_raise_plus_ground_ident_mpo, +) +from mesohops.tensor.tensor_eom_functions import ( + apply_system_operator, + tensor_matvec_prod, +) +from mesohops.trajectory.hops_trajectory import HopsTrajectory +from mesohops.util.dynamic_dict import Dict_wDefaults +from mesohops.util.exceptions import LockedException, TrajectoryError, UnsupportedRequest +from mesohops.util.physical_constants import precision + +__title__ = 'Tensor HOPS Trajectory' +__author__ = 'B. Z. Citty' +__maintainer__ = 'B. Z. Citty' + +# Default values for the tensor_param dictionary accepted by +# HopsTensorTrajectory: MPS representation, SVD truncation thresholds, MPO +# builder selection and IVP solver settings. FLAG_MPO_OPTIMIZE is read only +# on the number path. IVP_MAX_STEP of None leaves the step size unlimited. +TENSOR_DICT_DEFAULT = { + 'METHOD': 'fullstate', + 'MPS_EPSILON': 1e-2, + 'MPO_EPSILON': 0.0, + 'FLAG_MPO_OPTIMIZE': True, + 'BOND_DIM_MAX': 10, + 'TDVP_UPDATE_TYPE': 'krylov', + 'IVP_METHOD': 'BDF', + 'IVP_RTOL': 1e-7, + 'IVP_ATOL': 1e-9, + 'IVP_MAX_STEP': None, +} + +# Allowed types for each tensor_param key, used by DynamicDict for validation. +TENSOR_DICT_TYPES = { + 'METHOD': [str], + 'MPS_EPSILON': [float], + 'MPO_EPSILON': [float], + 'FLAG_MPO_OPTIMIZE': [bool], + 'BOND_DIM_MAX': [int], + 'TDVP_UPDATE_TYPE': [str], + 'IVP_METHOD': [str], + 'IVP_RTOL': [float], + 'IVP_ATOL': [float], + 'IVP_MAX_STEP': [float, type(None)], +} + + +VALID_INTEGRATORS = {'RUNGE_KUTTA', 'TDVP1', 'TDVP2'} + + +class HopsTensorTrajectory(HopsTrajectory): + """ + Subclass of HopsTrajectory for tensor-network (TadHOPS) calculations. + The wavefunction is represented as a matrix product state (MPS) via + HopsTensorWavefunction, enabling efficient simulation of large systems. + + Initialization, propagation, and inchworm integration are all overridden + to operate on the tensor representation. The parent class provides shared + infrastructure: noise, basis management, storage, and utility methods. + """ + + # NOTE: dsystem_dt is not set on tensor trajectories. The parent's dsystem_dt + # is a derivative closure rebuilt when the basis changes. The tensor EOM + # (self.tensor_basis.eom) is a two-phase object: build_generator constructs + # an MPO each step, then derivative applies it. The MPO rebuilds on the fly, + # so the EOM does not need to be rebuilt or passed through update_basis. + # + # NOTE: the parent's _phi slot is intentionally unused in the tensor path. + # The tensor wavefunction lives on self.wavefunction, and the phi property + # (below) returns self.wavefunction.list_cores_phi. Any inherited code path + # that accesses self._phi directly would raise AttributeError. + + __slots__ = ( + # --- Tensor wavefunctions --- + 'tensor_param', # Tensor configuration dict + 'wavefunction', # Tensor wavefunction + 'tensor_basis', # HopsTensorBasis (system/mode/noise_memory/eom) + # --- Tensor integrator --- + 'is_tdvp', # True when using TDVP integrator + # step and integration_var are inherited from HopsTrajectory + # and overridden during initialize to use tensor integrators + '_tdvp_method', # TDVP variant string ('1tdvp' or '2tdvp') + # --- Output --- + # psi_traj stored through self.storage (same as parent) + ) + + def __init__( + self, + system_param=None, + eom_param=None, + noise_param=None, + noise2_param=None, + hierarchy_param=None, + storage_param=None, + integration_param=None, + tensor_param=None, + ) -> None: + """ + Initializes the tensor trajectory. Calls the parent constructor for all + shared infrastructure, then sets up tensor-specific components. + + Parameters + ---------- + 1. system_param: dict + Dictionary of user-defined system parameters. + [see hops_system.py] + + 2. eom_param: dict + Dictionary of user-defined equation-of-motion parameters. + [see hops_eom.py] + + 3. noise_param: dict + Dictionary of user-defined noise parameters. + [see hops_noise.py] + + 4. noise2_param: dict + Dictionary of user-defined secondary noise parameters. + [see hops_noise.py] + + 5. hierarchy_param: dict + Dictionary of user-defined hierarchy parameters. + [see hops_hierarchy.py] + + 6. storage_param: dict + Dictionary of user-defined storage parameters. + [see hops_storage.py] + + 7. integration_param: dict + Dictionary of user-defined integration parameters. + [see integrator.py] + + 8. tensor_param: dict + Dictionary of tensor configuration parameters. + a. 'METHOD': str + Tensor representation type. + [options: 'fullstate', + 'number'] + b. 'MPS_EPSILON': float + SVD truncation threshold for MPS (wavefunction) + compression. + c. 'MPO_EPSILON': float + SVD truncation threshold on MPO construction, + live in two places. Where MPO parts are + summed: the hierarchy plus Hamiltonian MPO on + the FLAG_MPO_OPTIMIZE=False path, and the + low-temperature correction folds. And on the + coupling-block factorization behind the + general generator, where raising it drops + weak couplings and narrows the MPO. + Default 0.0 (no compression). + d. 'FLAG_MPO_OPTIMIZE': bool + True (default) builds the number-method + generator as one MPO, each bond carrying + as many channels as the rank of the + coupling block that crosses it. False + takes the reference path: a hierarchy MPO + plus a separate Hamiltonian MPO, added + and compressed. + Read only when METHOD is 'number'. + e. 'BOND_DIM_MAX': int + Maximum MPS bond dimension. + f. 'TDVP_UPDATE_TYPE': str + TDVP update scheme. Default: 'krylov'. + g. 'IVP_METHOD': str + IVP solver method. Default: 'BDF'. + h. 'IVP_RTOL': float + IVP relative tolerance. Default: 1e-7. + i. 'IVP_ATOL': float + IVP absolute tolerance. Default: 1e-9. + j. 'IVP_MAX_STEP': float | None + IVP maximum step size. Default: None. + + Returns + ------- + None + """ + # TODO: enable sparse Hamiltonian use (not necessary for this paper) + # Ensure dense Hamiltonian before parent constructor builds HopsSystem. + # Copy the dict to avoid mutating the caller's input. + if system_param is not None and 'HAMILTONIAN' in system_param: + if sparse.issparse(system_param['HAMILTONIAN']): + system_param = dict(system_param) + system_param['HAMILTONIAN'] = system_param['HAMILTONIAN'].toarray() + + # Register the tensor-specific max_tensor_complexity key through the + # base HopsStorage constructor rather than mutating storage_dic / + # dic_save / data post-hoc. Passing a callable as the value lets the + # adaptive.setter registration loop wire it into dic_save and data + # in one pass. + if storage_param is None: + storage_param = {} + else: + storage_param = dict(storage_param) + storage_param.setdefault( + 'max_tensor_complexity', save_max_tensor_complexity, + ) + + # Initialize all shared infrastructure via the parent constructor + super().__init__( + system_param=system_param, + eom_param=eom_param, + noise_param=noise_param, + noise2_param=noise2_param, + hierarchy_param=hierarchy_param, + storage_param=storage_param, + integration_param=integration_param, + ) + + # Read k_max from the authoritative source (parent already resolved + # defaults via HIERARCHY_DICT_DEFAULT). + k_max = self.basis.hierarchy.param['MAXHIER'] + + # Tensor configuration: fill missing keys from defaults and + # validate types, matching the parent's integration_param pattern + if tensor_param is None: + tensor_param = {} + self.tensor_param = Dict_wDefaults._initialize_dictionary( + tensor_param, + TENSOR_DICT_DEFAULT, + TENSOR_DICT_TYPES, + 'tensor_param in HopsTensorTrajectory', + ) + + # Construct HopsTensorBasis with shared references from self.basis. + # IMPORTANT: tensor_basis.system, tensor_basis.mode, and + # tensor_basis.noise_memory are the SAME objects as self.basis.system, + # self.basis.mode, and self.basis.noise_memory — not copies. + # Mutations through either path (self.basis.* or self.tensor_basis.*) + # affect the same underlying objects. Convention: inherited + # infrastructure accesses through self.basis, tensor-specific logic + # accesses through self.tensor_basis. + self.tensor_basis = HopsTensorBasis( + self.basis.system, self.basis.mode, self.basis.noise_memory, + ) + self.wavefunction = HopsTensorWavefunction( + k_max, self.tensor_param, self.integration_param, + self.basis.eom.param, + ) + + # Register tensor-aware storage functions that understand MPS data. + # phi_traj: store full MPS cores instead of a flat vector. + # phi_norm: compute hierarchy norm via MPS double-layer contraction. + if 'phi_traj' in self.storage.dic_save: + self.storage.dic_save['phi_traj'] = save_phi_traj_tensor + if 'phi_norm' in self.storage.dic_save: + self.storage.dic_save['phi_norm'] = save_phi_norm_tensor + # list_aux_norm: vector-HOPS diagnostic (per-auxiliary norms via + # reshape). Not meaningful for tensor HOPS where the hierarchy is + # compressed into MPS bond dimensions. + if 'list_aux_norm' in self.storage.dic_save: + warnings.warn( + 'list_aux_norm is not supported for tensor trajectories ' + '(hierarchy is compressed into MPS bond dimensions, not ' + 'enumerated as discrete auxiliary vectors). Disabling.', + stacklevel=2, + ) + del self.storage.dic_save['list_aux_norm'] + self.storage.data.pop('list_aux_norm', None) + + # max_tensor_complexity (per-timestep peak MPS complexity) was + # registered via storage_param above, before super().__init__, so + # the base HopsStorage adaptive.setter already populated + # storage_dic, dic_save, and data for this key in one pass. + + def make_adaptive(self, + delta_a: float = 1e-4, + delta_s: float = 1e-4, + update_step: int = 1, + f_discard: float = 0.01, + list_permanent_sites: list[int] | None = None, + adaptive_noise: bool = True) -> None: + """ + Configures this trajectory for adaptive HOPS. Overrides the parent + to reject `list_permanent_sites`, which `HopsTensorBasis` does not + honor — passing it to the parent would store it on + `system.param["list_permanent_sites"]` but the tensor adaptive + basis never reads that key, so the requested sites would silently + not be preserved. Failing here gives the caller a clear signal at + configuration time rather than wrong physics at run time. + See `HopsTrajectory.make_adaptive` for the parameter semantics. + """ + if list_permanent_sites is not None: + raise NotImplementedError( + 'list_permanent_sites is not supported for tensor ' + 'trajectories. HopsTensorBasis does not read ' + 'system.param["list_permanent_sites"], so passing this ' + 'argument would be silently ignored at runtime.' + ) + super().make_adaptive( + delta_a=delta_a, + delta_s=delta_s, + update_step=update_step, + f_discard=f_discard, + list_permanent_sites=list_permanent_sites, + adaptive_noise=adaptive_noise, + ) + + def _setup_integrator(self) -> None: + """ + Configures tensor-specific integration step function and variable gatherer. + + Overrides the parent to support TDVP integrators. + """ + if self.integrator == 'RUNGE_KUTTA': + self.step = runge_kutta_step_tensor + self.integration_var = runge_kutta_variables + self.integrator_step = 0.5 + self.is_tdvp = False + self._tdvp_method = None + elif self.integrator == 'TDVP1': + self.step = tdvp_step_tensor + self.integration_var = single_point_variables + self.integrator_step = 1.0 + self.is_tdvp = True + self._tdvp_method = '1tdvp' + elif self.integrator == 'TDVP2': + self.step = tdvp_step_tensor + self.integration_var = single_point_variables + self.integrator_step = 1.0 + self.is_tdvp = True + self._tdvp_method = '2tdvp' + else: + raise UnsupportedRequest( + f'Integrator {self.integrator!r}, expected one of ' + f'{sorted(VALID_INTEGRATORS)}', + type(self).__name__, + ) + + def initialize( + self, + psi_0: Sequence[complex] | np.ndarray, + timer_checkpoint: float | None = None, + ) -> None: + """ + Initializes the tensor trajectory. Sets up both the vector basis (shared + infrastructure used by inchworm) and the tensor wavefunction representation. + + Parameters + ---------- + 1. psi_0: np.ndarray(complex) + Wave function at initial time. + + 2. timer_checkpoint: float | None + Wall-clock time prior to initialization [units: s]. + If None, uses current wall-clock time. + + Returns + ------- + None + """ + if timer_checkpoint is None: + timer_checkpoint = timer.time() + + psi_0 = np.array(psi_0, dtype=np.complex128) + + if not self.__initialized__: + # --- Step 1: Initialize basis and EOM --- + # The parent calls self.basis.initialize(psi_0) which initializes + # hierarchy, system, mode, noise_memory, and builds the vector EOM + # derivative closure in one call. Tensor HOPS cherry-picks: + # + # - Hierarchy init is skipped because hierarchy depth is encoded in + # MPS core dimensions (k_max + 1 per mode), not through explicit + # auxiliary vector enumeration. + # - Vector EOM is skipped because tensor HOPS uses HopsTensorEOM + # (MPO-based) instead of a derivative closure. + # - The mode union with hierarchy modes is a no-op because the + # hierarchy object never populates its mode list in the tensor + # path. + self.basis.system.initialize(self.basis.adaptive_s, psi_0) + self.basis.mode.list_modeidx_abs = sorted( + self.basis.system.list_statemodeidx_abs + ) + self.basis.noise_memory.initialize() + self.tensor_basis.initialize(self.basis.eom.param.get('DELTA_S', 0)) + self.wavefunction.initialize( + psi_0, + self.tensor_basis.system, + ) + self.tensor_basis.eom = HopsTensorEOM( + self.wavefunction, + self.tensor_basis.system, + self.tensor_basis.mode, + self.tensor_basis.noise_memory, + self.tensor_basis.adaptive, + self.basis.eom.param, + ) + + # --- Step 2: z_mem --- + self.z_mem = np.zeros( + len(self.basis.noise_memory.list_zmemmodeidx_abs), + dtype=np.complex128, + ) + + # --- Step 3: storage.n_dim --- + self.storage.n_dim = self.basis.system.param['NSTATES'] + + # --- Step 4: Adaptive basis setup --- + if self.basis.adaptive: + if self.static_basis is not None: + raise NotImplementedError( + 'static_basis is not yet supported for tensor ' + 'trajectories.' + ) + # storage.adaptive is already set by make_adaptive(); + z_step = self._prepare_zstep(self.z_mem) + list_states_old, list_states_new = ( + self.tensor_basis.define_basis(self.wavefunction, z_step) + ) + self.wavefunction, self.z_mem = self.tensor_basis.update_basis( + self.wavefunction, self.z_mem, + list_states_old, list_states_new, + ) + + # --- Step 5: Set time --- + self.t = 0 + + # --- Step 6: Store initial state --- + # phi_new: the parent passes the full hierarchy vector phi. + # We pass psi because save_psi_traj (always active) slices + # phi_new[:len(state_list)], which is a no-op on psi. + # save_phi_traj_tensor and save_phi_norm_tensor ignore + # phi_new and read wavefunction directly. + self.storage.store_step( + phi_new=self.wavefunction.psi, + wavefunction=self.wavefunction, + state_list=list(self.tensor_basis.system.state_list), + t_new=0, + aux_list=self.auxiliary_list, + z_mem_new=self.z_mem, + list_zmemmodeidx_abs=( + self.basis.noise_memory.list_zmemmodeidx_abs + ), + max_tensor_complexity=0, + ) + + # --- Step 7: Metadata and lock --- + self.storage.metadata['INITIALIZATION_TIME'] = ( + timer.time() - timer_checkpoint + ) + self.__initialized__ = True + else: + raise LockedException('initialize', 'HopsTensorTrajectory') + + def propagate( + self, + t_advance: float, + tau: float, + timer_checkpoint: float | None = None, + ) -> None: + """ + Propagates the tensor wavefunction forward in time. At each step the MPS + is advanced via the configured integrator. + + Parameters + ---------- + 1. t_advance: float + How far out in time the calculation will run [units: fs]. + + 2. tau: float + Time step [units: fs]. + + 3. timer_checkpoint: float | None + System time prior to propagation [units: s]. + If None, uses current system time. + + Returns + ------- + None + """ + if timer_checkpoint is None: + timer_checkpoint = timer.time() + + # Construct the time axis + # NOTE: t_axis is only defined inside this branch. If TAU is None + # and INTERPOLATE is True, t_axis will be undefined and line + # `np.max(t_axis)` below will raise UnboundLocalError. Same issue + # exists in the parent HopsTrajectory.propagate. Awaiting team + # review before fixing. + t0 = self.t + if (self.noise1.param['TAU'] is not None) or not ( + self.noise1.param['INTERPOLATE'] + ): + if self._check_tau_step(tau, precision): + n_steps = int(np.ceil(t_advance / tau)) + t_axis = t0 + np.arange(1, 1 + n_steps) * tau + print('Integration from ', t0, ' to ', np.max(t_axis)) + else: + raise TrajectoryError( + 'Timesteps (' + + str(tau * self.integrator_step) + + ") that do not match noise.param['TAU'] (" + + str(self.noise1.param['TAU']) + + ')' + ) + + if np.max(t_axis) > self.noise1.param['TLEN']: + raise TrajectoryError( + "Trajectory times longer than noise.param['TLEN'] (" + + str(self.noise1.param['TLEN']) + + ')' + ) + + # Tracks the system timescale so the timestep warning fires only once + tau_sys = None + + store_step_timing = self.integration_param['STORE_STEP_TIMING'] + + for idx_t, t in enumerate(t_axis): + if store_step_timing: + t_step_start = timer.time() + # Check that timestep is resolved by system timescale + if tau > self.basis.system.system_timescale and ( + tau_sys is None or tau_sys > self.basis.system.system_timescale + ): + tau_sys = self.basis.system.system_timescale + + # Tensor step + dict_var = self.integration_var( + self.z_mem, + self.t, + self.noise1, + self.noise2, + tau, + self.basis.mode.list_l2idx_abs, + self.effective_noise_integration, + ) + # Parent does: phi, z_mem = self.step(self.dsystem_dt, **var_list). + # Tensor _step returns just z_mem because the wavefunction is + # mutated in place by the integrator. Peak uncompressed-MPS + # size during the step is published on eom.max_complexity_step + # (0 for TDVP — see step function docstrings) and read here. + z_mem = self._step(dict_var) + max_complexity_step = self.tensor_basis.eom.max_complexity_step + # Parent does: phi = self.normalize(phi). Tensor version mutates + # wavefunction class instance mutated in place by the integrator. + self.normalize() + + # (C) Adaptive basis update + if self.basis.adaptive: + if self.use_early_integrator: + print(f'Early Integration: Using {self.early_integrator}') + # The parent checks both INCH_WORM and STATIC hierarchy + # options. Tensor HOPS has no hierarchy adaptivity, so + # only the INCH_WORM path applies. + if self.early_integrator == 'INCH_WORM': + z_step = self._prepare_zstep(z_mem) + # TODO: make output match parent class (state + # update abstraction) + list_states_old, list_states_new = ( + self.tensor_basis.define_basis(self.wavefunction, z_step) + ) + # Deep copy: statenumber has nested lists (groups + # of arrays), so a flat list comprehension would + # only shallow-copy the outer list. + if self.wavefunction.method == 'number': + list_cores_checkpoint = [ + [arr.copy() for arr in g] + for g in self.wavefunction.list_cores_phi + ] + else: + list_cores_checkpoint = [ + c.copy() for c in self.wavefunction.list_cores_phi + ] + # Iterate until basis converges (define_basis proposes + # no further changes) or the inchworm cap is reached. + step_num = 0 + while list_states_old != [] or list_states_new != []: + ( + z_mem, + max_complexity_inch, + list_states_old, + list_states_new, + list_cores_checkpoint, + ) = self.inchworm_integrate( + tau, + list_cores_checkpoint, + list_states_old, + list_states_new, + ) + # Accumulate the per-timestep max across + # inchworm iterations. + if max_complexity_inch > max_complexity_step: + max_complexity_step = max_complexity_inch + step_num += 1 + if step_num >= self.inchworm_cap: + break + # Parent returns (phi, z_mem, self.dsystem_dt). Tensor omits + # dsystem_dt — tensor EOM rebuilds its MPO on the fly. + # TODO: adaptive path needs dsystem_dt in form of + # calc_deriv_cores being passed in + self.wavefunction, z_mem = self.tensor_basis.update_basis( + self.wavefunction, z_mem, list_states_old, list_states_new + ) + else: + raise UnsupportedRequest( + self.early_integrator, + 'early time integrator clause of the tensor propagate', + ) + self._early_step_counter += 1 + + # Standard adaptive integration: check every update_step + # steps whether states should be added or removed + elif (idx_t + 1) % self.update_step == 0: + z_step = self._prepare_zstep(z_mem) + list_states_old, list_states_new = ( + self.tensor_basis.define_basis(self.wavefunction, z_step) + ) + # Parent returns (phi, z_mem, self.dsystem_dt). Tensor omits + # dsystem_dt — tensor EOM rebuilds its MPO on the fly. + # TODO: adaptive path needs dsystem_dt in form of + # calc_deriv_cores being passed in + self.wavefunction, z_mem = self.tensor_basis.update_basis( + self.wavefunction, z_mem, list_states_old, list_states_new + ) + + # Parent also does: self.phi = phi. Tensor omits this because + # wavefunction is already on self, mutated in place. + self.z_mem = z_mem + self.t = t + + if self.storage.check_storage_time(t): + # phi_new receives psi, not the full hierarchy vector; + # see the comment in initialize() step 6. + self.storage.store_step( + phi_new=self.wavefunction.psi, + wavefunction=self.wavefunction, + state_list=list(self.tensor_basis.system.state_list), + t_new=t, + aux_list=self.auxiliary_list, + z_mem_new=self.z_mem, + list_zmemmodeidx_abs=self.basis.noise_memory.list_zmemmodeidx_abs, + max_tensor_complexity=max_complexity_step, + ) + + if store_step_timing: + self.storage.metadata['LIST_PROPAGATION_TIME'].append( + (t, timer.time() - t_step_start) + ) + + # Store propagation time + if not store_step_timing: + self.storage.metadata['LIST_PROPAGATION_TIME'].append( + timer.time() - timer_checkpoint + ) + + if tau_sys is not None: + warnings.warn( + f'At some point during propagation, the time step ({tau} fs)' + f' was larger than the estimated timescale associated with ' + f'the system Hamiltonian ({tau_sys} fs). A smaller time step ' + f'may be necessary to correctly resolve dynamics.' + ) + + def _step(self, dict_var: dict) -> np.ndarray: + """ + Dispatches a single tensor integration step (RK4 or TDVP). + + Mutates self.tensor_basis.eom.wavefunction in place. + + Parameters + ---------- + 1. dict_var : dict + Variables from integration_var. + + Returns + ------- + 1. z_mem : np.ndarray(complex) + Updated noise memory drift terms [units: cm^-1]. + + Side effects + ------------ + The underlying step function sets + `self.tensor_basis.eom.max_complexity_step` to the peak + uncompressed-MPS complexity observed (RK4 tracks across + sub-stages; TDVP sets 0). propagate() reads that attribute + for storage rather than threading the scalar through a tuple. + """ + if self.is_tdvp: + # TDVP requires additional solver configuration (Krylov + # subspace size, solver type, IVP tolerances) beyond the + # base noise/memory variables. + tp = self.tensor_param + return self.step( + self.tensor_basis.eom, + dict_var['z_mem'], + dict_var['z_rnd'], + dict_var['z_rnd2'], + dict_var['tau'], + method=self._tdvp_method, + # TODO: give krylov_conv_tol its own tensor_param key instead + # of coupling it to mps_epsilon (they control different things). + krylov_conv_tol=tp['MPS_EPSILON'] / 10, + update_type=tp['TDVP_UPDATE_TYPE'], + ivp_method=tp['IVP_METHOD'], + ivp_rtol=tp['IVP_RTOL'], + ivp_atol=tp['IVP_ATOL'], + ivp_max_step=tp['IVP_MAX_STEP'], + ) + # RK4 path: only needs the EOM and noise/memory variables + return self.step( + self.tensor_basis.eom, + dict_var['z_mem'], + dict_var['z_rnd'], + dict_var['z_rnd2'], + dict_var['tau'], + ) + + def _operator(self, op: np.ndarray | sparse.spmatrix) -> None: + """ + Applies an operator to the tensor wavefunction. Mirrors the parent + HopsTrajectory._operator: expands the adaptive basis to include all + states coupled by the operator, trims to the active basis, applies + the operator, then cleans up the basis afterward. + + Parameters + ---------- + 1. op: np.ndarray | sparse.spmatrix + The operator as a full system-space matrix, + shape (n_state_full, n_state_full). + + Returns + ------- + None + """ + if sparse.issparse(op): + op = op.tocsr() + + # Validate operator dimensions against full system size. Without + # this guard, too-small operators leak an IndexError from the + # np.ix_ trim below, and too-large operators are silently sliced + # to the first n_state_full x n_state_full block — neither is + # what a caller who passed a mismatched operator actually wants. + n_state_full = self.basis.system.param['NSTATES'] + if op.shape != (n_state_full, n_state_full): + raise ValueError( + f'op must have shape ({n_state_full}, {n_state_full}) ' + f'(full system size); got {op.shape}' + ) + + # TODO: adaptive basis expansion/cleanup around operator application + # is not fully tested and may not correctly handle all edge cases + # (e.g., states that become depopulated after the operator). + # Expand adaptive basis if operator couples to new states + if self.tensor_basis.adaptive: + list_states_operator = np.unique( + np.nonzero(op[:, self.tensor_basis.system.state_list])[0] + ) + list_states_new = sorted( + set(list_states_operator) - set(self.tensor_basis.system.state_list) + ) + self.wavefunction, self.z_mem = self.tensor_basis.update_basis( + self.wavefunction, + self.z_mem, + [], + list_states_new, + ) + + # Trim to active basis and apply + state_list = self.tensor_basis.system.state_list + if sparse.issparse(op): + O2_trimmed = op[np.ix_(state_list, state_list)].toarray() + else: + O2_trimmed = np.asarray(op)[np.ix_(state_list, state_list)] + + self.wavefunction.list_cores_phi = apply_system_operator( + self.wavefunction.list_cores_phi, + O2_trimmed, + self.wavefunction.method, + self.wavefunction.k_max, + self.wavefunction.M1_modes_per_state, + self.wavefunction.mps_epsilon, + self.wavefunction.bond_dim_max, + ) + + # Post-operator basis cleanup + if self.tensor_basis.adaptive: + z_step = self._prepare_zstep(self.z_mem) + list_states_old, list_states_new = ( + self.tensor_basis.define_basis(self.wavefunction, z_step) + ) + self.wavefunction, self.z_mem = self.tensor_basis.update_basis( + self.wavefunction, + self.z_mem, + list_states_old, + list_states_new, + ) + self.reset_early_time_integrator() + + def apply_dipole_raise(self, list_mu: np.ndarray) -> None: + """ + Applies the bond-dim-2 dipole-raise MPO + sum_k list_mu[k] * a_k^dagger + to the wavefunction in the vacuum convention. Used by the + fluorescence path, which matches gs_core fluorescence's raise + (no |g> preservation). For the absorption path, use + apply_dipole_raise_plus_ground_ident instead. + + Parameters + ---------- + 1. list_mu: np.ndarray(complex) + Per-site dipole amplitudes, shape (n_state,). + Entries set to 0 mark sites excluded from the + dipole's site selection. + + Returns + ------- + None + """ + n_state = len(self.tensor_basis.system.state_list) + list_cores_mpo = build_statenumber_dipole_mpo( + list_mu, n_state, self.wavefunction.k_max, + self.wavefunction.M1_modes_per_state, 'raise', + ) + new_cores, _ = tensor_matvec_prod( + self.wavefunction.list_cores_phi, list_cores_mpo, + self.wavefunction.mps_epsilon, self.wavefunction.bond_dim_max, + ) + self.wavefunction.list_cores_phi = new_cores + + def apply_dipole_raise_plus_ground_ident( + self, list_mu: np.ndarray, + ) -> None: + """ + Applies the bond-dim-3 MPO + sum_k list_mu[k] * a_k^dagger + I_g + to the wavefunction in the vacuum convention. The +I_g term + (projector onto the all-zeros configuration) preserves the GS + amplitude after the raise, matching the gs_core absorption + raise's "(0,0)=1 keep |g>" behavior. Used by the absorption + path; the fluorescence raise (no |g> preservation) uses + apply_dipole_raise instead. + + Parameters + ---------- + 1. list_mu: np.ndarray(complex) + Per-site dipole amplitudes, shape (n_state,). + Entries set to 0 mark sites excluded from the + dipole's site selection. + + Returns + ------- + None + """ + n_state = len(self.tensor_basis.system.state_list) + list_cores_mpo = build_statenumber_dipole_raise_plus_ground_ident_mpo( + list_mu, n_state, self.wavefunction.k_max, + self.wavefunction.M1_modes_per_state, + ) + new_cores, _ = tensor_matvec_prod( + self.wavefunction.list_cores_phi, list_cores_mpo, + self.wavefunction.mps_epsilon, self.wavefunction.bond_dim_max, + ) + self.wavefunction.list_cores_phi = new_cores + + def apply_dipole_lower_plus_ident(self, list_mu: np.ndarray) -> None: + """ + Applies the bond-dim-4 MPO + sum_k list_mu[k] * a_k + I_excited + to the wavefunction in the vacuum convention. The +I_excited + term preserves the single-excitation manifold content; the + sum_k a_k term creates the GS amplitude on the all-zeros + configuration of the MPS. Used by the fluorescence path as the + de-excitation step: applied with E_sig after the t2 waiting + period to inject the GS amplitude that forms the G/E coherence + measured during the detection phase. + + Parameters + ---------- + 1. list_mu: np.ndarray(complex) + Per-site dipole amplitudes, shape (n_state,). + Entries set to 0 mark sites excluded from the + dipole's site selection. + + Returns + ------- + None + """ + n_state = len(self.tensor_basis.system.state_list) + list_cores_mpo = build_statenumber_dipole_lower_plus_ident_mpo( + list_mu, n_state, self.wavefunction.k_max, + self.wavefunction.M1_modes_per_state, + ) + new_cores, _ = tensor_matvec_prod( + self.wavefunction.list_cores_phi, list_cores_mpo, + self.wavefunction.mps_epsilon, self.wavefunction.bond_dim_max, + ) + self.wavefunction.list_cores_phi = new_cores + + def normalize(self) -> None: + """ + Normalizes the tensor wavefunction in place if the EOM requires it. + + The parent HopsTrajectory.normalize takes phi as an argument and returns + the normalized phi — a value-passing pattern. The tensor version operates + on self.wavefunction in place because the wavefunction is a class instance + (HopsTensorWavefunction) rather than a bare numpy array. This difference + will shrink when the parent gets a wavefunction class. + + The normalization policy (whether to normalize based on EOM type) lives + here in the trajectory, not in the wavefunction class — matching the + parent where HopsTrajectory.normalize checks self.basis.eom.normalized. + """ + if self.basis.eom.normalized: + self.wavefunction.normalize() + + def inchworm_integrate( + self, + tau: float, + list_cores_checkpoint: list, + list_states_old=None, + list_states_new=None, + ) -> tuple: + """ + Performs one inchworm iteration for the tensor trajectory. + + The tensor basis is expanded by the proposed (list_states_old, + list_states_new) update. wavefunction is restored to the checkpoint + state in the new basis, then a full integration step is taken. The + resulting wavefunction is used to propose the next basis update. + + Note: the signature and return tuple differ from the parent's + inchworm_integrate because tensor HOPS has no auxiliary vectors + (the hierarchy is compressed into MPS bond dimensions) and uses + MPS cores instead of a flat phi vector. + + Parameters + ---------- + 1. tau: float + Time step [units: fs]. + + 2. list_cores_checkpoint: list(np.ndarray) + Saved list_cores_phi representing the wavefunction at + the start of the current timestep, expressed in + the basis prior to this iteration's update. + + 3. list_states_old: list | None + Tensor basis state indices to remove in + this iteration's basis update. + + 4. list_states_new: list | None + New tensor states to add in this iteration's + basis update. + + Returns + ------- + 1. z_mem: np.array(complex) + Noise memory drift after the tensor step [units: cm^-1]. + + 2. max_complexity: int + Peak uncompressed-MPS complexity across this + iteration's integrator call (0 for TDVP). + Read from eom.max_complexity_step which the + step function publishes. + + 3. list_states_old: list + Tensor state indices proposed for removal in the + next basis update. + + 4. list_states_new: list + New tensor states proposed for addition in the next + basis update. + + 5. list_cores_checkpoint: list(np.ndarray) + Updated checkpoint cores in the new (expanded) + basis, for use in the next inchworm iteration. + """ + # Restore wavefunction to the checkpoint, then expand the basis. + # After update_basis, wavefunction.list_cores_phi holds the checkpoint state + # expressed in the new (expanded) basis. + self.wavefunction.restore_phi(list_cores_checkpoint) + # TODO: for adaptive path, z_mem changes, so this logic should + # update to give the updated z_mem + self.wavefunction, _ = self.tensor_basis.update_basis( + self.wavefunction, self.z_mem, list_states_old, list_states_new + ) + # Deep copy: checkpoint must be independent of subsequent + # integrator mutations to list_cores_phi. Statenumber has nested + # lists (groups of arrays) requiring inner-level copy. + if self.wavefunction.method == 'number': + list_cores_checkpoint = [ + [arr.copy() for arr in g] + for g in self.wavefunction.list_cores_phi + ] + else: + list_cores_checkpoint = [ + c.copy() for c in self.wavefunction.list_cores_phi + ] + + dict_var = self.integration_var( + self.z_mem, + self.t, + self.noise1, + self.noise2, + tau, + self.basis.mode.list_l2idx_abs, + self.effective_noise_integration, + ) + z_mem = self._step(dict_var) + max_complexity = self.tensor_basis.eom.max_complexity_step + self.normalize() + + # Define next basis proposal from the stepped state + z_step = self._prepare_zstep(z_mem) + list_states_old, list_states_new = self.tensor_basis.define_basis( + self.wavefunction, z_step, + ) + + return ( + z_mem, + max_complexity, + list_states_old, + list_states_new, + list_cores_checkpoint, + ) + + # TODO: when implemented, match parent signatures and use cls(...) + # for subclassability: + # save_checkpoint(filepath: str | os.PathLike, ...) + # load_checkpoint( + # filename: str | os.PathLike, + # add_seed1: int | str | os.PathLike | np.ndarray | None, + # add_seed2: int | str | os.PathLike | np.ndarray | None, + # add_system_param: str | os.PathLike | None, + # ) -> HopsTensorTrajectory (via cls(...)) + def save_checkpoint( + self, + filepath: str, + compress: bool = True, + drop_seed: bool = False, + ) -> None: + """ + Not yet supported for tensor trajectories. + + Raises + ------ + NotImplementedError + """ + raise NotImplementedError( + 'Tensor trajectory checkpointing is not yet supported.' + ) + + @classmethod + def load_checkpoint( + cls, + filename: str, + add_seed1: int | None = None, + add_seed2: int | None = None, + add_system_param: str | None = None, + ) -> HopsTensorTrajectory: + """ + Not yet supported for tensor trajectories. + + Raises + ------ + NotImplementedError + """ + raise NotImplementedError( + 'Tensor trajectory checkpointing is not yet supported.' + ) + + @property + def psi(self) -> np.ndarray: + """Current physical wavefunction (compact, active states only). + + Matches the parent HopsTrajectory.psi which returns + phi[:n_state]. Storage handles reconstruction to full-size + arrays via state_list when the user accesses storage['psi_traj']. + """ + return self.wavefunction.psi + + @property + def phi(self) -> list: + """Current hierarchy wavefunction as MPS cores.""" + return self.wavefunction.list_cores_phi + + @phi.setter + def phi(self, value: list) -> None: + self.wavefunction.list_cores_phi = value diff --git a/src/mesohops/trajectory/hops_trajectory.py b/src/mesohops/trajectory/hops_trajectory.py index 51278e1..22c3f13 100644 --- a/src/mesohops/trajectory/hops_trajectory.py +++ b/src/mesohops/trajectory/hops_trajectory.py @@ -38,6 +38,7 @@ "INCHWORM_CAP": 5, "STATIC_BASIS": None, "EFFECTIVE_NOISE_INTEGRATION": False, + "STORE_STEP_TIMING": False, } INTEGRATION_DICT_TYPES = { @@ -47,6 +48,7 @@ "INCHWORM_CAP": [int], "STATIC_BASIS": [type(None), list, np.ndarray], "EFFECTIVE_NOISE_INTEGRATION": [bool], + "STORE_STEP_TIMING": [bool], } @@ -153,7 +155,7 @@ def __init__( 6. integration_param: dict Dictionary of user-defined integration parameters. - [see integrator_rk.py] + [see integrator.py] a. INTEGRATOR b. EARLY_ADAPTIVE_INTEGRATOR c. EARLY_INTEGRATOR_STEPS @@ -241,8 +243,19 @@ def __init__( f'calculations. The number of early integrator steps has ' f'been reset to the default of {self.early_steps}.') self._early_step_counter = 0 + self._setup_integrator() + + # LOCKING VARIABLE + self.__initialized__ = False + + def _setup_integrator(self) -> None: + """ + Configures the integration step function and variable gatherer. + + Subclasses override this to support additional integrators (e.g. TDVP). + """ if self.integrator == "RUNGE_KUTTA": - from mesohops.integrator.integrator_rk import ( + from mesohops.integrator.integrator import ( runge_kutta_step, runge_kutta_variables, ) @@ -256,9 +269,6 @@ def __init__( type(self).__name__, ) - # LOCKING VARIABLE - self.__initialized__ = False - def initialize( self, psi_0: Sequence[complex] | np.ndarray, @@ -473,7 +483,11 @@ def propagate( # Performs integration # ------------------- + store_step_timing = self.integration_param["STORE_STEP_TIMING"] + for (index_t, t) in enumerate(t_axis): + if store_step_timing: + t_step_start = timer.time() # Check that timestep is resolved if (tau > self.basis.system.system_timescale and (tau_sys is None or tau_sys > self.basis.system.system_timescale)): @@ -566,11 +580,17 @@ def propagate( list_zmemmodeidx_abs=self.basis.noise_memory.list_zmemmodeidx_abs, ) + if store_step_timing: + self.storage.metadata["LIST_PROPAGATION_TIME"].append( + (t, timer.time() - t_step_start) + ) + # Stores propagation time # -------------------------- - self.storage.metadata["LIST_PROPAGATION_TIME"].append(timer.time() - - timer_checkpoint) + if not store_step_timing: + self.storage.metadata["LIST_PROPAGATION_TIME"].append(timer.time() - + timer_checkpoint) # Warn the user if overly aggressive time steps were detected. if tau_sys is not None: diff --git a/src/mesohops/util/nondyadic_spectroscopy.py b/src/mesohops/util/nondyadic_spectroscopy.py new file mode 100644 index 0000000..aed87ca --- /dev/null +++ b/src/mesohops/util/nondyadic_spectroscopy.py @@ -0,0 +1,608 @@ +from typing import NamedTuple + +import numpy as np +from scipy import sparse + +from mesohops.util.exceptions import UnsupportedRequest + +__title__ = 'nondyadic_spectroscopy' +__author__ = 'A. Hartzell' +__maintainer__ = 'A. Hartzell' + + +class SpectroscopyDispatch(NamedTuple): + traj_kind: str + tensor_method: str | None + hilbert_tag: str + eom_tag: str + + +def _spectroscopy_key(traj, n_site): + """ + Canonical dispatch tuple for the spectroscopy code path. + + Encodes (in order) the trajectory type, tensor method (if tensor), + Hilbert-space convention, and EOM tag. The Hilbert-space convention + means: + - embedded: the ground state is an explicit basis slot in the + Hilbert space, so the trajectory dimension is n_site + 1. + - vacuum: the ground state is tracked implicitly as the all-zero + MPS configuration, so the trajectory dimension is n_site and + the tensor representation must be ``number``. + - excited_only: only the excited manifold is represented; this + is used for the LINEAR EOM for vector and fullstate + trajectories. + + Validates that the trajectory is uninitialized, the EOM is in the + supported set, and the Hilbert dimension is either n_site or + n_site + 1. + + Example return values: + SpectroscopyDispatch('vector', None, 'embedded', 'NL') + SpectroscopyDispatch('tensor', 'fullstate', 'embedded', 'LINEAR') + SpectroscopyDispatch('tensor', 'number', 'vacuum', 'NL') + + Parameters + ---------- + 1. traj : uninitialized trajectory object + 2. n_site : int + Number of physical sites (from len(list_transition_dipoles)). + + Returns + ------- + 1. spectroscopy_key : SpectroscopyDispatch + Canonical dispatch tuple. + + Raises + ------ + ValueError + If the trajectory is already initialized or the Hilbert + dimension is not in {n_site, n_site + 1}. + + UnsupportedRequest + On unsupported EOM or tensor method / convention combinations. + """ + if traj.__initialized__: + raise ValueError('initialized trajectory is not valid for spectroscopy dispatch') + eom = traj.basis.eom.param['EQUATION_OF_MOTION'] + if eom == 'NONLINEAR': + eom_tag = 'NL' + elif eom == 'LINEAR': + eom_tag = 'LINEAR' + else: + raise UnsupportedRequest(eom, 'nondyadic_spectroscopy') + + n_state_traj = traj.basis.system.param['NSTATES'] + if n_state_traj == n_site + 1: + hilbert_tag = 'embedded' + elif n_state_traj == n_site: + hilbert_tag = 'excited_only' + else: + raise ValueError( + f'trajectory dim {n_state_traj} != n_site ({n_site}) ' + f'or n_site+1 ({n_site + 1})', + ) + + wf = getattr(traj, 'wavefunction', None) + if wf is None: + return SpectroscopyDispatch('vector', None, hilbert_tag, eom_tag) + + method = traj.tensor_param['METHOD'] + if method == 'number': + # Number representation always uses the vacuum convention (ground is + # the all-zeros config, dim n_site); the embedded layout (n_site+1) + # is not supported. + if hilbert_tag == 'excited_only': + wf.flag_gs_vacuum = True + return SpectroscopyDispatch('tensor', 'number', 'vacuum', eom_tag) + raise UnsupportedRequest( + _format_dispatch( + SpectroscopyDispatch('tensor', 'number', hilbert_tag, eom_tag) + ), + 'nondyadic_spectroscopy', + ) + + if method != 'fullstate': + raise UnsupportedRequest(method, 'nondyadic_spectroscopy') + + if hilbert_tag in ('embedded', 'excited_only'): + return SpectroscopyDispatch('tensor', 'fullstate', hilbert_tag, eom_tag) + + raise UnsupportedRequest( + _format_dispatch( + SpectroscopyDispatch('tensor', 'fullstate', hilbert_tag, eom_tag) + ), + 'nondyadic_spectroscopy', + ) + + +def _format_dispatch(dispatch): + if dispatch.tensor_method is None: + return f'{dispatch.traj_kind}_{dispatch.hilbert_tag}_{dispatch.eom_tag}' + return ( + f'{dispatch.traj_kind}_{dispatch.tensor_method}_' + f'{dispatch.hilbert_tag}_{dispatch.eom_tag}' + ) + + +def _readout_prefactor(traj, norm_ratio, list_norm_sq): + """ + Returns the prefactor that multiplies in the + spectroscopy readout, branching on the EOM. + + NORMALIZED NONLINEAR (the validated dyadic reference) and plain + NONLINEAR are rescale-equivalent: operator_expectation divides + by , so is invariant under psi -> c(t) psi, and + plain NONLINEAR is just NORMALIZED NONLINEAR multiplied by a + complex scalar c(t) with |c(t)|^2 = norm_ratio / ||psi(t)||^2. The + dyadic readout norm_ratio * / ||psi||^2 is therefore + identical per-realization between the two EOMs. + + LINEAR uses raw noise (no Girsanov shift, no z_mem evolution) + and is NOT a rescaling of NORMALIZED NONLINEAR — different + measure entirely. The (N+1)-embedded operators decouple |g> + from the bath in H and L under all EOMs, and neither LINEAR + nor plain NONLINEAR has a rescaling term, so psi[0]=1 holds + deterministically under both — the difference is in the noise + measure, not psi[0] dynamics. Under LINEAR's raw measure the + unbiased per-trajectory estimator of C(t) is + directly; the dyadic norm_ratio/||psi||^2 prefactor has no scaling + identity to support it under LINEAR and would distort the + readout (||psi(t)||^2 random-walks freely without the Girsanov + shift). + + Parameters + ---------- + 1. traj : trajectory object + An initialized trajectory. + + 2. norm_ratio : float + Product of operator-norm changes from each + raise/lower step (>= 1 for the operators built here). + + 3. list_norm_sq : np.ndarray(float) + ||psi(t)||^2 at each detection-phase timestep. + + Returns + ------- + 1. prefactor : float or np.ndarray(float) + Scalar 1.0 under LINEAR; per-timestep + norm_ratio / list_norm_sq otherwise. + """ + if traj.basis.eom.param['EQUATION_OF_MOTION'] == 'LINEAR': + return 1.0 + return norm_ratio / list_norm_sq + + +def _build_operators_abs(list_transition_dipoles, E_1): + """ + Builds the raising and response operators for absorption. + + Parameters + ---------- + 1. list_transition_dipoles : np.ndarray, shape (n_site, 3) + Transition dipole moments per chromophore. + + 2. E_1 : np.ndarray, shape (3,) or (3, 1) + Field polarization (E_sig = E_1 for absorption). + + Returns + ------- + 1. O2_raise : sparse.coo_matrix, shape (n_state, n_state) + Excitation operator: maps |g> to sum_i (mu_i . E)|e_i>, plus + ground-state identity to preserve |g>. + + 2. F2_dense : np.ndarray, shape (n_state, n_state) + Response operator for the expectation value + / ||psi||^2. + Only row 0 is nonzero: F[0, i+1] = mu_i . E. + """ + # Convention: in the wavefunction picture, this operator excites: + # O|g> = sum_i (mu_i . E)|e_i>. In the density matrix picture, + # |e> -> |e_i>, (0, 0) keeps |g> + raise_data = np.concatenate([list_mu_dot_E, [1.0]]) + raise_row = np.concatenate([np.arange(1, n_state), [0]]) + raise_col = np.zeros(n_state, dtype=int) + O2_raise = sparse.coo_matrix( + (raise_data, (raise_row, raise_col)), + shape=(n_state, n_state), + ) + + # Response operator: only row 0 nonzero, F[0, i+1] = mu_i . E + F2_dense = np.zeros((n_state, n_state)) + F2_dense[0, 1:] = list_mu_dot_E + + return O2_raise, F2_dense + + +def _build_operators_fluor(list_transition_dipoles, E_1, E_sig): + """ + Builds the raising, lowering+identity, and response operators for + fluorescence. + + Parameters + ---------- + 1. list_transition_dipoles : np.ndarray, shape (n_site, 3) + Transition dipole moments per chromophore. + + 2. E_1 : np.ndarray, shape (3,) or (3, 1) + Field polarization for excitation. + + 3. E_sig : np.ndarray, shape (3,) or (3, 1) + Signal field polarization (also E_3 for lowering). + + Returns + ------- + 1. O2_raise : sparse.coo_matrix, shape (n_state, n_state) + Excitation operator: maps |g> to sum_i (mu_i . E_1)|e_i>. + + 2. O2_lower_ident : sparse.coo_matrix, shape (n_state, n_state) + Lowering |g>. + Only row 0 is nonzero: F[0, i+1] = mu_i . E_sig. + """ + # Convention: see _build_operators_abs for wavefunction vs density + # matrix naming convention for raising/lowering operators. + E_1 = np.asarray(E_1).ravel() + E_sig = np.asarray(E_sig).ravel() + n_site = len(list_transition_dipoles) + n_state = n_site + 1 + list_mu_dot_E1 = list_transition_dipoles @ E_1 + list_mu_dot_Esig = list_transition_dipoles @ E_sig + + # Raising operator: nonzero at (i+1, 0), maps |g> -> |e_i> + O2_raise = sparse.coo_matrix( + (list_mu_dot_E1, (np.arange(1, n_state), np.zeros(n_site, dtype=int))), + shape=(n_state, n_state), + ) + + # Lowering operator: row 0 has |g> of + traj.psi, which already carries the GS amplitude at slot 0. + + Parameters + ---------- + 1. traj : trajectory object + An initialized trajectory. + + 2. apply_operator : callable + No-arg callable that performs the operator + application on traj in place, e.g. + ``lambda: traj._operator(O2.toarray())`` for the + embedded key or + ``lambda: traj.apply_dipole_raise(list_mu)`` for the + vacuum key. + + 3. spectroscopy_key : SpectroscopyDispatch + Spectroscopy dispatch tuple from + _spectroscopy_key. + + Returns + ------- + 1. norm_ratio : float + Norm-squared ratio ||O psi||^2 / ||psi||^2. + """ + if spectroscopy_key.tensor_method == 'number' and spectroscopy_key.hilbert_tag == 'vacuum': + norm_sq_pre = traj.wavefunction.manifold_norm_sq + apply_operator() + norm_sq_post = traj.wavefunction.manifold_norm_sq + else: + psi_pre = traj.psi + norm_sq_pre = np.dot(np.conj(psi_pre), psi_pre).real + apply_operator() + psi_post = traj.psi + norm_sq_post = np.dot(np.conj(psi_post), psi_post).real + return norm_sq_post / norm_sq_pre + + +def calc_absorption_response(traj, list_transition_dipoles, E_1, t_max, t_step): + """ + Non-dyadic linear absorption C(t) for a single trajectory. + + Parameters + ---------- + 1. traj : uninitialized trajectory object + Any type supporting .initialize(), ._operator(), + .propagate(), and .storage['psi_traj']. Hilbert + dimension must equal len(list_transition_dipoles) (vacuum tensor or + excited-only LINEAR) or len(list_transition_dipoles) + 1 (embedded). + Conventions: + - embedded: an explicit ground-state basis slot is + present in the Hilbert space. + - vacuum: the ground state is implicit in the number + representation and corresponds to the all-zero MPS + configuration. + - excited_only: only the excited manifold is + propagated, and only for the LINEAR shortcut. + + 2. list_transition_dipoles : np.ndarray, shape (n_site, 3) + Transition dipole moments per chromophore. + + 3. E_1 : np.ndarray, shape (3,) + Field polarization (also used as E_sig for absorption). + + 4. t_max : float [units: fs] + Propagation time after excitation. + + 5. t_step : float [units: fs] + Time step. + + Returns + ------- + 1. C1_corr_t : np.ndarray(complex) + Absorption correlation function C(t) sampled at + t = t_step, 2*t_step, .... + """ + n_site = len(list_transition_dipoles) + spectroscopy_key = _spectroscopy_key(traj, n_site) + traj_kind, tensor_method, hilbert_tag, eom_tag = spectroscopy_key + + if hilbert_tag == 'embedded': + # vector + fullstate tensor + number tensor (GS embedded in + # the n_site+1 dim Hilbert space). + O2_raise, F2_dense = _build_operators_abs(list_transition_dipoles, E_1) + P1_psi_0 = np.zeros(n_site + 1, dtype=np.complex128) + P1_psi_0[0] = 1.0 + traj.initialize(P1_psi_0) + norm_ratio = _apply_op_and_track_norm( + traj, lambda: traj._operator(O2_raise.toarray()), spectroscopy_key, + ) + traj.propagate(t_max, t_step) + + # Skip psi_traj[0]: the ground state stored by initialize() + # before the raise. C(t) starts at t = t_step. + psi_traj_slice = np.asarray(traj.storage['psi_traj'])[1:] + list_norm_sq = np.sum(np.conj(psi_traj_slice) * psi_traj_slice, axis=1) + # psi @ F.T computes F @ psi[t] for all timesteps simultaneously. + psi_f_traj = psi_traj_slice @ F2_dense.T + list_expectation = np.sum(np.conj(psi_traj_slice) * psi_f_traj, axis=1) + return 2 * _readout_prefactor(traj, norm_ratio, list_norm_sq) * list_expectation + + if tensor_method == 'number' and hilbert_tag == 'vacuum': + # Vacuum tensor: |g> is the all-zeros MPS configuration; + # raise is applied as a bond-dim MPO. + E_1 = np.asarray(E_1).ravel() + list_mu_dot_E = list_transition_dipoles @ E_1 + P1_psi_0 = np.zeros(n_site, dtype=np.complex128) + traj.initialize(P1_psi_0) + # |psi_g|^2 = 1 by absorption's decoupled-GS assumption, so + # manifold_norm_sq mirrors the embedded key's + # ||O psi||^2 / ||psi||^2. + GS_AMP_SQ = 1.0 + norm_ratio = _apply_op_and_track_norm( + traj, + lambda: traj.apply_dipole_raise_plus_ground_ident(list_mu_dot_E), + spectroscopy_key, + ) + traj.propagate(t_max, t_step) + psi_traj_slice = np.asarray(traj.storage['psi_traj'])[1:] + list_norm_sq = np.sum( + np.conj(psi_traj_slice) * psi_traj_slice, axis=1, + ).real + list_response = psi_traj_slice @ list_mu_dot_E + return 2 * _readout_prefactor(traj, norm_ratio, list_norm_sq + GS_AMP_SQ) * list_response + + if eom_tag == 'LINEAR' and hilbert_tag == 'excited_only': + # Excited-only LINEAR optimization: raise is absorbed into the + # init via psi-linearity. psi_unnorm(t) = sqrt(norm_ratio) * psi(t), + # so C(t) = 2 * = 2 * sqrt(norm_ratio) * (psi(t) . V*). + E_1 = np.asarray(E_1).ravel() + list_mu_dot_E = list_transition_dipoles @ E_1 + # This is the norm of the initial excited-only seed vector, not the + # nonlinear readout prefactor used in the other branches. + mu_dot_e_norm_sq = float(np.dot(np.conj(list_mu_dot_E), list_mu_dot_E).real) + if mu_dot_e_norm_sq == 0.0: + raise ValueError( + 'mu . E has zero norm; cannot initialize excited-only ' + 'LINEAR absorption trajectory' + ) + P1_psi_0 = np.zeros(n_site, dtype=np.complex128) + P1_psi_0[:] = list_mu_dot_E / np.sqrt(mu_dot_e_norm_sq) + traj.initialize(P1_psi_0) + traj.propagate(t_max, t_step) + # Skip psi_traj[0] so the output time grid matches the + # embedded key: C(t) starts at t = t_step. + psi_traj_slice = np.asarray(traj.storage['psi_traj'])[1:] + return 2 * np.sqrt(mu_dot_e_norm_sq) * ( + psi_traj_slice @ np.conj(list_mu_dot_E) + ) + + raise UnsupportedRequest( + f'unsupported key {_format_dispatch(spectroscopy_key)}', + 'calc_absorption_response', + ) + + +def calc_fluorescence_response(traj, list_transition_dipoles, E_1, E_sig, t2, t3_max, t_step): + """ + Non-dyadic fluorescence C(t) for a single trajectory. + + Excites from the ground state with mu.E_1, waits for t2, applies + the (lower + I_excited) operator weighted by mu.E_sig, propagates + for t3_max, and reads out the fluorescence correlation function + from the bilinear / . The returned array is + sampled at t = t_step, 2*t_step, ... + + Parameters + ---------- + 1. traj : uninitialized trajectory object + Any type supporting .initialize(), ._operator(), + .propagate(), and .storage['psi_traj']. Hilbert + dimension must equal len(list_transition_dipoles) + 1 (embedded) or + len(list_transition_dipoles) (vacuum tensor). Conventions: + - embedded: an explicit ground-state basis slot is + present in the Hilbert space. + - vacuum: the ground state is implicit in the number + representation and corresponds to the all-zero MPS + configuration. + - excited_only: only the excited manifold is + propagated, and only for the LINEAR shortcut. + + 2. list_transition_dipoles : np.ndarray, shape (n_site, 3) + Transition dipole moments per chromophore. + + 3. E_1 : np.ndarray, shape (3,) + Field polarization for excitation. + + 4. E_sig : np.ndarray, shape (3,) + Signal field polarization (also used as E_3 for lowering). + + 5. t2 : float [units: fs] + Waiting time in the excited manifold before detection. + + 6. t3_max : float [units: fs] + Detection time after de-excitation. + + 7. t_step : float [units: fs] + Time step. + + Returns + ------- + 1. C1_corr_t : np.ndarray(complex) + Fluorescence correlation function C(t) sampled at + t = t_step, 2*t_step, .... + """ + n_site = len(list_transition_dipoles) + spectroscopy_key = _spectroscopy_key(traj, n_site) + traj_kind, tensor_method, hilbert_tag, eom_tag = spectroscopy_key + + if t2 % t_step > 1e-10 * t_step: + raise ValueError( + f't2={t2} is not an integer multiple of t_step={t_step}.' + ) + idx_t2 = round(t2 / t_step) + + if hilbert_tag == 'embedded': + # vector + fullstate tensor + number tensor (GS embedded in + # the n_site+1 dim Hilbert space). + O2_raise, O2_lower_ident, F2_dense = _build_operators_fluor( + list_transition_dipoles, E_1, E_sig + ) + P1_psi_0 = np.zeros(n_site + 1, dtype=np.complex128) + P1_psi_0[0] = 1.0 + traj.initialize(P1_psi_0) + + norm_ratio_1 = _apply_op_and_track_norm( + traj, lambda: traj._operator(O2_raise.toarray()), spectroscopy_key, + ) + traj.propagate(t2, t_step) + norm_ratio_2 = _apply_op_and_track_norm( + traj, lambda: traj._operator(O2_lower_ident.toarray()), spectroscopy_key, + ) + traj.propagate(t3_max, t_step) + + psi_traj_slice = np.asarray(traj.storage['psi_traj'])[idx_t2 + 1:] + list_norm_sq = np.sum(np.conj(psi_traj_slice) * psi_traj_slice, axis=1) + # psi @ F.T computes F @ psi[t] for all timesteps simultaneously. + psi_f_traj = psi_traj_slice @ F2_dense.T + list_expectation = np.sum(np.conj(psi_traj_slice) * psi_f_traj, axis=1) + return 4 * _readout_prefactor( + traj, norm_ratio_1 * norm_ratio_2, list_norm_sq, + ) * list_expectation + + if tensor_method == 'number' and hilbert_tag == 'vacuum': + # Vacuum tensor: |g> is the all-zeros MPS configuration; + # raise and (lower + I) are bond-dim MPOs. + if t3_max % t_step > 1e-10 * t_step: + raise ValueError( + f't3_max={t3_max} is not an integer multiple of ' + f't_step={t_step}.' + ) + # Capture the per-step GS amplitude for the detection-phase readout. + if not traj.storage.storage_dic.get('psi_g_traj', False): + traj.storage.storage_dic['psi_g_traj'] = True + traj.storage.adaptive = traj.storage.adaptive + + list_mu_dot_E1 = list_transition_dipoles @ np.asarray(E_1).ravel() + list_mu_dot_Esig = list_transition_dipoles @ np.asarray(E_sig).ravel() + # Response operator: only F[0, k+1] = mu_k . E_sig is nonzero. + F2_dense = np.zeros((n_site + 1, n_site + 1)) + F2_dense[0, 1:] = list_mu_dot_Esig + + P1_psi_0 = np.zeros(n_site, dtype=np.complex128) + traj.initialize(P1_psi_0) + + # raise(E_1): |g> -> sum_k mu_k(E_1) |e_k>. + norm_ratio_1 = _apply_op_and_track_norm( + traj, lambda: traj.apply_dipole_raise(list_mu_dot_E1), spectroscopy_key, + ) + traj.propagate(t2, t_step) + # (I + mu^-) with E_sig: +I preserves the single-ex content + # of psi(t2); the mu^- part injects the GS amplitude that + # creates the G/E coherence the detection-phase response + # measures. + norm_ratio_2 = _apply_op_and_track_norm( + traj, + lambda: traj.apply_dipole_lower_plus_ident(list_mu_dot_Esig), + spectroscopy_key, + ) + traj.propagate(t3_max, t_step) + + # Stitch psi_g (slot 0) and the single-excitation amplitudes + # (slots 1..n_site) into a length-(n_site+1) state vector per + # timestep. + psi_traj_slice = np.asarray(traj.storage['psi_traj'])[idx_t2 + 1:] + gs_amp_traj = np.asarray(traj.storage['psi_g_traj'])[idx_t2 + 1:] + psi_full_traj = np.column_stack([gs_amp_traj, psi_traj_slice]) + + list_norm_sq = np.sum( + np.conj(psi_full_traj) * psi_full_traj, axis=1, + ).real + # psi @ F.T computes F @ psi[t] for all timesteps simultaneously. + psi_f_traj = psi_full_traj @ F2_dense.T + list_expectation = np.sum(np.conj(psi_full_traj) * psi_f_traj, axis=1) + return 4 * _readout_prefactor( + traj, norm_ratio_1 * norm_ratio_2, list_norm_sq, + ) * list_expectation + + raise UnsupportedRequest( + f'unsupported key {_format_dispatch(spectroscopy_key)}', + 'calc_fluorescence_response', + ) diff --git a/src/mesohops/util/tensor_operations.py b/src/mesohops/util/tensor_operations.py new file mode 100644 index 0000000..f6781f2 --- /dev/null +++ b/src/mesohops/util/tensor_operations.py @@ -0,0 +1,909 @@ +""" +Core MPS arithmetic and extraction routines used by HopsTensorWavefunction. + +Functions +--------- +tensor_add(list_tensor_cores_1, list_tensor_cores_2, epsilon, bond_dim_max) + Add two MPS by block-concatenating their cores, then compress. + +tensor_compress(list_cores, epsilon, bond_dim_max) + Right-orthogonalize then left-sweep with SVD truncation (Oseledets + rounding algorithm). + +extract_psi(list_cores_phi, method, M1_modes_per_state) + Extract the physical wavefunction (zero-auxiliary slice) from an MPS. + +extract_gs_amp(list_cores_phi, method) + Extract the amplitude of the all-phys-zero configuration of an MPS. + In the vacuum convention this is the ground-state amplitude. + +phi_aux(list_cores_phi, method, M1_modes_per_state, indices) + Extract a specific auxiliary-state vector from an MPS given its + per-mode occupation indices. + +tensor_to_array(list_cores_phi, method, M1_modes_per_state, system, mode) + Flatten the ground and first-order auxiliary wavefunctions into a + single adHOPS-style array, for debugging purposes. + +contract_down(list_cores_phi, method, n_state) + Approximate per-state norm squared (sum over hierarchy) by contracting + squared core elements independently. + +contract_down_exact(list_cores_phi, method, n_state) + Exact per-state norm squared via a double-layer right-to-left sweep. +""" +from __future__ import annotations + +import numpy as np +from scipy.linalg import svd as scipy_svd + +from mesohops.basis.basis_functions import determine_error_thresh +from mesohops.basis.hops_modes import HopsModes +from mesohops.basis.hops_system import HopsSystem + +__title__ = 'Tensor Operations' +__author__ = 'B. Z. Citty' +__maintainer__ = 'B. Z. Citty' + + +def flatten_cores(list_cores_phi: list[list[np.ndarray]]) -> list[np.ndarray]: + """Flatten a statenumber list-of-lists MPS into a flat list of cores. + + Parameters + ---------- + 1. list_cores_phi: list(list(np.ndarray)) + Nested MPS: list_cores_phi[s] = [state_core_s, + mode_core_s0, mode_core_s1, ...] + + Returns + ------- + 1. list_cores_flat: list(np.ndarray) + 1D list of cores (one ndarray per element) in + statenumber MPS representation order: + state_0, mode_00, mode_01, ..., state_1, mode_10, ... + """ + return [core for group in list_cores_phi for core in group] + + +def unflatten_cores( + list_cores_flat: list[np.ndarray], + list_modes_per_site: list[int], +) -> list[list[np.ndarray]]: + """Restore list-of-lists MPS structure from a flat core list. + + Parameters + ---------- + 1. list_cores_flat: list(np.ndarray) + 1D list of cores (one ndarray per element) in + statenumber MPS representation order. + 2. list_modes_per_site: list(int) + Number of mode cores per site s + (i.e. len(list_cores_phi[s]) - 1). + + Returns + ------- + 1. list_cores_phi: list(list(np.ndarray)) + Nested MPS: list_cores_phi[s] = [state_core_s, + mode_core_s0, mode_core_s1, ...] + """ + result = [] + idx = 0 + for n_modes in list_modes_per_site: + result.append(list_cores_flat[idx: idx + 1 + int(n_modes)]) + idx += 1 + int(n_modes) + return result + + +def _flat_cores_with_labels( + list_cores_phi: list[list[np.ndarray]], +) -> list[tuple[np.ndarray, bool, int]]: + """Return a list of (core, is_state_core, state_idx) tuples in + statenumber MPS representation order. + + Parameters + ---------- + 1. list_cores_phi: list(list(np.ndarray)) + Nested MPS: list_cores_phi[s] = [state_core_s, + mode_core_s0, mode_core_s1, ...] + + Returns + ------- + 1. list_labeled: list(tuple(np.ndarray, bool, int)) + Each entry is (core, is_state_core, state_idx) where + is_state_core is True for state cores and False for + mode cores, and state_idx is the site index s. + """ + result = [] + for s, group in enumerate(list_cores_phi): + result.append((group[0], True, s)) + for core_m in group[1:]: + result.append((core_m, False, s)) + return result + + +def _statenumber_offsets(M1_modes_per_state: np.ndarray) -> np.ndarray: + """Compute the MPS core index of each state core. + + In number representation, the MPS layout is: + [state_0][mode_0_0]...[mode_0_M0][state_1][mode_1_0]... + This function returns the index of each state core. + + Parameters + ---------- + 1. M1_modes_per_state: array-like(int) + Number of mode cores per state. + + Returns + ------- + 1. M1_offsets: np.ndarray(int) + Core index of each state core. + """ + # Stride per state is 1 (the state core) + the number of mode cores; + # offsets are the cumulative stride prefix with a leading 0. + M1_strides = 1 + np.asarray(M1_modes_per_state, dtype=int) + M1_offsets = np.zeros(len(M1_strides), dtype=int) + if len(M1_strides) > 0: + M1_offsets[1:] = np.cumsum(M1_strides[:-1]) + return M1_offsets + + +def extract_gs_amp( + list_cores_phi: list[np.ndarray] | list[list[np.ndarray]], + method: str, +) -> np.complex128: + """ + Amplitude of the all-phys-zero configuration of the MPS. + + For number representation this is the amplitude of the + configuration with every state core in |0> and every mode core in + the hierarchy-ground state. In the ground-state-as-vacuum + convention this is the physical ground-state amplitude; in the + GS-as-state-core convention it is an unphysical "no state core + occupied" configuration whose amplitude is typically ~0 for + single-occupancy physical states. + + Parameters + ---------- + 1. list_cores_phi : list(np.ndarray) | list(list(np.ndarray)) + MPS cores. Nested statenumber groups are + flattened before contraction. + + 2. method : str + 'number' or 'fullstate'. + + Returns + ------- + 1. amp : np.complex128 + <0,0,...,0|psi>. + """ + if method == 'number': + list_cores = [] + for group in list_cores_phi: + list_cores.extend(group) + elif method == 'fullstate': + list_cores = list_cores_phi + else: + raise ValueError(f'Unknown method {method!r}.') + + if not list_cores: + raise ValueError( + 'extract_gs_amp requires at least one core; ' + 'got empty list_cores after flattening.' + ) + + M2_env = np.array([[1.0]], dtype=np.complex128) + for core in list_cores: + # core shape (l, p, r); slice at p=0 to get (l, r) and chain. + M2_env = M2_env @ core[:, 0, :] + return M2_env[0, 0] + + +def extract_psi( + list_cores_phi: list[np.ndarray] | list[list[np.ndarray]], + method: str, + M1_modes_per_state: np.ndarray, +) -> np.ndarray: + """ + Extracts the physical (zero-auxiliary) wavefunction from an MPS. + + Slices each mode core at occupation index 0 and contracts the + resulting bond matrices. For number representation, each state + core is sliced at physical index 1 (occupied) or 0 (unoccupied) + according to a one-hot encoding. + + Parameters + ---------- + 1. list_cores_phi: list(np.ndarray) | list(list(np.ndarray)) + MPS cores of the current wavefunction. For + fullstate, a flat list + [state_core, mode_core_0, ...]. For + number, nested: + list_cores_phi[s] = [state_core_s, mode_core_s0, ...] + 2. method: str + Tensor encoding type ('fullstate' or 'number'). + 3. M1_modes_per_state: np.ndarray(int) + Number of bath modes per state. + + Returns + ------- + 1. V1_phi_result: np.ndarray(complex) + Physical wavefunction, shape (n_state,). + """ + M1_modes_per_state = np.asarray(M1_modes_per_state, dtype=int) + n_total_modes = int(np.sum(M1_modes_per_state)) + + if method == 'fullstate': + list_cores_sliced = [list_cores_phi[0]] + for i in range(n_total_modes): + core_m = list_cores_phi[i + 1] + list_cores_sliced.append(core_m[:, 0, :]) + M2_env = list_cores_sliced[0] + for i in range(n_total_modes): + M2_env = np.tensordot(M2_env, list_cores_sliced[i + 1], 1) + # M2_env shape: (1, n_state, 1) — collapse trivial OBC boundary dims + V1_phi_result = M2_env[0, :, 0] + + elif method == 'number': + # list_cores_phi[s] = [state_core_s, mode_core_s0, ...] + n_state = len(list_cores_phi) + + # Pre-contract each group's mode-core chain (occ=0 slices) into a + # bond matrix. These products don't depend on target_state, so + # computing them once drops the per-target work from + # O(n_total_cores * chi^2) to O(n_state * chi^3). + list_mode_envs = [] + for group in list_cores_phi: + dim = group[0].shape[-1] + M2_mode_env = np.eye(dim, dtype=np.complex128) + for core_m in group[1:]: + M2_mode_env = M2_mode_env @ core_m[:, 0, :] + list_mode_envs.append(M2_mode_env) + + # Per-target contraction only touches the state cores and the + # pre-contracted mode envs. The one-hot slicing selects phys=1 + # at target_state's axis and phys=0 elsewhere. + V1_phi_result = np.zeros(n_state, dtype=np.complex128) + for target_state in range(n_state): + phys = 1 if target_state == 0 else 0 + M2_env = list_cores_phi[0][0][:, phys, :] @ list_mode_envs[0] + for s in range(1, n_state): + phys = 1 if s == target_state else 0 + M2_env = ( + M2_env + @ list_cores_phi[s][0][:, phys, :] + @ list_mode_envs[s] + ) + V1_phi_result[target_state] = M2_env[0, 0] + else: + raise ValueError(f'Unknown method {method!r}.') + return V1_phi_result + +def phi_aux( + list_cores_phi: list[np.ndarray] | list[list[np.ndarray]], + method: str, + M1_modes_per_state: np.ndarray, + indices: list[int], +) -> np.ndarray: + """ + Extracts an auxiliary-state wavefunction from an MPS. + + Slices each mode core at the occupation number given by indices, + then contracts the resulting bond matrices to yield a state vector. + For indices = [0, 0, ..., 0], this returns phi_0. + + Parameters + ---------- + 1. list_cores_phi: list(np.ndarray) | list(list(np.ndarray)) + MPS cores of the current wavefunction. For + fullstate, a flat list + [state_core, mode_core_0, ...]. For + number, nested: + list_cores_phi[s] = [state_core_s, mode_core_s0, ...] + 2. method: str + Tensor encoding type ('fullstate' or + 'number'). + 3. M1_modes_per_state: np.ndarray(int) + Number of bath modes per state. + 4. indices: list(int) + Per-mode occupation numbers selecting which auxiliary + to extract. Length must equal total number of modes. + + Returns + ------- + 1. phi_aux: np.ndarray(complex) + Auxiliary wavefunction, shape (n_state,). + """ + M1_modes_per_state = np.asarray(M1_modes_per_state, dtype=int) + n_total_modes = int(np.sum(M1_modes_per_state)) + + if method == 'fullstate': + list_cores_sliced = [list_cores_phi[0]] + for i in range(n_total_modes): + core_m = list_cores_phi[i + 1] + list_cores_sliced.append(core_m[:, indices[i], :]) + M2_env = list_cores_sliced[0] + for i in range(n_total_modes): + M2_env = np.tensordot(M2_env, list_cores_sliced[i + 1], 1) + # M2_env shape: (1, n_state, 1) — collapse trivial OBC boundary dims + V1_phi_result = M2_env[0, :, 0] + + elif method == 'number': + # list_cores_phi[s] = [state_core_s, mode_core_s0, ...] + n_state = len(list_cores_phi) + + # Pad indices with zeros for padded mode cores (non-bath states + # have extra cores beyond the real mode count). + n_total_padded = sum(len(g) - 1 for g in list_cores_phi) + if len(indices) < n_total_padded: + indices = list(indices) + [0] * (n_total_padded - len(indices)) + + # Build sliced list from the labeled-core view: state cores kept + # as-is (Dl, 2, Dr), mode cores sliced at the requested occupation + # index. mode_flat_idx tracks the running position into `indices` + # across all mode cores. + list_cores_sliced = [] + mode_flat_idx = 0 + for core, is_state_core, _ in _flat_cores_with_labels(list_cores_phi): + if is_state_core: + list_cores_sliced.append(core) + else: + list_cores_sliced.append(core[:, indices[mode_flat_idx], :]) + mode_flat_idx += 1 + + # Contract full chain. State cores contribute a dim-2 axis each; + # mode slices are 2D and contract purely over bond indices. + # Pre-collapse shape: (1, 2, 2, ..., 2, 1). The trailing indexing + # drops the trivial OBC boundary dims, leaving (2, 2, ..., 2) with + # one axis per system state. + M2_env = list_cores_sliced[0] + for core in list_cores_sliced[1:]: + M2_env = np.tensordot(M2_env, core, 1) + M2_env = M2_env[0, ..., 0] + + # Extract amplitude for each state by selecting index 1 at that state's + # axis and 0 everywhere else. Reset the previous state's index before + # advancing to avoid stale 1s — starting from i=1 so i=0 is never + # "reset" (it is freshly set on the first iteration). + list_state_idx = [0] * n_state + V1_phi_result = np.zeros(n_state, dtype=np.complex128) + for i in range(n_state): + if i > 0: + list_state_idx[i - 1] = 0 + list_state_idx[i] = 1 + V1_phi_result[i] = M2_env[tuple(list_state_idx)] + else: + raise ValueError(f'Unknown method {method!r}.') + + return V1_phi_result + +def tensor_add( + list_tensor_cores_1: list, + list_tensor_cores_2: list, + epsilon: float, + bond_dim_max: int, +) -> list: + """ + Adds two MPS (or MPO) tensors and compresses the result. + + Accepts flat lists of cores or nested statenumber MPS (list-of-lists). + For nested input the structure is flattened internally and restored on + output. For MPO inputs (rank-4 cores) the two physical indices are fused + before addition and restored afterward. + + Block-concatenates corresponding cores, then calls tensor_compress. + + Reference: Oseledets, 'Tensor-Train Decomposition', SIAM J. Sci. + Comput. 33(5), 2295-2317 (2011), Section 2.3. + + Parameters + ---------- + 1. list_tensor_cores_1: list + First MPS/MPO. Either a flat list of cores + or a nested list-of-lists (statenumber format). + 2. list_tensor_cores_2: list + Second MPS/MPO, same format as list_tensor_cores_1. + 3. epsilon: float + SVD truncation threshold for compression. + 4. bond_dim_max: int + Maximum bond dimension after compression. + + Returns + ------- + 1. list_cores_sum: list + Compressed cores of the sum, same structure as inputs. + """ + # Detect nested statenumber MPS (list-of-lists) and flatten before algebra. + is_nested = isinstance(list_tensor_cores_1[0], list) + if is_nested: + list_modes_per_site = [len(g) - 1 for g in list_tensor_cores_1] + list_tensor_cores_1 = flatten_cores(list_tensor_cores_1) + list_tensor_cores_2 = flatten_cores(list_tensor_cores_2) + + if len(list_tensor_cores_1) != len(list_tensor_cores_2): + raise ValueError( + f'MPS/MPO core count mismatch: ' + f'{len(list_tensor_cores_1)} != {len(list_tensor_cores_2)}.' + ) + # Verify open boundary conditions: left bond of first core and right + # bond of last core must each be 1 (no dangling virtual indices). + if list_tensor_cores_1[0].shape[0] != 1: + raise ValueError( + f'list_tensor_cores_1: left bond of first core must be 1, ' + f'got {list_tensor_cores_1[0].shape[0]}.' + ) + if list_tensor_cores_1[-1].shape[-1] != 1: + raise ValueError( + f'list_tensor_cores_1: right bond of last core must be 1, ' + f'got {list_tensor_cores_1[-1].shape[-1]}.' + ) + if list_tensor_cores_2[0].shape[0] != 1: + raise ValueError( + f'list_tensor_cores_2: left bond of first core must be 1, ' + f'got {list_tensor_cores_2[0].shape[0]}.' + ) + if list_tensor_cores_2[-1].shape[-1] != 1: + raise ValueError( + f'list_tensor_cores_2: right bond of last core must be 1, ' + f'got {list_tensor_cores_2[-1].shape[-1]}.' + ) + list_cores_sum = [] + # Direct-sum (block-diagonal) concatenation of MPS cores. + # Boundary cores (first/last) keep bond dim 1 on the open end; + # interior cores are block-diagonal in both bond indices. + for core_idx in range(len(list_tensor_cores_1)): + core_1 = list_tensor_cores_1[core_idx] + core_2 = list_tensor_cores_2[core_idx] + # Bond dims are always the first and last axes; physical dims are + # everything in between (one axis for MPS, two for MPO, etc.) + bond_left_1, bond_right_1 = core_1.shape[0], core_1.shape[-1] + bond_left_2, bond_right_2 = core_2.shape[0], core_2.shape[-1] + phys_shape = core_1.shape[1:-1] + if core_1.shape[1:-1] != core_2.shape[1:-1]: + raise ValueError( + f'Physical dimension mismatch at core {core_idx}: ' + f'{core_1.shape[1:-1]} != {core_2.shape[1:-1]}.' + ) + if core_idx == 0: + # Left boundary: left bond stays 1, right bond concatenated + core = np.zeros( + (1, *phys_shape, bond_right_1 + bond_right_2), + dtype=np.complex128, + ) + core[..., :bond_right_1] = core_1 + core[..., bond_right_1:] = core_2 + elif core_idx == len(list_tensor_cores_1) - 1: + # Right boundary: left bond concatenated, right bond stays 1 + core = np.zeros( + (bond_left_1 + bond_left_2, *phys_shape, 1), + dtype=np.complex128, + ) + core[:bond_left_1, ...] = core_1 + core[bond_left_1:, ...] = core_2 + else: + # Interior: block-diagonal in both bond indices + core = np.zeros( + (bond_left_1 + bond_left_2, *phys_shape, + bond_right_1 + bond_right_2), + dtype=np.complex128, + ) + core[:bond_left_1, ..., :bond_right_1] = core_1 + core[bond_left_1:, ..., bond_right_1:] = core_2 + list_cores_sum.append(core) + + list_cores_sum = tensor_compress(list_cores_sum, epsilon, bond_dim_max) + if is_nested: + list_cores_sum = unflatten_cores(list_cores_sum, list_modes_per_site) + return list_cores_sum + + +def calc_mps_complexity(list_cores: list[np.ndarray]) -> int: + """ + Computes a scalar complexity proxy for an MPS. + + Each core of shape (D_left, d_phys, D_right) contributes + D_left * D_right * max(D_left, D_right) * d_phys + and the total is summed across cores. This matches the leading cost + of a truncated SVD on the core reshaped as a (D_left * d_phys) x D_right + matrix (or its transpose, whichever orientation is larger), and serves + as a proxy for the overall work per matvec-then-compress cycle. + + Parameters + ---------- + 1. list_cores: list(np.ndarray) + MPS cores, each shaped (D_left, d_phys, D_right). + + Returns + ------- + 1. complexity: int + Sum of per-core complexity scores. + """ + complexity = 0 + for core in list_cores: + D_left, d_phys, D_right = core.shape + complexity += D_left * D_right * max(D_left, D_right) * d_phys + return complexity + + +def scale_mps( + cores: list[np.ndarray] | list[list[np.ndarray]], + factor: complex, +) -> None: + """Scale an MPS in-place by multiplying its first core by factor. + + Works for both flat and nested (statenumber list-of-lists) MPS structures. + Multiplying a single boundary core is equivalent to scaling the whole MPS + because bond-contracted MPS cores form a product. + + Parameters + ---------- + 1. cores: list(np.ndarray) | list(list(np.ndarray)) + MPS cores. For fullstate, a flat list + [state_core, mode_core_0, ...]. For + number, nested: + cores[s] = [state_core_s, mode_core_s0, ...] + 2. factor: complex + Scalar multiplier. + + Returns + ------- + None + """ + if isinstance(cores[0], list): + cores[0][0] = cores[0][0] * factor + else: + cores[0] = cores[0] * factor + + +def tensor_compress( + list_cores: list[np.ndarray], + epsilon: float, + bond_dim_max: int, +) -> list[np.ndarray]: + """ + Compresses an MPS via the Oseledets rounding algorithm. + + Right-orthogonalizes, then left-sweeps with truncated SVD to + reduce bond dimensions while preserving accuracy up to epsilon. + + Reference: Oseledets, 'Tensor-Train Decomposition', SIAM J. Sci. + Comput. 33(5), 2295-2317 (2011), Algorithm 2 (TT-rounding), p. 2305. + + Parameters + ---------- + 1. list_cores: list(np.ndarray) + MPS cores to compress. + 2. epsilon: float + SVD truncation threshold (applied to normalized + singular values). + 3. bond_dim_max: int + Maximum bond dimension after compression. + + Returns + ------- + 1. list_cores_compressed: list(np.ndarray) + Compressed MPS cores. + """ + n_cores = len(list_cores) + list_cores_compressed = [] + list_cores_ortho = [] + core_cur = list_cores[n_cores - 1] + + # === Step 1: Right-orthogonalization (Oseledets Alg. 2, p. 2305, line 1) === + # Sweep right-to-left, factoring each core into R @ Q via RQ + # decomposition. Q (right-orthogonal) is stored; R is absorbed + # into the neighboring core to the left. + # + # Right-orthogonalizing first ensures that during the subsequent + # left-to-right SVD sweep, the singular values at each bond give the + # exact truncation error for that bipartition of the chain. Without + # this gauge fix, the SVD threshold would not have a well-defined + # relationship to the overall approximation error. + # + # RQ for complex matrices uses QR of the conjugate transpose: + # core^H = Qt @ Rt (np.linalg.qr) + # core = Rt^H @ Qt^H = R @ Q + # Qt^H has orthonormal rows: (Qt^H)(Qt^H)^H = Qt^H Qt* = I + # because Qt has orthonormal columns (Qt^H Qt = I). + for i in range(n_cores - 1, 0, -1): + shape = core_cur.shape + # Reshape (Dl, ..., Dr) → (Dl, phys*Dr) to treat as a matrix; + # the left bond index is separated for the RQ factorization + core_cur = core_cur.reshape(shape[0], -1) + Qt, Rt = np.linalg.qr(core_cur.conj().T, mode='reduced') + R = Rt.conj().T # (Dl, chi): factor passed left + Q = Qt.conj().T # (chi, phys*Dr): right-orthogonal factor + # Restore physical indices: (chi, ..., Dr) + Q = Q.reshape(Q.shape[0], *shape[1:]) + list_cores_ortho.append(Q) + # Absorb R into the core to the left, propagating gauge freedom + # leftward so the next iteration operates on the updated core + core_cur = np.tensordot(list_cores[i - 1], R, 1) + list_cores_ortho.append(core_cur) + # Reverse: list was built right-to-left, need left-to-right order + list_cores_ortho.reverse() + + # === Step 2: SVD compression sweep (Oseledets Alg. 2, p. 2305, line 2) === + # Sweep left-to-right through the right-orthogonalized MPS. At each + # bond, the current core is unfolded into a matrix and its SVD + # gives the Schmidt decomposition across that bipartition. Singular + # values below the threshold (or beyond bond_dim_max) are discarded, + # reducing the bond dimension. U is kept as a new left-orthogonal + # core; the remaining weight S*Vh is absorbed into the next core + # to maintain the gauge and propagate the truncation error rightward. + core_cur = list_cores_ortho[0] + for i in range(n_cores - 1): + shape = core_cur.shape + # Reshape (Dl, ..., Dr) → (Dl*phys, Dr) for SVD bipartition; + # left bond and physical indices are merged into the row index + core_cur = core_cur.reshape(-1, shape[-1]) + # scipy.linalg.svd returns the decomposition A = U diag(S) Vh, + # where Vh = V† is the Hermitian conjugate (conjugate transpose) + # of the right singular vector matrix V (Oseledets eq. 2.1, + # p. 2296). The gesvd driver is used for better numerical + # stability than the default gesdd. + U, S, Vh = scipy_svd(core_cur, full_matrices=False, lapack_driver='gesvd') + # Truncation step 1: drop singular values below a relative + # threshold. S is normalized so that the threshold is scale- + # invariant; epsilon^2 is used because the error in the state + # norm is quadratic in the singular values (Frobenius norm). + # determine_error_thresh finds the largest cutoff such that + # the discarded squared norm stays within epsilon^2. + normalized_S = S / np.linalg.norm(S) + singular_value_threshold = determine_error_thresh( + np.flip(normalized_S), epsilon * epsilon, + ) + S[normalized_S <= singular_value_threshold] = 0.0 + # Truncation step 2: hard cap at bond_dim_max to prevent + # bond dimensions from growing beyond the MPS budget, + # independent of the accuracy-based threshold above + if S.shape[0] > bond_dim_max: + S[bond_dim_max:] = 0.0 + rank = np.count_nonzero(S) + U_trunc = U[:, :rank] + Vh_trunc = Vh[:rank, :] + # Reshape U back: (Dl, ..., rank) — left-orthogonal core + list_cores_compressed.append(U_trunc.reshape(*shape[:-1], rank)) + # Absorb singular values into Vh and contract with the next + # right-orthogonal core to form the new center core. This keeps + # all discarded weight local to the current bond. + M2_weighted_vh = S[:rank, None] * Vh_trunc + core_cur = np.tensordot( + M2_weighted_vh, list_cores_ortho[i + 1], 1, + ) + # Last core carries all remaining bond weights accumulated from the + # left sweep; it is appended as-is without further decomposition + list_cores_compressed.append(core_cur) + return list_cores_compressed + +def contract_down( + list_cores_phi: list[np.ndarray] | list[list[np.ndarray]], + method: str, + n_state: int, +) -> np.ndarray: + """ + Computes approximate per-state norm squared by independent core contraction. + + Parameters + ---------- + 1. list_cores_phi: list(np.ndarray) | list(list(np.ndarray)) + MPS cores. For fullstate representation, a flat + list [state_core, mode_core_0, ...]. For + number representation, nested: + list_cores_phi[s] = [state_core_s, mode_core_s0, ...] + 2. method: str + 'fullstate' or 'number'. + 3. n_state: int + Number of active system states. + + Returns + ------- + 1. V1_approx_norm_sq: np.ndarray(complex), shape (n_state,). + """ + + if method == 'fullstate': + # Approximate per-state norm by treating each core independently: + # replace each mode core A_m with sum_k |A_m[:,k,:]|^2, dropping + # cross-bond interference terms. This is cheaper than the exact + # double-layer contraction in contract_down_exact but overestimates + # the norm when bond correlations are significant. + # Sweep right-to-left, accumulating the squared transfer matrices. + for i in range(len(list_cores_phi) - 1, 0, -1): + outer_core = np.abs(list_cores_phi[i]) ** 2 + # Sum over physical index to get a bond-to-bond transfer matrix + outer_core_contracted = np.sum(outer_core, axis=1) + if i == len(list_cores_phi) - 1: + result = outer_core_contracted + else: + result = outer_core_contracted @ result + # Contract the state core (keeps physical index) with the + # accumulated mode transfer matrices to get per-state values + outer_core = np.abs(list_cores_phi[0]) ** 2 + result = outer_core @ result + V1_approx_norm_sq = np.sum(result, axis=-1)[0] + + elif method == 'number': + # list_cores_phi[s] = [state_core_s, mode_core_s0, ...] + # Approximation: replace each mode core A_m with sum(|A_m|^2, axis=phys), + # dropping cross-bond interference terms. + list_labeled = _flat_cores_with_labels(list_cores_phi) + # Precompute target_state-independent mode-core transfer matrices + # (|A_m|^2 summed over the physical axis). State cores stay 3-D + # to be sliced per target_state below. + list_mode_transfer = [ + np.sum(np.abs(core) ** 2, axis=1) if not is_state_core else None + for core, is_state_core, _ in list_labeled + ] + + V1_approx_norm_sq = np.zeros(n_state, dtype=np.complex128) + for target_state in range(n_state): + M2_env = None + for idx, (core, is_state_core, state_idx) in enumerate( + list_labeled, + ): + if is_state_core: + phys = 1 if state_idx == target_state else 0 + M2_slice = np.abs(core[:, phys, :]) ** 2 + M2_env = ( + M2_slice if M2_env is None + else np.tensordot(M2_env, M2_slice, 1) + ) + else: + M2_env = np.tensordot( + M2_env, list_mode_transfer[idx], 1, + ) + V1_approx_norm_sq[target_state] = M2_env[0][0] + else: + raise ValueError(f'Unknown method {method!r}.') + + return V1_approx_norm_sq + +def contract_down_exact( + list_cores_phi: list[np.ndarray] | list[list[np.ndarray]], + method: str, + n_state: int, +) -> np.ndarray: + """ + Computes exact per-state norm squared via double-layer contraction. + + Parameters + ---------- + 1. list_cores_phi: list(np.ndarray) | list(list(np.ndarray)) + MPS cores. For fullstate representation, a flat + list [state_core, mode_core_0, ...]. For + number representation, nested: + list_cores_phi[s] = [state_core_s, mode_core_s0, ...] + 2. method: str + 'fullstate' or 'number'. + 3. n_state: int + Number of active system states. + + Returns + ------- + 1. V1_norm_sq: np.ndarray(complex), shape (n_state,). + """ + + if method not in ('fullstate', 'number'): + raise ValueError( + f'Unknown method {method!r}. Expected ' + f"'fullstate' or 'number'." + ) + + if method == 'fullstate': + M2_boundary = np.sum(list_cores_phi[-1], axis=-1) + M2_boundary_conj = np.conj(M2_boundary) + M2_env = np.einsum( + 'ij, kj -> ik', M2_boundary, M2_boundary_conj, + ) + for i in range(len(list_cores_phi) - 2, 0, -1): + T3_core = list_cores_phi[i] + T3_core_conj = np.conj(T3_core) + M2_env = np.einsum( + 'ijk, il, mjl -> im', T3_core, M2_env, T3_core_conj, + optimize=True, + ) + M2_first_core = np.sum(list_cores_phi[0], axis=0) + M2_first_core_conj = np.conj(M2_first_core) + V1_norm_sq = np.diag(np.einsum( + 'ij,jl,kl -> ik', + M2_first_core, M2_env, M2_first_core_conj, + optimize=True, + )) + elif method == 'number': + # list_cores_phi[s] = [state_core_s, mode_core_s0, ...] + # One right-to-left double-layer sweep per target state. + # list_labeled is precomputed once; the inner loop reuses it for + # each target_state to avoid rebuilding per iteration. + if len(list_cores_phi) < n_state: + raise ValueError( + f'contract_down_exact number representation expects at least ' + f'{n_state} groups, got {len(list_cores_phi)}.' + ) + for s in range(len(list_cores_phi) - 1): + right_bond = list_cores_phi[s][-1].shape[2] + left_bond = list_cores_phi[s + 1][0].shape[0] + if right_bond != left_bond: + raise ValueError( + f'Environment/bond mismatch at group {s}: right bond ' + f'{right_bond} != left bond {left_bond} of group {s + 1}.' + ) + list_labeled = _flat_cores_with_labels(list_cores_phi) + V1_norm_sq = np.zeros(n_state, dtype=np.complex128) + + for target_state in range(n_state): + # R is the right environment matrix, shape (chi, chi). + # Initialized to scalar 1 for the right boundary (OBC). + R = np.ones((1, 1), dtype=np.complex128) + + for core, is_state_core, state_idx in reversed(list_labeled): + core_conj = np.conj(core) + if is_state_core: + # One-hot selection: project onto occupied (1) or + # unoccupied (0) physical index for this state. + phys = 1 if state_idx == target_state else 0 + M2_slice = core[:, phys, :] + R = np.einsum( + 'ai,ij,bj->ab', + M2_slice, R, core_conj[:, phys, :], + optimize=True, + ) + else: + # Mode core: sum over physical index in double layer. + R = np.einsum( + 'ami,ij,bmj->ab', core, R, core_conj, + optimize=True, + ) + + V1_norm_sq[target_state] = R[0, 0] + return V1_norm_sq + + +def tensor_to_array( + list_cores_phi: list[np.ndarray] | list[list[np.ndarray]], + method: str, + M1_modes_per_state: np.ndarray, + system: HopsSystem, + mode: HopsModes, +) -> np.ndarray: + """ + Returns the ground and first-order auxiliary wavefunctions as a flat array + in adHOPS form, for debugging purposes. + + Parameters + ---------- + 1. list_cores_phi: list(np.ndarray) | list(list(np.ndarray)) + MPS cores of the current wavefunction. For + fullstate, a flat list + [state_core, mode_core_0, ...]. For + number, nested: + list_cores_phi[s] = [state_core_s, mode_core_s0, ...] + + 2. method: str + Tensor encoding type ('fullstate' or + 'number'). + + 3. M1_modes_per_state: np.ndarray(int) + Number of bath modes per state. + + 4. system: HopsSystem + System object providing current state count. + + 5. mode: HopsModes + Mode object providing list_modeidx_abs. + + Returns + ------- + 1. V1_arrayform: np.array(complex) + Array of length n_state * (n_mode + 1) containing phi_0 + followed by each first-order auxiliary wavefunction. + """ + n_state = system.size + n_mode = len(mode.list_modeidx_abs) + V1_arrayform = np.zeros(n_state * (n_mode + 1), dtype=np.complex128) + V1_arrayform[0:n_state] = extract_psi(list_cores_phi, method, M1_modes_per_state) + for i_mode in range(n_mode): + list_occ_idx = [0] * n_mode + list_occ_idx[i_mode] = 1 + V1_arrayform[(i_mode + 1) * n_state:(i_mode + 2) * n_state] = ( + phi_aux(list_cores_phi, method, M1_modes_per_state, list_occ_idx) + ) + return V1_arrayform diff --git a/style_guide.md b/style_guide.md index 5d78eb1..cc1d073 100644 --- a/style_guide.md +++ b/style_guide.md @@ -192,6 +192,11 @@ Brief description of class. - Numbered list: `1. name: type1 | type2 [units: cm^-1]` - Note: we use brackets, not parens, for units + - When to annotate units: + * Every parameter, return, or attribute representing a dimensioned physical quantity carries `[units: ...]`. + * Explicitly dimensionless physical quantities (ratios, normalized amplitudes, relative errors) carry `[units: dimensionless]`. + * The bracket is omitted only for non-physical values (flags, indices, counts). + * For scalars and homogeneous containers the units bracket has a single entry; for heterogeneous structures the bracket matches the repeating element shown in the type — e.g., `list(tuple(complex, complex))` takes `[units: (cm^-2, cm^-1)]`. - Description on the next indented line, starting aligned with the first character after the colon. @@ -219,6 +224,15 @@ Parameters 2. rizzler: str [options: 'rizz_yes', 'rizz_no'] Brief description. +3. list_memory_terms: list(complex) [units: cm^-1] + Brief description (units describe each element). +4. list_gw: list(tuple(complex, complex)) [units: (cm^-2, cm^-1)] + Correlation-function (g, w) mode pairs. +5. rel_error: float [units: dimensionless] + Brief description. +6. n_modes: int + Mode count (non-physical: no units bracket). + Returns ------- 1. fojangle: float [units: fs^-1] diff --git a/tests/integrated_tests/test_tensor_LTC_eom.py b/tests/integrated_tests/test_tensor_LTC_eom.py new file mode 100644 index 0000000..d327a8e --- /dev/null +++ b/tests/integrated_tests/test_tensor_LTC_eom.py @@ -0,0 +1,366 @@ +import numpy as np +import pytest +import scipy as sp + +from mesohops.trajectory.exp_noise import bcf_exp +from mesohops.trajectory.hops_tensor_trajectory import HopsTensorTrajectory +from mesohops.trajectory.hops_trajectory import HopsTrajectory as HOPS + +__title__ = 'Test of tensor low-temperature correction' +__author__ = 'A. Hartzell' +__maintainer__ = 'A. Hartzell' + +# ============================================================ +# System Parameters (duplicated from test_LTC_eom.py) +# ============================================================ + +noise_param = { + 'SEED': 0, + 'MODEL': 'FFT_FILTER', + 'TLEN': 50.0, + 'TAU': 1.0, +} + +# --- 3-site system --- + +T3_loperator = np.zeros([3, 3, 3], dtype=np.float64) +T3_loperator[0, 0, 0] = 1.0 +T3_loperator[1, 1, 1] = 1.0 +T3_loperator[2, 2, 2] = 1.0 + +sys_param_ltc = { + 'HAMILTONIAN': np.array([[0, 1, 0], [1, 0, 1], [0, 1, 0]], + dtype=np.float64), + 'GW_SYSBATH': [[10.0, 10.0], [5.0, 5.0], [10.0, 10.0], [5.0, 5.0], + [10.0, 10.0], [5.0, 5.0]], + 'L_HIER': [T3_loperator[0], T3_loperator[0], T3_loperator[1], + T3_loperator[1], T3_loperator[2], T3_loperator[2]], + 'L_NOISE1': [T3_loperator[0], T3_loperator[0], T3_loperator[1], + T3_loperator[1], T3_loperator[2], T3_loperator[2], + T3_loperator[0], T3_loperator[1], T3_loperator[2]], + 'ALPHA_NOISE1': bcf_exp, + 'PARAM_NOISE1': [[10.0, 10.0], [5.0, 5.0], [10.0, 10.0], [5.0, 5.0], + [10.0, 10.0], [5.0, 5.0], [250.0, 1000.0], + [250.0, 2000.0], [250.0, 3000.0]], + 'PARAM_LT_CORR': [250.0 / 1000.0, 250.0 / 2000.0, 250.0 / 3000.0], + 'L_LT_CORR': [T3_loperator[0], T3_loperator[1], T3_loperator[2]], +} + +sys_param_no_ltc = { + 'HAMILTONIAN': np.array([[0, 1, 0], [1, 0, 1], [0, 1, 0]], + dtype=np.float64), + 'GW_SYSBATH': [[10.0, 10.0], [5.0, 5.0], [10.0, 10.0], [5.0, 5.0], + [10.0, 10.0], [5.0, 5.0]], + 'L_HIER': [T3_loperator[0], T3_loperator[0], T3_loperator[1], + T3_loperator[1], T3_loperator[2], T3_loperator[2]], + 'L_NOISE1': [T3_loperator[0], T3_loperator[0], T3_loperator[1], + T3_loperator[1], T3_loperator[2], T3_loperator[2], + T3_loperator[0], T3_loperator[1], T3_loperator[2]], + 'ALPHA_NOISE1': bcf_exp, + 'PARAM_NOISE1': [[10.0, 10.0], [5.0, 5.0], [10.0, 10.0], [5.0, 5.0], + [10.0, 10.0], [5.0, 5.0], [250.0, 1000.0], + [250.0, 2000.0], [250.0, 3000.0]], +} + +# --- 4-site system (complex LTC coefficients, multi-state L-operators) --- + +T3_loperator_4site = np.zeros([4, 4, 4], dtype=np.float64) +T3_loperator_4site[0, 0, 0] = 1.0 +T3_loperator_4site[0, 1, 1] = 1.0 +T3_loperator_4site[1, 0, 0] = 1.0 +T3_loperator_4site[1, 2, 2] = 1.0 +T3_loperator_4site[2, 2, 2] = 1.0 +T3_loperator_4site[2, 3, 3] = 1.0 +T3_loperator_4site[3, 1, 1] = 1.0 +T3_loperator_4site[3, 3, 3] = 1.0 + +sys_param_ltc_4site = { + 'HAMILTONIAN': np.array( + [[0, 1, 0, 0], [1, 0, 1, 0], [0, 1, 0, 1], [0, 0, 1, 0]], + dtype=np.float64), + 'GW_SYSBATH': [[10.0, 10.0], [5.0, 5.0], [10.0, 10.0], [5.0, 5.0], + [10.0, 10.0], [5.0, 5.0], [10.0, 10.0], [5.0, 5.0]], + 'L_HIER': [T3_loperator_4site[0], T3_loperator_4site[0], + T3_loperator_4site[1], T3_loperator_4site[1], + T3_loperator_4site[2], T3_loperator_4site[2], + T3_loperator_4site[3], T3_loperator_4site[3]], + 'L_NOISE1': [T3_loperator_4site[0], T3_loperator_4site[0], + T3_loperator_4site[1], T3_loperator_4site[1], + T3_loperator_4site[2], T3_loperator_4site[2], + T3_loperator_4site[3], T3_loperator_4site[3], + T3_loperator_4site[0], T3_loperator_4site[1], + T3_loperator_4site[2], T3_loperator_4site[3]], + 'ALPHA_NOISE1': bcf_exp, + 'PARAM_NOISE1': [[10.0, 10.0], [5.0, 5.0], [10.0, 10.0], [5.0, 5.0], + [10.0, 10.0], [5.0, 5.0], [10.0, 10.0], [5.0, 5.0], + [(250.0 + 10j), 1000.0], [(250.0 + 10j), 2000.0], + [(250.0 + 10j), 3000.0], [(250.0 + 10j), 4000.0]], + 'PARAM_LT_CORR': [(250.0 + 10j) / 1000.0, (250.0 + 10j) / 2000.0, + (250.0 + 10j) / 3000.0, (250.0 + 10j) / 4000.0], + 'L_LT_CORR': [T3_loperator_4site[0], T3_loperator_4site[1], + T3_loperator_4site[2], T3_loperator_4site[3]], +} + +sys_param_no_ltc_4site = { + 'HAMILTONIAN': np.array( + [[0, 1, 0, 0], [1, 0, 1, 0], [0, 1, 0, 1], [0, 0, 1, 0]], + dtype=np.float64), + 'GW_SYSBATH': [[10.0, 10.0], [5.0, 5.0], [10.0, 10.0], [5.0, 5.0], + [10.0, 10.0], [5.0, 5.0], [10.0, 10.0], [5.0, 5.0]], + 'L_HIER': [T3_loperator_4site[0], T3_loperator_4site[0], + T3_loperator_4site[1], T3_loperator_4site[1], + T3_loperator_4site[2], T3_loperator_4site[2], + T3_loperator_4site[3], T3_loperator_4site[3]], + 'L_NOISE1': [T3_loperator_4site[0], T3_loperator_4site[0], + T3_loperator_4site[1], T3_loperator_4site[1], + T3_loperator_4site[2], T3_loperator_4site[2], + T3_loperator_4site[3], T3_loperator_4site[3], + T3_loperator_4site[0], T3_loperator_4site[1], + T3_loperator_4site[2], T3_loperator_4site[3]], + 'ALPHA_NOISE1': bcf_exp, + 'PARAM_NOISE1': [[10.0, 10.0], [5.0, 5.0], [10.0, 10.0], [5.0, 5.0], + [10.0, 10.0], [5.0, 5.0], [10.0, 10.0], [5.0, 5.0], + [(250.0 + 10j), 1000.0], [(250.0 + 10j), 2000.0], + [(250.0 + 10j), 3000.0], [(250.0 + 10j), 4000.0]], +} + +# --- Shared parameters --- + +hier_param = {'MAXHIER': 1} +integrator_param = { + 'INTEGRATOR': 'RUNGE_KUTTA', + 'EARLY_ADAPTIVE_INTEGRATOR': 'INCH_WORM', + 'EARLY_INTEGRATOR_STEPS': 5, + 'INCHWORM_CAP': 5, + 'STATIC_BASIS': None, +} +tensor_param = { + 'MPS_EPSILON': 1e-12, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, +} + +t_max = 20.0 +t_step = 2.0 +H1_psi_0_3site = np.array([1.0, 0.0, 0.0], dtype=np.complex128) +H1_psi_0_4site = np.array([0.7071, 0.5773, 0.4082, 0.1872], + dtype=np.complex128) +H1_psi_0_4site = H1_psi_0_4site / np.linalg.norm(H1_psi_0_4site) + + +# ============================================================ +# Helpers +# ============================================================ + +def _run_tensor(sys_param, eom_param, psi_0, adaptive=False): + """Run tensor HOPS and return psi trajectory as array.""" + traj = HopsTensorTrajectory( + system_param=sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integrator_param, + tensor_param=tensor_param, + ) + if adaptive: + traj.make_adaptive(1e-5, 1e-5) + traj.initialize(psi_0) + traj.propagate(t_max, t_step) + return np.array(traj.storage.data['psi_traj']) + + +def _run_vector(sys_param, eom_param, psi_0, adaptive=False): + """Run vector HOPS and return psi trajectory as array.""" + hops = HOPS( + sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integrator_param, + ) + if adaptive: + hops.make_adaptive(1e-5, 1e-5) + hops.initialize(psi_0) + hops.propagate(t_max, t_step) + return np.array(hops.storage.data['psi_traj']) + + +# ============================================================ +# TEST SUITE: LTC has effect on tensor dynamics (Group A) +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: LTC modifies tensor dynamics — LINEAR +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_LTC_tensor_has_effect_linear(): + # This case tests that the low-temperature correction produces a + # measurable change in the tensor psi trajectory under the LINEAR + # equation of motion. Propagates with and without LTC on the same + # 3-site system and asserts the trajectories diverge. + eom_param = { + 'TIME_DEPENDENCE': False, + 'EQUATION_OF_MOTION': 'LINEAR', + } + psi_ltc = _run_tensor(sys_param_ltc, eom_param, H1_psi_0_3site) + psi_no_ltc = _run_tensor(sys_param_no_ltc, eom_param, H1_psi_0_3site) + + assert not np.allclose(psi_ltc, psi_no_ltc, atol=1e-6), ( + 'LTC had no measurable effect on tensor LINEAR dynamics' + ) + + +# ------------------------------------------------------------ +# TEST: LTC modifies tensor dynamics — NONLINEAR +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_LTC_tensor_has_effect_nonlinear(): + # This case tests that the low-temperature correction produces a + # measurable change in the tensor psi trajectory under the NONLINEAR + # equation of motion. + eom_param = {'EQUATION_OF_MOTION': 'NONLINEAR'} + psi_ltc = _run_tensor(sys_param_ltc, eom_param, H1_psi_0_3site) + psi_no_ltc = _run_tensor(sys_param_no_ltc, eom_param, H1_psi_0_3site) + + assert not np.allclose(psi_ltc, psi_no_ltc, atol=1e-6), ( + 'LTC had no measurable effect on tensor NONLINEAR dynamics' + ) + + +# ------------------------------------------------------------ +# TEST: LTC modifies tensor dynamics — NORMALIZED NONLINEAR +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_LTC_tensor_has_effect_normalized_nonlinear(): + # This case tests that the low-temperature correction produces a + # measurable change in the tensor psi trajectory under the + # NORMALIZED NONLINEAR equation of motion. + eom_param = {'EQUATION_OF_MOTION': 'NORMALIZED NONLINEAR'} + psi_ltc = _run_tensor(sys_param_ltc, eom_param, H1_psi_0_3site) + psi_no_ltc = _run_tensor(sys_param_no_ltc, eom_param, H1_psi_0_3site) + + assert not np.allclose(psi_ltc, psi_no_ltc, atol=1e-6), ( + 'LTC had no measurable effect on tensor NORMALIZED NONLINEAR dynamics' + ) + + +# ------------------------------------------------------------ +# TEST: LTC modifies tensor dynamics — adaptive, 3-site +# ------------------------------------------------------------ +@pytest.mark.level(2) +@pytest.mark.xfail(reason='tensor adaptive basis not yet working') +def test_LTC_tensor_has_effect_adaptive(): + # This case tests that the low-temperature correction produces a + # measurable change in the tensor psi trajectory with adaptive + # basis on the 3-site system. + eom_param = {'EQUATION_OF_MOTION': 'NORMALIZED NONLINEAR'} + psi_ltc = _run_tensor(sys_param_ltc, eom_param, H1_psi_0_3site, + adaptive=True) + psi_no_ltc = _run_tensor(sys_param_no_ltc, eom_param, H1_psi_0_3site, + adaptive=True) + + assert not np.allclose(psi_ltc, psi_no_ltc, atol=1e-6), ( + 'LTC had no measurable effect on adaptive tensor dynamics' + ) + + +# ------------------------------------------------------------ +# TEST: LTC modifies tensor dynamics — adaptive, 4-site multiparticle +# ------------------------------------------------------------ +@pytest.mark.level(2) +@pytest.mark.xfail(reason='tensor adaptive basis not yet working') +def test_LTC_tensor_has_effect_adaptive_multiparticle(): + # This case tests that the low-temperature correction produces a + # measurable change in the tensor psi trajectory with adaptive + # basis on the 4-site system with complex LTC coefficients and + # multi-state L-operators. + eom_param = {'EQUATION_OF_MOTION': 'NORMALIZED NONLINEAR'} + psi_ltc = _run_tensor(sys_param_ltc_4site, eom_param, H1_psi_0_4site, + adaptive=True) + psi_no_ltc = _run_tensor(sys_param_no_ltc_4site, eom_param, + H1_psi_0_4site, adaptive=True) + + assert not np.allclose(psi_ltc, psi_no_ltc, atol=1e-6), ( + 'LTC had no measurable effect on adaptive multiparticle tensor dynamics' + ) + + +# ============================================================ +# TEST SUITE: Tensor LTC matches vector LTC (Group B) +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Tensor LTC matches vector — LINEAR +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_LTC_tensor_vs_vector_linear(): + # This case tests that tensor HOPS with LTC produces the same psi + # trajectory as vector HOPS with LTC under the LINEAR equation of + # motion. Uses the 3-site system with real LTC coefficients. + eom_param = { + 'TIME_DEPENDENCE': False, + 'EQUATION_OF_MOTION': 'LINEAR', + } + psi_tensor = _run_tensor(sys_param_ltc, eom_param, H1_psi_0_3site) + psi_vector = _run_vector(sys_param_ltc, eom_param, H1_psi_0_3site) + + n_steps = min(len(psi_tensor), len(psi_vector)) + max_err = max( + np.linalg.norm(psi_tensor[i] - psi_vector[i]) + for i in range(n_steps) + ) + # Relaxed to 2e-9: this 3-site system with 9 noise modes has a + # baseline tensor-vs-vector discrepancy of ~1.7e-9 even without LTC, + # caused by accumulated SVD truncation (epsilon=1e-12) over 10 steps. + assert max_err < 2e-9, ( + f'Tensor LTC diverged from vector LTC (LINEAR): ' + f'max wf error = {max_err:.2e}' + ) + + +# ------------------------------------------------------------ +# TEST: Tensor LTC matches vector — NONLINEAR +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_LTC_tensor_vs_vector_nonlinear(): + # This case tests that tensor HOPS with LTC produces the same psi + # trajectory as vector HOPS with LTC under the NONLINEAR equation + # of motion. + eom_param = {'EQUATION_OF_MOTION': 'NONLINEAR'} + psi_tensor = _run_tensor(sys_param_ltc, eom_param, H1_psi_0_3site) + psi_vector = _run_vector(sys_param_ltc, eom_param, H1_psi_0_3site) + + n_steps = min(len(psi_tensor), len(psi_vector)) + max_err = max( + np.linalg.norm(psi_tensor[i] - psi_vector[i]) + for i in range(n_steps) + ) + assert max_err < 1e-9, ( + f'Tensor LTC diverged from vector LTC (NONLINEAR): ' + f'max wf error = {max_err:.2e}' + ) + + +# ------------------------------------------------------------ +# TEST: Tensor LTC matches vector — NORMALIZED NONLINEAR +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_LTC_tensor_vs_vector_normalized_nonlinear(): + # This case tests that tensor HOPS with LTC produces the same psi + # trajectory as vector HOPS with LTC under the NORMALIZED NONLINEAR + # equation of motion. + eom_param = {'EQUATION_OF_MOTION': 'NORMALIZED NONLINEAR'} + psi_tensor = _run_tensor(sys_param_ltc, eom_param, H1_psi_0_3site) + psi_vector = _run_vector(sys_param_ltc, eom_param, H1_psi_0_3site) + + n_steps = min(len(psi_tensor), len(psi_vector)) + max_err = max( + np.linalg.norm(psi_tensor[i] - psi_vector[i]) + for i in range(n_steps) + ) + assert max_err < 1e-9, ( + f'Tensor LTC diverged from vector LTC (NORMALIZED NONLINEAR): ' + f'max wf error = {max_err:.2e}' + ) diff --git a/tests/test_dimer_of_dimers_tensor.py b/tests/test_dimer_of_dimers_tensor.py new file mode 100644 index 0000000..0e12734 --- /dev/null +++ b/tests/test_dimer_of_dimers_tensor.py @@ -0,0 +1,729 @@ +import numpy as np +import pytest +import scipy as sp + +from mesohops.trajectory.exp_noise import bcf_exp +from mesohops.trajectory.hops_tensor_trajectory import HopsTensorTrajectory +from mesohops.trajectory.hops_trajectory import HopsTrajectory as HOPS +from mesohops.util.bath_corr_functions import bcf_convert_dl_to_exp + +__title__ = 'Dimer of Dimers: Tensor HOPS vs Vector HOPS' +__author__ = 'N. Covalsen' +__maintainer__ = 'N. Covalsen' + +# ============================================================ +# Dimer of Dimers System Setup +# ============================================================ +# Identical to test_dimer_of_dimers.py: 4 sites, 2 modes per site +# (Drude-Lorentz + LTC correction at 500 cm^-1) +noise_param = { + 'SEED': 0, + 'MODEL': 'FFT_FILTER', + 'TLEN': 25000.0, # Units: fs + 'TAU': 1.0, # Units: fs + 'STORE_RAW_NOISE': True, + 'RAND_MODEL': 'BOX_MULLER', +} + +nsite = 4 +e_lambda = 20.0 +gamma = 50.0 +temp = 140.0 +(g_0, w_0) = bcf_convert_dl_to_exp(e_lambda, gamma, temp) + +loperator = np.zeros([4, 4, 4], dtype=np.float64) +gw_sysbath = [] +lop_list = [] +for i in range(nsite): + loperator[i, i, i] = 1.0 + gw_sysbath.append([g_0, w_0]) + lop_list.append(sp.sparse.coo_matrix(loperator[i])) + gw_sysbath.append([-1j * np.imag(g_0), 500.0]) + lop_list.append(loperator[i]) + +hs = np.zeros([nsite, nsite]) +hs[0, 1] = 40 +hs[1, 0] = 40 +hs[1, 2] = 10 +hs[2, 1] = 10 +hs[2, 3] = 40 +hs[3, 2] = 40 + +sys_param = { + 'HAMILTONIAN': np.array(hs, dtype=np.complex128), + 'GW_SYSBATH': gw_sysbath, + 'L_HIER': lop_list, + 'L_NOISE1': lop_list, + 'ALPHA_NOISE1': bcf_exp, + 'PARAM_NOISE1': gw_sysbath, +} + +eom_param = {'EQUATION_OF_MOTION': 'NORMALIZED NONLINEAR'} +integrator_param = {'INTEGRATOR': 'RUNGE_KUTTA'} +hier_param = {'MAXHIER': 2, 'TRUNCATION_METHOD': 'rectangular'} + +psi_0 = np.array([0.0] * nsite, dtype=np.complex128) +psi_0[2] = 1.0 +psi_0 = psi_0 / np.linalg.norm(psi_0) + +t_max = 200.0 +t_step = 4.0 + + +# ============================================================ +# Helpers +# ============================================================ + +def _run_vector_hops(): + """Runs standard vector HOPS and returns psi_traj as array.""" + hops = HOPS( + sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integrator_param, + ) + hops.initialize(psi_0) + hops.propagate(t_max, t_step) + return np.array(hops.storage['psi_traj']) + + +def _run_tensor_hops(method, bond_dim_max=20, mps_epsilon=1e-10, + flag_mpo_optimize=True): + """Runs tensor HOPS and returns psi_traj as array.""" + tensor_param = { + 'MPS_EPSILON': mps_epsilon, + 'METHOD': method, + 'BOND_DIM_MAX': bond_dim_max, + 'FLAG_MPO_OPTIMIZE': flag_mpo_optimize, + } + traj = HopsTensorTrajectory( + system_param=sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integrator_param, + tensor_param=tensor_param, + ) + traj.initialize(psi_0) + traj.propagate(t_max, t_step) + return np.array(traj.storage['psi_traj']) + + +# ============================================================ +# Shared Fixtures +# ============================================================ +# Each propagation runs once per test session and is reused across tests. + +@pytest.fixture(scope='module') +def psi_vector(): + """Vector HOPS trajectory (computed once per module).""" + return _run_vector_hops() + + +@pytest.fixture(scope='module') +def psi_fullstate(): + """Fullstate tensor HOPS trajectory (computed once per module).""" + return _run_tensor_hops(method='fullstate') + + +@pytest.fixture( + scope='module', params=[True, False], ids=['optimized', 'general'], +) +def psi_statenumber_nn(request): + """Statenumber NN tensor HOPS trajectory, once per FLAG_MPO_OPTIMIZE + setting: the chain combined generator and the general two-MPO path. + """ + return _run_tensor_hops( + method='number', flag_mpo_optimize=request.param, + ) + + +# ============================================================ +# TEST SUITE: propagate() — vector HOPS vs tensor HOPS agreement +# ============================================================ +# These tests verify that the tensor HOPS EOM produces the same +# physical wavefunction trajectory as vector HOPS on the dimer-of-dimers +# system, with tight SVD convergence (epsilon=1e-10, bond_dim_max=20). +# At these parameters the hierarchy is small enough that SVD compression +# is essentially lossless, giving tensor-vs-vector agreement to ~1e-10. +# Tolerance set to 1e-9 to catch coupling-prefactor or MPO-wiring bugs +# that would accumulate over 50 RK4 steps. + +# ------------------------------------------------------------ +# TEST: Fullstate representation matches vector HOPS +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_fullstate_matches_vector_hops(psi_vector, psi_fullstate): + # This case tests that the fullstate tensor HOPS + # produces the same physical wavefunction trajectory as standard + # vector HOPS on the dimer-of-dimers system. + assert len(psi_vector) == len(psi_fullstate), "Trajectory lengths differ" + n_steps = len(psi_vector) + + # SVD compression errors accumulate over time steps, so we check the + # maximum wavefunction error across the full trajectory. + max_err = max( + np.linalg.norm(psi_fullstate[i] - psi_vector[i]) + for i in range(n_steps) + ) + assert max_err < 1e-9, ( + f'Fullstate tensor HOPS diverged from vector HOPS: ' + f'max wf error = {max_err:.2e}' + ) + + +# ------------------------------------------------------------ +# TEST: Statenumber NN representation matches vector HOPS +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_statenumber_nn_matches_vector_hops(psi_vector, psi_statenumber_nn): + # This case tests that the number with + # nearest-neighbor Hamiltonian MPO produces the same dynamics + # as standard vector HOPS. + assert len(psi_vector) == len(psi_statenumber_nn), "Trajectory lengths differ" + n_steps = len(psi_vector) + + max_err = max( + np.linalg.norm(psi_statenumber_nn[i] - psi_vector[i]) + for i in range(n_steps) + ) + assert max_err < 1e-9, ( + f'Statenumber NN tensor HOPS diverged from vector HOPS: ' + f'max wf error = {max_err:.2e}' + ) + + +# ============================================================ +# TEST SUITE: propagate() — non-nearest-neighbor Hamiltonian +# ============================================================ +# Long-range couplings H[0,2]=5, H[0,3]=3 break the NN MPO path and +# exercise the general (statenumber) MPO builder end-to-end. + +H2_ham_nonnn = np.zeros([nsite, nsite]) +H2_ham_nonnn[0, 1] = 40 +H2_ham_nonnn[1, 0] = 40 +H2_ham_nonnn[1, 2] = 10 +H2_ham_nonnn[2, 1] = 10 +H2_ham_nonnn[2, 3] = 40 +H2_ham_nonnn[3, 2] = 40 +H2_ham_nonnn[0, 2] = 5 +H2_ham_nonnn[2, 0] = 5 +H2_ham_nonnn[0, 3] = 3 +H2_ham_nonnn[3, 0] = 3 + +_sys_param_nonnn = dict(sys_param) +_sys_param_nonnn['HAMILTONIAN'] = np.array(H2_ham_nonnn, dtype=np.complex128) + + +def _run_vector_hops_nonnn(): + """Runs standard vector HOPS with non-NN Hamiltonian.""" + hops = HOPS( + _sys_param_nonnn, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integrator_param, + ) + hops.initialize(psi_0) + hops.propagate(t_max, t_step) + return np.array(hops.storage['psi_traj']) + + +def _run_tensor_hops_nonnn(method, bond_dim_max=30, mps_epsilon=1e-12): + """Runs tensor HOPS with non-NN Hamiltonian. + + The MPS budget is tighter than BOND_DIM_MAX=20 / MPS_EPSILON=1e-10 because + at that budget the gate measures MPS truncation rather than the MPO: the + number method sits 1.4e-9 from vector HOPS and the fullstate 2.3e-10, + against a 1e-9 threshold, and which side of it they fall on depends on the + MPO's bond profile rather than on the operator it encodes. At these + settings every comparison in the non-nearest-neighbor group is 1e-11 or + below. + """ + tensor_param = { + 'MPS_EPSILON': mps_epsilon, + 'METHOD': method, + 'BOND_DIM_MAX': bond_dim_max, + } + traj = HopsTensorTrajectory( + system_param=_sys_param_nonnn, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integrator_param, + tensor_param=tensor_param, + ) + traj.initialize(psi_0) + traj.propagate(t_max, t_step) + return np.array(traj.storage['psi_traj']) + + +@pytest.fixture(scope='module') +def psi_vector_nonnn(): + """Vector HOPS trajectory with non-NN Hamiltonian (computed once per module).""" + return _run_vector_hops_nonnn() + + +@pytest.fixture(scope='module') +def psi_fullstate_nonnn(): + """Fullstate tensor HOPS trajectory with non-NN Hamiltonian.""" + return _run_tensor_hops_nonnn(method='fullstate') + + +@pytest.fixture(scope='module') +def psi_statenumber_general_nonnn(): + """Statenumber general tensor HOPS trajectory with non-NN Hamiltonian.""" + return _run_tensor_hops_nonnn(method='number') + + +# ------------------------------------------------------------ +# TEST: Non-NN fullstate matches vector HOPS +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_fullstate_nonnn_matches_vector_hops( + psi_vector_nonnn, psi_fullstate_nonnn, +): + # This case tests that fullstate tensor HOPS handles long-range + # couplings correctly (fullstate embeds the Hamiltonian into the + # first MPO core, so non-NN structure is handled implicitly). + assert len(psi_vector_nonnn) == len(psi_fullstate_nonnn), "Trajectory lengths differ" + n_steps = len(psi_vector_nonnn) + + max_err = max( + np.linalg.norm(psi_fullstate_nonnn[i] - psi_vector_nonnn[i]) + for i in range(n_steps) + ) + assert max_err < 1e-9, ( + f'Fullstate tensor HOPS (non-NN Ham) diverged from vector HOPS: ' + f'max wf error = {max_err:.2e}' + ) + + +# ------------------------------------------------------------ +# TEST: Non-NN statenumber general representation matches vector HOPS +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_statenumber_general_nonnn_matches_vector_hops( + psi_vector_nonnn, psi_statenumber_general_nonnn, +): + # This case tests that the statenumber general MPO handles a non-nearest- + # neighbor Hamiltonian correctly. Long-range couplings (H[0,2]=5, H[0,3]=3) + # require the general MPO builder path; this verifies the result still + # matches standard vector HOPS to confirm correctness. + assert len(psi_vector_nonnn) == len(psi_statenumber_general_nonnn), "Trajectory lengths differ" + n_steps = len(psi_vector_nonnn) + + max_err = max( + np.linalg.norm(psi_statenumber_general_nonnn[i] - psi_vector_nonnn[i]) + for i in range(n_steps) + ) + assert max_err < 1e-9, ( + f'Statenumber general tensor HOPS (non-NN Ham) diverged from vector HOPS: ' + f'max wf error = {max_err:.2e}' + ) + + +# ------------------------------------------------------------ +# TEST: Non-NN fullstate vs statenumber agree +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_fullstate_vs_statenumber_nonnn( + psi_fullstate_nonnn, psi_statenumber_general_nonnn, +): + # This case tests that both representations agree for the non-NN + # Hamiltonian. The fullstate MPO embeds the full Hamiltonian in + # one core while the statenumber general MPO uses daisy-chained + # transfer matrices — both must produce the same result. + assert len(psi_fullstate_nonnn) == len(psi_statenumber_general_nonnn), "Trajectory lengths differ" + n_steps = len(psi_fullstate_nonnn) + + max_err = max( + np.linalg.norm(psi_fullstate_nonnn[i] - psi_statenumber_general_nonnn[i]) + for i in range(n_steps) + ) + assert max_err < 1e-9, ( + f'Fullstate vs statenumber disagree for non-NN Ham: ' + f'max wf error = {max_err:.2e}' + ) + + +# ============================================================ +# TEST SUITE: propagate() — linear EOM +# ============================================================ +# For linear HOPS the wavefunction is not renormalized, so its norm decays +# over time. The tensor and vector trajectories must agree element-wise on +# the raw (unnormalized) wavefunction at each stored timestep. + +eom_param_linear = {'EQUATION_OF_MOTION': 'LINEAR'} + + +def _run_vector_linear(): + """Runs vector HOPS with linear EOM and returns psi_traj.""" + hops = HOPS( + sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param_linear, + integration_param=integrator_param, + ) + hops.initialize(psi_0) + hops.propagate(t_max, t_step) + return np.array(hops.storage['psi_traj']) + + +def _run_tensor_linear(method, bond_dim_max=20, mps_epsilon=1e-10): + """Runs tensor HOPS with linear EOM and returns psi_traj.""" + tensor_param = { + 'MPS_EPSILON': mps_epsilon, + 'METHOD': method, + 'BOND_DIM_MAX': bond_dim_max, + } + traj = HopsTensorTrajectory( + system_param=sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param_linear, + integration_param=integrator_param, + tensor_param=tensor_param, + ) + traj.initialize(psi_0) + traj.propagate(t_max, t_step) + return np.array(traj.storage['psi_traj']) + + +@pytest.fixture(scope='module') +def psi_vector_linear(): + """Vector linear HOPS trajectory (computed once per module).""" + return _run_vector_linear() + + +@pytest.fixture(scope='module') +def psi_fullstate_linear(): + """Fullstate tensor linear HOPS trajectory (computed once per module).""" + return _run_tensor_linear(method='fullstate') + + +@pytest.fixture(scope='module') +def psi_statenumber_linear(): + """Statenumber tensor linear HOPS trajectory (computed once per module).""" + return _run_tensor_linear(method='number') + + +# ------------------------------------------------------------ +# TEST: Fullstate linear tensor HOPS matches vector linear HOPS +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_linear_fullstate_matches_vector(psi_vector_linear, psi_fullstate_linear): + # This case tests that fullstate tensor HOPS with the + # linear EOM produces the same unnormalized wavefunction trajectory + # as vector HOPS. The norm is allowed to decay; what matters is + # element-wise agreement of the raw psi_traj. The tighter tolerance + # (vs 1e-6 for nonlinear) reflects the absence of nonlinear norm + # correction errors. + assert len(psi_vector_linear) == len(psi_fullstate_linear), "Trajectory lengths differ" + n_steps = len(psi_vector_linear) + + max_err = max( + np.linalg.norm(psi_fullstate_linear[i] - psi_vector_linear[i]) + for i in range(n_steps) + ) + assert max_err < 1e-10, ( + f'Fullstate linear tensor HOPS diverged from vector linear HOPS: ' + f'max wf error = {max_err:.2e}' + ) + + # This case tests that the unnormalized wavefunction norm decays, + # confirming the linear EOM is not accidentally renormalizing. + norm_first = np.linalg.norm(psi_fullstate_linear[0]) + norm_last = np.linalg.norm(psi_fullstate_linear[-1]) + assert norm_last < norm_first, ( + f'Linear EOM norm should decay: first={norm_first:.6f}, ' + f'last={norm_last:.6f}' + ) + + +# ------------------------------------------------------------ +# TEST: Statenumber linear tensor HOPS matches vector linear HOPS +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_linear_statenumber_matches_vector(psi_vector_linear, psi_statenumber_linear): + # This case tests that number tensor HOPS with the + # linear EOM produces the same unnormalized wavefunction trajectory + # as vector HOPS. + assert len(psi_vector_linear) == len(psi_statenumber_linear), "Trajectory lengths differ" + n_steps = len(psi_vector_linear) + + max_err = max( + np.linalg.norm(psi_statenumber_linear[i] - psi_vector_linear[i]) + for i in range(n_steps) + ) + assert max_err < 1e-10, ( + f'Statenumber linear tensor HOPS diverged from vector linear HOPS: ' + f'max wf error = {max_err:.2e}' + ) + + +# ============================================================ +# TEST SUITE: propagate() — NONLINEAR EOM (no norm correction) +# ============================================================ +# NONLINEAR has feedback but no norm correction. This is a distinct +# code path from both LINEAR (no feedback, no norm) and NORMALIZED +# NONLINEAR (feedback + norm). The norm decays because there is no +# renormalization. + +eom_param_nonlinear = {'EQUATION_OF_MOTION': 'NONLINEAR'} + + +def _run_vector_nonlinear(): + """Runs vector HOPS with NONLINEAR EOM.""" + hops = HOPS( + sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param_nonlinear, + integration_param=integrator_param, + ) + hops.initialize(psi_0) + hops.propagate(t_max, t_step) + return np.array(hops.storage['psi_traj']) + + +def _run_tensor_nonlinear(method, bond_dim_max=20, mps_epsilon=1e-10): + """Runs tensor HOPS with NONLINEAR EOM.""" + tensor_param = { + 'MPS_EPSILON': mps_epsilon, + 'METHOD': method, + 'BOND_DIM_MAX': bond_dim_max, + } + traj = HopsTensorTrajectory( + system_param=sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param_nonlinear, + integration_param=integrator_param, + tensor_param=tensor_param, + ) + traj.initialize(psi_0) + traj.propagate(t_max, t_step) + return np.array(traj.storage['psi_traj']) + + +@pytest.fixture(scope='module') +def psi_vector_nonlinear(): + """Vector NONLINEAR HOPS trajectory (computed once per module).""" + return _run_vector_nonlinear() + + +@pytest.fixture(scope='module') +def psi_fullstate_nonlinear(): + """Fullstate tensor NONLINEAR HOPS trajectory (computed once per module).""" + return _run_tensor_nonlinear(method='fullstate') + + +@pytest.fixture(scope='module') +def psi_statenumber_nonlinear(): + """Statenumber tensor NONLINEAR HOPS trajectory (computed once per module).""" + return _run_tensor_nonlinear(method='number') + + +# ------------------------------------------------------------ +# TEST: Fullstate NONLINEAR tensor matches vector HOPS +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_nonlinear_fullstate_matches_vector( + psi_vector_nonlinear, psi_fullstate_nonlinear, +): + # This case tests that fullstate tensor HOPS with the NONLINEAR EOM + # ( feedback, no norm correction) matches vector HOPS. + assert len(psi_vector_nonlinear) == len(psi_fullstate_nonlinear), "Trajectory lengths differ" + n_steps = len(psi_vector_nonlinear) + + max_err = max( + np.linalg.norm(psi_fullstate_nonlinear[i] - psi_vector_nonlinear[i]) + for i in range(n_steps) + ) + assert max_err < 1e-9, ( + f'Fullstate NONLINEAR tensor HOPS diverged from vector HOPS: ' + f'max wf error = {max_err:.2e}' + ) + + +# ------------------------------------------------------------ +# TEST: Statenumber NONLINEAR tensor matches vector HOPS +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_nonlinear_statenumber_matches_vector( + psi_vector_nonlinear, psi_statenumber_nonlinear, +): + # This case tests that statenumber tensor HOPS with the NONLINEAR EOM + # matches vector HOPS. + assert len(psi_vector_nonlinear) == len(psi_statenumber_nonlinear), "Trajectory lengths differ" + n_steps = len(psi_vector_nonlinear) + + max_err = max( + np.linalg.norm(psi_statenumber_nonlinear[i] - psi_vector_nonlinear[i]) + for i in range(n_steps) + ) + assert max_err < 1e-9, ( + f'Statenumber NONLINEAR tensor HOPS diverged from vector HOPS: ' + f'max wf error = {max_err:.2e}' + ) + + +# ============================================================ +# TEST SUITE: propagate() — fullstate vs statenumber agreement +# ============================================================ +# Cross-comparison between representations catches bugs that affect +# both equally (where both diverge from vector HOPS by the same amount). + +# ------------------------------------------------------------ +# TEST: Fullstate and statenumber NN agree for NORMALIZED NONLINEAR +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_fullstate_vs_statenumber_normalized_nonlinear( + psi_fullstate, psi_statenumber_nn, +): + # This case tests that both tensor representations produce identical + # trajectories, independent of the vector HOPS reference. + assert len(psi_fullstate) == len(psi_statenumber_nn), "Trajectory lengths differ" + n_steps = len(psi_fullstate) + + max_err = max( + np.linalg.norm(psi_fullstate[i] - psi_statenumber_nn[i]) + for i in range(n_steps) + ) + assert max_err < 1e-9, ( + f'Fullstate vs statenumber disagree for NORMALIZED NONLINEAR: ' + f'max wf error = {max_err:.2e}' + ) + + +# ------------------------------------------------------------ +# TEST: Fullstate and statenumber agree for NONLINEAR +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_fullstate_vs_statenumber_nonlinear( + psi_fullstate_nonlinear, psi_statenumber_nonlinear, +): + # This case tests that both representations agree for the NONLINEAR + # EOM ( feedback, no norm correction). + assert len(psi_fullstate_nonlinear) == len(psi_statenumber_nonlinear), "Trajectory lengths differ" + n_steps = len(psi_fullstate_nonlinear) + + max_err = max( + np.linalg.norm(psi_fullstate_nonlinear[i] - psi_statenumber_nonlinear[i]) + for i in range(n_steps) + ) + assert max_err < 1e-9, ( + f'Fullstate vs statenumber disagree for NONLINEAR: ' + f'max wf error = {max_err:.2e}' + ) + + +# ============================================================ +# TEST SUITE: propagate() — ring and star Hamiltonians +# ============================================================ +# Each combined-generator MPO builder is selected by the coupling graph of +# the Hamiltonian, so ring and star topologies exercise code paths that the +# chain and general Hamiltonians above never reach. + +H2_ham_ring = np.zeros([nsite, nsite]) +for _i in range(nsite - 1): + H2_ham_ring[_i, _i + 1] = 40 + H2_ham_ring[_i + 1, _i] = 40 +# The bond that closes the ring carries a different amplitude from the chain +# bonds. The ring builder opens that bond at site 0, relays it the length of +# the chain and closes it at the last site against H[0, n-1] and H[n-1, 0]; +# giving it the neighbour amplitude would let a builder that picked up the +# wrong entry, or swapped the two directions, still pass. +H2_ham_ring[0, nsite - 1] = 25 +H2_ham_ring[nsite - 1, 0] = 25 + +# Hub at site 1 rather than site 0, so the builder's leaves-left and +# leaves-right branches are both used. +H2_ham_star = np.zeros([nsite, nsite]) +for _i in range(nsite): + if _i != 1: + H2_ham_star[1, _i] = 30 + H2_ham_star[_i, 1] = 30 + + +def _run_hops_topology(H2_ham_topology, method=None, bond_dim_max=40, + mps_epsilon=1e-13, flag_mpo_optimize=True): + """Runs vector HOPS (method=None) or tensor HOPS for a topology. + + The MPS budget is tighter than elsewhere in this module because the ring + closes an extra bond and so entangles more than the chain: at + BOND_DIM_MAX=20 / MPS_EPSILON=1e-10 the ring trajectory sits ~6e-8 from + vector HOPS through MPS truncation alone, for the general MPO path as + well as for the combined one. These settings put the wavefunction error + below the level being tested so the comparison probes the MPO. + """ + sys_param_topology = dict(sys_param) + sys_param_topology['HAMILTONIAN'] = np.array( + H2_ham_topology, dtype=np.complex128, + ) + if method is None: + traj_vector = HOPS( + sys_param_topology, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integrator_param, + ) + traj_vector.initialize(psi_0) + traj_vector.propagate(t_max, t_step) + return np.array(traj_vector.storage['psi_traj']) + traj_tensor = HopsTensorTrajectory( + system_param=sys_param_topology, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integrator_param, + tensor_param={ + 'MPS_EPSILON': mps_epsilon, 'METHOD': method, + 'BOND_DIM_MAX': bond_dim_max, + 'FLAG_MPO_OPTIMIZE': flag_mpo_optimize, + }, + ) + traj_tensor.initialize(psi_0) + traj_tensor.propagate(t_max, t_step) + return np.array(traj_tensor.storage['psi_traj']) + + +# ------------------------------------------------------------ +# TEST: Ring and star statenumber MPOs match vector HOPS +# ------------------------------------------------------------ +@pytest.mark.level(2) +@pytest.mark.parametrize('topology', ['ring', 'star']) +@pytest.mark.parametrize('flag_mpo_optimize', [True, False]) +def test_statenumber_topology_matches_vector_hops(topology, flag_mpo_optimize): + # This case tests that ring and star Hamiltonians reproduce vector HOPS + # both through their combined-generator MPOs (FLAG_MPO_OPTIMIZE True) and + # through the general hierarchy-plus-Hamiltonian path (False). Both are + # long-ranged Hamiltonians that the nearest-neighbor MPO cannot + # represent, and the combined builders reach them at a bond dimension + # fixed by the topology rather than by the interaction range. + H2_ham_topology = { + 'ring': H2_ham_ring, 'star': H2_ham_star, + }[topology] + psi_vector = _run_hops_topology(H2_ham_topology) + psi_tensor = _run_hops_topology( + H2_ham_topology, method='number', + flag_mpo_optimize=flag_mpo_optimize, + ) + + assert len(psi_vector) == len(psi_tensor), ( + 'Trajectory lengths differ' + ) + max_err = max( + np.linalg.norm(psi_tensor[i] - psi_vector[i]) + for i in range(len(psi_vector)) + ) + # Largest error observed across the two topologies is 3.8e-12 with the + # optimization on and 2.2e-12 with it off. + assert max_err < 1e-11, ( + f'Statenumber {topology} tensor HOPS diverged from vector HOPS: ' + f'max wf error = {max_err:.2e}' + ) diff --git a/tests/test_eom_hops_ksuper.py b/tests/test_eom_hops_ksuper.py index f7c9179..7cef809 100644 --- a/tests/test_eom_hops_ksuper.py +++ b/tests/test_eom_hops_ksuper.py @@ -323,7 +323,7 @@ def test_add_self_interaction_remove_aux(): Tests that _add_self_interaction() produces the correct time-independent self-interaction time-evolution matrix a) in general and b) when we remove auxiliaries from the basis and add them back in (that is, check accuracy and - self-consistency of the pyHOPS architecture that manages the self-interaction + self-consistency of the MesoHOPS architecture that manages the self-interaction terms). """ # Prepare Constants diff --git a/tests/test_hierarchy_class.py b/tests/test_hops_hierarchy.py similarity index 58% rename from tests/test_hierarchy_class.py rename to tests/test_hops_hierarchy.py index b36be59..641aa5d 100644 --- a/tests/test_hierarchy_class.py +++ b/tests/test_hops_hierarchy.py @@ -7,10 +7,16 @@ from mesohops.util.exceptions import UnsupportedRequest -__title__ = "test of hops hierarchy" -__author__ = "D. I. G. Bennett, L. Varvelo, J. K. Lynd" -__version__ = "1.6" -__date__ = "Aug. 13, 2025" +__title__ = 'Unit Tests for HopsHierarchy' +__author__ = 'D. I. G. Bennett, L. Varvelo, J. K. Lynd' +__version__ = '1.6' + + +def test_maxhier_overflow_warning(): + # This case tests that constructing a hierarchy with MAXHIER > 255 + # emits a warning about integer overflow. + with pytest.warns(UserWarning, match='integer overflow'): + HHier({'MAXHIER': 256}, {'N_HMODES': 2}) def test_hierarchy_initialize_true(): @@ -190,22 +196,8 @@ def test_aux_index_relative(): assert HH._aux_index(test_aux) == 10 -def test_const_aux_edge(): - """ - Tests whether const_aux_edge is properly creating an auxiliary index tuple for - an edge node at a particular depth along a given mode. - """ - - hierarchy_param = {"MAXHIER": 4} - system_param = {"N_HMODES": 4} - HH = HHier(hierarchy_param, system_param) - HH.initialize(True) - tmp = HH._const_aux_edge(2, 1, 4) - known_index_tuple = AuxVec([(2, 1)], 4) - assert tmp == known_index_tuple - -# a function created to be used to test define_triangular_hierarchy +# a helper used to test the hierarchy builder functions def map_to_auxvec(list_aux, n_hmodes): """ Helper function that maps a list of auxiliary indexing vectors to a list of @@ -215,12 +207,14 @@ def map_to_auxvec(list_aux, n_hmodes): ---------- 1. list_aux : list(list(int)) List of auxiliary indexing vectors in list form + 2. n_hmodes : int + Number of modes in the hierarchy RETURNS ------- 1. list_aux_vec : list(list(AuxVec)) List of AuxVec objects corresponding to the indexing vectors in - list_aux in a hierarchy with 4 modes + list_aux in a hierarchy with n_hmodes modes """ list_aux_vec = [] @@ -522,3 +516,273 @@ def test_add_connections(): assert vector_030._dict_aux_p1 == {} assert vector_003._dict_aux_m1 == {2:vector_002} assert vector_003._dict_aux_p1 == {} + + +# ============================================================ +# TEST SUITE: define_rectangular_hierarchy() +# ============================================================ + +# ------------------------------------------------------------ +# TEST: Correct auxiliaries for two hierarchy modes +# ------------------------------------------------------------ +def test_define_rect_hier_two_modes(): + # This case tests n_hmodes=2, maxhier=2 producing all 9 Cartesian + # product combinations. Uses sorted list comparison (not set) to + # catch duplicates in the output. + list_aux = HHier.define_rectangular_hierarchy(2, 2) + known_list = map_to_auxvec([ + [0, 0], [0, 1], [0, 2], + [1, 0], [1, 1], [1, 2], + [2, 0], [2, 1], [2, 2], + ], 2) + assert sorted(list_aux) == sorted(known_list), ( + 'Two-mode rectangular hierarchy does not match expected Cartesian product' + ) + + +# ------------------------------------------------------------ +# TEST: Rectangular is a strict superset of triangular for +# n_hmodes > 1 +# ------------------------------------------------------------ +@pytest.mark.parametrize('n_hmodes, maxhier', [(3, 2), (2, 3)]) +def test_define_rect_hier_superset_of_triangular(n_hmodes, maxhier): + # This case tests that the rectangular hierarchy is a strict superset + # of the triangular hierarchy for n_hmodes > 1. + list_rect = HHier.define_rectangular_hierarchy(n_hmodes, maxhier) + list_tri = HHier.define_triangular_hierarchy(n_hmodes, maxhier) + set_rect = set(list_rect) + set_tri = set(list_tri) + assert set_tri.issubset(set_rect), ( + 'Triangular hierarchy should be a subset of rectangular hierarchy' + ) + assert len(set_rect) > len(set_tri), ( + 'Rectangular hierarchy should be strictly larger than triangular ' + f'for n_hmodes > 1 (rect={len(set_rect)}, tri={len(set_tri)})' + ) + + +# ------------------------------------------------------------ +# TEST: Single-mode rectangular and triangular are identical +# ------------------------------------------------------------ +def test_define_rect_hier_single_mode_matches_triangular(): + # This case tests that for n_hmodes=1, rectangular and triangular + # truncation produce identical hierarchies. + n_hmodes = 1 + maxhier = 4 + list_rect = HHier.define_rectangular_hierarchy(n_hmodes, maxhier) + list_tri = HHier.define_triangular_hierarchy(n_hmodes, maxhier) + assert set(list_rect) == set(list_tri), ( + 'Rectangular and triangular hierarchies should be identical for ' + 'n_hmodes=1' + ) + + +# ------------------------------------------------------------ +# TEST: Zero vector is always present in rectangular hierarchy +# ------------------------------------------------------------ +def test_define_rect_hier_contains_zero_vector(): + # This case tests that the zero auxiliary vector (vacuum state) + # is always present in the rectangular hierarchy. + n_hmodes = 3 + maxhier = 2 + list_aux = HHier.define_rectangular_hierarchy(n_hmodes, maxhier) + zero_aux = AuxVec([], n_hmodes) + assert zero_aux in list_aux, ( + 'Zero auxiliary vector should always be present in rectangular hierarchy' + ) + + +# ============================================================ +# TEST SUITE: HopsHierarchy.initialize() with rectangular truncation +# ============================================================ + +# ------------------------------------------------------------ +# TEST: Non-adaptive initialization with rectangular truncation +# produces the rectangular hierarchy +# ------------------------------------------------------------ +def test_initialize_rect_trunc_nonadaptive(): + # This case tests that initializing with TRUNCATION_METHOD='rectangular' + # produces the same hierarchy as define_rectangular_hierarchy. + hierarchy_param = {'MAXHIER': 2, 'TRUNCATION_METHOD': 'rectangular'} + system_param = {'N_HMODES': 2} + HH = HHier(hierarchy_param, system_param) + HH.initialize(False) + expected = HHier.define_rectangular_hierarchy(2, 2) + assert len(HH.auxiliary_list) == len(expected), ( + f'Expected {len(expected)} auxiliaries, got {len(HH.auxiliary_list)}' + ) + assert set(HH.auxiliary_list) == set(expected), ( + 'Initialized rect hierarchy does not match define_rectangular_hierarchy' + ) + + +# ------------------------------------------------------------ +# TEST: Rectangular truncation with Markovian filter warns and +# produces correct filtered hierarchy +# ------------------------------------------------------------ +def test_rect_trunc_with_markovian_filter(): + # This case tests that rectangular truncation with a Markovian filter + # warns about tensor-specific filters, applies the filter, and + # produces the correct hierarchy. + hierarchy_param = { + 'MAXHIER': 2, + 'TRUNCATION_METHOD': 'rectangular', + 'STATIC_FILTERS': [('Markovian', [False, True])], + } + system_param = {'N_HMODES': 2} + HH = HHier(hierarchy_param, system_param) + with pytest.warns(UserWarning, match='tensor-specific filters'): + HH.initialize(False) + # Markovian on mode 1 removes any aux with mode 1 active at total + # depth > 1. From the 9-element rectangular hierarchy, only 4 survive. + known = [ + AuxVec([], 2), + AuxVec([(0, 1)], 2), + AuxVec([(1, 1)], 2), + AuxVec([(0, 2)], 2), + ] + assert sorted(HH.auxiliary_list) == sorted(known) + + +# ------------------------------------------------------------ +# TEST: Rectangular + Markovian filter matches triangular + +# Markovian filter (filtered modes produce same hierarchy) +# ------------------------------------------------------------ +def test_rect_trunc_markovian_matches_triangular(): + # This case tests that applying the same Markovian filter to + # rectangular and triangular hierarchies produces the same result, + # since filtering collapses the extra rectangular elements. + markov_filter = [('Markovian', [False, True])] + system_param = {'N_HMODES': 2} + rect_param = { + 'MAXHIER': 2, + 'TRUNCATION_METHOD': 'rectangular', + 'STATIC_FILTERS': list(markov_filter), + } + tri_param = { + 'MAXHIER': 2, + 'TRUNCATION_METHOD': 'triangular', + 'STATIC_FILTERS': list(markov_filter), + } + HH_rect = HHier(rect_param, system_param) + HH_tri = HHier(tri_param, system_param) + with pytest.warns(UserWarning, match='tensor-specific filters'): + HH_rect.initialize(False) + HH_tri.initialize(False) + assert sorted(HH_rect.auxiliary_list) == sorted(HH_tri.auxiliary_list) + + +# ------------------------------------------------------------ +# TEST: Rectangular truncation with Triangular filter produces +# correct filtered hierarchy +# ------------------------------------------------------------ +def test_rect_trunc_with_triangular_filter(): + # This case tests that applying a Triangular filter to a rectangular + # hierarchy correctly prunes auxiliaries. The Triangular filter on + # mode 1 with kmax=1 restricts the total depth in filtered modes + # to be <= 1. + hierarchy_param = { + 'MAXHIER': 2, + 'TRUNCATION_METHOD': 'rectangular', + 'STATIC_FILTERS': [('Triangular', [[False, True], 1])], + } + system_param = {'N_HMODES': 2} + HH = HHier(hierarchy_param, system_param) + with pytest.warns(UserWarning, match='tensor-specific filters'): + HH.initialize(False) + # Triangular filter on mode 1 with kmax_2=1: keeps only auxiliaries + # where the sum of depth in filtered modes (mode 1) is <= 1. + for aux in HH.auxiliary_list: + depth_mode1 = aux.get(1, 0) + assert depth_mode1 <= 1, ( + f'Triangular filter failed: aux {aux} has mode-1 depth ' + f'{depth_mode1}, expected <= 1' + ) + # Should be smaller than unfiltered rectangular + rect_unfiltered = HHier.define_rectangular_hierarchy(2, 2) + assert len(HH.auxiliary_list) < len(rect_unfiltered) + + +# ------------------------------------------------------------ +# TEST: Rectangular truncation with LongEdge filter produces +# correct filtered hierarchy +# ------------------------------------------------------------ +def test_rect_trunc_with_longedge_filter(): + # This case tests that applying a LongEdge filter to a rectangular + # hierarchy correctly prunes auxiliaries. LongEdge with kdepth=1 + # on mode 1: beyond total depth 1, only edge terms (single-mode + # depth) are kept for filtered modes. + hierarchy_param = { + 'MAXHIER': 2, + 'TRUNCATION_METHOD': 'rectangular', + 'STATIC_FILTERS': [('LongEdge', [[False, True], 1])], + } + system_param = {'N_HMODES': 2} + HH = HHier(hierarchy_param, system_param) + with pytest.warns(UserWarning, match='tensor-specific filters'): + HH.initialize(False) + known = [ + AuxVec([], 2), + AuxVec([(0, 1)], 2), + AuxVec([(1, 1)], 2), + AuxVec([(0, 2)], 2), + AuxVec([(1, 2)], 2), + ] + assert sorted(HH.auxiliary_list) == sorted(known) + + +# ------------------------------------------------------------ +# TEST: Rectangular + LongEdge matches triangular + LongEdge +# ------------------------------------------------------------ +def test_rect_trunc_longedge_matches_triangular(): + # This case tests that applying the same LongEdge filter to + # rectangular and triangular hierarchies produces the same result. + longedge_filter = [('LongEdge', [[False, True], 1])] + system_param = {'N_HMODES': 2} + rect_param = { + 'MAXHIER': 2, + 'TRUNCATION_METHOD': 'rectangular', + 'STATIC_FILTERS': list(longedge_filter), + } + tri_param = { + 'MAXHIER': 2, + 'TRUNCATION_METHOD': 'triangular', + 'STATIC_FILTERS': list(longedge_filter), + } + HH_rect = HHier(rect_param, system_param) + HH_tri = HHier(tri_param, system_param) + with pytest.warns(UserWarning, match='tensor-specific filters'): + HH_rect.initialize(False) + HH_tri.initialize(False) + assert sorted(HH_rect.auxiliary_list) == sorted(HH_tri.auxiliary_list) + + +# ------------------------------------------------------------ +# TEST: Rectangular truncation with multiple filters +# (Markovian + Triangular) +# ------------------------------------------------------------ +def test_rect_trunc_with_combined_filters(): + # This case tests that applying multiple filters to a rectangular + # hierarchy correctly chains the filtering. Markovian on mode 1 + # followed by Triangular on mode 0 with kmax=1. + hierarchy_param = { + 'MAXHIER': 2, + 'TRUNCATION_METHOD': 'rectangular', + 'STATIC_FILTERS': [ + ('Markovian', [False, True]), + ('Triangular', [[True, False], 1]), + ], + } + system_param = {'N_HMODES': 2} + HH = HHier(hierarchy_param, system_param) + with pytest.warns(UserWarning, match='tensor-specific filters'): + HH.initialize(False) + # Markovian on mode 1 gives: {}, {0:1}, {1:1}, {0:2} + # Triangular on mode 0 with kmax=1 keeps only depth <= 1 in mode 0 + # That removes {0:2}, leaving: {}, {0:1}, {1:1} + known = [ + AuxVec([], 2), + AuxVec([(0, 1)], 2), + AuxVec([(1, 1)], 2), + ] + assert sorted(HH.auxiliary_list) == sorted(known) diff --git a/tests/test_hops_storage.py b/tests/test_hops_storage.py index d9b46ba..8f71353 100644 --- a/tests/test_hops_storage.py +++ b/tests/test_hops_storage.py @@ -1,13 +1,20 @@ +import copy + import numpy as np import pytest +from unittest.mock import MagicMock, patch import mesohops.storage.storage_functions as sf from mesohops.basis.hops_aux import AuxiliaryVector from mesohops.storage.hops_storage import HopsStorage +from mesohops.storage.storage_functions import ( + save_max_tensor_complexity, + save_phi_traj_tensor, + save_phi_norm_tensor, +) from mesohops.trajectory.exp_noise import bcf_exp from mesohops.trajectory.hops_trajectory import HopsTrajectory from mesohops.util.exceptions import UnsupportedRequest -from unittest.mock import patch # New Hops_Storage tests @@ -553,3 +560,114 @@ def test_git_commit_hash_error_in_metadata(mock_get_hash): assert "GIT_COMMIT_HASH" in storage.metadata assert storage.metadata["GIT_COMMIT_HASH"] == ("A gigantic walrus ate the git " "repository here too") + + +# ============================================================ +# TEST SUITE: tensor storage functions +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: save_phi_traj_tensor returns deep copy of MPS cores +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_save_phi_traj_tensor_returns_deep_copy(): + # This case tests that save_phi_traj_tensor returns a deep copy + # of the MPS cores, not a reference to the original list. + mock_wfn = MagicMock() + core_0 = np.array([[[1.0 + 0j, 2.0], [3.0, 4.0]]]) + core_1 = np.array([[[5.0 + 0j], [6.0]], [[7.0], [8.0]]]) + mock_wfn.list_cores_phi = [core_0, core_1] + + result = save_phi_traj_tensor(wavefunction=mock_wfn) + + assert len(result) == 2 + np.testing.assert_array_equal(result[0], core_0) + np.testing.assert_array_equal(result[1], core_1) + # Verify deep copy: mutating original should not affect result + mock_wfn.list_cores_phi[0][0, 0, 0] = 999.0 + assert result[0][0, 0, 0] != 999.0 + + +# ------------------------------------------------------------ +# TEST: save_phi_norm_tensor computes hierarchy norm via contraction +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_save_phi_norm_tensor(): + # This case tests that save_phi_norm_tensor computes the correct + # hierarchy norm using contract_down_exact. For a simple 2-state + # fullstate MPS with known cores, verify the norm matches the + # analytical value. + mock_wfn = MagicMock() + # Simple 2-state fullstate MPS: state core + one mode core + # State core: (1, 2, 1), mode core: (1, 2, 1) + core_state = np.zeros((1, 2, 1), dtype=np.complex128) + core_state[0, 0, 0] = 1.0 # psi on state 0 + core_state[0, 1, 0] = 0.5 # psi on state 1 + core_mode = np.zeros((1, 2, 1), dtype=np.complex128) + core_mode[0, 0, 0] = 1.0 # k=0 occupation + core_mode[0, 1, 0] = 0.3 # k=1 occupation + mock_wfn.list_cores_phi = [core_state, core_mode] + mock_wfn.method = 'fullstate' + # M1_modes_per_site is the attribute save_phi_norm_tensor reads + mock_wfn.M1_modes_per_site = np.array([1, 1]) + + result = save_phi_norm_tensor(wavefunction=mock_wfn) + + # Analytical: per-state norm^2 from double-layer contraction: + # state 0: |1.0*1.0|^2 + |1.0*0.3|^2 = 1.09 + # state 1: |0.5*1.0|^2 + |0.5*0.3|^2 = 0.2725 + # Total norm = sqrt(1.09 + 0.2725) = sqrt(1.3625) + expected = np.sqrt(1.3625) + np.testing.assert_allclose(result, expected, atol=1e-12) + + +# ------------------------------------------------------------ +# TEST: per-state norms are individually correct +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_save_phi_norm_tensor_per_state(): + # This case tests that the per-state norm-squared values from + # contract_down_exact match the analytical values, not just + # the total norm. + from mesohops.util.tensor_operations import contract_down_exact + mock_wfn = MagicMock() + core_state = np.zeros((1, 2, 1), dtype=np.complex128) + core_state[0, 0, 0] = 1.0 + core_state[0, 1, 0] = 0.5 + core_mode = np.zeros((1, 2, 1), dtype=np.complex128) + core_mode[0, 0, 0] = 1.0 + core_mode[0, 1, 0] = 0.3 + mock_wfn.list_cores_phi = [core_state, core_mode] + mock_wfn.method = 'fullstate' + mock_wfn.M1_modes_per_site = np.array([1, 1]) + + V1_norm_sq = contract_down_exact( + mock_wfn.list_cores_phi, + mock_wfn.method, + len(mock_wfn.M1_modes_per_site), + ) + # Analytical per-state norm^2: + # state 0: |1.0*1.0|^2 + |1.0*0.3|^2 = 1.09 + # state 1: |0.5*1.0|^2 + |0.5*0.3|^2 = 0.2725 + np.testing.assert_allclose(V1_norm_sq[0].real, 1.09, atol=1e-12) + np.testing.assert_allclose(V1_norm_sq[1].real, 0.2725, atol=1e-12) + + +# ------------------------------------------------------------ +# TEST: save_max_tensor_complexity returns the scalar verbatim +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_save_max_tensor_complexity(): + # This case tests that save_max_tensor_complexity is an identity + # function over its `max_tensor_complexity` kwarg. Propagate passes + # the scalar from eom.max_complexity_step verbatim, and the storage + # callable forwards it unchanged into storage.data. + assert save_max_tensor_complexity(max_tensor_complexity=42) == 42 + assert save_max_tensor_complexity(max_tensor_complexity=0) == 0 + # Other kwargs are accepted (storage passes wavefunction etc. to + # every save function uniformly) and ignored here. + result = save_max_tensor_complexity( + max_tensor_complexity=7, wavefunction=None, t_new=1.5, + ) + assert result == 7 diff --git a/tests/test_hops_system.py b/tests/test_hops_system.py index 4ea82f3..8e86fbc 100644 --- a/tests/test_hops_system.py +++ b/tests/test_hops_system.py @@ -716,3 +716,38 @@ def __getitem__(self, key): with pytest.raises(TypeError): HSystem.reduce_sparse_matrix(BrokenDict(), [0, 1], True, filter_nz=True) + + +# ============================================================ +# TEST SUITE: _is_nearest_neighbor() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Auto-detects NN vs non-NN Hamiltonians +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_is_nearest_neighbor_auto_detection(): + # This case tests that a tridiagonal Hamiltonian is detected as NN. + H2_nn_dense = np.zeros([4, 4], dtype=np.complex128) + H2_nn_dense[0, 1] = 40 + H2_nn_dense[1, 0] = 40 + H2_nn_dense[1, 2] = 10 + H2_nn_dense[2, 1] = 10 + H2_nn_dense[2, 3] = 40 + H2_nn_dense[3, 2] = 40 + H2_nn = sp.sparse.coo_matrix(H2_nn_dense) + assert HSystem._is_nearest_neighbor(H2_nn) is True + + # This case tests that a Hamiltonian with long-range coupling is not NN. + H2_non_nn_dense = np.array(H2_nn_dense, dtype=np.complex128) + H2_non_nn_dense[0, 2] = 5.0 + H2_non_nn_dense[2, 0] = 5.0 + H2_non_nn = sp.sparse.coo_matrix(H2_non_nn_dense) + assert HSystem._is_nearest_neighbor(H2_non_nn) is False + + # This case tests that explicit sparse zeros on far off-diagonals + # do not break NN detection. + H2_sparse_zero = sp.sparse.lil_matrix(H2_nn_dense) + H2_sparse_zero[0, 3] = 0.0 + assert HSystem._is_nearest_neighbor(H2_sparse_zero.tocoo()) is True diff --git a/tests/test_hops_tensor_basis.py b/tests/test_hops_tensor_basis.py new file mode 100644 index 0000000..6927ccd --- /dev/null +++ b/tests/test_hops_tensor_basis.py @@ -0,0 +1,250 @@ +import numpy as np +import pytest +import scipy as sp + +from mesohops.basis.hops_modes import HopsModes +from mesohops.basis.hops_noise_memory import HopsNoiseMemory +from mesohops.basis.hops_system import HopsSystem +from mesohops.tensor.hops_tensor_wavefunction import HopsTensorWavefunction +from mesohops.tensor.hops_tensor_basis import HopsTensorBasis +from mesohops.trajectory.exp_noise import bcf_exp +from mesohops.util.bath_corr_functions import bcf_convert_dl_to_exp + +__title__ = 'Unit Tests for HopsTensorBasis' +__author__ = 'N. Covalsen' +__maintainer__ = 'N. Covalsen' + +# ============================================================ +# Shared Setup +# ============================================================ + +nsite = 4 +e_lambda = 20.0 +gamma = 50.0 +temp = 140.0 +(g_0, w_0) = bcf_convert_dl_to_exp(e_lambda, gamma, temp) + +loperator = np.zeros([4, 4, 4], dtype=np.float64) +gw_sysbath = [] +lop_list = [] +for i in range(nsite): + loperator[i, i, i] = 1.0 + gw_sysbath.append([g_0, w_0]) + lop_list.append(sp.sparse.coo_matrix(loperator[i])) + gw_sysbath.append([-1j * np.imag(g_0), 500.0]) + lop_list.append(loperator[i]) + +hs = np.zeros([nsite, nsite]) +hs[0, 1] = 40 +hs[1, 0] = 40 +hs[1, 2] = 10 +hs[2, 1] = 10 +hs[2, 3] = 40 +hs[3, 2] = 40 + +sys_param = { + 'HAMILTONIAN': np.array(hs, dtype=np.complex128), + 'GW_SYSBATH': gw_sysbath, + 'L_HIER': lop_list, + 'L_NOISE1': lop_list, + 'ALPHA_NOISE1': bcf_exp, + 'PARAM_NOISE1': gw_sysbath, +} + +psi_0 = np.array([0.0] * nsite, dtype=np.complex128) +psi_0[2] = 1.0 + +state_list = list(np.arange(nsite)) + +k_max = 4 + +tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, +} + +integrator_param = {'INTEGRATOR': 'RUNGE_KUTTA'} + +eom_param = {'EQUATION_OF_MOTION': 'NORMALIZED NONLINEAR'} + + +class _MockEOM: + """Minimal mock for HopsTensorEOM with a no-op refresh_builder.""" + + def refresh_builder(self): + pass + + +def _make_basis_objects(sp=sys_param): + '''Creates HopsSystem, HopsModes, HopsNoiseMemory directly.''' + system = HopsSystem(sp) + mode = HopsModes(system, hierarchy=None) + noise_memory = HopsNoiseMemory(system, mode) + return system, mode, noise_memory + + +def _make_tensor_basis(): + '''Creates a HopsTensorBasis with directly-constructed objects.''' + system, mode, noise_memory = _make_basis_objects() + return HopsTensorBasis(system, mode, noise_memory) + + +def _make_initialized_tensor_basis(delta_s=0, sl=None, psi=None): + '''Creates and initializes a HopsTensorBasis.''' + if sl is None: + sl = state_list + if psi is None: + psi = psi_0 + system, mode, noise_memory = _make_basis_objects() + tb = HopsTensorBasis(system, mode, noise_memory) + # Manually initialize the shared objects (normally done by trajectory) + system.initialize(delta_s > 0, psi) + mode.list_modeidx_abs = sorted(system.list_statemodeidx_abs) + noise_memory.initialize() + system.state_list = sl + tb.initialize(delta_s) + return tb + + +def _make_tensor_for_tb(tb, method='fullstate'): + '''Creates and initializes a HopsTensorWavefunction tied to a basis.''' + tp = dict(tensor_param, METHOD=method) + ht = HopsTensorWavefunction(k_max, tp, integrator_param, eom_param) + ht.initialize( + psi_0[tb.system.state_list], + tb.system, + ) + return ht + + +# ============================================================ +# TEST SUITE: HopsTensorBasis() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Constructor stores references and defaults +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_init_stores_references(): + '''HopsTensorBasis constructor stores refs and sets defaults.''' + # This case tests that __init__ stores system, mode, noise_memory refs + system = HopsSystem(sys_param) + mode = HopsModes(system, hierarchy=None) + noise_memory = HopsNoiseMemory(system, mode) + tb = HopsTensorBasis(system, mode, noise_memory) + assert tb.system is system + assert tb.mode is mode + assert tb.noise_memory is noise_memory + # This case tests that defaults are set correctly + assert tb.eom is None + assert tb.adaptive is False + assert tb.delta_s == 0 + + +# ------------------------------------------------------------ +# TEST: initialize() with delta_s=0 sets non-adaptive +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_initialize_nonadaptive(): + '''initialize with delta_s=0 leaves adaptive=False.''' + # Analytical: delta_s=0 means non-adaptive + tb = _make_initialized_tensor_basis() + assert tb.adaptive is False + assert tb.delta_s == 0 + + +# ------------------------------------------------------------ +# TEST: initialize() with delta_s>0 sets adaptive +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_initialize_adaptive(): + '''initialize with delta_s>0 sets adaptive=True.''' + # Analytical: delta_s > 0 means adaptive + tb = _make_initialized_tensor_basis(delta_s=0.1) + assert tb.adaptive is True + assert tb.delta_s == 0.1 + + +# ------------------------------------------------------------ +# TEST: define_basis non-adaptive returns empty lists +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_define_basis_nonadaptive_returns_empty(): + '''define_basis returns ([], []) when not adaptive.''' + # Analytical: non-adaptive define_basis returns empty lists immediately + tb = _make_initialized_tensor_basis(delta_s=0) + ht = _make_tensor_for_tb(tb) + z_step = [np.zeros(1)] * len(tb.mode.list_modeidx_abs) + list_old, list_new = tb.define_basis(ht, z_step) + assert list_old == [] + assert list_new == [] + + +# ------------------------------------------------------------ +# TEST: update_basis adding a state grows state_list +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_update_basis_add_grows_state_list(): + '''update_basis with list_states_new grows system.state_list.''' + # Setup: 3-state basis [0, 1, 2], add state 3 + sl_initial = [0, 1, 2] + psi_3 = np.array([0.0, 0.0, 1.0, 0.0], dtype=np.complex128) + tb = _make_initialized_tensor_basis(sl=sl_initial, psi=psi_3) + ht = _make_tensor_for_tb(tb) + tb.eom = _MockEOM() + z_mem = np.zeros(len(tb.mode.list_modeidx_abs), dtype=np.complex128) + state_list_before = sorted(tb.system.state_list) + # This case tests that adding state 3 grows the state_list + wf, z_out = tb.update_basis(ht, z_mem, [], [3]) + # Analytical: state_list should grow by 1 + assert 3 in tb.system.state_list + assert len(tb.system.state_list) == len(state_list_before) + 1 + # Analytical: z_mem passes through unchanged + assert z_out is z_mem + + +# ------------------------------------------------------------ +# TEST: update_basis removing a state shrinks state_list +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_update_basis_remove_shrinks_state_list(): + '''update_basis with list_states_old shrinks system.state_list.''' + # Setup: full 4-state basis, remove state 3 + tb = _make_initialized_tensor_basis() + ht = _make_tensor_for_tb(tb) + tb.eom = _MockEOM() + z_mem = np.zeros(len(tb.mode.list_modeidx_abs), dtype=np.complex128) + state_list_before = sorted(tb.system.state_list) + # This case tests that removing state 3 shrinks the state_list + wf, z_out = tb.update_basis(ht, z_mem, [3], []) + # Analytical: state_list should shrink by 1 + assert 3 not in tb.system.state_list + assert len(tb.system.state_list) == len(state_list_before) - 1 + # Analytical: wavefunction is returned (same object, mutated in place) + assert wf is ht + + +# ------------------------------------------------------------ +# TEST: update_basis roundtrip (add then remove) restores state_list +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_update_basis_roundtrip(): + '''Adding then removing a state returns state_list to original.''' + # Setup: 3-state basis [0, 1, 2] + sl_initial = [0, 1, 2] + psi_3 = np.array([0.0, 0.0, 1.0, 0.0], dtype=np.complex128) + tb = _make_initialized_tensor_basis(sl=sl_initial, psi=psi_3) + ht = _make_tensor_for_tb(tb) + tb.eom = _MockEOM() + z_mem = np.zeros(len(tb.mode.list_modeidx_abs), dtype=np.complex128) + state_list_original = sorted(tb.system.state_list) + # This case tests that add followed by remove restores original state_list + # Step 1: add state 3 + wf, z_out = tb.update_basis(ht, z_mem, [], [3]) + assert 3 in tb.system.state_list + # Step 2: remove state 3 + wf, z_out = tb.update_basis(wf, z_out, [3], []) + # Analytical: state_list should match original + assert sorted(tb.system.state_list) == state_list_original diff --git a/tests/test_hops_tensor_eom.py b/tests/test_hops_tensor_eom.py new file mode 100644 index 0000000..c96fd9d --- /dev/null +++ b/tests/test_hops_tensor_eom.py @@ -0,0 +1,1137 @@ +# tests/test_hops_tensor_eom.py +import numpy as np +import pytest +import scipy as sp + +from mesohops.basis.hops_modes import HopsModes +from mesohops.basis.hops_noise_memory import HopsNoiseMemory +from mesohops.basis.hops_system import HopsSystem +from mesohops.tensor.hops_tensor_wavefunction import HopsTensorWavefunction +from mesohops.tensor.hops_tensor_basis import HopsTensorBasis +from mesohops.eom.eom_functions import calc_delta_zmem, operator_expectation +from mesohops.tensor.hops_tensor_eom import HopsTensorEOM +from mesohops.tensor.tensor_eom_functions import tensor_matvec_prod +from mesohops.trajectory.exp_noise import bcf_exp +from mesohops.util.bath_corr_functions import bcf_convert_dl_to_exp +from mesohops.util.exceptions import UnsupportedRequest +from mesohops.util.tensor_operations import ( + _statenumber_offsets, + extract_psi, + tensor_add, +) + +__title__ = 'Unit Tests for HopsTensorEOM' +__author__ = 'N. Covalsen' +__maintainer__ = 'N. Covalsen' + +# ---- shared fixtures (dimer-of-dimers, same as test_hops_tensor_unit.py) ---- +nsite = 4 +e_lambda = 20.0 +gamma = 50.0 +temp = 140.0 +(g_0, w_0) = bcf_convert_dl_to_exp(e_lambda, gamma, temp) + +loperator = np.zeros([4, 4, 4], dtype=np.float64) +gw_sysbath = [] +lop_list = [] +for i in range(nsite): + loperator[i, i, i] = 1.0 + gw_sysbath.append([g_0, w_0]) + lop_list.append(sp.sparse.coo_matrix(loperator[i])) + gw_sysbath.append([-1j * np.imag(g_0), 500.0]) + lop_list.append(loperator[i]) + +hs = np.zeros([nsite, nsite]) +hs[0, 1] = 40 +hs[1, 0] = 40 +hs[1, 2] = 10 +hs[2, 1] = 10 +hs[2, 3] = 40 +hs[3, 2] = 40 + +sys_param = { + 'HAMILTONIAN': np.array(hs, dtype=np.complex128), + 'GW_SYSBATH': gw_sysbath, + 'L_HIER': lop_list, + 'L_NOISE1': lop_list, + 'ALPHA_NOISE1': bcf_exp, + 'PARAM_NOISE1': gw_sysbath, +} +psi_0 = np.array([0.0] * nsite, dtype=np.complex128) +psi_0[2] = 1.0 +k_max = 2 +state_list = np.arange(nsite) +delta_s = 0 +eom_param = {'EQUATION_OF_MOTION': 'NORMALIZED NONLINEAR'} + + +def _make_tb(sp=sys_param, ds=delta_s, psi=psi_0, sl=state_list): + """Creates an initialized HopsTensorBasis from sys_param dict.""" + system = HopsSystem(sp) + mode = HopsModes(system, hierarchy=None) + noise_memory = HopsNoiseMemory(system, mode) + tb = HopsTensorBasis(system, mode, noise_memory) + system.initialize(ds > 0, psi) + mode.list_modeidx_abs = sorted(system.list_statemodeidx_abs) + noise_memory.initialize() + system.state_list = sl + tb.initialize(ds) + return tb + + +def _make_eom(method='fullstate', flag_mpo_optimize=True): + """Creates an initialized HopsTensorEOM for the dimer-of-dimers.""" + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': method, + 'BOND_DIM_MAX': 20, + 'FLAG_MPO_OPTIMIZE': flag_mpo_optimize, + } + tb = _make_tb() + ht = HopsTensorWavefunction( + k_max, tensor_param, {'INTEGRATOR': 'RUNGE_KUTTA'}, eom_param, + ) + ht.initialize(psi_0, tb.system) + return HopsTensorEOM(ht, tb.system, tb.mode, tb.noise_memory, tb.adaptive, eom_param) + + +def _make_noise(eom): + """Returns non-trivial z_mem, z_rnd, z_rnd2 arrays sized for eom. + + Each array uses a different complex prefactor times np.arange to break + all symmetries and ensure z_rnd2 is never zero. + """ + n_mem = len(eom.noise_memory.list_zmemmodeidx_abs) + n_l2 = len(eom.mode.list_l2idx_abs) + z_mem = (0.01 + 0.02j) * np.arange(n_mem, dtype=np.complex128) + z_rnd = (0.03 - 0.01j) * np.arange(n_l2, dtype=np.complex128) + z_rnd2 = (-0.02 + 0.04j) * np.arange(n_l2, dtype=np.complex128) + return z_mem, z_rnd, z_rnd2 + + +def _make_noise_from_basis(tb): + """Returns non-trivial z_mem, z_rnd, z_rnd2 from a HopsTensorBasis. + + Same prefactors as _make_noise, for use in tests that construct EOM + instances manually from a shared basis. + """ + n_mem = len(tb.noise_memory.list_zmemmodeidx_abs) + n_l2 = len(tb.mode.list_l2idx_abs) + z_mem = (0.01 + 0.02j) * np.arange(n_mem, dtype=np.complex128) + z_rnd = (0.03 - 0.01j) * np.arange(n_l2, dtype=np.complex128) + z_rnd2 = (-0.02 + 0.04j) * np.arange(n_l2, dtype=np.complex128) + return z_mem, z_rnd, z_rnd2 + + +# ============================================================ +# TEST SUITE: build_generator() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Return value has correct shape and type +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_generator_returns_correct_shape(): + # This case tests that dz_dt has the same shape as z_mem and is complex. + eom = _make_eom() + z_mem, z_rnd, z_rnd2 = _make_noise(eom) + dz_dt = eom.build_generator(z_mem, z_rnd, z_rnd2) + assert dz_dt.shape == z_mem.shape + assert np.iscomplexobj(dz_dt) + # Cross-validation: fullstate and statenumber dz_dt should agree + eom_sn = _make_eom('number') + z_mem_sn, z_rnd_sn, z_rnd2_sn = _make_noise(eom_sn) + dz_dt_sn = eom_sn.build_generator(z_mem_sn, z_rnd_sn, z_rnd2_sn) + np.testing.assert_allclose( + dz_dt, dz_dt_sn, atol=1e-8, + err_msg='dz_dt should agree between fullstate and statenumber', + ) + + +# ------------------------------------------------------------ +# TEST: MPO cores populated for fullstate representation +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_generator_populates_mpo_cores_fullstate(): + # This case tests that mpo_cores is empty before and non-empty after. + # The fullstate MPO is a single chain: one system core followed by one + # core per bath mode, so build_generator should produce n_lop*modes_per_state + 1 + # 4-D cores. + eom = _make_eom('fullstate') + assert eom.mpo_cores == [] + z_mem, z_rnd, z_rnd2 = _make_noise(eom) + eom.build_generator(z_mem, z_rnd, z_rnd2) + # Analytical: fullstate MPO has 1 state core + n_total_modes mode cores + n_modes = sum(eom.wavefunction.M1_modes_per_state) + expected_n_cores = 1 + n_modes + assert len(eom.mpo_cores) == expected_n_cores, ( + f'Expected {expected_n_cores} MPO cores, got {len(eom.mpo_cores)}' + ) + # Each core is a 4-index MPO tensor (bL, d_out, d_in, bR) + assert eom.mpo_cores[0].ndim == 4 + + +# ------------------------------------------------------------ +# TEST: MPO cores populated for statenumber representation +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_generator_populates_mpo_cores_statenumber(): + # This case tests that mpo_cores is populated for statenumber method. + # The statenumber MPO interleaves state and mode cores, so build_generator + # produces n_state * (modes_per_state + 1) 4-D cores after adding the + # hierarchy and Hamiltonian MPOs together. + eom = _make_eom('number') + z_mem, z_rnd, z_rnd2 = _make_noise(eom) + eom.build_generator(z_mem, z_rnd, z_rnd2) + # Analytical: statenumber MPO has n_state + sum(modes_per_state) cores + n_state = len(eom.wavefunction.M1_modes_per_site) + n_modes = sum(eom.wavefunction.M1_modes_per_state) + expected_n_cores = n_state + n_modes + assert len(eom.mpo_cores) == expected_n_cores, ( + f'Expected {expected_n_cores} MPO cores, got {len(eom.mpo_cores)}' + ) + # Each core is a 4-index MPO tensor (bL, d_out, d_in, bR) + assert eom.mpo_cores[0].ndim == 4 + + +# ------------------------------------------------------------ +# TEST: Deterministic output for fullstate representation +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_generator_matches_across_instances_fullstate(): + # This case tests that two EOM instances on identical state produce + # identical dz_dt, mpo_cores, and derivative outputs. + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, + } + tb = _make_tb() + + def make(): + ht = HopsTensorWavefunction( + k_max, tensor_param, {'INTEGRATOR': 'RUNGE_KUTTA'}, eom_param, + ) + ht.initialize(psi_0, tb.system) + return HopsTensorEOM(ht, tb.system, tb.mode, tb.noise_memory, tb.adaptive, eom_param) + + z_mem, z_rnd, z_rnd2 = _make_noise_from_basis(tb) + + eom1 = make() + eom2 = make() + dz1 = eom1.build_generator(z_mem, z_rnd, z_rnd2) + dz2 = eom2.build_generator(z_mem, z_rnd, z_rnd2) + + assert np.allclose(dz1, dz2) + for c1, c2 in zip(eom1.mpo_cores, eom2.mpo_cores): + assert np.allclose(c1, c2) + cores1 = eom1.derivative() + cores2 = eom2.derivative() + assert np.allclose( + np.array([c.sum() for c in cores1]), + np.array([c.sum() for c in cores2]), + ) + + +# ------------------------------------------------------------ +# TEST: Deterministic output for statenumber representation +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_generator_matches_across_instances_statenumber(): + # This case tests that two EOM instances on identical state produce + # identical dz_dt and mpo_cores for statenumber representation. + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'number', + 'BOND_DIM_MAX': 20, + } + tb = _make_tb() + + def make(): + ht = HopsTensorWavefunction( + k_max, tensor_param, {'INTEGRATOR': 'RUNGE_KUTTA'}, eom_param, + ) + ht.initialize(psi_0, tb.system) + return HopsTensorEOM(ht, tb.system, tb.mode, tb.noise_memory, tb.adaptive, eom_param) + + z_mem, z_rnd, z_rnd2 = _make_noise_from_basis(tb) + + eom1 = make() + eom2 = make() + dz1 = eom1.build_generator(z_mem, z_rnd, z_rnd2) + dz2 = eom2.build_generator(z_mem, z_rnd, z_rnd2) + + assert np.allclose(dz1, dz2) + for c1, c2 in zip(eom1.mpo_cores, eom2.mpo_cores): + assert np.allclose(c1, c2) + + +# ------------------------------------------------------------ +# TEST: Linear EOM returns dz_dt = 0 and builds valid MPO +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_generator_linear_eom(): + # CASE: With EQUATION_OF_MOTION = LINEAR, build_generator should + # return dz_dt = 0 (no memory feedback), build a valid MPO, and + # omit z_mem from the noise field (z_hat uses only z_rnd, z_rnd2). + eom_linear_param = {'EQUATION_OF_MOTION': 'LINEAR'} + for method in ['fullstate', 'number']: + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': method, + 'BOND_DIM_MAX': 20, + } + tb = _make_tb() + ht = HopsTensorWavefunction( + k_max, tensor_param, {'INTEGRATOR': 'RUNGE_KUTTA'}, + eom_linear_param, + ) + ht.initialize(psi_0, tb.system) + eom = HopsTensorEOM( + ht, tb.system, tb.mode, tb.noise_memory, tb.adaptive, + eom_linear_param, + ) + assert eom.flag_linear is True + z_mem, z_rnd, z_rnd2 = _make_noise(eom) + dz_dt = eom.build_generator(z_mem, z_rnd, z_rnd2) + # dz_dt must be exactly zero for linear EOM + np.testing.assert_array_equal( + dz_dt, np.zeros_like(z_mem), + err_msg=f'dz_dt should be zero for LINEAR EOM ({method})', + ) + # MPO should still be populated + assert len(eom.mpo_cores) > 0, ( + f'mpo_cores should be populated for LINEAR EOM ({method})' + ) + + +# ------------------------------------------------------------ +# TEST: dz_dt matches by-hand formula +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_generator_dz_dt_values(): + # CASE: Verify dz_dt values against the analytical formula + # dz/dt[i] = * conj(g) - conj(w) * z_mem[i] + # for a non-adaptive, non-linear EOM. + eom = _make_eom('fullstate') + z_mem, z_rnd, z_rnd2 = _make_noise(eom) + dz_dt = eom.build_generator(z_mem, z_rnd, z_rnd2) + # Compute expected dz_dt by hand using the same formula + psi = eom.wavefunction.psi + list_L2_coo = eom.mode.list_L2_coo + list_expect_L2 = [ + operator_expectation(list_L2_coo[idx], psi) + for idx in range(len(list_L2_coo)) + ] + dz_expected = calc_delta_zmem( + z_mem, + list_expect_L2, + eom.noise_memory.list_zmemg_abs, + eom.noise_memory.list_zmemw_abs, + eom.mode.list_index_L2_by_hmode, + eom.mode.list_modeidx_abs, + eom.noise_memory.list_zmemmodeidx_abs, + eom.mode.list_l2idx_abs, + eom.system.list_activel2idx_abs, + ) + np.testing.assert_allclose( + dz_dt, dz_expected, atol=1e-12, + err_msg='dz_dt does not match by-hand calc_delta_zmem formula', + ) + + +# ============================================================ +# TEST SUITE: derivative() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Returns list of cores with correct length +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_derivative_returns_cores_list(): + # This case tests that derivative() returns a list of MPS cores + # with the same length as list_cores_phi. + eom = _make_eom() + z_mem, z_rnd, z_rnd2 = _make_noise(eom) + eom.build_generator(z_mem, z_rnd, z_rnd2) + cores = eom.derivative() + assert isinstance(cores, list) + assert len(cores) == len(eom.wavefunction.list_cores_phi) + # Cross-validation: derivative from fullstate and statenumber + # should produce the same physical state (psi component) + eom_sn = _make_eom('number') + z_mem_sn, z_rnd_sn, z_rnd2_sn = _make_noise(eom_sn) + eom_sn.build_generator(z_mem_sn, z_rnd_sn, z_rnd2_sn) + cores_sn = eom_sn.derivative() + psi_deriv = extract_psi( + cores, eom.wavefunction.method, eom.wavefunction.M1_modes_per_state, + ) + psi_deriv_sn = extract_psi( + cores_sn, + eom_sn.wavefunction.method, + eom_sn.wavefunction.M1_modes_per_state, + ) + np.testing.assert_allclose( + psi_deriv, psi_deriv_sn, atol=1e-8, + err_msg='Derivative psi should agree between representations', + ) + + +# ------------------------------------------------------------ +# TEST: No side effects on list_cores_phi +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_derivative_no_side_effects_on_list_cores_phi(): + # This case tests that derivative() does not mutate wavefunction.list_cores_phi. + # Note: only covers the non-TDVP path. In the TDVP path, build_generator + # itself recenters list_cores_phi (documented side effect), but derivative() + # remains pure regardless. + eom = _make_eom() + z_mem, z_rnd, z_rnd2 = _make_noise(eom) + eom.build_generator(z_mem, z_rnd, z_rnd2) + snapshot = [c.copy() for c in eom.wavefunction.list_cores_phi] + eom.derivative() + for before, after in zip(snapshot, eom.wavefunction.list_cores_phi): + assert np.allclose(before, after) + + +# ============================================================ +# TEST SUITE: _construct_MPO() — statenumber round-trip +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: _construct_MPO fuse-add-unfuse round-trip produces rank-4 cores +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_construct_mpo_statenumber_roundtrip(): + # This case tests that _construct_MPO for statenumber representation + # produces rank-4 MPO cores whose physical dimensions are consistent + # with the original unfused cores (d_out == d_in == sqrt(fused_dim)). + eom = _make_eom('number') + z_mem, z_rnd, z_rnd2 = _make_noise(eom) + eom.build_generator(z_mem, z_rnd, z_rnd2) + for i, core in enumerate(eom.mpo_cores): + assert core.ndim == 4, f'Core {i} should be rank-4 after unfuse' + w_l, d_out, d_in, w_r = core.shape + assert d_out == d_in, ( + f'Core {i} physical dims mismatch: d_out={d_out}, d_in={d_in}' + ) + # Also verify _list_cores_op and _list_cores_ham were unfused back to rank-4 + for i, core in enumerate(eom._list_cores_op): + assert core.ndim == 4, f'_list_cores_op[{i}] should be rank-4 after unfuse' + for i, core in enumerate(eom._list_cores_ham): + assert core.ndim == 4, f'_list_cores_ham[{i}] should be rank-4 after unfuse' + + +# ============================================================ +# TEST SUITE: _build_mpo_parts() — error paths +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Invalid method raises UnsupportedRequest +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_mpo_parts_invalid_method_raises(): + # This case tests that _build_mpo_parts raises UnsupportedRequest when + # wavefunction.method is not a recognized representation. + eom = _make_eom('fullstate') + # Monkey-patch the method to an invalid value after setup, + # then call _build_mpo_parts directly (build_generator fails earlier + # at phi_0 extraction before reaching the method dispatch). + eom.wavefunction.method = 'bogus' + with pytest.raises(UnsupportedRequest): + eom._build_mpo_parts( + np.zeros(eom.system.size, dtype=np.complex128), + np.zeros(eom.system.size, dtype=np.complex128), + 0.0, + ) + + +# ------------------------------------------------------------ +# TEST: Invalid normalization raises UnsupportedRequest +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_mpo_parts_invalid_normalization_raises(): + # This case tests that _build_mpo_parts raises UnsupportedRequest + # when an unsupported normalization is passed. + eom = _make_eom('fullstate') + z_mem, z_rnd, z_rnd2 = _make_noise(eom) + # Set invalid normalization on the instance and verify _build_mpo_parts raises + eom.normalization = 'unknown' + with pytest.raises(UnsupportedRequest): + eom._build_mpo_parts( + np.zeros(eom.system.size, dtype=np.complex128), + np.zeros(eom.system.size, dtype=np.complex128), + 0.0, + ) + + +# ------------------------------------------------------------ +# TEST: nearest_neighbors_ham branch uses nn MPO builder +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_mpo_parts_nn_ham_branch(): + # This case tests that _build_mpo_parts takes the + # nearest_neighbors_ham=True branch for the dimer-of-dimers + # Hamiltonian (tridiagonal). The nn MPO has bond dimension 4 + # for the Hamiltonian cores, while the general MPO has + # bond dimension 4 + 2*(n_state - 2). We verify that the + # nn branch produces a smaller Hamiltonian bond dimension. + # The separate Hamiltonian MPO under test is built only without the + # topology optimization, which otherwise absorbs it into one generator. + eom_nn = _make_eom('number', flag_mpo_optimize=False) + assert eom_nn.system.flag_nearest_neighbor_ham is True + z_conj = (0.01 + 0.02j) * np.arange( + eom_nn.mode.n_l2, dtype=np.complex128, + ) + L_conj_avg = np.zeros(eom_nn.mode.n_l2, dtype=np.complex128) + eom_nn._build_mpo_parts(z_conj, L_conj_avg, 0.0) + nn_ham_bond = max(c.shape[-1] for c in eom_nn._list_cores_ham) + # Force general path by patching nearest_neighbors_ham + eom_gen = _make_eom('number', flag_mpo_optimize=False) + eom_gen.system.flag_nearest_neighbor_ham = False + eom_gen.mpo_builder.flag_nearest_neighbor_ham = False + eom_gen._build_mpo_parts(z_conj, L_conj_avg, 0.0) + gen_ham_bond = max(c.shape[-1] for c in eom_gen._list_cores_ham) + # nn bond dim (4) should be strictly less than general (4 + 2*(n-2)) + assert nn_ham_bond < gen_ham_bond, ( + f'nn ham bond {nn_ham_bond} should be < general {gen_ham_bond}' + ) + + +# ============================================================ +# TEST SUITE: build_generator() — flag_norm=False +# ============================================================ + + +def _make_eom_unnormalized(method='fullstate'): + """Create a HopsTensorEOM with unnormalized (Linear) EOM.""" + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': method, + 'BOND_DIM_MAX': 20, + } + eom_param_linear = {'EQUATION_OF_MOTION': 'LINEAR'} + tb = _make_tb() + ht = HopsTensorWavefunction( + k_max, tensor_param, {'INTEGRATOR': 'RUNGE_KUTTA'}, eom_param_linear, + ) + ht.initialize(psi_0, tb.system) + return HopsTensorEOM(ht, tb.system, tb.mode, tb.noise_memory, tb.adaptive, eom_param_linear) + + +# ------------------------------------------------------------ +# TEST: flag_norm=False sets norm_corr=0 and produces valid MPO +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_generator_flag_norm_false_runs(): + # This case tests that build_generator completes without error + # when flag_norm is False (Linear EOM), producing valid MPO cores. + eom = _make_eom_unnormalized('fullstate') + assert not eom.wavefunction.flag_norm, 'flag_norm should be False' + z_mem, z_rnd, z_rnd2 = _make_noise(eom) + dz_dt = eom.build_generator(z_mem, z_rnd, z_rnd2) + # MPO cores should be populated + assert len(eom.mpo_cores) > 0 + # dz_dt should have the right shape + assert dz_dt.shape == z_mem.shape + # Differential: unnormalized MPO cores should differ from normalized. + # Note: dz_dt (the z_mem derivative) is identical between representations + # because norm_corr only enters the MPO, not the memory term. + eom_norm = _make_eom('fullstate') + z_mem_norm, z_rnd_norm, z_rnd2_norm = _make_noise(eom_norm) + eom_norm.build_generator(z_mem_norm, z_rnd_norm, z_rnd2_norm) + any_core_differs = any( + not np.allclose(c1, c2) + for c1, c2 in zip(eom.mpo_cores, eom_norm.mpo_cores) + ) + assert any_core_differs, ( + 'Unnormalized and normalized MPO cores should differ when norm_corr != 0' + ) + + +# ------------------------------------------------------------ +# TEST: Unnormalized MPO differs from normalized MPO +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_generator_norm_vs_unnorm_differ(): + # This case tests that the norm correction term actually changes + # the MPO. With identical inputs, normalized and unnormalized + # EOMs should produce different MPO cores. + eom_norm = _make_eom('fullstate') + eom_unnorm = _make_eom_unnormalized('fullstate') + z_mem_n, z_rnd_n, z_rnd2_n = _make_noise(eom_norm) + z_mem_u, z_rnd_u, z_rnd2_u = _make_noise(eom_unnorm) + eom_norm.build_generator(z_mem_n, z_rnd_n, z_rnd2_n) + eom_unnorm.build_generator(z_mem_u, z_rnd_u, z_rnd2_u) + # At least one core should differ (norm_corr != 0 changes the MPO) + any_differ = any( + not np.allclose(c1, c2) + for c1, c2 in zip(eom_norm.mpo_cores, eom_unnorm.mpo_cores) + ) + assert any_differ, ( + 'Normalized and unnormalized MPOs should differ when norm correction is nonzero' + ) + + +# ============================================================ +# TEST SUITE: _build_mpo_parts() — Hamiltonian MPO structure +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Hamiltonian MPO has correct number of cores +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_mpo_parts_ham_structure(): + # This case tests that the Hamiltonian MPO has the correct number + # of cores: n_state site cores + n_state * modes_per_state mode cores. + # A separate Hamiltonian MPO exists only without the optimization. + eom = _make_eom('number', flag_mpo_optimize=False) + z_conj_t = np.zeros(nsite, dtype=np.complex128) + L_conj_avg = np.zeros(nsite, dtype=np.complex128) + eom._build_mpo_parts(z_conj_t, L_conj_avg, 0.0) + V1_mps = eom.wavefunction.M1_modes_per_state + expected = nsite + int(np.sum(V1_mps)) + assert len(eom._list_cores_ham) == expected + + +# ------------------------------------------------------------ +# TEST: Hamiltonian MPO site cores have correct physical dimensions +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_mpo_parts_ham_site_dims(): + # This case tests that site cores in the Hamiltonian MPO have + # physical dimension 2x2 (binary occupation per site). + # A separate Hamiltonian MPO exists only without the optimization. + eom = _make_eom('number', flag_mpo_optimize=False) + z_conj_t = np.zeros(nsite, dtype=np.complex128) + L_conj_avg = np.zeros(nsite, dtype=np.complex128) + eom._build_mpo_parts(z_conj_t, L_conj_avg, 0.0) + offsets = _statenumber_offsets(eom.wavefunction.M1_modes_per_state) + for site in range(nsite): + idx = offsets[site] + core = eom._list_cores_ham[idx] + assert core.shape[1] == 2 and core.shape[2] == 2, ( + f'Site core {site} has shape {core.shape}, expected (*, 2, 2, *)' + ) + + +# ------------------------------------------------------------ +# TEST: Fullstate _list_cores_ham is empty after build +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_mpo_parts_fullstate_list_cores_ham_empty(): + # CASE: For fullstate representation, the Hamiltonian is embedded + # directly into _list_cores_op by build_fullstate_mpo. _list_cores_ham should + # remain an empty list. + eom = _make_eom('fullstate') + z_conj_t = np.zeros(nsite, dtype=np.complex128) + L_conj_avg = np.zeros(nsite, dtype=np.complex128) + eom._build_mpo_parts(z_conj_t, L_conj_avg, 0.0) + assert eom._list_cores_ham == [], ( + f'_list_cores_ham should be empty for fullstate, got {len(eom._list_cores_ham)} cores' + ) + assert len(eom._list_cores_op) > 0, '_list_cores_op should be populated' + + +# ------------------------------------------------------------ +# TEST: Hamiltonian MPO mode cores are identity +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_mpo_parts_ham_mode_cores_identity(): + # This case tests that mode cores in the Hamiltonian MPO act as + # identity operators (the Hamiltonian doesn't couple to hierarchy modes). + # A separate Hamiltonian MPO exists only without the optimization. + eom = _make_eom('number', flag_mpo_optimize=False) + z_conj_t = np.zeros(nsite, dtype=np.complex128) + L_conj_avg = np.zeros(nsite, dtype=np.complex128) + eom._build_mpo_parts(z_conj_t, L_conj_avg, 0.0) + V1_mps = eom.wavefunction.M1_modes_per_state + offsets = _statenumber_offsets(V1_mps) + for site in range(nsite): + idx_base = offsets[site] + for m in range(V1_mps[site]): + idx = idx_base + 1 + m + core = eom._list_cores_ham[idx] + # For each bond index pair (i, i), the mode core should be identity + for bond_idx in range(core.shape[0]): + np.testing.assert_allclose( + core[bond_idx, :, :, bond_idx], + np.eye(k_max + 1), + atol=1e-12, + err_msg=f'Mode core {idx} not identity at bond index {bond_idx}', + ) + + +# ============================================================ +# TEST SUITE: __init__() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Constructor stores references and builds MpoBuilder +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_init_stores_references(): + eom = _make_eom() + # This case tests that constructor stores all required references + # with correct types + assert isinstance(eom.wavefunction, HopsTensorWavefunction) + assert isinstance(eom.system, HopsSystem) + assert isinstance(eom.mode, HopsModes) + assert isinstance(eom.noise_memory, HopsNoiseMemory) + assert eom.adaptive is not None + # This case tests that normalization defaults to 'homps' + assert eom.normalization == 'homps' + # This case tests that MpoBuilder is constructed + assert eom.mpo_builder is not None + # This case tests that MPO storage starts empty + assert eom.mpo_cores == [] + assert eom._list_cores_op == [] + assert eom._list_cores_ham == [] + # This case tests that flag_linear is derived from eom_param + assert eom.flag_linear is False + + +# ------------------------------------------------------------ +# TEST: flag_linear is True for LINEAR EOM +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_init_flag_linear_true(): + # CASE: When EQUATION_OF_MOTION is LINEAR, flag_linear should be True. + tb = _make_tb() + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, + } + eom_linear = {'EQUATION_OF_MOTION': 'LINEAR'} + ht = HopsTensorWavefunction( + k_max, tensor_param, {'INTEGRATOR': 'RUNGE_KUTTA'}, eom_linear, + ) + ht.initialize(psi_0, tb.system) + eom = HopsTensorEOM( + ht, tb.system, tb.mode, tb.noise_memory, tb.adaptive, eom_linear, + ) + assert eom.flag_linear is True + + +# ============================================================ +# TEST SUITE: refresh_builder() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: refresh_builder updates builder state count +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_refresh_builder_updates_state_count(): + eom = _make_eom() + original_n = eom.mpo_builder.n_state + eom.system.state_list = [0, 2] + eom.refresh_builder() + # Analytical: n_state should match the new state list length + assert eom.mpo_builder.n_state == 2 + assert eom.mpo_builder.n_state != original_n + + +# ============================================================ +# TEST SUITE: _build_mpo_parts() — hierarchy MPO structure +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Fullstate hierarchy MPO has correct number of cores +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_mpo_parts_fullstate_n_cores(): + # This case tests that the fullstate hierarchy MPO has 1 system core + # + n_modes mode cores. + eom = _make_eom('fullstate') + z_conj_t = np.zeros(nsite, dtype=np.complex128) + L_conj_avg = np.zeros(nsite, dtype=np.complex128) + eom._build_mpo_parts(z_conj_t, L_conj_avg, 0.0) + n_total_modes = int(np.sum(eom.wavefunction.M1_modes_per_state)) + assert len(eom._list_cores_op) == 1 + n_total_modes + + +# ------------------------------------------------------------ +# TEST: Statenumber hierarchy MPO has correct number of cores +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_mpo_parts_statenumber_n_cores(): + # This case tests that the statenumber hierarchy MPO has + # n_state * (1 + modes_per_state) cores. + eom = _make_eom('number') + z_conj_t = np.zeros(nsite, dtype=np.complex128) + L_conj_avg = np.zeros(nsite, dtype=np.complex128) + eom._build_mpo_parts(z_conj_t, L_conj_avg, 0.0) + expected = nsite + int(np.sum(eom.wavefunction.M1_modes_per_state)) + assert len(eom._list_cores_op) == expected + + +# ------------------------------------------------------------ +# TEST: Fullstate system core has correct bond dimension +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_mpo_parts_fullstate_bondsize(): + # This case tests that the fullstate system core has the expected + # MPO bond dimension of n_lop_full + 2. + eom = _make_eom('fullstate') + z_conj_t = np.zeros(nsite, dtype=np.complex128) + L_conj_avg = np.zeros(nsite, dtype=np.complex128) + eom._build_mpo_parts(z_conj_t, L_conj_avg, 0.0) + expected_bond = eom.mode.n_l2 + 2 + assert eom._list_cores_op[0].shape[3] == expected_bond, ( + f'Fullstate system core bond dim = {eom._list_cores_op[0].shape[3]}, ' + f'expected {expected_bond}' + ) + + +# ------------------------------------------------------------ +# TEST: Statenumber site cores have bond dimension 5 +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_mpo_parts_statenumber_site_bond(): + # This case tests that statenumber site cores have the fixed + # MPO bond dimension of 5 (except first core left bond = 1). This is the + # hierarchy-only MPO, so the optimization is off. + eom = _make_eom('number', flag_mpo_optimize=False) + z_conj_t = np.zeros(nsite, dtype=np.complex128) + L_conj_avg = np.zeros(nsite, dtype=np.complex128) + eom._build_mpo_parts(z_conj_t, L_conj_avg, 0.0) + offsets = _statenumber_offsets(eom.wavefunction.M1_modes_per_state) + for site in range(nsite): + idx = offsets[site] + core = eom._list_cores_op[idx] + if site == 0: + assert core.shape[0] == 1, ( + f'First site core left bond should be 1, got {core.shape[0]}' + ) + else: + assert core.shape[0] == 5 + if site == nsite - 1: + # Last site group: final core (last mode core) should have Dr=1 + last_idx = offsets[site] + eom.wavefunction.M1_modes_per_state[site] + last_core = eom._list_cores_op[last_idx] + assert last_core.shape[3] == 1, ( + f'Last core right bond should be 1, got {last_core.shape[3]}' + ) + else: + assert core.shape[3] == 5 + + +# ------------------------------------------------------------ +# TEST: Statenumber hierarchy site cores have correct physical dims +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_mpo_parts_statenumber_site_phys_dims(): + # This case tests that statenumber site cores have physical + # dimension 2 (binary occupied/unoccupied) and mode cores have + # physical dimension k_max + 1. + eom = _make_eom('number') + z_conj_t = np.zeros(nsite, dtype=np.complex128) + L_conj_avg = np.zeros(nsite, dtype=np.complex128) + eom._build_mpo_parts(z_conj_t, L_conj_avg, 0.0) + offsets = _statenumber_offsets(eom.wavefunction.M1_modes_per_state) + d_mode = k_max + 1 + for site in range(nsite): + idx_site = offsets[site] + core_site = eom._list_cores_op[idx_site] + assert core_site.shape[1] == 2, ( + f'Site {site} core d_out should be 2, got {core_site.shape[1]}' + ) + assert core_site.shape[2] == 2, ( + f'Site {site} core d_in should be 2, got {core_site.shape[2]}' + ) + for m in range(eom.wavefunction.M1_modes_per_state[site]): + idx_mode = idx_site + 1 + m + core_mode = eom._list_cores_op[idx_mode] + assert core_mode.shape[1] == d_mode, ( + f'Mode core {idx_mode} d_out should be {d_mode}, ' + f'got {core_mode.shape[1]}' + ) + assert core_mode.shape[2] == d_mode, ( + f'Mode core {idx_mode} d_in should be {d_mode}, ' + f'got {core_mode.shape[2]}' + ) + + +# ============================================================ +# TEST SUITE: _construct_MPO() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Fullstate _construct_MPO assigns cores_op directly +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_construct_mpo_fullstate_direct(): + # This case tests that for fullstate, _construct_MPO sets + # mpo_cores = _list_cores_op (no tensor addition or compression). + eom = _make_eom('fullstate') + z_conj_t = np.zeros(nsite, dtype=np.complex128) + L_conj_avg = np.zeros(nsite, dtype=np.complex128) + eom._construct_MPO(z_conj_t, L_conj_avg, 0.0) + assert eom.mpo_cores is eom._list_cores_op + + +# ------------------------------------------------------------ +# TEST: MPO bond dim is not capped by tight MPS bond_dim_max +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_construct_mpo_statenumber_not_capped_by_mps_bond_dim(): + # This case tests that _construct_MPO uses its own bond_dim_max + # (sum of hierarchy + Hamiltonian bond dims) rather than + # wavefunction.bond_dim_max. With a tight MPS cap of 5, the old code + # would silently truncate the 9-dim combined MPO; the new code should + # preserve the full MPO regardless. + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'number', + 'BOND_DIM_MAX': 5, # tight MPS cap — must NOT affect MPO compression + 'EOM': 'Normalized_Nonlinear', + # The two MPOs whose sum is compressed here are built only without + # the topology optimization. + 'FLAG_MPO_OPTIMIZE': False, + } + tb = _make_tb() + ht = HopsTensorWavefunction( + k_max, tensor_param, {'INTEGRATOR': 'RUNGE_KUTTA'}, eom_param, + ) + ht.initialize(psi_0, tb.system) + eom = HopsTensorEOM(ht, tb.system, tb.mode, tb.noise_memory, tb.adaptive, eom_param) + z_mem, z_rnd, z_rnd2 = _make_noise(eom) + eom.build_generator(z_mem, z_rnd, z_rnd2) + # Combined MPO has dim 9 (hierarchy 5 + Ham NN 4); MPS cap of 5 + # must not truncate it below this. + max_bond = max(c.shape[-1] for c in eom.mpo_cores[:-1]) + assert max_bond > 5, ( + f'MPO bond dim = {max_bond} was capped by MPS bond_dim_max=5' + ) + + +# ------------------------------------------------------------ +# TEST: Compressed MPO is numerically equivalent to uncompressed sum +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_construct_mpo_statenumber_exact_compression_preserves_operator(): + # This case tests that exact MPO compression (epsilon=0) produces the + # same wavefunction derivative as the uncompressed block-concatenated + # sum of hierarchy + Hamiltonian MPOs. The compressed and uncompressed + # operators represent the same linear map. Both MPOs exist only without + # the topology optimization. + eom = _make_eom('number', flag_mpo_optimize=False) + z_mem, z_rnd, z_rnd2 = _make_noise(eom) + eom.build_generator(z_mem, z_rnd, z_rnd2) + deriv_compressed, _ = tensor_matvec_prod( + eom.wavefunction.flat_cores, + eom.mpo_cores, + eom.wavefunction.mps_epsilon, + eom.wavefunction.bond_dim_max, + ) + + # Build the uncompressed block-concatenated sum (no SVD truncation) + # by calling _build_mpo_parts and fusing/adding without compression. + eom2 = _make_eom('number', flag_mpo_optimize=False) + eom2.build_generator(z_mem, z_rnd, z_rnd2) + # Fuse physical dims to rank-3, form uncompressed block sum, unfuse + cores_op = [c.reshape(c.shape[0], c.shape[1]*c.shape[2], c.shape[3]) + for c in eom2._list_cores_op] + cores_ham = [c.reshape(c.shape[0], c.shape[1]*c.shape[2], c.shape[3]) + for c in eom2._list_cores_ham] + bond_max = (max(c.shape[-1] for c in cores_op) + + max(c.shape[-1] for c in cores_ham)) + cores_uncompressed = tensor_add(cores_op, cores_ham, 0, bond_max) + cores_uncompressed_r4 = [] + for core in cores_uncompressed: + phys = int(round(np.sqrt(core.shape[1]))) + cores_uncompressed_r4.append( + core.reshape(core.shape[0], phys, phys, core.shape[2]) + ) + deriv_uncompressed, _ = tensor_matvec_prod( + eom2.wavefunction.flat_cores, + cores_uncompressed_r4, + eom2.wavefunction.mps_epsilon, + eom2.wavefunction.bond_dim_max, + ) + + for i, (dc, du) in enumerate(zip(deriv_compressed, deriv_uncompressed)): + np.testing.assert_allclose( + dc, du, atol=1e-10, + err_msg=f'Derivative core {i}: compressed vs uncompressed mismatch', + ) + + +# ============================================================ +# TEST SUITE: Linear EOM (EQUATION_OF_MOTION = 'LINEAR') +# ============================================================ + +eom_param_linear = {'EQUATION_OF_MOTION': 'LINEAR'} + + +def _make_eom_linear(method='fullstate'): + """Creates an initialized HopsTensorEOM with the linear EOM.""" + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': method, + 'BOND_DIM_MAX': 20, + } + tb = _make_tb() + ht = HopsTensorWavefunction( + k_max, tensor_param, {'INTEGRATOR': 'RUNGE_KUTTA'}, eom_param_linear, + ) + ht.initialize(psi_0, tb.system) + return HopsTensorEOM(ht, tb.system, tb.mode, tb.noise_memory, tb.adaptive, + eom_param_linear) + + +# ------------------------------------------------------------ +# TEST: MPO is independent of z_mem for the linear EOM +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_linear_mpo_independent_of_zmem(): + # This case tests that the linear EOM ignores z_mem when building the MPO, + # consistent with the scalar EOM where z_hat = conj(z_rnd) only. + for method in ['fullstate', 'number']: + eom = _make_eom_linear(method) + z_mem, z_rnd, z_rnd2 = _make_noise(eom) + eom.build_generator(np.zeros_like(z_mem), z_rnd, z_rnd2) + mpo_zero_zmem = [c.copy() for c in eom.mpo_cores] + eom.build_generator(z_mem, z_rnd, z_rnd2) + for i, (c_zero, c_nonzero) in enumerate(zip(mpo_zero_zmem, eom.mpo_cores)): + np.testing.assert_allclose( + c_zero, c_nonzero, atol=1e-14, + err_msg=( + f'Linear EOM: MPO core {i} changed with z_mem ' + f'for {method} — z_mem must be suppressed' + ), + ) + + +# ============================================================ +# LTC fixture +# ============================================================ + +# sys_param with L_LT_CORR / PARAM_LT_CORR populated. The LTC +# L-operators reuse the site-diagonal projectors from L_HIER so +# that system_functions maps them to the same unique L2 indices. +sys_param_ltc = { + 'HAMILTONIAN': np.array(hs, dtype=np.complex128), + 'GW_SYSBATH': gw_sysbath, + 'L_HIER': lop_list, + 'L_NOISE1': lop_list, + 'ALPHA_NOISE1': bcf_exp, + 'PARAM_NOISE1': gw_sysbath, + 'L_LT_CORR': [loperator[i] for i in range(nsite)], + 'PARAM_LT_CORR': [250.0 / 1000.0] * nsite, +} + + +def _make_eom_ltc(method='fullstate', + eom_p=eom_param): + """Creates a HopsTensorEOM with LTC parameters populated.""" + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': method, + 'BOND_DIM_MAX': 20, + } + tb = _make_tb(sp=sys_param_ltc) + ht = HopsTensorWavefunction( + k_max, tensor_param, {'INTEGRATOR': 'RUNGE_KUTTA'}, eom_p, + ) + ht.initialize(psi_0, tb.system) + return HopsTensorEOM( + ht, tb.system, tb.mode, tb.noise_memory, tb.adaptive, eom_p, + ) + + +# ============================================================ +# TEST SUITE: build_generator() — LTC +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: build_generator sets _has_lt_corr = True (nonlinear) +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_generator_ltc_sets_flag_nonlinear(): + # This case tests that _has_lt_corr is True after build_generator + # when the system has nonzero LTC parameters (nonlinear EOM). + for method in ['fullstate', 'number']: + eom = _make_eom_ltc(method) + z_mem, z_rnd, z_rnd2 = _make_noise(eom) + eom.build_generator(z_mem, z_rnd, z_rnd2) + assert eom._has_lt_corr is True, ( + f'{method}: _has_lt_corr should be True with LTC params' + ) + + +# ------------------------------------------------------------ +# TEST: build_generator sets _has_lt_corr = True (linear) +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_generator_ltc_sets_flag_linear(): + # This case tests that _has_lt_corr is True after build_generator + # when the system has nonzero LTC parameters (linear EOM). + linear_eom_param = {'EQUATION_OF_MOTION': 'LINEAR'} + for method in ['fullstate', 'number']: + eom = _make_eom_ltc(method, eom_p=linear_eom_param) + z_mem, z_rnd, z_rnd2 = _make_noise(eom) + eom.build_generator(z_mem, z_rnd, z_rnd2) + assert eom._has_lt_corr is True, ( + f'{method}: _has_lt_corr should be True with LTC params (linear)' + ) + + +# ------------------------------------------------------------ +# TEST: no-LTC fixture has _has_lt_corr = False +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_generator_no_ltc_flag_false(): + # This case tests that _has_lt_corr remains False when the + # system has no LTC parameters. + for method in ['fullstate', 'number']: + eom = _make_eom(method) + z_mem, z_rnd, z_rnd2 = _make_noise(eom) + eom.build_generator(z_mem, z_rnd, z_rnd2) + assert eom._has_lt_corr is False, ( + f'{method}: _has_lt_corr should be False without LTC params' + ) + + +# ------------------------------------------------------------ +# TEST: LTC produces different mpo_cores than no-LTC (nonlinear) +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_generator_ltc_changes_mpo_nonlinear(): + # This case tests that the MPO cores differ between LTC and + # no-LTC fixtures. The LTC norm correction alters the state + # core of the MPO via the Hamiltonian channel. + for method in ['fullstate', 'number']: + eom_no_ltc = _make_eom(method) + z_mem, z_rnd, z_rnd2 = _make_noise(eom_no_ltc) + eom_no_ltc.build_generator(z_mem, z_rnd, z_rnd2) + cores_no_ltc = [c.copy() for c in eom_no_ltc.mpo_cores] + + eom_ltc = _make_eom_ltc(method) + z_mem_ltc, z_rnd_ltc, z_rnd2_ltc = _make_noise(eom_ltc) + eom_ltc.build_generator(z_mem_ltc, z_rnd_ltc, z_rnd2_ltc) + cores_ltc = eom_ltc.mpo_cores + + # At least one core must differ (shape or values) + any_diff = ( + len(cores_no_ltc) != len(cores_ltc) + or any( + c1.shape != c2.shape or not np.allclose(c1, c2, atol=1e-14) + for c1, c2 in zip(cores_no_ltc, cores_ltc) + ) + ) + assert any_diff, ( + f'{method}: MPO cores should differ between LTC and no-LTC' + ) diff --git a/tests/test_hops_tensor_trajectory.py b/tests/test_hops_tensor_trajectory.py new file mode 100644 index 0000000..4debe1d --- /dev/null +++ b/tests/test_hops_tensor_trajectory.py @@ -0,0 +1,1966 @@ +import numpy as np +import pytest +import scipy as sp + +from mesohops.trajectory.exp_noise import bcf_exp +from mesohops.trajectory.hops_tensor_trajectory import HopsTensorTrajectory +from mesohops.trajectory.hops_trajectory import HopsTrajectory as HOPS +from mesohops.util.bath_corr_functions import bcf_convert_dl_to_exp +from mesohops.util.tensor_operations import extract_gs_amp +from mesohops.util.exceptions import ( + LockedException, + TrajectoryError, + UnsupportedRequest, +) + +__title__ = 'Test Tensor HOPS vs Vector HOPS' +__author__ = 'N. Covalsen' +__maintainer__ = 'N. Covalsen' + +# ============================================================ +# Dimer of Dimers System Setup +# ============================================================ +# Identical to test_dimer_of_dimers.py: 4 sites, 2 modes per site +# (Drude-Lorentz + LTC correction at 500 cm^-1) +noise_param = { + 'SEED': 0, + 'MODEL': 'FFT_FILTER', + 'TLEN': 25000.0, # Units: fs + 'TAU': 1.0, # Units: fs + 'STORE_RAW_NOISE': True, + 'RAND_MODEL': 'BOX_MULLER', +} + +nsite = 4 +e_lambda = 20.0 +gamma = 50.0 +temp = 140.0 +(g_0, w_0) = bcf_convert_dl_to_exp(e_lambda, gamma, temp) + +T3_loperator = np.zeros([4, 4, 4], dtype=np.float64) +list_gw_sysbath = [] +list_lop = [] +for i in range(nsite): + T3_loperator[i, i, i] = 1.0 + list_gw_sysbath.append([g_0, w_0]) + list_lop.append(sp.sparse.coo_matrix(T3_loperator[i])) + list_gw_sysbath.append([-1j * np.imag(g_0), 500.0]) + list_lop.append(T3_loperator[i]) + +H2_hamiltonian = np.zeros([nsite, nsite]) +H2_hamiltonian[0, 1] = 40 +H2_hamiltonian[1, 0] = 40 +H2_hamiltonian[1, 2] = 10 +H2_hamiltonian[2, 1] = 10 +H2_hamiltonian[2, 3] = 40 +H2_hamiltonian[3, 2] = 40 + +sys_param = { + 'HAMILTONIAN': np.array(H2_hamiltonian, dtype=np.complex128), + 'GW_SYSBATH': list_gw_sysbath, + 'L_HIER': list_lop, + 'L_NOISE1': list_lop, + 'ALPHA_NOISE1': bcf_exp, + 'PARAM_NOISE1': list_gw_sysbath, +} + +eom_param = {'EQUATION_OF_MOTION': 'NORMALIZED NONLINEAR'} +integrator_param = {'INTEGRATOR': 'RUNGE_KUTTA'} +hier_param = {'MAXHIER': 2, 'TRUNCATION_METHOD': 'rectangular'} + +psi_0 = np.array([0.0] * nsite, dtype=np.complex128) +psi_0[2] = 1.0 +psi_0 = psi_0 / np.linalg.norm(psi_0) + +t_max = 200.0 +t_step = 4.0 + + +# ============================================================ +# Helper: run vector HOPS +# ============================================================ +def _run_vector_hops(): + """Runs standard vector HOPS and returns psi_traj as array.""" + hops = HOPS( + sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integrator_param, + ) + hops.initialize(psi_0) + hops.propagate(t_max, t_step) + return np.array(hops.storage.data['psi_traj']) + + +# ============================================================ +# Helper: run tensor HOPS +# ============================================================ +def _run_tensor_hops(method, bond_dim_max=20, mps_epsilon=1e-10): + """Runs tensor HOPS and returns psi_traj as array.""" + tensor_param = { + 'MPS_EPSILON': mps_epsilon, + 'METHOD': method, + 'BOND_DIM_MAX': bond_dim_max, + } + traj = HopsTensorTrajectory( + system_param=sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integrator_param, + tensor_param=tensor_param, + ) + traj.initialize(psi_0) + traj.propagate(t_max, t_step) + return np.array(traj.storage['psi_traj']) + + +# ============================================================ +# TEST SUITE: __init__() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: tensor_param defaults are filled when None is passed +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_init_tensor_param_defaults(): + # This case tests that passing tensor_param=None fills in all + # defaults from TENSOR_DICT_DEFAULT. + traj = HopsTensorTrajectory( + system_param=sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integrator_param, + tensor_param=None, + ) + assert traj.tensor_param['METHOD'] == 'fullstate' + assert traj.tensor_param['MPS_EPSILON'] == 1e-2 + assert traj.tensor_param['MPO_EPSILON'] == 0.0 + assert traj.tensor_param['BOND_DIM_MAX'] == 10 + + +# ------------------------------------------------------------ +# TEST: sparse Hamiltonian is coerced to dense +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_init_sparse_hamiltonian_coerced(): + # This case tests that a sparse Hamiltonian is converted to a dense + # array during construction without mutating the caller's dict. + sp_sys_param = dict(sys_param) + sp_sys_param['HAMILTONIAN'] = sp.sparse.csr_matrix( + sys_param['HAMILTONIAN'] + ) + original_ham = sp_sys_param['HAMILTONIAN'] + traj = HopsTensorTrajectory( + system_param=sp_sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integrator_param, + tensor_param={ + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, + }, + ) + # Caller's dict should not be mutated + assert sp.sparse.issparse(original_ham), ( + 'Constructor mutated the caller\'s system_param dict' + ) + + + +# ------------------------------------------------------------ +# TEST: storage functions are registered for tensor HOPS +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_init_storage_functions_registered(): + # This case tests that phi_traj and phi_norm storage functions + # are replaced with tensor-aware versions during construction. + traj = HopsTensorTrajectory( + system_param=sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integrator_param, + tensor_param={ + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, + }, + ) + # The tensor-specific functions have 'tensor' in their name + if 'phi_traj' in traj.storage.dic_save: + assert 'tensor' in traj.storage.dic_save['phi_traj'].__name__ + if 'phi_norm' in traj.storage.dic_save: + assert 'tensor' in traj.storage.dic_save['phi_norm'].__name__ + + +# ------------------------------------------------------------ +# TEST: list_aux_norm warning and removal +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_init_list_aux_norm_warning(): + # This case tests that list_aux_norm in storage triggers a warning + # and is removed, since it's not meaningful for tensor HOPS. + storage_param = {'list_aux_norm': True} + with pytest.warns(UserWarning, match='list_aux_norm'): + traj = HopsTensorTrajectory( + system_param=sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integrator_param, + storage_param=storage_param, + tensor_param={ + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, + }, + ) + assert 'list_aux_norm' not in traj.storage.dic_save + + +# ============================================================ +# TEST SUITE: initialize() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: tensor_basis.adaptive agrees with eom_param +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_initialize_nonadaptive_sets_adaptive_false(): + # This case tests that tensor_basis.adaptive is False after + # initializing with the default DELTA_S=0. + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, + } + traj = HopsTensorTrajectory( + system_param=sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integrator_param, + tensor_param=tensor_param, + ) + traj.initialize(psi_0) + assert traj.tensor_basis.adaptive is False + + +@pytest.mark.level(1) +def test_initialize_adaptive_sets_adaptive_true(): + # This case tests that tensor_basis.adaptive is True when + # constructed with DELTA_S > 0. The adaptive flag is set by + # tensor_basis.initialize() (called inside HopsTensorTrajectory + # .initialize()). We call tensor_basis.initialize() directly + # because the full initialize() path crashes in the adaptive + # define_basis step (noise index mismatch — pre-existing issue). + # Full adaptive initialization is tested in + # test_dimer_of_dimers_tensor.py. + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, + } + eom_adaptive = { + 'EQUATION_OF_MOTION': 'NORMALIZED NONLINEAR', + 'ADAPTIVE_S': True, + 'DELTA_S': 0.01, + } + traj_adap = HopsTensorTrajectory( + system_param=sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_adaptive, + integration_param={ + 'INTEGRATOR': 'RUNGE_KUTTA', + 'EARLY_INTEGRATOR_STEPS': 5, + 'INCHWORM_CAP': 5, + }, + tensor_param=tensor_param, + ) + traj_adap.tensor_basis.initialize(0.01) + assert traj_adap.tensor_basis.adaptive is True + + +# ------------------------------------------------------------ +# TEST: make_adaptive rejects list_permanent_sites +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_make_adaptive_rejects_list_permanent_sites(): + # This case tests that passing list_permanent_sites to a tensor + # trajectory's make_adaptive raises NotImplementedError. The + # parent class accepts the argument and stores it on + # system.param["list_permanent_sites"], but HopsTensorBasis does + # not read that key, so the requested sites would silently not be + # preserved in the adaptive basis. Failing fast at configuration + # time gives the caller a clear signal. + traj = HopsTensorTrajectory( + system_param=sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integrator_param, + tensor_param={ + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, + }, + ) + with pytest.raises(NotImplementedError, match='list_permanent_sites'): + traj.make_adaptive( + delta_a=1e-3, delta_s=1e-3, list_permanent_sites=[0], + ) + + +# ============================================================ +# TEST SUITE: initialize() — wavefunction encoding +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Initial wavefunction is preserved +# ------------------------------------------------------------ +@pytest.mark.level(1) +@pytest.mark.parametrize('method', [ + 'fullstate', 'number', +]) +def test_initialize_preserves_psi0(method): + # This case tests that the initial physical wavefunction is correctly + # encoded and recovered from the tensor representation. + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': method, + 'BOND_DIM_MAX': 20, + } + traj = HopsTensorTrajectory( + system_param=sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integrator_param, + tensor_param=tensor_param, + ) + traj.initialize(psi_0) + + np.testing.assert_allclose( + traj.storage['psi_traj'][0], + psi_0, + atol=1e-12, + err_msg=f'Initial wavefunction not preserved for {method}', + ) + + +# ------------------------------------------------------------ +# TEST: Auto-detection of nearest-neighbor Hamiltonian +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_initialize_nearest_neighbor_autodetect(): + # This case tests that HopsSystem.flag_nearest_neighbor_ham is set correctly + # during system initialization. The dimer-of-dimers Hamiltonian is + # nearest-neighbor, so it should be True. + traj = HopsTensorTrajectory( + system_param=sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integrator_param, + tensor_param={ + 'MPS_EPSILON': 1e-10, + 'METHOD': 'number', + 'BOND_DIM_MAX': 20, + }, + ) + traj.initialize(psi_0) + assert traj.tensor_basis.system.flag_nearest_neighbor_ham is True, ( + 'HopsSystem should identify the dimer-of-dimers Hamiltonian as nearest-neighbor' + ) + + # Non-NN case: add a long-range coupling (site 0 ↔ site 3) + H2_nonnn = np.array(sys_param['HAMILTONIAN'], dtype=np.complex128) + H2_nonnn[0, 3] = 5.0 + H2_nonnn[3, 0] = 5.0 + nonnn_sys_param = dict(sys_param) + nonnn_sys_param['HAMILTONIAN'] = H2_nonnn + traj_nonnn = HopsTensorTrajectory( + system_param=nonnn_sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integrator_param, + tensor_param={ + 'MPS_EPSILON': 1e-10, + 'METHOD': 'number', + 'BOND_DIM_MAX': 20, + }, + ) + traj_nonnn.initialize(psi_0) + assert traj_nonnn.tensor_basis.system.flag_nearest_neighbor_ham is False, ( + 'Hamiltonian with long-range coupling should not be identified as NN' + ) + + +# ------------------------------------------------------------ +# TEST: Double initialization raises LockedException +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_initialize_double_call_raises(): + # This case tests that calling initialize() twice raises LockedException. + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, + } + traj = HopsTensorTrajectory( + system_param=sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integrator_param, + tensor_param=tensor_param, + ) + traj.initialize(psi_0) + + with pytest.raises(LockedException): + traj.initialize(psi_0) + + +# ------------------------------------------------------------ +# TEST: storage.n_dim is set after initialization +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_storage_n_dim_set(): + # This case tests that storage.n_dim is set to NSTATES after + # initialization, so that storage['psi_traj'] can reconstruct + # the full dense wavefunction in adaptive mode. + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, + } + traj = HopsTensorTrajectory( + system_param=sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integrator_param, + tensor_param=tensor_param, + ) + traj.initialize(psi_0) + assert traj.storage.n_dim == nsite + + +# ------------------------------------------------------------ +# TEST: tensor_basis.eom is constructed during initialization +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_initialize_constructs_eom(): + # This case tests that tensor_basis.eom is set to a HopsTensorEOM + # instance after initialization. + from mesohops.tensor.hops_tensor_eom import HopsTensorEOM + traj = _make_initialized_traj() + assert isinstance(traj.tensor_basis.eom, HopsTensorEOM) + + +# ------------------------------------------------------------ +# TEST: z_mem initialized to zeros +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_initialize_zmem_zeros(): + # This case tests that z_mem is initialized to a complex zero array + # with length matching the noise memory mode indices. + traj = _make_initialized_traj() + assert traj.z_mem.dtype == np.complex128 + np.testing.assert_array_equal(traj.z_mem, 0.0) + + +# ------------------------------------------------------------ +# TEST: self.t = 0 after initialization +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_initialize_time_zero(): + # This case tests that t is set to 0 after initialization. + traj = _make_initialized_traj() + assert traj.t == 0 + + +# ------------------------------------------------------------ +# TEST: timer_checkpoint metadata is stored +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_initialize_stores_timer_metadata(): + # This case tests that INITIALIZATION_TIME is recorded in + # storage metadata after initialization. + traj = _make_initialized_traj() + assert 'INITIALIZATION_TIME' in traj.storage.metadata + assert traj.storage.metadata['INITIALIZATION_TIME'] >= 0 + + +# ============================================================ +# TEST SUITE: propagate() — norm preservation +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Physical wavefunction stays normalized +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_propagate_norm_preserved(): + # This case tests that the physical wavefunction norm stays + # close to 1 throughout propagation for the normalized nonlinear + # equation of motion. + psi_tensor = _run_tensor_hops( + method='fullstate', + ) + + for i_step in range(len(psi_tensor)): + norm = np.linalg.norm(psi_tensor[i_step]) + np.testing.assert_allclose( + norm, + 1.0, + atol=1e-10, + err_msg=f'Norm drifted to {norm} at step {i_step}', + ) + + +# ------------------------------------------------------------ +# TEST: Statenumber norm preservation +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_propagate_norm_preserved_statenumber(): + # This case tests norm preservation in statenumber representation. + psi_tensor = _run_tensor_hops(method='number') + for i_step in range(len(psi_tensor)): + norm = np.linalg.norm(psi_tensor[i_step]) + np.testing.assert_allclose( + norm, 1.0, atol=1e-10, + err_msg=f'Statenumber norm drifted to {norm} at step {i_step}', + ) + + +# ============================================================ +# TEST SUITE: propagate() — bond dimension sensitivity +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Larger bond dimension gives more accurate result +# ------------------------------------------------------------ +@pytest.mark.level(2) +@pytest.mark.parametrize('method', [ + 'fullstate', 'number', +]) +def test_propagate_bond_dim_convergence(method): + # This case tests that increasing bond dimension brings the tensor + # result closer to the vector HOPS result. + psi_vector = _run_vector_hops() + + # Fix mps_epsilon so only bond_dim_max varies; otherwise a failure + # could be caused by either parameter. + psi_small_bond = _run_tensor_hops( + method=method, bond_dim_max=2, mps_epsilon=1e-10, + ) + psi_large_bond = _run_tensor_hops( + method=method, bond_dim_max=20, mps_epsilon=1e-10, + ) + + n_steps = min(len(psi_vector), len(psi_small_bond), len(psi_large_bond)) + err_small = np.mean( + [np.linalg.norm(psi_small_bond[i] - psi_vector[i]) for i in range(n_steps)] + ) + err_large = np.mean( + [np.linalg.norm(psi_large_bond[i] - psi_vector[i]) for i in range(n_steps)] + ) + + assert err_large <= err_small, ( + f'{method}: larger bond dimension gave worse average result: ' + f'err_large={err_large:.2e} > err_small={err_small:.2e}' + ) + assert err_large < 0.1 * err_small or err_large < 1e-12, ( + f'{method}: larger bond dim should substantially improve accuracy: ' + f'err_large={err_large:.2e}, err_small={err_small:.2e}' + ) + + +# ============================================================ +# TEST SUITE: propagate() — TDVP1 integrator +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: TDVP1 approaches RK4 result +# ------------------------------------------------------------ +@pytest.mark.level(2) +@pytest.mark.parametrize('method', [ + 'fullstate', 'number', +]) +def test_propagate_tdvp1_approaches_rk4(method): + # TODO: TDVP1-vs-RK4 error (~0.03 at t_step=4.0) is larger than + # expected and does not converge below ~1e-4 even with small steps. + # Likely due to fundamental algorithmic differences (tangent-space + # projection vs MPO contraction + SVD compression). Verify accuracy + # independently before using TDVP in production. + # + # This case tests that the TDVP1 trajectory is qualitatively + # consistent with the tensor RK4 trajectory. Compares against + # tensor RK4 (not vector HOPS) to isolate integrator error from + # representation error. Uses chi=4; error is identical from chi=4 + # to chi=20 for this system (MPS is low-rank). + t_max_tdvp = 60.0 + chi_tdvp = 4 + + # Tensor RK4 reference (same representation and parameters) + psi_rk4 = _run_tensor_hops( + method=method, + bond_dim_max=chi_tdvp, + mps_epsilon=1e-10, + ) + + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': method, + 'BOND_DIM_MAX': chi_tdvp, + } + traj = HopsTensorTrajectory( + system_param=sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param={'INTEGRATOR': 'TDVP1'}, + tensor_param=tensor_param, + ) + traj.initialize(psi_0) + assert traj.is_tdvp is True, 'TDVP1 integrator was not activated' + traj.propagate(t_max_tdvp, t_step) + psi_tdvp = np.array(traj.storage['psi_traj']) + + n_steps = min(len(psi_rk4), len(psi_tdvp)) + max_err = max(np.linalg.norm(psi_tdvp[i] - psi_rk4[i]) for i in range(n_steps)) + assert max_err < 0.05, ( + f'TDVP1 diverged from tensor RK4: max wf error = {max_err:.2e}' + ) + + +# ============================================================ +# TEST SUITE: TDVP2 propagation +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: TDVP2 approaches RK4 result +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_propagate_tdvp2_approaches_rk4(): + # TODO: same accuracy concern as TDVP1 test above. See that TODO. + # + # This case tests that the TDVP2 trajectory is qualitatively + # consistent with the tensor RK4 trajectory. Compares against + # tensor RK4 (not vector HOPS) to isolate integrator error. + # TDVP2 two-site updates are ~4x more expensive than TDVP1, + # so we use a shorter propagation (20 fs). + t_max_tdvp = 20.0 + chi_tdvp = 4 + + # Tensor RK4 reference (same representation and parameters) + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': chi_tdvp, + } + traj_rk4 = HopsTensorTrajectory( + system_param=sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integrator_param, + tensor_param=tensor_param, + ) + traj_rk4.initialize(psi_0) + traj_rk4.propagate(t_max_tdvp, t_step) + psi_rk4 = np.array(traj_rk4.storage['psi_traj']) + + traj = HopsTensorTrajectory( + system_param=sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param={'INTEGRATOR': 'TDVP2'}, + tensor_param=tensor_param, + ) + traj.initialize(psi_0) + assert traj.is_tdvp is True, 'TDVP2 integrator was not activated' + traj.propagate(t_max_tdvp, t_step) + psi_tdvp2 = np.array(traj.storage['psi_traj']) + + n_steps = min(len(psi_rk4), len(psi_tdvp2)) + max_err = max(np.linalg.norm(psi_tdvp2[i] - psi_rk4[i]) for i in range(n_steps)) + assert max_err < 0.05, ( + f'TDVP2 diverged from tensor RK4: max wf error = {max_err:.2e}' + ) + + +# ============================================================ +# TEST SUITE: HopsTensorTrajectory integrator selection +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: RUNGE_KUTTA sets correct integrator attributes +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_init_rk4_integrator_attributes(): + # This case tests that RUNGE_KUTTA sets TDVP=False and the + # correct step/variable functions. + from mesohops.integrator.tensor_integrator import ( + runge_kutta_step_tensor, + runge_kutta_variables, + ) + + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, + } + traj = HopsTensorTrajectory( + system_param=sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param={'INTEGRATOR': 'RUNGE_KUTTA'}, + tensor_param=tensor_param, + ) + assert traj.is_tdvp is False + assert traj.step is runge_kutta_step_tensor + assert traj.integration_var is runge_kutta_variables + + +# ------------------------------------------------------------ +# TEST: TDVP1 sets correct integrator attributes +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_init_tdvp1_integrator_attributes(): + # This case tests that TDVP1 sets is_tdvp=True and the correct + # step/variable functions. + from mesohops.integrator.tensor_integrator import ( + single_point_variables, + tdvp_step_tensor, + ) + + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, + } + traj = HopsTensorTrajectory( + system_param=sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param={'INTEGRATOR': 'TDVP1'}, + tensor_param=tensor_param, + ) + assert traj.is_tdvp is True + assert traj.step is tdvp_step_tensor + assert traj.integration_var is single_point_variables + assert traj._tdvp_method == '1tdvp' + + +# ------------------------------------------------------------ +# TEST: TDVP2 sets correct integrator attributes +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_init_tdvp2_integrator_attributes(): + # This case tests that TDVP2 sets is_tdvp=True and the correct + # step/variable functions. + from mesohops.integrator.tensor_integrator import ( + single_point_variables, + tdvp_step_tensor, + ) + + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, + } + traj = HopsTensorTrajectory( + system_param=sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param={'INTEGRATOR': 'TDVP2'}, + tensor_param=tensor_param, + ) + assert traj.is_tdvp is True + assert traj.step is tdvp_step_tensor + assert traj.integration_var is single_point_variables + assert traj._tdvp_method == '2tdvp' + + +# ------------------------------------------------------------ +# TEST: Invalid integrator raises UnsupportedRequest +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_init_invalid_integrator_raises(): + # This case tests that an unsupported integrator name raises + # UnsupportedRequest. + from mesohops.util.exceptions import UnsupportedRequest + + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, + } + with pytest.raises(UnsupportedRequest): + HopsTensorTrajectory( + system_param=sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param={'INTEGRATOR': 'INVALID'}, + tensor_param=tensor_param, + ) + + +# ============================================================ +# TEST SUITE: _operator() +# ============================================================ + + +def _make_initialized_traj(): + """Helper: creates and initializes a fullstate tensor trajectory.""" + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, + } + traj = HopsTensorTrajectory( + system_param=sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integrator_param, + tensor_param=tensor_param, + ) + traj.initialize(psi_0) + return traj + + +# ------------------------------------------------------------ +# TEST: Sparse operator input works correctly +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_operator_sparse_input(): + # This case tests that _operator accepts a sparse matrix and + # produces the same result as the dense equivalent. + # Two separate trajectories are created so that the in-place mutation + # from one _operator call does not affect the second comparison. + traj_dense = _make_initialized_traj() + traj_sparse = _make_initialized_traj() + # Build dense and sparse versions of the same operator + H2_op_dense = np.eye(nsite, dtype=np.complex128) * 0.5 + H2_op_sparse = sp.sparse.csr_matrix(H2_op_dense) + # Apply each operator to its own trajectory + traj_dense._operator(H2_op_dense) + traj_sparse._operator(H2_op_sparse) + V1_wf_dense = traj_dense.wavefunction.psi + V1_wf_sparse = traj_sparse.wavefunction.psi + np.testing.assert_allclose( + V1_wf_sparse, + V1_wf_dense, + atol=1e-12, + err_msg='Sparse and dense ops should give same result', + ) + + +# ------------------------------------------------------------ +# TEST: Off-diagonal operator transfers population +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_operator_off_diagonal(): + # This case tests that an off-diagonal operator (swap sites 2 and 3) + # moves population from site 2 to site 3. Unlike the identity and + # projection tests, this exercises the einsum contraction with + # nonzero off-diagonal entries. + traj = _make_initialized_traj() + # Swap operator: |2><3| + |3><2| + identity on sites 0,1 + H2_swap = np.eye(nsite, dtype=np.complex128) + H2_swap[2, 2] = 0.0 + H2_swap[3, 3] = 0.0 + H2_swap[2, 3] = 1.0 + H2_swap[3, 2] = 1.0 + traj._operator(H2_swap) + V1_wf_after = traj.wavefunction.psi + V1_wf_full = np.zeros(nsite, dtype=np.complex128) + V1_wf_full[traj.tensor_basis.system.state_list] = V1_wf_after + # psi_0 had all amplitude on site 2; after swap it should be on site 3 + expected = np.zeros(nsite, dtype=np.complex128) + expected[3] = psi_0[2] + np.testing.assert_allclose( + V1_wf_full, + expected, + atol=1e-12, + err_msg='Swap operator did not transfer population', + ) + + +# ------------------------------------------------------------ +# TEST: Projection onto non-occupied site zeros wavefunction +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_operator_projection_orthogonal(): + # This case tests that projecting onto a site with zero initial + # amplitude zeros out the wavefunction. This is a stronger test + # than the existing projection test (which projects onto the + # occupied site and is effectively identity). + traj = _make_initialized_traj() + H2_proj_0 = np.zeros((nsite, nsite), dtype=np.complex128) + H2_proj_0[0, 0] = 1.0 # project onto site 0, but psi_0 is on site 2 + traj._operator(H2_proj_0) + V1_wf_after = traj.wavefunction.psi + np.testing.assert_allclose( + V1_wf_after, + 0.0, + atol=1e-12, + err_msg='Projection onto empty site should zero wf', + ) + + +# ------------------------------------------------------------ +# TEST: _operator does not reset early time integrator for non-adaptive +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_operator_no_reset_nonadaptive(): + # This case tests that _operator does NOT call reset_early_time_integrator + # for non-adaptive trajectories, matching the parent's behavior where the + # reset only happens inside the adaptive branch. + traj = _make_initialized_traj() + # Propagate a few steps to consume early integrator steps + traj.propagate(8.0, 4.0) + counter_before = traj._early_step_counter + H2_identity = np.eye(nsite, dtype=np.complex128) + traj._operator(H2_identity) + # Counter should be unchanged — _operator does not reset for non-adaptive + assert traj._early_step_counter == counter_before, ( + f'Expected _early_step_counter={counter_before}, ' + f'got {traj._early_step_counter}' + ) + + +# ------------------------------------------------------------ +# TEST: Statenumber off-diagonal operator transfers population +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_operator_statenumber_off_diagonal(): + # This case tests the statenumber _operator path with an off-diagonal + # swap operator on the initial state, isolating statenumber bugs + # from propagation. + traj = _make_initialized_traj_statenumber() + H2_swap = np.eye(nsite, dtype=np.complex128) + H2_swap[2, 2] = 0.0 + H2_swap[3, 3] = 0.0 + H2_swap[2, 3] = 1.0 + H2_swap[3, 2] = 1.0 + traj._operator(H2_swap) + V1_wf_full = np.zeros(nsite, dtype=np.complex128) + V1_wf_full[traj.tensor_basis.system.state_list] = traj.wavefunction.psi + expected = np.zeros(nsite, dtype=np.complex128) + expected[3] = psi_0[2] + np.testing.assert_allclose( + V1_wf_full, expected, atol=1e-10, + err_msg='Statenumber swap did not transfer population', + ) + + +# ------------------------------------------------------------ +# TEST: Statenumber projection onto non-occupied site +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_operator_statenumber_projection_orthogonal(): + # This case tests that projecting onto an unoccupied site zeros the + # wavefunction in statenumber representation. + traj = _make_initialized_traj_statenumber() + H2_proj_0 = np.zeros((nsite, nsite), dtype=np.complex128) + H2_proj_0[0, 0] = 1.0 + traj._operator(H2_proj_0) + np.testing.assert_allclose( + traj.wavefunction.psi, 0.0, atol=1e-12, + err_msg='Statenumber projection onto empty site should zero wf', + ) + + +# ------------------------------------------------------------ +# TEST: Sparse off-diagonal operator exercises CSR indexing +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_operator_sparse_off_diagonal(): + # This case tests that a sparse Hermitian operator with complex + # off-diagonal entries produces the same result as its dense + # equivalent. The 0.5*I test doesn't exercise sparse indexing on + # off-diagonal entries; this tests CSR conversion + np.ix_ trimming + # on non-trivial sparsity patterns. + H2_op = np.array([ + [0.5, 0.1+0.2j, 0.0, 0.05-0.1j], + [0.1-0.2j, 0.3, 0.15+0.1j, 0.0], + [0.0, 0.15-0.1j, 0.7, 0.2+0.05j], + [0.05+0.1j, 0.0, 0.2-0.05j, 0.6], + ], dtype=np.complex128) + traj_dense = _make_initialized_traj() + traj_sparse = _make_initialized_traj() + traj_dense._operator(H2_op) + traj_sparse._operator(sp.sparse.csr_matrix(H2_op)) + np.testing.assert_allclose( + traj_sparse.wavefunction.psi, traj_dense.wavefunction.psi, + atol=1e-12, + err_msg='Sparse Hermitian operator differs from dense', + ) + + +# ------------------------------------------------------------ +# TEST: Statenumber swap on initial state +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_operator_statenumber_swap_initial(): + # This case tests the statenumber _operator path on the initial state + # (before propagation), isolating statenumber bugs from propagation. + traj = _make_initialized_traj() + traj_sn = HopsTensorTrajectory( + system_param=sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integrator_param, + tensor_param={ + 'MPS_EPSILON': 1e-10, + 'METHOD': 'number', + 'BOND_DIM_MAX': 20, + }, + ) + traj_sn.initialize(psi_0) + # Swap sites 2 and 3 + H2_swap = np.zeros((nsite, nsite), dtype=np.complex128) + H2_swap[2, 3] = 1.0 + H2_swap[3, 2] = 1.0 + traj._operator(H2_swap) + traj_sn._operator(H2_swap) + np.testing.assert_allclose( + traj_sn.psi, traj.psi, atol=1e-10, + err_msg='Statenumber swap on initial state differs from fullstate', + ) + + +def _make_initialized_traj_statenumber(): + """Helper: creates and initializes a statenumber tensor trajectory.""" + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'number', + 'BOND_DIM_MAX': 20, + } + traj = HopsTensorTrajectory( + system_param=sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integrator_param, + tensor_param=tensor_param, + ) + traj.initialize(psi_0) + return traj + + +def _make_vacuum_traj(phys_init): + """Helper: number-method trajectory in the vacuum convention. + + Sets flag_gs_vacuum before initialize so the all-zeros MPS + configuration carries the physical ground state, matching the + absorption / fluorescence dipole pathways. + """ + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'number', + 'BOND_DIM_MAX': 20, + } + traj = HopsTensorTrajectory( + system_param=sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integrator_param, + tensor_param=tensor_param, + ) + traj.wavefunction.flag_gs_vacuum = True + traj.initialize(phys_init) + return traj + + +# ------------------------------------------------------------ +# TEST: apply_dipole_lower_plus_ident gives mu_k |g> + |e_k> +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_apply_dipole_lower_plus_ident_action(): + # The lower+ident dipole (sum_k mu_k a_k + I_excited) maps a + # single-excitation state sum_k c_k |e_k> to + # (sum_k mu_k c_k) |g> + sum_k c_k |e_k>: the single-excitation content + # is preserved and the all-zeros configuration picks up the mu-weighted + # sum. The single-excitation state is built the way the spectroscopy + # pathways do -- raise from |g> -- so the MPS stays in the + # single-excitation manifold. A direct multi-site initialize() would + # instead encode a product state carrying spurious double-excitation + # amplitudes c_j*c_k that the lower term would fold back into the + # single-excitation readback. + V1_c = np.zeros(nsite, dtype=np.complex128) + V1_c[1] = 0.6 + V1_c[2] = 0.8j + traj = _make_vacuum_traj(np.zeros(nsite, dtype=np.complex128)) + traj.apply_dipole_raise(V1_c) # |g> -> sum_k c_k |e_k> + + list_mu = np.array([0.5, 1.5 - 0.2j, 0.0, 0.7j], dtype=np.complex128) + traj.apply_dipole_lower_plus_ident(list_mu) + + psi = traj.wavefunction.psi + gs_amp = extract_gs_amp( + traj.wavefunction.list_cores_phi, traj.wavefunction.method, + ) + np.testing.assert_allclose( + psi, V1_c, atol=1e-12, + err_msg='lower+ident must preserve the single-excitation amplitudes', + ) + np.testing.assert_allclose( + gs_amp, np.sum(list_mu * V1_c), atol=1e-12, + err_msg='lower+ident must place sum_k mu_k c_k on the all-zeros config', + ) + + +# ------------------------------------------------------------ +# TEST: apply_dipole_raise_plus_ground_ident gives |g> + sum_k mu_k |e_k> +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_apply_dipole_raise_plus_ground_ident_action(): + # The raise+ground-ident dipole (sum_k mu_k a_k^dagger + I_g) maps the + # ground state |g> to |g> + sum_k mu_k |e_k>: the all-zeros amplitude + # is preserved and each excited site picks up mu_k. Verifies the + # trajectory-level MPO wiring. + traj = _make_vacuum_traj(np.zeros(nsite, dtype=np.complex128)) + + list_mu = np.array([0.5, 1.5 - 0.2j, 0.0, 0.7j], dtype=np.complex128) + traj.apply_dipole_raise_plus_ground_ident(list_mu) + + psi = traj.wavefunction.psi + gs_amp = extract_gs_amp( + traj.wavefunction.list_cores_phi, traj.wavefunction.method, + ) + np.testing.assert_allclose( + psi, list_mu, atol=1e-12, + err_msg='raise+ground-ident must place mu_k on each excited site', + ) + np.testing.assert_allclose( + gs_amp, 1.0, atol=1e-12, + err_msg='raise+ground-ident must preserve the ground-state amplitude', + ) + + +# ------------------------------------------------------------ +# TEST: Propagated statenumber swap matches fullstate swap +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_operator_statenumber_propagated_swap_vs_fullstate(): + # This case tests that a swap operator applied after propagation gives + # the same phi_0 in both representations. This is the strongest + # correctness test: the propagated state has multi-site amplitude and + # populated hierarchy, so the MPO must handle nontrivial bond dimensions, + # mode core pass-throughs, and cross-site transfer matrices. + H2_swap = np.eye(nsite, dtype=np.complex128) + H2_swap[2, 2] = 0.0 + H2_swap[3, 3] = 0.0 + H2_swap[2, 3] = 1.0 + H2_swap[3, 2] = 1.0 + + traj_full = _make_initialized_traj() + traj_full.propagate(20.0, t_step) + traj_full._operator(H2_swap) + + traj_snum = _make_initialized_traj_statenumber() + traj_snum.propagate(20.0, t_step) + traj_snum._operator(H2_swap) + + np.testing.assert_allclose( + traj_snum.wavefunction.psi, + traj_full.wavefunction.psi, + atol=1e-6, + err_msg='Statenumber operator result diverges from fullstate after propagation', + ) + + +# ------------------------------------------------------------ +# TEST: Propagated statenumber raise matches fullstate raise +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_operator_statenumber_propagated_raise_vs_fullstate(): + # This case tests a fluorescence-style raise operator on a propagated + # state. The raise operator |3><0| couples distant sites, exercising + # the long-range transfer-matrix channels in the MPO. + op_raise = np.zeros((nsite, nsite), dtype=np.complex128) + op_raise[3, 0] = 1.0 + + traj_full = _make_initialized_traj() + traj_full.propagate(20.0, t_step) + traj_full._operator(op_raise) + + traj_snum = _make_initialized_traj_statenumber() + traj_snum.propagate(20.0, t_step) + traj_snum._operator(op_raise) + + np.testing.assert_allclose( + traj_snum.wavefunction.psi, + traj_full.wavefunction.psi, + atol=1e-6, + err_msg='Statenumber raise operator diverges from fullstate after propagation', + ) + + +# ------------------------------------------------------------ +# TEST: Propagated statenumber general operator matches fullstate +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_operator_statenumber_propagated_general_vs_fullstate(): + # This case tests a dense operator with all entries nonzero on a + # propagated state. This is the worst case for the MPO bond dimension + # and exercises every channel in the transfer-matrix structure. + H2_general = np.array( + [ + [0.5, 0.1, 0.2, 0.0], + [0.1, 0.3, 0.0, 0.4], + [0.2, 0.0, 0.7, 0.1], + [0.0, 0.4, 0.1, 0.6], + ], + dtype=np.complex128, + ) + + traj_full = _make_initialized_traj() + traj_full.propagate(20.0, t_step) + traj_full._operator(H2_general) + + traj_snum = _make_initialized_traj_statenumber() + traj_snum.propagate(20.0, t_step) + traj_snum._operator(H2_general) + + np.testing.assert_allclose( + traj_snum.wavefunction.psi, + traj_full.wavefunction.psi, + atol=1e-6, + err_msg=( + 'Statenumber general operator diverges from fullstate after propagation' + ), + ) + + +# ------------------------------------------------------------ +# TEST: Tensor _operator matches vector _operator +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_operator_tensor_matches_vector(): + # This case tests that applying a dense Hermitian operator to the + # same initial state via tensor _operator and vector _operator + # produces the same physical wavefunction. Uses complex off-diagonal + # entries coupling all four sites to exercise the full contraction. + H2_op = np.array([ + [0.5, 0.1+0.2j, 0.0, 0.05-0.1j], + [0.1-0.2j, 0.3, 0.15+0.1j, 0.0], + [0.0, 0.15-0.1j, 0.7, 0.2+0.05j], + [0.05+0.1j, 0.0, 0.2-0.05j, 0.6], + ], dtype=np.complex128) + + # Tensor path + traj_tensor = _make_initialized_traj() + traj_tensor._operator(H2_op) + + # Vector path + hops_vector = HOPS( + sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integrator_param, + ) + hops_vector.initialize(psi_0) + hops_vector._operator(H2_op) + + np.testing.assert_allclose( + traj_tensor.psi, hops_vector.psi, atol=1e-10, + err_msg='Tensor _operator result differs from vector _operator', + ) + + +# ------------------------------------------------------------ +# TEST: Zero operator zeroes out the wavefunction cleanly +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_operator_zero_operator(): + # This case tests that applying a zero matrix via _operator produces + # psi = 0 across every state without crashing. Edge case for a + # degenerate operator. + traj = _make_initialized_traj() + H2_zero = np.zeros((nsite, nsite), dtype=np.complex128) + traj._operator(H2_zero) + V1_wf_after = traj.wavefunction.psi + np.testing.assert_allclose( + V1_wf_after, 0.0, atol=1e-12, + err_msg='Zero operator should zero out the wavefunction', + ) + + +# ------------------------------------------------------------ +# TEST: Mis-sized operator raises a clear ValueError +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_operator_wrong_dimensions_raises(): + # This case tests that _operator validates the operator shape + # against the full system size. Without the guard, too-small + # operators leak an IndexError from np.ix_ and too-large operators + # are silently sliced to the first n_state_full x n_state_full + # block — both are invalid silent-failure modes. + # Too-small (2x2 on a 4-state system) must raise. + traj_small = _make_initialized_traj() + H2_small = np.eye(2, dtype=np.complex128) + with pytest.raises(ValueError, match='shape'): + traj_small._operator(H2_small) + # Too-large (5x5 on a 4-state system) must also raise — no silent + # slicing to the first 4x4 block. + traj_large = _make_initialized_traj() + H2_large = np.eye(nsite + 1, dtype=np.complex128) + with pytest.raises(ValueError, match='shape'): + traj_large._operator(H2_large) + + +# ============================================================ +# TEST SUITE: Statenumber propagation +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Statenumber propagation produces valid trajectory +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_propagate_statenumber_valid_trajectory(): + # This case tests that propagation with number + # runs without error and produces a trajectory of the expected + # length with normalized wavefunctions. + psi_tensor = _run_tensor_hops( + method='number', + ) + # Trajectory should have initial + propagation steps + n_expected = int(t_max / t_step) + 1 + assert len(psi_tensor) == n_expected, ( + f'Expected {n_expected} steps, got {len(psi_tensor)}' + ) + # Norm should stay close to 1 + for i_step in range(len(psi_tensor)): + norm = np.linalg.norm(psi_tensor[i_step]) + np.testing.assert_allclose( + norm, + 1.0, + atol=1e-4, + err_msg=f'Statenumber norm drifted to {norm} at step {i_step}', + ) + + +# ------------------------------------------------------------ +# TEST: max_tensor_complexity is populated end-to-end after propagate +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_propagate_populates_max_tensor_complexity(): + # This case tests the full max_complexity round-trip through the + # side-channel: derivative() sets eom.last_matvec_complexity, the + # RK4 step function publishes the per-step max on + # eom.max_complexity_step, propagate reads it and forwards to + # storage via the registered save_max_tensor_complexity callable. + # The storage trace must have an entry per timestep with positive + # values (any RK4 matvec on a non-trivial MPS gives complexity > 0). + traj = _make_initialized_traj() + traj.propagate(t_max, t_step) + list_complexities = traj.storage['max_tensor_complexity'] + n_expected = int(t_max / t_step) + 1 + assert len(list_complexities) == n_expected, ( + f'Expected {n_expected} max_tensor_complexity entries, ' + f'got {len(list_complexities)}' + ) + for i, c in enumerate(list_complexities): + assert isinstance(c, (int, np.integer)), ( + f'Step {i}: expected int complexity, got {type(c).__name__}' + ) + # Step 0 captures the initial state before any RK4 matvec, so 0 is + # valid there. Every subsequent step ran the RK4 derivative chain + # on a non-trivial MPS, so the published max must be positive. + assert list_complexities[0] >= 0 + for i in range(1, len(list_complexities)): + assert list_complexities[i] > 0, ( + f'Step {i}: expected positive complexity, got ' + f'{list_complexities[i]}' + ) + + +# ------------------------------------------------------------ +# TEST: STORE_STEP_TIMING default off keeps per-call totals +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_propagate_step_timing_flag_default_off(): + # This case tests that without the STORE_STEP_TIMING flag, + # LIST_PROPAGATION_TIME retains its pre-flag behavior on the + # tensor path: one float per propagate() call. + traj = _make_initialized_traj() + traj.propagate(t_max, t_step) + traj.propagate(t_max, t_step) + + list_prop_time = traj.storage.metadata['LIST_PROPAGATION_TIME'] + assert len(list_prop_time) == 2 + for entry in list_prop_time: + assert isinstance(entry, float) + assert entry >= 0 + + +# ------------------------------------------------------------ +# TEST: STORE_STEP_TIMING on populates per-step (t, dt) tuples +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_propagate_step_timing_flag_on(): + # This case tests that with STORE_STEP_TIMING=True on the tensor + # path, LIST_PROPAGATION_TIME holds (t_fs, wall_seconds) tuples + # — one per integration step, sim-time stamps match the + # propagation grid, wall times are non-negative, and entries + # concatenate across propagate() calls with a monotonic time + # axis. + integration_param = { + 'INTEGRATOR': 'RUNGE_KUTTA', + 'STORE_STEP_TIMING': True, + } + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, + } + traj = HopsTensorTrajectory( + system_param=sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integration_param, + tensor_param=tensor_param, + ) + traj.initialize(psi_0) + n_steps = int(np.ceil(t_max / t_step)) + traj.propagate(t_max, t_step) + + list_prop_time = traj.storage.metadata['LIST_PROPAGATION_TIME'] + assert len(list_prop_time) == n_steps + for entry in list_prop_time: + assert isinstance(entry, tuple) and len(entry) == 2 + list_t = [t for t, _ in list_prop_time] + list_dt = [dt for _, dt in list_prop_time] + assert all(dt >= 0 for dt in list_dt) + np.testing.assert_allclose(list_t, t_step * np.arange(1, n_steps + 1)) + + # Second propagate call: entries append, time axis stays monotonic. + traj.propagate(t_max, t_step) + list_prop_time = traj.storage.metadata['LIST_PROPAGATION_TIME'] + assert len(list_prop_time) == 2 * n_steps + list_t = [t for t, _ in list_prop_time] + list_dt = [dt for _, dt in list_prop_time] + assert all(dt >= 0 for dt in list_dt) + assert np.all(np.diff(list_t) > 0) + + +# ------------------------------------------------------------ +# TEST: inchworm max-tracking takes max across iterations +# ------------------------------------------------------------ +@pytest.mark.level(2) +@pytest.mark.xfail( + reason='Adaptive tensor trajectory does not yet propagate cleanly ' + 'past define_basis (pre-existing noise-index mismatch noted ' + 'in test_initialize_adaptive_sets_adaptive_true). When that ' + 'root cause is fixed, this test will start passing and ' + 'strict=True will flag the xfail mark for removal.', + strict=True, +) +def test_inchworm_max_complexity_tracking(): + # This case tests that propagate's inchworm loop publishes the MAX + # of the per-iteration max_complexity values to storage, not the + # last one or the first one. The inchworm path runs only when the + # trajectory is adaptive AND the early-integrator counter is below + # the cap, so the test needs an adaptive trajectory. + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, + } + eom_adaptive = { + 'EQUATION_OF_MOTION': 'NORMALIZED NONLINEAR', + 'ADAPTIVE_S': True, + 'DELTA_S': 0.01, + } + traj_adap = HopsTensorTrajectory( + system_param=sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_adaptive, + integration_param={ + 'INTEGRATOR': 'RUNGE_KUTTA', + 'EARLY_INTEGRATOR_STEPS': 5, + 'INCHWORM_CAP': 5, + }, + tensor_param=tensor_param, + ) + # Currently crashes in initialize() at the adaptive define_basis + # step (noise index mismatch — pre-existing issue). xfail catches + # the crash; when fixed, propagate runs and the assertion below is + # the actual contract being verified. + traj_adap.initialize(psi_0) + traj_adap.propagate(t_max, t_step) + list_complexities = traj_adap.storage['max_tensor_complexity'] + n_expected = int(t_max / t_step) + 1 + assert len(list_complexities) == n_expected + # Inchworm contract: the per-step published value is the MAX across + # the inchworm iterations for that step, so post-step-0 entries must + # be positive (every step ran at least one matvec on a non-trivial + # MPS). + for i in range(1, len(list_complexities)): + assert list_complexities[i] > 0 + + + + +# ============================================================ +# TEST SUITE: Checkpoint guards +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: save_checkpoint raises NotImplementedError +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_save_checkpoint_raises(): + # This case tests that save_checkpoint raises NotImplementedError + # because tensor trajectory checkpointing is not yet supported. + traj = _make_initialized_traj() + with pytest.raises(NotImplementedError): + traj.save_checkpoint('/tmp/dummy_checkpoint.npz') + + +# ------------------------------------------------------------ +# TEST: load_checkpoint raises NotImplementedError +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_load_checkpoint_raises(): + # This case tests that load_checkpoint raises NotImplementedError + # because tensor trajectory checkpointing is not yet supported. + with pytest.raises(NotImplementedError): + HopsTensorTrajectory.load_checkpoint('/tmp/dummy_checkpoint.npz') + + +# ============================================================ +# TEST SUITE: statenumber checkpoint deep copy +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: statenumber MPS checkpoint is independent of original +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_statenumber_checkpoint_deep_copy(): + # This case tests that the checkpoint copy pattern used in propagate + # and inchworm_integrate produces an independent copy for statenumber + # representation. Mutating the original wavefunction after checkpoint + # creation should not affect the checkpoint, and vice versa. + traj = _make_initialized_traj_statenumber() + cores = traj.wavefunction.list_cores_phi + # Make a checkpoint using the same pattern as propagate + if traj.wavefunction.method == 'number': + checkpoint = [ + [arr.copy() for arr in g] for g in cores + ] + else: + checkpoint = [c.copy() for c in cores] + # Save a reference value from the checkpoint + val_before = checkpoint[0][0][0, 0, 0].copy() + # Mutate the original wavefunction's inner array + cores[0][0][0, 0, 0] *= 999.0 + # Checkpoint should be unchanged + assert checkpoint[0][0][0, 0, 0] == val_before, ( + 'Statenumber checkpoint was corrupted by mutation of original cores' + ) + + +# ============================================================ +# TEST SUITE: psi property returns compact wavefunction +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: psi returns compact array matching active state count +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_psi_returns_compact(): + # This case tests that traj.psi returns a compact array of + # length n_active (number of active states), not NSTATES. + # This matches the parent HopsTrajectory.psi behavior. + traj = _make_initialized_traj() + psi = traj.psi + # Non-adaptive: all states active, so n_active == NSTATES + assert len(psi) == nsite + # Analytical: psi_0 = [0, 0, 1, 0], so psi should match at init + np.testing.assert_allclose(psi, psi_0, atol=1e-12) + + +# ------------------------------------------------------------ +# TEST: psi works in statenumber representation +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_psi_returns_compact_statenumber(): + # This case tests extract_psi in statenumber representation. + traj = _make_initialized_traj_statenumber() + psi = traj.psi + assert len(psi) == nsite + np.testing.assert_allclose(psi, psi_0, atol=1e-12) + + +# ------------------------------------------------------------ +# TEST: storage psi_traj reconstructs correctly +# ------------------------------------------------------------ +@pytest.mark.level(2) +@pytest.mark.parametrize('method', [ + 'fullstate', 'number', +]) +def test_storage_psi_traj_after_propagate(method): + # This case tests that storage['psi_traj'] returns correctly + # shaped arrays after propagation. + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': method, + 'BOND_DIM_MAX': 20, + } + traj = HopsTensorTrajectory( + system_param=sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integrator_param, + tensor_param=tensor_param, + ) + traj.initialize(psi_0) + traj.propagate(8.0, 4.0) + psi_traj = np.array(traj.storage['psi_traj']) + # Should have shape (n_steps, NSTATES) + assert psi_traj.shape[1] == nsite + # Initial wavefunction should match psi_0 + np.testing.assert_allclose(psi_traj[0], psi_0, atol=1e-12) + # Later time steps should differ from initial (dynamics happened) + assert not np.allclose(psi_traj[-1], psi_0, atol=1e-6), ( + 'Final state identical to initial — propagation may not have run' + ) + + +# ============================================================ +# TEST SUITE: phi property +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: phi setter stores and retrieves cores +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_phi_setter(): + # This case tests that the phi setter stores the provided list of + # cores and the getter returns the exact same object (no copy). + traj = _make_initialized_traj() + original_phi = traj.phi + new_phi = [c * 2.0 for c in original_phi] + traj.phi = new_phi + # Invariant: getter returns what setter stored (same object) + assert traj.phi is new_phi + + +# ------------------------------------------------------------ +# TEST: phi after initialization reproduces psi_0 +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_phi_initial_reproduces_psi0(): + # This case tests that contracting the MPS cores from traj.phi + # reproduces the initial wavefunction psi_0. + traj = _make_initialized_traj() + np.testing.assert_allclose(traj.psi, psi_0, atol=1e-12, + err_msg='phi cores do not reproduce psi_0 after initialization') + + +# ------------------------------------------------------------ +# TEST: phi works in statenumber representation +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_phi_statenumber(): + # This case tests that traj.phi returns valid MPS cores in + # statenumber representation (nested list of core groups). + traj = _make_initialized_traj_statenumber() + phi = traj.phi + assert isinstance(phi, list) + assert len(phi) > 0 + # Statenumber phi is a list of lists (groups per state) + assert isinstance(phi[0], list) + # Verify psi is recoverable + np.testing.assert_allclose(traj.psi, psi_0, atol=1e-12) + + +# ------------------------------------------------------------ +# TEST: phi property shares reference with wavefunction +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_phi_shared_reference(): + # This case tests that traj.phi and traj.wavefunction.list_cores_phi + # are the same object — mutations through one path are visible + # through the other. + traj = _make_initialized_traj() + assert traj.phi is traj.wavefunction.list_cores_phi + # Mutate through wavefunction, read through phi + traj.wavefunction.list_cores_phi[0] = traj.wavefunction.list_cores_phi[0] * 2.0 + assert traj.phi[0] is traj.wavefunction.list_cores_phi[0] + + +# ============================================================ +# TEST SUITE: normalize() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: normalize is no-op when basis.eom.normalized is False +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_normalize_unnormalized_noop(): + # This case tests that normalize() leaves the wavefunction unchanged + # when the EOM does not require normalization. Constructs with + # NONLINEAR EOM (not NORMALIZED NONLINEAR) so basis.eom.normalized + # is False without monkey-patching. + traj = HopsTensorTrajectory( + system_param=sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param={'EQUATION_OF_MOTION': 'NONLINEAR'}, + integration_param=integrator_param, + tensor_param={ + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, + }, + ) + traj.initialize(psi_0) + assert traj.basis.eom.normalized is False + # Scale wavefunction so any accidental normalize() would be detectable + traj.phi = [c * 3.0 for c in traj.phi] + psi_scaled = traj.psi.copy() + traj.normalize() + # Invariant: wavefunction unchanged when basis.eom.normalized is False + np.testing.assert_allclose(traj.psi, psi_scaled, atol=1e-14) + + +# ------------------------------------------------------------ +# TEST: normalize actually normalizes when EOM requires it +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_normalize_rescales_wavefunction(): + # This case tests that normalize() actually restores unit norm + # when the EOM requires normalization. + traj = _make_initialized_traj() + # Scale wavefunction away from unit norm + traj.phi = [c * 3.0 for c in traj.phi] + assert abs(np.linalg.norm(traj.psi) - 1.0) > 0.1, 'Precondition: norm should differ from 1' + traj.normalize() + np.testing.assert_allclose( + np.linalg.norm(traj.psi), 1.0, atol=1e-10, + err_msg='normalize() did not restore unit norm', + ) + + +# ------------------------------------------------------------ +# TEST: normalize works in statenumber representation +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_normalize_statenumber(): + # This case tests the statenumber normalize path which indexes + # list_cores_phi[0][0] instead of list_cores_phi[0]. + traj = _make_initialized_traj_statenumber() + traj.phi = [[core * 3.0 for core in group] for group in traj.phi] + traj.normalize() + np.testing.assert_allclose( + np.linalg.norm(traj.psi), 1.0, atol=1e-10, + err_msg='Statenumber normalize() did not restore unit norm', + ) + + +# ------------------------------------------------------------ +# TEST: single RK4 step keeps norm within 1e-10 before normalize() +# ------------------------------------------------------------ +@pytest.mark.level(1) +@pytest.mark.parametrize('method', [ + 'fullstate', 'number', +]) +def test_normalize_single_step_norm_drift(method): + # This case tests that a single RK4 step preserves the physical + # wavefunction norm to within 1e-10 of unity before normalize() is + # applied. Uses a 2-site NORMALIZED NONLINEAR system (same physics as + # the vector counterpart) so both code paths are tested against + # identical physics. If the norm correction term in the tensor EOM + # derivative is wrong (e.g. wrong prefactor), norm drifts measurably + # even in one step. + T3_loperator = np.zeros([2, 2, 2], dtype=np.float64) + T3_loperator[0, 0, 0] = 1.0 + T3_loperator[1, 1, 1] = 1.0 + local_sys_param = { + 'HAMILTONIAN': np.array([[0, 10.0], [10.0, 0]], dtype=np.float64), + 'GW_SYSBATH': [[10.0, 10.0], [5.0, 5.0], [10.0, 10.0], [5.0, 5.0]], + 'L_HIER': [T3_loperator[0], T3_loperator[0], + T3_loperator[1], T3_loperator[1]], + 'L_NOISE1': [T3_loperator[0], T3_loperator[0], + T3_loperator[1], T3_loperator[1]], + 'ALPHA_NOISE1': bcf_exp, + 'PARAM_NOISE1': [[10.0, 10.0], [5.0, 5.0], [10.0, 10.0], [5.0, 5.0]], + } + local_noise_param = { + 'SEED': 0, + 'MODEL': 'FFT_FILTER', + 'TLEN': 10.0, + 'TAU': 1.0, + } + local_eom_param = { + 'TIME_DEPENDENCE': False, + 'EQUATION_OF_MOTION': 'NORMALIZED NONLINEAR', + } + local_hier_param = {'MAXHIER': 4} + local_integrator_param = {'INTEGRATOR': 'RUNGE_KUTTA'} + tensor_param = { + 'MPS_EPSILON': 1e-12, + 'METHOD': method, + 'BOND_DIM_MAX': 20, + } + H1_psi_0 = np.array([1.0 + 0.0j, 0.0 + 0.0j]) + + traj = HopsTensorTrajectory( + system_param=local_sys_param, + noise_param=local_noise_param, + hierarchy_param=local_hier_param, + eom_param=local_eom_param, + integration_param=local_integrator_param, + tensor_param=tensor_param, + ) + traj.initialize(H1_psi_0) + tau = 2.0 # 2 * noise TAU so RK4 samples [0, 1, 2] land on the noise grid + + # Gather integration variables for one step + dict_var = traj.integration_var( + traj.z_mem, + traj.t, + traj.noise1, + traj.noise2, + tau, + traj.basis.mode.list_l2idx_abs, + traj.effective_noise_integration, + ) + + # Call RK4 directly — mutates wavefunction in place, bypasses normalize() + traj._step(dict_var) + + # Norm of the physical wavefunction after RK4 (no normalize applied) + norm_psi = np.linalg.norm(traj.psi) + + np.testing.assert_allclose( + norm_psi, 1.0, atol=1e-10, + err_msg=( + f'Single-step norm drift {abs(norm_psi - 1.0):.2e} exceeds 1e-10. ' + 'The norm correction term in the tensor EOM derivative may be wrong.' + ), + ) + + +# ============================================================ +# TEST SUITE: propagate() — error guards +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: propagate raises when t_axis exceeds noise length +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_propagate_t_axis_exceeds_tlen(): + # This case tests that propagate raises TrajectoryError when the + # requested propagation time exceeds noise.param['TLEN']. The guard + # is at np.max(t_axis) > self.noise1.param['TLEN']. + traj = _make_initialized_traj() + # Try to propagate way beyond the noise length (TLEN=25000 fs) + with pytest.raises(TrajectoryError, match='longer than'): + traj.propagate(100000.0, 4.0) + + +# ------------------------------------------------------------ +# TEST: propagate raises on timestep/noise TAU mismatch +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_propagate_tau_mismatch_raises(): + # This case tests that propagate raises TrajectoryError when the + # timestep does not align with noise.param['TAU']. + traj = _make_initialized_traj() + # noise TAU is 1.0 fs; use a non-divisor timestep + with pytest.raises(TrajectoryError, match='TAU'): + traj.propagate(10.0, 0.7) + + +# ------------------------------------------------------------ +# TEST: propagate raises when TAU is None and INTERPOLATE is True +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_propagate_tau_none_interpolate_raises(): + # This case tests that propagate raises UnboundLocalError when + # noise TAU is None and INTERPOLATE is True. Under these conditions + # the t_axis variable is never defined, but is used downstream. + # This is a known bug (see NOTE in hops_tensor_trajectory.propagate); + # this test documents the failure mode until the guard is fixed. + traj = _make_initialized_traj() + traj.noise1.param['TAU'] = None + traj.noise1.param['INTERPOLATE'] = True + with pytest.raises(UnboundLocalError): + traj.propagate(8.0, 4.0) + + +# ------------------------------------------------------------ +# TEST: unsupported early integrator raises UnsupportedRequest +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_propagate_unsupported_early_integrator_raises(): + # This case tests that an unsupported EARLY_ADAPTIVE_INTEGRATOR + # value raises UnsupportedRequest during adaptive propagation. + # We initialize non-adaptive then patch the integration_param to + # trigger the unsupported early integrator path, because the full + # adaptive initialization path has a pre-existing noise indexing bug. + traj = _make_initialized_traj() + traj.basis.eom.param['ADAPTIVE'] = True + traj.basis.eom.param['ADAPTIVE_S'] = True + traj.tensor_basis.adaptive = True + traj.integration_param['EARLY_ADAPTIVE_INTEGRATOR'] = 'INVALID' + traj.integration_param['EARLY_INTEGRATOR_STEPS'] = 5 + traj._early_step_counter = 0 + with pytest.raises(UnsupportedRequest, match='does not support'): + traj.propagate(4.0, 4.0) + + +# ------------------------------------------------------------ +# TEST: system timescale warning fires when tau is too large +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_propagate_system_timescale_warning(): + # This case tests that propagate emits a warning when tau exceeds + # the system Hamiltonian timescale. + traj = _make_initialized_traj() + # Use a very large timestep that exceeds system timescale + # (system_timescale ~ 1/(max eigenvalue spread) ~ few fs) + large_tau = 1000.0 + with pytest.warns(UserWarning, match='timescale'): + traj.propagate(large_tau, large_tau) diff --git a/tests/test_hops_tensor_wavefunction.py b/tests/test_hops_tensor_wavefunction.py new file mode 100644 index 0000000..0168414 --- /dev/null +++ b/tests/test_hops_tensor_wavefunction.py @@ -0,0 +1,1567 @@ +import numpy as np +import pytest +import scipy as sp + +from mesohops.basis.hops_modes import HopsModes +from mesohops.basis.hops_noise_memory import HopsNoiseMemory +from mesohops.basis.hops_system import HopsSystem +from mesohops.tensor.hops_tensor_wavefunction import HopsTensorWavefunction +from mesohops.tensor.hops_tensor_basis import HopsTensorBasis +from mesohops.trajectory.exp_noise import bcf_exp +from mesohops.trajectory.hops_tensor_trajectory import HopsTensorTrajectory +from mesohops.util.bath_corr_functions import bcf_convert_dl_to_exp +from mesohops.util.exceptions import LockedException, UnsupportedRequest +from mesohops.basis.basis_functions import determine_error_thresh +from mesohops.util.tensor_operations import ( + extract_psi, + phi_aux as extract_phi_aux, + unflatten_cores, +) + +__title__ = 'Unit Tests for HopsTensorWavefunction' +__author__ = 'N. Covalsen' +__maintainer__ = 'N. Covalsen' + +# ============================================================ +# Shared Setup +# ============================================================ +# Dimer-of-dimers system (4 sites, 2 modes per site) +# identical to test_hops_tensor_trajectory.py and test_dimer_of_dimers.py + +nsite = 4 +e_lambda = 20.0 +gamma = 50.0 +temp = 140.0 +(g_0, w_0) = bcf_convert_dl_to_exp(e_lambda, gamma, temp) + +T3_loperator = np.zeros([4, 4, 4], dtype=np.float64) +list_gw_sysbath = [] +list_lop = [] +for i in range(nsite): + T3_loperator[i, i, i] = 1.0 + list_gw_sysbath.append([g_0, w_0]) + list_lop.append(sp.sparse.coo_matrix(T3_loperator[i])) + list_gw_sysbath.append([-1j * np.imag(g_0), 500.0]) + list_lop.append(T3_loperator[i]) + +H2_hamiltonian = np.zeros([nsite, nsite]) +H2_hamiltonian[0, 1] = 40 +H2_hamiltonian[1, 0] = 40 +H2_hamiltonian[1, 2] = 10 +H2_hamiltonian[2, 1] = 10 +H2_hamiltonian[2, 3] = 40 +H2_hamiltonian[3, 2] = 40 + +sys_param = { + 'HAMILTONIAN': np.array(H2_hamiltonian, dtype=np.complex128), + 'GW_SYSBATH': list_gw_sysbath, + 'L_HIER': list_lop, + 'L_NOISE1': list_lop, + 'ALPHA_NOISE1': bcf_exp, + 'PARAM_NOISE1': list_gw_sysbath, +} + +psi_0 = np.array([0.0] * nsite, dtype=np.complex128) +psi_0[2] = 1.0 + +k_max = 2 +state_list = np.arange(nsite) +delta_s = 0 # non-adaptive + + +eom_param = {'EQUATION_OF_MOTION': 'NORMALIZED NONLINEAR'} +noise_param = { + 'SEED': 0, + 'MODEL': 'FFT_FILTER', + 'TLEN': 25000.0, + 'TAU': 1.0, +} +hier_param = {'MAXHIER': 2} +integrator_param = {'INTEGRATOR': 'RUNGE_KUTTA'} + + +def _make_basis_objects(param=None): + """Creates HopsSystem, HopsModes, HopsNoiseMemory directly.""" + if param is None: + param = sys_param + system = HopsSystem(param) + mode = HopsModes(system, hierarchy=None) + noise_memory = HopsNoiseMemory(system, mode) + return system, mode, noise_memory + + +def _make_initialized_tensor_basis(param=None, ds=None): + """Creates and initializes a HopsTensorBasis.""" + if ds is None: + ds = delta_s + system, mode, noise_memory = _make_basis_objects(param) + tb = HopsTensorBasis(system, mode, noise_memory) + system.initialize(ds > 0, psi_0) + mode.list_modeidx_abs = sorted(system.list_statemodeidx_abs) + noise_memory.initialize() + system.state_list = state_list + tb.initialize(ds) + return tb + + +def _make_propagated_tensor(method, t_max=8.0, t_step=4.0): + """Creates a tensor trajectory, propagates, and returns the wavefunction. + + After propagation the MPS has physically realistic entanglement and + complex-valued cores with non-trivial bond dimensions — unlike the + trivial product state from _make_tensor. + """ + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': method, + 'BOND_DIM_MAX': 20, + } + traj = HopsTensorTrajectory( + system_param=sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integrator_param, + tensor_param=tensor_param, + ) + traj.initialize(psi_0) + traj.propagate(t_max, t_step) + return traj.wavefunction + + +def _make_tensor(method): + """Creates and initializes a HopsTensorWavefunction for the dimer-of-dimers system. + + NOTE: produces a trivial product state (single-site occupied, bond + dim 1). Tests requiring entangled or delocalized states should + use _make_propagated_tensor instead. + """ + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': method, + 'BOND_DIM_MAX': 20, + } + integrator_param = {'INTEGRATOR': 'RUNGE_KUTTA'} + tb = _make_initialized_tensor_basis() + ht = HopsTensorWavefunction(k_max, tensor_param, integrator_param, eom_param) + ht.initialize(psi_0, tb.system) + return ht + + +def _make_tensor_uninit(method): + """Creates a HopsTensorWavefunction WITHOUT calling initialize.""" + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': method, + 'BOND_DIM_MAX': 20, + } + integrator_param = {'INTEGRATOR': 'RUNGE_KUTTA'} + return HopsTensorWavefunction(k_max, tensor_param, integrator_param, eom_param) + + +def _make_tensor_basis_init(): + """Creates and initializes a HopsTensorBasis.""" + return _make_initialized_tensor_basis() + + +# ============================================================ +# TEST SUITE: __init__() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Config parameters are stored correctly +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_init_stores_config(): + """HopsTensorWavefunction stores tensor_param and integrator_param on init.""" + ht = _make_tensor_uninit('fullstate') + assert ht.mps_epsilon == 1e-10 + assert ht.method == 'fullstate' + assert ht.bond_dim_max == 20 + assert ht.k_max == 2 + assert ht.flag_norm is True + assert ht.flag_tdvp is False + # Default values before initialize + assert ht.list_cores_phi == [] + assert ht.__initialized__ is False + # system/mode are NOT stored on HopsTensorWavefunction + assert not hasattr(ht, 'system') + assert not hasattr(ht, 'mode') + assert not hasattr(ht, 'noise_memory') + + +# ------------------------------------------------------------ +# TEST: flag_gs_vacuum property defaults off and guards method +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_gs_vacuum_property_requires_number(): + # The flag defaults off and is set via its property. The all-zeros vacuum + # config only exists in number representation, so setting it True must + # raise for fullstate and succeed for number. + ht_full = _make_tensor_uninit('fullstate') + assert ht_full.flag_gs_vacuum is False + with pytest.raises(UnsupportedRequest): + ht_full.flag_gs_vacuum = True + + ht_num = _make_tensor_uninit('number') + assert ht_num.flag_gs_vacuum is False + ht_num.flag_gs_vacuum = True + assert ht_num.flag_gs_vacuum is True + + +def _number_mps_from_dense(dense, list_modes_per_site, epsilon=1e-12): + """Build a nested number MPS from a dense config tensor via exact TT-SVD. + + The dense axes are the MPS physical legs in order (site_0, mode_00, ...); + the TT round-trip is lossless, so the MPS represents the dense state + exactly. + """ + size = list(dense.shape) + flat_cores = [] + C = dense + rank = 1 + for i in range(len(size) - 1): + C = np.reshape(C, (int(rank * size[i]), int(C.size / (rank * size[i])))) + U, S, Vt = np.linalg.svd(C, full_matrices=False) + normalized_S = S / np.linalg.norm(S) + thr = determine_error_thresh(np.flip(normalized_S), epsilon * epsilon) + S[normalized_S <= thr] = 0.0 + prev_rank, rank = rank, int(np.count_nonzero(S)) + flat_cores.append( + U[:, :rank].astype(np.complex128).reshape(prev_rank, size[i], rank) + ) + C = np.diag(S[:rank]).astype(np.complex128) @ Vt[:rank, :] + flat_cores.append(C.reshape(C.shape[0], C.shape[1], 1)) + return unflatten_cores(flat_cores, list_modes_per_site) + + +# ------------------------------------------------------------ +# TEST: manifold_norm_sq sums excited norm and vacuum amplitude +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_manifold_norm_sq_includes_vacuum(): + # Build a random number MPS via TT and wrap it in a wavefunction. + # manifold_norm_sq must equal sum_s |psi_s|^2 for the physical wavefunction, + # i.e. site s occupation at k=0, plus |vacuum|^2 from the all-zeros config. + rng = np.random.default_rng(1) + k_max_local = 1 + list_modes_per_site = [1, 1] + size, site_axes, pos = [], [], 0 + for n_modes in list_modes_per_site: + site_axes.append(pos) + size += [2] + [k_max_local + 1] * n_modes + pos += 1 + n_modes + dense = rng.standard_normal(size) + 1j * rng.standard_normal(size) + n_state = len(list_modes_per_site) + + vac_idx = (0,) * len(size) + expected = np.abs(dense[vac_idx]) ** 2 + for s in range(n_state): + idx = [0] * len(size) + idx[site_axes[s]] = 1 + expected += np.abs(dense[tuple(idx)]) ** 2 + + ht = _make_tensor_uninit('number') + ht.M1_modes_per_site = np.array(list_modes_per_site, dtype=int) + ht.list_cores_phi = _number_mps_from_dense(dense, list_modes_per_site) + np.testing.assert_allclose(ht.manifold_norm_sq, expected, atol=1e-10) + + +# ------------------------------------------------------------ +# TEST: EOM flags are set correctly for different configurations +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_init_eom_flags(): + # This case tests that EOM-related flags are set correctly + # for different EOM and integrator configurations. + + # flag_norm=False when EOM is NONLINEAR + tensor_param_local = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, + } + integrator_param = {'INTEGRATOR': 'RUNGE_KUTTA'} + eom_nonlinear = {'EQUATION_OF_MOTION': 'NONLINEAR'} + ht = HopsTensorWavefunction( + k_max, tensor_param_local, integrator_param, eom_nonlinear, + ) + assert ht.flag_norm is False + assert ht.flag_tdvp is False + + # flag_tdvp=True when INTEGRATOR is 'TDVP1' + integrator_param_tdvp = {'INTEGRATOR': 'TDVP1'} + ht_tdvp = HopsTensorWavefunction( + k_max, tensor_param_local, integrator_param_tdvp, eom_param + ) + assert ht_tdvp.flag_norm is True + assert ht_tdvp.flag_tdvp is True + + +# ------------------------------------------------------------ +# TEST: flag_norm derives from eom_param, not tensor_param +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_init_flag_norm_from_eom_param(): + """HopsTensorWavefunction.flag_norm is set from eom_param, not tensor_param['EOM']. + """ + # tensor_param has no 'EOM' key — the EOM lives in eom_param only + tensor_param_no_eom = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, + } + integrator_param = {'INTEGRATOR': 'RUNGE_KUTTA'} + + # Case 1: normalized nonlinear → flag_norm=True + eom_normalized = {'EQUATION_OF_MOTION': 'NORMALIZED NONLINEAR'} + ht = HopsTensorWavefunction( + k_max, tensor_param_no_eom, integrator_param, eom_normalized, + ) + assert ht.flag_norm is True + + # Case 2: nonlinear → flag_norm=False + eom_nonlinear = {'EQUATION_OF_MOTION': 'NONLINEAR'} + ht2 = HopsTensorWavefunction( + k_max, tensor_param_no_eom, integrator_param, eom_nonlinear, + ) + assert ht2.flag_norm is False + + # Case 3: linear → flag_norm=False + eom_linear = {'EQUATION_OF_MOTION': 'LINEAR'} + ht3 = HopsTensorWavefunction( + k_max, tensor_param_no_eom, integrator_param, eom_linear, + ) + assert ht3.flag_norm is False + + +# ============================================================ +# TEST SUITE: initialize() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Dimension bookkeeping is correct +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_initialize_dimensions(): + # This case tests that M1_modes_per_state is computed + # correctly from system parameters after initialize. + ht = _make_tensor('number') + modes_per_site = len(list_gw_sysbath) // nsite + np.testing.assert_array_equal( + ht.M1_modes_per_state, [modes_per_site] * nsite + ) + + +# ------------------------------------------------------------ +# TEST: Statenumber MPS has correct group structure +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_initialize_statenumber_group_structure(): + # This case tests that each state group in a statenumber MPS has + # 1 state core + M1_modes_per_state[s] mode cores. + ht = _make_tensor('number') + for s, group in enumerate(ht.list_cores_phi): + expected_len = 1 + ht.M1_modes_per_state[s] + assert len(group) == expected_len, ( + f'State {s} group has {len(group)} cores, expected {expected_len}' + ) + # State core has physical dimension 2 (occupied/unoccupied) + assert group[0].shape[1] == 2 + # Mode cores have physical dimension k_max + 1 + for core_m in group[1:]: + assert core_m.shape[1] == k_max + 1 + + +# ------------------------------------------------------------ +# TEST: Sparse Hamiltonian is converted to dense +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_initialize_sparse_hamiltonian_converted(): + # This case tests that a sparse Hamiltonian is converted to dense during + # initialize (HopsTensorBasis/HopsSystem handles the conversion). + sp_param = dict(sys_param) + sp_param['HAMILTONIAN'] = sp.sparse.coo_matrix(H2_hamiltonian) + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, + } + integrator_param = {'INTEGRATOR': 'RUNGE_KUTTA'} + tb = _make_initialized_tensor_basis(param=sp_param) + ht = HopsTensorWavefunction(k_max, tensor_param, integrator_param, eom_param) + ht.initialize(psi_0, tb.system) + # Verify the stored Hamiltonian is dense (MpoBuilder requires dense input). + # HopsTensorTrajectory converts sparse → dense in __init__; here we bypass + # that, so we convert manually to confirm the test exercises the right path. + H2_ham = tb.system.param['HAMILTONIAN'] + if sp.sparse.issparse(H2_ham): + tb.system.param['HAMILTONIAN'] = np.asarray(H2_ham.todense()) + assert isinstance(tb.system.param['HAMILTONIAN'], np.ndarray), ( + 'Hamiltonian should be dense ndarray for tensor code' + ) + # Analytical: fullstate MPS has 1 state core + sum(M1_modes_per_state) mode cores + expected_n_cores = 1 + sum(ht.M1_modes_per_state) + assert len(ht.list_cores_phi) == expected_n_cores + # Analytical: psi should recover the initial state + np.testing.assert_allclose(ht.psi, psi_0, atol=1e-12) + + +# ------------------------------------------------------------ +# TEST: M1_modes_per_site is set from state_list +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_initialize_modes_per_site(): + # This case tests that M1_modes_per_site has one entry per state in state_list. + ht = _make_tensor('fullstate') + assert len(ht.M1_modes_per_site) == len(state_list) + + +# ============================================================ +# TEST SUITE: build_list_cores_phi() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Fullstate MPS has correct structure +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_list_cores_phi_fullstate_structure(): + # This case tests that the fullstate MPS has 1 system core + n_mode + # mode cores, with correct physical dimensions. + ht = _make_tensor('fullstate') + n_modes = len(list_gw_sysbath) + # 1 system core + n_modes mode cores + assert len(ht.list_cores_phi) == 1 + n_modes + # System core physical dimension = n_state + assert ht.list_cores_phi[0].shape[1] == nsite + # Mode cores physical dimension = k_max + 1 + for i in range(1, len(ht.list_cores_phi)): + assert ht.list_cores_phi[i].shape[1] == k_max + 1 + + +# ------------------------------------------------------------ +# TEST: Statenumber MPS has correct structure +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_list_cores_phi_statenumber_structure(): + # This case tests that the statenumber MPS has n_state groups, + # each containing 1 state core (physical dim 2) + mode cores (physical dim k_max+1). + ht = _make_tensor('number') + # list_cores_phi[s] = [state_core_s, mode_core_s0, ...] — one group per state + assert len(ht.list_cores_phi) == nsite + for site in range(nsite): + group = ht.list_cores_phi[site] + # First element is the state core: physical dim 2 + assert group[0].shape[1] == 2 + # Remaining elements are mode cores: physical dim k_max + 1 + assert len(group) == 1 + ht.M1_modes_per_state[site] + for mode_core in group[1:]: + assert mode_core.shape[1] == k_max + 1 + + +# ------------------------------------------------------------ +# TEST: phi_0 recovers the initial wavefunction +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_list_cores_phi_recovers_psi0(): + # This case tests that contracting the MPS back to a vector + # recovers the input wavefunction, for both representations. + for method in ['fullstate', 'number']: + ht = _make_tensor(method) + np.testing.assert_allclose( + ht.psi, + psi_0, + atol=1e-12, + err_msg=f'phi_0 does not match psi_0 for {method}', + ) + + +# ------------------------------------------------------------ +# TEST: All hierarchy modes start in ground state +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_list_cores_phi_hierarchy_ground_state(): + # This case tests that all first-order auxiliary wavefunctions + # are zero at initialization (all modes in ground state |0>). + for method in ['fullstate', 'number']: + ht = _make_tensor(method) + n_modes = len(list_gw_sysbath) + for i_mode in range(n_modes): + indices = [0] * n_modes + indices[i_mode] = 1 + phi_1 = extract_phi_aux( + ht.list_cores_phi, ht.method, ht.M1_modes_per_state, indices + ) + np.testing.assert_allclose( + phi_1, + 0.0, + atol=1e-12, + err_msg=f'Mode {i_mode} not in ground state for {method}', + ) + + +# ============================================================ +# TEST SUITE: inflate_bonds_to() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Bond dimensions reach target after inflation +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_inflate_bonds_to_target(): + # This case tests that all internal bond dimensions equal chi_target + # after explicitly calling inflate_bonds_to (inflate is not called + # automatically during initialize for non-TDVP integrators). + ht = _make_tensor('fullstate') + chi = ht.bond_dim_max + ht.inflate_bonds_to(chi, eps=0.0) + for i in range(len(ht.list_cores_phi) - 1): + assert ht.list_cores_phi[i].shape[2] == chi, ( + f'Bond {i} right dim = {ht.list_cores_phi[i].shape[2]}, expected {chi}' + ) + assert ht.list_cores_phi[i + 1].shape[0] == chi, ( + f'Bond {i} left dim = {ht.list_cores_phi[i + 1].shape[0]}, expected {chi}' + ) + # Verify no pathological output + for i, core in enumerate(ht.list_cores_phi): + assert not np.any(np.isnan(core)), f'Inflate produced NaN in core {i}' + assert core.shape[0] >= 1, f'Zero left bond in core {i}' + assert core.shape[2] >= 1, f'Zero right bond in core {i}' + + +# ------------------------------------------------------------ +# TEST: Bond consistency (left dim matches right dim) +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_check_bondsize_both_representations(): + # This case tests that adjacent cores have compatible bond dimensions + # for statenumber representation. + ht = _make_tensor('number') + assert ht.check_bondsize(), 'Bond dimensions are inconsistent' + + # This case tests check_bondsize for fullstate representation. + ht_full = _make_tensor('fullstate') + assert ht_full.check_bondsize(), 'Fullstate bond dimensions inconsistent at init' + + # This case tests check_bondsize after inflate_bonds_to. + ht_inflated = _make_tensor('fullstate') + ht_inflated.inflate_bonds_to(5, eps=0.0) + assert ht_inflated.check_bondsize(), 'Bond dimensions inconsistent after inflation' + + +# ------------------------------------------------------------ +# TEST: inflate_bonds_to with eps > 0 adds random noise padding +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_inflate_bonds_eps_nonzero(): + # This case tests that inflate_bonds_to with eps > 0 produces + # non-zero padding entries (random noise rather than zeros). + ht = _make_tensor('fullstate') + # Record original bond dims (all 1 for freshly built MPS) + original_shapes = [core.shape for core in ht.list_cores_phi] + chi_target = 5 + # eps > 0 fills padded bond entries with random noise scaled by eps, + # ensuring the inflated MPS is not rank-deficient. Any eps > 0 + # guarantees nonzero padding that survives subsequent SVD compression. + ht.inflate_bonds_to(chi_target, eps=0.1) + # Check that padded region (beyond original bond dim) contains nonzero entries + has_nonzero_pad = False + for core, orig_shape in zip(ht.list_cores_phi, original_shapes): + orig_left, _, orig_right = orig_shape + # Padded entries are those with left index >= orig_left or right >= orig_right + if core.shape[0] > orig_left: + if np.any(np.abs(core[orig_left:, :, :]) > 1e-15): + has_nonzero_pad = True + if core.shape[2] > orig_right: + if np.any(np.abs(core[:, :, orig_right:]) > 1e-15): + has_nonzero_pad = True + assert has_nonzero_pad, 'eps>0 padding should contain nonzero entries' + + +# ------------------------------------------------------------ +# TEST: inflate_bonds_to is a no-op when bonds already at target +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_inflate_bonds_noop_at_target(): + # This case tests that calling inflate_bonds_to with chi_target + # equal to the current bond dimension does not change the cores. + # First inflate to 4, then inflate again to 4 — the second call + # should be a no-op since bonds are already at target. + ht = _make_tensor('fullstate') + ht.inflate_bonds_to(4, eps=0.0) + cores_at_4 = [c.copy() for c in ht.list_cores_phi] + ht.inflate_bonds_to(4, eps=0.0) + for i, (before, after) in enumerate(zip(cores_at_4, ht.list_cores_phi)): + np.testing.assert_array_equal( + before, + after, + err_msg=f'Core {i} changed when chi_target == current bond dim', + ) + + +# ------------------------------------------------------------ +# TEST: Fullstate inflate preserves phi_0 +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_inflate_bonds_preserves_phi0_fullstate(): + # This case tests that inflating bonds does not change the physical + # wavefunction. Uses a propagated state with naturally complex-valued + # cores and non-trivial bond structure. + ht = _make_propagated_tensor('fullstate') + V1_phi_before = ht.psi.copy() + ht.inflate_bonds_to(ht.bond_dim_max) + np.testing.assert_allclose( + ht.psi, + V1_phi_before, + atol=1e-12, + err_msg='inflate_bonds_to changed phi_0 in fullstate', + ) + + +# ------------------------------------------------------------ +# TEST: Mixed bond dimensions — only undersized bonds are padded +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_inflate_bonds_mixed_bond_dims(): + # CASE: Inflate to 3, then to 5. Verify that the content in the + # bond-3 region is preserved after the second inflate (not + # overwritten), and that the new padding region is distinct. + ht = _make_tensor('fullstate') + # Inject complex values to avoid trivial special case + ht.list_cores_phi[0][0, :, 0] = np.array( + [0.3 + 0.1j, 0.5 - 0.2j, 0.6 + 0.4j, 0.1 - 0.3j], + dtype=np.complex128, + ) + ht.inflate_bonds_to(3, eps=0.0) + V1_psi_at_3 = ht.psi.copy() + cores_at_3 = [c.copy() for c in ht.list_cores_phi] + ht.inflate_bonds_to(5, eps=0.0) + # All interior bonds should now be 5 + for i in range(len(ht.list_cores_phi) - 1): + assert ht.list_cores_phi[i].shape[2] == 5, ( + f'Bond {i} right dim should be 5 after second inflate' + ) + # Physical wavefunction must be preserved through both inflations + np.testing.assert_allclose( + ht.psi, V1_psi_at_3, atol=1e-12, + err_msg='Second inflate changed psi', + ) + # The bond-3 subregion of each core should be preserved (not + # overwritten by the inflate-to-5 padding). + for i, (c3, c5) in enumerate(zip(cores_at_3, ht.list_cores_phi)): + dl3, d, dr3 = c3.shape + np.testing.assert_allclose( + c5[:dl3, :, :dr3], c3, atol=1e-14, + err_msg=f'Core {i}: bond-3 subregion altered by inflate to 5', + ) + + +# ------------------------------------------------------------ +# TEST: Double inflate preserves phi_0 +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_inflate_bonds_double_inflate_preserves_phi0(): + # CASE: Two successive inflations should still preserve the + # physical wavefunction. Uses complex multi-site state. + ht = _make_tensor('fullstate') + ht.list_cores_phi[0][0, :, 0] = np.array( + [0.4 - 0.3j, 0.2 + 0.5j, 0.1 + 0.1j, 0.6 - 0.2j], + dtype=np.complex128, + ) + V1_phi_before = ht.psi.copy() + ht.inflate_bonds_to(3, eps=0.0) + ht.inflate_bonds_to(10, eps=0.0) + np.testing.assert_allclose( + ht.psi, + V1_phi_before, + atol=1e-12, + err_msg='Double inflate changed phi_0', + ) + + +# ============================================================ +# TEST SUITE: normalize() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Normalize returns the norm before normalization +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_normalize_returns_norm(): + # This case tests that normalize() returns the norm of phi_0 + # before dividing. + ht = _make_tensor('fullstate') + # Scale phi to have non-unit norm + ht.list_cores_phi[0] = ht.list_cores_phi[0] * 2.0 + norm_before = np.linalg.norm(ht.psi) + returned_norm = ht.normalize() + np.testing.assert_allclose(returned_norm, norm_before, atol=1e-12) + + +# ------------------------------------------------------------ +# TEST: Normalize makes phi_0 unit norm +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_normalize_unit_norm(): + # This case tests that after normalize(), phi_0 has unit norm. + ht = _make_tensor('fullstate') + ht.list_cores_phi[0] = ht.list_cores_phi[0] * 3.7 + ht.normalize() + np.testing.assert_allclose( + np.linalg.norm(ht.psi), + 1.0, + atol=1e-12, + ) + + +# ------------------------------------------------------------ +# TEST: Normalize only divides first core +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_normalize_only_first_core(): + # This case tests that normalization only modifies the first core, + # leaving all other cores unchanged. + ht = _make_tensor('fullstate') + ht.list_cores_phi[0] = ht.list_cores_phi[0] * 2.0 + cores_before = [c.copy() for c in ht.list_cores_phi[1:]] + ht.normalize() + for i, (before, after) in enumerate(zip(cores_before, ht.list_cores_phi[1:])): + np.testing.assert_array_equal( + before, + after, + err_msg=f'Core {i + 1} was modified by normalize()', + ) + + +# ------------------------------------------------------------ +# TEST: Normalize always normalizes (even when EOM is NONLINEAR) +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_normalize_always_normalizes(): + # This case tests that normalize() always normalizes the MPS, + # even when the EOM is NONLINEAR (flag_norm is False). The + # flag_norm guard was removed from normalize(); the policy now + # lives in the trajectory. + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, + } + integrator_param = {'INTEGRATOR': 'RUNGE_KUTTA'} + eom_nonlinear = {'EQUATION_OF_MOTION': 'NONLINEAR'} + tb_la = _make_initialized_tensor_basis() + ht = HopsTensorWavefunction(k_max, tensor_param, integrator_param, eom_nonlinear) + ht.initialize(psi_0, tb_la.system) + ht.list_cores_phi[0] = ht.list_cores_phi[0] * 5.0 + core0_before = ht.list_cores_phi[0].copy() + ht.normalize() + # The core SHOULD have changed (normalize divides by norm) + assert not np.array_equal(ht.list_cores_phi[0], core0_before) + # phi_0 should now have unit norm + np.testing.assert_allclose(np.linalg.norm(ht.psi), 1.0, atol=1e-12) + + +# ------------------------------------------------------------ +# TEST: normalize works correctly for statenumber representation +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_normalize_statenumber(): + # This case tests that normalize produces unit-norm psi for statenumber + ht = _make_tensor('number') + # Scale first core to make norm != 1 + ht.list_cores_phi[0][0] = ht.list_cores_phi[0][0] * 3.0 + psi_before = ht.psi.copy() + norm_returned = ht.normalize() + # Invariant: returned norm equals pre-normalize norm + np.testing.assert_allclose( + norm_returned, np.linalg.norm(psi_before), atol=1e-12, + ) + # Invariant: post-normalize norm is 1 + np.testing.assert_allclose(np.linalg.norm(ht.psi), 1.0, atol=1e-12) + # Invariant: psi direction unchanged + np.testing.assert_allclose( + ht.psi, psi_before / np.linalg.norm(psi_before), atol=1e-12, + ) + + +# ------------------------------------------------------------ +# TEST: Statenumber normalize only modifies first state core +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_normalize_only_first_core_statenumber(): + # This case tests that normalization only modifies list_cores_phi[0][0] + # (the first state core), leaving other cores in the first group and + # all other groups unchanged. + ht = _make_tensor('number') + ht.list_cores_phi[0][0] = ht.list_cores_phi[0][0] * 2.0 + # Save copies of all cores except [0][0] + list_first_group_rest = [c.copy() for c in ht.list_cores_phi[0][1:]] + list_other_groups = [ + [c.copy() for c in group] for group in ht.list_cores_phi[1:] + ] + ht.normalize() + # Other cores in the first group should be unchanged + for i, (before, after) in enumerate( + zip(list_first_group_rest, ht.list_cores_phi[0][1:]) + ): + np.testing.assert_array_equal( + before, after, + err_msg=f'Group 0, core {i + 1} was modified by normalize()', + ) + # All other groups should be unchanged + for g, (group_before, group_after) in enumerate( + zip(list_other_groups, ht.list_cores_phi[1:]) + ): + for i, (before, after) in enumerate(zip(group_before, group_after)): + np.testing.assert_array_equal( + before, after, + err_msg=f'Group {g + 1}, core {i} was modified by normalize()', + ) + + +# ------------------------------------------------------------ +# TEST: normalize raises UnsupportedRequest for unknown method +# ------------------------------------------------------------ +def test_normalize_unsupported_method(): + # This case tests that normalize raises an error when the tensor + # method is not recognized. extract_psi (called via self.psi) + # raises ValueError before the method dispatch in normalize. + ht = _make_tensor('fullstate') + ht.method = 'bogus' + with pytest.raises((UnsupportedRequest, ValueError)): + ht.normalize() + + +# ============================================================ +# TEST SUITE: Properties (psi, phi_aux, flat_cores) +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: phi_aux returns correct length +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_phi_aux_returns_correct_length(): + # This case tests that phi_aux returns a vector of length n_state. + for method in ['fullstate', 'number']: + ht = _make_tensor(method) + n_modes = len(list_gw_sysbath) + indices = [0] * n_modes + indices[0] = 1 + phi_1 = extract_phi_aux( + ht.list_cores_phi, ht.method, ht.M1_modes_per_state, indices + ) + assert len(phi_1) == nsite + # Analytical: at initialization, all hierarchy occupations are zero, + # so any first-order auxiliary wavefunction must be zero. + np.testing.assert_allclose( + phi_1, 0.0, atol=1e-12, + err_msg=f'First-order auxiliary should be zero at init for {method}', + ) + + +# ============================================================ +# TEST SUITE: MPS utilities (restore_phi, tensor_compress, +# check_bondsize, linksize, get_core_shapes) +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: restore_phi from list replaces cores +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_restore_phi_from_list(): + # This case tests that restore_phi with a list copies the cores into + # list_cores_phi (independent copy, not a direct reference) and that + # subsequent mutation of the source list does not affect the tensor. + for method in ['fullstate', 'number']: + ht = _make_tensor(method) + if method == 'number': + # list_cores_phi is a list-of-lists for statenumber; + # create a scaled copy with the same group structure. + new_cores = [[core * 2.0 for core in group] for group in ht.list_cores_phi] + ht.restore_phi(new_cores) + # Values should match the source + for g_set, g_src in zip(ht.list_cores_phi, new_cores): + for c_set, c_src in zip(g_set, g_src): + np.testing.assert_array_equal(c_set, c_src) + # list_cores_phi must be an independent copy, not the same object + assert ht.list_cores_phi is not new_cores + # Mutating the source array must not affect the tensor + core_before = ht.list_cores_phi[0][0].copy() + new_cores[0][0] *= 0.0 + np.testing.assert_array_equal(ht.list_cores_phi[0][0], core_before) + else: + new_cores = [c * 2.0 for c in ht.list_cores_phi] + ht.restore_phi(new_cores) + for c_set, c_src in zip(ht.list_cores_phi, new_cores): + np.testing.assert_array_equal(c_set, c_src) + assert ht.list_cores_phi is not new_cores + core_before = ht.list_cores_phi[0].copy() + new_cores[0] *= 0.0 + np.testing.assert_array_equal(ht.list_cores_phi[0], core_before) + + +# ------------------------------------------------------------ +# TEST: restore_phi from another HopsTensorWavefunction copies cores +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_restore_phi_from_tensor(): + # This case tests that restore_phi from another HopsTensorWavefunction copies + # the cores (not just a reference). Mutating ht2 after restore_phi + # must not affect ht1. Tested for both representations. + for method in ['fullstate', 'number']: + ht1 = _make_tensor(method) + ht2 = _make_tensor(method) + if method == 'number': + # list_cores_phi[0] is a group (list); scale the state core of group 0 + ht2.list_cores_phi[0][0] = ht2.list_cores_phi[0][0] * 3.0 + else: + ht2.list_cores_phi[0] = ht2.list_cores_phi[0] * 3.0 + # Copy cores from ht2 into ht1 + ht1.restore_phi(ht2) + # Verify restore_phi actually transferred the correct values + np.testing.assert_allclose( + ht1.psi, ht2.psi, atol=1e-12, + err_msg=f'restore_phi did not transfer correct psi ({method})', + ) + phi0_after_set = ht1.psi.copy() + # Mutate ht2 — this must not leak into ht1 + if method == 'number': + ht2.list_cores_phi[0][0] = ht2.list_cores_phi[0][0] * 0.0 + else: + ht2.list_cores_phi[0] = ht2.list_cores_phi[0] * 0.0 + # ht1 should still have the pre-mutation values + np.testing.assert_allclose( + ht1.psi, + phi0_after_set, + atol=1e-12, + err_msg=f'Mutating source tensor leaked into restore_phi copy ({method})', + ) + + +# ------------------------------------------------------------ +# TEST: get_core_shapes returns correct shapes +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_get_core_shapes(): + # This case tests that get_core_shapes returns one shape tuple per core + # and each shape matches the actual core array dimensions. + ht = _make_tensor('fullstate') + shapes = ht.get_core_shapes() + assert len(shapes) == len(ht.list_cores_phi) + for i, shape in enumerate(shapes): + assert shape == ht.list_cores_phi[i].shape + + # This case tests get_core_shapes for statenumber (must flatten groups) + ht_sn = _make_tensor('number') + shapes_sn = ht_sn.get_core_shapes() + flat_cores = ht_sn.flat_cores + assert len(shapes_sn) == len(flat_cores) + for i, shape in enumerate(shapes_sn): + assert shape == flat_cores[i].shape + + +# ============================================================ +# TEST SUITE: __init__() — flag_tdvp edge cases +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: flag_tdvp is True for TDVP2 integrator +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_init_flag_tdvp_true_for_tdvp2(): + # This case tests that flag_tdvp is True for TDVP2 (not just TDVP1). + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, + } + ht = HopsTensorWavefunction(k_max, tensor_param, {'INTEGRATOR': 'TDVP2'}, eom_param) + assert ht.flag_tdvp is True + + +# ============================================================ +# TEST SUITE: initialize() guards and edge cases +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Double initialize raises LockedException +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_initialize_double_raises(): + # This case tests that calling initialize() twice raises LockedException. + # _make_tensor already calls initialize once internally. + ht = _make_tensor('fullstate') + tb = _make_tensor_basis_init() + with pytest.raises(LockedException): + ht.initialize(psi_0, tb.system) + + +# ------------------------------------------------------------ +# TEST: TDVP path inflates bond dimensions +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_initialize_tdvp_inflates_bonds(): + # This case tests that when flag_tdvp is True, initialize + # inflates bond dimensions up to bond_dim_max. + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, + } + integrator_param = {'INTEGRATOR': 'TDVP1'} + tb = _make_initialized_tensor_basis() + ht = HopsTensorWavefunction(k_max, tensor_param, integrator_param, eom_param) + ht.initialize(psi_0, tb.system) + # Analytical: TDVP inflate sets all internal bonds to bond_dim_max + for i, core in enumerate(ht.list_cores_phi[:-1]): + assert core.shape[2] == 20, ( + f'Core {i} right bond should be bond_dim_max=20, got {core.shape[2]}' + ) + + +# ------------------------------------------------------------ +# TEST: k_max=0 produces valid single-level hierarchy +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_initialize_k_max_zero(): + # This case tests that k_max=0 (single hierarchy level) produces + # valid MPS cores where mode cores have local dimension 1. + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, + } + integrator_param = {'INTEGRATOR': 'RUNGE_KUTTA'} + tb = _make_initialized_tensor_basis() + ht = HopsTensorWavefunction(0, tensor_param, integrator_param, eom_param) + ht.initialize(psi_0, tb.system) + # Analytical: fullstate MPS has 1 state core + n_total_modes mode cores + expected_n_cores = 1 + sum(ht.M1_modes_per_state) + assert len(ht.list_cores_phi) == expected_n_cores + # Mode cores should have local dimension k_max+1 = 1 + # First core is the system core, subsequent are mode cores + for core in ht.list_cores_phi[1:]: + assert core.shape[1] == 1, ( + f'Mode core should have local dim 1 for k_max=0, got {core.shape[1]}' + ) + + +# ============================================================ +# TEST SUITE: build_list_cores_phi() — error path +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Invalid method raises UnsupportedRequest +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_list_cores_phi_invalid_method_raises(): + # This case tests that passing an unknown method to + # build_list_cores_phi raises UnsupportedRequest. + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'bogus_method', + 'BOND_DIM_MAX': 20, + } + integrator_param = {'INTEGRATOR': 'RUNGE_KUTTA'} + ht = HopsTensorWavefunction(k_max, tensor_param, integrator_param, eom_param) + tb = _make_initialized_tensor_basis() + with pytest.raises(UnsupportedRequest): + ht.build_list_cores_phi(psi_0, len(tb.system.state_list)) + + +# ============================================================ +# TEST SUITE: restore_phi() — error path +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Invalid type raises TypeError +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_restore_phi_invalid_type_raises(): + # This case tests that passing a non-list, non-HopsTensorWavefunction + # argument to restore_phi raises TypeError. + ht = _make_tensor('fullstate') + with pytest.raises(TypeError, match='Expected list or HopsTensorWavefunction'): + ht.restore_phi('not_a_list_or_tensor') + + +# ============================================================ +# TEST SUITE: inflate_bonds_to() — statenumber coverage (Ritesh T1) +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Statenumber inflate reaches target bond dim +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_inflate_bonds_to_statenumber(): + # This case tests that inflate_bonds_to works for statenumber + # representation and all internal bonds reach chi_target. + ht = _make_tensor('number') + chi_target = 5 + ht.inflate_bonds_to(chi_target) + flat = ht.flat_cores + for i in range(len(flat) - 1): + bond_right = flat[i].shape[2] + bond_left_next = flat[i + 1].shape[0] + assert bond_right == chi_target, ( + f'Core {i} right bond {bond_right} != target {chi_target}' + ) + assert bond_left_next == chi_target, ( + f'Core {i + 1} left bond {bond_left_next} != target {chi_target}' + ) + + +# ------------------------------------------------------------ +# TEST: Statenumber inflate preserves phi_0 +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_inflate_bonds_preserves_phi0_statenumber(): + # This case tests that inflating bonds in statenumber representation + # does not change the physical wavefunction. + ht = _make_tensor('number') + V1_phi_before = ht.psi.copy() + ht.inflate_bonds_to(5) + np.testing.assert_allclose( + ht.psi, + V1_phi_before, + atol=1e-12, + err_msg='inflate_bonds_to changed phi_0 in statenumber', + ) + + +# ------------------------------------------------------------ +# TEST: Statenumber inflate with eps > 0 adds random noise padding +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_inflate_bonds_eps_nonzero_statenumber(): + # This case tests that inflate_bonds_to with eps > 0 produces + # non-zero padding entries in statenumber representation. + ht = _make_tensor('number') + original_shapes = [core.shape for core in ht.flat_cores] + chi_target = 5 + ht.inflate_bonds_to(chi_target, eps=0.1) + has_nonzero_pad = False + for core, orig_shape in zip(ht.flat_cores, original_shapes): + orig_left, _, orig_right = orig_shape + if core.shape[0] > orig_left: + if np.any(np.abs(core[orig_left:, :, :]) > 1e-15): + has_nonzero_pad = True + if core.shape[2] > orig_right: + if np.any(np.abs(core[:, :, orig_right:]) > 1e-15): + has_nonzero_pad = True + assert has_nonzero_pad, 'eps>0 padding should contain nonzero entries' + + +# ------------------------------------------------------------ +# TEST: Statenumber inflate is a no-op when bonds already at target +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_inflate_bonds_noop_at_target_statenumber(): + # This case tests that calling inflate_bonds_to with chi_target equal + # to the current bond dimension does not change the cores in + # statenumber representation. + ht = _make_tensor('number') + cores_before = [c.copy() for c in ht.flat_cores] + current_chi = ht.flat_cores[0].shape[2] + ht.inflate_bonds_to(current_chi, eps=0.0) + for i, (before, after) in enumerate(zip(cores_before, ht.flat_cores)): + np.testing.assert_array_equal( + before, + after, + err_msg=f'Core {i} changed when chi_target == current bond dim', + ) + + +# ============================================================ +# TEST SUITE: add_state_cores() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: add_state_cores fullstate increases core count +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_add_state_cores_fullstate_count(): + # This case tests that add_state_cores increases the number of mode + # cores by modes_per_state per new state (fullstate representation). + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, + } + integrator_param = {'INTEGRATOR': 'RUNGE_KUTTA'} + psi_2 = np.array([0.0, 1.0], dtype=np.complex128) + system, mode, noise_memory = _make_basis_objects() + system.initialize(False, psi_2) + mode.list_modeidx_abs = sorted(system.list_statemodeidx_abs) + noise_memory.initialize() + system.state_list = [0, 1] + ht = HopsTensorWavefunction(k_max, tensor_param, integrator_param, eom_param) + ht.initialize(psi_2, system) + cores_before = len(ht.list_cores_phi) + n_new_modes = sum(ht.M1_modes_per_state[s] for s in [2]) + ht.add_state_cores([2], list(system.state_list)) + assert len(ht.list_cores_phi) == cores_before + n_new_modes + assert ht.list_cores_phi[0].shape[1] == 3 + assert len(ht.M1_modes_per_site) == 3 + + +# ------------------------------------------------------------ +# TEST: add_state_cores statenumber increases group count +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_add_state_cores_statenumber_count(): + # This case tests that add_state_cores adds one new group per new state + # in statenumber representation, and the state count increases by one. + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'number', + 'BOND_DIM_MAX': 20, + } + integrator_param = {'INTEGRATOR': 'RUNGE_KUTTA'} + psi_2 = np.array([0.0, 1.0], dtype=np.complex128) + system, mode, noise_memory = _make_basis_objects() + system.initialize(False, psi_2) + mode.list_modeidx_abs = sorted(system.list_statemodeidx_abs) + noise_memory.initialize() + system.state_list = [0, 1] + ht = HopsTensorWavefunction(k_max, tensor_param, integrator_param, eom_param) + ht.initialize(psi_2, system) + groups_before = len(ht.list_cores_phi) + ht.add_state_cores([2], list(system.state_list)) + # Statenumber: one new group added + assert len(ht.list_cores_phi) == groups_before + 1 + assert len(ht.M1_modes_per_site) == 3 + # Analytical: adding a zero-amplitude state preserves existing amplitudes + # and the new state has zero amplitude + phi_after = extract_psi( + ht.list_cores_phi, + 'number', + ht.M1_modes_per_state, + ) + # This case tests that the newly added state has zero amplitude + # psi is ordered by sorted state list: [0, 1, 2] + # State 2 was just added and should have zero amplitude + np.testing.assert_allclose( + phi_after[2], 0.0, atol=1e-12, + err_msg='Newly added state should have zero amplitude', + ) + + +# ------------------------------------------------------------ +# TEST: add_state_cores preserves phi_0 for existing states +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_add_state_cores_fullstate_preserves_phi0(): + # This case tests that adding a new state leaves the existing phi_0 amplitudes + # unchanged and places zero amplitude on the newly added state. + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, + } + integrator_param = {'INTEGRATOR': 'RUNGE_KUTTA'} + psi_2 = np.array([0.0, 1.0], dtype=np.complex128) + system, mode, noise_memory = _make_basis_objects() + system.initialize(False, psi_2) + mode.list_modeidx_abs = sorted(system.list_statemodeidx_abs) + noise_memory.initialize() + system.state_list = [0, 1] + ht = HopsTensorWavefunction(k_max, tensor_param, integrator_param, eom_param) + ht.initialize(psi_2, system) + phi_0_before = ht.psi.copy() + ht.add_state_cores([2], list(system.state_list)) + phi_0_after = ht.psi + assert len(phi_0_after) == 3 + np.testing.assert_allclose(phi_0_after[:2], phi_0_before, atol=1e-12) + np.testing.assert_allclose(phi_0_after[2], 0.0, atol=1e-12) + + +# ------------------------------------------------------------ +# TEST: add_state_cores statenumber preserves existing amplitudes +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_add_state_cores_statenumber_preserves_phi0(): + # This case tests that existing amplitudes are unchanged after + # adding a new state in statenumber representation. + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'number', + 'BOND_DIM_MAX': 20, + } + integrator_param = {'INTEGRATOR': 'RUNGE_KUTTA'} + psi_2 = np.array([0.0, 1.0], dtype=np.complex128) + system, mode, noise_memory = _make_basis_objects() + system.initialize(False, psi_2) + mode.list_modeidx_abs = sorted(system.list_statemodeidx_abs) + noise_memory.initialize() + system.state_list = [0, 1] + ht = HopsTensorWavefunction(k_max, tensor_param, integrator_param, eom_param) + ht.initialize(psi_2, system) + psi_before = ht.psi.copy() + ht.add_state_cores([2], list(system.state_list)) + psi_after = ht.psi + # Analytical: existing states keep their amplitudes, new state has zero + np.testing.assert_allclose(psi_after[:2], psi_before, atol=1e-12) + np.testing.assert_allclose(psi_after[2], 0.0, atol=1e-12) + + +# ============================================================ +# TEST SUITE: remove_state_cores() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: remove_state_cores fullstate decreases core count +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_remove_state_cores_fullstate_count(): + # This case tests that remove_state_cores reduces the core count by + # modes_per_state for the removed state, shrinks the system core physical + # dimension, and decreases the site count by one. + ht = _make_tensor('fullstate') + cores_before = len(ht.list_cores_phi) + state_list = list(range(nsite)) + n_modes_removed = sum(ht.M1_modes_per_state[s] for s in [0]) + ht.remove_state_cores([0], state_list) + assert len(ht.list_cores_phi) == cores_before - n_modes_removed + assert ht.list_cores_phi[0].shape[1] == nsite - 1 + assert len(ht.M1_modes_per_site) == nsite - 1 + + +# ------------------------------------------------------------ +# TEST: remove_state_cores statenumber decreases group count +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_remove_state_cores_statenumber_count(): + # This case tests that remove_state_cores removes exactly one group + # in statenumber representation and the site count decreases by one. + ht = _make_tensor('number') + groups_before = len(ht.list_cores_phi) + state_list = list(range(nsite)) + ht.remove_state_cores([0], state_list) + assert len(ht.list_cores_phi) == groups_before - 1 + assert len(ht.M1_modes_per_site) == nsite - 1 + # Analytical: removing state 0 (zero amplitude in psi_0=[0,0,1,0]) + # should preserve total probability + phi_after = extract_psi( + ht.list_cores_phi, + 'number', + ht.M1_modes_per_state, + ) + # This case tests that total probability is preserved after removing + # a zero-amplitude state + assert np.sum(np.abs(phi_after) ** 2) > 0.5, ( + 'Total probability collapsed after removing zero-amplitude state' + ) + + +# ------------------------------------------------------------ +# TEST: remove_state_cores statenumber leftmost group (known bug) +# ------------------------------------------------------------ +@pytest.mark.level(1) +@pytest.mark.xfail( + reason=( + 'Known bug: remove_state_cores statenumber group_idx==0 uses `pass` ' + 'instead of absorbing |0> slices into the next group. The bug is visible ' + 'only when the left bond dimension Dr0 > 1 (post-propagation states). ' + 'Simple product-state initializations have Dr0==1 so the `else: pass` ' + 'path happens to be harmless. This xfail documents the structural defect.' + ), + strict=True, +) +def test_remove_state_cores_statenumber_leftmost_bug(): + # This case tests that removing the leftmost state in statenumber + # preserves the wavefunction correctly when the removed state has zero + # amplitude but the next state carries amplitude. + # psi starts on state 1 — state 0 is the leftmost group (zero amplitude). + # Removing state 0 must absorb its |0> slice into state 1's group. + # The known bug (else: pass in remove_state_cores) discards the |0> slice + # instead of absorbing it. Bonds are inflated to bond dim > 1 so the |0> + # slices carry information and the bug actually manifests. + psi_state1 = np.array([0.0, 1.0, 0.0, 0.0], dtype=np.complex128) + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'number', + 'BOND_DIM_MAX': 20, + } + integrator_param = {'INTEGRATOR': 'RUNGE_KUTTA'} + tb = _make_initialized_tensor_basis() + ht = HopsTensorWavefunction(k_max, tensor_param, integrator_param, eom_param) + ht.initialize(psi_state1, tb.system) + # Inflate bonds so the leftmost-removal bug actually manifests. + # After inflation, the inter-group bonds have dim > 1. The |0> slice of + # the removed group's state core is shape (1, Dr0) with Dr0=4. The bug + # discards this slice instead of absorbing it into the next group. This + # leaves the new first group with an orphaned left bond of dim Dr0 instead + # of 1, which is detectable by inspecting the core shape. + ht.inflate_bonds_to(4, eps=0.01) + full_state_list = list(range(nsite)) + ht.remove_state_cores([0], full_state_list) + # After correct removal, the new first group must have left bond dim == 1 + # (open left boundary). The bug leaves it at Dr0 == 4. + new_first_state_core = ht.list_cores_phi[0][0] + assert new_first_state_core.shape[0] == 1, ( + f'Leftmost group after removal has left bond dim ' + f'{new_first_state_core.shape[0]}, expected 1. ' + f'The |0> slice was discarded instead of absorbed.' + ) + + +# ------------------------------------------------------------ +# TEST: remove then re-add state preserves phi_0 values +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_remove_then_add_state_preserves_phi0(): + # This case tests that after removing state 0 (zero amplitude), the remaining + # phi_0 entries match the original values at indices 1 onward. + ht = _make_tensor('fullstate') + state_list = list(range(nsite)) + phi_0_orig = ht.psi.copy() + # Remove state 0 (which has zero amplitude in psi_0) + ht.remove_state_cores([0], state_list) + phi_0_after_remove = ht.psi + assert len(phi_0_after_remove) == nsite - 1 + np.testing.assert_allclose( + phi_0_after_remove, phi_0_orig[1:], atol=1e-12, + ) + + +# ------------------------------------------------------------ +# TEST: remove all but one state +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_remove_all_but_one(): + # This case tests that removing all states except the one carrying + # amplitude leaves a single-site MPS with non-zero psi[0]. + ht = _make_tensor('fullstate') + state_list = list(range(nsite)) + # Remove states 0, 1, 3 — keep only state 2 (which has the amplitude) + ht.remove_state_cores([0, 1, 3], state_list) + assert ht.list_cores_phi[0].shape[1] == 1 + assert len(ht.M1_modes_per_site) == 1 + assert len(ht.psi) == 1 + # Analytical: psi_0 = [0,0,1,0], keeping only state 2 (which had + # amplitude 1.0), so the single remaining amplitude should be 1.0 + np.testing.assert_allclose( + abs(ht.psi[0]), 1.0, atol=1e-10, + err_msg='Amplitude on kept state should be 1.0', + ) + + +# ------------------------------------------------------------ +# TEST: add_state_cores raises on invalid method +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_add_state_cores_invalid_method_raises(): + # This case tests that add_state_cores raises UnsupportedRequest + # when method is not a recognized representation. + ht = _make_tensor('fullstate') + ht.method = 'bogus' + with pytest.raises(UnsupportedRequest): + ht.add_state_cores([nsite], list(state_list)) + + +# ------------------------------------------------------------ +# TEST: remove_state_cores raises on invalid method +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_remove_state_cores_invalid_method_raises(): + # This case tests that remove_state_cores raises UnsupportedRequest + # when method is not a recognized representation. + ht = _make_tensor('fullstate') + ht.method = 'bogus' + with pytest.raises(UnsupportedRequest): + ht.remove_state_cores([state_list[0]], list(state_list)) + + +# ============================================================ +# TEST SUITE: update_phi_from_flat() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Flatten then update_phi_from_flat recovers original psi +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_update_phi_from_flat_recovers_psi(): + # This case tests that extracting flat cores, then writing them + # back via update_phi_from_flat, recovers the original psi. + # Uses a propagated state with naturally complex cores and + # non-trivial bond structure. + ht_full = _make_propagated_tensor('fullstate') + psi_before = ht_full.psi.copy() + flat = [c.copy() for c in ht_full.list_cores_phi] + ht_full.update_phi_from_flat(flat) + np.testing.assert_allclose(ht_full.psi, psi_before, atol=1e-14) + + # This case tests the same flatten->update->recover cycle for + # statenumber representation with propagated state. + ht_sn = _make_propagated_tensor('number') + psi_before_sn = ht_sn.psi.copy() + flat_sn = [c.copy() for c in ht_sn.flat_cores] + ht_sn.update_phi_from_flat(flat_sn) + np.testing.assert_allclose(ht_sn.psi, psi_before_sn, atol=1e-14) + + +# ------------------------------------------------------------ +# TEST: update_phi_from_flat makes an independent copy +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_update_phi_from_flat_independent_copy(): + # CASE: Mutating the input list after update_phi_from_flat must + # not affect list_cores_phi. Uses propagated state for non-trivial + # bond structure. Mirrors restore_phi independence tests. + for method in ['fullstate', 'number']: + ht = _make_propagated_tensor(method) + flat = [c.copy() for c in ht.flat_cores] + # Scale to distinguish from initial state + flat[0] = flat[0] * 2.0 + ht.update_phi_from_flat(flat) + psi_after_update = ht.psi.copy() + cores_after_update = [c.copy() for c in ht.flat_cores] + # Mutate the source list + flat[0] *= 0.0 + # ht should be unaffected — check both psi and raw cores + np.testing.assert_allclose( + ht.psi, psi_after_update, atol=1e-14, + err_msg=f'Mutating source leaked into psi ({method})', + ) + for i, (c_now, c_saved) in enumerate(zip(ht.flat_cores, cores_after_update)): + np.testing.assert_array_equal( + c_now, c_saved, + err_msg=f'Core {i} mutated by source modification ({method})', + ) + + +# ------------------------------------------------------------ +# TEST: flat_cores returns views into list_cores_phi +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_flat_cores_are_views(): + # CASE: For statenumber representation, flat_cores should return + # views into list_cores_phi. Mutating an entry in flat_cores should + # be visible through list_cores_phi and affect psi. This documents + # the view semantics — flat_cores is NOT a deep copy. + ht = _make_tensor('number') + flat = ht.flat_cores + psi_before = ht.psi.copy() + # Mutate the first flat core (which is list_cores_phi[0][0]) + flat[0][0, :, 0] *= 2.0 + psi_after = ht.psi + # psi should have changed because flat_cores are views + assert not np.allclose(psi_before, psi_after, atol=1e-14), ( + 'Mutating flat_cores should affect psi (view semantics)' + ) diff --git a/tests/test_hops_trajectory.py b/tests/test_hops_trajectory.py index 817e05a..454d334 100644 --- a/tests/test_hops_trajectory.py +++ b/tests/test_hops_trajectory.py @@ -68,6 +68,7 @@ 'INCHWORM_CAP': 5, 'STATIC_BASIS': None, 'EFFECTIVE_NOISE_INTEGRATION': False, + 'STORE_STEP_TIMING': False, } integrator_param_empty = {} integrator_param_partial = { @@ -521,6 +522,56 @@ def test_normalize_else(): assert np.allclose(norm, known_norm) +# ------------------------------------------------------------ +# TEST: single RK4 step keeps norm within 1e-10 before normalize() +# ------------------------------------------------------------ +def test_normalize_single_step_norm_drift(): + """ + Test + ---- + Tests that a single RK4 step preserves the physical wavefunction norm to + within 1e-10 of unity, before normalize() is applied. + + Case + ---- + 2-site system with NORMALIZED NONLINEAR EOM. The norm correction term in + the derivative should keep the dynamics approximately norm-preserving. If + the correction is wrong (e.g. wrong prefactor), the pre-normalization norm + drifts measurably even in a single step. + """ + hops = HOPS( + sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integrator_param, + ) + hops.initialize(psi_0) + tau = 2.0 # 2 * noise TAU so RK4 samples [0, 1, 2] land on the noise grid + + # Gather integration variables (noise, z_mem, etc.) for one step + var_list = hops.integration_var( + hops.phi, hops.z_mem, hops.t, hops.noise1, hops.noise2, + tau, hops.storage, hops.basis.mode.list_l2idx_abs, + hops.effective_noise_integration, + ) + + # Call RK4 directly — bypasses normalize() + phi_raw, _ = hops.step(hops.dsystem_dt, **var_list) + + # Norm of the physical wavefunction (first n_state elements) + n_state = hops.n_state + norm_psi = np.linalg.norm(phi_raw[:n_state]) + + np.testing.assert_allclose( + norm_psi, 1.0, atol=1e-10, + err_msg=( + f'Single-step norm drift {abs(norm_psi - 1.0):.2e} exceeds 1e-10. ' + 'The norm correction term in the EOM derivative may be wrong.' + ), + ) + + def test_inchworm_aux(): """ test for inchworm_integrate to make sure the aux are properly being added @@ -1299,6 +1350,76 @@ def test_propagation_timing(): ) +def test_propagation_step_timing_flag_default_off(): + """ + Checks that with STORE_STEP_TIMING absent (default False), + LIST_PROPAGATION_TIME retains its pre-flag behavior: one float + per propagate() call. + """ + hops = HOPS( + sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integrator_param_empty, + ) + hops.initialize(psi_0) + hops.propagate(4.0, 2.0) + hops.propagate(4.0, 2.0) + + list_prop_time = hops.storage.metadata["LIST_PROPAGATION_TIME"] + assert len(list_prop_time) == 2 + for entry in list_prop_time: + assert isinstance(entry, float) + assert entry >= 0 + + +def test_propagation_step_timing_flag_on(): + """ + Checks that with STORE_STEP_TIMING=True, LIST_PROPAGATION_TIME + holds (t_fs, wall_seconds) tuples — one per integration step, + sim-time stamps match the propagation grid, wall times are + non-negative, and entries concatenate across propagate() calls + with a monotonic time axis. + """ + integration_param = dict(integrator_param_empty) + integration_param["STORE_STEP_TIMING"] = True + hops = HOPS( + sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integration_param, + ) + hops.initialize(psi_0) + + tau = 2.0 + t_first = 4.0 + n_first = int(np.ceil(t_first / tau)) + hops.propagate(t_first, tau) + + list_prop_time = hops.storage.metadata["LIST_PROPAGATION_TIME"] + assert len(list_prop_time) == n_first + for entry in list_prop_time: + assert isinstance(entry, tuple) and len(entry) == 2 + list_t = [t for t, _ in list_prop_time] + list_dt = [dt for _, dt in list_prop_time] + assert all(dt >= 0 for dt in list_dt) + np.testing.assert_allclose(list_t, tau * np.arange(1, n_first + 1)) + + # Second propagate call: entries append, time axis stays monotonic. + t_second = 4.0 + n_second = int(np.ceil(t_second / tau)) + hops.propagate(t_second, tau) + + list_prop_time = hops.storage.metadata["LIST_PROPAGATION_TIME"] + assert len(list_prop_time) == n_first + n_second + list_t = [t for t, _ in list_prop_time] + list_dt = [dt for _, dt in list_prop_time] + assert all(dt >= 0 for dt in list_dt) + assert np.all(np.diff(list_t) > 0) + + def test_operator(): """ Tests management of operation in both non-adaptive and adaptive frameworks. diff --git a/tests/test_integrator_rk.py b/tests/test_integrator_rk.py index 72f1e81..4ef1d0a 100644 --- a/tests/test_integrator_rk.py +++ b/tests/test_integrator_rk.py @@ -1,5 +1,5 @@ import numpy as np -from mesohops.integrator.integrator_rk import runge_kutta_variables +from mesohops.integrator.integrator import runge_kutta_variables from mesohops.noise.hops_noise import HopsNoise from mesohops.trajectory.exp_noise import bcf_exp @@ -42,6 +42,13 @@ "CORR_PARAM": sys_param["PARAM_NOISE1"], } +# ============================================================ +# TEST SUITE: runge_kutta_variables() +# ============================================================ + +# ------------------------------------------------------------ +# TEST: effective noise integration averages noise correctly +# ------------------------------------------------------------ def test_effective_noise_integration(): """ Tests that the effective noise integration that averages the noise over all time diff --git a/tests/test_mpo_constructors.py b/tests/test_mpo_constructors.py new file mode 100644 index 0000000..66488cf --- /dev/null +++ b/tests/test_mpo_constructors.py @@ -0,0 +1,2241 @@ +import numpy as np +import pytest +import scipy.sparse as sparse + +from mesohops.basis.hops_modes import HopsModes +from mesohops.basis.hops_noise_memory import HopsNoiseMemory +from mesohops.basis.hops_system import HopsSystem +from mesohops.tensor.hops_tensor_wavefunction import HopsTensorWavefunction +from mesohops.tensor.hops_tensor_basis import HopsTensorBasis +from mesohops.basis.basis_functions import determine_error_thresh +from mesohops.tensor.mpo_constructors import ( + MpoBuilder, + build_statenumber_dipole_lower_plus_ident_mpo, + build_statenumber_dipole_mpo, + build_statenumber_dipole_raise_plus_ground_ident_mpo, + build_statenumber_operator_mpo, +) +from mesohops.tensor.tensor_eom_functions import tensor_matvec_prod +from mesohops.trajectory.exp_noise import bcf_exp +from mesohops.util.bath_corr_functions import bcf_convert_dl_to_exp +from mesohops.util.tensor_operations import extract_psi, unflatten_cores +from mesohops.util.tensor_operations import tensor_add + +__title__ = 'Unit Tests for MPO Constructors' +__author__ = 'N. Covalsen' +__maintainer__ = 'N. Covalsen' + +# ============================================================ +# Shared Setup +# ============================================================ + +nsite = 4 +e_lambda = 20.0 +gamma = 50.0 +temp = 140.0 +(g_0, w_0) = bcf_convert_dl_to_exp(e_lambda, gamma, temp) + +loperator = np.zeros([4, 4, 4], dtype=np.float64) +gw_sysbath = [] +lop_list = [] +for i in range(nsite): + loperator[i, i, i] = 1.0 + gw_sysbath.append([g_0, w_0]) + lop_list.append(sparse.coo_matrix(loperator[i])) + gw_sysbath.append([-1j * np.imag(g_0), 500.0]) + lop_list.append(loperator[i]) + +hs = np.zeros([nsite, nsite]) +hs[0, 1] = 40 +hs[1, 0] = 40 +hs[1, 2] = 10 +hs[2, 1] = 10 +hs[2, 3] = 40 +hs[3, 2] = 40 + +sys_param = { + 'HAMILTONIAN': np.array(hs, dtype=np.complex128), + 'GW_SYSBATH': gw_sysbath, + 'L_HIER': lop_list, + 'L_NOISE1': lop_list, + 'ALPHA_NOISE1': bcf_exp, + 'PARAM_NOISE1': gw_sysbath, +} + +psi_0 = np.array([0.0] * nsite, dtype=np.complex128) +psi_0[2] = 1.0 + +k_max = 2 +state_list = np.arange(nsite) +delta_s = 0 +eom_param = {'EQUATION_OF_MOTION': 'NORMALIZED NONLINEAR'} + + +def _make_tb(sp=sys_param, ds=delta_s, psi=psi_0, sl=state_list): + """Creates an initialized HopsTensorBasis from sys_param dict.""" + system = HopsSystem(sp) + mode = HopsModes(system, hierarchy=None) + noise_memory = HopsNoiseMemory(system, mode) + tb = HopsTensorBasis(system, mode, noise_memory) + system.initialize(ds > 0, psi) + mode.list_modeidx_abs = sorted(system.list_statemodeidx_abs) + noise_memory.initialize() + system.state_list = sl + tb.initialize(ds) + return tb + + +def _make_tensor_pair(method): + """Returns (HopsTensorWavefunction, HopsTensorBasis), both initialized.""" + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': method, + 'BOND_DIM_MAX': 20, + } + integrator_param = {'INTEGRATOR': 'RUNGE_KUTTA'} + tb = _make_tb() + ht = HopsTensorWavefunction(k_max, tensor_param, integrator_param, eom_param) + ht.initialize(psi_0, tb.system) + return ht, tb + + +def _apply_ham_mpo_and_extract(ht, tb, list_cores_ham): + """Apply a Hamiltonian MPO to the MPS and extract the physical wavefunction. + + Pads the Hamiltonian MPO with identity mode cores to match the + full MPS length, applies via tensor_matvec_prod, then extracts + phi_0 from the result. + """ + # Build identity mode cores matching MPS mode core dimensions + list_cores_full_mpo = [] + mps_cores_flat = ht.flat_cores + ham_idx = 0 + for idx, mps_core in enumerate(mps_cores_flat): + if ham_idx < len(list_cores_ham): + ham_core = list_cores_ham[ham_idx] + # Check if this ham core's physical dims match the MPS core + if ham_core.shape[1] == mps_core.shape[1]: + list_cores_full_mpo.append(ham_core) + ham_idx += 1 + continue + # Identity mode core: (1, d, d, 1) with I on the diagonal + d = mps_core.shape[1] + identity_core = np.zeros((1, d, d, 1), dtype=np.complex128) + for k in range(d): + identity_core[0, k, k, 0] = 1.0 + list_cores_full_mpo.append(identity_core) + # Append remaining ham cores if any + while ham_idx < len(list_cores_ham): + list_cores_full_mpo.append(list_cores_ham[ham_idx]) + ham_idx += 1 + + result_cores_flat, _ = tensor_matvec_prod( + mps_cores_flat, + list_cores_full_mpo, + ht.mps_epsilon, + ht.bond_dim_max, + ) + # For statenumber, phi_0 expects list-of-lists; restore structure + if ht.method == 'number': + result_cores = unflatten_cores(result_cores_flat, ht.M1_modes_per_state) + else: + result_cores = result_cores_flat + return extract_psi(result_cores, ht.method, ht.M1_modes_per_state) + + +# ============================================================ +# TEST SUITE: MpoBuilder.__init__() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: homps normalization produces correct ladder operators +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_elementary_operators_homps_ladder(): + # This case tests that homps normalization produces the correct + # ladder operators B2_lower, B2_raise, and occupation number N2_occ. + k_max = 2 + + # build a minimal mock mode + class MockMode: + list_g = np.array([1.0 + 0j, 2.0 + 0j]) + list_w = np.array([10.0, 20.0]) + n_l2 = 1 + n_hmodes = 2 + list_L2_coo = [] + list_L2_masks = [[[0], [0], None]] + list_index_L2_by_hmode = [0, 0] + + mode = MockMode() + elem = MpoBuilder( + k_max, + 2, + 1, + 1, + np.eye(2, dtype=np.complex128), + np.array([0, 1]), + mode, + 'homps', + ) + # B2_lower[i][i+1] = sqrt(i+1) for i in range(k_max) + assert elem.B2_lower[0, 1] == pytest.approx(1.0) + assert elem.B2_lower[1, 2] == pytest.approx(np.sqrt(2)) + # B2_raise is transpose of B2_lower for homps + assert elem.B2_raise[1, 0] == pytest.approx(1.0) + assert elem.B2_raise[2, 1] == pytest.approx(np.sqrt(2)) + # Analytical: N2_occ is the number operator, N2_occ[n,n] = n + assert elem.N2_occ[0, 0] == pytest.approx(0.0) + assert elem.N2_occ[1, 1] == pytest.approx(1.0) + assert elem.N2_occ[2, 2] == pytest.approx(2.0) + # This case tests that off-diagonal elements are zero + assert elem.N2_occ[0, 1] == pytest.approx(0.0) + assert elem.N2_occ[1, 0] == pytest.approx(0.0) + + +# ------------------------------------------------------------ +# TEST: adhops normalization produces correct ladder operators +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_elementary_operators_adhops_ladder(): + # This case tests that adhops normalization produces the correct + # unit ladder operators B2_lower, B2_raise, diagonal N2_occ, and + # coupling vectors C1_coupling_raise = list_w and C1_coupling_lower = list_g / list_w. + k_max = 2 + + class MockMode: + list_g = np.array([1.0 + 0j, 2.0 + 0j]) + list_w = np.array([10.0, 20.0]) + n_l2 = 1 + n_hmodes = 2 + list_L2_coo = [] + list_L2_masks = [[[0], [0], None]] + list_index_L2_by_hmode = [0, 0] + + mode = MockMode() + elem = MpoBuilder( + k_max, + 2, + 1, + 1, + np.eye(2, dtype=np.complex128), + np.array([0, 1]), + mode, + 'adhops', + ) + # adhops: B2_lower[i][i+1] = 1 (unit) + assert elem.B2_lower[0, 1] == pytest.approx(1.0) + assert elem.B2_lower[1, 2] == pytest.approx(1.0) + # N2_occ[i+1][i+1] = i+1 (explicit diagonal) + assert elem.N2_occ[1, 1] == pytest.approx(1.0) + assert elem.N2_occ[2, 2] == pytest.approx(2.0) + # C1_coupling_raise = list_w, C1_coupling_lower = list_g / list_w + np.testing.assert_allclose(elem.C1_coupling_raise, mode.list_w) + np.testing.assert_allclose(elem.C1_coupling_lower, mode.list_g / mode.list_w) + + +# ------------------------------------------------------------ +# TEST: Invalid normalization raises ValueError +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_mpo_builder_invalid_normalization_raises(): + class MockMode: + list_g = np.array([1.0 + 0j]) + list_w = np.array([10.0]) + n_l2 = 1 + n_hmodes = 1 + list_L2_coo = [] + list_L2_masks = [[[0], [0], None]] + list_index_L2_by_hmode = [0] + + with pytest.raises(ValueError, match='Unknown normalization'): + MpoBuilder( + 2, 2, 1, 1, np.eye(2, dtype=np.complex128), + np.array([0, 1]), MockMode(), 'bogus', + ) + + +# ------------------------------------------------------------ +# TEST: Scalar and derived attributes stored correctly +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_mpo_builder_stores_scalars(): + ht, tb = _make_tensor_pair('fullstate') + builder = MpoBuilder( + k_max=k_max, + n_state=tb.system.size, + modes_per_state=ht.M1_modes_per_state, + n_lop_full=tb.mode.n_l2, + ham=tb.system.param['HAMILTONIAN'], + state_list=tb.system.state_list, + mode=tb.mode, + normalization='homps', + ) + assert builder.k_max == k_max + assert builder.n_state == tb.system.size + assert builder.n_lop_full == tb.mode.n_l2 + np.testing.assert_array_equal( + builder.M1_modes_per_state, ht.M1_modes_per_state, + ) + # M1_mode_offset is cumulative sum with leading zero + expected_offset = np.concatenate( + [[0], np.cumsum(ht.M1_modes_per_state)] + ) + np.testing.assert_array_equal(builder.M1_mode_offset, expected_offset) + + +# ------------------------------------------------------------ +# TEST: State-dimension operators have correct shape and values +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_mpo_builder_state_operators(): + class MockMode: + list_g = np.array([1.0 + 0j]) + list_w = np.array([10.0]) + n_l2 = 1 + n_hmodes = 1 + list_L2_coo = [] + list_L2_masks = [[[0], [0], None]] + list_index_L2_by_hmode = [0] + + builder = MpoBuilder( + 2, 2, 1, 1, np.eye(2, dtype=np.complex128), + np.array([0, 1]), MockMode(), 'homps', + ) + # All state operators are (1, 2, 2, 1) + for name in ['T4_plus', 'T4_minus', 'Q4_site', 'P4_site']: + op = getattr(builder, name) + assert op.shape == (1, 2, 2, 1), f'{name} shape should be (1,2,2,1)' + # P = |1><1|, Q = |0><0|, T_left = |1><0|, T_right = |0><1| + np.testing.assert_array_equal( + builder.P4_site.reshape(2, 2), np.array([[0, 0], [0, 1]]), + ) + np.testing.assert_array_equal( + builder.Q4_site.reshape(2, 2), np.array([[1, 0], [0, 0]]), + ) + np.testing.assert_array_equal( + builder.T4_plus.reshape(2, 2), np.array([[0, 0], [1, 0]]), + ) + np.testing.assert_array_equal( + builder.T4_minus.reshape(2, 2), np.array([[0, 1], [0, 0]]), + ) + # I2_mode is identity of size k_max + 1 + np.testing.assert_array_equal( + builder.I2_mode, np.eye(3, dtype=np.complex128), + ) + + +# ============================================================ +# TEST SUITE: _build_statenumber_ham_general_mpo() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Ham general MPO has correct bond dimension +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_statenumber_ham_general_mpo_shape(): + # This case tests that the general Hamiltonian MPO bond dimension + # equals 4 + 2*(n_state-2) for a non-nearest-neighbor Hamiltonian. + # Build a 3-state non-NN system: state 0 couples to state 2 (skips 1). + n = 3 + ham_gen = np.zeros((n, n), dtype=np.complex128) + ham_gen[0, 2] = 5.0 + ham_gen[2, 0] = 5.0 + # Site-diagonal L-operators (projectors onto each site) + lop_gen = np.zeros((n, n, n), dtype=np.float64) + lop_list_gen = [] + for i in range(n): + lop_gen[i, i, i] = 1.0 + lop_list_gen.append(sparse.coo_matrix(lop_gen[i])) + gw_gen = [[g_0, w_0]] * n + sys_param_gen = { + 'HAMILTONIAN': ham_gen, + 'GW_SYSBATH': gw_gen, + 'L_HIER': lop_list_gen, + 'L_NOISE1': lop_list_gen, + 'ALPHA_NOISE1': bcf_exp, + 'PARAM_NOISE1': gw_gen, + } + tensor_param = { + 'METHOD': 'number', + 'MPS_EPSILON': 1e-10, + 'BOND_DIM_MAX': 20, + } + integrator_param = {'INTEGRATOR': 'RUNGE_KUTTA'} + psi_gen = np.array([1.0, 0.0, 0.0], dtype=np.complex128) + state_list_gen = np.arange(n) + tb = _make_tb(sp=sys_param_gen, ds=0, psi=psi_gen, sl=state_list_gen) + ht = HopsTensorWavefunction(k_max, tensor_param, integrator_param, eom_param) + ht.initialize(psi_gen, tb.system) + elem = MpoBuilder( + ht.k_max, + tb.system.size, + ht.M1_modes_per_state, + tb.mode.n_l2, + tb.system.param['HAMILTONIAN'], + tb.system.state_list, + tb.mode, + 'homps', + n_states_full=n, + flag_nearest_neighbor_ham=False, + ) + cores = elem._build_statenumber_ham_general_mpo() + bonddim = int(4 + 2 * (n - 2)) + assert cores[0].shape[3] == bonddim + # This case tests that the MPO can be contracted without producing NaN + flat_cores = [ + c for group in ht.list_cores_phi + for c in (group if isinstance(group, list) else [group]) + ] + result_cores, _ = tensor_matvec_prod(flat_cores, cores, 1e-10, 50) + for c in result_cores: + assert not np.any(np.isnan(c)), 'MPO contraction produced NaN' + + +# ------------------------------------------------------------ +# TEST: General Hamiltonian MPO gives H @ psi (statenumber) +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_statenumber_ham_general_mpo_value(): + # This case tests the general (non-NN) Hamiltonian MPO builder. + # Uses the same 4-site system — build_statenumber_ham_general_mpo + # should produce the same result as the NN builder for NN Hamiltonians. + ht, tb = _make_tensor_pair('number') + elem = MpoBuilder( + k_max, + tb.system.size, + ht.M1_modes_per_state, + tb.mode.n_l2, + tb.system.param['HAMILTONIAN'], + tb.system.state_list, + tb.mode, + 'homps', + n_states_full=tb.system.param['NSTATES'], + flag_nearest_neighbor_ham=False, + ) + list_cores_ham = elem._build_statenumber_ham_general_mpo() + psi_result = _apply_ham_mpo_and_extract(ht, tb, list_cores_ham) + psi_input = extract_psi( + ht.list_cores_phi, + ht.method, + ht.M1_modes_per_state, + ) + psi_expected = hs @ psi_input + np.testing.assert_allclose( + psi_result, + psi_expected, + atol=1e-8, + err_msg='General Hamiltonian MPO does not match H @ psi', + ) + + +# ------------------------------------------------------------ +# TEST: General ham MPO handles n_state=1 without crashing +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_statenumber_ham_general_single_site(): + # This case tests that the general Hamiltonian MPO builder handles + # a single-site system (n_state=1) without crashing. With one site + # there are no off-diagonal terms — only the diagonal energy. + E = 3.5 + ham_1 = np.array([[E]], dtype=np.complex128) + k_max_1 = 2 + n_modes = 1 + + class MockMode: + list_g = np.array([1.0 + 0j]) + list_w = np.array([10.0]) + n_l2 = 1 + n_hmodes = 1 + list_L2_coo = [sparse.coo_matrix(np.array([[1.0]]))] + list_L2_masks = [[[0], [0], None]] + list_index_L2_by_hmode = [0] + + mode = MockMode() + elem = MpoBuilder( + k_max_1, + 1, + np.array([n_modes]), + 1, + ham_1, + np.array([0]), + mode, + 'homps', + n_states_full=1, + flag_nearest_neighbor_ham=False, + ) + cores = elem.build_statenumber_ham_mpo() + + # Should produce 1 state core + n_modes identity mode cores = 2 cores + assert len(cores) == 1 + n_modes + + # State core shape: (1, 2, 2, 1) — single bond on each side + state_core = cores[0] + assert state_core.shape[0] == 1 + assert state_core.shape[3] == 1 + + # The occupied state should carry the diagonal energy E + # state_core[0, 1, 1, 0] = E (from H[0,0] * P4_site[0,1,1,0]) + assert state_core[0, 1, 1, 0] == pytest.approx(E) + # The unoccupied state should be identity pass-through + # state_core[0, 0, 0, 0] = 1.0 (from Q4_site[0,0,0,0]) + assert state_core[0, 0, 0, 0] == pytest.approx(1.0) + + +# ------------------------------------------------------------ +# TEST: General MPO with non-NN Ham reproduces H @ psi via contraction +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_statenumber_ham_general_mpo_nonnn_value(): + # Analytical: MPO applied to MPS should give H @ psi for a + # Hamiltonian with long-range coupling H[0,2] = 5.0 (skips site 1). + # The nearest-neighbor builder cannot represent this coupling, so this + # test specifically exercises the general MPO builder. + n = 3 + ham_gen = np.zeros((n, n), dtype=np.complex128) + ham_gen[0, 1] = 2.0 # NN coupling (covered by NN builder too) + ham_gen[1, 0] = 2.0 + ham_gen[0, 2] = 5.0 # Long-range coupling — requires general builder + ham_gen[2, 0] = 5.0 + ham_gen[1, 1] = 1.0 # On-site energies + ham_gen[2, 2] = 3.0 + + # Site-diagonal L-operators (projectors onto each site) + lop_gen = np.zeros((n, n, n), dtype=np.float64) + lop_list_gen = [] + for i in range(n): + lop_gen[i, i, i] = 1.0 + lop_list_gen.append(sparse.coo_matrix(lop_gen[i])) + gw_gen = [[g_0, w_0]] * n + sys_param_gen = { + 'HAMILTONIAN': ham_gen, + 'GW_SYSBATH': gw_gen, + 'L_HIER': lop_list_gen, + 'L_NOISE1': lop_list_gen, + 'ALPHA_NOISE1': bcf_exp, + 'PARAM_NOISE1': gw_gen, + } + tensor_param = { + 'METHOD': 'number', + 'MPS_EPSILON': 1e-10, + 'BOND_DIM_MAX': 20, + } + integrator_param = {'INTEGRATOR': 'RUNGE_KUTTA'} + # Delocalized initial state with complex phases so that all Hamiltonian + # matrix elements contribute non-trivially to H @ psi. + psi_gen = np.array([0.5 + 0.3j, -0.4 + 0.2j, 0.6 - 0.1j], dtype=np.complex128) + psi_gen /= np.linalg.norm(psi_gen) + state_list_gen = np.arange(n) + + tb = _make_tb(sp=sys_param_gen, ds=0, psi=psi_gen, sl=state_list_gen) + ht = HopsTensorWavefunction(k_max, tensor_param, integrator_param, eom_param) + ht.initialize(psi_gen, tb.system) + + elem = MpoBuilder( + ht.k_max, + tb.system.size, + ht.M1_modes_per_state, + tb.mode.n_l2, + tb.system.param['HAMILTONIAN'], + tb.system.state_list, + tb.mode, + 'homps', + n_states_full=n, + flag_nearest_neighbor_ham=False, + ) + list_cores_ham = elem._build_statenumber_ham_general_mpo() + + psi_result = _apply_ham_mpo_and_extract(ht, tb, list_cores_ham) + # Use psi_gen directly — avoids extract_psi dependency and matches + # the known input state (no SVD truncation at bond dim 1). + psi_expected = ham_gen @ psi_gen + + np.testing.assert_allclose( + psi_result, + psi_expected, + atol=1e-8, + err_msg='General MPO with non-NN coupling does not match H @ psi', + ) + + +# ============================================================ +# TEST SUITE: MpoBuilder.build_elementary_ops() +# ============================================================ + + +def _mock_mode(list_g=None, list_w=None): + """Helper: build a minimal mock mode object.""" + + class MockMode: + pass + + m = MockMode() + m.list_g = np.array(list_g or [1.0 + 0j, 2.0 + 0j]) + m.list_w = np.array(list_w or [10.0, 20.0]) + m.n_l2 = 1 + m.n_hmodes = len(m.list_w) + m.list_L2_coo = [] + # Rows/cols/ix_ per L2 in the active basis; [0][0] is the L2's site. + m.list_L2_masks = [[[i], [i], None] for i in range(m.n_l2)] + m.list_index_L2_by_hmode = [0] * m.n_hmodes + return m + + +def _make_elem(k_max_val=2, normalization='homps', mode=None): + """Helper: build and return an MpoBuilder instance.""" + if mode is None: + mode = _mock_mode() + return MpoBuilder( + k_max_val, + 2, + 1, + 1, + np.eye(2, dtype=np.complex128), + np.array([0, 1]), + mode, + normalization, + n_states_full=2, + flag_nearest_neighbor_ham=True, + ) + + +# ------------------------------------------------------------ +# TEST: Invalid normalization raises ValueError +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_elementary_ops_invalid_normalization(): + # This case tests that an unknown normalization string raises ValueError. + mode = _mock_mode() + with pytest.raises(ValueError, match='Unknown normalization'): + MpoBuilder( + 2, + 2, + 1, + 1, + np.eye(2, dtype=np.complex128), + np.array([0, 1]), + mode, + 'invalid', + n_states_full=2, + flag_nearest_neighbor_ham=True, + ) + + +# ------------------------------------------------------------ +# TEST: State-dimension operators have correct values +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_elementary_ops_state_dimension_operators(): + # This case tests P4_site, Q4_site, T4_plus, T4_minus, I2_state. + elem = _make_elem() + + # P4_site = |1><1| projector in 4-D core shape + expected_P4 = np.array([[0, 0], [0, 1]], dtype=np.complex128).reshape(1, 2, 2, 1) + np.testing.assert_allclose( + elem.P4_site, expected_P4, atol=1e-14, err_msg='P4_site should be |1><1|' + ) + + # Q4_site = |0><0| projector in 4-D core shape + expected_Q4 = np.array([[1, 0], [0, 0]], dtype=np.complex128).reshape(1, 2, 2, 1) + np.testing.assert_allclose( + elem.Q4_site, expected_Q4, atol=1e-14, err_msg='Q4_site should be |0><0|', + ) + + # T4_plus = |1><0| in 4-D core shape + expected_T4_plus = np.array([[0, 0], [1, 0]], dtype=np.complex128).reshape(1, 2, 2, 1) + np.testing.assert_allclose( + elem.T4_plus, expected_T4_plus, atol=1e-14, err_msg='T4_plus should have [1,0]=1', + ) + + # T4_minus = |0><1| in 4-D core shape + expected_T4_minus = np.array([[0, 1], [0, 0]], dtype=np.complex128).reshape(1, 2, 2, 1) + np.testing.assert_allclose( + elem.T4_minus, expected_T4_minus, atol=1e-14, err_msg='T4_minus should have [0,1]=1', + ) + + # I2_state is identity of size n_state=2 + np.testing.assert_allclose( + elem.I2_state, + np.eye(2, dtype=np.complex128), + atol=1e-14, + err_msg='I2_state should be identity', + ) + + +# ------------------------------------------------------------ +# TEST: k_max=0 produces valid degenerate operators +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_elementary_ops_k_max_zero(): + # This case tests that k_max=0 produces (1,1) mode operators + # with no ladder steps. + elem = _make_elem(k_max_val=0) + + # All mode operators should be (1, 1) + assert elem.B2_raise.shape == (1, 1) + assert elem.B2_lower.shape == (1, 1) + assert elem.N2_occ.shape == (1, 1) + assert elem.I2_mode.shape == (1, 1) + + # Ladder operators should be zero (no transitions possible) + np.testing.assert_allclose( + elem.B2_raise, [[0]], atol=1e-14, err_msg='B2_raise should be zero for k_max=0' + ) + np.testing.assert_allclose( + elem.B2_lower, [[0]], atol=1e-14, err_msg='B2_lower should be zero for k_max=0' + ) + np.testing.assert_allclose( + elem.N2_occ, [[0]], atol=1e-14, err_msg='N2_occ should be zero for k_max=0' + ) + + # Identity should still be [[1]] + np.testing.assert_allclose( + elem.I2_mode, [[1]], atol=1e-14, err_msg='I2_mode should be [[1]] for k_max=0' + ) + + +# ------------------------------------------------------------ +# TEST: homps C1_coupling_raise and C1_coupling_lower have correct values +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_elementary_ops_homps_coupling_vectors(): + # Convention of Gao et al., Phys. Rev. A 105, L030202 (2022), Eq. (10): + # C1_coupling_raise pairs with B2_raise (b†): V_m^+ = g/sqrt(|g|) + # C1_coupling_lower pairs with B2_lower (b): V_m^- = sqrt(|g|) + # A complex g is required: for real positive g the two orderings coincide. + list_g = [3.0 + 0j, 4.0 + 1j] + mode = _mock_mode(list_g=list_g) + elem = _make_elem(normalization='homps', mode=mode) + + for i, g in enumerate(list_g): + expected_raise = g / np.sqrt(np.abs(g)) + expected_lower = np.sqrt(np.abs(g)) + np.testing.assert_allclose( + elem.C1_coupling_raise[i], + expected_raise, + atol=1e-12, + err_msg=f'C1_coupling_raise[{i}] should be g/sqrt(|g|) (V_m^+, pairs with b†)', + ) + np.testing.assert_allclose( + elem.C1_coupling_lower[i], + expected_lower, + atol=1e-12, + err_msg=f'C1_coupling_lower[{i}] should be sqrt(|g|) (V_m^-, pairs with b)', + ) + # The split must reproduce the physical coupling either way round. + np.testing.assert_allclose( + elem.C1_coupling_raise[i] * elem.C1_coupling_lower[i], + g, + atol=1e-12, + err_msg=f'V_m^+ * V_m^- should equal g for mode {i}', + ) + + +# ------------------------------------------------------------ +# TEST: Core entries match hand-computed values for minimal system +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_statenumber_hierarchy_mpo_values_minimal(): + # This case tests that the statenumber hierarchy MPO core entries + # match hand-computed values for a 1-state, 1-mode, k_max=1 system + # with zero noise, using adhops normalization with known g and w. + # Setup: adhops with known g, w so entries are hand-computable + g, w = 2.0 + 0j, 10.0 + mode = _mock_mode(list_g=[g], list_w=[w]) + mode.n_hmodes = 1 + elem = _make_elem(k_max_val=1, normalization='adhops', mode=mode) + # Override n_state and modes_per_state for 1-site system + elem.n_state = 1 + elem.M1_modes_per_state = np.array([1]) + + z_conj_t = np.zeros(1, dtype=np.complex128) + L_conj_avg = np.zeros(1, dtype=np.complex128) + cores = elem.build_statenumber_hierarchy_mpo( + z_conj_t, + L_conj_avg, + 0.0, + ) + assert len(cores) == 2 + + # --- Core 0: state site (1, 2, 2, 5) --- + site = cores[0] + # This case tests the [0,:,:,0] block is Q4_site = |0><0|. + np.testing.assert_allclose( + site[0, :, :, 0], + np.array([[1, 0], [0, 0]], dtype=np.complex128), + atol=1e-14, + err_msg='[0,:,:,0] should be |0><0|', + ) + # This case tests [0,:,:,1] and [0,:,:,2] are 1j * |1><1|. + expected_proj = 1j * np.array([[0, 0], [0, 1]], dtype=np.complex128) + np.testing.assert_allclose( + site[0, :, :, 1], + expected_proj, + atol=1e-14, + err_msg='[0,:,:,1] should be 1j * |1><1|', + ) + # Channel 2 (damp1) at site 0 is 1j * |1><1| in the default (non-vacuum) + # convention; the site-0 widening to 1j * I_site only fires when + # flag_gs_vacuum is set (see the gated-widening test below). + np.testing.assert_allclose( + site[0, :, :, 2], + expected_proj, + atol=1e-14, + err_msg='[0,:,:,2] at site 0 should be 1j * |1><1| (non-vacuum)', + ) + # This case tests remaining bond slots [0,:,:,3] and [0,:,:,4] are zero. + np.testing.assert_allclose( + site[0, :, :, 3], + 0.0, + atol=1e-14, + err_msg='[0,:,:,3] should be zero', + ) + np.testing.assert_allclose( + site[0, :, :, 4], + 0.0, + atol=1e-14, + err_msg='[0,:,:,4] should be zero', + ) + + # --- Core 1: mode site (5, 2, 2, 1), last mode --- + mode_core = cores[1] + # adhops k_max=1: B2_raise=[[0,0],[1,0]], B2_lower=[[0,1],[0,0]] + # C1_coupling_raise=w=10, C1_coupling_lower=g/w=0.2 + # This case tests [1,:,:,0] = C1_coupling_raise*B2_raise - C1_coupling_lower*B2_lower. + expected_coupling = np.array( + [[0, -0.2], [10, 0]], + dtype=np.complex128, + ) + np.testing.assert_allclose( + mode_core[1, :, :, 0], + expected_coupling, + atol=1e-14, + err_msg='[1,:,:,0] coupling block incorrect', + ) + # This case tests [4,:,:,0] = identity (pass-through). + np.testing.assert_allclose( + mode_core[4, :, :, 0], + np.eye(2, dtype=np.complex128), + atol=1e-14, + err_msg='[4,:,:,0] should be identity', + ) + # This case tests [2,:,:,0] = -w * N2_occ = -10 * diag(0,1). + expected_damping = np.array( + [[0, 0], [0, -10]], + dtype=np.complex128, + ) + np.testing.assert_allclose( + mode_core[2, :, :, 0], + expected_damping, + atol=1e-14, + err_msg='[2,:,:,0] damping block incorrect', + ) + # This case tests remaining bond slots [0,:,:,0] and [3,:,:,0] are zero. + np.testing.assert_allclose( + mode_core[0, :, :, 0], + 0.0, + atol=1e-14, + err_msg='[0,:,:,0] should be zero for last mode', + ) + np.testing.assert_allclose( + mode_core[3, :, :, 0], + 0.0, + atol=1e-14, + err_msg='[3,:,:,0] should be zero for last mode', + ) + + +# ------------------------------------------------------------ +# TEST: site-0 damp1 widening is a no-op on the single-excitation manifold +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_statenumber_hierarchy_mpo_site0_widening_noop_on_manifold(): + # The site-0 damp1 widening (gated by flag_gs_vacuum) opens the + # vacuum-drift channel on the electronic-ground sector and must be a + # bit-exact no-op on the single-excitation manifold. This builds a + # 2-site / 2-mode hierarchy MPO (so the interior state core also fires), + # applies it to a manifold wavefunction via the generic + # vector -> MPS -> apply -> vector path, and checks the widened and + # un-widened actions agree on the single-excitation inputs while the + # ground-sector action differs. + layout = _dipole_layout(2, np.array([1, 1]), 1) + + def _hierarchy_mpo(flag_gs_vacuum): + mode = _mock_mode(list_g=[1.5 + 0j, 0.8 + 0j], list_w=[5.0, 15.0]) + mode.n_hmodes = 2 + mode.n_l2 = 2 + mode.list_L2_masks = [[[0], [0], None], [[1], [1], None]] + elem = MpoBuilder( + 1, 2, np.array([1, 1]), 2, np.eye(2, dtype=np.complex128), + np.array([0, 1]), mode, 'adhops', + n_states_full=2, + flag_nearest_neighbor_ham=True, + flag_gs_vacuum=flag_gs_vacuum, + ) + z_conj_t = np.array([0.2 + 0.1j, -0.1 + 0.3j], dtype=np.complex128) + L_conj_avg = np.array([0.3 - 0.05j, 0.4 + 0.2j], dtype=np.complex128) + return elem.build_statenumber_hierarchy_mpo(z_conj_t, L_conj_avg, 0.5) + + def _act(V2_phi, flag): + flat_cores = _vector_to_mps_flat(V2_phi, layout) + result_flat, _ = tensor_matvec_prod( + flat_cores, _hierarchy_mpo(flag), 1e-12, 64, + ) + return _mps_flat_to_vector(result_flat, layout) + + # Manifold vector form phi[state, aux]: rows = [ground, e_site0, e_site1], + # cols = [k=vac, k=mode0, k=mode1]. + rng = np.random.default_rng(0) + V2_phi = rng.standard_normal((3, 3)) + 1j * rng.standard_normal((3, 3)) + + # Single-excitation manifold input: zero the electronic-ground row. + V2_single = V2_phi.copy() + V2_single[0, :] = 0.0 + np.testing.assert_allclose( + _act(V2_single, True), _act(V2_single, False), atol=1e-12, + err_msg='site-0 widening must be a no-op on the single-excitation ' + 'manifold', + ) + + # The widening is live: with the ground sector populated (the + # Phi[vac, k=e_n] drift configurations) the action differs on/off. + assert not np.allclose( + _act(V2_phi, True), _act(V2_phi, False), atol=1e-12 + ), ( + 'site-0 widening should alter the ground-sector (vacuum-drift) action' + ) + + +# ============================================================ +# TEST SUITE: build_statenumber_ham_mpo() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: NN Hamiltonian MPO gives H @ psi (statenumber) +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_statenumber_ham_nn_mpo_value(): + # This case tests that the nearest-neighbor Hamiltonian MPO, + # when applied to the MPS, produces H @ psi. Uses a delocalized + # psi with complex phases so all four components contribute to H @ psi, + # and a negative coupling (-10) on bond 1-2 to verify sign handling. + # + # The input psi is used directly for the expected value — we compare + # psi_result against H_neg @ psi_test without going through extract_psi + # on the input MPS, which would only recover the normalized vector. + psi_test = np.array([0.3 + 0.1j, -0.5, 0.2 - 0.4j, 0.6 + 0.2j], dtype=np.complex128) + psi_test = psi_test / np.linalg.norm(psi_test) + + # Build Hamiltonian with negative coupling on bond 1-2 + H2_ham_neg = np.zeros([nsite, nsite], dtype=np.complex128) + H2_ham_neg[0, 1] = 40 + H2_ham_neg[1, 0] = 40 + H2_ham_neg[1, 2] = -10 + H2_ham_neg[2, 1] = -10 + H2_ham_neg[2, 3] = 40 + H2_ham_neg[3, 2] = 40 + + sys_param_neg = dict(sys_param) + sys_param_neg['HAMILTONIAN'] = H2_ham_neg + + tb = _make_tb(sp=sys_param_neg, psi=psi_test) + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'number', + 'BOND_DIM_MAX': 20, + } + integrator_param = {'INTEGRATOR': 'RUNGE_KUTTA'} + ht = HopsTensorWavefunction(k_max, tensor_param, integrator_param, eom_param) + ht.initialize(psi_test, tb.system) + + elem = MpoBuilder( + k_max, + tb.system.size, + ht.M1_modes_per_state, + tb.mode.n_l2, + tb.system.param['HAMILTONIAN'], + tb.system.state_list, + tb.mode, + 'homps', + n_states_full=tb.system.param['NSTATES'], + flag_nearest_neighbor_ham=True, + ) + list_cores_ham = elem.build_statenumber_ham_mpo() + + # This case tests interior state core shapes. + # For 4 sites with 2 modes each: state0(0), mode0a(1), mode0b(2), + # state1(3), mode1a(4), mode1b(5), state2(6), ..., state3(9), ... + # Interior state cores (sites 1, 2) should be (4, 2, 2, 4). + modes_per_state = ht.M1_modes_per_state[0] # 2 modes per state + interior_state_core_indices = [ + 1 + modes_per_state, # state1 core index + 1 + modes_per_state + (1 + modes_per_state), # state2 core index + ] + for core_idx in interior_state_core_indices: + assert list_cores_ham[core_idx].shape == (4, 2, 2, 4), ( + f'Interior state core at index {core_idx} should be (4, 2, 2, 4), ' + f'got {list_cores_ham[core_idx].shape}' + ) + + # Apply MPO to MPS and extract result + psi_result = _apply_ham_mpo_and_extract(ht, tb, list_cores_ham) + # Compare against H_neg @ psi_test directly (not via extract_psi on input MPS) + psi_expected = H2_ham_neg @ psi_test + np.testing.assert_allclose( + psi_result, + psi_expected, + atol=1e-8, + err_msg='NN Hamiltonian MPO does not match H @ psi', + ) + + +# ------------------------------------------------------------ +# TEST: NN Ham MPO mode cores are block-diagonal identity +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_nn_ham_mpo_mode_cores_identity(): + # This case tests that mode cores in the NN Hamiltonian MPO are + # block-diagonal identity: for each bond index j, + # T4_core_mode[j, :, :, j] == eye(k_max+1) and all off-diagonal + # bond blocks are zero. This reflects that the Hamiltonian acts only + # on state cores and passes the hierarchy dimension through unchanged. + ht, tb = _make_tensor_pair('number') + elem = MpoBuilder( + k_max, + tb.system.size, + ht.M1_modes_per_state, + tb.mode.n_l2, + tb.system.param['HAMILTONIAN'], + tb.system.state_list, + tb.mode, + 'homps', + n_states_full=tb.system.param['NSTATES'], + flag_nearest_neighbor_ham=True, + ) + list_cores_ham = elem.build_statenumber_ham_mpo() + + # Core layout: state0, mode0a, mode0b, state1, mode1a, mode1b, ... + # Identify mode core indices by skipping state cores (one per site). + I2_mode = np.eye(k_max + 1, dtype=np.complex128) + core_idx = 0 + for site in range(nsite): + # Skip the state core for this site + core_idx += 1 + n_modes = ht.M1_modes_per_state[site] + for _ in range(n_modes): + T4_core = list_cores_ham[core_idx] + bond_dim = T4_core.shape[0] + assert T4_core.shape == (bond_dim, k_max + 1, k_max + 1, bond_dim), ( + f'Mode core at index {core_idx} has unexpected shape {T4_core.shape}' + ) + # This case tests that diagonal bond blocks are identity + for j in range(bond_dim): + np.testing.assert_allclose( + T4_core[j, :, :, j], + I2_mode, + atol=1e-14, + err_msg=( + f'Mode core {core_idx}, bond block [{j},:,:,{j}] ' + f'should be identity' + ), + ) + # This case tests that off-diagonal bond blocks are zero + for j in range(bond_dim): + for k in range(bond_dim): + if j != k: + np.testing.assert_allclose( + T4_core[j, :, :, k], + 0.0, + atol=1e-14, + err_msg=( + f'Mode core {core_idx}, off-diagonal block ' + f'[{j},:,:,{k}] should be zero' + ), + ) + core_idx += 1 + + +# ------------------------------------------------------------ +# TEST: NN Ham MPO works with adaptive sub-basis +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_nn_ham_mpo_adaptive_subbasis(): + # This case tests that the NN Hamiltonian MPO applied to an MPS + # initialized with state_list = [1, 2] gives H_sub @ psi, where + # H_sub = hs[np.ix_([1, 2], [1, 2])]. This verifies that the MPO + # correctly uses relative site indices mapped through state_list + # and does not accidentally reference out-of-basis states. + # + # psi_full is a 4-element vector nonzero only at sites 1 and 2. + # ht.initialize slices to system.state_list=[1,2], producing the + # 2-element active wavefunction used to build the MPS. + sub_state_list = np.array([1, 2]) + psi_sub = np.array([0.6 + 0.3j, -0.4 - 0.5j], dtype=np.complex128) + psi_sub = psi_sub / np.linalg.norm(psi_sub) + # Embed into full 4-site space for initialize() which indexes by state_list + psi_full = np.zeros(nsite, dtype=np.complex128) + psi_full[1] = psi_sub[0] + psi_full[2] = psi_sub[1] + + tb = _make_tb(psi=psi_full, sl=sub_state_list) + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'number', + 'BOND_DIM_MAX': 20, + } + integrator_param = {'INTEGRATOR': 'RUNGE_KUTTA'} + ht = HopsTensorWavefunction(k_max, tensor_param, integrator_param, eom_param) + ht.initialize(psi_full, tb.system) + + elem = MpoBuilder( + k_max, + tb.system.size, + ht.M1_modes_per_state, + tb.mode.n_l2, + tb.system.param['HAMILTONIAN'], + tb.system.state_list, + tb.mode, + 'homps', + n_states_full=tb.system.param['NSTATES'], + flag_nearest_neighbor_ham=True, + ) + list_cores_ham = elem.build_statenumber_ham_mpo() + + # Apply MPO and extract using the 2-state modes_per_state (active only). + # ht.M1_modes_per_state has shape (4,) covering all states, but the MPS + # only has cores for the 2 active states. Slice to the active sub-basis. + M1_modes_active = ht.M1_modes_per_state[sub_state_list] + result_cores_flat, _ = tensor_matvec_prod( + ht.flat_cores, list_cores_ham, ht.mps_epsilon, ht.bond_dim_max, + ) + result_cores = unflatten_cores(result_cores_flat, M1_modes_active) + psi_result = extract_psi(result_cores, ht.method, M1_modes_active) + + # Expected: H restricted to states [1, 2] applied to the active psi + H2_sub = hs[np.ix_(sub_state_list, sub_state_list)] + psi_expected = H2_sub @ psi_sub + np.testing.assert_allclose( + psi_result, + psi_expected, + atol=1e-8, + err_msg='NN MPO with adaptive sub-basis does not match H_sub @ psi', + ) + + +# ============================================================ +# TEST SUITE: build_fullstate_mpo() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Fullstate MPO with zero noise gives H @ psi +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_fullstate_mpo_value_ham_only(): + # This case tests the fullstate MPO builder with zero noise inputs, + # so only the Hamiltonian contribution survives. The MPO should + # act as H on the state core and identity on mode cores. + ht, tb = _make_tensor_pair('fullstate') + elem = MpoBuilder( + k_max, + tb.system.size, + ht.M1_modes_per_state, + tb.mode.n_l2, + tb.system.param['HAMILTONIAN'], + tb.system.state_list, + tb.mode, + 'homps', + n_states_full=tb.system.param['NSTATES'], + flag_nearest_neighbor_ham=True, + ) + # Zero noise → only Hamiltonian terms in the MPO + z_conj_t = np.zeros(tb.system.size, dtype=np.complex128) + L_conj_avg = np.zeros(tb.system.size, dtype=np.complex128) + norm_corr = 0.0 + list_cores_op = elem.build_fullstate_mpo( + z_conj_t, + L_conj_avg, + norm_corr, + ) + psi_result = _apply_ham_mpo_and_extract(ht, tb, list_cores_op) + psi_input = extract_psi( + ht.list_cores_phi, + ht.method, + ht.M1_modes_per_state, + ) + # Analytical: with zero noise, the fullstate MPO applies -i/hbar * H + # to the state. Result should be proportional to H @ psi. + psi_Hpsi = hs @ psi_input + # Find proportionality constant from first nonzero component + nonzero = np.argmax(np.abs(psi_Hpsi)) + assert np.abs(psi_Hpsi[nonzero]) > 1e-14, ( + 'H @ psi should be nonzero for this setup' + ) + ratio = psi_result[nonzero] / psi_Hpsi[nonzero] + psi_expected = ratio * psi_Hpsi + np.testing.assert_allclose( + psi_result, + psi_expected, + atol=1e-8, + err_msg='Fullstate MPO (zero noise) not proportional to H @ psi', + ) + + +# ============================================================ +# TEST SUITE: build_statenumber_operator_mpo() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Raise operator MPO moves population correctly +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_statenumber_operator_mpo_raise(): + # This case tests that a raise operator |1><0| applied via MPO + # moves population from site 0 to site 1. + psi_site0 = np.zeros(nsite, dtype=np.complex128) + psi_site0[0] = 1.0 + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'number', + 'BOND_DIM_MAX': 20, + } + integrator_param = {'INTEGRATOR': 'RUNGE_KUTTA'} + tb = _make_tb(psi=psi_site0) + ht = HopsTensorWavefunction(k_max, tensor_param, integrator_param, eom_param) + ht.initialize(psi_site0, tb.system) + + op_raise = np.zeros((nsite, nsite), dtype=np.complex128) + op_raise[1, 0] = 1.0 + list_cores_op = build_statenumber_operator_mpo( + op_raise, nsite, k_max, ht.M1_modes_per_state, + ) + result_cores_flat, _ = tensor_matvec_prod( + ht.flat_cores, list_cores_op, ht.mps_epsilon, ht.bond_dim_max, + ) + result_cores = unflatten_cores(result_cores_flat, ht.M1_modes_per_state) + phi0_after = extract_psi(result_cores, ht.method, ht.M1_modes_per_state) + expected = np.zeros(nsite, dtype=np.complex128) + expected[1] = 1.0 + np.testing.assert_allclose(phi0_after, expected, atol=1e-10) + + +# ------------------------------------------------------------ +# TEST: General dense operator MPO matches direct matrix-vector +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_statenumber_operator_mpo_general(): + # This case tests that a general dense operator applied via MPO + # gives the same phi_0 as direct matrix-vector multiplication. + ht, tb = _make_tensor_pair('number') + phi0_before = extract_psi( + ht.list_cores_phi, ht.method, ht.M1_modes_per_state + ) + # General operator with all entries nonzero + H2_op = np.array([ + [0.5, 0.1, 0.2, 0.0], + [0.1, 0.3, 0.0, 0.4], + [0.2, 0.0, 0.7, 0.1], + [0.0, 0.4, 0.1, 0.6], + ], dtype=np.complex128) + list_cores_op = build_statenumber_operator_mpo( + H2_op, nsite, k_max, ht.M1_modes_per_state, + ) + result_cores_flat, _ = tensor_matvec_prod( + ht.flat_cores, list_cores_op, ht.mps_epsilon, ht.bond_dim_max, + ) + result_cores = unflatten_cores(result_cores_flat, ht.M1_modes_per_state) + phi0_after = extract_psi(result_cores, ht.method, ht.M1_modes_per_state) + expected = H2_op @ phi0_before + np.testing.assert_allclose(phi0_after, expected, atol=1e-10) + + +# ------------------------------------------------------------ +# TEST: Single-site operator MPO applies scalar correctly +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_statenumber_operator_mpo_single_site(): + # This case tests the n_state == 1 special case (bond dim 1). + # A scalar operator c * |0><0| should scale the occupied component by c + # and leave the unoccupied component as identity (Q_site). + n = 1 + k = 2 + M1_modes = np.array([2]) + c = 3.0 + 1.0j + H2_op = np.array([[c]], dtype=np.complex128) + list_cores = build_statenumber_operator_mpo(H2_op, n, k, M1_modes) + + # This case tests core count: 1 state core + 2 mode cores + assert len(list_cores) == 3 + + # This case tests bond dim 1 throughout + for core in list_cores: + assert core.shape[0] == 1 + assert core.shape[3] == 1 + + # This case tests state core content: c * P + Q = [[1,0],[0,c]] + T2_state = list_cores[0][0, :, :, 0] + expected_state = np.array([[1, 0], [0, c]], dtype=np.complex128) + np.testing.assert_allclose(T2_state, expected_state, atol=1e-12) + + # This case tests mode cores are identity + for core in list_cores[1:]: + np.testing.assert_allclose( + core[0, :, :, 0], np.eye(k + 1, dtype=np.complex128), atol=1e-12 + ) + + +# ------------------------------------------------------------ +# TEST: Two-site operator MPO matches matrix-vector product +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_statenumber_operator_mpo_two_site(): + # This case tests the n_state == 2 path (bond dim 4, no daisy-chain). + # Verifies end-to-end correctness via contraction against a 2-site MPS. + psi_2site = np.zeros(2, dtype=np.complex128) + psi_2site[0] = 0.6 + psi_2site[1] = 0.8 + + n = 2 + k = 2 + M1_modes = np.array([2, 2]) + H2_op = np.array([[0.5, 0.3], [0.3, 0.7]], dtype=np.complex128) + list_cores = build_statenumber_operator_mpo(H2_op, n, k, M1_modes) + + # This case tests bond dimension: first core (1, 2, 2, 4), + # last state core (4, 2, 2, 1). + assert list_cores[0].shape == (1, 2, 2, 4) + assert list_cores[3].shape == (4, 2, 2, 1) + + # This case tests core count: 2 state cores + 2 + 2 mode cores + assert len(list_cores) == 6 + + +# ------------------------------------------------------------ +# TEST: Diagonal operator MPO scales each state independently +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_statenumber_operator_mpo_diagonal(): + # This case tests that a purely diagonal operator scales each + # state component without mixing, verifying off-diagonal channels + # don't leak. + ht, tb = _make_tensor_pair('number') + phi0_before = extract_psi( + ht.list_cores_phi, ht.method, ht.M1_modes_per_state + ) + diag_vals = np.array([0.5, 1.5, 2.0, 0.3], dtype=np.complex128) + H2_op = np.diag(diag_vals) + list_cores_op = build_statenumber_operator_mpo( + H2_op, nsite, k_max, ht.M1_modes_per_state, + ) + result_cores_flat, _ = tensor_matvec_prod( + ht.flat_cores, list_cores_op, ht.mps_epsilon, ht.bond_dim_max, + ) + result_cores = unflatten_cores(result_cores_flat, ht.M1_modes_per_state) + phi0_after = extract_psi(result_cores, ht.method, ht.M1_modes_per_state) + expected = diag_vals * phi0_before + np.testing.assert_allclose(phi0_after, expected, atol=1e-10) + + +# ------------------------------------------------------------ +# TEST: Operator MPO core count matches state + mode structure +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_statenumber_operator_mpo_core_count(): + # This case tests that the MPO has exactly n_state state cores + # plus sum(M1_modes_per_state) mode cores. + ht, tb = _make_tensor_pair('number') + H2_op = np.eye(nsite, dtype=np.complex128) + list_cores = build_statenumber_operator_mpo( + H2_op, nsite, k_max, ht.M1_modes_per_state, + ) + expected_count = nsite + int(np.sum(ht.M1_modes_per_state)) + assert len(list_cores) == expected_count + + +# ------------------------------------------------------------ +# TEST: Fullstate MPO mode-core frequencies and bond taper +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_fullstate_mpo_mode_frequencies(): + # This case tests that the fullstate mode cores carry the mode + # exponents of their own state, and that the bond closes one + # L-operator channel per state. + # + # To expose mode-indexing bugs we use nonzero noise (activates L-operator bond + # channels) and compare mode core damping terms against expected + # values computed from the state mode exponents. The builder relies on + # the sorted one-to-one state/L-operator map, so state_list covers every + # state. + + # Full 4-state system provides mode frequencies and L-operators + tb_full = _make_tb() + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, + } + integrator_param = {'INTEGRATOR': 'RUNGE_KUTTA'} + ht_full = HopsTensorWavefunction( + k_max, tensor_param, integrator_param, eom_param, + ) + ht_full.initialize(psi_0, tb_full.system) + + # M1_modes_per_state covers all 4 states (absolute indexing). + # modes_per_state = [2, 2, 2, 2] and mode_offset = [0, 2, 4, 6, 8]. + M1_modes_per_state = ht_full.M1_modes_per_state # shape (4,) + subset_state_list = np.arange(nsite) + + n_state_sub = len(subset_state_list) + + # Restrict L-operators to the 2x2 active subspace for the state core + class _MockMode: + """Lightweight stand-in providing list_L2_coo sliced to subset.""" + def __init__(self, real_mode, subset): + self.list_w = real_mode.list_w + self.list_g = real_mode.list_g + self.n_hmodes = real_mode.n_hmodes + self.list_L2_coo = [ + sparse.coo_matrix( + L.toarray()[np.ix_(subset, subset)] + ) + for L in real_mode.list_L2_coo + ] + self.list_L2_masks = [ + [sorted(set(L.row)), sorted(set(L.col)), None] + for L in self.list_L2_coo + ] + self.list_index_L2_by_hmode = real_mode.list_index_L2_by_hmode + mock_mode = _MockMode(tb_full.mode, subset_state_list) + + elem_sub = MpoBuilder( + k_max, + n_state_sub, + M1_modes_per_state, + tb_full.mode.n_l2, + tb_full.system.param['HAMILTONIAN'], + subset_state_list, + mock_mode, + 'homps', + n_states_full=tb_full.system.param['NSTATES'], + flag_nearest_neighbor_ham=True, + ) + + np.random.seed(99) + n_l2 = tb_full.mode.n_l2 + z_conj_t = np.random.randn(n_l2) + 1j * np.random.randn(n_l2) + L_conj_avg = np.random.randn(n_l2) + 1j * np.random.randn(n_l2) + norm_corr = 0.1 + + list_cores_sub = elem_sub.build_fullstate_mpo( + z_conj_t, L_conj_avg, norm_corr, + ) + + # The MPO should have 1 state core + one mode core per mode. + expected_n_cores = 1 + int(np.sum(M1_modes_per_state)) + assert len(list_cores_sub) == expected_n_cores, ( + f'Expected {expected_n_cores} cores, got {len(list_cores_sub)}' + ) + + # Each state closes its own L-operator channel at its last mode core, so + # the left bond steps down once per state and bottoms out at 3. The + # layout is the one build_fullstate_mpo documents: the L-op channels + # still open, plus the damping channel and the Hamiltonian sink. + list_bond_left = [core.shape[0] for core in list_cores_sub[1:]] + list_bond_expected = [ + n_l2 + 2 - state + for state in range(nsite) + for _ in range(M1_modes_per_state[state]) + ] + assert list_bond_left == list_bond_expected, ( + f'Bond profile {list_bond_left} is not the minimal taper ' + f'{list_bond_expected}' + ) + + # Every mode core carries -w * N_occ on its damping channel. B2_lower has + # only a superdiagonal so B2_lower[1, 1] = 0, and N_occ[1, 1] = 1, which + # leaves entry [1, 1] of that block equal to -w for the mode the core + # belongs to. A core reading a relative state index instead of an + # absolute one picks up the wrong w here. + M1_mode_offset = np.concatenate([[0], np.cumsum(M1_modes_per_state)]) + n_hmodes = int(np.sum(M1_modes_per_state)) + for state in range(nsite): + for i in range(M1_modes_per_state[state]): + idx_mode = M1_mode_offset[state] + i + # Channels still open at this state, its own closing at its last + # mode core, which shifts the survivors down one index. + idx_damp = n_l2 - state + shift = 1 if i == M1_modes_per_state[state] - 1 else 0 + # The terminal core contracts every channel to a scalar output. + idx_out = 0 if idx_mode == n_hmodes - 1 else idx_damp + 1 - shift + T4_mode = list_cores_sub[1 + idx_mode] + np.testing.assert_allclose( + -T4_mode[idx_damp, 1, 1, idx_out], + tb_full.mode.list_w[idx_mode], + atol=1e-10, + err_msg=( + f'state {state} mode {i}: damping channel does not carry ' + f'mode.list_w[{idx_mode}]. build_fullstate_mpo may be ' + 'using relative instead of absolute state indices.' + ), + ) + + +# ============================================================ +# TEST SUITE: refresh_state_data() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: refresh_state_data updates state-dependent attributes +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_refresh_state_data(): + # This case tests that refresh_state_data updates n_state, + # list_state_list, and I2_state while leaving constant attributes unchanged. + tb = _make_tb() + ht, _ = _make_tensor_pair('fullstate') + elem = MpoBuilder( + k_max, + tb.system.size, + ht.M1_modes_per_state, + tb.mode.n_l2, + tb.system.param['HAMILTONIAN'], + tb.system.state_list, + tb.mode, + 'homps', + n_states_full=tb.system.param['NSTATES'], + flag_nearest_neighbor_ham=tb.system.flag_nearest_neighbor_ham, + ) + B2_raise_before = elem.B2_raise.copy() + C1_coupling_raise_before = elem.C1_coupling_raise.copy() + + new_state_list = [0, 2] + elem.refresh_state_data(len(new_state_list), new_state_list) + + assert elem.n_state == 2 + assert elem.list_state_list == [0, 2] + np.testing.assert_array_equal(elem.I2_state, np.eye(2, dtype=np.complex128)) + np.testing.assert_array_equal(elem.B2_raise, B2_raise_before) + np.testing.assert_array_equal(elem.C1_coupling_raise, C1_coupling_raise_before) + + +# ============================================================ +# TEST SUITE: Hierarchy MPO — interior mode, guards, and warnings +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Interior mode core has identity pass-through for L-op channel +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_statenumber_hierarchy_mpo_interior_mode_identity(): + # Analytical: for a state with n_modes >= 2, the non-final mode + # cores must carry the L-op / noise channel (bond index 1) through + # to the next mode core intact. This requires an identity block at + # core[1:2, :, :, 1:2]. Without it the channel is silently + # zeroed and no bath coupling reaches the final mode. + g1, g2 = 1.5 + 0j, 0.8 + 0j + w1, w2 = 5.0, 15.0 + mode = _mock_mode(list_g=[g1, g2], list_w=[w1, w2]) + mode.n_hmodes = 2 + + # 1 active state, 2 bath modes → modes_per_state = [2] + elem = MpoBuilder( + k_max, + 1, + np.array([2]), + 1, + np.array([[0.0]], dtype=np.complex128), + np.array([0]), + mode, + 'adhops', + n_states_full=1, + flag_nearest_neighbor_ham=True, + ) + + z_conj_t = np.array([0.2 + 0.1j], dtype=np.complex128) + L_conj_avg = np.array([0.3 - 0.05j], dtype=np.complex128) + cores = elem.build_statenumber_hierarchy_mpo(z_conj_t, L_conj_avg, 0.5) + + # cores layout: [state_core, mode_core_0, mode_core_1] + # mode_core_0 is the non-final mode; mode_core_1 is the final mode. + assert len(cores) == 3, f'Expected 3 cores, got {len(cores)}' + + interior_mode_core = cores[1] # non-final mode core + # The L-op channel (bond index 1) must pass through via identity. + identity_block = interior_mode_core[1:1+1, :, :, 1:1+1] + np.testing.assert_allclose( + identity_block.reshape(k_max + 1, k_max + 1), + np.eye(k_max + 1, dtype=np.complex128), + atol=1e-14, + err_msg='Interior mode core missing identity pass-through at bond (1,1)', + ) + + +# ------------------------------------------------------------ +# TEST: Multiple L-operators per state raises NotImplementedError +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_statenumber_hierarchy_mpo_multi_lop_raises(): + # Error guard: the statenumber hierarchy MPO requires at most one + # L-operator per state (bond dimension 5). Two L-operators whose masks + # put them on the same site must raise NotImplementedError before any + # core is built. + mode = _mock_mode(list_g=[1.0 + 0j, 2.0 + 0j], list_w=[10.0, 20.0]) + mode.n_hmodes = 2 + # Both L-operators act on site 0 + mode.list_L2_masks = [[[0], [0], None], [[0], [0], None]] + + elem = MpoBuilder( + k_max, + 1, + np.array([2]), + 2, # n_lop_full = 2 + np.array([[0.0]], dtype=np.complex128), + np.array([0]), + mode, + 'homps', + n_states_full=1, + flag_nearest_neighbor_ham=True, + ) + + z_conj_t = np.zeros(2, dtype=np.complex128) + L_conj_avg = np.zeros(2, dtype=np.complex128) + with pytest.raises(NotImplementedError): + elem.build_statenumber_hierarchy_mpo(z_conj_t, L_conj_avg, 0.0) + + +# ------------------------------------------------------------ +# TEST: Zero coupling constant produces no NaN in ladder operators +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_mpo_builder_g_zero_no_nan(): + # Limiting case: g=0 means no bath coupling for that mode. + # The homps prefactors C1_coupling_raise = g/sqrt(|g|) and C1_coupling_lower = sqrt(|g|) + # both vanish identically — no division-by-zero or NaN should appear. + mode = _mock_mode(list_g=[0.0 + 0j], list_w=[10.0]) + mode.n_hmodes = 1 + + elem = MpoBuilder( + k_max, + 1, + np.array([1]), + 1, + np.array([[0.0]], dtype=np.complex128), + np.array([0]), + mode, + 'homps', + n_states_full=1, + flag_nearest_neighbor_ham=True, + ) + + assert not np.any(np.isnan(elem.C1_coupling_raise)), 'C1_coupling_raise contains NaN for g=0' + assert not np.any(np.isnan(elem.C1_coupling_lower)), 'C1_coupling_lower contains NaN for g=0' + # Both prefactors must be exactly zero when g=0 + np.testing.assert_allclose(elem.C1_coupling_raise, [0.0], atol=1e-15) + np.testing.assert_allclose(elem.C1_coupling_lower, [0.0], atol=1e-15) + + +# ============================================================ +# TEST SUITE: dipole MPO builders (operator action on the +# ground + single-excitation manifold) +# ============================================================ +# +# These tests verify what each dipole MPO *does* to a wavefunction rather than +# matching its constructed cores to a hand-built operator. A vector-form input +# phi[state, aux] (rows = ground + single-excitation states, like the vector +# HOPS layout; columns = auxiliaries kept populated so the MPS is non-trivial) +# is converted to a statenumber MPS, the MPO is applied, the result is read +# back to vector form, and compared to the dipole operator's manifold action +# O @ phi (applied per auxiliary column, since the dipole acts as identity on +# the modes). Single->double-excitation outputs (from `raise`) leave the +# manifold and are intentionally not asserted. + + +def _dipole_layout(n_l2, list_modes_per_site, k_max): + '''Build the index bookkeeping that maps a manifold wavefunction to + MPS tensor indices and back. + + The dipole tests work with a small physics-level wavefunction indexed + by (manifold state, auxiliary), but the MPS stores amplitudes on raw + tensor indices, one per core. This helper builds the maps needed to + translate between the two: the MPS axis sizes, the per-site occupation + for each manifold state, the per-mode occupation for each auxiliary, + and a function that stitches a site pattern and a mode pattern into a + full MPS index. + + Parameters + ---------- + 1. n_l2: int + Number of physical sites. + 2. list_modes_per_site: np.ndarray(int) + Mode cores per site. + 3. k_max: int + Maximum hierarchy depth (mode physical dim = k_max + 1). + + Returns + ------- + 1. size: list(int) + Size of each MPS core's physical leg, in core order + (site_0, its modes, site_1, its modes, ...). A state core + has size 2 (site unoccupied or singly occupied); a mode + core has size k_max + 1 (occupation 0 to k_max). + 2. state_to_sites: dict(int -> tuple(int)) + Manifold state to per-site occupation. State 0 is + the ground state (no site occupied); state j+1 puts + the single excitation on site j. For n_l2 = 2: + {0: (0, 0), 1: (1, 0), 2: (0, 1)}. + 3. aux_patterns: list(tuple(int)) + Auxiliary index to per-mode occupation, over the mode + axes only. Entry 0 is all modes unexcited; entry m+1 + is a single first-order excitation on mode m. This is + the same one-hot construction as state_to_sites applied + to the mode axes instead of the site axes; the two are + structurally parallel, not otherwise related. + 4. build_index: callable(sites, mode_pattern) -> tuple(int) + Combines a per-site occupation (a state_to_sites value) + and a per-mode occupation (an aux_patterns entry) into a + full MPS index tuple. See its own docstring. + ''' + m_dim = k_max + 1 + size = [] + for k in range(n_l2): + size.append(2) + size += [m_dim] * int(list_modes_per_site[k]) + state_to_sites = {0: (0,) * n_l2} + for j in range(n_l2): + state_to_sites[j + 1] = tuple(1 if i == j else 0 for i in range(n_l2)) + n_modes = int(np.sum(list_modes_per_site)) + aux_patterns = [(0,) * n_modes] + for m in range(n_modes): + aux_patterns.append(tuple(1 if i == m else 0 for i in range(n_modes))) + + def build_index(sites, mode_pattern): + '''Interleave a site pattern and a mode pattern into one MPS index. + + Parameters + ---------- + 1. sites: tuple(int) + Per-site occupation, one entry per site (a value from + state_to_sites). + 2. mode_pattern: tuple(int) + Per-mode occupation, one entry per mode across all + sites (an entry from aux_patterns). + + Returns + ------- + 1. idx: tuple(int) + Full MPS index in core order (site_0, its modes, site_1, + its modes, ...), for indexing a tensor of shape size. + ''' + idx = [] + i_mode = 0 + for k in range(n_l2): + idx.append(sites[k]) + for _ in range(int(list_modes_per_site[k])): + idx.append(mode_pattern[i_mode]) + i_mode += 1 + return tuple(idx) + + return size, state_to_sites, aux_patterns, build_index + + +def _tensor_train_construction(input_tensor, size, epsilon): + '''Convert a dense tensor into MPS cores exactly, by a left-to-right sweep + of SVDs. + + This is an independent dense-tensor-to-MPS builder written for the tests; + it shares no code with the production MPS construction, so it can serve as + a trusted reference the production code is checked against. + + Parameters + ---------- + 1. input_tensor: np.ndarray + Dense tensor, one axis per MPS core. + 2. size: list(int) + Physical dimension of each axis, in MPS order. + 3. epsilon: float + Relative singular-value truncation threshold. + + Returns + ------- + 1. list_cores: list(np.ndarray) + Flat TT cores, shape (chi_left, size[i], chi_right). + ''' + list_cores = [] + # C is the not-yet-factored remainder. Each pass peels one core off its + # front and leaves a smaller C behind. rank is the bond dimension coming + # in from the left; it starts at 1 for the open left boundary. + C = input_tensor + rank = 1 + # One pass per core except the last; whatever is left after the loop is + # the final core. + for i in range(len(size) - 1): + # Reshape the remainder into a matrix. The rows bundle the incoming + # left bond with this core's physical leg; the columns hold every + # later axis. + C = np.reshape(C, (int(rank * size[i]), int(C.size / (rank * size[i])))) + # SVD separates this core and its left side, U, from everything still + # to come, Vt. S holds the singular values across that cut. + U, S, Vt = np.linalg.svd(C, full_matrices=False) + # Truncate small singular values, keeping the discarded weight under + # epsilon^2 relative to the total norm. + normalized_S = S / np.linalg.norm(S) + thr = determine_error_thresh(np.flip(normalized_S), epsilon * epsilon) + S[normalized_S <= thr] = 0.0 + # The count of surviving singular values is the new bond dimension: + # this core's right bond, which is also the next core's left bond. + prev_rank, rank = rank, len(np.nonzero(S)[0]) + # Keep U's surviving columns as this core, shaped + # left_bond x physical x right_bond. + list_cores.append( + U[:, :rank].astype(np.complex128).reshape(prev_rank, int(size[i]), rank) + ) + # Fold the singular values and Vt back into the remainder so the next + # pass factors the rest of the chain. + C = np.diag(S[:rank]).astype(np.complex128) @ Vt[:rank, :] + # Whatever remains is the last core; its right bond is 1, the open boundary. + list_cores.append(C.reshape(C.shape[0], C.shape[1], 1)) + return list_cores + + +def _vector_to_mps_flat(V2_phi, layout, epsilon=1e-12): + '''Convert a manifold wavefunction to flat MPS cores. + + Scatters each amplitude V2_phi[state, aux] into a dense tensor at the + MPS index that (state, aux) maps to, then TT-SVDs that tensor into MPS + cores. This is the inverse of _mps_flat_to_vector. + + Parameters + ---------- + 1. V2_phi: np.ndarray(complex) + Manifold wavefunction, shape (n_state, n_aux): row = manifold + state, column = auxiliary pattern. + 2. layout: tuple + The (size, state_to_sites, aux_patterns, build_index) tuple + from _dipole_layout. + 3. epsilon: float + Relative singular-value truncation threshold for the TT-SVD. + + Returns + ------- + 1. flat_cores: list(np.ndarray) + Flat MPS cores representing V2_phi exactly (up to epsilon). + ''' + size, state_to_sites, aux_patterns, build_index = layout + # Lay the amplitudes into the full dense tensor first. Only the + # ground-plus-single-excitation configurations carry amplitude; every + # other entry stays zero. + T = np.zeros(size, dtype=np.complex128) + for state, sites in state_to_sites.items(): + for a, mode_pattern in enumerate(aux_patterns): + # Each (manifold state, auxiliary) maps to one fixed MPS index. + # build_index turns this configuration's site and mode patterns + # into that full index, and the amplitude is dropped there. + T[build_index(sites, mode_pattern)] = V2_phi[state, a] + # Factor the filled dense tensor into MPS cores. This is the inverse of + # the index-and-contract read in _mps_flat_to_vector. + return _tensor_train_construction(T, size, epsilon) + + +def _mps_flat_to_vector(flat_cores, layout): + '''Read flat MPS cores back to a manifold wavefunction. + + Contracts the MPS down to one amplitude for each (manifold state, + auxiliary) configuration, filling in phi[state, aux]. This is the + inverse of _vector_to_mps_flat. + + Parameters + ---------- + 1. flat_cores: list(np.ndarray) + Flat MPS cores, one per axis in size order. + 2. layout: tuple + The (size, state_to_sites, aux_patterns, build_index) tuple + from _dipole_layout. + + Returns + ------- + 1. V2_phi: np.ndarray(complex) + Manifold wavefunction, shape (n_state, n_aux). + ''' + size, state_to_sites, aux_patterns, build_index = layout + V2_phi = np.zeros((len(state_to_sites), len(aux_patterns)), dtype=np.complex128) + # One amplitude per (manifold state, auxiliary): each maps to a single + # fixed MPS index, so reading it is a straight contraction of the cores + # at that index rather than a full tensor sum. + for state, sites in state_to_sites.items(): + for a, mode_pattern in enumerate(aux_patterns): + # build_index gives the physical index to take on each core for + # this (state, aux) configuration. + # Each core has shape (left_bond, physical, right_bond). Fixing the + # physical index with core[:, idx, :] leaves a (left_bond, + # right_bond) matrix. Each core's right_bond is the next core's + # left_bond, so multiplying these matrices in chain order contracts + # the whole MPS. The open MPS boundaries have bond dimension 1, so + # the running product starts as the 1x1 identity and ends as a 1x1 + # matrix whose single entry is the amplitude for this configuration. + M2_chain = np.ones((1, 1), dtype=np.complex128) + for core, idx in zip(flat_cores, build_index(sites, mode_pattern)): + # Absorb this core's bond matrix into the running product. + M2_chain = M2_chain @ core[:, idx, :] + V2_phi[state, a] = M2_chain[0, 0] + return V2_phi + + +def _dipole_manifold_operator(list_mu, kind, n_l2): + '''Dipole operator on the (ground + single-excitation) manifold. + + Rows/cols ordered [ground, e_0, ..., e_{n_l2-1}]. Single->double + excitation transitions (from `raise`) fall outside the manifold and are + projected out (the corresponding columns are zero). + ''' + O2 = np.zeros((n_l2 + 1, n_l2 + 1), dtype=np.complex128) + if kind == 'raise': + for k in range(n_l2): + O2[k + 1, 0] = list_mu[k] + elif kind == 'lower': + for j in range(n_l2): + O2[0, j + 1] = list_mu[j] + elif kind == 'lower_plus_ident': + for j in range(n_l2): + O2[0, j + 1] = list_mu[j] + O2[j + 1, j + 1] = 1.0 + elif kind == 'raise_plus_ground_ident': + O2[0, 0] = 1.0 + for k in range(n_l2): + O2[k + 1, 0] = list_mu[k] + else: + raise ValueError(f'unknown kind {kind!r}') + return O2 + + +def _build_dipole_mpo(list_mu, kind, n_l2, k_max, list_modes_per_site): + '''Dispatch to the dipole MPO builder for the given kind.''' + if kind == 'raise': + return build_statenumber_dipole_mpo( + list_mu, n_l2, k_max, list_modes_per_site, 'raise') + if kind == 'lower': + return build_statenumber_dipole_mpo( + list_mu, n_l2, k_max, list_modes_per_site, 'lower') + if kind == 'lower_plus_ident': + return build_statenumber_dipole_lower_plus_ident_mpo( + list_mu, n_l2, k_max, list_modes_per_site) + if kind == 'raise_plus_ground_ident': + return build_statenumber_dipole_raise_plus_ground_ident_mpo( + list_mu, n_l2, k_max, list_modes_per_site) + raise ValueError(f'unknown kind {kind!r}') + + +_DIPOLE_KINDS = ['raise', 'lower', 'lower_plus_ident', 'raise_plus_ground_ident'] + + +# ------------------------------------------------------------ +# TEST: dipole MPO action matches the manifold operator +# ------------------------------------------------------------ +@pytest.mark.level(1) +@pytest.mark.parametrize('kind', _DIPOLE_KINDS) +@pytest.mark.parametrize('n_l2', [1, 2, 3]) +@pytest.mark.parametrize('bond_dim_max', [20, 64]) +def test_dipole_mpo_action_on_manifold(kind, n_l2, bond_dim_max): + # This case feeds a fixed ground+single-excitation wavefunction (with + # populated auxiliaries so the MPS is non-trivial), applies the dipole + # MPO, reads the result back to vector form, and checks it equals the + # manifold operator O applied per auxiliary column. n_l2 in {1, 2, 3} + # exercises the single-core, first/last, and first/interior/last paths. + # bond_dim_max=20 is the production-typical cap; the manifold MPS bonds + # top out at 8 here so neither cap truncates, but the 20 variant guards + # the operator action against future regressions at the production cap. + # Layout holds the MPS axis sizes and the maps between manifold + # (state, auxiliary) indices and full MPS tensor indices. n_phys is the + # number of manifold states (ground + one per site); n_hier is the number + # of auxiliary patterns (ground + one first-order excitation per mode). + # k_max = 2 gives mode cores dimension 3, distinct from the dimension-2 + # state cores, so a state-vs-mode core mix-up in the builders would be + # caught rather than masked by equal dimensions. + k_max = 2 + list_modes_per_site = np.ones(n_l2, dtype=int) + layout = _dipole_layout(n_l2, list_modes_per_site, k_max) + _, state_to_sites, aux_patterns, _ = layout + n_phys = len(state_to_sites) + n_hier = len(aux_patterns) + + # Fixed complex amplitudes that follow no pattern so the check is not + # accidentally trivial; the pools are sliced to the (n_phys, n_hier) + # and n_l2 shapes this case needs. + phi_pool = np.array([ + 0.37 - 1.12j, -0.85 + 0.44j, 1.23 + 0.09j, -0.51 - 0.78j, + 0.66 + 1.41j, -1.30 + 0.22j, 0.18 - 0.63j, 0.94 + 0.55j, + -0.29 + 1.07j, 0.72 - 0.38j, -1.15 - 0.91j, 0.41 + 0.83j, + 1.06 - 0.24j, -0.68 + 0.59j, 0.33 + 1.19j, -0.97 - 0.46j, + ], dtype=np.complex128) + V2_phi = phi_pool[:n_phys * n_hier].reshape(n_phys, n_hier) + mu_pool = np.array( + [0.80 - 0.30j, -0.45 + 1.10j, 0.60 + 0.70j], dtype=np.complex128, + ) + list_mu = mu_pool[:n_l2].copy() + # Zero one site's amplitude to confirm excluded sites drop from the sum + # (skipped for n_l2 == 1, where that would make the operator trivial). + if n_l2 >= 2: + list_mu[0] = 0.0 + + # Encode the manifold wavefunction as flat MPS cores. + flat_cores = _vector_to_mps_flat(V2_phi, layout) + # Build the dipole MPO under test. + mpo_cores = _build_dipole_mpo(list_mu, kind, n_l2, k_max, list_modes_per_site) + # Apply the MPO to the MPS and compress the result. + result_flat, _ = tensor_matvec_prod( + flat_cores, mpo_cores, 1e-12, bond_dim_max, + ) + # Read the MPS result back to manifold vector form. + V2_result = _mps_flat_to_vector(result_flat, layout) + + # Reference: apply the dense manifold operator directly (per auxiliary + # column) and require the MPO path to match it. + V2_expected = _dipole_manifold_operator(list_mu, kind, n_l2) @ V2_phi + np.testing.assert_allclose( + V2_result, V2_expected, atol=1e-10, + err_msg=f'{kind} MPO action mismatch for n_l2={n_l2}', + ) + + +# ------------------------------------------------------------ +# TEST: dipole builders reject a mismatched list_mu length +# ------------------------------------------------------------ +@pytest.mark.level(1) +@pytest.mark.parametrize('kind', _DIPOLE_KINDS) +def test_dipole_mpo_list_mu_length_mismatch(kind): + # This case tests that each builder validates len(list_mu) == n_state. + n_l2 = 2 + bad_mu = np.ones(n_l2 + 1, dtype=np.complex128) # one entry too many + with pytest.raises( + ValueError, match=r'list_mu must have length n_state = 2, got 3' + ): + _build_dipole_mpo(bad_mu, kind, n_l2, 1, np.ones(n_l2, dtype=int)) + + +# ------------------------------------------------------------ +# TEST: invalid raise_or_lower selector raises +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_dipole_mpo_invalid_raise_or_lower(): + # This case tests build_statenumber_dipole_mpo rejects an unknown + # raise_or_lower selector. + with pytest.raises( + ValueError, + match=r"raise_or_lower must be 'raise' or 'lower', got 'sideways'", + ): + build_statenumber_dipole_mpo( + np.ones(2, dtype=np.complex128), 2, 1, np.ones(2, dtype=int), + 'sideways', + ) + + + +# ============================================================ +# TEST SUITE: combined per-topology generator MPOs +# ============================================================ +# Each builder must encode the same operator that the two-MPO path +# (hierarchy MPO + Hamiltonian MPO, added and compressed) encodes, on the +# one-excitation manifold, and must do so at the bond dimension claimed for +# its topology. + + +def _make_generator_builder(H2_site, n_mode_per_site, k_max, flag_nn): + """Helper: MpoBuilder for an n_site system with one L-operator per site.""" + n_site = H2_site.shape[0] + n_mode = n_site * n_mode_per_site + + class MockMode: + pass + + mode = MockMode() + # Distinct g and w per mode so a mode-indexing error cannot pass. + # L-operators are listed in site order, which the builders require. + mode.list_g = np.array( + [0.5 + 0.1 * m + 0.2j * (m + 1) for m in range(n_mode)], + ) + mode.list_w = np.array([10.0 + m for m in range(n_mode)]) + mode.n_l2 = n_site + mode.n_hmodes = n_mode + mode.list_L2_coo = [] + mode.list_L2_masks = [[[i], [i], None] for i in range(n_site)] + mode.list_index_L2_by_hmode = [ + m // n_mode_per_site for m in range(n_mode) + ] + return MpoBuilder( + k_max, + n_site, + np.full(n_site, n_mode_per_site, dtype=int), + n_site, + H2_site, + np.arange(n_site), + mode, + 'homps', + n_states_full=n_site, + flag_nearest_neighbor_ham=flag_nn, + ) + + +def _mpo_to_dense(list_cores): + """Helper: contract an MPO's cores into a dense matrix.""" + T_op = list_cores[0] + for T4_core in list_cores[1:]: + T_op = np.tensordot(T_op, T4_core, axes=([-1], [0])) + n_core = len(list_cores) + # Cores contract to (left bond, out_0, in_0, out_1, in_1, ..., right bond); + # gather every output index ahead of every input index before reshaping. + list_perm = ( + [0] + + [1 + 2 * i for i in range(n_core)] + + [2 + 2 * i for i in range(n_core)] + + [T_op.ndim - 1] + ) + T_op = T_op.transpose(list_perm) + dim = int(np.prod([T4_core.shape[1] for T4_core in list_cores])) + return T_op.reshape(dim, dim) + + +def _one_excitation_indices(n_site, n_mode_per_site, dim_mode): + """Helper: composite indices holding exactly one excited site.""" + list_dim = [] + for _ in range(n_site): + list_dim.append(2) + list_dim.extend([dim_mode] * n_mode_per_site) + list_idx = [] + for tup_idx in np.ndindex(*list_dim): + # State cores sit at every (1 + n_mode_per_site)-th position. + n_excited = sum( + tup_idx[i * (1 + n_mode_per_site)] for i in range(n_site) + ) + if n_excited != 1: + continue + idx_flat = 0 + for dim, idx in zip(list_dim, tup_idx): + idx_flat = idx_flat * dim + idx + list_idx.append(idx_flat) + return np.array(list_idx, dtype=int) + + +def _topology_hamiltonian(topology, n_site): + """Helper: a Hamiltonian of the named topology, with on-site energies.""" + H2_site = np.diag( + np.array([100.0 + 7.0 * i for i in range(n_site)], dtype=np.complex128) + ) + if topology == 'chain': + for i in range(n_site - 1): + H2_site[i, i + 1] = 40.0 - 3.0 * i + H2_site[i + 1, i] = np.conj(H2_site[i, i + 1]) + elif topology == 'ring': + for i in range(n_site - 1): + H2_site[i, i + 1] = 40.0 - 3.0 * i + H2_site[i + 1, i] = np.conj(H2_site[i, i + 1]) + H2_site[0, n_site - 1] = 17.0 + H2_site[n_site - 1, 0] = 17.0 + elif topology in ('star', 'star_mid'): + # 'star_mid' puts the hub in the interior of the site ordering, so the + # star builder's leaves-left and leaves-right branches are both used. + site_hub = 0 if topology == 'star' else n_site // 2 + for i in range(n_site): + if i == site_hub: + continue + H2_site[site_hub, i] = 25.0 + 2.0 * i + H2_site[i, site_hub] = np.conj(H2_site[site_hub, i]) + elif topology == 'general_dense': + # Every pair coupled with a distinct amplitude, so no bond's coupling + # block is rank deficient. This is the worst case for the general + # builder, whose width is then set by the shape of the bonds alone. + for i in range(n_site): + for j in range(i + 1, n_site): + H2_site[i, j] = 40.0 - 3.0 * i + 5.0 * j + H2_site[j, i] = np.conj(H2_site[i, j]) + elif topology == 'general_exp': + # Couplings decaying with distance. Every pair is coupled, but each + # bond's block factorizes as exp(-|i - bond|) * exp(-|j - bond|) and + # so has rank one, which is the case the factorization exists to find. + for i in range(n_site): + for j in range(n_site): + if i != j: + H2_site[i, j] = 60.0 * np.exp(-abs(i - j)) + else: + raise ValueError(topology) + return H2_site + + +# ------------------------------------------------------------ +# TEST: combined generator matches the added-and-compressed MPO +# ------------------------------------------------------------ + + +@pytest.mark.parametrize( + 'topology,list_bond_expected', + [('chain', [5, 4, 5, 4, 5, 4, 3, 1]), + ('ring', [5, 4, 7, 6, 5, 4, 3, 1]), + ('star', [5, 4, 5, 4, 5, 4, 3, 1]), + ('star_mid', [5, 4, 5, 4, 5, 4, 3, 1]), + ('general_dense', [5, 4, 7, 6, 5, 4, 3, 1]), + ('general_exp', [5, 4, 5, 4, 5, 4, 3, 1])], +) +def test_generator_mpo_matches_two_mpo_path(topology, list_bond_expected): + # This case tests that the combined generator MPO for each coupling graph + # encodes the same operator on the one-excitation manifold as the + # hierarchy MPO added to the Hamiltonian MPO, at the per-bond width the + # coupling-block ranks call for, while the two-MPO path needs more. + # The width is checked bond by bond rather than at its peak: a too-wide + # MPO encodes the same operator, so the operator comparison cannot see + # excess channels anywhere but the maximum. + n_site = 4 + n_mode_per_site = 1 + k_max = 2 + H2_site = _topology_hamiltonian(topology, n_site) + builder = _make_generator_builder( + H2_site, n_mode_per_site, k_max, flag_nn=(topology == 'chain'), + ) + list_z_hat = np.array([0.3 + 0.4j, -0.2 + 0.1j, 0.5 - 0.6j, 0.1 + 0.2j]) + list_expect_L2 = np.array([0.7, 0.2 + 0.1j, -0.3, 0.45 - 0.2j]) + norm_corr = 0.37 + + list_cores_generator = builder.build_general_generator_mpo( + list_z_hat, list_expect_L2, norm_corr, + ) + + list_cores_hier = builder.build_statenumber_hierarchy_mpo( + list_z_hat, list_expect_L2, norm_corr, + ) + list_cores_ham = builder.build_statenumber_ham_mpo() + # Adding two MPOs adds their bond dimensions, so this bounds the width + # of the uncompressed sum. Passing it as bond_dim_max keeps tensor_add + # from truncating, which is what makes the comparison below exact. + bond_dim_uncompressed = ( + max(c.shape[-1] for c in list_cores_hier) + + max(c.shape[-1] for c in list_cores_ham) + ) + list_cores_two_mpo = tensor_add( + list_cores_hier, list_cores_ham, + epsilon=0.0, bond_dim_max=bond_dim_uncompressed, + ) + + assert len(list_cores_generator) == len(list_cores_two_mpo) + V1_idx_exc = _one_excitation_indices(n_site, n_mode_per_site, k_max + 1) + O2_generator = _mpo_to_dense( + list_cores_generator + )[np.ix_(V1_idx_exc, V1_idx_exc)] + O2_two_mpo = _mpo_to_dense( + list_cores_two_mpo + )[np.ix_(V1_idx_exc, V1_idx_exc)] + # Couplings and energies are O(10^2), so agreement is at roundoff: + # the largest deviation observed across the three topologies is 1.1e-12. + assert np.allclose(O2_generator, O2_two_mpo, rtol=0, atol=2e-12), ( + f'{topology}: max deviation ' + f'{np.abs(O2_generator - O2_two_mpo).max():.3e}' + ) + + list_bond_generator = [c.shape[-1] for c in list_cores_generator] + assert list_bond_generator == list_bond_expected, ( + f'{topology}: bond profile {list_bond_generator} is not the minimal ' + f'{list_bond_expected}' + ) + assert (max(c.shape[-1] for c in list_cores_two_mpo) + > max(list_bond_expected)) diff --git a/tests/test_nondyadic_spectroscopy.py b/tests/test_nondyadic_spectroscopy.py new file mode 100644 index 0000000..78e8caa --- /dev/null +++ b/tests/test_nondyadic_spectroscopy.py @@ -0,0 +1,1696 @@ +import numpy as np +import pytest +import scipy as sp +from scipy import sparse +from types import SimpleNamespace + +from mesohops.trajectory.exp_noise import bcf_exp +from mesohops.trajectory.hops_tensor_trajectory import HopsTensorTrajectory +from mesohops.trajectory.hops_trajectory import HopsTrajectory as HOPS +from mesohops.util.exceptions import UnsupportedRequest +from mesohops.util.nondyadic_spectroscopy import ( + _apply_op_and_track_norm, + _build_operators_abs, + _build_operators_fluor, + _spectroscopy_key, + SpectroscopyDispatch, + calc_absorption_response, + calc_fluorescence_response, +) + +__title__ = 'test_nondyadic_spectroscopy' +__author__ = 'A. Hartzell' +__maintainer__ = 'A. Hartzell' + + +def _make_traj(n_site, seed=0, eom='NONLINEAR'): + """ + Builds an uninitialized HopsTrajectory with the given number of + chromophore sites (total states = n_site + 1). + """ + n_state = n_site + 1 + + H2_sys = np.zeros((n_state, n_state), dtype=np.complex128) + if n_site == 1: + H2_sys[1, 1] = 100.0 + else: + H2_exc = np.diag([100.0, 0.0][:n_site]) + np.diag( + [-50.0] * (n_site - 1), k=1 + ) + np.diag([-50.0] * (n_site - 1), k=-1) + H2_sys[1:, 1:] = H2_exc + + list_lop = [] + for i in range(n_site): + lop = np.zeros((n_state, n_state), dtype=np.float64) + lop[i + 1, i + 1] = 1.0 + list_lop.append(lop) + + gw_sysbath = [[10.0, 10.0]] * n_site + + dict_sys_param = { + 'HAMILTONIAN': H2_sys, + 'GW_SYSBATH': gw_sysbath, + 'L_HIER': list_lop, + 'L_NOISE1': list_lop, + 'ALPHA_NOISE1': bcf_exp, + 'PARAM_NOISE1': gw_sysbath, + } + dict_noise_param = { + 'SEED': seed, + 'MODEL': 'FFT_FILTER', + 'TLEN': 200.0, + 'TAU': 1.0, + } + dict_hier_param = {'MAXHIER': 2} + dict_eom_param = {'EQUATION_OF_MOTION': eom} + + return HOPS( + dict_sys_param, + noise_param=dict_noise_param, + hierarchy_param=dict_hier_param, + eom_param=dict_eom_param, + ) + + +def _make_dyadic_params(n_site, t_max, t_step, max_hier, seed): + """ + Builds shared physical parameters for dyadic vs nondyadic comparisons. + Returns the Hamiltonian, dipoles, field, bath modes, l-operators, and + the chromophore/convergence/noise dicts needed by both sides. + """ + from mesohops.trajectory.dyadic_spectra import ( + prepare_chromophore_input_dict, + prepare_convergence_parameter_dict, + ) + from mesohops.util.bath_corr_functions import bcf_convert_dl_to_exp + + n_state = n_site + 1 + + H2_sys = np.zeros((n_state, n_state), dtype=np.complex128) + H2_exc = np.diag([100.0, 0.0][:n_site]) + np.diag( + [-50.0] * (n_site - 1), k=1 + ) + np.diag([-50.0] * (n_site - 1), k=-1) + H2_sys[1:, 1:] = H2_exc + + list_transition_dipoles = np.array([[0.6, 0.0, 0.3], [0.0, 0.5, 0.4]]) + E_1 = np.array([0.0, 0.0, 1.0]) + + list_modes = bcf_convert_dl_to_exp(50.0, 50.0, 300.0) + list_lop = [ + sparse.coo_matrix(([1], ([i + 1], [i + 1])), shape=(n_state, n_state)) + for i in range(n_site) + ] + + chromophore_dict = prepare_chromophore_input_dict( + list_transition_dipoles, H2_sys, {'list_lop': list_lop, 'list_modes': list_modes} + ) + convergence_dict = prepare_convergence_parameter_dict( + t_step=t_step, max_hier=max_hier + ) + + t_total = 1000.0 + t_max + dict_noise_param = { + 'SEED': seed, + 'MODEL': 'FFT_FILTER', + 'TLEN': t_total, + 'TAU': 0.5, + } + + return (H2_sys, list_transition_dipoles, E_1, n_state, n_site, + chromophore_dict, convergence_dict, dict_noise_param) + + +# ============================================================ +# TEST SUITE: _check_eom() +# ============================================================ + +# ------------------------------------------------------------ +# TEST: rejects NORMALIZED NONLINEAR EOM +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_check_eom_rejects_normalized_nonlinear_abs(): + # This case tests that absorption raises for NORMALIZED NONLINEAR + traj = _make_traj(2, eom='NORMALIZED NONLINEAR') + list_transition_dipoles = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + with pytest.raises(UnsupportedRequest, match='NORMALIZED NONLINEAR'): + calc_absorption_response(traj, list_transition_dipoles, E_1, 4.0, 2.0) + + +@pytest.mark.level(1) +def test_check_eom_rejects_normalized_nonlinear_fluor(): + # This case tests that fluorescence raises for NORMALIZED NONLINEAR + traj = _make_traj(2, eom='NORMALIZED NONLINEAR') + list_transition_dipoles = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + E_sig = np.array([1.0, 0.0, 0.0]) + with pytest.raises(UnsupportedRequest, match='NORMALIZED NONLINEAR'): + calc_fluorescence_response(traj, list_transition_dipoles, E_1, E_sig, 4.0, 6.0, 2.0) + + +# ------------------------------------------------------------ +# TEST: rejects initialized trajectory +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_check_eom_rejects_initialized_traj_abs(): + # This case tests that absorption raises if trajectory is already initialized + traj = _make_traj(2) + n_state = 3 + P1_psi_0 = np.zeros(n_state, dtype=np.complex128) + P1_psi_0[0] = 1.0 + traj.initialize(P1_psi_0) + list_transition_dipoles = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + with pytest.raises(ValueError, match='initialized trajectory'): + calc_absorption_response(traj, list_transition_dipoles, E_1, 4.0, 2.0) + + +@pytest.mark.level(1) +def test_check_eom_rejects_initialized_traj_fluor(): + # This case tests that fluorescence raises if trajectory is already initialized + traj = _make_traj(2) + n_state = 3 + P1_psi_0 = np.zeros(n_state, dtype=np.complex128) + P1_psi_0[0] = 1.0 + traj.initialize(P1_psi_0) + list_transition_dipoles = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + E_sig = np.array([1.0, 0.0, 0.0]) + with pytest.raises(ValueError, match='initialized trajectory'): + calc_fluorescence_response(traj, list_transition_dipoles, E_1, E_sig, 4.0, 6.0, 2.0) + + +# ============================================================ +# TEST SUITE: _build_operators_abs() +# ============================================================ + +# ------------------------------------------------------------ +# TEST: raise operator preserves ground state +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_operators_abs_raise_keeps_ground(): + # This case tests that O2_raise has a 1 at (0, 0) to preserve |g> + list_transition_dipoles = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + O2_raise, _ = _build_operators_abs(list_transition_dipoles, E_1) + O2_dense = O2_raise.toarray() + n_total = O2_dense.shape[0] + # Analytical: ground row preserves ground state + assert O2_dense[0, 0] == 1.0 + # This case tests that ground row has no off-diagonal entries + for j in range(1, n_total): + assert O2_dense[0, j] == 0.0, f'O2_raise[0,{j}] should be 0' + # This case tests that excited rows have entries only in column 0 + for i in range(1, n_total): + for j in range(1, n_total): + assert O2_dense[i, j] == 0.0, f'O2_raise[{i},{j}] should be 0' + + +# ------------------------------------------------------------ +# TEST: raise operator has correct excited entries +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_operators_abs_raise_excited_entries(): + # This case tests that (i+1, 0) entries equal mu_i . E + list_transition_dipoles = np.array([[3.0, 0.0, 0.0], [0.0, 5.0, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + O2_raise, _ = _build_operators_abs(list_transition_dipoles, E_1) + O2_dense = O2_raise.toarray() + # Analytical: O2_raise[i+1, 0] = mu_i . E_1 + # list_transition_dipoles = [[3,0,0],[0,5,0]], E_1 = [1,0,0] + # mu_0 . E_1 = 3.0, mu_1 . E_1 = 0.0 + assert O2_dense[1, 0] == pytest.approx(3.0) + assert O2_dense[2, 0] == pytest.approx(0.0) + + +# ------------------------------------------------------------ +# TEST: F2 response operator structure +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_operators_abs_F2_structure(): + # This case tests that F2 has nonzero entries only in row 0, columns 1: + list_transition_dipoles = np.array([[2.0, 0.0, 0.0], [0.0, 4.0, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + _, F2_dense = _build_operators_abs(list_transition_dipoles, E_1) + assert F2_dense[0, 0] == 0.0 + assert F2_dense[0, 1] == 2.0 + assert F2_dense[0, 2] == 0.0 + assert np.all(F2_dense[1:, :] == 0.0) + + +# ============================================================ +# TEST SUITE: _build_operators_fluor() +# ============================================================ + +# ------------------------------------------------------------ +# TEST: Fluorescence operators have correct structure +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_operators_fluor_structure(): + # This case tests that the raising, lowering+identity, and response + # operators constructed for fluorescence have analytically correct entries. + # list_transition_dipoles = [[1,0,0],[0,1,0]], E_1 = [1,0,0], E_sig = [0,1,0] + # mu_0.E_1 = 1.0, mu_1.E_1 = 0.0 + # mu_0.E_sig = 0.0, mu_1.E_sig = 1.0 + list_transition_dipoles = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + E_sig = np.array([0.0, 1.0, 0.0]) + O2_raise, O2_lower_ident, F2 = _build_operators_fluor(list_transition_dipoles, E_1, E_sig) + O2_raise_dense = O2_raise.toarray() + # Analytical: O2_raise[i+1, 0] = mu_i . E_1 + assert O2_raise_dense[1, 0] == pytest.approx(1.0) # mu_0 . E_1 = 1.0 + assert O2_raise_dense[2, 0] == pytest.approx(0.0) # mu_1 . E_1 = 0.0 + # Analytical: O2_lower_ident has |g> [3, 1, 0], so the norm ratio is 10 / 2 = 5. + n_site = 2 + n_state = n_site + 1 + traj = _make_traj(n_site, seed=0) + P1_psi_0 = np.zeros(n_state, dtype=np.complex128) + P1_psi_0[:] = np.array([1.0, 1.0, 0.0], dtype=np.complex128) + traj.initialize(P1_psi_0) + op = np.array([ + [1.0, 2.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + ], dtype=np.complex128) + norm_ratio = _apply_op_and_track_norm( + traj, + lambda: traj._operator(op), + SpectroscopyDispatch('vector', None, 'embedded', 'NL'), + ) + np.testing.assert_allclose(norm_ratio, 5.0, atol=1e-10) + + +class _FakeVacuumTrajectory: + def __init__(self, full_state): + self._full_state = np.array(full_state, dtype=np.complex128) + self.basis = SimpleNamespace( + eom=SimpleNamespace(param={'EQUATION_OF_MOTION': 'NONLINEAR'}) + ) + self.wavefunction = self + + @property + def psi(self): + return self._full_state[1:] + + @property + def manifold_norm_sq(self): + return np.sum(np.conj(self._full_state) * self._full_state).real + + def _operator(self, op): + self._full_state = op @ self._full_state + + +# ------------------------------------------------------------ +# TEST: Vacuum branch uses manifold norm rather than psi norm +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_apply_op_and_track_norm_vacuum_branch(): + # This case ensures the vacuum-convention branch uses the manifold + # norm helper before and after the operator application. + traj = _FakeVacuumTrajectory([1.0, 1.0, 0.0]) + op = np.array([ + [1.0, 1.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + ], dtype=np.complex128) + norm_ratio = _apply_op_and_track_norm( + traj, + lambda: traj._operator(op), + SpectroscopyDispatch('tensor', 'number', 'vacuum', 'NL'), + ) + # Full-state norm changes from ||[1,1,0]||^2 = 2 to ||[2,1,0]||^2 = 5. + np.testing.assert_allclose(norm_ratio, 2.5, atol=1e-10) + + +# ============================================================ +# TEST SUITE: calc_absorption_response() +# ============================================================ + +# ------------------------------------------------------------ +# TEST: correct output shape +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_calc_absorption_response_shape(): + # This case tests output length matches number of propagation time steps + traj = _make_traj(2, seed=42) + list_transition_dipoles = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + C1_corr_t = calc_absorption_response(traj, list_transition_dipoles, E_1, 10.0, 2.0) + assert len(C1_corr_t) == 5 + + +# ------------------------------------------------------------ +# TEST: deterministic with same seed +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_calc_absorption_response_deterministic(): + list_transition_dipoles = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + + traj1 = _make_traj(2, seed=7) + C1_a = calc_absorption_response(traj1, list_transition_dipoles, E_1, 10.0, 2.0) + + traj2 = _make_traj(2, seed=7) + C1_b = calc_absorption_response(traj2, list_transition_dipoles, E_1, 10.0, 2.0) + + np.testing.assert_allclose(C1_a, C1_b) + + +# ------------------------------------------------------------ +# TEST: C(0) is nonzero +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_calc_absorption_response_c0_value(): + traj = _make_traj(2, seed=0) + list_transition_dipoles = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + C1_corr_t = calc_absorption_response(traj, list_transition_dipoles, E_1, 4.0, 2.0) + # Analytical: C(t_step) ≈ 2 * sum_i |mu_i . E_1|^2 for short times. + # Factor of 2 accounts for the complex conjugate pathway. + # list_transition_dipoles = [[1,0,0],[0,1,0]], E_1 = [1,0,0] + # mu_0 . E_1 = 1.0, mu_1 . E_1 = 0.0 + # C(t_step) ≈ 2 * (|1|^2 + |0|^2) = 2.0 within stochastic tolerance + np.testing.assert_allclose( + abs(C1_corr_t[0]), 2.0, atol=1e-2, + err_msg='C(t_step) should be close to 2 * sum(|mu_i . E|^2) = 2.0', + ) + + +# ------------------------------------------------------------ +# TEST: single-site system (n_site=1) +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_calc_absorption_response_single_site(): + traj = _make_traj(1, seed=0) + list_transition_dipoles = np.array([[1.0, 0.0, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + C1_corr_t = calc_absorption_response(traj, list_transition_dipoles, E_1, 4.0, 2.0) + assert len(C1_corr_t) == 2 + assert abs(C1_corr_t[0]) > 0.0 + + +# ------------------------------------------------------------ +# TEST: LINEAR EOM short-time amplitude +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_calc_absorption_response_linear_c0_value(): + # Locks in the LINEAR readout convention (see _readout_prefactor). + # Under LINEAR psi[0]=1 exactly, so (t_step) ≈ |mu|^2 + # and the function should return ≈ 2*|mu|^2 = 2.0 for this dipole + # choice. A regression that re-applied the dyadic + # dyadic normalization prefactor under LINEAR would return ≈ 4.0 + # (extra (1+|mu|^2) factor at t≈0). + traj = _make_traj(2, seed=0, eom='LINEAR') + list_transition_dipoles = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + C1_corr_t = calc_absorption_response(traj, list_transition_dipoles, E_1, 4.0, 2.0) + np.testing.assert_allclose( + abs(C1_corr_t[0]), 2.0, atol=1e-2, + err_msg='LINEAR C(t_step) should match NONLINEAR analytical limit ' + '2 * |mu|^2 = 2.0; a (1+|mu|^2) prefactor leak would give 4.0', + ) + + +# ------------------------------------------------------------ +# TEST: LINEAR readout matches the analytical correlator form +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_calc_absorption_response_linear_correlator_form(): + # The (N+1)-embedded operators decouple |g> from the bath in H + # and in every L (operator-construction property, true under all + # EOMs). Under LINEAR — raw noise, no rescaling term — psi(t)[0] + # = 1 holds deterministically, so reduces to + # sum_i (mu_i.E) * psi(t)[i+1], the absorption correlator + # directly. The function should return 2x this with no extra + # norm machinery. Catches any regression that re-applies the + # dyadic normalization prefactor (or any extra scalar factor) under + # LINEAR, at machine precision with no ensemble floor. + list_transition_dipoles = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + traj = _make_traj(2, seed=7, eom='LINEAR') + C1_func = calc_absorption_response(traj, list_transition_dipoles, E_1, 20.0, 2.0) + + # psi[0] = 1 invariant under LINEAR — structural property of the setup + psi_traj = np.asarray(traj.storage['psi_traj']) + np.testing.assert_allclose( + psi_traj[:, 0], 1.0, atol=1e-10, + err_msg='LINEAR should hold psi[0]=1 (embedded |g> decoupling + ' + 'no rescaling term)', + ) + + # Manual correlator: 2 * sum_i (mu_i.E) * psi(t)[i+1] + list_mu_dot_E = list_transition_dipoles @ E_1 + C1_manual = 2 * (psi_traj[1:, 1:] @ list_mu_dot_E) + np.testing.assert_allclose(C1_func, C1_manual, atol=1e-10) + + +# ------------------------------------------------------------ +# TEST: psi[0]=1 invariant under LINEAR and plain NONLINEAR +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_psi_g_invariant_under_linear_and_nonlinear(): + # Structural property of the (N+1)-embedded picture: |g> is + # decoupled from the bath in H and in every L. Combined with + # the absence of any global-rescaling term (which only + # NORMALIZED NONLINEAR has), this fixes psi[0]=1 deterministically + # under BOTH LINEAR and plain NONLINEAR. The two EOMs differ in + # the noise measure (raw vs Girsanov-shifted), not in psi[0] + # dynamics — pinning this for both guards against any future + # change to H, the L construction, or the raise/lower operators + # that would silently couple |g> to the bath and break the + # readout invariants both branches of _readout_prefactor depend + # on. + list_transition_dipoles = np.array([[0.6, 0.0, 0.3], [0.0, 0.5, 0.4]]) + E_1 = np.array([0.0, 0.0, 1.0]) + for eom in ('LINEAR', 'NONLINEAR'): + traj = _make_traj(2, seed=7, eom=eom) + calc_absorption_response(traj, list_transition_dipoles, E_1, 20.0, 2.0) + psi_traj = np.asarray(traj.storage['psi_traj']) + np.testing.assert_allclose( + psi_traj[:, 0], 1.0, atol=1e-10, + err_msg=f'{eom}: psi[0] should stay at 1 in the embedded picture', + ) + + +# ------------------------------------------------------------ +# TEST: NORMALIZED NONLINEAR is a uniform rescaling of plain NONLINEAR +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_normalized_nonlinear_rescales_plain_nonlinear(): + # Central claim of the corrected readout rationale: under the + # same noise seed, ψ_NORM(t) = c(t) · ψ_NL(t) where c(t) is a + # complex scalar with |c(t)|^2 equal to the post-raise norm ratio + # divided by ||ψ_NL(t)||^2. This + # holds because operator_expectation in the EOM is normalized + # (divides by <ψ|ψ>), so is rescale-invariant and the + # NL/NORM-NL EOMs differ only by a global scale factor c(t) + # satisfying dc/dt = -norm_corr · c. The rescaling collapses the + # dyadic readout normalization prefactor times <ψ|F|ψ>/||ψ||^2 to the same per-realization + # value for both EOMs, which is what makes the existing dyadic- + # vs-nondyadic comparison pass at rtol=5e-3 and what justifies + # using the dyadic prefactor for plain NL inside _readout_prefactor. + # Pinning this identity guards against any change to + # operator_expectation or the EOM that would silently break it. + n_site = 2 + n_state = n_site + 1 + seed = 7 + t_max, t_step = 10.0, 2.0 + + list_transition_dipoles = np.array([[0.6, 0.0, 0.3], [0.0, 0.5, 0.4]]) + E_1 = np.array([0.0, 0.0, 1.0]) + O2_raise, _ = _build_operators_abs(list_transition_dipoles, E_1) + + P1_psi_0 = np.zeros(n_state, dtype=np.complex128) + P1_psi_0[0] = 1.0 + + traj_nl = _make_traj(n_site, seed=seed, eom='NONLINEAR') + traj_nl.initialize(P1_psi_0) + traj_nl._operator(O2_raise.toarray()) + traj_nl.propagate(t_max, t_step) + psi_nl = np.asarray(traj_nl.storage['psi_traj']) + + traj_norm = _make_traj(n_site, seed=seed, eom='NORMALIZED NONLINEAR') + traj_norm.initialize(P1_psi_0) + traj_norm._operator(O2_raise.toarray()) + traj_norm.propagate(t_max, t_step) + psi_norm = np.asarray(traj_norm.storage['psi_traj']) + + # ψ_NL[t, 0] = 1 (|g> decoupled, no rescaling under plain NL), so + # c(t) = ψ_NORM[t, 0] / ψ_NL[t, 0] = ψ_NORM[t, 0]. Verify all + # other components rescale by the same scalar. + coeff_traj = psi_norm[:, 0] + psi_predicted = coeff_traj[:, None] * psi_nl + np.testing.assert_allclose( + psi_norm, psi_predicted, atol=1e-6, + err_msg='NORMALIZED NL should be a uniform (scalar) rescaling ' + 'of plain NL — non-uniformity means is no longer ' + 'rescale-invariant', + ) + + # Norm-preservation under NORM-NL plus the rescaling identity + # forces |c(t)|^2 = post-raise norm ratio / ||ψ_NL(t)||^2. + coeff_post_op = np.dot(np.conj(psi_nl[0]), psi_nl[0]).real + list_norm_sq_nl = np.sum(np.conj(psi_nl) * psi_nl, axis=1).real + np.testing.assert_allclose( + np.abs(coeff_traj) ** 2, coeff_post_op / list_norm_sq_nl, atol=1e-6, + err_msg='|c(t)|^2 should equal the post-raise norm ratio / ||ψ_NL(t)||^2', + ) + + +# ============================================================ +# TEST SUITE: calc_fluorescence_response() +# ============================================================ + +# ------------------------------------------------------------ +# TEST: correct output shape +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_calc_fluorescence_response_shape(): + traj = _make_traj(2, seed=42) + list_transition_dipoles = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + E_sig = np.array([1.0, 0.0, 0.0]) + C1_corr_t = calc_fluorescence_response( + traj, list_transition_dipoles, E_1, E_sig, 4.0, 10.0, 2.0 + ) + assert len(C1_corr_t) == 5 + + +# ------------------------------------------------------------ +# TEST: deterministic with same seed +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_calc_fluorescence_response_deterministic(): + list_transition_dipoles = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + E_sig = np.array([1.0, 0.0, 0.0]) + + traj1 = _make_traj(2, seed=7) + C1_a = calc_fluorescence_response(traj1, list_transition_dipoles, E_1, E_sig, 4.0, 10.0, 2.0) + + traj2 = _make_traj(2, seed=7) + C1_b = calc_fluorescence_response(traj2, list_transition_dipoles, E_1, E_sig, 4.0, 10.0, 2.0) + + np.testing.assert_allclose(C1_a, C1_b) + + +# ------------------------------------------------------------ +# TEST: non-trivial dipole weighting produces nonzero result +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_calc_fluorescence_response_two_coefficients(): + traj = _make_traj(2, seed=0) + list_transition_dipoles = np.array([[1.0, 0.5, 0.0], [0.5, 1.0, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + E_sig = np.array([0.0, 1.0, 0.0]) + C1_corr_t = calc_fluorescence_response( + traj, list_transition_dipoles, E_1, E_sig, 4.0, 6.0, 2.0 + ) + # Invariant: fluorescence response should be complex and non-trivial + assert np.iscomplexobj(C1_corr_t), 'Response should be complex' + assert np.any(np.abs(C1_corr_t) > 1e-12), ( + 'Fluorescence response should be nonzero for nonzero dipoles' + ) + + +# ------------------------------------------------------------ +# TEST: LINEAR fluorescence readout matches analytical correlator +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_calc_fluorescence_response_linear_correlator_form(): + # The (N+1)-embedded operators decouple |g> from the bath under + # all EOMs, and LINEAR has no global-rescaling term, so after + # the lowering+identity operator at the end of t2 lifts psi[0] + # to c_g, psi[0] stays at c_g exactly through the t3 phase. The + # detection readout therefore reduces to + # 4 * conj(c_g) * sum_i (mu_i.E_sig) * psi(t3)[i+1]. Catches + # any regression that re-applies the dyadic normalization prefactor + # (or any extra scalar factor) under LINEAR. + list_transition_dipoles = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + E_sig = np.array([1.0, 0.0, 0.0]) + t2 = 4.0 + t3_max = 10.0 + t_step = 2.0 + traj = _make_traj(2, seed=7, eom='LINEAR') + C1_func = calc_fluorescence_response( + traj, list_transition_dipoles, E_1, E_sig, t2, t3_max, t_step + ) + + # Pull the stored wavefunction trajectory; the detection-phase + # window starts one step after de-excitation (t = t2 + t_step). + psi_traj = np.asarray(traj.storage['psi_traj']) + idx_t2 = round(t2 / t_step) + psi_t3 = psi_traj[idx_t2 + 1:] + + # |g> decoupled under LINEAR -> psi[0] constant during t3 + np.testing.assert_allclose( + psi_t3[:, 0], psi_t3[0, 0], atol=1e-10, + err_msg='psi[0] should be constant during t3 under LINEAR', + ) + + # Manual correlator: 4 * conj(psi[0]) * sum_i (mu_i.E_sig) * psi[i+1] + list_mu_dot_Esig = list_transition_dipoles @ E_sig + C1_manual = 4 * np.conj(psi_t3[:, 0]) * (psi_t3[:, 1:] @ list_mu_dot_Esig) + np.testing.assert_allclose(C1_func, C1_manual, atol=1e-10) + + +# ============================================================ +# TEST SUITE: dyadic vs nondyadic comparison +# ============================================================ + +# ------------------------------------------------------------ +# TEST: absorption C(t) matches DyadicSpectra +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_absorption_matches_dyadic_spectra(): + """ + Compares non-dyadic absorption against DyadicSpectra with identical + physical parameters and the same noise seed. Dipoles are scaled so + S = sum((mu_i.E)^2) = 1. The dyadic formulation uses NORMALIZED + NONLINEAR while the nondyadic uses NONLINEAR, so a small tolerance + accounts for normalization drift. + """ + from mesohops.trajectory.dyadic_spectra import ( + DyadicSpectra as DHOPS, + ) + from mesohops.trajectory.dyadic_spectra import ( + prepare_spectroscopy_input_dict, + ) + + n_site = 2 + seed = 42 + t_max = 50.0 + t_step = 1.0 + max_hier = 4 + + (H2_sys, list_transition_dipoles, E_1, n_state, n_site, + chromophore_dict, convergence_dict, dict_noise_param) = _make_dyadic_params( + n_site, t_max, t_step, max_hier, seed + ) + + # Dyadic calculation + spec_dict = prepare_spectroscopy_input_dict( + 'ABSORPTION', + {'t_1': t_max}, + {'E_1': E_1}, + {'list_ket_sites': np.arange(1, n_site + 1)}, + ) + dhops = DHOPS(spec_dict, chromophore_dict, convergence_dict, seed) + C1_dyadic = dhops.calculate_spectrum() + + # Non-dyadic calculation with matching parameters + dict_sys_param = { + 'HAMILTONIAN': H2_sys, + 'GW_SYSBATH': chromophore_dict['gw_sysbath_hier'], + 'L_HIER': chromophore_dict['lop_list_hier'], + 'L_NOISE1': chromophore_dict['lop_list_noise'], + 'ALPHA_NOISE1': bcf_exp, + 'PARAM_NOISE1': chromophore_dict['gw_sysbath_noise'], + } + traj = HOPS( + dict_sys_param, + noise_param=dict_noise_param, + hierarchy_param={'MAXHIER': max_hier}, + eom_param={'EQUATION_OF_MOTION': 'NONLINEAR'}, + ) + C1_nondyadic = calc_absorption_response(traj, list_transition_dipoles, E_1, t_max, t_step) + + np.testing.assert_allclose(C1_nondyadic, C1_dyadic, rtol=5e-3) + + +# ------------------------------------------------------------ +# TEST: fluorescence C(t) matches DyadicSpectra +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_fluorescence_matches_dyadic_spectra(): + """ + Compares non-dyadic fluorescence against DyadicSpectra with identical + physical parameters and the same noise seed. + """ + from mesohops.trajectory.dyadic_spectra import ( + DyadicSpectra as DHOPS, + ) + from mesohops.trajectory.dyadic_spectra import ( + prepare_spectroscopy_input_dict, + ) + + n_site = 2 + seed = 42 + t2 = 50.0 + t3_max = 50.0 + t_step = 1.0 + max_hier = 4 + + (H2_sys, list_transition_dipoles, E_1, n_state, n_site, + chromophore_dict, convergence_dict, dict_noise_param) = _make_dyadic_params( + n_site, t2 + t3_max, t_step, max_hier, seed + ) + E_sig = E_1 + + # Dyadic calculation + spec_dict = prepare_spectroscopy_input_dict( + 'FLUORESCENCE', + {'t_2': t2, 't_3': t3_max}, + {'E_1': E_1, 'E_sig': E_sig}, + { + 'list_ket_sites': np.arange(1, n_site + 1), + 'list_bra_sites': np.arange(1, n_site + 1), + }, + ) + dhops = DHOPS(spec_dict, chromophore_dict, convergence_dict, seed) + C1_dyadic = dhops.calculate_spectrum() + + # Non-dyadic calculation with matching parameters + dict_sys_param = { + 'HAMILTONIAN': H2_sys, + 'GW_SYSBATH': chromophore_dict['gw_sysbath_hier'], + 'L_HIER': chromophore_dict['lop_list_hier'], + 'L_NOISE1': chromophore_dict['lop_list_noise'], + 'ALPHA_NOISE1': bcf_exp, + 'PARAM_NOISE1': chromophore_dict['gw_sysbath_noise'], + } + traj = HOPS( + dict_sys_param, + noise_param=dict_noise_param, + hierarchy_param={'MAXHIER': max_hier}, + eom_param={'EQUATION_OF_MOTION': 'NONLINEAR'}, + ) + C1_nondyadic = calc_fluorescence_response( + traj, list_transition_dipoles, E_1, E_sig, t2, t3_max, t_step + ) + + np.testing.assert_allclose(C1_nondyadic, C1_dyadic, rtol=5e-3) + + +# ============================================================ +# TEST SUITE: HopsTensorTrajectory compatibility +# ============================================================ + +def _make_tensor_traj(n_site, seed=0, eom='NONLINEAR'): + """ + Builds an uninitialized HopsTensorTrajectory with the given number + of chromophore sites (total states = n_site + 1). + """ + n_state = n_site + 1 + + H2_sys = np.zeros((n_state, n_state), dtype=np.complex128) + H2_exc = np.diag([100.0, 0.0][:n_site]) + np.diag( + [-50.0] * (n_site - 1), k=1 + ) + np.diag([-50.0] * (n_site - 1), k=-1) + H2_sys[1:, 1:] = H2_exc + + list_lop = [] + for i in range(n_site): + lop = np.zeros((n_state, n_state), dtype=np.float64) + lop[i + 1, i + 1] = 1.0 + list_lop.append(lop) + + gw_sysbath = [[10.0, 10.0]] * n_site + + dict_sys_param = { + 'HAMILTONIAN': H2_sys, + 'GW_SYSBATH': gw_sysbath, + 'L_HIER': list_lop, + 'L_NOISE1': list_lop, + 'ALPHA_NOISE1': bcf_exp, + 'PARAM_NOISE1': gw_sysbath, + } + dict_noise_param = { + 'SEED': seed, + 'MODEL': 'FFT_FILTER', + 'TLEN': 200.0, + 'TAU': 1.0, + } + dict_hier_param = {'MAXHIER': 2} + dict_eom_param = {'EQUATION_OF_MOTION': eom} + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, + } + + return HopsTensorTrajectory( + dict_sys_param, + noise_param=dict_noise_param, + hierarchy_param=dict_hier_param, + eom_param=dict_eom_param, + tensor_param=tensor_param, + ) + + +# ------------------------------------------------------------ +# TEST: tensor absorption produces correct shape +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_tensor_absorption_shape(): + # This case tests that calc_absorption_response works with + # HopsTensorTrajectory and returns the correct output length. + traj = _make_tensor_traj(2, seed=42) + list_transition_dipoles = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + C1_corr_t = calc_absorption_response(traj, list_transition_dipoles, E_1, 10.0, 2.0) + assert len(C1_corr_t) == 5 + + +# ------------------------------------------------------------ +# TEST: tensor and vector absorption agree +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_tensor_absorption_matches_vector(): + # This case tests that tensor and vector trajectories produce + # the same absorption C(t) for the same seed and parameters. + list_transition_dipoles = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + + traj_vec = _make_traj(2, seed=7) + C1_vec = calc_absorption_response(traj_vec, list_transition_dipoles, E_1, 10.0, 2.0) + + traj_tensor = _make_tensor_traj(2, seed=7) + C1_tensor = calc_absorption_response(traj_tensor, list_transition_dipoles, E_1, 10.0, 2.0) + + np.testing.assert_allclose(C1_tensor, C1_vec, atol=1e-6) + + +# ------------------------------------------------------------ +# TEST: tensor fluorescence produces correct shape +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_tensor_fluorescence_shape(): + # This case tests that calc_fluorescence_response works with + # HopsTensorTrajectory and returns the correct output length. + traj = _make_tensor_traj(2, seed=42) + list_transition_dipoles = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + E_sig = np.array([1.0, 0.0, 0.0]) + C1_corr_t = calc_fluorescence_response( + traj, list_transition_dipoles, E_1, E_sig, 4.0, 10.0, 2.0 + ) + assert len(C1_corr_t) == 5 + + +# ------------------------------------------------------------ +# TEST: tensor and vector fluorescence agree +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_tensor_fluorescence_matches_vector(): + # This case tests that tensor and vector trajectories produce + # the same fluorescence C(t) for the same seed and parameters. + list_transition_dipoles = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + E_sig = np.array([1.0, 0.0, 0.0]) + + traj_vec = _make_traj(2, seed=7) + C1_vec = calc_fluorescence_response(traj_vec, list_transition_dipoles, E_1, E_sig, 4.0, 10.0, 2.0) + + traj_tensor = _make_tensor_traj(2, seed=7) + C1_tensor = calc_fluorescence_response( + traj_tensor, list_transition_dipoles, E_1, E_sig, 4.0, 10.0, 2.0 + ) + + np.testing.assert_allclose(C1_tensor, C1_vec, atol=1e-6) + + +# ------------------------------------------------------------ +# TEST: tensor LINEAR absorption per-trajectory identity +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_tensor_absorption_response_linear_correlator_form(): + # _readout_prefactor branches on the EOM string, which is + # populated identically on HopsTrajectory and HopsTensorTrajectory. + # Asserts the same per-realization identity as the vector LINEAR + # test (test_calc_absorption_response_linear_correlator_form): + # under the (N+1)-embedded picture |g> is decoupled in H and L, + # LINEAR has no rescaling term, so psi(t)[0]=1 deterministically + # and 2* = 2 * sum_i (mu_i.E) * psi(t)[i+1]. + list_transition_dipoles = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + traj = _make_tensor_traj(2, seed=7, eom='LINEAR') + C1_func = calc_absorption_response(traj, list_transition_dipoles, E_1, 10.0, 2.0) + + psi_traj = np.asarray(traj.storage['psi_traj']) + np.testing.assert_allclose( + psi_traj[:, 0], 1.0, atol=1e-10, + err_msg='LINEAR (tensor) should hold psi[0]=1', + ) + + list_mu_dot_E = list_transition_dipoles @ E_1 + C1_manual = 2 * (psi_traj[1:, 1:] @ list_mu_dot_E) + np.testing.assert_allclose(C1_func, C1_manual, atol=1e-10) + + +# ------------------------------------------------------------ +# TEST: tensor LINEAR fluorescence per-trajectory identity +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_tensor_fluorescence_response_linear_correlator_form(): + # Tensor analog of test_calc_fluorescence_response_linear_correlator_form. + # After the lowering+identity operator at the end of t2, psi[0] + # picks up a nonzero amplitude c_g and stays there during t3 + # under LINEAR (|g> decoupled in the embedded picture, no + # rescaling term). Detection readout reduces to + # 4 * conj(c_g) * sum_i (mu_i.E_sig) * psi(t3)[i+1]. + list_transition_dipoles = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + E_sig = np.array([1.0, 0.0, 0.0]) + t2 = 4.0 + t3_max = 10.0 + t_step = 2.0 + traj = _make_tensor_traj(2, seed=7, eom='LINEAR') + C1_func = calc_fluorescence_response( + traj, list_transition_dipoles, E_1, E_sig, t2, t3_max, t_step + ) + + psi_traj = np.asarray(traj.storage['psi_traj']) + idx_t2 = round(t2 / t_step) + psi_t3 = psi_traj[idx_t2 + 1:] + + np.testing.assert_allclose( + psi_t3[:, 0], psi_t3[0, 0], atol=1e-10, + err_msg='psi[0] should be constant during t3 under LINEAR (tensor)', + ) + + list_mu_dot_Esig = list_transition_dipoles @ E_sig + C1_manual = 4 * np.conj(psi_t3[:, 0]) * ( + psi_t3[:, 1:] @ list_mu_dot_Esig + ) + np.testing.assert_allclose(C1_func, C1_manual, atol=1e-10) + + +def _make_tensor_traj_sn(n_site, seed=0): + """ + Builds an uninitialized number-representation HopsTensorTrajectory in the + vacuum convention: NSTATES = n_site, the ground state is the all-zeros + MPS configuration. + """ + n_state = n_site + + H2_sys = ( + np.diag([100.0, 0.0][:n_site]) + + np.diag([-50.0] * (n_site - 1), k=1) + + np.diag([-50.0] * (n_site - 1), k=-1) + ).astype(np.complex128) + + list_lop = [] + for i in range(n_site): + lop = np.zeros((n_state, n_state), dtype=np.float64) + lop[i, i] = 1.0 + list_lop.append(lop) + + gw_sysbath = [[10.0, 10.0]] * n_site + + dict_sys_param = { + 'HAMILTONIAN': H2_sys, + 'GW_SYSBATH': gw_sysbath, + 'L_HIER': list_lop, + 'L_NOISE1': list_lop, + 'ALPHA_NOISE1': bcf_exp, + 'PARAM_NOISE1': gw_sysbath, + } + dict_noise_param = { + 'SEED': seed, + 'MODEL': 'FFT_FILTER', + 'TLEN': 200.0, + 'TAU': 1.0, + } + dict_hier_param = {'MAXHIER': 2} + dict_eom_param = {'EQUATION_OF_MOTION': 'NONLINEAR'} + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'number', + 'BOND_DIM_MAX': 20, + } + + return HopsTensorTrajectory( + dict_sys_param, + noise_param=dict_noise_param, + hierarchy_param=dict_hier_param, + eom_param=dict_eom_param, + tensor_param=tensor_param, + ) + + +# ------------------------------------------------------------ +# TEST: statenumber absorption matches vector +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_statenumber_absorption_matches_vector(): + # CASE: Statenumber tensor trajectory should produce the same + # absorption C(t) as vector HOPS for the same seed. + list_transition_dipoles = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + + traj_vec = _make_traj(2, seed=7) + C1_vec = calc_absorption_response(traj_vec, list_transition_dipoles, E_1, 10.0, 2.0) + + traj_sn = _make_tensor_traj_sn(2, seed=7) + C1_sn = calc_absorption_response(traj_sn, list_transition_dipoles, E_1, 10.0, 2.0) + + np.testing.assert_allclose(C1_sn, C1_vec, atol=1e-6) + + +# ------------------------------------------------------------ +# TEST: statenumber absorption matches fullstate (tight tolerance) +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_statenumber_absorption_matches_fullstate(): + # CASE: Both representations use the same tensor EOM — differences + # are only SVD truncation, which is negligible for this system. + list_transition_dipoles = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + + traj_fs = _make_tensor_traj(2, seed=7) + C1_fs = calc_absorption_response(traj_fs, list_transition_dipoles, E_1, 10.0, 2.0) + + traj_sn = _make_tensor_traj_sn(2, seed=7) + C1_sn = calc_absorption_response(traj_sn, list_transition_dipoles, E_1, 10.0, 2.0) + + np.testing.assert_allclose(C1_sn, C1_fs, atol=1e-10) + + +# ------------------------------------------------------------ +# TEST: statenumber fluorescence matches vector +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_statenumber_fluorescence_matches_vector(): + # CASE: Statenumber tensor trajectory should produce the same + # fluorescence C(t) as vector HOPS for the same seed. + list_transition_dipoles = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + E_sig = np.array([1.0, 0.0, 0.0]) + + traj_vec = _make_traj(2, seed=7) + C1_vec = calc_fluorescence_response( + traj_vec, list_transition_dipoles, E_1, E_sig, 4.0, 10.0, 2.0, + ) + + traj_sn = _make_tensor_traj_sn(2, seed=7) + C1_sn = calc_fluorescence_response( + traj_sn, list_transition_dipoles, E_1, E_sig, 4.0, 10.0, 2.0, + ) + + np.testing.assert_allclose(C1_sn, C1_vec, atol=1e-6) + + +# ------------------------------------------------------------ +# TEST: statenumber fluorescence matches fullstate (tight tolerance) +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_statenumber_fluorescence_matches_fullstate(): + # CASE: Both representations use the same tensor EOM — differences + # are only SVD truncation, which is negligible for this system. + list_transition_dipoles = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + E_sig = np.array([1.0, 0.0, 0.0]) + + traj_fs = _make_tensor_traj(2, seed=7) + C1_fs = calc_fluorescence_response( + traj_fs, list_transition_dipoles, E_1, E_sig, 4.0, 10.0, 2.0, + ) + + traj_sn = _make_tensor_traj_sn(2, seed=7) + C1_sn = calc_fluorescence_response( + traj_sn, list_transition_dipoles, E_1, E_sig, 4.0, 10.0, 2.0, + ) + + np.testing.assert_allclose(C1_sn, C1_fs, atol=1e-10) + +# ============================================================ + +def _make_traj_excited_only(n_site, seed=0, eom='LINEAR'): + """ + Builds an uninitialized HopsTrajectory in the excited-only (N-dim) + layout — H, L_HIER, L_NOISE1 are (n_site, n_site) with no + embedded |g> state. Mirrors _make_traj otherwise so the per-site + excited-block dynamics match between the two layouts under LINEAR + with a shared seed. + """ + n_state = n_site + + if n_site == 1: + H2_sys = np.array([[100.0]], dtype=np.complex128) + else: + H2_sys = ( + np.diag([100.0, 0.0][:n_site]) + + np.diag([-50.0] * (n_site - 1), k=1) + + np.diag([-50.0] * (n_site - 1), k=-1) + ).astype(np.complex128) + + list_lop = [] + for i in range(n_site): + lop = np.zeros((n_state, n_state), dtype=np.float64) + lop[i, i] = 1.0 + list_lop.append(lop) + + gw_sysbath = [[10.0, 10.0]] * n_site + + dict_sys_param = { + 'HAMILTONIAN': H2_sys, + 'GW_SYSBATH': gw_sysbath, + 'L_HIER': list_lop, + 'L_NOISE1': list_lop, + 'ALPHA_NOISE1': bcf_exp, + 'PARAM_NOISE1': gw_sysbath, + } + dict_noise_param = { + 'SEED': seed, + 'MODEL': 'FFT_FILTER', + 'TLEN': 200.0, + 'TAU': 1.0, + } + dict_hier_param = {'MAXHIER': 2} + dict_eom_param = {'EQUATION_OF_MOTION': eom} + + return HOPS( + dict_sys_param, + noise_param=dict_noise_param, + hierarchy_param=dict_hier_param, + eom_param=dict_eom_param, + ) + + +def _make_tensor_traj_excited_only(n_site, seed=0, eom='LINEAR'): + """ + Tensor analog of _make_traj_excited_only. + """ + n_state = n_site + + if n_site == 1: + H2_sys = np.array([[100.0]], dtype=np.complex128) + else: + H2_sys = ( + np.diag([100.0, 0.0][:n_site]) + + np.diag([-50.0] * (n_site - 1), k=1) + + np.diag([-50.0] * (n_site - 1), k=-1) + ).astype(np.complex128) + + list_lop = [] + for i in range(n_site): + lop = np.zeros((n_state, n_state), dtype=np.float64) + lop[i, i] = 1.0 + list_lop.append(lop) + + gw_sysbath = [[10.0, 10.0]] * n_site + + dict_sys_param = { + 'HAMILTONIAN': H2_sys, + 'GW_SYSBATH': gw_sysbath, + 'L_HIER': list_lop, + 'L_NOISE1': list_lop, + 'ALPHA_NOISE1': bcf_exp, + 'PARAM_NOISE1': gw_sysbath, + } + dict_noise_param = { + 'SEED': seed, + 'MODEL': 'FFT_FILTER', + 'TLEN': 200.0, + 'TAU': 1.0, + } + dict_hier_param = {'MAXHIER': 2} + dict_eom_param = {'EQUATION_OF_MOTION': eom} + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, + } + + return HopsTensorTrajectory( + dict_sys_param, + noise_param=dict_noise_param, + hierarchy_param=dict_hier_param, + eom_param=dict_eom_param, + tensor_param=tensor_param, + ) + + +# ------------------------------------------------------------ +# TEST: N-dim and (N+1)-dim LINEAR keys agree per-realization +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_calc_absorption_response_linear_keys_match(): + # Keystone correctness test for the excited-only dispatch key. + # Under LINEAR the (N+1) embedding has psi[0]=1 deterministically + # (|g> decoupled from H and every L), and the noise is keyed per + # mode-index, so with the same SEED, TLEN, TAU, MAXHIER, and + # matched bath modes per site the excited-block dynamics are + # bitwise identical between the two layouts. Their C(t) outputs + # must therefore agree at machine precision. A regression in the + # excited-only readout (wrong seed normalization, wrong conj convention, + # off-by-one in the time slice) breaks this immediately. + list_transition_dipoles = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + + traj_embed = _make_traj(2, seed=7, eom='LINEAR') + C1_embed = calc_absorption_response(traj_embed, list_transition_dipoles, E_1, 20.0, 2.0) + + traj_excited = _make_traj_excited_only(2, seed=7, eom='LINEAR') + C1_excited = calc_absorption_response( + traj_excited, list_transition_dipoles, E_1, 20.0, 2.0 + ) + + np.testing.assert_allclose(C1_embed, C1_excited, atol=1e-10) + + +# ------------------------------------------------------------ +# TEST: rejects NONLINEAR with N-dim trajectory +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_calc_absorption_response_rejects_nonlinear_excited_only(): + # The N-dim key is LINEAR-only because the NONLINEAR mean-field + # term = / requires the + # |g> bra-norm denominator from the embedded layout. Combining + # NONLINEAR with an excited-only trajectory must raise. + traj = _make_traj_excited_only(2, eom='NONLINEAR') + list_transition_dipoles = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + with pytest.raises(UnsupportedRequest, match='vector_excited_only_NL'): + calc_absorption_response(traj, list_transition_dipoles, E_1, 4.0, 2.0) + + +# ------------------------------------------------------------ +# TEST: rejects trajectory dim that is neither n_site nor n_site+1 +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_calc_absorption_response_rejects_bad_dim(): + # Trajectory NSTATES=2 with list_transition_dipoles of n_site=4 means accepted + # dims are {4, 5}; 2 is in neither bucket. Confirm the dispatcher + # raises rather than silently picking a key. + traj = _make_traj_excited_only(2, eom='LINEAR') # NSTATES = 2 + list_transition_dipoles = np.array([ + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.5, 0.5, 0.0], + [0.0, 0.0, 1.0], + ]) # n_site = 4, accepted dims = {4, 5} + E_1 = np.array([1.0, 0.0, 0.0]) + with pytest.raises(ValueError, match='trajectory dim'): + calc_absorption_response(traj, list_transition_dipoles, E_1, 4.0, 2.0) + + +# ------------------------------------------------------------ +# TEST: tensor and vector excited-only LINEAR agree +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_tensor_absorption_excited_only_matches_vector(): + # The N-dim key runs through HopsTensorTrajectory's tensor EOM + # without the |g> leg the existing tensor tests exercise. Confirm + # tensor and vector backends produce the same C(t) for the + # excited-only LINEAR layout, mirroring test_tensor_absorption_ + # matches_vector for the embedded key. + list_transition_dipoles = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + + traj_vec = _make_traj_excited_only(2, seed=7, eom='LINEAR') + C1_vec = calc_absorption_response(traj_vec, list_transition_dipoles, E_1, 10.0, 2.0) + + traj_tensor = _make_tensor_traj_excited_only(2, seed=7, eom='LINEAR') + C1_tensor = calc_absorption_response( + traj_tensor, list_transition_dipoles, E_1, 10.0, 2.0 + ) + + np.testing.assert_allclose(C1_tensor, C1_vec, atol=1e-6) + + +# Shared builder for the vacuum-vs-gs-core regression tests below. +# The leading underscore keeps pytest from collecting it as a test. +def _make_tensor_traj_for_compare(n_site, use_gs, seed=0, eom='NONLINEAR', + noise_model='FFT_FILTER'): + """ + Tensor trajectory in either GS-as-state-core or vacuum convention, + with caller-selected EOM and noise model so the absorption (LINEAR + + ZERO) and fluorescence (NL + FFT_FILTER) regressions can share a + single builder. + + use_gs=True : NSTATES = n_site + 1, H placed at [1:, 1:], L-ops at + site+1, method='fullstate'. The GS slot + rides as the index-0 physical state. + use_gs=False: NSTATES = n_site, H placed at the bare excited block, + L-ops at site, method='number' with + flag_gs_vacuum=True. Same physics, no GS slot. + """ + # Site energies: first site elevated, rest at 0. Generalizes the + # earlier [100.0, 0.0][:n_site] pattern to arbitrary n_site so the + # trimer regression below can use the same helper. + list_site_energies = np.array( + [100.0] + [0.0] * (n_site - 1), dtype=np.complex128, + ) + + if use_gs: + n_state = n_site + 1 + H2_sys = np.zeros((n_state, n_state), dtype=np.complex128) + H2_exc = ( + np.diag(list_site_energies) + + np.diag(np.full(n_site - 1, -50.0, dtype=np.complex128), k=1) + + np.diag(np.full(n_site - 1, -50.0, dtype=np.complex128), k=-1) + ) + H2_sys[1:, 1:] = H2_exc + + list_lop = [] + for i in range(n_site): + lop = np.zeros((n_state, n_state), dtype=np.float64) + lop[i + 1, i + 1] = 1.0 + list_lop.append(lop) + else: + n_state = n_site + H2_sys = ( + np.diag(list_site_energies) + + np.diag(np.full(n_site - 1, -50.0, dtype=np.complex128), k=1) + + np.diag(np.full(n_site - 1, -50.0, dtype=np.complex128), k=-1) + ) + + list_lop = [] + for i in range(n_site): + lop = np.zeros((n_state, n_state), dtype=np.float64) + lop[i, i] = 1.0 + list_lop.append(lop) + + gw_sysbath = [[10.0, 10.0]] * n_site + + dict_sys_param = { + 'HAMILTONIAN': H2_sys, + 'GW_SYSBATH': gw_sysbath, + 'L_HIER': list_lop, + 'L_NOISE1': list_lop, + 'ALPHA_NOISE1': bcf_exp, + 'PARAM_NOISE1': gw_sysbath, + } + dict_noise_param = { + 'SEED': seed, + 'MODEL': noise_model, + 'TLEN': 200.0, + 'TAU': 1.0, + } + dict_hier_param = {'MAXHIER': 2} + dict_eom_param = {'EQUATION_OF_MOTION': eom} + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': ( + 'fullstate' if use_gs + else 'number' + ), + 'BOND_DIM_MAX': 20, + } + + return HopsTensorTrajectory( + dict_sys_param, + noise_param=dict_noise_param, + hierarchy_param=dict_hier_param, + eom_param=dict_eom_param, + tensor_param=tensor_param, + ) + + +# ------------------------------------------------------------ +# TEST: absorption — vacuum convention matches gs_core dimer +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_absorption_vacuum_matches_gs_core_dimer(): + """ + LINEAR + ZERO-noise absorption on a dimer. The vacuum-convention + trajectory (NSTATES=2, all-zeros |g>) and the gs_core trajectory + (NSTATES=3, psi[0]=1 |g>) represent the same physical system and + must produce the same C_abs(t) within MPS truncation noise. + + Uses asymmetric non-orthogonal dipoles so both site W-blocks of + the raise MPO contribute non-trivially: orthogonal mu = [1, 0] + would zero the mu_2 column and silently hide site-2 bugs. + """ + n_site = 2 + list_transition_dipoles = np.array([[0.6, 0.0, 0.3], [0.4, 0.5, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + t_max = 10.0 + t_step = 2.0 + + traj_gs = _make_tensor_traj_for_compare( + n_site, use_gs=True, seed=11, eom='LINEAR', noise_model='ZERO', + ) + C_gs = calc_absorption_response(traj_gs, list_transition_dipoles, E_1, t_max, t_step) + + traj_vac = _make_tensor_traj_for_compare( + n_site, use_gs=False, seed=11, eom='LINEAR', noise_model='ZERO', + ) + C_vac = calc_absorption_response( + traj_vac, list_transition_dipoles, E_1, t_max, t_step, + ) + + assert C_gs.shape == C_vac.shape + np.testing.assert_allclose(C_vac, C_gs, atol=1e-12, rtol=1e-10) + + +# ------------------------------------------------------------ +# TEST: absorption NL — vacuum convention matches gs_core dimer +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_absorption_vacuum_matches_gs_core_dimer_nl(): + """ + NL absorption on a dimer with a fixed-seed FFT_FILTER noise + realization. Same noise -> same trajectory dynamics in both + conventions; C_abs(t) must agree within MPS truncation noise. + + Exercises the flag_gs_vacuum=True branch of + HopsTensorEOM.compute_z_mem_update (adds |gs_amp|^2 to the + mean-field denominator under NL), which the LINEAR + ZERO + sibling above does not touch. + """ + n_site = 2 + list_transition_dipoles = np.array([[0.6, 0.0, 0.3], [0.4, 0.5, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + t_max = 10.0 + t_step = 2.0 + + traj_gs = _make_tensor_traj_for_compare(n_site, use_gs=True, seed=17) + C_gs = calc_absorption_response(traj_gs, list_transition_dipoles, E_1, t_max, t_step) + + traj_vac = _make_tensor_traj_for_compare(n_site, use_gs=False, seed=17) + C_vac = calc_absorption_response( + traj_vac, list_transition_dipoles, E_1, t_max, t_step, + ) + + assert C_gs.shape == C_vac.shape + np.testing.assert_allclose(C_vac, C_gs, atol=1e-12, rtol=1e-10) + + +# ------------------------------------------------------------ +# TEST: fluorescence — vacuum convention matches gs_core dimer +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_fluorescence_vacuum_matches_gs_core_dimer(): + """ + NL fluorescence on a dimer with a fixed-seed FFT_FILTER noise + realization. Same noise -> same trajectory dynamics in both + conventions; C_fl(t) must agree within MPS truncation noise. + """ + n_site = 2 + list_transition_dipoles = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + E_1 = np.array([1.0, 0.0, 0.0]) + E_sig = np.array([1.0, 0.0, 0.0]) + t2 = 4.0 + t3_max = 10.0 + t_step = 2.0 + + traj_gs = _make_tensor_traj_for_compare(n_site, use_gs=True, seed=13) + C_gs = calc_fluorescence_response( + traj_gs, list_transition_dipoles, E_1, E_sig, t2, t3_max, t_step, + ) + + traj_vac = _make_tensor_traj_for_compare(n_site, use_gs=False, seed=13) + C_vac = calc_fluorescence_response( + traj_vac, list_transition_dipoles, E_1, E_sig, t2, t3_max, t_step, + ) + + assert C_gs.shape == C_vac.shape + np.testing.assert_allclose(C_vac, C_gs, atol=1e-12, rtol=1e-10) + + +# ------------------------------------------------------------ +# TEST: fluorescence on a trimer with non-uniform dipoles. +# Exercises the interior state-core W-matrices of both new +# MPOs (bond-dim-2 raise and bond-dim-3 lower+ident) — N=2 +# has no interior core, so this case is the first one that +# actually contracts a (3, 2, 2, 3) interior block. +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_fluorescence_vacuum_matches_gs_core_trimer(): + """ + Trimer fluorescence regression with multi-cartesian dipoles and + distinct E_1 != E_sig, so that: + - y / z columns of list_transition_dipoles contribute (not just x), + - the raise pathway (mu . E_1) and the lower pathway + (mu . E_sig) use distinct effective dipole vectors. + Catches bugs in non-x cartesian channels and in any asymmetry + between the raise and lower paths. + """ + n_site = 3 + list_transition_dipoles = np.array([ + [1.0, 0.2, 0.3], + [0.7, 0.5, 0.1], + [0.4, 0.6, 0.2], + ]) + E_1 = np.array([1.0, 0.0, 0.0]) + E_sig = np.array([0.0, 1.0, 0.0]) + t2 = 4.0 + t3_max = 10.0 + t_step = 2.0 + + traj_gs = _make_tensor_traj_for_compare(n_site, use_gs=True, seed=29) + C_gs = calc_fluorescence_response( + traj_gs, list_transition_dipoles, E_1, E_sig, t2, t3_max, t_step, + ) + + traj_vac = _make_tensor_traj_for_compare(n_site, use_gs=False, seed=29) + C_vac = calc_fluorescence_response( + traj_vac, list_transition_dipoles, E_1, E_sig, t2, t3_max, t_step, + ) + + assert C_gs.shape == C_vac.shape + np.testing.assert_allclose(C_vac, C_gs, atol=1e-12, rtol=1e-10) + + +# ============================================================ +# TEST SUITE: _spectroscopy_key() +# ============================================================ + +# ------------------------------------------------------------ +# TEST: vector trajectory at dim n_site+1 maps to the embedded key +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_spectroscopy_key_vector_embedded(): + traj = _make_traj(2) + assert _spectroscopy_key(traj, 2) == SpectroscopyDispatch( + 'vector', None, 'embedded', 'NL' + ) + + +# ------------------------------------------------------------ +# TEST: vector trajectory at dim n_site maps to the excited-only key +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_spectroscopy_key_vector_excited_only(): + traj = _make_traj_excited_only(2) + assert _spectroscopy_key(traj, 2) == SpectroscopyDispatch( + 'vector', None, 'excited_only', 'LINEAR' + ) + + +# ------------------------------------------------------------ +# TEST: fullstate tensor trajectory maps to the embedded key +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_spectroscopy_key_fullstate_embedded(): + traj = _make_tensor_traj(2) + assert _spectroscopy_key(traj, 2) == SpectroscopyDispatch( + 'tensor', 'fullstate', 'embedded', 'NL' + ) + + +# ------------------------------------------------------------ +# TEST: number trajectory takes the vacuum key and flags the wavefunction +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_spectroscopy_key_number_vacuum_sets_flag(): + traj = _make_tensor_traj_sn(2) + assert _spectroscopy_key(traj, 2) == SpectroscopyDispatch( + 'tensor', 'number', 'vacuum', 'NL' + ) + assert traj.wavefunction.flag_gs_vacuum is True + + +# ------------------------------------------------------------ +# TEST: number representation rejects the embedded convention +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_spectroscopy_key_number_embedded_raises(): + # Dim 2 read as the embedded layout of a 1-site system; number trajectories + # only support the vacuum convention. + traj = _make_tensor_traj_sn(2) + with pytest.raises(UnsupportedRequest, match='tensor_number_embedded'): + _spectroscopy_key(traj, 1) + + +# ------------------------------------------------------------ +# TEST: rejects an EOM outside the supported set +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_spectroscopy_key_rejects_unsupported_eom(): + traj = _make_traj(2, eom='NORMALIZED NONLINEAR') + with pytest.raises(UnsupportedRequest, match='NORMALIZED NONLINEAR'): + _spectroscopy_key(traj, 2) + + +# ------------------------------------------------------------ +# TEST: rejects a Hilbert dimension that is neither n_site nor n_site+1 +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_spectroscopy_key_rejects_bad_dimension(): + traj = _make_traj(2) + with pytest.raises(ValueError, match='!='): + _spectroscopy_key(traj, 5) + + +# ------------------------------------------------------------ +# TEST: rejects an already-initialized trajectory +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_spectroscopy_key_rejects_initialized_traj(): + traj = _make_traj(2) + P1_psi_0 = np.zeros(3, dtype=np.complex128) + P1_psi_0[0] = 1.0 + traj.initialize(P1_psi_0) + with pytest.raises(ValueError, match='initialized trajectory'): + _spectroscopy_key(traj, 2) diff --git a/tests/test_tdvp.py b/tests/test_tdvp.py new file mode 100644 index 0000000..79b9426 --- /dev/null +++ b/tests/test_tdvp.py @@ -0,0 +1,2522 @@ +import numpy as np +import pytest + +from mesohops.tensor.tdvp import ( + _apply_heff_bond, + _apply_heff_site, + _apply_heff_twosite, + _arnoldi_expm, + _ivp_solve, + _lanczos_expm, + _solve_local, + _split_qr, + _split_rq, + _split_svd, + contract_left, + contract_right, + initialize, + recenter_to_zero, + sweep_left_1tdvp, + sweep_left_2tdvp, + sweep_right_1tdvp, + sweep_right_2tdvp, + timestep, +) + +__title__ = 'Unit Tests for TDVP Internals' +__author__ = 'N. Covalsen' +__maintainer__ = 'N. Covalsen' + + +# ============================================================ +# TEST SUITE: _split_qr() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Q is left-orthogonal and Q @ R recovers input +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_split_qr_orthogonal_and_recovers(): + # This case tests that Q has orthonormal columns (Q†Q = I) + # and Q @ R recovers the original tensor. + np.random.seed(10) + core_input = np.random.randn(2, 3, 4) + 1j * np.random.randn(2, 3, 4) + q, r = _split_qr(core_input) + # Q†Q = I (left-orthogonal over combined left-physical index) + Dl, d, Dr_q = q.shape + Q_mat = q.reshape(Dl * d, Dr_q) + np.testing.assert_allclose( + Q_mat.conj().T @ Q_mat, + np.eye(Dr_q), + atol=1e-12, + err_msg='Q is not left-orthogonal', + ) + # Q @ R recovers original + recovered = np.tensordot(q, r, axes=((2,), (0,))) + np.testing.assert_allclose( + recovered, core_input, atol=1e-12, err_msg='Q @ R does not recover input' + ) + + +# ============================================================ +# TEST SUITE: _split_rq() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Q is right-orthogonal and R @ Q recovers input +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_split_rq_orthogonal_and_recovers(): + # This case tests that Q has orthonormal rows (QQ† = I over + # combined physical-right index) and R @ Q recovers the original. + np.random.seed(11) + core_input = np.random.randn(4, 3, 2) + 1j * np.random.randn(4, 3, 2) + r, q = _split_rq(core_input) + # Q is right-orthogonal: reshape (Dl_q, d, Dr) -> (Dl_q, d*Dr), check M M† = I + Dl_q, d, Dr = q.shape + Q_mat = q.reshape(Dl_q, d * Dr) + np.testing.assert_allclose( + Q_mat @ Q_mat.conj().T, + np.eye(Dl_q), + atol=1e-12, + err_msg='Q is not right-orthogonal', + ) + # R @ Q recovers original + recovered = np.tensordot(r, q, axes=((1,), (0,))) + np.testing.assert_allclose( + recovered, core_input, atol=1e-12, err_msg='R @ Q does not recover input' + ) + + +# ============================================================ +# TEST SUITE: _split_svd() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: eps threshold keeps only SVs above cutoff +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_split_svd_eps_threshold(): + # This case tests that eps=0.5 drops the singular value 0.01 < 0.5, + # leaving only the single SV 10.0 → chi_new=1. + # Build theta by embedding a known 2x2 diagonal matrix in shape (1, 4, 1): + # _split_svd reshapes (1, 4, 1) → (2, 2), giving SVs [10, 0.01]. + d = 2 + mat = np.diag([10.0, 0.01]).astype(np.complex128) + theta = mat.reshape(1, d * d, 1) + U, S, Vt, chi_new = _split_svd(theta, 0, 0.5) + assert chi_new == 1, f'Expected chi_new=1, got {chi_new}' + assert U.shape == (1, d, 1) + assert Vt.shape == (1, d, 1) + np.testing.assert_allclose(S[0], 10.0, atol=1e-12) + + +# ------------------------------------------------------------ +# TEST: chi_max caps rank even when SVs are above eps +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_split_svd_rank_cap(): + # This case tests that chi_max=2 limits the rank even when + # all singular values exceed eps. + np.random.seed(13) + d = 3 + theta = np.random.randn(1, d * d, 1) + 1j * np.random.randn(1, d * d, 1) + U, S, Vt, chi_new = _split_svd(theta, 2, 0.0) + assert chi_new == 2, f'Expected chi_new=2, got {chi_new}' + assert U.shape == (1, d, 2) + assert Vt.shape == (2, d, 1) + # Analytical: the retained singular values should be the 2 largest + _, S_full, _, _ = _split_svd(theta, 0, 0.0) + S_sorted = np.sort(S_full)[::-1] + np.testing.assert_allclose( + np.sort(S)[::-1], S_sorted[:2], atol=1e-12, + err_msg='Retained SVs should be the 2 largest', + ) + + +# ------------------------------------------------------------ +# TEST: U @ diag(S) @ Vt recovers the input tensor +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_split_svd_recovers_input(): + # This case tests that contracting U, S, Vt back together + # recovers the original two-site tensor (up to numerical precision). + np.random.seed(14) + d = 3 + Dl, Dr = 2, 2 + theta = np.random.randn(Dl, d * d, Dr) + 1j * np.random.randn(Dl, d * d, Dr) + U, S, Vt, chi_new = _split_svd(theta, 0, 1e-14) + # U: (Dl, d, chi), S: (chi,), Vt: (chi, d, Dr) + # recovery: sum_k U[a,s1,k] * S[k] * Vt[k,s2,b] → (Dl, d, d, Dr) → reshape + recovered = np.tensordot(U * S[None, None, :], Vt, axes=((2,), (0,))).reshape( + Dl, d * d, Dr + ) + np.testing.assert_allclose( + recovered, theta, atol=1e-10, err_msg='U @ diag(S) @ Vt does not recover input' + ) + + +# ------------------------------------------------------------ +# TEST: U is left-orthogonal after truncation +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_split_svd_u_left_orthogonal(): + # This case tests that the U factor satisfies U†U = I (left-orthogonal). + np.random.seed(15) + d = 3 + Dl, Dr = 2, 2 + theta = np.random.randn(Dl, d * d, Dr) + 1j * np.random.randn(Dl, d * d, Dr) + U, S, Vt, chi_new = _split_svd(theta, 0, 1e-14) + U_mat = U.reshape(Dl * d, chi_new) + np.testing.assert_allclose( + U_mat.conj().T @ U_mat, + np.eye(chi_new), + atol=1e-12, + err_msg='U is not left-orthogonal', + ) + + +# ------------------------------------------------------------ +# TEST: Vt factor from _split_svd is right-orthogonal +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_split_svd_vt_right_orthogonal(): + # This case tests that the Vt factor satisfies Vt Vt† = I. + np.random.seed(16) + d = 3 + Dl, Dr = 2, 2 + theta = np.random.randn(Dl, d * d, Dr) + 1j * np.random.randn(Dl, d * d, Dr) + U, S, Vt, chi_new = _split_svd(theta, 0, 1e-14) + Vt_mat = Vt.reshape(chi_new, d * Dr) + np.testing.assert_allclose( + Vt_mat @ Vt_mat.conj().T, + np.eye(chi_new), + atol=1e-12, + err_msg='Vt is not right-orthogonal', + ) + + +# ============================================================ +# TEST SUITE: _ivp_solve() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Identity RHS scales state by exp(dt) +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_ivp_solve_identity_rhs(): + # This case tests that dy/dt = y gives y(t) = exp(t) * y(0). + state = np.array([1.0 + 0j, 0.5 + 0j]) + dt = 0.1 + result = _ivp_solve(state, lambda v: v, dt) + np.testing.assert_allclose(result, np.exp(dt) * state, atol=1e-6) + + +# ============================================================ +# TEST SUITE: _arnoldi_expm() +# ============================================================ + + +def _trivial_envs(): + """Build trivial (1,1,1) left and right environment tensors.""" + env_left = np.ones((1, 1, 1), dtype=np.complex128) + env_right = np.ones((1, 1, 1), dtype=np.complex128) + return env_left, env_right + + +def _identity_mpo_core(d): + """Build an identity MPO core of shape (1, d, d, 1).""" + core_mpo = np.zeros((1, d, d, 1), dtype=np.complex128) + for k in range(d): + core_mpo[0, k, k, 0] = 1.0 + return core_mpo + + +def _diagonal_mpo_core(diag): + """Build a diagonal MPO core from a 1-D array of eigenvalues.""" + d = len(diag) + core_mpo = np.zeros((1, d, d, 1), dtype=np.complex128) + for k in range(d): + core_mpo[0, k, k, 0] = diag[k] + return core_mpo + + +# ------------------------------------------------------------ +# TEST: Identity MPO with trivial environments scales by exp(dt) +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_arnoldi_expm_identity_mpo_scales_by_exp(): + # This case tests that the effective operator is the identity, so + # exp(dt * I) * state = exp(dt) * state. + d = 3 + state = np.random.RandomState(31).randn(1, d, 1) + 0j + env_left, env_right = _trivial_envs() + core_mpo = _identity_mpo_core(d) + dt = 0.05 + result = _arnoldi_expm( + state, + lambda v: _apply_heff_site(v, env_left, core_mpo, env_right), + dt, + 1e-12, + ) + assert result.shape == state.shape + np.testing.assert_allclose(result, np.exp(dt) * state, atol=1e-10) + + +# ------------------------------------------------------------ +# TEST: Bond update with trivial environments scales by exp(dt) +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_arnoldi_expm_bond_update_trivial_envs(): + # This case tests the bond-update path. With trivial (ones) environments, + # _apply_heff_bond is the identity on (1,1) bond tensors, so + # exp(dt * I) * state = exp(dt) * state. + state = np.array([[3.0 + 1j]], dtype=np.complex128) + env_left, env_right = _trivial_envs() + dt = 0.1 + result = _arnoldi_expm( + state, + lambda v: _apply_heff_bond(v, env_left, env_right), + dt, + 1e-12, + ) + assert result.shape == state.shape + np.testing.assert_allclose(result, np.exp(dt) * state, atol=1e-10) + + +# ------------------------------------------------------------ +# TEST: Known diagonal Hamiltonian matches analytical result +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_arnoldi_expm_known_diagonal_hamiltonian(): + # This case tests _arnoldi_expm against an analytical result for a diagonal + # Hamiltonian (Pauli Z). The effective operator with trivial environments + # is diag(+1, -1). For state with both components nonzero, the Krylov + # subspace spans the full 2D space, giving an exact result: + # exp(dt * Z) @ state = [exp(dt)*a, exp(-dt)*b] + d = 2 + rng = np.random.RandomState(34) + state = (rng.randn(1, d, 1) + 1j * rng.randn(1, d, 1)).astype(np.complex128) + env_left, env_right = _trivial_envs() + core_mpo = _diagonal_mpo_core([1.0, -1.0]) + dt = 0.05 + result = _arnoldi_expm( + state, + lambda v: _apply_heff_site(v, env_left, core_mpo, env_right), + dt, + 1e-12, + ) + expected = state.copy() + expected[0, 0, 0] *= np.exp(dt) + expected[0, 1, 0] *= np.exp(-dt) + np.testing.assert_allclose(result, expected, atol=1e-10) + + +# ------------------------------------------------------------ +# TEST: Convergence on rank-1 effective Hamiltonian +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_arnoldi_expm_early_termination(): + # This case tests that the convergence check terminates iteration + # after one step for the identity MPO (1D Krylov space): h_{2,1} = 0 + # → the error estimate is exactly zero → conv_tol check passes + # immediately. Result must equal exp(dt) * state. + d = 3 + state = np.random.RandomState(35).randn(1, d, 1) + 0j + env_left, env_right = _trivial_envs() + core_mpo = _identity_mpo_core(d) + dt = 0.02 + matvec_heff = lambda v: _apply_heff_site(v, env_left, core_mpo, env_right) # noqa: E731 + result = _arnoldi_expm(state, matvec_heff, dt, 1e-12) + np.testing.assert_allclose(result, np.exp(dt) * state, atol=1e-10) + + +# ------------------------------------------------------------ +# TEST: Full Krylov loop with non-Hermitian off-diagonal MPO +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_arnoldi_expm_full_krylov_loop(): + # This case tests _arnoldi_expm with a non-trivial off-diagonal MPO + # that does NOT trigger early breakdown, forcing the full Krylov loop. + # We use a 3x3 nilpotent-shift matrix as the effective Hamiltonian. + d = 3 + # Build an MPO core that acts as a shift: |i> -> |i+1 mod d> + core_mpo = np.zeros((1, d, d, 1), dtype=complex) + for i in range(d): + core_mpo[0, (i + 1) % d, i, 0] = 1.0 + env_left = np.ones((1, 1, 1), dtype=complex) + env_right = np.ones((1, 1, 1), dtype=complex) + state = np.zeros((1, d, 1), dtype=complex) + state[0, 0, 0] = 1.0 + dt = 0.1 + matvec = lambda v: _apply_heff_site(v, env_left, core_mpo, env_right) # noqa: E731 + result = _arnoldi_expm(state, matvec, dt, conv_tol=1e-14) + # Reference: build the d x d shift matrix and compute expm directly + H_shift = np.zeros((d, d), dtype=complex) + for i in range(d): + H_shift[(i + 1) % d, i] = 1.0 + from scipy.linalg import expm as scipy_expm + expected_vec = scipy_expm(dt * H_shift) @ np.array([1, 0, 0], dtype=complex) + expected = expected_vec.reshape(1, d, 1) + np.testing.assert_allclose(result, expected, atol=1e-10) + + +# ------------------------------------------------------------ +# TEST: _arnoldi_expm returns zeros for zero-vector input +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_arnoldi_expm_zero_vector(): + # Analytical: if input is zero, output must be zero (beta < tol path) + phys_dims = [2, 3] + list_cores_mpo = _make_list_mpo(phys_dims) + L = np.ones((1, 1, 1), dtype=np.complex128) + R = np.ones((1, 1, 1), dtype=np.complex128) + W = list_cores_mpo[0] + v = np.zeros((1, phys_dims[0], 1), dtype=np.complex128) + matvec = lambda x, _L=L, _W=W, _R=R: _apply_heff_site(x, _L, _W, _R) # noqa: E731 + result = _arnoldi_expm(v, matvec, 0.01, 1e-12) + np.testing.assert_allclose(result, 0.0, atol=1e-14) + + +# ============================================================ +# TEST SUITE: contract_left() / contract_right() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Output shape is correct for both directions +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_contract_right_output_shape_and_value(): + # contract_right(R_next, W, B): R_next (Dr,bR,Dr), W (bL,d,d,bR), + # B (Dl,d,Dr) → output (Dl, bL, Dl) + d = 2 + Dl, Dr, bL, bR = 3, 4, 2, 2 + rng = np.random.RandomState(41) + core_mps = rng.randn(Dl, d, Dr) + 1j * rng.randn(Dl, d, Dr) + core_mpo = rng.randn(bL, d, d, bR) + 1j * rng.randn(bL, d, d, bR) + env_init = rng.randn(Dr, bR, Dr) + 1j * rng.randn(Dr, bR, Dr) + # Compute via contract_right + result = contract_right(env_init, core_mpo, core_mps) + assert result.shape == (Dl, bL, Dl), f'Expected {(Dl, bL, Dl)}, got {result.shape}' + # Verify value against direct einsum + expected = np.einsum( + 'ijk,lmi,nmoj,pok->lnp', + env_init, + core_mps.conj(), + core_mpo, + core_mps, + ) + np.testing.assert_allclose( + result, + expected, + atol=1e-12, + err_msg='contract_right value mismatch vs direct einsum', + ) + + +# ------------------------------------------------------------ +# TEST: Output shape and value are correct for contract_left +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_contract_left_output_shape_and_value(): + # contract_left(L_prev, W, A): L_prev (Dl,bL,Dl), W (bL,d,d,bR), + # A (Dl,d,Dr) → output (Dr, bR, Dr) + d = 2 + Dl, Dr, bL, bR = 3, 4, 2, 2 + rng = np.random.RandomState(42) + core_mps = rng.randn(Dl, d, Dr) + 1j * rng.randn(Dl, d, Dr) + core_mpo = rng.randn(bL, d, d, bR) + 1j * rng.randn(bL, d, d, bR) + env_init = rng.randn(Dl, bL, Dl) + 1j * rng.randn(Dl, bL, Dl) + # Compute via contract_left + result = contract_left(env_init, core_mpo, core_mps) + assert result.shape == (Dr, bR, Dr), f'Expected {(Dr, bR, Dr)}, got {result.shape}' + # Verify value against direct einsum + expected = np.einsum( + 'ijk,ilm,jlno,knp->mop', + env_init, + core_mps.conj(), + core_mpo, + core_mps, + ) + np.testing.assert_allclose( + result, + expected, + atol=1e-12, + err_msg='contract_left value mismatch vs direct einsum', + ) + + +# ------------------------------------------------------------ +# TEST: Full sweep builds consistent environments +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_contract_right_full_sweep_consistency(): + # This case tests that a right-to-left sweep with identity MPO + # on a random 3-site MPS produces a final environment equal to + # the MPS norm squared. + rng = np.random.RandomState(44) + phys_dims = [2, 3, 2] + list_cores = [rng.randn(1, d, 1) + 1j * rng.randn(1, d, 1) for d in phys_dims] + list_cores_mpo = [_identity_mpo_core(d) for d in phys_dims] + + phi = np.ones((1, 1, 1), dtype=np.complex128) + for i in range(len(list_cores) - 1, -1, -1): + phi = contract_right(phi, list_cores_mpo[i], list_cores[i]) + + # Compute norm squared by full contraction + state = list_cores[0] + for c in list_cores[1:]: + state = np.tensordot(state, c, axes=([-1], [0])) + norm_sq = np.vdot(state, state) + np.testing.assert_allclose( + phi.item(), + norm_sq, + atol=1e-10, + err_msg='Full R-to-L sweep should give norm squared', + ) + + +def _identity_mpo_core_2site(d): + """Build identity MPO list_cores for two adjacent sites of dim d.""" + core_mpo_site1 = np.zeros((1, d, d, 1), dtype=np.complex128) + core_mpo_site2 = np.zeros((1, d, d, 1), dtype=np.complex128) + for k in range(d): + core_mpo_site1[0, k, k, 0] = 1.0 + core_mpo_site2[0, k, k, 0] = 1.0 + return core_mpo_site1, core_mpo_site2 + + +def _make_list_cores(phys_dims, seed=0): + """Build random bond-dim-1 MPS list_cores as a plain list.""" + rng = np.random.RandomState(seed) + return [rng.randn(1, d, 1) + 1j * rng.randn(1, d, 1) for d in phys_dims] + + +def _make_list_mpo(phys_dims): + """Build an identity MPO as a plain list of list_cores.""" + return [_identity_mpo_core(d) for d in phys_dims] + + +def _mps_norm(cores): + """Compute from MPS cores via full contraction.""" + env = np.ones((1, 1), dtype=np.complex128) + for c in cores: + env = np.einsum('ij,ikl,jkm->lm', env, c, np.conj(c)) + return np.sqrt(np.abs(env[0, 0])) + + +# ============================================================ +# TEST SUITE: _lanczos_expm() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Identity MPO with trivial environments scales by exp(dt) +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_lanczos_expm_identity_mpo_scales_by_exp(): + # This case tests that the Hermitian identity operator gives + # exp(dt * I) * state = exp(dt) * state. + d = 3 + state = np.random.RandomState(71).randn(1, d, 1) + 0j + env_left, env_right = _trivial_envs() + core_mpo = _identity_mpo_core(d) + dt = 0.05 + result = _lanczos_expm( + state, + lambda v: _apply_heff_site(v, env_left, core_mpo, env_right), + dt, + 1e-12, + ) + np.testing.assert_allclose(result, np.exp(dt) * state, atol=1e-10) + + +# ------------------------------------------------------------ +# TEST: Known diagonal Hamiltonian matches analytical result +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_lanczos_expm_known_diagonal_hamiltonian(): + # This case tests _lanczos_expm against an analytical result for + # the Hermitian diagonal Hamiltonian diag(+1, -1) (Pauli Z). + # exp(dt * Z) @ [a, b] = [exp(dt)*a, exp(-dt)*b] + d = 2 + rng = np.random.RandomState(72) + state = (rng.randn(1, d, 1) + 1j * rng.randn(1, d, 1)).astype(np.complex128) + env_left, env_right = _trivial_envs() + core_mpo = _diagonal_mpo_core([1.0, -1.0]) + dt = 0.05 + result = _lanczos_expm( + state, + lambda v: _apply_heff_site(v, env_left, core_mpo, env_right), + dt, + 1e-12, + ) + expected = state.copy() + expected[0, 0, 0] *= np.exp(dt) + expected[0, 1, 0] *= np.exp(-dt) + np.testing.assert_allclose(result, expected, atol=1e-10) + + +# ------------------------------------------------------------ +# TEST: Convergence on rank-1 effective Hamiltonian +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_lanczos_expm_early_termination(): + # This case tests that the convergence check terminates iteration + # after one step for the identity MPO (1D Krylov space): β_1 = 0 + # → the error estimate is exactly zero → conv_tol check passes + # immediately. Result must equal exp(dt) * state. + d = 3 + state = np.random.RandomState(73).randn(1, d, 1) + 0j + env_left, env_right = _trivial_envs() + core_mpo = _identity_mpo_core(d) + dt = 0.02 + matvec_heff = lambda v: _apply_heff_site(v, env_left, core_mpo, env_right) # noqa: E731 + result = _lanczos_expm(state, matvec_heff, dt, 1e-12) + np.testing.assert_allclose(result, np.exp(dt) * state, atol=1e-10) + + +# ------------------------------------------------------------ +# TEST: _lanczos_expm returns zeros for zero-vector input +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_lanczos_expm_zero_vector(): + # Analytical: if input is zero, output must be zero (beta < tol path) + phys_dims = [2, 3] + list_cores_mpo = _make_list_mpo(phys_dims) + L = np.ones((1, 1, 1), dtype=np.complex128) + R = np.ones((1, 1, 1), dtype=np.complex128) + W = list_cores_mpo[0] + v = np.zeros((1, phys_dims[0], 1), dtype=np.complex128) + matvec = lambda x, _L=L, _W=W, _R=R: _apply_heff_site(x, _L, _W, _R) # noqa: E731 + result = _lanczos_expm(v, matvec, 0.01, 1e-12) + np.testing.assert_allclose(result, 0.0, atol=1e-14) + + +# ============================================================ +# TEST SUITE: _solve_local() — dispatch branches +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Lanczos dispatch produces correct result +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_solve_local_lanczos_dispatch(): + # This case tests that solver='lanczos' dispatches to _lanczos_expm + # and produces exp(dt * I) * state = exp(dt) * state with identity operator. + d = 2 + state = np.random.RandomState(80).randn(1, d, 1) + 0j + env_left, env_right = _trivial_envs() + core_mpo = _identity_mpo_core(d) + dt = 0.05 + result = _solve_local( + state, + lambda v: _apply_heff_site(v, env_left, core_mpo, env_right), + dt, + solver='lanczos', + ) + np.testing.assert_allclose(result, np.exp(dt) * state, atol=1e-10) + + +# ------------------------------------------------------------ +# TEST: IVP dispatch produces correct result +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_solve_local_ivp_dispatch(): + # This case tests that solver='ivp' dispatches to _ivp_solve + # and produces exp(dt) * state with identity operator. + d = 2 + state = np.random.RandomState(81).randn(1, d, 1) + 0j + dt = 0.05 + result = _solve_local(state, lambda v: v, dt, solver='ivp') + np.testing.assert_allclose(result, np.exp(dt) * state, atol=1e-6) + + +# ------------------------------------------------------------ +# TEST: Unknown solver raises ValueError +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_solve_local_unknown_solver_raises(): + # This case tests that an invalid solver string raises ValueError. + d = 2 + state = np.random.RandomState(82).randn(1, d, 1) + 0j + with pytest.raises(ValueError, match='Unknown solver'): + _solve_local(state, lambda v: v, 0.01, solver='bogus') + + +# ============================================================ +# TEST SUITE: _ivp_solve() — negative dt and custom kwargs +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Negative dt gives backward integration +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_ivp_solve_negative_dt(): + # This case tests that dt < 0 activates the sign = -1.0 branch, + # giving dy/dt = -y → y(|dt|) = exp(-|dt|) * y(0), equivalent + # to exp(dt) * y(0) with dt < 0. + state = np.array([1.0 + 0j, 0.5 + 0j]) + dt = -0.1 + result = _ivp_solve(state, lambda v: v, dt) + np.testing.assert_allclose(result, np.exp(dt) * state, atol=1e-6) + + +# ------------------------------------------------------------ +# TEST: Custom kwargs forwarded to solve_ivp +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_ivp_solve_custom_kwargs(): + # This case tests that method, rtol, atol kwargs are forwarded + # to solve_ivp. Uses RK45 with tighter tolerances. + state = np.array([1.0 + 0j, 0.5 + 0j]) + dt = 0.1 + result = _ivp_solve(state, lambda v: v, dt, method='RK45', rtol=1e-10, atol=1e-12) + np.testing.assert_allclose(result, np.exp(dt) * state, atol=1e-8) + + +# ============================================================ +# TEST SUITE: _split_svd() — edge cases +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Non-square middle axis raises AssertionError +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_split_svd_nonsquare_middle_axis_raises(): + # This case tests that a middle axis that is not a perfect square + # (d1 != d2) raises ValueError when d1 is not provided. + theta = np.random.RandomState(83).randn(1, 6, 1) + 0j + with pytest.raises(ValueError, match='not a perfect square'): + _split_svd(theta, 0, 0.0) + + +# ------------------------------------------------------------ +# TEST: Non-square middle axis succeeds when d1 is provided +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_split_svd_nonsquare_with_d1(): + # This case tests that _split_svd handles d1 != d2 when d1 is + # explicitly provided. This is required for tensor HOPS where + # adjacent MPS sites have different physical dimensions. + np.random.seed(84) + d1, d2, Dl, Dr = 2, 3, 4, 5 + theta_4d = np.random.randn(Dl, d1, d2, Dr) + 1j * np.random.randn(Dl, d1, d2, Dr) + theta = theta_4d.reshape(Dl, d1 * d2, Dr) + U, S, Vt, chi_new = _split_svd(theta, 0, 0.0, d1=d1) + assert U.shape[0] == Dl + assert U.shape[1] == d1 + assert Vt.shape[1] == d2 + assert Vt.shape[2] == Dr + # Verify recovery: U @ diag(S) @ Vt ≈ theta + recovered = np.einsum('ijk,k,klm->ijlm', U, S, Vt).reshape(Dl, d1 * d2, Dr) + np.testing.assert_allclose(recovered, theta, atol=1e-12, + err_msg='SVD with d1 provided does not recover input') + + +# ------------------------------------------------------------ +# TEST: All SVs below eps still keeps chi_new=1 +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_split_svd_all_svs_below_eps_floor(): + # This case tests that when all singular values fall below eps, + # chi_new is floored to 1 (max(0, 1) = 1) rather than 0. + d = 2 + mat = np.diag([0.001, 0.002]).astype(np.complex128) + theta = mat.reshape(1, d * d, 1) + U, S, Vt, chi_new = _split_svd(theta, 0, 1.0) + assert chi_new == 1, f'Expected chi_new=1 (floor), got {chi_new}' + assert U.shape == (1, d, 1) + assert Vt.shape == (1, d, 1) + # Analytical: when floored to chi=1, the retained SV should be the largest + assert S[0] == pytest.approx(0.002, abs=1e-14), ( + 'Retained SV should be the largest (0.002)' + ) + + +# ============================================================ +# TEST SUITE: timestep() — unknown method +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Unknown method raises ValueError +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_timestep_unknown_method_raises(): + # This case tests that an invalid method string raises ValueError. + phys_dims = [2, 2] + list_cores = _make_list_cores(phys_dims, seed=84) + list_cores_mpo = _make_list_mpo(phys_dims) + core_M, list_cores_B, L0, list_envs_R = initialize(list_cores, list_cores_mpo) + with pytest.raises(ValueError, match='Unknown method'): + timestep( + 0.01, + L0, + list_envs_R, + list_cores_mpo, + core_M, + list_cores_B, + method='3tdvp', + ) + + +# ============================================================ +# TEST SUITE: initialize() — correctness +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: B list_cores are right-orthogonal after initialization +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_initialize_b_cores_right_orthogonal(): + # This case tests that each B core returned by initialize() satisfies + # the right-orthogonality condition: Q @ Q† = I. + phys_dims = [2, 3, 2] + list_cores = _make_list_cores(phys_dims, seed=85) + list_cores_mpo = _make_list_mpo(phys_dims) + core_M, list_cores_B, L0, list_envs_R = initialize(list_cores, list_cores_mpo) + for idx, B in enumerate(list_cores_B): + chi_l, d, Dr = B.shape + Q_mat = B.reshape(chi_l, d * Dr) + np.testing.assert_allclose( + Q_mat @ Q_mat.conj().T, + np.eye(chi_l), + atol=1e-12, + err_msg=f'B core {idx} is not right-orthogonal', + ) + + +# ------------------------------------------------------------ +# TEST: Norm preserved through initialization +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_initialize_norm_preserved(): + # This case tests that the MPS norm is preserved through the + # gauge transformation in initialize(). Contracts all list_cores + # before and after and compares ⟨ψ|ψ⟩. + phys_dims = [2, 3, 2] + list_cores = _make_list_cores(phys_dims, seed=86) + list_cores_mpo = _make_list_mpo(phys_dims) + + # norm before + state_before = list_cores[0] + for c in list_cores[1:]: + state_before = np.tensordot(state_before, c, axes=([-1], [0])) + norm_sq_before = np.vdot(state_before, state_before) + + # norm after + core_M, list_cores_B, L0, list_envs_R = initialize(list_cores, list_cores_mpo) + all_cores = [core_M] + list(list_cores_B) + state_after = all_cores[0] + for c in all_cores[1:]: + state_after = np.tensordot(state_after, c, axes=([-1], [0])) + norm_sq_after = np.vdot(state_after, state_after) + + np.testing.assert_allclose( + norm_sq_after, + norm_sq_before, + atol=1e-10, + err_msg='Norm not preserved through initialization', + ) + + +# ------------------------------------------------------------ +# TEST: Environment tensors consistent with contract_right +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_initialize_envs_consistent_with_contract_right(): + # This case tests that each right environment R_j returned by + # initialize() equals contract_right(R_{j+1}, W_j, B_j), verifying + # the environment sweep is self-consistent. + phys_dims = [2, 3, 2] + list_cores = _make_list_cores(phys_dims, seed=88) + list_cores_mpo = _make_list_mpo(phys_dims) + core_M, list_cores_B, L0, list_envs_R = initialize( + list_cores, list_cores_mpo, + ) + n_sites = len(phys_dims) + # list_envs_R[n_sites-1] = R_{L+1} = boundary (1,1,1) + # list_envs_R[j-1] = R_j for j=2..L+1 + # Verify: list_envs_R[j-1] == contract_right(list_envs_R[j], W_j, B_{j-1}) + for j in range(n_sites - 1, 0, -1): + R_next = list_envs_R[j] + W_j = list_cores_mpo[j] + B_j = list_cores_B[j - 1] + R_expected = contract_right(R_next, W_j, B_j) + np.testing.assert_allclose( + list_envs_R[j - 1], + R_expected, + atol=1e-12, + err_msg=f'Environment R_{j} inconsistent with contract_right', + ) + + +# ============================================================ +# TEST SUITE: L=2 boundary condition +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: 1TDVP timestep with L=2 +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_timestep_1tdvp_l2(): + # This case tests that 1TDVP runs without error on L=2 and + # produces correct output structure. + phys_dims = [2, 2] + list_cores = _make_list_cores(phys_dims, seed=88) + list_cores_mpo = _make_list_mpo(phys_dims) + core_M, list_cores_B, L0, list_envs_R = initialize(list_cores, list_cores_mpo) + norm_before = _mps_norm(list_cores) + core_M, list_cores_B, L0, list_envs_R = timestep( + 0.01, + L0, + list_envs_R, + list_cores_mpo, + core_M, + list_cores_B, + method='1tdvp', + ) + all_cores = [core_M] + list(list_cores_B) + assert len(all_cores) == 2 + assert all_cores[0].shape[0] == 1, 'Left boundary bond dim should be 1' + assert all_cores[-1].shape[2] == 1, 'Right boundary bond dim should be 1' + # Invariant: TDVP preserves the norm of the state + norm_after = _mps_norm(all_cores) + np.testing.assert_allclose( + norm_after, norm_before, atol=1e-10, + err_msg='1TDVP should preserve norm', + ) + + +# ------------------------------------------------------------ +# TEST: 2TDVP timestep with L=2 and R_2 fixup +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_timestep_2tdvp_l2_r2_fixup(): + # This case tests that 2TDVP on L=2 exercises the R_2 fixup branch + # in sweep_left_2tdvp (source:1100-1102). Verifies list_envs_R[0] + # is not None and has the correct shape after the timestep. + phys_dims = [2, 2] + list_cores = _make_list_cores(phys_dims, seed=89) + list_cores_mpo = _make_list_mpo(phys_dims) + core_M, list_cores_B, L0, list_envs_R = initialize(list_cores, list_cores_mpo) + core_M, list_cores_B, L0, list_envs_R = timestep( + 0.01, + L0, + list_envs_R, + list_cores_mpo, + core_M, + list_cores_B, + method='2tdvp', + ) + all_cores = [core_M] + list(list_cores_B) + assert len(all_cores) == 2 + assert all_cores[0].shape[0] == 1, 'Left boundary bond dim should be 1' + assert all_cores[-1].shape[2] == 1, 'Right boundary bond dim should be 1' + # Verify R_2 fixup: list_envs_R[0] must not be None + assert list_envs_R[0] is not None, 'R_2 fixup failed: list_envs_R[0] is None' + assert list_envs_R[0].ndim == 3, 'R_2 should be a rank-3 environment tensor' + + +# ============================================================ +# TEST SUITE: timestep() — alternate solvers +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: 1TDVP with lanczos solver +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_timestep_1tdvp_lanczos_solver(): + # This case tests that the lanczos solver propagates through + # the full timestep pipeline without error. + phys_dims = [2, 3, 2] + list_cores = _make_list_cores(phys_dims, seed=90) + list_cores_mpo = _make_list_mpo(phys_dims) + core_M, list_cores_B, L0, list_envs_R = initialize(list_cores, list_cores_mpo) + norm_before = _mps_norm(list_cores) + core_M, list_cores_B, L0, list_envs_R = timestep( + 0.01, + L0, + list_envs_R, + list_cores_mpo, + core_M, + list_cores_B, + method='1tdvp', + solver='lanczos', + ) + all_cores = [core_M] + list(list_cores_B) + assert len(all_cores) == len(phys_dims) + assert all_cores[0].shape[0] == 1, 'Left boundary bond dim should be 1' + assert all_cores[-1].shape[2] == 1, 'Right boundary bond dim should be 1' + # Invariant: TDVP preserves the norm of the state + norm_after = _mps_norm(all_cores) + np.testing.assert_allclose( + norm_after, norm_before, atol=1e-10, + err_msg='1TDVP should preserve norm', + ) + + +# ------------------------------------------------------------ +# TEST: 1TDVP with ivp solver raises on complex dt +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_timestep_1tdvp_ivp_solver_complex_dt_raises(): + # This case tests that the ivp solver cannot handle the complex + # dt = -1j * delta that the TDVP sweeps pass to _solve_local. + # _ivp_solve compares dt >= 0 which fails for complex numbers. + # This documents a known limitation of the ivp backend. + phys_dims = [2, 3, 2] + list_cores = _make_list_cores(phys_dims, seed=91) + list_cores_mpo = _make_list_mpo(phys_dims) + core_M, list_cores_B, L0, list_envs_R = initialize(list_cores, list_cores_mpo) + with pytest.raises(TypeError): + timestep( + 0.01, + L0, + list_envs_R, + list_cores_mpo, + core_M, + list_cores_B, + method='1tdvp', + solver='ivp', + ) + + +# ============================================================ +# TEST SUITE: 2TDVP chi_max truncation +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: chi_max caps bond dimension +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_timestep_2tdvp_chi_max_caps_bond(): + # This case tests that chi_max=2 limits all output bond dimensions + # to at most 2, even when the input MPS has bond dimension 4. + d = 2 + rng = np.random.RandomState(92) + # Build chi=4 MPS: (1,d,4), (4,d,4), (4,d,1) + list_cores = [ + rng.randn(1, d, 4) + 1j * rng.randn(1, d, 4), + rng.randn(4, d, 4) + 1j * rng.randn(4, d, 4), + rng.randn(4, d, 1) + 1j * rng.randn(4, d, 1), + ] + list_cores_mpo = _make_list_mpo([d, d, d]) + core_M, list_cores_B, L0, list_envs_R = initialize(list_cores, list_cores_mpo) + core_M, list_cores_B, L0, list_envs_R = timestep( + 0.01, + L0, + list_envs_R, + list_cores_mpo, + core_M, + list_cores_B, + method='2tdvp', + chi_max=2, + ) + all_cores = [core_M] + list(list_cores_B) + for idx, c in enumerate(all_cores): + Dl, d_phys, Dr = c.shape + assert Dl <= 2, f'Core {idx}: left bond dim {Dl} exceeds chi_max=2' + assert Dr <= 2, f'Core {idx}: right bond dim {Dr} exceeds chi_max=2' + + +# ============================================================ +# TEST SUITE: Complex dt for Krylov solvers +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: _arnoldi_expm with imaginary dt +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_arnoldi_expm_imaginary_dt(): + # This case tests _arnoldi_expm with dt = -0.05j (the actual TDVP + # use case). For diagonal H = diag(+1, -1): + # exp(-0.05j * H) @ [a, b] = [exp(-0.05j)*a, exp(+0.05j)*b] + # This is a unitary rotation — magnitudes preserved, phases shifted. + d = 2 + rng = np.random.RandomState(93) + state = (rng.randn(1, d, 1) + 1j * rng.randn(1, d, 1)).astype(np.complex128) + env_left, env_right = _trivial_envs() + core_mpo = _diagonal_mpo_core([1.0, -1.0]) + dt = -0.05j + result = _arnoldi_expm( + state, + lambda v: _apply_heff_site(v, env_left, core_mpo, env_right), + dt, + 1e-12, + ) + expected = state.copy() + expected[0, 0, 0] *= np.exp(dt) + expected[0, 1, 0] *= np.exp(-dt) + np.testing.assert_allclose(result, expected, atol=1e-10) + # Verify magnitudes preserved (unitary evolution) + np.testing.assert_allclose( + np.abs(result), + np.abs(state), + atol=1e-10, + err_msg='Imaginary dt should preserve magnitudes', + ) + + +# ------------------------------------------------------------ +# TEST: _lanczos_expm with imaginary dt +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_lanczos_expm_imaginary_dt(): + # This case tests _lanczos_expm with dt = -0.05j for the same + # Hermitian diagonal H = diag(+1, -1). Same analytical result. + d = 2 + rng = np.random.RandomState(94) + state = (rng.randn(1, d, 1) + 1j * rng.randn(1, d, 1)).astype(np.complex128) + env_left, env_right = _trivial_envs() + core_mpo = _diagonal_mpo_core([1.0, -1.0]) + dt = -0.05j + result = _lanczos_expm( + state, + lambda v: _apply_heff_site(v, env_left, core_mpo, env_right), + dt, + 1e-12, + ) + expected = state.copy() + expected[0, 0, 0] *= np.exp(dt) + expected[0, 1, 0] *= np.exp(-dt) + np.testing.assert_allclose(result, expected, atol=1e-10) + # Verify magnitudes preserved (unitary evolution) + np.testing.assert_allclose( + np.abs(result), + np.abs(state), + atol=1e-10, + err_msg='Imaginary dt should preserve magnitudes', + ) + + +# ============================================================ +# TEST SUITE: _apply_heff_site() — non-trivial MPO +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Diagonal MPO applies eigenvalues to input +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_apply_heff_site_diagonal_mpo_applies_eigenvalues(): + # This case tests that _apply_heff_site with a diagonal MPO diag(2, -1) + # and trivial environments gives H|v⟩ = [2*a, -1*b]. + d = 2 + Dl, Dr = 1, 1 + state = np.array([[[3.0 + 1j]], [[0.5 - 2j]]], dtype=np.complex128) + state = state.reshape(Dl, d, Dr) + env_left, env_right = _trivial_envs() + core_mpo = _diagonal_mpo_core([2.0, -1.0]) + result = _apply_heff_site(state, env_left, core_mpo, env_right) + expected = state.copy() + expected[0, 0, 0] *= 2.0 + expected[0, 1, 0] *= -1.0 + np.testing.assert_allclose( + result, expected, atol=1e-12, err_msg='Diagonal MPO should scale each component' + ) + + +# ============================================================ +# TEST SUITE: _apply_heff_bond() — non-trivial environments +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Scaled environments multiply bond tensor +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_apply_heff_bond_scaled_envs(): + # This case tests that _apply_heff_bond with L=[[[3.0]]] and R=[[[2.0]]] + # gives HC = 3 * 2 * C = 6 * C. Verifies both environments + # contribute multiplicatively. + H2_bond = np.array([[1.5 + 0.5j]], dtype=np.complex128) + env_left_scaled = np.array([[[3.0]]], dtype=np.complex128) + env_right_scaled = np.array([[[2.0]]], dtype=np.complex128) + result = _apply_heff_bond(H2_bond, env_left_scaled, env_right_scaled) + np.testing.assert_allclose( + result, + 6.0 * H2_bond, + atol=1e-12, + err_msg='Scaled envs should multiply bond tensor', + ) + + +# ============================================================ +# TEST SUITE: _apply_heff_bond() — multi-bond-dimension +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Bond contraction with chi > 1 and w > 1 +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_apply_heff_bond_multi_bond_dim(): + # This case tests that _apply_heff_bond contracts correctly when + # the bond tensor and environments have dimensions > 1. + # Uses chi_l=2, chi_r=3, w=2 and verifies against direct einsum. + chi_l, chi_r, w = 2, 3, 2 + rng = np.random.RandomState(100) + H2_bond = rng.randn(chi_l, chi_r) + 1j * rng.randn(chi_l, chi_r) + env_left = rng.randn(chi_l, w, chi_l) + 1j * rng.randn(chi_l, w, chi_l) + env_right = rng.randn(chi_r, w, chi_r) + 1j * rng.randn(chi_r, w, chi_r) + # Compute via function + result = _apply_heff_bond(H2_bond, env_left, env_right) + # Compute via direct einsum for reference + expected = np.einsum('ijk,ljm,km->il', env_left, env_right, H2_bond) + np.testing.assert_allclose( + result, + expected, + atol=1e-12, + err_msg='_apply_heff_bond with chi>1 does not match einsum', + ) + + +# ============================================================ +# TEST SUITE: _apply_heff_twosite() — non-trivial MPO +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Diagonal MPO scales two-site tensor by eigenvalues +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_apply_heff_twosite_diagonal_mpo(): + # This case tests _apply_heff_twosite with diagonal MPO diag(2, -1) + # on both sites. With trivial environments, the effective operator + # is diag(2,-1) ⊗ diag(2,-1), so each (s1,s2) component is + # scaled by eigenvalue[s1] * eigenvalue[s2]. + d = 2 + eigenvalues = [2.0, -1.0] + rng = np.random.RandomState(101) + # Build theta as (1, d*d, 1) with known components + theta = rng.randn(1, d * d, 1) + 1j * rng.randn(1, d * d, 1) + core_mpo_site1 = _diagonal_mpo_core(eigenvalues) + core_mpo_site2 = _diagonal_mpo_core(eigenvalues) + env_left = np.ones((1, 1, 1), dtype=np.complex128) + env_right = np.ones((1, 1, 1), dtype=np.complex128) + # Compute result + result = _apply_heff_twosite( + theta, + env_left, + core_mpo_site1, + core_mpo_site2, + env_right, + ) + # Build expected: each (s1, s2) scaled by eigenvalues[s1]*eigenvalues[s2] + theta_4d = theta.reshape(1, d, d, 1) + expected_4d = theta_4d.copy() + for s1 in range(d): + for s2 in range(d): + expected_4d[0, s1, s2, 0] *= eigenvalues[s1] * eigenvalues[s2] + expected = expected_4d.reshape(1, d * d, 1) + np.testing.assert_allclose(result, expected, atol=1e-12) + + +# ------------------------------------------------------------ +# TEST: Two-site contraction with chi > 1 and w > 1 +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_apply_heff_twosite_multi_bond_dim(): + # This case tests _apply_heff_twosite with non-trivial bond and MPO + # dimensions: Dl=2, Dr=3, d=2 (so d*d=4), w_l=2, w_m=2, w_r=2. + # Verifies against direct einsum. + Dl, Dr, d = 2, 3, 2 + w_l, w_m, w_r = 2, 2, 2 + rng = np.random.RandomState(201) + theta = rng.randn(Dl, d * d, Dr) + 1j * rng.randn(Dl, d * d, Dr) + L = rng.randn(Dl, w_l, Dl) + 1j * rng.randn(Dl, w_l, Dl) + W1 = rng.randn(w_l, d, d, w_m) + 1j * rng.randn(w_l, d, d, w_m) + W2 = rng.randn(w_m, d, d, w_r) + 1j * rng.randn(w_m, d, d, w_r) + R = rng.randn(Dr, w_r, Dr) + 1j * rng.randn(Dr, w_r, Dr) + result = _apply_heff_twosite(theta, L, W1, W2, R) + # Direct einsum reference + theta_4d = theta.reshape(Dl, d, d, Dr) + expected_4d = np.einsum( + 'ijk,jlmn,nopq,rqs,kmps->ilor', + L, W1, W2, R, theta_4d, + ) + expected = expected_4d.reshape(Dl, d * d, Dr) + np.testing.assert_allclose( + result, expected, atol=1e-12, + err_msg='_apply_heff_twosite with chi>1/w>1 does not match einsum', + ) + + +# ------------------------------------------------------------ +# TEST: Two-site contraction with non-uniform physical dims (d1 != d2) +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_apply_heff_twosite_nonuniform_phys_dims(): + # This case tests _apply_heff_twosite where the two sites have + # different physical dimensions (d1=2, d2=3), as occurs in tensor + # HOPS with statenumber representation. + Dl, Dr = 2, 3 + d1, d2 = 2, 3 + w_l, w_m, w_r = 2, 2, 2 + rng = np.random.RandomState(202) + theta = rng.randn(Dl, d1 * d2, Dr) + 1j * rng.randn(Dl, d1 * d2, Dr) + L = rng.randn(Dl, w_l, Dl) + 1j * rng.randn(Dl, w_l, Dl) + W1 = rng.randn(w_l, d1, d1, w_m) + 1j * rng.randn(w_l, d1, d1, w_m) + W2 = rng.randn(w_m, d2, d2, w_r) + 1j * rng.randn(w_m, d2, d2, w_r) + R = rng.randn(Dr, w_r, Dr) + 1j * rng.randn(Dr, w_r, Dr) + result = _apply_heff_twosite(theta, L, W1, W2, R) + # Direct einsum reference + theta_4d = theta.reshape(Dl, d1, d2, Dr) + expected_4d = np.einsum( + 'ijk,jlmn,nopq,rqs,kmps->ilor', + L, W1, W2, R, theta_4d, + ) + expected = expected_4d.reshape(Dl, d1 * d2, Dr) + np.testing.assert_allclose( + result, expected, atol=1e-12, + err_msg='_apply_heff_twosite with d1!=d2 does not match einsum', + ) + + +# ============================================================ +# TEST SUITE: _apply_heff_site() — multi-bond-dimension +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Site contraction with chi > 1 and MPO bond w > 1 +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_apply_heff_site_multi_bond_dim(): + # This case tests _apply_heff_site with non-trivial bond dimensions: + # Dl=2, Dr=3, d=2, w_l=2, w_r=2. Verifies against direct einsum. + Dl, Dr, d, w_l, w_r = 2, 3, 2, 2, 2 + rng = np.random.RandomState(102) + core_M = rng.randn(Dl, d, Dr) + 1j * rng.randn(Dl, d, Dr) + env_left = rng.randn(Dl, w_l, Dl) + 1j * rng.randn(Dl, w_l, Dl) + core_mpo = rng.randn(w_l, d, d, w_r) + 1j * rng.randn(w_l, d, d, w_r) + env_right = rng.randn(Dr, w_r, Dr) + 1j * rng.randn(Dr, w_r, Dr) + # Compute via function + result = _apply_heff_site(core_M, env_left, core_mpo, env_right) + # Compute via direct einsum + expected = np.einsum( + 'ijk,jlmn,onp,kmp->ilo', + env_left, + core_mpo, + env_right, + core_M, + ) + np.testing.assert_allclose( + result, + expected, + atol=1e-12, + err_msg='_apply_heff_site with chi>1 does not match einsum', + ) + + +# ============================================================ +# TEST SUITE: _split_svd() — combined truncation +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: chi_max and eps both active simultaneously +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_split_svd_chi_max_and_eps_combined(): + # This case tests that chi_max AND eps both constrain together. + # SVs = [10, 5, 0.01]. With eps=0.5 → keeps [10, 5] (drops 0.01). + # With chi_max=1 → keeps only [10]. Both active → chi_new=1. + d = 3 + H2_diag = np.diag([10.0, 5.0, 0.01]).astype(np.complex128) + theta = H2_diag.reshape(1, d * d, 1) + U, S, Vt, chi_new = _split_svd(theta, chi_max=1, eps=0.5) + # chi_max=1 is the binding constraint + assert chi_new == 1 + np.testing.assert_allclose(S[0], 10.0, atol=1e-12) + + # Now chi_max=2, eps=0.5 → eps drops 0.01, chi_max allows 2 → chi_new=2 + U, S, Vt, chi_new = _split_svd(theta, chi_max=2, eps=0.5) + assert chi_new == 2 + np.testing.assert_allclose(S[0], 10.0, atol=1e-12) + np.testing.assert_allclose(S[1], 5.0, atol=1e-12) + + +# ============================================================ +# TEST SUITE: _ivp_solve() — rank-3 input and max_step +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Rank-3 input reshape round-trip +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_ivp_solve_rank3_input(): + # This case tests that _ivp_solve correctly handles rank-3 tensors + # by flattening to 1D, integrating, and reshaping back. + # dy/dt = y → y(t) = exp(t) * y(0). + d = 3 + state = np.random.RandomState(103).randn(1, d, 1) + 0j + dt = 0.05 + result = _ivp_solve(state, lambda v: v, dt) + assert result.shape == (1, d, 1), 'Output shape must match input' + np.testing.assert_allclose(result, np.exp(dt) * state, atol=1e-6) + + +# ============================================================ +# TEST SUITE: _arnoldi_expm() — negative real dt +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Negative real dt shrinks state +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_arnoldi_expm_negative_real_dt(): + # This case tests exp(-|dt| * I) * state = exp(-|dt|) * state. + # With identity MPO and trivial envs, the operator is I. + d = 2 + rng = np.random.RandomState(104) + state = (rng.randn(1, d, 1) + 1j * rng.randn(1, d, 1)).astype( + np.complex128, + ) + env_left, env_right = _trivial_envs() + core_mpo = _identity_mpo_core(d) + dt = -0.05 + result = _arnoldi_expm( + state, + lambda v: _apply_heff_site(v, env_left, core_mpo, env_right), + dt, + 1e-12, + ) + np.testing.assert_allclose(result, np.exp(dt) * state, atol=1e-10) + + +# ============================================================ +# TEST SUITE: _lanczos_expm() — negative real dt +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Negative real dt shrinks state +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_lanczos_expm_negative_real_dt(): + # This case tests exp(-|dt| * I) * state = exp(-|dt|) * state. + d = 2 + rng = np.random.RandomState(105) + state = (rng.randn(1, d, 1) + 1j * rng.randn(1, d, 1)).astype( + np.complex128, + ) + env_left, env_right = _trivial_envs() + core_mpo = _identity_mpo_core(d) + dt = -0.05 + result = _lanczos_expm( + state, + lambda v: _apply_heff_site(v, env_left, core_mpo, env_right), + dt, + 1e-12, + ) + np.testing.assert_allclose(result, np.exp(dt) * state, atol=1e-10) + + +# ============================================================ +# TEST SUITE: _solve_local() — custom kwargs passthrough +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: conv_tol forwarded through arnoldi dispatch +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_solve_local_custom_conv_tol(): + # This case tests that conv_tol propagates through _solve_local to + # _arnoldi_expm. With identity MPO the result must be exp(dt) * state + # regardless of conv_tol (convergence is trivial for rank-1 operators). + d = 2 + state = np.random.RandomState(106).randn(1, d, 1) + 0j + env_left, env_right = _trivial_envs() + core_mpo = _identity_mpo_core(d) + dt = 0.05 + result = _solve_local( + state, + lambda v: _apply_heff_site(v, env_left, core_mpo, env_right), + dt, + solver='arnoldi', + conv_tol=1e-8, + ) + np.testing.assert_allclose(result, np.exp(dt) * state, atol=1e-8) + + +# ============================================================ +# TEST SUITE: initialize() — boundary and bond-dim cases +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: L=1 single-site chain +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_initialize_l1_single_site(): + # This case tests the degenerate L=1 case. With one site: + # - list_cores_B should be empty + # - list_envs_R should have 1 element (the boundary R_2 = [[1]]) + # - core_M equals the input core (no RQ sweep needed) + phys_dims = [3] + list_cores = _make_list_cores(phys_dims, seed=107) + list_cores_mpo = _make_list_mpo(phys_dims) + core_M, list_cores_B, L0, list_envs_R = initialize( + list_cores, + list_cores_mpo, + ) + assert len(list_cores_B) == 0, 'L=1 should have no B cores' + assert len(list_envs_R) == 1, 'L=1 should have 1 env (boundary)' + # core_M should equal the input (no gauge transform needed) + np.testing.assert_allclose(core_M, list_cores[0], atol=1e-12) + # Boundary environment should be [[1]] + np.testing.assert_allclose( + list_envs_R[0], + np.ones((1, 1, 1)), + atol=1e-12, + ) + + +# ------------------------------------------------------------ +# TEST: Higher bond-dimension MPS +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_initialize_high_bond_dim(): + # This case tests initialize with a 3-site MPS with bond dims > 1: + # (1,d,3), (3,d,2), (2,d,1). Verifies norm preservation and + # B cores right-orthogonal after the RQ gauge sweep. + d = 2 + rng = np.random.RandomState(108) + list_cores = [ + rng.randn(1, d, 3) + 1j * rng.randn(1, d, 3), + rng.randn(3, d, 2) + 1j * rng.randn(3, d, 2), + rng.randn(2, d, 1) + 1j * rng.randn(2, d, 1), + ] + list_cores_mpo = _make_list_mpo([d, d, d]) + # Compute norm before + state_before = list_cores[0] + for c in list_cores[1:]: + state_before = np.tensordot(state_before, c, axes=([-1], [0])) + norm_sq_before = np.vdot(state_before, state_before) + # Initialize + core_M, list_cores_B, L0, list_envs_R = initialize( + list_cores, + list_cores_mpo, + ) + # Verify norm preservation + all_cores = [core_M] + list(list_cores_B) + state_after = all_cores[0] + for c in all_cores[1:]: + state_after = np.tensordot(state_after, c, axes=([-1], [0])) + norm_sq_after = np.vdot(state_after, state_after) + np.testing.assert_allclose( + norm_sq_after, + norm_sq_before, + atol=1e-10, + err_msg='Norm not preserved for high bond-dim MPS', + ) + # Verify B cores are right-orthogonal + for idx, B in enumerate(list_cores_B): + chi_l, d_phys, Dr = B.shape + Q_mat = B.reshape(chi_l, d_phys * Dr) + np.testing.assert_allclose( + Q_mat @ Q_mat.conj().T, + np.eye(chi_l), + atol=1e-12, + err_msg=f'B core {idx} not right-orthogonal', + ) + + +# ============================================================ +# TEST SUITE: timestep() — delta=0, eps, diagonal MPO +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: delta=0 returns state unchanged +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_timestep_delta_zero_unchanged(): + # This case tests that timestep with delta=0 returns the + # state and environments unchanged (exp(0) = I). + phys_dims = [2, 3, 2] + list_cores = _make_list_cores(phys_dims, seed=109) + list_cores_mpo = _make_list_mpo(phys_dims) + core_M, list_cores_B, L0, list_envs_R = initialize( + list_cores, + list_cores_mpo, + ) + # Save copies of the initial state + core_M_before = core_M.copy() + list_B_before = [b.copy() for b in list_cores_B] + # Timestep with delta=0 + core_M_out, list_B_out, L0_out, list_R_out = timestep( + 0.0, + L0, + list_envs_R, + list_cores_mpo, + core_M, + list_cores_B, + method='1tdvp', + ) + # Core M should be unchanged + np.testing.assert_allclose(core_M_out, core_M_before, atol=1e-12) + # B cores should be unchanged + for idx in range(len(list_B_before)): + np.testing.assert_allclose( + list_B_out[idx], + list_B_before[idx], + atol=1e-12, + err_msg=f'B core {idx} changed with delta=0', + ) + + +# ------------------------------------------------------------ +# TEST: 2TDVP with eps > 0 truncates small singular values +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_timestep_2tdvp_eps_truncation(): + # This case tests that eps > 0 in 2TDVP triggers SVD truncation. + # Uses a high bond-dim MPS so there are singular values to drop. + d = 2 + rng = np.random.RandomState(110) + list_cores = [ + rng.randn(1, d, 4) + 1j * rng.randn(1, d, 4), + rng.randn(4, d, 4) + 1j * rng.randn(4, d, 4), + rng.randn(4, d, 1) + 1j * rng.randn(4, d, 1), + ] + list_cores_mpo = _make_list_mpo([d, d, d]) + core_M, list_cores_B, L0, list_envs_R = initialize( + list_cores, + list_cores_mpo, + ) + # Run with large eps to force aggressive truncation + core_M_out, list_B_out, L0_out, list_R_out = timestep( + 0.01, + L0, + list_envs_R, + list_cores_mpo, + core_M, + list_cores_B, + method='2tdvp', + eps=1.0, + ) + # With eps=1.0 many SVs should be dropped, reducing bond dims + all_cores = [core_M_out] + list(list_B_out) + max_bond = max(max(c.shape[0], c.shape[2]) for c in all_cores) + assert max_bond < 4, f'eps=1.0 should truncate bond dim below 4, got {max_bond}' + + +# ------------------------------------------------------------ +# TEST: Diagonal MPO through full 1TDVP pipeline +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_timestep_1tdvp_diagonal_mpo(): + # This case tests the full pipeline with a non-identity Hamiltonian. + # Uses diagonal on-site energies diag(ε₁, ε₂) on each site. + # After time δ, each component picks up exp(-i δ εₖ) per site. + d = 2 + delta = 0.01 + eigenvalues = [1.0, -0.5] + # Build product state: site 1 in state 0, site 2 in state 1 + core_site1 = np.zeros((1, d, 1), dtype=np.complex128) + core_site1[0, 0, 0] = 1.0 # state |0⟩ + core_site2 = np.zeros((1, d, 1), dtype=np.complex128) + core_site2[0, 1, 0] = 1.0 # state |1⟩ + list_cores = [core_site1, core_site2] + # Diagonal MPO: each site has diag(1.0, -0.5) + list_cores_mpo = [_diagonal_mpo_core(eigenvalues) for _ in range(2)] + core_M, list_cores_B, L0, list_envs_R = initialize( + list_cores, + list_cores_mpo, + ) + # Timestep + core_M_out, list_B_out, L0_out, list_R_out = timestep( + delta, + L0, + list_envs_R, + list_cores_mpo, + core_M, + list_cores_B, + method='1tdvp', + ) + # For product state |0⟩⊗|1⟩ under diagonal H: + # E_total = ε[0] + ε[1] = 1.0 + (-0.5) = 0.5 + # TDVP Strang splitting convention: dt = -1j*delta passed to + # _solve_local → net evolution is exp(+i delta E_total) + expected_phase = np.exp(+1j * delta * (eigenvalues[0] + eigenvalues[1])) + # Contract output MPS to get scalar amplitude + all_cores_out = [core_M_out] + list(list_B_out) + psi_out = all_cores_out[0] + for c in all_cores_out[1:]: + psi_out = np.tensordot(psi_out, c, axes=([-1], [0])) + # The (0,1) component should carry the phase + amplitude = psi_out[0, 0, 1, 0] # |0⟩⊗|1⟩ component + np.testing.assert_allclose( + amplitude, + expected_phase, + atol=1e-6, + err_msg='Diagonal MPO did not produce correct phase', + ) + + +# ============================================================ +# TEST SUITE: timestep() — L=4 multi-iteration 2TDVP +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: 2TDVP with L=4 exercises multi-iteration sweep body +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_timestep_2tdvp_l4_multi_iteration(): + # This case tests that 2TDVP with L=4 correctly exercises + # the multi-iteration backward evolution inside both sweeps. + # Verifies output structure and norm preservation. + d = 2 + phys_dims = [d, d, d, d] + list_cores = _make_list_cores(phys_dims, seed=111) + list_cores_mpo = _make_list_mpo(phys_dims) + core_M, list_cores_B, L0, list_envs_R = initialize( + list_cores, + list_cores_mpo, + ) + # Compute norm before + all_before = [core_M] + list(list_cores_B) + psi_before = all_before[0] + for c in all_before[1:]: + psi_before = np.tensordot(psi_before, c, axes=([-1], [0])) + norm_sq_before = np.vdot(psi_before, psi_before) + # Timestep + core_M_out, list_B_out, L0_out, list_R_out = timestep( + 0.01, + L0, + list_envs_R, + list_cores_mpo, + core_M, + list_cores_B, + method='2tdvp', + ) + # Verify structure: 4 cores, boundary bond dims = 1 + all_out = [core_M_out] + list(list_B_out) + assert len(all_out) == 4 + assert all_out[0].shape[0] == 1, 'Left boundary should be 1' + assert all_out[-1].shape[2] == 1, 'Right boundary should be 1' + # Verify norm preservation (identity MPO → unitary evolution) + psi_after = all_out[0] + for c in all_out[1:]: + psi_after = np.tensordot(psi_after, c, axes=([-1], [0])) + norm_sq_after = np.vdot(psi_after, psi_after) + np.testing.assert_allclose( + norm_sq_after, + norm_sq_before, + rtol=1e-6, + err_msg='Norm not preserved for L=4 2TDVP', + ) + + +# ============================================================ +# TEST SUITE: timestep() — 2TDVP diagonal MPO physics +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: 2TDVP diagonal MPO produces correct phase and preserves norm +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_timestep_2tdvp_diagonal_mpo_norm(): + # Invariant: norm should be preserved under diagonal MPO evolution + phys_dims = [2, 2] + list_cores = _make_list_cores(phys_dims, seed=200) + list_cores_mpo = [] + for d in phys_dims: + mpo = np.zeros((1, d, d, 1), dtype=np.complex128) + for k in range(d): + mpo[0, k, k, 0] = float(k + 1) + list_cores_mpo.append(mpo) + core_M, list_cores_B, L0, list_envs_R = initialize( + list_cores, list_cores_mpo, + ) + norm_before = _mps_norm([core_M] + list(list_cores_B)) + dt = 0.005 + core_M2, list_cores_B2, L02, list_envs_R2 = timestep( + dt, L0, list_envs_R, list_cores_mpo, core_M, list_cores_B, + method='2tdvp', chi_max=4, eps=1e-12, + ) + # Invariant: norm preserved + norm_after = _mps_norm([core_M2] + list(list_cores_B2)) + np.testing.assert_allclose(norm_after, norm_before, atol=1e-8) + + +# ============================================================ +# TEST SUITE: timestep() — multi-step norm preservation +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Two consecutive 1TDVP timesteps preserve norm +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_timestep_1tdvp_multi_step_norm(): + # Invariant: norm preserved across multiple steps + phys_dims = [2, 3, 2] + list_cores = _make_list_cores(phys_dims, seed=210) + list_cores_mpo = _make_list_mpo(phys_dims) + core_M, list_cores_B, L0, list_envs_R = initialize( + list_cores, list_cores_mpo, + ) + norm_init = _mps_norm([core_M] + list(list_cores_B)) + for _ in range(2): + core_M, list_cores_B, L0, list_envs_R = timestep( + 0.005, L0, list_envs_R, list_cores_mpo, + core_M, list_cores_B, method='1tdvp', + ) + norm_final = _mps_norm([core_M] + list(list_cores_B)) + np.testing.assert_allclose(norm_final, norm_init, atol=1e-8) + + +# ============================================================ +# TEST SUITE: sweep_right_1tdvp() — direct tests +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: A cores are left-orthogonal after right sweep +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_sweep_right_1tdvp_a_cores_left_orthogonal(): + # This case tests that A cores from sweep_right_1tdvp satisfy + # the left-orthogonality condition: Q†Q = I. + phys_dims = [2, 3, 2] + list_cores = _make_list_cores(phys_dims, seed=112) + list_cores_mpo = _make_list_mpo(phys_dims) + core_M, list_cores_B, L0, list_envs_R = initialize( + list_cores, + list_cores_mpo, + ) + list_envs_L, list_cores_A, core_M_last = sweep_right_1tdvp( + 0.01, + L0, + list_envs_R, + list_cores_mpo, + core_M, + list_cores_B, + ) + # Verify each A core is left-orthogonal + for idx, A in enumerate(list_cores_A): + Dl, d, Dr = A.shape + Q_mat = A.reshape(Dl * d, Dr) + np.testing.assert_allclose( + Q_mat.conj().T @ Q_mat, + np.eye(Dr), + atol=1e-10, + err_msg=f'A core {idx} not left-orthogonal after right sweep', + ) + + +# ------------------------------------------------------------ +# TEST: Left environments consistent with A cores +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_sweep_right_1tdvp_envs_consistent(): + # This case tests that list_envs_L[j+1] = contract_left( + # list_envs_L[j], W_j, A_j) for each j. + phys_dims = [2, 3, 2] + list_cores = _make_list_cores(phys_dims, seed=113) + list_cores_mpo = _make_list_mpo(phys_dims) + core_M, list_cores_B, L0, list_envs_R = initialize( + list_cores, + list_cores_mpo, + ) + list_envs_L, list_cores_A, core_M_last = sweep_right_1tdvp( + 0.01, + L0, + list_envs_R, + list_cores_mpo, + core_M, + list_cores_B, + ) + # Check consistency: L_{j+1} = contract_left(L_j, W_j, A_j) + for j in range(len(list_cores_A)): + L_rebuilt = contract_left( + list_envs_L[j], + list_cores_mpo[j], + list_cores_A[j], + ) + np.testing.assert_allclose( + list_envs_L[j + 1], + L_rebuilt, + atol=1e-10, + err_msg=f'Left env at j={j + 1} inconsistent with A core', + ) + + +# ------------------------------------------------------------ +# TEST: Norm preserved through right sweep +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_sweep_right_1tdvp_norm_preserved(): + # This case tests that the MPS norm is preserved through the + # right sweep. Contracts A cores + core_M_last and compares. + phys_dims = [2, 3, 2] + list_cores = _make_list_cores(phys_dims, seed=114) + list_cores_mpo = _make_list_mpo(phys_dims) + core_M, list_cores_B, L0, list_envs_R = initialize( + list_cores, + list_cores_mpo, + ) + # Norm before + all_before = [core_M] + list(list_cores_B) + psi_before = all_before[0] + for c in all_before[1:]: + psi_before = np.tensordot(psi_before, c, axes=([-1], [0])) + norm_sq_before = np.vdot(psi_before, psi_before) + # Right sweep + list_envs_L, list_cores_A, core_M_last = sweep_right_1tdvp( + 0.01, + L0, + list_envs_R, + list_cores_mpo, + core_M, + list_cores_B, + ) + # Norm after: A cores + core_M_last + all_after = list(list_cores_A) + [core_M_last] + psi_after = all_after[0] + for c in all_after[1:]: + psi_after = np.tensordot(psi_after, c, axes=([-1], [0])) + norm_sq_after = np.vdot(psi_after, psi_after) + np.testing.assert_allclose( + norm_sq_after, + norm_sq_before, + rtol=1e-6, + err_msg='Norm not preserved through right sweep', + ) + + +# ============================================================ +# TEST SUITE: sweep_left_1tdvp() — direct tests +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: B cores right-orthogonal and norm preserved +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_sweep_left_1tdvp_b_cores_and_norm(): + # This case tests that sweep_left_1tdvp produces right-orthogonal + # B cores and preserves the MPS norm. + phys_dims = [2, 3, 2] + list_cores = _make_list_cores(phys_dims, seed=115) + list_cores_mpo = _make_list_mpo(phys_dims) + core_M, list_cores_B, L0, list_envs_R = initialize( + list_cores, + list_cores_mpo, + ) + # Norm before + all_before = [core_M] + list(list_cores_B) + psi_before = all_before[0] + for c in all_before[1:]: + psi_before = np.tensordot(psi_before, c, axes=([-1], [0])) + norm_sq_before = np.vdot(psi_before, psi_before) + # Right sweep first (needed as input) + list_envs_L, list_cores_A, core_M_last = sweep_right_1tdvp( + 0.005, + L0, + list_envs_R, + list_cores_mpo, + core_M, + list_cores_B, + ) + R_boundary = list_envs_R[-1] + # Left sweep + core_M_first, list_B_out, list_R_out = sweep_left_1tdvp( + 0.005, + list_envs_L, + R_boundary, + list_cores_mpo, + list_cores_A, + core_M_last, + ) + # Verify B cores are right-orthogonal + for idx, B in enumerate(list_B_out): + chi_l, d, Dr = B.shape + Q_mat = B.reshape(chi_l, d * Dr) + np.testing.assert_allclose( + Q_mat @ Q_mat.conj().T, + np.eye(chi_l), + atol=1e-10, + err_msg=f'B core {idx} not right-orthogonal after left sweep', + ) + # Verify norm preserved through full right+left sweep + all_after = [core_M_first] + list(list_B_out) + psi_after = all_after[0] + for c in all_after[1:]: + psi_after = np.tensordot(psi_after, c, axes=([-1], [0])) + norm_sq_after = np.vdot(psi_after, psi_after) + np.testing.assert_allclose( + norm_sq_after, + norm_sq_before, + rtol=1e-6, + err_msg='Norm not preserved through right+left sweep', + ) + + +# ------------------------------------------------------------ +# TEST: sweep_left_1tdvp with lanczos solver runs without error +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_sweep_left_1tdvp_lanczos_solver(): + # This case tests that sweep_left_1tdvp works with the lanczos + # solver (non-default). Identity MPO makes the problem Hermitian, + # which is the Lanczos requirement. + phys_dims = [2, 3, 2] + list_cores = _make_list_cores(phys_dims, seed=116) + list_cores_mpo = _make_list_mpo(phys_dims) + core_M, list_cores_B, L0, list_envs_R = initialize( + list_cores, list_cores_mpo, + ) + list_envs_L, list_cores_A, core_M_last = sweep_right_1tdvp( + 0.005, L0, list_envs_R, list_cores_mpo, core_M, list_cores_B, + solver='lanczos', + ) + R_boundary = list_envs_R[-1] + core_M_first, list_B_out, list_R_out = sweep_left_1tdvp( + 0.005, list_envs_L, R_boundary, list_cores_mpo, + list_cores_A, core_M_last, solver='lanczos', + ) + # Basic sanity: output shapes match input + assert core_M_first.shape[1] == phys_dims[0] + assert len(list_B_out) == len(phys_dims) - 1 + # Invariant: B cores from left sweep should be right-orthogonal + for i, B in enumerate(list_B_out): + Dl, d, Dr = B.shape + M = B.reshape(Dl, d * Dr) + eye_check = M @ M.conj().T + np.testing.assert_allclose( + eye_check, np.eye(Dl), atol=1e-10, + err_msg=f'B core {i} not right-orthogonal after Lanczos left sweep', + ) + + +# ============================================================ +# TEST SUITE: sweep_right_2tdvp() / sweep_left_2tdvp() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: 2TDVP right sweep produces left-orthogonal A cores +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_sweep_right_2tdvp_a_cores_left_orthogonal(): + # This case tests that A cores from sweep_right_2tdvp satisfy + # the left-orthogonality condition. + d = 2 + phys_dims = [d, d, d] + list_cores = _make_list_cores(phys_dims, seed=116) + list_cores_mpo = _make_list_mpo(phys_dims) + core_M, list_cores_B, L0, list_envs_R = initialize( + list_cores, + list_cores_mpo, + ) + list_envs_L, list_cores_A, core_M_last = sweep_right_2tdvp( + 0.01, + L0, + list_envs_R, + list_cores_mpo, + core_M, + list_cores_B, + chi_max=0, + eps=0.0, + ) + for idx, A in enumerate(list_cores_A): + Dl, d_phys, Dr = A.shape + Q_mat = A.reshape(Dl * d_phys, Dr) + np.testing.assert_allclose( + Q_mat.conj().T @ Q_mat, + np.eye(Dr), + atol=1e-10, + err_msg=f'A core {idx} not left-orthogonal (2TDVP)', + ) + + +# ------------------------------------------------------------ +# TEST: 2TDVP full sweep preserves norm +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_sweep_2tdvp_full_norm_preserved(): + # This case tests norm preservation through a full 2TDVP + # right+left sweep pair. + d = 2 + phys_dims = [d, d, d] + list_cores = _make_list_cores(phys_dims, seed=117) + list_cores_mpo = _make_list_mpo(phys_dims) + core_M, list_cores_B, L0, list_envs_R = initialize( + list_cores, + list_cores_mpo, + ) + # Norm before + all_before = [core_M] + list(list_cores_B) + psi_before = all_before[0] + for c in all_before[1:]: + psi_before = np.tensordot(psi_before, c, axes=([-1], [0])) + norm_sq_before = np.vdot(psi_before, psi_before) + # Right sweep + list_envs_L, list_cores_A, core_M_last = sweep_right_2tdvp( + 0.005, + L0, + list_envs_R, + list_cores_mpo, + core_M, + list_cores_B, + chi_max=0, + eps=0.0, + ) + R_boundary = list_envs_R[-1] + # Left sweep + core_M_first, list_B_out, list_R_out = sweep_left_2tdvp( + 0.005, + list_envs_L, + R_boundary, + list_cores_mpo, + list_cores_A, + core_M_last, + chi_max=0, + eps=0.0, + ) + # Norm after + all_after = [core_M_first] + list(list_B_out) + psi_after = all_after[0] + for c in all_after[1:]: + psi_after = np.tensordot(psi_after, c, axes=([-1], [0])) + norm_sq_after = np.vdot(psi_after, psi_after) + np.testing.assert_allclose( + norm_sq_after, + norm_sq_before, + rtol=1e-6, + err_msg='Norm not preserved through 2TDVP right+left sweep', + ) + + +# ------------------------------------------------------------ +# TEST: 2TDVP left sweep produces right-orthogonal B cores and preserves norm +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_sweep_left_2tdvp_b_cores_and_norm(): + # This case tests that sweep_left_2tdvp produces right-orthogonal + # B cores and preserves the MPS norm through a full right+left sweep. + # 2TDVP requires uniform physical dimensions (d1==d2 for two-site merge). + phys_dims = [2, 2, 2] + list_cores = _make_list_cores(phys_dims, seed=120) + list_cores_mpo = _make_list_mpo(phys_dims) + core_M, list_cores_B, L0, list_envs_R = initialize( + list_cores, + list_cores_mpo, + ) + # Norm before + all_before = [core_M] + list(list_cores_B) + psi_before = all_before[0] + for c in all_before[1:]: + psi_before = np.tensordot(psi_before, c, axes=([-1], [0])) + norm_sq_before = np.vdot(psi_before, psi_before) + # Right sweep first (needed as input to left sweep) + list_envs_L, list_cores_A, core_M_last = sweep_right_2tdvp( + 0.005, + L0, + list_envs_R, + list_cores_mpo, + core_M, + list_cores_B, + chi_max=0, + eps=0.0, + ) + R_boundary = list_envs_R[-1] + # Left sweep + core_M_first, list_B_out, list_R_out = sweep_left_2tdvp( + 0.005, + list_envs_L, + R_boundary, + list_cores_mpo, + list_cores_A, + core_M_last, + chi_max=0, + eps=0.0, + ) + # Verify B cores are right-orthogonal + for idx, B in enumerate(list_B_out): + chi_l, d, Dr = B.shape + Q_mat = B.reshape(chi_l, d * Dr) + np.testing.assert_allclose( + Q_mat @ Q_mat.conj().T, + np.eye(chi_l), + atol=1e-10, + err_msg=f'B core {idx} not right-orthogonal after left sweep', + ) + # Verify norm preserved through full right+left sweep + all_after = [core_M_first] + list(list_B_out) + psi_after = all_after[0] + for c in all_after[1:]: + psi_after = np.tensordot(psi_after, c, axes=([-1], [0])) + norm_sq_after = np.vdot(psi_after, psi_after) + np.testing.assert_allclose( + norm_sq_after, + norm_sq_before, + rtol=1e-6, + err_msg='Norm not preserved through 2TDVP right+left sweep', + ) + + +# ============================================================ +# Helpers for recenter_to_zero tests +# ============================================================ + + +def _make_simple_mps(n_cores, phys_dims, bond_dims): + """Build a simple MPS with specified structure. + + Parameters + ---------- + 1. n_cores: int + Number of MPS cores (sites). + 2. phys_dims: list(int) + Physical dimension of each core. + 3. bond_dims: list(int) + Bond dimensions, length n_cores + 1. bond_dims[0] and + bond_dims[-1] should be 1 for open boundary conditions. + + Returns + ------- + 1. cores: list(np.ndarray) + MPS cores, each shaped (bond_left, phys_dim, bond_right). + """ + cores = [] + for i in range(n_cores): + core = np.zeros( + (bond_dims[i], phys_dims[i], bond_dims[i + 1]), + dtype=np.complex128, + ) + core += 0.01 * ( + np.random.randn(*core.shape) + 1j * np.random.randn(*core.shape) + ) + cores.append(core) + return cores + + +def _contract_mps(c_list): + """Full contraction of an MPS into a state vector.""" + s = c_list[0] + for c in c_list[1:]: + s = np.tensordot(s, c, axes=([-1], [0])) + return s.squeeze().ravel() + + +# ============================================================ +# TEST SUITE: recenter_to_zero() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Sites 1..L-1 become left-canonical +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_recenter_to_zero_left_canonical(): + # This case tests that after recentering, all sites n >= 1 satisfy + # M_n^dagger M_n = I (left-canonical form). + np.random.seed(42) + cores = _make_simple_mps(4, [2, 3, 3, 2], [1, 3, 4, 3, 1]) + result = recenter_to_zero(cores) + for n in range(1, len(result)): + Dl, d, Dr = result[n].shape + M = result[n].reshape(Dl * d, Dr) + eye_check = M.conj().T @ M + np.testing.assert_allclose( + eye_check, + np.eye(Dr, dtype=np.complex128), + atol=1e-12, + err_msg=f'Site {n} is not left-canonical', + ) + + +# ------------------------------------------------------------ +# TEST: Recentering preserves the represented state +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_recenter_to_zero_preserves_state(): + # This case tests that the full contraction of the MPS gives the same + # state vector before and after recentering. + np.random.seed(123) + cores = _make_simple_mps(3, [2, 2, 2], [1, 2, 2, 1]) + state_before = _contract_mps([c.copy() for c in cores]) + result = recenter_to_zero(cores) + state_after = _contract_mps(result) + np.testing.assert_allclose(state_after, state_before, atol=1e-12) + + +# ------------------------------------------------------------ +# TEST: ValueError when last core Dr != 1 +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_recenter_to_zero_bad_boundary(): + # This case tests that a ValueError is raised when the last core + # has Dr != 1 (violating open boundary conditions). + bad_core = np.zeros((1, 2, 2), dtype=np.complex128) + with pytest.raises(ValueError, match='open boundaries'): + recenter_to_zero([bad_core]) + + +# ------------------------------------------------------------ +# TEST: recenter_to_zero raises on internal bond mismatch +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_recenter_to_zero_bond_mismatch(): + # This case tests that a ValueError is raised when adjacent cores + # have incompatible bond dimensions. The upfront validation catches + # mismatches for any number of cores, including 2-core MPS. + + # 2-core case: right bond of core_0 (3) != left bond of core_1 (2) + core_0 = np.ones((1, 2, 3), dtype=np.complex128) + core_1 = np.ones((2, 2, 1), dtype=np.complex128) + with pytest.raises(ValueError, match='Bond mismatch'): + recenter_to_zero([core_0, core_1]) + + # 3-core case: right bond of core_1 (4) != left bond of core_2 (2) + core_0 = np.ones((1, 2, 3), dtype=np.complex128) + core_1 = np.ones((3, 2, 4), dtype=np.complex128) + core_2 = np.ones((2, 2, 1), dtype=np.complex128) + with pytest.raises(ValueError, match='Bond mismatch'): + recenter_to_zero([core_0, core_1, core_2]) + + +# ------------------------------------------------------------ +# TEST: normalize=True makes ||A[0]||_F = 1 +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_recenter_to_zero_normalize(): + # This case tests that normalize=True rescales the MPS so that + # the Frobenius norm of A[0] equals 1 and the state direction is preserved. + np.random.seed(7) + cores = _make_simple_mps(3, [2, 3, 2], [1, 3, 3, 1]) + state_orig = _contract_mps(cores) + result = recenter_to_zero(cores, normalize=True) + norm_A0 = np.linalg.norm(result[0].ravel()) + np.testing.assert_allclose(norm_A0, 1.0, atol=1e-12) + # The contracted state should be proportional to the original + state_norm = _contract_mps(result) + ratio = state_orig / state_norm + np.testing.assert_allclose( + ratio, ratio[0] * np.ones_like(ratio), atol=1e-12, + err_msg='Normalized state is not proportional to original', + ) + + +# ------------------------------------------------------------ +# TEST: recenter_to_zero copy=False reuses the input list +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_recenter_to_zero_no_copy(): + # This case tests that copy=False returns the same list object and + # that the cores are actually left-canonical after the call. + np.random.seed(55) + cores = _make_simple_mps(3, [2, 3, 2], [1, 4, 4, 1]) + result = recenter_to_zero(cores, copy=False) + assert result is cores + # Verify cores 1..L-1 are left-canonical (Q^dQ = I) + for n in range(1, len(result)): + Dl, d, Dr = result[n].shape + Q = result[n].reshape(Dl * d, Dr) + np.testing.assert_allclose( + Q.conj().T @ Q, np.eye(Dr), atol=1e-12, + err_msg=f'Core {n} is not left-canonical after copy=False', + ) + + # copy=True should return an independent list + np.random.seed(56) + cores2 = _make_simple_mps(3, [2, 3, 2], [1, 4, 4, 1]) + result2 = recenter_to_zero(cores2, copy=True) + assert result2 is not cores2 + + +# ------------------------------------------------------------ +# TEST: recenter_to_zero with single core preserves state +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_recenter_to_zero_single_core(): + # Limiting case: single core MPS, sweep is empty, state preserved + core = np.array([[[1.0 + 0j], [0.5 + 0.1j]]]) # shape (1, 2, 1) + result = recenter_to_zero([core]) + np.testing.assert_allclose(result[0], core, atol=1e-14) + + +# ------------------------------------------------------------ +# TEST: recenter_to_zero with two cores (R goes directly to scalar) +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_recenter_to_zero_two_cores(): + # This case tests the two-core MPS where the sweep runs exactly once + # and R folds directly into site 0 (no absorb-into-next-core branch). + np.random.seed(42) + cores = _make_simple_mps(2, [3, 2], [1, 4, 1]) + state_before = _contract_mps(cores) + result = recenter_to_zero(cores) + state_after = _contract_mps(result) + np.testing.assert_allclose(state_after, state_before, atol=1e-12) + + +# ------------------------------------------------------------ +# TEST: recenter_to_zero on already-canonical MPS is near-no-op +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_recenter_to_zero_already_canonical(): + # This case tests that passing an already-canonical MPS returns the + # same state (within numerical precision). + np.random.seed(99) + cores = _make_simple_mps(3, [2, 3, 2], [1, 3, 3, 1]) + canonical = recenter_to_zero(cores) + state_first = _contract_mps(canonical) + result = recenter_to_zero(canonical) + state_second = _contract_mps(result) + np.testing.assert_allclose(state_second, state_first, atol=1e-12) + + +# ------------------------------------------------------------ +# TEST: recenter_to_zero with all bond dims = 1 +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_recenter_to_zero_trivial_bonds(): + # This case tests the degenerate geometry where every bond is 1, + # so QR reduces to scalar extraction at every site. + np.random.seed(11) + cores = _make_simple_mps(4, [2, 3, 2, 2], [1, 1, 1, 1, 1]) + state_before = _contract_mps(cores) + result = recenter_to_zero(cores) + state_after = _contract_mps(result) + np.testing.assert_allclose(state_after, state_before, atol=1e-12) + + +# ------------------------------------------------------------ +# TEST: recenter_to_zero normalize with zero-norm MPS +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_recenter_to_zero_zero_norm(): + # This case tests the norm_phi > 0 guard: a zero MPS should be + # returned without division by zero. The QR sweep may introduce + # non-zero Q factors, but A[0] should remain zero (R[0,0] = 0 + # is folded into it), so the overall state is still zero. + cores = [ + np.zeros((1, 2, 3), dtype=np.complex128), + np.zeros((3, 3, 1), dtype=np.complex128), + ] + result = recenter_to_zero(cores, normalize=True) + # The orthogonality center (site 0) should be zero + np.testing.assert_allclose(result[0], 0.0, atol=1e-15) + + +# ------------------------------------------------------------ +# TEST: empty list returns empty +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_recenter_to_zero_empty_list(): + # This case tests that an empty list input returns an empty list + # without error. + result = recenter_to_zero([]) + assert result == [] + + +# ============================================================ +# TEST SUITE: Krylov edge cases — breakdown and fallback +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Arnoldi explicit breakdown guard (conv_tol=0) +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_arnoldi_expm_explicit_breakdown_guard(): + # This case tests the explicit breakdown guard (lines 412-416) + # by setting conv_tol=0 so the convergence check (h * |phi| < 0) + # can never fire. The rank-1 projector P = |e_0> _KRYLOV_DIM_MAX=60. + from scipy.linalg import expm as scipy_expm + rng = np.random.default_rng(42) + n = 70 + H = rng.standard_normal((n, n)) + 1j * rng.standard_normal((n, n)) + H = (H + H.conj().T) / 2 # symmetrize for Lanczos + v = rng.standard_normal(n).astype(complex) + matvec = lambda x: H @ x + dt = 0.01 # small dt for accuracy with truncated Krylov + result = _lanczos_expm(v, matvec, dt=dt, conv_tol=0.0) + expected = scipy_expm(dt * H) @ v + np.testing.assert_allclose(result, expected, rtol=1e-4) diff --git a/tests/test_tensor_basis_shared_refs.py b/tests/test_tensor_basis_shared_refs.py new file mode 100644 index 0000000..9248673 --- /dev/null +++ b/tests/test_tensor_basis_shared_refs.py @@ -0,0 +1,135 @@ +import numpy as np +import pytest +import scipy as sp + +from mesohops.trajectory.exp_noise import bcf_exp +from mesohops.trajectory.hops_tensor_trajectory import HopsTensorTrajectory +from mesohops.util.bath_corr_functions import bcf_convert_dl_to_exp + +__title__ = 'Integration Tests for Tensor Basis Shared References' +__author__ = 'A. Hartzell' +__maintainer__ = 'A. Hartzell' + +# ============================================================ +# Shared Setup +# ============================================================ + +nsite = 4 +e_lambda = 20.0 +gamma = 50.0 +temp = 140.0 +(g_0, w_0) = bcf_convert_dl_to_exp(e_lambda, gamma, temp) + +T3_loperator = np.zeros([4, 4, 4], dtype=np.float64) +list_gw_sysbath = [] +list_lop = [] +for i in range(nsite): + T3_loperator[i, i, i] = 1.0 + list_gw_sysbath.append([g_0, w_0]) + list_lop.append(sp.sparse.coo_matrix(T3_loperator[i])) + list_gw_sysbath.append([-1j * np.imag(g_0), 500.0]) + list_lop.append(T3_loperator[i]) + +sys_param = { + 'HAMILTONIAN': np.array(np.zeros([nsite, nsite]), dtype=np.complex128), + 'GW_SYSBATH': list_gw_sysbath, + 'L_HIER': list_lop, + 'L_NOISE1': list_lop, + 'ALPHA_NOISE1': bcf_exp, + 'PARAM_NOISE1': list_gw_sysbath, +} +sys_param['HAMILTONIAN'][0, 1] = 40 +sys_param['HAMILTONIAN'][1, 0] = 40 +sys_param['HAMILTONIAN'][1, 2] = 10 +sys_param['HAMILTONIAN'][2, 1] = 10 +sys_param['HAMILTONIAN'][2, 3] = 40 +sys_param['HAMILTONIAN'][3, 2] = 40 + +noise_param = { + 'SEED': 0, + 'MODEL': 'FFT_FILTER', + 'TLEN': 25000.0, + 'TAU': 1.0, +} + +hier_param = {'MAXHIER': 2} +eom_param = {'EQUATION_OF_MOTION': 'NORMALIZED NONLINEAR'} +integrator_param = {'INTEGRATOR': 'RUNGE_KUTTA'} +tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, +} + +psi_0 = np.array([0.0] * nsite, dtype=np.complex128) +psi_0[2] = 1.0 +psi_0 = psi_0 / np.linalg.norm(psi_0) + + +def _make_traj(): + return HopsTensorTrajectory( + system_param=sys_param, + noise_param=noise_param, + hierarchy_param=hier_param, + eom_param=eom_param, + integration_param=integrator_param, + tensor_param=tensor_param, + ) + + +# ============================================================ +# TEST SUITE: HopsTensorTrajectory() — shared reference identity +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: tensor_basis.system is basis.system +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_system_is_shared(): + '''tensor_basis.system must be the same object as basis.system.''' + # This case tests that tensor_basis holds a reference to the same + # HopsSystem instance as the HOPS basis, not a copy. + traj = _make_traj() + assert traj.tensor_basis.system is traj.basis.system + + +# ------------------------------------------------------------ +# TEST: tensor_basis.mode is basis.mode +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_mode_is_shared(): + '''tensor_basis.mode must be the same object as basis.mode.''' + # This case tests that tensor_basis holds a reference to the same + # HopsModes instance as the HOPS basis, not a copy. + traj = _make_traj() + assert traj.tensor_basis.mode is traj.basis.mode + + +# ------------------------------------------------------------ +# TEST: tensor_basis.noise_memory is basis.noise_memory +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_noise_memory_is_shared(): + '''tensor_basis.noise_memory must be the same object as basis.noise_memory.''' + # This case tests that tensor_basis holds a reference to the same + # HopsNoiseMemory instance as the HOPS basis, not a copy. + traj = _make_traj() + assert traj.tensor_basis.noise_memory is traj.basis.noise_memory + + +# ------------------------------------------------------------ +# TEST: State list mutation through tensor_basis is visible through basis +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_state_list_mutation_visibility(): + # This case tests that mutating state_list through the tensor_basis path + # is immediately visible through the basis path, confirming the two paths + # reference the same underlying HopsSystem object. + traj = _make_traj() + # Mutate state_list through tensor_basis path + traj.tensor_basis.system.state_list = [0, 1] + # Invariant: visible through basis path (same object) + np.testing.assert_array_equal(traj.basis.system.state_list, [0, 1]) + + diff --git a/tests/test_tensor_eom_functions.py b/tests/test_tensor_eom_functions.py new file mode 100644 index 0000000..35ea880 --- /dev/null +++ b/tests/test_tensor_eom_functions.py @@ -0,0 +1,1849 @@ +import numpy as np +import pytest +import scipy as sp + +from mesohops.basis.hops_modes import HopsModes +from mesohops.basis.hops_noise_memory import HopsNoiseMemory +from mesohops.basis.hops_system import HopsSystem +from mesohops.eom.eom_functions import ( + calc_norm_corr, + compress_zmem, + operator_expectation, +) +from mesohops.tensor.hops_tensor_wavefunction import HopsTensorWavefunction +from mesohops.tensor.hops_tensor_basis import HopsTensorBasis +from mesohops.tensor.mpo_constructors import build_statenumber_operator_mpo +from mesohops.tensor.tdvp import recenter_to_zero +from mesohops.tensor.tensor_eom_functions import ( + apply_system_operator, + calc_norm_corr_tensor, + tensor_matvec_prod, +) +from mesohops.trajectory.exp_noise import bcf_exp +from mesohops.util.bath_corr_functions import bcf_convert_dl_to_exp +from mesohops.util.tensor_operations import ( + calc_mps_complexity, + extract_psi, + phi_aux, + unflatten_cores, +) + +__title__ = 'Unit Tests for Tensor EOM Functions' +__author__ = 'N. Covalsen' +__maintainer__ = 'N. Covalsen' + + +def _build_mode_l2_map(system, mode): + """Build the (mode_idx, l2_idx) mapping for calc_norm_corr_tensor tests.""" + dict_state_to_idx = { + s: i for i, s in enumerate(sorted(system.state_list)) + } + list_mode_l2_map = [] + for mode_idx in range(len(mode.list_modeidx_abs)): + abs_hmode_idx = mode.list_modeidx_abs[mode_idx] + sys_state = system.param['LIST_STATE_INDICES_BY_HMODE'][abs_hmode_idx] + sys_state_key = int(np.asarray(sys_state).ravel()[0]) + if sys_state_key not in dict_state_to_idx: + continue + ordered_state_idx = dict_state_to_idx[sys_state_key] + list_l2_for_state = system.param[ + 'LIST_INDEX_L2_BY_STATE_INDICES' + ][ordered_state_idx] + list_mode_l2_map.append((mode_idx, list_l2_for_state[0])) + return list_mode_l2_map + + +# ============================================================ +# Shared Setup +# ============================================================ + +nsite = 4 +e_lambda = 20.0 +gamma = 50.0 +temp = 140.0 +(g_0, w_0) = bcf_convert_dl_to_exp(e_lambda, gamma, temp) + +T3_loperator = np.zeros([4, 4, 4], dtype=np.float64) +list_gw_sysbath = [] +list_lop = [] +for i in range(nsite): + T3_loperator[i, i, i] = 1.0 + list_gw_sysbath.append([g_0, w_0]) + list_lop.append(sp.sparse.coo_matrix(T3_loperator[i])) + list_gw_sysbath.append([-1j * np.imag(g_0), 500.0]) + list_lop.append(T3_loperator[i]) + +H2_hamiltonian = np.zeros([nsite, nsite]) +H2_hamiltonian[0, 1] = 40 +H2_hamiltonian[1, 0] = 40 +H2_hamiltonian[1, 2] = 10 +H2_hamiltonian[2, 1] = 10 +H2_hamiltonian[2, 3] = 40 +H2_hamiltonian[3, 2] = 40 + +sys_param = { + 'HAMILTONIAN': np.array(H2_hamiltonian, dtype=np.complex128), + 'GW_SYSBATH': list_gw_sysbath, + 'L_HIER': list_lop, + 'L_NOISE1': list_lop, + 'ALPHA_NOISE1': bcf_exp, + 'PARAM_NOISE1': list_gw_sysbath, +} + +psi_0 = np.array([0.0] * nsite, dtype=np.complex128) +psi_0[2] = 1.0 + +k_max = 2 +state_list = np.arange(nsite) +delta_s = 0 +eom_param = {'EQUATION_OF_MOTION': 'NORMALIZED NONLINEAR'} + + +def _make_tb(sp=sys_param, ds=delta_s, psi=psi_0, sl=state_list): + """Creates an initialized HopsTensorBasis from sys_param dict.""" + system = HopsSystem(sp) + mode = HopsModes(system, hierarchy=None) + noise_memory = HopsNoiseMemory(system, mode) + tb = HopsTensorBasis(system, mode, noise_memory) + system.initialize(ds > 0, psi) + mode.list_modeidx_abs = sorted(system.list_statemodeidx_abs) + noise_memory.initialize() + system.state_list = sl + tb.initialize(ds) + return tb + + +def _make_tensor(method): + """Creates and initializes a HopsTensorWavefunction for the dimer-of-dimers system. + """ + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': method, + 'BOND_DIM_MAX': 20, + } + integrator_param = {'INTEGRATOR': 'RUNGE_KUTTA'} + tb = _make_tb() + ht = HopsTensorWavefunction(k_max, tensor_param, integrator_param, eom_param) + ht.initialize(psi_0, tb.system) + return ht + + +def _make_tensor_with_psi(method, psi): + """Same as _make_tensor but initializes with a caller-provided psi. + + Lets tests that need a non-default initial state reuse the standard + fixture without copy-pasting the HopsTensorWavefunction setup. + """ + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': method, + 'BOND_DIM_MAX': 20, + } + integrator_param = {'INTEGRATOR': 'RUNGE_KUTTA'} + tb = _make_tb(psi=psi) + ht = HopsTensorWavefunction(k_max, tensor_param, integrator_param, eom_param) + ht.initialize(psi, tb.system) + return ht + + +def _make_simple_mps(n_cores, phys_dims, bond_dims): + """Build a simple MPS with specified structure. + + Parameters + ---------- + 1. n_cores: int + Number of MPS cores (sites). + 2. phys_dims: list(int) + Physical dimension of each core. + 3. bond_dims: list(int) + Bond dimensions, length n_cores + 1. bond_dims[0] and + bond_dims[-1] should be 1 for open boundary conditions. + + Returns + ------- + 1. cores: list(np.ndarray) + MPS cores, each shaped (bond_left, phys_dim, bond_right). + """ + cores = [] + for i in range(n_cores): + core = np.zeros( + (bond_dims[i], phys_dims[i], bond_dims[i + 1]), + dtype=np.complex128, + ) + core += 0.01 * ( + np.random.randn(*core.shape) + 1j * np.random.randn(*core.shape) + ) + cores.append(core) + return cores + + +# ============================================================ +# TEST SUITE: recenter_to_zero() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Sites 1..L-1 become left-canonical +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_recenter_to_zero_left_canonical(): + # This case tests that after recentering, all sites n >= 1 satisfy + # M_n^dagger M_n = I (left-canonical form). + np.random.seed(42) + cores = _make_simple_mps(4, [2, 3, 3, 2], [1, 3, 4, 3, 1]) + result = recenter_to_zero(cores) + for n in range(1, len(result)): + Dl, d, Dr = result[n].shape + M = result[n].reshape(Dl * d, Dr) + eye_check = M.conj().T @ M + np.testing.assert_allclose( + eye_check, + np.eye(Dr, dtype=np.complex128), + atol=1e-12, + err_msg=f'Site {n} is not left-canonical', + ) + + +# ------------------------------------------------------------ +# TEST: Recentering preserves the represented state +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_recenter_to_zero_preserves_state(): + # This case tests that the full contraction of the MPS gives the same + # state vector before and after recentering. + np.random.seed(123) + cores = _make_simple_mps(3, [2, 2, 2], [1, 2, 2, 1]) + + def _contract(cs): + result = cs[0] + for c in cs[1:]: + result = np.tensordot(result, c, axes=([-1], [0])) + return result.squeeze() + + state_before = _contract([c.copy() for c in cores]) + result = recenter_to_zero(cores) + state_after = _contract(result) + np.testing.assert_allclose(state_after, state_before, atol=1e-12) + + +# ------------------------------------------------------------ +# TEST: ValueError when last core Dr != 1 +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_recenter_to_zero_bad_boundary(): + # This case tests that a ValueError is raised when the last core + # has Dr != 1 (violating open boundary conditions). + bad_core = np.zeros((1, 2, 2), dtype=np.complex128) + with pytest.raises(ValueError, match='open boundaries'): + recenter_to_zero([bad_core]) + + +# ------------------------------------------------------------ +# TEST: recenter_to_zero raises on internal bond mismatch +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_recenter_to_zero_bond_mismatch(): + # This case tests that a ValueError is raised when adjacent cores + # have incompatible bond dimensions. The upfront validation catches + # mismatches for any number of cores, including 2-core MPS. + + # 2-core case: right bond of core_0 (3) != left bond of core_1 (2) + core_0 = np.ones((1, 2, 3), dtype=np.complex128) + core_1 = np.ones((2, 2, 1), dtype=np.complex128) + with pytest.raises(ValueError, match='Bond mismatch'): + recenter_to_zero([core_0, core_1]) + + # 3-core case: right bond of core_1 (4) != left bond of core_2 (2) + core_0 = np.ones((1, 2, 3), dtype=np.complex128) + core_1 = np.ones((3, 2, 4), dtype=np.complex128) + core_2 = np.ones((2, 2, 1), dtype=np.complex128) + with pytest.raises(ValueError, match='Bond mismatch'): + recenter_to_zero([core_0, core_1, core_2]) + + +# ------------------------------------------------------------ +# TEST: normalize=True makes ||A[0]||_F = 1 +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_recenter_to_zero_normalize(): + # This case tests that normalize=True rescales the MPS so that + # the Frobenius norm of A[0] equals 1 and the state direction is preserved. + np.random.seed(7) + cores = _make_simple_mps(3, [2, 3, 2], [1, 3, 3, 1]) + + def contract(c_list): + s = c_list[0] + for c in c_list[1:]: + s = np.tensordot(s, c, axes=([-1], [0])) + return s.squeeze().ravel() + + state_orig = contract(cores) + result = recenter_to_zero(cores, normalize=True) + norm_A0 = np.linalg.norm(result[0].ravel()) + np.testing.assert_allclose(norm_A0, 1.0, atol=1e-12) + # The contracted state should be proportional to the original + state_norm = contract(result) + ratio = state_orig / state_norm + np.testing.assert_allclose( + ratio, ratio[0] * np.ones_like(ratio), atol=1e-12, + err_msg='Normalized state is not proportional to original', + ) + + +# ------------------------------------------------------------ +# TEST: recenter_to_zero copy=False reuses the input list +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_recenter_to_zero_no_copy(): + # This case tests that copy=False returns the same list object and + # that the cores are actually left-canonical after the call. + ht = _make_tensor('fullstate') + ht.inflate_bonds_to(4, eps=0.01) + cores = ht.list_cores_phi + result = recenter_to_zero(cores, copy=False) + assert result is cores + # Verify cores 1..L-1 are left-canonical (Q^dQ = I) + for n in range(1, len(result)): + Dl, d, Dr = result[n].shape + Q = result[n].reshape(Dl * d, Dr) + np.testing.assert_allclose( + Q.conj().T @ Q, np.eye(Dr), atol=1e-12, + err_msg=f'Core {n} is not left-canonical after copy=False', + ) + + +# ------------------------------------------------------------ +# TEST: recenter_to_zero copy=True returns independent cores +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_recenter_to_zero_copy_independence(): + # This case tests that copy=True returns cores that are genuinely + # independent from the input: mutating the result must not affect + # the original. A shallow-copy bug (new list object but shared + # ndarray references) would pass the list-identity check yet leak + # mutations through; this test catches that. + ht = _make_tensor('fullstate') + ht.inflate_bonds_to(4, eps=0.01) + cores = ht.list_cores_phi + # Snapshot every input core so we can detect any leaked mutation + cores_snapshot = [c.copy() for c in cores] + result = recenter_to_zero(cores, copy=True) + # Outer list must be a new object + assert result is not cores + # Mutate every core of the result; the originals must be unaffected + for core in result: + core[...] = 0.0 + for i, (before, after) in enumerate(zip(cores_snapshot, cores)): + np.testing.assert_array_equal( + after, before, + err_msg=( + f'copy=True leaked: mutating result core {i} also ' + f'modified original cores[{i}]' + ), + ) + + +# ------------------------------------------------------------ +# TEST: recenter_to_zero with single core preserves state +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_recenter_to_zero_single_core(): + # Limiting case: single core MPS, sweep is empty, state preserved + core = np.array([[[1.0 + 0j], [0.5 + 0.1j]]]) # shape (1, 2, 1) + result = recenter_to_zero([core]) + np.testing.assert_allclose(result[0], core, atol=1e-14) + + +# ------------------------------------------------------------ +# TEST: recenter_to_zero with two cores (R goes directly to scalar) +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_recenter_to_zero_two_cores(): + # This case tests the two-core MPS where the sweep runs exactly once + # and R folds directly into site 0 (no absorb-into-next-core branch). + np.random.seed(42) + cores = _make_simple_mps(2, [3, 2], [1, 4, 1]) + + def contract(c_list): + s = c_list[0] + for c in c_list[1:]: + s = np.tensordot(s, c, axes=([-1], [0])) + return s.squeeze().ravel() + + state_before = contract(cores) + result = recenter_to_zero(cores) + state_after = contract(result) + np.testing.assert_allclose(state_after, state_before, atol=1e-12) + + +# ------------------------------------------------------------ +# TEST: recenter_to_zero on already-canonical MPS is near-no-op +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_recenter_to_zero_already_canonical(): + # This case tests that passing an already-canonical MPS is a near-no-op: + # both the contracted physical state and the individual cores are + # preserved (within numerical precision). + np.random.seed(99) + cores = _make_simple_mps(3, [2, 3, 2], [1, 3, 3, 1]) + canonical = recenter_to_zero(cores) + + def contract(c_list): + s = c_list[0] + for c in c_list[1:]: + s = np.tensordot(s, c, axes=([-1], [0])) + return s.squeeze().ravel() + + state_first = contract(canonical) + result = recenter_to_zero(canonical) + state_second = contract(result) + # State-level idempotency: the contracted quantum state is preserved + np.testing.assert_allclose(state_second, state_first, atol=1e-12) + # Core-level idempotency: each core is preserved element-wise. Catches + # per-core drift (e.g. sign/phase shifts) that would cancel at contraction. + for n in range(len(canonical)): + np.testing.assert_allclose( + result[n], canonical[n], atol=1e-12, + err_msg=f'Core {n} was modified by re-canonicalization', + ) + + +# ------------------------------------------------------------ +# TEST: recenter_to_zero with all bond dims = 1 +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_recenter_to_zero_trivial_bonds(): + # This case tests the degenerate geometry where every bond is 1, + # so QR reduces to scalar extraction at every site. + np.random.seed(11) + cores = _make_simple_mps(4, [2, 3, 2, 2], [1, 1, 1, 1, 1]) + + def contract(c_list): + s = c_list[0] + for c in c_list[1:]: + s = np.tensordot(s, c, axes=([-1], [0])) + return s.squeeze().ravel() + + state_before = contract(cores) + result = recenter_to_zero(cores) + state_after = contract(result) + np.testing.assert_allclose(state_after, state_before, atol=1e-12) + + +# ------------------------------------------------------------ +# TEST: recenter_to_zero normalize with zero-norm MPS +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_recenter_to_zero_zero_norm(): + # This case tests the norm_phi > 0 guard: a zero MPS should be + # returned without division by zero. The QR sweep may introduce + # non-zero Q factors, but A[0] should remain zero (R[0,0] = 0 + # is folded into it), so the overall state is still zero. + cores = [ + np.zeros((1, 2, 3), dtype=np.complex128), + np.zeros((3, 3, 1), dtype=np.complex128), + ] + result = recenter_to_zero(cores, normalize=True) + # The orthogonality center (site 0) should be zero + np.testing.assert_allclose(result[0], 0.0, atol=1e-15) + + +# ============================================================ +# TEST SUITE: tensor_matvec_prod() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Mismatched core counts raise ValueError +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tensor_matvec_prod_length_mismatch(): + # This case tests that a ValueError is raised when the MPS and MPO + # have different numbers of cores. + ht = _make_tensor('fullstate') + list_cores_short_mpo = [np.zeros((1, 2, 2, 1), dtype=np.complex128)] + with pytest.raises(ValueError, match='same number of cores'): + tensor_matvec_prod( + ht.list_cores_phi, + list_cores_short_mpo, + ht.mps_epsilon, + ht.bond_dim_max, + ) + + +# ------------------------------------------------------------ +# TEST: Non-trivial MPO (diagonal scaling) produces correct result +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tensor_matvec_prod_diagonal_scaling(): + # This case tests that a diagonal scaling MPO (multiply every physical + # index by 2) correctly doubles the physical wavefunction. This exercises + # the einsum contraction with non-identity operator entries. + ht = _make_tensor('fullstate') + V1_phi_before = ht.psi.copy() + # Build a 2*I MPO: each core applies 2*identity on the physical index + list_cores_scale = [] + for core in ht.list_cores_phi: + phys_dim = core.shape[1] + mpo_core = np.zeros((1, phys_dim, phys_dim, 1), dtype=np.complex128) + mpo_core[0, :, :, 0] = 2.0 * np.eye(phys_dim) + list_cores_scale.append(mpo_core) + # Apply the scaling MPO + list_cores_result, _ = tensor_matvec_prod( + ht.list_cores_phi, + list_cores_scale, + ht.mps_epsilon, + ht.bond_dim_max, + ) + V1_phi_after = extract_psi( + list_cores_result, + 'fullstate', + ht.M1_modes_per_state, + ) + # The system core carries the physical dimension; mode cores get identity + # scaled by 2 each, so phi_0 is multiplied by 2^(n_cores). + n_cores = len(ht.list_cores_phi) + expected = V1_phi_before * (2.0**n_cores) + np.testing.assert_allclose( + V1_phi_after, + expected, + atol=1e-10, + err_msg='Diagonal scaling MPO did not produce correct result', + ) + + +# ------------------------------------------------------------ +# TEST: Compression respects bond_dim_max +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tensor_matvec_prod_compression(): + # This case tests that the SVD compression enforces bond_dim_max. + # We inflate the MPS bonds first so the contracted result would + # exceed bond_dim_max=2 without compression. + ht = _make_tensor('fullstate') + ht.inflate_bonds_to(4, eps=0.01) + # Build identity MPO with bond dim 2 to further inflate contracted bonds + list_cores_mpo = [] + for core in ht.list_cores_phi: + phys_dim = core.shape[1] + mpo_core = np.zeros((2, phys_dim, phys_dim, 2), dtype=np.complex128) + mpo_core[0, :, :, 0] = np.eye(phys_dim) + mpo_core[1, :, :, 1] = np.eye(phys_dim) * 0.001 + list_cores_mpo.append(mpo_core) + # Fix boundary bonds to 1 + list_cores_mpo[0] = list_cores_mpo[0][:1, :, :, :] + list_cores_mpo[-1] = list_cores_mpo[-1][:, :, :, :1] + # Apply with tight bond_dim_max + small_bond = 3 + list_cores_result, _ = tensor_matvec_prod( + ht.list_cores_phi, + list_cores_mpo, + 1e-10, + small_bond, + ) + # Verify all internal bonds respect bond_dim_max + for i, core in enumerate(list_cores_result): + assert core.shape[0] <= small_bond, ( + f'Core {i} left bond {core.shape[0]} exceeds bond_dim_max={small_bond}' + ) + assert core.shape[2] <= small_bond, ( + f'Core {i} right bond {core.shape[2]} exceeds bond_dim_max={small_bond}' + ) + # Invariant: compressed result should approximate the uncompressed contraction + list_cores_uncompressed, _ = tensor_matvec_prod( + ht.list_cores_phi, + list_cores_mpo, + 1e-14, # very tight epsilon + 999, # no bond cap + ) + psi_compressed = extract_psi( + list_cores_result, ht.method, ht.M1_modes_per_state, + ) + psi_uncompressed = extract_psi( + list_cores_uncompressed, ht.method, ht.M1_modes_per_state, + ) + np.testing.assert_allclose( + psi_compressed, psi_uncompressed, atol=1e-4, + err_msg='Compressed matvec should approximate uncompressed', + ) + + +# ------------------------------------------------------------ +# TEST: Single-site MPS edge case +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tensor_matvec_prod_single_site(): + # This case tests that tensor_matvec_prod works correctly when the + # MPS has exactly one core (the loop body runs once). + phys_dim = 3 + # Single MPS core: (1, 3, 1) + list_cores_vec = [np.array([[[1.0 + 0j], [0.5 + 0.3j], [0.0 + 0.2j]]])] + # Single MPO core: 2*I on phys_dim=3 + mpo_core = np.zeros((1, phys_dim, phys_dim, 1), dtype=np.complex128) + mpo_core[0, :, :, 0] = 2.0 * np.eye(phys_dim) + list_cores_mpo = [mpo_core] + list_cores_result, _ = tensor_matvec_prod( + list_cores_vec, + list_cores_mpo, + 1e-10, + 20, + ) + # Should be a single core with the same shape + assert len(list_cores_result) == 1 + assert list_cores_result[0].shape == (1, phys_dim, 1) + # Values should be doubled + np.testing.assert_allclose( + list_cores_result[0], + 2.0 * list_cores_vec[0], + atol=1e-14, + ) + + +# ------------------------------------------------------------ +# TEST: Rectangular MPO (dim_out != dim_in) reshapes physical dimension +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tensor_matvec_prod_rectangular_mpo(): + # This case tests that tensor_matvec_prod correctly handles an MPO + # whose physical output dimension differs from its physical input + # dimension (rectangular operator). The einsum 'LoiR,lir->LloRr' + # contracts only the i index, so dim_out and dim_in need not match; + # the output core's physical dim equals the MPO's dim_out. + dim_in = 2 + dim_out = 3 + # Single MPS core with phys_dim = dim_in + V1_vec = np.array([1.0 + 0j, 0.5 + 0.2j], dtype=np.complex128) + list_cores_vec = [V1_vec.reshape(1, dim_in, 1)] + # Non-trivial rectangular MPO core: embeds the 2-dim vector into 3 + # dimensions. The last row sums the two input components so an + # index swap or scalar factor bug would fail the value check. + M2_embed = np.array( + [[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]], dtype=np.complex128, + ) + mpo_core = M2_embed.reshape(1, dim_out, dim_in, 1) + list_cores_mpo = [mpo_core] + list_cores_result, _ = tensor_matvec_prod( + list_cores_vec, list_cores_mpo, 1e-10, 20, + ) + # Output shape carries dim_out as the physical dimension + assert len(list_cores_result) == 1 + assert list_cores_result[0].shape == (1, dim_out, 1) + # Output values match the direct rectangular matrix product M @ v + V1_expected = M2_embed @ V1_vec + np.testing.assert_allclose( + list_cores_result[0].reshape(dim_out), V1_expected, atol=1e-12, + err_msg='Rectangular MPO did not match direct matrix-vector product', + ) + + +# ------------------------------------------------------------ +# TEST: Input cores are not mutated by tensor_matvec_prod +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tensor_matvec_prod_input_not_mutated(): + # This case tests that tensor_matvec_prod produces a new result list + # without modifying the input MPS cores. Element-wise equality against + # a snapshot of the inputs catches any in-place mutation (e.g. a + # downstream helper reshaping input cores instead of copies). + ht = _make_tensor('fullstate') + # Snapshot copies of every input core so any in-place mutation by + # tensor_matvec_prod would diverge the originals from the snapshot. + list_cores_vec_snapshot = [c.copy() for c in ht.list_cores_phi] + # Build a non-trivial MPO so the contraction path is fully exercised + list_cores_mpo = [] + for core in ht.list_cores_phi: + phys_dim = core.shape[1] + mpo_core = np.zeros((1, phys_dim, phys_dim, 1), dtype=np.complex128) + mpo_core[0, :, :, 0] = 2.0 * np.eye(phys_dim) + list_cores_mpo.append(mpo_core) + _, _ = tensor_matvec_prod( + ht.list_cores_phi, list_cores_mpo, ht.mps_epsilon, ht.bond_dim_max, + ) + # Verify every input core is unchanged from its pre-call snapshot + for i, (before, after) in enumerate( + zip(list_cores_vec_snapshot, ht.list_cores_phi) + ): + np.testing.assert_array_equal( + after, before, + err_msg=f'Input core {i} was mutated by tensor_matvec_prod', + ) + + +# ------------------------------------------------------------ +# TEST: Nested input cores are not mutated by tensor_matvec_prod +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tensor_matvec_prod_input_not_mutated_nested(): + # This case tests that the nested (number) path + # also preserves input cores. flatten_cores returns a new list but + # reuses ndarray references, so any in-place mutation downstream + # would be visible through the original nested structure. + ht = _make_tensor('number') + # Snapshot every core inside every group before the call + list_cores_phi_snapshot = [ + [c.copy() for c in group] for group in ht.list_cores_phi + ] + # Build a non-trivial MPO matching the flat core count + list_cores_mpo = [] + for core in ht.flat_cores: + phys_dim = core.shape[1] + mpo_core = np.zeros((1, phys_dim, phys_dim, 1), dtype=np.complex128) + mpo_core[0, :, :, 0] = 2.0 * np.eye(phys_dim) + list_cores_mpo.append(mpo_core) + _, _ = tensor_matvec_prod( + ht.list_cores_phi, list_cores_mpo, ht.mps_epsilon, ht.bond_dim_max, + ) + # Verify every nested input core is unchanged from its pre-call snapshot + for s, (group_before, group_after) in enumerate( + zip(list_cores_phi_snapshot, ht.list_cores_phi) + ): + assert len(group_before) == len(group_after), ( + f'Group {s} size changed from {len(group_before)} to {len(group_after)}' + ) + for i, (before, after) in enumerate(zip(group_before, group_after)): + np.testing.assert_array_equal( + after, before, + err_msg=( + f'Nested input core at group {s}, position {i} ' + f'was mutated by tensor_matvec_prod' + ), + ) + + +# ------------------------------------------------------------ +# TEST: Empty input raises a clear ValueError +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tensor_matvec_prod_empty_list(): + # This case tests that tensor_matvec_prod guards against empty input + # with a clear ValueError rather than leaking an IndexError from a + # downstream helper (tensor_compress previously indexed into the + # empty list). + with pytest.raises(ValueError, match='at least one core'): + tensor_matvec_prod([], [], 1e-10, 20) + + +# ------------------------------------------------------------ +# TEST: Off-diagonal MPO mixes physical states correctly +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tensor_matvec_prod_off_diagonal_mpo(): + # This case tests that tensor_matvec_prod correctly handles an + # off-diagonal MPO that mixes physical states. Uses a Pauli-X + # (NOT gate) on the state core of a fullstate MPS: + # X = [[0,1],[1,0]] swaps |0⟩ ↔ |1⟩. + # For psi = [a, b], X|psi⟩ = [b, a]. + # + # We build a 2-state MPS, apply an MPO with X on the state core + # and identity on mode cores, then verify the physical wavefunction + # components are swapped. + nsite_local = 2 + sys_param_local = { + 'HAMILTONIAN': np.zeros([nsite_local, nsite_local], dtype=np.complex128), + 'GW_SYSBATH': list_gw_sysbath[:2], + 'L_HIER': list_lop[:1], + 'L_NOISE1': list_lop[:1], + 'ALPHA_NOISE1': bcf_exp, + 'PARAM_NOISE1': list_gw_sysbath[:2], + } + tensor_param_local = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, + } + psi_local = np.array([0.3 + 0.1j, 0.7 - 0.2j], dtype=np.complex128) + tb = _make_tb(sp=sys_param_local, ds=0, psi=psi_local, sl=np.arange(nsite_local)) + ht = HopsTensorWavefunction( + k_max, tensor_param_local, {'INTEGRATOR': 'RUNGE_KUTTA'}, eom_param, + ) + ht.initialize(psi_local, tb.system) + # Extract phi before applying MPO + V1_phi_before = ht.psi.copy() + # Build an MPO with Pauli-X on state core, identity on mode cores + list_cores_mpo = [] + for core in ht.list_cores_phi: + phys_dim = core.shape[1] + mpo_core = np.zeros((1, phys_dim, phys_dim, 1), dtype=np.complex128) + if phys_dim == nsite_local: + # State core: apply Pauli-X (swap) + pauli_x = np.array([[0, 1], [1, 0]], dtype=np.complex128) + mpo_core[0, :, :, 0] = pauli_x + else: + # Mode core: identity + mpo_core[0, :, :, 0] = np.eye(phys_dim) + list_cores_mpo.append(mpo_core) + # Apply the MPO + list_cores_result, _ = tensor_matvec_prod( + ht.list_cores_phi, + list_cores_mpo, + ht.mps_epsilon, + ht.bond_dim_max, + ) + V1_phi_after = extract_psi( + list_cores_result, + 'fullstate', + ht.M1_modes_per_state, + ) + # Pauli-X swaps components: [a, b] → [b, a] + V1_expected = np.array( + [V1_phi_before[1], V1_phi_before[0]], + dtype=np.complex128, + ) + np.testing.assert_allclose( + V1_phi_after, + V1_expected, + atol=1e-10, + err_msg='Off-diagonal MPO (Pauli-X) did not swap physical states', + ) + + +# ------------------------------------------------------------ +# TEST: Nested statenumber MPS path (flatten/contract/unflatten) +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tensor_matvec_prod_nested_statenumber(): + # This case tests the is_nested branch: passing a list-of-lists + # statenumber MPS and verifying the result is correctly unflattened. + ht = _make_tensor('number') + V1_phi_before = ht.psi.copy() + # Build identity MPO matching the flattened core count + from mesohops.util.tensor_operations import flatten_cores + flat = flatten_cores(ht.list_cores_phi) + list_cores_mpo = [] + for core in flat: + phys_dim = core.shape[1] + mpo_core = np.zeros((1, phys_dim, phys_dim, 1), dtype=np.complex128) + mpo_core[0, :, :, 0] = np.eye(phys_dim) + list_cores_mpo.append(mpo_core) + # Pass the nested (list-of-lists) MPS directly + result, _ = tensor_matvec_prod( + ht.list_cores_phi, list_cores_mpo, ht.mps_epsilon, ht.bond_dim_max, + ) + # Result should be nested (list-of-lists) with same structure + assert isinstance(result[0], list), ( + 'Nested input should produce nested output' + ) + assert len(result) == len(ht.list_cores_phi) + for g_res, g_orig in zip(result, ht.list_cores_phi): + assert len(g_res) == len(g_orig), ( + 'Group size should be preserved after flatten/unflatten' + ) + # Identity MPO should preserve the wavefunction + V1_phi_after = extract_psi(result, 'number', + ht.M1_modes_per_state) + np.testing.assert_allclose(V1_phi_after, V1_phi_before, atol=1e-10) + + +# ------------------------------------------------------------ +# TEST: Nested input with MPO length mismatch raises ValueError +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tensor_matvec_prod_length_mismatch_nested(): + # This case tests that the length-mismatch guard fires on the nested + # trajectory too: a statenumber MPS is flattened first, so an MPO + # whose length matches neither the outer-group count nor the flat + # core count must still trip the "same number of cores" check. + ht = _make_tensor('number') + list_cores_short_mpo = [np.zeros((1, 2, 2, 1), dtype=np.complex128)] + with pytest.raises(ValueError, match='same number of cores'): + tensor_matvec_prod( + ht.list_cores_phi, + list_cores_short_mpo, + ht.mps_epsilon, + ht.bond_dim_max, + ) + + +# ------------------------------------------------------------ +# TEST: Tuple-of-tuples nested input is routed through the nested path +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tensor_matvec_prod_tuple_nested_input(): + # This case tests that a tuple-of-tuples MPS (not list-of-lists) is + # detected as nested and produces the same result as the equivalent + # list-of-lists input. Catches regressions of the isinstance check + # that previously accepted only `list`, silently misclassifying + # tuple-wrapped nested MPSs as flat and failing in the einsum. + ht = _make_tensor('number') + # Build identity MPO matching the flat core count + list_cores_mpo = [] + for core in ht.flat_cores: + phys_dim = core.shape[1] + mpo_core = np.zeros((1, phys_dim, phys_dim, 1), dtype=np.complex128) + mpo_core[0, :, :, 0] = np.eye(phys_dim) + list_cores_mpo.append(mpo_core) + # Apply with the list-of-lists input (baseline) + list_result, _ = tensor_matvec_prod( + ht.list_cores_phi, list_cores_mpo, ht.mps_epsilon, ht.bond_dim_max, + ) + # Apply with the tuple-of-tuples copy of the same data + tuple_input = tuple(tuple(group) for group in ht.list_cores_phi) + tuple_result, _ = tensor_matvec_prod( + tuple_input, list_cores_mpo, ht.mps_epsilon, ht.bond_dim_max, + ) + # Both results should be nested and structurally identical + assert isinstance(list_result[0], list), 'list input should stay nested' + assert isinstance(tuple_result[0], list), 'tuple input should unflatten to lists' + assert len(list_result) == len(tuple_result) + for group_list, group_tuple in zip(list_result, tuple_result): + assert len(group_list) == len(group_tuple) + for core_list, core_tuple in zip(group_list, group_tuple): + np.testing.assert_allclose(core_list, core_tuple, atol=1e-14) + + +# ------------------------------------------------------------ +# TEST: Per-site different operators catches mispairing +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tensor_matvec_prod_per_site_scaling(): + # This case tests that different operators on different sites are + # applied to the correct cores. A uniform scaling (2*I on every site) + # would mask a site-pairing bug; here we use 2*I on the state core + # and 3*I on mode cores to verify each factor is applied correctly. + ht = _make_tensor('fullstate') + V1_phi_before = ht.psi.copy() + list_cores_mpo = [] + for i, core in enumerate(ht.list_cores_phi): + phys_dim = core.shape[1] + mpo_core = np.zeros((1, phys_dim, phys_dim, 1), dtype=np.complex128) + if i == 0: + mpo_core[0, :, :, 0] = 2.0 * np.eye(phys_dim) + else: + mpo_core[0, :, :, 0] = 3.0 * np.eye(phys_dim) + list_cores_mpo.append(mpo_core) + result, _ = tensor_matvec_prod( + ht.list_cores_phi, list_cores_mpo, ht.mps_epsilon, ht.bond_dim_max, + ) + V1_phi_after = extract_psi( + result, 'fullstate', ht.M1_modes_per_state, + ) + # State core scaled by 2, each of n_mode mode cores scaled by 3 + n_modes = len(ht.list_cores_phi) - 1 + expected = V1_phi_before * 2.0 * (3.0 ** n_modes) + np.testing.assert_allclose(V1_phi_after, expected, atol=1e-10, + err_msg='Per-site scaling factors not applied to correct cores') + + +# ------------------------------------------------------------ +# TEST: Non-trivial Hamiltonian-style MPO vs dense H2_op @ psi +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tensor_matvec_prod_general_operator(): + # This case tests tensor_matvec_prod against the actual statenumber + # operator MPO from mpo_constructors. Unlike the identity / scaling / + # Pauli-X tests, this MPO has: + # - Non-trivial bond dimension (> 1) carrying state across sites + # - Off-diagonal entries that mix physical indices across sites + # via daisy-chained transfer matrices + # The physical wavefunction after contraction is cross-validated + # against the direct dense H2_op @ psi product, providing a reference + # independent of the MPS contraction logic itself. + ht = _make_tensor('number') + V1_phi_before = extract_psi( + ht.list_cores_phi, ht.method, ht.M1_modes_per_state, + ) + # Non-trivial dense system operator: every entry is nonzero so every + # bond channel of the MPO (diagonal, left/right transfer, daisy-chain) + # is exercised. + H2_op = np.array([ + [1.0 + 0.1j, 0.3 + 0.1j, 0.2 - 0.1j, 0.1 + 0.05j], + [0.3 - 0.1j, 0.8 + 0.0j, 0.4 + 0.2j, 0.2 + 0.0j], + [0.2 + 0.1j, 0.4 - 0.2j, 0.5 + 0.0j, 0.3 + 0.1j], + [0.1 - 0.05j, 0.2 + 0.0j, 0.3 - 0.1j, 0.9 + 0.0j], + ], dtype=np.complex128) + list_cores_mpo = build_statenumber_operator_mpo( + H2_op, nsite, k_max, ht.M1_modes_per_state, + ) + # Confirm the MPO is genuinely non-trivial (bond dim > 1 somewhere) + max_bond = max( + max(c.shape[0] for c in list_cores_mpo), + max(c.shape[3] for c in list_cores_mpo), + ) + assert max_bond > 1, f'Expected non-trivial MPO bond, got {max_bond}' + # Apply the MPO via flat-core path (nested-path coverage is in a + # dedicated test upstream). + list_cores_result_flat, _ = tensor_matvec_prod( + ht.flat_cores, list_cores_mpo, ht.mps_epsilon, ht.bond_dim_max, + ) + list_cores_result = unflatten_cores( + list_cores_result_flat, ht.M1_modes_per_state, + ) + V1_phi_after = extract_psi( + list_cores_result, ht.method, ht.M1_modes_per_state, + ) + # Dense cross-validation: the MPO applies H2_op to the system space + # and identity to mode cores, so the resulting physical wavefunction + # equals H2_op @ V1_phi_before. + V1_expected = H2_op @ V1_phi_before + np.testing.assert_allclose( + V1_phi_after, V1_expected, atol=1e-10, + err_msg='MPS contraction does not match dense H2_op @ psi reference', + ) + + +# ------------------------------------------------------------ +# TEST: Returned complexity scalar matches calc_mps_complexity +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tensor_matvec_prod_complexity_scalar_matches_formula(): + # This case tests that the complexity scalar returned by + # tensor_matvec_prod equals calc_mps_complexity(list_cores_compressed) + # — i.e. the per-core sum of D_left * D_right * max(D_left, D_right) * + # d_phys evaluated on the contracted-but-uncompressed cores. + # Manually compute the post-contraction core shapes so the assertion + # is independent of the implementation's internal variable name. + ht = _make_tensor('fullstate') + # Build a non-trivial 2*I MPO (bond dim 1) so each contracted core + # has shape (mpo_bond * mps_bond, dim_phys, mpo_bond * mps_bond); + # for bond-1 inputs that's identical to the input bond shape. + list_cores_mpo = [] + for core in ht.list_cores_phi: + phys_dim = core.shape[1] + mpo_core = np.zeros((1, phys_dim, phys_dim, 1), dtype=np.complex128) + mpo_core[0, :, :, 0] = 2.0 * np.eye(phys_dim) + list_cores_mpo.append(mpo_core) + # Reproduce the contracted-but-uncompressed core shapes the + # implementation builds internally (see tensor_eom_functions L92-95): + # shape (mpo_dl * vec_dl, dim_out, mpo_dr * vec_dr). + list_uncompressed = [] + for core_mpo, core_vec in zip(list_cores_mpo, ht.list_cores_phi): + mpo_dl, dim_out, _, mpo_dr = core_mpo.shape + vec_dl, _, vec_dr = core_vec.shape + list_uncompressed.append( + np.zeros( + (mpo_dl * vec_dl, dim_out, mpo_dr * vec_dr), + dtype=np.complex128, + ) + ) + expected_complexity = calc_mps_complexity(list_uncompressed) + _, complexity = tensor_matvec_prod( + ht.list_cores_phi, list_cores_mpo, ht.mps_epsilon, ht.bond_dim_max, + ) + assert complexity == expected_complexity, ( + f'tensor_matvec_prod complexity {complexity} does not match ' + f'calc_mps_complexity of pre-compress cores {expected_complexity}' + ) + + +# ------------------------------------------------------------ +# TEST: Output cores are 3-D with compatible bonds between neighbors +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tensor_matvec_prod_output_bond_compatibility(): + # This case tests that the output MPS has a valid bond structure: + # every core is 3-dimensional, and the right bond of each core + # matches the left bond of its neighbor. Catches bugs where the + # contraction or compression produces mismatched adjacent cores, + # which would only surface later when the result is used in any + # downstream tensor operation (contract, extract_psi, etc.). + ht = _make_tensor('fullstate') + # Build a non-trivial 2*I MPO so the full contraction + compression + # pipeline runs (not a degenerate bond-1 pass-through). + list_cores_mpo = [] + for core in ht.list_cores_phi: + phys_dim = core.shape[1] + mpo_core = np.zeros((1, phys_dim, phys_dim, 1), dtype=np.complex128) + mpo_core[0, :, :, 0] = 2.0 * np.eye(phys_dim) + list_cores_mpo.append(mpo_core) + list_cores_result, _ = tensor_matvec_prod( + ht.list_cores_phi, list_cores_mpo, ht.mps_epsilon, ht.bond_dim_max, + ) + # Every output core must be 3-D (dim_left, dim_phys, dim_right) + for i, core in enumerate(list_cores_result): + assert core.ndim == 3, ( + f'Output core {i} is not 3-D (got shape {core.shape})' + ) + # Bond compatibility: right bond of core i matches left bond of core i+1 + for i in range(len(list_cores_result) - 1): + dim_right_i = list_cores_result[i].shape[2] + dim_left_next = list_cores_result[i + 1].shape[0] + assert dim_right_i == dim_left_next, ( + f'Bond mismatch between cores {i} and {i + 1}: ' + f'right bond of core {i} is {dim_right_i} but left bond ' + f'of core {i + 1} is {dim_left_next}' + ) + # Left boundary of the first core and right boundary of the last + # core should both be size 1 (open-boundary MPS) + assert list_cores_result[0].shape[0] == 1, ( + f'First core left bond is not 1 (got {list_cores_result[0].shape[0]})' + ) + assert list_cores_result[-1].shape[2] == 1, ( + f'Last core right bond is not 1 (got {list_cores_result[-1].shape[2]})' + ) + + +# ------------------------------------------------------------ +# TEST: Identity MPO preserves auxiliary-core slices, not just phi_0 +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tensor_matvec_prod_identity_preserves_aux_slices(): + # This case tests that an identity MPO preserves the full quantum + # state across the entire system x modes product basis, not just the + # phi_0 projection that extract_psi (which slices each mode core at + # index 0) returns. A contraction bug that corrupted auxiliary + # slices (mode-core indices > 0) while leaving the [0] slice intact + # would be invisible to extract_psi-based tests; comparing the full + # tensor contraction catches it. + # Individual cores aren't compared directly because SVD compression + # has gauge freedom (sign/phase flips on adjacent bonds that cancel + # in the contracted state), so element-wise core equality is too + # strict; the full contraction is the gauge-invariant observable. + ht = _make_tensor('fullstate') + # Inflate bonds so mode cores carry non-trivial auxiliary amplitudes + # (not just [0] = 1). Without this the auxiliary slices are all + # zero and corruption would be undetectable. + ht.inflate_bonds_to(4, eps=0.05) + # Build an identity MPO: one core per MPS core, bond dim 1 + list_cores_mpo = [] + for core in ht.list_cores_phi: + phys_dim = core.shape[1] + mpo_core = np.zeros((1, phys_dim, phys_dim, 1), dtype=np.complex128) + mpo_core[0, :, :, 0] = np.eye(phys_dim) + list_cores_mpo.append(mpo_core) + list_cores_result, _ = tensor_matvec_prod( + ht.list_cores_phi, list_cores_mpo, ht.mps_epsilon, ht.bond_dim_max, + ) + + def _contract_cores(list_cores): + # Sequentially tensordot adjacent bonds to build the full state + # across every physical index (system + all modes). Output shape + # collapses to the product of physical dimensions. + state = list_cores[0] + for core in list_cores[1:]: + state = np.tensordot(state, core, axes=([-1], [0])) + return state.squeeze() + + state_before = _contract_cores(ht.list_cores_phi) + state_after = _contract_cores(list_cores_result) + np.testing.assert_allclose( + state_after, state_before, atol=1e-10, + err_msg=( + 'Full state contraction changed under identity MPO; ' + 'auxiliary-slice contributions may be corrupted' + ), + ) + + +# ============================================================ +# TEST SUITE: calc_norm_corr_tensor() +# ============================================================ + + +def _make_norm_corr_args( + method='fullstate', z_scale=0.0, psi=None, + z_rnd_func=None, aux_populate=None, +): + """Helper: builds all arguments for calc_norm_corr_tensor. + + Parameters + ---------- + 1. method: str + 2. z_scale: float + Scale factor for uniform noise. Ignored if z_rnd_func given. + 3. psi: np.ndarray or None + Initial state. Defaults to module-level psi_0. + 4. z_rnd_func: callable or None + f(n_l2) -> np.ndarray of noise values. Overrides z_scale. + 5. aux_populate: list of (mode_idx, value) or None + Manually populate first-order auxiliaries in the MPS. + + Returns + ------- + 1. dict with keys: 'wavefunction', 'z_hat', 'list_avg_L2', 'mode', + 'list_mode_l2_map', and 'system' (for cross-validation). + """ + if psi is None: + psi = psi_0 + + # Build the tensor wavefunction that calc_norm_corr_tensor consumes + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': method, + 'BOND_DIM_MAX': 20, + } + tb = _make_tb(psi=psi) + ht = HopsTensorWavefunction( + k_max, tensor_param, {'INTEGRATOR': 'RUNGE_KUTTA'}, eom_param, + ) + ht.initialize(psi, tb.system) + + # Optionally inject amplitude into first-order mode cores. Use + # flat_cores + update_phi_from_flat so the same (mode_idx, value) + # pair targets the correct core in both representations even + # though the flat layouts differ (fullstate interleaves no state + # cores among modes; statenumber interleaves one state core per + # site group). + if aux_populate is not None: + flat = [c.copy() for c in ht.flat_cores] + for (m_idx, val) in aux_populate: + if method == 'fullstate': + core_idx = 1 + m_idx + else: + # Walk through state groups until m_idx lands in one + modes_per_state = np.asarray( + ht.M1_modes_per_state, dtype=int, + ) + state_idx = 0 + remaining = int(m_idx) + while remaining >= modes_per_state[state_idx]: + remaining -= modes_per_state[state_idx] + state_idx += 1 + offset = sum( + 1 + int(modes_per_state[s]) for s in range(state_idx) + ) + core_idx = offset + 1 + remaining + flat[core_idx][0, 1, 0] = val + ht.update_phi_from_flat(flat) + + # Pull mode/system indexers needed for downstream construction + mode = tb.mode + system = tb.system + list_l2idx_abs = mode.list_l2idx_abs + list_index_L2_by_hmode = mode.list_index_L2_by_hmode + + # Compute expectation values — second input to calc_norm_corr_tensor + V1_psi = ht.psi + list_avg_L2 = [ + operator_expectation(mode.list_L2_coo[i], V1_psi) + for i in range(len(mode.list_L2_coo)) + ] + + # Stochastic noise vector — uniform z_scale unless a generator is passed + if z_rnd_func is not None: + z_rnd = z_rnd_func(len(list_l2idx_abs)) + else: + z_rnd = np.ones(len(list_l2idx_abs), dtype=np.complex128) * z_scale + + # z_hat: conjugate noise plus compressed memory — the drive for norm correction + z_mem = np.zeros( + len(tb.noise_memory.list_zmemmodeidx_abs), dtype=np.complex128, + ) + z_hat = np.conj(z_rnd[list_l2idx_abs]) + compress_zmem( + z_mem, list_index_L2_by_hmode, + tb.noise_memory.list_zmemactivemodeidx_rel, + ) + + result = dict( + wavefunction=ht, + psi=V1_psi, + z_hat=z_hat, + list_avg_L2=list_avg_L2, + mode=mode, + list_index_L2_by_mode=mode.list_index_L2_by_hmode, + ) + # Store system for cross-validation but not as a calc_norm_corr_tensor arg + result['_system'] = system + return result + + +def _flat_norm_corr(kwargs): + """Cross-validate calc_norm_corr_tensor against flat calc_norm_corr.""" + ht = kwargs['wavefunction'] + mode = kwargs['mode'] + system = kwargs['_system'] + z_hat = kwargs['z_hat'] + list_avg_L2 = kwargs['list_avg_L2'] + V1_psi = ht.psi + n_state = len(V1_psi) + n_modes_total = len(mode.list_modeidx_abs) + phi_flat = np.zeros(n_state * (1 + n_modes_total), dtype=np.complex128) + phi_flat[:n_state] = V1_psi + dict_mode_to_l2 = dict(_build_mode_l2_map(system, mode)) + list_index_phi_L2_mode_flat = [] + for m_idx in range(n_modes_total): + list_aux_idx_m = [0] * n_modes_total + list_aux_idx_m[m_idx] = 1 + phi_flat[n_state * (1 + m_idx):n_state * (2 + m_idx)] = phi_aux( + ht.list_cores_phi, ht.method, ht.M1_modes_per_state, + list_aux_idx_m, + ) + if m_idx in dict_mode_to_l2: + list_index_phi_L2_mode_flat.append( + (1 + m_idx, dict_mode_to_l2[m_idx], m_idx) + ) + return calc_norm_corr( + phi_flat, z_hat, list_avg_L2, mode.list_L2_coo, + n_state, list_index_phi_L2_mode_flat, mode.list_g, mode.list_w, + ) + + +# ------------------------------------------------------------ +# TEST: Nonzero noise gives nonzero correction +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_calc_norm_corr_tensor_nonzero_noise(): + # Scope: z-component-only regime. With the initial localized state + # (only zeroth auxiliary populated) phi_aux returns zeros, so the + # per-mode correction loop contributes nothing and the result + # collapses to Re(z_hat . list_avg_L2). This test pins the + # z-component behavior; per-mode loop coverage is in + # test_calc_norm_corr_tensor_populated_hierarchy and + # test_calc_norm_corr_tensor_all_modes_active. + kwargs = _make_norm_corr_args(z_scale=1.0) + tensor_kwargs = {k: v for k, v in kwargs.items() if k != '_system'} + result = calc_norm_corr_tensor(**tensor_kwargs) + # Compute expected leading-order term + expected = np.real(np.dot(kwargs['z_hat'], kwargs['list_avg_L2'])) + np.testing.assert_allclose( + result, + expected, + atol=1e-12, + err_msg=( + 'Norm correction should equal Re(z_hat . list_avg_L2) for initial state' + ), + ) + + +# ------------------------------------------------------------ +# TEST: Statenumber representation gives same result as fullstate +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_calc_norm_corr_tensor_statenumber(): + # This case tests that number produces the same + # norm correction as fullstate for identical inputs. + kwargs_full = _make_norm_corr_args( + method='fullstate', + z_scale=0.5, + ) + kwargs_sn = _make_norm_corr_args( + method='number', + z_scale=0.5, + ) + tensor_full = {k: v for k, v in kwargs_full.items() if k != '_system'} + tensor_sn = {k: v for k, v in kwargs_sn.items() if k != '_system'} + result_full = calc_norm_corr_tensor(**tensor_full) + result_sn = calc_norm_corr_tensor(**tensor_sn) + np.testing.assert_allclose( + result_sn, + result_full, + atol=1e-10, + err_msg='statenumber and fullstate gave different norm corrections', + ) + + +# ------------------------------------------------------------ +# TEST: Multi-site initial state (superposition) +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_calc_norm_corr_tensor_superposition(): + # This case tests that a superposition initial state (amplitude on + # multiple sites) exercises the L2 operator application more + # meaningfully than a single-site localized state. + psi_super = np.array([0.5, 0.5, 0.5, 0.5], dtype=np.complex128) + psi_super = psi_super / np.linalg.norm(psi_super) + kwargs = _make_norm_corr_args( + psi=psi_super, + z_rnd_func=lambda n: (0.03 - 0.01j) * np.arange(n, dtype=np.complex128), + ) + tensor_kwargs = {k: v for k, v in kwargs.items() if k != '_system'} + result = calc_norm_corr_tensor(**tensor_kwargs) + assert np.isreal(result) + expected = _flat_norm_corr(kwargs) + np.testing.assert_allclose( + result, expected, atol=1e-12, + err_msg='Superposition: tensor norm corr does not match flat calc_norm_corr', + ) + + +# ------------------------------------------------------------ +# TEST: Populated hierarchy exercises per-mode loop +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_calc_norm_corr_tensor_populated_hierarchy(): + # This case tests the per-mode correction loop by manually populating + # the first-order auxiliary in the MPS. With zeroth-order only, phi_aux + # returns zeros and the loop contributes nothing. Here we set a nonzero + # first-order auxiliary so that - + * + # is actually computed. + # Mode 4 corresponds to site 2 (modes_per_state=2, site 2 → mode 2*2=4). + kwargs = _make_norm_corr_args( + z_rnd_func=lambda n: (0.05 + 0.02j) * np.arange(n, dtype=np.complex128), + aux_populate=[(4, 0.3 + 0.1j)], + ) + tensor_kwargs = {k: v for k, v in kwargs.items() if k != '_system'} + result = calc_norm_corr_tensor(**tensor_kwargs) + assert np.isreal(result) + + # Verify that the per-mode loop actually contributed by comparing against + # the z-component alone (delta = z_hat . list_avg_L2) + z_component_only = np.real(np.dot(kwargs['z_hat'], kwargs['list_avg_L2'])) + assert result != z_component_only, ( + 'Per-mode loop contributed nothing despite populated auxiliary' + ) + + # Cross-validate against the flat-vector calc_norm_corr + expected = _flat_norm_corr(kwargs) + np.testing.assert_allclose( + result, expected, atol=1e-12, + err_msg='calc_norm_corr_tensor does not match flat calc_norm_corr', + ) + + +# ------------------------------------------------------------ +# TEST: Statenumber representation with populated hierarchy +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_calc_norm_corr_tensor_statenumber_populated(): + # This case tests that the per-mode correction loop works correctly + # in statenumber representation by populating a first-order auxiliary + # and comparing against the fullstate result. Both representations + # are built from _make_norm_corr_args with identical inputs, which + # produces the same z_hat deterministically so the cross-check is + # meaningful without rebuilding the fixture by hand. + z_rnd_func = lambda n: (0.05 + 0.02j) * np.arange(n, dtype=np.complex128) + aux_populate = [(4, 0.3 + 0.1j)] + + kwargs_full = _make_norm_corr_args( + method='fullstate', + z_rnd_func=z_rnd_func, + aux_populate=aux_populate, + ) + kwargs_sn = _make_norm_corr_args( + method='number', + z_rnd_func=z_rnd_func, + aux_populate=aux_populate, + ) + tensor_full = {k: v for k, v in kwargs_full.items() if k != '_system'} + tensor_sn = {k: v for k, v in kwargs_sn.items() if k != '_system'} + result_full = calc_norm_corr_tensor(**tensor_full) + result_sn = calc_norm_corr_tensor(**tensor_sn) + + # Per-mode loop must have contributed (otherwise the comparison + # degenerates to the z-component alone and covers nothing new) + z_only = np.real( + np.dot(kwargs_full['z_hat'], kwargs_full['list_avg_L2']) + ) + assert result_full != z_only, ( + 'Fullstate per-mode loop contributed nothing' + ) + np.testing.assert_allclose( + result_full, result_sn, atol=1e-10, + err_msg='Statenumber norm correction diverges from fullstate with ' + 'populated auxiliary', + ) + + +# ------------------------------------------------------------ +# TEST: calc_norm_corr_tensor with all modes active is nonzero +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_calc_norm_corr_tensor_all_modes_active(): + # This case tests the full norm correction with every state-mode + # exercised: one first-order auxiliary is populated per system state + # so the per-mode correction loop contributes across all states + # (not just a single mode), alongside nonzero noise that drives the + # z-component term. The result is cross-validated against the + # flat-vector calc_norm_corr reference at atol=1e-12 — strictly + # stronger than the prior `abs(result) > 1e-6` check, which only + # verified the output wasn't trivially zero. + # Populate one first-order mode auxiliary per system state: with + # modes_per_state = 2 and nsite = 4, the first mode of each state + # sits at flat mode indices 0, 2, 4, 6. + kwargs = _make_norm_corr_args( + z_rnd_func=lambda n: (0.04 + 0.02j) * np.arange(n, dtype=np.complex128), + aux_populate=[ + (0, 0.10 + 0.05j), + (2, 0.15 - 0.03j), + (4, 0.20 + 0.08j), + (6, 0.12 - 0.06j), + ], + ) + tensor_kwargs = {k: v for k, v in kwargs.items() if k != '_system'} + result = calc_norm_corr_tensor(**tensor_kwargs) + assert np.isreal(result) + # Per-mode loop must actually contribute; otherwise the populated + # auxiliaries didn't feed through and the test would be z-only. + z_component_only = np.real( + np.dot(kwargs['z_hat'], kwargs['list_avg_L2']) + ) + assert result != z_component_only, ( + 'Per-mode loop contributed nothing despite populated auxiliaries ' + 'across all states' + ) + # Cross-validate against the flat-vector reference implementation + expected = _flat_norm_corr(kwargs) + np.testing.assert_allclose( + result, expected, atol=1e-12, + err_msg=( + 'all_modes_active: calc_norm_corr_tensor does not match ' + 'flat calc_norm_corr with all state auxiliaries populated' + ), + ) + + +# ------------------------------------------------------------ +# TEST: Sign of correction can be negative +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_calc_norm_corr_tensor_sign(): + # Scope: z-component sign-flip property. The leading-order term + # Re(z_hat . ) is linear in z_scale, so flipping z_scale's sign + # flips the correction sign. This test uses the initial state + # (per-mode loop inactive) because sign-flipping is a property of + # the z-component alone; populated-auxiliary behavior is covered by + # test_calc_norm_corr_tensor_populated_hierarchy and + # test_calc_norm_corr_tensor_all_modes_active. + kwargs_pos = _make_norm_corr_args(z_scale=1.0) + kwargs_neg = _make_norm_corr_args(z_scale=-1.0) + tensor_pos = {k: v for k, v in kwargs_pos.items() if k != '_system'} + tensor_neg = {k: v for k, v in kwargs_neg.items() if k != '_system'} + result_pos = calc_norm_corr_tensor(**tensor_pos) + result_neg = calc_norm_corr_tensor(**tensor_neg) + # With opposite z_scales, the leading z_hat . list_avg_L2 term flips sign + assert result_pos * result_neg < 0, ( + f'Expected opposite signs: z_scale=1.0 gave {result_pos}, ' + f'z_scale=-1.0 gave {result_neg}' + ) + # Analytical: the leading term Re(z_hat . ) flips sign with z_scale, + # so magnitudes should be approximately equal + np.testing.assert_allclose( + abs(result_pos), abs(result_neg), atol=1e-10, + err_msg='Magnitude should be equal for opposite z_scales', + ) + + +# ------------------------------------------------------------ +# TEST: Single-mode system (n_modes = 1) +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_calc_norm_corr_tensor_single_mode(): + # Scope: minimal-system edge case (1 site, 1 bath mode) in the + # z-component-only regime. Per-mode loop is not exercised here + # because the goal is to confirm the function runs on a degenerate + # single-mode fixture without indexing errors; multi-mode per-mode + # loop coverage lives in test_calc_norm_corr_tensor_populated_hierarchy + # and test_calc_norm_corr_tensor_all_modes_active. + H2_hamiltonian_1 = np.array([[0.0]], dtype=np.complex128) + lop_1 = [sp.sparse.coo_matrix(np.array([[1.0]]))] + sys_param_1 = { + 'HAMILTONIAN': H2_hamiltonian_1, + 'GW_SYSBATH': [[g_0, w_0]], + 'L_HIER': lop_1, + 'L_NOISE1': lop_1, + 'ALPHA_NOISE1': bcf_exp, + 'PARAM_NOISE1': [[g_0, w_0]], + } + psi_1 = np.array([1.0], dtype=np.complex128) + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, + } + tb = _make_tb(sp=sys_param_1, ds=0, psi=psi_1, sl=np.array([0])) + ht = HopsTensorWavefunction( + k_max, tensor_param, {'INTEGRATOR': 'RUNGE_KUTTA'}, eom_param, + ) + ht.initialize(psi_1, tb.system) + + mode = tb.mode + system = tb.system + list_L2_coo = mode.list_L2_coo + + V1_psi = ht.psi + list_avg_L2 = [ + operator_expectation(list_L2_coo[i], V1_psi) for i in range(len(list_L2_coo)) + ] + z_hat = np.array([0.5 + 0.3j], dtype=np.complex128) + + result = calc_norm_corr_tensor( + ht, + V1_psi, + z_hat, + list_avg_L2, + mode, + mode.list_index_L2_by_hmode, + ) + assert np.isreal(result) + # With 1 mode and initial MPS (zeroth-order only), the per-mode + # loop contributes nothing, so result should equal Re(z_hat . avg_L2) + expected = np.real(np.dot(z_hat, list_avg_L2)) + np.testing.assert_allclose(result, expected, atol=1e-12) + + +# ============================================================ +# TEST SUITE: apply_system_operator() +# ============================================================ + +# Shared swap matrices +# Swaps sites 0<->2 and 1<->3 +O2_SWAP_02_13 = np.array( + [[0, 0, 1, 0], [0, 0, 0, 1], [1, 0, 0, 0], [0, 1, 0, 0]], + dtype=np.complex128, +) +# Swaps sites 0<->1 and 2<->3 +O2_SWAP_01_23 = np.array( + [[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]], + dtype=np.complex128, +) + + +# ------------------------------------------------------------ +# TEST: Projection operator zeros out components +# ------------------------------------------------------------ +@pytest.mark.level(1) +@pytest.mark.parametrize('method', ['fullstate', 'number']) +def test_apply_system_operator_projection(method): + # This case tests that projecting onto site 0 zeros out the + # entire wavefunction, because psi_0 lives entirely on site 2. + ht = _make_tensor(method) + # |0><0| projector: keeps only site 0 amplitude + O2_proj = np.zeros((nsite, nsite), dtype=np.complex128) + O2_proj[0, 0] = 1.0 + ht.list_cores_phi = apply_system_operator( + ht.list_cores_phi, O2_proj, ht.method, ht.k_max, + ht.M1_modes_per_state, ht.mps_epsilon, ht.bond_dim_max, + ) + # psi_0 had zero weight on site 0, so projection gives zero everywhere + np.testing.assert_allclose(ht.psi, 0.0, atol=1e-10) + + +# ------------------------------------------------------------ +# TEST: Sparse operator is handled correctly +# ------------------------------------------------------------ +@pytest.mark.level(1) +@pytest.mark.parametrize('method', ['fullstate', 'number']) +def test_apply_system_operator_sparse(method): + # This case tests that passing a sparse operator directly gives + # the same result as a dense one. apply_system_operator converts + # sparse to dense internally via .toarray(). + ht_dense = _make_tensor(method) + ht_sparse = _make_tensor(method) + O2_swap_sparse = sp.sparse.coo_matrix(O2_SWAP_01_23) + ht_dense.list_cores_phi = apply_system_operator( + ht_dense.list_cores_phi, O2_SWAP_01_23, ht_dense.method, ht_dense.k_max, + ht_dense.M1_modes_per_state, ht_dense.mps_epsilon, ht_dense.bond_dim_max, + ) + # Pass sparse directly — apply_system_operator handles conversion + ht_sparse.list_cores_phi = apply_system_operator( + ht_sparse.list_cores_phi, O2_swap_sparse, + ht_sparse.method, ht_sparse.k_max, + ht_sparse.M1_modes_per_state, ht_sparse.mps_epsilon, ht_sparse.bond_dim_max, + ) + np.testing.assert_allclose(ht_dense.psi, ht_sparse.psi, atol=1e-12) + + +# ------------------------------------------------------------ +# TEST: Off-diagonal operator swaps populations +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_apply_system_operator_offdiagonal(): + # This case tests that an off-diagonal swap operator correctly + # permutes the state core entries. psi_0 = [0,0,1,0], so after + # swapping states 0<->2 and 1<->3, we expect [1,0,0,0]. + ht = _make_tensor('fullstate') + ht.list_cores_phi = apply_system_operator( + ht.list_cores_phi, O2_SWAP_02_13, ht.method, ht.k_max, + ht.M1_modes_per_state, ht.mps_epsilon, ht.bond_dim_max, + ) + expected = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.complex128) + np.testing.assert_allclose(ht.psi, expected, atol=1e-12) + + +# ------------------------------------------------------------ +# TEST: Statenumber off-diagonal swap matches fullstate +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_apply_system_operator_cross_representation(): + # This case tests that an off-diagonal swap operator gives the same + # result in statenumber and fullstate representations. + ht_full = _make_tensor('fullstate') + ht_snum = _make_tensor('number') + ht_full.list_cores_phi = apply_system_operator( + ht_full.list_cores_phi, O2_SWAP_02_13, ht_full.method, ht_full.k_max, + ht_full.M1_modes_per_state, ht_full.mps_epsilon, ht_full.bond_dim_max, + ) + ht_snum.list_cores_phi = apply_system_operator( + ht_snum.list_cores_phi, O2_SWAP_02_13, ht_snum.method, ht_snum.k_max, + ht_snum.M1_modes_per_state, ht_snum.mps_epsilon, ht_snum.bond_dim_max, + ) + np.testing.assert_allclose(ht_snum.psi, ht_full.psi, atol=1e-10) + + +# ------------------------------------------------------------ +# TEST: Raise operator moves population (both representations) +# ------------------------------------------------------------ +@pytest.mark.level(1) +@pytest.mark.parametrize( + 'method', ['fullstate', 'number'], +) +def test_apply_system_operator_raise(method): + # This case tests a fluorescence-style raise operator |1><0| + # applied to the initial wavefunction. + # Custom psi_0 on site 0 (rather than the module-level default at + # site 2) is needed because |1><0| only acts non-trivially when + # the input has amplitude on site 0 — starting at site 2 would + # give a zero result and not distinguish a working operator from + # a broken one. + # |1><0| is the simplest off-diagonal operator that moves + # population between non-adjacent MPS entries. In statenumber + # representation this forces the operator MPO to route amplitude + # through the daisy-chained transfer-matrix channels between + # state cores 0 and 1, exercising cross-site coupling that + # diagonal operators cannot. + psi_site0 = np.zeros(nsite, dtype=np.complex128) + psi_site0[0] = 1.0 + ht = _make_tensor_with_psi(method, psi_site0) + # Raise operator: |1><0| + O2_raise = np.zeros((nsite, nsite), dtype=np.complex128) + O2_raise[1, 0] = 1.0 + ht.list_cores_phi = apply_system_operator( + ht.list_cores_phi, O2_raise, ht.method, ht.k_max, + ht.M1_modes_per_state, ht.mps_epsilon, ht.bond_dim_max, + ) + # Expected: site 0 population moves to site 1, all other sites + # stay at zero. This specific check distinguishes a working + # raise from silent no-ops (e.g., identity applied in place of + # the actual operator) and from destination-mislocation bugs. + V1_expected = np.zeros(nsite, dtype=np.complex128) + V1_expected[1] = 1.0 + np.testing.assert_allclose(ht.psi, V1_expected, atol=1e-10) + + + +# ------------------------------------------------------------ +# TEST: Operator preserves norm (both representations) +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_apply_system_operator_preserves_norm(): + # This case tests that a unitary operator preserves the norm + # of the wavefunction in both representations. + for method in ['fullstate', 'number']: + ht = _make_tensor(method) + norm_before = np.linalg.norm(ht.psi) + # Permutation matrix (unitary): swap sites 0<->2 and 1<->3 + O2_perm = np.array( + [[0, 0, 1, 0], [0, 0, 0, 1], [1, 0, 0, 0], [0, 1, 0, 0]], + dtype=np.complex128, + ) + ht.list_cores_phi = apply_system_operator( + ht.list_cores_phi, O2_perm, ht.method, ht.k_max, + ht.M1_modes_per_state, ht.mps_epsilon, ht.bond_dim_max, + ) + norm_after = np.linalg.norm(ht.psi) + np.testing.assert_allclose( + norm_after, + norm_before, + rtol=1e-8, + err_msg=f'Norm not preserved for {method}', + ) + + +# ------------------------------------------------------------ +# TEST: Superposition initial state +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_apply_system_operator_superposition(): + # This case tests apply_system_operator on a delocalized + # superposition state, not just a single-site basis state. + # psi = (|0> + |2>) / sqrt(2) + tb = _make_tb() + psi_super = np.array([1, 0, 1, 0], dtype=np.complex128) / np.sqrt(2) + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, + } + ht = HopsTensorWavefunction( + k_max, tensor_param, {'INTEGRATOR': 'RUNGE_KUTTA'}, eom_param, + ) + ht.initialize(psi_super, tb.system) + # Apply |2><2| projector: should keep only the |2> component + O2_proj_2 = np.zeros((nsite, nsite), dtype=np.complex128) + O2_proj_2[2, 2] = 1.0 + ht.list_cores_phi = apply_system_operator( + ht.list_cores_phi, O2_proj_2, ht.method, ht.k_max, + ht.M1_modes_per_state, ht.mps_epsilon, ht.bond_dim_max, + ) + expected = np.array([0, 0, 1, 0], dtype=np.complex128) / np.sqrt(2) + np.testing.assert_allclose( + ht.psi, + expected, + atol=1e-12, + err_msg='Projection on superposition state failed', + ) + + +# ------------------------------------------------------------ +# TEST: Unknown method raises NotImplementedError +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_apply_system_operator_invalid_method(): + # This case tests the NotImplementedError guard for unrecognized + # method strings. Without it, an invalid method would fall through + # the function with undefined behavior rather than raising a clear + # error at the API boundary. + ht = _make_tensor('fullstate') + O2_identity = np.eye(nsite, dtype=np.complex128) + with pytest.raises(NotImplementedError, match='not implemented'): + apply_system_operator( + ht.list_cores_phi, O2_identity, 'invalid_method', + ht.k_max, ht.M1_modes_per_state, + ht.mps_epsilon, ht.bond_dim_max, + ) + + +# ------------------------------------------------------------ +# TEST: Complex-valued operator is applied with correct phase +# ------------------------------------------------------------ +@pytest.mark.level(1) +@pytest.mark.parametrize( + 'method', ['fullstate', 'number'], +) +def test_apply_system_operator_complex_operator(method): + # This case tests that a complex-valued operator (1j * permutation) + # is applied with the correct phase. psi_0 = [0,0,1,0] on site 2; + # swapping states 0<->2 and 1<->3 with a 1j prefactor should give + # [1j, 0, 0, 0]. Catches any silent complex -> real casting bug. + ht = _make_tensor(method) + O2_complex_swap = 1.0j * O2_SWAP_02_13 + ht.list_cores_phi = apply_system_operator( + ht.list_cores_phi, O2_complex_swap, ht.method, ht.k_max, + ht.M1_modes_per_state, ht.mps_epsilon, ht.bond_dim_max, + ) + V1_expected = np.array([1.0j, 0.0, 0.0, 0.0], dtype=np.complex128) + np.testing.assert_allclose( + ht.psi, V1_expected, atol=1e-10, + err_msg=( + f'Complex operator (1j * swap) produced wrong result in ' + f'{method}; phase handling may be broken' + ), + ) + + +# ------------------------------------------------------------ +# TEST: General dense complex operator — fullstate vs statenumber +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_apply_system_operator_general_dense_cross_representation(): + # This case tests that fullstate and statenumber representations + # produce the same result for a general dense complex operator + # (every entry nonzero, not a permutation). This is the strongest + # cross-check for the statenumber MPO path, which must reconstruct + # the full dense action through daisy-chained transfer matrices. + rng = np.random.default_rng(0) + O2_op = ( + rng.standard_normal((nsite, nsite)) + + 1.0j * rng.standard_normal((nsite, nsite)) + ).astype(np.complex128) + + ht_full = _make_tensor('fullstate') + ht_snum = _make_tensor('number') + ht_full.list_cores_phi = apply_system_operator( + ht_full.list_cores_phi, O2_op, ht_full.method, ht_full.k_max, + ht_full.M1_modes_per_state, + ht_full.mps_epsilon, ht_full.bond_dim_max, + ) + ht_snum.list_cores_phi = apply_system_operator( + ht_snum.list_cores_phi, O2_op, ht_snum.method, ht_snum.k_max, + ht_snum.M1_modes_per_state, + ht_snum.mps_epsilon, ht_snum.bond_dim_max, + ) + np.testing.assert_allclose( + ht_snum.psi, ht_full.psi, atol=1e-10, + err_msg=( + 'Statenumber and fullstate diverged on a general dense ' + 'complex operator; MPO cross-site channels may be wrong' + ), + ) + + diff --git a/tests/test_tensor_functions_adaptive.py b/tests/test_tensor_functions_adaptive.py new file mode 100644 index 0000000..4cbd411 --- /dev/null +++ b/tests/test_tensor_functions_adaptive.py @@ -0,0 +1,393 @@ +__title__ = 'Unit Tests for tensor_functions_adaptive' +__author__ = 'N. Covalsen' +__maintainer__ = 'N. Covalsen' + +import numpy as np +import pytest +import scipy as sp + +from mesohops.basis.hops_modes import HopsModes +from mesohops.basis.hops_noise_memory import HopsNoiseMemory +from mesohops.basis.hops_system import HopsSystem +from mesohops.tensor.hops_tensor_wavefunction import HopsTensorWavefunction +from mesohops.tensor.hops_tensor_basis import HopsTensorBasis +from mesohops.tensor.tensor_functions_adaptive import ( + tensor_state_adaptive_check_add_state, + tensor_state_adaptive_check_remove_state, +) +from mesohops.trajectory.exp_noise import bcf_exp +from mesohops.util.bath_corr_functions import bcf_convert_dl_to_exp +from mesohops.util.exceptions import UnsupportedRequest + +# ============================================================ +# Shared Setup +# ============================================================ + +nsite = 4 +e_lambda = 20.0 +gamma = 50.0 +temp = 140.0 +(g_0, w_0) = bcf_convert_dl_to_exp(e_lambda, gamma, temp) + +loperator = np.zeros([4, 4, 4], dtype=np.float64) +gw_sysbath = [] +lop_list = [] +for i in range(nsite): + loperator[i, i, i] = 1.0 + gw_sysbath.append([g_0, w_0]) + lop_list.append(sp.sparse.coo_matrix(loperator[i])) + gw_sysbath.append([-1j * np.imag(g_0), 500.0]) + lop_list.append(loperator[i]) + +# Dimer-of-dimers: 0-1 and 2-3 coupled strongly, 1-2 weakly +hs = np.zeros([nsite, nsite]) +hs[0, 1] = 40 +hs[1, 0] = 40 +hs[1, 2] = 10 +hs[2, 1] = 10 +hs[2, 3] = 40 +hs[3, 2] = 40 +H2_sys_hamiltonian = np.array(hs, dtype=np.complex128) + +sys_param = { + 'HAMILTONIAN': H2_sys_hamiltonian, + 'GW_SYSBATH': gw_sysbath, + 'L_HIER': lop_list, + 'L_NOISE1': lop_list, + 'ALPHA_NOISE1': bcf_exp, + 'PARAM_NOISE1': gw_sysbath, +} + +psi_0 = np.array([0.0, 0.0, 1.0, 0.0], dtype=np.complex128) + +k_max = 4 + +tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': 'fullstate', + 'BOND_DIM_MAX': 20, +} + +integrator_param = {'INTEGRATOR': 'RUNGE_KUTTA'} +eom_param = {'EQUATION_OF_MOTION': 'NORMALIZED NONLINEAR'} + + +def _make_basis_objects(sp=sys_param): + '''Creates HopsSystem, HopsModes, HopsNoiseMemory directly.''' + system = HopsSystem(sp) + mode = HopsModes(system, hierarchy=None) + noise_memory = HopsNoiseMemory(system, mode) + return system, mode, noise_memory + + +def _make_initialized_tensor_basis(delta_s=0, sl=None, psi=None): + '''Creates and initializes a HopsTensorBasis.''' + if sl is None: + sl = list(np.arange(nsite)) + if psi is None: + psi = psi_0 + system, mode, noise_memory = _make_basis_objects() + tb = HopsTensorBasis(system, mode, noise_memory) + system.initialize(delta_s > 0, psi) + mode.list_modeidx_abs = sorted(system.list_statemodeidx_abs) + noise_memory.initialize() + system.state_list = sl + tb.initialize(delta_s) + return tb + + +def _make_wavefunction(state_list, psi_on_states, method='fullstate'): + ''' + Builds an initialized HopsTensorWavefunction for the given state_list + and initial wavefunction. + + Parameters + ---------- + 1. state_list : list(int) + Absolute state indices to include. + 2. psi_on_states : np.ndarray(complex) + Initial wavefunction amplitudes, one per state in state_list. + 3. method : str + Tensor encoding type. + + Returns + ------- + 1. ht : HopsTensorWavefunction + 2. M1_modes_per_state : np.ndarray(int) + Number of bath modes per absolute state index + (length = n_state_full). + ''' + tp = dict(tensor_param, METHOD=method) + # Build psi in full nsite space so system.initialize accepts it + psi_full = np.zeros(nsite, dtype=np.complex128) + for i, s in enumerate(state_list): + psi_full[s] = psi_on_states[i] + tb = _make_initialized_tensor_basis(delta_s=0, sl=state_list, psi=psi_full) + ht = HopsTensorWavefunction(k_max, tp, integrator_param, eom_param) + # HopsTensorWavefunction.initialize does phi_0[system.state_list] internally, + # so we must pass the full-length psi_full here. + ht.initialize( + psi_full, + tb.system, + ) + return ht, ht.M1_modes_per_state + + +# ============================================================ +# TEST SUITE: tensor_state_adaptive_check_add_state() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: fullstate adds Hamiltonian-coupled states +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_check_add_state_fullstate_adds_coupled_states(): + '''check_add_state (fullstate) adds H-coupled states outside basis. + + Setup: 2-state basis [1, 2], psi localized on state 2 (index 1 in basis). + State 2 couples to state 1 (already in basis) and state 3 (outside). + With delta_s small but nonzero, state 3 should be added. + ''' + # Analytical: H has nonzero coupling 2<->3 (hs[2,3]=40) and 1<->2 (hs[1,2]=10). + # With psi on state 2, flux into state 3 should exceed threshold for small delta_s. + sl = [1, 2] + psi_local = np.array([0.0, 1.0], dtype=np.complex128) # psi on state 2 + ht, M1_modes_per_state = _make_wavefunction(sl, psi_local) + + n_state = len(sl) + n_state_full = nsite + + list_new = tensor_state_adaptive_check_add_state( + list_cores_phi=ht.list_cores_phi, + ham=H2_sys_hamiltonian, + old_states=sl, + n_state_full=n_state_full, + n_state=n_state, + delta_s=1e-6, # very small threshold — all significant flux should trigger + state_list=sl, + method='fullstate', + M1_modes_per_state=M1_modes_per_state, + ) + # State 3 couples to state 2 via H[2,3]=40; should be detected + assert 3 in list_new + # States already in basis must not be returned + for s in sl: + assert s not in list_new + + +# ------------------------------------------------------------ +# TEST: number matches fullstate +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_check_add_state_statenumber_matches_fullstate(): + '''statenumber and fullstate return the same new-state set. + + Analytical: both methods estimate the same boundary flux; results + should agree for a localized initial state. + ''' + sl = [1, 2] + psi_local = np.array([0.0, 1.0], dtype=np.complex128) + + ht_full, M1_modes_per_state = _make_wavefunction( + sl, psi_local, method='fullstate' + ) + ht_num, _ = _make_wavefunction( + sl, psi_local, method='number' + ) + + kwargs = dict( + ham=H2_sys_hamiltonian, + old_states=sl, + n_state_full=nsite, + n_state=len(sl), + delta_s=1e-6, + state_list=sl, + M1_modes_per_state=M1_modes_per_state, + ) + list_new_full = tensor_state_adaptive_check_add_state( + list_cores_phi=ht_full.list_cores_phi, + method='fullstate', + **kwargs, + ) + list_new_num = tensor_state_adaptive_check_add_state( + list_cores_phi=ht_num.list_cores_phi, + method='number', + **kwargs, + ) + # Both should agree on which states to add (sets may differ in order) + assert set(list_new_full) == set(list_new_num) + + +# ------------------------------------------------------------ +# TEST: invalid method raises UnsupportedRequest +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_check_add_state_invalid_method_raises(): + '''check_add_state raises UnsupportedRequest for unknown method.''' + sl = [1, 2] + psi_local = np.array([0.0, 1.0], dtype=np.complex128) + ht, M1_modes_per_state = _make_wavefunction(sl, psi_local) + + with pytest.raises(UnsupportedRequest): + tensor_state_adaptive_check_add_state( + list_cores_phi=ht.list_cores_phi, + ham=H2_sys_hamiltonian, + old_states=sl, + n_state_full=nsite, + n_state=len(sl), + delta_s=0.1, + state_list=sl, + method='invalidmethod', + M1_modes_per_state=M1_modes_per_state, + ) + + +# ------------------------------------------------------------ +# TEST: all states present → returns empty list +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_check_add_state_full_basis_returns_empty(): + '''check_add_state returns empty list when all states are in the basis. + + Analytical: V1_flux[state_list] is zeroed out by the function; when + state_list = all states, no state can have positive flux. + ''' + sl = list(np.arange(nsite)) + ht, M1_modes_per_state = _make_wavefunction(sl, psi_0) + + list_new = tensor_state_adaptive_check_add_state( + list_cores_phi=ht.list_cores_phi, + ham=H2_sys_hamiltonian, + old_states=sl, + n_state_full=nsite, + n_state=len(sl), + delta_s=0.1, + state_list=sl, + method='fullstate', + M1_modes_per_state=M1_modes_per_state, + ) + assert list_new == [] + + +# ------------------------------------------------------------ +# TEST: delta_s=0 adds all coupled states outside basis +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_check_add_state_delta_s_zero_adds_all_coupled(): + '''With delta_s=0 threshold is zero, every state with nonzero flux is added. + + Analytical: determine_error_thresh with max_error=0 returns 0; all states + with V1_flux > 0 are added. H applied to psi (on state 2 in basis [1,2]) + gives nonzero flux exactly for states directly coupled to state 2. + H[2,3]=40 so state 3 gets flux; H[2,0]=0 so state 0 gets no flux. + With basis [1,2], states 0 and 3 are outside, but only 3 couples directly. + ''' + sl = [1, 2] + psi_local = np.array([0.0, 1.0], dtype=np.complex128) + ht, M1_modes_per_state = _make_wavefunction(sl, psi_local) + + list_new = tensor_state_adaptive_check_add_state( + list_cores_phi=ht.list_cores_phi, + ham=H2_sys_hamiltonian, + old_states=sl, + n_state_full=nsite, + n_state=len(sl), + delta_s=0.0, + state_list=sl, + method='fullstate', + M1_modes_per_state=M1_modes_per_state, + ) + # With zero threshold, all directly coupled states outside basis appear. + # H[2,3]=40 gives flux to state 3; H[2,0]=0 so state 0 gets none. + assert 3 in list_new + # State 0 is not directly coupled to state 2, so no flux → not added + assert 0 not in list_new + # States in basis must not appear + for s in sl: + assert s not in list_new + + +# ============================================================ +# TEST SUITE: tensor_state_adaptive_check_remove_state() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: fullstate removes low-flux states +# ------------------------------------------------------------ +@pytest.mark.xfail( + reason=( + 'Bug in tensor_functions_adaptive: contract_down_exact returns a ' + 'read-only diag view; V1_error += ... raises ValueError. ' + 'Fix: replace V1_error += with V1_error = V1_error + ...' + ), + strict=True, +) +@pytest.mark.level(1) +def test_check_remove_state_fullstate_removes_decoupled(): + '''check_remove_state (fullstate) identifies low-flux states to remove. + + Setup: full 4-state basis, psi on state 2. With a mock dsystem_dt that + returns zero derivative and large delta_s, states with small coupling + flux are below threshold and should be candidates for removal. + + Known bug: V1_error returned by contract_down_exact is read-only + (numpy diag view), so V1_error += ... raises ValueError. Marked xfail. + ''' + sl = list(np.arange(nsite)) + ht, M1_modes_per_state = _make_wavefunction(sl, psi_0) + + def _mock_dsystem_dt(z_mem, z_rnd, z_rnd2): + return [np.zeros_like(c) for c in ht.list_cores_phi] + + z_rnd = np.zeros(len(sl), dtype=np.complex128) + z_step = [z_rnd, z_rnd.copy(), np.zeros(len(sl), dtype=np.complex128)] + + list_old = tensor_state_adaptive_check_remove_state( + list_cores_phi=ht.list_cores_phi, + ham=H2_sys_hamiltonian, + z_step=z_step, + n_state_full=nsite, + n_state=len(sl), + delta_s=1.0, # large threshold — most states should be removal candidates + state_list=sl, + method='fullstate', + M1_modes_per_state=M1_modes_per_state, + dsystem_dt=_mock_dsystem_dt, + ) + # Analytical: psi is on state 2. States with no or weak coupling to state + # 2 (e.g., state 0 which only couples to state 1) should be removable. + # Result must be relative indices into sl, so values are in [0, n_state). + assert all(0 <= idx < len(sl) for idx in list_old) + assert len(list_old) >= 0 # basic sanity: no exception + + +# ------------------------------------------------------------ +# TEST: invalid method raises UnsupportedRequest +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_check_remove_state_invalid_method_raises(): + '''check_remove_state raises UnsupportedRequest for unknown method.''' + sl = list(np.arange(nsite)) + ht, M1_modes_per_state = _make_wavefunction(sl, psi_0) + + def _mock_dsystem_dt(z_mem, z_rnd, z_rnd2): + return [np.zeros_like(c) for c in ht.list_cores_phi] + + z_rnd = np.zeros(len(sl), dtype=np.complex128) + z_step = [z_rnd, z_rnd.copy(), np.zeros(len(sl), dtype=np.complex128)] + + with pytest.raises(UnsupportedRequest): + tensor_state_adaptive_check_remove_state( + list_cores_phi=ht.list_cores_phi, + ham=H2_sys_hamiltonian, + z_step=z_step, + n_state_full=nsite, + n_state=len(sl), + delta_s=0.1, + state_list=sl, + method='invalidmethod', + M1_modes_per_state=M1_modes_per_state, + dsystem_dt=_mock_dsystem_dt, + ) diff --git a/tests/test_tensor_integrator.py b/tests/test_tensor_integrator.py new file mode 100644 index 0000000..eec79cd --- /dev/null +++ b/tests/test_tensor_integrator.py @@ -0,0 +1,1310 @@ +import numpy as np +import pytest +import scipy as sp + +from mesohops.integrator.tensor_integrator import ( + _build_tdvp_solver_kwargs, + runge_kutta_step_tensor, + runge_kutta_variables, + single_point_variables, + tdvp1_step_tensor, + tdvp2_step_tensor, +) +from mesohops.noise.hops_noise import HopsNoise +from mesohops.tensor.hops_tensor_wavefunction import HopsTensorWavefunction +from mesohops.tensor.hops_tensor_basis import HopsTensorBasis +from mesohops.tensor.hops_tensor_eom import HopsTensorEOM +from mesohops.trajectory.exp_noise import bcf_exp +from mesohops.util.bath_corr_functions import bcf_convert_dl_to_exp +from mesohops.util.physical_constants import hbar + +__title__ = 'Unit Tests for Tensor Integrators' +__author__ = 'N. Covalsen' +__maintainer__ = 'N. Covalsen' + + +# ============================================================ +# Shared Setup: Noise +# ============================================================ +# Mirrors test_integrator_rk.py noise configuration exactly. + +noise_param = { + 'SEED': np.array([np.arange(-10, 10.5, 0.5), -1 * np.arange(-10, 10.5, 0.5)]), + 'MODEL': 'PRE_CALCULATED', + 'TLEN': 10.0, # [fs] + 'TAU': 0.25, # [fs] +} + +noise_param_two = { + 'SEED': np.array( + [np.arange(-10, 10.5, 0.5) / 2, -1 * np.arange(-10, 10.5, 0.5) / 2] + ), + 'MODEL': 'PRE_CALCULATED', + 'TLEN': 10.0, # [fs] + 'TAU': 0.25, # [fs] +} + +T3_loperator_noise = np.zeros([2, 2, 2], dtype=np.float64) +T3_loperator_noise[0, 0, 0] = 1.0 +T3_loperator_noise[1, 1, 1] = 1.0 + +sys_param_noise = { + 'HAMILTONIAN': np.array([[0, 10.0], [10.0, 0]], dtype=np.float64), + 'GW_SYSBATH': [[10.0, 10.0], [5.0, 5.0]], + 'L_HIER': T3_loperator_noise, + 'ALPHA_NOISE1': bcf_exp, + 'PARAM_NOISE1': [[10.0, 10.0], [5.0, 5.0]], + 'L_NOISE1': T3_loperator_noise, +} + +sys_param_noise['NSITE'] = len(sys_param_noise['HAMILTONIAN'][0]) +sys_param_noise['NMODES'] = len(sys_param_noise['GW_SYSBATH'][0]) +sys_param_noise['N_L2'] = 2 +sys_param_noise['L_IND_BY_NMODE1'] = [0, 1] +sys_param_noise['LIND_DICT'] = { + 0: T3_loperator_noise[0, :, :], + 1: T3_loperator_noise[1, :, :], +} + +noise_corr = { + 'CORR_FUNCTION': sys_param_noise['ALPHA_NOISE1'], + 'N_L2': sys_param_noise['N_L2'], + 'LIND_BY_NMODE': sys_param_noise['L_IND_BY_NMODE1'], + 'CORR_PARAM': sys_param_noise['PARAM_NOISE1'], +} + + +# ============================================================ +# Shared Setup: Tensor System (4-site dimer of dimers) +# ============================================================ + +nsite = 4 +e_lambda = 20.0 +gamma = 50.0 +temp = 140.0 +(g_0, w_0) = bcf_convert_dl_to_exp(e_lambda, gamma, temp) + +T3_loperator = np.zeros([4, 4, 4], dtype=np.float64) +list_gw_sysbath = [] +list_lop = [] +for i in range(nsite): + T3_loperator[i, i, i] = 1.0 + list_gw_sysbath.append([g_0, w_0]) + list_lop.append(sp.sparse.coo_matrix(T3_loperator[i])) + list_gw_sysbath.append([-1j * np.imag(g_0), 500.0]) + list_lop.append(T3_loperator[i]) + +H2_hamiltonian = np.zeros([nsite, nsite]) +H2_hamiltonian[0, 1] = 40 +H2_hamiltonian[1, 0] = 40 +H2_hamiltonian[1, 2] = 10 +H2_hamiltonian[2, 1] = 10 +H2_hamiltonian[2, 3] = 40 +H2_hamiltonian[3, 2] = 40 + +sys_param = { + 'HAMILTONIAN': np.array(H2_hamiltonian, dtype=np.complex128), + 'GW_SYSBATH': list_gw_sysbath, + 'L_HIER': list_lop, + 'L_NOISE1': list_lop, + 'ALPHA_NOISE1': bcf_exp, + 'PARAM_NOISE1': list_gw_sysbath, +} + +psi_0 = np.array([0.0] * nsite, dtype=np.complex128) +psi_0[2] = 1.0 + +k_max = 2 +state_list = np.arange(nsite) +delta_s = 0 +eom_param = {'EQUATION_OF_MOTION': 'NORMALIZED NONLINEAR'} + + +def _make_eom(method): + """Returns (HopsTensorWavefunction, HopsTensorEOM, HopsTensorBasis), all initialized.""" + from mesohops.basis.hops_modes import HopsModes + from mesohops.basis.hops_noise_memory import HopsNoiseMemory + from mesohops.basis.hops_system import HopsSystem + + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': method, + 'BOND_DIM_MAX': 20, + } + integrator_param = {'INTEGRATOR': 'RUNGE_KUTTA'} + # Construct shared basis objects directly (same pattern as test_hops_tensor_basis) + system = HopsSystem(sys_param) + mode = HopsModes(system, hierarchy=None) + noise_memory = HopsNoiseMemory(system, mode) + tb = HopsTensorBasis(system, mode, noise_memory) + # Manually initialize shared objects (normally done by trajectory) + system.initialize(delta_s > 0, psi_0) + mode.list_modeidx_abs = sorted(system.list_statemodeidx_abs) + noise_memory.initialize() + system.state_list = state_list + tb.initialize(delta_s) + ht = HopsTensorWavefunction(k_max, tensor_param, integrator_param, eom_param) + ht.initialize(psi_0, tb.system) + eom = HopsTensorEOM(ht, tb.system, tb.mode, tb.noise_memory, tb.adaptive, eom_param) + return ht, eom, tb + + +# ============================================================ +# TEST SUITE: runge_kutta_variables() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: Effective noise integration produces correct averages +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_runge_kutta_variables_effective_noise_integration(): + # This case tests that runge_kutta_variables produces correct averaged + # noise when effective_noise_integration=True, and point-sampled noise + # when False. Mirrors test_integrator_rk.test_effective_noise_integration + # but calls the tensor_integrator copy of runge_kutta_variables. + test_noise = HopsNoise(noise_param, noise_corr) + test_noise2 = HopsNoise(noise_param_two, noise_corr) + + rk_var_control = runge_kutta_variables( + 'b', + 5.0, + test_noise, + test_noise2, + 1.5, + [0, 1], + effective_noise_integration=False, + ) + rk_var_integrated = runge_kutta_variables( + 'b', + 5.0, + test_noise, + test_noise2, + 1.5, + [0, 1], + effective_noise_integration=True, + ) + + # Noise fine-step index arithmetic: the noise arrays are sampled at + # `noise_TAU = 0.25` resolution, so t=5.0 lands at index t / noise_TAU + # = 20, and the span `[t, t + 1.5*tau)` with `tau = 1.5` advances + # 1.5*1.5/0.25 = 9 fine steps to index 29 — but runge_kutta_variables + # grabs an extra 3-step buffer for the half-step averaging, landing + # the slice end at 20 + 12 = 32. Hence `[5*4 : 8*4]`. + known_noise_1 = noise_param['SEED'][:, 5 * 4 : 8 * 4] + known_control_1 = known_noise_1[:, np.array([0, 3, 6])] + known_integrated_1 = np.array( + [ + np.mean(known_noise_1[:, :3], axis=1), + np.mean(known_noise_1[:, 3:6], axis=1), + np.mean(known_noise_1[:, 6:9], axis=1), + ] + ).T + + known_noise_2 = noise_param_two['SEED'][:, 5 * 4 : 8 * 4] + known_control_2 = known_noise_2[:, np.array([0, 3, 6])] + known_integrated_2 = np.array( + [ + np.mean(known_noise_2[:, :3], axis=1), + np.mean(known_noise_2[:, 3:6], axis=1), + np.mean(known_noise_2[:, 6:9], axis=1), + ] + ).T + + assert np.allclose(rk_var_control['z_rnd'], known_control_1) + assert np.allclose(rk_var_control['z_rnd2'], known_control_2) + assert np.allclose(rk_var_integrated['z_rnd'], known_integrated_1) + assert np.allclose(rk_var_integrated['z_rnd2'], known_integrated_2) + # Passthrough values: z_mem ('b') and tau (1.5) are forwarded + # untouched regardless of the effective_noise_integration flag. + assert rk_var_control['z_mem'] == 'b' + assert rk_var_integrated['z_mem'] == 'b' + assert rk_var_control['tau'] == 1.5 + assert rk_var_integrated['tau'] == 1.5 + + +# ============================================================ +# TEST SUITE: runge_kutta_step_tensor() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: RK4 step produces finite output and preserves norm +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_runge_kutta_step_tensor_values(): + # CASE: Run a single RK4 step with a real EOM and verify the + # output z_mem is finite and the wavefunction norm stays O(1). + ht, eom, tb = _make_eom('fullstate') + n_modes = len(tb.mode.list_modeidx_abs) + n_l2 = tb.mode.n_l2 + tau = 1.0 + z_mem = (0.01 + 0.02j) * np.arange(n_modes, dtype=np.complex128) + z_rnd_3pt = (0.03 - 0.01j) * np.arange( + n_l2, + dtype=np.complex128, + ) + z_rnd = np.column_stack( + [z_rnd_3pt, z_rnd_3pt, z_rnd_3pt], + ) + z_rnd2_3pt = (-0.02 + 0.04j) * np.arange( + n_l2, + dtype=np.complex128, + ) + z_rnd2 = np.column_stack( + [z_rnd2_3pt, z_rnd2_3pt, z_rnd2_3pt], + ) + z_mem_rk = runge_kutta_step_tensor( + eom, + z_mem.copy(), + z_rnd, + z_rnd2, + tau, + ) + # z_mem should be finite and changed from the input + assert np.all(np.isfinite(z_mem_rk)), 'RK4 z_mem contains NaN or Inf' + assert not np.allclose(z_mem_rk, z_mem), 'RK4 z_mem unchanged from input' + # Wavefunction norm should be approximately preserved + psi_after = eom.wavefunction.psi + norm_after = np.linalg.norm(psi_after) + assert 0.5 < norm_after < 2.0, ( + f'RK4 step produced extreme norm: {norm_after}' + ) + + +# ------------------------------------------------------------ +# TEST: RK4 step runs with statenumber representation +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_runge_kutta_step_tensor_statenumber(): + # CASE: Run a single RK4 step with statenumber representation. + # The checkpoint copy has a nested-list branch for statenumber + # that is not exercised by fullstate tests. + ht, eom, tb = _make_eom('number') + n_modes = len(tb.mode.list_modeidx_abs) + n_l2 = tb.mode.n_l2 + tau = 1.0 + z_mem = (0.01 + 0.02j) * np.arange(n_modes, dtype=np.complex128) + z_rnd_3pt = (0.03 - 0.01j) * np.arange( + n_l2, + dtype=np.complex128, + ) + z_rnd = np.column_stack( + [z_rnd_3pt, z_rnd_3pt, z_rnd_3pt], + ) + z_rnd2_3pt = (-0.02 + 0.04j) * np.arange( + n_l2, + dtype=np.complex128, + ) + z_rnd2 = np.column_stack( + [z_rnd2_3pt, z_rnd2_3pt, z_rnd2_3pt], + ) + z_mem_rk = runge_kutta_step_tensor( + eom, + z_mem.copy(), + z_rnd, + z_rnd2, + tau, + ) + assert np.all(np.isfinite(z_mem_rk)), 'RK4 z_mem contains NaN or Inf' + assert not np.allclose(z_mem_rk, z_mem), 'RK4 z_mem unchanged from input' + psi_after = eom.wavefunction.psi + norm_after = np.linalg.norm(psi_after) + assert 0.5 < norm_after < 2.0, ( + f'RK4 step produced extreme norm: {norm_after}' + ) + + +# ============================================================ +# TEST SUITE: tdvp1_step_tensor() +# ============================================================ + + +def _make_tb(sp=sys_param, ds=delta_s, psi=psi_0, sl=state_list): + """Creates an initialized HopsTensorBasis from sys_param dict.""" + from mesohops.basis.hops_modes import HopsModes + from mesohops.basis.hops_noise_memory import HopsNoiseMemory + from mesohops.basis.hops_system import HopsSystem + + system = HopsSystem(sp) + mode = HopsModes(system, hierarchy=None) + noise_memory = HopsNoiseMemory(system, mode) + tb = HopsTensorBasis(system, mode, noise_memory) + system.initialize(ds > 0, psi) + mode.list_modeidx_abs = sorted(system.list_statemodeidx_abs) + noise_memory.initialize() + system.state_list = sl + tb.initialize(ds) + return tb + + +def _make_eom_tdvp(method): + """Returns (HopsTensorWavefunction, HopsTensorEOM, HopsTensorBasis) for TDVP1.""" + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': method, + 'BOND_DIM_MAX': 20, + } + integrator_param = {'INTEGRATOR': 'TDVP1'} + tb = _make_tb() + ht = HopsTensorWavefunction(k_max, tensor_param, integrator_param, eom_param) + ht.initialize(psi_0, tb.system) + eom = HopsTensorEOM(ht, tb.system, tb.mode, tb.noise_memory, tb.adaptive, eom_param) + return ht, eom, tb + + +def _make_eom_tdvp2(method): + """Returns (HopsTensorWavefunction, HopsTensorEOM, HopsTensorBasis) for TDVP2.""" + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': method, + 'BOND_DIM_MAX': 20, + } + integrator_param = {'INTEGRATOR': 'TDVP2'} + tb = _make_tb() + ht = HopsTensorWavefunction(k_max, tensor_param, integrator_param, eom_param) + ht.initialize(psi_0, tb.system) + eom = HopsTensorEOM(ht, tb.system, tb.mode, tb.noise_memory, tb.adaptive, eom_param) + return ht, eom, tb + + +# ------------------------------------------------------------ +# TEST: TDVP1 runs without error on fullstate representation +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tdvp1_step_tensor_runs_without_error(): + # This case tests that tdvp1_step_tensor runs to completion + # on the fullstate representation. + ht, eom, tb = _make_eom_tdvp('fullstate') + n_modes = len(tb.mode.list_modeidx_abs) + n_l2 = tb.mode.n_l2 + tau = 0.5 + z_mem = np.zeros(n_modes, dtype=np.complex128) + z_rnd = np.column_stack( + [ + np.zeros(n_l2, dtype=np.complex128), + ] + ) + z_rnd2 = np.column_stack( + [ + np.zeros(n_l2, dtype=np.complex128), + ] + ) + psi_before = eom.wavefunction.psi.copy() + norm_before = np.linalg.norm(psi_before) + z_mem_out = tdvp1_step_tensor(eom, z_mem, z_rnd, z_rnd2, tau) + assert z_mem_out.shape == z_mem.shape + assert not np.any(np.isnan(z_mem_out)) + # Invariant: TDVP approximately preserves norm + norm_after = np.linalg.norm(eom.wavefunction.psi) + np.testing.assert_allclose( + norm_after, norm_before, atol=0.1, + err_msg='TDVP step should approximately preserve norm', + ) + + +# ------------------------------------------------------------ +# TEST: TDVP1 updates phi (MPS cores change after step) +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tdvp1_step_tensor_updates_phi(): + # This case tests that tdvp1_step_tensor actually modifies + # the MPS wavefunction. + ht, eom, tb = _make_eom_tdvp('fullstate') + n_modes = len(tb.mode.list_modeidx_abs) + n_l2 = tb.mode.n_l2 + tau = 0.5 + z_mem = np.zeros(n_modes, dtype=np.complex128) + z_rnd = np.column_stack( + [ + 0.01 * np.arange(n_l2, dtype=np.complex128), + ] + ) + z_rnd2 = np.column_stack( + [ + 0.01 * np.arange(n_l2, dtype=np.complex128), + ] + ) + cores_before = [c.copy() for c in ht.list_cores_phi] + tdvp1_step_tensor(eom, z_mem.copy(), z_rnd, z_rnd2, tau) + cores_after = ht.list_cores_phi + # Shapes may change (TDVP re-canonicalizes), so compare contracted MPS + psi_before = cores_before[0] + for c in cores_before[1:]: + psi_before = np.tensordot(psi_before, c, axes=([-1], [0])) + psi_after = cores_after[0] + for c in cores_after[1:]: + psi_after = np.tensordot(psi_after, c, axes=([-1], [0])) + assert not np.allclose(psi_before, psi_after), ( + 'TDVP1 step did not modify the MPS wavefunction' + ) + + +# ------------------------------------------------------------ +# TEST: TDVP1 on statenumber representation +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tdvp1_step_tensor_statenumber(): + # This case tests that tdvp1_step_tensor works with + # statenumber representation. + ht, eom, tb = _make_eom_tdvp('number') + n_modes = len(tb.mode.list_modeidx_abs) + n_l2 = tb.mode.n_l2 + tau = 0.5 + z_mem = np.zeros(n_modes, dtype=np.complex128) + z_rnd = np.column_stack( + [ + np.zeros(n_l2, dtype=np.complex128), + ] + ) + z_rnd2 = np.column_stack( + [ + np.zeros(n_l2, dtype=np.complex128), + ] + ) + psi_before = eom.wavefunction.psi.copy() + norm_before = np.linalg.norm(psi_before) + z_mem_out = tdvp1_step_tensor(eom, z_mem, z_rnd, z_rnd2, tau) + assert z_mem_out.shape == z_mem.shape + assert not np.any(np.isnan(z_mem_out)) + # Invariant: TDVP approximately preserves norm + norm_after = np.linalg.norm(eom.wavefunction.psi) + np.testing.assert_allclose( + norm_after, norm_before, atol=0.1, + err_msg='TDVP step should approximately preserve norm', + ) + + +# ============================================================ +# TEST SUITE: tdvp2_step_tensor() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: TDVP2 runs without error on fullstate representation +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tdvp2_step_tensor_runs_without_error(): + # This case tests that tdvp2_step_tensor runs to completion + # on the fullstate representation. + ht, eom, tb = _make_eom_tdvp2('fullstate') + n_modes = len(tb.mode.list_modeidx_abs) + n_l2 = tb.mode.n_l2 + tau = 0.5 + z_mem = np.zeros(n_modes, dtype=np.complex128) + z_rnd = np.column_stack( + [ + np.zeros(n_l2, dtype=np.complex128), + ] + ) + z_rnd2 = np.column_stack( + [ + np.zeros(n_l2, dtype=np.complex128), + ] + ) + psi_before = eom.wavefunction.psi.copy() + norm_before = np.linalg.norm(psi_before) + z_mem_out = tdvp2_step_tensor(eom, z_mem, z_rnd, z_rnd2, tau) + assert z_mem_out.shape == z_mem.shape + assert not np.any(np.isnan(z_mem_out)) + # Invariant: TDVP approximately preserves norm + norm_after = np.linalg.norm(eom.wavefunction.psi) + np.testing.assert_allclose( + norm_after, norm_before, atol=0.1, + err_msg='TDVP step should approximately preserve norm', + ) + + +# ------------------------------------------------------------ +# TEST: TDVP2 on statenumber representation +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tdvp2_step_tensor_statenumber(): + # This case tests that tdvp2_step_tensor works with + # statenumber representation (non-uniform physical dims). + ht, eom, tb = _make_eom_tdvp2('number') + n_modes = len(tb.mode.list_modeidx_abs) + n_l2 = tb.mode.n_l2 + tau = 0.5 + z_mem = np.zeros(n_modes, dtype=np.complex128) + z_rnd = np.column_stack( + [ + np.zeros(n_l2, dtype=np.complex128), + ] + ) + z_rnd2 = np.column_stack( + [ + np.zeros(n_l2, dtype=np.complex128), + ] + ) + psi_before = eom.wavefunction.psi.copy() + norm_before = np.linalg.norm(psi_before) + z_mem_out = tdvp2_step_tensor(eom, z_mem, z_rnd, z_rnd2, tau) + assert z_mem_out.shape == z_mem.shape + assert not np.any(np.isnan(z_mem_out)) + # Invariant: TDVP approximately preserves norm + norm_after = np.linalg.norm(eom.wavefunction.psi) + np.testing.assert_allclose( + norm_after, norm_before, atol=0.1, + err_msg='TDVP step should approximately preserve norm', + ) + + +# ------------------------------------------------------------ +# TEST: TDVP1 and TDVP2 z_mem agree (both Euler z_mem update) +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_tdvp1_tdvp2_zmem_agree(): + # This case tests that TDVP1 and TDVP2 produce the same + # z_mem update (both use Euler), verifying the z_mem path + # is identical. + ht1, eom1, tb1 = _make_eom_tdvp('fullstate') + ht2, eom2, tb2 = _make_eom_tdvp2('fullstate') + n_modes = len(tb1.mode.list_modeidx_abs) + n_l2 = tb1.mode.n_l2 + tau = 0.5 + z_mem = np.zeros(n_modes, dtype=np.complex128) + z_rnd = np.column_stack( + [ + 0.01 * np.arange(n_l2, dtype=np.complex128), + ] + ) + z_rnd2 = np.column_stack( + [ + -0.005 * np.arange(n_l2, dtype=np.complex128), + ] + ) + z_mem_1 = tdvp1_step_tensor(eom1, z_mem.copy(), z_rnd, z_rnd2, tau) + z_mem_2 = tdvp2_step_tensor(eom2, z_mem.copy(), z_rnd, z_rnd2, tau) + np.testing.assert_allclose( + z_mem_1, + z_mem_2, + atol=1e-10, + err_msg='TDVP1 and TDVP2 z_mem should agree (both Euler update)', + ) + + +# ============================================================ +# TEST SUITE: runge_kutta_step_tensor() — RK4 formula verification +# ============================================================ + + +def _contract_mps_to_vector(list_cores): + """Contract a list of MPS cores into a single state vector.""" + state = list_cores[0] + for core in list_cores[1:]: + state = np.tensordot(state, core, axes=([-1], [0])) + return state.reshape(-1) + + +class _MockPhiTensor: + """Minimal mock for HopsTensorWavefunction with 2-core bond-dim-1 MPS. + + Uses a trivial scalar first core (value 1) and a second core encoding + the physical state. The bond between the two cores is always rank-1 + after tensor_add compression, so OBC is preserved and the mock + remains analytically tractable. + """ + + __slots__ = ('list_cores_phi', 'mps_epsilon', 'bond_dim_max') + + def __init__(self, n_phys): + # Core 0: trivial scalar 1; core 1: physical state (initially zero) + self.list_cores_phi = [ + np.ones((1, 1, 1), dtype=np.complex128), + np.zeros((1, n_phys, 1), dtype=np.complex128), + ] + self.mps_epsilon = 1e-10 + self.bond_dim_max = 20 + + @property + def flat_cores(self): + return self.list_cores_phi + + def restore_phi(self, cores): + self.list_cores_phi = [c.copy() for c in cores] + + def update_phi_from_flat(self, cores): + self.list_cores_phi = [c.copy() for c in cores] + + +class _MockEOM: + """Mock EOM that returns constant derivatives. + + build_generator always returns dz = V1_dz_const. + derivative always returns a 2-core MPS encoding V1_dphi_const. + This makes the RK4 output analytically predictable. + """ + + __slots__ = ( + 'wavefunction', 'V1_dphi_const', 'V1_dz_const', + # Side-channel attributes the real HopsTensorEOM publishes; + # integrator tests read these to track complexity across stages. + 'last_matvec_complexity', 'max_complexity_step', + ) + + def __init__(self, n_phys, n_zmem): + self.wavefunction = _MockPhiTensor(n_phys) + # Constant derivatives with different complex values to + # break symmetry and catch index/weighting errors + self.V1_dphi_const = (0.1 + 0.2j) * np.arange( + 1, n_phys + 1, dtype=np.complex128 + ) + self.V1_dz_const = (0.3 - 0.1j) * np.arange(1, n_zmem + 1, dtype=np.complex128) + # Match the HopsTensorEOM side-channel contract; tests may + # override `last_matvec_complexity` inside `derivative` to pin + # the RK4 max-tracking behavior. + self.last_matvec_complexity = 0 + self.max_complexity_step = 0 + + def build_generator(self, z_mem, z_rnd, z_rnd2): + return self.V1_dz_const.copy() + + def derivative(self): + # Side-channel: real derivative() sets this; mock sets 0 as a + # sentinel so max-tracking tests can override via subclassing. + self.last_matvec_complexity = 0 + return [ + np.ones((1, 1, 1), dtype=np.complex128), + self.V1_dphi_const.reshape(1, -1, 1).copy(), + ] + + +# ------------------------------------------------------------ +# TEST: RK4 z_mem matches analytical formula with constant dz +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_runge_kutta_step_tensor_zmem_formula(): + # With constant dz at every stage, the RK4 z_mem formula reduces to: + # z_new = z_old + tau/hbar * dz * (1/6 + 2/6 + 2/6 + 1/6) + # = z_old + tau/hbar * dz + # because the weights sum to 1. + n_phys = 3 + n_zmem = 2 + mock_eom = _MockEOM(n_phys, n_zmem) + tau = 2.0 + z_mem_0 = (0.05 + 0.01j) * np.arange(n_zmem, dtype=np.complex128) + # z_rnd shape: (n_modes, 3) — 3 time points for RK4 + z_rnd = np.zeros((n_zmem, 3), dtype=np.complex128) + z_rnd2 = np.zeros((n_zmem, 3), dtype=np.complex128) + + z_mem_new = runge_kutta_step_tensor( + mock_eom, + z_mem_0.copy(), + z_rnd, + z_rnd2, + tau, + ) + # Constant dz → weights (1+2+2+1)/6 = 1 → z_new = z_old + tau/hbar * dz + z_mem_expected = z_mem_0 + tau / hbar * mock_eom.V1_dz_const + np.testing.assert_allclose( + z_mem_new, + z_mem_expected, + atol=1e-12, + err_msg='RK4 z_mem does not match formula for constant dz', + ) + + +# ------------------------------------------------------------ +# TEST: RK4 phi matches analytical formula with constant dphi +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_runge_kutta_step_tensor_phi_formula(): + # With constant dphi at every stage, the RK4 phi formula is: + # phi_new = phi_old + (tau/hbar) * dphi + # * (1/6 + 1/3 + 1/3 + 1/6) + # = phi_old + (tau/hbar) * dphi + # The final phi lives in mock_eom.wavefunction.list_cores_phi. + # For bond-dim-1 single-core MPS, this is exact (no SVD truncation). + n_phys = 3 + n_zmem = 2 + mock_eom = _MockEOM(n_phys, n_zmem) + tau = 2.0 + # Set initial phi to a known state + V1_phi_0 = (1.0 + 0.5j) * np.arange(1, n_phys + 1, dtype=np.complex128) + mock_eom.wavefunction.list_cores_phi = [ + np.ones((1, 1, 1), dtype=np.complex128), + V1_phi_0.reshape(1, -1, 1).copy(), + ] + z_mem_0 = np.zeros(n_zmem, dtype=np.complex128) + z_rnd = np.zeros((n_zmem, 3), dtype=np.complex128) + z_rnd2 = np.zeros((n_zmem, 3), dtype=np.complex128) + + runge_kutta_step_tensor( + mock_eom, + z_mem_0.copy(), + z_rnd, + z_rnd2, + tau, + ) + # Contract the MPS to a state vector + V1_phi_new = _contract_mps_to_vector(mock_eom.wavefunction.list_cores_phi) + # Constant dphi → weights sum to 1 → + # phi_new = phi_old + (tau/hbar) * dphi + V1_phi_expected = V1_phi_0 + tau / hbar * mock_eom.V1_dphi_const + np.testing.assert_allclose( + V1_phi_new, + V1_phi_expected, + atol=1e-10, + err_msg='RK4 phi does not match formula for constant dphi', + ) + + +# ------------------------------------------------------------ +# TEST: RK4 resets phi to original state before each stage +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_runge_kutta_step_tensor_state_reset(): + # This test verifies that each RK4 stage starts from the + # original phi, not from a mutated intermediate. We use a + # mock EOM that records the phi it sees at each build_generator + # call. With constant derivatives, stages 1-3 should see + # phi_0 + c_rk[i] * k[i-1] (scaled), but the key check is + # that after the full step, phi is built from phi_0 + weighted sum + # (not accumulated from intermediates). + n_phys = 2 + n_zmem = 1 + + # Track what phi the EOM sees at each stage + list_phi_seen = [] + + class _RecordingEOM(_MockEOM): + __slots__ = () + + def build_generator(self, z_mem, z_rnd, z_rnd2): + # Record the current phi state by contracting MPS + V1_phi_cur = _contract_mps_to_vector( + self.wavefunction.list_cores_phi, + ) + list_phi_seen.append(V1_phi_cur.copy()) + return self.V1_dz_const.copy() + + mock_eom = _RecordingEOM(n_phys, n_zmem) + V1_phi_0 = np.array([1.0 + 0j, 2.0 + 0j], dtype=np.complex128) + mock_eom.wavefunction.list_cores_phi = [ + np.ones((1, 1, 1), dtype=np.complex128), + V1_phi_0.reshape(1, -1, 1).copy(), + ] + tau = 1.0 + z_mem_0 = np.zeros(n_zmem, dtype=np.complex128) + z_rnd = np.zeros((n_zmem, 3), dtype=np.complex128) + z_rnd2 = np.zeros((n_zmem, 3), dtype=np.complex128) + + runge_kutta_step_tensor( + mock_eom, + z_mem_0.copy(), + z_rnd, + z_rnd2, + tau, + ) + # Stage 0 (i=0): should see phi_0 unchanged + np.testing.assert_allclose( + list_phi_seen[0], + V1_phi_0, + atol=1e-12, + err_msg='Stage 0 did not start from original phi', + ) + # Stages 1-3 (i=1,2,3): should see phi_0 + c_rk[i]*k[i-1]*(tau/hbar) + # not an accumulated state from previous stages + c_rk = [0.0, 0.5, 0.5, 1.0] + for stage in range(1, 4): + V1_phi_expected = ( + V1_phi_0 + c_rk[stage] * tau / hbar * mock_eom.V1_dphi_const + ) + np.testing.assert_allclose( + list_phi_seen[stage], + V1_phi_expected, + atol=1e-10, + err_msg=f'Stage {stage} did not reset to phi_0 + c*k', + ) + + +# ------------------------------------------------------------ +# TEST: RK4 weights are individually correct (state-dependent derivative) +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_runge_kutta_step_tensor_rk4_weights(): + # CASE: With dphi/dt = phi (derivative returns current state), + # the RK4 formula for one step reduces to the 4th-order Taylor + # expansion of exp(a) where a = tau/hbar: + # phi_new = phi_0 * (1 + a + a^2/2 + a^3/6 + a^4/24) + # If any individual RK4 weight (1/6, 1/3, 1/3, 1/6) were wrong, + # the polynomial coefficients would differ. This test catches + # weight errors that the constant-derivative tests miss. + n_phys = 2 + n_zmem = 1 + + class _StateDependentEOM(_MockEOM): + __slots__ = () + + def derivative(self): + self.last_matvec_complexity = 0 + return [c.copy() for c in self.wavefunction.list_cores_phi] + + mock_eom = _StateDependentEOM(n_phys, n_zmem) + V1_phi_0 = np.array([1.0 + 0j, 2.0 + 0j], dtype=np.complex128) + mock_eom.wavefunction.list_cores_phi = [ + np.ones((1, 1, 1), dtype=np.complex128), + V1_phi_0.reshape(1, -1, 1).copy(), + ] + tau = 1.0 + z_mem_0 = np.zeros(n_zmem, dtype=np.complex128) + z_rnd = np.zeros((n_zmem, 3), dtype=np.complex128) + z_rnd2 = np.zeros((n_zmem, 3), dtype=np.complex128) + + runge_kutta_step_tensor( + mock_eom, + z_mem_0.copy(), + z_rnd, + z_rnd2, + tau, + ) + V1_phi_new = _contract_mps_to_vector(mock_eom.wavefunction.list_cores_phi) + # RK4 applied to dphi/dt = phi gives the 4th-order exponential Taylor series + a = tau / hbar + V1_phi_expected = V1_phi_0 * (1 + a + a**2 / 2 + a**3 / 6 + a**4 / 24) + np.testing.assert_allclose( + V1_phi_new, + V1_phi_expected, + atol=1e-10, + err_msg='RK4 weights are incorrect (state-dependent derivative test)', + ) + + +# ------------------------------------------------------------ +# TEST: RK4 z_mem weights are individually correct (z-dependent generator) +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_runge_kutta_step_tensor_rk4_weights_zmem(): + # CASE: With dz/dt = z (build_generator returns current z_mem), + # the RK4 formula for one step reduces to the 4th-order Taylor + # expansion of exp(a) where a = tau/hbar. This pins the individual + # z_mem weights (1, 2, 2, 1)/6 on the z_mem path (source L122-133) + # independently of the phi path — a constant-dz test passes for + # any weight set summing to 1, so it can't detect wrong individual + # weights here. + n_phys = 2 + n_zmem = 2 + + class _ZmemDependentEOM(_MockEOM): + __slots__ = () + + def build_generator(self, z_mem, z_rnd, z_rnd2): + # dz/dt = z — the z-analog of the state-dependent phi test + return z_mem.copy() + + mock_eom = _ZmemDependentEOM(n_phys, n_zmem) + V1_z_mem_0 = np.array([1.0 + 0.5j, -0.25 + 0.75j], dtype=np.complex128) + tau = 1.0 + z_rnd = np.zeros((n_zmem, 3), dtype=np.complex128) + z_rnd2 = np.zeros((n_zmem, 3), dtype=np.complex128) + + z_mem_new = runge_kutta_step_tensor( + mock_eom, + V1_z_mem_0.copy(), + z_rnd, + z_rnd2, + tau, + ) + # RK4 applied to dz/dt = z gives the 4th-order exponential Taylor series + a = tau / hbar + V1_z_mem_expected = V1_z_mem_0 * ( + 1 + a + a**2 / 2 + a**3 / 6 + a**4 / 24 + ) + np.testing.assert_allclose( + z_mem_new, + V1_z_mem_expected, + atol=1e-10, + err_msg='RK4 z_mem weights are incorrect (z-dependent generator test)', + ) + + +# ------------------------------------------------------------ +# TEST: RK4 publishes max complexity across the four substeps +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_runge_kutta_step_tensor_max_complexity_across_substeps(): + # This case tests that runge_kutta_step_tensor publishes the MAX of + # the four per-substep complexities on eom.max_complexity_step, not + # the first/last/sum/etc. With derivative reporting [5, 10, 3, 7] + # across the four RK4 stages, the published max must be 10. The + # constant-complexity mocks elsewhere always report 0, so this test + # is the only one that pins the max-tracking semantics. + n_phys = 2 + n_zmem = 1 + + class _MaxTrackingEOM(_MockEOM): + __slots__ = ('list_complexities', '_call_count') + + def __init__(self, n_phys, n_zmem, list_complexities): + super().__init__(n_phys, n_zmem) + self.list_complexities = list_complexities + self._call_count = 0 + + def derivative(self): + cores = super().derivative() + self.last_matvec_complexity = ( + self.list_complexities[self._call_count] + ) + self._call_count += 1 + return cores + + list_complexities = [5, 10, 3, 7] + mock_eom = _MaxTrackingEOM(n_phys, n_zmem, list_complexities) + tau = 1.0 + z_mem_0 = np.zeros(n_zmem, dtype=np.complex128) + z_rnd = np.zeros((n_zmem, 3), dtype=np.complex128) + z_rnd2 = np.zeros((n_zmem, 3), dtype=np.complex128) + + runge_kutta_step_tensor(mock_eom, z_mem_0, z_rnd, z_rnd2, tau) + assert mock_eom._call_count == 4, ( + 'RK4 should call derivative exactly 4 times' + ) + assert mock_eom.max_complexity_step == max(list_complexities), ( + f'Expected max_complexity_step={max(list_complexities)}, ' + f'got {mock_eom.max_complexity_step}' + ) + + +# ------------------------------------------------------------ +# TEST: TDVP publishes max_complexity_step = 0 +# ------------------------------------------------------------ +@pytest.mark.level(2) +def test_tdvp_step_tensor_max_complexity_zero(): + # This case tests that tdvp_step_tensor unconditionally publishes 0 + # on eom.max_complexity_step. TDVP doesn't call tensor_matvec_prod, + # so the peak-size proxy doesn't apply; the contract is "0 always". + # Stuffing a non-zero value into eom.max_complexity_step before the + # call lets us verify the step function actually overwrites — a no-op + # that simply leaves the previous value untouched would silently + # leak data from the previous step. + _, eom, _ = _make_eom_tdvp(method='fullstate') + eom.max_complexity_step = 999 # sentinel — must be overwritten + n_zmem = len(eom.noise_memory.list_zmemmodeidx_abs) + z_mem = np.zeros(n_zmem, dtype=np.complex128) + z_rnd = np.zeros((n_zmem, 3), dtype=np.complex128) + z_rnd2 = np.zeros((n_zmem, 3), dtype=np.complex128) + eom.build_generator( + z_mem, z_rnd[:, 0], z_rnd2[:, 0], + ) # populate eom.mpo_cores + tdvp1_step_tensor(eom, z_mem, z_rnd, z_rnd2, tau=1.0) + assert eom.max_complexity_step == 0, ( + f'TDVP should publish max_complexity_step=0, ' + f'got {eom.max_complexity_step}' + ) + + +# ------------------------------------------------------------ +# TEST: RK4 phi formula with statenumber representation +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_runge_kutta_step_tensor_phi_formula_statenumber(): + # This case tests that the RK4 phi update works correctly with + # statenumber (nested list-of-lists) MPS, exercising the checkpoint + # deep-copy branch (isinstance(g, list)) and the nested-aware + # tensor_add/scale_mps paths. With constant dphi, the formula is + # the same as fullstate: phi_new = phi_old + (tau/hbar) * dphi. + n_state = 2 + n_zmem = 1 + k_max = 1 # mode core phys dim = k_max + 1 = 2 + + # Statenumber MPS: each state group has [state_core, mode_core] + # State core shape: (1, 2, 1) — binary occupied/unoccupied + # Mode core shape: (1, 2, 1) — k=0,1 occupation + state_0 = np.zeros((1, 2, 1), dtype=np.complex128) + state_0[0, 1, 0] = 1.0 # state 0 occupied + mode_0 = np.zeros((1, 2, 1), dtype=np.complex128) + mode_0[0, 0, 0] = 1.0 # k=0 + + state_1 = np.zeros((1, 2, 1), dtype=np.complex128) + state_1[0, 0, 0] = 1.0 # state 1 unoccupied + mode_1 = np.zeros((1, 2, 1), dtype=np.complex128) + mode_1[0, 0, 0] = 1.0 # k=0 + + list_cores_phi_0 = [[state_0, mode_0], [state_1, mode_1]] + + # Constant derivative MPS with same nested structure + dstate_0 = np.zeros((1, 2, 1), dtype=np.complex128) + dstate_0[0, 1, 0] = 0.1 + 0.2j + dmode_0 = np.zeros((1, 2, 1), dtype=np.complex128) + dmode_0[0, 0, 0] = 1.0 + + dstate_1 = np.zeros((1, 2, 1), dtype=np.complex128) + dstate_1[0, 0, 0] = 0.0 # unoccupied stays zero + dmode_1 = np.zeros((1, 2, 1), dtype=np.complex128) + dmode_1[0, 0, 0] = 1.0 + + list_cores_dphi = [[dstate_0, dmode_0], [dstate_1, dmode_1]] + + class _MockPhiStatenumber: + __slots__ = ('list_cores_phi', 'mps_epsilon', 'bond_dim_max') + + def __init__(self): + self.list_cores_phi = [ + [c.copy() for c in g] for g in list_cores_phi_0 + ] + self.mps_epsilon = 1e-10 + self.bond_dim_max = 20 + + @property + def flat_cores(self): + return [c for g in self.list_cores_phi for c in g] + + def restore_phi(self, cores): + self.list_cores_phi = [ + [c.copy() for c in g] for g in cores + ] + + def update_phi_from_flat(self, cores): + self.list_cores_phi = [ + [c.copy() for c in g] for g in cores + ] + + class _MockEOMStatenumber: + __slots__ = ( + 'wavefunction', 'V1_dz_const', + 'last_matvec_complexity', 'max_complexity_step', + ) + + def __init__(self): + self.wavefunction = _MockPhiStatenumber() + self.V1_dz_const = np.array([0.3 - 0.1j], dtype=np.complex128) + self.last_matvec_complexity = 0 + self.max_complexity_step = 0 + + def build_generator(self, z_mem, z_rnd, z_rnd2): + return self.V1_dz_const.copy() + + def derivative(self): + self.last_matvec_complexity = 0 + return [[c.copy() for c in g] for g in list_cores_dphi] + + mock_eom = _MockEOMStatenumber() + tau = 2.0 + z_mem_0 = np.zeros(n_zmem, dtype=np.complex128) + z_rnd = np.zeros((n_zmem, 3), dtype=np.complex128) + z_rnd2 = np.zeros((n_zmem, 3), dtype=np.complex128) + + # Record psi before + from mesohops.util.tensor_operations import extract_psi + V1_psi_before = extract_psi( + mock_eom.wavefunction.list_cores_phi, + 'number', + np.array([1, 1]), + ) + + runge_kutta_step_tensor( + mock_eom, + z_mem_0.copy(), + z_rnd, + z_rnd2, + tau, + ) + + V1_psi_after = extract_psi( + mock_eom.wavefunction.list_cores_phi, + 'number', + np.array([1, 1]), + ) + + # Constant dphi → weights sum to 1 → + # psi_new = psi_old + (tau/hbar) * dpsi_const + V1_dpsi = extract_psi( + list_cores_dphi, + 'number', + np.array([1, 1]), + ) + V1_psi_expected = V1_psi_before + tau / hbar * V1_dpsi + np.testing.assert_allclose( + V1_psi_after, + V1_psi_expected, + atol=1e-10, + err_msg='RK4 phi formula failed for statenumber representation', + ) + + +# ------------------------------------------------------------ +# TEST: RK4 stages receive correct noise time-point indices +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_runge_kutta_step_tensor_noise_indexing(): + # CASE: Verify that RK4 stages 0-3 receive noise columns + # [0, 1, 1, 2] (corresponding to t, t+tau/2, t+tau/2, t+tau). + n_phys = 2 + n_zmem = 1 + list_z_rnd_seen = [] + list_z_rnd2_seen = [] + + class _NoiseRecordingEOM(_MockEOM): + __slots__ = () + + def build_generator(self, z_mem, z_rnd, z_rnd2): + list_z_rnd_seen.append(z_rnd.copy()) + list_z_rnd2_seen.append(z_rnd2.copy()) + return self.V1_dz_const.copy() + + mock_eom = _NoiseRecordingEOM(n_phys, n_zmem) + tau = 1.0 + z_mem_0 = np.zeros(n_zmem, dtype=np.complex128) + # Distinct values per column to identify which was passed + z_rnd = np.array([[1.0 + 0j, 2.0 + 0j, 3.0 + 0j]]) + z_rnd2 = np.array([[4.0 + 0j, 5.0 + 0j, 6.0 + 0j]]) + + runge_kutta_step_tensor( + mock_eom, + z_mem_0.copy(), + z_rnd, + z_rnd2, + tau, + ) + # Stages should receive noise columns: [0, 1, 1, 2] + list_expected_col = [0, 1, 1, 2] + for i in range(4): + np.testing.assert_array_equal( + list_z_rnd_seen[i], z_rnd[:, list_expected_col[i]], + err_msg=f'Stage {i} received wrong z_rnd noise column', + ) + np.testing.assert_array_equal( + list_z_rnd2_seen[i], z_rnd2[:, list_expected_col[i]], + err_msg=f'Stage {i} received wrong z_rnd2 noise column', + ) + + +# ------------------------------------------------------------ +# TEST: RK4 does not mutate the input z_mem array +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_runge_kutta_step_tensor_zmem_not_mutated(): + # CASE: The input z_mem array should not be modified in-place. + # At i=0 the code aliases z_mem_tmp = z_mem (not a copy), so + # this test guards against in-place mutation by build_generator + # or the final z_mem update using += instead of +. + n_phys = 3 + n_zmem = 2 + mock_eom = _MockEOM(n_phys, n_zmem) + tau = 2.0 + z_mem_0 = (0.05 + 0.01j) * np.arange(n_zmem, dtype=np.complex128) + z_mem_input = z_mem_0.copy() + z_rnd = np.zeros((n_zmem, 3), dtype=np.complex128) + z_rnd2 = np.zeros((n_zmem, 3), dtype=np.complex128) + + runge_kutta_step_tensor( + mock_eom, + z_mem_input, + z_rnd, + z_rnd2, + tau, + ) + np.testing.assert_array_equal( + z_mem_input, z_mem_0, + err_msg='runge_kutta_step_tensor mutated the input z_mem array', + ) + + +# ============================================================ +# TEST SUITE: _build_tdvp_solver_kwargs() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: krylov alias resolves to arnoldi +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_tdvp_solver_kwargs_krylov_alias(): + # The 'krylov' update_type is a convenience alias for 'arnoldi'. + # Verify that the returned dict has solver='arnoldi' and the + # correct conv_tol. + result = _build_tdvp_solver_kwargs('krylov', 1e-4) + assert result['solver'] == 'arnoldi' + assert result['conv_tol'] == 1e-4 + + +# ------------------------------------------------------------ +# TEST: lanczos returns correct kwargs +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_tdvp_solver_kwargs_lanczos(): + # 'lanczos' is passed through unchanged with conv_tol set. + result = _build_tdvp_solver_kwargs('lanczos', 1e-6) + assert result['solver'] == 'lanczos' + assert result['conv_tol'] == 1e-6 + + +# ------------------------------------------------------------ +# TEST: ivp returns method, rtol, atol, and optional max_step +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_build_tdvp_solver_kwargs_ivp(): + # With all ivp kwargs provided, max_step is included. + result = _build_tdvp_solver_kwargs( + 'ivp', 1e-6, ivp_method='RK45', ivp_rtol=1e-5, + ivp_atol=1e-7, ivp_max_step=0.01, + ) + assert result['solver'] == 'ivp' + assert result['method'] == 'RK45' + assert result['rtol'] == 1e-5 + assert result['atol'] == 1e-7 + assert result['max_step'] == 0.01 + # When ivp_max_step is not provided, max_step must be absent. + result_no_max = _build_tdvp_solver_kwargs('ivp', 1e-6) + assert 'max_step' not in result_no_max + + +# ============================================================ +# TEST SUITE: single_point_variables() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: single_point_variables raises NotImplementedError for effective noise +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_single_point_variables_effective_noise_raises(): + # effective_noise_integration is not implemented for single_point_variables; + # verify it raises NotImplementedError. + test_noise = HopsNoise(noise_param, noise_corr) + test_noise2 = HopsNoise(noise_param_two, noise_corr) + with pytest.raises(NotImplementedError): + single_point_variables( + None, + 0.0, + test_noise, + test_noise2, + 0.5, + effective_noise_integration=True, + ) + + +# ------------------------------------------------------------ +# TEST: single_point_variables returns correct dict structure +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_single_point_variables_normal_path(): + # This case tests that single_point_variables returns a dict + # with keys z_mem, z_rnd, z_rnd2, tau and that the noise arrays + # have shape (n_l2, 1) — one time point sampled per L2 operator. + test_noise = HopsNoise(noise_param, noise_corr) + test_noise2 = HopsNoise(noise_param_two, noise_corr) + n_l2 = sys_param_noise['N_L2'] + z_mem_in = np.array([0.1 + 0.2j, -0.3 + 0.0j]) + tau = 0.5 + t = 1.0 + + result = single_point_variables( + z_mem_in, + t, + test_noise, + test_noise2, + tau, + effective_noise_integration=False, + ) + + # Returned dict must have exactly these keys + assert set(result.keys()) == {'z_mem', 'z_rnd', 'z_rnd2', 'tau'} + # z_mem is passed through unchanged + np.testing.assert_array_equal(result['z_mem'], z_mem_in) + # tau is passed through unchanged + assert result['tau'] == tau + # Noise arrays: shape (n_l2, 1) — one time point + assert result['z_rnd'].shape == (n_l2, 1) + assert result['z_rnd2'].shape == (n_l2, 1) diff --git a/tests/test_tensor_nonuniform_modes.py b/tests/test_tensor_nonuniform_modes.py new file mode 100644 index 0000000..e8cf667 --- /dev/null +++ b/tests/test_tensor_nonuniform_modes.py @@ -0,0 +1,138 @@ +import numpy as np +import pytest +import scipy as sp + +from mesohops.basis.hops_modes import HopsModes +from mesohops.basis.hops_noise_memory import HopsNoiseMemory +from mesohops.basis.hops_system import HopsSystem +from mesohops.tensor.hops_tensor_wavefunction import HopsTensorWavefunction +from mesohops.tensor.hops_tensor_basis import HopsTensorBasis +from mesohops.trajectory.exp_noise import bcf_exp +from mesohops.util.bath_corr_functions import bcf_convert_dl_to_exp + +__title__ = 'test_tensor_nonuniform_modes' +__author__ = 'A. Hartzell' +__maintainer__ = 'A. Hartzell' + +# ============================================================ +# Shared Setup: 3-state spectroscopy system +# ============================================================ +# Ground state (index 0) has no bath coupling. +# Excited states (indices 1, 2) each have one L-operator with 2 modes. + +n_site = 2 +n_state = n_site + 1 +e_lambda = 20.0 +gamma = 50.0 +temp = 140.0 +(g_0, w_0) = bcf_convert_dl_to_exp(e_lambda, gamma, temp) + +H2_sys = np.zeros((n_state, n_state), dtype=np.complex128) +H2_sys[1:, 1:] = np.array([[100, -50], [-50, 0]], dtype=np.complex128) + +# L-operators only for excited states (no ground-state bath coupling) +list_lop = [] +list_gw_sysbath = [] +for i in range(n_site): + lop = np.zeros((n_state, n_state), dtype=np.float64) + lop[i + 1, i + 1] = 1.0 + list_lop.append(sp.sparse.coo_matrix(lop)) + list_gw_sysbath.append([g_0, w_0]) + list_lop.append(lop) + list_gw_sysbath.append([-1j * np.imag(g_0), 500.0]) + +sys_param = { + 'HAMILTONIAN': H2_sys, + 'GW_SYSBATH': list_gw_sysbath, + 'L_HIER': list_lop, + 'L_NOISE1': list_lop, + 'ALPHA_NOISE1': bcf_exp, + 'PARAM_NOISE1': list_gw_sysbath, +} + +psi_0 = np.zeros(n_state, dtype=np.complex128) +psi_0[0] = 1.0 + +k_max = 2 +state_list = np.arange(n_state) +delta_s = 0 + + +def _make_spectroscopy_tensor(method): + """Creates and initializes a HopsTensorWavefunction for the spectroscopy system.""" + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': method, + 'BOND_DIM_MAX': 20, + } + eom_param = {'EQUATION_OF_MOTION': 'NORMALIZED NONLINEAR'} + integrator_param = {'INTEGRATOR': 'RUNGE_KUTTA'} + system = HopsSystem(sys_param) + mode = HopsModes(system, hierarchy=None) + noise_memory = HopsNoiseMemory(system, mode) + tb = HopsTensorBasis(system, mode, noise_memory) + system.initialize(delta_s > 0, psi_0) + mode.list_modeidx_abs = sorted(system.list_statemodeidx_abs) + noise_memory.initialize() + system.state_list = state_list + tb.initialize(delta_s) + ht = HopsTensorWavefunction(k_max, tensor_param, integrator_param, eom_param) + ht.initialize(psi_0, tb.system) + return ht + + +# ============================================================ +# TEST SUITE: M1_modes_per_state computation +# ============================================================ + +# ------------------------------------------------------------ +# TEST: non-uniform mode counts computed correctly +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_modes_per_state_nonuniform(): + # This case tests that the ground state gets 0 modes and each excited + # state gets the correct number of modes from LIST_HMODE_INDICES_BY_STATE + ht = _make_spectroscopy_tensor('fullstate') + # Ground state: 0 modes, excited states: 2 modes each (g_0/w_0 pair) + expected = np.array([0, 2, 2]) + np.testing.assert_array_equal(ht.M1_modes_per_state, expected) + + +# ------------------------------------------------------------ +# TEST: mode offset array is cumulative sum +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_mode_offset_nonuniform(): + # This case tests that M1_mode_offset is [0, cumsum(M1_modes_per_state)] + ht = _make_spectroscopy_tensor('fullstate') + expected = np.array([0, 0, 2, 4]) + np.testing.assert_array_equal(ht.M1_mode_offset, expected) + + +# ============================================================ +# TEST SUITE: MPS core layout with non-uniform modes +# ============================================================ + +# ------------------------------------------------------------ +# TEST: fullstate MPS has correct number of cores +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_fullstate_core_count_nonuniform(): + # This case tests that the MPS has 1 state core + total_modes mode cores + ht = _make_spectroscopy_tensor('fullstate') + # 1 state core + 4 mode cores (2 per excited state, 0 for ground) + assert len(ht.list_cores_phi) == 5 + + +# ------------------------------------------------------------ +# TEST: phi_0 extraction works with non-uniform modes +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_phi_0_extraction_nonuniform(): + # This case tests that the physical wavefunction is recovered correctly + ht = _make_spectroscopy_tensor('fullstate') + V1_phi = ht.psi + # Initial state is |g> = [1, 0, 0] + np.testing.assert_allclose(V1_phi[0], 1.0, atol=1e-12) + np.testing.assert_allclose(V1_phi[1], 0.0, atol=1e-12) + np.testing.assert_allclose(V1_phi[2], 0.0, atol=1e-12) diff --git a/tests/test_tensor_operations.py b/tests/test_tensor_operations.py new file mode 100644 index 0000000..3b672bc --- /dev/null +++ b/tests/test_tensor_operations.py @@ -0,0 +1,1385 @@ +import numpy as np +import pytest +import scipy as sp + +from mesohops.basis.hops_modes import HopsModes +from mesohops.basis.hops_noise_memory import HopsNoiseMemory +from mesohops.basis.hops_system import HopsSystem +from mesohops.tensor.hops_tensor_wavefunction import HopsTensorWavefunction +from mesohops.tensor.hops_tensor_basis import HopsTensorBasis +from mesohops.trajectory.exp_noise import bcf_exp +from mesohops.util.bath_corr_functions import bcf_convert_dl_to_exp +from mesohops.basis.basis_functions import determine_error_thresh +from mesohops.storage.storage_functions import save_psi_g_traj_tensor +from mesohops.util.tensor_operations import ( + calc_mps_complexity, + contract_down, + contract_down_exact, + extract_gs_amp, + extract_psi, + flatten_cores, + scale_mps, + tensor_add, + tensor_compress, + tensor_to_array, + unflatten_cores, +) +from mesohops.util.tensor_operations import ( + phi_aux as extract_phi_aux, +) + +__title__ = 'Unit Tests for Tensor Operations' +__author__ = 'N. Covalsen' +__maintainer__ = 'N. Covalsen' + +# ============================================================ +# Shared Setup +# ============================================================ + +nsite = 4 +e_lambda = 20.0 +gamma = 50.0 +temp = 140.0 +(g_0, w_0) = bcf_convert_dl_to_exp(e_lambda, gamma, temp) + +loperator = np.zeros([4, 4, 4], dtype=np.float64) +gw_sysbath = [] +lop_list = [] +for i in range(nsite): + loperator[i, i, i] = 1.0 + gw_sysbath.append([g_0, w_0]) + lop_list.append(sp.sparse.coo_matrix(loperator[i])) + gw_sysbath.append([-1j * np.imag(g_0), 500.0]) + lop_list.append(loperator[i]) + +hs = np.zeros([nsite, nsite]) +hs[0, 1] = 40 +hs[1, 0] = 40 +hs[1, 2] = 10 +hs[2, 1] = 10 +hs[2, 3] = 40 +hs[3, 2] = 40 + +sys_param = { + 'HAMILTONIAN': np.array(hs, dtype=np.complex128), + 'GW_SYSBATH': gw_sysbath, + 'L_HIER': lop_list, + 'L_NOISE1': lop_list, + 'ALPHA_NOISE1': bcf_exp, + 'PARAM_NOISE1': gw_sysbath, +} + +psi_0 = np.array([0.0] * nsite, dtype=np.complex128) +psi_0[2] = 1.0 + +k_max = 2 +state_list = np.arange(nsite) +delta_s = 0 +eom_param = {'EQUATION_OF_MOTION': 'NORMALIZED NONLINEAR'} + + +def _make_tensor(method): + """Creates and initializes a HopsTensorWavefunction for the dimer-of-dimers system.""" + tensor_param = { + 'MPS_EPSILON': 1e-10, + 'METHOD': method, + 'BOND_DIM_MAX': 20, + } + integrator_param = {'INTEGRATOR': 'RUNGE_KUTTA'} + system = HopsSystem(sys_param) + mode = HopsModes(system, hierarchy=None) + noise_memory = HopsNoiseMemory(system, mode) + tb = HopsTensorBasis(system, mode, noise_memory) + system.initialize(delta_s > 0, psi_0) + mode.list_modeidx_abs = sorted(system.list_statemodeidx_abs) + noise_memory.initialize() + system.state_list = state_list + tb.initialize(delta_s) + ht = HopsTensorWavefunction(k_max, tensor_param, integrator_param, eom_param) + ht.initialize(psi_0, tb.system) + return ht + + +def _make_tb(): + """Creates and returns a HopsTensorBasis for the dimer-of-dimers system.""" + system = HopsSystem(sys_param) + mode = HopsModes(system, hierarchy=None) + noise_memory = HopsNoiseMemory(system, mode) + tb = HopsTensorBasis(system, mode, noise_memory) + system.initialize(delta_s > 0, psi_0) + mode.list_modeidx_abs = sorted(system.list_statemodeidx_abs) + noise_memory.initialize() + system.state_list = state_list + tb.initialize(delta_s) + return tb + + +def _make_simple_mps(n_cores, phys_dims, bond_dims): + """Build a simple MPS with deterministic complex values.""" + cores = [] + offset = 0 + for i in range(n_cores): + size = bond_dims[i] * phys_dims[i] * bond_dims[i + 1] + real_part = np.arange(offset, offset + size, dtype=np.float64) * 0.01 + imag_part = np.arange(offset + size, offset + 2 * size, dtype=np.float64) * 0.01 + core = (real_part + 1j * imag_part).reshape( + bond_dims[i], phys_dims[i], bond_dims[i + 1], + ) + cores.append(core) + offset += 2 * size + return cores + + +# ============================================================ +# TEST SUITE: extract_psi() +# ============================================================ + +# ------------------------------------------------------------ +# TEST: extract_psi recovers initial wavefunction +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_extract_psi_recovers_initial_wavefunction(): + # This case tests that extract_psi on the initialized MPS returns + # the original psi_0 = [0, 0, 1, 0] for both representations. + for method in ['fullstate', 'number']: + # Setup: create an initialized HopsTensorWavefunction + ht = _make_tensor(method) + # Action: extract phi_0 from the MPS + result = extract_psi( + ht.list_cores_phi, method, ht.M1_modes_per_state, + ) + # Assertion: result matches psi_0 + np.testing.assert_allclose( + result, psi_0, atol=1e-10, + err_msg=f'phi_0 does not recover psi_0 for {method}', + ) + + +# ------------------------------------------------------------ +# TEST: statenumber extract_psi on a bond-dim-2 superposition +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_extract_psi_statenumber_superposition_bond_dim_2(): + # This case tests that extract_psi correctly recovers a dense + # superposition in statenumber representation, which requires bond + # dim > 1. Also exercises the per-target-state one-hot selection + # with non-trivial coefficients (not zero/one). + a = 0.6 + 0.2j + b = -0.3 + 0.5j + # Bond dim 2 encodes two branches: + # k=0: state 0 occupied, state 1 unoccupied -> amplitude a for |1,0> + # k=1: state 0 unoccupied, state 1 occupied -> amplitude b for |0,1> + state_0 = np.zeros((1, 2, 2), dtype=np.complex128) + state_0[0, 1, 0] = a # branch 0 carries a into phys=1 (occupied) + state_0[0, 0, 1] = b # branch 1 carries b into phys=0 (unoccupied) + state_1 = np.zeros((2, 2, 1), dtype=np.complex128) + state_1[0, 0, 0] = 1.0 # branch 0: state 1 unoccupied + state_1[1, 1, 0] = 1.0 # branch 1: state 1 occupied + list_cores_phi = [[state_0], [state_1]] + M1_modes = np.array([0, 0], dtype=int) + # Extract the physical wavefunction; expect [a, b] + result = extract_psi(list_cores_phi, 'number', M1_modes) + np.testing.assert_allclose( + result, np.array([a, b], dtype=np.complex128), atol=1e-12, + err_msg='Bond-dim-2 superposition amplitudes not recovered', + ) + + +# ------------------------------------------------------------ +# TEST: statenumber extract_psi with non-uniform modes per state +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_extract_psi_statenumber_nonuniform_modes(): + # This case tests that extract_psi correctly traverses varying group + # sizes in the statenumber branch. Builds a 3-site MPS with + # M1_modes_per_state = [1, 3, 2] (non-uniform) and verifies the + # extracted wavefunction on a product state [0, 1, 0]. + # Use k_max_local = 2 so the state core (dim 2) and mode core + # (dim k_max_local + 1 = 3) are structurally distinct. + k_max_local = 2 + n_state = 3 + # Build an unoccupied state core (amp in phys=0) and an occupied + # state core (amp in phys=1), both with trivial bond dim 1. + def _state_core(occupied): + core = np.zeros((1, 2, 1), dtype=np.complex128) + core[0, 1 if occupied else 0, 0] = 1.0 + return core + # Build a mode core in the ground-state slot (phys=0). + def _mode_core(): + core = np.zeros((1, k_max_local + 1, 1), dtype=np.complex128) + core[0, 0, 0] = 1.0 + return core + # Only state 1 is occupied so expected phi = [0, 1, 0] + list_cores_phi = [ + [_state_core(False)] + [_mode_core() for _ in range(1)], + [_state_core(True)] + [_mode_core() for _ in range(3)], + [_state_core(False)] + [_mode_core() for _ in range(2)], + ] + # Sanity: non-uniform group sizes as flagged + assert [len(g) - 1 for g in list_cores_phi] == [1, 3, 2] + M1_modes = np.array([1, 3, 2], dtype=int) + result = extract_psi(list_cores_phi, 'number', M1_modes) + V1_expected = np.zeros(n_state, dtype=np.complex128) + V1_expected[1] = 1.0 + np.testing.assert_allclose( + result, V1_expected, atol=1e-12, + err_msg='Non-uniform modes per state not handled correctly', + ) + + +# ------------------------------------------------------------ +# TEST: statenumber extract_psi on a 1-state system +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_extract_psi_statenumber_single_state(): + # This case tests the degenerate 1-state edge case, where the + # one-hot selection collapses to "always pick occupied for the only + # state". 1 state + 1 mode. + k_max_local = 2 + state_core = np.zeros((1, 2, 1), dtype=np.complex128) + state_core[0, 1, 0] = 1.0 # only state occupied + mode_core = np.zeros((1, k_max_local + 1, 1), dtype=np.complex128) + mode_core[0, 0, 0] = 1.0 # mode in ground state + list_cores_phi = [[state_core, mode_core]] + M1_modes = np.array([1], dtype=int) + result = extract_psi(list_cores_phi, 'number', M1_modes) + np.testing.assert_allclose( + result, np.array([1.0], dtype=np.complex128), atol=1e-12, + err_msg='1-state extract_psi did not return unit amplitude', + ) + + +# ------------------------------------------------------------ +# TEST: extract_psi raises on invalid method +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_extract_psi_invalid_method_raises(): + # This case tests that extract_psi raises ValueError for an unrecognized + # method string, ensuring the dispatcher does not silently return wrong data. + ht = _make_tensor('fullstate') + with pytest.raises(ValueError): + extract_psi(ht.list_cores_phi, 'bogus', ht.M1_modes_per_state) + + +# ============================================================ +# TEST SUITE: extract_gs_amp() +# ============================================================ + +def _number_mps_from_dense(dense, list_modes_per_site, epsilon=1e-12): + """Build a nested number MPS from a dense config tensor via exact TT-SVD. + + The dense axes are the MPS physical legs in order (site_0, mode_00, ..., + site_1, ...); the TT round-trip is lossless, so the MPS represents the + dense state exactly. + """ + size = list(dense.shape) + flat_cores = [] + C = dense + rank = 1 + for i in range(len(size) - 1): + C = np.reshape(C, (int(rank * size[i]), int(C.size / (rank * size[i])))) + U, S, Vt = np.linalg.svd(C, full_matrices=False) + normalized_S = S / np.linalg.norm(S) + thr = determine_error_thresh(np.flip(normalized_S), epsilon * epsilon) + S[normalized_S <= thr] = 0.0 + prev_rank, rank = rank, int(np.count_nonzero(S)) + flat_cores.append( + U[:, :rank].astype(np.complex128).reshape(prev_rank, size[i], rank) + ) + C = np.diag(S[:rank]).astype(np.complex128) @ Vt[:rank, :] + flat_cores.append(C.reshape(C.shape[0], C.shape[1], 1)) + return unflatten_cores(flat_cores, list_modes_per_site) + + +def _fullstate_mps_from_dense(dense, epsilon=1e-12): + """Build a flat fullstate MPS from a dense config tensor via exact TT-SVD.""" + size = list(dense.shape) + flat_cores = [] + C = dense + rank = 1 + for i in range(len(size) - 1): + C = np.reshape(C, (int(rank * size[i]), int(C.size / (rank * size[i])))) + U, S, Vt = np.linalg.svd(C, full_matrices=False) + normalized_S = S / np.linalg.norm(S) + thr = determine_error_thresh(np.flip(normalized_S), epsilon * epsilon) + S[normalized_S <= thr] = 0.0 + prev_rank, rank = rank, int(np.count_nonzero(S)) + flat_cores.append( + U[:, :rank].astype(np.complex128).reshape(prev_rank, size[i], rank) + ) + C = np.diag(S[:rank]).astype(np.complex128) @ Vt[:rank, :] + flat_cores.append(C.reshape(C.shape[0], C.shape[1], 1)) + return flat_cores + + +# ------------------------------------------------------------ +# TEST: extract_gs_amp recovers the all-zeros (vacuum) amplitude +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_extract_gs_amp_recovers_vacuum_amplitude(): + # A random number MPS is built from a dense hierarchy state via TT; the + # vacuum amplitude extract_gs_amp returns must equal the dense state's + # all-zeros entry. k_max_local = 2 gives mode cores dimension 3 vs state + # cores dimension 2 (size [2, 3, 2, 3]), so state and mode cores are + # distinguishable and a core-ordering bug would surface. + rng = np.random.default_rng(0) + k_max_local = 2 + list_modes_per_site = [1, 1] + size = [] + for n_modes in list_modes_per_site: + size += [2] + [k_max_local + 1] * n_modes + dense = rng.standard_normal(size) + 1j * rng.standard_normal(size) + + mps_number = _number_mps_from_dense(dense, list_modes_per_site) + amp_number = extract_gs_amp(mps_number, 'number') + np.testing.assert_allclose(amp_number, dense[(0,) * len(size)], atol=1e-10) + + mps_fullstate = _fullstate_mps_from_dense(dense) + amp_fullstate = extract_gs_amp(mps_fullstate, 'fullstate') + np.testing.assert_allclose(amp_fullstate, dense[(0,) * len(size)], atol=1e-10) + + +# ------------------------------------------------------------ +# TEST: save_psi_g_traj_tensor returns the vacuum amplitude +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_save_psi_g_traj_tensor_returns_vacuum_amplitude(): + # The storage hook delegates to extract_gs_amp; on a TT-built number MPS + # it must return the all-zeros (vacuum) amplitude. + rng = np.random.default_rng(2) + list_modes_per_site = [1, 1] + size = [] + for n_modes in list_modes_per_site: + size += [2] + [2] * n_modes + dense = rng.standard_normal(size) + 1j * rng.standard_normal(size) + + ht = _make_tensor('number') + ht.list_cores_phi = _number_mps_from_dense(dense, list_modes_per_site) + amp = save_psi_g_traj_tensor(ht) + np.testing.assert_allclose(amp, dense[(0,) * len(size)], atol=1e-10) + + +# ============================================================ +# TEST SUITE: contract_down_exact() +# ============================================================ + +# ------------------------------------------------------------ +# TEST: Per-state norms match |psi_0[i]|^2 for ground-state MPS +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_contract_down_exact_per_state(): + # This case tests that for the initialized MPS (all hierarchy in ground + # state), contract_down_exact returns |psi_0[i]|^2 for each state. + for method in ['fullstate', 'number']: + ht = _make_tensor(method) + n_state = nsite + result = contract_down_exact( + ht.list_cores_phi, method, n_state, + ) + expected = np.abs(psi_0) ** 2 + np.testing.assert_allclose( + result.real, expected, atol=1e-12, + err_msg=f'Per-state norm mismatch for {method}', + ) + # Imag part must be negligible — a conjugation error in the + # double-layer contraction could match real but drift imag. + np.testing.assert_allclose( + result.imag, 0.0, atol=1e-12, + err_msg=f'Per-state norm imag part nonzero for {method}', + ) + + +# ------------------------------------------------------------ +# TEST: Statenumber bond-dim-2 superposition — per-state norms match |a|^2, |b|^2 +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_contract_down_exact_superposition_bond_dim_2(): + # This case tests contract_down_exact on a dense 2-site superposition + # a|s0> + b|s1> encoded with bond dim 2 in statenumber representation. + # Exercises both the einsum double-layer contraction with non-trivial + # matrices and the per-target-state one-hot selection at nonzero + # amplitudes (not just zero/one), closing the bond-dim-1 limitation of + # test_contract_down_exact_per_state. + a = 0.6 + 0.2j + b = -0.3 + 0.5j + state_0 = np.zeros((1, 2, 2), dtype=np.complex128) + state_0[0, 1, 0] = a + state_0[0, 0, 1] = b + state_1 = np.zeros((2, 2, 1), dtype=np.complex128) + state_1[0, 0, 0] = 1.0 + state_1[1, 1, 0] = 1.0 + list_cores_phi = [[state_0], [state_1]] + result = contract_down_exact( + list_cores_phi, 'number', 2, + ) + V1_expected = np.array( + [np.abs(a) ** 2, np.abs(b) ** 2], dtype=np.complex128, + ) + np.testing.assert_allclose( + result.real, V1_expected.real, atol=1e-12, + err_msg='Superposition per-state norms not recovered', + ) + np.testing.assert_allclose( + result.imag, 0.0, atol=1e-12, + err_msg='Superposition per-state norm imag part nonzero', + ) + + +# ------------------------------------------------------------ +# TEST: Fullstate with bond dim > 1 still returns |psi[i]|^2 +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_contract_down_exact_fullstate_bond_dim_4(): + # This case tests that contract_down_exact still returns |psi[i]|^2 + # when the MPS is inflated to a non-trivial bond dimension. Using + # eps=0 makes the inflation exactly lossless (no noise added to pad + # out new bond channels), so the contracted per-state norms must + # match the ground-state values to machine precision. Exercises the + # einsum double-layer contraction on fullstate cores with chi > 1. + ht = _make_tensor('fullstate') + ht.inflate_bonds_to(4, eps=0.0) + result = contract_down_exact( + ht.list_cores_phi, 'fullstate', nsite, + ) + expected = np.abs(psi_0) ** 2 + np.testing.assert_allclose( + result.real, expected, atol=1e-12, + err_msg='Inflated fullstate did not recover |psi[i]|^2', + ) + np.testing.assert_allclose( + result.imag, 0.0, atol=1e-12, + err_msg='Inflated fullstate imag part nonzero', + ) + + +# ------------------------------------------------------------ +# TEST: Statenumber raises ValueError on too-few cores +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_contract_down_exact_too_few_cores(): + # This case tests that contract_down_exact raises a ValueError when + # given fewer groups than n_state for number. + # 1 group (list-of-lists) for n_state=2 is too few. + short_cores = [[np.zeros((1, 2, 1), dtype=np.complex128), + np.zeros((1, 3, 1), dtype=np.complex128), + np.zeros((1, 3, 1), dtype=np.complex128)]] + with pytest.raises(ValueError, match='expects at least'): + contract_down_exact( + short_cores, 'number', 2, + ) + + +# ------------------------------------------------------------ +# TEST: Invalid method raises ValueError +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_contract_down_exact_invalid_method(): + # This case tests that contract_down_exact raises a ValueError when + # given an unrecognized method string. + ht = _make_tensor('fullstate') + with pytest.raises(ValueError, match='Unknown method'): + contract_down_exact( + ht.list_cores_phi, 'badmethod', nsite, + ) + + +# ------------------------------------------------------------ +# TEST: Statenumber raises ValueError on bond mismatch +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_contract_down_exact_bond_mismatch(): + # This case tests that contract_down_exact raises a ValueError when + # adjacent group bond dimensions are incompatible. + # Group 0 has right bond 3; group 1 has left bond 1 — mismatch. + bad_cores = [ + [np.zeros((1, 2, 3), dtype=np.complex128), + np.zeros((3, 3, 3), dtype=np.complex128), + np.zeros((3, 3, 3), dtype=np.complex128)], + [np.zeros((1, 2, 1), dtype=np.complex128), + np.zeros((1, 3, 1), dtype=np.complex128), + np.zeros((1, 3, 1), dtype=np.complex128)], + ] + with pytest.raises(ValueError, match='Environment/bond mismatch'): + contract_down_exact(bad_cores, 'number', 2) + + +# ============================================================ +# TEST SUITE: contract_down() +# ============================================================ + +# ------------------------------------------------------------ +# TEST: Approximate contraction agrees with exact for bond-dim-1 MPS +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_contract_down_agrees_with_exact_bonddim1(): + # This case tests that for a bond-dim-1 MPS (no entanglement between + # cores), the approximate contraction matches the exact contraction. + for method in ['fullstate', 'number']: + ht = _make_tensor(method) + n_state = nsite + approx = contract_down( + ht.list_cores_phi, method, n_state, + ) + exact = contract_down_exact( + ht.list_cores_phi, method, n_state, + ) + np.testing.assert_allclose( + approx.real, exact.real, atol=1e-10, + err_msg=f'Approximate vs exact mismatch for {method}', + ) + + # This case tests with bond dim > 1 (entangled MPS) where the + # approximation may differ but invariants still hold + ht_inflated = _make_tensor('fullstate') + ht_inflated.inflate_bonds_to(4, eps=0.1) + approx_infl = contract_down( + ht_inflated.list_cores_phi, 'fullstate', + nsite, + ) + # Invariant: per-state norms should be non-negative + assert np.all(approx_infl.real >= -1e-10), ( + 'Per-state norms should be non-negative' + ) + # Invariant: sum should approximate total norm squared + psi_infl = extract_psi( + ht_inflated.list_cores_phi, 'fullstate', + ht_inflated.M1_modes_per_state, + ) + expected_norm_sq = np.sum(np.abs(psi_infl) ** 2) + np.testing.assert_allclose( + np.sum(approx_infl).real, expected_norm_sq.real, atol=1e-4, + err_msg='Sum of contract_down should approximate norm squared', + ) + + +# ------------------------------------------------------------ +# TEST: Both representations agree +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_contract_down_both_representations_agree(): + # This case tests that contract_down() returns the same per-state + # values for both MPS representations. + # Setup: create tensors with both methods + ht_full = _make_tensor('fullstate') + ht_state = _make_tensor('number') + # Action: compute per-state norms from each + result_full = contract_down( + ht_full.list_cores_phi, 'fullstate', + nsite, + ) + result_state = contract_down( + ht_state.list_cores_phi, 'number', + nsite, + ) + # Assertion: both representations give the same result + np.testing.assert_allclose( + result_full, result_state, atol=1e-10, + err_msg='contract_down differs between fullstate and statenumber', + ) + # Imag part must be negligible in both representations + np.testing.assert_allclose( + result_full.imag, 0.0, atol=1e-12, + err_msg='Fullstate contract_down imag part nonzero', + ) + np.testing.assert_allclose( + result_state.imag, 0.0, atol=1e-12, + err_msg='Statenumber contract_down imag part nonzero', + ) + + # This case tests invariant for inflated (entangled) MPS + ht_infl = _make_tensor('fullstate') + ht_infl.inflate_bonds_to(4, eps=0.1) + result_infl = contract_down( + ht_infl.list_cores_phi, 'fullstate', + nsite, + ) + # Invariant: sum of per-state norms equals norm squared + psi_infl = extract_psi( + ht_infl.list_cores_phi, 'fullstate', + ht_infl.M1_modes_per_state, + ) + np.testing.assert_allclose( + np.sum(result_infl).real, + np.sum(np.abs(psi_infl) ** 2).real, + atol=1e-4, + ) + + +# ------------------------------------------------------------ +# TEST: Sum equals norm squared +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_contract_down_sum_equals_norm_squared(): + # This case tests that the sum of per-state values from contract_down() + # equals the total MPS norm squared. + for method in ['fullstate', 'number']: + # Setup: create an initialized HopsTensorWavefunction + ht = _make_tensor(method) + # Action: compute per-state norms + result = contract_down( + ht.list_cores_phi, method, nsite, + ) + # Compute expected norm squared from phi_0 (for bond-dim-1 MPS, + # only the zeroth auxiliary contributes) + phi = extract_psi( + ht.list_cores_phi, method, ht.M1_modes_per_state, + ) + expected_norm_sq = np.sum(np.abs(phi) ** 2) + # Assertion: sum of per-state values equals norm squared + np.testing.assert_allclose( + np.sum(result).real, expected_norm_sq.real, atol=1e-10, + err_msg=f'Sum of contract_down != norm squared for {method}', + ) + # Imag part must be negligible — a conjugation error would match + # the real sum but drift imag across per-state values. + np.testing.assert_allclose( + result.imag, 0.0, atol=1e-12, + err_msg=f'Per-state imag part nonzero for {method}', + ) + + # This case tests the invariant for inflated (bond dim > 1) MPS + ht_infl = _make_tensor('fullstate') + ht_infl.inflate_bonds_to(4, eps=0.1) + result_infl = contract_down( + ht_infl.list_cores_phi, 'fullstate', + nsite, + ) + phi_infl = extract_psi( + ht_infl.list_cores_phi, 'fullstate', + ht_infl.M1_modes_per_state, + ) + expected_infl = np.sum(np.abs(phi_infl) ** 2) + np.testing.assert_allclose( + np.sum(result_infl).real, expected_infl.real, atol=1e-4, + err_msg='Sum of contract_down != norm squared for inflated MPS', + ) + + +# ------------------------------------------------------------ +# TEST: Invalid method raises ValueError +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_contract_down_invalid_method_raises(): + # This case tests that contract_down raises ValueError for an + # unrecognized method string, matching the parallel test on + # contract_down_exact. + ht = _make_tensor('fullstate') + with pytest.raises(ValueError, match='Unknown method'): + contract_down(ht.list_cores_phi, 'bogus', nsite) + + +# ------------------------------------------------------------ +# TEST: Statenumber bond-dim-2 superposition per-state norms +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_contract_down_superposition_statenumber_bond_dim_2(): + # This case tests contract_down on a dense 2-site superposition + # a|s0> + b|s1> encoded with bond dim 2 in statenumber representation. + # Covers both "statenumber with bond dim > 1" and "superposition + # state" in one test — a dense statenumber superposition naturally + # requires bond > 1, so the inline traversal + per-target + # contraction is exercised on non-trivial bond matrices. + a = 0.6 + 0.2j + b = -0.3 + 0.5j + state_0 = np.zeros((1, 2, 2), dtype=np.complex128) + state_0[0, 1, 0] = a + state_0[0, 0, 1] = b + state_1 = np.zeros((2, 2, 1), dtype=np.complex128) + state_1[0, 0, 0] = 1.0 + state_1[1, 1, 0] = 1.0 + list_cores_phi = [[state_0], [state_1]] + result = contract_down( + list_cores_phi, 'number', 2, + ) + V1_expected = np.array( + [np.abs(a) ** 2, np.abs(b) ** 2], dtype=np.complex128, + ) + np.testing.assert_allclose( + result.real, V1_expected.real, atol=1e-12, + err_msg='Superposition per-state norms not recovered by contract_down', + ) + np.testing.assert_allclose( + result.imag, 0.0, atol=1e-12, + err_msg='Superposition contract_down imag part nonzero', + ) + + +# ============================================================ +# TEST SUITE: phi_aux() +# ============================================================ + +# ------------------------------------------------------------ +# TEST: First-order excitation is zero for initialized MPS +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_phi_aux_first_order_zero_at_init(): + # This case tests that all first-order auxiliary wavefunctions + # are zero for the initial MPS (all modes start in |0>). + for method in ['fullstate', 'number']: + ht = _make_tensor(method) + n_modes = len(gw_sysbath) + for i_mode in range(n_modes): + indices = [0] * n_modes + indices[i_mode] = 1 + phi_1 = extract_phi_aux( + ht.list_cores_phi, method, ht.M1_modes_per_state, indices, + ) + np.testing.assert_allclose( + phi_1, 0.0, atol=1e-12, + err_msg=f'Mode {i_mode} first-order not zero for {method}', + ) + + +# ------------------------------------------------------------ +# TEST: phi_aux recovers nonzero auxiliary after manual injection +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_phi_aux_nonzero_injection(): + # This case tests that phi_aux returns the exact analytically + # computable auxiliary amplitude after a known injection. A + # function that returned random nonzero garbage would pass the + # prior "any nonzero" check, so we assert the exact value. + # For a bond-dim-1 MPS with psi_0 = [0,0,1,0] and an injection of + # z = 0.5+0.1j into mode 0's occupation-1 slice, phi_aux at + # indices=[1,0,0,...] contracts to z * psi_0 = [0, 0, z, 0]. + z_inject = 0.5 + 0.1j + for method in ['fullstate', 'number']: + ht = _make_tensor(method) + n_modes = int(np.sum(ht.M1_modes_per_state)) + # Inject z_inject into mode 0's occupation-1 slice. + # For fullstate: list_cores_phi[1][:, 1, :] is the occ-1 slice. + # For statenumber: list_cores_phi[0][1][:, 1, :] is the first mode. + if method == 'fullstate': + ht.list_cores_phi[1][:, 1, :] = z_inject + else: + ht.list_cores_phi[0][1][:, 1, :] = z_inject + indices = [0] * n_modes + indices[0] = 1 + V1_phi_1 = extract_phi_aux( + ht.list_cores_phi, method, ht.M1_modes_per_state, indices, + ) + V1_expected = z_inject * psi_0 + np.testing.assert_allclose( + V1_phi_1, V1_expected, atol=1e-12, + err_msg=f'phi_aux injection value wrong for {method}', + ) + + +# ------------------------------------------------------------ +# TEST: phi_aux with all-zero indices equals extract_psi +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_phi_aux_zero_indices_equals_extract_psi(): + # This case pins the docstring guarantee that phi_aux with + # all-zero indices returns phi_0 — identical to extract_psi. + for method in ['fullstate', 'number']: + ht = _make_tensor(method) + n_modes = int(np.sum(ht.M1_modes_per_state)) + V1_aux_zero = extract_phi_aux( + ht.list_cores_phi, method, ht.M1_modes_per_state, + [0] * n_modes, + ) + V1_phi_0 = extract_psi( + ht.list_cores_phi, method, ht.M1_modes_per_state, + ) + np.testing.assert_allclose( + V1_aux_zero, V1_phi_0, atol=1e-12, + err_msg=f'phi_aux([0]*n_modes) != extract_psi for {method}', + ) + + +# ------------------------------------------------------------ +# TEST: phi_aux with higher-order and multi-mode indices +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_phi_aux_higher_order_and_multi_mode(): + # This case tests phi_aux on auxiliary indices beyond the common + # [0,...,0,1,0,...] pattern: a single mode at occupation 2 + # (second-order) and two modes simultaneously at occupation 1 + # (multi-mode). For a bond-dim-1 MPS with psi_0 = [0,0,1,0], each + # injection contributes multiplicatively, so the expected aux is + # (product of injections) * psi_0. + z_0 = 0.4 - 0.2j # injected into mode 0 at occ 2 (second-order) + z_1 = 0.3 + 0.1j # injected into mode 1 at occ 1 + for method in ['fullstate', 'number']: + ht = _make_tensor(method) + n_modes = int(np.sum(ht.M1_modes_per_state)) + if method == 'fullstate': + ht.list_cores_phi[1][:, 2, :] = z_0 + ht.list_cores_phi[2][:, 1, :] = z_1 + else: + ht.list_cores_phi[0][1][:, 2, :] = z_0 + ht.list_cores_phi[0][2][:, 1, :] = z_1 + # Sub-case A: higher-order on a single mode + indices_second_order = [0] * n_modes + indices_second_order[0] = 2 + V1_aux_A = extract_phi_aux( + ht.list_cores_phi, method, ht.M1_modes_per_state, + indices_second_order, + ) + np.testing.assert_allclose( + V1_aux_A, z_0 * psi_0, atol=1e-12, + err_msg=f'phi_aux second-order wrong for {method}', + ) + # Sub-case B: two modes at occupation 1 simultaneously + indices_multi = [0] * n_modes + indices_multi[0] = 1 # mode 0 — but we injected at occ 2, not 1 + indices_multi[1] = 1 + # Mode 0's occ-1 slice is still zero (we only populated occ-2), + # so the product includes a zero factor and the aux vanishes. + V1_aux_B = extract_phi_aux( + ht.list_cores_phi, method, ht.M1_modes_per_state, + indices_multi, + ) + np.testing.assert_allclose( + V1_aux_B, 0.0, atol=1e-12, + err_msg=( + f'phi_aux multi-mode with one zero factor should be ' + f'zero for {method}' + ), + ) + + +# ------------------------------------------------------------ +# TEST: Statenumber phi_aux zero-pads short indices +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_phi_aux_statenumber_padding(): + # This case tests the statenumber-path padding logic: when the + # caller passes an `indices` list shorter than the total mode + # count, the function zero-pads via the n_total_padded check + # rather than crashing. Short all-zero indices must yield the + # same result as [0]*n_total (which equals extract_psi by the + # docstring guarantee). + ht = _make_tensor('number') + n_total = int(np.sum(ht.M1_modes_per_state)) + # Pass half the indices — forces the padding branch to fire. + indices_short = [0] * (n_total // 2) + V1_aux_short = extract_phi_aux( + ht.list_cores_phi, 'number', + ht.M1_modes_per_state, indices_short, + ) + V1_phi_0 = extract_psi( + ht.list_cores_phi, 'number', + ht.M1_modes_per_state, + ) + np.testing.assert_allclose( + V1_aux_short, V1_phi_0, atol=1e-12, + err_msg='Short indices should zero-pad to the full mode count', + ) + + +# ============================================================ +# TEST SUITE: tensor_add() +# ============================================================ + +# ------------------------------------------------------------ +# TEST: Mismatched core counts raise ValueError +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tensor_add_length_mismatch(): + # This case tests that tensor_add raises ValueError for mismatched + # core counts. + cores_a = [np.zeros((1, 2, 1), dtype=np.complex128)] + cores_b = [np.zeros((1, 2, 1), dtype=np.complex128)] * 2 + with pytest.raises(ValueError, match='core count mismatch'): + tensor_add(cores_a, cores_b, 1e-10, 20) + + +# ------------------------------------------------------------ +# TEST: Physical dimension mismatch raises ValueError (Jacob T6) +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tensor_add_phys_dim_mismatch(): + # This case tests that tensor_add raises ValueError when + # corresponding cores have different physical dimensions. + cores_a = [np.zeros((1, 2, 1), dtype=np.complex128)] + cores_b = [np.zeros((1, 3, 1), dtype=np.complex128)] + with pytest.raises(ValueError, match='Physical dimension mismatch'): + tensor_add(cores_a, cores_b, 1e-10, 20) + + +# ------------------------------------------------------------ +# TEST: Adding known MPS gives correct sum (Jacob T5) +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tensor_add_known_sum(): + # This case tests tensor_add with two initialized HopsTensorWavefunction MPS + # and verifies that the physical wavefunction sums correctly. + # Uses the shared dimer system: psi_0 = [0, 0, 1, 0]. + ht = _make_tensor('fullstate') + V1_phi_before = ht.psi.copy() + # Scale the second MPS by 0.5+0.3j to break symmetry + scale = 0.5 + 0.3j + list_scaled = [c.copy() for c in ht.list_cores_phi] + list_scaled[0] = list_scaled[0] * scale + list_cores_sum = tensor_add( + ht.list_cores_phi, list_scaled, ht.mps_epsilon, ht.bond_dim_max, + ) + # Extract phi_0 from the summed MPS + V1_phi_sum = extract_psi( + list_cores_sum, 'fullstate', ht.M1_modes_per_state, + ) + expected = V1_phi_before * (1.0 + scale) + np.testing.assert_allclose( + V1_phi_sum, expected, atol=1e-10, + err_msg='tensor_add known sum does not match expected', + ) + + +# ------------------------------------------------------------ +# TEST: Statenumber (nested) tensor_add gives correct sum +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tensor_add_statenumber_nested(): + # This case tests tensor_add with nested statenumber MPS (list-of-lists). + # The nested code path flattens, adds, compresses, then unflattens. + ht = _make_tensor('number') + V1_phi_before = ht.psi.copy() + scale = 0.7 - 0.2j + list_scaled = [ + [c.copy() for c in group] for group in ht.list_cores_phi + ] + list_scaled[0][0] = list_scaled[0][0] * scale + list_cores_sum = tensor_add( + ht.list_cores_phi, list_scaled, ht.mps_epsilon, ht.bond_dim_max, + ) + # Result should be nested (list-of-lists) + assert isinstance(list_cores_sum[0], list), ( + 'tensor_add should return nested structure for nested input' + ) + V1_phi_sum = extract_psi( + list_cores_sum, 'number', ht.M1_modes_per_state, + ) + expected = V1_phi_before * (1.0 + scale) + np.testing.assert_allclose( + V1_phi_sum, expected, atol=1e-10, + err_msg='tensor_add statenumber sum does not match expected', + ) + + +# ------------------------------------------------------------ +# TEST: Large epsilon does not catastrophically lose information +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tensor_add_large_epsilon(): + # This case tests that tensor_add with a deliberately large epsilon + # still produces an approximate result rather than silently dropping + # most information during SVD truncation. + ht = _make_tensor('fullstate') + V1_psi_before = ht.psi.copy() + other_cores = [c.copy() for c in ht.list_cores_phi] + epsilon = 0.5 + list_cores_sum = tensor_add( + ht.list_cores_phi, other_cores, epsilon, ht.bond_dim_max, + ) + V1_psi_sum = extract_psi( + list_cores_sum, 'fullstate', ht.M1_modes_per_state, + ) + expected = 2.0 * V1_psi_before + np.testing.assert_allclose( + V1_psi_sum, expected, atol=epsilon, + err_msg='tensor_add with large epsilon lost too much information', + ) + + +# ============================================================ +# TEST SUITE: tensor_compress() +# ============================================================ + +# ------------------------------------------------------------ +# TEST: Compression respects bond_dim_max +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tensor_compress_bond_dim_max(): + # This case tests that after compression, no bond dimension + # exceeds bond_dim_max. + cores = _make_simple_mps(4, [2, 3, 3, 2], [1, 8, 8, 8, 1]) + max_bond = 3 + compressed = tensor_compress(cores, 1e-10, max_bond) + for i, core in enumerate(compressed): + assert core.shape[0] <= max_bond, ( + f'Core {i} left bond {core.shape[0]} > {max_bond}' + ) + assert core.shape[2] <= max_bond, ( + f'Core {i} right bond {core.shape[2]} > {max_bond}' + ) + + +# ------------------------------------------------------------ +# TEST: Compression accuracy within epsilon +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tensor_compress_accuracy(): + # This case tests that the compressed MPS approximates the original + # by contracting both to full state vectors and comparing. + cores = _make_simple_mps(3, [2, 2, 2], [1, 4, 4, 1]) + compressed = tensor_compress(cores, 1e-10, 20) + # Contract original + state_orig = cores[0] + for c in cores[1:]: + state_orig = np.tensordot(state_orig, c, axes=([-1], [0])) + # Contract compressed + state_comp = compressed[0] + for c in compressed[1:]: + state_comp = np.tensordot(state_comp, c, axes=([-1], [0])) + np.testing.assert_allclose( + state_comp, state_orig, atol=1e-8, + err_msg='Compressed MPS does not approximate original', + ) + + +# ------------------------------------------------------------ +# TEST: tensor_compress with large epsilon reduces bond dims +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tensor_compress_epsilon_truncation(): + # This case tests that compressing with a large epsilon discards small + # singular values and reduces bond dimensions below the inflated baseline. + ht = _make_tensor('fullstate') + ht.inflate_bonds_to(8, eps=0.01) + cores_inflated = [c.copy() for c in ht.list_cores_phi] + psi_before = extract_psi( + cores_inflated, 'fullstate', ht.M1_modes_per_state, + ) + # Compress with large epsilon but no hard bond cap + cores_compressed = tensor_compress(cores_inflated, 0.5, 999) + # Bound: bonds should be reduced + max_bond_before = max(c.shape[2] for c in cores_inflated[:-1]) + max_bond_after = max(c.shape[2] for c in cores_compressed[:-1]) + assert max_bond_after <= max_bond_before + # Invariant: state approximately preserved + psi_after = extract_psi( + cores_compressed, 'fullstate', ht.M1_modes_per_state, + ) + np.testing.assert_allclose(psi_after, psi_before, atol=0.5) + + +# ============================================================ +# TEST SUITE: tensor_to_array() +# ============================================================ + +# ------------------------------------------------------------ +# TEST: tensor_to_array recovers full state vector with nonzero auxiliaries +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tensor_to_array_recovers_vector(): + # Analytical: returns phi_0 followed by first-order auxiliaries. + # Test both representations with zero auxiliaries (initial state) + # and fullstate with nonzero auxiliaries after injection. + for method in ['fullstate', 'number']: + ht = _make_tensor(method) + tb = _make_tb() + result = tensor_to_array( + ht.list_cores_phi, method, ht.M1_modes_per_state, + tb.system, tb.mode, + ) + np.testing.assert_allclose(result[:nsite], psi_0, atol=1e-12) + np.testing.assert_allclose(result[nsite:], 0.0, atol=1e-12) + + # Inject nonzero occupation-1 slice into the first mode core + # and verify both phi_0 and auxiliaries are correct. + ht = _make_tensor('fullstate') + tb = _make_tb() + ht.list_cores_phi[1][:, 1, :] = 0.3 + 0.2j + result = tensor_to_array( + ht.list_cores_phi, 'fullstate', ht.M1_modes_per_state, + tb.system, tb.mode, + ) + # Physical wavefunction should still be correct after injection + np.testing.assert_allclose(result[:nsite], psi_0, atol=1e-12) + # First-order auxiliary for mode 0 should now be nonzero + aux_block = result[nsite:2 * nsite] + assert np.any(np.abs(aux_block) > 1e-10), ( + 'First-order auxiliary should be nonzero after injection' + ) + + +# ============================================================ +# TEST SUITE: flatten_cores() / unflatten_cores() +# ============================================================ + +# ------------------------------------------------------------ +# TEST: flatten then unflatten is identity +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_flatten_unflatten_roundtrip(): + # Algebraic property: unflatten(flatten(x)) == x for any list of core groups. + ht = _make_tensor('number') + original = ht.list_cores_phi + flat = flatten_cores(original) + modes_per_site = [len(g) - 1 for g in original] + restored = unflatten_cores(flat, modes_per_site) + assert len(restored) == len(original) + for s in range(len(original)): + assert len(restored[s]) == len(original[s]) + for i in range(len(original[s])): + np.testing.assert_array_equal(restored[s][i], original[s][i]) + + # Second pass with non-uniform modes_per_site [1, 3, 2] to exercise + # unflatten_cores's int(n_modes) indexing with varying group sizes. + # The _make_tensor fixture uses uniform 2 modes per state, which + # never exercises that path. + original_nonuniform = [ + [ + np.zeros((1, 2, 1), dtype=np.complex128), + np.zeros((1, 3, 1), dtype=np.complex128), + ], + [ + np.zeros((1, 2, 1), dtype=np.complex128), + np.zeros((1, 3, 1), dtype=np.complex128), + np.zeros((1, 3, 1), dtype=np.complex128), + np.zeros((1, 3, 1), dtype=np.complex128), + ], + [ + np.zeros((1, 2, 1), dtype=np.complex128), + np.zeros((1, 3, 1), dtype=np.complex128), + np.zeros((1, 3, 1), dtype=np.complex128), + ], + ] + # Mark each core with a unique pattern so a mis-indexed restore + # would be detectable element-wise. + marker = 1 + for group in original_nonuniform: + for core in group: + core[0, 0, 0] = marker + marker += 1 + modes_nonuniform = [1, 3, 2] + flat_nonuniform = flatten_cores(original_nonuniform) + restored_nonuniform = unflatten_cores(flat_nonuniform, modes_nonuniform) + assert len(restored_nonuniform) == len(original_nonuniform) + for s in range(len(original_nonuniform)): + assert len(restored_nonuniform[s]) == len(original_nonuniform[s]) + for i in range(len(original_nonuniform[s])): + np.testing.assert_array_equal( + restored_nonuniform[s][i], original_nonuniform[s][i], + ) + + +# ============================================================ +# TEST SUITE: tensor_add() — OBC guards +# ============================================================ + +# ------------------------------------------------------------ +# TEST: tensor_add raises on non-OBC boundaries +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tensor_add_obc_guards(): + # This case tests that tensor_add raises ValueError with the + # expected error message when either MPS violates open boundary + # conditions (left bond != 1 or right bond != 1). Covers all four + # combinations — bad left / right bond on the first / second MPS — + # with match strings so a wrong-ValueError would not silently pass. + cores_ok = [ + np.ones((1, 2, 3), dtype=np.complex128), + np.ones((3, 2, 1), dtype=np.complex128), + ] + # Left bond != 1 on first MPS + cores_bad_left = [ + np.ones((2, 2, 3), dtype=np.complex128), + np.ones((3, 2, 1), dtype=np.complex128), + ] + with pytest.raises(ValueError, match='left bond of first core must be 1'): + tensor_add(cores_bad_left, cores_ok, 1e-10, 10) + # Right bond != 1 on first MPS + cores_bad_right = [ + np.ones((1, 2, 3), dtype=np.complex128), + np.ones((3, 2, 2), dtype=np.complex128), + ] + with pytest.raises(ValueError, match='right bond of last core must be 1'): + tensor_add(cores_bad_right, cores_ok, 1e-10, 10) + # Left bond != 1 on second MPS + with pytest.raises(ValueError, match='left bond of first core must be 1'): + tensor_add(cores_ok, cores_bad_left, 1e-10, 10) + # Right bond != 1 on second MPS + with pytest.raises(ValueError, match='right bond of last core must be 1'): + tensor_add(cores_ok, cores_bad_right, 1e-10, 10) + + +# ------------------------------------------------------------ +# TEST: tensor_add doubles phi_0 +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tensor_add_doubles_phi0(): + # This case tests that adding a tensor to itself doubles the physical + # wavefunction phi_0, verifying tensor_add produces the correct sum. + ht = _make_tensor('fullstate') + phi_0_before = ht.psi.copy() + other_cores = [c.copy() for c in ht.list_cores_phi] + ht.list_cores_phi = tensor_add( + ht.list_cores_phi, other_cores, ht.mps_epsilon, ht.bond_dim_max, + ) + np.testing.assert_allclose(ht.psi, 2 * phi_0_before, atol=1e-12) + + +# ------------------------------------------------------------ +# TEST: tensor_add bond dimensions are compatible +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tensor_add_bond_dims_compatible(): + # This case tests that after tensor_add, the resulting MPS has + # compatible bond dimensions between adjacent cores. + ht = _make_tensor('fullstate') + other_cores = [c.copy() for c in ht.list_cores_phi] + ht.list_cores_phi = tensor_add( + ht.list_cores_phi, other_cores, ht.mps_epsilon, ht.bond_dim_max, + ) + assert ht.check_bondsize(), 'Bond dims incompatible after tensor_add' + + +# ------------------------------------------------------------ +# TEST: tensor_add respects bond_dim_max after compression +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tensor_add_respects_bond_dim_max(): + # This case tests that after tensor_add with compression, + # no bond dimension exceeds bond_dim_max. + ht = _make_tensor('fullstate') + other_cores = [c.copy() for c in ht.list_cores_phi] + ht.list_cores_phi = tensor_add( + ht.list_cores_phi, other_cores, ht.mps_epsilon, ht.bond_dim_max, + ) + for i, core in enumerate(ht.list_cores_phi): + assert core.shape[0] <= ht.bond_dim_max, ( + f'Core {i} left bond {core.shape[0]} > bond_dim_max {ht.bond_dim_max}' + ) + assert core.shape[2] <= ht.bond_dim_max, ( + f'Core {i} right bond {core.shape[2]} > bond_dim_max {ht.bond_dim_max}' + ) + + +# ------------------------------------------------------------ +# TEST: tensor_add works in statenumber representation +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tensor_add_statenumber(): + # This case tests that adding a tensor to itself doubles phi_0 + # in statenumber representation. + ht = _make_tensor('number') + V1_phi_before = ht.psi.copy() + list_other = [c.copy() for c in ht.flat_cores] + ht.update_phi_from_flat( + tensor_add(ht.flat_cores, list_other, ht.mps_epsilon, ht.bond_dim_max), + ) + np.testing.assert_allclose( + ht.psi, + 2 * V1_phi_before, + atol=1e-10, + err_msg='tensor_add did not double phi_0 in statenumber', + ) + + +# ------------------------------------------------------------ +# TEST: tensor_add result has no NaN or zero-dim bonds +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_tensor_add_no_pathological_output(): + # This case tests that the output of tensor_add has no NaN values + # and all bond dimensions are at least 1 (no degenerate cores). + ht = _make_tensor('fullstate') + list_other = [c.copy() for c in ht.list_cores_phi] + ht.list_cores_phi = tensor_add( + ht.list_cores_phi, list_other, ht.mps_epsilon, ht.bond_dim_max, + ) + for i, core in enumerate(ht.list_cores_phi): + assert not np.any(np.isnan(core)), f'NaN found in core {i} after tensor_add' + assert core.shape[0] >= 1, f'Core {i} left bond is 0' + assert core.shape[2] >= 1, f'Core {i} right bond is 0' + + +# ============================================================ +# TEST SUITE: scale_mps() +# ============================================================ + + +# ------------------------------------------------------------ +# TEST: scale_mps with nested statenumber MPS +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_scale_mps_statenumber(): + # This case tests that scale_mps correctly handles nested + # list-of-lists statenumber MPS by scaling cores[0][0]. + ht = _make_tensor('number') + V1_psi_before = ht.psi.copy() + scale_mps(ht.list_cores_phi, 3.0) + np.testing.assert_allclose( + ht.psi, + 3.0 * V1_psi_before, + atol=1e-10, + err_msg='scale_mps did not scale statenumber MPS correctly', + ) + + +# ============================================================ +# TEST SUITE: calc_mps_complexity() +# ============================================================ +# Per-core formula: D_left * D_right * max(D_left, D_right) * d_phys +# Total: sum over all cores. + +# ------------------------------------------------------------ +# TEST: Empty list returns 0 +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_calc_mps_complexity_empty(): + # This case tests that an empty core list yields complexity 0. + assert calc_mps_complexity([]) == 0 + + +# ------------------------------------------------------------ +# TEST: Single core matches per-core formula +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_calc_mps_complexity_single_core_symmetric(): + # This case pins the symmetric (D_left == D_right) formula: + # core shape (3, 4, 3) -> 3 * 3 * max(3,3) * 4 = 108. + core = np.zeros((3, 4, 3), dtype=np.complex128) + expected = 3 * 3 * 3 * 4 + assert calc_mps_complexity([core]) == expected + + +# ------------------------------------------------------------ +# TEST: Single core with asymmetric bonds exercises the max() branch +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_calc_mps_complexity_single_core_asymmetric(): + # This case pins the max(D_left, D_right) branch when bonds differ. + # core shape (5, 2, 3) -> 5 * 3 * max(5,3) * 2 = 150 + # core shape (3, 2, 5) -> 3 * 5 * max(3,5) * 2 = 150 (transpose-equivalent) + left_heavy = np.zeros((5, 2, 3), dtype=np.complex128) + right_heavy = np.zeros((3, 2, 5), dtype=np.complex128) + assert calc_mps_complexity([left_heavy]) == 5 * 3 * 5 * 2 + assert calc_mps_complexity([right_heavy]) == 3 * 5 * 5 * 2 + + +# ------------------------------------------------------------ +# TEST: Multi-core complexity sums per-core scores +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_calc_mps_complexity_multi_core_sum(): + # This case pins additivity: total = sum of per-core scores. + # Three cores with mixed shapes: (1, 2, 3), (3, 2, 4), (4, 2, 1). + list_cores = [ + np.zeros((1, 2, 3), dtype=np.complex128), + np.zeros((3, 2, 4), dtype=np.complex128), + np.zeros((4, 2, 1), dtype=np.complex128), + ] + expected = ( + 1 * 3 * 3 * 2 # max(1,3) = 3 + + 3 * 4 * 4 * 2 # max(3,4) = 4 + + 4 * 1 * 4 * 2 # max(4,1) = 4 + ) + assert calc_mps_complexity(list_cores) == expected + + +# ------------------------------------------------------------ +# TEST: Complexity is invariant under value changes (depends only on shape) +# ------------------------------------------------------------ +@pytest.mark.level(1) +def test_calc_mps_complexity_shape_only(): + # This case tests that the formula depends purely on shapes; element + # values do not affect the score. + core_shape = (2, 3, 4) + zeros_core = np.zeros(core_shape, dtype=np.complex128) + random_core = np.random.default_rng(0).standard_normal(core_shape) \ + + 1j * np.random.default_rng(1).standard_normal(core_shape) + assert calc_mps_complexity([zeros_core]) == calc_mps_complexity([random_core]) diff --git a/tests/test_timing_tests.py b/tests/test_timing_tests.py index ea22d1b..d2686d5 100644 --- a/tests/test_timing_tests.py +++ b/tests/test_timing_tests.py @@ -1,4 +1,5 @@ import os +import sys import pytest import importlib.resources as resources from mesohops.timing.timing_models import * @@ -15,7 +16,7 @@ def test_absorption(): # Locate the absorption.py file within the mesohops.timing package with resources.path("mesohops.timing.timing_models", "absorption.py") as absorption_path: # Tests that the absorption file runs without error - output = os.system(f"python3 {absorption_path} {seed} {nstate}") + output = os.system(f"{sys.executable} {absorption_path} {seed} {nstate}") assert output == 0 @pytest.mark.level(3) @@ -27,7 +28,7 @@ def test_fluorescence(): with resources.path("mesohops.timing.timing_models", "fluorescence.py") as fluorescence_path: # Tests that the fluorescence file runs without error - output = os.system(f"python3 {fluorescence_path} {seed} {nstate}") + output = os.system(f"{sys.executable} {fluorescence_path} {seed} {nstate}") assert output == 0 @pytest.mark.level(3) @@ -39,7 +40,7 @@ def test_holstein_1_particle(): with resources.path("mesohops.timing.timing_models", "holstein_1_particle.py") as holstein_1_particle_path: # Tests that the holstein_1_particle file runs without error - output = os.system(f"python3 {holstein_1_particle_path} {seed} {nstate}") + output = os.system(f"{sys.executable} {holstein_1_particle_path} {seed} {nstate}") assert output == 0 @pytest.mark.level(3) @@ -51,7 +52,7 @@ def test_holstein_2_particle(): with resources.path("mesohops.timing.timing_models", "holstein_2_particle.py") as holstein_2_particle_path: # Tests that the holstein_2_particle file runs without error - output = os.system(f"python3 {holstein_2_particle_path} {seed} {nstate}") + output = os.system(f"{sys.executable} {holstein_2_particle_path} {seed} {nstate}") assert output == 0 @pytest.mark.level(3) @@ -63,7 +64,7 @@ def test_markovian_filter(): with resources.path("mesohops.timing.timing_models", "markovian_filter.py") as markovian_filter_path: # Tests that the markovian_filter file runs without error - output = os.system(f"python3 {markovian_filter_path} {seed} {nstate}") + output = os.system(f"{sys.executable} {markovian_filter_path} {seed} {nstate}") assert output == 0 @pytest.mark.level(3) @@ -75,7 +76,7 @@ def test_longedge_filter(): with resources.path("mesohops.timing.timing_models", "longedge_filter.py") as longedge_filter_path: # Tests that the longedge_filter file runs without error - output = os.system(f"python3 {longedge_filter_path} {seed} {nstate}") + output = os.system(f"{sys.executable} {longedge_filter_path} {seed} {nstate}") assert output == 0 @pytest.mark.level(3) @@ -87,7 +88,7 @@ def test_triangular_filter(): with resources.path("mesohops.timing.timing_models", "triangular_filter.py") as triangular_filter_path: # Tests that the triangular_filter file runs without error - output = os.system(f"python3 {triangular_filter_path} {seed} {nstate}") + output = os.system(f"{sys.executable} {triangular_filter_path} {seed} {nstate}") assert output == 0 @pytest.mark.level(3) @@ -99,6 +100,6 @@ def test_peierls(): with resources.path("mesohops.timing.timing_models", "peierls.py") as peierls_path: # Tests that the peierls file runs without error - output = os.system(f"python3 {peierls_path} {seed} {nstate}") + output = os.system(f"{sys.executable} {peierls_path} {seed} {nstate}") assert output == 0