diff --git a/lib_q_computer_math.py b/lib_q_computer_math.py
index 2853673..2f3826e 100644
--- a/lib_q_computer_math.py
+++ b/lib_q_computer_math.py
@@ -284,14 +284,23 @@ class State(Vector):
         if not (0 < qubit_n <= max_qubits):
             raise Exception("Partial measurement of qubit_n must be between 1 and {}".format(max_qubits))
         format_str = self.get_fmt_of_element()
+        # e.g. for state |000>:
+        # ['000', '001', '010', '011', '100', '101', '110', '111']
         bin_repr = [format_str.format(i) for i in range(len(self))]
+        # e.g. for qbit_n = 2
+        # [0, 0, 1, 1, 0, 0, 1, 1]
         partial_measurement_of_qbit = [int(b[qubit_n - 1]) for b in bin_repr]
+        # [0, 1, 4, 5]
         indexes_for_p_0 = [i for i, index in enumerate(partial_measurement_of_qbit) if index==0]
         weights_0 = [self.get_prob(j) for j in indexes_for_p_0]
         choices_0 = [format_str.format(i) for i in indexes_for_p_0]
         measurement_result = random.choices(choices_0, weights_0)[0]
-        # TODO: collapse the state and change the probabilities of the rest
+        # TODO: Verify if this is the correct collapse to lower dimension after partial measurement
         #       https://www.youtube.com/watch?v=MG_9JWsrKtM&list=PL1826E60FD05B44E4&index=16
+        normalization_factor = np.sqrt(np.sum([self.m[i][0] for i in indexes_for_p_0]))
+        self.m = [
+            [(self.m[i][0]**2)/normalization_factor] for i in indexes_for_p_0
+        ]
         return measurement_result
 
     def pretty_print(self):