[Feature] Add density matrix support for initial state (#640)

docs/content/state_init.md

@@ -122,6 +122,22 @@ final_state = run(CNOT(0, 1), state=init_state)
 print(f"Final state = {final_state}") # markdown-exec: hide
 ```
 
+## Density matrices conversion
+
+It is also possible to obtain density matrices from statevectors. They can be passed as inputs to quantum programs performing density matrix based operations such as noisy simulations, when the backend allows such as PyQTorch.
+
+```python exec="on" source="material-block" result="json" session="states"
+from qadence import product_state, density_mat
+
+init_state = product_state("10")
+init_density_matrix = density_mat(init_state)
+print(f"Initial = {init_density_matrix}") # markdown-exec: hide
+
+final_density_matrix = run(CNOT(0, 1), state=init_density_matrix)
+print(f"Final = {final_density_matrix}") # markdown-exec: hide
+
+```
+
 ## Block initialization
 
 Not all backends support custom statevector initialization, however previous utility functions have their counterparts to initialize the respective blocks:

qadence/backends/utils.py

@@ -110,9 +110,23 @@ def to_list_of_dicts(param_values: ParamDictType) -> list[ParamDictType]:
 
 
 def pyqify(state: Tensor, n_qubits: int = None) -> ArrayLike:
-    """Convert a state of shape (batch_size, 2**n_qubits) to [2] * n_qubits + [batch_size]."""
+    """Convert a state of shape (batch_size, 2**n_qubits) to [2] * n_qubits + [batch_size].
+
+    Or set the batch_size of a density matrix as the last dimension for PyQTorch.
+    """
     if n_qubits is None:
         n_qubits = int(log2(state.shape[1]))
+    if isinstance(state, DensityMatrix):
+        if (
+            len(state.shape) != 3
+            or (state.shape[1] != 2**n_qubits)
+            or (state.shape[1] != state.shape[2])
+        ):
+            raise ValueError(
+                "The initial state must be composed of tensors/arrays of size "
+                f"(batch_size, 2**n_qubits, 2**n_qubits). Found: {state.shape = }."
+            )
+        return torch.einsum("kij->ijk", state)
     if len(state.shape) != 2 or (state.shape[1] != 2**n_qubits):
         raise ValueError(
             "The initial state must be composed of tensors/arrays of size "

qadence/states.py

@@ -6,6 +6,7 @@
 
 import torch
 from numpy.typing import ArrayLike
+from pyqtorch.utils import DensityMatrix
 from torch import Tensor, concat
 from torch.distributions import Categorical, Distribution
 
@@ -37,6 +38,8 @@
     "is_normalized",
     "rand_bitstring",
     "equivalent_state",
+    "DensityMatrix",
+    "density_mat",
 ]
 
 ATOL_64 = 1e-14  # 64 bit precision
@@ -319,6 +322,24 @@ def random_state(
     return state
 
 
+# DENSITY MATRIX
+
+
+def density_mat(state: Tensor) -> DensityMatrix:
+    """
+    Computes the density matrix from a pure state vector.
+
+    Arguments:
+        state: The pure state vector :math:`|\\psi\\rangle`.
+
+    Returns:
+        Tensor: The density matrix :math:`\\rho = |\psi \\rangle \\langle\\psi|`.
+    """
+    if isinstance(state, DensityMatrix):
+        return state
+    return DensityMatrix(torch.einsum("bi,bj->bij", (state, state.conj())))
+
+
 # BLOCKS
 
 

tests/backends/pyq/test_quantum_pyq.py

@@ -40,6 +40,7 @@
     CRX,
     CRY,
     CRZ,
+    CZ,
     RX,
     RY,
     RZ,
@@ -56,7 +57,7 @@
     Z,
 )
 from qadence.parameters import FeatureParameter, Parameter
-from qadence.states import random_state, uniform_state, zero_state
+from qadence.states import DensityMatrix, density_mat, random_state, uniform_state, zero_state
 from qadence.transpile import set_trainable
 from qadence.types import PI, BackendName, DiffMode
 from qadence.utils import P0, P1
@@ -168,6 +169,14 @@ def test_raise_error_for_ill_dimensioned_initial_state() -> None:
         backend.run(backend.circuit(circuit), state=initial_state)
 
 
+def test_raise_error_for_ill_dimensioned_density_matrix() -> None:
+    circuit = QuantumCircuit(2, X(0) @ X(1))
+    backend = Backend()
+    initial_state = DensityMatrix(torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=torch.complex128))
+    with pytest.raises(ValueError):
+        backend.run(backend.circuit(circuit), state=initial_state)
+
+
 @pytest.mark.parametrize(
     "gate, state",
     [
@@ -192,6 +201,13 @@ def test_run_with_nonparametric_single_qubit_gates(
     wf = backend.run(pyqtorch_circ, state=initial_state)
     assert torch.allclose(wf, state)
 
+    # test with density matrix
+    initial_state = density_mat(initial_state)
+    dm = backend.run(pyqtorch_circ, state=initial_state)
+    assert isinstance(dm, DensityMatrix)
+    expected_dm = density_mat(state.unsqueeze(0) if len(state.shape) == 1 else state)
+    assert torch.allclose(dm, expected_dm)
+
 
 @pytest.mark.parametrize(
     "gate, matrix",
@@ -245,10 +261,20 @@ def test_run_with_nonparametric_single_qubit_gates_and_random_initial_state(
     theta2 = random.uniform(0.0, 2.0 * PI)
     complex2 = complex(np.cos(theta2), np.sin(theta2))
     initial_state = torch.tensor([[complex1, complex2]], dtype=torch.complex128)
-    wf = backend.run(backend.circuit(circuit), state=initial_state)
+    pyqtorch_circ = backend.circuit(circuit)
+    wf = backend.run(pyqtorch_circ, state=initial_state)
     expected_state = torch.matmul(matrix, initial_state[0])
     assert torch.allclose(wf, expected_state)
 
+    # test with density matrix
+    initial_state = density_mat(initial_state)
+    dm = backend.run(pyqtorch_circ, state=initial_state)
+    assert isinstance(dm, DensityMatrix)
+    expected_dm = density_mat(
+        expected_state.unsqueeze(0) if len(expected_state.shape) == 1 else expected_state
+    )
+    assert torch.allclose(dm, expected_dm)
+
 
 @pytest.mark.parametrize(
     "parametric_gate, state",
@@ -355,6 +381,15 @@ def test_run_with_parametric_single_qubit_gates_and_random_initial_state(
     expected_state = torch.matmul(matrix, initial_state[0])
     assert torch.allclose(wf, expected_state)
 
+    # test with density matrix
+    initial_state = density_mat(initial_state)
+    dm = backend.run(pyqtorch_circ, embed(params, {}), state=initial_state)
+    assert isinstance(dm, DensityMatrix)
+    expected_dm = density_mat(
+        expected_state.unsqueeze(0) if len(expected_state.shape) == 1 else expected_state
+    )
+    assert torch.allclose(dm, expected_dm)
+
 
 @pytest.mark.parametrize(
     "parametric_gate, state",
@@ -395,6 +430,13 @@ def test_run_with_parametric_two_qubit_gates(
     wf = backend.run(pyqtorch_circ, embed(params, {}), state=initial_state)
     assert torch.allclose(wf, state)
 
+    # test with density matrix
+    initial_state = density_mat(initial_state)
+    dm = backend.run(pyqtorch_circ, embed(params, {}), state=initial_state)
+    assert isinstance(dm, DensityMatrix)
+    expected_dm = density_mat(state.unsqueeze(0) if len(state.shape) == 1 else state)
+    assert torch.allclose(dm, expected_dm)
+
 
 @pytest.mark.parametrize(
     "parametric_gate, matrix",
@@ -463,6 +505,15 @@ def test_run_with_parametric_two_qubit_gates_and_random_state(
     expected_state = torch.matmul(matrix, initial_state[0])
     assert torch.allclose(wf, expected_state)
 
+    # test with density matrix
+    initial_state = density_mat(initial_state)
+    dm = backend.run(pyqtorch_circ, embed(params, {}), state=initial_state)
+    assert isinstance(dm, DensityMatrix)
+    expected_dm = density_mat(
+        expected_state.unsqueeze(0) if len(expected_state.shape) == 1 else expected_state
+    )
+    assert torch.allclose(dm, expected_dm)
+
 
 @pytest.mark.parametrize(
     "gate, state",
@@ -539,13 +590,18 @@ def test_expectation_with_pauli_gates_and_random_state(
     expectation_value = backend.expectation(
         pyqtorch_circ, pyqtorch_obs, embed(params, {}), state=initial_state
     )
+    expectation_value_init_dm = backend.expectation(
+        pyqtorch_circ, pyqtorch_obs, embed(params, {}), state=density_mat(initial_state)
+    )
+
     Z_matrix = torch.tensor(
         [[1.0 + 0.0j, 0.0 + 0.0j], [0.0 + 0.0j, -1.0 + 0.0j]], dtype=torch.complex128
     )
     final_state = torch.matmul(Z_matrix, torch.matmul(matrix, initial_state[0]))
     probas = torch.square(torch.abs(final_state))
     expected_value = probas[0] - probas[1]
     assert torch.allclose(expectation_value, expected_value)
+    assert torch.allclose(expectation_value, expectation_value_init_dm)
 
 
 @pytest.mark.flaky(max_runs=5)
@@ -594,7 +650,7 @@ def test_controlled_rotation_gates_with_heterogeneous_parameters() -> None:
     [
         X(0),
         RZ(1, 0.5),
-        # CRY(0,1,0.2) write proper test for this
+        CRY(0, 1, 0.2),
     ],
 )
 def test_scaled_operation(block: AbstractBlock) -> None:
@@ -631,10 +687,10 @@ def test_scaled_featureparam_batching(batch_size: int) -> None:
         X(0),
         Y(0),
         Z(0),
-        # S(0), # TODO implement SDagger in PyQ
+        S(0),
         # T(0), # TODO implement TDagger in PyQ
         CNOT(0, 1),
-        # CZ(0, 1), # TODO implement CZ in PyQ?
+        CZ(0, 1),
         SWAP(0, 1),
         H(0),
         I(0),

tests/qadence/test_states.py

@@ -4,11 +4,15 @@
 
 import numpy as np
 import pytest
+import torch
 from torch import Tensor
 
 from qadence.backends.jax_utils import jarr_to_tensor
+from qadence.backends.utils import pyqify
 from qadence.execution import run
 from qadence.states import (
+    DensityMatrix,
+    density_mat,
     equivalent_state,
     ghz_block,
     ghz_state,
@@ -69,3 +73,25 @@ def test_product_state(n_qubits: int, backend: str) -> None:
     assert is_normalized(state_direct)
     assert is_normalized(state_block)
     assert equivalent_state(state_direct, state_block)
+
+
+def test_density_mat() -> None:
+    state_direct = product_state("00")
+    state_dm = density_mat(state_direct)
+    assert len(state_dm.shape) == 3
+    assert isinstance(state_dm, DensityMatrix)
+    assert state_dm.shape[0] == 1
+    assert state_dm.shape[1] == state_dm.shape[2] == 4
+
+    state_dm2 = density_mat(state_dm)
+    assert isinstance(state_dm2, DensityMatrix)
+    assert torch.allclose(state_dm2, state_dm)
+
+    with pytest.raises(ValueError):
+        pyqify(state_dm2.unsqueeze(0))
+
+    with pytest.raises(ValueError):
+        pyqify(state_dm2.view((1, 2, 8)))
+
+    with pytest.raises(ValueError):
+        pyqify(state_dm2.view((2, 4, 2)))