Skip to content

Commit

Permalink
added videos notebook
Browse files Browse the repository at this point in the history
  • Loading branch information
@KeKsBoTer
KeKsBoTer committed Jun 7, 2024
1 parent bd2eb0b commit 939e984
Showing 1 changed file with 328 additions and 0 deletions.
328 changes: 328 additions & 0 deletions videos.ipynb
View file
Original file line number Diff line number Diff line change
@@ -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
}

0 comments on commit 939e984

Please sign in to comment.