{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "provenance": []
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Day 26: Constrained Optimisation & Lagrange Multipliers\n",
    "\n",
    "*Beginner to Ninja: Math for Causal Inference*\n",
    "\n",
    "Optimising subject to budgets and adding-up constraints, the shadow price, and the experimental-design and weighting problems that use it.\n",
    "\n",
    "**Learning outcomes**\n",
    "\n",
    "- Set up a Lagrangian, solve the first-order conditions and interpret the multiplier as a shadow price.\n",
    "- Prove the envelope theorem and handle simple inequality constraints (KKT).\n",
    "- Derive consumer demand, cost-minimising inputs and minimum-variance weights.\n",
    "- Derive optimal sample allocation across treatment arms and strata, and the form of entropy-balancing weights."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Theory\n",
    "\n",
    "### The Lagrangian\n",
    "\n",
    "To maximise or minimise $f(x)$ subject to $g(x) = c$, form\n",
    "\n",
    "$$\n",
    "\\mathcal{L}(x, \\lambda) = f(x) - \\lambda\\,(g(x) - c)\n",
    "$$\n",
    "\n",
    "and solve $\\nabla_x\\mathcal{L} = 0$ together with the constraint. At the optimum, $\\nabla f = \\lambda\\nabla g$: the objective's gradient is parallel to the constraint's, so no feasible direction improves $f$ (geometrically, a level curve is tangent to the constraint). With several constraints, use one multiplier per constraint.\n",
    "\n",
    "### The shadow price\n",
    "\n",
    "Let $V(c)$ be the optimal value. The **envelope theorem** says\n",
    "\n",
    "$$\n",
    "\\frac{dV}{dc} = \\lambda^*,\n",
    "$$\n",
    "\n",
    "so the multiplier is the marginal value of relaxing the constraint: the marginal utility of income, the marginal cost of output, the value of one more unit of budget.\n",
    "\n",
    "### Inequality constraints (KKT)\n",
    "\n",
    "For $\\min f(x)$ subject to $h(x) \\le 0$: $\\nabla f + \\mu\\nabla h = 0$, $\\mu \\ge 0$, $h(x) \\le 0$ and **complementary slackness** $\\mu\\,h(x) = 0$. Either the constraint binds, or its multiplier is zero.\n",
    "\n",
    "### Where constraints appear in causal inference\n",
    "\n",
    "- Allocating a fixed budget across treatment and control arms, or across strata.\n",
    "- **Synthetic control:** weights that are non-negative and sum to one (Day 85).\n",
    "- **Entropy balancing / calibration:** weights closest to uniform that exactly match covariate means (Day 84).\n",
    "- Restricted least squares: imposing $R\\beta = r$.\n",
    "\n",
    "## Part 1 · Warm-up"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 1 · A classic\n",
    "\n",
    "Maximise $xy$ subject to $x + y = 10$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "$\\mathcal L = xy - \\lambda(x + y - 10)$: $y = \\lambda$, $x = \\lambda$, so $x = y = 5$, $\\lambda = 5$, value 25.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 2 · Closest point on a line\n",
    "\n",
    "Minimise $x^2 + y^2$ subject to $x + 2y = 5$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Lagrangian: $\\mathcal L = x^2 + y^2 - \\lambda(x + 2y - 5)$.\n",
    "\n",
    "**Step 2.** FOCs: $2x = \\lambda$ and $2y = 2\\lambda$, so $y = 2x$.\n",
    "\n",
    "**Step 3.** Constraint: $x + 2(2x) = 5 \\Rightarrow x = 1$, $y = 2$, $\\lambda = 2$.\n",
    "\n",
    "**Answer:** $(1, 2)$ with value 5: the point on the line closest to the origin (a projection, Day 20).\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 3 · Cobb–Douglas demand\n",
    "\n",
    "Maximise $U = x^{0.3}y^{0.7}$ with prices $p_x = 2$, $p_y = 7$ and income $M = 100$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Cobb–Douglas spends the share $\\alpha$ of income on $x$ and $1 - \\alpha$ on $y$ (Exercise 6).\n",
    "\n",
    "**Step 2.** $x = 0.3 \\times 100/2 = 15$.\n",
    "\n",
    "**Step 3.** $y = 0.7 \\times 100/7 = 10$.\n",
    "\n",
    "**Answer:** 15 units of $x$ and 10 of $y$ (spending \\$30 and \\$70).\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 4 · Shadow price check\n",
    "\n",
    "In Exercise 1 with constraint $x + y = c$, compute $V(c)$ and verify $dV/dc = \\lambda$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "$x = y = c/2$, $V(c) = c^2/4$, $dV/dc = c/2 = 5 = \\lambda$ at $c = 10$. ✓\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 5 · A binding inequality\n",
    "\n",
    "Minimise $(x - 3)^2$ subject to $x \\le 2$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "The unconstrained minimum $x = 3$ is infeasible, so the constraint binds: $x = 2$. KKT: $2(x - 3) + \\mu = 0 \\Rightarrow \\mu = 2 \\ge 0$. ✓\n",
    "\n",
    "</details>\n",
    "\n",
    "## Part 2 · Proofs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 6 · Cobb–Douglas demand\n",
    "\n",
    "Derive the demand for $x$ when maximising $x^{\\alpha}y^{1-\\alpha}$ subject to $p_xx + p_yy = M$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Maximise the log, $\\alpha\\ln x + (1 - \\alpha)\\ln y$.\n",
    "\n",
    "**Step 2.** FOCs: $\\alpha/x = \\lambda p_x$ and $(1 - \\alpha)/y = \\lambda p_y$.\n",
    "\n",
    "**Step 3.** So $p_xx = \\alpha/\\lambda$ and $p_yy = (1 - \\alpha)/\\lambda$; adding, $M = 1/\\lambda$.\n",
    "\n",
    "**Step 4.** Hence $x = \\alpha M/p_x$.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 7 · The envelope theorem\n",
    "\n",
    "Prove $dV/dc = \\lambda^*$ for $V(c) = \\max f(x)$ s.t. $g(x) = c$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** $V(c) = f(x^*(c))$, so $V'(c) = \\nabla f^{\\top}\\frac{dx^*}{dc} = \\lambda^*\\nabla g^{\\top}\\frac{dx^*}{dc}$ using the FOC.\n",
    "\n",
    "**Step 2.** Differentiating the constraint $g(x^*(c)) = c$ gives $\\nabla g^{\\top}\\frac{dx^*}{dc} = 1$.\n",
    "\n",
    "**Step 3.** So $V'(c) = \\lambda^*$.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 8 · Minimum-variance weights\n",
    "\n",
    "Minimise $w^{\\top}\\Sigma w$ subject to $\\mathbf 1^{\\top}w = 1$ ($\\Sigma$ PD).\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Lagrangian: $\\mathcal L = w^{\\top}\\Sigma w - \\lambda(\\mathbf 1^{\\top}w - 1)$.\n",
    "\n",
    "**Step 2.** FOC (Day 24 rules, $\\Sigma$ symmetric): $2\\Sigma w - \\lambda\\mathbf 1 = 0 \\Rightarrow w = \\frac\\lambda2\\Sigma^{-1}\\mathbf 1$.\n",
    "\n",
    "**Step 3.** Impose the constraint: $\\mathbf 1^{\\top}w = \\frac\\lambda2\\mathbf 1^{\\top}\\Sigma^{-1}\\mathbf 1 = 1 \\Rightarrow \\frac\\lambda2 = \\frac{1}{\\mathbf 1^{\\top}\\Sigma^{-1}\\mathbf 1}$.\n",
    "\n",
    "**Step 4.** Substitute back:\n",
    "\n",
    "$$\n",
    "w^* = \\frac{\\Sigma^{-1}\\mathbf 1}{\\mathbf 1^{\\top}\\Sigma^{-1}\\mathbf 1}.\n",
    "$$\n",
    "\n",
    "**Step 5.** With independent estimates, $\\Sigma = \\operatorname{diag}(v_j)$ and $w_j^* \\propto 1/v_j$: inverse-variance weighting (Day 8), the efficient way to pool estimates.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 9 · Restricted least squares\n",
    "\n",
    "Minimise $\\|y - Xb\\|^2$ subject to $Rb = r$ and show\n",
    "\n",
    "$$\n",
    "\\tilde b = \\hat\\beta - (X^{\\top}X)^{-1}R^{\\top}\\big[R(X^{\\top}X)^{-1}R^{\\top}\\big]^{-1}(R\\hat\\beta - r).\n",
    "$$\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Lagrangian (the factor 2 on $\\lambda$ keeps the algebra tidy): $\\mathcal L = \\|y - Xb\\|^2 + 2\\lambda^{\\top}(Rb - r)$.\n",
    "\n",
    "**Step 2.** FOC in $b$: $-2X^{\\top}(y - Xb) + 2R^{\\top}\\lambda = 0$, i.e. $X^{\\top}Xb = X^{\\top}y - R^{\\top}\\lambda$.\n",
    "\n",
    "**Step 3.** Solve for $b$: $\\tilde b = \\hat\\beta - (X^{\\top}X)^{-1}R^{\\top}\\lambda$, where $\\hat\\beta$ is unrestricted OLS.\n",
    "\n",
    "**Step 4.** Impose $R\\tilde b = r$: $R\\hat\\beta - R(X^{\\top}X)^{-1}R^{\\top}\\lambda = r$.\n",
    "\n",
    "**Step 5.** Solve for the multiplier: $\\lambda = \\big[R(X^{\\top}X)^{-1}R^{\\top}\\big]^{-1}(R\\hat\\beta - r)$.\n",
    "\n",
    "**Step 6.** Substitute into Step 3:\n",
    "\n",
    "$$\n",
    "\\tilde b = \\hat\\beta - (X^{\\top}X)^{-1}R^{\\top}\\big[R(X^{\\top}X)^{-1}R^{\\top}\\big]^{-1}(R\\hat\\beta - r).\n",
    "$$\n",
    "\n",
    "The correction is large when $\\hat\\beta$ badly violates the restriction: the idea behind Wald and LR tests.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 10 · Optimal allocation between arms\n",
    "\n",
    "With outcome sds $\\sigma_1, \\sigma_0$ and per-unit costs $c_1, c_0$, minimise $\\operatorname{Var}(\\hat\\tau) = \\sigma_1^2/n_1 + \\sigma_0^2/n_0$ subject to $c_1n_1 + c_0n_0 = B$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** FOCs: $\\sigma_1^2/n_1^2 = \\lambda c_1$ and $\\sigma_0^2/n_0^2 = \\lambda c_0$.\n",
    "\n",
    "**Step 2.** Dividing,\n",
    "\n",
    "$$\n",
    "\\frac{n_1}{n_0} = \\frac{\\sigma_1}{\\sigma_0}\\sqrt{\\frac{c_0}{c_1}}.\n",
    "$$\n",
    "\n",
    "**Step 3.** Put more units where outcomes are noisier and where they are cheaper.\n",
    "\n",
    "**Step 4.** With equal costs and variances, split evenly.\n",
    "\n",
    "</details>\n",
    "\n",
    "## Part 3 · Applied"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 11 · An expensive treatment\n",
    "\n",
    "Treatment costs \\$40 per unit, control \\$10, the budget is \\$12,000 and outcome variances are equal. How many units go to each arm?\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Optimal ratio (Exercise 10) with equal variances: $n_1/n_0 = \\sqrt{c_0/c_1} = \\sqrt{10/40} = 0.5$.\n",
    "\n",
    "**Step 2.** Budget: $40n_1 + 10n_0 = 12{,}000$ with $n_1 = 0.5n_0$: $20n_0 + 10n_0 = 30n_0 = 12{,}000$.\n",
    "\n",
    "**Step 3.** $n_0 = 400$, $n_1 = 200$. Check: $40(200) + 10(400) = 12{,}000$. ✓\n",
    "\n",
    "**Answer:** 200 treated and 400 controls.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 12 · Neyman allocation across strata\n",
    "\n",
    "Two strata have sizes 6,000 and 4,000 and outcome sds 10 and 30. Allocate a sample of 300 to minimise the variance of the stratified mean (allocation $n_h \\propto N_h\\sigma_h$).\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Allocation weights $N_h\\sigma_h$: $6{,}000 \\times 10 = 60{,}000$ and $4{,}000 \\times 30 = 120{,}000$.\n",
    "\n",
    "**Step 2.** Shares: $60{,}000/180{,}000 = 1/3$ and $2/3$.\n",
    "\n",
    "**Step 3.** Sample: $300/3 = 100$ and $200$.\n",
    "\n",
    "**Answer:** 100 from the large, homogeneous stratum and 200 from the smaller but noisier one.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 13 · Cost minimisation\n",
    "\n",
    "Produce $q = K^{0.5}L^{0.5} = 10$ at minimum cost with $w = 4$, $r = 1$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Tangency (MRTS = input-price ratio): $\\frac{MP_L}{MP_K} = \\frac{0.5K^{0.5}L^{-0.5}}{0.5K^{-0.5}L^{0.5}} = \\frac KL = \\frac wr = 4$, so $K = 4L$.\n",
    "\n",
    "**Step 2.** Output constraint: $\\sqrt{4L \\cdot L} = 2L = 10 \\Rightarrow L = 5$, $K = 20$.\n",
    "\n",
    "**Step 3.** Cost: $4(5) + 1(20) = 40$.\n",
    "\n",
    "**Answer:** 5 units of labour, 20 of capital, minimum cost 40.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 14 · Synthetic-control weights are not always unique\n",
    "\n",
    "A treated region has pre-period outcome 3; donor regions have 1, 5 and 6. Find two different weight vectors with $w \\ge 0$, $\\sum w = 1$ that reproduce 3 exactly.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** We need $w_1(1) + w_2(5) + w_3(6) = 3$ with $w \\ge 0$ and $w_1 + w_2 + w_3 = 1$.\n",
    "\n",
    "**Step 2.** Try $w = (0.5, 0.5, 0)$: $0.5 + 2.5 = 3$. ✓\n",
    "\n",
    "**Step 3.** Try $w = (0.6, 0, 0.4)$: $0.6 + 2.4 = 3$. ✓\n",
    "\n",
    "**Answer:** both fit perfectly. With one pre-period outcome, many weight vectors match exactly, which is why synthetic control matches many pre-periods and why penalised versions exist.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 15 · Entropy-balancing weights\n",
    "\n",
    "Minimise $\\sum_iw_i\\ln(w_i/q_i)$ subject to $\\sum_iw_i = 1$ and $\\sum_iw_ix_i = \\bar x_{\\text{target}}$. Show the solution has the form $w_i \\propto q_i\\,e^{\\lambda x_i}$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Lagrangian:\n",
    "\n",
    "$$\n",
    "\\mathcal L = \\sum_iw_i\\ln(w_i/q_i) - \\lambda_0\\Big(\\sum_iw_i - 1\\Big) - \\lambda\\Big(\\sum_iw_ix_i - \\bar x_{\\text{target}}\\Big).\n",
    "$$\n",
    "\n",
    "**Step 2.** FOC for $w_i$: $\\ln(w_i/q_i) + 1 - \\lambda_0 - \\lambda x_i = 0$.\n",
    "\n",
    "**Step 3.** Solve: $w_i = q_i\\,e^{\\lambda_0 - 1}\\,e^{\\lambda x_i}$; the factor $e^{\\lambda_0 - 1}$ is the same for everyone, so $w_i \\propto q_ie^{\\lambda x_i}$.\n",
    "\n",
    "**Step 4.** Choose $\\lambda$ so the weighted mean of $x$ equals the target, and $\\lambda_0$ so the weights sum to one.\n",
    "\n",
    "**Answer:** entropy-balancing weights tilt the base weights exponentially in the covariate (Day 84).\n",
    "\n",
    "</details>\n",
    "\n",
    "## Verify in Python"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy import optimize\n",
    "\n",
    "res = optimize.minimize(lambda v: -(v[0]**0.3 * v[1]**0.7), [10, 10],\n",
    "                        constraints={\"type\": \"eq\", \"fun\": lambda v: 2 * v[0] + 7 * v[1] - 100})\n",
    "print(\"Cobb-Douglas demand:\", res.x.round(3))\n",
    "\n",
    "Sig = np.array([[100, 60], [60, 400.]]); one = np.ones(2)\n",
    "w = np.linalg.solve(Sig, one); w /= w.sum()\n",
    "print(\"minimum-variance weights:\", w.round(4))\n",
    "\n",
    "n0 = 12_000 / (40 * 0.5 + 10); print(\"arm sizes (n1, n0):\", n0 * 0.5, n0)\n",
    "N, s = np.array([6000, 4000]), np.array([10, 30]); print(\"Neyman allocation:\", 300 * N * s / (N * s).sum())\n",
    "\n",
    "# Entropy balancing: tilt uniform weights on x so the weighted mean is 4\n",
    "x = np.array([1, 2, 3, 4, 5, 6.]); q = np.ones(6) / 6\n",
    "lam = optimize.brentq(lambda l: (q * np.exp(l * x) * x).sum() / (q * np.exp(l * x)).sum() - 4, -5, 5)\n",
    "wts = q * np.exp(lam * x); wts /= wts.sum()\n",
    "print(\"entropy-balancing weights:\", wts.round(3), \" weighted mean:\", round(wts @ x, 4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Key takeaways\n",
    "\n",
    "- At a constrained optimum, the objective's gradient is a multiple of the constraint's; the multiple is the shadow price.\n",
    "- Budget allocation in experiments and stratified samples is a Lagrange problem: favour noisy and cheap arms.\n",
    "- Weighting estimators (minimum variance, synthetic control, entropy balancing) are constrained optimisation problems."
   ]
  }
 ]
}