From 823bdc378f1fefd77dbbd070b4351e1e0cc9845a Mon Sep 17 00:00:00 2001
From: KuangYu <windwhisper.yu@gmail.com>
Date: Fri, 17 Jun 2022 10:43:53 +0800
Subject: [PATCH] Modify demos, add jit

---
 examples/classical/demo.ipynb                 |  9 +++++----
 .../demo.ipynb                                | 20 ++++++++++++++-----
 2 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/examples/classical/demo.ipynb b/examples/classical/demo.ipynb
index 6c739e123..0de988b5b 100644
--- a/examples/classical/demo.ipynb
+++ b/examples/classical/demo.ipynb
@@ -134,7 +134,8 @@
     "pairs = nbList.pairs\n",
     "nbfunc = potentials[-1]\n",
     "energy = nbfunc(positions, box, pairs, ff.getGenerators()[-1].params)\n",
-    "print(energy)"
+    "print(energy)\n",
+    "print(pairs)"
    ]
   },
   {
@@ -198,9 +199,9 @@
    "hash": "44fe82502fda871be637af1aa98d2b3ddaac01204dd30f1519cbec4e95000815"
   },
   "kernelspec": {
-   "display_name": "dmff",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
-   "name": "dmff"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
@@ -212,7 +213,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.12"
+   "version": "3.7.12"
   }
  },
  "nbformat": 4,
diff --git a/examples/fluctuated_leading_term_waterff/demo.ipynb b/examples/fluctuated_leading_term_waterff/demo.ipynb
index abc0b5982..8453e1da0 100644
--- a/examples/fluctuated_leading_term_waterff/demo.ipynb
+++ b/examples/fluctuated_leading_term_waterff/demo.ipynb
@@ -196,7 +196,7 @@
     "            dScales\n",
     "        )\n",
     "        return E_pme \n",
-    "    return admp_calculator\n"
+    "    return jax.jit(admp_calculator)\n"
    ]
   },
   {
@@ -231,20 +231,30 @@
     "print(force)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "c74539d4",
+   "metadata": {},
+   "source": [
+    "The running speed of the first pass is slow because JAX is trying to track the data flow and compile the code. Once the code is compiled, it is run much faster, until the shapes of the input parameters change, trigerring a recompilation.  "
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
    "id": "d7a0fd1b-a363-40a0-9017-febe0fe76f01",
    "metadata": {},
    "outputs": [],
-   "source": []
+   "source": [
+    "print(-force_fn(positions, box, pairs))"
+   ]
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "dmff",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
-   "name": "dmff"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
@@ -256,7 +266,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.12"
+   "version": "3.7.12"
   }
  },
  "nbformat": 4,