density_file -> density

adding xlim, ylim to dst.subplot
adding r, rp to dst class as indirect variables

ess/TraceWin.py

@@ -18,6 +18,7 @@ class dst:
         self.filename = filename
         # used to create dict behaviour..
         self._columns = ["x", "xp", "y", "yp", "phi", "E"]
+        self._indirect_columns = ["r", "rp"]
         if filename:
             # read in the file..
             self._readBinaryFile()
@@ -119,23 +120,31 @@ class dst:
         self.mass = Table[-1]
 
     def keys(self):
-        return self._columns[:]
+        return self._columns[:] + self._indirect_columns[:]
 
     def __getitem__(self, key):
         # makes the class function as a dictionary
         # e.g. dst['x'] returns the x array..
+        import numpy
+
         try:
-            i = self._columns.index(key)
-            return self._data[:, i]
-        except KeyError:
-            raise KeyError("Available keys: " + str(self._columns))
+            if key == "r":
+                return numpy.sqrt(self["x"] ** 2 + self["y"] ** 2)
+            elif key == "rp":
+                return numpy.sqrt(self["xp"] ** 2 + self["yp"] ** 2)
+            else:
+                i = self._columns.index(key)
+                return self._data[:, i]
+
+        except ValueError:
+            raise KeyError(f"Unknown key '{key}', available keys: {self.keys()}")
 
     def __setitem__(self, key, value):
         try:
             i = self._columns.index(key)
             self._data[:, i] = value
-        except KeyError:
-            raise KeyError("Available keys: " + str(self._columns))
+        except ValueError:
+            raise KeyError(f"Unknown key '{key}', available keys: {self.keys()}")
 
     def save(self, filename, toutatis=False):
         """
@@ -180,7 +189,7 @@ class dst:
         fout.write(pack("d", self.mass))
         fout.close()
 
-    def subplot(self, index, x, y=None, nb=100, mask=None):
+    def subplot(self, index, x, y=None, nb=100, mask=None, xlim=None, ylim=None):
         """
         Create a subplot histogram similar to TraceWin.
 
@@ -189,12 +198,12 @@ class dst:
             import numpy as np
             from ess import TraceWin
             from matplotlib import pyplot as plt
-            data=TraceWin.dst('part_dtl1.dst')
-            m=np.where(data['E']>3.5)
-            data.subplot(221,'x','xp',mask=m)
-            data.subplot(222,'y','yp',mask=m)
-            data.subplot(223,'phi','E',mask=m)
-            data.subplot(224,'x','y',mask=m)
+            data = TraceWin.dst('part_dtl1.dst')
+            m = np.where(data['E']>3.5)
+            data.subplot(221, 'x', 'xp', mask=m)
+            data.subplot(222, 'y', 'yp', mask=m)
+            data.subplot(223, 'phi', 'E', mask=m)
+            data.subplot(224, 'x', 'y', mask=m)
             plt.show()
 
         """
@@ -242,17 +251,29 @@ class dst:
             plt.title("{} [{}] - {} [{}]".format(x, units[x], y, units[y]))
             hist, bin_edges = np.histogram(dx, bins=nb)
             b = bin_edges[:-1] + 0.5 * (bin_edges[1] - bin_edges[0])
+            if ylim:
+                miny = ylim[0]
+                maxy = ylim[1]
+            else:
+                miny = min(dy)
+                maxy = max(dy)
             plt.plot(
                 b,
-                hist * 0.2 * (max(dy) - min(dy)) / max(hist) + min(dy),
+                hist * 0.2 * (maxy - miny) / max(hist) + miny,
                 "k",
                 lw=1.5,
                 drawstyle="steps",
             )
             hist, bin_edges = np.histogram(dy, bins=nb)
             b = bin_edges[:-1] + 0.0 * (bin_edges[1] - bin_edges[0])
+            if xlim:
+                minx = xlim[0]
+                maxx = xlim[1]
+            else:
+                minx = min(dx)
+                maxx = max(dx)
             plt.plot(
-                hist * 0.2 * (max(dx) - min(dx)) / max(hist) + min(dx),
+                hist * 0.2 * (maxx - minx) / max(hist) + minx,
                 b,
                 "k",
                 lw=1.5,
@@ -262,6 +283,10 @@ class dst:
             # plot a simple 1D histogram..
             plt.hist(dx, bins=nb)
             plt.title("{} [{}]".format(x, units[x]))
+        if xlim:
+            plt.xlim(xlim)
+        if ylim:
+            plt.ylim(ylim)
 
 
 class plt:
@@ -567,7 +592,7 @@ class plt:
         return _dst
 
 
-class density_file:
+class density:
     """
     Simple class to read a TraceWin density file
     into a pythonized object
@@ -1053,6 +1078,12 @@ class density_file:
         fout.close()
 
 
+class density_file(density):
+    def __init__(self, filename, envelope=None):
+        print("Deprecated, use TraceWin.density() instead")
+        super().__init__(filename, envelope)
+
+
 class remote_data_merger:
     def __init__(self, base="."):
         self._base = base