LM_SCAD_MCP_np.memmap_MODULE_parallel.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Set up the class fundementals "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-10-31T03:39:53.383463Z",
     "start_time": "2022-10-31T03:39:53.202002Z"
    }
   },
   "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.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",
    "import multiprocess as _mp\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-10-31T03:39:53.533418Z",
     "start_time": "2022-10-31T03:39:53.384617Z"
    }
   },
   "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-10-31T03:39:53.558225Z",
     "start_time": "2022-10-31T03:39:53.534633Z"
    },
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "# @_jit(nopython=True, cache=True, parallel=True, fastmath=True, nogil=True)\n",
    "def _memmap_update_smooth_grad_convex_LM_parallel(N, p, X, beta_md, y, _dtype,\n",
    "                                                  _order, core_num, multp):\n",
    "    '''\n",
    "    Update the gradient of the smooth convex objective component.\n",
    "    '''\n",
    "    _itemsize = _np.dtype(_dtype).itemsize\n",
    "    # first calcualte _=X@beta_md-y\n",
    "    if _order == \"F\":\n",
    "\n",
    "        def __parallel_plus(_ind):\n",
    "            import numpy as _np\n",
    "            _X = _np.memmap(X,\n",
    "                            dtype=_dtype,\n",
    "                            mode='r',\n",
    "                            offset=_ind[0] * _itemsize * N,\n",
    "                            shape=(N, len(_ind)))\n",
    "            return _X @ beta_md[_ind]\n",
    "#             __ = _np.zeros(N)\n",
    "#             for j in _ind:\n",
    "#                 _X = _np.memmap(X,\n",
    "#                                 dtype=_dtype,\n",
    "#                                 mode='r',\n",
    "#                                 offset=j * _itemsize * N,\n",
    "#                                 shape=(N, ))\n",
    "#                 __ += _X * beta_md[j]\n",
    "#             return __\n",
    "\n",
    "# multiprocessing starts here\n",
    "\n",
    "        ind = _np.arange(p)\n",
    "        _splited_array = _np.array_split(ind, core_num * multp)\n",
    "        _splited_array = [\n",
    "            __array for __array in _splited_array if __array.size != 0\n",
    "        ]\n",
    "        with _mp.Pool(core_num) as pl:\n",
    "            _ = pl.map(__parallel_plus, _splited_array)\n",
    "        _ = _np.array(_).sum(0)\n",
    "    elif _order == \"C\":\n",
    "\n",
    "        def __parallel_assign(_ind):\n",
    "            import numpy as _np\n",
    "            _X = _np.memmap(X,\n",
    "                            dtype=_dtype,\n",
    "                            mode='r',\n",
    "                            offset=_ind[0] * _itemsize * p,\n",
    "                            shape=(p, len(_ind)))\n",
    "            return beta_md @ _X\n",
    "#             k = 0\n",
    "#             __ = _np.zeros(len(_ind))\n",
    "#             for j in _ind:\n",
    "#                 _X = _np.memmap(X,\n",
    "#                                 dtype=_dtype,\n",
    "#                                 mode='r',\n",
    "#                                 offset=j * _itemsize * p,\n",
    "#                                 shape=(p, ))\n",
    "#                 __[k] = _X @ beta_md\n",
    "#                 k += 1\n",
    "#             return __\n",
    "\n",
    "# multiprocessing starts here\n",
    "\n",
    "        ind = _np.arange(N)\n",
    "        _splited_array = _np.array_split(ind, core_num * multp)\n",
    "        _splited_array = [\n",
    "            __array for __array in _splited_array if __array.size != 0\n",
    "        ]\n",
    "        with _mp.Pool(core_num) as pl:\n",
    "            _ = pl.map(__parallel_assign, _splited_array)\n",
    "        _ = _np.hstack(_)\n",
    "    _ -= y\n",
    "    # then calculate _XTXbeta = X.T@X@beta_md = X.T@_\n",
    "    if _order == \"F\":\n",
    "\n",
    "        def __parallel_assign(_ind):\n",
    "            import numpy as _np\n",
    "            _X = _np.memmap(X,\n",
    "                            dtype=_dtype,\n",
    "                            mode='r',\n",
    "                            offset=_ind[0] * _itemsize * N,\n",
    "                            shape=(N, len(_ind)))\n",
    "            return _ @ _X\n",
    "#             k = 0\n",
    "#             __ = _np.zeros(len(_ind))\n",
    "#             for j in _ind:\n",
    "#                 _X = _np.memmap(X,\n",
    "#                                 dtype=_dtype,\n",
    "#                                 mode='r',\n",
    "#                                 offset=j * _itemsize * N,\n",
    "#                                 shape=(N, ))\n",
    "#                 __[k] = _X @ _\n",
    "#                 k += 1\n",
    "#             return __\n",
    "\n",
    "# multiprocessing starts here\n",
    "\n",
    "        ind = _np.arange(p)\n",
    "        _splited_array = _np.array_split(ind, core_num * multp)\n",
    "        _splited_array = [\n",
    "            __array for __array in _splited_array if __array.size != 0\n",
    "        ]\n",
    "        with _mp.Pool(core_num) as pl:\n",
    "            _XTXbeta = pl.map(__parallel_assign, _splited_array)\n",
    "        _XTXbeta = _np.hstack(_XTXbeta)\n",
    "    elif _order == \"C\":\n",
    "\n",
    "        def __parallel_plus(_ind):\n",
    "            import numpy as _np\n",
    "            _X = _np.memmap(X,\n",
    "                            dtype=_dtype,\n",
    "                            mode='r',\n",
    "                            offset=_ind[0] * _itemsize * p,\n",
    "                            shape=(p, len(_ind)))\n",
    "            return _X @ _[_ind]\n",
    "#             __ = _np.zeros(p)\n",
    "#             for j in _ind:\n",
    "#                 _X = _np.memmap(X,\n",
    "#                                 dtype=_dtype,\n",
    "#                                 mode='r',\n",
    "#                                 offset=j * _itemsize * p,\n",
    "#                                 shape=(p, ))\n",
    "#                 __ += _X * _[j]\n",
    "#             return __\n",
    "\n",
    "# multiprocessing starts here\n",
    "\n",
    "        ind = _np.arange(N)\n",
    "        _splited_array = _np.array_split(ind, core_num * multp)\n",
    "        _splited_array = [\n",
    "            __array for __array in _splited_array if __array.size != 0\n",
    "        ]\n",
    "        with _mp.Pool(core_num) as pl:\n",
    "            _XTXbeta = pl.map(__parallel_plus, _splited_array)\n",
    "        _XTXbeta = _np.array(_XTXbeta).sum(0)\n",
    "    del _\n",
    "    return 1 / N * _XTXbeta\n",
    "\n",
    "\n",
    "# @_jit(nopython=True, cache=True, parallel=True, fastmath=True, nogil=True)\n",
    "def _memmap_update_smooth_grad_SCAD_LM_parallel(N, p, X, beta_md, y, _lambda,\n",
    "                                                a, _dtype, _order, core_num,\n",
    "                                                multp):\n",
    "    '''\n",
    "    Update the gradient of the smooth objective component for SCAD penalty.\n",
    "    '''\n",
    "    return _memmap_update_smooth_grad_convex_LM_parallel(\n",
    "        N=N,\n",
    "        p=p,\n",
    "        X=X,\n",
    "        beta_md=beta_md,\n",
    "        y=y,\n",
    "        _dtype=_dtype,\n",
    "        _order=_order,\n",
    "        core_num=core_num,\n",
    "        multp=multp) + SCAD_concave_grad(x=beta_md, lambda_=_lambda, a=a)\n",
    "\n",
    "\n",
    "# @_jit(nopython=True, cache=True, parallel=True, fastmath=True, nogil=True)\n",
    "def _memmap_update_smooth_grad_MCP_LM_parallel(N, p, X, beta_md, y, _lambda,\n",
    "                                               gamma, _dtype, _order, core_num,\n",
    "                                               multp):\n",
    "    '''\n",
    "    Update the gradient of the smooth objective component for MCP penalty.\n",
    "    '''\n",
    "    return _memmap_update_smooth_grad_convex_LM_parallel(\n",
    "        N=N,\n",
    "        p=p,\n",
    "        X=X,\n",
    "        beta_md=beta_md,\n",
    "        y=y,\n",
    "        _dtype=_dtype,\n",
    "        _order=_order,\n",
    "        core_num=core_num,\n",
    "        multp=multp) + 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 memmap_lambda_max_LM_parallel(X,\n",
    "                                  y,\n",
    "                                  N,\n",
    "                                  p,\n",
    "                                  _dtype,\n",
    "                                  _order,\n",
    "                                  core_num=\"NOT DECLARED\",\n",
    "                                  multp=1):\n",
    "    \"\"\"\n",
    "    Calculate the lambda_max, i.e., the minimum lambda to nullify all penalized betas.\n",
    "    \"\"\"\n",
    "    #     X_temp = X.copy()\n",
    "    #     X_temp = X_temp[:,1:]\n",
    "    #     X_temp -= _np.mean(X_temp,0).reshape(1,-1)\n",
    "    #     X_temp /= _np.std(X_temp,0)\n",
    "    #     y_temp = y.copy()\n",
    "    #     y_temp -= _np.mean(y)\n",
    "    #     y_temp /= _np.std(y)\n",
    "    if core_num == \"NOT DECLARED\":\n",
    "        core_num = _mp.cpu_count()\n",
    "    else:\n",
    "        assert core_num <= _mp.cpu_count(\n",
    "        ), \"Declared number of cores used for multiprocessing should not exceed number of cores on this machine.\"\n",
    "    assert core_num >= 2, \"Multiprocessing should not be used on single-core machines.\"\n",
    "\n",
    "    grad_at_0 = _memmap_update_smooth_grad_convex_LM_parallel(\n",
    "        N=N,\n",
    "        p=p,\n",
    "        X=X,\n",
    "        beta_md=_np.zeros(p),\n",
    "        y=y,\n",
    "        _dtype=_dtype,\n",
    "        _order=_order,\n",
    "        core_num=core_num,\n",
    "        multp=multp)\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 memmap_UAG_LM_SCAD_MCP_parallel(design_matrix,\n",
    "                                    outcome,\n",
    "                                    N,\n",
    "                                    p,\n",
    "                                    L_convex,\n",
    "                                    _dtype,\n",
    "                                    _order,\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",
    "                                    core_num=\"NOT DECLARED\",\n",
    "                                    multp=1):\n",
    "    '''\n",
    "    Carry out the optimization for penalized LM for a fixed lambda.\n",
    "    '''\n",
    "    if core_num == \"NOT DECLARED\":\n",
    "        core_num = _mp.cpu_count()\n",
    "    else:\n",
    "        assert core_num <= _mp.cpu_count(\n",
    "        ), \"Declared number of cores used for multiprocessing should not exceed number of cores on this machine.\"\n",
    "    assert core_num >= 2, \"Multiprocessing should not be used on single-core machines.\"\n",
    "\n",
    "    X = design_matrix\n",
    "    y = outcome\n",
    "    _itemsize = _np.dtype(_dtype).itemsize\n",
    "    if _np.all(beta_0 == _np.ones(1)):\n",
    "        if _order == \"F\":\n",
    "\n",
    "            def __parallel_assign(_ind):\n",
    "                import numpy as _np\n",
    "                _X = _np.memmap(X,\n",
    "                                dtype=_dtype,\n",
    "                                mode='r',\n",
    "                                offset=_ind[0] * _itemsize * N,\n",
    "                                shape=(N, len(_ind)))\n",
    "                return y @ _X / _np.var(_X, 0).reshape(1, -1) / len(y)\n",
    "#                 k = 0\n",
    "#                 __ = _np.zeros(len(_ind))\n",
    "#                 for j in _ind:\n",
    "#                     _X = _np.memmap(X,\n",
    "#                                     dtype=_dtype,\n",
    "#                                     mode='r',\n",
    "#                                     offset=j * _itemsize * N,\n",
    "#                                     shape=(N, ))\n",
    "#                     __[k] = _X @ y / _np.var(_X) / len(y)\n",
    "#                     k += 1\n",
    "#                 return __\n",
    "\n",
    "# multiprocessing starts here\n",
    "\n",
    "            ind = _np.arange(p)\n",
    "            _splited_array = _np.array_split(ind, core_num * multp)\n",
    "            _splited_array = [\n",
    "                __array for __array in _splited_array if __array.size != 0\n",
    "            ]\n",
    "            with _mp.Pool(core_num) as pl:\n",
    "                _XTy = pl.map(__parallel_assign, _splited_array)\n",
    "            _XTy = _np.hstack(_XTy)\n",
    "        elif _order == \"C\":\n",
    "\n",
    "            def __parallel_plus(_ind):\n",
    "                import numpy as _np\n",
    "                _X = _np.memmap(X,\n",
    "                                dtype=_dtype,\n",
    "                                mode='r',\n",
    "                                offset=_ind[0] * _itemsize * p,\n",
    "                                shape=(p, len(_ind)))\n",
    "                return _X @ y[_ind]\n",
    "#                 __ = _np.zeros(p)\n",
    "#                 for j in _ind:\n",
    "#                     _X = _np.memmap(X,\n",
    "#                                     dtype=_dtype,\n",
    "#                                     mode='r',\n",
    "#                                     offset=j * _itemsize * p,\n",
    "#                                     shape=(p, ))\n",
    "#                     __ += _X * y[j]\n",
    "#                 return __\n",
    "\n",
    "# multiprocessing starts here\n",
    "\n",
    "            ind = _np.arange(N)\n",
    "            _splited_array = _np.array_split(ind, core_num * multp)\n",
    "            _splited_array = [\n",
    "                __array for __array in _splited_array if __array.size != 0\n",
    "            ]\n",
    "            with _mp.Pool(core_num) as pl:\n",
    "                _XTy = pl.map(__parallel_plus, _splited_array)\n",
    "            _XTy = _np.array(_XTy).sum(0)\n",
    "        beta = _np.sign(_XTy)\n",
    "    else:\n",
    "        beta = beta_0\n",
    "    # passing other parameters\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",
    "    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 = _memmap_update_smooth_grad_SCAD_LM_parallel(\n",
    "                N=N,\n",
    "                p=p,\n",
    "                X=X,\n",
    "                beta_md=beta_md,\n",
    "                y=y,\n",
    "                _lambda=_lambda,\n",
    "                a=a,\n",
    "                _dtype=_dtype,\n",
    "                _order=_order,\n",
    "                core_num=core_num,\n",
    "                multp=multp)\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 = _memmap_update_smooth_grad_MCP_LM_parallel(\n",
    "                N=N,\n",
    "                p=p,\n",
    "                X=X,\n",
    "                beta_md=beta_md,\n",
    "                y=y,\n",
    "                _lambda=_lambda,\n",
    "                gamma=gamma,\n",
    "                _dtype=_dtype,\n",
    "                _order=_order,\n",
    "                core_num=core_num,\n",
    "                multp=multp)\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",
    "# @_jit(nopython=True, cache=True, parallel=True, fastmath=True, nogil=True)\n",
    "def memmap_solution_path_LM_parallel(design_matrix,\n",
    "                                     outcome,\n",
    "                                     lambda_,\n",
    "                                     L_convex,\n",
    "                                     N,\n",
    "                                     p,\n",
    "                                     beta_0=_np.ones(1),\n",
    "                                     tol=1e-2,\n",
    "                                     maxit=500,\n",
    "                                     penalty=\"SCAD\",\n",
    "                                     a=3.7,\n",
    "                                     gamma=2.,\n",
    "                                     _dtype='float32',\n",
    "                                     _order=\"F\",\n",
    "                                     core_num=\"NOT DECLARED\",\n",
    "                                     multp=1):\n",
    "    '''\n",
    "    Carry out the optimization for the solution path without the strong rule.\n",
    "    '''\n",
    "    if core_num == \"NOT DECLARED\":\n",
    "        core_num = _mp.cpu_count()\n",
    "    else:\n",
    "        assert core_num <= _mp.cpu_count(\n",
    "        ), \"Declared number of cores used for multiprocessing should not exceed number of cores on this machine.\"\n",
    "    assert core_num >= 2, \"Multiprocessing should not be used on single-core machines.\"\n",
    "\n",
    "    beta_mat = _np.zeros((len(lambda_) + 1, p))\n",
    "    for j in range(len(lambda_)):\n",
    "        beta_mat[j + 1, :] = memmap_UAG_LM_SCAD_MCP_parallel(\n",
    "            design_matrix=design_matrix,\n",
    "            outcome=outcome,\n",
    "            N=N,\n",
    "            p=p,\n",
    "            L_convex=L_convex,\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",
    "            _dtype=_dtype,\n",
    "            _order=_order,\n",
    "            core_num=core_num,\n",
    "            multp=multp)[1]\n",
    "    return beta_mat[1:, :]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-10-31T03:41:28.455384Z",
     "start_time": "2022-10-31T03:39:53.573122Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 1.96867704e-01  0.00000000e+00  0.00000000e+00 ...  0.00000000e+00\n",
      "   0.00000000e+00  0.00000000e+00]\n",
      " [ 2.10435948e-01  0.00000000e+00  0.00000000e+00 ...  0.00000000e+00\n",
      "   0.00000000e+00  0.00000000e+00]\n",
      " [ 2.18271253e-01  0.00000000e+00  0.00000000e+00 ...  0.00000000e+00\n",
      "   0.00000000e+00  0.00000000e+00]\n",
      " ...\n",
      " [ 7.35216570e-02 -4.25017007e-01  6.52911611e-04 ...  0.00000000e+00\n",
      "  -7.55299733e-01  0.00000000e+00]\n",
      " [ 7.24373140e-02 -4.26819656e-01  8.50617125e-03 ... -2.56520069e-04\n",
      "  -7.54529186e-01  0.00000000e+00]\n",
      " [ 8.13181946e-02 -4.57818252e-01  1.97498914e-01 ... -4.56403924e-02\n",
      "  -8.06451332e-01  0.00000000e+00]]\n"
     ]
    }
   ],
   "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",
    "from tempfile import mkdtemp\n",
    "import os.path as path\n",
    "\n",
    "np.random.seed(1)\n",
    "N = 2000\n",
    "SNR = 5.\n",
    "true_beta = np.array([2., -2, 8, -8] + [0] * 1000)\n",
    "X_cov = toeplitz(.6**np.arange(true_beta.shape[0]))\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([\n",
    "    1.23\n",
    "]), true_beta))  # here just define the intercept to be 1.23 for simulated data\n",
    "epsilon = np.random.normal(0, true_sigma_sim, N)\n",
    "y_sim = X_sim @ true_beta_intercept + epsilon\n",
    "\n",
    "L_convex = 1 / N * (np.linalg.eigvalsh(X_sim @ X_sim.T)[-1])\n",
    "\n",
    "filename = path.join(mkdtemp(), 'newfile.dat')\n",
    "fp = np.memmap(filename,\n",
    "               dtype='float32',\n",
    "               mode='w+',\n",
    "               shape=(N, 1005),\n",
    "               order=\"C\")\n",
    "fp[:] = X_sim[:]\n",
    "fp.flush()\n",
    "\n",
    "lambda_seq = np.arange(40.) / 80\n",
    "lambda_seq = lambda_seq[1:]\n",
    "lambda_seq = lambda_seq[::-1]\n",
    "\n",
    "# do NOT include the design matrix intercept column\n",
    "LM_beta = memmap_solution_path_LM_parallel(design_matrix=filename,\n",
    "                                           N=N,\n",
    "                                           p=1005,\n",
    "                                           outcome=y_sim,\n",
    "                                           lambda_=lambda_seq,\n",
    "                                           L_convex=L_convex,\n",
    "                                           beta_0=np.ones(1),\n",
    "                                           tol=1e-2,\n",
    "                                           maxit=500,\n",
    "                                           penalty=\"SCAD\",\n",
    "                                           a=3.7,\n",
    "                                           gamma=2.,\n",
    "                                           _dtype=\"float32\",\n",
    "                                           _order=\"C\")\n",
    "\n",
    "print(LM_beta)  # to make sure strong rule runs correctly\n",
    "\n",
    "# %timeit solution_path_LM_strongrule(design_matrix=X_sim, outcome=y_sim, lambda_=lambda_seq, beta_0 = np.ones(1), tol=1e-2, maxit=500, penalty=\"SCAD\", a=3.7, gamma=2., add_intercept_column=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-10-31T03:42:50.056452Z",
     "start_time": "2022-10-31T03:41:28.461989Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[-0.25883078  0.          0.         ...  0.          0.\n",
      "   0.        ]\n",
      " [-0.30222458  0.          0.         ...  0.          0.\n",
      "   0.        ]\n",
      " [-0.30577746  0.          0.         ...  0.          0.\n",
      "   0.        ]\n",
      " ...\n",
      " [-0.45085665  0.          0.29066327 ...  0.30676908  0.00166282\n",
      "   0.        ]\n",
      " [-0.40006552 -0.00843024  0.2413313  ...  0.36358958  0.07762319\n",
      "   0.        ]\n",
      " [-0.39970425 -0.01441658  0.24154475 ...  0.36639908  0.08280139\n",
      "   0.        ]]\n"
     ]
    }
   ],
   "source": [
    "import matplotlib.markers as markers\n",
    "import matplotlib.pyplot as plt\n",
    "import timeit\n",
    "from scipy.linalg import toeplitz, block_diag\n",
    "from tempfile import mkdtemp\n",
    "import os.path as path\n",
    "\n",
    "np.random.seed(1)\n",
    "N = 2000\n",
    "SNR = 5.\n",
    "true_beta = np.array([2., -2, 8, -8] + [0] * 1000)\n",
    "X_cov = toeplitz(.6**np.arange(true_beta.shape[0]))\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([\n",
    "    1.23\n",
    "]), true_beta))  # here just define the intercept to be 1.23 for simulated data\n",
    "epsilon = np.random.normal(0, true_sigma_sim, N)\n",
    "y_sim = X_sim @ true_beta_intercept + epsilon\n",
    "\n",
    "L_convex = 1 / N * (np.linalg.eigvalsh(X_sim @ X_sim.T)[-1])\n",
    "\n",
    "filename = path.join(mkdtemp(), 'newfile.dat')\n",
    "fp = np.memmap(filename,\n",
    "               dtype='float32',\n",
    "               mode='w+',\n",
    "               shape=(N, 1005),\n",
    "               order=\"F\")\n",
    "fp[:] = X_sim[:]\n",
    "fp.flush()\n",
    "\n",
    "lambda_seq = np.arange(40.) / 80\n",
    "lambda_seq = lambda_seq[1:]\n",
    "lambda_seq = lambda_seq[::-1]\n",
    "\n",
    "# do NOT include the design matrix intercept column\n",
    "LM_beta = memmap_solution_path_LM_parallel(design_matrix=filename,\n",
    "                                           N=N,\n",
    "                                           p=1005,\n",
    "                                           outcome=y_sim,\n",
    "                                           lambda_=lambda_seq,\n",
    "                                           L_convex=L_convex,\n",
    "                                           beta_0=np.ones(1),\n",
    "                                           tol=1e-2,\n",
    "                                           maxit=500,\n",
    "                                           penalty=\"SCAD\",\n",
    "                                           a=3.7,\n",
    "                                           gamma=2.,\n",
    "                                           _dtype=\"float32\",\n",
    "                                           _order=\"F\")\n",
    "\n",
    "print(LM_beta)  # to make sure strong rule runs correctly\n",
    "\n",
    "# %timeit solution_path_LM_strongrule(design_matrix=X_sim, outcome=y_sim, lambda_=lambda_seq, beta_0 = np.ones(1), tol=1e-2, maxit=500, penalty=\"SCAD\", a=3.7, gamma=2., add_intercept_column=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}