squish/scripts/density_of_states.py

64 lines
1.8 KiB
Python

from __future__ import annotations
from typing import List, Tuple, Dict
import argparse, math, numpy as np, os, pickle, itertools
import matplotlib.pyplot as plt
import matplotlib.ticker as mtick
from pathlib import Path
from squish import Simulation, DomainParams
from squish.common import OUTPUT_DIR
from script_tools import RC_SETTINGS, get_args, get_data, get_simulation_data
def main():
sims_path, regen = get_args(
"Density of states of various N across alpha",
"folders that contains various N simulations to plot",
)
packages = []
for fol in sims_path.iterdir():
if fol.is_file():
continue
data, n, r = get_data(
fol / "package.pkl", get_simulation_data, args=(fol,), regen=regen
)
domain = DomainParams(n, 1, 1, r)
packages.append([data, domain])
packages.sort(key=lambda x: x[1].n)
plt.rcParams.update(RC_SETTINGS)
fig = plt.figure(figsize=(18, 15))
gs = fig.add_gridspec(1, 1)
ax = fig.add_subplot(gs[0])
for j, package in enumerate(packages):
data, domain = package
distinct_disordered = []
for ordered in data["distinct"]["Ordered"]:
distinct_disordered.append(np.count_nonzero(ordered == False))
ax.plot(data["all"]["alpha"], distinct_disordered, label=f"N={domain.n}")
alphas = packages[0][0]["all"]["alpha"]
ax.set_xticks([round(w, 2) for w in alphas[::10]])
ax.set_xticklabels([f"{round(w, 3):.2f}" for w in alphas[::10]], rotation=90)
ax.set_xlim(0.3, 1.0)
ax.set_xlabel("Aspect Ratio")
ax.set_ylabel("Number of Distinct Equilibria")
ax.grid(zorder=0)
ax.legend(loc="center right", fancybox=True, bbox_to_anchor=(1.34, 0.5))
fig.savefig(OUTPUT_DIR / "DoS.png")
print(f"Wrote to {OUTPUT_DIR / 'DoS.png'}")
if __name__ == "__main__":
main()