use random state vector and MPS for tests

python/ffsim/tenpy/__init__.py

@@ -24,7 +24,8 @@
 from ffsim.tenpy.gates.orbital_rotation import apply_orbital_rotation
 from ffsim.tenpy.gates.ucj import apply_ucj_op_spin_balanced
 from ffsim.tenpy.hamiltonians.molecular_hamiltonian import MolecularHamiltonianMPOModel
-from ffsim.tenpy.util import bitstring_to_mps
+from ffsim.tenpy.random.random import random_mps
+from ffsim.tenpy.util import bitstring_to_mps, mps_to_statevector, statevector_to_mps
 
 __all__ = [
     "apply_ucj_op_spin_balanced",
@@ -35,7 +36,10 @@
     "bitstring_to_mps",
     "givens_rotation",
     "MolecularHamiltonianMPOModel",
+    "mps_to_statevector",
     "num_interaction",
     "num_num_interaction",
     "on_site_interaction",
+    "random_mps",
+    "statevector_to_mps",
 ]

python/ffsim/tenpy/random/__init__.py

@@ -0,0 +1,9 @@
+# (C) Copyright IBM 2025.
+#
+# This code is licensed under the Apache License, Version 2.0. You may
+# obtain a copy of this license in the LICENSE.txt file in the root directory
+# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# Any modifications or derivative works of this code must retain this
+# copyright notice, and modified files need to carry a notice indicating
+# that they have been altered from the originals.

python/ffsim/tenpy/random/random.py

