Snakefile

import sys
import time
import os
from os import path
import pandas as pd
from glob import glob
import re
from threading import Lock
from itertools import zip_longest
from collections import OrderedDict, namedtuple

if not workflow.overwrite_configfiles:
    configfile: "config.yml"

WORKDIR = path.abspath(path.join(config["workdir_top"], config["pipeline"]))
workdir: WORKDIR
SNAKEDIR = path.dirname(workflow.snakefile)

include: "snakelib/utils.snake"

in_fastq = config["reads_fastq"]
if not path.isabs(in_fastq):
	in_fastq = path.join(SNAKEDIR, in_fastq)


class Node:
    def __init__(self, Id, File, Left, Right, Parent, Level):
        self.Id = Id
        self.File = File
        self.Left = Left
        self.Right = Right
        self.Parent = Parent
        self.Level = Level
        self.Done = False
        self.RightSide = False
    def __repr__(self):
        return "Node:{} Level: {} File: {} Done: {} Left: {} Right: {} Parent: {}".format(self.Id, self.Level, self.File, self.Done, self.Left.Id if self.Left is not None else None, self.Right.Id if self.Right is not None else None, self.Parent.Id if self.Parent is not None else None)

def grouper(n, iterable, fillvalue=None):
    "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
    args = [iter(iterable)] * n
    return zip_longest(fillvalue=fillvalue, *args)

def build_job_tree(batch_dir):
    batches = glob("{}/isONbatch_*.cer".format(batch_dir))
    batch_ids = [int(re.search('/isONbatch_(.*)\.cer$', x).group(1)) for x in batches]
    global LEVELS
    global JOB_TREE
    LEVELS = OrderedDict()
    LEVELS[0] = []
    for Id, bf in sorted(zip(batch_ids, batches), key=lambda x: x[0]):
        n = Node(Id, "clusters/isONcluster_{}.cer".format(Id), None, None, None, 0)
        n.Done = True
        JOB_TREE[Id] = n
        LEVELS[0].append(n)
    level = 0
    max_id = LEVELS[0][-1].Id
    while len(LEVELS[level]) != 1:
        next_level = level + 1
        LEVELS[next_level] = []
        for l, r in grouper(2, LEVELS[level]):
            if r is None:
                LEVELS[level].pop()
                l.Level += 1
                LEVELS[next_level].append(l)
                continue
            max_id += 1
            new_batch = "clusters/isONcluster_{}.cer".format(max_id)
            new_node = Node(max_id, new_batch, l, r, None, next_level)
            l.Parent = new_node
            r.Parent =  new_node
            r.RightSide = True
            LEVELS[next_level].append(new_node)
            JOB_TREE[max_id] = new_node
        level = next_level
    global ROOT
    ROOT = JOB_TREE[len(JOB_TREE)-1].Id
    JOB_TREE[ROOT].RightSide = True
    print("Merge clustering job tree nodes:",file=sys.stderr)
    for n in JOB_TREE.values():
        print("\t{}".format(n),file=sys.stderr)


def generate_rules(levels, snk):
    init_template = """
rule cluster_job_%d:
    input:
        left = "sorted/batches/isONbatch_%d.cer",
    output: "clusters/isONcluster_%d.cer"
    shell: "isONclust2 cluster -x %s -v -Q -l %s -o %s %s; sync"

    """
    template = """
rule cluster_job_%d:
    input:
        left = "clusters/isONcluster_%d.cer",
        right = "clusters/isONcluster_%d.cer",
    output: "clusters/isONcluster_%d.cer"
    shell: "isONclust2 cluster -x %s -v -Q -l %s -r %s -o %s %s; sync"

    """

    link_template="""
rule link_root:
    input: "clusters/isONcluster_%d.cer",
    output: "clusters/isONcluster_ROOT.cer"
    shell: "ln -s `realpath {input}` {output}"

    """

    fh = open(snk, "w")
    for nr, l in levels.items():
        for n in l:
            purge = "-z" if n.RightSide else ""
            if nr == 0 or n.Left is None or n.Right is None:
                jr = init_template % (n.Id, n.Id, n.Id, config["cls_mode"], "{input.left}", "{output}", purge)
                fh.write(jr)
            else:
                jr = template % (n.Id, n.Left.Id, n.Right.Id, n.Id, config["cls_mode"], "{input.left}", "{input.right}", "{output}", purge)
                fh.write(jr)
    global ROOT
    fh.write(link_template % ROOT)
    fh.write("\nROOT = %d" % ROOT)
    fh.flush()
    fh.close()

