From bc5b578724c1f0e9e5ac36663729e5d1b70e53c6 Mon Sep 17 00:00:00 2001
From: Sergei Isakov <54642992+sergeisakov@users.noreply.github.com>
Date: Tue, 13 Feb 2024 12:59:15 +0100
Subject: [PATCH] Add method for adding to opstring.

---
 pybind_interface/pybind_main.cpp |  6 ++++++
 pybind_interface/pybind_main.h   | 12 ++++++++++--
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/pybind_interface/pybind_main.cpp b/pybind_interface/pybind_main.cpp
index 63a03e4a..2520d054 100644
--- a/pybind_interface/pybind_main.cpp
+++ b/pybind_interface/pybind_main.cpp
@@ -356,6 +356,12 @@ void add_gate_to_opstring(const Cirq::GateKind gate_kind,
   opstring->ops.push_back(create_gate(gate_kind, 0, qubits, params));
 }
 
+void add_matrix_gate_to_opstring(const std::vector<unsigned>& qubits,
+                                 const std::vector<float>& matrix,
+                                 OpString<Cirq::GateCirq<float>>* opstring) {
+  opstring->ops.push_back(create_matrix_gate(0, qubits, matrix));
+}
+
 // Methods for simulating amplitudes.
 
 std::vector<std::complex<float>> qsim_simulate(const py::dict &options) {
diff --git a/pybind_interface/pybind_main.h b/pybind_interface/pybind_main.h
index 3263fdd8..1e439547 100644
--- a/pybind_interface/pybind_main.h
+++ b/pybind_interface/pybind_main.h
@@ -22,6 +22,7 @@
 #include <pybind11/stl_bind.h>
 namespace py = pybind11;
 
+#include <map>
 #include <vector>
 
 #include "../lib/circuit.h"
@@ -75,12 +76,17 @@ void add_channel(const unsigned time,
                      prob_matrix_unitary_triples,
                  qsim::NoisyCircuit<qsim::Cirq::GateCirq<float>>* ncircuit);
 
-// Method for populating opstrings.
+// Methods for populating opstrings.
 void add_gate_to_opstring(
     const qsim::Cirq::GateKind gate_kind,
     const std::vector<unsigned>& qubits,
     qsim::OpString<qsim::Cirq::GateCirq<float>>* opstring);
 
+void add_matrix_gate_to_opstring(
+    const std::vector<unsigned>& qubits,
+    const std::vector<float>& matrix,
+    qsim::OpString<qsim::Cirq::GateCirq<float>>* opstring);
+
 // Methods for simulating noiseless circuits.
 std::vector<std::complex<float>> qsim_simulate(const py::dict &options);
 
@@ -376,7 +382,9 @@ std::vector<std::complex<float>> qsimh_simulate(const py::dict &options);
             "Adds a channel to the given noisy circuit.");                            \
                                                                                       \
       m.def("add_gate_to_opstring", &add_gate_to_opstring,                            \
-            "Adds a gate to the given opstring.");
+            "Adds a gate to the given opstring.");                                    \
+      m.def("add_matrix_gate_to_opstring", &add_matrix_gate_to_opstring,              \
+            "Adds a matrix gate to the given opstring.");
 
 #define GPU_MODULE_BINDINGS                                                                  \
       m.doc() = "pybind11 plugin";  /* optional module docstring */                   \