block_sparse_toy_opt.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy\n",
    "import itertools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "b = 5\n",
    "n = 3\n",
    "\n",
    "ta = b * n\n",
    "\n",
    "max_dist = 3\n",
    "\n",
    "dense = numpy.full((ta, ta), max_dist, int)\n",
    "for i in range(max_dist):\n",
    "    dense[numpy.arange(ta - i), numpy.arange(ta - i) + i] = i\n",
    "    dense[numpy.arange(ta - i) + i, numpy.arange(ta - i)] = i\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "bmap = numpy.full((b, b), 0, int)\n",
    "null_block = numpy.full((n, n), max_dist, int)\n",
    "\n",
    "blocks = [null_block]\n",
    "\n",
    "for bi, bj in itertools.product(range(b), range(b)):\n",
    "    sub_block = dense[bi * n:(bi+1) * n, bj * n : (bj+1) * n]\n",
    "    \n",
    "    for prev_block_i, prev_block in enumerate(blocks):\n",
    "        if numpy.all(sub_block == prev_block):\n",
    "            bmap[bi, bj] = prev_block_i\n",
    "            break\n",
    "    else:\n",
    "        bmap[bi, bj] = len(blocks)\n",
    "        blocks.append(sub_block)\n",
    "\n",
    "blocks = numpy.concatenate(\n",
    "    [numpy.expand_dims(b, 0) for b in blocks],\n",
    "    axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0, 1, 2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3],\n",
       "       [1, 0, 1, 2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3],\n",
       "       [2, 1, 0, 1, 2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3],\n",
       "       [3, 2, 1, 0, 1, 2, 3, 3, 3, 3, 3, 3, 3, 3, 3],\n",
       "       [3, 3, 2, 1, 0, 1, 2, 3, 3, 3, 3, 3, 3, 3, 3],\n",
       "       [3, 3, 3, 2, 1, 0, 1, 2, 3, 3, 3, 3, 3, 3, 3],\n",
       "       [3, 3, 3, 3, 2, 1, 0, 1, 2, 3, 3, 3, 3, 3, 3],\n",
       "       [3, 3, 3, 3, 3, 2, 1, 0, 1, 2, 3, 3, 3, 3, 3],\n",
       "       [3, 3, 3, 3, 3, 3, 2, 1, 0, 1, 2, 3, 3, 3, 3],\n",
       "       [3, 3, 3, 3, 3, 3, 3, 2, 1, 0, 1, 2, 3, 3, 3],\n",
       "       [3, 3, 3, 3, 3, 3, 3, 3, 2, 1, 0, 1, 2, 3, 3],\n",
       "       [3, 3, 3, 3, 3, 3, 3, 3, 3, 2, 1, 0, 1, 2, 3],\n",
       "       [3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 2, 1, 0, 1, 2],\n",
       "       [3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 2, 1, 0, 1],\n",
       "       [3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 2, 1, 0]])"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dense"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[1, 2, 0, 0, 0],\n",
       "       [3, 1, 2, 0, 0],\n",
       "       [0, 3, 1, 2, 0],\n",
       "       [0, 0, 3, 1, 2],\n",
       "       [0, 0, 0, 3, 1]])"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bmap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[3, 3, 3],\n",
       "        [3, 3, 3],\n",
       "        [3, 3, 3]],\n",
       "\n",
       "       [[0, 1, 2],\n",
       "        [1, 0, 1],\n",
       "        [2, 1, 0]],\n",
       "\n",
       "       [[3, 3, 3],\n",
       "        [2, 3, 3],\n",
       "        [1, 2, 3]],\n",
       "\n",
       "       [[3, 2, 1],\n",
       "        [3, 3, 2],\n",
       "        [3, 3, 3]]])"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "blocks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}