def test_particle_ids(self):
p1 = ParticleSource("test", Box(), 120, 5, np.array((1, 0, 0)), 200,
-1, 2)
a1 = p1.generate_initial_particles()
a2 = p1.generate_each_step()
assert_array_equal(a1.ids, range(120))
assert_array_equal(a2.ids, range(5))
def test_export_h5(self):
f = BytesIO()
p1 = ParticleSource()
with h5py.File(f, mode="w") as h5file:
p1.export_h5(h5file.create_group(p1.name))
with h5py.File(f, mode="r") as h5file:
p2 = ParticleSource.import_h5(h5file[p1.name])
assert_dataclass_eq(p1, p2)
def test_generate_momentums(self, backend):
ps = ParticleSource(mean_momentum=(3, -2, 0), temperature=10, mass=.3)
ps._generator = RandomState(123)
p = ps.generate_num_of_particles(1000000)
assert_almost_ae = p.xp.testing.assert_array_almost_equal
assert_almost_ae(p.momentums.mean(axis=0), (3, -2, 0), 2)
assert_almost_ae(((p.momentums - p.xp.array([3, -2, 0])).std(axis=0)),
np.full(3, sqrt(3)), 2)
def test_generate_positions(self):
ps = ParticleSource()
ps._generator = RandomState(123)
p = ps.generate_num_of_particles(100)
assert_ae = p.xp.testing.assert_array_equal
assert_ae(
p.positions,
Box().generate_uniform_random_posititons(RandomState(123), 100))
def test_h5(self, tmpdir):
fname = tmpdir.join('test_particle_source.h5')
p1 = ParticleSource("test", Box(), 120, 5, np.array((1, 0, 0)), 0, -1,
2)
with h5py.File(fname, mode="w") as h5file:
p1.save_h5(h5file)
with h5py.File(fname, mode="r") as h5file:
p2 = ParticleSource.load_h5(h5file)
assert p1 == p2
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_init(self):
p = ParticleSource()
assert p.name == "particle_source"
assert_dataclass_eq(p.shape, Box())
assert p.initial_number_of_particles == 0
assert p.particles_to_generate_each_step == 0
assert_array_equal(p.mean_momentum, np.zeros(3))
assert p.temperature == 0
assert p.charge == 0
assert p.mass == 1
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 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_generate(self):
p1 = ParticleSource("test", Box(), 120, 5, np.array((1, 0, 0)), 200,
-1, 2)
p2 = ParticleSource("test", Box(), 120, 5, np.array((1, 0, 0)), 200,
-1, 2)
p3 = ParticleSource("test", Box(), 120, 5, np.array((1, 0, 0)), 200,
-1, 2)
p3._generator.set_state(p1._generator.get_state())
a1 = p1.generate_initial_particles()
a2 = p2.generate_initial_particles()
a3 = p3.generate_initial_particles()
assert a1 != a2
assert a1 == a3
a32 = p3.generate_each_step()
assert a1 != a32
a12 = p1.generate_each_step()
assert a12 == a32