{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Structured prediction\n",
    "\n",
    "\n",
    "\n",
    "In this example$\\newcommand{\\reals}{\\mathbf{R}}$$\\newcommand{\\ones}{\\mathbf{1}}$, we fit a regression model to structured data, using an LLCP.\n",
    "The training dataset $\\mathcal D$ contains $N$ input-output pairs $(x, y)$,\n",
    "where $x \\in \\reals^{n}_{++}$ is an input and $y \\in \\reals^{m}_{++}$ is an\n",
    "outputs. The entries of each output $y$ are sorted in ascending order, meaning\n",
    "$y_1 \\leq y_2 \\leq \\cdots y_m$.\n",
    "\n",
    "Our regression model $\\phi : \\reals^{n}_{++} \\to \\reals^{m}_{++}$ takes as\n",
    "input a vector $x \\in \\reals^{n}_{++}$, and solves an LLCP to produce a\n",
    "prediction $\\hat y \\in \\reals^{m}_{++}$. In particular, the solution of the\n",
    "LLCP is model's prediction. The model is of the form\n",
    "$$\n",
    "\\begin{equation}\n",
    "\\begin{array}{lll}\n",
    "\\phi(x) = &\n",
    "\\mbox{argmin} & \\ones^T (z/y + y / z) \\\\\n",
    "& \\mbox{subject to} &  y_i \\leq y_{i+1}, \\quad i=1, \\ldots, m-1 \\\\\n",
    "&& z_i = c_i x_1^{A_{i1}}x_2^{A_{i2}}\\cdots x_n^{A_{in}}, \\quad i = 1, \\ldots, m.\n",
    "\\end{array}\\label{e-model}\n",
    "\\end{equation}\n",
    "$$\n",
    "Here, the minimization is over $y \\in \\reals^{m}_{++}$ and an auxiliary\n",
    "variable $z \\in \\reals^{m}_{++}$, $\\phi(x)$ is the optimal value of $y$, and\n",
    "the parameters are $c \\in \\reals^{m}_{++}$ and $A \\in \\reals^{m \\times n}$. The\n",
    "ratios in the objective are meant elementwise, as is the inequality $y \\leq z$, and\n",
    "$\\ones$ denotes the vector of all ones. Given a vector $x$, this model finds a\n",
    "sorted vector $\\hat y$ whose entries are close to monomial functions of $x$\n",
    "(which are the entries of $z$), as measured by the fractional error.\n",
    "\n",
    "\n",
    "The training loss\n",
    "$\\mathcal{L}(\\phi)$ of the model on the training set is the mean squared loss\n",
    "$$\n",
    "\\mathcal{L}(\\phi) = \\frac{1}{N}\\sum_{(x, y) \\in \\mathcal D} \\|y - \\phi(x)\\|_2^2.\n",
    "$$\n",
    "We emphasize that $\\mathcal{L}(\\phi)$ depends on $c$ and $A$.\n",
    "In this example, we fit the parameters $c$ and $A$ in the LLCP \n",
    "to minimize the training loss $\\mathcal{L}(\\phi)$.\n",
    "\n",
    "**Fitting.** We fit the parameters by an iterative projected\n",
    "gradient descent method on $\\mathcal L(\\phi)$. In each iteration, we first\n",
    "compute predictions $\\phi(x)$ for each input in the training set; this requires\n",
    "solving $N$ LLCPs. Next, we evaluate the training loss $\\mathcal L(\\phi)$. To\n",
    "update the parameters, we compute the gradient $\\nabla \\mathcal L(\\phi)$ of the\n",
    "training loss with respect to the parameters $c$ and $A$. This\n",
    "requires differentiating through the solution map of the LLCP. We can compute this gradient efficiently, using the ``backward`` method in CVXPY\n",
    "(or CVXPY Layers). Finally, we subtract\n",
    "a small multiple of the gradient from the parameters. Care must be taken to\n",
    "ensure that $c$ is strictly positive; this can be done by clamping the entries\n",
    "of $c$ at some small threshold slightly above zero. We run this method for\n",
    "a fixed number of iterations.\n",
    "\n",
    "This example is described in the paper [Differentiating through Log-Log Convex Programs](http://web.stanford.edu/~boyd/papers/pdf/diff_llcvx.pdf).\n",
    "\n",
    "Shane Barratt formulated the idea of using an optimization layer to regress on sorted vectors.\n",
    "\n",
    "**Requirements.**\n",
    "This example requires PyTorch and CvxpyLayers >= v0.1.3."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from cvxpylayers.torch import CvxpyLayer\n",
    "\n",
    "\n",
    "import cvxpy as cp\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "torch.set_default_tensor_type(torch.DoubleTensor)\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 20\n",
    "m = 10\n",
    "\n",
    "# Number of training input-output pairs\n",
    "N = 100\n",
    "\n",
    "# Number of validation pairs\n",
    "N_val = 50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.random.manual_seed(243)\n",
    "np.random.seed(243)\n",
    "\n",
    "normal = torch.distributions.multivariate_normal.MultivariateNormal(torch.zeros(n), torch.eye(n))\n",
    "lognormal = lambda batch: torch.exp(normal.sample(torch.tensor([batch])))\n",
    "\n",
    "A_true = torch.randn((m, n)) / 10\n",
    "c_true = np.abs(torch.randn(m))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_data(num_points, seed):\n",
    "    torch.random.manual_seed(seed)\n",
    "    np.random.seed(seed)\n",
    "    \n",
    "    latent = lognormal(num_points)\n",
    "    noise = lognormal(num_points)\n",
    "    inputs = noise + latent\n",
    "\n",
    "    input_cp = cp.Parameter(pos=True, shape=(n,))\n",
    "    prediction = cp.multiply(c_true.numpy(), cp.gmatmul(A_true.numpy(), input_cp))\n",
    "    y = cp.Variable(pos=True, shape=(m,))\n",
    "    objective_fn = cp.sum(prediction / y + y/prediction)\n",
    "    constraints = []\n",
    "    for i in range(m-1):\n",
    "        constraints += [y[i] <= y[i+1]]\n",
    "    problem = cp.Problem(cp.Minimize(objective_fn), constraints)\n",
    "    \n",
    "    outputs = []\n",
    "    for i in range(num_points):\n",
    "        input_cp.value = inputs[i, :].numpy()\n",
    "        problem.solve(cp.SCS, gp=True)\n",
    "        outputs.append(y.value)\n",
    "    return inputs, torch.stack([torch.tensor(t) for t in outputs])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x12b367cd0>]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 9036 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "train_inputs, train_outputs = generate_data(N, 243)\n",
    "plt.plot(train_outputs[0, :].numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x12da7e410>]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 8272 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "val_inputs, val_outputs = generate_data(N_val, 0)\n",
    "plt.plot(val_outputs[0, :].numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Monomial fit to each component\n",
    "\n",
    "We will initialize the parameters in our LLCP model by fitting monomials to the training data, without enforcing the monotonicity constraint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "log_c = cp.Variable(shape=(m,1))\n",
    "theta = cp.Variable(shape=(n, m))\n",
    "inputs_np = train_inputs.numpy()\n",
    "log_outputs_np = np.log(train_outputs.numpy()).T\n",
    "log_inputs_np = np.log(inputs_np).T\n",
    "offsets = cp.hstack([log_c]*N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "cp_preds = theta.T @ log_inputs_np + offsets\n",
    "objective_fn = (1/N) * cp.sum_squares(cp_preds - log_outputs_np)\n",
    "lstq_problem = cp.Problem(cp.Minimize(objective_fn))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lstq_problem.is_dcp()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-----------------------------------------------------------------\n",
      "           OSQP v0.6.0  -  Operator Splitting QP Solver\n",
      "              (c) Bartolomeo Stellato,  Goran Banjac\n",
      "        University of Oxford  -  Stanford University 2019\n",
      "-----------------------------------------------------------------\n",
      "problem:  variables n = 1210, constraints m = 1000\n",
      "          nnz(P) + nnz(A) = 23000\n",
      "settings: linear system solver = qdldl,\n",
      "          eps_abs = 1.0e-05, eps_rel = 1.0e-05,\n",
      "          eps_prim_inf = 1.0e-04, eps_dual_inf = 1.0e-04,\n",
      "          rho = 1.00e-01 (adaptive),\n",
      "          sigma = 1.00e-06, alpha = 1.60, max_iter = 10000\n",
      "          check_termination: on (interval 25),\n",
      "          scaling: on, scaled_termination: off\n",
      "          warm start: on, polish: on, time_limit: off\n",
      "\n",
      "iter   objective    pri res    dua res    rho        time\n",
      "   1   0.0000e+00   3.30e+00   1.22e+04   1.00e-01   3.06e-03s\n",
      "  50   1.0014e-02   1.72e-07   1.64e-07   1.75e-03   7.37e-03s\n",
      "plsh   1.0014e-02   1.56e-15   1.17e-14   --------   9.68e-03s\n",
      "\n",
      "status:               solved\n",
      "solution polish:      successful\n",
      "number of iterations: 50\n",
      "optimal objective:    0.0100\n",
      "run time:             9.68e-03s\n",
      "optimal rho estimate: 8.77e-05\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0.010014212812318733"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lstq_problem.solve(verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "c = torch.exp(torch.tensor(log_c.value)).squeeze()\n",
    "lstsq_val_preds = []\n",
    "for i in range(N_val):\n",
    "    inp = val_inputs[i, :].numpy()\n",
    "    pred = cp.multiply(c,cp.gmatmul(theta.T.value, inp))\n",
    "    lstsq_val_preds.append(pred.value)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fitting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A_param = cp.Parameter(shape=(m, n))\n",
    "c_param = cp.Parameter(pos=True, shape=(m,))\n",
    "x_slack = cp.Variable(pos=True, shape=(n,))\n",
    "x_param = cp.Parameter(pos=True, shape=(n,))\n",
    "y = cp.Variable(pos=True, shape=(m,))\n",
    "\n",
    "prediction = cp.multiply(c_param, cp.gmatmul(A_param, x_slack))\n",
    "objective_fn = cp.sum(prediction / y + y / prediction)\n",
    "constraints = [x_slack == x_param]\n",
    "for i in range(m-1):\n",
    "    constraints += [y[i] <= y[i+1]]\n",
    "problem = cp.Problem(cp.Minimize(objective_fn), constraints)\n",
    "problem.is_dgp(dpp=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "A_param.value = np.random.randn(m, n)\n",
    "x_param.value = np.abs(np.random.randn(n))\n",
    "c_param.value = np.abs(np.random.randn(m))\n",
    "\n",
    "layer = CvxpyLayer(problem, parameters=[A_param, c_param, x_param], variables=[y], gp=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(epoch 0) train / val (0.0018 / 0.0014) \n",
      "(epoch 1) train / val (0.0017 / 0.0014) \n",
      "(epoch 2) train / val (0.0017 / 0.0014) \n",
      "(epoch 3) train / val (0.0017 / 0.0014) \n",
      "(epoch 4) train / val (0.0017 / 0.0014) \n",
      "(epoch 5) train / val (0.0017 / 0.0014) \n",
      "(epoch 6) train / val (0.0016 / 0.0014) \n",
      "(epoch 7) train / val (0.0016 / 0.0014) \n",
      "(epoch 8) train / val (0.0016 / 0.0014) \n",
      "(epoch 9) train / val (0.0016 / 0.0014) \n"
     ]
    }
   ],
   "source": [
    "torch.random.manual_seed(1)\n",
    "A_tch = torch.tensor(theta.T.value)\n",
    "A_tch.requires_grad_(True)\n",
    "c_tch = torch.tensor(np.squeeze(np.exp(log_c.value)))\n",
    "c_tch.requires_grad_(True)\n",
    "train_losses = []\n",
    "val_losses = []\n",
    "\n",
    "lam1 = torch.tensor(1e-1)\n",
    "lam2 = torch.tensor(1e-1)\n",
    "\n",
    "opt = torch.optim.SGD([A_tch, c_tch], lr=5e-2)\n",
    "for epoch in range(10):\n",
    "    preds = layer(A_tch, c_tch, train_inputs, solver_args={'acceleration_lookback': 0})[0]\n",
    "    loss = (preds - train_outputs).pow(2).sum(axis=1).mean(axis=0)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        val_preds = layer(A_tch, c_tch, val_inputs, solver_args={'acceleration_lookback': 0})[0]\n",
    "        val_loss = (val_preds - val_outputs).pow(2).sum(axis=1).mean(axis=0)\n",
    "\n",
    "    print('(epoch {0}) train / val ({1:.4f} / {2:.4f}) '.format(epoch, loss, val_loss))\n",
    "    train_losses.append(loss.item())\n",
    "    val_losses.append(val_loss.item())\n",
    "    \n",
    "    opt.zero_grad()\n",
    "    loss.backward()\n",
    "    opt.step()\n",
    "    with torch.no_grad():\n",
    "        c_tch = torch.max(c_tch, torch.tensor(1e-8))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    train_preds_tch = layer(A_tch, c_tch, train_inputs)[0]\n",
    "    train_preds = [t.detach().numpy() for t in train_preds_tch]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    val_preds_tch = layer(A_tch, c_tch, val_inputs)[0]\n",
    "    val_preds = [t.detach().numpy() for t in val_preds_tch]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 15160 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure()\n",
    "\n",
    "\n",
    "i = 0\n",
    "plt.plot(val_preds[i], label='LLCP', color='teal')\n",
    "plt.plot(lstsq_val_preds[i], label='least squares', linestyle='--', color='gray')\n",
    "plt.plot(val_outputs[i], label='true', linestyle='-.', color='orange')\n",
    "w, h = 8, 3.5\n",
    "plt.xlabel(r'$i$')\n",
    "plt.ylabel(r'$y_i$')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
