logistic_SCAD_MCP_np_MODULE.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Set up the class fundementals "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-08-09T22:25:15.841045Z",
     "start_time": "2022-08-09T22:25:15.568621Z"
    }
   },
   "outputs": [],
   "source": [
    "import warnings\n",
    "\n",
    "warnings.filterwarnings('ignore')\n",
    "# import os, sys\n",
    "# import collections\n",
    "import numpy as _np\n",
    "# import matplotlib.markers as markers\n",
    "# import matplotlib.pyplot as plt\n",
    "# import timeit\n",
    "# import collections\n",
    "# from scipy.linalg import toeplitz, block_diag\n",
    "# from scipy.stats import median_abs_deviation as mad\n",
    "# import multiprocessing\n",
    "# import cProfile\n",
    "# import itertools\n",
    "from numba import jit as _jit\n",
    "from numba import njit as _njit\n",
    "# import warnings\n",
    "# warnings.filterwarnings('ignore') # this is just to hide all the warnings\n",
    "# import rpy2.robjects as robjects\n",
    "# import matplotlib.pyplot as plt # change font globally to Times\n",
    "# plt.style.use('ggplot')\n",
    "# plt.rcParams.update({\n",
    "#     \"text.usetex\": True,\n",
    "#     \"font.family\": \"Times New Roman\",\n",
    "#     \"font.sans-serif\": [\"Times New Roman\"],\n",
    "#     \"font.size\": 12})\n",
    "\n",
    "# os.chdir(sys.path[0]) # ensure working direcotry is set same as the file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-08-09T22:25:16.043734Z",
     "start_time": "2022-08-09T22:25:15.845095Z"
    }
   },
   "outputs": [],
   "source": [
    "######################################  some SCAD and MCP things  #######################################\n",
    "@_jit(nopython=True, cache=True, parallel=True, fastmath=True, nogil=True)\n",
    "def soft_thresholding(x, lambda_):\n",
    "    '''\n",
    "    To calculate soft-thresholding mapping of a given ONE-DIMENSIONAL tensor, BESIDES THE FIRST TERM (so beta_0 will not be penalized). \n",
    "    This function is to be used for calculation involving L1 penalty term later. \n",
    "    '''\n",
    "    return _np.hstack((_np.array([x[0]]),\n",
    "                       _np.where(\n",
    "                           _np.abs(x[1:]) > lambda_,\n",
    "                           x[1:] - _np.sign(x[1:]) * lambda_, 0)))\n",
    "\n",
    "\n",
    "soft_thresholding(_np.random.rand(20), 3.1)\n",
    "\n",
    "\n",
    "@_jit(nopython=True, cache=True, parallel=True, fastmath=True, nogil=True)\n",
    "def SCAD(x, lambda_, a=3.7):\n",
    "    '''\n",
    "    To calculate SCAD penalty value;\n",
    "    #x can be a multi-dimensional tensor;\n",
    "    lambda_, a are scalars;\n",
    "    Fan and Li suggests to take a as 3.7 \n",
    "    '''\n",
    "    # here I notice the function is de facto a function of absolute value of x, therefore take absolute value first to simplify calculation\n",
    "    x = _np.abs(x)\n",
    "    temp = _np.where(\n",
    "        x <= lambda_, lambda_ * x,\n",
    "        _np.where(x < a * lambda_,\n",
    "                  (2 * a * lambda_ * x - x**2 - lambda_**2) / (2 * (a - 1)),\n",
    "                  lambda_**2 * (a + 1) / 2))\n",
    "    temp[0] = 0.  # this is to NOT penalize intercept beta later\n",
    "    return temp\n",
    "\n",
    "\n",
    "@_jit(nopython=True, cache=True, parallel=True, fastmath=True, nogil=True)\n",
    "def SCAD_grad(x, lambda_, a=3.7):\n",
    "    '''\n",
    "    To calculate the gradient of SCAD wrt. input x; \n",
    "    #x can be a multi-dimensional tensor. \n",
    "    '''\n",
    "    # here decompose x to sign and its absolute value for easier calculation\n",
    "    sgn = _np.sign(x)\n",
    "    x = _np.abs(x)\n",
    "    temp = _np.where(\n",
    "        x <= lambda_, lambda_ * sgn,\n",
    "        _np.where(x < a * lambda_, (a * lambda_ * sgn - sgn * x) / (a - 1), 0))\n",
    "    temp[0] = 0.  # this is to NOT penalize intercept beta later\n",
    "    return temp\n",
    "\n",
    "\n",
    "@_jit(nopython=True, cache=True, parallel=True, fastmath=True, nogil=True)\n",
    "def MCP(x, lambda_, gamma):\n",
    "    '''\n",
    "    To calculate MCP penalty value; \n",
    "    #x can be a multi-dimensional tensor. \n",
    "    '''\n",
    "    # the function is a function of absolute value of x\n",
    "    x = _np.abs(x)\n",
    "    temp = _np.where(x <= gamma * lambda_, lambda_ * x - x**2 / (2 * gamma),\n",
    "                     .5 * gamma * lambda_**2)\n",
    "    temp[0] = 0.  # this is to NOT penalize intercept beta later\n",
    "    return temp\n",
    "\n",
    "\n",
    "@_jit(nopython=True, cache=True, parallel=True, fastmath=True, nogil=True)\n",
    "def MCP_grad(x, lambda_, gamma):\n",
    "    '''\n",
    "    To calculate MCP gradient wrt. input x; \n",
    "    #x can be a multi-dimensional tensor. \n",
    "    '''\n",
    "    temp = _np.where(\n",
    "        _np.abs(x) < gamma * lambda_,\n",
    "        lambda_ * _np.sign(x) - x / gamma, _np.zeros_like(x))\n",
    "    temp[0] = 0.  # this is to NOT penalize intercept beta later\n",
    "    return temp\n",
    "\n",
    "\n",
    "@_jit(nopython=True, cache=True, parallel=True, fastmath=True, nogil=True)\n",
    "def SCAD_concave(x, lambda_, a=3.7):\n",
    "    '''\n",
    "    The value of concave part of SCAD penalty; \n",
    "    #x can be a multi-dimensional tensor. \n",
    "    '''\n",
    "    x = _np.abs(x)\n",
    "    temp = _np.where(\n",
    "        x <= lambda_, 0.,\n",
    "        _np.where(x < a * lambda_,\n",
    "                  (lambda_ * x - (x**2 + lambda_**2) / 2) / (a - 1),\n",
    "                  (a + 1) / 2 * lambda_**2 - lambda_ * x))\n",
    "    temp[0] = 0.  # this is to NOT penalize intercept beta later\n",
    "    return temp\n",
    "\n",
    "\n",
    "@_jit(nopython=True, cache=True, parallel=True, fastmath=True, nogil=True)\n",
    "def SCAD_concave_grad(x, lambda_, a=3.7):\n",
    "    '''\n",
    "    The gradient of concave part of SCAD penalty wrt. input x; \n",
    "    #x can be a multi-dimensional tensor. \n",
    "    '''\n",
    "    sgn = _np.sign(x)\n",
    "    x = _np.abs(x)\n",
    "    temp = _np.where(\n",
    "        x <= lambda_, 0.,\n",
    "        _np.where(x < a * lambda_, (lambda_ * sgn - sgn * x) / (a - 1),\n",
    "                  -lambda_ * sgn))\n",
    "    temp[0] = 0.  # this is to NOT penalize intercept beta later\n",
    "    return temp\n",
    "\n",
    "\n",
    "@_jit(nopython=True, cache=True, parallel=True, fastmath=True, nogil=True)\n",
    "def MCP_concave(x, lambda_, gamma):\n",
    "    '''\n",
    "    The value of concave part of MCP penalty; \n",
    "    #x can be a multi-dimensional tensor. \n",
    "    '''\n",
    "    # similiar as in MCP\n",
    "    x = _np.abs(x)\n",
    "    temp = _np.where(x <= gamma * lambda_, -(x**2) / (2 * gamma),\n",
    "                     (gamma * lambda_**2) / 2 - lambda_ * x)\n",
    "    temp[0] = 0.  # this is to NOT penalize intercept beta later\n",
    "    return temp\n",
    "\n",
    "\n",
    "@_jit(nopython=True, cache=True, parallel=True, fastmath=True, nogil=True)\n",
    "def MCP_concave_grad(x, lambda_, gamma):\n",
    "    '''\n",
    "    The gradient of concave part of MCP penalty wrt. input x; \n",
    "    #x can be a multi-dimensional tensor. \n",
    "    '''\n",
    "    temp = _np.where(\n",
    "        _np.abs(x) < gamma * lambda_, -x / gamma, -lambda_ * _np.sign(x))\n",
    "    temp[0] = 0.  # this is to NOT penalize intercept beta later\n",
    "    return temp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-08-09T22:25:16.119938Z",
     "start_time": "2022-08-09T22:25:16.045251Z"
    }
   },
   "outputs": [],
   "source": [
    "@_jit(nopython=True, cache=True, parallel=True, fastmath=True, nogil=True)\n",
    "def _update_smooth_grad_convex_logistic(N, X, beta_md, y):\n",
    "    '''\n",
    "    Update the gradient of the smooth convex objective component.\n",
    "    '''\n",
    "    return (X.T @ (_np.tanh(X @ beta_md / 2.) / 2. - y + .5)) / (2. * N)\n",
    "\n",
    "\n",
    "@_jit(nopython=True, cache=True, parallel=True, fastmath=True, nogil=True)\n",
    "def _update_smooth_grad_SCAD_logistic(N, X, beta_md, y, _lambda, a):\n",
    "    '''\n",
    "    Update the gradient of the smooth objective component for SCAD penalty.\n",
    "    '''\n",
    "    return _update_smooth_grad_convex_logistic(\n",
    "        N=N, X=X, beta_md=beta_md, y=y) + SCAD_concave_grad(\n",
    "            x=beta_md, lambda_=_lambda, a=a)\n",
    "\n",
    "\n",
    "@_jit(nopython=True, cache=True, parallel=True, fastmath=True, nogil=True)\n",
    "def _update_smooth_grad_MCP_logistic(N, X, beta_md, y, _lambda, gamma):\n",
    "    '''\n",
    "    Update the gradient of the smooth objective component for MCP penalty.\n",
    "    '''\n",
    "    return _update_smooth_grad_convex_logistic(\n",
    "        N=N, X=X, beta_md=beta_md, y=y) + MCP_concave_grad(\n",
    "            x=beta_md, lambda_=_lambda, gamma=gamma)\n",
    "\n",
    "\n",
    "@_jit(nopython=True, cache=True, parallel=True, fastmath=True, nogil=True)\n",
    "def _eval_obj_SCAD_logistic(N, X, beta_md, y, _lambda, a, x_temp):\n",
    "    '''\n",
    "    evaluate value of the objective function.\n",
    "    '''\n",
    "    error = y - X @ x_temp\n",
    "    return (error.T @ error) / (2. * N) + _np.sum(\n",
    "        SCAD(x_temp, lambda_=_lambda, a=a))\n",
    "\n",
    "\n",
    "@_jit(nopython=True, cache=True, parallel=True, fastmath=True, nogil=True)\n",
    "def _eval_obj_MCP_logistic(N, X, beta_md, y, _lambda, gamma, x_temp):\n",
    "    '''\n",
    "    evaluate value of the objective function.\n",
    "    '''\n",
    "    error = y - X @ x_temp\n",
    "    return (error.T @ error) / (2 * N) + _np.sum(\n",
    "        SCAD(x_temp, lambda_=_lambda, gamma=gamma))\n",
    "\n",
    "\n",
    "@_jit(nopython=True, cache=True, parallel=True, fastmath=True, nogil=True)\n",
    "def lambda_max_logistic(X, y):\n",
    "    \"\"\"\n",
    "    Calculate the lambda_max, i.e., the minimum lambda to nullify all penalized betas.\n",
    "    \"\"\"\n",
    "    grad_at_0 = (y - _np.mean(y)) @ X_temp / (2 * len(y))\n",
    "    lambda_max = _np.linalg.norm(grad_at_0[1:], ord=_np.infty)\n",
    "    return lambda_max\n",
    "\n",
    "\n",
    "@_jit(nopython=True, cache=True, parallel=True, fastmath=True, nogil=True)\n",
    "def _strong_rule_seq_logistic(X, y, beta_old, lambda_new, lambda_old):\n",
    "    \"\"\"\n",
    "    Use sequential strong to determine which betas to be nullified next.\n",
    "    \"\"\"\n",
    "    grad = _np.abs(\n",
    "        (y - _np.tanh(X @ beta_old / 2) / 2 - .5) @ X_temp / (2 * len(y)))\n",
    "    eliminated = (grad < 2 * lambda_new - lambda_old\n",
    "                  )  # True means the value gets eliminated\n",
    "    eliminated = _np.hstack(\n",
    "        (_np.array([False]),\n",
    "         eliminated))  # because intercept coefficient is not penalized\n",
    "    return eliminated\n",
    "\n",
    "\n",
    "@_jit(nopython=True, cache=True, parallel=True, fastmath=True, nogil=True)\n",
    "def UAG_logistic_SCAD_MCP(design_matrix,\n",
    "                          outcome,\n",
    "                          beta_0=_np.ones(1),\n",
    "                          tol=1e-2,\n",
    "                          maxit=500,\n",
    "                          _lambda=.5,\n",
    "                          penalty=\"SCAD\",\n",
    "                          a=3.7,\n",
    "                          gamma=2.,\n",
    "                          L_convex=1.1,\n",
    "                          add_intercept_column=True):\n",
    "    '''\n",
    "    Carry out the optimization for penalized logistic model for a fixed lambda.\n",
    "    '''\n",
    "    X = design_matrix.copy()\n",
    "    y = outcome.copy()\n",
    "    N = X.shape[0]\n",
    "    if _np.all(beta_0 == _np.ones(1)):\n",
    "        center_X = (X - 1 / N * _np.sum(X, 0).reshape(1, -1))\n",
    "        cov = (y - _np.mean(y)) @ center_X\n",
    "        X_var = _np.sum(X**2, 0)\n",
    "        beta = cov / X_var\n",
    "    else:\n",
    "        beta = beta_0\n",
    "#     add design matrix column for the intercept, if it's not there already\n",
    "    if add_intercept_column == True:\n",
    "        if _np.any(\n",
    "                X[:, 0] != X[0, 0]\n",
    "        ):  # check if design matrix has included a column for intercept or not\n",
    "            intercept_design = _np.ones(N).reshape(-1, 1)\n",
    "            X = _np.hstack((intercept_design, X))\n",
    "            beta = _np.hstack((_np.array([0.]), beta))\n",
    "    # passing other parameters\n",
    "    p = X.shape[1]  # so here p includes the intercept design matrix column\n",
    "    smooth_grad = _np.ones(p)\n",
    "    beta_ag = beta.copy()\n",
    "    beta_md = beta.copy()\n",
    "    k = 0\n",
    "    converged = False\n",
    "    opt_alpha = 1.\n",
    "    #     L_convex = 1/N*_np.max(_np.linalg.eigvalsh(X@X.T)[-1]).item()\n",
    "    if L_convex == 1.1:\n",
    "        L_convex = 1 / N * (_np.linalg.eigvalsh(X @ X.T)[-1])\n",
    "    else:\n",
    "        pass\n",
    "    old_speed_norm = 1.\n",
    "    speed_norm = 1.\n",
    "    restart_k = 0\n",
    "\n",
    "    if penalty == \"SCAD\":\n",
    "        #         L = _np.max(_np.array([L_convex, 1./(a-1)]))\n",
    "        L = _np.linalg.norm(_np.array([L_convex, 1. / (a - 1)]), ord=_np.infty)\n",
    "        opt_beta = .99 / L\n",
    "        while ((not converged) or (k < 3)) and k <= maxit:\n",
    "            k += 1\n",
    "            if old_speed_norm > speed_norm and k - restart_k >= 3:  # in this case, restart\n",
    "                opt_alpha = 1.  # restarting\n",
    "                restart_k = k  # restarting\n",
    "            else:  # restarting\n",
    "                opt_alpha = 2 / (\n",
    "                    1 + (1 + 4. / opt_alpha**2)**.5\n",
    "                )  #parameter settings based on minimizing Ghadimi and Lan's rate of convergence error upper bound\n",
    "            opt_lambda = opt_beta / opt_alpha  #parameter settings based on minimizing Ghadimi and Lan's rate of convergence error upper bound\n",
    "            beta_md_old = beta_md.copy()  # restarting\n",
    "            beta_md = (1 - opt_alpha) * beta_ag + opt_alpha * beta\n",
    "            old_speed_norm = speed_norm  # restarting\n",
    "            speed_norm = _np.linalg.norm(beta_md - beta_md_old,\n",
    "                                         ord=2)  # restarting\n",
    "            converged = (_np.linalg.norm(beta_md - beta_md_old, ord=_np.infty)\n",
    "                         < tol)\n",
    "            smooth_grad = _update_smooth_grad_SCAD_logistic(N=N,\n",
    "                                                            X=X,\n",
    "                                                            beta_md=beta_md,\n",
    "                                                            y=y,\n",
    "                                                            _lambda=_lambda,\n",
    "                                                            a=a)\n",
    "            beta = soft_thresholding(x=beta - opt_lambda * smooth_grad,\n",
    "                                     lambda_=opt_lambda * _lambda)\n",
    "            beta_ag = soft_thresholding(x=beta_md - opt_beta * smooth_grad,\n",
    "                                        lambda_=opt_beta * _lambda)\n",
    "#             converged = _np.all(_np.max(_np.abs(beta_md - beta_ag)/opt_beta) < tol).item()\n",
    "#             converged = (_np.linalg.norm(beta_md - beta_ag, ord=_np.infty) < (tol*opt_beta))\n",
    "    else:\n",
    "        #         L = _np.max(_np.array([L_convex, 1./(gamma)]))\n",
    "        L = _np.linalg.norm(_np.array([L_convex, 1. / (gamma)]), ord=_np.infty)\n",
    "        opt_beta = .99 / L\n",
    "        while ((not converged) or (k < 3)) and k <= maxit:\n",
    "            k += 1\n",
    "            if old_speed_norm > speed_norm and k - restart_k >= 3:  # in this case, restart\n",
    "                opt_alpha = 1.  # restarting\n",
    "                restart_k = k  # restarting\n",
    "            else:  # restarting\n",
    "                opt_alpha = 2 / (\n",
    "                    1 + (1 + 4. / opt_alpha**2)**.5\n",
    "                )  #parameter settings based on minimizing Ghadimi and Lan's rate of convergence error upper bound\n",
    "            opt_lambda = opt_beta / opt_alpha  #parameter settings based on minimizing Ghadimi and Lan's rate of convergence error upper bound\n",
    "            beta_md_old = beta_md.copy()  # restarting\n",
    "            beta_md = (1 - opt_alpha) * beta_ag + opt_alpha * beta\n",
    "            old_speed_norm = speed_norm  # restarting\n",
    "            speed_norm = _np.linalg.norm(beta_md - beta_md_old,\n",
    "                                         ord=2)  # restarting\n",
    "            converged = (_np.linalg.norm(beta_md - beta_md_old, ord=_np.infty)\n",
    "                         < tol)\n",
    "            smooth_grad = _update_smooth_grad_MCP_logistic(N=N,\n",
    "                                                           X=X,\n",
    "                                                           beta_md=beta_md,\n",
    "                                                           y=y,\n",
    "                                                           _lambda=_lambda,\n",
    "                                                           gamma=gamma)\n",
    "            beta = soft_thresholding(x=beta - opt_lambda * smooth_grad,\n",
    "                                     lambda_=opt_lambda * _lambda)\n",
    "            beta_ag = soft_thresholding(x=beta_md - opt_beta * smooth_grad,\n",
    "                                        lambda_=opt_beta * _lambda)\n",
    "#             converged = _np.all(_np.max(_np.abs(beta_md - beta_ag)/opt_beta) < tol).item()\n",
    "#             converged = (_np.linalg.norm(beta_md - beta_ag, ord=_np.infty) < (tol*opt_beta))\n",
    "    return k, beta_md\n",
    "\n",
    "\n",
    "# def vanilla_proximal(self):\n",
    "#     '''\n",
    "#     Carry out optimization using vanilla gradient descent.\n",
    "#     '''\n",
    "#     if self.penalty == \"SCAD\":\n",
    "#         L = max([self.L_convex, 1/(self.a-1)])\n",
    "#         self.vanilla_stepsize = 1/L\n",
    "#         self._eval_obj_SCAD_logistic(self.beta_md, self.obj_value)\n",
    "#         self._eval_obj_SCAD_logistic(self.beta, self.obj_value_ORIGINAL)\n",
    "#         self._eval_obj_SCAD_logistic(self.beta_ag, self.obj_value_AG)\n",
    "#         self.old_beta = self.beta_md - 10.\n",
    "#         while not self.converged:\n",
    "#             self.k += 1\n",
    "#             if self.k <= self.maxit:\n",
    "#                 self._update_smooth_grad_SCAD_logistic()\n",
    "#                 self.beta_md = self.soft_thresholding(self.beta_md - self.vanilla_stepsize*self.smooth_grad, self.vanilla_stepsize*self._lambda)\n",
    "#                 self.converged = _np.all(_np.max(_np.abs(self.beta_md - self.old_beta)) < self.tol).item()\n",
    "#                 self.old_beta = self.beta_md.copy()\n",
    "#                 self._eval_obj_SCAD_logistic(self.beta_md, self.obj_value)\n",
    "#                 self._eval_obj_SCAD_logistic(self.beta, self.obj_value_ORIGINAL)\n",
    "#                 self._eval_obj_SCAD_logistic(self.beta_ag, self.obj_value_AG)\n",
    "#             else:\n",
    "#                 break\n",
    "#     else:\n",
    "#         L = max([self.L_convex, 1/self.gamma])\n",
    "#         self.vanilla_stepsize = 1/L\n",
    "#         self._eval_obj_MCP_logistic(self.beta_md, self.obj_value)\n",
    "#         self._eval_obj_MCP_logistic(self.beta, self.obj_value_ORIGINAL)\n",
    "#         self._eval_obj_MCP_logistic(self.beta_ag, self.obj_value_AG)\n",
    "#         self.old_beta = self.beta_md - 10.\n",
    "#         while not self.converged:\n",
    "#             self.k += 1\n",
    "#             if self.k <= self.maxit:\n",
    "#                 self._update_smooth_grad_MCP_logistic()\n",
    "#                 self.beta_md = self.soft_thresholding(self.beta_md - self.vanilla_stepsize*self.smooth_grad, self.vanilla_stepsize*self._lambda)\n",
    "#                 self.converged = _np.all(_np.max(_np.abs(self.beta_md - self.old_beta)) < self.tol).item()\n",
    "#                 self.old_beta = self.beta_md.copy()\n",
    "#                 self._eval_obj_MCP_logistic(self.beta_md, self.obj_value)\n",
    "#                 self._eval_obj_MCP_logistic(self.beta, self.obj_value_ORIGINAL)\n",
    "#                 self._eval_obj_MCP_logistic(self.beta_ag, self.obj_value_AG)\n",
    "#             else:\n",
    "#                 break\n",
    "#     return self.report_results()\n",
    "\n",
    "\n",
    "@_jit(nopython=True, cache=True, parallel=True, fastmath=True, nogil=True)\n",
    "def solution_path_logistic(design_matrix,\n",
    "                           outcome,\n",
    "                           lambda_,\n",
    "                           beta_0=_np.ones(1),\n",
    "                           tol=1e-2,\n",
    "                           maxit=500,\n",
    "                           penalty=\"SCAD\",\n",
    "                           a=3.7,\n",
    "                           gamma=2.,\n",
    "                           add_intercept_column=True):\n",
    "    '''\n",
    "    Carry out the optimization for the solution path without the strong rule.\n",
    "    '''\n",
    "    #     add design matrix column for the intercept, if it's not there already\n",
    "    if add_intercept_column == True:\n",
    "        if _np.any(\n",
    "                design_matrix[:, 0] != design_matrix[0, 0]\n",
    "        ):  # check if design matrix has included a column for intercept or not\n",
    "            intercept_design = _np.ones(N).reshape(-1, 1)\n",
    "            _design_matrix = design_matrix.copy()\n",
    "            _design_matrix = _np.hstack((intercept_design, _design_matrix))\n",
    "        else:\n",
    "            _design_matrix = design_matrix\n",
    "    else:\n",
    "        _design_matrix = design_matrix\n",
    "    beta_mat = _np.zeros((len(lambda_) + 1, _design_matrix.shape[1]))\n",
    "    for j in range(len(lambda_)):\n",
    "        beta_mat[j + 1, :] = UAG_logistic_SCAD_MCP(\n",
    "            design_matrix=_design_matrix,\n",
    "            outcome=outcome,\n",
    "            beta_0=beta_mat[j, :],\n",
    "            tol=tol,\n",
    "            maxit=maxit,\n",
    "            _lambda=lambda_[j],\n",
    "            penalty=penalty,\n",
    "            a=a,\n",
    "            gamma=gamma,\n",
    "            add_intercept_column=False)[1]\n",
    "    return beta_mat[1:, :]\n",
    "\n",
    "\n",
    "# with strong rule\n",
    "\n",
    "\n",
    "@_jit(nopython=True, cache=True, parallel=True, fastmath=True, nogil=True)\n",
    "def _UAG_logistic_SCAD_MCP_strongrule(design_matrix,\n",
    "                                      outcome,\n",
    "                                      beta_0=_np.ones(1),\n",
    "                                      tol=1e-2,\n",
    "                                      maxit=500,\n",
    "                                      _lambda=.5,\n",
    "                                      penalty=\"SCAD\",\n",
    "                                      a=3.7,\n",
    "                                      gamma=2.,\n",
    "                                      L_convex=1.1,\n",
    "                                      add_intercept_column=True,\n",
    "                                      strongrule=True):\n",
    "    '''\n",
    "    Carry out the optimization for a fixed lambda with strong rule.\n",
    "    '''\n",
    "    X = design_matrix.copy()\n",
    "    y = outcome.copy()\n",
    "    N = X.shape[0]\n",
    "    if _np.all(beta_0 == _np.ones(1)):\n",
    "        center_X = (X - 1 / N * _np.sum(X, 0).reshape(1, -1))\n",
    "        cov = (y - _np.mean(y)) @ center_X\n",
    "        X_var = _np.sum(X**2, 0)\n",
    "        beta = cov / X_var\n",
    "    else:\n",
    "        beta = beta_0\n",
    "#     add design matrix column for the intercept, if it's not there already\n",
    "    if add_intercept_column == True:\n",
    "        if _np.any(\n",
    "                X[:, 0] != X[0, 0]\n",
    "        ):  # check if design matrix has included a column for intercept or not\n",
    "            intercept_design = _np.ones(N).reshape(-1, 1)\n",
    "            X = _np.hstack((intercept_design, X))\n",
    "            beta = _np.hstack((_np.array([0.]), beta))\n",
    "    if strongrule == True:\n",
    "        _lambda_max = lambda_max_logistic(X, y)\n",
    "        p_original = X.shape[1]\n",
    "        elim = _strong_rule_seq_logistic(X,\n",
    "                                         y,\n",
    "                                         beta_old=_np.zeros(p_original),\n",
    "                                         lambda_new=_lambda,\n",
    "                                         lambda_old=_lambda_max)\n",
    "        X = X[:, _np.logical_not(elim)]\n",
    "        beta = beta[_np.logical_not(elim)]\n",
    "\n",
    "    # passing other parameters\n",
    "    p = X.shape[1]  # so here p includes the intercept design matrix column\n",
    "    smooth_grad = _np.ones(p)\n",
    "    beta_ag = beta.copy()\n",
    "    beta_md = beta.copy()\n",
    "    k = 0\n",
    "    converged = False\n",
    "    opt_alpha = 1.\n",
    "    #     L_convex = 1/N*_np.max(_np.linalg.eigvalsh(X@X.T)[-1]).item()\n",
    "    if L_convex == 1.1:\n",
    "        L_convex = 1 / N * (_np.linalg.eigvalsh(X @ X.T)[-1])\n",
    "    else:\n",
    "        pass\n",
    "    old_speed_norm = 1.\n",
    "    speed_norm = 1.\n",
    "    restart_k = 0\n",
    "\n",
    "    if penalty == \"SCAD\":\n",
    "        #         L = _np.max(_np.array([L_convex, 1./(a-1)]))\n",
    "        L = _np.linalg.norm(_np.array([L_convex, 1. / (a - 1)]), ord=_np.infty)\n",
    "        opt_beta = .99 / L\n",
    "        while ((not converged) or (k < 3)) and k <= maxit:\n",
    "            k += 1\n",
    "            if old_speed_norm > speed_norm and k - restart_k >= 3:  # in this case, restart\n",
    "                opt_alpha = 1.  # restarting\n",
    "                restart_k = k  # restarting\n",
    "            else:  # restarting\n",
    "                opt_alpha = 2. / (\n",
    "                    1. + (1. + 4. / opt_alpha**2)**.5\n",
    "                )  #parameter settings based on minimizing Ghadimi and Lan's rate of convergence error upper bound\n",
    "            opt_lambda = opt_beta / opt_alpha  #parameter settings based on minimizing Ghadimi and Lan's rate of convergence error upper bound\n",
    "            beta_md_old = beta_md.copy()  # restarting\n",
    "            beta_md = (1. - opt_alpha) * beta_ag + opt_alpha * beta\n",
    "            old_speed_norm = speed_norm  # restarting\n",
    "            speed_norm = _np.linalg.norm(beta_md - beta_md_old,\n",
    "                                         ord=2)  # restarting\n",
    "            converged = (_np.linalg.norm(beta_md - beta_md_old, ord=_np.infty)\n",
    "                         < tol)\n",
    "            smooth_grad = _update_smooth_grad_SCAD_logistic(N=N,\n",
    "                                                            X=X,\n",
    "                                                            beta_md=beta_md,\n",
    "                                                            y=y,\n",
    "                                                            _lambda=_lambda,\n",
    "                                                            a=a)\n",
    "            beta = soft_thresholding(x=beta - opt_lambda * smooth_grad,\n",
    "                                     lambda_=opt_lambda * _lambda)\n",
    "            beta_ag = soft_thresholding(x=beta_md - opt_beta * smooth_grad,\n",
    "                                        lambda_=opt_beta * _lambda)\n",
    "#             converged = _np.all(_np.max(_np.abs(beta_md - beta_ag)/opt_beta) < tol).item()\n",
    "#             converged = (_np.linalg.norm(beta_md - beta_ag, ord=_np.infty) < (tol*opt_beta))\n",
    "    else:\n",
    "        #         L = _np.max(_np.array([L_convex, 1./(gamma)]))\n",
    "        L = _np.linalg.norm(_np.array([L_convex, 1. / (gamma)]), ord=_np.infty)\n",
    "        opt_beta = .99 / L\n",
    "        while ((not converged) or (k < 3)) and k <= maxit:\n",
    "            k += 1\n",
    "            if old_speed_norm > speed_norm and k - restart_k >= 3:  # in this case, restart\n",
    "                opt_alpha = 1.  # restarting\n",
    "                restart_k = k  # restarting\n",
    "            else:  # restarting\n",
    "                opt_alpha = 2 / (\n",
    "                    1. + (1. + 4. / opt_alpha**2)**.5\n",
    "                )  #parameter settings based on minimizing Ghadimi and Lan's rate of convergence error upper bound\n",
    "            opt_lambda = opt_beta / opt_alpha  #parameter settings based on minimizing Ghadimi and Lan's rate of convergence error upper bound\n",
    "            beta_md_old = beta_md.copy()  # restarting\n",
    "            beta_md = (1. - opt_alpha) * beta_ag + opt_alpha * beta\n",
    "            old_speed_norm = speed_norm  # restarting\n",
    "            speed_norm = _np.linalg.norm(beta_md - beta_md_old,\n",
    "                                         ord=2)  # restarting\n",
    "            converged = (_np.linalg.norm(beta_md - beta_md_old, ord=_np.infty)\n",
    "                         < tol)\n",
    "            smooth_grad = _update_smooth_grad_MCP_logistic(N=N,\n",
    "                                                           X=X,\n",
    "                                                           beta_md=beta_md,\n",
    "                                                           y=y,\n",
    "                                                           _lambda=_lambda,\n",
    "                                                           gamma=gamma)\n",
    "            beta = soft_thresholding(x=beta - opt_lambda * smooth_grad,\n",
    "                                     lambda_=opt_lambda * _lambda)\n",
    "            beta_ag = soft_thresholding(x=beta_md - opt_beta * smooth_grad,\n",
    "                                        lambda_=opt_beta * _lambda)\n",
    "\n",
    "\n",
    "#             converged = _np.all(_np.max(_np.abs(beta_md - beta_ag)/opt_beta) < tol).item()\n",
    "#             converged = (_np.linalg.norm(beta_md - beta_ag, ord=_np.infty) < (tol*opt_beta))\n",
    "#     if strongrule == True:\n",
    "#         _beta_output = _np.zeros((p_original))\n",
    "# #         _ = _np.argwhere(_np.logical_not(elim)).flatten()\n",
    "# #         print(_)\n",
    "# #         for j in range(len(_)):\n",
    "# #             if j<10:\n",
    "# #                 print(j)\n",
    "# #                 print(_[j])\n",
    "# #             _beta_output[_[j]] = beta_md[j]\n",
    "# #             if j<10:\n",
    "# #                 print(_beta_output[_[j]])\n",
    "#         _beta_output[~elim] = beta_md  # this line of code can't compile\n",
    "#     else:\n",
    "#         _beta_output = beta_md\n",
    "    return k, beta_md, elim\n",
    "\n",
    "\n",
    "@_jit(nopython=True, cache=True, parallel=True, fastmath=True, nogil=True)\n",
    "def UAG_logistic_SCAD_MCP_strongrule(design_matrix,\n",
    "                                     outcome,\n",
    "                                     beta_0=_np.ones(1),\n",
    "                                     tol=1e-2,\n",
    "                                     maxit=500,\n",
    "                                     _lambda=.5,\n",
    "                                     penalty=\"SCAD\",\n",
    "                                     a=3.7,\n",
    "                                     gamma=2.,\n",
    "                                     L_convex=1.1,\n",
    "                                     add_intercept_column=True,\n",
    "                                     strongrule=True):\n",
    "    \"\"\"\n",
    "    Carry out the optimization for a fixed lambda for penanlized logistic model with strong rule.\n",
    "    \"\"\"\n",
    "    _k, _beta_md, _elim = _UAG_logistic_SCAD_MCP_strongrule(\n",
    "        design_matrix=design_matrix,\n",
    "        outcome=outcome,\n",
    "        beta_0=beta_0,\n",
    "        tol=tol,\n",
    "        maxit=maxit,\n",
    "        _lambda=_lambda,\n",
    "        penalty=penalty,\n",
    "        a=a,\n",
    "        gamma=gamma,\n",
    "        L_convex=L_convex,\n",
    "        add_intercept_column=add_intercept_column,\n",
    "        strongrule=strongrule)\n",
    "    output_beta = _np.zeros(len(_elim))\n",
    "    output_beta[_np.logical_not(_elim)] = _beta_md\n",
    "    return _k, output_beta\n",
    "\n",
    "\n",
    "@_jit(nopython=True, cache=True, parallel=True, fastmath=True, nogil=True)\n",
    "def solution_path_logistic_strongrule(design_matrix,\n",
    "                                      outcome,\n",
    "                                      lambda_,\n",
    "                                      beta_0=_np.ones(1),\n",
    "                                      tol=1e-2,\n",
    "                                      maxit=500,\n",
    "                                      penalty=\"SCAD\",\n",
    "                                      a=3.7,\n",
    "                                      gamma=2.,\n",
    "                                      add_intercept_column=True):\n",
    "    '''\n",
    "    Carry out the optimization for the solution path of a penalized logistic model with strong rule.\n",
    "    '''\n",
    "    #     add design matrix column for the intercept, if it's not there already\n",
    "    _design_matrix = design_matrix.copy()\n",
    "    if add_intercept_column == True:\n",
    "        if _np.any(\n",
    "                design_matrix[:, 0] != design_matrix[0, 0]\n",
    "        ):  # check if design matrix has included a column for intercept or not\n",
    "            intercept_design = _np.ones(N).reshape(-1, 1)\n",
    "            _design_matrix = _np.hstack((intercept_design, _design_matrix))\n",
    "    beta_mat = _np.empty((len(lambda_) + 1, _design_matrix.shape[1]))\n",
    "    beta_mat[0, :] = 0.\n",
    "    _lambda_max = lambda_max_logistic(_design_matrix, outcome)\n",
    "    lambda_ = _np.hstack((_np.array([_lambda_max]), lambda_))\n",
    "    elim = _np.array([False] * _design_matrix.shape[1])\n",
    "    for j in range(len(lambda_) - 1):\n",
    "        _elim = _strong_rule_seq_logistic(X=_design_matrix,\n",
    "                                          y=outcome,\n",
    "                                          beta_old=beta_mat[j, :],\n",
    "                                          lambda_new=lambda_[j + 1],\n",
    "                                          lambda_old=lambda_[j])\n",
    "        elim = _np.logical_and(elim, _elim)\n",
    "        _beta_0 = beta_mat[j, :]\n",
    "        _new_beta = _np.zeros(_design_matrix.shape[1])\n",
    "        _new_beta[_np.logical_not(elim)] = UAG_logistic_SCAD_MCP(\n",
    "            design_matrix=_design_matrix[:, _np.logical_not(elim)],\n",
    "            outcome=outcome,\n",
    "            beta_0=_beta_0[_np.logical_not(elim)],\n",
    "            tol=tol,\n",
    "            maxit=maxit,\n",
    "            _lambda=lambda_[j],\n",
    "            penalty=penalty,\n",
    "            a=a,\n",
    "            gamma=gamma,\n",
    "            add_intercept_column=False)[1]\n",
    "        beta_mat[j + 1, :] = _new_beta\n",
    "    return beta_mat[1:, :]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-08-09T22:25:58.168832Z",
     "start_time": "2022-08-09T22:25:16.345592Z"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.markers as markers\n",
    "import matplotlib.pyplot as plt\n",
    "import timeit\n",
    "from scipy.linalg import toeplitz, block_diag\n",
    "\n",
    "# this cell is for profiling the function\n",
    "np.random.seed(0)\n",
    "N = 1000\n",
    "SNR = 5.\n",
    "true_beta = np.array([.5, -.5, .8, -.8] + [0] * 2000)\n",
    "X_cov = toeplitz(.5**np.arange(2004))\n",
    "mean = np.zeros(true_beta.shape[0])\n",
    "X = np.random.multivariate_normal(mean, X_cov, N)\n",
    "X -= np.mean(X, 0).reshape(1, -1)\n",
    "X /= np.std(X, 0)\n",
    "intercept_design_column = np.ones(N).reshape(N, 1)\n",
    "X_sim = np.concatenate((intercept_design_column, X), 1)\n",
    "true_sigma_sim = np.sqrt(true_beta.T @ X_cov @ true_beta / SNR)\n",
    "true_beta_intercept = np.concatenate((np.array([0.5]), true_beta))\n",
    "signal = X_sim @ true_beta_intercept + np.random.normal(0, true_sigma_sim, N)\n",
    "y_sim = np.random.binomial(1, np.tanh(signal / 2) / 2 + .5)\n",
    "\n",
    "fit1 = UAG_logistic_SCAD_MCP(design_matrix=X_sim,\n",
    "                             outcome=y_sim,\n",
    "                             tol=1e-2,\n",
    "                             maxit=500,\n",
    "                             _lambda=.2,\n",
    "                             penalty=\"SCAD\",\n",
    "                             a=3.7,\n",
    "                             gamma=2.)\n",
    "\n",
    "fit2 = solution_path_logistic(design_matrix=X_sim,\n",
    "                              outcome=y_sim,\n",
    "                              tol=1e-2,\n",
    "                              maxit=500,\n",
    "                              lambda_=np.linspace(.005, .08, 60)[::-1],\n",
    "                              penalty=\"SCAD\",\n",
    "                              a=3.7,\n",
    "                              gamma=2.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-08-09T22:25:58.177529Z",
     "start_time": "2022-08-09T22:25:58.170734Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[4.18288649e-01, 0.00000000e+00, 0.00000000e+00, ...,\n",
       "        0.00000000e+00, 0.00000000e+00, 0.00000000e+00],\n",
       "       [4.19652110e-01, 0.00000000e+00, 0.00000000e+00, ...,\n",
       "        0.00000000e+00, 0.00000000e+00, 0.00000000e+00],\n",
       "       [4.20959440e-01, 0.00000000e+00, 0.00000000e+00, ...,\n",
       "        0.00000000e+00, 0.00000000e+00, 0.00000000e+00],\n",
       "       ...,\n",
       "       [4.50533893e-01, 4.79884573e-02, 0.00000000e+00, ...,\n",
       "        0.00000000e+00, 0.00000000e+00, 0.00000000e+00],\n",
       "       [4.51245655e-01, 5.41205703e-02, 0.00000000e+00, ...,\n",
       "        3.57689149e-04, 0.00000000e+00, 0.00000000e+00],\n",
       "       [4.52171425e-01, 5.98629351e-02, 0.00000000e+00, ...,\n",
       "        1.09629533e-03, 0.00000000e+00, 0.00000000e+00]])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fit2"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.13"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  },
  "nbTranslate": {
   "displayLangs": [
    "*"
   ],
   "hotkey": "alt-t",
   "langInMainMenu": true,
   "sourceLang": "en",
   "targetLang": "fr",
   "useGoogleTranslate": true
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "position": {
    "height": "877.844px",
    "left": "2188px",
    "right": "20px",
    "top": "120px",
    "width": "352px"
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}