{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5bfb0f88",
   "metadata": {},
   "source": [
    "## Figure 2 From Paper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7cf5c6fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import symd\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "import numpy as np\n",
    "import pickle\n",
    "import pandas as pd\n",
    "import skunk\n",
    "import svglib\n",
    "import seaborn as sns\n",
    "import gzip\n",
    "import urllib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f12e496f",
   "metadata": {},
   "outputs": [],
   "source": [
    "base_colors = [\n",
    "    \"f94144\",\n",
    "    \"f3722c\",\n",
    "    \"f8961e\",\n",
    "    \"f9844a\",\n",
    "    \"f9c74f\",\n",
    "    \"90be6d\",\n",
    "    \"43aa8b\",\n",
    "    \"4d908e\",\n",
    "    \"577590\",\n",
    "    \"277da1\",\n",
    "]\n",
    "colors = [\"#\" + c for c in base_colors]\n",
    "sns.set_style(\"white\")\n",
    "sns.set_style(\"ticks\")\n",
    "sns.set(\n",
    "    rc={\n",
    "        \"axes.facecolor\": \"#f5f4e9\",\n",
    "        \"grid.color\": \"#AAAAAA\",\n",
    "        \"axes.edgecolor\": \"#333333\",\n",
    "        \"figure.facecolor\": \"#FFFFFF\",\n",
    "        \"axes.grid\": False,\n",
    "        \"axes.prop_cycle\": plt.cycler(\"color\", plt.cm.Dark2.colors),\n",
    "        \"font.family\": \"monospace\",\n",
    "    }\n",
    ")\n",
    "titles = [\n",
    "    \"p1\",\n",
    "    \"p2\",\n",
    "    \"pm\",\n",
    "    \"pg\",\n",
    "    \"cm\",\n",
    "    \"pmm\",\n",
    "    \"pmg\",\n",
    "    \"pgg\",\n",
    "    \"cmm\",\n",
    "    \"p4\",\n",
    "    \"p4m\",\n",
    "    \"p4g\",\n",
    "    \"p3\",\n",
    "    \"p3m1\",\n",
    "    \"p31m\",\n",
    "    \"p6\",\n",
    "    \"p6m\",\n",
    "]\n",
    "print(symd.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea0505b8",
   "metadata": {},
   "source": [
    "### Make Atlas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b8457aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "def scale(x, ib):\n",
    "    s = np.apply_along_axis(lambda xi: ib @ xi, 1, x)\n",
    "    return np.fmod(s, 1.0)\n",
    "\n",
    "\n",
    "def compute_symm(positions, gnum, cell, ndim, n):\n",
    "    group = symd.groups.load_group(gnum, ndim)\n",
    "    cell = np.array(cell).reshape(ndim, ndim)\n",
    "    ib = np.linalg.inv(cell)\n",
    "    s = scale(positions[:, :ndim], ib)\n",
    "    members = [symd.groups.str2mat(e) for e in group[\"genpos\"]]\n",
    "    folded_positions = np.zeros_like(s)\n",
    "    for i in range(n):\n",
    "        folded_positions[i, :] = s[i, :]\n",
    "        for j in range(1, len(members)):\n",
    "            k = j * n + i\n",
    "            im = np.linalg.inv(members[j])\n",
    "            w = im[:ndim, :ndim]\n",
    "            folded_positions[i, :] += w @ s[k] + im[ndim, :ndim]\n",
    "        folded_positions[i, :] /= len(members)\n",
    "        for j in range(1, len(members)):\n",
    "            k = j * n + i\n",
    "            w = members[j][:ndim, :ndim]\n",
    "            folded_positions[k] = w @ folded_positions[i] + members[j][ndim, :ndim]\n",
    "    return np.mean((s[:k] - folded_positions[:k]) ** 2)\n",
    "\n",
    "\n",
    "def rmsd(p1, p2):\n",
    "    return np.mean((p1 - p2) ** 2, axis=(1, 2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cefcbe23",
   "metadata": {},
   "outputs": [],
   "source": [
    "def crystal(\n",
    "    n,\n",
    "    group,\n",
    "    w=None,\n",
    "    retries=5,\n",
    "    steps=10**6,\n",
    "    steps2=5 * 10**3,\n",
    "    ndims=2,\n",
    "    starting_density=0.2,\n",
    "    positions=False,\n",
    "):\n",
    "    # trying to have n be number in UNIT CELLL\n",
    "    # so have to adjust for group size\n",
    "    m = len(symd.groups.load_group(group, ndims).genpos)\n",
    "    n = max(2, n // m)\n",
    "    if w is not None:\n",
    "        n += sum(w)\n",
    "        name = f\"{group}-{n}-{sum(w)}\"\n",
    "    else:\n",
    "        name = f\"{group}-{n}\"\n",
    "    print(\"Simulating\", n, \"particles:\", name)\n",
    "    # break out the try/except because we will accept failed NPT (because it jams so hard)\n",
    "    for i in range(retries):\n",
    "        np.random.seed(i)\n",
    "        cell = symd.groups.get_cell(starting_density, group, 2, n, w)\n",
    "        # NPT\n",
    "        md = symd.Symd(\n",
    "            nparticles=n,\n",
    "            cell=cell,\n",
    "            ndims=ndims,\n",
    "            images=2,\n",
    "            force=\"lj\",\n",
    "            wyckoffs=w,\n",
    "            group=group,\n",
    "            steps=steps,\n",
    "            exeDir=f\"crystal-{name}\",\n",
    "            pressure=0.25,\n",
    "            temperature=0.1,\n",
    "            start_temperature=0.5,\n",
    "        )\n",
    "        try:\n",
    "            md.remove_overlap()\n",
    "        except RuntimeError as e:\n",
    "            continue\n",
    "        if positions:\n",
    "            md.log_positions(period=250)\n",
    "        else:\n",
    "            md.log_positions()\n",
    "        try:\n",
    "            md.run()\n",
    "        except RuntimeError as e:\n",
    "            d = md.number_density()\n",
    "            if d < 0.5:\n",
    "                print(\"Not dense enough, retrying\", d)\n",
    "                continue\n",
    "\n",
    "        # NVT\n",
    "        md.runParams[\"start_temperature\"] = 0.05\n",
    "        md.runParams[\"temperature\"] = 1e-4\n",
    "        md.runParams[\"box_update_period\"] = 0\n",
    "        md.runParams[\"langevin_gamma\"] = 0.5\n",
    "        md.runParams[\"steps\"] = steps // 4\n",
    "        md.log_positions(filename=\"equil.xyz\")\n",
    "        try:\n",
    "            md.run()\n",
    "        except RuntimeError as e:\n",
    "            continue\n",
    "        config = md.positions[-1]\n",
    "\n",
    "        # Stability\n",
    "        fp = np.loadtxt(md.runParams[\"final_positions\"])\n",
    "        # changing group, so need to read projected cell\n",
    "        cell = md.read_cell(bravais=True)\n",
    "        m = fp.shape[0]\n",
    "        md2 = symd.Symd(\n",
    "            nparticles=m,\n",
    "            cell=cell,\n",
    "            ndims=2,\n",
    "            images=2,\n",
    "            force=\"lj\",\n",
    "            wyckoffs=None,\n",
    "            group=1,\n",
    "            steps=steps2,\n",
    "            exeDir=f\"melt-{name}\",\n",
    "            temperature=None,\n",
    "            start_temperature=0.0,\n",
    "        )\n",
    "        # run once to get melting traj\n",
    "        # then again for longer with longer period\n",
    "        md2.log_positions(period=10)\n",
    "        md2.runParams[\"start_positions\"] = md.runParams[\"final_positions\"]\n",
    "        try:\n",
    "            md2.run()\n",
    "        except RuntimeError as e:\n",
    "            continue\n",
    "        traj = md2.positions\n",
    "        csm = rmsd(md2.positions[:, :m], md2.positions[0, :m])\n",
    "        # csm = []\n",
    "        # for i in range(md2.positions.shape[0]):\n",
    "        #    csm.append(compute_symm(md2.positions[i], group, md2.read_cell(), ndims, n))\n",
    "        if positions:\n",
    "            return np.concatenate((md.positions, md2.positions))\n",
    "        return (\n",
    "            config,\n",
    "            md2.positions[-1],\n",
    "            md2.number_density(),\n",
    "            csm,\n",
    "            traj,\n",
    "            np.arange(0, steps2, 10) * md2.runParams[\"time_step\"],\n",
    "        )\n",
    "    return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4f1bbdb",
   "metadata": {},
   "outputs": [],
   "source": [
    "config, config2, nd, csm, traj, time = crystal(10, 6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c247332",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(8, 8))\n",
    "plt.title(f\"{nd}\")\n",
    "plt.plot(config[:, 0], config[:, 1], \".\")\n",
    "plt.plot(config2[:, 0], config2[:, 1], \".\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed63c371",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "plt.plot(time, csm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "412f19dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Turned off\n",
    "if False:\n",
    "    from multiprocessing import Pool\n",
    "\n",
    "    cdf = None\n",
    "    results = []\n",
    "    trajs = {}\n",
    "\n",
    "    with Pool() as pool:\n",
    "        for N in [8, 16, 32, 128]:\n",
    "            for i, t in enumerate(titles):\n",
    "                W = len(symd.groups.load_group(i + 1, 2)[\"specpos\"])\n",
    "                for j in range(1 + W):\n",
    "                    wycks = None if j == 0 else [1] * j\n",
    "                    name = f\"{t}-w{j}-n{N}\"\n",
    "                    job = pool.apply_async(crystal, (N, i + 1, wycks))\n",
    "                    # job = crystal(N, i+1, wycks)\n",
    "                    results.append((t, name, N, j, job))\n",
    "\n",
    "        for r in results:\n",
    "            t, name, N, j, ar = r\n",
    "            print(\"Getting result for \", name)\n",
    "            res = ar.get()\n",
    "            # res = ar\n",
    "            if res is None:\n",
    "                continue\n",
    "            config, config2, nd, csm, traj, time = res\n",
    "            T = len(csm)\n",
    "\n",
    "            df2 = pd.DataFrame(\n",
    "                {\n",
    "                    \"Group\": [t] * T,\n",
    "                    \"Traj\": [name] * T,\n",
    "                    \"rho\": [nd] * T,\n",
    "                    \"N\": [N] * T,\n",
    "                    \"Wyckoffs\": [str(j)] * T,\n",
    "                    \"RMSD\": csm,\n",
    "                    \"Time\": time,\n",
    "                }\n",
    "            )\n",
    "            if cdf is None:\n",
    "                cdf = df2\n",
    "            else:\n",
    "                cdf = pd.concat((cdf, df2))\n",
    "            trajs[name] = traj\n",
    "\n",
    "    cdf.reset_index(inplace=True)\n",
    "    cdf.to_pickle(\"atlas2d.pkl.gz\")\n",
    "    with open(\"atlas2d.traj.pkl\", \"wb\") as f:\n",
    "        pickle.dump(trajs, f, pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23369349",
   "metadata": {},
   "source": [
    "### Plot Figure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a4a4428",
   "metadata": {},
   "outputs": [],
   "source": [
    "urllib.request.urlretrieve(\n",
    "    \"https://www.dropbox.com/s/m6gi1ecv66ylm06/atlas2d.traj.pkl.gz?dl=1\",\n",
    "    \"atlas2d.traj.pkl.gz\",\n",
    ")\n",
    "urllib.request.urlretrieve(\n",
    "    \"https://www.dropbox.com/s/smn8tvljxhlp75f/atlas2d.pkl.gz?dl=1\", \"atlas2d.pkl.gz\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8cc5a12",
   "metadata": {},
   "outputs": [],
   "source": [
    "cdf = pd.read_pickle(\"atlas2d.pkl.gz\")\n",
    "with gzip.open(\"atlas2d.traj.pkl.gz\", \"rb\") as f:\n",
    "    trajs = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08460bc8",
   "metadata": {},
   "outputs": [],
   "source": [
    "cdf.query(\"rho > 0.5\").Traj.unique().shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94be7fdc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_config(pos, figsize=(1.5, 1.5), color=\"#333333\"):\n",
    "    N, D = pos.shape\n",
    "    fig, ax = plt.subplots(figsize=figsize)\n",
    "    ax.plot(\n",
    "        pos[:, 0],\n",
    "        pos[:, 1],\n",
    "        color=color,\n",
    "        marker=\"o\",\n",
    "        markersize=8,\n",
    "        markeredgewidth=0.0,\n",
    "        linestyle=\"None\",\n",
    "    )\n",
    "    ax.set_facecolor(\"#f5f4e9\")\n",
    "    fig.patch.set_facecolor(\"#f5f4e9\")\n",
    "    ax.axis(\"off\")\n",
    "    # want about 64 points\n",
    "    # q = max(0.05, min(0.5, 64 / N))\n",
    "    # xlim = np.quantile(pos[:,0], [0.5 - q, 0.5 + q])\n",
    "    # ylim = np.quantile(pos[:,1], [0.5 - q, 0.5 + q])\n",
    "    xlim = (-3, 3)\n",
    "    ylim = (-3, 3)\n",
    "    ax.set_aspect(\"equal\")\n",
    "    ax.set_xlim(*xlim)\n",
    "    ax.set_ylim(*ylim)\n",
    "    plt.tight_layout()\n",
    "    return skunk.pltsvg(fig=fig)\n",
    "\n",
    "\n",
    "skunk.display(plot_config(trajs[\"pg-w0-n16\"][0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2631a13b",
   "metadata": {},
   "outputs": [],
   "source": [
    "top = cdf[cdf.Time == cdf.Time.max()].sort_values(by=[\"RMSD\"])[:25]\n",
    "for r, t in zip(top.RMSD.values, top.Traj.values):\n",
    "    plot_config(trajs[t][0])\n",
    "    plt.title(f\"{t} = {r}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73755a4b",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "replaces = []\n",
    "\n",
    "\n",
    "def annotate_config(rmsd, traj, time, wyckoffs, color):\n",
    "    late_rmsd = rmsd.where(time > 15)\n",
    "    idx = late_rmsd.argmin()\n",
    "    label = traj.iloc[idx]\n",
    "    w = wyckoffs.iloc[idx]\n",
    "    x = time.iloc[idx]\n",
    "    y = rmsd.iloc[idx]\n",
    "    if y > 1:\n",
    "        x = 0\n",
    "        y = 0\n",
    "    replaces.append((label, f\"C{w}\"))\n",
    "    ax = plt.gca()\n",
    "    box = skunk.Box(25, 25, label)\n",
    "    ab = mpl.offsetbox.AnnotationBbox(\n",
    "        box,\n",
    "        (x, y),\n",
    "        pad=0,\n",
    "        bboxprops=dict(edgecolor=\"#333\", linewidth=1),\n",
    "        xybox=(0.7 if x == 0 else 0.2, 0.7),\n",
    "        xycoords=\"data\",\n",
    "        boxcoords=\"axes fraction\",\n",
    "        arrowprops=dict(arrowstyle=\"->\", color=\"#333\"),\n",
    "    )\n",
    "\n",
    "    ax.add_artist(ab)\n",
    "\n",
    "\n",
    "g = sns.relplot(\n",
    "    data=cdf.query(\"rho > 0.5\"),\n",
    "    x=\"Time\",\n",
    "    y=\"RMSD\",\n",
    "    kind=\"line\",\n",
    "    hue=\"Wyckoffs\",\n",
    "    col=\"Group\",\n",
    "    col_wrap=6,\n",
    "    aspect=1,\n",
    "    linewidth=1,\n",
    "    style=\"N\",\n",
    "    height=1.75,\n",
    "    palette=\"Dark2\",\n",
    "    hue_order=[str(i) for i in range(9)],\n",
    ")\n",
    "plt.ylim(-0.1, 1)\n",
    "# sns.move_legend(g, \"lower right\", bbox_to_anchor=(0.9,0.1), ncol=3)\n",
    "g.map(annotate_config, \"RMSD\", \"Traj\", \"Time\", \"Wyckoffs\").set_axis_labels(\n",
    "    r\"Time [$\\tau$]\", \"RMSD\"\n",
    ")\n",
    "\n",
    "main_svg = skunk.pltsvg()\n",
    "svg = skunk.insert(\n",
    "    {l: plot_config(trajs[l][0], color=c) for l, c in replaces}, svg=main_svg\n",
    ")\n",
    "skunk.display(svg)\n",
    "with open(\"atlas.svg\", \"w\") as f:\n",
    "    f.write(svg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a93f8a1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from cairosvg import svg2png\n",
    "\n",
    "svg2png(bytestring=svg, write_to=\"atlas_raster.png\", dpi=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecebae8b",
   "metadata": {},
   "source": [
    "### TOC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4b338dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(titles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7a83e87",
   "metadata": {},
   "outputs": [],
   "source": [
    "name = \"p3m1-w0-n128\"\n",
    "group_i = titles.index(name.split(\"-\")[0]) + 1\n",
    "n_i = int(name.split(\"-n\")[1])\n",
    "group = symd.groups.load_group(group_i, 2)\n",
    "gp = len(group.genpos)\n",
    "pos = trajs[name][0]\n",
    "\n",
    "N, D = pos.shape\n",
    "\n",
    "# build out colors\n",
    "M = max(2, n_i // gp)\n",
    "print(n_i, gp, N, M, N // M)\n",
    "c = [colors[i % len(colors)] for i in range(M)]\n",
    "for i in range(1, N // M):\n",
    "    c.extend([colors[i % len(colors)] for i in range(M)])\n",
    "#    # which particle is being duplicated\n",
    "#\n",
    "#    c.extend([colors[((i - 1) % M) % len(colors)]] * M)\n",
    "#    #c.extend(['#333'] * M)\n",
    "print(len(c))\n",
    "\n",
    "figsize = (3.33, 1.88)\n",
    "aspect = figsize[0] / figsize[1]\n",
    "fig, ax = plt.subplots(figsize=figsize, dpi=300)\n",
    "ax.scatter(pos[:, 0], pos[:, 1], marker=\"o\", s=2, c=c)\n",
    "# want about 64 points\n",
    "q = 0.4\n",
    "xlim = np.quantile(pos[:, 0], [0.5 - q, 0.5 + q])\n",
    "ylim = np.quantile(pos[:, 1], [0.5 - q, 0.5 + q])\n",
    "yd = (ylim[1] - ylim[0]) / aspect\n",
    "yc = (ylim[1] + ylim[0]) / 2\n",
    "ylim = (yc - yd, yc + yd)\n",
    "# ax.set_xlim(*xlim)\n",
    "# ax.set_ylim(*ylim)\n",
    "x = 20\n",
    "ax.set_xlim(-x, x)\n",
    "ax.set_ylim(-x / aspect, x / aspect)\n",
    "ax.set_aspect(\"equal\")\n",
    "ax.set_facecolor(\"#f5f4e9\")\n",
    "fig.patch.set_facecolor(\"#f5f4e9\")\n",
    "ax.axis(\"off\")\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"toc.tiff\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae522963",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}