Working version with focusing and defocusing. It seems the code is also more stable

Known problems: change kx so it doesn't have 0. Make fourier plots consistent.