Add cs_invariant to plotting functions

Add cs_invariant to plot_beamdata, plot_bunchdata and streak_beamdata

mbtrack2/tracking/monitors/plotting.py

@@ -21,11 +21,12 @@ def plot_beamdata(filenames, dataset="mean", dimension="tau", stat_var="mean",
     ----------
     filenames : str or list of str
         Names of the HDF5 files to be plotted.
-    dataset : {"current","emit","mean","std"}
+    dataset : {"current","emit","mean","std","cs_invariant"}
         HDF5 file's dataset to be plotted. The default is "mean".
     dimension : str
          The dimension of the dataset to plot:
             for "emit", dimension = {"x","y","s"},
+            for "cs_invariant", dimension = {"x","y"},
             for "mean" and "std", dimension = {"x","xp","y","yp","tau","delta"}.
             not used if "current".
         The default is "tau".
@@ -76,6 +77,15 @@ def plot_beamdata(filenames, dataset="mean", dimension="tau", stat_var="mean",
                 elif stat_var == "std":
                     y = np.nanstd(data[axis,bunch_index,:],0)
                 y_label = stat_var + " " + label[axis]
+            elif dataset == "cs_invariant":
+                dimension_dict = {"x":0, "y":1}
+                axis = dimension_dict[dimension]
+                label = ['$J_x$ (m)', '$J_y$ (m)']
+                if stat_var == "mean":
+                    y = np.nanmean(data[axis,bunch_index,:],0)
+                elif stat_var == "std":
+                    y = np.nanstd(data[axis,bunch_index,:],0)
+                y_label = stat_var + " " + label[axis]
             elif dataset == "mean" or dataset == "std":
                 dimension_dict = {"x":0, "xp":1, "y":2, "yp":3, "tau":4, 
                                   "delta":5}
@@ -112,6 +122,12 @@ def plot_beamdata(filenames, dataset="mean", dimension="tau", stat_var="mean",
                          "$\\epsilon_{s}$ (m.rad)"]
                 y = data[axis,:,idx]
                 y_label = label[axis]
+            elif dataset == "cs_invariant":
+                dimension_dict = {"x":0, "y":1}
+                axis = dimension_dict[dimension]
+                label = ['$J_x$ (m)', '$J_y$ (m)']
+                y = data[axis,:,idx]
+                y_label = label[axis]
             elif dataset == "mean" or dataset == "std":
                 dimension_dict = {"x":0, "xp":1, "y":2, "yp":3, "tau":4, 
                                   "delta":5}
@@ -144,11 +160,12 @@ def streak_beamdata(filename, dataset="mean", dimension="tau", cm_lim=None):
     ----------
     filename : str
         Name of the HDF5 file that contains the data.
-    dataset : {"current","emit","mean","std"}
+    dataset : {"current","emit","mean","std","cs_invariant"}
         HDF5 file's dataset to be plotted. The default is "mean".
     dimension : str
          The dimension of the dataset to plot:
             for "emit", dimension = {"x","y","s"},
+            for "cs_invariant", dimension = {"x","y"},
             for "mean" and "std", dimension = {"x","xp","y","yp","tau","delta"}.
             not used if "current".
         The default is "tau".
@@ -184,6 +201,13 @@ def streak_beamdata(filename, dataset="mean", dimension="tau", cm_lim=None):
         z = np.array(data["emit"][axis,:,:]).T
         z_label = label[axis]
         title = z_label
+    elif dataset == "cs_invariant":
+        dimension_dict = {"x":0, "y":1}
+        axis = dimension_dict[dimension]
+        label = ['$J_x$ (m)', '$J_y$ (m)']
+        z = np.array(data["cs_invariant"][axis,:,:]).T
+        z_label = label[axis]
+        title = z_label
     else:
         dimension_dict = {"x":0, "xp":1, "y":2, "yp":3, "tau":4, 
                               "delta":5}
@@ -237,6 +261,7 @@ def plot_bunchdata(filenames, bunch_number, dataset, dimension="x",
         The dimension of the dataset to plot. Use "None" for "current",
         otherwise use the following : 
             for "emit", dimension = {"x","y","s"},
+            for "cs_invariant", dimension = {"x","y"},
             for "mean" and "std", dimension = {"x","xp","y","yp","tau","delta"},
             for "action", dimension = {"x","y"}.
     legend : list of str, optional