add u1 helper functions in quantum

Author: refraction-ray

CHANGELOG.md

@@ -4,12 +4,14 @@
 
 ### Added
 
-- Add jax jitted function load/save utilities in experimental module
+- Add jax jitted function load/save utilities in experimental.py
 
 - Add `circuit.to_openqasm_file` function for compatibility of qiskit>1
 
 - Add `tc.cite()` to get citation information
 
+- Add `u1_inds`, `u1_mask`, `u1_project`, and `u1_enlarge` functions in quantum.py as utilities in charged conservation systems
+
 ### Fixed
 
 - Fix customized jax eigh operator by noting the return is a namedtuple

tensorcircuit/backends/jax_backend.py

@@ -609,14 +609,16 @@ def f_jax(*args: Any, **kws: Any) -> Any:
         return carry
 
     def scatter(self, operand: Tensor, indices: Tensor, updates: Tensor) -> Tensor:
-        rank = len(operand.shape)
-        dnums = libjax.lax.ScatterDimensionNumbers(
-            update_window_dims=(),
-            inserted_window_dims=tuple([i for i in range(rank)]),
-            scatter_dims_to_operand_dims=tuple([i for i in range(rank)]),
-        )
-        r = libjax.lax.scatter(operand, indices, updates, dnums)
-        return r
+        updates = jnp.reshape(updates, indices.shape)
+        return operand.at[indices].set(updates)
+        # rank = len(operand.shape)
+        # dnums = libjax.lax.ScatterDimensionNumbers(
+        #     update_window_dims=(),
+        #     inserted_window_dims=tuple([i for i in range(rank)]),
+        #     scatter_dims_to_operand_dims=tuple([i for i in range(rank)]),
+        # )
+        # r = libjax.lax.scatter(operand, indices, updates, dnums)
+        # return r
 
     def coo_sparse_matrix(
         self, indices: Tensor, values: Tensor, shape: Tensor

tensorcircuit/quantum.py

@@ -11,6 +11,7 @@
 # pylint: disable=invalid-name
 
 import logging
+import math
 import os
 from functools import partial, reduce
 from operator import matmul, mul, or_
@@ -1160,6 +1161,96 @@ def quimb2qop(qb_mpo: Any) -> QuOperator:
     return qop
 
 
+def u1_inds(n: int, m: int) -> Tensor:
+    """
+    Generate all the combination index of m down spins in n sites.
+
+    .. code-block:: python
+
+        print(u1_inds(5, 1))
+        # [1, 2, 4, 8, 16]
+
+
+    :param n: number of total sites
+    :type n: int
+    :param m: number of down spins (1 in 0, 1)
+    :type m: int
+    :return: index tensor
+    :rtype: Tensor
+    """
+    # m down spins
+    num_combinations = math.comb(n, m)
+    inds = np.zeros([num_combinations], dtype="int64")
+    if m == 0:
+        inds[0] = 0
+        return inds
+    combination = (1 << m) - 1
+
+    for i in range(num_combinations):
+        inds[i] = combination
+
+        # Find the next combination using Gosper's Hack
+        u = combination & -combination
+        v = u + combination
+        combination = v + (((v ^ combination) // u) >> 2)
+    return backend.convert_to_tensor(inds)
+
+
+def u1_mask(n: int, m: int) -> Tensor:
+    """
+    Return the 1d array of size 2**n filled with zero,
+    one only in elements corresponding to the m down spins
+
+    :param n: number of total sites
+    :type n: int
+    :param m: number of down spins (1 in 0, 1)
+    :type m: int
+    :return: _description_
+    :rtype: Tensor
+    """
+    inds = u1_inds(n, m)
+    m = backend.scatter(
+        backend.zeros([2**n]),
+        backend.reshape(inds, [-1, 1]),
+        backend.ones([math.comb(n, m)]),
+    )
+    return m
+
+
+def u1_project(s: Tensor, n: int, m: int) -> Tensor:
+    """
+    Project a state s to the subspace with m down spins in n sites
+
+    :param s: input state of size 2**n
+    :type s: Tensor
+    :param n: number of total sites
+    :type n: int
+    :param m: number of down spins (1 in 0, 1)
+    :type m: int
+    :return: projected state of size C_n^m
+    :rtype: Tensor
+    """
+    return backend.gather1d(s, u1_inds(n, m))
+
+
+def u1_enlarge(s: Tensor, n: int, m: int) -> Tensor:
+    """
+    Enlarge a state s in the subspace with m down spins in n sites to
+    the full Hilbert space wavefunction of size 2**n
+
+    :param s: input state of size C_n^m
+    :type s: Tensor
+    :param n: number of total sites
+    :type n: int
+    :param m: number of down spins (1 in 0, 1)
+    :type m: int
+    :return: enlarged state of size 2**n
+    :rtype: Tensor
+    """
+    inds = u1_inds(n, m)
+    return backend.scatter(backend.zeros([2**n]), backend.reshape(inds, [-1, 1]), s)
+
+
 def heisenberg_hamiltonian(
     g: Graph,
     hzz: float = 1.0,

tests/test_quantum.py

@@ -524,3 +524,26 @@ def test_reduced_wavefunction(backend):
     c1.h(0)
     c1.cnot(0, 1)
     np.testing.assert_allclose(s1, c1.state(), atol=1e-5)
+
+
+@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
+def test_u1_mask(backend):
+    g = tc.templates.graphs.Line1D(8)
+    sumz = tc.quantum.heisenberg_hamiltonian(g, hzz=0, hxx=0, hyy=0, hz=1)
+    for i in range(9):
+        s = tc.quantum.u1_mask(8, i)
+        s /= tc.backend.norm(s)
+        c = tc.Circuit(8, inputs=s)
+        zexp = tc.templates.measurements.operator_expectation(c, sumz)
+        np.testing.assert_allclose(zexp, 8 - 2 * i, atol=1e-6)
+
+
+@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
+def test_u1_project(backend):
+    c = tc.Circuit(8)
+    c.x([0, 2, 4])
+    c.exp1(0, 1, unitary=tc.gates._swap_matrix, theta=0.6)
+    s = c.state()
+    s1 = tc.quantum.u1_project(s, 8, 3)
+    assert s1.shape[-1] == 56
+    np.testing.assert_allclose(tc.quantum.u1_enlarge(s1, 8, 3), s)