Add SetBasisState and SetState to lightning.kokkos

.github/CHANGELOG.md

@@ -63,8 +63,9 @@
 * Add a Catalyst-specific wrapping class for Lightning Kokkos.
   [(#770)](https://github.com/PennyLaneAI/pennylane-lightning/pull/770)
 
-* Add `initial_state_prep` option to Catalyst TOML file.
+* Add `initial_state_prep` option to TOML files and Lightning Kokkos.
   [(#826)](https://github.com/PennyLaneAI/pennylane-lightning/pull/826)
+  [(#838)](https://github.com/PennyLaneAI/pennylane-lightning/pull/838)
 
 ### Documentation
 

pennylane_lightning/core/_version.py

@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.38.0-dev25"
+__version__ = "0.38.0-dev26"

pennylane_lightning/core/src/simulators/lightning_kokkos/catalyst/LightningKokkosSimulator.cpp

@@ -137,6 +137,19 @@ void LightningKokkosSimulator::PrintState() {
 }
 /// LCOV_EXCL_STOP
 
+void LightningKokkosSimulator::SetState(
+    DataView<std::complex<double>, 1> &data_view) {
+    std::vector<Kokkos::complex<double>> data_vector(data_view.begin(),
+                                                     data_view.end());
+    std::vector<std::size_t> count(data_vector.size());
+    std::iota(std::begin(count), std::end(count), 0);
+    this->device_sv->setStateVector(count, data_vector);
+}
+
+void LightningKokkosSimulator::SetBasisState(const std::size_t index) {
+    this->device_sv->setBasisState(index);
+}
+
 auto LightningKokkosSimulator::Zero() const -> Result {
     return const_cast<Result>(&GLOBAL_RESULT_FALSE_CONST);
 }

pennylane_lightning/core/src/simulators/lightning_kokkos/catalyst/tests/Test_LightningKokkosSimulator.cpp

@@ -659,4 +659,32 @@ TEST_CASE("LightningKokkosSimulator::GateSet", "[GateSet]") {
             expected{3, 0, 1, {"Hadamard", "Hadamard", "IsingZZ"}, {}};
         REQUIRE(LKsim->CacheManagerInfo() == expected);
     }
+
+    SECTION("setBasisState") {
+        const size_t num_qubits = 1;
+        std::unique_ptr<LKSimulator> sv1 = std::make_unique<LKSimulator>();
+        std::vector<intptr_t> Qs = sv1->AllocateQubits(num_qubits);
+        sv1->setBasisState(0);
+        std::vector<std::complex<PrecisionT>> expected(
+            {{1.0, 0.0}, {0.0, 0.0}});
+        CHECK(sv1->getDataVector() == expected);
+        sv1->setBasisState(1);
+        std::vector<std::complex<PrecisionT>> expected2(
+            {{0.0, 0.0}, {1.0, 0.0}});
+        CHECK(sv1->getDataVector() == expected2);
+    }
+
+    SECTION("setStateVector") {
+        const size_t num_qubits = 1;
+        std::unique_ptr<LKSimulator> sv1 = std::make_unique<LKSimulator>();
+        std::vector<intptr_t> Qs = sv1->AllocateQubits(num_qubits);
+        sv1->setStateVector({0}, {{0.5, 0.5}});
+        std::vector<std::complex<PrecisionT>> expected(
+            {{0.5, 0.5}, {0.0, 0.0}});
+        CHECK(sv1->getDataVector() == expected);
+        sv1->setStateVector({1}, {{0.5, 0.5}});
+        std::vector<std::complex<PrecisionT>> expected2(
+            {{0.5, 0.5}, {0.5, 0.5}});
+        CHECK(sv1->getDataVector() == expected2);
+    }
 }