<commit_before>#!/usr/bin/env python

import numpy as np
import chainer
import chainer.functions as F
import chainer.links as L
import chainer.datasets.image_dataset as ImageDataset
import six
import os
from PIL import Image

from chainer import cuda, optimizers, serializers, Variable
import cv2

class ImageAndRefDataset(chainer.dataset.DatasetMixin):

    def __init__(self, paths, root1='./input', root2='./ref', dtype=np.float32):
        self._paths = paths
        self._root1 = root1
        self._root2 = root2
        self._dtype = dtype

    def __len__(self):
        return len(self._paths)

    def get_name(self, i):
        return self._paths[i]

    def get_example(self, i, minimize=False, blur=0, s_size=128):
        path1 = os.path.join(self._root1, self._paths[i])
        #image1 = ImageDataset._read_image_as_array(path1, self._dtype)
        image1 = cv2.imread(path1, cv2.IMREAD_GRAYSCALE)
        image1 = np.asarray(image1, self._dtype)

        _image1 = image1.copy()
        if minimize:
            if image1.shape[0] < image1.shape[1]:
                s0 = s_size
                s1 = int(image1.shape[1] * (s_size / image1.shape[0]))
                s1 = s1 - s1 % 16
            else:
                s1 = s_size
                s0 = int(image1.shape[0] * (s_size / image1.shape[1]))
                s0 = s0 - s0 % 16

            _image1 = image1.copy()
            _image1 = cv2.resize(_image1, (4 * s1, 4 * s0),
                                 interpolation=cv2.INTER_AREA)
            #noise = np.random.normal(0,5*np.random.rand(),_image1.shape).astype(self._dtype)

            if blur > 0:
                blured = cv2.blur(_image1, ksize=(blur, blur))
                image1 = _image1 + blured - 255

            image1 = cv2.resize(image1, (s1, s0), interpolation=cv2.INTER_AREA)

        # image is grayscale
        if image1.ndim == 2:
            image1 = image1[:, :, np.newaxis]
        if _image1.ndim == 2:
            _image1 = _image1[:, :, np.newaxis]

        image1 = np.insert(image1, 1, -512, axis=2)
        image1 = np.insert(image1, 2, 128, axis=2)
        image1 = np.insert(image1, 3, 128, axis=2)

        # add color ref image
        path_ref = os.path.join(self._root2, self._paths[i])

        if minimize:
            image_ref = cv2.imread(path_ref, cv2.IMREAD_UNCHANGED)
            image_ref = cv2.resize(image_ref, (image1.shape[1], image1.shape[
                                   0]), interpolation=cv2.INTER_NEAREST)
            b, g, r, a = cv2.split(image_ref)
            image_ref = cv2.cvtColor(cv2.merge((b, g, r)), cv2.COLOR_BGR2YUV)

            for x in range(image1.shape[0]):
                for y in range(image1.shape[1]):
                    if a[x][y] != 0:
                        for ch in range(3):
                            image1[x][y][ch + 1] = image_ref[x][y][ch]

        else:
            image_ref = cv2.imread(path_ref, cv2.IMREAD_COLOR)
            image_ref = cv2.cvtColor(image_ref, cv2.COLOR_BGR2YUV)
            image1 = cv2.resize(
                image1, (4 * image_ref.shape[1], 4 * image_ref.shape[0]), interpolation=cv2.INTER_AREA)
            image_ref = cv2.resize(image_ref, (image1.shape[1], image1.shape[
                                   0]), interpolation=cv2.INTER_AREA)

            for x in range(image1.shape[0]):
                for y in range(image1.shape[1]):
                    for ch in range(3):
                        image1[x][y][ch + 1] = image_ref[x][y][ch]

        return image1.transpose(2, 0, 1),  _image1.transpose(2, 0, 1)

