430 lines
14 KiB
Python
430 lines
14 KiB
Python
from typing import Tuple
|
|
import multiprocessing
|
|
import logging
|
|
import copy
|
|
from itertools import chain
|
|
|
|
import pyopencl, meanas, gridlock
|
|
import numpy
|
|
from numpy import pi, sin, cos, exp
|
|
from numpy.linalg import norm
|
|
|
|
from meanas import fdtd, fdfd
|
|
from meanas.fdmath import vec
|
|
from meanas.fdfd.waveguide_3d import compute_source, compute_overlap_e
|
|
from meanas.fdfd import operators
|
|
from meanas.fdtd import maxwell_e, maxwell_h, cpml_params, updates_with_cpml, poynting
|
|
|
|
|
|
numpy.set_printoptions(linewidth=int(1e10))
|
|
|
|
logging.basicConfig(level=logging.DEBUG)
|
|
logger = logging.getLogger(__name__)
|
|
logging.getLogger('matplotlib').setLevel(logging.WARNING)
|
|
logging.getLogger('pyopencl').setLevel(logging.WARNING)
|
|
logging.getLogger('pytools').setLevel(logging.WARNING)
|
|
|
|
|
|
fh = logging.FileHandler('opt.log')
|
|
fh.setLevel(logging.INFO)
|
|
logger.addHandler(fh)
|
|
|
|
|
|
def saveplot_2d_mp(*args):
|
|
multiprocessing.Process(target=saveplot_2d, args=args).start()
|
|
|
|
|
|
def saveplot_2d(val, name, xz_coords):
|
|
val = numpy.squeeze(val)
|
|
pyplot.figure(figsize=(8, 6))
|
|
xz_grids = numpy.meshgrid(*xz_coords, indexing='ij')
|
|
vmax = numpy.abs(val).max()
|
|
if (val < 0).any():
|
|
args = {'vmin': -vmax, 'vmax': vmax, 'cmap': 'seismic'}
|
|
else:
|
|
args = {'vmin': 0, 'vmax': vmax, 'cmap': 'hot'}
|
|
|
|
pyplot.pcolormesh(*xz_grids, val, **args)
|
|
pyplot.colorbar(orientation='horizontal')
|
|
pyplot.title(f'{name}')
|
|
pyplot.gca().set_aspect('equal', adjustable='box')
|
|
pyplot.savefig(f'{name}.png', dpi=240)
|
|
pyplot.close()
|
|
|
|
|
|
def pulse(wl, dwl, dt, turn_on=1e-10):
|
|
# dt * dw = 4 * ln(2)
|
|
w = 2 * pi / wl
|
|
freq = 1 / wl
|
|
fwhm = dwl * w * w / (2 * pi)
|
|
alpha = (fwhm * fwhm) / 8 * numpy.log(2)
|
|
delay = numpy.sqrt(-numpy.log(turn_on) / alpha)
|
|
delay = numpy.ceil(delay * freq) / freq # force delay to integer number of periods to maintain phase
|
|
logger.info(f'src_time {2 * delay / dt}')
|
|
|
|
n = numpy.floor(pi / (w * dt))
|
|
logger.info(f'save timestep would be {n} * dt = {n * dt}')
|
|
|
|
# nrm = numpy.exp(-w * w / alpha) / 2
|
|
|
|
def source_phasor(i):
|
|
t0 = i * dt - delay
|
|
envelope = numpy.sqrt(numpy.sqrt(2 * alpha / pi)) * numpy.exp(-alpha * t0**2)
|
|
# if t0 < 0:
|
|
# envelope = numpy.exp(-alpha * t0**2)
|
|
# else:
|
|
# envelope = 1
|
|
return envelope, numpy.cos(w * t0), numpy.sin(w * t0)
|
|
return source_phasor, delay, n #, nrm
|
|
|
|
|
|
def get_wgmode_xp(half_dims, polarity, grid, epsilon, wl, dxes):
|
|
dims = [-half_dims, half_dims]
|
|
dims[0][0] = dims[1][0]
|
|
ind_dims = (grid.pos2ind(dims[0], which_shifts=None).astype(int),
|
|
grid.pos2ind(dims[1], which_shifts=None).astype(int))
|
|
wg_slices = tuple(slice(i, f+1) for i, f in zip(*ind_dims))
|
|
wg_args = {
|
|
'omega': 2 * pi / wl,
|
|
'slices': wg_slices,
|
|
'dxes': dxes,
|
|
'axis': 0,
|
|
'polarity': polarity,
|
|
}
|
|
|
|
wg_results = fdfd.waveguide_3d.solve_mode(mode_number=0, **wg_args, epsilon=epsilon)
|
|
return wg_args, wg_results
|
|
|
|
|
|
def get_gaussian(m, grid, dxes, wl):
|
|
def grid2gaussian(xyz, center, w0=4600, tilt=numpy.deg2rad(-8)):
|
|
xs, ys, zs = xyz
|
|
xs -= center[0]
|
|
ys -= center[1]
|
|
zs -= center[2]
|
|
xg, yg, zg = numpy.meshgrid(xs, ys, zs, indexing='ij')
|
|
|
|
rot = numpy.array([[ cos(tilt), 0, sin(tilt)],
|
|
[ 0, 1, 0],
|
|
[-sin(tilt), 0, cos(tilt)]])
|
|
|
|
x, y, z = (rot @ numpy.stack((xg.ravel(), yg.ravel(), zg.ravel()))).reshape(3, *grid.shape)
|
|
r2 = x * x + y * y # sq. distance from beam center along tilted plane
|
|
z2 = z * z # sq. distance from waist along centerline
|
|
|
|
zr = pi * w0 * w0 / wl
|
|
zr2 = zr * zr
|
|
wz2 = w0 * w0 * (1 + z2 / zr2)
|
|
wz = numpy.sqrt(wz2)
|
|
|
|
k = 2 * pi / wl
|
|
Rz = z * (1 + zr2 / z2)
|
|
gouy = numpy.arctan(z / zr)
|
|
|
|
gaussian = w0 / wz * exp(-r2 / wz2) * exp(1j * (k * z + k * r2 / 2 / Rz - gouy))
|
|
# window_x = scipy.signal.windows.kaiser(xs.size, 14)
|
|
# gaussian *= window_x[:, None, None]
|
|
return gaussian
|
|
|
|
zsEy = grid.shifted_xyz(1)[2]
|
|
gaussianEy = grid2gaussian(grid.shifted_xyz(1), [0, 0, zsEy[m[2]]])
|
|
|
|
normEy = gaussianEy[m[0]:-m[0], :, m[2]]
|
|
gaussianEy /= numpy.sqrt((normEy[1].conj() * normEy[1]).sum())
|
|
|
|
return gaussianEy
|
|
|
|
|
|
def run(pml=(10, 0, 10), dx=20, wl=1310, dwl=130, wg_zh=400, wg_x=-7500, fiber_z=1000, max_t=int(10e3)):
|
|
omega = 2 * pi / wl
|
|
|
|
x_min = -10e3 - pml[0] * dx
|
|
x_max = 10e3 + pml[0] * dx
|
|
z_min = -600 - pml[2] * dx
|
|
z_max = 1400 + pml[2] * dx
|
|
|
|
ex = numpy.arange(x_min, x_max + dx / 2, dx)
|
|
ez = numpy.arange(z_min, z_max + dx / 2, dx)
|
|
exyz = [ex, [-dx / 2, dx / 2], ez]
|
|
grid = gridlock.Grid(exyz, periodic=True)
|
|
epsilon = grid.allocate(1.45**2)
|
|
|
|
def unvec(f):
|
|
return meanas.fdmath.unvec(f, grid.shape)
|
|
|
|
# grid.draw_slab(epsilon, surface_normal=2, center=[0, 0, 0], thickness=160, eps=3.5**2)
|
|
|
|
|
|
e = numpy.zeros_like(epsilon, dtype=numpy.float32)
|
|
h = numpy.zeros_like(epsilon, dtype=numpy.float32)
|
|
|
|
dxes = [grid.dxyz, grid.autoshifted_dxyz()]
|
|
min_dx = min(min(dxn) for dxn in chain(*dxes))
|
|
dt = min_dx * .99 / numpy.sqrt(3)
|
|
|
|
source_phasor, delay, n_fft = pulse(wl, dwl, dt)
|
|
if 2 * delay / dt > max_t:
|
|
raise Exception('Source extends beyond end of simulation')
|
|
|
|
|
|
m = numpy.array(pml) + 10
|
|
m[2] = grid.pos2ind([0, 0, fiber_z], which_shifts=0)[2] - grid.shape[2]
|
|
|
|
ey_gauss = numpy.zeros_like(epsilon, dtype=complex)
|
|
ey_gauss = get_gaussian(m, grid, dxes, wl / 1.45)
|
|
e_gauss = numpy.zeros_like(epsilon, dtype=numpy.complex64)
|
|
e_gauss[1] = ey_gauss
|
|
mask = numpy.zeros_like(epsilon, dtype=int)
|
|
mask[..., :m[2]] = 1
|
|
src_op = operators.e_boundary_source(mask=vec(mask), omega=omega, dxes=dxes, epsilon=vec(epsilon))
|
|
|
|
def zero_pmls(c):
|
|
for a in range(3):
|
|
c[a][:pml[0]+1, :, :] = 0
|
|
c[a][-pml[0]-1:, :, :] = 0
|
|
c[a][:, :, :pml[2]+1] = 0
|
|
c[a][:, :, -pml[2]-1:] = 0
|
|
return c
|
|
|
|
# J = unvec(src_op @ vec(e_gauss))
|
|
# J[:, :12, :, :] = 0
|
|
# J[:, -12:, :, :] = 0
|
|
# J[:, :, :, :12] = 0
|
|
# J[:, :, :, -12:] = 0
|
|
# zero_pmls(J)
|
|
|
|
J = numpy.zeros_like(epsilon, dtype=complex)
|
|
J[1, 500, 0, 60] = 1
|
|
zero_pmls(J)
|
|
|
|
half_dims = numpy.array([wg_x, dx, wg_zh])
|
|
wg_args, wg_results = get_wgmode_xp(half_dims, -1, grid, wl, dxes)
|
|
|
|
E_out = compute_overlap_e(E=wg_results['E'], wavenumber=wg_results['wavenumber'],
|
|
dxes=dxes, axis=0, polarity=+1, slices=wg_args['slices'])
|
|
|
|
|
|
jr = (J.real / epsilon).astype(numpy.float32)
|
|
ji = (J.imag / epsilon).astype(numpy.float32)
|
|
|
|
|
|
eph = numpy.zeros_like(e, dtype=numpy.complex64)
|
|
ephm = numpy.zeros_like(e, dtype=numpy.complex64)
|
|
# powers = numpy.zeros((max_t, 5))
|
|
p_ph = 0
|
|
|
|
pml_params = [[cpml_params(axis=dd, polarity=pp, dt=dt,
|
|
thickness=pml[dd], epsilon_eff=1.0**2)
|
|
if pml[dd] > 0 else None
|
|
for pp in (-1, +1)]
|
|
for dd in range(3)]
|
|
update_E, update_H = updates_with_cpml(cpml_params=pml_params, dt=dt,
|
|
dxes=dxes, epsilon=epsilon)
|
|
|
|
mov_interval = 10
|
|
mov = numpy.empty((max_t // mov_interval, e.shape[1], e.shape[3]), dtype=numpy.float32)
|
|
|
|
for t in range(max_t):
|
|
update_E(e, h, epsilon)
|
|
_, cm5, sm5 = source_phasor(t - 0.5)
|
|
ephm += (cm5 - 1j * sm5) * e
|
|
|
|
a, c, s = source_phasor(t)
|
|
p_ph += a * c * c
|
|
e -= (a * c) * jr - (a * s) * ji
|
|
|
|
update_H(e, h)
|
|
|
|
_, cp5, sp5 = source_phasor(t + 0.5)
|
|
eph += (cp5 - 1j * sp5) * e
|
|
|
|
# S = poynting(e, h, epsilon)
|
|
#
|
|
# powers[t, :] = (
|
|
# numpy.sum(S[2, m[0]+3:-m[0]-2, :, m[2]-6]), # below src
|
|
# numpy.sum(S[2, m[0]+3:-m[0]-2, :, m[2]+4]), # above src
|
|
# numpy.sum(S[2, m[0]+3:-m[0]-2, :, pml[2]+2]), # bottom
|
|
# numpy.sum(S[0, +m[0]+2, :, pml[2]+3:m[2]+4]), # left
|
|
# numpy.sum(S[0, -m[0]-2, :, pml[2]+3:m[2]+4]), # right
|
|
# )
|
|
|
|
if t % mov_interval == 0:
|
|
mov[t // mov_interval] = e[1, :, 0, :].real
|
|
|
|
eph *= dt / p_ph
|
|
ephm *= dt / p_ph
|
|
|
|
src_power = -(J * eph).real.sum() / 2 * dx ** 3
|
|
|
|
hph = meanas.fdfd.functional.e2h(omega=omega, dxes=dxes)(eph)
|
|
sph = meanas.fdtd.poynting(e=eph, h=hph.conj(), dxes=dxes)
|
|
planes_powers = numpy.array((
|
|
-sph[0, 11, :, 11:-12].sum(),
|
|
+sph[0, -12, :, 11:-12].sum(),
|
|
-sph[2, 11:-12, :, 11].sum(),
|
|
+sph[2, 11:-12, :, -12].sum(),
|
|
)).real / 2
|
|
planes_power = planes_powers.sum()
|
|
|
|
print(f'{src_power=}, {planes_power=}')
|
|
|
|
|
|
# Verify
|
|
A = meanas.fdfd.operators.e_full(omega=omega, dxes=dxes, epsilon=vec(epsilon))
|
|
b = -1j * omega * vec(J) #* numpy.exp(1j * dt / 2 * omega)
|
|
c = A @ vec(eph)
|
|
logger.info('FWD inaccuracy: |Ax-b|/|b| = {}'.format(norm(c-b) / norm(b)))
|
|
normdiv = norm(b) / norm(c)
|
|
logger.info(f'{normdiv=}')
|
|
logger.info('FWD renormed inaccuracy: |Ax-b|/|b| = {}'.format(norm(c * normdiv - b) / norm(b)))
|
|
|
|
b = -1j * omega * vec(J)
|
|
logger.info('FWD base inaccuracy: |Ax-b|/|b| = {}'.format(norm(c-b) / norm(b)))
|
|
|
|
from scipy.optimize import minimize
|
|
def resid(x):
|
|
b = -1j * omega * vec(J) * numpy.exp(1j * dt * x * omega)
|
|
return norm(c - b) / norm(b)
|
|
print('min', minimize(resid, 0.25, options={'xatol': 1e-14, 'fatol': 1e-14}))
|
|
|
|
# fig, ax, anim = plot_movie(mov, balanced=True, interval=300)
|
|
# anim.save('output.mp4')
|
|
|
|
print('solving...')
|
|
cdxes = copy.deepcopy(dxes)
|
|
for axis in range(3):
|
|
thickness = pml[axis]
|
|
if not thickness:
|
|
continue
|
|
for pp, polarity in enumerate((-1, 1)):
|
|
print(axis, polarity, thickness)
|
|
cdxes = fdfd.scpml.stretch_with_scpml(cdxes, axis=axis, polarity=polarity,
|
|
omega=omega, epsilon_effective=1.0**2,
|
|
thickness=thickness)
|
|
eph2v = meanas.fdfd.solvers.generic(
|
|
omega=omega, dxes=cdxes, J=vec(J), epsilon=vec(epsilon),
|
|
matrix_solver_opts={'atol': 1e-3, 'tol': 1e-3, 'x0': vec(eph)})
|
|
eph2 = unvec(eph2v)
|
|
|
|
pyplot.figure()
|
|
pyplot.pcolormesh(numpy.abs(eph/eph2)[1, 11:-11, 0, 11:-11].real.T)
|
|
pyplot.colorbar()
|
|
pyplot.title('mag')
|
|
pyplot.figure()
|
|
pyplot.pcolormesh(numpy.angle(eph/eph2)[1, 11:-11, 0, 11:-11].real.T)
|
|
pyplot.colorbar()
|
|
pyplot.title('angle')
|
|
pyplot.show()
|
|
breakpoint()
|
|
|
|
|
|
import matplotlib
|
|
from matplotlib import cycler, animation, colors, ticker, pyplot
|
|
|
|
def set_pyplot_cycle() -> None:
|
|
pyplot.rc('lines', linewidth=2.5)
|
|
pyplot.rc('axes', prop_cycle(
|
|
cycler('color', 'krbgcm')
|
|
* cycler('linestyle', ['-', '--', ':', '-.'])
|
|
))
|
|
|
|
|
|
def pcm(x, y, z, pca={}, cba={}, bare=False, eq=True) -> Tuple:
|
|
z = numpy.array(z)
|
|
|
|
if numpy.any(z < 0):
|
|
vmax = numpy.abs(z).max()
|
|
pcolor_args = {'vmin': -vmax, 'vmax': vmax, 'cmap': 'seismic', **pca}
|
|
else:
|
|
pcolor_args = {'cmap': 'seismic', **pca}
|
|
|
|
xe = centers2edges(x)
|
|
ye = centers2edges(y)
|
|
|
|
if bare:
|
|
fig = pyplot.gcf()
|
|
ax = pyplot.gca()
|
|
else:
|
|
fig, ax = pyplot.subplot()
|
|
|
|
im = ax.pcolormesh(xe, ye, z.T, **pcolor_args)
|
|
if eq:
|
|
ax.set_aspect('equal', adjustable='box')
|
|
|
|
if not bare:
|
|
ax.format_coord = lambda xx, yy: format_coord(xx, yy, xe, ye, z.T)
|
|
fig.colorbar(im, ax=ax, **cba)
|
|
|
|
return fig, ax, im
|
|
|
|
|
|
def pcc(x, y, z, cfa={}, cba={}, n_levels: int = 15, bare: bool = False, eq: bool = True) -> Tuple:
|
|
z = numpy.array(z)
|
|
|
|
if numpy.any(z < 0):
|
|
vmax = numpy.abs(z).max()
|
|
pcolor_args = {'vmin': -vmax, 'vmax': vmax, 'cmap': 'seismic', **cfa}
|
|
else:
|
|
pcolor_args = {'cmap': 'hot', **cfa}
|
|
|
|
xe = centers2edges(x)
|
|
ye = centers2edges(y)
|
|
|
|
if bare:
|
|
fig = pyplot.gcf()
|
|
ax = pyplot.gca()
|
|
else:
|
|
fig, ax = pyplot.subplot()
|
|
|
|
levels = ticker.MaxNLocator(nbins=n_levels).tick_values(z.min(), z.max())
|
|
cmap = pyplot.get_cmap(pcolor_args['cmap'])
|
|
norm = color.BoundaryNorm(levels, ncolors=cmap.N, clip=True)
|
|
|
|
im = ax.contourf(x, y, z.T, levels=levels, **pcolor_args)
|
|
if eq:
|
|
ax.set_aspect('equal', adjustable='box')
|
|
|
|
if not bare:
|
|
ax.format_coord = lambda xx, yy: format_coord(xx, yy, xe, ye, z.T)
|
|
fig.colorbar(im, ax=ax, **cba)
|
|
|
|
return fig, ax, im
|
|
|
|
|
|
def centers2edges(centers):
|
|
d = numpy.diff(centers) / 2
|
|
e = numpy.hstack((centers[0] - d[0], centers[:-1] + d, centers[-1] + d[-1]))
|
|
return e
|
|
|
|
|
|
def format_coord(x, y, xs, ys, vs):
|
|
col = numpy.digitize(x, xs)
|
|
row = numpy.digitize(y, ys)
|
|
if 0 < row <= vs.shape[0] and 0 < col <= vs.shape[1]:
|
|
z = vs[row - 1, col - 1]
|
|
return f'x={x:1.4g}, y={y:1.4g}, z={z:1.4g}'
|
|
else:
|
|
return f'x={x:1.4g}, y={y:1.4g}'
|
|
|
|
|
|
def plot_movie(arr, balanced=True, interval=300, pca={}):
|
|
if balanced:
|
|
vmax = numpy.abs(arr).max()
|
|
pcolor_args = {'vmin': -vmax, 'vmax': vmax, 'cmap': 'seismic', **pca}
|
|
else:
|
|
pcolor_args = {'cmap': 'seismic', **pca}
|
|
|
|
fig, ax = pyplot.subplots()
|
|
im = ax.pcolormesh(arr[0, :, :].T, **pcolor_args)
|
|
ax.set_aspect('equal', adjustable='box')
|
|
|
|
def animate(ii):
|
|
im.set_array(arr[ii, :, :].T.ravel())
|
|
|
|
anim = animation.FuncAnimation(fig, animate, frames=arr.shape[0], repeat=True, interval=interval)
|
|
return fig, im, anim
|
|
|
|
|
|
if __name__ == '__main__':
|
|
run()
|