From 43b816b6785775ed25accea9039733fdedcb37b2 Mon Sep 17 00:00:00 2001 From: Felix Koehler Date: Mon, 8 Apr 2024 16:22:52 +0200 Subject: [PATCH] Example showase for 2d --- examples/solver_showcase_2d.ipynb | 46468 ++++++++++++++++++++++++++++ 1 file changed, 46468 insertions(+) create mode 100644 examples/solver_showcase_2d.ipynb diff --git a/examples/solver_showcase_2d.ipynb b/examples/solver_showcase_2d.ipynb new file mode 100644 index 0000000..47f6533 --- /dev/null +++ b/examples/solver_showcase_2d.ipynb @@ -0,0 +1,46468 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Presenting all the built-in solvers working in 2d" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "import jax\n", + "import jax.numpy as jnp\n", + "import matplotlib.pyplot as plt" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [], + "source": [ + "import exponax as ex\n", + "from IPython.display import HTML" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Linear PDEs" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Advection" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "$$\n", + "\\frac{\\partial u}{\\partial t} + \\vec{c} \\cdot \\nabla u = 0\n", + "$$" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "2024-04-08 15:33:37.695864: W external/xla/xla/service/gpu/nvptx_compiler.cc:679] The NVIDIA driver's CUDA version is 12.2 which is older than the ptxas CUDA version (12.3.52). Because the driver is older than the ptxas version, XLA is disabling parallel compilation, which may slow down compilation. You should update your NVIDIA driver or use the NVIDIA-provided CUDA forward compatibility packages.\n" + ] + }, + { + "data": { + "text/html": [ + "" + ], + "text/plain": [ + "" + ] + }, + "execution_count": 3, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "DOMAIN_EXTENT = 1.0\n", + "NUM_POINTS = 100\n", + "DT = 0.01\n", + "# Can also supply a scalar to use the same value for both dimensions\n", + "VELOCITY = jnp.array([0.4, 1.0])\n", + "\n", + "advection_stepper = ex.stepper.Advection(\n", + " 2, DOMAIN_EXTENT, NUM_POINTS, DT, velocity=VELOCITY\n", + ")\n", + "\n", + "grid = ex.make_grid(2, DOMAIN_EXTENT, NUM_POINTS)\n", + "u_0 = jnp.sin(2 * jnp.pi * grid[0:1] / DOMAIN_EXTENT) * jnp.cos(\n", + " 2 * 2 * jnp.pi * grid[1:2] / DOMAIN_EXTENT\n", + ")\n", + "\n", + "advection_trj = ex.rollout(advection_stepper, 40, include_init=True)(u_0)\n", + "\n", + "advection_ani = ex.viz.animate_state_2d(advection_trj)\n", + "\n", + "HTML(advection_ani.to_html5_video())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Diffusion" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "$$\n", + "\\frac{\\partial u}{\\partial t} = \\nu \\Delta u\n", + "$$" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "" + ], + "text/plain": [ + "" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "DOMAIN_EXTENT = 1.0\n", + "NUM_POINTS = 100\n", + "DT = 0.01\n", + "# See next section for anisotropic diffusion\n", + "NU = 0.01\n", + "\n", + "anisotropic_diffusion_stepper = ex.stepper.Diffusion(\n", + " 2, DOMAIN_EXTENT, NUM_POINTS, DT, diffusivity=NU\n", + ")\n", + "\n", + "grid = ex.make_grid(2, DOMAIN_EXTENT, NUM_POINTS)\n", + "u_0 = jnp.sin(2 * jnp.pi * grid[0:1] / DOMAIN_EXTENT) * jnp.cos(\n", + " 2 * 2 * jnp.pi * grid[1:2] / DOMAIN_EXTENT\n", + ")\n", + "\n", + "diffusion_trj = ex.rollout(anisotropic_diffusion_stepper, 40, include_init=True)(u_0)\n", + "\n", + "diffusion_ani = ex.viz.animate_state_2d(diffusion_trj)\n", + "\n", + "HTML(diffusion_ani.to_html5_video())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Anisotropic Diffusion" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "$$\n", + "\\frac{\\partial u}{\\partial t} = \\nabla \\cdot \\left( A \\nabla u \\right)\n", + "$$" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "" + ], + "text/plain": [ + "" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "DOMAIN_EXTENT = 1.0\n", + "NUM_POINTS = 100\n", + "DT = 0.01\n", + "# Can also supply a 2d vector for diagonal diffusivity. For full anisotropy, the\n", + "# matrix must be positive definite.\n", + "NU = jnp.array([[0.005, 0.003], [0.003, 0.04]])\n", + "\n", + "anisotropic_diffusion_stepper = ex.stepper.Diffusion(\n", + " 2, DOMAIN_EXTENT, NUM_POINTS, DT, diffusivity=NU\n", + ")\n", + "\n", + "grid = ex.make_grid(2, DOMAIN_EXTENT, NUM_POINTS)\n", + "u_0 = jnp.sin(2 * jnp.pi * grid[0:1] / DOMAIN_EXTENT) * jnp.cos(\n", + " 2 * 2 * jnp.pi * grid[1:2] / DOMAIN_EXTENT\n", + ")\n", + "\n", + "anisotropic_diffusion_trj = ex.rollout(\n", + " anisotropic_diffusion_stepper, 40, include_init=True\n", + ")(u_0)\n", + "\n", + "anisotropic_diffusion_ani = ex.viz.animate_state_2d(anisotropic_diffusion_trj)\n", + "\n", + "HTML(anisotropic_diffusion_ani.to_html5_video())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Advection-Diffusion" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "$$\n", + "\\frac{\\partial u}{\\partial t} + \\vec{c} \\cdot \\nabla u = \\nu \\Delta u\n", + "$$" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "" + ], + "text/plain": [ + "" + ] + }, + "execution_count": 6, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "DOMAIN_EXTENT = 1.0\n", + "NUM_POINTS = 100\n", + "DT = 0.01\n", + "# Can supply up to a vector for the advection speed, and up to a matrix for\n", + "# anisotropic diffusion\n", + "velocity = 1.0\n", + "diffusivity = 0.03\n", + "\n", + "advection_diffusion_stepper = ex.stepper.AdvectionDiffusion(\n", + " 2, DOMAIN_EXTENT, NUM_POINTS, DT, velocity=velocity, diffusivity=diffusivity\n", + ")\n", + "\n", + "grid = ex.make_grid(2, DOMAIN_EXTENT, NUM_POINTS)\n", + "u_0 = jnp.sin(2 * jnp.pi * grid[0:1] / DOMAIN_EXTENT) * jnp.cos(\n", + " 2 * 2 * jnp.pi * grid[1:2] / DOMAIN_EXTENT\n", + ")\n", + "\n", + "advection_diffusion_trj = ex.rollout(\n", + " advection_diffusion_stepper, 40, include_init=True\n", + ")(u_0)\n", + "\n", + "advection_diffusion_ani = ex.viz.animate_state_2d(advection_diffusion_trj)\n", + "\n", + "HTML(advection_diffusion_ani.to_html5_video())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Dispersion" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "$$\n", + "\\frac{\\partial u}{\\partial t} + \\vec{\\xi} \\cdot (\\nabla \\odot \\nabla \\odot \\nabla) u = 0\n", + "$$" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "" + ], + "text/plain": [ + "" + ] + }, + "execution_count": 7, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "DOMAIN_EXTENT = 1.0\n", + "NUM_POINTS = 100\n", + "DT = 0.01\n", + "# Can also supply a vector for different dispersivity per dimension\n", + "DISPERSIVITY = 0.01\n", + "\n", + "dispersion_stepper = ex.stepper.Dispersion(\n", + " 2, DOMAIN_EXTENT, NUM_POINTS, DT, dispersivity=DISPERSIVITY\n", + ")\n", + "\n", + "grid = ex.make_grid(2, DOMAIN_EXTENT, NUM_POINTS)\n", + "u_0 = jnp.sin(2 * jnp.pi * grid[0:1] / DOMAIN_EXTENT) * jnp.cos(\n", + " 2 * 2 * jnp.pi * grid[1:2] / DOMAIN_EXTENT\n", + ") + jnp.sin(3 * 2 * jnp.pi * grid[0:1] / DOMAIN_EXTENT)\n", + "\n", + "dispersion_trj = ex.rollout(dispersion_stepper, 40, include_init=True)(u_0)\n", + "\n", + "dispersion_ani = ex.viz.animate_state_2d(dispersion_trj)\n", + "\n", + "HTML(dispersion_ani.to_html5_video())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Hyper-Diffusion" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "$$\n", + "\\frac{\\partial u}{\\partial t} = \\nu \\Delta^2 u\n", + "$$" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "" + ], + "text/plain": [ + "" + ] + }, + "execution_count": 8, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "DOMAIN_EXTENT = 1.0\n", + "NUM_POINTS = 100\n", + "DT = 0.01\n", + "HYPER_DIFFUSIVITY = 0.0001\n", + "\n", + "hyper_diffusion_stepper = ex.stepper.HyperDiffusion(\n", + " 2, DOMAIN_EXTENT, NUM_POINTS, DT, hyper_diffusivity=HYPER_DIFFUSIVITY\n", + ")\n", + "\n", + "grid = ex.make_grid(2, DOMAIN_EXTENT, NUM_POINTS)\n", + "u_0 = jnp.sin(2 * jnp.pi * grid[0:1] / DOMAIN_EXTENT) * jnp.cos(\n", + " 2 * 2 * jnp.pi * grid[1:2] / DOMAIN_EXTENT\n", + ")\n", + "\n", + "hyper_diffusion_trj = ex.rollout(hyper_diffusion_stepper, 40, include_init=True)(u_0)\n", + "\n", + "hyper_diffusion_ani = ex.viz.animate_state_2d(hyper_diffusion_trj)\n", + "\n", + "HTML(hyper_diffusion_ani.to_html5_video())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Nonlinear PDEs" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Burgers" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "$$\n", + "\\frac{\\partial u}{\\partial t} + \\frac{1}{2} \\nabla \\cdot \\left( u \\otimes u \\right) = \\nu \\Delta u\n", + "$$" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "" + ], + "text/plain": [ + "" + ] + }, + "execution_count": 10, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "DOMAIN_EXTENT = 1.0\n", + "NUM_POINTS = 100\n", + "DT = 0.01\n", + "NU = 0.01\n", + "\n", + "burgers_stepper = ex.stepper.Burgers(2, DOMAIN_EXTENT, NUM_POINTS, DT, diffusivity=NU)\n", + "\n", + "grid = ex.make_grid(2, DOMAIN_EXTENT, NUM_POINTS)\n", + "\n", + "# Burgers has two channels!\n", + "u_0 = jnp.concatenate(\n", + " [\n", + " jnp.sin(2 * jnp.pi * grid[0:1] / DOMAIN_EXTENT)\n", + " * jnp.cos(2 * 2 * jnp.pi * grid[1:2] / DOMAIN_EXTENT),\n", + " jnp.cos(2 * jnp.pi * grid[0:1] / DOMAIN_EXTENT)\n", + " * jnp.sin(2 * 2 * jnp.pi * grid[1:2] / DOMAIN_EXTENT),\n", + " ]\n", + ")\n", + "\n", + "burgers_trj = ex.rollout(burgers_stepper, 40, include_init=True)(u_0)\n", + "\n", + "burgers_ani = ex.viz.animate_state_2d_facet(\n", + " burgers_trj,\n", + " grid=(1, 2),\n", + " figsize=(7, 3),\n", + ")\n", + "\n", + "HTML(burgers_ani.to_html5_video())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Single-Channel Burgers" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This is a hack to not have the channel dimension grow together with the spatial\n", + "dimension.\n", + "\n", + "$$ \\frac{\\partial u}{\\partial t} + \\frac{1}{2} (\\vec{1} \\cdot \\nabla)\n", + "(u^2) = \\nu \\Delta u $$" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "" + ], + "text/plain": [ + "" + ] + }, + "execution_count": 11, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "DOMAIN_EXTENT = 1.0\n", + "NUM_POINTS = 100\n", + "DT = 0.01\n", + "NU = 0.01\n", + "\n", + "single_channel_burgers_stepper = ex.stepper.Burgers(\n", + " 2, DOMAIN_EXTENT, NUM_POINTS, DT, diffusivity=NU, single_channel=True\n", + ")\n", + "\n", + "grid = ex.make_grid(2, DOMAIN_EXTENT, NUM_POINTS)\n", + "\n", + "u_0 = jnp.sin(2 * jnp.pi * grid[0:1] / DOMAIN_EXTENT) * jnp.cos(\n", + " 2 * 2 * jnp.pi * grid[1:2] / DOMAIN_EXTENT\n", + ")\n", + "\n", + "single_channel_burgers_trj = ex.rollout(\n", + " single_channel_burgers_stepper, 40, include_init=True\n", + ")(u_0)\n", + "\n", + "single_channel_burgers_ani = ex.viz.animate_state_2d(single_channel_burgers_trj)\n", + "\n", + "HTML(single_channel_burgers_ani.to_html5_video())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Kuramoto-Sivashinsky (KS)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The combustion format (using the gradient norm) generalizes nicely to higher dimensions\n", + "\n", + "$$\n", + "\\frac{\\partial u}{\\partial t} + \\frac{1}{2} \\| \\nabla u \\|_2^2 + \\Delta u + \\Delta^2 u = 0\n", + "$$" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "" + ], + "text/plain": [ + "" + ] + }, + "execution_count": 13, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "DOMAIN_EXTENT = 30.0\n", + "NUM_POINTS = 100\n", + "DT = 0.3\n", + "\n", + "ks_stepper = ex.stepper.KuramotoSivashinsky(2, DOMAIN_EXTENT, NUM_POINTS, DT)\n", + "\n", + "grid = ex.make_grid(2, DOMAIN_EXTENT, NUM_POINTS)\n", + "\n", + "\n", + "# IC is irrelevant\n", + "u_0 = jax.random.normal(jax.random.PRNGKey(0), (1, NUM_POINTS, NUM_POINTS))\n", + "\n", + "warmed_up_u_0 = ex.repeat(ks_stepper, 500)(u_0)\n", + "\n", + "ks_trj = ex.rollout(ks_stepper, 40, include_init=True)(warmed_up_u_0)\n", + "\n", + "ks_ani = ex.viz.animate_state_2d(ks_trj, vlim=(-6, 6))\n", + "\n", + "HTML(ks_ani.to_html5_video())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Korteweg-de Vries (KdV)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Works best with single channel hack\n", + "\n", + "$$\n", + "\\frac{\\partial u}{\\partial t} + \\frac{1}{2} (\\vec{1} \\cdot \\nabla) u^2 + \\vec{1} \\cdot (\\nabla \\odot \\nabla \\odot \\nabla) u = 0\n", + "$$" + ] + }, + { + "cell_type": "code", + "execution_count": 22, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "" + ], + "text/plain": [ + "" + ] + }, + "execution_count": 22, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "DOMAIN_EXTENT = 20.0\n", + "NUM_POINTS = 100\n", + "DT = 0.05\n", + "\n", + "kdv_stepper = ex.stepper.KortewegDeVries(\n", + " 2,\n", + " DOMAIN_EXTENT,\n", + " NUM_POINTS,\n", + " DT,\n", + " single_channel=True,\n", + ")\n", + "\n", + "grid = ex.make_grid(2, DOMAIN_EXTENT, NUM_POINTS)\n", + "u_0 = jnp.sin(2 * jnp.pi * grid[0:1] / DOMAIN_EXTENT) * jnp.cos(\n", + " 2 * 2 * jnp.pi * grid[1:2] / DOMAIN_EXTENT\n", + ")\n", + "\n", + "kdv_trj = ex.rollout(kdv_stepper, 40, include_init=True)(u_0)\n", + "\n", + "kdv_ani = ex.viz.animate_state_2d(kdv_trj)\n", + "\n", + "HTML(kdv_ani.to_html5_video())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Decaying turbulence" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Navier-Stokes in streamfunction-vorticity formulation\\\n", + "\n", + "$$\n", + "\\frac{\\partial u}{\\partial t} + \\left( \\begin{bmatrix} 1 \\\\ -1 \\end{bmatrix} \\odot \\nabla (\\Delta^{-1} u)\\right) \\cdot \\nabla u = \\nu \\Delta u\n", + "$$" + ] + }, + { + "cell_type": "code", + "execution_count": 35, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "" + ], + "text/plain": [ + "" + ] + }, + "execution_count": 35, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "DOMAIN_EXTENT = 1.0\n", + "NUM_POINTS = 100\n", + "DT = 1.0\n", + "NU = 0.0003\n", + "\n", + "decaying_turbulence_ns_stepper = ex.RepeatedStepper(\n", + " ex.stepper.NavierStokesVorticity(\n", + " 2, DOMAIN_EXTENT, NUM_POINTS, DT / 10, diffusivity=NU\n", + " ),\n", + " 10,\n", + ")\n", + "\n", + "u_0 = ex.ic.DiffusedNoise(2, max_one=True)(NUM_POINTS, key=jax.random.PRNGKey(0))\n", + "\n", + "decaying_turbulence_ns_trj = ex.rollout(\n", + " decaying_turbulence_ns_stepper, 40, include_init=True\n", + ")(u_0)\n", + "\n", + "decaying_turbulence_ns_ani = ex.viz.animate_state_2d(decaying_turbulence_ns_trj)\n", + "\n", + "HTML(decaying_turbulence_ns_ani.to_html5_video())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Kolmogorow Flow" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Uses streamline vorticity formulation\n", + "\n", + "$$\n", + "\\frac{\\partial u}{\\partial t} + \\left( \\begin{bmatrix} 1 \\\\ -1 \\end{bmatrix} \\odot \\nabla (\\Delta^{-1} u)\\right) \\cdot \\nabla u = \\lambda u + \\nu \\Delta u + f\n", + "$$\n", + "\n", + "with a forcing on a specific mode over y direction" + ] + }, + { + "cell_type": "code", + "execution_count": 33, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "" + ], + "text/plain": [ + "" + ] + }, + "execution_count": 33, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "DOMAIN_EXTENT = 2 * jnp.pi\n", + "NUM_POINTS = 100\n", + "DT = 0.5\n", + "NU = 0.001\n", + "DRAG = -0.1\n", + "INJECTION_MODE = 4\n", + "INJECTION_SCALE = 1.0\n", + "\n", + "kolmogorow_flow_ns_stepper = ex.RepeatedStepper(\n", + " ex.stepper.KolmogorovFlowVorticity(\n", + " 2,\n", + " DOMAIN_EXTENT,\n", + " NUM_POINTS,\n", + " DT / 50,\n", + " diffusivity=NU,\n", + " drag=DRAG,\n", + " injection_mode=INJECTION_MODE,\n", + " injection_scale=INJECTION_SCALE,\n", + " ),\n", + " 50,\n", + ")\n", + "\n", + "u_0 = ex.ic.DiffusedNoise(2, max_one=True, zero_mean=True)(\n", + " NUM_POINTS, key=jax.random.PRNGKey(0)\n", + ")\n", + "\n", + "warmed_up_u_0 = ex.repeat(kolmogorow_flow_ns_stepper, 500)(u_0)\n", + "\n", + "kolmogorow_flow_ns_trj = ex.rollout(kolmogorow_flow_ns_stepper, 40, include_init=True)(\n", + " warmed_up_u_0\n", + ")\n", + "\n", + "kolmogorow_flow_ns_ani = ex.viz.animate_state_2d(\n", + " kolmogorow_flow_ns_trj, vlim=(-6.0, 6.0)\n", + ")\n", + "\n", + "HTML(kolmogorow_flow_ns_ani.to_html5_video())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Reaction-Diffusion PDEs" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Fisher-KPP" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "$$\n", + "\\frac{\\partial u}{\\partial t} = \\nu \\Delta u + r u (1 - u)\n", + "$$" + ] + }, + { + "cell_type": "code", + "execution_count": 41, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "" + ], + "text/plain": [ + "" + ] + }, + "execution_count": 41, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "DOMAIN_EXTENT = 10.0\n", + "NUM_POINTS = 100\n", + "DT = 0.01\n", + "DIFFUSIVITY = 0.01\n", + "REACTIVITY = 10.0\n", + "\n", + "fisher_kpp_stepper = ex.reaction.FisherKPP(\n", + " 2, DOMAIN_EXTENT, NUM_POINTS, DT, diffusivity=DIFFUSIVITY, reactivity=REACTIVITY\n", + ")\n", + "\n", + "ic_gen = ex.ic.ClampingICGenerator(ex.ic.RandomTruncatedFourierSeries(2), limits=(0, 1))\n", + "u_0 = ic_gen(100, key=jax.random.PRNGKey(0))\n", + "\n", + "fisher_kpp_trj = ex.rollout(fisher_kpp_stepper, 40, include_init=True)(u_0)\n", + "\n", + "fisher_kpp_ani = ex.viz.animate_state_2d(fisher_kpp_trj)\n", + "\n", + "HTML(fisher_kpp_ani.to_html5_video())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Gray-Scott" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "$$\n", + "\\begin{aligned}\n", + "\\frac{\\partial u_0}{\\partial t} &= \\nu_0 \\Delta u_0 - u_0 u_1^2 + f (1 - u_0) \\\\\n", + "\\frac{\\partial u_1}{\\partial t} &= \\nu_1 \\Delta u_1 + u_0 u_1^2 - (f + k) u_1\n", + "\\end{aligned}\n", + "$$" + ] + }, + { + "cell_type": "code", + "execution_count": 44, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "" + ], + "text/plain": [ + "" + ] + }, + "execution_count": 44, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "DOMAIN_EXTENT = 1.0\n", + "NUM_POINTS = 100\n", + "DT = 30.0\n", + "DIFFUSIVITY_0 = 2e-5\n", + "DIFFUSIVITY_1 = 1e-5\n", + "FEED_RATE = 0.04\n", + "KILL_RATE = 0.06\n", + "\n", + "gray_scott_stepper = ex.RepeatedStepper(\n", + " ex.reaction.GrayScott(\n", + " 2,\n", + " DOMAIN_EXTENT,\n", + " NUM_POINTS,\n", + " DT / 30,\n", + " diffusivity_1=DIFFUSIVITY_0,\n", + " diffusivity_2=DIFFUSIVITY_1,\n", + " feed_rate=FEED_RATE,\n", + " kill_rate=KILL_RATE,\n", + " ),\n", + " 30,\n", + ")\n", + "\n", + "u_0 = ex.ic.RandomMultiChannelICGenerator(\n", + " [\n", + " ex.ic.RandomGaussianBlobs(2, one_complement=True),\n", + " ex.ic.RandomGaussianBlobs(2),\n", + " ]\n", + ")(NUM_POINTS, key=jax.random.PRNGKey(0))\n", + "\n", + "gray_scott_trj = ex.rollout(gray_scott_stepper, 40, include_init=True)(u_0)\n", + "\n", + "gray_scott_ani = ex.viz.animate_state_2d_facet(\n", + " gray_scott_trj,\n", + " grid=(1, 2),\n", + " figsize=(7, 3),\n", + ")\n", + "\n", + "HTML(gray_scott_ani.to_html5_video())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Swift-Hohenberg" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "$$\n", + "\\frac{\\partial u}{\\partial t} = r u - (1 + \\Delta)^2 u + u^2 - u^3\n", + "$$" + ] + }, + { + "cell_type": "code", + "execution_count": 46, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "" + ], + "text/plain": [ + "" + ] + }, + "execution_count": 46, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "DOMAIN_EXTENT = 20.0 * jnp.pi\n", + "NUM_POINTS = 100\n", + "DT = 1.0\n", + "\n", + "swift_hohenberg_stepper = ex.RepeatedStepper(\n", + " ex.reaction.SwiftHohenberg(2, DOMAIN_EXTENT, NUM_POINTS, DT / 10),\n", + " 10,\n", + ")\n", + "\n", + "u_0 = ex.ic.RandomTruncatedFourierSeries(2, max_one=True)(\n", + " NUM_POINTS, key=jax.random.PRNGKey(0)\n", + ")\n", + "\n", + "swift_hohenberg_trj = ex.rollout(swift_hohenberg_stepper, 40, include_init=True)(u_0)\n", + "\n", + "swift_hohenberg_ani = ex.viz.animate_state_2d(swift_hohenberg_trj)\n", + "\n", + "HTML(swift_hohenberg_ani.to_html5_video())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# All together" + ] + }, + { + "cell_type": "code", + "execution_count": 48, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "" + ], + "text/plain": [ + "" + ] + }, + "execution_count": 48, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "joint_trj = jnp.concatenate(\n", + " [\n", + " advection_trj,\n", + " diffusion_trj,\n", + " anisotropic_diffusion_trj,\n", + " advection_diffusion_trj,\n", + " dispersion_trj,\n", + " hyper_diffusion_trj,\n", + " burgers_trj,\n", + " single_channel_burgers_trj,\n", + " kdv_trj,\n", + " ks_trj,\n", + " decaying_turbulence_ns_trj,\n", + " kolmogorow_flow_ns_trj,\n", + " fisher_kpp_trj,\n", + " gray_scott_trj,\n", + " swift_hohenberg_trj,\n", + " ],\n", + " axis=1,\n", + ")\n", + "\n", + "joint_ani = ex.viz.animate_state_2d_facet(\n", + " joint_trj,\n", + " titles=[\n", + " \"Advection\",\n", + " \"Diffusion\",\n", + " \"Anisotropic Diffusion\",\n", + " \"Advection-Diffusion\",\n", + " \"Dispersion\",\n", + " \"Hyper-Diffusion\",\n", + " \"Burgers channel 1\",\n", + " \"Burgers channel 2\",\n", + " \"Burgers single channel\",\n", + " \"KdV\",\n", + " \"KS\",\n", + " \"Decaying Turbulence\",\n", + " \"Kolmogorov Flow\",\n", + " \"Fisher-KPP\",\n", + " \"Gray-Scott 1\",\n", + " \"Gray-Scott 2\",\n", + " \"Swift-Hohenberg\",\n", + " \"\",\n", + " ],\n", + " grid=(3, 6),\n", + " figsize=(17, 10),\n", + ")\n", + "\n", + "HTML(joint_ani.to_html5_video())" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "jax_fresh", + "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.10.13" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +}