def SphericalClipping(A, dr, dz, cells):
rho = np.outer(grid_tools.cell_centers(0, dr * A.shape[0], A.shape[0]),
np.ones(A.shape[1]))
z = np.outer(
np.ones(A.shape[0]),
grid_tools.cell_centers(-dz * A.shape[1] / 2, dz * A.shape[1] / 2,
A.shape[1]))
clip_rgn = np.where(rho**2 + z**2 > (dr * (A.shape[0] - cells))**2)
A[clip_rgn] = 0.0
def GetGridInfo(self): w_nodes = self.center[0] + grid_tools.cyclic_nodes(-self.size[0]/2,self.size[0]/2,self.pts[0]) x_nodes = grid_tools.cell_centers(-self.size[1]/2,self.size[1]/2,self.pts[1]) y_nodes = grid_tools.cell_centers(-self.size[2]/2,self.size[2]/2,self.pts[2]) w_walls = grid_tools.cell_walls(w_nodes[0],w_nodes[-1],self.pts[0],self.default_band) x_walls = grid_tools.cell_walls(x_nodes[0],x_nodes[-1],self.pts[1]) y_walls = grid_tools.cell_walls(y_nodes[0],y_nodes[-1],self.pts[2]) plot_ext = np.array([w_walls[0],w_walls[-1],x_walls[0],x_walls[-1],y_walls[0],y_walls[-1]]) return w_nodes,x_nodes,y_nodes,plot_ext
def test_full_slab_edge(self):
a = np.ones((4, 4, 4))
a[2, 2, 3] = 2.0
wn = grid_tools.cyclic_nodes(0.0, 4.0, 4)
xn = grid_tools.cell_centers(0.0, 4.0, 4)
yn = grid_tools.cell_centers(0.0, 4.0, 4)
w = np.array([2.0, 2.0, 2.0])
x = np.array([2.5, 2.5, 2.5])
y = np.array([4.0, 3.5, 3.0])
data = grid_tools.DataFromGrid(w, x, y, wn, xn, yn, a)
assert np.allclose(data, [2.5, 2.0, 1.5])
def LoadMap(self,input_dict):
N = input_dict['mesh points']
N = (N[0]+4,N[1]+4)
self.dr = self.Rd/(N[0]-4)
self.dz = self.Lz/(N[1]-4)
rho = grid_tools.cell_centers(-2*self.dr,self.Rd+2*self.dr,N[0])
z = grid_tools.cell_centers(-2*self.dz-self.Lz/2,2*self.dz+self.Lz/2,N[1])
coeff = self.vol_dict['radial coefficients']
fr = coeff[0] + coeff[1]*rho**2 + coeff[2]*rho**4 + coeff[3]*rho**6
fz = np.ones(N[1])
#fz = np.exp(-z**2/self.Lz**2)
self.ne = input_dict['density multiplier']*np.outer(fr,fz)
def LoadMap(self,input_dict):
N = input_dict['mesh points']
N = (N[0]+4,N[1]+4,N[2]+4)
self.dx = self.size[0]/(N[0]-4)
self.dy = self.size[1]/(N[1]-4)
self.dz = self.size[2]/(N[2]-4)
x = grid_tools.cell_centers(-2*self.dx-self.size[0]/2,2*self.dx+self.size[0]/2,N[0])
y = grid_tools.cell_centers(-2*self.dy-self.size[1]/2,2*self.dy+self.size[1]/2,N[1])
z = grid_tools.cell_centers(-2*self.dz-self.size[2]/2,2*self.dz+self.size[2]/2,N[2])
rho2 = np.outer(x**2,np.ones(N[1])) + np.outer(np.ones(N[0]),y**2)
coeff = self.vol_dict['radial coefficients']
fr = coeff[0] + coeff[1]*rho2 + coeff[2]*rho2**2 + coeff[3]*rho2**4
self.ne = input_dict['density multiplier']*np.einsum('ij,k->ijk',fr,np.ones(N[2]))
def spherical_to_cylindrical(A,q_list,r_list,rho_pts,z_pts): rmax = r_list[-1] + 0.5*(r_list[1]-r_list[0]) B = np.zeros((rho_pts,z_pts)).astype(np.complex) rho_list = grid_tools.cell_centers(0,rmax,rho_pts) z_list = grid_tools.cell_centers(-rmax,rmax,z_pts) rho = np.outer(rho_list,np.ones(z_pts)) z = np.outer(np.ones(rho_pts),z_list) fr = scipy.interpolate.RectBivariateSpline(q_list,r_list,np.real(A),kx=3,ky=3) fi = scipy.interpolate.RectBivariateSpline(q_list,r_list,np.imag(A),kx=3,ky=3) r = np.sqrt(rho**2 + z**2) theta = np.arccos(z/r) B = fr.ev(theta,r) + 1j*fi.ev(theta,r) clip_rgn = np.where(rho**2+z**2 > rmax**2) B[clip_rgn] = 0 return rho_list,z_list,B
def GetGridInfo(self): w_nodes = self.center[0] + grid_tools.cyclic_nodes(-self.size[0]/2,self.size[0]/2,self.pts[0]) rho_nodes = grid_tools.cell_centers(0.0,self.size[1]/2,self.pts[1]) phi_nodes = grid_tools.cyclic_nodes(-np.pi,np.pi,self.pts[2]) w_walls = grid_tools.cell_walls(w_nodes[0],w_nodes[-1],self.pts[0],self.default_band) rho_walls = grid_tools.cell_walls(rho_nodes[0],rho_nodes[-1],self.pts[1]) phi_walls = grid_tools.cell_walls(phi_nodes[0],phi_nodes[-1],self.pts[2]) plot_ext = np.array([w_walls[0],w_walls[-1],rho_walls[0],rho_walls[-1],phi_walls[0],phi_walls[-1]]) return w_nodes,rho_nodes,phi_nodes,plot_ext
def GetFields(self,dz,A):
'''Compute Maxwell fields in a sequence of planes in any geometry.
The geometry is encapsulated in the TransverseModeTool self.T.
dz = distance from eikonal plane to center of interrogation region
A = complex amplitude (any Cartesian component) in eikonal plane'''
w_nodes,x1_nodes,x2_nodes,plot_ext = self.GetGridInfo()
z_nodes = grid_tools.cell_centers(dz-self.size[3]/2,dz+self.size[3]/2,self.pts[3])
ans = np.einsum('ijk,l->ijkl',A,np.ones(self.pts[3]))
ans = self.T.kspace(ans)
phase_adv = np.ones(ans.shape).astype(np.complex) * z_nodes[np.newaxis,np.newaxis,np.newaxis,:]
kz = np.sqrt(0j + np.zeros(ans.shape[:3]) + w_nodes[...,np.newaxis,np.newaxis]**2 - self.T.kr2()[np.newaxis,...])
phase_adv *= kz[...,np.newaxis]
# Following applies galilean transformation from z,t to z,t-z/c
phase_adv -= np.outer(w_nodes,z_nodes)[:,np.newaxis,np.newaxis,:]
phase_adv[np.where(np.imag(phase_adv)<0)] *= -1
ans *= np.exp(1j*phase_adv)
ans = self.T.rspace(ans)
z_ext = np.array([-self.size[3]/2,self.size[3]/2])
dom4d = np.concatenate((plot_ext,z_ext))
return ans,dom4d
def load_rays_xw(xp, bundle_radius, N, box, loading_coordinates):
'''Load the rays in the z=0 plane in a regular pattern.'''
num_bundles = N[0] * N[1] * N[2] * N[3]
o0 = np.ones(N[0])
o1 = np.ones(N[1])
o2 = np.ones(N[2])
o3 = np.ones(N[3])
# load frequencies to respect FFT conventions
grid0 = grid_tools.cyclic_nodes(box[0], box[1], N[0])
grid1 = grid_tools.cell_centers(box[2], box[3], N[1])
if loading_coordinates == 'cartesian':
grid2 = grid_tools.cell_centers(box[4], box[5], N[2])
else:
grid2 = grid_tools.cyclic_nodes(box[4], box[5], N[2])
grid3 = grid_tools.cell_centers(box[6], box[7], N[3])
if box[0] == 0.0:
grid0 = grid0[1:]
o0 = o0[1:]
num_bundles -= N[1] * N[2] * N[3]
# Load the primary rays in configuration+w space
if loading_coordinates == 'cartesian':
xp[:, 0, 0] = 0.0
xp[:, 0, 1] = np.einsum('i,j,k,l', o0, grid1, o2,
o3).reshape(num_bundles)
xp[:, 0, 2] = np.einsum('i,j,k,l', o0, o1, grid2,
o3).reshape(num_bundles)
xp[:, 0, 3] = np.einsum('i,j,k,l', o0, o1, o2,
grid3).reshape(num_bundles)
xp[:, 0, 4] = np.einsum('i,j,k,l', grid0, o1, o2,
o3).reshape(num_bundles)
else:
xp[:, 0, 0] = 0.0
xp[:, 0, 1] = np.einsum('i,j,k,l', o0, grid1, np.cos(grid2),
o3).reshape(num_bundles)
xp[:, 0, 2] = np.einsum('i,j,k,l', o0, grid1, np.sin(grid2),
o3).reshape(num_bundles)
xp[:, 0, 3] = np.einsum('i,j,k,l', o0, o1, o2,
grid3).reshape(num_bundles)
xp[:, 0, 4] = np.einsum('i,j,k,l', grid0, o1, o2,
o3).reshape(num_bundles)
# Load the satellite rays in configuration+w space
xp[:, 1, :] = xp[:, 0, :]
xp[:, 2, :] = xp[:, 0, :]
xp[:, 3, :] = xp[:, 0, :]
xp[:, 4, :] = xp[:, 0, :]
xp[:, 5, :] = xp[:, 0, :]
xp[:, 6, :] = xp[:, 0, :]
xp[:, 1, 1] += bundle_radius[1]
xp[:, 2, 1] -= bundle_radius[1]
xp[:, 3, 2] += bundle_radius[2]
xp[:, 4, 2] -= bundle_radius[2]
xp[:, 5, 3] += bundle_radius[3]
xp[:, 6, 3] -= bundle_radius[3]
def track(cl, xp, eikonal, vg, vol_dict):
'''Propagate unidirectional fully dispersive waves using eikonal data as a boundary condition.
The volume must be oriented so the polarization axis is x (linear polarization only).
:param numpy.array xp: ray phase space with shape (bundles,rays,8)
:param numpy.array eikonal: ray eikonal data with shape (bundles,4)
:param numpy.array vg: ray group velocity with shape (bundles,rays,4)
:param dictionary vol_dict: input file dictionary for the volume'''
band = vol_dict['frequency band']
size = (band[1] - band[0], ) + vol_dict['size']
N = vol_dict['wave grid']
# N[3] is the number of diagnostic planes, including the initial plane
diagnostic_steps = N[3] - 1
subcycles = vol_dict['subcycles']
steps = diagnostic_steps * subcycles
field_planes = steps + 1
powersof2 = [2**i for i in range(32)]
if N[0] - 1 not in powersof2:
raise ValueError('UPPE propagator requires 2**n+1 w-nodes')
if N[1] not in powersof2:
raise ValueError('UPPE propagator requires 2**n x-nodes')
if N[2] not in powersof2:
raise ValueError('UPPE propagator requires 2**n y-nodes')
try:
window_speed = vol_dict['window speed']
except:
window_speed = 1.0
try:
chi3 = vol_dict['chi3']
except:
chi3 = 0.0
# Capture the rays
if vol_dict['wave coordinates'] == 'cartesian':
field_tool = caustic_tools.FourierTool(N, band, (0, 0, 0), size[1:],
cl)
else:
field_tool = caustic_tools.BesselBeamTool(N, band, (0, 0, 0), size[1:],
cl)
w_nodes, x1_nodes, x2_nodes, plot_ext = field_tool.GetGridInfo()
A = np.zeros(N).astype(np.complex)
J = np.zeros(N).astype(np.complex)
ne = np.zeros(N).astype(np.complex)
A0, dom3d = field_tool.GetBoundaryFields(xp[:, 0, :], eikonal, 1)
# Setup the wave propagation domain
chi = vol_dict['dispersion inside'].chi(w_nodes)
try:
ionizer = vol_dict['ionizer']
except KeyError:
ionizer = ionization.Ionization(1.0, 1.0, 1.0, 1.0)
dens_nodes = grid_tools.cell_centers(-size[3] / 2, size[3] / 2, steps)
field_walls = grid_tools.cell_walls(dens_nodes[0], dens_nodes[-1], steps)
diagnostic_walls = np.linspace(-size[3] / 2, size[3] / 2, N[3])
dz = field_walls[1] - field_walls[0]
Dz = diagnostic_walls[1] - diagnostic_walls[0]
dom4d = np.concatenate((dom3d, [field_walls[0], field_walls[-1]]))
# Step through the domain
# Strategy to get density plane is to re-use ray gather system
# This works as long as the shape of xp is (*,*,8)
xp_eff = np.zeros((N[1], N[2], 8))
if vol_dict['wave coordinates'] == 'cartesian':
xp_eff[..., 1] = np.outer(x1_nodes, np.ones(N[2]))
xp_eff[..., 2] = np.outer(np.ones(N[1]), x2_nodes)
else:
xp_eff[..., 1] = np.outer(x1_nodes, np.cos(x2_nodes))
xp_eff[..., 2] = np.outer(x1_nodes, np.sin(x2_nodes))
A[..., 0] = A0
J[..., 0] = 0.0
ne[..., 0] = 0.0
for k in range(diagnostic_steps):
print('Advancing to diagnostic plane', k + 1)
rhs_evals = 0
for s in range(subcycles):
if subcycles > 1:
if s == 0:
print(' subcycling .', end='', flush=True)
else:
print('.', end='', flush=True)
xp_eff[..., 3] = dens_nodes[k * subcycles + s]
dens = vol_dict['object'].GetDensity(xp_eff)
A0, J0, ne0, evals = propagator(cl, field_tool, window_speed, A0,
chi, chi3, dens, ionizer, dz)
A0[:4, ...] = 0.0
rhs_evals += evals
print('', rhs_evals, 'evaluations of j(w,kx,ky)')
A[..., k + 1] = A0
J[..., k + 1] = J0
ne[..., k + 1] = ne0
# Finish by relaunching rays and returning UPPE data
field_tool.RelaunchRays(xp, eikonal, vg, A[..., -1], size[3])
return A, J, ne, dom4d
def track(xp, eikonal, vg, vol_dict):
'''Propagate unidirectional fully dispersive waves using eikonal data as a boundary condition.
The volume must be oriented so the polarization axis is x (linear polarization only).
:param numpy.array xp: ray phase space with shape (bundles,rays,8)
:param numpy.array eikonal: ray eikonal data with shape (bundles,4)
:param numpy.array vg: ray group velocity with shape (bundles,rays,4)
:param dictionary vol_dict: input file dictionary for the volume'''
band = vol_dict['frequency band']
size = (band[1] - band[0], ) + vol_dict['size']
N = vol_dict['wave grid']
# N[3] is the number of diagnostic planes, including the initial plane
diagnostic_steps = N[3] - 1
subcycles = vol_dict['subcycles']
steps = diagnostic_steps * subcycles
field_planes = steps + 1
Vol = vol_dict['object']
try:
window_speed = vol_dict['window speed']
except:
window_speed = 1.0
try:
chi3 = vol_dict['chi3']
except:
chi3 = 0.0
# Capture the rays
if vol_dict['wave coordinates'] == 'cartesian':
field_tool = caustic_tools.FourierTool(N, band, (0, 0, 0), size[1:],
Vol.queue, Vol.transform_k)
else:
field_tool = caustic_tools.BesselBeamTool(N, band, (0, 0, 0), size[1:],
Vol.queue, Vol.transform_k)
w_nodes, x1_nodes, x2_nodes, plot_ext = field_tool.GetGridInfo()
A = np.zeros(N).astype(np.complex)
J = np.zeros(N).astype(np.complex)
ne = np.zeros(N).astype(np.complex)
A0, dom3d = field_tool.GetBoundaryFields(xp[:, 0, :], eikonal, 1)
# Setup the wave propagation domain
chi = vol_dict['dispersion inside'].chi(w_nodes)
try:
ionizer = vol_dict['ionizer']
except KeyError:
ionizer = ionization.Ionization(1.0, 1.0, 1.0, 1.0)
dens_nodes = grid_tools.cell_centers(-size[3] / 2, size[3] / 2, steps)
field_walls = grid_tools.cell_walls(dens_nodes[0], dens_nodes[-1], steps)
diagnostic_walls = np.linspace(-size[3] / 2, size[3] / 2, N[3])
dz = field_walls[1] - field_walls[0]
Dz = diagnostic_walls[1] - diagnostic_walls[0]
dom4d = np.concatenate((dom3d, [field_walls[0], field_walls[-1]]))
# Step through the domain
# Strategy to get density plane is to re-use ray gather system
# This works as long as the shape of xp is (*,*,8)
xp_eff = np.zeros((N[1], N[2], 8))
if vol_dict['wave coordinates'] == 'cartesian':
xp_eff[..., 1] = np.outer(x1_nodes, np.ones(N[2]))
xp_eff[..., 2] = np.outer(np.ones(N[1]), x2_nodes)
else:
xp_eff[..., 1] = np.outer(x1_nodes, np.cos(x2_nodes))
xp_eff[..., 2] = np.outer(x1_nodes, np.sin(x2_nodes))
A[..., 0] = A0
J[..., 0] = 0.0
ne[..., 0] = 0.0
for k in range(diagnostic_steps):
print('Advancing to diagnostic plane', k + 1)
for s in range(subcycles):
xp_eff[..., 3] = dens_nodes[k * subcycles + s]
dens = vol_dict['object'].GetDensity(xp_eff)
A0, J0, ne0 = propagator(field_tool, window_speed, A0, chi, chi3,
dens, ionizer, dz)
try:
A0 *= vol_dict['damping filter'](w_nodes)[:, np.newaxis,
np.newaxis]
except KeyError:
A0 = A0
A[..., k + 1] = A0
J[..., k + 1] = J0
ne[..., k + 1] = ne0
# Return the wave amplitude
# Rays are re-launched externally
return A, J, ne, dom4d
def track(cl, xp, eikonal, vg, vol_dict):
'''Propagate unidirectional fully dispersive waves using eikonal data as a boundary condition.
The volume must be oriented so the polarization axis is x (linear polarization only).
:param numpy.array xp: ray phase space with shape (bundles,rays,8)
:param numpy.array eikonal: ray eikonal data with shape (bundles,4)
:param numpy.array vg: ray group velocity with shape (bundles,rays,4)
:param dictionary vol_dict: input file dictionary for the volume'''
band = vol_dict['frequency band']
NL_band = vol_dict['nonlinear band']
size = (band[1] - band[0], ) + vol_dict['size']
N = vol_dict['wave grid']
# N[3] is the number of diagnostic planes, including the initial plane
diagnostic_steps = N[3] - 1
subcycles = vol_dict['subcycles']
steps = diagnostic_steps * subcycles
field_planes = steps + 1
powersof2 = [2**i for i in range(32)]
if N[0] - 1 not in powersof2:
raise ValueError('UPPE propagator requires 2**n+1 w-nodes')
if N[1] not in powersof2:
raise ValueError('UPPE propagator requires 2**n x-nodes')
if N[2] not in powersof2:
raise ValueError('UPPE propagator requires 2**n y-nodes')
try:
window_speed = vol_dict['window speed']
except KeyError:
window_speed = 1.0
try:
chi3 = vol_dict['chi3']
except KeyError:
chi3 = 0.0
try:
full_relaunch = vol_dict['full relaunch']
except KeyError:
full_relaunch = False
# Capture the rays
if vol_dict['wave coordinates'] == 'cartesian':
field_tool = caustic_tools.FourierTool(N, band, (0, 0, 0), size[1:],
cl)
else:
field_tool = caustic_tools.BesselBeamTool(N, band, (0, 0, 0), size[1:],
cl, vol_dict['radial modes'])
warnings.warn(
'Polarization information is lost upon entering paraxial region.')
w_nodes, x1_nodes, x2_nodes, plot_ext = field_tool.GetGridInfo()
A = np.zeros(N).astype(np.complex)
J = np.zeros(N).astype(np.complex)
ne = np.zeros(N).astype(np.complex)
A[..., 0], dom3d = field_tool.GetBoundaryFields(xp[:, 0, :], eikonal, 1)
# Setup the wave propagation medium
chi = vol_dict['dispersion inside'].chi(w_nodes).astype(np.complex)
try:
ionizer = ionization.Ionizer(vol_dict['ionizer'])
except KeyError:
ionizer = None
mat = Material(cl.q, field_tool, N, vol_dict['wave coordinates'],
vol_dict['object'], chi, chi3,
vol_dict['density reference'], window_speed)
# Setup the wave propagation domain
dens_nodes = grid_tools.cell_centers(-size[3] / 2, size[3] / 2, steps)
field_walls = grid_tools.cell_walls(dens_nodes[0], dens_nodes[-1], steps)
diagnostic_walls = np.linspace(-size[3] / 2, size[3] / 2, N[3])
dzmin = vol_dict['minimum step']
Dz = diagnostic_walls[1] - diagnostic_walls[0]
dom4d = np.concatenate((dom3d, [field_walls[0], field_walls[-1]]))
# Step through the domain
A0 = np.copy(A[..., 0])
for k in range(diagnostic_steps):
zi = diagnostic_walls[k]
zf = diagnostic_walls[k + 1]
print('Advancing to diagnostic plane', k + 1)
A0, J0, ne0, evals = propagator(cl, field_tool, window_speed, A0, mat,
ionizer, NL_band, subcycles, zi, zf,
dzmin)
print('', evals, 'evaluations of j(w,kx,ky)')
A[..., k + 1] = A0
J[..., k + 1] = J0
ne[..., k + 1] = ne0
# Finish by relaunching rays and returning UPPE data
if full_relaunch:
field_tool.RelaunchRays1(xp, eikonal, vg, A[..., -1], size[3],
vol_dict['dispersion inside'])
else:
field_tool.RelaunchRays(xp, eikonal, vg, A[..., -1], size[3],
vol_dict['dispersion inside'])
return A, J, ne, dom4d
def test_ghost_cells(self):
a = np.zeros((4, 4, 4))
nodes = grid_tools.cell_centers(0.0, 4.0, 4)
xa, xnodes = grid_tools.AddGhostCells(a, nodes, 1)
assert np.allclose(xa, np.zeros((4, 5, 4)))
assert np.allclose(xnodes, [-0.5, 0.5, 1.5, 2.5, 3.5])
def test_normal(self):
nodes = grid_tools.cell_centers(0.0, 5.0, 5)
assert np.allclose(nodes, [0.5, 1.5, 2.5, 3.5, 4.5])
walls = grid_tools.cell_walls(0.5, 4.5, 5)
assert np.allclose(walls, [0.0, 1.0, 2.0, 3.0, 4.0, 5.0])
def track(cl, xp, eikonal, vg, vol_dict):
'''Propagate paraxial waves using eikonal data as a boundary condition.
The volume must be oriented so the polarization axis is x (linear polarization only).
:param numpy.array xp: ray phase space with shape (bundles,rays,8)
:param numpy.array eikonal: ray eikonal data with shape (bundles,4)
:param numpy.array vg: ray group velocity with shape (bundles,rays,4)
:param dictionary vol_dict: input file dictionary for the volume'''
band = vol_dict['frequency band']
size = (band[1] - band[0], ) + vol_dict['size']
N = vol_dict['wave grid']
# N[3] is the number of diagnostic planes, including the initial plane
diagnostic_steps = N[3] - 1
subcycles = vol_dict['subcycles']
steps = diagnostic_steps * subcycles
field_planes = steps + 1
Vol = vol_dict['object']
powersof2 = [2**i for i in range(32)]
if N[0] not in powersof2:
raise ValueError('Paraxial propagator requires 2**n w-nodes')
if N[1] not in powersof2:
raise ValueError('Paraxial propagator requires 2**n x-nodes')
if N[2] not in powersof2:
raise ValueError('Paraxial propagator requires 2**n y-nodes')
try:
window_speed = vol_dict['window speed']
except:
window_speed = 1.0
try:
chi3 = vol_dict['chi3']
except:
chi3 = 0.0
# Capture the rays
if vol_dict['wave coordinates'] == 'cartesian':
field_tool = caustic_tools.FourierTool(N, band, (0, 0, 0), size[1:],
cl)
else:
field_tool = caustic_tools.BesselBeamTool(N, band, (0, 0, 0), size[1:],
cl)
w_nodes, x1_nodes, x2_nodes, plot_ext = field_tool.GetGridInfo()
A = np.zeros(N).astype(np.complex)
A0, dom3d = field_tool.GetBoundaryFields(xp[:, 0, :], eikonal, 1)
n0 = np.mean(
np.sqrt(np.einsum('...i,...i', xp[:, 0, 5:8], xp[:, 0, 5:8])) /
xp[:, 0, 4])
ng = 1.0 / np.mean(
np.sqrt(np.einsum('...i,...i', vg[:, 0, 1:4], vg[:, 0, 1:4])))
# Setup the wave propagation domain
chi = vol_dict['dispersion inside'].chi(w_nodes)
dens_nodes = grid_tools.cell_centers(-size[3] / 2, size[3] / 2, steps)
field_walls = grid_tools.cell_walls(dens_nodes[0], dens_nodes[-1], steps)
diagnostic_walls = np.linspace(-size[3] / 2, size[3] / 2, N[3])
dz = field_walls[1] - field_walls[0]
Dz = diagnostic_walls[1] - diagnostic_walls[0]
dom4d = np.concatenate((dom3d, [field_walls[0], field_walls[-1]]))
# Step through the domain
# Strategy to get density plane is to re-use ray gather system
# This works as long as the shape of xp is (*,*,8)
xp_eff = np.zeros((N[1], N[2], 8))
if vol_dict['wave coordinates'] == 'cartesian':
xp_eff[..., 1] = np.outer(x1_nodes, np.ones(N[2]))
xp_eff[..., 2] = np.outer(np.ones(N[1]), x2_nodes)
else:
xp_eff[..., 1] = np.outer(x1_nodes, np.cos(x2_nodes))
xp_eff[..., 2] = np.outer(x1_nodes, np.sin(x2_nodes))
A[..., 0] = A0
for k in range(diagnostic_steps):
print('Advancing to diagnostic plane', k + 1)
for s in range(subcycles):
xp_eff[..., 3] = dens_nodes[k * subcycles + s]
dens = vol_dict['object'].GetDensity(xp_eff)
A0 = propagator(field_tool, A0, chi, dens, n0, ng, dz, True)
try:
A0 *= vol_dict['damping filter'](w_nodes)[:, np.newaxis,
np.newaxis]
except KeyError:
A0 = A0
A[..., k + 1] = A0
# Finish by relaunching rays and returning wave data
field_tool.RelaunchRays(xp, eikonal, vg, A[..., -1], size[3])
return A, dom4d
