LKJ.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "import pyro\n",
    "import pyro.distributions as dist\n",
    "from pyro.infer.mcmc import NUTS, MCMC\n",
    "\n",
    "pyro.enable_validation()\n",
    "pyro.set_rng_seed(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Sample: 100%|██████████| 300/300 [00:10<00:00, 26.56it/s, step size=4.35e-01, acc. rate=0.950]\n"
     ]
    }
   ],
   "source": [
    "def model(y):\n",
    "    d = y.shape[1]\n",
    "    N = y.shape[0]\n",
    "    # Vector of variances for each of the d variables\n",
    "    theta = pyro.sample(\"theta\", dist.HalfCauchy(torch.full((d,), 1, dtype=y.dtype)))\n",
    "    # Lower cholesky factor of a correlation matrix\n",
    "    eta = y.new_full((), 1)\n",
    "    L_omega = pyro.sample(\"L_omega\", dist.LKJCorrCholesky(d, eta))\n",
    "    # Lower cholesky factor of the covariance matrix\n",
    "    L_Omega = torch.mm(torch.diag(theta.sqrt()), L_omega)\n",
    "\n",
    "    # Vector of expectations\n",
    "    mu = y.new_zeros(d)\n",
    "\n",
    "    with pyro.plate(\"observations\", N):\n",
    "        obs = pyro.sample(\"obs\", dist.MultivariateNormal(mu, scale_tril=L_Omega),\n",
    "                          obs=y)\n",
    "    return obs\n",
    "\n",
    "\n",
    "y = torch.randn(500, 5)\n",
    "nuts_kernel = NUTS(model)\n",
    "mcmc = MCMC(nuts_kernel, num_samples=200, warmup_steps=100).run(y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "support = mcmc.marginal([\"theta\", \"L_omega\"]).support(flatten=True)\n",
    "theta_mean = support[\"theta\"].mean(dim=0)\n",
    "L_omega_mean = support[\"L_omega\"].mean(dim=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.9573, 0.9238, 1.1333, 1.0575, 1.0755])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "theta_mean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 1.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
       "        [-0.0180,  0.9988,  0.0000,  0.0000,  0.0000],\n",
       "        [-0.0226,  0.0146,  0.9972,  0.0000,  0.0000],\n",
       "        [-0.0077,  0.0518, -0.0322,  0.9953,  0.0000],\n",
       "        [-0.0596, -0.0175,  0.0340, -0.0224,  0.9936]])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "L_omega_mean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.9164, -0.0159, -0.0245, -0.0078, -0.0613],\n",
       "        [-0.0159,  0.8516,  0.0157,  0.0507, -0.0163],\n",
       "        [-0.0245,  0.0157,  1.2782, -0.0374,  0.0426],\n",
       "        [-0.0078,  0.0507, -0.0374,  1.1121, -0.0271],\n",
       "        [-0.0613, -0.0163,  0.0426, -0.0271,  1.1484]])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "scale_tril = (theta_mean.unsqueeze(1) * L_omega_mean)\n",
    "covariance = scale_tril.matmul(scale_tril.t())\n",
    "covariance"
   ]
  }
 ],
 "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}