Corrected errors in input+output plane coordinate systems and added lens example

[villadsegede]

There is a mistake in the coordinate system definitions of the current version which means that for example a lens does not focus exactly on the optical axis. If you try and run the supplied example in an old version of the code you would get:

And with the new version you would instead get the following on-axis responses:

P.S. I have also updated the calculation of kz in order to include non-propagating modes. If not, the supplied example will fail

P.P.S. Changes only made to the monochromatic simulator

[rafael-fuente]

Thank you for correcting this mistake. I was aware there was an error with the FFT bin-center frequency, but since I thought in the scale I was testing it was small, I didn't take the time to fix it. But in this example, by plotting the PSF, you clearly show how it can be relevant!

About including the evanescent modes, it can also be done using where method:

The naive method raises a warning since it's evaluating a negative square root using float data type:

argument = (2 * bd.pi / self.λ) ** 2 - kx ** 2 - ky ** 2
kz = bd.where(argument >= 0, bd.sqrt(argument), 1j*bd.sqrt(-argument) )

But using this simple trick it can be solved:

argument = (2 * bd.pi / self.λ) ** 2 - kx ** 2 - ky ** 2
tmp = bd.sqrt(bd.abs(argument))
kz = bd.where(argument >= 0, tmp, 1j*tmp)

It's slightly faster than the implementation you provide and when using GPU doesn't force Cuda stream synchronization.

Anyway, the ideal implementation should be done using a compiled extension, for example:

import numba
from numba import jit

@numba.jit("complex128[:,:](f8[:,:], i8, i8)", nopython=True, nogil=True)
def get_kz(argument, Nx, Ny):

    kz = bd.zeros((Ny, Nx),dtype='complex128')

    for i in range(0,Ny):
        for j in range(0,Nx):
            if argument[i,j] >= 0:
                kz[i,j] = bd.sqrt(argument[i,j] )
            else:
                kz[i,j] = 1j*bd.sqrt(-argument[i,j])
    return kz

But here the where option is enough for simplicity.

There are also a lot of other further improvements that can be done in this repo for serious applications, for example, it makes sense to implement a plane scaling factor m for each propagation step. I'll test this if I find some time in the future.