make unitary private

horqrux/analog.py

@@ -21,7 +21,7 @@ class _HamiltonianEvolution(Primitive):
     target: QubitSupport
     control: QubitSupport
 
-    def unitary(self, values: dict[str, Array] = dict()) -> Array:
+    def _unitary(self, values: dict[str, Array] = dict()) -> Array:
         return expm(values["hamiltonian"] * (-1j * values["time_evolution"]))
 
 

horqrux/parametric.py

@@ -65,7 +65,7 @@ def __iter__(self) -> Iterable:
     def tree_unflatten(cls, aux_data: Any, children: Any) -> Any:
         return cls(*children, *aux_data)
 
-    def unitary(self, values: dict[str, float] = dict()) -> Array:
+    def _unitary(self, values: dict[str, float] = dict()) -> Array:
         return _unitary(OPERATIONS_DICT[self.generator_name], self.parse_values(values))
 
     def jacobian(self, values: dict[str, float] = dict()) -> Array:
@@ -141,7 +141,7 @@ def RZ(
 
 
 class _PHASE(Parametric):
-    def unitary(self, values: dict[str, float] = dict()) -> Array:
+    def _unitary(self, values: dict[str, float] = dict()) -> Array:
         u = jnp.eye(2, 2, dtype=default_dtype)
         u = u.at[(1, 1)].set(jnp.exp(1.0j * self.parse_values(values)))
         return u

horqrux/primitive.py

@@ -61,7 +61,7 @@ def tree_flatten(self) -> tuple[tuple, tuple[str, TargetQubits, ControlQubits, N
     def tree_unflatten(cls, aux_data: Any, children: Any) -> Any:
         return cls(*children, *aux_data)
 
-    def unitary(self, values: dict[str, float] = dict()) -> Array:
+    def _unitary(self, values: dict[str, float] = dict()) -> Array:
         """Obtain the base unitary from `generator_name`.
 
         Args:
@@ -81,7 +81,7 @@ def dagger(self, values: dict[str, float] = dict()) -> Array:
         Returns:
             Array: The base unitary daggered from `generator_name`.
         """
-        return _dagger(self.unitary(values))
+        return _dagger(self._unitary(values))
 
     def tensor(self, values: dict[str, float] = dict()) -> Array:
         """Obtain the unitary taking into account the qubit support for controlled operations.
@@ -92,7 +92,7 @@ def tensor(self, values: dict[str, float] = dict()) -> Array:
         Returns:
             Array: Unitary representation taking into account the qubit support.
         """
-        base_unitary = self.unitary(values)
+        base_unitary = self._unitary(values)
         if is_controlled(self.control):
             return controlled(base_unitary, self.target, self.control)
         return base_unitary

horqrux/shots.py

@@ -29,7 +29,7 @@ def to_matrix(
         observable.control == observable.parse_idx(none_like(observable.target)),
         "Controlled gates cannot be promoted from observables to operations on the whole state vector",
     )
-    unitary = observable.unitary(values=values)
+    unitary = observable._unitary(values=values)
     target = observable.target[0][0]
     identity = jnp.eye(2, dtype=unitary.dtype)
     ops = [identity for _ in range(n_qubits)]

horqrux/utils.py

@@ -99,7 +99,7 @@ def list(cls) -> list[str]:
 
 
 class OperationType(StrEnum):
-    UNITARY = "unitary"
+    UNITARY = "_unitary"
     DAGGER = "dagger"
     JACOBIAN = "jacobian"
 

tests/test_gates.py

@@ -10,7 +10,7 @@
 from horqrux.apply import apply_gate, apply_operator
 from horqrux.parametric import PHASE, RX, RY, RZ
 from horqrux.primitive import NOT, SWAP, H, I, S, T, X, Y, Z
-from horqrux.utils import density_mat, equivalent_state, product_state, random_state
+from horqrux.utils import OperationType, density_mat, equivalent_state, product_state, random_state
 
 MAX_QUBITS = 7
 PARAMETRIC_GATES = (RX, RY, RZ, PHASE)
@@ -31,7 +31,7 @@ def test_primitive(gate_fn: Callable) -> None:
     # test density matrix is similar to pure state
     dm = apply_operator(
         density_mat(orig_state),
-        gate.unitary(),
+        gate._unitary(),
         gate.target[0],
         gate.control[0],
     )
@@ -54,7 +54,7 @@ def test_controlled_primitive(gate_fn: Callable) -> None:
     # test density matrix is similar to pure state
     dm = apply_operator(
         density_mat(orig_state),
-        gate.unitary(),
+        gate._unitary(),
         gate.target[0],
         gate.control[0],
     )
@@ -75,7 +75,7 @@ def test_parametric(gate_fn: Callable) -> None:
     # test density matrix is similar to pure state
     dm = apply_operator(
         density_mat(orig_state),
-        gate.unitary(values),
+        gate._unitary(values),
         gate.target[0],
         gate.control[0],
     )
@@ -99,7 +99,7 @@ def test_controlled_parametric(gate_fn: Callable) -> None:
     # test density matrix is similar to pure state
     dm = apply_operator(
         density_mat(orig_state),
-        gate.unitary(values),
+        gate._unitary(values),
         gate.target[0],
         gate.control[0],
     )
@@ -149,10 +149,20 @@ def test_merge_gates() -> None:
         "c": np.random.uniform(0.1, 2 * np.pi),
     }
     state_grouped = apply_gate(
-        product_state("0000"), gates, values, "unitary", group_gates=True, merge_ops=True
+        product_state("0000"),
+        gates,
+        values,
+        OperationType.UNITARY,
+        group_gates=True,
+        merge_ops=True,
     )
     state = apply_gate(
-        product_state("0000"), gates, values, "unitary", group_gates=False, merge_ops=False
+        product_state("0000"),
+        gates,
+        values,
+        OperationType.UNITARY,
+        group_gates=False,
+        merge_ops=False,
     )
     assert jnp.allclose(state_grouped, state)