From d76fe53585662fce81212dc64aee854b4392fcfe Mon Sep 17 00:00:00 2001
From: ChristopherMayes <31023527+ChristopherMayes@users.noreply.github.com>
Date: Wed, 23 Mar 2022 11:20:37 -0700
Subject: [PATCH 1/2] Add ParticleStatus enum, fix astra initial particle
 writing

---
 pmd_beamphysics/__init__.py         | 3 ++-
 pmd_beamphysics/interfaces/astra.py | 8 ++++++--
 pmd_beamphysics/status.py           | 9 +++++++++
 3 files changed, 17 insertions(+), 3 deletions(-)
 create mode 100644 pmd_beamphysics/status.py

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..af3c2d2 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[ int(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

From cc02a9d731756d2a962324497cf7d2988f192c8e Mon Sep 17 00:00:00 2001
From: ChristopherMayes <31023527+ChristopherMayes@users.noreply.github.com>
Date: Wed, 23 Mar 2022 11:32:15 -0700
Subject: [PATCH 2/2] fix array mask

---
 pmd_beamphysics/interfaces/astra.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pmd_beamphysics/interfaces/astra.py b/pmd_beamphysics/interfaces/astra.py
index af3c2d2..c87002f 100644
--- a/pmd_beamphysics/interfaces/astra.py
+++ b/pmd_beamphysics/interfaces/astra.py
@@ -176,7 +176,7 @@ def vprint(*a, **k):
     astra_status = status.copy()
     astra_status[ np.where(status==1) ] = 5 # Astra normal (alive)
     astra_status[ np.where(status==2) ] = 1 # Astra passive
-    astra_status[ int(ParticleStatus.CATHODE) ] = -1 # Astra cathode
+    astra_status[ np.where(status==ParticleStatus.CATHODE)] = -1 # Astra cathode
     
     data['status'][i_start:] = astra_status