From ae5be5e966f227d56cdf4b84f6202dffc36f139f Mon Sep 17 00:00:00 2001 From: Arthur Lu Date: Wed, 13 Nov 2024 22:44:15 +0000 Subject: [PATCH] basic hyperbolic distance implementation using FHE in python --- basic.ipynb | 241 ++++++++++++++++++++++++++++++++++++++++++++++++++ venv_setup.sh | 6 ++ 2 files changed, 247 insertions(+) create mode 100644 basic.ipynb create mode 100755 venv_setup.sh diff --git a/basic.ipynb b/basic.ipynb new file mode 100644 index 0000000..682f8a4 --- /dev/null +++ b/basic.ipynb @@ -0,0 +1,241 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Client Setup" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "from Pyfhel import Pyfhel, PyCtxt" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[Client] Initializing Pyfhel session and data...\n" + ] + } + ], + "source": [ + "print(f\"[Client] Initializing Pyfhel session and data...\")\n", + "HE_client = Pyfhel() # Creating empty Pyfhel object\n", + "ckks_params = {\n", + " 'scheme': 'CKKS', # can also be 'ckks'\n", + " 'n': 2**14, # Polynomial modulus degree. For CKKS, n/2 values can be\n", + " # encoded in a single ciphertext. \n", + " # Typ. 2^D for D in [10, 15]\n", + " 'scale': 2**30, # All the encodings will use it for float->fixed point\n", + " # conversion: x_fix = round(x_float * scale)\n", + " # You can use this as default scale or use a different\n", + " # scale on each operation (set in HE.encryptFrac)\n", + " 'qi_sizes': [60, 30, 30, 30, 60] # Number of bits of each prime in the chain. \n", + " # Intermediate values should be close to log2(scale)\n", + " # for each operation, to have small rounding errors.\n", + "}\n", + "HE_client.contextGen(**ckks_params) # Generate context for bfv scheme\n", + "HE_client.keyGen() # Generates both a public and a private key\n", + "HE_client.relinKeyGen()\n", + "HE_client.rotateKeyGen()" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [], + "source": [ + "# Generate and encrypt query vector\n", + "x = np.array([1.5, 2, 3.3, 4])\n", + "cx = np.array([HE_client.encrypt(x[j]) for j in range(len(x))])" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[Client] sending HE_client= and cx=[\n", + " \n", + " \n", + " ]\n" + ] + } + ], + "source": [ + "# Serializing data and public context information\n", + "\n", + "s_context = HE_client.to_bytes_context()\n", + "s_public_key = HE_client.to_bytes_public_key()\n", + "s_relin_key = HE_client.to_bytes_relin_key()\n", + "s_rotate_key = HE_client.to_bytes_rotate_key()\n", + "s_cx = [cx[j].to_bytes() for j in range(len(cx))]\n", + "\n", + "print(f\"[Client] sending HE_client={HE_client} and cx={cx}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Server Mock" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [], + "source": [ + "def hyperbolic_distance_parts(u, v): # returns only the numerator and denominator of the hyperbolic distance formula\n", + " diff = u - v\n", + " du = -(1 - u @ u) # for some reason we need to negate this\n", + " dv = -(1 - v @ v) # for some reason we need to negate this\n", + " return diff @ diff, du * dv # returns the numerator and denominator\n" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [], + "source": [ + "# document matrix containing rows of document encoding vectors\n", + "D = [\n", + " [0.5, -1.5, 4, 5],\n", + " [1.0, 1.5, 4, 5]\n", + "]" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[Server] received HE_server= and cx=[\n", + " \n", + " \n", + " ]\n", + "[Server] Distances computed! Responding: res=[(, ), (, )]\n" + ] + } + ], + "source": [ + "HE_server = Pyfhel()\n", + "HE_server.from_bytes_context(s_context)\n", + "HE_server.from_bytes_public_key(s_public_key)\n", + "HE_server.from_bytes_relin_key(s_relin_key)\n", + "HE_server.from_bytes_rotate_key(s_rotate_key)\n", + "cx = np.array([PyCtxt(pyfhel=HE_server, bytestring=s_cx[j]) for j in range(len(s_cx))])\n", + "print(f\"[Server] received HE_server={HE_server} and cx={cx}\")\n", + "\n", + "# Encode each document weights in plaintext\n", + "res = []\n", + "\n", + "for i in range(len(D)):\n", + " d = np.array(D[i])\n", + " cd = np.array([HE_server.encrypt(d[j]) for j in range(len(d))])\n", + " # Compute distance bewteen recieved query and D[i]\n", + " res.append(hyperbolic_distance_parts(cx, cd))\n", + "\n", + "s_res = [(res[j][0].to_bytes(), res[j][1].to_bytes()) for j in range(len(res))]\n", + "\n", + "print(f\"[Server] Distances computed! Responding: res={res}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Client Parse Response" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [], + "source": [ + "def hyperbolic_distance(u, v):\n", + " num = ((u - v) @ (u - v))\n", + " den = (1 - (u @ u)) * (1 - (v @ v))\n", + " return np.arccosh(1 + 2 * (num / den))" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[Client] Response received! \n", + "Result is [np.float64(0.20736836029574815), np.float64(0.07564769989949309)] \n", + "Should be [np.float64(0.20738785895993414), np.float64(0.07565488467449914)]\n", + "Diff [1.94986642e-05 7.18477501e-06]\n" + ] + } + ], + "source": [ + "#res = np.array([HE_client.decrypt(c_res[j]) for j in range(len(c_res))])[:,0]\n", + "#res = HE_client.decrypt(c_res)\n", + "c_res = []\n", + "for i in range(len(s_res)):\n", + " c_num = PyCtxt(pyfhel=HE_server, bytestring=s_res[i][0])\n", + " c_den = PyCtxt(pyfhel=HE_server, bytestring=s_res[i][1])\n", + " p_num = HE_client.decrypt(c_num)[0]\n", + " p_den = HE_client.decrypt(c_den)[0]\n", + " dist = np.arccosh(1 + 2 * (p_num / p_den))\n", + " c_res.append(dist)\n", + "\n", + "# Checking result\n", + "expected = [hyperbolic_distance(x, np.array(w)) for w in D]\n", + "print(f\"[Client] Response received! \\nResult is {c_res} \\nShould be {expected}\\nDiff {np.abs(np.array(c_res) - np.array(expected))}\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "cs239", + "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.9.20" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/venv_setup.sh b/venv_setup.sh new file mode 100755 index 0000000..0d46621 --- /dev/null +++ b/venv_setup.sh @@ -0,0 +1,6 @@ +apt install -y python3.9-venv python3.9-dev +mkdir ~/.venvs +rm -rf ~/.venvs/cs239 +python3.9 -m venv ~/.venvs/cs239 +source ~/.venvs/cs239/bin/activate +pip install pyfhel \ No newline at end of file