def adaptive(t,
y,
h,
prop_1,
trans,
prop_2,
inv_trans,
err_tol=1e-3,
err_fun=None):
if err_fun is None:
err_fun = lambda yd, ys: (mathx.abs_sqd(yd - ys).sum() / mathx.abs_sqd(
yd).sum())**0.5
yd = double(t, y, h, prop_1, trans, prop_2, inv_trans)
ys = single(t, y, h, prop_1, trans, prop_2, inv_trans)
# Per channel - maybe TODO make option
#rel_err=(abs(yd-ys)/np.maximum(np.maximum(abs(yd),abs(ys)),sys.float_info.min)).max()
err = err_fun(yd, ys)
eps = err_tol / err
pow = 1. / 3
ok = eps >= 1
if eps >= 1:
h_next = h * min(0.9 * eps**pow, 10)
else:
h_next = h * max(0.9 * eps**pow, 0.1)
return yd, ok, h_next, None, None, h, False
def spectrogram(t, sign, Et, Gt, omega, bound_cond='pad'):
assert bound_cond in ('pad', 'odd-cyclic', 'cyclic')
if not type(Gt) == np.array:
nG = Gt
ga = np.arange(nG) - (nG - 1.0) / 2
Gt = np.exp(-(5 * ga / nG)**2)
# plt.figure()
# plt.plot(ga,Gt)
else:
nG = Gt.size
Dt = t[1] - t[0]
t_G = np.arange(nG) * Dt
if bound_cond == 'pad':
t, Et = usv.pad_sampled(t, Et, nG, nG)
omega_s = conj_axis(t_G, omega[0], 'start')
ft = FTD(sign, x=t_G, k=omega_s)
if bound_cond in ('cyclic', 'odd-cyclic'):
tsis = np.arange(t.size) - int(nG / 2)
else:
tsis = np.arange(t.size - Gt.size)
t_s = t[0] + (tsis + nG / 2) * Dt
S = np.zeros((omega_s.size, t_s.size))
for tsii in range(len(tsis)):
tsi = tsis[tsii]
inds = np.arange(tsi, tsi + nG)
if bound_cond == 'odd-cyclic':
sign = (-1)**np.floor(inds / t.size)
else:
sign = 1
if bound_cond in ('cyclic', 'odd-cyclic'):
inds = inds % t.size
S[:, tsii] = mx.abs_sqd(ft.trans(Et[inds] * Gt * sign))
return (t_s, omega_s, S)
def refract_field_gradient(normal, k1, Er, gradxyEr1, k2, Ir=None):
"""Apply local Snell's law at an interface given derivatives of the field.
Args:
normal (3-tuple): arrays of surface normal components. Should have positive z component.
k1: wavenumber in initial medium.
Er: field.
gradxyEr1 (2-tuple): derivatives of Er w.r.t x and y.
k2: wavenumber in final medium.
Returns:
Igradphi2 (3-tuple): product of intensity |Er|^2 and gradient of phase w.r.t. coordinate axes.
Ir: |Er|^2 (for reuse).
"""
assert np.all(normal[2] > 0)
assert len(gradxyEr1) == 2
if Ir is None:
Ir = mathx.abs_sqd(Er)
# Calculate intensity-scaled phase gradient (3 vector).
Igradphi = calc_Igradphi(k1, Er, gradxyEr1, Ir)
Igradphi_tangent = mathx.project_onto_plane(Igradphi, normal)[0]
Igradphi_normal = np.maximum((Ir * k2)**2 - mathx.dot(Igradphi_tangent),
0)**0.5
Igradphi2 = [
tc + nc * Igradphi_normal for tc, nc in zip(Igradphi_tangent, normal)
]
return Igradphi2
def calc_Igradphi(k, Er, gradxyEr, Ir=None) -> list:
if Ir is None:
Ir = mathx.abs_sqd(Er)
# Calculate phase gradient times intensity.
Igradxyphi = [(component * Er.conj()).imag for component in gradxyEr]
# The phase gradient is the eikonal (with scalar k included). In the short wavelength approximation (geometrical optics)
# the length of the eikonal is k. (See Born & Wolf eq. 15b in sec. 3.1.)
Igradzphi = np.maximum((Ir * k)**2 - matseq.dot(Igradxyphi), 0)**0.5
Igradphi = Igradxyphi + [Igradzphi]
return Igradphi
def __init__(self,
lamb,
n,
z,
rs_support,
Er,
gradxyE,
rs_center=(0, 0),
qs_center=(0, 0),
polarizationxy=(1, 0)):
assert np.isscalar(z)
Profile.__init__(self, lamb, n, z, rs_support, Er, gradxyE, rs_center,
qs_center, polarizationxy)
self.z = z
# Flat beam is beam with quadratic phase removed.
self.Er_flat = self.Er * self.calc_quadratic_phase_factor(
self.x, self.y).conj()
self.Eq_flat = math.fft2(self.Er_flat)
Iq_flat = mathx.abs_sqd(self.Eq_flat)
sumIq_flat = Iq_flat.sum()
self.qs_support = 2 * np.pi * np.asarray(Er.shape) / self.rs_support
kx, ky = sa.calc_kxky(rs_support, Er.shape, qs_center)
self.kx = kx
self.ky = ky
self.q_center_indices = abs(kx - qs_center[0]).argmin(), abs(
ky - qs_center[1]).argmin()
mean_kx_flat = mathx.moment(kx, Iq_flat, 1, sumIq_flat)
mean_ky_flat = mathx.moment(ky, Iq_flat, 1, sumIq_flat)
self.centroid_qs_flat = mean_kx_flat, mean_ky_flat
var_kx_flat = mathx.moment(kx - qs_center[0], Iq_flat, 2, sumIq_flat)
var_ky_flat = mathx.moment(ky - qs_center[1], Iq_flat, 2, sumIq_flat)
# Calculate angular variance (of plane beam) from flat beam.
var_kx = bvar.infer_angular_variance_spherical(self.var_rs[0],
self.phi_cs[0],
var_kx_flat)
var_ky = bvar.infer_angular_variance_spherical(self.var_rs[1],
self.phi_cs[1],
var_ky_flat)
self.var_qs = np.asarray((var_kx, var_ky))
dz_waists, self.var_r_waists, self.Msqds, self.z_R = np.asarray([
bvar.calc_waist(self.k, var_r, phi_c,
var_q) for var_r, phi_c, var_q in zip(
self.var_rs, self.phi_cs, self.var_qs)
]).T
self.z_waists = dz_waists + self.z
self.z_waist = np.mean(dz_waists) + self.z
def calc_refracted_propagation_spherical(r_support,
z12,
normal12,
k1,
Er1,
gradEr1,
k2,
z,
r1_centers=(0, 0),
qs_center=(0, ),
f_nexts=(np.inf, np.inf),
qfds=(1, 1)):
"""Refract"""
Ir1 = mathx.abs_sqd(Er1)
Igradphi1 = refract_field_gradient(normal12, k1, Er1, gradEr1, k2, Ir1)
return calc_propagation_spherical(k2, r_support, z12, Er1, z, Igradphi1,
r1_centers, qs_center, f_nexts, qfds)
def __init__(self, omega, Ef, omega_0, Et=None, It=None, phit=None, fwhm=None, *args, **kwargs):
super().__init__(omega=omega, *args, **kwargs)
if Et is None:
Et = ft.trans(self.omega, -1, Ef, self.t, AXIS_TEMP) * np.exp(1j * omega_0 * self.t)
if It is None:
It = mathx.abs_sqd(Et)
if phit is None:
phit = mathx.unwrap(np.angle(Et), abs(self.t).argmin(), axis=AXIS_TEMP)
self.Ef = Ef
self.omega_0 = omega_0
self.Et = Et
self.It = It
self.phit = phit
if fwhm is None:
t_half_max = mathx.peak_crossings(self.t, self.It, 0.5, AXIS_TEMP)
fwhm = t_half_max[1] - t_half_max[0]
self.fwhm = fwhm
def to_mode(key):
mode = modes[key]
polarization = self.frame[0, :]
new_polarization = np.einsum('...i, ...ij', polarization,
mode.matrix)
pol_Er_factor = (v4hb.dot(new_polarization)**0.5).reshape(
self.profile.Er.shape)
# Take mean polarization.
new_polarization = v4hb.normalize((new_polarization.reshape(
(-1, 4)) * mathx.abs_sqd(self.profile.Er).reshape(
(-1, 1))).sum(axis=0))
Er_factor = pol_Er_factor * (mode.n / self.profile.n)**0.5
if mode.direction == rt1.Directions.TRANSMITTED:
return self.refract(normal, mode.n, Er_factor,
new_polarization)
else:
return self.reflect(normal, mode.n, Er_factor,
new_polarization)
def propagate_curved_to_plane_spherical(k,
rs_support,
Eri,
z,
ms,
ri_centers=(0, 0),
qs_center=(0, 0),
ro_centers=None,
kz_mode='local_xy',
max_iterations=None,
tol=None):
"""
Args:
k:
rs_support:
Eri:
zx:
ms:
ri_centers:
qs_center:
zy:
ro_centers:
max_iterations:
tol:
Returns:
"""
assert kz_mode in ('local_xy', 'paraxial')
if ro_centers is None:
Iri = mathx.abs_sqd(Eri)
sumIri = Iri.sum()
z_center = mathx.moment(z, Iri, 1, sumIri)
ro_centers = math.adjust_r(k, ri_centers, z, qs_center,
kz_mode) # TODO should be z_center?
Ero, propagator = invert_plane_to_curved_spherical(k, rs_support * ms, Eri,
-z, 1 / ms, ro_centers,
qs_center, ri_centers,
kz_mode, max_iterations,
tol)
return Ero
def mask(self, f: np.ndarray, gradxyf: tuple, n: float = None):
"""Return self with real-space mask applied.
Args:
f: Mask amplitude sampled at same points as self.Er.
gradxyf: Mask gradients along x and y.
n: New refractive index (defaults to self.n).
"""
Er = self.Er * f
gradxyE = [
gradE * f + self.Er * gradf
for gradE, gradf in zip(self.gradxyE, gradxyf)
]
Ir = mathx.abs_sqd(Er)
Igradphi = fsq.calc_Igradphi(self.k, Er, gradxyE, Ir)
# Mean transverse wavenumber is intensity-weighted average of transverse gradient of phase.
qs_center = np.asarray([component.sum()
for component in Igradphi[:2]]) / Ir.sum()
return self.change(Er=Er, gradxyE=gradxyE, n=n, qs_center=qs_center)
def calc_propagation_spherical(k,
r_support,
zi,
Er,
z,
Igradphi=None,
rs_center=(0, 0),
qs_center=(0, 0),
f_nexts=(np.inf, np.inf),
qfds=(1, 1)):
assert np.isscalar(z)
Ir = mathx.abs_sqd(Er)
if Igradphi is None:
Ek = math.fft2(Er)
kx, ky = sa.calc_kxky(r_support, Ek.shape, qs_center)
gradxEx = math.ifft2(Ek * 1j * kx)
gradyEy = math.ifft2(Ek * 1j * ky)
Igradphi = calc_Igradphi(k, Er, (gradxEx, gradyEy), Ir)
zi_mean, rs_center, var_rs, qs_center, var_qs, phi_cs = calc_beam_properties(
k, r_support, zi, Er, Igradphi, rs_center)
if isinstance(z, str) and z == 'waist':
return_z = True
z, var_r0, Msqd, z_R = bvar.calc_waist(k, var_rs, phi_cs, var_qs)
z = np.mean(z)
else:
return_z = False
m = bvar.calc_propagation_ms(k, r_support, var_rs, phi_cs, var_qs,
z - zi_mean, Er.shape, f_nexts, qfds)
if return_z:
return m, z, rs_center, qs_center, zi_mean, rs_center + qs_center / k * (
z - zi_mean)
else:
return m, rs_center, qs_center, zi_mean, rs_center + qs_center / k * (
z - zi_mean)
def __init__(self,
lamb: float,
n: float,
z_center: float,
rs_support,
Er,
gradxyE,
rs_center=(0, 0),
qs_center=(0, 0),
polarizationxy=(1, 0)):
self.lamb = float(lamb)
self.n = float(n)
self.z_center = float(z_center)
rs_support = sa.to_scalar_pair(rs_support)
assert (rs_support > 0).all()
self.rs_support = rs_support
Er = np.asarray(Er).astype(complex)
self.Er = Er
assert not np.isnan(Er).any()
assert len(gradxyE) == 2
self.gradxyE = gradxyE
self.rs_center = sa.to_scalar_pair(rs_center)
self.qs_center = sa.to_scalar_pair(qs_center)
self.Eq = math.fft2(Er)
self.k = 2 * np.pi * self.n / self.lamb
self.Ir = mathx.abs_sqd(Er)
sumIr = self.Ir.sum()
if np.isclose(sumIr, 0):
raise NullProfileError()
self.Igradphi = fsq.calc_Igradphi(self.k, Er, gradxyE, self.Ir)
x, y = sa.calc_xy(rs_support, Er.shape, rs_center)
self.x = x
self.y = y
self.r_center_indices = abs(x - rs_center[0]).argmin(), abs(
y - rs_center[1]).argmin()
self.delta_x, self.delta_y = self.rs_support / Er.shape
self.power = sumIr * self.delta_x * self.delta_y
mean_x = mathx.moment(x, self.Ir, 1, sumIr)
mean_y = mathx.moment(y, self.Ir, 1, sumIr)
self.centroid_rs = mean_x, mean_y
var_x = mathx.moment(x - rs_center[0], self.Ir, 2, sumIr)
var_y = mathx.moment(y - rs_center[1], self.Ir, 2, sumIr)
self.var_rs = np.asarray((var_x, var_y))
# Calculate phase of quadratic component at RMS distance from center. Proportional to A in Siegman IEE J. Quantum Electronics Vol. 27
# 1991. Positive means diverging.
phi_cx = 0.5 * (
(x - rs_center[0]) *
(self.Igradphi[0] - self.Ir * qs_center[0])).sum() / sumIr
phi_cy = 0.5 * (
(y - rs_center[1]) *
(self.Igradphi[1] - self.Ir * qs_center[1])).sum() / sumIr
self.phi_cs = np.asarray((phi_cx, phi_cy))
self.rocs = mathx.divide0(self.k * self.var_rs, 2 * self.phi_cs,
np.inf)
xi, yi = np.unravel_index(self.Ir.argmax(), self.Ir.shape)
self.peak_indices = xi, yi
self.peak_Er = self.Er[xi, yi]
self.peak_rs = np.asarray((self.x[xi], self.y[yi]))
self.peak_qs = np.asarray([
mathx.divide0((gradxE[xi, yi] * Er[xi, yi].conj()).imag,
self.Ir[xi, yi]) for gradxE in gradxyE
])
self.kz_center = math.calc_kz(self.k, *self.qs_center)
# Calculate 3D vectors.
vector_center = v4hb.normalize(
v4hb.stack_xyzw(*self.qs_center, self.kz_center, 0))
polarizationz = -(polarizationxy *
vector_center[:2]).sum() / vector_center[2]
origin_center = v4hb.stack_xyzw(*self.rs_center, self.z_center, 1)
polarization = polarizationxy[0], polarizationxy[1], polarizationz, 0
y = v4hb.cross(vector_center, polarization)
self.frame = np.c_[polarization, y, vector_center, origin_center].T
def norm(self, E, conj=False):
sc = self.Dk if conj else self.Dx
return np.sum(mx.abs_sqd(E), axis=self.axis, keepdims=True) * sc
def propagate_arbitrary_curved_to_plane_spherical(k,
xi,
yi,
Eri,
roc_xi,
roc_yi,
zi,
ro_supports,
num_pointsso,
ri_centers=(0, 0),
qs_center=(0, 0),
ro_centers=None,
kz_mode='local_xy',
invert_kwargs=None):
"""
Args:
k:
xi (Mx1 array):
yi (N array):
Eri (MxN array):
roc_x (M*N array): Input radius of curvature along x.
roc_y (M*N array): Input radius of curvature along y.
zi (MxN array):
ro_supports:
num_pointsso:
ri_centers:
qs_center:
ro_centers:
kz_mode:
invert_kwargs:
Returns:
Ero (array of size num_pointsso):
"""
assert xi.shape == (Eri.shape[0], 1)
assert yi.shape == (Eri.shape[1], )
assert Eri.ndim == 2
assert roc_xi.shape == Eri.shape
assert roc_yi.shape == Eri.shape
assert zi.shape == Eri.shape
if ro_centers is None:
z_center = mathx.moment(zi, mathx.abs_sqd(Eri), 1)
ro_centers = math.adjust_r(k, ri_centers, z_center, qs_center, kz_mode)
z = -zi
roc_x = roc_xi + zi
roc_y = roc_yi + zi
Ero, propagator = invert_plane_to_curved_spherical_arbitrary(
k, ro_supports, num_pointsso, Eri, z, xi, yi, roc_x, roc_y, ro_centers,
qs_center, ri_centers, kz_mode, invert_kwargs)
return Ero
# def propagate_plane_to_waist_spherical_paraxial(k, rs_support, Er, z, rs_center = (0, 0), qs_center = (0, 0), rocs = np.inf):
# rs_support = to_pair(rs_support)
# rocs = to_pair(rocs)
# rs_center = np.asarray(rs_center)
# qs_center = np.asarray(qs_center)
# num_pointss = Er.shape[-2:]
# ms_flat, zs_flat, ms = np.asarray([calc_curved_propagation(k, r_support, num_points, roc, z) for r_support, num_points, roc in zip(rs_support, num_pointss, rocs)]).T
# Er_flat, r_centers_flat = propagate_plane_to_plane_spherical_paraxial(k, rs_support, Er, -zs_flat, 1/ms_flat, rs_center, qs_center, carrier = False)
# return Er_flat, zs_flat, ms_flat, ms, r_centers_flat
#
# def propagate_plane_to_plane_spherical(k, rs_support, Er, z, rs_center = (0, 0), qs_center = (0, 0), rocs = np.inf):
# rs_support = to_pair(rs_support)
# rs_center = np.asarray(rs_center)
# qs_center = np.asarray(qs_center)
# rocs = to_pair(rocs)
# num_pointss = Er.shape[-2:]
# Er_flat, zs_flat, ms_flat, ms, r_centers_flat = propagate_plane_to_waist_spherical_paraxial(k, rs_support, Er, z, rs_center, qs_center, rocs)
# r_supports_flat = rs_support/ms_flat
#
# # See Dane's notebook 2 p105.
# denominator = (k**2-qs_center[0]**2-qs_center[1]**2)**(3/2)
# z_paraxial_x = z*k*(k**2-qs_center[1]**2)/denominator
# z_paraxial_y = z*k*(k**2-qs_center[0]**2)/denominator
# zs_paraxial = np.asarray((z_paraxial_x, z_paraxial_y))
#
# kx_flat, ky_flat = [calc_q(r_support_flat, num_points, q_center, axis) for r_support_flat, num_points, q_center, axis in zip(r_supports_flat, num_pointss, qs_center, (-2, -1))]
# extra_propagator = mathx.expj(calc_kz(k, kx_flat, ky_flat)*z+z_paraxial_x*kx_flat**2/(2*k)+z_paraxial_y*ky_flat**2/(2*k))
#
# Eq_flat = fft2(Er_flat)
# Eq_flat *= extra_propagator
# Er_flat = ifft2(Eq_flat)
#
# Er, _ = propagate_plane_to_plane_spherical_paraxial(k, r_supports_flat, Er_flat, zs_paraxial+zs_flat, ms*ms_flat, r_centers_flat, qs_center, carrier = False)
# rp_centers = rs_center+z*qs_center/(k**2-qs_center**2)**0.5
#
# return Er, ms, rp_centers
def calc_beam_properties(rs_support, z, Er, Igradphi, rs_center=(0, 0)):
"""Calculate real space and angular centroids, and curvature.
Assumes that propagation is to +z. Positive ROC means diverging i.e. center of curvature is to left.
Args:
k:
x:
y:
Er:
Igradphi (tuple): I times gradient of phase along the x, y and z.
rs_center:
Ir:
Returns:
"""
x, y = sa.calc_xy(rs_support, Er.shape, rs_center)
Ir = mathx.abs_sqd(Er)
sumIr = Ir.sum()
# Calculate centroid. Improve the input estimate rs_center.
rs_center[0], rs_center[1], varx, vary, _ = mathx.mean_and_variance2(
x, y, Ir, sumIr)
# Mean transverse wavenumber is intensity-weighted average of transverse gradient of phase.
qs_center = np.asarray([component.sum()
for component in Igradphi[:2]]) / sumIr
meanz = mathx.moment(z, Ir, 1, sumIr)
# Correct spatial phase for surface curvature - approximate propagation from z to meanz.
Er = Er * mathx.expj((meanz - z) * mathx.divide0(Igradphi[2], Ir))
# Do this twice, since on the first pass we are using qs_center from Igradphi which is just an estimate.
for _ in range(2):
# Calculate phase of quadratic component at RMS distance from center. Proportional to A in Siegman IEE J. Quantum Electronics Vol. 27
# 1991. Positive means diverging.
phi_cx = 0.5 * ((x - rs_center[0]) *
(Igradphi[0] - Ir * qs_center[0])).sum() / sumIr
phi_cy = 0.5 * ((y - rs_center[1]) *
(Igradphi[1] - Ir * qs_center[1])).sum() / sumIr
# Fourier transform Er with quadratic phase removed.
Ek = math.fft2(Er *
mathx.expj(-(x - rs_center[0])**2 * phi_cx / varx) *
mathx.expj(-(y - rs_center[1])**2 * phi_cy / vary))
kx, ky = sa.calc_kxky(rs_support, Er.shape, qs_center)
# Calculate mean square sizes of Fourier transform with quadratic phase removed.
#qs_center[0], qs_center[1], varkxp, varkyp, _ = mathx.mean_and_variance2(kx, ky, mathx.abs_sqd(Ek))
qs_center[0] = mathx.moment(kx, abs(Ek), 1)
qs_center[1] = mathx.moment(ky, abs(Ek), 1)
Ik = mathx.abs_sqd(Ek)
sumIk = Ik.sum()
varkxp = mathx.moment(kx - qs_center[0], Ik, 2, sumIk)
varkyp = mathx.moment(ky - qs_center[1], Ik, 2, sumIk)
# Calculate angular variance.
varkx = bvar.infer_angular_variance_spherical(varx, phi_cx, varkxp)
varky = bvar.infer_angular_variance_spherical(vary, phi_cy, varkyp)
return meanz, rs_center, np.asarray((varx, vary)), qs_center, np.asarray(
(varkx, varky)), np.asarray((phi_cx, phi_cy))