def test_binary_field(self):
d = Config().make()
d.particle_arrays = [
ParticleArray(1, -1, 1, [(1, 2, 3)], [(-2, 2, 0)], False)
]
assert_array_almost_equal(
d.binary_electric_field_at_positions((1, 2, 3)), [(0, 0, 0)])
assert_array_almost_equal(
d.binary_electric_field_at_positions((1, 2, 4)), [(0, 0, -1)])
assert_array_almost_equal(
d.binary_electric_field_at_positions((0, 2, 3)), [(1, 0, 0)])
assert_array_almost_equal(
d.binary_electric_field_at_positions((0, 1, 2)),
[(1 / sqrt(27), 1 / sqrt(27), 1 / sqrt(27))])
d.particle_arrays = [
ParticleArray(2, -1, 1, [(1, 2, 3), (1, 2, 3)], [(-2, 2, 0),
(0, 0, 0)],
False),
ParticleArray(2, -1, 1, [(1, 2, 3), (1, 2, 3)], [(-2, 2, 0),
(0, 0, 0)], False)
]
assert_array_almost_equal(
d.binary_electric_field_at_positions([(1, 2, 3), (1, 2, 4),
(0, 2, 3), (0, 1, 2)]),
[(0, 0, 0), (0, 0, -4), (4, 0, 0),
(4 / sqrt(27), 4 / sqrt(27), 4 / sqrt(27))])
def test_weight_particles_charge_to_mesh(self):
mesh = SpatialMeshConf((2, 4, 8),
(1, 2, 4)).make(BoundaryConditionsConf())
particle_arrays = [ParticleArray(1, -2, 4, [(1, 1, 3)], [(0, 0, 0)])]
mesh.weight_particles_charge_to_mesh(particle_arrays)
assert_array_equal(
mesh.charge_density,
np.array([[[0, 0, 0], [0, 0, 0], [0, 0, 0]],
[[-0.25 / 8, -0.75 / 8, 0], [-0.25 / 8, -0.75 / 8, 0],
[0, 0, 0]], [[0, 0, 0], [0, 0, 0], [0, 0, 0]]]))
particle_arrays = [
ParticleArray([1, 2], -2, 4, [(1, 1, 3), (1, 1, 3)],
np.zeros((2, 3)))
]
mesh.clear_old_density_values()
mesh.weight_particles_charge_to_mesh(particle_arrays)
assert_array_equal(
mesh.charge_density,
np.array([[[0, 0, 0], [0, 0, 0], [0, 0, 0]],
[[-0.25 / 4, -0.75 / 4, 0], [-0.25 / 4, -0.75 / 4, 0],
[0, 0, 0]], [[0, 0, 0], [0, 0, 0], [0, 0, 0]]]))
mesh.clear_old_density_values()
particle_arrays = [ParticleArray(1, -2, 4, [(2, 4, 8)], [(0, 0, 0)])]
mesh.weight_particles_charge_to_mesh(particle_arrays)
assert_array_equal(
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], [0, 0, 0], [0, 0, -0.25]]]))
particle_arrays = [ParticleArray(1, -2, 4, [(1, 2, 8.1)], [(0, 0, 0)])]
with pytest.raises(ValueError,
match="Position is out of meshgrid bounds"):
mesh.weight_particles_charge_to_mesh(particle_arrays)
def test_generate_for_simulation(self):
ps = ParticleSource('test', Box(6, 0), 17, 13, (4, 4, 4), 0, -2, 6)
assert_dataclass_eq(
ps.generate_initial_particles(),
ParticleArray(range(17), -2, 6, np.full((17, 3), 6),
np.full((17, 3), 4), False))
assert_dataclass_eq(
ps.generate_each_step(),
ParticleArray(range(13), -2, 6, np.full((13, 3), 6),
np.full((13, 3), 4), False))
def test_generate_particles(self):
ps = ParticleSource('test', Box((1., 2., 3.), 0), 17, 13, (-2, 3, 1),
0, -2, 6)
assert_dataclass_eq(
ps.generate_num_of_particles(3),
ParticleArray(range(3), -2, 6, [(1, 2, 3)] * 3, [(-2, 3, 1)] * 3,
False))
assert_dataclass_eq(
ps.generate_num_of_particles(1),
ParticleArray([0], -2, 6, [(1, 2, 3)], [(-2, 3, 1)], False))
assert_dataclass_eq(
ps.generate_num_of_particles(0),
ParticleArray([], -2, 6, np.empty((0, 3)), np.empty((0, 3)),
False))
def test_update_positions(self):
p = ParticleArray([123], -1.0, 2.0, [(0., 0., 1.)], [(1., 0., 3.)])
p.update_positions(10.0)
assert_array_equal(p.positions, [(5., 0., 16.)])
p = ParticleArray((1, 2), -1.0, 2.0, [(0., 0., 1.), (1, 2, 3)],
[(1., 0., 3.), (-1, -0.5, 0)])
p.update_positions(10.0)
assert_array_equal(p.positions, [(5., 0., 16.), (-4, -0.5, 3)])
def test_field_at_point(self):
p = ParticleArray([1], -16.0, 2.0, [(0., 0., 1.)], [(1., 0., 3.)])
assert_array_equal(p.field_at_points((2., 0., 1.)), (-4, 0, 0))
assert_array_equal(p.field_at_points((2., 0., 1.)), np.array(
(-4, 0, 0)))
assert_array_equal(p.field_at_points(np.array((2., 0., 1.))),
(-4, 0, 0))
assert_array_equal(p.field_at_points((0., 0., 1.)),
np.array([np.nan, np.nan, np.nan]))
p = ParticleArray('12', -16.0, 2.0, [(0, 0, 1), (0, 0, 0)],
np.zeros((2, 3)))
assert_array_equal(p.field_at_points((0, 0, 0.5)), (0, 0, 0))
assert_array_equal(p.field_at_points((0, 0, 2)), (0, 0, -20))
assert_array_equal(p.field_at_points((0., 0., 0)),
np.array([np.nan, np.nan, np.nan])) # todo: fix!
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_update_momentums_function(self):
assert_array_equal(
ParticleArray._boris_update_momentums(-1, 2, (1, 0, 3), 0.1,
(-1.0, 2.0, 3.0), (0, 0, 0)),
(1.1, -0.2, 2.7))
assert_array_equal(
ParticleArray._boris_update_momentums(-1, 2, (1, 0, 3), 2,
(-1.0, 2.0, 3.0),
(2 * speed_of_light, 0, 0)),
(3, -2, -5))
assert_array_equal(
ParticleArray._boris_update_momentums(-1, 2, (1, 0, 3), 2,
(-1.0, 2.0, 3.0),
(2 * speed_of_light, 0, 0)),
(3, -2, -5))
assert_array_equal(
ParticleArray._boris_update_momentums(
charge=-1,
mass=2,
momentum_arr=self.xp.array([(1, 0, 3)] * 10),
dt=2,
el_field_arr=[(-1.0, 2.0, 3.0)] * 10,
mgn_field_arr=[(2 * speed_of_light, 0, 0)] * 10),
self.xp.array([(3, -2, -5)] * 10))
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))
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_update_momentums_no_mgn(self):
p = ParticleArray(123, -1.0, 2.0, (0., 0., 1.), (1., 0., 3.))
p.boris_update_momentum_no_mgn(0.1, (-1.0, 2.0, 3.0))
assert_array_equal(p.momentums, (1.1, -0.2, 2.7))
def do_write(self, sim: 'Simulation', h5file: File) -> None:
gg = h5file.create_group('SpatialMesh')
sim.mesh.export_h5(gg)
for i, c in enumerate('xyz'):
gg[f'electric_field_{c}'] = sim.electric_field.array.data[
..., i].flatten()
gg['charge_density'] = sim.charge_density.data.flatten()
gg['potential'] = sim.potential.data.flatten()
sim.time_grid.export_h5(h5file.create_group('TimeGrid'))
g = h5file.create_group('ParticleSources')
g.attrs['number_of_sources'] = [len(sim.particle_sources)]
for s in sim.particle_sources:
s.export_h5(g.create_group(s.name))
for p in sim.particle_arrays:
s = next(s for s in sim.particle_sources
if s.charge == p.charge and s.mass == p.mass)
p.export_h5(g[s.name])
for s in sim.particle_sources:
if 'particle_id' not in g[s.name]:
ParticleArray([], s.charge, s.mass, np.empty((0, 3)),
np.empty((0, 3)), True).export_h5(g[s.name])
g = h5file.create_group('InnerRegions')
g.attrs['number_of_regions'] = [len(sim.inner_regions)]
for s in sim.inner_regions:
s.export_h5(g.create_group(s.name))
g = h5file.create_group('ExternalFields')
if sim.electric_fields.__class__.__name__ == "FieldZero":
ff = []
elif sim.electric_fields.__class__.__name__ == "FieldSum":
ff = sim.electric_fields.fields
else:
ff = [sim.electric_fields]
g.attrs['number_of_electric_fields'] = len(ff)
for s in ff:
sg = g.create_group(s.name)
if s.__class__ is FieldUniform:
ft = 'electric_uniform'
for i, c in enumerate('xyz'):
sg.attrs[
f'electric_uniform_field_{c}'] = s.uniform_field_vector[
i]
elif s.__class__ is FieldExpression:
ft = 'electric_tinyexpr'
for i, c in enumerate('xyz'):
expr = getattr(s, f'expression_{c}')
expr = np.string_(expr.encode('utf8')) + b'\x00'
sg.attrs[f'electric_tinyexpr_field_{c}'] = expr
elif s.__class__ is FieldFromCSVFile:
ft = 'electric_on_regular_grid'
sg.attrs['electric_h5filename'] = np.string_(
s.field_filename.encode('utf8') + b'\x00')
sg.attrs['field_type'] = np.string_(ft.encode('utf8') + b'\x00')
if sim.magnetic_fields.__class__.__name__ == "FieldZero":
ff = []
elif sim.magnetic_fields.__class__.__name__ == "FieldSum":
ff = sim.magnetic_fields.fields
else:
ff = [sim.magnetic_fields]
g.attrs['number_of_magnetic_fields'] = len(ff)
for s in ff:
sg = g.create_group(s.name)
if s.__class__ is FieldUniform:
ft = 'magnetic_uniform'
sg.attrs['speed_of_light'] = speed_of_light
for i, c in enumerate('xyz'):
sg.attrs[
f'magnetic_uniform_field_{c}'] = s.uniform_field_vector[
i]
elif s.__class__ is FieldExpression:
ft = 'magnetic_tinyexpr'
sg.attrs['speed_of_light'] = speed_of_light
for i, c in enumerate('xyz'):
expr = getattr(s, f'expression_{c}')
expr = np.string_(expr.encode('utf8')) + b'\x00'
sg.attrs[f'magnetic_tinyexpr_field_{c}'] = expr
elif s.__class__ is FieldFromCSVFile:
ft = 'magnetic_on_regular_grid'
sg.attrs['magnetic_h5filename'] = np.string_(
s.field_filename.encode('utf8') + b'\x00')
sg.attrs['field_type'] = np.string_(ft.encode('utf8') + b'\x00')
g = h5file.create_group('ParticleInteractionModel')
g.attrs['particle_interaction_model'] = \
np.string_(sim.particle_interaction_model.name.encode('utf8') + b'\x00')
def test_field_at_point(self):
p = ParticleArray([1], -16.0, 2.0, [(0., 0., 1.)], [(1., 0., 3.)])
assert_array_equal(p.field_at_points((2., 0., 1.)), [(-4, 0, 0)])
assert_array_equal(p.field_at_points((2., 0., 1.)),
self.xp.array([(-4, 0, 0)]))
assert_array_equal(p.field_at_points(self.xp.array((2., 0., 1.))),
[(-4, 0, 0)])
assert_array_equal(p.field_at_points((0., 0., 1.)),
self.xp.array([[0, 0, 0]]))
p = ParticleArray((1, 2), -16.0, 2.0, [(0, 0, 1), (0, 0, 0)],
self.xp.zeros((2, 3)))
assert_array_equal(p.field_at_points((0, 0, 0.5)), [(0, 0, 0)])
assert_array_equal(p.field_at_points((0, 0, 2)), [(0, 0, -20)])
assert_array_equal(p.field_at_points((0., 0., 0)),
self.xp.array([[0, 0, 16]]))
assert_array_equal(
p.field_at_points([(0, 0, 0.5), (0, 0, 2), (0, 0, 2)]),
[(0, 0, 0), (0, 0, -20), (0, 0, -20)])