def Initialize(self, input_dict):
super().Initialize(input_dict)
km = 2 * np.pi / input_dict['inner scale']
k0 = 2 * np.pi / input_dict['outer scale']
Ns = input_dict['grid']
ds = (self.Lx / (Ns[0] - 1), self.Ly / (Ns[1] - 1)
) # make it so outer nodes lie on boundary
kc = (np.pi / ds[0], np.pi / ds[1])
kx = np.fft.ifftshift(grid_tools.cyclic_nodes(-kc[0], kc[0], Ns[0]))
ky = np.fft.ifftshift(grid_tools.cyclic_nodes(-kc[1], kc[1], Ns[1]))
k2 = np.outer(kx**2, np.ones(Ns[1])) + np.outer(np.ones(Ns[0]), ky**2)
screen = np.exp(-k2 / km**2) * (0j + k2 + k0**2)**(-11 / 6)
pos = lambda i: slice(1, int(Ns[i] / 2))
neg = lambda i: slice(int(Ns[i] / 2) + 1, Ns[i])
crit = lambda i: int(Ns[i] / 2)
screen *= np.random.random(Ns) - 0.5 + 1j * np.random.random(Ns) - 0.5j
screen[neg(0), neg(1)] = np.conj(screen[pos(0), pos(1)][::-1, ::-1])
screen[neg(0), pos(1)] = np.conj(screen[pos(0), neg(1)][::-1, ::-1])
screen[0, :] = 0
screen[:, 0] = 0
screen[crit(0), :] = np.real(screen[crit(0), :])
screen[:, crit(1)] = np.real(screen[:, crit(1)])
screen = np.real(np.fft.ifft(np.fft.ifft(screen, axis=0), axis=1))
self.screen = 1 + input_dict['amplitude'] * screen / np.max(
np.abs(screen))
self.xn = np.linspace(-self.Lx / 2, self.Lx / 2, Ns[0])
self.yn = np.linspace(-self.Ly / 2, self.Ly / 2, Ns[1])
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 PulseSpectrum(xp, box, N, pulse_length, w0):
'''Weight the rays based their frequency. Assume transform limited pulse.
Spectral amplitude is matched with time domain amplitude by performing a test FFT.
If the lower frequency bound is zero, real carrier resolved fields are used.
:param numpy.ndarray xp: Ray phase space, primary ray frequency must be loaded
:param tuple box: The first two elements are the frequency bounds
:param tuple N: The first element is the number of frequency nodes
:param pulse_length: The Gaussian pulse width (1/e amplitude)
:param w0: The central frequency of the wave
:returns: weight factors with shape (bundles,)'''
sigma_w = 2.0 / pulse_length
freq_range = grid_tools.cyclic_nodes(box[0], box[1], N[0])
# Create a unit-height Gaussian spectral amplitude
envelope = np.exp(-(freq_range - w0)**2 / sigma_w**2)
if box[0] == 0.0:
# Intepret as real carrier resolved field
envelope = np.fft.irfft(envelope)
else:
# Interpret as spectral envelope putting negative frequencies last
envelope = np.fft.ifftshift(envelope)
# Get the corresponding time domain envelope
envelope = np.fft.ifft(envelope)
coeff = 1 / np.max(np.abs(envelope))
return coeff * np.exp(-(xp[:, 0, 4] - w0)**2 / sigma_w**2)
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 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 test_cyclic(self):
nodes = grid_tools.cyclic_nodes(0.0, 5.0, 5)
assert np.allclose(nodes, [0.0, 1.0, 2.0, 3.0, 4.0])
center, width = grid_tools.cyclic_center_and_width(0.0, 5.0)
assert center == pytest.approx(2.5)
assert width == pytest.approx(5)