diff --git a/pmd_beamphysics/__init__.py b/pmd_beamphysics/__init__.py index 26ad5ca..5a06a0c 100644 --- a/pmd_beamphysics/__init__.py +++ b/pmd_beamphysics/__init__.py @@ -1,4 +1,5 @@ -from .particles import ParticleGroup, single_particle +from .particles import ParticleGroup, single_particle +from .status import ParticleStatus from .fields.fieldmesh import FieldMesh from .readers import particle_paths from .writers import pmd_init diff --git a/pmd_beamphysics/interfaces/astra.py b/pmd_beamphysics/interfaces/astra.py index 5b1b89c..c87002f 100644 --- a/pmd_beamphysics/interfaces/astra.py +++ b/pmd_beamphysics/interfaces/astra.py @@ -1,4 +1,5 @@ import numpy as np +from pmd_beamphysics.status import ParticleStatus from pmd_beamphysics.readers import component_alias import os @@ -173,9 +174,12 @@ def vprint(*a, **k): # Revese: 1->5, 2->1 status = particle_group.status astra_status = status.copy() - astra_status[ np.where(status==1) ] = 5 - astra_status[ np.where(status==2) ] = 1 + astra_status[ np.where(status==1) ] = 5 # Astra normal (alive) + astra_status[ np.where(status==2) ] = 1 # Astra passive + astra_status[ np.where(status==ParticleStatus.CATHODE)] = -1 # Astra cathode + data['status'][i_start:] = astra_status + # Handle reference particle. If any -1 are found, assume we are starting at the cathode if -1 in astra_status: ref_particle['status']= -1 # At cathode diff --git a/pmd_beamphysics/status.py b/pmd_beamphysics/status.py new file mode 100644 index 0000000..cd3ea57 --- /dev/null +++ b/pmd_beamphysics/status.py @@ -0,0 +1,9 @@ +from enum import IntEnum + +class ParticleStatus(IntEnum): + """ + Particle Status Enum + This is defined by the openPMD-beamphysics standard as integers. + """ + CATHODE = 0 + ALIVE = 1 \ No newline at end of file