Code refactor about color_map

mourisl · Sep 8, 2024 · bab5594 · bab5594
1 parent ddb7137
commit bab5594
Showing 1 changed file with 43 additions and 19 deletions.
diff --git a/vizMSA_dev.ipynb b/vizMSA_dev.ipynb
@@ -44,7 +44,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 259,
    "id": "f8508020-34de-4cd6-8988-8bc1d57990c9",
    "metadata": {},
    "outputs": [],
@@ -58,18 +58,41 @@
     "    if (end == None):\n",
     "        end = length - 1 \n",
     "    elif (end < 0):\n",
-    "        end = length + end "
+    "        end = length + end \n",
+    "    return start, end"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 273,
+   "id": "e2add348-467e-4ad3-8166-c70220d67b24",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def GetColorMap(msa, color_order = None, palette = None):\n",
+    "    color_map = {}\n",
+    "    if (color_order != None):\n",
+    "        for i,c in enumerate(color_order):\n",
+    "            color_map[c] = sns.color_palette(palette)[i]\n",
+    "            \n",
+    "    # The case some of the alphabet color is not specified in the order     \n",
+    "    for a in msa:\n",
+    "        for c in a:\n",
+    "            if c not in color_map:\n",
+    "                size = len(color_map)\n",
+    "                color_map[c] = sns.color_palette(palette)[size]\n",
+    "    return color_map"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 253,
+   "execution_count": 274,
    "id": "37d12119-b410-453a-8ce4-ee254fb4b957",
    "metadata": {},
    "outputs": [],
    "source": [
     "def DrawMSA(msa, seq_names = None, start = None, end = None,\n",
-    "            axlim = None, color_order = None, palette=None, ax=None, fig=None):\n",
+    "            axlim = None, color_map = None, color_order = None, palette=None, ax=None, fig=None):\n",
     "    # Get the canvas attributes.\n",
     "    ax = ax or plt.gca()\n",
     "    \n",
@@ -87,17 +110,8 @@
     "        fontsize = CalculateFontsize(ax.get_xlim(), ax.get_ylim(), ax, fig, height, length)\n",
     "    else:\n",
     "        fontsize = CalculateFontsize(axlim[0], axlim[1], ax, fig, height, length)\n",
-    "    chrColor = {}\n",
     "    \n",
-    "    if color_order == None:\n",
-    "        for a in msa:\n",
-    "            for c in a[start:end+1]:\n",
-    "                if c not in chrColor:\n",
-    "                    size = len(chrColor)\n",
-    "                    chrColor[c] = sns.color_palette(palette)[size]\n",
-    "    else:\n",
-    "        for i,c in enumerate(color_order):\n",
-    "            chrColor[c] = sns.color_palette(palette)[i]\n",
+    "    color_map = GetColorMap(msa, color_order, palette)\n",
     "    \n",
     "    lengthUnit = 1 / (end - start + 1)\n",
     "    heightUnit = 1 / height \n",
@@ -112,7 +126,7 @@
     "                transform=ax.transAxes if axlim == None else ax.transData)\n",
     "            ax.add_patch( patches.Rectangle(xy=(j * lengthUnit, i * heightUnit),\n",
     "                                           width = lengthUnit, height=heightUnit,\n",
-    "                                          facecolor=chrColor[c], linewidth=fontsize/50, edgecolor=\"white\",\n",
+    "                                          facecolor=color_map[c], linewidth=fontsize/50, edgecolor=\"white\",\n",
     "                                          transform=ax.transAxes if axlim == None else ax.transData))\n",
     "    if (axlim == None):\n",
     "        ax.set_xlim(-0.5, end - start + 1 - 0.5)\n",
@@ -132,7 +146,7 @@
     "            tickLabels.append(i if seq_names == None else seq_names[i])\n",
     "        ax.set_yticks(ticks, tickLabels)\n",
     "            \n",
-    "    return ax"
+    "    return ax, color_map"
    ]
   },
   {
@@ -143,6 +157,8 @@
    "outputs": [],
    "source": [
     "def GetConsensus(msa, start = None, end = None):\n",
+    "    start, end = GetStartEnd(ms, start, end)\n",
+    "    \n",
     "    "
    ]
   },
@@ -160,7 +176,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 255,
+   "execution_count": 277,
    "id": "9857702b-904f-4aa8-a330-d867bb91afd0",
    "metadata": {},
    "outputs": [
@@ -178,10 +194,18 @@
    "source": [
     "ax = plt.gca()\n",
     "ax.set(xlim=[-20,20], ylim=[-20, 20])\n",
-    "ax=DrawMSA(msa=[\"AC-G-T\", \"A-CGT-\"], seq_names = [\"r1\", \"r2\"],\n",
-    "           start = 1, end = 3, color_order=\"ACGT-\")"
+    "ax, color_map=DrawMSA(msa=[\"AC-G-T\", \"A-CGT-\"], seq_names = [\"r1\", \"r2\"],\n",
+    "           start = 1, end = 3, color_order=\"ACGT\")"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "92c38ded-6965-402d-ad2e-6908dd38dca6",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
   {
    "cell_type": "code",
    "execution_count": 252,