Discussion for speed up in wavefront.power

[syhaffert]

I have played around a bit with timing the calculation of the power of a wavefront. I have found that we can actually increase the speed of that part of the code by a factor of about 2.5. The following code tests 4 different approaches to calculating the total power. HCIPy currently uses the method from function h()

    import numpy as np
    import numexpr as ne
    import time

    npix = 256
    z = np.random.randn(npix * npix) + 1j * np.random.randn(npix * npix)

    def f():
        return np.abs(z)**2
    
    def g():
        test = np.real(z)**2 
        test += np.imag(z)**2
        return test

    def h():
        variables = {'field': z}
        return ne.evaluate('real(abs(field))**2 ', local_dict=variables)

    def k():
        variables = {'field': z}
        return ne.evaluate('real(field)**2 + imag(field)**2', local_dict=variables)

    nit = 10000
    labels = ['f', 'g', 'h', 'k']
    for fi, func in enumerate([f, g, h, k]):
        
        start = time.time()
        for i in range(nit):
            func()
        end = time.time()
        elapsed = (end - start) / nit
        print("Func {:s} Elapsed: {:0.2f} us".format(labels[fi], 1e6 * elapsed))

I measure the following timings Func f Elapsed: 238 +- 40 us Func g Elapsed: 144 +- 14 us Func h Elapsed: 257 +- 14 us Func k Elapsed: 167 +- 16 us

So it seems like we can speed up the power calculation by either implement function k or g. I would prefer function g because that is a pure numpy function that also immediately makes it compatible for the other back ends.