workflow/finemapping_results_wrangler.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "kernel": "SoS"
   },
   "source": [
    "# Fine-mapping result post processing\n",
    "\n",
    "This pipeline consolidates results from various fine-mapping tools to uniform format, add rsID as necessary, and perform a simple \"liftover\" via rsID (not the formal UCSC `liftOver`) to generate output in HG37 and HG38 builds."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "kernel": "SoS"
   },
   "source": [
    "This pipeline was devised by Gao Wang (UChicago), with core implemention done by Kushal Dey (Harvard)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "kernel": "SoS"
   },
   "source": [
    "## Input data\n",
    "\n",
    "Input are results of fine-mapping pipeline `summary_statistics_finemapping.ipynb` in R's `RDS` format for SuSiE and CAVIAR, and `pkl` format for DAP.\n",
    "\n",
    "## Output data\n",
    "\n",
    "columns are:\n",
    "\n",
    "```\n",
    "chr pos ref alt variant_id locus_id ...\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "kernel": "SoS"
   },
   "source": [
    "where `snp_id` will be rsID if some annotation files on rsID are provided and the `update_variant_id` workflow below was executed. Otherwise it will take the default format of `chr:pos:ref:alt`.\n",
    "\n",
    "Because each fine-mapping software can have different output, the columns after the first 6 are arbitary. For SuSiE for example they are\n",
    "\n",
    "```\n",
    "chr pos ref alt variant_id locus_id pip cs cs_size cs_purity\n",
    "```\n",
    "\n",
    "For CAVIAR they are\n",
    "```\n",
    "chr pos ref alt variant_id locus_id pip cs_size\n",
    "```\n",
    "\n",
    "Naturally per locus `cs` related information will be redundant. This is one draw-back to using a flat table output format.\n",
    "\n",
    "## Additional data processing\n",
    "\n",
    "Here we also provide additional routines to process the data,\n",
    "\n",
    "1. Update the `variant_id` column using external annotations, for example by rsID.\n",
    "2. \"liftover\" to other builds -- we only support it via rsID matching\n",
    "\n",
    "Notice that only the first 5 columns are necessary for these additional operations. The columns after the fifth can be arbitary and will be kept during the process.\n",
    "\n",
    "- To trigger optional step 1, parameter `--id-map-prefix` and `--id-map-suffix` have to be valid.\n",
    "- To trigger optional step 2, parameter `--coordinate-map-prefix` and `--coordinate-map-suffix` have to be valid."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "kernel": "SoS"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/scratch/midway2/gaow/GIT/github/fine-mapping"
     ]
    }
   ],
   "source": [
    "%cd ~/GIT/github/fine-mapping"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "kernel": "SoS"
   },
   "source": [
    "## The workflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "kernel": "Bash"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "usage: sos run workflow/finemapping_results_wrangler.ipynb\n",
      "               [workflow_name | -t targets] [options] [workflow_options]\n",
      "  workflow_name:        Single or combined workflows defined in this script\n",
      "  targets:              One or more targets to generate\n",
      "  options:              Single-hyphen sos parameters (see \"sos run -h\" for details)\n",
      "  workflow_options:     Double-hyphen workflow-specific parameters\n",
      "\n",
      "Workflows:\n",
      "  default\n",
      "  consolidate\n",
      "  update_variant_id\n",
      "  update_coordinate\n",
      "\n",
      "Global Workflow Options:\n",
      "  --loci . (as path)\n",
      "                        Loci file\n",
      "  --ss-data-prefix . (as path)\n",
      "                        summary statistics file prefix\n",
      "  --pattern 'uniform.SuSiE_B.L_5.prior_0p005.res_var_false'\n",
      "                        identifier for fine-mapping results to be extracted\n",
      "\n",
      "Sections\n",
      "  default:\n",
      "  consolidate:          Consolidate fine-mapping results to a single file\n",
      "    Workflow Options:\n",
      "      --pip-thresh 0.05 (as float)\n",
      "                        Keep PIP above these thresholds\n",
      "      --round-off 6 (as int)\n",
      "                        Round PIP to given digits\n",
      "  update_variant_id:    Update variant ID based on chrom and pos, from per chrom\n",
      "                        files (optional)\n",
      "    Workflow Options:\n",
      "      --id-map-prefix . (as path)\n",
      "                        Path containing files for variant ID update rule Each\n",
      "                        file is a separate chromsome\n",
      "      --id-map-suffix bim\n",
      "      --columns 2 4 (as list)\n",
      "                        columns first element for variant ID, 2nd element for\n",
      "                        genomic position is [2,4] for BIM files\n",
      "      --chroms 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 (as list)\n",
      "                        chromosome identifiers\n",
      "  update_coordinate:    Update genomic coordinates based on variant ID, from per\n",
      "                        chrom files (optional)\n",
      "    Workflow Options:\n",
      "      --coordinate-map-prefix . (as path)\n",
      "                        Path containing files for coordinate update rule Each\n",
      "                        file is a separate chromsome\n",
      "      --coordinate-map-suffix bim\n",
      "      --coordinate-version-id hgX\n",
      "                        coordinate identifier\n",
      "      --columns 2 4 5 6 (as list)\n",
      "                        columns first element for variant ID, 2nd element for\n",
      "                        genomic position 3rd for reference allele and 4th for\n",
      "                        alternative allele is [2,4,5,6] for BIM files\n",
      "      --chroms 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 (as list)\n",
      "                        chromosome identifiers\n"
     ]
    }
   ],
   "source": [
    "sos run workflow/finemapping_results_wrangler.ipynb -h"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "kernel": "SoS"
   },
   "outputs": [],
   "source": [
    "[global]\n",
    "# Loci file\n",
    "parameter: loci = path()\n",
    "# summary statistics file prefix\n",
    "parameter: ss_data_prefix = path()\n",
    "# identifier for fine-mapping results to be extracted\n",
    "parameter: pattern = \"uniform.SuSiE_B.L_5.prior_0p005.res_var_false\"\n",
    "\n",
    "fail_if(not loci.is_file(), msg = 'Please specify valid path for --loci')\n",
    "fail_if(ss_data_prefix.is_file(), msg = '--ss-data-prefix should be a path not a file (usually file without extension, if using input from my data wrangling pipeline)')\n",
    "\n",
    "ss_data_prefix = ss_data_prefix.absolute()\n",
    "chunks = [x.strip().split() for x in open(f'{loci:a}').readlines() if not x.strip().startswith('#')]\n",
    "chunks = [x[3] if len(x) == 4 else \"%s_%s_%s\" % (x[0], x[1], x[2]) for x in chunks]\n",
    "data = [f'{ss_data_prefix}/{x}/{x}.summary_stats.gz' for x in chunks]\n",
    "\n",
    "\n",
    "if 'SuSiE' in pattern:\n",
    "    source = 'susie'\n",
    "elif 'CAVIAR' in pattern:\n",
    "    source = 'caviar'\n",
    "elif 'DAP' in pattern:\n",
    "    source = 'dap'\n",
    "elif 'FINEMAP' in pattern:\n",
    "    source = 'finemap'\n",
    "else:\n",
    "    raise ValueError(\"Invalid --pattern specification\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "kernel": "SoS"
   },
   "outputs": [],
   "source": [
    "[default]\n",
    "sos_run('consolidate+update_variant_id+update_coordinate')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "kernel": "SoS"
   },
   "source": [
    "## Fine-mapping results consolidation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "kernel": "SoS"
   },
   "outputs": [],
   "source": [
    "# Consolidate fine-mapping results to a single file\n",
    "[consolidate]\n",
    "depends: R_library('rlang>=0.3.0'), R_library('dplyr'), R_library('data.table'), R_library('R.utils'), R_library('dscrutils') # can be installed via `devtools::install_github(\"stephenslab/dsc\",subdir = \"dscrutils\", force = TRUE)`\n",
    "# Keep PIP above these thresholds\n",
    "parameter: pip_thresh = 0.05\n",
    "# Round PIP to given digits\n",
    "parameter: round_off = 6\n",
    "\n",
    "ext = 'pkl' if source in ['dap'] else 'rds'\n",
    "\n",
    "input: [x for x in paths([f'{ss_data_prefix}/{c}/{c}.{pattern}.{ext}' for c in chunks]) if x.is_file()]\n",
    "output: f\"{ss_data_prefix}.{pattern}.{loci:bn}.gz\"\n",
    "fail_if(len(_input) == 0, msg = f'Cannot find valid input files for given loci')\n",
    "R: expand = \"${ }\", workdir = ss_data_prefix, stdout = f'{_output:n}.log'\n",
    "    # Here we define get_*_output functions for different output format\n",
    "    get_susie_output = function(unit, rds_file) {\n",
    "        cs_id = cs_size = cs_purity = rep(NA, length(rds_file$var_names))\n",
    "        num_cs = length(rds_file$sets$cs)\n",
    "        for(id in 1:num_cs){\n",
    "            idx = rds_file$sets$cs[[id]]\n",
    "            cs_id[idx] = names(rds_file$sets$cs)[id]\n",
    "            cs_size[idx] = length(rds_file$sets$cs[[id]])\n",
    "            cs_purity[idx] = rds_file$sets$purity[id,1]\n",
    "        }\n",
    "        out = cbind.data.frame(rep(unit, length(rds_file$var_names)),\n",
    "                                rds_file$var_names,\n",
    "                                rds_file$pip, cs_id, cs_size, cs_purity)\n",
    "        colnames(out) = c(\"locus_id\", \"variant_id\", \"pip\", \"cs\", \"cs_size\", \"cs_purity\")\n",
    "        out[which(out[,3] >= ${pip_thresh} | !is.na(out[,4])), ]\n",
    "    }\n",
    "   \n",
    "    get_caviar_output = function(unit, rds_file) {\n",
    "      snp = rds_file$snp\n",
    "      snp = snp[match(rds_file$var_names, snp$snp),]\n",
    "      cs_annot = rep(0, length(rds_file$var_names))\n",
    "      cs_annot[match(rds_file$set, rds_file$snp$snp)] = 1 \n",
    "      out = cbind.data.frame(rep(unit, length(rds_file$var_names)),\n",
    "                              rds_file$var_names,\n",
    "                              snp$snp_prob,\n",
    "                              cs_annot)\n",
    "      colnames(out) = c(\"gene_id\", \"var_id\", \"pip\", \"cs\")\n",
    "      out[which(out[,3] >= ${pip_thresh} | out[,4] > 0), ]\n",
    "    }\n",
    "\n",
    "    get_dap_output = function(unit, rds_file) {\n",
    "      out = cbind.data.frame(rep(unit, length(rds_file$var_names)),\n",
    "                              rds_file$var_names,\n",
    "                              rds_file$snp$snp_prob,\n",
    "                              rds_file$snp$cluster)\n",
    "      colnames(out) = c(\"gene_id\", \"var_id\", \"pip\", \"cs\")\n",
    "      out[which(out[,3] >= ${pip_thresh} | out[,4] > 0),]\n",
    "    }\n",
    "  \n",
    "    get_finemap_output =  function(unit, rds_file) {\n",
    "      snp = rds_file$snp\n",
    "      snp = snp[match(rds_file$var_names, snp$snp),]\n",
    "      n = sum(cumsum(rds_file$set$config_prob) < 0.95) + 1\n",
    "      if(n > nrow(rds_file$set)){\n",
    "        rds_file$set = 0\n",
    "      }else{\n",
    "        rds_file$set = unlist(lapply(1:n, function(i) strsplit(as.character(rds_file$set[i,2]), \",\")[[1]]))\n",
    "        rds_file$set = unique(as.vector(unlist(rds_file$set)))\n",
    "      }\n",
    "  \n",
    "      cs_annot = rep(0, length(rds_file$var_names))\n",
    "      cs_annot[match(rds_file$set, rds_file$snp$snp)] = 1 \n",
    "      out = cbind.data.frame(rep(unit, length(rds_file$var_names)),\n",
    "                              rds_file$var_names,\n",
    "                              snp$snp_prob,\n",
    "                              cs_annot)\n",
    "      colnames(out) = c(\"gene_id\", \"var_id\", \"pip\", \"cs\")\n",
    "      out[which(out[,3] >= ${pip_thresh} | out[,4] > 0), ]\n",
    "    }\n",
    "  \n",
    "    is_float = function(x) {is.numeric(x) && !is.integer(x)}\n",
    "\n",
    "    # Data extraction script\n",
    "    suppressMessages(library(dplyr))\n",
    "    suppressMessages(library(data.table))\n",
    "    files = c(${paths([x.relative_to(ss_data_prefix) for x in _input]):r,})\n",
    "    processed_dat = list()\n",
    "    idx = 0\n",
    "    for (f in files) {\n",
    "      idx = idx + 1\n",
    "      rds_file = dscrutils::read_dsc(f)\n",
    "      unit = dirname(f)\n",
    "      processed_dat[[idx]] = get_${source}_output(unit, rds_file)\n",
    "      cat(\"We are at unit #\", idx, \":\", unit, \"\\n\")\n",
    "    }\n",
    "    processed_dat = data.frame(rbindlist(processed_dat))\n",
    "    extract_coord = data.frame(do.call(rbind, lapply(processed_dat[,2], function(x) return(strsplit(as.character(x), \":\")[[1]]))))\n",
    "    df = data.frame(\"chr\" = extract_coord[,1], \"pos\" = extract_coord[,2], \"ref\" = extract_coord[,3], \"alt\" = extract_coord[,4],\n",
    "                    \"variant_id\" = processed_dat[,2], \"locus_id\" = processed_dat[,1])\n",
    "    df = cbind.data.frame(df, processed_dat[, 3:ncol(processed_dat)])\n",
    "    df %>% mutate_if(is_float, ~round(., ${round_off})) -> df\n",
    "    df_sorted = df [order(df[,1], df[,2]),]\n",
    "    fwrite(df_sorted, file = \"${_output:n}.tmp\", sep = \"\\t\", quote=FALSE, row.names=FALSE)\n",
    "    R.utils::gzip(\"${_output:n}.tmp\", destname=${_output:r}, overwrite=TRUE, remove=TRUE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "kernel": "SoS"
   },
   "source": [
    "## Update variant ID via genomic coordinantes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "kernel": "SoS"
   },
   "outputs": [],
   "source": [
    "# Update variant ID based on chrom and pos, from per chrom files (optional)\n",
    "[update_variant_id]\n",
    "depends: R_library('rlang>=0.3.0'), R_library('dplyr'), R_library('R.utils'), R_library('data.table')\n",
    "# Path containing files for variant ID update rule\n",
    "# Each file is a separate chromsome\n",
    "parameter: id_map_prefix = path()\n",
    "parameter: id_map_suffix = 'bim'\n",
    "# columns first element for variant ID, 2nd element for genomic position\n",
    "# is [2,4] for BIM files\n",
    "parameter: columns = [2,4]\n",
    "# chromosome identifiers\n",
    "parameter: chroms = [x+1 for x in range(22)]\n",
    "import glob\n",
    "skip_if(len(glob.glob(f\"{id_map_prefix:a}.*.{id_map_suffix}\")) == 0, msg = 'Variant ID are not updated because no valid file is found using --id-map-prefix and --id-map-suffix')\n",
    "output: f'{_input:n}.id_updated.gz'\n",
    "R: expand = \"${ }\", stdout = f'{_output:n}.log'\n",
    "    suppressMessages(library(dplyr))\n",
    "    suppressMessages(library(data.table))\n",
    "    out =  data.frame(fread(\"zcat ${_input}\"))\n",
    "    out %>% mutate_if(is.factor, as.character) -> out\n",
    "    chroms = c(${paths(chroms):r,})\n",
    "    for(numchr in chroms){\n",
    "        which_chr = which(out$chr == numchr)\n",
    "        out_sub = out[which_chr, ]\n",
    "        dbfile = data.frame(fread(paste0(\"${id_map_prefix}.\", numchr, \".${id_map_suffix}\")))\n",
    "        out_sub_new = out_sub\n",
    "        idx1 = match(out_sub$pos, dbfile$V${columns[1]})\n",
    "        idx2 = idx1[which(!is.na(idx1))]\n",
    "        idx3 = 1:length(out_sub$pos)\n",
    "        idx4 = idx3[which(!is.na(idx1))]\n",
    "        out_sub_new[idx4, \"variant_id\"] = dbfile[idx2,\"V${columns[0]}\"]\n",
    "        out[which_chr, ] = out_sub_new\n",
    "        cat(\"Variant IDs updated for chromosome\", numchr, \"\\n\")\n",
    "    }\n",
    "    fwrite(out, file = \"${_output:n}.tmp\", sep = \"\\t\", quote=FALSE, row.names=FALSE)\n",
    "    R.utils::gzip(\"${_output:n}.tmp\", destname=${_output:r}, overwrite=TRUE, remove=TRUE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "kernel": "SoS"
   },
   "source": [
    "## Update genomic coordinates via variant ID"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "kernel": "SoS"
   },
   "outputs": [],
   "source": [
    "# Update genomic coordinates based on variant ID, from per chrom files (optional)\n",
    "[update_coordinate]\n",
    "depends: R_library('rlang>=0.3.0'), R_library('dplyr'), R_library('data.table'), R_library('R.utils')\n",
    "# Path containing files for coordinate update rule\n",
    "# Each file is a separate chromsome\n",
    "parameter: coordinate_map_prefix = path()\n",
    "parameter: coordinate_map_suffix = 'bim'\n",
    "# coordinate identifier\n",
    "parameter: coordinate_version_id = 'hgX'\n",
    "# columns first element for variant ID, 2nd element for genomic position\n",
    "# 3rd for reference allele and 4th for alternative allele\n",
    "# is [2,4,5,6] for BIM files\n",
    "parameter: columns = [2,4,5,6]\n",
    "# chromosome identifiers\n",
    "parameter: chroms = [x+1 for x in range(22)]\n",
    "import glob\n",
    "skip_if(len(glob.glob(f\"{coordinate_map_prefix:a}.*.{coordinate_map_suffix}\")) == 0, msg = 'Genomic coordinates are not updated because no valid file is found using --coordinate-map-prefix and --coordinate-map-suffix')\n",
    "output: f'{_input:n}.{coordinate_version_id}.gz'.replace('.id_updated.', '.')\n",
    "R: expand = \"${ }\", stdout = f'{_output:n}.log'\n",
    "    suppressMessages(library(dplyr))\n",
    "    suppressMessages(library(data.table))\n",
    "    out =  data.frame(fread(\"zcat ${_input}\"))\n",
    "    out %>% mutate_if(is.factor, as.character) -> out\n",
    "    chroms = c(${paths(chroms):r,})\n",
    "    out2 = c()\n",
    "    for(numchr in chroms){\n",
    "        dbfile = data.frame(fread(paste0(\"${coordinate_map_prefix}.\", numchr, \".${coordinate_map_suffix}\")))\n",
    "        out2_sub = out[which(out$chr == numchr), ]\n",
    "        idx1 = match(out2_sub$variant_id, dbfile$V${columns[0]})\n",
    "        idx2 = idx1[which(!is.na(idx1))]\n",
    "        idx3 = 1:length(out2_sub$variant_id)\n",
    "        idx4 = idx3[which(!is.na(idx1))]\n",
    "        out2_sub[idx4, c(\"pos\", \"ref\", \"alt\")] = dbfile[idx2, c(\"V${columns[1]}\", \"V${columns[2]}\", \"V${columns[3]}\")]\n",
    "        out2 = rbind.data.frame(out2, out2_sub[idx4,])\n",
    "        cat(\"Genomic coordinate updated for chromosome\", numchr, \"\\n\")\n",
    "    }\n",
    "    fwrite(out2, file = \"${_output:n}.tmp\", sep = \"\\t\", quote=FALSE, row.names=FALSE)\n",
    "    R.utils::gzip(\"${_output:n}.tmp\", destname=${_output:r}, overwrite=TRUE, remove=TRUE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "kernel": "SoS"
   },
   "source": [
    "## Example run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "kernel": "Bash"
   },
   "outputs": [],
   "source": [
    "sos run workflow/finemapping_results_wrangler.ipynb \\\n",
    "    --ss-data-prefix ~/tmp/01-Jan-2019 \\\n",
    "    --pattern uniform.SuSiE_B.L_5.prior_0p005.res_var_false \\\n",
    "    --id-map-prefix /project2/mstephens/SuSIE_gtex_CPP/rsID_map/Hg38/1000G.EUR.hg38 \\\n",
    "    --coordinate-map-prefix /project2/mstephens/SuSIE_gtex_CPP/rsID_map/Hg37/1000G.EUR.QC \\\n",
    "    --coordinate-version-id hg37"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "SoS",
   "language": "sos",
   "name": "sos"
  },
  "language_info": {
   "codemirror_mode": "sos",
   "file_extension": ".sos",
   "mimetype": "text/x-sos",
   "name": "sos",
   "nbconvert_exporter": "sos_notebook.converter.SoS_Exporter",
   "pygments_lexer": "sos"
  },
  "sos": {
   "kernels": [
    [
     "Bash",
     "bash",
     "Bash",
     "#E6EEFF"
    ],
    [
     "SoS",
     "sos",
     "",
     ""
    ]
   ],
   "version": "0.19.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}