def load_distribution_function(self, file_name):
"""
This function is used to load the distribution function from the
dump file that was created by dump_distribution_function.
Parameters
----------
file_name : The distribution_function array will be loaded from this
provided file name.
Examples
--------
>> solver.load_distribution_function('distribution_function')
The above statemant will load the distribution function data stored in the file
distribution_function.h5 into self.f
"""
viewer = PETSc.Viewer().createHDF5(file_name + '.h5',
PETSc.Viewer.Mode.READ)
self._glob_f.load(viewer)
self.Y[:, :, :, 0] = 2 * af.fft2(af.to_array(self._glob_f_value[:])) \
/ (self.N_q1 * self.N_q2)
return
def adjoint(self, residual, cache):
V_obj_af, field_layer_in_or_grad,\
flag_gpu_inout = cache
#back-propagte to volume center
field_bp = residual
if self.focus_at_center:
#propagate to the last layer
field_bp = self._propagationInplace(field_bp,
self.distance_end_to_center)
for layer in range(self.shape[2] - 1, -1, -1):
field_bp_1 = af.ifft2(
af.fft2(field_bp) *
af.conjg(self.green_kernel_2d)) * self.pixel_size_z
if layer > 0:
field_bp = self._propagationInplace(
field_bp, self.slice_separation[layer], adjoint=True)
field_bp[:, :] += field_bp_1 * af.conjg(V_obj_af[:, :, layer])
field_layer_in_or_grad[:, :, layer] = field_bp_1 * af.conjg(
field_layer_in_or_grad[:, :, layer])
#Unbinning
grad = self._meanObject(field_layer_in_or_grad, adjoint=True)
if flag_gpu_inout:
return {'gradient': grad}
else:
return {'gradient': np.array(grad)}
def adjoint(self, field):
"""Adjoint operator for pupil (and estimate pupil if selected)"""
field_adj_f = af.fft2(field)
field_pupil_adj = af.ifft2(af.conjg(self.getPupil()) * field_adj_f)
#Pupil recovery
if self.flag_update:
self._update(field_adj_f)
return field_pupil_adj
def forward(self, field, **kwargs):
"""Apply pupil"""
self.field_f = af.fft2(field)
if self.flag_update:
if self.update_method == "GaussNewton":
self.approx_hessian[:, :] += self.field_f * af.conjg(
self.field_f)
field = af.ifft2(self.getPupil() * self.field_f)
return field
def af_padded_fft2(img, pad_val=0., pad_periods=0):
#side = img.dims()[0]
#padded_img = af.constant(0., (1+pad_periods)*side, (1+pad_periods)*side)
#padded_img[:side, :side] = img
#fft = af.fft2(padded_img)
fft = af.fft2(img)
return fft
def fft2(array):
# Reorder from (N_p, N_s, N_q1, N_q2) --> (N_q1, N_q2, N_p, N_s)
array = af.reorder(array, 2, 3, 0, 1)
array = af.fft2(array)
# Reorder back from (N_q1, N_q2, N_p, N_s) --> (N_p, N_s, N_q1, N_q2)
array = af.reorder(array, 2, 3, 0, 1)
af.eval(array)
return(array)
def compute_electrostatic_fields(self, f=None, f_hat=None):
"""
Computes the electrostatic fields by making use of the
Poisson equation: div^2 phi = rho
"""
if (self.single_mode_evolution == True):
# Intializing for the electrostatic Case:
delta_rho_hat = self.compute_moments('density', f)
delta_phi_hat = self.physical_system.params.charge_electron \
* delta_rho_hat/( self.physical_system.params.k_q1**2
+ self.physical_system.params.k_q2**2
)
self.delta_E1_hat = -delta_phi_hat * (1j *
self.physical_system.params.k_q1)
self.delta_E2_hat = -delta_phi_hat * (1j *
self.physical_system.params.k_q2)
self.delta_E3_hat = 0
self.delta_B1_hat = 0
self.delta_B2_hat = 0
self.delta_B3_hat = 0
else:
rho_hat = 2 * af.fft2(self.compute_moments('density', f=f, f_hat=f_hat)) \
/ (self.N_q1 * self.N_q2) # Scaling Appropriately
# Defining lambda functions to perform broadcasting operations:
# This is done using af.broadcast, which allows us to perform
# batched operations when operating on arrays of different sizes:
divide = lambda a, b: a / b
multiply = lambda a, b: a * b
# af.broadcast(function, *args) performs batched operations on
# function(*args):
phi_hat = self.physical_system.params.charge_electron * \
af.broadcast(divide, rho_hat,
(self.k_q1**2 + self.k_q2**2)
)
# Setting the background electric potential to zero:
phi_hat[0, 0] = 0
# Tiling to make phi_hat of positionsExpanded form:
phi_hat = af.tile(phi_hat, 1, 1, self.N_p1 * self.N_p2 * self.N_p3)
self.E1_hat = af.broadcast(multiply, -phi_hat, (1j * self.k_q1))
self.E2_hat = af.broadcast(multiply, -phi_hat, (1j * self.k_q2))
af.eval(self.E1_hat, self.E2_hat)
return
def forward(self, field, propagation_distances):
"""defocus with angular spectrum"""
field_defocus = self.pupil_support * af.fft2(field)
field_defocus = af.tile(field_defocus, 1, 1,
len(propagation_distances))
for z_idx, propagation_distance in enumerate(propagation_distances):
propagation_kernel = af.exp(self.prop_kernel_phase *
propagation_distance)
field_defocus[:, :, z_idx] *= propagation_kernel
af.ifft2_inplace(field_defocus)
return field_defocus
def fft_poisson(self, f=None):
"""
Solves the Poisson Equation using the FFTs:
Used as a backup solver in case of low resolution runs
(ie. used on a single node) with periodic boundary
conditions.
"""
if (self.performance_test_flag == True):
tic = af.time()
if (self._comm.size != 1):
raise Exception('FFT solver can only be used when run in serial')
else:
N_g = self.N_ghost
rho = af.reorder( self.physical_system.params.charge_electron \
* self.compute_moments('density', f)[:, N_g:-N_g, N_g:-N_g],
1, 2, 0
)
k_q1 = fftfreq(rho.shape[0], self.dq1)
k_q2 = fftfreq(rho.shape[1], self.dq2)
k_q2, k_q1 = np.meshgrid(k_q2, k_q1)
k_q1 = af.to_array(k_q1)
k_q2 = af.to_array(k_q2)
rho_hat = af.fft2(rho)
potential_hat = rho_hat / (4 * np.pi**2 * (k_q1**2 + k_q2**2))
potential_hat[0, 0] = 0
E1_hat = -1j * 2 * np.pi * k_q1 * potential_hat
E2_hat = -1j * 2 * np.pi * k_q2 * potential_hat
# Non-inclusive of ghost-zones:
E1_physical = af.reorder(af.real(af.ifft2(E1_hat)), 2, 0, 1)
E2_physical = af.reorder(af.real(af.ifft2(E2_hat)), 2, 0, 1)
self.cell_centered_EM_fields[0, N_g:-N_g, N_g:-N_g] = E1_physical
self.cell_centered_EM_fields[1, N_g:-N_g, N_g:-N_g] = E2_physical
af.eval(self.cell_centered_EM_fields)
if (self.performance_test_flag == True):
af.sync()
toc = af.time()
self.time_fieldsolver += toc - tic
return
def adjoint(self, residual, propagation_distances):
"""adjoint operator for defocus with angular spectrum"""
field_focus = af.constant(0.0,
self.shape[0],
self.shape[1],
dtype=af_complex_datatype)
for z_idx, propagation_distance in enumerate(propagation_distances):
propagation_kernel_conj = af.exp(-1.0 * self.prop_kernel_phase *
propagation_distances[z_idx])
field_focus[:, :] += propagation_kernel_conj * af.fft2(
residual[:, :, z_idx])
field_focus *= self.pupil_support
af.ifft2_inplace(field_focus)
return field_focus
def __init__(self):
self.N_q1 = 32
self.N_q2 = 32
self.N_p1 = 4
self.N_p2 = 5
self.N_p3 = 6
self.p1_start = self.p2_start = self.p3_start = 0
self.dp1 = 2 / self.N_p1
self.dp2 = 2 / self.N_p2
self.dp3 = 2 / self.N_p3
self.p1, self.p2, self.p3 = self._calculate_p_center()
# Creating an object with required attributes for the test:
self.physical_system = type('obj', (object, ), {
'moment_exponents': moment_exponents,
'moment_coeffs': moment_coeffs
})
self.single_mode_evolution = False
self.f = af.randu(self.N_q1,
self.N_q2,
self.N_p1 * self.N_p2 * self.N_p3,
dtype=af.Dtype.f64)
self.Y = 2 * af.fft2(self.f) / (self.N_q1 * self.N_q2)
self._da_dump_f = PETSc.DMDA().create(
[self.N_q1, self.N_q2],
dof=(self.N_p1 * self.N_p2 * self.N_p3),
)
self._da_dump_moments = PETSc.DMDA().create([self.N_q1, self.N_q2],
dof = len(self.physical_system.\
moment_exponents)
)
self._glob_f = self._da_dump_f.createGlobalVec()
self._glob_f_value = self._da_dump_f.getVecArray(self._glob_f)
self._glob_moments = self._da_dump_moments.createGlobalVec()
self._glob_moments_value = self._da_dump_moments.\
getVecArray(self._glob_moments)
PETSc.Object.setName(self._glob_f, 'distribution_function')
PETSc.Object.setName(self._glob_moments, 'moments')
def RK5_step(self, dt):
self.Y = integrators.RK5(dY_dt, self.Y, dt, self)
# Solving for tau = 0 systems
if (self.single_mode_evolution == False and af.any_true(
self.physical_system.params.tau(self.q1_center, self.q2_center,
self.p1, self.p2, self.p3) == 0)):
f_hat = self.Y[:, :, :, 0]
f = af.real(af.ifft2(0.5 * self.N_q2 * self.N_q1 * f_hat))
self.Y[:, :, :, 0] = 2 * af.fft2(
self._source(f, self.q1_center, self.q2_center, self.p1, self.p2,
self.p3, self.compute_moments,
self.physical_system.params,
True)) / (self.N_q2 * self.N_q1)
return
def test_1V():
obj = test()
obj.single_mode_evolution = False
f_generalized = MB_dist(obj.q1_center, obj.q2_center, obj.p1, obj.p2,
obj.p3, 1)
C_f_hat_generalized = 2 * af.fft2(
collision_operator.BGK(
f_generalized, obj.q1_center, obj.q2_center, obj.p1, obj.p2,
obj.p3, obj.compute_moments,
obj.physical_system.params)) / (obj.N_q2 * obj.N_q1)
# Background
C_f_hat_generalized[0, 0, :] = 0
# Finding the indices of the mode excited:
i_q1_max = np.unravel_index(
af.imax(af.abs(C_f_hat_generalized))[1],
(obj.N_q1, obj.N_q2, obj.N_p1 * obj.N_p2 * obj.N_p3),
order='F')[0]
i_q2_max = np.unravel_index(
af.imax(af.abs(C_f_hat_generalized))[1],
(obj.N_q1, obj.N_q2, obj.N_p1 * obj.N_p2 * obj.N_p3),
order='F')[1]
obj.p1 = np.array(af.reorder(obj.p1, 1, 2, 3, 0))
obj.p2 = np.array(af.reorder(obj.p2, 1, 2, 3, 0))
obj.p3 = np.array(af.reorder(obj.p3, 1, 2, 3, 0))
delta_f_hat = 0.01 * (1 / (2 * np.pi))**(1 / 2) \
* np.exp(-0.5 * obj.p1**2)
obj.single_mode_evolution = True
C_f_hat_single_mode = collision_operator.linearized_BGK(
delta_f_hat, obj.p1, obj.p2, obj.p3, obj.compute_moments,
obj.physical_system.params)
assert (af.mean(
af.abs(
af.flat(C_f_hat_generalized[i_q1_max, i_q2_max]) -
af.to_array(C_f_hat_single_mode.flatten()))) < 1e-14)
def __fftn__(a, s, axes, direction='forward'):
if len(s) != 3 and len(s) != 2 and len(s) != 1:
raise NotImplementedError
if axes is not None:
raise NotImplementedError
if(direction == 'forward'):
if len(s) == 3:
fa = arrayfire.fft3(a.d_array, s[2], s[1], s[0])
elif len(s) == 2:
fa = arrayfire.fft2(a.d_array, s[1], s[0])
elif len(s) == 1:
fa = arrayfire.fft(a.d_array, s[0])
elif direction == 'inverse':
if len(s) == 3:
fa = arrayfire.ifft3(a.d_array, s[2], s[1], s[0])
elif len(s) == 2:
fa = arrayfire.ifft2(a.d_array, s[1], s[0])
elif len(s) == 1:
fa = arrayfire.ifft(a.d_array, s[0])
else:
raise ValueError('Wrong FFT direction')
return ndarray(pu.af_shape(fa), dtype=pu.typemap(fa.dtype()), af_array=fa)
def __fftn__(a, s, axes, direction='forward'):
if len(s) != 3 and len(s) != 2 and len(s) != 1:
raise NotImplementedError
if axes is not None:
raise NotImplementedError
if(direction == 'forward'):
if len(s) == 3:
fa = arrayfire.fft3(a.d_array, s[2], s[1], s[0])
elif len(s) == 2:
fa = arrayfire.fft2(a.d_array, s[1], s[0])
elif len(s) == 1:
fa = arrayfire.fft(a.d_array, s[0])
elif direction == 'inverse':
if len(s) == 3:
fa = arrayfire.ifft3(a.d_array, s[2], s[1], s[0])
elif len(s) == 2:
fa = arrayfire.ifft2(a.d_array, s[1], s[0])
elif len(s) == 1:
fa = arrayfire.ifft(a.d_array, s[0])
else:
raise ValueError('Wrong FFT direction')
return ndarray(a.shape, dtype=pu.typemap(fa.dtype()), af_array=fa)
def forward(self, V_obj, fx_illu, fy_illu):
if (type(V_obj).__module__ == np.__name__):
V_obj_af = af.to_array(V_obj)
flag_gpu_inout = False
else:
V_obj_af = V_obj
flag_gpu_inout = True
#Binning
obj_af = self._meanObject(V_obj_af)
#compute illumination
field, fx_illu, fy_illu, fz_illu = self._genIllumination(
fx_illu, fy_illu)
field[:, :] *= np.exp(1.0j * 2.0 * np.pi * fz_illu *
self.initial_z_position)
field_layer_in = af.constant(0.0,
self.shape[0],
self.shape[1],
self.shape[2],
dtype=af_complex_datatype)
for layer in range(self.shape[2]):
field_layer_in[:, :, layer] = field
field = self._propagationInplace(field,
self.slice_separation[layer])
field_scat = af.ifft2(af.fft2(field_layer_in[:, :, layer] * obj_af[:, :, layer])*\
self.green_kernel_2d) * self.pixel_size_z
field[:, :] += field_scat
#store intermediate variables for adjoint operation
cache = (obj_af, field_layer_in, flag_gpu_inout)
if self.focus_at_center:
#propagate to volume center
field = self._propagationInplace(field,
self.distance_end_to_center,
adjoint=True)
return {'forward_scattered_field': field, 'cache': cache}
def bench(A, iters=100):
start = time()
for t in range(iters):
B = af.fft2(A)
af.sync()
return (time() - start) / iters
def fft_poisson(self, rho):
"""
Solves the Poisson Equation using FFTs:
Used as a backup solver for low resolution runs
(ie. used on a single node) with periodic boundary
conditions.
Parameters
----------
rho : af.Array
Array that holds the charge density for each species
"""
if(self.performance_test_flag == True):
tic = af.time()
if (self._comm.size != 1):
raise Exception('FFT solver can only be used when run in serial')
else:
N_g = self.N_g
# Reorder from (1, N_s, N_q1, N_q2) --> (N_q1, N_q2, 1, N_s)
rho = af.reorder(rho[:, :, N_g:-N_g, N_g:-N_g] , 2, 3, 0, 1)
# Summing for all the species:
rho = af.sum(rho, 3)
k_q1 = np.fft.fftfreq(rho.shape[0], self.dq1)
k_q2 = np.fft.fftfreq(rho.shape[1], self.dq2)
k_q2, k_q1 = np.meshgrid(k_q2, k_q1)
k_q1 = af.to_array(k_q1)
k_q2 = af.to_array(k_q2)
rho_hat = af.fft2(rho)
potential_hat = rho_hat / (4 * np.pi**2 * (k_q1**2 + k_q2**2))
potential_hat[0, 0] = 0
E1_hat = -1j * 2 * np.pi * k_q1 * potential_hat
E2_hat = -1j * 2 * np.pi * k_q2 * potential_hat
# Non-inclusive of ghost-zones:
E1_ifft = af.ifft2(E1_hat, scale=1)/(E1_hat.shape[0] * E1_hat.shape[1])
E2_ifft = af.ifft2(E2_hat, scale=1)/(E2_hat.shape[0] * E2_hat.shape[1])
E1_physical = af.reorder(af.real(E1_ifft), 2, 3, 0, 1)
E2_physical = af.reorder(af.real(E2_ifft), 2, 3, 0, 1)
self.cell_centered_EM_fields[0, 0, N_g:-N_g, N_g:-N_g] = E1_physical
self.cell_centered_EM_fields[1, 0, N_g:-N_g, N_g:-N_g] = E2_physical
af.eval(self.cell_centered_EM_fields)
if(self.performance_test_flag == True):
af.sync()
toc = af.time()
self.time_fieldsolver += toc - tic
return
def simple_signal(verbose=False):
display_func = _util.display_func(verbose)
print_func = _util.print_func(verbose)
a = af.randu(10, 1)
pos0 = af.randu(10) * 10
display_func(af.approx1(a, pos0))
a = af.randu(3, 3)
pos0 = af.randu(3, 3) * 10
pos1 = af.randu(3, 3) * 10
display_func(af.approx2(a, pos0, pos1))
a = af.randu(8, 1)
display_func(a)
display_func(af.fft(a))
display_func(af.dft(a))
display_func(af.real(af.ifft(af.fft(a))))
display_func(af.real(af.idft(af.dft(a))))
a = af.randu(4, 4)
display_func(a)
display_func(af.fft2(a))
display_func(af.dft(a))
display_func(af.real(af.ifft2(af.fft2(a))))
display_func(af.real(af.idft(af.dft(a))))
a = af.randu(4, 4, 2)
display_func(a)
display_func(af.fft3(a))
display_func(af.dft(a))
display_func(af.real(af.ifft3(af.fft3(a))))
display_func(af.real(af.idft(af.dft(a))))
a = af.randu(10, 1)
b = af.randu(3, 1)
display_func(af.convolve1(a, b))
display_func(af.fft_convolve1(a, b))
display_func(af.convolve(a, b))
display_func(af.fft_convolve(a, b))
a = af.randu(5, 5)
b = af.randu(3, 3)
display_func(af.convolve2(a, b))
display_func(af.fft_convolve2(a, b))
display_func(af.convolve(a, b))
display_func(af.fft_convolve(a, b))
a = af.randu(5, 5, 3)
b = af.randu(3, 3, 2)
display_func(af.convolve3(a, b))
display_func(af.fft_convolve3(a, b))
display_func(af.convolve(a, b))
display_func(af.fft_convolve(a, b))
b = af.randu(3, 1)
x = af.randu(10, 1)
a = af.randu(2, 1)
display_func(af.fir(b, x))
display_func(af.iir(b, a, x))
def run(iters):
for t in range(iters):
B = af.fft2(A)
af.sync()
def __init__(self):
# Creating a physical_system object with
# attributes required for the test
self.physical_system = type('obj', (object, ), {
'moment_exponents': moment_exponents,
'moment_coeffs': moment_coeffs
})
self.single_mode_evolution = False
self.p1_start = -10
self.p2_start = -10
self.p3_start = -10
self.N_p1 = 32
self.N_p2 = 32
self.N_p3 = 32
self.dp1 = (-2 * self.p1_start) / self.N_p1
self.dp2 = (-2 * self.p2_start) / self.N_p2
self.dp3 = (-2 * self.p3_start) / self.N_p3
self.N_q1 = 16
self.N_q2 = 16
self.N_ghost = 3
self.p1 = self.p1_start + (0.5 + np.arange(self.N_p1)) * self.dp1
self.p2 = self.p2_start + (0.5 + np.arange(self.N_p2)) * self.dp2
self.p3 = self.p3_start + (0.5 + np.arange(self.N_p3)) * self.dp3
self.p2, self.p1, self.p3 = np.meshgrid(self.p2, self.p1, self.p3)
self.p1 = af.reorder(af.flat(af.to_array(self.p1)), 2, 3, 0, 1)
self.p2 = af.reorder(af.flat(af.to_array(self.p2)), 2, 3, 0, 1)
self.p3 = af.reorder(af.flat(af.to_array(self.p3)), 2, 3, 0, 1)
self.q1 = (-self.N_ghost + 0.5 +
np.arange(self.N_q1 + 2 * self.N_ghost)) / self.N_q1
self.q2 = (-self.N_ghost + 0.5 +
np.arange(self.N_q2 + 2 * self.N_ghost)) / self.N_q2
self.q2, self.q1 = np.meshgrid(self.q2, self.q1)
self.q1 = af.to_array(self.q1)
self.q2 = af.to_array(self.q2)
rho = (1 + 0.01 * af.sin(2 * np.pi * self.q1 + 4 * np.pi * self.q2))
T = (1 + 0.01 * af.cos(2 * np.pi * self.q1 + 4 * np.pi * self.q2))
p1_b = 0.01 * af.exp(-10 * self.q1**2 - 10 * self.q2**2)
p2_b = 0.01 * af.exp(-10 * self.q1**2 - 10 * self.q2**2)
p3_b = 0.01 * af.exp(-10 * self.q1**2 - 10 * self.q2**2)
self.f = maxwell_boltzmann(rho, T, p1_b, p2_b, p3_b, self.p1, self.p2,
self.p3)
self.Y = 2 * af.fft2(self.f) / (self.N_q1 * self.N_q2)
def fft_poisson(self, rho):
"""
Solves the Poisson Equation using FFTs:
Used as a backup solver for low resolution runs
(ie. used on a single node) with periodic boundary
conditions.
Parameters
----------
rho : af.Array
Array that holds the charge density for each species
"""
if (self.performance_test_flag == True):
tic = af.time()
if (self._comm.size != 1):
raise Exception('FFT solver can only be used when run in serial')
else:
N_g = self.N_g
# Reorder from (1, 1, N_q1, N_q2) --> (N_q1, N_q2, 1, 1)
rho = af.reorder(rho[:, :, N_g:-N_g, N_g:-N_g], 2, 3, 0, 1)
k_q1 = np.fft.fftfreq(rho.shape[0], self.dq1)
k_q2 = np.fft.fftfreq(rho.shape[1], self.dq2)
k_q2, k_q1 = np.meshgrid(k_q2, k_q1)
k_q1 = af.to_array(k_q1)
k_q2 = af.to_array(k_q2)
# TEMP: Removing background for the time being to not break existing examples:
rho = rho - af.mean(rho)
rho_hat = af.fft2(rho)
# Ensuring that there is no background charge density:
# try:
# assert(af.sum(af.abs(rho_hat[0, 0]) < 1))
# except:
# raise Exception('Cannot pass rho with a background charge density \
# when solving using periodic boundary conditions'
# )
potential_hat = rho_hat / (4 * np.pi**2 * (k_q1**2 + k_q2**2))
potential_hat[0, 0] = 0
E1_hat = -1j * 2 * np.pi * k_q1 * potential_hat
E2_hat = -1j * 2 * np.pi * k_q2 * potential_hat
# Non-inclusive of ghost-zones:
E1_ifft = af.ifft2(E1_hat,
scale=1) / (E1_hat.shape[0] * E1_hat.shape[1])
E2_ifft = af.ifft2(E2_hat,
scale=1) / (E2_hat.shape[0] * E2_hat.shape[1])
E1_physical = af.reorder(af.real(E1_ifft), 2, 3, 0, 1)
E2_physical = af.reorder(af.real(E2_ifft), 2, 3, 0, 1)
self.cell_centered_EM_fields[0, 0, N_g:-N_g,
N_g:-N_g] = E1_physical / self.params.eps
self.cell_centered_EM_fields[1, 0, N_g:-N_g,
N_g:-N_g] = E2_physical / self.params.eps
af.eval(self.cell_centered_EM_fields)
if (self.performance_test_flag == True):
af.sync()
toc = af.time()
self.time_fieldsolver += toc - tic
return
pos1 = af.randu(3, 3) * 10 af.display(af.approx2(a, pos0, pos1)) a = af.randu(8, 1) af.display(a) af.display(af.fft(a)) af.display(af.dft(a)) af.display(af.real(af.ifft(af.fft(a)))) af.display(af.real(af.idft(af.dft(a)))) a = af.randu(4, 4) af.display(a) af.display(af.fft2(a)) af.display(af.dft(a)) af.display(af.real(af.ifft2(af.fft2(a)))) af.display(af.real(af.idft(af.dft(a)))) a = af.randu(4, 4, 2) af.display(a) af.display(af.fft3(a)) af.display(af.dft(a)) af.display(af.real(af.ifft3(af.fft3(a)))) af.display(af.real(af.idft(af.dft(a)))) a = af.randu(10, 1) b = af.randu(3, 1) af.display(af.convolve1(a, b))
def simple_signal(verbose=False):
display_func = _util.display_func(verbose)
print_func = _util.print_func(verbose)
signal = af.randu(10)
x_new = af.randu(10)
x_orig = af.randu(10)
display_func(af.approx1(signal, x_new, xp = x_orig))
signal = af.randu(3, 3)
x_new = af.randu(3, 3)
x_orig = af.randu(3, 3)
y_new = af.randu(3, 3)
y_orig = af.randu(3, 3)
display_func(af.approx2(signal, x_new, y_new, xp = x_orig, yp = y_orig))
a = af.randu(8, 1)
display_func(a)
display_func(af.fft(a))
display_func(af.dft(a))
display_func(af.real(af.ifft(af.fft(a))))
display_func(af.real(af.idft(af.dft(a))))
b = af.fft(a)
af.ifft_inplace(b)
display_func(b)
af.fft_inplace(b)
display_func(b)
b = af.fft_r2c(a)
c = af.fft_c2r(b)
display_func(b)
display_func(c)
a = af.randu(4, 4)
display_func(a)
display_func(af.fft2(a))
display_func(af.dft(a))
display_func(af.real(af.ifft2(af.fft2(a))))
display_func(af.real(af.idft(af.dft(a))))
b = af.fft2(a)
af.ifft2_inplace(b)
display_func(b)
af.fft2_inplace(b)
display_func(b)
b = af.fft2_r2c(a)
c = af.fft2_c2r(b)
display_func(b)
display_func(c)
a = af.randu(4, 4, 2)
display_func(a)
display_func(af.fft3(a))
display_func(af.dft(a))
display_func(af.real(af.ifft3(af.fft3(a))))
display_func(af.real(af.idft(af.dft(a))))
b = af.fft3(a)
af.ifft3_inplace(b)
display_func(b)
af.fft3_inplace(b)
display_func(b)
b = af.fft3_r2c(a)
c = af.fft3_c2r(b)
display_func(b)
display_func(c)
a = af.randu(10, 1)
b = af.randu(3, 1)
display_func(af.convolve1(a, b))
display_func(af.fft_convolve1(a, b))
display_func(af.convolve(a, b))
display_func(af.fft_convolve(a, b))
a = af.randu(5, 5)
b = af.randu(3, 3)
display_func(af.convolve2(a, b))
display_func(af.fft_convolve2(a, b))
display_func(af.convolve(a, b))
display_func(af.fft_convolve(a, b))
a = af.randu(5, 5, 3)
b = af.randu(3, 3, 2)
display_func(af.convolve3(a, b))
display_func(af.fft_convolve3(a, b))
display_func(af.convolve(a, b))
display_func(af.fft_convolve(a, b))
b = af.randu(3, 1)
x = af.randu(10, 1)
a = af.randu(2, 1)
display_func(af.fir(b, x))
display_func(af.iir(b, a, x))
display_func(af.medfilt1(a))
display_func(af.medfilt2(a))
display_func(af.medfilt(a))
def dY_dt_multimode_evolution(Y, self):
"""
Returns the value of the derivative of the fourier mode quantities
of the distribution function, and the field quantities with
respect to time. This is used to evolve the system in time.
Input:
------
Y : The array Y is the state of the system as given by the result of
the last time-step's integration. The elements of Y, hold the
following data:
f_hat = Y[:, :, :, 0]
E1_hat = Y[:, :, :, 1]
E2_hat = Y[:, :, :, 2]
E3_hat = Y[:, :, :, 3]
B1_hat = Y[:, :, :, 4]
B2_hat = Y[:, :, :, 5]
B3_hat = Y[:, :, :, 6]
At t = 0 the initial state of the system is passed to this function:
Output:
-------
dY_dt : The time-derivatives of all the quantities stored in Y
"""
f_hat = Y[:, :, :, 0]
self.E1_hat = Y[:, :, :, 1]
self.E2_hat = Y[:, :, :, 2]
self.E3_hat = Y[:, :, :, 3]
self.B1_hat = Y[:, :, :, 4]
self.B2_hat = Y[:, :, :, 5]
self.B3_hat = Y[:, :, :, 6]
# Scaling Appropriately:
f = af.real(af.ifft2(0.5 * self.N_q2 * self.N_q1 * f_hat))
C_f_hat = 2 * af.fft2(self._source(f, self.q1_center, self.q2_center,
self.p1, self.p2, self.p3,
self.compute_moments,
self.physical_system.params
)
)/(self.N_q2 * self.N_q1)
if( self.physical_system.params.fields_solver == 'electrostatic'
or self.physical_system.params.fields_solver == 'fft'
):
compute_electrostatic_fields(self, f_hat=f_hat)
# When method is FDTD, this function returns the timederivatives
# of the field quantities which is stepped using a numerical integrator:
elif(self.physical_system.params.fields_solver == 'fdtd'):
pass
else:
raise NotImplementedError('Method invalid/not-implemented')
mom_bulk_p1 = self.compute_moments('mom_p1_bulk', f_hat=f_hat)
mom_bulk_p2 = self.compute_moments('mom_p2_bulk', f_hat=f_hat)
mom_bulk_p3 = self.compute_moments('mom_p3_bulk', f_hat=f_hat)
J1_hat = 2 * af.fft2( self.physical_system.params.charge_electron
* mom_bulk_p1
)/(self.N_q1 * self.N_q2)
J2_hat = 2 * af.fft2( self.physical_system.params.charge_electron
* mom_bulk_p2
)/(self.N_q1 * self.N_q2)
J3_hat = 2 * af.fft2( self.physical_system.params.charge_electron
* mom_bulk_p3
)/(self.N_q1 * self.N_q2)
# Defining lambda functions to perform broadcasting operations:
# This is done using af.broadcast, which allows us to perform
# batched operations when operating on arrays of different sizes:
multiply = lambda a,b:a * b
addition = lambda a,b:a + b
# af.broadcast(function, *args) performs batched operations on
# function(*args):
dE1_hat_dt = af.broadcast(addition,
af.broadcast(multiply, self.B3_hat, 1j * self.k_q2),
- J1_hat
)
dE2_hat_dt = af.broadcast(addition,
af.broadcast(multiply,-self.B3_hat, 1j * self.k_q1),
- J2_hat
)
dE3_hat_dt = af.broadcast(addition,
af.broadcast(multiply, self.B2_hat, 1j * self.k_q1)
- af.broadcast(multiply, self.B1_hat, 1j * self.k_q2),
- J3_hat
)
dB1_hat_dt = af.broadcast(multiply, -self.E3_hat, 1j * self.k_q2)
dB2_hat_dt = af.broadcast(multiply, self.E3_hat, 1j * self.k_q1)
dB3_hat_dt = af.broadcast(multiply, self.E1_hat, 1j * self.k_q2) \
- af.broadcast(multiply, self.E2_hat, 1j * self.k_q1)
(A_p1, A_p2, A_p3) = af.broadcast(self._A_p, self.q1_center, self.q2_center,
self.p1, self.p2, self.p3,
self.E1_hat, self.E2_hat, self.E3_hat,
self.B1_hat, self.B2_hat, self.B3_hat,
self.physical_system.params
)
df_hat_dt = -1j * ( af.broadcast(multiply, self.k_q1, self._A_q1)
+ af.broadcast(multiply, self.k_q2, self._A_q2)
) * f_hat
# Adding the fields term only when charge is non-zero
if(self.physical_system.params.charge_electron != 0):
fields_term = af.broadcast(multiply, A_p1, self.dfdp1_background) \
+ af.broadcast(multiply, A_p2, self.dfdp2_background) \
+ af.broadcast(multiply, A_p3, self.dfdp3_background)
df_hat_dt -= fields_term
# Avoiding addition of the collisional term when tau != inf
tau = self.physical_system.params.tau(self.q1_center,
self.q2_center,
self.p1, self.p2, self.p3
)
df_hat_dt += af.select(tau != np.inf,\
C_f_hat,\
0
)
# Obtaining the dY_dt vector by joining the derivative quantities of
# the individual distribution function and field modes:
dY_dt = af.join(3, af.join(3, df_hat_dt, dE1_hat_dt, dE2_hat_dt, dE3_hat_dt),
dB1_hat_dt, dB2_hat_dt, dB3_hat_dt
)
af.eval(dY_dt)
return(dY_dt)
def fft(arr):
return af.fft2(arr)
def run(iters):
for t in range(iters):
B = af.fft2(A)
af.sync()
def simple_signal(verbose=False):
display_func = _util.display_func(verbose)
signal = af.randu(10)
x_new = af.randu(10)
x_orig = af.randu(10)
display_func(af.approx1(signal, x_new, xp=x_orig))
signal = af.randu(3, 3)
x_new = af.randu(3, 3)
x_orig = af.randu(3, 3)
y_new = af.randu(3, 3)
y_orig = af.randu(3, 3)
display_func(af.approx2(signal, x_new, y_new, xp=x_orig, yp=y_orig))
a = af.randu(8, 1)
display_func(a)
display_func(af.fft(a))
display_func(af.dft(a))
display_func(af.real(af.ifft(af.fft(a))))
display_func(af.real(af.idft(af.dft(a))))
b = af.fft(a)
af.ifft_inplace(b)
display_func(b)
af.fft_inplace(b)
display_func(b)
b = af.fft_r2c(a)
c = af.fft_c2r(b)
display_func(b)
display_func(c)
a = af.randu(4, 4)
display_func(a)
display_func(af.fft2(a))
display_func(af.dft(a))
display_func(af.real(af.ifft2(af.fft2(a))))
display_func(af.real(af.idft(af.dft(a))))
b = af.fft2(a)
af.ifft2_inplace(b)
display_func(b)
af.fft2_inplace(b)
display_func(b)
b = af.fft2_r2c(a)
c = af.fft2_c2r(b)
display_func(b)
display_func(c)
a = af.randu(4, 4, 2)
display_func(a)
display_func(af.fft3(a))
display_func(af.dft(a))
display_func(af.real(af.ifft3(af.fft3(a))))
display_func(af.real(af.idft(af.dft(a))))
b = af.fft3(a)
af.ifft3_inplace(b)
display_func(b)
af.fft3_inplace(b)
display_func(b)
b = af.fft3_r2c(a)
c = af.fft3_c2r(b)
display_func(b)
display_func(c)
a = af.randu(10, 1)
b = af.randu(3, 1)
display_func(af.convolve1(a, b))
display_func(af.fft_convolve1(a, b))
display_func(af.convolve(a, b))
display_func(af.fft_convolve(a, b))
a = af.randu(5, 5)
b = af.randu(3, 3)
display_func(af.convolve2(a, b))
display_func(af.fft_convolve2(a, b))
display_func(af.convolve(a, b))
display_func(af.fft_convolve(a, b))
c = af.convolve2NN(a, b)
display_func(c)
in_dims = c.dims()
incoming_grad = af.constant(1, in_dims[0], in_dims[1])
g = af.convolve2GradientNN(incoming_grad, a, b, c)
display_func(g)
a = af.randu(5, 5, 3)
b = af.randu(3, 3, 2)
display_func(af.convolve3(a, b))
display_func(af.fft_convolve3(a, b))
display_func(af.convolve(a, b))
display_func(af.fft_convolve(a, b))
b = af.randu(3, 1)
x = af.randu(10, 1)
a = af.randu(2, 1)
display_func(af.fir(b, x))
display_func(af.iir(b, a, x))
display_func(af.medfilt1(a))
display_func(af.medfilt2(a))
display_func(af.medfilt(a))
def simple_signal(verbose=False):
display_func = _util.display_func(verbose)
print_func = _util.print_func(verbose)
a = af.randu(10, 1)
pos0 = af.randu(10) * 10
display_func(af.approx1(a, pos0))
a = af.randu(3, 3)
pos0 = af.randu(3, 3) * 10
pos1 = af.randu(3, 3) * 10
display_func(af.approx2(a, pos0, pos1))
a = af.randu(8, 1)
display_func(a)
display_func(af.fft(a))
display_func(af.dft(a))
display_func(af.real(af.ifft(af.fft(a))))
display_func(af.real(af.idft(af.dft(a))))
a = af.randu(4, 4)
display_func(a)
display_func(af.fft2(a))
display_func(af.dft(a))
display_func(af.real(af.ifft2(af.fft2(a))))
display_func(af.real(af.idft(af.dft(a))))
a = af.randu(4, 4, 2)
display_func(a)
display_func(af.fft3(a))
display_func(af.dft(a))
display_func(af.real(af.ifft3(af.fft3(a))))
display_func(af.real(af.idft(af.dft(a))))
a = af.randu(10, 1)
b = af.randu(3, 1)
display_func(af.convolve1(a, b))
display_func(af.fft_convolve1(a, b))
display_func(af.convolve(a, b))
display_func(af.fft_convolve(a, b))
a = af.randu(5, 5)
b = af.randu(3, 3)
display_func(af.convolve2(a, b))
display_func(af.fft_convolve2(a, b))
display_func(af.convolve(a, b))
display_func(af.fft_convolve(a, b))
a = af.randu(5, 5, 3)
b = af.randu(3, 3, 2)
display_func(af.convolve3(a, b))
display_func(af.fft_convolve3(a, b))
display_func(af.convolve(a, b))
display_func(af.fft_convolve(a, b))
b = af.randu(3, 1)
x = af.randu(10, 1)
a = af.randu(2, 1)
display_func(af.fir(b, x))
display_func(af.iir(b, a, x))
def _initialize(self, params):
"""
Called when the solver object is declared. This function is
used to initialize the distribution function, and the field
quantities using the options as provided by the user. The
quantities are then mapped to the fourier basis by taking FFTs.
The independant modes are then evolved by using the linear
solver.
"""
# af.broadcast(function, *args) performs batched operations on
# function(*args):
f = af.broadcast(self.physical_system.initial_conditions.\
initialize_f, self.q1_center, self.q2_center,
self.p1, self.p2, self.p3, params
)
# Taking FFT:
f_hat = af.fft2(f)
# Since (k_q1, k_q2) = (0, 0) will give the background distribution:
# The division by (self.N_q1 * self.N_q2) is performed since the FFT
# at (0, 0) returns (amplitude * (self.N_q1 * self.N_q2))
self.f_background = af.abs(f_hat[0, 0, :])/ (self.N_q1 * self.N_q2)
# Calculating derivatives of the background distribution function:
self._calculate_dfdp_background()
# Scaling Appropriately:
# Except the case of (0, 0) the FFT returns
# (0.5 * amplitude * (self.N_q1 * self.N_q2)):
f_hat = 2 * f_hat / (self.N_q1 * self.N_q2)
if(self.single_mode_evolution == True):
# Background
f_hat[0, 0, :] = 0
# Finding the indices of the perturbation introduced:
i_q1_max = np.unravel_index(af.imax(af.abs(f_hat))[1],
(self.N_q1, self.N_q2,
self.N_p1 * self.N_p2 * self.N_p3
),order = 'F'
)[0]
i_q2_max = np.unravel_index(af.imax(af.abs(f_hat))[1],
(self.N_q1, self.N_q2,
self.N_p1 * self.N_p2 * self.N_p3
),order = 'F'
)[1]
# Taking sum to get a scalar value:
params.k_q1 = af.sum(self.k_q1[i_q1_max, i_q2_max])
params.k_q2 = af.sum(self.k_q2[i_q1_max, i_q2_max])
self.f_background = np.array(self.f_background).\
reshape(self.N_p1, self.N_p2, self.N_p3)
self.p1 = np.array(self.p1).\
reshape(self.N_p1, self.N_p2, self.N_p3)
self.p2 = np.array(self.p2).\
reshape(self.N_p1, self.N_p2, self.N_p3)
self.p3 = np.array(self.p3).\
reshape(self.N_p1, self.N_p2, self.N_p3)
self._A_q1 = np.array(self._A_q1).\
reshape(self.N_p1, self.N_p2, self.N_p3)
self._A_q2 = np.array(self._A_q2).\
reshape(self.N_p1, self.N_p2, self.N_p3)
params.rho_background = self.compute_moments('density',
self.f_background
)
params.p1_bulk_background = self.compute_moments('mom_p1_bulk',
self.f_background
) / params.rho_background
params.p2_bulk_background = self.compute_moments('mom_p2_bulk',
self.f_background
) / params.rho_background
params.p3_bulk_background = self.compute_moments('mom_p3_bulk',
self.f_background
) / params.rho_background
params.T_background = ((1 / params.p_dim) * \
self.compute_moments('energy',
self.f_background
)
- params.rho_background *
params.p1_bulk_background**2
- params.rho_background *
params.p2_bulk_background**2
- params.rho_background *
params.p3_bulk_background**2
) / params.rho_background
self.dfdp1_background = np.array(self.dfdp1_background).\
reshape(self.N_p1, self.N_p2, self.N_p3)
self.dfdp2_background = np.array(self.dfdp2_background).\
reshape(self.N_p1, self.N_p2, self.N_p3)
self.dfdp3_background = np.array(self.dfdp3_background).\
reshape(self.N_p1, self.N_p2, self.N_p3)
# Unable to recover pert_real and pert_imag accurately from the input.
# There seems to be some error in recovering these quantities which
# is dependant upon the resolution. Using user-defined quantities instead
delta_f_hat = params.pert_real * self.f_background \
+ params.pert_imag * self.f_background * 1j
self.Y = np.array([delta_f_hat])
compute_electrostatic_fields(self)
self.Y = np.array([delta_f_hat,
self.delta_E1_hat, self.delta_E2_hat, self.delta_E3_hat,
self.delta_B1_hat, self.delta_B2_hat, self.delta_B3_hat
]
)
else:
# Using a vector Y to evolve the system:
self.Y = af.constant(0, self.N_q1, self.N_q2,
self.N_p1 * self.N_p2 * self.N_p3,
7, dtype = af.Dtype.c64
)
# Assigning the 0th indice along axis 3 to the f_hat:
self.Y[:, :, :, 0] = f_hat
# Initializing the EM field quantities:
self.E3_hat = af.constant(0, self.N_q1, self.N_q2,
self.N_p1 * self.N_p2 * self.N_p3,
dtype = af.Dtype.c64
)
self.B1_hat = af.constant(0, self.N_q1, self.N_q2,
self.N_p1 * self.N_p2 * self.N_p3,
dtype = af.Dtype.c64
)
self.B2_hat = af.constant(0, self.N_q1, self.N_q2,
self.N_p1 * self.N_p2 * self.N_p3,
dtype = af.Dtype.c64
)
self.B3_hat = af.constant(0, self.N_q1, self.N_q2,
self.N_p1 * self.N_p2 * self.N_p3,
dtype = af.Dtype.c64
)
# Initializing EM fields using Poisson Equation:
if(self.physical_system.params.fields_initialize == 'electrostatic' or
self.physical_system.params.fields_initialize == 'fft'
):
compute_electrostatic_fields(self)
# If option is given as user-defined:
elif(self.physical_system.params.fields_initialize == 'user-defined'):
E1, E2, E3 = \
self.physical_system.initial_conditions.initialize_E(self.q1_center, self.q2_center, self.physical_system.params)
B1, B2, B3 = \
self.physical_system.initial_conditions.initialize_B(self.q1_center, self.q2_center, self.physical_system.params)
# Scaling Appropriately
self.E1_hat = af.tile(2 * af.fft2(E1) / (self.N_q1 * self.N_q2), 1, 1, self.N_p1 * self.N_p2 * self.N_p3)
self.E2_hat = af.tile(2 * af.fft2(E2) / (self.N_q1 * self.N_q2), 1, 1, self.N_p1 * self.N_p2 * self.N_p3)
self.E3_hat = af.tile(2 * af.fft2(E3) / (self.N_q1 * self.N_q2), 1, 1, self.N_p1 * self.N_p2 * self.N_p3)
self.B1_hat = af.tile(2 * af.fft2(B1) / (self.N_q1 * self.N_q2), 1, 1, self.N_p1 * self.N_p2 * self.N_p3)
self.B2_hat = af.tile(2 * af.fft2(B2) / (self.N_q1 * self.N_q2), 1, 1, self.N_p1 * self.N_p2 * self.N_p3)
self.B3_hat = af.tile(2 * af.fft2(B3) / (self.N_q1 * self.N_q2), 1, 1, self.N_p1 * self.N_p2 * self.N_p3)
else:
raise NotImplementedError('Method invalid/not-implemented')
# Assigning other indices along axis 3 to be
# the EM field quantities:
self.Y[:, :, :, 1] = self.E1_hat
self.Y[:, :, :, 2] = self.E2_hat
self.Y[:, :, :, 3] = self.E3_hat
self.Y[:, :, :, 4] = self.B1_hat
self.Y[:, :, :, 5] = self.B2_hat
self.Y[:, :, :, 6] = self.B3_hat
af.eval(self.Y)
return
pos1 = af.randu(3, 3) * 10 af.display(af.approx2(a, pos0, pos1)) a = af.randu(8, 1) af.display(a) af.display(af.fft(a)) af.display(af.dft(a)) af.display(af.real(af.ifft(af.fft(a)))) af.display(af.real(af.idft(af.dft(a)))) a = af.randu(4, 4) af.display(a) af.display(af.fft2(a)) af.display(af.dft(a)) af.display(af.real(af.ifft2(af.fft2(a)))) af.display(af.real(af.idft(af.dft(a)))) a = af.randu(4, 4, 2) af.display(a) af.display(af.fft3(a)) af.display(af.dft(a)) af.display(af.real(af.ifft3(af.fft3(a)))) af.display(af.real(af.idft(af.dft(a)))) a = af.randu(10, 1) b = af.randu(3, 1) af.display(af.convolve1(a, b))
def bench(A, iters = 100):
start = time()
for t in range(iters):
B = af.fft2(A)
af.sync()
return (time() - start) / iters
