Update plotting in examples

Author: austin-hoover

examples/Envelope/danilov_envelope/benchmark_danilov_envelope_20.py

@@ -28,6 +28,8 @@
 from utils import make_fodo_lattice
 from utils import BunchMonitor
 
+plt.style.use("style.mplstyle")
+
 
 # Parse arguments
 # --------------------------------------------------------------------------------------
@@ -123,31 +125,36 @@
 # Plot comparison
 # --------------------------------------------------------------------------------------
 
-figwidth = 4.0 * args.periods
-figwidth = min(figwidth, 10.0)
-
-fig, ax = plt.subplots(figsize=(figwidth, 2.5), constrained_layout=True)
-for i, key in enumerate(histories):
-    history = histories[key]
-
-    plot_kws = {}
-    if key == "envelope":
-        plot_kws["ls"] = "-"
-        plot_kws["lw"] = 2.5
-        plot_kws["marker"] = None
-    if key == "bunch":
-        plot_kws["ls"] = "-"
-        plot_kws["lw"] = 0.0
-        plot_kws["marker"] = "."
-        plot_kws["ms"] = 3.0
-        plot_kws["color"] = "black"
-
-    ax.plot(history["s"], history["xrms"] * 1000.0, **plot_kws)
-    ax.plot(history["s"], history["yrms"] * 1000.0, **plot_kws)
-
-ax.set_ylim(0.0, ax.get_ylim()[1])
-ax.set_xlabel("Distance [m]")
-ax.set_ylabel("Size [mm]")
+figwidth = 3.0 * args.periods
+figwidth = min(figwidth, 7.0)
+
+fig, axs = plt.subplots(nrows=2, figsize=(figwidth, 4.0), sharex=True, sharey=True)
+for i, ax in enumerate(axs):
+    param = ["xrms", "yrms"][i]
+    for j, key in enumerate(histories):
+        history = histories[key]
+        if key == "envelope":
+            ax.plot(
+                history["s"],
+                np.multiply(history[param], 1000.0),
+                color="black",
+                lw=1.5,
+            )
+        else:
+            stride = 10
+            ax.plot(
+                history["s"][::stride],
+                np.multiply(history[param][::stride], 1000.0),
+                marker=".",
+                lw=0,
+                color="red",
+            )
+
+for ax in axs:
+    ax.set_ylim(0.0, ax.get_ylim()[1])
+axs[1].set_xlabel("Distance [m]")
+axs[0].set_ylabel("RMS x [mm]")
+axs[1].set_ylabel("RMS y [mm]")
 
 filename = "fig_benchmark_rms.png"
 filename = os.path.join(output_dir, filename)

examples/Envelope/danilov_envelope/benchmark_danilov_envelope_22.py

@@ -28,6 +28,8 @@
 from utils import make_fodo_lattice
 from utils import BunchMonitor
 
+plt.style.use("style.mplstyle")
+
 
 # Parse arguments
 # --------------------------------------------------------------------------------------
@@ -121,37 +123,38 @@
 # Plot comparison
 # --------------------------------------------------------------------------------------
 
-figwidth = 4.0 * args.periods
-figwidth = min(figwidth, 6.0)
-
-fig, axs = plt.subplots(figsize=(figwidth, 4.0), nrows=2, sharex=True, constrained_layout=True)
-for i, key in enumerate(histories):
-    history = histories[key]
-
-    plot_kws = {}
-    if key == "envelope":
-        plot_kws["ls"] = "-"
-        plot_kws["lw"] = 2.5
-        plot_kws["marker"] = None
-    if key == "bunch":
-        plot_kws["ls"] = "-"
-        plot_kws["lw"] = 0.0
-        plot_kws["marker"] = "+"
-        plot_kws["ms"] = 3.0
-        plot_kws["color"] = "black"
-
-    axs[0].plot(history["s"], history["xrms"] * 1.00e03, **plot_kws)
-    axs[0].plot(history["s"], history["yrms"] * 1.00e03, **plot_kws)
-    axs[1].plot(history["s"], history["rxy"], **plot_kws)
-
-axs[0].set_ylim(0.0, axs[0].get_ylim()[1])
-axs[1].set_ylim(-1.0, 1.0)
-
-axs[0].set_xlabel("Distance [m]")
-axs[0].set_ylabel("Size [mm]")
-axs[1].set_ylabel("rxy")
-
-filename = "fig_benchmark.png"
+figwidth = 3.0 * args.periods
+figwidth = min(figwidth, 7.0)
+
+fig, axs = plt.subplots(nrows=2, figsize=(figwidth, 4.0), sharex=True, sharey=True)
+for i, ax in enumerate(axs):
+    param = ["xrms", "yrms"][i]
+    for j, key in enumerate(histories):
+        history = histories[key]
+        if key == "envelope":
+            ax.plot(
+                history["s"],
+                np.multiply(history[param], 1000.0),
+                color="black",
+                lw=1.5,
+            )
+        else:
+            stride = 10
+            ax.plot(
+                history["s"][::stride],
+                np.multiply(history[param][::stride], 1000.0),
+                marker=".",
+                lw=0,
+                color="red",
+            )
+
+for ax in axs:
+    ax.set_ylim(0.0, ax.get_ylim()[1])
+axs[1].set_xlabel("Distance [m]")
+axs[0].set_ylabel("RMS x [mm]")
+axs[1].set_ylabel("RMS y [mm]")
+
+filename = "fig_benchmark_rms.png"
 filename = os.path.join(output_dir, filename)
 plt.savefig(filename, dpi=300)
 plt.show()

examples/Envelope/danilov_envelope/match_danilov_envelope_20.py

@@ -16,6 +16,8 @@
 
 from utils import make_fodo_lattice
 
+plt.style.use("style.mplstyle")
+
 
 # Parse arguments
 # --------------------------------------------------------------------------------------

examples/Envelope/danilov_envelope/match_danilov_envelope_22.py

@@ -16,6 +16,8 @@
 
 from utils import make_fodo_lattice
 
+plt.style.use("style.mplstyle")
+
 
 # Parse arguments
 # --------------------------------------------------------------------------------------

examples/Envelope/danilov_envelope/plot_danilov_envelope_20.py

@@ -3,6 +3,8 @@
 
 from orbit.danilov_envelope import DanilovEnvelope20
 
+plt.style.use("style.mplstyle")
+
 
 envelope = DanilovEnvelope20(
     eps_x=10.00e-06,

examples/Envelope/danilov_envelope/plot_danilov_envelope_22.py

@@ -3,6 +3,8 @@
 
 from orbit.danilov_envelope import DanilovEnvelope22
 
+plt.style.use("style.mplstyle")
+
 
 envelope = DanilovEnvelope22(
     intrinsic_emittance=20.00e-06,

examples/Envelope/danilov_envelope/style.mplstyle

@@ -0,0 +1,4 @@
+axes.linewidth: 1.25
+figure.constrained_layout.use: True
+xtick.minor.visible: True
+ytick.minor.visible: True

examples/Envelope/danilov_envelope/utils.py

@@ -61,7 +61,7 @@ def make_fodo_lattice(
         fill_frac : Fraction of the lattice occupied by quadrupoles.
         start : str
             "drift":
-                O-F-F-O-O-D-D-O. 
+                O-F-F-O-O-D-D-O.
             "quad":
                 F-O-O-D-D-O-O-F
         fringe: Toggles fringe fields before/after quads.