35 lines
1.3 KiB
Python
35 lines
1.3 KiB
Python
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
import os
|
|
import typer
|
|
|
|
app = typer.Typer()
|
|
|
|
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(energy_files: list[str]=[], labels: list[str]=[], mass: float=1.0):
|
|
"""Plot kinetic energy from velocity 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()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
app() |