fix controlled ops with dm

horqrux/apply.py

@@ -67,7 +67,7 @@ def apply_operator(
         return jnp.moveaxis(a=state, source=new_state_dims, destination=state_dims)
 
     # Apply operator to density matrix: ρ' = O ρ O†
-    support_perm = target + tuple(set(tuple(range(state.ndim // 2))) - set(target))
+    support_perm = state_dims + tuple(set(tuple(range(state.ndim // 2))) - set(state_dims))
     state = permute_basis(state, support_perm, False)
     state = jnp.tensordot(a=operator_reshaped, b=state, axes=(op_out_dims, new_state_dims))
 

tests/test_gates.py

@@ -52,6 +52,16 @@ def test_controlled_primitive(gate_fn: Callable) -> None:
         apply_operator(state, gate.dagger(), gate.target[0], gate.control[0]), orig_state
     )
 
+    # test density matrix is similar to pure state
+    dm = apply_operator(
+        density_mat(orig_state),
+        gate.unitary(),
+        gate.target[0],
+        gate.control[0],
+        is_state_densitymat=True,
+    )
+    assert jnp.allclose(dm, density_mat(state))
+
 
 @pytest.mark.parametrize("gate_fn", PARAMETRIC_GATES)
 def test_parametric(gate_fn: Callable) -> None:
@@ -89,6 +99,16 @@ def test_controlled_parametric(gate_fn: Callable) -> None:
         apply_operator(state, gate.dagger(values), gate.target[0], gate.control[0]), orig_state
     )
 
+    # test density matrix is similar to pure state
+    dm = apply_operator(
+        density_mat(orig_state),
+        gate.unitary(values),
+        gate.target[0],
+        gate.control[0],
+        is_state_densitymat=True,
+    )
+    assert jnp.allclose(dm, density_mat(state))
+
 
 @pytest.mark.parametrize(
     ["bitstring", "expected_state"],