Matrix_ONN_converters.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "673de73c-3a58-4b75-bca2-9917309967d4",
   "metadata": {},
   "source": [
    "# From NN to ONN, different strategies comparison\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ab2b6bd-2c74-4312-8e55-8d1ba9f5a4d4",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "\n",
    "Given Y=WX a standard trained NN layer, with Y the prediction, X input data, and W ALREADY TRAINED weights.\n",
    "\n",
    "Convert the weights W into a sequence of 2x2 rotation matrices (describing MZI) and scaling vector (photonic attenuator).\n",
    "\n",
    "**Strategy 1 TT-decompotion**\n",
    "\n",
    "[1] https://aip.scitation.org/doi/10.1063/5.0070913# Inspiration for designing the strategy 1\n",
    "\n",
    "[2] https://github.com/Bihaqo/t3f/ T3F framework for TensorTrain arithmetics based on Tensorflow. I implement with numpy.\n",
    "\n",
    "**Strategy 2 simple 2x2 decomposition**\n",
    "\n",
    "**Strategy 3 Teacher-Student**\n",
    "\n",
    "[3] https://arxiv.org/pdf/1503.02531.pdf Knowledge Distilling\n",
    "\n",
    "**Strategy 4 Teacher-Student with prior decompotion**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0ff7a498-bdb2-4ac7-9ed6-fde52b40cd74",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "X\n",
      "[[-0.95956321]\n",
      " [ 0.66523969]\n",
      " [ 0.5563135 ]\n",
      " [ 0.7400243 ]]\n",
      "weights\n",
      "[[ 0.09762701  0.43037873  0.20552675  0.08976637]\n",
      " [-0.1526904   0.29178823 -0.12482558  0.783546  ]\n",
      " [ 0.92732552 -0.23311696  0.58345008  0.05778984]\n",
      " [ 0.13608912  0.85119328 -0.85792788 -0.8257414 ]]\n",
      "Y\n",
      "[[ 0.37339233]\n",
      " [ 0.85102612]\n",
      " [-0.67755906]\n",
      " [-0.65268413]]\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "np.random.seed(0)\n",
    "weights = np.random.uniform(-1., +1., (4, 4))\n",
    "X=np.random.uniform(-1., +1., (4, 1))\n",
    "\n",
    "Y=np.dot(weights, X)\n",
    "print(\"X\")\n",
    "print(X)\n",
    "print(\"weights\")\n",
    "print(weights)\n",
    "print(\"Y\")\n",
    "print(Y)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "145470fb-6548-4deb-932a-2453df6275f5",
   "metadata": {},
   "source": [
    "# Strategy 1: rank-2 TensorTrain decomposition\n",
    "\n",
    "Illustration of the strategy:\n",
    "\n",
    "![alt text](SVD.png \"svd decomposition\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3129319-7748-468e-bb0e-f56076bc00f5",
   "metadata": {},
   "source": [
    "### Util functions:\n",
    "* from_arr_to_tt : Converts np.ndarray into TensorTrain\n",
    "* from_tt_to_arr : Converts TensorTrain format into np.ndarray\n",
    "* tt_dot : Dot product between TensorTrain tensors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9b40d221-bff9-4c69-8c28-6720ef17a12f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def ranks(tt_cores):\n",
    "    ranks = []\n",
    "    for i in range(len(tt_cores)):\n",
    "        s = tt_cores[i].shape[0]\n",
    "        ranks.append(s)\n",
    "    s = tt_cores[-1].shape[-1]\n",
    "    ranks.append(s)\n",
    "    return np.stack(ranks, axis=0)\n",
    "\n",
    "\n",
    "class TensorTrain:\n",
    "    def __init__(self, tt_cores, tt_shapes, tt_ranks):\n",
    "        self.tt_cores = tt_cores\n",
    "        self.tt_shapes = tt_shapes\n",
    "        self.tt_ranks = tt_ranks\n",
    "\n",
    "\n",
    "def from_tt_to_arr(tt:TensorTrain, original_shape:tuple) -> np.ndarray:\n",
    "    \"\"\"Converts a TensorTrain into a regular tensor or matrix.\"\"\"\n",
    "    tt_ranks = tt.tt_ranks\n",
    "    res = tt.tt_cores[0]\n",
    "    for i in range(1, len(tt.tt_cores)):\n",
    "        res = np.reshape(res, (-1, tt_ranks[i]))\n",
    "        curr_core = np.reshape(tt.tt_cores[i], (tt_ranks[i], -1))\n",
    "        res = np.matmul(res, curr_core)\n",
    "    return np.reshape(res, original_shape)\n",
    "\n",
    "\n",
    "def _from_nd_arr_to_tt(arr:np.ndarray, max_tt_rank:int=10) -> TensorTrain:\n",
    "    \"\"\"Converts a given Numpy array to a TT-tensor of the same shape.\"\"\"\n",
    "    static_shape = list(arr.shape)\n",
    "    dynamic_shape = arr.shape\n",
    "    d = static_shape.__len__()\n",
    "    max_tt_rank = np.array(max_tt_rank).astype(np.int32)\n",
    "    if max_tt_rank.size == 1:\n",
    "        max_tt_rank = (max_tt_rank * np.ones(d + 1)).astype(np.int32)\n",
    "    ranks = [1] * (d + 1)\n",
    "    tt_cores = []\n",
    "    are_tt_ranks_defined = True\n",
    "    for core_idx in range(d - 1):\n",
    "        curr_mode = static_shape[core_idx]\n",
    "        if curr_mode is None:\n",
    "            curr_mode = dynamic_shape[core_idx]\n",
    "        rows = ranks[core_idx] * curr_mode\n",
    "        arr = np.reshape(arr, [rows, -1])\n",
    "        columns = arr.shape[1]\n",
    "        if columns is None:\n",
    "            columns = np.shape(arr)[1]\n",
    "\n",
    "        u, s, vT = np.linalg.svd(arr, full_matrices=False)\n",
    "        v = vT.T.T.T  # anti-transpose\n",
    "\n",
    "        # arr == u @ diag(s) @ vT\n",
    "        if max_tt_rank[core_idx + 1] == 1:\n",
    "            ranks[core_idx + 1] = 1\n",
    "        else:\n",
    "            ranks[core_idx + 1] = min(max_tt_rank[core_idx + 1], rows, columns)\n",
    "        u = u[:, 0:ranks[core_idx + 1]]\n",
    "        s = s[0:ranks[core_idx + 1]]\n",
    "        v = v[:, 0:ranks[core_idx + 1]]\n",
    "        core_shape = (ranks[core_idx], curr_mode, ranks[core_idx + 1])\n",
    "        tt_cores.append(np.reshape(u, core_shape))\n",
    "        arr = np.matmul(np.diag(s), np.transpose(v))\n",
    "    last_mode = static_shape[-1]\n",
    "    if last_mode is None:\n",
    "        last_mode = dynamic_shape[-1]\n",
    "    core_shape = (ranks[d - 1], last_mode, ranks[d])\n",
    "    tt_cores.append(np.reshape(arr, core_shape))\n",
    "    if not are_tt_ranks_defined:\n",
    "        ranks = None\n",
    "    return TensorTrain(tt_cores, static_shape, ranks)\n",
    "\n",
    "\n",
    "def from_arr_to_tt(mat:np.ndarray, shape:tuple, max_tt_rank:int=10) -> TensorTrain:\n",
    "    \"\"\"Converts a given matrix or vector to a TT-matrix.\"\"\"\n",
    "\n",
    "    # transpose\n",
    "    shape = np.array(shape)\n",
    "    tens = np.reshape(mat, shape.flatten())  # Warning there:\n",
    "    d = len(shape[0])\n",
    "    transpose_idx = np.arange(2 * d).reshape(2, d).T.flatten()\n",
    "    transpose_idx = list(transpose_idx.astype(int))\n",
    "    while len(transpose_idx) < len(tens.shape):\n",
    "        transpose_idx.append(len(transpose_idx))\n",
    "    tens = np.transpose(tens, transpose_idx)\n",
    "\n",
    "    new_shape = np.prod(shape, axis=0)\n",
    "    tens = np.reshape(tens, new_shape)\n",
    "    tt_tens = _from_nd_arr_to_tt(tens, max_tt_rank)\n",
    "\n",
    "    tt_cores = []\n",
    "    static_tt_ranks = list(tt_tens.tt_ranks)\n",
    "    dynamic_tt_ranks = ranks(tt_tens.tt_cores)\n",
    "    for core_idx in range(d):\n",
    "        curr_core = tt_tens.tt_cores[core_idx]\n",
    "        curr_rank = static_tt_ranks[core_idx]\n",
    "        if curr_rank is None:\n",
    "            curr_rank = dynamic_tt_ranks[core_idx]\n",
    "        next_rank = static_tt_ranks[core_idx + 1]\n",
    "        if next_rank is None:\n",
    "            next_rank = dynamic_tt_ranks[core_idx + 1]\n",
    "        curr_core_new_shape = [curr_rank, shape[0, core_idx], shape[1, core_idx], next_rank]\n",
    "\n",
    "        # patch:\n",
    "        # if max_tt_rank==2:\n",
    "        # while np.prod(curr_core_new_shape) < np.prod(curr_core.shape):\n",
    "        #  curr_core_new_shape.insert(1, 2)\n",
    "        try:\n",
    "            curr_core = np.reshape(curr_core, curr_core_new_shape)\n",
    "        except:\n",
    "            print(\"Error\")\n",
    "\n",
    "        tt_cores.append(curr_core)\n",
    "    return TensorTrain(tt_cores, shape, tt_tens.tt_ranks)\n",
    "\n",
    "\n",
    "def tt_dot(a: TensorTrain, b: TensorTrain) -> TensorTrain:\n",
    "    \"\"\"Multiplies two TT-matrices and returns the TT-matrix of the result.\"\"\"\n",
    "    ndims = len(a.tt_cores)\n",
    "    einsum_str = 'aijb,cjkd->acikbd'\n",
    "    result_cores = []\n",
    "    for core_idx in range(ndims):\n",
    "        a_core = a.tt_cores[core_idx]\n",
    "        b_core = b.tt_cores[core_idx]\n",
    "\n",
    "        try:\n",
    "            curr_res_core = np.einsum(einsum_str, a_core, b_core) #<------------ 2x2 multiplication\n",
    "        except ValueError:\n",
    "            print(\"Einstein Sum error\")\n",
    "\n",
    "        res_left_rank = a.tt_ranks[core_idx] * b.tt_ranks[core_idx]\n",
    "        res_right_rank = a.tt_ranks[core_idx + 1] * b.tt_ranks[core_idx + 1]\n",
    "        left_mode = a.tt_shapes[0][core_idx]\n",
    "        right_mode = b.tt_shapes[1][core_idx]\n",
    "\n",
    "        core_shape = [res_left_rank, left_mode, right_mode, res_right_rank]\n",
    "        # while np.prod(core_shape) < np.prod(curr_res_core.shape):\n",
    "        #  core_shape.insert(1, 2)\n",
    "        curr_res_core = np.reshape(curr_res_core, core_shape)\n",
    "\n",
    "        result_cores.append(curr_res_core)\n",
    "    res_shape = (a.tt_shapes[0], b.tt_shapes[1])\n",
    "    out_ranks = [a_r * b_r for a_r, b_r in zip(a.tt_ranks, b.tt_ranks)]\n",
    "    return TensorTrain(result_cores, res_shape, out_ranks)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c65cde2a-c076-46c4-88f3-ba6af43c5df8",
   "metadata": {},
   "source": [
    "### Conversion of weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7b9c8c7c-59e1-4965-aac1-32d55c1205e3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Weigths reconstruction after SVD MSE: 1.8262931624205774e-31\n",
      "Weigths reconstruction after 2x2 TT decomp. MSE: 0.13211937943719954\n"
     ]
    }
   ],
   "source": [
    "# SVD decomposition\n",
    "u, s, vT = np.linalg.svd(weights, full_matrices=False)\n",
    "V = vT.T.T.T # anti-transpose\n",
    "s_diag=np.diag(s)\n",
    "reconstructed_weights=np.dot(np.dot(u , s_diag) , vT)\n",
    "print(f\"Weigths reconstruction after SVD MSE: {np.mean((weights-reconstructed_weights)**2)}\")\n",
    "\n",
    "# from array to rank2 tt format\n",
    "X_tt=from_arr_to_tt(X, ((2, 2), (1, 1)), max_tt_rank=2)\n",
    "vt_tt=from_arr_to_tt(vT, ((2, 2), (2, 2)), max_tt_rank=2)\n",
    "s_tt=from_arr_to_tt(s_diag, ((2, 2), (2, 2)), max_tt_rank=2)\n",
    "u_tt=from_arr_to_tt(u, ((2, 2), (2, 2)), max_tt_rank=2)\n",
    "\n",
    "# compute TT weights\n",
    "w_tt = tt_dot(tt_dot(u_tt, s_tt), vt_tt)\n",
    "#w_tt=from_arr_to_tt(weights, ((2, 2), (2, 2)), max_tt_rank=2)\n",
    "\n",
    "# Only for checking\n",
    "weights_reconstructed=from_tt_to_arr(w_tt, weights.shape)\n",
    "print(f\"Weigths reconstruction after 2x2 TT decomp. MSE: {np.mean((weights-weights_reconstructed)**2)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a24a3d69-85d7-4c33-b8af-f9f4ce848e4c",
   "metadata": {},
   "source": [
    "### Prediction with TT cores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e95266dd-4752-4994-9444-4a28d6b804a1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Expected Y: [[ 0.37339233]\n",
      " [ 0.85102612]\n",
      " [-0.67755906]\n",
      " [-0.65268413]]\n",
      "Reconstructed Y: [[ 0.47324387]\n",
      " [ 0.67756381]\n",
      " [-0.49314283]\n",
      " [-1.02924294]]\n",
      "Prediction MSE: 0.0539663482169657\n"
     ]
    }
   ],
   "source": [
    "Y_tt=tt_dot(w_tt ,X_tt)\n",
    "Y_reconstructed=from_tt_to_arr(Y_tt, Y.shape)\n",
    "\n",
    "print(f\"Expected Y:\", Y)\n",
    "print(f\"Reconstructed Y:\", Y_reconstructed)\n",
    "print(f\"Prediction MSE: {np.mean((Y-Y_reconstructed)**2)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7015a123-8ccb-44bd-9dee-434d1fdc7309",
   "metadata": {},
   "source": [
    "TODO: convert unitary TT-cores into 2x2 rotation TT-cores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "9adb0c72-e53d-4746-bcba-6aaff617f1c4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[array([[[[ 0.04914496,  0.04631135,  0.30888377, -0.21194745,\n",
       "           -0.13345392,  0.14266277,  0.48326384, -0.37430949],\n",
       "          [ 0.13339556,  0.25923827, -0.16043367, -0.12646277,\n",
       "            0.1707655 ,  0.23297152, -0.33580986, -0.34692864]],\n",
       " \n",
       "         [[-0.5378256 ,  0.16894952, -0.05201302, -0.07182773,\n",
       "            0.29717143, -0.06882413,  0.14954386, -0.04710726],\n",
       "          [-0.11797425, -0.478231  , -0.18648136, -0.35399926,\n",
       "            0.11389008,  0.34985751,  0.03254413,  0.13251827]]]]),\n",
       " array([[[[-3.75884716e-01],\n",
       "          [ 1.91689594e-01]],\n",
       " \n",
       "         [[-7.98015332e-01],\n",
       "          [-4.11100709e-01]]],\n",
       " \n",
       " \n",
       "        [[[-7.02982874e-02],\n",
       "          [ 1.18352830e-01]],\n",
       " \n",
       "         [[ 1.18956505e+00],\n",
       "          [ 2.37145632e+00]]],\n",
       " \n",
       " \n",
       "        [[[-4.58305819e-02],\n",
       "          [ 2.33721810e-02]],\n",
       " \n",
       "         [[-9.72997982e-02],\n",
       "          [ 9.10835572e-02]]],\n",
       " \n",
       " \n",
       "        [[[-8.57127540e-03],\n",
       "          [ 1.44304326e-02]],\n",
       " \n",
       "         [[-8.60548677e-02],\n",
       "          [-7.68257853e-02]]],\n",
       " \n",
       " \n",
       "        [[[ 1.40253271e+00],\n",
       "          [-1.20519290e+00]],\n",
       " \n",
       "         [[-3.27418217e-01],\n",
       "          [ 6.21050767e-02]]],\n",
       " \n",
       " \n",
       "        [[[ 1.06412650e+00],\n",
       "          [ 8.28189194e-01]],\n",
       " \n",
       "         [[ 1.10388891e-01],\n",
       "          [ 3.74880801e-01]]],\n",
       " \n",
       " \n",
       "        [[[ 1.71006927e-01],\n",
       "          [-6.23753522e-02]],\n",
       " \n",
       "         [[-3.99211960e-02],\n",
       "          [ 2.56738290e-02]]],\n",
       " \n",
       " \n",
       "        [[[-8.65866139e-03],\n",
       "          [-1.18203871e-01]],\n",
       " \n",
       "         [[-1.61648258e-02],\n",
       "          [-1.20585018e-03]]]])]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "w_tt.tt_cores"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05599699-8e05-4192-af75-4f2dfd62f29c",
   "metadata": {},
   "source": [
    "## Strategy 2: Tiled 2x2 matrices decomposition"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3895b40-4377-4108-b5e1-80c317fdbc39",
   "metadata": {},
   "source": [
    "Steps:\n",
    "* Split NxN matrix into 2x2 tiles\n",
    "* Using eiven values/vectors decomposition to compute 2x2 rotation matrices 2d scaling vectors\n",
    "* Using those matrices/vectors, to compute phase shift in MZIs/attenuators based on arccos/arcsin"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23040adb-4b91-4c95-9e15-cbf985595126",
   "metadata": {},
   "source": [
    "## Utils\n",
    "* from_arr_to_tiles\n",
    "* from_tiles_to_arr\n",
    "* dot_tile_vec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "id": "b30a2e84-8bd0-46ee-ba91-640416dc3a30",
   "metadata": {},
   "outputs": [],
   "source": [
    "def from_arr_to_tiles(matrix):\n",
    "    n = matrix.shape[0]\n",
    "    if n % 2 != 0:\n",
    "        raise ValueError(\"Matrix must be of even size\")\n",
    "\n",
    "    block_matrices = []\n",
    "    for i in range(0, n, 2):\n",
    "        for j in range(0, n, 2):\n",
    "            block = matrix[i:i+2, j:j+2]\n",
    "            block_matrices.append(block)\n",
    "    return block_matrices\n",
    "\n",
    "def from_tiles_to_arr(block_matrices):\n",
    "    n = len(block_matrices)\n",
    "    size = int(np.sqrt(n))\n",
    "    if size ** 2 != n:\n",
    "        raise ValueError(\"Number of block matrices must be a perfect square\")\n",
    "\n",
    "    C = np.zeros((size*2, size*2))\n",
    "    for i in range(size):\n",
    "        for j in range(size):\n",
    "            C[i*2:i*2+2, j*2:j*2+2] = block_matrices[i*size+j]\n",
    "    return C\n",
    "\n",
    "def dot_tile_vec(block_matrices, x):\n",
    "    \"\"\"tiled matrix and normal vector multiplication \"\"\"\n",
    "    n = len(block_matrices)\n",
    "    size = int(np.sqrt(n))\n",
    "    if size ** 2 != n:\n",
    "        raise ValueError(\"Number of block matrices must be a perfect square\")\n",
    "    if x.shape[0] != size*2:\n",
    "        raise ValueError(\"Input vector must have the same size as the matrix\")\n",
    "    y = x.copy()\n",
    "    result = np.zeros(y.shape)\n",
    "    for i in range(size):\n",
    "        for j in range(size):\n",
    "            b1=block_matrices[i*size+j]\n",
    "            b2=y[j*2:j*2+2]\n",
    "            result[i*2:i*2+2] += np.dot(b1, b2) #<----- 2x2 dot between matrix and vector\n",
    "    return result\n",
    "\n",
    "def tiled_prediction(tiled_mat_a, tiled_vec_b, vec_x):\n",
    "    \"\"\"tiled computation between: (a.x)+b. \n",
    "    With a tiled matrix format and b a tiled vector format. x is a standard vector. \"\"\"\n",
    "    n = len(tiled_mat_a)\n",
    "    size = int(np.sqrt(n))\n",
    "    y = vec_x.copy()\n",
    "    result = np.zeros(y.shape)\n",
    "    for i in range(size):\n",
    "        for j in range(size):\n",
    "            b1=tiled_mat_a[i * size + j] # <---- 2x2 matrix\n",
    "            b2=y[j*2:j*2+2] # <------ 1x2 numpy array\n",
    "            b3=np.array([tiled_vec_b[i * size + j]]) # <------- 1x2 numpy array\n",
    "            result[i*2:i*2+2] += np.dot(b1, b2) + np.diag(b3) #<----- The core computing is here\n",
    "    return result\n",
    "\n",
    "\n",
    "def from_random_to_rotation(x):\n",
    "    eigenvalues, eigenvectors = np.linalg.eig(x)\n",
    "    \n",
    "    # rotation matrix are not unique. Different runs may produces different rotation. However, the dot product of them is unique.\n",
    "    rot0=[[eigenvectors[0][0], -eigenvectors[1][0]],\n",
    "          [eigenvectors[1][0], eigenvectors[0][0]]]\n",
    "    rot1=[[eigenvectors[0][1], -eigenvectors[1][1]],\n",
    "          [eigenvectors[1][1], eigenvectors[0][1]]]\n",
    "    rot=np.dot(rot0,rot1)\n",
    "    \n",
    "    make_real=lambda x: x.real\n",
    "    rot[0][0]=make_real(rot[0][0])\n",
    "    rot[1][0]=make_real(rot[1][0])\n",
    "    rot[0][1]=make_real(rot[0][1])\n",
    "    rot[1][1]=make_real(rot[1][1])\n",
    "    \n",
    "    eigenvalues[0]=make_real(eigenvalues[0])\n",
    "    eigenvalues[1]=make_real(eigenvalues[1])\n",
    "    \n",
    "    return rot.astype(np.float32), eigenvalues.astype(np.float32)\n",
    "\n",
    "def check_and_validate_2x2_rot_mat(x):\n",
    "    epsilon=1e-2\n",
    "    if x.shape != (2, 2):\n",
    "        raise ValueError(f\"Matrix {x} is not 2x2\")\n",
    "    det = np.linalg.det(x)\n",
    "    if abs(det - 1) > epsilon: #The determinant of a rotation matrix is always 1, because the matrix preserves the volume of the space it acts on. \n",
    "        raise ValueError(f\"Matrix {x} is not a rotation matrix. Determinant: {det}\")\n",
    "    if np.allclose(np.dot(x, x.T),np.eye(2),atol=epsilon): #\n",
    "        raise ValueError(f\"Matrix {x} is not orthogonal. x.xT = {I}\")\n",
    "\n",
    "def rotation_mzi_angle(x):\n",
    "    #check_and_validate_2x2_rot_mat(x)\n",
    "\n",
    "    clipped_x=np.clip(x, -1, 1)\n",
    "    angle = np.arccos(clipped_x[0][0])\n",
    "    # revert the angle if needed\n",
    "    if clipped_x[1,0] < 0:\n",
    "        angle = 2*np.pi - angle \n",
    "    return angle\n",
    "\n",
    "def rotation_2attenuators_angle(x):\n",
    "    clipped_x=np.clip(x, -1, 1)\n",
    "    return np.arccos(clipped_x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79e54907-127f-4c77-bc50-c8d87433da35",
   "metadata": {},
   "source": [
    "### Converts NxN random weights -> 2x2 rotation matrices and 2d scaling vectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "c3a69895-6371-4800-8ff1-46c896a3af2c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_28226/3285656845.py:77: ComplexWarning: Casting complex values to real discards the imaginary part\n",
      "  return rot.astype(np.float32), eigenvalues.astype(np.float32)\n"
     ]
    }
   ],
   "source": [
    "# Tiling weights 4x4 into 4 2x2 matrices\n",
    "w_t=from_arr_to_tiles(weights)\n",
    "\n",
    "# From 2x2 matrix to 2x2 rotation matrix\n",
    "w_rot_t=[]\n",
    "w_val_t=[]\n",
    "for w_ti in w_t:\n",
    "    rot_matrix, eigen_values=from_random_to_rotation(w_ti)\n",
    "    w_rot_t.append(rot_matrix)\n",
    "    w_val_t.append(eigen_values)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6363cde-c408-4479-9206-b1eb80bfc614",
   "metadata": {},
   "source": [
    "### Prediction with the 2x2 rotation matrices and 2D scaling vec."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "db2a8ce8-1a4f-41fa-aa87-ccfe58c595c0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Expected Y: [[ 0.37339233]\n",
      " [ 0.85102612]\n",
      " [-0.67755906]\n",
      " [-0.65268413]]\n",
      "Reconstructed Y: [[-1.03054864]\n",
      " [ 0.92809649]\n",
      " [ 0.67841055]\n",
      " [ 2.26272296]]\n",
      "Prediction MSE: 0.0539663482169657\n"
     ]
    }
   ],
   "source": [
    "# Predicting with 2x2 rotations and 2d scaling vectors \n",
    "reconstructed_y=tiled_prediction(w_rot_t, w_val_t, X)\n",
    "print(f\"Expected Y:\", Y)\n",
    "print(f\"Reconstructed Y:\", reconstructed_y)\n",
    "print(f\"Prediction MSE: {np.mean((Y-Y_reconstructed)**2)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "513eacac-fde5-43cf-8476-7e6a38884aec",
   "metadata": {},
   "source": [
    "### MZI and attenuators angles of the phase shift (radian)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "id": "25732c70-754a-4190-a140-3b35a8ea0991",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MZI#0 rotation angle = 1.074408 , associated attenuators: [1.3748369 1.3748369]\n",
      "MZI#1 rotation angle = 1.6313761 , associated attenuators: [1.3432273 0.7021385]\n",
      "MZI#2 rotation angle = 1.3048719 , associated attenuators: [0.4750728 0.4750728]\n",
      "MZI#3 rotation angle = 1.0543903 , associated attenuators: [0.9916114 2.4810073]\n"
     ]
    }
   ],
   "source": [
    "# MZI angles\n",
    "for i in range(len(w_rot_t)):\n",
    "    mzi_theta=rotation_mzi_angle(w_rot_t[i])\n",
    "    atts_theta=rotation_2attenuators_angle(w_val_t[i])\n",
    "    print(f\"MZI#{i} rotation angle = {str(round(mzi_theta,7))} , associated attenuators: {np.round(atts_theta,7)}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1163547-d52e-4b5b-9054-8d301d2c9478",
   "metadata": {},
   "source": [
    "## Strategy 3: Teacher-Student approach \n",
    "\n",
    "The teacher is a standard ANN (NxN weights) and the student ONN (set of 2x2 rotation matrices). This method is a data driven procedure."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57c417ed-d400-4ed9-807c-46a9ce572ab4",
   "metadata": {},
   "source": [
    "### Generate random data to get the teacher predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 107,
   "id": "b9983406-b661-4eb7-ae8e-ea648973243b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def teacher(W, nb_data):\n",
    "    synt_data=np.random.uniform(-1., +1., (nb_data , 4, 1))\n",
    "    teacher_pred=[]\n",
    "    for x in synt_data:\n",
    "        pred_Y=np.dot(W,x)# <-- teacher model\n",
    "        teacher_pred.append( pred_Y )\n",
    "    \n",
    "    #reformating by squeezing\n",
    "    synt_data=synt_data.squeeze()\n",
    "    teacher_pred=np.array(teacher_pred).squeeze()\n",
    "    return synt_data, teacher_pred\n",
    "teacher_X, teacher_pred = teacher(weights, 1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bcf63b49-1761-48e9-9460-0ff844cd289d",
   "metadata": {},
   "source": [
    "### The student is trained to mimic the teacher"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "id": "3c966b57-47fe-4e6b-a134-386cf121c4eb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.10698833\n",
      "0.07487212\n",
      "0.07161769\n",
      "0.07159352\n",
      "0.07159182\n"
     ]
    }
   ],
   "source": [
    "from ONN import ONN # contains photonic neural network\n",
    "hp={\"lr\":1., \"lr_decay\":10., \"layers\":[4], \"epochs\": 5}\n",
    "student_onn=ONN(hp)\n",
    "student_onn.initialize()   \n",
    "student_onn.fit(teacher_X, teacher_pred, teacher_X, teacher_pred)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "410f0c69-1485-4199-aa3a-2fc7034cf901",
   "metadata": {},
   "source": [
    "### Prediction with the student model (ONN)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "id": "64e3b86b-cbc6-487b-9e4f-7be732a7706e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Y expected: [ 0.37339233  0.85102612 -0.67755906 -0.65268413]\n",
      "Y reconstructed: [ 0.93145835  0.62140274 -0.75870734 -0.62532663]\n",
      "Prediction MSE: 0.09287450462579727\n"
     ]
    }
   ],
   "source": [
    "Y_reconstructed=student_onn.predict(np.array([X.squeeze()])).squeeze()\n",
    "Y_expected=Y.squeeze()\n",
    "\n",
    "print(f\"Y expected: {Y_expected}\")\n",
    "print(f\"Y reconstructed: {Y_reconstructed}\")\n",
    "print(f\"Prediction MSE: {np.mean((Y_expected-Y_reconstructed)**2)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38b432b4-8245-4387-bc8d-e3f59603fbd7",
   "metadata": {},
   "source": [
    "### Theta of phase shifters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "id": "77b6ec30-7a9d-4056-b853-62f79fdc8735",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Array([-3.1486027,  0.282853 ], dtype=float32), Array([1.6477056], dtype=float32), Array([ 1.4729869, -2.0515711], dtype=float32), Array([-0.7082593], dtype=float32)]\n"
     ]
    }
   ],
   "source": [
    "print(model.W)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff76b168-616b-4c04-a907-b261f9b3038b",
   "metadata": {},
   "source": [
    "NB: The teacher is 4x4 matrix, and the student contains 4 columns of respectively 2,1,2,1 MZIs. The teacher contains 16 parameters and the student only 6 params."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eed5b45b-257b-4cae-b4d1-84cad05cc553",
   "metadata": {},
   "source": [
    "## Strategy 4: Teacher-Student with prior SVD decomposition\n",
    "\n",
    "After SVD decomposition, we have two teachers and two associated students. The procedure is similar than strategy 3 but applied indpedantly on S and V."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aaabbe13-02ab-4209-a4b5-89a1bee3c7b2",
   "metadata": {},
   "source": [
    "teacher_u, teacher_v produce preds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "id": "88f33278-a3d8-4c31-b40a-94df2fc26631",
   "metadata": {},
   "outputs": [],
   "source": [
    "u, s, vT = np.linalg.svd(weights, full_matrices=False)\n",
    "s_diag=np.diag(s)\n",
    "teacher_u_X, teacher_u_pred = teacher(u, 1000)\n",
    "teacher_v_X, teacher_v_pred = teacher(vT, 1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29c75609-469d-4120-8021-8c85c306d385",
   "metadata": {},
   "source": [
    "student_u, student_v are trained to mimic their teacher"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "id": "8b27a1b9-2a07-4070-b24c-c54c4baff987",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Teacher_u -> Student_u training\n",
      "0.33702737\n",
      "0.33489755\n",
      "0.3336078\n",
      "0.33354986\n",
      "0.33354458\n",
      "Teacher_v -> Student_v training\n",
      "0.34730762\n",
      "0.34290767\n",
      "0.34273097\n",
      "0.34270787\n",
      "0.34270602\n"
     ]
    }
   ],
   "source": [
    "from ONN import ONN # contains photonic neural network\n",
    "hp={\"lr\":1., \"lr_decay\":10., \"layers\":[4], \"epochs\": 5}\n",
    "student_u=ONN(hp)\n",
    "student_u.initialize() \n",
    "print(\"Teacher_u -> Student_u training\")\n",
    "student_u.fit(teacher_u_X, teacher_u_pred, teacher_u_X, teacher_u_pred)\n",
    "\n",
    "hp={\"lr\":1., \"lr_decay\":10., \"layers\":[4], \"epochs\": 5}\n",
    "student_v=ONN(hp)\n",
    "student_v.initialize() \n",
    "print(\"Teacher_v -> Student_v training\")\n",
    "student_v.fit(teacher_v_X, teacher_v_pred, teacher_v_X, teacher_v_pred)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d8cd291-3ce1-4ebd-95e8-3f6991dc6a66",
   "metadata": {},
   "source": [
    "Prediction on new data with student_u, student_v and diag_s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "id": "030272d1-f7c5-46dc-8f0b-8d77b2d2e52a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Y expected: [ 0.37339233  0.85102612 -0.67755906 -0.65268413]\n",
      "Y reconstructed: [ 1.60870897  0.63157953 -0.01950631  0.29593791]\n",
      "Prediction MSE: 0.7267703008886807\n"
     ]
    }
   ],
   "source": [
    "# compute v.T*x\n",
    "vTx=student_u.predict(np.array([X.squeeze()])).squeeze()\n",
    "# compute u*v.T*x\n",
    "uvTx=student_v.predict(np.array([vTx.squeeze()])).squeeze()\n",
    "# compute diag_s*u*v.T*x\n",
    "Y_reconstructed=np.dot(s_diag , uvTx)\n",
    "\n",
    "Y_expected=Y.squeeze()\n",
    "print(f\"Y expected: {Y_expected}\")\n",
    "print(f\"Y reconstructed: {Y_reconstructed}\")\n",
    "print(f\"Prediction MSE: {np.mean((Y_expected-Y_reconstructed)**2)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5acc7269-e6a7-42d9-a3fe-8b0a782431e8",
   "metadata": {},
   "source": [
    "### Theta of phase shifters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "id": "62c58068-7f0a-4315-bb9c-774088221829",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "U phase shifters:\n",
      "[Array([-2.5794332, -5.7190742], dtype=float32), Array([5.4699335], dtype=float32), Array([-0.9637948,  0.1765687], dtype=float32), Array([4.623848], dtype=float32)]\n",
      "V phase shifters:\n",
      "[Array([-1.0001415, -3.054548 ], dtype=float32), Array([1.2980582], dtype=float32), Array([-4.163986, -8.60073 ], dtype=float32), Array([9.581584], dtype=float32)]\n",
      "S attenuators:\n",
      "[1.57967707 1.09579685 0.77645471 0.32089482]\n"
     ]
    }
   ],
   "source": [
    "print(\"U phase shifters:\")\n",
    "print(student_u.W)\n",
    "print(\"V phase shifters:\")\n",
    "print(student_v.W)\n",
    "print(\"S attenuators:\")\n",
    "print(s)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7655d68-027f-4a78-b7de-e6322b3531c6",
   "metadata": {},
   "source": [
    "Conclusion:  Different random seeds show that strat3 is more accurate than strat4. \"y_pred=approx(w).x\" seems more accurate than successive errors \"y_pred=s.approx(u).approx(v).x\" ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5e4ea36-23de-4fda-9714-40a5237c2777",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}