diff --git a/scripts/width_diagrams.py b/scripts/width_diagrams.py new file mode 100644 index 0000000..69e64e9 --- /dev/null +++ b/scripts/width_diagrams.py @@ -0,0 +1,267 @@ +from __future__ import annotations +from typing import List, Tuple, Dict +import argparse, math, numpy as np, os +import matplotlib.pyplot as plt +import matplotlib.ticker as mtick +from multiprocessing import Pool, cpu_count +from pathlib import Path + +import squish.ordered as order +from squish import Simulation, DomainParams +from squish.common import OUTPUT_DIR + + +def order_process(domain: DomainParams) -> Tuple[float, float, float]: + energies = [] + configs = order.configurations(domain) + for config in configs: + energies.append( + 2 * domain.w * domain.h + + 2 + * math.pi + * domain.n + * (domain.r ** 2 - 2 * domain.r * order.avg_radius(domain, config)) + ) + + return domain.w, min(energies), max(energies) + + +def get_ordered_energies(orig_domain: DomainParams, widths: np.ndarray) -> Dict: + data = {} + domains = [ + DomainParams(orig_domain.n, w, orig_domain.h, orig_domain.r) for w in widths + ] + + with Pool(cpu_count()) as pool: + mins, maxes = {}, {} + for i, res in enumerate(pool.imap_unordered(order_process, domains)): + mins[res[0]] = res[1] + maxes[res[0]] = res[2] + + hashes = int(21 * i / len(widths)) + print( + f'Generating at width {res[0]:.02f}... |{"#"*hashes}{" "*(20-hashes)}|' + + f" {i+1}/{len(widths)} completed.", + flush=True, + end="\r", + ) + + print(flush=True) + + data["min"] = list([x[1] for x in sorted(mins.items())]) + data["max"] = list([x[1] for x in sorted(maxes.items())]) + + return data + + +def eq_file_process(file: Path) -> Tuple[float, List[float], List[float]]: + sim, frames = Simulation.load(file) + + alls = [] + for frame_info in frames: + alls.append( + [ + frame_info["energy"], + np.var(frame_info["stats"]["avg_radius"]) <= 1e-8, + np.count_nonzero(frame_info["stats"]["site_edge_count"] != 6), + ] + ) + + sim, frames = Simulation.load(file) + sim.frames = list(frames) + counts = sim.get_distinct() + + distincts = [] + for j, frame_info in enumerate(sim.frames): + distincts.append( + [ + frame_info["energy"], + np.var(frame_info["stats"]["avg_radius"]) <= 1e-8, + np.count_nonzero(frame_info["stats"]["site_edge_count"] != 6), + counts[j], + ] + ) + + return sim.domain.w, alls, distincts + + +def get_equilibria_data(filepath: Path) -> Tuple[Dict, numpy.ndarray, DomainParams]: + data = {"all": {}, "distinct": {}} + files = list(Path(filepath).iterdir()) + + with Pool(cpu_count()) as pool: + for i, res in enumerate(pool.imap_unordered(eq_file_process, files)): + data["all"][res[0]] = res[1] + data["distinct"][res[0]] = res[2] + + hashes = int(21 * i / len(files)) + print( + f'Loading simulations... |{"#"*hashes}{" "*(20-hashes)}|' + + f" {i+1}/{len(files)} simulations loaded.", + flush=True, + end="\r", + ) + print(flush=True) + + sim, frames = Simulation.load(files[0]) + widths = np.asarray(sorted(data["all"])) + domain = DomainParams(sim.domain.n, widths[-1], sim.domain.h, sim.domain.r) + return data, widths, domain + + +def axis_settings(ax, widths): + ax.invert_xaxis() + ax.grid(zorder=0) + ax.set_xticks([round(w, 2) for w in widths[::-2]]) + ax.set_xticklabels(ax.get_xticks(), rotation=90) + plt.subplots_adjust(0.07, 0.12, 0.97, 0.9) + + +def main(): + # Loading arguments. + parser = argparse.ArgumentParser("Outputs width search data into diagrams") + parser.add_argument( + "sims_path", + metavar="path/to/data", + help="folder that contains simulation files, or cached data file.", + ) + parser.add_argument( + "-q", + "--quiet", + dest="quiet", + action="store_true", + default=False, + help="suppress all normal output", + ) + + args = parser.parse_args() + + data, widths, domain = get_equilibria_data(Path(args.sims_path)) + order_data = get_ordered_energies(domain, widths) + + fig_folder = OUTPUT_DIR / Path(f"ShrinkEnergyComparison - N{domain.n}") + fig_folder.mkdir(exist_ok=True) + + # Torus minimum energies used as reference. + + # Probability of disorder diagram. + fig, ax = plt.subplots(figsize=(16, 8)) + all_disorder_count = [] + for width in widths: + equal_shape = list([c[1] for c in data["all"][width]]) + all_disorder_count.append( + 100 * equal_shape.count(False) / len(data["all"][width]) + ) + + ax.plot(widths, all_disorder_count) + axis_settings(ax, widths) + + ax.yaxis.set_major_formatter(mtick.PercentFormatter()) + ax.title.set_text(f"Probability of Disorder - N{domain.n}") + ax.set_xlabel("Width") + ax.set_ylabel("Disordered Equilibria") + boa_y_min = round(min(all_disorder_count) / 20) * 20 - 5 + ax.set_yticks(np.arange(boa_y_min, 100.01, 2.5)) + fig.savefig(fig_folder / "Probability of Disorder.png") + + # Density of States diagram. + fig, ax = plt.subplots(figsize=(16, 8)) + distinct_ordered, distinct_unordered = [], [] + for width in widths: + equal_shape = list([c[1] for c in data["distinct"][width]]) + distinct_ordered.append(equal_shape.count(True)) + distinct_unordered.append(equal_shape.count(False)) + + ax2 = ax.twinx() + ax.plot(widths, distinct_unordered, label="Unordered Equilibria", color="C0") + ax2.plot(widths, distinct_ordered, label="Ordered Equilibria", color="C1") + axis_settings(ax, widths) + ax.title.set_text(f"Density of States - N{domain.n}") + ax.set_xlabel("Width") + ax.set_ylabel("Number of States (Disordered)", color="C0") + ax2.set_ylabel("Number of States (Ordered)", color="C1") + + dos_y_max_unorder = 1.05 * max(distinct_unordered) + dos_y_max_order = 1.05 * max(distinct_ordered) + ax.set_yticks(np.linspace(0, dos_y_max_unorder, 20).astype(int)) + # ax.set_yticks(np.arange(0, dos_y_max_unorder, round(dos_y_max_unorder/200, 1)*10)) + ax2.set_yticks(np.arange(0, dos_y_max_order)) + + fig.savefig(fig_folder / "Density Of States.png") + + # Defect density diagram + fig, ax = plt.subplots(figsize=(16, 8)) + + defects = [] + for width in widths: + defects.append( + sum([c[2] for c in data["all"][width] if not c[1]]) + / len(data["all"][width]) + ) + + ax.plot(widths, defects) + axis_settings(ax, widths) + ax.title.set_text(f"Average Defects - N{domain.n}") + ax.set_xlabel("Width") + ax.set_ylabel("Defects") + ax.set_yticks(np.arange(0, 1 + max(defects), 0.5)) + fig.savefig(fig_folder / "Defects.png") + + # Bifurcation diagram + fig, ax = plt.subplots(figsize=(16, 8)) + + ordered_energies, unordered_energies = [], [] + for width in widths: + ordered_energies.append([c[0] for c in data["distinct"][width] if c[1]]) + unordered_energies.append([c[0] for c in data["distinct"][width] if not c[1]]) + + for i in range(len(order_data["min"])): + ordered_energies[i].append(order_data["min"][i]) + ordered_energies[i].append(order_data["max"][i]) + + null_unorder = [] + for i, energies in enumerate(unordered_energies): + if len(energies) == 0: + null_unorder.append(i) + energies.append(order_data["min"][i]) + + min_order = np.asarray([min(width) for width in ordered_energies]) + max_order = np.asarray([max(width) for width in ordered_energies]) + min_unorder = np.asarray([min(width) for width in unordered_energies]) + max_unorder = np.asarray([max(width) for width in unordered_energies]) + + offset = np.array(order_data["min"]) + # offset = np.array(min_order) + + min_unorder_off = min_unorder - offset + max_unorder_off = max_unorder - offset + ax.plot(widths, min_order - offset, color="C1") + # ax.plot(widths, max_order - offset, color='C1', linestyle='dotted') + ax.plot(widths, min_unorder_off, color="C0") + ax.plot(widths, max_unorder_off, color="C0", linestyle="dotted") + axis_settings(ax, widths) + + for i in null_unorder: + ax.scatter(widths[i], 0, marker="X", color="blue", s=50, zorder=4) + # ax.scatter(widths[i], max_unorder[i] - offset[i], + # marker='X', edgecolors="blue", facecolors='none', s=100, zorder=4) + + ax.title.set_text(f"Reduced Energy vs. Width - N{domain.n}") + ax.set_xlabel("Width") + ax.set_ylabel("Reduced Energy") + bif_y_max = np.max(np.abs(np.concatenate((min_unorder_off, max_unorder_off)))) + bif_top = np.arange( + 0, bif_y_max, round(bif_y_max / 20, -math.floor(math.log10(bif_y_max / 20))) + ) + ax.set_yticks(np.concatenate((-bif_top[1:][::-1], bif_top))) + fig.savefig(fig_folder / "Bifurcation.png") + + print(f"Wrote to {fig_folder}.") + + +if __name__ == "__main__": + os.environ["QT_LOGGING_RULES"] = "*=false" + try: + main() + except KeyboardInterrupt: + print("Program terminated by user.")