03c. Using MKL and Eigen.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e44b3d7b",
   "metadata": {
    "papermill": {
     "duration": 0.014011,
     "end_time": "2023-07-13T16:18:29.283118",
     "exception": false,
     "start_time": "2023-07-13T16:18:29.269107",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Using MKL and Eigen"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac865d86",
   "metadata": {
    "papermill": {
     "duration": 0.006005,
     "end_time": "2023-07-13T16:18:29.298131",
     "exception": false,
     "start_time": "2023-07-13T16:18:29.292126",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Table of contents**<a id='toc0_'></a>    \n",
    "- 1. [Imports](#toc1_)    \n",
    "- 2. [Setup and settings](#toc2_)    \n",
    "- 3. [Compile C++](#toc3_)    \n",
    "- 4. [Matrix multiplication](#toc4_)    \n",
    "- 5. [Solve linear equation system](#toc5_)    \n",
    "- 6. [Sparse matrices](#toc6_)    \n",
    "- 7. [Clean up](#toc7_)    \n",
    "\n",
    "<!-- vscode-jupyter-toc-config\n",
    "\tnumbering=true\n",
    "\tanchor=true\n",
    "\tflat=false\n",
    "\tminLevel=2\n",
    "\tmaxLevel=6\n",
    "\t/vscode-jupyter-toc-config -->\n",
    "<!-- THIS CELL WILL BE REPLACED ON TOC UPDATE. DO NOT WRITE YOUR TEXT IN THIS CELL -->"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f27105e-5ab8-4fee-9c7e-0b40be4e6d79",
   "metadata": {
    "papermill": {
     "duration": 0.007006,
     "end_time": "2023-07-13T16:18:29.311141",
     "exception": false,
     "start_time": "2023-07-13T16:18:29.304135",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "This notebooks shows how to speed up matrix multiplication and solving (sparse) linear equation systems with the **Math Kernal Library (MKL)** from Intel and **[Eigen](https://eigen.tuxfamily.org/index.php?title=Main_Page)**. Using explicit parallization with OpenMP in particular provide a significant speed-up."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cd9f66d",
   "metadata": {
    "papermill": {
     "duration": 0.007005,
     "end_time": "2023-07-13T16:18:29.324151",
     "exception": false,
     "start_time": "2023-07-13T16:18:29.317146",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Timing computer:** Windows 10 computer with 2x Intel(R) Xeon(R) Gold 6254 3.10 GHz CPUs (18 cores, 36 logical processes each) and 768 GB of RAM."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "346693a2",
   "metadata": {
    "papermill": {
     "duration": 0.006005,
     "end_time": "2023-07-13T16:18:29.336161",
     "exception": false,
     "start_time": "2023-07-13T16:18:29.330156",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## 1. <a id='toc1_'></a>[Imports](#toc0_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0aabed03",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-13T16:18:29.351173Z",
     "iopub.status.busy": "2023-07-13T16:18:29.351173Z",
     "iopub.status.idle": "2023-07-13T16:18:29.741490Z",
     "shell.execute_reply": "2023-07-13T16:18:29.741490Z"
    },
    "papermill": {
     "duration": 0.40332,
     "end_time": "2023-07-13T16:18:29.745493",
     "exception": false,
     "start_time": "2023-07-13T16:18:29.342173",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import os\n",
    "import shutil\n",
    "import urllib.request\n",
    "import numpy as np\n",
    "import scipy.linalg\n",
    "import scipy.sparse\n",
    "import scipy.sparse.linalg\n",
    "import scipy.io\n",
    "\n",
    "import time\n",
    "\n",
    "from EconModel import cpptools"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ec47644",
   "metadata": {
    "papermill": {
     "duration": 0.007007,
     "end_time": "2023-07-13T16:18:29.766511",
     "exception": false,
     "start_time": "2023-07-13T16:18:29.759504",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## 2. <a id='toc2_'></a>[Setup and settings](#toc0_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "dc574d96-451d-4d4d-aa51-1639a4619505",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-13T16:18:29.781523Z",
     "iopub.status.busy": "2023-07-13T16:18:29.780523Z",
     "iopub.status.idle": "2023-07-13T16:18:29.845576Z",
     "shell.execute_reply": "2023-07-13T16:18:29.845576Z"
    },
    "papermill": {
     "duration": 0.076063,
     "end_time": "2023-07-13T16:18:29.848579",
     "exception": false,
     "start_time": "2023-07-13T16:18:29.772516",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Eigen succesfully installed\n"
     ]
    }
   ],
   "source": [
    "cpptools.setup_Eigen(do_print=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "846d6606-3e3e-4cce-b2ec-afe02a7be263",
   "metadata": {
    "papermill": {
     "duration": 0.007006,
     "end_time": "2023-07-13T16:18:29.869596",
     "exception": false,
     "start_time": "2023-07-13T16:18:29.862590",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Settings:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8c7bc56b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-13T16:18:29.884608Z",
     "iopub.status.busy": "2023-07-13T16:18:29.883607Z",
     "iopub.status.idle": "2023-07-13T16:18:29.945702Z",
     "shell.execute_reply": "2023-07-13T16:18:29.945702Z"
    },
    "papermill": {
     "duration": 0.074104,
     "end_time": "2023-07-13T16:18:29.949705",
     "exception": false,
     "start_time": "2023-07-13T16:18:29.875601",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "threads_vec = [(1,8),(2,4),(4,4)]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3947fe3f",
   "metadata": {
    "papermill": {
     "duration": 0.007006,
     "end_time": "2023-07-13T16:18:29.970678",
     "exception": false,
     "start_time": "2023-07-13T16:18:29.963672",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## 3. <a id='toc3_'></a>[Compile C++](#toc0_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7dfb74ba",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-13T16:18:29.985690Z",
     "iopub.status.busy": "2023-07-13T16:18:29.984690Z",
     "iopub.status.idle": "2023-07-13T16:18:42.496867Z",
     "shell.execute_reply": "2023-07-13T16:18:42.496867Z"
    },
    "papermill": {
     "duration": 12.524188,
     "end_time": "2023-07-13T16:18:42.500870",
     "exception": false,
     "start_time": "2023-07-13T16:18:29.976682",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "flags = '/LD /EHsc /O3 /openmp /std:c++17 /Qmkl /Icppfuncs/'\n",
    "#flags = '/LD /EHsc /O3 /openmp /std:c++17 /Qmkl /QxCOMMON-AVX512 /Icppfuncs/'\n",
    "example_MKL = cpptools.link_to_cpp('cppfuncs/example_MKL.cpp',\n",
    "                                   options={'compiler':'intel','flags':flags})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1e6f1fa",
   "metadata": {
    "papermill": {
     "duration": 0.010008,
     "end_time": "2023-07-13T16:18:42.523889",
     "exception": false,
     "start_time": "2023-07-13T16:18:42.513881",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## 4. <a id='toc4_'></a>[Matrix multiplication](#toc0_)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccbf6456-b70b-4322-be01-025caa4a1951",
   "metadata": {
    "papermill": {
     "duration": 0.007006,
     "end_time": "2023-07-13T16:18:42.536900",
     "exception": false,
     "start_time": "2023-07-13T16:18:42.529894",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "$$\n",
    "C = \\alpha A B + \\beta C\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c742bf68-a3c0-45f4-aa33-3d9e7572de67",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-13T16:18:42.550911Z",
     "iopub.status.busy": "2023-07-13T16:18:42.550911Z",
     "iopub.status.idle": "2023-07-13T16:18:42.609959Z",
     "shell.execute_reply": "2023-07-13T16:18:42.609959Z"
    },
    "papermill": {
     "duration": 0.071057,
     "end_time": "2023-07-13T16:18:42.613962",
     "exception": false,
     "start_time": "2023-07-13T16:18:42.542905",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "alpha = 1.0\n",
    "beta = 0.0\n",
    "    \n",
    "def allocate_mat(n=300,nrep=500):\n",
    "    \n",
    "    # varying sizes to test code\n",
    "    \n",
    "    np.random.seed(17)\n",
    "    \n",
    "    A = np.random.uniform(size=(nrep,2*n,n))\n",
    "    B = np.random.uniform(size=(nrep,n,3*n))\n",
    "    C = np.random.uniform(size=(nrep,2*n,3*n))\n",
    "    \n",
    "    return A,B,C\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "43a973a3-8c67-4b79-ab1b-5eef9830bab7",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-13T16:18:42.635979Z",
     "iopub.status.busy": "2023-07-13T16:18:42.634978Z",
     "iopub.status.idle": "2023-07-13T16:18:42.691024Z",
     "shell.execute_reply": "2023-07-13T16:18:42.691024Z"
    },
    "papermill": {
     "duration": 0.067097,
     "end_time": "2023-07-13T16:18:42.694070",
     "exception": false,
     "start_time": "2023-07-13T16:18:42.626973",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def matmul_np(alpha,beta,A,B,C):\n",
    "        \n",
    "    t0 = time.perf_counter()\n",
    "    C[:,:,:] = alpha*(A@B)+beta*C\n",
    "    secs = time.perf_counter()-t0\n",
    "    \n",
    "    return secs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "92681220-829b-4fbd-9bd3-242f1d67b979",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-13T16:18:42.716044Z",
     "iopub.status.busy": "2023-07-13T16:18:42.716044Z",
     "iopub.status.idle": "2023-07-13T16:18:42.771089Z",
     "shell.execute_reply": "2023-07-13T16:18:42.771089Z"
    },
    "papermill": {
     "duration": 0.068055,
     "end_time": "2023-07-13T16:18:42.775092",
     "exception": false,
     "start_time": "2023-07-13T16:18:42.707037",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def matmul_BLAS(alpha,beta,A,B,C,threads_internal,threads_omp):\n",
    "    \n",
    "    nrep,m,k = A.shape\n",
    "    nrep,k,n = B.shape\n",
    "    nrep,m,n = C.shape\n",
    "    \n",
    "    secs = example_MKL.matmul_BLAS(nrep,m,n,k,alpha,beta,A,B,C,threads_internal,threads_omp)\n",
    "    \n",
    "    return secs\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e6d1a551-8003-40c7-8b94-a15f6fcdd52d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-13T16:18:42.799112Z",
     "iopub.status.busy": "2023-07-13T16:18:42.799112Z",
     "iopub.status.idle": "2023-07-13T16:18:42.855157Z",
     "shell.execute_reply": "2023-07-13T16:18:42.855157Z"
    },
    "papermill": {
     "duration": 0.071057,
     "end_time": "2023-07-13T16:18:42.859160",
     "exception": false,
     "start_time": "2023-07-13T16:18:42.788103",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def matmul_Eigen(alpha,beta,A,B,C,threads_internal,threads_omp):\n",
    "    \n",
    "    nrep,m,k = A.shape\n",
    "    nrep,k,n = B.shape\n",
    "    nrep,m,n = C.shape\n",
    "    \n",
    "    secs = example_MKL.matmul_Eigen(nrep,m,n,k,alpha,beta,A,B,C,threads_internal,threads_omp)\n",
    "    \n",
    "    return secs\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d182b534-303d-42e9-9259-476a667c5534",
   "metadata": {
    "papermill": {
     "duration": 0.007964,
     "end_time": "2023-07-13T16:18:42.879176",
     "exception": false,
     "start_time": "2023-07-13T16:18:42.871212",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Test with single matrix multiplication:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "3a8d50fd-8c4a-43a9-bc13-b4ca001d8b0e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-13T16:18:42.893188Z",
     "iopub.status.busy": "2023-07-13T16:18:42.892187Z",
     "iopub.status.idle": "2023-07-13T16:18:51.904524Z",
     "shell.execute_reply": "2023-07-13T16:18:51.903522Z"
    },
    "papermill": {
     "duration": 9.023345,
     "end_time": "2023-07-13T16:18:51.908526",
     "exception": false,
     "start_time": "2023-07-13T16:18:42.885181",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "n = 500:\n",
      "numpy:                               98 ms\n",
      "mkl (no threads, no omp):            50 ms\n",
      "mkl (all internal, no omp):          21 ms\n",
      "Eigen (no threads, no omp):          71 ms\n",
      "Eigen (all internal, no omp):        39 ms\n",
      "\n",
      "n = 1000:\n",
      "numpy:                              233 ms\n",
      "mkl (no threads, no omp):           377 ms\n",
      "mkl (all internal, no omp):         140 ms\n",
      "Eigen (no threads, no omp):         387 ms\n",
      "Eigen (all internal, no omp):       204 ms\n",
      "\n",
      "n = 2000:\n",
      "numpy:                             1897 ms\n",
      "mkl (no threads, no omp):          2459 ms\n",
      "mkl (all internal, no omp):        1423 ms\n",
      "Eigen (no threads, no omp):        3140 ms\n",
      "Eigen (all internal, no omp):      1572 ms\n",
      "\n"
     ]
    }
   ],
   "source": [
    "threads = 4\n",
    "for n in [500,1000,2000]:\n",
    "    \n",
    "    print(f'{n = }:') \n",
    "    \n",
    "    # numpy\n",
    "    A,B,C = allocate_mat(n,1)\n",
    "    name = 'numpy:'\n",
    "    secs = matmul_np(alpha,beta,A,B,C)\n",
    "    print(f'{name:30s} {secs*1000:8.0f} ms')\n",
    "\n",
    "    # mkl\n",
    "    A,B,C = allocate_mat(n,1)\n",
    "    name = f'mkl (no threads, no omp):'\n",
    "    secs = matmul_BLAS(alpha,beta,A,B,C,1,1)\n",
    "    print(f'{name:30s} {secs*1000:8.0f} ms')\n",
    "\n",
    "    A,B,C = allocate_mat(n,1)\n",
    "    name = f'mkl (all internal, no omp):'\n",
    "    secs = matmul_BLAS(alpha,beta,A,B,C,threads,1)\n",
    "    print(f'{name:30s} {secs*1000:8.0f} ms')\n",
    "    \n",
    "    # Eigen\n",
    "    A,B,C = allocate_mat(n,1)\n",
    "    name = f'Eigen (no threads, no omp):'\n",
    "    secs = matmul_Eigen(alpha,beta,A,B,C,1,1)\n",
    "    print(f'{name:30s} {secs*1000:8.0f} ms')\n",
    "\n",
    "    A,B,C = allocate_mat(n,1)\n",
    "    name = f'Eigen (all internal, no omp):'\n",
    "    secs = matmul_Eigen(alpha,beta,A,B,C,threads,1)\n",
    "    print(f'{name:30s} {secs*1000:8.0f} ms')    \n",
    "    \n",
    "    print('')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44023737",
   "metadata": {
    "papermill": {
     "duration": 0.007008,
     "end_time": "2023-07-13T16:18:51.929544",
     "exception": false,
     "start_time": "2023-07-13T16:18:51.922536",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Test with multiple matrix multiplications:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "8a5bb951",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-13T16:18:51.945559Z",
     "iopub.status.busy": "2023-07-13T16:18:51.944558Z",
     "iopub.status.idle": "2023-07-13T16:19:18.455115Z",
     "shell.execute_reply": "2023-07-13T16:19:18.454114Z"
    },
    "papermill": {
     "duration": 26.525573,
     "end_time": "2023-07-13T16:19:18.461121",
     "exception": false,
     "start_time": "2023-07-13T16:18:51.935548",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "n = 200, nrep = 400:\n",
      "numpy:                             3063 ms\n",
      "mkl (internal= 1,omp= 8):           296 ms\n",
      "mkl (internal= 2,omp= 4):           513 ms\n",
      "mkl (internal= 4,omp= 4):           442 ms\n",
      "\n",
      "n = 400, nrep = 200:\n",
      "numpy:                            22113 ms\n",
      "mkl (internal= 1,omp= 8):          5986 ms\n",
      "mkl (internal= 2,omp= 4):          3385 ms\n",
      "mkl (internal= 4,omp= 4):          4630 ms\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for (n,nrep) in [(200,400),(400,200)]:\n",
    "    \n",
    "    print(f'{n = }, {nrep = }:') \n",
    "    \n",
    "    A,B,C = allocate_mat(n,nrep)\n",
    "    \n",
    "    # numpy\n",
    "    name = 'numpy:'\n",
    "    secs = matmul_np(alpha,beta,A,B,C)\n",
    "    print(f'{name:30s} {secs*1000:8.0f} ms')\n",
    "    C_np = C.copy()\n",
    "    \n",
    "    # mkl\n",
    "    for (threads_internal,threads_omp) in threads_vec:\n",
    "        \n",
    "        name = f'mkl (internal={threads_internal:2d},omp={threads_omp:2d}):'\n",
    "        C[:] = 0\n",
    "        secs = matmul_BLAS(alpha,beta,A,B,C,threads_internal,threads_omp)\n",
    "        print(f'{name:30s} {secs*1000:8.0f} ms')\n",
    "        \n",
    "    # check results are the same\n",
    "    max_abs_diff = np.max(np.abs(C_np-C))\n",
    "    assert np.isclose(max_abs_diff,0.0)   \n",
    "    \n",
    "    print('')\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fc1b15b",
   "metadata": {
    "papermill": {
     "duration": 0.011009,
     "end_time": "2023-07-13T16:19:18.487143",
     "exception": false,
     "start_time": "2023-07-13T16:19:18.476134",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## 5. <a id='toc5_'></a>[Solve linear equation system](#toc0_)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2fac8e7-2c4a-4677-8814-345e0075f4bf",
   "metadata": {
    "papermill": {
     "duration": 0.010008,
     "end_time": "2023-07-13T16:19:18.507159",
     "exception": false,
     "start_time": "2023-07-13T16:19:18.497151",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "$$\n",
    "Ax = b\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "abb147a9-003a-4846-9db6-c71df172c4b9",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-13T16:19:18.534179Z",
     "iopub.status.busy": "2023-07-13T16:19:18.532183Z",
     "iopub.status.idle": "2023-07-13T16:19:18.605238Z",
     "shell.execute_reply": "2023-07-13T16:19:18.605238Z"
    },
    "papermill": {
     "duration": 0.090072,
     "end_time": "2023-07-13T16:19:18.608240",
     "exception": false,
     "start_time": "2023-07-13T16:19:18.518168",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def allocate_solve_mat(n=500,nrhs=100,nrep=500):\n",
    "    \n",
    "    np.random.seed(17)\n",
    "\n",
    "    A = np.random.uniform(size=(nrep,n,n))\n",
    "    x = np.random.uniform(size=(nrep,n,nrhs))\n",
    "    b = A@x\n",
    "\n",
    "    return A,b,x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "0598b06c-df96-4877-aa33-7642de2d1165",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-13T16:19:18.633262Z",
     "iopub.status.busy": "2023-07-13T16:19:18.633262Z",
     "iopub.status.idle": "2023-07-13T16:19:18.701316Z",
     "shell.execute_reply": "2023-07-13T16:19:18.701316Z"
    },
    "papermill": {
     "duration": 0.080065,
     "end_time": "2023-07-13T16:19:18.704318",
     "exception": false,
     "start_time": "2023-07-13T16:19:18.624253",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def solve_np(A,b):\n",
    "    \n",
    "    t0 = time.perf_counter()\n",
    "    x = np.zeros(b.shape)\n",
    "    for i in range(A.shape[0]):\n",
    "        x[i,:,:] = np.linalg.solve(A[i],b[i])\n",
    "    secs = time.perf_counter()-t0\n",
    "    \n",
    "    return x,secs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "d3f5ec5a-ab84-4602-a5a5-298de0ab4344",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-13T16:19:18.729340Z",
     "iopub.status.busy": "2023-07-13T16:19:18.729340Z",
     "iopub.status.idle": "2023-07-13T16:19:18.787386Z",
     "shell.execute_reply": "2023-07-13T16:19:18.787386Z"
    },
    "papermill": {
     "duration": 0.072058,
     "end_time": "2023-07-13T16:19:18.791388",
     "exception": false,
     "start_time": "2023-07-13T16:19:18.719330",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def solve_LAPACKE(A,b,x,threads_internal,threads_omp):\n",
    "    \n",
    "    nrep,n,_ = A.shape\n",
    "    _,_,nrh = b.shape\n",
    "    \n",
    "    b_old = b.copy()\n",
    "    secs = example_MKL.solve_LAPACKE(nrep,n,nrh,A,b,threads_internal,threads_omp)\n",
    "    x[:] = b.copy()\n",
    "    b[:] = b_old\n",
    "    \n",
    "    return x,secs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "e0a50318-cc2e-4c51-985d-1e1a5097cd2a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-13T16:19:18.816413Z",
     "iopub.status.busy": "2023-07-13T16:19:18.815409Z",
     "iopub.status.idle": "2023-07-13T16:19:18.884465Z",
     "shell.execute_reply": "2023-07-13T16:19:18.884465Z"
    },
    "papermill": {
     "duration": 0.081066,
     "end_time": "2023-07-13T16:19:18.888469",
     "exception": false,
     "start_time": "2023-07-13T16:19:18.807403",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def solve_Eigen(A,b,x,threads_internal,threads_omp):\n",
    "    \n",
    "    nrep,n,_ = A.shape\n",
    "    _,_,nrh = b.shape\n",
    "    \n",
    "    b_old = b.copy()\n",
    "    secs = example_MKL.solve_Eigen(nrep,n,nrh,A,b,threads_internal,threads_omp)\n",
    "    x[:] = b.copy()\n",
    "    b[:] = b_old\n",
    "    \n",
    "    return x,secs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c81015b-937b-44e0-acd8-53e61ea0f317",
   "metadata": {
    "papermill": {
     "duration": 0.010007,
     "end_time": "2023-07-13T16:19:18.917493",
     "exception": false,
     "start_time": "2023-07-13T16:19:18.907486",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Test solving multiple equations systems at once:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "89192387",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-13T16:19:18.937512Z",
     "iopub.status.busy": "2023-07-13T16:19:18.936511Z",
     "iopub.status.idle": "2023-07-13T16:19:57.833147Z",
     "shell.execute_reply": "2023-07-13T16:19:57.833147Z"
    },
    "papermill": {
     "duration": 38.910649,
     "end_time": "2023-07-13T16:19:57.837150",
     "exception": false,
     "start_time": "2023-07-13T16:19:18.926501",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "numpy                              4737 ms\n",
      "mkl (internal= 1,omp= 8):          3072 ms\n",
      "mkl (internal= 2,omp= 4):          3806 ms\n",
      "mkl (internal= 4,omp= 4):          3147 ms\n",
      "Eigen (internal= 1,omp= 8):        2828 ms\n",
      "Eigen (internal= 2,omp= 4):        2911 ms\n",
      "Eigen (internal= 4,omp= 4):        5110 ms\n"
     ]
    }
   ],
   "source": [
    "# numpy\n",
    "name = 'numpy'\n",
    "A,b,x = allocate_solve_mat()\n",
    "x,secs = solve_np(A,b)\n",
    "print(f'{name:30s} {secs*1000:8.0f} ms')\n",
    "x_np = x.copy()\n",
    "\n",
    "# mkl\n",
    "for (threads_internal,threads_omp) in threads_vec:\n",
    "\n",
    "    name = f'mkl (internal={threads_internal:2d},omp={threads_omp:2d}):'\n",
    "    A,b,x = allocate_solve_mat()\n",
    "    x,secs = solve_LAPACKE(A,b,x,threads_internal,threads_omp)\n",
    "    print(f'{name:30s} {secs*1000:8.0f} ms')\n",
    "\n",
    "# check results are the same\n",
    "max_abs_diff = np.max(np.abs(x_np-x))\n",
    "assert max_abs_diff < 1e-7\n",
    "    \n",
    "# Eigen\n",
    "for (threads_internal,threads_omp) in threads_vec:\n",
    "\n",
    "    name = f'Eigen (internal={threads_internal:2d},omp={threads_omp:2d}):'\n",
    "    A,b,x = allocate_solve_mat()\n",
    "    x,secs = solve_Eigen(A,b,x,threads_internal,threads_omp)\n",
    "    print(f'{name:30s} {secs*1000:8.0f} ms')     \n",
    "    \n",
    "# check results are the same\n",
    "max_abs_diff = np.max(np.abs(x_np-x))\n",
    "assert max_abs_diff < 1e-7    \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a09de5e3",
   "metadata": {
    "papermill": {
     "duration": 0.009008,
     "end_time": "2023-07-13T16:19:57.860170",
     "exception": false,
     "start_time": "2023-07-13T16:19:57.851162",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## 6. <a id='toc6_'></a>[Sparse matrices](#toc0_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "8b975489-356c-44d0-838a-94a6e6315830",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-13T16:19:57.878186Z",
     "iopub.status.busy": "2023-07-13T16:19:57.877190Z",
     "iopub.status.idle": "2023-07-13T16:19:59.914837Z",
     "shell.execute_reply": "2023-07-13T16:19:59.914837Z"
    },
    "papermill": {
     "duration": 2.050664,
     "end_time": "2023-07-13T16:19:59.918840",
     "exception": false,
     "start_time": "2023-07-13T16:19:57.868176",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "nrep = 100\n",
    "\n",
    "# A\n",
    "url = 'https://math.nist.gov/pub/MatrixMarket2/SPARSKIT/tokamak/utm1700b.mtx.gz'\n",
    "urllib.request.urlretrieve(url,'matrices/utm1700b.mtx.gz')\n",
    "A = scipy.io.mmread('matrices/utm1700b.mtx.gz')\n",
    "nrow,ncol = A.shape\n",
    "\n",
    "# numpy\n",
    "A_np = np.broadcast_to(A.toarray(),(nrep,nrow,ncol)) \n",
    "\n",
    "# csr\n",
    "A_csr = scipy.sparse.csr_array(A)\n",
    "\n",
    "nnz = A_csr.data.size\n",
    "data = np.broadcast_to(A_csr.data,(nrep,nnz))\n",
    "indices = np.broadcast_to(A_csr.indices,(nrep,nnz))\n",
    "indptr = np.broadcast_to(A_csr.indptr,(nrep,ncol+1))\n",
    "\n",
    "# b\n",
    "url = 'https://math.nist.gov/pub/MatrixMarket2/SPARSKIT/tokamak/utm1700b_rhs1.mtx.gz'\n",
    "urllib.request.urlretrieve(url,'matrices/utm1700b_rhs1.mtx.gz')\n",
    "b = scipy.io.mmread('matrices/utm1700b_rhs1.mtx.gz')\n",
    "b = np.broadcast_to(b.ravel(),(nrep,nrow))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a77071a-6e2a-446d-af15-b6d102d1fd4c",
   "metadata": {
    "papermill": {
     "duration": 0.008006,
     "end_time": "2023-07-13T16:19:59.941859",
     "exception": false,
     "start_time": "2023-07-13T16:19:59.933853",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Test solving multiple equations systems at once:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "66229fb1",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-13T16:19:59.959876Z",
     "iopub.status.busy": "2023-07-13T16:19:59.958876Z",
     "iopub.status.idle": "2023-07-13T16:20:11.846546Z",
     "shell.execute_reply": "2023-07-13T16:20:11.846546Z"
    },
    "papermill": {
     "duration": 11.900683,
     "end_time": "2023-07-13T16:20:11.850549",
     "exception": false,
     "start_time": "2023-07-13T16:19:59.949866",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "numpy                         :    12682 ms\n",
      "spsolve                       :     2366 ms\n",
      "Eigen (internal= 1,omp= 8):   :      962 ms\n",
      "Eigen (internal= 2,omp= 4):   :     1259 ms\n",
      "Eigen (internal= 4,omp= 4):   :     1047 ms\n"
     ]
    }
   ],
   "source": [
    "# numpy\n",
    "name = 'numpy' \n",
    "x_np = np.zeros((nrep,nrow))\n",
    "\n",
    "t0 = time.perf_counter()\n",
    "for i in range(nrep):\n",
    "    x_np[i] = np.linalg.solve(A_np[i],b[i])\n",
    "secs = time.perf_counter()-t0\n",
    "print(f'{name:30s}: {secs*1000:8.0f} ms')\n",
    "\n",
    "for i in range(nrep): assert np.max(np.abs(A_np[i]@x_np[i]-b[i])) < 1e-7\n",
    "\n",
    "# spsolve\n",
    "name = 'spsolve' \n",
    "x_sp = np.zeros((nrep,nrow))\n",
    "\n",
    "t0 = time.perf_counter()\n",
    "for i in range(nrep):\n",
    "    x_sp[i] = scipy.sparse.linalg.spsolve(A_csr,b[i])\n",
    "secs = time.perf_counter()-t0\n",
    "print(f'{name:30s}: {secs*1000:8.0f} ms')\n",
    "\n",
    "for i in range(nrep): assert np.max(np.abs(A_csr@x_sp[i]-b[i])) < 1e-7\n",
    "\n",
    "# Eigen\n",
    "name = 'Eigen_ParadisoLU'\n",
    "x_Eigen = np.zeros((nrep,nrow))\n",
    "\n",
    "for (threads_internal,threads_omp) in threads_vec:\n",
    "    \n",
    "    name = f'Eigen (internal={threads_internal:2d},omp={threads_omp:2d}):'\n",
    "    secs = example_MKL.solve_sparse_Eigen(nrep,nrow,ncol,nnz,indptr,indices,data,b,x_Eigen,\n",
    "                                          threads_internal,threads_omp)\n",
    "    print(f'{name:30s}: {secs*1000:8.0f} ms')\n",
    "\n",
    "for i in range(nrep): assert np.max(np.abs(A_csr@x_Eigen[i]-b[i])) < 1e-7"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12e7dc37",
   "metadata": {
    "papermill": {
     "duration": 0.009007,
     "end_time": "2023-07-13T16:20:11.873568",
     "exception": false,
     "start_time": "2023-07-13T16:20:11.864561",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## 7. <a id='toc7_'></a>[Clean up](#toc0_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "f0e0274e-7ef5-409b-b739-67ddb2d93d38",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-13T16:20:11.891586Z",
     "iopub.status.busy": "2023-07-13T16:20:11.890584Z",
     "iopub.status.idle": "2023-07-13T16:20:11.978653Z",
     "shell.execute_reply": "2023-07-13T16:20:11.978653Z"
    },
    "papermill": {
     "duration": 0.101081,
     "end_time": "2023-07-13T16:20:11.982656",
     "exception": false,
     "start_time": "2023-07-13T16:20:11.881575",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "example_MKL.clean_up()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "cef23703-66ae-46e4-8e03-3b8ad6668121",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-13T16:20:12.008679Z",
     "iopub.status.busy": "2023-07-13T16:20:12.007678Z",
     "iopub.status.idle": "2023-07-13T16:20:12.289905Z",
     "shell.execute_reply": "2023-07-13T16:20:12.289905Z"
    },
    "papermill": {
     "duration": 0.295238,
     "end_time": "2023-07-13T16:20:12.293908",
     "exception": false,
     "start_time": "2023-07-13T16:20:11.998670",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "os.remove(f'cppfuncs/Eigen-main.zip')\n",
    "shutil.rmtree(f'cppfuncs/Eigen/')"
   ]
  }
 ],
 "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.12.4"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 106.086289,
   "end_time": "2023-07-13T16:20:13.235673",
   "environment_variables": {},
   "exception": null,
   "input_path": "03c. Using MKL and Eigen.ipynb",
   "output_path": "temp_notebooks/03c. Using MKL and Eigen.ipynb",
   "parameters": {},
   "start_time": "2023-07-13T16:18:27.149384",
   "version": "2.4.0"
  },
  "toc-autonumbering": true
 },
 "nbformat": 4,
 "nbformat_minor": 5
}