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)