rename and cleanup thermodynamic_linear_algebra + add readme for ex…

…amples (#18) * rename and cleanup `thermodynamic_linear_algebra` notebook, add readme for examples * fix names in notebooks
normal-computing · May 8, 2024 · 397e7f7 · 397e7f7
1 parent d727e21
commit 397e7f7
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 40 deletions.
diff --git a/examples/README.md b/examples/README.md
@@ -0,0 +1,6 @@
+# Examples
+
+You will find two Jupyter notebooks in this folder:
+
+- `diffrax_comparison.ipynb` runs a simple OU process using `thermox` and [`diffrax`](https://github.com/patrick-kidger/diffrax) and compares runtimes (showing a large benefit from using `thermox` for long simulation times)
+- `thermodynamic_linear_algebra.ipynb` is a small tutorial on how to use functions from the `thermox.linalg` module.
diff --git a/examples/thermo_linear_algebra.ipynb → examples/thermodynamic_linear_algebra.ipynb b/examples/thermo_linear_algebra.ipynb → examples/thermodynamic_linear_algebra.ipynb
@@ -16,7 +16,8 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### In this notebook we show how to run three basic thermodynamic algorithms:\n",
+    "In this notebook we show how to run three basic thermodynamic algorithms, using functions from `thermox.linalg`:\n",
+    "\n",
     "1. Thermodynamic linear solver: find $x$ such that $Ax = b$,\n",
     "2. Thermodynamic matrix inverse: find $A^{-1}$,\n",
     "3. Thermodynamic matrix exponential: find $\\exp{(A)}$.\n",
@@ -36,15 +37,7 @@
    "cell_type": "code",
    "execution_count": 2,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "key = jax.random.PRNGKey(42) # random PRNG key\n",
     "dimension = 50 # problem size\n",
@@ -80,7 +73,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "||x_s - x*|| =  0.2386516\n"
+      "||x_s - x*|| =  0.14343248\n"
      ]
     }
    ],
@@ -116,7 +109,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "||A^{-1} - C(t,t')|| = 1.5911188\n"
+      "||A^{-1} - C(t,t')|| = 1.7223996\n"
      ]
     }
    ],
@@ -140,7 +133,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "||A^{-1} - C(t,t)|| = 0.46243933\n"
+      "||A^{-1} - C(t,t)|| = 0.5807177\n"
      ]
     }
    ],
@@ -153,7 +146,14 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "It went down! And gathering 10 times more samples only took about twice the time.\n",
+    "It went down! And gathering 10 times more samples only took about twice the time."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Thermodynamic matrix exponential\n",
     "\n",
     "Let us now consider matrix exponentials. Due to the way we obtain the matrix exponentials, the negative exponential $\\exp{(-A)}$ is more easily gathered. This is because the autocovariance function is equal to $\\exp{(-A)}$, when we have $A$ as the drift term of the SDE.\n",
     "\n",
@@ -166,14 +166,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "thermo_negexp = thermox.linalg.negexpm(A)\n"
+    "thermo_negexp = thermox.linalg.expnegm(A)\n"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": []
-  },
   {
    "cell_type": "code",
    "execution_count": 9,
@@ -183,7 +178,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "||exp(-A) - C(t+tau,t)||= 0.6903533\n"
+      "||exp(-A) - C(t+tau,t)||= 0.6535645\n"
      ]
     }
    ],
@@ -207,31 +202,14 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "||exp(-A) - C(t+tau,t)||= 0.21998222\n"
+      "||exp(-A) - C(t+tau,t)||= 0.21104243\n"
      ]
     }
    ],
    "source": [
-    "thermo_negexp = thermox.linalg.negexpm(A, num_samples=100000, dt=1)\n",
+    "thermo_negexp = thermox.linalg.expnegm(A, num_samples=100000, dt=1)\n",
     "print(r\"||exp(-A) - C(t+tau,t)||=\", jnp.linalg.norm(expm(-A)- thermo_negexp))"
    ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## TODO?? : add notes and code about non-positive matrices for matrix exponentials"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": []
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": []
   }
  ],
  "metadata": {