diff_game_of_life.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# convolution based attempt of the game of life\n",
    "import torch\n",
    "import numpy as np\n",
    "import cv2\n",
    "import os\n",
    "import gc\n",
    "import tqdm\n",
    "from functools import partial\n",
    "torch.set_num_threads(2)\n",
    "video_output_path = 'ca-experiment.mp4'\n",
    "in_channels = 3\n",
    "out_channels = 3\n",
    "groups = 1\n",
    "batchsize = 1\n",
    "fps = 30\n",
    "length_seconds = 60\n",
    "grid_resolution = (384, 216)\n",
    "output_resolution = (1280, 720)\n",
    "rule = [np.exp(1)/3,np.exp(1),0.98*np.exp(1),1.02*np.exp(1)]\n",
    "\n",
    "conv_out = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tocv2_img(src, shape):\n",
    "    img = src.permute(0,2,3,1).mean(0).squeeze().cpu().numpy()\n",
    "    img /= img.max()\n",
    "    img = (255*img).astype('uint8')\n",
    "    img = cv2.resize(img, shape, interpolation=cv2.INTER_NEAREST)\n",
    "    return img\n",
    "\n",
    "def set_seed(grid, seed, loc, index=None):\n",
    "    shp = seed.shape[-2:]\n",
    "    if index is None:\n",
    "        grid[:, :, loc[0]:shp[0]+loc[0], loc[1]:shp[1]+loc[1]] = seed\n",
    "    else:\n",
    "        grid[index[0], index[1], loc[0]:shp[0]+loc[0], log[1]:shp[1]+loc[1]] = seed\n",
    "\n",
    "def step(grid, kernels, rules, detach=True, circular_padding=True):\n",
    "    global conv_out\n",
    "    # pad input  \n",
    "    pad = (kernels.shape[-2]/2, kernels.shape[-2]/2,\n",
    "           kernels.shape[-1]/2, kernels.shape[-1]/2)\n",
    "    inps = torch.nn.functional.pad(grid, pad)\n",
    "    if circular_padding:\n",
    "        # sides\n",
    "        inps[:, :, 0:pad[0], pad[2]:-pad[3]] = grid[:, :, -pad[0]:, :]\n",
    "        inps[:, :, -pad[1]:, pad[2]:-pad[3]] = grid[:, :, 0:pad[1]:, :]\n",
    "        inps[:, :, pad[0]:-pad[1], 0:pad[2]] = grid[:, :, :, -pad[2]:]\n",
    "        inps[:, :, pad[0]:-pad[1], -pad[3]:] = grid[:, :, :, 0:pad[3]:]\n",
    "        # corners\n",
    "        inps[:, :, 0:pad[0], 0:pad[2]] = grid[:, :, -pad[0]:, -pad[2]:]\n",
    "        inps[:, :, -pad[1]:, 0:pad[2]] = grid[:, :, 0:pad[1]:, -pad[2]:]\n",
    "        inps[:, :, 0:pad[0], -pad[3]:] = grid[:, :, -pad[0]:, 0:pad[3]]\n",
    "        inps[:, :, -pad[1]:, -pad[3]:] = grid[:, :, 0:pad[1], 0:pad[3]:]\n",
    "\n",
    "    # convolve grid with kernels\n",
    "    conv_out = torch.conv1d(inps, kernels, groups=groups)\n",
    "    \n",
    "    # update grid from rules. this will do a logical or between the list of rules\n",
    "    update = -grid\n",
    "    for rule in rules:\n",
    "        update = update + rule(conv_out, grid)\n",
    "    grid = grid + update\n",
    "    return grid\n",
    "\n",
    "def soft_ge(x, value, scale=16, bias=4):\n",
    "    return ((x-value)*scale+bias).sigmoid()\n",
    "\n",
    "def soft_le(x, value, scale=16, bias=4):\n",
    "    return (-(x-value)*scale+bias).sigmoid()\n",
    "\n",
    "def soft_or(x, y, scale=16, bias=4):\n",
    "    return ((x+y)*scale-bias).sigmoid()\n",
    "\n",
    "def game_of_life_rule(sum_grid, grid, rule=[2,3,3,3]):\n",
    "    diff_grid = sum_grid - grid\n",
    "    g_img = tocv2_img(grid, output_resolution)\n",
    "    dg_img = tocv2_img(diff_grid, output_resolution)\n",
    "    population_cond = grid*(soft_ge(diff_grid, rule[0])*(soft_le(diff_grid, rule[1])))\n",
    "    reprod_cond = (1-grid)*(soft_ge(diff_grid, rule[2])*(soft_le(diff_grid, rule[3])))\n",
    "    surviving_cells = soft_or(population_cond, reprod_cond)\n",
    "    return surviving_cells\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the format is a bit weird, to conform to the minibatch processing API of the conv2d function. \n",
    "grid = torch.zeros((batchsize, in_channels, grid_resolution[-1], grid_resolution[-2]))\n",
    "# convolution with this filter results in adding up the  number of ones in a 3x3 neighborhood\n",
    "# dimensions are (input channels, output channels, height width)\n",
    "def build_filters():\n",
    "    sum_filter_3_3 = torch.ones(in_channels, out_channels, 3, 3)\n",
    "    ind_sum_filter_3_3 = torch.stack([(i==j)*torch.ones(3, 3)\n",
    "                                      for i in range(1, in_channels+1)\n",
    "                                      for j in range(1, out_channels+1)]).reshape(in_channels, out_channels, 3, 3)\n",
    "    weighted_channel_filter_3_3 = ind_sum_filter_3_3.clone().uniform_()\n",
    "    weighted_channel_filter_3_3 /= weighted_channel_filter_3_3.sum(0) \n",
    "    return [sum_filter_3_3, ind_sum_filter_3_3, weighted_channel_filter_3_3]\n",
    "\n",
    "# the seed as a sub image\n",
    "seed = torch.tensor([[[0,1,1,0],\n",
    "                      [1,1,0,0],\n",
    "                      [0,1,0,0],\n",
    "                      [0,0,0,0]],\n",
    "                     [[0,0,0,0],\n",
    "                      [0,1,1,0],\n",
    "                      [0,0,1,1],\n",
    "                      [0,0,1,0]],\n",
    "                     [[0,0,0,0],\n",
    "                      [0,0,1,0],\n",
    "                      [0,0,1,1],\n",
    "                      [0,1,1,0]]]).float()\n",
    "seed = seed.repeat(batchsize, 1, 1, 1)\n",
    "\n",
    "# initialize grid\n",
    "set_seed(grid, seed, ((grid.shape[-2]-seed.shape[-2])/2, (grid.shape[-1]-seed.shape[-1])/2))\n",
    "\n",
    "locs = np.random.uniform([8,8], [grid.shape[-2]-8, grid.shape[-1]-8], size=(10,2)).astype('uint32')\n",
    "for loc in locs:\n",
    "    set_seed(grid, (torch.rand_like(seed) > 0.5).float(), loc)\n",
    "    \n",
    "\n",
    "kernels_list = build_filters()\n",
    "\n",
    "gol_rule = partial(game_of_life_rule, rule=rule)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 26%|██▌       | 459/1800 [00:32<01:34, 14.22it/s]"
     ]
    }
   ],
   "source": [
    "cv2.destroyAllWindows()\n",
    "try:\n",
    "    fourcc = cv2.VideoWriter_fourcc(*'H264')\n",
    "    vw = cv2.VideoWriter(video_output_path, fourcc, fps, output_resolution)\n",
    "except Exception as e:\n",
    "    print e\n",
    "    vw = None\n",
    "        \n",
    "if torch.cuda.is_available():\n",
    "    grid = grid.cuda()\n",
    "    kernels_list = [kernels.cuda() for kernels in kernels_list]\n",
    "    \n",
    "conv_out = grid.clone()\n",
    "for i in tqdm.tqdm(range(int(fps*length_seconds))):\n",
    "    # every once in a while, change the way neighbours are counted\n",
    "    if i%61 == 1:\n",
    "        kernels_list = build_filters()\n",
    "        if torch.cuda.is_available():\n",
    "            kernels_list = [kernels.cuda() for kernels in kernels_list]\n",
    "    # every once in a while, drop random seeds at random locations\n",
    "    if i%31 == 1:\n",
    "        locs = np.random.uniform([8,8], [grid.shape[-2]-8, grid.shape[-1]-8], size=(3,2)).astype('uint32')\n",
    "        for loc in locs:\n",
    "            set_seed(grid, (torch.rand_like(seed) > 0.5).float(), loc)\n",
    "    kernels = kernels_list[2]\n",
    "    # convert to image with channels as last dimension, and desired output resolution\n",
    "    grid_img = tocv2_img(grid, output_resolution)\n",
    "    \n",
    "    # show current grid state\n",
    "    cv2.imshow('big bang orig', grid_img)\n",
    "    cv2.waitKey(int(1000.0/fps))\n",
    "    \n",
    "    # write output\n",
    "    if vw is not None:\n",
    "        vw.write(grid_img)\n",
    "    \n",
    "    # step the cellular automaton\n",
    "    grid = step(grid, kernels, (gol_rule,))\n",
    "    \n",
    "    # clean up memory\n",
    "    gc.collect()\n",
    "    torch.cuda.empty_cache()\n",
    "\n",
    "vw.release()\n",
    "cv2.destroyAllWindows()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}