Matrix_Lemmas.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Matrix Lemmas",
      "provenance": [],
      "collapsed_sections": [
        "40U3GlYaKBv_"
      ],
      "authorship_tag": "ABX9TyPXpW7mwyb7labeSyHMolRi",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/TomMaullin/BLMM-sandbox/blob/master/Matrix_Lemmas.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CyTmrSzIJuep",
        "colab_type": "text"
      },
      "source": [
        "# Matrix reshape lemmas\n",
        "\n",
        "This notebook contains all matrix reshaping lemmas found during the BLMM project."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1T1lYptfC6Wf",
        "colab_type": "text"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HOZrPOy9JqgH",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import numpy as np\n",
        "import time\n",
        "import scipy.sparse"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eN5WVXEvtBm5",
        "colab_type": "text"
      },
      "source": [
        "### Helper functions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E7vaH4gltHCU",
        "colab_type": "text"
      },
      "source": [
        "This section contains helper functions used later in the notebook."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tG1v-2spZSJ2",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def mat2vec(matrix):\n",
        "  \n",
        "  #Return vectorised matrix\n",
        "  return(matrix.transpose().reshape(matrix.shape[0]*matrix.shape[1],1))\n",
        "\n",
        "def vec2mat(vec): # Only for square matrices\n",
        "  \n",
        "  # Return matrix\n",
        "  return(vec.reshape(np.int64(np.sqrt(vec.shape[0])),np.int64(np.sqrt(vec.shape[0]))).transpose())\n",
        "\n",
        "# Block vec with row-wise chunks of size n\n",
        "def matb2vecb(matb,n):\n",
        "  \n",
        "  m = matb.shape[0]\n",
        "  k = matb.shape[1]\n",
        "\n",
        "  vecb = matb.reshape(m//n, n, k).transpose((1, 0, 2)).reshape(n, m*k//n).transpose().reshape(m//n,n*k)\n",
        "\n",
        "  #Return vectorised matrix\n",
        "  return(vecb)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QeuHq8YFT_wQ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def comMat(a, b):\n",
        "\n",
        "  # Example matrix with unique elements\n",
        "  m = np.arange(a*b).reshape(a,b)\n",
        "\n",
        "  # Vec(m) and Vec(m')\n",
        "  vecm = mat2vec(m)\n",
        "  vecmt = mat2vec(m.transpose())\n",
        "\n",
        "  # Work out mapping between them.\n",
        "  k=(vecm.transpose()==vecmt).astype(int).transpose()\n",
        "\n",
        "  return(k)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "thBJGbQfC15c",
        "colab_type": "text"
      },
      "source": [
        "## Lemmas"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "40U3GlYaKBv_",
        "colab_type": "text"
      },
      "source": [
        "### **Lemma 1: Sum of $A_iB_i'$**\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3ouMxkVxKSjE",
        "colab_type": "text"
      },
      "source": [
        "Let $A$ and $B$ be matrices of size $(m,k)$, partitioned row-wise into chunks $A_1,...A_{m/n}$ and $B_1,...B_{l}$ of size $(n,k)$ where $m|n$ and $l=m/n$. I.e.\n",
        "\n",
        "$$A = \\begin{bmatrix} A_1 \\\\ A_2 \\\\ \\vdots \\\\ A_{l} \\end{bmatrix}, B = \\begin{bmatrix} B_1 \\\\ B_2 \\\\ \\vdots \\\\ B_{l} \\end{bmatrix}$$\n",
        "\n",
        "Then:\n",
        "\n",
        "$$\\sum_{i=1}^l A_iB_i' = ((A')^{r(kl,n)})'(B')^{r(kl,n)}$$"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pX9zRkOxJssh",
        "colab_type": "code",
        "outputId": "ff7c4c05-072f-4075-8d56-50f5cee8bf5d",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 369
        }
      },
      "source": [
        "m = 8\n",
        "n=2\n",
        "l=m//n\n",
        "k=4\n",
        "\n",
        "print('Example Shapes:')\n",
        "print('m: ', m)\n",
        "print('k: ', k)\n",
        "print('l: ', l)\n",
        "print('n: ', n)\n",
        "\n",
        "# A and B are random matrices\n",
        "A = np.random.randn(m,k)\n",
        "B = np.random.randn(m,k)\n",
        "\n",
        "t1 = time.time()\n",
        "\n",
        "# First we evaluate the LHS\n",
        "sumAB = np.zeros((n,n))\n",
        "for i in np.arange(m//n):\n",
        "\n",
        "  # Work out the row-wise chunks\n",
        "  Ai = A[n*i:n*(i+1),:]\n",
        "  Bi = B[n*i:n*(i+1),:]\n",
        "\n",
        "  # Perform the summation\n",
        "  sumAB = sumAB + Ai @ Bi.transpose()\n",
        "\n",
        "t2 = time.time()\n",
        "\n",
        "# We print the time\n",
        "time1=t2-t1\n",
        "print('Time (LHS):')\n",
        "print(time1)\n",
        "\n",
        "# And the result\n",
        "print('Result (LHS):')\n",
        "print(sumAB)\n",
        "\n",
        "# Now the RHS\n",
        "t1 = time.time()\n",
        "sumAB2 = A.transpose().reshape((m*k//n,n)).transpose() @ B.transpose().reshape((m*k//n,n))\n",
        "t2 = time.time()\n",
        "\n",
        "# We print the time\n",
        "time2=t2-t1\n",
        "print('Time (RHS):')\n",
        "print(time2)\n",
        "\n",
        "# And the result\n",
        "print('Result (RHS):')\n",
        "print(sumAB2)\n",
        "\n",
        "# And the difference\n",
        "print('Difference:')\n",
        "print(sumAB-sumAB2)\n",
        "\n",
        "\n",
        "# And the ratio between times\n",
        "print('time ratio:')\n",
        "print(time2/time1)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Example Shapes:\n",
            "m:  8\n",
            "k:  4\n",
            "l:  4\n",
            "n:  2\n",
            "Time (LHS):\n",
            "0.0005023479461669922\n",
            "Result (LHS):\n",
            "[[6.94992855 1.12747496]\n",
            " [9.38470056 3.32729359]]\n",
            "Time (RHS):\n",
            "9.799003601074219e-05\n",
            "Result (RHS):\n",
            "[[6.94992855 1.12747496]\n",
            " [9.38470056 3.32729359]]\n",
            "Difference:\n",
            "[[-8.88178420e-16  2.22044605e-16]\n",
            " [ 0.00000000e+00 -4.44089210e-16]]\n",
            "time ratio:\n",
            "0.1950640721404841\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "P9z6eOrfnM36",
        "colab_type": "code",
        "outputId": "dfa85aa2-c669-457d-c91f-725555eddb7f",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 263
        }
      },
      "source": [
        "print(A)\n",
        "\n",
        "print(A.transpose().reshape((m*k//n,n)).transpose())"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "[[-1.03758992 -0.9929396   2.25993744  1.45764701]\n",
            " [-0.55193597  1.42300798 -0.17094288  1.57698125]\n",
            " [ 2.22205484  0.21382045  0.44550954 -1.29899934]\n",
            " [ 1.40566434 -0.28190786  1.2002049  -2.1912631 ]\n",
            " [-0.66768399  1.16896045 -1.38153112 -0.7878413 ]\n",
            " [-1.63726715  1.76841371 -1.38325086 -1.63711714]\n",
            " [ 0.00706326  1.08959998 -0.74409399 -1.02894232]\n",
            " [-0.2313653  -0.90833872  1.7759244   1.07333532]]\n",
            "[[-1.03758992  2.22205484 -0.66768399  0.00706326 -0.9929396   0.21382045\n",
            "   1.16896045  1.08959998  2.25993744  0.44550954 -1.38153112 -0.74409399\n",
            "   1.45764701 -1.29899934 -0.7878413  -1.02894232]\n",
            " [-0.55193597  1.40566434 -1.63726715 -0.2313653   1.42300798 -0.28190786\n",
            "   1.76841371 -0.90833872 -0.17094288  1.2002049  -1.38325086  1.7759244\n",
            "   1.57698125 -2.1912631  -1.63711714  1.07333532]]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "foGGlFiwOTrz",
        "colab_type": "text"
      },
      "source": [
        "### **Lemma 2: Sum of $A_iCB_i'$**\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UkUjXFvCOcc7",
        "colab_type": "text"
      },
      "source": [
        "Let $A$ and $B$ be matrices of sizes $(m,k_1)$ and $(m,k_2)$ respectively, partitioned row-wise into chunks $A_1,...A_{m/n}$ and $B_1,...B_{l}$ of sizes $(n,k_1)$ and $(m,k_2)$ respectively where $m|n$ and $l=m/n$. I.e.\n",
        "\n",
        "$$A = \\begin{bmatrix} A_1 \\\\ A_2 \\\\ \\vdots \\\\ A_{l} \\end{bmatrix}, B = \\begin{bmatrix} B_1 \\\\ B_2 \\\\ \\vdots \\\\ B_{l} \\end{bmatrix}$$\n",
        "\n",
        "Let $C$ be a matrix of size $(k_1, k_2)$. Then:\n",
        "\n",
        "$$\\sum_{i=1}^l A_iCB_i' = ((A')^{r(k_1l,n)})'(CB')^{r(k_1l,n)} = ((C'A')^{r(k_2l,n)})'(B')^{r(k_2l,n)}$$"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Hmmbv1kBJt6U",
        "colab_type": "code",
        "outputId": "923e77d6-ca9a-4928-97be-fef63a86cc1b",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 562
        }
      },
      "source": [
        "m = 1000\n",
        "n=2\n",
        "l=m//n\n",
        "k1=100\n",
        "k2=3\n",
        "\n",
        "print('Example Shapes:')\n",
        "print('m: ', m)\n",
        "print('k1, k2: ', k1, k2)\n",
        "print('l: ', l)\n",
        "print('n: ', n)\n",
        "\n",
        "# A, B and C are random matrices\n",
        "A = np.random.randn(m,k1)\n",
        "B = np.random.randn(m,k2)\n",
        "C = np.random.randn(k1,k2)\n",
        "\n",
        "t1 = time.time()\n",
        "\n",
        "# First we evaluate the LHS\n",
        "sumACB = np.zeros((n,n))\n",
        "for i in np.arange(m//n):\n",
        "\n",
        "  # Work out the row-wise chunks\n",
        "  Ai = A[n*i:n*(i+1),:]\n",
        "  Bi = B[n*i:n*(i+1),:]\n",
        "\n",
        "  # Perform the summation\n",
        "  sumACB = sumACB + Ai @ C @ Bi.transpose()\n",
        "\n",
        "t2 = time.time()\n",
        "\n",
        "# We print the time\n",
        "time1=t2-t1\n",
        "print('Time (LHS):')\n",
        "print(time1)\n",
        "\n",
        "# And the result\n",
        "print('Result (LHS):')\n",
        "print(sumACB)\n",
        "\n",
        "# Now the middle\n",
        "t1 = time.time()\n",
        "sumACB2 = A.transpose().reshape((m*k1//n,n)).transpose() @ (C @ B.transpose()).reshape((m*k1//n,n))\n",
        "t2 = time.time()\n",
        "\n",
        "# We print the time\n",
        "time2=t2-t1\n",
        "print('Time (middle):')\n",
        "print(time2)\n",
        "\n",
        "# And the result\n",
        "print('Result (middle):')\n",
        "print(sumACB2)\n",
        "\n",
        "# Now the RHS\n",
        "t1 = time.time()\n",
        "sumACB3 = (C.transpose() @ A.transpose()).reshape((m*k2//n,n)).transpose() @ B.transpose().reshape((m*k2//n,n))\n",
        "t2 = time.time()\n",
        "\n",
        "# We print the time\n",
        "time3=t2-t1\n",
        "print('Time (RHS):')\n",
        "print(time3)\n",
        "\n",
        "# And the result\n",
        "print('Result (RHS):')\n",
        "print(sumACB3)\n",
        "\n",
        "# And the difference\n",
        "print('Differences:')\n",
        "print(sumACB-sumACB2)\n",
        "print(sumACB2-sumACB3)\n",
        "print(sumACB3-sumACB)\n",
        "\n",
        "\n",
        "# And the ratio between times\n",
        "print('time ratios:')\n",
        "print(time2/time1)\n",
        "print(time3/time2)\n",
        "print(time1/time3)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Example Shapes:\n",
            "m:  1000\n",
            "k1, k2:  100 3\n",
            "l:  500\n",
            "n:  2\n",
            "Time (LHS):\n",
            "0.004046916961669922\n",
            "Result (LHS):\n",
            "[[  75.30258318 -686.5630488 ]\n",
            " [  16.42764312 -148.29578741]]\n",
            "Time (middle):\n",
            "0.002064943313598633\n",
            "Result (middle):\n",
            "[[  75.30258318 -686.5630488 ]\n",
            " [  16.42764312 -148.29578741]]\n",
            "Time (RHS):\n",
            "0.0009284019470214844\n",
            "Result (RHS):\n",
            "[[  75.30258318 -686.5630488 ]\n",
            " [  16.42764312 -148.29578741]]\n",
            "Differences:\n",
            "[[-1.98951966e-13 -2.27373675e-13]\n",
            " [ 1.42108547e-13 -1.98951966e-13]]\n",
            "[[ 1.84741111e-13  2.27373675e-13]\n",
            " [-3.01980663e-13  1.42108547e-13]]\n",
            "[[1.42108547e-14 0.00000000e+00]\n",
            " [1.59872116e-13 5.68434189e-14]]\n",
            "time ratios:\n",
            "0.5102509720749382\n",
            "0.4496016626255629\n",
            "4.359013867488444\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LkLSurGMPzSm",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l-O0kE2dR-hq",
        "colab_type": "text"
      },
      "source": [
        "### **Lemma 3: Sum of $(A_i-B_iC)(A_i-B_iC)'$**\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V8BLX63KSTuB",
        "colab_type": "text"
      },
      "source": [
        "Let $A$ and $B$ be matrices of sizes $(m,k_1)$ and $(m,k_2)$ respectively, partitioned row-wise into chunks $A_1,...A_{m/n}$ and $B_1,...B_{l}$ of sizes $(n,k_1)$ and $(m,k_2)$ respectively where $m|n$ and $l=m/n$. I.e.\n",
        "\n",
        "$$A = \\begin{bmatrix} A_1 \\\\ A_2 \\\\ \\vdots \\\\ A_{l} \\end{bmatrix}, B = \\begin{bmatrix} B_1 \\\\ B_2 \\\\ \\vdots \\\\ B_{l} \\end{bmatrix}$$\n",
        "\n",
        "Let $C$ be a matrix of size $(k_1, k_2)$. Then:\n",
        "\n",
        "$$\\sum_{i=1}^l (A_i-B_iC)(A_i-B_iC)'=((A')^{r(k_1l,n)})'(A')^{r(k_1l,n)}$$"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "dwRV5b3ENSTw"
      },
      "source": [
        "### **Lemma 4: Sum of $A_i \\otimes B_i$**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W6DqHthUNVZ4",
        "colab_type": "text"
      },
      "source": [
        "Let $A$ and $B$ be matrices of sizes $(m,k_1)$ and $(m,k_2)$ respectively, partitioned row-wise into chunks $A_1,...A_{m/n}$ and $B_1,...B_{l}$ of sizes $(n,k_1)$ and $(n,k_2)$ respectively where $m|n$ and $l=m/n$. I.e.\n",
        "\n",
        "$$A = \\begin{bmatrix} A_1 \\\\ A_2 \\\\ \\vdots \\\\ A_{l} \\end{bmatrix}, B = \\begin{bmatrix} B_1 \\\\ B_2 \\\\ \\vdots \\\\ B_{l} \\end{bmatrix}$$\n",
        "\n",
        "Let $a_i=vec(A_i)$ and $b_i=vec(B_i)$ for all $i \\in \\{1,...,l\\}$ and define:\n",
        "\n",
        "$$\\tilde{a} = \\begin{bmatrix} a_1 \\\\ a_2 \\\\ \\vdots \\\\ a_{l} \\end{bmatrix}, \\tilde{b} = \\begin{bmatrix} b_1 \\\\ b_2 \\\\ \\vdots \\\\ b_{l} \\end{bmatrix}$$\n",
        "\n",
        "Then we have the following relationship:\n",
        "\n",
        "$$\\sum_{i=1}^{l} A_i \\otimes B_i = ((I_{k_1} \\otimes P_{n,k_2} \\otimes I_{n})vec(\\tilde{b}'\\tilde{a}))^{r(n^2,k_1k_2)}$$"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3i-5n7qUDqON",
        "colab_type": "code",
        "outputId": "807eed2e-c2b4-4f99-b1b3-4bd2bab4fca5",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 510
        }
      },
      "source": [
        "m = 4000\n",
        "n=2\n",
        "l=m//n\n",
        "k1=4000\n",
        "k2=1\n",
        "\n",
        "print('Example Shapes:')\n",
        "print('m: ', m)\n",
        "print('k1, k2: ', k1, k2)\n",
        "print('l: ', l)\n",
        "print('n: ', n)\n",
        "\n",
        "# A, B and C are random matrices\n",
        "A = np.random.randn(m,k1)\n",
        "B = np.random.randn(m,k2)\n",
        "\n",
        "t1 = time.time()\n",
        "\n",
        "# First we evaluate the LHS\n",
        "sumAkB = np.zeros((n**2,k1*k2))\n",
        "for i in np.arange(m//n):\n",
        "\n",
        "  # Work out the row-wise chunks\n",
        "  Ai = A[n*i:n*(i+1),:]\n",
        "  Bi = B[n*i:n*(i+1),:]\n",
        "\n",
        "  # Perform the summation\n",
        "  sumAkB = sumAkB + np.kron(Ai, Bi)\n",
        "\n",
        "t2 = time.time()\n",
        "print('Time (LHS):')\n",
        "print(t2-t1)\n",
        "print('Shape (LHS):')\n",
        "print(sumAkB.shape)\n",
        "\n",
        "# Now for RHS\n",
        "\n",
        "t1 = time.time()\n",
        "# This matrix only needs be calculated once\n",
        "K = comMat(n, k2)\n",
        "IPI = np.kron(np.kron(np.eye(k1), K), np.eye(n))\n",
        "\n",
        "t2 = time.time()\n",
        "atilde = matb2vecb(A,n)\n",
        "btilde = matb2vecb(B,n)\n",
        "\n",
        "vecba = mat2vec(btilde.transpose() @ atilde)\n",
        "\n",
        "RHS_noreshape = IPI @ vecba \n",
        "\n",
        "RHS_final = RHS_noreshape.reshape(k1*k2,n**2).transpose()\n",
        "\n",
        "\n",
        "t3 = time.time()\n",
        "print(t3-t1)\n",
        "print(t3-t2)\n",
        "\n",
        "print(np.allclose(RHS_final, sumAkB))\n",
        "print(RHS_final)\n",
        "print(sumAkB)\n",
        "\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Example Shapes:\n",
            "m:  4000\n",
            "k1, k2:  4000 1\n",
            "l:  2000\n",
            "n:  2\n",
            "Time (LHS):\n",
            "4.968502044677734\n",
            "Shape (LHS):\n",
            "(4, 4000)\n",
            "9.945883750915527\n",
            "0.28938937187194824\n",
            "True\n",
            "[[  3.42441572  18.17508019  63.62925992 ...  55.98095479  28.59916152\n",
            "  -32.34109229]\n",
            " [ -8.39253504 -42.67972678 -38.29225863 ...  -9.02363504  40.82507971\n",
            "  -14.20268779]\n",
            " [ 63.37303062 -71.32816464 -61.47922823 ...  24.33618291 -90.81564465\n",
            "  -43.89668294]\n",
            " [-72.79101746  -1.62471214   9.52570803 ... -46.98355562  48.78224011\n",
            "  -49.16646219]]\n",
            "[[  3.42441572  18.17508019  63.62925992 ...  55.98095479  28.59916152\n",
            "  -32.34109229]\n",
            " [ -8.39253504 -42.67972678 -38.29225863 ...  -9.02363504  40.82507971\n",
            "  -14.20268779]\n",
            " [ 63.37303062 -71.32816464 -61.47922823 ...  24.33618291 -90.81564465\n",
            "  -43.89668294]\n",
            " [-72.79101746  -1.62471214   9.52570803 ... -46.98355562  48.78224011\n",
            "  -49.16646219]]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BXAQ5QhY1CXl",
        "colab_type": "text"
      },
      "source": [
        "### **Lemma 5: Double sum of $A_iB_j' \\otimes C_iD_j'$**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QiXaNKC1NZMe",
        "colab_type": "text"
      },
      "source": [
        "\n",
        "Let $A$ and $C$ be matrices of sizes $(m_1,k_1)$ and $(m_1,k_2)$ respectively, partitioned row-wise into chunks $A_1,...A_{l_1}$ and $C_1,...C_{l_1}$ of sizes $(n_1,k_1)$ and $(n_1,k_2)$ respectively where $m_1|n_1$ and $l_1=m_1/n_1$. I.e.\n",
        "\n",
        "$$A = \\begin{bmatrix} A_1 \\\\ A_2 \\\\ \\vdots \\\\ A_{l_1} \\end{bmatrix}, C = \\begin{bmatrix} C_1 \\\\ C_2 \\\\ \\vdots \\\\ C_{l_1} \\end{bmatrix}$$\n",
        "\n",
        "Let $B$ and $D$ be matrices of sizes $(m_2,k_1)$ and $(m_2,k_2)$ respectively, partitioned row-wise into chunks $B_1,...B_{l_2}$ and $D_1,...D_{l_2}$ of sizes $(n_2,k_1)$ and $(n_2,k_2)$ respectively where $m_2|n_2$ and $l_2=m_2/n_2$. I.e.\n",
        "\n",
        "$$B = \\begin{bmatrix} B_1 \\\\ \\vdots \\\\ B_{l_2} \\end{bmatrix}, D = \\begin{bmatrix} D_1\\\\ \\vdots \\\\ D_{l_2} \\end{bmatrix}$$\n",
        "\n",
        "Then:\n",
        "\n",
        "$$\\sum_{j=1}^{l_2} \\sum_{i=1}^{l_1}A_iB_j' \\otimes C_iD_j' = ((I_{k_1} \\otimes P_{n_1,k_2} \\otimes I_{n_1})vec(\\tilde{c}'\\tilde{a}))^{r(n_1^2,k_1k_2)}\\big[((I_{k_1} \\otimes P_{n_2,k_2} \\otimes I_{n_2})vec(\\tilde{d}'\\tilde{b}))^{r(n_2^2,k_1k_2)}\\big]'$$\n",
        "\n",
        "(This follows directly from lemma 4 and the mixed product property of the kronecker product)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4jZVCwvq49Sa",
        "colab_type": "code",
        "outputId": "cd09fe85-2dfc-442a-821a-369fb5c9f297",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 140
        }
      },
      "source": [
        "# Dimensions\n",
        "import scipy.sparse\n",
        "\n",
        "# Partition sizes\n",
        "n1 = 2\n",
        "n2 = 2\n",
        "\n",
        "# Number of rows\n",
        "m1 = 1000\n",
        "m2 = 1200\n",
        "\n",
        "# Number of columns\n",
        "k1 = 1000\n",
        "k2 = 1000\n",
        "\n",
        "# Matrices \n",
        "\n",
        "A = np.random.randn(m1,k1)\n",
        "B = np.random.randn(m2,k1)\n",
        "C = np.random.randn(m1,k2)\n",
        "D = np.random.randn(m2,k2)\n",
        "\n",
        "runningSum = np.zeros((n1**2,n2**2))\n",
        "t1 = time.time()\n",
        "# Work out LHS\n",
        "for j in np.arange(m2//n2):\n",
        "\n",
        "  Bj = B[n2*j:n2*(j+1),:]\n",
        "  Dj = D[n2*j:n2*(j+1),:]\n",
        "\n",
        "  for i in np.arange(m1//n1):\n",
        "\n",
        "    Ai = A[n1*i:n1*(i+1),:]\n",
        "    Ci = C[n1*i:n1*(i+1),:]\n",
        "\n",
        "    runningSum = runningSum + np.kron(Ai @ Bj.transpose(), Ci @ Dj.transpose())\n",
        "\n",
        "t2 = time.time()\n",
        "print('Time (LHS)')\n",
        "print(t2-t1)\n",
        "\n",
        "# Now for RHS\n",
        "\n",
        "t1 = time.time()\n",
        "# This matrix only needs be calculated once\n",
        "Kmat1 = scipy.sparse.csr_matrix(comMat(n1, k2))\n",
        "IPI1 = scipy.sparse.kron(scipy.sparse.kron(scipy.sparse.identity(k1), Kmat1), scipy.sparse.identity(n1))\n",
        "Kmat2 = scipy.sparse.csr_matrix(comMat(n2, k2))\n",
        "IPI2 = scipy.sparse.kron(scipy.sparse.kron(scipy.sparse.identity(k1), Kmat2), scipy.sparse.identity(n2))\n",
        "\n",
        "t2 = time.time()\n",
        "atilde = matb2vecb(A,n1)\n",
        "btilde = matb2vecb(B,n2)\n",
        "ctilde = matb2vecb(C,n1)\n",
        "dtilde = matb2vecb(D,n2)\n",
        "\n",
        "vecca = mat2vec(ctilde.transpose() @ atilde)\n",
        "vecdb = mat2vec(dtilde.transpose() @ btilde)\n",
        "\n",
        "RHS1_noreshape = IPI1 @ vecca \n",
        "RHS2_noreshape = IPI2 @ vecdb \n",
        "\n",
        "RHS1_final = RHS1_noreshape.reshape(k1*k2,n1**2).transpose()\n",
        "RHS2_final = RHS2_noreshape.reshape(k1*k2,n2**2)\n",
        "\n",
        "RHS_final = RHS1_final @ RHS2_final\n",
        "\n",
        "t3 = time.time()\n",
        "print('Time (preprocessing)')\n",
        "print(t2-t1)\n",
        "print('Time (RHS)')\n",
        "print(t3-t2)\n",
        "\n",
        "print(np.allclose(RHS_final, runningSum))\n"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Time (LHS)\n",
            "11.228847980499268\n",
            "Time (preprocessing)\n",
            "0.36844754219055176\n",
            "Time (RHS)\n",
            "0.45260167121887207\n",
            "True\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5nZUsKzwzG7X",
        "colab_type": "text"
      },
      "source": [
        "### **Lemma 6: Double sum of $A_{i,j} \\otimes B_iC_j'$**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G8ZGn9uizVMo",
        "colab_type": "text"
      },
      "source": [
        "\n",
        "Let $A$ be a matrix of size $(m_1,m_2)$, partitioned block-wise into chunks $\\{A_{ij}\\}_{i=1,...l_1}^{j=1...l_2}$ of size $(n_1,n_2)$ where $m_1|n_1$ with $l_1=m_1/n_1$ $m_2|n_2$ with $l_2=m_2/n_2$. I.e.\n",
        "\n",
        "$$A = \\begin{bmatrix} A_{11} & A_{12} & ... & A_{1l_2} \\\\ A_{21} & A_{22} & ... & A_{2l_2} \\\\ \\vdots \\\\ A_{l_11} & A_{l_12} & ... & A_{l_1l_2} \\end{bmatrix}$$\n",
        "\n",
        "Let $B$ and $C$ be matrices of sizes $(m_1,k)$ and $(m_2,k)$ respectively, partitioned row-wise into chunks $B_1,...B_{l_1}$ and $C_1,...C_{l_2}$ of sizes $(n_1,k)$ and $(n_2,k)$ respectively. I.e.\n",
        "\n",
        "$$B = \\begin{bmatrix} B_1 \\\\ \\vdots \\\\ B_{l_1} \\end{bmatrix}, C = \\begin{bmatrix} C_1 \\\\ C_2 \\\\ \\vdots \\\\ C_{l_2} \\end{bmatrix}$$\n",
        "\n",
        "Then:\n",
        "\n",
        "$$\\sum_{j=1}^{l_2} \\sum_{i=1}^{l_1}A_{ij} \\otimes B_iC_j' = ((I_{n_2} \\otimes P_{n_1,n_2} \\otimes I_{n_1})vec(\\widetilde{(bc')}'\\tilde{a}))^{r(n_1^2,n_2^2)}$$\n",
        "\n",
        "$$\\widetilde{(bc')}=vecb(BC')$$\n",
        "$$\\widetilde{a}=vecb(A)$$"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "e738BtrQ877X",
        "colab_type": "code",
        "outputId": "3812e4b3-615a-42b0-c168-fdffbd69d8eb",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 228
        }
      },
      "source": [
        "n1 = 2\n",
        "n2 = 2\n",
        "m1 = 400\n",
        "m2 = 800\n",
        "k=100\n",
        "l1 = m1//n1\n",
        "l2 = m2//n2\n",
        "\n",
        "# Create A\n",
        "A = np.random.randn(m1,m2)\n",
        "\n",
        "B = np.random.randn(m1,k)\n",
        "C = np.random.randn(m2,k)\n",
        "\n",
        "runningSum = np.zeros((n1**2,n2**2))\n",
        "\n",
        "# LHS calculation\n",
        "t1 = time.time()\n",
        "for i in np.arange(l1):\n",
        "  for j in np.arange(l2):\n",
        "\n",
        "    Aij = A[n1*i:n1*(i+1),n2*j:n2*(j+1)]\n",
        "    Bi = B[n1*i:n1*(i+1),:]\n",
        "    Cj = C[n2*j:n2*(j+1),:]\n",
        "\n",
        "    runningSum = runningSum + np.kron(Aij, Bi @ Cj.transpose())\n",
        "\n",
        "t2 = time.time()\n",
        "print(t2-t1)\n",
        "\n",
        "# RHS calculation\n",
        "t1 = time.time()\n",
        "\n",
        "# Change A to stacked form\n",
        "As = A.reshape((m1//n1,n1,m2//n2,n2)).transpose(0,2,1,3).reshape(m1*m2//n2,n2)\n",
        "\n",
        "# B times C transpose\n",
        "BCt = B @ C.transpose()\n",
        "\n",
        "# B times C transpose (stacked)\n",
        "BCts = BCt.reshape((m1//n1,n1,m2//n2,n2)).transpose(0,2,1,3).reshape(m1*m2//n2,n2)\n",
        "\n",
        "# This matrix only needs be calculated once\n",
        "K = comMat(n1, n2)\n",
        "IPI = scipy.sparse.kron(scipy.sparse.kron(scipy.sparse.identity(n2), K), scipy.sparse.identity(n1))\n",
        "\n",
        "t2 = time.time()\n",
        "atilde = matb2vecb(As,n1)\n",
        "bctilde = matb2vecb(BCts,n1)\n",
        "\n",
        "vecbca = mat2vec(bctilde.transpose() @ atilde)\n",
        "\n",
        "RHS_noreshape = IPI @ vecbca \n",
        "\n",
        "RHS_final = RHS_noreshape.reshape(n1**2,n2**2).transpose()\n",
        "\n",
        "\n",
        "t3 = time.time()\n",
        "print(t3-t1)\n",
        "print(t3-t2)\n",
        "\n",
        "print(np.allclose(RHS_final, runningSum))\n",
        "print(RHS_final)\n",
        "print(runningSum)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "2.628631353378296\n",
            "0.020653724670410156\n",
            "0.009497404098510742\n",
            "True\n",
            "[[ 1550.21361498  2590.53767701 -2047.52861912   213.8205423 ]\n",
            " [  268.72422262   977.72926385   -52.94880813  1720.64550173]\n",
            " [  742.71205499  3316.38039641  4157.54250268  2200.37850575]\n",
            " [ -886.16548126   950.98653872  5493.9735444   3485.59525685]]\n",
            "[[ 1550.21361498  2590.53767701 -2047.52861912   213.8205423 ]\n",
            " [  268.72422262   977.72926385   -52.94880813  1720.64550173]\n",
            " [  742.71205499  3316.38039641  4157.54250268  2200.37850575]\n",
            " [ -886.16548126   950.98653872  5493.9735444   3485.59525685]]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qFgyuR-mvJbs",
        "colab_type": "text"
      },
      "source": [
        "## Comments"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tNF39wnNC900",
        "colab_type": "text"
      },
      "source": [
        "This section contains comments relating the lemmas above to BLMM"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "65-yTZC7tNd3",
        "colab_type": "text"
      },
      "source": [
        "### Derivative of $l$ with respect to $D_k$"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bgvFn9suDDvS",
        "colab_type": "text"
      },
      "source": [
        "We already have that; $$\\frac{\\partial l}{\\partial D_k} = \\frac{1}{2}\\sum_{j=1}^{l_k}(T_{(k,j)}u)(T_{(k,j)}u)' - \\frac{1}{2}\\sum_{j=1}^kT_{(k,j)}T_{(k,j)}'$$\n",
        "\n",
        "\n",
        "Consider the first term, given by:\n",
        "$$\\frac{1}{2}\\sum_{j=1}^{l_k}(T_{(k,j)}u)(T_{(k,j)}u)'$$\n",
        "$$= \\sigma^{-2}\\frac{1}{2}\\sum_{j=1}^{l_k}(Z'_{(k,j)}e-Z'_{(k,j)}ZD(I+Z'ZD)^{-1}Z'e)(e'Z_{(k,j)}-e'ZD(I+Z'ZD)^{-1}Z'Z_{(k,j)})$$\n",
        "$$ = \\sigma^{-2}\\frac{1}{2}\\sum_{j=1}^{l_k} (A_j - B_j)(A_j-B_j)'$$\n",
        "\n",
        "Upon letting:\n",
        " - $A_j = Z'_{(k,j)}e$ with $A = Z'_{(k)}e$ \n",
        " - $B_j = Z'_{(k,j)}ZD(I+Z'ZD)^{-1}Z'e$ with $B = Z'_{(k)}ZD(I+Z'ZD)^{-1}Z'e$ \n",
        "\n",
        "Noting, now, by Lemma 1, that this is equal to:\n",
        "\n",
        "$$= \\sigma^{-2}\\frac{1}{2}((A')^{r(l_k,q_k)}-(B')^{r(l_k,q_k)})'((A')^{r(l_k,q_k)}-(B')^{r(l_k,q_k)})$$\n",
        "\n",
        "$$= \\sigma^{-2}\\frac{1}{2}\\bigg((e'Z_{(k)})^{r(l_k,q_k)}-(e'Z(I+DZ'Z)^{-1}DZ'Z_{(k)})^{r(l_k,q_k)}\\bigg)'\\bigg((e'Z_{(k)})^{r(l_k,q_k)}-(e'Z(I+DZ'Z)^{-1}DZ'Z_{(k)})^{r(l_k,q_k)}\\bigg)$$\n",
        "\n",
        "This is how the first term is computed in a computationally efficient manner. The second term can be computed by first noting:\n",
        "\n",
        "$$\\frac{1}{2}\\sum_{j=1}^{l_k}T_{(k,j)}T_{(k,j)}'=\\frac{1}{2}\\sum_{j=1}^{l_k}\\bigg(Z_{(k,j)}'Z_{(k,j)} - Z_{(k,j)}'ZD(I+Z'ZD)^{-1}Z_{(k,j)}\\bigg)$$\n",
        "$$=\\frac{1}{2}\\sum_{j=1}^{l_k}Z_{(k,j)}'Z_{(k,j)} - \\frac{1}{2}\\sum_{j=1}^{l_k}\\bigg( Z_{(k,j)}'ZD(I+Z'ZD)^{-1}Z_{(k,j)}\\bigg)$$\n",
        "$$=\\frac{1}{2}\\sum_{j=1}^{l_k}Z_{(k,j)}'Z_{(k,j)} - \\frac{1}{2}\\sum_{j=1}^{l_k}A_jB_j'$$\n",
        "\n",
        "Upon letting:\n",
        " - $A_j = Z'_{(k,j)}Z$ with $A = Z'_{(k)}Z$ \n",
        " - $B_j = Z'_{(k,j)}Z(I+DZ'Z)^{-1}D$ with $B = Z'_{(k)}Z(I+DZ'Z)^{-1}D$ \n",
        "\n",
        "Upon applying lemma 1, we see that this equals;\n",
        "\n",
        " $$=\\frac{1}{2}\\sum_{j=1}^{l_k}Z_{(k,j)}'Z_{(k,j)} - ((A')^{r(l_k\\sum_il_iq_i,q_k)})'(B')^{r(l_k\\sum_il_iq_i,q_k)}$$\n",
        " $$=\\frac{1}{2}\\sum_{j=1}^{l_k}Z_{(k,j)}'Z_{(k,j)} - ((Z'Z_{(k,j)})^{r(l_k\\sum_il_iq_i,q_k)})'(D(I+Z'ZD)^{-1}Z_{(k,j)})^{r(l_k\\sum_il_iq_i,q_k)}$$\n",
        "\n",
        "Note that the first term can be computed prior to any iteration as it does not depend on any parameter estimates and therefore can be ignored here."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RJITQSX9tRFP",
        "colab_type": "text"
      },
      "source": [
        "### Covariance of derivative of $l$ with respect to $D_k$ and $l$ with respect to $\\sigma^2$ \n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iS3J_HGYtYox",
        "colab_type": "text"
      },
      "source": [
        "The covariance of the $\\frac{\\partial l}{\\partial \\sigma^2}$ and $\\frac{\\partial l}{\\partial D_k}$ is given by:\n",
        "\n",
        "$$cov\\bigg(\\frac{\\partial l}{\\partial \\sigma^2},\\frac{\\partial l}{\\partial D_k}\\bigg) = \\frac{1}{2\\sigma^2}\\text{vec}'\\bigg(\\sum_{j=1}^{l_k}T_{(k,j)}T_{(k,j)}'\\bigg)\\mathcal{D^{+'}}$$\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "The second term can be computed by first noting:\n",
        "\n",
        "$$\\sum_{j=1}^{l_k}T_{(k,j)}T_{(k,j)}'=\\sum_{j=1}^{l_k}\\bigg(Z_{(k,j)}'Z_{(k,j)} - Z_{(k,j)}'ZD(I+Z'ZD)^{-1}Z_{(k,j)}\\bigg)$$\n",
        "$$=\\sum_{j=1}^{l_k}Z_{(k,j)}'Z_{(k,j)} - \\sum_{j=1}^{l_k}\\bigg( Z_{(k,j)}'ZD(I+Z'ZD)^{-1}Z_{(k,j)}\\bigg)$$\n",
        "$$=\\sum_{j=1}^{l_k}Z_{(k,j)}'Z_{(k,j)} - \\sum_{j=1}^{l_k}A_jB_j'$$\n",
        "\n",
        "Upon letting:\n",
        " - $A_j = Z'_{(k,j)}Z$ with $A = Z'_{(k)}Z$ \n",
        " - $B_j = Z'_{(k,j)}Z(I+DZ'Z)^{-1}D$ with $B = Z'_{(k)}Z(I+DZ'Z)^{-1}D$ \n",
        "\n",
        "Upon applying lemma 1, we see that this equals;\n",
        "\n",
        " $$=\\sum_{j=1}^{l_k}Z_{(k,j)}'Z_{(k,j)} - ((A')^{r(l_k\\sum_il_iq_i,q_k)})'(B')^{r(l_k\\sum_il_iq_i,q_k)}$$\n",
        " $$=\\sum_{j=1}^{l_k}Z_{(k,j)}'Z_{(k,j)} - ((Z'Z_{(k)})^{r(l_k\\sum_il_iq_i,q_k)})'(D(I+Z'ZD)^{-1}Z_{(k)})^{r(l_k\\sum_il_iq_i,q_k)}$$\n",
        "\n",
        "Note that the first term can be computed prior to any iteration as it does not depend on any parameter estimates and therefore can be ignored here. Therefore, we have:\n",
        "\n",
        "$$cov\\bigg(\\frac{\\partial l}{\\partial \\sigma^2},\\frac{\\partial l}{\\partial D_k}\\bigg) = \\frac{1}{2\\sigma^2}\\text{vec}'\\bigg(\\sum_{j=1}^{l_k}Z_{(k,j)}'Z_{(k,j)}\\bigg)\\mathcal{D^{+'}} - \\frac{1}{2\\sigma^2}\\text{vec}'\\bigg(((Z'Z_{(k)})^{r(l_k\\sum_il_iq_i,q_k)})'(D(I+Z'ZD)^{-1}Z_{(k)})^{r(l_k\\sum_il_iq_i,q_k)}\\bigg)\\mathcal{D^{+'}}$$"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MOI0oEmUw9b6",
        "colab_type": "text"
      },
      "source": [
        "### Covariance of derivative of $l$ with respect to $D_k$"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FiXwt-uxxDAA",
        "colab_type": "text"
      },
      "source": [
        "The covariance of the $\\frac{\\partial l}{\\partial D_{k_1}}$ and $\\frac{\\partial l}{\\partial D_{k_2}}$ is given by:\n",
        "\n",
        "$$\\text{cov}\\bigg(\\frac{\\delta l}{\\delta \\text{vech}(D_{k_1})},\\frac{\\delta l}{\\delta \\text{vech}(D_{k_2})}\\bigg)=\\frac{1}{2}\\mathcal{D}_{k_1}^+\\sum_{j=1}^{l_{k_2}}\\sum_{i=1}^{l_{k_1}}(R_{(k_1,k_2,i,j)}\\otimes R_{(k_1,k_2, i,j)}\n",
        ")\\mathcal{D}_{k_2}^{+'}$$\n",
        "\n",
        "Where $R_{(k_1,k_2,i,j)}=Z'_{(k_1,i)}Z_{(k_2,j)}-Z_{(k_1,i)}'ZD(I+Z'ZD)^{-1}Z'Z_{(k_2,j)}$. Let:\n",
        "\n",
        " - $A_{i,j}=Z'_{(k_1,i)}Z_{(k_2,j)}$ where $A=Z'Z_{(k_1,k_2)}$.\n",
        " - $B_{i,j}=Z_{(k_1,i)}'ZD(I+Z'ZD)^{-1}Z'Z_{(k_2,j)}$\n",
        "\n",
        "It follows that:\n",
        "\n",
        " $$\\text{cov}\\bigg(\\frac{\\delta l}{\\delta \\text{vech}(D_{k_1})},\\frac{\\delta l}{\\delta \\text{vech}(D_{k_2})}\\bigg)=\\frac{1}{2}\\mathcal{D}_{k_1}^+\\sum_{j=1}^{l_{k_2}}\\sum_{i=1}^{l_{k_1}}\\bigg((A_{i,j} - B_{i,j})\\otimes (A_{i,j} - B_{i,j})\\bigg)\\mathcal{D}_{k_2}^{+'}$$\n",
        " $$=\\frac{1}{2}\\mathcal{D}_{k_1}^+\\bigg(\\sum_{j=1}^{l_{k_2}}\\sum_{i=1}^{l_{k_1}}(A_{i,j}\\otimes A_{i,j}) - \\sum_{j=1}^{l_{k_2}}\\sum_{i=1}^{l_{k_1}}(A_{i,j}\\otimes B_{i,j}) - \\sum_{j=1}^{l_{k_2}}\\sum_{i=1}^{l_{k_1}}(B_{i,j}\\otimes A_{i,j}) - \\sum_{j=1}^{l_{k_2}}\\sum_{i=1}^{l_{k_1}}(B_{i,j}\\otimes B_{i,j})\\bigg)\\mathcal{D}_{k_2}^{+'}$$\n",
        "\n",
        "Each of these can be calculated in the below mannor via a variant of lemma 6:\n",
        "\n",
        "\n",
        "$$\\sum_{j=1}^{l_2} \\sum_{i=1}^{l_1}A_{i,j} \\otimes A_{i,j} = ((I_{q_2} \\otimes P_{q_1,q_2} \\otimes I_{q_1})vec(\\tilde{a}'\\tilde{a}))^{r(q_1^2,q_2^2)}$$\n",
        "\n",
        "$$\\sum_{j=1}^{l_2} \\sum_{i=1}^{l_1}B_{i,j} \\otimes A_{i,j} = ((I_{q_2} \\otimes P_{q_1,q_2} \\otimes I_{q_1})vec(\\tilde{b}'\\tilde{a}))^{r(q_1^2,q_2^2)}$$\n",
        "\n",
        "$$\\sum_{j=1}^{l_2} \\sum_{i=1}^{l_1}A_{i,j} \\otimes B_{i,j} = ((I_{q_2} \\otimes P_{q_1,q_2} \\otimes I_{q_1})vec(\\tilde{a}'\\tilde{b}))^{r(q_1^2,q_2^2)}$$\n",
        "\n",
        "$$\\sum_{j=1}^{l_2} \\sum_{i=1}^{l_1}B_{i,j} \\otimes B_{i,j} = ((I_{q_2} \\otimes P_{q_1,q_2} \\otimes I_{q_1})vec(\\tilde{b}'\\tilde{b}))^{r(q_1^2,q_2^2)}$$\n",
        "\n",
        "Where the lower case letter and tilde notation represents the $vecb(\\cdot)$ operator. I.e. if $X$ is partitioned into blocks of size $(n_1 \\times n_2)$ like so:\n",
        "\n",
        "$$X = \\begin{bmatrix} X_{1,1} & X_{1,2} & ... & X_{1,l_2} \\\\ X_{2,1} & X_{2,2} & ... & X_{2,l_2} \\\\ \\vdots \\\\ X_{l_1,1} & X_{l_1,2} & ... & X_{l_1,l_2} \\end{bmatrix}$$\n",
        "\n",
        "Then $\\tilde{x}=vecb^{(n_1,n_2)}(X)$ is given by:\n",
        "\n",
        "$$\\tilde{x} = \\begin{bmatrix} vec'(X_{1,1}) \\\\ vec'(X_{1,2}) \\\\ \\vdots \\\\ vec'(X_{1,l_2}) \\\\ vec'(X_{2,1}) \\\\ \\vdots \\\\ vec'(X_{l_1,l_2}) \\end{bmatrix}$$\n",
        "\n",
        "It is also worth noting that the first sum given above does not involve any estimated parameters so does not need to be recomputed for each iteration and can be computed beforehand."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Znm81lAfRrEW",
        "colab_type": "text"
      },
      "source": [
        "## Final Functions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "00x_ZkSSRt7a",
        "colab_type": "text"
      },
      "source": [
        "This section contains functions, and tests of those functions, intended for use in BLMM, based on the previous lemmas."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "et8qrPakTYi4"
      },
      "source": [
        "### Function: ``sumAijBijt``"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2sPHtrOiTaed",
        "colab_type": "text"
      },
      "source": [
        "This function takes in the following inputs:\n",
        "\n",
        " - A: A 2D matrix of dimension $(m_1 \\times m_2)$.\n",
        " - B: A 2D matrix of dimension $(m_3 \\times m_4)$.\n",
        " - pA: The size of the block partitions of $A$, e.g. if $A_{i,j}$ is of dimension $(n_1 \\times n_2)$ the ``pA=[n1, n2]``.\n",
        " - pB: The size of the block partitions of $B$, e.g. if $B_{i,j}$ is of dimension $(n_3 \\times n_4)$ the ``pB=[n3, n4]``.\n",
        "\n",
        "And gives the following output:\n",
        "\n",
        " - S: The sum of the partitions of $A$ multiplied by the transpose of the partitions of $B$; i.e. \n",
        " \n",
        " $$\\sum_{j=1}^{l_1}\\sum_{i=1}^{l_2} A_{i,j}B_{i,j}'$$\n",
        "  \n",
        "  where $l_1:=\\frac{m_1}{n_1}=\\frac{m_3}{n_3}$ and $l_2:=\\frac{m_2}{n_2}=\\frac{m_4}{n_4}$."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JIE1JH7FTZjN",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0yBU9wFrRjrC",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "n1 = 2\n",
        "n2 = 2\n",
        "m1 = 400\n",
        "m2 = 800\n",
        "k=100\n",
        "l1 = m1//n1\n",
        "l2 = m2//n2\n",
        "\n",
        "# Create A\n",
        "A = np.random.randn(m1,m2)\n",
        "\n",
        "B = np.random.randn(m1,k)\n",
        "C = np.random.randn(m2,k)\n",
        "\n",
        "\n",
        "# RHS calculation\n",
        "t1 = time.time()\n",
        "\n",
        "# Change A to stacked form\n",
        "As = A.reshape((m1//n1,n1,m2//n2,n2)).transpose(0,2,1,3).reshape(m1*m2//n2,n2)"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}