diff --git a/vqls_prototype/solver/vqls.py b/vqls_prototype/solver/vqls.py
index ff395b9..1ae0522 100644
--- a/vqls_prototype/solver/vqls.py
+++ b/vqls_prototype/solver/vqls.py
@@ -208,6 +208,9 @@ def construct_circuit(  # pylint: disable=too-many-branches
             # ensure the vector is double
             vector = vector.astype("float64")
 
+            if vector.ndim == 2:
+                vector = vector.flatten()
+
             # create the circuit
             nqbit = int(np.log2(len(vector)))
             self.vector_circuit = QuantumCircuit(nqbit, name="Ub")
diff --git a/vqls_prototype/tomography/htree_qst.py b/vqls_prototype/tomography/htree_qst.py
index e1337e2..4e79f43 100644
--- a/vqls_prototype/tomography/htree_qst.py
+++ b/vqls_prototype/tomography/htree_qst.py
@@ -45,10 +45,10 @@ def get_tree(self):
         """Compute the tree"""
 
         def init_tree():
-            """_summary_
+            """initialize the tree
 
             Returns:
-                _type_: _description_
+                tree: tree
             """
             trees = []
             level_root, level_leaf = [], []
@@ -64,10 +64,10 @@ def init_tree():
             return trees, level_root, level_leaf
 
         def link_trees(trees):
-            """_summary_
+            """link multiple tress
 
             Args:
-                trees (_type_): _description_
+                trees (list): list of trees
             """
             ntree = len(trees)
             level_root, level_leaf = [], []
@@ -95,7 +95,7 @@ def link_trees(trees):
         return tree
 
     def get_path(self):
-        """_summary_"""
+        """Create the paths between the root and all the leaves"""
         paths = []
         for inode in range(self.size):
             paths.append(list(self.tree.rsearch(inode)))
@@ -105,7 +105,7 @@ def get_path_sparse_matrix(self):
         """transforms the path into a sparse matrix
 
         Returns:
-            _type_: _description_
+            coo matrix: sparse matrix of the path
         """
         row_idx, col_idx, vals = [], [], []
         for ip, path in enumerate(self.path_to_node):
@@ -118,10 +118,10 @@ def get_path_sparse_matrix(self):
         )
 
     def get_circuits(self):
-        """_summary_
+        """Create the circuits containing a single H on a given qubit after the circuit
 
         Args:
-            circuits (_type_): _description_
+            circuits (list): List of circuits
         """
         list_circuits = [self.circuit.measure_all(inplace=False)]
 
@@ -133,10 +133,10 @@ def get_circuits(self):
         return list_circuits
 
     def get_samples(self, parameters):
-        """_summary_
+        """Sample the circuit
 
         Args:
-            sampler (_type_): _description_
+            parameters (np.array): values of the variational parameters of the circuit
         """
         results = (
             self.sampler.run(self.list_circuits, [parameters] * self.ncircuits)
@@ -152,10 +152,10 @@ def get_samples(self, parameters):
         return samples
 
     def get_weight(self, samples):
-        """_summary_
+        """Get the relative sign between parent/child node
 
         Args:
-            samples (_type_): _description_
+            samples (list): lit of samples of circuits
         """
         # root
         weights = np.zeros_like(samples[0])
@@ -173,10 +173,10 @@ def get_weight(self, samples):
         return weights
 
     def get_signs(self, weights):
-        """Agregate the signs of the statevector
+        """Compute the signs of each components
 
         Args:
-            weights (np.array):
+            weights (np.array): relative sign between parent/child in the tree
         """
 
         # if the path is not known
@@ -191,23 +191,27 @@ def get_signs(self, weights):
         return np.multiply.reduceat(mat.data, self.idx_path_matrix[:-1])
 
     def get_relative_amplitude_sign(self, parameters):
-        """_summary_
+        """Get the relative amplitude of each components relative to the root
 
         Args:
-            parameters (_type_): _description_
+            parameters (np.array): values of the variational parameters of the circuit
         """
         samples = self.get_samples(parameters)
         weights = self.get_weight(samples)
         return self.get_signs(weights)
 
     def get_statevector(self, parameters):
-        """_summary_
+        """Get the statevector of the circuit
 
         Args:
-            parameters (_type_): _description_
+            parameters (np.array): values of the variational parameters of the circuit
         """
         samples = self.get_samples(parameters)
-        amplitudes = np.sqrt(samples[0])
+        if np.any(samples[0] < 0):
+            print("Warning : Negative sampling values found in HTree")
+            amplitudes = np.sqrt(np.abs(samples[0]))
+        else:
+            amplitudes = np.sqrt(samples[0])
         weights = self.get_weight(samples)
         signs = self.get_signs(weights)
         return amplitudes * signs