{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Basic Matrix"
   ],
   "id": "site-title"
  },
  {
   "cell_type": "markdown",
   "id": "23baedf9",
   "metadata": {},
   "source": "## Basic Matrix Operations with JAX"
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d63d1dc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax as J\n",
    "import jax.numpy as jnp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1aefad8d",
   "metadata": {},
   "source": "### Addition and Subtraction"
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "38d8d670",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[4, 5, 6, 8, 9],\n",
       "       [3, 5, 7, 8, 3],\n",
       "       [1, 4, 6, 0, 2]], dtype=int32)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A = jnp.array(\n",
    "    [\n",
    "        [4, 5, 6, 8, 9],\n",
    "        [3, 5, 7, 8, 3],\n",
    "        [1, 4, 6, 0, 2]\n",
    "    ]\n",
    ")\n",
    "\n",
    "A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9dbf9339",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[10,  1,  6,  8,  9],\n",
       "       [ 3, 12,  7,  8,  3],\n",
       "       [ 1,  4,  6, -1,  2]], dtype=int32)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "B = jnp.array(\n",
    "    [\n",
    "        [10, 1, 6, 8, 9],\n",
    "        [3, 12, 7, 8, 3],\n",
    "        [1, 4, 6, -1, 2]\n",
    "    ]\n",
    ")\n",
    "\n",
    "B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "99bdd2d0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[14,  6, 12, 16, 18],\n",
       "       [ 6, 17, 14, 16,  6],\n",
       "       [ 2,  8, 12, -1,  4]], dtype=int32)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# addition\n",
    "\n",
    "A_plus_B = jnp.add(A, B)\n",
    "A_plus_B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3f7916e4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[-6,  4,  0,  0,  0],\n",
       "       [ 0, -7,  0,  0,  0],\n",
       "       [ 0,  0,  0,  1,  0]], dtype=int32)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# subtraction\n",
    "\n",
    "A_minus_B = jnp.subtract(A, B)\n",
    "A_minus_B"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40b6ab88",
   "metadata": {},
   "source": "### Multiplication"
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "1ae63254",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "A : (3, 5)  B : (3, 5)\n"
     ]
    }
   ],
   "source": [
    "# check shapes first!\n",
    "print(f\"A : {A.shape}  B : {B.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52f8dda9",
   "metadata": {},
   "source": "The dimensions of A and B don't satisfy the rules of Matrix multiplication. 3X5 does not multiply with 3X5.\nSo let's create a new matrix here  with 5X3 dimension."
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "73e2523c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[1, 2, 3],\n",
       "       [2, 3, 5],\n",
       "       [3, 4, 6],\n",
       "       [4, 5, 7],\n",
       "       [5, 6, 8]], dtype=int32)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "C = jnp.array(\n",
    "    [\n",
    "        [1, 2, 3],\n",
    "        [2, 3, 5],\n",
    "        [3, 4, 6],\n",
    "        [4, 5, 7],\n",
    "        [5, 6, 8],\n",
    "    ]\n",
    ")\n",
    "\n",
    "C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "8a737f43",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[109, 141, 201],\n",
       "       [ 81, 107, 156],\n",
       "       [ 37,  50,  75]], dtype=int32)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A_mult_C = jnp.dot(A, C)\n",
    "A_mult_C"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b58c0ed9",
   "metadata": {},
   "source": "### Inverting\n\nErm, we need a invertible matrix here. You can't just invert every matrix out there! An invertible matrix has two properties:\n\n1. It's a square matrix (equal number of rows and columns)\n2. The determinant of the matrix is NOT 0\n\n[What is an invertible matrix](https://www.studypug.com/algebra-help/2-x-2-invertible-matrix#:~:text=An%20invertible%20matrix%20is%20a,the%20matrix%20is%20not%200.)\n\nLuckily `A_mult_C` is a square matrix. Let's see if this can be inverted. "
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "89172e7e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Invertible\n"
     ]
    }
   ],
   "source": [
    "D = A_mult_C\n",
    "\n",
    "determinant = jnp.linalg.det(D)\n",
    "\n",
    "if determinant != 0:\n",
    "    print(\"Invertible\")\n",
    "else:\n",
    "    print(\"Not Invertible\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "0eac1d08",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[ 2.4193413 , -5.645129  ,  5.258035  ],\n",
       "       [-3.2580476 ,  7.9354434 , -7.7741566 ],\n",
       "       [ 0.97848994, -2.5053654 ,  2.6021404 ]], dtype=float32)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# time to invert\n",
    "inverse = jnp.linalg.inv(D)\n",
    "inverse"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}