def plot_r_q_polar(self, flat=False, tilt=False, show=True):
"""Plot amplitude and phase in real and angular space.
To see self.Er and its transform exactly, set flat=False and tilt=True.
Args:
flat (bool): Plot Er_flat instead of Er (both domains).
tilt (bool): Include the central tilt in the real and angular space phase plots.
show:
Returns:
GraphicsLayoutWidget: Contains the four plots and a heading.
plots (RQ tuple of AbsPhi tuples): The AlignedPlotItems.
"""
Er = self.Er_flat if flat else self.Er
Eq = math.fft2(Er)
if not tilt:
Er = Er * mathx.expj(-(self.qs_center[0] * self.x +
self.qs_center[1] * self.y))
Eq = Eq * mathx.expj(self.rs_center[0] * self.kx +
self.rs_center[1] * self.ky)
glw = pg.GraphicsLayoutWidget()
glw.addLabel(self.title_str)
glw.nextRow()
gl = glw.addLayout()
plots = plotting.plot_r_q_polar(self.rs_support, Er, self.rs_center,
self.qs_center, gl, Eq)
glw.resize(830, 675)
if show:
glw.show()
return glw, plots
def calc_plane_to_curved_flat_arbitrary_factors(k,
rs_support,
num_pointss,
z,
xo,
yo,
qs_center=(0, 0),
kz_mode='local_xy'):
"""Calculate factors for propagation from plane to arbitrarily sampled curved surface.
Args:
k (scalar): Wavenumber.
rs_support (scalar or pair of): Input aperture along x and y.
num_pointss (int or tuple of): Number of input samples along x and y. Referred to as K and L below.
z (M*N array): Propagation distances.
xo (Mx1 array): Output x values.
yo (M array): Output y values.
qs_center (tuple or pair of): Center of transverse wavenumber support.
kz_mode (str): 'paraxial' or 'local_xy'.
Returns:
invTx (M*K array): Inverse transform factor.
gradxinvTx (M*K array): X-derivative of inverse transform factor.
invTy (L*N array): Y inverse transform factor.
gradyinvTy (L*N array): Y-derivative of inverse transform factor.
Px (M*N*K array): X propagation.
Py (M*N*L array): Y propagation.
Tx (K*K array): X transform factor.
Ty (L*L array): Y transform factor.
"""
assert kz_mode in ('local_xy', 'paraxial')
Tx = make_fft_matrix(num_pointss[0])
if num_pointss[1] == num_pointss[0]:
Ty = Tx
else:
Ty = make_fft_matrix(num_pointss[1])
if kz_mode == 'local_xy':
fx, fy, delta_kz, delta_gxkz, delta_gykz = calc_quadratic_kz_correction(
k, *qs_center)
zx = z * fx
zy = z * fy
else:
zx = z
zy = z
delta_kz = k
kx, ky = sa.calc_kxky(rs_support, num_pointss, qs_center)
# kx & ky are the right-most indices of the propagators.
Px = calc_propagator_quadratic_1d(
k, kx[:, 0], zx[:, :, None]) * mathx.expj(delta_kz * z[:, :, None])
Py = calc_propagator_quadratic_1d(k, ky, zy[:, :, None])
if kz_mode == 'local_xy':
Px *= mathx.expj(delta_gxkz * kx[:, 0] * z[:, :, None])
Py *= mathx.expj(delta_gykz * ky * z[:, :, None])
invTx = make_ifft_arbitrary_matrix(rs_support[0], num_pointss[0],
qs_center[0], xo[:, 0])
invTy = make_ifft_arbitrary_matrix(rs_support[1], num_pointss[1],
qs_center[1], yo)
gradxinvTx = 1j * kx[:, 0] * invTx
gradyinvTy = 1j * ky * invTy
return invTx, gradxinvTx, invTy, gradyinvTy, Px, Py, Tx, Ty
def plot_r_q_polar(self, flat=False, tilt=False, show=True):
"""Plot approximate plane profile and surface z relative to z_mean."""
app = self.app
Er = app.Er_flat if flat else app.Er
Eq = math.fft2(Er)
if not tilt:
Er = Er * mathx.expj(-(app.qs_center[0] * app.x +
app.qs_center[1] * app.y))
Eq = Eq * mathx.expj(app.rs_center[0] * app.kx +
app.rs_center[1] * app.ky)
glw = pg.GraphicsLayoutWidget()
glw.addLabel(self.title_str)
glw.nextRow()
gl = glw.addLayout()
plot = gl.addAlignedPlot(labels={'left': 'y (mm)', 'bottom': 'x (mm)'})
x, y, zu = sa.unroll_r(self.rs_support, self.z, self.rs_center)
image = plot.image((zu - self.mean_z) * 1e3,
rect=pg.axes_to_rect(x * 1e3, y * 1e3),
lut=pg.get_colormap_lut('bipolar'))
gl.addHorizontalSpacer(10)
gl.addColorBar(image=image, rel_row=2, label='Relative z (mm)')
glw.nextRow()
glw.addLabel('Approximate planar profile')
glw.nextRow()
gl = glw.addLayout()
plots = plotting.plot_r_q_polar(app.rs_support, Er, app.rs_center,
app.qs_center, gl, Eq)
glw.resize(830, 675)
if show:
glw.show()
return glw, plots
def prepare_curved_paraxial_propagation_1d(k,
r_support,
Er,
z,
m,
r_center=0,
q_center=0,
axis=-1,
carrier=True):
"""Modifies Er."""
assert not (np.isclose(z, 0) ^ np.isclose(m, 1))
if np.isclose(z, 0):
return 1, 1, r_center
num_points = Er.shape[axis]
roc = z / (m - 1)
r = sa.calc_r(r_support, num_points, r_center, axis)
q = sa.calc_q(r_support, num_points, q_center, axis) #-q_center
Er *= math.calc_quadratic_phase_1d(-k, r - r_center, roc)
propagator = math.calc_propagator_quadratic_1d(k * m,
q - k * r_center / roc, z)
ro_center = r_center + q_center / k * z
ro = sa.calc_r(r_support * m, num_points, ro_center, axis)
post_factor = math.calc_quadratic_phase_1d(k, ro, roc + z) * mathx.expj(
k * (r_center**2 / (2 * roc) - r_center * ro / (roc * m))) / m**0.5
if carrier:
post_factor *= mathx.expj(k * z)
return propagator, post_factor, ro_center
def test_curved_interface_collimate_offset():
lamb = 587.6e-9
waist0 = 30e-6
num_points = 2**7
k = 2 * np.pi / lamb
r0_support = (np.pi * num_points)**0.5 * waist0
# Define x lateral offset of beam.
r_offsets = np.asarray((4e-3, -2e-3))
# We center the numerical window on the beam.
rs_center = r_offsets
x0, y0 = asbp.calc_xy(r0_support, num_points, rs_center)
# Create input beam.
Er0 = asbp.calc_gaussian(k, x0, y0, waist0, r0s=r_offsets)
# We will propagate it a distance f.
f = 100e-3
z_R = np.pi * waist0**2 / lamb
m = (1 + (f / z_R)**2)**0.5
# We collimate it with a spherical interface.
n = 1.5
roc = f * (n - 1)
# At the interface, support is expanded by the curved wavefront propagation.
r1_support = m * r0_support
x1, y1 = asbp.calc_xy(r1_support, num_points, rs_center)
# Calculate and plot interface sag.
sag = functions.calc_sphere_sag_xy(roc, x1, y1)
xu, yu, sagu = asbp.unroll_r(r1_support, sag)
# Propagate to curved surface.
Er1, _ = asbp.propagate_plane_to_curved_spherical(k, r0_support, Er0,
f + sag, m, rs_center)
xu, yu, Er1u = asbp.unroll_r(r1_support, Er1)
#
# x2, y2=asbp.calc_xy(r1_support, num_points)
qs_centers2 = -r_offsets / f * k
Er2, propagator = asbp.invert_plane_to_curved_spherical(
k * n,
r1_support,
Er1,
-f * n + sag,
1, (0, 0),
qs_centers2,
max_iterations=10) # -x_offset/f*k
xu, yu, Er2u = asbp.unroll_r(r1_support, Er2)
tilt_factor = mathx.expj(-(xu * r_offsets[0] + yu * r_offsets[1]) / f * k)
##
waist2 = f * lamb / (np.pi * waist0)
Er2_theory = asbp.calc_gaussian(k, x1, y1, waist2) * tilt_factor
Er2_theory *= mathx.expj(np.angle(Er2[0, 0] / Er2_theory[0, 0]))
assert mathx.allclose(abs(Er2), abs(Er2_theory), 2e-2)
if 0:
##
Er2_fig = asbp.plot_r_q_polar(r1_support, Er2, qs_center=qs_centers2)
Er2_theory_fig = asbp.plot_r_q_polar(r1_support,
Er2_theory,
qs_center=qs_centers2)
def calc_plane_to_curved_flat_factors(k,
rs_support,
num_pointss,
z,
qs_center=(0, 0),
kz_mode='local_xy'):
"""Calculate factors for propagation from plane to uniformly sampled curved surface.
Args:
k:
rs_support:
num_pointss:
z:
qs_center:
kz_mode (str): 'paraxial' or 'local_xy'.
Returns:
invTx (2D array):
gradxinvTx (2D array):
invTy (2D array):
gradyinvTy (2D array):
Px (3D array):
Py (3D array):
Tx (2D array):
Ty (2D array):
"""
assert kz_mode in ('local_xy', 'paraxial')
Tx = make_fft_matrix(num_pointss[0])
if num_pointss[1] == num_pointss[0]:
Ty = Tx
else:
Ty = make_fft_matrix(num_pointss[1])
if kz_mode == 'local_xy':
fx, fy, delta_kz, delta_gxkz, delta_gykz = calc_quadratic_kz_correction(
k, *qs_center)
zx = z * fx
zy = z * fy
else:
zx = z
zy = z
delta_kz = k
kx, ky = sa.calc_kxky(rs_support, num_pointss, qs_center)
# kx & ky are the right-most indices of the propagators.
Px = calc_propagator_quadratic_1d(
k, kx[:, 0], zx[:, :, None]) * mathx.expj(delta_kz * z[:, :, None])
Py = calc_propagator_quadratic_1d(k, ky, zy[:, :, None])
if kz_mode == 'local_xy':
Px *= mathx.expj(delta_gxkz * kx[:, 0] * z[:, :, None])
Py *= mathx.expj(delta_gykz * ky * z[:, :, None])
invTx = Tx.conj()
invTy = Ty.conj()
gradxinvTx = 1j * kx[:, 0] * invTx
gradyinvTy = 1j * ky * invTy
return invTx, gradxinvTx, invTy, gradyinvTy, Px, Py, Tx, Ty
def test_pbg_calc_field():
origin = v4hb.stack_xyzw(0, 0, 0, 1)
vector = v4hb.stack_xyzw(0, 0, 1, 0)
axis1 = v4hb.normalize(v4hb.stack_xyzw(1, 1j, 0.5j, 0))
lamb = 860e-9
k = 2*np.pi/lamb
waist = 100e-6
z_R = np.pi*waist**2/lamb
z_waist = 1*z_R
mode = pbg.Mode.make(rt1.Line(origin, vector), axis1, lamb, waist, z_waist)[0]
z = np.linspace(-z_R*16, z_R*16, 200)[:, None]
x = np.arange(3)[:, None, None]*waist
#zv, xv = [array.ravel() for array in np.broadcast_arrays(zs, xs)]
points = v4hb.concatenate_xyzw(x, 0, z, 1)
field, phi_axis, grad_phi = mode.calc_field(k, points, calc_grad_phi=True)
psis = field*mathx.expj(-z*k)
psis_true = beams.FundamentalGaussian(lamb, w_0=waist, flux=1).E(z - z_waist, x)*np.exp(
-1j*np.arctan(z_waist/z_R))
assert np.allclose(psis, psis_true)
if 0: # Useful for testing in IPython.
glw = pg.GraphicsLayoutWidget()
absz_plot = glw.addAlignedPlot()
glw.nextRows()
phase_plot = glw.addAlignedPlot()
for x, color, psi, psi_true in zip(xs, pg.tableau10, psis.T, psis_true.T):
absz_plot.plot(zs[:, 0]*1e3, abs(psi), pen=color)
absz_plot.plot(zs[:, 0]*1e3, abs(psi_true), pen=pg.mkPen(color, style=pg.DashLine))
phase_plot.plot(zs[:, 0]*1e3, np.angle(psi), pen=color)
phase_plot.plot(zs[:, 0]*1e3, np.angle(psi_true), pen=pg.mkPen(color, style=pg.DashLine))
glw.show()
def __init__(self,
lamb,
n,
zfun,
rs_support,
Er,
gradxyE,
rs_center=(0, 0),
qs_center=(0, 0),
polarizationxy=(1, 0)):
z_center = zfun(*rs_center)
Profile.__init__(self, lamb, n, z_center, rs_support, Er, gradxyE,
rs_center, qs_center, polarizationxy)
self.zfun = zfun
self.z = zfun(self.x, self.y)
self.valid = np.isfinite(self.z)
self.z[~self.valid] = self.z[self.valid].mean()
assert np.allclose(Er[~self.valid], 0)
sumIr = self.Ir.sum()
self.mean_z = (self.z * self.Ir).sum() / sumIr
app_propagator = mathx.expj(
(self.mean_z - self.z) * mathx.divide0(self.Igradphi[2], self.Ir))
app_propagator[np.isnan(self.z)] = 0
Er_plane = Er * app_propagator
gradxyE_plane = tuple(c * app_propagator for c in gradxyE)
# Approximate plane profile.
self.app = PlaneProfile(lamb, n, self.mean_z, rs_support, Er_plane,
gradxyE_plane, rs_center, qs_center)
def propagate_plane_to_plane_spherical(k,
rs_support,
Er,
z,
ms,
rs_center=(0, 0),
qs_center=(0, 0),
ro_centers=None,
kz_mode='local_xy'):
"""Propagate from one plane to another using Sziklas-Siegman.
Args:
k (scalar): Wavenumber.
rs_support (scalar or pair): Real space full support in x and y.
Er (2D array): Input field.
z (scalar): Propagation distance.
ms (scalar or pair): Sziklas-Siegman magnification in x and y.
rs_center (pair): Center of domain and center of wavefront curvature.
qs_center (pair): Center of transverse wavenumber domain.
ro_centers (pair): Center of output domain. Default is what follows from rs_center, qs_center and z.
kz_mode (str): 'paraxial' or 'local_xy'.
Returns:
2D array: Output field.
"""
assert kz_mode in ('local_xy', 'paraxial')
rs_support = sa.to_scalar_pair(rs_support)
ms = sa.to_scalar_pair(ms)
if kz_mode == 'paraxial':
zx = z
zy = z
elif kz_mode == 'local_xy':
fx, fy, delta_kz, delta_gxkz, delta_gykz = math.calc_quadratic_kz_correction(
k, *qs_center)
zx = fx * z
zy = fy * z
else:
raise ValueError('Unknown kz_mode %s.' % kz_mode)
x, y = sa.calc_xy(rs_support, Er.shape, rs_center)
roc_x = zx / (ms[0] - 1)
roc_y = zy / (ms[1] - 1)
kx, ky = sa.calc_kxky(rs_support, Er.shape, qs_center)
kxp = kx - k * rs_center[0] / roc_x
kyp = ky - k * rs_center[1] / roc_y
Er = Er * math.calc_quadratic_phase_1d(
-k, x - rs_center[0], roc_x) * math.calc_quadratic_phase_1d(
-k, y - rs_center[1], roc_y)
Eq = math.fft2(Er)
Eq *= math.calc_propagator_quadratic_1d(
k * ms[0], kxp, zx) * math.calc_propagator_quadratic_1d(
k * ms[1], kyp, zy)
# If local_xy mode, need translation correction factor.
if kz_mode == 'local_xy':
Eq *= mathx.expj(delta_gxkz * kx / ms[0] * z +
delta_gykz * ky / ms[1] * z)
Er = math.ifft2(Eq) * math.calc_spherical_post_factor(
k, rs_support, Er.shape, z, ms, rs_center, qs_center, ro_centers,
kz_mode)
return Er
def propagate_plane_to_plane_spherical_1d(k,
r_support,
Er,
z,
r_center=0,
q_center=0,
roc=np.inf,
axis=-1):
num_points = Er.shape[axis]
Er_flat, z_flat, m_flat, m, r_center_flat = propagate_plane_to_waist_spherical_paraxial_1d(
k, r_support, Er, z, r_center, q_center, roc)
r_support_flat = r_support / m_flat
z_paraxial = z / (1 - (q_center / k)**2)**1.5
q_flat = sa.calc_q(r_support_flat, num_points, q_center, axis)
extra_propagator = mathx.expj((k**2 - q_flat**2)**0.5 * z -
(k - q_flat**2 / (2 * k)) * z_paraxial)
Eq_flat = math.fft(Er_flat, axis)
Eq_flat *= extra_propagator
Er_flat = math.ifft(Eq_flat, axis)
Er, _ = propagate_plane_to_plane_spherical_paraxial_1d(
k, r_support_flat, Er_flat, z_paraxial + z_flat, m * m_flat,
r_center_flat, q_center)
rp_center = r_center + z * q_center / (k**2 - q_center**2)**0.5
return Er, m, rp_center
def propagate_plane_to_plane_flat_1d(k,
r_support,
Er,
z,
q_center=0,
axis=-1,
paraxial=False):
num_points = Er.shape[axis]
Eq = math.fft(Er, axis)
q = sa.calc_q(r_support, num_points, q_center, axis)
if paraxial:
Eq *= mathx.expj((k - q**2 / (2 * k)) * z)
else:
Eq *= mathx.expj((k**2 - q**2)**0.5 * z)
Er = math.ifft(Eq, axis)
return Er
def shift_r_center_1d(r_support,
num_points,
Er,
delta_r_center,
q_center,
k_on_roc=0,
axis=-1):
r = sa.calc_r(r_support, num_points, axis=axis)
q = sa.calc_q(r_support, num_points, q_center, axis=axis)
# Remove quadratic phase and tilt.
Er *= mathx.expj(-(q[0] * r + k_on_roc * r**2 / 2))
Ek = math.fft(Er, axis)
Ek *= mathx.expj(delta_r_center * q)
Er = math.ifft(Ek, axis)
Er *= mathx.expj(q[0] * r + k_on_roc * (r + delta_r_center)**2 / 2)
return Er
def test_curved_interface_collimate():
lamb = 587.6e-9
waist0 = 150e-6
num_points = 2**7
k = 2 * np.pi / lamb
r0_support = (np.pi * num_points)**0.5 * waist0
x0, y0 = asbp.calc_xy(r0_support, num_points)
Er0 = asbp.calc_gaussian(k, x0, y0, waist0)
f = 100e-3
n = 1.5
roc = f * (n - 1)
z_R = np.pi * waist0**2 / lamb
m = (1 + (f / z_R)**2)**0.5
r1_support = m * r0_support
x1, y1 = asbp.calc_xy(r1_support, num_points)
sag = functions.calc_sphere_sag_xy(roc, x1, y1)
Er1, _ = asbp.propagate_plane_to_curved_spherical(k, r0_support, Er0,
f + sag, m)
Er2, propagator = asbp.invert_plane_to_curved_spherical(
k * n, r1_support, Er1, -f * n + sag, 1)
##
waist2 = f * lamb / (np.pi * waist0)
Er2_theory = asbp.calc_gaussian(k, x1, y1, waist2)
Er2_theory *= mathx.expj(np.angle(Er2[0, 0] / Er2_theory[0, 0]))
assert mathx.allclose(Er2, Er2_theory, 1e-3)
def plot_q_(images, r_support, Eq, qs_center, rs_center):
kxu, kyu, Equ = asbp.unroll_q(r_support, Eq, qs_center)
Equ = Equ * mathx.expj((rs_center[0] * kxu + rs_center[1] * kyu))
images[0].setImage(abs(Equ))
images[0].setRect(pg.axes_to_rect(kxu / 1e3, kyu / 1e3))
images[1].setImage(np.angle(Equ) / (2 * np.pi))
images[1].setRect(pg.axes_to_rect(kxu / 1e3, kyu / 1e3))
def interpolate(self, rs_support, num_pointss, rs_center=None):
assert self.zfun is not None
appi = self.app.interpolate(rs_support, num_pointss, rs_center)
z = self.zfun(appi.x, appi.y)
# Propagate from interpolated planar profile to resampled z.
Ir = appi.Ir
propagator = mathx.expj(
(z - appi.z) * mathx.divide0(appi.Igradphi[2], Ir))
Er = appi.Er * propagator
gradxyE = tuple(c * propagator for c in appi.gradxyE)
return type(self)(self.lamb, self.n, self.zfun, appi.rs_support, Er,
gradxyE, appi.rs_center, appi.qs_center)
def propagate_plane_to_plane_flat(k,
rs_support,
Er,
z,
qs_center=(0, 0),
kz_mode='local_xy'):
"""Regularly (non-Sziklas-Siegman) propagate a field between two flat planes.
Args:
k (scalar): Wavenumber.
rs_support (scalar or pair): Real space support.
Er (2D array): Initial field.
z (scalar): Propagation distance.
qs_center (pair): Center of transverse wavenumber space.
kz_mode (str): 'paraxial', 'local_xy', or 'exact'.
Returns:
2D array: Propagated field.
"""
rs_support = sa.to_scalar_pair(rs_support)
qs_center = sa.to_scalar_pair(qs_center)
Eq = math.fft2(Er)
kx, ky = sa.calc_kxky(rs_support, Eq.shape, qs_center)
if kz_mode == 'paraxial':
propagator = math.calc_propagator_paraxial(k, kx, ky, z)
elif kz_mode in ('local_xy', 'local'):
kxp = kx - qs_center[0]
kyp = ky - qs_center[1]
kz, gxkz, gykz, gxxkz, gyykz, gxykz = math.expand_kz(k, *qs_center)
propagator = mathx.expj((kz + gxkz * kxp + gykz * kyp +
gxxkz * kxp**2 / 2 + gyykz * kyp**2 / 2) * z)
if kz_mode == 'local':
propagator *= mathx.expj(gxykz * kxp * kyp * z)
elif kz_mode == 'exact':
propagator = math.calc_propagator_exact(k, kx, ky, z)
else:
raise ValueError('Unknown kz_mode %s.' % kz_mode)
Eq *= propagator
Er = math.ifft2(Eq)
return Er
def test_pbg_simple():
waist = 100e-6
n = 1
lamb = 800e-9
k = 2*np.pi*n/lamb
# Totally arbitrary origin and axes.
origin = np.asarray([1, 2, 3, 1])
vector = v4hb.normalize(np.asarray([1, -2, -1, 0]))
axis_x = v4hb.normalize(np.asarray([1, 1, 0, 0]))
mode, (axis_x, axis_y) = pbg.Mode.make(otk.rt1._lines.Line(origin, vector), axis_x, lamb / n, waist)
z_R = waist**2*k/2
z = np.asarray([-10, 1, 0, 1, 10])[:, None]*z_R*0
x = np.arange(3)[:, None, None]*waist
y = np.arange(3)[:, None, None, None]*waist
#z = np.asarray([[z_R]])
#x = np.asarray([[waist]])
#y = np.asarray([[waist]])
rho = (x**2 + y**2)**0.5
xhat = np.asarray([1, 0, 0, 0])
yhat = np.asarray([0, 1, 0, 0])
rhohat = mathx.divide0(xhat*x + yhat*y, rho)
mode_matrix = np.stack([axis_x, axis_y, vector, origin], 0)
points = v4hb.concatenate_xyzw(x, y, z, 1).dot(mode_matrix)
field, phi_axis, grad_phi = mode.calc_field(k, points, calc_grad_phi=True)
beam = beams.FundamentalGaussian(lamb/n, w_0=waist, flux=1)
field_true = beam.E(z, rho)*mathx.expj(z*k)
field_true_axis = beam.E(z, 0)*mathx.expj(z*k)
grad_phi_true = (beam.dphidz(z, rho) + k)*vector + beam.dphidrho(z, rho)*rhohat.dot(mode_matrix)
testing.assert_allclose(abs(field), abs(field_true))
testing.assert_allclose(field, field_true)
testing.assert_allclose(abs(mathx.wrap_to_pm(phi_axis - np.angle(field_true_axis), np.pi)), 0, atol=1e-6)
testing.assert_allclose(grad_phi, grad_phi_true, atol=1e-6)
def propagate_plane_to_plane_spherical_paraxial_1dE(k,
r_support,
Eri,
z,
m,
r_center=0,
q_center=0):
"""Propagate field with spherical wavefront from plane to plane.
kz is included.
Args:
k:
rs_support:
Eri:
z (2D array): propagation distance vs (x, y)
m:
rs_center:
qs_center:
Returns:
"""
assert Eri.ndim == 1
num_points = len(Eri)
ri = sa.calc_r(r_support, num_points, r_center)
qi = sa.calc_q(r_support, num_points, q_center)
roc = z / (m - 1)
Qir = mathx.expj(-k * ri**2 / (2 * roc[:, None]))
T = math.make_fft_matrix(num_points)
P = mathx.expj((k - qi**2 / (2 * k * m)) * z[:, None])
invT = T.conj()
ro = ri * m
Qor = mathx.expj(k * ro**2 / (2 * (roc + z))) / m**0.5
# ro: i, qi: j, ri: k
Ero = opt_einsum.contract('i, ij, ij, jk, ik, k->i', Qor, invT, P, T, Qir,
Eri)
return Ero
def fourier_transform(self, f, rs0=(0, 0), z=None):
if z is None:
z = self.z + 2 * f
# I *think* that there is an efficient implementation for highly curved wavefronts, analagous to
# https://doi.org/10.1364/OE.24.025974
# Derivation of normalization factor Dane's logbook 3 p81.
norm_fac = self.delta_x * self.delta_y * self.k / (
2 * np.pi * f) * np.prod(self.Er.shape)**0.5
Er = math.fft2(self.Er) * mathx.expj(rs0[0] * self.kx +
rs0[1] * self.ky) * norm_fac
rs_support = self.qs_support * f / self.k
rs_center = self.qs_center * f / self.k
qs_center = (self.rs_center - rs0) * self.k / f
gradxyE = fsq.calc_gradxyE(rs_support, Er, qs_center)
return PlaneProfile(self.lamb, self.n, z, rs_support, Er, gradxyE,
rs_center, qs_center)
def make_ifft_arbitrary_matrix(r_support0, num_points0, q_center0, r1):
"""
Warning: The matrix does not enforce a particular subset of the infinite periodic r1 domain. So if the range of r1
is greater than 2*pi/r_support0, then the output will be repeated.
Args:
r_support0:
num_points0:
q_center0:
r1:
Returns:
"""
assert r1.ndim == 1
q0 = sa.calc_q(r_support0, num_points0, q_center0)
return mathx.expj(q0 * r1[:, None]) / len(q0)**0.5
def calc_spherical_post_factor(k,
rs_support,
num_pointss,
z,
ms,
rs_center=(0, 0),
qs_center=(0, 0),
ro_centers=None,
kz_mode='local_xy'):
assert np.isscalar(z)
ro_supports = rs_support * ms
qs_center = sa.to_scalar_pair(qs_center)
if kz_mode == 'paraxial':
zx = z
zy = z
kz_center = k
delta_kz = k
elif kz_mode == 'local_xy':
fx, fy, delta_kz, delta_gxkz, delta_gykz = calc_quadratic_kz_correction(
k, *qs_center)
zx = fx * z
zy = fy * z
kz_center = (k**2 - (qs_center**2).sum())**0.5
else:
raise ValueError('Unknown kz_ mode %s.', kz_mode)
if ro_centers is None:
ro_centers = rs_center + qs_center / kz_center * z
xo, yo = sa.calc_xy(ro_supports, num_pointss, ro_centers)
if kz_mode == 'local_xy':
xo += delta_gxkz * z
yo += delta_gykz * z
roc_x = zx / (ms[0] - 1)
roc_y = zy / (ms[1] - 1)
roc_xo = roc_x + zx
roc_yo = roc_y + zy
# See derivation page 114 Dane's logbook 2.
Qo = calc_quadratic_phase_1d(k, xo, roc_xo) * calc_quadratic_phase_1d(
k, yo, roc_yo) * mathx.expj(delta_kz * z + k *
(rs_center[0]**2 /
(2 * roc_x) + rs_center[1]**2 /
(2 * roc_y) - rs_center[0] * xo /
(roc_x * ms[0]) - rs_center[1] * yo /
(roc_y * ms[1]))) / (ms[0] * ms[1])**0.5
return Qo
def plot_r(images, beam, scale=1):
profile = beam.profile
xu, yu, Eru = asbp.unroll_r(profile.rs_support, profile.Er * scale,
profile.rs_center)
tilt_mode = self.tilt_mode_widget.currentIndex()
if not isinstance(beam.profile, asbp.PlaneProfile):
q_profile = profile.app
else:
q_profile = profile
if tilt_mode == 0:
q0s = q_profile.centroid_qs_flat
else:
q0s = q_profile.peak_qs
q0s += np.asarray(
(self.tilt_nudge_x_widget.value(),
self.tilt_nudge_y_widget.value())) / 1e3 * profile.k
Eru = Eru * mathx.expj(-(q0s[0] *
(xu - profile.peak_rs[0]) + q0s[1] *
(yu - profile.peak_rs[1]) +
np.angle(profile.peak_Er)))
images[0].setImage(abs(Eru))
images[0].setRect(pg.axes_to_rect(xu * 1e3, yu * 1e3))
images[1].setImage(np.angle(Eru) / (2 * np.pi))
images[1].setRect(pg.axes_to_rect(xu * 1e3, yu * 1e3))
def calc_propagator_paraxial(k, kx, ky, l):
return mathx.expj(calc_kz_paraxial(k, kx, ky) * l)
def calc_propagator_quadratic(k, kx, ky, l):
return mathx.expj(calc_kz_quadratic(k, kx, ky) * l)
def calc_propagator_quadratic_1d(k, q, l):
"""kz not included."""
return mathx.expj(calc_kz_quadratic_1d(k, q) * l)
def calc_quadratic_phase_mask(x, y, k_on_f):
m = mathx.expj(-0.5 * (x**2 + y**2) * k_on_f)
gradxm = -m * x * k_on_f
gradym = -m * y * k_on_f
return m, (gradxm, gradym)
def calc_quadratic_phase_factor(self, x, y):
rs_center = self.rs_center
phi_cx, phi_cy = self.phi_cs
var_x, var_y = self.var_rs
return mathx.expj((x - rs_center[0])**2 * phi_cx / var_x) * mathx.expj(
(y - rs_center[1])**2 * phi_cy / var_y)
def calc_plane_to_curved_spherical_arbitrary_factors(k,
rs_support,
num_pointss,
z,
xo,
yo,
roc_x,
roc_y,
ri_centers=(0, 0),
qs_center=(0, 0),
ro_centers=None,
kz_mode='local_xy'):
"""Calculate factors for 2D Sziklas-Siegman propagation (paraxial) from flat to curved surface(s).
The zero-oth order phase (k_z) is (arbitrarily) included in QinvTPx.
Compared to regular Sziklas-Siegman procedure, the factors are complicated by a transformation that allows nonzero
real and angular space centers. See derivation page 114 Dane's logbook 2.
Possible optimizations: could probably shave 50% off by more aggressive use of numba to avoid calculation of intermediates.
But typically the time is spent using (tensor contract) rather than calculating the factors, so it won't bring
a great speedup.
Args:
k (scalar): Wavenumber.
rs_support (scalar or pair): Real space support.
num_pointss (scalar int or pair): Number of sampling points in input.
z (M*N array): Propagation for distance.
xo (M*1 array): Output x sample values.
yo (N array): Output y sample values.
roc_x (M*N array): Radius of curvature in x, at the input, but sampled at the output points.
roc_y (M*N array): Radius of curvature in y, at the input, but sampled at the output points.
ri_centers (pair of scalars): Center of initial real-space aperture.
qs_center (pair of scalars): Center of angular space aperture.
kz_mode (str): 'paraxial' or 'local_xy'.
Returns:
QinvTPx (3D array): Product propagator, inverse DFT and quadratic output factor along x. Axes are (xo, yo, kx).
gradxQinvTPx (3D array): Derivative of above w.r.t output x. Note that z is constant - the derivative is along
a transverse plane rather than the surface.
QinvTPy (3D array): Product propagator, inverse DFT and quadratic output factor along x. Axes are (xo, yo, kx).
gradyQinvTPy (3D array): Derivative of above w.r.t output x. As with gradxQinvTPx, z is constant.
Tx (2D array): DFT factor. Axes are (kx, xi).
Ty (2D array): DFT factor. Axes are (ky, yi).
Qix (3D array): Input x quadratic phase factor vs x. Axes are (xo, yo, xi).
Qiy (3D array): Input y quadratic phase factor. Axes are (xo, yo, yi).
"""
assert np.isscalar(k)
rs_support = sa.to_scalar_pair(rs_support)
assert kz_mode in ('local_xy', 'paraxial')
assert np.ndim(z) == 2
assert np.shape(z) == np.shape(roc_x)
assert np.shape(z) == np.shape(roc_y)
assert np.shape(xo) == (np.shape(z)[0], 1)
assert np.shape(yo) == (np.shape(z)[1], )
num_pointss = sa.to_scalar_pair(num_pointss)
qs_center = sa.to_scalar_pair(qs_center)
if ro_centers is None:
z_center = np.mean(z)
ro_centers = adjust_r(k, ri_centers, z_center, qs_center, kz_mode)
# If local_xy mode is requested then propagation distances must be scaled. Calculate scaled propagation distances
# and required kz component.
if kz_mode == 'paraxial':
zx = z
zy = z
delta_kz = k
else:
fx, fy, delta_kz, delta_gxkz, delta_gykz = calc_quadratic_kz_correction(
k, *qs_center)
zx = fx * z
zy = fy * z
mx = zx / roc_x + 1
my = zy / roc_y + 1
xi, yi = sa.calc_xy(rs_support, num_pointss, ri_centers)
kxi, kyi = sa.calc_kxky(rs_support, num_pointss,
qs_center) # +q_curvatures)
# Quadratic phase is centered at origin. Numerically, this requires evaluation of complex exponentials of size N^3.
# Significant fraction of total cost of this function, but unavoidable.
Qix = calc_quadratic_phase_1d(-k, xi[:, 0] - ri_centers[0], roc_x[:, :,
None])
Qiy = calc_quadratic_phase_1d(-k, yi - ri_centers[1], roc_y[:, :, None])
Tx = make_fft_matrix(num_pointss[0])
if num_pointss[1] == num_pointss[0]:
Ty = Tx
else:
Ty = make_fft_matrix(num_pointss[1])
roc_xo = roc_x + zx
roc_yo = roc_y + zy
# If in local_xy mode, then the result of the Sziklas-Siegman step needs to be translated. This means that the
# factors which depend on xo and yo need to be translated too. (Note that the x and y factors in the inverse DFT
# are not translated.) We define xop and yop as the translated output coordinates.
if kz_mode == 'local_xy':
xop = xo + delta_gxkz * z
yop = yo + delta_gykz * z
else:
xop = xo
yop = yo[None, :]
# Effective kx and ky for propagation takes into account center of curvature.
kxip = kxi[:, 0] - k * ri_centers[0] / roc_x[:, :, None]
kyip = kyi - k * ri_centers[1] / roc_y[:, :, None]
r2_support_x = mx * rs_support[0]
r2_support_y = my * rs_support[1]
valid_x = abs(xo - ro_centers[0]) < r2_support_x * 0.5
valid_y = abs(yo - ro_centers[1]) < r2_support_y * 0.5
# Calculate product of propagator and inverse transform along x axis. The x axis (arbitrarily) includes delta_kz.
# Could combine all exponents before exponentiation? Won't help much because the time will be dominated by the propagator
# which is N^3 - the others are N^2.
phi = delta_kz * z + k * (ri_centers[0]**2 /
(2 * roc_x) - ri_centers[0] * xop / (roc_x * mx))
QinvTPx = (
calc_propagator_quadratic_1d(k * mx[:, :, None], kxip, zx[:, :, None])
* # Propagation phase scaled by magnification.
calc_quadratic_phase_1d(k, xop, roc_xo)[:, :, None]
* # Normal Sziklas-Siegman final quadratic phase.
mathx.expj(phi[:, :, None] + kxi[:, 0] * (xo / mx)[:, :, None]) /
num_pointss[0]**0.5 / # Correction phases and inverse DFT.
abs(mx[:, :, None])**0.5 * # Magnification correction to amplitude.
valid_x[:, :, None])
# Calculate product of propagator and inverse transform along x axis.
phi = k * (ri_centers[1]**2 / (2 * roc_y) - ri_centers[1] * yop /
(roc_y * my))
QinvTPy = (
calc_propagator_quadratic_1d(k * my[:, :, None], kyip, zy[:, :, None])
* # Propagation phase scaled by magnification.
calc_quadratic_phase_1d(k, yop, roc_yo)[:, :, None]
* # Normal Sziklas-Siegman final quadratic phase.
mathx.expj(phi[:, :, None] + kyi * (yo / my)[:, :, None]) /
num_pointss[1]**0.5 / # Correction phases and inverse DFT.
abs(my[:, :, None])**0.5 * # Magnification correction to amplitude.
valid_y[:, :, None])
# If local_xy mode, need translation correction factor. Could combine this with above to reduce number of N^3 expj
# calls but this would only give a marginal performance improvement.
if kz_mode == 'local_xy':
QinvTPx *= mathx.expj(delta_gxkz * kxi[:, 0] * (z / mx)[:, :, None])
QinvTPy *= mathx.expj(delta_gykz * kyi * (z / my)[:, :, None])
# Evaluate derivatives of the invTP factors with respect to output x and y (but constant z - not along the curved
# surface).
gradxQinvTPx = 1j * (k * (xop / roc_xo - ri_centers[0] /
(roc_x * mx))[:, :, None] +
kxi[:, 0] / mx[:, :, None]) * QinvTPx
gradyQinvTPy = 1j * (k * (yop / roc_yo - ri_centers[1] /
(roc_y * my))[:, :, None] +
kyi / my[:, :, None]) * QinvTPy
factors = QinvTPx, gradxQinvTPx, QinvTPy, gradyQinvTPy, Tx, Ty, Qix, Qiy
return factors
def calc_plane_to_curved_spherical_factors(k,
rs_support,
num_pointss,
z,
ms,
ri_centers=(0, 0),
qs_center=(0, 0),
ro_centers=(0, 0),
kz_mode='local_xy'):
"""Calculate factors for 2D Sziklas-Siegman propagation (paraxial) from flat to curved surface(s).
The magnifications must be the same for all output points, but can be different for the x and y planes.
The zero-oth order phase (k_z) is (arbitrarily) included in QinvTPx.
Compared to regular Sziklas-Siegman procedure, the factors are complicated by a transformation that allows nonzero
real and angular space centers. See derivation page 114 Dane's logbook 2.
Possible optimizations: could probably shave 50% off by more aggressive use of numba to avoid calculation of intermediates.
But typically the time is spent using (tensor contract) rather than calculating the factors, so it won't bring
a great speedup.
Args:
k (scalar): Wavenumber.
rs_support (scalar or pair): Real space support.
num_pointss (scalar int or pair): Number of sampling points.
z (array of size num_pointss): Propagation for distance.
ms (scalar or pair): Magnification(s).
ri_centers (pair of scalars): Center of initial real-space aperture.
qs_center (pair of scalars): Center of angular space aperture.
ro_centers (pair of scalars): Center of final real-space aperture. If not given, the shift from ri_centers is
inferred from the propagation distance and qs_center.
kz_mode (str): 'paraxial' or 'local_xy'.
Returns:
QinvTPx (3D array): Product propagator, inverse DFT and quadratic output factor along x. Axes are (xo, yo, kx).
gradxQinvTPx (3D array): Derivative of above w.r.t output x. Note that z is constant - the derivative is along
a transverse plane rather than the surface.
QinvTPy (3D array): Product propagator, inverse DFT and quadratic output factor along x. Axes are (xo, yo, kx).
gradyQinvTPy (3D array): Derivative of above w.r.t output x. As with gradxQinvTPx, z is constant.
Tx (2D array): DFT factor. Axes are (kx, xi).
Ty (2D array): DFT factor. Axes are (ky, yi).
Qix (3D array): Input x quadratic phase factor vs x. Axes are (xo, yo, xi).
Qiy (3D array): Input y quadratic phase factor. Axes are (xo, yo, yi).
"""
assert kz_mode in ('local_xy', 'paraxial')
num_pointss = sa.to_scalar_pair(num_pointss)
ms = sa.to_scalar_pair(ms)
if np.isclose(ms, 1).any():
logger.debug('At least one magnification (%g, %g) close to 1.', ms)
qs_center = sa.to_scalar_pair(qs_center)
# TODO (maybe) factor out adjust_r.
if kz_mode == 'paraxial':
kz_center = k
else:
kz_center = (k**2 - (qs_center**2).sum())**0.5
z_center = np.mean(z)
if ro_centers is None:
ro_centers = ri_centers + qs_center / kz_center * z_center
# If local_xy mode is requested then propagation distances must be scaled. Calculate scaled propagation distances
# and required kz component.
if kz_mode == 'paraxial':
zx = z
zy = z
delta_kz = k
else:
fx, fy, delta_kz, delta_gxkz, delta_gykz = calc_quadratic_kz_correction(
k, *qs_center)
zx = fx * z
zy = fy * z
ro_supports = rs_support * ms
xi, yi = sa.calc_xy(rs_support, num_pointss, ri_centers)
kxi, kyi = sa.calc_kxky(rs_support, num_pointss,
qs_center) # +q_curvatures)
roc_x = mathx.divide0(zx, ms[0] - 1, np.inf)
roc_y = mathx.divide0(zy, ms[1] - 1, np.inf)
# Quadratic phase is centered at origin. Numerically, this requires evaluation of complex exponentials of size N^3.
# Significant fraction of total cost of this function, but unavoidable.
Qix = calc_quadratic_phase_1d(-k, xi[:, 0] - ri_centers[0], roc_x[:, :,
None])
Qiy = calc_quadratic_phase_1d(-k, yi - ri_centers[1], roc_y[:, :, None])
Tx = make_fft_matrix(num_pointss[0])
if num_pointss[1] == num_pointss[0]:
Ty = Tx
else:
Ty = make_fft_matrix(num_pointss[1])
if 0:
# This was an experiment - we calculate where a given pencil beam lands based on initial position and
# momentum. It didn't improve the stability of the algorithm. In a ray picture, clip factor should be 0.5. But
# I found that anything below 3 gave unacceptable artefacts.
clip_factor = 3
mean_rocs = z_centers / (ms - 1)
xo_kx_xi = xi[:, 0] * ms[0] + (
kxi - k * ri_centers[0] / mean_rocs[0]) * z_centers[0] / k
yo_ky_yi = yi * ms[1] + (
kyi[:, None] - k * ri_centers[1] / mean_rocs[1]) * z_centers[1] / k
in_xo = abs(xo_kx_xi - ro_centers[0]) < ro_supports[0] * clip_factor
in_yo = abs(yo_ky_yi - ro_centers[1]) < ro_supports[1] * clip_factor
Txp = Tx * in_xo
Typ = Ty * in_yo
xo, yo = sa.calc_xy(ro_supports, num_pointss, ro_centers)
roc_xo = roc_x + zx
roc_yo = roc_y + zy
# If in local_xy mode, then the factors which depend on xo and yo use transformed quantities, which we subsitute here.
if kz_mode == 'local_xy':
xo = xo + delta_gxkz * z
yo = yo + delta_gykz * z
# Effective kx and ky for propagation takes into account center of curvature.
kxip = kxi[:, 0] - k * ri_centers[0] / roc_x[:, :, None]
kyip = kyi - k * ri_centers[1] / roc_y[:, :, None]
# Calculate product of propagator and inverse transform along x axis. The x axis (arbitrarily) includes delta_kz.
# Could combine all exponents before exponentiation? Won't help much because the time will be dominated by the propagator
# which is N^3 - the others are N^2.
phi = delta_kz * z + k * (ri_centers[0]**2 /
(2 * roc_x) - ri_centers[0] * xo /
(roc_x * ms[0]))
QinvTPx = (
calc_propagator_quadratic_1d(k * ms[0], kxip, zx[:, :, None])
* # Propagation phase scaled by magnification.
calc_quadratic_phase_1d(k, xo, roc_xo)[:, :, None]
* # Normal Sziklas-Siegman final quadratic phase.
mathx.expj(phi[:, :, None]) * Tx.conj()[:, None, :] / # Inverse DFT.
abs(ms[0])**0.5) # Magnification correction to amplitude.
# Calculate product of propagator and inverse transform along x axis. Result depends on
phi = k * (ri_centers[1]**2 / (2 * roc_y) - ri_centers[1] * yo /
(roc_y * ms[1]))
QinvTPy = (
calc_propagator_quadratic_1d(k * ms[1], kyip, zy[:, :, None])
* # Propagation phase scaled by magnification.
calc_quadratic_phase_1d(k, yo, roc_yo)[:, :, None]
* # Normal Sziklas-Siegman final quadratic phase.
mathx.expj(phi[:, :, None]) * Ty.conj() / # Inverse DFT.
abs(ms[1])**0.5) # Magnification correction to amplitude.
# If local_xy mode, need translation correction factor. Could combine this with above to reduce number of N^3 expj
# calls.
if kz_mode == 'local_xy':
QinvTPx *= mathx.expj(delta_gxkz * kxi[:, 0] / ms[0] * z[:, :, None])
QinvTPy *= mathx.expj(delta_gykz * kyi / ms[1] * z[:, :, None])
# Evaluate derivatives of the invTP factors with respect to output x and y (but constant z - not along the curved
# surface).
gradxQinvTPx = 1j * (k * (xo / roc_xo)[:, :, None] - ri_centers[0] * k /
(roc_x[:, :, None] * ms[0]) +
kxi[:, 0] / ms[0]) * QinvTPx
gradyQinvTPy = 1j * (k * (yo / roc_yo)[:, :, None] - ri_centers[1] * k /
(roc_y[:, :, None] * ms[1]) + kyi / ms[1]) * QinvTPy
# Won't use z gradients for now - will keep for future.
# gradzPx=1j*calc_kz_paraxial_1d(k*ms[0], kxip)*Px
factors = QinvTPx, gradxQinvTPx, QinvTPy, gradyQinvTPy, Tx, Ty, Qix, Qiy
return factors
def calc_quadratic_phase_1d(k, r, roc):
return mathx.expj(k * r**2 / (2 * roc))
