Skip to content

Commit

Permalink
Formatting
Browse files Browse the repository at this point in the history
  • Loading branch information
@AngelFP
AngelFP committed Jun 18, 2024
1 parent 8f1d43f commit 9e2d5e6
Show file tree
Hide file tree
Showing 12 changed files with 941 additions and 502 deletions.
2 changes: 1 addition & 1 deletion wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/__init__.py
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,3 @@
from .wakefield import Quasistatic2DWakefieldIon

__all__ = ['Quasistatic2DWakefieldIon']
__all__ = ["Quasistatic2DWakefieldIon"]
177 changes: 112 additions & 65 deletions wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/adaptive_grid.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
from .utils import longitudinal_gradient, radial_gradient


class AdaptiveGrid():
class AdaptiveGrid:
"""Grid whose size dynamically adapts to the extent of a particle bunch.
The number of radial cells is fixed, but its transverse extent is
Expand Down Expand Up @@ -49,6 +49,7 @@ class AdaptiveGrid():
deposit to and gather from the base grid (if they haven't escaped
from it too).
"""

def __init__(
self,
x: np.ndarray,
Expand All @@ -59,7 +60,7 @@ def __init__(
nxi: int,
xi_plasma: np.ndarray,
r_max: Optional[float] = None,
r_lim: Optional[float] = None
r_lim: Optional[float] = None,
):
self.bunch_name = bunch_name
self.xi_plasma = xi_plasma
Expand All @@ -72,7 +73,7 @@ def __init__(
self._r_max_hist = []

self._update(x, y, xi)

@property
def r_min_cell(self):
"""Radial position of first cell (ignoring guard cells)."""
Expand Down Expand Up @@ -111,9 +112,8 @@ def update_if_needed(self, x, y, xi, n_p, pp_hist):
# Only trigger radial update if the radial size is not fixed.
if self._r_max is None:
r_max_beam = np.max(np.sqrt(x**2 + y**2))
update_r = (
(r_max_beam > self.r_max_cell) or
(r_max_beam < self.r_max_cell * 0.9)
update_r = (r_max_beam > self.r_max_cell) or (
r_max_beam < self.r_max_cell * 0.9
)
# It a radial limit is set, update only if limit has not been
# reached.
Expand All @@ -127,9 +127,8 @@ def update_if_needed(self, x, y, xi, n_p, pp_hist):

xi_min_beam = np.min(xi)
xi_max_beam = np.max(xi)
update_xi = (
(xi_min_beam < self.xi_grid[0 + self.nxi_border]) or
(xi_max_beam > self.xi_grid[-1 - self.nxi_border])
update_xi = (xi_min_beam < self.xi_grid[0 + self.nxi_border]) or (
xi_max_beam > self.xi_grid[-1 - self.nxi_border]
)
if update_r or update_xi:
self._update(x, y, xi)
Expand All @@ -152,19 +151,30 @@ def calculate_fields(self, n_p, pp_hist, reset_fields=True):
self._reset_fields()
s_d = ge.plasma_skin_depth(n_p * 1e-6)
calculate_fields_on_grid(
self.i_grid, self.r_grid, s_d,
self.psi_grid, self.b_t, pp_hist['r_hist'], pp_hist['log_r_hist'],
pp_hist['sum_1_hist'], pp_hist['sum_2_hist'],
pp_hist['a_0_hist'], pp_hist['a_i_hist'], pp_hist['b_i_hist'])
self.i_grid,
self.r_grid,
s_d,
self.psi_grid,
self.b_t,
pp_hist["r_hist"],
pp_hist["log_r_hist"],
pp_hist["sum_1_hist"],
pp_hist["sum_2_hist"],
pp_hist["a_0_hist"],
pp_hist["a_i_hist"],
pp_hist["b_i_hist"],
)

E_0 = ge.plasma_cold_non_relativisct_wave_breaking_field(n_p * 1e-6)
longitudinal_gradient(
self.psi_grid[2:-2, 2:-2], self.dxi/s_d, self.e_z[2:-2, 2:-2])
self.psi_grid[2:-2, 2:-2], self.dxi / s_d, self.e_z[2:-2, 2:-2]
)
radial_gradient(
self.psi_grid[2:-2, 2:-2], self.dr/s_d, self.e_r[2:-2, 2:-2])
self.psi_grid[2:-2, 2:-2], self.dr / s_d, self.e_r[2:-2, 2:-2]
)
self.e_r -= self.b_t
self.e_z *= - E_0
self.e_r *= - E_0
self.e_z *= -E_0
self.e_r *= -E_0
self.b_t *= E_0 / ct.c

def calculate_bunch_source(self, bunch, n_p, p_shape):
Expand All @@ -179,12 +189,23 @@ def calculate_bunch_source(self, bunch, n_p, p_shape):
p_shape : str
The particle shape.
"""
self.b_t_bunch[:] = 0.
self.q_bunch[:] = 0.
self.b_t_bunch[:] = 0.0
self.q_bunch[:] = 0.0
all_deposited = deposit_bunch_charge(
bunch.x, bunch.y, bunch.xi, bunch.q, n_p,
self.nr-self.nr_border, self.nxi, self.r_grid, self.xi_grid,
self.dr, self.dxi, p_shape, self.q_bunch)
bunch.x,
bunch.y,
bunch.xi,
bunch.q,
n_p,
self.nr - self.nr_border,
self.nxi,
self.r_grid,
self.xi_grid,
self.dr,
self.dxi,
p_shape,
self.q_bunch,
)
calculate_bunch_source(self.q_bunch, self.nr, self.nxi, self.b_t_bunch)
return all_deposited

Expand All @@ -199,9 +220,25 @@ def gather_fields(self, x, y, z, ex, ey, ez, bx, by, bz):
The arrays where the gathered field components will be stored.
"""
return gather_main_fields_cyl_linear(
self.e_r, self.e_z, self.b_t, self.xi_min, self.xi_max,
self.r_min_cell, self.r_max_cell, self.dxi, self.dr, x, y, z,
ex, ey, ez, bx, by, bz)
self.e_r,
self.e_z,
self.b_t,
self.xi_min,
self.xi_max,
self.r_min_cell,
self.r_max_cell,
self.dxi,
self.dr,
x,
y,
z,
ex,
ey,
ez,
bx,
by,
bz,
)

