compression.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb286944",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax.config import config\n",
    "config.update(\"jax_enable_x64\", True)\n",
    "\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.animation import FFMpegFileWriter\n",
    "from tqdm.auto import tqdm\n",
    "from functools import lru_cache\n",
    "from plot_tools import texnum\n",
    "\n",
    "from jarzynski import init_piston, forward, imap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd1d4512",
   "metadata": {},
   "outputs": [],
   "source": [
    "init_piston_j = jax.jit(init_piston, static_argnums=1)\n",
    "forward_j = jax.jit(forward)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c3e71f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "state = init_piston_j(jax.random.PRNGKey(0), 20, 0.0)\n",
    "forward_j(1.0, state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65daddc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def experiment(key, xs, n, vel, r):\n",
    "    state = init_piston_j(key, n, r)\n",
    "    state['walls']['v'] = jnp.array([\n",
    "        [0.0, 0.0, vel/2],\n",
    "        [0.0, 0.0, -vel/2],\n",
    "    ])\n",
    "    t = 0\n",
    "    ws = [0]\n",
    "    for x in xs[1:]:\n",
    "        dt = x / vel - t\n",
    "        n, state, work = forward_j(dt, state)\n",
    "        t += dt\n",
    "\n",
    "        ws += [ws[-1] + work]\n",
    "\n",
    "    return jnp.array(ws)\n",
    "\n",
    "xs = jnp.linspace(0.0, 1.5, 100)\n",
    "\n",
    "@lru_cache\n",
    "def foobar(n, vel, r, rep):\n",
    "    keys = jax.random.split(jax.random.PRNGKey(1), num=rep)\n",
    "    return jnp.stack([experiment(key, xs, n, vel, r) for key in tqdm(keys)], axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23a6d726",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "Va = 2 * jnp.pi\n",
    "Vb = (2 - xs) * jnp.pi\n",
    "\n",
    "def Z(n, V, r):\n",
    "    vol = 4/3 * jnp.pi * r**3\n",
    "    return V**n * (1 - n * (n-1) / 2 * vol / V)\n",
    "\n",
    "plt.style.use('classic')\n",
    "\n",
    "fig, axs = plt.subplots(1, 3, dpi=100, figsize=(8, 3))\n",
    "\n",
    "exps = [(5, 0.01, 0.1), (5, 2.0, 0.1), (50, 0.01, 0.1)]\n",
    "# exps = [(10, 0.01, 1e-6)]\n",
    "\n",
    "for ax, (n, vel, r) in zip(axs, exps):\n",
    "    plt.sca(ax)\n",
    "    works = foobar(n, vel, r)\n",
    "\n",
    "    for m in [1, 10, works.shape[1]]:\n",
    "        plt.plot(\n",
    "            Z(n, Vb, r) / Z(n, Va, r), \n",
    "            jnp.exp(-works[:, :m]).mean(1), \n",
    "            '-', \n",
    "            label=fr'$N_{{\\mathrm{{avg}}}} = {m}$'\n",
    "        )\n",
    "\n",
    "    plt.legend(loc=0, fontsize=12, labelspacing=0, handlelength=1.5, handletextpad=0.2, frameon=False)\n",
    "    plt.plot([0,1], [0,1], '--k')\n",
    "    plt.xscale('log')\n",
    "    plt.yscale('log')\n",
    "    plt.title(fr'$n={n}$, $v_p={vel}$', fontsize=13)\n",
    "\n",
    "plt.sca(axs[0])\n",
    "plt.xlim(1e-4, 1)\n",
    "plt.ylim(1e-4, 1)\n",
    "plt.yticks([])\n",
    "\n",
    "plt.sca(axs[1])\n",
    "plt.xlim(1e-4, 1)\n",
    "plt.ylim(1e-4, 1)\n",
    "plt.yticks([])\n",
    "\n",
    "plt.sca(axs[2])\n",
    "plt.xlim(3e-7, 1)\n",
    "plt.ylim(3e-7, 1)\n",
    "plt.yticks([])\n",
    "\n",
    "plt.sca(axs[1])\n",
    "plt.xlabel(r'$e^{-\\Delta F / T} = Z_B/Z_A$')\n",
    "plt.sca(axs[0])\n",
    "plt.ylabel(r'$\\langle e^{-W / T} \\rangle$')\n",
    "\n",
    "plt.tight_layout(pad=0.1)\n",
    "\n",
    "plt.savefig('je.pgf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14f66d65",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e31680c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def random_piston_positions(key, z, n, r):\n",
    "    def sample(key):\n",
    "        key, k = jax.random.split(key)\n",
    "        x = jax.random.uniform(k, (3,), minval=-1.0, maxval=1.0)\n",
    "        x = x * jnp.array([1 - r, 1 - r, z - r])\n",
    "        return key, x\n",
    "\n",
    "    def body_for(i, args):\n",
    "        key, pos, ok = args\n",
    "\n",
    "        def cond(args):\n",
    "            key, x, ok = args\n",
    "            test = (jnp.linalg.norm(x[:2]) < 1.0 - r) & jnp.all(jnp.linalg.norm(pos - x, axis=1) > 2 * r)\n",
    "            return ~test\n",
    "\n",
    "        def body(args):\n",
    "            key, x, ok = args\n",
    "            return sample(key) + (0,)\n",
    "\n",
    "        key, x = sample(key)\n",
    "        key, x, ok = jax.lax.while_loop(cond, body, (key, x, ok))\n",
    "        return key, pos.at[i].set(x), ok\n",
    "\n",
    "    pos = 2.0 * jnp.ones((n, 3))\n",
    "    key, pos, ok = jax.lax.fori_loop(0, n, body_for, (key, pos, 1))\n",
    "\n",
    "    return key, pos, ok\n",
    "\n",
    "\n",
    "def partition_function(key, z, n, r, num):\n",
    "    key = jax.random.split(key, num=num)\n",
    "    key, pos, ok = imap(random_piston_positions, (0, None, None, None), 0)(key, z, n, r)\n",
    "    V = (2 - 2 * r)**2 * (2 * z - 2 * r)\n",
    "    return V**n * jnp.mean(ok)\n",
    "\n",
    "\n",
    "partition_function = jax.jit(partition_function, static_argnums=(2, 4))\n",
    "\n",
    "\n",
    "@lru_cache\n",
    "def partition_function2(n, r, num):\n",
    "    return jnp.stack([\n",
    "        partition_function(jax.random.PRNGKey(0), 1 - x / 2, n, r, num)\n",
    "        for x in tqdm(xs)\n",
    "    ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39bd9674",
   "metadata": {},
   "outputs": [],
   "source": [
    "partition_function2(8, 0.1, 200_000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "434e8684",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.style.use('classic')\n",
    "\n",
    "Va = 2 * jnp.pi\n",
    "Vb = (2 - xs) * jnp.pi\n",
    "\n",
    "def Z(n, V, r):\n",
    "    vol = 4/3 * jnp.pi * r**3\n",
    "    return V**n * (1 - n * (n-1) / 2 * vol / V)\n",
    "\n",
    "fig, axs = plt.subplots(1, 2, dpi=100, figsize=(8, 3))\n",
    "\n",
    "exps = [(8, 0.01, 1e-6, 2000), (8, 0.01, 0.1, 2000), ]\n",
    "\n",
    "for ax, (n, vel, r, rep) in zip(axs, exps):\n",
    "    plt.sca(ax)\n",
    "    works = foobar(n, vel, r, rep)\n",
    "    zs = partition_function2(n, r, 200_000)\n",
    "\n",
    "    \n",
    "    plt.plot(\n",
    "        Z(n, Vb, r) / Z(n, Va, r), \n",
    "        jnp.exp(-works).mean(1), \n",
    "        '-b', \n",
    "        label=fr'approximate Z'\n",
    "    )\n",
    "#     for rep in jnp.logspace(1, jnp.log10(works.shape[1]), 10):\n",
    "    plt.plot(\n",
    "        zs / zs[0], \n",
    "        jnp.exp(-works[:, :]).mean(1), \n",
    "        '-r', \n",
    "        label=fr'numerical Z'\n",
    "    )\n",
    "\n",
    "    plt.legend(loc=0, fontsize=12, labelspacing=0, handlelength=1.5, handletextpad=0.2, frameon=False)\n",
    "    plt.plot([0,1], [0,1], '--k')\n",
    "    plt.xscale('log')\n",
    "    plt.yscale('log')\n",
    "    plt.title(fr'$n={n}$, $v_p={texnum(vel)}$, $r={texnum(r)}$', fontsize=13)\n",
    "    plt.yticks([])\n",
    "    plt.xlabel(r'$e^{-\\Delta F / T} = Z_B/Z_A$')\n",
    "\n",
    "plt.sca(axs[0])\n",
    "plt.xlim(1e-5, 1)\n",
    "plt.ylim(1e-5, 1)\n",
    "plt.sca(axs[1])\n",
    "plt.xlim(1e-7, 1)\n",
    "plt.ylim(1e-7, 1)\n",
    "\n",
    "plt.sca(axs[0])\n",
    "plt.ylabel(r'$\\langle e^{-W / T} \\rangle$')\n",
    "\n",
    "plt.tight_layout(pad=0.1)\n",
    "\n",
    "plt.savefig('je.pgf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed66d1c1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cdd8fd86",
   "metadata": {},
   "outputs": [],
   "source": [
    "works.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afecf35b",
   "metadata": {},
   "outputs": [],
   "source": [
    "zs.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ea39f32",
   "metadata": {},
   "outputs": [],
   "source": [
    "works.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed2b0c30",
   "metadata": {},
   "outputs": [],
   "source": [
    "zs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c38a13bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "i = -1\n",
    "\n",
    "rep = 1 + jnp.arange(works.shape[1])\n",
    "error = jnp.exp(-works[i, :]).cumsum() / rep - (zs[i] / zs[0])\n",
    "\n",
    "plt.plot(rep, error, '.-')\n",
    "# plt.xscale('log')\n",
    "# plt.yscale('log')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b5bbc1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(works[-1], bins=40);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d84d5f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "(works < 0.0).sum(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fcd6f880",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7d5ebe2",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e154e0b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "Va = 2 * jnp.pi\n",
    "Vb = (2 - xs) * jnp.pi\n",
    "\n",
    "def Z(n, V, r):\n",
    "    vol = 4/3 * jnp.pi * r**3\n",
    "    return V**n * (1 - n * (n-1) / 2 * vol / V)\n",
    "\n",
    "plt.style.use('classic')\n",
    "\n",
    "fig, axs = plt.subplots(1, 2, dpi=100, figsize=(8, 3))\n",
    "\n",
    "exps = [(8, 0.01, 1e-6, 1000), (8, 0.01, 0.1, 1000), ]\n",
    "\n",
    "for ax, (n, vel, r, rep) in zip(axs, exps):\n",
    "    plt.sca(ax)\n",
    "    works = foobar(n, vel, r, rep)\n",
    "\n",
    "    for m in [1, 10, works.shape[1]]:\n",
    "        plt.plot(\n",
    "            Z(n, Vb, r) / Z(n, Va, r), \n",
    "            jnp.exp(-works[:, :m]).mean(1), \n",
    "            '-', \n",
    "            label=fr'$N_{{\\mathrm{{avg}}}} = {m}$'\n",
    "        )\n",
    "\n",
    "    plt.legend(loc=0, fontsize=12, labelspacing=0, handlelength=1.5, handletextpad=0.2, frameon=False)\n",
    "    plt.plot([0,1], [0,1], '--k')\n",
    "    plt.xscale('log')\n",
    "    plt.yscale('log')\n",
    "    plt.title(fr'$n={n}$, $v_p={vel}$', fontsize=13)\n",
    "\n",
    "    plt.xlim(1e-7, 1)\n",
    "    plt.ylim(1e-7, 1)\n",
    "    plt.yticks([])\n",
    "\n",
    "plt.sca(axs[1])\n",
    "plt.xlabel(r'$e^{-\\Delta F / T} = Z_B/Z_A$')\n",
    "plt.sca(axs[0])\n",
    "plt.ylabel(r'$\\langle e^{-W / T} \\rangle$')\n",
    "\n",
    "plt.tight_layout(pad=0.1)\n",
    "\n",
    "plt.savefig('je.pgf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94be75af",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b3109eb",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90931638",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27d9070e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def cond(x):\n",
    "    return x == 1\n",
    "\n",
    "def body(x):\n",
    "    return 1\n",
    "\n",
    "jax.lax.while_loop(cond, body, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4272aecb",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b657e0c8",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "234bd18a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_piston(state, proj, draw_walls=True):\n",
    "    plt.axis('off')\n",
    "    plt.axis('square')\n",
    "    plt.xlim(-1.1, 1.1)\n",
    "    plt.ylim(-1.1, 1.1)\n",
    "\n",
    "    r = state['balls']['r']\n",
    "    p = proj(state['balls']['x'])\n",
    "    phi = jnp.linspace(0, 2 * jnp.pi, 100)\n",
    "    \n",
    "    x = p[:, 0] + r * jnp.cos(phi[:, None])\n",
    "    y = p[:, 1] + r * jnp.sin(phi[:, None])\n",
    "    plt.plot(x, y)\n",
    "    \n",
    "    if draw_walls:\n",
    "        x = proj(state['walls']['x'])\n",
    "        j = proj(state['walls']['j'])\n",
    "        k = proj(state['walls']['k'])\n",
    "        path = [x, x + j, x + j + k, x + k, x]\n",
    "        plt.plot([x[:, 0] for x in path], [x[:, 1] for x in path], 'white')\n",
    "\n",
    "\n",
    "def xy(pos):\n",
    "    return pos[..., [0, 1]]\n",
    "\n",
    "def xz(pos):\n",
    "    return pos[..., [0, 2]]\n",
    "\n",
    "def zy(pos):\n",
    "    return pos[..., [2, 1]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f64d09e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.style.use('dark_background')\n",
    "\n",
    "fig, [axs1, axs2] = plt.subplots(2, 3, figsize=(9, 6))\n",
    "\n",
    "def view(axs, states):\n",
    "    ax1, ax2, ax3 = axs\n",
    "    \n",
    "    plt.sca(ax1)\n",
    "    plt.cla()\n",
    "    plot_piston(state, xz)\n",
    "\n",
    "    plt.sca(ax2)\n",
    "    plt.cla()\n",
    "    plot_piston(state, xy, False)\n",
    "    phi = jnp.linspace(0, 2 * jnp.pi, 100)\n",
    "    plt.plot(jnp.cos(phi), jnp.sin(phi), 'white')\n",
    "\n",
    "    plt.sca(ax3)\n",
    "    plt.cla()\n",
    "    plot_piston(state, zy)\n",
    "    \n",
    "state = init_piston_j(jax.random.PRNGKey(0), 50, 1e-1)\n",
    "view(axs1, state)\n",
    "\n",
    "vel = 15\n",
    "state['walls']['v'] = jnp.array([\n",
    "    [0.0, 0.0, vel/2],\n",
    "    [0.0, 0.0, -vel/2],\n",
    "])\n",
    "_, state, _ = forward_j(1.5 / vel, state)\n",
    "view(axs2, state)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig('compression.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4aa3292",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5145adc",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7efff768",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "plt.style.use('dark_background')\n",
    "\n",
    "fig, [ax1, ax2, ax3] = plt.subplots(1, 3, figsize=(9, 3))\n",
    "\n",
    "def view(state):\n",
    "    plt.sca(ax1)\n",
    "    plt.cla()\n",
    "    plot_piston(state, xz)\n",
    "\n",
    "    plt.sca(ax2)\n",
    "    plt.cla()\n",
    "    plot_piston(state, xy, False)\n",
    "    phi = jnp.linspace(0, 2 * jnp.pi, 100)\n",
    "    plt.plot(jnp.cos(phi), jnp.sin(phi), 'white')\n",
    "\n",
    "    plt.sca(ax3)\n",
    "    plt.cla()\n",
    "    plot_piston(state, zy)\n",
    "\n",
    "    plt.tight_layout()\n",
    "\n",
    "fps = 10\n",
    "moviewriter = FFMpegFileWriter(fps=fps)\n",
    "\n",
    "sout = 2.0\n",
    "sin = 8.0\n",
    "vel = 1.0\n",
    "\n",
    "dt = (1.5 / vel) / (sin * fps)\n",
    "state = init_piston_j(jax.random.PRNGKey(0), 50, 1e-1)\n",
    "\n",
    "def capture(state, sec):\n",
    "    for _ in tqdm(range(round(sec * fps))):\n",
    "        view(state)\n",
    "        moviewriter.grab_frame()\n",
    "        _, state, _ = forward_j(dt, state)\n",
    "        \n",
    "    return state\n",
    "\n",
    "with moviewriter.saving(fig, 'myfile.mp4', dpi=200):\n",
    "    state = capture(state, sout)\n",
    "\n",
    "#     state['walls']['v'] = jnp.array([\n",
    "#         [0.0, 0.0, vel/2],\n",
    "#         [0.0, 0.0, -vel/2],\n",
    "#     ])\n",
    "    \n",
    "#     state = capture(state, sin)\n",
    "\n",
    "#     state['walls']['v'] = jnp.array([\n",
    "#         [0.0, 0.0, 0],\n",
    "#         [0.0, 0.0, 0],\n",
    "#     ])\n",
    "\n",
    "#     state = capture(state, sout)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b3c1a13",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}