plt class was modified to include a flag for removing/keeping lost particles....

plt class was modified to include a flag for removing/keeping lost particles. The default remains to be removing.

ess/TraceWin.py

@@ -299,9 +299,11 @@ class plt:
               print(data['x'])
     """
 
-    def __init__(self, filename):
+    def __init__(self, filename, flag_remove_loss=True):
         # easy storage..
         self.filename = filename
+        # option to remove lost particles, default True
+        self.flag_remove_loss = flag_remove_loss
         # used to create dict behaviour..
         self._columns = ["x", "xp", "y", "yp", "phi", "E", "l"]
         # read in the file..
@@ -371,14 +373,23 @@ class plt:
             i = self.Nelp.index(key)
 
             ret = {}
-            # some particles are lost, exclude those:
-            lost_mask = self._data[i]["l"] == 0
-            for key in self._data[i]:
-                if isinstance(self._data[i][key], numpy.ndarray):
-                    ret[key] = self._data[i][key][lost_mask]
-                else:
-                    ret[key] = self._data[i][key]
+
+            # if want to EXCLUDE lost particles (default):
+            if self.flag_remove_loss == True:
+                lost_mask = self._data[i]["l"] == 0
+                for para in self._data[i]:
+                    if isinstance(self._data[i][para], numpy.ndarray):
+                        ret[para] = self._data[i][para][lost_mask]
+                    else:
+                        ret[para] = self._data[i][para]
+
+            # if want to KEEP lost particles:
+            else:
+                for para in self._data[i]:
+                    ret[para] = self._data[i][para]
+
             return ret
+
         else:
             print("No data to plot at element", key)