examples/colab/PRAUC_training.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i4vFDhvn_Bdn"
      },
      "source": [
        "##### Copyright 2018 The TensorFlow Constrained Optimization Authors. All Rights Reserved.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\"); you may not use this file except in compliance with the License. You may obtain a copy of the License at\n",
        "\n",
        "\u003e http://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RpUmH2nk_Bdo"
      },
      "source": [
        "## PR-AUC Maximization\n",
        "In this colab, we'll show how to use the TF Constrained Optimization (TFCO) library to train a model to maximize the *Area Under the Precision-Recall Curve (PR-AUC)*. We'll show how to train the model both with (i) plain TensorFlow (in eager mode), and (ii) with a custom tf.Estimator.\n",
        "\n",
        "We start by importing the relevant modules."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FoYVEXPA_Bdp"
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import os\n",
        "import pandas as pd\n",
        "import shutil\n",
        "from sklearn import metrics\n",
        "from sklearn import model_selection\n",
        "import tensorflow.compat.v2 as tf"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ea1KIDRgC4eq"
      },
      "outputs": [],
      "source": [
        "# Tensorflow constrained optimization library\n",
        "!pip install git+https://github.com/google-research/tensorflow_constrained_optimization\n",
        "import tensorflow_constrained_optimization as tfco"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BLxZD5uH_Bds"
      },
      "source": [
        "## Communities and Crimes\n",
        "\n",
        "We will use the  *Communities and Crimes* dataset from the UCI Machine Learning repository for our illustration. This dataset contains various demographic and racial distribution details (aggregated from census and law enforcement data sources) about different communities in the US, along with the per capita crime rate in each commmunity.\n",
        "\n",
        "\n",
        "We begin by downloading and preprocessing the dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MgoeyhS0_Bds"
      },
      "outputs": [],
      "source": [
        "# List of column names in the dataset.\n",
        "column_names = [\"state\", \"county\", \"community\", \"communityname\", \"fold\", \"population\", \"householdsize\", \"racepctblack\", \"racePctWhite\", \"racePctAsian\", \"racePctHisp\", \"agePct12t21\", \"agePct12t29\", \"agePct16t24\", \"agePct65up\", \"numbUrban\", \"pctUrban\", \"medIncome\", \"pctWWage\", \"pctWFarmSelf\", \"pctWInvInc\", \"pctWSocSec\", \"pctWPubAsst\", \"pctWRetire\", \"medFamInc\", \"perCapInc\", \"whitePerCap\", \"blackPerCap\", \"indianPerCap\", \"AsianPerCap\", \"OtherPerCap\", \"HispPerCap\", \"NumUnderPov\", \"PctPopUnderPov\", \"PctLess9thGrade\", \"PctNotHSGrad\", \"PctBSorMore\", \"PctUnemployed\", \"PctEmploy\", \"PctEmplManu\", \"PctEmplProfServ\", \"PctOccupManu\", \"PctOccupMgmtProf\", \"MalePctDivorce\", \"MalePctNevMarr\", \"FemalePctDiv\", \"TotalPctDiv\", \"PersPerFam\", \"PctFam2Par\", \"PctKids2Par\", \"PctYoungKids2Par\", \"PctTeen2Par\", \"PctWorkMomYoungKids\", \"PctWorkMom\", \"NumIlleg\", \"PctIlleg\", \"NumImmig\", \"PctImmigRecent\", \"PctImmigRec5\", \"PctImmigRec8\", \"PctImmigRec10\", \"PctRecentImmig\", \"PctRecImmig5\", \"PctRecImmig8\", \"PctRecImmig10\", \"PctSpeakEnglOnly\", \"PctNotSpeakEnglWell\", \"PctLargHouseFam\", \"PctLargHouseOccup\", \"PersPerOccupHous\", \"PersPerOwnOccHous\", \"PersPerRentOccHous\", \"PctPersOwnOccup\", \"PctPersDenseHous\", \"PctHousLess3BR\", \"MedNumBR\", \"HousVacant\", \"PctHousOccup\", \"PctHousOwnOcc\", \"PctVacantBoarded\", \"PctVacMore6Mos\", \"MedYrHousBuilt\", \"PctHousNoPhone\", \"PctWOFullPlumb\", \"OwnOccLowQuart\", \"OwnOccMedVal\", \"OwnOccHiQuart\", \"RentLowQ\", \"RentMedian\", \"RentHighQ\", \"MedRent\", \"MedRentPctHousInc\", \"MedOwnCostPctInc\", \"MedOwnCostPctIncNoMtg\", \"NumInShelters\", \"NumStreet\", \"PctForeignBorn\", \"PctBornSameState\", \"PctSameHouse85\", \"PctSameCity85\", \"PctSameState85\", \"LemasSwornFT\", \"LemasSwFTPerPop\", \"LemasSwFTFieldOps\", \"LemasSwFTFieldPerPop\", \"LemasTotalReq\", \"LemasTotReqPerPop\", \"PolicReqPerOffic\", \"PolicPerPop\", \"RacialMatchCommPol\", \"PctPolicWhite\", \"PctPolicBlack\", \"PctPolicHisp\", \"PctPolicAsian\", \"PctPolicMinor\", \"OfficAssgnDrugUnits\", \"NumKindsDrugsSeiz\", \"PolicAveOTWorked\", \"LandArea\", \"PopDens\", \"PctUsePubTrans\", \"PolicCars\", \"PolicOperBudg\", \"LemasPctPolicOnPatr\", \"LemasGangUnitDeploy\", \"LemasPctOfficDrugUn\", \"PolicBudgPerPop\", \"ViolentCrimesPerPop\"]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 246
        },
        "id": "gJ_JcV-V_Bdu",
        "outputId": "f0ac9073-186e-47ca-c8ab-85d8012ec0ce",
        "scrolled": true
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "\u003cdiv\u003e\n",
              "\u003cstyle scoped\u003e\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "\u003c/style\u003e\n",
              "\u003ctable border=\"1\" class=\"dataframe\"\u003e\n",
              "  \u003cthead\u003e\n",
              "    \u003ctr style=\"text-align: right;\"\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003estate\u003c/th\u003e\n",
              "      \u003cth\u003ecounty\u003c/th\u003e\n",
              "      \u003cth\u003ecommunity\u003c/th\u003e\n",
              "      \u003cth\u003ecommunityname\u003c/th\u003e\n",
              "      \u003cth\u003efold\u003c/th\u003e\n",
              "      \u003cth\u003epopulation\u003c/th\u003e\n",
              "      \u003cth\u003ehouseholdsize\u003c/th\u003e\n",
              "      \u003cth\u003eracepctblack\u003c/th\u003e\n",
              "      \u003cth\u003eracePctWhite\u003c/th\u003e\n",
              "      \u003cth\u003eracePctAsian\u003c/th\u003e\n",
              "      \u003cth\u003eracePctHisp\u003c/th\u003e\n",
              "      \u003cth\u003eagePct12t21\u003c/th\u003e\n",
              "      \u003cth\u003eagePct12t29\u003c/th\u003e\n",
              "      \u003cth\u003eagePct16t24\u003c/th\u003e\n",
              "      \u003cth\u003eagePct65up\u003c/th\u003e\n",
              "      \u003cth\u003enumbUrban\u003c/th\u003e\n",
              "      \u003cth\u003epctUrban\u003c/th\u003e\n",
              "      \u003cth\u003emedIncome\u003c/th\u003e\n",
              "      \u003cth\u003epctWWage\u003c/th\u003e\n",
              "      \u003cth\u003epctWFarmSelf\u003c/th\u003e\n",
              "      \u003cth\u003epctWInvInc\u003c/th\u003e\n",
              "      \u003cth\u003epctWSocSec\u003c/th\u003e\n",
              "      \u003cth\u003epctWPubAsst\u003c/th\u003e\n",
              "      \u003cth\u003epctWRetire\u003c/th\u003e\n",
              "      \u003cth\u003emedFamInc\u003c/th\u003e\n",
              "      \u003cth\u003eperCapInc\u003c/th\u003e\n",
              "      \u003cth\u003ewhitePerCap\u003c/th\u003e\n",
              "      \u003cth\u003eblackPerCap\u003c/th\u003e\n",
              "      \u003cth\u003eindianPerCap\u003c/th\u003e\n",
              "      \u003cth\u003eAsianPerCap\u003c/th\u003e\n",
              "      \u003cth\u003eOtherPerCap\u003c/th\u003e\n",
              "      \u003cth\u003eHispPerCap\u003c/th\u003e\n",
              "      \u003cth\u003eNumUnderPov\u003c/th\u003e\n",
              "      \u003cth\u003ePctPopUnderPov\u003c/th\u003e\n",
              "      \u003cth\u003ePctLess9thGrade\u003c/th\u003e\n",
              "      \u003cth\u003ePctNotHSGrad\u003c/th\u003e\n",
              "      \u003cth\u003ePctBSorMore\u003c/th\u003e\n",
              "      \u003cth\u003ePctUnemployed\u003c/th\u003e\n",
              "      \u003cth\u003ePctEmploy\u003c/th\u003e\n",
              "      \u003cth\u003ePctEmplManu\u003c/th\u003e\n",
              "      \u003cth\u003e...\u003c/th\u003e\n",
              "      \u003cth\u003eRentMedian\u003c/th\u003e\n",
              "      \u003cth\u003eRentHighQ\u003c/th\u003e\n",
              "      \u003cth\u003eMedRent\u003c/th\u003e\n",
              "      \u003cth\u003eMedRentPctHousInc\u003c/th\u003e\n",
              "      \u003cth\u003eMedOwnCostPctInc\u003c/th\u003e\n",
              "      \u003cth\u003eMedOwnCostPctIncNoMtg\u003c/th\u003e\n",
              "      \u003cth\u003eNumInShelters\u003c/th\u003e\n",
              "      \u003cth\u003eNumStreet\u003c/th\u003e\n",
              "      \u003cth\u003ePctForeignBorn\u003c/th\u003e\n",
              "      \u003cth\u003ePctBornSameState\u003c/th\u003e\n",
              "      \u003cth\u003ePctSameHouse85\u003c/th\u003e\n",
              "      \u003cth\u003ePctSameCity85\u003c/th\u003e\n",
              "      \u003cth\u003ePctSameState85\u003c/th\u003e\n",
              "      \u003cth\u003eLemasSwornFT\u003c/th\u003e\n",
              "      \u003cth\u003eLemasSwFTPerPop\u003c/th\u003e\n",
              "      \u003cth\u003eLemasSwFTFieldOps\u003c/th\u003e\n",
              "      \u003cth\u003eLemasSwFTFieldPerPop\u003c/th\u003e\n",
              "      \u003cth\u003eLemasTotalReq\u003c/th\u003e\n",
              "      \u003cth\u003eLemasTotReqPerPop\u003c/th\u003e\n",
              "      \u003cth\u003ePolicReqPerOffic\u003c/th\u003e\n",
              "      \u003cth\u003ePolicPerPop\u003c/th\u003e\n",
              "      \u003cth\u003eRacialMatchCommPol\u003c/th\u003e\n",
              "      \u003cth\u003ePctPolicWhite\u003c/th\u003e\n",
              "      \u003cth\u003ePctPolicBlack\u003c/th\u003e\n",
              "      \u003cth\u003ePctPolicHisp\u003c/th\u003e\n",
              "      \u003cth\u003ePctPolicAsian\u003c/th\u003e\n",
              "      \u003cth\u003ePctPolicMinor\u003c/th\u003e\n",
              "      \u003cth\u003eOfficAssgnDrugUnits\u003c/th\u003e\n",
              "      \u003cth\u003eNumKindsDrugsSeiz\u003c/th\u003e\n",
              "      \u003cth\u003ePolicAveOTWorked\u003c/th\u003e\n",
              "      \u003cth\u003eLandArea\u003c/th\u003e\n",
              "      \u003cth\u003ePopDens\u003c/th\u003e\n",
              "      \u003cth\u003ePctUsePubTrans\u003c/th\u003e\n",
              "      \u003cth\u003ePolicCars\u003c/th\u003e\n",
              "      \u003cth\u003ePolicOperBudg\u003c/th\u003e\n",
              "      \u003cth\u003eLemasPctPolicOnPatr\u003c/th\u003e\n",
              "      \u003cth\u003eLemasGangUnitDeploy\u003c/th\u003e\n",
              "      \u003cth\u003eLemasPctOfficDrugUn\u003c/th\u003e\n",
              "      \u003cth\u003ePolicBudgPerPop\u003c/th\u003e\n",
              "      \u003cth\u003eViolentCrimesPerPop\u003c/th\u003e\n",
              "    \u003c/tr\u003e\n",
              "  \u003c/thead\u003e\n",
              "  \u003ctbody\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e0\u003c/th\u003e\n",
              "      \u003ctd\u003e8\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eLakewoodcity\u003c/td\u003e\n",
              "      \u003ctd\u003e1\u003c/td\u003e\n",
              "      \u003ctd\u003e0.19\u003c/td\u003e\n",
              "      \u003ctd\u003e0.33\u003c/td\u003e\n",
              "      \u003ctd\u003e0.02\u003c/td\u003e\n",
              "      \u003ctd\u003e0.90\u003c/td\u003e\n",
              "      \u003ctd\u003e0.12\u003c/td\u003e\n",
              "      \u003ctd\u003e0.17\u003c/td\u003e\n",
              "      \u003ctd\u003e0.34\u003c/td\u003e\n",
              "      \u003ctd\u003e0.47\u003c/td\u003e\n",
              "      \u003ctd\u003e0.29\u003c/td\u003e\n",
              "      \u003ctd\u003e0.32\u003c/td\u003e\n",
              "      \u003ctd\u003e0.20\u003c/td\u003e\n",
              "      \u003ctd\u003e1.0\u003c/td\u003e\n",
              "      \u003ctd\u003e0.37\u003c/td\u003e\n",
              "      \u003ctd\u003e0.72\u003c/td\u003e\n",
              "      \u003ctd\u003e0.34\u003c/td\u003e\n",
              "      \u003ctd\u003e0.60\u003c/td\u003e\n",
              "      \u003ctd\u003e0.29\u003c/td\u003e\n",
              "      \u003ctd\u003e0.15\u003c/td\u003e\n",
              "      \u003ctd\u003e0.43\u003c/td\u003e\n",
              "      \u003ctd\u003e0.39\u003c/td\u003e\n",
              "      \u003ctd\u003e0.40\u003c/td\u003e\n",
              "      \u003ctd\u003e0.39\u003c/td\u003e\n",
              "      \u003ctd\u003e0.32\u003c/td\u003e\n",
              "      \u003ctd\u003e0.27\u003c/td\u003e\n",
              "      \u003ctd\u003e0.27\u003c/td\u003e\n",
              "      \u003ctd\u003e0.36\u003c/td\u003e\n",
              "      \u003ctd\u003e0.41\u003c/td\u003e\n",
              "      \u003ctd\u003e0.08\u003c/td\u003e\n",
              "      \u003ctd\u003e0.19\u003c/td\u003e\n",
              "      \u003ctd\u003e0.10\u003c/td\u003e\n",
              "      \u003ctd\u003e0.18\u003c/td\u003e\n",
              "      \u003ctd\u003e0.48\u003c/td\u003e\n",
              "      \u003ctd\u003e0.27\u003c/td\u003e\n",
              "      \u003ctd\u003e0.68\u003c/td\u003e\n",
              "      \u003ctd\u003e0.23\u003c/td\u003e\n",
              "      \u003ctd\u003e...\u003c/td\u003e\n",
              "      \u003ctd\u003e0.35\u003c/td\u003e\n",
              "      \u003ctd\u003e0.38\u003c/td\u003e\n",
              "      \u003ctd\u003e0.34\u003c/td\u003e\n",
              "      \u003ctd\u003e0.38\u003c/td\u003e\n",
              "      \u003ctd\u003e0.46\u003c/td\u003e\n",
              "      \u003ctd\u003e0.25\u003c/td\u003e\n",
              "      \u003ctd\u003e0.04\u003c/td\u003e\n",
              "      \u003ctd\u003e0.0\u003c/td\u003e\n",
              "      \u003ctd\u003e0.12\u003c/td\u003e\n",
              "      \u003ctd\u003e0.42\u003c/td\u003e\n",
              "      \u003ctd\u003e0.50\u003c/td\u003e\n",
              "      \u003ctd\u003e0.51\u003c/td\u003e\n",
              "      \u003ctd\u003e0.64\u003c/td\u003e\n",
              "      \u003ctd\u003e0.03\u003c/td\u003e\n",
              "      \u003ctd\u003e0.13\u003c/td\u003e\n",
              "      \u003ctd\u003e0.96\u003c/td\u003e\n",
              "      \u003ctd\u003e0.17\u003c/td\u003e\n",
              "      \u003ctd\u003e0.06\u003c/td\u003e\n",
              "      \u003ctd\u003e0.18\u003c/td\u003e\n",
              "      \u003ctd\u003e0.44\u003c/td\u003e\n",
              "      \u003ctd\u003e0.13\u003c/td\u003e\n",
              "      \u003ctd\u003e0.94\u003c/td\u003e\n",
              "      \u003ctd\u003e0.93\u003c/td\u003e\n",
              "      \u003ctd\u003e0.03\u003c/td\u003e\n",
              "      \u003ctd\u003e0.07\u003c/td\u003e\n",
              "      \u003ctd\u003e0.1\u003c/td\u003e\n",
              "      \u003ctd\u003e0.07\u003c/td\u003e\n",
              "      \u003ctd\u003e0.02\u003c/td\u003e\n",
              "      \u003ctd\u003e0.57\u003c/td\u003e\n",
              "      \u003ctd\u003e0.29\u003c/td\u003e\n",
              "      \u003ctd\u003e0.12\u003c/td\u003e\n",
              "      \u003ctd\u003e0.26\u003c/td\u003e\n",
              "      \u003ctd\u003e0.20\u003c/td\u003e\n",
              "      \u003ctd\u003e0.06\u003c/td\u003e\n",
              "      \u003ctd\u003e0.04\u003c/td\u003e\n",
              "      \u003ctd\u003e0.9\u003c/td\u003e\n",
              "      \u003ctd\u003e0.5\u003c/td\u003e\n",
              "      \u003ctd\u003e0.32\u003c/td\u003e\n",
              "      \u003ctd\u003e0.14\u003c/td\u003e\n",
              "      \u003ctd\u003e0.20\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e1\u003c/th\u003e\n",
              "      \u003ctd\u003e53\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eTukwilacity\u003c/td\u003e\n",
              "      \u003ctd\u003e1\u003c/td\u003e\n",
              "      \u003ctd\u003e0.00\u003c/td\u003e\n",
              "      \u003ctd\u003e0.16\u003c/td\u003e\n",
              "      \u003ctd\u003e0.12\u003c/td\u003e\n",
              "      \u003ctd\u003e0.74\u003c/td\u003e\n",
              "      \u003ctd\u003e0.45\u003c/td\u003e\n",
              "      \u003ctd\u003e0.07\u003c/td\u003e\n",
              "      \u003ctd\u003e0.26\u003c/td\u003e\n",
              "      \u003ctd\u003e0.59\u003c/td\u003e\n",
              "      \u003ctd\u003e0.35\u003c/td\u003e\n",
              "      \u003ctd\u003e0.27\u003c/td\u003e\n",
              "      \u003ctd\u003e0.02\u003c/td\u003e\n",
              "      \u003ctd\u003e1.0\u003c/td\u003e\n",
              "      \u003ctd\u003e0.31\u003c/td\u003e\n",
              "      \u003ctd\u003e0.72\u003c/td\u003e\n",
              "      \u003ctd\u003e0.11\u003c/td\u003e\n",
              "      \u003ctd\u003e0.45\u003c/td\u003e\n",
              "      \u003ctd\u003e0.25\u003c/td\u003e\n",
              "      \u003ctd\u003e0.29\u003c/td\u003e\n",
              "      \u003ctd\u003e0.39\u003c/td\u003e\n",
              "      \u003ctd\u003e0.29\u003c/td\u003e\n",
              "      \u003ctd\u003e0.37\u003c/td\u003e\n",
              "      \u003ctd\u003e0.38\u003c/td\u003e\n",
              "      \u003ctd\u003e0.33\u003c/td\u003e\n",
              "      \u003ctd\u003e0.16\u003c/td\u003e\n",
              "      \u003ctd\u003e0.30\u003c/td\u003e\n",
              "      \u003ctd\u003e0.22\u003c/td\u003e\n",
              "      \u003ctd\u003e0.35\u003c/td\u003e\n",
              "      \u003ctd\u003e0.01\u003c/td\u003e\n",
              "      \u003ctd\u003e0.24\u003c/td\u003e\n",
              "      \u003ctd\u003e0.14\u003c/td\u003e\n",
              "      \u003ctd\u003e0.24\u003c/td\u003e\n",
              "      \u003ctd\u003e0.30\u003c/td\u003e\n",
              "      \u003ctd\u003e0.27\u003c/td\u003e\n",
              "      \u003ctd\u003e0.73\u003c/td\u003e\n",
              "      \u003ctd\u003e0.57\u003c/td\u003e\n",
              "      \u003ctd\u003e...\u003c/td\u003e\n",
              "      \u003ctd\u003e0.38\u003c/td\u003e\n",
              "      \u003ctd\u003e0.40\u003c/td\u003e\n",
              "      \u003ctd\u003e0.37\u003c/td\u003e\n",
              "      \u003ctd\u003e0.29\u003c/td\u003e\n",
              "      \u003ctd\u003e0.32\u003c/td\u003e\n",
              "      \u003ctd\u003e0.18\u003c/td\u003e\n",
              "      \u003ctd\u003e0.00\u003c/td\u003e\n",
              "      \u003ctd\u003e0.0\u003c/td\u003e\n",
              "      \u003ctd\u003e0.21\u003c/td\u003e\n",
              "      \u003ctd\u003e0.50\u003c/td\u003e\n",
              "      \u003ctd\u003e0.34\u003c/td\u003e\n",
              "      \u003ctd\u003e0.60\u003c/td\u003e\n",
              "      \u003ctd\u003e0.52\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003e0.02\u003c/td\u003e\n",
              "      \u003ctd\u003e0.12\u003c/td\u003e\n",
              "      \u003ctd\u003e0.45\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003e0.00\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003e0.67\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e2\u003c/th\u003e\n",
              "      \u003ctd\u003e24\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eAberdeentown\u003c/td\u003e\n",
              "      \u003ctd\u003e1\u003c/td\u003e\n",
              "      \u003ctd\u003e0.00\u003c/td\u003e\n",
              "      \u003ctd\u003e0.42\u003c/td\u003e\n",
              "      \u003ctd\u003e0.49\u003c/td\u003e\n",
              "      \u003ctd\u003e0.56\u003c/td\u003e\n",
              "      \u003ctd\u003e0.17\u003c/td\u003e\n",
              "      \u003ctd\u003e0.04\u003c/td\u003e\n",
              "      \u003ctd\u003e0.39\u003c/td\u003e\n",
              "      \u003ctd\u003e0.47\u003c/td\u003e\n",
              "      \u003ctd\u003e0.28\u003c/td\u003e\n",
              "      \u003ctd\u003e0.32\u003c/td\u003e\n",
              "      \u003ctd\u003e0.00\u003c/td\u003e\n",
              "      \u003ctd\u003e0.0\u003c/td\u003e\n",
              "      \u003ctd\u003e0.30\u003c/td\u003e\n",
              "      \u003ctd\u003e0.58\u003c/td\u003e\n",
              "      \u003ctd\u003e0.19\u003c/td\u003e\n",
              "      \u003ctd\u003e0.39\u003c/td\u003e\n",
              "      \u003ctd\u003e0.38\u003c/td\u003e\n",
              "      \u003ctd\u003e0.40\u003c/td\u003e\n",
              "      \u003ctd\u003e0.84\u003c/td\u003e\n",
              "      \u003ctd\u003e0.28\u003c/td\u003e\n",
              "      \u003ctd\u003e0.27\u003c/td\u003e\n",
              "      \u003ctd\u003e0.29\u003c/td\u003e\n",
              "      \u003ctd\u003e0.27\u003c/td\u003e\n",
              "      \u003ctd\u003e0.07\u003c/td\u003e\n",
              "      \u003ctd\u003e0.29\u003c/td\u003e\n",
              "      \u003ctd\u003e0.28\u003c/td\u003e\n",
              "      \u003ctd\u003e0.39\u003c/td\u003e\n",
              "      \u003ctd\u003e0.01\u003c/td\u003e\n",
              "      \u003ctd\u003e0.27\u003c/td\u003e\n",
              "      \u003ctd\u003e0.27\u003c/td\u003e\n",
              "      \u003ctd\u003e0.43\u003c/td\u003e\n",
              "      \u003ctd\u003e0.19\u003c/td\u003e\n",
              "      \u003ctd\u003e0.36\u003c/td\u003e\n",
              "      \u003ctd\u003e0.58\u003c/td\u003e\n",
              "      \u003ctd\u003e0.32\u003c/td\u003e\n",
              "      \u003ctd\u003e...\u003c/td\u003e\n",
              "      \u003ctd\u003e0.29\u003c/td\u003e\n",
              "      \u003ctd\u003e0.27\u003c/td\u003e\n",
              "      \u003ctd\u003e0.31\u003c/td\u003e\n",
              "      \u003ctd\u003e0.48\u003c/td\u003e\n",
              "      \u003ctd\u003e0.39\u003c/td\u003e\n",
              "      \u003ctd\u003e0.28\u003c/td\u003e\n",
              "      \u003ctd\u003e0.00\u003c/td\u003e\n",
              "      \u003ctd\u003e0.0\u003c/td\u003e\n",
              "      \u003ctd\u003e0.14\u003c/td\u003e\n",
              "      \u003ctd\u003e0.49\u003c/td\u003e\n",
              "      \u003ctd\u003e0.54\u003c/td\u003e\n",
              "      \u003ctd\u003e0.67\u003c/td\u003e\n",
              "      \u003ctd\u003e0.56\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003e0.01\u003c/td\u003e\n",
              "      \u003ctd\u003e0.21\u003c/td\u003e\n",
              "      \u003ctd\u003e0.02\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003e0.00\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003e0.43\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e3\u003c/th\u003e\n",
              "      \u003ctd\u003e34\u003c/td\u003e\n",
              "      \u003ctd\u003e5.0\u003c/td\u003e\n",
              "      \u003ctd\u003e81440.0\u003c/td\u003e\n",
              "      \u003ctd\u003eWillingborotownship\u003c/td\u003e\n",
              "      \u003ctd\u003e1\u003c/td\u003e\n",
              "      \u003ctd\u003e0.04\u003c/td\u003e\n",
              "      \u003ctd\u003e0.77\u003c/td\u003e\n",
              "      \u003ctd\u003e1.00\u003c/td\u003e\n",
              "      \u003ctd\u003e0.08\u003c/td\u003e\n",
              "      \u003ctd\u003e0.12\u003c/td\u003e\n",
              "      \u003ctd\u003e0.10\u003c/td\u003e\n",
              "      \u003ctd\u003e0.51\u003c/td\u003e\n",
              "      \u003ctd\u003e0.50\u003c/td\u003e\n",
              "      \u003ctd\u003e0.34\u003c/td\u003e\n",
              "      \u003ctd\u003e0.21\u003c/td\u003e\n",
              "      \u003ctd\u003e0.06\u003c/td\u003e\n",
              "      \u003ctd\u003e1.0\u003c/td\u003e\n",
              "      \u003ctd\u003e0.58\u003c/td\u003e\n",
              "      \u003ctd\u003e0.89\u003c/td\u003e\n",
              "      \u003ctd\u003e0.21\u003c/td\u003e\n",
              "      \u003ctd\u003e0.43\u003c/td\u003e\n",
              "      \u003ctd\u003e0.36\u003c/td\u003e\n",
              "      \u003ctd\u003e0.20\u003c/td\u003e\n",
              "      \u003ctd\u003e0.82\u003c/td\u003e\n",
              "      \u003ctd\u003e0.51\u003c/td\u003e\n",
              "      \u003ctd\u003e0.36\u003c/td\u003e\n",
              "      \u003ctd\u003e0.40\u003c/td\u003e\n",
              "      \u003ctd\u003e0.39\u003c/td\u003e\n",
              "      \u003ctd\u003e0.16\u003c/td\u003e\n",
              "      \u003ctd\u003e0.25\u003c/td\u003e\n",
              "      \u003ctd\u003e0.36\u003c/td\u003e\n",
              "      \u003ctd\u003e0.44\u003c/td\u003e\n",
              "      \u003ctd\u003e0.01\u003c/td\u003e\n",
              "      \u003ctd\u003e0.10\u003c/td\u003e\n",
              "      \u003ctd\u003e0.09\u003c/td\u003e\n",
              "      \u003ctd\u003e0.25\u003c/td\u003e\n",
              "      \u003ctd\u003e0.31\u003c/td\u003e\n",
              "      \u003ctd\u003e0.33\u003c/td\u003e\n",
              "      \u003ctd\u003e0.71\u003c/td\u003e\n",
              "      \u003ctd\u003e0.36\u003c/td\u003e\n",
              "      \u003ctd\u003e...\u003c/td\u003e\n",
              "      \u003ctd\u003e0.70\u003c/td\u003e\n",
              "      \u003ctd\u003e0.77\u003c/td\u003e\n",
              "      \u003ctd\u003e0.89\u003c/td\u003e\n",
              "      \u003ctd\u003e0.63\u003c/td\u003e\n",
              "      \u003ctd\u003e0.51\u003c/td\u003e\n",
              "      \u003ctd\u003e0.47\u003c/td\u003e\n",
              "      \u003ctd\u003e0.00\u003c/td\u003e\n",
              "      \u003ctd\u003e0.0\u003c/td\u003e\n",
              "      \u003ctd\u003e0.19\u003c/td\u003e\n",
              "      \u003ctd\u003e0.30\u003c/td\u003e\n",
              "      \u003ctd\u003e0.73\u003c/td\u003e\n",
              "      \u003ctd\u003e0.64\u003c/td\u003e\n",
              "      \u003ctd\u003e0.65\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003e0.02\u003c/td\u003e\n",
              "      \u003ctd\u003e0.39\u003c/td\u003e\n",
              "      \u003ctd\u003e0.28\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003e0.00\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003e0.12\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e4\u003c/th\u003e\n",
              "      \u003ctd\u003e42\u003c/td\u003e\n",
              "      \u003ctd\u003e95.0\u003c/td\u003e\n",
              "      \u003ctd\u003e6096.0\u003c/td\u003e\n",
              "      \u003ctd\u003eBethlehemtownship\u003c/td\u003e\n",
              "      \u003ctd\u003e1\u003c/td\u003e\n",
              "      \u003ctd\u003e0.01\u003c/td\u003e\n",
              "      \u003ctd\u003e0.55\u003c/td\u003e\n",
              "      \u003ctd\u003e0.02\u003c/td\u003e\n",
              "      \u003ctd\u003e0.95\u003c/td\u003e\n",
              "      \u003ctd\u003e0.09\u003c/td\u003e\n",
              "      \u003ctd\u003e0.05\u003c/td\u003e\n",
              "      \u003ctd\u003e0.38\u003c/td\u003e\n",
              "      \u003ctd\u003e0.38\u003c/td\u003e\n",
              "      \u003ctd\u003e0.23\u003c/td\u003e\n",
              "      \u003ctd\u003e0.36\u003c/td\u003e\n",
              "      \u003ctd\u003e0.02\u003c/td\u003e\n",
              "      \u003ctd\u003e0.9\u003c/td\u003e\n",
              "      \u003ctd\u003e0.50\u003c/td\u003e\n",
              "      \u003ctd\u003e0.72\u003c/td\u003e\n",
              "      \u003ctd\u003e0.16\u003c/td\u003e\n",
              "      \u003ctd\u003e0.68\u003c/td\u003e\n",
              "      \u003ctd\u003e0.44\u003c/td\u003e\n",
              "      \u003ctd\u003e0.11\u003c/td\u003e\n",
              "      \u003ctd\u003e0.71\u003c/td\u003e\n",
              "      \u003ctd\u003e0.46\u003c/td\u003e\n",
              "      \u003ctd\u003e0.43\u003c/td\u003e\n",
              "      \u003ctd\u003e0.41\u003c/td\u003e\n",
              "      \u003ctd\u003e0.28\u003c/td\u003e\n",
              "      \u003ctd\u003e0.00\u003c/td\u003e\n",
              "      \u003ctd\u003e0.74\u003c/td\u003e\n",
              "      \u003ctd\u003e0.51\u003c/td\u003e\n",
              "      \u003ctd\u003e0.48\u003c/td\u003e\n",
              "      \u003ctd\u003e0.00\u003c/td\u003e\n",
              "      \u003ctd\u003e0.06\u003c/td\u003e\n",
              "      \u003ctd\u003e0.25\u003c/td\u003e\n",
              "      \u003ctd\u003e0.30\u003c/td\u003e\n",
              "      \u003ctd\u003e0.33\u003c/td\u003e\n",
              "      \u003ctd\u003e0.12\u003c/td\u003e\n",
              "      \u003ctd\u003e0.65\u003c/td\u003e\n",
              "      \u003ctd\u003e0.67\u003c/td\u003e\n",
              "      \u003ctd\u003e...\u003c/td\u003e\n",
              "      \u003ctd\u003e0.36\u003c/td\u003e\n",
              "      \u003ctd\u003e0.38\u003c/td\u003e\n",
              "      \u003ctd\u003e0.38\u003c/td\u003e\n",
              "      \u003ctd\u003e0.22\u003c/td\u003e\n",
              "      \u003ctd\u003e0.51\u003c/td\u003e\n",
              "      \u003ctd\u003e0.21\u003c/td\u003e\n",
              "      \u003ctd\u003e0.00\u003c/td\u003e\n",
              "      \u003ctd\u003e0.0\u003c/td\u003e\n",
              "      \u003ctd\u003e0.11\u003c/td\u003e\n",
              "      \u003ctd\u003e0.72\u003c/td\u003e\n",
              "      \u003ctd\u003e0.64\u003c/td\u003e\n",
              "      \u003ctd\u003e0.61\u003c/td\u003e\n",
              "      \u003ctd\u003e0.53\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003e0.04\u003c/td\u003e\n",
              "      \u003ctd\u003e0.09\u003c/td\u003e\n",
              "      \u003ctd\u003e0.02\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003e0.00\u003c/td\u003e\n",
              "      \u003ctd\u003eNaN\u003c/td\u003e\n",
              "      \u003ctd\u003e0.03\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "  \u003c/tbody\u003e\n",
              "\u003c/table\u003e\n",
              "\u003cp\u003e5 rows × 128 columns\u003c/p\u003e\n",
              "\u003c/div\u003e"
            ],
            "text/plain": [
              "   state  county  ...  PolicBudgPerPop ViolentCrimesPerPop\n",
              "0      8     NaN  ...             0.14                0.20\n",
              "1     53     NaN  ...              NaN                0.67\n",
              "2     24     NaN  ...              NaN                0.43\n",
              "3     34     5.0  ...              NaN                0.12\n",
              "4     42    95.0  ...              NaN                0.03\n",
              "\n",
              "[5 rows x 128 columns]"
            ]
          },
          "execution_count": 4,
          "metadata": {
            "tags": []
          },
          "output_type": "execute_result"
        }
      ],
      "source": [
        "dataset_url = \"http://archive.ics.uci.edu/ml/machine-learning-databases/communities/communities.data\"\n",
        "\n",
        "# Read dataset from the UCI web repository and assign column names.\n",
        "data_df = pd.read_csv(dataset_url, sep=\",\", names=column_names,\n",
        "                      na_values=\"?\")\n",
        "data_df.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MWHnYDmL_Bdx"
      },
      "source": [
        "The 'ViolentCrimesPerPop' column contains the per capita crime rate for each community. We label the communities with a crime rate above the 70-th percentile as 'high crime' and the others as 'low crime'. These would serve as our binary target labels."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fJtOkt90_Bdy"
      },
      "outputs": [],
      "source": [
        "# Make sure there are no missing values in the \"ViolentCrimesPerPop\" column.\n",
        "assert(not data_df[\"ViolentCrimesPerPop\"].isna().any())\n",
        "\n",
        "# Binarize the \"ViolentCrimesPerPop\" column and obtain labels.\n",
        "crime_rate_70_percentile = data_df[\"ViolentCrimesPerPop\"].quantile(q=0.7)\n",
        "labels_df = (data_df[\"ViolentCrimesPerPop\"] \u003e= crime_rate_70_percentile)\n",
        "\n",
        "# Now that we have assigned binary labels,\n",
        "# we drop the \"ViolentCrimesPerPop\" column from the data frame.\n",
        "data_df.drop(columns=\"ViolentCrimesPerPop\", inplace=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n8-rQaGH_Bd2"
      },
      "source": [
        "We drop all categorical columns, and use only the numerical/boolean features."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WC12I50e_Bd3"
      },
      "outputs": [],
      "source": [
        "data_df.drop(columns=[\"state\", \"county\", \"community\", \"communityname\", \"fold\"],\n",
        "             inplace=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2eJcNA1B_Bd5"
      },
      "source": [
        "Some of the numerical columns contain missing values (denoted by a NaN). For each feature that has at least one value missing, we append an additional boolean \"is_missing\" feature indicating that the value was missing, and fill the missing value with 0."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8_RgukXn_Bd5"
      },
      "outputs": [],
      "source": [
        "feature_names = data_df.columns\n",
        "for feature_name in feature_names:\n",
        "    # Which rows have missing values?\n",
        "    missing_rows = data_df[feature_name].isna()\n",
        "    if missing_rows.any():  # Check if at least one row has a missing value.\n",
        "        data_df[feature_name].fillna(0.0, inplace=True)  # Fill NaN with 0.\n",
        "        missing_rows.rename(feature_name + \"_is_missing\", inplace=True)\n",
        "        data_df = data_df.join(missing_rows)  # Append \"is_missing\" feature."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S8U5yENt_Bd-"
      },
      "source": [
        "Finally, we divide the dataset randomly into two-thirds for training and one-thirds for testing."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bixDM2G0Bspg"
      },
      "outputs": [],
      "source": [
        "# Set random seed so that the results are reproducible.\n",
        "np.random.seed(123456)\n",
        "\n",
        "# Train and test indices.\n",
        "train_indices, test_indices = model_selection.train_test_split(\n",
        "    np.arange(data_df.shape[0]), test_size=1./3.)\n",
        "\n",
        "# Train and test data.\n",
        "x_train_df = data_df.loc[train_indices].astype(np.float32)\n",
        "y_train_df = labels_df.loc[train_indices].astype(np.float32)\n",
        "x_test_df = data_df.loc[test_indices].astype(np.float32)\n",
        "y_test_df = labels_df.loc[test_indices].astype(np.float32)\n",
        "\n",
        "# Convert data frames to NumPy arrays.\n",
        "x_train = x_train_df.values\n",
        "y_train = y_train_df.values\n",
        "x_test = x_test_df.values\n",
        "y_test = y_test_df.values"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gXyyoSrG_BeA"
      },
      "source": [
        "## (i) PR-AUC Training with Plain TF"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hoqCCpf0oTIY"
      },
      "outputs": [],
      "source": [
        "batch_size = 128  # first fix the batch size for mini-batch training"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wqo2mpRi_BeA"
      },
      "source": [
        "\n",
        "We will work with a linear classification model and define the data and model tensors.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "i1fxMQmy_BeB"
      },
      "outputs": [],
      "source": [
        "# Create linear Keras model.\n",
        "layers = []\n",
        "layers.append(tf.keras.Input(shape=(x_train.shape[-1],)))\n",
        "layers.append(tf.keras.layers.Dense(1))\n",
        "model = tf.keras.Sequential(layers)\n",
        "\n",
        "# Create nullary functions that return labels and logits from the current\n",
        "# batch. In eager mode, TFCO requires these to be provided via nullary function.\n",
        "# We will maintain a running array of batch indices.\n",
        "batch_indices = np.arange(batch_size)\n",
        "labels_fn = lambda: tf.constant(y_train[batch_indices], dtype=tf.float32)\n",
        "logits_fn = lambda: model(x_train[batch_indices, :])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gmFBVhsQ_BeD"
      },
      "source": [
        "We next set up the constraint optimization problem to optimize PR-AUC.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8Swn_y0T_BeD"
      },
      "outputs": [],
      "source": [
        "# Create context with labels and predictions.\n",
        "context = tfco.rate_context(logits_fn, labels_fn)\n",
        "\n",
        "# Create optimization problem with PR-AUC as the objective. The library\n",
        "# expects a minimization objective, so we negate the PR-AUC.\n",
        "\n",
        "# We use the pr_auc rate helper which uses a Riemann approximation to the area\n",
        "# under the precision-recall curve (recall on the horizontal axis, precision on\n",
        "# the vertical axis). We would need to specify the the number of bins\n",
        "# (\"rectangles\") to use for the Riemann approximation. We also can optionally\n",
        "# specify the surrogate to be used to approximate the PR-AUC.\n",
        "pr_auc_rate = tfco.pr_auc(\n",
        "    context, bins=10, penalty_loss=tfco.SoftmaxCrossEntropyLoss())\n",
        "problem = tfco.RateMinimizationProblem(-pr_auc_rate)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OwpTHgqg_BeG"
      },
      "source": [
        "We then create a loss function from the `problem` and optimize it to train the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nkq2aVIn_BeG"
      },
      "outputs": [],
      "source": [
        "# Create Lagrangian loss for `problem`. What we get back is a loss function, a\n",
        "# a nullary function that returns a list of update_ops that need to be run\n",
        "# before every gradient update, and the Lagrange multiplier variables internally\n",
        "# maintained by the loss function. The argument `dual_scale` is a\n",
        "# hyper-parameter that specifies the relative importance placed on updates on\n",
        "# the Lagrange multipliers.\n",
        "loss_fn, update_ops_fn, multipliers = tfco.create_lagrangian_loss(\n",
        "    problem, dual_scale=1.0)\n",
        "\n",
        "# Set up optimizer and the list of variables to optimize.\n",
        "optimizer = tf.keras.optimizers.Adagrad(learning_rate=0.1)\n",
        "var_list = (\n",
        "    model.trainable_weights + list(problem.trainable_variables) + [multipliers])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ra-Cog9C_BeI"
      },
      "source": [
        "Before proceeding to solving the training problem, we write an evaluation function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gFUyqbVx_BeI"
      },
      "outputs": [],
      "source": [
        "def pr_auc(model, features, labels):\n",
        "    # Returns the PR-AUC for given model, features and binary labels.\n",
        "    scores = model.predict(features)\n",
        "    return metrics.average_precision_score(labels, scores)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MID6ChJn_BeK"
      },
      "source": [
        "We are now ready to train our model. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 261
        },
        "id": "ZWUWtSp8NzK5",
        "outputId": "bd5488b6-030c-4cb7-e276-8788ec81a3d5"
      },
      "outputs": [
        {
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAfcAAAD0CAYAAACGlm89AAAABHNCSVQICAgIfAhkiAAAAAlwSFlz\nAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjEsIGh0\ndHA6Ly9tYXRwbG90bGliLm9yZy+j8jraAAAgAElEQVR4nO3de5xcVZnv/8/T91uS7lwJuQMJEEQD\nBFCZUUCBgDOCRw4mzgxhxp/8OArO4Iw/4OAgwwz+cM454ujhOIOKeEGQQS7RQREFdFSCJBAuCeYC\nISGdkHun7911ec4fe1eyafpS1V3dVbX7+3696lVV+9ZPVbLq2Wvttdcyd0dERETio6zQAYiIiEh+\nKbmLiIjEjJK7iIhIzCi5i4iIxIySu4iISMwouYuIiMSMknuMmdlPzWxloeMQEZGxpeReZMysPfJI\nm1lX5P2f5XIsd7/Q3b8zzDhej/zt3WZ2t5k1hOueMrPucN0+M3vQzGZmccynzOygmVX3s/z/6bPs\nbDPbEXlvZvYZM3vZzDrMbIeZ/buZnTyczydSDPJZ3sPjva0s9Vk/38w88jdeN7PrI+s9LF/tZtZs\nZl82s/Ih/mZDuP1P+1nnZnZcn2U3m9n3I+8nmtlXzGx7eJxXw/dTc/v0EqXkXmTcvSHzALYDfxpZ\ndk9mOzOrGINw/jSM41RgKfD5yLqrw3XHAQ3A/xzsQGY2H/hjwIEPDyOWfwH+GvgMMBlYBDwMfGgY\nxxIpCtmW91HQGP7NFcBNZrYssu5d4br3Ax8D/mqIY30U6AHOM7OjcgnCzKqAXwInAcuAicB7gP3A\nGbkcS95Kyb1EZGqyZnadmb0JfNvMmszsJ2a2N6wR/8TMZkf2OXwWb2ZXmNlvzOx/httuNbMLs/nb\n7t4M/BR4Rz/rWgiS7JIhDnM5sBq4G8jpUoGZLQQ+Daxw9yfcvcfdO939Hne/LZdjiZQCMyszs+vD\nWux+M7vfzCaH62rM7Pvh8hYze9bMZpjZrQQn0P87rAH/76H+jrs/Dayn/7K9BfgtQ5ftlcC/Ai8C\nf57bJ+VyYC7wEXff4O5pd9/j7v/o7o/meCyJUHIvLUcR1FrnAVcS/Pt9O3w/F+gCBivQZwIbganA\nPwPfMjMb6o+a2RzgIuD5ftZNAf4LsGWIw1wO3BM+LjCzGUP93YgPADvc/fc57CNSyq4BLiGoPR8N\nHATuCNetBCYBc4ApwFVAl7vfCPwnYauau1892B8IL3WdRVBr7q9sn0BwsjBg2TazecDZHCnbl2f/\nEQH4IPAzd2/PcT8ZgpJ7aUkDXwhrrl3uvt/dfxTWYtuAWwl+DAayzd2/4e4p4DvATGCwJPuwmbUA\nvwF+BXwxsu6rZnYI2EdwsnDNQAcxsz8iOAG5393XAq8CHx/y0x4xBdiVw/Yipe4q4EZ33+HuPcDN\nwKXh5bgEQZk4zt1T7r7W3VtzPP4+4ADwTeB6d/9lZN1zZtYBvAI8BfyfQY7zF8CL7r4BuA84ycxO\nySEOle1RouReWva6e3fmjZnVmdm/mdk2M2sFfg00DtIB5s3MC3fvDF82DPL3LnH3Rnef5+6fcveu\nyLrPuPsk4J1AEzA7jOm/Rzrr/Gu47Urg5+6+L3z/A97aNJ8EKvv87UqCHzEIrr8N2WFPJEbmAQ+F\nze4tBIk2RXAy/j3gMeA+M9tpZv9sZn3Lz1CmunuTu5/o7l/ts+5Ugt+FjxG09tXD4btv+nb2y7TI\nZS7f/Yq3lu0UKtsFoeReWvpO4fe3wPHAme4+EXhfuHzIpva8BeT+EvBPwB1mZu7+xUiHoKvMrBa4\nDHi/mb0Z9he4FniXmb0rPMx2YH6fQy8AtoWvfwnMNrOlo/6BRIrDG8CF4cl15lHj7s3unnD3f3D3\nxcB7gT/hSHN4Xqb59MD9wNPATeGyC6Od/czsvcBC4IZI2T4T+Hikw+9QZfsXBJfp6vMRtxyh5F7a\nJhBcZ28JO9t8oUBxfIegRtFfL/hLCM7eFxN0zFkCnEhwbTDzg/RD4C/N7IzwOuAighOA+wDcfTNB\n0+C9YcfCqrBT0fLobTwiMfKvwK3hNW3MbJqZXRy+PsfMTg5b6FoJasHpcL/dwDF5jOM24JMD9IJf\nCTzOW8v2O4BaINNZ94fA581sdthJ8IPAnwIPhOu/R3Ai8yMzOyHcZkrYAnhRHj/HuKPkXtq+QlCQ\n9hH0RP9ZIYJw916CW9X+vp/VK4Fvu/t2d38z8yDo+PdnZlbh7o8B1xN0DjwEPEpwwnBn5DifCfe5\nA2ghuG7/EeDHo/SxRArpX4BVwM/NrI2gfJ8ZrjuKIDm2EjTX/4ogSWb2uzS8I6Zvc3vOwpa5XwOf\niy43sxqCFrmvRcu1u28NY8k0zd8C/I6g385Bgo68f+buL4fH7yHoVPcHghOFVuD3BP14nhlp/OOZ\nueelFUdERESKhGruIiIiMaPkLiIiEjNK7iIiIjGj5C4iIhIzYzH5SE6mTp3q8+fPL3QYIiVh7dq1\n+9x9WqHjyJXKuUj2hlPOiy65z58/nzVr1hQ6DJGSYGbbht6q+Kici2RvOOVczfIiIiIxo+QuIiIS\nM0ruIiIiMaPkLjLOmNkyM9toZlv6G5vfzOaa2ZNm9ryZvRgd49vMbgj322hmF2R7TBEZW0ruIuNI\nONnIHQQTeywGVpjZ4j6bfR64391PAZYTzucdbrccOAlYBvwfMyvP8pgiMoaU3EWKUDKVJpUelXkf\nzgC2uPtr4YQ/9wEX99nGgYnh60nAzvD1xcB97t4TThCyJTxeNscUkT46epKM1vwuRXcrnEixSaWd\nrkSKrt4U3YnU4ddd4eueRIruRPrwuszr7mSK7t4jy4Ln8DjJFD2JNL2pNIlk8NyTTJNIpelNpkk7\n/NtfnMYFJ/U30+aIzCKYYjNjB0dmG8u4mWA2smuAeoJZuzL7ru6z76zw9VDHFBmXuhMptu3vZOu+\nDrbu6+D1fR1s3R8872nr4bm/P4/J9VV5/7tK7lLSUmmnN5mmJ5kKn4NHZll3Ik1HT5L2niRtPUna\nu5O09yTo6EnR1p2kszd5JOkm0vQcTtCZJBwcaziqKsqorSwPHlXlVFeUUVsVvJ9YW0l1RRlVFWVU\nlZdRGT5nllWWl3HstIY8f1tZWwHc7e7/y8zeA3zPzN4x0oOa2ZXAlQBz584d6eFkHHAPTqxbu5Kk\n3EmnHXdIu4ePYJue5JFynnl0HC7vKdLulJlRZlBeZpgZ5WVQbgZmdPUG23X2Bvt19KTo6A2O05tM\nH963rMwoNwvelwXHcA9+h9LupNLhw510GnpTafa19xCtnE9tqGL+lHrev2ga86fWU2aj890pucuo\nS6SCgtfZmyk8QcHpDJ87elKHC2ZHn8LZ0RPWjpOpIHGHtd2eRPA+OYym6/Iyo6G6gobqCuqqgsRb\nU1HOpNpKaiZUH35fU1lGbVVFmJyDRF0TJupM0q4On2sqy6jJrK8MEnnZaJXakWkG5kTezw6XRX2C\n4Jo67v50OHf31CH2HeqYuPudwJ0AS5cu1VzT40wylaalK8HBjl4OdPRysLOXAx2J8LmXls4Eh7qC\n55auxOH3idTw/6uYQX1VBWUWJuEwAWdODDKXvmoqy2iorqC+uoK6qgoaqsuZXF/FnKY6qivKcDic\ntP1wEg9OMowjSb+8PHwuC04AKsqMoxtrmT+1jgVT65k/tZ6JNZX5+UKHoOQuWXF3OnpTtHYlaO1O\n0NqVjLxOcLAzQUtnLwc7g8La0pkIC2wvHb2prP9OfVU59WHibaipoL6qgmkTqqmuOFKrra4oP1Lr\nrQiSalV5GdWVYe038r66oowJ1ZXBsarLmVBdSU1lGWZFmXjHwrPAQjNbQJCAlwMf77PNduADwN1m\ndiJQA+wFVgE/MLMvA0cDC4HfA5bFMWWcONSZYPOeNjbvaWfz7vbg9e523mztHnCfuqpymuqqmFRb\nSWNdJYtmNDCp9sj7CTUVVJaVYUZYaw6eLaxRV5aXHT5hr6+uYEJNmKgry4c8yXb3WP4eKLkL6bSz\nr72H5pYudh3qZmdLFztbutl1qIudh7rZ1dLFvvYehqokT6ypoKm+iqa6KqY2VLFwegONdVU01lWG\nha6cuqrIc1UFddXl1IfL6qsqirW2GxvunjSzq4HHgHLgLndfb2a3AGvcfRXwt8A3zOxags51V3jQ\n62e9md0PbACSwKfdPQXQ3zHH/MPJmEim0rzZ2k3zwS52Huqi+WAXzS1dbD/Qyebd7exp6zm8bW1l\nOcdNb+C9x05h9uQ6ptRX0VRfxeS6KibXB4/GukpqKssL9nnimNgBbLR66g3X0qVLXWNO51cq7Rzs\n7KX5YBdvHOxkx8Eu3jjQyRsHu9hxoJMdLV1vu65cW1nOzMYaZjXWMnNSDdMmVDOptpKJNZVMPPxc\nwcSa4Kx6Um0lFeW6+WKsmdlad19a6DhypXJevJKpNLsOdQe/FQe62BH+ZgSPTt5s7X7bif7Uhipm\nNdWxcHpD8JjRwMLpE5jVWKsT9jwYTjlXzb1E9CRT7Gvvpas3SVdv+nBHsEyv7c7eFIe6Euxv72V/\nRw/72nvY397LvvYeDnT0vq0wNtZVMqepjhNmTuC8xTOY3VTL0Y21zJxUy9GNNUyqrYztGa3IeOfu\n7O/oZfPudrbsaWPLnnY272ln2/4geUdvwywzmDmplllNtbz7mCnMaqplVmPwe5F5Xciat/RPyb1I\n9CbTbNvfwRsHO2k+2MWOliPNXc0Hu9jbp8flQOqrypk6oZop9VXMmVzHKXObmNpQxZT6Ko5urGXO\n5DpmN9UyYYw6dYhI4SRSad44cOQ2rFf3drAlvB7e0pk4vF1DdQXHTm/g9PlNh38j5jTVMbupjpmN\nNVSqVa7kKLmPsc7eJK/u6WBzeLaceWw70PmWs+XK8qCX5azGWt6/aBqzmmo5amINdWEnkdqqI722\n68LnCTWV1FbpDFpkvEinnX0dPew+1MPu1qCfzNZ9nby+P0jm2/v8rjTVVbJw+gQuOnkmx00Lms+P\nm97AURNr1FIXM0ruoyiVdjbvaWPttoOsff0ga7cfZNv+zsPrK8qM+VPrWTRjAh9650yOndZw+Kx5\nWkO1rlWJCBBcB39uewtPbtzD1r0d7G7rZvehbva09bztdtDaynLmT61n8cyJfOjkmSyYWs+CafUs\nmFJP0ygMliLFSck9j9p7kqzb3hIk8+0HeX7bQdp6kkDQ4eS0eU1ceursw2fL86bUq7lLRPrV1p3g\n15v28ctXdvPkxj0c7EwcrhDMmFjNu4+dwlETazhqUg0zJgaPmZNqVDEQIMvkbmbLgH8huM3lm+5+\nW5/1c4HvAI3hNte7+6NmNh94BdgYbrra3a/KT+iFt7u1mzWvH+TZ1w+wZtsBNuxsJe3BwAnHz5jA\nh5cczWnzmjhtXhNzJ9ep2UtEBrWzpYufr3+TX/5hD6tf208i5TTWVXLO8dP54IkzeN+iqeovI1kZ\nMrlHZnw6j2DM6GfNbJW7b4hslplF6uvhbFCPAvPDda+6+5L8hl0YBzp6eWz9mzy79QDPbjvAGwe6\ngKAZbMmcRq4+5zhOmz+ZU+Y2jtkoRCJS2t481M2jL+3iJy/u5LntLQAcM62evzprAR84cQanzm3U\nbaaSs2xq7odnfAIws8yMT9HkPtAsUiUvmUrzq017+fc1O/jlH3aTSDlTG6pZOq+Jle+Zz+nzJ7P4\n6IlqXheRrO1p7eanL7/JT17cybOvHwTgxJkT+dwFx3PhO47imMLNKyAxkU1yH8ksUgALzOx5oBX4\nvLv/5/DDHTtb9rTz72vf4MHnmtnb1sOU+iouf898PnrqbE6cOUFN7CKSk+aWLh5f/yY/W/8mz2w9\ngHtw+e6z5y063KFWJF/y1aFuoFmkdgFz3X2/mZ0GPGxmJ7l7a3TnYpktKpV2Hnq+mXue2cbz21so\nLzPOOX46/3XpbM45fjpVFaqdi0h23J0Nu1p5fMNufr5+Nxt2BT97C6c38JlzF/In75zJwhkTChyl\nxFU2yX3Ys0i5+x6gJ1y+1sxeBRYBbxl3shhmi3ppxyE+/8jLvPBGC4tmNHDjRSdyySmzmDahuhDh\niEgJytwx84tXdvP4ht00t3RhBqfNbeK/X3QC5y0+igVT6wsdpowD2ST3Yc8iZWbTgAPunjKzYwhm\nkXotb9HnwaGuBF/++Ua+t3obk+ur+ZflS/jwu45Ws7uIDCqVdrbsaWfdGwd5fnsL695oYdPuNtIO\n1RVl/PHCqfz1BxZy7onTmdqgSoKMrSGT+0hmkTKz9wG3mFkCSANXufuBUfs0OXB3Hlm3k3/6j1c4\n0NHD5e+Zz2fPX6Re7hJrWdzWejtwTvi2Dpju7o1mdg5we2TTE4Dl7v6wmd0NvB84FK67wt3XjeLH\nKJjdrd08+Fwz/7l5Ly/uOER7OI7FpNpKlsxpZNk7jmLJnEbOWDCZuioNIyKFk9X/Pnd/lOD2tuiy\nmyKvNwBn9bPfj4AfjTDGvNuyp43PP/wyq187wLvmNHL3X57OO2ZNKnRYIqMqm9ta3f3ayPbXAKeE\ny58EloTLJwNbgJ9HDv85d39g1D9EAfQm0zzxh93cv2YHT23cQ9rhpKMn8pFTZnHK3EaWzGlkwdR6\ntfZJURl3p5b3/n47Nz3yMnVVFXzxIyez/PQ5Gs1JxotsbmuNWgF8oZ/llwI/dffOftbFxsY327h/\nzRs89HwzBzp6OWpiDZ86+zguPW0283XdXIrcuEruO1u6+Icfr+eMBZP56vJTmKLrYDK+ZHNbKwBm\nNg9YADzRz+rlwJf7LLvVzG4CfkkwQmVPP8csirtiBtPS2cuPX9jJA2t38MKOQ1SWGx88cQaXnT6H\n9y2cRrkqAlIixlVy/9LP/kDa4bb/8k4ldpHBLQcecPdUdKGZzQROJuiDk3ED8CZQRXDXy3XALX0P\nWAx3xfSnN5nmqY17ePC55sMDVZ1w1AT+/k8Wc8mSo/VbISVp3CT3tdsO8Mi6nVxz7nHMmVxX6HBE\nCiGb21ozlgOf7mf5ZcBD7n54MnB33xW+7DGzbwN/l4dYR5W781LzIR58rplVL+zkQEcvUxuq+It3\nz+ejp81i8cyJuoYuJW1cJPd02vmHH2/gqIk1/Lezjy10OCKFks1trZjZCUAT8HQ/x1hBUFOPbj/T\n3XdZkA0vAV7Od+D59KtNe/nif7zCxt1tVFWUcd7iGXz01Fn88cJpGkZaYmNcJPcfPbeDF3cc4isf\nW6LbU2TcyvK2VgiS/n3u/pam83CWxznAr/oc+p5wTAsD1gFFOfPjntZubvnJBn7y4i4WTK3nix85\nmQ+dPJNJdbr9VeIn9pmurTvBl362kVPmNnLxkqMLHY5IQQ11W2v4/uYB9n2doFNe3+Xn5i/C/Eul\nnXue2cb/+NlGelJprv3gIq46+xiqK8oLHZrIqIl9cr/jyVfZ197DN1cu1TU0kXHm5eZD3PhwMKz0\nWcdN4Z8uOVnDv8q4EOvk/vq+Du76zVY+eupslsxpLHQ4IjJG2nuS3P74Jr79261Mrq/iKx9bwsVL\nNKy0jB+xTu63PvoKleXGdcuOL3QoIjJGtuxp46/uXsP2A518/My5XHfBCbquLuNObJP7f27ey+Mb\ndvP/LTue6RNrCh2OiIyB32zex3+7Zy3VFWXc//++hzMWTC50SCIFEcvknkylueXHG5g7uY6/OmtB\nocMRkTHwg2e28/ePvMxx0xr41hVLmd2k8Sxk/Iplcr/nme1s3tPOv/75adRUqkesSJyl0s7//+gr\nfPM3Wzn7+Gl8bcUpTNDsjjLOxS65H+zo5cuPb+K9x07hgpNmFDocERlFHT1J/vq+dfzild1c8d75\nfP5DJ1KhgWhE4pfcn9q0h0NdCf7uguPVM1YkxnYd6uITd6/hD2+28g8fPomV751f6JBEikbskvuh\nzmDI63kaP14ktrbsaePPvvkMHT0pvnXF6Zxz/PRChyRSVGKX3Nu6kwC65iYSY99fvZ1DXQke/vRZ\nnHDUxEKHI1J0Yndxqq0nSXVFGVUVsftoIhJq6exl+oQaJXaRAcQuA7Z1J1VrF4m5oJzHruFRJG9i\nmNwTTFShF+mXmS0zs41mtsXMru9n/e1mti58bDKzlsi6VGTdqsjyBWb2THjMH5pZ1Wh/DiV3kcHF\nMLknaVChF3kbMysH7gAuBBYDK8xscXQbd7/W3Ze4+xLga8CDkdVdmXXu/uHI8i8Bt7v7ccBB4BOj\n+kGA1u6EWuhEBhHD5J7QGb1I/84Atrj7a+7eC9wHXDzI9iuAewc7oAX3m54LPBAu+g5wSR5iHZRq\n7iKDi2FyTzKhWmf0Iv2YBbwReb+DfuZnBzCzecAC4InI4hozW2Nmq80sk8CnAC3unszimFeG+6/Z\nu3fvSD4Hrd0JJqrmLjKg2J36tvfojF4kD5YDD7h7KrJsnrs3m9kxwBNm9hJwKNsDuvudwJ0AS5cu\n9eEGlk67yrnIEOJZc9cZvUh/moE5kfezw2X9WU6fJnl3bw6fXwOeAk4B9gONZpbJtIMdMy86epO4\no+QuMohYJfdUeEavDnUi/XoWWBj2bq8iSOCr+m5kZicATcDTkWVNZlYdvp4KnAVscHcHngQuDTdd\nCTwymh9CA1WJDC2r5J7F7TNzzexJM3vezF40s4si624I99toZhfkM/i+2nuCQq9b4UTeLrwufjXw\nGPAKcL+7rzezW8ws2vt9OXBfmLgzTgTWmNkLBMn8NnffEK67DvismW0huAb/rdH8HJnkrmvuIgMb\nMgtGbp85j6CzzLNmtipSsAE+T/BD8fXw1ppHgfnh6+XAScDRwC/MbFGf63h509YdjCuv5jqR/rn7\nowTlM7rspj7vb+5nv98BJw9wzNcIeuKPCZVzkaFlU3PP5vYZBzLjQE4CdoavLyaoAfS4+1ZgC6P4\nI6DmOpH4O1LOldxFBpJNcs/m9pmbgT83sx0EtYJrctg3b7fIZJrlVehF4qv1cM1dJ/EiA8lXh7oV\nwN3uPhu4CPiemWV9bHe/092XuvvSadOmDTuINhV6kdg7cs1dJ/EiA8mmdGRz+8wngGUA7v60mdUA\nU7PcN28yhb6hWoVeJK50+U1kaNnUrrO5fWY78AEAMzsRqAH2htstN7NqM1sALAR+n6/g+2rVGb1I\n7LV1J6goM2oqY3Unr0heDZkF3T1pZpnbZ8qBuzK3zwBr3H0V8LfAN8zsWoLOdVeEt9GsN7P7gQ1A\nEvj0aPWUBzXLi4wHmXHlg2HtRaQ/WVVxh7p9Jrwt7qwB9r0VuHUEMWatvTupM3qRmGvTjHAiQ4pV\nFtQZvUj8aUY4kaHFLLknNPSsSMwpuYsMLWbJXdO9isRdq5rlRYYUv+SuM3qRWFM5FxlavJJ7j6Z7\nFRlMFpNA3W5m68LHJjNrCZcvMbOnzWx9ODnUxyL73G1mWyP7LRnNz9DWndCkMSJDiNXpb1DoJxQ6\nDJGilM0kUO5+bWT7awjmbAfoBC53981mdjSw1swec/eWcP3n3P2B0f4M7sG0zqq5iwwuXjX3bs3l\nLjKIbCaBiloB3Avg7pvcfXP4eiewBxj+WNHD1NGbIu2aP0JkKLFJ7jqjFxlSVhM5AZjZPGAB8EQ/\n684AqoBXI4tvDZvrbzez6gGOOeIJojRQlUh2YpPcO3tTpNKuQi+SH8uBB/qOKGlmM4HvAX/p7ulw\n8Q3ACcDpwGTguv4OmI8JojTdq0h2YpPcVehFhpTLRE7LCZvkM8xsIvAfwI3uvjqz3N13eaAH+DZB\n8/+oUM1dJDuxSe7tPSr0IkPIZhIozOwEoAl4OrKsCngI+G7fjnNhbR4Lhoa8BHh5tD5Aq07iRbIS\nmxKiQi8yuCwngYIg6d8XTv6UcRnwPmCKmV0RLrvC3dcB95jZNMCAdcBVo/UZNJe7SHZiU0ION8tr\nLneRAQ01CVT4/uZ+9vs+8P0BjnluHkMclJrlRbITm2Z5FXqR+FPfGpHsxCi5q9CLxF1bd4LyMqO2\nsrzQoYgUtdgk93Yld5HYa+3StM4i2YhNcm/rTmAG9VVK7iJx1dad0Am8SBZik9xbu5M0VFVQVqYz\nepG40rTOItmJTXLXNJAi8adyLpKdGCX3hHrKi8Rcq8q5SFZik9w1aYxI/LV1JzWAjUgWYpPc1Vwn\nEn/qUCeSnRgldzXXicTZkWmdVc5FhhKj5J6kQWf0IrHV0Zsi7RrLQiQbsUruKvQi8aUhpkWyF4vk\n3pNM0ZtKM1GFXmRQZrbMzDaa2RYzu76f9beb2brwscnMWiLrVprZ5vCxMrL8NDN7KTzmV22Uho/T\nENMi2csquY/wByEVWfe2uaPzQYVeZGhmVg7cAVwILAZWmNni6Dbufq27L3H3JcDXgAfDfScDXwDO\nBM4AvmBmTeFuXwc+CSwMH8tGI/4jNXeVc5GhDFlKIj8I5wE7gGfNbJW7b8hs4+7XRra/Bjglcoiu\n8Idi1Ci5i2TlDGCLu78GYGb3ARcDGwbYfgVBQge4AHjc3Q+E+z4OLDOzp4CJ7r46XP5d4BLgp/kO\nvvVwOVcLnchQsqm5H/5BcPdeIPODMJAVwL35CC5bmTP6Bg1LKTKYWcAbkfc7wmVvY2bzgAXAE0Ps\nOyt8nc0xrzSzNWa2Zu/evTkHnzmJ133uIkPLJrmP5AcBoCYs0KvN7JIB9stLoVfNXSRvlgMPuHsq\nXwd09zvdfam7L502bVrO+6tDnUj28t2hrr8fhHnuvhT4OPAVMzu27075K/RK7iKDaAbmRN7PDpf1\nZzlvbYEbaN/m8HU2xxyRwzX3WpVzkaFkk9xH8oOAuzeHz68BT/HW6/F5caS5Tmf0IoN4FlhoZgvM\nrIqgvL6tk6uZnQA0AU9HFj8GnG9mTWFHuvOBx9x9F9BqZu8Oe8lfDjwyGsG3dScoLzNqK8tH4/Ai\nsZJNch/2D0L4Q1Advp4KnMXAnXeGTc3yIkNz9yRwNUGifgW4393Xm9ktZvbhyKbLgfvc3SP7HgD+\nkeD34FnglkznOuBTwDeBLcCrjEJnOjgylsUo3WknEitDZkN3T5pZ5gehHLgr84MArHH3TKJ/2w8C\ncCLwb2aWJjiRuC3ayz5fMsm9vlrJXWQw7v4o8GifZTf1eX/zAPveBdzVz/I1wDvyF2X/NFCVSPay\nKinD/UFw998BJ48gvqy0dV9BZOwAABKtSURBVCeorSynsjwWY/KISD/auhNM0B0xIlmJRTbUGb1I\n/LWqnItkLRbJXXO5i8RfcBKvmrtINmKR3Fs13atI7LV1JzSAjUiWYpHc1SwvEn8q5yLZi0lyT6jQ\ni8SYu4eX39RCJ5KNmCT3pHrRisRYZ2+KVNp1Ei+SpVgkd3WoE4m3Ns0IJ5KTkk/uyVSazt6UCr1I\njGn+CJHclHxyb+/R0LMicdeq5C6Sk5JP7pnmugYVepHYalWzvEhOSj65Z87odf+ryNDMbJmZbTSz\nLWZ2/QDbXGZmG8xsvZn9IFx2jpmtizy6zeyScN3dZrY1sm5JvuM+MvOjyrlINkq+pKijjUh2zKwc\nuAM4D9gBPGtmq6KTOZnZQuAG4Cx3P2hm0wHc/UlgSbjNZIIZ4H4eOfzn3P2B0Yr9yDV3lXORbJR8\nzb1d072KZOsMYIu7v+buvcB9wMV9tvkkcIe7HwRw9z39HOdS4Kfu3jmq0UZoWmeR3JR8cm/r0Rm9\nSJZmAW9E3u8Il0UtAhaZ2W/NbLWZLevnOMuBe/ssu9XMXjSz282sOn8hB9q6E5SXGXVV5fk+tEgs\nlX5y1xm9SD5VAAuBs4EVwDfMrDGz0sxmEkzj/FhknxuAE4DTgcnAdf0d2MyuNLM1ZrZm7969OQXV\n1p2koboCM8tpP5HxKjbJvaFayV1kCM3AnMj72eGyqB3AKndPuPtWYBNBss+4DHjI3ROZBe6+ywM9\nwLcJmv/fxt3vdPel7r502rRpOQWuceVFclPyyb21O0FVeRk1lWquExnCs8BCM1tgZlUEzeur+mzz\nMEGtHTObStBM/1pk/Qr6NMmHtXksqFZfAryc78DbNPOjSE5K/lS4XWf0Illx96SZXU3QpF4O3OXu\n683sFmCNu68K151vZhuAFEEv+P0AZjafoOb/qz6HvsfMpgEGrAOuynfsrSrnIjkp+dKi5jqR7Ln7\no8CjfZbdFHntwGfDR999X+ftHfBw93PzHmgfbd1JZjXWjPafEYmNkm+WV3OdSPypnIvkJgbJPanO\ndCIxpxY6kdzEIrmr0IvEl7trWmeRHJV8cg8KvZrrROKqszdFKu0q5yI5KPnk3tqd0Bm9SIxpoCqR\n3JV0ck+ng+Y6zRQlEl+aNEYkdyWd3Dt6k7ir0IvEWatq7iI5yyq5DzUHdDhZRGYu501m1hJZt9LM\nNoePlfkM/vDQsyr0IrGVqbmrhU4ke0OWlmzmgHb3ayPbXwOcEr6eDHwBWAo4sDbc92A+gm/v0Rm9\nSNwdueauFjqRbGVTc89mDuio6NjTFwCPu/uBMKE/DvQ3heSw6FqcSPypQ51I7rJJ7tnMAQ2Amc0D\nFgBP5LLvcKeC1LU4kfjTSbxI7vLdoW458IC7p3LZabhTQWbO6HUtTiS+2rqTlBnUV2nmR5FsZZPc\ns5kDOmM5b50OMpd9c5Y5o2+o1hm9SDaG6hwbbnOZmW0ws/Vm9oPI8lSk4+yqyPIFZvZMeMwfhtPJ\n5k1bd4KG6gqCGWVFJBvZJPds5oDGzE4AmoCnI4sz00c2mVkTcH64LC90LU4ke5HOsRcCi4EVZra4\nzzYLgRuAs9z9JOBvIqu73H1J+PhwZPmXgNvd/TjgIPCJfMYdDDGtE3iRXAyZ3N09CWTmgH4FuD8z\nB7SZRQv4cuC+cMrIzL4HgH8kOEF4FrglXJYX7d1JysuMOjXXiWQjm86xnwTuyNzR4u57BjugBdXp\nc4EHwkXfAS7JZ9Cay10kd1mVmKHmgA7f3zzAvncBdw0zvkGpuU4kJ/11cD2zzzaLAMzst0A5cLO7\n/yxcV2Nma4AkcJu7PwxMAVrCSkDmmAN1uL0SuBJg7ty5WQfd1p1gYq1q7iK5KOnTYc0IJ5J3FcBC\n4GyCPjK/NrOT3b0FmOfuzWZ2DPCEmb0EHMr2wO5+J3AnwNKlS32IzQ9r605ydGNNDh9BREp6+NlW\nzeUukotsOrjuAFa5e8LdtwKbCJI97t4cPr8GPEUwWNV+oNHMKgY55oi09SR0zV0kRyWd3Nu6E0xU\noRfJVjadYx8mqLVjZlMJmulfCzvFVkeWnwVsCPvYPAlcGu6/Engkn0G3dqmFTiRXJZ3cg7ncVehF\nspFl59jHgP1mtoEgaX/O3fcDJwJrzOyFcPltkSGorwM+a2ZbCK7BfyuPMauciwxDSZeYtu4kC6eX\n9EcQGVNDdY4Na+KfDR/RbX4HnDzAMV8j6Imfd529KVJpV7O8SI5Kuube1q1rcSJxprEsRIanZJO7\nu6u3vEjMaVx5keEp2eTenUiTTLvmcheJMU0OJTI8JZvc23p0Ri8Sd5mauyaHEslN6SZ3zQgnEntH\nrrnrJF4kFyWf3NVcJxJfKuciw1PCyV3N8iJxp3IuMjwlnNyDM3oNPysSX23dScoM6jXzo0hOSji5\nZ87oldxF4kozP4oMTwknd3W0EYm7YCwLlXGRXJV8clezvEh8tWqgKpFhKenk3lBdQXmZmutE4koz\nP4oMTwkn94TO6EVyZGbLzGyjmW0xs+sH2OYyM9tgZuvN7AfhsiVm9nS47EUz+1hk+7vNbKuZrQsf\nS/IVr4aYFhmeki01mZq7iGTHzMqBO4DzgB3As2a2KjJ1K2a2ELgBOMvdD5rZ9HBVJ3C5u282s6OB\ntWb2mLu3hOs/5+4P5Dvmtp4Ei2oa8n1Ykdgr2Zq75ngWydkZwBZ3f83de4H7gIv7bPNJ4A53Pwjg\n7nvC503uvjl8vRPYA0wb7YDVoU5keEo2uWu6V5GczQLeiLzfES6LWgQsMrPfmtlqM1vW9yBmdgZQ\nBbwaWXxr2Fx/u5lV9/fHzexKM1tjZmv27t07ZLCa+VFk+Eo4uavQi4yCCmAhcDawAviGmTVmVprZ\nTOB7wF+6ezpcfANwAnA6MBm4rr8Du/ud7r7U3ZdOmzZ0pb8rkSKVdp3EiwxDySb3VjXXieSqGZgT\neT87XBa1A1jl7gl33wpsIkj2mNlE4D+AG919dWYHd9/lgR7g2wTN/yOmceVFhq9kk7t6y4vk7Flg\noZktMLMqYDmwqs82DxPU2jGzqQTN9K+F2z8EfLdvx7mwNo8Fw8hdArycj2A1CqXI8JVkqelNpulJ\nppmg3vIiWXP3pJldDTwGlAN3uft6M7sFWOPuq8J155vZBiBF0At+v5n9OfA+YIqZXREe8gp3Xwfc\nY2bTAAPWAVflI97Ww9M6q4VOJFclmR3be9RcJzIc7v4o8GifZTdFXjvw2fAR3eb7wPcHOOa5+Y9U\nzfIiI5FVs/xwB74Il6cig1v0bQIcFk0DKRJ/KuciwzfkKfEIB74A6HL3vI1YBTqjFxkPVM5Fhi+b\nmvuwB74YLa06oxeJPXWoExm+bJL7SAe+qAkHrlhtZpf09wdyHdyiXWf0IrHX1p3EDOqrVM5FcpWv\nUhMd+GI28GszOzkcd3qeuzeb2THAE2b2krtHR7bC3e8E7gRYunSpD/XH1FwnEn+tXQkaqiso08yP\nIjnLpuY+ooEv3L05fH4NeAo4ZYQxq6ONyDjQ1p3UbXAiw5RNch/JwBdNmXGmw+VnARsYIdXcReKv\nVUNMiwzbkCVnhANfvBf4NzNLE5xI3BbtZT9cbT1JairLqCwv2QH2RGQIGoVSZPiyKjkjGPjid8DJ\nIw/zrdq6EzRUq7lOJM7aupMcNamm0GGIlKSSrPoG1+J0Ri8SZ209qrmLDFfJJncVepF4UzkXGb4S\nTe4J9ZQXGYYRDiW90sw2h4+VkeWnmdlL4TG/Gs4ONyLuHiZ3lXOR4SjR5K4zepFcRYaSvhBYDKww\ns8V9tokOJX0S8Dfh8snAF4AzCUat/IKZNYW7fZ1glMqF4SM6iNWwdCVSpNKuW+FEhqlkk3uDpnsV\nydVIhpK+AHjc3Q+E6x4HloVzuU9099Vhx9rvEszpPiK63VVkZEoyubf3qLlOZBhGMpT0QPvOCl8P\ndkwgt2GmNa68yMiUXHJPpT1M7ir0IqMgOpT0CuAbZtaYjwO7+53uvtTdl06bNm3QbVvDmrua5UWG\np+SSe3uPmutEhmkkQ0kPtG9z+HqwY+ZMzfIiI1Nyyb2uqpxVV5/Fn77r6EKHIlJqhj2UNEdGoWwK\nO9KdDzzm7ruAVjN7d9hL/nLgkZEGeurcRh761Hs5YebEkR5KZFwqudPiyvIy3jk7L62EIuPKSIaS\nBjCzfyQ4QQC4xd0PhK8/BdwN1AI/DR8jMqGmklPmNg29oYj0q+SSu4gM33CHkg7X3QXc1c/yNcA7\n8h6siAxbyTXLi4iIyOCU3EVERGJGyV1ERCRmlNxFRERiRsldREQkZizoHFs8zGwvsG2IzaYC+8Yg\nnFwVa1xQvLEVa1xQGrHNc/fBh3srQirno6ZYYyvWuKB4Y4vGlXM5L7rkng0zW+PuSwsdR1/FGhcU\nb2zFGhcotkIr1s9YrHFB8cZWrHFB8cY20rjULC8iIhIzSu4iIiIxU6rJ/c5CBzCAYo0Lije2Yo0L\nFFuhFetnLNa4oHhjK9a4oHhjG1FcJXnNXURERAZWqjV3ERERGYCSu4iISMyUVHI3s2VmttHMtpjZ\n9UUQz+tm9pKZrTOzNeGyyWb2uJltDp/HZN5KM7vLzPaY2cuRZf3GYoGvht/ji2Z26hjHdbOZNYff\n2zozuyiy7oYwro1mdsEoxjXHzJ40sw1mtt7M/jpcXgzf2UCxFfx7Gwsq5wPGUZRlfJDYCv7/dVyX\nc3cviQfB/NOvAscAVcALwOICx/Q6MLXPsn8Grg9fXw98aYxieR9wKvDyULEAFxHMuW3Au4Fnxjiu\nm4G/62fbxeG/azWwIPz3Lh+luGYCp4avJwCbwr9fDN/ZQLEV/Hsbg//HKucDx1GUZXyQ2Ar+/3U8\nl/NSqrmfAWxx99fcvRe4D7i4wDH152LgO+Hr7wCXjMUfdfdfAweyjOVi4LseWA00mtnMMYxrIBcD\n97l7j7tvBbYQ/LuPRly73P258HUb8Aowi+L4zgaKbSBj9r2NAZXzARRrGR8ktoGonA8e20By+t5K\nKbnPAt6IvN/B4F/EWHDg52a21syuDJfNcPdd4es3gRmFCW3QWIrhu7w6bPa6K9KkWZC4zGw+cArw\nDEX2nfWJDYroexslxfhZirmcF9X/134Uzf/X8VbOSym5F6M/cvdTgQuBT5vZ+6IrPWhLKYp7DYsp\nFuDrwLHAEmAX8L8KFYiZNQA/Av7G3Vuj6wr9nfUTW9F8b+NMSZTzYokjomj+v47Hcl5Kyb0ZmBN5\nPztcVjDu3hw+7wEeImgi2Z1pxgmf9xQuwgFjKeh36e673T3l7mngGxxpWhrTuMyskqBQ3ePuD4aL\ni+I76y+2YvneRlnRfZYiL+dF8f+1P8Xy/3W8lvNSSu7PAgvNbIGZVQHLgVWFCsbM6s1sQuY1cD7w\nchjTynCzlcAjhYkQBollFXB52DP03cChSBPVqOtzDesjBN9bJq7lZlZtZguAhcDvRykGA74FvOLu\nX46sKvh3NlBsxfC9jQGV89wU/P/rQIrh/+u4Luf56PU3Vg+CnoybCHoJ3ljgWI4h6Ln4ArA+Ew8w\nBfglsBn4BTB5jOK5l6AJJ0FwLeYTA8VC0BP0jvB7fAlYOsZxfS/8uy+G/2FnRra/MYxrI3DhKMb1\nRwRNcS8C68LHRUXynQ0UW8G/tzH6v6xy3n8sRVnGB4mt4P9fx3M51/CzIiIiMVNKzfIiIiKSBSV3\nERGRmFFyFxERiRkldxERkZhRchcREYkZJfdxysxuDGciejGceehMM/sbM6srdGwikh8q5+OXboUb\nh8zsPcCXgbPdvcfMphLMwPU7gvs69xU0QBEZMZXz8U019/FpJrDP3XsAwkJ+KXA08KSZPQlgZueb\n2dNm9pyZ/Xs4BnJmfut/tmCO69+b2XHh8v9qZi+b2Qtm9uvCfDQRCamcj2OquY9DYeH9DVBHMDrT\nD939V2b2OuEZfXiW/yDBKEgdZnYdUO3ut4TbfcPdbzWzy4HL3P1PzOwlYJm7N5tZo7u3FOQDiojK\n+Tinmvs45O7twGnAlcBe4IdmdkWfzd4NLAZ+a2brCMZfnhdZf2/k+T3h698Cd5vZJ4Hy0YleRLKh\ncj6+VRQ6ACkMd08BTwFPhWfiK/tsYsDj7r5ioEP0fe3uV5nZmcCHgLVmdpq7789v5CKSLZXz8Us1\n93HIzI43s4WRRUuAbUAbMCFctho4K3Kdrd7MFkX2+Vjk+elwm2Pd/Rl3v4mgphCdnlBExpDK+fim\nmvv41AB8zcwagSSwhaDpbgXwMzPb6e7nhE1495pZdbjf5wlm6wJoMrMXgZ5wP4D/Ef6YGMGMSy+M\nyacRkf6onI9j6lAnOYt2yCl0LCIyOlTOS5ua5UVERGJGNXcREZGYUc1dREQkZpTcRUREYkbJXURE\nJGaU3EVERGJGyV1ERCRm/i+lo6zMIX4KhwAAAABJRU5ErkJggg==\n",
            "text/plain": [
              "\u003cFigure size 504x252 with 2 Axes\u003e"
            ]
          },
          "metadata": {
            "tags": []
          },
          "output_type": "display_data"
        }
      ],
      "source": [
        "num_steps = 250\n",
        "num_examples = x_train.shape[0]\n",
        "\n",
        "train_objectives = []\n",
        "test_objectives = []\n",
        "\n",
        "for ii in range(num_steps):\n",
        "  # Indices for current batch; cycle back once we reach the end of stream.\n",
        "  batch_indices = np.arange(ii * batch_size, (ii + 1) * batch_size)\n",
        "  batch_indices = [ind % num_examples for ind in batch_indices]\n",
        "\n",
        "  # First run update ops, and then gradient update.\n",
        "  update_ops_fn()\n",
        "  optimizer.minimize(loss_fn, var_list=var_list)\n",
        "\n",
        "  # Record train and test objectives once every 10 steps.\n",
        "  if ii % 10 == 0:\n",
        "    train_objectives.append(pr_auc(model, x_train, y_train))\n",
        "    test_objectives.append(pr_auc(model, x_test, y_test))\n",
        "\n",
        "# Plot training and test objective as a function of steps.\n",
        "fig, ax = plt.subplots(1, 2, figsize=(7, 3.5))\n",
        "ax[0].plot(np.arange(1, num_steps + 1, 10), train_objectives)\n",
        "ax[0].set_title('Train PR-AUC')\n",
        "ax[0].set_xlabel('Steps')\n",
        "ax[1].plot(np.arange(1, num_steps + 1, 10), test_objectives)\n",
        "ax[1].set_title('Test PR-AUC')\n",
        "ax[1].set_xlabel('Steps')\n",
        "fig.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BTGOcqe0pO5a"
      },
      "source": [
        "# (ii) PR-AUC Training with Custom Estimators"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uppYKix2-SdA"
      },
      "source": [
        "We next show how one can use TFCO to optimize PR-AUC using custom tf.Estimators.\n",
        "\n",
        "We first create `feature_columns` to convert the dataset into a format that can be processed by an estimator."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_klY8Dqueag_"
      },
      "outputs": [],
      "source": [
        "feature_columns = []\n",
        "for feature_name in x_train_df.columns:\n",
        "  feature_columns.append(\n",
        "      tf.feature_column.numeric_column(feature_name, dtype=tf.float32))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jyaRvuOXAGt-"
      },
      "source": [
        "We next construct the input functions that return the data to be used by the estimator for training/evaluation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HoEGUpg9pTdD"
      },
      "outputs": [],
      "source": [
        "def make_input_fn(\n",
        "    data_df, label_df, num_epochs=10, shuffle=True, batch_size=32):\n",
        "  def input_fn():\n",
        "    ds = tf.data.Dataset.from_tensor_slices((dict(data_df), label_df))\n",
        "    if shuffle:\n",
        "      ds = ds.shuffle(1000)\n",
        "    ds = ds.batch(batch_size).repeat(num_epochs)\n",
        "    return ds\n",
        "  return input_fn\n",
        "\n",
        "train_input_fn = make_input_fn(x_train_df, y_train_df, num_epochs=25)\n",
        "test_input_fn = make_input_fn(x_test_df, y_test_df, num_epochs=1, shuffle=False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QSw7IHgKA5NC"
      },
      "source": [
        "We then write the model function that is used by the estimator to create the model, loss, optimizers and metrics."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fUD8PAddptO7"
      },
      "outputs": [],
      "source": [
        "def make_model_fn(feature_columns):\n",
        "  # Returns model_fn.\n",
        "\n",
        "  def model_fn(features, labels, mode):\n",
        "    # Create model from features.\n",
        "    layers = []\n",
        "    layers.append(tf.keras.layers.DenseFeatures(feature_columns))\n",
        "    layers.append(tf.keras.layers.Dense(1))\n",
        "    model = tf.keras.Sequential(layers)\n",
        "    logits = model(features)\n",
        "\n",
        "    # Baseline cross-entropy loss.\n",
        "    baseline_loss_fn = tf.keras.losses.BinaryCrossentropy(from_logits=True)\n",
        "    baseline_loss = baseline_loss_fn(labels, logits)\n",
        "\n",
        "    # As a slight variant from the above previous training, we will optimize a\n",
        "    # weighted combination of PR-AUC and the baseline loss.\n",
        "    baseline_coef = 0.2\n",
        "\n",
        "    train_op = None\n",
        "    if mode == tf.estimator.ModeKeys.TRAIN:\n",
        "      # Set up PR-AUC optimization problem.\n",
        "      # Create context with labels and predictions.\n",
        "      context = tfco.rate_context(logits, labels)\n",
        "\n",
        "      # Create optimization problem with PR-AUC as the objective. The library\n",
        "      # expects a minimization objective, so we negate the PR-AUC. We optimize\n",
        "      # a convex combination of (negative) PR-AUC and the baseline loss (wrapped\n",
        "      # in a rate object).\n",
        "      pr_auc_rate = tfco.pr_auc(\n",
        "          context, bins=10, penalty_loss=tfco.SoftmaxCrossEntropyLoss())\n",
        "      problem = tfco.RateMinimizationProblem((1 - baseline_coef) *\n",
        "                                             (-pr_auc_rate) + baseline_coef *\n",
        "                                             tfco.wrap_rate(baseline_loss))\n",
        "\n",
        "      # Create Lagrangian loss for `problem`. What we get back is a loss\n",
        "      # function, a nullary function that returns a list of update_ops that\n",
        "      # need to be run before every gradient update, and the Lagrange\n",
        "      # multipliers maintained internally by the loss.\n",
        "      # The argument `dual_scale` is a hyper-parameter that specifies the\n",
        "      # relative importance placed on updates on the Lagrange multipliers.\n",
        "      loss_fn, update_ops_fn, multipliers = tfco.create_lagrangian_loss(\n",
        "          problem, dual_scale=1.0)\n",
        "\n",
        "      # Set up optimizer and the list of variables to optimize the loss.\n",
        "      optimizer = tf.keras.optimizers.Adagrad(learning_rate=0.1)\n",
        "      optimizer.iterations = tf.compat.v1.train.get_or_create_global_step()\n",
        "\n",
        "      # Get minimize op and group with update_ops.\n",
        "      var_list = (\n",
        "          model.trainable_weights + list(problem.trainable_variables) +\n",
        "          [multipliers])\n",
        "      minimize_op = optimizer.get_updates(loss_fn(), var_list)\n",
        "      update_ops = update_ops_fn()\n",
        "      train_op = tf.group(*update_ops, minimize_op)\n",
        "\n",
        "    # Evaluate PR-AUC.\n",
        "    pr_auc_metric = tf.keras.metrics.AUC(curve='PR')\n",
        "    pr_auc_metric.update_state(labels, tf.sigmoid(logits))\n",
        "\n",
        "    # We do not use the Lagrangian loss for evaluation/bookkeeping\n",
        "    # purposes as it depends on some internal variables that may not be\n",
        "    # set properly during evaluation time. We instead pass loss=baseline_loss.\n",
        "    return tf.estimator.EstimatorSpec(\n",
        "        mode=mode,\n",
        "        predictions=logits,\n",
        "        loss=baseline_loss,\n",
        "        train_op=train_op,\n",
        "        eval_metric_ops={'PR-AUC': pr_auc_metric})\n",
        "\n",
        "  return model_fn"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4kc3_ftIBC6Y"
      },
      "source": [
        "We are now ready to train the estimator."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 676
        },
        "id": "TImVz7WMp-Nb",
        "outputId": "d8f64525-dba6-46b2-a7e6-39d05c0a5d85"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Using default config.\n",
            "INFO:tensorflow:Using config: {'_model_dir': 'tfco_tmp', '_tf_random_seed': None, '_save_summary_steps': 100, '_save_checkpoints_steps': None, '_save_checkpoints_secs': 600, '_session_config': allow_soft_placement: true\n",
            "graph_options {\n",
            "  rewrite_options {\n",
            "    meta_optimizer_iterations: ONE\n",
            "  }\n",
            "}\n",
            ", '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': 100, '_train_distribute': None, '_device_fn': None, '_protocol': None, '_eval_distribute': None, '_experimental_distribute': None, '_experimental_max_worker_delay_secs': None, '_session_creation_timeout_secs': 7200, '_service': None, '_cluster_spec': ClusterSpec({}), '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': '', '_evaluation_master': '', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1}\n",
            "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/resource_variable_ops.py:1666: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.\n",
            "Instructions for updating:\n",
            "If using Keras pass *_constraint arguments to layers.\n",
            "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/training/training_util.py:236: Variable.initialized_value (from tensorflow.python.ops.variables) is deprecated and will be removed in a future version.\n",
            "Instructions for updating:\n",
            "Use Variable.read_value. Variables in 2.X are initialized automatically both in eager and graph (inside tf.defun) contexts.\n",
            "INFO:tensorflow:Calling model_fn.\n",
            "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/optimizer_v2/adagrad.py:106: calling Constant.__init__ (from tensorflow.python.ops.init_ops) with dtype is deprecated and will be removed in a future version.\n",
            "Instructions for updating:\n",
            "Call initializer instance with the dtype argument instead of passing it to the constructor\n",
            "INFO:tensorflow:Done calling model_fn.\n",
            "INFO:tensorflow:Create CheckpointSaverHook.\n",
            "INFO:tensorflow:Graph was finalized.\n",
            "INFO:tensorflow:Running local_init_op.\n",
            "INFO:tensorflow:Done running local_init_op.\n",
            "INFO:tensorflow:Calling checkpoint listeners before saving checkpoint 0...\n",
            "INFO:tensorflow:Saving checkpoints for 0 into tfco_tmp/model.ckpt.\n",
            "INFO:tensorflow:Calling checkpoint listeners after saving checkpoint 0...\n",
            "INFO:tensorflow:loss = 0.66444886, step = 0\n",
            "INFO:tensorflow:global_step/sec: 60.6059\n",
            "INFO:tensorflow:loss = 0.24953337, step = 100 (1.651 sec)\n",
            "INFO:tensorflow:global_step/sec: 225.562\n",
            "INFO:tensorflow:loss = 0.26060176, step = 200 (0.443 sec)\n",
            "INFO:tensorflow:Calling checkpoint listeners before saving checkpoint 250...\n",
            "INFO:tensorflow:Saving checkpoints for 250 into tfco_tmp/model.ckpt.\n",
            "INFO:tensorflow:Calling checkpoint listeners after saving checkpoint 250...\n",
            "INFO:tensorflow:Loss for final step: 0.17999612.\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "\u003ctensorflow_estimator.python.estimator.estimator.EstimatorV2 at 0x7f1f61287cf8\u003e"
            ]
          },
          "execution_count": 21,
          "metadata": {
            "tags": []
          },
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# Create a temporary model directory.\n",
        "model_dir = \"tfco_tmp\"\n",
        "if os.path.exists(model_dir):\n",
        "  shutil.rmtree(model_dir)\n",
        "\n",
        "# Train estimator.\n",
        "estimator_lin = tf.estimator.Estimator(\n",
        "    make_model_fn(feature_columns), model_dir=model_dir)\n",
        "estimator_lin.train(train_input_fn, steps=250) "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iYkWYCgvBGJA"
      },
      "source": [
        "Finally, we evaluate the trained model on the test set."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 230
        },
        "id": "iiB6oG2fqC7S",
        "outputId": "043ce405-99f7-48a7-d2de-d60708017a30"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Calling model_fn.\n",
            "INFO:tensorflow:Done calling model_fn.\n",
            "INFO:tensorflow:Starting evaluation at 2020-04-08T18:36:00Z\n",
            "INFO:tensorflow:Graph was finalized.\n",
            "INFO:tensorflow:Restoring parameters from tfco_tmp/model.ckpt-250\n",
            "INFO:tensorflow:Running local_init_op.\n",
            "INFO:tensorflow:Done running local_init_op.\n",
            "INFO:tensorflow:Inference Time : 0.56312s\n",
            "INFO:tensorflow:Finished evaluation at 2020-04-08-18:36:01\n",
            "INFO:tensorflow:Saving dict for global step 250: PR-AUC = 0.7938361, global_step = 250, loss = 0.38235107\n",
            "INFO:tensorflow:Saving 'checkpoint_path' summary for global step 250: tfco_tmp/model.ckpt-250\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "{'PR-AUC': 0.7938361, 'global_step': 250, 'loss': 0.38235107}"
            ]
          },
          "execution_count": 22,
          "metadata": {
            "tags": []
          },
          "output_type": "execute_result"
        }
      ],
      "source": [
        "estimator_lin.evaluate(test_input_fn)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6xU1eK4GVKd_"
      },
      "source": [
        "## Closing Remarks"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zjzlHyFnC7d3"
      },
      "source": [
        "Before closing, we point out that there are three main hyper-paramters you may want to tune to improve the PR-AUC training:\n",
        "\n",
        "- `learning_rate`\n",
        "- `dual_scale`\n",
        "- `baseline_coeff`\n",
        "\n",
        "You may also be interested in exploring helpers for other similar metrics that TFCO allows you to optimize:\n",
        "- `tfco.precision_at_recall`\n",
        "- `tfco.recall_at_precision`\n",
        "- `tfco.inverse_precision_at_recall`"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "last_runtime": {
        "build_target": "",
        "kind": "local"
      },
      "name": "PRAUC_training.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.5.4rc1"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}