From 939e98402343e7cec48b206616ca5bbb6deb10ba Mon Sep 17 00:00:00 2001 From: Simon Niedermayr Date: Fri, 7 Jun 2024 13:30:41 +0200 Subject: [PATCH] added videos notebook --- videos.ipynb | 328 +++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 328 insertions(+) create mode 100644 videos.ipynb diff --git a/videos.ipynb b/videos.ipynb new file mode 100644 index 0000000..f62d76c --- /dev/null +++ b/videos.ipynb @@ -0,0 +1,328 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from IPython.display import HTML\n", + "from matplotlib import pyplot as plt\n", + "import jax.numpy as jnp\n", + "import numpy as np\n", + "import exponax as ex\n", + "import jax\n", + "import cmasher\n", + "\n", + "import seaborn as sns\n", + "sns.set_theme()\n", + "\n", + "from vape import diverging_alpha,render,viewer" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "cmap_linear = cmasher.watermelon\n", + "cmap_nonlinear = sns.color_palette(\"icefire\", as_cmap=True)\n", + "cmap_diff = cmasher.copper_s\n", + "cmap_diff" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "resolution = 1024\n", + "fps = 30" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "import copy\n", + "from matplotlib.colors import LinearSegmentedColormap, ListedColormap\n", + "\n", + "\n", + "def triangle_wave(x,p):\n", + " return 2*np.abs(x/p-np.floor(x/p+0.5)) \n", + "\n", + "\n", + "def zigzag_alpha(cmap,min_alpha=0.2):\n", + " \"\"\"changes the alpha channel of a colormap to be linear (0->0, 1->1)\n", + "\n", + " Args:\n", + " cmap (Colormap): colormap\n", + "\n", + " Returns:a\n", + " Colormap: new colormap\n", + " \"\"\"\n", + " if isinstance(cmap, ListedColormap):\n", + " colors = copy.deepcopy(cmap.colors)\n", + " for i, a in enumerate(colors):\n", + " a.append((triangle_wave(i / (cmap.N - 1),0.5)*(1-min_alpha))+min_alpha)\n", + " return ListedColormap(colors, cmap.name)\n", + " elif isinstance(cmap, LinearSegmentedColormap):\n", + " segmentdata = copy.deepcopy(cmap._segmentdata)\n", + " segmentdata[\"alpha\"] = np.array([\n", + " [0.0, 0.0, 0.0],\n", + " [0.25, 1.0, 1.0],\n", + " [0.5, 0.0, 0.0],\n", + " [0.75, 1.0, 1.0],\n", + " [1.0, 0.0, 0.0]]\n", + " )\n", + " return LinearSegmentedColormap(cmap.name,segmentdata)\n", + " else:\n", + " raise TypeError(\n", + " \"cmap must be either a ListedColormap or a LinearSegmentedColormap\"\n", + " )\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from IPython.display import Video\n", + "import imageio.v2 as iio\n", + "import os\n", + "from tqdm import tqdm\n", + "os.environ[\"IMAGEIO_FFMPEG_EXE\"] = \"ffmpeg\" \n", + "\n", + "def symmetric_min_max(arr):\n", + " vmin = arr.min().item()\n", + " vmax = arr.max().item()\n", + " absmax = max(abs(vmin), abs(vmax))\n", + " return -absmax, absmax\n", + "\n", + "def chunk_list(lst, n):\n", + " for i in range(0, len(lst), n):\n", + " yield lst[i:i + n]\n", + " \n", + "def create_video(out_file:str,volume, cmap,resolution=1024, duration=10, fps=30, vrange=None) -> Video:\n", + " \"\"\"\n", + " Create a video from a volume\n", + " \n", + " Args:\n", + " out_file (str): output file\n", + " volume (np.ndarray): volume of shape [N,C,H,W,D]\n", + " cmap (Colormap): colormap\n", + " resolution (int): video width and height in pixels\n", + " duration (int): duration in seconds\n", + " fps (int): frames per second\n", + " vrange (float): range of values to display\n", + " \"\"\"\n", + " if vrange is None:\n", + " vmin, vmax = symmetric_min_max(volume)\n", + " elif isinstance(vrange, (int, float)):\n", + " vmin, vmax = -vrange, vrange\n", + " elif isinstance(vrange, (tuple, list)):\n", + " vmin, vmax = vrange\n", + " else:\n", + " raise ValueError(\"vrange must be None, a number, or a tuple of two numbers\")\n", + " \n", + " num_channels = volume.shape[1]\n", + " n = duration*fps\n", + " \n", + " print(\"num_channels=\",num_channels,\", vmin=\",vmin,\",vmax=\",vmax)\n", + "\n", + " with iio.get_writer(out_file, format='FFMPEG', mode='I', fps=fps,codec='h264_nvenc',) as w:\n", + " batch_size = 64 # how many frames to store in memory at once\n", + " for time_steps in tqdm(list(chunk_list(list(range(n)), batch_size))):\n", + " frames = np.zeros((len(time_steps), resolution, resolution*num_channels, 3), dtype=np.uint8)\n", + " for c in range(num_channels):\n", + " imgs = render(\n", + " np.array(volume[:, c]),\n", + " cmap,\n", + " [i / (n-1) for i in time_steps],\n", + " background=(0, 0, 0, 255),\n", + " distance_scale=10,\n", + " vmin=vmin,\n", + " vmax=vmax,\n", + " width=resolution,\n", + " height=resolution,\n", + " )\n", + " # gamma correction\n", + " imgs = ((imgs/255.)**(2.4)*255).astype(np.uint8)\n", + " frames[:,:,resolution*c:resolution*(c+1)] = imgs[..., :3]\n", + " for img in frames:\n", + " w.append_data(img)\n", + "\n", + " w.close()\n", + " return Video(url=out_file)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "from os import makedirs\n", + "\n", + "\n", + "def symmetric_min_max(arr):\n", + " vmin = arr.min().item()\n", + " vmax = arr.max().item()\n", + " absmax = max(abs(vmin), abs(vmax))\n", + " return -absmax,absmax\n", + "#\n", + "\n", + "\n", + "DOMAIN_EXTENT = 1.0\n", + "NUM_POINTS = 64\n", + "DT = 0.01\n", + "NU = 0.01\n", + "\n", + "burgers_stepper = ex.stepper.Burgers(3, DOMAIN_EXTENT, NUM_POINTS, DT, diffusivity=NU)\n", + "\n", + "grid = ex.make_grid(3, DOMAIN_EXTENT, NUM_POINTS)\n", + "\n", + "\n", + "ic_gen = ex.ic.RandomTruncatedFourierSeries(3, cutoff=2, max_one=True)\n", + "multi_channel_ic_gen = ex.ic.RandomMultiChannelICGenerator([ic_gen, ic_gen, ic_gen])\n", + "u_0 = multi_channel_ic_gen(NUM_POINTS, key=jax.random.PRNGKey(1))\n", + "\n", + "burgers_trj_3d = ex.rollout(burgers_stepper, 128, include_init=True)(u_0)\n", + "\n", + "makedirs(\"videos\", exist_ok=True)\n", + "create_video(\"videos/burgers.mp4\",burgers_trj_3d,zigzag_alpha(cmap_nonlinear, 0.1),resolution=resolution, duration=10, fps=fps, vrange=0.4)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "DOMAIN_EXTENT = 1.0\n", + "NUM_POINTS = 64\n", + "DT = 0.01\n", + "VELOCITY = 1.0\n", + "\n", + "advection_stepper = ex.stepper.Advection(\n", + " 3,\n", + " DOMAIN_EXTENT,\n", + " NUM_POINTS,\n", + " DT,\n", + " velocity=np.array([VELOCITY * 2, VELOCITY * 0.4, VELOCITY]),\n", + ")\n", + "\n", + "u_0 = ex.ic.DiffusedNoise(3, max_one=True, zero_mean=True)(\n", + " NUM_POINTS, key=jax.random.PRNGKey(0)\n", + ")\n", + "advection_trj_3d = ex.rollout(advection_stepper, 64, include_init=True)(u_0)\n", + "\n", + "create_video(\"videos/advection.mp4\",advection_trj_3d,zigzag_alpha(cmap_linear, 0.1),resolution=resolution, duration=10, fps=fps, vrange=0.5)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "DOMAIN_EXTENT = 30.0\n", + "NUM_POINTS = 64\n", + "DT = 0.1\n", + "\n", + "ks_stepper = ex.stepper.KuramotoSivashinsky(3, DOMAIN_EXTENT, NUM_POINTS, DT)\n", + "\n", + "# IC is irrelevant\n", + "u_0 = jax.random.normal(jax.random.PRNGKey(0), (1, NUM_POINTS, NUM_POINTS, NUM_POINTS))\n", + "warmed_up_u_0 = ex.repeat(ks_stepper, 500)(u_0)\n", + "ks_trj_3d = ex.rollout(ks_stepper, 64, include_init=True)(warmed_up_u_0)\n", + "\n", + "create_video(\"videos/ks.mp4\",ks_trj_3d,zigzag_alpha(cmap_nonlinear, 0.1),resolution=resolution, duration=10, fps=fps, vrange=3.0)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "DOMAIN_EXTENT = 1.0\n", + "NUM_POINTS = 64\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", + " 3,\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", + " 15,\n", + ")\n", + "\n", + "u_0 = ex.ic.RandomMultiChannelICGenerator(\n", + " [\n", + " ex.ic.RandomGaussianBlobs(3, one_complement=True, num_blobs=1),\n", + " ex.ic.RandomGaussianBlobs(3, num_blobs=1),\n", + " ]\n", + ")(NUM_POINTS, key=jax.random.PRNGKey(0))\n", + "\n", + "gray_scott_trj_3d = ex.rollout(gray_scott_stepper, 128, include_init=True)(u_0)\n", + "\n", + "\n", + "create_video(\"videos/gray_scott.mp4\",gray_scott_trj_3d,zigzag_alpha(cmap_diff, 0.0),resolution=resolution, duration=10, fps=fps, vrange=(0,1))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "ano", + "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.3" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +}