Reduced density matrix to pyqrack (#444)

Co-authored-by: David Plankensteiner <[email protected]>

Author: jon-wurtz

src/bloqade/pyqrack/device.py

@@ -4,6 +4,7 @@
 import numpy as np
 from kirin import ir
 from kirin.passes import fold
+from kirin.dialects.ilist import IList
 
 from bloqade.squin import noise as squin_noise
 from pyqrack.pauli import Pauli
@@ -18,6 +19,7 @@
     _default_pyqrack_args,
 )
 from bloqade.pyqrack.task import PyQrackSimulatorTask
+from pyqrack.qrack_simulator import QrackSimulator
 from bloqade.squin.noise.rewrite import RewriteNoiseStmts
 from bloqade.analysis.address.lattice import AnyAddress
 from bloqade.analysis.address.analysis import AddressAnalysis
@@ -26,6 +28,55 @@
 Params = ParamSpec("Params")
 
 
+def _pyqrack_reduced_density_matrix(
+    inds: tuple[int, ...], sim_reg: QrackSimulator, tol: float = 1e-12
+) -> "np.linalg._linalg.EighResult":
+    """
+    Extract the reduced density matrix representing the state of a list
+    of qubits from a PyQRack simulator register.
+
+    Inputs:
+        inds: A list of integers labeling the qubit registers to extract the reduced density matrix for
+        sim_reg: The PyQRack simulator register to extract the reduced density matrix from
+        tol: The tolerance for density matrix eigenvalues to be considered non-zero.
+    Outputs:
+        An eigh result containing the eigenvalues and eigenvectors of the reduced density matrix.
+    """
+    # Identify the rest of the qubits in the register
+    N = sim_reg.num_qubits()
+    other = tuple(set(range(N)).difference(inds))
+
+    if len(set(inds)) != len(inds):
+        raise ValueError("Qubits must be unique.")
+
+    if max(inds) > N - 1:
+        raise ValueError(
+            f"Qubit indices {inds} exceed the number of qubits in the register {N}."
+        )
+
+    reordering = inds + other
+    # Fix pyqrack edannes to be consistent with Cirq.
+    reordering = tuple(N - 1 - x for x in reordering)
+    # Extract the statevector from the PyQRack qubits
+    statevector = np.array(sim_reg.out_ket())
+    # Reshape into a (2,2,2, ..., 2) tensor
+    vec_f = np.reshape(statevector, (2,) * N)
+    # Reorder the indexes to obey the order of the qubits
+    vec_p = np.transpose(vec_f, reordering)
+    # Rehape into a 2^N by 2^M matrix to compute the singular value decomposition
+    vec_svd = np.reshape(vec_p, (2 ** len(inds), 2 ** len(other)))
+    # The singular values and vectors are the eigenspace of the reduced density matrix
+    s, v, d = np.linalg.svd(vec_svd, full_matrices=False)
+
+    # Remove the negligable singular values
+    nonzero_inds = np.where(np.abs(v) > tol)[0]
+    s = s[:, nonzero_inds]
+    v = v[nonzero_inds] ** 2
+    # Forge into the correct result type
+    result = np.linalg._linalg.EighResult(eigenvalues=v, eigenvectors=s)
+    return result
+
+
 @dataclass
 class PyQrackSimulatorBase(AbstractSimulatorDevice[PyQrackSimulatorTask]):
     """PyQrack simulation device base class."""
@@ -50,7 +101,6 @@ def new_task(
         kwargs: dict[str, Any],
         memory: MemoryType,
     ) -> PyQrackSimulatorTask[Params, RetType, MemoryType]:
-
         if squin_noise in mt.dialects:
             # NOTE: rewrite noise statements
             mt_ = mt.similar(mt.dialects)
@@ -112,6 +162,51 @@ def pauli_expectation(pauli: list[Pauli], qubits: list[PyQrackQubit]) -> float:
 
         return sim_reg.pauli_expectation(qubit_ids, pauli)
 
+    @staticmethod
+    def quantum_state(
+        qubits: list[PyQrackQubit] | IList[PyQrackQubit, Any], tol: float = 1e-12
+    ) -> "np.linalg._linalg.EighResult":
+        """
+        Extract the reduced density matrix representing the state of a list
+        of qubits from a PyQRack simulator register.
+
+        Inputs:
+            qubits: A list of PyQRack qubits to extract the reduced density matrix for
+            tol: The tolerance for density matrix eigenvalues to be considered non-zero.
+        Outputs:
+            An eigh result containing the eigenvalues and eigenvectors of the reduced density matrix.
+        """
+        if len(qubits) == 0:
+            return np.linalg._linalg.EighResult(
+                eigenvalues=np.array([]), eigenvectors=np.array([]).reshape(0, 0)
+            )
+        sim_reg = qubits[0].sim_reg
+
+        if not all([x.sim_reg is sim_reg for x in qubits]):
+            raise ValueError("All qubits must be from the same simulator register.")
+        inds: tuple[int, ...] = tuple(qubit.addr for qubit in qubits)
+
+        return _pyqrack_reduced_density_matrix(inds, sim_reg, tol)
+
+    @classmethod
+    def reduced_density_matrix(
+        cls, qubits: list[PyQrackQubit] | IList[PyQrackQubit, Any], tol: float = 1e-12
+    ) -> np.ndarray:
+        """
+        Extract the reduced density matrix representing the state of a list
+        of qubits from a PyQRack simulator register.
+
+        Inputs:
+            qubits: A list of PyQRack qubits to extract the reduced density matrix for
+            tol: The tolerance for density matrix eigenvalues to be considered non-zero.
+        Outputs:
+            A dense 2^n x 2^n numpy array representing the reduced density matrix.
+        """
+        rdm = cls.quantum_state(qubits, tol)
+        return np.einsum(
+            "ax,x,bx", rdm.eigenvectors, rdm.eigenvalues, rdm.eigenvectors.conj()
+        )
+
 
 @dataclass
 class StackMemorySimulator(PyQrackSimulatorBase):

src/bloqade/pyqrack/task.py

@@ -2,6 +2,7 @@
 from dataclasses import dataclass
 
 from bloqade.task import AbstractSimulatorTask
+from bloqade.pyqrack.reg import QubitState, PyQrackQubit
 from bloqade.pyqrack.base import (
     MemoryABC,
     PyQrackInterpreter,
@@ -36,3 +37,17 @@ def state_vector(self) -> list[complex]:
         """Returns the state vector of the simulator."""
         self.run()
         return self.state.sim_reg.out_ket()
+
+    def qubits(self) -> list[PyQrackQubit]:
+        """Returns the qubits in the simulator."""
+        try:
+            N = self.state.sim_reg.num_qubits()
+            return [
+                PyQrackQubit(
+                    addr=i, sim_reg=self.state.sim_reg, state=QubitState.Active
+                )
+                for i in range(N)
+            ]
+        except AttributeError:
+            Warning("Task has not been run, there are no qubits!")
+            return []

test/pyqrack/runtime/test_qrack.py

@@ -1,10 +1,13 @@
 import math
 from unittest.mock import Mock, call
 
+import cirq
+import numpy as np
 from kirin import ir
 
-from bloqade import qasm2
+from bloqade import qasm2, squin
 from pyqrack.pauli import Pauli
+from bloqade.pyqrack import StackMemorySimulator
 from bloqade.pyqrack.base import MockMemory, PyQrackInterpreter
 
 
@@ -91,7 +94,6 @@ def program():
 
 
 def test_basic_control_gates():
-
     @qasm2.main
     def program():
         q = qasm2.qreg(3)
@@ -162,3 +164,205 @@ def program():
             call.r(3, 0.5, 1),
         ]
     )
+
+
+def test_rdm1():
+    """
+    Is extracting the exact state vector consistent with cirq?
+    This test also validates the ordering of the qubit basis.
+    """
+
+    @squin.kernel
+    def program():
+        q = squin.qubit.new(5)
+        squin.gate.h(q[1])
+        return q
+
+    emulator = StackMemorySimulator(min_qubits=6)
+    task = emulator.task(program)
+    qubits = task.run()
+    rho = emulator.quantum_state(qubits)
+
+    assert all(np.isclose(rho.eigenvalues, [1]))
+
+    circuit = cirq.Circuit()
+    qbs = cirq.LineQubit.range(5)
+    circuit.append(cirq.H(qbs[1]))
+    for i in range(5):
+        circuit.append(cirq.I(qbs[i]))
+    state = cirq.Simulator().simulate(circuit).state_vector()
+    assert cirq.equal_up_to_global_phase(state, rho[1][:, 0])
+
+
+def test_rdm1b():
+    """
+    Is extracting the exact state vector consistent with cirq?
+    This test also validates the ordering of the qubit basis.
+    Same as test_rdm1, but with the total qubits equal to the number of qubits in the program.
+    """
+
+    @squin.kernel
+    def program():
+        q = squin.qubit.new(5)
+        squin.gate.h(q[1])
+        return q
+
+    emulator = StackMemorySimulator(min_qubits=5)
+    task = emulator.task(program)
+    qubits = task.run()
+    rho = emulator.quantum_state(qubits)
+
+    assert all(np.isclose(rho.eigenvalues, [1]))
+
+    circuit = cirq.Circuit()
+    qbs = cirq.LineQubit.range(5)
+    circuit.append(cirq.H(qbs[1]))
+    for i in range(5):
+        circuit.append(cirq.I(qbs[i]))
+    state = cirq.Simulator().simulate(circuit).state_vector()
+    assert cirq.equal_up_to_global_phase(state, rho[1][:, 0])
+
+
+def test_rdm2():
+    """
+    Does the RDM project correctly?
+    """
+
+    @squin.kernel
+    def program():
+        """
+        Creates a GHZ state on qubits 0,1,3,4 on a total of 6 qubits.
+        """
+        q = squin.qubit.new(6)
+        squin.gate.h(q[0])
+        squin.gate.cx(q[0], q[1])
+        squin.gate.cx(q[0], q[3])
+        squin.gate.cx(q[0], q[4])
+        return q
+
+    emulator = StackMemorySimulator(min_qubits=6)
+    task = emulator.task(program)
+    qubits = task.run()
+    rho = emulator.quantum_state([qubits[x] for x in [0, 1, 3, 4]])
+    target = np.array([1] + [0] * (14) + [1]) / np.sqrt(2) + 0j
+    assert all(np.isclose(rho.eigenvalues, [1]))
+    assert cirq.equal_up_to_global_phase(rho[1][:, 0], target)
+
+    rho2 = emulator.quantum_state([qubits[x] for x in [0, 1, 3]])
+    assert all(np.isclose(rho2.eigenvalues, [0.5, 0.5]))
+    assert rho2.eigenvectors.shape == (8, 2)
+
+
+def test_rdm3():
+    """
+    Out-of-order indexing is consistent with cirq.
+    """
+
+    @squin.kernel
+    def program():
+        """
+        Random unitaries on 3 qubits.
+        """
+        q = squin.qubit.new(3)
+        squin.gate.rx(0.1, q[0])
+        squin.gate.ry(0.2, q[1])
+        squin.gate.rx(0.3, q[2])
+        return q
+
+    emulator = StackMemorySimulator(min_qubits=6)
+    task = emulator.task(program)
+    qubits = task.run()
+
+    # Canonical ordering
+    rho = emulator.quantum_state([qubits[x] for x in [0, 1, 2]])
+    circuit = cirq.Circuit()
+    qbs = cirq.LineQubit.range(3)
+    circuit.append(cirq.rx(0.1)(qbs[0]))
+    circuit.append(cirq.ry(0.2)(qbs[1]))
+    circuit.append(cirq.rx(0.3)(qbs[2]))
+    state = cirq.Simulator().simulate(circuit).state_vector()
+    assert cirq.equal_up_to_global_phase(state, rho[1][:, 0])
+
+    # Reverse ordering 0->2, 1->, 2->0
+    rho = emulator.quantum_state([qubits[x] for x in [2, 1, 0]])
+    circuit = cirq.Circuit()
+    qbs = cirq.LineQubit.range(3)
+    circuit.append(cirq.rx(0.1)(qbs[2]))
+    circuit.append(cirq.ry(0.2)(qbs[1]))
+    circuit.append(cirq.rx(0.3)(qbs[0]))
+    state = cirq.Simulator().simulate(circuit).state_vector()
+    assert cirq.equal_up_to_global_phase(state, rho[1][:, 0])
+
+    # Other ordering
+    rho = emulator.quantum_state([qubits[x] for x in [1, 2, 0]])
+    circuit = cirq.Circuit()
+    qbs = cirq.LineQubit.range(3)
+    circuit.append(cirq.rx(0.1)(qbs[2]))
+    circuit.append(cirq.ry(0.2)(qbs[0]))
+    circuit.append(cirq.rx(0.3)(qbs[1]))
+    state = cirq.Simulator().simulate(circuit).state_vector()
+    assert cirq.equal_up_to_global_phase(state, rho[1][:, 0])
+
+
+def test_rdm4():
+    rho = StackMemorySimulator.quantum_state([])
+    assert rho.eigenvalues.shape == (0,)
+    assert rho.eigenvectors.shape == (0, 0)
+
+
+def test_rdm5():
+    @squin.kernel
+    def program():
+        """
+        Random unitaries on 3 qubits.
+        """
+        q = squin.qubit.new(3)
+        return q
+
+    emulator = StackMemorySimulator(min_qubits=6)
+    task = emulator.task(program)
+    qubits = task.run()
+    rho = emulator.reduced_density_matrix(qubits)
+    assert rho.shape == (8, 8)
+
+
+def test_rdm_failures():
+    @squin.kernel
+    def program():
+        q = squin.qubit.new(3)
+        return q
+
+    emulator = StackMemorySimulator(min_qubits=6)
+    task = emulator.task(program)
+    qbsA = task.qubits()
+    qubits = task.run()
+    qubits2 = task.run()
+    qbsB = task.qubits()
+    assert len(qbsA) == 0
+    assert len(qbsB) == 6
+
+    try:
+        emulator.quantum_state([qubits[0], qubits[0]])
+        assert False, "Should have failed; qubits must be unique"
+    except ValueError as e:
+        assert str(e) == "Qubits must be unique."
+
+    try:
+        emulator.quantum_state([qubits[0], qubits2[1]])
+        assert False, "Should have failed; qubits must be from the same register"
+    except ValueError as e:
+        assert str(e) == "All qubits must be from the same simulator register."
+
+
+def test_get_qubits():
+    @squin.kernel
+    def program():
+        q = squin.qubit.new(3)
+        return q
+
+    emulator = StackMemorySimulator(min_qubits=6)
+    task = emulator.task(program)
+    task.run()
+
+    qubits2 = task.qubits()
+    assert len(qubits2) == 6