From 7dde547cbbb4ffabc706b93b4ee725992a70d1dd Mon Sep 17 00:00:00 2001
From: WrathfulSpatula <stranoj@gmail.com>
Date: Sat, 19 Oct 2024 21:32:26 -0400
Subject: [PATCH] Debug

---
 src/common/qengine.cl | 33 +++++++++++++++++++++++++++------
 src/qengine/state.cpp | 23 +++++++++++++++++++----
 2 files changed, 46 insertions(+), 10 deletions(-)

diff --git a/src/common/qengine.cl b/src/common/qengine.cl
index 636093c75..092eef321 100644
--- a/src/common/qengine.cl
+++ b/src/common/qengine.cl
@@ -587,8 +587,8 @@ void kernel decomposeprob(global cmplx* stateVec, constant bitCapIntOcl4* bitCap
         for (bitCapIntOcl k = 0U; k < partPower; k++) {
             bitCapIntOcl l = j | (k << start);
 
-            cmplx amp = stateVec[l];
-            real1 nrm = dot(amp, amp);
+            const cmplx amp = stateVec[l];
+            const real1 nrm = dot(amp, amp);
             partProb += nrm;
             partStateAngle[k] += arg(amp) * nrm;
         }
@@ -606,14 +606,28 @@ void kernel decomposeprob(global cmplx* stateVec, constant bitCapIntOcl4* bitCap
             l |= (k ^ l) << len;
             l = j | l;
 
-            cmplx amp = stateVec[l];
-            real1 nrm = dot(amp, amp);
+            const cmplx amp = stateVec[l];
+            const real1 nrm = dot(amp, amp);
             partProb += nrm;
             remainderStateAngle[k] += arg(amp) * nrm;
         }
 
         partStateProb[lcv] = partProb;
     }
+
+    for (bitCapIntOcl lcv = ID; lcv < remainderPower; lcv += Nthreads) {
+        const real1 prob = remainderStateProb[lcv];
+        if (prob > REAL1_EPSILON) {
+            remainderStateAngle[lcv] /= prob;
+        }
+    }
+
+    for (bitCapIntOcl lcv = ID; lcv < partPower; lcv += Nthreads) {
+        const real1 prob = partStateProb[lcv];
+        if (prob > REAL1_EPSILON) {
+            partStateAngle[lcv] /= prob;
+        }
+    }
 }
 
 void kernel decomposeamp(
@@ -647,7 +661,7 @@ void kernel disposeprob(global cmplx* stateVec, constant bitCapIntOcl4* bitCapIn
         for (bitCapIntOcl k = 0U; k < partPower; k++) {
             bitCapIntOcl l = j | (k << start);
 
-            cmplx amp = stateVec[l];
+            const cmplx amp = stateVec[l];
             partProb += dot(amp, amp);
         }
 
@@ -662,10 +676,17 @@ void kernel disposeprob(global cmplx* stateVec, constant bitCapIntOcl4* bitCapIn
             l |= (k ^ l) << len;
             l = j | l;
 
-            cmplx amp = stateVec[l];
+            const cmplx amp = stateVec[l];
             remainderStateAngle[k] += arg(amp) * dot(amp, amp);
         }
     }
+
+    for (bitCapIntOcl lcv = ID; lcv < remainderPower; lcv += Nthreads) {
+        const real1 prob = remainderStateProb[lcv];
+        if (prob > REAL1_EPSILON) {
+            remainderStateAngle[lcv] /= prob;
+        }
+    }
 }
 
 void kernel dispose(global cmplx* stateVec, constant bitCapIntOcl4* bitCapIntOclPtr, global cmplx* nStateVec)
diff --git a/src/qengine/state.cpp b/src/qengine/state.cpp
index 7cfdd9cdb..696f5c0ac 100644
--- a/src/qengine/state.cpp
+++ b/src/qengine/state.cpp
@@ -1290,6 +1290,18 @@ void QEngineCPU::DecomposeDispose(bitLenInt start, bitLenInt length, QEngineCPUP
                 remainderStateAngle[k] += arg(amp) * nrm;
             }
         });
+        par_for(0U, remainderPower, [&](const bitCapIntOcl& lcv, const unsigned& cpu) {
+            const real1 prob = remainderStateProb[lcv];
+            if (prob > amplitudeFloor) {
+                remainderStateAngle[lcv] /= prob;
+            }
+        });
+        par_for(0U, partPower, [&](const bitCapIntOcl& lcv, const unsigned& cpu) {
+            const real1 prob = partStateProb[lcv];
+            if (prob > amplitudeFloor) {
+                partStateAngle[lcv] /= prob;
+            }
+        });
     } else {
         par_for(0U, remainderPower, [&](const bitCapIntOcl& lcv, const unsigned& cpu) {
             bitCapIntOcl j;
@@ -1298,10 +1310,7 @@ void QEngineCPU::DecomposeDispose(bitLenInt start, bitLenInt length, QEngineCPUP
 
             for (bitCapIntOcl k = 0U; k < partPower; ++k) {
                 bitCapIntOcl l = j | (k << start);
-
-                const complex amp = stateVec->read(l);
-                const real1 nrm = norm(amp);
-                remainderStateProb[lcv] += nrm;
+                remainderStateProb[lcv] += norm(stateVec->read(l));
             }
         });
 
@@ -1317,6 +1326,12 @@ void QEngineCPU::DecomposeDispose(bitLenInt start, bitLenInt length, QEngineCPUP
                 remainderStateAngle[k] += arg(amp) * norm(amp);
             }
         });
+        par_for(0U, remainderPower, [&](const bitCapIntOcl& lcv, const unsigned& cpu) {
+            const real1 prob = remainderStateProb[lcv];
+            if (prob > amplitudeFloor) {
+                remainderStateAngle[lcv] /= prob;
+            }
+        });
     }
 
     if (destination) {