np_array_over_websocket.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How to send numpy arrays over a websocket\n",
    "A friend approached me with a problem:\n",
    "> How can I send numpy arrays between two different python processes running two different versions on the same machine?\n",
    "\n",
    "Some kind of loopback socket makes sense. To reduce code complexity, websockets specifically make sense. This leaves the question of how to package the data on the websocket. The array could be packed and unpacked with custom `json` code or maybe custom `msgpack` code. A simpler solution is leveraging the `.npy` format that should exist between both versions of Python with their respective NumPy libraries.\n",
    "\n",
    "The below code shows how to pack and unpack numpy arrays that could be sent as a binary string using a websocket.\n",
    "\n",
    "This notebook does NOT show how to use the actual websockets to do the sending. Maybe that's for a future notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "metadata": {},
   "outputs": [],
   "source": [
    "import io\n",
    "from numpy.random import randint, standard_normal\n",
    "dtypes = ('float32', 'float64', 'int64', 'int32', 'int16')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test the send/receive IO logic\n",
    "1. Make fake array\n",
    "2. Pack fake array in to byte buffer using `io.BytesIO` and `np.save`\n",
    "    * Note: this byte buffer is something that could be send over a websocket\n",
    "3. Unpack byte buffer and populate fake array using `io.BytesIO` and `np.load`\n",
    "4. Confirm that the 'sent' array and 'received' array are equal. They should be bit equal."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 127,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test array 1 of 10\n",
      "dtype = float32\n",
      "shape = [3 1 1 2]\n",
      "byte_arr length = 152\n",
      "YAY! : shape + elements equal\n",
      "YAY! : dtype equal\n",
      "\n",
      "Test array 2 of 10\n",
      "dtype = float32\n",
      "shape = [2 4 2 2]\n",
      "byte_arr length = 256\n",
      "YAY! : shape + elements equal\n",
      "YAY! : dtype equal\n",
      "\n",
      "Test array 3 of 10\n",
      "dtype = int64\n",
      "shape = [4 2 3 2]\n",
      "byte_arr length = 512\n",
      "YAY! : shape + elements equal\n",
      "YAY! : dtype equal\n",
      "\n",
      "Test array 4 of 10\n",
      "dtype = float32\n",
      "shape = [4 4 3]\n",
      "byte_arr length = 320\n",
      "YAY! : shape + elements equal\n",
      "YAY! : dtype equal\n",
      "\n",
      "Test array 5 of 10\n",
      "dtype = int64\n",
      "shape = [3 1 2 3]\n",
      "byte_arr length = 272\n",
      "YAY! : shape + elements equal\n",
      "YAY! : dtype equal\n",
      "\n",
      "Test array 6 of 10\n",
      "dtype = int64\n",
      "shape = [2 1 3]\n",
      "byte_arr length = 176\n",
      "YAY! : shape + elements equal\n",
      "YAY! : dtype equal\n",
      "\n",
      "Test array 7 of 10\n",
      "dtype = int16\n",
      "shape = [2 2 1]\n",
      "byte_arr length = 136\n",
      "YAY! : shape + elements equal\n",
      "YAY! : dtype equal\n",
      "\n",
      "Test array 8 of 10\n",
      "dtype = int16\n",
      "shape = [4 3]\n",
      "byte_arr length = 152\n",
      "YAY! : shape + elements equal\n",
      "YAY! : dtype equal\n",
      "\n",
      "Test array 9 of 10\n",
      "dtype = int64\n",
      "shape = [4 2 1 2]\n",
      "byte_arr length = 256\n",
      "YAY! : shape + elements equal\n",
      "YAY! : dtype equal\n",
      "\n",
      "Test array 10 of 10\n",
      "dtype = int64\n",
      "shape = [2]\n",
      "byte_arr length = 144\n",
      "YAY! : shape + elements equal\n",
      "YAY! : dtype equal\n",
      "\n"
     ]
    }
   ],
   "source": [
    "num_test_arrays = 10\n",
    "for ix in range(num_test_arrays):\n",
    "    print(\"Test array {} of {}\".format(ix+1, num_test_arrays))\n",
    "    \n",
    "    s = randint(1, 5, size=randint(1,5))\n",
    "    d = dtypes[randint(0, len(dtypes))]\n",
    "\n",
    "    asend = None\n",
    "    if d.startswith('float'):\n",
    "        asend = standard_normal(s).astype(d)\n",
    "    elif d.startswith('int'):\n",
    "        asend = randint(0, high=2**15, size=s, dtype=d)\n",
    "\n",
    "    print(\"dtype = {}\".format(d))\n",
    "    print(\"shape = {}\".format(s))\n",
    "    # print(asend)\n",
    "\n",
    "    # sending side\n",
    "    fake_file_write = io.BytesIO()\n",
    "    np.save(fake_file_write, asend, allow_pickle=False, fix_imports=False)\n",
    "    byte_arr = fake_file_write.getvalue()\n",
    "    print(\"byte_arr length = {}\".format(len(byte_arr)))\n",
    "    \n",
    "    # receiving side\n",
    "    fake_file_read = io.BytesIO(byte_arr)\n",
    "    arecv = np.load(fake_file_read, allow_pickle=False, fix_imports=False)\n",
    "    \n",
    "    # test shape and elements with np.array_equal\n",
    "    if asend.shape == arecv.shape:\n",
    "        print('YAY! : shape + elements equal')\n",
    "    else:\n",
    "        print('OH NO!!!!!! : shape + elements ***********NOT*********** EQUAL')\n",
    "\n",
    "    # test if dtype is equal\n",
    "    if asend.dtype == arecv.dtype:\n",
    "        print('YAY! : dtype equal')\n",
    "    else:\n",
    "        print('OH NO!!!!!! : dtype ***********NOT*********** EQUAL')\n",
    "    print()"
   ]
  },
  {
   "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.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}