ipynb/mizuhiro_suzuki-jump_estimation_basic.ipynb

{
    "cells": [
        {
            "cell_type": "markdown",
            "id": "4decb0a8-edfd-49c8-abbc-b8c82cf97016",
            "metadata": {},
            "source": [
                "# OLS and ML estimations with Julia JuMP"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "2f92e842-2c97-4051-891d-e0ec96aedf17",
            "metadata": {},
            "source": [
                "This notebook showcases how to estimate parameters based on OLS or ML methods using [Julia's JuMP package](https://jump.dev/JuMP.jl/stable/).\n",
                "The main purpose of this document is not to argue that JuMP is a perfect tool for simple estimations but rather to show how to use JuMP with data."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 1,
            "id": "4f693d92-ac7c-422c-b640-c2f8d2cf2be4",
            "metadata": {},
            "outputs": [],
            "source": [
                "# using Pkg\n",
                "\n",
                "# Pkg.add(\"JuMP\")\n",
                "# Pkg.add(\"DataFrames\")\n",
                "# Pkg.add(\"Ipopt\")\n",
                "# Pkg.add(\"BenchmarkTools\")\n",
                "# Pkg.add(\"Optim\")\n",
                "# Pkg.add(\"ForwardDiff\")"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 2,
            "id": "3c1b23a8-024d-4864-bd71-ae0406fe60c9",
            "metadata": {},
            "outputs": [],
            "source": [
                "using JuMP, Ipopt\n",
                "using Random\n",
                "using BenchmarkTools\n",
                "using Optim\n",
                "# using DataFrames\n",
                "# using ForwardDiff"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "c82253e1-c0af-450d-8a3f-c9699370ca82",
            "metadata": {},
            "source": [
                "## Data generation"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "bae7f17b-1404-4a5f-b35b-36be60129ca9",
            "metadata": {},
            "source": [
                "I consider the following equation for estimation:\n",
                "$$\n",
                "    y_i = \\alpha + X_i' \\beta + u_i.\n",
                "$$\n",
                "The error term follows the standard normal distribution.\n",
                "The data are simulated with $X_i$ and all coefficients following the standard normal distribution."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 3,
            "id": "c6f89bfd-8929-4c43-abae-b386c9158724",
            "metadata": {},
            "outputs": [],
            "source": [
                "Random.seed!(3);"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 4,
            "id": "f0432699-fcf1-42a1-b47a-25ec3c0ed89e",
            "metadata": {},
            "outputs": [],
            "source": [
                "N = 1_000;\n",
                "K = 3;"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 5,
            "id": "0aa3b9e1-aacf-4722-be1b-5db0d6377d94",
            "metadata": {},
            "outputs": [],
            "source": [
                "β = randn(Float64, K + 1);"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 6,
            "id": "88e15558-7e5e-46fa-aea6-2b0d68c63295",
            "metadata": {},
            "outputs": [],
            "source": [
                "X = hcat(ones(N), randn(Float64, (N, K)));\n",
                "u = randn(Float64, N);\n",
                "y = X * β .+ u;"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "4f9fb26f-874d-4e73-8ac8-552f0aa61d07",
            "metadata": {},
            "source": [
                "# OLS"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "67ff77f8-b5fb-494a-9441-05aeb2b84eaa",
            "metadata": {},
            "source": [
                "## OLS with matrix inversion"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 7,
            "id": "cb9ca8e0-818f-475c-a714-8fdabde0b55f",
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "estimate_ols (generic function with 1 method)"
                        ]
                    },
                    "execution_count": 7,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "function estimate_ols(y, X)\n",
                "    βhat = (X \\ X) * (X \\ y);\n",
                "    return βhat\n",
                "end"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 8,
            "id": "0478007d-82d2-4a26-bb7f-dcad0d308360",
            "metadata": {},
            "outputs": [],
            "source": [
                "βhat_ols = estimate_ols(y, X);"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 9,
            "id": "67b4f794-88a7-4263-9822-073684729516",
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "BenchmarkTools.Trial: 10000 samples with 1 evaluation.\n",
                            " Range \u001b[90m(\u001b[39m\u001b[36m\u001b[1mmin\u001b[22m\u001b[39m … \u001b[35mmax\u001b[39m\u001b[90m):  \u001b[39m\u001b[36m\u001b[1m39.458 μs\u001b[22m\u001b[39m … \u001b[35m 11.431 ms\u001b[39m  \u001b[90m┊\u001b[39m GC \u001b[90m(\u001b[39mmin … max\u001b[90m): \u001b[39m0.00% … 98.41%\n",
                            " Time  \u001b[90m(\u001b[39m\u001b[34m\u001b[1mmedian\u001b[22m\u001b[39m\u001b[90m):     \u001b[39m\u001b[34m\u001b[1m41.625 μs               \u001b[22m\u001b[39m\u001b[90m┊\u001b[39m GC \u001b[90m(\u001b[39mmedian\u001b[90m):    \u001b[39m0.00%\n",
                            " Time  \u001b[90m(\u001b[39m\u001b[32m\u001b[1mmean\u001b[22m\u001b[39m ± \u001b[32mσ\u001b[39m\u001b[90m):   \u001b[39m\u001b[32m\u001b[1m55.592 μs\u001b[22m\u001b[39m ± \u001b[32m151.815 μs\u001b[39m  \u001b[90m┊\u001b[39m GC \u001b[90m(\u001b[39mmean ± σ\u001b[90m):  \u001b[39m9.66% ±  7.36%\n",
                            "\n",
                            "  \u001b[34m█\u001b[39m\u001b[39m▄\u001b[39m▄\u001b[39m▃\u001b[32m▂\u001b[39m\u001b[39m▃\u001b[39m▂\u001b[39m▂\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m▁\n",
                            "  \u001b[34m█\u001b[39m\u001b[39m█\u001b[39m█\u001b[39m█\u001b[32m█\u001b[39m\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m▇\u001b[39m▆\u001b[39m▆\u001b[39m▆\u001b[39m▅\u001b[39m▅\u001b[39m▃\u001b[39m▅\u001b[39m▅\u001b[39m▃\u001b[39m▃\u001b[39m▅\u001b[39m▅\u001b[39m▄\u001b[39m▅\u001b[39m▅\u001b[39m▄\u001b[39m▅\u001b[39m▃\u001b[39m▄\u001b[39m▄\u001b[39m▄\u001b[39m▄\u001b[39m▁\u001b[39m▄\u001b[39m▄\u001b[39m▄\u001b[39m▃\u001b[39m▃\u001b[39m▁\u001b[39m▅\u001b[39m▄\u001b[39m▃\u001b[39m▄\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▃\u001b[39m▃\u001b[39m▁\u001b[39m▄\u001b[39m▁\u001b[39m▃\u001b[39m▃\u001b[39m \u001b[39m█\n",
                            "  39.5 μs\u001b[90m       \u001b[39m\u001b[90mHistogram: \u001b[39m\u001b[90m\u001b[1mlog(\u001b[22m\u001b[39m\u001b[90mfrequency\u001b[39m\u001b[90m\u001b[1m)\u001b[22m\u001b[39m\u001b[90m by time\u001b[39m       306 μs \u001b[0m\u001b[1m<\u001b[22m\n",
                            "\n",
                            " Memory estimate\u001b[90m: \u001b[39m\u001b[33m172.95 KiB\u001b[39m, allocs estimate\u001b[90m: \u001b[39m\u001b[33m55\u001b[39m."
                        ]
                    },
                    "execution_count": 9,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "@benchmark estimate_ols(y, X)"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "e52c54eb-d6bf-43a3-ba60-f958922bb4f9",
            "metadata": {},
            "source": [
                "## OLS with JuMP"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "70536770-9350-4753-8569-32552449fb85",
            "metadata": {},
            "source": [
                "Here I use JuMP for OLS estimation.\n",
                "I use the `Ipopt` solver."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 10,
            "id": "2f663d22-0ba3-43bf-85a9-9460f8cffbb8",
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "estimate_ols_jump (generic function with 1 method)"
                        ]
                    },
                    "execution_count": 10,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "function estimate_ols_jump(y, X; K = K, N = N)\n",
                "\n",
                "    model = Model(Ipopt.Optimizer)\n",
                "    set_silent(model)\n",
                "    set_string_names_on_creation(model, false)\n",
                "    @variable(model, β[1:(K + 1)])\n",
                "    @objective(\n",
                "        model,\n",
                "        Min,\n",
                "        sum((y[i] - sum(X[i, j] * β[j] for j in 1:(K + 1)))^2 for i in 1:N)\n",
                "    )\n",
                "    optimize!(model)\n",
                "    @assert is_solved_and_feasible(model)\n",
                "\n",
                "    return value.(β)\n",
                "end"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 11,
            "id": "981ca359-851b-4268-9341-173e9c0b5a0e",
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "\n",
                        "******************************************************************************\n",
                        "This program contains Ipopt, a library for large-scale nonlinear optimization.\n",
                        " Ipopt is released as open source code under the Eclipse Public License (EPL).\n",
                        "         For more information visit https://github.com/coin-or/Ipopt\n",
                        "******************************************************************************\n",
                        "\n"
                    ]
                }
            ],
            "source": [
                "βhat_ols_jump = estimate_ols_jump(y, X);"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 12,
            "id": "8c2a781f-2765-4ab3-a080-63732c7445e3",
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "BenchmarkTools.Trial: 895 samples with 1 evaluation.\n",
                            " Range \u001b[90m(\u001b[39m\u001b[36m\u001b[1mmin\u001b[22m\u001b[39m … \u001b[35mmax\u001b[39m\u001b[90m):  \u001b[39m\u001b[36m\u001b[1m3.728 ms\u001b[22m\u001b[39m … \u001b[35m37.804 ms\u001b[39m  \u001b[90m┊\u001b[39m GC \u001b[90m(\u001b[39mmin … max\u001b[90m): \u001b[39m0.00% … 68.12%\n",
                            " Time  \u001b[90m(\u001b[39m\u001b[34m\u001b[1mmedian\u001b[22m\u001b[39m\u001b[90m):     \u001b[39m\u001b[34m\u001b[1m4.758 ms              \u001b[22m\u001b[39m\u001b[90m┊\u001b[39m GC \u001b[90m(\u001b[39mmedian\u001b[90m):    \u001b[39m0.00%\n",
                            " Time  \u001b[90m(\u001b[39m\u001b[32m\u001b[1mmean\u001b[22m\u001b[39m ± \u001b[32mσ\u001b[39m\u001b[90m):   \u001b[39m\u001b[32m\u001b[1m5.577 ms\u001b[22m\u001b[39m ± \u001b[32m 2.406 ms\u001b[39m  \u001b[90m┊\u001b[39m GC \u001b[90m(\u001b[39mmean ± σ\u001b[90m):  \u001b[39m5.66% ±  8.31%\n",
                            "\n",
                            "  \u001b[39m \u001b[39m \u001b[39m▃\u001b[39m▆\u001b[39m█\u001b[39m▆\u001b[34m▃\u001b[39m\u001b[39m▁\u001b[39m \u001b[39m \u001b[39m \u001b[32m \u001b[39m\u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \n",
                            "  \u001b[39m▂\u001b[39m▆\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m█\u001b[34m█\u001b[39m\u001b[39m█\u001b[39m▇\u001b[39m▄\u001b[39m▅\u001b[32m▃\u001b[39m\u001b[39m▄\u001b[39m▄\u001b[39m▄\u001b[39m▃\u001b[39m▃\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▃\u001b[39m▂\u001b[39m▂\u001b[39m▃\u001b[39m▃\u001b[39m▃\u001b[39m▃\u001b[39m▄\u001b[39m▄\u001b[39m▃\u001b[39m▃\u001b[39m▃\u001b[39m▃\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▁\u001b[39m▂\u001b[39m▁\u001b[39m▂\u001b[39m▂\u001b[39m▁\u001b[39m▁\u001b[39m▂\u001b[39m▁\u001b[39m▂\u001b[39m▁\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▁\u001b[39m▂\u001b[39m \u001b[39m▃\n",
                            "  3.73 ms\u001b[90m        Histogram: frequency by time\u001b[39m        13.6 ms \u001b[0m\u001b[1m<\u001b[22m\n",
                            "\n",
                            " Memory estimate\u001b[90m: \u001b[39m\u001b[33m3.90 MiB\u001b[39m, allocs estimate\u001b[90m: \u001b[39m\u001b[33m36453\u001b[39m."
                        ]
                    },
                    "execution_count": 12,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "@benchmark estimate_ols_jump(y, X)"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "a8cb159d-9419-4bd4-81dc-2b0995bd43bd",
            "metadata": {},
            "source": [
                "#### Performance comparisons: Access global variables"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "e01ebda1-d37e-4234-be30-2c4df4ba2aef",
            "metadata": {},
            "source": [
                "As [in the general Julia performance tips](https://docs.julialang.org/en/v1/manual/performance-tips/index.html#Avoid-untyped-global-variables), it is a good practice to avoid to call untyped global variables in functions.\n",
                "In `estimate_ols_jump`, I added `K` and `N` as keyword arguments of the function.\n",
                "Not doing this (= accessing the untyped global variables from the function) can increase the allocated memory and the runtime:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 13,
            "id": "cc51ac4f-4fd3-40f0-9992-785bc7e3d868",
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "estimate_ols_jump_global (generic function with 1 method)"
                        ]
                    },
                    "execution_count": 13,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "function estimate_ols_jump_global(y, X)\n",
                "\n",
                "    model = Model(Ipopt.Optimizer)\n",
                "    set_silent(model)\n",
                "    @variable(model, β[1:(K + 1)])\n",
                "    @objective(\n",
                "        model,\n",
                "        Min,\n",
                "        sum((y[i] - sum(X[i, j] * β[j] for j in 1:(K + 1)))^2 for i in 1:N)\n",
                "    )\n",
                "    optimize!(model)\n",
                "    @assert is_solved_and_feasible(model)\n",
                "\n",
                "    return value.(β)\n",
                "end"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 14,
            "id": "e387beff-17ed-4fa7-a093-b79f284ef633",
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "BenchmarkTools.Trial: 818 samples with 1 evaluation.\n",
                            " Range \u001b[90m(\u001b[39m\u001b[36m\u001b[1mmin\u001b[22m\u001b[39m … \u001b[35mmax\u001b[39m\u001b[90m):  \u001b[39m\u001b[36m\u001b[1m4.208 ms\u001b[22m\u001b[39m … \u001b[35m36.129 ms\u001b[39m  \u001b[90m┊\u001b[39m GC \u001b[90m(\u001b[39mmin … max\u001b[90m): \u001b[39m0.00% … 59.37%\n",
                            " Time  \u001b[90m(\u001b[39m\u001b[34m\u001b[1mmedian\u001b[22m\u001b[39m\u001b[90m):     \u001b[39m\u001b[34m\u001b[1m5.259 ms              \u001b[22m\u001b[39m\u001b[90m┊\u001b[39m GC \u001b[90m(\u001b[39mmedian\u001b[90m):    \u001b[39m0.00%\n",
                            " Time  \u001b[90m(\u001b[39m\u001b[32m\u001b[1mmean\u001b[22m\u001b[39m ± \u001b[32mσ\u001b[39m\u001b[90m):   \u001b[39m\u001b[32m\u001b[1m6.106 ms\u001b[22m\u001b[39m ± \u001b[32m 2.598 ms\u001b[39m  \u001b[90m┊\u001b[39m GC \u001b[90m(\u001b[39mmean ± σ\u001b[90m):  \u001b[39m5.70% ±  8.52%\n",
                            "\n",
                            "  \u001b[39m \u001b[39m▄\u001b[39m█\u001b[39m█\u001b[34m█\u001b[39m\u001b[39m▆\u001b[39m▅\u001b[39m▃\u001b[32m▂\u001b[39m\u001b[39m▂\u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m▃\u001b[39m▂\u001b[39m▁\u001b[39m▁\u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \n",
                            "  \u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m█\u001b[34m█\u001b[39m\u001b[39m█\u001b[39m█\u001b[39m█\u001b[32m█\u001b[39m\u001b[39m█\u001b[39m▇\u001b[39m█\u001b[39m█\u001b[39m▇\u001b[39m▇\u001b[39m▆\u001b[39m▄\u001b[39m▆\u001b[39m▆\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m▆\u001b[39m▇\u001b[39m▆\u001b[39m▅\u001b[39m▆\u001b[39m▅\u001b[39m▁\u001b[39m▅\u001b[39m▄\u001b[39m▁\u001b[39m▁\u001b[39m▄\u001b[39m▄\u001b[39m▁\u001b[39m▄\u001b[39m▁\u001b[39m▄\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▄\u001b[39m▄\u001b[39m▄\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▄\u001b[39m \u001b[39m█\n",
                            "  4.21 ms\u001b[90m      \u001b[39m\u001b[90mHistogram: \u001b[39m\u001b[90m\u001b[1mlog(\u001b[22m\u001b[39m\u001b[90mfrequency\u001b[39m\u001b[90m\u001b[1m)\u001b[22m\u001b[39m\u001b[90m by time\u001b[39m     17.7 ms \u001b[0m\u001b[1m<\u001b[22m\n",
                            "\n",
                            " Memory estimate\u001b[90m: \u001b[39m\u001b[33m4.30 MiB\u001b[39m, allocs estimate\u001b[90m: \u001b[39m\u001b[33m52459\u001b[39m."
                        ]
                    },
                    "execution_count": 14,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "@benchmark estimate_ols_jump_global(y, X)"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "8f73bcfa-80bb-44d2-8a81-0ee66ee89941",
            "metadata": {},
            "source": [
                "## Estimates comparison"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 34,
            "id": "c1c2ffde-daa8-4443-b4bb-7453fef392b8",
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "4×3 Matrix{Float64}:\n",
                            " -0.627644  -0.627644  -0.663764\n",
                            "  1.36376    1.36376    1.36321\n",
                            "  0.756496   0.756496   0.717411\n",
                            "  0.660766   0.660766   0.678743"
                        ]
                    },
                    "execution_count": 34,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "[βhat_ols βhat_ols_jump β]"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "2d5bc36c-631b-4b93-b83e-290fffe7a7ab",
            "metadata": {},
            "source": [
                "# Maximum Likelihood"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "ab35992f-e6c6-45db-a93f-e1c58d90d85e",
            "metadata": {},
            "source": [
                "Next, I estimate the parameters with the maximum likelihood (ML) method.\n",
                "Here, I also estimate the standard deviation of the error term (the true s.d. is 1)."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 16,
            "id": "f35c96f8-50ef-40a1-b7ae-c8c575c72086",
            "metadata": {},
            "outputs": [],
            "source": [
                "initial_x = ones(K + 2);"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 17,
            "id": "afd6cfb6-24b2-493d-b949-8786cc48794c",
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "calculate_log_likelihood (generic function with 1 method)"
                        ]
                    },
                    "execution_count": 17,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "function calculate_log_likelihood(param, y, X; K = K, N = N)\n",
                "    β = @view param[1:(K + 1)]\n",
                "    σ = param[K + 2]\n",
                "    log_likelihood = (\n",
                "        N / 2 * log(1 / (2 * π * σ^2)) -\n",
                "        sum((y - X * β).^2) / (2 * σ^2)\n",
                "    )\n",
                "\n",
                "    return log_likelihood\n",
                "end"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "76ec2706-8403-478b-9ad7-6c7509090d89",
            "metadata": {},
            "source": [
                "## With `Optim`"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "1a05d3d5-b2ba-4281-ba55-15fd062c5149",
            "metadata": {},
            "source": [
                "First, I do the ML estimation with Julia's `Optim` package, which is one of the most popular optimization package."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 18,
            "id": "b566c8b6-8181-46e2-bac0-c09875d144dd",
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "estimate_ml_optim (generic function with 1 method)"
                        ]
                    },
                    "execution_count": 18,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "function estimate_ml_optim(y, X, initial_x)\n",
                "\n",
                "    ml_res_optim = optimize(\n",
                "        param -> - calculate_log_likelihood(param, y, X), initial_x\n",
                "    )\n",
                "    \n",
                "    return ml_res_optim.minimizer\n",
                "end"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 19,
            "id": "5c87f89d-797d-45cf-b88b-0c855333e785",
            "metadata": {},
            "outputs": [],
            "source": [
                "est_ml_optim = estimate_ml_optim(y, X, initial_x);\n",
                "βhat_ml_optim = est_ml_optim[1:(K + 1)];\n",
                "σhat_ml_optim = est_ml_optim[K + 2];"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 20,
            "id": "d9d87da6-b2ee-4b89-8772-b2e90535e78b",
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "BenchmarkTools.Trial: 3807 samples with 1 evaluation.\n",
                            " Range \u001b[90m(\u001b[39m\u001b[36m\u001b[1mmin\u001b[22m\u001b[39m … \u001b[35mmax\u001b[39m\u001b[90m):  \u001b[39m\u001b[36m\u001b[1m751.208 μs\u001b[22m\u001b[39m … \u001b[35m25.035 ms\u001b[39m  \u001b[90m┊\u001b[39m GC \u001b[90m(\u001b[39mmin … max\u001b[90m): \u001b[39m 0.00% …  0.00%\n",
                            " Time  \u001b[90m(\u001b[39m\u001b[34m\u001b[1mmedian\u001b[22m\u001b[39m\u001b[90m):     \u001b[39m\u001b[34m\u001b[1m829.042 μs              \u001b[22m\u001b[39m\u001b[90m┊\u001b[39m GC \u001b[90m(\u001b[39mmedian\u001b[90m):    \u001b[39m 0.00%\n",
                            " Time  \u001b[90m(\u001b[39m\u001b[32m\u001b[1mmean\u001b[22m\u001b[39m ± \u001b[32mσ\u001b[39m\u001b[90m):   \u001b[39m\u001b[32m\u001b[1m  1.306 ms\u001b[22m\u001b[39m ± \u001b[32m 1.174 ms\u001b[39m  \u001b[90m┊\u001b[39m GC \u001b[90m(\u001b[39mmean ± σ\u001b[90m):  \u001b[39m25.41% ± 23.56%\n",
                            "\n",
                            "  \u001b[39m█\u001b[34m▇\u001b[39m\u001b[39m▄\u001b[39m▄\u001b[39m▄\u001b[39m▁\u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[32m \u001b[39m\u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m▂\u001b[39m▂\u001b[39m▃\u001b[39m▄\u001b[39m▅\u001b[39m▄\u001b[39m▃\u001b[39m▂\u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m▁\n",
                            "  \u001b[39m█\u001b[34m█\u001b[39m\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m▇\u001b[39m▇\u001b[39m▆\u001b[39m▃\u001b[32m▆\u001b[39m\u001b[39m▇\u001b[39m▄\u001b[39m▆\u001b[39m▅\u001b[39m▃\u001b[39m▃\u001b[39m▃\u001b[39m▆\u001b[39m▅\u001b[39m▅\u001b[39m▆\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m▇\u001b[39m▆\u001b[39m▇\u001b[39m▆\u001b[39m▅\u001b[39m▆\u001b[39m▅\u001b[39m▄\u001b[39m▅\u001b[39m▅\u001b[39m▁\u001b[39m▃\u001b[39m▃\u001b[39m▅\u001b[39m▅\u001b[39m▄\u001b[39m▁\u001b[39m▁\u001b[39m▅\u001b[39m▄\u001b[39m▃\u001b[39m▃\u001b[39m▃\u001b[39m▁\u001b[39m▄\u001b[39m▃\u001b[39m▄\u001b[39m▁\u001b[39m▄\u001b[39m \u001b[39m█\n",
                            "  751 μs\u001b[90m        \u001b[39m\u001b[90mHistogram: \u001b[39m\u001b[90m\u001b[1mlog(\u001b[22m\u001b[39m\u001b[90mfrequency\u001b[39m\u001b[90m\u001b[1m)\u001b[22m\u001b[39m\u001b[90m by time\u001b[39m      3.96 ms \u001b[0m\u001b[1m<\u001b[22m\n",
                            "\n",
                            " Memory estimate\u001b[90m: \u001b[39m\u001b[33m8.19 MiB\u001b[39m, allocs estimate\u001b[90m: \u001b[39m\u001b[33m4577\u001b[39m."
                        ]
                    },
                    "execution_count": 20,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "@benchmark estimate_ml_optim(y, X, initial_x)"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "b702694b-0495-4dfd-8f01-d3a374ff82da",
            "metadata": {},
            "source": [
                "#### Performance comparison: Copy param, not view"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "54cd6a3d-c62f-46f6-8815-2156554e3d99",
            "metadata": {},
            "source": [
                "In `calculate_log_likelihood`, $\\beta$ is extracted from the parameter input with `@view` macro.\n",
                "Without `@view`, Julia creates a copy when making the $\\beta$ variable, which results in additional memory allocations."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 21,
            "id": "1a91984c-1d20-498f-8d66-f96e3640e3ea",
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "estimate_ml_optim_no_view (generic function with 1 method)"
                        ]
                    },
                    "execution_count": 21,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "function calculate_log_likelihood_no_view(param, y, X; K = K, N = N)\n",
                "    β = param[1:(K + 1)] # <======== Not using @view\n",
                "    σ = param[K + 2]\n",
                "    log_likelihood = (\n",
                "        N / 2 * log(1 / (2 * π * σ^2)) -\n",
                "        sum((y - X * β).^2) / (2 * σ^2)\n",
                "    )\n",
                "\n",
                "    return log_likelihood\n",
                "end\n",
                "\n",
                "function estimate_ml_optim_no_view(y, X, initial_x)\n",
                "\n",
                "    ml_res_optim = optimize(\n",
                "        param -> - calculate_log_likelihood_no_view(param, y, X), initial_x\n",
                "    )\n",
                "    \n",
                "    return ml_res_optim.minimizer\n",
                "end"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 22,
            "id": "e5d23ef3-2734-4eac-97ec-0981eb2855e1",
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "BenchmarkTools.Trial: 3741 samples with 1 evaluation.\n",
                            " Range \u001b[90m(\u001b[39m\u001b[36m\u001b[1mmin\u001b[22m\u001b[39m … \u001b[35mmax\u001b[39m\u001b[90m):  \u001b[39m\u001b[36m\u001b[1m767.041 μs\u001b[22m\u001b[39m … \u001b[35m 18.438 ms\u001b[39m  \u001b[90m┊\u001b[39m GC \u001b[90m(\u001b[39mmin … max\u001b[90m): \u001b[39m 0.00% … 64.76%\n",
                            " Time  \u001b[90m(\u001b[39m\u001b[34m\u001b[1mmedian\u001b[22m\u001b[39m\u001b[90m):     \u001b[39m\u001b[34m\u001b[1m874.417 μs               \u001b[22m\u001b[39m\u001b[90m┊\u001b[39m GC \u001b[90m(\u001b[39mmedian\u001b[90m):    \u001b[39m 0.00%\n",
                            " Time  \u001b[90m(\u001b[39m\u001b[32m\u001b[1mmean\u001b[22m\u001b[39m ± \u001b[32mσ\u001b[39m\u001b[90m):   \u001b[39m\u001b[32m\u001b[1m  1.328 ms\u001b[22m\u001b[39m ± \u001b[32m941.587 μs\u001b[39m  \u001b[90m┊\u001b[39m GC \u001b[90m(\u001b[39mmean ± σ\u001b[90m):  \u001b[39m25.65% ± 23.51%\n",
                            "\n",
                            "  \u001b[39m▇\u001b[39m█\u001b[34m▅\u001b[39m\u001b[39m▄\u001b[39m▄\u001b[39m▂\u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[32m \u001b[39m\u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m▁\u001b[39m▂\u001b[39m▄\u001b[39m▄\u001b[39m▄\u001b[39m▃\u001b[39m▂\u001b[39m▁\u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m▁\n",
                            "  \u001b[39m█\u001b[39m█\u001b[34m█\u001b[39m\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m▇\u001b[39m▇\u001b[39m▅\u001b[39m▆\u001b[32m▅\u001b[39m\u001b[39m▆\u001b[39m▅\u001b[39m▅\u001b[39m▅\u001b[39m▅\u001b[39m▆\u001b[39m▃\u001b[39m▅\u001b[39m▃\u001b[39m▆\u001b[39m▇\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m▇\u001b[39m▇\u001b[39m▆\u001b[39m▆\u001b[39m▅\u001b[39m▃\u001b[39m▅\u001b[39m▅\u001b[39m▄\u001b[39m▄\u001b[39m▃\u001b[39m▄\u001b[39m▄\u001b[39m▂\u001b[39m▃\u001b[39m▃\u001b[39m▄\u001b[39m▄\u001b[39m▂\u001b[39m▃\u001b[39m▂\u001b[39m▃\u001b[39m▃\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▃\u001b[39m \u001b[39m█\n",
                            "  767 μs\u001b[90m        \u001b[39m\u001b[90mHistogram: \u001b[39m\u001b[90m\u001b[1mlog(\u001b[22m\u001b[39m\u001b[90mfrequency\u001b[39m\u001b[90m\u001b[1m)\u001b[22m\u001b[39m\u001b[90m by time\u001b[39m        3.9 ms \u001b[0m\u001b[1m<\u001b[22m\n",
                            "\n",
                            " Memory estimate\u001b[90m: \u001b[39m\u001b[33m8.22 MiB\u001b[39m, allocs estimate\u001b[90m: \u001b[39m\u001b[33m5285\u001b[39m."
                        ]
                    },
                    "execution_count": 22,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "@benchmark estimate_ml_optim_no_view(y, X, initial_x)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "bb79e960-e9c8-4ad9-9338-910b06659e2f",
            "metadata": {},
            "outputs": [],
            "source": []
        },
        {
            "cell_type": "markdown",
            "id": "5ddde928-95a4-4e17-9307-4124e0355952",
            "metadata": {},
            "source": [
                "## With `Optim` with a constraint"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "1b3ccb0b-be8b-4ed0-95d1-c034db9c6be3",
            "metadata": {},
            "source": [
                "In the optimization above, there was no constraint for $\\sigma$, which could have returned an error due to negative standard deviations.\n",
                "To prevent this, I add a positive constraint for $\\sigma$ to the model."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 23,
            "id": "1284978a-a9cc-4c26-816b-706c53218624",
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "estimate_ml_optim_constrainted (generic function with 1 method)"
                        ]
                    },
                    "execution_count": 23,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "function estimate_ml_optim_constrainted(y, X, initial_x)\n",
                "\n",
                "    ml_res_optim = optimize(\n",
                "        param -> - calculate_log_likelihood(param, y, X), \n",
                "        vcat(repeat([-Inf], K + 1), [1e-10]),\n",
                "        repeat([Inf], K + 2),\n",
                "        initial_x,\n",
                "        Fminbox(GradientDescent())\n",
                "    )\n",
                "    \n",
                "    return ml_res_optim.minimizer\n",
                "end"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 24,
            "id": "d4c2c2ef-f109-43ae-9033-cf9b49de37be",
            "metadata": {},
            "outputs": [],
            "source": [
                "est_ml_optim_constrained = estimate_ml_optim_constrainted(y, X, initial_x);\n",
                "βhat_ml_optim_constrained = est_ml_optim[1:(K + 1)];\n",
                "σhat_ml_optim_constrained = est_ml_optim[K + 2];"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 25,
            "id": "f3d22cec-b2dd-4f2b-aaa8-722035de5a88",
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "BenchmarkTools.Trial: 942 samples with 1 evaluation.\n",
                            " Range \u001b[90m(\u001b[39m\u001b[36m\u001b[1mmin\u001b[22m\u001b[39m … \u001b[35mmax\u001b[39m\u001b[90m):  \u001b[39m\u001b[36m\u001b[1m3.310 ms\u001b[22m\u001b[39m … \u001b[35m32.614 ms\u001b[39m  \u001b[90m┊\u001b[39m GC \u001b[90m(\u001b[39mmin … max\u001b[90m): \u001b[39m 0.00% … 0.00%\n",
                            " Time  \u001b[90m(\u001b[39m\u001b[34m\u001b[1mmedian\u001b[22m\u001b[39m\u001b[90m):     \u001b[39m\u001b[34m\u001b[1m4.863 ms              \u001b[22m\u001b[39m\u001b[90m┊\u001b[39m GC \u001b[90m(\u001b[39mmedian\u001b[90m):    \u001b[39m24.65%\n",
                            " Time  \u001b[90m(\u001b[39m\u001b[32m\u001b[1mmean\u001b[22m\u001b[39m ± \u001b[32mσ\u001b[39m\u001b[90m):   \u001b[39m\u001b[32m\u001b[1m5.296 ms\u001b[22m\u001b[39m ± \u001b[32m 1.766 ms\u001b[39m  \u001b[90m┊\u001b[39m GC \u001b[90m(\u001b[39mmean ± σ\u001b[90m):  \u001b[39m25.92% ± 6.06%\n",
                            "\n",
                            "  \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m█\u001b[34m█\u001b[39m\u001b[39m▅\u001b[39m▂\u001b[32m▁\u001b[39m\u001b[39m \u001b[39m \u001b[39m \u001b[39m▃\u001b[39m▃\u001b[39m▂\u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \n",
                            "  \u001b[39m█\u001b[39m▇\u001b[39m▁\u001b[39m▄\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m█\u001b[34m█\u001b[39m\u001b[39m█\u001b[39m█\u001b[32m█\u001b[39m\u001b[39m█\u001b[39m▆\u001b[39m▇\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m▇\u001b[39m▇\u001b[39m▆\u001b[39m▅\u001b[39m▆\u001b[39m▄\u001b[39m▅\u001b[39m▄\u001b[39m▄\u001b[39m▅\u001b[39m▅\u001b[39m▅\u001b[39m▄\u001b[39m▁\u001b[39m▁\u001b[39m▄\u001b[39m▁\u001b[39m▄\u001b[39m▁\u001b[39m▄\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▄\u001b[39m▁\u001b[39m▄\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▄\u001b[39m▁\u001b[39m▁\u001b[39m▄\u001b[39m▄\u001b[39m \u001b[39m█\n",
                            "  3.31 ms\u001b[90m      \u001b[39m\u001b[90mHistogram: \u001b[39m\u001b[90m\u001b[1mlog(\u001b[22m\u001b[39m\u001b[90mfrequency\u001b[39m\u001b[90m\u001b[1m)\u001b[22m\u001b[39m\u001b[90m by time\u001b[39m     12.9 ms \u001b[0m\u001b[1m<\u001b[22m\n",
                            "\n",
                            " Memory estimate\u001b[90m: \u001b[39m\u001b[33m33.90 MiB\u001b[39m, allocs estimate\u001b[90m: \u001b[39m\u001b[33m19486\u001b[39m."
                        ]
                    },
                    "execution_count": 25,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "@benchmark estimate_ml_optim_constrainted(y, X, initial_x)"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "93a2cb0c-4f75-458d-afe5-555774b7eaf0",
            "metadata": {},
            "source": [
                "## With JuMP"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "2851faca-c99a-4b3f-bb93-faf200ad8a40",
            "metadata": {},
            "source": [
                "Here, I estimate the parameters with ML with JuMP.\n",
                "As shown below, the run time is a bit slower but pretty comparable to the `Optim` implementation with constraints."
            ]
        },
        {
            "cell_type": "markdown",
            "id": "927b03c6-e598-4850-b04b-aef5ac129345",
            "metadata": {},
            "source": [
                "### Directly specify the ML objective function within JuMP"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "5f0b0cbd-d55e-4a6e-84da-3740c238f24b",
            "metadata": {},
            "source": [
                "First, I directly calculate the objective function within JuMP.\n",
                "This part uses [the code in the official JuMP document](calculate_log_likelihood_no_view) with simplifications."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 26,
            "id": "d3198c54-716b-4efe-9e26-c04e9d162a7e",
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "estimate_ml_jump_direct (generic function with 1 method)"
                        ]
                    },
                    "execution_count": 26,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "function estimate_ml_jump_direct(y, X; K = K, N = N)\n",
                "        \n",
                "    model = Model(Ipopt.Optimizer)\n",
                "    set_silent(model)\n",
                "    @variable(model, β[1:(K + 1)])\n",
                "    @variable(model, σ >= 0.0)\n",
                "    @objective(\n",
                "        model,\n",
                "        Max,\n",
                "        N / 2 * log(1 / (2 * π * σ^2)) -\n",
                "        sum((y[i] - sum(X[i, j] * β[j] for j in 1:(K + 1)))^2 for i in 1:N) / (2 * σ^2)\n",
                "    )\n",
                "    optimize!(model)\n",
                "    @assert is_solved_and_feasible(model)\n",
                "\n",
                "    return vcat(value.(β), value(σ))\n",
                "\n",
                "end"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 27,
            "id": "8f351611-288f-4f98-a2d0-be8a682ee375",
            "metadata": {},
            "outputs": [],
            "source": [
                "est_ml_jump_direct = estimate_ml_jump_direct(y, X);\n",
                "βhat_ml_jump_direct = est_ml_jump_direct[1:(K + 1)];\n",
                "σhat_ml_jump_direct = est_ml_jump_direct[K + 2];"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 28,
            "id": "e9975e54-2c28-46c9-a557-4424f4d88144",
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "BenchmarkTools.Trial: 626 samples with 1 evaluation.\n",
                            " Range \u001b[90m(\u001b[39m\u001b[36m\u001b[1mmin\u001b[22m\u001b[39m … \u001b[35mmax\u001b[39m\u001b[90m):  \u001b[39m\u001b[36m\u001b[1m5.867 ms\u001b[22m\u001b[39m … \u001b[35m41.218 ms\u001b[39m  \u001b[90m┊\u001b[39m GC \u001b[90m(\u001b[39mmin … max\u001b[90m): \u001b[39m0.00% … 49.84%\n",
                            " Time  \u001b[90m(\u001b[39m\u001b[34m\u001b[1mmedian\u001b[22m\u001b[39m\u001b[90m):     \u001b[39m\u001b[34m\u001b[1m7.011 ms              \u001b[22m\u001b[39m\u001b[90m┊\u001b[39m GC \u001b[90m(\u001b[39mmedian\u001b[90m):    \u001b[39m0.00%\n",
                            " Time  \u001b[90m(\u001b[39m\u001b[32m\u001b[1mmean\u001b[22m\u001b[39m ± \u001b[32mσ\u001b[39m\u001b[90m):   \u001b[39m\u001b[32m\u001b[1m7.980 ms\u001b[22m\u001b[39m ± \u001b[32m 2.892 ms\u001b[39m  \u001b[90m┊\u001b[39m GC \u001b[90m(\u001b[39mmean ± σ\u001b[90m):  \u001b[39m4.37% ±  7.41%\n",
                            "\n",
                            "  \u001b[39m \u001b[39m \u001b[39m▄\u001b[39m█\u001b[39m█\u001b[34m▄\u001b[39m\u001b[39m▁\u001b[39m \u001b[39m \u001b[39m \u001b[32m \u001b[39m\u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \n",
                            "  \u001b[39m▄\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m█\u001b[34m█\u001b[39m\u001b[39m█\u001b[39m▄\u001b[39m▅\u001b[39m▅\u001b[32m▅\u001b[39m\u001b[39m▄\u001b[39m▄\u001b[39m▃\u001b[39m▄\u001b[39m▃\u001b[39m▂\u001b[39m▃\u001b[39m▃\u001b[39m▃\u001b[39m▂\u001b[39m▃\u001b[39m▂\u001b[39m▃\u001b[39m▃\u001b[39m▃\u001b[39m▃\u001b[39m▄\u001b[39m▃\u001b[39m▂\u001b[39m▃\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▁\u001b[39m▁\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▁\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▂\u001b[39m▂\u001b[39m▁\u001b[39m▁\u001b[39m▂\u001b[39m \u001b[39m▃\n",
                            "  5.87 ms\u001b[90m        Histogram: frequency by time\u001b[39m        18.5 ms \u001b[0m\u001b[1m<\u001b[22m\n",
                            "\n",
                            " Memory estimate\u001b[90m: \u001b[39m\u001b[33m3.99 MiB\u001b[39m, allocs estimate\u001b[90m: \u001b[39m\u001b[33m38468\u001b[39m."
                        ]
                    },
                    "execution_count": 28,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "@benchmark estimate_ml_jump_direct(y, X)"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "3586cc21-6781-4abf-9a39-77aa4676dd72",
            "metadata": {},
            "source": [
                "### Call the function to calculate log likelihood"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "55727962-ae61-4ab3-b4f1-a5f803ea90a7",
            "metadata": {},
            "source": [
                "Here, I create a function to calculate log-likelihood and call it within JuMP.\n",
                "In JuMP, this kind of function called from JuMP models is called \"user-defined function\".\n",
                "Since user-defined functions do not accept vector inputs ([document](https://jump.dev/JuMP.jl/stable/manual/nonlinear/#User-defined-operators-with-vector-inputs)), I have to use Julia's splatting syntax.\n",
                "Interestingly, using a user-defined function does not increase the run-time, as shown below."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 29,
            "id": "7c0c5e5e-3022-4f9d-b7ba-a930d7e9288b",
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "calculate_log_likelihood_jump (generic function with 1 method)"
                        ]
                    },
                    "execution_count": 29,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "function calculate_log_likelihood_jump(args...; K = K, N = N)\n",
                "\n",
                "    param = args[1]\n",
                "    y = args[2]\n",
                "    X = args[3]\n",
                "\n",
                "    β = @view param[1:(K + 1)]\n",
                "    σ = param[K + 2]\n",
                "    log_likelihood = (\n",
                "        N / 2 * log(1 / (2 * π * σ^2)) -\n",
                "        sum((y - X * β).^2) / (2 * σ^2)\n",
                "    )\n",
                "\n",
                "    return log_likelihood\n",
                "end"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 30,
            "id": "03c18caf-8c6d-4dde-b9cd-81e72f7a65ca",
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "estimate_ml_jump_call_ll (generic function with 1 method)"
                        ]
                    },
                    "execution_count": 30,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "function estimate_ml_jump_call_ll(y, X; K = K, N = N)\n",
                "        \n",
                "    model = Model(Ipopt.Optimizer)\n",
                "    set_silent(model)\n",
                "    @variable(model, param[1:(K + 2)])\n",
                "    set_lower_bound(param[K + 2], 0)\n",
                "    @objective(\n",
                "        model,\n",
                "        Max,\n",
                "        calculate_log_likelihood_jump(param, y, X)\n",
                "    )\n",
                "    optimize!(model)\n",
                "    @assert is_solved_and_feasible(model)\n",
                "\n",
                "    return value.(param)\n",
                "\n",
                "end"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 31,
            "id": "e1bd2d81-96d7-429e-bdce-38dc7cd0aee5",
            "metadata": {},
            "outputs": [],
            "source": [
                "est_ml_jump_call_ll = estimate_ml_jump_call_ll(y, X);\n",
                "βhat_ml_jump_call_ll = est_ml_jump_call_ll[1:(K + 1)];\n",
                "σhat_ml_jump_call_ll = est_ml_jump_call_ll[K + 2];"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 32,
            "id": "dad88bc9-7c4f-4618-8264-823401f8c088",
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "BenchmarkTools.Trial: 623 samples with 1 evaluation.\n",
                            " Range \u001b[90m(\u001b[39m\u001b[36m\u001b[1mmin\u001b[22m\u001b[39m … \u001b[35mmax\u001b[39m\u001b[90m):  \u001b[39m\u001b[36m\u001b[1m5.830 ms\u001b[22m\u001b[39m … \u001b[35m76.622 ms\u001b[39m  \u001b[90m┊\u001b[39m GC \u001b[90m(\u001b[39mmin … max\u001b[90m): \u001b[39m0.00% … 74.09%\n",
                            " Time  \u001b[90m(\u001b[39m\u001b[34m\u001b[1mmedian\u001b[22m\u001b[39m\u001b[90m):     \u001b[39m\u001b[34m\u001b[1m6.945 ms              \u001b[22m\u001b[39m\u001b[90m┊\u001b[39m GC \u001b[90m(\u001b[39mmedian\u001b[90m):    \u001b[39m0.00%\n",
                            " Time  \u001b[90m(\u001b[39m\u001b[32m\u001b[1mmean\u001b[22m\u001b[39m ± \u001b[32mσ\u001b[39m\u001b[90m):   \u001b[39m\u001b[32m\u001b[1m8.020 ms\u001b[22m\u001b[39m ± \u001b[32m 3.937 ms\u001b[39m  \u001b[90m┊\u001b[39m GC \u001b[90m(\u001b[39mmean ± σ\u001b[90m):  \u001b[39m4.88% ±  7.58%\n",
                            "\n",
                            "  \u001b[39m \u001b[39m \u001b[39m█\u001b[39m█\u001b[39m█\u001b[34m▅\u001b[39m\u001b[39m▁\u001b[39m \u001b[39m \u001b[39m \u001b[32m \u001b[39m\u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \n",
                            "  \u001b[39m▅\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m█\u001b[34m█\u001b[39m\u001b[39m█\u001b[39m▆\u001b[39m▆\u001b[39m▅\u001b[32m▄\u001b[39m\u001b[39m▄\u001b[39m▃\u001b[39m▃\u001b[39m▄\u001b[39m▃\u001b[39m▄\u001b[39m▃\u001b[39m▃\u001b[39m▃\u001b[39m▃\u001b[39m▃\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▁\u001b[39m▁\u001b[39m▂\u001b[39m▂\u001b[39m▃\u001b[39m▃\u001b[39m▃\u001b[39m▃\u001b[39m▃\u001b[39m▃\u001b[39m▃\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▃\u001b[39m▂\u001b[39m▁\u001b[39m▁\u001b[39m▂\u001b[39m▁\u001b[39m▂\u001b[39m▁\u001b[39m▁\u001b[39m▂\u001b[39m▁\u001b[39m▁\u001b[39m▂\u001b[39m▂\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▃\u001b[39m \u001b[39m▃\n",
                            "  5.83 ms\u001b[90m        Histogram: frequency by time\u001b[39m        18.2 ms \u001b[0m\u001b[1m<\u001b[22m\n",
                            "\n",
                            " Memory estimate\u001b[90m: \u001b[39m\u001b[33m4.01 MiB\u001b[39m, allocs estimate\u001b[90m: \u001b[39m\u001b[33m38495\u001b[39m."
                        ]
                    },
                    "execution_count": 32,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "@benchmark estimate_ml_jump_call_ll(y, X)"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "8634077b-5b19-41c7-bea7-d2e0706091f9",
            "metadata": {},
            "source": [
                "## Estimates comparison"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 33,
            "id": "fc7e20df-80f1-41d8-89e0-dda971dd02f8",
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "4×5 Matrix{Float64}:\n",
                            " -0.627644  -0.627644  -0.627644  -0.627644  -0.663764\n",
                            "  1.36376    1.36376    1.36376    1.36376    1.36321\n",
                            "  0.756495   0.756495   0.756496   0.756496   0.717411\n",
                            "  0.660767   0.660767   0.660766   0.660766   0.678743"
                        ]
                    },
                    "execution_count": 33,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "[βhat_ml_optim βhat_ml_optim_constrained βhat_ml_jump_direct βhat_ml_jump_call_ll β]"
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "Julia 1.11.1",
            "language": "julia",
            "name": "julia-1.11"
        },
        "language_info": {
            "file_extension": ".jl",
            "mimetype": "application/julia",
            "name": "julia",
            "version": "1.11.1"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 5
}