QuarksLab / Strange /strange /exp /plot_results.py
TwoQuarks's picture
Upload 100 files
ccba775 verified
Raw
History Blame Contribute Delete
2.2 kB
import os
import sys
from pathlib import Path
THIS_DIR = Path(__file__).resolve().parent
PROJECT_ROOT = THIS_DIR.parent
sys.path.append(str(PROJECT_ROOT))
import pandas as pd
import matplotlib.pyplot as plt
def plot_from_csv(csv_path: Path, out_dir: Path, prefix: str) -> None:
df = pd.read_csv(csv_path)
out_dir.mkdir(parents=True, exist_ok=True)
# Episode return
if {"episode","reward"}.issubset(df.columns):
ep_ret = df.groupby("episode")["reward"].sum().reset_index()
plt.figure()
plt.plot(ep_ret["episode"], ep_ret["reward"])
plt.xlabel("episode")
plt.ylabel("episode_return")
plt.title(f"{prefix} episode return")
plt.tight_layout()
plt.savefig(out_dir / f"{prefix.lower()}_episode_return.png", dpi=160)
plt.close()
# Stability / diversity if present
for col in ["stability","diversity","phase","progress","w_t","H_swarm","S_t","T"]:
if col in df.columns:
series = df.groupby("episode")[col].mean().reset_index()
plt.figure()
plt.plot(series["episode"], series[col])
plt.xlabel("episode")
plt.ylabel(col)
plt.title(f"{prefix} mean {col}")
plt.tight_layout()
plt.savefig(out_dir / f"{prefix.lower()}_{col}.png", dpi=160)
plt.close()
def main():
import argparse
ap = argparse.ArgumentParser()
ap.add_argument("--csv", required=True, help="Path to results CSV (relative to project root is ok).")
ap.add_argument("--out", default="graphics", help="Output directory for plots (relative ok).")
ap.add_argument("--prefix", default="Strange", help="Plot filename prefix.")
args = ap.parse_args()
csv_path = Path(args.csv)
if not csv_path.is_absolute():
csv_path = (PROJECT_ROOT / csv_path).resolve()
out_dir = Path(args.out)
if not out_dir.is_absolute():
out_dir = (PROJECT_ROOT / out_dir).resolve()
if not csv_path.exists():
raise FileNotFoundError(f"CSV not found: {csv_path}")
plot_from_csv(csv_path, out_dir, args.prefix)
print(f"[ok] plots saved to: {out_dir}")
if __name__ == "__main__":
main()