{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "cc3ea947",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from sklearn.linear_model import LinearRegression\n",
    "from sklearn.cluster import KMeans\n",
    "\n",
    "\n",
    "class CorrelationNetwork:\n",
    "    def fit(self, X):\n",
    "        self.p = X.shape[-1]\n",
    "        self.regs = [[LinearRegression() for _ in range(self.p)] for _ in range(self.p)]\n",
    "        for i in range(self.p):\n",
    "            for j in range(self.p):\n",
    "                self.regs[i][j].fit(X[:, j, np.newaxis], X[:, i, np.newaxis])\n",
    "        return self\n",
    "    \n",
    "    def predict(self, n):\n",
    "        betas = np.zeros((self.p, self.p))\n",
    "        for i in range(self.p):\n",
    "            for j in range(self.p):\n",
    "                betas[i, j] = self.regs[i][j].coef_.squeeze()\n",
    "        corrs = betas * betas.T\n",
    "        return np.tile(np.expand_dims(corrs, axis=0), (n, 1, 1))\n",
    "    \n",
    "    def mses(self, X):\n",
    "        mses = np.zeros(len(X))\n",
    "        for i in range(self.p):\n",
    "            for j in range(self.p):\n",
    "                residual = self.regs[i][j].predict(X[:, j, np.newaxis]) - X[:, i, np.newaxis]\n",
    "                residual = residual[:, 0]\n",
    "                mses += (residual ** 2) / self.p**2\n",
    "        return mses\n",
    "    \n",
    "\n",
    "class MarkovNetwork:\n",
    "    def fit(self, X):\n",
    "        self.p = X.shape[-1]\n",
    "        self.regs = [LinearRegression() for _ in range(self.p)]\n",
    "        for i in range(self.p):\n",
    "            mask = np.ones_like(X)\n",
    "            mask[:, i] = 0\n",
    "            self.regs[i].fit(X * mask, X[:, i, np.newaxis])\n",
    "        return self\n",
    "    \n",
    "    def predict(self, n):\n",
    "        betas = np.zeros((self.p, self.p))\n",
    "        for i in range(self.p):\n",
    "            betas[i] = self.regs[i].coef_.squeeze()\n",
    "            betas[i, i] = 0\n",
    "        precision = - np.sign(betas) * np.sqrt(np.abs(betas * betas.T))\n",
    "        return np.tile(np.expand_dims(precision, axis=0), (n, 1, 1))\n",
    "    \n",
    "    def mses(self, X):\n",
    "        mses = np.zeros(len(X))\n",
    "        for i in range(self.p):\n",
    "            mask = np.ones_like(X)\n",
    "            mask[:, i] = 0\n",
    "            residual = self.regs[i].predict(X * mask) - X[:, i, np.newaxis]\n",
    "            residual = residual[:, 0]\n",
    "            mses += (residual ** 2) / self.p\n",
    "        return mses\n",
    "\n",
    "\n",
    "class GroupedNetworks:\n",
    "    def __init__(self, model_class):\n",
    "        self.model_class = model_class\n",
    "    \n",
    "    def fit(self, X, labels):\n",
    "        self.models = {}\n",
    "        self.p = X.shape[-1]\n",
    "        for label in np.unique(labels):\n",
    "            label_idx = labels == label\n",
    "            X_label = X[label_idx]\n",
    "            model = self.model_class().fit(X_label)\n",
    "            self.models[label] = model\n",
    "        return self\n",
    "    \n",
    "    def predict(self, labels):\n",
    "        networks = np.zeros((len(labels), self.p, self.p))\n",
    "        for label in np.unique(labels):\n",
    "            label_idx = labels == label\n",
    "            networks[label_idx] = self.models[label].predict(label_idx.sum())\n",
    "        return networks\n",
    "    \n",
    "    def mses(self, X, labels):\n",
    "        mses = np.zeros(len(X))\n",
    "        for label in np.unique(labels):\n",
    "            label_idx = labels == label\n",
    "            X_label = X[label_idx]\n",
    "            mses[label_idx] = self.models[label].mses(X_label)\n",
    "        return mses\n",
    "\n",
    "\n",
    "class ClusteredCorrelation:\n",
    "    def __init__(self, K, clusterer=None):\n",
    "        self.K = K\n",
    "        if clusterer is None:\n",
    "            self.kmeans = KMeans(n_clusters=K)\n",
    "            self.prefit = False\n",
    "        else:\n",
    "            self.kmeans = clusterer\n",
    "            self.prefit = True\n",
    "        self.models = {k: PopulationCorrelation() for k in range(K)}\n",
    "    \n",
    "    def fit(self, C, X):\n",
    "        self.p = X.shape[-1]\n",
    "        if not self.prefit:\n",
    "            self.kmeans.fit(C)\n",
    "        labels = self.kmeans.predict(C)\n",
    "        for k in range(self.K):\n",
    "            k_idx = labels == k\n",
    "            X_k, C_k = X[k_idx], C[k_idx]\n",
    "            self.models[k].fit(C_k, X_k)\n",
    "        return self\n",
    "    \n",
    "    def predict(self, C):\n",
    "        labels = self.kmeans.predict(C)\n",
    "        corrs = np.zeros((len(C), self.p, self.p))\n",
    "        for label in np.unique(labels):\n",
    "            l_idx = labels == label\n",
    "            C_l = C[l_idx]\n",
    "            corrs[l_idx] = self.models[label].predict(C_l)\n",
    "        return corrs\n",
    "    \n",
    "    def mses(self, C, X):\n",
    "        labels = self.kmeans.predict(C)\n",
    "        mses = np.zeros(len(C))\n",
    "        for label in np.unique(labels):\n",
    "            l_idx = labels == label\n",
    "            l_count = sum(l_idx)\n",
    "            C_l, X_l = C[l_idx], X[l_idx]\n",
    "            mses[l_idx] = self.models[label].mses(C_l, X_l)\n",
    "        return mses\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "09171c99",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.3065607347398675\n",
      "0.2646665813126492\n",
      "0.2830745940131997\n",
      "0.02622722360951478\n"
     ]
    }
   ],
   "source": [
    "n, x_dim = 100, 20\n",
    "labels = np.random.randint(0, 5, (n,))\n",
    "X = np.random.uniform(-1, 1, (n, x_dim))\n",
    "\n",
    "corr = CorrelationNetwork().fit(X)\n",
    "corr.predict(n)\n",
    "print(corr.mses(X).mean())\n",
    "\n",
    "mark = MarkovNetwork().fit(X)\n",
    "mark.predict(n)\n",
    "print(mark.mses(X).mean())\n",
    "\n",
    "grouped_corr = GroupedNetworks(CorrelationNetwork).fit(X, labels)\n",
    "grouped_corr.predict(labels)\n",
    "print(grouped_corr.mses(X, labels).mean())\n",
    "\n",
    "grouped_mark = GroupedNetworks(MarkovNetwork).fit(X, labels)\n",
    "grouped_mark.predict(labels)\n",
    "print(grouped_mark.mses(X, labels).mean())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "82fa4c82",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 8048 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "plt.imshow(mark.predict(1)[0])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "3310d6d2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 6892 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.imshow(corr.predict(1)[0])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cabccd0c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:contextualized]",
   "language": "python",
   "name": "conda-env-contextualized-py"
  },
  "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.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
