Lab_0_NumericsCheatSheet.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyPyFolE+2ueFaT5zqY/S/jY",
      "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_0_NumericsCheatSheet.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# That's a nice calculation you got there. Be a real shame if something happened to it, or: Numerics Cheat Sheet\n",
        "This document is meant to illustrate how the numerics of common calculations can fail, and how \"built-in\" functions that replace the direct operations can improve numerical stability.\n",
        "\n",
        "As usual for this course, we will be operating in [JAX](https://jax.readthedocs.io/en/latest/) using its `numpy` and `scipy` submodules."
      ],
      "metadata": {
        "id": "bP0KUR2Ao9T7"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vF7NzB4_iPc0"
      },
      "outputs": [],
      "source": [
        "# let's import our necessary modules before getting started\n",
        "\n",
        "import jax\n",
        "import jax.numpy as jnp\n",
        "import jax.scipy as jsp\n",
        "import jax.scipy.special as jspec"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## log(1 + p) and exp(p) - 1\n",
        "These functions typically arise in the calculation of [cross-entropy](https://en.wikipedia.org/wiki/Cross-entropy#Cross-entropy_loss_function_and_logistic_regression) for a [binary classifier](https://en.wikipedia.org/wiki/Binary_classification) or the [log-likelihood](https://en.wikipedia.org/wiki/Likelihood_function) for [logistic regression](https://en.wikipedia.org/wiki/Logistic_regression) under a [Bernoulli](https://en.wikipedia.org/wiki/Bernoulli_distribution) model.\n",
        "\n",
        "[log1p](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.log1p.html) calculates `jnp.log(1 + p)` in a numerically stable manner. Calulating this directly can result in catastrophic roundoff when `p` is very small and produce `0.0` when a positive value should exist.\n",
        "\n",
        "Similarly, [expm1](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.expm1.html) calculates `jnp.exp(x) - 1` in a numerically stable manner and direct calculation can produce `0.0` when `p` is very small."
      ],
      "metadata": {
        "id": "WvJjkP81inHX"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "for pow in jnp.arange(10, 0, -1):\n",
        "  p = 10. ** (-pow)\n",
        "  safe_res = jnp.log1p(p)\n",
        "  unsafe_res = jnp.log(1. + p)\n",
        "  print(f\"log(1 + {p:.3e}) | Safe = {safe_res:.5e} | Unsafe = {unsafe_res:.5e}\")\n",
        "\n",
        "print(\"-\"*62)\n",
        "for pow in jnp.arange(10, 0, -1):\n",
        "  p = 10. ** (-pow)\n",
        "  safe_res = jnp.expm1(p)\n",
        "  unsafe_res = jnp.exp(p) - 1.\n",
        "  print(f\"exp({p:.3e}) - 1 | Safe = {safe_res:.5e} | Unsafe = {unsafe_res:.5e}\")"
      ],
      "metadata": {
        "id": "gHb0K8k1ioy8",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "a4d1deb1-2094-413b-d660-6040d03cbbec"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "log(1 + 1.000e-10) | Safe = 1.00000e-10 | Unsafe = 0.00000e+00\n",
            "log(1 + 1.000e-09) | Safe = 1.00000e-09 | Unsafe = 0.00000e+00\n",
            "log(1 + 1.000e-08) | Safe = 1.00000e-08 | Unsafe = 0.00000e+00\n",
            "log(1 + 1.000e-07) | Safe = 1.00000e-07 | Unsafe = 1.19209e-07\n",
            "log(1 + 1.000e-06) | Safe = 1.00000e-06 | Unsafe = 9.53674e-07\n",
            "log(1 + 1.000e-05) | Safe = 9.99995e-06 | Unsafe = 1.00135e-05\n",
            "log(1 + 1.000e-04) | Safe = 9.99950e-05 | Unsafe = 1.00012e-04\n",
            "log(1 + 1.000e-03) | Safe = 9.99500e-04 | Unsafe = 9.99547e-04\n",
            "log(1 + 1.000e-02) | Safe = 9.95033e-03 | Unsafe = 9.95032e-03\n",
            "log(1 + 1.000e-01) | Safe = 9.53102e-02 | Unsafe = 9.53102e-02\n",
            "--------------------------------------------------------------\n",
            "exp(1.000e-10) - 1 | Safe = 1.00000e-10 | Unsafe = 0.00000e+00\n",
            "exp(1.000e-09) - 1 | Safe = 1.00000e-09 | Unsafe = 0.00000e+00\n",
            "exp(1.000e-08) - 1 | Safe = 1.00000e-08 | Unsafe = 0.00000e+00\n",
            "exp(1.000e-07) - 1 | Safe = 1.00000e-07 | Unsafe = 1.19209e-07\n",
            "exp(1.000e-06) - 1 | Safe = 1.00000e-06 | Unsafe = 9.53674e-07\n",
            "exp(1.000e-05) - 1 | Safe = 1.00000e-05 | Unsafe = 1.00136e-05\n",
            "exp(1.000e-04) - 1 | Safe = 1.00005e-04 | Unsafe = 1.00017e-04\n",
            "exp(1.000e-03) - 1 | Safe = 1.00050e-03 | Unsafe = 1.00052e-03\n",
            "exp(1.000e-02) - 1 | Safe = 1.00502e-02 | Unsafe = 1.00502e-02\n",
            "exp(1.000e-01) - 1 | Safe = 1.05171e-01 | Unsafe = 1.05171e-01\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [],
      "metadata": {
        "id": "ecb2Gg2eSnzv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "for pow in jnp.arange(10, 0, -1):\n",
        "  p = 10. ** (-pow)\n",
        "  safe_res = jnp.log1p(p)\n",
        "  unsafe_res = jnp.log(1. + p)\n",
        "  print(f\"log(1 + {p:.3e}) | Safe = {safe_res:.5e} | Unsafe = {unsafe_res:.5e}\")"
      ],
      "metadata": {
        "id": "PzHdm3OVSoiU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "xlogy and xlog1py"
      ],
      "metadata": {
        "id": "MT6Hf0YWiuYu"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "jspec.xlogy\n",
        "jspec.xlog1py"
      ],
      "metadata": {
        "id": "s1B37Lyui1dW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "logsumexp, logaddexp, and softmax"
      ],
      "metadata": {
        "id": "JYuJo_vujRYb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "jspec.logsumexp\n",
        "jnp.logaddexp"
      ],
      "metadata": {
        "id": "Q0ryX8NsjSrQ"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}