{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "developmental-reserve",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "data_tr = torch.load(\"training-test-data/training_data.txt\")\n",
    "target_tr = torch.load(\"training-test-data/training_targets.txt\")\n",
    "data_test = torch.load(\"training-test-data/test_data.txt\")\n",
    "target_test = torch.load(\"training-test-data/test_targets.txt\")\n",
    "\n",
    "num_samples = 512 # number of training samples\n",
    "num_samples_test = 100 # number of test samples\n",
    "new_dim1 = 32 * 1 # first dimension\n",
    "new_dim2 = 32 * 1 # second dimension\n",
    "old_dim = 32 # CIFAR original dimension\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "contained-usage",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# find the minimum norm solution \n",
    "\n",
    "import cvxpy as cp\n",
    "import math\n",
    "\n",
    "\n",
    "# beta_cv_pred = np.zeros((new_dim1, new_dim2))\n",
    "\n",
    "def solve_cvxpy(option):\n",
    "    global beta_cv_pred_0;\n",
    "    global beta_cv_pred_1;\n",
    "    global beta_cv_pred_2;\n",
    "    data_tr_numpy =data_tr.data.numpy()\n",
    "    target_tr_numpy =target_tr.data.numpy()\n",
    "\n",
    "    print(data_tr.shape)\n",
    "    print(target_tr.shape)\n",
    "\n",
    "    dft_matrix1 = np.fft.fft(np.eye(new_dim1)) / np.sqrt(new_dim1)\n",
    "    dft_matrix2 = np.fft.fft(np.eye(new_dim2)) / np.sqrt(new_dim2)\n",
    "\n",
    "    beta_cv_0 = cp.Variable((new_dim1, new_dim2))\n",
    "    beta_cv_1 = cp.Variable((new_dim1, new_dim2))\n",
    "    beta_cv_2 = cp.Variable((new_dim1, new_dim2))\n",
    "    beta_cv = cp.Variable((new_dim1 * new_dim2, 3)) # Only going to used for ell2ell1 \n",
    "    function = 0 \n",
    "\n",
    "  # ell2-ell1 norm\n",
    "\n",
    "    if option == \"ell2ell2\":\n",
    "        function += 2 * cp.norm(beta_cv, 'fro')\n",
    "\n",
    "    if option == \"ell2ell1\":\n",
    "        for z in range(new_dim1 * new_dim2):\n",
    "            function += 2 * cp.norm(beta_cv[z]) \n",
    "\n",
    "  # Sum of the ell1 norms\n",
    "    if option == \"ell1\":\n",
    "        function += 2 * cp.sum(cp.abs(dft_matrix1 @ beta_cv_0 @ dft_matrix2.T))\n",
    "        function += 2 * cp.sum(cp.abs(dft_matrix1 @ beta_cv_1 @ dft_matrix2.T))\n",
    "        function += 2 * cp.sum(cp.abs(dft_matrix1 @ beta_cv_2 @ dft_matrix2.T))\n",
    "    objective = cp.Minimize(function)\n",
    "  \n",
    "  # Set up constraints\n",
    "    constraints = [target_tr[i] * cp.sum(cp.multiply(beta_cv_0, data_tr_numpy[i][0]) + cp.multiply(beta_cv_1, data_tr_numpy[i][1]) + cp.multiply(beta_cv_2, data_tr_numpy[i][2])) >= 1 for i in range(num_samples)] # margin >=1 \n",
    "    if option == \"ell2ell1\" or option == \"ell2ell2\":\n",
    "        beta_cv_fft_0 = cp.vec(cp.abs(dft_matrix1 @ beta_cv_0 @ dft_matrix2.T))\n",
    "        beta_cv_fft_1 = cp.vec(cp.abs(dft_matrix1 @ beta_cv_1 @ dft_matrix2.T))\n",
    "        beta_cv_fft_2 = cp.vec(cp.abs(dft_matrix1 @ beta_cv_2 @ dft_matrix2.T))\n",
    "        beta_cv_tr = beta_cv.T\n",
    "        constraints = constraints + [beta_cv_fft_0 <= beta_cv_tr[0]]\n",
    "        constraints = constraints + [beta_cv_fft_1 <= beta_cv_tr[1]]\n",
    "        constraints = constraints + [beta_cv_fft_2 <= beta_cv_tr[2]]\n",
    "\n",
    "    prob = cp.Problem(objective, constraints)\n",
    "    print(\"set up problem\")\n",
    "    result = prob.solve(solver=cp.SCS)\n",
    "    print(\"solved problem\")\n",
    "    print(result)\n",
    "    beta_cv_pred0 = beta_cv_0.value\n",
    "    beta_cv_pred1 = beta_cv_1.value\n",
    "    beta_cv_pred2 = beta_cv_2.value\n",
    "    betas = [beta_cv_pred0, beta_cv_pred1, beta_cv_pred2]\n",
    "    return betas\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "# # beta_cv_pred_ell1 = solve_cvxpy(\"ell1\")\n",
    "# # data_np =data_tr.data.numpy()\n",
    "# # preds = [np.tensordot(data_np[i], beta_cv_pred) for i in range(num_samples)]\n",
    "# # margins = [target_tr[i] * preds[i] for i in range(num_samples)]\n",
    "# # max_margin = max(margins).data.numpy()\n",
    "# # print(max_margin)\n",
    "\n",
    "# plt.imshow(beta_cv_pred_ell2, cmap='gray', interpolation='none')\n",
    "# plt.show()\n",
    "\n",
    "# beta_cv_pred_ell1 = solve_cvxpy(\"ell1\")\n",
    "\n",
    "# plt.imshow(beta_cv_pred_ell1 , cmap='gray', interpolation='none')\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "brutal-atlas",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([512, 3, 32, 32])\n",
      "torch.Size([512])\n",
      "set up problem\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.9/site-packages/cvxpy/problems/problem.py:1267: UserWarning: Solution may be inaccurate. Try another solver, adjusting the solver settings, or solve with verbose=True for more information.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "solved problem\n",
      "82.84792164886227\n",
      "torch.Size([512, 3, 32, 32])\n",
      "torch.Size([512])\n",
      "set up problem\n",
      "solved problem\n",
      "4.753747102981253\n",
      "torch.Size([512, 3, 32, 32])\n",
      "torch.Size([512])\n",
      "set up problem\n",
      "solved problem\n",
      "122.14803636375322\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "\n",
    "betas = solve_cvxpy(\"ell2ell1\")\n",
    "for i in range(3):\n",
    "    name = \"experiments-data/cvxpy-ell2ell1-inputchannel\" + str(i)\n",
    "    pd.DataFrame(betas[i]).to_csv(name, header=False, index=False)\n",
    "    \n",
    "betas  = solve_cvxpy(\"ell2ell2\")\n",
    "for i in range(3):\n",
    "    name = \"experiments-data/cvxpy-ell2ell2-inputchannel\" + str(i)\n",
    "    pd.DataFrame(betas[i]).to_csv(name, header=False, index=False)\n",
    "\n",
    "betas  = solve_cvxpy(\"ell1\")\n",
    "for i in range(3):\n",
    "    name = \"experiments-data/cvxpy-ell1-inputchannel\" + str(i)\n",
    "    pd.DataFrame(betas[i]).to_csv(name, header=False, index=False)\n",
    "\n",
    "# plt.imshow(beta_cv_pred_ell1 , cmap='gray', interpolation='none')\n",
    "# plt.show()\n",
    "\n",
    "# beta_cv_pred_ell1 = solve_cvxpy(\"ell1\")\n",
    "# print(beta_cv_pred_ell1)\n",
    "# name = \"experiments-data/cvxpy-ell1\"\n",
    "# pd.DataFrame(beta_cv_pred_ell1).to_csv(name, header=False, index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "phantom-april",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "beta_to_plot_time = []\n",
    "beta_to_plot_freq = []\n",
    "beta_to_plot_time.append(beta_ell2)\n",
    "#         print(beta.shape)\n",
    "#         print(np.fft.fft2(beta, norm='ortho').shape)\n",
    "beta_to_plot_freq.append(np.abs(np.fft.fft2(beta_ell2, norm='ortho')))\n",
    "\n",
    "beta_to_plot_time.append(beta_ell1)\n",
    "#         print(beta.shape)\n",
    "#         print(np.fft.fft2(beta, norm='ortho').shape)\n",
    "beta_to_plot_freq.append(np.abs(np.fft.fft2(beta_ell1, norm='ortho')))\n",
    "\n",
    "# Compute beta for linear CNNs\n",
    "sns.set_context(\"paper\",font_scale=3, rc={\"lines.linewidth\":1})\n",
    "fig, axs = plt.subplots(2, 2, figsize=(10,10))\n",
    "for k in range(2):\n",
    "    axs[0][k].imshow(beta_to_plot_time[k],cmap='gray')\n",
    "    axs[0][k].set_xticklabels([])\n",
    "    axs[0][k].set_yticklabels([])\n",
    "\n",
    "    axs[1][k].imshow(beta_to_plot_freq[k],cmap='gray', norm=LogNorm(vmin=0.0001, vmax=0.08))\n",
    "    axs[1][k].set_xticklabels([])\n",
    "    axs[1][k].set_yticklabels([])\n",
    "\n",
    "plt.tight_layout() # for some reason not using tight_layout before savefig messes some borders up\n",
    "name = \"cvxpy\" + str(Cout[cout]) + \".pdf\"\n",
    "plt.savefig(name, format='pdf')\n",
    "\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
