Statistics in Julia - Maximum Likelihood Estimation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Intended Audience\n",
    "\n",
    "This notebook is intended to be read by people with an interest in statistical computing that already have:\n",
    "* An intermediate knowledge of statistics. For example, the reader should already understand most of the material in a book like [\"All of Statistics\"](https://www.stat.cmu.edu/~larry/all-of-statistics/).\n",
    "* An intermediate knowledge of programming. For example, the reader should already understand most of the material in a book like [\"Advanced R\"](https://adv-r.hadley.nz/).\n",
    "\n",
    "In contrast, little to no knowledge of Julia is assumed. Characteristics of Julia that are not shared with languages like Python and R will be explicitly called out."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Mastering Julia for Statistical Computing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The core Julia language has exceptionally [good documentation](https://docs.julialang.org/en/v1/), but the documentation is focused on describing the language itself rather than on how to use the language to solve problems. Because of this gap, I find that many potential Julia users wish there were learning materials that (a) work through solving specific problems in detail rather than through language features in the abstract and (b) discuss best practices for programming in Julia as part of the exposition of solving specific problems. For statistics, this kind of documentation for Julia is, as of April 2020, still largely underdeveloped, and some of what does exist is not freely available online.\n",
    "\n",
    "This notebook represents my attempt to help just a little bit by writing up one specific problem in great detail. My hope is that reading through this notebook and internalizing the main ideas will help bring a user from \"able to write Julia that mostly works\" to \"writing Julia code that is competitive with code written by experts\". To do that, I'll describe how to implement maximum likelihood estimation for the logistic regression model. The final result is not intended to serve as the state of the art implementation of logistic regression for Julia, but, despite that, it should be clear to readers how to make the final implementation competitive with any other implementation in terms of numerical accuracy or execution performance.\n",
    "\n",
    "One major goal I have for this document is to put down in writing the many alternative ways of writing code that are plausible and help the reader understand the pros and cons of each. I find that too little of the educational material that programmers are exposed to focuses on comparisons between implementations, even though such comparisons often form the essence of how effective mentors educate their mentees. Learning to program well requires that one learn a causal mental model about how small changes in code lead to large changes in accuracy or performance; developing such a causal mental model requires that one see many small changes alongside the effects those changes have.\n",
    "\n",
    "In addition to these goals, the approach I'll take in this notebook is intended to emphasize broadly applicable principles that apply to model fitting for any probabilistic model. As such, we will not dig into special propreties of the logistic regression model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Review of the Mathematics of Logistic Regression via MLE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we begin, let's review the logistic regression in purely mathematical terms, which I assume you've seen before. The formulation we'll use here works as follows:\n",
    "\n",
    "* There's a constant (e.g. non-random) design matrix, $X$, that has $n$ rows and $d + 1$ columns. Each row describes an observation; each observation consists of $d$ features and a constant term that defines the intercept term of the logistic regression model.\n",
    "* There's a random, but observed, vector, $y$, of length $n$ that contains the binary outcome of the logistic regression model.\n",
    "* There's a constant (e.g. non-random) vector of parameters, $\\beta$, of length $d + 1$ that defines the coefficients of the logistic regression model.\n",
    "* Given the design matrix and the coefficients, we can construct the matrix-vector product $X \\beta$ and call it $z$. This vector defines the probabilities that $y_i = 1$ for all $i$, but the probabilities are in the logit-space. To get probabilities in probability space, we map $z_i \\to y_i$ via the inverse logit function, $\\text{logit}^{-1}(z) = (1 + \\exp(-z))^{-1}$.\n",
    "\n",
    "Summarizing all of that in distributional notation:\n",
    "\n",
    "$$\n",
    "X \\in \\mathbb{R}^{n, d + 1} \\\\\n",
    "\\beta \\in \\mathbb{R}^{d + 1} \\\\\n",
    "z = X \\beta \\\\\n",
    "p_i = \\text{logit}^{-1}(z_i) \\\\\n",
    "y_i \\sim \\text{Bernoulli}(p_i)\n",
    "$$\n",
    "\n",
    "The likelihood function for the full dataset of $n$ observations is therefore:\n",
    "\n",
    "$$\n",
    "L(\\beta) = \\prod_{i = 1}^{n} p_i^{y_i} (1 - p_i)^{1 - y_i}\n",
    "$$\n",
    "\n",
    "The log likelihood function, which is both mathematically and numerically better suited for use, is:\n",
    "\n",
    "$$\n",
    "\\mathcal{L}(\\beta) = \\sum_{i = 1}^{n} y_i \\log(p_i) + (1 - y_i) \\log(1 - p_i)\n",
    "$$\n",
    "\n",
    "We find the maximum likelihood estimate of $\\beta$ by maximizing the log likelihood. This is equivalent to minimizing the negative log likelihood. Because blackbox optimization API's usually default to minimization, we'll work with the negative log likelihood in our implementation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 1: Import Relevant Julia Libraries"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A common belief among programmers is that it's usually better to re-use existing code (especially public libraries that have been used by many other people) than to write code from scratch. For learning purposes, we will write several pieces of code from scratch in this notebook. But we will also try to use existing libraries as much as possible.\n",
    "\n",
    "In what follows, I'm going to assume you're using Julia 1.4, which is what I have installed on my machine. In addition to making use of the basic Julia installation, we'll use a few packages in this example that need to be manually installed by the user. If you do not already have these required packages installed, change `false` to `true` below and execute this block of code to install the missing packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_install = false\n",
    "\n",
    "if new_install\n",
    "    import Pkg\n",
    "\n",
    "    Pkg.add(\"Distributions\")\n",
    "    Pkg.add(\"ForwardDiff\")\n",
    "    Pkg.add(\"Optim\")\n",
    "    Pkg.add(\"StatsFuns\")\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this example, we're going to use the following packages:\n",
    "\n",
    "* *Standard Library*\n",
    "    * *LinearAlgebra*: The package provides functions and types for linear algebra, including computing dot products and matrix-vector multiplication.\n",
    "    * *Statistics*: This package provides functions for computing the most fundamental statistics like `mean` and `var.\n",
    "* *User-Installed Libraries*\n",
    "    * *Distributions*: This package provides functions and types for working with probability distributions. We'll use it to define Bernoulli and Normal distributions.\n",
    "    * *ForwardDiff*: This package provides functions for automatically differentation quasi-arbitrary Julia functions.\n",
    "    * *Optim*: This packages provides functions for optimization of quasi-arbitrary Julia functions.\n",
    "    * *StatsFuns*: This package provides implementations of common statistical functions like the CDF of the logistic distribution (i.e. the inverse logit function).\n",
    "\n",
    "We'll going to pull in each of these packages using Julia's `import` keyword, which doesn't import any names into the local scope except for the names explicitly specified by the user. If you prefer to get access to everything in the package in the way that Python's `from foo import *` works or R's `library(foo)`, you can do `using Foo` instead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import LinearAlgebra: diag, dot, mul!\n",
    "import Statistics: cov, mean, var\n",
    "\n",
    "import Distributions: Bernoulli, Normal, cquantile\n",
    "import ForwardDiff: hessian\n",
    "import Optim: LBFGS, minimizer, optimize\n",
    "import StatsFuns: logistic, log1pexp, logit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 2: Write a Data Generation Function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we implement a function for fitting a logistic regression via maximum likelhood, we're going to implement a function to generate a set of samples from the model given the design matrix, $X$, and the parameter vector, $\\beta$. I see starting this way as a broad principle for writing correct code for probabilistic models, so let's call it out as an explicit principle:\n",
    "\n",
    "**When implementing a probabilistic model, write the data generating function first.**\n",
    "\n",
    "More broadly, whenever possible, I like to work in the following order:\n",
    "\n",
    "* _Step 1: Write the generative model code._\n",
    "* _Step 2: Use that code to generate data that is truly generated by the model you're estimating._\n",
    "* _Step 3: Write the model fitting code._\n",
    "* _Step 4: Evaluate the model fitting code by checking how it behaves on data generated by the model. Exploit the fact that the parameters are fully known when you assess performance using generated data._\n",
    "\n",
    "The virtue of taking this approach is that it becomes far easier to test your model fitting code; all of the frequentist statistical theory about the quality of estimated parameters can be applied directly to your code as unit tests. If, in contrast, you work with an existing dataset whose data generating process you don't know, you can only check whether your code produces the same answers as existing software or analytic calculations. Since analytic calculations for logistic regressions are not generally tractable, that appproach is not particularly fulfilling."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With all that in mind, let's start to implement our data generation function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "generate! (generic function with 1 method)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function generate!(y, X, β)\n",
    "    for i in eachindex(y)\n",
    "        zᵢ = dot(X[i, :], β)\n",
    "        pᵢ = logistic(zᵢ)\n",
    "        y[i] = rand(Bernoulli(pᵢ))\n",
    "    end\n",
    "    y\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's walk through this function line-by-line to understand what's happening and why we've implemented things this way."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Line 1\n",
    "\n",
    "```\n",
    "function generate!(y, X, β)\n",
    "```\n",
    "\n",
    "The first line indicates that we're defining a function called `generate!` that takes in three arguments: `y`, `X` and `β`. The function name has an exclamation mark at the end to indicate that the function mutates at least one of its arguments. In particular, the function mutates `y`, which is the first argument because there's a convention in Julia of placing the arguments that will be mutated at the start of the argument list.\n",
    "\n",
    "Why are we writing this function so that it operates via mutation? Because mutating functions generally have better performance since they make it possible to remove memory allocations from the function body. If we want a pure function, we can write a wrapper for this mutating function that allocates a new arrray for `y`, calls this mutating function and then return the newly allocated `y`. This argument derives from another design principle:\n",
    "\n",
    "* **Write a mutating function that performs no allocations first and then write a pure wrapper for it that automatically allocates memory.**\n",
    "\n",
    "This principle is itself a special case of a broader principle:\n",
    "\n",
    "* **When faced with a tradeoff between performance and safety between functions X and Y, consider whether X can be built on top of Y or Y can be built on top of X. If, for example, Y can be built from X but X cannot be built from Y, implement X and provide it -- then expose Y built on top of X to users who prefer a safer, slower approach.**\n",
    "\n",
    "Note that we did not specify the types of any arguments. We could have written something like this instead:\n",
    "\n",
    "```\n",
    "function generate!(y::Vector{Float64}, X::Matrix{Float64}, β::Vector{Float64})\n",
    "```\n",
    "\n",
    "Why did we not specify types in that way? Because we want to keep our code generic. This reflects a general tension in how types are used in Julia. We can specify types either because:\n",
    "\n",
    "* We want to block certain types from being passed to the function by forcing a method error.\n",
    "* We want to make use of multiple dispatch to overload the function name to do slightly different things for different argument types.\n",
    "\n",
    "We don't want to do either of those things right now. Later on we might want to restrict our input types a bit more after understanding the space of valid inputs we want to allow. But for now we want to keep our code generic as we explore the design space. To give a sense of inputs we probably want to accept, all of the following seem reasonable:\n",
    "\n",
    "* `y::Vector{Int8}`\n",
    "* `y::Vector{Int32}`\n",
    "* `y::Vector{Int64}`\n",
    "* `y::Vector{Float32}`\n",
    "\n",
    "By not writing any input types, we allow all of these. If we had chosen the monomorphic `y::Vector{Float64}` declaration, we would have banned all of them from use.\n",
    "\n",
    "Note also that the absence of types will not introduce any performance problems. There are important cases in writing Julia code in which type information should be specified for maximum performance. The argument definitions for a function is never one of those cases. This is important enough to deserve being called out explicitly:\n",
    "\n",
    "**You do not need to annotate the types of function parameters to write fast Julia code.**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Line 2\n",
    "\n",
    "```\n",
    "for i in eachindex(y)\n",
    "```\n",
    "\n",
    "In this line, we start a loop over the indices for `y`. The `eachindex(y)` function is a popular way to extract indices for arrays in Julia because it can provide the most efficient indexing strategy depending on the type of array you're using, whereas a more straightforward `for i in 1:length(y)` idiom might be sub-optimal for more complex array types like sparse arrays. If you're interested to see more, try running `eachindex` on a 10x10 dense matrix like `rand(10, 10)` versus running it on `SparseArrays.spzeros(10, 10)` from the `SparseArrays` package that's part of Julia's standard library.\n",
    "\n",
    "One thing to note about this line is that we're going to implicitly assume in the body of the looop that the following invariants hold, but we will not be testing them: `length(y) == size(X, 1)` and `size(X, 2) == length(β)`. We are skipping these checks for the same reason we're mutating our inputs: they can be hoisted out of this part of the code and tested before we enter any loop that invokes this code. We make code like this fast by not doing work we don't have to do, although the result is code that's less safe. In some settings, Julia might be smart enough to remove redundant checks, but we're going to take a bit more control for ourselves.\n",
    "\n",
    "Stated as a design principle: **it's better to write an inner function that assumes invariants hold and then write outer functions to enforce them than it is to constantly recheck them inside the inner function, especially if the inner function could only throw an exception if the invariants failed in the outer functions**. This principle and the previous principle of starting with a mutating function reflect a broader principle: acknowledge asymmetries in what can be built on top of what. It's easy to build a safe, slower function on top of an unsafe, faster function -- but the reverse is not true. Likewise it's easy to build a pure wrapper that allocates outputs on top of a mutating function, but it's not possible to avoid allocations if your pure function always performs them.\n",
    "\n",
    "Having set the loop start and end, we enter the loop body where we'll generate our observations one-by-one. For each observation, we generate some scalar intermediates. (We could also generate all of them at once using a mutating operation. We can even get away without needing to allocate memory because we can reuse the same vector.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Line 3\n",
    "\n",
    "```\n",
    "zᵢ = dot(X[i, :], β)\n",
    "```\n",
    "\n",
    "Here we extract the i-th row of `X` and take its dot product with `β` using the `LinearAlgebra.dot` function, which takes in two vectors and produces a scalar. The result is stored in the scalar variable `zᵢ`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Lines 4-5\n",
    "\n",
    "```\n",
    "pᵢ = logistic(zᵢ)\n",
    "y[i] = rand(Bernoulli(pᵢ))\n",
    "```\n",
    "\n",
    "Here we use the `logistic` function from the StatsFuns package to transform `zᵢ` into `pᵢ`. The benefit of using the version of this function from the StatsFuns package is that it's been written to support the largest set of inputs possible; a naive `logistic(z) = 1 / (1 + exp(-z))` implementation will generate exact `0.0` probabilities earlier than it should.\n",
    "\n",
    "After computing `pᵢ`, we generate the observed `y[i]` Bernoulli outcomes. We do this by constructing a Bernoulli distribution object and calling the `rand` function on it. This is very efficient because the `Bernoulli` object is immutable."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Lines 6-8\n",
    "\n",
    "These lines mostly contain the ends of blocks, which Julia marks with `end`.\n",
    "\n",
    "The interesting part is Line 7, where we return `y` at the end to make it easy to use our function with inputs that aren't named variables."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Some Small Possible Variations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are a couple of minor changes we could make that might be important in some settings. For example, we can use views instead of copies to replace `X[i, :]` with `@view(X[i, :])`. For very large arrays, this can help reduce the amount of memory we allocate, but it's worth noting that complex views can cause downstream code to be slow because they have to constantly perform quirky indexing operations that are not cache-friendly.\n",
    "\n",
    "We could also replace `LinearAlgebra.dot(X[i, :], β)` with `X[i, :]' * β` using Julia's lazy transpose operation.\n",
    "\n",
    "Finally we could draw random samples from the logistic distribution and compare them with `zᵢ` instead of ever generating `pᵢ`. This is the latent variable representation of logistic regression, but it's not clear that it would provide performance benefits to make this change.\n",
    "\n",
    "In what follows, we make a few of these changes and also use the builtin Julia macro `@inbounds` to turn off bounds checking inside of a code block. This changes the code from throwing a bounds-checking exception if the invariants we mentioned earlier are false to segfaulting. There are non-trivial performance benefits to doing this in some cases since the Julia compiler isn't always able to convince itself that bounds-checks are safe to eliminate automatically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "generate! (generic function with 1 method)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function generate!(y, X, β)\n",
    "    @inbounds for i in eachindex(y)\n",
    "        zᵢ = @view(X[i, :])' * β\n",
    "        pᵢ = logistic(zᵢ)\n",
    "        y[i] = rand(Bernoulli(pᵢ))\n",
    "    end\n",
    "    y\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you're interested in deciding which of the many combinatorial variants is best, you should use the `@btime` macro from the BenchmarkTools package to compare them explicitly in terms of speed and memory usage."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A Different Coding Style"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All of the variants we've seen so far have explicit iterative loops over the data. This is often the easiest way to write code like this and it's quite fast because of Julia's language design. But there are other approaches:\n",
    "\n",
    "* We could exploit the embarassingly parallel nature of the sum we're computing: there's no relationship between the i-th term in our inputs or outputs to any other term. We could employ threads that operate on disjoint set of observations to take advantage of this independence, for example.\n",
    "* Write a \"vectorized\" solution that operates on the entire dataset at once at the function call level (but where implementation still has to think about how to process individual elements)\n",
    "* Make use of BLAS calls to get some peformance improvements at the cost of potentially requiring allocating memory. In this example, there is no such cost because we get away with writing the value of `z` temporarily into the `y` array and then writing over that array with the value of `y`.\n",
    "\n",
    "Below, we'll use this BLAS approach by calling `LinearAlgebra.mul!`. At that same time, we'll show a \"vectorized\" approach to evaluating `logistic` and calling `rand` by using Julia's dot broadcasting notation, which explicitly \"vectorizes\" any scalar function we already have. Dot broadcasting is very special part of Julia because it uses syntax to lift scalar functions to vectorized functions, but it strictly more performant than traditional vectorization because it sees a whole sequence of operations at once and can perform loop fusion to avoid having to loop over an array multiple times."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "generate! (generic function with 1 method)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function generate!(y, X, β)\n",
    "    mul!(y, X, β)\n",
    "    y .= rand.(Bernoulli.(logistic.(y)))\n",
    "    y\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Testing our implementations of the `generate!` function to ensure they're equivalent is really hard. We can check they generate similar data according to summary statistics and use frequentist bounds to ensure the summaries are credible, but that's all I know how to do. In this example, I'm very confident all the modifications are safe."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 3: Generate Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have a function to sample data, let's create some data with it and use it to test the model fitting code we'll write. We'll start with a decent number of observations so that our estimates are tolerably precise without being so accurate that confidence intervals are uninteresting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "simulate_data (generic function with 1 method)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function simulate_data(n, d)\n",
    "    X = hcat(ones(n), rand(Normal(0, 1), n, d));\n",
    "    β = rand(Normal(0, 1), d + 1)\n",
    "    y = Array{Float64}(undef, n)\n",
    "    generate!(y, X, β)\n",
    "    y, X, β\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "([0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0  …  0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0], [1.0 0.10781238613877225 0.9185577169564151; 1.0 0.2875710507117643 0.29447235851120307; … ; 1.0 1.4970548123092562 0.25385694441532525; 1.0 -2.371265555206343 0.8649145476870038], [-1.1382399429248256, -0.25633479736844866, -0.29936682624813693])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y, X, β = simulate_data(10_000, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 4: Implementing the Log Likelihood Function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have data in hand, let's code up the log likelihood function. I like to do this by copying the body of the `generate!` function and then changing it appropriately so the two pieces of code are maximally similar. A formal PPL would make it easier to reuse code between them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "log_likelihood_naive (generic function with 1 method)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function log_likelihood_naive(X, y, β)\n",
    "    ll = 0.0\n",
    "    @inbounds for i in eachindex(y)\n",
    "        zᵢ = dot(X[i, :], β)\n",
    "        pᵢ = logistic(zᵢ)\n",
    "        ll += y[i] * log(pᵢ) + (1 - y[i]) * log(1 - pᵢ)\n",
    "    end\n",
    "    ll\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This code has some problems of numerical accuracy and it also perform computations we don't really need to perform to produce the correct output. The essence of the problem with this naive translation of the log likelihoood equation is that we're doing work to compute `logistic(zᵢ)` when we really only need `log(logistic(zᵢ))`. We can improve on this using the `log1pexp` method we imported earlier, which also provides substantially better numerical accuracy when any of the `zᵢ < -710.0`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "log_likelihood (generic function with 1 method)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function log_likelihood(X, y, β)\n",
    "    ll = 0.0\n",
    "    @inbounds for i in eachindex(y)\n",
    "        zᵢ = dot(X[i, :], β)\n",
    "        c = -log1pexp(-zᵢ) # Conceptually equivalent to log(1 / (1 + exp(-zᵢ))) == -log(1 + exp(-zᵢ))\n",
    "        ll += y[i] * c + (1 - y[i]) * (-zᵢ + c) # Conceptually equivalent to log(exp(-zᵢ) / (1 + exp(-zᵢ)))\n",
    "    end\n",
    "    ll\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can reassure ourselves that our changes are correct:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(-5490.842415218151, -5490.842415218151)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(\n",
    "    log_likelihood_naive(X, y, β),\n",
    "    log_likelihood(X, y, β),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The log likelihood as we've written it is a function of both the data and the parameters, but mathematically it should only depend on the parameters, $\\beta$. In addition to that mathetical reason for creating a new function, we want a function only of the parameters because the optimization algorithms in Optim assume the inputs have that property. To achieve both goals, we'll construct a closure that partially applies the log likelihood function for us and negates it to give us the negative log likelihood we want to minimize."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Log Likelihood Function Should Be a Closure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "make_closures (generic function with 1 method)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "make_closures(X, y) = β -> -log_likelihood(X, y, β)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "#3 (generic function with 1 method)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nll = make_closures(X, y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 5: Minimizing the Negative Log Likelihood Function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have the negative log likelihood we'll want to minimize it starting from some point. It's common to initialize all of the parameters to zero, so let's start there:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "β₀ = zeros(2 + 1); # d = 2 and we want an intercept term"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can then check whether the negative log likelihood evaluated relative to the zero parameter function gives a value that seems plausible:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6931.471805600547"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nll(β₀)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we want to minimixe the negative log likelihood. To do that, we'll use the Optim.jl library, which provides an `optimize` method for minimization of blackbox functions. We'll pass two options to `optimize` to improve our results:\n",
    "1. We'll use the L-BFGS algorithm to exploit the gradient that can be compute for the negative log likelihood function rather than let the algorithm default to Nelder-Mead.\n",
    "2. We'll pass in an argument to use forward-mode automatic differentation to ensure that we get exact gradients rather than approximate ones. Without this, the algorithm will sometimes (or even often) fail to converge to a highly precise result because the finite-difference gradients that calculated by default will become inaccurate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       " * Status: success\n",
       "\n",
       " * Candidate solution\n",
       "    Minimizer: [-1.13e+00, -2.44e-01, -3.23e-01]\n",
       "    Minimum:   5.490175e+03\n",
       "\n",
       " * Found with\n",
       "    Algorithm:     L-BFGS\n",
       "    Initial Point: [0.00e+00, 0.00e+00, 0.00e+00]\n",
       "\n",
       " * Convergence measures\n",
       "    |x - x'|               = 1.21e-08 ≰ 0.0e+00\n",
       "    |x - x'|/|x'|          = 1.07e-08 ≰ 0.0e+00\n",
       "    |f(x) - f(x')|         = 3.09e-11 ≰ 0.0e+00\n",
       "    |f(x) - f(x')|/|f(x')| = 5.63e-15 ≰ 0.0e+00\n",
       "    |g(x)|                 = 5.16e-10 ≤ 1.0e-08\n",
       "\n",
       " * Work counters\n",
       "    Seconds run:   1  (vs limit Inf)\n",
       "    Iterations:    7\n",
       "    f(x) calls:    28\n",
       "    ∇f(x) calls:   28\n"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res = optimize(nll, β₀, LBFGS(), autodiff=:forward)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I personally like to initialize the parameters to values that I'm hopeful will let the algorithm converge faster without being costly to compute. One heuristic I've used for logistic regression is the following:\n",
    "\n",
    "* Set the intercept to be exactly right if there were no other parameters.\n",
    "* Set the other coefficients based on doing a standard univariate OLS fit to the logit-transformed data after replacing 0's with $\\epsilon$ and 1's with $1 - \\epsilon$. I use a large $\\epsilon = 0.1$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3-element Array{Float64,1}:\n",
       " -1.0932860443854355\n",
       " -0.1939078725087912\n",
       " -0.2582119475758624"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function initialize!(β₀, X, y, ϵ = 0.1)\n",
    "    β₀[1] = logit(mean(y))\n",
    "    logit_y = [ifelse(y_i == 1.0, logit(1 - ϵ), logit(ϵ)) for y_i in y]\n",
    "    for j in 2:length(β₀)\n",
    "         β₀[j] = cov(logit_y, @view(X[:, j])) / var(@view(X[:, j]))\n",
    "    end\n",
    "    β₀\n",
    "end\n",
    "\n",
    "initialize!(β₀, X, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       " * Status: success\n",
       "\n",
       " * Candidate solution\n",
       "    Minimizer: [-1.13e+00, -2.44e-01, -3.23e-01]\n",
       "    Minimum:   5.490175e+03\n",
       "\n",
       " * Found with\n",
       "    Algorithm:     L-BFGS\n",
       "    Initial Point: [-1.09e+00, -1.94e-01, -2.58e-01]\n",
       "\n",
       " * Convergence measures\n",
       "    |x - x'|               = 5.79e-11 ≰ 0.0e+00\n",
       "    |x - x'|/|x'|          = 5.10e-11 ≰ 0.0e+00\n",
       "    |f(x) - f(x')|         = 2.36e-11 ≰ 0.0e+00\n",
       "    |f(x) - f(x')|/|f(x')| = 4.31e-15 ≰ 0.0e+00\n",
       "    |g(x)|                 = 1.39e-13 ≤ 1.0e-08\n",
       "\n",
       " * Work counters\n",
       "    Seconds run:   0  (vs limit Inf)\n",
       "    Iterations:    6\n",
       "    f(x) calls:    17\n",
       "    ∇f(x) calls:   17\n"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res = optimize(nll, β₀, LBFGS(), autodiff=:forward)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you compare the work counters, you can see that the optimization procedure had to compute the log likelihood and its gradient fewer times. I don't know how to prove that my initialization approach will always have this effect and think it's very possible that this initialization won't always find an optimum faster. But in practice I've found it can make things a bit faster because the initialization step costs much less than evaluating `f(x)` and `∇f(x)` a few times."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Given the results of optimization, we can extract our estimates using the `Optim.minimizer` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3-element Array{Float64,1}:\n",
       " -1.1341732805409002\n",
       " -0.24414707752560674\n",
       " -0.3229958057866628"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "β̂ = minimizer(res)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 6: Testing Our Estimates via Confidence Interval Coverage Checks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we have estimates, but how do we know if the estimates are good enough?\n",
    "\n",
    "Given results about the convergence in probability of logistic regression coefficients, we know that as `n` goes to infinity, β̂ converges to β. Unfortunately, we can't simulate infinite data. In finite sample the convergence isn't complete, so there's some error.\n",
    "\n",
    "So the question for evaluating our code becomes: how do we know the error is reasonable? This is a place where frequentist statistics is very useful -- if the model is true (which we're trying to ensure occurs by construction), then asymptotically we can use the Fisher Information Matrix to compute confidence intervals for β̂ and check whether they contain β. We'll follow a standard of using the observed Fisher information matrix instead, since that only requires us to evaluate the Hesssian of the negative log likelihood function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "compute_ses (generic function with 1 method)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function compute_ses(nll, β̂)\n",
    "    H = hessian(nll, β̂)\n",
    "    ses = sqrt.(diag(inv(H)))\n",
    "    ses\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "See Chapter 9 of Wasserman's All of Statistics for details on the math we're using here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "compute_cis (generic function with 1 method)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function compute_cis(nll, β̂, α)\n",
    "    ses = compute_ses(nll, β̂)\n",
    "    τ = cquantile(Normal(0, 1), α)\n",
    "    lower = β̂ - τ * ses\n",
    "    upper = β̂ + τ * ses\n",
    "    lower, upper\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Standard CI computation using quantiles from the normal distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "check_cis (generic function with 1 method)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "check_cis(β, lower, upper) = all(lower .<= β .<= upper)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "true"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "α = 0.001\n",
    "check_cis(β, compute_cis(nll, β̂, α)...)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Conclusion"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Hopefully this short tutorial gives a flavor of how to fit models via MLE in Julia. There's many more topics that we could have explored, but I wanted to keep things relatively short. A few topics that I'd encourage the reader to investigate:\n",
    "\n",
    "* Is it better to use the analytic gradient in `optimize` in terms of accuracy or performance?\n",
    "* Do our results match the results from Julia's GLM package or R's `glm` function?\n",
    "* Should we modify the log_likelihood function to automatically scale inputs so they have a standard deviation of 1?\n",
    "* How do we construct robust standard errors when the model is misspecified? There's an intro in Julia to robust standard errors [here](https://github.com/PaulSoderlind/FinancialEconometrics/blob/master/Ch12_MLE.ipynb)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.4.1",
   "language": "julia",
   "name": "julia-1.4"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.5.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}