Lab_1_Introduction.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyOG0OtidY1MTXvhDw1rDY+J",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/USCbiostats/PM520/blob/main/Lab_1_Introduction.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# PM520, Lab 1: Introduction to Python\n",
        "Much of this content is based on the excellent resource [Learn Python in Y Minutes](https://learnxinyminutes.com/docs/python/).\n",
        "\n",
        "## 1. Primitive Datatypes and Operators\n",
        "Python basics and built-in datatypes"
      ],
      "metadata": {
        "id": "clhG00dyC4_b"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "YQi-hEOdBIlg",
        "outputId": "6a206f43-912a-40fd-b6a8-61a89ab2e686"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "metadata": {},
          "execution_count": 32
        }
      ],
      "source": [
        "# Single line comments start with a number symbol.\n",
        "\n",
        "\"\"\" Multiline strings can be written\n",
        "    using three \"s, and are often used\n",
        "    as documentation.\n",
        "\"\"\"\n",
        "\n",
        "\n",
        "\n",
        "# You have numbers\n",
        "3  # => 3\n",
        "\n",
        "# Math is what you would expect\n",
        "1 + 1   # => 2\n",
        "8 - 1   # => 7\n",
        "10 * 2  # => 20\n",
        "35 / 5  # => 7.0\n",
        "\n",
        "# Integer division rounds down for both positive and negative numbers.\n",
        "5 // 3       # => 1\n",
        "-5 // 3      # => -2\n",
        "5.0 // 3.0   # => 1.0 # works on floats too\n",
        "-5.0 // 3.0  # => -2.0\n",
        "\n",
        "# The result of division is always a float\n",
        "10.0 / 3  # => 3.3333333333333335\n",
        "\n",
        "# Modulo operation\n",
        "7 % 3   # => 1\n",
        "# i % j have the same sign as j, unlike C\n",
        "-7 % 3  # => 2\n",
        "\n",
        "# Exponentiation (x**y, x to the yth power)\n",
        "2**3  # => 8\n",
        "\n",
        "# Enforce precedence with parentheses\n",
        "1 + 3 * 2    # => 7\n",
        "(1 + 3) * 2  # => 8\n",
        "\n",
        "# Boolean values are primitives (Note: the capitalization)\n",
        "True   # => True\n",
        "False  # => False\n",
        "\n",
        "# negate with not\n",
        "not True   # => False\n",
        "not False  # => True\n",
        "\n",
        "# Boolean Operators\n",
        "# Note \"and\" and \"or\" are case-sensitive\n",
        "True and False  # => False\n",
        "False or True   # => True\n",
        "\n",
        "# True and False are actually 1 and 0 but with different keywords\n",
        "True + True # => 2\n",
        "True * 8    # => 8\n",
        "False - 5   # => -5\n",
        "\n",
        "# Comparison operators look at the numerical value of True and False\n",
        "0 == False  # => True\n",
        "2 > True    # => True\n",
        "2 == True   # => False\n",
        "-5 != False # => True\n",
        "\n",
        "# None, 0, and empty strings/lists/dicts/tuples/sets all evaluate to False.\n",
        "# All other values are True\n",
        "bool(0)     # => False\n",
        "bool(\"\")    # => False\n",
        "bool([])    # => False\n",
        "bool({})    # => False\n",
        "bool(())    # => False\n",
        "bool(set()) # => False\n",
        "bool(4)     # => True\n",
        "bool(-6)    # => True\n",
        "\n",
        "# Using boolean logical operators on ints casts them to booleans for evaluation,\n",
        "# but their non-cast value is returned. Don't mix up with bool(ints) and bitwise\n",
        "# and/or (&,|)\n",
        "bool(0)     # => False\n",
        "bool(2)     # => True\n",
        "0 and 2     # => 0\n",
        "bool(-5)    # => True\n",
        "bool(2)     # => True\n",
        "-5 or 0     # => -5\n",
        "\n",
        "# Equality is ==\n",
        "1 == 1  # => True\n",
        "2 == 1  # => False\n",
        "\n",
        "# Inequality is !=\n",
        "1 != 1  # => False\n",
        "2 != 1  # => True\n",
        "\n",
        "# More comparisons\n",
        "1 < 10  # => True\n",
        "1 > 10  # => False\n",
        "2 <= 2  # => True\n",
        "2 >= 2  # => True\n",
        "\n",
        "# Seeing whether a value is in a range\n",
        "1 < 2 and 2 < 3  # => True\n",
        "2 < 3 and 3 < 2  # => False\n",
        "# Chaining makes this look nicer\n",
        "1 < 2 < 3  # => True\n",
        "2 < 3 < 2  # => False\n",
        "\n",
        "# (is vs. ==) is checks if two variables refer to the same object, but == checks\n",
        "# if the objects pointed to have the same values.\n",
        "a = [1, 2, 3, 4]  # Point a at a new list, [1, 2, 3, 4]\n",
        "b = a             # Point b at what a is pointing to\n",
        "b is a            # => True, a and b refer to the same object\n",
        "b == a            # => True, a's and b's objects are equal\n",
        "b = [1, 2, 3, 4]  # Point b at a new list, [1, 2, 3, 4]\n",
        "b is a            # => False, a and b do not refer to the same object\n",
        "b == a            # => True, a's and b's objects are equal\n",
        "\n",
        "# Strings are created with \" or '\n",
        "\"This is a string.\"\n",
        "'This is also a string.'\n",
        "\n",
        "# Strings can be added too\n",
        "\"Hello \" + \"world!\"  # => \"Hello world!\"\n",
        "# String literals (but not variables) can be concatenated without using '+'\n",
        "\"Hello \" \"world!\"    # => \"Hello world!\"\n",
        "\n",
        "# A string can be treated like a list of characters\n",
        "\"Hello world!\"[0]  # => 'H'\n",
        "\n",
        "# You can find the length of a string\n",
        "len(\"This is a string\")  # => 16\n",
        "\n",
        "# Since Python 3.6, you can use f-strings or formatted string literals.\n",
        "name = \"Reiko\"\n",
        "f\"She said her name is {name}.\" # => \"She said her name is Reiko\"\n",
        "# Any valid Python expression inside these braces is returned to the string.\n",
        "f\"{name} is {len(name)} characters long.\" # => \"Reiko is 5 characters long.\"\n",
        "\n",
        "# None is an object\n",
        "None  # => None\n",
        "\n",
        "# Don't use the equality \"==\" symbol to compare objects to None\n",
        "# Use \"is\" instead. This checks for equality of object identity.\n",
        "\"etc\" is None  # => False\n",
        "None is None   # => True\n",
        "x = None\n",
        "x is None"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 2. Variables and Collections\n",
        "Storing state in variables and collection datatypes"
      ],
      "metadata": {
        "id": "3HLKfEs8DUL8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Python has a print function\n",
        "print(\"I'm Python. Nice to meet you!\")  # => I'm Python. Nice to meet you!\n",
        "\n",
        "# By default the print function also prints out a newline at the end.\n",
        "# Use the optional argument end to change the end string.\n",
        "print(\"Hello, World\", end=\"!\")  # => Hello, World!\n",
        "\n",
        "# Accessing a previously unassigned variable is an exception.\n",
        "# See Control Flow to learn more about exception handling.\n",
        "some_unknown_var  # Raises a NameError\n",
        "\n",
        "# if can be used as an expression\n",
        "# Equivalent of C's '?:' ternary operator\n",
        "some_var = \"yay!\" if 0 < 1 else \"nay!\"  # => \"nay!\"\n",
        "\n",
        "# Lists store sequences\n",
        "li = []\n",
        "# You can start with a prefilled list\n",
        "other_li = [4, 5, 6]\n",
        "\n",
        "# Add stuff to the end of a list with append\n",
        "li.append(1)    # li is now [1]\n",
        "li.append(2)    # li is now [1, 2]\n",
        "li.append(4)    # li is now [1, 2, 4]\n",
        "li.append(3)    # li is now [1, 2, 4, 3]\n",
        "\n",
        "# Access a list like you would any array\n",
        "li[0]   # => 1\n",
        "# Look at the last element\n",
        "li[-1]  # => 3\n",
        "\n",
        "# Looking out of bounds is an IndexError\n",
        "li[4]  # Raises an IndexError\n",
        "\n",
        "# You can look at ranges with slice syntax.\n",
        "# The start index is included, the end index is not\n",
        "# (It's a closed/open range for you mathy types.)\n",
        "li[1:3]   # Return list from index 1 to 3 => [2, 4]\n",
        "li[2:]    # Return list starting from index 2 => [4, 3]\n",
        "li[:3]    # Return list from beginning until index 3  => [1, 2, 4]\n",
        "li[::2]   # Return list selecting elements with a step size of 2 => [1, 4]\n",
        "li[::-1]  # Return list in reverse order => [3, 4, 2, 1]\n",
        "# Use any combination of these to make advanced slices\n",
        "# li[start:end:step]\n",
        "\n",
        "# Make a one layer deep copy using slices\n",
        "li2 = li[:]  # => li2 = [1, 2, 4, 3] but (li2 is li) will result in false.\n",
        "\n",
        "# Check for existence in a list with \"in\"\n",
        "1 in li  # => True\n",
        "\n",
        "# Examine the length with \"len()\"\n",
        "len(li)  # => 4\n",
        "\n",
        "# Tuples are like lists but are immutable.\n",
        "tup = (1, 2, 3)\n",
        "tup[0]      # => 1\n",
        "tup[0] = 3  # Raises a TypeError\n",
        "\n",
        "# Note that a tuple of length one has to have a comma after the last element but\n",
        "# tuples of other lengths, even zero, do not.\n",
        "type((1))   # => <class 'int'>\n",
        "type((1,))  # => <class 'tuple'>\n",
        "type(())    # => <class 'tuple'>\n",
        "\n",
        "# You can do most of the list operations on tuples too\n",
        "len(tup)         # => 3\n",
        "tup + (4, 5, 6)  # => (1, 2, 3, 4, 5, 6)\n",
        "tup[:2]          # => (1, 2)\n",
        "2 in tup         # => True\n",
        "\n",
        "# You can unpack tuples (or lists) into variables\n",
        "a, b, c = (1, 2, 3)  # a is now 1, b is now 2 and c is now 3\n",
        "# You can also do extended unpacking\n",
        "a, *b, c = (1, 2, 3, 4)  # a is now 1, b is now [2, 3] and c is now 4\n",
        "# Tuples are created by default if you leave out the parentheses\n",
        "d, e, f = 4, 5, 6  # tuple 4, 5, 6 is unpacked into variables d, e and f\n",
        "# respectively such that d = 4, e = 5 and f = 6\n",
        "# Now look how easy it is to swap two values\n",
        "e, d = d, e  # d is now 5 and e is now 4"
      ],
      "metadata": {
        "id": "NxxDIb3gB5TF",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "33e04e6c-b9e8-4411-e253-696be9e1cfb8"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[2, 3]"
            ]
          },
          "metadata": {},
          "execution_count": 63
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 3. Control Flow and Iterables"
      ],
      "metadata": {
        "id": "33L5xvv7D01i"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Let's just make a variable\n",
        "some_var = 10\n",
        "\n",
        "# Here is an if statement. Indentation is significant in Python!\n",
        "# Convention is to use four spaces, not tabs.\n",
        "# This prints \"some_var is smaller than 10\"\n",
        "if some_var > 10:\n",
        "    print(\"some_var is totally bigger than 10.\")\n",
        "elif some_var < 10:    # This elif clause is optional.\n",
        "    print(\"some_var is smaller than 10.\")\n",
        "else:                  # This is optional too.\n",
        "    print(\"some_var is indeed 10.\")\n",
        "\n",
        "\n",
        "\"\"\"\n",
        "For loops iterate over lists\n",
        "prints:\n",
        "    dog is a mammal\n",
        "    cat is a mammal\n",
        "    mouse is a mammal\n",
        "\"\"\"\n",
        "for animal in [\"dog\", \"cat\", \"mouse\"]:\n",
        "    # You can use format() to interpolate formatted strings\n",
        "    print(f\"{animal} is a mammal\")\n",
        "\n",
        "\"\"\"\n",
        "\"range(number)\" returns an iterable of numbers\n",
        "from zero up to (but excluding) the given number\n",
        "prints:\n",
        "    we're at index=0\n",
        "    we're at index=1\n",
        "    we're at index=2\n",
        "    we're at index=3\n",
        "\"\"\"\n",
        "for i in range(4):\n",
        "    print(f\"we're at index={i}\")\n",
        "\n",
        "\"\"\"\n",
        "\"range(lower, upper)\" returns an iterable of numbers\n",
        "from the lower number to the upper number\n",
        "prints:\n",
        "    we're at index=4\n",
        "    we're at index=5\n",
        "    we're at index=6\n",
        "    we're at index=7\n",
        "\"\"\"\n",
        "for i in range(4, 8):\n",
        "    print(f\"we're at index={i}\")\n",
        "\n",
        "\"\"\"\n",
        "\"range(lower, upper, step)\" returns an iterable of numbers\n",
        "from the lower number to the upper number, while incrementing\n",
        "by step. If step is not indicated, the default value is 1.\n",
        "prints:\n",
        "    4\n",
        "    6\n",
        "\"\"\"\n",
        "for i in range(4, 8, 2):\n",
        "    print(i)\n",
        "\n",
        "\"\"\"\n",
        "Loop over a list to retrieve both the index and the value of each list item:\n",
        "    0 dog\n",
        "    1 cat\n",
        "    2 mouse\n",
        "\"\"\"\n",
        "animals = [\"dog\", \"cat\", \"mouse\"]\n",
        "for i, value in enumerate(animals):\n",
        "    print(i, value)\n",
        "\n",
        "\"\"\"\n",
        "While loops go until a condition is no longer met.\n",
        "prints:\n",
        "    0\n",
        "    1\n",
        "    2\n",
        "    3\n",
        "\"\"\"\n",
        "x = 0\n",
        "while x < 4:\n",
        "    print(x)\n",
        "    x += 1  # Shorthand for x = x + 1\n",
        "\n",
        "# Handle exceptions with a try/except block\n",
        "try:\n",
        "    # Use \"raise\" to raise an error\n",
        "    raise IndexError(\"This is an index error\")\n",
        "except IndexError as e:\n",
        "    pass                 # Refrain from this, provide a recovery (next example).\n",
        "except (TypeError, NameError):\n",
        "    pass                 # Multiple exceptions can be processed jointly.\n",
        "else:                    # Optional clause to the try/except block. Must follow\n",
        "                         # all except blocks.\n",
        "    print(\"All good!\")   # Runs only if the code in try raises no exceptions\n",
        "finally:                 # Execute under all circumstances\n",
        "    print(\"We can clean up resources here\")"
      ],
      "metadata": {
        "id": "gS10hhFWCJlQ",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "2c45ec0f-e036-47eb-bb61-69f7a1a7fa8c"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "We can clean up resources here\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 4. Functions"
      ],
      "metadata": {
        "id": "ORZll6wOEFry"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Use \"def\" to create new functions\n",
        "def add(x, y):\n",
        "    print(f\"x is {x} and y is {y}\")\n",
        "    return x + y  # Return values with a return statement\n",
        "\n",
        "# Calling functions with parameters\n",
        "add(5, 6)  # => prints out \"x is 5 and y is 6\" and returns 11\n",
        "\n",
        "# Another way to call functions is with keyword arguments\n",
        "add(y=6, x=5)  # Keyword arguments can arrive in any order.\n",
        "\n",
        "\n",
        "# Returning multiple values (with tuple assignments)\n",
        "def swap(x, y):\n",
        "    return y, x  # Return multiple values as a tuple without the parenthesis.\n",
        "                 # (Note: parenthesis have been excluded but can be included)\n",
        "\n",
        "x = 1\n",
        "y = 2\n",
        "x, y = swap(x, y)     # => x = 2, y = 1\n",
        "# (x, y) = swap(x,y)  # Again the use of parenthesis is optional.\n",
        "\n",
        "# global scope\n",
        "x = 5\n",
        "\n",
        "def set_x(num):\n",
        "    # local scope begins here\n",
        "    # local var x not the same as global var x\n",
        "    x = num    # => 43\n",
        "    print(x)   # => 43\n",
        "\n",
        "def set_global_x(num):\n",
        "    # global indicates that particular var lives in the global scope\n",
        "    global x\n",
        "    print(x)   # => 5\n",
        "    x = num    # global var x is now set to 6\n",
        "    print(x)   # => 6\n",
        "\n",
        "set_x(43)\n",
        "set_global_x(6)\n",
        "\"\"\"\n",
        "prints:\n",
        "    43\n",
        "    5\n",
        "    6\n",
        "\"\"\"\n",
        "\n",
        "# Python has first class functions\n",
        "def create_adder(x):\n",
        "    def adder(y):\n",
        "        return x + y\n",
        "    return adder\n",
        "\n",
        "add_10 = create_adder(10)\n",
        "add_10(3)   # => 13\n",
        "\n",
        "# There are also anonymous functions\n",
        "# !!! note: no 'return' key used !!!\n",
        "(lambda x: x > 2)(3)                  # => True\n",
        "(lambda x, y: x ** 2 + y ** 2)(2, 1)  # => 5\n",
        "\n",
        "# you can store lambdas in variables\n",
        "is_greater_than_two = lambda x: x > 2\n",
        "is_greater_than_two(3)\n",
        "\n",
        "# There are built-in higher order functions\n",
        "list(map(add_10, [1, 2, 3]))          # => [11, 12, 13]\n",
        "list(map(max, [1, 2, 3], [4, 2, 1]))  # => [4, 2, 3]\n",
        "\n",
        "list(filter(lambda x: x > 5, [3, 4, 5, 6, 7]))  # => [6, 7]\n",
        "\n",
        "# We can use list comprehensions for nice maps and filters\n",
        "# List comprehension stores the output as a list (which itself may be nested).\n",
        "[add_10(i) for i in [1, 2, 3]]         # => [11, 12, 13]\n",
        "[x for x in [3, 4, 5, 6, 7] if x > 5]  # => [6, 7]\n",
        "\n",
        "# You can construct set and dict comprehensions as well.\n",
        "{x for x in 'abcddeef' if x not in 'abc'}  # => {'d', 'e', 'f'}\n",
        "{x: x**2 for x in range(5)}  # => {0: 0, 1: 1, 2: 4, 3: 9, 4: 16}"
      ],
      "metadata": {
        "id": "2i6qB6ARC0mQ",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "b883801f-469a-4c92-b3d3-ba0c9ce716f0"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[11, 12, 13]"
            ]
          },
          "metadata": {},
          "execution_count": 89
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 5. Modules"
      ],
      "metadata": {
        "id": "ucWYq3g2ErPn"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# You can import modules\n",
        "import math\n",
        "print(math.sqrt(16))  # => 4.0\n",
        "\n",
        "# You can get specific functions from a module\n",
        "from math import ceil, floor\n",
        "print(ceil(3.7))   # => 4\n",
        "print(floor(3.7))  # => 3\n",
        "\n",
        "# You can import all functions from a module.\n",
        "# Warning: this is not recommended\n",
        "from math import *\n",
        "\n",
        "# You can shorten module names\n",
        "import math as m\n",
        "math.sqrt(16) == m.sqrt(16)  # => True\n",
        "\n",
        "# Python modules are just ordinary Python files. You\n",
        "# can write your own, and import them. The name of the\n",
        "# module is the same as the name of the file.\n",
        "\n",
        "# You can find out which functions and attributes\n",
        "# are defined in a module.\n",
        "import math\n",
        "dir(math)\n",
        "\n",
        "# If you have a Python script named math.py in the same\n",
        "# folder as your current script, the file math.py will\n",
        "# be loaded instead of the built-in Python module.\n",
        "# This happens because the local folder has priority\n",
        "# over Python's built-in libraries.\n",
        "\n",
        "\n",
        "# Decorators are a form of syntactic sugar.\n",
        "# They make code easier to read while accomplishing clunky syntax.\n",
        "\n",
        "# Wrappers are one type of decorator.\n",
        "# They're really useful for adding logging to existing functions without needing to modify them.\n",
        "from functools import wraps\n",
        "\n",
        "def log_function(func):\n",
        "    @wraps(func) # this ensures docstring, function name, arguments list, etc. are all copied\n",
        "                 # to the wrapped function - instead of being replaced with wrapper's info\n",
        "    def wrapper(*args, **kwargs):\n",
        "        print(\"Entering function\", func.__name__)\n",
        "        result = func(*args, **kwargs)\n",
        "        print(\"Exiting function\", func.__name__)\n",
        "        return result\n",
        "    return wrapper\n",
        "\n",
        "@log_function\n",
        "def my_function(x, y):\n",
        "    return x+y\n",
        "\n",
        "my_function(1,2) # => \"Entering function my_function\"\n",
        "                 # => \"3\"\n",
        "                 # => \"Exiting function my_function\"\n",
        "\n",
        "print(my_function.__name__) # => 'my_function'\n",
        "print(my_function.__code__.co_argcount) # => 2"
      ],
      "metadata": {
        "id": "thk9dDVzEpIj",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "44d04b5d-a2ec-41ae-a2a7-afacf4c4b44a"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Entering function my_function\n",
            "Exiting function my_function\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "3"
            ]
          },
          "metadata": {},
          "execution_count": 93
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# JAX and JAX.Numpy\n",
        "[JAX](https://github.com/google/jax) is a Google-backed library to enable automatic differentiation of Python code, while supporting ultra-fast runtime due to \"Just-In-Time\" (i.e. JIT) compilation from their custom bytecode (i.e. XLA). Hence JAX = JIT + AutoDiff + XLA.\n",
        "\n",
        "However, before we can use JAX we'll need to install it. To do that, we'll use the Python package management tool `pip`. To call `pip` in colab (or any terminal function) we need to prepend the command with a `!`.\n"
      ],
      "metadata": {
        "id": "iz7dK2s8GKcx"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install jax"
      ],
      "metadata": {
        "id": "CJQIOPlDGNcZ",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "d5f74a1c-7d07-4be1-c781-9309b3d34737"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: jax in /usr/local/lib/python3.10/dist-packages (0.4.23)\n",
            "Requirement already satisfied: ml-dtypes>=0.2.0 in /usr/local/lib/python3.10/dist-packages (from jax) (0.2.0)\n",
            "Requirement already satisfied: numpy>=1.22 in /usr/local/lib/python3.10/dist-packages (from jax) (1.23.5)\n",
            "Requirement already satisfied: opt-einsum in /usr/local/lib/python3.10/dist-packages (from jax) (3.3.0)\n",
            "Requirement already satisfied: scipy>=1.9 in /usr/local/lib/python3.10/dist-packages (from jax) (1.11.4)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's practice importing JAX and using the numpy implementation backed by JAX. numpy is a Python library for n-dimensional arrays. Here we are using JAX's implementation, which will enable us to take advantage of all of JAX's features."
      ],
      "metadata": {
        "id": "zzehFq5gIn4R"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import jax\n",
        "import jax.numpy as jnp\n",
        "import jax.random as rdm\n",
        "\n",
        "# readr::read_tsv\n",
        "# let's practice some numpy tricks\n",
        "x = jnp.arange(9)\n",
        "y = jnp.ones(9)\n",
        "print(f\"x = {x} | y = {y}\")\n",
        "\n",
        "z = x + y\n",
        "print(f\"z = {z} | x + 1 = {x + 1}\")\n",
        "\n",
        "P = 4\n",
        "i = jnp.eye(P)\n",
        "a = 2 * jnp.ones(P)\n",
        "print(f\"i = {i} | a = {a}\")\n",
        "\n",
        "# is this mat/vec mult?\n",
        "b = i * a\n",
        "print(f\"b = {b}\")\n",
        "\n",
        "A = jnp.array([[5., 1], [1, 5]])\n",
        "a = 2 * jnp.ones(2)\n",
        "print(f\"A = {A}\")\n",
        "b = A * a\n",
        "print(f\"b = {b}\")\n",
        "\n",
        "# nope! b is matrix; mat/vec mult => vec\n",
        "b = A @ a\n",
        "print(f\"b = {b}\")\n",
        "b = jnp.dot(A, a)\n",
        "print(f\"b = {b}\")"
      ],
      "metadata": {
        "id": "gUOdJ5sFI4jv",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "701bfcd7-b5e8-4dd9-eb59-73f2680e0402"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "x = [0 1 2 3 4 5 6 7 8] | y = [1. 1. 1. 1. 1. 1. 1. 1. 1.]\n",
            "z = [1. 2. 3. 4. 5. 6. 7. 8. 9.] | x + 1 = [1 2 3 4 5 6 7 8 9]\n",
            "i = [[1. 0. 0. 0.]\n",
            " [0. 1. 0. 0.]\n",
            " [0. 0. 1. 0.]\n",
            " [0. 0. 0. 1.]] | a = [2. 2. 2. 2.]\n",
            "b = [[2. 0. 0. 0.]\n",
            " [0. 2. 0. 0.]\n",
            " [0. 0. 2. 0.]\n",
            " [0. 0. 0. 2.]]\n",
            "A = [[5. 1.]\n",
            " [1. 5.]]\n",
            "b = [[10.  2.]\n",
            " [ 2. 10.]]\n",
            "b = [12. 12.]\n",
            "b = [12. 12.]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# indexing arrays\n",
        "X = jnp.arange(9).reshape((3,3))\n",
        "\n",
        "# what is the shape of x?\n",
        "print(f\"shape(x) = {X.shape}\")\n",
        "\n",
        "# indexing and 'slicing'\n",
        "print(f\"1st row of x {X[0]}\")\n",
        "print(f\"1st row of x {X[0,:]}\")\n",
        "print(f\"1st col of x {X[:,0]}\")\n",
        "\n",
        "# arrays can be n-dimensional and not just vectors/matrices!\n",
        "X = jnp.arange(27).reshape((3,3,3))\n",
        "print(f\"shape(x) = {X.shape}\")\n",
        "print(f\"1st matrix of x {X[0]}\")\n",
        "print(f\"1st matrix of x {X[0,:]}\")\n",
        "print(f\"1st row of each matrix of x {X[:,0,:]}\")"
      ],
      "metadata": {
        "id": "jThdMTR8ME7p",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "5a8e38a2-0d4b-4e28-f73d-e7d1bf216752"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "shape(x) = (3, 3)\n",
            "1st row of x [0 1 2]\n",
            "1st row of x [0 1 2]\n",
            "1st col of x [0 3 6]\n",
            "shape(x) = (3, 3, 3)\n",
            "1st matrix of x [[0 1 2]\n",
            " [3 4 5]\n",
            " [6 7 8]]\n",
            "1st matrix of x [[0 1 2]\n",
            " [3 4 5]\n",
            " [6 7 8]]\n",
            "1st row of each matrix of x [[ 0  1  2]\n",
            " [ 9 10 11]\n",
            " [18 19 20]]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# JIT warm up\n",
        "def my_func(x):\n",
        "  return jnp.sum(x ** 2)\n",
        "\n",
        "# `jax.jit` takes as input a function and returns the JIT-compiled function\n",
        "my_func_jit = jax.jit(my_func)\n",
        "\n",
        "# results should be the same\n",
        "is_same = jnp.allclose(my_func(jnp.ones(4)), my_func_jit(jnp.ones(4)))\n",
        "print(f\"Results are same? {is_same}\")\n",
        "\n",
        "%timeit my_func(jnp.ones(4)) # let's measure time\n",
        "%timeit my_func_jit(jnp.ones(4)).block_until_ready() # measure using JIT; need to block until result is returned\n",
        "\n",
        "# results computed faster in the JIT compiled function! We did no extra work\n",
        "# except wrap our function using a JAX command! Now let's see how to\n",
        "# use the decorator sytax to handle that automatically for us\n",
        "\n",
        "@jax.jit\n",
        "def my_new_func(x):\n",
        "  return jnp.sum(x ** 2)\n",
        "\n",
        "# the @jax.jit above the function definition informs the Python interpreter\n",
        "# to \"decorate\" `my_new_func` with the `jax.jit` function, which will automatically\n",
        "# wrap my_new_func in the JIT compiled version. That is, anytime we call `my_new_func`\n",
        "# we're actually calling the same thing as `jax.jit(my_new_func)`\n",
        "%timeit my_new_func(jnp.ones(4)).block_until_ready()\n",
        "\n",
        "# the average time is similar to the above `my_func_jit` which shows that we're\n",
        "# calling the JIT'd version."
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "E-eExdElehhL",
        "outputId": "bae58e3d-dcd5-447c-9ae0-2f131cdf1192"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Results are same? True\n",
            "687 µs ± 171 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n",
            "286 µs ± 18.9 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n",
            "279 µs ± 22 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
          ]
        }
      ]
    }
  ]
}