marchantia_stress.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "marchantia_stress.ipynb",
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/tqiaowen/marchantia-stress/blob/main/marchantia_stress.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f4f_0f6taBn9"
      },
      "source": [
        "# Code to replicate analysis\n",
        "\n",
        "## Chapters\n",
        "### 1. Miscellaneous preparation steps\n",
        "1. Install non-default modules and upgrade modules\n",
        "1. Mount Google Drive\n",
        "1. Set paths (point to google drive folder to work in)\n",
        "1. Download files (initial set up, skip if continuing)\n",
        "1. Import modules, initialise paths\n",
        "\n",
        "### 2. Generating data for analysis\n",
        "1. Differential Gene Expression (DESeq2)\n",
        "1. Supp. Fig 4: Comparison of DEGs between two controls\n",
        "1. Diurnal Gene Expression (JTK_cycle)\n",
        "\n",
        "### 3. Analysis and plotting\n",
        "1. Figure 1 & Supp. Fig 1: Measurements and Student's t-test\n",
        "1. Figure 2: Interspecies comparison (Biological processes)\n",
        "1. Supp. Fig 5: Interspecies comparison (Gene families)\n",
        "1. Figure 3: Stress responsiveness\n",
        "1. Figure 4: Upset plot and summary of DEGs in Marchantia\n",
        "1. Figure 5, Supp. Fig. 6 & 7: Inter-stress (Marchantia only) comparison\n",
        "1. Figure 6: Diurnal gene expression\n",
        "1. Supp. Fig 2: QC of RNA-seq data\n",
        "1. Supp. Fig 3: Volcano plots (DESeq2)\n",
        "\n",
        "1. Supp. Fig 8: Overview of diurnal data\n",
        "\n",
        "### 4. Experimental\n",
        "1. Download RNA-seq data\n",
        "1. Mapping and generating expression matrix"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "63O3M3hPgwyv"
      },
      "source": [
        "# 1. Miscellaneous preparation steps"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uGwkg7hkhMYY"
      },
      "source": [
        "### 1.1 Install non-default modules and upgrade modules"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9THYinX_gpJ4"
      },
      "source": [
        "# install non-default colab modules\n",
        "# Restart runtime after installation and skip to next step\n",
        "!pip install upsetplot\n",
        "!pip install matplotlib --upgrade"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DutthK8QhAkJ"
      },
      "source": [
        "### 1.2 Mount Google Drive"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gLAcF40pZUM4"
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')\n",
        "!rm -rf /content/sample_data"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "50GdOAqr8D6D"
      },
      "source": [
        "#@title 1.3 Set path {display-mode: \"form\"}\n",
        "\n",
        "#@markdown Enter the path of the directory you want to work in.\n",
        "\n",
        "drive_path = '/content/gdrive/My Drive/' #@param {type: 'string'}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5-ZWhJDzo7e6"
      },
      "source": [
        "### 1.4 Download files (first time only)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SFhEwOSCGeF_"
      },
      "source": [
        "# Downloads necessary files to perform analyses [only need to be done once]\n",
        "# https://gist.github.com/iamtekeste/3cdfd0366ebfd2c0d805 download raw files directtly from Google Drive\n",
        "!wget --no-check-certificate -r \"https://drive.google.com/uc?id=1cbKgWbEWtstl_2_rb06tI_D-vseprPnT&export=download\" -O marchantia_stress.zip\n",
        "\n",
        "dir_path = drive_path + 'marchantia_stress/'\n",
        "dir_path_safe = dir_path.replace(' ', '\\ ')\n",
        "!mkdir $dir_path_safe\n",
        "!unzip marchantia_stress.zip -d $dir_path_safe"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u6GBaG6BQBR6"
      },
      "source": [
        "###1.5 Import modules, set paths"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "m3F4Mqz3P__H"
      },
      "source": [
        "# import modules\n",
        "import os\n",
        "import string\n",
        "%load_ext rpy2.ipython\n",
        "import pandas as pd\n",
        "import math\n",
        "from matplotlib_venn import venn2\n",
        "from matplotlib import pyplot as plt\n",
        "from ast import literal_eval\n",
        "import seaborn as sns\n",
        "from collections import Counter\n",
        "import random\n",
        "from statsmodels.stats.multitest import multipletests\n",
        "import numpy as np\n",
        "from scipy.cluster.hierarchy import dendrogram, linkage\n",
        "from scipy.spatial.distance import squareform\n",
        "from mpl_toolkits.axes_grid1.inset_locator import inset_axes\n",
        "\n",
        "dir_path = drive_path + 'marchantia_stress/'\n",
        "dir_path_safe = dir_path.replace(' ', '\\ ')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dqb8m7ATg19d"
      },
      "source": [
        "# 2. Generating data for analysis"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bzX8KZDXIMXV"
      },
      "source": [
        "### 2.1 Differential Gene Expression"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DRBwbtjMVsSF"
      },
      "source": [
        "# Making necessary directories for outputs\n",
        "mpo_path_safe = dir_path_safe + \"prep_files/mpo/deseq/\"\n",
        "osa_path_safe = dir_path_safe + \"prep_files/osa/deseq/\"\n",
        "if not os.path.exists(mpo_path_safe):\n",
        "    !mkdir -p $mpo_path\n",
        "    print(\"Directories made: \" + mpo_path.replace('\\\\', ''))\n",
        "if not os.path.exists(osa_path_safe):\n",
        "    !mkdir -p $osa_path\n",
        "    print(\"Directories made: \" + osa_path.replace('\\\\', ''))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "a_uzqH_PUYQ-"
      },
      "source": [
        "# Installing DESeq2\n",
        "%%R\n",
        "if (!requireNamespace(\"BiocManager\", quietly = TRUE))\n",
        "    install.packages(\"BiocManager\")\n",
        "\n",
        "BiocManager::install(\"DESeq2\", ask = FALSE)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rjetqXevenIe"
      },
      "source": [
        "# To pull python variables\n",
        "mpo_path = dir_path + \"prep_files/mpo/deseq/\"\n",
        "osa_path = dir_path + \"prep_files/osa/deseq/\"\n",
        "\n",
        "%R -i dir_path\n",
        "%Rget dir_path\n",
        "\n",
        "%R -i dir_path_safe\n",
        "%Rget dir_path_safe\n",
        "\n",
        "%R -i mpo_path\n",
        "%Rget mpo_path\n",
        "\n",
        "%R -i osa_path\n",
        "%Rget osa_path"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0HBke3tVIPk5"
      },
      "source": [
        "# adapted from DESeq2_stressonly_phase1n2.R\n",
        "%%R\n",
        "# DESeq2 (Marchantia)\n",
        "library('DESeq2')\n",
        "library('RColorBrewer')\n",
        "\n",
        "sink(paste0(mpo_path, \"phase1n2_sum.txt\"), type=\"output\")\n",
        "\n",
        "raw_counts <- read.table(file = paste0(dir_path, 'prep_files/all_stress_raw.tsv'), sep = '\\t', header = TRUE)\n",
        "raw_counts <- data.frame(raw_counts, row.names = 1)\n",
        "\n",
        "stresses <- c(\"controlH2\", \"controlD2\", \"H\", \"C\", \"M\", \"S\", \"L\", \"D\", \"N\",\n",
        "              \"HS\", \"HM\", \"HN\", \"CS\", \"CM\", \"CN\", \"SM\", \"ML\", \"NL\", \"MN\",\n",
        "              \"SD\", \"MD\", \"ND\", \"HD\", \"CD\", \"CL\", \"LS\", \"SN\")\n",
        "colData = read.csv(paste0(dir_path, 'summary_files/all_stress.txt'), sep = '\\t', row.names=1, header = FALSE)\n",
        "names(colData) <- c('condition')\n",
        "\n",
        "dds = DESeqDataSetFromMatrix(countData=raw_counts,\n",
        "                             colData=colData,\n",
        "                             design=~condition)\n",
        "dds = DESeq(dds)\n",
        "\n",
        "y = 2\n",
        "for (x in 1:2){\n",
        "  for (i in y:length(stresses)){\n",
        "    if (stresses[i] != stresses[x]){\n",
        "      res = results(dds, contrast=c(\"condition\", stresses[i], stresses[x]))\n",
        "      res = res[order(res$pvalue),]\n",
        "      resSig = subset(res, res$padj < 0.05 & abs(res$log2FoldChange) > 1)\n",
        "      resSig = resSig[ order(resSig$padj), ]\n",
        "      print(paste(stresses[i], 'vs', stresses[x]))\n",
        "      summary(res)\n",
        "      summary(resSig)\n",
        "      write.table(as.data.frame(res), file=paste(mpo_path, stresses[i], stresses[x], '_res.tsv', sep = ''),\n",
        "                  quote=FALSE, sep='\\t', col.names = NA)\n",
        "      write.table(as.data.frame(resSig), file=paste(mpo_path, stresses[i], stresses[x], '_resSig.tsv', sep = ''),\n",
        "                  quote=FALSE, sep='\\t', col.names = NA)\n",
        "    }\n",
        "  }\n",
        "  y = y + 1\n",
        "}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yW_QBfG4tW0N"
      },
      "source": [
        "%%R\n",
        "# DESeq2 (Rice) adapted from DESeq2_Osa.R\n",
        "\n",
        "library('DESeq2')\n",
        "library('RColorBrewer')\n",
        "\n",
        "raw_counts <- read.table(file = paste0(dir_path, 'prep_files/expmat_Osa_raw.tsv'), sep = '\\t', header = TRUE)\n",
        "raw_counts <- data.frame(raw_counts, row.names = 1)\n",
        "\n",
        "annotations <- c(\"1913_salt\", \"1913_control\",\n",
        "              \"5941_cold\", \"5941_control\",\n",
        "              \"ERP003982_salt\", \"ERP003982_control\",\n",
        "              \"GSE57950_drought\", \"GSE57950_control\")\n",
        "\n",
        "md1913 = c(\"1913_salt\",\n",
        "           \"1913_control\",\n",
        "           \"1913_control\",\n",
        "           \"1913_salt\",\n",
        "           \"1913_salt\",\n",
        "           \"1913_control\")\n",
        "md5941 = c(rep(c(\"5941_control\"), 3),\n",
        "           rep(c(\"5941_cold\"), 3),\n",
        "           rep(c(\"5941_control\"), 3),\n",
        "           rep(c(\"5941_cold\"), 3))\n",
        "md3982 = c(\"ERP003982_salt\",\n",
        "           \"ERP003982_control\",\n",
        "           rep(c(\"ERP003982_salt\"), 2),\n",
        "           rep(c(\"ERP003982_control\"), 2))\n",
        "md57950 = c(rep(c(\"GSE57950_control\"), 6),\n",
        "             rep(c(\"GSE57950_drought\"), 6))\n",
        "\n",
        "mdlist <- list(md1913, md5941, md3982, md57950)\n",
        "\n",
        "for (i in seq(1,length(annotations), by = 2)){\n",
        "  df = raw_counts[, grep(strsplit(annotations[i], \"_\")[[1]][1], names(raw_counts))]\n",
        "  sampleMetaData <- data.frame(condition = mdlist[[i - (i-1)/2]])\n",
        "  dds = DESeqDataSetFromMatrix(countData=df,\n",
        "                               colData=sampleMetaData,\n",
        "                               design=~condition)\n",
        "  dds = DESeq(dds)\n",
        "  res = results(dds, contrast=c(\"condition\", annotations[i], annotations[i+1]))\n",
        "  res = res[order(res$pvalue),]\n",
        "  resSig = subset(res, res$padj < 0.05 & abs(res$log2FoldChange) > 1)\n",
        "  resSig = resSig[ order(resSig$padj), ]\n",
        "  print(paste(annotations[i], 'vs', annotations[i+1]))\n",
        "  summary(res)\n",
        "  summary(resSig)\n",
        "  write.table(as.data.frame(res), file=paste(osa_path, annotations[i], strsplit(annotations[i+1], \"_\")[[1]][2], '_res.tsv', sep = ''),\n",
        "              quote=FALSE, sep='\\t', col.names = NA)\n",
        "  write.table(as.data.frame(resSig), file=paste(osa_path, annotations[i], strsplit(annotations[i+1], \"_\")[[1]][2], '_resSig.tsv', sep = ''),\n",
        "              quote=FALSE, sep='\\t', col.names = NA)\n",
        "}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OZ8K1xcl7dfZ"
      },
      "source": [
        "### 2.2 Supp. Fig 4: Comparison of DEGs between two controls"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-Bz6wm_ABLin"
      },
      "source": [
        "wdir = dir_path + 'prep_files/mpo/deseq/'\n",
        "deseqouts = [x for x in os.listdir(wdir) if \"resSig.tsv\" in x]\n",
        "deseqouts.remove('controlD2controlH2_resSig.tsv')\n",
        "\n",
        "# create subplots\n",
        "xlen = 4\n",
        "ylen = math.ceil(len(deseqouts)/8)\n",
        "figw = xlen * 4\n",
        "figh = ylen * 2.5\n",
        "\n",
        "stress_list = ['H', 'C', 'HM', 'CM', 'M', 'CL', 'ML', 'L', 'HS', 'CS', 'SM', 'LS', 'S',\n",
        "\t\t\t   'HN', 'CN', 'MN', 'NL', 'SN', 'N', 'HD', 'CD', 'MD', 'SD', 'ND', 'D']\n",
        "f_axes = string.ascii_uppercase[:len(stress_list)]\n",
        "axd = plt.figure(constrained_layout=True,\n",
        "\t\t\t\t figsize=(figw, figh)).subplot_mosaic(\n",
        "\t\t\t\t\t \"\"\"\n",
        "\t\t\t\t\t A......\n",
        "\t\t\t\t\t .B.....\n",
        "\t\t\t\t\t CDE....\n",
        "\t\t\t\t\t .FGH...\n",
        "\t\t\t\t\t IJKLM..\n",
        "\t\t\t\t\t NOPQRS.\n",
        "\t\t\t\t\t TUV.WXY\n",
        "\t\t\t\t\t \"\"\",\n",
        "\t\t\t\t\t gridspec_kw = {'hspace' : 0.3}\n",
        "\t\t\t\t\t )\n",
        "\n",
        "counter = 0\n",
        "for i in stress_list:\n",
        "\tfiles = [x for x in deseqouts if x.startswith(i+'control')]\n",
        "\tfiles.sort()\n",
        "\tfileD2 = pd.read_csv(wdir + files[0],\n",
        "\t\t\t\t\t  sep = \"\\t\", header = 0, index_col = 0)\n",
        "\tfileH2 = pd.read_csv(wdir + files[1],\n",
        "\t\t\t\t\t  sep = \"\\t\", header = 0, index_col = 0)\n",
        "\tD2index, H2index = set(fileD2.index.tolist()), set(fileH2.index.tolist())\n",
        "\tstatus = []\n",
        "\t\n",
        "\t# Create sets\n",
        "\tD2only = D2index - H2index\n",
        "\tH2only = H2index - D2index\n",
        "\tD2H2 = D2index & H2index\n",
        "\t\n",
        "\t# Subsets of df for D2only, H2only and D2H2\n",
        "\tD2onlydf = fileD2[fileD2.index.isin(list(D2only))]\n",
        "\tH2onlydf = fileH2[fileH2.index.isin(list(H2only))]\n",
        "\tD2H2df = fileD2.append(fileH2)\n",
        "\tD2H2df = D2H2df[D2H2df.index.isin(list(D2H2))]\n",
        "\tD2H2df.sort_index(inplace=True)\n",
        "\t\n",
        "\t# Create list of lists of all differentially expressed genes with corresponding status\n",
        "\tstress = \"Mpo_\" + files[0].split(\"controlD2\")[0]\n",
        "\tfor j in D2only:\n",
        "\t\tstatus.append([j, stress, 'D2', str(fileD2.loc[j, \"log2FoldChange\"]), str(fileD2.loc[j, \"padj\"]), \"N/A\", \"N/A\"])\n",
        "\tfor k in H2only:\n",
        "\t\tstatus.append([k, stress, 'H2', \"N/A\", \"N/A\", str(fileH2.loc[k, \"log2FoldChange\"]), str(fileH2.loc[k, \"padj\"])])\n",
        "\tfor m in D2H2:\n",
        "\t\tstatus.append([m, stress, 'D2H2', str(fileD2.loc[m, \"log2FoldChange\"]), str(fileD2.loc[m, \"padj\"]), str(fileH2.loc[m, \"log2FoldChange\"]), str(fileH2.loc[m, \"padj\"])])\n",
        "\tstatus.sort()\n",
        "\t\n",
        "\t# Output file with genes and status: D2, H2 or D2H2\n",
        "\twith open(wdir + \"sets/results/\" + stress + \".txt\", \"w+\") as filo:\n",
        "\t\tfilo.write(\"gene\\tstress\\tstatus\\tL2FC_D2\\tpadj_D2\\tL2FC_H2\\tpadj_H2\\n\")\n",
        "\t\tfor n in status:\n",
        "\t\t\tfilo.write(\"\\t\".join(n) + \"\\n\")\n",
        "\t\t\t\n",
        "\t# Venn diagram\n",
        "\tsp_ax = axd[f_axes[stress_list.index(stress.split('_')[1])]]\n",
        "\tvenn2(subsets=(len(D2only), len(H2only), len(D2H2)),\n",
        "\t   set_labels = ('D2', 'H2'),\n",
        "\t   ax = sp_ax)\n",
        "\tsp_ax.set_title(stress.split('_')[1], size=14)\n",
        "\tplt.savefig(dir_path + \"figures/\" + \"suppfig4.png\", dpi = 600)\n",
        "\t\n",
        "\t# increase counter\n",
        "\tcounter += 1"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SFx5kliR9PQE"
      },
      "source": [
        "# Sort genes according to whether they are the same in both controls\n",
        "# adapted from compile_sigGenes_phase1n2.py\n",
        "deseqdir = dir_path + 'prep_files/mpo/deseq/'\n",
        "setdir = deseqdir + 'sets/results/'\n",
        "setdir_safe = setdir.replace(' ', '\\ ')\n",
        "if not os.path.exists(setdir):\n",
        "\t!mkdir -p $setdir_safe\n",
        "\n",
        "merfile = open(dir_path + 'mercator/MpoProt.results.txt', 'r')\n",
        "ofile = open(deseqdir + 'resSig_compiled.txt', 'w+')\n",
        "efile = open(deseqdir + 'resSig_failed.txt', 'w+')\n",
        "setfiles = [x for x in os.listdir(setdir) if '.txt' in x]\n",
        "\n",
        "ofile.write(\"\\t\".join(['gene', 'stress', 'L2FC_D2', 'L2FC_H2', 'annotation']) + \"\\n\")\n",
        "efile.write(\"\\t\".join(['gene', 'stress', 'L2FC_D2', 'L2FC_H2', 'annotation']) + \"\\n\")\n",
        "\n",
        "def get_anno(gene):\n",
        "\tgene = gene.lower()\n",
        "\treturn meranno[gene]\n",
        "\n",
        "def up_down(val):\n",
        "\tif val < 0:\n",
        "\t\tstat = \"DOWN\"\n",
        "\telif val > 0:\n",
        "\t\tstat = \"UP\"\n",
        "\telif math.isnan(val):\n",
        "\t\tstat = \"NaN\"\n",
        "\treturn stat\n",
        "\n",
        "meranno = {}\n",
        "\n",
        "for line in merfile:\n",
        "\tif len(line.rstrip().split(\"\\t\")) == 5:\n",
        "\t\tbincode, name, identifier, desc, ptype = line.rstrip().replace(\"'\", \"\").split(\"\\t\")\n",
        "\t\tif identifier not in meranno:\n",
        "\t\t\tmeranno[identifier] = [[bincode, desc]]\n",
        "\t\telse:\n",
        "\t\t\tmeranno[identifier].append([bincode, desc])\n",
        "\n",
        "for i in setfiles:\n",
        "\tcontent = pd.read_csv(setdir + i, sep = \"\\t\", header = 0)\n",
        "\tsigGenes = content[content['status'] == \"D2H2\"]\n",
        "\tsigGenes['annotation'] = sigGenes['gene'].apply(get_anno)\n",
        "\tsigGenes['L2FC_D2'] = sigGenes['L2FC_D2'].apply(up_down)\n",
        "\tsigGenes['L2FC_H2'] = sigGenes['L2FC_H2'].apply(up_down)\n",
        "\tsigGenes = sigGenes.drop(columns = ['status', 'padj_D2', 'padj_H2'])\n",
        "\tfor index, row in sigGenes.iterrows():\n",
        "\t\tif row['L2FC_D2'] != row['L2FC_H2']:\n",
        "\t\t\tefile.write(\"\\t\".join([str(z) for z in row]) + \"\\n\")\n",
        "\t\telse:\n",
        "\t\t\tofile.write(\"\\t\".join([str(z) for z in row]) + \"\\n\")\n",
        "\n",
        "ofile.close()\n",
        "efile.close()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jQ3r_Ze15mZf"
      },
      "source": [
        "### 2.3 Diurnal gene expression"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KtgutT8H3MQY"
      },
      "source": [
        "jtk_dir = dir_path_safe + \"JTK/\"\n",
        "%cd $jtk_dir"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "k0usa5wb5pxj"
      },
      "source": [
        "%%R\n",
        "# https://towardsdatascience.com/how-to-install-packages-in-r-google-colab-423e8928cd2e\n",
        "#system(paste(\"cd\", paste0(dir_path_safe, \"JTK/\")))\n",
        "source(\"JTK_CYCLEv3.1.R\")\n",
        "\n",
        "project <- \"Mpo_JTK\"\n",
        "\n",
        "options(stringsAsFactors=FALSE)\n",
        "annot <- read.delim(\"annot_diur.txt\")\n",
        "data <- read.delim(\"expmat_diur.txt\")\n",
        "\n",
        "rownames(data) <- data[,1]\n",
        "data <- data[,-1]\n",
        "jtkdist(6, 3)       # 6 total time points, 3 replicates per time point\n",
        "\n",
        "periods <- 6:6       # looking for rhythms between 0-23 hours (i.e. between 1 and 6 time points per cycle).\n",
        "jtk.init(periods,4)  # 4 is the number of hours between time points\n",
        "\n",
        "cat(\"JTK analysis started on\",date(),\"\\n\")\n",
        "flush.console()\n",
        "\n",
        "st <- system.time({\n",
        "  res <- apply(data,1,function(z) {\n",
        "    jtkx(z)\n",
        "    c(JTK.ADJP,JTK.PERIOD,JTK.LAG,JTK.AMP)\n",
        "  })\n",
        "  res <- as.data.frame(t(res))\n",
        "  bhq <- p.adjust(unlist(res[,1]),\"BH\")\n",
        "  res <- cbind(bhq,res)\n",
        "  colnames(res) <- c(\"BH.Q\",\"ADJ.P\",\"PER\",\"LAG\",\"AMP\")\n",
        "  results <- cbind(annot,res,data)\n",
        "  results <- results[order(res$ADJ.P,-res$AMP),]\n",
        "})\n",
        "print(st)\n",
        "\n",
        "save(results,file=paste(\"JTK\",project,\"rda\",sep=\".\"))\n",
        "write.table(results,file=paste(\"JTK\",project,\"txt\",sep=\".\"),row.names=F,col.names=T,quote=F,sep=\"\\t\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-vX6y-Do3yOB"
      },
      "source": [
        "%cd $dir_path_safe"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AoC7_qrA4iOf"
      },
      "source": [
        "# adapted from clean_Mpo.py\n",
        "# To prepare and format Mpo JTK results to Camilla supp standard.\n",
        "# NR -- Not rhythmic genes ADJ.P <0.05\n",
        "# NE -- row[1:].max() > 1; no expression of TPM > 1 across all timepoints and replicates\n",
        "\n",
        "expanno = dir_path + \"summary_files/diurnal_exp.txt\"\n",
        "\n",
        "# label conversion to experiment annotation\n",
        "annodict = {}\n",
        "with open(expanno, \"r\") as expannof:\n",
        "\tcontent = expannof.readlines()\n",
        "\tfor line in content:\n",
        "\t\tlabel, actual = line.strip().split(\"\\t\")\n",
        "\t\tannodict[label] = actual + '_' + label.split('_')[1]\n",
        "\n",
        "diurlabels = [\"gene\"]\n",
        "diurlabels.extend(list(annodict.keys()))\n",
        "\n",
        "# select only diurnal experiments\n",
        "diurexpmat = dir_path + 'prep_files/diurnal_exp.tsv'\n",
        "diuronly = pd.read_csv(diurexpmat, sep='\\t')\n",
        "mpogenes = diuronly.gene.to_list()\n",
        "diuronly.set_index(\"gene\", inplace=True)\n",
        "diuronly.columns = [annodict[x] for x in diuronly.columns.to_list()]\n",
        "\n",
        "# prepping annotation file for JTK_Cycle/supp.\n",
        "meranno = {}\n",
        "merp = dir_path + 'mercator/MpoProt.results.txt'\n",
        "merfile = open(merp, 'r')\n",
        "merfile.readline()\n",
        "for line in merfile:\n",
        "\tif len(line.rstrip().split(\"\\t\")) == 5:\n",
        "\t\tbincode, name, identifier, desc, ptype = line.rstrip().replace(\"'\", \"\").split(\"\\t\")\n",
        "\t\tif identifier not in meranno:\n",
        "\t\t\tmeranno[identifier] = name\n",
        "\n",
        "# JTK output with expmat_diur\n",
        "JTKout = pd.read_csv(dir_path + \"JTK/JTK.Mpo_JTK.txt\", sep = \"\\t\")\n",
        "JTKgenes = JTKout.Probe.to_list()\n",
        "NEgenes = [x for x in mpogenes if x not in JTKgenes]\n",
        "\n",
        "# for supp\n",
        "colused = list(JTKout.columns)\n",
        "colunwanted = ['BH.Q', 'PER','AMP']\n",
        "for i in colunwanted:\n",
        "\tcolused.remove(i)\n",
        "\n",
        "# to format it to similar format of Camilla's supp material\n",
        "forsupp = JTKout[colused]\n",
        "# defnitions to change not significantly rthymic genes to NR instead of default output values\n",
        "def NRcheck(num):\n",
        "\tif num >= 0.05:\n",
        "\t\treturn \"NR\"\n",
        "\telse:\n",
        "\t\treturn \"{:.2E}\".format(num)\n",
        "forsupp['ADJ.P'] = forsupp['ADJ.P'].apply(lambda x: NRcheck(x))\n",
        "\n",
        "def phaseCheck(adjval, lagval):\n",
        "\tif adjval == \"NR\":\n",
        "\t\tnewval = \"NR\"\n",
        "\telse:\n",
        "\t\tnewval = lagval + 2\n",
        "\t\tif newval >= 24:\n",
        "\t\t\tnewval = newval - 24\n",
        "\treturn newval\n",
        "\n",
        "forsupp[\"LAG\"] = forsupp.apply(lambda row: phaseCheck(row[\"ADJ.P\"], row[\"LAG\"]), axis = 1)\n",
        "\n",
        "# to format genes that are not expressed (NE) and excluded in JTK analysis to supp file\n",
        "NEcollect = {}\n",
        "for j in NEgenes:\n",
        "\tNEcollect[j] = [meranno[j.lower()], \"NE\", \"NE\"] + diuronly.loc[j, colused[4:]].to_list()\n",
        "\n",
        "NEdf = pd.DataFrame(NEcollect, index = colused[1:])\n",
        "NEdf = NEdf.transpose()\n",
        "NEdf.reset_index(inplace=True)\n",
        "NEdf.columns = colused\n",
        "# combine formatted JTK output and NE genes\n",
        "combined = forsupp.append(NEdf, ignore_index = True)\n",
        "combined.sort_values(\"Probe\", inplace=True, ignore_index = True)\n",
        "# write to directory and ready for use (for analysis)\n",
        "cleaned = dir_path + \"diurnal/\"\n",
        "combined.to_csv(cleaned + \"Mpo_supp.txt\", index = False, sep = \"\\t\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V9HG7DByg6Pg"
      },
      "source": [
        "# 3. Analysis and plotting"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4qA5pIrH7VZW"
      },
      "source": [
        "### Figure 1 & Supp. Fig 1: Measurements and Student's t-test"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DXpM_NKa743D"
      },
      "source": [
        "# adated from measurements_forsupp.py\n",
        "wdir = dir_path + 'prep_files/'\n",
        "odir = dir_path + 'figures/'\n",
        "odir_safe = dir_path_safe + 'figures/'\n",
        "if not os.path.exists(odir):\n",
        "    !mkdir $odir_safe\n",
        "infile = 'phase1n2_measurements_nooutliers.txt'\n",
        "\n",
        "measurements = pd.read_csv(wdir + infile, sep='\\t')\n",
        "# Single letter to full single stress description\n",
        "singled = {'C':'Cold',\n",
        "\t\t   'H':'Heat',\n",
        "\t\t   'S':'Salt',\n",
        "\t\t   'M':'Mannitol',\n",
        "\t\t   'L':'Light',\n",
        "\t\t   'D':'Dark',\n",
        "\t\t   'N':'Nitrogen'\n",
        "\t\t   }\n",
        "\n",
        "# For plotting all controls\n",
        "areatype = ['Parea', 'Earea']\n",
        "titletype = ['15', '21']\n",
        "for i in range(0,2):\n",
        "\tarea = areatype[i]\n",
        "\ttitle = titletype[i]\n",
        "\tcontrol_m = measurements[measurements.Stress == 'Control'][[\"Batch\", area]].groupby('Batch', sort = False).mean()\n",
        "\tcontrol_s = measurements[measurements.Stress == 'Control'][[\"Batch\", area]].groupby('Batch', sort = False).std()\n",
        "\tcontrol_m.plot.bar(yerr=[list(control_s[area]), list(control_s[area])[::-1]], legend=False, title='Control (Day '+ title + ')', capsize=4)\n",
        "\tplt.savefig(odir + 'Control_Day' + title + '.png', dpi = 600, bbox_inches='tight')\n",
        "\tplt.show()\n",
        "\n",
        "# df with only single stress measurements\n",
        "ss_meas = measurements[(measurements.Condition != 'None') & (measurements.Condition != 'mixed')]\n",
        "# df with only crossed stress measurements\n",
        "cs_meas = measurements[measurements.Condition == 'mixed']\n",
        "# df with only controls\n",
        "control_meas = measurements[measurements.Stress == 'Control']\n",
        "\n",
        "controltype = ['Stress']\n",
        "controltitle = ['_merged']\n",
        "\n",
        "xaxislabel = {'Heat': 'Temperature (\\u00B0C)',\n",
        "\t\t\t  'Cold': 'Temperature (\\u00B0C)',\n",
        "\t\t\t  'Mannitol': 'Mannitol (mM)',\n",
        "\t\t\t  'Salt': 'NaCl (mM)',\n",
        "\t\t\t  'Light': 'Light intensity (\\u03bcEm\\u207b\\u00b2s\\u207b\\u00b9)',\n",
        "\t\t\t  'Dark': 'Days',\n",
        "\t\t\t  'Nitrogen': 'KNO\\u2083 (%)'}\n",
        "\n",
        "# Supp. Fig. 1\n",
        "# t-test (control as b, following test, a)\n",
        "tout = open(wdir + 'ttest.txt', 'w+') \n",
        "from scipy import stats as st\n",
        "\n",
        "for ss in list(ss_meas.Stress.unique()):\n",
        "\tfor i, c in enumerate(controltype):\n",
        "\t\tcontrol_batches = ss_meas[ss_meas.Stress == ss].Batch.unique()\n",
        "\t\tcontrol_mean = control_meas[control_meas.Batch.isin(control_batches)].groupby(c, sort = False).mean()\n",
        "\t\tcontrol_std = control_meas[control_meas.Batch.isin(control_batches)].groupby(c, sort = False).std()\n",
        "\t\tstress_mean = ss_meas[ss_meas.Stress == ss].groupby('Condition', sort=False).mean()\n",
        "\t\tstress_std = ss_meas[ss_meas.Stress == ss].groupby('Condition', sort=False).std()\n",
        "\t\t\n",
        "\t\tif c == 'Stress': # t-test\n",
        "\t\t\tcontrol_df = control_meas[control_meas.Batch.isin(control_batches)][['Parea', 'Earea']]\n",
        "\t\t\tstress_conds = ss_meas[ss_meas.Stress == ss].Condition.unique()\n",
        "\t\t\tfor k, a in enumerate(areatype):\n",
        "\t\t\t\tfor scond in stress_conds:\n",
        "\t\t\t\t\tstress_df = ss_meas[(ss_meas.Stress == ss) & (ss_meas.Condition == scond)]\n",
        "\t\t\t\t\ttstat, pval = st.ttest_ind(stress_df[a], control_df[a])\n",
        "\t\t\t\t\ttout.write(('\\t').join(['Day '+ titletype[k], ss + '_' + scond, 'control_merged', str(tstat), str(pval)]) + \"\\n\")\n",
        "\t\t\n",
        "\t\tlabels = control_mean.index.to_list() + stress_mean.index.to_list()\n",
        "\t\tfor j, a in enumerate(areatype): # Day 15 or 21 area\n",
        "\t\t\tcoll_mean = list(control_mean[a]) + list(stress_mean[a])\n",
        "\t\t\tcoll_std = list(control_std[a]) + list(stress_std[a])\n",
        "\t\t\tplt.bar(labels, coll_mean, yerr = coll_std, capsize=4)\n",
        "\t\t\tplt.title(ss + ' (Day ' + titletype[j] + ')')\n",
        "\t\t\tplt.xlabel(xaxislabel[ss]) \n",
        "\t\t\tplt.ylabel('Area (mm\\u00b2)')\n",
        "\t\t\tplt.savefig(odir + ss + '_Day' + titletype[j] + controltitle[i] + '.png', dpi = 600, bbox_inches='tight')\n",
        "\t\t\tplt.show()\n",
        "\n",
        "# cross_stress plot\n",
        "for i, c in enumerate(controltype):\n",
        "\tcs_control_batches = cs_meas.Batch.unique()\n",
        "\tcs_control_mean = control_meas[control_meas.Batch.isin(cs_control_batches)].groupby(c, sort = False).mean()\n",
        "\tcs_control_std = control_meas[control_meas.Batch.isin(cs_control_batches)].groupby(c, sort = False).std()\n",
        "\tcs_stress_mean = cs_meas.groupby('Stress', sort=False).mean()\n",
        "\tcs_stress_std = cs_meas.groupby('Stress', sort=False).std()\n",
        "\tcs_labels = cs_control_mean.index.to_list() + cs_stress_mean.index.to_list()\n",
        "\t\n",
        "\tif c == 'Stress': #t-test\n",
        "\t\tcontrol_df = control_meas[control_meas.Batch.isin(cs_control_batches)][['Parea', 'Earea']]\n",
        "\t\tfor k, a in enumerate(areatype):\n",
        "\t\t\tfor ss in list(cs_meas.Stress.unique()):\n",
        "\t\t\t\tstress_df = cs_meas[(cs_meas.Stress == ss)]\n",
        "\t\t\t\ttstat, pval = st.ttest_ind(stress_df[a], control_df[a])\n",
        "\t\t\t\ttout.write(('\\t').join(['Day '+ titletype[k], ss + '_' + 'mixed', 'control_merged', str(tstat), str(pval)]) + \"\\n\")\n",
        "\t\t\t\t\t\n",
        "\tfor j, a in enumerate(areatype): # Day 15 or 21 area\n",
        "\t\t\tcoll_mean = list(cs_control_mean[a]) + list(cs_stress_mean[a])\n",
        "\t\t\tcoll_std = list(cs_control_std[a]) + list(cs_stress_std[a])\n",
        "\t\t\tplt.bar(cs_labels, coll_mean, yerr = coll_std, capsize=4)\n",
        "\t\t\tplt.title('Cross stress (Day ' + titletype[j] + ')')\n",
        "\t\t\tplt.xticks(rotation=90)\n",
        "\t\t\tplt.xlabel('Experiment') \n",
        "\t\t\tplt.ylabel('Area (mm\\u00b2)')\n",
        "\t\t\tplt.savefig(odir + 'Cross_stress_Day' + titletype[j] + controltitle[i] + '.png', dpi = 600, bbox_inches='tight')\n",
        "\t\t\tplt.show()\n",
        "\t\t\t\n",
        "# single stress reps and cross stress (control - merged)\n",
        "def ss_grab(stress, condition):\n",
        "\t\"\"\"\n",
        "\tSlice the relevant condition for \n",
        "\tParameters\n",
        "\t----------\n",
        "\tstress : string\n",
        "\t\tStress of interest.\n",
        "\tcondition : string\n",
        "\t\tCondition of interest.\n",
        "\n",
        "\tReturns\n",
        "\t-------\n",
        "\tsssub : dataframe\n",
        "\t\tdf of single stress.\n",
        "\n",
        "\t\"\"\"\n",
        "\tsssub = measurements[(measurements.Stress == stress) & (measurements.Condition == condition)]\n",
        "\treturn sssub\n",
        "\n",
        "srep_keys = [['Cold', '3'],\n",
        "\t\t\t ['Heat', '33'],\n",
        "\t\t\t ['Salt', '40'],\n",
        "\t\t\t ['Mannitol', '100'],\n",
        "\t\t\t ['Light', '435'],\n",
        "\t\t\t ['Dark', '3'],\n",
        "\t\t\t ['Nitrogen', '0']]\n",
        "srepdf = measurements[(measurements.Stress == 'Cold') & (measurements.Condition == '3')]\n",
        "\n",
        "for s, c in srep_keys[1:]:\n",
        "\tsrepdf = pd.concat([srepdf, ss_grab(s, c)])\n",
        "\t\n",
        "s_cs_meas = pd.concat([srepdf, cs_meas])\n",
        "\t\n",
        "s_cs_control_batches = list(cs_meas.Batch.unique()) + list(srepdf.Batch.unique())\n",
        "s_cs_control_mean = control_meas[control_meas.Batch.isin(s_cs_control_batches)].groupby('Stress', sort = False).mean()\n",
        "s_cs_control_std = control_meas[control_meas.Batch.isin(s_cs_control_batches)].groupby('Stress', sort = False).std()\n",
        "s_cs_stress_mean = s_cs_meas.groupby('Stress', sort=False).mean()\n",
        "s_cs_stress_std = s_cs_meas.groupby('Stress', sort=False).std()\n",
        "\n",
        "s_cs_meas_label = [x + ' (' + x[0] + ')' if len(x) > 2 else x for x in s_cs_stress_mean.index]\n",
        "s_cs_meas_label[s_cs_meas_label.index('Light (L)')] = 'High light (L)'\n",
        "s_cs_meas_label[s_cs_meas_label.index('Dark (D)')] = 'Darkness (D)'\n",
        "s_cs_labels = s_cs_control_mean.index.to_list() + s_cs_meas_label\n",
        "\n",
        "# Fig1\n",
        "#t-test\n",
        "control_df = control_meas[control_meas.Batch.isin(s_cs_control_batches)][['Parea', 'Earea']]\n",
        "## cross-stress\n",
        "for k, a in enumerate(areatype):\n",
        "\tfor ss in list(cs_meas.Stress.unique()):\n",
        "\t\tstress_df = cs_meas[(cs_meas.Stress == ss)]\n",
        "\t\tfor singleS in ss:\n",
        "\t\t\tsinglecontrol = srepdf[srepdf.Stress == singled[singleS]][a]\n",
        "\t\t\ttstat, pval = st.ttest_ind(stress_df[a], singlecontrol)\n",
        "\t\t\ttout.write(('\\t').join(['Day '+ titletype[k], ss + '_' + 'mixed', 'control_' + singled[singleS], str(tstat), str(pval)]) + \"\\n\")\n",
        "## single stress\n",
        "for k, a in enumerate(areatype):\n",
        "\tfor ss in list(srepdf.Stress.unique()):\n",
        "\t\tcond = srepdf[(srepdf.Stress == ss)].Condition.unique()[0]\n",
        "\t\tstress_df = srepdf[(srepdf.Stress == ss)][a]\n",
        "\t\ttstat, pval = st.ttest_ind(stress_df, control_df[a])\n",
        "\t\ttout.write(('\\t').join(['Day '+ titletype[k], ss + '_' + cond, 'control', str(tstat), str(pval)]) + \"\\n\")\n",
        "\t\t\n",
        "tout.close()\t\n",
        "# plotting\n",
        "colour_seq = ['tomato'] +  ['mediumseagreen']*7 + ['cornflowerblue']*20\n",
        "for j, a in enumerate(areatype): # Day 15 or 21 area\n",
        "\t\tcoll_mean = list(s_cs_control_mean[a]) + list(s_cs_stress_mean[a])\n",
        "\t\tcoll_std = list(s_cs_control_std[a]) + list(s_cs_stress_std[a])\n",
        "\t\tplt.bar(s_cs_labels, coll_mean, yerr = coll_std, capsize=4, color = colour_seq)\n",
        "\t\tplt.title('Area (Day ' + titletype[j] + ')')\n",
        "\t\tplt.xticks(rotation=90)\n",
        "\t\tplt.xlabel('Experiment') \n",
        "\t\tplt.ylabel('Area (mm\\u00b2)')\n",
        "\t\tplt.savefig(odir + 'fig1_Day' + titletype[j] + '.png', dpi = 600, bbox_inches='tight')\n",
        "\t\tplt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DXF0cNS57aAn"
      },
      "source": [
        "### Figure 2: Interspecies comparison (Biological processes)\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PcYVfVEpH0Jc"
      },
      "source": [
        "# adapted from cross_spe_mapman.py\n",
        "### FUNCTION ###\n",
        "def anno_split(row):\n",
        "\treturn int(row['annotation'][0][0].split(\".\")[0])\n",
        "def bin_count(row):\n",
        "\tcount = row['rel_count']\n",
        "\tif count >= 0.5:\n",
        "\t\treturn 0.5\n",
        "\telif count >= 0.25:\n",
        "\t\treturn 0.35\n",
        "\telif count > 0.0:\n",
        "\t\treturn 0.2\n",
        "\telse:\n",
        "\t\treturn 0\n",
        "def label_color(xlabel):\n",
        "\tif \"heat\" in xlabel:\n",
        "\t\treturn \"firebrick\"\n",
        "\telif \"cold\" in xlabel:\n",
        "\t\treturn \"steelblue\"\n",
        "\telif \"light\" in xlabel:\n",
        "\t\treturn \"darkorange\"\n",
        "\telif \"dark\" in xlabel:\n",
        "\t\treturn \"black\"\n",
        "\telif \"salt\" in xlabel:\n",
        "\t\treturn \"rebeccapurple\"\n",
        "\telif \"mannitol\" in xlabel:\n",
        "\t\treturn \"mediumvioletred\"\n",
        "\telif \"nitrogen\" in xlabel:\n",
        "\t\treturn \"forestgreen\"\n",
        "\telse:\n",
        "\t\treturn \"slategrey\"\n",
        "\n",
        "def species_color(xlabel):\n",
        "\tif \"Ath\" in xlabel:\n",
        "\t\treturn \"firebrick\"\n",
        "\telif \"Cpa\" in xlabel:\n",
        "\t\treturn \"steelblue\"\n",
        "\telif \"Cre\" in xlabel:\n",
        "\t\treturn \"darkorange\"\n",
        "\telif \"Osa\" in xlabel:\n",
        "\t\treturn \"rebeccapurple\"\n",
        "\telif \"Mpo\" in xlabel:\n",
        "\t\treturn \"forestgreen\"\n",
        "\n",
        "### PATHS ###\n",
        "wdir = dir_path + 'prep_files/'\n",
        "setres = wdir + 'Figure2_alldata_compiled_updated.txt'\n",
        "jdir = wdir + 'proteomes/'\n",
        "mdict = wdir + 'merdict.txt'\n",
        "\n",
        "### DICTIONARY OF MERCATOR BINS ###\n",
        "dicto = literal_eval(open(mdict, 'r').read())\n",
        "\n",
        "### LOAD GENE PER SPECIES AND COUNT\n",
        "# initialise species\n",
        "spedicto = {'ARATH' : 'Ath',\n",
        "\t\t\t'CHLRE' : 'Cre',\n",
        "\t\t\t'CYAPA' : 'Cpa',\n",
        "\t\t\t'MARPO' : 'Mpo',\n",
        "\t\t\t'ORYSA' : 'Osa'}\n",
        "species_list = list(spedicto.values())\n",
        "\t\t\n",
        "# initialise gene count in species [for % of DGEs]\n",
        "Gdicto = {}\n",
        "# {\"species\" : [\"gene1\", \"gene2\"...]}\n",
        "for pepfile in [x for x in os.listdir(jdir) if '.ini' not in x]:\n",
        "\twith open(jdir + pepfile, \"r\") as peppy:\n",
        "\t\tspecies, genes = pepfile.split('.fa')[0], []\n",
        "\t\tfor lini in peppy:\n",
        "\t\t\tif '>' in lini:\n",
        "\t\t\t\tgenes.append(lini.strip().split('>')[1])\n",
        "\t\tGdicto[spedicto[species]] = len(genes)\n",
        "\t\t\n",
        "# Get only Ath genes (for name conversion, mercator output)\n",
        "athdict = {}\n",
        "with open(jdir + \"ARATH.fa\", \"r\") as athgenes:\n",
        "\tfor lini in athgenes:\n",
        "\t\tif '>' in lini:\n",
        "\t\t\tgenename = lini.strip().split('>')[1]\n",
        "\t\t\tathdict[genename.lower()] = genename\n",
        "\t\t\t\n",
        "athdict2 = {} # for name conversion (DGE table)\n",
        "with open(jdir + \"ARATH.fa\", \"r\") as athgenes:\n",
        "\tfor lini in athgenes:\n",
        "\t\tif '>' in lini:\n",
        "\t\t\tgenename = lini.strip().split('>')[1]\n",
        "\t\t\tathdict2[genename.lower().capitalize()] = genename\n",
        "\n",
        "### LOAD SIGNIFICANTLY DIFFERENTIAL GENE TABLE ###\n",
        "sigtable = pd.read_csv(setres, sep='\\t', header=0, index_col=0)\n",
        "sigtable = sigtable.reset_index()\n",
        "sigtable[\"gene\"].replace(athdict2, inplace=True)\n",
        "sigtable = sigtable.set_index(\"gene\")\n",
        "\n",
        "### DICTIONARY OF MERCATOR ANNOTATION ###\n",
        "merdir = wdir + 'mercator_results/'\n",
        "merlist = [x for x in os.listdir(merdir) if '.results.txt' in x]\n",
        "# Read mercator annotations (list of lists) as lists instead of string\n",
        "sigtable['annotation'] = sigtable['annotation'].apply(literal_eval)\n",
        "\n",
        "meranno = {}\n",
        "map2anno = {}\n",
        "for i in merlist:\n",
        "\tsp = i.split(\"Prot\")[0]\n",
        "\tmerfile = open(merdir + i, 'r')\n",
        "\tmerfile.readline()\n",
        "\tfor line in merfile:\n",
        "\t\tlinecon = line.rstrip().replace(\"'\", \"\").split(\"\\t\")\n",
        "\t\tif len(linecon) == 5:\n",
        "\t\t\tbincode, name, identifier, desc, ptype = linecon\n",
        "\t\t\tif identifier not in meranno:\n",
        "\t\t\t\tmeranno[identifier] = [sp, [bincode.split('.')[0]]]\n",
        "\t\t\telse:\n",
        "\t\t\t\tmeranno[identifier][1].append(bincode.split('.')[0])\n",
        "\t\tif len(linecon[0].split('.')) == 2:\n",
        "\t\t\tmap2anno[linecon[0]] = linecon[1]\n",
        "\n",
        "merdf = pd.DataFrame.from_dict(meranno, orient = 'index', columns = ['species', 'code'])\n",
        "merdf = merdf.reset_index()\n",
        "merdf[\"index\"].replace(athdict, inplace=True)\n",
        "merdf = merdf.set_index(\"index\")\n",
        "\n",
        "all_s = list(set(sigtable.stress.to_list()))\n",
        "\n",
        "def bin_collate(updown):\n",
        "\t'''\n",
        "\tTo generate df of collated bins per stress\n",
        "\n",
        "\tParameters\n",
        "\t----------\n",
        "\tupdown : str\n",
        "\t\tchoice of whether to construct for upregulated or downregulated genes.\n",
        "\n",
        "\tReturns\n",
        "\t-------\n",
        "\tNone.\n",
        "\n",
        "\t'''\n",
        "\tdicto = {}\n",
        "\tfor stress in all_s:\n",
        "\t\tgenes = [x for x in sigtable[(sigtable.stress == stress) & (sigtable.L2FC_D2 == updown)].index.to_list()]\n",
        "\t\tbeans = [y[0].split('.')[0] for x in sigtable[(sigtable.stress == stress) & (sigtable.L2FC_D2 == updown)].annotation.to_list() for y in x]\n",
        "\t\tdicto[stress] = [genes, beans]\n",
        "\tdf_beans = pd.DataFrame.from_dict(dicto, orient='index', columns=['gene', 'bins'])\n",
        "\treturn df_beans\n",
        "\t\n",
        "def sig_df(df, sigcol, mapbins):\n",
        "\t\"\"\"\n",
        "\tCalculates and correct mapman bin enrichment p-value for all stresses\n",
        "\tReturns dataframe\n",
        "\n",
        "\tParameters\n",
        "\t----------\n",
        "\tdf : dataframe\n",
        "\t\tdf containing genes and corresponding mapman bins of DEGs.\n",
        "\tsigcol : str\n",
        "\t\tcolumn name to use for enrichment\n",
        "\tmapbins : list\n",
        "\t\tlist of mapman annotation/bins to use\n",
        "\n",
        "\tReturns\n",
        "\t-------\n",
        "\tdf_sig : dataframe\n",
        "\t\tdf summarising enrichment (corrected p-value) for each mapman bin (row)\n",
        "\t\tand each stress (column).\n",
        "\n",
        "\t\"\"\"\n",
        "\tsig_sum = {}\n",
        "\tfor s in all_s:\n",
        "\t\ts_count = Counter(df.loc[s][sigcol])\n",
        "\t\tvalid_bins = list(s_count.keys()) # bins found in stress\n",
        "\t\t# initialise count dicitonary\n",
        "\t\tsig_count = {}\n",
        "\t\tfor key in valid_bins:\n",
        "\t\t\tsig_count[key] = 1\n",
        "\t\t# initilaise values for mercator by species\n",
        "\t\tsp = s.split('_')[0]\n",
        "\t\t# random simulations\n",
        "\t\tsimno = 1000\n",
        "\t\tfor i in range(simno):\n",
        "\t\t\tshuffle = merdf[merdf.species == sp].code.to_list()\n",
        "\t\t\trandom.shuffle(shuffle)\n",
        "\t\t\tsub = shuffle[:len(df.loc[s].gene)]\n",
        "\t\t\tsub_count = Counter([y for x in sub for y in x])\n",
        "\t\t\tfor mapman in valid_bins:\n",
        "\t\t\t\tif sub_count[mapman] >= s_count[mapman]:\n",
        "\t\t\t\t\tsig_count[mapman] += 1\n",
        "\t\t# p-value calculation\n",
        "\t\tpval_coll = []\n",
        "\t\tfor mapman in valid_bins:\n",
        "\t\t\tpval = sig_count[mapman]/simno\n",
        "\t\t\t# correction for pval > 1\n",
        "\t\t\tif pval <= 1:\n",
        "\t\t\t\tpval_coll.append(pval)\n",
        "\t\t\telse:\n",
        "\t\t\t\tpval_coll.append(float(round(pval)))\n",
        "\t\t\n",
        "\t\t# BH correction for multiple testing\n",
        "\t\ty = multipletests(pvals=pval_coll, alpha=0.05, method=\"fdr_bh\")[1]\n",
        "\t\tall_bins_corr_pval = []\n",
        "\t\tfor mapman in mapbins:\n",
        "\t\t\tif mapman in valid_bins:\n",
        "\t\t\t\tall_bins_corr_pval.append(y[valid_bins.index(mapman)])\n",
        "\t\t\telse:\n",
        "\t\t\t\tall_bins_corr_pval.append(np.nan)\n",
        "\t\tsig_sum[s] = all_bins_corr_pval\n",
        "\n",
        "\tdf_sig = pd.DataFrame.from_dict(sig_sum, orient='index', columns=mapbins)\n",
        "\treturn df_sig\n",
        "\n",
        "def chunk(uval, dval):\n",
        "\tif math.isnan(uval) and math.isnan(dval):\n",
        "\t\t# not differentially regulated\n",
        "\t\tcat = 0\n",
        "\telif  uval >= 0.05 and (dval >= 0.05 or math.isnan(dval)):\n",
        "\t\t# not enriched\n",
        "\t\tcat = 0\n",
        "\telif  dval >= 0.05 and (uval >= 0.05 or math.isnan(uval)):\n",
        "\t\t# not enriched\n",
        "\t\tcat = 0\n",
        "\telif  uval < 0.05 and dval < 0.05:\n",
        "\t\t# differentially up and downregulated in bin\n",
        "\t\tcat = 2\n",
        "\telif dval < 0.05:\n",
        "\t\t# differentially downregualted\n",
        "\t\tcat = 1\n",
        "\telif uval < 0.05:\n",
        "\t\t# differentially upregulated\n",
        "\t\tcat = 3\n",
        "\treturn cat\n",
        "\n",
        "# initialisation for enrichment\n",
        "#mapbins = list(map2anno.keys()) #level 2\n",
        "mapbins = [str(x) for x in list(dicto.keys())] # level 1\n",
        "\n",
        "# segregate stress and associated genes and mapman bins into up and downregulated df respectively\n",
        "df_U = bin_collate('UP')\n",
        "df_D = bin_collate('DOWN')\n",
        "\n",
        "# df of significance values\n",
        "df_sig_U = sig_df(df_U, 'bins', mapbins)\n",
        "df_sig_D = sig_df(df_D, 'bins', mapbins)\n",
        "\n",
        "cat_dict = {}\n",
        "for mapman in list(df_sig_U.columns):\n",
        "\tcat_col = []\n",
        "\tfor stress in list(df_sig_U.index):\n",
        "\t\tuval, dval = df_sig_U.loc[stress, mapman], df_sig_D.loc[stress, mapman]\n",
        "\t\tcat_col.append(chunk(uval,dval))\n",
        "\tcat_dict[mapman] = cat_col\n",
        "\n",
        "df_combined_sig = pd.DataFrame.from_dict(cat_dict, orient='index', columns=list(df_sig_U.index))\n",
        "df_combined_sig = df_combined_sig.loc[df_combined_sig.max(axis=1) > 0,:] # remove bins w/o enrichment\n",
        "df_combined_sig = df_combined_sig.loc[(df_combined_sig > 0).sum(axis=1) >2,:] # select for at least 2 enrichment\n",
        "df_combined_sig = df_combined_sig.loc[:,df_combined_sig.max() > 0] # remove stresses w/o enrichment\n",
        "\t\n",
        "def ji_cal(a, b):\n",
        "\t# jaccard index calculation\n",
        "\treturn len(a&b) / len(a|b)\n",
        "\n",
        "def jdistprep(df, axis):\n",
        "\t'''\n",
        "\tConvert df to sets (for calculation of JD of X axis)\n",
        "\n",
        "\tParameters\n",
        "\t----------\n",
        "\tdf : dataframe\n",
        "\t\tdataframe of categorical variables to be converted to sets.\n",
        "\taxis : int \n",
        "\t\taxis to do sets on, 0 by column (default), 1 by row\n",
        "\n",
        "\tReturns\n",
        "\t-------\n",
        "\tdicto : dict\n",
        "\t\tdictionary containing list of column values.\n",
        "\n",
        "\t'''\n",
        "\tif axis == 1:\n",
        "\t\tdf = df.T\n",
        "\tdxkeys = df.columns.to_list()\n",
        "\tdykeys = df.index.to_list()\n",
        "\tdicto = {}\n",
        "\tfor col in dxkeys:\n",
        "\t\tdicto[col] = [dykeys[i] + '_' + str(x) for i, x in enumerate(df[col].to_list())]\n",
        "\treturn [dicto, dxkeys]\n",
        "\n",
        "def jdist(df, axis=0):\n",
        "\t'''\n",
        "\tConstruct jaccard distance square matrix\n",
        "\n",
        "\tParameters\n",
        "\t----------\n",
        "\tdf : df\n",
        "\t\tdataframe to be used for jiprep/ jdist calculation.\n",
        "\taxis : int\n",
        "\t\taxis to do sets on, 0 by column (default), 1 by row\n",
        "\n",
        "\tReturns\n",
        "\t-------\n",
        "\tlinkage_matrix : list\n",
        "\t\tcondensed jaccard distance matrix.\n",
        "\tjlist : list\n",
        "\t\tlist of list (jaccard distance square matrix)\n",
        "\tdicto : dict\n",
        "\t\tdictionary of list\n",
        "\n",
        "\t'''\n",
        "\tdicto, dxkeys = jdistprep(df, axis)\n",
        "\tjlist = []\n",
        "\tfor key in dxkeys:\n",
        "\t\tcol = []\n",
        "\t\tfor key2 in dxkeys:\n",
        "\t\t\tset1, set2 = dicto[key], dicto[key2]\n",
        "\t\t\t\n",
        "\t\t\tset1x = set([x for x in set1 if x.split('_')[1] != '0'])\n",
        "\t\t\tset2x = set([x for x in set2 if x.split('_')[1] != '0'])\n",
        "\t\t\tcol.append(1 - ji_cal(set1x, set2x))\n",
        "\t\tjlist.append(col)\n",
        "\tdists = squareform(jlist)\n",
        "\tlinkage_matrix = linkage(dists, \"single\")\n",
        "\treturn linkage_matrix, jlist, dicto\n",
        "\n",
        "def plot_dendro(linkage_matrix, ax, orient):\n",
        "\t'''\n",
        "\tPlots dendrogram into subplot\n",
        "\n",
        "\tParameters\n",
        "\t----------\n",
        "\tmat : list of lists\n",
        "\t\tContains the square matrix of jaccard distances.\n",
        "\tax : axes\n",
        "\t\taxis of subplot to plot to.\n",
        "\torient : str\n",
        "\t\torientation of dendrogram to be plotted.\n",
        "\n",
        "\tReturns\n",
        "\t-------\n",
        "\tNone.\n",
        "\n",
        "\t'''\n",
        "\t\n",
        "\tdendrogram(linkage_matrix, no_labels=True, ax=ax, orientation=orient, color_threshold=0, above_threshold_color='#000000')\n",
        "\n",
        "xmat, xlist, xdict = jdist(df_combined_sig)\n",
        "ymat, ylist , ydict = jdist(df_combined_sig, axis=1)\n",
        "\n",
        "yden = dendrogram(ymat, labels=df_combined_sig.index.to_list(), orientation='left')\n",
        "plt.show()\n",
        "xden = dendrogram(xmat, labels=df_combined_sig.columns.to_list(), orientation='top')\n",
        "plt.show()\n",
        "\n",
        "yorder = yden['ivl']\n",
        "xorder = xden['ivl']\n",
        "df_sig_reordered = df_combined_sig[xorder]\n",
        "df_sig_reordered = df_sig_reordered.reindex(yorder[::-1])\n",
        "\n",
        "\"\"\"\n",
        "Custom plot\n",
        "\"\"\"\n",
        "\n",
        "# Create plot with subplot\n",
        "fig, ax = plt.subplots(6,2, constrained_layout=True,\n",
        "\t\t\t\t\t   figsize=(16.3, 18), # (width, height)\n",
        "\t\t\t\t\t   gridspec_kw={'width_ratios': [1, 8.3],\n",
        "\t\t\t\t\t 'height_ratios': [0.5,0.5,0.5,0.5,1,5.3]})\n",
        "plt.rcParams['font.size'] = '16'\n",
        "ax_1, ax_2, ax_3, ax_4, ax_5, ax_6, ax_7, ax_8, ax1, ax2, ax3, ax4 = ax.flatten()\n",
        "\n",
        "for ax in [ax_1, ax_2, ax_3, ax_4, ax_5, ax_6, ax_7, ax_8, ax1, ax2, ax3, ax4]:\n",
        "\tax.tick_params(axis='both', which='major', labelsize=16)\n",
        "\n",
        "ax_1.axis('off') # empty\n",
        "#ax_2.axis('off') # DGE % of genes\n",
        "ax_3.axis('off') # empty\n",
        "#ax_4.axis('off') # DGE % of TFs\n",
        "ax_5.axis('off') # empty\n",
        "#ax_6.axis('off') #  DGE % of kinases\n",
        "ax_7.axis('off') # empty\n",
        "#ax_8.axis('off') # %DGE up/down reg\n",
        "ax1.axis('off') # cbar\n",
        "ax2.axis('off') # dendrogram row\n",
        "ax3.axis('off') # dendrogram column\n",
        "\n",
        "# ax_2 DGE % of genes\n",
        "### STATISTICS OF DGEs ###\n",
        "DGEperdict = {}\n",
        "for sx in sigtable.stress.unique():\n",
        "\tDGEperdict[sx] = (len(sigtable[sigtable.stress == sx])/Gdicto[sx.split(\"_\")[0]])*100\n",
        "DGEper = pd.DataFrame(DGEperdict, index = [\"DGEs\"])\n",
        "DGEper = DGEper.transpose()\n",
        "DGEper[\"NotDGE\"] = DGEper.apply(lambda row: 100 - row, axis=0)\n",
        "DGEper = DGEper.reindex(xorder)\n",
        "\n",
        "DGEper.plot.bar(stacked = True,\n",
        "\t\t\t\tcolor = {\"DGEs\":\"firebrick\", \"NotDGE\":\"darkgrey\"},\n",
        "\t\t\t\tedgecolor = \"black\",\n",
        "\t\t\t\tylim = [0, 50],\n",
        "\t\t\t\tax = ax_2)\n",
        "handles, labels = ax_2.get_legend_handles_labels()\n",
        "ax_2.legend(handles=handles[:-1], labels=labels[:-1],\n",
        "          loc='center left', bbox_to_anchor=(1, 0.5))\n",
        "ax_2.set_ylabel(\"% genes\", rotation = 90, fontsize=16)\n",
        "ax_2.yaxis.set_label_coords(-0.06,0.36)\n",
        "ax_2.axes.get_xaxis().set_visible(False)\n",
        "\n",
        "# ax_4 DGE % of DGEs that are TFs\n",
        "tfdir = dir_path + 'tf_kinases/'\n",
        "tfpaths = [x for x in os.listdir(tfdir) if \".ini\" not in x and \".TF.\" in x]\n",
        "tfdict = {}\n",
        "\n",
        "for file in tfpaths:\n",
        "\ttempspe = spedicto[file.split(\".\")[0]]\n",
        "\tcontent = open(tfdir + file, \"r\")\n",
        "\tfor line in content:\n",
        "\t\tgene, anno = line.strip().split(\"\\t\")\n",
        "\t\tif anno!= \"NoFunction\":\n",
        "\t\t\ttfdict[gene] = anno\n",
        "\t\n",
        "sigmod = sigtable.reset_index()\n",
        "tfdf = pd.DataFrame(columns = list(sigmod.columns))\n",
        "for sx in sigtable.stress.unique():\n",
        "\ttfsubset = sigmod[(sigmod.stress == sx) & (sigmod.apply(lambda row: row[\"gene\"] in tfdict, axis = 1))]\n",
        "\ttfdf = tfdf.append(tfsubset, ignore_index = True)\n",
        "\t\n",
        "TFperdict = {}\n",
        "for sx in sigtable.stress.unique():\n",
        "\tTFperdict[sx] = (len(tfdf[tfdf.stress == sx])/len(sigmod[sigmod.stress == sx]))*100\n",
        "TFper = pd.DataFrame(TFperdict, index = [\"TFs\"])\n",
        "TFper = TFper.transpose()\n",
        "TFper[\"NotTFs\"] = TFper.apply(lambda row: 100 - row, axis=0)\n",
        "TFper = TFper.reindex(xorder)\n",
        "\n",
        "TFper.plot.bar(stacked = True,\n",
        "\t\t\t\tcolor = {\"TFs\":\"forestgreen\", \"NotTFs\":\"darkgrey\"},\n",
        "\t\t\t\tedgecolor = \"black\",\n",
        "\t\t\t\tax = ax_4,\n",
        "\t\t\t\tylim = [0,25])\n",
        "TFhandles, TFlabels = ax_4.get_legend_handles_labels()\n",
        "ax_4.legend(handles=TFhandles[:-1], labels=TFlabels[:-1],\n",
        "          loc='center left', bbox_to_anchor=(1, 0.5))\n",
        "ax_4.set_ylabel(\"% DEGs\", rotation = 90, fontsize=16)\n",
        "ax_4.yaxis.set_label_coords(-0.06,0.36)\n",
        "ax_4.axes.get_xaxis().set_visible(False)\n",
        "\n",
        "# ax_6 DGE % of DGEs that are kinases\n",
        "kindir = dir_path + 'tf_kinases/'\n",
        "kinpaths = [x for x in os.listdir(kindir) if \".ini\" not in x and \".kinases.\" in x]\n",
        "kindict = {}\n",
        "\n",
        "for file in kinpaths:\n",
        "\ttempspe = spedicto[file.split(\".\")[0]]\n",
        "\tcontent = open(kindir + file, \"r\")\n",
        "\tfor line in content:\n",
        "\t\tgene, anno, anno1, anno2 = line.strip().split(\"\\t\")\n",
        "\t\tif anno!= \"NoFunction\":\n",
        "\t\t\tkindict[gene] = [anno, anno1, anno2]\n",
        "\t\n",
        "kindf = pd.DataFrame(columns = list(sigmod.columns))\n",
        "for sx in sigtable.stress.unique():\n",
        "\tif sx.split(\"_\")[0] == \"Ath\":\n",
        "\t\tkinsubset = sigmod[(sigmod.stress == sx) & (sigmod.apply(lambda row: row[\"gene\"].upper() in kindict, axis = 1))]\n",
        "\telse:\n",
        "\t\tkinsubset = sigmod[(sigmod.stress == sx) & (sigmod.apply(lambda row: row[\"gene\"] in kindict, axis = 1))]\n",
        "\tkindf = kindf.append(kinsubset, ignore_index = True)\n",
        "\t\n",
        "kinperdict = {}\n",
        "for sx in sigtable.stress.unique():\n",
        "\tkinperdict[sx] = (len(kindf[kindf.stress == sx])/len(sigmod[sigmod.stress == sx]))*100\n",
        "kinper = pd.DataFrame(kinperdict, index = [\"kinases\"])\n",
        "kinper = kinper.transpose()\n",
        "kinper[\"NotKinases\"] = kinper.apply(lambda row: 100 - row, axis=0)\n",
        "kinper = kinper.reindex(xorder)\n",
        "\n",
        "kinper.plot.bar(stacked = True,\n",
        "\t\t\t\tcolor = {\"kinases\":\"darkgoldenrod\", \"NotKinases\":\"darkgrey\"},\n",
        "\t\t\t\tedgecolor = \"black\",\n",
        "\t\t\t\tax = ax_6,\n",
        "\t\t\t\tylim = [0,25])\n",
        "kinhandles, kinlabels = ax_6.get_legend_handles_labels()\n",
        "ax_6.legend(handles=kinhandles[:-1], labels=kinlabels[:-1],\n",
        "          loc='center left', bbox_to_anchor=(1, 0.5))\n",
        "ax_6.set_ylabel(\"% DEGs\", rotation = 90, fontsize=16)\n",
        "ax_6.yaxis.set_label_coords(-0.06,0.36)\n",
        "ax_6.axes.get_xaxis().set_visible(False)\n",
        "\n",
        "# ax_8 DGE, up/down ratio\n",
        "uddict = {}\n",
        "for sx in sigtable.stress.unique():\n",
        "\tupcount = len(sigtable[(sigtable.stress == sx) & (sigtable.L2FC_D2 == \"UP\")])\n",
        "\tdowncount = len(sigtable[(sigtable.stress == sx) & (sigtable.L2FC_D2 == \"DOWN\")])\n",
        "\ttotal = upcount + downcount\n",
        "\tuddict[sx] = [downcount/(total)*100, upcount/(total)*100]\n",
        "udper = pd.DataFrame(uddict, index = [\"downregulated\", \"upregulated\"])\n",
        "udper = udper.transpose()\n",
        "udper = udper.reindex(xorder)\n",
        "udper.plot.bar(stacked = True,\n",
        "\t\t\t\tcolor = {\"upregulated\":\"firebrick\", \"downregulated\":\"navy\"},\n",
        "\t\t\t\tedgecolor = \"white\",\n",
        "\t\t\t\tyticks = [0, 50, 100],\n",
        "\t\t\t\tax = ax_8)\n",
        "dgehandles, dgelabels = ax_8.get_legend_handles_labels()\n",
        "ax_8.legend(handles=dgehandles[::-1], labels=dgelabels[::-1],\n",
        "\t\t\tloc='center left', bbox_to_anchor=(1, 0.5))\n",
        "ax_8.set_ylabel(\"% DEGs\", rotation = 90, fontsize=16)\n",
        "ax_8.yaxis.set_label_coords(-0.06,0.36)\n",
        "ax_8.axes.get_xaxis().set_visible(False)\n",
        "\n",
        "# ax2/3 Dendrogram\n",
        "\n",
        "plot_dendro(xmat, ax2, 'top')\n",
        "plot_dendro(ymat, ax3, 'left')\n",
        "\n",
        "# ax4 Heatmap\n",
        "from matplotlib.colors import ListedColormap\n",
        "cmap = ListedColormap([\"lightgray\", \"royalblue\", \"violet\", \"firebrick\"])\n",
        "catno = 4\n",
        "\n",
        "hplot = ax4.imshow(df_sig_reordered, cmap=cmap)\n",
        "ax4.yaxis.tick_right()\n",
        "ax4.set_ylabel(\"\")\n",
        "ax4.set_xticks(np.arange(0, len(df_sig_reordered.columns), 1))\n",
        "ax4.set_yticks(np.arange(0, len(df_sig_reordered), 1))\n",
        "\n",
        "xcolour = [species_color(x) for x in xorder]\n",
        "newlabel = df_sig_reordered.columns.to_list()\n",
        "longlabel = ['heat', 'cold', 'light', 'dark', 'salt', 'mannitol', 'nitrogen']\n",
        "shortlabel = ['H', 'C', 'L', 'D', 'S', 'M', 'N']\n",
        "for i, y in enumerate(longlabel):\n",
        "\tnewlabel = [x.replace(y, shortlabel[i]) for x in newlabel]\n",
        "newlabel = [x.replace('_', ' ') for x in newlabel]\n",
        "ax4.set_xticklabels(newlabel, rotation=90, fontsize=18)\n",
        "\n",
        "for i, tick_label in enumerate(ax4.get_xticklabels()):\n",
        "\ttick_text = tick_label.get_text()\n",
        "\ttick_label.set_color(xcolour[i])\n",
        "\n",
        "anno_long = ['annotated', 'cellulose', 'biosynthesis', 'hemicellulose', 'pectin', 'channels', 'degradation']\n",
        "ax4.set_yticklabels([dicto[int(x)] for x in df_sig_reordered.index.to_list()], fontsize=18)\n",
        "\n",
        "# colourbar\n",
        "cbarticks = [(x/(catno*2))*(catno-1) for x in range(1,catno*2,2)]\n",
        "axins = inset_axes(ax1,\n",
        "\t\t\t\t\twidth=\"40%\",\n",
        "\t\t\t\t\theight=\"90%\", \n",
        "\t\t\t\t\tloc = 'center')\n",
        "cbar = fig.colorbar(hplot, cax=axins, ticks = cbarticks)\n",
        "cbar.ax.set_yticklabels(['N', 'D', 'UD', 'U'])\n",
        "\n",
        "#plt.tight_layout()\n",
        "plt.savefig(dir_path + 'figures/fig2.png',\n",
        "\t\t\t\tdpi=600,\n",
        "\t\t\t\tbbox_inches='tight')\n",
        "\n",
        "# reset rcparams\n",
        "import matplotlib as mpl\n",
        "mpl.rcParams.update(mpl.rcParamsDefault)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EGvetfN37wGO"
      },
      "source": [
        "### Supp. Fig 5: Interspecies comparison (Gene families)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gaI3DxHSPhD2"
      },
      "source": [
        "wdir = dir_path + 'prep_files/'\n",
        "jdir = wdir + 'proteomes/'\n",
        "OFfile = 'Orthogroups.txt'\n",
        "DGEfile = wdir + 'Figure2_alldata_compiled_updated.txt'\n",
        "jhdir = wdir + 'interspeGF/'\n",
        "jhdir_safe = dir_path_safe + 'prep_files/interspeGF/'\n",
        "if not os.path.exists(jhdir):\n",
        "\t!mkdir $jhdir_safe\n",
        "\n",
        "# initialise species\n",
        "spedicto = {'ARATH' : 'Ath',\n",
        "\t\t\t'CHLRE' : 'Cre',\n",
        "\t\t\t'CYAPA' : 'Cpa',\n",
        "\t\t\t'MARPO' : 'Mpo',\n",
        "\t\t\t'ORYSA' : 'Osa'}\n",
        "species_list = list(spedicto.values())\n",
        "\t\t\n",
        "# initialise genes in species\n",
        "Gdicto = {}\n",
        "# {\"species\" : [\"gene1\", \"gene2\"...]}\n",
        "for pepfile in [x for x in os.listdir(jdir) if '.ini' not in x]:\n",
        "\twith open(jdir + pepfile, \"r\") as peppy:\n",
        "\t\tspecies, genes = pepfile.split('.fa')[0], []\n",
        "\t\tfor lini in peppy:\n",
        "\t\t\tif '>' in lini:\n",
        "\t\t\t\tgenes.append(lini.strip().split('>')[1])\n",
        "\t\tGdicto[spedicto[species]] = genes\n",
        "\n",
        "# initialise orthofinder groups and corresponding genes by species\n",
        "OFdicto = {}\n",
        "OG_list = []\n",
        "#{\"OGx_spe\":[\"gene1\", \"gene2\"]}\n",
        "with open(wdir+OFfile, \"r\") as content:\n",
        "\tfor line in content:\n",
        "\t\tog, val = line.split(': ')[0], line.rstrip().split(': ')[1]\n",
        "\t\tOG_list.append(og)\n",
        "\t\tfor spe in spedicto.values():\n",
        "\t\t\tOFdicto[og + \"_\" + spe] = [x for x in val.split(' ') if x in Gdicto[spe]]\n",
        "\n",
        "        \n",
        "# initialise DGEs\n",
        "DGEdicto = {}\n",
        "# {\"gene1_stress\" : \"UP/DOWN\"}\n",
        "spe_stress = []\n",
        "with open(DGEfile, \"r\") as dgecon:\n",
        "\tdgecon.readline()\n",
        "\tfor lino in dgecon:\n",
        "\t\tgene, stress, L2FC_D2, L2FC_H2, annotation = lino.strip().split(\"\\t\")\n",
        "\t\t# to account for difference in gene names in this file and in Orthogroups.txt and ARATH.fa\n",
        "\t\tif stress.split(\"_\")[0] == \"Ath\":\n",
        "\t\t\tgene = gene.upper()\n",
        "\t\telif stress.split(\"_\")[0] == \"Cre\":\n",
        "\t\t\tgene = gene.split(\".t\")[0]\n",
        "\t\tDGEdicto[gene +  '_' + stress] = L2FC_D2\n",
        "\t\tif stress not in spe_stress:\n",
        "\t\t\tspe_stress.append(stress)\n",
        "\n",
        "# Functions\n",
        "def get_ortho_status(og, spe, stress):\n",
        "\t\"\"\"\n",
        "\tParameters\n",
        "\t----------\n",
        "\tog : str\n",
        "\t\torthogroup name.\n",
        "\tspe : str\n",
        "\t\tspecies code (3 letters).\n",
        "\tstress : str\n",
        "\t\ttype of stress.\n",
        "\n",
        "\tReturns\n",
        "\t-------\n",
        "\tstatus : str/None\n",
        "\t\tog + \"_UP\": consistently UP\n",
        "\t\tog + \"_DOWN\": consistently DOWN\n",
        "\t\tog + \"_AMB\": ambiguous, UP and DOWN detected\n",
        "\t\tog + \"_NC\" : all genes present have no signicant DGEs\n",
        "\t\tNone: No gene present in orthogroup\n",
        "\n",
        "\t\"\"\"\n",
        "\tstatus = []\n",
        "\tif len(OFdicto[og + \"_\" + spe]) > 0:\n",
        "\t\tfor gene in OFdicto[og + \"_\" + spe]:\n",
        "\t\t\tif gene + \"_\" + spe + \"_\"+ stress in DGEdicto:\n",
        "\t\t\t\tstatus.append(DGEdicto[gene + \"_\" + spe + \"_\"+ stress])\n",
        "\t\t\telse:\n",
        "\t\t\t\tstatus.append(\"NC\") # need to account for unchanged genes\n",
        "\t\tstatus = list(set(status))\n",
        "\t\tif len(status) > 1:\t\t\t# more than 1 type, AMB: ambiguous\n",
        "\t\t\tif \"UP\" in status and \"DOWN\" in status:\n",
        "\t\t\t\tstatus = \"AMB\"\n",
        "\t\t\telif \"UP\" in status:\n",
        "\t\t\t\tstatus = \"UP\"\n",
        "\t\t\telif \"DOWN\" in status:\n",
        "\t\t\t\tstatus = \"DOWN\"\n",
        "\t\telif  status[0] == \"NC\":\t# only 1 type, NC: no change\n",
        "\t\t\tstatus = \"NC\"\n",
        "\t\telif status[0] == \"DOWN\":\t# only 1 type, DOWN: downregulated\n",
        "\t\t\tstatus = \"DOWN\"\n",
        "\t\telif status[0] == \"UP\":\t\t# only 1 type, UP: upregulated\n",
        "\t\t\tstatus = \"UP\"\n",
        "\telse:\n",
        "\t\tstatus = \"None\" # no gene in the orthogroup\n",
        "\treturn status\n",
        "\n",
        "all_statuses = {}\n",
        "for item in spe_stress:\n",
        "\tspecies, stress_type = item.split(\"_\", 1)\n",
        "\t# container for collecting all the orthogroup status for item in spe_stress\n",
        "\tspe_stress_stat = []\n",
        "\tfor orthogroup in OG_list:\n",
        "\t\tspe_stress_stat.append(get_ortho_status(orthogroup, species, stress_type))\n",
        "\tall_statuses[item] = spe_stress_stat\n",
        "\n",
        "# Initialize dataframe for OG (rows)/spe_stress (cols)\n",
        "OGstats = pd.DataFrame(all_statuses, index = OG_list, columns = spe_stress)\n",
        "OGstats.to_csv(jhdir + 'OGstats.txt', sep=\"\\t\")\n",
        "\n",
        "# Calculate Jaccard distance\n",
        "\n",
        "# Functions\n",
        "def cal_jd(cond1, cond2):\n",
        "\t\"\"\"\n",
        "\tParameters\n",
        "\t----------\n",
        "\tstress1 : pandas Series\n",
        "\t\tOG stat of first stress\n",
        "\tstress2 : pandas Series\n",
        "\t\tOG stat of second stress\n",
        "\n",
        "\tReturns\n",
        "\t-------\n",
        "\tstatus : float\n",
        "\t\tmodified jaccard distance\n",
        "\t\"\"\"\n",
        "\t# Score container\n",
        "\tscore = \"\"\n",
        "\tif \"None\" in cond1 or \"None\" in cond2: \t# 1 or more OG absent\n",
        "\t\tscore = None\n",
        "\telif cond1 != cond2:\t\t\t\t\t# status do not match \n",
        "\t\tscore = 0\n",
        "\telif cond1 == \"NC\" and cond2 == \"NC\":\t# match but NC\n",
        "\t\tscore = None\n",
        "\telif cond1 == cond2:\n",
        "\t\tscore = 1\t\t\t\t\t\t\t# match\n",
        "\treturn score\n",
        "\n",
        "jd_dict = {}\n",
        "jd_counts = {}\n",
        "for i in range(len(spe_stress)):\n",
        "\t# container for jaccard distances\n",
        "\tjd_con = []\n",
        "\t# container for counts of OG per comparison\n",
        "\tcount_con = []\n",
        "\tfor j in range(len(spe_stress)):\n",
        "\t\tif spe_stress[i] == spe_stress[j]:\n",
        "\t\t\tjd_con.append(0)\n",
        "\t\t\t# reflects the number of OG != \"None\"\n",
        "\t\t\tcount_con.append(sum(OGstats[spe_stress[i]] != \"None\"))\n",
        "\t\telse:\n",
        "\t\t\tinterdf = OGstats[[spe_stress[i], spe_stress[j]]]\n",
        "\t\t\tscoreseries = interdf.apply(lambda row: cal_jd(row[spe_stress[i]], row[spe_stress[j]]), axis=1)\n",
        "\t\t\tjd_con.append(1 - (scoreseries.sum()/scoreseries.count()))\n",
        "\t\t\tcount_con.append(scoreseries.count())\n",
        "\tjd_dict[spe_stress[i]] = jd_con\n",
        "\tjd_counts[spe_stress[i]] = count_con\n",
        "\t\n",
        "# Initialize dataframe for spe_stress (rows)/spe_stress (cols) [jaccard distance]\n",
        "JDstats = pd.DataFrame(jd_dict, index = spe_stress, columns = spe_stress)\n",
        "JDstats.to_csv(jhdir + 'JDstats.txt', sep=\"\\t\")\n",
        "# Initialize dataframe for spe_stress (rows)/ spe_stress (cols)\n",
        "# [OG count used to calculate jaccard distance]\n",
        "JDcounts = pd.DataFrame(jd_counts, index = spe_stress, columns = spe_stress)\n",
        "JDcounts.to_csv(jhdir + 'JDcounts.txt', sep=\"\\t\")\n",
        "\n",
        "# clustermap -- distance matrix\n",
        "\n",
        "def label_color(xlabel):\n",
        "\tif \"heat\" in xlabel or 'Mpo_H' in xlabel:\n",
        "\t\treturn \"firebrick\"\n",
        "\telif \"cold\" in xlabel or 'Mpo_C' in xlabel:\n",
        "\t\treturn \"steelblue\"\n",
        "\telif \"light\" in xlabel or 'Mpo_L' in xlabel:\n",
        "\t\treturn \"darkorange\"\n",
        "\telif \"dark\" in xlabel or 'Mpo_D' in xlabel:\n",
        "\t\treturn \"black\"\n",
        "\telif \"salt\" in xlabel or 'Mpo_S' in xlabel:\n",
        "\t\treturn \"rebeccapurple\"\n",
        "\telif \"mannitol\" in xlabel or 'Mpo_M' in xlabel:\n",
        "\t\treturn \"mediumvioletred\"\n",
        "\telif \"nitrogen\" in xlabel or 'Mpo_N' in xlabel:\n",
        "\t\treturn \"forestgreen\"\n",
        "\telse:\n",
        "\t\treturn \"slategrey\"\n",
        "\n",
        "import scipy.spatial as sp, scipy.cluster.hierarchy as hc\n",
        "\n",
        "colnames = JDstats.index.to_series().apply(lambda row: row.split(\"_\")[0])\n",
        "specoldict = dict(zip(colnames.unique(), \"rgbcy\"))\n",
        "specol = colnames.map(specoldict)\n",
        "\n",
        "linkage = hc.linkage(sp.distance.squareform(JDstats), method='single')\n",
        "g = sns.clustermap(JDstats,\n",
        "\t\t\t\t   row_linkage = linkage,\n",
        "\t\t\t\t   col_linkage = linkage,\n",
        "\t\t\t\t   row_colors = specol,\n",
        "\t\t\t\t   xticklabels = True,\n",
        "\t\t\t\t   yticklabels=True)\n",
        "newlabel = [x.get_text() for x in g.ax_heatmap.axes.get_xticklabels()]\n",
        "longlabel = ['heat', 'cold', 'light', 'dark', 'salt', 'mannitol', 'nitrogen']\n",
        "shortlabel = ['H', 'C', 'L', 'D', 'S', 'M', 'N']\n",
        "for i, y in enumerate(longlabel):\n",
        "\tnewlabel = [x.replace(y, shortlabel[i]) for x in newlabel]\n",
        "newlabel = [x.replace('_', ' ') for x in newlabel]\n",
        "g.ax_heatmap.axes.set_xticklabels(newlabel, rotation=90, fontsize=16)\n",
        "\n",
        "for tick_label in g.ax_heatmap.axes.get_yticklabels():\n",
        "\ttick_text = tick_label.get_text()\n",
        "\ttick_label.set_color(label_color(tick_text))\n",
        "g.ax_heatmap.axes.set_yticklabels(newlabel, fontsize=16)\n",
        "\n",
        "plt.savefig(dir_path + 'figures/suppfig5.png')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3oBtH2a37ckp"
      },
      "source": [
        "### Figure 3: Stress responsiveness"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XsnoIc2NYqKh"
      },
      "source": [
        "# adapted from stress_res_og.py\n",
        "OFpath = dir_path + 'prep_files/Orthogroups.txt'\n",
        "jhdir = dir_path + 'prep_files/interspeGF/'\n",
        "prefix = {'Cpa|' : 'Cpa',\n",
        "\t\t  'Cre' : 'Cre',\n",
        "\t\t  'Mp' : 'Mpo',\n",
        "\t\t  'ChrUn' : 'Osa',\n",
        "\t\t  'LOC_Os' : 'Osa',\n",
        "\t\t  'AT' : 'Ath'}\n",
        "spelist = list(prefix.values())\n",
        "\n",
        "def spe_finder(gene):\n",
        "\t'''\n",
        "\tFinds the species the gene belongs to\n",
        "\n",
        "\tParameters\n",
        "\t----------\n",
        "\tgene : str\n",
        "\t\tGene ID.\n",
        "\n",
        "\tReturns\n",
        "\t-------\n",
        "\tspestat : str\n",
        "\t\tCorresponding species of gene.\n",
        "\n",
        "\t'''\n",
        "\tspestat = 'Other'\n",
        "\tfor k in list(prefix.keys()):\n",
        "\t\tif gene.startswith(k):\n",
        "\t\t\tspestat = prefix[k]\n",
        "\treturn spestat\n",
        "\n",
        "# order of species\n",
        "spe_order = ['Cpa', 'Cre', 'Mpo', 'Osa', 'Ath']\n",
        "spe_class = [['Angiosperm', ['Osa', 'Ath']],\n",
        "\t\t\t ['Embryophyte', ['Mpo', 'Osa', 'Ath']],\n",
        "\t\t\t ['Viridiplantae', spe_order[1:]],\n",
        "\t\t\t ['Archaeplastida', spe_order]]\n",
        "\n",
        "def speclass(spelist):\n",
        "\tspec = []\n",
        "\tfor spe in spelist:\n",
        "\t\tfor i, c in enumerate(spe_class):\n",
        "\t\t\tif spe in c[1]:\n",
        "\t\t\t\tspec.append(i)\n",
        "\t\t\t\tbreak\n",
        "\treturn spe_class[max(spec)][0]\n",
        "\n",
        "# initialise OG stats df\n",
        "OGstats = pd.read_csv(jhdir + 'OGstats.txt', sep=\"\\t\", index_col=0)\n",
        "newlabel = OGstats.columns.to_list()\n",
        "longlabel = ['heat', 'cold', 'light', 'dark', 'salt', 'mannitol', 'nitrogen', 'drought']\n",
        "shortlabel = ['H', 'C', 'L', 'D', 'S', 'M', 'N', 'M']\n",
        "for i, y in enumerate(longlabel):\n",
        "\tnewlabel = [x.replace(y, shortlabel[i]) for x in newlabel]\n",
        "OGstats.columns = newlabel\n",
        "\n",
        "\n",
        "# initialise DGEs\n",
        "DGEfile = dir_path + 'prep_files/Figure2_alldata_compiled_updated.txt'\n",
        "DGElist = []\n",
        "DGEbins = {}\n",
        "# {\"gene1_stress\" : \"UP/DOWN\"}\n",
        "spe_stress = []\n",
        "with open(DGEfile, \"r\") as dgecon:\n",
        "\tdgecon.readline()\n",
        "\tfor lino in dgecon:\n",
        "\t\tgene, stress, L2FC_D2, L2FC_H2, annotation = lino.strip().split(\"\\t\")\n",
        "\t\t# to account for difference in gene names in this file and in Orthogroups.txt and ARATH.fa\n",
        "\t\tif stress.split(\"_\")[0] == \"Ath\":\n",
        "\t\t\tgene = gene.upper()\n",
        "\t\telif stress.split(\"_\")[0] == \"Cre\":\n",
        "\t\t\tgene = gene.split(\".t\")[0]\n",
        "\t\tDGElist.append(gene)\n",
        "\t\tbinlist = [int(x[0].replace(\"'\", \"\").split('.')[0]) for x in literal_eval(annotation)]\n",
        "\t\tDGEbins[gene] = list(set(binlist))\n",
        "DGElist = set(DGElist)\n",
        "\n",
        "og_genes = {}\n",
        "spespec_og = {} # list of species specific OGs, excludes OGs of species not included in analysis\n",
        "other_og = {} # # list of non-species specific OGs, excludes OGs of species not included in analysis\n",
        "with open(OFpath, 'r') as OFfile:\n",
        "\tfor line in OFfile:\n",
        "\t\tog, val = line.split(': ')[0], line.rstrip().split(': ')[1]\n",
        "\t\tog_species = list(set([spe_finder(x) for x in val.split(' ')]))\n",
        "\t\tif len(og_species) == 1 and og_species[0] in spelist:\n",
        "\t\t\tspespec_og[og] = og_species[0]\n",
        "\t\t\tog_genes[og] = list(set(val.split(' ')) & DGElist)\n",
        "\t\telif len(og_species) > 1:\n",
        "\t\t\togclass = speclass(og_species)\n",
        "\t\t\tother_og[og] = ogclass\n",
        "\t\t\tog_genes[og] = list(set(val.split(' ')) & DGElist)\n",
        "\n",
        "\n",
        "# =============================================================================\n",
        "# \n",
        "# Stress-responsive OGs\n",
        "# \n",
        "# =============================================================================\n",
        "from collections import Counter, defaultdict\n",
        "\n",
        "df_coln = ['Archaeplastida', 'Viridiplantae', 'Embryophyte', 'Angiosperm']\n",
        "\n",
        "# omit OGs that contain only 'None' across all stresses\n",
        "oglist = OGstats.index.to_list()\n",
        "suboglist = [x for x in oglist if set(OGstats.loc[x].to_list()) != {'None'}] \n",
        "val_spespec_og = list(set(spespec_og) & set(suboglist))\n",
        "val_spespec_og.sort()\n",
        "val_other_og = list(set(suboglist) - set(val_spespec_og))\n",
        "val_other_og.sort()\n",
        "\n",
        "# df that contains only OGs that are not made up of 'None'\n",
        "subdf = OGstats.loc[suboglist]\n",
        "\n",
        "# dictionary to contain all statuses\n",
        "og_stress_stat = defaultdict(list)\n",
        "\n",
        "def update_stress_stat(newstat, phyla, dfcount, og, stresstype):\n",
        "\tdfcount.loc[newstat, phyla] += 1\n",
        "\tog_stress_stat[og].append(stresstype + '_' + newstat)\n",
        "\t\n",
        "# to get counts per stress\n",
        "stresslist = shortlabel[:-1]\n",
        "def counts_per_stress(slabel):\n",
        "\tstresstype = '_' + slabel\n",
        "\tvalid_exp = [x for x in OGstats.columns.to_list() if stresstype in x]\n",
        "\tvalid_spe = [x.split('_')[0] for x in valid_exp]\n",
        "\tordered_exp = [y for x in spe_order for y in valid_exp if x in y]\n",
        "\tunique_spe = [x for x in spe_order if x in valid_spe]\n",
        "\tstresssub = subdf[ordered_exp] # df containing only exps of required stress type\n",
        "\t\n",
        "\tdumdict = {'UP': [0 for x in range(4 + len(unique_spe))],\n",
        "\t\t\t   'DOWN': [0 for x in range(4 + len(unique_spe))],\n",
        "\t\t\t   'AMB': [0 for x in range(4 + len(unique_spe))],\n",
        "\t\t\t   'MIXED': [0 for x in range(4 + len(unique_spe))],\n",
        "\t\t\t   'NR': [0 for x in range(4 + len(unique_spe))]}\n",
        "\t\n",
        "\tdfcount = pd.DataFrame.from_dict(dumdict, orient='index', columns = df_coln + unique_spe)\n",
        "\tnogroup = []\n",
        "\tfor og, ogclass in spespec_og.items(): # species specific OGs\n",
        "\t\tkcount = Counter(stresssub.loc[og].to_list())\n",
        "\t\t# species specific OGs\n",
        "\t\tif kcount['None'] != len(valid_spe): # ignore if OG not valid for stress\n",
        "\t\t\tif slabel == 'S' and ogclass == 'Osa':\n",
        "\t\t\t\tif kcount['NC'] == 2: # for species specific OG that is 'NC' [Osa]\n",
        "\t\t\t\t\tupdate_stress_stat('NR', ogclass, dfcount, og, slabel)\n",
        "\t\t\t\telse: # kcount['NC']!= 2\n",
        "\t\t\t\t\tstat = [x for x in kcount if x != 'None' and 'NC']\n",
        "\t\t\t\t\tif len(stat) == 1: # only one type of UP/DOWN/AMB\n",
        "\t\t\t\t\t\tupdate_stress_stat(stat[0], ogclass, dfcount, og, slabel)\n",
        "\t\t\t\t\telse: # mixture of UP/DOWN/AMB\n",
        "\t\t\t\t\t\tupdate_stress_stat('MIXED', ogclass, dfcount, og, slabel)\n",
        "\t\t\telse:\n",
        "\t\t\t\tif 'NC' not in kcount: # for species specific OG that is not 'NC'\n",
        "\t\t\t\t\tstat = [x for x in kcount if kcount[x] == 1]\n",
        "\t\t\t\t\tupdate_stress_stat(stat[0], ogclass, dfcount, og, slabel)\n",
        "\t\t\t\telse: # for species specific OG that is 'NC'\n",
        "\t\t\t\t\tupdate_stress_stat('NR', ogclass, dfcount, og, slabel)\n",
        "\t\n",
        "\tfor og, ogclass in other_og.items(): # non-species specific OGs\n",
        "\t\tkcount = Counter(stresssub.loc[og].to_list())\t\n",
        "\t\tif kcount['None'] != len(valid_spe): # ignore if OG not valid for stress\n",
        "\t\t\tnotnil = [x for x in list(kcount.keys()) if x != 'None' and x != 'NC']\n",
        "\t\t\tif len(notnil) > 0:\n",
        "\t\t\t\tif len(notnil) > 1: # contains combination of UP/DOWN/AMB\n",
        "\t\t\t\t\tupdate_stress_stat('MIXED', ogclass, dfcount, og, slabel)\n",
        "\t\t\t\telif len(notnil) == 1: # contains only one type of status apart from 'NR'\n",
        "\t\t\t\t\tif kcount[notnil[0]] > 1: # if UP/DOWN/AMB appear more than once\n",
        "\t\t\t\t\t\tupdate_stress_stat(notnil[0], ogclass, dfcount, og, slabel)\n",
        "\t\t\t\t\telif kcount[notnil[0]] == 1: # if UP/DOWN/AMB only appear once\n",
        "\t\t\t\t\t\tupdate_stress_stat('NR', ogclass, dfcount, og, slabel)\n",
        "\t\t\t\t\telse:\n",
        "\t\t\t\t\t\tnogroup.append(og)\n",
        "\t\t\telse: # OGs that only have 'None' and 'NC'\t\n",
        "\t\t\t\tupdate_stress_stat('NR', ogclass, dfcount, og, slabel)\n",
        "\t\telse: # all 'None', meaning that OG is valid in other species not present in this analysis\n",
        "\t\t\tupdate_stress_stat('NR', ogclass, dfcount, og, slabel)\n",
        "\treturn dfcount.T\n",
        "\n",
        "dH, dC, dL, dD, dS, dM, dN = [counts_per_stress(x) for x in stresslist]\n",
        "dflist = dH, dC, dL, dD, dS, dM, dN\n",
        "wdir = dir_path + 'phylostrata/'\n",
        "wdir_safe = dir_path_safe + 'phylostrata/'\n",
        "if not os.path.exists(wdir):\n",
        "    !mkdir $wdir_safe\n",
        "for i, df in enumerate(dflist):\n",
        "\tdf.to_csv(wdir + stresslist[i] + '_df.txt', sep=\"\\t\")\n",
        "\n",
        "# =============================================================================\n",
        "# \n",
        "# Quantifying stress responsiveness of Orthogroups\n",
        "# \n",
        "# =============================================================================\n",
        "\n",
        "# intermediate container for new df\n",
        "resog_dict = {}\n",
        "for dicto in [other_og, spespec_og]:\n",
        "\tfor og, phyla in dicto.items():\n",
        "\t\treslist = [x for x in og_stress_stat[og] if 'NR' not in x]\n",
        "\t\tresog_dict[og] = [phyla, reslist, len(reslist), og_genes[og]]\n",
        "\n",
        "resog_df = pd.DataFrame.from_dict(resog_dict, orient='index', columns=['Phylostrata', 'Responsive in', 'Count', 'Genes'])\n",
        "resog_df.to_csv(wdir + 'resog_df.txt', sep=\"\\t\")\n",
        "\n",
        "sorder = df_coln + spe_order\n",
        "grouped_count = resog_df.groupby(['Phylostrata','Count']).count()['Genes'].unstack().reindex(sorder)\n",
        "grouped_per = grouped_count.copy()\n",
        "for row in sorder:\n",
        "\tgrouped_per.loc[row] = grouped_per.loc[row].apply(lambda x: (x/grouped_count.loc[row].sum())*100)\n",
        "\n",
        "sorder.remove('Angiosperm')\n",
        "\n",
        "g = grouped_count.loc[sorder].plot.bar(stacked=True, ylabel = 'Count')\n",
        "g.legend(bbox_to_anchor=(1, 0.75))\n",
        "\n",
        "g2 = grouped_per.loc[sorder,[i for i in range(1,8)]].plot.bar(stacked=True, ylabel='Percentage (%)')\n",
        "g2.legend(bbox_to_anchor=(1, 0.75))\n",
        "\n",
        "# Percentage of OGs from various phylostrata that are responsive in respective number of stresses (x-axis)\n",
        "grouped_per_bycount = grouped_count.copy()\n",
        "for col in grouped_per_bycount.columns.to_list():\n",
        "\tgrouped_per_bycount[col] = grouped_per_bycount[col].apply(lambda x: (x/grouped_count[col].sum())*100)\n",
        "g3 = grouped_per_bycount.loc[sorder,[i for i in range(1,8)]].T.plot.bar(stacked=True, ylabel='Percentage (%)')\n",
        "g3.legend(bbox_to_anchor=(1, 1))\n",
        "\n",
        "# log y of number og OGs responsive in respective number of stresses (x axis)\n",
        "g = grouped_count.loc[sorder].T.plot.bar(logy=True, ylabel = 'Number of OGs')\n",
        "g.legend(bbox_to_anchor=(1, 0.75))\n",
        "\n",
        "countbysres = grouped_count.loc[sorder].T\n",
        "countbysres.to_csv(wdir + 'countbystressres.txt', sep='\\t')\n",
        "\n",
        "# Mapman bins\n",
        "import seaborn as sns\n",
        "import math\n",
        "merdict = literal_eval(open(dir_path + 'prep_files/merdict.txt', 'r').read())\n",
        "# By Phylo\n",
        "catdict = {} # Mapman bin count for Phylostrata that are stress responsive (Count > 0)\n",
        "for cat in sorder:\n",
        "\tcatbins = [z for x in resog_df[(resog_df.Phylostrata == cat) & (resog_df.Count > 0)].Genes.to_list() for y in x for z in DGEbins[y]]\n",
        "\tcatdict[cat] = Counter(catbins)\n",
        "\n",
        "catdf = pd.DataFrame.from_dict(catdict)\n",
        "catdf.sort_index(inplace=True)\n",
        "catdf.reset_index(inplace=True)\n",
        "catdf.columns = ['Mapman bins'] + sorder\n",
        "catdf['Mapman bins'] = catdf['Mapman bins'].apply(lambda x: merdict[x])\n",
        "catdf.set_index('Mapman bins', inplace=True)\n",
        "\n",
        "catperdf = catdf.copy()\n",
        "catperlogdf = catdf.copy()\n",
        "for x in catperdf.columns.to_list():\n",
        "\ttotal = catdf[x].sum()\n",
        "\tcatperdf[x] = catdf[x].apply(lambda x: (x/total)*100)\n",
        "\tcatperlogdf[x] = catdf[x].apply(lambda x: math.log((x/total)*100,2))\n",
        "\n",
        "catperdf.fillna(float(0), inplace=True)\n",
        "\n",
        "f = sns.clustermap(catperdf, yticklabels=True, col_cluster=False) # to get linkage for logged values (percentages can be filled 0 but cannot fill NaN with 0 for logged values)\n",
        "row_linkage = f.dendrogram_row.linkage\n",
        "\n",
        "sns.clustermap(catperlogdf, yticklabels=True, col_cluster=False, row_linkage = row_linkage, cmap='coolwarm')\n",
        "\n",
        "# By Stress responsiveness\n",
        "countdict = {} # Mapman bin count for Phylostrata that are stress responsive (Count > 0)\n",
        "for count in [i for i in range(1,8)]:\n",
        "\tcountbins = [z for x in resog_df[(resog_df.Count == count)].Genes.to_list() for y in x for z in DGEbins[y]]\n",
        "\tcountdict[str(count)] = Counter(countbins)\n",
        "countdf = pd.DataFrame.from_dict(countdict)\n",
        "countdf.sort_index(inplace=True)\n",
        "countdf.reset_index(inplace=True)\n",
        "countdf.columns = ['Mapman bins'] + [i for i in range(1,8)]\n",
        "countdf['Mapman bins'] = countdf['Mapman bins'].apply(lambda x: merdict[x])\n",
        "countdf.set_index('Mapman bins', inplace=True)\n",
        "\n",
        "# column normalised\n",
        "countperdf = countdf.copy()\n",
        "countperlogdf = countdf.copy()\n",
        "for x in countperdf.columns.to_list():\n",
        "\ttotal = countdf[x].sum()\n",
        "\tcountperdf[x] = countdf[x].apply(lambda x: (x/total)*100)\n",
        "\tcountperlogdf[x] = countdf[x].apply(lambda x: math.log((x/total)*100,2))\n",
        "\n",
        "countperdf.fillna(float(0), inplace=True)\n",
        "\n",
        "f2 = sns.clustermap(countperdf, yticklabels=True, col_cluster=False, cmap='coolwarm') # to get linkage for logged values (percentages can be filled 0 but cannot fill NaN with 0 for logged values)\n",
        "row_linkage2 = f2.dendrogram_row.linkage\n",
        "\n",
        "sns.clustermap(countperlogdf, yticklabels=True, col_cluster=False, row_linkage = row_linkage2, cmap='coolwarm')\n",
        "\n",
        "# row normalised\n",
        "countper_rownorm_df = countdf.copy()\n",
        "countper_rownorm_logdf = countdf.copy()\n",
        "for x in countper_rownorm_df.index.to_list():\n",
        "\ttotal = countdf.loc[x].sum()\n",
        "\tcountper_rownorm_df.loc[x] = countdf.loc[x].apply(lambda x: (x/total)*100)\n",
        "\tcountper_rownorm_logdf.loc[x] = countdf.loc[x].apply(lambda x: math.log((x/total)*100, 2))\n",
        "\n",
        "countper_rownorm_df.fillna(float(0), inplace=True)\n",
        "\n",
        "f3 = sns.clustermap(countper_rownorm_df, yticklabels=True, col_cluster=False,cmap='coolwarm') # to get linkage for logged values (percentages can be filled 0 but cannot fill NaN with 0 for logged values)\n",
        "row_linkage3 = f3.dendrogram_row.linkage\n",
        "\n",
        "sns.clustermap(countper_rownorm_logdf, yticklabels=True, col_cluster=False,\n",
        "\t\t\t   row_linkage = row_linkage3, cmap='coolwarm', figsize=(5,6))\n",
        "plt.savefig(dir_path + 'figures/fig3a', dpi=600)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "exDaka827fIh"
      },
      "source": [
        "### Figure 4: Upset plot and summary of DEGs in Marchantia"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A2YBe6KzehDL"
      },
      "source": [
        "# Fig 4A and B (adapted from DGE_count_sizecorr.py)\n",
        "cross = pd.read_csv(dir_path + 'prep_files/mpo/deseq/resSig_compiled.txt', sep='\\t')\n",
        "\n",
        "cross.stress = [x.split('_')[1] for x in list(cross.stress)]\n",
        "stress_l = list(cross.stress.unique())\n",
        "stress_l.sort(key=lambda x: len(x))\n",
        "c_dgecount = cross.groupby(['stress', 'L2FC_D2']).count().gene.to_frame(name='count')\n",
        "\n",
        "unstacked = c_dgecount.unstack().reindex(stress_l)\n",
        "ax = unstacked.plot.bar(figsize=(7,3), stacked=True, ylabel='Number of DEGs', color=['navy', 'firebrick'])\n",
        "\n",
        "handles, labels = ax.get_legend_handles_labels()\n",
        "ax.legend(handles=handles[::-1], labels=[x.split(', ')[1].split(')')[0] for x in labels][::-1])\n",
        "plt.savefig(dir_path + 'figures/fig4a.png', dpi=600)\n",
        "\n",
        "wdir = dir_path + 'prep_files/'\n",
        "infile = 'phase1n2_measurements_nooutliers.txt'\n",
        "\n",
        "measurements = pd.read_csv(wdir + infile, sep='\\t')\n",
        "\n",
        "def ss_grab(stress, condition):\n",
        "\t\"\"\"\n",
        "\tSlice the relevant condition for \n",
        "\tParameters\n",
        "\t----------\n",
        "\tstress : string\n",
        "\t\tStress of interest.\n",
        "\tcondition : string\n",
        "\t\tCondition of interest.\n",
        "\n",
        "\tReturns\n",
        "\t-------\n",
        "\tsssub : dataframe\n",
        "\t\tdf of single stress.\n",
        "\n",
        "\t\"\"\"\n",
        "\tsssub = measurements[(measurements.Stress == stress) & (measurements.Condition == condition)]\n",
        "\treturn sssub\n",
        "\n",
        "srep_keys = [['Cold', '3'],\n",
        "\t\t\t ['Heat', '33'],\n",
        "\t\t\t ['Salt', '40'],\n",
        "\t\t\t ['Mannitol', '100'],\n",
        "\t\t\t ['Light', '435'],\n",
        "\t\t\t ['Dark', '3'],\n",
        "\t\t\t ['Nitrogen', '0']]\n",
        "srepdf = measurements[(measurements.Stress == 'Cold') & (measurements.Condition == '3')]\n",
        "\n",
        "for s, c in srep_keys[1:]:\n",
        "\tsrepdf = pd.concat([srepdf, ss_grab(s, c)])\n",
        "\n",
        "crepdf = measurements[measurements.Condition == 'mixed']\n",
        "m_nocon = pd.concat([srepdf, crepdf])\n",
        "m_nocon.Stress = [x[0] if len(x) > 2 else x for x in m_nocon.Stress]\n",
        "noHL = m_nocon[m_nocon.Stress != 'HL']\n",
        "\n",
        "avg_meas = noHL.groupby('Stress').mean()[['Parea', 'Earea']]\n",
        "avg_meas.reindex(stress_l)\n",
        "\n",
        "totaldeg = cross.groupby(['stress']).count()['L2FC_D2']\n",
        "totaldeg.reindex(stress_l)\n",
        "avg_meas['totaldeg'] = totaldeg\n",
        "avg_meas = avg_meas[['totaldeg', 'Parea', 'Earea']]\n",
        "\n",
        "# size plots by df\n",
        "ax = avg_meas.plot.scatter(x='totaldeg', y='Parea', color='orange', label='Area (Day 15)')\n",
        "\n",
        "for ind, dat in avg_meas.iterrows():\n",
        "    ax.annotate(ind, (dat['totaldeg'], dat['Parea']),\n",
        "\t\t\t\txytext=(-4,-12), textcoords='offset points')\n",
        "\n",
        "avg_meas.plot.scatter(x='totaldeg', y='Earea', color='navy', label='Area (Day 21)', ax=ax)\n",
        "for ind, dat in avg_meas.iterrows():\n",
        "    ax.annotate(ind, (dat['totaldeg'], dat['Earea']),\n",
        "\t\t\t\txytext=(-4,-12), textcoords='offset points')\n",
        "\t\n",
        "# size plot with regression\n",
        "from scipy import stats\n",
        "\n",
        "q_colnames = avg_meas.columns.to_list()\n",
        "dicto = {'Parea' : 'Day 15', 'Earea' : 'Day 21'}\n",
        "\n",
        "def plot_reg(df, title):\n",
        "\tlabels = []\n",
        "\tcol = q_colnames[0]\n",
        "\tfor col2 in q_colnames[1:]:\n",
        "\t\tplt.scatter(col, col2, data=df)\n",
        "\t\t#m, c = np.polyfit(df[col], df[col2], 1)\n",
        "\t\tm, c, r_value, p_value, std_err = stats.linregress(df[col], df[col2])\n",
        "\t\tfor ind, dat in avg_meas.iterrows():\n",
        "\t\t\tplt.annotate(ind, (dat[col], dat[col2]),\n",
        "\t\t\t  xytext=(-4,-12), textcoords='offset points')\n",
        "\t\tplt.plot(df[col], m*df[col] + c)\n",
        "\t\tlabels.append(dicto[col2] + ' (R\\u00b2: ' + str(round(r_value**2,2)) + ', p: ' + str(round(p_value, 2)) + ')')\n",
        "\tplt.legend(labels)\n",
        "\tplt.xlabel('DEG count')\n",
        "\tplt.ylabel('Size (mm\\u00b2)')\n",
        "\tplt.savefig(dir_path + 'figures/fig4b.png', dpi=600)\n",
        "\t\n",
        "plot_reg(avg_meas, 'Number of DEGs vs Size')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0NoknTLri0z9"
      },
      "source": [
        "# Fig 4C and D (adapted from upset.py)\n",
        "import upsetplot\n",
        "from collections import defaultdict\n",
        "\n",
        "wdir = dir_path + 'prep_files/mpo/deseq/'\n",
        "odir = wdir + 'upset/'\n",
        "odir_safe = dir_path_safe + 'prep_files/mpo/deseq/' + 'upset/'\n",
        "\n",
        "if not os.path.exists(odir):\n",
        "\t!mkdir $odir_safe\n",
        "\n",
        "data = pd.read_csv(wdir + 'resSig_compiled.txt', sep = '\\t')\n",
        "\n",
        "### FUNCTIONS ###\n",
        "\n",
        "def upset_matrix(set_dict, stress_types):\n",
        "\tupset_data_sub = upsetplot.from_contents({k: v for k, v in set_dict.items() if k in stress_types})\n",
        "\treturn upset_data_sub # , fig=None\n",
        "\n",
        "\n",
        "def plot_selected(cond_dict, cond_list, filename, title, orient = \"horizontal\", cutoff=50):\n",
        "\tdf_set = upset_matrix(cond_dict, cond_list)\n",
        "\tdf_set = df_set.sort_index()\n",
        "\t\n",
        "\t# preparation to output all data\n",
        "\tindex_names = list(df_set.index.names)\n",
        "\tindex_list = df_set.index.to_list()\n",
        "\tset_count = Counter(index_list)\n",
        "\tcounter_list = [[k, v] for k, v in set_count.items()]\n",
        "\tcounter_list.sort(key = lambda x: x[1], reverse=True)\n",
        "\t\n",
        "\t# writing output to file\n",
        "\twith open(filename + \"_matrix.txt\", \"w+\") as ofile:\n",
        "\t\tofile.write(\"\\t\".join(index_names + ['count', 'genes']) + \"\\n\")\n",
        "\t\tfor i in counter_list:\n",
        "\t\t\tglist = df_set.loc[i[0],:].id.to_list()\n",
        "\t\t\tofile.write(\"\\t\".join([str(int(x)) for x in i[0]] + [str(i[1]), str(glist)]) + \"\\n\")\n",
        "\t\n",
        "\t# writing top 50 to file\n",
        "\tset_cutoff = set_count.most_common(cutoff)\n",
        "\tselection = [x[0] for x in set_cutoff]\n",
        "\twith open(filename + \"_top50.txt\", \"w+\") as cfile:\n",
        "\t\tcfile.write(\"\\t\".join([str(index_names), 'count', 'genes']) + \"\\n\")\n",
        "\t\tfor i in range(len(set_cutoff)):\n",
        "\t\t\tglist = df_set.loc[set_cutoff[i][0],:].id.to_list()\n",
        "\t\t\tcfile.write(\"\\t\".join([str(set_cutoff[i][0]), str(set_cutoff[i][1]), str(glist)]) + \"\\n\")\n",
        "\t# selection for plotting\n",
        "\tsel_matrix = df_set.loc[selection[0], :]\t\n",
        "\tfor i in range(1, len(selection)):\n",
        "\t\tsel_matrix = sel_matrix + df_set.loc[selection[i],:]\n",
        "\tupsetplot.plot(sel_matrix, orientation = orient, sort_by = 'cardinality')\n",
        "\tplt.title(title, size=20)\n",
        "\tif \"upreg\" in filename:\n",
        "\t\tfigname = 'c'\n",
        "\telse:\n",
        "\t\tfigname = 'd'\n",
        "\tplt.savefig(dir_path + 'figures/Fig4' + figname + '.png', dpi=600)\n",
        "\n",
        "### END ###\n",
        "\n",
        "# Reshape data to have for every category,\n",
        "cond_dict_U = defaultdict(list) # genres_movies\n",
        "cond_dict_D = defaultdict(list)\n",
        "for index, row in data.iterrows():\n",
        "\tif row['L2FC_D2'] == 'UP':\n",
        "\t\tcond_dict_U[row['stress'].split(\"_\")[1]].append(row['gene'])\n",
        "\telif row['L2FC_D2'] == 'DOWN':\n",
        "\t\tcond_dict_D[row['stress'].split(\"_\")[1]].append(row['gene'])\n",
        "\t\t\n",
        "all_stress_list = [x.split('_')[1] for x in data.stress.unique()]\n",
        "\n",
        "# initialise dictionaries of up and downregulated genes for each condition\n",
        "cond_dict_U_set = dict()\n",
        "cond_dict_D_set = dict()\n",
        "for k, v in cond_dict_U.items():\n",
        "    cond_dict_U_set[k] = set(v)\n",
        "for k, v in cond_dict_D.items():\n",
        "    cond_dict_D_set[k] = set(v)\n",
        "\n",
        "# Plot horizontal (default)\n",
        "plot_selected(cond_dict = cond_dict_D_set,\n",
        "\t\t\t  cond_list = all_stress_list,\n",
        "\t\t\t  filename = odir + \"all_downreg\",\n",
        "\t\t\t  title = \"Downregulated genes\")\n",
        "\n",
        "plot_selected(cond_dict = cond_dict_U_set,\n",
        "\t\t\t  cond_list = all_stress_list,\n",
        "\t\t\t  filename = odir + \"all_upreg\",\n",
        "\t\t\t  title = \"Upregulated genes\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_ENdRwbI7iDw"
      },
      "source": [
        "### Supp Figs 6 & 7, Figure 5: Inter-stress (Marchantia only) comparison"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2DTxA1RKpwvR"
      },
      "source": [
        "# Supp. figs 6 & 7 (adapted from indivenn_hm.py)\n",
        "from collections import defaultdict\n",
        "from matplotlib_venn import venn3\n",
        "\n",
        "wdir = dir_path + 'prep_files/mpo/deseq/'\n",
        "odir = dir_path + 'figures/'\n",
        "data = pd.read_csv(wdir + 'resSig_compiled.txt', sep = '\\t')\n",
        "\n",
        "# Mercator bin conversion\n",
        "dicto = literal_eval(open(dir_path + 'prep_files/merdict.txt', 'r').read())\n",
        "\n",
        "all_s = [x.split(\"_\")[1] for x in data.stress.unique()]\n",
        "single = [x.split(\"_\")[1] for x in data.stress.unique() if len(x.split(\"_\")[1]) == 1]\n",
        "cross = [x.split(\"_\")[1] for x in data.stress.unique() if len(x.split(\"_\")[1]) == 2]\n",
        "\n",
        "data.annotation = data.annotation.apply(literal_eval)\n",
        "data[\"mername\"] = data.annotation.apply(lambda x: dicto[int(x[0][0].split('.')[0])])\n",
        "\n",
        "dict_A = defaultdict(list)\n",
        "dict_U = defaultdict(list)\n",
        "dict_D = defaultdict(list)\n",
        "\n",
        "def sum_to_dict(dicto, stress, reg):\n",
        "\tif reg == \"ALL\":\n",
        "\t\tsubset = data[(data.stress == \"Mpo_\" + stress)]\n",
        "\telse:\n",
        "\t\tsubset = data[(data.stress == \"Mpo_\" + stress) & (data.L2FC_D2 == reg)]\n",
        "\tdicto[stress].append(set(subset.gene.to_list()))\n",
        "\tdicto[stress].append(subset.mername.to_list())\n",
        "\n",
        "def dict_to_df(dicto):\n",
        "\tdf = pd.DataFrame.from_dict(dicto, orient='index', columns=[\"gene\", \"mername\"])\n",
        "\treturn df\n",
        "\n",
        "for s in all_s:\n",
        "\t#sum_to_dict(dict_A, s, \"ALL\")\n",
        "\tsum_to_dict(dict_U, s, \"UP\")\n",
        "\tsum_to_dict(dict_D, s, \"DOWN\")\n",
        "\n",
        "#df_A = dict_to_df(dict_A)\n",
        "df_U = dict_to_df(dict_U)\n",
        "df_D = dict_to_df(dict_D)\n",
        "\n",
        "def plot_venn(df, s1, s2, c1, title, axis):\n",
        "\tvenn3([df.loc[s1].gene, df.loc[s2].gene, df.loc[c1].gene],\n",
        "\t   (s1, s2, c1),\n",
        "\t   ax = axis)\n",
        "\taxis.set_title(title, size=20)\n",
        "\n",
        "# create subplots\n",
        "xlen = 4\n",
        "ylen = math.ceil(len(all_s)/4)\n",
        "figw = xlen * 4\n",
        "figh = ylen * 3.5\n",
        "\n",
        "a_axes = string.ascii_uppercase[:len(all_s)]\n",
        "def plot_subplot(df, title_ext):\n",
        "\taxa = plt.figure(constrained_layout=True,\n",
        "\t\t\t\t figsize=(figw, figh)).subplot_mosaic(\n",
        "\t\t\t\t\t \"\"\"\n",
        "\t\t\t\t\t ABCD\n",
        "\t\t\t\t\t EFGH\n",
        "\t\t\t\t\t IJKL\n",
        "\t\t\t\t\t MNOP\n",
        "\t\t\t\t\t QR..\n",
        "\t\t\t\t\t \"\"\"\n",
        "\t\t\t\t\t )\n",
        "\n",
        "\tfor c in range(len(cross)):\n",
        "\t\tst = cross[c]\n",
        "\t\tplot_venn(df, st[0], st[1], st, st + title_ext, axa[a_axes[c]])\n",
        "\tplt.savefig(odir+'supp_fig6or7' + title_ext +'.png', dpi=600)\n",
        "\n",
        "\n",
        "#df_col = [[df_A, ''], [df_U, \"_upregulated\"], [df_D, \"_downregulated\"]]\n",
        "df_col = [[df_U, \"_upregulated\"], [df_D, \"_downregulated\"]]\n",
        "\n",
        "for x in df_col:\n",
        "\tdf_type, ext = x[0], x[1]\n",
        "\tplot_subplot(df_type, ext)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8EibR5DSslOq"
      },
      "source": [
        "# Figure 5 (adapted from plot_venn_sum.py)\n",
        "\n",
        "from collections import defaultdict, Counter\n",
        "from matplotlib_venn import venn3, venn3_circles\n",
        "import random\n",
        "from scipy import stats\n",
        "\n",
        "wdir = dir_path + 'prep_files/mpo/deseq/'\n",
        "odir = wdir + 'indivenn_hm/'\n",
        "data = pd.read_csv(wdir + 'resSig_compiled.txt', sep = '\\t')\n",
        "\n",
        "# Mercator\n",
        "### DICTIONARY OF MERCATOR BINS ###\n",
        "mfile = dir_path + 'mercator/MpoProt.results.txt'\n",
        "\n",
        "meranno = defaultdict(list)\n",
        "merbin = defaultdict(list)\n",
        "map2anno = {}\n",
        "\n",
        "merfile = open(mfile, 'r')\n",
        "merfile.readline()\n",
        "for line in merfile:\n",
        "\tlinecon = line.rstrip().replace(\"'\", \"\").split(\"\\t\")\n",
        "\tif len(linecon) == 5:\n",
        "\t\tbincode, name, identifier, desc, ptype = linecon\n",
        "\t\tmeranno[identifier].append(dicto[int(bincode.split('.')[0])])\n",
        "\t\tmerbin[identifier].append('.'.join(bincode.split('.')[:2]))\n",
        "\tif len(linecon[0].split('.')) == 2:\n",
        "\t\tmap2anno[linecon[0]] = linecon[1]\n",
        "\t\t\n",
        "\n",
        "all_s = [x.split(\"_\")[1] for x in data.stress.unique()]\n",
        "single = [x.split(\"_\")[1] for x in data.stress.unique() if len(x.split(\"_\")[1]) == 1]\n",
        "cross = [x.split(\"_\")[1] for x in data.stress.unique() if len(x.split(\"_\")[1]) == 2]\n",
        "\n",
        "data.annotation = data.annotation.apply(literal_eval)\n",
        "data[\"mername\"] = data.annotation.apply(lambda x: [dicto[int(y[0].split('.')[0])] for y in x]) # different from cell above, hence the repetitive code\n",
        "\n",
        "dict_A = defaultdict(list)\n",
        "dict_U = defaultdict(list)\n",
        "dict_D = defaultdict(list)\n",
        "\n",
        "def sum_to_dict(dicto, stress, reg):\n",
        "\tif reg == \"ALL\":\n",
        "\t\tsubset = data[(data.stress == \"Mpo_\" + stress)]\n",
        "\telse:\n",
        "\t\tsubset = data[(data.stress == \"Mpo_\" + stress) & (data.L2FC_D2 == reg)]\n",
        "\tdicto[stress].append(set(subset.gene.to_list()))\n",
        "\tdicto[stress].append([y for x in subset.mername.to_list() for y in x])\n",
        "\tdicto[stress].append(['.'.join(y[0].split('.')[:2]) for x in subset.annotation.to_list() for y in x])\n",
        "\n",
        "def dict_to_df(dicto):\n",
        "\tdf = pd.DataFrame.from_dict(dicto, orient='index', columns=[\"gene\", \"mername\", \"mapbin2\"])\n",
        "\treturn df\n",
        "\n",
        "for s in all_s:\n",
        "\tsum_to_dict(dict_A, s, \"ALL\")\n",
        "\tsum_to_dict(dict_U, s, \"UP\")\n",
        "\tsum_to_dict(dict_D, s, \"DOWN\")\n",
        "\n",
        "df_A = dict_to_df(dict_A)\n",
        "df_U = dict_to_df(dict_U)\n",
        "df_D = dict_to_df(dict_D)\n",
        "\n",
        "# =============================================================================\n",
        "# \n",
        "# # Summary of stress response\n",
        "# \n",
        "# =============================================================================\n",
        "\n",
        "# Q1 : ji_cal(a, b) [%]\n",
        "def ji_cal(a, b):\n",
        "\t# jaccard index calculation\n",
        "\treturn len(a&b) / len(a|b)\n",
        "# Q2: |(A − AB)/A − (B − AB)/B| [% difference]\n",
        "def suppInX(a, b, ab):\n",
        "\treturn len((a-ab))/len(a) - len((b-ab))/len(b)\n",
        "\n",
        "# Q3: (AB - A - B) / AB [%]\n",
        "def novel(a, b, ab):\n",
        "\treturn len(ab - a - b) / len(ab)\n",
        "\n",
        "def q_col(df, colnames):\n",
        "\t\"\"\"\n",
        "\tCollates the params for each cross stress and output in df\n",
        "\t\n",
        "\tParameters\n",
        "\t----------\n",
        "\tdf : dataframe\n",
        "\t\tdataframe to use (all genes, upreg/downreg only).\n",
        "\n",
        "\tReturns\n",
        "\t-------\n",
        "\tq_df : dataframe\n",
        "\t\tdatafram containing JI of all cross stress.\n",
        "\n",
        "\t\"\"\"\n",
        "\tq_dict = {}\n",
        "\tfor c in range(len(cross)):\n",
        "\t\tst = cross[c]\n",
        "\t\ta = df.loc[st[0]].gene\n",
        "\t\tb = df.loc[st[1]].gene\n",
        "\t\tab = df.loc[st].gene\n",
        "\t\tq_dict[st] = [\n",
        "\t\t\tji_cal(a,b),\n",
        "\t\t\t#perXInAB(a,ab),\n",
        "\t\t\t#perXInAB(b,ab),\n",
        "\t\t\tsuppInX(a,b,ab),\n",
        "\t\t\t#suppInAB(a,b,ab),\n",
        "\t\t\tnovel(a,b,ab)\n",
        "\t\t\t\t]\n",
        "\tq_df = pd.DataFrame.from_dict(q_dict, orient=\"index\", columns = colnames)\n",
        "\treturn q_df\n",
        "\n",
        "q_colnames = [\"similarity\", \"suppression\", \"novel interaction\"]\n",
        "\n",
        "q_A, q_U, q_D = [q_col(df_A, q_colnames), q_col(df_U, q_colnames), q_col(df_D, q_colnames)]\n",
        "qdf_col = [[q_U, \"Upregulated DEGs\"], [q_D, \"Downregulated DEGs\"]]\n",
        "\n",
        "def plot_q_subplots(df, outerax):\n",
        "\tax = [axe[x] for x in outerax]\n",
        "\t#plt.suptitle(title, fontsize=14)\n",
        "\tfor i, axis in enumerate(ax):\n",
        "\t\tif q_colnames[i] == \"suppression\":\n",
        "\t\t\tsns.heatmap(df[q_colnames[i]].to_frame().transpose(), cmap='coolwarm', ax=axis)\n",
        "\t\telse:\n",
        "\t\t\tsns.heatmap(df[q_colnames[i]].to_frame().transpose(), cmap='Blues', ax=axis)\n",
        "\t\taxis.set_yticklabels([q_colnames[i]], rotation=0)\n",
        "\t\tcbar = axis.collections[0].colorbar\n",
        "\t\tminval = round(df[q_colnames[i]].min(),2)\n",
        "\t\tmaxval = round(df[q_colnames[i]].max(),2)\n",
        "\t\n",
        "\t\twhile round(minval*100,2) % 5 != 0:\n",
        "\t\t\tminval += 0.01\n",
        "\t\twhile round(maxval*100,2) % 5 != 0:\n",
        "\t\t\tmaxval -= 0.01\n",
        "\t\tcbar.set_ticks([minval, maxval])\n",
        "\n",
        "def plot_reg(df, outerax):\n",
        "\tlabels = []\n",
        "\tstatscol = []\n",
        "\tfor i, col in enumerate(q_colnames[:-1]):\n",
        "\t\tfor j, col2 in enumerate(q_colnames[i+1:]):\n",
        "\t\t\touterax.scatter(col, col2, data=df)\n",
        "\t\t\tm, c, r_value, p_value, std_err = stats.linregress(df[col], df[col2])\n",
        "\t\t\tstatscol.append([m, c, r_value, p_value, std_err])\n",
        "\t\t\t# m, c = np.polyfit(df[col], df[col2], 1)\n",
        "\t\t\touterax.plot(df[col], m*df[col] + c)\n",
        "\t\t\tlabels.append(col[:3] + ' v ' + col2[:3] + ' ($\\mathregular{R^{2}}$: '+str(round(r_value**2,1))+', p: ' + str('{:.2f}'.format(round(p_value,2))+')'))\n",
        "\touterax.legend(labels, fontsize=\"x-small\")\n",
        "\treturn statscol\n",
        "\t\n",
        "def dum_venn(a_b, c_a, c_b, c_ab, ax, col, title, ac=20, bc=20, cc=20):\n",
        "\t'''\n",
        "\n",
        "\tParameters\n",
        "\t----------\n",
        "\ta_b : int\n",
        "\t\tSize of A&B.\n",
        "\tc_a : int\n",
        "\t\tSize of A&C-B.\n",
        "\tc_b : int\n",
        "\t\tSize of B&C-A.\n",
        "\tc_ab : int\n",
        "\t\tSzie of C&(A&B).\n",
        "\tax : axis handle\n",
        "\t\tAxis handle of subplot to plot into.\n",
        "\tcol : list\n",
        "\t\tList containing lists of patch id and corresponding colour.\n",
        "\ttitle:\n",
        "\t\t\n",
        "\tac : int, optional\n",
        "\t\tSize of set a. The default is 20.\n",
        "\tbc : int, optional\n",
        "\t\tSize of set b. The default is 20.\n",
        "\tcc : int, optional\n",
        "\t\tSize of set c. The default is 20.\n",
        "\n",
        "\tReturns\n",
        "\t-------\n",
        "\tNone.\n",
        "\n",
        "\t'''\n",
        "\t\n",
        "# =============================================================================\n",
        "# \ta_b = 8 # A&B\n",
        "# \tc_a_b = 4 # C&A-B/ C&B-A\n",
        "# \tc_ab = 3   # C&(A&B)\n",
        "# =============================================================================\n",
        "\tdum = list(string.ascii_uppercase + string.ascii_lowercase)\n",
        "\trandom.shuffle(dum)\n",
        "\ta = set(dum[:20])\n",
        "\tb = set(list(a)[:a_b] + [x for x in dum if x not in a][:bc-a_b])\n",
        "\tc = set(list(a-b)[:c_a] +\n",
        "\t\t list(b-a)[:c_b] + list(a&b)[:c_ab] +\n",
        "\t\t [x for x in dum if x not in a and x not in b][:cc-c_a-c_b-c_ab])\n",
        "\n",
        "\tv = venn3([a, b, c],\n",
        "\t   ('A', 'B', 'AB'),\n",
        "\t   ax = ax) # ax = axis\n",
        "\tvenn3_circles([a, b, c], linewidth=1, color='k', ax=ax)\n",
        "\tfor i in col:\n",
        "\t\tv.get_patch_by_id(i[0]).set_color(i[1])\n",
        "\tfor idx, subset in enumerate(v.subset_labels):\n",
        "\t\tv.subset_labels[idx].set_visible(False)\n",
        "\tax.set_title(title, fontsize=16)\n",
        "\n",
        "# =============================================================================\n",
        "# #\n",
        "# # Initialising subplot\n",
        "# #\n",
        "# =============================================================================\n",
        "#figsize=(figw, figh)\n",
        "\n",
        "top_mosaic = [[\"v1\", \"v2\", \"v3\"]]\n",
        "eq_mosaic = [\t[\"e1\", \"e2\", \"e3\"]\t]\n",
        "middle_mosaic = [\n",
        "\t[\"u1\", \"d1\"],\n",
        "\t[\"u2\", \"d2\"],\n",
        "\t[\"u3\", \"d3\"]\n",
        "]\n",
        "bottom_mosaic = [[\"r1\", \"r2\"]]\n",
        "\n",
        "figw, figh = 11, 9\n",
        "fig = plt.figure(figsize=(figw, figh))\n",
        "axc = fig.subplot_mosaic(\n",
        "\ttop_mosaic,\n",
        "\tgridspec_kw={\n",
        "\t\t\"bottom\": 0.75,\n",
        "\t\t\"top\": 1,\n",
        "\t\t#\"wspace\": 0.5,\n",
        "\t\t#\"hspace\": 0.5,\n",
        "\t\t}\n",
        "\t)\n",
        "axd = fig.subplot_mosaic(\n",
        "\teq_mosaic,\n",
        "\tgridspec_kw={\n",
        "\t\t\"bottom\": 0.55,\n",
        "\t\t\"top\": 0.8,\n",
        "\t\t#\"wspace\": 0.5,\n",
        "\t\t#\"hspace\": 0.5,\n",
        "\t\t}\n",
        "\t)\n",
        "\n",
        "axe = fig.subplot_mosaic(\n",
        "\tmiddle_mosaic,\n",
        "\tgridspec_kw={\n",
        "\t\t\"bottom\": 0.38,\n",
        "\t\t\"top\": 0.6,\n",
        "\t\t#\"wspace\": 0.5,\n",
        "\t\t\"hspace\": 0.2,\n",
        "\t\t}\n",
        "\t)\n",
        "axf = fig.subplot_mosaic(\n",
        "\tbottom_mosaic,\n",
        "\tgridspec_kw={\n",
        "\t\t\"bottom\": 0,\n",
        "\t\t\"top\": 0.3,\n",
        "\t\t#\"wspace\": 0.5,\n",
        "\t\t#\"hspace\": 0.5,\n",
        "\t\t}\n",
        "\t)\n",
        "\n",
        "for axy in ['e1', 'e2', 'e3']:\n",
        "\taxd[axy].axis('off')\n",
        "for axy in [\"v1\", \"v2\", \"v3\"]:\n",
        "\taxc[axy].set_anchor('N')\n",
        "\n",
        "axd['e1'].text(0.39, 0.45, r\"$\\frac{A \\cap B}{A \\cup B}$\", fontsize=20)\n",
        "axd['e2'].text(0.04, 0.45, r\"$\\frac{A-B-AB}{A}-\\frac{B-A-AB}{B}$\", fontsize=20)\n",
        "axd['e3'].text(0.29, 0.45, r\"$\\frac{AB-A-B}{AB}$\", fontsize=20)\n",
        "\n",
        "for seq, x in enumerate(qdf_col):\n",
        "\tdf_type, title = x[0], x[1]\n",
        "\tplot_q_subplots(df_type, [x[seq] for x in middle_mosaic])\n",
        "for axy in ['u1', 'u2', 'd1', 'd2']:\n",
        "\taxe[axy].set_xticklabels([])\n",
        "\taxe[axy].xaxis.set_visible(False)\n",
        "for axy in ['d1', 'd2', 'd3']:\n",
        "\taxe[axy].set_yticklabels([])\n",
        "\taxe[axy].yaxis.set_visible(False)\n",
        "\n",
        "v1col = [['100', 'white'], ['110', 'limegreen'],\n",
        "\t\t ['101', 'white'], ['111', 'limegreen'],\n",
        "\t\t ['010', 'white'], ['011', 'white'],\n",
        "\t\t ['001', 'white']]\n",
        "v2col = [['100', 'red'], ['110', 'white'],\n",
        "\t\t ['101', 'white'], ['111', 'white'],\n",
        "\t\t ['010', 'cornflowerblue'], ['011', 'white'],\n",
        "\t\t ['001', 'white']]\n",
        "v3col = [['100', 'white'], ['110', 'white'],\n",
        "\t\t ['101', 'white'], ['111', 'white'],\n",
        "\t\t ['010', 'white'], ['011', 'white'],\n",
        "\t\t ['001', 'darkorchid']]\n",
        "dum_venn(a_b=8,c_a=4, c_b = 4,c_ab=3,\n",
        "\t\t ax=axc['v1'], col=v1col, title='Similarity')\n",
        "dum_venn(a_b=8, c_a=10, c_b=6, c_ab=3,\n",
        "\t\t ax=axc['v2'], col=v2col, title='Suppression')\n",
        "dum_venn(a_b=8, c_a=4, c_b=4, c_ab=3,\n",
        "\t\t ax=axc['v3'], col=v3col, title='Novel interaction')\n",
        "\n",
        "reg_stats = []\n",
        "for i, x in enumerate(qdf_col):\n",
        "\tdf_type, title = x[0], x[1]\n",
        "\tdf_abs = df_type[:]\n",
        "\tdf_abs.suppression = abs(df_abs.suppression)\n",
        "\treg_stats.append(plot_reg(df_abs, axf[bottom_mosaic[0][i]]))\n",
        "\n",
        "plt.savefig(dir_path+'figures/Fig5A_E.png', dpi=600)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zd8VylDwyLWn"
      },
      "source": [
        "# Fig 5F refer to l2_en_jaccard_hm.png\n",
        "# =============================================================================\n",
        "# \n",
        "#  Summary: What is the dominant effect of each stress?\n",
        "# \n",
        "# =============================================================================\n",
        "def sum_df(df):\n",
        "\tdicto_sum = {}\n",
        "\tfor ss in single:\n",
        "\t\tori = [x for x in cross if ss in x]\n",
        "\t\trelcross = [y for x in cross if ss in x for y in x if ss not in y]\n",
        "\t\tsim = [ji_cal(df.loc[x[0]].gene, df.loc[x[1]].gene) for x in ori]\n",
        "\t\tnov = [novel(df.loc[x[0]].gene, df.loc[x[1]].gene, df.loc[x].gene) for x in ori]\n",
        "\t\tsup = [suppInX(df.loc[ss].gene, df.loc[x].gene, df.loc[ori[i]].gene) for i, x in enumerate(relcross)]\n",
        "\t\tfor i, x in enumerate(relcross):\n",
        "\t\t\tdicto_sum[ss+x] = [ss, sim[i], sup[i], nov[i]]\n",
        "\t\t\t\n",
        "\tdf_sum= pd.DataFrame.from_dict(dicto_sum, orient='index', columns=['stress', 'similarity', 'suppression', 'novel'])\n",
        "\treturn df_sum\n",
        "\n",
        "cond=['similarity', 'suppression', 'novel']\n",
        "df_list = [sum_df(df_U), sum_df(df_D)]\n",
        "\n",
        "# figs, axs = plt.subplots(3,2,\n",
        "# \t\t\t\t\t\t sharex=True,\n",
        "# \t\t\t\t\t\t sharey='row',\n",
        "# \t\t\t\t\t\t constrained_layout=True,\n",
        "# \t\t\t\t\t\t figsize=(7,6))\n",
        "\n",
        "# for i, x in enumerate(df_list):\n",
        "# \tfor j, y in enumerate(cond):\n",
        "# \t\tsns.violinplot(x='stress', y= y, data=x, ax= axs[j][i])\n",
        "# for k in range(3):\n",
        "# \taxs[k][1].set_ylabel('')\n",
        "# for l in range(2):\n",
        "# \tfor m in range(2):\n",
        "# \t\taxs[l][m].set_xlabel('')\n",
        "# axs[0][0].set_title('Upregulated', fontsize=14)\n",
        "# axs[0][1].set_title('Downregulated', fontsize=14)\n",
        "# plt.savefig(odir + 'venn_sum.png', dpi=600)\n",
        "\n",
        "# =============================================================================\n",
        "# \n",
        "# Enrichment\n",
        "# \n",
        "# =============================================================================\n",
        "\n",
        "from statsmodels.stats.multitest import multipletests\n",
        "import math\n",
        "import numpy as np\n",
        "\n",
        "ori_count = Counter([y for x in list(meranno.values()) for y in x])\n",
        "mapbins = list(dicto.values())\n",
        "\n",
        "def sig_df(df, sigcol, merdict, mapbins):\n",
        "\t\"\"\"\n",
        "\tCalculates and correct mapman bin enrichment p-value for all stresses\n",
        "\tReturns dataframe\n",
        "\n",
        "\tParameters\n",
        "\t----------\n",
        "\tdf : dataframe\n",
        "\t\tdf containing genes and corresponding mapman bins of DEGs.\n",
        "\tsigcol : str\n",
        "\t\tcolumn name to use for enrichment\n",
        "\tmerdict : dict\n",
        "\t\tcorresponding dictionary of mapman annotation/ 2nd level bins to use\n",
        "\tmapbins : list\n",
        "\t\tlist of mapman annotation/bins to use\n",
        "\n",
        "\tReturns\n",
        "\t-------\n",
        "\tdf_sig : dataframe\n",
        "\t\tdf summarising enrichment (corrected p-value) for each mapman bin (row)\n",
        "\t\tand each stress (column).\n",
        "\n",
        "\t\"\"\"\n",
        "\tsig_sum = {}\n",
        "\tfor s in all_s:\n",
        "\t\ts_count = Counter(df.loc[s][sigcol])\n",
        "\t\tvalid_bins = list(s_count.keys()) # bins found in stress\n",
        "\t\t# initialise count dicitonary\n",
        "\t\tsig_count = {}\n",
        "\t\tfor key in valid_bins:\n",
        "\t\t\tsig_count[key] = 1\n",
        "\t\t# random simulations\n",
        "\t\tfor i in range(1000):\n",
        "\t\t\tshuffle = list(merdict.values())\n",
        "\t\t\trandom.shuffle(shuffle)\n",
        "\t\t\tsub = shuffle[:len(df.loc[s].gene)]\n",
        "\t\t\tsub_count = Counter([y for x in sub for y in x])\n",
        "\t\t\tfor mapman in valid_bins:\n",
        "\t\t\t\tif sub_count[mapman] >= s_count[mapman]:\n",
        "\t\t\t\t\tsig_count[mapman] += 1\n",
        "\t\t# p-value calculation\n",
        "\t\tpval_coll = []\n",
        "\t\tfor mapman in valid_bins:\n",
        "\t\t\tpval = sig_count[mapman]/1000\n",
        "\t\t\t# correction for pval > 1\n",
        "\t\t\tif pval <= 1:\n",
        "\t\t\t\tpval_coll.append(pval)\n",
        "\t\t\telse:\n",
        "\t\t\t\tpval_coll.append(float(round(pval)))\n",
        "\t\t\n",
        "\t\t# BH correction for multiple testing\n",
        "\t\ty = multipletests(pvals=pval_coll, alpha=0.05, method=\"fdr_bh\")[1]\n",
        "\t\tall_bins_corr_pval = []\n",
        "\t\tfor mapman in mapbins:\n",
        "\t\t\tif mapman in valid_bins:\n",
        "\t\t\t\tall_bins_corr_pval.append(y[valid_bins.index(mapman)])\n",
        "\t\t\telse:\n",
        "\t\t\t\tall_bins_corr_pval.append(None)\n",
        "\t\tsig_sum[s] = all_bins_corr_pval\n",
        "\n",
        "\tdf_sig = pd.DataFrame.from_dict(sig_sum, orient='index', columns=mapbins)\n",
        "\treturn df_sig\n",
        "\t\n",
        "def chunk(uval, dval):\n",
        "\tif math.isnan(uval) and math.isnan(dval):\n",
        "\t\t# not differentially regulated\n",
        "\t\tcat = 0\n",
        "\telif  uval >= 0.05 and (dval >= 0.05 or math.isnan(dval)):\n",
        "\t\t# not enriched\n",
        "\t\tcat = 0\n",
        "\telif  dval >= 0.05 and (uval >= 0.05 or math.isnan(uval)):\n",
        "\t\t# not enriched\n",
        "\t\tcat = 0\n",
        "\telif  uval < 0.05 and dval < 0.05:\n",
        "\t\t# differentially up and downregulated in bin\n",
        "\t\tcat = 2\n",
        "\telif dval < 0.05:\n",
        "\t\t# differentially downregualted\n",
        "\t\tcat = 1\n",
        "\telif uval < 0.05:\n",
        "\t\t# differentially upregulated\n",
        "\t\tcat = 3\n",
        "\treturn cat\n",
        "\n",
        "from matplotlib.colors import ListedColormap\n",
        "cmap = ListedColormap([\"lightgray\", \"royalblue\", \"violet\", \"firebrick\"])\n",
        "catno = 4\n",
        "cbarticks = [(x/(catno*2))*(catno-1) for x in range(1,catno*2,2)]\n",
        "\n",
        "# =============================================================================\n",
        "# \n",
        "# Enrichment (Part 2: 2nd level Mapman)\n",
        "# \n",
        "# =============================================================================\n",
        "mapbins2 = list(set([y for x in list(merbin.values()) for y in x]))\n",
        "mapbins2.sort(key=lambda x: (int(x.split('.')[0]), int(x.split('.')[1])))\n",
        "\n",
        "df_sig_U2 = sig_df(df_U, 'mapbin2', merbin, mapbins2)\n",
        "df_sig_D2 = sig_df(df_D, 'mapbin2', merbin, mapbins2)\n",
        "df_sig_U2 = df_sig_U2.fillna(value=np.nan)\n",
        "df_sig_D2 = df_sig_D2.fillna(value=np.nan)\n",
        "\n",
        "cat_dict2 = {}\n",
        "for mapman in list(df_sig_U2.columns):\n",
        "\tcat_col = []\n",
        "\tfor stress in list(df_sig_U2.index):\n",
        "\t\tuval, dval = df_sig_U2.loc[stress, mapman], df_sig_D2.loc[stress, mapman]\n",
        "\t\tcat_col.append(chunk(uval,dval))\n",
        "\tcat_dict2[mapman] = cat_col\n",
        "\n",
        "df_combined_sig2 = pd.DataFrame.from_dict(cat_dict2, orient='index', columns=list(df_sig_U2.index))\n",
        "df_combined_sig2 = df_combined_sig2.loc[df_combined_sig2.max(axis=1) > 0,:]\n",
        "df_combined_sig2 = df_combined_sig2.loc[(df_combined_sig2 > 0).sum(axis=1) >2,:]\n",
        "df_combined_sig2.reset_index(inplace=True)\n",
        "df_combined_sig2['index'] = df_combined_sig2['index'].apply(lambda x: map2anno[x])\n",
        "df_combined_sig2.set_index('index', inplace=True)\n",
        "\n",
        "# =============================================================================\n",
        "# \n",
        "# Plotting 2nd level mapman enrichment (df_combined_sig2)\n",
        "# \n",
        "# =============================================================================\n",
        "from scipy.cluster.hierarchy import dendrogram, linkage\n",
        "from scipy.spatial.distance import squareform\n",
        "from mpl_toolkits.axes_grid1.inset_locator import inset_axes\n",
        "\n",
        "def jdistprep(df, axis):\n",
        "\t'''\n",
        "\tConvert df to sets (for calculation of JD of X axis)\n",
        "\n",
        "\tParameters\n",
        "\t----------\n",
        "\tdf : dataframe\n",
        "\t\tdataframe of categorical variables to be converted to sets.\n",
        "\taxis : int \n",
        "\t\taxis to do sets on, 0 by column (default), 1 by row\n",
        "\n",
        "\tReturns\n",
        "\t-------\n",
        "\tdicto : dict\n",
        "\t\tdictionary containing list of column values.\n",
        "\n",
        "\t'''\n",
        "\tif axis == 1:\n",
        "\t\tdf = df.T\n",
        "\tdxkeys = df.columns.to_list()\n",
        "\tdykeys = df.index.to_list()\n",
        "\tdicto = {}\n",
        "\tfor col in dxkeys:\n",
        "\t\tdicto[col] = [dykeys[i] + '_' + str(x) for i, x in enumerate(df[col].to_list())]\n",
        "\treturn [dicto, dxkeys]\n",
        "\n",
        "def jdist(df, axis=0):\n",
        "\t'''\n",
        "\tConstruct jaccard distance square matrix\n",
        "\n",
        "\tParameters\n",
        "\t----------\n",
        "\tdf : df\n",
        "\t\tdataframe to be used for jiprep/ jdist calculation.\n",
        "\taxis : int\n",
        "\t\taxis to do sets on, 0 by column (default), 1 by row\n",
        "\n",
        "\tReturns\n",
        "\t-------\n",
        "\tlinkage_matrix : list\n",
        "\t\tcondensed jaccard distance matrix.\n",
        "\tjlist : list\n",
        "\t\tlist of list (jaccard distance square matrix)\n",
        "\tdicto : dict\n",
        "\t\tdictionary of list\n",
        "\n",
        "\t'''\n",
        "\tdicto, dxkeys = jdistprep(df, axis)\n",
        "\tjlist = []\n",
        "\tfor key in dxkeys:\n",
        "\t\tcol = []\n",
        "\t\tfor key2 in dxkeys:\n",
        "\t\t\tset1, set2 = dicto[key], dicto[key2]\n",
        "\t\t\tset1x = set([x for x in set1 if x.split('_')[1] != '0'])\n",
        "\t\t\tset2x = set([x for x in set2 if x.split('_')[1] != '0'])\n",
        "\t\t\tcol.append(1 - ji_cal(set1x, set2x))\n",
        "\t\tjlist.append(col)\n",
        "\tdists = squareform(jlist)\n",
        "\tlinkage_matrix = linkage(dists, \"single\")\n",
        "\treturn linkage_matrix, jlist, dicto\n",
        "\n",
        "def plot_dendro(linkage_matrix, ax, orient):\n",
        "\t'''\n",
        "\tPlots dendrogram into subplot\n",
        "\n",
        "\tParameters\n",
        "\t----------\n",
        "\tmat : list of lists\n",
        "\t\tContains the square matrix of jaccard distances.\n",
        "\tax : axes\n",
        "\t\taxis of subplot to plot to.\n",
        "\torient : str\n",
        "\t\torientation of dendrogram to be plotted.\n",
        "\n",
        "\tReturns\n",
        "\t-------\n",
        "\tNone.\n",
        "\n",
        "\t'''\n",
        "\t\n",
        "\tdendrogram(linkage_matrix, no_labels=True, ax=ax, orientation=orient, color_threshold=0, above_threshold_color='#000000')\n",
        "\n",
        "xmat, xlist, xdict = jdist(df_combined_sig2)\n",
        "ymat, ylist , ydict = jdist(df_combined_sig2, axis=1)\n",
        "\n",
        "yden = dendrogram(ymat, labels=df_combined_sig2.index.to_list(), orientation='left') #, color_threshold=0, above_threshold_color='#000000'\n",
        "plt.show()\n",
        "xden = dendrogram(xmat, labels=df_combined_sig2.columns.to_list(), orientation='top') #, color_threshold=0, above_threshold_color='#000000'\n",
        "plt.show()\n",
        "\n",
        "yorder = yden['ivl']\n",
        "xorder = xden['ivl']\n",
        "df_sig2_reordered = df_combined_sig2[xorder]\n",
        "df_sig2_reordered = df_sig2_reordered.reindex(yorder[::-1])\n",
        "\n",
        "fig, ax = plt.subplots(2,2,\n",
        "\t\t\t\t\t   figsize=(7.5,8.5), # (width, height)\n",
        "\t\t\t\t\t   constrained_layout=True,\n",
        "\t\t\t\t\t   gridspec_kw={'width_ratios': [1.5, 5],'height_ratios': [1, 5]}) # constrained_layout=True,\n",
        "ax0, ax1, ax2, ax3 = ax.flatten()\n",
        "for i in [ax0, ax1, ax2]:\n",
        "\ti.axis('off')\n",
        "\n",
        "plot_dendro(xmat, ax1, 'top')\n",
        "plot_dendro(ymat, ax2, 'left')\n",
        "# heatmap, tick and tick labels\n",
        "hplot = ax3.imshow(df_sig2_reordered, cmap=cmap)\n",
        "ax3.yaxis.tick_right()\n",
        "ax3.set_ylabel(\"\")\n",
        "ax3.set_xticks(np.arange(0, len(df_sig2_reordered.columns), 1))\n",
        "ax3.set_yticks(np.arange(0, len(df_sig2_reordered), 1))\n",
        "\n",
        "\n",
        "xcolour = ['k'] + ['firebrick']*4 + ['gray']*6 + ['mediumseagreen']*4 + ['k']*2 + ['darkorange']*3 + ['k']*2 + ['royalblue']*3\n",
        "ax3.set_xticklabels(df_sig2_reordered.columns.to_list(), rotation=90)\n",
        "\n",
        "for i, tick_label in enumerate(ax3.get_xticklabels()):\n",
        "\ttick_text = tick_label.get_text()\n",
        "\ttick_label.set_color(xcolour[i])\n",
        "\t\n",
        "anno_long = ['annotated', 'cellulose', 'biosynthesis', 'hemicellulose', 'pectin', 'channels', 'degradation']\n",
        "ax3.set_yticklabels([(lambda x: x.split('.')[1].lower() if x.split('.')[1] not in anno_long else x.lower())(x) for x in df_sig2_reordered.index.to_list()])\n",
        "# cbar plotting and control\n",
        "axins = inset_axes(ax0,\n",
        "\t\t\t\t\twidth=\"40%\",  # width = 50% of parent_bbox width\n",
        "\t\t\t\t\theight=\"90%\",  # height : 5%\n",
        "\t\t\t\t\tloc = 'center')\n",
        "cbar = fig.colorbar(hplot, cax=axins, ticks = cbarticks)\n",
        "cbar.ax.set_yticklabels(['N', 'D', 'UD', 'U'])\n",
        "\n",
        "plt.savefig(dir_path+'figures/Fig5F.png', dpi=600, bbox_inches='tight') # no N"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xuuwbZc27lDH"
      },
      "source": [
        "### Figure 6: Diurnal gene expression"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GAgDmV7kdQXn"
      },
      "source": [
        "#Fig 6 A to D, adapted from Mpo_panel1.py\n",
        "from scipy.stats import zscore\n",
        "\n",
        "wdir = dir_path + \"diurnal/\"\n",
        "Mpodf = pd.read_csv(wdir + \"Mpo_supp.txt\", sep = \"\\t\", index_col = 0)\n",
        "\n",
        "Mpo_exp_only = Mpodf[Mpodf.LAG != \"NE\"]\n",
        "Mpo_rhy_only = Mpo_exp_only[Mpo_exp_only.LAG != \"NR\"]\n",
        "\n",
        "perall = (len(Mpo_rhy_only) / len(Mpodf))*100\n",
        "perexp = (len(Mpo_rhy_only) / len(Mpo_exp_only))*100\n",
        "\n",
        "# Subset rhythmic genes only\n",
        "Mpo_rhy_only.LAG = Mpo_rhy_only.LAG.astype(int)\n",
        "Mpo_rhy_only[\"ADJ.P\"] = Mpo_rhy_only[\"ADJ.P\"].astype(np.float)\n",
        "Mpo_rhy_only.sort_values([\"LAG\", \"ADJ.P\"], inplace=True)\n",
        "\n",
        "# normalisation of rhythmic gene expression\n",
        "rhy_zscore = Mpo_rhy_only[Mpo_rhy_only.columns.to_list()[Mpo_rhy_only.columns.to_list().index(\"ZT2_1\"):]]\n",
        "rhy_zscore = rhy_zscore.transpose()\n",
        "rhy_zscore = rhy_zscore.apply(zscore)\n",
        "rhy_zscore = rhy_zscore.transpose()\n",
        "plt.figure(figsize=(10,20))\n",
        "\n",
        "\n",
        "\"\"\"\n",
        "Panel 1a) plot\n",
        "\"\"\"\n",
        "colnames = [x.split(\"_\")[0] for x in rhy_zscore.columns.to_list()]\n",
        "condcoldict = {}\n",
        "for x in list(set(colnames)):\n",
        "\tif int(x.split(\"ZT\")[1]) < 12:\n",
        "\t\tcondcoldict[x] = (255,255,153) #\"khaki\"\n",
        "\telse:\n",
        "\t\tcondcoldict[x] = (160,160,160) #\"lightslategrey\"\n",
        "\t\t\n",
        "condcol = np.array([[condcoldict[x] for x in colnames]])\n",
        "\n",
        "fig, ax = plt.subplots(2,1,\n",
        "\t\t\t\t\t   figsize=(4,8), # (width, height)\n",
        "\t\t\t\t\t   gridspec_kw={'height_ratios': [0.091, 3.9]})\n",
        "fig.subplots_adjust(hspace=0.01)\n",
        "ax1, ax2= ax.flatten()\n",
        "\n",
        "ax1.imshow(condcol)\n",
        "# Set gridlines\n",
        "ax1.set_xticks(np.arange(-.5, 18, 3))\n",
        "ax1.set_yticks(np.arange(-.5, 1, 1))\n",
        "ax1.grid(color='k', linestyle='-', linewidth=1)\n",
        "ax1.set_xticklabels([])\n",
        "ax1.set_yticklabels([])\n",
        "ax1.xaxis.set_ticks_position('none')\n",
        "ax1.yaxis.set_ticks_position('none')\n",
        "ax1.set_anchor('W')\n",
        "\n",
        "sns.heatmap(rhy_zscore, cmap = \"coolwarm\", ax = ax2, yticklabels=False, xticklabels=True, center=0)\n",
        "ax2.set_ylabel(\"Genes\")\n",
        "plt.savefig(dir_path + 'figures/Fig6A.png', dpi = 600)\n",
        "plt.show()\n",
        "\n",
        "\"\"\"\n",
        "Panel 1B) plot\n",
        "\"\"\"\n",
        "rhycount = len(Mpo_rhy_only)\n",
        "LAGcount = Mpo_rhy_only.groupby(\"LAG\").count().annotation.to_frame()\n",
        "LAGcount.columns = [\"count\"]\n",
        "LAGcount[\"percent\"] = LAGcount.apply(lambda x: (x/rhycount)*100)\n",
        "g = LAGcount.percent.plot(xticks = LAGcount.index.to_list(),\n",
        "\t\t\t\t\t\t  yticks = [0, 18],\n",
        "\t\t\t\t\t\t  ylim = [0,18],\n",
        "\t\t\t\t\t\t  ylabel = \"% rhythmic genes\",\n",
        "\t\t\t\t\t\t  xlabel = \"Phase\",\n",
        "\t\t\t\t\t\t  color = \"k\")\n",
        "g.axvline(12, color = \"k\")\n",
        "plt.savefig(dir_path + 'figures/Fig6B.png', dpi = 600)\n",
        "plt.show()\n",
        "\n",
        "\"\"\"\n",
        "Panel 1C) plot\n",
        "\"\"\"\n",
        "fp = dir_path + \"diurnal/Ferrari_2019/SD14_compat.txt\"\n",
        "qp = dir_path + \"diurnal/Ferrari_2019/OF_20210623_compat.tsv\"\n",
        "camortho = pd.read_csv(fp, sep=\"\\t\", index_col=0)\n",
        "camortho.Ath = camortho.Ath.str.upper()\n",
        "qortho = pd.read_csv(qp, sep=\"\\t\", index_col=0)\n",
        "qortho.Osa = qortho.Osa.str.replace(\"\\.[0-9]*\", \"\", regex=True)\n",
        "\n",
        "camgrps = [\"OG0000156\", \"OG0000215\", \"OG0000679\", \"OG0004739\", \"OG0004944\", \"OG0005370\"]\n",
        "qgrps = [\"OG0000167\", \"OG0000301\", \"OG0000399\", \"OG0003516\", \"OG0004855\", \"OG0004502\"]\n",
        "\n",
        "cgrpdict = {\"OG0000156\":\"OG0000167 (Cyclin A, B)\",\n",
        "\t\t\t\"OG0000215\":\"OG0000301 (Cyclin D)\",\n",
        "\t\t\t\"OG0000679\":\"OG0000399 (CDK)\",\n",
        "\t\t\t\"OG0004739\":\"OG0003516 (Timeless)\",\n",
        "\t\t\t\"OG0004944\":\"OG0004855 (DNA primase)\",\n",
        "\t\t\t\"OG0005370\":\"OG0004502 (DNA polymerase)\"}\n",
        "\n",
        "mpogrpgenes = {}\n",
        "for i in range(len(camgrps)):\n",
        "\ttestc = camortho.loc[camgrps[i],:].to_list()\n",
        "\ttestc = [x.split(\", \") for x in testc if type(x) != float]\n",
        "\ttestcs = [x for a in testc for x in a]\n",
        "\ttestq = qortho.loc[qgrps[i],:].to_list()\n",
        "\ttestq = [x.split(\", \") for x in testq if type(x) != float]\n",
        "\ttestqs = [x for a in testq for x in a]\n",
        "\tc_s = set(testcs) - set(testqs)\n",
        "\tprint(camgrps[i])\n",
        "\tprint(str(list(c_s)))\n",
        "\tq_s = list(set(testqs) - set(testcs))\n",
        "\tmpogrpgenes[cgrpdict[camgrps[i]]] = [x for x in q_s if \"Mp\" in x]\n",
        "\t\n",
        "mpocyclegenes = [x for a in list(mpogrpgenes.values()) for x in a if x in rhy_zscore.index]\n",
        "timepoints = [\"ZT2\", \"ZT6\", \"ZT10\", \"ZT14\", \"ZT18\", \"ZT22\"]\n",
        "# dataframe normalised timepoints for mpocyclegenes only\n",
        "cycle_zscore = rhy_zscore.loc[mpocyclegenes].transpose()\n",
        "# dictionary of mpocyclegenes and their corresponding OG information\n",
        "cycle_grp = {}\n",
        "for k, v in mpogrpgenes.items():\n",
        "\tfor item in v:\n",
        "\t\tcycle_grp[item] = k\n",
        "qwOGgrps = list(mpogrpgenes.keys())\n",
        "# to transpose and create new df that contains the average zscore of replicates\n",
        "cycle_dict = {}\n",
        "for t in timepoints:\n",
        "\tcycle_dict[t] = cycle_zscore.loc[[t+\"_1\", t+\"_2\", t+\"_3\"],:].mean().to_list()\n",
        "cycle_df = pd.DataFrame(cycle_dict, columns = timepoints, index = [x + \": \" + cycle_grp[x] for x in cycle_zscore.columns.to_list()]).transpose()\n",
        "\n",
        "colours = [\"maroon\", \"orangered\", \"forestgreen\", \"midnightblue\", \"mediumorchid\", \"steelblue\", \"darkseagreen\"]\n",
        "\n",
        "cycle_cols = {}\n",
        "for keys in cycle_grp.keys():\n",
        "\tcycle_cols[keys + \": \" + cycle_grp[keys]] = colours[qwOGgrps.index(cycle_grp[keys])]\n",
        "\n",
        "cycle_df.plot().legend(bbox_to_anchor=(0.81, -0.1))\n",
        "cycle_df.plot(color = cycle_cols).legend(bbox_to_anchor=(0.81, -0.1))\n",
        "plt.savefig(dir_path + 'figures/Fig6C.png',\n",
        "\t\t\tdpi = 600,\n",
        "\t\t\tbbox_inches='tight')\n",
        "\n",
        "\"\"\"\n",
        "Panel 1D) Mercator by phase\n",
        "\"\"\"\n",
        "import seaborn as sns\n",
        "\n",
        "Mpo_rhy_only[\"MapMan bins\"] = Mpo_rhy_only.apply(lambda x: x.annotation.split(\".\")[0].capitalize(), axis=1)\n",
        "mer_grouped = Mpo_rhy_only.groupby([\"MapMan bins\", \"LAG\"]).count().annotation.unstack(fill_value=0)\n",
        "phases = mer_grouped.columns.to_list()\n",
        "binsum = mer_grouped.sum(axis=1)\n",
        "for phase in phases:\n",
        "\tmer_grouped[phase] = mer_grouped[phase]/binsum\n",
        "g = sns.clustermap(mer_grouped)\n",
        "plt.show()\n",
        "g_ytick = [str(x).split(\"'\")[1] for x in g.ax_heatmap.get_yticklabels()]\n",
        "drow = g.dendrogram_row.linkage\n",
        "mer_grouped_reordered = mer_grouped.reindex(g_ytick)\n",
        "sns.heatmap(mer_grouped_reordered, yticklabels=True)\n",
        "mer_grouped.transpose().plot().legend(bbox_to_anchor=(0.72, -0.1))\n",
        "\n",
        "# subplots with linkage\n",
        "mer_grouped_dendro = mer_grouped.reindex(g_ytick[::-1])\n",
        "from scipy.cluster.hierarchy import dendrogram\n",
        "figii, axii = plt.subplots(1,2,\n",
        "\t\t\t\t\t   figsize=(10,6), # (width, height\n",
        "\t\t\t\t\t   constrained_layout=True,\n",
        "\t\t\t\t\t   gridspec_kw={'width_ratios': [1.9, 8.1]})\n",
        "ax1ii, ax2ii= axii.flatten()\n",
        "ax1ii.axis(\"off\")\n",
        "dendrogram(drow, no_labels=True, ax=ax1ii, orientation='left', color_threshold=0, above_threshold_color='#000000')\n",
        "sns.heatmap(mer_grouped_dendro, yticklabels=True, ax = ax2ii)\n",
        "ax2ii.set_ylabel(\"\")\n",
        "plt.show()\n",
        "\n",
        "# to plot heatmap by chunks\n",
        "def chunk(num):\n",
        "\tif  num == 0:\n",
        "\t\tcat = 0\n",
        "\telif num < 0.1:\n",
        "\t\tcat = 1\n",
        "\telif num < 0.2:\n",
        "\t\tcat = 2\n",
        "\telif num < 0.3:\n",
        "\t\tcat = 3\n",
        "\telif num < 0.4:\n",
        "\t\tcat = 4\n",
        "\telse:\n",
        "\t\tcat = 5\n",
        "\treturn cat\n",
        "\n",
        "mer_grouped_chunk = mer_grouped.reindex(g_ytick)\n",
        "for col in mer_grouped_chunk:\n",
        "\tmer_grouped_chunk[col] = mer_grouped_chunk[col].apply(lambda x: chunk(x))\n",
        "sns.heatmap(mer_grouped_chunk, yticklabels=True)\n",
        "\n",
        "\"\"\"\n",
        "Mercator count binned by percentage with custom colormap\n",
        "\"\"\"\n",
        "from matplotlib.colors import ListedColormap\n",
        "figm, (axm1, axm2) = plt.subplots(1,2,\n",
        "\t\t\t\t\t\t\t\t  figsize=(6.3,6), # (width, height\n",
        "\t\t\t\t\t\t\t\t  constrained_layout=True,\n",
        "\t\t\t\t\t\t\t\t  gridspec_kw={'width_ratios': [1.6, 8.4]})\n",
        "\n",
        "axm1.axis(\"off\")\n",
        "dendrogram(drow,\n",
        "\t\t   no_labels=True,\n",
        "\t\t   ax=axm1, orientation='left',\n",
        "\t\t   color_threshold=0,\n",
        "\t\t   above_threshold_color='#000000')\n",
        "cmap = ListedColormap([\"gray\", \"lightsteelblue\", \"lightgreen\", \"palegoldenrod\", \"coral\", \"indianred\"])\n",
        "mplot = axm2.imshow(mer_grouped_chunk, cmap=cmap, interpolation=\"none\")\n",
        "\n",
        "axm2.set_xticks(np.arange(0, len(mer_grouped_chunk.columns), 1))\n",
        "axm2.set_yticks(np.arange(0, len(mer_grouped_chunk), 1))\n",
        "axm2.set_xticklabels(mer_grouped_chunk.columns.to_list())\n",
        "axm2.set_yticklabels(mer_grouped_chunk.index.to_list())\n",
        "cbar = figm.colorbar(mplot,\n",
        "\t\t\t\t\t ax=axm2,\n",
        "\t\t\t\t\t ticks = [x/12 for x in np.arange(5,6*10,10)],\n",
        "\t\t\t\t\t label=\"% rhythmic genes in bin\")\n",
        "#cbar = fig.colorbar(cax, ticks=[-1, 0, 1])\n",
        "cbar.ax.set_yticklabels(['None', '0-9%', '10-19%', '20-29%', '30-39%', '>40%'])  # vertically oriented colorbar\n",
        "plt.savefig(dir_path + 'figures/Fig6D.png', dpi = 600)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DXuBGD65hn7R"
      },
      "source": [
        "# Fig 6E and F, adapted from 1to1ortho.py\n",
        "Mpo_exp_only = Mpodf[Mpodf.LAG != \"NE\"]\n",
        "Mpo_rhy_only = Mpo_exp_only[Mpo_exp_only.LAG != \"NR\"]\n",
        "Mpo_rhy_genes = Mpo_rhy_only.index.to_list()\n",
        "\n",
        "species = [\"Cpa\", \"Ppu\", \"Cre\", \"Kni\", \"Ppa\", \"Smo\", \"Pab\", \"Osa\", \"Ath\"]\n",
        "night = [8, 12, 12, 12, 8, 12, 8, 6, 6]\n",
        "daylength = [16, 12, 12, 12, 16, 12, 16, 6, 6]\n",
        "freq = [1, 1, 1, 1, 1, 1, 1, 2, 2]\n",
        "odir = dir_path + \"diurnal/Orthologues_Mpo/\"\n",
        "ofiles = [\"Mpo__v__\" + x + \".tsv\" for x in species]\n",
        "\n",
        "camdir = dir_path + \"diurnal/\"\n",
        "camfiles = [x + \"_supp.txt\" for x in species]\n",
        "\n",
        "### FUNCTION ###\n",
        "def lag_diff(a, b):\n",
        "\t\"\"\"\n",
        "\n",
        "\tParameters\n",
        "\t----------\n",
        "\ta : int\n",
        "\t\tLAG value of species X.\n",
        "\tb : int\n",
        "\t\tLAG value of Mpo.\n",
        "\n",
        "\tReturns\n",
        "\t-------\n",
        "\tdiff : int\n",
        "\t\tsmallest LAG diff.\n",
        "\n",
        "\t\"\"\"\n",
        "\tdiff = a - b\n",
        "\tif abs(diff) > 12:\n",
        "\t\tif diff < 0:\n",
        "\t\t\tdiff = diff + 24\n",
        "\t\telse:\n",
        "\t\t\tdiff = diff - 24\n",
        "\treturn diff\n",
        "### END ###\n",
        "\n",
        "f_axes = string.ascii_uppercase[:len(species)]\n",
        "d_axes = string.ascii_lowercase[:len(species)]\n",
        "axd = plt.figure(constrained_layout=True,\n",
        "\t\t\t\t figsize=(27,6)).subplot_mosaic(\n",
        "    \"\"\"\n",
        "\tabcdefghi\n",
        "\tABCDEFGHI\n",
        "    \"\"\"\n",
        ")\n",
        "\n",
        "#for spe in species:\n",
        "# get 1 to 1 orthologue\n",
        "for z in range(len(species)):\n",
        "\tol_df = pd.read_csv(odir + ofiles[z], sep=\"\\t\", index_col=0)\n",
        "\tol_df = ol_df[ol_df.Mpo.apply(lambda row: len(row.split(\", \")) ==1) & ol_df[species[z]].apply(lambda row: len(row.split(\", \")) ==1)]\n",
        "\tif species[z] == \"Osa\":\n",
        "\t\tosa_dict = {}\n",
        "\t\tosa_genes = ol_df.Osa.to_list()\n",
        "\t\tfor gene in osa_genes:\n",
        "\t\t\tosa_dict[gene] = gene.split(\".\")[0]\n",
        "\t\tol_df.Osa.replace(osa_dict, inplace=True)\t\n",
        "\t# get LAGs\n",
        "\tcam_f = pd.read_csv(camdir + camfiles[z], sep=\"\\t\", index_col=0)\n",
        "\t# Mpo LAG\n",
        "\tol_df[\"Mpo_LAG\"] = ol_df.apply(lambda row: Mpodf.LAG.loc[row.Mpo], axis=1)\n",
        "\tol_df[species[z] + \"_LAG\"] = ol_df.apply(lambda row: cam_f.phase.get(row[species[z]], None), axis=1)\n",
        "\t# exclude NE and NR in either Mpo or species[z] LAG   \n",
        "\tfor i in [\"Mpo\", species[z]]:\n",
        "\t\tfor j in [\"NE\", \"NR\"]:\n",
        "\t\t\tol_df = ol_df[ol_df[i + \"_LAG\"] != j]\n",
        "\t\tol_df = ol_df[ol_df[i + \"_LAG\"].notna()]\n",
        "\t# correcting LAG value 24 to 0 and converting to numeric\n",
        "\tol_df[species[z] + \"_LAG\"] = ol_df[species[z] + \"_LAG\"].replace({\"24\":\"0\"})\n",
        "\tol_df.Mpo_LAG = pd.to_numeric(ol_df.Mpo_LAG)\n",
        "\tol_df[species[z] + \"_LAG\"] = pd.to_numeric(ol_df[species[z] + \"_LAG\"])\n",
        "\t\n",
        "\t# calculating smallest lag diff\n",
        "\tol_df[\"LAG_diff\"] = ol_df.apply(lambda row: lag_diff(row[species[z] + \"_LAG\"], row.Mpo_LAG), axis=1)\n",
        "\tdiff_ser = ol_df.groupby(\"LAG_diff\").count().Mpo_LAG\n",
        "\tmax_diff = diff_ser[diff_ser == diff_ser.max()].index.to_list()\n",
        "\t\n",
        "\t# Plot LAG diff\n",
        "\taxd[f_axes[z]].set_xticks(np.arange(-.5,len(ol_df.LAG_diff.unique())-1))\n",
        "\tdiff_ser_index = diff_ser.index.to_list()\n",
        "\tdiff_xticks = []\n",
        "\tif len(diff_ser_index) > 13:\n",
        "\t\tfor i in range(len(diff_ser_index)):\n",
        "\t\t\tif i%2 == 0:\n",
        "\t\t\t\tdiff_xticks.append(str(diff_ser_index[i]))\n",
        "\t\t\telse:\n",
        "\t\t\t\tdiff_xticks.append(\"\")\n",
        "\telse:\n",
        "\t\tdiff_xticks = diff_ser_index\n",
        "\n",
        "\tsns.histplot(ol_df.LAG_diff,\n",
        "\t\t\t  #x=diff_xticks,\n",
        "\t\t\t  bins=len(ol_df.LAG_diff.unique()),\n",
        "\t\t\t  kde=True,\n",
        "\t\t\t  ax=axd[d_axes[z]],)\n",
        "\tif z != 0:\n",
        "\t\taxd[d_axes[z]].set_ylabel(\"\")\n",
        "\telif z == 0:\n",
        "\t\taxd[d_axes[z]].set_ylabel(\"Count\", fontsize=14)\n",
        "\taxd[d_axes[z]].set_xlabel(\"\")\n",
        "\n",
        "\t# Plot for 1 to 1 ortho\n",
        "\tmpo_tp = list(range(0,24,2)) # x-axis\n",
        "\tfull_other_tp = list(ol_df[species[z] + \"_LAG\"].unique())\n",
        "\tother_tp = list(range(0,24,freq[z])) # y-axis\n",
        "\tsum_dict = {}\n",
        "\tfor o in other_tp:\n",
        "\t\to_col = []\n",
        "\t\tfor m in mpo_tp:\n",
        "\t\t\to_col.append(sum((ol_df.Mpo_LAG == m) & (ol_df[species[z] + \"_LAG\"] == o)))\n",
        "\t\tsum_dict[\"ZT\" + str(o)] = o_col\n",
        "\tsum_df = pd.DataFrame(sum_dict, columns = [\"ZT\" + str(y) for y in other_tp], index = [\"ZT\" + str(x) for x in mpo_tp])\n",
        "\t\n",
        "\taxd[f_axes[z]].imshow(sum_df, cmap=\"Blues\", aspect=\"auto\")\n",
        "\taxd[f_axes[z]].set_xticks(np.arange(-.5,len(other_tp)-1))\n",
        "\taxd[f_axes[z]].set_yticks(np.arange(-.5,len(mpo_tp)-1))\n",
        "\txticklist = [other_tp[0]] + [\"\" for x in range(0,daylength[z]-1)] + [other_tp[daylength[z]]] + [\"\" for x in range(0,night[z]-2)] + [other_tp[-1]]\n",
        "\tyticklist = [mpo_tp[0]] + [\"\" for x in range(0,6-1)] + [mpo_tp[6]] + [\"\" for x in range(0,6-2)] + [mpo_tp[-1]]\n",
        "\taxd[f_axes[z]].set_xticklabels(xticklist)\n",
        "\taxd[f_axes[z]].set_yticklabels(yticklist)\n",
        "\taxd[f_axes[z]].axvline(daylength[z]-0.5, color=\"k\")\n",
        "\taxd[f_axes[z]].axhline(6-0.5, color=\"k\")\n",
        "\taxd[f_axes[z]].set_xlabel(species[z], fontsize=14)\n",
        "\tif z != 0:\n",
        "\t\taxd[f_axes[z]].set_yticks([])\n",
        "\t\taxd[f_axes[z]].set_yticklabels([])\n",
        "\telif z == 0:\n",
        "\t\taxd[f_axes[z]].set_ylabel(\"Mpo\", fontsize=14)\n",
        "plt.savefig(dir_path + 'figures/Fig6E_F.png', dpi = 600)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RdmYRmdc7oA4"
      },
      "source": [
        "### Supp. Fig 2: QC of RNA-seq data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZUjfF_uuYTGW"
      },
      "source": [
        "# adapted from QC_scaled_updated.py\n",
        "from sklearn.preprocessing import StandardScaler\n",
        "from scipy.stats import pearsonr\n",
        "\n",
        "o_dir = dir_path + 'figures/'\n",
        "sumdir = dir_path + 'summary_files/'\n",
        "\n",
        "expdesc = ['all_stress', 'diurnal_exp', 'single_stress']\n",
        "targetexp = expdesc[0]\n",
        "targetp = sumdir + targetexp + '.txt'\n",
        "expmatp = dir_path + 'prep_files/' + targetexp + '.tsv'\n",
        "exps = [x.split(\"\\t\")[0] for x in open(targetp, \"r\").readlines()]\n",
        "labels = [x.strip().split(\"\\t\")[1] + '_' + x.split(\"\\t\")[0].split('_')[1] for x in open(targetp, \"r\").readlines()]\n",
        "\n",
        "df = pd.read_csv(expmatp, index_col = 0, sep = \"\\t\", header = 0)\n",
        "df.columns = labels\n",
        "\n",
        "# Standard Scaling\n",
        "scaled_features = StandardScaler().fit_transform(df.values)\n",
        "df_scaled = pd.DataFrame(scaled_features, index = df.index, columns = df.columns)\n",
        "\n",
        "# plot cluster map\n",
        "sns.set(font_scale=1.6)\n",
        "\n",
        "methods = \"average\"\n",
        "\n",
        "g1 = sns.clustermap(df_scaled.corr(),\n",
        "\t\t\t\t method = methods,\n",
        "\t\t\t\t figsize=(20,20),\n",
        "\t\t\t\t xticklabels=True,\n",
        "\t\t\t\t yticklabels=True)\n",
        "plt.title(\"All stress (scaled): \" + methods)\n",
        "plt.savefig(o_dir + \"SuppFig2\" + '.png')\n",
        "\n",
        "\n",
        "# PCC of experiments\n",
        "pcc_out = open(dir_path + \"prep_files/mpo/all_stress_PCC.txt\", \"w+\")\n",
        "pcc_out.write(\"exp1\\texp2\\tpcc_val\\tp_value\\n\")\n",
        "exps = list(df_scaled.columns)\n",
        "for exp1 in range(len(exps)):\n",
        "\tfor exp2 in range(exp1):\n",
        "\t\tif exps[exp1].split(\"_\")[0] == exps[exp2].split(\"_\")[0]:\n",
        "\t\t\tpcc_val, p_value = pearsonr(df_scaled[exps[exp1]], df_scaled[exps[exp2]])\n",
        "\t\t\tpcc_out.write(exps[exp1] + \"\\t\" + exps[exp2] + \"\\t\" + str(pcc_val) + \"\\t\" + str(p_value) + \"\\n\")\n",
        "pcc_out.close()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cTKSAfIj7rcq"
      },
      "source": [
        "### Supp. Fig 3: Volcano plots (DESeq2)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5Cp2Q3EKrCD3"
      },
      "source": [
        "# adapted from deseq_volcano.py\n",
        "\n",
        "wdir = dir_path + 'prep_files/mpo/deseq/'\n",
        "odir = wdir + 'volcano/'\n",
        "deseqouts = [x for x in os.listdir(wdir) if \"res.tsv\" in x] # controlD2controlH2_res.tsv\n",
        "control = 'controlD2controlH2_res.tsv'\n",
        "deseqouts.pop(deseqouts.index(control))\n",
        "all_stress = list(set([x.split('control')[0] for i, x in enumerate(deseqouts)]))\n",
        "all_stress.sort()\n",
        "s_stress = [x for x in all_stress if len(x) == 1]\n",
        "c_stress = [x for x in all_stress if len(x) == 2]\n",
        "all_stress = s_stress + c_stress\n",
        "\n",
        "# plot control\n",
        "controls = pd.read_csv(wdir + control,\n",
        "\t\t\t  sep = \"\\t\", header = 0, index_col = 0)\n",
        "sns.scatterplot(x = controls['log2FoldChange'],\n",
        "\t\t\t\t y = -np.log10(controls[\"padj\"]),\n",
        "\t\t\t\t #ax = axs[axcord[0][0], axcord[0][1]],\n",
        "\t\t\t\t alpha = 0.2,\n",
        "\t\t\t\t marker = '.',\n",
        "\t\t\t\t legend = False,\n",
        "\t\t\t\t edgecolor = \"none\",\n",
        "\t\t\t\t hue = np.logical_and(abs(controls['log2FoldChange']) > 1,\n",
        "\t\t\t\t\t\t  -np.log10(controls[\"padj\"]) > -np.log10(0.05)))\n",
        "plt.title(\"control \" + control.split(\"control\")[1] + \" vs control H2\")\n",
        "\n",
        "# plot everything else\n",
        "xlen = 5\n",
        "ylen = math.ceil(len(deseqouts)/5)\n",
        "fig, axs = plt.subplots(ylen, xlen, figsize=(30, 37.5), sharex='col', sharey='row')\n",
        "#sns.set(font_scale=1.6)\n",
        "axcord = []\n",
        "for a in range(ylen):\n",
        "\tfor b in range(xlen):\n",
        "\t\taxcord.append([a, b])\n",
        "\n",
        "for i, z in enumerate(all_stress):\n",
        "\tfiles = [x for x in deseqouts if x.startswith(z+'control')]\n",
        "\tfiles.sort()\n",
        "\tfileD2 = pd.read_csv(wdir + files[0],\n",
        "\t\t\t\t\t  sep = \"\\t\", header = 0, index_col = 0)\n",
        "\tfileH2 = pd.read_csv(wdir + files[1],\n",
        "\t\t\t\t\t  sep = \"\\t\", header = 0, index_col = 0)\n",
        "\tD2ax = int(((i*2)//10)*10 + ((i*2)%10)/2)\n",
        "\tH2ax = int(D2ax + 5)\n",
        "\t# Volcano plots\n",
        "\t# against control D2\n",
        "\tsns.scatterplot(x = fileD2['log2FoldChange'],\n",
        "\t\t\t\t y = -np.log10(fileD2[\"padj\"]),\n",
        "\t\t\t\t ax = axs[axcord[D2ax][0], axcord[D2ax][1]],\n",
        "\t\t\t\t alpha = 0.2,\n",
        "\t\t\t\t marker = '.',\n",
        "\t\t\t\t legend = False,\n",
        "\t\t\t\t edgecolor = \"none\",\n",
        "\t\t\t\t hue = np.logical_and(abs(fileD2['log2FoldChange']) > 1,\n",
        "\t\t\t\t\t\t  -np.log10(fileD2[\"padj\"]) > -np.log10(0.05)))\n",
        "\taxs[axcord[D2ax][0], axcord[D2ax][1]].set_title(z + \" vs control D2\")\n",
        "\t\n",
        "\t# against control H2\n",
        "\tsns.scatterplot(x= fileH2['log2FoldChange'],\n",
        "\t\t\t\t y = -np.log10(fileH2[\"padj\"]),\n",
        "\t\t\t\t ax = axs[axcord[H2ax][0], axcord[H2ax][1]],\n",
        "\t\t\t\t alpha = 0.2,\n",
        "\t\t\t\t marker = '.',\n",
        "\t\t\t\t legend = False,\n",
        "\t\t\t\t edgecolor = \"none\",\n",
        "\t\t\t\t hue = np.logical_and(abs(fileH2['log2FoldChange']) > 1,\n",
        "\t\t\t\t\t\t  -np.log10(fileH2[\"padj\"]) > -np.log10(0.05)))\n",
        "\taxs[axcord[H2ax][0], axcord[H2ax][1]].set_title(z + \" vs control H2\")\n",
        "\t\n",
        "for ax in axs.flat:\n",
        "    ax.set(xlabel='log2FoldChange', ylabel='-log10 padj')\n",
        "\n",
        "# Hide x labels and tick labels for top plots and y ticks for right plots.\n",
        "for ax in axs.flat:\n",
        "    ax.label_outer()\n",
        "\n",
        "plt.savefig(dir_path + \"figures/SuppFig3.png\", dpi = 600)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ag4VOlpF72Dh"
      },
      "source": [
        "### Supp. Fig 8: Overview of diurnal data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JJ8VHomNbC3R"
      },
      "source": [
        "# adapted from QC_scaled_updated.py\n",
        "\n",
        "from sklearn.decomposition import PCA\n",
        "from sklearn.preprocessing import StandardScaler\n",
        "\n",
        "o_dir = dir_path + 'figures/'\n",
        "sumdir = dir_path + 'summary_files/'\n",
        "\n",
        "expdesc = ['all_stress', 'diurnal_exp', 'single_stress']\n",
        "targetexp = expdesc[1]\n",
        "targetp = sumdir + targetexp + '.txt'\n",
        "expmatp = dir_path + 'prep_files/' + targetexp + '.tsv'\n",
        "exps = [x.split(\"\\t\")[0] for x in open(targetp, \"r\").readlines()]\n",
        "labels = [x.strip().split(\"\\t\")[1] + '_' + x.split(\"\\t\")[0].split('_')[1] for x in open(targetp, \"r\").readlines()]\n",
        "\n",
        "df = pd.read_csv(expmatp, index_col = 0, sep = \"\\t\", header = 0)\n",
        "df.columns = labels\n",
        "\n",
        "# =============================================================================\n",
        "#\n",
        "# PCA for diurnal by genes (Panel A)\n",
        "#\n",
        "# =============================================================================\n",
        "pca = PCA(n_components=2)\n",
        "diurnal_transformed = StandardScaler().fit_transform(df.values)\n",
        "pcomp = pca.fit_transform(diurnal_transformed.T)\n",
        "p_df = pd.DataFrame(data = pcomp, columns = ['PC1', 'PC2'])\n",
        "pc1, pc2 = pca.explained_variance_ratio_\n",
        "\n",
        "finalDf = p_df.copy()\n",
        "finalDf['target'] = [x.split('_')[0] for x in df.columns.to_list()]\n",
        "#plot PCA figure\n",
        "PCA_plot = o_dir + 'SuppFig8A.png'\n",
        "\n",
        "fig = plt.figure(figsize = (8,5))\n",
        "ax = fig.add_subplot(1,1,1) \n",
        "ax.set_xlabel('Principal Component 1: ' + str(round(pc1, 2)), fontsize = 15)\n",
        "ax.set_ylabel('Principal Component 2: ' + str(round(pc2, 2)), fontsize = 15)\n",
        "\n",
        "targets = [x.split('_')[0] for i, x in enumerate(df.columns.to_list()) if i%3 == 0]\n",
        "colors = ['r', 'y', 'g', 'b', 'c', 'm']\n",
        "for target, color in zip(targets,colors):\n",
        "    indicesToKeep = finalDf['target'] == target\n",
        "    ax.scatter(finalDf.loc[indicesToKeep, 'PC1']\n",
        "               , finalDf.loc[indicesToKeep, 'PC2']\n",
        "               , c = color\n",
        "               , s = 50\n",
        "\t\t\t   , alpha = 0.5\n",
        "\t\t\t   , edgecolors = 'k')\n",
        "\n",
        "ax.legend(targets, loc = 'upper right', fontsize='xx-small')\n",
        "ax.grid()\n",
        "\n",
        "plt.savefig(PCA_plot, dpi=600)\n",
        "\n",
        "# =============================================================================\n",
        "#\n",
        "# PCA for diurnal by genes (Panel B)\n",
        "#\n",
        "# =============================================================================\n",
        "\n",
        "# JTK_output info\n",
        "wdir = dir_path + \"diurnal/\"\n",
        "Mpodf = pd.read_csv(wdir + \"Mpo_supp.txt\", sep = \"\\t\", index_col = 0)\n",
        "\n",
        "#PCA part\n",
        "df['target'] = [Mpodf.loc[x].LAG for x in df.index.to_list()]\n",
        "diurnal_filt = df[df.target != 'NE']\n",
        "diurnal_scaled = StandardScaler().fit_transform(diurnal_filt.iloc[:,:-1].T.values)\n",
        "\n",
        "pca = PCA(n_components=2)\n",
        "pcomp = pca.fit_transform(diurnal_scaled.T)\n",
        "p_df = pd.DataFrame(data = pcomp, columns = ['PC1', 'PC2'])\n",
        "pc1, pc2 = pca.explained_variance_ratio_\n",
        "\n",
        "finalDf = p_df.copy()\n",
        "finalDf['target'] = diurnal_filt.target.to_list()\n",
        "\n",
        "#plot PCA figure\n",
        "PCA_plot = o_dir + 'SuppFig8B.png'\n",
        "\n",
        "fig = plt.figure(figsize = (8,5))\n",
        "ax = fig.add_subplot(1,1,1) \n",
        "ax.set_xlabel('Principal Component 1: ' + str(round(pc1, 2)), fontsize = 15)\n",
        "ax.set_ylabel('Principal Component 2: ' + str(round(pc2, 2)), fontsize = 15)\n",
        "\n",
        "targets = list(finalDf.target.unique())\n",
        "targets.sort()\n",
        "num_only = [int(x) for x in targets[:-1]]\n",
        "num_only.sort()\n",
        "new_targets = [str(x) for x in num_only] + [targets[-1]]\n",
        "colors = [(1,1,0), (1,0.75,0), (1,0.5,0),\n",
        "\t\t  (1,0.25,0), (1,0,0.25), (1,0,0.5),\n",
        "\t\t  (1,0,0.75), (1,0,1), (0.75,0,1),\n",
        "\t\t  (0.5,0,1), (0.25,0,1), (0,0,1),\n",
        "\t\t  (0.75,0.75,0.75)]\n",
        "for target, color in zip(new_targets,colors):\n",
        "    indicesToKeep = finalDf['target'] == target\n",
        "    ax.scatter(finalDf.loc[indicesToKeep, 'PC1']\n",
        "               , finalDf.loc[indicesToKeep, 'PC2']\n",
        "               , color = color\n",
        "               , s = 50\n",
        "\t\t\t   , alpha = 0.5\n",
        "\t\t\t   , edgecolors = None)\n",
        "\n",
        "ax.legend(new_targets, bbox_to_anchor=(1, 1), fontsize='x-small')\n",
        "ax.grid()\n",
        "\n",
        "plt.savefig(PCA_plot, dpi=600)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HWapwnXyoItc"
      },
      "source": [
        "# 4. Experimental"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2DT-AxvEZtOw"
      },
      "source": [
        "### 2.1 Download RNA-seq experiments !experimental"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2hVebO7OZsP3"
      },
      "source": [
        "kal_dir = dir_path + 'kal_out/'\n",
        "def kal_index():\n",
        "\n",
        "def get_ftp_links(RunID):\n",
        "\t'''(str)->(lst,str)\n",
        "\tReturn ftp link in the paired and unpaired format for the RunID specified\n",
        "\t'''\n",
        "\tdir2 = \"\"\n",
        "\tif 9 < len(RunID) <= 12:\n",
        "\t\tdir2 = \"0\"*(12 - len(RunID)) + RunID[-(len(RunID)-9):] + \"/\"\n",
        "\t\tdirs = RunID[:6] + \"/\" + dir2 + RunID\n",
        "\t\tftp_link_paired = [dirs + \"/\" + RunID + \"_1.fastq.gz\",\n",
        "\t\t\t\t\t dirs + \"/\" + RunID + \"_2.fastq.gz\"]\n",
        "\t\tftp_link_unpaired = dirs + \"/\" + RunID + \".fastq.gz\"\n",
        "\telif len(RunID) == 9:\n",
        "\t\tdirs = RunID[:6] + \"/\" + RunID\n",
        "\t\tftp_link_paired = [dirs + \"/\" + RunID + \"_1.fastq.gz\",\n",
        "\t\t\t\t\t dirs + \"/\" + RunID + \"_2.fastq.gz\"]\n",
        "\t\tftp_link_unpaired = dirs + \"/\" + RunID + \".fastq.gz\"\n",
        "\treturn ftp_link_paired, ftp_link_unpaired\n",
        "\n",
        "def kal_single(outname, index, SpotLen, flink):\n",
        "\t!kallisto quant -i $index -o $outname --single -l $SpotLen -s 20 -t 2 <(curl $flink)\n",
        "\n",
        "def kal_paired(outname, index, flink1, flink2):\n",
        "\t!kallisto quant -i $index -o $outname -t 2 <(curl $flink1 $flink2)\n",
        "\n",
        "# Download Rice experiments\n",
        "kal_osa = kal_dir + 'osa/'\n",
        "if not os.path.exists(kal_osa):\n",
        "\t!mkdir $kal_osa\n",
        "RunTable = pd.read_csv(sum_dir + \"selected_Osa.txt\",\n",
        "\t\t\t  sep = \"\\t\", header = 0)\n",
        "\n",
        "for i in range(len(RunTable)):\n",
        "\trunid = RunTable[\"Run\"][i]\n",
        "\tstudy = RunTable[\"Study\"][i]\n",
        "\tliblay = RunTable[\"Layout\"][i]\n",
        "\tspotlen = RunTable[\"Spot_length\"][i]\n",
        "\tif study + \"_\" + runid not in completed:\n",
        "\t\tpath_paired, path_single = get_ftp_links(runid)\n",
        "\t\tprint(str(i) + \"\\t\" + path_single.split(\"/\")[-1].split(\".fastq.gz\")[0] + \"\\t\" + liblay + \"\\n\")\n",
        "\t\tif liblay == \"SINGLE\":\n",
        "\t\t\tkal_single(kal_osa + study+'_'+runid, osa_idx, spotlen, pathsingle)\n",
        "\t\telif liblay == \"PAIRED\":\n",
        "\t\t\tkal_paired(kal_osa + study+'_'+runid, osa_idx, path_paired[0], path_paired[1])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PbX3ba-a2uNC"
      },
      "source": [
        "### 2.2 Generate expression matrix !experimental"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Eko_SXvA2mOn"
      },
      "source": [
        "# Generation of gene expression matrix and kallisto statistics\n",
        "def kal_extract(kout, exps):\n",
        "\t'''(str,list)->(dict,dict)\n",
        "\tReturn dictionary containing tpm and raw expression value\n",
        "\t'''\n",
        "\tdicto = {}\n",
        "\tdicto_raw = {}\n",
        "\n",
        "\toutput_header = 'gene\\t'\n",
        "\toutput_content = ''\n",
        "\tfor folder in exps:\n",
        "\t\tfilep = kout + folder + '/abundance.tsv'\n",
        "\t\tif os.path.exists(filep):\n",
        "\t\t\tprint('In directory ' + folder)\n",
        "\t\t\toutput_header += folder + '\\t'\n",
        "\t\t\tcontent = open(filep, 'r')\n",
        "\t\t\tcontent.readline()\n",
        "\t\t\tfor item in content:\n",
        "\t\t\t\titem, tpm = item.rstrip().split('\\t')\n",
        "\t\t\t\traw = str(round(float(values[-2])))\n",
        "\t\t\t\tif item in dicto:\n",
        "\t\t\t\t\tdicto[item].append(tpm)\n",
        "\t\t\t\telse:\n",
        "\t\t\t\t\tdicto[item] = [tpm]\n",
        "\t\t\t\tif item in dicto_raw:\n",
        "\t\t\t\t\tdicto_raw[item].append(raw)\n",
        "\t\t\t\telse:\n",
        "\t\t\t\t\tdicto_raw[item] = [raw]\n",
        "\tif '' in dicto:\n",
        "\t\tdicto.pop('')\n",
        "\tif '' in dicto_raw:\n",
        "\t\tdicto.pop('')\n",
        "\treturn dicto, dicto_raw\n",
        " \n",
        "def write_expmat(filepath, dicttouse):\n",
        "\t'''(str, dict)->(None)\n",
        "\tWrites expression matrix to file from dictionary\n",
        "\t'''\n",
        "\twith open(filepath, \"w+\") as output_file:\n",
        "\t\toutput_file.write(output_header[:-1] + \"\\n\")\n",
        "\t\tfor key, value in dicttouse.items():\n",
        "\t\t\tline = ''\n",
        "\t\t\tline += key + '\\t'\n",
        "\t\t\tfor item in value:\n",
        "\t\t\t\tline += item + '\\t'\n",
        "\t\t\toutput_file.write(line[:-1] + \"\\n\")\n",
        "\n",
        "def kal_stats(kout):\n",
        "\t'''(str)->(None)\n",
        "\tWrites summary of mapping statistics of kallisto runs to file\n",
        "\t'''\n",
        "\tkal_dirs = [x for x in os.listdir(kout)]\n",
        "\n",
        "\twith open(kout + \"kallisto_stats.txt\", \"w+\") as output_file:\n",
        "\t\toutput_file.write(\"experiment\\tn_processed\\tn_pseudoaligned\\tn_unique\\tp_pseudoaligned\\tp_unique\\n\")\n",
        "\t\tfor folder in kal_dirs:\n",
        "\t\t\tkallisto_json = ast.literal_eval(open(kout + folder + '/run_info.json', 'r').read())\n",
        "\t\t\tprocessed = kallisto_json[\"n_processed\"]\n",
        "\t\t\tpseudoaligned = kallisto_json[\"n_pseudoaligned\"]\n",
        "\t\t\tunique = kallisto_json[\"n_unique\"]\n",
        "\t\t\tppseudoaligned = kallisto_json[\"p_pseudoaligned\"]\n",
        "\t\t\tpunique = kallisto_json[\"p_unique\"]\n",
        "\t\t\toutput_file.write(folder + \"\\t\" +\n",
        "\t\t\t\t\t\tstr(processed) + \"\\t\" +\n",
        "\t\t\t\t\t\tstr(pseudoaligned) + \"\\t\" +\n",
        "\t\t\t\t\t\tstr(unique) + \"\\t\" +\n",
        "\t\t\t\t\t\tstr(ppseudoaligned) + \"\\t\" +\n",
        "\t\t\t\t\t\tstr(punique) + \"\\n\")\n",
        "\n",
        "# Marchantia\n",
        "sum_dir = dir_path + 'summary_files/'\n",
        "\n",
        "expdesc = ['all_stress', 'diurnal_exp', 'single_stress', 'cross_stress']\n",
        "for targetexp in expdesc:\n",
        "\ttargetp = sum_dir + targetexp + '.txt'\n",
        "\texpmatp = dir_path + 'prep_files/' + targetexp + '.tsv'\n",
        "\texpmatrawp = dir_path + 'prep_files/' + targetexp + '_raw.tsv'\n",
        "\n",
        "\tmpo_exps = [x.split(\"\\t\")[0] for x in open(targetp, \"r\").readlines()]\n",
        "\tmpo_tpm, mpo_raw = kal_extract(kal_dir + 'mpo/', mpo_exps)\n",
        "\tfor i in [expmatp, expmatrawp]:\n",
        "\t\twrite_expmat(i, mpo_tpm)\n",
        "\t\twrite_expmat(i, mpo_raw)\n",
        "kal_stats(kal_dir + 'mpo/', \"kallisto_stats.txt\")\n",
        "\n",
        "# Rice\n",
        "RunTable = pd.read_csv(sum_dir + \"selected_Osa.txt\",\n",
        "\t\t\t  sep = \"\\t\", header = 0)\n",
        "osa_runs = RunTable.Run.to_list()\n",
        "osa_study = RunTable.Study.to_list()\n",
        "osa_exps = [osa_study[i] + '_' + x for i, x in enumerate(osa_runs)]\n",
        "expmatp = dir_path + 'prep_files/' + 'expmat_Osa.tsv'\n",
        "expmatrawp = dir_path + 'prep_files/' + 'expmat_Osa_raw.tsv'\n",
        "osa_tpm, osa_raw = kal_extract(kal_dir + 'osa/', osa_exps)\n",
        "for i in [expmatp, expmatrawp]:\n",
        "\twrite_expmat(i, osa_tpm)\n",
        "\twrite_expmat(i, osa_raw)\n",
        "kal_stats(kal_dir + 'osa/', \"kallisto_stats.txt\")"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}