{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "773f49af",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sage.all import *\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "a9fd2992",
   "metadata": {},
   "outputs": [],
   "source": [
    "class NBRW:\n",
    "    \"\"\"A class to store several relevant attributes of a graph G in Non-Backtracking Random Walks\"\"\"\n",
    "    \n",
    "    def __init__(self, G):\n",
    "        \"\"\"Accepts a a sage graph\"\"\"\n",
    "        \n",
    "        self.G = G  # store graph\n",
    "        self.m = len(G.edges())  # number of edges\n",
    "        self.n = len(G.vertices())  # number of vertices\n",
    "        \n",
    "        self.S = self.S_mat()\n",
    "        self.T = self.T_mat()\n",
    "        self.tau = self.tau_mat()\n",
    "        self.B = self.B_mat()\n",
    "        self.Dinv = self.D_inv()\n",
    "        \n",
    "        self.Pnb = self.nb_trans_mat()\n",
    "        self.Znb = self.Znb_mat()\n",
    "        self.Wnb = self.Wnb_mat()\n",
    "        self.Mev = self.M_ev_NBRW()\n",
    "        self.Znb_e = self.Znb_e_mat()\n",
    "        \n",
    "    \n",
    "    \n",
    "    def S_mat(self):\n",
    "        \"\"\"Returns matrix S such that ST = C (TS=A)\n",
    "        S maps from edge space to vertex space\n",
    "        S is a 2m x n matrix\"\"\"\n",
    "        \n",
    "        edges = list(self.G.edges())\n",
    "        \n",
    "        S = matrix.zero(2*self.m, self.n)\n",
    "        for x in range(self.n):\n",
    "            # iterate through nodes\n",
    "            for j in range(self.m):\n",
    "                # iterate through edges\n",
    "                if edges[j][1] == x:\n",
    "                    S[j, x] = 1\n",
    "                if edges[j][0] == x:\n",
    "                    S[j+self.m, x] = 1\n",
    "        return S\n",
    "    \n",
    "    def T_mat(self):\n",
    "        \"\"\"Returns T matrix such that ST = C (TS=A)\n",
    "        Maps from vertex to edge\n",
    "        n x 2m\"\"\"\n",
    "        \n",
    "        edges = list(self.G.edges())\n",
    "        \n",
    "        T = matrix.zero(self.n, 2*self.m)\n",
    "        for x in range(self.n):\n",
    "            # iterate through nodes\n",
    "            for j in range(self.m):\n",
    "                # iterate through edges\n",
    "                if edges[j][0] == x:\n",
    "                    T[x,j] = 1\n",
    "                if edges[j][1] == x:\n",
    "                    T[x, j+self.m] = 1\n",
    "        \n",
    "        return T\n",
    "    \n",
    "    def tau_mat(self):\n",
    "        \"\"\"Return matrix tau such that ST - tau = B\n",
    "        2m x 2m\n",
    "        Switches directed edges\n",
    "        \"\"\"\n",
    "        \n",
    "        edges = list(G.edges())\n",
    "        \n",
    "        zero = matrix.zero(self.m)  # m x m zeros\n",
    "        I = identity_matrix(self.m)  # m x m identity\n",
    "        \n",
    "        t = block_matrix([[zero, I], [I, zero]])\n",
    "        \n",
    "        return t  # like an identity with opposite diagonal\n",
    "    \n",
    "    def B_mat(self):\n",
    "        \"\"\"Return matrix B = ST - tau\n",
    "        non-backtracking edge adjacency matrix\"\"\"\n",
    "        \n",
    "        return self.S * self.T - self.tau\n",
    "    \n",
    "    def nb_trans_mat(self):\n",
    "        \"\"\"return NBRW edge space transition probability matrix\n",
    "        P_{nb} = (D-I)^{-1}B\"\"\"\n",
    "        b = self.B\n",
    "        row_sums = sum(b)\n",
    "        row_sums = [i ^ -1 for i in row_sums]\n",
    "        return b * diagonal_matrix(row_sums)\n",
    "    \n",
    "    def D_inv(self):\n",
    "        \"\"\"D^{-1}\"\"\"\n",
    "        return np.diag([1 / d for d in self.G.degree()])\n",
    "    \n",
    "    def trans(self):\n",
    "        \"\"\"vertex space transition matrix of a simple random walk\"\"\"\n",
    "        D = self.D\n",
    "        return D @ self.G.adjacency_matrix()\n",
    "    \n",
    "    def nb_hitting_times(self, j):\n",
    "        \"\"\"\n",
    "        return list (np.array()) m_{i, j} where m is\n",
    "        NB hitting time from i to j for all i, with j fixed.\n",
    "\n",
    "        THIS IS ACTUALLY I THINK nb hitting time m_{e, j}. 2m vector of hitting times from\n",
    "        the edge e to the vertex j. (SEE Theorem 4.3 I believe.)\n",
    "\n",
    "        i.e. the columns of M_{ev}\n",
    "        \"\"\"\n",
    "        # Ask Adam what's going on here!\n",
    "        m = len(self.G.edges())\n",
    "        n = len(self.G)\n",
    "        P = self.nb_trans_mat()\n",
    "        Gedges = [(g[0], g[1]) for g in self.G.edges()]\n",
    "        Gedges.extend([(g[1], g[0]) for g in self.G.edges()])\n",
    "        entries_to_delete = []\n",
    "\n",
    "        # Find the rows/columns to delete\n",
    "        for i in range(len(Gedges)):\n",
    "            if Gedges[i][0] == j:\n",
    "                entries_to_delete.append(i)\n",
    "        P_new = np.delete(P, entries_to_delete, axis=0)\n",
    "        P_new = np.delete(P_new, entries_to_delete, axis=1)\n",
    "\n",
    "        # Keep a list of what was NOT deleted, to construct the full vector later\n",
    "        remaining_entries = [i for i in range(2 * m)]\n",
    "        for k in entries_to_delete:\n",
    "            remaining_entries.remove(k)\n",
    "\n",
    "        # Solve (I - P)x = 1\n",
    "        ones = np.array([1] * len(P_new))\n",
    "        tau_vec = np.linalg.solve(np.identity(len(P_new)) - P_new, ones)\n",
    "\n",
    "        # recreate full tau_vec, putting the 0's back in where necessary\n",
    "        full_tau = np.array([0.0] * (2 * m))\n",
    "        for i in range(len(remaining_entries)):\n",
    "            full_tau[remaining_entries[i]] = tau_vec[i]\n",
    "\n",
    "        return full_tau\n",
    "    \n",
    "    def M_ev_NBRW(self):\n",
    "        \"\"\"\n",
    "        return the matrix $M_{ev}$, that is, the hitting times (mean first passage times)\n",
    "        of the nonbacktracking random walk\n",
    "        \"\"\"\n",
    "        m = len(self.G.edges())\n",
    "        n = len(self.G)\n",
    "        Mev = np.zeros((2 * m, n))\n",
    "        for i in range(n):\n",
    "            Mvec = self.nb_hitting_times(i)\n",
    "            for j in range(2 * m):\n",
    "                Mev[j][i] = Mvec[j]\n",
    "\n",
    "        return Mev\n",
    "    \n",
    "    def P_hat(self):\n",
    "        \"\"\"\n",
    "        (D-I)^{-1}A.\n",
    "        I should probably not take an inverse since its so easy but whatever. change as needed\n",
    "        \"\"\"\n",
    "        return np.diag([1 / (d - 1) for d in G.degree()]) @ G.adjacency_matrix()\n",
    "\n",
    "\n",
    "    def P_mat(self):\n",
    "        \"\"\"\n",
    "        D^{-1}A\n",
    "        \"\"\"\n",
    "        return np.diag([1 / d for d in G.degree()]) @ G.adjacency_matrix()\n",
    "\n",
    "    \n",
    "    def Znb_mat(self):\n",
    "        \"\"\"\n",
    "        A guess at what Znb should be. Seems to work pretty close to what we want numerically. Matches exactly with\n",
    "        edge transitive graphs. Usually close with all other graphs we've tried.\n",
    "        \"\"\"\n",
    "        n = len(self.G)\n",
    "        m = len(self.G.edges())\n",
    "        Dinv = np.diag([1 / d for d in self.G.degree()])\n",
    "        T = self.T\n",
    "        S = self.S\n",
    "        Wnb = np.ones((2 * m, 2 * m)) / (2 * m)\n",
    "        Wv = np.outer(np.ones(n), [d / (2 * m) for d in self.G.degree()])\n",
    "\n",
    "        return np.eye(n) + Dinv @ T @ (np.linalg.inv(np.eye(2 * m) - self.nb_trans_mat() + Wnb)) @ S - Wv\n",
    "\n",
    "    \n",
    "    def Wnb_mat(self):\n",
    "        return np.ones((2 * self.m, 2 * self.m)) / (2 * self.m)\n",
    "    \n",
    "    \n",
    "    def Znb_e_mat(self):\n",
    "        I = identity_matrix(2*self.m)\n",
    "        return np.linalg.inv(I - self.Pnb + self.Wnb)\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "df5c3f5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "G = graphs.FolkmanGraph()\n",
    "G = NBRW(G)\n",
    "I = identity_matrix(2*G.m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "b421b76b",
   "metadata": {},
   "outputs": [],
   "source": [
    "first = G.Znb @ G.Dinv @ G.T @ (I - G.Pnb) @ G.Mev"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5bde710b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "SageMath 10.2",
   "language": "sage",
   "name": "sagemath"
  },
  "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.11.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
