Merge pull request #521 from qiboteam/staterepr

State in computational basis

src/qibo/abstractions/states.py

@@ -68,6 +68,28 @@ def tensor(self, x):
                                     "".format(len(x), self.nqubits))
         self._tensor = x
 
+    @abstractmethod
+    def symbolic(self, decimals=5, cutoff=1e-10, max_terms=20):  # pragma: no cover
+        """Dirac notation representation of the state in the computational basis.
+
+        Args:
+            decimals (int): Number of decimals for the amplitudes.
+                Default is 5.
+            cutoff (float): Amplitudes with absolute value smaller than the
+                cutoff are ignored from the representation.
+                Default is 1e-10.
+            max_terms (int): Maximum number of terms to print. If the state
+                contains more terms they will be ignored.
+                Default is 20.
+
+        Returns:
+            A string representing the state in the computational basis.
+        """
+        raise_error(NotImplementedError)
+
+    def __repr__(self):
+        return self.symbolic()
+
     @abstractmethod
     def __array__(self): # pragma: no cover
         """State's tensor representation as an array."""
@@ -78,14 +100,31 @@ def numpy(self): # pragma: no cover
         """State's tensor representation as a numpy array."""
         raise_error(NotImplementedError)
 
-    def state(self, numpy=False):
+    def state(self, numpy=False, decimals=-1, cutoff=1e-10, max_terms=20):
         """State's tensor representation as an backend tensor.
 
         Args:
             numpy (bool): If ``True`` the returned tensor will be a numpy array,
                 otherwise it will follow the backend tensor type.
                 Default is ``False``.
+            decimals (int): If positive the Diract representation of the state
+                in the computational basis will be returned as a string.
+                ``decimals`` will be the number of decimals of each amplitude.
+                Default is -1.
+            cutoff (float): Amplitudes with absolute value smaller than the
+                cutoff are ignored from the Dirac representation.
+                Ignored if ``decimals < 0``. Default is 1e-10.
+            max_terms (int): Maximum number of terms in the Dirac representation.
+                If the state contains more terms they will be ignored.
+                Ignored if ``decimals < 0``. Default is 20.
+
+        Returns:
+            If ``decimals < 0`` a tensor representing the state in the computational
+            basis, otherwise a string with the Dirac representation of the state
+            in the computational basis.
         """
+        if decimals >= 0:
+            return self.symbolic(decimals, cutoff, max_terms)
         if numpy:
             return self.numpy()
         return self.tensor

src/qibo/core/states.py

@@ -38,6 +38,19 @@ def __array__(self):
     def numpy(self):
         return self.__array__()
 
+    def symbolic(self, decimals=5, cutoff=1e-10, max_terms=20):
+        state = self.numpy()
+        terms = []
+        for i in K.np.nonzero(state)[0]:
+            b = bin(i)[2:].zfill(self.nqubits)
+            if K.np.abs(state[i]) >= cutoff:
+                x = round(state[i], decimals)
+                terms.append(f"{x}|{b}>")
+            if len(terms) >= max_terms:
+                terms.append("...")
+                break
+        return " + ".join(terms)
+
     @classmethod
     def zero_state(cls, nqubits):
         state = cls(nqubits)
@@ -132,6 +145,21 @@ def expectation(self, hamiltonian, normalize=False):
 
 class MatrixState(VectorState):
 
+    def symbolic(self, decimals=5, cutoff=1e-10, max_terms=20):
+        state = self.numpy()
+        terms = []
+        indi, indj = K.np.nonzero(state)
+        for i, j in zip(indi, indj):
+            bi = bin(i)[2:].zfill(self.nqubits)
+            bj = bin(j)[2:].zfill(self.nqubits)
+            if K.np.abs(state[i, j]) >= cutoff:
+                x = round(state[i, j], decimals)
+                terms.append(f"{x}|{bi}><{bj}|")
+            if len(terms) >= max_terms:
+                terms.append("...")
+                break
+        return " + ".join(terms)
+
     @property
     def shape(self):
         return (self.nstates, self.nstates)

src/qibo/tests/test_core_states.py

@@ -68,6 +68,49 @@ def test_vector_state_to_density_matrix(backend):
         state.to_density_matrix()
 
 
+@pytest.mark.parametrize("target", range(5))
+@pytest.mark.parametrize("density_matrix", [False, True])
+def test_state_representation(target, density_matrix):
+    from qibo import models, gates
+    c = models.Circuit(5, density_matrix=density_matrix)
+    c.add(gates.H(target))
+    result = c()
+    bstring = target * "0" + "1" + (4 - target) * "0"
+    if density_matrix:
+        target_str = 3 * [f"(0.5+0j)|00000><00000| + (0.5+0j)|00000><{bstring}| + (0.5+0j)|{bstring}><00000| + (0.5+0j)|{bstring}><{bstring}|"]
+    else:
+        target_str = [f"(0.70711+0j)|00000> + (0.70711+0j)|{bstring}>",
+                      f"(0.7+0j)|00000> + (0.7+0j)|{bstring}>",
+                      f"(0.71+0j)|00000> + (0.71+0j)|{bstring}>"]
+    assert str(result) == target_str[0]
+    assert result.state(decimals=5) == target_str[0]
+    assert result.symbolic(decimals=1) == target_str[1]
+    assert result.symbolic(decimals=2) == target_str[2]
+
+
+@pytest.mark.parametrize("density_matrix", [False, True])
+def test_state_representation_max_terms(density_matrix):
+    from qibo import models, gates
+    c = models.Circuit(5, density_matrix=density_matrix)
+    c.add(gates.H(i) for i in range(5))
+    result = c()
+    if density_matrix:
+        assert result.symbolic(max_terms=3) == "(0.03125+0j)|00000><00000| + (0.03125+0j)|00000><00001| + (0.03125+0j)|00000><00010| + ..."
+        assert result.symbolic(max_terms=5) == "(0.03125+0j)|00000><00000| + (0.03125+0j)|00000><00001| + (0.03125+0j)|00000><00010| + (0.03125+0j)|00000><00011| + (0.03125+0j)|00000><00100| + ..."
+    else:
+        assert result.symbolic(max_terms=3) == "(0.17678+0j)|00000> + (0.17678+0j)|00001> + (0.17678+0j)|00010> + ..."
+        assert result.symbolic(max_terms=5) == "(0.17678+0j)|00000> + (0.17678+0j)|00001> + (0.17678+0j)|00010> + (0.17678+0j)|00011> + (0.17678+0j)|00100> + ..."
+
+
+def test_state_representation_cutoff():
+    from qibo import models, gates
+    c = models.Circuit(2)
+    c.add(gates.RX(0, theta=0.1))
+    result = c()
+    assert result.state(decimals=5) == "(0.99875+0j)|00> + -0.04998j|10>"
+    assert result.state(decimals=5, cutoff=0.1) == "(0.99875+0j)|00>"
+
+
 @pytest.mark.parametrize("state_type", ["VectorState", "MatrixState"])
 @pytest.mark.parametrize("use_gate", [False, True])
 def test_state_probabilities(backend, state_type, use_gate):