{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Projection onto intersection of sets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We try to find the projection of a vector into the intersection of simple sets (whose projection can be computed easily) - See Appendix A of *A Convex Approach to Minimal Partitions Antonin Chambolle, Daniel Cremers, Thomas Pock*\n",
    "\n",
    "\n",
    "$$\n",
    "proj_K(x) = \\bigcap_{1 \\leq i_1 < i_2 \\leq k} K_{i_1,i_2} \\quad  K_{i_1,i_2}= \\{ x: |x_{i_2} - x_{i_1}| \\leq \\sigma_{i1, i2} \\quad \\forall i1<i2 \\}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from pyproximal.projection import *\n",
    "from pyproximal.proximal import *\n",
    "\n",
    "np.random.seed(42)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create vector x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[22 22  8]\n",
      " [22 24  8]\n",
      " [ 8  8  3]]\n",
      "0 1 22 31.74892270912087\n",
      "0 2 8 7.548719254472989\n",
      "1 2 8 39.297641963593854\n"
     ]
    }
   ],
   "source": [
    "k = 3\n",
    "xtrue = np.random.normal(0, 50, k)\n",
    "x = xtrue.copy()\n",
    "x1 = xtrue.copy()\n",
    "sigma = np.array([[3,2,1], [2,2,1], [3,4,1]])\n",
    "sigma = sigma.T @ sigma\n",
    "print(sigma)\n",
    "\n",
    "for i1 in range(k-1):\n",
    "    for i2 in range(i1+1, k):\n",
    "        print(i1, i2, sigma[i1, i2], np.linalg.norm(x[i1] - x[i2]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Projection of vector x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "25\n"
     ]
    }
   ],
   "source": [
    "niter = 40\n",
    "tol = 1e-20\n",
    "\n",
    "x12 = np.zeros((k,k))\n",
    "for iiter in range(niter):\n",
    "    xold = x.copy()\n",
    "    for i1 in range(k-1):\n",
    "        for i2 in range(i1+1, k):\n",
    "            xtilde = x[i2] - x[i1] + x12[i1, i2]\n",
    "            xtildeabs = np.abs(xtilde)\n",
    "            xdtilde = np.maximum(0, xtildeabs - sigma[i1, i2]) * xtilde / xtildeabs\n",
    "            x[i1] = x[i1] + 0.5 * (xdtilde - x12[i1, i2])\n",
    "            x[i2] = x[i2] - 0.5 * (xdtilde - x12[i1, i2])\n",
    "            x12[i1, i2] = xdtilde\n",
    "    if max(np.abs(x - xold)) < tol:\n",
    "        break\n",
    "print(iiter)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check projected vector satisfy the condition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 1 22 31.74892270912087\n",
      "0 2 8 7.548719254472989\n",
      "1 2 8 39.297641963593854\n",
      "0 1 22 16.000000000000004\n",
      "0 2 8 8.0\n",
      "1 2 8 8.000000000000004\n"
     ]
    }
   ],
   "source": [
    "for i1 in range(k-1):\n",
    "    for i2 in range(i1+1, k):\n",
    "        print(i1, i2, sigma[i1, i2], np.abs(xtrue[i1] - xtrue[i2]))\n",
    "        \n",
    "for i1 in range(k-1):\n",
    "    for i2 in range(i1+1, k):\n",
    "        print(i1, i2, sigma[i1, i2], np.abs(x[i1] - x[i2]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ -0.06673448,  15.68218822, -15.61545374])"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ic = IntersectionProj(k, 1, sigma, niter, tol)\n",
    "x1 = ic(x1)\n",
    "x1 - xtrue"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "False\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "ic = Intersection(k, 1, sigma, niter, tol)\n",
    "print(ic(xtrue))\n",
    "x = ic.prox(xtrue, 1)\n",
    "print(ic(x))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Repeat the same, now with a matrix with n columns (algorithm works on each column indipendently)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[22 22  8]\n",
      " [22 24  8]\n",
      " [ 8  8  3]]\n",
      "0 1 22 [ 99.62521212  38.83567092  11.46403679 102.24712845  26.27362288]\n",
      "0 2 8 [171.81550505  74.53822289  16.40752861 129.60219679  22.65936983]\n",
      "1 2 8 [ 72.19029294 113.3738938    4.94349182  27.35506834   3.61425305]\n"
     ]
    }
   ],
   "source": [
    "k = 3\n",
    "n = 5\n",
    "xtrue = np.random.normal(0, 50, (k, n))\n",
    "x = xtrue.copy()\n",
    "x1 = xtrue.copy()\n",
    "sigma = np.array([[3,2,1], [2,2,1], [3,4,1]])\n",
    "sigma = sigma.T @ sigma\n",
    "print(sigma)\n",
    "\n",
    "for i1 in range(k-1):\n",
    "    for i2 in range(i1+1, k):\n",
    "        print(i1, i2, sigma[i1, i2], np.abs(x[i1] - x[i2]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "36\n"
     ]
    }
   ],
   "source": [
    "niter = 50\n",
    "tol = 1e-20\n",
    "\n",
    "x12 = np.zeros((k,k,n))\n",
    "for iiter in range(niter):\n",
    "    xold = x.copy()\n",
    "    for i1 in range(k-1):\n",
    "        for i2 in range(i1+1, k):\n",
    "            xtilde = x[i2] - x[i1] + x12[i1, i2]\n",
    "            xtildeabs = np.abs(xtilde)\n",
    "            xdtilde = np.maximum(0, xtildeabs - sigma[i1, i2]) * xtilde / xtildeabs\n",
    "            x[i1] = x[i1] + 0.5 * (xdtilde - x12[i1, i2])\n",
    "            x[i2] = x[i2] - 0.5 * (xdtilde - x12[i1, i2])\n",
    "            x12[i1, i2] = xdtilde\n",
    "    if max(np.sum(np.abs(x-xold), axis=0)) < tol:\n",
    "        break\n",
    "print(iiter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 1 22 [1.60000000e+01 3.55271368e-15 7.26027249e+00 1.60000000e+01\n",
      " 1.60000000e+01]\n",
      "0 2 8 [8. 8. 8. 8. 8.]\n",
      "1 2 8 [8.         8.         0.73972751 8.         8.        ]\n"
     ]
    }
   ],
   "source": [
    "for i1 in range(k-1):\n",
    "    for i2 in range(i1+1, k):\n",
    "        print(i1, i2, sigma[i1, i2], np.abs(x[i1] - x[i2]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Same using the projection operator in PyProximal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 1 22 [1.60000000e+01 2.25490737e-11 7.26027249e+00 1.60000000e+01\n",
      " 1.60000000e+01]\n",
      "0 2 8 [8. 8. 8. 8. 8.]\n",
      "1 2 8 [8.         8.         0.73972751 8.         8.        ]\n"
     ]
    }
   ],
   "source": [
    "ic = IntersectionProj(k, n, sigma, niter, tol)\n",
    "x1 = ic(x1)\n",
    "x1 = x1.reshape(k,n)\n",
    "\n",
    "for i1 in range(k-1):\n",
    "    for i2 in range(i1+1, k):\n",
    "        print(i1, i2, sigma[i1, i2], np.abs(x1[i1] - x1[i2]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and the proximal operator in PyProximal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "False\n",
      "True\n",
      "0 1 22 [1.60000000e+01 2.25490737e-11 7.26027249e+00 1.60000000e+01\n",
      " 1.60000000e+01]\n",
      "0 2 8 [8. 8. 8. 8. 8.]\n",
      "1 2 8 [8.         8.         0.73972751 8.         8.        ]\n"
     ]
    }
   ],
   "source": [
    "ic = Intersection(k, n, sigma, niter, tol)\n",
    "print(ic(xtrue, 1e-3))\n",
    "x = ic.prox(xtrue, 1)\n",
    "print(ic(x, 1e-3))\n",
    "x = x.reshape(k,n)\n",
    "\n",
    "for i1 in range(k-1):\n",
    "    for i2 in range(i1+1, k):\n",
    "        print(i1, i2, sigma[i1, i2], np.abs(x[i1] - x[i2]))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pylops",
   "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.13.5"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}