GMM_EM_vs_PPEM.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/HanaHasan04/PPEM-for-GMM/blob/main/GMM_EM_vs_PPEM.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# **Evaluating Privacy-Preserving vs. Non-Privacy-Preserving Expectation Maximization (PPEM vs. EM) for Gaussian Mixture Models (GMM)**"
      ],
      "metadata": {
        "id": "hn-41OG8_L0a"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Given $n$ participants (= $n$ data points), and a Gaussian Mixture Model (GMM) with $c$ Gaussian components, where each component $j$ is characterized by its mean $\\mu_{j}$, covariance $\\Sigma_{j}$, and coefficient $\\beta_{j}$. The log-likelihood of the GMM is defined by:  \n",
        " $$\\mathcal{L}=\\sum_{i=1}^{n}\\log \\left(\\sum_{j=1}^{c}\\beta_{j} \\cdot \\mathcal{N}\\left(\\mathbf{x}_{i} | \\mu_{j}, \\Sigma_{j}\\right)\\right)$$\n",
        "\n",
        "At each iteration, the log-likelihood is calculated under the current estimates of the GMM parameters, and the log-likelihood of each method is plotted in order to compare them."
      ],
      "metadata": {
        "id": "52EmyZa34zsF"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **Imports**"
      ],
      "metadata": {
        "id": "u32vE8sq3AOU"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JnR_TACI4EtM"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "from numpy import random\n",
        "import matplotlib.pyplot as plt\n",
        "from matplotlib.patches import Ellipse\n",
        "import matplotlib.transforms as transforms\n",
        "from scipy.stats import multivariate_normal\n",
        "import copy\n",
        "!pip install tenseal\n",
        "import tenseal as ts"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **Classes**"
      ],
      "metadata": {
        "id": "SH-m3FxC2xd8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class KMS:\n",
        "  \"\"\"\n",
        "  This class represents a Key Management Service (KMS).\n",
        "  Our KMS is used to generate, distribute, and manage cryptographic keys for the CKKS scheme via TenSEAL library.\n",
        "\n",
        "  Attributes:\n",
        "  -----------\n",
        "  poly_modulus_degree:\n",
        "      The degree of the polynomial modulus used in the CKKS scheme.\n",
        "\n",
        "  coeff_mod_bit_sizes:\n",
        "      The bit sizes of the coefficient modulus used in the CKKS scheme.\n",
        "\n",
        "  context: TenSEALContext\n",
        "      The parameters of the CKKS scheme in TenSEAL.\n",
        "\n",
        "  clients:\n",
        "      The list of clients (parties).\n",
        "  \"\"\"\n",
        "\n",
        "  def __init__(self, clients, poly_modulus_degree = 8192, coeff_mod_bit_sizes = [60, 40, 40, 60], global_scale = 2 ** 40):\n",
        "    self.poly_modulus_degree = poly_modulus_degree\n",
        "    self.coeff_mod_bit_sizes = coeff_mod_bit_sizes\n",
        "    self.global_scale = global_scale\n",
        "    self.context = None\n",
        "    self.clients = clients\n",
        "\n",
        "  def gen_context(self):\n",
        "    \"\"\"Generates a TenSEALContext.\"\"\"\n",
        "    self.context = ts.context(ts.SCHEME_TYPE.CKKS, poly_modulus_degree=self.poly_modulus_degree, coeff_mod_bit_sizes=self.coeff_mod_bit_sizes)\n",
        "    self.context.generate_galois_keys()\n",
        "    self.context.global_scale = self.global_scale\n",
        "\n",
        "  def dist_context(self):\n",
        "    \"\"\"Distributes the current context to the list of clients.\"\"\"\n",
        "    for client in self.clients:\n",
        "      client.context = self.context"
      ],
      "metadata": {
        "id": "hxoOfN4s4-HA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class Server:\n",
        "  \"\"\"\n",
        "  This class represents the server - an untrusted third party.\n",
        "\n",
        "  Attributes:\n",
        "  -----------\n",
        "  num_of_clients:\n",
        "      The total number of clients (data points).\n",
        "\n",
        "  vectors:\n",
        "      The vectors sent by clients.\n",
        "  \"\"\"\n",
        "\n",
        "  def __init__(self, num_of_clients):\n",
        "    self.num_of_clients = num_of_clients\n",
        "    self.vectors = []\n",
        "\n",
        "  def add_vec(self, vector):\n",
        "    \"\"\"Adds a vector to the server's list of vectors.\"\"\"\n",
        "    self.vectors.append(vector)\n",
        "\n",
        "  def calc_sum_enc(self):\n",
        "    \"\"\"Calculates the sum of the encrypted vectors stored on the server.\"\"\"\n",
        "    if len(self.vectors) != self.num_of_clients:\n",
        "      raise ValueError(\"The number of vectors stored on the server does not match the number of clients.\")\n",
        "    return sum(self.vectors)\n",
        "\n",
        "  def calc_sum(self):\n",
        "    \"\"\"Calculates the sum of the vectors stored on the server.\"\"\"\n",
        "    if len(self.vectors) != self.num_of_clients:\n",
        "      raise ValueError(\"The number of vectors stored on the server does not match the number of clients.\")\n",
        "    return [sum(x) for x in zip(*self.vectors)]\n",
        "\n",
        "  def clear_server(self):\n",
        "    \"\"\"Clears the server's list of vectors.\"\"\"\n",
        "    self.vectors = []"
      ],
      "metadata": {
        "id": "3JlO3dyV5BoF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class Client:\n",
        "  \"\"\"\n",
        "  This class represents a single party (a client) that holds a data point belonging to a Gaussian Mixture Model (GMM). \n",
        "  The data point is represented as a two-dimensional point.\n",
        "\n",
        "  Attributes:\n",
        "  -----------\n",
        "  data: \n",
        "      The two-dimensional data point.\n",
        "\n",
        "  num_of_gaussians:\n",
        "      Number of gaussian components.\n",
        "\n",
        "  context: TenSEALContext\n",
        "      The parameters of the CKKS scheme in TenSEAL.\n",
        "  \n",
        "  a:\n",
        "      The intermediate updates for the coefficients of the GMM. size: num_of_gaussians.\n",
        "\n",
        "  b: \n",
        "      The intermediate updates for the means of the GMM. size: (num_of_gaussians, 2).\n",
        "\n",
        "  c: \n",
        "      The intermediate updates for the covariances of the GMM. size: (num_of_gaussians, 2, 2).\n",
        "  \"\"\"\n",
        "\n",
        "  def __init__(self, data, num_of_gaussians):\n",
        "    self.data = data\n",
        "    self.num_of_gaussians = num_of_gaussians\n",
        "    self.context = None \n",
        "    # for intermediate updates\n",
        "    self.a = np.zeros(num_of_gaussians)      \n",
        "    self.b = np.zeros((num_of_gaussians, 2))\n",
        "    self.c = np.zeros((num_of_gaussians, 2, 2))\n",
        "\n",
        "  def e_step(self, gmm):\n",
        "    \"\"\"Performs the E-Step of the EM algorithm. Computes the conditional probabilities that the client's data belongs to each Gaussian model.\"\"\"\n",
        "    self.a[:] = [gmm.coefficients[j] * multivariate_normal.pdf(self.data, mean=gmm.means[j], cov=gmm.covariances[j]) for j in range(self.num_of_gaussians)]\n",
        "    self.a[:] /= self.a[:].sum()\n",
        "\n",
        "  def inter_m_step(self, j, gmm):\n",
        "    \"\"\"Performs the intermediate M-Step calcultaions of the EM algorithm. Computes the local updates for the coefficients, means, and covariances of the GMM.\"\"\"\n",
        "    a_j = self.a[j]\n",
        "    b_j_0 = self.a[j] * self.data[0]\n",
        "    b_j_1 = self.a[j] * self.data[1]\n",
        "    c_j_00 = self.a[j] * (self.data[0] - gmm.means[j][0])**2\n",
        "    c_j_01 = self.a[j] * (self.data[0] - gmm.means[j][0]) * (self.data[1] - gmm.means[j][1])\n",
        "    c_j_10 = self.a[j] * (self.data[1] - gmm.means[j][1]) * (self.data[0] - gmm.means[j][0])\n",
        "    c_j_11 = self.a[j] * (self.data[1] - gmm.means[j][1])**2\n",
        "\n",
        "    v_j = [a_j, b_j_0, b_j_1, c_j_00, c_j_01, c_j_10, c_j_11] \n",
        "    return v_j\n",
        "\n",
        "  def pp_inter_m_step(self, j, gmm, context):\n",
        "    \"\"\"Performs the privacy-preserving intermediate M-Step calcultaions of the EM algorithm. Computes the local updates for the coefficients, means, and covariances of the GMM.\"\"\"\n",
        "    a_j = self.a[j]\n",
        "    b_j_0 = self.a[j] * self.data[0]\n",
        "    b_j_1 = self.a[j] * self.data[1]\n",
        "    c_j_00 = self.a[j] * (self.data[0] - gmm.means[j][0])**2\n",
        "    c_j_01 = self.a[j] * (self.data[0] - gmm.means[j][0]) * (self.data[1] - gmm.means[j][1])\n",
        "    c_j_10 = self.a[j] * (self.data[1] - gmm.means[j][1]) * (self.data[0] - gmm.means[j][0])\n",
        "    c_j_11 = self.a[j] * (self.data[1] - gmm.means[j][1])**2\n",
        "\n",
        "    v_j = [a_j, b_j_0, b_j_1, c_j_00, c_j_01, c_j_10, c_j_11]\n",
        "    enc_v_j = self.encrypt(v_j, context)  \n",
        "    return enc_v_j\n",
        "\n",
        "  def encrypt(self, vec, context):\n",
        "    \"\"\"Encrypt a vector using the CKKS scheme in TenSEAL.\"\"\"\n",
        "    return ts.ckks_vector(context, vec)\n",
        "\n",
        "  def decrypt(self, vec):\n",
        "    \"\"\"Decrypt a CKKSVector using the CKKS scheme in TenSEAL.\"\"\"\n",
        "    return vec.decrypt()"
      ],
      "metadata": {
        "id": "bAfvgcwZ5IS1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class GMM:\n",
        "  \"\"\"\n",
        "  This class represents a Gaussian Mixture Model (GMM). We wish to fit a GMM to the full dataset of all parties.\n",
        "\n",
        "  Attributes:\n",
        "  -----------\n",
        "  parties:\n",
        "      The parties (clients) that the data is distributed among them.\n",
        "\n",
        "  num_of_parties:\n",
        "      Number of parties (which is the number of data points).\n",
        "\n",
        "  server: Server\n",
        "      The server (untrusted third party) to perform computations on the encrypted data.\n",
        "    \n",
        "  kms: KMS\n",
        "      The key management system that is used to generate, distribute, and manage cryptographic keys.\n",
        "\n",
        "  num_of_gaussians: \n",
        "      Number of gaussian components.\n",
        "\n",
        "  colors_of_gaussians: \n",
        "      Colors of the gaussian clusters for poltting. size: (num_of_gaussians, 3).\n",
        "\n",
        "  means:\n",
        "      The means of the Gaussian components. size: (num_of_gaussians, 2).\n",
        "      \n",
        "  covariances:\n",
        "      The covariances of the Gaussian components. size: (num_of_gaussians, 2, 2).\n",
        "      \n",
        "  coefficients:\n",
        "      The coefficients of the Gaussian components. size: num_of_gaussians. \n",
        "  \"\"\"\n",
        "\n",
        "  def __init__(self, parties, num_of_parties, num_of_gaussians):\n",
        "    self.parties = parties\n",
        "    self.num_of_parties = num_of_parties\n",
        "    self.server = Server(num_of_parties)\n",
        "    self.kms = KMS(parties, poly_modulus_degree = 8192, coeff_mod_bit_sizes = [60, 40, 40, 60], global_scale = 2 ** 40)\n",
        "    self.num_of_gaussians = num_of_gaussians\n",
        "    self.colors_of_gaussians = random.rand(num_of_gaussians, 3)\n",
        "    # parameters estimates\n",
        "    self.means = random.rand(num_of_gaussians, 2)*20 - 10     # initial means: random from uniform[-10, 10] \n",
        "    covariances = np.zeros((num_of_gaussians, 2, 2))\n",
        "    for j in range(num_of_gaussians):\n",
        "        covariances[j] = np.eye(2)\n",
        "    self.covariances = covariances\n",
        "    self.coefficients = np.ones(num_of_gaussians)/num_of_gaussians\n",
        "\n",
        "  def e_step(self):\n",
        "    \"\"\"E-step of EM algorithm.\"\"\"\n",
        "    for party in self.parties:\n",
        "      party.e_step(self)\n",
        "\n",
        "  def pp_m_step(self):\n",
        "    \"\"\"privacy-preserving M-step of EM algorithm.\"\"\"\n",
        "    self.kms.gen_context()\n",
        "    self.kms.dist_context()\n",
        "\n",
        "    for j in range(self.num_of_gaussians):\n",
        "      for party in self.parties:\n",
        "        enc_v_ij = party.pp_inter_m_step(j, self, self.kms.context)\n",
        "        self.server.add_vec(enc_v_ij)\n",
        "\n",
        "      enc_sum_j = self.server.calc_sum_enc()\n",
        "      self.server.clear_server()\n",
        "\n",
        "      sum_j = self.parties[random.randint(0, self.num_of_parties - 1)].decrypt(enc_sum_j)\n",
        "      a_j = sum_j[0]; b_j_0 = sum_j[1]; b_j_1 = sum_j[2]; c_j_00 = sum_j[3]; c_j_01 = sum_j[4]; c_j_10 = sum_j[5]; c_j_11 = sum_j[6];  \n",
        "      self.coefficients[j] = a_j / self.num_of_parties\n",
        "      self.means[j][0] = b_j_0 / a_j\n",
        "      self.means[j][1] = b_j_1 / a_j\n",
        "      self.covariances[j][0][0] = c_j_00 / a_j\n",
        "      self.covariances[j][0][1] = c_j_01 / a_j\n",
        "      self.covariances[j][1][0] = c_j_10 / a_j\n",
        "      self.covariances[j][1][1] = c_j_11 / a_j    \n",
        "\n",
        "  def m_step(self):\n",
        "    \"\"\"M-step of EM algorithm.\"\"\"\n",
        "    for j in range(self.num_of_gaussians):\n",
        "      for party in self.parties:\n",
        "        v_ij = party.inter_m_step(j, self)\n",
        "        self.server.add_vec(v_ij)\n",
        "\n",
        "      sum_j = self.server.calc_sum()\n",
        "      self.server.clear_server()\n",
        "\n",
        "      a_j = sum_j[0]; b_j_0 = sum_j[1]; b_j_1 = sum_j[2]; c_j_00 = sum_j[3]; c_j_01 = sum_j[4]; c_j_10 = sum_j[5]; c_j_11 = sum_j[6];  \n",
        "      self.coefficients[j] = a_j / self.num_of_parties\n",
        "      self.means[j][0] = b_j_0 / a_j\n",
        "      self.means[j][1] = b_j_1 / a_j\n",
        "      self.covariances[j][0][0] = c_j_00 / a_j\n",
        "      self.covariances[j][0][1] = c_j_01 / a_j\n",
        "      self.covariances[j][1][0] = c_j_10 / a_j\n",
        "      self.covariances[j][1][1] = c_j_11 / a_j    \n",
        "\n",
        "  def plot_gaussian(self, mean, cov, ax, n_std=3.0, facecolor='none', **kwargs):\n",
        "    \"\"\"Utility function to plot one Gaussian from mean and covariance.\"\"\"\n",
        "    pearson = cov[0, 1]/np.sqrt(cov[0, 0] * cov[1, 1])\n",
        "    ell_radius_x = np.sqrt(1 + pearson)\n",
        "    ell_radius_y = np.sqrt(1 - pearson)\n",
        "    ellipse = Ellipse((0, 0), width=ell_radius_x * 2, height=ell_radius_y * 2, facecolor=facecolor, **kwargs)\n",
        "    scale_x = np.sqrt(cov[0, 0]) * n_std\n",
        "    mean_x = mean[0]\n",
        "    scale_y = np.sqrt(cov[1, 1]) * n_std\n",
        "    mean_y = mean[1]\n",
        "    transf = transforms.Affine2D().rotate_deg(45).scale(scale_x, scale_y).translate(mean_x, mean_y)\n",
        "    ellipse.set_transform(transf + ax.transData)\n",
        "    return ax.add_patch(ellipse)\n",
        "\n",
        "  def draw(self, ax, n_std=2.0, facecolor='none', **kwargs):\n",
        "    \"\"\"Function to draw the Gaussians.\"\"\"\n",
        "    for j in range(self.num_of_gaussians):\n",
        "        self.plot_gaussian(self.means[j], self.covariances[j], ax, n_std=n_std, edgecolor=self.colors_of_gaussians[j], **kwargs)\n",
        "\n",
        "  def log_likelihood(self, X):\n",
        "    \"\"\"Computes the log-likelihood of X under current parameters.\"\"\"\n",
        "    ll = []\n",
        "    for d in X:\n",
        "        tot = 0\n",
        "        for j in range(self.num_of_gaussians):\n",
        "            tot += self.coefficients[j] * multivariate_normal.pdf(d, mean=self.means[j], cov=self.covariances[j])\n",
        "        ll.append(np.log(tot))\n",
        "    return np.sum(ll)"
      ],
      "metadata": {
        "id": "Ud1DJZeY5SSM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **Utility Functions**"
      ],
      "metadata": {
        "id": "uMtTIDWi1LPt"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def gen_data(num_of_gaussians=3, points_per_gaussian=200, mean_range=[-10, 10]):\n",
        "  \"\"\"\n",
        "  Generates 2D points from a mixture of Gaussians.\n",
        "\n",
        "  Args:\n",
        "  - num_of_gaussians: Number of Gaussian components\n",
        "  - points_per_gaussian: Number of points in each Gaussian component\n",
        "  - mean_range: Range of mean values\n",
        "\n",
        "  Returns:\n",
        "  2D points data: Generated 2D points from a mixture of Gaussians\n",
        "  num_of_points: Number of the generated points = Number of Gaussians * Number of points in each component\n",
        "  \"\"\"\n",
        "  x = []\n",
        "  mean = random.rand(num_of_gaussians, 2)*(mean_range[1]-mean_range[0]) + mean_range[0]\n",
        "  for i in range(num_of_gaussians):\n",
        "      cov = random.rand(2, 12)\n",
        "      cov = np.matmul(cov, cov.T)\n",
        "      _x = np.random.multivariate_normal(mean[i], cov, points_per_gaussian)\n",
        "      x += list(_x)\n",
        "  x = np.array(x)\n",
        "  fig = plt.figure()\n",
        "  ax = fig.gca()\n",
        "  ax.scatter(x[:,0], x[:,1], s=3, alpha=0.4)\n",
        "  ax.autoscale(enable=True) \n",
        "  tot_points = num_of_gaussians*points_per_gaussian\n",
        "  return x, tot_points\n",
        "\n",
        "def plot(title):\n",
        "    \"\"\"Draw the data points and the fitted mixture model.\"\"\"\n",
        "    fig = plt.figure(figsize=(8, 8))\n",
        "    ax = fig.gca()\n",
        "    ax.scatter(X[:, 0], X[:, 1], s=3, alpha=0.4)\n",
        "    ax.scatter(gmm.means[:, 0], gmm.means[:, 1], c=gmm.colors_of_gaussians)\n",
        "    gmm.draw(ax, lw=3)\n",
        "    plt.title(title)\n",
        "    plt.show()\n",
        "    plt.clf()"
      ],
      "metadata": {
        "id": "_ic4Qw7N7bUk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **Applying EM and PPEM for a GMM**"
      ],
      "metadata": {
        "id": "8eTOI8r91Wuc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# TODO: change the parameters of the GMM\n",
        "NUM_OF_CLUSTERS = 3\n",
        "POINTS_PER_CLUSTER = 1000\n",
        "MEAN_RANGE = [-10, 10]\n",
        "X, NUM_OF_POINTS = gen_data(num_of_gaussians=NUM_OF_CLUSTERS, points_per_gaussian=POINTS_PER_CLUSTER, mean_range=MEAN_RANGE)\n",
        "\n",
        "clients = []\n",
        "for i in range(NUM_OF_POINTS):\n",
        "  x = Client(data=X[i], num_of_gaussians=NUM_OF_CLUSTERS)\n",
        "  clients.append(x)\n",
        "\n",
        "gmm = GMM(parties=clients, num_of_parties=NUM_OF_POINTS, num_of_gaussians=NUM_OF_CLUSTERS)\n",
        "gmm2 = copy.deepcopy(gmm)\n",
        "\n",
        "NUM_OF_ITERATIONS = 20    # TODO: change the number of iterations\n",
        "\n",
        "# non privacy-preserving\n",
        "ll = []\n",
        "for e in range(NUM_OF_ITERATIONS):\n",
        "    gmm.e_step()\n",
        "    gmm.m_step()\n",
        "    ll.append(gmm.log_likelihood(X))\n",
        "\n",
        "# privacy-preserving\n",
        "ll2 = []\n",
        "for e in range(NUM_OF_ITERATIONS):\n",
        "    gmm2.e_step()\n",
        "    gmm2.pp_m_step()\n",
        "    ll2.append(gmm2.log_likelihood(X))\n",
        "\n",
        "# plot the log-likelihood for both methods\n",
        "min_ll = np.min([np.min(ll), np.min(ll2)])\n",
        "max_ll = np.max([np.max(ll), np.max(ll2)])\n",
        "plt.plot(ll2, label='privacy-preserving', color='#00FF00', linewidth=5)\n",
        "plt.plot(ll, label='not privacy-preserving', color='#FF3131', linewidth=2.5, linestyle='dotted')\n",
        "plt.xlabel('Iteration')\n",
        "plt.ylabel('log-likelihood')\n",
        "plt.legend()\n",
        "plt.ylim(min_ll-500, max_ll+800)\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "tELGdY5j1Ftq"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}