Refactor coherent state implementation

Signed-off-by: Thien Nguyen <[email protected]>

python/cudaq/operator/cudm_state.py

@@ -37,18 +37,21 @@ def ket2dm(ket: cupy.ndarray) -> cupy.ndarray:
 
 # Helper to create a 'coherent' state as a state vector.
 def coherent_state(N: int, alpha: float):
-    sqrtn = cupy.sqrt(cupy.arange(0, N, dtype=cupy.complex128))
-    sqrtn[0] = 1
-    data = alpha / sqrtn
-    data[0] = cupy.exp(-cupy.abs(alpha)**2 / 2.0)
-    cupy.cumprod(data, out=sqrtn)
-    return sqrtn
+    # Computes coherent state amplitudes in the Fock basis
+    # Ref: https://en.wikipedia.org/wiki/Coherent_state
+    fock_amplitudes = cupy.zeros(N, dtype=cupy.complex128)
+    amplitude = numpy.exp(-numpy.abs(alpha)**2 / 2.0)
+    for n in range(N):
+        fock_amplitudes[n] = amplitude
+        amplitude *= (alpha / numpy.sqrt(n + 1))
+    return fock_amplitudes
 
 
 # Helper to create a coherent state as a density matrix.
 def coherent_dm(N: int, alpha: float):
     return ket2dm(coherent_state(N, alpha))
 
+
 # A Python wrapper of `CuDensityMatState` state.
 class CuDensityMatState(object):
     __ctx = None