231 lines
6.7 KiB
Python
231 lines
6.7 KiB
Python
import argparse, numpy as np, os
|
|
import matplotlib.pyplot as plt
|
|
import matplotlib.ticker as mtick
|
|
from multiprocessing import Pool, cpu_count
|
|
from pathlib import Path
|
|
|
|
from squish import Simulation, DomainParams, Energy, ordered, Diagram
|
|
from squish.common import OUTPUT_DIR
|
|
from script_tools import (
|
|
RC_SETTINGS,
|
|
get_args,
|
|
get_data,
|
|
get_simulation_data,
|
|
get_ordered_data,
|
|
)
|
|
|
|
NAME = "Continuation"
|
|
|
|
|
|
def proc(combo):
|
|
dia, min_order, min_disorder, hex_ratios, alphas, sim_alphas, sim_energies, offset, length, num_frames = (
|
|
combo
|
|
)
|
|
sim = dia.sim
|
|
out_folder = sim.path / "continuation"
|
|
|
|
plt.rcParams.update(RC_SETTINGS)
|
|
plt.rcParams.update({"figure.constrained_layout.use": False})
|
|
fig = plt.figure(figsize=(30, 15))
|
|
|
|
e_hex = ordered.e_hex(sim.domains[0])
|
|
|
|
for i in range(length):
|
|
gs = fig.add_gridspec(1, 2)
|
|
gs.update(left=0, right=0.98, top=0.93, bottom=0.12, wspace=0.08)
|
|
ax = fig.add_subplot(gs[1])
|
|
|
|
curr_alpha = sim.domains[i].w / sim.domains[i].h
|
|
curr_vee = 100 * (sim.energies[i] / sim.domains[i].n - e_hex)
|
|
|
|
ax.plot(
|
|
alphas,
|
|
100 * min_order,
|
|
color="C1",
|
|
label="Min Ordered",
|
|
zorder=10,
|
|
linestyle="dotted",
|
|
)
|
|
ax.plot(
|
|
alphas,
|
|
100 * min_disorder,
|
|
color="C0",
|
|
label="Min Disordered",
|
|
zorder=11,
|
|
linestyle="dotted",
|
|
)
|
|
ax.plot(
|
|
sim_alphas[:175],
|
|
100 * sim_energies[:175],
|
|
color="C2",
|
|
label="_nolegend_",
|
|
linestyle="dashed",
|
|
)
|
|
ax.plot(
|
|
sim_alphas[174:], 100 * sim_energies[174:], color="C2", label="Continuation"
|
|
)
|
|
|
|
ax.scatter(
|
|
hex_ratios, [0] * len(hex_ratios), color="C2", s=120, marker="H", zorder=50
|
|
)
|
|
|
|
ax.scatter(
|
|
curr_alpha,
|
|
curr_vee,
|
|
s=250,
|
|
facecolors="none",
|
|
edgecolors="C6",
|
|
linewidth=4,
|
|
zorder=100,
|
|
)
|
|
|
|
ax.set_xlim(0.3, 1)
|
|
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)
|
|
|
|
# start, end = ax.get_ylim()
|
|
# space = np.linspace(0, end, 20)
|
|
|
|
space = np.arange(0, 10, 0.5)
|
|
ax.set_ylim(-0.15, 9.5)
|
|
ax.set_yticks(space[::-1])
|
|
ax.ticklabel_format(axis="y", style="sci")
|
|
|
|
ax.set_xlabel("Aspect Ratio")
|
|
ax.set_ylabel(r"VEE $\left[\times 10^{2}\right]$")
|
|
|
|
ax.legend(loc="upper center")
|
|
# ax.legend()
|
|
ax.grid(zorder=0)
|
|
|
|
props = dict(boxstyle="round", facecolor="white", alpha=0.8, zorder=20)
|
|
ax.text(
|
|
0.873,
|
|
0.97,
|
|
f"N={sim.domains[i].n}",
|
|
transform=ax.transAxes,
|
|
verticalalignment="top",
|
|
bbox=props,
|
|
)
|
|
|
|
ax1 = fig.add_subplot(gs[0])
|
|
dia.voronoi_plot(i, ax1)
|
|
fig.savefig(out_folder / f"img{i+offset:05}.png")
|
|
fig.clear()
|
|
|
|
j = len(list((sim.path / "continuation").iterdir()))
|
|
hashes = int(21 * j / num_frames)
|
|
print(
|
|
f'Generating frames... |{"#"*hashes}{" "*(20-hashes)}|'
|
|
+ f" {j}/{num_frames} frames rendered.",
|
|
flush=True,
|
|
end="\r",
|
|
)
|
|
|
|
print(flush=True)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Makes video of analytic continuation."
|
|
)
|
|
parser.add_argument(
|
|
"sims_path",
|
|
metavar="sim_dir",
|
|
help="folder that contains simulation data at various aspect ratios for some N",
|
|
)
|
|
parser.add_argument(
|
|
"shrink_sim",
|
|
metavar="continuation_sim",
|
|
help="simulation that contains the continuation data",
|
|
)
|
|
args = parser.parse_args()
|
|
sims_path = Path(args.sims_path)
|
|
|
|
data, n, r = get_data(
|
|
sims_path / "package.pkl", get_simulation_data, args=(sims_path,)
|
|
)
|
|
domain, alphas = DomainParams(n, 1, 1, r), data["all"]["alpha"]
|
|
ordered_data = get_data(
|
|
OUTPUT_DIR / "OrderedCache" / f"{n}.pkl",
|
|
get_ordered_data,
|
|
args=(domain, alphas),
|
|
)
|
|
|
|
min_disorder, max_disorder = [], []
|
|
for i, energies in enumerate(data["distinct"]["Energy"]):
|
|
disorder_energies = []
|
|
for j, energy in enumerate(energies):
|
|
if not data["distinct"]["Ordered"][i][j]:
|
|
disorder_energies.append(energy)
|
|
min_disorder.append(min(disorder_energies))
|
|
max_disorder.append(max(disorder_energies))
|
|
|
|
min_order, min_order_coer = [], []
|
|
for i, energies in enumerate(ordered_data["Energy"]):
|
|
min_order.append(energies[0])
|
|
min_order_coer.append(ordered_data["Coercivity"][i][0])
|
|
e_hex = ordered.e_hex(domain)
|
|
min_disorder = np.array(min_disorder) / domain.n - e_hex
|
|
max_disorder = np.array(max_disorder) / domain.n - e_hex
|
|
min_order = np.array(min_order) / domain.n - e_hex
|
|
|
|
hex_ratios = ordered.hexagon_alpha(domain.n)
|
|
|
|
sim = Simulation.from_file(Path(args.shrink_sim))
|
|
sim_alphas = np.array([x.w / x.h for x in sim.frames])
|
|
sim_energies = np.array([x.energy for x in sim.frames]) / sim.domain.n - e_hex
|
|
dia = Diagram(sim, ["voronoi"])
|
|
|
|
out_folder = sim.path / "continuation"
|
|
out_folder.mkdir(exist_ok=True)
|
|
|
|
combo_list = []
|
|
for i in range(cpu_count()):
|
|
start, end = (
|
|
int(i * len(sim) / cpu_count()),
|
|
int((i + 1) * len(sim) / cpu_count()),
|
|
)
|
|
new_dia = Diagram(None, ["voronoi"])
|
|
new_dia.sim = dia.sim.slice(list(range(start, end)))
|
|
combo_list.append(
|
|
(
|
|
new_dia,
|
|
min_order,
|
|
min_disorder,
|
|
hex_ratios,
|
|
alphas,
|
|
sim_alphas,
|
|
sim_energies,
|
|
start,
|
|
len(sim.frames[start:end]),
|
|
len(sim),
|
|
)
|
|
)
|
|
|
|
print("Starting figure generation...", flush=True)
|
|
with Pool(cpu_count()) as pool:
|
|
for _ in pool.imap_unordered(proc, combo_list):
|
|
pass
|
|
|
|
video_path = sim.path / (NAME + ".mp4")
|
|
|
|
fps = 30
|
|
print("Assembling MP4...", flush=True)
|
|
os.system(
|
|
f"ffmpeg -hide_banner -loglevel error -r {fps} -i"
|
|
+ f' "{sim.path}/continuation/img%05d.png"'
|
|
+ f" -c:v libx264 -crf 18 -preset slow -pix_fmt yuv420p -vf"
|
|
+ f' "scale=trunc(iw/2)*2:trunc(ih/2)*2" -f mp4 "{video_path}"'
|
|
)
|
|
|
|
print(f"Wrote to {video_path}.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
os.environ["QT_LOGGING_RULES"] = "*=false"
|
|
try:
|
|
main()
|
|
except KeyboardInterrupt:
|
|
print("Program terminated by user.")
|