Merge 9646a23 into fd41d8f

PennyLaneAI · Mar 19, 2021 · 9e0d783 · 9e0d783
2 parents fd41d8f + 9646a23
commit 9e0d783
Show file tree

Hide file tree

Showing 5 changed files with 108 additions and 6 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -6,6 +6,9 @@
 
 ### Improvements
 
+* All the gates supported by `default.qubit` are now supported.
+  [(#91)](https://github.com/PennyLaneAI/pennylane-lightning/pull/91)
+
 * Add new lightweight backend with performance improvements.
   [(#57)](https://github.com/PennyLaneAI/pennylane-lightning/pull/57)
 

diff --git a/pennylane_lightning/lightning_qubit.py b/pennylane_lightning/lightning_qubit.py
@@ -49,9 +49,7 @@ class LightningQubit(DefaultQubit):
     version = __version__
     author = "Xanadu Inc."
 
-    operations = {
-        "BasisState",
-        "QubitStateVector",
+    kernel_operations = {
         "PauliX",
         "PauliY",
         "PauliZ",
@@ -133,9 +131,24 @@ def apply_lightning(self, state, operations):
         Returns:
             array[complex]: the output state tensor
         """
-        op_names = [o.name for o in operations]
-        op_wires = [self.wires.indices(o.wires) for o in operations]
-        op_param = [o.parameters for o in operations]
+        op_names = []
+        op_wires = []
+        op_param = []
+
+        for o in operations:
+            op_wires.append(self.wires.indices(o.wires))
+            if o.name in self.kernel_operations:
+                op_names.append(o.name)
+                op_param.append(o.parameters)
+            else:
+                op_names.append("QubitUnitary")
+
+                mat = o.matrix
+                if not np.iscomplex(o.matrix).all():
+                    mat = mat.astype(np.complex128)
+
+                unitary_view = mat.ravel().view(np.float64)
+                op_param.append(unitary_view)
 
         state_vector = np.ravel(state)
         apply(state_vector, op_names, op_wires, op_param, self.num_wires)

diff --git a/pennylane_lightning/src/Gates.cpp b/pennylane_lightning/src/Gates.cpp
@@ -556,6 +556,52 @@ void Pennylane::CSWAPGate::applyKernel(const StateVector& state, const std::vect
 
 // -------------------------------------------------------------------------------------------------------------
 
+const string Pennylane::Unitary::label = "QubitUnitary";
+
+Pennylane::Unitary Pennylane::Unitary::create(const vector<double>& parameters) {
+    auto size = parameters.size();
+    unsigned int log2size = 0;
+    while (size >>= 1) ++log2size;
+    unsigned int numQubits = (log2size - 1) / 2;
+
+    return Pennylane::Unitary(numQubits, parameters);
+}
+
+Pennylane::Unitary::Unitary(const int numQubits, const std::vector<double>& real_matrix)
+    : AbstractGate(numQubits)
+    , real_matrix{real_matrix}
+{}
+
+void Pennylane::Unitary::applyKernel(const StateVector& state, const std::vector<size_t>& indices, const std::vector<size_t>& externalIndices) {
+    const vector<double>& real_matrix = asRealMatrix();
+    assert(indices.size() == length);
+
+    vector<CplxType> v(indices.size());
+    for (const size_t& externalIndex : externalIndices) {
+        CplxType* shiftedState = state.arr + externalIndex;
+        // Gather
+        size_t pos = 0;
+        for (const size_t& index : indices) {
+            v[pos] = shiftedState[index];
+            pos++;
+        }
+
+        // Apply + scatter
+        for (size_t i = 0; i < indices.size(); i++) {
+            size_t index = indices[i];
+            shiftedState[index] = 0;
+            size_t baseIndex = i * indices.size();
+            for (size_t j = 0; j < indices.size(); j++) {
+                unsigned int indx = 2 * (baseIndex + j);
+                CplxType matrix_elem(real_matrix[indx], real_matrix[indx + 1]);
+                shiftedState[index] += matrix_elem * v[j];
+            }
+        }
+    }
+}
+
+// -------------------------------------------------------------------------------------------------------------
+
 template<class GateType>
 static void addToDispatchTable(map<string, function<unique_ptr<Pennylane::AbstractGate>(const vector<double>&)>>& dispatchTable) {
     dispatchTable.emplace(GateType::label, [](const vector<double>& parameters) { return make_unique<GateType>(GateType::create(parameters)); });
@@ -583,6 +629,7 @@ static map<string, function<unique_ptr<Pennylane::AbstractGate>(const vector<dou
     addToDispatchTable<Pennylane::CGeneralRotationGate>(dispatchTable);
     addToDispatchTable<Pennylane::ToffoliGate>(dispatchTable);
     addToDispatchTable<Pennylane::CSWAPGate>(dispatchTable);
+    addToDispatchTable<Pennylane::Unitary>(dispatchTable);
     return dispatchTable;
 }
 

diff --git a/pennylane_lightning/src/Gates.hpp b/pennylane_lightning/src/Gates.hpp
@@ -323,6 +323,23 @@ namespace Pennylane {
         void applyKernel(const StateVector& state, const std::vector<size_t>& indices, const std::vector<size_t>& externalIndices);
     };
 
+    class Unitary : public AbstractGate {
+    private:
+        const std::vector<CplxType> matrix;
+        const std::vector<double> real_matrix;
+    public:
+        static const std::string label;
+        static Unitary create(const std::vector<double>& parameters);
+        Unitary(const int numQubits, const std::vector<double>& real_matrix);
+        inline const std::vector<CplxType>& asMatrix() {
+            return matrix;
+        }
+        inline const std::vector<double>& asRealMatrix() {
+            return real_matrix;
+        }
+        void applyKernel(const StateVector& state, const std::vector<size_t>& indices, const std::vector<size_t>& externalIndices);
+    };
+
     /**
      * Produces the requested gate, defined by a label and the list of parameters
      *

diff --git a/tests/test_apply.py b/tests/test_apply.py
@@ -19,6 +19,7 @@
 
 # pylint: disable=protected-access,cell-var-from-loop
 from unittest import mock
+from scipy.stats import unitary_group, ortho_group
 
 import numpy as np
 import pennylane as qml
@@ -1212,3 +1213,24 @@ def test_pauliz_hadamard(self, theta, phi, varphi, monkeypatch, qubit_device_3_w
             - 2 * np.cos(theta) * np.sin(phi) * np.sin(2 * varphi)
         ) / 4
         assert np.allclose(var, expected, atol=tol, rtol=0)
+
+
+@pytest.mark.parametrize("real", [True, False])
+@pytest.mark.parametrize("wires", range(2, 10))
+def test_qubit_unitary(wires, real):
+    """Test if lightning.qubit successful applies an arbitrary unitary"""
+
+    if real:
+        unitary = ortho_group.rvs(2 ** wires, random_state=0)
+    else:
+        unitary = unitary_group.rvs(2 ** wires, random_state=0)
+
+    dev = qml.device("lightning.qubit", wires=wires)
+
+    @qml.qnode(dev)
+    def f():
+        qml.QubitUnitary(unitary, wires=range(wires))
+        return qml.state()
+
+    target_state = unitary[:, 0]
+    assert np.allclose(f(), target_state)