{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lLMDGooqcWGg"
   },
   "source": [
    "# Problem Set 0: Introduction to Python, Matrices and NumPy\n",
    "\n",
    "**Release Date:** 15 August 2023\n",
    "\n",
    "**Due Date:** 23:59, 26 August 2023"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gFWBI2aWcsOM"
   },
   "source": [
    "# Overview\n",
    "\n",
    "Python is the most common language used for modern AI applications. Modern machine learning also often involves linear algebra. We will not be teaching either Python or linear algebra explictly in CS2109S. Instead, we have designed this problem set to help you to become familiar with Python and review (or learn) some fundamental matrix operations. We will also learn how to use some useful functions in NumPy, a Python package that allows us to easily manipulate multidimensional arrays (like vectors and matrices).\n",
    "\n",
    "This problem set will not be fully comprehensive. There will likely still be more things that you will need to learn on your own. However, this is a feature and not a bug. The high-level goal of CS2109S is to learn how to learn new things on your own! However, this problem set will hopefully be a helpful primer for your learning journey.\n",
    "\n",
    "Welcome to CS2109S!\n",
    "\n",
    "Required Files:\n",
    "\n",
    "* OxCGRT_2020.csv\n",
    "* prepare_data.py\n",
    "\n",
    "Optional File:\n",
    "* ps0.py\n",
    "\n",
    "**Honour Code**: Note that plagiarism will not be condoned! You may discuss with your classmates and check the internet for references, but you MUST NOT submit code/report that is copied directly from other sources!\n",
    "\n",
    "**IMPORTANT**: While it is possible to write and run Python code directly in Jupyter notebook, we recommend that you do this Problem set with an iDE using the `.py` file provided. An IDE will make debugging significantly easier. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3zYCsHsde_Dc"
   },
   "source": [
    "# Part 1: Basic Python and Matrix Operations\n",
    "We will begin this part of the homework by taking a look at Python’s basic features and syntax. Then, we shall put what we have learnt into practice by implementing several common matrix operations. You will never be required to implement matrix operations for real work. We are making you do this to give you some practice with Python. However, a good understanding of matrices and how their operations actually work will likely be useful in understanding the material to be covered in the later part of this module!\n",
    "\n",
    "**Note**: feel free to skip to Matrix Operations in Python if you already have a good grasp of Python."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LNTITqxCfS8c"
   },
   "source": [
    "## 1.0 Introduction to Python\n",
    "In this section, we will look at the basics of Python. If you are keen on finding out more about it, do take a look at https://docs.python.org/3.9/tutorial/index.html."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "v1sxzgdefXjs"
   },
   "source": [
    "### 1.0.1 Basic Data Types"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vD3FHQsMggMG"
   },
   "source": [
    "*Number*\n",
    "\n",
    "Floats and integers in Python behave similarly to those in other programming languages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "7stw_ROEdd1K"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2110\n",
      "2108\n",
      "4218\n",
      "1054.5\n",
      "1054\n",
      "1\n"
     ]
    }
   ],
   "source": [
    "x = 2109        # Declares and assigns a value to the variable x\n",
    "print(x + 1)   # Addition ; prints 2110\n",
    "print(x - 1)   # Subtraction ; prints 2108\n",
    "print(x * 2)   # Multiplication ; prints 4218\n",
    "print(x / 2)   # Floating point division ; prints 1054.5\n",
    "print(x // 2)  # Integer division ; prints 1054\n",
    "print(x % 2)   # Modulo division ; prints 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "T9k52oX1f6MN"
   },
   "source": [
    "*Boolean*\n",
    "\n",
    "The Boolean operators are as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "aYakFamqfltm"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "False\n",
      "True\n",
      "False\n"
     ]
    }
   ],
   "source": [
    "a = True        # Assigns true to variable a\n",
    "b = False       # Assigns false to variable b\n",
    "print(a and b) # Logical and; prints False\n",
    "print(a or b)  # Logical or; prints True\n",
    "print(not a)   # Logical negation ; prints False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "chDjvRWvgDjk"
   },
   "source": [
    "*Strings*\n",
    "\n",
    "Similar to numbers, strings in Python behave similarly to those in other programming languages. Note that in Python, single and double quotations can be used to indicate that a sequence of characters is a string. The following are both valid strings:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "wCyeFOzBgDDy"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hello\n",
      "world\n"
     ]
    }
   ],
   "source": [
    "s1 = 'hello'\n",
    "s2 = \"world\"\n",
    "print(s1)\n",
    "print(s2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NMC1anYogXfk"
   },
   "source": [
    "For more information on the types of operations we can perform on strings in Python, please refer to https://docs.python.org/3.9/library/stdtypes.html#string-methods."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DodDX1vagkhu"
   },
   "source": [
    "### 1.0.2 A Few Other Data Types\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6h4Cx3e0gt0X"
   },
   "source": [
    "*Lists*\n",
    "\n",
    "One of the data structures which we will use is lists. Note that they behave more like dynamic arrays than linked lists in other programming languages. The following shows examples of common operations on lists in Python.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "XYgsMsS7gUhL"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3\n",
      "3\n",
      "4 foo\n",
      "[None, None, None]\n"
     ]
    }
   ],
   "source": [
    "arr = [1, 2, 3]     # Creates a list\n",
    "print(arr[2])       # Accesses the element at index 2 (0- indexed); prints 3\n",
    "print(arr[-1])      # Accesses the element at the last index ; prints 3\n",
    "arr[1] = 'foo'     # Re - assigns the value at index 1 to 'foo'\n",
    "arr.append(4)       # Adds a new element 4 to the end of the list\n",
    "x = arr.pop()      # Removes and returns the last element\n",
    "y = arr.pop(1)     # Removes and returns the element at index 1\n",
    "print(x, y)         # Prints '4 foo'\n",
    "arr = [ None ] * 3  # Creates the list [None, None, None]\n",
    "print(arr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qtU7hCzlhnqq"
   },
   "source": [
    "Note that it is possible to access and assign a sublist in python by slicing. For example,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "QO2f7b54hETZ"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2, 3]\n",
      "[3, 4, 5]\n",
      "[1, 2, 3]\n",
      "[1, 2, 5]\n"
     ]
    }
   ],
   "source": [
    "arr = [1, 2, 3, 4, 5] # Creates a list\n",
    "print(arr[1:3])       # Prints [2, 3]\n",
    "print(arr[2:])        # Prints [3, 4, 5]\n",
    "print(arr[:3])        # Prints [1, 2, 3]\n",
    "arr[2:] = [5]        \n",
    "print(arr)            # Prints [1, 2, 5]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Yf4NB01OhpJa"
   },
   "source": [
    "To find out more about lists in Python, refer to https://docs.python.org/3.9/tutorial/datastructures.html#more-on-lists.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "L4c--tvKhpOP"
   },
   "source": [
    "*Tuples*\n",
    "\n",
    "As you work with Python, you may encounter tuples. A tuple behaves like a list in some ways. However, unlike a list, it is immutable. You will unlikely be using tuples when you write code, but tuples are sometimes returned by some python functions. You can think of them as lists and use them like lists. Just don’t try to modify them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "id": "NzoAX57miDR_"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 ... cool\n"
     ]
    }
   ],
   "source": [
    "t = (1 , 'cool')                  # Declares a tuple containing two elements\n",
    "print (t[0], t[1], sep=\" ... \") # Prints \"1 ... cool \""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LrwA3lL1i9Dt"
   },
   "source": [
    "### 1.0.3 Loops"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4shMhxbBjDU9"
   },
   "source": [
    "*While Loops*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "id": "epwhKYnWi5ua"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "01234"
     ]
    }
   ],
   "source": [
    "i = 0\n",
    "while (i < 5):\n",
    "    print(i, end=\"\")\n",
    "    i += 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "oqGq0fY5jYRx"
   },
   "source": [
    "Take note of the lack of braces used as compared to other programming languages. Instead, indentations are used to indicate the scope of a block of code.\n",
    "\n",
    "*Side note*: print in python automatically adds a space between arguments, and a newline character at the end. Adding the keywords sep = \"...\" replaces the space with 3 periods, and end = \"\" replaces the newline character with an empty string."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9SDs0Jjrjl2r"
   },
   "source": [
    "*For Loops*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "id": "RmjdPs0FjQ9p"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "01234"
     ]
    }
   ],
   "source": [
    "for i in range(5):\n",
    "    print(i, end=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-trDyt8TjpN8"
   },
   "source": [
    "Instead of the conventential (initialisation; test; update) used in other programming languages, python uses a library function `range` to iterate through elements. By default, using range with one parameter will iterate through values starting from 0 up to and not including the value in the parameter. For more freedom, the range function can be called with extra parameters as well:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "id": "yTJmES_ujlYk"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "234"
     ]
    }
   ],
   "source": [
    "for i in range (2, 5):\n",
    "    print(i, end=\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "id": "Jdow039Cj1wB"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3579"
     ]
    }
   ],
   "source": [
    "for i in range(3, 10, 2):\n",
    "    print(i, end=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "limrwJVoj7Sf"
   },
   "source": [
    "For more information about `range`, you may want to refer to: https://docs.python.org/3/library/stdtypes.html#range"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1t4qAfZxkR0L"
   },
   "source": [
    "### 1.0.4 Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "g6ZZIThGkTwr"
   },
   "source": [
    "A function in Python with control flow looks something like this"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "id": "hNnmQNnUj5HO"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Zero here!\n",
      "One here!\n",
      "2\n",
      "hello\n",
      "3\n",
      "hello\n"
     ]
    }
   ],
   "source": [
    "def foo(x):\n",
    "    for i in range(4):\n",
    "        if i == 0:\n",
    "            print('Zero here!')\n",
    "        elif i == 1:\n",
    "            print('One here!')\n",
    "        else:\n",
    "            print(i)\n",
    "            print(x)\n",
    "foo('hello')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "95ATQp3BknlI"
   },
   "source": [
    "Recall that indentation in Python is important as it is used to indicate the scope of a block of code. Incorrect level of indentation can cause syntax errors and unintended behaviours, even errors. For instance, if we are to indent the last print statement in the previous code snippet, we will get a different result. The next code snippet illustrates this. Make sure you understand why this is so."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "id": "YYze3lmKkjHQ"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Zero here!\n",
      "hello\n",
      "One here!\n",
      "hello\n",
      "2\n",
      "hello\n",
      "3\n",
      "hello\n"
     ]
    }
   ],
   "source": [
    "def foo(x):\n",
    "    for i in range(4):\n",
    "        if i == 0:\n",
    "            print('Zero here!')\n",
    "        elif i == 1:\n",
    "            print('One here!')\n",
    "        else:\n",
    "            print(i)\n",
    "        print(x) # Additional level of indentation\n",
    "foo('hello')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "r-2meP8gk0vo"
   },
   "source": [
    "For more information about functions, you may want to refer to https://docs.python.org/3.9/tutorial/controlflow.html#defining-functions.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0Y8RMpHbk2m-"
   },
   "source": [
    "### 1.0.5 Aliasing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OxLQG0SSk7Ea"
   },
   "source": [
    "When two variables refer to the same object in Python, aliasing occurs. We can check whether to variables are aliases of each other by using `is`.\n",
    "\n",
    "Note that if two variables a and b are aliases of each other, i.e. a `is` b returns `True`, a == b will return `True`. However, if a == b returns `True`, it does not mean that the variables are aliases of each other. For example,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "id": "ogMC9R_Fkx3Q"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "False\n",
      "True\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "a = [1, 2, 3]\n",
    "b = [1, 2, 3]\n",
    "\n",
    "print(a == b) # True\n",
    "print(a is b) # False\n",
    "\n",
    "c = a         # Now , c points to the same object as a\n",
    "print(a == c) # True\n",
    "print(a is c) # True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GQ7OvFRWlc7_"
   },
   "source": [
    "Think about what this means. When aliasing occurs, unintended side effects may surface. Consider the previous example. Suppose we want to modify c. We might do something like"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "id": "bfF3lsNqlcUI"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['hello', 2, 3]\n"
     ]
    }
   ],
   "source": [
    "c[0] = 'hello'\n",
    "print(a) # Prints ['hello', 2, 3]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Dl2X6vhql4-S"
   },
   "source": [
    "To avoid aliasing then, we can use copy that is provided by Python. Returning to the previous example, if we want c to contain elements that are identical to a, except for the zeroth one, we can do"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "id": "OTjT9qfMllTg"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1, 2, 3]\n",
      "['hello', 2, 3]\n"
     ]
    }
   ],
   "source": [
    "a = [1, 2, 3]\n",
    "c = a.copy()\n",
    "c[0] = 'hello'\n",
    "print(a) # Prints [1, 2, 3]\n",
    "print(c) # Prints ['hello', 2, 3]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0XOXfFxbmNWi"
   },
   "source": [
    "Note however that `copy` returns what is called a *shallow* copy, i.e. that the copy only done at the first level. To do a deep copy, we should recursively copy the nested lists. We can conveniently do this using use [copy.deepcopy]( https://docs.python.org/3.9/library/copy.html). In the case of a list of lists, simply doing copy may not work. Consider the following example to understand the difference between shallow copy and deep copy:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "id": "M6S6hMvAmHO4"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shallow copy\n",
      "[[5, 2], [3, 4]]\n",
      "[[5, 2], [3, 4]]\n",
      "Deep copy\n",
      "[[1, 2], [3, 4]]\n",
      "[[5, 2], [3, 4]]\n"
     ]
    }
   ],
   "source": [
    "import copy\n",
    "\n",
    "print('Shallow copy')\n",
    "a = [[1, 2], [3, 4]]\n",
    "b = a.copy() # Performs a shallow copy of variable a\n",
    "b[0][0] = 5  # Modifies both a and b\n",
    "print(a)     # Prints [[5, 2], [3, 4]]\n",
    "print(b)     # Prints [[5, 2], [3, 4]]\n",
    "\n",
    "print('Deep copy')\n",
    "x = [[1, 2], [3, 4]]\n",
    "y = copy.deepcopy(x) # Performs a deep copy of variable x\n",
    "y[0][0] = 5          # Modifies y only\n",
    "print(x)             # Prints [[1, 2], [3, 4]]\n",
    "print(y)             # Prints [[5, 2], [3, 4]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DQKm3UeGnVry"
   },
   "source": [
    "### 1.0.6 Swapping Variables"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VYs8aZyQnX6v"
   },
   "source": [
    "Instead of swapping two elements with a temporary variable, as we do in other programming languages, it is possible to do something like this in Python."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "id": "0tyZ-eG8nNtR"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2\n",
      "1\n"
     ]
    }
   ],
   "source": [
    "a = 1\n",
    "b = 2\n",
    "a, b = b, a\n",
    "print(a) # Prints 2\n",
    "print(b) # Prints 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Qb0XbrcHng_f"
   },
   "source": [
    "### 1.0.7 Lambda Functions\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2z1eTBrtnith"
   },
   "source": [
    "Lambda functions are anonymous functions. Here is an example of a lambda function in python:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "id": "JxxJHznJncDj"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2109\n",
      "2109\n"
     ]
    }
   ],
   "source": [
    "def increment_by_one(x):\n",
    "    return x + 1\n",
    "print(increment_by_one(2108))   # Prints 2109\n",
    "print((lambda x : x + 1)(2108)) # Prints 2109"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yIDa-K3JnzHv"
   },
   "source": [
    "### 1.0.8 Map\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YokpZBBLn1LG"
   },
   "source": [
    "Transforming a list of elements by applying a function on each element is a common operation. For example, we may want to increment each element in an array by 1. Perhaps, the first thing that comes to mind is to use a for-loop. However, in Python, we can use `map` instead. The following code snippet illustrates this.\n",
    "\n",
    "**Note**: As `map` is an iterator, we need to call `list` if we want the result to be returned as a list."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "id": "LZIIf_OFnvDL"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "a = [1, 2, 3, 4]\n",
    "b = []\n",
    "for i in range(len(a)):\n",
    "    b.append(a[i] + 1)\n",
    "\n",
    "c = list(map(lambda x : x + 1, a)) # Equivalent to for loop above\n",
    "print(b == c)                      # Prints True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Sy5Wp04EoWoO"
   },
   "source": [
    "In fact, we can pass in more than one list—or more generally, iterable—into `map`. The number of arguments that the function, which is passed into `map`, just has to match the number of lists—or more generally, iterables—passed into `map`. For instance,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "id": "CAGwgRjioR9c"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['helloworld', 'byecat']\n"
     ]
    }
   ],
   "source": [
    "a = [\"hello\", \"bye\"]\n",
    "b = [\"world\", \"cat\"]\n",
    "\n",
    "c = list(map(lambda x, y: x + y, a, b))\n",
    "print(c) # Prints ['helloworld', 'byecat']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UxjKFV9yozO7"
   },
   "source": [
    "### 1.0.9 Filter\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tfiGPKpwo1VG"
   },
   "source": [
    "Another useful function is `filter`. Like what its name suggests, we can use it to filter a list—or more generally, an iterable—of elements. An example of how it works is as shown."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "id": "Os-QeB-Fosks"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "a = [1, 2, 3, 4]\n",
    "filtered_1 = []\n",
    "for i in range(4):\n",
    "    if a[i] % 2 == 0:\n",
    "        filtered_1.append(a[i])\n",
    "\n",
    "filtered_2 = list(filter(lambda x : x % 2 == 0, a)) # Equivalent as for loop above\n",
    "print(filtered_1 == filtered_2)                     # Prints True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.0.10 Dictionary"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One of the data structures which are commonly used is hash tables or hash maps. It is a data structure that implements associative arrays, i.e., a mapping from a certain key to a value.\n",
    "\n",
    "In Python, this data structure is called _dictionary_. Examples of how it works are given below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{}\n",
      "{'foo': 'bar', 'one': 1}\n",
      "bar\n",
      "{'foo': 'cs2109s', 'two': 2}\n",
      "True\n",
      "False\n",
      "{'foo': 'cs2109s', 'two': 2, 1: 'one', 2: 'two'}\n"
     ]
    }
   ],
   "source": [
    "# Initialize an empty dictionary\n",
    "a = {}\n",
    "# or\n",
    "a = dict()\n",
    "\n",
    "print(a) # Prints {}\n",
    "\n",
    "# We can also initialize a dictionary with some items\n",
    "a = {'foo': 'bar', 'one': 1}\n",
    "\n",
    "print(a) # Prints a as initialized\n",
    "\n",
    "print(a['foo']) # Prints \"bar\"\n",
    "\n",
    "a['two'] = 2         # Adds a new key \"two\" with value 2\n",
    "a['foo'] = 'cs2109s' # Updates the value of key \"foo\" with \"cs2109s\"\n",
    "del a['one']         # Deletes the key \"one\"\n",
    "\n",
    "print(a) # Prints the updated a\n",
    "\n",
    "print('two' in a)   # True\n",
    "print('three' in a) # False\n",
    "\n",
    "b = {1: 'one', 2: 'two'} # Creates a new dictionary\n",
    "\n",
    "c = {**a, **b} # Merges dictionaries a and b\n",
    "print(c)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Hashable Objects\n",
    "\n",
    "Not all objects can be used as a key. Only objects whose hash values can be computed are useable for keys. We call these objects hashable objects. Some commonly used keys are strings, numbers, and tuples of strings and numbers. In contrast, a list is not hashable and therefore can not be used as keys. In general, immutable objects are usually hashable and mutable objects are usually unhashable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'test': 1}\n",
      "{'test': 1, 0: 1}\n",
      "{'test': 1, 0: 1, ('test', 0): 1}\n",
      "unhashable type: 'list'\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    a = {}\n",
    "    a['test'] = 1      # String: OK\n",
    "    print(a)\n",
    "    \n",
    "    a[0] = 1           # Number: OK\n",
    "    print(a)\n",
    "    \n",
    "    a[('test', 0)] = 1 # Tuple of string and number: OK\n",
    "    print(a)\n",
    "    \n",
    "    a[['test', 0]] = 1 # List: FAIL\n",
    "    print(a)\n",
    "except Exception as e:\n",
    "    print(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.0.11 Set"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A set is a collection of unique data. That is, elements of a set cannot be duplicate. Similar to a dictionary keys, the elements of a set must be hashable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "set()\n",
      "{0, 1, 2}\n",
      "{0, 1, 2, 3}\n",
      "{1, 2, 3}\n",
      "True\n",
      "False\n",
      "{3, 4, 5}\n",
      "{3}\n",
      "{3}\n",
      "{1, 2, 3, 4, 5}\n",
      "{1, 2, 3, 4, 5}\n",
      "{1, 2, 4, 5}\n",
      "{1, 2, 4, 5}\n",
      "{1, 2}\n",
      "{4, 5}\n",
      "True\n",
      "False\n"
     ]
    }
   ],
   "source": [
    "# Initialize an empty set\n",
    "a = set()\n",
    "\n",
    "print(a) # Prints set()\n",
    "\n",
    "# Initialize a set with some elements\n",
    "a = set([0, 1, 2])\n",
    "\n",
    "print(a)    # Prints {0, 1, 2}\n",
    "\n",
    "a.add(3)    # Adds 3\n",
    "\n",
    "print(a)    # Prints {0, 1, 2, 3}\n",
    "\n",
    "a.remove(0) # Remove 0\n",
    "\n",
    "print(a)    # Prints {1, 2, 3}\n",
    "\n",
    "print(1 in a) # True\n",
    "print(4 in a) # False\n",
    "\n",
    "a2 = set([3, 4, 5])\n",
    "\n",
    "print(a2)\n",
    "\n",
    "print(a.intersection(a2))          # Intersection of two sets\n",
    "print(a & a2)                      # Intersection of two sets\n",
    "print(a.union(a2))                 # Union of two sets\n",
    "print(a | a2)                      # Union of two sets\n",
    "print(a.symmetric_difference(a2))  # Symmetric difference of two sets\n",
    "print(a ^ a2)                      # Symmetric difference of two sets\n",
    "print(a - a2)                      # Difference of two sets\n",
    "print(a2 - a)                      # Difference of two sets\n",
    "\n",
    "print(a == set([3, 2, 1]))    # True\n",
    "print(a == set([3, 2, 1, 0])) # False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nLSj1KTtpYAr"
   },
   "source": [
    "## 1.1 Matrix Operators in Python"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "US7_thWVpfcN"
   },
   "source": [
    "Now, let us do some programming in Python!\n",
    "\n",
    "**IMPORTANT**: For the tasks in part 1 of this problem, you need to ensure that none of the matrix operations which you implement modify the input matrix (or matrices). In addition, you are not allowed to import any packages that have not already been imported (because the whole point is for you to learn Python by implementing matrix operations)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SkU-Uy7Tpfkh"
   },
   "source": [
    "### Matrix\n",
    "\n",
    "Recall that each matrix has a dimension $n \\times m$, where $n$ is the number of rows and $m$ is the number of columns in the matrix. For instance, the following matrix $X$ has a dimension of 2 × 3.\n",
    "  \n",
    "$$\n",
    "X = \n",
    "\\begin{bmatrix}\n",
    "5 & 7 & 9 \\\\\n",
    "1 & 4 & 3\n",
    "\\end{bmatrix}\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_ed3mXFkqEVh"
   },
   "source": [
    "In addition, we can refer to each entry in the matrix by indexing it using the row and column where it is located. For example, the entry with value ’3’ in $X$ is the (1, 2) entry while the entry with value ’7’ is the (0, 1) entry. More generally, the entry in the $i$-th row and $j$-th column in a matrix $A$ is the ($i$, $j$) entry, and its value is denoted by $A_{i,j}$.\n",
    "\n",
    "**Note**: to be consistent with indexing in Python, we will be using zero-based index.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_eOX5GtkqnBV"
   },
   "source": [
    "Now, this begs the question of how a matrix in Python should be represented. For\n",
    "simplicity, we shall do it with a list of lists. We can define matrix *X* in two ways:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "id": "3XoRJChhpJMa"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Method 1\n",
      "[[5, 7, 9], [1, 4, 3]]\n",
      "Method 2\n",
      "[[5, 7, 9], [1, 4, 3]]\n"
     ]
    }
   ],
   "source": [
    "print('Method 1')\n",
    "X = [None] * 2\n",
    "X[0] = [5, 7, 9]\n",
    "X[1] = [1, 4, 3]\n",
    "print(X)\n",
    "\n",
    "print('Method 2')\n",
    "X = [[5, 7, 9],\n",
    "     [1, 4, 3]]\n",
    "print(X)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9l1gJmz6q8EA"
   },
   "source": [
    "In the following tasks, you **may assume** that the input matrices are\n",
    "* list of lists\n",
    "* the number of elements in each of the ’inner’ lists can be assumed to be the same (because the number of columns in a matrix is the same for each row)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8ODILekjrCeH"
   },
   "source": [
    "### Task 1.1 Scalar Multiplication"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WJFIy6IBrG5C"
   },
   "source": [
    "Our first task is to implement `mult_scalar`. This function takes in two arguments, a matrix $A$ and a scalar $c$, and returns a new matrix $A$ obtained by multiplying each element by $c$.\n",
    "\n",
    "Scalar multiplication of an $n \\times m$ matrix $A$ by $c$ is done by multiplying $c$ to each *A<sub>i,j</sub>*, i.e.\n",
    "\n",
    "\\begin{equation}\n",
    "    cA =\n",
    "    \\begin{bmatrix}\n",
    "    cA_{0, 0} & cA_{0, 1} & ... & cA_{0, m-1}\\\\\n",
    "    cA_{1, 0} & cA_{1, 1} & ... & cA_{1, m-1}\\\\\n",
    "    \\vdots & \\vdots & \\ddots & \\vdots\\\\\n",
    "    cA_{n-1, 0} & cA_{n-1, 1} & ... & cA_{n-1, m-1}\\\\\n",
    "    \\end{bmatrix}\n",
    "\\end{equation}\n",
    "\n",
    "For example if $c = 2$, then\n",
    "\\begin{equation*}\n",
    "    2X = \n",
    "    \\begin{bmatrix}\n",
    "    10 & 14 & 18\\\\\n",
    "    2 & 8 & 6\n",
    "    \\end{bmatrix}\n",
    "\\end{equation*}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*IMPORTANT NOTE:* While you can write and run your code directly in this Jupyter notebook, it is difficult to debug your code directly here. We would recommend that you code using an IDE with the `.py` file provided. **Do not use NumPy for this question**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "id": "Aa_BHke5q5sR"
   },
   "outputs": [],
   "source": [
    "def mult_scalar(A, c):\n",
    "    \"\"\"\n",
    "    Returns a new matrix created by multiplying elements of matrix A by a scalar c.\n",
    "    Note\n",
    "    ----\n",
    "    Do not use numpy for this question.\n",
    "    \"\"\"\n",
    "    # TODO: add your solution here and remove `raise NotImplementedError`\n",
    "    return [[i * c for i in row] for row in A]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[10, 14, 18], [2, 8, 6]]\n"
     ]
    }
   ],
   "source": [
    "# Test cases for Task 1.1\n",
    "A = [[5, 7, 9], [1, 4, 3]]\n",
    "A_copy = copy.deepcopy(A)\n",
    "\n",
    "actual = mult_scalar(A_copy, 2)\n",
    "expected = [[10, 14, 18], [2, 8, 6]]\n",
    "assert(A == A_copy) # check for aliasing\n",
    "print(actual)\n",
    "assert(actual == expected)\n",
    "\n",
    "\n",
    "A2 = [[6, 5, 5], [8, 6, 0], [1, 5, 8]]\n",
    "A2_copy = copy.deepcopy(A2)\n",
    "\n",
    "actual2 = mult_scalar(A2_copy, 5)\n",
    "expected2 = [[30, 25, 25], [40, 30, 0], [5, 25, 40]]\n",
    "assert(A2 == A2_copy) # check for aliasing\n",
    "assert(actual2 == expected2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FL7S2pHhrfWK"
   },
   "source": [
    "### Task 1.2 Matrix Addition\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ncvdVPbvrike"
   },
   "source": [
    "Our next task is to implement `add_matrices`. This function takes in two arguments, a matrix $A$ and a matrix $B$, and returns a new matrix that is the result of adding $B$ to $A$.\n",
    "\n",
    "For this operation to be valid, $A$ and $B$ must have the same dimensions. If this is the case, then\n",
    "\\begin{equation}\n",
    "    A + B = \n",
    "    \\begin{bmatrix}\n",
    "    A_{0, 0} + B_{0, 0} & A_{0, 1} + B_{0, 1} & ... & A_{0, m-1} + B_{0, m-1}\\\\\n",
    "    A_{1, 0} + B_{1, 0} & A_{1, 1} + B_{1, 1} & ... & A_{1, m-1} + B_{1, m-1}\\\\\n",
    "    \\vdots & \\vdots & \\ddots & \\vdots\\\\\n",
    "    A_{n-1, 0} + B_{n-1, 0} & A_{n-1, 1} + B_{n-1, 1} & ... & A_{n-1, m-1} + B_{n-1, m-1}\\\\\n",
    "    \\end{bmatrix}\n",
    "\\end{equation}\n",
    "\n",
    "To exemplify this, consider $X$ and $Y$ below.\n",
    "\n",
    "\\begin{equation}\n",
    "    X = \n",
    "    \\begin{bmatrix}\n",
    "    5 & 7 & 9\\\\\n",
    "    1 & 4 & 3\n",
    "    \\end{bmatrix}\n",
    "\\end{equation}\n",
    "\n",
    "\\begin{equation}\n",
    "    Y = \n",
    "    \\begin{bmatrix}\n",
    "    2 & 3 & 4\\\\\n",
    "    5 & 6 & 7\n",
    "    \\end{bmatrix}\n",
    "\\end{equation}\n",
    "\n",
    "Then, we have\n",
    "\\begin{equation*}\n",
    "    Y + X = \n",
    "    \\begin{bmatrix}\n",
    "    2 + 5 & 3 + 7 & 4 + 9 \\\\\n",
    "    5 + 1 & 6 + 4 & 7 + 3\n",
    "    \\end{bmatrix}\n",
    "\\end{equation*}\n",
    "\n",
    "Observe that because addition of scalar values is commutative (i.e. x + y = y + x), we always have $A + B = B + A$. However, note that, as we shall see later, not all matrix operations are commutative.\n",
    "\n",
    "*IMPORTANT NOTE:* **Do not use NumPy for this question**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "id": "gVdOsgvhrWcT"
   },
   "outputs": [],
   "source": [
    "def add_matrices(A, B):\n",
    "    \"\"\"\n",
    "    Returns a new matrix that is the result of adding matrix B to matrix A.\n",
    "    Note\n",
    "    ----\n",
    "    Do not use numpy for this question.\n",
    "    \"\"\"\n",
    "    if len(A) != len(B) or len(A[0]) != len(B[0]):\n",
    "        raise Exception('A and B cannot be added as they have incompatible dimensions!')\n",
    "    result = [[0] * len(A[0]) for _ in A]\n",
    "    for i in range(len(A)):\n",
    "        for j in range(len(A[0])):\n",
    "            result[i][j] = A[i][j] + B[i][j]\n",
    "    return result\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[7, 10, 13], [6, 10, 10]]\n"
     ]
    }
   ],
   "source": [
    "# Test case for Task 1.2\n",
    "A = [[5, 7, 9], [1, 4, 3]]\n",
    "B = [[2, 3, 4], [5, 6, 7]]\n",
    "A_copy = copy.deepcopy(A)\n",
    "B_copy = copy.deepcopy(B)\n",
    "\n",
    "actual = add_matrices(A_copy, B_copy)\n",
    "print(actual)\n",
    "expected = [[7, 10, 13], [6, 10, 10]]\n",
    "assert(A == A_copy) # check for aliasing\n",
    "assert(B == B_copy) # check for aliasing\n",
    "assert(actual == expected)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sYRR2xWkr0pF"
   },
   "source": [
    "### Task 1.3 Transpose a Matrix"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vzpIHNX7r7Dh"
   },
   "source": [
    "Our third task is to implement `transpose_matrix`. This function takes in one argument, a matrix $A$, and returns a new matrix that is the transpose of $A$. \n",
    "\n",
    "Less formally, the transpose of a matrix $A$ can be found by changing the rows in the matrix to columns, and this transpose is denoted by $A^T$. For example, \n",
    "for matrix X\n",
    "\n",
    "\\begin{equation}\n",
    "    X = \n",
    "    \\begin{bmatrix}\n",
    "    5 & 7 & 9\\\\\n",
    "    1 & 4 & 3\n",
    "    \\end{bmatrix}\n",
    "\\end{equation}\n",
    "\n",
    "\n",
    "then the transpose of matrix $X$ is given by\n",
    "\n",
    "\\begin{equation*}\n",
    "    X^T = \n",
    "    \\begin{bmatrix}\n",
    "    5 & 1\\\\\n",
    "    7 & 4\\\\\n",
    "    9 & 3\n",
    "    \\end{bmatrix}\n",
    "\\end{equation*}\n",
    "\n",
    "More formally, suppose $\\mathbf{r_i}$ represents the i-th row in an $n \\times m$ matrix $A$, i.e.\n",
    "\\begin{equation}\n",
    "    A = \n",
    "    \\begin{bmatrix}\n",
    "    \\mathbf{r_0}\\\\\n",
    "    \\mathbf{r_1}\\\\\n",
    "    \\vdots\\\\\n",
    "    \\mathbf{r_{n-2}}\\\\\n",
    "    \\mathbf{r_{n-1}}\n",
    "    \\end{bmatrix}\n",
    "\\end{equation}\n",
    "\n",
    "Then, $A^T$ is given by\n",
    "\\begin{equation}\n",
    "    A^T = \n",
    "    \\begin{bmatrix}\n",
    "    \\mathbf{r_0} & \\mathbf{r_1} & ... & \\mathbf{r_{n-2}} & \\mathbf{r_{n-1}}\n",
    "    \\end{bmatrix}\n",
    "\\end{equation}\n",
    "\n",
    "Note that this means that $A^T$ has dimension $m \\times n$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our third task is to implement `transpose_matrix`. This function takes in one argument, a matrix $A$, and transposes $A$. \n",
    "\n",
    "Less formally, the transpose of a matrix $A$ can be found by changing the rows in the matrix to columns, and this transpose is denoted by $A^T$. For example, \n",
    "for matrix X\n",
    "\n",
    "\\begin{equation}\n",
    "    X = \n",
    "    \\begin{bmatrix}\n",
    "    5 & 7 & 9\\\\\n",
    "    1 & 4 & 3\n",
    "    \\end{bmatrix}\n",
    "\\end{equation}\n",
    "\n",
    "\n",
    "then the transpose of matrix $X$ is given by\n",
    "\n",
    "\\begin{equation*}\n",
    "    X^T = \n",
    "    \\begin{bmatrix}\n",
    "    5 & 1\\\\\n",
    "    7 & 4\\\\\n",
    "    9 & 3\n",
    "    \\end{bmatrix}\n",
    "\\end{equation*}\n",
    "\n",
    "More formally, suppose $\\mathbf{r_i}$ represents the i-th row in an $n \\times m$ matrix $A$, i.e.\n",
    "\\begin{equation}\n",
    "    A = \n",
    "    \\begin{bmatrix}\n",
    "    \\mathbf{r_0}\\\\\n",
    "    \\mathbf{r_1}\\\\\n",
    "    \\vdots\\\\\n",
    "    \\mathbf{r_{n-2}}\\\\\n",
    "    \\mathbf{r_{n-1}}\n",
    "    \\end{bmatrix}\n",
    "\\end{equation}\n",
    "\n",
    "Then, $A^T$ is given by\n",
    "\\begin{equation}\n",
    "    A^T = \n",
    "    \\begin{bmatrix}\n",
    "    \\mathbf{r_0} & \\mathbf{r_1} & ... & \\mathbf{r_{n-2}} & \\mathbf{r_{n-1}}\n",
    "    \\end{bmatrix}\n",
    "\\end{equation}\n",
    "\n",
    "Note that this means that $A^T$ has dimension $m \\times n$.\n",
    "\n",
    "*IMPORTANT NOTE:* **Do not use NumPy for this question**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "id": "K1OahnevrtIV"
   },
   "outputs": [],
   "source": [
    "def transpose_matrix(A):\n",
    "    \"\"\"\n",
    "    Returns a new matrix that is the transpose of matrix A.\n",
    "    Note\n",
    "    ----\n",
    "    Do not use numpy for this question.\n",
    "    \"\"\"\n",
    "    # return list([list(a) for a in zip(*A)])\n",
    "    rows = len(A)\n",
    "    cols = len(A[0])\n",
    "    result = [[0] * rows for _ in range(cols)]\n",
    "    for i in range(cols):\n",
    "        for j in range(rows):\n",
    "            result[i][j] = A[j][i]\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[5, 1], [7, 4], [9, 3]]\n"
     ]
    }
   ],
   "source": [
    "# Test case for Task 1.3\n",
    "A = [[5, 7, 9], [1, 4, 3]]\n",
    "A_copy = copy.deepcopy(A)\n",
    "\n",
    "actual = transpose_matrix(A_copy)\n",
    "print(actual)\n",
    "expected = [[5, 1], [7, 4], [9, 3]]\n",
    "assert(A == A_copy)\n",
    "assert(actual == expected)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "umaW01qxz3RF"
   },
   "source": [
    "### Vector\n",
    "\n",
    "A vector is an $n \\times 1$ matrix. One mathematical operation that is defined for vectors of the same dimension is the \\emph{dot product}. Suppose we have\n",
    "$\\mathbf{p} = \n",
    "    \\begin{bmatrix}\n",
    "    3 & 2 & 3\\\\\n",
    "    \\end{bmatrix}^{T}\n",
    "$ and\n",
    "$\\mathbf{q} = \n",
    "    \\begin{bmatrix}\n",
    "    4 & 5 & 6\\\\\n",
    "    \\end{bmatrix}^{T}\n",
    "$. Then, the dot product of $\\mathbf{p}$ and $\\mathbf{q}$ is denoted by $\\mathbf{p} \\cdot \\mathbf{q}$, and is given by\n",
    "\n",
    "\\begin{equation*}\n",
    "\\mathbf{p} \\cdot \\mathbf{q} = 3 \\times 4 + 2 \\times 5 + 3 \\times 6 = 40\n",
    "\\end{equation*}\n",
    "\n",
    "More generally, for two vectors $\\mathbf{u}$ and $\\mathbf{v}$, where both have a dimension of $n \\times 1$,\n",
    "\\begin{equation}\n",
    "    \\mathbf{u} \\cdot \\mathbf{v} = \\sum_{i = 0}^{n - 1} u_{i, 1} v_{i, 1}\n",
    "\\end{equation}\n",
    "\n",
    "Observe that this operation gives us a scalar value, i.e. $\\mathbf{u} \\cdot \\mathbf{v} \\in \\mathbb{R}$, not a matrix. \n",
    "\n",
    "We shall see how dot product relates to \\emph{matrix multiplication}, which we are going to implement in the next task."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "y9xDe72euNh3"
   },
   "source": [
    "### Task 1.4 Multiply Two Matrices"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JqWvt6u6uTuq"
   },
   "source": [
    "Our fourth task is to implement `mult_matrix`. This function takes in two arguments, a matrix $A$ and a matrix $B$, and returns a new matrix that is the result of multiplying them.\n",
    "\n",
    "For this operation to be valid, the number of columns $m$ in $A$ must be equal to the number of rows $n'$ in $B$, i.e. $m = n'$. If this is the case, then the value of each entry $(i, j)$ in $A \\times B$, or $AB$, is given by the dot product of the transpose of the $i$-th row in $A$ and the $j$-th column in $B$.\n",
    "\n",
    "For example, to compute $XY^T$, where $X$ and $Y$ are as defined in Task 1.2, respectively,\n",
    "\n",
    "<img src = 'imgs/matrix_vector_multiplication.png'>\n",
    "\n",
    "More formally, suppose that $A$ has dimension $n \\times m$ and $B$ has dimension $n' \\times m'$, where $m = n'$. Moreover, the $i$-th row in $A$ is denoted by $\\mathbf{r_i}$ and the $j$-th column in $B$ is denoted by $\\mathbf{c_j}$ such that\n",
    "\n",
    "\\begin{equation*}\n",
    "    B = \n",
    "    \\begin{bmatrix}\n",
    "    \\mathbf{c_0} & \\mathbf{c_1} & ... & \\mathbf{c_{m'-2}} & \\mathbf{c_{m'-1}}\n",
    "    \\end{bmatrix}\n",
    "\\end{equation*}\n",
    "\n",
    "Then,\n",
    "\\begin{equation}\n",
    "    AB =\n",
    "    \\begin{bmatrix}\n",
    "    \\mathbf{r_0}^T \\cdot \\mathbf{c_0} & \\mathbf{r_0}^T \\cdot \\mathbf{c_1} & ... & \\mathbf{r_0}^T \\cdot \\mathbf{c_{m'-1}}\\\\\n",
    "    \\mathbf{r_1}^T \\cdot \\mathbf{c_0} & \\mathbf{r_1}^T \\cdot \\mathbf{c_1} & ... & \\mathbf{r_1}^T \\cdot \\mathbf{c_{m'-1}}\\\\\n",
    "    \\vdots & \\vdots & \\ddots & \\vdots\\\\\n",
    "    \\mathbf{r_{n-1}}^T \\cdot \\mathbf{c_0} & \\mathbf{r_{n-1}}^T \\cdot \\mathbf{c_1} & ... & \\mathbf{r_{n-1}}^T \\cdot \\mathbf{c_{m'-1}}\\\\\n",
    "    \\end{bmatrix}\n",
    "\\end{equation}\n",
    "\n",
    "Food for thought: what is the dimension of $AB$? As a follow-up question, what does this tell us about the commutativity (or the lack of it) of matrix multiplication?\n",
    "\n",
    "*IMPORTANT NOTE:* **Do not use NumPy for this question**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "id": "3wUd3W6SrxiB"
   },
   "outputs": [],
   "source": [
    "def dot_prod(A, B):\n",
    "    if len(A) != len(B):\n",
    "        raise Exception('A and B cannot be multiplied as they have incompatible dimensions!')\n",
    "    return sum([A[i] * B[i] for i in range(len(A))])\n",
    "\n",
    "def mult_matrices(A, B):\n",
    "    \"\"\"\n",
    "    Multiplies matrix A by matrix B, giving AB.\n",
    "    Note\n",
    "    ----\n",
    "    Do not use numpy for this question.\n",
    "    \"\"\"\n",
    "    if len(A[0]) != len(B):\n",
    "        raise Exception('Incompatible dimensions for matrix multiplication of A and B')\n",
    "    res_rows = len(A)\n",
    "    res_cols = len(B[0])\n",
    "    result = [[0] * res_cols for _ in range(res_rows)]\n",
    "    trans_B = transpose_matrix(B)\n",
    "    for i in range(res_rows):\n",
    "        for j in range(res_cols):\n",
    "            result[i][j] = dot_prod(A[i], trans_B[j])\n",
    "    return result\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test Cases for Task 1.4\n",
    "A = [[5, 7, 9], [1, 4, 3]]\n",
    "B = [[2, 5], [3, 6], [4, 7]]\n",
    "A_copy = copy.deepcopy(A)\n",
    "B_copy = copy.deepcopy(B)\n",
    "\n",
    "actual = mult_matrices(A, B)\n",
    "expected = [[67, 130], [26, 50]]\n",
    "assert(A == A_copy and B == B_copy)\n",
    "assert(actual == expected)\n",
    "\n",
    "A2 = [[-13, -10], [-24, 14]]\n",
    "B2 = [[1, 0], [0, 1]]\n",
    "A2_copy = copy.deepcopy(A2)\n",
    "B2_copy = copy.deepcopy(B2)\n",
    "\n",
    "actual2 = mult_matrices(A2, B2)\n",
    "expected2 = [[-13, -10], [-24, 14]]\n",
    "assert(A2 == A2_copy and B2 == B2_copy)\n",
    "assert(actual2 == expected2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uMfIOl_TvSuT"
   },
   "source": [
    "### Task 1.5 Inverse of a Matrix"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0xaIDHhz1-x9"
   },
   "source": [
    "This is the last task for this part of the homework. In this task, we are supposed to implement `invert_matrix`. This function takes in one argument, a matrix $A$, and tries to find the inverse of $A$. Yes, we said 'tries'.\n",
    "\n",
    "For the inverse of an $n \\times m$ matrix $A$ to exist, $A$ has to be a **square matrix**, i.e. $n = m$. Moreover, $A$ has to satisfy some [properties](https://en.wikipedia.org/wiki/Invertible_matrix) if you are interested in the technicalities of it. However, these are not required to complete this task. Therefore, it is not always possible to find the inverse of a matrix.\n",
    "\n",
    "Suppose the inverse of $A$ exists, then it is denoted by $A^{-1}$ and is such that\n",
    "\\begin{equation}\n",
    "    AA^{-1} = I = A^{-1}A\n",
    "\\end{equation}\n",
    "where $I$ is the identity matrix. In other words, the dimension of $I$ depends on the context, e.g. if an invertible matrix $Q$ has a dimension of $2 \\times 2$, then $\n",
    "QQ^{-1} = I =\n",
    "\\begin{bmatrix}\n",
    "1 & 0\\\\\n",
    "0 & 1\n",
    "\\end{bmatrix}\n",
    "$ \n",
    "with dimension $n \\times n$, diagonal entries that are all 1s, and off-diagonal entries that are all 0s, i.e.\n",
    "\\begin{equation}\n",
    "    I_{i, j} =\n",
    "    \\begin{cases}\n",
    "        1 & \\text{if $i = j$}\\\\\n",
    "        0 & \\text{otherwise}\n",
    "    \\end{cases}       \n",
    "\\end{equation}\n",
    "\n",
    "For example, if we have\n",
    "\\begin{equation*}\n",
    "    P = \n",
    "    \\begin{bmatrix}\n",
    "    1 & 0 & 0\\\\\n",
    "    0 & 1 & 0\\\\\n",
    "    0 & -4 & 1\n",
    "    \\end{bmatrix}\n",
    "\\end{equation*}\n",
    "then\n",
    "\\begin{equation*}\n",
    "    P^{-1} = \n",
    "    \\begin{bmatrix}\n",
    "    1 & 0 & 0\\\\\n",
    "    0 & 1 & 0\\\\\n",
    "    0 & 4 & 1\n",
    "    \\end{bmatrix}\n",
    "\\end{equation*}\n",
    "\n",
    "We can check that this is true by computing $PP^{-1}$ and $P^{-1}P$. If the results evaluate to \n",
    "$\n",
    "\\begin{bmatrix}\n",
    "    1 & 0 & 0\\\\\n",
    "    0 & 1 & 0\\\\\n",
    "    0 & 0 & 1\n",
    "\\end{bmatrix}\n",
    "$, then $P^{-1}$ is correctly computed.\n",
    "\n",
    "Quick question: what is the inverse of $P^{-1}$? In fact, more generally, given any *invertible matrix* $X$, i.e. $X^{-1}$ exists, what is the inverse of $X^{-1}$?\n",
    "\n",
    "The description of **our algorithm for finding the inverse of a matrix is as follows**.\n",
    "\n",
    "Given any matrix $A$, in order to attempt finding $A^{-1}$, we need to check whether $A$ is a square matrix. If it isn't, we should return False to indicate that it is not possible; otherwise, we will construct an adjacency matrix as shown below.\n",
    "\n",
    "\\begin{equation}\n",
    "    \\begin{bmatrix}\n",
    "    A & I\n",
    "    \\end{bmatrix}\n",
    "\\end{equation}\n",
    "\n",
    "Note that here, the adjacency matrix is of dimension $n \\times 2n$. Since both $A$ and $I$ have a dimension of $n \\times n$, placing their entries side-by-side will give $2n$ columns.\n",
    "\n",
    "For instance, suppose we have\n",
    "\n",
    "\\begin{equation*}\n",
    "    Z =\n",
    "    \\begin{bmatrix}\n",
    "    0 & 3 & 2\\\\\n",
    "    0 & 0 & 1\\\\\n",
    "    1 & 5 & 3\n",
    "    \\end{bmatrix}\n",
    "\\end{equation*}\n",
    "\n",
    "Then, the adjacency matrix is given by\n",
    "\\begin{equation*}\n",
    "    \\begin{bmatrix}\n",
    "    0 & 3 & 2 & 1 & 0 & 0\\\\\n",
    "    0 & 0 & 1 & 0 & 1 & 0\\\\\n",
    "    1 & 5 & 3 & 0 & 0 & 1\n",
    "    \\end{bmatrix}\n",
    "\\end{equation*}\n",
    "\n",
    "Next, for each $i$ th row in the augmented matrix, do the following\n",
    "\n",
    "1. Find a row $k$, where $i \\leq k \\leq (n - 1)$, that has a non-zero-valued entry in its $i$ th column, and swap it with the $i$ th row. If no such row exists, return False. The matrix is not invertible.\n",
    "2. Multiply the new $i$ th row by a scalar such that the value in the $i$ th column becomes 1.\n",
    "3. Add multiples of the new $i$ th row to all other rows such that the value in their $i$ th column becomes 0.\n",
    "\n",
    "Finally, get the last $n$ columns that is found in the resulting adjacency matrix. These columns form $A^{-1}$, i.e. if $\\mathbf{c_i}$ denotes the $i$ th column in the adjacency matrix, then\n",
    "\\begin{equation}\n",
    "    A^{-1} =\n",
    "    \\begin{bmatrix}\n",
    "    \\mathbf{c_n} & \\mathbf{c_{n + 1}} & ... & \\mathbf{c_{2n - 1}}\n",
    "    \\end{bmatrix}\n",
    "\\end{equation}\n",
    "\n",
    "Let us run through this algorithm to find $Z^{-1}$.\n",
    "1. We have the adjacency matrix \n",
    "$\n",
    "    \\begin{bmatrix}\n",
    "    0 & 3 & 2 & 1 & 0 & 0\\\\\n",
    "    0 & 0 & 1 & 0 & 1 & 0\\\\\n",
    "    1 & 5 & 3 & 0 & 0 & 1\n",
    "    \\end{bmatrix}\n",
    "$\n",
    "    We shall swap the 2nd row with the 0th row, giving\n",
    "    $\n",
    "    \\begin{bmatrix}\n",
    "    1 & 5 & 3 & 0 & 0 & 1\\\\\n",
    "    0 & 0 & 1 & 0 & 1 & 0\\\\\n",
    "    0 & 3 & 2 & 1 & 0 & 0\\\\\n",
    "    \\end{bmatrix}\n",
    "    $\n",
    "2. Since the value of $(0, 0)$ entry is 1, we don't need to multiply the 0th row by any scalar. Moreover, we do not need to add multiples of this row to the other rows because the value of the entry in the 0th column of the other rows are all 0s.\n",
    "3. Moving on to $i = 1$, we see that we need to swap the last row with the $1$st row, giving \n",
    "$\n",
    "    \\begin{bmatrix}\n",
    "    1 & 5 & 3 & 0 & 0 & 1\\\\\n",
    "    0 & 3 & 2 & 1 & 0 & 0\\\\\n",
    "    0 & 0 & 1 & 0 & 1 & 0\\\\\n",
    "    \\end{bmatrix}\n",
    "$\n",
    "4. Since the value of the $(1, 1)$ entry is 3, we need to divide the 1st row by $3$, giving\n",
    "$\n",
    "    \\begin{bmatrix}\n",
    "    1 & 5 & 3 & 0 & 0 & 1\\\\\n",
    "    0 & 1 & \\frac{2}{3} & \\frac{1}{3} & 0 & 0\\\\\n",
    "    0 & 0 & 1 & 0 & 1 & 0\\\\\n",
    "    \\end{bmatrix}\n",
    "$\n",
    "5. Given that the value of the $(0, 1)$ entry is 5, not 0, we need to add $-5 \\times \\mathbf{a_1}$, where $\\mathbf{a_1}$ represents the 1st row in the current adjacency matrix, to the 0th row. This leaves us with\n",
    "$\n",
    "    \\begin{bmatrix}\n",
    "    1 & 0 & -\\frac{1}{3} & -\\frac{5}{3} & 0 & 1\\\\\n",
    "    0 & 1 & \\frac{2}{3} & \\frac{1}{3} & 0 & 0\\\\\n",
    "    0 & 0 & 1 & 0 & 1 & 0\\\\\n",
    "    \\end{bmatrix}\n",
    "$\n",
    "6. Now, $i = 2$. Observe that there are no other rows that can be swapped with the last row, since in this case substituting $i = 2$, we have $2 \\leq k \\leq 2 \\implies k = 2$. Moreover, we don't need to multiply the 2nd row by a scalar since the value in its last column is 1.\n",
    "7. To get 0s in the last column of the other rows, we need to add $\\frac{1}{3} \\times \\mathbf{a_2}$ and $-\\frac{2}{3} \\times \\mathbf{a_2}$, where $\\mathbf{a_2}$ represents the last row in the adjacency matrix shown in step 6, to the 0th and 1st row, respectively. The result is \n",
    "$\n",
    "    \\begin{bmatrix}\n",
    "    1 & 0 & 0 & -\\frac{5}{3} & \\frac{1}{3} & 1\\\\\n",
    "    0 & 1 & 0 & \\frac{1}{3} & -\\frac{2}{3} & 0\\\\\n",
    "    0 & 0 & 1 & 0 & 1 & 0\\\\\n",
    "    \\end{bmatrix}\n",
    "$\n",
    "8. Finally, we can get $Z^{-1}$ by taking the last 3 columns.\n",
    "\\begin{equation}\n",
    "    Z^{-1} = \n",
    "    \\begin{bmatrix}\n",
    "    -\\frac{5}{3} & \\frac{1}{3} & 1\\\\\n",
    "    \\frac{1}{3} & -\\frac{2}{3} & 0\\\\\n",
    "    0 & 1 & 0\\\\\n",
    "    \\end{bmatrix}\n",
    "\\end{equation}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Implement the function `invert_matrix` that takes in a matrix $A$ and returns the inverse matrix. If $A$ is not invertible, `invert_matrix` should return `False`.\n",
    "\n",
    "*IMPORTANT NOTE:* **Do not use NumPy for this question**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "id": "m1XBH6qwrytj"
   },
   "outputs": [],
   "source": [
    "def invert_matrix(A):\n",
    "    \"\"\"\n",
    "    Returns the inverse of matrix A, if it exists; otherwise, returns False\n",
    "    \"\"\"\n",
    "    if len(A[0]) != len(A):\n",
    "        return False\n",
    "    A_len = len(A)\n",
    "    result = copy.deepcopy(A)\n",
    "    # Step 0\n",
    "    for i in range(A_len):\n",
    "        result[i].extend([0] * A_len)\n",
    "        result[i][i + A_len] = 1\n",
    "    result = [[float(i) for i in row] for row in result]\n",
    "    # Step 1\n",
    "    for i in range(A_len):\n",
    "        # Step 1\n",
    "        for k in range(i, A_len):\n",
    "            if result[k][i] != 0:\n",
    "                break\n",
    "        if result[k][i] == 0:\n",
    "            return False\n",
    "        result[i], result[k] = result[k], result[i]\n",
    "        # Step 2\n",
    "        scalar = 1 / result[i][i]\n",
    "        result[i] = [scalar * x for x in result[i]]\n",
    "        # Step 3: Add multiples of the new ith row to all other rows such that the value in their ith column becomes 0\n",
    "        for k in range(A_len):\n",
    "            if k == i:\n",
    "                continue\n",
    "            scalar = -result[k][i]\n",
    "            result[k] = [result[k][j] + scalar * result[i][j] for j in range(2 * A_len)]\n",
    "        \n",
    "    for i in range(A_len):\n",
    "        result[i] = result[i][A_len:]\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test case for Task 1.5\n",
    "A = [[1, 0 ,0], [0, 1, 0], [0, -4, 1]]\n",
    "A_copy = copy.deepcopy(A)\n",
    "\n",
    "actual = invert_matrix(A)\n",
    "expected = [[1, 0 ,0], [0, 1, 0], [0, 4, 1]]\n",
    "assert(A == A_copy)\n",
    "for i in range(len(A)):\n",
    "    for j in range(len(A[0])):\n",
    "        assert(round(actual[i][j], 11) == round(expected[i][j], 11))\n",
    "        \n",
    "        \n",
    "A2 = [[0, 3, 2], [0, 0, 1], [1, 5, 3]]\n",
    "A2_copy = copy.deepcopy(A2)\n",
    "\n",
    "actual2 = invert_matrix(A2)\n",
    "expected2 = [[-5/3, 1/3 ,1], [1/3, -2/3, 0], [0, 1, 0]]\n",
    "assert(A2 == A2_copy)\n",
    "for i in range(len(A2)):\n",
    "    for j in range(len(A2[0])):\n",
    "        assert(round(actual2[i][j], 11) == round(expected2[i][j], 11))\n",
    "        \n",
    "        \n",
    "A3 = [[1, 0, 0], [0, 1, 0], [0, 0, 0]] # non-invertible matrix\n",
    "actual3 = invert_matrix(A3)\n",
    "expected3 = False\n",
    "assert actual3 == expected3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "J4pQSFqk5pi6"
   },
   "source": [
    "# Part 2: Introduction to NumPy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "v7NPTROk6Dp9"
   },
   "source": [
    "Now that we have a better grasp of Python, let us take a look at NumPy, a commonly used package for multidimensional array manipulation! This will come in handy when we work with certain machine learning models later in the semester. Besides, in machine learning problems, large amount of data is usually present, necessitating approaches that make manipulating them easy and fast; and one such approach is to represent our data as multidimensional arrays before manipulating them.\n",
    "\n",
    "The tasks in this section are designed to give us the opportunity to apply what we are about to learn about NumPy to analyse [COVID-19 data](https://github.com/OxCGRT/covid-policy-tracker)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Invert array without NumPy takes 0.24375009536743164s\n",
      "Invert array with NumPy takes 0.0015101432800292969s\n"
     ]
    }
   ],
   "source": [
    "import random\n",
    "import time\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "random.seed(2109)\n",
    "A = [[random.random() for j in range(100)] for i in range(100)] # 100 x 100\n",
    "start = time.time()\n",
    "_ = invert_matrix(A)\n",
    "end = time.time()\n",
    "print(f'Invert array without NumPy takes {end - start}s')\n",
    "\n",
    "A_numpy = np.array(A)\n",
    "start = time.time()\n",
    "_ = np.linalg.inv(A)\n",
    "end = time.time()\n",
    "print(f'Invert array with NumPy takes {end - start}s')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that **each task is independent**. In other words, even if you are stuck at a particular question, it is still possible to get full marks for the later questions.\n",
    "\n",
    "**IMPORTANT**: It is critical that we realise that the underlying implementation of the operations available in NumPy are often optimised for performance. As a result, using these operations, as opposed to iterative approaches, to manipulate elements in the array will often lead to more performant code. This, coupled with the fact that we want to practise using NumPy in this part of the homework, means that **your implementation in the following tasks should not involve any iteration, including `map` and `filter`, recursion, or any iterative approaches like for-loops**. Instead, please work with the operations available in NumPy. \n",
    "\n",
    "You are allowed to use any [mathematical functions](https://numpy.org/doc/stable/reference/routines.math.html), but this **does not mean that you are allowed to use any NumPy functions** (there are NumPy functions that aren't mathematical functions). For example,[ np.vectorize ](https://numpy.org/doc/stable/reference/generated/numpy.vectorize.html) is not allowed since it is iterative. Please refer to the forum post \"List of allowed numpy functions for assignments\" under \"Homework\" for a (non-exhaustive) list of functions that are considered non-iterative for PS0 and for the CS2109S course. If you are in doubt about an unlisted function, please ask in the forum. **Solutions that violate this will be heavily penalised**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z0uLebHS7W6u"
   },
   "source": [
    "## NumPy Basics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yVL9SBwa7Y00"
   },
   "source": [
    "### Import NumPy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "R9eik6gP7c9r"
   },
   "source": [
    "In order to use NumPy, we need to import it. This can be done by"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "id": "GGPjViIWry89"
   },
   "outputs": [],
   "source": [
    "import numpy\n",
    "a = numpy.arange(5) # Returns a NumPy array [0, 1, 2, 3, 4]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dA4YI8LT7kCB"
   },
   "source": [
    "However, it is customary to use the shorthand np instead. Therefore, we will do"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "id": "V5P8ESNy7lye"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "a = np.arange(5) # Note that here we use `np` instead of `numpy`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UQaoJky47t6s"
   },
   "source": [
    "### NumPy Arrays"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "W3NDseCE7-WB"
   },
   "source": [
    "NumPy arrays are objects in NumPy that represent multidimensional arrays. Such\n",
    "arrays have the data type `numpy.ndarray`, or `np.ndarray` for short. Unlike Python lists, NumPy arrays cannot store data that are of different types.\n",
    "\n",
    "To create a NumPy array from a Python list, we can do"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {
    "id": "1unOnNqB72_3"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(3,)\n",
      "1 2 3\n",
      "[9 2 3]\n",
      "(2, 3)\n",
      "1 5 3\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "a = np.array([1, 2, 3]) # Create 1D array vector\n",
    "print(a.shape)          # Prints(3, )\n",
    "print(a[0], a[1], a[2]) # Prints 1, 2, 3\n",
    "a[0] = 9                # Change the zeroth element to 9\n",
    "print(a)                # Prints[9 2 3]\n",
    "\n",
    "b = np.array([[1, 2, 3], [4, 5, 6]])  # Creates 2D array (matrix)\n",
    "print(b.shape)                        # Prints(2, 3)\n",
    "print(b[0, 0], b[1, 1], b[0, 2])      # Prints 1, 5, 3\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tW_eQ0CN8uvW"
   },
   "source": [
    "Alternatively, we can also create pre-filled arrays with a specified shape with certain NumPy functions. Several examples of these are shown below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "id": "ZlO57Ak38pX3"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Object a\n",
      "[[0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]]\n",
      "Object b\n",
      "[1. 1.]\n",
      "Object c\n",
      "[[1. 1. 1.]\n",
      " [1. 1. 1.]\n",
      " [1. 1. 1.]]\n",
      "Object d\n",
      "[[False False False]\n",
      " [False False False]]\n",
      "Object e\n",
      "[0 1 2 3 4]\n"
     ]
    }
   ],
   "source": [
    "a = np.zeros((3, 3))        # Create 3x3 matrix with all zeros\n",
    "b = np.ones(2)              # Create vector of size 2 with all ones\n",
    "c = np.ones((3, 3))         # Create 3x3 matrix with all ones\n",
    "d = np.full((2, 3), False)  # Create a constant array\n",
    "e = np.arange(5)            # Creates a 1D array with values [0, 5)\n",
    "\n",
    "print('Object a')\n",
    "print(a)\n",
    "print('Object b')\n",
    "print(b)\n",
    "print('Object c')\n",
    "print(c)\n",
    "print('Object d')\n",
    "print(d)\n",
    "print('Object e')\n",
    "print(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "oc4X5elN9oNp"
   },
   "source": [
    "Note that we can use functions like `np.zeros` and `np.ones` to pre-allocate arrays. We can then assign values to the entries using indexing.\n",
    "More information about array creation in NumPy can be found [here](https://numpy.org/doc/stable/user/basics.creation.html#arrays-creation)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TQ-gm6-I9yxd"
   },
   "source": [
    "### Basic Indexing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MKuUe8Dj96Mn"
   },
   "source": [
    "Accessing the value of an element in a NumPy array is similar to that of lists in Python. However, observe that in the following, instead of doing something like a[i][j], we have `a[i, j]`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "id": "aVtAsMWM-Bsi"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\n",
      "6\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "a = np.array([[1, 2, 3], [4, 5, 6]])\n",
    "print(a[0, 0])                       # Prints 1\n",
    "print(a[1, 2])                       # Prints 6"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Eghq5yqQ-IEN"
   },
   "source": [
    "Slicing works very similarly too! For example,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "id": "WH1UVm53-JRA"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2 3]\n",
      "[[2 3]\n",
      " [5 6]]\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "a = np.array([[1, 2, 3, 4], [4, 5, 6, 7]])\n",
    "print(a[0, 1:3])                            # Prints [2, 3]\n",
    "print(a[:, 1:3])                            # Prints [[2 3]\n",
    "                                            #         [5 6]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YEC-5kzB96Qe"
   },
   "source": [
    "However, do be careful of aliasing. The following illustrates an example of a side-effect that is caused by aliasing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {
    "id": "b4aPSc9V9FqU"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5\n"
     ]
    }
   ],
   "source": [
    "b = a                                       # Aliasing\n",
    "b[0, 0] = 5\n",
    "print(a[0, 0])                              # Prints \"5\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "B7icHqWF-Sy3"
   },
   "source": [
    "### Element-Wise Math Operation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NISyJMKR-XWO"
   },
   "source": [
    "NumPy provides several *element-wise math operations*. For example,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "id": "vNx9Wz_c-ROf"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 6  8]\n",
      " [10 12]]\n",
      "[[ 5 12]\n",
      " [21 32]]\n"
     ]
    }
   ],
   "source": [
    "x = np.array([[1,2],[3,4]])\n",
    "y = np.array([[5,6],[7,8]])\n",
    "\n",
    "print(x + y)                   # Prints [[ 6  8]\n",
    "                               #         [10 12]]\n",
    "\n",
    "print(x * y)                   # Prints [[ 5 12]\n",
    "                               #         [21 32]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "81ptLkxU-a-w"
   },
   "source": [
    "Similarly, operations for subtraction and division are also provided, and are given by `-` and `/`, respectively.\n",
    "\n",
    "**Note**: `*` is an element-wise multiplication of two matrices, not a matrix multiplication of\n",
    "two matrices, as covered in task 1.4."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### General Math Operation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "NumPy provides many *math operations*. We give some examples below. For a full list of what Numpy can do, please have a look at the following links:\n",
    "\n",
    "* [Math](https://numpy.org/doc/stable/reference/routines.math.html)\n",
    "* [Statistics](https://numpy.org/doc/stable/reference/routines.statistics.html)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Sum\n",
    "\n",
    "`numpy.sum` computes the sum of array elements over a given axis.\n",
    "\n",
    "The following is a sample execution of how this function works:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.0\n",
      "2.0\n",
      "6\n",
      "[0 6]\n",
      "[1 5]\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "print(np.sum([]))                       # 0.0\n",
    "print(np.sum([0.5, 1.5]))               # 2.0\n",
    "print(np.sum([[0, 1], [0, 5]]))         # 6\n",
    "print(np.sum([[0, 1], [0, 5]], axis=0)) # array([0, 6])\n",
    "print(np.sum([[0, 1], [0, 5]], axis=1)) # array([1, 5])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Mean\n",
    "\n",
    "`numpy.mean` computes the average of the array elements. The average is taken over the flattened array by default, otherwise over the specified axis. \n",
    "\n",
    "The following is a sample execution of how this function works:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.5\n",
      "[2. 3.]\n",
      "[1.5 3.5]\n"
     ]
    }
   ],
   "source": [
    "a = np.array([[1, 2], [3, 4]])\n",
    "print(np.mean(a))         # 2.5\n",
    "print(np.mean(a, axis=0)) # array([2., 3.])\n",
    "print(np.mean(a, axis=1)) # array([1.5, 3.5])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Floor & Ceil\n",
    "\n",
    "`numpy.floor` returns the floor of the input, element-wise. The floor of the scalar $x$ is the largest integer $i$, such that $i \\le x$. It is often denoted as $\\lfloor x \\rfloor$.\n",
    "\n",
    "`numpy.ceil` returns the ceiling of the input, element-wise. The ceil of the scalar x is the smallest integer i, such that i >= x. It is often denoted as $\\lceil x \\rceil$.\n",
    "\n",
    "The following is a sample execution of how this function works:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[-2. -2. -1.  0.  1.  1.  2.]\n",
      "[-1. -1. -0.  1.  2.  2.  2.]\n"
     ]
    }
   ],
   "source": [
    "a = np.array([-1.7, -1.5, -0.2, 0.2, 1.5, 1.7, 2.0])\n",
    "print(np.floor(a)) # array([-2., -2., -1.,  0.,  1.,  1.,  2.])\n",
    "print(np.ceil(a))  # array([-1., -1., -0.,  1.,  2.,  2.,  2.])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "F-qWcOyh-q-w"
   },
   "source": [
    "## COVID-19 Data as NumPy Arrays"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nu8iuUGZ-trh"
   },
   "source": [
    "Now that we understand the representation of NumPy arrays, we will try to learn how to use NumPy by working on some practical problems. In particular, we shall take a look at the COVID-19 data which will be used in the following tasks. These data are [time series data](https://en.wikipedia.org/wiki/Time_series) obtained from 1 January 2020 to 31 December 2020. \n",
    "\n",
    "For the convenience of notation, we shall refer to 1 January 2020 as the $1$st day and 31 December 2019 as the $0$th day. Observe that this means that the following data **do not** contain measurements on the 0th day, i.e. 31 December 2019.} In addition, in the following tasks, assume that $i, j \\geq 0$, whenever they are used as indexes.\n",
    "\n",
    "First, let us describe the data set that we will be working with. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uLwkuhfF3PiC"
   },
   "source": [
    "### **`cases_cumulative`**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HUbvGwPK3UPD"
   },
   "source": [
    "It is a 2D `np.ndarray` with each row representing the data of a country while the columns of each row represent the time series data of the cumulative number of confirmed cases in that country, i.e. the $i$-th row of `cases_cumulative` contains the data of the $i$-th country, and the $(i, j)$ entry of `cases_cumulative` is the *cumulative number* of confirmed cases on the $(j + 1)$th day in the $i$th country.\n",
    "\n",
    "For example, suppose"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {
    "id": "3gbHqkGa-fkB"
   },
   "outputs": [],
   "source": [
    "cases_cumulative = np.array([[0, 1, 2, 3], [0, 20, 21, 35]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kGeE-oed3ofk"
   },
   "source": [
    "This means that the 0th country has recorded 0, 1, 2 and 3 cases on the 1st, 2nd, 3rd and 4th day, respectively. Similarly, the 1st country has recorded 0, 20, 21 and 35 cases on the 1st, 2nd, 3rd and 4th day.\n",
    "\n",
    "**Important**:\n",
    "\n",
    "Observe that ${(i, j)}$ entry corresponds to the data for the ${i}$-th country but ${(j + 1)}$-th day. For day 0, assume that the number of cases recorded is 0 for all countries.\n",
    "\n",
    "Here, *cumulative number* means that the recorded data is the total number of cases detected since day 0, not only the new cases that have been detected on a particular day. Hence, generally, `cases_cumulative[i, j]` $\\leq$ `cases_cumulative[i, j + 1]`.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "K-bF65CK3tBV"
   },
   "source": [
    "### **`cases_top_cumulative`**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "oV4X8qcY4N8u"
   },
   "source": [
    "This variable is similar to `cases_cumulative`. However, it only contains data for the three countries that recorded the most number of confirmed COVID-19 cases in 2020. These data might come in handy during debugging and testing.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XPBq32dV4SpE"
   },
   "source": [
    "### **`deaths_cumulative`**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hATkVnyx4ViC"
   },
   "source": [
    "`deaths_cumulative` is similar to `cases_cumulative`, but instead of recording the cumulative number of confirmed cases, it records the cumulative number of deaths caused by COVID-19.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hdnlmCc04kzp"
   },
   "source": [
    "### **`healthcare_spending`**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SjvPjJ9i4k28"
   },
   "source": [
    "This is a 2D `np.ndarray` with each row representing the data of a country while the columns of each row represent the time series data of the emergency healthcare spending made by that country, i.e. the $i$-th row of `healthcare_spending` contains the data of the $i$-th country, and the $(i, j)$ entry of `healthcare_spending` is the amount which the $i$-th country spent on healthcare on the $(j + 1)$-th day. Assume that the healthcare spending on the 0-th day is \\$0 for all countries.\n",
    "\n",
    "For example, suppose"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "L5Dg8kPE3rid"
   },
   "outputs": [],
   "source": [
    "healthcare_spending = np.array([[0, 100, 0, 200], [0, 0, 0, 1000]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AIKQ6qCb46hu"
   },
   "source": [
    "Then, the 0th country only made an emergency expenditure on healthcare on the 2nd and 4th days that amounted to \\$100 and \\$200, respectively; similarly, the 1st country only made an emergency expenditure on the 4th day which amounted to \\$1000.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mtbqt7uC5Bv3"
   },
   "source": [
    "### **`mask_prices`**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ce-xxi_t5Dlo"
   },
   "source": [
    "`mask_prices` is a 1D `np.ndarray` such that the $j$-th entry represents the global average cost of 100 masks on the $(j + 1)$-th day.\n",
    "\n",
    "For example, if"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "4x0dxMTJ4-X8"
   },
   "outputs": [],
   "source": [
    "mask_prices = np.array([4, 5, 20, 18])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WCnTULNz5R8y"
   },
   "source": [
    "then the global average cost of 100 masks on the 1st, 2nd, 3rd and 4th day is \\$4, \\$5, \\$20 and \\$18, respectively.\n",
    "\n",
    "**Note**: The prices are not real but the overall fluctuations are chosen to simulate those observed in 2020."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "n70RPHKx5fne"
   },
   "source": [
    "### **`stringency_values`**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_4yj4Ts55hsE"
   },
   "source": [
    "This variable is a 3D `np.ndarray` with each row representing the data of a country, and the columns of each row representing the time series data of the stringency values as a vector. To be specific, on each day, there are four different stringency indicators: 'school closing', 'workplace closing', 'stay at home requirements' and 'international travel controls'. \n",
    "\n",
    "For instance, suppose"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JY_K3jkh5Yo_"
   },
   "outputs": [],
   "source": [
    "stringency_values = np.array([[[0, 0, 0, 0], [1, 0, 0, 0]],\\\n",
    "    [[0, 0, 0, 0], [0, 1, 2, 0]]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8Beq-_7r5hxI"
   },
   "source": [
    "This means that for the 0th country, on the 1st day, the stringency values for 'school closing', 'workplace closing', 'stay at home requirements' and 'international travel controls' were all 0. On the 2nd day, all the stringency values remained the same, except for 'school closing', which rose to 1.\n",
    "\n",
    "As for the 1st country, on the 1st day, the stringency values were the same as the 0th country's. However, on the 2nd day, the stringency values for 'workplace closing' and 'stay at home requirements' rose to 1 and 2, respectively.\n",
    "\n",
    "In this context, the higher the stringency value, the more restrictive the policy is."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get the data from CSV"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As shown in the last few lines of `ps0.py`, we can get the data for these variables by doing the following"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from prepare_data import * # loads the `get_...` helper funtions\n",
    "\n",
    "df = get_data()\n",
    "cases_cumulative = get_n_cases_cumulative(df)\n",
    "deaths_cumulative = get_n_deaths_cumulative(df)\n",
    "healthcare_spending = get_healthcare_spending(df)\n",
    "mask_prices = get_mask_prices(healthcare_spending.shape[1])\n",
    "stringency_values = get_stringency_values(df)\n",
    "cases_top_cumulative = get_n_cases_top_cumulative(df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iUOLLiiK52kC"
   },
   "source": [
    "In the upcoming tasks, we will be manipulating these data, and finding answers to certain questions that we might have about issues that are related to COVID-19.\n",
    "\n",
    "**IMPORTANT**: You are allowed to use any of the [mathematical functions](https://numpy.org/doc/stable/reference/routines.math.html) from NumPy in all of the following tasks.\n",
    "For your convenience, you may use some of the non-mathematical functions provided by NumPy, but not all of them. For example, [`np.vectorize`](https://numpy.org/doc/stable/reference/generated/numpy.vectorize.html) isn't allowed since it is iterative. If in doubt, do ask in the forum :) With this said, do note that it is absolutely possible to complete these tasks without the non-mathematical functions that are not introduced in this document.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XtMITSIQ6wXR"
   },
   "source": [
    "## Task 2.1: Computing Death Rates"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9BTFV_PG6yvu"
   },
   "source": [
    "Our first task in this part of the homework is to implement `compute_death_rate_first_n_days`, which takes in three arguments, `n`, `cases_cumulative` and `deaths_cumulative`, and **computes the average number of deaths recorded for every confirmed case that has been recorded from the first day to the nth day (inclusive)**. This should be done for each country.\n",
    "\n",
    "The return value should be a `np.ndarray` such that the entry in the $i$-th row corresponds to the death rate in the $i$-th country as represented in `cases_cumulative` and `deaths_cumulative`.\n",
    "\n",
    "For instance, if the returned value is `np.array([0.5, 0.2])`, it means that in the 0th country, for every 2 individuals who contracted the virus, one of them will, on average, die. In contrast, in the 1st country, for every 5 individuals who contracted the virus, only one of them will, on average, die.\n",
    "\n",
    "**Note**:\n",
    "You may assume that the $i$-th row in `cases_cumulative` represents the same country as the $i$-th row in `deaths_cumulative`. Moreover, if there are no confirmed cases for a particular country, its average death rate should be zero. If the data includes less than `n` days of data, then you just return the result for all the days given in the data provided. \n",
    "\n",
    "Your implementation should not involve any iteration, including `map` and `filter`, recursion, or any iterative approaches like for-loops.\n",
    "\n",
    "In this task, the goal is to learn how to deal with `nan` values using `np.nan_to_num`, which has the following signature: `numpy.nan_to_num(x, copy=True, nan=0.0, posinf=None, neginf=None)`. You should also check out the detailed description [here](https://numpy.org/doc/stable/reference/generated/numpy.nan_to_num.html) because there are many functions of the same nature. Essentially, we can use `np.nan_to_num` to convert all the `nan` values in an `np.ndarray` to a specified value. Some examples of how `np.nan_to_num` can be used is shown below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(np.nan_to_num(np.inf)) #1.7976931348623157e+308\n",
    "print(np.nan_to_num(-np.inf)) #-1.7976931348623157e+308\n",
    "print(np.nan_to_num(np.nan)) #0.0\n",
    "x = np.array([np.inf, -np.inf, np.nan, -128, 128])\n",
    "np.nan_to_num(x)\n",
    "np.nan_to_num(x, nan=-9999, posinf=33333333, neginf=33333333)  # Specify nan to be -9999, and both posinf and neginf to be 33333333\n",
    "np.nan_to_num(y, nan=111111, posinf=222222) # Specify nan to be 111111, and both posinf and neginf to be 222222"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SOF2JEoW5xJE"
   },
   "outputs": [],
   "source": [
    "def compute_death_rate_first_n_days(n, cases_cumulative, deaths_cumulative):\n",
    "    '''\n",
    "    Computes the average number of deaths recorded for every confirmed case\n",
    "    that is recorded from the first day to the nth day (inclusive).\n",
    "    Parameters\n",
    "    ----------\n",
    "    n: int\n",
    "        How many days of data to return in the final array.\n",
    "    cases_cumulative: np.ndarray\n",
    "        2D `ndarray` with each row representing the data of a country, and the columns\n",
    "        of each row representing the time series data of the cumulative number of\n",
    "        confirmed cases in that country, i.e. the ith row of `cases_cumulative`\n",
    "        contains the data of the ith country, and the (i, j) entry of\n",
    "        `cases_cumulative` is the cumulative number of confirmed cases on the\n",
    "        (j + 1)th day in the ith country.\n",
    "    deaths_cumulative: np.ndarray\n",
    "        2D `ndarray` with each row representing the data of a country, and the columns\n",
    "        of each row representing the time series data of the cumulative number of\n",
    "        confirmed deaths (as a result of COVID-19) in that country, i.e. the ith\n",
    "        row of `n_deaths_cumulative` contains the data of the ith country, and\n",
    "        the (i, j) entry of `n_deaths_cumulative` is the cumulative number of\n",
    "        confirmed deaths on the (j + 1)th day in the ith country.\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    Average number of deaths recorded for every confirmed case from the first day\n",
    "    to the nth day (inclusive) for each country as a 1D `ndarray` such that the\n",
    "    entry in the ith row corresponds to the death rate in the ith country as\n",
    "    represented in `cases_cumulative` and `deaths_cumulative`.\n",
    "    Note\n",
    "    ----\n",
    "    `cases_cumulative` and `deaths_cumulative` are such that the ith row in the \n",
    "    former and that in the latter contain data of the same country. In addition,\n",
    "    if there are no confirmed cases for a particular country, the expected death\n",
    "    rate for that country should be zero. (Hint: to deal with NaN look at\n",
    "    `np.nan_to_num`)\n",
    "    Your implementation should not involve any iteration, including `map` and `filter`, \n",
    "    recursion, or any iterative approaches like for-loops.\n",
    "    '''\n",
    "    \n",
    "    # TODO: add your solution here and remove `raise NotImplementedError`\n",
    "    raise NotImplementedError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test cases for Task 2.1\n",
    "n_cases_cumulative = cases_cumulative[:3, :] #Using data from CSV. Make sure to run relevant cell above\n",
    "n_deaths_cumulative = deaths_cumulative[:3, :]\n",
    "expected = np.array([0.0337837838, 0.0562347188, 0.1410564226])\n",
    "np.testing.assert_allclose(compute_death_rate_first_n_days(100, n_cases_cumulative, n_deaths_cumulative), expected)\n",
    "\n",
    "sample_cumulative = np.array([[1,2,3,4,8,8,10,10,10,10], [1,2,3,4,8,8,10,10,10,10]])\n",
    "sample_death = np.array([[0,0,0,1,2,2,2,2,5,5], [0,0,0,1,2,2,2,2,5,5]])\n",
    "\n",
    "expected2 = np.array([0.5, 0.5])\n",
    "assert(np.all(compute_death_rate_first_n_days(10, sample_cumulative, sample_death) == expected2))\n",
    "\n",
    "sample_cumulative2 = np.array([[1,2,3,4,8,8,10,10,10,10]])\n",
    "sample_death2 = np.array([[0,0,0,1,2,2,2,2,5,5]])\n",
    "\n",
    "expected3 = np.array([0.5])\n",
    "assert(compute_death_rate_first_n_days(10, sample_cumulative2, sample_death2) == expected3)\n",
    "expected4 = np.array([0.25])\n",
    "assert(compute_death_rate_first_n_days(5, sample_cumulative2, sample_death2) == expected4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uCQ7_juB7ks0"
   },
   "source": [
    "## Task 2.2: Computing Daily Increase in Cases"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "E7W2-0HI72VS"
   },
   "source": [
    "Our second task requires us to implement `compute_increase_in_cases` which accepts one argument `cases_cumulative` and **computes the daily increase in confirmed cases for each country, starting from the first day**.\n",
    "\n",
    "This function should return the daily increase in cases for each country as a 2D `np.ndarray` such that the $(i, j)$ entry corresponds to the increase in confirmed cases in the $i$-th country on the $(j + 1)$-th day, where $j$ is non-negative.\n",
    "\n",
    "Recall that as we have previously mentioned, the number of cases on the 0th day can be assumed to be 0 for all countries.\n",
    "\n",
    "To be specific, suppose `cases_cumulative` only contains data for 4 days such that it is equal to `np.array([[1, 2, 3, 4], [1, 3, 6, 10]])`. Then, the return value of this function should be `np.array([[1, 1, 1, 1], [1, 2, 3, 4]])`. This means that from days 1 to 4, in country 0, there was one new recorded case on each day. In contrast, in country 1, there were 2, 3 and 4 new cases on days 2, 3 and 4, respectively.\n",
    "\n",
    "**Note**:\n",
    "Your implementation should not involve any iteration, including `map` and `filter`, recursion, or any iterative approaches like for-loops.\n",
    "\n",
    "In this task, the goal is to learn how to compute differences between adjacent elements in an array using `np.diff`, which has the following signature: `numpy.diff(a, n=1, axis=-1, prepend=<no value>, append=<no value>)`. `np.diff` calculates the n-th discrete difference along the given axis. The first difference is given by `out[i] = a[i+1] - a[i]` along the given axis; higher differences are calculated by using diff recursively. The following are some examples of how `np.diff` can be used:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.array([1, 2, 4, 7, 0])\n",
    "np.diff(x) # array([1, 2, 3, -7])\n",
    "np.diff(x, n=2) # array([1, 1, -10])\n",
    "\n",
    "x = np.array([[1, 3, 6, 10], [0, 5, 6, 8]])\n",
    "np.diff(x) # array([[2, 3, 4], [5, 1, 2]])\n",
    "np.diff(x, axis=0) # array([[-1, 2, 0, -2]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Le8aWyr17etU"
   },
   "outputs": [],
   "source": [
    "def compute_increase_in_cases(n, cases_cumulative):\n",
    "    '''\n",
    "    Computes the daily increase in confirmed cases for each country for the first n days, starting\n",
    "    from the first day.\n",
    "    Parameters\n",
    "    ----------    \n",
    "    n: int\n",
    "        How many days of data to return in the final array. If the input data has fewer\n",
    "        than n days of data then we just return whatever we have for each country up to n. \n",
    "    cases_cumulative: np.ndarray\n",
    "        2D `ndarray` with each row representing the data of a country, and the columns\n",
    "        of each row representing the time series data of the cumulative number of\n",
    "        confirmed cases in that country, i.e. the ith row of `cases_cumulative`\n",
    "        contains the data of the ith country, and the (i, j) entry of\n",
    "        `cases_cumulative` is the cumulative number of confirmed cases on the\n",
    "        (j + 1)th day in the ith country.\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    Daily increase in cases for each country as a 2D `ndarray` such that the (i, j)\n",
    "    entry corresponds to the increase in confirmed cases in the ith country on\n",
    "    the (j + 1)th day, where j is non-negative.\n",
    "    Note\n",
    "    ----\n",
    "    The number of cases on the zeroth day is assumed to be 0, and we want to\n",
    "    compute the daily increase in cases starting from the first day.\n",
    "    Your implementation should not involve any iteration, including `map` and `filter`, \n",
    "    recursion, or any iterative approaches like for-loops.\n",
    "    '''\n",
    "    \n",
    "    # TODO: add your solution here and remove `raise NotImplementedError`\n",
    "    raise NotImplementedError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test case for Task 2.2\n",
    "cases_cumulative = np.zeros((100, 20))\n",
    "cases_cumulative[:, :] = np.arange(1, 21)\n",
    "actual = compute_increase_in_cases(100, cases_cumulative)\n",
    "assert(np.all(actual == np.ones((100, 20))))\n",
    "\n",
    "sample_cumulative = np.array([[1,2,3,4,8,8,10,10,10,10],[1,1,3,5,8,10,15,20,25,30]])\n",
    "expected = np.array([[1, 1, 1, 1, 4.], [1, 0, 2, 2, 3]])\n",
    "assert(np.all(compute_increase_in_cases(5,sample_cumulative) == expected))\n",
    "\n",
    "expected2 = np.array([[1, 1, 1, 1, 4, 0, 2, 0, 0, 0],[1, 0, 2, 2, 3, 2, 5, 5, 5, 5]])\n",
    "assert(np.all(compute_increase_in_cases(10,sample_cumulative) == expected2))\n",
    "assert(np.all(compute_increase_in_cases(20,sample_cumulative) == expected2))\n",
    "\n",
    "sample_cumulative2 = np.array([[51764, 51848, 52007, 52147, 52330, 52330],\\\n",
    "                            [55755, 56254, 56572, 57146, 57727, 58316],\\\n",
    "                            [97857, 98249, 98631, 98988, 99311, 99610]])\n",
    "expected3 = np.array([\\\n",
    "            [51764, 84, 159, 140, 183, 0],\\\n",
    "            [55755, 499, 318, 574, 581, 589],\\\n",
    "            [97857, 392, 382, 357, 323, 299]])\n",
    "assert(np.all(compute_increase_in_cases(6,sample_cumulative2) == expected3))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yx-hn5dC8RA6"
   },
   "source": [
    "## Axes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ggXC9VM-8YFf"
   },
   "source": [
    "When using certain NumPy functions, we might encounter the axis argument. [`np.sum`](https://numpy.org/doc/stable/reference/generated/numpy.sum.html) is an example.\n",
    "\n",
    "For a 3D (or 2D, if we ignore the blue slice at the back which contains the values 2, 4, 6 and 8), the axes are defined as shown in figure below.\n",
    "\n",
    "<img src = 'imgs/numpy_axes.png'>\n",
    "\n",
    "*Note*: NumPy arrays can have an arbitrary positive integer number of axes. If we have an N-D array with N > 3, we just need to (mentally) project it onto a smaller dimensional space so that we can have a handle of what each axis means.}\n",
    "\n",
    "Therefore, returning to `np.sum`, we see that"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GtCGMCpp8QQT"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "a = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])\n",
    "print(np.sum(a, axis=0))      # Prints [[ 6,  8],\n",
    "                              #         [10, 12]]\n",
    "print(np.sum(a, axis=1))      # Prints [[ 4,  6],\n",
    "                              #         [12, 14]]\n",
    "print(np.sum(a, axis=2))      # Prints [[ 3,  7],\n",
    "                              #         [11, 15]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aF5DV48W9AKp"
   },
   "source": [
    "## Task 2.3: Finding Maximum Daily Increase in Cases"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nKwWcvbp9CiU"
   },
   "source": [
    "Our third task is to implement `find_max_increase_in_cases` which takes in one argument `n_cases_increase` and **finds the maximum daily increase in confirmed cases for each country**.\n",
    "\n",
    "In this case, `n_cases_increase` is the output obtained from calling `compute_increase_in_cases(n_cases_cumulative)`.\n",
    "\n",
    "The return value should be a 1D `np.ndarray` that represents the maximum daily increase in cases for each country. In particular, the $i$-th entry of this array should correspond to the increase in confirmed cases in the $i$-th country as represented in `n_cases_increase`.\n",
    "\n",
    "Returning to our previous example in task 2.2, suppose the daily increase in cases is given by `np.array([[1, 1, 1, 1], [1, 2, 3, 4]])`. Clearly, the maximum daily increase in cases is 1 and 4 for country 0 and 1, respectively. Therefore, we should expect the output of this function to be `np.array([1, 4])`.\n",
    "\n",
    "**Note**:\n",
    "Your implementation should not involve any iteration, including `map` and `filter`, recursion, or any iterative approaches like for-loops.\n",
    "\n",
    "In this task, the goal is to learn how to use `np.max`, or the equivalent `np.amax`, to find the max value in a column of a matrix. The following is a sample execution of how these functions work:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = np.arange(4).reshape((2,2)) # array([[0, 1], [2, 3]])\n",
    "np.amax(a)  # 3 ->  Maximum of the flattened array\n",
    "np.amax(a, axis=0)  # array([2, 3]) -> Maxima along the first axis (first column) \n",
    "np.amax(a, axis=1)  # array([1, 3]) -> Maxima along the second axis (second column)\n",
    "np.amax(a, where=[False, True], initial=-1, axis=0) # array([-1,  3])\n",
    "b = np.arange(5, dtype=float) # array([0., 1., 2., 3., 4.])\n",
    "b[2] = np.NaN  # array([ 0., 1., nan, 3., 4.])\n",
    "np.amax(b) # nan"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TnYgBS0m9B63"
   },
   "outputs": [],
   "source": [
    "def find_max_increase_in_cases(n_cases_increase):\n",
    "    '''\n",
    "    Finds the maximum daily increase in confirmed cases for each country.\n",
    "    Parameters\n",
    "    ----------\n",
    "    n_cases_increase: np.ndarray\n",
    "        2D `ndarray` with each row representing the data of a country, and the columns\n",
    "        of each row representing the time series data of the daily increase in the\n",
    "        number of confirmed cases in that country, i.e. the ith row of \n",
    "        `n_cases_increase` contains the data of the ith country, and the (i, j) entry of\n",
    "        `n_cases_increase` is the daily increase in the number of confirmed cases on the\n",
    "        (j + 1)th day in the ith country.\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    Maximum daily increase in cases for each country as a 1D `ndarray` such that the\n",
    "    ith entry corresponds to the increase in confirmed cases in the ith country as\n",
    "    represented in `n_cases_increase`.\n",
    "    Your implementation should not involve any iteration, including `map` and `filter`, \n",
    "    recursion, or any iterative approaches like for-loops.\n",
    "    '''\n",
    "    \n",
    "    # TODO: add your solution here and remove `raise NotImplementedError`\n",
    "    raise NotImplementedError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test case for Task 2.3\n",
    "n_cases_increase = np.ones((100, 20))\n",
    "actual = find_max_increase_in_cases(n_cases_increase)\n",
    "expected = np.ones(100)\n",
    "assert(np.all(actual == expected))\n",
    "\n",
    "sample_increase = np.array([[1,2,3,4,8,8,10,10,10,10],[1,1,3,5,8,10,15,20,25,30]])\n",
    "expected2 = np.array([10, 30]) # max of [1,2,3,4,8,8,10,10,10,10] => 10, max of [1,1,3,5,8,10,15,20,25,30] => 30\n",
    "assert(np.all(find_max_increase_in_cases(sample_increase) == expected2))\n",
    "\n",
    "sample_increase2 = np.array([\\\n",
    "            [51764, 84, 159, 140, 183, 0],\\\n",
    "            [55755, 499, 318, 574, 581, 589],\\\n",
    "            [97857, 392, 382, 357, 323, 299]])\n",
    "expected3 = np.array([51764, 55755, 97857])\n",
    "assert(np.all(find_max_increase_in_cases(sample_increase2) == expected3))\n",
    "\n",
    "n_cases_increase2 = compute_increase_in_cases(cases_top_cumulative.shape[1], cases_top_cumulative)\n",
    "expected4 = np.array([ 68699.,  97894., 258110.])\n",
    "assert(np.all(find_max_increase_in_cases(n_cases_increase2) == expected4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "e10BJ74p9eS-"
   },
   "source": [
    "## Broadcasting"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PPuJv8kC9hJb"
   },
   "source": [
    "*Broadcasting* makes it possible for us to apply arithmetic operations to arrays with different shapes. This is especially helpful when we have a smaller array and larger array such that we want to copy the smaller array multiple times such that its shape (i.e. the number of elements in each axis) is the same as the larger array's, before performing arithmetic operations on both arrays. Specifically, with broadcasting, we do not need to explicitly make copies of the smaller array for the arithmetic operation to work. \n",
    "\n",
    "For example,\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5danymoS9cqj"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "a = np.array([[1, 2, 3], [4, 5, 6]])\n",
    "b = np.array([[0, 1, 2]])\n",
    "c = np.array([[0, 1, 2], [0, 1, 2]])\n",
    "d = np.full((2, 3), 5)\n",
    "\n",
    "print(a + b)                           # Prints [[1 3 5]\n",
    "                                       #         [4 6 8]]\n",
    "                                       # Equivalent to a + c\n",
    "print(a + 5)                           # Prints [[ 6,  7,  8],\n",
    "                                       #         [ 9, 10, 11]]\n",
    "                                       # Equivalent to a + d"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3wZr_D4z9s3C"
   },
   "source": [
    "The details of broadcasting can be found [here](https://numpy.org/doc/stable/user/basics.broadcasting.html).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EeYafm-090kt"
   },
   "source": [
    "At times, we might need to add dimensions to one of the arrays so that broadcasting gives us our desired result. To do so, we can use `None`. \n",
    "\n",
    "For example, suppose we want to create a $3 \\times 2$ matrix, represented by `c`, from `a = np.array([4, 5, 6])` and `b = np.array([1, 2])` such that `c[i, j] = a[i] + b[j]`. Then, we can do the following"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {
    "id": "nbe68bm--FuQ"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[4 5 6]\n",
      "[1 2]\n",
      "[[4]\n",
      " [5]\n",
      " [6]]\n",
      "[[1 2]]\n",
      "[[5 6]\n",
      " [6 7]\n",
      " [7 8]]\n"
     ]
    }
   ],
   "source": [
    "a = np.array([4, 5, 6])\n",
    "print(a)                        # Prints [4, 5, 6]\n",
    "b = np.array([1, 2])\n",
    "print(b)                        # Prints [1, 2]\n",
    "c = a[:, None] \n",
    "print(c)                        # Prints [[4]\n",
    "                                #         [5]\n",
    "                                #         [6]]\n",
    "d = b[None, :]\n",
    "print(d)                        # Prints [[1 2]]\n",
    "e = c + d\n",
    "print(e)                        # Prints [[5 6]\n",
    "                                #         [6 7]\n",
    "                                #         [7 8]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qhIjMRK1-JM-"
   },
   "source": [
    "Alternatively, we can use `reshape` to get the same results.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fulFjQf8-Icw"
   },
   "outputs": [],
   "source": [
    "c = a.reshape((3, 1)) + b\n",
    "print(c) # Prints [[5 6]\n",
    "         #         [6 7]\n",
    "         #         [7 8]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ts_b7Hrk-QZ3"
   },
   "source": [
    "## Task 2.4: Computing Number of Purchaseable Masks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6DhPv99S-Tin"
   },
   "source": [
    "Our fourth task is to implement `compute_n_masks_purchaseable`. This function takes two arguments, namely `healthcare_spending` and `mask_prices`, and it should **compute the total number of masks that each country can purchase if she spends all her emergency healthcare spending on masks**.\n",
    "\n",
    "Assume that all countries bought the masks at the global average costs, and to reduce administrative hassle, they only buy masks on day $(j + 1)$ with their emergency funds on day $(j + 1)$, i.e. the masks that are bought on a particular day are not purchased with funding from the previous days.\n",
    "\n",
    "The return value should be the total number of masks, which each country can purchase, represented as a 1D `np.ndarray` such that the $i$-th entry corresponds to the total number of masks purchaseable by the $i$-th country as represented in `healthcare_spending`.\n",
    "\n",
    "For example, if we have `healthcare_spending = np.array([[0, 100, 0], [100, 0, 200]])` and `mask_prices = np.array([4, 3, 20])`. Then, we expect the return value to be `np.array([3300. 3500.])`. This is because country 0 can only buy $33 \\times 100 = 3300$ masks on day 2 with \\$100 since the masks are priced at \\$3 and are sold in batches of 100. Similarly, we find that country 1 can only buy $(25 + 10) \\times 100 = 3500$ masks. \n",
    "\n",
    "In this task, you may use the `np.sum` and the `np.floor` functions.\n",
    "\n",
    "**Note**:\n",
    "Your implementation should not involve any iteration, including `map` and `filter`, recursion, or any iterative approaches like for-loops."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lscWd3LX-SR7"
   },
   "outputs": [],
   "source": [
    "def compute_n_masks_purchaseable(healthcare_spending, mask_prices):\n",
    "    '''\n",
    "    Computes the total number of masks that each country can purchase if she\n",
    "    spends all her emergency healthcare spending on masks.\n",
    "    Parameters\n",
    "    ----------\n",
    "    healthcare_spending: np.ndarray\n",
    "        2D `ndarray` with each row representing the data of a country, and the columns\n",
    "        of each row representing the time series data of the emergency healthcare\n",
    "        spending made by that country, i.e. the ith row of `healthcare_spending`\n",
    "        contains the data of the ith country, and the (i, j) entry of\n",
    "        `healthcare_spending` is the amount which the ith country spent on healthcare\n",
    "        on (j + 1)th day.\n",
    "    mask_prices: np.ndarray\n",
    "        1D `ndarray` such that the jth entry represents the cost of 100 masks on the\n",
    "        (j + 1)th day.\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    Total number of masks which each country can purchase as a 1D `ndarray` such\n",
    "    that the ith entry corresponds to the total number of masks purchaseable by the\n",
    "    ith country as represented in `healthcare_spending`.\n",
    "    Note\n",
    "    ----\n",
    "    The masks can only be bought in batches of 100s.\n",
    "    Your implementation should not involve any iteration, including `map` and `filter`, \n",
    "    recursion, or any iterative approaches like for-loops.\n",
    "    '''\n",
    "    \n",
    "    # TODO: add your solution here and remove `raise NotImplementedError`\n",
    "    raise NotImplementedError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test case for Task 2.4\n",
    "prices_constant = np.ones(5)\n",
    "healthcare_spending_constant = np.ones((7, 5))\n",
    "actual = compute_n_masks_purchaseable(healthcare_spending_constant, prices_constant)\n",
    "expected = np.ones(7) * 500\n",
    "assert(np.all(actual == expected))\n",
    "\n",
    "healthcare_spending1 = healthcare_spending[:3, :]  #Using data from CSV\n",
    "expected2 = [3068779300, 378333500, 6208321700]\n",
    "assert(np.all(compute_n_masks_purchaseable(healthcare_spending1, mask_prices)==expected2))\n",
    "\n",
    "healthcare_spending2 = np.array([[0, 100, 0], [100, 0, 200]])\n",
    "mask_prices2 = np.array([4, 3, 20])\n",
    "expected3 = np.array([3300, 3500])\n",
    "assert(np.all(compute_n_masks_purchaseable(healthcare_spending2, mask_prices2)==expected3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "h-S8a06N_GP6"
   },
   "source": [
    "## Matrix Multiplication"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_lN2l9BV_INc"
   },
   "source": [
    "To perform matrix multiplication on two NumPy arrays, we use the syntax `@`. In particular, if we want to multiply array `A` and `B`, we will do `A @ B`. For example,\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = np.array([[1,1],[1,2]]) # array([[1, 1], [1, 2]])\n",
    "b = np.array([[1],[1]])  # array([[1], [1]])\n",
    "c = a @ b\n",
    "print(c) # Prints [[2]\n",
    "         #         [3]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The specifics of this operations is as follow:\n",
    "\n",
    "- if A and B are both 2D, they are multiplied like conventional matrices.\n",
    "- if one of them is N-D, with N > 2, that matrix is treated as a stack of matrices. Broadcasting will be done where necessary.\n",
    "\n",
    "More details can be found [here](https://numpy.org/doc/stable/user/basics.broadcasting.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZMRfyy8s_gCg"
   },
   "source": [
    "## Task 2.5: Computing Stringency Index"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DdY664xo_hhn"
   },
   "source": [
    "Our fifth last task for this homework is to implement `compute_stringency_index`. This function takes `stringency_values` as argument, and **computes the daily stringency index for each country**.\n",
    "\n",
    "This function returns the daily stringency index for each country as a 2D `np.ndarray` such that the $(i, j)$ entry corresponds to the stringency index in the $i$-th country on the $(j + 1)$-th day. In this case, the higher the stringency index, the more restrictive the measures are.\n",
    "\n",
    "Recall that on each day, each country has four stringency values for ’school  closing’,  ’workplace  closing’, ’stay  at  home  requirements’  and ’international travel controls’. In this case, we shall assume that 'stay at home requirements' is the most restrictive regulation among the other regulations; 'international travel controls' is more restrictive than 'school closing' and 'workplace closing'; and 'school closing' and 'workplace closing' are equally restrictive. Thus, to compute the stringency index, we shall weigh each stringency value by 1, 1, 3 and 2 for 'school closing', 'workplace closing', 'stay at home requirements' and 'international travel controls', respectively. Then, **the stringency index for the $i$th country on the $(j + 1)$th day is given by**:\n",
    "\n",
    "`stringency_values[i, j, 0] + stringency_values[i, j, 1] + 3 * stringency_values[i, j, 2] + 2 * stringency_values[i, j, 3]`\n",
    "\n",
    "**Note**: Use the matrix multiplication operator (`@` operator) to compute the stringency index. Please do not use iterative approaches like for-loops."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TnEqpaNo-z2v"
   },
   "outputs": [],
   "source": [
    "def compute_stringency_index(stringency_values):\n",
    "    '''\n",
    "    Computes the daily stringency index for each country.\n",
    "    Parameters\n",
    "    ----------\n",
    "    stringency_values: np.ndarray\n",
    "        3D `ndarray` with each row representing the data of a country, and the columns\n",
    "        of each row representing the time series data of the stringency values as a\n",
    "        vector. To be specific, on each day, there are four different stringency\n",
    "        values for 'school closing', 'workplace closing', 'stay at home requirements'\n",
    "        and 'international travel controls', respectively. For instance, the (i, j, 0)\n",
    "        entry represents the `school closing` stringency value for the ith country\n",
    "        on the (j + 1)th day.\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    Daily stringency index for each country as a 2D `ndarray` such that the (i, j)\n",
    "    entry corresponds to the stringency index in the ith country on the (j + 1)th\n",
    "    day.\n",
    "    In this case, we shall assume that 'stay at home requirements' is the most\n",
    "    restrictive regulation among the other regulations, 'international travel\n",
    "    controls' is more restrictive than 'school closing' and 'workplace closing',\n",
    "    and 'school closing' and 'workplace closing' are equally restrictive. Thus,\n",
    "    to compute the stringency index, we shall weigh each stringency value by 1,\n",
    "    1, 3 and 2 for 'school closing', 'workplace closing', 'stay at home\n",
    "    requirements' and 'international travel controls', respectively. Then, the \n",
    "    index for the ith country on the (j + 1)th day is given by\n",
    "    `stringency_values[i, j, 0] + stringency_values[i, j, 1] +\n",
    "    3 * stringency_values[i, j, 2] + 2 * stringency_values[i, j, 3]`.\n",
    "    Note\n",
    "    ----\n",
    "    Use matrix operations and broadcasting to complete this question. Please do\n",
    "    not use iterative approaches like for-loops.\n",
    "    '''\n",
    "    \n",
    "    # TODO: add your solution here and remove `raise NotImplementedError`\n",
    "    raise NotImplementedError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test case for Task 2.5\n",
    "stringency_values = np.ones((10, 20, 4))\n",
    "stringency_values[:, 10:, :] *= 2\n",
    "actual = compute_stringency_index(stringency_values)\n",
    "expected = np.ones((10, 20)) * (1 + 1 + 3 + 2)\n",
    "expected[:, 10:] *= 2\n",
    "assert(np.all(actual == expected))\n",
    "\n",
    "stringency_values2 = np.array([[[0, 0, 0, 0], [1, 0, 0, 0]], [[0, 0, 0, 0], [0, 1, 2, 0]]])\n",
    "actual2 = compute_stringency_index(stringency_values2)\n",
    "expected2 = np.array([[0, 1], [0, 7]])\n",
    "assert(np.all(actual2 == expected2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nqtOvepLARkQ"
   },
   "source": [
    "## More Indexing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "F5OtHQYwAU_q"
   },
   "source": [
    "In fact, indexing in NumPy can be more sophisticated than what we have seen previously! \n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZMxvRCGCAZe9"
   },
   "source": [
    "### Integer Array Indexing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "17ngFFwGAddt"
   },
   "source": [
    "We can use a Python list or a NumPy array to index a NumPy array. For example,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "te4V8yif_6JG"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "a = np.array([[1, 2, 3], [4, 5, 6]])\n",
    "print(a[:, [0, 2]])                   # Prints [[1 3]\n",
    "                                      #         [4 6]]\n",
    "print(a[[1, 0], [1, 2]])              # Prints [5 3]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7J5iATyJAffY"
   },
   "source": [
    "Note that `a[[1, 0], [1, 2]]` essentially returns `a[1, 1]` and  `a[0, 2]`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PO3JUERDAp-V"
   },
   "source": [
    "### Boolean Array Indexing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wofAGm5VAq84"
   },
   "source": [
    "An array of Boolean can be used to determine which elements of a given array should be left out. For instance,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RPG5cp5EAlpr"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "a = np.array([4, 3, 1, 5, 10])\n",
    "desired_indices = a > 3\n",
    "print(a[desired_indices])        # Selects values in `a` that are\n",
    "                                 # greater than 3; prints [ 4  5 10]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction to Matplotlib\n",
    "\n",
    "Matplotlib is a powerful and widely-used Python library for creating visualisations. We can use it to generate a wide variety of plots, charts, and graphs. We will make use of this library often in future problem sets.\n",
    "\n",
    "Below is an example of how to use Matplotlib to plot two different functions in the same graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 39884 characters]",
      "text/plain": [
       "<Figure size 800x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# Generate x values from -5 to 5\n",
    "x = np.linspace(-5, 5, 400)\n",
    "\n",
    "# Define two functions\n",
    "y1 = x**2\n",
    "y2 = np.sin(x)\n",
    "\n",
    "# Create a new figure and axis\n",
    "plt.figure(figsize=(8, 6))\n",
    "plt.title('Plotting Two Functions')\n",
    "plt.xlabel('x-axis')\n",
    "plt.ylabel('y-axis')\n",
    "\n",
    "# Plot the first function in blue\n",
    "plt.plot(x, y1, label='y = x^2', color='blue')\n",
    "\n",
    "# Plot the second function in red\n",
    "plt.plot(x, y2, label='y = sin(x)', color='red')\n",
    "\n",
    "# Add a legend\n",
    "plt.legend()\n",
    "\n",
    "# Show the plot\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CdeS5eOoAuvF"
   },
   "source": [
    "## Task 2.6: Average Daily Increase in Cases"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6eOmPKKwAxaL"
   },
   "source": [
    "In this task, we need to implement `average_increase_in_cases`. This function takes `n_cases_increase` and `n_adj_entries_avg` as arguments, and **averages the increase in cases for each day using data from the previous and next `n_adj_entries_avg` number of days**.\n",
    "\n",
    "Similar to task 2.3, `n_cases_increase` is the output obtained from calling `compute_increase_in_cases(n_cases_cumulative)`.\n",
    "\n",
    "In this case, you can assume that **`n_adj_entries_avg` is a positive integer**.\n",
    "\n",
    "For example, if we find that the daily increase in cases for **one country** is given by `a = np.array([0, 5, 10, 15, 20, 25, 30])` and `n_adj_entries_avg = 2`, then\n",
    "- the average increase on day 3 is $(0 + 5 + 10 + 15 + 20) / 5 = 10$\n",
    "- the average increase on day 4 is $(5 + 10 + 15 + 20 + 25) / 5 = 15$\n",
    "- the average increase on day 5 is $(10 + 15 + 20 + 25 + 30) / 5 = 20$\n",
    "- it is not possible to compute the average for the other days since there will not be sufficient values to use for computing the mean\n",
    "\n",
    "Therefore, for **this** country, the average daily increase should be `np.array([np.nan, np.nan, 10, 15, 20, np.nan, np.nan])`.\n",
    "\n",
    "In other words, the average increase in cases for a particular country on the $(j + 1)$-th day is given by the mean of the daily increase in cases over the interval `[-n_adj_entries_avg + j, n_adj_entries_avg + j]` (Note: this interval includes the endpoints).\n",
    "\n",
    "The return value should be the mean increase in cases for each day -- using data from the previous and next `n_adj_entries_avg` number of days -- represented as a 2D `np.ndarray` such that the $(i, j)$ entry represents the average increase in daily cases on the $(j + 1)$-th day in the $i$th country, **rounded down to the smallest integer**.\n",
    "\n",
    "**Note:** In this task, you may use some of the Numpy functions that you have learned earlier, such as `numpy.mean` and `numpy.floor`.\n",
    "Your implementation should not involve any iteration, including `map` and `filter`, recursion, or any iterative approaches like for-loops.In addition, you may benefit from using the following additional functions:\n",
    "\n",
    "* [`numpy.lib.stride_tricks.sliding_window_view`](https://numpy.org/devdocs/reference/generated/numpy.lib.stride_tricks.sliding_window_view.html): Create a sliding window view into the array with the given window shape.\n",
    "* [`numpy.ndarray.fill`](https://numpy.org/doc/stable/reference/generated/numpy.ndarray.fill.html): Fill the array with a scalar value.\n",
    "\n",
    "The basic usage of these functions is given below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[1 2 3]\n",
      " [2 3 4]\n",
      " [3 4 5]\n",
      " [4 5 6]]\n",
      "[0 0 3 4 5 6]\n"
     ]
    }
   ],
   "source": [
    "a = np.array([1, 2, 3, 4, 5, 6])\n",
    "print(np.lib.stride_tricks.sliding_window_view(a, 3)) # Create a sliding window of length 3\n",
    "'''\n",
    "[[1, 2, 3],\n",
    " [2, 3, 4],\n",
    " [3, 4, 5],\n",
    " [4, 5, 6]]\n",
    "'''\n",
    "\n",
    "a[:2].fill(0) # Fill the first two elements of the array a with 0\n",
    "\n",
    "print(a) # array([0, 0, 3, 4, 5, 6])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-33YHZWUAw7A"
   },
   "outputs": [],
   "source": [
    "def average_increase_in_cases(n_cases_increase, n_adj_entries_avg=7):\n",
    "    '''\n",
    "    Averages the increase in cases for each day using data from the previous\n",
    "    `n_adj_entries_avg` number of days and the next `n_adj_entries_avg` number\n",
    "    of days.\n",
    "    Parameters\n",
    "    ----------\n",
    "    n_cases_increase: np.ndarray\n",
    "        2D `ndarray` with each row representing the data of a country, and the columns\n",
    "        of each row representing the time series data of the daily increase in the\n",
    "        number of confirmed cases in that country, i.e. the ith row of \n",
    "        `n_cases_increase` contains the data of the ith country, and the (i, j) entry of\n",
    "        `n_cases_increase` is the daily increase in the number of confirmed cases on the\n",
    "        (j + 1)th day in the ith country.\n",
    "    n_adj_entries_avg: int\n",
    "        Number of days from which data will be used to compute the average increase\n",
    "        in cases. This should be a positive integer.\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    Mean increase in cases for each day, using data from the previous\n",
    "    `n_adj_entries_avg` number of days and the next `n_adj_entries_avg` number\n",
    "    of days, as a 2D `ndarray` such that the (i, j) entry represents the\n",
    "    average increase in daily cases on the (j + 1)th day in the ith country,\n",
    "    rounded down to the smallest integer.\n",
    "    \n",
    "    The average increase in cases for a particular country on the (j + 1)th day\n",
    "    is given by the mean of the daily increase in cases over the interval\n",
    "    [-`n_adj_entries_avg` + j, `n_adj_entries_avg` + j]. (Note: this interval\n",
    "    includes the endpoints).\n",
    "    Note\n",
    "    ----\n",
    "    Since this computation requires data from the previous `n_adj_entries_avg`\n",
    "    number of days and the next `n_adj_entries_avg` number of days, it is not\n",
    "    possible to compute the average for the first and last `n_adj_entries_avg`\n",
    "    number of days. Therefore, set the average increase in cases for these days\n",
    "    to `np.nan` for all countries.\n",
    "    Your implementation should not involve any iteration, including `map` and `filter`, \n",
    "    recursion, or any iterative approaches like for-loops.\n",
    "    '''\n",
    "    \n",
    "    # TODO: add your solution here and remove `raise NotImplementedError`\n",
    "    raise NotImplementedError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test case for Task 2.6\n",
    "n_cases_increase = np.array([[0, 5, 10, 15, 20, 25, 30]])\n",
    "actual = average_increase_in_cases(n_cases_increase, n_adj_entries_avg=2)\n",
    "expected = np.array([[np.nan, np.nan, 10, 15, 20, np.nan, np.nan]])\n",
    "assert(np.array_equal(actual, expected, equal_nan=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualising the Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To view the results of averaging, you can use the helper function `visualise_increase`. This function utilises the `Matplotlib` library to display data of different coutries in the same graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def visualise_increase(n_cases_increase, n_cases_increase_avg=None):\n",
    "    '''\n",
    "    Visualises the increase in cases for each country that is represented in\n",
    "    `n_cases_increase`. If `n_cases_increase_avg` is passed into the\n",
    "    function as well, visualisation will also be done for the average increase in\n",
    "    cases for each country.\n",
    "\n",
    "    NOTE: If more than 5 countries are represented, only the plots for the first 5\n",
    "    countries will be shown.\n",
    "    '''\n",
    "    days = np.arange(1, n_cases_increase.shape[1] + 1)  # Our x axis will be \"days\"\n",
    "    plt.figure() # Start a new graph\n",
    "    for i in range(min(5, n_cases_increase.shape[0])):   # A curve for each row (country)\n",
    "        plt.plot(days, n_cases_increase[i, :], label='country {}'.format(i))\n",
    "    plt.legend()\n",
    "    plt.title('Increase in Cases')\n",
    "\n",
    "    if n_cases_increase_avg is None:\n",
    "        plt.show()\n",
    "        return\n",
    "    \n",
    "    plt.figure() # Start a new graph     \n",
    "    for i in range(min(5, n_cases_increase_avg.shape[0])): # A curve for each row (country)\n",
    "        plt.plot(days, n_cases_increase_avg[i, :], label='country {}'.format(i))\n",
    "    plt.legend()\n",
    "    plt.title('Average Increase in Cases')\n",
    "    plt.show() # Show all graphs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To help you understand how this helper function works, the following is an example of using the helper function `visualise_increase` on a dummy data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_cases_increase = np.array([[0, 2, 5, 3, 11, 9, 12, 1, 15, 30], [20, 12, 1, 7, 12, 9, 9, 28, 4, 16]])\n",
    "visualise_increase(n_cases_increase, average_increase_in_cases(n_cases_increase, n_adj_entries_avg=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FblvT7rbAxyh"
   },
   "source": [
    "The results we obtained when using `n_cases_top_cumulative` to compute `n_cases_increase` and `n_cases_increase_avg`, with the default value of 7 for `n_adj_entries_avg`, are shown in two figures below respectively.\n",
    "\n",
    "<img src = 'imgs/daily_increase.png'>\n",
    "\n",
    "<img src = 'imgs/daily_increase_avg.png'>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pU2uVb4RIObo"
   },
   "source": [
    "## Task 2.7: Finding Peaks in Daily Increase"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Wo2FuuQoISZu"
   },
   "source": [
    "Our last task is to implement `is_peak`. This function takes `n_cases_increase_avg` and `n_adj_entries_peak` as arguments, and **determines whether the $(j + 1)$-th day was a day when the increase in cases peaked in the $i$th country**.\n",
    "\n",
    "In this case, `n_cases_increase_avg` is given by `average_increase_in_cases(n_cases_increase, n_adj_entries_avg)`.\n",
    "\n",
    "Suppose `a` is the average daily increase in cases, with the $(i, j)$ entry indicating the **average increase** in cases on the $(j + 1)$-th day in the $i$-th country. Moreover, let `n_adj_entries_peak` be denoted by `m`.\n",
    "\n",
    "In addition, an increase on the $(j + 1)$-th day is deemed *significant* in the $i$-th country if `a[i, j]` is greater than 10 percent of the mean of all **average daily increases** in the country. `np.nan` is, as the name suggests, not a number and cannot contribute to the mean. Therefore, if we have the set $\\{2, 3, `np.nan`\\}$, the mean should be $(2 + 3) / 2 = 2.5$.\n",
    "\n",
    "Now, to determine whether there is a *peak* on the $(j + 1)$-th day in the $i$-th country, check whether `a[i, j]` is *maximum* in`(a[i, j - m], a[i, j - m + 1], ..., a[i, j + m - 1], a[i, j + m])`. \n",
    "\n",
    "In the event that there are non-unique values such that there exists some value `a[i, k] = a[i, j]`, where $k \\neq j$ and $j - m \\leq k \\leq j + m$, we consider `a[i, j]` a maximum if and only if $k > j$. For example, in $(5, 10, 10)$, 10 is considered a maximum (since the second 10 appears to the right of the value 10 in the middle); but in $(10, 10, 5)$, 10 is not considered a maximum (since there exists another 10 to the left of the value in the middle).\n",
    "\n",
    "Then, **if `a[i, j]` is maximum and significant, then there is a peak on the $(j + 1)$-th day in the ith country; otherwise, there is no peak.**\n",
    "\n",
    "In this case, the return value of this function should be a 2D `np.ndarray` with the $(i, j)$ entry indicating whether there is a peak in the daily increase in cases on the $(j + 1)$-th day in the $i$th country.\n",
    "\n",
    "While imperfect, this approach gives a pretty decent result especially with the default value of 7 for `n_adj_entries_avg` and `n_adj_entries_peak`, as shown in figure below.\n",
    "\n",
    "<img src = 'imgs/peaks.png'>\n",
    "\n",
    "**Note:** In this task, you may use some of the Numpy functions that you have learned earlier, including `np.lib.stride_tricks.sliding_window_view`. In addition, you may benefit from using the following additional functions:\n",
    "- [`numpy.count_nonzero`](https://numpy.org/doc/stable/reference/generated/numpy.count_nonzero.html): This function (recursively) counts how many elements in an array (and in sub-arrays thereof) that have non-zero values.\n",
    "- [`numpy.isnan`](https://numpy.org/doc/stable/reference/generated/numpy.isnan.html): Test element-wise for NaN and return result as a boolean array.\n",
    "- [`numpy.nanmean`](https://numpy.org/doc/stable/reference/generated/numpy.nanmean.html): Compute the arithmetic mean along the specified axis, ignoring NaNs. Returns the average of the array elements. The average is taken over the flattened array by default, otherwise over the specified axis. \n",
    "\n",
    "Your implementation should not involve any iteration, including `map` and `filter`, recursion, or any iterative approaches like for-loops.\n",
    "\n",
    "The basic usage of these functions is given below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "a = np.array([[0, 1, 7, 0], [3, 0, 2, 19]])\n",
    "print(np.count_nonzero(a)) # 5\n",
    "\n",
    "a = np.array([1, 2, np.nan])\n",
    "print(np.isnan(a)) # [False, False, True]\n",
    "\n",
    "a = np.array([[1, np.nan], [3, 5]])\n",
    "print(np.nanmean(a)) # 3.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2noFaxZLIB1I"
   },
   "outputs": [],
   "source": [
    "def is_peak(n_cases_increase_avg, n_adj_entries_peak=7):\n",
    "    '''\n",
    "    Determines whether the (j + 1)th day was a day when the increase in cases\n",
    "    peaked in the ith country.\n",
    "    Parameters\n",
    "    ----------\n",
    "    n_cases_increase_avg: np.ndarray\n",
    "        2D `ndarray` with each row representing the data of a country, and the columns\n",
    "        of each row representing the time series data of the average daily increase in the\n",
    "        number of confirmed cases in that country, i.e. the ith row of \n",
    "        `n_cases_increase` contains the data of the ith country, and the (i, j) entry of\n",
    "        `n_cases_increase` is the average daily increase in the number of confirmed\n",
    "        cases on the (j + 1)th day in the ith country. In this case, the 'average'\n",
    "        is computed using the output from `average_increase_in_cases`.\n",
    "    n_adj_entries_peak: int\n",
    "        Number of days that determines the size of the window in which peaks are\n",
    "        to be detected. \n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    2D `ndarray` with the (i, j) entry indicating whether there is a peak in the\n",
    "    daily increase in cases on the (j + 1)th day in the ith country.\n",
    "    Suppose `a` is the average daily increase in cases, with the (i, j) entry\n",
    "    indicating the average increase in cases on the (j + 1)th day in the ith\n",
    "    country. Moreover, let `n_adj_entries_peak` be denoted by `m`.\n",
    "    In addition, an increase on the (j + 1)th day is deemed significant in the\n",
    "    ith country if `a[i, j]` is greater than 10 percent of the mean of all\n",
    "    average daily increases in the country.\n",
    "    Now, to determine whether there is a peak on the (j + 1)th day in the ith\n",
    "    country, check whether `a[i, j]` is maximum in {`a[i, j - m]`, `a[i, j - m + 1]`,\n",
    "    ..., `a[i, j + m - 1]`, `a[i, j + m]`}. If it is and `a[i, j]` is significant,\n",
    "    then there is a peak on the (j + 1)th day in the ith country; otherwise,\n",
    "    there is no peak.\n",
    "    Note\n",
    "    ----\n",
    "    Let d = `n_adj_entries_avg` + `n_adj_entries_peak`, where `n_adj_entries_avg`\n",
    "    is that used to compute `n_cases_increase_avg`. Observe that it is not\n",
    "    possible to detect a peak in the first and last d days, i.e. these days should\n",
    "    not be peaks.\n",
    "    \n",
    "    As described in `average_increase_in_cases`, to compute the average daily\n",
    "    increase, we need data from the previous and the next `n_adj_entries_avg`\n",
    "    number of days. Hence, we won't have an average for these days, precluding\n",
    "    the computation of peaks during the first and last `n_adj_entries_avg` days.\n",
    "    Moreover, similar to `average_increase_in_cases`, we need the data over the\n",
    "    interval [-`n_adj_entries_peak` + j, `n_adj_entries_peak` + j] to determine\n",
    "    whether the (j + 1)th day is a peak.\n",
    "    Hint: to determine `n_adj_entries_avg` from `n_cases_increase_avg`,\n",
    "    `np.count_nonzero` and `np.isnan` may be helpful.\n",
    "\n",
    "    Your implementation should not involve any iteration, including `map` and `filter`, \n",
    "    recursion, or any iterative approaches like for-loops.\n",
    "    '''\n",
    "\n",
    "    # TODO: add your solution here and remove `raise NotImplementedError`\n",
    "    raise NotImplementedError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test case for Task 2.7\n",
    "\n",
    "n_cases_increase_avg = np.array([[np.nan, np.nan, 10, 10, 5, 20, 7, np.nan, np.nan], [np.nan, np.nan, 15, 5, 16, 17, 17, np.nan, np.nan]])\n",
    "n_adj_entries_peak = 1\n",
    "\n",
    "actual = is_peak(n_cases_increase_avg, n_adj_entries_peak=n_adj_entries_peak)\n",
    "expected = np.array([[False, False, False, False, False, True, False, False, False],\n",
    "                     [False, False, False, False, False, True, False, False, False]])\n",
    "assert np.all(actual == expected)\n",
    "\n",
    "n_cases_increase_avg2 = np.array([[np.nan, np.nan, 10, 20, 20, 20, 20, np.nan, np.nan], [np.nan, np.nan, 20, 20, 20, 20, 10, np.nan, np.nan]])\n",
    "n_adj_entries_peak2 = 1\n",
    "\n",
    "actual2 = is_peak(n_cases_increase_avg2, n_adj_entries_peak=n_adj_entries_peak2)\n",
    "expected2 = np.array([[False, False, False, True, False, False, False, False, False],\n",
    "                    [False, False, False, False, False, False, False, False, False]])\n",
    "assert np.all(actual2 == expected2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Explanation of the example test case\n",
    "\n",
    "Below is the process on how the function should compute the above example test case.\n",
    "\n",
    "Consider:\n",
    "- `n_cases_increase_avg = np.array([[np.nan, np.nan, 10, 10, 5, 20, 7, np.nan, np.nan], [np.nan, np.nan, 15, 5, 16, 17, 17, np.nan, np.nan]])`\n",
    "- `n_adj_entries_peak = 1`\n",
    "\n",
    "Let's get the solution for this instance of `is_peak` function parameter values.\n",
    "\n",
    "1. Work with only the non-nan portion of the n_cases_increase_avg array -> [[10, 10, 5, 20, 7], [15, 5, 16, 17, 17]]\n",
    "2. First, we check for max for each window of size 3 (1 left, middle, 1 right) -> 1 is the value of `n_adj_entries_peak`\n",
    "    * We have to look at the following windows for country 1:\n",
    "      1. [10, 10, 5] for 10: False as there is an equal max element to its left in the window\n",
    "      2. [10, 5, 20] for 5: False as it is not max\n",
    "      3. [5, 20,  7] for 20: True as it is max\n",
    "      \n",
    "      So, for country 1, the max checking boolean result is: [False, False, True]\n",
    "\n",
    "    * We have to look at the following windows for country 2: \n",
    "      1. [15, 5, 16] for 5: False as it is not max\n",
    "      2. [5, 16, 17] for 16: False as it is not max\n",
    "      3. [16, 17,  17] for 17: True as it is max and the equal element is on its right side in the window. \n",
    "      \n",
    "      So, for country 2, the max checking boolean result is: [False, False, True]\n",
    "      \n",
    "    * So, final result for max checking is: [[False, False, True], [False, False, True]]\n",
    "\n",
    "\n",
    "3. Now, we check for significance -> whether an element of interest is greater than 10% of the mean of avg daily increases in a country\n",
    "    - For country 1,\n",
    "      - The threshold value is: ((10 + 10 + 5 + 20 + 7) / 5) * 0.1 = 1.04\n",
    "      - 10, 5 and 20 -> all are greater than this value. \n",
    "      - So, the significance checking boolean result for country 1 is: [True, True, True]\n",
    "    - For country 2,\n",
    "      - The threshold value is: ((15 + 5 + 16 + 17 + 17) / 5) * 0.1 = 1.40\n",
    "      - 5, 16, 17 -> all are greater than this value. \n",
    "      - So, the significance checking boolean result for country 2 is: [True, True, True]\n",
    "    \n",
    "    - So, final result for significance checking is: [[True, True, True], [True, True, True]]\n",
    "\n",
    "4. The condition for being a peak is to satisfy both the max and the significance criteria.\n",
    "So, we now have \"and\" the final resultant boolean arrays for max and significance checking that we derived a while ago\n",
    "The result we get after the operation is: [[False, False, True], [False, False, True]]\n",
    "\n",
    "5. So, is it done? \n",
    "Not yet. Our final result should have the same shape as the n_cases_increase_avg array, which in this\n",
    "case is 2 X 9. But the result we have currently is of shape 2 X 3.\n",
    "So, for each of the 2 rows, we have to pad with 6 False values; 3 on the left and 3 on the right which results in:\n",
    "[[False, False, False, False, False, True, False, False, False], \n",
    "    [False, False, False, False, False, True, False, False, False]]\n",
    "  \n",
    "6. At last! We got our final answer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualising the peaks\n",
    "\n",
    "The peaks in daily increase can be highlighted on the graph using our helper function `visualise_peaks`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def visualise_peaks(n_cases_increase_avg, peaks):\n",
    "    '''\n",
    "    Visualises peaks for each of the country that is represented in\n",
    "    `n_cases_increase_avg` according to variable `peaks`.\n",
    "    \n",
    "    NOTE: If there are more than 5 countries, only the plots for the first 5\n",
    "    countries will be shown.\n",
    "    '''\n",
    "    days = np.arange(1, n_cases_increase_avg.shape[1] + 1) # Days will be our x-coordinates\n",
    "\n",
    "    plt.figure() # Start a graph\n",
    "    \n",
    "    for i in range(min(5, n_cases_increase_avg.shape[0])): # A curve for each row (country) \n",
    "        plt.plot(days, n_cases_increase_avg[i, :], label='country {}'.format(i)) # Plot the daily increase curve\n",
    "        peak = (np.nonzero(peaks[i, :]))[0]\n",
    "        peak_days = peak + 1 # since data starts from day 1, not 0\n",
    "        plt.scatter(peak_days, n_cases_increase_avg[i, peak]) # Scatterplot of peak(s) that lay on top of the curve\n",
    "    \n",
    "    plt.legend()\n",
    "    plt.show() # Display graph"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To view the results of your solution, you can use `visualise_peaks` as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualise the results on the test case\n",
    "visualise_peaks(n_cases_increase_avg, is_peak(n_cases_increase_avg, n_adj_entries_peak=n_adj_entries_peak))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gv3IzLY8ISfY"
   },
   "source": [
    "# Submission"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kIPn-5PNJ8k5"
   },
   "source": [
    "Once you are done, please submit your work to Coursemology, by copying the right snippets of code into the corresponding box that says 'Your answer', and click 'Save'.  After you save, you can make changes to your\n",
    "submission.\n",
    "\n",
    "Once you are satisfied with what you have uploaded, click 'Finalize submission.'  **Note that once your submission is finalized, it is considered to be submitted for grading and cannot be changed**. If you need to undo\n",
    "this action, you will have to email your assigned tutor for help. Please do not finalize your submission until you are sure that you want to submit your solutions for grading. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Appendix"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below is a non-exhaustive list of NumPy functions allowed and now allowed: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Allowed functions:\n",
    "\n",
    "numpy.all\n",
    "numpy.any\n",
    "numpy.append\n",
    "numpy.arange\n",
    "numpy.argmax\n",
    "numpy.compress\n",
    "numpy.concatenate\n",
    "numpy.cumsum\n",
    "numpy.diff\n",
    "numpy.divide\n",
    "numpy.einsum\n",
    "numpy.equal\n",
    "numpy.full\n",
    "numpy.greater\n",
    "numpy.greater_equal\n",
    "numpy.hstack\n",
    "numpy.insert\n",
    "numpy.lib.stride_tricks\n",
    "numpy.logical_and\n",
    "numpy.logical_not\n",
    "numpy.ndarray.flatten\n",
    "numpy.ndarray.reshape \n",
    "numpy.newaxis\n",
    "numpy.ones\n",
    "numpy.pad\n",
    "numpy.shape\n",
    "numpy.squeeze\n",
    "numpy.take\n",
    "np.ufunc.reduce\n",
    "numpy.where\n",
    "numpy.zeros\n",
    "numpy.zeros_like\n",
    "\n",
    "# Plus all functions mentioned in NumPy Basics\n",
    "\n",
    "# Not allowed:\n",
    "\n",
    "numpy.apply_along_axis\n",
    "numpy.apply_over_axes\n",
    "numpy.vectorize"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "ps0.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "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.9"
  },
  "vscode": {
   "interpreter": {
    "hash": "b1bf6cc60825bc0168f0daef984b080cea2a9fe0c964c898af2495b9f96ac9e2"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
