{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\operatorname{tr}\\left(\\left(0.5 \\left(AA G_{AA} + 0.5 AB G_{BA} + 0.5 AR G_{RA}\\right) AB + \\left(AA G_{AB} + 0.5 AB G_{BB} + 0.5 AR G_{RB}\\right) BB + 0.5 \\left(AA G_{AR} + 0.5 AB G_{BR} + 0.5 AR G_{RR}\\right) RB\\right) G_{BA} \\right) + \\operatorname{tr}\\left(\\left(0.25 BA G_{AA} AB + 0.5 BA G_{AB} BB + 0.25 BA G_{AR} RB\\right) G_{BB} \\right) + \\operatorname{tr}\\left(\\left(0.25 RA G_{AA} AB + 0.5 RA G_{AB} BB + 0.25 RA G_{AR} RB\\right) G_{BR} \\right) + 0.5 \\operatorname{tr}\\left(\\left(AA G_{AB} + 0.5 AB G_{BB} + 0.5 AR G_{RB}\\right) BA G_{AA} \\right) + 0.5 \\operatorname{tr}\\left(\\left(AA G_{AB} + 0.5 AB G_{BB} + 0.5 AR G_{RB}\\right) BR G_{RA} \\right) + 0.25 \\operatorname{tr}\\left(BA G_{AB} BA G_{AB} \\right) + 0.25 \\operatorname{tr}\\left(BA G_{AB} BR G_{RB} \\right) + 0.25 \\operatorname{tr}\\left(RA G_{AB} BA G_{AR} \\right) + 0.25 \\operatorname{tr}\\left(RA G_{AB} BR G_{RR} \\right)$"
      ],
      "text/plain": [
       "Trace((0.5*(AA*G_AA + 0.5*AB*G_BA + 0.5*AR*G_RA)*AB + (AA*G_AB + 0.5*AB*G_BB + 0.5*AR*G_RB)*BB + 0.5*(AA*G_AR + 0.5*AB*G_BR + 0.5*AR*G_RR)*RB)*G_BA) + Trace((0.25*BA*G_AA*AB + 0.5*BA*G_AB*BB + 0.25*BA*G_AR*RB)*G_BB) + Trace((0.25*RA*G_AA*AB + 0.5*RA*G_AB*BB + 0.25*RA*G_AR*RB)*G_BR) + 0.5*Trace((AA*G_AB + 0.5*AB*G_BB + 0.5*AR*G_RB)*BA*G_AA) + 0.5*Trace((AA*G_AB + 0.5*AB*G_BB + 0.5*AR*G_RB)*BR*G_RA) + 0.25*Trace(BA*G_AB*BA*G_AB) + 0.25*Trace(BA*G_AB*BR*G_RB) + 0.25*Trace(RA*G_AB*BA*G_AR) + 0.25*Trace(RA*G_AB*BR*G_RR)"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from sympy import MatrixSymbol, BlockMatrix, ZeroMatrix, trace, block_collapse\n",
    "\n",
    "# Define MatrixSymbols with specific sizes\n",
    "# A = 2x2, B = 3x3, R = 4x4\n",
    "AA = MatrixSymbol(\"AA\", 2, 2)  # 2x2\n",
    "AB = MatrixSymbol(\"AB\", 2, 3)  # 2x3\n",
    "BA = MatrixSymbol(\"BA\", 3, 2)  # 3x2\n",
    "AR = MatrixSymbol(\"AR\", 2, 4)  # 2x4\n",
    "RA = MatrixSymbol(\"RA\", 4, 2)  # 4x2\n",
    "BB = MatrixSymbol(\"BB\", 3, 3)  # 3x3\n",
    "BR = MatrixSymbol(\"BR\", 3, 4)  # 3x4\n",
    "RB = MatrixSymbol(\"RB\", 4, 3)  # 4x3\n",
    "RR = MatrixSymbol(\"RR\", 4, 4)  # 4x4\n",
    "\n",
    "# Create block matrices VA and VB using BlockMatrix\n",
    "# We need to use ZeroMatrix with appropriate dimensions for padding\n",
    "VA = BlockMatrix(\n",
    "    [\n",
    "        [AA, 0.5 * AB, 0.5 * AR],\n",
    "        [0.5 * BA, ZeroMatrix(3, 3), ZeroMatrix(3, 4)],\n",
    "        [0.5 * RA, ZeroMatrix(4, 3), ZeroMatrix(4, 4)],\n",
    "    ]\n",
    ")\n",
    "\n",
    "VB = BlockMatrix(\n",
    "    [\n",
    "        [ZeroMatrix(2, 2), 0.5 * AB, ZeroMatrix(2, 4)],\n",
    "        [0.5 * BA, BB, 0.5 * BR],\n",
    "        [ZeroMatrix(4, 2), 0.5 * RB, ZeroMatrix(4, 4)],\n",
    "    ]\n",
    ")\n",
    "\n",
    "# Define G matrix symbols with matching dimensions\n",
    "G_AA = MatrixSymbol(\"G_AA\", 2, 2)\n",
    "G_AB = MatrixSymbol(\"G_AB\", 2, 3)\n",
    "G_BA = MatrixSymbol(\"G_BA\", 3, 2)\n",
    "G_AR = MatrixSymbol(\"G_AR\", 2, 4)\n",
    "G_RA = MatrixSymbol(\"G_RA\", 4, 2)\n",
    "G_BB = MatrixSymbol(\"G_BB\", 3, 3)\n",
    "G_BR = MatrixSymbol(\"G_BR\", 3, 4)\n",
    "G_RB = MatrixSymbol(\"G_RB\", 4, 3)\n",
    "G_RR = MatrixSymbol(\"G_RR\", 4, 4)\n",
    "\n",
    "# Create G matrix as a BlockMatrix\n",
    "G = BlockMatrix([[G_AA, G_AB, G_AR], [G_BA, G_BB, G_BR], [G_RA, G_RB, G_RR]])\n",
    "\n",
    "\n",
    "# Calculate the trace of VA@G@VB@G\n",
    "# First, let's calculate the product step by step\n",
    "product = block_collapse(VA @ G @ VB @ G)\n",
    "\n",
    "# Calculate the trace\n",
    "result = trace(product)\n",
    "result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Trace((0.5*(AA*G_AA + 0.5*AB*G_BA)*AB + (AA*G_AB + 0.5*AB*G_BB)*BB)*G_BA) + Trace((0.25*BA*G_AA*AB + 0.5*BA*G_AB*BB)*G_BB) + 0.5*Trace((AA*G_AB + 0.5*AB*G_BB)*BA*G_AA) + 0.25*Trace(BA*G_AB*BA*G_AB)\n"
     ]
    }
   ],
   "source": [
    "# Define MatrixSymbols with specific sizes\n",
    "AA = MatrixSymbol(\"AA\", 2, 2)  # 2x2\n",
    "AB = MatrixSymbol(\"AB\", 2, 3)  # 2x3\n",
    "BA = MatrixSymbol(\"BA\", 3, 2)  # 3x2\n",
    "BB = MatrixSymbol(\"BB\", 3, 3)  # 3x3\n",
    "\n",
    "# Create block matrices VA and VB using BlockMatrix\n",
    "# R-related terms are replaced with zero matrices\n",
    "VA = BlockMatrix(\n",
    "    [\n",
    "        [AA, 0.5 * AB, ZeroMatrix(2, 4)],\n",
    "        [0.5 * BA, ZeroMatrix(3, 3), ZeroMatrix(3, 4)],\n",
    "        [ZeroMatrix(4, 2), ZeroMatrix(4, 3), ZeroMatrix(4, 4)],\n",
    "    ]\n",
    ")\n",
    "\n",
    "VB = BlockMatrix(\n",
    "    [\n",
    "        [ZeroMatrix(2, 2), 0.5 * AB, ZeroMatrix(2, 4)],\n",
    "        [0.5 * BA, BB, ZeroMatrix(3, 4)],\n",
    "        [ZeroMatrix(4, 2), ZeroMatrix(4, 3), ZeroMatrix(4, 4)],\n",
    "    ]\n",
    ")\n",
    "\n",
    "# Define G matrix symbols with matching dimensions (kept original)\n",
    "G_AA = MatrixSymbol(\"G_AA\", 2, 2)\n",
    "G_AB = MatrixSymbol(\"G_AB\", 2, 3)\n",
    "G_BA = MatrixSymbol(\"G_BA\", 3, 2)\n",
    "G_AR = MatrixSymbol(\"G_AR\", 2, 4)\n",
    "G_RA = MatrixSymbol(\"G_RA\", 4, 2)\n",
    "G_BB = MatrixSymbol(\"G_BB\", 3, 3)\n",
    "G_BR = MatrixSymbol(\"G_BR\", 3, 4)\n",
    "G_RB = MatrixSymbol(\"G_RB\", 4, 3)\n",
    "G_RR = MatrixSymbol(\"G_RR\", 4, 4)\n",
    "\n",
    "# Create G matrix as a BlockMatrix (kept original)\n",
    "G = BlockMatrix([[G_AA, G_AB, G_AR], [G_BA, G_BB, G_BR], [G_RA, G_RB, G_RR]])\n",
    "\n",
    "# Calculate the product and trace\n",
    "product = block_collapse(VA @ G @ VB @ G)\n",
    "result = trace(product)\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Trace(AA*G_AB*BB*G_BA)\n"
     ]
    }
   ],
   "source": [
    "from sympy import MatrixSymbol, BlockMatrix, ZeroMatrix, trace, block_collapse\n",
    "\n",
    "# Define MatrixSymbols with specific sizes\n",
    "AA = MatrixSymbol(\"AA\", 2, 2)  # 2x2\n",
    "BB = MatrixSymbol(\"BB\", 3, 3)  # 3x3\n",
    "\n",
    "# Create block matrices VA and VB using BlockMatrix with only AA and BB\n",
    "VA = BlockMatrix(\n",
    "    [\n",
    "        [AA, ZeroMatrix(2, 3), ZeroMatrix(2, 4)],\n",
    "        [ZeroMatrix(3, 2), ZeroMatrix(3, 3), ZeroMatrix(3, 4)],\n",
    "        [ZeroMatrix(4, 2), ZeroMatrix(4, 3), ZeroMatrix(4, 4)],\n",
    "    ]\n",
    ")\n",
    "\n",
    "VB = BlockMatrix(\n",
    "    [\n",
    "        [ZeroMatrix(2, 2), ZeroMatrix(2, 3), ZeroMatrix(2, 4)],\n",
    "        [ZeroMatrix(3, 2), BB, ZeroMatrix(3, 4)],\n",
    "        [ZeroMatrix(4, 2), ZeroMatrix(4, 3), ZeroMatrix(4, 4)],\n",
    "    ]\n",
    ")\n",
    "\n",
    "# Define G matrix symbols with matching dimensions\n",
    "G_AA = MatrixSymbol(\"G_AA\", 2, 2)\n",
    "G_AB = MatrixSymbol(\"G_AB\", 2, 3)\n",
    "G_BA = MatrixSymbol(\"G_BA\", 3, 2)\n",
    "G_AR = MatrixSymbol(\"G_AR\", 2, 4)\n",
    "G_RA = MatrixSymbol(\"G_RA\", 4, 2)\n",
    "G_BB = MatrixSymbol(\"G_BB\", 3, 3)\n",
    "G_BR = MatrixSymbol(\"G_BR\", 3, 4)\n",
    "G_RB = MatrixSymbol(\"G_RB\", 4, 3)\n",
    "G_RR = MatrixSymbol(\"G_RR\", 4, 4)\n",
    "\n",
    "# Create G matrix as a BlockMatrix\n",
    "G = BlockMatrix([[G_AA, G_AB, G_AR], [G_BA, G_BB, G_BR], [G_RA, G_RB, G_RR]])\n",
    "\n",
    "# Calculate the product and simplify\n",
    "product = block_collapse(VA @ G @ VB @ G)\n",
    "\n",
    "# Calculate the trace\n",
    "result = trace(product)\n",
    "simplified_result = result.simplify()\n",
    "\n",
    "print(simplified_result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}