Commit 1236de9a authored by Leon Merten Lohse's avatar Leon Merten Lohse
Browse files

code style

parent 902d39a3
Pipeline #174171 passed with stage
in 41 seconds
......@@ -64,9 +64,6 @@ class DiscreteHankelTransform:
self._matrix = hankelMatrix(N, n)
def __call__(self, x):
return self._matrix @ x
from pyfftw.interfaces.numpy_fft import fft,fftn,fftfreq,ifftn,ifftshift,fftshift
from pyfftw.interfaces.numpy_fft import fftn, ifftn, ifftshift, fftshift
import numpy as np
from functools import reduce
......@@ -7,7 +7,10 @@ from .misc import fftfreqn
_lambda0 = 1. # all lengths are in units of the wavelength
_k0 = 2 * np.pi / _lambda0
_compute_potential = lambda n_, k_=_k0 : k_*k_ * (1 - n_*n_)
def _compute_potential(n_, k_=_k0):
return k_*k_ * (1 - n_*n_)
def squaresum(a):
......@@ -40,10 +43,9 @@ class MultislicePropagator():
dtype = np.complex128
def __init__(self, u0, d, wl=1.):
ndim = len(d)
# ndim = len(d)
# TODO: check input
# assert u0 is array
......@@ -53,21 +55,20 @@ class MultislicePropagator():
self._dz = d[0]
self._dperp = np.array(d[1:])
self._ones = np.ones(u0.shape, dtype = self.dtype)
self._ones = np.ones(u0.shape, dtype=self.dtype)
self._k = _k0 / wl
self._fourierKernel = rayleighSommerfeldTF(self._ones.shape, self._dperp, self._k, self._dz)
self.u = ifftshift(u0)
def step(self, n):
nshift = ifftshift(n)
F = _compute_potential(nshift, self._k) * self._ones
self._realKernel = np.exp(-1j * self._dz / (2 * self._k) * F )
self._realKernel = np.exp(-1j * self._dz / (2 * self._k) * F)
up = self.u
self.u = ifftn(self._fourierKernel * fftn(self._realKernel * up))
......@@ -88,7 +89,6 @@ class FDPropagator2d(fd.Solver2d):
super().__init__(Az, Axx, F0, u0, dz, dx)
def step(self, n, boundary):
F = _compute_potential(n, self._k)
......@@ -112,7 +112,6 @@ class FDPropagatorCS(fd.Solver2dfull):
super().__init__(Az, Axx, Ax, F0, u0, dz, dx)
def step(self, n, boundary):
F = _compute_potential(n, self._k) * self._x
......@@ -133,10 +132,8 @@ class FDPropagator3d(fd.Solver3d):
super().__init__(Az, Axx, Ayy, F0, u0, dz, dy, dx)
def step(self, n, boundary):
F = _compute_potential(n, self._k)
return super().step(F, boundary)
......@@ -2,7 +2,7 @@ from . import hankel
from .misc import fftfreqn, gridn
import numpy as np
from pyfftw.interfaces.numpy_fft import fft,fftn,fftfreq,ifftn,ifftshift,fftshift
from pyfftw.interfaces.numpy_fft import fftn, ifftn
from scipy.signal import fftconvolve
_lambda0 = 1. # all lengths are in units of the wavelength
......@@ -12,7 +12,7 @@ _k0 = 2 * np.pi / _lambda0
def fresnelKernelCS(N, fresnelNumber):
hFreq = hankel.hankelFreq(N)
kern = np.exp(-1j * np.pi * hFreq**2 / fresnelNumber )
kern = np.exp(-1j * np.pi * hFreq**2 / fresnelNumber)
return kern
......@@ -27,8 +27,7 @@ class FresnelPropagatorCS:
kern = fresnelKernelCS(N, fresnelNumber)
# TODO: express this nicer (kern is a diagonal matrix)
self._matrix = Y @ (kern[:,None] * Y)
self._matrix = Y @ (kern[:, None] * Y)
def __call__(self, u):
......@@ -64,15 +63,13 @@ class FresnelTFPropagator:
self._kernel = fresnelTFKernel(shape, fresnelNumbers)
def __call__(self, u):
uprop = ifftn( self._kernel * fftn(u))
uprop = ifftn(self._kernel * fftn(u))
return uprop
def fresnelIRKernel(shape, fresnelNumbers):
ndim = len(shape)
......@@ -83,13 +80,12 @@ def fresnelIRKernel(shape, fresnelNumbers):
kernel = np.ones(2 * np.array(shape, dtype=int) - 1, dtype=np.complex128)
# phase factor
kernel *= np.prod(np.sqrt(np.abs(fresnelNumbers)) * np.exp((-1j * np.pi / 4) * np.sign(fresnelNumbers)) )
kernel *= np.prod(np.sqrt(np.abs(fresnelNumbers)) * np.exp((-1j * np.pi / 4) * np.sign(fresnelNumbers)))
x = gridn(shape)
for dim in range(ndim):
xdim = x[dim] # / shape[dim]
kernel *= np.exp(1j*np.pi*fresnelNumbers[dim] * xdim**2)
kernel *= np.exp(1j*np.pi*fresnelNumbers[dim] * x[dim]**2)
return kernel
......@@ -106,7 +102,6 @@ class FresnelIRPropagator:
self._kernel = fresnelIRKernel(shape, fresnelNumbers)
def __call__(self, u):
uprop = fftconvolve(u, self._kernel, mode='valid')
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment