def _transfer_ray_4x4_1(field,
wavenumbers,
layer,
dmat1,
dmat2,
beta=0,
phi=0,
nsteps=1,
betamax=BETAMAX,
out=None):
d, epsv, epsa = layer
kd = wavenumbers * d
alpha, f = alphaf(beta, phi, epsv, epsa)
p = phasem(alpha, kd[..., None, None])
e = E_mat(f, mode=None)
ei = inv(e)
for j in range(nsteps):
field = dotmf(dmat1, field, out=out)
field = ifft2(field, out=field)
field = dotmdmf(e, p, ei, field, out=field)
field = fft2(field, out=field)
field = dotmf(dmat2, field, out=out)
return field
def _transfer_ray_4x4_4(field,
wavenumbers,
layer,
beta=0,
phi=0,
nsteps=1,
dmat=None,
out=None):
d, epsv, epsa = layer
if dmat is None:
kd = wavenumbers * d
else:
kd = wavenumbers * d / 2
alpha, f, fi = alphaffi(beta, phi, epsv, epsa)
p = phasem(alpha, kd[..., None, None])
for j in range(nsteps):
if dmat is None:
field = dotmdmf(f, p, fi, field, out=out)
else:
field = dotmdmf(f, p, fi, field, out=out)
field = diffract(field, dmat, out=field)
field = dotmdmf(f, p, fi, field, out=field)
return field
def propagate_4x4_full(field,
wavenumbers,
layer,
nsteps=1,
betamax=BETAMAX,
out=None):
shape = field.shape[-2:]
d, epsv, epsa = layer
kd = wavenumbers * d / nsteps
if out is None:
out = np.empty_like(field)
out_af = None
pm = None
ii, jj = np.meshgrid(range(shape[0]),
range(shape[1]),
copy=False,
indexing="ij")
for step in range(nsteps):
for i in range(len(wavenumbers)):
ffield = fft2(field[..., i, :, :, :])
ofield = np.zeros_like(out[..., i, :, :, :])
b, p = betaphi(shape, wavenumbers[i])
mask = b < betamax
amplitude = ffield[..., mask]
betas = b[mask]
phis = p[mask]
iind = ii[mask]
jind = jj[mask]
for bp in sorted(zip(range(len(betas)), betas, phis, iind, jind),
key=lambda el: el[1],
reverse=False):
#for j,bp in enumerate(zip(betas,phis,iind,jind)):
j, beta, phi, ieig, jeig = bp
out_af = alphaffi(beta, phi, epsv, epsa, out=out_af)
alpha, f, fi = out_af
pm = phasem(alpha, kd[i], out=pm)
w = eigenwave(amplitude.shape[:-1] + shape,
ieig,
jeig,
amplitude=amplitude[..., j])
w = dotmdmf(f, pm, fi, w, out=w)
np.add(ofield, w, ofield)
out[..., i, :, :, :] = ofield
field = out
return out
def _transfer_ray_4x4_1_windows(windows, field, wavenumbers, layer, dmat1, dmat2, betas, phis,
nsteps = 1, out = None,
betamax = BETAMAX):
d, epsv, epsa = layer
kd = wavenumbers*d
if out is None:
out = np.zeros_like(field)
for j in range(nsteps):
field = dotmf(dmat1,field)
for window, beta, phi in zip(windows, betas, phis):
alpha, f = alphaf(beta,phi,epsv,epsa)
p = phasem(alpha,kd[...,None,None])
e = E_mat(f, mode = None)
ei = inv(e)
f = field * window
f = ifft2(f, out = f)
f = dotmdmf(e,p,ei,f, out = f)
f = fft2(field, out = f)
out += f
out = dotmf(dmat2,out, out = out)
return out
def _transfer_ray_4x4_2(field,
wavenumbers,
layer,
beta=0,
phi=0,
nsteps=1,
dmatpn=None,
out=None,
tmpdata=None):
_out = {} if tmpdata is None else tmpdata
d, epsv, epsa = layer
if dmatpn is None:
kd = wavenumbers * d
else:
dmatp, dmatn = dmatpn
kd = wavenumbers * d / 2
alpha, f, fi = alphaffi(beta, phi, epsv, epsa, out=_out.get("affi"))
p = phasem(alpha, kd[..., None, None], out=_out.get("p"))
if tmpdata is not None:
tmpdata["affi"] = (alpha, f, fi)
tmpdata["p"] = p
if dmatpn is not None:
e = E_mat(f, mode=None)
ei = inv(e, out=_out.get("ei"))
if tmpdata is not None:
tmpdata["ei"] = ei
for j in range(nsteps):
if dmatpn is None:
field = dotmdmf(f, p, fi, field, out=out)
else:
field = dotmdmf(e, p, fi, field, out=out)
diffract(field[..., 0::2, :, :], dmatp, out=field[..., 0::2, :, :])
diffract(field[..., 1::2, :, :], dmatn, out=field[..., 1::2, :, :])
field = dotmdmf(f, p, ei, field, out=field)
return field