Ttrenty/feature/partial trace

README.md

@@ -28,8 +28,8 @@ curl -sSf https://pixi.sh/install.sh | bash
 # Install all project dependencies:
 pixi install
 
-# Build and run the simulator:
-pixi run mojo build src/main.mojo && ./main
+# Build and run examples of the simulator:
+pixi run main
 ```
 
 

TODOs.md

@@ -1,26 +1,54 @@
 # TODOs
 
-A pure state is represented by a ket vector, e.g., $ |\psi\rangle = \alpha|0\rangle + \beta|1\rangle $, where $ \alpha $ and $ \beta $ are complex numbers satisfying the normalization condition $ |\alpha|^2 + |\beta|^2 = 1 $.
-
 ## Ordered by Priority (5 ↑) / Difficulty (5 ↓)
 
-- 5 / 1 : density matrix calculcation from state vectors
-
-- 5 / 5 : Extend qubitWiseMultiply() to 2 and multiple qubits gates (Test it)
+### Implementations
 
 - 4 / 3 : Implement measurement gates
 
-- 4 / 5 : Implement the computation of statistics: (6.5 and 6.6)
+- 4 / 5 : Implement the computation of statistics (6.5 and 6.6)
 
-- 3 / 2 :  Use a separate list for things that are not real gate to not slow down the main run logic
-
-- 3 / 3 : Implement naive implementation of the functions 
+- 3 / 3 : Implement naive implementation of the functions to compare performances
     - matrix multiplication (but starting from right or smart)
     - partial trace
 
-- 3 / 4 : Compile time circuit creation?
+- 5 / 5 : Start adding support for GPU in the base classes if needed (not possible to use SIMD(complexfloat64) anymore, or keep them but seperate them when moving data to GPU)
+    - struct PureBasisState
+    - struct ComplexMatrix
+    - struct Gate
+
+- 5 / ? : GPU implementation of:
+    - qubit_wise_multiply()
+    - apply_swap()
+    - partial_trace()
+
+
+### Tests
+
+- 5 / 2 : Test qubit_wise_multiply_extended() that can take multiple qubits gates (2 and more, iSWAP for example)
+
+- 5 / 2 : Test for everything that will be implement in GPU
+    - qubit_wise_multiply()
+    - apply_swap()
+    - partial_trace()
+    - struct PureBasisState's methods
+    - struct ComplexMatrix's methods
+    - struct Gate's Gate
+
+### Benchmarks
 
 - 3 / 2 : Reproduce table from page 10
 
-- 2 / 4 : qubitWiseMultiply() but for multiple qubits gates applied to non-adjacent qubits
+## Droped for now
+
+- 4 / 4 : Gradient computation with Finite Difference
+
+- 3 / 2 :  Use a separate list for things that are not real gate to not slow down the main run logic
+
+- 3 / 4 : Compile time circuit creation?
+
+- 3 / 4 : Gradient computation with Parameter-Shift
+
+- 3 / 100 : Gradient computation with Adjoint Method
 
+- 2 / 4 : qubit_wise_multiply_extended() but for gates applied to non-adjacent qubits

examples/main.mojo

@@ -380,6 +380,47 @@ fn presentation() -> None:
     print("Final quantum state:\n", final_state)
 
 
+fn test_density_matrix() -> None:
+    """
+    Returns the density matrix of the given quantum state.
+    If qubits is empty, returns the full density matrix.
+    """
+    num_qubits: Int = 2
+    qc: GateCircuit = GateCircuit(num_qubits)
+
+    qc.apply_gates(
+        Hadamard(0),
+        Hadamard(1, controls=[0]),
+        Z(0),
+        X(1),
+    )
+
+    print("Quantum circuit created:\n", qc)
+
+    qsimu = StateVectorSimulator(
+        qc,
+        initial_state=PureBasisState.from_bitstring("00"),
+        optimisation_level=0,  # No optimisations for now
+        verbose=True,
+        verbose_step_size=ShowAfterEachGate,  # ShowAfterEachGate, ShowOnlyEnd
+    )
+    final_state = qsimu.run()
+    print("Final quantum state:\n", final_state)
+
+    matrix = final_state.to_density_matrix()
+    print("Density matrix:\n", matrix)
+    other_matrix = partial_trace(final_state, [])  # Empty list means full trace
+    print("Partial trace matrix:\n", other_matrix)
+    other_matrix_0 = partial_trace(
+        final_state, [0]
+    )  # Empty list means full trace
+    print("Partial trace matrix qubit 0:\n", other_matrix_0)
+    other_matrix_1 = partial_trace(
+        final_state, [1]
+    )  # Empty list means full trace
+    print("Partial trace matrix qubit 1:\n", other_matrix_1)
+
+
 # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
 # MARK:         Tests                #
 # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
@@ -427,8 +468,8 @@ fn presentation() -> None:
 
 def main():
     args = argv()
-    number_qubits: Int = 15
-    number_layers: Int = 10
+    number_qubits: Int = 10
+    number_layers: Int = 20
     if len(args) == 3:
         try:
             number_qubits = Int(args[1])
@@ -441,14 +482,16 @@ def main():
     else:
         print("Usage: ./main [number_of_qubits] [number_of_layers]")
 
-    simulate_figure1_circuit()
+    # simulate_figure1_circuit()
 
-    # simulate_figure1_circuit_abstract()
+    simulate_figure1_circuit_abstract()
 
-    # simulate_random_circuit(number_qubits, number_layers)
+    simulate_random_circuit(number_qubits, number_layers)
 
     # simulate_figure4_circuit()
 
     # simulate_figure4_circuit_abstract()
 
     # presentation()
+
+    # test_density_matrix()

pixi.toml

@@ -13,17 +13,17 @@ cmd = ".github/scripts/check-format.sh"
 
 [tasks.format] # Format the code
 cmd = "pixi run mojo format ./src ./tests"
-inputs = ["./src/**/*.mojo", "./tests/**/*.mojo"]
+inputs = ["./examples/**/*.mojo", "./src/**/*.mojo", "./tests/**/*.mojo"]
 
 [tasks.create_build_dir]
 cmd = "mkdir -p build/"
 
 [tasks.build] # Compile any mojo file
 args = [
-    { "arg" = "full_file_path", "default" = "main" },
+    { "arg" = "full_file_path", "default" = "examples/main.mojo" },
     { "arg" = "executable_name", "default" = "main" },
 ]
-inputs = ["{{ full_file_path }}"]
+inputs = ["{{ full_file_path }}", "./src/**/*.mojo"]
 outputs = ["build/{{ executable_name }}"]
 cmd = "pixi run mojo build {{ full_file_path }} -o {{ executable_name }} && cp {{ executable_name }} build/{{ executable_name }} && rm {{ executable_name }}"
 depends-on = ["create_build_dir"]
@@ -70,6 +70,7 @@ depends-on = ["install"]
 
 [tasks]
 
+tests = [{ task = "test" }]
 p = [{ task = "clear" }, { task = "package" }]
 m = [{ task = "clear" }, { task = "main" }]
 # t = "clear && pixi run package && pixi run mojo test tests --filter"

src/abstractions/gate_circuit.mojo

@@ -112,7 +112,8 @@ struct GateCircuit(Movable, Stringable, Writable):
                 controls_flags: List[List[Int]] = gate.control_qubits_with_flags
 
                 for control in controls_flags:
-                    control_index, flag = control[0], control[1]
+                    # control_index, flag = control[0], control[1]
+                    control_index = control[0]
                     while (
                         wires_current_gate[qubit_index]
                         < wires_current_gate[control_index]

src/abstractions/simulator.mojo

@@ -211,15 +211,8 @@ struct StateVectorSimulator(Copyable, Movable):
         i: Int = 0
         layer_index: Int = 0
         for gate in self.circuit.gates:  # Iterate over the gates in the circuit
-            if gate.symbol not in [_SEPARATOR.symbol, SWAP.symbol]:
-                # Apply the next gate
-                quantum_state = qubit_wise_multiply_extended(
-                    len(gate.target_qubits),  # Number of target qubits
-                    gate.matrix,
-                    gate.target_qubits,  # Assuming single target qubit
-                    quantum_state,
-                    gate.control_qubits_with_flags,
-                )
+            if gate.symbol == _SEPARATOR.symbol:
+                continue
             elif gate.symbol == SWAP.symbol:
                 if len(gate.target_qubits) != 2:
                     print("Error: SWAP gate must have exactly 2 target qubits.")
@@ -231,11 +224,15 @@ struct StateVectorSimulator(Copyable, Movable):
                     quantum_state,
                     gate.control_qubits_with_flags,
                 )
-            elif gate.symbol == _SEPARATOR.symbol:
-                continue
             else:
-                print("Error: Unexpected gate symbol:", gate.symbol)
-                continue  # Skip unexpected symbols
+                # Apply the next gate
+                quantum_state = qubit_wise_multiply_extended(
+                    len(gate.target_qubits),  # Number of target qubits
+                    gate.matrix,
+                    gate.target_qubits,  # Assuming single target qubit
+                    quantum_state,
+                    gate.control_qubits_with_flags,
+                )
 
             i += 1
             if self.verbose:

src/base/qubits_operations.mojo

@@ -587,40 +587,48 @@ fn partial_trace[
     n = quantum_state.number_qubits()
     conj_quantum_state = quantum_state.conjugate()
 
-    is_traced_out: List[Bool] = [False] * n
-    for i in range(len(qubits_to_trace_out)):
-        is_traced_out[qubits_to_trace_out[i]] = True
-
-    qubits_to_keep: List[Int] = []
-    for i in range(n):
-        if not is_traced_out[i]:
-            qubits_to_keep.append(i)
-
     # is_traced_out: List[Bool] = [False] * n
-    # qubits_to_keep: List[Int] = []
+    # for i in range(len(qubits_to_trace_out)):
+    #     is_traced_out[qubits_to_trace_out[i]] = True
 
-    # current_index: Int = 0
-    # current_traced_index: Int = 0
-    # for _ in range(n):
-    #     if qubits_to_trace_out[current_traced_index] == current_index:
-    #         is_traced_out[current_index] = True
-    #         current_traced_index += 1
-    #         if current_traced_index >= len(qubits_to_trace_out):
-    #             break
-    #     else:
+    # qubits_to_keep: List[Int] = []
+    # for i in range(n):
+    #     if not is_traced_out[i]:
     #         qubits_to_keep.append(i)
-    #     current_index += 1
 
-    # # Ensure all qubits have been added to qubits_to_keep
-    # for i in range(current_index, n):
-    #     qubits_to_keep.append(i)
+    is_traced_out: List[Bool] = [False] * n
+    qubits_to_keep: List[Int] = []
+
+    current_index: Int = 0
+    current_traced_index: Int = 0
+    for _ in range(n):
+        if current_traced_index >= len(qubits_to_trace_out):
+            break
+        if qubits_to_trace_out[current_traced_index] == current_index:
+            is_traced_out[current_index] = True
+            current_traced_index += 1
+            current_index += 1
+            if current_traced_index >= len(qubits_to_trace_out):
+                break
+        else:
+            qubits_to_keep.append(current_index)
+            current_index += 1
+
+    # Ensure all qubits have been added to qubits_to_keep
+    for i in range(current_index, n):
+        qubits_to_keep.append(i)
 
     num_qubits_to_trace_out: Int = len(qubits_to_trace_out)
     num_qubits_to_keep: Int = len(qubits_to_keep)
     if num_qubits_to_trace_out + num_qubits_to_keep != n:
         print(
-            "Error: The total number of qubits to trace out and keep does not"
-            " match the total number of qubits."
+            "Error: The total number of qubits to trace out (",
+            num_qubits_to_trace_out,
+            ") and keep (",
+            num_qubits_to_keep,
+            ") does not match the total number of qubits (",
+            n,
+            ").",
         )
         return ComplexMatrix(0, 0)  # Return an empty matrix
 

src/base/state_and_matrix.mojo

@@ -36,7 +36,7 @@ struct PureBasisState(Copyable, Movable, Stringable, Writable):
         self.state_vector = CustomList[ComplexFloat64, hint_trivial_type=True](
             length=size, fill=ComplexFloat64(0.0, 0.0)
         )
-        self.state_vector.memset_zero()  # Initialize the state vector with zeros
+        # self.state_vector.memset_zero()
 
     @always_inline
     fn __getitem__(self, index: Int) -> ComplexFloat64:
@@ -64,9 +64,11 @@ struct PureBasisState(Copyable, Movable, Stringable, Writable):
             if amplitude_im == 0.0 and amplitude_re == 0.0:
                 amplitude_str: String = String(Int(amplitude_re))
             elif amplitude_im == 0.0:
-                amplitude_str = String(round(amplitude_re, 2))
+                # amplitude_str = String(round(amplitude_re, 2))
+                amplitude_str = String(amplitude_re)
             elif amplitude_re == 0.0:
-                amplitude_str = String(round(amplitude_im, 2)) + "i"
+                # amplitude_str = String(round(amplitude_im, 2)) + "i"
+                amplitude_str = String(amplitude_im) + "i"
             else:
                 amplitude_str = (
                     String(round(amplitude_re, 2))
@@ -179,7 +181,7 @@ struct PureBasisState(Copyable, Movable, Stringable, Writable):
             conjugated_state.state_vector[i] = self.state_vector[i].conjugate()
         return conjugated_state
 
-    fn get_density_matrix(self) -> ComplexMatrix:
+    fn to_density_matrix(self) -> ComplexMatrix:
         """Returns the density matrix of the pure state.
 
         The density matrix is computed as the outer product of the state vector with itself.