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