projektor_cifar10_example.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "dc46435c-77b1-449a-bb35-4f625111a8ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "import numpy as np \n",
    "import pandas as pd \n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import cm\n",
    "from sklearn.metrics import r2_score, mean_squared_error, mean_absolute_error\n",
    "import math\n",
    "from sklearn.linear_model import LinearRegression\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9653d71-ad9e-452e-a01a-631017853e8e",
   "metadata": {},
   "source": [
    "# Create Training Data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "612cc297-6c0a-461d-832f-04fccbf338f2",
   "metadata": {},
   "source": [
    "Run the prep_train_data.py file to generate fitting data.\n",
    "```\n",
    "    python prep_train_data.py --n {n0} --cnum {cuda_num}\n",
    "    python prep_train_data.py --n {n1} --cnum {cuda_num}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb1712f2-ab56-4261-8a1b-d40437b28a93",
   "metadata": {},
   "source": [
    "# Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "20bedf3b-3127-4667-b469-5102d20d0e08",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_qs(reserrlog, otlog, accs):\n",
    "\n",
    "    qsreserrlog = []\n",
    "    for i in range(len(reserrlog)):\n",
    "        cur_arr = np.array(reserrlog[i])\n",
    "        len_cur_arr = len(cur_arr)\n",
    "    #     print(cur_arr)\n",
    "    #     print(len_cur_arr)\n",
    "        if len_cur_arr:\n",
    "            len_to_add = 11 - len_cur_arr\n",
    "            if len_to_add:\n",
    "                to_add = np.zeros(len_to_add) + 0\n",
    "                cur_arr = np.concatenate((cur_arr, to_add))\n",
    "            qsreserrlog.append(cur_arr)\n",
    "\n",
    "\n",
    "    #         print(len(cur_arr))\n",
    "    qsreserrlog = np.array(qsreserrlog)\n",
    "\n",
    "    qsotlog = []\n",
    "    for i in range(len(otlog)):\n",
    "        cur_arr = np.array(otlog[i])\n",
    "        len_cur_arr = len(cur_arr)\n",
    "        if len_cur_arr:\n",
    "            len_to_add = 11 - len_cur_arr\n",
    "            if len_to_add:\n",
    "                to_add = np.zeros(len_to_add) + 0\n",
    "                cur_arr = np.concatenate((cur_arr, to_add))\n",
    "            qsotlog.append(cur_arr)\n",
    "\n",
    "    qsaccs = []\n",
    "    for i in range(len(accs)):\n",
    "        cur_arr = np.array(accs[i])\n",
    "        len_cur_arr = len(cur_arr)\n",
    "        if len_cur_arr:\n",
    "            len_to_add = 11 - len_cur_arr\n",
    "            if len_to_add:\n",
    "                to_add = np.zeros(len_to_add) + 0\n",
    "                cur_arr = np.concatenate((cur_arr, to_add))\n",
    "            qsaccs.append(cur_arr)\n",
    "\n",
    "    #         print(len(cur_arr))\n",
    "    qsotlog = np.array(qsotlog)\n",
    "    qsaccs = np.array(qsaccs)\n",
    "    \n",
    "    return qsreserrlog, qsotlog, qsaccs\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "891093af-62e2-4ed0-a16e-cc7a0bd9e4a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "size = 1000\n",
    "reserrlog, otlog, accs = pickle.load(open(f\"projektor_data/cif10_3sources_unbalanced_{size}_br_10_rep_5.res\", 'rb'))\n",
    "qsreserrlog10, qsotlog10, qsaccs10 = get_qs(reserrlog, otlog, accs) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcea1375-f59c-4b29-a799-ea785e50d750",
   "metadata": {},
   "outputs": [],
   "source": [
    "size = 1500\n",
    "reserrlog, otlog, accs = pickle.load(open(f\"projektor_data/cif10_3sources_unbalanced_{size}_br_10_rep_5.res\", 'rb'))\n",
    "qsreserrlog15, qsotlog15, qsaccs15 = get_qs(reserrlog, otlog, accs) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "514d07cb-0d26-4faf-ad30-b27f29885028",
   "metadata": {},
   "source": [
    "# Fit Data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69252e62-af6f-44cc-a673-7a7f81341868",
   "metadata": {},
   "source": [
    "## Create Mask for Test Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02714cba-c2bb-4efb-859f-aa8b4cdae9c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# q1 mask\n",
    "size1 = qsreserrlog.shape[0]\n",
    "size2 = qsreserrlog.shape[1]\n",
    "q1mask1 = np.zeros([size1,size1])\n",
    "q1mask2 = np.zeros([size1,size1])\n",
    "thr = 5\n",
    "for l1 in range(size1):\n",
    "    for l2 in range(size1):\n",
    "        if qsreserrlog[l1,l2] != 0:\n",
    "            if l2 < thr:\n",
    "                q1mask1[l1, l2] = 1\n",
    "            if l1+l2 > 10:\n",
    "                q1mask1[l1, l2] = 0\n",
    "            if l2 >= thr:\n",
    "                q1mask2[l1, l2] = 1\n",
    "            if l1+l2 > 10:\n",
    "                q1mask2[l1, l2] = 0\n",
    "            \n",
    "q1_mask_train_inds = q1mask1\n",
    "q1_mask_pred_inds = q1mask2\n",
    "\n",
    "q1mask_train = np.nonzero(q1_mask_train_inds != 0)\n",
    "q1mask_pred = np.nonzero(q1_mask_pred_inds != 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03a4b05f-e901-4c87-b7ff-49080d200cd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "A1 = np.zeros((size1,size2))\n",
    "A2 = np.zeros((size1,size2))\n",
    "\n",
    "# Fill each 1-dimension subarray with [0, 1, 2, 3, 4, 5]\n",
    "for i1 in range(size1):\n",
    "    for i2 in range(size2):\n",
    "        A1[i1, i2] = i1/(size1-1)\n",
    "        A2[i1, i2] = i2/(size1-1)\n",
    "print(A2)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "370743a5-3607-4cd3-90ac-ca7795889ebf",
   "metadata": {},
   "source": [
    "## Fit Data with $\\texttt{Projektor}$ for performance prediction for extrapolation\n",
    "$$\\hat{\\mathcal{L}}(\\mathcal{A}(\\mathcal{D}(N, \\mathbf{p})),D^\\text{val})=\\sum_{i=1}^m (b_2^i\\cdot p_i^2+b_1^i\\cdot p_i+b_0)\\cdot \\text{OT}(\\mathcal{D}(N, \\mathbf{p}), D^\\text{val})+\\sum_{i=1}^m (c_2^i\\cdot p_i^2+c_1^i\\cdot p_i+c_0)$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1167a6d5-7d22-4730-9159-495d18b76d9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.linear_model import LinearRegression\n",
    "npqsres = np.nan_to_num(np.array(qsreserrlog))\n",
    "npqsaccs = np.nan_to_num(np.array(qsaccs))\n",
    "npqsop = np.nan_to_num(np.array(qsotlog))\n",
    "\n",
    "# mask q1mask2 set to 0 q1mask2/q1mask_pred\n",
    "npqsres[q1mask_pred] = 0\n",
    "npqsop[q1mask_pred] = 0\n",
    "npqsaccs[q1mask_pred] = 0\n",
    "\n",
    "# npqsres[cover_train_mask] = 0\n",
    "\n",
    "res_non_zero_indices = np.nonzero(npqsres != 0)\n",
    "accs_non_zero_indices = np.nonzero(npqsaccs != 0)\n",
    "ot_non_zero_indices = np.nonzero(npqsop != 0)\n",
    "# Extract the non-zero elements\n",
    "res_non_zero_elements = np.log(npqsres[res_non_zero_indices])\n",
    "accs_non_zero_elements = npqsaccs[ot_non_zero_indices]\n",
    "ot_non_zero_elements = npqsop[ot_non_zero_indices]\n",
    "\n",
    "scalingmat1 = A1\n",
    "scalingmat2 = A2\n",
    "scalingmat3 = 1 - scalingmat1 - scalingmat2 \n",
    "\n",
    "qss1 = scalingmat1[ot_non_zero_indices]\n",
    "qss2 = scalingmat2[ot_non_zero_indices]\n",
    "qss3 = scalingmat3[ot_non_zero_indices]\n",
    "\n",
    "\n",
    "y2fit = np.array(accs_non_zero_elements.reshape(-1, 1))\n",
    "x2fit = np.zeros([np.shape(accs_non_zero_indices)[1], 8])\n",
    "x2fit[:, 0] = ot_non_zero_elements\n",
    "x2fit[:, 1] = qss1\n",
    "x2fit[:, 2] = qss2\n",
    "x2fit[:, 3] = qss1*ot_non_zero_elements\n",
    "x2fit[:, 4] = qss2*ot_non_zero_elements\n",
    "x2fit[:, 5] = qss2*qss2*ot_non_zero_elements\n",
    "x2fit[:, 6] = qss1*qss1*ot_non_zero_elements\n",
    "\n",
    "model = LinearRegression()\n",
    "model.fit(x2fit,y2fit)\n",
    "\n",
    "model_predicted_train = model.predict(x2fit).flatten()\n",
    "y_train = y2fit.flatten()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bae62ce7-ed3e-4220-8673-325f9940d82f",
   "metadata": {},
   "source": [
    "## Predict Performance of Unseen $\\textbf{p}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a098b9c-0d76-4b37-a6eb-36481c07918d",
   "metadata": {},
   "outputs": [],
   "source": [
    "npqsres = np.nan_to_num(np.array(qsreserrlog))\n",
    "npqsaccs = np.nan_to_num(np.array(qsaccs))\n",
    "npqsop = np.nan_to_num(np.array(qsotlog))\n",
    "\n",
    "# mask q1mask2 set to 0 q1mask2/q1mask_pred\n",
    "# npqsres[q1mask_pred] = 0\n",
    "# npqsop[q1mask_pred] = 0\n",
    "# npqsaccs[q1mask_pred] = 0\n",
    "\n",
    "# npqsres[cover_train_mask] = 0\n",
    "\n",
    "res_non_zero_indices = np.nonzero(npqsres != 0)\n",
    "accs_non_zero_indices = np.nonzero(npqsaccs != 0)\n",
    "ot_non_zero_indices = np.nonzero(npqsop != 0)\n",
    "# Extract the non-zero elements\n",
    "res_non_zero_elements = np.log(npqsres[res_non_zero_indices])\n",
    "accs_non_zero_elements = npqsaccs[ot_non_zero_indices]\n",
    "ot_non_zero_elements = npqsop[ot_non_zero_indices]\n",
    "\n",
    "scalingmat1 = A1\n",
    "scalingmat2 = A2\n",
    "scalingmat3 = 1 - scalingmat1 - scalingmat2 \n",
    "\n",
    "qss1 = scalingmat1[ot_non_zero_indices]\n",
    "qss2 = scalingmat2[ot_non_zero_indices]\n",
    "qss3 = scalingmat3[ot_non_zero_indices]\n",
    "\n",
    "y2fit = np.array(accs_non_zero_elements.reshape(-1, 1))\n",
    "x2fit = np.zeros([np.shape(accs_non_zero_indices)[1], 8])\n",
    "x2fit[:, 0] = ot_non_zero_elements\n",
    "x2fit[:, 1] = qss1\n",
    "x2fit[:, 2] = qss2\n",
    "x2fit[:, 3] = qss1*ot_non_zero_elements\n",
    "x2fit[:, 4] = qss2*ot_non_zero_elements\n",
    "x2fit[:, 5] = qss2*qss2*ot_non_zero_elements\n",
    "x2fit[:, 6] = qss1*qss1*ot_non_zero_elements\n",
    "\n",
    "otqfit = model.predict(x2fit)\n",
    "print(model.coef_)\n",
    "resot = np.zeros([size1,size1])\n",
    "resot[ot_non_zero_indices] = otqfit.flatten()\n",
    "\n",
    "model_predicted_test = resot[q1mask_pred]\n",
    "y_test = npqsaccs[q1mask_pred] # np.log(npqsres[q1mask_pred])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9b71e94-ffff-489f-8e1d-be8c133063a6",
   "metadata": {},
   "source": [
    "# Plot Fitted Data and Performance Prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a70458d-09eb-4f5d-a963-45ee1967f0df",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.scatter(model_predicted_train, y_train, marker='o', label='Training', color='blue', alpha=0.7)\n",
    "plt.scatter(model_predicted_test, y_test, marker='o', label='Prediction (PQ-OT)', color='orange', alpha=0.7)\n",
    "# plt.xlabel('pseudo-quadratic OT')\n",
    "plt.xlabel('pseudo-quadratic OT prediction')\n",
    "plt.ylabel('model err (log)')\n",
    "\n",
    "# calculate the slope and intercept of the linear regression line\n",
    "pltslope, pltintercept =[1, 0]\n",
    "# create a line based on the linear regression equation\n",
    "x_line = np.linspace(0, 4, 100)\n",
    "y_line = pltslope * x_line + pltintercept\n",
    "# plot the line\n",
    "plt.plot(x_line, y_line, c='orange')\n",
    "# plt.xlim([1, 3])\n",
    "# plt.ylim([1, 3])\n",
    "# plt.xlim([-3, 0.5])\n",
    "# plt.ylim([-3, 0.5])\n",
    "\n",
    "# show the plot\n",
    "plt.legend()\n",
    "# plt.title('Interpolation on Cifar 10 to Smaller Values by Centering Scaling')\n",
    "plt.title(f'Extrapolation on DataSource#3(>{thr/size1 * 100:.0f}) by PQ-OT')\n",
    "plt.show()\n",
    "\n",
    "Y_train = y_train\n",
    "Y_predict_train = model_predicted_train\n",
    "mse_train = mean_squared_error(Y_train, Y_predict_train)\n",
    "\n",
    "print(\"Train MSE: \", mse_train)\n",
    "\n",
    "train_model_mean_perf_diff = mean_absolute_error(Y_predict_train - Y_train)\n",
    "print(\"Train performance mean difference: \", train_model_mean_perf_diff)\n",
    "\n",
    "\n",
    "print(\"--------TEST-----------\")\n",
    "\n",
    "Y_test = y_test\n",
    "Y_predicted = model_predicted_test\n",
    "mse_pred = mean_squared_error(Y_test, Y_predicted)\n",
    "test_model_mean_perf_diff = mean_absolute_error(Y_predicted - Y_test)\n",
    "print(\"Test performance mean difference: \", test_model_mean_perf_diff)\n",
    "\n",
    "print(\"R2: \", r2_score(Y_test, Y_predicted))\n",
    "print(\"MAE:\", mean_absolute_error(Y_test, Y_predicted))\n",
    "print(\"MSE: \", mean_squared_error(Y_test, Y_predicted))\n",
    "print(\"RMSE: \", math.sqrt(mean_squared_error(Y_test, Y_predicted)))\n",
    "# pqot_mse.append(mean_squared_error(Y_test, Y_predicted))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "343a7490-4002-49ba-89ce-a717223b745c",
   "metadata": {},
   "source": [
    "# Performance Projection"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0483a3f-f6ad-4bd0-bd48-99075152de0a",
   "metadata": {},
   "source": [
    "Fit above function for two datasets of size $N_0$ and $N_1$.\n",
    "Then, project performance onto large data scales, using the equation below:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74a6556a-81d4-48bd-8739-88ae2fa80b25",
   "metadata": {},
   "source": [
    "$$\\hat{\\mathcal{L}}(\\mathcal{A}(\\mathcal{D}(N,\\mathbf{p}); D^\\text{val}) = \\left(\\log\\frac{N_1}{N_0}\\right)^{-1}\\left[ \\log\\frac{ N}{N_0}\\hat{\\mathcal{L}}(\\mathcal{A}(\\mathcal{D}(N_1,\\mathbf{p})); D^\\text{val})-\\log\\frac{ N}{N_1}\\hat{\\mathcal{L}}(\\mathcal{A}(\\mathcal{D}(N_0,\\mathbf{p}); D^\\text{val})\\right]$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3032d06-7f9e-415c-8d71-7d4d55607984",
   "metadata": {},
   "source": [
    "### Example for projecting performance onto N from $N_0$ and $N_1$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64885fed-f5ae-43dd-ac38-4c061e8963b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 10000\n",
    "n0 = 1000\n",
    "acc0 = A # Accuracy of a model trained on N_0 for a given p\n",
    "\n",
    "n1 = 1500\n",
    "acc1 = B # Accuracy of a model trained on N_1 for a given p\n",
    "\n",
    "# acc_10k = np.log(n1/n0) * (n/n0* acc1 -n/n1 * acc0)\n",
    "acc_10k =  (1/np.log(n1/n0)) * (np.log(n/n0)* acc1 - np.log(n/n1) * acc0)\n",
    "\n",
    "print(acc_10k)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f3f3c1b-45d9-44f2-bac9-66aa4794c0bd",
   "metadata": {},
   "source": [
    "#### To verify projection performance, we can actually train a model on $N$ points and have a direct comparison.\n",
    "#### To do so, we can generate data points for size $N$:\n",
    "    \n",
    "```\n",
    "    python prep_train_data.py --n {N} --cnum {cuda_num}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b45eba3-0658-4587-b81f-14d22a3dadac",
   "metadata": {},
   "source": [
    "# Optimal Data Source Selection"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2b80e44-a4b4-48a4-9ace-f7b38e541ef8",
   "metadata": {},
   "source": [
    "We proceed to optimal data source selection after we have fitted function for data scales $N_0$ and $N_1$. To find the optimal mixing ratio $\\textbf{p}^*$, we adopt a gradient-based method with each update as follows:\n",
    "\n",
    "$$\\mathbf{p^{t+1}}\\leftarrow\\mathbf{p^{t}} + d^t\\cdot \\left.\\frac{\\partial \\hat{\\mathcal{L}}(\\mathcal{A}(\\mathcal{D}(N_s,\\mathbf{p})),D^\\text{val})}{\\partial \\mathbf{p}}\\right|_{\\mathbf{p}=\\mathbf{p}^t},$$\n",
    "\n",
    "where\n",
    "\n",
    "$$\\frac{\\partial \\hat{\\mathcal{L}}(\\mathcal{A}(\\mathcal{D}(N,\\mathbf{p}); D^\\text{val})}{\\partial \\mathbf{p}}=\n",
    "     \\left(\\log\\frac{N_1}{N_0}\\right)^{-1}\\left[ \\log\\frac{ N}{N_0}\\frac{\\partial \\hat{\\mathcal{L}}(\\mathcal{A}(\\mathcal{D}(N_1,\\mathbf{p})); D^\\text{val})}{\\partial \\mathbf{p}}-\\log\\frac{ N}{N_1}\\frac{\\partial\\hat{\\mathcal{L}}(\\mathcal{A}(\\mathcal{D}(N_0,\\mathbf{p}); D^\\text{val})}{\\partial \\mathbf{p}}\\right].$$\n",
    "     \n",
    "Each gradient of $\\hat{\\mathcal{L}}$ can be defined as below:\n",
    "\n",
    "$$\\frac{\\partial \\hat{\\mathcal{L}}(\\mathcal{A}(\\mathcal{D}(N,\\mathbf{p}); D^\\text{val})}{\\partial p_i} =  (b_2^i\\cdot p_i^2+b_1^i\\cdot p_i+b_0)\\cdot\\frac{\\partial \\text{OT}(\\mathcal{A}(\\mathcal{D}(N,\\mathbf{p}); D^\\text{val})}{\\partial p_i} \n",
    "    +[b_2^i\\cdot (2p_i)+b_1^i]\\cdot \\text{OT}(\\mathcal{A}(\\mathcal{D}(N,\\mathbf{p}); D^\\text{val}) + [c_2^i\\cdot (2p_i)+c_1^i]$$\n",
    "    \n",
    "and each gradient of OT can be efficiently computed as follows:\n",
    "\n",
    "$$\\frac{\\partial \\text{OT}(\\mathcal{A}(\\mathcal{D}(N,\\mathbf{p}); D^\\text{val})}{\\partial p\n",
    "    _i}=\\frac{1}{n_i}\\left(\\sum_{j=1}^{n_i} f_i^j - \\frac{n_i}{N-n_i}\\sum_{x=\\{1...m\\}\\setminus i}\\sum_{y=1}^{n_x} f_x^y\\right).$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e9e39c0-0843-445d-a09e-ec6e75fbd9d3",
   "metadata": {},
   "source": [
    "### We provide a full code for the optimization step below (whereas all helper functions are provided in the prep_train_data.py file.\n",
    "This code follows the above example for $N=10,000$ and data scales $N_0 = 1,000$ and $N_1 = 1,500$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "928e9da4-641e-4924-9ec8-30985eae0eb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "n = 10000\n",
    "n0 = 1000\n",
    "n1 = 1500\n",
    "# q = 0.0\n",
    "batch_size = 256\n",
    "\n",
    "\n",
    "# make test dataloader\n",
    "\n",
    "test_loader = torch.utils.data.DataLoader(dataset=TensorDataset(torch.Tensor(test_features).permute(0,3,1,2), torch.LongTensor(test_labels)), \n",
    "                                                batch_size=batch_size, \n",
    "                                                shuffle=False)\n",
    "\n",
    "gdreserrlog = []\n",
    "gdresacc = []\n",
    "gdotlog = []\n",
    "gd1log = []\n",
    "gd2log = []\n",
    "\n",
    "q1 = 1/3\n",
    "q2 = 1/3\n",
    "q3 = 1-q1-q2\n",
    "\n",
    "q1log = [q1]\n",
    "q2log = [q2]\n",
    "\n",
    "lr = 5e-4\n",
    "reps = 1\n",
    "iters = 100\n",
    "all_qs = []\n",
    "\n",
    "print(\"iters: \", iters)\n",
    "print(\"reps: \", reps)\n",
    "for iteration in range(iters):\n",
    "    cacheerr = []\n",
    "    cacheacc = []\n",
    "    cacheot = []\n",
    "    cachegd1 = []\n",
    "    cachegd2 = []\n",
    "    cacheot_n1 = []\n",
    "    cachegd1_n1 = []\n",
    "    cachegd2_n1 = []\n",
    "    \n",
    "    start_t = time.time()\n",
    "    all_qs.append((q1,q2,q3))\n",
    "\n",
    "    for rep in range(reps):\n",
    "\n",
    "        # create dataset\n",
    "        train_x, train_y, ds1_len, ds2_len, ds3_len = dataset_q(q1, q2, n0, train_features, train_labels)\n",
    "        # make train dataloader\n",
    "        train_loader, train_loader_ot, test_loader = model_loaders(train_x, train_y,  test_features, test_labels, n0,transform)\n",
    "\n",
    "        \n",
    "        print(f\"Lengths: S1={ds1_len}, S2={ds2_len}, S3={ds3_len}\")\n",
    "        # get OT dist\n",
    "        ot_dist, dual_sol = get_ot_dist(train_loader_ot, test_loader, n=n0)\n",
    "        cacheot.append(ot_dist)\n",
    "\n",
    "        # create dataset\n",
    "        train_x_n1, train_y_n1, ds1_len_n1, ds2_len_n1, ds3_len_n1 = dataset_q(q1, q2, n1, train_features, train_labels)\n",
    "        # make train dataloader\n",
    "        train_loader_n1, train_loader_ot_n1, test_loader_n1 = model_loaders(train_x_n1, train_y_n1,  test_features, test_labels, n1, transform)\n",
    "\n",
    "        \n",
    "        print(f\"_n1 Lengths: S1={ds1_len_n1}, S2={ds2_len_n1}, S3={ds3_len_n1}\")\n",
    "        # get OT dist\n",
    "        ot_dist_n1, dual_sol_n1 = get_ot_dist(train_loader_ot_n1, test_loader_n1, n=n1)\n",
    "        cacheot_n1.append(ot_dist_n1)\n",
    "    \n",
    "    \n",
    "        train_x_n, train_y_n, _, _, _ = dataset_q(q1, q2, n, train_features, train_labels)\n",
    "        # make train dataloader\n",
    "        train_loader_n, train_loader_ot_n, test_loader_n = model_loaders(train_x_n, train_y_n,  test_features, test_labels, n,transform)\n",
    "\n",
    "    \n",
    "#         get model error (Optional: Can be commented our for faster optimization)\n",
    "        errloss, acc = get_model_log_err(train_loader_n, test_loader_n)\n",
    "        cacheerr.append(errloss)\n",
    "        cacheacc.append(acc)\n",
    "#         cacheerr.append(0) #get_model_log_err(train_loader, test_loader))\n",
    "        \n",
    "        \n",
    "        full_ds_len = np.shape(train_x)[0]\n",
    "        \n",
    "        full_ds_len_n1 = np.shape(train_x_n1)[0]\n",
    "        \n",
    "        \n",
    "        q1_end_pos = ds1_len\n",
    "        q2_end_pos = q1_end_pos + ds2_len\n",
    "        q3_end_pos = q2_end_pos + ds3_len\n",
    "        \n",
    "        q1_end_pos_n1 = ds1_len_n1\n",
    "        q2_end_pos_n1 = q1_end_pos_n1 + ds2_len_n1\n",
    "        q3_end_pos_n1 = q2_end_pos_n1 + ds3_len_n1\n",
    "        \n",
    "        fall = dual_sol\n",
    "        fds1= fall[0:q1_end_pos]\n",
    "        fds2 = fall[q1_end_pos:q2_end_pos]\n",
    "        fds3 = fall[q2_end_pos:q3_end_pos]\n",
    "        \n",
    "        fall_n1 = dual_sol_n1\n",
    "        fds1_n1= fall_n1[0:q1_end_pos_n1]\n",
    "        fds2_n1 = fall_n1[q1_end_pos_n1:q2_end_pos_n1]\n",
    "        fds3_n1 = fall_n1[q2_end_pos_n1:q3_end_pos_n1]\n",
    "        \n",
    "        # Calibrated Gradient for Q1\n",
    "        if full_ds_len-ds1_len != 0:\n",
    "            calib_q1 = np.sum(fds1) - (np.sum(fall) - np.sum(fds1))*(ds1_len/(full_ds_len-ds1_len))\n",
    "        else:\n",
    "            calib_q1 = 0\n",
    "\n",
    "        if full_ds_len_n1-ds1_len_n1 != 0:\n",
    "            calib_q1_n1 = np.sum(fds1_n1) - (np.sum(fall_n1) - np.sum(fds1_n1))*(ds1_len_n1/(full_ds_len_n1-ds1_len_n1))\n",
    "        else:\n",
    "            calib_q1_n1 = 0\n",
    "            \n",
    "        # Calibrated Gradient for Q2    \n",
    "        if full_ds_len-ds2_len != 0:\n",
    "            calib_q2 = np.sum(fds2) - (np.sum(fall) - np.sum(fds2))*(ds2_len/(full_ds_len-ds2_len))\n",
    "        else:\n",
    "            calib_q2 = 0\n",
    "            \n",
    "            \n",
    "        if full_ds_len_n1-ds2_len_n1 != 0:\n",
    "            calib_q2_n1 = np.sum(fds2_n1) - (np.sum(fall_n1) - np.sum(fds2_n1))*(ds2_len_n1/(full_ds_len_n1-ds2_len_n1))\n",
    "        else:\n",
    "            calib_q2_n1 = 0\n",
    "            \n",
    "        # Gradient for OT for Q1\n",
    "        if ds1_len != 0:\n",
    "            calib_q1 /= ds1_len\n",
    "            gd_q1 = (b2q1*q1*q1+b1q1*q1+b0)*calib_q1+(b2q1*2*q1+b1q1)*ot_dist+(c2q1*2*q1+c1q1)\n",
    "        else:\n",
    "            gd_q1 = 0\n",
    "            \n",
    "           \n",
    "        if ds1_len_n1 != 0:\n",
    "            calib_q1_n1 /= ds1_len_n1\n",
    "            gd_q1_n1 = (b2q1_n1*q1*q1+b1q1_n1*q1+b0_n1) *calib_q1_n1+(b2q1_n1*2*q1+b1q1_n1)*ot_dist_n1+(c2q1_n1*2*q1+c1q1_n1)\n",
    "        else:\n",
    "            gd_q1_n1 = 0\n",
    "        \n",
    "        # Gradient for OT for Q2\n",
    "        if ds2_len != 0:\n",
    "            calib_q2 /= ds2_len\n",
    "            gd_q2 = (b2q2*q2*q2+b1q2*q2+b0)*calib_q2+(b2q2*2*q2+b1q2)*ot_dist+(c2q2*2*q2+c1q2)\n",
    "        \n",
    "            \n",
    "        if ds2_len_n1 != 0:\n",
    "            calib_q2_n1 /= ds2_len_n1\n",
    "            gd_q2_n1 = (b2q2_n1*q2*q2+b1q2_n1*q2+b0_n1) *calib_q2_n1+(b2q2_n1*2*q2+b1q2_n1)*ot_dist_n1+(c2q2_n1*2*q2+c1q2_n1)\n",
    "        \n",
    "        \n",
    "        if (gd_q1 == 0) & (gd_q2 == 0):\n",
    "            print(\"zero gradient!\")\n",
    "            break\n",
    "        if (gd_q1_n1 == 0) & (gd_q2_n1 == 0):\n",
    "            print(\"zero gradient!\")\n",
    "            break\n",
    "        cachegd1.append(gd_q1)\n",
    "        cachegd2.append(gd_q2)\n",
    "        cachegd1_n1.append(gd_q1_n1)\n",
    "        cachegd2_n1.append(gd_q2_n1)\n",
    "        \n",
    "    term1 = 1/np.log(n1/n0)\n",
    "    term2 = np.log(n/n0)\n",
    "    term3 = np.log(n/n1)\n",
    "    \n",
    "    grad1 = term1 * (term2*np.median(cachegd1_n1) - term3*np.median(cachegd1))\n",
    "    grad2 = term1 * (term2*np.median(cachegd2_n1) - term3*np.median(cachegd2))\n",
    "    step = lr*np.power(0.90, iteration)\n",
    "    q1 = q1 + gd_q1*step\n",
    "    q2 = q2 + gd_q2*step\n",
    "    q3 = 1-q1-q2\n",
    "    q1log.append(q1)\n",
    "    q2log.append(q2)\n",
    "    gdreserrlog.append(np.median(cacheerr))\n",
    "    gdotlog.append(np.median(cacheot))\n",
    "    gdresacc.append(np.median(cacheacc))\n",
    "    gd1log.append(grad1)\n",
    "    gd2log.append(grad2)\n",
    "    print(np.median(cacheot))\n",
    "    print(\"cacheerr: \", cacheerr)\n",
    "    print(\"cacheot: \", cacheot)\n",
    "    print(\"cacheacc: \", cacheacc)   \n",
    "    print(\"iter: \", iteration, \" it took: \", time.time() - start_t)\n",
    "    \n",
    "    print(\"\\n New ITER\")\n",
    "    print(\"GRADIENT: \", gd_q1, gd_q2, step)\n",
    "    print(f'Q1 {q1}, Q2 {q2}, Q3 {q3}')\n",
    "    \n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "otdd",
   "language": "python",
   "name": "otdd"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}