Further development on project 3

src/project3/python/plot_energies.py

@@ -2,8 +2,41 @@ import matplotlib.pyplot as plt
 import numpy as np
 import os
 import typer
+from pathlib import Path
+from typing import List
 
 app = typer.Typer()
+SHOW_PLOTS = os.getenv("SHOW_PLOTS", "1") == "1"
+
+
+def get_rc_params():
+    colors = ["FF220C", "70D6FF", "8AAA79", "666370", "1C1F33"]
+    rcParams = plt.rcParams
+    # Use LaTeX for rendering
+    # Setup fonts
+    rcParams["text.usetex"] = True
+    rcParams["font.family"] = "serif"
+    rcParams["font.size"] = 12
+    rcParams["axes.labelsize"] = 12
+    rcParams["axes.titlesize"] = 12
+    rcParams["legend.fontsize"] = 10
+    rcParams["xtick.labelsize"] = 10
+    rcParams["ytick.labelsize"] = 10
+    # Figure size and resolution
+    rcParams["figure.figsize"] = (6, 4)
+    rcParams["figure.dpi"] = 300
+    # Use colors from the palette
+    rcParams["axes.prop_cycle"] = plt.cycler(color=[f"#{color}" for color in colors])
+    # Grid
+    rcParams["axes.grid"] = True
+    rcParams["grid.alpha"] = 0.5
+    rcParams["grid.linestyle"] = "--"
+    # Point ticks to the inside of the axes
+    rcParams["xtick.direction"] = "in"
+    rcParams["ytick.direction"] = "in"
+    rcParams["xtick.top"] = True
+    rcParams["ytick.right"] = True
+    return rcParams
 
 def load_data(total_energy_file):
     total_energy = np.loadtxt(total_energy_file)
@@ -15,20 +48,25 @@ def load_data(total_energy_file):
     return total_energy, kinetic_energy, potential_energy, step_count
 
 @app.command()
-def plot_kinetic_energy(energy_files: list[str]=[], labels: list[str]=[], mass: float=1.0):
+def plot_kinetic_energy(total_energy_files: List[Path], output_file: Path):
     """Plot kinetic energy from velocity files."""
+    plt.rcParams.update(get_rc_params())
+    fig, ax = plt.subplots()
+    energy_files = [str(fn) for fn in total_energy_files if fn.is_file()]
+    labels = [f"{fn.split('/')[-2]} ({fn.split('/')[-1].split('_')[4]})" for fn in energy_files]
     for i, energy_file in enumerate(energy_files):
         total_energy, kinetic_energy, potential_energy, step_count = load_data(energy_file)
         time = np.linspace(0, 50e-6, step_count)
         color = f"C{i % 10}"
-        plt.plot(time, total_energy, label=labels[i], color=color, linestyle="-")
-        plt.plot(time, potential_energy, color=color, linestyle="--")
-        plt.plot(time, kinetic_energy, color=color, linestyle=":")
-    plt.xlabel('Time (s)')
-    plt.ylabel('Energy (J)')
-    plt.title('Energy Evolution')
-    plt.legend()
-    plt.show()
+        ax.plot(time, total_energy, label=labels[i], color=color, linestyle="-")
+        ax.plot(time, potential_energy, color=color, linestyle="--")
+        ax.plot(time, kinetic_energy, color=color, linestyle=":")
+    ax.set_xlabel("Time (s)")
+    ax.set_ylabel(r"Energy (Da $\cdot$ m$^2$/s$^2$)")
+    ax.legend()
+    fig.savefig(output_file)
+    if SHOW_PLOTS:
+        fig.show()
 
 
 if __name__ == "__main__":