def sim_full():
return Simulation(
TimeGrid(200, 2, 20),
MeshGrid(5, 51),
particle_sources=[
ParticleSource('a', Box()),
ParticleSource('c', Cylinder()),
ParticleSource('d', Tube(start=(0, 0, 0), end=(0, 0, 1)))
],
inner_regions=[
InnerRegion('1', Box(), 1),
InnerRegion('2', Sphere(), -2),
InnerRegion('3', Cylinder(), 0),
InnerRegion('4', Tube(), 4)
],
particle_interaction_model=Model.binary,
electric_fields=[FieldUniform('x', 'electric', (-2, -2, 1))],
magnetic_fields=[
FieldExpression('y', 'magnetic', '0', '0', '3*x + sqrt(y) - z**2')
])
def test_init_rhs(self):
mesh = SpatialMeshConf((4, 3, 3)).make(BoundaryConditionsConf())
solver = FieldSolver(mesh, [])
solver.init_rhs_vector_in_full_domain(mesh)
assert_array_equal(solver.rhs, np.zeros(3 * 2 * 2))
mesh = SpatialMeshConf((4, 3, 3)).make(BoundaryConditionsConf(-2))
solver = FieldSolver(mesh, [])
solver.init_rhs_vector_in_full_domain(mesh)
assert_array_equal(solver.rhs,
[6, 4, 6, 6, 4, 6, 6, 4, 6, 6, 4, 6]) # what
mesh = SpatialMeshConf((4, 4, 5)).make(BoundaryConditionsConf(-2))
solver = FieldSolver(mesh, [])
solver.init_rhs_vector_in_full_domain(mesh)
assert_array_equal(solver.rhs, [
6, 4, 6, 4, 2, 4, 6, 4, 6, 4, 2, 4, 2, 0, 2, 4, 2, 4, 4, 2, 4, 2,
0, 2, 4, 2, 4, 6, 4, 6, 4, 2, 4, 6, 4, 6
]) # what
mesh = SpatialMeshConf((8, 12, 5),
(2, 3, 1)).make(BoundaryConditionsConf(-1))
solver = FieldSolver(mesh, [])
solver.init_rhs_vector_in_full_domain(mesh)
assert_array_equal(solver.rhs, [
49, 40, 49, 45, 36, 45, 49, 40, 49, 13, 4, 13, 9, 0, 9, 13, 4, 13,
13, 4, 13, 9, 0, 9, 13, 4, 13, 49, 40, 49, 45, 36, 45, 49, 40, 49
])
mesh = SpatialMeshConf((4, 6, 9),
(1, 2, 3)).make(BoundaryConditionsConf())
solver = FieldSolver(mesh, [])
mesh.charge_density = np.array([[[0, 0, 0, 0], [0, 0, 0, 0],
[0, 0, 0, 0], [0, 0, 0, 0]],
[[0, 0, 0, 0], [0, 1, 2, 0],
[0, -1, 0, 0], [0, 0, 0, 0]],
[[0, 0, 0, 0], [0, 3, 4, 0],
[0, 0, -1, 0], [0, 0, 0, 0]],
[[0, 0, 0, 0], [0, 5, 6, 0],
[0, -1, 0, 0], [0, 0, 0, 0]],
[[0, 0, 0, 0], [0, 0, 0, 0],
[0, 0, 0, 0], [0, 0, 0, 0]]])
solver.init_rhs_vector_in_full_domain(mesh)
assert_allclose(
solver.rhs, -np.array([1, 3, 5, -1, 0, -1, 2, 4, 6, 0, -1, 0]) *
np.pi * 4 * 36)
mesh = SpatialMeshConf((4, 6, 9),
(1, 2, 3)).make(BoundaryConditionsConf())
solver = FieldSolver(mesh, [])
region = InnerRegion('test', Box((1, 2, 3), (1, 2, 3)), 3)
solver.init_rhs_vector(mesh, [region])
assert_array_equal(solver.rhs, [3, 3, 0, 3, 3, 0, 3, 3, 0, 3, 3, 0])
def import_from_h5(h5file: h5py.File) -> Simulation:
fields = [Field.import_h5(g) for g in h5file['ExternalFields'].values()]
sources = [ParticleSource.import_h5(g) for g in h5file['ParticleSources'].values()]
particles = [ParticleArray.import_h5(g) for g in h5file['ParticleSources'].values()]
# cupy max has no `initial` argument
max_id = int(max([(p.ids.get() if hasattr(p.ids, 'get') else p.ids).max(initial=-1) for p in particles], default=-1))
g = h5file['SpatialMesh']
mesh = MeshGrid.import_h5(g)
charge = Reader.array_class(mesh, (), np.reshape(g['charge_density'], mesh.n_nodes))
potential = Reader.array_class(mesh, (), np.reshape(g['potential'], mesh.n_nodes))
field = FieldOnGrid('spatial_mesh', 'electric', Reader.array_class(mesh, 3, np.moveaxis(
np.array([np.reshape(g[f'electric_field_{c}'], mesh.n_nodes) for c in 'xyz']),
0, -1)))
return Simulation(
time_grid=TimeGrid.import_h5(h5file['TimeGrid']),
mesh=mesh, charge_density=charge, potential=potential, electric_field=field,
inner_regions=[InnerRegion.import_h5(g) for g in h5file['InnerRegions'].values()],
electric_fields=[f for f in fields if f.electric_or_magnetic == 'electric'],
magnetic_fields=[f for f in fields if f.electric_or_magnetic == 'magnetic'],
particle_interaction_model=Model[
h5file['ParticleInteractionModel'].attrs['particle_interaction_model'].decode('utf8')
],
particle_sources=sources, particle_arrays=particles, particle_tracker=ParticleTracker(max_id)
)
def test_absorb_charge_inverted(self):
particles = ParticleArray([1], -2.0, 1.0, (0, 0, 0), np.zeros(3))
ir = InnerRegion('test', Box(), inverted=True)
assert ir.total_absorbed_particles == 0
assert ir.total_absorbed_charge == 0
ir.collide_with_particles(particles)
assert ir.total_absorbed_particles == 0
assert ir.total_absorbed_charge == 0
assert_dataclass_eq(
particles,
ParticleArray([1], -2.0, 1.0, np.zeros((1, 3)), np.zeros((1, 3))))
particles = ParticleArray([1], -2.0, 1.0, (10, 10, 10), np.zeros(3))
ir = InnerRegion('test', Box(), inverted=True)
assert ir.total_absorbed_particles == 0
assert ir.total_absorbed_charge == 0
ir.collide_with_particles(particles)
assert ir.total_absorbed_particles == 1
assert ir.total_absorbed_charge == -2
assert_dataclass_eq(
particles,
ParticleArray([], -2.0, 1.0, np.zeros((0, 3)), np.zeros((0, 3))))
particles = ParticleArray([1, 2], -2.0, 1.0, [(0, 0, 0), (10, 10, 10)],
np.zeros((2, 3)))
ir = InnerRegion('test', Box(), inverted=True)
assert ir.total_absorbed_particles == 0
assert ir.total_absorbed_charge == 0
ir.collide_with_particles(particles)
assert ir.total_absorbed_particles == 1
assert ir.total_absorbed_charge == -2
assert_dataclass_eq(
particles,
ParticleArray([1], -2.0, 1.0, [(0, 0, 0)], np.zeros((1, 3))))
def test_absorb_charge(self):
particles = ParticleArray([1], -2.0, 1.0, (0, 0, 0), np.zeros(3))
ir = InnerRegion('test', Box())
assert ir.total_absorbed_particles == 0
assert ir.total_absorbed_charge == 0
ir.collide_with_particles(particles)
assert ir.total_absorbed_particles == 1
assert ir.total_absorbed_charge == -2
assert particles == ParticleArray([], -2.0, 1.0, np.zeros((0, 3)), np.zeros((0, 3)))
particles = ParticleArray([1], -2.0, 1.0, (10, 10, 10), np.zeros(3))
ir = InnerRegion('test', Box())
assert ir.total_absorbed_particles == 0
assert ir.total_absorbed_charge == 0
ir.collide_with_particles(particles)
assert ir.total_absorbed_particles == 0
assert ir.total_absorbed_charge == 0
assert particles == ParticleArray([1], -2.0, 1.0, [(10, 10, 10)], np.zeros((1, 3)))
particles = ParticleArray([1, 2], -2.0, 1.0, [(0, 0, 0), (10, 10, 10)], np.zeros((2, 3)))
ir = InnerRegion('test', Box())
assert ir.total_absorbed_particles == 0
assert ir.total_absorbed_charge == 0
ir.collide_with_particles(particles)
assert ir.total_absorbed_particles == 1
assert ir.total_absorbed_charge == -2
assert particles == ParticleArray([2], -2.0, 1.0, [(10, 10, 10)], np.zeros((1, 3)))
def test_zero_nondiag_inside_objects(self):
mesh = SpatialMeshConf((4, 6, 9),
(1, 2, 3)).make(BoundaryConditionsConf())
solver = FieldSolver(mesh, [])
region = InnerRegion('test', Box((1, 2, 3), (1, 2, 3)), 3)
a = csr_matrix(np.full((12, 12), 2))
assert_array_equal(a.toarray(), [[2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
[2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
[2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
[2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
[2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
[2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
[2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
[2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
[2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
[2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
[2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
[2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2]])
result = solver.zero_nondiag_for_nodes_inside_objects(
a, mesh, [region])
assert_array_equal(result.toarray(),
[[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
[0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
[2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
[0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
[2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
[2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2]])
# TODO: check algorithm if on-diagonal zeros should turn into ones
a = csr_matrix(
np.array([[4, 0, 3, 0, 0, 0, 0, 2, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 4, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 2, 0, 0, 3, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 2, 0, 0, 0, 0, 0, 0, 1, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 6, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]))
result = solver.zero_nondiag_for_nodes_inside_objects(
a, mesh, [region])
assert_array_equal(result.toarray(),
[[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 2, 0, 0, 3, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])
class Simulation(SerializableH5):
array_class: Type[ArrayOnGrid] = inject.attr(ArrayOnGrid)
def __init__(self,
time_grid: TimeGrid,
mesh: MeshGrid,
inner_regions: Sequence[InnerRegion] = (),
particle_sources: Sequence[ParticleSource] = (),
electric_fields: Sequence[Field] = (),
magnetic_fields: Sequence[Field] = (),
particle_interaction_model: Model = Model.PIC,
charge_density: Optional[ArrayOnGrid] = None,
potential: Optional[ArrayOnGrid] = None,
electric_field: Optional[FieldOnGrid] = None,
particle_tracker: Optional[ParticleTracker] = None,
particle_arrays: Sequence[ParticleArray] = ()):
super().__init__()
self.time_grid: TimeGrid = time_grid
self.mesh: MeshGrid = mesh
self.charge_density: ArrayOnGrid = self.array_class(
mesh) if charge_density is None else charge_density
self.potential: ArrayOnGrid = self.array_class(
mesh) if potential is None else potential
self.electric_field: FieldOnGrid = FieldOnGrid('spatial_mesh', 'electric', self.array_class(mesh, 3)) \
if electric_field is None else electric_field
self._domain = InnerRegion('simulation_domain',
shapes.Box(0, mesh.size),
inverted=True)
self.inner_regions: List[InnerRegion] = list(inner_regions)
self.particle_sources: List[ParticleSource] = list(particle_sources)
self.electric_fields: Field = FieldSum.factory(electric_fields,
'electric')
self.magnetic_fields: Field = FieldSum.factory(magnetic_fields,
'magnetic')
self.particle_interaction_model: Model = particle_interaction_model
self.particle_arrays: List[ParticleArray] = list(particle_arrays)
self.consolidate_particle_arrays()
if self.particle_interaction_model == Model.binary:
self._dynamic_field = FieldParticles('binary_particle_field',
self.particle_arrays)
if not is_trivial(self.potential, self.inner_regions):
self._dynamic_field += self.electric_field
elif self.particle_interaction_model == Model.noninteracting:
if not is_trivial(self.potential, self.inner_regions):
self._dynamic_field = self.electric_field
else:
self._dynamic_field = FieldZero('Uniform_potential_zero_field',
'electric')
else:
self._dynamic_field = self.electric_field
self.particle_tracker = ParticleTracker(
) if particle_tracker is None else particle_tracker
@property
def dict(self) -> dict:
d = super().dict
d['particle_interaction_model'] = d['particle_interaction_model'].name
return d
def advance_one_time_step(self, field_solver):
self.push_particles()
self.generate_and_prepare_particles(field_solver)
self.time_grid.update_to_next_step()
def eval_charge_density(self):
self.charge_density.reset()
for p in self.particle_arrays:
self.charge_density.distribute_at_positions(p.charge, p.positions)
def push_particles(self):
self.boris_integration(self.time_grid.time_step_size)
def generate_and_prepare_particles(self, field_solver, initial=False):
self.generate_valid_particles(initial)
if self.particle_interaction_model == Model.PIC:
self.eval_charge_density()
field_solver.eval_potential(self.charge_density, self.potential)
self.potential.gradient(self.electric_field.array)
self.shift_new_particles_velocities_half_time_step_back()
self.consolidate_particle_arrays()
def generate_valid_particles(self, initial=False):
# First generate then remove.
# This allows for overlap of source and inner region.
self.generate_new_particles(initial)
self.apply_domain_boundary_conditions()
self.remove_particles_inside_inner_regions()
def boris_integration(self, dt):
for particles in self.particle_arrays:
total_el_field, total_mgn_field = \
self.compute_total_fields_at_positions(particles.positions)
if self.magnetic_fields != 0 and total_mgn_field.any():
particles.boris_update_momentums(dt, total_el_field,
total_mgn_field)
else:
particles.boris_update_momentum_no_mgn(dt, total_el_field)
particles.update_positions(dt)
def prepare_boris_integration(self, minus_half_dt):
# todo: place newly generated particle_arrays into separate buffer
for particles in self.particle_arrays:
if not particles.momentum_is_half_time_step_shifted:
total_el_field, total_mgn_field = \
self.compute_total_fields_at_positions(particles.positions)
if self.magnetic_fields != 0 and total_mgn_field.any():
particles.boris_update_momentums(minus_half_dt,
total_el_field,
total_mgn_field)
else:
particles.boris_update_momentum_no_mgn(
minus_half_dt, total_el_field)
particles.momentum_is_half_time_step_shifted = True
def compute_total_fields_at_positions(self, positions):
total_el_field = self.electric_fields + self._dynamic_field
return total_el_field.get_at_points(positions, self.time_grid.current_time), \
self.magnetic_fields.get_at_points(positions, self.time_grid.current_time)
def shift_new_particles_velocities_half_time_step_back(self):
minus_half_dt = -1.0 * self.time_grid.time_step_size / 2.0
self.prepare_boris_integration(minus_half_dt)
def apply_domain_boundary_conditions(self):
for arr in self.particle_arrays:
self._domain.collide_with_particles(arr)
self.particle_arrays = [
a for a in self.particle_arrays if len(a.ids) > 0
]
def remove_particles_inside_inner_regions(self):
for region in self.inner_regions:
for p in self.particle_arrays:
region.collide_with_particles(p)
self.particle_arrays = [
a for a in self.particle_arrays if len(a.ids) > 0
]
def generate_new_particles(self, initial=False):
for src in self.particle_sources:
particles = src.generate_initial_particles(
) if initial else src.generate_each_step()
if len(particles.ids):
particles.ids = particles.xp.asarray(
self.particle_tracker.generate_particle_ids(
len(particles.ids)))
self.particle_arrays.append(particles)
def consolidate_particle_arrays(self):
particles_by_type = defaultdict(list)
for p in self.particle_arrays:
particles_by_type[(p.mass, p.charge,
p.momentum_is_half_time_step_shifted)].append(p)
self.particle_arrays = []
for k, v in particles_by_type.items():
mass, charge, shifted = k
ids = v[0].xp.concatenate([p.ids for p in v])
positions = v[0].xp.concatenate([p.positions for p in v])
momentums = v[0].xp.concatenate([p.momentums for p in v])
if len(ids):
self.particle_arrays.append(
ParticleArray(ids, charge, mass, positions, momentums,
shifted))