<commit_before>#!/usr/bin/env python
from __future__ import print_function
import screed
import docopt
import numpy as np
from sys import stderr, stdout
import sys
import json


CLI = '''
USAGE:
    add-barcodes.py [options] KEYFILE FASTQ [FASTQ2]

OPTIONS:
    -s SEED         Output file [default: 1234]
    -o OUTPUT       Output file [default: stdout]
    -G GAMMA_SHAPE  Shape of gamma distribution, source of sample frequencies.
                    [default: 2]
'''


def read_axe_key(path):
    samples = {}
    with open(path) as fh:
        for line in fh:
            line = line.strip()
            if line.lower().startswith('barcode') or \
                    line.startswith('#') or \
                    line.startswith(';'):
                continue
            fields = line.split('\t')
            barcodes = fields[:-1]
            if len(barcodes) != 2:
                barcodes.append('')
            samples[fields[-1]] = barcodes

    return [{'id': x, 'bcd': samples[x]} for x in sorted(samples)]


def samplefreq(samples, gamma_shape=2):
    # Sample frequencies are taken from a gamma distribution, and normalised to
    # sum to 1.
    freqs = np.random.gamma(gamma_shape, size=len(samples))
    freqs /= np.sum(freqs)

    return freqs.cumsum()


def mutrate(n, phredscore=25):
    '''Random number from the exponential distrbuiton <= n.'''
    error_rate = 10**(phredscore / 10.)
    x = n + 1
    while x > n:
        x = int(np.random.exponential(1./np.log(error_rate)))
    return x


def mutate(seq, dist, alphabet='ACGT'):
    dist = min(dist, len(seq))

    # distance of 0 is unchanged
    if dist < 1:
        return seq

    idx = np.arange(len(seq))
    np.random.shuffle(idx)
    seq = list(seq)
    for i in range(dist):
        replacement = seq[idx[i]]
        while seq[idx[i]] == replacement:
            replacement = alphabet[np.random.choice(len(alphabet))]
        seq[idx[i]] = replacement
    return "".join(seq)


def add_barcode_to_read(read_pair, samples, cumsum_prob, max_mismatch=0.5):
    r = np.random.uniform()
    idx = np.searchsorted(cumsum_prob, r)

    sample = samples[idx]

    def read_str(read, barcode):
        fake_qual = read.quality[:len(barcode)]
        if len(barcode) > 0:
            avg_qual = sum(ord(x)-33 for x in fake_qual) / float(len(barcode))
            mismatch = mutrate(int(max_mismatch * len(barcode)),
                               phredscore=avg_qual)
        else:
            mismatch = 0
        bcd_seq = mutate(barcode, mismatch)


        sample_tag = json.dumps({'id': sample['id'],
                                 'barcodes': sample['bcd'],
                                 'mismatches': mismatch})

        return '@{}\t{}\n{}{}\n+\n{}{}'.format(
                read.name, sample_tag,
                bcd_seq,  read.sequence,
                fake_qual, read.quality<commit_msg>add-barcodes: Add RE site insertion, use python3<commit_after>#!/usr/bin/env python3
from __future__ import print_function
import screed
import docopt
import numpy as np
from sys import stderr, stdout
import sys
import json


CLI = '''
USAGE:
    add-barcodes.py [options] KEYFILE FASTQ [FASTQ2]

OPTIONS:
    -s SEED         Output file [default: 1234]
    -o OUTPUT       Output file [default: stdout]
    -G GAMMA_SHAPE  Shape of gamma distribution, source of sample frequencies.
                    [default: 2]
'''


def read_axe_key(path):
    samples = {}
    with open(path) as fh:
        for line in fh:
            line = line.strip()
            if line.lower().startswith('barcode') or \
                    line.startswith('#') or \
                    line.startswith(';'):
                continue
            fields = line.split('\t')
            barcodes = fields[:-1]
            if len(barcodes) != 2:
                barcodes.append('')
            samples[fields[-1]] = barcodes

    return [{'id': x, 'bcd': samples[x]} for x in sorted(samples)]


def samplefreq(samples, gamma_shape=2):
    # Sample frequencies are taken from a gamma distribution, and normalised to
    # sum to 1.
    freqs = np.random.gamma(gamma_shape, size=len(samples))
    freqs /= np.sum(freqs)

    return freqs.cumsum()


def mutrate(n, phredscore=25):
    '''Random number from the exponential distrbuiton <= n.'''
    error_rate = 10**(phredscore / 10.)
    x = n + 1
    while x > n:
        x = int(np.random.exponential(1./np.log(error_rate)))
    return x


def mutate(seq, dist, alphabet='ACGT'):
    dist = min(dist, len(seq))

    # distance of 0 is unchanged
    if dist < 1:
        return seq

    idx = np.arange(len(seq))
    np.random.shuffle(idx)
    seq = list(seq)
    for i in range(dist):
        replacement = seq[idx[i]]
        while seq[idx[i]] == replacement:
            replacement = alphabet[np.random.choice(len(alphabet))]
        seq[idx[i]] = replacement
    return "".join(seq)


def add_barcode_to_read(read_pair, samples, cumsum_prob, max_mismatch=0.5,
                        re_site='TGCAG', gibberish_prob=0.01):
    r = np.random.uniform()
    idx = np.searchsorted(cumsum_prob, r)

    sample = samples[idx]

    def read_str(read, barcode):
        fake_qual = read.quality[:len(barcode) + len(re_site)]
        if barcode:
            avg_qual = sum(ord(x)-33 for x in fake_qual) / float(len(barcode))
            mismatch = mutrate(int(max_mismatch * len(barcode)),
                               phredscore=avg_qual)
        elif np.random.uniform() < gibberish_prob:
            # Mutate barcode fully to create random gibberish as a barcode
            mismatch = len(barcode)
        else:
            mismatch = 0

        if re_site:
            mismatch = mutrate(len(re_site), phredscore=avg_qual)
            re_seq = mutate(re_site, mismatch)
        else:
            re_seq = ''

        bcd_seq = mutate(barcode, mismatch)

        sample_tag = json.dumps({'id': sample['id'],
                                 'barcodes': sample['bcd'],
                                 'mismatches': mismatch})

        return '@{}\t{}{}\n{}{}\n+\n{}{}'.format(
                read.name, sample_tag,
                re_seq, bcd_seq, read.sequence,
                fake_qual, read.quality
{'fraction_non_alphanumeric': 0.12316857188858477, 'fraction_numerical': 0.009338270809853486, 'mean_word_length': 2.0631163708086784, 'pattern_counts': {'":': 0, '<': 8, '<?xml version=': 0, '>': 6, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '455527', 'n_tokens_mistral': 2065, 'n_tokens_neox': 1948, 'n_words': 493}