Merge branch 'master' of github.com:ksjdragon/squish

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.")