def BGK(f, t, q1, q2, v1, v2, v3, moments, params, flag=False):
"""
Return BGK operator -(f-f0)/tau.
Parameters:
-----------
f : Distribution function array
shape:(N_v, N_s, N_q1, N_q2)
t : Time elapsed
q1 : The array that holds data for the q1 dimension in q-space
shape:(1, 1, N_q1, N_q2)
q2 : The array that holds data for the q2 dimension in q-space
shape:(1, 1, N_q1, N_q2)
v1 : The array that holds data for the v1 dimension in v-space
shape:(N_v, N_s, 1, 1)
v2 : The array that holds data for the v2 dimension in v-space
shape:(N_v, N_s, 1, 1)
v3 : The array that holds data for the v3 dimension in v-space
shape:(N_v, N_s, 1, 1)
params: The parameters file/object that is originally declared by the user.
This can be used to inject other functions/attributes into the function
flag: Toggle used for evaluating tau = 0 cases need to be evaluated. When set to True, this
function is made to return f0, thus setting f = f0 wherever tau = 0
"""
n = moments('density', f)
m = params.mass
# Floor used to avoid 0/0 limit:
eps = 1e-30
v1_bulk = moments('mom_v1_bulk', f) / (n + eps)
v2_bulk = moments('mom_v2_bulk', f) / (n + eps)
v3_bulk = moments('mom_v3_bulk', f) / (n + eps)
T = (1 / params.p_dim) * (2 * multiply(moments('energy', f), m) - multiply(
n, m) * v1_bulk**2 - multiply(n, m) * v2_bulk**2 -
multiply(n, m) * v3_bulk**2) / (n + eps) + eps
f_MB = f0(v1, v2, v3, n, T, v1_bulk, v2_bulk, v3_bulk, params)
tau = params.tau(q1, q2, v1, v2, v3)
if (flag == False):
C_f = -(f - f_MB) / tau
return (C_f)
# Accounting for purely collisional cases(f = f0):
else:
return (f_MB)
def calculate_dfdp_background(self):
"""
Calculates the derivative of the background distribution
with respect to the variables p1, p2, p3. This is used to
solve for the contribution from the fields
"""
f_b = af.moddims(self.f_background, self.N_p1, self.N_p2, self.N_p3,
self.N_species)
# 8th order central diff -- just in case
# f_minus_4 = af.shift(f_b, 4); f_plus_4 = af.shift(f_b, -4)
# f_minus_3 = af.shift(f_b, 3); f_plus_3 = af.shift(f_b, -3)
# f_minus_2 = af.shift(f_b, 2); f_plus_2 = af.shift(f_b, -2)
# f_minus_1 = af.shift(f_b, 1); f_plus_1 = af.shift(f_b, -1)
#
# dfdp1_background = \
# ( 1/280)*f_minus_4 + (-4/105)*f_minus_3 + ( 1/5)*f_minus_2 + (-4/5)*f_minus_1 \
# + (-1/280)*f_plus_4 + ( 4/105)*f_plus_3 + (-1/5)*f_plus_2 + ( 4/5)*f_plus_1
#
# dfdp1_background = dfdp1_background/self.dp1
# Using a 4th order central difference stencil:
# Using moddims since in p-expanded dp needs to be (1, 1, 1, Ns)
dfdp1_background = multiply(
(-af.shift(f_b, -2) + 8 * af.shift(f_b, -1) + af.shift(f_b, 2) -
8 * af.shift(f_b, 1)),
1 / af.moddims(12 * self.dp1, 1, 1, 1, len(self.dp1)))
dfdp2_background = multiply(
(-af.shift(f_b, 0, -2) + 8 * af.shift(f_b, 0, -1) +
af.shift(f_b, 0, 2) - 8 * af.shift(f_b, 0, 1)),
1 / af.moddims(12 * self.dp2, 1, 1, 1, len(self.dp2)))
dfdp3_background = multiply(
(-af.shift(f_b, 0, 0, -2) + 8 * af.shift(f_b, 0, 0, -1) +
af.shift(f_b, 0, 0, 2) - 8 * af.shift(f_b, 0, 0, 1)),
1 / af.moddims(12 * self.dp3, 1, 1, 1, len(self.dp3)))
# Reshaping such that the variations in velocity are along axis 0:
self.dfdp1_background = af.moddims(dfdp1_background,
self.N_p1 * self.N_p2 * self.N_p3,
self.N_species)
self.dfdp2_background = af.moddims(dfdp2_background,
self.N_p1 * self.N_p2 * self.N_p3,
self.N_species)
self.dfdp3_background = af.moddims(dfdp3_background,
self.N_p1 * self.N_p2 * self.N_p3,
self.N_species)
af.eval(self.dfdp1_background, self.dfdp2_background,
self.dfdp3_background)
return
def compute_electrostatic_fields(self, rho_hat):
"""
Computes the electrostatic fields by making use of FFTs by solving
the Poisson equation: div^2 phi = rho to return the FT of the
fields.
Parameters
----------
rho_hat : af.Array
FT for the charge density for each of the species.
shape:(1, N_s, N_q1, N_q2)
"""
# Summing over all the species:
phi_hat = multiply(af.sum(rho_hat, 1), 1 / (self.k_q1**2 + self.k_q2**2)) # (1, 1, N_q1, N_q2)
# Setting the background electric potential to zero:
phi_hat[: , :, 0, 0] = 0
self.E1_hat = -phi_hat * 1j * self.k_q1 / self.params.eps
self.E2_hat = -phi_hat * 1j * self.k_q2 / self.params.eps
self.E3_hat = 0 * self.E1_hat
self.B1_hat = 0 * self.E1_hat
self.B2_hat = 0 * self.E1_hat
self.B3_hat = 0 * self.E1_hat
af.eval(self.E1_hat, self.E2_hat, self.E3_hat,
self.B1_hat, self.B2_hat, self.B3_hat
)
return
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.
Parameters:
-----------
params: The parameters file/object that is originally declared
by the user.
"""
# 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_center, self.p2_center, self.p3_center, params
)
# Taking FFT:
self.f_hat = 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(self.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)):
self.f_hat = 2 * self.f_hat / (self.N_q1 * self.N_q2)
rho_hat_initial = multiply(self.physical_system.params.charge,
self.compute_moments('density', f_hat=self.f_hat)
)
self.fields_solver = fields_solver(self.physical_system,
rho_hat_initial
)
return
def _initialize(self, params):
"""
Called when the solver object is declared. This function is
used to initialize the distribution function, using the options
as provided by the user.
Parameters
----------
params : module
params contains all details of which methods to use
in addition to useful physical constant. Additionally,
it can also be used to inject methods which need to be
used inside some solver routine
"""
# Initializing with the provided I.C's:
# af.broadcast, allows us to perform batched operations
# when operating on arrays of different sizes
# af.broadcast(function, *args) performs batched
# operations on function(*args)
self.f = af.broadcast(self.physical_system.initial_conditions.\
initialize_f, self.q1_center, self.q2_center,
self.p1_center, self.p2_center, self.p3_center, params
)
self.f_initial = self.f
if(self.physical_system.params.fields_enabled):
rho_initial = multiply(self.physical_system.params.charge,
self.compute_moments('density')
)
self.fields_solver = fields_solver(self.physical_system, rho_initial,
self.performance_test_flag
)
def op_fields(self, dt):
"""
Evolves the following part of the equations specified:
df/dt + A_p1 df/dp1 + A_p2 df/dp1 + A_p3 df/dp1 = 0
Parameters
----------
dt : double
Time-step size to evolve the system
"""
if(self.performance_test_flag == True):
tic = af.time()
if(self.physical_system.params.fields_solver == 'electrostatic'):
rho = multiply(self.physical_system.params.charge,
self.compute_moments('density')
)
self.fields_solver.compute_electrostatic_fields(rho)
# Evolving fields:
if(self.physical_system.params.fields_solver == 'fdtd'):
if(self.physical_system.params.hybrid_model_enabled == True):
communicate_fields(self.fields_solver, True)
B1 = self.fields_solver.yee_grid_EM_fields[3] # (i + 1/2, j)
B2 = self.fields_solver.yee_grid_EM_fields[4] # (i, j + 1/2)
B3 = self.fields_solver.yee_grid_EM_fields[5] # (i, j)
B1_plus_q2 = af.shift(B1, 0, 0, 0, -1)
B2_plus_q1 = af.shift(B2, 0, 0, -1, 0)
B3_plus_q1 = af.shift(B3, 0, 0, -1, 0)
B3_plus_q2 = af.shift(B3, 0, 0, 0, -1)
# curlB_x = dB3/dq2
curlB_1 = (B3_plus_q2 - B3) / self.dq2 # (i, j + 1/2)
# curlB_y = -dB3/dq1
curlB_2 = -(B3_plus_q1 - B3) / self.dq1 # (i + 1/2, j)
# curlB_z = (dB2/dq1 - dB1/dq2)
curlB_3 = (B2_plus_q1 - B2) / self.dq1 - (B1_plus_q2 - B1) / self.dq2 # (i + 1/2, j + 1/2)
# c --> inf limit: J = (∇ x B) / μ
mu = self.physical_system.params.mu
J1 = curlB_1 / mu # (i, j + 1/2)
J2 = curlB_2 / mu # (i + 1/2, j)
J3 = curlB_3 / mu # (i + 1/2, j + 1/2)
# Using Generalized Ohm's Law for electric field:
# (v X B)_x = B3 * v2 - B2 * v3
# (v X B)_x --> (i, j + 1/2)
v_cross_B_1 = 0.5 * (B3_plus_q2 + B3) * self.compute_moments('mom_v2_bulk', f = 0.5 * (self.f + af.shift(self.f, 0, 0, 1))) \
/ self.compute_moments('density', f = 0.5 * (self.f + af.shift(self.f, 0, 0, 1))) \
- B2 * self.compute_moments('mom_v3_bulk', f = 0.5 * (self.f + af.shift(self.f, 0, 0, 1))) \
/ self.compute_moments('density', f = 0.5 * (self.f + af.shift(self.f, 0, 0, 1)))
# (v X B)_y = B1 * v3 - B3 * v1
# (v X B)_y --> (i + 1/2, j)
v_cross_B_2 = B1 * self.compute_moments('mom_v3_bulk', f = 0.5 * (self.f + af.shift(self.f, 0, 0, 0, 1))) \
/ self.compute_moments('density', f = 0.5 * (self.f + af.shift(self.f, 0, 0, 0, 1))) \
- 0.5 * (B3_plus_q1 + B3) * self.compute_moments('mom_v1_bulk', f = 0.5 * (self.f + af.shift(self.f, 0, 0, 0, 1))) \
/ self.compute_moments('density', f = 0.5 * (self.f + af.shift(self.f, 0, 0, 0, 1)))
# (v X B)_z = B2 * v1 - B1 * v2
# (v X B)_z --> (i + 1/2, j + 1/2)
v_cross_B_3 = 0.5 * (B2_plus_q1 + B2) * self.compute_moments('mom_v1_bulk', f = self.f) \
/ self.compute_moments('density', f = self.f) \
- 0.5 * (B1_plus_q2 + B1) * self.compute_moments('mom_v2_bulk', f = self.f) \
/ self.compute_moments('density', f = self.f)
# (J X B)_x = B3 * J2 - B2 * J3
# (J X B)_x --> (i, j + 1/2)
J_cross_B_1 = 0.5 * (B3_plus_q2 + B3) * ( J2 + af.shift(J2, 0, 0, 0, -1)
+ af.shift(J2, 0, 0, 1) + af.shift(J2, 0, 0, 1, -1)
) * 0.25 \
- B2 * (af.shift(J3, 0, 0, 1) + J3) * 0.5
# (J X B)_y = B1 * J3 - B3 * J1
# (J X B)_y --> (i + 1/2, j)
J_cross_B_2 = B1 * (af.shift(J3, 0, 0, 0, 1) + J3) * 0.5 \
- 0.5 * (B3_plus_q1 + B3) * ( J1 + af.shift(J1, 0, 0, 0, 1)
+ af.shift(J1, 0, 0, -1) + af.shift(J1, 0, 0, -1, 1)
) * 0.25
# (J X B)_z = B2 * J1 - B1 * J2
# (J X B)_z --> (i + 1/2, j + 1/2)
J_cross_B_3 = 0.5 * (B2_plus_q1 + B2) * (af.shift(J1, 0, 0, -1) + J1) * 0.5 \
- 0.5 * (B1_plus_q2 + B1) * (af.shift(J2, 0, 0, 0, -1) + J2) * 0.5
n_i = self.compute_moments('density')
T_e = self.physical_system.params.fluid_electron_temperature
# Using a 4th order stencil:
dn_q1 = (- af.shift(n_i, 0, 0, -2) + 8 * af.shift(n_i, 0, 0, -1)
- 8 * af.shift(n_i, 0, 0, 1) + af.shift(n_i, 0, 0, 2)
) / (12 * self.dq1)
dn_q2 = (- af.shift(n_i, 0, 0, 0, -2) + 8 * af.shift(n_i, 0, 0, 0, -1)
- 8 * af.shift(n_i, 0, 0, 0, 1) + af.shift(n_i, 0, 0, 0, 2)
) / (12 * self.dq2)
# E = -(v X B) + (J X B) / (en) - T ∇n / (en)
E1 = -v_cross_B_1 + J_cross_B_1 \
/ (multiply(self.compute_moments('density', f = 0.5 * (self.f + af.shift(self.f, 0, 0, 1))),
self.physical_system.params.charge
)
) \
- 0.5 * T_e * (dn_q1 + af.shift(dn_q1, 0, 0, 1)) / multiply(self.physical_system.params.charge, n_i) # (i, j + 1/2)
E2 = -v_cross_B_2 + J_cross_B_2 \
/ (multiply(self.compute_moments('density', f = 0.5 * (self.f + af.shift(self.f, 0, 0, 0, 1))),
self.physical_system.params.charge
)
) \
- 0.5 * T_e * (dn_q2 + af.shift(dn_q2, 0, 0, 0, 1)) / multiply(self.physical_system.params.charge, n_i) # (i + 1/2, j)
E3 = -v_cross_B_3 + J_cross_B_3 \
/ (multiply(self.compute_moments('density', f = self.f),
self.physical_system.params.charge
)
) # (i + 1/2, j + 1/2)
self.fields_solver.yee_grid_EM_fields[0] = E1
self.fields_solver.yee_grid_EM_fields[1] = E2
self.fields_solver.yee_grid_EM_fields[2] = E3
af.eval(self.fields_solver.yee_grid_EM_fields)
else:
J1 = multiply(self.physical_system.params.charge,
self.compute_moments('mom_v1_bulk', f = 0.5 * (self.f + af.shift(self.f, 0, 0, 1)))
) # (i, j + 1/2)
J2 = multiply(self.physical_system.params.charge,
self.compute_moments('mom_v2_bulk', f = 0.5 * (self.f + af.shift(self.f, 0, 0, 0, 1)))
) # (i + 1/2, j)
J3 = multiply(self.physical_system.params.charge,
self.compute_moments('mom_v3_bulk', f = f)
) # (i + 1/2, j + 1/2)
self.fields_solver.evolve_electrodynamic_fields(J1, J2, J3, self.dt)
f_interp_p_3d(self, dt)
af.eval(self.f)
if(self.performance_test_flag == True):
af.sync()
toc = af.time()
self.time_fieldstep += toc - tic
return
def dfields_hat_dt(f_hat, fields_hat, self):
"""
Returns the value of the derivative of the fields_hat with respect to time
respect to time. This is used to evolve the fields with time.
NOTE:All the fields quantities are included in fields_hat as follows:
E1_hat = fields_hat[0]
E2_hat = fields_hat[1]
E3_hat = fields_hat[2]
B1_hat = fields_hat[3]
B2_hat = fields_hat[4]
B3_hat = fields_hat[5]
Input:
------
f_hat : Fourier mode values for the distribution function at which the slope is computed
At t = 0 the initial state of the system is passed to this function:
fields_hat : Fourier mode values for the fields at which the slope is computed
At t = 0 the initial state of the system is passed to this function:
Output:
-------
df_dt : The time-derivative of f_hat
"""
eps = self.physical_system.params.eps
mu = self.physical_system.params.mu
B1_hat = fields_hat[3]
B2_hat = fields_hat[4]
B3_hat = fields_hat[5]
if (self.physical_system.params.hybrid_model_enabled == True):
# curlB_x = dB3/dq2
curlB_1 = B3_hat * 1j * self.k_q2
# curlB_y = -dB3/dq1
curlB_2 = -B3_hat * 1j * self.k_q1
# curlB_z = (dB2/dq1 - dB1/dq2)
curlB_3 = (B2_hat * 1j * self.k_q1 - B1_hat * 1j * self.k_q2)
# c --> inf limit: J = (∇ x B) / μ
J1_hat = curlB_1 / mu
J2_hat = curlB_2 / mu
J3_hat = curlB_3 / mu
else:
J1_hat = multiply(self.physical_system.params.charge,
self.compute_moments('mom_v1_bulk', f_hat=f_hat))
J2_hat = multiply(self.physical_system.params.charge,
self.compute_moments('mom_v2_bulk', f_hat=f_hat))
J3_hat = multiply(self.physical_system.params.charge,
self.compute_moments('mom_v3_bulk', f_hat=f_hat))
# Summing along all species:
J1_hat = af.sum(J1_hat, 1)
J2_hat = af.sum(J2_hat, 1)
J3_hat = af.sum(J3_hat, 1)
# Checking that there is no mean field component:
# try:
# assert(af.mean(af.abs(B1_hat[:, 0, 0])) < 1e-12)
# assert(af.mean(af.abs(B2_hat[:, 0, 0])) < 1e-12)
# assert(af.mean(af.abs(B3_hat[:, 0, 0])) < 1e-12)
# except:
# raise SystemExit('Linear Solver cannot solve for non-zero mean magnetic fields')
# Equations Solved:
# dE1/dt = + dB3/dq2 - J1
# dE2/dt = - dB3/dq1 - J2
# dE3/dt = dB2/dq1 - dB1/dq2 - J3
if (self.physical_system.params.hybrid_model_enabled == True):
# Using Generalized Ohm's Law for electric field:
n_i_hat = self.compute_moments('density', f_hat=f_hat)
n_i = af.real(ifft2(0.5 * self.N_q2 * self.N_q1 * n_i_hat))
v1 = af.real(
ifft2(0.5 * self.N_q2 * self.N_q1 *
self.compute_moments('mom_v1_bulk', f_hat=f_hat))) / n_i
v2 = af.real(
ifft2(0.5 * self.N_q2 * self.N_q1 *
self.compute_moments('mom_v2_bulk', f_hat=f_hat))) / n_i
v3 = af.real(
ifft2(0.5 * self.N_q2 * self.N_q1 *
self.compute_moments('mom_v3_bulk', f_hat=f_hat))) / n_i
B1 = af.real(ifft2(0.5 * self.N_q2 * self.N_q1 * B1_hat))
B2 = af.real(ifft2(0.5 * self.N_q2 * self.N_q1 * B2_hat))
B3 = af.real(ifft2(0.5 * self.N_q2 * self.N_q1 * B3_hat))
J1 = af.real(ifft2(0.5 * self.N_q2 * self.N_q1 * J1_hat))
J2 = af.real(ifft2(0.5 * self.N_q2 * self.N_q1 * J2_hat))
J3 = af.real(ifft2(0.5 * self.N_q2 * self.N_q1 * J3_hat))
T_e = self.physical_system.params.fluid_electron_temperature
# Computing the quantities needed and then doing FT again:
# # (v X B)_x = B3 * v2 - B2 * v3
# v_cross_B_1 = B3 * v2 - B2 * v3
# # (v X B)_y = B1 * v3 - B3 * v1
# v_cross_B_2 = B1 * v3 - B3 * v1
# # (v X B)_z = B2 * v1 - B1 * v2
# v_cross_B_3 = B2 * v1 - B1 * v2
# # (J X B)_x = B3 * J2 - B2 * J3
# J_cross_B_1 = B3 * J2 - B2 * J3
# # (J X B)_y = B1 * J3 - B3 * J1
# J_cross_B_2 = B1 * J3 - B3 * J1
# # (J X B)_z = B2 * J1 - B1 * J2
# J_cross_B_3 = B2 * J1 - B1 * J2
# # Using a 4th order stencil:
# dn_q1 = (- af.shift(n_i, 0, 0, -2) + 8 * af.shift(n_i, 0, 0, -1)
# - 8 * af.shift(n_i, 0, 0, 1) + af.shift(n_i, 0, 0, 2)
# ) / (12 * self.dq1)
# dn_q2 = (- af.shift(n_i, 0, 0, 0, -2) + 8 * af.shift(n_i, 0, 0, 0, -1)
# - 8 * af.shift(n_i, 0, 0, 0, 1) + af.shift(n_i, 0, 0, 0, 2)
# ) / (12 * self.dq2)
# # E = -(v X B) + (J X B) / (en) - T ∇n / (en)
# E1 = - v_cross_B_1 + J_cross_B_1 / multiply(self.physical_system.params.charge, n_i) \
# - T_e * dn_q1 / multiply(self.physical_system.params.charge, n_i)
# E2 = - v_cross_B_2 + J_cross_B_2 / multiply(self.physical_system.params.charge, n_i) \
# - T_e * dn_q2 / multiply(self.physical_system.params.charge, n_i)
# E3 = - v_cross_B_3 + J_cross_B_3 / multiply(self.physical_system.params.charge, n_i)
# E1_hat = 2 * fft2(E1) / (self.N_q1 * self.N_q2)
# E2_hat = 2 * fft2(E2) / (self.N_q1 * self.N_q2) - T_e * n_i_hat * self.k_q2 / af.mean(n_i)
# E3_hat = 2 * fft2(E3) / (self.N_q1 * self.N_q2)
# background magnetic fields == 0 ALWAYS:
# (v X B)_x = B3 * v2 - B2 * v3
v_cross_B_1_hat = B3_hat * af.mean(v2) - B2_hat * af.mean(v3)
# (v X B)_y = B1 * v3 - B3 * v1
v_cross_B_2_hat = B1_hat * af.mean(v3) - B3_hat * af.mean(v1)
# (v X B)_z = B2 * v1 - B1 * v2
v_cross_B_3_hat = B2_hat * af.mean(v1) - B1_hat * af.mean(v2)
# (J X B)_x = B3 * J2 - B2 * J3
J_cross_B_1_hat = B3_hat * af.mean(J2) - B2_hat * af.mean(J3)
# (J X B)_y = B1 * J3 - B3 * J1
J_cross_B_2_hat = B1_hat * af.mean(J3) - B3_hat * af.mean(J1)
# (J X B)_z = B2 * J1 - B1 * J2
J_cross_B_3_hat = B2_hat * af.mean(J1) - B1_hat * af.mean(J2)
E1_hat = - v_cross_B_1_hat + multiply(J_cross_B_1_hat, 1 / self.physical_system.params.charge) / af.mean(n_i) \
- T_e * n_i_hat * self.k_q1 / af.mean(n_i)
E2_hat = - v_cross_B_2_hat + multiply(J_cross_B_2_hat, 1 / self.physical_system.params.charge) / af.mean(n_i) \
- T_e * n_i_hat * self.k_q2 / af.mean(n_i)
E3_hat = -v_cross_B_3_hat + multiply(
J_cross_B_3_hat,
1 / self.physical_system.params.charge) / af.mean(n_i)
self.fields_solver.fields_hat[0] = E1_hat
self.fields_solver.fields_hat[1] = E2_hat
self.fields_solver.fields_hat[2] = E3_hat
else:
E1_hat = fields_hat[0]
E2_hat = fields_hat[1]
E3_hat = fields_hat[2]
dE1_hat_dt = B3_hat * 1j * self.k_q2 / (mu * eps) - J1_hat / eps
dE2_hat_dt = -B3_hat * 1j * self.k_q1 / (mu * eps) - J2_hat / eps
dE3_hat_dt = (B2_hat * 1j * self.k_q1 - B1_hat * 1j * self.k_q2) / (mu * eps) \
-J3_hat / eps
# dB1/dt = - dE3/dq2
# dB2/dt = + dE3/dq1
# dB3/dt = - (dE2/dq1 - dE1/dq2)
dB1_hat_dt = -E3_hat * 1j * self.k_q2
dB2_hat_dt = E3_hat * 1j * self.k_q1
dB3_hat_dt = E1_hat * 1j * self.k_q2 - E2_hat * 1j * self.k_q1
dfields_hat_dt = af.join(0, af.join(0, dE1_hat_dt, dE2_hat_dt, dE3_hat_dt),
dB1_hat_dt, dB2_hat_dt, dB3_hat_dt)
af.eval(dfields_hat_dt)
return (dfields_hat_dt)
def df_dt_fvm(f, self, term_to_return='all'):
"""
Returns the expression for df/dt which is then
evolved by a timestepper.
Parameters
----------
f : af.Array
Array of the distribution function at which df_dt is to
be evaluated.
"""
# Giving shorter name references:
reconstruction_in_q = self.physical_system.params.reconstruction_method_in_q
reconstruction_in_p = self.physical_system.params.reconstruction_method_in_p
riemann_in_q = self.physical_system.params.riemann_solver_in_q
riemann_in_p = self.physical_system.params.riemann_solver_in_p
# Initializing df_dt
df_dt = 0
if (self.physical_system.params.solver_method_in_q == 'FVM'):
f_left_plus_eps, f_right_minus_eps = reconstruct(
self, f, 2, reconstruction_in_q)
f_bot_plus_eps, f_top_minus_eps = reconstruct(self, f, 3,
reconstruction_in_q)
# f_left_minus_eps of i-th cell is f_right_minus_eps of the (i-1)th cell
f_left_minus_eps = af.shift(f_right_minus_eps, 0, 0, 1)
# Extending the same to bot:
f_bot_minus_eps = af.shift(f_top_minus_eps, 0, 0, 0, 1)
# af.broadcast used to perform batched operations on arrays of different sizes:
self._C_q1 = af.broadcast(self._C_q, self.time_elapsed,
self.q1_left_center, self.q2_left_center,
self.p1_center, self.p2_center,
self.p3_center,
self.physical_system.params)[0]
self._C_q2 = af.broadcast(self._C_q, self.time_elapsed,
self.q1_center_bot, self.q2_center_bot,
self.p1_center, self.p2_center,
self.p3_center,
self.physical_system.params)[1]
f_left = riemann_solver(self, f_left_minus_eps, f_left_plus_eps,
self._C_q1)
f_bot = riemann_solver(self, f_bot_minus_eps, f_bot_plus_eps,
self._C_q2)
left_flux = multiply(self._C_q1, f_left)
bot_flux = multiply(self._C_q2, f_bot)
right_flux = af.shift(left_flux, 0, 0, -1)
top_flux = af.shift(bot_flux, 0, 0, 0, -1)
# First get the purely spatial sqrt_get_g
#g_tmp = sqrt_det_g(self.q1_center, self.q2_center, \
# self.physical_system.params.q1_start_local_left, \
# self.physical_system.params.q2_start_local_bottom)
g_tmp = self.physical_system.params.sqrt_det_g
# Now need to make it compatible with the higher dim data structure f
g = multiply(g_tmp, self.p1_center**0)
df_dt += - (right_flux - left_flux) / (g*self.dq1) \
- (top_flux - bot_flux ) / (g*self.dq2) \
if (self.physical_system.params.source_enabled == True):
df_dt += self._source(f, self.time_elapsed, self.q1_center,
self.q2_center, self.p1_center,
self.p2_center, self.p3_center,
self.compute_moments,
self.physical_system.params, False)
if (self.physical_system.params.solver_method_in_p == 'FVM'
and self.physical_system.params.fields_enabled == True):
if (self.physical_system.params.fields_type == 'electrodynamic'
and self.fields_solver.at_n == False):
#TODO: Remove hybrid_model from lib and put into src
if (self.physical_system.params.hybrid_model_enabled == True):
communicate_fields(self.fields_solver, True)
B1 = self.fields_solver.yee_grid_EM_fields[3] # (i + 1/2, j)
B2 = self.fields_solver.yee_grid_EM_fields[4] # (i, j + 1/2)
B3 = self.fields_solver.yee_grid_EM_fields[5] # (i, j)
B1_plus_q2 = af.shift(B1, 0, 0, 0, -1)
B2_plus_q1 = af.shift(B2, 0, 0, -1, 0)
B3_plus_q1 = af.shift(B3, 0, 0, -1, 0)
B3_plus_q2 = af.shift(B3, 0, 0, 0, -1)
# curlB_x = dB3/dq2
curlB_1 = (B3_plus_q2 - B3) / self.dq2 # (i, j + 1/2)
# curlB_y = -dB3/dq1
curlB_2 = -(B3_plus_q1 - B3) / self.dq1 # (i + 1/2, j)
# curlB_z = (dB2/dq1 - dB1/dq2)
curlB_3 = (B2_plus_q1 - B2) / self.dq1 - (
B1_plus_q2 - B1) / self.dq2 # (i + 1/2, j + 1/2)
# c --> inf limit: J = (∇ x B) / μ
mu = self.physical_system.params.mu
J1 = curlB_1 / mu # (i, j + 1/2)
J2 = curlB_2 / mu # (i + 1/2, j)
J3 = curlB_3 / mu # (i + 1/2, j + 1/2)
# Using Generalized Ohm's Law for electric field:
# (v X B)_x = B3 * v2 - B2 * v3
# (v X B)_x --> (i, j + 1/2)
v_cross_B_1 = 0.5 * (B3_plus_q2 + B3) * self.compute_moments('mom_v2_bulk', f = f_left) \
/ self.compute_moments('density', f = f_left) \
- B2 * self.compute_moments('mom_v3_bulk', f = f_left) \
/ self.compute_moments('density', f = f_left)
# (v X B)_y = B1 * v3 - B3 * v1
# (v X B)_y --> (i + 1/2, j)
v_cross_B_2 = B1 * self.compute_moments('mom_v3_bulk', f = f_bot) \
/ self.compute_moments('density', f = f_bot) \
- 0.5 * (B3_plus_q1 + B3) * self.compute_moments('mom_v1_bulk', f = f_bot) \
/ self.compute_moments('density', f = f_bot)
# (v X B)_z = B2 * v1 - B1 * v2
# (v X B)_z --> (i + 1/2, j + 1/2)
v_cross_B_3 = 0.5 * (B2_plus_q1 + B2) * self.compute_moments('mom_v1_bulk', f = f) \
/ self.compute_moments('density', f = f) \
- 0.5 * (B1_plus_q2 + B1) * self.compute_moments('mom_v2_bulk', f = f) \
/ self.compute_moments('density', f = f)
# (J X B)_x = B3 * J2 - B2 * J3
# (J X B)_x --> (i, j + 1/2)
J_cross_B_1 = 0.5 * (B3_plus_q2 + B3) * ( J2 + af.shift(J2, 0, 0, 0, -1)
+ af.shift(J2, 0, 0, 1) + af.shift(J2, 0, 0, 1, -1)
) * 0.25 \
- B2 * (af.shift(J3, 0, 0, 1) + J3) * 0.5
# (J X B)_y = B1 * J3 - B3 * J1
# (J X B)_y --> (i + 1/2, j)
J_cross_B_2 = B1 * (af.shift(J3, 0, 0, 0, 1) + J3) * 0.5 \
- 0.5 * (B3_plus_q1 + B3) * ( J1 + af.shift(J1, 0, 0, 0, 1)
+ af.shift(J1, 0, 0, -1) + af.shift(J1, 0, 0, -1, 1)
) * 0.25
# (J X B)_z = B2 * J1 - B1 * J2
# (J X B)_z --> (i + 1/2, j + 1/2)
J_cross_B_3 = 0.5 * (B2_plus_q1 + B2) * (af.shift(J1, 0, 0, -1) + J1) * 0.5 \
- 0.5 * (B1_plus_q2 + B1) * (af.shift(J2, 0, 0, 0, -1) + J2) * 0.5
n_i = self.compute_moments('density')
T_e = self.physical_system.params.fluid_electron_temperature
# Using a 4th order stencil:
dn_q1 = (-af.shift(n_i, 0, 0, -2) + 8 *
af.shift(n_i, 0, 0, -1) - 8 * af.shift(n_i, 0, 0, 1) +
af.shift(n_i, 0, 0, 2)) / (12 * self.dq1)
dn_q2 = (-af.shift(n_i, 0, 0, 0, -2) +
8 * af.shift(n_i, 0, 0, 0, -1) -
8 * af.shift(n_i, 0, 0, 0, 1) +
af.shift(n_i, 0, 0, 0, 2)) / (12 * self.dq2)
# E = -(v X B) + (J X B) / (en) - T ∇n / (en)
E1 = -v_cross_B_1 + J_cross_B_1 \
/ (multiply(self.compute_moments('density', f = f_left),
self.physical_system.params.charge
)
) \
- 0.5 * T_e * (dn_q1 + af.shift(dn_q1, 0, 0, 1)) / multiply(self.physical_system.params.charge, n_i) # (i, j + 1/2)
E2 = -v_cross_B_2 + J_cross_B_2 \
/ (multiply(self.compute_moments('density', f = f_bot),
self.physical_system.params.charge
)
) \
- 0.5 * T_e * (dn_q2 + af.shift(dn_q2, 0, 0, 0, 1)) / multiply(self.physical_system.params.charge, n_i) # (i + 1/2, j)
E3 = -v_cross_B_3 + J_cross_B_3 \
/ (multiply(self.compute_moments('density', f = f),
self.physical_system.params.charge
)
) # (i + 1/2, j + 1/2)
self.fields_solver.yee_grid_EM_fields[0] = E1
self.fields_solver.yee_grid_EM_fields[1] = E2
self.fields_solver.yee_grid_EM_fields[2] = E3
af.eval(self.fields_solver.yee_grid_EM_fields)
else:
J1 = multiply(self.physical_system.params.charge,
self.compute_moments('mom_v1_bulk',
f=f_left)) # (i, j + 1/2)
J2 = multiply(self.physical_system.params.charge,
self.compute_moments('mom_v2_bulk',
f=f_bot)) # (i + 1/2, j)
J3 = multiply(self.physical_system.params.charge,
self.compute_moments('mom_v3_bulk',
f=f)) # (i + 1/2, j + 1/2)
# This gets called only at the (n+1/2)-th step
# This means that J's are at (n+1/2)
# evolves E^n --> E^{n + 1}
# evolves B^{n+1/2} --> B^{n + 3 / 2}
# Updates cell_centered_EM_fields_at_{n/n_plus_half}:
# cell_centered_EM_fields_at_n from (E^{n-1}, B^{n-1}) ---> (E^{n}, B^{n})
# cell_centered_EM_fields_at_n_plus_half from (E^{n-1/2}, B^{n-1/2}) ---> (E^{n+1/2}, B^{n+1/2})
self.fields_solver.evolve_electrodynamic_fields(
J1, J2, J3, self.dt)
if (self.physical_system.params.fields_type == 'electrostatic'):
if (self.physical_system.params.fields_solver == 'fft'):
rho = multiply(self.physical_system.params.charge,
self.compute_moments('density', f=f))
self.fields_solver.compute_electrostatic_fields(rho)
#TODO : This is a hack. Fix.
E1 = 0. * self.fields_solver.cell_centered_EM_fields[0]
E2 = 0. * self.fields_solver.cell_centered_EM_fields[1]
E3 = 0. * self.fields_solver.cell_centered_EM_fields[2]
B1 = 0. * self.fields_solver.cell_centered_EM_fields[3]
B2 = 0. * self.fields_solver.cell_centered_EM_fields[4]
B3 = 0. * self.fields_solver.cell_centered_EM_fields[5]
# Fields solver object is passed to C_p where the get_fields method
# is used to get the electromagnetic fields. The fields returned are
# located at the center of the cell. The fields returned are in accordance with
# the time level of the simulation. i.e:
# On the n-th step it returns (E1^{n}, E2^{n}, E3^{n}, B1^{n}, B2^{n}, B3^{n})
# On the (n+1/2)-th step it returns (E1^{n+1/2}, E2^{n+1/2}, E3^{n+1/2}, B1^{n+1/2}, B2^{n+1/2}, B3^{n+1/2})
self._C_p1 = af.broadcast(self._C_p, self.time_elapsed, self.q1_center,
self.q2_center, self.p1_left, self.p2_left,
self.p3_left, E1, E2, E3, B1, B2, B3,
self.physical_system.params)[0]
self._C_p2 = af.broadcast(self._C_p, self.time_elapsed, self.q1_center,
self.q2_center, self.p1_bottom,
self.p2_bottom, self.p3_bottom, E1, E2, E3,
B1, B2, B3, self.physical_system.params)[1]
self._C_p3 = af.broadcast(self._C_p, self.time_elapsed, self.q1_center,
self.q2_center, self.p1_back, self.p2_back,
self.p3_back, E1, E2, E3, B1, B2, B3,
self.physical_system.params)[2]
self._C_p1 = self._convert_to_p_expanded(self._C_p1)
self._C_p2 = self._convert_to_p_expanded(self._C_p2)
self._C_p3 = self._convert_to_p_expanded(self._C_p3)
f = self._convert_to_p_expanded(f)
# Variation of p1 is along axis 0:
f_left_plus_eps, f_right_minus_eps = reconstruct(
self, f, 0, reconstruction_in_p)
# Variation of p2 is along axis 1:
f_bot_plus_eps, f_top_minus_eps = reconstruct(self, f, 1,
reconstruction_in_p)
# Variation of p3 is along axis 2:
f_back_plus_eps, f_front_minus_eps = reconstruct(
self, f, 2, reconstruction_in_p)
# f_left_minus_eps of i-th cell is f_right_minus_eps of the (i-1)th cell
f_left_minus_eps = af.shift(f_right_minus_eps, 1)
# Extending the same to bot:
f_bot_minus_eps = af.shift(f_top_minus_eps, 0, 1)
# Extending the same to back:
f_back_minus_eps = af.shift(f_front_minus_eps, 0, 0, 1)
# flipping due to strange error seen when working with
# multiple species on the CPU backend where f != flip(flip(f))
# Doesn't seem to be a problem on CUDA:
# Yet to test on OpenCL
# TODO: Find exact cause for this bug.
# f != flip(flip(f)) seems to happen after conversion to p_expanded
# flux_p1 = self._C_p1 * af.flip(af.flip(f))
# flux_p2 = self._C_p2 * af.flip(af.flip(f))
# flux_p3 = self._C_p3 * af.flip(af.flip(f))
f_left_p1 = riemann_solver(self, f_left_minus_eps, f_left_plus_eps,
self._C_p1)
f_bot_p2 = riemann_solver(self, f_bot_minus_eps, f_bot_plus_eps,
self._C_p2)
f_back_p3 = riemann_solver(self, f_back_minus_eps, f_back_plus_eps,
self._C_p3)
left_flux_p1 = multiply(self._C_p1, f_left_p1)
bot_flux_p2 = multiply(self._C_p2, f_bot_p2)
back_flux_p3 = multiply(self._C_p3, f_back_p3)
right_flux_p1 = af.shift(left_flux_p1, -1)
top_flux_p2 = af.shift(bot_flux_p2, 0, -1)
front_flux_p3 = af.shift(back_flux_p3, 0, 0, -1)
left_flux_p1 = self._convert_to_q_expanded(left_flux_p1)
right_flux_p1 = self._convert_to_q_expanded(right_flux_p1)
bot_flux_p2 = self._convert_to_q_expanded(bot_flux_p2)
top_flux_p2 = self._convert_to_q_expanded(top_flux_p2)
back_flux_p3 = self._convert_to_q_expanded(back_flux_p3)
front_flux_p3 = self._convert_to_q_expanded(front_flux_p3)
d_flux_p1_dp1 = multiply((right_flux_p1 - left_flux_p1), 1 / self.dp1)
d_flux_p2_dp2 = multiply((top_flux_p2 - bot_flux_p2), 1 / self.dp2)
d_flux_p3_dp3 = multiply((front_flux_p3 - back_flux_p3), 1 / self.dp3)
df_dt += -(d_flux_p1_dp1 + d_flux_p2_dp2 + d_flux_p3_dp3)
if (term_to_return == 'd_flux_p1_dp1'):
af.eval(d_flux_p1_dp1)
return (d_flux_p1_dp1)
elif (term_to_return == 'd_flux_p2_dp2'):
af.eval(d_flux_p2_dp2)
return (d_flux_p2_dp2)
if (term_to_return == 'd_flux_p3_dp3'):
af.eval(d_flux_p3_dp3)
return (d_flux_p3_dp3)
else:
af.eval(df_dt)
return (df_dt)
def df_hat_dt(f_hat, fields_hat, self):
"""
Returns the value of the derivative of the f_hat with respect to time
respect to time. This is used to evolve the system in time.
Input:
------
f_hat : Fourier mode values for the distribution function at which the slope is computed
At t = 0 the initial state of the system is passed to this function:
fields_hat : Fourier mode values for the fields at which the slope is computed
At t = 0 the initial state of the system is passed to this function:
Output:
-------
df_dt : The time-derivative of f_hat
"""
(A_q1, A_q2) = self._A_q(self.time_elapsed, self.q1_center, self.q2_center,
self.p1_center, self.p2_center, self.p3_center,
self.physical_system.params)
df_hat_dt = -1j * (multiply(self.k_q1, A_q1) +
multiply(self.k_q2, A_q2)) * f_hat
if (self.physical_system.params.source_enabled == True):
# Scaling Appropriately:
f = af.real(ifft2(0.5 * self.N_q2 * self.N_q1 * f_hat))
C_f_hat = 2 * fft2(
self._source(
f, self.time_elapsed, self.q1_center, self.q2_center,
self.p1_center, self.p2_center, self.p3_center,
self.compute_moments,
self.physical_system.params)) / (self.N_q2 * self.N_q1)
df_hat_dt += C_f_hat
if (self.physical_system.params.fields_enabled == True):
if (self.physical_system.params.fields_type == 'electrostatic'):
rho_hat = multiply(self.physical_system.params.charge,
self.compute_moments('density', f_hat=f_hat))
self.fields_solver.compute_electrostatic_fields(rho_hat)
elif (self.physical_system.params.fields_type == 'electrodynamic'):
# Handled by dfields_hat_dt
pass
# Used in debugging; advection tests for p-space where fields are to be held constant
elif (self.physical_system.params.fields_type == 'None'):
pass
else:
raise NotImplementedError('Invalid option for fields solver!')
# get_fields for linear solver returns the mode amplitudes of the fields
# So, we obtain A_p1_hat, A_p2_hat, A_p3_hat
(A_p1_hat, A_p2_hat, A_p3_hat) = af.broadcast(
self._A_p, self.time_elapsed, self.q1_center, self.q2_center,
self.p1_center, self.p2_center, self.p3_center, self.fields_solver,
self.physical_system.params)
fields_term = multiply(A_p1_hat, self.dfdp1_background) \
+ multiply(A_p2_hat, self.dfdp2_background) \
+ multiply(A_p3_hat, self.dfdp3_background)
# Including the mean magnetic field term:
# TODO: Maybe wrong to add mean magnetic fields since e^{ikx} term won't cancel out. Check again
# Multiplying by q_center**0 to get the values in the array of required dimension
# Dividing by 2 to normalize appropriately(look at the initialization sector):
B1_mean = af.mean(
self.fields_solver.fields_hat[3, 0, 0, 0]) * self.q1_center**0 / 2
B2_mean = af.mean(
self.fields_solver.fields_hat[4, 0, 0, 0]) * self.q1_center**0 / 2
B3_mean = af.mean(
self.fields_solver.fields_hat[5, 0, 0, 0]) * self.q1_center**0 / 2
e = self.physical_system.params.charge
m = self.physical_system.params.mass
# Converting delta_f array to velocity_expanded form:
f_hat_p_expanded = af.moddims(f_hat, self.N_p1, self.N_p2, self.N_p3,
self.N_species * self.N_q1 * self.N_q2)
# Computing ddelta_f_dp using a 4th order finite different stencil:
ddelta_f_dp1 = multiply((-af.shift(f_hat_p_expanded, -2) +
8 * af.shift(f_hat_p_expanded, -1) +
af.shift(f_hat_p_expanded, 2) -
8 * af.shift(f_hat_p_expanded, 1)),
1 / (12 * self.dp1))
ddelta_f_dp2 = multiply((-af.shift(f_hat_p_expanded, 0, -2) +
8 * af.shift(f_hat_p_expanded, 0, -1) +
af.shift(f_hat_p_expanded, 0, 2) -
8 * af.shift(f_hat_p_expanded, 0, 1)),
1 / (12 * self.dp2))
ddelta_f_dp3 = multiply((-af.shift(f_hat_p_expanded, 0, 0, -2) +
8 * af.shift(f_hat_p_expanded, 0, 0, -1) +
af.shift(f_hat_p_expanded, 0, 0, 2) -
8 * af.shift(f_hat_p_expanded, 0, 0, 1)),
1 / (12 * self.dp3))
# Converting back to positions expanded:
ddelta_f_dp1 = af.moddims(ddelta_f_dp1,
self.N_p1 * self.N_p2 * self.N_p3,
self.N_species, self.N_q1, self.N_q2)
ddelta_f_dp2 = af.moddims(ddelta_f_dp2,
self.N_p1 * self.N_p2 * self.N_p3,
self.N_species, self.N_q1, self.N_q2)
ddelta_f_dp3 = af.moddims(ddelta_f_dp3,
self.N_p1 * self.N_p2 * self.N_p3,
self.N_species, self.N_q1, self.N_q2)
fields_term_mean_magnetic_fields = \
multiply(e/m, ( (multiply(self.p2_center, B3_mean) - multiply(self.p3_center, B2_mean)) * ddelta_f_dp1
+ (multiply(self.p3_center, B1_mean) - multiply(self.p1_center, B3_mean)) * ddelta_f_dp2
+ (multiply(self.p1_center, B2_mean) - multiply(self.p2_center, B1_mean)) * ddelta_f_dp3
)
)
df_hat_dt -= fields_term + fields_term_mean_magnetic_fields
af.eval(df_hat_dt)
return (df_hat_dt)
def f_interp_p_3d(self, dt):
"""
Performs 3D interpolation in the p-space to solve for the equation:
df/dt + A_p1 df/dp1 + A_p2 df/dp2 + A_p3 df/dp3 = 0
Similar to the above function, this is done by backtracing the
characteristic curves and interpolating at the origin of the characteristics.
Parameters
----------
dt : double
Time-step size to evolve the system
NOTE: This function currently makes use of a Strang split approx1, approx2 with
reorders to apply along the intended axes. With implementation of approx3
complete this would be changed to make use of a single call of approx3.
Ref:https://github.com/arrayfire/arrayfire/issues/1837
"""
# Following Strang Splitting:
if (self.performance_test_flag == True):
tic = af.time()
A_p1 = af.broadcast(self._A_p, self.time_elapsed, self.q1_center,
self.q2_center, self.p1_left, self.p2_left,
self.p3_left, self.fields_solver,
self.physical_system.params)[0]
A_p2 = af.broadcast(self._A_p, self.time_elapsed, self.q1_center,
self.q2_center, self.p1_bottom, self.p2_bottom,
self.p3_bottom, self.fields_solver,
self.physical_system.params)[1]
A_p3 = af.broadcast(self._A_p, self.time_elapsed, self.q1_center,
self.q2_center, self.p1_back, self.p2_back,
self.p3_back, self.fields_solver,
self.physical_system.params)[2]
# # Using the add method wrapped with af.broadcast
# p1_new = add(self.p1_center, - dt * A_p1)
# p2_new = add(self.p2_center, - dt * A_p2)
# # Since the interpolation function are being performed in velocity space,
# # the arrays used in the computation need to be in p_expanded form.
# # Hence we will need to convert the same:
# p1_new = self._convert_to_p_expanded(p1_new)
# p2_new = self._convert_to_p_expanded(p2_new)
# # Transforming interpolant to go from [0, N_p - 1]:
# p1_lower_boundary = add(self.p1_start, 0.5 * self.dp1)
# p2_lower_boundary = add(self.p2_start, 0.5 * self.dp2)
# p1_interpolant = multiply(add(p1_new, -p1_lower_boundary), 1 / self.dp1)
# p2_interpolant = multiply(add(p2_new, -p2_lower_boundary), 1 / self.dp2)
# if(self.physical_system.params.p_dim == 3):
# p3_new = add(self.p3_center, - 0.5 * dt * A_p3)
# p3_new = self._convert_to_p_expanded(p3_new)
# p3_lower_boundary = add(self.p3_start, 0.5 * self.dp3)
# # Reordering from (N_p1, N_p2, N_p3, N_s * N_q) --> (N_p3, N_p1, N_p2, N_s * N_q)
# p3_interpolant = af.reorder(mulitiply(add(p3_new, -p3_lower_boundary), 1 / self.dp3), 2, 0, 1, 3)
# # We perform the 3d interpolation by performing individual 1d + 2d interpolations:
# self.f = self._convert_to_p_expanded(self.f)
# if(self.physical_system.params.p_dim == 3):
# # Reordering from (N_p1, N_p2, N_p3, N_s * N_q) --> (N_p3, N_p1, N_p2, N_s * N_q)
# self.f = af.approx1(af.reorder(self.f, 2, 0, 1, 3),
# p3_interpolant,
# af.INTERP.CUBIC_SPLINE
# )
# # Reordering back from (N_p1, N_p2, N_p3, N_s * N_q) --> (N_p3, N_p1, N_p2, N_s * N_q)
# self.f = af.reorder(self.f, 1, 2, 0, 3)
# self.f = af.approx2(self.f,
# p1_interpolant,
# p2_interpolant,
# af.INTERP.BICUBIC_SPLINE
# )
# if(self.physical_system.params.p_dim == 3):
# # Reordering from (N_p1, N_p2, N_p3, N_s * N_q) --> (N_p3, N_p1, N_p2, N_s * N_q)
# self.f = af.approx1(af.reorder(self.f, 2, 0, 1, 3),
# p3_interpolant,
# af.INTERP.CUBIC_SPLINE
# )
# # Reordering back from (N_p3, N_p1, N_p2, N_s * N_q) --> (N_p1, N_p2, N_p3, N_s * N_q)
# self.f = af.reorder(self.f, 1, 2, 0, 3)
# Using the add method wrapped with af.broadcast
p2_new = add(self.p2_center, -dt * A_p2)
p3_new = add(self.p3_center, -dt * A_p3)
# Since the interpolation function are being performed in velocity space,
# the arrays used in the computation need to be in p_expanded form.
# Hence we will need to convert the same:
p2_new = self._convert_to_p_expanded(p2_new)
p3_new = self._convert_to_p_expanded(p3_new)
# Transforming interpolant to go from [0, N_p - 1]:
p2_lower_boundary = add(self.p2_start, 0.5 * self.dp2)
p3_lower_boundary = add(self.p3_start, 0.5 * self.dp3)
p2_interpolant = multiply(add(p2_new, -p2_lower_boundary), 1 / self.dp2)
p3_interpolant = multiply(add(p3_new, -p3_lower_boundary), 1 / self.dp3)
self.f = self._convert_to_p_expanded(self.f)
self.f = af.approx2(af.reorder(self.f, 1, 2, 0, 3),
af.reorder(p2_interpolant, 1, 2, 0, 3),
af.reorder(p3_interpolant, 1, 2, 0, 3),
af.INTERP.BICUBIC_SPLINE)
# Reordering back from (N_p2, N_p3, N_p1, N_s * N_q) --> (N_p1, N_p2, N_p3, N_s * N_q)
self.f = af.reorder(self.f, 2, 0, 1, 3)
self.f = self._convert_to_q_expanded(self.f)
af.eval(self.f)
if (self.performance_test_flag == True):
af.sync()
toc = af.time()
self.time_interp3 += toc - tic
return