def get_openpmd_data(self, global_time, diags):
"""Get the field data at the grid to store in the openpmd diagnostics.
Expand All @@ -221,8 +258,8 @@ def get_openpmd_data(self, global_time, diags):

# Grid parameters.
grid_spacing = [self.dr, self.dxi]
grid_labels = ['r', 'z']
grid_global_offset = [0., global_time*ct.c + self.xi_min]
grid_labels = ["r", "z"]
grid_global_offset = [0.0, global_time * ct.c + self.xi_min]

# Initialize field diags lists.
names = []
Expand All @@ -231,45 +268,45 @@ def get_openpmd_data(self, global_time, diags):
arrays = []

# Add requested fields to lists.
if 'E' in diags:
names += ['E']
comps += [['r', 'z']]
if "E" in diags:
names += ["E"]
comps += [["r", "z"]]
attrs += [{}]
arrays += [
[np.ascontiguousarray(self.e_r.T[2:-2, 2:-2]),
np.ascontiguousarray(self.e_z.T[2:-2, 2:-2])]
[
np.ascontiguousarray(self.e_r.T[2:-2, 2:-2]),
np.ascontiguousarray(self.e_z.T[2:-2, 2:-2]),
]
]
if 'B' in diags:
names += ['B']
comps += [['t']]
if "B" in diags:
names += ["B"]
comps += [["t"]]
attrs += [{}]
arrays += [
[np.ascontiguousarray(self.b_t.T[2:-2, 2:-2])]
]
arrays += [[np.ascontiguousarray(self.b_t.T[2:-2, 2:-2])]]

# Create dictionary with all diagnostics data.
comp_pos = [[0.5, 0.]] * len(names)
comp_pos = [[0.5, 0.0]] * len(names)
fld_zip = zip(names, comps, attrs, arrays, comp_pos)
diag_data = {}
diag_data['fields'] = []
diag_data["fields"] = []
for fld, comps, attrs, arrays, pos in fld_zip:
fld += '_' + self.bunch_name
diag_data['fields'].append(fld)
fld += "_" + self.bunch_name
diag_data["fields"].append(fld)
diag_data[fld] = {}
if comps is not None:
diag_data[fld]['comps'] = {}
diag_data[fld]["comps"] = {}
for comp, arr in zip(comps, arrays):
diag_data[fld]['comps'][comp] = {}
diag_data[fld]['comps'][comp]['array'] = arr
diag_data[fld]['comps'][comp]['position'] = pos
diag_data[fld]["comps"][comp] = {}
diag_data[fld]["comps"][comp]["array"] = arr
diag_data[fld]["comps"][comp]["position"] = pos
else:
diag_data[fld]['array'] = arrays[0]
diag_data[fld]['position'] = pos
diag_data[fld]['grid'] = {}
diag_data[fld]['grid']['spacing'] = grid_spacing
diag_data[fld]['grid']['labels'] = grid_labels
diag_data[fld]['grid']['global_offset'] = grid_global_offset
diag_data[fld]['attributes'] = attrs
diag_data[fld]["array"] = arrays[0]
diag_data[fld]["position"] = pos
diag_data[fld]["grid"] = {}
diag_data[fld]["grid"]["spacing"] = grid_spacing
diag_data[fld]["grid"]["labels"] = grid_labels
diag_data[fld]["grid"]["global_offset"] = grid_global_offset
diag_data[fld]["attributes"] = attrs

return diag_data

Expand All @@ -285,14 +322,14 @@ def _update(self, x, y, xi):
self._r_max_hist.append(r_max)
self.dr = r_max / (self.nr - self.nr_border)
r_max += self.nr_border * self.dr
self.r_grid = np.linspace(self.dr/2, r_max - self.dr/2, self.nr)
self.r_grid = np.linspace(self.dr / 2, r_max - self.dr / 2, self.nr)

# Create grid in xi
xi_min_beam = np.min(xi)
xi_max_beam = np.max(xi)
self.i_grid = np.where(
(self.xi_plasma > xi_min_beam - self.dxi * (1 + self.nxi_border)) &
(self.xi_plasma < xi_max_beam + self.dxi * (1 + self.nxi_border))
(self.xi_plasma > xi_min_beam - self.dxi * (1 + self.nxi_border))
& (self.xi_plasma < xi_max_beam + self.dxi * (1 + self.nxi_border))
)[0]
self.xi_grid = self.xi_plasma[self.i_grid]
self.xi_max = self.xi_grid[-1]
Expand All @@ -309,17 +346,27 @@ def _update(self, x, y, xi):

def _reset_fields(self):
"""Reset value of the fields at the grid."""
self.psi_grid[:] = 0.
self.b_t[:] = 0.
self.e_r[:] = 0.
self.e_z[:] = 0.
self.psi_grid[:] = 0.0
self.b_t[:] = 0.0
self.e_r[:] = 0.0
self.e_z[:] = 0.0


@njit_serial()
def calculate_fields_on_grid(
i_grid, r_grid, s_d,
psi_grid, bt_grid, r_hist, log_r_hist, sum_1_hist, sum_2_hist,
a_0_hist, a_i_hist, b_i_hist):
i_grid,
r_grid,
s_d,
psi_grid,
bt_grid,
r_hist,
log_r_hist,
sum_1_hist,
sum_2_hist,
a_0_hist,
a_i_hist,
b_i_hist,
):
"""Compute the plasma fields on the grid.
Compiling this method in numba avoids significant overhead.
Expand All @@ -336,7 +383,7 @@ def calculate_fields_on_grid(
log_r=log_r_hist[j, :n_elec],
sum_1_arr=sum_1_hist[j, :n_elec + 1],
sum_2_arr=sum_2_hist[j, :n_elec + 1],
psi=psi
psi=psi,
)
calculate_psi_with_interpolation(
r_eval=r_grid / s_d,
Expand All @@ -353,5 +400,5 @@ def calculate_fields_on_grid(
a=a_i_hist[j],
b=b_i_hist[j],
r=r_hist[j, :n_elec],
b_theta=b_theta
b_theta=b_theta,
)
Loading

0 comments on commit 9e2d5e6

Please sign in to comment.