diff --git a/tests/python/test_dataframe.py b/tests/python/test_dataframe.py new file mode 100644 index 000000000..a209317b4 --- /dev/null +++ b/tests/python/test_dataframe.py @@ -0,0 +1,69 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 The ImpactX Community +# +# Authors: Axel Huebl +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +from impactx import ImpactX, RefPart, distribution, elements + + +def test_df_pandas(): + """ + This tests using ImpactX and Pandas Dataframes + """ + sim = ImpactX() + + sim.particle_shape = 2 + sim.slice_step_diagnostics = True + sim.init_grids() + + # init particle beam + kin_energy_MeV = 2.0e3 + bunch_charge_C = 1.0e-9 + npart = 10000 + + # reference particle + pc = sim.particle_container() + ref = pc.ref_particle() + ref.set_charge_qe(-1.0).set_mass_MeV(0.510998950).set_kin_energy_MeV(kin_energy_MeV) + + # particle bunch + distr = distribution.Waterbag( + sigmaX=3.9984884770e-5, + sigmaY=3.9984884770e-5, + sigmaT=1.0e-3, + sigmaPx=2.6623538760e-5, + sigmaPy=2.6623538760e-5, + sigmaPt=2.0e-3, + muxpx=-0.846574929020762, + muypy=0.846574929020762, + mutpt=0.0, + ) + sim.add_particles(bunch_charge_C, distr, npart) + + assert pc.TotalNumberOfParticles() == npart + + # init accelerator lattice + fodo = [ + elements.Drift(0.25), + elements.Quad(1.0, 1.0), + elements.Drift(0.5), + elements.Quad(1.0, -1.0), + elements.Drift(0.25), + ] + sim.lattice.extend(fodo) + + # simulate + sim.evolve() + + # compare number of global particles + df = pc.to_df(local=False) + if df is not None: + assert npart == len(df) + + +if __name__ == "__main__": + test_df_pandas()