def test_compute_electrostatic_fields():
test_obj = test()
compute_electrostatic_fields(test_obj)
E1 = 0.5 * test_obj.N_q1 * test_obj.N_q2 * af.ifft2(test_obj.E1_hat)
E2 = 0.5 * test_obj.N_q1 * test_obj.N_q2 * af.ifft2(test_obj.E2_hat)
E1_analytical = test_obj.physical_system.params.charge_electron \
* 2 * np.pi / (20 * np.pi**2) \
* ( 0.01 * af.sin(2 * np.pi * test_obj.q1 + 4 * np.pi * test_obj.q2)
+ 0.02 * af.cos(2 * np.pi * test_obj.q1 + 4 * np.pi * test_obj.q2)
)
E2_analytical = test_obj.physical_system.params.charge_electron \
* 4 * np.pi / (20 * np.pi**2) \
* ( 0.01 * af.sin(2 * np.pi * test_obj.q1 + 4 * np.pi * test_obj.q2)
+ 0.02 * af.cos(2 * np.pi * test_obj.q1 + 4 * np.pi * test_obj.q2)
)
add = lambda a,b:a+b
error_E1 = af.mean(af.abs(af.broadcast(add, E1_analytical, - E1)))
error_E2 = af.mean(af.abs(af.broadcast(add, E2_analytical, - E2)))
assert(error_E1 < 1e-14)
assert(error_E2 < 1e-14)
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, 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 ifft2(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.ifft2(array, scale=1)/(array.shape[0] * array.shape[1]) # fix for https://github.com/arrayfire/arrayfire/issues/2050
# 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 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 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 __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 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 compute_moments(self, moment_name, f=None, f_hat=None):
"""
Used in computing the moments of the distribution function.
The moment definitions which are passed to physical system
are used in computing these moment quantities.
Parameters
----------
moments_name : str
Pass the moment name which needs to be computed.
It must be noted that this needs to be defined by the
user under moment_defs under src and passed to the
physical_system object.
f/f_hat: np.ndarray
Pass this argument as well when you want to compute the
moments of the input array and not the one stored by the state vector
of the object.
Examples
--------
>> solver.compute_moments('density')
The above line will lookup the definition for 'density' under the dict
moments_exponents, and moments_coefficients and calculate the same
accordingly
When evolving for a single mode, this function returns moment_hat
"""
# Checking that the moment-name is defined by the user:
try:
moment_exponents = \
np.array(
self.physical_system.moment_exponents[moment_name]
)
moment_coeffs = \
np.array(
self.physical_system.moment_coeffs[moment_name]
)
except BaseException:
raise KeyError('moment_name not defined under physical system')
# This checks that the moment definition is of the form
# [[a, b, c], [d, e, f]...]. Alternatively if it isn't it checks
# whether the definition is of the form [a, b, c]
try:
moment_variable = 1
for i in range(moment_exponents.shape[0]):
moment_variable *= moment_coeffs[i, 0] \
* self.p1**(moment_exponents[i, 0]) \
+ moment_coeffs[i, 1] \
* self.p2**(moment_exponents[i, 1]) \
+ moment_coeffs[i, 2] \
* self.p3**(moment_exponents[i, 2])
except BaseException:
moment_variable = moment_coeffs[0] * self.p1**(moment_exponents[0]) \
+ moment_coeffs[1] * self.p2**(moment_exponents[1]) \
+ moment_coeffs[2] * self.p3**(moment_exponents[2])
if(self.single_mode_evolution == True):
if(f is None):
delta_moment_hat = np.sum(self.Y[0] * moment_variable) \
* self.dp3 * self.dp2 * self.dp1
return(delta_moment_hat)
else:
moment_hat = np.sum(f * moment_variable) \
* self.dp3 * self.dp2 * self.dp1
return(moment_hat)
# When evolving for several modes:
else:
# Since f_hat = Y[:, :, :, 0]:
# We sum along axis 2 which contains the variations in velocity:
# 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
# af.broadcast(function, *args) performs batched operations on
# function(*args):
if(f_hat is None and f is None):
moment_hat = af.sum(af.broadcast(multiply, self.Y[:, :, :, 0],
moment_variable
),
2
) * self.dp3 * self.dp2 * self.dp1
# Scaling Appropriately:
moment_hat = 0.5 * self.N_q2 * self.N_q1 * moment_hat
moment = af.real(af.ifft2(moment_hat))
elif(f_hat is not None and f is None):
moment_hat = af.sum(af.broadcast(multiply, f_hat,
moment_variable
),
2
) * self.dp3 * self.dp2 * self.dp1
# Scaling Appropriately:
moment_hat = 0.5 * self.N_q2 * self.N_q1 * moment_hat
moment = af.real(af.ifft2(moment_hat))
elif(f_hat is None and f is not None):
moment = af.sum(af.broadcast(multiply, f,
moment_variable
), 2
) * self.dp3 * self.dp2 * self.dp1
else:
raise BaseException('Invalid Option: Both f and f_hat cannot \
be provided as arguments'
)
af.eval(moment)
return(moment)
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))
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)) af.display(af.fft_convolve1(a, b)) af.display(af.convolve(a, b))
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
def dump_distribution_function(self, file_name):
"""
This function is used to dump distribution function to a file for
later usage.This dumps the complete 5D distribution function which
can be used for post-processing
Parameters
----------
file_name : The distribution_function array will be dumped to this
provided file name.
Returns
-------
This function returns None. However it creates a file 'file_name.h5',
containing the data of the distribution function
Examples
--------
>> solver.dump_distribution_function('distribution_function')
The above statement will create a HDF5 file which contains the
distribution function. The data is always stored with the key
'distribution_function'
This can later be accessed using
>> import h5py
>> h5f = h5py.File('distribution_function', 'r')
>> f = h5f['distribution_function'][:]
>> h5f.close()
"""
if (self.single_mode_evolution == True):
f_b = self.f_background.reshape(1, 1,
self.N_p1 * self.N_p2 * self.N_p3)
k_q1 = self.physical_system.params.k_q1
k_q2 = self.physical_system.params.k_q2
q1 = self.q1_center.to_ndarray().reshape(self.N_q1, self.N_q2, 1)
q2 = self.q2_center.to_ndarray().reshape(self.N_q1, self.N_q2, 1)
df = ( self.Y[0].reshape(1, 1, self.N_p1 * self.N_p2 * self.N_p3) \
* np.exp(1j * (k_q1 * q1 + k_q2 * q2))
).real
self._glob_f_value[:] = f_b + df
else:
# Scaling Appropriately:
self._glob_f_value[:] = 0.5 * self.N_q2 * self.N_q1 \
* np.array(af.ifft2(self.Y[:, :, :, 0])).real
viewer = PETSc.Viewer().createHDF5(file_name + '.h5', 'w')
viewer(self._glob_f)
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)
dt = params.N_cfl * min(nls.dq1, nls.dq2) \
/ max(domain.p1_end, domain.p2_end, domain.p3_end)
time_array = np.arange(0, params.t_final + dt, dt)
# Initializing Arrays used in storing the data:
E_data_ls = np.zeros_like(time_array)
E_data_nls = np.zeros_like(time_array)
for time_index, t0 in enumerate(time_array):
if (time_index % 100 == 0):
print('Computing For Time =', t0)
E_data_nls[time_index] = af.sum(nls.cell_centered_EM_fields[:, N_g:-N_g,
N_g:-N_g]**2)
E1_ls = af.real(0.5 * (ls.N_q1 * ls.N_q2) * af.ifft2(ls.E1_hat[:, :, 0]))
E_data_ls[time_index] = af.sum(E1_ls**2)
nls.strang_timestep(dt)
ls.RK4_timestep(dt)
h5f = h5py.File('data.h5', 'w')
h5f.create_dataset('electrical_energy_ls', data=E_data_ls)
h5f.create_dataset('electrical_energy_nls', data=E_data_nls)
h5f.create_dataset('time', data=time_array)
h5f.close()
pl.plot(time_array, E_data_ls, '--', color='black', label='Linear Solver')
pl.plot(time_array, E_data_nls, label='Nonlinear Solver')
pl.ylabel(r'SUM($|E|^2$)')
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 ifft(arr):
return af.ifft2(arr)
def af_unpadded_ifft2(fft, pad_periods=0):
side = fft.dims()[0] // (1 + pad_periods)
return af.ifft2(fft)[:side, :side]
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)) af.display(af.fft_convolve1(a, b)) af.display(af.convolve(a, b))
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))
