add more tests

tests/helpers.py

@@ -47,16 +47,25 @@ def calc_mat_vec_wavefunction(
     )
 
 
-def get_op_support(op: type[Primitive] | type[Parametric], n_qubits: int) -> tuple:
+def get_op_support(
+    op: type[Primitive] | type[Parametric], n_qubits: int, get_ordered: bool = False
+) -> tuple:
     """Decides a random qubit support for any gate, up to a some max n_qubits."""
     if op in OPS_1Q.union(OPS_PARAM_1Q):
         supp: tuple = (random.randint(0, n_qubits - 1),)
+        ordered_supp = supp
     elif op in OPS_2Q.union(OPS_PARAM_2Q):
         supp = tuple(random.sample(range(n_qubits), 2))
+        ordered_supp = tuple(sorted(supp))
     elif op in OPS_3Q:
         i, j, k = tuple(random.sample(range(n_qubits), 3))
+        a, b, c = tuple(sorted((i, j, k)))
         supp = ((i, j), k) if op == Toffoli else (i, (j, k))
-    return supp
+        ordered_supp = ((a, b), c) if op == Toffoli else (a, (b, c))
+    if get_ordered:
+        return supp, ordered_supp
+    else:
+        return supp
 
 
 def random_pauli_hamiltonian(

tests/test_tensor.py

@@ -10,17 +10,29 @@
 from pyqtorch.hamiltonians import GeneratorType, HamiltonianEvolution
 from pyqtorch.primitives import (
     CNOT,
+    CPHASE,
+    CRX,
+    CRY,
+    CRZ,
+    CSWAP,
+    CY,
+    CZ,
+    OPS_2Q,
+    OPS_3Q,
     OPS_DIGITAL,
     OPS_PARAM,
+    OPS_PARAM_2Q,
     SWAP,
     N,
     Parametric,
     Primitive,
     Projector,
+    Toffoli,
 )
 from pyqtorch.utils import (
     ATOL,
     RTOL,
+    permute_basis,
     random_state,
 )
 
@@ -240,3 +252,57 @@ def test_hevo_tensor_tensor(
     psi_star = operator(psi_init, values)
     psi_expected = calc_mat_vec_wavefunction(operator, psi_init, values, full_support)
     assert torch.allclose(psi_star, psi_expected, rtol=RTOL, atol=ATOL)
+
+
+@pytest.mark.parametrize("n_qubits", [3, 5])
+def test_permute_tensor(n_qubits: int) -> None:
+    for op in OPS_2Q.union(OPS_3Q):
+        supp, ordered_supp = get_op_support(op, n_qubits, get_ordered=True)
+
+        op_concrete1 = op(*supp)
+        op_concrete2 = op(*ordered_supp)
+
+        mat1 = op_concrete1.tensor()
+        mat2 = op_concrete2.tensor()
+
+        perm = op_concrete1._qubit_support.qubits
+
+        assert torch.allclose(mat1, permute_basis(mat2, perm))
+        assert torch.allclose(mat2, permute_basis(mat1, perm, inv=True))
+
+
+@pytest.mark.parametrize("n_qubits", [3, 5])
+@pytest.mark.parametrize("batch_size", [1, 5])
+def test_permute_tensor_parametric(n_qubits: int, batch_size: int) -> None:
+    for op in OPS_PARAM_2Q:
+        supp, ordered_supp = get_op_support(op, n_qubits, get_ordered=True)
+        params = [f"{op.__name__}_th{i}" for i in range(op.n_params)]
+        values = {param: torch.rand(batch_size) for param in params}
+
+        op_concrete1 = op(*supp, *params)
+        op_concrete2 = op(*ordered_supp, *params)
+
+        mat1 = op_concrete1.tensor(values=values)
+        mat2 = op_concrete2.tensor(values=values)
+
+        perm = op_concrete1._qubit_support.qubits
+
+        assert torch.allclose(mat1, permute_basis(mat2, perm))
+        assert torch.allclose(mat2, permute_basis(mat1, perm, inv=True))
+
+
+def test_tensor_symmetries() -> None:
+    assert not torch.allclose(CNOT(0, 1).tensor(), CNOT(1, 0).tensor())
+    assert not torch.allclose(CY(0, 1).tensor(), CY(1, 0).tensor())
+    assert not torch.allclose(CZ(0, 1).tensor(), CY(1, 0).tensor())
+    assert not torch.allclose(CRX(0, 1, 1.0).tensor(), CRX(1, 0, 1.0).tensor())
+    assert not torch.allclose(CRY(0, 1, 1.0).tensor(), CRY(1, 0, 1.0).tensor())
+    assert not torch.allclose(CRZ(0, 1, 1.0).tensor(), CRZ(1, 0, 1.0).tensor())
+    assert torch.allclose(CPHASE(0, 1, 1.0).tensor(), CPHASE(1, 0, 1.0).tensor())
+    assert torch.allclose(SWAP(0, 1).tensor(), SWAP(1, 0).tensor())
+    assert torch.allclose(CSWAP(0, (1, 2)).tensor(), CSWAP(0, (2, 1)).tensor())
+    assert torch.allclose(CSWAP(1, (0, 2)).tensor(), CSWAP(1, (2, 0)).tensor())  # F
+    assert torch.allclose(CSWAP(2, (0, 1)).tensor(), CSWAP(2, (1, 0)).tensor())  # F
+    assert torch.allclose(Toffoli((0, 1), 2).tensor(), Toffoli((1, 0), 2).tensor())
+    assert torch.allclose(Toffoli((0, 2), 1).tensor(), Toffoli((2, 0), 1).tensor())  # F
+    assert torch.allclose(Toffoli((1, 2), 0).tensor(), Toffoli((2, 1), 0).tensor())  # F