sandbox.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "cb008513-dcb9-4063-927f-b65070b43ca0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import h5py\n",
    "import torch_geometric as pyg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "485a812d-2ebe-43ee-a720-2af83006efae",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'3.5.0'"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "h5py.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "dbeb6517-b648-4f0f-b251-89f9bed68d6a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'6.0'"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import yaml\n",
    "yaml.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "8b3c52d5-26eb-4e84-9b1f-4554171c25c0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<KeysViewHDF5 ['filt_8', 'filt_grad_8', 'grad_filt_8']>\n"
     ]
    }
   ],
   "source": [
    "with h5py.File('/home/jupyter/the_hard_way/incorporated/forkane/data/raw/DNS1_00116000.h5') as file:\n",
    "    print(file.keys())\n",
    "    x = file['/grad_filt_8'][:]\n",
    "    y = file['/filt_8'][:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "a62c274f-9a63-40e6-a17f-b015930b682f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(65, 33, 33)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "grid_shape = x.shape\n",
    "grid_shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "c1eec6b6-333e-4450-bc40-f50c7bed3d79",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['DNS1_00125000.h5',\n",
       " 'DNS1_00121000.h5',\n",
       " 'DNS1_00119400.h5',\n",
       " 'DNS1_00120400.h5',\n",
       " 'DNS3_00128200.h5',\n",
       " 'DNS2_00119200.h5',\n",
       " 'DNS2_00121800.h5',\n",
       " 'DNS1_00121800.h5',\n",
       " 'DNS1_00126800.h5',\n",
       " 'DNS3_00127000.h5',\n",
       " 'DNS2_00122200.h5',\n",
       " 'DNS1_00119200.h5',\n",
       " 'DNS1_00120000.h5',\n",
       " 'DNS3_00129800.h5',\n",
       " 'DNS1_00117200.h5',\n",
       " 'DNS1_00116400.h5',\n",
       " 'DNS1_00124800.h5',\n",
       " 'DNS3_00128400.h5',\n",
       " 'DNS1_00118000.h5',\n",
       " 'DNS2_00117200.h5',\n",
       " 'DNS2_00122000.h5',\n",
       " 'DNS2_00116800.h5',\n",
       " 'DNS1_00124000.h5',\n",
       " 'DNS2_00125000.h5',\n",
       " 'DNS2_00118400.h5',\n",
       " 'DNS1_00122400.h5',\n",
       " 'DNS1_00125400.h5',\n",
       " 'DNS1_00121200.h5',\n",
       " 'DNS1_00121600.h5',\n",
       " 'DNS2_00117000.h5',\n",
       " 'DNS1_00117000.h5',\n",
       " 'DNS2_00118000.h5',\n",
       " 'DNS2_00125800.h5',\n",
       " 'DNS1_00123200.h5',\n",
       " 'DNS2_00116600.h5',\n",
       " 'DNS2_00117400.h5',\n",
       " 'DNS1_00127000.h5',\n",
       " 'DNS2_00121400.h5',\n",
       " 'DNS1_00116800.h5',\n",
       " 'DNS1_00123400.h5',\n",
       " 'DNS1_00119800.h5',\n",
       " 'DNS2_00120200.h5',\n",
       " 'DNS2_00124800.h5',\n",
       " 'DNS2_00118800.h5',\n",
       " 'DNS2_00124000.h5',\n",
       " 'DNS2_00125200.h5',\n",
       " 'DNS1_00116600.h5',\n",
       " 'DNS3_00129200.h5',\n",
       " 'DNS1_00122800.h5',\n",
       " 'DNS2_00116200.h5',\n",
       " 'DNS2_00118600.h5',\n",
       " 'DNS2_00124400.h5',\n",
       " 'DNS3_00127200.h5',\n",
       " 'DNS2_00117600.h5',\n",
       " 'DNS1_00126600.h5',\n",
       " 'DNS1_00125200.h5',\n",
       " 'DNS3_00128600.h5',\n",
       " 'DNS2_00116000.h5',\n",
       " 'DNS2_00122800.h5',\n",
       " 'DNS3_00127800.h5',\n",
       " 'DNS2_00118200.h5',\n",
       " 'DNS2_00123800.h5',\n",
       " 'DNS2_00120000.h5',\n",
       " 'DNS3_00129000.h5',\n",
       " 'DNS1_00123000.h5',\n",
       " 'DNS2_00119800.h5',\n",
       " 'DNS1_00118400.h5',\n",
       " 'DNS2_00126600.h5',\n",
       " 'DNS2_00119000.h5',\n",
       " 'DNS2_00123000.h5',\n",
       " 'DNS1_00123800.h5',\n",
       " 'DNS2_00125400.h5',\n",
       " 'DNS3_00128000.h5',\n",
       " 'DNS1_00124400.h5',\n",
       " 'DNS1_00122000.h5',\n",
       " 'DNS1_00126200.h5',\n",
       " 'DNS2_00123200.h5',\n",
       " 'DNS1_00118200.h5',\n",
       " 'DNS2_00126000.h5',\n",
       " 'DNS2_00124200.h5',\n",
       " 'DNS1_00125800.h5',\n",
       " 'DNS2_00122600.h5',\n",
       " 'DNS2_00121000.h5',\n",
       " 'DNS1_00122600.h5',\n",
       " 'DNS3_00127600.h5',\n",
       " 'DNS1_00118600.h5',\n",
       " 'DNS3_00130000.h5',\n",
       " 'DNS2_00124600.h5',\n",
       " 'DNS2_00126200.h5',\n",
       " 'DNS1_00117400.h5',\n",
       " 'DNS1_00124600.h5',\n",
       " 'DNS2_00125600.h5',\n",
       " 'DNS2_00119400.h5',\n",
       " 'DNS2_00121200.h5',\n",
       " 'DNS3_00127400.h5',\n",
       " 'DNS2_00120800.h5',\n",
       " 'DNS1_00126400.h5',\n",
       " 'DNS1_00126000.h5',\n",
       " 'DNS1_00120800.h5',\n",
       " 'DNS3_00129600.h5',\n",
       " 'DNS1_00120200.h5',\n",
       " 'DNS2_00119600.h5',\n",
       " 'DNS1_00123600.h5',\n",
       " 'DNS2_00123600.h5',\n",
       " 'DNS1_00121400.h5',\n",
       " 'DNS1_00119600.h5',\n",
       " 'DNS2_00116400.h5',\n",
       " 'DNS1_00116200.h5',\n",
       " 'DNS1_00119000.h5',\n",
       " 'DNS1_00125600.h5',\n",
       " 'DNS1_00118800.h5',\n",
       " 'DNS3_00129400.h5',\n",
       " 'DNS1_00116000.h5',\n",
       " 'DNS2_00121600.h5',\n",
       " 'DNS2_00120600.h5',\n",
       " 'DNS1_00124200.h5',\n",
       " 'DNS2_00126800.h5',\n",
       " 'DNS2_00127000.h5',\n",
       " 'DNS1_00117800.h5',\n",
       " 'DNS2_00120400.h5',\n",
       " 'DNS1_00117600.h5',\n",
       " 'DNS2_00122400.h5',\n",
       " 'DNS2_00126400.h5',\n",
       " 'DNS2_00123400.h5',\n",
       " 'DNS3_00128800.h5',\n",
       " 'DNS2_00117800.h5',\n",
       " 'DNS1_00120600.h5']"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "root_path = '/home/jupyter/the_hard_way/incorporated/forkane/data/'\n",
    "os.listdir(os.path.join(root_path))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "4ed9203c-3b14-4a7e-93b4-088c73cce4c1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['DNS1_00125000.h5', 'DNS1_00121000.h5', 'DNS1_00119400.h5', 'DNS1_00120400.h5', 'DNS3_00128200.h5', 'DNS2_00119200.h5', 'DNS2_00121800.h5', 'DNS1_00121800.h5', 'DNS1_00126800.h5', 'DNS3_00127000.h5', 'DNS2_00122200.h5', 'DNS1_00119200.h5', 'DNS1_00120000.h5', 'DNS3_00129800.h5', 'DNS1_00117200.h5', 'DNS1_00116400.h5', 'DNS1_00124800.h5', 'DNS3_00128400.h5', 'DNS1_00118000.h5', 'DNS2_00117200.h5', 'DNS2_00122000.h5', 'DNS2_00116800.h5', 'DNS1_00124000.h5', 'DNS2_00125000.h5', 'DNS2_00118400.h5', 'DNS1_00122400.h5', 'DNS1_00125400.h5', 'DNS1_00121200.h5', 'DNS1_00121600.h5', 'DNS2_00117000.h5', 'DNS1_00117000.h5', 'DNS2_00118000.h5', 'DNS2_00125800.h5', 'DNS1_00123200.h5', 'DNS2_00116600.h5', 'DNS2_00117400.h5', 'DNS1_00127000.h5', 'DNS2_00121400.h5', 'DNS1_00116800.h5', 'DNS1_00123400.h5', 'DNS1_00119800.h5', 'DNS2_00120200.h5', 'DNS2_00124800.h5', 'DNS2_00118800.h5', 'DNS2_00124000.h5', 'DNS2_00125200.h5', 'DNS1_00116600.h5', 'DNS3_00129200.h5', 'DNS1_00122800.h5', 'DNS2_00116200.h5', 'DNS2_00118600.h5', 'DNS2_00124400.h5', 'DNS3_00127200.h5', 'DNS2_00117600.h5', 'DNS1_00126600.h5', 'DNS1_00125200.h5', 'DNS3_00128600.h5', 'DNS2_00116000.h5', 'DNS2_00122800.h5', 'DNS3_00127800.h5', 'DNS2_00118200.h5', 'DNS2_00123800.h5', 'DNS2_00120000.h5', 'DNS3_00129000.h5', 'DNS1_00123000.h5', 'DNS2_00119800.h5', 'DNS1_00118400.h5', 'DNS2_00126600.h5', 'DNS2_00119000.h5', 'DNS2_00123000.h5', 'DNS1_00123800.h5', 'DNS2_00125400.h5', 'DNS3_00128000.h5', 'DNS1_00124400.h5', 'DNS1_00122000.h5', 'DNS1_00126200.h5', 'DNS2_00123200.h5', 'DNS1_00118200.h5', 'DNS2_00126000.h5', 'DNS2_00124200.h5', 'DNS1_00125800.h5', 'DNS2_00122600.h5', 'DNS2_00121000.h5', 'DNS1_00122600.h5', 'DNS3_00127600.h5', 'DNS1_00118600.h5', 'DNS3_00130000.h5', 'DNS2_00124600.h5', 'DNS2_00126200.h5', 'DNS1_00117400.h5', 'DNS1_00124600.h5', 'DNS2_00125600.h5', 'DNS2_00119400.h5', 'DNS2_00121200.h5', 'DNS3_00127400.h5', 'DNS2_00120800.h5', 'DNS1_00126400.h5', 'DNS1_00126000.h5', 'DNS1_00120800.h5', 'DNS3_00129600.h5', 'DNS1_00120200.h5', 'DNS2_00119600.h5', 'DNS1_00123600.h5', 'DNS2_00123600.h5', 'DNS1_00121400.h5', 'DNS1_00119600.h5', 'DNS2_00116400.h5', 'DNS1_00116200.h5', 'DNS1_00119000.h5', 'DNS1_00125600.h5', 'DNS1_00118800.h5', 'DNS3_00129400.h5', 'DNS1_00116000.h5', 'DNS2_00121600.h5', 'DNS2_00120600.h5', 'DNS1_00124200.h5', 'DNS2_00126800.h5', 'DNS2_00127000.h5', 'DNS1_00117800.h5', 'DNS2_00120400.h5', 'DNS1_00117600.h5', 'DNS2_00122400.h5', 'DNS2_00126400.h5', 'DNS2_00123400.h5', 'DNS3_00128800.h5', 'DNS2_00117800.h5', 'DNS1_00120600.h5']\n"
     ]
    }
   ],
   "source": [
    "with open(os.path.join(root_path, \"filenames.yaml\"), \"r\") as f:\n",
    "    filenames = yaml.safe_load(f)\n",
    "\n",
    "print(filenames)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "057877af-87ba-4781-975b-48d3789e6026",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Dataset(pyg.data.Dataset):\n",
    "    \n",
    "    def __init__(self, root, transform=None, pre_transform=None):\n",
    "        self._grid_shape = None\n",
    "        super().__init__(root, transform, pre_transform)\n",
    "\n",
    "    @property\n",
    "    def raw_file_names(self):\n",
    "        with open(os.path.join(self.root, \"filenames.yaml\"), \"r\") as stream:\n",
    "            filenames = yaml.safe_load(stream)\n",
    "        return filenames\n",
    "\n",
    "    @property\n",
    "    def processed_file_names(self):\n",
    "        return [f\"data-{idx}.pt\" for idx in range(self.len())]\n",
    "    \n",
    "    @property\n",
    "    def grid_shape(self):\n",
    "        if not self._grid_shape:\n",
    "            with h5py.File(self.raw_paths[0], 'r') as file:\n",
    "                field_val = file['filt_8'][:]\n",
    "            x_size, y_size, z_size = field_val.shape[0], field_val.shape[1], field_val.shape[2]\n",
    "            self._grid_shape = (z_size, y_size, x_size)\n",
    "        return self._grid_shape\n",
    "\n",
    "    def download(self):\n",
    "        raise RuntimeError(\n",
    "            'Data not found. Please download the data at {} and move all files in file.tgz/DATA in {}'.format(\n",
    "                'https://uftp.fz-juelich.de:9112/UFTP_Auth/rest/access/JUDAC/dc4eef36-1929-41f6-9eb9-c11417be1dcf',\n",
    "                self.raw_dir))\n",
    "\n",
    "    def process(self):\n",
    "        graph = pyg.utils.convert.from_networkx(\n",
    "            nx.grid_graph(dim=self.grid_shapae)\n",
    "        )\n",
    "        undirected_index = graph.edge_index\n",
    "        \n",
    "        i = 0\n",
    "        for raw_path in self.raw_paths:\n",
    "            \n",
    "            with h5py.File(raw_path, 'r') as file:\n",
    "                c = file[\"/filt_8\"][:]\n",
    "                sigma = file[\"/filt_grad_8\"][:]\n",
    "                \n",
    "            data = pyg.data.Data(\n",
    "                x=torch.tensor(c.reshape(-1,1), dtype=torch.float), \n",
    "                edge_index=torch.tensor(undirected_index, dtype=torch.long),\n",
    "                y=torch.tensor(sigma.reshape(-1,1), dtype=torch.float)\n",
    "            )\n",
    "\n",
    "            if self.pre_filter is not None and not self.pre_filter(data):\n",
    "                continue\n",
    "\n",
    "            if self.pre_transform is not None:\n",
    "                data = self.pre_transform(data)\n",
    "\n",
    "            torch.save(data, os.path.join(self.processed_dir, 'data-{}.pt'.format(i)))\n",
    "            i += 1\n",
    "\n",
    "    def get(self, idx):\n",
    "        data = torch.load(os.path.join(self.processed_dir, 'data-{}.pt'.format(idx)))\n",
    "        return data\n",
    "    \n",
    "    def len(self):\n",
    "        return len(self.raw_file_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2d8beb40-4938-4308-857d-dfa869360bf7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['hello', 'world!']"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"hello world!\".split()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "09fc6b73-0670-4115-9294-677562ba368a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0, 2, 2, 3]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hey = list((0,))\n",
    "hey.extend([2, 2, 3])\n",
    "hey"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "b77fc3c1-5a9d-401a-9b50-8417cab254b3",
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "'shape' is an invalid keyword argument for array()",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_5817/1039822120.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mnumpy\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 3\u001b[0;31m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0marray\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0marange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m16\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mshape\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m4\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m4\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m: 'shape' is an invalid keyword argument for array()"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "np.array(np.arange(16), shape=(4, 4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "id": "461dc7fb-ecc6-41ba-997d-14f72eb0ae46",
   "metadata": {},
   "outputs": [],
   "source": [
    "x_dim = 65\n",
    "y_dim = 33\n",
    "z_dim = 33"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "cd0094b2-66e0-475c-8859-6db07a96687e",
   "metadata": {},
   "outputs": [],
   "source": [
    "array = np.reshape(np.arange(x_dim * y_dim * z_dim), (x_dim, y_dim, z_dim))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "a39bfb59-59ee-4b2a-b541-dbe28ff7fe73",
   "metadata": {},
   "outputs": [],
   "source": [
    "arr = np.array([[[0, 1, 2, 3],\n",
    "                [4, 5, 6, 7],\n",
    "                [8, 9, 10, 11],\n",
    "                [12, 13, 14, 15]],\n",
    "                [[40, 41, 42, 43],\n",
    "                [44, 45, 46, 47],\n",
    "                [48, 49, 410, 411],\n",
    "                [412, 413, 414, 415]]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "2d176239-d299-4254-a737-68c74b9553e9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[  0,   1,   2,   3],\n",
       "        [  4,   5,   6,   7],\n",
       "        [  8,   9,  10,  11],\n",
       "        [ 12,  13,  14,  15]],\n",
       "\n",
       "       [[ 40,  41,  42,  43],\n",
       "        [ 44,  45,  46,  47],\n",
       "        [ 48,  49, 410, 411],\n",
       "        [412, 413, 414, 415]]])"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "arr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "016cb8c0-554c-47a8-862f-e1aeddc71b96",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 1  2  3  4]\n",
      " [ 5  6  7  8]\n",
      " [ 9 10 11 12]\n",
      " [13 14 15 16]]\n"
     ]
    }
   ],
   "source": [
    "arr_right = arr + 1\n",
    "print(arr_right)\n",
    "arr_right = arr_right[..., :-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "213efd3a-cda5-4b54-ad1b-98d419fb7a3b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[  0,   1,   2],\n",
       "        [  4,   5,   6],\n",
       "        [  8,   9,  10],\n",
       "        [ 12,  13,  14]],\n",
       "\n",
       "       [[ 40,  41,  42],\n",
       "        [ 44,  45,  46],\n",
       "        [ 48,  49, 410],\n",
       "        [412, 413, 414]]])"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "left = arr[..., :-1]\n",
    "left"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "4554ef71-ce1c-4a68-af10-42433fd5f13d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def grid_3d_connectivity_matrix(shape):\n",
    "    \n",
    "    x_dim, y_dim, z_dim = shape\n",
    "    \n",
    "    array = np.reshape(np.arange(x_dim * y_dim * z_dim), (x_dim, y_dim, z_dim))\n",
    "    \n",
    "    columnar = (-1, 1)\n",
    "    \n",
    "    left = np.reshape(array[..., :-1], columnar)\n",
    "    right = np.reshape(array[..., 1:], columnar)\n",
    "    up = np.reshape(array[:, :-1, :], columnar)\n",
    "    down = np.reshape(array[:, 1:, :], columnar)\n",
    "    sup = np.reshape(array[1:, ...], columnar)\n",
    "    inf = np.reshape(array[:-1, ...], columnar)\n",
    "    \n",
    "    connectivity = np.concatenate(\n",
    "        (np.hstack((left, right)),\n",
    "         np.hstack((right, left)),\n",
    "         np.hstack((up, down)),\n",
    "         np.hstack((down, up)),\n",
    "         np.hstack((sup, inf)),\n",
    "         np.hstack((inf, sup))),)\n",
    "    \n",
    "    return connectivity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "id": "4c67bb1a-428e-4c5b-903c-277f7ee43ef0",
   "metadata": {},
   "outputs": [],
   "source": [
    "c = grid_3d_connectivity_matrix((3, 3, 3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "id": "5f08b1f7-fb42-4da7-ab5a-db47cd5f2ba7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 0,  1],\n",
       "       [ 1,  2],\n",
       "       [ 3,  4],\n",
       "       [ 4,  5],\n",
       "       [ 6,  7],\n",
       "       [ 7,  8],\n",
       "       [ 9, 10],\n",
       "       [10, 11],\n",
       "       [12, 13],\n",
       "       [13, 14],\n",
       "       [15, 16],\n",
       "       [16, 17],\n",
       "       [18, 19],\n",
       "       [19, 20],\n",
       "       [21, 22],\n",
       "       [22, 23],\n",
       "       [24, 25],\n",
       "       [25, 26],\n",
       "       [ 1,  0],\n",
       "       [ 2,  1],\n",
       "       [ 4,  3],\n",
       "       [ 5,  4],\n",
       "       [ 7,  6],\n",
       "       [ 8,  7],\n",
       "       [10,  9],\n",
       "       [11, 10],\n",
       "       [13, 12],\n",
       "       [14, 13],\n",
       "       [16, 15],\n",
       "       [17, 16],\n",
       "       [19, 18],\n",
       "       [20, 19],\n",
       "       [22, 21],\n",
       "       [23, 22],\n",
       "       [25, 24],\n",
       "       [26, 25],\n",
       "       [ 0,  3],\n",
       "       [ 1,  4],\n",
       "       [ 2,  5],\n",
       "       [ 3,  6],\n",
       "       [ 4,  7],\n",
       "       [ 5,  8],\n",
       "       [ 9, 12],\n",
       "       [10, 13],\n",
       "       [11, 14],\n",
       "       [12, 15],\n",
       "       [13, 16],\n",
       "       [14, 17],\n",
       "       [18, 21],\n",
       "       [19, 22],\n",
       "       [20, 23],\n",
       "       [21, 24],\n",
       "       [22, 25],\n",
       "       [23, 26],\n",
       "       [ 3,  0],\n",
       "       [ 4,  1],\n",
       "       [ 5,  2],\n",
       "       [ 6,  3],\n",
       "       [ 7,  4],\n",
       "       [ 8,  5],\n",
       "       [12,  9],\n",
       "       [13, 10],\n",
       "       [14, 11],\n",
       "       [15, 12],\n",
       "       [16, 13],\n",
       "       [17, 14],\n",
       "       [21, 18],\n",
       "       [22, 19],\n",
       "       [23, 20],\n",
       "       [24, 21],\n",
       "       [25, 22],\n",
       "       [26, 23],\n",
       "       [ 9,  0],\n",
       "       [10,  1],\n",
       "       [11,  2],\n",
       "       [12,  3],\n",
       "       [13,  4],\n",
       "       [14,  5],\n",
       "       [15,  6],\n",
       "       [16,  7],\n",
       "       [17,  8],\n",
       "       [18,  9],\n",
       "       [19, 10],\n",
       "       [20, 11],\n",
       "       [21, 12],\n",
       "       [22, 13],\n",
       "       [23, 14],\n",
       "       [24, 15],\n",
       "       [25, 16],\n",
       "       [26, 17],\n",
       "       [ 0,  9],\n",
       "       [ 1, 10],\n",
       "       [ 2, 11],\n",
       "       [ 3, 12],\n",
       "       [ 4, 13],\n",
       "       [ 5, 14],\n",
       "       [ 6, 15],\n",
       "       [ 7, 16],\n",
       "       [ 8, 17],\n",
       "       [ 9, 18],\n",
       "       [10, 19],\n",
       "       [11, 20],\n",
       "       [12, 21],\n",
       "       [13, 22],\n",
       "       [14, 23],\n",
       "       [15, 24],\n",
       "       [16, 25],\n",
       "       [17, 26]])"
      ]
     },
     "execution_count": 78,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "03d8bc35-f922-4208-8dba-1d34dc2b79ba",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[  1,   2,   3],\n",
       "        [  5,   6,   7],\n",
       "        [  9,  10,  11],\n",
       "        [ 13,  14,  15]],\n",
       "\n",
       "       [[ 41,  42,  43],\n",
       "        [ 45,  46,  47],\n",
       "        [ 49, 410, 411],\n",
       "        [413, 414, 415]]])"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "right = arr[..., 1:]\n",
    "right"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "982579da-f520-4594-89b1-e5cb97ca1836",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[  0,   1,   2,   3],\n",
       "        [  4,   5,   6,   7],\n",
       "        [  8,   9,  10,  11]],\n",
       "\n",
       "       [[ 40,  41,  42,  43],\n",
       "        [ 44,  45,  46,  47],\n",
       "        [ 48,  49, 410, 411]]])"
      ]
     },
     "execution_count": 61,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "up = arr[:, :-1, :]\n",
    "up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "7282eac0-a645-4f1e-adbb-7be87be2fac5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[  4,   5,   6,   7],\n",
       "        [  8,   9,  10,  11],\n",
       "        [ 12,  13,  14,  15]],\n",
       "\n",
       "       [[ 44,  45,  46,  47],\n",
       "        [ 48,  49, 410, 411],\n",
       "        [412, 413, 414, 415]]])"
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "down = arr[:, 1:, :]\n",
    "down"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "1188be4e-ece9-46d0-b84d-33320b496eca",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[ 40,  41,  42,  43],\n",
       "        [ 44,  45,  46,  47],\n",
       "        [ 48,  49, 410, 411],\n",
       "        [412, 413, 414, 415]]])"
      ]
     },
     "execution_count": 67,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "z_up = arr[1:, ...]\n",
    "z_up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "18a7393e-4605-4c54-9795-da3c2a69358a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[ 0,  1,  2,  3],\n",
       "        [ 4,  5,  6,  7],\n",
       "        [ 8,  9, 10, 11],\n",
       "        [12, 13, 14, 15]]])"
      ]
     },
     "execution_count": 69,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "z_down = arr[:-1, ...]\n",
    "z_down"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "ee7bf312-b79e-4647-ab90-631463480e55",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 1],\n",
       "       [ 2],\n",
       "       [ 3],\n",
       "       [ 5],\n",
       "       [ 6],\n",
       "       [ 7],\n",
       "       [ 9],\n",
       "       [10],\n",
       "       [11],\n",
       "       [13],\n",
       "       [14],\n",
       "       [15]])"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "arr_right.reshape((-1, 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "38e0ac7f-b8f5-4a6d-83a0-ca35e06b0877",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 1,  2,  3],\n",
       "       [ 5,  6,  7],\n",
       "       [ 9, 10, 11],\n",
       "       [13, 14, 15]])"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "arr_right = arr[..., 1:]\n",
    "arr_right"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "8eadd4cd-6699-4294-8946-4ed5ac49a573",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 4,  5,  6,  7],\n",
       "       [ 8,  9, 10, 11],\n",
       "       [12, 13, 14, 15]])"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "arr_down = arr[1:, ...]\n",
    "arr_down"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "638c99d5-9f1a-4782-b3be-8181a20d6bbf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 0,  1,  2,  3],\n",
       "       [ 4,  5,  6,  7],\n",
       "       [ 8,  9, 10, 11]])"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "arr_up = arr[:-1, ...]\n",
    "arr_up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "240fa1b6-da56-4af9-9469-2a946fd0123e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 0,  1,  2],\n",
       "       [ 4,  5,  6],\n",
       "       [ 8,  9, 10],\n",
       "       [12, 13, 14]])"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "arr_left = arr[..., :-1]\n",
    "arr_left"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "d5f2bcea-db86-4fd1-a551-e9f55f00d316",
   "metadata": {},
   "outputs": [],
   "source": [
    "columnar = (-1, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "40c74d57-d3ea-4b47-b0f6-b0f546941ced",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 0],\n",
       "       [ 1],\n",
       "       [ 2],\n",
       "       [ 4],\n",
       "       [ 5],\n",
       "       [ 6],\n",
       "       [ 8],\n",
       "       [ 9],\n",
       "       [10],\n",
       "       [12],\n",
       "       [13],\n",
       "       [14]])"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "arr_left.reshape(columnar)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "a719296e-8747-4f5b-a0cf-93cda5ef73f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "connectivity = np.concatenate(\n",
    "    (np.hstack((arr_left.reshape(columnar), arr_right.reshape(columnar))),\n",
    "     np.hstack((arr_right.reshape(columnar), arr_left.reshape(columnar))),\n",
    "     np.hstack((arr_up.reshape(columnar), arr_down.reshape(columnar))),\n",
    "     np.hstack((arr_down.reshape(columnar), arr_up.reshape(columnar)))),)\n",
    "     "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "9f69dd38-7dc3-4899-939b-3ca0ace7b80b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(48, 2)"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "connectivity.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "ebe1a760-b9c7-47d3-a749-8dd21d97fde2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 0,  1],\n",
       "       [ 1,  2],\n",
       "       [ 2,  3],\n",
       "       [ 4,  5],\n",
       "       [ 5,  6],\n",
       "       [ 6,  7],\n",
       "       [ 8,  9],\n",
       "       [ 9, 10],\n",
       "       [10, 11],\n",
       "       [12, 13],\n",
       "       [13, 14],\n",
       "       [14, 15]])"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a = np.hstack((arr_left.reshape(columnar), arr_right.reshape(columnar)))\n",
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "3cae4dad-ea93-4df3-97fe-96c79cd7c02f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 0,  1],\n",
       "       [ 1,  2],\n",
       "       [ 2,  3],\n",
       "       [ 4,  5],\n",
       "       [ 5,  6],\n",
       "       [ 6,  7],\n",
       "       [ 8,  9],\n",
       "       [ 9, 10],\n",
       "       [10, 11],\n",
       "       [12, 13],\n",
       "       [13, 14],\n",
       "       [14, 15],\n",
       "       [ 0,  1],\n",
       "       [ 1,  2],\n",
       "       [ 2,  3],\n",
       "       [ 4,  5],\n",
       "       [ 5,  6],\n",
       "       [ 6,  7],\n",
       "       [ 8,  9],\n",
       "       [ 9, 10],\n",
       "       [10, 11],\n",
       "       [12, 13],\n",
       "       [13, 14],\n",
       "       [14, 15]])"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.concatenate((a, a))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5215ced7-0545-4e5c-a240-14e948f5468f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "python3",
   "name": "pytorch-gpu.1-9.m82",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/pytorch-gpu.1-9:m82"
  },
  "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.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}