@@ -0,0 +1,57 @@
+# (C) Copyright IBM 2025.
+#
+# This code is licensed under the Apache License, Version 2.0. You may
+# obtain a copy of this license in the LICENSE.txt file in the root directory
+# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# Any modifications or derivative works of this code must retain this
+# copyright notice, and modified files need to carry a notice indicating
+# that they have been altered from the originals.
+
+from tenpy.algorithms.tebd import RandomUnitaryEvolution
+from tenpy.networks.mps import MPS
+
+import ffsim
+from ffsim.tenpy.util import bitstring_to_mps
+
+
+def random_mps(
+    norb: int, nelec: tuple[int, int], n_steps: int = 10, chi_max: int = 100
+) -> MPS:
+    """Return a random MPS generated from a random unitary evolution.
+
+    Args:
+        norb: The number of orbitals.
+        nelec: The number of electrons.
+        n_steps: The number of steps in the random unitary evolution.
+        chi_max: The maximum bond dimension in the random unitary evolution.
+
+    Returns:
+        The random MPS.
+    """
+
+    # initialize Hartree-Fock state
+    dim = ffsim.dim(norb, nelec)
+    strings = ffsim.addresses_to_strings(
+        range(dim), norb=norb, nelec=nelec, bitstring_type=ffsim.BitstringType.STRING
+    )
+    string_tuples = [
+        (
+            int(string[len(string) // 2 :], base=2),
+            int(string[: len(string) // 2], base=2),
+        )
+        for string in strings
+    ]
+    mps = bitstring_to_mps(string_tuples[0], norb)
+
+    # apply random unitary evolution
+    tebd_params = {
+        "N_steps": n_steps,
+        "trunc_params": {"chi_max": chi_max},
+        "verbose": 0,
+    }
+    eng = RandomUnitaryEvolution(mps, tebd_params)
+    eng.run()
+    mps.canonical_form()
+
+    return mps

python/ffsim/tenpy/util.py

@@ -12,9 +12,16 @@
 
 from __future__ import annotations
 
+from copy import deepcopy
+
+import numpy as np
 import tenpy.linalg.np_conserved as npc
+from tenpy.algorithms.exact_diag import ExactDiag
+from tenpy.models.model import MPOModel
 from tenpy.networks.mps import MPS
-from tenpy.networks.site import SpinHalfFermionSite
+from tenpy.networks.site import FermionSite, SpinHalfFermionSite
+
+import ffsim
 
 
 def bitstring_to_mps(bitstring: tuple[int, int], norb: int) -> MPS:
@@ -35,22 +42,34 @@ def bitstring_to_mps(bitstring: tuple[int, int], norb: int) -> MPS:
 
     # relabel using TeNPy SpinHalfFermionSite convention
     product_state = []
-    swap_factor = 1
-    previous_site_occupation = None
     for i, site in enumerate(zip(reversed(string_b), reversed(string_a))):
         site_occupation = int("".join(site), base=2)
         product_state.append(site_occupation)
 
-        if site_occupation in [0b01, 0b11] and previous_site_occupation in [0b10, 0b11]:
-            swap_factor *= -1
-
-        previous_site_occupation = site_occupation
-
     # construct product state MPS
     shfs = SpinHalfFermionSite(cons_N="N", cons_Sz="Sz")
     mps = MPS.from_product_state([shfs] * norb, product_state)
 
-    # map from ffsim to TeNPy ordering
+    # map from TeNPy to ffsim ordering
+    fs = FermionSite(conserve="N")
+    alpha_sector = mps.expectation_value("Nu")
+    beta_sector = mps.expectation_value("Nd")
+    product_state_fs_tenpy = [
+        int(val) for pair in zip(alpha_sector, beta_sector) for val in pair
+    ]
+    mps_fs = MPS.from_product_state([fs] * 2 * norb, product_state_fs_tenpy)
+
+    tenpy_ordering = list(range(2 * norb))
+    midpoint = len(tenpy_ordering) // 2
+    mask1 = tenpy_ordering[:midpoint][::-1]
+    mask2 = tenpy_ordering[midpoint:][::-1]
+    ffsim_ordering = [int(val) for pair in zip(mask1, mask2) for val in pair]
+
+    mps_ref = deepcopy(mps_fs)
+    mps_ref.permute_sites(ffsim_ordering, swap_op=None)
+    mps_fs.permute_sites(ffsim_ordering, swap_op="auto")
+    swap_factor = mps_fs.overlap(mps_ref)
+
     if swap_factor == -1:
         minus_identity_npc = npc.Array.from_ndarray(
             -shfs.get_op("Id").to_ndarray(),
@@ -60,3 +79,166 @@ def bitstring_to_mps(bitstring: tuple[int, int], norb: int) -> MPS:
         mps.apply_local_op(0, minus_identity_npc)
 
     return mps
+
+
+def mps_to_statevector(mps: MPS, mpo_model: MPOModel) -> np.ndarray:
+    r"""Return the MPS as a state vector.
+
+    Args:
+        mps: The MPS.
+        mpo_model: The MPO model.
+
+    Returns:
+        The state vector.
+    """
+
+    # generate the (ffsim-ordered) list of product states
+    norb = mps.L
+    n_alpha = round(np.sum(mps.expectation_value("Nu")))
+    n_beta = round(np.sum(mps.expectation_value("Nd")))
+    nelec = (n_alpha, n_beta)
+    product_states = _generate_product_states(norb, nelec)
+
+    # initialize the TeNPy ExactDiag class instance
+    charge_sector = mps.get_total_charge(True)
+    exact_diag = ExactDiag(mpo_model, charge_sector=charge_sector)
+
+    # determine the mapping from TeNPy basis to ffsim basis
+    basis_ordering_ffsim, swap_factors_ffsim = _map_tenpy_to_ffsim_basis(
+        product_states, exact_diag, norb
+    )
+
+    # convert TeNPy MPS to ffsim statevector
+    statevector = exact_diag.mps_to_full(mps).to_ndarray()
+    statevector = np.multiply(swap_factors_ffsim, statevector[basis_ordering_ffsim])
+
+    return statevector
+
+
+def statevector_to_mps(
+    statevector: np.ndarray, mpo_model: MPOModel, norb: int, nelec: tuple[int, int]
+) -> MPS:
+    r"""Return the state vector as an MPS.
+
+    Args:
+        statevector: The state vector.
+        mpo_model: The MPO model.
+        norb: The number of orbitals.
+        nelec: The number of electrons.
+
+    Returns:
+        The MPS.
+    """
+
+    # generate the (ffsim-ordered) list of product states
+    product_states = _generate_product_states(norb, nelec)
+
+    # initialize the TeNPy ExactDiag class instance
+    mps_reference = MPS.from_product_state(mpo_model.lat.mps_sites(), product_states[0])
+    charge_sector = mps_reference.get_total_charge(True)
+    exact_diag = ExactDiag(mpo_model, charge_sector=charge_sector)
+    statevector_reference = exact_diag.mps_to_full(mps_reference)
+    leg_charge = statevector_reference.legs[0]
+
+    # determine the mapping from ffsim basis to TeNPy basis
+    basis_ordering_ffsim, swap_factors_ffsim = _map_tenpy_to_ffsim_basis(
+        product_states, exact_diag, norb
+    )
+    basis_ordering_tenpy = np.argsort(basis_ordering_ffsim)
+    swap_factors_tenpy = swap_factors_ffsim[np.argsort(basis_ordering_ffsim)]
+
+    # convert ffsim statevector to TeNPy MPS
+    statevector = np.multiply(swap_factors_tenpy, statevector[basis_ordering_tenpy])
+    statevector_npc = npc.Array.from_ndarray(statevector, [leg_charge])
+    mps = exact_diag.full_to_mps(statevector_npc)
+
+    return mps
+
+
+def _generate_product_states(norb: int, nelec: tuple[int, int]) -> list:
+    r"""Generate the ffsim-ordered list of product states in TeNPy notation.
+
+    Args:
+        norb: The number of orbitals.
+        nelec: The number of electrons.
+
+    Returns:
+        The ffsim-ordered list of product states in TeNPy notation.
+    """
+
+    # generate the strings
+    dim = ffsim.dim(norb, nelec)
+    strings = ffsim.addresses_to_strings(
+        range(dim), norb=norb, nelec=nelec, bitstring_type=ffsim.BitstringType.STRING
+    )
+    string_tuples = [
+        (
+            int(string[len(string) // 2 :], base=2),
+            int(string[: len(string) // 2], base=2),
+        )
+        for string in strings
+    ]
+
+    # convert strings to product states
+    product_states = []
+    for bitstring in string_tuples:
+        # unpack bitstrings
+        int_a, int_b = bitstring
+        string_a = format(int_a, f"0{norb}b")
+        string_b = format(int_b, f"0{norb}b")
+
+        # relabel using TeNPy SpinHalfFermionSite convention
+        product_state = []
+        for i, site in enumerate(zip(reversed(string_b), reversed(string_a))):
+            site_occupation = int("".join(site), base=2)
+            product_state.append(site_occupation)
+        product_states.append(product_state)
+
+    return product_states
+
+
+def _map_tenpy_to_ffsim_basis(
+    product_states: list, exact_diag: ExactDiag, norb: int
+) -> tuple[np.ndarray, np.ndarray]:
+    r"""Map from the TeNPy basis to the ffsim basis.
+
+    Args:
+        product_states: The ffsim-ordered list of product states in TeNPy notation.
+        exact_diag: The TeNPy ExactDiag class instance.
+        norb: The number of orbitals.
+
+    Returns:
+        basis_ordering_ffsim: The permutation to map from the TeNPy to ffsim basis.
+        swap_factors: The minus signs that are introduced due to this mapping.
+    """
+
+    basis_ordering_ffsim = []
+    swap_factors = []
+    for i, state in enumerate(product_states):
+        # basis_ordering_ffsim
+        prod_mps = MPS.from_product_state(exact_diag.model.lat.mps_sites(), state)
+        prod_statevector = list(exact_diag.mps_to_full(prod_mps).to_ndarray())
+        idx = prod_statevector.index(1)
+        basis_ordering_ffsim.append(idx)
+
+        # swap_factors
+        fs = FermionSite(conserve="N")
+        alpha_sector = prod_mps.expectation_value("Nu")
+        beta_sector = prod_mps.expectation_value("Nd")
+        product_state_fs_tenpy = [
+            int(val) for pair in zip(alpha_sector, beta_sector) for val in pair
+        ]
+        mps_fs = MPS.from_product_state([fs] * 2 * norb, product_state_fs_tenpy)
+        #
+        tenpy_ordering = list(range(2 * norb))
+        midpoint = len(tenpy_ordering) // 2
+        mask1 = tenpy_ordering[:midpoint][::-1]
+        mask2 = tenpy_ordering[midpoint:][::-1]
+        ffsim_ordering = [int(val) for pair in zip(mask1, mask2) for val in pair]
+        #
+        mps_ref = deepcopy(mps_fs)
+        mps_ref.permute_sites(ffsim_ordering, swap_op=None)
+        mps_fs.permute_sites(ffsim_ordering, swap_op="auto")
+        swap_factors.append(mps_fs.overlap(mps_ref))
+
+    return np.array(basis_ordering_ffsim), np.array(swap_factors)