Encoding and decoding in CKKS

encoding_decoding_ckks.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here for the rest of the notebook we choose to keep building upon the `CKKSEncoder` class we have defined earlier. Instead of redefining the class each time we want to add or change methods, we will simply use `patch_to` from the `fastcore` package from [Fastai](https://github.com/fastai/fastai). This allows to monkey patch objects that have already been defined. This is purely for conveniency, and you could just redefine the `CKKSEncoder` at each cell with the added methods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-23T17:09:55.999586Z",
     "start_time": "2020-05-23T17:09:55.617615Z"
    }
   },
   "outputs": [],
   "source": [
    "# !pip3 install fastcore\n",
    "\n",
    "from fastcore.foundation import patch_to"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-23T17:09:56.009762Z",
     "start_time": "2020-05-23T17:09:56.001929Z"
    }
   },
   "outputs": [],
   "source": [
    "@patch_to(CKKSEncoder)\n",
    "def pi(self, z: np.array) -> np.array:\n",
    "    \"\"\"Projects a vector of H into C^{N/2}.\"\"\"\n",
    "    \n",
    "    N = self.M // 4\n",
    "    return z[:N]\n",
    "\n",
    "@patch_to(CKKSEncoder)\n",
    "def pi_inverse(self, z: np.array) -> np.array:\n",
    "    \"\"\"Expands a vector of C^{N/2} by expanding it with its\n",
    "    complex conjugate.\"\"\"\n",
    "    \n",
    "    z_conjugate = z[::-1]\n",
    "    z_conjugate = [np.conjugate(x) for x in z_conjugate]\n",
    "    return np.concatenate([z, z_conjugate])\n",
    "\n",
    "# We can now initialize our encoder with the added methods\n",
    "encoder = CKKSEncoder(M)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-23T17:09:56.026175Z",
     "start_time": "2020-05-23T17:09:56.014037Z"
    }
   },
   "outputs": [],
   "source": [
    "z = np.array([0,1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-23T17:09:56.042695Z",
     "start_time": "2020-05-23T17:09:56.029673Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0, 1, 1, 0])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "encoder.pi_inverse(z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-23T17:09:56.055281Z",
     "start_time": "2020-05-23T17:09:56.046534Z"
    }
   },
   "outputs": [],
   "source": [
    "@patch_to(CKKSEncoder)\n",
    "def create_sigma_R_basis(self):\n",
    "    \"\"\"Creates the basis (sigma(1), sigma(X), ..., sigma(X** N-1)).\"\"\"\n",
    "\n",
    "    self.sigma_R_basis = np.array(self.vandermonde(self.xi, self.M)).T\n",
    "    \n",
    "@patch_to(CKKSEncoder)\n",
    "def __init__(self, M):\n",
    "    \"\"\"Initialize with the basis\"\"\"\n",
    "    self.xi = np.exp(2 * np.pi * 1j / M)\n",
    "    self.M = M\n",
    "    self.create_sigma_R_basis()\n",
    "    \n",
    "encoder = CKKSEncoder(M)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now have a look at the basis $\\sigma(1), \\sigma(X), \\sigma(X^2), \\sigma(X^3)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-23T17:09:56.075264Z",
     "start_time": "2020-05-23T17:09:56.059839Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 1.00000000e+00+0.j        ,  1.00000000e+00+0.j        ,\n",
       "         1.00000000e+00+0.j        ,  1.00000000e+00+0.j        ],\n",
       "       [ 7.07106781e-01+0.70710678j, -7.07106781e-01+0.70710678j,\n",
       "        -7.07106781e-01-0.70710678j,  7.07106781e-01-0.70710678j],\n",
       "       [ 2.22044605e-16+1.j        , -4.44089210e-16-1.j        ,\n",
       "         1.11022302e-15+1.j        , -1.38777878e-15-1.j        ],\n",
       "       [-7.07106781e-01+0.70710678j,  7.07106781e-01+0.70710678j,\n",
       "         7.07106781e-01-0.70710678j, -7.07106781e-01-0.70710678j]])"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "encoder.sigma_R_basis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we will check that elements of $\\mathbb{Z} \\{ \\sigma(1), \\sigma(X), \\sigma(X^2), \\sigma(X^3) \\}$ are encoded as integer polynomials."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-23T17:09:56.090570Z",
     "start_time": "2020-05-23T17:09:56.078736Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1.+2.41421356j, 1.+0.41421356j, 1.-0.41421356j, 1.-2.41421356j])"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Here we simply take a vector whose coordinates are (1,1,1,1) in the lattice basis\n",
    "coordinates = [1,1,1,1]\n",
    "\n",
    "b = np.matmul(encoder.sigma_R_basis.T, coordinates)\n",
    "b"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can check now that it does encode to an integer polynomial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-23T17:09:56.106987Z",
     "start_time": "2020-05-23T17:09:56.095166Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$x \\mapsto \\text{(1+2.220446049250313e-16j)} + (\\text{(1+0j)})\\,x + (\\text{(0.9999999999999998+2.7755575615628716e-17j)})\\,x^{2} + (\\text{(1+2.220446049250313e-16j)})\\,x^{3}$"
      ],
      "text/plain": [
       "Polynomial([1.+2.22044605e-16j, 1.+0.00000000e+00j, 1.+2.77555756e-17j,\n",
       "       1.+2.22044605e-16j], domain=[-1,  1], window=[-1,  1])"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "p = encoder.sigma_inverse(b)\n",
    "p"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-23T17:09:56.122465Z",
     "start_time": "2020-05-23T17:09:56.111492Z"
    }
   },
   "outputs": [],
   "source": [
    "@patch_to(CKKSEncoder)\n",
    "def compute_basis_coordinates(self, z):\n",
    "    \"\"\"Computes the coordinates of a vector with respect to the orthogonal lattice basis.\"\"\"\n",
    "    output = np.array([np.real(np.vdot(z, b) / np.vdot(b,b)) for b in self.sigma_R_basis])\n",
    "    return output\n",
    "\n",
    "def round_coordinates(coordinates):\n",
    "    \"\"\"Gives the integral rest.\"\"\"\n",
    "    coordinates = coordinates - np.floor(coordinates)\n",
    "    return coordinates\n",
    "\n",
    "def coordinate_wise_random_rounding(coordinates):\n",
    "    \"\"\"Rounds coordinates randonmly.\"\"\"\n",
    "    r = round_coordinates(coordinates)\n",
    "    f = np.array([np.random.choice([c, c-1], 1, p=[1-c, c]) for c in r]).reshape(-1)\n",
    "    \n",
    "    rounded_coordinates = coordinates - f\n",
    "    rounded_coordinates = [int(coeff) for coeff in rounded_coordinates]\n",
    "    return rounded_coordinates\n",
    "\n",
    "@patch_to(CKKSEncoder)\n",
    "def sigma_R_discretization(self, z):\n",
    "    \"\"\"Projects a vector on the lattice using coordinate wise random rounding.\"\"\"\n",
    "    coordinates = self.compute_basis_coordinates(z)\n",
    "    \n",
    "    rounded_coordinates = coordinate_wise_random_rounding(coordinates)\n",
    "    y = np.matmul(self.sigma_R_basis.T, rounded_coordinates)\n",
    "    return y\n",
    "\n",
    "encoder = CKKSEncoder(M)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, because there might be loss of precisions during the rounding step, we had the scale parameter $\\delta$, to allow a fixed level of precision."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-23T17:09:56.137266Z",
     "start_time": "2020-05-23T17:09:56.127876Z"
    }
   },
   "outputs": [],
   "source": [
    "@patch_to(CKKSEncoder)\n",
    "def __init__(self, M:int, scale:float):\n",
    "    \"\"\"Initializes with scale.\"\"\"\n",
    "    self.xi = np.exp(2 * np.pi * 1j / M)\n",
    "    self.M = M\n",
    "    self.create_sigma_R_basis()\n",
    "    self.scale = scale\n",
    "    \n",
    "@patch_to(CKKSEncoder)\n",
    "def encode(self, z: np.array) -> Polynomial:\n",
    "    \"\"\"Encodes a vector by expanding it first to H,\n",
    "    scale it, project it on the lattice of sigma(R), and performs\n",
    "    sigma inverse.\n",
    "    \"\"\"\n",
    "    pi_z = self.pi_inverse(z)\n",
    "    scaled_pi_z = self.scale * pi_z\n",
    "    rounded_scale_pi_zi = self.sigma_R_discretization(scaled_pi_z)\n",
    "    p = self.sigma_inverse(rounded_scale_pi_zi)\n",
    "    \n",
    "    # We round it afterwards due to numerical imprecision\n",
    "    coef = np.round(np.real(p.coef)).astype(int)\n",
    "    p = Polynomial(coef)\n",
    "    return p\n",
    "\n",
    "@patch_to(CKKSEncoder)\n",
    "def decode(self, p: Polynomial) -> np.array:\n",
    "    \"\"\"Decodes a polynomial by removing the scale, \n",
    "    evaluating on the roots, and project it on C^(N/2)\"\"\"\n",
    "    rescaled_p = p / self.scale\n",
    "    z = self.sigma(rescaled_p)\n",
    "    pi_z = self.pi(z)\n",
    "    return pi_z\n",
    "\n",
    "scale = 64\n",
    "\n",
    "encoder = CKKSEncoder(M, scale)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now see it on action, the full encoder used by CKKS : "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-23T17:09:56.154616Z",
     "start_time": "2020-05-23T17:09:56.142612Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([3.+4.j, 2.-1.j])"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "z = np.array([3 +4j, 2 - 1j])\n",
    "z"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that we now have an integer polynomial as our encoding."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-23T17:09:56.174750Z",
     "start_time": "2020-05-23T17:09:56.158397Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$x \\mapsto \\text{160.0} + \\text{90.0}\\,x + \\text{160.0}\\,x^{2} + \\text{45.0}\\,x^{3}$"
      ],
      "text/plain": [
       "Polynomial([160.,  90., 160.,  45.], domain=[-1,  1], window=[-1,  1])"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "p = encoder.encode(z)\n",
    "p"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And it actually decodes well ! "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-23T17:09:56.205155Z",
     "start_time": "2020-05-23T17:09:56.191067Z"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([2.99718446+3.99155337j, 2.00281554-1.00844663j])"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "encoder.decode(p)"
   ]
  }
 ],
 "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.8.2"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}