class Image2ImageDataset(chainer.dataset.DatasetMixin):

    def __init__(self, paths, root1='./input', root2='./terget', dtype=np.float32, leak=(0, 0), root_ref = None, train=False):
        if isinstance(paths, six.string_types):
            with open(paths) as paths_file:
                paths = [path.strip() for path in paths_file]
        self._paths = paths
        self._root1 = root1
        self._root2 = root2
        self._root_ref = root_ref
        self._dtype = dtype
        self._leak = leak
        self._img_dict = {}
        self._train = train

    def set_img_dict(self, img_dict):
        self._img_dict = img_dict

    def get_vec(self, name):
        tag_size = 1539
        v = np.zeros(tag_size).astype(np.int32)
        if name in self._img_dict.keys():
            for i in self._img_dict[name][3]:
                v[i] = 1
        return v

    def __len__(self):
        return len(self._paths)

    def get_name(self, i):
        return self._paths[i]

    def get_example(self, i, minimize=False, log=False, bin_r=0):
        if self._train:
            bin_r = 0.9

        readed = False
        if np.random.rand() < bin_r:
            if np.random.rand() < 0.3:
                path1 = os.path.join(self._root1 + "_b2r/", self._paths[i])
            else:
                path1 = os.path.join(self._root1 + "_cnn/", self._paths[i])
            path2 = os.path.join(self._root2 + "_b2r/", self._paths[i])
            image1 = cv2.imread(path1, cv2.IMREAD_GRAYSCALE)
            image2 = cv2.imread(path2, cv2.IMREAD_COLOR)
            if image1 is not None and image2 is not None:
                if image1.shape[0] > 0 and image1.shape[1] and image2.shape[0] > 0 and image2.shape[1]:
                    readed = True
        if not readed:
            path1 = os.path.join(self._root1, self._paths[i])
            path2 = os.path.join(self._root2, self._paths[i])
            image1 = cv2.imread(path1, cv2.IMREAD_GRAYSCALE)
            image2 = cv2.imread(path2, cv2.IMREAD_COLOR)

        image2 = cv2.cvtColor(image2, cv2.COLOR_BGR2YUV)
        name1 = os.path.basename(self._paths[i])

        if self._train and np.random.rand() < 0.2:
            ret, image1 = cv2.threshold(
                image1, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

        # add flip and noise
        if self._train:
            if np.random.rand() > 0.5:
                image1 = cv2.flip(image1, 1)
                image2 = cv2.flip(image2, 1)
            if np.random.rand() > 0.9:
                image1 = cv2.flip(image1, 0)
                image2 = cv2.flip(image2, 0)

        image1 = np.asarray(image1, self._dtype)
        image2 = np.asarray(image2, self._dtype)

        if self._train:
            noise = np.random.normal(
                0, 5 * np.random.rand(), image1.shape).astype(self._dtype)
            image1 += noise
            noise = np.random.normal(
                0, 5 * np.random.rand(), image2.shape).astype(self._dtype)
            image2 += noise
            noise = np.random.normal(0, 16)
            image1 += noise
            image1[image1 < 0] = 0

        # image is grayscale
        if image1.ndim == 2:
            image1 = image1[:, :, np.newaxis]
        if image2.ndim == 2:
            image2 = image2[:, :, np.newaxis]

        image1 = np.insert(image1, 1, -512, axis=2)
        image1 = np.insert(image1, 2, 128, axis=2)
        image1 = np.insert(image1, 3, 128, axis=2)

        # randomly add terget image px
        if self._leak[1] > 0:
            image0 = image1
            n = np.random.randint(16, self._leak[1])
            if self._train:
                r = np.random.rand()
                if r < 0.4:
                    n = 0
                elif r < 0.7:
                    n = np.random.randint(2, 16)

            x = np.random.randint(1, image1.shape[0] - 1, n)
            y = np.random.randint(1, image1.shape[1] - 1, n)
            for i in range(n):
                for ch in range(3):
                    d = 20
                    v = image2[x[i]][y[i]][ch] + np.random.normal(0, 5)
                    v = np.floor(v / d + 0.5) * d
                    image1[x[i]][y[i]][ch + 1] = v
                if np.random.rand() > 0.5:
                    for ch in range(3):
                        image1[x[i]][y[i] + 1][ch +
                                               1] = image1[x[i]][y[i]][ch + 1]
                        image1[x[i]][y[i] - 1][ch +
                                               1] = image1[x[i]][y[i]][ch + 1]
                if np.random.rand() > 0.5:
                    for ch in range(3):
                        image1[x[i] + 1][y[i]][ch +
                                               1] = image1[x[i]][y[i]][ch + 1]
                        image1[x[i] - 1][y[i]][ch +
                                               1] = image1[x[i]][y[i]][ch + 1]

        image1 = (image1.transpose(2, 0, 1))
        image2 = (image2.transpose(2, 0, 1))
        #image1 = (image1.transpose(2, 0, 1) -128) /128
        #image2 = (image2.transpose(2, 0, 1) -128) /128

        return image1, image2  # ,vec

class Image2ImageDatasetX2(Image2ImageDataset):

    def get_example(self, i, minimize=False, log=False, bin_r=0):
        path1 = os.path.join(self._root1, self._paths[i])
        path2 = os.path.join(self._root2, self._paths[i])
        #image1 = ImageDataset._read_image_as_array(path1, self._dtype)
        image1 = cv2.imread(path1, cv2.IMREAD_GRAYSCALE)
        image2 = cv2.imread(path2, cv2.IMREAD_COLOR)
        image2 = cv2.cvtColor(image2, cv2.COLOR_BGR2YUV)
        image2 = np.asarray(image2, self._dtype)
        name1 = os.path.basename(self._paths[i])
        vec = self.get_vec(name1)

        # add flip and noise
        if self._train:
            if np.random.rand() > 0.5:
                image1 = cv2.flip(image1, 1)
                image2 = cv2.flip(image2, 1)
            if np.random.rand() > 0.8:
                image1 = cv2.flip(image1, 0)
                image2 = cv2.flip(image2, 0)

        if self._train:
            bin_r = 0.3
        if np.random.rand() < bin_r:
            ret, image1 = cv2.threshold(
                image1, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

        _image1 = image1.copy()
        _image2 = image2.copy()
        image1 = cv2.resize(image1, (128, 128), interpolation=cv2.INTER_AREA)
        image2 = cv2.resize(image2, (128, 128), interpolation=cv2.INTER_AREA)

        image1 = np.asarray(image1, self._dtype)
        _image1 = np.asarray(_image1, self._dtype)

        if self._train:
            noise = np.random.normal(0, 5, image1.shape).astype(self._dtype)
            image1 = image1 + noise
            noise = np.random.normal(0, 5, image2.shape).astype(self._dtype)
            image2 = image2 + noise
            noise = np.random.normal(
                0, 4 * np.random.rand(), _image1.shape).astype(self._dtype)
            noise += np.random.normal(0, 24)
            _image1 = _image1 + noise
            _image1[_image1 < 0] = 0
            _image1[_image1 > 255] = 255

        # image is grayscale
        if image1.ndim == 2:
            image1 = image1[:, :, np.newaxis]
        if image2.ndim == 2:
            image2 = image2[:, :, np.newaxis]
        if _image1.ndim == 2:
            _image1 = _image1[:, :, np.newaxis]
        if _image2.ndim == 2:
            _image2 = _image2[:, :, np.newaxis]

        image1 = np.insert(image1, 1, -512, axis=2)
        image1 = np.insert(image1, 2, 128, axis=2)
        image1 = np.insert(image1, 3, 128, axis=2)

        # randomly add terget image px
        if self._leak[1] > 0:
            image0 = image1
            n = np.random.randint(self._leak[0], self._leak[1])
            x = np.random.randint(1, image1.shape[0] - 1, n)
            y = np.random.randint(1, image1.shape[1] - 1, n)
            for i in range(n):
                for ch in range(3):
                    d = 20
                    v = image2[x[i]][y[i]][ch] + np.random.normal(0, 5)
                    #v = np.random.normal(128,40)
                    v = np.floor(v / d + 0.5) * d
                    image1[x[i]][y[i]][ch + 1] = v
                    if np.random.rand() > 0.5:
                        image1[x[i]][y[i] + 1][ch + 1] = v
                        image1[x[i]][y[i] - 1][ch + 1] = v
                    if np.random.rand() > 0.5:
                        image1[x[i] + 1][y[i]][ch + 1] = v
                        image1[x[i] - 1][y[i]][ch + 1] = v

        return image1.transpose(2, 0, 1), image2.transpose(2, 0, 1), _image1.transpose(2, 0, 1), _image2.transpose(2, 0, 1)
<commit_msg>Remove PIL import<commit_after>#!/usr/bin/env python

import numpy as np
import chainer
import chainer.functions as F
import chainer.links as L
import chainer.datasets.image_dataset as ImageDataset
import six
import os

from chainer import cuda, optimizers, serializers, Variable
import cv2

class ImageAndRefDataset(chainer.dataset.DatasetMixin):

    def __init__(self, paths, root1='./input', root2='./ref', dtype=np.float32):
        self._paths = paths
        self._root1 = root1
        self._root2 = root2
        self._dtype = dtype

    def __len__(self):
        return len(self._paths)

    def get_name(self, i):
        return self._paths[i]

    def get_example(self, i, minimize=False, blur=0, s_size=128):
        path1 = os.path.join(self._root1, self._paths[i])
        #image1 = ImageDataset._read_image_as_array(path1, self._dtype)
        image1 = cv2.imread(path1, cv2.IMREAD_GRAYSCALE)
        image1 = np.asarray(image1, self._dtype)

        _image1 = image1.copy()
        if minimize:
            if image1.shape[0] < image1.shape[1]:
                s0 = s_size
                s1 = int(image1.shape[1] * (s_size / image1.shape[0]))
                s1 = s1 - s1 % 16
            else:
                s1 = s_size
                s0 = int(image1.shape[0] * (s_size / image1.shape[1]))
                s0 = s0 - s0 % 16

            _image1 = image1.copy()
            _image1 = cv2.resize(_image1, (4 * s1, 4 * s0),
                                 interpolation=cv2.INTER_AREA)
            #noise = np.random.normal(0,5*np.random.rand(),_image1.shape).astype(self._dtype)

            if blur > 0:
                blured = cv2.blur(_image1, ksize=(blur, blur))
                image1 = _image1 + blured - 255

            image1 = cv2.resize(image1, (s1, s0), interpolation=cv2.INTER_AREA)

        # image is grayscale
        if image1.ndim == 2:
            image1 = image1[:, :, np.newaxis]
        if _image1.ndim == 2:
            _image1 = _image1[:, :, np.newaxis]

        image1 = np.insert(image1, 1, -512, axis=2)
        image1 = np.insert(image1, 2, 128, axis=2)
        image1 = np.insert(image1, 3, 128, axis=2)

        # add color ref image
        path_ref = os.path.join(self._root2, self._paths[i])

        if minimize:
            image_ref = cv2.imread(path_ref, cv2.IMREAD_UNCHANGED)
            image_ref = cv2.resize(image_ref, (image1.shape[1], image1.shape[
                                   0]), interpolation=cv2.INTER_NEAREST)
            b, g, r, a = cv2.split(image_ref)
            image_ref = cv2.cvtColor(cv2.merge((b, g, r)), cv2.COLOR_BGR2YUV)

            for x in range(image1.shape[0]):
                for y in range(image1.shape[1]):
                    if a[x][y] != 0:
                        for ch in range(3):
                            image1[x][y][ch + 1] = image_ref[x][y][ch]

        else:
            image_ref = cv2.imread(path_ref, cv2.IMREAD_COLOR)
            image_ref = cv2.cvtColor(image_ref, cv2.COLOR_BGR2YUV)
            image1 = cv2.resize(
                image1, (4 * image_ref.shape[1], 4 * image_ref.shape[0]), interpolation=cv2.INTER_AREA)
            image_ref = cv2.resize(image_ref, (image1.shape[1], image1.shape[
                                   0]), interpolation=cv2.INTER_AREA)

            for x in range(image1.shape[0]):
                for y in range(image1.shape[1]):
                    for ch in range(3):
                        image1[x][y][ch + 1] = image_ref[x][y][ch]

        return image1.transpose(2, 0, 1),  _image1.transpose(2, 0, 1)

class Image2ImageDataset(chainer.dataset.DatasetMixin):

    def __init__(self, paths, root1='./input', root2='./terget', dtype=np.float32, leak=(0, 0), root_ref = None, train=False):
        if isinstance(paths, six.string_types):
            with open(paths) as paths_file:
                paths = [path.strip() for path in paths_file]
        self._paths = paths
        self._root1 = root1
        self._root2 = root2
        self._root_ref = root_ref
        self._dtype = dtype
        self._leak = leak
        self._img_dict = {}
        self._train = train

    def set_img_dict(self, img_dict):
        self._img_dict = img_dict

    def get_vec(self, name):
        tag_size = 1539
        v = np.zeros(tag_size).astype(np.int32)
        if name in self._img_dict.keys():
            for i in self._img_dict[name][3]:
                v[i] = 1
        return v

    def __len__(self):
        return len(self._paths)

    def get_name(self, i):
        return self._paths[i]

    def get_example(self, i, minimize=False, log=False, bin_r=0):
        if self._train:
            bin_r = 0.9

        readed = False
        if np.random.rand() < bin_r:
            if np.random.rand() < 0.3:
                path1 = os.path.join(self._root1 + "_b2r/", self._paths[i])
            else:
                path1 = os.path.join(self._root1 + "_cnn/", self._paths[i])
            path2 = os.path.join(self._root2 + "_b2r/", self._paths[i])
            image1 = cv2.imread(path1, cv2.IMREAD_GRAYSCALE)
            image2 = cv2.imread(path2, cv2.IMREAD_COLOR)
            if image1 is not None and image2 is not None:
                if image1.shape[0] > 0 and image1.shape[1] and image2.shape[0] > 0 and image2.shape[1]:
                    readed = True
        if not readed:
            path1 = os.path.join(self._root1, self._paths[i])
            path2 = os.path.join(self._root2, self._paths[i])
            image1 = cv2.imread(path1, cv2.IMREAD_GRAYSCALE)
            image2 = cv2.imread(path2, cv2.IMREAD_COLOR)

        image2 = cv2.cvtColor(image2, cv2.COLOR_BGR2YUV)
        name1 = os.path.basename(self._paths[i])

        if self._train and np.random.rand() < 0.2:
            ret, image1 = cv2.threshold(
                image1, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

        # add flip and noise
        if self._train:
            if np.random.rand() > 0.5:
                image1 = cv2.flip(image1, 1)
                image2 = cv2.flip(image2, 1)
            if np.random.rand() > 0.9:
                image1 = cv2.flip(image1, 0)
                image2 = cv2.flip(image2, 0)

        image1 = np.asarray(image1, self._dtype)
        image2 = np.asarray(image2, self._dtype)

        if self._train:
            noise = np.random.normal(
                0, 5 * np.random.rand(), image1.shape).astype(self._dtype)
            image1 += noise
            noise = np.random.normal(
                0, 5 * np.random.rand(), image2.shape).astype(self._dtype)
            image2 += noise
            noise = np.random.normal(0, 16)
            image1 += noise
            image1[image1 < 0] = 0

        # image is grayscale
        if image1.ndim == 2:
            image1 = image1[:, :, np.newaxis]
        if image2.ndim == 2:
            image2 = image2[:, :, np.newaxis]

        image1 = np.insert(image1, 1, -512, axis=2)
        image1 = np.insert(image1, 2, 128, axis=2)
        image1 = np.insert(image1, 3, 128, axis=2)

        # randomly add terget image px
        if self._leak[1] > 0:
            image0 = image1
            n = np.random.randint(16, self._leak[1])
            if self._train:
                r = np.random.rand()
                if r < 0.4:
                    n = 0
                elif r < 0.7:
                    n = np.random.randint(2, 16)

            x = np.random.randint(1, image1.shape[0] - 1, n)
            y = np.random.randint(1, image1.shape[1] - 1, n)
            for i in range(n):
                for ch in range(3):
                    d = 20
                    v = image2[x[i]][y[i]][ch] + np.random.normal(0, 5)
                    v = np.floor(v / d + 0.5) * d
                    image1[x[i]][y[i]][ch + 1] = v
                if np.random.rand() > 0.5:
                    for ch in range(3):
                        image1[x[i]][y[i] + 1][ch +
                                               1] = image1[x[i]][y[i]][ch + 1]
                        image1[x[i]][y[i] - 1][ch +
                                               1] = image1[x[i]][y[i]][ch + 1]
                if np.random.rand() > 0.5:
                    for ch in range(3):
                        image1[x[i] + 1][y[i]][ch +
                                               1] = image1[x[i]][y[i]][ch + 1]
                        image1[x[i] - 1][y[i]][ch +
                                               1] = image1[x[i]][y[i]][ch + 1]

        image1 = (image1.transpose(2, 0, 1))
        image2 = (image2.transpose(2, 0, 1))
        #image1 = (image1.transpose(2, 0, 1) -128) /128
        #image2 = (image2.transpose(2, 0, 1) -128) /128

        return image1, image2  # ,vec

class Image2ImageDatasetX2(Image2ImageDataset):

    def get_example(self, i, minimize=False, log=False, bin_r=0):
        path1 = os.path.join(self._root1, self._paths[i])
        path2 = os.path.join(self._root2, self._paths[i])
        #image1 = ImageDataset._read_image_as_array(path1, self._dtype)
        image1 = cv2.imread(path1, cv2.IMREAD_GRAYSCALE)
        image2 = cv2.imread(path2, cv2.IMREAD_COLOR)
        image2 = cv2.cvtColor(image2, cv2.COLOR_BGR2YUV)
        image2 = np.asarray(image2, self._dtype)
        name1 = os.path.basename(self._paths[i])
        vec = self.get_vec(name1)

        # add flip and noise
        if self._train:
            if np.random.rand() > 0.5:
                image1 = cv2.flip(image1, 1)
                image2 = cv2.flip(image2, 1)
            if np.random.rand() > 0.8:
                image1 = cv2.flip(image1, 0)
                image2 = cv2.flip(image2, 0)

        if self._train:
            bin_r = 0.3
        if np.random.rand() < bin_r:
            ret, image1 = cv2.threshold(
                image1, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

        _image1 = image1.copy()
        _image2 = image2.copy()
        image1 = cv2.resize(image1, (128, 128), interpolation=cv2.INTER_AREA)
        image2 = cv2.resize(image2, (128, 128), interpolation=cv2.INTER_AREA)

        image1 = np.asarray(image1, self._dtype)
        _image1 = np.asarray(_image1, self._dtype)

        if self._train:
            noise = np.random.normal(0, 5, image1.shape).astype(self._dtype)
            image1 = image1 + noise
            noise = np.random.normal(0, 5, image2.shape).astype(self._dtype)
            image2 = image2 + noise
            noise = np.random.normal(
                0, 4 * np.random.rand(), _image1.shape).astype(self._dtype)
            noise += np.random.normal(0, 24)
            _image1 = _image1 + noise
            _image1[_image1 < 0] = 0
            _image1[_image1 > 255] = 255

        # image is grayscale
        if image1.ndim == 2:
            image1 = image1[:, :, np.newaxis]
        if image2.ndim == 2:
            image2 = image2[:, :, np.newaxis]
        if _image1.ndim == 2:
            _image1 = _image1[:, :, np.newaxis]
        if _image2.ndim == 2:
            _image2 = _image2[:, :, np.newaxis]

        image1 = np.insert(image1, 1, -512, axis=2)
        image1 = np.insert(image1, 2, 128, axis=2)
        image1 = np.insert(image1, 3, 128, axis=2)

        # randomly add terget image px
        if self._leak[1] > 0:
            image0 = image1
            n = np.random.randint(self._leak[0], self._leak[1])
            x = np.random.randint(1, image1.shape[0] - 1, n)
            y = np.random.randint(1, image1.shape[1] - 1, n)
            for i in range(n):
                for ch in range(3):
                    d = 20
                    v = image2[x[i]][y[i]][ch] + np.random.normal(0, 5)
                    #v = np.random.normal(128,40)
                    v = np.floor(v / d + 0.5) * d
                    image1[x[i]][y[i]][ch + 1] = v
                    if np.random.rand() > 0.5:
                        image1[x[i]][y[i] + 1][ch + 1] = v
                        image1[x[i]][y[i] - 1][ch + 1] = v
                    if np.random.rand() > 0.5:
                        image1[x[i] + 1][y[i]][ch + 1] = v
                        image1[x[i] - 1][y[i]][ch + 1] = v

        return image1.transpose(2, 0, 1), image2.transpose(2, 0, 1), _image1.transpose(2, 0, 1), _image2.transpose(2, 0, 1)
<|endoftext|>


<commit_before>#!/usr/bin/env python
#
# tools.py
#
# Reimplementations of Santi's trim_longitudes.py and blastclust_lst2fasta.py
# scripts/functions, and other miscellaneous functions
#
# (c) The James Hutton Institute 2016
# Author: Leighton Pritchard and Peter Thorpe

import os

import gzip
from subprocess import check_output, CalledProcessError
import hashlib
import sys

from collections import defaultdict
from optparse import OptionParser

from Bio import SeqIO
from Bio.SeqRecord import SeqRecord

class NotExecutableError(Exception):
    """Exception raised when expected executable is not executable"""
    def __init__(self, message):
        self.message = message

def is_exe(filename):
    """Returns True if path is to an executable file"""
    if os.path.isfile(filename) and os.access(filename, os.X_OK):
        return True
    else:
        try:
            exefile = check_output(["which", filename]).strip()
        except CalledProcessError:
            raise NotExecutableError("{0} does not exist".format(filename))
    return os.path.isfile(exefile) and os.access(exefile, os.X_OK)

# Defaults  settings are for the current ITS1 phy
# primers and the development ITS1 Phy database. DON'T BLINDLY USE THESE
# The filenames, we now provide directly.
def trim_seq(infname, outfname, lclip=53, rclip=0, minlen=100):
    """Trims all FASTA sequences in infname by lclip and rclip on the
    'left'- and 'right'-hand ends, respectively and writes the results to
    outfname. Any clipped sequence with length < minlen is not written.

    Defaults are equivalent to the default settings are for the current ITS1
    phy primers and the development ITS1 Phy database. PLEASE DONT BLINDLY
    USE THESE. The filenames, we now provide directly.
    """
    # ensure these are int, as they may have been passed as a string.
    lclip = int(lclip)
    rclip = int(rclip)
    with open(infname, 'r') as fh:
        # rclip needs to be set up to allow for zero values, using logic:
        # rclip_coordinate = len(s) - rclip
        # Use generators to save memory
        s_trim = (s[lclip:(len(s) - rclip)] for s in SeqIO.parse(fh, 'fasta'))
        return SeqIO.write((s for s in s_trim if len(s) >= minlen),
                           outfname, 'fasta')

def convert_fq_to_fa(in_file, out_file):
    """ Function to convert fq to fa
    -i the fastq file to be converted
    -o the desired out fasta file name
    the function will auto detect if the file is .gz
    requires biopython
    """
    # if the file is .gz we need to deal with it
    if in_file.endswith(".gz"):
        in_file = gzip.open(in_file, mode='rt', compresslevel=9,
                            encoding=None, errors=None,
                            newline=None)
    SeqIO.convert(in_file, "fastq", out_file, "fasta")

def deduplicate_and_rename(seqlist):
    """Removes duplicates from the passed SeqRecords, and replaces sequence
    IDs with the md5 hash of the sequence, suffixed with the number of
    sequences that were replaced. Returns a tuple (sequences, names_old_to_new)
    where names_old_to_new is a list of tuples of (old_name, hash) pairs.

    - sequences   Iterable of Bio.SeqRecord objects
    """
    names_old_to_new = list()
    abundance = defaultdict(int)
    hash_to_seq = defaultdict(str)
    hash_to_name = defaultdict(str)
    # compile hashed sequence data - the loop is necessary because we need
    # to run through completely once to obtain abundance info. If memory
    # to store this becomes an issue, we might want to try an alternative
    # approach
    for seq in seqlist:
        seqhash = hashlib.md5(str(seq.seq).encode()).hexdigest()
        abundance[seqhash] += 1
        hash_to_seq[seqhash] = seq.seq
        hash_to_name[seqhash] = seq.id
        database_output = ("{0}\t{1}\n".format(seq.id, seqhash))
        names_old_to_new.append(database_output)
    # generate list of sequences, named for the hash, including abundance data,
    # and return with ID:hash lookup
    seqlist = list()
    for name, abundance_val in abundance.items():
        seqname = "{0}_{1}".format(name, abundance_val)
        seqlist.append(SeqRecord(id=seqname, description="",
                                 seq=hash_to_seq[name]))
    return(seqlist, names_old_to_new, hash_to_seq)

def dereplicate_name(fasta, database_out, out):
    """function to dereplicate the seq. Thus generating an
    abundance of that seq. Rename the seq to something
    swarm will work with.
    - fasta     - fasta file of assembled seq
    - database_out   - ouput old name to new file.
    -out        - outfile
    """
    # open files to write to
    fasta_out = open(out, 'w')
    name_out = open(database_out, "w")
    # convert the file to a list of Seqrecord objects
    seqlist = list(SeqIO.parse(fasta, 'fasta'))
    seqlist, names_old_to_new, hash_to_seq = deduplicate_and_rename(seqlist)
    for i in (names_old_to_new):
        name_out.write(i)
    for seq_record in (seqlist):
        SeqIO.write(seq_record, fasta_out, "fasta")
    # close the open files
    fasta_out.close()
    name_out.close()

def check_OTU_db_abundance_val(OTUfasta):
    """this is a function to check that the OTU database has _1
    at the end of the title, this is required for swamr clustering
    takes in a fasta and if required, writes out a new one"""
    # we will iterate through first to see if we need to rewrite it.
    bad_count = 0
    for seq_record in SeqIO.parse(OTUfasta, "fasta"):
        if not seq_record.id.endswith("_1"):
            # print ("OTU db %s missing abundance value" % seq_record.id)
            bad_count = bad_count + 1
    if bad_count > 0:
        # only if we need to write to a new file, we will.
        outfile = OTUfasta.split(".fa")[0] + "abundance.fasta"
        f_out = open(outfile, 'w')
        for seq_record in SeqIO.parse(OTUfasta, "fasta"):
            seq_record.id = seq_record.id + "_1"
            seq_record.description = ""
            SeqIO.write(seq_record, f_out, "fasta")
        f_out.close()
        return outfile
    else:
        return OTUfasta

# Function replacing Santi's blastclust_lst2fasta.py script
def blastclust_to_fasta(infname, seqfname, outdir):
    """Converts input BLASTCLUST output list to a subdirectory of FASTA files,
    each of which contains all sequences from a single cluster. The sequences
    matching the IDs listed in BLASTCLUST output should all be in the file
    seqfname.
    """
    outdirname = os.path.join(outdir, "blastclust_OTUs")
    if not os.path.exists(outdirname):
        os.makedirs(outdirname)
    seqdict = SeqIO.index(seqfname, 'fasta')
    with open(infname, 'r') as fh:
        otu_id = 0
        for line in fh:
            otu_id += 1
            outfname = os.path.join(outdirname,
                                    "blastclust_OTU_%06d.fasta" % otu_id)
            SeqIO.write((seqdict[key] for key in line.split()),
                        outfname, 'fasta')
    return outdirname

# the following three function are to rename the clusters back to their
# original names
def coded_name_to_species(database_file):
    """functiong takes the already generated tab separated
    database of coded name to species file. Returns a dic
    of coded_name to species"""
    with open(database_file) as file:
        data = file.read().split("\n")
    coded_name_to_species_dict = dict()
    for line in data:
        if not line.strip():
            continue  # if the last line is blank
        if line.startswith("#"):
            continue
        coded_name, species = line.split("\t")
        coded_name_to_species_dict[coded_name.rstrip()] = species.rstrip()
    return coded_name_to_species_dict

def parse_tab_file_get_clusters(filename1, database, out_file):
    """function to open up a tab or space separeted clustering
    output and rename according to the name in the database file.
    Abundance is also appended to the name"""
    # call the function to get the dictionary
    # populated with the database
    coded_name_to_species_dict = coded_name_to_species(database)
    cluster_file = open(filename1, "r")
    summary_out_file = open(out_file, "w")

    count = int(0)
    # iterate through the cluster file. One line per cluster
    for line in cluster_file:
        if not line.strip():
            continue  # if the last line is blank
        if line.startswith("#"):  # dont want comment lines
            continue
        output_str = ""
        if "\t" in line:
            cluster_line = line.rstrip("\n").split("\t")
        else:
            # different clustering program?
            cluster_line = line.rstrip("\n").split()
        count += 1
        for member in cluster_line:
            try:
                # data would be: 2f1454d16278fda2d44f26ebf7a0ed05_310
                # _abundance value
                split_name = member.split("_")[:-1]
                species = ("_").join(split_name)
                abundance = member.split("_")[-1]
                # print (abundance)
            except:
                KeyError
                errmsg = """something went wrong with decoding the names maybe
                your names were not separated in your database? If so, the
                block of code above this staement need adjusting accordingly.
                """
                sys.exit(os.system(errmsg))
            # add the info to a str, we will write at the
            # end of the cluster line
            cluster_summary = "%s_abundance=%s\t" % (species, abundance)
            output_str = output_str + cluster_summary
        summary_out_file.write(output_str+"\n")
    # close the files
    cluster_file.close()
    summary_out_file.close()
<commit_msg>altered way of tracking old to new names<commit_after>#!/usr/bin/env python
#
# tools.py
#
# Reimplementations of Santi's trim_longitudes.py and blastclust_lst2fasta.py
# scripts/functions, and other miscellaneous functions
#
# (c) The James Hutton Institute 2016
# Author: Leighton Pritchard and Peter Thorpe

import os

import gzip
from subprocess import check_output, CalledProcessError
import hashlib
import sys

from collections import defaultdict
from optparse import OptionParser

from Bio import SeqIO
from Bio.SeqRecord import SeqRecord

class NotExecutableError(Exception):
    """Exception raised when expected executable is not executable"""
    def __init__(self, message):
        self.message = message

def is_exe(filename):
    """Returns True if path is to an executable file"""
    if os.path.isfile(filename) and os.access(filename, os.X_OK):
        return True
    else:
        try:
            exefile = check_output(["which", filename]).strip()
        except CalledProcessError:
            raise NotExecutableError("{0} does not exist".format(filename))
    return os.path.isfile(exefile) and os.access(exefile, os.X_OK)

# Defaults  settings are for the current ITS1 phy
# primers and the development ITS1 Phy database. DON'T BLINDLY USE THESE
# The filenames, we now provide directly.
def trim_seq(infname, outfname, lclip=53, rclip=0, minlen=100):
    """Trims all FASTA sequences in infname by lclip and rclip on the
    'left'- and 'right'-hand ends, respectively and writes the results to
    outfname. Any clipped sequence with length < minlen is not written.

    Defaults are equivalent to the default settings are for the current ITS1
    phy primers and the development ITS1 Phy database. PLEASE DONT BLINDLY
    USE THESE. The filenames, we now provide directly.
    """
    # ensure these are int, as they may have been passed as a string.
    lclip = int(lclip)
    rclip = int(rclip)
    with open(infname, 'r') as fh:
        # rclip needs to be set up to allow for zero values, using logic:
        # rclip_coordinate = len(s) - rclip
        # Use generators to save memory
        s_trim = (s[lclip:(len(s) - rclip)] for s in SeqIO.parse(fh, 'fasta'))
        return SeqIO.write((s for s in s_trim if len(s) >= minlen),
                           outfname, 'fasta')

def convert_fq_to_fa(in_file, out_file):
    """ Function to convert fq to fa
    -i the fastq file to be converted
    -o the desired out fasta file name
    the function will auto detect if the file is .gz
    requires biopython
    """
    # if the file is .gz we need to deal with it
    if in_file.endswith(".gz"):
        in_file = gzip.open(in_file, mode='rt', compresslevel=9,
                            encoding=None, errors=None,
                            newline=None)
    SeqIO.convert(in_file, "fastq", out_file, "fasta")

def deduplicate_and_rename(seqlist):
    """Removes duplicates from the passed SeqRecords, and replaces sequence
    IDs with the md5 hash of the sequence, suffixed with the number of
    sequences that were replaced. Returns a tuple (sequences, names_old_to_new)
    where names_old_to_new is a list of tuples of (old_name, hash) pairs.

    - sequences   Iterable of Bio.SeqRecord objects
    """
    abundance = defaultdict(int)
    hash_to_seq = defaultdict(str)
    hash_to_name = defaultdict(str)
    new_to_old = defaultdict(list)
    # compile hashed sequence data - the loop is necessary because we need
    # to run through completely once to obtain abundance info. If memory
    # to store this becomes an issue, we might want to try an alternative
    # approach
    for seq in seqlist:
        seqhash = hashlib.md5(str(seq.seq).encode()).hexdigest()
        abundance[seqhash] += 1
        hash_to_seq[seqhash] = seq.seq
        hash_to_name[seqhash] = seq.id
        new_to_old[seqhash].append(seq.id)
    # generate list of sequences, named for the hash, including abundance data,
    # and return with ID:hash lookup
    seqlist = list()
    for name, abundance_val in abundance.items():
        seqname = "{0}_{1}".format(name, abundance_val)
        seqlist.append(SeqRecord(id=seqname, description="",
                                 seq=hash_to_seq[name]))
    return(seqlist, new_to_old, hash_to_seq)

def dereplicate_name(fasta, database_out, out):
    """function to dereplicate the seq. Thus generating an
    abundance of that seq. Rename the seq to something
    swarm will work with.
    - fasta     - fasta file of assembled seq
    - database_out   - ouput old name to new file.
    -out        - outfile
    """
    # open files to write to
    fasta_out = open(out, 'w')
    name_out = open(database_out, "w")
    # convert the file to a list of Seqrecord objects
    seqlist = list(SeqIO.parse(fasta, 'fasta'))
    seqlist, new_to_old, hash_to_seq = deduplicate_and_rename(seqlist)
    for key, vals in new_to_old.items():
        out_data = "%s\t%s\n" % (key, "\t".join(vals))
        name_out.write(out_data)
    for seq_record in (seqlist):
        SeqIO.write(seq_record, fasta_out, "fasta")
    # close the open files
    fasta_out.close()
    name_out.close()

def check_OTU_db_abundance_val(OTUfasta):
    """this is a function to check that the OTU database has _1
    at the end of the title, this is required for swamr clustering
    takes in a fasta and if required, writes out a new one"""
    # we will iterate through first to see if we need to rewrite it.
    bad_count = 0
    for seq_record in SeqIO.parse(OTUfasta, "fasta"):
        if not seq_record.id.endswith("_1"):
            # print ("OTU db %s missing abundance value" % seq_record.id)
            bad_count = bad_count + 1
    if bad_count > 0:
        # only if we need to write to a new file, we will.
        outfile = OTUfasta.split(".fa")[0] + "abundance.fasta"
        f_out = open(outfile, 'w')
        for seq_record in SeqIO.parse(OTUfasta, "fasta"):
            seq_record.id = seq_record.id + "_1"
            seq_record.description = ""
            SeqIO.write(seq_record, f_out, "fasta")
        f_out.close()
        return outfile
    else:
        return OTUfasta

# Function replacing Santi's blastclust_lst2fasta.py script
def blastclust_to_fasta(infname, seqfname, outdir):
    """Converts input BLASTCLUST output list to a subdirectory of FASTA files,
    each of which contains all sequences from a single cluster. The sequences
    matching the IDs listed in BLASTCLUST output should all be in the file
    seqfname.
    """
    outdirname = os.path.join(outdir, "blastclust_OTUs")
    if not os.path.exists(outdirname):
        os.makedirs(outdirname)
    seqdict = SeqIO.index(seqfname, 'fasta')
    with open(infname, 'r') as fh:
        otu_id = 0
        for line in fh:
            otu_id += 1
            outfname = os.path.join(outdirname,
                                    "blastclust_OTU_%06d.fasta" % otu_id)
            SeqIO.write((seqdict[key] for key in line.split()),
                        outfname, 'fasta')
    return outdirname

# the following three function are to rename the clusters back to their
# original names
def coded_name_to_species(database_file):
    """functiong takes the already generated tab separated
    database of coded name to species file. Returns a dic
    of coded_name to species"""
    with open(database_file) as file:
        data = file.read().split("\n")
    coded_name_to_species_dict = dict()
    for line in data:
        if not line.strip():
            continue  # if the last line is blank
        if line.startswith("#"):
            continue
        coded_name, species = line.split("\t")
        coded_name_to_species_dict[coded_name.rstrip()] = species.rstrip()
    return coded_name_to_species_dict

def parse_tab_file_get_clusters(filename1, database, out_file):
    """function to open up a tab or space separeted clustering
    output and rename according to the name in the database file.
    Abundance is also appended to the name"""
    # call the function to get the dictionary
    # populated with the database
    coded_name_to_species_dict = coded_name_to_species(database)
    cluster_file = open(filename1, "r")
    summary_out_file = open(out_file, "w")

    count = int(0)
    # iterate through the cluster file. One line per cluster
    for line in cluster_file:
        if not line.strip():
            continue  # if the last line is blank
        if line.startswith("#"):  # dont want comment lines
            continue
        output_str = ""
        if "\t" in line:
            cluster_line = line.rstrip("\n").split("\t")
        else:
            # different clustering program?
            cluster_line = line.rstrip("\n").split()
        count += 1
        for member in cluster_line:
            try:
                # data would be: 2f1454d16278fda2d44f26ebf7a0ed05_310
                # _abundance value
                split_name = member.split("_")[:-1]
                species = ("_").join(split_name)
                abundance = member.split("_")[-1]
                # print (abundance)
            except:
                KeyError
                errmsg = """something went wrong with decoding the names maybe
                your names were not separated in your database? If so, the
                block of code above this staement need adjusting accordingly.
                """
                sys.exit(os.system(errmsg))
            # add the info to a str, we will write at the
            # end of the cluster line
            cluster_summary = "%s_abundance=%s\t" % (species, abundance)
            output_str = output_str + cluster_summary
        summary_out_file.write(output_str+"\n")
    # close the files
    cluster_file.close()
    summary_out_file.close()
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Utilities relating to online services"""
from __future__ import (print_function, unicode_literals)

import requests

def url_check(url=None, method='get', timeout=None, expected_codes=200):
    """Check if a request for a URL gives an expected response"""
    if expected_codes:
        if not isinstance(expected_codes, list):
            expected_codes = expected_codes,
        response = None
        if method:
            if method.lower() == 'get':
                response = requests.get(url=url, timeout=timeout)
            if method.lower() == 'head':
                response = requests.head(url=url, timeout=timeout)
        if response.status_code in expected_codes:
            return True
<commit_msg>fix healthcheck (allow insecure https).<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Utilities relating to online services"""
from __future__ import (print_function, unicode_literals)

import time

import requests

def url_check(url=None, method='get', timeout=None, expected_codes=200, interval=None, healthy_threshold=None):
    """Check if a request for a URL gives an expected response"""
    if expected_codes:
        if not isinstance(expected_codes, list):
            expected_codes = expected_codes,
        response = None
        wait_timeout = time.time() + timeout
        healthy_count = 0
        while wait_timeout > time.time():
            try:
                if method:
                    if method.lower() == 'get':
                        response = requests.get(url=url, timeout=10, verify=False)
                    if method.lower() == 'head':
                        response = requests.head(url=url, timeout=10, verify=False)
                if response.status_code in expected_codes:
                    healthy_count += 1
                if healthy_count == healthy_threshold:
                    return True
            except:
                healthy_count = 0
                time.sleep(5)
<|endoftext|>


<commit_before>#!/usr/bin/env python
# Created by: Lee Bergstrand
# Descript: A Bio-Python program that takes a list of query proteins and uses local BLASTp to search
#           for highly similer proteins within a local blast database (usally a local db of a target
#           proteome). The program then BLASTps backward from the found subject protein to the proteome
#           for which the original query protein if found in order to confirm gene orthology.
#
# Requirements: - This program requires the Biopython module: http://biopython.org/wiki/Download
#               - This script requires BLAST+ 2.2.9 or later.
#               - All operations are done with protien sequences.
#               - All query proteins should be from sequenced genomes in order to facilitate backwards BLAST.
#               - MakeBlastDB must be used to create BLASTp databases for both query and subject proteomes.
#               - BLAST databases require the FASTA file they were made from to be in the same directory.
#
# Usage: BackBLAST.py <queryGeneList.faa> <queryProteomes.csv> <subject1.faa>
# Example: BackBLAST.py queryGeneList.faa queryProteomes.csv AUUJ00000000.faa
#----------------------------------------------------------------------------------------
#===========================================================================================================
#Imports:
import sys
import csv
import subprocess
from Bio import SeqIO
import time
#===========================================================================================================
# Functions:

# 1: Checks if in proper number of arguments are passed gives instructions on proper use.
def argsCheck():
	if len(sys.argv) < 4:
		print "Orthologous Gene Finder"
		print "By Lee Bergstrand\n"
		print "Please refer to source code for documentation\n"
		print "Usage: " + sys.argv[0] + " <queryProteomes.csv> <queryGeneList.faa> <subject1.faa>\n"
		print "Examples:" + sys.argv[0] + " queryProteomes.csv queryGeneList.faa AUUJ0000000.faa"
		exit(1) # Aborts program. (exit(1) indicates that an error occured)
#-------------------------------------------------------------------------------------------------
# 2: Runs BLAST, can either be sent a fasta formatted string or a file ...
def runBLAST(query, BLASTDBFile):

	query = query.strip()

	# If the string starts with > it is a fasta file and should be sent to blast via stdin.
	if query.startswith(">"):
		BlastArgs = "blastp -db "+ BLASTDBFile + " -num_threads 16 -outfmt \"10 qseqid sseqid pident evalue qcovhsp score\""
		blastp = subprocess.Popen(BlastArgs, shell=True, stdin=subprocess.PIPE,stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
		BLASTOut, err = blastp.communicate(query)

	# Else the string is a file name and should be submitted as a file argument.
	else:
		BLASTOut = subprocess.check_output(["blastp", "-db", BLASTDBFile, "-query", query, "-evalue", "1e-40", "-num_threads", "16", "-outfmt", "10 qseqid sseqid pident evalue qcovhsp score"]) # Runs BLASTp and save output to a string. Blastp is set to output csv which can be parsed.
	return BLASTOut
#-------------------------------------------------------------------------------------------------
# 3: Filters HSPs by Percent Identity...
def filtreBLASTCSV(BLASTOut, minIdent):

	BLASTCSVOut = BLASTOut.splitlines(True) # Converts raw BLAST csv output into list of csv rows.
	BLASTreader = csv.reader(BLASTCSVOut) # Reads BLAST csv rows as a csv.

	BLASTCSVOutFiltred = [] # Note should simply delete unwanted HSPs from curent list rather than making new list.
					        # Rather than making a new one.
	for HSP in BLASTreader:
		if HSP[2] >= minIdent: # Filtres by minimum ident.
			# Converts each HSP parameter that should be a number to a number.
			HSP[2] = float(HSP[2])
			HSP[3] = float(HSP[3])
			HSP[4] = int(HSP[4])
			HSP[5] = int(HSP[5])
			BLASTCSVOutFiltred.append(HSP) # Appends to output array.

	return BLASTCSVOutFiltred
#-------------------------------------------------------------------------------------------------
# 4: Finds Top Scoring Hit For Each Query Protien In BLAST result... Could be more elegantly done...
def getTopHits(BLASTCSVOut):

	topHits = []
	topHits.append(BLASTCSVOut[0])
	currentQuery = BLASTCSVOut[0]

	# If ties occur include these tied in the top hit list. (Ties should have the same score)
	for x in range(1, len(BLASTCSVOut)):
		if BLASTCSVOut[x][5] == currentQuery[5]:			topHits.append(BLASTCSVOut[x])
		else:
			break # Break out of loop if hit has lower score than top hit. # Should remove break as it is bad voodoo...

	return topHits
#-------------------------------------------------------------------------------------------------
# 4: Returns Accession of proteome for which the query protein is found
def getQueryProteome(proteomesCSV, queryProtein):

	for row in proteomesCSV:
		if row[1].strip() == queryProtein.strip():
			proteomeAccession = row[0]
			break # Should remove break as it is bad voodoo...
	return proteomeAccession
#===========================================================================================================
# Main program code:
print time.time()
# House keeping...
argsCheck() # Checks if the number of arguments are correct.

queryFile = sys.argv[1]
queryProteomesFile = sys.argv[2]

# File extension check
if not queryFile.endswith(".faa"):
	print "[Warning] " + queryFile + " may not be a amino acid fasta file!"
# File extension check
if not queryProteomesFile.endswith(".csv"):
	print "[Warning] " + queryProteomesFile + " may not be a csv file!"

BLASTDBFile = sys.argv[3]

print "Forward Blasting to subject proteome..."
BLASTOut = runBLAST(queryFile, BLASTDBFile) # Forward BLASTs from query protien to subject proteome
BLASTOut = filtreBLASTCSV(BLASTOut, 30) # Filtres BLAST results by PIdnet.

# Attemps to open csv contain a csv that links each query protien to its respective proteome.
try:
	QueryProteomes = [] # Stores this CSV in memory for later use
	inFile = open(queryProteomesFile, "r")
	reader = csv.reader(inFile) # opens file with csv module which takes into account verying csv formats and parses correctly
	for row in reader:
		QueryProteomes.append(row)
except IOError:
	print "Failed to open " + queryProteomesFile
	exit(1)

BackBlastOutput = []

print "Back-Blasting hits to query proteome..."
# For each top Hit...
for hit in BLASTOut:
	#print time.time()

	subjectProtein = hit[1]
	queryProtein = hit[0]

	# Takes the csv list of query proteomes and a queryProtien and find what query proteome the query protien is part of.
	CurrentQueryProteome = getQueryProteome(QueryProteomes, queryProtein) + ".faa"

	# Opens query proteome and extracts fasta formated sequence of the query protien.
	# This will be used as the query protien fasta will be used as a BLAST query.
	#print time.time()
	subjectProtienFASTA = subprocess.check_output(["grep", "-A", "1", subjectProtein, BLASTDBFile])
	#print time.time()
	BackBlastOut = runBLAST(subjectProtienFASTA, CurrentQueryProteome) #Backwards BLASTs from subject protien hit to query proteome.
	#print time.time()
	BackBlastOut = filtreBLASTCSV(BackBlastOut, 30) # Filtres BLAST results by PIdnet.
	BackHits = getTopHits(BackBlastOut) # Gets top hits from the BackBlast.

	match = False # By default set match to false
	for BackHit in BackHits:
		if BackHit[1] == queryProtein:
			match = True

	if match == True:
		BackBlastOutput.append(hit)
	#print time.time()
OutFile = BLASTDBFile.rstrip(".faa") + ".csv"

# Attempts to write reciprocal BLAST output to file.
try:
	writeFile = open(OutFile, "w")
	writer = csv.writer(writeFile)
	print ">> Output file created."
	print ">> Writing Data..."
except IOError:
	print "Failed to create " + outFile
	exit(1)

for row in BackBlastOutput:
	writer.writerow(row)

writeFile.close()
print time.time()
print "done"	<commit_msg>Updated script to use a python dictionary to store proteome of subject.<commit_after>#!/usr/bin/env python
# Created by: Lee Bergstrand
# Descript: A Bio-Python program that takes a list of query proteins and uses local BLASTp to search
#           for highly similer proteins within a local blast database (usally a local db of a target
#           proteome). The program then BLASTps backward from the found subject protein to the proteome
#           for which the original query protein if found in order to confirm gene orthology.
#
# Requirements: - This program requires the Biopython module: http://biopython.org/wiki/Download
#               - This script requires BLAST+ 2.2.9 or later.
#               - All operations are done with protien sequences.
#               - All query proteins should be from sequenced genomes in order to facilitate backwards BLAST.
#               - MakeBlastDB must be used to create BLASTp databases for both query and subject proteomes.
#               - BLAST databases require the FASTA file they were made from to be in the same directory.
#
# Usage: BackBLAST.py <queryGeneList.faa> <queryProteomes.csv> <subject1.faa>
# Example: BackBLAST.py queryGeneList.faa queryProteomes.csv AUUJ00000000.faa
#----------------------------------------------------------------------------------------
#===========================================================================================================
#Imports:
import sys
import csv
import subprocess
from Bio import SeqIO
import time
#===========================================================================================================
# Functions:

# 1: Checks if in proper number of arguments are passed gives instructions on proper use.
def argsCheck():
	if len(sys.argv) < 4:
		print "Orthologous Gene Finder"
		print "By Lee Bergstrand\n"
		print "Please refer to source code for documentation\n"
		print "Usage: " + sys.argv[0] + " <queryProteomes.csv> <queryGeneList.faa> <subject1.faa>\n"
		print "Examples:" + sys.argv[0] + " queryProteomes.csv queryGeneList.faa AUUJ0000000.faa"
		exit(1) # Aborts program. (exit(1) indicates that an error occured)
#-------------------------------------------------------------------------------------------------
# 2: Runs BLAST, can either be sent a fasta formatted string or a file ...
def runBLAST(query, BLASTDBFile):

	query = query.strip()

	# If the string starts with > it is a fasta file and should be sent to blast via stdin.
	if query.startswith(">"):
		BlastArgs = "blastp -db "+ BLASTDBFile + " -num_threads 16 -outfmt \"10 qseqid sseqid pident evalue qcovhsp score\""
		blastp = subprocess.Popen(BlastArgs, shell=True, stdin=subprocess.PIPE,stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
		BLASTOut, err = blastp.communicate(query)

	# Else the string is a file name and should be submitted as a file argument.
	else:
		BLASTOut = subprocess.check_output(["blastp", "-db", BLASTDBFile, "-query", query, "-evalue", "1e-40", "-num_threads", "16", "-outfmt", "10 qseqid sseqid pident evalue qcovhsp score"]) # Runs BLASTp and save output to a string. Blastp is set to output csv which can be parsed.
	return BLASTOut
#-------------------------------------------------------------------------------------------------
# 3: Filters HSPs by Percent Identity...
def filtreBLASTCSV(BLASTOut, minIdent):

	BLASTCSVOut = BLASTOut.splitlines(True) # Converts raw BLAST csv output into list of csv rows.
	BLASTreader = csv.reader(BLASTCSVOut) # Reads BLAST csv rows as a csv.

	BLASTCSVOutFiltred = [] # Note should simply delete unwanted HSPs from curent list rather than making new list.
					        # Rather than making a new one.
	for HSP in BLASTreader:
		if HSP[2] >= minIdent: # Filtres by minimum ident.
			# Converts each HSP parameter that should be a number to a number.
			HSP[2] = float(HSP[2])
			HSP[3] = float(HSP[3])
			HSP[4] = int(HSP[4])
			HSP[5] = int(HSP[5])
			BLASTCSVOutFiltred.append(HSP) # Appends to output array.

	return BLASTCSVOutFiltred
#-------------------------------------------------------------------------------------------------
# 4: Finds Top Scoring Hit For Each Query Protien In BLAST result... Could be more elegantly done...
def getTopHits(BLASTCSVOut):

	topHits = []
	topHits.append(BLASTCSVOut[0])
	currentQuery = BLASTCSVOut[0]

	# If ties occur include these tied in the top hit list. (Ties should have the same score)
	for x in range(1, len(BLASTCSVOut)):
		if BLASTCSVOut[x][5] == currentQuery[5]:			topHits.append(BLASTCSVOut[x])
		else:
			break # Break out of loop if hit has lower score than top hit. # Should remove break as it is bad voodoo...

	return topHits
#-------------------------------------------------------------------------------------------------
# 4: Returns Accession of proteome for which the query protein is found
def getQueryProteome(proteomesCSV, queryProtein):

	for row in proteomesCSV:
		if row[1].strip() == queryProtein.strip():
			proteomeAccession = row[0]
			break # Should remove break as it is bad voodoo...
	return proteomeAccession
#----------------------------------------------------------------------
# 5: Creates a python dictionary (hash table) that contains the the fasta for each protien in the proteome.
def createProteomeHash(ProteomeFile):
	ProteomeHash = dict()
	try:
		handle = open(ProteomeFile, "rU")
		proteome = SeqIO.parse(handle, "fasta")
		for record in proteome:
			ProteomeHash.update({ record.id : record.format("fasta") })
		handle.close()
	except IOError:
		print "Failed to open " + ProteomeFile
		exit(1)

	return ProteomeHash
#===========================================================================================================
# Main program code:
print time.time()
# House keeping...
argsCheck() # Checks if the number of arguments are correct.

queryFile = sys.argv[1]
queryProteomesFile = sys.argv[2]

# File extension check
if not queryFile.endswith(".faa"):
	print "[Warning] " + queryFile + " may not be a amino acid fasta file!"
# File extension check
if not queryProteomesFile.endswith(".csv"):
	print "[Warning] " + queryProteomesFile + " may not be a csv file!"

BLASTDBFile = sys.argv[3]

print "Forward Blasting to subject proteome..."
BLASTOut = runBLAST(queryFile, BLASTDBFile) # Forward BLASTs from query protien to subject proteome
BLASTOut = filtreBLASTCSV(BLASTOut, 30) # Filtres BLAST results by PIdnet.

# Attemps to open csv contain a csv that links each query protien to its respective proteome.
try:
	QueryProteomes = [] # Stores this CSV in memory for later use
	inFile = open(queryProteomesFile, "r")
	reader = csv.reader(inFile) # opens file with csv module which takes into account verying csv formats and parses correctly
	for row in reader:
		QueryProteomes.append(row)
except IOError:
	print "Failed to open " + queryProteomesFile
	exit(1)

BackBlastOutput = []
SubjectProteomeHash = createProteomeHash(BLASTDBFile) # Creates python dictionary for geneome.

print "Back-Blasting hits to query proteome..."
# For each top Hit...
for hit in BLASTOut:
	subjectProtein = hit[1]
	queryProtein = hit[0]

	# Takes the csv list of query proteomes and a queryProtien and find what query proteome the query protien is part of.
	CurrentQueryProteome = getQueryProteome(QueryProteomes, queryProtein) + ".faa"
	subjectProtienFASTA = SubjectProteomeHash.get(subjectProtein) # Extracts subjectProtien from python dictionary.
	BackBlastOut = runBLAST(subjectProtienFASTA, CurrentQueryProteome) #Backwards BLASTs from subject protien hit to query proteome.

	BackBlastOut = filtreBLASTCSV(BackBlastOut, 30) # Filtres BLAST results by PIdnet.
	BackHits = getTopHits(BackBlastOut) # Gets top hits from the BackBlast.

	match = False # By default set match to false
	for BackHit in BackHits:
		if BackHit[1] == queryProtein:
			match = True

	if match == True:
		BackBlastOutput.append(hit)

OutFile = BLASTDBFile.rstrip(".faa") + ".csv"

# Attempts to write reciprocal BLAST output to file.
try:
	writeFile = open(OutFile, "w")
	writer = csv.writer(writeFile)
	print ">> Output file created."
	print ">> Writing Data..."
except IOError:
	print "Failed to create " + outFile
	exit(1)

for row in BackBlastOutput:
	writer.writerow(row)

writeFile.close()
print "done"
print time.time()<|endoftext|>


<commit_before>#!/usr/bin/env python
# Author: Nick Zwart
# Date: 2016jun01
#   Backup all the projects of a git-hub style website via git mirroring.
#   https://www.garron.me/en/bits/backup-git-bare-repo.html

import os
import sys
import gitlab
import shutil
import hashlib
import optparse
import subprocess

class GitLabAPI:
    def __init__(self, token, url):
        self._token = token
        self._url = url

        # authenticate a gitlab session
        self._gl = gitlab.Gitlab(self._url, self._token)
        self._gl.auth()

        # list all projects
        self._projects = self._gl.projects.list(all=True)

    def numProjects(self):
        return len(self._projects)

    def projectPath(self, index):
        return self._projects[index].path_with_namespace

    def projectURL(self, index):
        return self._projects[index].ssh_url_to_repo

class GitBareMirror:
    '''A simple git interface for managing bare-mirroed repos that backup url
    accessible upstream repos.
    '''

    def __init__(self, path, url, overwrite=False, moveAside=False):
        self._path = path
        self._origin_url = url
        self._overwrite = overwrite
        self._moveAside = moveAside

        if self.validExistingRepo():
            self.update()
        else:
            self.createMirroredRepo()

    def validExistingRepo(self):

        try:
            assert os.path.isdir(self._path)

            # move to the existing repo and check if its bare
            os.chdir(self._path)
            cmd = subprocess.Popen('git rev-parse --is-bare-repository',
                    shell=True, stdout=subprocess.PIPE)
            cmd.wait()

            # Error checking
            assert cmd.returncode != 128, ('ERROR: The supplied directory '
                    'exists, but is not a git repo.')
            assert cmd.returncode == 0, 'ERROR: git error'
            firstline = cmd.stdout.readlines()[0].decode('utf8')
            assert 'false' not in firstline, ('ERROR: The supplied directory '
                    'is NOT a bare repo.')
            assert 'true' in firstline

            # check if the existing repo has the same origin url
            # -prevent name collision if group/org namespace isn't used
            cmd = subprocess.Popen('git config --get remote.origin.url',
                    shell=True, stdout=subprocess.PIPE)
            cmd.wait()
            firstline = cmd.stdout.readlines()[0].decode('utf8')

            assert self._origin_url in firstline, ('ERROR: The existing repo '
                    'has a url that differs from the supplied origin url.')

            return True

        except AssertionError as err:
            print('Existing repo failed via:', err)
            return False

    def update(self):
        cmd = subprocess.Popen('git remote update', shell=True,
                stdout=subprocess.PIPE)
        cmd.wait()

        assert cmd.returncode == 0, 'ERROR: git error'
        print('SUCCESS (updated)')

    def createMirroredRepo(self):

        # Handle existing directories based on user options:
        # move the dir to a unique name, remove it, or fail w/ exception
        if self._moveAside and os.path.exists(self._path):
            parentPath = os.path.dirname(self._path)
            dirContents = str(os.listdir(parentPath)).encode('utf8')
            newNameExt = hashlib.md5(dirContents).hexdigest()
            newName = self._path+'_'+newNameExt
            print('MOVING PATH', self._path, newName)
            shutil.move(self._path, newName)
        elif self._overwrite and os.path.exists(self._path):
            print('REMOVING PATH', self._path)
            shutil.rmtree(self._path)
        else:
            assert not os.path.exists(self._path), ('ERROR: the supplied path '
                    'already exists, unable to create mirror.')

        os.makedirs(self._path)
        os.chdir(self._path)
        cmd = subprocess.Popen('git clone --mirror ' + str(self._origin_url)
                + ' .', shell=True, stdout=subprocess.PIPE)
        cmd.wait()
        print('SUCCESS (new mirror)')

if __name__ == '__main__':

    # parse input args
    parser = optparse.OptionParser()
    parser.add_option('--path', dest='backupPath', action='store',
            type='string', default=os.path.expanduser('~/backup'),
            help='The directory to store the backups.')
    parser.add_option('--ignore-errors', dest='ignoreErrors',
            action='store_true', default=False,
            help='Continue to backup other repos if one has failed.')
    parser.add_option('--overwrite', dest='overwrite',
            action='store_true', default=False,
            help='Overwrite existing directories.')
    parser.add_option('--move-aside', dest='moveAside',
            action='store_true', default=False,
            help='Move existing directories aside with a tempfile extension.')
    parser.add_option('--token', dest='token', action='store',
            type='string', default=None,
            help='The token required to access the target git web api.')
    parser.add_option('--website', dest='website', action='store',
            type='string', default=None,
            help='The hub website where the git repos are stored.')

    options, args = parser.parse_args(sys.argv)

    assert options.token is not None
    assert options.website is not None

    # Check for existing backup directory and make one if it doesn't exist.
    if not os.path.isdir(options.backupPath):
        print('The specified backup path doesn\'t exist.')
        sys.exit(1)

    # Get the repository info from the git web api.
    webapi = GitLabAPI(options.token, options.website)

    # Display whats going on as the repos get either updated or newly mirrored.
    print('Repository:')
    for i in range(webapi.numProjects()):
        try:
            curPath = os.path.join(options.backupPath, webapi.projectPath(i))
            curURL = webapi.projectURL(i)
            print('\nSyncing: ', curURL, curPath)
            repo = GitBareMirror(curPath, curURL, overwrite=options.overwrite,
                    moveAside=options.moveAside)
        except Exception as err:
            if options.ignoreErrors:
                print(err)
            else:
                raise
<commit_msg>added timestamps to backup script<commit_after>#!/usr/bin/env python
# Author: Nick Zwart
# Date: 2016jun01
#   Backup all the projects of a git-hub style website via git mirroring.
#   https://www.garron.me/en/bits/backup-git-bare-repo.html

import os
import sys
import time
import gitlab
import shutil
import hashlib
import optparse
import subprocess

class GitLabAPI:
    def __init__(self, token, url):
        self._token = token
        self._url = url

        # authenticate a gitlab session
        self._gl = gitlab.Gitlab(self._url, self._token)
        self._gl.auth()

        # list all projects
        self._projects = self._gl.projects.list(all=True)

    def numProjects(self):
        return len(self._projects)

    def projectPath(self, index):
        return self._projects[index].path_with_namespace

    def projectURL(self, index):
        return self._projects[index].ssh_url_to_repo

class GitBareMirror:
    '''A simple git interface for managing bare-mirroed repos that backup url
    accessible upstream repos.
    '''

    def __init__(self, path, url, overwrite=False, moveAside=False):
        self._path = path
        self._origin_url = url
        self._overwrite = overwrite
        self._moveAside = moveAside

        if self.validExistingRepo():
            self.update()
        else:
            self.createMirroredRepo()

    def validExistingRepo(self):

        try:
            assert os.path.isdir(self._path)

            # move to the existing repo and check if its bare
            os.chdir(self._path)
            cmd = subprocess.Popen('git rev-parse --is-bare-repository',
                    shell=True, stdout=subprocess.PIPE)
            cmd.wait()

            # Error checking
            assert cmd.returncode != 128, ('ERROR: The supplied directory '
                    'exists, but is not a git repo.')
            assert cmd.returncode == 0, 'ERROR: git error'
            firstline = cmd.stdout.readlines()[0].decode('utf8')
            assert 'false' not in firstline, ('ERROR: The supplied directory '
                    'is NOT a bare repo.')
            assert 'true' in firstline

            # check if the existing repo has the same origin url
            # -prevent name collision if group/org namespace isn't used
            cmd = subprocess.Popen('git config --get remote.origin.url',
                    shell=True, stdout=subprocess.PIPE)
            cmd.wait()
            firstline = cmd.stdout.readlines()[0].decode('utf8')

            assert self._origin_url in firstline, ('ERROR: The existing repo '
                    'has a url that differs from the supplied origin url.')

            return True

        except AssertionError as err:
            print('Existing repo failed via:', err)
            return False

    def update(self):
        cmd = subprocess.Popen('git remote update', shell=True,
                stdout=subprocess.PIPE)
        cmd.wait()

        assert cmd.returncode == 0, 'ERROR: git error'
        print('SUCCESS (updated)')

    def createMirroredRepo(self):

        # Handle existing directories based on user options:
        # move the dir to a unique name, remove it, or fail w/ exception
        if self._moveAside and os.path.exists(self._path):
            parentPath = os.path.dirname(self._path)
            dirContents = str(os.listdir(parentPath)).encode('utf8')
            newNameExt = hashlib.md5(dirContents).hexdigest()
            newName = self._path+'_'+newNameExt
            print('MOVING PATH', self._path, newName)
            shutil.move(self._path, newName)
        elif self._overwrite and os.path.exists(self._path):
            print('REMOVING PATH', self._path)
            shutil.rmtree(self._path)
        else:
            assert not os.path.exists(self._path), ('ERROR: the supplied path '
                    'already exists, unable to create mirror.')

        os.makedirs(self._path)
        os.chdir(self._path)
        cmd = subprocess.Popen('git clone --mirror ' + str(self._origin_url)
                + ' .', shell=True, stdout=subprocess.PIPE)
        cmd.wait()
        print('SUCCESS (new mirror)')

if __name__ == '__main__':

    # parse input args
    parser = optparse.OptionParser()
    parser.add_option('--path', dest='backupPath', action='store',
            type='string', default=os.path.expanduser('~/backup'),
            help='The directory to store the backups.')
    parser.add_option('--ignore-errors', dest='ignoreErrors',
            action='store_true', default=False,
            help='Continue to backup other repos if one has failed.')
    parser.add_option('--overwrite', dest='overwrite',
            action='store_true', default=False,
            help='Overwrite existing directories.')
    parser.add_option('--move-aside', dest='moveAside',
            action='store_true', default=False,
            help='Move existing directories aside with a tempfile extension.')
    parser.add_option('--token', dest='token', action='store',
            type='string', default=None,
            help='The token required to access the target git web api.')
    parser.add_option('--website', dest='website', action='store',
            type='string', default=None,
            help='The hub website where the git repos are stored.')

    options, args = parser.parse_args(sys.argv)

    localtime = time.asctime( time.localtime(time.time()) )
    print("BackupHub Start:", localtime)

    assert options.token is not None
    assert options.website is not None

    # Check for existing backup directory and make one if it doesn't exist.
    if not os.path.isdir(options.backupPath):
        print('The specified backup path doesn\'t exist.')
        sys.exit(1)

    # Get the repository info from the git web api.
    webapi = GitLabAPI(options.token, options.website)

    # Display whats going on as the repos get either updated or newly mirrored.
    print('Repository:')
    for i in range(webapi.numProjects()):
        try:
            curPath = os.path.join(options.backupPath, webapi.projectPath(i))
            curURL = webapi.projectURL(i)
            print('\nSyncing: ', curURL, curPath)
            repo = GitBareMirror(curPath, curURL, overwrite=options.overwrite,
                    moveAside=options.moveAside)
        except Exception as err:
            if options.ignoreErrors:
                print(err)
            else:
                raise

    localtime = time.asctime( time.localtime(time.time()) )
    print("BackupHub Finished:", localtime)
<|endoftext|>


<commit_before>import random
import string
import re
import fileinput

from datetime import datetime
from dateutil.relativedelta import relativedelta

from itertools import dropwhile, chain, zip_longest

import multiprocessing as mp

def nchunks(xs, n):
    """Yields n chunks about the same size. """
    xs = list(xs)
    start = 0
    for i in range(n):
        if i + 1 == n:
            # last chunk
            yield xs[start:]
        else:
            chunksize = round(len(xs[start:]) / (n - i))
            yield xs[start:start + chunksize]
            start += chunksize

def prepend_header(filename, header=None, drop=1):
    """Drop n lines and prepend header

    Args:
        filename (str)
        header (str)
        drop (int)
    """
    for no, line in enumerate(fileinput.input(filename, inplace=True)):
        # it's meaningless to set drop to -1, -2, ...
        if no == 0 and drop == 0:
            if header:
                print(header)
            print(line, end='')
        # replace
        elif no + 1 == drop:
            if header:
                print(header)
        elif no >= drop:
            print(line, end='')
        else:
            # no + 1 < drop
            continue

def random_string(nchars=20):
    "Generates a random string of lengh 'n' with alphabets and digits. "
    return ''.join(random.SystemRandom().choice(string.ascii_letters +
                   string.digits) for _ in range(nchars))

def mpairs(seq1, seq2, key1, key2=None):
    """Generates a tuple of matching pairs
    key1 and key2 are functions (getters, sort of)

    seq1, seq2 must be sorted before being passed here
        and also each key value(which is returned by key funcs) must be UNIQUE
        otherwise you will see unexpected results
    """
    key2 = key2 or key1

    s1, s2 = next(seq1), next(seq2)
    k1, k2 = key1(s1), key2(s2)

    while True:
        try:
            if k1 == k2:
                yield (s1, s2)
                s1, s2 = next(seq1), next(seq2)
                k1, k2 = key1(s1), key2(s2)
            elif k1 < k2:
                s1 = next(dropwhile(lambda x: key1(x) < k2, seq1))
                k1 = key1(s1)
            else:
                s2 = next(dropwhile(lambda x: key2(x) < k1, seq2))
                k2 = key2(s2)

        except StopIteration:
            break

def camel2snake(name):
    """
    Args:
        name (str): camelCase
    Returns:
        str: snake_case
    """
    s1 = re.sub('(.)([A-Z][a-z]+)', r'\1_\2', name)
    return re.sub('([a-z0-9])([A-Z])', r'\1_\2', s1).lower()

def peek_first(seq):
    """
    Note:
        peeked first item is pushed back to the sequence
    Args:
        seq (Iter[type])
    Returns:
        Tuple(type, Iter[type])
    """
    seq = iter(seq)
    first_item = next(seq)
    return first_item, chain([first_item], seq)

def listify(colstr):
    """A comma or space separated string to a list of strings

    Args:
        colstr (str)
    Returns:
        List[str]

    Example:
        >>> _listify('a b c')
        ['a', 'b', 'c']

        >>> _listify('a, b, c')
        ['a', 'b', 'c']
    """
    if isinstance(colstr, str):
        if ',' in colstr:
            return [x.strip() for x in colstr.split(',')]
        else:
            return [x for x in colstr.split(' ') if x]
    else:
        return colstr

# copied from 'itertools'
def grouper(iterable, n, fillvalue=None):
    "Collect data into fixed-length chunks or blocks"
    # grouper('ABCDEFG', 3, 'x') --> ABC DEF Gxx"
    args = [iter(iterable)] * n
    return zip_longest(*args, fillvalue=fillvalue)

# Pool.iamp in multiprocessing doesn't allow
# it to pass locally defined functions
# I won't define unordered version although it is somewhat faster
# and easy to write and also doesn't cause a lot of troubles.
# Nevertheless I want to minimize any nondeterminism at all costs.
def pmap(func, seq, chunksize=1, processes=mp.cpu_count()):
    """
    parallel map, ordered version
    if you are curious about the parameters, refer to multiprocessing library
    documentation (Pool.imap)
    """
    the_end = random_string()

    # Opening multiple ques sounds dumb in a way
    # but this is a easier way to implement the ordered version of
    # parrallel map. It's just that there is a limit in the number of
    # ques in the OS. Of course you wouldn't make more than 20 processes.
    que1s = [mp.Queue(1) for _ in range(processes)]
    que2s = [mp.Queue(1) for _ in range(processes)]

    ws = []

    def insert1(seq, que1s):
        for chunks in grouper(grouper(seq, chunksize, the_end),
                              processes, the_end):
            for que1, chunk in zip(que1s, chunks):
                que1.put(chunk)
        for que1 in que1s:
            que1.put(the_end)

    w0 = mp.Process(target=insert1, args=(seq, que1s))
    w0.daemon = True
    w0.start()
    ws.append(w0)

    def insert2(que1, que2):
        while True:
            chunk = que1.get()
            if chunk == the_end:
                que2.put(the_end)
                return
            else:
                result = []
                for x in chunk:
                    if x != the_end:
                        try:
                            result.append(func(x))
                        except Exception as error:
                            que2.put(the_end)
                            que2.put('child process error: ' + repr(error))
                            return
                que2.put(result)

    for que1, que2 in zip(que1s, que2s):
        w = mp.Process(target=insert2, args=(que1, que2))
        w.daemon = True
        w.start()
        ws.append(w)

    while True:
        for que2 in que2s:
            result = que2.get()
            if result == the_end:
                if not que2.empty():
                    # you have an error message.
                    print(que2.get())
                for w in ws:
                    w.terminate()
                return
            else:
                yield from result

# The following guys are also going to be
# included in "extended sqlite functions" set

# If the return value is True it is converted to 1 or 0 in sqlite3
def isnum(x):
    return isinstance(x, float) or isinstance(x, int)

def istext(x):
    return isinstance(x, str)

def yyyymm(date, n):
    d1 = datetime.strptime(str(date), '%Y%m') + relativedelta(months=n)
    return int(d1.strftime('%Y%m'))

def yyyymmdd(date, n):
    d1 = datetime.strptime(str(date), '%Y%m%d') + relativedelta(days=n)
    return int(d1.strftime('%Y%m%d'))

<commit_msg>umerge<commit_after>import random
import string
import re
import fileinput

from datetime import datetime
from dateutil.relativedelta import relativedelta

from itertools import dropwhile, chain, zip_longest

import multiprocessing as mp
import threading as th
from queue import Queue

def nchunks(xs, n):
    """Yields n chunks about the same size. """
    xs = list(xs)
    start = 0
    for i in range(n):
        if i + 1 == n:
            # last chunk
            yield xs[start:]
        else:
            chunksize = round(len(xs[start:]) / (n - i))
            yield xs[start:start + chunksize]
            start += chunksize

def prepend_header(filename, header=None, drop=1):
    """Drop n lines and prepend header

    Args:
        filename (str)
        header (str)
        drop (int)
    """
    for no, line in enumerate(fileinput.input(filename, inplace=True)):
        # it's meaningless to set drop to -1, -2, ...
        if no == 0 and drop == 0:
            if header:
                print(header)
            print(line, end='')
        # replace
        elif no + 1 == drop:
            if header:
                print(header)
        elif no >= drop:
            print(line, end='')
        else:
            # no + 1 < drop
            continue

def random_string(nchars=20):
    "Generates a random string of lengh 'n' with alphabets and digits. "
    return ''.join(random.SystemRandom().choice(string.ascii_letters +
                   string.digits) for _ in range(nchars))

def mpairs(seq1, seq2, key1, key2=None):
    """Generates a tuple of matching pairs
    key1 and key2 are functions (getters, sort of)

    seq1, seq2 must be sorted before being passed here
        and also each key value(which is returned by key funcs) must be UNIQUE
        otherwise you will see unexpected results
    """
    key2 = key2 or key1

    s1, s2 = next(seq1), next(seq2)
    k1, k2 = key1(s1), key2(s2)

    while True:
        try:
            if k1 == k2:
                yield (s1, s2)
                s1, s2 = next(seq1), next(seq2)
                k1, k2 = key1(s1), key2(s2)
            elif k1 < k2:
                s1 = next(dropwhile(lambda x: key1(x) < k2, seq1))
                k1 = key1(s1)
            else:
                s2 = next(dropwhile(lambda x: key2(x) < k1, seq2))
                k2 = key2(s2)

        except StopIteration:
            break

def camel2snake(name):
    """
    Args:
        name (str): camelCase
    Returns:
        str: snake_case
    """
    s1 = re.sub('(.)([A-Z][a-z]+)', r'\1_\2', name)
    return re.sub('([a-z0-9])([A-Z])', r'\1_\2', s1).lower()

def peek_first(seq):
    """
    Note:
        peeked first item is pushed back to the sequence
    Args:
        seq (Iter[type])
    Returns:
        Tuple(type, Iter[type])
    """
    seq = iter(seq)
    first_item = next(seq)
    return first_item, chain([first_item], seq)

def listify(colstr):
    """A comma or space separated string to a list of strings

    Args:
        colstr (str)
    Returns:
        List[str]

    Example:
        >>> _listify('a b c')
        ['a', 'b', 'c']

        >>> _listify('a, b, c')
        ['a', 'b', 'c']
    """
    if isinstance(colstr, str):
        if ',' in colstr:
            return [x.strip() for x in colstr.split(',')]
        else:
            return [x for x in colstr.split(' ') if x]
    else:
        return colstr

# copied from 'itertools'
def grouper(iterable, n, fillvalue=None):
    "Collect data into fixed-length chunks or blocks"
    # grouper('ABCDEFG', 3, 'x') --> ABC DEF Gxx"
    args = [iter(iterable)] * n
    return zip_longest(*args, fillvalue=fillvalue)

# Pool.iamp in multiprocessing doesn't allow
# it to pass locally defined functions
# I won't define unordered version although it is somewhat faster
# and easy to write and also doesn't cause a lot of troubles.
# Nevertheless I want to minimize any nondeterminism at all costs.
def pmap(func, seq, chunksize=1, processes=mp.cpu_count()):
    """
    parallel map, ordered version
    if you are curious about the parameters, refer to multiprocessing library
    documentation (Pool.imap)
    """
    the_end = random_string()

    # Opening multiple ques sounds dumb in a way
    # but this is a easier way to implement the ordered version of
    # parrallel map. It's just that there is a limit in the number of
    # ques in the OS. Of course you wouldn't make more than 20 processes.
    que1s = [mp.Queue(1) for _ in range(processes)]
    que2s = [mp.Queue(1) for _ in range(processes)]

    ws = []

    def insert1(seq, que1s):
        for chunks in grouper(grouper(seq, chunksize, the_end),
                              processes, the_end):
            for que1, chunk in zip(que1s, chunks):
                que1.put(chunk)
        for que1 in que1s:
            que1.put(the_end)

    w0 = mp.Process(target=insert1, args=(seq, que1s))
    w0.daemon = True
    w0.start()
    ws.append(w0)

    def insert2(que1, que2):
        while True:
            chunk = que1.get()
            if chunk == the_end:
                que2.put(the_end)
                return
            else:
                result = []
                for x in chunk:
                    if x != the_end:
                        try:
                            result.append(func(x))
                        except Exception as error:
                            que2.put(the_end)
                            que2.put('child process error: ' + repr(error))
                            return
                que2.put(result)

    for que1, que2 in zip(que1s, que2s):
        w = mp.Process(target=insert2, args=(que1, que2))
        w.daemon = True
        w.start()
        ws.append(w)

    while True:
        for que2 in que2s:
            result = que2.get()
            if result == the_end:
                if not que2.empty():
                    # you have an error message.
                    print(que2.get())
                for w in ws:
                    w.terminate()
                return
            else:
                yield from result

# unordered merge
def umerge(*sequences, parallel=False):
    the_end = random_string() if parallel else object()
    que = mp.Queue() if parallel else Queue()
    worker = mp.Process if parallel else th.Thread
    n_sequences = len(sequences)

    def insert(sequence):
        for x in sequence:
            que.put(x)
        que.put(the_end)

    ws = []
    for seq in sequences:
        w = worker(target=insert, args=(seq,))
        w.daemon = True
        ws.append(w)
        w.start()

    the_end_cnt = 0
    while True:
        val = que.get()
        if val == the_end:
            the_end_cnt += 1
            if the_end_cnt == n_sequences:
                break
        else:
            yield val

    for w in ws:
        w.join()

# The following guys are also going to be
# included in "extended sqlite functions" set

# If the return value is True it is converted to 1 or 0 in sqlite3
def isnum(x):
    return isinstance(x, float) or isinstance(x, int)

def istext(x):
    return isinstance(x, str)

def yyyymm(date, n):
    d1 = datetime.strptime(str(date), '%Y%m') + relativedelta(months=n)
    return int(d1.strftime('%Y%m'))

def yyyymmdd(date, n):
    d1 = datetime.strptime(str(date), '%Y%m%d') + relativedelta(days=n)
    return int(d1.strftime('%Y%m%d'))

<|endoftext|>


<commit_before>#! /usr/bin/env python
# -*- coding: utf-8 -*-

from sigma_urls import get_chemistry_base_urls, get_chemistry_urls

if __name__ == '__main__':
    get_chemistry_base_urls()
    get_chemistry_urls()<commit_msg>add get materials<commit_after>#! /usr/bin/env python
# -*- coding: utf-8 -*-

from sigma_urls import get_materials_urls

if __name__ == '__main__':
    get_materials_urls()<|endoftext|>


<commit_before>"""Pyhole Utility Library"""

from threading import Thread

def admin(func):
    """Administration Decorator"""

    def f(self, *args, **kwargs):
        if self.irc.source == self.irc.admin:
            func(self, *args, **kwargs)
        else:
            self.irc.say("Sorry, you are not authorized to do that.")
    f.__doc__ = func.__doc__
    return f

def spawn(func):
    """Thread Spawning Decorator"""

    def f(self, *args, **kwargs):
        if args and args[0]:
            params = " ".join(args)
            t = Thread(target=func, args=(self, params), kwargs=kwargs)
        else:
            t = Thread(target=func, args=(self,), kwargs=kwargs)
        self.irc.log.debug("Spawning: %s" % t.name)
        t.start()
    f.__doc__ = func.__doc__
    return f
<commit_msg>Minor cleanup<commit_after>"""Pyhole Utility Library"""

import threading

def admin(func):
    """Administration Decorator"""

    def f(self, *args, **kwargs):
        if self.irc.source == self.irc.admin:
            func(self, *args, **kwargs)
        else:
            self.irc.say("Sorry, you are not authorized to do that.")
    f.__doc__ = func.__doc__
    return f

def spawn(func):
    """Thread Spawning Decorator"""

    def f(self, *args, **kwargs):
        if args and args[0]:
            params = " ".join(args)
            t = threading.Thread(
                target=func,
                args=(self, params),
                kwargs=kwargs)
        else:
            t = threading.Thread(target=func, args=(self,), kwargs=kwargs)
        self.irc.log.debug("Spawning: %s" % t.name)
        t.start()
    f.__doc__ = func.__doc__
    return f
<|endoftext|>