def test_has_refs():
class StructWRef(xo.Struct):
a = xo.Ref(xo.Float64[:])
assert StructWRef._has_refs
class StructNoRef(xo.Struct):
a = xo.Float64[:]
assert not StructNoRef._has_refs
class NestedWRef(xo.Struct):
s = StructWRef
assert NestedWRef._has_refs
class NestedNoRef(xo.Struct):
s = StructNoRef
assert not NestedNoRef._has_refs
ArrNoRef = xo.Float64[:]
assert not ArrNoRef._has_refs
ArrWRef = xo.Ref(xo.Float64)[:]
assert ArrWRef._has_refs
class StructArrRef(xo.Struct):
arr = ArrWRef
assert StructArrRef._has_refs
class StructArrNoRef(xo.Struct):
arr = ArrNoRef
assert not StructArrNoRef._has_refs
ArrOfStructRef = NestedWRef[:]
assert ArrOfStructRef._has_refs
class MyUnion(xo.UnionRef):
_ref = [xo.Float64, xo.Int32]
assert MyUnion._has_refs
class MyStruct2(xo.Struct):
a = xo.Float64[:]
sr = xo.Ref(MyStruct)
class StructWRef(xo.Struct):
a = xo.Ref(xo.Float64[:])
class MyStructRef(xo.Struct):
a = xo.Ref(xo.Float64[:], xo.String)
class SpaceCharge3D(xt.BeamElement):
"""
Simulates the effect of space charge on a bunch.
Args:
context (XfContext): identifies the :doc:`context <contexts>`
on which the computation is executed.
update_on_track (bool): If ``True`` the beam field map is update
at each interaction. If ``False`` the initial field map is
used at each interaction (frozen model). The default is
``True``.
length (float): the length of the space-charge interaction in
meters.
apply_z_kick (bool): If ``True``, the longitudinal kick on the
particles is applied.
x_range (tuple): Horizontal extent (in meters) of the
computing grid.
y_range (tuple): Vertical extent (in meters) of the
computing grid.
z_range (tuple): Longitudina extent (in meters) of
the computing grid.
nx (int): Number of cells in the horizontal direction.
ny (int): Number of cells in the vertical direction.
nz (int): Number of cells in the vertical direction.
dx (float): Horizontal cell size in meters. It can be
provided alternatively to ``nx``.
dy (float): Vertical cell size in meters. It can be
provided alternatively to ``ny``.
dz (float): Longitudinal cell size in meters.It can be
provided alternatively to ``nz``.
x_grid (np.ndarray): Equispaced array with the horizontal grid points
(cell centers).
It can be provided alternatively to ``x_range``, ``dx``/``nx``.
y_grid (np.ndarray): Equispaced array with the horizontal grid points
(cell centers).
It can be provided alternatively to ``y_range``, ``dy``/``ny``.
z_grid (np.ndarray): Equispaced array with the horizontal grid points
(cell centers).
It can be provided alternatively to ``z_range``, ``dz``/``nz``.
rho (np.ndarray): initial charge density at the grid points in
Coulomb/m^3.
phi (np.ndarray): initial electric potential at the grid points in
Volts. If not provided the ``phi`` is calculated from ``rho``
using the Poisson solver (if available).
solver (str or solver object): Defines the Poisson solver to be used
to compute phi from rho. Accepted values are ``FFTSolver3D`` and
``FFTSolver2p5D``. A Xfields solver object can also be provided.
In case ``update_on_track``is ``False`` and ``phi`` is provided
by the user, this argument can be omitted.
gamma0 (float): Relativistic gamma factor of the beam. This is required
only if the solver is ``FFTSolver3D``.
Returns:
(SpaceCharge3D): A space-charge 3D beam element.
"""
_xofields = {
'fieldmap': xo.Ref(TriLinearInterpolatedFieldMap._XoStruct),
'length': xo.Float64,
}
_extra_c_sources = [
_pkg_root.joinpath('headers/constants.h'),
_pkg_root.joinpath(
'fieldmaps/interpolated_src/linear_interpolators.h'),
_pkg_root.joinpath('beam_elements/spacecharge_src/spacecharge3d.h'),
]
def copy(self, _context=None, _buffer=None, _offset=None):
if _buffer is not self._buffer:
raise NotImplementedError
return SpaceCharge3D(_context=_context,
_buffer=_buffer,
_offset=_offset,
update_on_track=self.update_on_track,
length=self.length,
apply_z_kick=self.apply_z_kick,
fieldmap=self.fieldmap)
def __init__(self,
_context=None,
_buffer=None,
_offset=None,
update_on_track=True,
length=None,
apply_z_kick=True,
fieldmap=None,
x_range=None,
y_range=None,
z_range=None,
nx=None,
ny=None,
nz=None,
dx=None,
dy=None,
dz=None,
x_grid=None,
y_grid=None,
z_grid=None,
rho=None,
phi=None,
solver=None,
gamma0=None,
fftplan=None):
self.update_on_track = update_on_track
self.apply_z_kick = apply_z_kick
if solver == 'FFTSolver3D':
assert gamma0 is not None, ('To use FFTSolver3D '
'gamma0 must be provided')
if gamma0 is not None:
if not np.isscalar(gamma0):
raise ValueError('gamma0 needs to be a scalar')
scale_coordinates_in_solver = (1., 1., float(gamma0))
else:
scale_coordinates_in_solver = (1., 1., 1.)
if _buffer is not None:
_context = _buffer.context
if _context is None:
_context = xo.context_default
if fieldmap is None:
# I build the fieldmap on a temporary buffer
temp_buff = _context.new_buffer()
fieldmap = TriLinearInterpolatedFieldMap(
_buffer=temp_buff,
rho=rho,
phi=phi,
x_grid=z_grid,
y_grid=y_grid,
z_grid=z_grid,
x_range=x_range,
y_range=y_range,
z_range=z_range,
dx=dx,
dy=dy,
dz=dz,
nx=nx,
ny=ny,
nz=nz,
solver=solver,
scale_coordinates_in_solver=scale_coordinates_in_solver,
updatable=update_on_track,
fftplan=fftplan)
self.xoinitialize(_context=_context,
_buffer=_buffer,
_offset=_offset,
fieldmap=fieldmap,
length=length)
# temp_buff is deallocate here
@property
def iscollective(self):
return self.update_on_track
def track(self, particles):
"""
Computes and applies the space-charge forces for the provided set of
particles.
Args:
particles (Particles Object): Particles to be tracked.
"""
if self.update_on_track:
self.fieldmap.update_from_particles(particles=particles)
# call C tracking kernel
super().track(particles)