Generic spin scbk

packages/openfermion/quri_parts/openfermion/transforms/__init__.py

@@ -14,11 +14,12 @@
 
 from openfermion.ops import FermionOperator, InteractionOperator, MajoranaOperator
 from openfermion.transforms import bravyi_kitaev as of_bravyi_kitaev
-from openfermion.transforms import get_fermion_operator
+from openfermion.transforms import bravyi_kitaev_tree, get_fermion_operator
 from openfermion.transforms import jordan_wigner as of_jordan_wigner
-from openfermion.transforms import (
-    symmetry_conserving_bravyi_kitaev as of_symmetry_conserving_bravyi_kitaev,
-)
+from openfermion.transforms import reorder
+from openfermion.transforms.opconversions import edit_hamiltonian_for_spin
+from openfermion.transforms.opconversions.remove_symmetry_qubits import remove_indices
+from openfermion.utils import up_then_down
 from typing_extensions import TypeAlias
 
 from quri_parts.chem.transforms import (
@@ -372,7 +373,7 @@ class OpenFermionJordanWignerFactory(
 Example:
     You may create mappers out of `jordan_wigner`
 
-    >>> operator_mapper = jordan_wigner.get_of_operator_mapper()
+    >>> operator_mapper = jordan_wigner.get_of_operator_mapper(8)
     >>> operator_mapper(FermionOperator("1^ 1"))
     (0.5+0j)*I + (-0.5+0j)*Z1
 
@@ -448,7 +449,7 @@ class OpenFermionBravyiKitaevFactory(
 Example:
     You may create mappers out of `bravyi_kitaev`
 
-    >>> operator_mapper = bravyi_kitaev.get_of_operator_mapper()
+    >>> operator_mapper = bravyi_kitaev.get_of_operator_mapper(8)
     >>> print(operator_mapper(FermionOperator("1^ 1")))
     (0.5+0j)*I + (-0.5+0j)*Z0 Z1
 
@@ -478,6 +479,11 @@ class OpenFermionBravyiKitaevFactory(
 """
 
 
+def _get_scbk_parity_factor(n_fermions: int, sz: float) -> tuple[int, int]:
+    n_spin_ups = int(0.5 * (n_fermions + 2 * sz))
+    return (-1) ** n_spin_ups, (-1) ** n_fermions
+
+
 class OpenFermionSymmetryConservingBravyiKitaev(
     SymmetryConservingBravyiKitaev, OpenFermionQubitMapping
 ):
@@ -518,28 +524,49 @@ def of_operator_mapper(self) -> OpenFermionQubitOperatorMapper:
 
         if self.n_fermions is None:
             raise ValueError("n_fermions is required.")
-        if self.sz not in [0.0, 0.5]:
-            raise ValueError("Current implementation only supports sz = 0.0 or 0.5.")
+        if self.sz is None:
+            raise ValueError("sz is required.")
+
+        n_spin_orbitals = self.n_spin_orbitals
+        n_fermions = self.n_fermions
+        sz = self.sz
 
         def mapper(
             op: Union["FermionOperator", "InteractionOperator", "MajoranaOperator"]
         ) -> "Operator":
             if not isinstance(op, FermionOperator):
                 op = get_fermion_operator(op)
 
-            openfermion_op = FermionOperator()
-            for pauli_product, coef in op.terms.items():
-                openfermion_op += FermionOperator(pauli_product, coef)
-
-            operator = Operator()
-            for single_term in openfermion_op.get_operators():
-                if has_particle_number_symmetry(single_term, check_spin_symmetry=True):
-                    operator += operator_from_openfermion_op(
-                        of_symmetry_conserving_bravyi_kitaev(
-                            single_term, self.n_spin_orbitals, self.n_fermions
-                        )
-                    )
-            return operator
+            symmetry_op = FermionOperator()
+            for op_tuple, coeff in op.terms.items():
+                single_op = FermionOperator(op_tuple, coeff)
+                if has_particle_number_symmetry(single_op, check_spin_symmetry=True):
+                    symmetry_op += single_op
+
+            (
+                mid_parity_factor,
+                last_parity_factor,
+            ) = _get_scbk_parity_factor(n_fermions, sz)
+
+            transformed_op = bravyi_kitaev_tree(
+                reorder(symmetry_op, up_then_down, num_modes=n_spin_orbitals),
+                n_qubits=n_spin_orbitals,
+            )
+            operator_tappered = edit_hamiltonian_for_spin(
+                transformed_op,
+                spin_orbital=n_spin_orbitals,
+                orbital_parity=last_parity_factor,
+            )
+            operator_tappered = edit_hamiltonian_for_spin(
+                operator_tappered,
+                spin_orbital=n_spin_orbitals // 2,
+                orbital_parity=mid_parity_factor,
+            )
+            operator_tappered = remove_indices(
+                operator_tappered, [n_spin_orbitals // 2, n_spin_orbitals]
+            )
+
+            return operator_from_openfermion_op(operator_tappered)
 
         return mapper
 

packages/openfermion/tests/openfermion/transforms/test_transforms.py

@@ -14,6 +14,7 @@
     symmetry_conserving_bravyi_kitaev as of_symmetry_conserving_bravyi_kitaev,
 )
 
+from quri_parts.core.operator import PAULI_IDENTITY, Operator, pauli_label
 from quri_parts.core.state import ComputationalBasisState
 from quri_parts.openfermion.operator import (
     FermionOperator,
@@ -68,6 +69,73 @@ def test_scbk_op_mapper(self) -> None:
         assert total_transformed == expected
         assert of_total_transformed == expected
 
+        # n_spin_up = 2, n_fermions = 2
+        n_spin_orbitals = 4
+        n_fermions = 2
+        sz = 1.0
+
+        expected = Operator({PAULI_IDENTITY: 1})
+        expected.add_term(pauli_label("X0"), 0.25)  # (0.25+0j) [X0]
+        expected.add_term(pauli_label("X0"), (+1) * -0.25)  # (-0.25+0j) [X0 Z1]
+        expected.add_term(pauli_label("Y0"), 0.25j)  # 0.25j [Y0]
+        expected.add_term(pauli_label("Y0"), (+1) * -0.25j)  # -0.25j [Y0 Z1]
+        expected.add_term(pauli_label("X1"), (+1) * -0.25)  # (-0.25+0j) [Z1 X2 Z3]
+        expected.add_term(pauli_label("Y1"), (+1) * -0.25j)  # -0.25j [Z1 Y2 Z3]
+        expected.add_term(pauli_label("X1"), 0.25)  # (0.25+0j) [X2]
+        expected.add_term(pauli_label("Y1"), 0.25j)  # 0.25j [Y2]
+
+        op_mapper = symmetry_conserving_bravyi_kitaev(
+            n_spin_orbitals, n_fermions, sz
+        ).of_operator_mapper
+        assert op_mapper(op_total) == expected
+
+        # n_spin_up = 1, n_fermions = 3
+        n_spin_orbitals = 4
+        n_fermions = 3
+        sz = -0.5
+
+        expected = Operator({PAULI_IDENTITY: 1})
+        expected.add_term(pauli_label("X0"), 0.25)  # (0.25+0j) [X0]
+        expected.add_term(pauli_label("X0"), (-1) * -0.25)  # (-0.25+0j) [X0 Z1]
+        expected.add_term(pauli_label("Y0"), 0.25j)  # 0.25j [Y0]
+        expected.add_term(pauli_label("Y0"), (-1) * -0.25j)  # -0.25j [Y0 Z1]
+        expected.add_term(pauli_label("X1"), (+1) * -0.25)  # (-0.25+0j) [Z1 X2 Z3]
+        expected.add_term(pauli_label("Y1"), (+1) * -0.25j)  # -0.25j [Z1 Y2 Z3]
+        expected.add_term(pauli_label("X1"), 0.25)  # (0.25+0j) [X2]
+        expected.add_term(pauli_label("Y1"), 0.25j)  # 0.25j [Y2]
+
+        op_mapper = symmetry_conserving_bravyi_kitaev(
+            n_spin_orbitals, n_fermions, sz
+        ).of_operator_mapper
+        assert op_mapper(op_total) == expected
+
+        # n_spin_up = 2, n_fermions = 3
+        n_spin_orbitals = 4
+        n_fermions = 3
+        sz = 0.5
+
+        expected = Operator({PAULI_IDENTITY: 1})
+        expected.add_term(pauli_label("X0"), 0.25)  # (0.25+0j) [X0]
+        expected.add_term(pauli_label("X0"), (+1) * -0.25)  # (-0.25+0j) [X0 Z1]
+        expected.add_term(pauli_label("Y0"), 0.25j)  # 0.25j [Y0]
+        expected.add_term(pauli_label("Y0"), (+1) * -0.25j)  # -0.25j [Y0 Z1]
+        expected.add_term(pauli_label("X1"), (-1) * -0.25)  # (-0.25+0j) [Z1 X2 Z3]
+        expected.add_term(pauli_label("Y1"), (-1) * -0.25j)  # -0.25j [Z1 Y2 Z3]
+        expected.add_term(pauli_label("X1"), 0.25)  # (0.25+0j) [X2]
+        expected.add_term(pauli_label("Y1"), 0.25j)  # 0.25j [Y2]
+
+        op_mapper = symmetry_conserving_bravyi_kitaev(
+            n_spin_orbitals, n_fermions, sz
+        ).of_operator_mapper
+        assert op_mapper(op_total) == expected
+
+        expected = operator_from_openfermion_op(
+            of_symmetry_conserving_bravyi_kitaev(
+                of_op_symmetry_conserve, n_spin_orbitals, n_fermions
+            )
+        )
+        assert op_mapper(op_total) == expected
+
 
 class TestStateMapper:
     def test_jw_state_mapper(self) -> None:
@@ -394,3 +462,71 @@ def test_scbk_state_mapper_property(self) -> None:
             ComputationalBasisState(n_spin_orbitals - 2, bits=0b111111)
         )
         assert inv_mapped == [0, 1, 4, 5]
+
+        # Test for generic spin SCBK: integer spin
+        n_spin_orbitals = 8
+        n_fermions = 4
+        state_mapper = symmetry_conserving_bravyi_kitaev(
+            n_spin_orbitals, n_fermions, 1.0
+        ).state_mapper
+        inv_state_mapper = symmetry_conserving_bravyi_kitaev(
+            n_spin_orbitals, n_fermions, 1.0
+        ).inv_state_mapper
+
+        mapped = state_mapper([0, 1, 2, 4])
+        assert mapped == ComputationalBasisState(n_spin_orbitals - 2, bits=0b11101)
+
+        inv_mapped = inv_state_mapper(
+            ComputationalBasisState(n_spin_orbitals - 2, bits=0b11101)
+        )
+        assert inv_mapped == [0, 1, 2, 4]
+
+        mapped = state_mapper([0, 2, 3, 4])
+        assert mapped == ComputationalBasisState(n_spin_orbitals - 2, bits=0b10101)
+
+        inv_mapped = inv_state_mapper(
+            ComputationalBasisState(n_spin_orbitals - 2, bits=0b10101)
+        )
+        assert inv_mapped == [0, 2, 3, 4]
+
+        mapped = state_mapper([0, 1, 2, 6])
+        assert mapped == ComputationalBasisState(n_spin_orbitals - 2, bits=0b11001)
+
+        inv_mapped = inv_state_mapper(
+            ComputationalBasisState(n_spin_orbitals - 2, bits=0b11001)
+        )
+        assert inv_mapped == [0, 1, 2, 6]
+
+        # Test for generic spin SCBK: half integer spin
+        n_spin_orbitals = 8
+        n_fermions = 3
+        state_mapper = symmetry_conserving_bravyi_kitaev(
+            n_spin_orbitals, n_fermions, -1.5
+        ).state_mapper
+        inv_state_mapper = symmetry_conserving_bravyi_kitaev(
+            n_spin_orbitals, n_fermions, -1.5
+        ).inv_state_mapper
+
+        mapped = state_mapper([1, 3, 5])
+        assert mapped == ComputationalBasisState(n_spin_orbitals - 2, bits=0b101000)
+
+        inv_mapped = inv_state_mapper(
+            ComputationalBasisState(n_spin_orbitals - 2, bits=0b101000)
+        )
+        assert inv_mapped == [1, 3, 5]
+
+        mapped = state_mapper([1, 3, 7])
+        assert mapped == ComputationalBasisState(n_spin_orbitals - 2, bits=0b1000)
+
+        inv_mapped = inv_state_mapper(
+            ComputationalBasisState(n_spin_orbitals - 2, bits=0b1000)
+        )
+        assert inv_mapped == [1, 3, 7]
+
+        mapped = state_mapper([3, 5, 7])
+        assert mapped == ComputationalBasisState(n_spin_orbitals - 2, bits=0b110000)
+
+        inv_mapped = inv_state_mapper(
+            ComputationalBasisState(n_spin_orbitals - 2, bits=0b110000)
+        )
+        assert inv_mapped == [3, 5, 7]

packages/openfermion/tests/openfermion/utils/test_post_selection_filters.py

@@ -8,8 +8,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import pytest
-
 from quri_parts.openfermion.utils.post_selection_filters import (
     create_bk_electron_number_post_selection_filter_fn,
     create_jw_electron_number_post_selection_filter_fn,
@@ -82,11 +80,36 @@ def test_create_scbk_electron_number_post_selection_filter_fn() -> None:
     assert filter_fn(0b110)  # [2, 4, 7]
     assert not filter_fn(48)  # [3, 5, 6]
 
-    with pytest.raises(ValueError):
-        filter_fn = create_scbk_electron_number_post_selection_filter_fn(
-            qubit_count, 2, sz=1.0
-        )
-    with pytest.raises(ValueError):
-        filter_fn = create_scbk_electron_number_post_selection_filter_fn(
-            qubit_count, 3, sz=-0.5
-        )
+    filter_fn = create_scbk_electron_number_post_selection_filter_fn(
+        qubit_count, 2, sz=1.0
+    )
+    assert filter_fn(0b1)  # [0, 2]
+    assert filter_fn(0b111)  # [0, 4]
+    assert filter_fn(0b11)  # [0, 6]
+    assert filter_fn(0b110)  # [2, 4]
+    assert filter_fn(0b10)  # [2, 6]
+    assert filter_fn(0b100)  # [4, 6]
+    assert not filter_fn(0b11011)  # [0, 1]
+    assert not filter_fn(0b110000)  # [3, 5]
+    assert not filter_fn(0b10010)  # [2, 3]
+
+    filter_fn = create_scbk_electron_number_post_selection_filter_fn(
+        qubit_count, 3, sz=-0.5
+    )
+    assert filter_fn(0b1011)  # [0, 1, 3]
+    assert filter_fn(0b111011)  # [0, 1, 5]
+    assert filter_fn(0b11011)  # [0, 1, 7]
+    assert filter_fn(0b110011)  # [0, 3, 5]
+    assert filter_fn(0b10011)  # [0, 3, 7]
+    assert filter_fn(0b100011)  # [0, 5, 7]
+    assert filter_fn(0b1010)  # [1, 2, 3]
+    assert filter_fn(0b111010)  # [1, 2, 5]
+    assert filter_fn(0b11010)  # [1, 2, 7]
+    assert filter_fn(0b1100)  # [1, 3, 4]
+    assert filter_fn(0b1000)  # [1, 3, 6]
+    assert filter_fn(0b1100)  # [1, 4, 5]
+    assert not filter_fn(0b11001)  # [0, 1, 2]
+    assert not filter_fn(0b11111)  # [0, 1, 4]
+    assert not filter_fn(0b1)  # [0, 2, 7]
+    assert not filter_fn(0b111)  # [0, 4, 6]
+    assert not filter_fn(0b11110)  # [1, 2, 4]