helpers.py

import os.path

from gwf import *


class source(object):
    def __init__(self, path):
        self.path = path

    def __call__(self, func):
        def wrapper(*args, **kwargs):
            opts, spec = func(*args, **kwargs)
            new_spec = 'source {} && {}'.format(self.path, spec)
            return opts, new_spec
        return wrapper


class extendpath(object):
    def __init__(self, path):
        self.path = path

    def __call__(self, func):
        def wrapper(*args, **kwargs):
            opts, spec = func(*args, **kwargs)
            new_spec = 'export PATH={}:$PATH && {}'.format(self.path, spec)
            return opts, new_spec
        return wrapper


BIRC_LOAD_REGISTRY = []

class birc_load(object):

    def __init__(self, package, version):
        self.package = package
        self.version = version

        self._check_package()

        BIRC_LOAD_REGISTRY.append((self.package, self.version))

    @classmethod
    def packages(self):
        return iter(BIRC_LOAD_REGISTRY)

    def _check_package(self):
        if not os.path.exists(self._get_package_path()):
            raise Exception('Package "{}" version {} does not exist.'.format(
                  self.package, self.version))

    def _get_package_path(self):
        return '/com/extra/{}/{}/load.sh'.format(self.package, self.version)

    def __call__(self, func):
        def wrapper(*args, **kwargs):
            opts, spec = func(*args, **kwargs)
            new_spec = 'source {} && {}'.format(self._get_package_path(), spec)
            return opts, new_spec
        return wrapper


def with_ext(path, newext):
    """Replace the extension of `path` with `newext`"""
    basepath, _ = os.path.splitext(path)
    return basepath + '.' + newext


def pathiter(paths):
    for path in paths:
        base, ext = os.path.splitext(path)
        basedir = os.path.dirname(base)
        basename = os.path.basename(base)
        yield path, basedir, basename, ext

template_decorators.py

import os.path

from gwf import *

from helpers import birc_load, extendpath, with_ext, pathiter


@birc_load(package='mvapich2', version='2.1')
@extendpath('$(pwd)/bin/mrbayes-3.2.6/src')
def mrbayes(inputfile, outputdir):
    treename, _ = os.path.splitext(inputfile)
    logfile = os.path.join(outputdir, treename + '.log')
    treefile = os.path.join(outputdir, treename + '.nex.con.tre')

    opts = {
        'input': [inputfile],
        'output': [treefile, logfile],
        'walltime': '280:00:00',
        'queue': 'normal',
        'cores': 16,
        'memory': '16g'
    }

    spec = """
    cd {outputdir} && \
    MPICH_NEMESIS_NETMOD=tcp mpirun -np 16 mb {inputfile} > {logfile}
    """.format(inputfile=inputfile, logfile=logfile, outputdir=outputdir)

    return opts, spec