Fix signed int overflow in uccsd (#64)

### Fix some issues to allow compilation of `uccsd`

- Fix signed int overflow in `uccsd`
- allow specifying shots in `vqe`

Towards: NVIDIA/cuda-quantum#2357

---------

Signed-off-by: Anna Gringauze <agringauze@nvidia.com>

libs/solvers/lib/stateprep/uccsd.cpp

@@ -407,10 +407,10 @@ __qpu__ std::size_t getNumVirtualBeta(std::size_t numElectrons,
 __qpu__ void uccsd(cudaq::qview<> qubits, const std::vector<double> &thetas,
                    std::size_t numElectrons, std::size_t spin) {
 
-  auto numOccAlpha = getNumOccupiedAlpha(numElectrons, spin, qubits.size());
-  auto numOccBeta = getNumOccupiedBeta(numElectrons, spin, qubits.size());
-  auto numVirtAlpha = getNumVirtualAlpha(numElectrons, spin, qubits.size());
-  auto numVirtBeta = getNumVirtualBeta(numElectrons, spin, qubits.size());
+  int numOccAlpha = getNumOccupiedAlpha(numElectrons, spin, qubits.size());
+  int numOccBeta = getNumOccupiedBeta(numElectrons, spin, qubits.size());
+  int numVirtAlpha = getNumVirtualAlpha(numElectrons, spin, qubits.size());
+  int numVirtBeta = getNumVirtualBeta(numElectrons, spin, qubits.size());
   std::vector<std::size_t> occupiedAlpha(numOccAlpha),
       virtualAlpha(numVirtAlpha), occupiedBeta(numOccBeta),
       virtualBeta(numVirtBeta);
@@ -433,10 +433,11 @@ __qpu__ void uccsd(cudaq::qview<> qubits, const std::vector<double> &thetas,
       occupiedBeta[counter] = i * 2 + 1;
       counter++;
     }
-    counter = 0;
 
+    counter = 0;
     for (std::size_t i = 0; i < numVirtBeta; i++) {
       virtualBeta[counter] = i * 2 + numElectrons - 1;
+      counter++;
     }
 
   } else {
@@ -507,18 +508,18 @@ __qpu__ void uccsd(cudaq::qview<> qubits, const std::vector<double> &thetas,
         }
 
   counter = 0;
-  for (std::size_t p = 0; p < numOccAlpha - 1; p++)
-    for (std::size_t q = p + 1; q < numOccAlpha; q++)
-      for (std::size_t r = 0; r < numVirtAlpha - 1; r++)
-        for (std::size_t s = r + 1; s < numVirtAlpha; s++)
+  for (int p = 0; p < numOccAlpha - 1; p++)
+    for (int q = p + 1; q < numOccAlpha; q++)
+      for (int r = 0; r < numVirtAlpha - 1; r++)
+        for (int s = r + 1; s < numVirtAlpha; s++)
           counter++;
 
   std::vector<std::size_t> doublesAlpha(4 * counter);
   counter = 0;
-  for (std::size_t p = 0; p < numOccAlpha - 1; p++)
-    for (std::size_t q = p + 1; q < numOccAlpha; q++)
-      for (std::size_t r = 0; r < numVirtAlpha - 1; r++)
-        for (std::size_t s = r + 1; s < numVirtAlpha; s++) {
+  for (int p = 0; p < numOccAlpha - 1; p++)
+    for (int q = p + 1; q < numOccAlpha; q++)
+      for (int r = 0; r < numVirtAlpha - 1; r++)
+        for (int s = r + 1; s < numVirtAlpha; s++) {
           doublesAlpha[counter] = occupiedAlpha[p];
           counter++;
           doublesAlpha[counter] = occupiedAlpha[q];
@@ -530,17 +531,17 @@ __qpu__ void uccsd(cudaq::qview<> qubits, const std::vector<double> &thetas,
         }
 
   counter = 0;
-  for (std::size_t p = 0; p < numOccBeta - 1; p++)
-    for (std::size_t q = p + 1; q < numOccBeta; q++)
-      for (std::size_t r = 0; r < numVirtBeta - 1; r++)
-        for (std::size_t s = r + 1; s < numVirtBeta; s++)
+  for (int p = 0; p < numOccBeta - 1; p++)
+    for (int q = p + 1; q < numOccBeta; q++)
+      for (int r = 0; r < numVirtBeta - 1; r++)
+        for (int s = r + 1; s < numVirtBeta; s++)
           counter++;
   std::vector<std::size_t> doublesBeta(4 * counter);
   counter = 0;
-  for (std::size_t p = 0; p < numOccBeta - 1; p++)
-    for (std::size_t q = p + 1; q < numOccBeta; q++)
-      for (std::size_t r = 0; r < numVirtBeta - 1; r++)
-        for (std::size_t s = r + 1; s < numVirtBeta; s++) {
+  for (int p = 0; p < numOccBeta - 1; p++)
+    for (int q = p + 1; q < numOccBeta; q++)
+      for (int r = 0; r < numVirtBeta - 1; r++)
+        for (int s = r + 1; s < numVirtBeta; s++) {
           doublesBeta[counter] = occupiedBeta[p];
           counter++;
           doublesBeta[counter] = occupiedBeta[q];

libs/solvers/python/bindings/solvers/py_solvers.cpp

@@ -132,6 +132,9 @@ class PythonOptimizer : public optim::optimizer {
     if (kwargs.contains("verbose"))
       kwargs.attr("pop")("verbose");
 
+    if (kwargs.contains("shots"))
+      kwargs.attr("pop")("shots");
+
     if (initParams.empty())
       initParams.resize(dim);
 

libs/solvers/python/tests/test_uccsd.py

@@ -141,3 +141,41 @@ def ansatz(thetas: list[float]):
 
     print(energy)
     assert np.isclose(energy, -107.6059, 1e-2)
+
+
+def test_uccsd_loops():
+    repro_num_electrons = 2
+    repro_num_qubits = 8
+
+    repro_thetas = [
+        -0.00037043841404585794, 0.0003811110195084151, 0.2286823796532558,
+        -0.00037043841404585794, 0.0003811110195084151, 0.2286823796532558,
+        -0.00037043841404585794, 0.0003811110195084151, 0.2286823796532558,
+        -0.00037043841404585794, 0.0003811110195084151, 0.2286823796532558,
+        -0.00037043841404585794, 0.0003811110195084151, 0.2286823796532558,
+        -0.00037043841404585794, 0.0003811110195084151, 0.2286823796532558,
+        -0.00037043841404585794, 0.0003811110195084151, 0.2286823796532558,
+        -0.00037043841404585794, 0.0003811110195084151, 0.2286823796532558,
+        -0.00037043841404585794, 0.0003811110195084151, 0.2286823796532558
+    ]
+
+    @cudaq.kernel
+    def repro_trial_state(qubits: cudaq.qvector, num_electrons: int,
+                          thetas: list[float]):
+        for i in range(num_electrons):
+            x(qubits[i])
+        solvers.stateprep.uccsd(qubits, thetas, num_electrons, 0)
+
+    @cudaq.kernel
+    def repro():
+        repro_qubits = cudaq.qvector(repro_num_qubits)
+        repro_trial_state(repro_qubits, repro_num_electrons, repro_thetas)
+
+    counts = cudaq.sample(repro, shots_count=1000)
+    assert len(counts) == 6
+    assert '00000011' in counts
+    assert '00000110' in counts
+    assert '00010010' in counts
+    assert '01000010' in counts
+    assert '10000001' in counts
+    assert '11000000' in counts