how to iterate through an array of complex numbers

Question

I'm trying to calculate the energy of a complex valued signal. Passing an array of complex numbers into the energy function, it separates the real and imaginary parts of the number and converts them into their polar equivalents. It then returns the sum of the squares of the real parts of each complex number. Everytime I try to call the energy function it says that the arctan2 ufunc is not supported for the input types.

def toExponential(a, b):
  c = np.sqrt(a**2 + b**2)
  d = np.arctan2(b,a)
  return (c,d)

def energy(x):
  sum = 0
  for i in x:
    e = ((i + np.conj(i))/2)
    f = ((i - np.conj(i)/(1j * 2)))
    r,i = toExponential(e,f)
    sum = r**2 + sum
  return sum

Answer 1

I think you are passing e and f to the to np.arctan2(b,a) ,instead of the real and imaginary parts of the complex number magnitude, phase = np.abs(i), np.angle(i)

Try this out

def energy(x):
        sum = 0
        for i in x:
            magnitude, phase = np.abs(i), np.angle(i)
            sum += magnitude**2
        return sum

The magnitude and phase of each complex number are extracted in this example using the numpy.abs() and numpy.angle() methods, respectively. The energy of a complex signal is then determined by adding the squares of the complex numbers' magnitudes, which is the appropriate procedure.

Answer 2

If speed is important, here is a vectorized version of @MohamedFathallah's answer:

def energy(x):
    return np.sum(np.abs(x)**2)

or

def energy(x):
    return np.sum(np.real(x)**2 + np.imag(x)**2)