def count_fastq_bases(fname, size=128000000):
    fh = open(fname, "r")
    count = 0
    while True:
        b = fh.read(size)
        if not b:
            break
        count += b.count("A")
        count += b.count("T")
        count += b.count("G")
        count += b.count("C")
        count += b.count("U")
    fh.close()
    return count

def preprocess_reads(fq):
    pc_opts = config["pychopper_opts"]
    concat = config["concatenate"]
    thr = config["cores"]

    out_fq = "processed_reads/full_length_reads.fq"
    if os.path.isdir("processed_reads"):
        return out_fq

    shell("mkdir -p processed_reads")
    if concat:
        print("Concatenating reads under directory: " + fq)
        shell("find %s  -regextype posix-extended -regex '.*\.(fastq|fq)$' -exec cat {{}} \\; > processed_reads/input_reads.fq" % fq)
    else:
        shell("ln -s `realpath %s` processed_reads/input_reads.fq" % fq)

    if config["run_pychopper"]:
        print("Running pychopper of fastq file: processed_reads/input_reads.fq")
        shell("(cd processed_reads; cdna_classifier.py -t %d %s input_reads.fq full_length_reads.fq)" % (thr, pc_opts))
    else:
        shell("ln -s `realpath processed_reads/input_reads.fq` processed_reads/full_length_reads.fq")

    return out_fq


ROOT = None
DYNAMIC_RULES="job_rules.snk"
if ((not os.path.isfile(os.path.join(WORKDIR,"sorted","sorted_reads.fastq"))) or (not os.path.isfile(os.path.join(SNAKEDIR, DYNAMIC_RULES)))):
    print("Preprocessing read in fastq file:", in_fastq)
    proc_fastq = preprocess_reads(in_fastq)
    print("Counting records in input fastq:", proc_fastq)
    nr_bases = count_fastq_bases(proc_fastq)
    print("Bases in input: {} megabases".format(int(nr_bases/10**6)))
    if config['batch_size'] < 0:
        config['batch_size'] = int(nr_bases/1000/config["cores"])

    print("Batch size is: {}".format(config['batch_size']))

    init_cls_options = """ --batch-size {} --kmer-size {} --window-size {} --min-shared {} --min-qual {}\
                         --mapped-threshold {} --aligned-threshold {} --min-fraction {} --min-prob-no-hits {} -M {} -P {} -g {} -c {} -F {} """
    init_cls_options = init_cls_options.format(config["batch_size"], config["kmer_size"], config["window_size"], config["min_shared"], config["min_qual"], \
                        config["mapped_threshold"], config["aligned_threshold"], config["min_fraction"], config["min_prob_no_hits"], config["batch_max_seq"], config["consensus_period"],
    config["consensus_minimum"], config["consensus_maximum"], config["min_left_cls"])

    shell("""
            rm -fr clusters sorted
            mkdir -p sorted; isONclust2 sort {} -v -o sorted {};
            mkdir -p clusters;
        """.format(init_cls_options, proc_fastq))
    JOB_TREE = OrderedDict()
    LEVELS = None

    build_job_tree("sorted/batches")
    generate_rules(LEVELS, "{}/job_rules.snk".format(SNAKEDIR))

include: DYNAMIC_RULES



rule all:
    input: rules.link_root.output
    output: directory("final_clusters")
    shell:
        """ isONclust2 dump -v -i sorted/sorted_reads_idx.cer -o final_clusters {input}; sync """