Make spin_op.distribute_terms batch more efficiently

python/tests/builder/test_SpinOperator.py

@@ -378,12 +378,44 @@ def test_spin_op_from_word():
 def test_spin_op_serdes():
     for nq in range(1, 31):
         for nt in range(1, nq + 1):
-            h1 = cudaq.SpinOperator.random(qubit_count=nq, term_count=nt)
+            h1 = cudaq.SpinOperator.random(qubit_count=nq,
+                                           term_count=nt,
+                                           seed=13)
             h2 = h1.serialize()
             h3 = cudaq.SpinOperator(h2, nq)
             assert (h1 == h3)
 
 
+def test_spin_op_random():
+    # Make sure that the user gets all the random terms that they ask for.
+    qubit_count = 5
+    term_count = 7
+    for i in range(100):
+        hamiltonian = cudaq.SpinOperator.random(qubit_count, term_count, seed=i)
+        assert hamiltonian.get_term_count() == term_count
+
+
+def test_spin_op_random_too_many():
+    # Make sure that the user gets all the random terms that they ask for.
+    qubit_count = 3
+    term_count = 21  # too many because 6 choose 3 = 20
+    with pytest.raises(RuntimeError) as error:
+        hamiltonian = cudaq.SpinOperator.random(qubit_count, term_count)
+
+
+def test_spin_op_batch_efficiently():
+    qubit_count = 5
+    term_count = 7
+    num_of_gpus = 4
+    hamiltonian = cudaq.SpinOperator.random(qubit_count, term_count, seed=13)
+    batched = hamiltonian.distribute_terms(num_of_gpus)
+    assert len(batched) == 4
+    assert batched[0].get_term_count() == 2
+    assert batched[1].get_term_count() == 2
+    assert batched[2].get_term_count() == 2
+    assert batched[3].get_term_count() == 1
+
+
 # leave for gdb debugging
 if __name__ == "__main__":
     loc = os.path.abspath(__file__)

runtime/cudaq/spin/spin_op.cpp

@@ -23,6 +23,7 @@
 #include <iostream>
 #include <map>
 #include <random>
+#include <set>
 #include <utility>
 #include <vector>
 
@@ -295,11 +296,43 @@ spin_op spin_op::random(std::size_t nQubits, std::size_t nTerms,
   std::mt19937 gen(seed);
   std::vector<std::complex<double>> coeffs(nTerms, 1.0);
   std::vector<spin_op_term> randomTerms;
+  // Make sure we don't put duplicates into randomTerms by using dupCheckSet
+  std::set<std::vector<bool>> dupCheckSet;
+  if (nQubits <= 30) {
+    // For the given algorithm below that sets bool=true for 1/2 of the the
+    // termData, the maximum number of unique terms is n choose k, where n =
+    // 2*nQubits, and k=nQubits. For up to 30 qubits, we can calculate n choose
+    // k without overflows (i.e. 60 choose 30 = 118264581564861424) to validate
+    // that nTerms is reasonable. For anything larger, the user can't set nTerms
+    // large enough to run into actual problems because they would encounter
+    // memory limitations long before anything else.
+    // Note: use the multiplicative formula to evaluate n-choose-k. The
+    // arrangement of multiplications and divisions do not truncate any division
+    // remainders.
+    std::size_t maxTerms = 1;
+    for (std::size_t i = 1; i <= nQubits; i++) {
+      maxTerms *= 2 * nQubits + 1 - i;
+      maxTerms /= i;
+    }
+    if (nTerms > maxTerms)
+      throw std::runtime_error(
+          fmt::format("Unable to produce {} unique random terms for {} qubits",
+                      nTerms, nQubits));
+  }
   for (std::size_t i = 0; i < nTerms; i++) {
     std::vector<bool> termData(2 * nQubits);
-    std::fill_n(termData.begin(), termData.size() * (1 - .5), 1);
-    std::shuffle(termData.begin(), termData.end(), gen);
-    randomTerms.push_back(termData);
+    while (true) {
+      std::fill_n(termData.begin(), nQubits, true);
+      std::shuffle(termData.begin(), termData.end(), gen);
+      if (dupCheckSet.contains(termData)) {
+        // Prepare to loop again
+        std::fill(termData.begin(), termData.end(), false);
+      } else {
+        dupCheckSet.insert(termData);
+        break;
+      }
+    }
+    randomTerms.push_back(std::move(termData));
   }
 
   return spin_op(randomTerms, coeffs);
@@ -488,29 +521,24 @@ std::size_t spin_op::num_terms() const { return terms.size(); }
 std::vector<spin_op> spin_op::distribute_terms(std::size_t numChunks) const {
   // Calculate how many terms we can equally divide amongst the chunks
   auto nTermsPerChunk = num_terms() / numChunks;
+  auto leftover = num_terms() % numChunks;
 
   // Slice the given spin_op into subsets for each chunk
   std::vector<spin_op> spins;
-  for (auto i : cudaq::range(numChunks)) {
-    // lowerBound here is the start index
-    auto lowerBound = i * nTermsPerChunk;
-
-    // Get the iterator and advance to lowerBound
-    auto start = terms.begin();
-    std::advance(start, lowerBound);
-
-    // The number of terms we want is nTermsPerChunk, but if
-    // this is the last iteration of this loop, we'll add
-    // any run-over terms to the final chunk
-    auto count =
-        nTermsPerChunk + (i == numChunks - 1 ? (num_terms() % numChunks) : 0);
+  auto termIt = terms.begin();
+  for (std::size_t chunkIx = 0; chunkIx < numChunks; chunkIx++) {
+    // Evenly distribute any leftovers across the early chunks
+    auto count = nTermsPerChunk + (chunkIx < leftover ? 1 : 0);
 
     // Get the chunk from the terms list.
     std::unordered_map<spin_op_term, std::complex<double>> sliced;
-    std::copy_n(start, count, std::inserter(sliced, sliced.end()));
+    std::copy_n(termIt, count, std::inserter(sliced, sliced.end()));
 
     // Add to the return vector
     spins.emplace_back(sliced);
+
+    // Get ready for the next loop
+    std::advance(termIt, count);
   }
 
   // return the terms.

unittests/spin_op/SpinOpTester.cpp

@@ -240,6 +240,6 @@ TEST(SpinOpTester, checkDistributeTerms) {
   auto distributed = H.distribute_terms(2);
 
   EXPECT_EQ(distributed.size(), 2);
-  EXPECT_EQ(distributed[0].num_terms(), 2);
-  EXPECT_EQ(distributed[1].num_terms(), 3);
+  EXPECT_EQ(distributed[0].num_terms(), 3);
+  EXPECT_EQ(distributed[1].num_terms(), 2);
 }