{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Biostat 257 Homework 3\n",
    "\n",
    "**Due ~~May 6~~ May 13 @ 11:59PM**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "versioninfo()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Consider a linear mixed effects model\n",
    "$$\n",
    "    \\mathbf{Y}_i = \\mathbf{X}_i \\boldsymbol{\\beta} + \\mathbf{Z}_i \\boldsymbol{\\gamma} + \\boldsymbol{\\epsilon}_i, \\quad i=1,\\ldots,n,\n",
    "$$\n",
    "where   \n",
    "- $\\mathbf{Y}_i \\in \\mathbb{R}^{n_i}$ is the response vector of $i$-th individual,  \n",
    "- $\\mathbf{X}_i \\in \\mathbb{R}^{n_i \\times p}$ is the fixed effect predictor matrix of $i$-th individual,  \n",
    "- $\\mathbf{Z}_i \\in \\mathbb{R}^{n_i \\times q}$ is the random effect predictor matrix of $i$-th individual,  \n",
    "- $\\boldsymbol{\\epsilon}_i \\in \\mathbb{R}^{n_i}$ are multivariate normal $N(\\mathbf{0}_{n_i},\\sigma^2 \\mathbf{I}_{n_i})$,  \n",
    "- $\\boldsymbol{\\beta} \\in \\mathbb{R}^p$ are fixed effects, and  \n",
    "- $\\boldsymbol{\\gamma} \\in \\mathbb{R}^q$ are random effects assumed to be $N(\\mathbf{0}_q, \\boldsymbol{\\Sigma}_{q \\times q}$) independent of $\\boldsymbol{\\epsilon}_i$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Q1 Formula (10 pts)\n",
    "\n",
    "Write down the log-likelihood of the $i$-th datum $(\\mathbf{y}_i, \\mathbf{X}_i, \\mathbf{Z}_i)$ given parameters $(\\boldsymbol{\\beta}, \\boldsymbol{\\Sigma}, \\sigma^2)$. \n",
    "\n",
    "**Hint:** For non-statisticians, feel free to ask for help in class or office hours. Point of this exercise is computing not statistics."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Q2 Start-up code\n",
    "\n",
    "Use the following template to define a type `LmmObs` that holds an LMM datum $(\\mathbf{y}_i, \\mathbf{X}_i, \\mathbf{Z}_i)$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define a type that holds LMM datum\n",
    "struct LmmObs{T <: AbstractFloat}\n",
    "    # data\n",
    "    y :: Vector{T}\n",
    "    X :: Matrix{T}\n",
    "    Z :: Matrix{T}\n",
    "    # working arrays\n",
    "    # whatever intermediate vectors/arrays you may want to pre-allocate\n",
    "    storage_p  :: Vector{T}\n",
    "    storage_q  :: Vector{T}\n",
    "    xtx        :: Matrix{T}\n",
    "    ztx        :: Matrix{T}\n",
    "    ztz        :: Matrix{T}\n",
    "    storage_qq :: Matrix{T}\n",
    "end\n",
    "\n",
    "# constructor\n",
    "function LmmObs(\n",
    "        y::Vector{T}, \n",
    "        X::Matrix{T}, \n",
    "        Z::Matrix{T}\n",
    "        ) where T <: AbstractFloat\n",
    "    storage_p  = Vector{T}(undef, size(X, 2))\n",
    "    storage_q  = Vector{T}(undef, size(Z, 2))\n",
    "    xtx        = transpose(X) * X\n",
    "    ztx        = transpose(Z) * X\n",
    "    ztz        = transpose(Z) * Z\n",
    "    storage_qq = similar(ztz)\n",
    "    LmmObs(y, X, Z, storage_p, storage_q, xtx, ztx, ztz, storage_qq)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Write a function, with interface   \n",
    "```julia\n",
    "logl!(obs, β, L, σ²)\n",
    "```\n",
    "that evaluates the log-likelihood of the $i$-th datum. Here `L` is the lower triangular Cholesky factor from the Cholesky decomposition `Σ=LL'`. Make your code efficient in the $n_i \\gg q$ case. Think the intensive longitudinal measurement setting.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function logl!(\n",
    "        obs :: LmmObs{T}, \n",
    "        β   :: Vector{T}, \n",
    "        L   :: Matrix{T}, \n",
    "        σ²  :: T) where T <: AbstractFloat\n",
    "    n, p, q = size(obs.X, 1), size(obs.X, 2), size(obs.Z, 2)    \n",
    "    # TODO: compute and return the log-likelihood\n",
    "    obs.storage_qq = 0\n",
    "    sleep(1e-3) # wait 1 ms as if your code takes 1ms\n",
    "    return zero(T)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Hint**: This function shouldn't be very long. Mine, obeying 80-character rule, is 25 lines. If you find yourself writing very long code, you're on the wrong track. Think about algorithm first then use BLAS functions to reduce memory allocations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Q3 Correctness (15 pts)\n",
    "\n",
    "Compare your result (both accuracy and timing) to the [Distributions.jl](https://juliastats.org/Distributions.jl/stable/multivariate/#Distributions.AbstractMvNormal) package using following data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using BenchmarkTools, Distributions, LinearAlgebra, Random\n",
    "\n",
    "Random.seed!(257)\n",
    "# dimension\n",
    "n, p, q = 2000, 5, 3\n",
    "# predictors\n",
    "X  = [ones(n) randn(n, p - 1)]\n",
    "Z  = [ones(n) randn(n, q - 1)]\n",
    "# parameter values\n",
    "β  = [2.0; -1.0; rand(p - 2)]\n",
    "σ² = 1.5\n",
    "Σ  = fill(0.1, q, q) + 0.9I\n",
    "# generate y\n",
    "y  = X * β + Z * rand(MvNormal(Σ)) + sqrt(σ²) * randn(n)\n",
    "\n",
    "# form an LmmObs object\n",
    "obs = LmmObs(y, X, Z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is the standard way to evaluate log-density of a multivariate normal, using the Distributions.jl package. Let's evaluate the log-likelihood of this datum."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "μ  = X * β\n",
    "Ω  = Z * Σ * transpose(Z) +  σ² * I\n",
    "mvn = MvNormal(μ, Symmetric(Ω)) # MVN(μ, Σ)\n",
    "logpdf(mvn, y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check that your answer matches that from Distributions.jl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "L = Matrix(cholesky(Σ).L)\n",
    "logl!(obs, β, L, σ²)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**You will lose all 15 + 30 + 30 = 75 points** if the following statement throws `AssertionError`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@assert logl!(obs, β, Matrix(cholesky(Σ).L), σ²) ≈ logpdf(mvn, y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Q4 Efficiency (30 pts)\n",
    "\n",
    "Benchmarking your code and compare to the Distributions.jl function `logpdf`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# benchmark the `logpdf` function in Distribution.jl\n",
    "bm1 = @benchmark logpdf($mvn, $y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# benchmark your implementation\n",
    "L = Matrix(cholesky(Σ).L)\n",
    "bm2 = @benchmark logl!($obs, $β, $L, $σ²)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The points you will get is\n",
    "$$\n",
    "\\frac{x}{1000} \\times 30,\n",
    "$$\n",
    "where $x$ is the speedup of your program against the standard method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this is the points you'll get\n",
    "clamp(median(bm1).time / median(bm2).time / 1000 * 30, 0, 30)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Hint**: Apparently I am using 1000 as denominator because I expect your code to be at least $1000 \\times$ faster than the standard method."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Q5 Memory (30 pts)\n",
    "\n",
    "You want to avoid memory allocation in the \"hot\" function `logl!`. You will lose 1 point for each `1 KiB = 1024 bytes` memory allocation. In other words, the points you get for this question is"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "clamp(30 - median(bm2).memory / 1024, 0, 30)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Hint**: I am able to reduce the memory allocation to 0 bytes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Q6 Misc (15 pts)\n",
    "\n",
    "Coding style, Git workflow, etc. For reproducibity, make sure we (TA and myself) can run your Jupyter Notebook. That is how we grade Q4 and Q5. If we cannot run it, you will get zero points."
   ]
  }
 ],
 "metadata": {
  "@webio": {
   "lastCommId": null,
   "lastKernelId": null
  },
  "kernelspec": {
   "display_name": "Julia 1.7.1",
   "language": "julia",
   "name": "julia-1.7"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.7.1"
  },
  "toc": {
   "colors": {
    "hover_highlight": "#DAA520",
    "running_highlight": "#FF0000",
    "selected_highlight": "#FFD700"
   },
   "moveMenuLeft": true,
   "nav_menu": {
    "height": "87px",
    "width": "252px"
   },
   "navigate_menu": true,
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": true,
   "threshold": 4,
   "toc_cell": false,
   "toc_section_display": "block",
   "toc_window_display": false,
   "widenNotebook": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}