diff --git a/mbtrack2/impedance/impedance_model.py b/mbtrack2/impedance/impedance_model.py
index 7b6ff1fb4a73d5eaddb32086a1fea15760bc145e..96efc7fb265554b7d17c631f5731b48ba4c6d19e 100644
--- a/mbtrack2/impedance/impedance_model.py
+++ b/mbtrack2/impedance/impedance_model.py
@@ -8,7 +8,7 @@ from copy import deepcopy
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
-#from matplotlib import colormaps
+from matplotlib import colormaps
 from mpl_toolkits.axes_grid1.inset_locator import inset_axes
 from scipy.integrate import trapz
 from scipy.interpolate import interp1d
diff --git a/mbtrack2/tracking/monitors/plotting.py b/mbtrack2/tracking/monitors/plotting.py
index 5cfeeba9f0386e99009cc42a8cbdef521d2af05d..e407151aa17b7450205afcf80f6909c56ff5f970 100644
--- a/mbtrack2/tracking/monitors/plotting.py
+++ b/mbtrack2/tracking/monitors/plotting.py
@@ -1220,7 +1220,6 @@ def plot_cavitydata(filename,
                     bunch_number=0,
                     turn=None,
                     cm_lim=None,
-                    max = False,
                     show_objective=False):
     """
     Plot data recorded by CavityMonitor.
@@ -1274,10 +1273,6 @@ def plot_cavitydata(filename,
     cavity_data = file[cavity_name]
 
     time = np.array(cavity_data["time"])
-    if max :
-        maxT = np.max(np.nonzero(time))-1
-    else:
-        maxT = len(time)
 
     ph = {"cavity": 0, "beam": 1, "generator": 2, "ig": 3}
     labels = ["Cavity", "Beam", "Generator", "Generator"]
@@ -1312,21 +1307,21 @@ def plot_cavitydata(filename,
 
         fig, ax = plt.subplots()
         twin = ax.twinx()
-        p1, = ax.plot(time[:maxT],
-                      np.abs(data[ph[phasor]][:maxT]) * units_val[ph[phasor]],
+        p1, = ax.plot(time,
+                      np.abs(data[ph[phasor]]) * units_val[ph[phasor]],
                       color="r",
                       label=ylabel1)
-        p2, = twin.plot(time[:maxT],
-                        np.angle(data[ph[phasor]])[:maxT],
+        p2, = twin.plot(time,
+                        np.angle(data[ph[phasor]]),
                         color="b",
                         label=ylabel2)
         if show_objective:
-            o1, = ax.plot(time[:maxT],
-                          np.array(cavity_data["Vc"][:maxT]) * 1e-6,
+            o1, = ax.plot(time,
+                          np.array(cavity_data["Vc"]) * 1e-6,
                           "r--",
                           label="Cavity voltage objective value [MV]")
-            o2, = twin.plot(time[:maxT],
-                            np.array(cavity_data["theta"][:maxT]),
+            o2, = twin.plot(time,
+                            np.array(cavity_data["theta"]),
                             "b--",
                             label="Cavity phase objective value [rad]")
         ax.set_xlabel("Turn number")