{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ziwrREgJx5ov"
      },
      "outputs": [],
      "source": [
        "dname = 'dataset1'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EVoWoTNUbK2G"
      },
      "outputs": [],
      "source": [
        "TR_ITER = 10001\n",
        "TR_AVG_ITER = 2500"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "33zdXjdb2OZ0"
      },
      "outputs": [],
      "source": [
        "EXP_NAME = 'testing_pickle1'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bo8IFgCbT80v"
      },
      "outputs": [],
      "source": [
        "# Path to load data fro\n",
        "BASE_DIR = ''"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "we_9ZHOJ1u77"
      },
      "outputs": [],
      "source": [
        "SAVE_DIR = ''"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1s_i6Hr6a3QO"
      },
      "source": [
        "# Func"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Cu709diOzAUj"
      },
      "outputs": [],
      "source": [
        "from datetime import datetime\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eTgFZHaMEg4U"
      },
      "outputs": [],
      "source": [
        "import time\n",
        "start = time.time()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MHh5CY-k2Cq3"
      },
      "outputs": [],
      "source": [
        "import os\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EXxipvmSWDIi"
      },
      "outputs": [],
      "source": [
        "%matplotlib inline\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import scipy.stats\n",
        "import pandas as pd"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FDwgr2mvrr43"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "import tensorflow_probability as tfp"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MiA1dcJqpTKA"
      },
      "outputs": [],
      "source": [
        "from xgboost import XGBClassifier\n",
        "from xgboost import XGBRegressor\n",
        "from scipy.linalg import null_space\n",
        "from sklearn.cross_decomposition import CCA\n",
        "from sklearn.ensemble import RandomForestClassifier\n",
        "from scipy.spatial.distance import cdist\n",
        "from sklearn.linear_model import LogisticRegression"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZvQxVPBdwJfc"
      },
      "outputs": [],
      "source": [
        "from scipy.integrate import trapezoid\n",
        "from sklearn import metrics\n",
        "\n",
        "def auprc_threshs(confidences, Y, threshs, as_percentages=True):\n",
        "  areas = []\n",
        "  precisions, recalls, thresholds = metrics.precision_recall_curve(Y, confidences)\n",
        "  for t in threshs:\n",
        "    if t < 1.0:\n",
        "      interp_prec_t = np.interp(t, recalls[::-1], precisions[::-1])\n",
        "      recalls_thresh = np.concatenate(([t], recalls[recalls<t]))\n",
        "      precisions_thresh = np.concatenate(([interp_prec_t], precisions[recalls<t]))\n",
        "    else:\n",
        "      recalls_thresh = recalls\n",
        "      precisions_thresh = precisions\n",
        "    area_t = metrics.auc(recalls_thresh, precisions_thresh)\n",
        "    if as_percentages:\n",
        "      area_t = area_t/t\n",
        "    areas.append(area_t)\n",
        "  return np.array(areas)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8Id7d7pQs3Wd"
      },
      "outputs": [],
      "source": [
        "external_colstart = 4 if dname in set(['dataset1', 'dataset2', 'dataset3', 'dataset4']) else 2  # SHOULD BE 4 for 'dataset1' ... 'dataset4' otherwise 2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dguz7FBcW5mq"
      },
      "outputs": [],
      "source": [
        "with open(os.path.join(BASE_DIR, dname, 'train.csv'), 'r') as f:\n",
        "  dataX = np.float32(np.array([line.strip().split(',')[2:] for line in f])[1:])\n",
        "\n",
        "with open(os.path.join(BASE_DIR, dname, 'train.csv'), 'r') as f:\n",
        "  dataY = np.float32(np.array([line.strip().split(',')[1] for line in f])[1:])\n",
        "\n",
        "X = dataX\n",
        "Y = dataY"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "48e-tD-8W5mq"
      },
      "outputs": [],
      "source": [
        "with open(os.path.join(BASE_DIR, dname, 'val_id.csv'), 'r') as f:\n",
        "  dataX_val = np.float32(np.array([line.strip().split(',')[2:] for line in f])[1:])\n",
        "\n",
        "with open(os.path.join(BASE_DIR, dname, 'val_id.csv'), 'r') as f:\n",
        "  dataY_val = np.float32(np.array([line.strip().split(',')[1] for line in f])[1:])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ww0XXrrvW5mq"
      },
      "outputs": [],
      "source": [
        "with open(os.path.join(BASE_DIR, dname, 'test_id.csv'), 'r') as f:\n",
        "  dataX_tst = np.float32(np.array([line.strip().split(',')[2:] for line in f])[1:])\n",
        "\n",
        "with open(os.path.join(BASE_DIR, dname, 'test_id.csv'), 'r') as f:\n",
        "  dataY_tst = np.float32(np.array([line.strip().split(',')[1] for line in f])[1:])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-Zt3t0tqW5mr"
      },
      "outputs": [],
      "source": [
        "with open(os.path.join(BASE_DIR, dname, 'val_ood.csv'), 'r') as f:\n",
        "  external_X = np.float32(np.array([line.strip().split(',')[external_colstart:] for line in f])[1:])\n",
        "\n",
        "with open(os.path.join(BASE_DIR, dname, 'val_ood.csv'), 'r') as f:\n",
        "  external_Y = np.float32(np.array([line.strip().split(',')[1] for line in f])[1:])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vaojUtEdY2p2"
      },
      "outputs": [],
      "source": [
        "X = dataX\n",
        "Y = dataY\n",
        "\n",
        "# standardize the data\n",
        "mu_x = np.mean(X, 0, keepdims=True)\n",
        "#sigma_x = np.std(X, 0, keepdims=True)\n",
        "sigma_x = np.ones_like(mu_x)\n",
        "X = (X-mu_x)/sigma_x\n",
        "\n",
        "# PCA\n",
        "XtX = np.matmul(X.T, X)\n",
        "eigvals, eigvecs = np.linalg.eig(XtX)\n",
        "sortinds = np.argsort(-eigvals)\n",
        "V = eigvecs[:, sortinds]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "T8FwitHVZ0-W"
      },
      "outputs": [],
      "source": [
        "external_X = (external_X-mu_x)/sigma_x\n",
        "# held_X = (dataX_val-mu_x)/sigma_x\n",
        "# held_Y = dataY_val\n",
        "held_X = (np.concatenate((dataX_val, dataX_tst), 0)-mu_x)/sigma_x\n",
        "held_Y = np.concatenate((dataY_val, dataY_tst), 0)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2PuqxRwXJICh"
      },
      "outputs": [],
      "source": [
        "NDIRS = 128  # 256\n",
        "X = X@V[:, :NDIRS]\n",
        "held_X = held_X@V[:, :NDIRS]\n",
        "external_X = external_X@V[:, :NDIRS]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_U_l31CacP8F"
      },
      "outputs": [],
      "source": [
        "def get_preds(x, ests, mus=None, dirs=None, use_proba=False):\n",
        "  if mus is None:\n",
        "    mus = [0.0 for _ in range(len(ests))]\n",
        "  if dirs is not None:\n",
        "    if use_proba:\n",
        "      pred_array = [e.predict_proba((x-mu)@d)[:, 1] for mu, d, e in zip(mus, dirs, ests)]\n",
        "    else:\n",
        "      pred_array = [e.predict((x-mu)@d) for mu, d, e in zip(mus, dirs, ests)]\n",
        "  else:\n",
        "    if use_proba:\n",
        "      pred_array = [e.predict_proba((x-mu))[:, 1] for mu, e in zip(mus, ests)]\n",
        "    else:\n",
        "      pred_array = [e.predict((x-mu)) for mu, e in zip(mus, ests)]\n",
        "  return np.stack(pred_array, -1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s-1Nkhj3gdfx"
      },
      "outputs": [],
      "source": [
        "def aggregate_preds(preds, use_bernoulli_std=False):\n",
        "  mean_pred = np.mean(preds, 1, keepdims=False)\n",
        "  if use_bernoulli_std:\n",
        "    std_pred = mean_pred*(1-mean_pred)\n",
        "  else:\n",
        "    std_pred = np.std(preds, 1, keepdims=False)\n",
        "  return np.float32(mean_pred>0.5), mean_pred, std_pred"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NDuawjDXkq8C"
      },
      "source": [
        "Train on random directions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "31XMRekNlXBt"
      },
      "outputs": [],
      "source": [
        "from sklearn.metrics import precision_score"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KpPp9Psmlmk0"
      },
      "outputs": [],
      "source": [
        "def divide_points(percentage_threshold, X, y, index):\n",
        "  dot_products = (X @ X[index, None, :].T)[:, 0]\n",
        "  thresh = np.quantile(dot_products, 1.-percentage_threshold)\n",
        "  # id_x, id_y, ood_x, ood_y\n",
        "  return (X[dot_products<thresh], y[dot_products<thresh],\n",
        "    X[dot_products>=thresh], y[dot_products>=thresh])\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WnAOruvQp_iL"
      },
      "outputs": [],
      "source": [
        "from abc import ABC, abstractmethod"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0PlRAAV5p_iL"
      },
      "outputs": [],
      "source": [
        "class Ensemble(ABC):\n",
        "  \"\"\" General Ensembling Class \"\"\"\n",
        "\n",
        "  @abstractmethod\n",
        "  def __init__(self, *args, **kwargs):\n",
        "    self.nmodels = kwargs.get('nmodels', 256)\n",
        "\n",
        "  @property\n",
        "  @abstractmethod\n",
        "  def estimators(self):\n",
        "    pass\n",
        "\n",
        "  @abstractmethod\n",
        "  def add_estimator(self, est):\n",
        "    pass\n",
        "\n",
        "  @abstractmethod\n",
        "  def base_model_generator(self):\n",
        "    # Returns something trainable with fit(X,Y)\n",
        "    pass\n",
        "\n",
        "  def random_sample(self, X, Y):\n",
        "    # Potentially subsample X and Y\n",
        "    return X, Y\n",
        "\n",
        "  def fit(self, X, Y, verbose=True):\n",
        "    # Returns something trainable with fit(X,Y)\n",
        "    for t in range(self.nmodels):\n",
        "      sX, sY = self.random_sample(X, Y)\n",
        "      est = self.base_model_generator()\n",
        "      if verbose:\n",
        "        print('Fitting Model {}'.format(t))\n",
        "      est.fit(sX, sY)\n",
        "      self.add_estimator(est)\n",
        "\n",
        "  def predictions(self, x):\n",
        "    # Assumes estimators return (N, ) shape predict\n",
        "    pred_array = [e.predict(x) for e in self.estimators]\n",
        "    return np.stack(pred_array, -1)\n",
        "\n",
        "  def predict(self, x=None, preds=None, return_only_prediction=True):\n",
        "    assert (x is not None) or (preds is not None)\n",
        "    if preds is None:\n",
        "      preds = self.predictions(x)\n",
        "\n",
        "    mean_pred = np.mean(preds, 1, keepdims=False)\n",
        "    std_pred = np.std(preds, 1, keepdims=False)\n",
        "\n",
        "    if return_only_prediction:\n",
        "      return np.float32(mean_pred>0.5)\n",
        "    return np.float32(mean_pred>0.5), mean_pred, std_pred"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "keuiurs4p_iM"
      },
      "outputs": [],
      "source": [
        "class ProjectionEnsemble(Ensemble):\n",
        "  \"\"\" General Ensembling Class With Projections of Dimensions\"\"\"\n",
        "\n",
        "  @abstractmethod\n",
        "  def __init__(self, *args, **kwargs):\n",
        "    super(ProjectionEnsemble, self).__init__(*args, **kwargs)\n",
        "\n",
        "  @property\n",
        "  @abstractmethod\n",
        "  def directions(self):\n",
        "    pass\n",
        "\n",
        "  @property\n",
        "  @abstractmethod\n",
        "  def means(self):\n",
        "    pass\n",
        "\n",
        "  def predictions(self, x):\n",
        "    # Assumes estimators return (N, ) shape predict\n",
        "    pred_array = [\n",
        "      e.predict((x-mu)@d)\n",
        "      for mu, d, e in zip(self.means, self.directions, self.estimators)\n",
        "    ]\n",
        "    return np.stack(pred_array, -1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lKdrmysap_iM"
      },
      "outputs": [],
      "source": [
        "class PCAProjectionEnsemble(ProjectionEnsemble):\n",
        "  \"\"\" General Ensembling Class With Projections of Dimensions\"\"\"\n",
        "\n",
        "  def __init__(self, *args, **kwargs):\n",
        "    super(ProjectionEnsemble, self).__init__(*args, **kwargs)\n",
        "    self.ndirs = kwargs.get('ndirs', 8)\n",
        "    self.Vpca = kwargs['Vpca']  # d x k\n",
        "    self.Vorth = kwargs.get('Vorth', 8)\n",
        "    self.projection_split = kwargs.get('projection_split', False)\n",
        "    self.split_perc = kwargs.get('split_perc', .5)\n",
        "    self._splits_rand = []\n",
        "    self._dirs_inds = []\n",
        "    self._estimators = []\n",
        "\n",
        "  @property\n",
        "  def directions(self):\n",
        "    return [self.Vpca[:, dinds] for dinds in self._dirs_inds]\n",
        "\n",
        "  @property\n",
        "  def means(self):\n",
        "    return [np.zeros((1, self.Vpca.shape[0]), dtype=np.float32)] * len(self._dirs_inds)\n",
        "\n",
        "  @property\n",
        "  def estimators(self):\n",
        "    return self._estimators\n",
        "\n",
        "  def add_estimator(self, est):\n",
        "    self._estimators.append(est)\n",
        "\n",
        "  def base_model_generator(self):\n",
        "    # Returns something trainable with fit(X,Y)\n",
        "    # TODO: make init argument\n",
        "    return XGBClassifier(n_estimators=40)\n",
        "\n",
        "  def random_sample(self, X, Y, XV=None):\n",
        "    # Potentially subsample X and Y\n",
        "    N_tot, d = X.shape\n",
        "\n",
        "    if self.projection_split:\n",
        "      newdir = np.random.randn(d, 1)\n",
        "      newdir = newdir/np.sqrt(np.sum(newdir**2))\n",
        "      projs = (X@newdir)[:, 0]\n",
        "      medianp = np.median(projs)\n",
        "      in_range = projs <= medianp\n",
        "    else:\n",
        "      in_range = np.random.rand(N_tot) <= self.split_perc\n",
        "    self._splits_rand.append(in_range)\n",
        "\n",
        "    Y_train = Y[in_range]\n",
        "    rprm = np.random.permutation(self.Vpca.shape[1])\n",
        "    self._dirs_inds.append(rprm[:self.ndirs])\n",
        "    if XV is not None:\n",
        "      X_train = XV[in_range][:, self._dirs_inds[-1]]\n",
        "    else:\n",
        "      X_train = X[in_range]@self.Vpca[:, self._dirs_inds[-1]]\n",
        "\n",
        "    return X_train, Y_train\n",
        "\n",
        "  def fit(self, X, Y, verbose=True):\n",
        "    XV = X@self.Vpca\n",
        "    for t in range(self.nmodels):\n",
        "      sX, sY = self.random_sample(X, Y, XV)\n",
        "      est = self.base_model_generator()\n",
        "      if verbose:\n",
        "        print('Fitting Model {}'.format(t))\n",
        "      est.fit(sX, sY)\n",
        "      self.add_estimator(est)\n",
        "\n",
        "  def predictions(self, x):\n",
        "    xV = x@self.Vpca\n",
        "    # Assumes estimators return (N, ) shape predict\n",
        "    pred_array = [\n",
        "      e.predict(xV[:, di])\n",
        "      for di, e in zip(self._dirs_inds, self.estimators)\n",
        "    ]\n",
        "    return np.stack(pred_array, -1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Pop3TGAKo-hJ"
      },
      "outputs": [],
      "source": [
        "# with open(os.path.join(BASE_DIR +'/' + EXP_NAME + '/' + dname, 'pca_projens_1024.p'), 'rb') as f:\n",
        "#     pca_proj = pickle.load(f)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VhsbsS3b00fZ"
      },
      "outputs": [],
      "source": [
        "def knn_gibbs(Xsamp, X, nn_featinds, k):\n",
        "  Xf = X[:, nn_featinds]\n",
        "  Xsf = Xsamp[:, nn_featinds]\n",
        "  Xfnorm2 = np.sum(Xf**2, keepdims=True)\n",
        "  Xsfnorm2 = np.sum(Xsf**2, keepdims=True)\n",
        "  dotsf = Xsf@Xf.T\n",
        "  distsf = Xsfnorm2 + Xfnorm2.T - 2*dotsf\n",
        "  sortdistsf = np.argsort(distsf, axis=1)\n",
        "  # neighs = sortdistsf[:, np.random.choice(k, size=(X.shape[0], ))]\n",
        "  neighs = sortdistsf[np.arange(Xsamp.shape[0]), np.random.choice(k, size=(X.shape[0], ))]\n",
        "  Xsamp = np.copy(Xsamp)\n",
        "  Xsamp[:, np.logical_not(nn_featinds)] = X[neighs, :][:, np.logical_not(nn_featinds)]\n",
        "  return Xsamp\n",
        "\n",
        "def gibbs_sample(X, k, nrounds):\n",
        "  Xsamp = X\n",
        "  for r in range(nrounds):\n",
        "    featinds = np.random.rand(X.shape[1])<.5\n",
        "    Xsamp = knn_gibbs(Xsamp, X, featinds, k)\n",
        "    Xsamp = knn_gibbs(Xsamp, X, np.logical_not(featinds), k)\n",
        "  return Xsamp\n",
        "\n",
        "def permute_expand(X):\n",
        "  N, d = X.shape\n",
        "  return np.stack([X[np.random.permutation(N), i] for i in range(d)], -1)\n",
        "\n",
        "def convex_expand(X):\n",
        "  N, d = X.shape\n",
        "  X1 = X[np.random.permutation(N), :]\n",
        "  alphas = np.random.rand(N, 1)\n",
        "  return alphas*X + (1-alphas)*X1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nMFXxF3r1Ux8"
      },
      "outputs": [],
      "source": [
        "expand_funcs = [\n",
        "    # lambda x: gibbs_sample(x, 40, 2),\n",
        "    # lambda x: permute_expand(x),\n",
        "    # lambda x: convex_expand(x),\n",
        "    lambda x: x*(1+np.abs(np.random.randn(x.shape[0], 1)))\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wbZJQblJE35k"
      },
      "outputs": [],
      "source": [
        "class DatasetDiscriminator(tf.keras.Model):\n",
        "  def __init__(self, dim, nhidden_layers=2, hidden_size=256, out_size=64, net=None):\n",
        "    super(DatasetDiscriminator, self).__init__()\n",
        "\n",
        "    self.dim = dim\n",
        "    if net is not None:\n",
        "      self.net = net\n",
        "    else:\n",
        "      layers = [\n",
        "        tf.keras.layers.Dense(\n",
        "            hidden_size, activation=tf.nn.elu) for _ in range(nhidden_layers)] + \\\n",
        "            [tf.keras.layers.Dense(out_size)]\n",
        "      self.net = tf.keras.Sequential(layers)\n",
        "\n",
        "    # self.beta = tf.Variable(\n",
        "    #  tf.random.normal([out_size, 1], stddev=0.01), name='beta')\n",
        "\n",
        "  def get_features(self, x_in):\n",
        "    # print(x_in.shape)\n",
        "    feats = self.net(x_in)\n",
        "    # print(feats.shape)\n",
        "    return feats\n",
        "\n",
        "  def __call__(self, x_in):\n",
        "    return self.get_features(x_in)\n",
        "\n",
        "  def predictions(self, x_in):\n",
        "    return tf.nn.sigmoid(self.get_features(x_in))\n",
        "\n",
        "\n",
        "  def predictions_via_logits(self, x_in):\n",
        "    return tf.nn.sigmoid(np.mean(self.get_features(x_in), -1))\n",
        "\n",
        "  # def get_logits(self, x_in):\n",
        "  #   return tf.matmul(self.get_features(x_in), self.beta)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gbIGqE5nr86Z"
      },
      "outputs": [],
      "source": [
        "def smooth_maximum(v, temp=1.0):\n",
        "  softmaxv = tf.nn.softmax(v/temp, axis=-1)\n",
        "  return tf.reduce_sum(v*softmaxv, axis=-1, keepdims=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hMSLzFqAr9vb"
      },
      "outputs": [],
      "source": [
        "def sigmoid_entropy(logits, normalize=False):\n",
        "  NORM = 0.6931471805599453  # -np.log(.5)\n",
        "  p = tf.nn.sigmoid(logits)\n",
        "  entropy = p*tf.nn.softplus(-logits) + (1.-p)*tf.nn.softplus(logits)\n",
        "  if normalize:\n",
        "    return entropy/NORM\n",
        "  return entropy"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XAFjz88dGLeo"
      },
      "outputs": [],
      "source": [
        "def training_batch(data, llks, batchsize):\n",
        "  def generate_batch():\n",
        "    rprm_in = np.random.permutation(data.shape[0])\n",
        "    if len(llks.shape)==1:\n",
        "      return data[rprm_in[:batchsize], :], llks[rprm_in[:batchsize], None]\n",
        "    return data[rprm_in[:batchsize], :], llks[rprm_in[:batchsize]]\n",
        "  return generate_batch"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Gs9M6roQaIza"
      },
      "source": [
        "# NT EMOE"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Bjyvy8JzUNWo",
        "outputId": "2a6a0a0e-44b5-4c0f-908c-56158dd5a4b2"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Fitting Model 0\n",
            "Fitting Model 1\n",
            "Fitting Model 2\n",
            "Fitting Model 3\n",
            "Fitting Model 4\n",
            "Fitting Model 5\n",
            "Fitting Model 6\n",
            "Fitting Model 7\n",
            "Fitting Model 8\n",
            "Fitting Model 9\n",
            "Fitting Model 10\n",
            "Fitting Model 11\n",
            "Fitting Model 12\n",
            "Fitting Model 13\n",
            "Fitting Model 14\n",
            "Fitting Model 15\n",
            "Fitting Model 16\n",
            "Fitting Model 17\n",
            "Fitting Model 18\n",
            "Fitting Model 19\n",
            "Fitting Model 20\n",
            "Fitting Model 21\n",
            "Fitting Model 22\n",
            "Fitting Model 23\n",
            "Fitting Model 24\n",
            "Fitting Model 25\n",
            "Fitting Model 26\n",
            "Fitting Model 27\n",
            "Fitting Model 28\n",
            "Fitting Model 29\n",
            "Fitting Model 30\n",
            "Fitting Model 31\n",
            "Fitting Model 32\n",
            "Fitting Model 33\n",
            "Fitting Model 34\n",
            "Fitting Model 35\n",
            "Fitting Model 36\n",
            "Fitting Model 37\n",
            "Fitting Model 38\n",
            "Fitting Model 39\n",
            "Fitting Model 40\n",
            "Fitting Model 41\n",
            "Fitting Model 42\n",
            "Fitting Model 43\n",
            "Fitting Model 44\n",
            "Fitting Model 45\n",
            "Fitting Model 46\n",
            "Fitting Model 47\n",
            "Fitting Model 48\n",
            "Fitting Model 49\n",
            "Fitting Model 50\n",
            "Fitting Model 51\n",
            "Fitting Model 52\n",
            "Fitting Model 53\n",
            "Fitting Model 54\n",
            "Fitting Model 55\n",
            "Fitting Model 56\n",
            "Fitting Model 57\n",
            "Fitting Model 58\n",
            "Fitting Model 59\n",
            "Fitting Model 60\n",
            "Fitting Model 61\n",
            "Fitting Model 62\n",
            "Fitting Model 63\n",
            "Fitting Model 64\n",
            "Fitting Model 65\n",
            "Fitting Model 66\n",
            "Fitting Model 67\n",
            "Fitting Model 68\n",
            "Fitting Model 69\n",
            "Fitting Model 70\n",
            "Fitting Model 71\n",
            "Fitting Model 72\n",
            "Fitting Model 73\n",
            "Fitting Model 74\n",
            "Fitting Model 75\n",
            "Fitting Model 76\n",
            "Fitting Model 77\n",
            "Fitting Model 78\n",
            "Fitting Model 79\n",
            "Fitting Model 80\n",
            "Fitting Model 81\n",
            "Fitting Model 82\n",
            "Fitting Model 83\n",
            "Fitting Model 84\n",
            "Fitting Model 85\n",
            "Fitting Model 86\n",
            "Fitting Model 87\n",
            "Fitting Model 88\n",
            "Fitting Model 89\n",
            "Fitting Model 90\n",
            "Fitting Model 91\n",
            "Fitting Model 92\n",
            "Fitting Model 93\n",
            "Fitting Model 94\n",
            "Fitting Model 95\n",
            "Fitting Model 96\n",
            "Fitting Model 97\n",
            "Fitting Model 98\n",
            "Fitting Model 99\n",
            "Fitting Model 100\n",
            "Fitting Model 101\n",
            "Fitting Model 102\n",
            "Fitting Model 103\n",
            "Fitting Model 104\n",
            "Fitting Model 105\n",
            "Fitting Model 106\n",
            "Fitting Model 107\n",
            "Fitting Model 108\n",
            "Fitting Model 109\n",
            "Fitting Model 110\n",
            "Fitting Model 111\n",
            "Fitting Model 112\n",
            "Fitting Model 113\n",
            "Fitting Model 114\n",
            "Fitting Model 115\n",
            "Fitting Model 116\n",
            "Fitting Model 117\n",
            "Fitting Model 118\n",
            "Fitting Model 119\n",
            "Fitting Model 120\n",
            "Fitting Model 121\n",
            "Fitting Model 122\n",
            "Fitting Model 123\n",
            "Fitting Model 124\n",
            "Fitting Model 125\n",
            "Fitting Model 126\n",
            "Fitting Model 127\n",
            "Fitting Model 128\n",
            "Fitting Model 129\n",
            "Fitting Model 130\n",
            "Fitting Model 131\n",
            "Fitting Model 132\n",
            "Fitting Model 133\n",
            "Fitting Model 134\n",
            "Fitting Model 135\n",
            "Fitting Model 136\n",
            "Fitting Model 137\n",
            "Fitting Model 138\n",
            "Fitting Model 139\n",
            "Fitting Model 140\n",
            "Fitting Model 141\n",
            "Fitting Model 142\n",
            "Fitting Model 143\n",
            "Fitting Model 144\n",
            "Fitting Model 145\n",
            "Fitting Model 146\n",
            "Fitting Model 147\n",
            "Fitting Model 148\n",
            "Fitting Model 149\n",
            "Fitting Model 150\n",
            "Fitting Model 151\n",
            "Fitting Model 152\n",
            "Fitting Model 153\n",
            "Fitting Model 154\n",
            "Fitting Model 155\n",
            "Fitting Model 156\n",
            "Fitting Model 157\n",
            "Fitting Model 158\n",
            "Fitting Model 159\n",
            "Fitting Model 160\n",
            "Fitting Model 161\n",
            "Fitting Model 162\n",
            "Fitting Model 163\n",
            "Fitting Model 164\n",
            "Fitting Model 165\n",
            "Fitting Model 166\n",
            "Fitting Model 167\n",
            "Fitting Model 168\n",
            "Fitting Model 169\n",
            "Fitting Model 170\n",
            "Fitting Model 171\n",
            "Fitting Model 172\n",
            "Fitting Model 173\n",
            "Fitting Model 174\n",
            "Fitting Model 175\n",
            "Fitting Model 176\n",
            "Fitting Model 177\n",
            "Fitting Model 178\n",
            "Fitting Model 179\n",
            "Fitting Model 180\n",
            "Fitting Model 181\n",
            "Fitting Model 182\n",
            "Fitting Model 183\n",
            "Fitting Model 184\n",
            "Fitting Model 185\n",
            "Fitting Model 186\n",
            "Fitting Model 187\n",
            "Fitting Model 188\n",
            "Fitting Model 189\n",
            "Fitting Model 190\n",
            "Fitting Model 191\n",
            "Fitting Model 192\n",
            "Fitting Model 193\n",
            "Fitting Model 194\n",
            "Fitting Model 195\n",
            "Fitting Model 196\n",
            "Fitting Model 197\n",
            "Fitting Model 198\n",
            "Fitting Model 199\n",
            "Fitting Model 200\n",
            "Fitting Model 201\n",
            "Fitting Model 202\n",
            "Fitting Model 203\n",
            "Fitting Model 204\n",
            "Fitting Model 205\n",
            "Fitting Model 206\n",
            "Fitting Model 207\n",
            "Fitting Model 208\n",
            "Fitting Model 209\n",
            "Fitting Model 210\n",
            "Fitting Model 211\n",
            "Fitting Model 212\n",
            "Fitting Model 213\n",
            "Fitting Model 214\n",
            "Fitting Model 215\n",
            "Fitting Model 216\n",
            "Fitting Model 217\n",
            "Fitting Model 218\n",
            "Fitting Model 219\n",
            "Fitting Model 220\n",
            "Fitting Model 221\n",
            "Fitting Model 222\n",
            "Fitting Model 223\n",
            "Fitting Model 224\n",
            "Fitting Model 225\n",
            "Fitting Model 226\n",
            "Fitting Model 227\n",
            "Fitting Model 228\n",
            "Fitting Model 229\n",
            "Fitting Model 230\n",
            "Fitting Model 231\n",
            "Fitting Model 232\n",
            "Fitting Model 233\n",
            "Fitting Model 234\n",
            "Fitting Model 235\n",
            "Fitting Model 236\n",
            "Fitting Model 237\n",
            "Fitting Model 238\n",
            "Fitting Model 239\n",
            "Fitting Model 240\n",
            "Fitting Model 241\n",
            "Fitting Model 242\n",
            "Fitting Model 243\n",
            "Fitting Model 244\n",
            "Fitting Model 245\n",
            "Fitting Model 246\n",
            "Fitting Model 247\n",
            "Fitting Model 248\n",
            "Fitting Model 249\n",
            "Fitting Model 250\n",
            "Fitting Model 251\n",
            "Fitting Model 252\n",
            "Fitting Model 253\n",
            "Fitting Model 254\n",
            "Fitting Model 255\n",
            "Fitting Model 256\n",
            "Fitting Model 257\n",
            "Fitting Model 258\n",
            "Fitting Model 259\n",
            "Fitting Model 260\n",
            "Fitting Model 261\n",
            "Fitting Model 262\n",
            "Fitting Model 263\n",
            "Fitting Model 264\n",
            "Fitting Model 265\n",
            "Fitting Model 266\n",
            "Fitting Model 267\n",
            "Fitting Model 268\n",
            "Fitting Model 269\n",
            "Fitting Model 270\n",
            "Fitting Model 271\n",
            "Fitting Model 272\n",
            "Fitting Model 273\n",
            "Fitting Model 274\n",
            "Fitting Model 275\n",
            "Fitting Model 276\n",
            "Fitting Model 277\n",
            "Fitting Model 278\n",
            "Fitting Model 279\n",
            "Fitting Model 280\n",
            "Fitting Model 281\n",
            "Fitting Model 282\n",
            "Fitting Model 283\n",
            "Fitting Model 284\n",
            "Fitting Model 285\n",
            "Fitting Model 286\n",
            "Fitting Model 287\n",
            "Fitting Model 288\n",
            "Fitting Model 289\n",
            "Fitting Model 290\n",
            "Fitting Model 291\n",
            "Fitting Model 292\n",
            "Fitting Model 293\n",
            "Fitting Model 294\n",
            "Fitting Model 295\n",
            "Fitting Model 296\n",
            "Fitting Model 297\n",
            "Fitting Model 298\n",
            "Fitting Model 299\n",
            "Fitting Model 300\n",
            "Fitting Model 301\n",
            "Fitting Model 302\n",
            "Fitting Model 303\n",
            "Fitting Model 304\n",
            "Fitting Model 305\n",
            "Fitting Model 306\n",
            "Fitting Model 307\n",
            "Fitting Model 308\n",
            "Fitting Model 309\n",
            "Fitting Model 310\n",
            "Fitting Model 311\n",
            "Fitting Model 312\n",
            "Fitting Model 313\n",
            "Fitting Model 314\n",
            "Fitting Model 315\n",
            "Fitting Model 316\n",
            "Fitting Model 317\n",
            "Fitting Model 318\n",
            "Fitting Model 319\n",
            "Fitting Model 320\n",
            "Fitting Model 321\n",
            "Fitting Model 322\n",
            "Fitting Model 323\n",
            "Fitting Model 324\n",
            "Fitting Model 325\n",
            "Fitting Model 326\n",
            "Fitting Model 327\n",
            "Fitting Model 328\n",
            "Fitting Model 329\n",
            "Fitting Model 330\n",
            "Fitting Model 331\n",
            "Fitting Model 332\n",
            "Fitting Model 333\n",
            "Fitting Model 334\n",
            "Fitting Model 335\n",
            "Fitting Model 336\n",
            "Fitting Model 337\n",
            "Fitting Model 338\n",
            "Fitting Model 339\n",
            "Fitting Model 340\n",
            "Fitting Model 341\n",
            "Fitting Model 342\n",
            "Fitting Model 343\n",
            "Fitting Model 344\n",
            "Fitting Model 345\n",
            "Fitting Model 346\n",
            "Fitting Model 347\n",
            "Fitting Model 348\n",
            "Fitting Model 349\n",
            "Fitting Model 350\n",
            "Fitting Model 351\n",
            "Fitting Model 352\n",
            "Fitting Model 353\n",
            "Fitting Model 354\n",
            "Fitting Model 355\n",
            "Fitting Model 356\n",
            "Fitting Model 357\n",
            "Fitting Model 358\n",
            "Fitting Model 359\n",
            "Fitting Model 360\n",
            "Fitting Model 361\n",
            "Fitting Model 362\n",
            "Fitting Model 363\n",
            "Fitting Model 364\n",
            "Fitting Model 365\n",
            "Fitting Model 366\n",
            "Fitting Model 367\n",
            "Fitting Model 368\n",
            "Fitting Model 369\n",
            "Fitting Model 370\n",
            "Fitting Model 371\n",
            "Fitting Model 372\n",
            "Fitting Model 373\n",
            "Fitting Model 374\n",
            "Fitting Model 375\n",
            "Fitting Model 376\n",
            "Fitting Model 377\n",
            "Fitting Model 378\n",
            "Fitting Model 379\n",
            "Fitting Model 380\n",
            "Fitting Model 381\n",
            "Fitting Model 382\n",
            "Fitting Model 383\n",
            "Fitting Model 384\n",
            "Fitting Model 385\n",
            "Fitting Model 386\n",
            "Fitting Model 387\n",
            "Fitting Model 388\n",
            "Fitting Model 389\n",
            "Fitting Model 390\n",
            "Fitting Model 391\n",
            "Fitting Model 392\n",
            "Fitting Model 393\n",
            "Fitting Model 394\n",
            "Fitting Model 395\n",
            "Fitting Model 396\n",
            "Fitting Model 397\n",
            "Fitting Model 398\n",
            "Fitting Model 399\n",
            "Fitting Model 400\n",
            "Fitting Model 401\n",
            "Fitting Model 402\n",
            "Fitting Model 403\n",
            "Fitting Model 404\n",
            "Fitting Model 405\n",
            "Fitting Model 406\n",
            "Fitting Model 407\n",
            "Fitting Model 408\n",
            "Fitting Model 409\n",
            "Fitting Model 410\n",
            "Fitting Model 411\n",
            "Fitting Model 412\n",
            "Fitting Model 413\n",
            "Fitting Model 414\n",
            "Fitting Model 415\n",
            "Fitting Model 416\n",
            "Fitting Model 417\n",
            "Fitting Model 418\n",
            "Fitting Model 419\n",
            "Fitting Model 420\n",
            "Fitting Model 421\n",
            "Fitting Model 422\n",
            "Fitting Model 423\n",
            "Fitting Model 424\n",
            "Fitting Model 425\n",
            "Fitting Model 426\n",
            "Fitting Model 427\n",
            "Fitting Model 428\n",
            "Fitting Model 429\n",
            "Fitting Model 430\n",
            "Fitting Model 431\n",
            "Fitting Model 432\n",
            "Fitting Model 433\n",
            "Fitting Model 434\n",
            "Fitting Model 435\n",
            "Fitting Model 436\n",
            "Fitting Model 437\n",
            "Fitting Model 438\n",
            "Fitting Model 439\n",
            "Fitting Model 440\n",
            "Fitting Model 441\n",
            "Fitting Model 442\n",
            "Fitting Model 443\n",
            "Fitting Model 444\n",
            "Fitting Model 445\n",
            "Fitting Model 446\n",
            "Fitting Model 447\n",
            "Fitting Model 448\n",
            "Fitting Model 449\n",
            "Fitting Model 450\n",
            "Fitting Model 451\n",
            "Fitting Model 452\n",
            "Fitting Model 453\n",
            "Fitting Model 454\n",
            "Fitting Model 455\n",
            "Fitting Model 456\n",
            "Fitting Model 457\n",
            "Fitting Model 458\n",
            "Fitting Model 459\n",
            "Fitting Model 460\n",
            "Fitting Model 461\n",
            "Fitting Model 462\n",
            "Fitting Model 463\n",
            "Fitting Model 464\n",
            "Fitting Model 465\n",
            "Fitting Model 466\n",
            "Fitting Model 467\n",
            "Fitting Model 468\n",
            "Fitting Model 469\n",
            "Fitting Model 470\n",
            "Fitting Model 471\n",
            "Fitting Model 472\n",
            "Fitting Model 473\n",
            "Fitting Model 474\n",
            "Fitting Model 475\n",
            "Fitting Model 476\n",
            "Fitting Model 477\n",
            "Fitting Model 478\n",
            "Fitting Model 479\n",
            "Fitting Model 480\n",
            "Fitting Model 481\n",
            "Fitting Model 482\n",
            "Fitting Model 483\n",
            "Fitting Model 484\n",
            "Fitting Model 485\n",
            "Fitting Model 486\n",
            "Fitting Model 487\n",
            "Fitting Model 488\n",
            "Fitting Model 489\n",
            "Fitting Model 490\n",
            "Fitting Model 491\n",
            "Fitting Model 492\n",
            "Fitting Model 493\n",
            "Fitting Model 494\n",
            "Fitting Model 495\n",
            "Fitting Model 496\n",
            "Fitting Model 497\n",
            "Fitting Model 498\n",
            "Fitting Model 499\n",
            "Fitting Model 500\n",
            "Fitting Model 501\n",
            "Fitting Model 502\n",
            "Fitting Model 503\n",
            "Fitting Model 504\n",
            "Fitting Model 505\n",
            "Fitting Model 506\n",
            "Fitting Model 507\n",
            "Fitting Model 508\n",
            "Fitting Model 509\n",
            "Fitting Model 510\n",
            "Fitting Model 511\n",
            "Fitting Model 512\n",
            "Fitting Model 513\n",
            "Fitting Model 514\n",
            "Fitting Model 515\n",
            "Fitting Model 516\n",
            "Fitting Model 517\n",
            "Fitting Model 518\n",
            "Fitting Model 519\n",
            "Fitting Model 520\n",
            "Fitting Model 521\n",
            "Fitting Model 522\n",
            "Fitting Model 523\n",
            "Fitting Model 524\n",
            "Fitting Model 525\n",
            "Fitting Model 526\n",
            "Fitting Model 527\n",
            "Fitting Model 528\n",
            "Fitting Model 529\n",
            "Fitting Model 530\n",
            "Fitting Model 531\n",
            "Fitting Model 532\n",
            "Fitting Model 533\n",
            "Fitting Model 534\n",
            "Fitting Model 535\n",
            "Fitting Model 536\n",
            "Fitting Model 537\n",
            "Fitting Model 538\n",
            "Fitting Model 539\n",
            "Fitting Model 540\n",
            "Fitting Model 541\n",
            "Fitting Model 542\n",
            "Fitting Model 543\n",
            "Fitting Model 544\n",
            "Fitting Model 545\n",
            "Fitting Model 546\n",
            "Fitting Model 547\n",
            "Fitting Model 548\n",
            "Fitting Model 549\n",
            "Fitting Model 550\n",
            "Fitting Model 551\n",
            "Fitting Model 552\n",
            "Fitting Model 553\n",
            "Fitting Model 554\n",
            "Fitting Model 555\n",
            "Fitting Model 556\n",
            "Fitting Model 557\n",
            "Fitting Model 558\n",
            "Fitting Model 559\n",
            "Fitting Model 560\n",
            "Fitting Model 561\n",
            "Fitting Model 562\n",
            "Fitting Model 563\n",
            "Fitting Model 564\n",
            "Fitting Model 565\n",
            "Fitting Model 566\n",
            "Fitting Model 567\n",
            "Fitting Model 568\n",
            "Fitting Model 569\n",
            "Fitting Model 570\n",
            "Fitting Model 571\n",
            "Fitting Model 572\n",
            "Fitting Model 573\n",
            "Fitting Model 574\n",
            "Fitting Model 575\n",
            "Fitting Model 576\n",
            "Fitting Model 577\n",
            "Fitting Model 578\n",
            "Fitting Model 579\n",
            "Fitting Model 580\n",
            "Fitting Model 581\n",
            "Fitting Model 582\n",
            "Fitting Model 583\n",
            "Fitting Model 584\n",
            "Fitting Model 585\n",
            "Fitting Model 586\n",
            "Fitting Model 587\n",
            "Fitting Model 588\n",
            "Fitting Model 589\n",
            "Fitting Model 590\n",
            "Fitting Model 591\n",
            "Fitting Model 592\n",
            "Fitting Model 593\n",
            "Fitting Model 594\n",
            "Fitting Model 595\n",
            "Fitting Model 596\n",
            "Fitting Model 597\n",
            "Fitting Model 598\n",
            "Fitting Model 599\n",
            "Fitting Model 600\n",
            "Fitting Model 601\n",
            "Fitting Model 602\n",
            "Fitting Model 603\n",
            "Fitting Model 604\n",
            "Fitting Model 605\n",
            "Fitting Model 606\n",
            "Fitting Model 607\n",
            "Fitting Model 608\n",
            "Fitting Model 609\n",
            "Fitting Model 610\n",
            "Fitting Model 611\n",
            "Fitting Model 612\n",
            "Fitting Model 613\n",
            "Fitting Model 614\n",
            "Fitting Model 615\n",
            "Fitting Model 616\n",
            "Fitting Model 617\n",
            "Fitting Model 618\n",
            "Fitting Model 619\n",
            "Fitting Model 620\n",
            "Fitting Model 621\n",
            "Fitting Model 622\n",
            "Fitting Model 623\n",
            "Fitting Model 624\n",
            "Fitting Model 625\n",
            "Fitting Model 626\n",
            "Fitting Model 627\n",
            "Fitting Model 628\n",
            "Fitting Model 629\n",
            "Fitting Model 630\n",
            "Fitting Model 631\n",
            "Fitting Model 632\n",
            "Fitting Model 633\n",
            "Fitting Model 634\n",
            "Fitting Model 635\n",
            "Fitting Model 636\n",
            "Fitting Model 637\n",
            "Fitting Model 638\n",
            "Fitting Model 639\n",
            "Fitting Model 640\n",
            "Fitting Model 641\n",
            "Fitting Model 642\n",
            "Fitting Model 643\n",
            "Fitting Model 644\n",
            "Fitting Model 645\n",
            "Fitting Model 646\n",
            "Fitting Model 647\n",
            "Fitting Model 648\n",
            "Fitting Model 649\n",
            "Fitting Model 650\n",
            "Fitting Model 651\n",
            "Fitting Model 652\n",
            "Fitting Model 653\n",
            "Fitting Model 654\n",
            "Fitting Model 655\n",
            "Fitting Model 656\n",
            "Fitting Model 657\n",
            "Fitting Model 658\n",
            "Fitting Model 659\n",
            "Fitting Model 660\n",
            "Fitting Model 661\n",
            "Fitting Model 662\n",
            "Fitting Model 663\n",
            "Fitting Model 664\n",
            "Fitting Model 665\n",
            "Fitting Model 666\n",
            "Fitting Model 667\n",
            "Fitting Model 668\n",
            "Fitting Model 669\n",
            "Fitting Model 670\n",
            "Fitting Model 671\n",
            "Fitting Model 672\n",
            "Fitting Model 673\n",
            "Fitting Model 674\n",
            "Fitting Model 675\n",
            "Fitting Model 676\n",
            "Fitting Model 677\n",
            "Fitting Model 678\n",
            "Fitting Model 679\n",
            "Fitting Model 680\n",
            "Fitting Model 681\n",
            "Fitting Model 682\n",
            "Fitting Model 683\n",
            "Fitting Model 684\n",
            "Fitting Model 685\n",
            "Fitting Model 686\n",
            "Fitting Model 687\n",
            "Fitting Model 688\n",
            "Fitting Model 689\n",
            "Fitting Model 690\n",
            "Fitting Model 691\n",
            "Fitting Model 692\n",
            "Fitting Model 693\n",
            "Fitting Model 694\n",
            "Fitting Model 695\n",
            "Fitting Model 696\n",
            "Fitting Model 697\n",
            "Fitting Model 698\n",
            "Fitting Model 699\n",
            "Fitting Model 700\n",
            "Fitting Model 701\n",
            "Fitting Model 702\n",
            "Fitting Model 703\n",
            "Fitting Model 704\n",
            "Fitting Model 705\n",
            "Fitting Model 706\n",
            "Fitting Model 707\n",
            "Fitting Model 708\n",
            "Fitting Model 709\n",
            "Fitting Model 710\n",
            "Fitting Model 711\n",
            "Fitting Model 712\n",
            "Fitting Model 713\n",
            "Fitting Model 714\n",
            "Fitting Model 715\n",
            "Fitting Model 716\n",
            "Fitting Model 717\n",
            "Fitting Model 718\n",
            "Fitting Model 719\n",
            "Fitting Model 720\n",
            "Fitting Model 721\n",
            "Fitting Model 722\n",
            "Fitting Model 723\n",
            "Fitting Model 724\n",
            "Fitting Model 725\n",
            "Fitting Model 726\n",
            "Fitting Model 727\n",
            "Fitting Model 728\n",
            "Fitting Model 729\n",
            "Fitting Model 730\n",
            "Fitting Model 731\n",
            "Fitting Model 732\n",
            "Fitting Model 733\n",
            "Fitting Model 734\n",
            "Fitting Model 735\n",
            "Fitting Model 736\n",
            "Fitting Model 737\n",
            "Fitting Model 738\n",
            "Fitting Model 739\n",
            "Fitting Model 740\n",
            "Fitting Model 741\n",
            "Fitting Model 742\n",
            "Fitting Model 743\n",
            "Fitting Model 744\n",
            "Fitting Model 745\n",
            "Fitting Model 746\n",
            "Fitting Model 747\n",
            "Fitting Model 748\n",
            "Fitting Model 749\n",
            "Fitting Model 750\n",
            "Fitting Model 751\n",
            "Fitting Model 752\n",
            "Fitting Model 753\n",
            "Fitting Model 754\n",
            "Fitting Model 755\n",
            "Fitting Model 756\n",
            "Fitting Model 757\n",
            "Fitting Model 758\n",
            "Fitting Model 759\n",
            "Fitting Model 760\n",
            "Fitting Model 761\n",
            "Fitting Model 762\n",
            "Fitting Model 763\n",
            "Fitting Model 764\n",
            "Fitting Model 765\n",
            "Fitting Model 766\n",
            "Fitting Model 767\n",
            "Fitting Model 768\n",
            "Fitting Model 769\n",
            "Fitting Model 770\n",
            "Fitting Model 771\n",
            "Fitting Model 772\n",
            "Fitting Model 773\n",
            "Fitting Model 774\n",
            "Fitting Model 775\n",
            "Fitting Model 776\n",
            "Fitting Model 777\n",
            "Fitting Model 778\n",
            "Fitting Model 779\n",
            "Fitting Model 780\n",
            "Fitting Model 781\n",
            "Fitting Model 782\n",
            "Fitting Model 783\n",
            "Fitting Model 784\n",
            "Fitting Model 785\n",
            "Fitting Model 786\n",
            "Fitting Model 787\n",
            "Fitting Model 788\n",
            "Fitting Model 789\n",
            "Fitting Model 790\n",
            "Fitting Model 791\n",
            "Fitting Model 792\n",
            "Fitting Model 793\n",
            "Fitting Model 794\n",
            "Fitting Model 795\n",
            "Fitting Model 796\n",
            "Fitting Model 797\n",
            "Fitting Model 798\n",
            "Fitting Model 799\n",
            "Fitting Model 800\n",
            "Fitting Model 801\n",
            "Fitting Model 802\n",
            "Fitting Model 803\n",
            "Fitting Model 804\n",
            "Fitting Model 805\n",
            "Fitting Model 806\n",
            "Fitting Model 807\n",
            "Fitting Model 808\n",
            "Fitting Model 809\n",
            "Fitting Model 810\n",
            "Fitting Model 811\n",
            "Fitting Model 812\n",
            "Fitting Model 813\n",
            "Fitting Model 814\n",
            "Fitting Model 815\n",
            "Fitting Model 816\n",
            "Fitting Model 817\n",
            "Fitting Model 818\n",
            "Fitting Model 819\n",
            "Fitting Model 820\n",
            "Fitting Model 821\n",
            "Fitting Model 822\n",
            "Fitting Model 823\n",
            "Fitting Model 824\n",
            "Fitting Model 825\n",
            "Fitting Model 826\n",
            "Fitting Model 827\n",
            "Fitting Model 828\n",
            "Fitting Model 829\n",
            "Fitting Model 830\n",
            "Fitting Model 831\n",
            "Fitting Model 832\n",
            "Fitting Model 833\n",
            "Fitting Model 834\n",
            "Fitting Model 835\n",
            "Fitting Model 836\n",
            "Fitting Model 837\n",
            "Fitting Model 838\n",
            "Fitting Model 839\n",
            "Fitting Model 840\n",
            "Fitting Model 841\n",
            "Fitting Model 842\n",
            "Fitting Model 843\n",
            "Fitting Model 844\n",
            "Fitting Model 845\n",
            "Fitting Model 846\n",
            "Fitting Model 847\n",
            "Fitting Model 848\n",
            "Fitting Model 849\n",
            "Fitting Model 850\n",
            "Fitting Model 851\n",
            "Fitting Model 852\n",
            "Fitting Model 853\n",
            "Fitting Model 854\n",
            "Fitting Model 855\n",
            "Fitting Model 856\n",
            "Fitting Model 857\n",
            "Fitting Model 858\n",
            "Fitting Model 859\n",
            "Fitting Model 860\n",
            "Fitting Model 861\n",
            "Fitting Model 862\n",
            "Fitting Model 863\n",
            "Fitting Model 864\n",
            "Fitting Model 865\n",
            "Fitting Model 866\n",
            "Fitting Model 867\n",
            "Fitting Model 868\n",
            "Fitting Model 869\n",
            "Fitting Model 870\n",
            "Fitting Model 871\n",
            "Fitting Model 872\n",
            "Fitting Model 873\n",
            "Fitting Model 874\n",
            "Fitting Model 875\n",
            "Fitting Model 876\n",
            "Fitting Model 877\n",
            "Fitting Model 878\n",
            "Fitting Model 879\n",
            "Fitting Model 880\n",
            "Fitting Model 881\n",
            "Fitting Model 882\n",
            "Fitting Model 883\n",
            "Fitting Model 884\n",
            "Fitting Model 885\n",
            "Fitting Model 886\n",
            "Fitting Model 887\n",
            "Fitting Model 888\n",
            "Fitting Model 889\n",
            "Fitting Model 890\n",
            "Fitting Model 891\n",
            "Fitting Model 892\n",
            "Fitting Model 893\n",
            "Fitting Model 894\n",
            "Fitting Model 895\n",
            "Fitting Model 896\n",
            "Fitting Model 897\n",
            "Fitting Model 898\n",
            "Fitting Model 899\n",
            "Fitting Model 900\n",
            "Fitting Model 901\n",
            "Fitting Model 902\n",
            "Fitting Model 903\n",
            "Fitting Model 904\n",
            "Fitting Model 905\n",
            "Fitting Model 906\n",
            "Fitting Model 907\n",
            "Fitting Model 908\n",
            "Fitting Model 909\n",
            "Fitting Model 910\n",
            "Fitting Model 911\n",
            "Fitting Model 912\n",
            "Fitting Model 913\n",
            "Fitting Model 914\n",
            "Fitting Model 915\n",
            "Fitting Model 916\n",
            "Fitting Model 917\n",
            "Fitting Model 918\n",
            "Fitting Model 919\n",
            "Fitting Model 920\n",
            "Fitting Model 921\n",
            "Fitting Model 922\n",
            "Fitting Model 923\n",
            "Fitting Model 924\n",
            "Fitting Model 925\n",
            "Fitting Model 926\n",
            "Fitting Model 927\n",
            "Fitting Model 928\n",
            "Fitting Model 929\n",
            "Fitting Model 930\n",
            "Fitting Model 931\n",
            "Fitting Model 932\n",
            "Fitting Model 933\n",
            "Fitting Model 934\n",
            "Fitting Model 935\n",
            "Fitting Model 936\n",
            "Fitting Model 937\n",
            "Fitting Model 938\n",
            "Fitting Model 939\n",
            "Fitting Model 940\n",
            "Fitting Model 941\n",
            "Fitting Model 942\n",
            "Fitting Model 943\n",
            "Fitting Model 944\n",
            "Fitting Model 945\n",
            "Fitting Model 946\n",
            "Fitting Model 947\n",
            "Fitting Model 948\n",
            "Fitting Model 949\n",
            "Fitting Model 950\n",
            "Fitting Model 951\n",
            "Fitting Model 952\n",
            "Fitting Model 953\n",
            "Fitting Model 954\n",
            "Fitting Model 955\n",
            "Fitting Model 956\n",
            "Fitting Model 957\n",
            "Fitting Model 958\n",
            "Fitting Model 959\n",
            "Fitting Model 960\n",
            "Fitting Model 961\n",
            "Fitting Model 962\n",
            "Fitting Model 963\n",
            "Fitting Model 964\n",
            "Fitting Model 965\n",
            "Fitting Model 966\n",
            "Fitting Model 967\n",
            "Fitting Model 968\n",
            "Fitting Model 969\n",
            "Fitting Model 970\n",
            "Fitting Model 971\n",
            "Fitting Model 972\n",
            "Fitting Model 973\n",
            "Fitting Model 974\n",
            "Fitting Model 975\n",
            "Fitting Model 976\n",
            "Fitting Model 977\n",
            "Fitting Model 978\n",
            "Fitting Model 979\n",
            "Fitting Model 980\n",
            "Fitting Model 981\n",
            "Fitting Model 982\n",
            "Fitting Model 983\n",
            "Fitting Model 984\n",
            "Fitting Model 985\n",
            "Fitting Model 986\n",
            "Fitting Model 987\n",
            "Fitting Model 988\n",
            "Fitting Model 989\n",
            "Fitting Model 990\n",
            "Fitting Model 991\n",
            "Fitting Model 992\n",
            "Fitting Model 993\n",
            "Fitting Model 994\n",
            "Fitting Model 995\n",
            "Fitting Model 996\n",
            "Fitting Model 997\n",
            "Fitting Model 998\n",
            "Fitting Model 999\n",
            "Fitting Model 1000\n",
            "Fitting Model 1001\n",
            "Fitting Model 1002\n",
            "Fitting Model 1003\n",
            "Fitting Model 1004\n",
            "Fitting Model 1005\n",
            "Fitting Model 1006\n",
            "Fitting Model 1007\n",
            "Fitting Model 1008\n",
            "Fitting Model 1009\n",
            "Fitting Model 1010\n",
            "Fitting Model 1011\n",
            "Fitting Model 1012\n",
            "Fitting Model 1013\n",
            "Fitting Model 1014\n",
            "Fitting Model 1015\n",
            "Fitting Model 1016\n",
            "Fitting Model 1017\n",
            "Fitting Model 1018\n",
            "Fitting Model 1019\n",
            "Fitting Model 1020\n",
            "Fitting Model 1021\n",
            "Fitting Model 1022\n",
            "Fitting Model 1023\n"
          ]
        }
      ],
      "source": [
        "pca_proj = PCAProjectionEnsemble(nmodels=2**10, ndirs=64, Vpca=np.eye(NDIRS, dtype=np.float32))\n",
        "pca_proj.fit(X, Y)\n",
        "# with open(os.path.join(BASE_DIR +'/' + EXP_NAME + '/' + dname, 'EMOE_pca_projens.p'), 'wb') as f:\n",
        "#     pickle.dump(pca_proj, f)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hi49WbHHaExb",
        "outputId": "1c04a077-6af7-4f45-e027-a378151c1e55"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "2025-01-22_02:32:20\n"
          ]
        }
      ],
      "source": [
        "from datetime import datetime\n",
        "now = datetime.now()\n",
        "dt_string = now.strftime('%Y-%m-%d_%H:%M:%S')\n",
        "print(dt_string)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b2QZCRh04rnj"
      },
      "outputs": [],
      "source": [
        "X_train = X  # X[~train_low_conf_ind]\n",
        "labels_train = Y[:, None]  # Y[~train_low_conf_ind, None]\n",
        "\n",
        "train_predictions = pca_proj.predictions(X)\n",
        "train_ens_preds, train_mp, train_sp = pca_proj.predict(preds=train_predictions, return_only_prediction=False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "k4v2lSfWgq7v"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qnS-1z_sgrbE"
      },
      "outputs": [],
      "source": [
        "expand_funcs = [\n",
        "    # # lambda x: gibbs_sample(x, 40, 2),\n",
        "    # lambda x: permute_expand(x),\n",
        "    # lambda x: convex_expand(x),\n",
        "    lambda x: x*(1+np.abs(np.random.randn(x.shape[0], 1)))\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rj3Yvj1BxaZI",
        "outputId": "7a7e52a5-d12e-40e0-adfe-5597015b1439"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "[[[[[[[[[[[[[[[[[[[[[[  0  ]]]]]]]]]]]]]]]]]]]]]]\n",
            ">>>>>>>>0\n",
            ">>>>>>>>1\n",
            ">>>>>>>>2\n",
            "[[[[[[[[[[[[[[[[[[[[[[  1  ]]]]]]]]]]]]]]]]]]]]]]\n",
            ">>>>>>>>0\n",
            ">>>>>>>>1\n",
            ">>>>>>>>2\n",
            "[[[[[[[[[[[[[[[[[[[[[[  2  ]]]]]]]]]]]]]]]]]]]]]]\n",
            ">>>>>>>>0\n",
            ">>>>>>>>1\n",
            ">>>>>>>>2\n",
            "[[[[[[[[[[[[[[[[[[[[[[  3  ]]]]]]]]]]]]]]]]]]]]]]\n",
            ">>>>>>>>0\n",
            ">>>>>>>>1\n",
            ">>>>>>>>2\n",
            "[[[[[[[[[[[[[[[[[[[[[[  4  ]]]]]]]]]]]]]]]]]]]]]]\n",
            ">>>>>>>>0\n",
            ">>>>>>>>1\n",
            ">>>>>>>>2\n"
          ]
        }
      ],
      "source": [
        "conf_range = .1\n",
        "conf_expX = []\n",
        "conf_exppreds = []\n",
        "unconf_expX = []\n",
        "unconf_exppreds = []\n",
        "all_expX = []\n",
        "all_exppreds = []\n",
        "for dexptrl in range(5):\n",
        "  print('[[[[[[[[[[[[[[[[[[[[[[  {}  ]]]]]]]]]]]]]]]]]]]]]]'.format(dexptrl))\n",
        "  for trnd, expfunc in enumerate(expand_funcs):\n",
        "    print('>>>>>>>>{}'.format(trnd))\n",
        "\n",
        "    expanded_X = expfunc(X)\n",
        "    preds_exp = pca_proj.predictions(expanded_X)\n",
        "    ens_preds_exp, mp_exp, sp_exp = pca_proj.predict(preds=preds_exp, return_only_prediction=False)\n",
        "\n",
        "    high_conf_ind = np.logical_or(\n",
        "      np.greater(mp_exp, 1-conf_range/2), np.less(mp_exp, conf_range/2))\n",
        "    conf_expX.append(expanded_X[high_conf_ind])\n",
        "    conf_exppreds.append(preds_exp[high_conf_ind])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oQTlC7F44xe1"
      },
      "outputs": [],
      "source": [
        "conf_expX = np.concatenate(conf_expX, 0)\n",
        "conf_exppreds = np.concatenate(conf_exppreds, 0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ApCYqpVbGamZ"
      },
      "source": [
        "## RUN Tiny EMEO"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JmqpeS-CGOQf",
        "outputId": "db50574e-90c5-4bbf-9cc1-bba7fce70b78"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "0.5\n",
            "Loss at step -01: 1.515\n",
            "Loss at step -01: 1.518\n",
            "Loss at step 000: 1.523\n",
            "Loss at step 1000: 0.230\n",
            "Loss at step 2000: 0.179\n",
            "Getting Average Model (T=0)\n",
            "Loss at step 3000: 0.219\n",
            "Loss at step 4000: 0.159\n",
            "Loss at step 5000: 0.136\n",
            "Getting Average Model (T=1)\n",
            "Loss at step 6000: 0.136\n",
            "Loss at step 7000: 0.160\n",
            "Getting Average Model (T=2)\n",
            "Loss at step 8000: 0.182\n",
            "Loss at step 9000: 0.205\n",
            "Loss at step 10000: 0.162\n",
            "Getting Average Model (T=3)\n"
          ]
        }
      ],
      "source": [
        "def loss(model, inputs, labels, inputs_conf, labels_conf, alpha=0.5):\n",
        "  # inputs_unconf_logits = model.get_features(inputs_unconf)\n",
        "  # inputs_unconf_sigmoids = tf.nn.sigmoid(inputs_unconf_logits)\n",
        "  # Encourage diversity by minimizing probability distribution innerproducts\n",
        "  # p*q + (1-p)(1-q) = 1 - p - q + 2*p*q\n",
        "  # Note: encouraging diversity for low confidence instances, to minimize dependencies on errors\n",
        "  # WRT existing previous experts\n",
        "  # pw_prods = 1. - inputs_unconf_sigmoids - labels_unconf + 2.*inputs_unconf_sigmoids*labels_unconf  # N x nexps\n",
        "  # diversity_loss = tf.reduce_mean(tf.reduce_mean(pw_prods, -1))\n",
        "  diversity_loss = 0.0\n",
        "\n",
        "  # uncertainty_loss = .5*tf.reduce_mean(\n",
        "  #   tf.abs(tf.reduce_mean(inputs_unconf_sigmoids, -1)-.5)) + \\\n",
        "  #   .5*tf.reduce_mean(\n",
        "  #     tf.abs(tf.reduce_mean(inputs_unconf_sigmoids, 0)-.5))\n",
        "  # mean_axis1 = tf.reduce_mean(inputs_unconf_sigmoids, -1, keepdims=True)\n",
        "  # var_axis1 = tf.reduce_mean(\n",
        "  #   tf.reduce_mean((inputs_unconf_sigmoids-mean_axis1)**2, -1))\n",
        "  # mean_axis0 = tf.reduce_mean(inputs_unconf_sigmoids, 0, keepdims=True)\n",
        "  # var_axis0 = tf.reduce_mean(\n",
        "  #   tf.reduce_mean((inputs_unconf_sigmoids-mean_axis0)**2, 0))\n",
        "  # uncertainty_loss = -.5*(var_axis0 + var_axis1)  # increase variances\n",
        "  uncertainty_loss = 0.0\n",
        "\n",
        "  conf_logits = model.get_features(inputs_conf)\n",
        "  conf_ce_loss = tf.reduce_mean(\n",
        "    tf.keras.losses.binary_crossentropy(labels_conf, conf_logits, from_logits=True))\n",
        "  conf_prob_loss = tf.reduce_mean(tf.abs(\n",
        "    tf.reduce_mean(tf.nn.sigmoid(conf_logits), 1, keepdims=True) - tf.reduce_mean(labels_conf, 1, keepdims=True)\n",
        "  ))\n",
        "  conf_loss = conf_ce_loss  # + conf_prob_loss\n",
        "\n",
        "  in_logits = model.get_features(inputs)\n",
        "  in_ce_loss = tf.reduce_mean(\n",
        "    tf.keras.losses.binary_crossentropy(labels, in_logits, from_logits=True))\n",
        "  in_prob_loss = tf.reduce_mean(tf.abs(\n",
        "    tf.reduce_mean(tf.nn.sigmoid(in_logits), 1, keepdims=True) - tf.reduce_mean(labels, 1, keepdims=True)\n",
        "  ))\n",
        "  in_loss = in_ce_loss + in_prob_loss\n",
        "  return in_loss + alpha*conf_loss + diversity_loss + uncertainty_loss\n",
        "\n",
        "def grad(model, inputs, labels, inputs_conf, labels_conf, alpha=0.5):\n",
        "  with tf.GradientTape() as tape:\n",
        "    loss_value = loss(model, inputs, labels, inputs_conf, labels_conf, alpha=alpha)\n",
        "  return tape.gradient(loss_value, model.trainable_variables)\n",
        "\n",
        "steps = TR_ITER\n",
        "#optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)\n",
        "bsize = 256\n",
        "# model = DatasetDiscriminator(dim=X.shape[0], hidden_size=512, nhidden_layers=2, out_size=train_predictions.shape[1])\n",
        "\n",
        "avgint = TR_AVG_ITER\n",
        "\n",
        "conf_thresh = .4\n",
        "alpha_rnd = .5\n",
        "print(alpha_rnd)\n",
        "# critical_X = critical_ood[np.random.permutation(critical_ood.shape[0])[:X.shape[0]]]\n",
        "# critical_X = critical_ood[np.random.permutation(critical_ood.shape[0])[:1000]]\n",
        "# critical_X = critical_ood\n",
        "expanded_X = X*(1+np.abs(np.random.randn(X.shape[0], 1)))\n",
        "expanded_predictions = pca_proj.predictions(expanded_X)\n",
        "expanded_ens_preds, expanded_mp, expanded_sp = pca_proj.predict(expanded_X, preds=expanded_predictions, return_only_prediction=False)\n",
        "expanded_X = expanded_X[np.abs(expanded_mp-.5)>conf_thresh]\n",
        "expanded_predictions = expanded_predictions[np.abs(expanded_mp-.5)>conf_thresh]\n",
        "\n",
        "batch = training_batch(\n",
        "  X_train,  # np.concatenate((X_train, conf_expX), 0),\n",
        "  np.float32(train_predictions),  #np.concatenate((np.float32(labels_train), np.float32(conf_exppreds)), 0),\n",
        "  bsize)\n",
        "batch_conf = training_batch(\n",
        "  np.concatenate((conf_expX, expanded_X), 0), # expanded_X,\n",
        "  np.float32(np.concatenate((conf_exppreds, expanded_predictions), 0)),\n",
        "  bsize)\n",
        "batch_unconf = training_batch(\n",
        "  unconf_expX,\n",
        "  np.float32(unconf_exppreds),\n",
        "  bsize)\n",
        "# conf_lbls = np.concatenate((train_predictions, expanded_predictions), 0)\n",
        "# batch_conf = training_batch(\n",
        "#   np.concatenate((X, expanded_X), 0),\n",
        "#   conf_lbls,  # (conf_lbls-np.min(conf_lbls))/(np.max(conf_lbls)-np.min(conf_lbls)),\n",
        "#   bsize)\n",
        "\n",
        "xb, xlbls = batch()\n",
        "xb_conf, xlbls_conf = batch_conf()\n",
        "# xb_unconf, xlbls_unconf = batch_unconf()\n",
        "# TODO: consider new model each round?\n",
        "T = 0\n",
        "model = DatasetDiscriminator(dim=X.shape[0], hidden_size=32, nhidden_layers=2, out_size=train_predictions.shape[1])\n",
        "average_model = DatasetDiscriminator(dim=X.shape[0], hidden_size=32, nhidden_layers=2, out_size=train_predictions.shape[1])\n",
        "optimizer = tf.keras.optimizers.Adam(learning_rate=0.0005)\n",
        "print(\"Loss at step {:03d}: {:.3f}\".format(-1, loss(model, xb, xlbls, xb_conf, xlbls_conf, alpha=alpha_rnd)))\n",
        "print(\"Loss at step {:03d}: {:.3f}\".format(-1, loss(average_model, xb, xlbls, xb_conf, xlbls_conf, alpha=alpha_rnd)))\n",
        "for i in range(steps):\n",
        "  xb, xlbls = batch()\n",
        "  xb_conf, xlbls_conf = batch_conf()\n",
        "  # xb_unconf, xlbls_unconf = batch_unconf()\n",
        "  grads = grad(model, xb, xlbls, xb_conf, xlbls_conf, alpha=alpha_rnd)\n",
        "  optimizer.apply_gradients(zip(grads, model.trainable_variables))\n",
        "  if i % 1000 == 0:\n",
        "    print(\"Loss at step {:03d}: {:.3f}\".format(i, loss(model, xb, xlbls, xb_conf, xlbls_conf, alpha=alpha_rnd)))\n",
        "  if i > 0 and i % avgint == 0:\n",
        "    print('Getting Average Model (T={})'.format(T))\n",
        "    for avgweights, weights in zip(average_model.trainable_variables, model.trainable_variables):\n",
        "      avgweights.assign((T/(T+1.))*avgweights + (1./(T+1))*weights)\n",
        "    T = T + 1\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rgHHl4BRfNS0"
      },
      "outputs": [],
      "source": [
        "AUCTHRESHS = [0.1, 0.2, 0.3, 0.4, 1]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SZQhmMBpEwLN"
      },
      "outputs": [],
      "source": [
        "end = time.time()\n",
        "length = end - start"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "r7CPg3GeE16I",
        "outputId": "9f33f36a-a762-46e4-dfc4-667848b43cae"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "It took 1061.5244574546814 seconds!\n"
          ]
        }
      ],
      "source": [
        "print(\"It took\", length, \"seconds!\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Fa-pqZk0hcHR"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yHEKdzIlhrHl"
      },
      "outputs": [],
      "source": [
        "save_path = os.path.join(SAVE_DIR, dname, dt_string)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1bVBhiABhrHl"
      },
      "outputs": [],
      "source": [
        "os.makedirs(save_path, exist_ok=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "l2WKovbZhrHl"
      },
      "outputs": [],
      "source": [
        "models_dict = {}\n",
        "models_dict['base'] = pca_proj\n",
        "models_dict['model_net'] = model.net\n",
        "models_dict['avg_model_net'] = average_model.net"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IKsHL1IDhrHm"
      },
      "outputs": [],
      "source": [
        "with open(os.path.join(save_path, 'models.p'), 'wb') as f:\n",
        "    pickle.dump(models_dict, f)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "s1r1vnZxhrHm",
        "outputId": "d0cf24d2-ef71-46c4-b5ec-ba8f3de3efca"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "TensorShape([7392, 1024])"
            ]
          },
          "execution_count": 520,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "model.get_features(external_X).shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3-B0d0INhrHm"
      },
      "outputs": [],
      "source": [
        "predictions_dict = {}\n",
        "predictions_dict['heldid_Y'] = held_Y\n",
        "predictions_dict['heldid_base'] = pca_proj.predictions(held_X)\n",
        "predictions_dict['heldid_model_net'] = model.predictions(held_X)\n",
        "predictions_dict['heldid_avg_model_net'] = average_model.predictions(held_X)\n",
        "predictions_dict['heldood_Y'] = external_Y\n",
        "predictions_dict['heldood_base'] = pca_proj.predictions(external_X)\n",
        "predictions_dict['heldood_model_net'] = model.predictions(external_X)\n",
        "predictions_dict['heldood_avg_model_net'] = average_model.predictions(external_X)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "y5tx3UiAhrHm"
      },
      "outputs": [],
      "source": [
        "with open(os.path.join(save_path, 'predictions.p'), 'wb') as f:\n",
        "    pickle.dump(predictions_dict, f)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c1YcWqG0hrHm"
      },
      "outputs": [],
      "source": [
        "results_dict = {}\n",
        "results_dict['threshs'] = AUCTHRESHS"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 651
        },
        "id": "FTsa8JlqhrHn",
        "outputId": "bfec10d0-bbfc-4f5e-de24-2f42b3323def"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Base:\n",
            "[0.99934591 0.99869182 0.99803773 0.99738363 0.98615538] (AUROC:0.8792977845915576,ACC:0.9092458884877658)\n",
            "Model:\n",
            "[0.99685077 0.99636725 0.99602944 0.99541868 0.98322942] (AUROC:0.868830309384008,ACC:0.9108503810669876)\n",
            "Average Model:\n",
            "[0.9941475  0.9941168  0.99380902 0.99356101 0.98232302] (AUROC:0.8674111376797674,ACC:0.9122543120738067)\n",
            "Base+Model:\n",
            "[0.99664381 0.99645522 0.99645914 0.99625005 0.98539759] (AUROC:0.8830108864584355,ACC:0.9119534697152025)\n",
            "Base+Average Model:\n",
            "[0.99454299 0.99503082 0.99512763 0.99509648 0.98457102] (AUROC:0.8802508841220532,ACC:0.9107501002807862)\n"
          ]
        },
        {
          "data": {
            "image/png": "[encoded data, 46528 characters]",
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "# plot id\n",
        "###########\n",
        "base_predictions = predictions_dict['heldid_base']\n",
        "model_predictions = predictions_dict['heldid_model_net']\n",
        "average_model_predictions = predictions_dict['heldid_avg_model_net']\n",
        "true_Y = predictions_dict['heldid_Y']\n",
        "\n",
        "plt.figure()\n",
        "\n",
        "# base model\n",
        "confs = np.mean(base_predictions, -1)\n",
        "preds = (confs >.5)\n",
        "precision, recall, thresholds = metrics.precision_recall_curve(true_Y, confs)\n",
        "plt.plot(recall, precision, c='brown')\n",
        "BASE_AUPRC = auprc_threshs(confs, true_Y, AUCTHRESHS)\n",
        "BASE_AUROC = metrics.roc_auc_score(true_Y, confs)\n",
        "BASE_ACC = metrics.accuracy_score(true_Y, preds)\n",
        "\n",
        "print('Base:\\n{} (AUROC:{},ACC:{})'.format(BASE_AUPRC, BASE_AUROC,BASE_ACC))\n",
        "\n",
        "# student model\n",
        "confs = np.mean(model_predictions, -1)\n",
        "preds = (confs >.5)\n",
        "precision, recall, thresholds = metrics.precision_recall_curve(true_Y, confs)\n",
        "plt.plot(recall, precision, c='red')\n",
        "MOD_AUPRC = auprc_threshs(confs, true_Y, AUCTHRESHS)\n",
        "MOD_AUROC = metrics.roc_auc_score(true_Y, confs)\n",
        "MOD_ACC = metrics.accuracy_score(true_Y, preds)\n",
        "print('Model:\\n{} (AUROC:{},ACC:{})'.format(MOD_AUPRC, MOD_AUROC,MOD_ACC))\n",
        "\n",
        "# average student model\n",
        "confs = np.mean(average_model_predictions, -1)\n",
        "preds = (confs >.5)\n",
        "precision, recall, thresholds = metrics.precision_recall_curve(true_Y, confs)\n",
        "plt.plot(recall, precision, c='pink')\n",
        "AVGMOD_AUPRC = auprc_threshs(confs, true_Y, AUCTHRESHS)\n",
        "AVGMOD_AUROC = metrics.roc_auc_score(true_Y, confs)\n",
        "AVGMOD_ACC = metrics.accuracy_score(true_Y, preds)\n",
        "print('Average Model:\\n{} (AUROC:{},ACC:{})'.format(AVGMOD_AUPRC, AVGMOD_AUROC,AVGMOD_ACC))\n",
        "\n",
        "# base+student model\n",
        "confs = (np.mean(base_predictions, -1)+np.mean(model_predictions, -1))/2.\n",
        "preds = (confs >.5)\n",
        "precision, recall, thresholds = metrics.precision_recall_curve(true_Y, confs)\n",
        "plt.plot(recall, precision, c='blue')\n",
        "EMOE_AUPRC = auprc_threshs(confs, true_Y, AUCTHRESHS)\n",
        "EMOE_AUROC = metrics.roc_auc_score(true_Y, confs)\n",
        "EMOE_ACC = metrics.accuracy_score(true_Y, preds)\n",
        "print('Base+Model:\\n{} (AUROC:{},ACC:{})'.format(EMOE_AUPRC, EMOE_AUROC,EMOE_ACC))\n",
        "\n",
        "# base+average student model\n",
        "confs = (np.mean(base_predictions, -1)+np.mean(average_model_predictions, -1))/2.\n",
        "preds = (confs >.5)\n",
        "precision, recall, thresholds = metrics.precision_recall_curve(true_Y, confs)\n",
        "plt.plot(recall, precision, c='cyan')\n",
        "AVGEMOE_AUPRC = auprc_threshs(confs, true_Y, AUCTHRESHS)\n",
        "AVGEMOE_AUROC = metrics.roc_auc_score(true_Y, confs)\n",
        "AVGEMOE_ACC = metrics.accuracy_score(true_Y, preds)\n",
        "print('Base+Average Model:\\n{} (AUROC:{},ACC:{})'.format(AVGEMOE_AUPRC, AVGEMOE_AUROC,AVGEMOE_ACC))\n",
        "\n",
        "plt.xticks(np.arange(0, 1.01, step=0.1))\n",
        "plt.xticks(np.arange(0, 1.01, step=0.05), minor=True)\n",
        "plt.yticks(np.arange(0.6, 1.01, step=0.05))\n",
        "plt.grid(True, which='both')\n",
        "plt.xlabel('Recall ({} Positive)'.format(int(np.sum(true_Y))))\n",
        "plt.ylabel('Precision')\n",
        "plt.title('ID Precision vs Recall')\n",
        "plt.legend(['Base', 'NNet', 'AVGNNET', 'EMOE', 'AVGEMOE'])\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U0vHW1hqhrHn"
      },
      "outputs": [],
      "source": [
        "results_dict['heldid_base_AUPRC'] = BASE_AUPRC\n",
        "results_dict['heldid_model_AUPRC'] = MOD_AUPRC\n",
        "results_dict['heldid_avg_model_AUPRC'] = AVGMOD_AUPRC\n",
        "results_dict['heldid_emoe_AUPRC'] = EMOE_AUPRC\n",
        "results_dict['heldid_avg_emoe_AUPRC'] = AVGEMOE_AUPRC\n",
        "\n",
        "results_dict['heldid_base_AUROC'] = BASE_AUROC\n",
        "results_dict['heldid_model_AUROC'] = MOD_AUROC\n",
        "results_dict['heldid_avg_model_AUROC'] = AVGMOD_AUROC\n",
        "results_dict['heldid_emoe_AUROC'] = EMOE_AUROC\n",
        "results_dict['heldid_avg_emoe_AUROC'] = AVGEMOE_AUROC"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 651
        },
        "id": "4vK71YGzhrHn",
        "outputId": "58ba17e6-4528-4717-96e6-33d805f6af94"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Base:\n",
            "[0.98978626 0.97957253 0.97123953 0.96371758 0.91160392] (AUROC:0.7001263008870153,ACC:0.8158820346320347)\n",
            "Model:\n",
            "[0.97229201 0.96298288 0.95304496 0.94442465 0.89961754] (AUROC:0.6806948730566995,ACC:0.8058712121212122)\n",
            "Average Model:\n",
            "[0.97738618 0.96631571 0.95793361 0.95015054 0.90271913] (AUROC:0.6849178923024337,ACC:0.8110119047619048)\n",
            "Base+Model:\n",
            "[0.9879569  0.97780064 0.9681153  0.95912635 0.9084658 ] (AUROC:0.6960897633252283,ACC:0.8158820346320347)\n",
            "Base+Average Model:\n",
            "[0.98755189 0.97715961 0.96852404 0.9602149  0.90913937] (AUROC:0.6971072612101128,ACC:0.8168290043290043)\n"
          ]
        },
        {
          "data": {
            "image/png": "[encoded data, 57464 characters]",
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "# plot ood\n",
        "###########\n",
        "base_predictions = predictions_dict['heldood_base']\n",
        "model_predictions = predictions_dict['heldood_model_net']\n",
        "average_model_predictions = predictions_dict['heldood_avg_model_net']\n",
        "true_Y = predictions_dict['heldood_Y']\n",
        "\n",
        "plt.figure()\n",
        "\n",
        "# base model\n",
        "confs = np.mean(base_predictions, -1)\n",
        "preds = (confs >.5)\n",
        "precision, recall, thresholds = metrics.precision_recall_curve(true_Y, confs)\n",
        "plt.plot(recall, precision, c='brown')\n",
        "BASE_AUPRC = auprc_threshs(confs, true_Y, AUCTHRESHS)\n",
        "BASE_AUROC = metrics.roc_auc_score(true_Y, confs)\n",
        "BASE_ACC = metrics.accuracy_score(true_Y, preds)\n",
        "\n",
        "print('Base:\\n{} (AUROC:{},ACC:{})'.format(BASE_AUPRC, BASE_AUROC,BASE_ACC))\n",
        "\n",
        "# student model\n",
        "confs = np.mean(model_predictions, -1)\n",
        "preds = (confs >.5)\n",
        "precision, recall, thresholds = metrics.precision_recall_curve(true_Y, confs)\n",
        "plt.plot(recall, precision, c='red')\n",
        "MOD_AUPRC = auprc_threshs(confs, true_Y, AUCTHRESHS)\n",
        "MOD_AUROC = metrics.roc_auc_score(true_Y, confs)\n",
        "MOD_ACC = metrics.accuracy_score(true_Y, preds)\n",
        "print('Model:\\n{} (AUROC:{},ACC:{})'.format(MOD_AUPRC, MOD_AUROC,MOD_ACC))\n",
        "\n",
        "# average student model\n",
        "confs = np.mean(average_model_predictions, -1)\n",
        "preds = (confs >.5)\n",
        "precision, recall, thresholds = metrics.precision_recall_curve(true_Y, confs)\n",
        "plt.plot(recall, precision, c='pink')\n",
        "AVGMOD_AUPRC = auprc_threshs(confs, true_Y, AUCTHRESHS)\n",
        "AVGMOD_AUROC = metrics.roc_auc_score(true_Y, confs)\n",
        "AVGMOD_ACC = metrics.accuracy_score(true_Y, preds)\n",
        "print('Average Model:\\n{} (AUROC:{},ACC:{})'.format(AVGMOD_AUPRC, AVGMOD_AUROC,AVGMOD_ACC))\n",
        "\n",
        "# base+student model\n",
        "confs = (np.mean(base_predictions, -1)+np.mean(model_predictions, -1))/2.\n",
        "preds = (confs >.5)\n",
        "precision, recall, thresholds = metrics.precision_recall_curve(true_Y, confs)\n",
        "plt.plot(recall, precision, c='blue')\n",
        "EMOE_AUPRC = auprc_threshs(confs, true_Y, AUCTHRESHS)\n",
        "EMOE_AUROC = metrics.roc_auc_score(true_Y, confs)\n",
        "EMOE_ACC = metrics.accuracy_score(true_Y, preds)\n",
        "print('Base+Model:\\n{} (AUROC:{},ACC:{})'.format(EMOE_AUPRC, EMOE_AUROC,EMOE_ACC))\n",
        "\n",
        "# base+average student model\n",
        "confs = (np.mean(base_predictions, -1)+np.mean(average_model_predictions, -1))/2.\n",
        "preds = (confs >.5)\n",
        "precision, recall, thresholds = metrics.precision_recall_curve(true_Y, confs)\n",
        "plt.plot(recall, precision, c='cyan')\n",
        "AVGEMOE_AUPRC = auprc_threshs(confs, true_Y, AUCTHRESHS)\n",
        "AVGEMOE_AUROC = metrics.roc_auc_score(true_Y, confs)\n",
        "AVGEMOE_ACC = metrics.accuracy_score(true_Y, preds)\n",
        "print('Base+Average Model:\\n{} (AUROC:{},ACC:{})'.format(AVGEMOE_AUPRC, AVGEMOE_AUROC,AVGEMOE_ACC))\n",
        "\n",
        "plt.xticks(np.arange(0, 1.01, step=0.1))\n",
        "plt.xticks(np.arange(0, 1.01, step=0.05), minor=True)\n",
        "plt.yticks(np.arange(0.6, 1.01, step=0.05))\n",
        "plt.grid(True, which='both')\n",
        "plt.xlabel('Recall ({} Positive)'.format(int(np.sum(true_Y))))\n",
        "plt.ylabel('Precision')\n",
        "plt.title('ID Precision vs Recall')\n",
        "plt.legend(['Base', 'NNet', 'AVGNNET', 'EMOE', 'AVGEMOE'])\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vPUQWf5ohrHn"
      },
      "outputs": [],
      "source": [
        "results_dict['heldood_base_AUPRC'] = BASE_AUPRC\n",
        "results_dict['heldood_model_AUPRC'] = MOD_AUPRC\n",
        "results_dict['heldood_avg_model_AUPRC'] = AVGMOD_AUPRC\n",
        "results_dict['heldood_emoe_AUPRC'] = EMOE_AUPRC\n",
        "results_dict['heldood_avg_emoe_AUPRC'] = AVGEMOE_AUPRC\n",
        "\n",
        "results_dict['heldood_base_AUROC'] = BASE_AUROC\n",
        "results_dict['heldood_model_AUROC'] = MOD_AUROC\n",
        "results_dict['heldood_avg_model_AUROC'] = AVGMOD_AUROC\n",
        "results_dict['heldood_emoe_AUROC'] = EMOE_AUROC\n",
        "results_dict['heldood_avg_emoe_AUROC'] = AVGEMOE_AUROC\n",
        "\n",
        "results_dict['heldood_base_ACC'] = BASE_ACC\n",
        "results_dict['heldood_model_ACC'] = MOD_ACC\n",
        "results_dict['heldood_avg_model_ACC'] = AVGMOD_ACC\n",
        "results_dict['heldood_emoe_ACC'] = EMOE_ACC\n",
        "results_dict['heldood_avg_emoe_ACC'] = AVGEMOE_ACC"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GG1XFcaEhrHn"
      },
      "outputs": [],
      "source": [
        "with open(os.path.join(save_path, 'results.p'), 'wb') as f:\n",
        "    pickle.dump(results_dict, f)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nvURo8ZwFTuX"
      },
      "source": [
        "## RUN FULL EMOE"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "qqydmkQoGfcR",
        "outputId": "496646b5-b3e3-4714-8cf9-cefab14afd73"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "0.5\n",
            "Loss at step -01: 1.579\n",
            "Loss at step -01: 1.581\n",
            "Loss at step 000: 1.508\n",
            "Loss at step 1000: 0.227\n",
            "Loss at step 2000: 0.162\n",
            "Loss at step 3000: 0.141\n",
            "Loss at step 4000: 0.136\n",
            "Loss at step 5000: 0.106\n",
            "Getting Average Model (T=0)\n",
            "Loss at step 6000: 0.081\n",
            "Loss at step 7000: 0.100\n",
            "Loss at step 8000: 0.082\n",
            "Loss at step 9000: 0.073\n",
            "Loss at step 10000: 0.094\n",
            "Getting Average Model (T=1)\n",
            "Loss at step 11000: 0.090\n",
            "Loss at step 12000: 0.067\n",
            "Loss at step 13000: 0.075\n",
            "Loss at step 14000: 0.077\n",
            "Loss at step 15000: 0.077\n",
            "Getting Average Model (T=2)\n",
            "Loss at step 16000: 0.067\n",
            "Loss at step 17000: 0.080\n",
            "Loss at step 18000: 0.083\n",
            "Loss at step 19000: 0.074\n",
            "Loss at step 20000: 0.072\n",
            "Getting Average Model (T=3)\n"
          ]
        }
      ],
      "source": [
        "steps = 20001\n",
        "#optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)\n",
        "bsize = 256\n",
        "# model = DatasetDiscriminator(dim=X.shape[0], hidden_size=512, nhidden_layers=2, out_size=train_predictions.shape[1])\n",
        "\n",
        "avgint = 5000\n",
        "\n",
        "conf_thresh = .4\n",
        "alpha_rnd = .5\n",
        "print(alpha_rnd)\n",
        "# critical_X = critical_ood[np.random.permutation(critical_ood.shape[0])[:X.shape[0]]]\n",
        "# critical_X = critical_ood[np.random.permutation(critical_ood.shape[0])[:1000]]\n",
        "\n",
        "expanded_X = X*(1+np.abs(np.random.randn(X.shape[0], 1)))\n",
        "expanded_predictions = pca_proj.predictions(expanded_X)\n",
        "expanded_ens_preds, expanded_mp, expanded_sp = pca_proj.predict(expanded_X, preds=expanded_predictions, return_only_prediction=False)\n",
        "expanded_X = expanded_X[np.abs(expanded_mp-.5)>conf_thresh]\n",
        "expanded_predictions = expanded_predictions[np.abs(expanded_mp-.5)>conf_thresh]\n",
        "\n",
        "batch = training_batch(\n",
        "  X_train,  # np.concatenate((X_train, conf_expX), 0),\n",
        "  np.float32(train_predictions),  #np.concatenate((np.float32(labels_train), np.float32(conf_exppreds)), 0),\n",
        "  bsize)\n",
        "batch_conf = training_batch(\n",
        "  np.concatenate((conf_expX, expanded_X), 0), # expanded_X,\n",
        "  np.float32(np.concatenate((conf_exppreds, expanded_predictions), 0)),\n",
        "  bsize)\n",
        "batch_unconf = training_batch(\n",
        "  unconf_expX,\n",
        "  np.float32(unconf_exppreds),\n",
        "  bsize)\n",
        "# conf_lbls = np.concatenate((train_predictions, expanded_predictions), 0)\n",
        "# batch_conf = training_batch(\n",
        "#   np.concatenate((X, expanded_X), 0),\n",
        "#   conf_lbls,  # (conf_lbls-np.min(conf_lbls))/(np.max(conf_lbls)-np.min(conf_lbls)),\n",
        "#   bsize)\n",
        "\n",
        "xb, xlbls = batch()\n",
        "xb_conf, xlbls_conf = batch_conf()\n",
        "# xb_unconf, xlbls_unconf = batch_unconf()\n",
        "# TODO: consider new model each round?\n",
        "T = 0\n",
        "model = DatasetDiscriminator(dim=X.shape[0], hidden_size=512, nhidden_layers=2, out_size=train_predictions.shape[1])\n",
        "average_model = DatasetDiscriminator(dim=X.shape[0], hidden_size=512, nhidden_layers=2, out_size=train_predictions.shape[1])\n",
        "optimizer = tf.keras.optimizers.Adam(learning_rate=0.0005)\n",
        "print(\"Loss at step {:03d}: {:.3f}\".format(-1, loss(model, xb, xlbls, xb_conf, xlbls_conf, alpha=alpha_rnd)))\n",
        "print(\"Loss at step {:03d}: {:.3f}\".format(-1, loss(average_model, xb, xlbls, xb_conf, xlbls_conf, alpha=alpha_rnd)))\n",
        "for i in range(steps):\n",
        "  xb, xlbls = batch()\n",
        "  xb_conf, xlbls_conf = batch_conf()\n",
        "  # xb_unconf, xlbls_unconf = batch_unconf()\n",
        "  grads = grad(model, xb, xlbls, xb_conf, xlbls_conf, alpha=alpha_rnd)\n",
        "  optimizer.apply_gradients(zip(grads, model.trainable_variables))\n",
        "  if i % 1000 == 0:\n",
        "    print(\"Loss at step {:03d}: {:.3f}\".format(i, loss(model, xb, xlbls, xb_conf, xlbls_conf, alpha=alpha_rnd)))\n",
        "  if i > 0 and i % avgint == 0:\n",
        "    print('Getting Average Model (T={})'.format(T))\n",
        "    for avgweights, weights in zip(average_model.trainable_variables, model.trainable_variables):\n",
        "      avgweights.assign((T/(T+1.))*avgweights + (1./(T+1))*weights)\n",
        "    T = T + 1\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "98-4lUyBGfcR"
      },
      "outputs": [],
      "source": [
        "end = time.time()\n",
        "length = end - start"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "nJsUzA0VGfcR",
        "outputId": "b8210529-d0fa-4b57-961c-15e0d68aaec1"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "It took 2228.165778875351 seconds!\n"
          ]
        }
      ],
      "source": [
        "print(\"It took\", length, \"seconds!\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jbz6Vk2kGfcR"
      },
      "outputs": [],
      "source": [
        "# from datetime import datetime\n",
        "# now = datetime.now()\n",
        "# dt_string = now.strftime('%Y-%m-%d_%H:%M:%S')\n",
        "# print(dt_string)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FXk5DZKwGfcR"
      },
      "outputs": [],
      "source": [
        "save_path = os.path.join(SAVE_DIR, 'full_EMOE' , dname, dt_string)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gU_lDDN8GfcR"
      },
      "outputs": [],
      "source": [
        "os.makedirs(save_path, exist_ok=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TWkWll2YGfcR"
      },
      "outputs": [],
      "source": [
        "models_dict = {}\n",
        "models_dict['base'] = pca_proj\n",
        "models_dict['model_net'] = model.net\n",
        "models_dict['avg_model_net'] = average_model.net"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Hs09LbTeGfcR"
      },
      "outputs": [],
      "source": [
        "with open(os.path.join(save_path, 'models.p'), 'wb') as f:\n",
        "    pickle.dump(models_dict, f)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Xhho6yD2GfcR"
      },
      "outputs": [],
      "source": [
        "predictions_dict = {}\n",
        "predictions_dict['heldid_Y'] = held_Y\n",
        "predictions_dict['heldid_base'] = pca_proj.predictions(held_X)\n",
        "predictions_dict['heldid_model_net'] = model.predictions(held_X)\n",
        "predictions_dict['heldid_avg_model_net'] = average_model.predictions(held_X)\n",
        "predictions_dict['heldood_Y'] = external_Y\n",
        "predictions_dict['heldood_base'] = pca_proj.predictions(external_X)\n",
        "predictions_dict['heldood_model_net'] = model.predictions(external_X)\n",
        "predictions_dict['heldood_avg_model_net'] = average_model.predictions(external_X)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AgoP2JGpGfcS"
      },
      "outputs": [],
      "source": [
        "with open(os.path.join(save_path, 'predictions.p'), 'wb') as f:\n",
        "    pickle.dump(predictions_dict, f)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "H-JFD9nZGfcS"
      },
      "outputs": [],
      "source": [
        "results_dict = {}\n",
        "results_dict['threshs'] = AUCTHRESHS"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 651
        },
        "id": "n77FJ6qrGfcS",
        "outputId": "021df153-2fc6-40ee-dffb-0735382d2282"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Base:\n",
            "[0.99934591 0.99869182 0.99803773 0.99738363 0.98615538] (AUROC:0.8792977845915576,ACC:0.9092458884877658)\n",
            "Model:\n",
            "[0.99448297 0.99598816 0.99640948 0.99633526 0.98624324] (AUROC:0.8898776229355682,ACC:0.9127557160048135)\n",
            "Average Model:\n",
            "[0.99422974 0.99537798 0.99599513 0.99605778 0.98601166] (AUROC:0.8887797038213702,ACC:0.9127557160048135)\n",
            "Base+Model:\n",
            "[0.99447217 0.99596969 0.99639197 0.99660406 0.98728677] (AUROC:0.8964646558924108,ACC:0.9104492579221821)\n",
            "Base+Average Model:\n",
            "[0.99419869 0.99557323 0.9963682  0.99672765 0.98738597] (AUROC:0.8968792433492885,ACC:0.910950661853189)\n"
          ]
        },
        {
          "data": {
            "image/png": "[encoded data, 45000 characters]",
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "# plot id\n",
        "###########\n",
        "base_predictions = predictions_dict['heldid_base']\n",
        "model_predictions = predictions_dict['heldid_model_net']\n",
        "average_model_predictions = predictions_dict['heldid_avg_model_net']\n",
        "true_Y = predictions_dict['heldid_Y']\n",
        "\n",
        "plt.figure()\n",
        "\n",
        "# base model\n",
        "confs = np.mean(base_predictions, -1)\n",
        "preds = (confs >.5)\n",
        "precision, recall, thresholds = metrics.precision_recall_curve(true_Y, confs)\n",
        "plt.plot(recall, precision, c='brown')\n",
        "BASE_AUPRC = auprc_threshs(confs, true_Y, AUCTHRESHS)\n",
        "BASE_AUROC = metrics.roc_auc_score(true_Y, confs)\n",
        "BASE_ACC = metrics.accuracy_score(true_Y, preds)\n",
        "\n",
        "print('Base:\\n{} (AUROC:{},ACC:{})'.format(BASE_AUPRC, BASE_AUROC,BASE_ACC))\n",
        "\n",
        "# student model\n",
        "confs = np.mean(model_predictions, -1)\n",
        "preds = (confs >.5)\n",
        "precision, recall, thresholds = metrics.precision_recall_curve(true_Y, confs)\n",
        "plt.plot(recall, precision, c='red')\n",
        "MOD_AUPRC = auprc_threshs(confs, true_Y, AUCTHRESHS)\n",
        "MOD_AUROC = metrics.roc_auc_score(true_Y, confs)\n",
        "MOD_ACC = metrics.accuracy_score(true_Y, preds)\n",
        "print('Model:\\n{} (AUROC:{},ACC:{})'.format(MOD_AUPRC, MOD_AUROC,MOD_ACC))\n",
        "\n",
        "# average student model\n",
        "confs = np.mean(average_model_predictions, -1)\n",
        "preds = (confs >.5)\n",
        "precision, recall, thresholds = metrics.precision_recall_curve(true_Y, confs)\n",
        "plt.plot(recall, precision, c='pink')\n",
        "AVGMOD_AUPRC = auprc_threshs(confs, true_Y, AUCTHRESHS)\n",
        "AVGMOD_AUROC = metrics.roc_auc_score(true_Y, confs)\n",
        "AVGMOD_ACC = metrics.accuracy_score(true_Y, preds)\n",
        "print('Average Model:\\n{} (AUROC:{},ACC:{})'.format(AVGMOD_AUPRC, AVGMOD_AUROC,AVGMOD_ACC))\n",
        "\n",
        "# base+student model\n",
        "confs = (np.mean(base_predictions, -1)+np.mean(model_predictions, -1))/2.\n",
        "preds = (confs >.5)\n",
        "precision, recall, thresholds = metrics.precision_recall_curve(true_Y, confs)\n",
        "plt.plot(recall, precision, c='blue')\n",
        "EMOE_AUPRC = auprc_threshs(confs, true_Y, AUCTHRESHS)\n",
        "EMOE_AUROC = metrics.roc_auc_score(true_Y, confs)\n",
        "EMOE_ACC = metrics.accuracy_score(true_Y, preds)\n",
        "print('Base+Model:\\n{} (AUROC:{},ACC:{})'.format(EMOE_AUPRC, EMOE_AUROC,EMOE_ACC))\n",
        "\n",
        "# base+average student model\n",
        "confs = (np.mean(base_predictions, -1)+np.mean(average_model_predictions, -1))/2.\n",
        "preds = (confs >.5)\n",
        "precision, recall, thresholds = metrics.precision_recall_curve(true_Y, confs)\n",
        "plt.plot(recall, precision, c='cyan')\n",
        "AVGEMOE_AUPRC = auprc_threshs(confs, true_Y, AUCTHRESHS)\n",
        "AVGEMOE_AUROC = metrics.roc_auc_score(true_Y, confs)\n",
        "AVGEMOE_ACC = metrics.accuracy_score(true_Y, preds)\n",
        "print('Base+Average Model:\\n{} (AUROC:{},ACC:{})'.format(AVGEMOE_AUPRC, AVGEMOE_AUROC,AVGEMOE_ACC))\n",
        "\n",
        "plt.xticks(np.arange(0, 1.01, step=0.1))\n",
        "plt.xticks(np.arange(0, 1.01, step=0.05), minor=True)\n",
        "plt.yticks(np.arange(0.6, 1.01, step=0.05))\n",
        "plt.grid(True, which='both')\n",
        "plt.xlabel('Recall ({} Positive)'.format(int(np.sum(true_Y))))\n",
        "plt.ylabel('Precision')\n",
        "plt.title('ID Precision vs Recall')\n",
        "plt.legend(['Base', 'NNet', 'AVGNNET', 'EMOE', 'AVGEMOE'])\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XMO0PYR5GfcS"
      },
      "outputs": [],
      "source": [
        "results_dict['heldid_base_AUPRC'] = BASE_AUPRC\n",
        "results_dict['heldid_model_AUPRC'] = MOD_AUPRC\n",
        "results_dict['heldid_avg_model_AUPRC'] = AVGMOD_AUPRC\n",
        "results_dict['heldid_emoe_AUPRC'] = EMOE_AUPRC\n",
        "results_dict['heldid_avg_emoe_AUPRC'] = AVGEMOE_AUPRC\n",
        "\n",
        "results_dict['heldid_base_AUROC'] = BASE_AUROC\n",
        "results_dict['heldid_model_AUROC'] = MOD_AUROC\n",
        "results_dict['heldid_avg_model_AUROC'] = AVGMOD_AUROC\n",
        "results_dict['heldid_emoe_AUROC'] = EMOE_AUROC\n",
        "results_dict['heldid_avg_emoe_AUROC'] = AVGEMOE_AUROC"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 651
        },
        "id": "HUPjeXDAGfcS",
        "outputId": "c47df412-90c7-4436-d02f-d9526e7d130a"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Base:\n",
            "[0.98978626 0.97957253 0.97123953 0.96371758 0.91160392] (AUROC:0.7001263008870153,ACC:0.8158820346320347)\n",
            "Model:\n",
            "[0.99526782 0.98351106 0.9717482  0.96185949 0.91175095] (AUROC:0.7037347245016448,ACC:0.8129058441558441)\n",
            "Average Model:\n",
            "[0.9943343  0.98104533 0.96850787 0.95951256 0.91021444] (AUROC:0.7008675004302957,ACC:0.8112824675324676)\n",
            "Base+Model:\n",
            "[0.99475699 0.98678272 0.97869601 0.96982057 0.91588921] (AUROC:0.7109260570123659,ACC:0.8157467532467533)\n",
            "Base+Average Model:\n",
            "[0.99437379 0.98658581 0.97890026 0.96988944 0.91573644] (AUROC:0.7102346142004863,ACC:0.8156114718614719)\n"
          ]
        },
        {
          "data": {
            "image/png": "[encoded data, 52416 characters]",
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "# plot ood\n",
        "###########\n",
        "base_predictions = predictions_dict['heldood_base']\n",
        "model_predictions = predictions_dict['heldood_model_net']\n",
        "average_model_predictions = predictions_dict['heldood_avg_model_net']\n",
        "true_Y = predictions_dict['heldood_Y']\n",
        "\n",
        "plt.figure()\n",
        "\n",
        "# base model\n",
        "confs = np.mean(base_predictions, -1)\n",
        "preds = (confs >.5)\n",
        "precision, recall, thresholds = metrics.precision_recall_curve(true_Y, confs)\n",
        "plt.plot(recall, precision, c='brown')\n",
        "BASE_AUPRC = auprc_threshs(confs, true_Y, AUCTHRESHS)\n",
        "BASE_AUROC = metrics.roc_auc_score(true_Y, confs)\n",
        "BASE_ACC = metrics.accuracy_score(true_Y, preds)\n",
        "\n",
        "print('Base:\\n{} (AUROC:{},ACC:{})'.format(BASE_AUPRC, BASE_AUROC,BASE_ACC))\n",
        "\n",
        "# student model\n",
        "confs = np.mean(model_predictions, -1)\n",
        "preds = (confs >.5)\n",
        "precision, recall, thresholds = metrics.precision_recall_curve(true_Y, confs)\n",
        "plt.plot(recall, precision, c='red')\n",
        "MOD_AUPRC = auprc_threshs(confs, true_Y, AUCTHRESHS)\n",
        "MOD_AUROC = metrics.roc_auc_score(true_Y, confs)\n",
        "MOD_ACC = metrics.accuracy_score(true_Y, preds)\n",
        "print('Model:\\n{} (AUROC:{},ACC:{})'.format(MOD_AUPRC, MOD_AUROC,MOD_ACC))\n",
        "\n",
        "# average student model\n",
        "confs = np.mean(average_model_predictions, -1)\n",
        "preds = (confs >.5)\n",
        "precision, recall, thresholds = metrics.precision_recall_curve(true_Y, confs)\n",
        "plt.plot(recall, precision, c='pink')\n",
        "AVGMOD_AUPRC = auprc_threshs(confs, true_Y, AUCTHRESHS)\n",
        "AVGMOD_AUROC = metrics.roc_auc_score(true_Y, confs)\n",
        "AVGMOD_ACC = metrics.accuracy_score(true_Y, preds)\n",
        "print('Average Model:\\n{} (AUROC:{},ACC:{})'.format(AVGMOD_AUPRC, AVGMOD_AUROC,AVGMOD_ACC))\n",
        "\n",
        "# base+student model\n",
        "confs = (np.mean(base_predictions, -1)+np.mean(model_predictions, -1))/2.\n",
        "preds = (confs >.5)\n",
        "precision, recall, thresholds = metrics.precision_recall_curve(true_Y, confs)\n",
        "plt.plot(recall, precision, c='blue')\n",
        "EMOE_AUPRC = auprc_threshs(confs, true_Y, AUCTHRESHS)\n",
        "EMOE_AUROC = metrics.roc_auc_score(true_Y, confs)\n",
        "EMOE_ACC = metrics.accuracy_score(true_Y, preds)\n",
        "print('Base+Model:\\n{} (AUROC:{},ACC:{})'.format(EMOE_AUPRC, EMOE_AUROC,EMOE_ACC))\n",
        "\n",
        "# base+average student model\n",
        "confs = (np.mean(base_predictions, -1)+np.mean(average_model_predictions, -1))/2.\n",
        "preds = (confs >.5)\n",
        "precision, recall, thresholds = metrics.precision_recall_curve(true_Y, confs)\n",
        "plt.plot(recall, precision, c='cyan')\n",
        "AVGEMOE_AUPRC = auprc_threshs(confs, true_Y, AUCTHRESHS)\n",
        "AVGEMOE_AUROC = metrics.roc_auc_score(true_Y, confs)\n",
        "AVGEMOE_ACC = metrics.accuracy_score(true_Y, preds)\n",
        "print('Base+Average Model:\\n{} (AUROC:{},ACC:{})'.format(AVGEMOE_AUPRC, AVGEMOE_AUROC,AVGEMOE_ACC))\n",
        "\n",
        "plt.xticks(np.arange(0, 1.01, step=0.1))\n",
        "plt.xticks(np.arange(0, 1.01, step=0.05), minor=True)\n",
        "plt.yticks(np.arange(0.6, 1.01, step=0.05))\n",
        "plt.grid(True, which='both')\n",
        "plt.xlabel('Recall ({} Positive)'.format(int(np.sum(true_Y))))\n",
        "plt.ylabel('Precision')\n",
        "plt.title('ID Precision vs Recall')\n",
        "plt.legend(['Base', 'NNet', 'AVGNNET', 'EMOE', 'AVGEMOE'])\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "T2gDhB4RGfcS"
      },
      "outputs": [],
      "source": [
        "results_dict['heldood_base_AUPRC'] = BASE_AUPRC\n",
        "results_dict['heldood_model_AUPRC'] = MOD_AUPRC\n",
        "results_dict['heldood_avg_model_AUPRC'] = AVGMOD_AUPRC\n",
        "results_dict['heldood_emoe_AUPRC'] = EMOE_AUPRC\n",
        "results_dict['heldood_avg_emoe_AUPRC'] = AVGEMOE_AUPRC\n",
        "\n",
        "results_dict['heldood_base_AUROC'] = BASE_AUROC\n",
        "results_dict['heldood_model_AUROC'] = MOD_AUROC\n",
        "results_dict['heldood_avg_model_AUROC'] = AVGMOD_AUROC\n",
        "results_dict['heldood_emoe_AUROC'] = EMOE_AUROC\n",
        "results_dict['heldood_avg_emoe_AUROC'] = AVGEMOE_AUROC\n",
        "\n",
        "results_dict['heldood_base_ACC'] = BASE_ACC\n",
        "results_dict['heldood_model_ACC'] = MOD_ACC\n",
        "results_dict['heldood_avg_model_ACC'] = AVGMOD_ACC\n",
        "results_dict['heldood_emoe_ACC'] = EMOE_ACC\n",
        "results_dict['heldood_avg_emoe_ACC'] = AVGEMOE_ACC"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "caAfTf29GfcS"
      },
      "outputs": [],
      "source": [
        "with open(os.path.join(save_path, 'results.p'), 'wb') as f:\n",
        "    pickle.dump(results_dict, f)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "L4",
      "machine_shape": "hm",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
