Sampling_Using_Cholesky.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "## Question 2\n",
        "Implement from scratch a sampling method to draw samples from a multivariate Normal (MVN) distribution in JAX. [10 Marks]\n",
        "\n",
        "- Your code should work for any number of dimensions but please set the number of dimensions (random variables of MVN) to 10 for this task.\n",
        "- You are only allowed to use jax.random.uniform. You are especially not allowed to use jax.random.normal.\n",
        "- You should randomly create the mean and covariance matrix to fully specify an MVN distribution.\n",
        "- Implement a sampling method from scratch using which you can draw samples from the specified MVN distribution.\n",
        "- Use your sampling method to draw multiple samples from the MVN distribution and reconstruct the parameters of your MVN distribution (mean and covariance matrix) to confirm that your sampling method is working correctly.\n",
        "\n",
        "\n",
        "---\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "G9kxe0djUEnZ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Approach Used\n",
        "\n",
        "To generate samples $\\mathbf{x} \\sim \\mathcal{N}(\\mathbf{m}, K)$ with arbitrary mean $\\mathbf{m}$ and covariance matrix $K$ using a scalar Gaussian generator (which is readily available in many programming environments) we proceed as follows: first, compute the Cholesky decomposition (also known as the \"matrix square root\") $L$ of the positive definite symmetric covariance matrix $K=L L^{\\top}$, where $L$ is a lower triangular matrix. Then generate $\\mathbf{u} \\sim \\mathcal{N}(\\mathbf{0}, I)$ by multiple separate calls to the scalar Gaussian generator. Compute $\\mathbf{x}=\\mathbf{m}+L \\mathbf{u}$, which has the desired distribution with mean $\\mathrm{m}$ and covariance $L \\mathbb{E}\\left[\\mathbf{u u}^{\\top}\\right] L^{\\top}=L L^{\\top}=K$ (by the independence of the elements of $\\mathbf{u}$ ).\n",
        "\n",
        "In practice it may be necessary to add a small multiple of the identity matrix $\\epsilon I$ to the covariance matrix for numerical reasons. This is because the eigenvalues of the matrix $K$ can decay very rapidly and without this stabilization the Cholesky decomposition fails. The effect on the generated samples is to add additional independent noise of variance $\\epsilon$. From the context $\\epsilon$ can usually be chosen to have inconsequential effects on the samples, while ensuring numerical stability.\n",
        "\n",
        "Refrence : <a href=\"http://gaussianprocess.org/gpml/chapters/RWA.pdf\">Gaussian Process for Machine Learning | Appendix A | Gaussian Identities</a>\n",
        "\n"
      ],
      "metadata": {
        "id": "u46ZFU2wUZPb"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "UNrWHq361BvT"
      },
      "outputs": [],
      "source": [
        "import jax.numpy as jnp\n",
        "from jax.numpy import linalg as JLA\n",
        "import jax"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "boltlLCoztiS"
      },
      "source": [
        "## Declare Dimension, Means , Covariance Matrix\n",
        "\n",
        "\n",
        "Let $M$ be an $n \\times n$ Hermitian matrix (this includes real symmetric matrices).$M$ is positive definite if and only if all of its eigenvalues are positive.\n",
        "\n",
        "Refrence : <a href=\"https://en.wikipedia.org/wiki/\">Wikipedia | Definite matrix</a>\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "geGiwyg20nFn",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "6264f127-bcbd-4efc-fd27-1848600d2934"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
          ]
        }
      ],
      "source": [
        "# # of Random Variables\n",
        "d = 10\n",
        "\n",
        "# epsilon to stabalize cholesky\n",
        "epsilon = 0.1\n",
        "\n",
        "# Radom Mean\n",
        "mean = jax.random.uniform(jax.random.PRNGKey(0), shape=(d, 1), minval=0, maxval=1, dtype=jnp.float32)\n",
        "\n",
        "# Random matrix for covariance\n",
        "mat = jax.random.uniform(jax.random.PRNGKey(2), shape=(d, d), minval=0, maxval=1, dtype=jnp.float32)\n",
        "\n",
        "# Make a symmetric matrix\n",
        "sym_mat = (mat+mat.T)/2\n",
        "\n",
        "# Make matrix positive definite\n",
        "cov = jnp.dot(sym_mat, sym_mat.T)\n",
        "\n",
        "# Add identity matrix for statbility\n",
        "cov = cov + epsilon*jnp.identity(d)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BUJTPcdeAFP7"
      },
      "source": [
        "### Verify if Covariance Matrix is Positive Definite"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "c7YVQ8D5-Vwg"
      },
      "outputs": [],
      "source": [
        "# Calculate Eigen values,all of them must be positive\n",
        "cov_eignvals = JLA.eigvals(cov)\n",
        "if(jnp.any(cov_eignvals < 0)):\n",
        "    print(\"The covariance matrix is not positive definite, generate again\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Cholesky decomposition\n",
        "\n",
        "In linear algebra, the Cholesky decomposition or Cholesky factorization is a decomposition of a Hermitian, positive-definite matrix into the product of a lower triangular matrix and its conjugate transpose.\n",
        "\n",
        "Refrences:\n",
        "\n",
        "- <a href=\"https://www.youtube.com/watch?v=xloCwiVDkho&list=PLYdroRCLMg5MgczmIkeY_XVJiJ5LVDFuh&index=44\">Linear Algebra: Cholesky Decomposition</a>\n",
        "\n",
        "- <a href=\"https://www.youtube.com/watch?v=NppyUqgQqd0&list=PLYdroRCLMg5MgczmIkeY_XVJiJ5LVDFuh&index=44\">Linear Algebra: Cholesky Decomposition Example</a>\n",
        "\n",
        "- <a href=\"https://ocw.mit.edu/courses/10-34-numerical-methods-applied-to-chemical-engineering-fall-2005/resources/lecturenotes142/\"> Cholesky MIT OCW</a>"
      ],
      "metadata": {
        "id": "oxCYVcD_JYHa"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "g09stjT_KVwZ"
      },
      "outputs": [],
      "source": [
        "def Chol_decomp(mat):\n",
        "    '''\n",
        "        Perform Cholesky decomposition of input matrix and return L i.e\n",
        "        lower triangular matrix\n",
        "\n",
        "        Args:\n",
        "            mat - (d, d) JAX Numpy array containing covariance matrix\n",
        "                  (d is the number of random variables/dimensions)\n",
        "        Returns:\n",
        "            L - (d, d) JAX Numpy array containing the lower trainagular matrix\n",
        "    '''\n",
        "\n",
        "    shape = cov.shape[0]\n",
        "\n",
        "    # Declare a dxd zero matrix\n",
        "    L = jnp.zeros((d, d))\n",
        "\n",
        "    # Loop over covariance matrix\n",
        "    for i in range(shape):\n",
        "        for j in range(i + 1):\n",
        "            sum = 0\n",
        "            # For Variance values\n",
        "            if (j == i):\n",
        "                for k in range(j):\n",
        "                    sum += (L[j][k])**2\n",
        "                L = L.at[i, j].set(jnp.sqrt(mat[i][j] - sum))\n",
        "            # For covariace values\n",
        "            else:\n",
        "                for k in range(j):\n",
        "                    sum += L[i][k] * L[j][k]\n",
        "                if(L[j][j] > 0):\n",
        "                    L = L.at[i, j].set((mat[i][j] - sum) / (L[j][j]))\n",
        "    return L\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8AV03SbD05lG"
      },
      "source": [
        "### Validate Cholsky Decomposition\n",
        "\n",
        "**NOTE**: If the Decomposition did not work correctly.You might want to check the following:\n",
        "- If you have rounded of the values while comparing covariance matrix $K$ with $L L^{\\top}$.\n",
        "- If rounding is not the issue, you might want to increase the value of the epsilon, this will ensure that Cholesky decomposition works correctly. Increasing the value of epsilon will increase the variance since we add $\\epsilon I$ to the covariance matrix $K$ to stabilize cholesky hence you must also increase the sample_size while calling the gen_std_MVN(<a href=\"https://ocw.mit.edu/courses/res-6-012-introduction-to-probability-spring-2018/resources/the-weak-law-of-large-numbers/\"> Weak Law of Large Numbers </a>) to ensure that the covariance and mean obtained after sampling are close to the original values. Example values `d=30,epsilon=10,sample_size=100000`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "NSffipzBLImQ",
        "outputId": "232a587e-897a-48f4-e0d9-d858315b191d"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Decomposition worked Correctly\n"
          ]
        }
      ],
      "source": [
        "L = Chol_decomp(cov)\n",
        "\n",
        "# Verify if the dot product of L and L transpose return covariance matrix\n",
        "decomp_val = jnp.round(cov, 4) == jnp.round(jnp.dot(L, L.T), 4)\n",
        "if(jnp.any(decomp_val == False)):\n",
        "    print(\"Seems like the Decomposition did not work correctly.\")\n",
        "else:\n",
        "    print(\"Decomposition worked Correctly\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Generating Standard Normal Using CLT\n",
        "\n",
        "Suppose $X_{1}, X_{2}, \\ldots, X_{n}$ are independent random variables with the same underlying distribution. In this case, we say that the $X_{i}$ are independent and identically-distributed, or i.i.d. In particular, the $X_{i}$ all have the same mean $\\mu$ and standard deviation $\\sigma$.\n",
        "Let $\\bar{X}_{n}$ be the average of $X_{1}, \\ldots, X_{n}$ :\n",
        "$$\n",
        "\\bar{X}_{n}=\\frac{X_{1}+X_{2}+\\cdots+X_{n}}{n}=\\frac{1}{n} \\sum_{i=1}^{n} X_{i}\n",
        "$$\n",
        "Note that $\\bar{X}_{n}$ is itself a random variable. The law of large numbers and central limit theorem tell us about the value and distribution of $\\bar{X}_{n}$, respectively.\n",
        "\n",
        "CLT: As $n$ grows, the distribution of $\\bar{X}_{n}$ converges to the normal distribution $N\\left(\\mu, \\sigma^{2} / n\\right)$.\n",
        "\n",
        "Refrences:\n",
        "\n",
        "- <a href=\"https://ocw.mit.edu/courses/18-05-introduction-to-probability-and-statistics-spring-2014/resources/mit18_05s14_reading6b/\">MIT OCW | Introduction to Probability and Statistics</a>\n",
        "\n",
        "- <a href=\"https://ocw.mit.edu/courses/6-041-probabilistic-systems-analysis-and-applied-probability-fall-2010/resources/mit6_041f10_l19/\"> MIT OCW | Probabilistic Systems Analysis and Applied Probability</a>"
      ],
      "metadata": {
        "id": "l5DNecuzQ-3-"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "kp93bGPA4Ubq"
      },
      "outputs": [],
      "source": [
        "def gen_std_MVN(sample_size=10000, d=10):\n",
        "    '''\n",
        "        Returns a standard Multivariate gaussian based on the dimesions/\n",
        "        number of random variables and sample size. This is based on the\n",
        "        idea of Central Limit theorem\n",
        "\n",
        "        Args:\n",
        "            d -  integer | the number of random variables\n",
        "            sample_size - integer | # of samples to be added to generate a gaussian\n",
        "\n",
        "        Returns:\n",
        "            std_MVN - (d, sample_size) JAX Numpy array containing the standard Multivariate\n",
        "                       gaussian matrix\n",
        "    '''\n",
        "\n",
        "    std_MVN = jnp.zeros((d, sample_size))\n",
        "\n",
        "    for i in range(d):\n",
        "\n",
        "        # Draw samples from Uniform Distribution\n",
        "        unf = jax.random.uniform(jax.random.PRNGKey(i), shape=(sample_size, d), minval=0, maxval=100, dtype=jnp.float32)\n",
        "        # Add them to create a scalar normal distribution\n",
        "        normal = jnp.sum(unf, axis=1)\n",
        "        # Covert to standard Normal\n",
        "        std_normal = (normal - jnp.mean(normal)) / jnp.sqrt(jnp.var(normal))\n",
        "        std_MVN = std_MVN.at[i, :].set(std_normal)\n",
        "\n",
        "    return std_MVN\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "std_MVN= gen_std_MVN(d=d)\n",
        "samples = mean + jnp.dot(L, std_MVN)"
      ],
      "metadata": {
        "id": "RXH41e6VdSs-"
      },
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Validating the results"
      ],
      "metadata": {
        "id": "hi_roZQ4FE_h"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "absolut_diff_mean = jnp.sum(jnp.abs(jnp.mean(samples, axis=1) - mean.squeeze()))\n",
        "absolut_diff_cov = jnp.sum(jnp.abs(jnp.cov(samples) - cov))"
      ],
      "metadata": {
        "id": "lT-CTHJnF-4E"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(f\"The absolute difference between all the values of the original mean vs the mean obtained from sampling is {absolut_diff_mean}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "SLr310JLFlA1",
        "outputId": "2176caae-0f06-4ca5-9bbf-444d4750e2f3"
      },
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "The absolute difference between all the values of the original mean vs the mean obtained from sampling is 1.6391277313232422e-06\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(f\"The average of absolute difference between all the values of the original covariance vs the covariance obtained from sampling is {absolut_diff_cov/(d*d)}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1xXz5jDCGnZU",
        "outputId": "45443b9b-a094-42c0-8cbb-0fd80bdfd74d"
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "The average of absolute difference between all the values of the original covariance vs the covariance obtained from sampling is 0.012451081536710262\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(\"Original Mean \\n\",mean.squeeze())\n",
        "print(\"\\n\")\n",
        "print(\"Sample Mean \\n\",jnp.mean(samples,axis=1))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2fhUc0UXIDzO",
        "outputId": "5092395e-489e-4f6c-f7ed-53ea146ebbb1"
      },
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Original Mean \n",
            " [0.35490513 0.60419905 0.4275843  0.23061597 0.32985854 0.43953657\n",
            " 0.25099766 0.27730572 0.7678207  0.71474564]\n",
            "\n",
            "\n",
            "Sample Mean \n",
            " [0.35490483 0.60419923 0.42758447 0.230616   0.32985836 0.43953642\n",
            " 0.2509974  0.27730566 0.76782095 0.7147457 ]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(\"Original Covariance Matrix \\n\",cov)\n",
        "print(\"\\n\")\n",
        "print(\"Sample Covariance Matrix \\n\",jnp.cov(samples))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6E9ZubE1Ifm_",
        "outputId": "8d1ed96f-e954-4989-d86b-9204eb0bee6f"
      },
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Original Covariance Matrix \n",
            " [[2.8252301 2.163516  1.8826021 2.1715758 2.6237981 2.8255095 2.7453787\n",
            "  2.6538415 1.9432937 2.3962338]\n",
            " [2.163516  2.6559076 2.0175126 2.2711544 2.4282575 2.3687296 2.4298837\n",
            "  2.4429803 1.9467937 2.2675002]\n",
            " [1.8826021 2.0175126 2.778164  2.2047598 2.5771775 2.5989413 2.5783846\n",
            "  2.4079242 1.9919269 2.4645398]\n",
            " [2.1715758 2.2711544 2.2047598 2.5554626 2.573891  2.616433  2.574296\n",
            "  2.430259  1.9815549 2.3292205]\n",
            " [2.6237981 2.4282575 2.5771775 2.573891  3.2553551 3.0281265 3.1051197\n",
            "  2.992762  2.156797  2.6898634]\n",
            " [2.8255095 2.3687296 2.5989413 2.616433  3.0281265 3.5964012 3.1709008\n",
            "  2.9781177 2.277362  2.7988725]\n",
            " [2.7453787 2.4298837 2.5783846 2.574296  3.1051197 3.1709008 3.5501003\n",
            "  3.177404  2.230322  3.0325356]\n",
            " [2.6538415 2.4429803 2.4079242 2.430259  2.992762  2.9781177 3.177404\n",
            "  3.294318  2.2810152 2.6716766]\n",
            " [1.9432937 1.9467937 1.9919269 1.9815549 2.156797  2.277362  2.230322\n",
            "  2.2810152 2.1277618 2.0248706]\n",
            " [2.3962338 2.2675002 2.4645398 2.3292205 2.6898634 2.7988725 3.0325356\n",
            "  2.6716766 2.0248706 3.1371326]]\n",
            "\n",
            "\n",
            "Sample Covariance Matrix \n",
            " [[2.8255126 2.1478322 1.8859898 2.155817  2.6087596 2.8374686 2.7429662\n",
            "  2.638953  1.9417595 2.400996 ]\n",
            " [2.1478322 2.631821  2.0086052 2.253092  2.4016716 2.3637989 2.4153802\n",
            "  2.420702  1.9404138 2.24824  ]\n",
            " [1.8859898 2.0086052 2.7815945 2.1957161 2.5691214 2.6150327 2.5777519\n",
            "  2.4030457 1.9932505 2.46446  ]\n",
            " [2.155817  2.253092  2.1957161 2.5359015 2.546936  2.6112163 2.556969\n",
            "  2.4057417 1.969562  2.3109517]\n",
            " [2.6087596 2.4016716 2.5691214 2.546936  3.2260442 3.0257413 3.0783837\n",
            "  2.9606335 2.141022  2.6657379]\n",
            " [2.8374686 2.3637989 2.6150327 2.6112163 3.0257413 3.62479   3.1854417\n",
            "  2.978844  2.288485  2.812861 ]\n",
            " [2.7429662 2.4153802 2.5777519 2.556969  3.0783837 3.1854417 3.5384793\n",
            "  3.1551738 2.2344491 3.0316358]\n",
            " [2.638953  2.420702  2.4030457 2.4057417 2.9606335 2.978844  3.1551738\n",
            "  3.2641642 2.2736814 2.6553135]\n",
            " [1.9417595 1.9404138 1.9932505 1.969562  2.141022  2.288485  2.2344491\n",
            "  2.2736814 2.1278377 2.0231688]\n",
            " [2.400996  2.24824   2.46446   2.3109517 2.6657379 2.812861  3.0316358\n",
            "  2.6553135 2.0231688 3.1455264]]\n"
          ]
        }
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "Sampling Using Cholesky.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}