tutorials/3d_nonlinear_regression.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fdb3af1b",
   "metadata": {},
   "source": [
    "## 3D Nonlinear Regression"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c1f057c",
   "metadata": {},
   "source": [
    "**Copyright (c) Meta Platforms, Inc. and affiliates.**\n",
    "\n",
    "This source code is licensed under the MIT license found in the\n",
    "LICENSE file in the root directory of this source tree."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2514d65",
   "metadata": {},
   "source": [
    "In this notebook let's use DiffKt to fit a nonlinear regression to a dataset of two input variables $ x_1 $ and $ x_2 $ and output variable $ y $. Bring in the DiffKt library for the tensor library and automatic differentiation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d1b23168",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "@file:DependsOn(\"../kotlin/api/build/libs/api.jar\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be712fe4",
   "metadata": {},
   "source": [
    "The dataset can be found [here](https://bit.ly/35ReT3i) and here it is visualized as a scatterplot below. \n",
    "\n",
    "![](./resources/ybgAGvOQXT.mp4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bec502f3",
   "metadata": {},
   "source": [
    "We see a 3D parabola-like shape above, and we are going to fit the following function:\n",
    "\n",
    "$$\n",
    "y = ax_1^2 + bx_2^2 + c \n",
    "$$\n",
    "\n",
    "$ x_1 $ and $ x_2 $ are the input variables, $ y $ is the output variable, and $ a $, $ b $, and $ c $ are the coefficients we will use gradient descent to solve for. \n",
    "\n",
    "First let's `import` three libraries: `Random`, `URL`,ad `diffkt`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0dc563c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import java.net.URL\n",
    "import kotlin.random.Random\n",
    "import org.diffkt.*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c7e0722",
   "metadata": {},
   "source": [
    "Next let's declare a `Point` class that will hold the two input variables $ x_1 $ and $ x_2 $ and the output variable $ y $. We will then use the `URL` function to read the CSV stored [here](https://bit.ly/3ty5BRZ), split the lines using a regular expression, and process each line as a `Sequence`. Note we drop the first line with column names and split the comma-separated values, and then convert them into `Float` values. Then we can package each trio of values into a `Point` and collect into a `List<Point>`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "befb1079",
   "metadata": {},
   "outputs": [],
   "source": [
    "data class Point(val x1: Float, val x2: Float, val y: Float)\n",
    "\n",
    "val points = URL(\"https://bit.ly/35ebET5\")    // read CSV\n",
    "    .readText().split(Regex(\"\\\\r?\\\\n\"))       // split lines using regular expression\n",
    "    .filter { it.matches(Regex(\"[-,.0-9]+\")) }  // filter only numeric records using regular expression\n",
    "    .map { it.split(\",\").map{ it.toFloat()} } // split commas into columns\n",
    "    .map { (x1,x2,y) -> Point(x1,x2,y) }      // map to Point objects"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed14aec6",
   "metadata": {},
   "source": [
    "We are going to need to map these points to DiffKt tensors. We will use the `tensorOf()` function and map the `points` inside it. Now when we map to the inputs `x1` and `x2` and the output tensor `y`, we use lambda functions as arguments to specify what columns we want to generate and on what values. However notice on the `x` tensor below we add a third column simply returning a $ 1 $. This is going to add a column of 1's next to our `x1` and `x2` input variables. Why is this necessary? It will serve as a placeholder to generate the intercept coefficient. Without it, we would only generate the slopes for `x1` and `x2` without any intercept. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3fbcdc45",
   "metadata": {},
   "outputs": [],
   "source": [
    "// map variables to input and output variable tensors\n",
    "// add a placeholder \"1\" column to generate intercept on input tensor\n",
    "val x = tensorOf(points.flatMap { listOf(it.x1, it.x2, 1f) }.map(::FloatScalar) ).reshape(points.size, 3)\n",
    "val y = tensorOf(points.map { it.y }.map(::FloatScalar) ).reshape(points.size, 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f7e410f",
   "metadata": {},
   "source": [
    "To represent our three coefficients $ a $, $ b $, and $ c $ we will use a float tensor holding these three values. Let's initialize them as random values between $ 0 $ and $ 1 $.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4177db84",
   "metadata": {},
   "outputs": [],
   "source": [
    "// initialize coefficients\n",
    "var coeffs: DTensor = FloatTensor.random(Random,Shape(3,1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8fe52503",
   "metadata": {},
   "source": [
    "To visualize the tensor operations, let's say our coefficients $ a $ were initialized with the following values. \n",
    "\n",
    "$ A = \\left[\\begin{matrix}0.1\\\\0.2\\\\0.5\\end{matrix}\\right] $ \n",
    "\n",
    "And let's say we have 3 records of $ X $ inputs with the added column of 1's. \n",
    "\n",
    "$ X = \\left[\\begin{matrix}2 & 10 & 1\\\\4 & 20 & 1\\\\10 & 30 & 1\\end{matrix}\\right] $ \n",
    "\n",
    "To get the predicted $ \\hat{Y} $ values, we apply matrix multiplication (dot products) between the squared input $ X $ variables (with the additional column of 1's) and the coefficients $ A $.  \n",
    "\n",
    "$ \\hat{Y} = X^2 \\cdot A $ \n",
    "\n",
    "$ \\hat{Y} = \\left[\\begin{matrix}2 & 10 & 1\\\\4 & 20 & 1\\\\10 & 30 & 1\\end{matrix}\\right]^2 \\cdot \\left[\\begin{matrix}0.1\\\\0.2\\\\0.5\\end{matrix}\\right] $ \n",
    "\n",
    "$ \\hat{Y} = \\left[\\begin{matrix}4 & 100 & 1\\\\16 & 400 & 1\\\\100 & 900 & 1\\end{matrix}\\right] \\cdot \\left[\\begin{matrix}0.1\\\\0.2\\\\0.5\\end{matrix}\\right] $ \n",
    "\n",
    "$ \\hat{Y} =  \\left[\\begin{matrix}(4 \\times 0.1) + (100 \\times 0.2) + (1 \\times 0.5) \\\\(16 \\times 0.1) + (400 \\times 0.2) + (1 \\times 0.5) \\\\(100 \\times 0.1) + (900 \\times 0.2) + (1 \\times 0.5) \\end{matrix}\\right] $\n",
    "\n",
    "$ \\hat{Y} = \\left[\\begin{matrix}20.9\\\\82.1\\\\190.5\\end{matrix}\\right] $\n",
    "\n",
    "So that would yield predictions of $ 20.9 $, $ 82.1 $, and $ 190.5 $. \n",
    "\n",
    "To get predictions on all data given the current coefficients, use DiffKt's `*` operator: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4a94c1af",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[1.4780331], [0.73806214], [-1.8771539], [0.031777054], [0.00975579], [-0.34466755], [-0.92360955], [0.8195014], [2.3792126], [1.0538011], [-0.036050916], [1.991267], [1.378245], [0.10144332], [-0.74854106], [2.379893], [-1.6779041], [1.4739765], [-1.050663], [0.12192413], [-1.1324165], [-0.49822634], [0.2588521], [-0.53845006], [0.10402456], [-0.67420214], [1.4083693], [0.54015714], [0.5197551], [-1.0881097], [-1.2567005], [0.28458583], [0.4415309], [0.89086825], [0.1377117], [2.2032223], [-1.2418356], [0.12989569], [0.053967], [2.2094414], [0.41328257], [1.9897099], [1.9527221], [-0.3773461], [1.0414197]]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "val yPredictions = x.matmul(coeffs)\n",
    "yPredictions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e52411c",
   "metadata": {},
   "source": [
    "To calculate the total loss, let's use a sum of squared loss. Subtract the actual $ Y $ values from the predicted $ \\hat{Y} $ values. Take those differences, square them, and sum them. \n",
    "\n",
    "$ E = \\sum{(Y - \\hat{Y})^2 } $ \n",
    "\n",
    "Let's say we have these predicted $ \\hat{Y} $ and actual $ Y $ values. \n",
    "\n",
    "$ \\hat{Y} = \\left[\\begin{matrix}20.9\\\\82.1\\\\190.5\\end{matrix}\\right] $\n",
    "\n",
    "$ Y = \\left[\\begin{matrix}21.2\\\\85.3\\\\189.1\\end{matrix}\\right] $\n",
    "\n",
    "Here is how we would calculate the sum of squares. \n",
    "\n",
    "$ E = \\sum{(Y - \\hat{Y})^2 } $ \n",
    "\n",
    "$ E = \\sum{( \\left[\\begin{matrix}21.2\\\\85.3\\\\189.1\\end{matrix}\\right] - \\left[\\begin{matrix}20.9\\\\82.1\\\\190.5\\end{matrix}\\right])^2 } $ \n",
    "\n",
    "$ E = \\sum{(\\left[\\begin{matrix}0.3\\\\3.2\\\\-1.4\\end{matrix}\\right])^2} $ \n",
    "\n",
    "$ E = \\sum{\\left[\\begin{matrix}0.09\\\\10.24\\\\0.0625\\end{matrix}\\right]} $ \n",
    "\n",
    "$ E = 10.3925 $ \n",
    "\n",
    "We can implement this as a `loss()` function in Kotlin using DiffKt as shown below. Remember that the predicted $ \\hat{Y} $ values are the dot products of `x` and the coefficients. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "6c7bd874",
   "metadata": {},
   "outputs": [],
   "source": [
    "// calculate sum of squares of the error with given slope and intercept for a line\n",
    "fun loss(coeffs: DTensor): DScalar =\n",
    "    (y - x.pow(2).matmul(coeffs)).pow(2).sum()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5bde0817",
   "metadata": {},
   "source": [
    "Let's perform gradient descent. For $ 2,000 $ iterations, we will use a learning rate of $ .0001 $ and take the reverse derivative of the `loss()` function with regards to the `coeffs` tensor. This will return the gradient for each $ a $, $ b $, and $ c $ coefficient respectively which we multiply by the learning rate and subtract from the `coeffs` tensor. We subtract because we want to descend on the gradients. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "42e8abb0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "betas=[[0.022516534], [-0.046410125], [-1.1387229]]"
     ]
    }
   ],
   "source": [
    "// The learning rate\n",
    "val lr = .0001F\n",
    "\n",
    "// The number of iterations to perform gradient descent\n",
    "val iterations = 2000\n",
    "\n",
    "// Perform gradient descent\n",
    "for (i in 0..iterations) {\n",
    "\n",
    "    // get gradients for line slope and intercept\n",
    "    val betaGradients = reverseDerivative(coeffs, ::loss)\n",
    "\n",
    "    // update m and b by subtracting the (learning rate) * (slope)\n",
    "    coeffs -= betaGradients * lr\n",
    "}\n",
    "print(\"betas=$coeffs\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "704097bb",
   "metadata": {},
   "source": [
    "Now we have fitted our function to the data! Here is a visualization of the gradient descent in action.\n",
    "\n",
    "\n",
    "![](./resources/nipuzkhOdO.mp4)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Kotlin",
   "language": "kotlin",
   "name": "kotlin"
  },
  "language_info": {
   "codemirror_mode": "text/x-kotlin",
   "file_extension": ".kt",
   "mimetype": "text/x-kotlin",
   "name": "kotlin",
   "nbconvert_exporter": "",
   "pygments_lexer": "kotlin",
   "version": "1.7.0-dev-3303"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}