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