added HH learning code
This commit is contained in:
N_T committed
1 parent
245a81bfa5
commit
9bd9f531ea
3 files changed
+902
-5
No files matched your search
@@ -17,6 +17,7 @@ parts:
|
||||
- caption: Physical Losses
|
||||
chapters:
|
||||
- file: physicalloss.md
|
||||
- file: physicalloss-div.ipynb
|
||||
- file: physicalloss-code.ipynb
|
||||
- file: physicalloss-discuss.md
|
||||
- caption: Differentiable Physics
|
||||
|
||||
@@ -0,0 +1,893 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "fiQvU0_fCXwj"
|
||||
},
|
||||
"source": [
|
||||
"# Learning the Helmholtz-Hodge Decomposition\n",
|
||||
"\n",
|
||||
"In the following notebook we'll following the aforementioned paper by Tompson et al. {cite}`tompson2017` and train a neural network that can essentially perform a Helmholtz-Hodge decomposition. This is a very classic and time consuming part of many numerical solvers, and enables splitting an aribtrary vector field into a solenoidal (divergence-free) and irrotational part (the pressure gradient). Because this is traditionally veru time consuming, it's an interesting goal for a learned approach. As a stepping stone towards integrating full solvers, we'll formulate a physics-based loss via a discretized PDE-constraint.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "VF33GyuyCXwk"
|
||||
},
|
||||
"source": [
|
||||
"## Solving Navier-Stokes\n",
|
||||
"\n",
|
||||
"To motivate the topic, let's briefly revisit how the incompressible Navier-Stokes equations are often solved: a fundamental variant employs operator splitting to separately solve advection and pressure correction (giving first-order accuracy in time). Thus, given an arbitrary flow field $\\mathbf{u}$, we compute a self advection step solving $\\frac{\\partial \\mathbf{u}}{\\partial{t}} + \\mathbf{u} \\cdot \\nabla \\mathbf{u} = 0$. This step is usually efficient, with a complexity linear in the number of discretization points, and gives an advected flow field $\\tilde{\\mathbf{u}}$. \n",
|
||||
"\n",
|
||||
"The pressure-projection (or _Chorin_ projection) step using the aforementioned Helmholtz-Hodge decomposition requires solving a Poisson problem $\\nabla ^ 2 p = \\nabla \\cdot \\tilde{\\mathbf{u}}$, an elliptic PDE that is non-trivial to solve in the general case. In the Navier-Stokes setting it gives an instantaneous scalar pressure field $p$, the gradient of which happens to capture the divergent parts of the velocity field. Hence, once we subtract $\\nabla p$ we obtain a divergence free field that satisfies $\\nabla \\cdot \\mathbf{u} = 0$."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "ikG8LK18CXwl"
|
||||
},
|
||||
"source": [
|
||||
"Our goal below will be to obtain $p$ for new velocity inputs $\\mathbf{u}$, and the neat property of this problem is\n",
|
||||
"that we can perform this training without a supervised setup. I.e., we don't have to precompute a large number of $(\\mathbf{u},p)$ pairs to train this, but rather we'll rely on the discretized PDE-constraint. This could be called an _unsupervised_ training, but this labeling is misleading, as the training target in all \"unsupervised\" cases is simply computed on the fly.\n",
|
||||
"\n",
|
||||
"The more important conceptual aspect of the approach is that we have a target PDE $\\nabla \\cdot \\mathbf{u} = 0$, that we can discretize on our computational domain with a suitable finite-difference operator. Once we do this in a differentiable way, we can directly use the FD operator to train our neural network: as we aim for minimizing $\\nabla \\cdot \\tilde{\\tilde{\\mathbf{u}}}$, this effectively gives us an $L_2$ loss with a target of 0. From this loss, we can back-propagate through the discretized divergence operator to obtain a gradient for $\\mathbf{u}$, which is computed (as outlined above), as $\\tilde{\\tilde{\\mathbf{u}}} = \\tilde{\\mathbf{u}} - \\nabla p$, where $p$ is the output of our neural network $f$ with parameters $\\theta$. \n",
|
||||
"As we know that the divergence is all that is required to uniquely determine the pressure field, we can pass the divergence to $f$ to simplify the inference task.\n",
|
||||
"Thus $f( \\nabla \\cdot \\tilde{\\mathbf{u}} ; \\theta)$ will receive a gradient backpropagated from the loss through all steps to update its state $\\theta$ such that the velocity field above will end up divergence free.\n",
|
||||
"\n",
|
||||
"Putting these steps together, we aim for solving the minimization problem\n",
|
||||
"$$\n",
|
||||
"\\text{arg min}_{\\theta} \\sum_i \\Big( \\nabla \\cdot (\\tilde{\\mathbf{u}} - \\nabla f(\\nabla \\cdot \\tilde{\\mathbf{u}} ; \\theta)) \\Big)^2 ,\n",
|
||||
"$$\n",
|
||||
"where the $i$ subscript denotes an arbitrary number of different inputs, e.g., from a minibatch at training time.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "KTZ4m8rnCXwl"
|
||||
},
|
||||
"source": [
|
||||
"## Setting up the Discrete PDE\n",
|
||||
"\n",
|
||||
"As before, we'll use the `PhiFlow` framework. \n",
|
||||
"The cell below installs it via `pip`, and imports the PyTorch backend together with `tqdm` for tracking the progress. Note that this notebook can be switched to TensorFlow quite easily by changing the import below."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"metadata": {
|
||||
"id": "EK8xTOQlCXwm"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"!pip install --quiet phiflow==3.1 tqdm\n",
|
||||
"from tqdm import tqdm\n",
|
||||
"from phiml import nn\n",
|
||||
"\n",
|
||||
"from phi.torch.flow import *\n",
|
||||
"# this notebook largely works identically with tensorflow - try replacing the line above with this one\n",
|
||||
"#from phi.tf.flow import * "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "opKU6RVrCXwm"
|
||||
},
|
||||
"source": [
|
||||
"Let's define some basic constants, the grid resolution (which we'll reuse as _physical_ dimensions in PhiFlow), the number of different flow fields we'll consider for training, and the batch size."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"metadata": {
|
||||
"id": "hs8v9vEECXwn"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"RES_X = 32\n",
|
||||
"RES_Y = 32\n",
|
||||
"\n",
|
||||
"NUM_SAMPLES_TRAIN = 100\n",
|
||||
"BATCH_SIZE = 10"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "-6mIFTtQCXwn"
|
||||
},
|
||||
"source": [
|
||||
"We will use random, periodic flow fields generated via PhiFlow's frequency-based synthesis. The next cell will prec-compute all these flow fields (100 in total), and visualize an example. Note, we're simply taking random, divergent flow fields here. We don't know what the correct divergence-free counterparts are."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 3,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 463
|
||||
},
|
||||
"id": "CDCkr16bCXwn",
|
||||
"outputId": "43131d6b-83c0-438e-9790-b40120f3b427"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Initial L2 of divergence: \u001b[94m195.28241\u001b[0m, max=1.8969224691390991\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"image/png": "iVBORw0KGgoAAAANSUhEUgAABCEAAAGtCAYAAAAoOG76AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjkuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8hTgPZAAAACXBIWXMAAA9hAAAPYQGoP6dpAAEAAElEQVR4nOzdd3QU5dfA8e9ueu+9AQktdELvvQpIUaoUUVHBgh0r6k+xvXbFgopSBOkgvffeewmQkEZI72133j+WBMJmZxNaot7POTkHpj4zuzu7c+d57tUoiqIghBBCCCGEEEIIcY9pK7sBQgghhBBCCCGE+G+QIIQQQgghhBBCCCHuCwlCCCGEEEIIIYQQ4r6QIIQQQgghhBBCCCHuCwlCCCGEEEIIIYQQ4r6QIIQQQgghhBBCCCHuCwlCCCGEEEIIIYQQ4r6wrOwGCCGEEEIIIYQQ/1RpaWlkZWVVdjNKODo64urqWtnNMEmCEEIIIYQQQgghxG1IS0ujbm1PEhJ1ld2UElqtluTk5CobiJAghBBCCCGEEEIIcRuysrJISNQRdbAazk6Vn+0gI1NPSMRlsrKyJAghhBBCCCGEEEL8Gzk6aXB00lR2M9BT+W0wR4IQQgghhBBCCCHEHdApenRKZbfC0I6qrvL7iwghhBBCCCGEEOI/QXpCCCGEEEIIIYQQd0CPgp7K7wpRFdpgjgQhhBBCCCGEEEKIO6BHT1UYCFE1WqFOhmMIIYQQQgghhBDivpCeEEIIIYQQQgghxB3QKQo6pfKHQlSFNpgjPSGEEEIIIYQQQghxX0hPCCGEEEIIIYQQ4g5IYsrykyCEEEIIIYQQQghxB/Qo6KpAAOCfEISQ4RhCCCGEEEIIIYS4L6QnhBBCCCGEEEIIcQdkOEb5SRBCCCGEEEIIIYS4A1Ido/xkOIYQQgghhBBCCCHuC+kJIYQQQgghhBBC3AH99b/KVhXaYI70hBBCCCGEEEIIIcR9IT0hhBBCCCGEEEKIO6CrIiU6q0IbzJEghBBCCCGEEEIIcQd0iuGvslWFNpgjwzGEEEIIIYQQQoj/oOHDh2NnZ4dGo0Gj0eDk5MR77713T/cpQQghhBBCCCGEEOIO6KvQX0WEhYXx2muvsW7dOtasWUPDhg155513WLZs2W2chfKR4RhCCCGEEEIIIcQd0KNBh6aym4H+ehvi4+NLTXd2dsbZ2dlo+ffff7/U/3v27IlWq2XJkiUMGDDgnrRRghBCCCGEEEIIIcS/SIsWLUr9v2PHjmzZskV1nYKCAl566SUURWHQoEH3rG0ShBBCCCGEEEIIIe6AXjH8VbbiNuzbtw8/P7+S6WX1gii2aNEihgwZUvL/qVOn0r9//3vWRglCCCGEEEIIIYQQ/yJ+fn4EBgaWa9mePXuyYcMG4uPj+e6773j33Xdp0qTJPQtEaBRFqQLxGiGEEEIIIYQQ4p8lJiaGoKAg9p70xdGp8us+ZGXqaVkvgStXrpQ7CHErd3d3fHx8OH369F1unYH0hBBCCCGEEEIIIe6ArookprwbbVAUhcLCwrvQmrJJEEIIIYQQQgghhPgPat26NSNGjKBJkyZcvXqVzz77jLS0NF599dV7tk8JQgghhBBCCCGEEHdAr2jQK5XfE6KibUhOTmby5MnodDo0Gg2urq5MmzaN11577R61UIIQQgghhBBCCCHEHfmnDsc4d+7cPWqJaZWfOUMIIYQQQgghhBD/CdITQgghhBBCCCGEuAM6tOiqwDN+XWU3oBwq/ywJIYQQQgghhBDiP0F6QgghhBBCCCGEEHdAqSKJKZUq0AZzJAghhBBCCCGEEELcgX9qYsrKIMMxhBBCCCGEEEIIcV9ITwghhBBCCCGEEOIO6BQtOqXyn/HrlMpugXkShBBCCCGEEEIIIe6AHg36KjDQQE/Vj0JU/lkSQgghhBBCCCHEf4L0hBBCCCGEEEIIIe6AJKYsP+kJIYQQQgghhBBCiPtCekIIIYQQQgghhBB3oOokpqz6OSEkCCGEEEIIIYQQQtwBQ2LKyh8KURXaYE7lh2qEEEIIIYQQQgjxnyA9IYQQQgghhBBCiDugR4uuCjzj/yeU6JQghBBCCCGEEEIIcQckJ0T5Vf5ZEkIIIYQQQgghxH+C9IQQQgghhBBCCCHugB4t+irwjP+fMByj8s+SEEIIIYQQQggh/hMkCCH+NTQaDVOnTr0n2758+TIajYaZM2fek+0LIf657uW1pyqbOnUqGs39KwP2Xz3PQvyT3O/rghBViU7RVJm/qk6CEKJS9O/fH3t7ezIzM00uM3LkSKytrUlOTr6PLSu/VatWyQ9iIf5B/g3Xnaruww8/ZOnSpZXdDCHEXTJz5kw0Gk3Jn62tLf7+/vTs2ZOvv/5a9XoqxH+N7np1jKrwV9VV/RaKf6WRI0eSm5vLkiVLypyfk5PDsmXL6NWrFx4eHve5dcZCQkLIzc3lkUceKZm2atUq3n333UpslRCiIv5p152q7s033yQ3N7fUNAlCCPHv9N577zFr1iymT5/OM888A8Dzzz9PgwYNOHbsWMlyZV0XhBDiVhKEEJWif//+ODk5MXfu3DLnL1u2jOzsbEaOHHmfW1a24ui/hYVFZTdFCHGb/mnXnarO0tISW1vbym6GEOI+6N27N6NGjWLcuHFMmTKFtWvXsmHDBhITE+nfv39J4KEyrguKokjgQ1QJekVbZf6quqrfQvGvZGdnx6BBg9i4cSOJiYlG8+fOnYuTkxP9+/cnLS2N559/nqCgIGxsbAgLC+Pjjz9Gr9eb3c/hw4fp3bs3zs7OODo60rVrV/bs2WO0XFpaGpMnT6ZatWrY2NgQGBjI6NGjSUpKAoxzQowdO5bvvvsOoFQ3RUVRqFatGgMGDDDaR15eHi4uLkyYMKEip0oIcZdU5LoD3NG1p7w0Gg2TJk1iwYIFhIeHY2dnR+vWrTl+/DgAP/74I2FhYdja2tKpUycuX75cav3t27fz0EMPERwcjI2NDUFBQUyePLnMH+TF+7C1taV+/fosWbKEsWPHUq1atZJliq91n332GT/99BOhoaHY2NjQvHlz9u/fX2p7t4791mg0ZGdn8/vvv5dcE8eOHQtgtB9T2wDIz89n8uTJeHl5lbweMTExZZ6/2NhYHn30UXx8fLCxsaFevXr8+uuvpk63EOIu6tKlC2+99RZRUVHMnj0bMP5M169fn86dOxutq9frCQgIYMiQIaWmffnll9SrVw9bW1t8fHyYMGECqamppdatVq0aDzzwAGvXrqVZs2bY2dnx448/AhAVFUX//v1xcHDA29ubyZMns3btWjQaDVu2bCm1nb1799KrVy9cXFywt7enY8eO7Ny5s9Qyxcdz4cIFxo4di6urKy4uLowbN46cnByj45o9ezYtWrTA3t4eNzc3OnTowLp160ots3r1atq3b4+DgwNOTk707duXkydPluOMi6qusodg/JOGY0iJTlFpRo4cye+//85ff/3FpEmTSqanpKSwdu1ahg8fjqIodOzYkdjYWCZMmEBwcDC7du1iypQpxMfH8+WXX5rc/smTJ2nfvj3Ozs688sorWFlZ8eOPP9KpUye2bt1Ky5YtAcjKyqJ9+/acPn2aRx99lKZNm5KUlMTy5cuJiYnB09PTaNsTJkwgLi6O9evXM2vWrJLpGo2GUaNG8cknn5CSkoK7u3vJvBUrVpCRkcGoUaPuwtkTQtyO8lx37OzsyMnJue1rT0Vt376d5cuXM3HiRACmTZvGAw88wCuvvML333/P008/TWpqKp988gmPPvoomzZtKll3wYIF5OTk8NRTT+Hh4cG+ffv45ptviImJYcGCBSXLrVy5kqFDh9KgQQOmTZtGamoq48ePJyAgoMw2zZ07l8zMTCZMmIBGo+GTTz5h0KBBXLx4ESsrqzLXmTVrFo899hgtWrTgiSeeACA0NLTC5+Oxxx5j9uzZjBgxgjZt2rBp0yb69u1rtNzVq1dp1apVSSDHy8uL1atXM378eDIyMnj++ecrvG8hRMU88sgjvP7666xbt47HH3/caP7QoUOZOnUqCQkJ+Pr6lkzfsWMHcXFxDBs2rGTahAkTmDlzJuPGjePZZ5/l0qVLfPvttxw+fJidO3eWuvacPXuW4cOHM2HCBB5//HFq165NdnY2Xbp0IT4+nueeew5fX1/mzp3L5s2bjdq1adMmevfuTUREBO+88w5arZbffvuNLl26sH37dlq0aFFq+Ycffpjq1aszbdo0Dh06xIwZM/D29ubjjz8uWebdd99l6tSptGnThvfeew9ra2v27t3Lpk2b6NGjB2C4To4ZM4aePXvy8ccfk5OTw/Tp02nXrh2HDx8uM1grxL+SIkQlKSoqUvz8/JTWrVuXmv7DDz8ogLJ27Vrl/fffVxwcHJRz586VWua1115TLCwslOjo6JJpgPLOO++U/P/BBx9UrK2tlcjIyJJpcXFxipOTk9KhQ4eSaW+//bYCKIsXLzZqo16vVxRFUS5duqQAym+//VYyb+LEiUpZH6GzZ88qgDJ9+vRS0/v3769Uq1atZJtCiPuvPNcdRVHu6NpTEYBiY2OjXLp0qWTajz/+qACKr6+vkpGRUTJ9ypQpClBq2ZycHKNtTps2TdFoNEpUVFTJtAYNGiiBgYFKZmZmybQtW7YogBISElIyrfha5+HhoaSkpJRMX7ZsmQIoK1asKJn2zjvvGF0DHRwclDFjxhi1acyYMaX2Y2obR44cUQDl6aefLrXciBEjjM7z+PHjFT8/PyUpKanUssOGDVNcXFzKPDdCiIr57bffFEDZv3+/yWVcXFyUJk2aKIpi/Jku/k30zTfflFrn6aefVhwdHUs+p9u3b1cAZc6cOaWWW7NmjdH0kJAQBVDWrFlTatn/+7//UwBl6dKlJdNyc3OVOnXqKICyefNmRVEMv+1q1qyp9OzZs9RvspycHKV69epK9+7dS6YVH8+jjz5aal8DBw5UPDw8Sv5//vxLine truncated
|
||||
"text/plain": [
|
||||
"<Figure size 1200x500 with 5 Axes>"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"vel = field.StaggeredGrid(\n",
|
||||
" values=field.Noise(batch(batch=NUM_SAMPLES_TRAIN)),\n",
|
||||
" extrapolation=math.extrapolation.PERIODIC,\n",
|
||||
" bounds=geom.Box(x=RES_X, y=RES_Y), \n",
|
||||
" resolution=math.spatial(x=RES_X, y=RES_Y),\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"# plot the velocity field\n",
|
||||
"plot = vis.plot({ \"velocity\": vel.batch[0], \"vel. magnitude\": math.vec_length(vel.batch[0].at_centers().values), \"divergence\": field.divergence(vel.batch[0]), })\n",
|
||||
"\n",
|
||||
"# for initial assessments, let's print the typical divergence of an input\n",
|
||||
"print( f\"Initial L2 of divergence: {math.l2_loss(field.divergence(vel.batch[0]))}, max={field.divergence(vel.batch[0]).values.max}\" )"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "p2zXNACNCXwo"
|
||||
},
|
||||
"source": [
|
||||
"Here you can see the flowlines together with velocity magnitudes and the divergence per cell. The latter is exactly what we're aiming for removing. This visualization shows that the divergence is smaller than the actual magnitude of the velocities, with an average of around 0.4, as indicated by the L2 output right above the images.\n",
|
||||
"\n",
|
||||
"Next, we will define a Navier-Stokes simulation step. Given our reduced setup without external forces or obstacles, it's very simple: a call to an advection function in PhiFlow, followed by `fluid.make_incompressible()` to invoke the Poisson solver. We'll directly annotate this function and the following ones for JIT compilation with `@jit_compile`. This is important for good performance on GPUs, but it makes debugging much harder. So when changing the code, it's highly recommended to remove them. The code will work just as well without, just slower. Once everything's running as it should, re-activate JIT compilation for the _real_ training runs. \n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 4,
|
||||
"metadata": {
|
||||
"id": "HxhmDFPOCXwo"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"/home/thuerey/anaconda3/envs/torch24/lib/python3.12/site-packages/phiml/math/_optimize.py:631: UserWarning: Possible rank deficiency detected. Matrix might be singular which can lead to convergence problems. Please specify using Solve(rank_deficiency=...).\n",
|
||||
" warnings.warn(\"Possible rank deficiency detected. Matrix might be singular which can lead to convergence problems. Please specify using Solve(rank_deficiency=...).\")\n",
|
||||
"/home/thuerey/anaconda3/envs/torch24/lib/python3.12/site-packages/phiml/backend/torch/_torch_backend.py:800: UserWarning: Sparse CSR tensor support is in beta state. If you miss a functionality in the sparse tensor support, please submit a feature request to https://github.com/pytorch/pytorch/issues. (Triggered internally at ../aten/src/ATen/SparseCsrTensorImpl.cpp:53.)\n",
|
||||
" return torch.sparse_csr_tensor(row_pointers, column_indices, values, shape, device=values.device)\n",
|
||||
"/home/thuerey/anaconda3/envs/torch24/lib/python3.12/site-packages/phiml/math/_optimize.py:631: UserWarning: Possible rank deficiency detected. Matrix might be singular which can lead to convergence problems. Please specify using Solve(rank_deficiency=...).\n",
|
||||
" warnings.warn(\"Possible rank deficiency detected. Matrix might be singular which can lead to convergence problems. Please specify using Solve(rank_deficiency=...).\")\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"@jit_compile\n",
|
||||
"def step(v, dt = 1.0):\n",
|
||||
" v = advect.mac_cormack(v, v, dt)\n",
|
||||
" v, p = fluid.make_incompressible(v, [])\n",
|
||||
" return v, p\n",
|
||||
"\n",
|
||||
"v,p = step(vel)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "QYJ3jdLeCXwo"
|
||||
},
|
||||
"source": [
|
||||
"This cell directly runs the solver on our random velocity field from before. If you're curious, you can plot it via `vis.plot()`. However, as we've demonstrated in earlier chapters that PhiFlow's NS solver works, we'll focus on the NN training. Our goal below is to train the network without having to precompute many pressure examples with a function like `make_incompressible()` above. Instead, we'll use the differentiable divergence operator to set up our training.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "IQuy2EYpfqXp"
|
||||
},
|
||||
"source": [
|
||||
"## Neural Network Training\n",
|
||||
"\n",
|
||||
"As NN architecture for an elliptic problem, an architecture with global communication is suitable. Below, we initialize a U-Net, but feel free to try the ResNet variant, which works less well due to its lack of a wide receptive field. (Given that property, it's doing surprisingly well.) As we're dealing with a periodic domain for simplicity, the NN likewise needs to be configured for periodic processing via `periodic=True`. It's input is a single channel (the divergence), and the output a very different content, the pressure. However, for the network this is likewise simply a single, scalar channel. The `filters=24` determine the total number of parameters. Feel free to increase this to improve accuracy (and reduce computational performance of the NN inference). This is the classic accuracy vs performance trade-off that NNs share with all classic numerical methods.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 5,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "n2gn4poGCXwp",
|
||||
"outputId": "b3cbf653-3681-4ba1-d83e-efe853479fa3"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"UNet(\n",
|
||||
" (inc): DoubleConv(\n",
|
||||
" (double_conv): Sequential(\n",
|
||||
" (0): Conv2d(1, 24, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), padding_mode=circular)\n",
|
||||
" (1): Identity()\n",
|
||||
" (2): ReLU()\n",
|
||||
" (3): Conv2d(24, 24, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), padding_mode=circular)\n",
|
||||
" (4): Identity()\n",
|
||||
" (5): ReLU(inplace=True)\n",
|
||||
" )\n",
|
||||
" )\n",
|
||||
" (down1): Down(\n",
|
||||
" (maxpool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
|
||||
" (conv): ResNetBlock(\n",
|
||||
" (sample_input): Identity()\n",
|
||||
" (bn_sample): Identity()\n",
|
||||
" (bn1): Identity()\n",
|
||||
" (conv1): Conv2d(24, 24, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
|
||||
" (bn2): Identity()\n",
|
||||
" (conv2): Conv2d(24, 24, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
|
||||
" )\n",
|
||||
" )\n",
|
||||
" (up1): Up(\n",
|
||||
" (up): Upsample(scale_factor=2.0, mode='bilinear')\n",
|
||||
" (conv): DoubleConv(\n",
|
||||
" (double_conv): Sequential(\n",
|
||||
" (0): Conv2d(48, 24, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), padding_mode=circular)\n",
|
||||
" (1): Identity()\n",
|
||||
" (2): ReLU()\n",
|
||||
" (3): Conv2d(24, 24, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), padding_mode=circular)\n",
|
||||
" (4): Identity()\n",
|
||||
" (5): ReLU(inplace=True)\n",
|
||||
" )\n",
|
||||
" )\n",
|
||||
" )\n",
|
||||
" (down2): Down(\n",
|
||||
" (maxpool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
|
||||
" (conv): ResNetBlock(\n",
|
||||
" (sample_input): Identity()\n",
|
||||
" (bn_sample): Identity()\n",
|
||||
" (bn1): Identity()\n",
|
||||
" (conv1): Conv2d(24, 24, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
|
||||
" (bn2): Identity()\n",
|
||||
" (conv2): Conv2d(24, 24, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
|
||||
" )\n",
|
||||
" )\n",
|
||||
" (up2): Up(\n",
|
||||
" (up): Upsample(scale_factor=2.0, mode='bilinear')\n",
|
||||
" (conv): DoubleConv(\n",
|
||||
" (double_conv): Sequential(\n",
|
||||
" (0): Conv2d(48, 24, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), padding_mode=circular)\n",
|
||||
" (1): Identity()\n",
|
||||
" (2): ReLU()\n",
|
||||
" (3): Conv2d(24, 24, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), padding_mode=circular)\n",
|
||||
" (4): Identity()\n",
|
||||
" (5): ReLU(inplace=True)\n",
|
||||
" )\n",
|
||||
" )\n",
|
||||
" )\n",
|
||||
" (down3): Down(\n",
|
||||
" (maxpool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
|
||||
" (conv): ResNetBlock(\n",
|
||||
" (sample_input): Identity()\n",
|
||||
" (bn_sample): Identity()\n",
|
||||
" (bn1): Identity()\n",
|
||||
" (conv1): Conv2d(24, 24, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
|
||||
" (bn2): Identity()\n",
|
||||
" (conv2): Conv2d(24, 24, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
|
||||
" )\n",
|
||||
" )\n",
|
||||
" (up3): Up(\n",
|
||||
" (up): Upsample(scale_factor=2.0, mode='bilinear')\n",
|
||||
" (conv): DoubleConv(\n",
|
||||
" (double_conv): Sequential(\n",
|
||||
" (0): Conv2d(48, 24, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), padding_mode=circular)\n",
|
||||
" (1): Identity()\n",
|
||||
" (2): ReLU()\n",
|
||||
" (3): Conv2d(24, 24, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), padding_mode=circular)\n",
|
||||
" (4): Identity()\n",
|
||||
" (5): ReLU(inplace=True)\n",
|
||||
" )\n",
|
||||
" )\n",
|
||||
" )\n",
|
||||
" (outc): Conv2d(24, 1, kernel_size=(1, 1), stride=(1, 1))\n",
|
||||
")\n",
|
||||
"Total number of trainable parameters: 83521\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from phiml import nn \n",
|
||||
"\n",
|
||||
"network = nn.u_net(in_channels=1, out_channels=1, levels=4, periodic=True, batch_norm=False, in_spatial=(RES_X, RES_Y), filters=24)\n",
|
||||
"\n",
|
||||
"# uncomment this for an alternative architecture that could generalize to different resolutions\n",
|
||||
"#network = nn.res_net(in_channels=1,out_channels=1,layers=[32,32,32])\n",
|
||||
"\n",
|
||||
"# print network and parameter summary\n",
|
||||
"from phiml.backend import BACKENDS\n",
|
||||
"if any([b.name == 'torch' for b in BACKENDS]):\n",
|
||||
" print(network)\n",
|
||||
" print(\"Total number of trainable parameters: \"+ str( sum(p.numel() for p in network.parameters()) ))\n",
|
||||
"elif any([b.name == 'tf' for b in BACKENDS]):\n",
|
||||
" network.summary() \n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "q8uvCyqsCXwp"
|
||||
},
|
||||
"source": [
|
||||
"As can be seen here, the NN has around 83k trainable parameters.\n",
|
||||
"\n",
|
||||
"Next we can setup the training process: the `eval_nn()` function below takes a batch of flow fields, computes their divergence, and calls the NN with these inputs. The resulting scalar field is converted to a PhiFlow tensor and interpreted as a pressure field. Its gradient is subtracted from the input flow following Helmholtz-Hodge. Our goal is to make the resulting flow field `v` incompressible, i.e., it's divergence should be minimal.\n",
|
||||
"\n",
|
||||
"To achieve this goal, we'll define a loss function that computes the $L^2$ of the per-cell divergence `loss_div()`. For convenience, we'll also define a helper functions `loss_func` that will be used by PhiFlow. This function evaluates the NN for an input, and computes the resulting loss. PhiFlow will call this function and then make sure the gradient backpropagates from the single loss value through the discrete divergence operator to the weights of the NN. Note that in the incompressible setting we always remove the divergence present in the flow, and hence the time step does not play a role.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 6,
|
||||
"metadata": {
|
||||
"id": "46wOU55hCXwp"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"@jit_compile\n",
|
||||
"def eval_nn(v):\n",
|
||||
" nn_input = field.divergence(field=v, order=2)\n",
|
||||
" p = math.native_call(network, nn_input.values)\n",
|
||||
" p = field.CenteredGrid(\n",
|
||||
" values=p,\n",
|
||||
" extrapolation=math.extrapolation.PERIODIC, \n",
|
||||
" bounds=geom.Box(x=RES_X, y=RES_Y),\n",
|
||||
" resolution=spatial(x=RES_X, y=RES_Y),\n",
|
||||
" )\n",
|
||||
" grad_pres = field.spatial_gradient(p, boundary=math.extrapolation.PERIODIC, at=v.sampled_at) \n",
|
||||
" v = v - grad_pres\n",
|
||||
" return v, p\n",
|
||||
"\n",
|
||||
"# loss functions\n",
|
||||
"\n",
|
||||
"@jit_compile\n",
|
||||
"def loss_div(v):\n",
|
||||
" div = field.divergence(field=v,order=2)\n",
|
||||
" div_sum = 2 * math.l2_loss(div)\n",
|
||||
" return div_sum, div\n",
|
||||
"\n",
|
||||
"@jit_compile\n",
|
||||
"def loss_func(v):\n",
|
||||
" v, p = eval_nn(v)\n",
|
||||
" loss, div = loss_div(v) \n",
|
||||
" return math.mean(loss,dim=batch), v, p"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "nymq_eSDCXwp"
|
||||
},
|
||||
"source": [
|
||||
"Before training the network, let's look at what a randomly initialized network produces. Below we evaluate the untrained network on the first flow field input:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 7,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 551
|
||||
},
|
||||
"id": "wa9T3aXqCXwp",
|
||||
"outputId": "b17a0492-7746-4dbe-ab3a-b3a59912fa03"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Loss for untrained network: \u001b[94m390.21857\u001b[0m\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"image/png": "iVBORw0KGgoAAAANSUhEUgAAA0oAAAHvCAYAAABqlC+ZAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjkuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8hTgPZAAAACXBIWXMAAA9hAAAPYQGoP6dpAABnhklEQVR4nO3de3hU1dn+8XtymAQISQgkmUQCCYhykIMNJcSzkBKQUi2pFaXlYIRqE6tibaWlIGpL9bVqVZTXVlF/JdViPVSqKILAqwTQKKKIERAFCQmHkIQk5Dj794dl6oQE9kIys4Hv57r2VTNz74dnZ3Loylp7jcuyLEsAAAAAAJ+QYDcAAAAAAE7DQAkAAAAAWmCgBAAAAAAtMFACAAAAgBYYKAEAAABACwyUAAAAAKAFBkoAAAAA0AIDJQAAAABogYESAAAAALTAQAlA0D311FNyuVz64osvAv5vf/HFF3K5XHrqqacC/m8DAADnYqAEwJY//OEPeumll4LdBgAAQEAwUAJgS3sOlH7605/q0KFD6tmzZ7vUBwAAMMVACWgHXq9XdXV1wW4jaGpqaozyoaGhioyMlMvlaqeOAAAAzDBQAtpwxx13yOVy6dNPP9WPf/xjRUdHq2vXrrrpppuOGAS5XC7l5+dr0aJFGjBggCIiIrR06VJJ0q5du3TttdcqMTFRERERGjBggJ588skj/r2HH35YAwYMUMeOHdWlSxcNHTpUBQUFxn0f7Z4bl8ulO+6444hr3Lp1q6ZMmaLY2FjFxMRo6tSpqq2t9TuvpqZGTz/9tFwul1wul6ZMmeJX45NPPtE111yjLl266IILLpAkbdy4UVOmTFGvXr0UGRkpj8eja6+9Vvv37/frq7V7lFJTU/X9739fb7/9toYNG6bIyEj16tVLzzzzzBHXVVFRoZtvvlkpKSmKiIjQmWeeqXvuuUder/eI3JQpUxQTE6PY2FhNnjxZFRUVZp9gAABwWggLdgOA0/34xz9Wamqq5s2bp7Vr1+qhhx7SgQMHjvg/7CtWrNA//vEP5efnq1u3bkpNTVVZWZmGDx/uG0jFx8frtddeU25urqqqqnTzzTdLkv7yl7/oF7/4hX70ox/5BmIbN27UunXrdM011wTkGtPS0jRv3jy9//77+utf/6qEhATdc889kqT/9//+n6677joNGzZM06dPlyT17t3br8aVV16pPn366A9/+IMsy5IkLVu2TJ9//rmmTp0qj8ejTZs26fHHH9emTZu0du3aY84gbd26VT/60Y+Um5uryZMn68knn9SUKVOUnp6uAQMGSJJqa2t18cUXa9euXfrZz36mHj16aM2aNZo5c6Z2796tBx98UJJkWZYuv/xyvf3227r++uvVr18/vfjii5o8efKJ/FQCAIBThQWgVXPmzLEkWT/4wQ/8Hv/5z39uSbI+/PBD32OSrJCQEGvTpk1+2dzcXCspKcnat2+f3+MTJkywYmJirNraWsuyLOvyyy+3BgwYcEL63r59uyXJWrhw4RHPSbLmzJnj+/jwNV577bV+uR/+8IdW165d/R7r1KmTNXny5CNqHq5x9dVXH/Hc4ev7pr///e+WJGv16tW+xxYuXGhJsrZv3+57rGfPnkfk9uzZY0VERFi33nqr77G77rrL6tSpk/XZZ5/5/Tu33367FRoaau3YscOyLMt66aWXLEnWvffe68s0NTVZF154YZufLwAAcPpi6R1wDHl5eX4f33jjjZKkV1991e/xiy++WP379/d9bFmW/vnPf2rcuHGyLEv79u3zHdnZ2aqsrNT7778vSYqNjdVXX32ld999t52vpnXXX3+938cXXnih9u/fr6qqquOuIUkdOnTw/XddXZ327dun4cOHS5Lv2o+mf//+uvDCC30fx8fH6+yzz9bnn3/ue2zx4sW68MIL1aVLF7/PcVZWlpqbm7V69WpJX79eYWFhuuGGG3znhoaG+l5PAACAb2LpHXAMffr08fu4d+/eCgkJOeI9f9LS0vw+3rt3ryoqKvT444/r8ccfb7X2nj17JEm//vWv9eabb2rYsGE688wzNWrUKF1zzTU6//zzT9yFHEWPHj38Pu7SpYsk6cCBA4qOjrZVo+X1S1J5ebnmzp2rZ5991neth1VWVhr3dbi3AwcO+D7esmWLNm7cqPj4+FZrHP53v/zySyUlJSkqKsrv+bPPPvuYfQAAgNMPAyXAUFv31Xxz9kSSbyOBn/zkJ23eBzNo0CBJUr9+/VRcXKwlS5Zo6dKl+uc//6lHH31Us2fP1ty5c09If83NzW2eExoa2urj1n/uNbKj5fVLX9/7tGbNGt12220aMmSIoqKi5PV6NXr06CM2Wjjevrxer773ve/pV7/6VavZs846y+YVAAAA/BcDJeAYtmzZ4jdbsnXrVnm9XqWmph71vPj4eHXu3FnNzc3Kyso65r/TqVMnXXXVVbrqqqvU0NCg8ePH6/e//71mzpypyMhI2/0eng1quZvbl19+abtGa0y37j5w4ICWL1+uuXPnavbs2b7Ht2zZ8q36aKl3796qrq4+5ue4Z8+eWr58uaqrq/1mlYqLi09oPwAA4NTAPUrAMcyfP9/v44cffliSNGbMmKOeFxoaqpycHP3zn//Uxx9/fMTze/fu9f13y+2y3W63+vfvL8uy1NjYKOnr3d0+/fRT7du376j/bnR0tLp16+a7N+ewRx999KjnHUunTp2MttI+PBvUclbq8C50J8qPf/xjFRYW6vXXXz/iuYqKCjU1NUmSLrvsMjU1Nemxxx7zPd/c3Ox7PQEAAL6JGSXgGLZv364f/OAHGj16tAoLC/W3v/1N11xzjQYPHnzMc//4xz/qrbfeUkZGhqZNm6b+/furvLxc77//vt58802Vl5dLkkaNGiWPx6Pzzz9fiYmJ2rx5sx555BGNHTtWnTt3liStX79el156qebMmeP3Xkitue666/THP/5R1113nYYOHarVq1frs88++1afh/T0dL355pu6//77lZycrLS0NGVkZLSZj46O1kUXXaR7771XjY2NOuOMM/TGG29o+/bt36qPlm677Tb961//0ve//33f1uE1NTX66KOP9Pzzz+uLL75Qt27dNG7cOJ1//vm6/fbb9cUXX6h///564YUXbN0rBQAATj8MlIBjeO655zR79mzdfvvtCgsLU35+vv7nf/7H1rmJiYlav3697rzzTr3wwgt69NFH1bVrVw0YMMD3HkWS9LOf/UyLFi3S/fffr+rqanXv3l2/+MUvNGvWrOPqefbs2dq7d6+ef/55/eMf/9CYMWP02muvKSEh4bjqSdL999+v6dOna9asWTp06JAmT5581IGSJBUUFOjGG2/U/PnzZVmWRo0apddee03JycnH3UdLHTt21KpVq/SHP/xBixcv1jPPPKPo6GidddZZmjt3rmJiYiRJISEh+te//qWbb75Zf/vb3+RyufSDH/xAf/rTn3TuueeesH4AAMCpwWWZ3K0NnEbuuOMOzZ07V3v37lW3bt2C3Q4AAAACiHuUAAAAAKAFBkoAAAAA0AIDJQAAAABogXuUAAAAAKAFZpQAAAAAoAUGSgAAAADQwin/Pkper1clJSXq3LmzXC5XsNsBAJxiLMvSwYMHlZycrJAQ/v4I4Njq6urU0NAQ7DZ83G63IiMjg92G45zyA6WSkhKlpKQEuw0AwClu586d6t69e7DbAOBwdXV1SusZpdI9zcFuxcfj8Wj79u0Mllo45QdKnTt3lvT1L7Do6OggdwMAONVUVVUpJSXF9/sGAI6moaFBpXuatb2op6I7B38WuuqgV2npX6qhoYGBUgun/EDp8HK76OhoBkoAgHbD8m4AJqI7hzhioIS2nfIDJQAAAMBpmi2vmh3wJj3NljfYLTgWw1gAAAAAaIEZJQAAACDAvLLkVfCnlJzQg1MxUAIAAAACzCuvnLDozRldOBNL7wAAAACgBWaUAAAAgABrtiw1W8Ff9uaEHpyKgRIAAAAQYNyj5HwsvQMAAACAFphRAgAAAALMK0vNDpjNYUapbcwoAQAAAEALzCgBAAAAAcY9Ss7HQAkAAAAIMHa9cz6W3gEAAABAC8woAQAAAAHm/c8RbE7owakYKAEAAAAB1uyQXe+c0INTMVAyNGLkPKP851PsZxPecBvVboh2GeWrU+xnww6Z1U4d8YVR/vO9XW1nY6MOGdUODWm/v40sGfA3o/wvd42ynV29eqBR7biPjeK6dVaB7WytN8Ko9tx3fmCUj0uosp2t2mz/a0WSmjuavf6W237eZZCVpLBSs+/p2M32s+WDzH6xeQbsMcqXlUfbzoZ92tGodq9nvjLKb/nZGfZrv1BtVPuNdbON8gCA0wcDJQAAACDAmq2vj2BzQg9OxWYOAAAAANACM0oAAABAgLGZg/MxUAIAAAACzCuXmmV2T3h79YHWsfQOAAAAAFpgRgkAAAAIMK/19RFsTujBqRgoAQAAAAHW7JCld07owalYegcAAAAALTBQAgAAAALs8IySEw5T8+fPV2pqqiIjI5WRkaH169cfNb948WL17dtXkZGRGjhwoF599VW/5y3L0uzZs5WUlKQOHTooKytLW7Zs8T3/xRdfKDc3V2lpaerQoYN69+6tOXPmqKGhwbh3EwyUAAAAANjy3HPPacaMGZozZ47ef/99DR48WNnZ2dqzZ0+r+TVr1ujqq69Wbm6uPvjgA11xxRW64oor9PHHH/sy9957rx566CEtWLBA69atU6dOnZSdna26ujpJ0qeffiqv16v//d//1aZNm/TAAw9owYIF+s1vftOu18pACQAAAAgwr+VyzGHi/vvv17Rp0zR16lT1799fCxYsUMeOHfXkk0+2mv/zn/+s0aNH67bbblO/fv1011136Tvf+Y4eeeQRSV/PJj344IOaNWuWLr/8cg0aNEjPPPOMSkpK9NJLL0mSRo8erYULF2rUqFHq1auXfvCDH+iXv/ylXnjhhW/Line truncated
|
||||
"text/plain": [
|
||||
"<Figure size 1200x500 with 2 Axes>"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"vel_untrained, pres_untrained = eval_nn(vel.batch[0])\n",
|
||||
"\n",
|
||||
"# optionally, visualize the outputs: this doesnt look much different from before as the NN is untrained\n",
|
||||
"#plot = vis.plot( {\"vel untrained\": vel_untrained.batch[0], \"vel len. untrained\": math.vec_length(vel_untrained.batch[0].at_centers().values), \"div. untrained\": field.divergence(vel_untrained.batch[0]), })\n",
|
||||
"\n",
|
||||
"# print the loss and divergence sum of the corrected velocity from untrained NN\n",
|
||||
"loss, div_untrained = loss_div(vel_untrained)\n",
|
||||
"print(f\"Loss for untrained network: {loss}\")\n",
|
||||
"\n",
|
||||
"# also, we visualize the pressure field\n",
|
||||
"plot = vis.plot(pres_untrained.batch[0],title=\"pres. untrained\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "ZNEQUTu_CXwq"
|
||||
},
|
||||
"source": [
|
||||
"Not surprisingly, the pressure field looks quite random, and the loss in terms of divergence is large."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Zf2iFVn8gZ6K"
|
||||
},
|
||||
"source": [
|
||||
"## Training\n",
|
||||
"\n",
|
||||
"With our setup so far, training is very simple: we simply choose a few of the random flow fields, and evaulate the network, and change its weights so that the divergence loss becomes smaller. We simply loop over our precomputed flow fields without randomization below for a fixed number of epochs. Our network converges surprisingly fast. After a few epochs the loss should have decreased by more than two orders of magnitude."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 8,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "D7Ip2Da3CXwq",
|
||||
"outputId": "5b056c46-266b-4068-d207-1538a244d6e1"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
" 0%| | 0/15 [00:00<?, ?it/s]/home/thuerey/anaconda3/envs/torch24/lib/python3.12/site-packages/phiml/backend/torch/_torch_backend.py:1091: RuntimeWarning: PyTorch does not support nested tracing. The inner JIT of native(eval_nn) will be ignored.\n",
|
||||
" warnings.warn(f\"PyTorch does not support nested tracing. The inner JIT of {self.f.__name__} will be ignored.\", RuntimeWarning)\n",
|
||||
"/home/thuerey/anaconda3/envs/torch24/lib/python3.12/site-packages/phiml/backend/torch/_torch_backend.py:1091: RuntimeWarning: PyTorch does not support nested tracing. The inner JIT of native(loss_div) will be ignored.\n",
|
||||
" warnings.warn(f\"PyTorch does not support nested tracing. The inner JIT of {self.f.__name__} will be ignored.\", RuntimeWarning)\n",
|
||||
"/home/thuerey/anaconda3/envs/torch24/lib/python3.12/site-packages/phiml/backend/torch/_torch_backend.py:1091: RuntimeWarning: PyTorch does not support nested tracing. The inner JIT of native(eval_nn) will be ignored.\n",
|
||||
" warnings.warn(f\"PyTorch does not support nested tracing. The inner JIT of {self.f.__name__} will be ignored.\", RuntimeWarning)\n",
|
||||
"/home/thuerey/anaconda3/envs/torch24/lib/python3.12/site-packages/phiml/backend/torch/_torch_backend.py:1091: RuntimeWarning: PyTorch does not support nested tracing. The inner JIT of native(loss_div) will be ignored.\n",
|
||||
" warnings.warn(f\"PyTorch does not support nested tracing. The inner JIT of {self.f.__name__} will be ignored.\", RuntimeWarning)\n",
|
||||
"100%|██████████| 15/15 [00:16<00:00, 1.12s/it]\n",
|
||||
"/home/thuerey/anaconda3/envs/torch24/lib/python3.12/site-packages/phiml/backend/torch/_torch_backend.py:1091: RuntimeWarning: PyTorch does not support nested tracing. The inner JIT of native(eval_nn) will be ignored.\n",
|
||||
" warnings.warn(f\"PyTorch does not support nested tracing. The inner JIT of {self.f.__name__} will be ignored.\", RuntimeWarning)\n",
|
||||
"/home/thuerey/anaconda3/envs/torch24/lib/python3.12/site-packages/phiml/backend/torch/_torch_backend.py:1091: RuntimeWarning: PyTorch does not support nested tracing. The inner JIT of native(loss_div) will be ignored.\n",
|
||||
" warnings.warn(f\"PyTorch does not support nested tracing. The inner JIT of {self.f.__name__} will be ignored.\", RuntimeWarning)\n",
|
||||
"/home/thuerey/anaconda3/envs/torch24/lib/python3.12/site-packages/phiml/backend/torch/_torch_backend.py:1091: RuntimeWarning: PyTorch does not support nested tracing. The inner JIT of native(eval_nn) will be ignored.\n",
|
||||
" warnings.warn(f\"PyTorch does not support nested tracing. The inner JIT of {self.f.__name__} will be ignored.\", RuntimeWarning)\n",
|
||||
"/home/thuerey/anaconda3/envs/torch24/lib/python3.12/site-packages/phiml/backend/torch/_torch_backend.py:1091: RuntimeWarning: PyTorch does not support nested tracing. The inner JIT of native(loss_div) will be ignored.\n",
|
||||
" warnings.warn(f\"PyTorch does not support nested tracing. The inner JIT of {self.f.__name__} will be ignored.\", RuntimeWarning)\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Final loss=0.6785882711410522\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"optimizer = nn.adam(network, learning_rate=5e-3)\n",
|
||||
"for epoch in tqdm(range(15)):\n",
|
||||
" for b in range(int(NUM_SAMPLES_TRAIN / BATCH_SIZE)):\n",
|
||||
" vel_input = vel.batch[b*BATCH_SIZE:b*BATCH_SIZE + BATCH_SIZE]\n",
|
||||
" loss, pred_v, pred_p = nn.update_weights(\n",
|
||||
" network, optimizer, loss_func, vel_input\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
"loss, _,_ = loss_func(vel)\n",
|
||||
"print(f\"Final loss={loss.numpy()}\") "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "paM92jVECXwq"
|
||||
},
|
||||
"source": [
|
||||
"It's good to see the loss going down, but of course the big question now is: how does this network fare with new inputs. I.e., how well does it generalize to different, arbitrary flows given that it was only trained on the synthetic, randomly sampled flows."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "eNOoy9MkCXwr"
|
||||
},
|
||||
"source": [
|
||||
"## Testing with New Inputs\n",
|
||||
"\n",
|
||||
"We can check this by producing a few new inputs. Below, we'll likewise use PhiFlow's noise generation to get new fields, but to make things interesting we're increasing the scale by a factor of $2 \\times$. Hence, the network will receive divergence inputs with magnitudes it hasn't seen before. These samples are effectivly _out of the distribution_ of the training inputs."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 9,
|
||||
"metadata": {
|
||||
"id": "wrY6VqgmCXwr"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"NUM_SAMPLES_TEST = 10\n",
|
||||
"\n",
|
||||
"vel_test = field.StaggeredGrid(\n",
|
||||
" values=field.Noise(math.batch(batch=NUM_SAMPLES_TEST),scale=2.), \n",
|
||||
" extrapolation=math.extrapolation.PERIODIC,\n",
|
||||
" bounds=geom.Box(x=RES_X, y=RES_Y),\n",
|
||||
" resolution=math.spatial(x=RES_X, y=RES_Y),\n",
|
||||
")\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Now we run our trained network and the solver (for comparison) on these inputs:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 10,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 1000
|
||||
},
|
||||
"id": "Xa2lzWKpCXwr",
|
||||
"outputId": "d29687f2-b967-438f-9632-92bed69f3e90"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"/home/thuerey/anaconda3/envs/torch24/lib/python3.12/site-packages/phiml/math/_optimize.py:631: UserWarning: Possible rank deficiency detected. Matrix might be singular which can lead to convergence problems. Please specify using Solve(rank_deficiency=...).\n",
|
||||
" warnings.warn(\"Possible rank deficiency detected. Matrix might be singular which can lead to convergence problems. Please specify using Solve(rank_deficiency=...).\")\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Original, mean divergence=2059.143, div. max=4.946\n",
|
||||
"NN, mean divergence=11.510, div. max=0.417\n",
|
||||
"Solver, mean divergence=0.000, div. max=0.000\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"image/png": "iVBORw0KGgoAAAANSUhEUgAABKUAAAG+CAYAAABClLe0AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjkuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8hTgPZAAAACXBIWXMAAA9hAAAPYQGoP6dpAAEAAElEQVR4nOzdZXgUV9vA8f9K3N09uAQN7lrcCrTQlrZQ6k/9qXv71N2gUKQUL+7uwSUkxEPc3bM274ctlDTu4e35XVc/dHZ2Zjbsnjlzn3PfRyZJkoQgCIIgCIIgCIIgCIIgtCB5a1+AIAiCIAiCIAiCIAiC8O8jglKCIAiCIAiCIAiCIAhCixNBKUEQBEEQBEEQBEEQBKHFiaCUIAiCIAiCIAiCIAiC0OJEUEoQBEEQBEEQBEEQBEFocSIoJQiCIAiCIAiCIAiCILQ4EZQSBEEQBEEQBEEQBEEQWpwISgmCIAiCIAiCIAiCIAgtTgSlBEEQBEEQBEEQBEEQhBYnglKC8P+ETCbj3Xffbe3LEARBEOpBtN2CIAh3p7bSfreV6xCEhhJBKUFoBVOmTMHU1JTCwsJq95k3bx6GhoZkZ2e34JUJgiAIVWmL7ba3tzcymYxnnnmm0mvHjh1DJpOxefPm29tWrlyJTCbD2NiY5OTkSu8ZPnw4Xbt2bdZrFgRBaGltsf0WBOFvIiglCK1g3rx5lJaWsnXr1ipfLykpYfv27YwfPx47O7sWvjpBEAThn9pyu/3rr7+SkpJS5/3Ly8v55JNPmvGKBEEQ2o623H4LgiCCUoLQKqZMmYKFhQVr166t8vXt27dTXFzMvHnzWvjKBEEQhKq01Xa7S5cuaLXaegWZevToUe9AliAIwt2qrbbfd7Pi4uLWvgTh/xERlBKEVmBiYsKMGTM4fPgwGRkZlV5fu3YtFhYWTJkyhby8PJ577jk8PDwwMjLC39+fTz/9FJ1O1wpXLgiC8O9Un3YbaLG229vbmwcffLBeQabXX3+93oEsQRCEu1VbbL8PHjzI4MGDsba2xtzcnA4dOvD6669X2CcjI4NHH30UJycnjI2NCQgIYNWqVTUed/PmzchkMo4fP17ptSVLliCTyQgJCbm9LTw8nFmzZmFra4uxsTF9+vRhx44dFd53K/X7+PHjPPnkkzg6OuLu7t6ITy8IFYmglCC0knnz5qHRaNi4cWOF7Tk5Oezfv5/p06cjSRLDhg1jzZo1PPjgg3z33XcMGjSI1157jRdeeKGVrlwQBOHfqS7ttomJCSUlJS3adr/xxhtoNJo6B5l8fHzqHcgSBEG4m7Wl9js0NJRJkyZRXl7O+++/z5dffsmUKVM4ffr07X1KS0sZPnw4v//+O/PmzePzzz/HysqKBQsW8O2331Z77IkTJ2Jubl7pcwJs2LCBLl263K4dGBoaSv/+/QkLC+PVV1/lyy+/xMzMjGnTplWZ6vjkk09y48YN3n77bV599dUm+EsIwl8kQRBahUajkVxcXKQBAwZU2P7LL79IgLR//37pgw8+kMzMzKTIyMgK+7z66quSQqGQEhISbm8DpHfeeaclLl0QBOFfqS7ttiRJLdZ2e3l5SRMnTpQkSZIefvhhydjYWEpJSZEkSZKOHj0qAdKmTZtu779ixQoJkC5cuCDFxMRISqVSevbZZ2+/PmzYMKlLly4NuhZBEIS2rC21319//bUESJmZmdXu880330iAtGbNmtvbVCqVNGDAAMnc3FwqKCio9jruu+8+ydHRUdJoNLe3paamSnK5XHr//fdvbxs1apTUrVs3qays7PY2nU4nDRw4UGrXrt3tbbfuHYMHD65wTEFoKmKmlCC0EoVCwdy5cwkKCiIuLu729rVr1+Lk5MSoUaPYtGkTQ4YMwcbGhqysrNv/jR49Gq1Wy4kTJ1rvAwiCIPzL1KXdBlql7X7zzTfrNVvK19eXBx54gKVLl5Kamtrk1yMIgtCWtKX229raGtDXsqouJXDPnj04Oztz33333d5mYGDAs88+S1FRUZXpebfMmTOHjIwMjh07dnvb5s2b0el0zJkzB9DPEDty5AizZ8+msLDw9ufMzs5m3LhxREVFVVqlddGiRSgUigZ+akGonghKCUIrulVQ8VbhxaSkJE6ePMncuXNRKBRERUWxb98+HBwcKvw3evRogCrz4gVBEITmU1u7DbRK292QIFN9A1mCIAh3s7bSfs+ZM4dBgwaxcOFCnJycmDt3Lhs3bqwQoIqPj6ddu3bI5RUf1zt16nT79eqMHz8eKysrNmzYcHvbhg0b6NGjB+3btwcgOjoaSZJ46623Kn3Wd955p8rP6uPj07gPLgjVULb2BQjCv1nv3r3p2LEj69at4/XXX2fdunVIknT7pqnT6RgzZgyvvPJKle+/dWMRBEEQWkZt7Ta0Xtv9xhtv8Pvvv/Ppp58ybdq0Wvf39fVl/vz5LF26VNQHEQTh/7220n6bmJhw4sQJjh49yu7du9m3bx8bNmxg5MiRHDhwoNGzkYyMjG7Xhfrpp59IT0/n9OnTfPzxx7f3uRUAe+mllxg3blyVx/H396903YLQHERQShBa2bx583jrrbcIDg5m7dq1tGvXjr59+wLg5+dHUVHR7dEZQRAEofXV1G5D67Xdfn5+zJ8/nyVLltCvX786vefNN99kzZo1fPrpp818dYIgCK2vrbTfcrmcUaNGMWrUKL766is+/vhj3njjDY4ePcro0aPx8vIiODgYnU5XYbZUeHg4AF5eXjUef86cOaxatYrDhw8TFhaGJEm3U/dAPygB+pRA8ZwhtDaRvicIrezW6Mzbb7/N1atXK4zWzJ49m6CgIPbv31/pfXl5eWg0mmqPW1JSQnh4OFlZWU1/0YIgCP9iNbXb0Lpt95tvvolareazzz6r0/53BrLS0tIadE5BEIS7RVtov3Nycipt69GjBwDl5eUATJgwgbS0tAopeBqNhu+//x5zc3OGDRtW4zlGjx6Nra0tGzZsYMOGDQQGBlZIv3N0dGT48OEsWbKkypTvzMzMGo8vCE1JzJQShFbm4+PDwIED2b59O0CFm+PLL7/Mjh07mDRpEgsWLKB3794UFxdz/fp1Nm/eTFxcHPb29lUe9/z584wYMYJ33nmHd999tyU+iiAIwr9CTe02tG7bfSvItGrVqjq/51baX0REBF26dKn3OQVBEO4WbaH9fv/99zlx4gQTJ07Ey8uLjIwMfvrpJ9zd3Rk8eDAAjz32GEuWLGHBggVcunQJb29vNm/ezOnTp/nmm2+wsLCo8XMaGBgwY8YM1q9fT3FxMV988UWlfX788UcGDx5Mt27dWLRoEb6+vqSnpxMUFERSUhLXrl2r8RyC0FREUEoQ2oB58+Zx5swZAgMDK+Rvm5qacvz4cT7++GM2bdrE6tWrsbS0pH379rz33ntYWVm14lULgiD8e1XXbkPrt923UvK0Wm2d9vf39693IEsQBOFu1drt95QpU4iLi+O3334jKysLe3t7hg0bVuH4JiYmHDt2jFdffZVVq1ZRUFBAhw4dWLFiBQsWLKjTeebMmcOyZcuQyWTMnj270uudO3fm4sWLvPfee6xcuZLs7GwcHR3p2bMnb7/9dqM/pyDUlUySJKm1L0IQBEEQBEEQBEEQBEH4dxE1pQRBEARBEARBEARBEIQWJ4JSgiAIgiAIgiAIgiAIQosTQSlBEARBEARBEARBEAShxYmglCAIgiAIgiAIgiAIgtDiRFBKEARBEARBEARBEARBaHEiKCUIgiAIgiAIgiAIgiC0OGVrX0Bz0+l0pKSkYGFhgUwma+3LEQRBEO4gSRKFhYW4uroilzd+nES0+YIgCG2XaPMFQRD+Pera5v+/D0qlpKTg4eHR2pchCIIg1CAxMRF3d/dGH0e0+YIgCG2faPMFQRD+PWpr8//fB6UsLCwA/R/C0tKyla9GEARBuFNBQQEeHh632+rGEm2+IAhC2yXafEEQhH+Purb5/++DUrem8lpaWoqblSAIQhvVVGkXos0XBEFo+0SbLwiC8O9RW5svCp0LgiAIgiAIgiAIgiAILU4EpQRBEARBEARBEARBEIQWJ4JSgiAIgiAIgiAIgiAIQosTQSlBEARBEARBEARBEAShxYmglCAIgiAIgiAIgiAIgtDiRFBKEARBEARBEARBEARBaHEiKCUIgiAIgiAIgiAIgiC0OBGUEgRBEARBEARBEARBEFqcCEoJgiAIgiAIgiAIgiAILU4EpQRBEARBEARBEARBEIQWJ4JSgiAIgiAIgiAIgiAIQosTQSlBEARBEARBEARBEAShxYmglCAIgiAIgiAIgiAIgtDiRFBKEARBEARBEARBEARBaHEiKCUIgiAIgiAIgiAIgiC0OBGUEgRBEARBEARBEARBEFqcCEoJgiAIgiAIgiAIgiAILU4EpZqATicRm57dqteQkVOITifVuM/6I1c4ExrXMhfUSsq0xdW+Jkk1/32aSnOeR5IkjqfEtthnaSiVWkNIXFqrnV+j06G7428kSRJqrbbJz6PSaskoLmry49ZFcXE5BQWlrXJuofUkF+WTWlzQ2pdRq4z8IrQ6XZMcq623d82tXKtp7Uto87Q6HXGFOc1y7IvpyWSXljTLsesivbCowv2suel0un/9b646Cdl5xGY2/fespr/3xYwkitTlTX7OplSqVpNe2LR9oX/zd1An1e83mJCe24xX03blqUrRNFE/o6Ul/kv/zWoiglKNlF1YwgPfrWfe1+vJLqw+INKcvlp5hBnPLuPU5Zhq90nOyufrP0/y9HdbeerbLUQnZ91+7WR0HJ8eOEFRuaolLrdZ/RrzGt9HPktaadztbbuSLzLz5Gd8E7Gz2c9fpC7iq8iPWR23DJW29k7E3AN/MGbnUi5lJtXp+CvDL/HQ4Y28ErSnzd6wdTodb68+wKKvNhKekN4q13AkMYaAP77jraCDfHPxDH1W/cSG8OuNOuaG69cZsXw5/zt+HIAzSQkM/30Line truncated
|
||||
"text/plain": [
|
||||
"<Figure size 1200x500 with 3 Axes>"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"image/png": "iVBORw0KGgoAAAANSUhEUgAABCAAAAGtCAYAAADH+gXEAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjkuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8hTgPZAAAACXBIWXMAAA9hAAAPYQGoP6dpAABZmUlEQVR4nO3deZRU1bn//8+pnpseoGnopmUGRVCRxAFRYxCJSAzRiF68NwMah5t7gYjIzQrfOGHMF5O4oiZOv+s3AZPIlWA0ijdKHCFGUEGJooBImBS6mXuCnqrO7w9ixw7w7AOnijrdvF9r1dKu5wzP2WfXrmb33vt4vu/7AgAAAAAASKFYuhMAAAAAAAAdHx0QAAAAAAAg5eiAAAAAAAAAKUcHBAAAAAAASDk6IAAAAAAAQMrRAQEAAAAAAFKODggAAAAAAJBymelOAAAAAACA9qqyslK7d+9OdxqtunTpovLy8nSncVB0QAAAAAAAcAQqKyv1uVN7qnJbPN2ptLF169ZIdkLQAQEAAAAAwBHYvXu3KrfFtXF5XxUVpn+Fg5rahPqctkG7d++mAwIAAAAAgI6moNBTQaGX7jSUUPpzsNABAQAAAABACHE/obif7iz25xFl6R8jAgAAAAAAOjxGQAAAAAAAEEJCvhJK/xCIKORgoQMCAAAAAIAQEkooCpMfopHFoTEFAwAAAAAApBwjIAAAAAAACCHu+4r76Z/+EIUcLIyAAAAAAAAAKccICAAAAAAAQmARymDogAAAAAAAIISEfMUj8I//qHdAMAUDAAAAAACkHCMgAAAAAAAIgSkYwdABAQAAAABACDwFIximYAAAAAAAgJRjBAQAAAAAACEk/v5KtyjkYGEEBAAAAAAASDlGQAAAAAAAEEI8Io/hjEIOFjogAAAAAAAIIe7vf6VbFHKwMAUDAAAAAACkHCMgAAAAAAAIgUUog6EDAgAAAACAEBLyFJeX7jSUiEAOFqZgAAAAAACAlGMEBAAAAAAAIST8/a90i0IOFkZAAAAAAACAlKMDAgAAAACAEOJ/XwMiCq8jdcIJJ8jzPHXr1i2JJdMWUzAAAAAAAAgh7D/+k5nHkbjtttu0du1aeV5qr4EREAAAAAAAHKPWrVunO+64Q9dee61isdR2ETACAgAAAACAEBK+p4Sf/hEQn+awbt26Nu+XlpYecmrFsGHD1K1bNz3yyCOaPXt2SvOjAwIAAAAAgBCiNgVj3Lhxbd7v06ePNmzYcMD2F110kerr67V27dqjkR4dEAAAAAAAdCQLFizQgAEDWn8uLS09YJtnnnlGCxcu1N13363y8vKjkhcdEAAAAAAAhBBXTPEILLEY//t/BwwYoMGDB5vbPvroo5Kk6dOna/r06a3v79ixQ57nae/evcrLy0tqfnRAAAAAAABwjJk1a5a+8IUvtHnvpptuUl5enu68886kdz5IdEAAAAAAABCKH5FFKP3DyOGEE07QCSec0Oa96dOnKy8vT1OnTk1yZvvRAQEAAAAAQAhRW4QyquiAAAAAAAAAamlpSenx6YAAAAAAACCEuB9T3I/AIpR+ujOw0QEBAAAAAEAICXlKROApGAlFuwci/SUERNjtt98uz4v2PCoAHRvtEIAoaU9tUt++fXXVVVelOw0An0EHBI4pc+bMked5ra/c3FxVVFRozJgx+vnPf67a2tp0pwigA4tiGzRy5Eh5nqdx48YdENuwYYM8z9Pdd9/d+t6rr77amv/y5csP2Oeqq65SQUFBSnMGkBxRbJOA9urTRSij8IoyOiBwTLrjjjv0m9/8Rg899JCmTJkiSZo6dapOOeUUvfvuu63b3Xzzzdq3b1+60gTQQQVtg6Sj1w49++yzB+1QsNx+++2pSQbAURXFNglAx8QaEDgmjR07VqeffnrrzzNmzNDLL7+sr3zlK/rqV7+qVatWKS8vT5mZmcrM5GMCILmCtkGSjko71Lt3b9XW1mrmzJl65plnAu0zbNgwPfvss3r77bf1+c9/PqX5AUitqLVJ7VlDQ4Oys7MVi/F33mNNdBahZA0IoF0YNWqUbrnlFm3cuFG//e1vJR04z/Hkk0/W+eeff8C+iURCxx13nC6//PKjli+AjuVgbZB0dNqhwsJC3XjjjVqwYIHefvvtQPtMmTJFXbp0YRQE0EGlq01atmyZxowZo9LSUuXl5alfv3769re/3Wab+vp63XTTTerVq5dycnI0aNAg3X333fKNf3gtW7ZMnufp0UcfPSC2cOFCeZ6nZ599tvW9Tz75RN/+9rdVVlamnJwcnXTSSfrVr37VZr9Pp6Q9/vjjuvnmm3XccccpPz9fNTU1h3XN6Bj2L0IZjVeU0QEBfMY3v/lNSdKf/vSng8YnTJigxYsXq7Kyss37r732mrZs2aIrr7wy5TkC6LhcbZCUunbohhtuOKwOhaKiosPutADQvhztNmnbtm268MILtWHDBn3/+9/XL37xC33961/X0qVLW7fxfV9f/epXdc899+iiiy7Sz372Mw0aNEj/9V//pWnTph3y2Keffrr69++v3/3udwfE5s2bpy5dumjMmDGSpKqqKp111ll68cUXNXnyZN13330aOHCgrrnmGt17770H7P/DH/5Q//u//6vp06fr//7f/6vs7OzA1wwca+iAAD6jZ8+eKi4u1rp16w4anzBhghKJhJ544ok278+bN08FBQW6+OKLj0aaADooVxskpa4dKioq0tSpUw+rQ+G73/2uunTpopkzZx7ROQFE29Fuk15//XXt3r1bjz32mKZPn65rr71Wd955pz744IPWbZ555hm9/PLL+uEPf6hHHnlEkyZN0jPPPKPLL79c9913nzPXF154Qbt37259r6mpSU899ZS+9rWvKSsrS5L0gx/8QPF4XO+8845uueUWfec739HTTz+tK6+8UrfffvsBa2A0NDTo9ddf14033qjvf//7ys/PD3zN6DgSiikegVcUHgVqiXZ2QBoUFBQcctXnE044QcOGDdO8efNa34vH43riiSc0bty41vmRAHCkrDZISm079OkoiKAdCsXFxZo6daqeeeYZvfPOO0d8XgDRdTTbpM6dO0vavyhuc3PzQbf54x//qIyMDH33u99t8/5NN90k3/f13HPPHfL4EyZMUHNzs5588snW9/70pz9pz549mjBhgqT9Iyx+//vfa9y4cfJ9Xzt27Gh9jRkzRtXV1Qd00k6cOJHfAdG6BkQUXlEW7eyANKirq1NhYeEh4xMmTNBf/vIXffLJJ5L2z//btm1b6xcXAIThaoOk1LVDR9KhcMMNN6hz586sBQF0UEezTfriF7+o8ePHa+bMmSotLdUll1yi2bNnq7GxsXWbjRs3qqKi4oCcBg8e3Bo/lFNPPVUnnnhim86SefPmqbS0VKNGjZIkbd++XXv27NF///d/q1u3bm1eV199taT9U0U+q1+/fod1ncCxjA4I4DM+/vhjVVdXa+DAgYfcZsKECfJ9X/Pnz5ck/e53v1NxcbEuuuiio5UmgA4qSBskpbYd+rRDgVEQAI52m+R5np544gktWbJEkydPbl0I8rTTTlNdXd0RX8c/5/rKK69ox44damxs1DPPPKPx48e3PtkjkUhIkr7xjW/ohRdeOOjrnHPOaXNMRj9A2j8FIyqvKIt2dsBR9pvf/EaSWhchOph+/frpzDPP1Lx589TS0qInn3xSl156qXJyco5WmgA6qCBtkJTadujTDoWnn346cIfC1KlTD6vTAkD7kK426ayzztKPfvQjLVu2TI899pjef/99Pf7445KkPn36aMuWLQdMC1m9enVr3DJhwgS1tLTo97//vZ577jnV1NS0WSizW7duKiwsVDwe1+jRow/66t69+xFdFwA6IIBWny5o1K9fP3396183t50wYYKWLl2qX/3qV9qxY0ebIYbNzc1avXq1tm7dmuqUAXQgh9MGSalthz7tULjjjjsCbf/ZTosVK1Yc0TkBREs62qTdu3cf8CjNYcOGSVLrNIwvf/nLisfjuv/++9tsd88998jzPI0dO9Y8x+DBg3XKKado3rx5mjdvnnr06KHzzjuvNZ6RkaHx48fr97//vVauXHnA/tu3bzePj2NX3Pci84qyzHQnAKTDc889p9WrV6ulpUVVVVV6+eWX9cILL6hPnz565plnlJuba+7/L//yL5o+fbqmT5+ukpISjR49ujX2ySefaPDgwZo4caLmzJmT4isB0B6FbYOk1LZDxcXFuuGGGw5rRMMNN9yge+65R3/961/VqVOnwz4ngPSJSpv06KOP6sEHH9TXvvY1DRgwQLW1tXrkkUdUVFSkL3/5y5KkcePG6fzzz9cPfvADbdiwQaeeeqr+9Kc/6emnn9bUqVM1YMAAZ64TJkzQrbfeqtzcXF1zzTWKxdr+Tfauu+7SK6+8ouHDh+u6667TkCFDtGvXLr399tt68cUXtWvXLuc5cOz59CkU6RaX794ojeiAwDHp1ltvlSRlZ2erpKREp5xyiu69915dffXVzoWWpP2PpTr77LP1l7/8Rddee23rY5sAIIiwbZCU+nZo6tSpuvfee1VdXR1o+86dO2vq1KlMwwDaoai0SV/84hf15ptv6vHHH1dVVZWKi4t15pln6rHHHmtd6DEWi+mZZ57Rrbfeqnnz5mn27Nnq27evfvrTn+qmm24KdJ4JEybo5ptv1t69ew+6UGZZWZnefPNN3XHHHXryySf14IMPqmvXrjrppJP04x//+LCvC8A/eP4/j3MCAAAAAABOq1at0pAhQ/Srtz+n/MKMdKejvbVxffvz7+iDDz5ofTpMlDACAgAAAACAEJiCEUz6SwgAAAAAAHR4jIAAAAAAACCEhBSJJ1Ak0p2AAyMgAAAAAABAyjECAgAAAACAEBKKKRGBv+9HIQdLh++ASCQS2rJliwoLC+V56R8SA+DY4vu+amtrVVFRccBzxoOgDQOQbrRjANqzsG1YUHE/prif/n/8RyEHS4fLine truncated
|
||||
"text/plain": [
|
||||
"<Figure size 1200x500 with 6 Axes>"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"vel_nn, p_nn = eval_nn(vel_test)\n",
|
||||
"vel_solver, p_solver = fluid.make_incompressible(vel_test, [])\n",
|
||||
"\n",
|
||||
"loss, div = loss_div(vel_test)\n",
|
||||
"print(f\"Original, mean divergence={loss.mean:.3f}, div. max={math.max(div.values).mean:.3f}\")\n",
|
||||
"\n",
|
||||
"loss, div = loss_div(vel_nn)\n",
|
||||
"print(f\"NN, mean divergence={loss.mean:.3f}, div. max={math.max(div.values).mean:.3f}\")\n",
|
||||
"\n",
|
||||
"loss, div = loss_div(vel_solver)\n",
|
||||
"print(f\"Solver, mean divergence={loss.mean:.3f}, div. max={math.max(div.values).mean:.3f}\")\n",
|
||||
"\n",
|
||||
"plot = vis.plot({ \"vel.\": vel_test.batch[0], \"vel. NN\": vel_nn.batch[0], \"vel. solver\": vel_solver.batch[0] })\n",
|
||||
"plot = vis.plot({ \"div.\": field.divergence(vel_test.batch[0]), \"div. NN\": field.divergence(vel_nn.batch[0]), \"div. solver\": field.divergence(vel_solver.batch[0]), })"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"THe quantative evalation confirms that our network has learned to reduce the divergence. The remaining one is larger than the one from PhiFlow's solver, but nonetheless orders of magnitude smaller than the original one. The images at the bottom nicely visualize this.\n",
|
||||
"\n",
|
||||
"Next, let's check what pressure field the network has actually learned to produce."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 12,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 541
|
||||
},
|
||||
"id": "lAA8BKXrCXwr",
|
||||
"outputId": "7eedc994-e28b-4cbf-8c36-231add1377fd"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"image/png": "iVBORw0KGgoAAAANSUhEUgAABGEAAAHqCAYAAACk31OOAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjkuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8hTgPZAAAACXBIWXMAAA9hAAAPYQGoP6dpAABxnklEQVR4nO3deXyV9Zn///fJck72hJCQgIRdRVSQUkXQKggK2HG0WqvTdirWUeuAU4VOBzodFqdOOmrtatF2qNiOe8flp/NVi4pYLbigiBsICCQsSdiyJ+ck59y/P6gZj7J8Ljg5S/J6Ph7345FzcuU6n3u/8jmf+759nud5AgAAAAAAQLdKS3QDAAAAAAAAegM6YQAAAAAAAOKAThgAAAAAAIA4oBMGAAAAAAAgDuiEAQAAAAAAiAM6YQAAAAAAAOKAThgAAAAAAIA4oBMGAAAAAAAgDuiEAQAAAAAAiAM6YQDE1MyZMzVkyJBENwMAAAAAkg6dMEAMLVu2TD6fr2vKysrSCSecoNmzZ6u2tjbRzQMAAIipVK19Fi1aJJ/Pp7KyMrW2tn7u90OGDNHf/M3fRL33yTz+5Cc/+Vz8J8vhzTff7LY2A+gZ6IQBusEtt9yiP/zhD/rVr36liRMnasmSJZowYcJBT/IAAACpLlVrn7q6Oi1ZssT0N7fffnvSzxeA5EUnDNANZsyYoW9+85v6h3/4By1btkw33XSTtmzZoieffPKQf9PS0hLHFvYskUhE7e3tiW4GAAC9VqrWPqeddppuv/12tbW1OcfX1tbq7rvv7uaWAeip6IQB4uC8886TJG3ZskXSgfum5OXlafPmzbrwwguVn5+vb3zjG5IOdCj87Gc/08knn6ysrCyVlZXp+uuv1/79+6Nyvvnmm5o2bZpKSkqUnZ2toUOH6tvf/vZRta+mpkZXX321Bg4cqEAgoP79++viiy/W1q1bo+J+/etf6+STT1YgENCAAQM0a9Ys1dfXHzJvR0eHiouLdfXVV3/ud42NjcrKytL3vve9rveCwaAWLlyoESNGKBAIqKKiQt///vcVDAaj/tbn82n27Nm6//77u9rz7LPPHtW8AwCA2Ev22ucTCxYsUG1trfNomLPOOkvnnXeebrvtNueOGwD4tIxENwDoDTZv3ixJ6tu3b9d7nZ2dmjZtms4++2zdcccdysnJkSRdf/31WrZsma6++mr90z/9k7Zs2aJf/epXevvtt/Xqq68qMzNTdXV1uuCCC1RaWqp58+apqKhIW7du1WOPPXZU7bvsssv0/vvv68Ybb9SQIUNUV1en5cuXq6qqqusmu4sWLdLixYs1depU3XDDDdqwYYOWLFmiN954o6tdn5WZmamvfOUreuyxx3TPPffI7/d3/e6JJ55QMBjUlVdeKelAAfa3f/u3euWVV3TdddfppJNO0rvvvquf/vSn+uijj/TEE09E5X7xxRf1yCOPaPbs2SopKeFmwAAAJJFkr30+8aUvfamrU+WGG25Qdnb2Ef9m0aJFOuecc7RkyRLNmTPnmD4fQC/kAYiZe++915PkPf/8897u3bu96upq76GHHvL69u3rZWdne9u3b/c8z/OuuuoqT5I3b968qL//85//7Eny7r///qj3n3322aj3H3/8cU+S98Ybbxxzm/fv3+9J8m6//fZDxtTV1Xl+v9+74IILvHA43PX+r371K0+S97vf/a7rvauuusobPHhw1+vnnnvOk+Q99dRTUTkvvPBCb9iwYV2v//CHP3hpaWnen//856i4u+++25Pkvfrqq13vSfLS0tK8999/3zy/AAAgdlKx9vE8z1u4cKEnydu9e7e3cuVKT5J35513dv1+8ODB3pe//OWov5HkzZo1y/M8z5s8ebJXXl7utba2Ri2HWLUPQM/F5UhAN5g6dapKS0tVUVGhK6+8Unl5eXr88cd13HHHRcXdcMMNUa8fffRRFRYW6vzzz9eePXu6pnHjxikvL08rVqyQJBUVFUmSnn76aXV0dBxTW7Ozs+X3+/XSSy99btjvJ55//nmFQiHddNNNSkv7v8PGtddeq4KCAv3v//7vIfOfd955Kikp0cMPP9z13v79+7V8+XJdccUVXe89+uijOumkkzRy5Mioef9kOPMn8/6Jc889V6NGjTqqeQYAALGVSrXPZ51zzjmaPHmy6RKjRYsWqaamhnvDADCjEwboBnfddZeWL1+uFStW6IMPPtDHH3+sadOmRcVkZGRo4MCBUe9t3LhRDQ0N6tevn0pLS6Om5uZm1dXVSTrQAXHZZZdp8eLFKikp0cUXX6x77733c/dOcREIBPSf//mfeuaZZ1RWVqZzzjlHt912m2pqarpitm3bJkk68cQTo/7W7/dr2LBhXb8/mIyMDF122WV68sknu9r32GOPqaOjI6oTZuPGjXr//fc/N98nnHCCJHXN+yeGDh1qnlcAANA9Uqn2ORhrp8rRdNwAgMQ9YYBuccYZZ+iLX/ziYWMCgUDUqBLpwH1R+vXrp/vvv/+gf1NaWirpwI1p//jHP2r16tV66qmn9Nxzz+nb3/62fvKTn2j16tXKy8sztfemm27SRRddpCeeeELPPfec/u3f/k2VlZV68cUXNXbsWFOug7nyyit1zz336JlnntEll1yiRx55RCNHjtSYMWO6YiKRiE499VTdeeedB81RUVER9drlmm0AABAfqVb7fNY555yjSZMm6bbbbtN3vvMdp79ZuHChJk2apHvuuadrpA4AHAmdMEASGT58uJ5//nmdddZZTp0MZ555ps4880zdeuuteuCBB/SNb3xDDz30kP7hH/7hqD577ty5mjt3rjZu3KjTTjtNP/nJT/Tf//3fGjx4sCRpw4YNGjZsWNffhEIhbdmyRVOnTj1s7nPOOUf9+/fXww8/rLPPPlsvvvii/vVf//Vzn//OO+9oypQp8vl85vYDAIDUk8ja57MWLVrU1ani4txzz9WkSZP0n//5n1qwYMExfz6A3oHLkYAk8rWvfU3hcFj//u///rnfdXZ2dj0Oev/+/fI8L+r3p512miRFDcvdvHlz19MJDqW1tVXt7e1R7w0fPlz5+flduaZOnSq/369f/OIXUZ+7dOlSNTQ06Mtf/vJhPyMtLU1f/epX9dRTT+kPf/iDOjs7oy5F+mTed+zYod/+9ref+/u2tja1tLQc9jMAAEDqSUTtcyif7lT5bG10KJ9cxvSb3/zmqD4TQO/DSBggiZx77rm6/vrrVVlZqbVr1+qCCy5QZmamNm7cqEcffVQ///nP9dWvflX33Xeffv3rX+srX/mKhg8frqamJv32t79VQUGBLrzwwq58U6ZMkSRt3br1kJ/50UcfacqUKfra176mUaNGKSMjQ48//rhqa2u7Hh9dWlqq+fPna/HixZo+fbr+9m//Vhs2bNCvf/1rnX766frmN795xHm74oor9Mtf/lILFy7UqaeeqpNOOinq93//93+vRx55RN/5zne0YsUKnXXWWQqHw1q/fr0eeeQRPffcc0cc5gwAAFJLImqfw1m4cKEmT55sav+5556rlStXHtXnAeh96IQBkszdd9+tcePG6Z577tEPfvADZWRkaMiQIfrmN7+ps846S9KBE/7rr7+uhx56SLW1tSosLNQZZ5yh+++/33zD2oqKCv3d3/2dXnjhBf3hD39QRkaGRo4cqUceeUSXXXZZV9yiRYtUWlqqX/3qV7r55ptVXFys6667Tv/xH/+hzMzMI37OxIkTVVFRoerq6s+NgpEOjJZ54okn9NOf/lS///3v9fjjjysnJ0fDhg3Td7/73a4b9AIAgJ4l3rXP4UyaNMncqbJo0SJTxw2A3s3nfXZcHwAAAAAAAGKOe8IAAAAAAADEAZ0wAAAAAAAAcUAnDAAAAAAAQBzQCQMAAAAAABAHdMIAAAAAAADEAZ0wAAAAAAAAcZCR6AZ0t0gkop07dyo/P18+ny/RzQEApCjP89TU1KQBAwYoLS1+32G0t7crFArFNKff71dWVlZMc6J7UMcAAGKBOiZ59PhOmJ07d6qioiLRzQAA9BDV1dUaOHBgXD6rvb1dQwfnqaYuHNO85eXl2rJlS0oXML0FdQwAIJaoYxKvx3fC5OfnSzqwsRUUFCS4NQCAVNXY2KiKioqu80o8hEIh1dSFtW3NEBXkx+Zbq8amiAaP26pQKJSyxUtvQh0DAIgF6pjk0eM7YT4ZultQUEDxAgA4Zom4JCQv36e8/Nh8bkRc0pJKqGMAALFEHZN4Pb4TBgCAVBf2Igp7scsFAAAQL9Qx0Xg6EgAAAAAAQBwwEgYAgCQXkaeIYvMVUqzyAAAAuKCOicZIGAAAAAAAgDhgJAwAAEkuoohidQV07DIBAAAcGXVMNDphAABIcmHPU9iLzfDbWOUBAABwQR0TjcuRAAAAAAAA4oCRMAAAJDluaAcAAFIVdUw0RsIAAJDkIvIUjtHUE4oXAACQOlKljrn11ls1ceJE5eTkqKio6JBxy5Yt0+jRo5WVlaV+/fpp1qxZps9hJAwAAAAAAOjVQqGQLr/8ck2YMEFLly49aMydd96pn/zkJ7r99ts1fvx4tbS0aOvWrabPoRMGAIAkxzBeAACQqlKljlm8eLGkAyNdDmb//v364Q9/qKeeekpTpkzpen/06NGmz+FyJAAAAAAAgMNYvny5IpGIduzYoZNOOkkDBw7U1772NVVXV5vyMBIGAIAkx6MdAQBAquqOOqaxsTHq/UAgoEAgEJPPOJSPP/5YkUhE//Ef/6Gf//znKiws1A9/+EOdf/75Wrdunfx+v1MeRsIAAJDkIjGeAAAA4qU76piKigoVFhZ2TZWVlQf97Hnz5snn8x12Wr9+vdt8RCLq6OjQL37xC02bNk1nnnmmHnzwQW3cuFErVqxwXh6MhAEAAAAAACmjurpaBQUFXa8PNQpm7ty5mjlz5mFzDRs2zOkz+/fvL0kaNWpU13ulpaUqKSlRVVWVUw6JThiTkQt/6hyLine truncated
|
||||
"text/plain": [
|
||||
"<Figure size 1200x500 with 4 Axes>"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"# plot the pressure on the training data, compared with the pressure given by traditional solver.\n",
|
||||
"plot = vis.plot({\"pres. solver\": p_solver.batch[0], \"pres. NN\": p_nn.batch[0]}, same_scale=False)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "xJP_D4ruCXws"
|
||||
},
|
||||
"source": [
|
||||
"Interestingly, these fields share some similarities in terms of their structures, but are clearly not identical. Also, they typically have strongly varying magnitudes (the NN training is stochastic, so the differences can vary). Often, the NN learns to generate values that are 10 times larger, while PhiFlows pressure solvers is constrained to produce values with a zero mean.\n",
|
||||
"\n",
|
||||
"Nonetheless, as measured above, the pressure field of the network successfully removes most of the divergence, and mathematically, the Poisson problem here has a null-space in the form of a constant offset: the offset does not play a role once we take the derivatives, and hence all offsets are perfectly valid solutions. Our differentiable loss used for training does not have any constraint on absolute values, and hence the network is free to use any offset it happens to produce from its random initialization."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "bwOdfQABZTOi"
|
||||
},
|
||||
"source": [
|
||||
"## Tougher Tests: Fluid Simulations with Obstacles\n",
|
||||
"\n",
|
||||
"Our initial goal were Navier-Stokes solvers, and hence it's import to check whether our network can actually cope with inputs in the form of flow fields produced by a fluid solver. After all, it was only trained on the synthetic, randomized fields that are quite different from the motions of fluids. To make things more interesting, we'll directly include a buoyant smoke plume to drive the simulation, and an obstacle. \n",
|
||||
"\n",
|
||||
"The next cell sets up a simple domain with a smoke source and a square obstacle right above it."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 13,
|
||||
"metadata": {
|
||||
"id": "JlcMjRduZPh4"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"smoke = CenteredGrid(0, extrapolation.PERIODIC, x=RES_X, y=RES_Y, bounds=Box(x=RES_X, y=RES_Y)) \n",
|
||||
"vel_smk = StaggeredGrid(0, extrapolation.PERIODIC, x=RES_X, y=RES_Y, bounds=Box(x=RES_X, y=RES_Y)) \n",
|
||||
"\n",
|
||||
"RECTANGLE = Obstacle(Box(x=(12, 20), y=(18, 20)))\n",
|
||||
"INFLOW_LOCATION = tensor((16, 5), channel(vector='x,y'))\n",
|
||||
"INFLOW = 0.6 * CenteredGrid(Sphere(center=INFLOW_LOCATION, radius=3), extrapolation.PERIODIC, x=RES_X, y=RES_Y, bounds=Box(x=RES_X, y=RES_Y))\n",
|
||||
"CYLINDER = Obstacle(geom.infinite_cylinder(x=16, y=20, radius=3, inf_dim=None))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "4CpNw9lZhiXI"
|
||||
},
|
||||
"source": [
|
||||
"In our NN-supported solver, we now have some options regarding how to impose the zero-velocity boundary conditions for the obstacle. The safest way to handle them is to set velocities both before and after invoking the NN-based pressure correction step. The former ensure the network sees the solid obstacle, and the second one ensures the boundaries are fulfilled even if the network has made slight errors at the obstacle boundary. Hence, we'll call PhiFlow's `apply_boundary_conditions()` two times in the `step_nn()` function below.\n",
|
||||
"\n",
|
||||
"In parallel, we'll also define a `step_obs()` function that runs a classic flow solver without NN for comparison."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 14,
|
||||
"metadata": {
|
||||
"id": "-uiB8iYZCXwx"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"dt = 0.1\n",
|
||||
"\n",
|
||||
"@jit_compile\n",
|
||||
"def step_nn(v, dt, f, obstacle):\n",
|
||||
" v = v + f * dt\n",
|
||||
" v = advect.mac_cormack(v, v, dt)\n",
|
||||
" v = fluid.apply_boundary_conditions(v, obstacle)\n",
|
||||
" nn_input = field.divergence(field=v, order=2)\n",
|
||||
" p = math.native_call(network, nn_input.values)\n",
|
||||
"\n",
|
||||
" p = field.CenteredGrid( values=p,\n",
|
||||
" extrapolation=math.extrapolation.PERIODIC,\n",
|
||||
" bounds=geom.Box(x=RES_X, y=RES_Y),\n",
|
||||
" resolution=spatial(x=RES_X, y=RES_Y),\n",
|
||||
" )\n",
|
||||
" grad_pres = field.spatial_gradient(p, at=v.sampled_at) \n",
|
||||
" v = v - grad_pres\n",
|
||||
" v = fluid.apply_boundary_conditions(v, obstacle)\n",
|
||||
" return v, p\n",
|
||||
"\n",
|
||||
"@jit_compile\n",
|
||||
"def step_obs(v, dt, f, obstacle):\n",
|
||||
" v = v + f * dt\n",
|
||||
" v = advect.mac_cormack(v, v, dt)\n",
|
||||
" v, p = fluid.make_incompressible(v, obstacle)\n",
|
||||
" return v, p"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "k3FfUh61Z1d0"
|
||||
},
|
||||
"source": [
|
||||
"Now we're ready to run and compare the simulations. The code below does this for `STEPS` simulation steps, running both solvers one after the other. The advected marker densities are tracked in the `traj_smk_X` lists, and allow for an intuitive, qualitative check regarding the results. The marker density is both driving and following the flow velocity, and thus nicely highlights changes in the motion of the simulated fluid."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 15,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 626
|
||||
},
|
||||
"id": "6_5wdcNUCXwx",
|
||||
"outputId": "49a9b5d7-e9cf-40ac-f7af-77f4152968d8"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
" 0%| | 0/50 [00:00<?, ?it/s]/home/thuerey/anaconda3/envs/torch24/lib/python3.12/site-packages/phiml/math/_optimize.py:631: UserWarning: Possible rank deficiency detected. Matrix might be singular which can lead to convergence problems. Please specify using Solve(rank_deficiency=...).\n",
|
||||
" warnings.warn(\"Possible rank deficiency detected. Matrix might be singular which can lead to convergence problems. Please specify using Solve(rank_deficiency=...).\")\n",
|
||||
"/home/thuerey/anaconda3/envs/torch24/lib/python3.12/site-packages/phiml/math/_optimize.py:631: UserWarning: Possible rank deficiency detected. Matrix might be singular which can lead to convergence problems. Please specify using Solve(rank_deficiency=...).\n",
|
||||
" warnings.warn(\"Possible rank deficiency detected. Matrix might be singular which can lead to convergence problems. Please specify using Solve(rank_deficiency=...).\")\n",
|
||||
"100%|██████████| 50/50 [01:37<00:00, 1.96s/it]\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"image/png": "iVBORw0KGgoAAAANSUhEUgAABAYAAAHvCAYAAADHKd00AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjkuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8hTgPZAAAACXBIWXMAAA9hAAAPYQGoP6dpAABG+klEQVR4nO3deXxU5dn/8e8kIQmQjTWLSTCAgMpiS4GmKqAgi1ZF0KKtiv4sbsBTRB/bqEWxWpRa64ZaUUFrkaWPQF2KRTTwWEEE2VUUigqSBEGzEEggc87vD8r8fiPDODdMcs6c+bxfr/PSOXPNPdfJQG5y5brv47Nt2xYAAAAAAIhLCU4nAAAAAAAAnENhAAAAAACAOEZhAAAAAACAOEZhAAAAAACAOEZhAAAAAACAOEZhAAAAAACAOEZhAAAAAACAOEZhAAAAAACAOEZhAAAAAACAOEZhAIhzAwcO1MCBA51OAwAAAIBDKAwAMWjjxo269NJL1aFDB6Wmpuqkk07Seeedp8cff9zp1AAAcNSsWbPk8/mUmpqqr7766qjnBw4cqO7duzuQGQC4F4UBIMa89957+tGPfqT169dr7NixeuKJJ/TLX/5SCQkJevTRR51ODwAAV6ivr9cDDzzgdBoAEBOSnE4AgJn7779fmZmZ+uCDD5SVlRX03O7du51JylBDQ4Msy1JycrLTqQAAPOqMM87QjBkzVFJSory8PKfTAQBXo2MAiDHbtm3T6aefflRRQJLat28f+P+Ghgb97ne/U6dOnZSSkqKTTz5Zd9xxh+rr6485dkVFhZKSkjRlypSjntuyZYt8Pp+eeOKJwLnKykpNnDhRBQUFSklJUefOnfXggw/KsqxAzOeffy6fz6eHHnpIjzzySCCfjz766Di/AgAAfL877rhDfr//e7sGfD6fxo8fr4ULF6p79+5KSUnR6aefrsWLFzdRpgDgPAoDQIzp0KGD1qxZo02bNoWN++Uvf6nJkyfrhz/8of70pz9pwIABmjp1qi6//PJjviY7O1sDBgzQvHnzjnpu7ty5SkxM1GWXXSZJ2r9/vwYMGKCXXnpJV199tR577DGdeeaZKikp0aRJk456/cyZM/X444/r+uuv1x//+Ee1bt3a8MoBAIhcUVGRrr76as2YMUO7du0KG/vuu+/q5ptv1uWXX65p06aprq5Oo0aN0t69e5soWwBwFksJgBhz2223afjw4TrjjDPUt29fnX322Ro0aJDOOeccNWvWTJK0fv16vfDCC/rlL3+pGTNmSJJuvvlmtW/fXg899JDeeecdnXPOOSHHHz16tG644QZt2rQpaHOmuXPnasCAAcrOzpYkPfzww9q2bZvWrl2rU045RZJ0ww03KC8vT3/4wx906623qqCgIPD6nTt3auvWrWrXrl2jfF0AAPiuO++8Uy+++KIefPDBsPvwfPzxx/roo4/UqVMnSdI555yjXr166eWXX9b48eObKl0AcAwdA0CMOe+887RixQpddNFFWr9+vaZNm6ahQ4fqpJNO0t///ndJ0htvvCFJR/3m/tZbb5Ukvf7668ccf+TIkUpKStLcuXMD5zZt2qSPPvpIo0ePDpybP3++zj77bLVq1Up79uwJHIMHD5bf79fy5cuDxh01ahRFAQBAk+rYsaOuuuoqPfPMMyorKztm3ODBgwNFAUnq2bOnMjIy9O9//7sp0gQAx1EYAGJQnz599Morr+jbb7/VqlWrVFJSopqaGl166aX66KOP9MUXXyghIUGdO3cOel1OTo6ysrL0xRdfHHPstm3batCgQUHLCebOnaukpCSNHDkycO6zzz7T4sWL1a5du6Bj8ODBko7eCLGoqCgalw4AgJG77rpLDQ0NYfcaKCwsPOpcq1at9O233zZmagDgGiwlAGJYcnKy+vTpoz59+qhLly669tprNX/+/MDzPp/vuMa9/PLLde2112rdunU644wzNG/ePA0aNEht27YNxFiWpfPOO0+33357yDG6dOkS9Lh58+bHlQsAACeiY8eOuvLKK/XMM8/oN7/5TciYxMTEkOdt227M1ADANSgMAB7xox/9SJJUVlamDh06yLIsffbZZzr11FMDMRUVFaqsrFSHDh3CjjVixAjdcMMNgeUEn376qUpKSoJiOnXqpH379gU6BAAAcKu77rpLL730kh588EGnUwEAV2IpARBj3nnnnZC/wTiyr0DXrl11/vnnS5IeeeSRoJiHH35YknTBBReEfY+srCwNHTpU8+bN05w5c5ScnKwRI0YExfzsZz/TihUr9Oabbx71+srKSjU0NER6SQAANKpOnTrpyiuv1J///GeVl5c7nQ4AuA4dA0CMmTBhgvbv369LLrlE3bp108GDB/Xee+9p7ty5Ovnkk3XttdcqKytLY8aM0TPPPKPKykoNGDBAq1at0gsvvKARI0Yc844E/7/Ro0fryiuv1JNPPqmhQ4cqKysr6Pn//u//1t///nf99Kc/1TXXXKPevXurtrZWGzdu1N/+9jd9/vnnQUsPAABw0p133qm//OUv2rJli04//XSn0wEAV6EwAMSYhx56SPPnz9cbb7yhZ555RgcPHlRhYaFuvvlm3XXXXYEf4J999ll17NhRs2bN0oIFC5STk6OSkhLdfffdEb3PRRddpObNm6umpibobgRHtGjRQsuWLdPvf/97zZ8/Xy+++KIyMjLUpUsXTZkyRZmZmdG8bAAATkjnzp115ZVX6oUXXnA6FQBwHZ/NrioAAAAAAMQt9hgAAAAAACCOURgAAAAAACCOURgAAAAAACCOURgAAAAAACCOURgAAAAAACCOURgAAAAAACCOJTmdQGOzLEu7du1Senq6fD6f0+kAADzGtm3V1NQoLy9PCQnO1NuZ6wAAjckNc92JqKys1L59+5xOIyAtLU1ZWVlOpxHE84WBXbt2qaCgwOk0AAAet2PHDuXn5zvy3sx1AICm4ORcd7wqKyt1ate2Kt/tdzqVgISEBO3du9dVxQHPFwbS09MlHf5DnJGR4XA2AACvqa6uVkFBQWC+cQJzHQCgMblhrjte+/btU/luv7av6aCMdOe7HaprLBX1/kL79u2jMNCUjrRUZmRk8I8lAECjcbKFn7kOANAUYnm5WkZ6gisKA27l+cIAAAAAACC++W1LftvpLA7n4UaUTAAAAAAAiGN0DAAAAAAAPM2SLUvOtwy4IYdQKAwAAAAAADzNkiU3NPG7I4ujsZQAAAAAAIA4RscAAAAAAMDT/LYtv+18G78bcgiFwgAAAAAAwNPYYyA8lhIAAAAAABDH6BgAAAAAAHiaJVt+F/y2no4BAAAAAADgOnQMAAAAAAA8jT0GwqMwAAAAAADwNO5KEB5LCQAAAAAAiGN0DAAAAAAAPM36z+E0N+QQCoUBAAAAAICn+V1yVwI35BAKhQFD5yVc5nQKAIAoWWLNdzoF12GeAwDvYJ5DpCgMAAAAAAA8zW8fPpzmhhxCYfNBAAAAAADiGB0DAAAAAABPY/PB8CgMAAAAAAA8zZJPfvmcTkOWC3IIhaUEAAAAAADEMToGAAAAAACeZtmHD6e5IYdQKAwAAAAAADzN75KlBG7IIRSWEgAAAAAAEMfoGAAAAAAAeBodA+HRMQAAAAAAQByjYwAAAAAA4GmW7ZNlO//bejfkEAqFAQAAAACAp7GUIDyWEgAAAAAAEMcoDAAAAAAAPM2vBNccph5//HFlZ2crMTFRPp9PJSUlgef279+vfv36KTU1VT6fT4mJierUqZM+/PBDo/egMAAAAAAA8DT7P3sMOH3Yx7HHwLfffqtTTjlFt91221HP7dmzR1u3btVNN92kN954Q3/+859VUVGhs88+2+g92GMAAAAAAACXmjx5siZPnixJmjZtWtBzhYWF2rt3b9C5Zs2a6ZprrtGKFStUXFwc0XtQGAAAAAAAeJrbNh8sKysLOp+RkaGMjIyovEdFRYUkqaCgIOLXUBgAAAAAAKAJ9e3bN+jxgAEDVFpaesLjVlZW6p577lGHDh2Un58f8esoDAAAAAAAPM1vJ8hvO7/Fnt8+/N9Vq1YpNzc3cD4a3QL79+9Xt27dJEnvvvuu0Wsd/co89dRT6tmzZ6Btori4WP/4xz8Cz9fV1WncuHFq06aN0tLSNGrUqEBbBAAAAAAAkbDkk6UEFxyHlxLk5uYqPz8/cJxoYWD//v3q3LmzqqqqtH79eqNuAcnhwkB+fr4eeOABrVmzRqtXr9a5556riy++WJs3b5Yk3XLLLXr11Vc1f/58LVu2TLt27dLIkSOdTBkAAAAAANc4UhT45ptv9OGHH+qUU04xHsPRpQQXXnhh0OP7779fTz31lFauXKn8/Hw999xzmj17ts4991xJ0syZM3Xqqadq5cqV+vGPf+xEygAAAACAGOO2zQdNlJeXa9myZYHHH3/8sebOnavCwkL16tVLnTp10tdff63Zs2fr0KFD2rBhgySpY8eOSktLi+g9XLPHgN/v1/z581VbW6vi4mKtWbNGhw4d0uDBgwMx3bp1U2FhoVasWHHMwkB9fb3q6+sDj6urqxs9dwAAmhJzHQAAZtyzx4Bt/Jo5c+bolltuCTxetGiRFi1apE6dOmnWrFkqLy+XJI0ePTrodX/60580ceLEiN7D8cLAxo0bVVxcrLq6OqWlpWnBggU67bTTtG7dOiUnJysrKysoPjs7O3DhoUydOlVTpkxp5KwBAHAOcx0AAPFj4sSJYX/At4+j2PBdjpdMunbtqnXr1un999/XTTfdpDFjxuijjz467vFKSkpUVVUVOHbs2BHFbAEAcB5zHQAAZg5vPuiOw40c7xhITk5W586dJUm9e/fWBx98oEcffVSjR4/WwYMHVVlZGdQ1UFFRoZycnGOOl5KSopSUlMZOGwAAxzDXAQCAaHK8Y+C7LMtSfX29evfurWbNmmnp0qWB57Zs2aIvv/xSxcXFDmYIAAAAAIgllhLkd8Fhue9HcEkOdwyUlJRLine truncated
|
||||
"text/plain": [
|
||||
"<Figure size 1200x500 with 4 Axes>"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"STEPS = 50\n",
|
||||
"smk_solv = smoke\n",
|
||||
"smk_nn = smoke\n",
|
||||
"vel_smk_solv = vel_smk\n",
|
||||
"vel_smk_nn = vel_smk\n",
|
||||
"traj_smk_solv = [smk_solv]\n",
|
||||
"traj_smk_nn = [smk_nn]\n",
|
||||
"\n",
|
||||
"# flexible backends require a small workaround for PyTorch's no_grad context here; by default (eg for TF) use a null-context\n",
|
||||
"from contextlib import nullcontext \n",
|
||||
"context = nullcontext() \n",
|
||||
"if any([b.name == 'torch' for b in BACKENDS]):\n",
|
||||
" context = torch.no_grad()\n",
|
||||
"\n",
|
||||
"with context:\n",
|
||||
" for i in tqdm(range(STEPS)):\n",
|
||||
" smk_solv = advect.mac_cormack(smk_solv, vel_smk_solv, dt=dt) + INFLOW\n",
|
||||
" smk_nn = advect.mac_cormack(smk_nn, vel_smk_nn, dt=dt) + INFLOW\n",
|
||||
" \n",
|
||||
" buoyancy_force_solv = smk_solv * (0, 1.0) @ vel_smk_solv\n",
|
||||
" buoyancy_force_nn = smk_nn * (0, 1.0) @ vel_smk_nn\n",
|
||||
" \n",
|
||||
" vel_smk_solv, _ = step_obs(vel_smk_solv, dt, buoyancy_force_solv, [RECTANGLE])\n",
|
||||
" vel_smk_nn, _ = step_nn(vel_smk_nn, dt, buoyancy_force_nn, [RECTANGLE])\n",
|
||||
" \n",
|
||||
" traj_smk_solv.append(smk_solv)\n",
|
||||
" traj_smk_nn.append(smk_nn)\n",
|
||||
"\n",
|
||||
"traj_smk_solv = field.stack(traj_smk_solv, batch('time'))\n",
|
||||
"traj_smk_nn = field.stack(traj_smk_nn, batch('time'))\n",
|
||||
"\n",
|
||||
"plot = vis.plot({\"Solver\": traj_smk_solv.time[STEPS-1], \"NN\": traj_smk_nn.time[STEPS-1]}) # show last frame"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "dnd8eqT9iRXO"
|
||||
},
|
||||
"source": [
|
||||
"The plot above shows the last frame of the simulations, and they're remarkably simiar. This results shows that the NN has successfully learned to reduce the divergence of arbitrary flows, despite having only seen the synthetic noisy inputs at training time.\n",
|
||||
"\n",
|
||||
"Moreover, it highlights the success of the differentiable discrete operator used for training. The training converges very quickly, and the trained network generalizaes in a _zero-shot_ fashion to completely new inputs. The network hasn't even seen any obstcles at training time, and can still infer (mostly correct) pressure fields to handle them. This aspect is nonetheless an interesting point for improvements of this implementation. The NN could receive the obstacle geometry as an additional input, and could be trained to pay special attention to the boundary region via increased loss values.\n",
|
||||
"\n",
|
||||
"Lastly, if you're interested in watching the full evolution of the simulated trajectories, you can comment out and run the following cell. It generates a movie that can be watched in the browser via PhiFlow's `plot(...,animate=X)` function."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"#vis.plot({\"solver\": traj_smk_solv, \"nn\": traj_smk_nn}, animate='time') # uncomment to view animation"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "vrfGtqwGkoML"
|
||||
},
|
||||
"source": [
|
||||
"# Next Steps\n",
|
||||
"\n",
|
||||
"- Improve the architecture of the neural network to boost its performance; obvious aspects to improve on are depth and features per layer, but other architectures are likewise worth a try.\n",
|
||||
"- Especially fully convolutional architectures like the `res_net` above are worth a try: evaluate how well a trained NN with these architectures can handle new and larger input grids.\n",
|
||||
"- Add support for obstacles with an additional input and an updated training data set.\n",
|
||||
"\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"gpuType": "A100",
|
||||
"machine_shape": "hm",
|
||||
"provenance": []
|
||||
},
|
||||
"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.12.4"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 4
|
||||
}
|
||||
+8
-5
@@ -109,7 +109,6 @@ in {doc}`diffphys` and after, we'll focus on direct NN representations (variant
|
||||
The second variant of employing physical residuals as soft constraints
|
||||
instead uses fully connected NNs to represent $\mathbf{u}$. This _physics-informed_ approach was popularized by Raissi et al. {cite}`raissi2019pinn`, and has some interesting pros and cons that we'll outline in the following. We will target this physics-informed version (variant 2) in the following code examples and discussions.
|
||||
|
||||
|
||||
The central idea here is that the aforementioned general function $f$ that we're after in our learning problems
|
||||
can also be used to obtain a representation of a physical field, e.g., a field $\mathbf{u}$ that satisfies $R=0$. This means $\mathbf{u}(\mathbf{x})$ will
|
||||
be turned into $\mathbf{u}(\mathbf{x}, \theta)$ where we choose the NN parameters $\theta$ such that a desired $\mathbf{u}$ is
|
||||
@@ -144,9 +143,13 @@ For higher order derivatives, such as $\frac{\partial^2 u}{\partial x^2}$, we ca
|
||||
|
||||
## Summary so far
|
||||
|
||||
The approach above gives us a method to include physical equations into DL learning as a soft constraint: the residual loss.
|
||||
The v2 approach above gives us a method to include physical equations into DL learning as a soft constraint: the residual loss.
|
||||
Typically, this setup is suitable for _inverse problems_, where we have certain measurements or observations
|
||||
for which we want to find a PDE solution. Because of the high cost of the reconstruction (to be
|
||||
demonstrated in the following), the solution manifold shouldn't be overly complex. E.g., it is typically not possible
|
||||
to capture a wide range of solutions, such as with the previous supervised airfoil example, by only using a physical residual loss.
|
||||
for which we want to find a PDE solution. Because of the ill-posedness of the optimization and learning problem,
|
||||
and the high cost of the reconstruction (to be
|
||||
demonstrated in the following), the solution manifold shouldn't be overly complex for these PINN approaches.
|
||||
E.g., it is typically very involved to capture a wide range of solutions, such as with the previous supervised airfoil example.
|
||||
|
||||
Next, we'll demonstrate these concepts with code: first, we'll show how learning the Helmholtz decomposition works out in
|
||||
practice with a v1-approach. Afterwards, we'll illustrate the PINN-approaches with a practical example.
|
||||
|
||||
Reference in new issue
Block a user