Merge pull request #44 from KuangYu/devel

Modify demo
deepmodeling · Jun 17, 2022 · 4bd556d · 4bd556d
2 parents bc12ab5 + 6c7516e
commit 4bd556d
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 9 deletions.
diff --git 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
@@ -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,