Repository: quantumlib/OpenFermion
Pull request #591 (2020-04-25 00:04:38)

Spin orbital 2pos constraints


Patch:
diff --git a/src/openfermion/third_party/representability/constraints/__init__.py b/src/openfermion/third_party/representability/constraints/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/openfermion/third_party/representability/constraints/spin_orbital_2pos_constraints.py b/src/openfermion/third_party/representability/constraints/spin_orbital_2pos_constraints.py
new file mode 100644
index 0000000..9f09fbc
--- /dev/null
+++ b/src/openfermion/third_party/representability/constraints/spin_orbital_2pos_constraints.py
@@ -0,0 +1,352 @@
+from typing import List, Optional, Union
+from itertools import product
+from openfermion.third_party.representability._dualbasis import \
+    DualBasisElement, DualBasis
+from openfermion.utils._rdm_mapping_functions import kronecker_delta
+
+
+def tpdm_trace_constraint(dim: int, normalization: float) -> DualBasisElement:
+    """
+    Generate the trace constraint on the 2-PDM
+
+    Args:
+        dim:  Dimension of the spin-orbital single-particle Hilbert space
+        normalization: Desired trace value.
+
+    Returns: A DualBasisElement (Row of the constraint matrix) encoding
+             the trace constraint on the 2-RDM.
+    """
+    tensor_elements = [(i, j, i, j) for i, j in product(range(dim), repeat=2)]
+    tensor_names = ['cckk'] * (dim**2)
+    tensor_coeffs = [1.0] * (dim**2)
+    bias = 0
+    return DualBasisElement(tensor_names=tensor_names,
+                            tensor_elements=tensor_elements,
+                            tensor_coeffs=tensor_coeffs,
+                            bias=bias,
+                            scalar=normalization)
+
+
+def tpdm_antisymmetry_constraint(dim: int) -> DualBasis:
+    """
+    The dual basis elements representing the antisymmetry constraints
+
+    :param dim: spinless Fermion basis rank
+    :return: the dual basis of antisymmetry_constraints
+    :rtype: DualBasis
+    """
+    # dual_basis = DualBasis()
+    dbe_list = []
+    for p, q, r, s in product(range(dim), repeat=4):
+        if p * dim + q <= r * dim + s:
+            if p < q and r < s:
+                tensor_elements = [
+                    tuple(indices) for indices in _coord_generator(p, q, r, s)
+                ]
+                tensor_names = ['cckk'] * len(tensor_elements)
+                tensor_coeffs = [0.5] * len(tensor_elements)
+                dbe = DualBasisElement()
+                for n, e, c in zip(tensor_names, tensor_elements,
+                                   tensor_coeffs):
+                    dbe.add_element(n, e, c)
+
+                # dual_basis += dbe
+                dbe_list.append(dbe)
+
+    return DualBasis(elements=dbe_list)
+
+
+def tpdm_to_opdm_mapping(dim: int,
+                         normalization: Union[float, int]) -> DualBasis:
+    """
+    Construct the DualBasis for mapping of the tpdm to the opdm
+
+    Args:
+        dim: dimension of the spin-orbital basis.
+        normalization: Scalar for mapping tpdm to opdm.  Generally, this is
+                       1 / (N - 1) where N is the number of electrons.
+    Returns:
+        DualBasis for all dim^2
+    """
+    db_basis = DualBasis()
+    for i in range(dim):
+        for j in range(i, dim):
+            dbe = DualBasisElement()
+            # contract over tpdm terms
+            for r in range(dim):
+                # duplicate entries get summed in DualBasisElement
+                dbe.add_element('cckk', (i, r, j, r), 0.5)
+                dbe.add_element('cckk', (j, r, i, r), 0.5)
+
+            # opdm terms
+            dbe.add_element('ck', (i, j), -0.5 * normalization)
+            dbe.add_element('ck', (j, i), -0.5 * normalization)
+            dbe.simplify()
+            db_basis += dbe
+
+    return db_basis
+
+
+def opdm_to_ohdm_mapping(dim: int) -> DualBasis:
+    """
+    Map the ck to kc
+
+    D1 + Q1 = I
+
+    Args:
+        dim: dimension of the spin-orbital basis
+    Returns:
+        DualBasis for the 1-RDM representability constraint
+    """
+    dbe_list = []
+    for i in range(dim):
+        for j in range(i, dim):
+            dbe = DualBasisElement()
+            if i != j:
+                dbe.add_element('ck', (i, j), 0.5)
+                dbe.add_element('ck', (j, i), 0.5)
+                dbe.add_element('kc', (j, i), 0.5)
+                dbe.add_element('kc', (i, j), 0.5)
+                dbe.dual_scalar = 0.0
+            else:
+                dbe.add_element('ck', (i, j), 1.0)
+                dbe.add_element('kc', (i, j), 1.0)
+                dbe.dual_scalar = 1.0
+
+            # db += dbe
+            dbe_list.append(dbe)
+
+    return DualBasis(elements=dbe_list)  # db
+
+
+def sz_constraint(dim: int, sz: Union[float, int]) -> DualBasis:
+    """
+    Constraint on the 1-RDM
+
+    Args:
+        dim: dimension of the spin-orbital basis.
+        sz: expectation value of the magnetic quantum number.
+
+    Returns:
+        DualBasis
+    """
+    dbe = DualBasisElement()
+    for i in range(dim // 2):
+        dbe.add_element('ck', (2 * i, 2 * i), 0.5)
+        dbe.add_element('ck', (2 * i + 1, 2 * i + 1), -0.5)
+    dbe.dual_scalar = sz
+    return DualBasis(elements=[dbe])
+
+
+def na_constraint(dim: int, na: Union[float, int]) -> DualBasis:
+    """
+    Constraint the trace of the alpha block of the opdm to equal the number
+    of spin-up electrons
+
+    Args:
+        dim: Dimension of the spin-orbital basis.
+        na: Number of spin up electrons
+
+    Returns:
+        DualBasis representing the cosntraint that is length 1
+    """
+
+    dbe = DualBasisElement()
+    for i in range(dim // 2):
+        dbe.add_element('ck', (2 * i, 2 * i), 1.0)
+    dbe.dual_scalar = na
+    return DualBasis(elements=[dbe])
+
+
+def nb_constraint(dim, nb):
+    """
+    Constraint the trace of the alpha block of the opdm to equal the number
+    of spin-down electrons
+
+    Args:
+        dim: Dimension of the spin-orbital basis.
+        na: Number of spin down electrons
+
+    Returns:
+        DualBasis representing the cosntraint that is length 1
+    """
+    dbe = DualBasisElement()
+    for i in range(dim // 2):
+        dbe.add_element('ck', (2 * i + 1, 2 * i + 1), 1.0)
+    dbe.dual_scalar = nb
+    return DualBasis(elements=[dbe])
+
+
+def tpdm_to_thdm_mapping(dim: int) -> DualBasis:
+    """
+    Generate the dual basis elements for a mapping of the 2-RDM to the
+    two-hole-RDM
+
+    Args:
+        dim: Dimension of the spin-orbital basis
+
+    Returns:
+        DualBasis representing the equality constraint.
+    """
+    dbe_list = []
+
+    def d2q2element(p: int, q: int, r: int, s: int, factor: Union[float, int])\
+            -> DualBasisElement:
+        """
+        Build the dual basis element for symmetric form of 2-marginal
+
+        :param p: tensor index
+        :param q: tensor index
+        :param r: tensor index
+        :param s: tensor index
+        :param factor: scaling coeff for a symmetric constraint
+        :return: the dual basis of the mapping
+        """
+        dbe = DualBasisElement()
+        dbe.add_element('cckk', (p, q, r, s), -1.0 * factor)
+        dbe.add_element('kkcc', (r, s, p, q), +1.0 * factor)
+        if q == s:
+            dbe.add_element('ck', (p, r), factor)
+        if p == r:
+            dbe.add_element('ck', (q, s), factor)
+        if q == r:
+            dbe.add_element('ck', (p, s), -1. * factor)
+        if p == s:
+            dbe.add_element('ck', (q, r), -1. * factor)
+
+        dbe.dual_scalar = (
+            kronecker_delta(q, s) * kronecker_delta(p, r) -
+            kronecker_delta(q, r) * kronecker_delta(p, s)) * factor
+        return dbe
+
+    for i, j, k, l in product(range(dim), repeat=4):
+        if i * dim + j <= k * dim + l:
+            db_element = d2q2element(i, j, k, l, 0.5)
+            db_element_2 = d2q2element(k, l, i, j, 0.5)
+            dbe_list.append(db_element.join_elements(db_element_2))
+
+    return DualBasis(elements=dbe_list)
+
+
+def tpdm_to_phdm_mapping(dim: int) -> DualBasis:
+    """
+
+    Args:
+        dim: Dimension of the spin-orbital basis
+
+    Returns:
+        DualBasis representing the equality constraint
+    """
+    dbe_list = []
+
+    def g2d2map(p: int,
+                q: int,
+                r: int,
+                s: int,
+                factor: Optional[Union[float, int]] = 1) -> DualBasisElement:
+        """
+        Build the dual basis element for a symmetric 2-marginal
+
+        :param p: tensor index
+        :param q: tensor index
+        :param r: tensor index
+        :param s: tensor index
+        :param factor: weighting of the element
+        :return: the dual basis element
+        """
+        dbe = DualBasisElement()
+        if q == s:
+            dbe.add_element('ck', (p, r), -1. * factor)
+        dbe.add_element('ckck', (p, s, r, q), 1.0 * factor)
+        dbe.add_element('cckk', (p, q, r, s), 1.0 * factor)
+        dbe.dual_scalar = 0
+
+        return dbe
+
+    for i, j, k, l in product(range(dim), repeat=4):
+        if i * dim + j <= k * dim + l:
+            db_element = g2d2map(i, j, k, l, factor=0.5)
+            db_element_2 = g2d2map(k, l, i, j, factor=0.5)
+            dbe_list.append(db_element.join_elements(db_element_2))
+
+    return DualBasis(elements=dbe_list)
+
+
+def spin_orbital_linear_constraints(dim: int,
+                                    num_alpha: Union[int, float],
+                                    num_beta: Union[int, float],
+                                    constraint_list: List[str],
+                                    sz: Optional[Union[None, float, int]] = None
+                                   ) -> DualBasis:
+    """
+    Construct dual basis constraints for 2-positivity in a spin-orbital
+    basis
+
+    Args:
+        dim: Dimension of the spin-orbital basis
+        num_alpha: number of spin-up electrons.
+        num_beta: number of spin-down electrons.
+        constraint_list: list of which matrices to include in the 2-pos set.
+        sz: target magnetic quantum number.
+
+    Returns:
+        DualBasis element.
+    """
+    N = num_alpha + num_beta
+    dual_basis = DualBasis()
+    if 'cckk' in constraint_list:
+        # Natural constraints on particle-conserving 2-RDMs
+        # trace constraint on D2
+        dual_basis += tpdm_trace_constraint(dim, N * (N - 1))
+        # antisymmetry constraint
+        dual_basis += tpdm_antisymmetry_constraint(dim)
+
+    if 'ck' in constraint_list:
+        dual_basis += tpdm_to_opdm_mapping(dim, N - 1)
+        if sz is not None:
+            dual_basis += sz_constraint(dim, sz)
+
+        dual_basis += na_constraint(dim, num_alpha)
+        dual_basis += nb_constraint(dim, num_beta)
+
+        if 'kc' in constraint_list:
+            dual_basis += opdm_to_ohdm_mapping(dim)
+
+    # d2 -> q2
+    if 'kkcc' in constraint_list:
+        print('tqdm constraints')
+        dual_basis += tpdm_to_thdm_mapping(dim)
+
+    # d2 -> g2
+    if "ckck" in constraint_list:
+        print('phdm constraints')
+        dual_basis += tpdm_to_phdm_mapping(dim)
+
+    return dual_basis
+
+
+def _coord_generator(i, j, k, l):
+    """
+    Generator for equivalent spin-orbital indices given a set of four
+    spin-orbitals. Indices are in chemist notation so for real-valued chemical
+    spinless Hamiltonians the integrals have 8-fold symmetry.
+
+    using this function and iterating over the following:
+        i >= j && k >= l && ij >= kl
+        where ij = i*(i + 1)/2 + j
+              kl = k*(k + 1)/2 + l
+
+    spatial real-values
+    i, j, k, l
+    j, i, k, l
+    i, j, l, k
+    j, i, l, k
+    k, l, i, j
+    k, l, j, i
+    l, k, i, j
+    l, k, j, i
+    """
+    unique_set = {(i, j, k, l), (j, i, k, l), (i, j, l, k), (j, i, l, k),
+                  (k, l, i, j), (k, l, j, i), (l, k, i, j), (l, k, j, i)}
+    for index_element in unique_set:
+        yield index_element


Test patch:
diff --git a/src/openfermion/third_party/representability/constraints/spin_orbital_2pos_constraints_test.py b/src/openfermion/third_party/representability/constraints/spin_orbital_2pos_constraints_test.py
new file mode 100644
index 0000000..6cf24e1
--- /dev/null
+++ b/src/openfermion/third_party/representability/constraints/spin_orbital_2pos_constraints_test.py
@@ -0,0 +1,167 @@
+from itertools import product
+import os
+import numpy as np
+from openfermion.third_party.representability.constraints.spin_orbital_2pos_constraints import (  # pylint: disable=line-too-long
+    tpdm_antisymmetry_constraint, tpdm_trace_constraint, _coord_generator,
+    tpdm_to_opdm_mapping, opdm_to_ohdm_mapping, sz_constraint, na_constraint,
+    nb_constraint, tpdm_to_thdm_mapping, tpdm_to_phdm_mapping,
+    spin_orbital_linear_constraints)
+from openfermion.third_party.representability._dualbasis import \
+    DualBasisElement, DualBasis
+from openfermion.third_party.representability._namedtensor import Tensor
+from openfermion.third_party.representability._multitensor import MultiTensor
+from openfermion.config import DATA_DIRECTORY
+from openfermion import MolecularData
+from openfermion.utils import map_two_pdm_to_two_hole_dm, \
+    map_two_pdm_to_particle_hole_dm
+
+
+def test_trace_constraint():
+    dbe = tpdm_trace_constraint(4, 10)
+    assert dbe.primal_tensors_names == ['cckk'] * 4**2
+    assert dbe.primal_elements == [
+        (i, j, i, j) for i, j in product(range(4), repeat=2)
+    ]
+    assert np.isclose(dbe.constant_bias, 0)
+    assert np.isclose(dbe.dual_scalar, 10)
+
+
+def test_coord_generator():
+    i, j, k, l = 0, 1, 2, 3
+    true_set = {(i, j, k, l), (j, i, k, l), (i, j, l, k), (j, i, l, k),
+                (k, l, i, j), (k, l, j, i), (l, k, i, j), (l, k, j, i)}
+
+    assert true_set == set(_coord_generator(i, j, k, l))
+
+    i, j, k, l = 1, 1, 2, 3
+    true_set = {(i, j, k, l), (j, i, k, l), (i, j, l, k), (j, i, l, k),
+                (k, l, i, j), (k, l, j, i), (l, k, i, j), (l, k, j, i)}
+    assert true_set == set(_coord_generator(i, j, k, l))
+
+
+def test_d2_antisymm():
+    db = tpdm_antisymmetry_constraint(4)
+    for dbe in db:
+        i, j, k, l = dbe.primal_elements[0]
+        assert len(list(_coord_generator(i, j, k,
+                                         l))) == len(dbe.primal_elements)
+        assert np.isclose(dbe.constant_bias, 0)
+        assert np.isclose(dbe.dual_scalar, 0)
+        assert np.unique(dbe.primal_coeffs)
+
+
+def test_tpdm_opdm_mapping():
+    db = tpdm_to_opdm_mapping(6, 1 / 2)
+    for dbe in db:
+        assert isinstance(dbe, DualBasisElement)
+        assert set(dbe.primal_tensors_names) == {'cckk', 'ck'}
+        assert len(dbe.primal_elements) == 7 or len(dbe.primal_elements) == 14
+        assert np.isclose(dbe.dual_scalar, 0)
+        assert np.isclose(dbe.constant_bias, 0)
+        print(vars(dbe))
+        for idx, element in enumerate(dbe.primal_elements):
+            if len(element) == 4:
+                assert element[1] == element[3]
+                if element[0] == element[2]:
+                    assert np.isclose(dbe.primal_coeffs[idx], 1.0)
+                else:
+                    assert np.isclose(dbe.primal_coeffs[idx], 0.5)
+                assert dbe.primal_coeffs
+            elif len(element) == 2:
+                gem_idx = [len(x) for x in dbe.primal_elements].index(4)
+                assert sorted(element) == sorted([
+                    dbe.primal_elements[gem_idx][0],
+                    dbe.primal_elements[gem_idx][2]
+                ])
+
+
+def test_opdm_to_ohdm_mapping():
+    db = opdm_to_ohdm_mapping(6)
+    for dbe in db:
+        assert isinstance(dbe, DualBasisElement)
+        assert set(dbe.primal_tensors_names) == {'ck', 'kc'}
+        if len(dbe.primal_tensors_names) == 4:
+            assert np.allclose(dbe.primal_coeffs, 0.5)
+            assert np.isclose(dbe.dual_scalar, 0.0)
+        elif len(dbe.primal_tensors_names) == 2:
+            assert np.allclose(dbe.primal_coeffs, 1.0)
+            assert np.isclose(dbe.dual_scalar, 1.0)
+
+
+def test_sz_constraint():
+    db = sz_constraint(6, 3.5)
+    test_dbe = db.elements[0]
+    for i in range(3):
+        np.isclose(test_dbe.primal_coeffs[2 * i], 0.5)
+        np.isclose(test_dbe.primal_coeffs[2 * i + 1], 0.5)
+        assert test_dbe.primal_tensors_names[2 * i] == 'ck'
+        assert test_dbe.primal_tensors_names[2 * i + 1] == 'ck'
+    assert np.isclose(test_dbe.dual_scalar, 3.5)
+
+
+def test_na_constraint():
+    db = na_constraint(6, 3)
+    assert len(db) == 1
+    test_dbe = db.elements[0]
+    for i in range(3):
+        np.isclose(test_dbe.primal_coeffs[i], 1.0)
+        assert test_dbe.primal_elements[i] == (2 * i, 2 * i)
+    assert np.isclose(test_dbe.dual_scalar, 3)
+
+
+def test_nb_constraint():
+    db = nb_constraint(6, 3)
+    assert len(db) == 1
+    test_dbe = db.elements[0]
+    for i in range(3):
+        np.isclose(test_dbe.primal_coeffs[i], 1.0)
+        assert test_dbe.primal_elements[i] == (2 * i + 1, 2 * i + 1)
+    assert np.isclose(test_dbe.dual_scalar, 3)
+
+
+def test_d2_to_q2():
+    # this should test for correctness
+    filename = os.path.join(DATA_DIRECTORY, "H1-Li1_sto-3g_singlet_1.45.hdf5")
+    molecule = MolecularData(filename=filename)
+    molecule.load()
+    opdm = molecule.fci_one_rdm
+    tpdm = molecule.fci_two_rdm
+    thdm = map_two_pdm_to_two_hole_dm(tpdm, opdm)
+
+    db = tpdm_to_thdm_mapping(6)
+    topdm = Tensor(tensor=opdm, name='ck')
+    ttpdm = Tensor(tensor=np.einsum('ijlk', tpdm), name='cckk')
+    tthdm = Tensor(tensor=np.einsum('ijlk', thdm), name='kkcc')
+    mt = MultiTensor(tensors=[topdm, ttpdm, tthdm], dual_basis=db)
+    A, _, b = mt.synthesize_dual_basis()
+    vec_rdms = mt.vectorize_tensors()
+    assert np.isclose(np.linalg.norm(A.dot(vec_rdms) - b), 0)
+
+
+def test_d2_to_g2():
+    # this should test for correctness
+    filename = os.path.join(DATA_DIRECTORY, "H1-Li1_sto-3g_singlet_1.45.hdf5")
+    molecule = MolecularData(filename=filename)
+    molecule.load()
+    opdm = molecule.fci_one_rdm
+    tpdm = molecule.fci_two_rdm
+    phdm = map_two_pdm_to_particle_hole_dm(tpdm, opdm)
+
+    db = tpdm_to_phdm_mapping(molecule.n_qubits)
+    topdm = Tensor(tensor=opdm, name='ck')
+    ttpdm = Tensor(tensor=np.einsum('ijlk', tpdm), name='cckk')
+    tphdm = Tensor(tensor=np.einsum('ijlk', phdm), name='ckck')
+    mt = MultiTensor(tensors=[topdm, ttpdm, tphdm], dual_basis=db)
+    A, _, b = mt.synthesize_dual_basis()
+    vec_rdms = mt.vectorize_tensors()
+    assert np.isclose(np.linalg.norm(A.dot(vec_rdms) - b), 0)
+
+
+def test_spin_orbital_dual_basis_construction():
+    db = spin_orbital_linear_constraints(
+        dim=6,
+        num_alpha=3,
+        num_beta=3,
+        constraint_list=['ck', 'kc', 'cckk', 'kkcc', 'ckck'],
+        sz=0)
+    assert isinstance(db, DualBasis)