73 lines
2.6 KiB
Python
73 lines
2.6 KiB
Python
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)
|
|
kin_file = total_energy_file.replace("total_energies", "kinetic_energies")
|
|
pot_file = total_energy_file.replace("total_energies", "potential_energies")
|
|
kinetic_energy = np.loadtxt(kin_file)
|
|
potential_energy = np.loadtxt(pot_file)
|
|
step_count = len(total_energy)
|
|
return total_energy, kinetic_energy, potential_energy, step_count
|
|
|
|
@app.command()
|
|
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}"
|
|
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__":
|
|
app() |