added helicity and status read write

heparchy/read/_base.py

@@ -91,7 +91,7 @@ def com_energy(self) -> dict:
         """
 
     @abstractmethod
-    def get_custom(self, name: str):
+    def get_custom_meta(self, name: str):
         """Returns user-defined piece of metadata."""
 
     @abstractmethod

heparchy/read/hdf.py

@@ -43,6 +43,14 @@ def color(self) -> np.ndarray:
     def pdg(self) -> np.ndarray:
         return self._grp['pdg'][...]
 
+    @property
+    def status(self) -> np.ndarray:
+        return self._grp['status'][...]
+
+    @property
+    def helicity(self) -> np.ndarray:
+        return self._grp['helicity'][...]
+
     @property
     def available(self) -> list:
         dataset_names = list()
@@ -55,6 +63,9 @@ def mask(self, name: str) -> np.ndarray:
     def get_custom(self, name: str):
         return self._grp[name][...]
 
+    def get_custom_meta(self, name: str):
+        return self._grp.attrs[name]
+
     def copy(self):
         return deepcopy(self)
 
@@ -96,7 +107,7 @@ def com_energy(self) -> dict:
             'unit': self._meta['e_unit'],
             }
 
-    def get_custom(self, name: str):
+    def get_custom_meta(self, name: str):
         return self._meta[name]
 
     def read_event(self, evt_num):

heparchy/write/hdf.py

@@ -152,6 +152,41 @@ def set_pdg(self, data: np.ndarray) -> None:
                 shape=(self.num_pcls,),
                 dtype=self.__types.int,
                 )
+
+    def set_status(self, data: np.ndarray) -> None:
+        """Write status codes for all particles to event.
+        
+        Parameters
+        ----------
+        data : iterable or 1d numpy array
+            Iterable of ints containing status codes for every
+            particle in the event.
+        """
+        self.__set_num_pcls(data)
+        self.__mk_dset(
+                name='status',
+                data=data,
+                shape=(self.num_pcls,),
+                dtype=self.__types.int,
+                )
+
+    def set_helicity(self, data: np.ndarray) -> None:
+        """Write helicity values for all particles to event.
+
+        Parameters
+        ----------
+        data : iterable or 1d numpy array
+            Iterable of floats containing helicity values for every
+            particle in the event.
+        """
+        self.__set_num_pcls(data)
+        self.__mk_dset(
+                name='helicity',
+                data=data,
+                shape=(self.num_pcls,),
+                dtype=self.__types.helicity,
+                )
+
 
     def set_mask(self, name: str, data: np.ndarray) -> None:
         """Write bool mask for all particles in event.
@@ -201,6 +236,19 @@ def set_custom_dataset(
                 dtype=dtype,
                 )
 
+    def set_custom_meta(self, name: str, metadata: Any) -> None:
+        """Store custom metadata to the event.
+
+        Parameters
+        ----------
+        name : str
+            Handle to access the metadata at read time.
+        metadata : str, int, float, or iterables thereof
+            The data you wish to store.
+        """
+        self.__evt.attrs[name] = metadata
+
+
 class _ProcessWriter(ProcessWriterBase):
     def __init__(self, file_obj: WriterBase, key: str):
         self.__file_obj = file_obj