Standardize figure settings

Author: MilesCranmer

gui/plots.py

@@ -20,7 +20,6 @@ def plot_pareto_curve(df: pd.DataFrame, maxsize: int):
     if len(df) == 0 or "Equation" not in df.columns:
         return fig
 
-    # Plotting the data
     ax.loglog(
         df["Complexity"],
         df["Loss"],
@@ -31,23 +30,12 @@ def plot_pareto_curve(df: pd.DataFrame, maxsize: int):
         markersize=6,
     )
 
-    # Set the axis limits
     ax.set_xlim(0.5, maxsize + 1)
     ytop = 2 ** (np.ceil(np.log2(df["Loss"].max())))
     ybottom = 2 ** (np.floor(np.log2(df["Loss"].min() + 1e-20)))
     ax.set_ylim(ybottom, ytop)
 
-    ax.grid(True, which="both", ls="--", linewidth=0.5, color="gray", alpha=0.5)
-    ax.spines["top"].set_visible(False)
-    ax.spines["right"].set_visible(False)
-
-    # Range-frame the plot
-    for direction in ["bottom", "left"]:
-        ax.spines[direction].set_position(("outward", 10))
-
-    # Delete far ticks
-    ax.tick_params(axis="both", which="major", labelsize=10, direction="out", length=5)
-    ax.tick_params(axis="both", which="minor", labelsize=8, direction="out", length=3)
+    stylize_axis(ax)
 
     ax.set_xlabel("Complexity")
     ax.set_ylabel("Loss")
@@ -57,14 +45,23 @@ def plot_pareto_curve(df: pd.DataFrame, maxsize: int):
 
 
 def plot_example_data(test_equation, num_points, noise_level, data_seed):
+    fig, ax = plt.subplots(figsize=(6, 6), dpi=100)
+
     X, y = generate_data(test_equation, num_points, noise_level, data_seed)
     x = X["x"]
 
-    plt.rcParams["font.family"] = "IBM Plex Mono"
-    fig, ax = plt.subplots(figsize=(6, 6), dpi=100)
-
     ax.scatter(x, y, alpha=0.7, edgecolors="w", s=50)
 
+    stylize_axis(ax)
+
+    ax.set_xlabel("x")
+    ax.set_ylabel("y")
+    fig.tight_layout(pad=2)
+
+    return fig
+
+
+def stylize_axis(ax):
     ax.grid(True, which="both", ls="--", linewidth=0.5, color="gray", alpha=0.5)
     ax.spines["top"].set_visible(False)
     ax.spines["right"].set_visible(False)
@@ -76,9 +73,3 @@ def plot_example_data(test_equation, num_points, noise_level, data_seed):
     # Delete far ticks
     ax.tick_params(axis="both", which="major", labelsize=10, direction="out", length=5)
     ax.tick_params(axis="both", which="minor", labelsize=8, direction="out", length=3)
-
-    ax.set_xlabel("x")
-    ax.set_ylabel("y")
-    fig.tight_layout(pad=2)
-
-    return fig