{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# __Matrix__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In group elastic net problems, the matrix object plays a crucial role in the performance of the solver.\n",
    "It becomes apparent in our optimization algorithm (and our benchmark analysis) \n",
    "that most of the runtime lies in interacting with the matrix object, e.g. computing inner-products.\n",
    "Hence, a highly performant matrix object implementing a select set of methods that the solver requires\n",
    "will yield tremendous speed gains overall.\n",
    "In addition, we have found numerous examples where a matrix class admits some special structure\n",
    "that can be exploited for further speed and memory gains.\n",
    "One simple example is a large sparse matrix, which cannot fit in memory as a dense matrix.\n",
    "Another example is genomics datasets which are not only sparse, but only take on 3 possible integer values\n",
    "(see [Examples](./examples.ipynb)), and generally have over 160 billion entries with 30\\% non-zero entries.\n",
    "\n",
    "For these reasons, we found it fruitful to abstract out the matrix class.\n",
    "`adelie` provides a few matrix classes in the `ad.matrix` submodule.\n",
    "We discuss below some ways a user may interact with these classes\n",
    "as well as define a class of their own to plug into our solver."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "import adelie as ad\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## __Naive Matrix Class__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The _naive matrix class_ refers to matrix classes that abstract out the feature matrix $X$.\n",
    "The simplest example of such a matrix is simply a dense matrix.\n",
    "Let us first construct a dense matrix and wrap it using `ad.matrix.dense`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 100\n",
    "p = 1000\n",
    "seed = 0\n",
    "\n",
    "np.random.seed(seed)\n",
    "X = np.random.normal(0, 1, (n, p))\n",
    "X_wrap = ad.matrix.dense(X, method=\"naive\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`X_wrap` can be thought of a simple wrapper of `X`, only exposing a few methods that our solver requires.\n",
    "For example, `.bmul()` is a method that computes `X[:, i:i+q].T @ (w * v)`.\n",
    "The canonical application of `.bmul()` is in computing the correlation between the feature matrix and the residual with observation weights `w`.\n",
    "It is worth mentioning that `adelie` also \n",
    "relies on such member functions when computing diagnostic quantities \n",
    "in `ad.diagnostic`.\n",
    "\n",
    "As an example, we generate data below and show the equivalence of calling `.bmul()` \n",
    "and the equivalent numpy code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "i = 10      # starting column index\n",
    "q = 3       # column block size\n",
    "\n",
    "w = np.random.uniform(0, 1, n)\n",
    "v = np.random.normal(0, 1, n)\n",
    "out = np.empty(q)\n",
    "X_wrap.bmul(i, q, v, w, out)\n",
    "assert np.allclose(\n",
    "    out,\n",
    "    X[:, i:i+q].T @ (w * v),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the full set of methods, we refer the readers to \n",
    "[MatrixNaiveBase64](https://jamesyang007.github.io/adelie/generated/adelie.matrix.MatrixNaiveBase64.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## __Covariance Matrix Class__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The _covariance matrix class_ refers to matrix classes that abstract out the covariance matrix \n",
    "$A = X^\\top X$.\n",
    "This matrix is currently only used in the context of `ad.gaussian_cov` solver.\n",
    "Nonetheless, like the naive matrix class, it exposes its own set of member functions that\n",
    "the covariance method solver requires.\n",
    "For example, it also exposes `.bmul()` but computes a different quantity: `A[:, subset].T @ v` where `subset` is a subset of column indices.\n",
    "\n",
    "We take the same data as above and show the equivalence of calling `.bmul()` and the equivalent numpy code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "A = X.T @ X\n",
    "A_wrap = ad.matrix.dense(A, method=\"cov\")\n",
    "\n",
    "i, j = 2, 5     # starting (i, j) position of the block of A\n",
    "p, q = 3, 7     # number of rows/cols of the block\n",
    "\n",
    "subset = np.arange(j, j + q)\n",
    "indices = np.arange(i, i + p)\n",
    "values = np.random.normal(0, 1, p)\n",
    "out = np.empty(q)\n",
    "A_wrap.bmul(subset, indices, values, out)\n",
    "assert np.allclose(\n",
    "    out,\n",
    "    A[i:i+p, j:j+q].T @ values,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the full set of methods, we refer the readers to \n",
    "[MatrixCovBase64](https://jamesyang007.github.io/adelie/generated/adelie.matrix.MatrixCovBase64.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## __Custom Matrix Class__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One of the defining features of `adelie` is the flexibility for the user to specify her own matrix classes.\n",
    "The custom matrix class can be implemented in either C++ _or_ Python!\n",
    "In this section, we demonstrate this feature by defining a custom naive matrix class \n",
    "equivalent to `ad.matrix.dense` in Python.\n",
    "All of the discussion carries through for the covariance matrix class\n",
    "and any important differences will be mentioned in passing.\n",
    "\n",
    "We first show the full code for our custom matrix class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Dense(ad.matrix.MatrixNaiveBase64):\n",
    "    def __init__(self, mat):\n",
    "        self.mat = mat\n",
    "        # MUST call base class __init__!\n",
    "        ad.matrix.MatrixNaiveBase64.__init__(self)\n",
    "    def bmul(self, j, q, v, w, out):\n",
    "        out[...] = self.mat[:, j:j+q].T @ (w * v)\n",
    "    def btmul(self, j, q, v, out):\n",
    "        out[...] += self.mat[:, j:j+q] @ v\n",
    "    def cmul(self, j, v, w):\n",
    "        return self.mat[:, j] @ (w * v)\n",
    "    def ctmul(self, j, v, out):\n",
    "        out[...] += self.mat[:, j] * v\n",
    "    def rows(self):\n",
    "        return self.mat.shape[0]\n",
    "    def cols(self):\n",
    "        return self.mat.shape[1]\n",
    "    def cov(self, j, q, sqrt_weights, out, buffer):\n",
    "        buffer = sqrt_weights[:, None] * self.mat[:, j:j+q] # just to demonstrate use of buffer\n",
    "        out[...] = buffer.T @ buffer\n",
    "    def mul(self, v, w, out):\n",
    "        out[...] = self.mat.T @ (w * v)\n",
    "    def sp_btmul(self, v, out):\n",
    "        out[...] = v @ self.mat"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We remark on a few important points:\n",
    "\n",
    "- The custom (naive) matrix class must inherit from the provided base class `ad.matrix.MatrixNaiveBase64`.\n",
    "  For users interested in using 32-bit floats, inherit from `ad.matrix.MatrixNaiveBase32`,\n",
    "  however beware of the numerical instability in using 32-bit floats!\n",
    "  You may observe strange behaviors in the optimizer, so _use with caution_!\n",
    "  Users who are implementing a covariance matrix class must inherit from `ad.matrix.MatrixCovBase64`\n",
    "  or `ad.matrix.MatrixCovBase32`.\n",
    "- The base class constructor must get called _exactly_ as shown above _without the use of_ `super()`!\n",
    "  This is a quirk of the `pybind11` package which is what we rely on for exporting C++ classes to Python.\n",
    "- Many of the member functions are given the output container to store the result of computing a quantity.\n",
    "  Hence, we use the syntax `out[...] = expression` to modify the container _in-place_.\n",
    "  A common pitfall is to write `out = expression`, which will only redirect the local variable `out` \n",
    "  to point to a different location in memory.\n",
    "  This low-level interface is done for memory efficiency."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now show that the matrix can be passed into our existing solver `ad.grpnet` \n",
    "with no further changes to the code.\n",
    "We first generate a response vector and call `ad.grpnet`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = X[:, -1] * np.random.normal(0, 1) + np.random.normal(0, 1, n)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 52%|\u001b[1;32m█\u001b[0m\u001b[1;32m█\u001b[0m\u001b[1;32m█\u001b[0m\u001b[1;32m█\u001b[0m\u001b[1;32m█\u001b[0m     | 52/100 [00:00:00<00:00:00, 535.72it/s] [dev:90.2%]\n"
     ]
    }
   ],
   "source": [
    "state = ad.grpnet(\n",
    "    X=Dense(X),\n",
    "    glm=ad.glm.gaussian(y=y),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can compare solutions by comparing with the solutions from passing in a dense matrix,\n",
    "which will internally get wrapped using `ad.matrix.dense`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 52%|\u001b[1;32m█\u001b[0m\u001b[1;32m█\u001b[0m\u001b[1;32m█\u001b[0m\u001b[1;32m█\u001b[0m\u001b[1;32m█\u001b[0m     | 52/100 [00:00:00<00:00:00, 13808.06it/s] [dev:90.2%]\n"
     ]
    }
   ],
   "source": [
    "state_exp = ad.grpnet(\n",
    "    X=X,\n",
    "    glm=ad.glm.gaussian(y=y),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert np.allclose(\n",
    "    state.betas.toarray(),\n",
    "    state_exp.betas.toarray(),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The convenience of defining a matrix class in Python has a performance cost.\n",
    "If the matrix size is sufficiently large that the member functions become expensive,\n",
    "the cost of dispatching to the Python interpreter becomes negligible.\n",
    "For the best performance, users may be interested in porting the code to C++ and exporting the bindings.\n",
    "We refer the readers to the [matrix](https://github.com/JamesYang007/adelie/tree/main/adelie/src/include/adelie_core/matrix) \n",
    "directory in our C++ backend for the numerous examples of how to extend the base classes in C++.\n",
    "The associated binding code can be found in [matrix.cpp](https://github.com/JamesYang007/adelie/blob/main/adelie/src/matrix.cpp)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "adelie",
   "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.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}