text:
#!/usr/bin/env python3
from __future__ import print_function
import pandas as pd
import tensorflow as tf
import numpy as np
import gzip
import time
from collections import OrderedDict
from datetime import datetime
import sys
import os
import subprocess
import glob
import argparse
import scipy.optimize
import scipy.stats as stats
from scipy.special import loggamma

sys.path.insert(1, os.path.dirname(__file__))
import genotypeio

has_rpy2 = False
e = subprocess.call('which R', shell=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
try:
    import rpy2
    import rfunc
    if e==0:
        has_rpy2 = True
except:
    pass
if not has_rpy2:
    print("Warning: 'rfunc' cannot be imported. R and the 'rpy2' Python package are needed.")

class SimpleLogger(object):
    def __init__(self, logfile=None, verbose=True):
        self.console = sys.stdout
        self.verbose = verbose
        if logfile is not None:
            self.log = open(logfile, 'w')
        else:
            self.log = None

    def write(self, message):
        if self.verbose:
            self.console.write(message+'\n')
        if self.log is not None:
            self.log.write(message+'\n')
            self.log.flush()

output_dtype_dict = {
    'num_var':np.int32,
    'beta_shape1':np.float32,
    'beta_shape2':np.float32,
    'true_df':np.float32,
    'pval_true_df':np.float64,
    'variant_id':str,
    'tss_distance':np.int32,
    'ma_samples':np.int32,
    'ma_count':np.int32,
    'maf':np.float32,
    'ref_factor':np.int32,
    'pval_nominal':np.float64,
    'slope':np.float32,
    'slope_se':np.float32,
    'pval_perm':np.float64,
    'pval_beta':np.float64,
}

#------------------------------------------------------------------------------
#  Core functions for mapping associations on GPU
#------------------------------------------------------------------------------
class Residualizer(object):
    def __init__(self, C_t):
        # center and orthogonalize
        self.Q_t, _ = tf.qr(C_t - tf.reduce_mean(C_t, 0), full_matrices=False, name='qr')

    def transform(self, M_t, center=True):
        """Residualize rows of M wrt columns of C"""
        if center:
            M0_t = M_t - tf.reduce_mean(M_t, axis=1, keepdims=True)
        else:
            M0_t = M_t
        return M_t - tf.matmul(tf.matmul(M0_t, self.Q_t), self.Q_t, transpose_b=True)  # keep original mean

def residualize(M_t, C_t):
    """Residualize M wrt columns of C"""

    # center and orthogonalize
    Q_t, _ = tf.qr(C_t - tf.reduce_mean(C_t, 0), full_matrices=False, name='qr')

    # residualize M relative to C
    M0_t = M_t - tf.reduce_mean(M_t, axis=1, keepdims=True)
    return M_t - tf.matmul(tf.matmul(M0_t, Q_t), Q_t, transpose_b=True)  # keep original mean

def center_normalize(M_t, axis=0):
    """Center and normalize M"""
    if axis == 0:
        N_t = M_t - tf.reduce_mean(M_t, 0)
        return tf.divide(N_t, tf.sqrt(tf.reduce_sum(tf.pow(N_t, 2), 0)))
    elif axis == 1:
        N_t = M_t - tf.reduce_mean(M_t, axis=1, keepdims=True)
        return tf.divide(N_t, tf.sqrt(tf.reduce_sum(tf.pow(N_t, 2), axis=1, keepdims=True)))

def calculate_maf(genotype_t):
    """Calculate minor allele frequency"""
    af_t = tf.reduce_sum(genotype_t,1) / (2*tf.cast(tf.shape(genotype_t)[1], tf.float32))
    return tf.where(af_t>0.5, 1-af_t, af_t)

def _calculate_corr(genotype_t, phenotype_t, covariates_t, return_sd=False):
    """Calculate correlation between normalized residual genotypes and phenotypes"""
    # residualize
    genotype_res_t = residualize(genotype_t, covariates_t)  # variants x samples
    phenotype_res_t = residualize(phenotype_t, covariates_t)  # phenotypes x samples

    if return_sd:
        _, gstd = tf.nn.moments(genotype_res_t, axes=1)
        _, pstd = tf.nn.moments(phenotype_res_t, axes=1)

    # center and normalize
    genotype_res_t = center_normalize(genotype_res_t, axis=1)
    phenotype_res_t = center_normalize(phenotype_res_t, axis=1)

    # correlation
    if return_sd:
        return tf.squeeze(tf.matmul(genotype_res_t, phenotype_res_t, transpose_b=True)), tf.sqrt(pstd / gstd)
    else:
        return tf.squeeze(tf.matmul(genotype_res_t, phenotype_res_t, transpose_b=True))

def _calculate_max_r2(genotypes_t, phenotypes_t, permutations_t, covariates_t, maf_threshold=0.05):
    maf_t = calculate_maf(genotypes_t)
    ix = tf.where(maf_t>=maf_threshold)

    g2 = tf.squeeze(tf.gather(genotypes_t, ix))
    r2_nom_t = tf.pow(_calculate_corr(g2, phenotypes_t, covariates_t), 2)
    r2_emp_t = tf.pow(_calculate_corr(g2, permutations_t, covariates_t), 2)

    return tf.squeeze(tf.reduce_max(r2_nom_t, axis=0)), tf.squeeze(tf.reduce_max(r2_emp_t, axis=0)), tf.gather(tf.squeeze(ix), tf.argmax(r2_nom_t, axis=0))

def calculate_pval(r2_t, dof, maf_t=None, return_sparse=True, r2_threshold=0, return_r2=False):
    """Calculate p-values from squared correlations"""
    dims = r2_t.get_shape()
    if return_sparse:
        ix = tf.where(r2_t>=r2_threshold, name='threshold_r2')
        r2_t = tf.gather_nd(r2_t, ix)

    r2_t = tf.cast(r2_t, tf.float64)

    tstat = tf.sqrt(tf.divide(tf.scalar_mul(dof, r2_t), 1 - r2_t), name='tstat')
    tdist = tf.contrib.distributions.StudentT(np.float64(dof), loc=np.float64(0.0), scale=np.float64(1.0))

    if return_sparse:
        pval_t = tf.SparseTensor(ix, tf.scalar_mul(2, tdist.cdf(-tf.abs(tstat))), dims)
        if maf_t is not None:
            maf_t = tf.gather(maf_t, ix[:,0])
    else:
        pval_t = tf.scalar_mul(2, tdist.cdf(-tf.abs(tstat)))

    if maf_t is not None:
        if return_r2:
            return pval_t, maf_t, r2_t
        else:
            return pval_t, maf_t
    else:
        return pval_t

def _interaction_assoc_row(genotype_t, phenotype_t, icovariates_t, return_sd=False):
    """
    genotype_t must be a 1D tensor
    icovariates_t: [covariates_t, interaction_t]
    """
    gi_covariates_t = tf.concat([icovariates_t, tf.reshape(genotype_t, [-1,1])], axis=1)
    ix_t = tf.reshape(tf.multiply(genotype_t, icovariates_t[:,-1]), [1,-1])  # must be 1 x N
    return _calculate_corr(ix_t, phenotype_t, gi_covariates_t, return_sd=return_sd)

def calculate_association(genotype_t, phenotype_t, covariates_t, interaction_t=None, return_sparse=True, r2_threshold=None, return_r2=False):
    """Calculate genotype-phenotype associations"""
    maf_t = calculate_maf(genotype_t)

    if interaction_t is None:
        r2_t = tf.pow(_calculate_corr(genotype_t, phenotype_t, covariates_t), 2)
        dof = genotype_t.shape[1].value - 2 - covariates_t.shape[1].value
    else:
        icovariates_t = tf.concat([covariates_t, interaction_t], axis=1)
        r2_t = tf.pow(tf.map_fn(lambda x: _interaction_assoc_row(x, phenotype_t, icovariates_t), genotype_t, infer_shape=False), 2)
        dof = genotype_t.shape[1].value - 4 - covariates_t.shape[1].value

    return calculate_pval(r2_t, dof, maf_t, return_sparse=return_sparse, r2_threshold=r2_threshold, return_r2=return_r2)

def get_sample_indexes(vcf_sample_ids, phenotype_df):
    """Get index of sample IDs in VCF"""
    return tf.constant([vcf_sample_ids.index(i) for i in phenotype_df.columns])

def initialize_data(phenotype_df, covariates_df, batch_size, interaction_s=None, dtype=tf.float32):
    """Generate placeholders"""
    num_samples = phenotype_df.shape[1]
    genotype_t = tf.placeholder(dtype, shape=[batch_size, num_samples])
    phenotype_t = tf.constant(phenotype_df.values, dtype=dtype)
    phenotype_t = tf.reshape(phenotype_t, shape=[-1, num_samples])
    covariates_t = tf.constant(covariates_df.values, dtype=dtype)
    covariates_t = tf.reshape(covariates_t, shape=[-1, covariates_df.shape[1]])
    if interaction_s is None:
        return genotype_t, phenotype_t, covariates_t
    else:
        interaction_t = tf.constant(interaction_s.values, dtype=dtype)
        interaction_t = tf.reshape(interaction_t, [-1,1])
        return genotype_t, phenotype_t, covariates_t, interaction_t

#------------------------------------------------------------------------------
#  Functions for beta-approximating empirical p-values
#------------------------------------------------------------------------------
def pval_from_corr(r2, dof):
    tstat2 = dof * r2 / (1 - r2)
    return 2*stats.t.cdf(-np.abs(np.sqrt(tstat2)), dof)

def df_cost(r2, dof):
    """minimize abs(1-alpha) as a function of M_eff"""
    pval = pval_from_corr(r2, dof)
    mean = np.mean(pval)
    var = np.var(pval)
    return mean * (mean * (1.0-mean) / var - 1.0) - 1.0

def beta_log_likelihood(x, shape1, shape2):
    """negative log-likelihood of beta distribution"""
    logbeta = loggamma(shape1) + loggamma(shape2) - loggamma(shape1+shape2)
    return (1.0-shape1)*np.sum(np.log(x)) + (1.0-shape2)*np.sum(np.log(1.0-x)) + len(x)*logbeta

def fit_beta_parameters(r2_perm, dof, tol=1e-4, return_minp=False):
    """
      r2_perm:    array of max. r2 values from permutations
      dof:        degrees of freedom
    """
    try:
        true_dof = scipy.optimize.newton(lambda x: df_cost(r2_perm, x), dof, tol=tol, maxiter=50)
    except:
        print('WARNING: scipy.optimize.newton failed to converge (running scipy.optimize.minimize)')
        res = scipy.optimize.minimize(lambda x: np.abs(df_cost(r2_perm, x)), dof, method='Nelder-Mead', tol=tol)
        true_dof = res.x[0]

    pval = pval_from_corr(r2_perm, true_dof)
    mean, var = np.mean(pval), np.var(pval)
    beta_shape1 = mean * (mean * (1 - mean) / var - 1)
    beta_shape2 = beta_shape1 * (1/mean - 1)
    res = scipy.optimize.minimize(lambda s: beta_log_likelihood(pval, s[0], s[1]), [beta_shape1, beta_shape2], method='Nelder-Mead', tol=tol)
    beta_shape1, beta_shape2 = res.x
    if return_minp:
        return beta_shape1, beta_shape2, true_dof, pval
    else:
        return beta_shape1, beta_shape2, true_dof

def calculate_beta_approx_pval(r2_perm, r2_nominal, dof, tol=1e-4):
    """
      r2_nominal: nominal max. r2 (scalar or array)
      r2_perm:    array of max. r2 values from permutations
      dof:        degrees of freedom
    """
    beta_shape1, beta_shape2, true_dof = fit_beta_parameters(r2_perm, dof, tol)
    pval_true_dof = pval_from_corr(r2_nominal, true_dof)
    pval_beta = stats.beta.cdf(pval_true_dof, beta_shape1, beta_shape2)
    return pval_beta, beta_shape1, beta_shape2, true_dof, pval_true_dof

#------------------------------------------------------------------------------
#  Top-level functions for running cis-/trans-QTL mapping
#------------------------------------------------------------------------------
def calculate_cis_permutations(genotypes_t, range_t, phenotype_t, covariates_t, permutation_ix_t):
    """Calculate nominal and empirical correlations"""
    permutations_t = tf.gather(phenotype_t, permutation_ix_t)

    r_nominal_t, std_ratio_t = _calculate_corr(genotypes_t, tf.reshape(phenotype_t, [1,-1]), covariates_t, return_sd=True)

    corr_t = tf.pow(_calculate_corr(genotypes_t, permutations_t, covariates_t), 2)
    corr_t.set_shape([None,None])
    r2_perm_t = tf.cast(tf.reduce_max(tf.boolean_mask(corr_t, ~tf.reduce_any(tf.is_nan(corr_t), 1)), axis=0), tf.float64)

    ix = tf.argmax(tf.pow(r_nominal_t, 2))
    return r_nominal_t[ix], std_ratio_t[ix], range_t[ix], r2_perm_t, genotypes_t[ix], tf.shape(r_nominal_t)[0]

def process_cis_permutations(r2_perm, r_nominal, std_ratio, g, num_var, dof, n_samples, nperm=10000):
    """Calculate beta-approximated empirical p-value and annotate phenotype"""
    r2_nominal = r_nominal*r_nominal
    pval_perm = (np.sum(r2_perm>=r2_nominal)+1) / (nperm+1)
    pval_beta, beta_shape1, beta_shape2, true_dof, pval_true_dof = calculate_beta_approx_pval(r2_perm, r2_nominal, dof)

    maf = np.sum(g) / (2*n_samples)
    if maf <= 0.5:
        ref_factor = 1
        ma_samples = np.sum(g>0.5)
        ma_count = np.sum(g[g>0.5])
    else:
        maf = 1-maf
        ref_factor = -1
        ma_samples = np.sum(g<1.5)
        ma_count = np.sum(g[g<1.5])

    slope = r_nominal * std_ratio
    tstat2 = dof * r2_nominal / (1 - r2_nominal)
    slope_se = np.abs(slope) / np.sqrt(tstat2)

    return pd.Series(OrderedDict([
        ('num_var', num_var),
        ('beta_shape1', beta_shape1),
        ('beta_shape2', beta_shape2),
        ('true_df', true_dof),
        ('pval_true_df', pval_true_dof),
        ('variant_id', np.NaN),
        ('tss_distance', np.NaN),
        ('ma_samples', ma_samples),
        ('ma_count', ma_count),
        ('maf', maf),
        ('ref_factor', ref_factor),
        ('pval_nominal', pval_from_corr(r2_nominal, dof)),
        ('slope', slope),
        ('slope_se', slope_se),
        ('pval_perm', pval_perm),
        ('pval_beta', pval_beta),
    ]))

def _process_group_permutations(buf):
    """
    Merge results for grouped phenotypes

    buf: [r_nom, s_r, var_ix, r2_perm, g, ng, nid]
    """
    # select phenotype with strongest nominal association
    max_ix = np.argmax(np.abs([b[0] for b in buf]))
    r_nom, s_r, var_ix = buf[max_ix][:3]
    g, ng, nid = buf[max_ix][4:]
    # select best phenotype correlation for each permutation
    r2_perm = np.max([b[3] for b in buf], 0)
    return r_nom, s_r, var_ix, r2_perm, g, ng, nid

def map_cis(plink_reader, phenotype_df, phenotype_pos_df, covariates_df, group_s=None, nperm=10000, logger=None):
    """Run cis-QTL mapping"""
    assert np.all(phenotype_df.columns==covariates_df.index)
    if logger is None:
        logger = SimpleLogger()

    logger.write('cis-QTL mapping: empirical p-values for phenotypes')
    logger.write('  * {} samples'.format(phenotype_df.shape[1]))
    logger.write('  * {} phenotypes'.format(phenotype_df.shape[0]))
    if group_s is not None:
        logger.write('  * {} phenotype groups'.format(len(group_s.unique())))
        group_dict = group_s.to_dict()
    logger.write('  * {} covariates'.format(covariates_df.shape[1]))
    logger.write('  * {} variants'.format(plink_reader.bed.shape[0]))

    dof = phenotype_df.shape[1] - 2 - covariates_df.shape[1]

    # permutation indices
    n_samples = phenotype_df.shape[1]
    ix = np.arange(n_samples)
    permutation_ix_t = tf.convert_to_tensor(np.array([np.random.permutation(ix) for i in range(nperm)]))

    # placeholders
    covariates_t = tf.constant(covariates_df.values, dtype=tf.float32)
    genotype_t = tf.placeholder(dtype=tf.float32, shape=(None))
    phenotype_t = tf.placeholder(dtype=tf.float32, shape=(None))

    # iterate over chromosomes
    res_df = []
    start_time = time.time()
    with tf.Session() as sess:
        for chrom in phenotype_pos_df.loc[phenotype_df.index, 'chr'].unique():
            logger.write('  Mapping chromosome {}'.format(chrom))
            igc = genotypeio.InputGeneratorCis(plink_reader, phenotype_df.loc[phenotype_pos_df['chr']==chrom], phenotype_pos_df)

            dataset = tf.data.Dataset.from_generator(igc.generate_data, output_types=(tf.float32, tf.float32, tf.int32, tf.string))
            dataset = dataset.prefetch(1)
            iterator = dataset.make_one_shot_iterator()
            next_phenotype, next_genotypes, next_range, next_id = iterator.get_next()

            r_nominal_t, std_ratio_t, varpos_t, r2_perm_t, g_t, ng_t = calculate_cis_permutations(
                next_genotypes, next_range, next_phenotype, covariates_t, permutation_ix_t)

            if group_s is None:
                for i in range(1, igc.n_phenotypes+1):
                    r_nom, s_r, var_ix, r2_perm, g, ng, nid = sess.run([r_nominal_t, std_ratio_t, varpos_t, r2_perm_t, g_t, ng_t, next_id])

                    # post-processing (on CPU)
                    res_s = process_cis_permutations(r2_perm, r_nom, s_r, g, ng, dof, phenotype_df.shape[1], nperm=nperm)
                    res_s.name = nid.decode()
                    res_s['variant_id'] = igc.chr_variant_pos.index[var_ix]
                    res_s['tss_distance'] = igc.chr_variant_pos[res_s['variant_id']] - igc.phenotype_tss[res_s.name]
                    res_df.append(res_s)
                    print('\r  * computing permutations for phenotype {}/{}'.format(i, igc.n_phenotypes), end='')
                print()
            else:
                n_groups = len(igc.phenotype_df.index.map(group_dict).unique())
                buf = []
                processed_groups = 0
                previous_group = None
                for i in range(0, igc.n_phenotypes):
                    group_id = group_dict[igc.phenotype_df.index[i]]
                    ires = sess.run([r_nominal_t, std_ratio_t, varpos_t, r2_perm_t, g_t, ng_t, next_id])
                    if (group_id != previous_group and len(buf)>0):  # new group, process previous
                        # post-processing (on CPU)
                        r_nom, s_r, var_ix, r2_perm, g, ng, nid = _process_group_permutations(buf)
                        res_s = process_cis_permutations(r2_perm, r_nom, s_r, g, ng, dof, phenotype_df.shape[1], nperm=nperm)
                        res_s.name = nid.decode()
                        res_s['variant_id'] = igc.chr_variant_pos.index[var_ix]
                        res_s['tss_distance'] = igc.chr_variant_pos[res_s['variant_id']] - igc.phenotype_tss[res_s.name]
                        res_s['group_id'] = group_id
                        res_s['group_size'] = len(buf)
                        res_df.append(res_s)
                        processed_groups += 1
                        print('\r  * computing permutations for phenotype group {}/{}'.format(processed_groups, n_groups), end='')
                        # reset
                        buf = [ires]
                    else:
                        buf.append(ires)
                    previous_group = group_id
                # process last group
                r_nom, s_r, var_ix, r2_perm, g, ng, nid = _process_group_permutations(buf)
                res_s = process_cis_permutations(r2_perm, r_nom, s_r, g, ng, dof, phenotype_df.shape[1], nperm=nperm)
                res_s.name = nid.decode()
                res_s['variant_id'] = igc.chr_variant_pos.index[var_ix]
                res_s['tss_distance'] = igc.chr_variant_pos[res_s['variant_id']] - igc.phenotype_tss[res_s.name]
                res_s['group_id'] = group_id
                res_s['group_size'] = len(buf)
                res_df.append(res_s)
                processed_groups += 1
                print('\r  * computing permutations for phenotype group {}/{}'.format(processed_groups, n_groups), end='\n')

    res_df = pd.concat(res_df, axis=1).T
    res_df.index.name = 'phenotype_id'
    logger.write('  Time elapsed: {:.2f} min'.format((time.time()-start_time)/60))
    logger.write('done.')
    return res_df.astype(output_dtype_dict)

def map_cis_independent(plink_reader, summary_df, phenotype_df, phenotype_pos_df, covariates_df, fdr=0.05, fdr_col='qval', nperm=10000, logger=None):
    """
    Run independent cis-QTL mapping (forward-backward regression)

    summary_df: output from map_cis, annotated with q-values (calculate_qvalues)
    """
    assert np.all(phenotype_df.index==phenotype_pos_df.index)
    if logger is None:
        logger = SimpleLogger()

    signif_df = summary_df[summary_df[fdr_col]<=fdr].copy()
    signif_df = signif_df[[
        'num_var', 'beta_shape1', 'beta_shape2', 'true_df', 'pval_true_df',
        'variant_id', 'tss_distance', 'ma_samples', 'ma_count', 'maf', 'ref_factor',
        'pval_nominal', 'slope', 'slope_se', 'pval_perm', 'pval_beta',
    ]]
    signif_threshold = signif_df['pval_beta'].max()
    # subset significant phenotypes
    ix = signif_df.index[signif_df.index.isin(phenotype_df.index)]
    phenotype_df = phenotype_df.loc[ix]
    phenotype_pos_df = phenotype_pos_df.loc[ix]

    logger.write('cis-QTL mapping: conditionally independent variants')
    logger.write('  * {} samples'.format(phenotype_df.shape[1]))
    logger.write('  * {} significant phenotypes'.format(signif_df.shape[0]))
    logger.write('  * {} covariates'.format(covariates_df.shape[1]))
    logger.write('  * {} variants'.format(plink_reader.bed.shape[0]))

    # print('Significance threshold: {}'.format(signif_threshold))

    # permutation indices
    n_samples = phenotype_df.shape[1]
    ix = np.arange(n_samples)
    permutation_ix_t = tf.convert_to_tensor(np.array([np.random.permutation(ix) for i in range(nperm)]))

    # placeholders
    genotypes_t = tf.placeholder(dtype=tf.float32, shape=[None,None])
    range_t = tf.placeholder(dtype=tf.int32, shape=[None])
    phenotype_t = tf.placeholder(dtype=tf.float32, shape=[None])
    covariates_t = tf.placeholder(dtype=tf.float32, shape=[None, None])

    # graph
    r_nominal_t, std_ratio_t, varpos_t, r2_perm_t, g_t, ng_t = calculate_cis_permutations(
        genotypes_t, range_t, phenotype_t, covariates_t, permutation_ix_t)

    # iterate over chromosomes
    res_df = []
    start_time = time.time()
    with tf.Session() as sess:
        for chrom in phenotype_pos_df['chr'].unique():
            logger.write('  Mapping chromosome {}'.format(chrom))
            igc = genotypeio.InputGeneratorCis(plink_reader, phenotype_df.loc[phenotype_pos_df['chr']==chrom], phenotype_pos_df)
            ix_dict = {i:k for k,i in enumerate(plink_reader.bim.loc[plink_reader.bim['chrom']==chrom, 'snp'])}

            igc.chr_genotypes, igc.chr_variant_pos = igc.plink_reader.get_region(chrom, verbose=False)
            igc.loaded_chrom = chrom

            # iterate through significant phenotypes, fetch associated genotypes
            for j, (phenotype, genotypes, genotype_range, phenotype_id) in enumerate(igc.generate_data(), 1):

                # 1) forward pass
                forward_df = [signif_df.loc[phenotype_id]]  # initialize results with top variant
                covariates = covariates_df.values.copy()  # initialize covariates
                dosage_dict = {}
                while True:
                    # add variant to covariates
                    variant_id = forward_df[-1]['variant_id']
                    ig = igc.chr_genotypes[ix_dict[variant_id]]
                    dosage_dict[variant_id] = ig
                    covariates = np.hstack([covariates, ig.reshape(-1,1)]).astype(np.float32)
                    dof = phenotype_df.shape[1] - 2 - covariates.shape[1]
                    # find next variant
                    r_nom, s_r, var_ix, r2_perm, g, ng = sess.run([r_nominal_t, std_ratio_t, varpos_t, r2_perm_t, g_t, ng_t],
                        feed_dict={genotypes_t:genotypes, range_t:genotype_range, phenotype_t:phenotype, covariates_t:covariates})
                    res_s = process_cis_permutations(r2_perm, r_nom, s_r, g, ng, dof, phenotype_df.shape[1], nperm=nperm)
                    # add to list if significant
                    if res_s['pval_beta'] <= signif_threshold:
                        res_s.name = phenotype_id
                        res_s['variant_id'] = igc.chr_variant_pos.index[var_ix]
                        res_s['tss_distance'] = igc.chr_variant_pos[res_s['variant_id']] - igc.phenotype_tss[res_s.name]
                        forward_df.append(res_s)
                    else:
                        break
                forward_df = pd.concat(forward_df, axis=1).T
                dosage_df = pd.DataFrame(dosage_dict)
                # print(forward_df)

                # 2) backward pass
                if forward_df.shape[0]>1:
                    back_df = []
                    variant_set = set()
                    for k,i in enumerate(forward_df['variant_id'], 1):
                        covariates = np.hstack([
                            covariates_df.values,
                            dosage_df[np.setdiff1d(forward_df['variant_id'], i)].values,
                        ])
                        r_nom, s_r, var_ix, r2_perm, g, ng = sess.run([r_nominal_t, std_ratio_t, varpos_t, r2_perm_t, g_t, ng_t],
                            feed_dict={genotypes_t:genotypes, range_t:genotype_range, phenotype_t:phenotype, covariates_t:covariates})
                        dof = phenotype_df.shape[1] - 2 - covariates.shape[1]
                        res_s = process_cis_permutations(r2_perm, r_nom, s_r, g, ng, dof, phenotype_df.shape[1], nperm=nperm)
                        res_s['variant_id'] = igc.chr_variant_pos.index[var_ix]
                        if res_s['pval_beta'] <= signif_threshold and res_s['variant_id'] not in variant_set:
                            res_s.name = phenotype_id
                            res_s['tss_distance'] = igc.chr_variant_pos[res_s['variant_id']] - igc.phenotype_tss[res_s.name]
                            res_s['rank'] = k
                            back_df.append(res_s)
                            variant_set.add(res_s['variant_id'])
                    if len(back_df)>0:
                        res_df.append(pd.concat(back_df, axis=1).T)
                    # print('back')
                    # print(pd.concat(back_df, axis=1).T)
                else:  # single variant
                    forward_df['rank'] = 1
                    res_df.append(forward_df)
                print('\r  * computing independent QTL for phenotype {}/{}'.format(j, igc.n_phenotypes), end='')
            print()
    res_df = pd.concat(res_df, axis=0)
    res_df.index.name = 'phenotype_id'
    logger.write('  Time elapsed: {:.2f} min'.format((time.time()-start_time)/60))
    logger.write('done.')
    return res_df.reset_index().astype(output_dtype_dict)

def calculate_qvalues(res_df, fdr=0.05, qvalue_lambda=None):
    """Annotate permutation results with q-values, p-value threshold"""

    print('Computing q-values')
    print('  * Number of phenotypes tested: {}'.format(res_df.shape[0]))
    print('  * Correlation between Beta-approximated and empirical p-values: : {:.4f}'.format(
        stats.pearsonr(res_df['pval_perm'], res_df['pval_beta'])[0]))

    # calculate q-values
    if qvalue_lambda is None:
        qval, pi0 = rfunc.qvalue(res_df['pval_beta'])
    else:
        print('  * Calculating q-values with lambda = {:.3f}'.format(qvalue_lambda))
        qval, pi0 = rfunc.qvalue(res_df['pval_beta'], qvalue_lambda)
    res_df['qval'] = qval
    print('  * Proportion of significant phenotypes (1-pi0): {:.2f}'.format(1 - pi0))
    print('  * QTL phenotypes @ FDR {:.2f}: {}'.format(fdr, np.sum(res_df['qval']<=fdr)))

    # determine global min(p) significance threshold and calculate nominal p-value threshold for each gene
    ub = res_df.loc[res_df['qval']>fdr, 'pval_beta'].sort_values()[0]
    lb = res_df.loc[res_df['qval']<=fdr, 'pval_beta'].sort_values()[-1]
    pthreshold = (lb+ub)/2
    print('  * min p-value threshold @ FDR {}: {:.6g}'.format(fdr, pthreshold))
    res_df['pval_nominal_threshold'] = stats.beta.ppf(pthreshold, res_df['beta_shape1'], res_df['beta_shape2'])

def annotate_genes(gene_df, annotation_gtf, lookup_df=None):
    """
    Add gene and variant annotations (e.g., gene_name, rs_id, etc.) to gene-level output

    gene_df:        output from map_cis()
    annotation_gtf: gene annotation in GTF format
    lookup_df:      DataFrame with variant annotations, indexed by 'variant_id'
    """
    gene_dict = {}
    print('['+datetime.now().strftime("%b %d %H:%M:%S")+'] Adding gene and variant annotations', flush=True)
    print('  * parsing GTF', flush=True)
    with open(annotation_gtf) as gtf:
        for row in gtf:
            row = row.strip().split('\t')
            if row[0][0]=='#' or row[2]!='gene': continue
            # get gene_id and gene_name from attributes
            attr = dict([i.split() for i in row[8].replace('"','').split(';') if i!=''])
            # gene_name, gene_chr, gene_start, gene_end, strand
            gene_dict[attr['gene_id']] = [attr['gene_name'], row[0], row[3], row[4], row[6]]

    print('  * annotating genes', flush=True)
    if 'group_id' in gene_df:
        gene_info = pd.DataFrame(data=[gene_dict[i] for i in gene_df['group_id']],
                                 columns=['gene_name', 'gene_chr', 'gene_start', 'gene_end', 'strand'],
                                 index=gene_df.index)
    else:
        gene_info = pd.DataFrame(data=[gene_dict[i] for i in gene_df.index],
                                 columns=['gene_name', 'gene_chr', 'gene_start', 'gene_end', 'strand'],
                                 index=gene_df.index)
    gene_df = pd.concat([gene_info, gene_df], axis=1)
    assert np.all(gene_df.index==gene_info.index)

    col_order = ['gene_name', 'gene_chr', 'gene_start', 'gene_end', 'strand',
        'num_var', 'beta_shape1', 'beta_shape2', 'true_df', 'pval_true_df', 'variant_id', 'tss_distance']
    if lookup_df is not None:
        print('  * adding variant annotations from lookup table', flush=True)
        gene_df = gene_df.join(lookup_df, on='variant_id')  # add variant information
        col_order += list(lookup_df.columns)
    col_order += ['ma_samples', 'ma_count', 'maf', 'ref_factor',
        'pval_nominal', 'slope', 'slope_se', 'pval_perm', 'pval_beta']
    if 'group_id' in gene_df:
        col_order += ['group_id', 'group_size']
    col_order += ['qval', 'pval_nominal_threshold']
    gene_df = gene_df[col_order]
    print('done.', flush=True)
    return gene_df

def get_significant_pairs(res_df, nominal_prefix, fdr=0.05):
    """Significant variant-phenotype pairs based on nominal p-value threshold for each phenotype"""
    print('['+datetime.now().strftime("%b %d %H:%M:%S")+'] tensorQTL: filtering significant variant-phenotype pairs', flush=True)
    assert 'qval' in res_df

    # significant phenotypes (apply FDR threshold)
    df = res_df.loc[res_df['qval']<=fdr, ['pval_nominal_threshold', 'pval_nominal', 'pval_beta']].copy()
    df.rename(columns={'pval_nominal': 'min_pval_nominal'}, inplace=True)
    signif_phenotype_ids = set(df.index)
    threshold_dict = df['pval_nominal_threshold'].to_dict()

    nominal_files = {os.path.basename(i).split('.')[-2]:i for i in glob.glob(nominal_prefix+'*.parquet')}
    chroms = sorted(nominal_files.keys())
    signif_df = []
    for k,c in enumerate(chroms, 1):
        print('  * parsing significant variant-phenotype pairs for chr. {}/{}'.format(k, len(chroms)), end='\r', flush=True)
        nominal_df = pd.read_parquet(nominal_files[c])
        nominal_df = nominal_df[nominal_df['phenotype_id'].isin(signif_phenotype_ids)]

        m = nominal_df['pval_nominal']<nominal_df['phenotype_id'].apply(lambda x: threshold_dict[x])
        signif_df.append(nominal_df[m])
    print()
    signif_df = pd.concat(signif_df, axis=0)
    signif_df = signif_df.merge(df, left_on='phenotype_id', right_index=True)
    print('['+datetime.now().strftime("%b %d %H:%M:%S")+'] done', flush=True)
    return signif_df
    # signif_df.to_parquet(nominal_prefix.rsplit('.',1)[0]+'.cis_qtl_significant_pairs.parquet')

def calculate_cis_nominal(genotypes_t, phenotype_t, covariates_t, dof):
    """
    Calculate nominal associations

    genotypes_t: genotypes x samples
    phenotype_t: single phenotype
    covariates_t: covariates matrix, samples x covariates
    """
    p = tf.reshape(phenotype_t, [1,-1])

    r_nominal_t, std_ratio_t = _calculate_corr(genotypes_t, p, covariates_t, return_sd=True)
    r2_nominal_t = tf.pow(r_nominal_t, 2)
    pval_t = calculate_pval(r2_nominal_t, dof, maf_t=None, return_sparse=False)
    slope_t = tf.multiply(r_nominal_t, std_ratio_t)
    slope_se_t = tf.divide(tf.abs(slope_t), tf.sqrt(tf.divide(tf.scalar_mul(dof, r2_nominal_t), 1 - r2_nominal_t)))

    # calculate MAF
    n = covariates_t.shape[0].value
    n2 = 2*n
    af_t = tf.reduce_sum(genotypes_t,1) / n2
    ix_t = af_t<=0.5
    maf_t = tf.where(ix_t, af_t, 1-af_t)
    # calculate MA samples and counts
    m = tf.cast(genotypes_t>0.5, tf.float32)
    a = tf.reduce_sum(m, 1)
    b = tf.reduce_sum(tf.cast(genotypes_t<1.5, tf.float32), 1)
    ma_samples_t = tf.where(ix_t, a, b)
    m = tf.multiply(m, genotypes_t)
    a = tf.reduce_sum(m, 1)
    ma_count_t = tf.where(ix_t, a, n2-a)

    return pval_t, slope_t, slope_se_t, maf_t, ma_samples_t, ma_count_t

def map_cis_nominal(plink_reader, phenotype_df, phenotype_pos_df, covariates_df, prefix,
                    output_dir='.', logger=None):
    """
    cis-QTL mapping: nominal associations for all variant-phenotype pairs

    Association results for each chromosome are written to parquet files
    in the format <output_dir>/<prefix>.cis_qtl_pairs.<chr>.parquet
    """
    if logger is None:
        logger = SimpleLogger()

    logger.write('cis-QTL mapping: nominal associations for all variant-phenotype pairs')
    logger.write('  * {} samples'.format(phenotype_df.shape[1]))
    logger.write('  * {} phenotypes'.format(phenotype_df.shape[0]))
    logger.write('  * {} covariates'.format(covariates_df.shape[1]))
    logger.write('  * {} variants'.format(plink_reader.bed.shape[0]))

    # placeholders
    covariates_t = tf.constant(covariates_df.values, dtype=tf.float32)
    genotype_t = tf.placeholder(dtype=tf.float32, shape=(None))
    phenotype_t = tf.placeholder(dtype=tf.float32, shape=(None))
    dof = phenotype_df.shape[1] - 2 - covariates_df.shape[1]

    with tf.Session() as sess:
        # iterate over chromosomes
        start_time = time.time()
        for chrom in phenotype_pos_df.loc[phenotype_df.index, 'chr'].unique():
            logger.write('  Mapping chromosome {}'.format(chrom))
            igc = genotypeio.InputGeneratorCis(plink_reader, phenotype_df.loc[phenotype_pos_df['chr']==chrom], phenotype_pos_df)

            dataset = tf.data.Dataset.from_generator(igc.generate_data, output_types=(tf.float32, tf.float32, tf.int32, tf.string))
            dataset = dataset.prefetch(1)
            iterator = dataset.make_one_shot_iterator()
            next_phenotype, next_genotypes, _, next_id = iterator.get_next()
            x = calculate_cis_nominal(next_genotypes, next_phenotype, covariates_t, dof)

            chr_res_df = []
            for i in range(1, igc.n_phenotypes+1):
                (pval_nominal, slope, slope_se, maf, ma_samples, ma_count), phenotype_id = sess.run([x, next_id])
                phenotype_id = phenotype_id.decode()
                r = igc.cis_ranges[phenotype_id]
                variant_ids = plink_reader.variant_pos[chrom].index[r[0]:r[1]+1]
                nv = len(variant_ids)
                tss_distance = np.int32(plink_reader.variant_pos[chrom].values[r[0]:r[1]+1] - igc.phenotype_tss[phenotype_id])

                chr_res_df.append(pd.DataFrame(OrderedDict([
                    ('phenotype_id', [phenotype_id]*nv),
                    ('variant_id', variant_ids),
                    ('tss_distance', tss_distance),
                    ('maf', maf),
                    ('ma_samples', np.int32(ma_samples)),
                    ('ma_count', np.int32(ma_count)),
                    ('pval_nominal', pval_nominal),
                    ('slope', slope),
                    ('slope_se', slope_se),
                ])))
                print('\r    computing associations for phenotype {}/{}'.format(i, igc.n_phenotypes), end='')
            print()
            logger.write('    time elapsed: {:.2f} min'.format((time.time()-start_time)/60))
            print('  * writing output')
            pd.concat(chr_res_df, copy=False).to_parquet(os.path.join(output_dir, '{}.cis_qtl_pairs.{}.parquet'.format(prefix, chrom)))
    logger.write('done.')

def calculate_nominal_interaction(genotypes_t, phenotype_t, interaction_t, dof, residualizer,
                                  interaction_mask_t=None, maf_threshold_interaction=0.05,
                                  return_sparse=False, tstat_threshold=None):
    """"""
    phenotype_t = tf.reshape(phenotype_t, [1,-1])
    # filter monomorphic sites (to avoid colinearity in X)
    mask_t = ~(tf.reduce_all(tf.equal(genotypes_t, 0), axis=1) |
               tf.reduce_all(tf.equal(genotypes_t, 1), axis=1) |
               tf.reduce_all(tf.equal(genotypes_t, 2), axis=1)
               )

    if interaction_mask_t is not None:
        upper_t = tf.boolean_mask(genotypes_t, interaction_mask_t, axis=1)
        lower_t = tf.boolean_mask(genotypes_t, ~interaction_mask_t, axis=1)
        mask_t = (mask_t &
                  (calculate_maf(upper_t) >= maf_threshold_interaction) &
                  (calculate_maf(lower_t) >= maf_threshold_interaction)
                  )

    mask_t.set_shape([None])  # required for tf.boolean_mask
    genotypes_t = tf.boolean_mask(genotypes_t, mask_t)

    s = tf.shape(genotypes_t)
    ng = s[0]
    ns = tf.cast(s[1], tf.float32)

    g0_t = genotypes_t - tf.reduce_mean(genotypes_t, axis=1, keepdims=True)
    gi_t = tf.multiply(genotypes_t, interaction_t)
    gi0_t = gi_t - tf.reduce_mean(gi_t, axis=1, keepdims=True)

    i0_t = interaction_t - tf.reduce_mean(interaction_t)
    p_t = tf.reshape(phenotype_t, [1,-1])
    p0_t = p_t - tf.reduce_mean(p_t, axis=1, keepdims=True)

    # residualize rows
    g0_t = residualizer.transform(g0_t, center=False)
    gi0_t = residualizer.transform(gi0_t, center=False)
    p0_t = residualizer.transform(p0_t, center=False)
    i0_t = residualizer.transform(i0_t, center=False)
    i0_t = tf.tile(i0_t, [ng, 1])

    # regression
    X_t = tf.stack([g0_t, i0_t, gi0_t], axis=2)
    Xinv = tf.linalg.inv(tf.matmul(X_t, X_t, transpose_a=True))
    b_t = tf.matmul(tf.matmul(Xinv, X_t, transpose_b=True), tf.tile(tf.reshape(p0_t, [1,1,-1]), [ng,1,1]), transpose_b=True)

    # calculate b, b_se
    r_t = tf.squeeze(tf.matmul(X_t, b_t)) - p0_t
    rss_t = tf.reduce_sum(tf.multiply(r_t, r_t), axis=1)
    Cx = Xinv * tf.reshape(rss_t, [-1,1,1]) / dof
    b_se_t = tf.sqrt(tf.matrix_diag_part(Cx))
    b_t = tf.squeeze(b_t)
    tstat_t = tf.divide(tf.cast(b_t, tf.float64), tf.cast(b_se_t, tf.float64))
    # weird tf bug? without cast/copy, divide appears to modify b_se_t??
    # calculate pval
    tdist = tf.contrib.distributions.StudentT(np.float64(dof), loc=np.float64(0.0), scale=np.float64(1.0))
    pval_t = tf.scalar_mul(2, tdist.cdf(-tf.abs(tstat_t)))

    # calculate MAF
    n2 = 2*ns
    af_t = tf.reduce_sum(genotypes_t,1) / n2
    ix_t = af_t<=0.5
    maf_t = tf.where(ix_t, af_t, 1-af_t)
    # calculate MA samples and counts
    m = tf.cast(genotypes_t>0.5, tf.float32)
    a = tf.reduce_sum(m, 1)
    b = tf.reduce_sum(tf.cast(genotypes_t<1.5, tf.float32), 1)
    ma_samples_t = tf.where(ix_t, a, b)
    m = tf.multiply(m, genotypes_t)
    a = tf.reduce_sum(m, 1)
    ma_count_t = tf.where(ix_t, a, n2-a)

    return pval_t, b_t, b_se_t, maf_t, ma_samples_t, ma_count_t, mask_t

def map_cis_interaction_nominal(plink_reader, phenotype_df, phenotype_pos_df, covariates_df, interaction_s,
                                prefix, maf_threshold_interaction=0.05, best_only=False, output_dir='.', logger=None):
    """
    cis-QTL mapping: nominal associations for all variant-phenotype pairs

    Association results for each chromosome are written to parquet files
    in the format <output_dir>/<prefix>.cis_qtl_pairs.<chr>.parquet
    """
    if logger is None:
        logger = SimpleLogger()

    logger.write('cis-QTL mapping: nominal associations for all variant-phenotype pairs')
    logger.write('  * {} samples'.format(phenotype_df.shape[1]))
    logger.write('  * {} phenotypes'.format(phenotype_df.shape[0]))
    logger.write('  * {} covariates'.format(covariates_df.shape[1]))
    logger.write('  * {} variants'.format(plink_reader.bed.shape[0]))
    logger.write('  * including interaction term')

    covariates_t = tf.constant(covariates_df.values, dtype=tf.float32)
    dof = phenotype_df.shape[1] - 4 - covariates_df.shape[1]
    interaction_t = tf.constant(interaction_s.values.reshape(1,-1), dtype=tf.float32)  # 1 x n
    interaction_mask_t = tf.constant(interaction_s >= interaction_s.median())
    residualizer = Residualizer(covariates_t)

    with tf.Session() as sess:
        # iterate over chromosomes
        start_time = time.time()
        best_assoc = []
        for chrom in phenotype_pos_df.loc[phenotype_df.index, 'chr'].unique():
            logger.write('  Mapping chromosome {}'.format(chrom))
            igc = genotypeio.InputGeneratorCis(plink_reader, phenotype_df.loc[phenotype_pos_df['chr']==chrom], phenotype_pos_df)

            dataset = tf.data.Dataset.from_generator(igc.generate_data, output_types=(tf.float32, tf.float32, tf.int32, tf.string))
            dataset = dataset.prefetch(1)
            iterator = dataset.make_one_shot_iterator()
            next_phenotype, next_genotypes, _, next_id = iterator.get_next()

            x = calculate_nominal_interaction(next_genotypes, next_phenotype, interaction_t, dof, residualizer,
                                      interaction_mask_t=interaction_mask_t, maf_threshold_interaction=0.05)

            chr_res_df = []
            for i in range(1, igc.n_phenotypes+1):
                (pval, b, b_se, maf, ma_samples, ma_count, maf_mask), phenotype_id = sess.run([x, next_id])

                phenotype_id = phenotype_id.decode()
                r = igc.cis_ranges[phenotype_id]
                variant_ids = plink_reader.variant_pos[chrom].index[r[0]:r[1]+1]
                tss_distance = np.int32(plink_reader.variant_pos[chrom].values[r[0]:r[1]+1] - igc.phenotype_tss[phenotype_id])
                if maf_mask is not None:
                    variant_ids = variant_ids[maf_mask]
                    tss_distance = tss_distance[maf_mask]
                nv = len(variant_ids)
                df = pd.DataFrame(OrderedDict([
                    ('phenotype_id', [phenotype_id]*nv),
                    ('variant_id', variant_ids),
                    ('tss_distance', tss_distance),
                    ('maf', maf),
                    ('ma_samples', np.int32(ma_samples)),
                    ('ma_count', np.int32(ma_count)),
                    ('pval_g', pval[:,0]),
                    ('b_g', b[:,0]),
                    ('b_g_se', b_se[:,0]),
                    ('pval_i', pval[:,1]),
                    ('b_i', b[:,1]),
                    ('b_i_se', b_se[:,1]),
                    ('pval_gi', pval[:,2]),
                    ('b_gi', b[:,2]),
                    ('b_gi_se', b_se[:,2]),
                ]))
                if best_only:
                    best_assoc.append(df.loc[df['pval_gi'].idxmin()])
                else:
                    chr_res_df.append(df)
                print('\r    computing associations for phenotype {}/{}'.format(i, igc.n_phenotypes), end='')
            print()
            logger.write('    time elapsed: {:.2f} min'.format((time.time()-start_time)/60))
            if not best_only:
                print('  * writing output')
                pd.concat(chr_res_df, copy=False).to_parquet(os.path.join(output_dir, '{}.cis_qtl_pairs.{}.parquet'.format(prefix, chrom)))
        if best_only:
            pd.concat(best_assoc, axis=1).T.set_index('phenotype_id').to_csv(
                    os.path.join(output_dir, '{}.cis_qtl_top_assoc.txt.gz'.format(prefix)),
                    sep='\t', float_format='%.6g', compression='gzip'
                )
    logger.write('done.')

def map_trans(genotype_df, phenotype_df, covariates_df, interaction_s=None,
              return_sparse=True, pval_threshold=1e-5, maf_threshold=0.05,
              return_r2=False, batch_size=20000, logger=None, verbose=True):
    """Run trans-QTL mapping from genotypes in memory"""
    if logger is None:
        logger = SimpleLogger(verbose=verbose)
    assert np.all(phenotype_df.columns==covariates_df.index)

    variant_ids = genotype_df.index.tolist()
    variant_dict = {i:j for i,j in enumerate(variant_ids)}
    n_variants = len(variant_ids)
    n_samples = phenotype_df.shape[1]

    logger.write('trans-QTL mapping')
    logger.write('  * {} samples'.format(n_samples))
    logger.write('  * {} phenotypes'.format(phenotype_df.shape[0]))
    logger.write('  * {} covariates'.format(covariates_df.shape[1]))
    logger.write('  * {} variants'.format(n_variants))
    if interaction_s is not None:
        logger.write('  * including interaction term')

    # with tf.device('/cpu:0'):
    ggt = genotypeio.GenotypeGeneratorTrans(genotype_df.values, batch_size=batch_size, dtype=np.float32)
    dataset_genotypes = tf.data.Dataset.from_generator(ggt.generate_data, output_types=tf.float32)
    dataset_genotypes = dataset_genotypes.prefetch(10)
    iterator = dataset_genotypes.make_one_shot_iterator()
    next_element = iterator.get_next()

    # index of VCF samples corresponding to phenotypes
    ix_t = get_sample_indexes(genotype_df.columns.tolist(), phenotype_df)
    next_element = tf.gather(next_element, ix_t, axis=1)

    # calculate correlation threshold for sparse output
    if return_sparse:
        dof = n_samples - 2 - covariates_df.shape[1]
        tstat_threshold = stats.t.ppf(pval_threshold/2, dof)
        r2_threshold = tstat_threshold**2 / (dof + tstat_threshold**2)
    else:
        r2_threshold = None
        tstat_threshold = None

    if interaction_s is None:
        genotypes, phenotypes, covariates = initialize_data(phenotype_df, covariates_df,
                                                            batch_size=batch_size, dtype=tf.float32)
        # with tf.device('/gpu:0'):
        x = calculate_association(genotypes, phenotypes, covariates, return_sparse=return_sparse,
                                  r2_threshold=r2_threshold, return_r2=return_r2)
    else:
        genotypes, phenotypes, covariates, interaction = initialize_data(phenotype_df, covariates_df,
                                                                         batch_size=batch_size, interaction_s=interaction_s)
        x = calculate_association(genotypes, phenotypes, covariates, interaction_t=interaction,
                                  return_sparse=return_sparse, r2_threshold=r2_threshold)

    init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())

    start_time = time.time()
    if verbose:
        print('  Mapping batches')
    with tf.Session() as sess:
        # run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
        # run_metadata = tf.RunMetadata()
        # writer = tf.summary.FileWriter('logs', sess.graph, session=sess)
        sess.run(init_op)
        pval_list = []
        maf_list = []
        r2_list = []
        for i in range(1, ggt.num_batches+1):
            if verbose:
                sys.stdout.write('\r  * processing batch {}/{}'.format(i, ggt.num_batches))
                sys.stdout.flush()

            g_iter = sess.run(next_element)
            # x: p_values, maf, {r2}
            p_ = sess.run(x, feed_dict={genotypes:g_iter})#, options=run_options, run_metadata=run_metadata)
            # writer.add_run_metadata(run_metadata, 'batch{}'.format(i))

            pval_list.append(p_[0])
            maf_list.append(p_[1])
            if return_r2:
                r2_list.append(p_[2])

        if return_sparse:
            pval = tf.sparse_concat(0, pval_list).eval()
        else:
            pval = tf.concat(pval_list, 0).eval()
        maf = tf.concat(maf_list, 0).eval()
        if return_r2:
            r2 = tf.concat(r2_list, 0).eval()
        if verbose:
            print()
        # writer.close()

    logger.write('    time elapsed: {:.2f} min'.format((time.time()-start_time)/60))

    if return_sparse:
        ix = pval.indices[:,0]<n_variants  # truncate last batch
        v = [variant_dict[i] for i in pval.indices[ix,0]]
        if phenotype_df.shape[0]>1:
            phenotype_ids = phenotype_df.index[pval.indices[ix,1]]
        else:
            phenotype_ids = phenotype_df.index.tolist()*len(pval.values)
        pval_df = pd.DataFrame(
            np.array([v, phenotype_ids, pval.values[ix], maf[ix]]).T,
            columns=['variant_id', 'phenotype_id', 'pval', 'maf']
        )
        pval_df['pval'] = pval_df['pval'].astype(np.float64)
        pval_df['maf'] = pval_df['maf'].astype(np.float32)
        if return_r2:
            pval_df['r2'] = r2[ix].astype(np.float32)
    else:
        # truncate last batch
        pval = pval[:n_variants]
        maf = maf[:n_variants]
        # add indices
        pval_df = pd.DataFrame(pval, index=variant_ids, columns=[i for i in phenotype_df.index])
        pval_df['maf'] = maf
        pval_df.index.name = 'variant_id'

    if maf_threshold is not None and maf_threshold>0:
        logger.write('  * filtering output by MAF >= {}'.format(maf_threshold))
        pval_df = pval_df[pval_df['maf']>=maf_threshold]

    logger.write('done.')
    return pval_df

def map_trans_permutations(genotype_input, phenotype_df, covariates_df, permutations=None,
                           split_chr=True, pval_df=None, nperms=10000, maf_threshold=0.05,
                           batch_size=20000, logger=None):
    """
    Warning: this function requires that all phenotypes are normally distributed,
             e.g., inverse normal transformed
    """
    if logger is None:
        logger = SimpleLogger()
    assert np.all(phenotype_df.columns==covariates_df.index)

    if split_chr:
        plink_reader = genotype_input
        # assert isinstance(plink_reader, genotypeio.PlinkReader)
        if pval_df is not None:
            assert 'phenotype_chr' in pval_df and 'pval' in pval_df and np.all(pval_df.index==phenotype_df.index)
        variant_ids = plink_reader.bim['snp'].tolist()
        # index of VCF samples corresponding to phenotypes
        ix_t = get_sample_indexes(plink_reader.fam['iid'].tolist(), phenotype_df)
    else:
        genotype_df = genotype_input
        assert isinstance(genotype_df, pd.DataFrame)
        variant_ids = genotype_df.index.tolist()
        # index of VCF samples corresponding to phenotypes
        ix_t = get_sample_indexes(genotype_df.columns.tolist(), phenotype_df)

    n_variants = len(variant_ids)
    n_samples = phenotype_df.shape[1]
    dof = phenotype_df.shape[1] - 2 - covariates_df.shape[1]

    logger.write('trans-QTL mapping (FDR)')
    logger.write('  * {} samples'.format(n_samples))
    logger.write('  * {} phenotypes'.format(phenotype_df.shape[0]))
    logger.write('  * {} covariates'.format(covariates_df.shape[1]))
    logger.write('  * {} variants'.format(n_variants))

    if permutations is None:  # generate permutations
        q = stats.norm.ppf(np.arange(1,n_samples+1)/(n_samples+1))
        qv = np.tile(q,[nperms,1])
        for i in np.arange(nperms):
            np.random.shuffle(qv[i,:])
    else:
        assert permutations.shape[1]==n_samples
        nperms = permutations.shape[0]
        qv = permutations
        logger.write('  * {} permutations'.format(nperms))
    permutations_t = tf.constant(qv, dtype=tf.float32)
    permutations_t = tf.reshape(permutations_t, shape=[-1, n_samples])

    genotypes_t, phenotypes_t, covariates_t = initialize_data(phenotype_df, covariates_df, batch_size=batch_size, dtype=tf.float32)
    max_r2_nominal_t, max_r2_permuted_t, idxmax_t = _calculate_max_r2(genotypes_t, phenotypes_t, permutations_t, covariates_t, maf_threshold=maf_threshold)
    init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())

    if split_chr:
        start_time = time.time()
        chr_max_r2_nominal = OrderedDict()
        chr_max_r2_empirical = OrderedDict()
        print('  Mapping batches')
        with tf.Session() as sess:
            sess.run(init_op)

            for chrom in plink_reader.bim['chrom'].unique():
                genotypes, variant_pos = plink_reader.get_region(chrom, verbose=True)
                ggt = genotypeio.GenotypeGeneratorTrans(genotypes, batch_size=batch_size, dtype=np.float32)
                dataset_genotypes = tf.data.Dataset.from_generator(ggt.generate_data, output_types=tf.float32)
                dataset_genotypes = dataset_genotypes.prefetch(10)
                iterator = dataset_genotypes.make_one_shot_iterator()
                next_element = iterator.get_next()
                next_element = tf.gather(next_element, ix_t, axis=1)

                nominal_list = []
                perms_list = []
                nominal_idx_list = []
                for i in range(1, ggt.num_batches+1):
                    sys.stdout.write('\r  * {}: processing batch {}/{}  '.format(chrom, i, ggt.num_batches))
                    sys.stdout.flush()

                    g_iter = sess.run(next_element)
                    max_r2_nominal, max_r2_permuted = sess.run([max_r2_nominal_t, max_r2_permuted_t], feed_dict={genotypes_t:g_iter})
                    nominal_list.append(max_r2_nominal)
                    perms_list.append(max_r2_permuted)

                chr_max_r2_nominal[chrom] = np.max(np.array(nominal_list), 0)
                chr_max_r2_empirical[chrom] = np.max(np.array(perms_list), 0)
        logger.write('    time elapsed: {:.2f} min'.format((time.time()-start_time)/60))

        chr_max_r2_nominal = pd.DataFrame(chr_max_r2_nominal, index=phenotype_df.index)
        chr_max_r2_empirical = pd.DataFrame(chr_max_r2_empirical)

        # compute leave-one-out max
        max_r2_nominal = OrderedDict()
        max_r2_empirical = OrderedDict()
        for c in chr_max_r2_nominal:
            max_r2_nominal[c] = chr_max_r2_nominal[np.setdiff1d(chr_max_r2_nominal.columns, c)].max(1)
            max_r2_empirical[c] = chr_max_r2_empirical[np.setdiff1d(chr_max_r2_empirical.columns, c)].max(1)
        max_r2_nominal = pd.DataFrame(max_r2_nominal, index=phenotype_df.index)
        max_r2_empirical = pd.DataFrame(max_r2_empirical)  # nperms x chrs

        if pval_df is not None:
            # nominal p-value (sanity check, matches pval_df['pval'])
            r2_nominal = max_r2_nominal.lookup(pval_df['phenotype_id'], pval_df['phenotype_chr'])
            tstat = np.sqrt( dof*r2_nominal / (1-r2_nominal) )
            minp_nominal = pd.Series(2*stats.t.cdf(-np.abs(tstat), dof), index=pval_df['phenotype_id'])

        # empirical p-values
        tstat = np.sqrt( dof*max_r2_empirical / (1-max_r2_empirical) )
        minp_empirical = pd.DataFrame(2*stats.t.cdf(-np.abs(tstat), dof), columns=tstat.columns)
        if pval_df is not None:
            pval_perm = np.array([(np.sum(minp_empirical[chrom]<=p)+1)/(nperms+1) for p, chrom in zip(pval_df['pval'], pval_df['phenotype_chr'])])

        beta_shape1 = {}
        beta_shape2 = {}
        true_dof = {}
        for c in max_r2_empirical:
            beta_shape1[c], beta_shape2[c], true_dof[c] = fit_beta_parameters(max_r2_empirical[c], dof)
        if pval_df is not None:
            beta_shape1 = [beta_shape1[c] for c in pval_df['phenotype_chr']]
            beta_shape2 = [beta_shape2[c] for c in pval_df['phenotype_chr']]
            true_dof = [true_dof[c] for c in pval_df['phenotype_chr']]
        else:
            chroms = plink_reader.bim['chrom'].unique()
            beta_shape1 = [beta_shape1[c] for c in chroms]
            beta_shape2 = [beta_shape2[c] for c in chroms]
            true_dof = [true_dof[c] for c in chroms]

        if pval_df is not None:
            pval_true_dof = pval_from_corr(r2_nominal, true_dof)
            pval_beta = stats.beta.cdf(pval_true_dof, beta_shape1, beta_shape2)
            variant_id = [np.NaN]*len(pval_beta)

    else:
        ggt = genotypeio.GenotypeGeneratorTrans(genotype_df.values, batch_size=batch_size, dtype=np.float32)
        dataset_genotypes = tf.data.Dataset.from_generator(ggt.generate_data, output_types=tf.float32)
        dataset_genotypes = dataset_genotypes.prefetch(10)
        iterator = dataset_genotypes.make_one_shot_iterator()
        next_element = iterator.get_next()
        next_element = tf.gather(next_element, ix_t, axis=1)

        start_time = time.time()
        max_r2_nominal = []
        max_r2_nominal_idx = []
        max_r2_empirical = []
        with tf.Session() as sess:
            sess.run(init_op)
            for i in range(ggt.num_batches):
                sys.stdout.write('\rProcessing batch {}/{}'.format(i+1, ggt.num_batches))
                sys.stdout.flush()
                g_iter = sess.run(next_element)
                res = sess.run([max_r2_nominal_t, max_r2_permuted_t, idxmax_t], feed_dict={genotypes_t:g_iter})
                max_r2_nominal.append(res[0])
                max_r2_nominal_idx.append(res[2] + i*batch_size)
                max_r2_empirical.append(res[1])
            print()
        max_r2_nominal = np.array(max_r2_nominal)
        max_r2_nominal_idx = np.array(max_r2_nominal_idx)
        max_r2_empirical = np.array(max_r2_empirical)
        logger.write('    time elapsed: {:.2f} min'.format((time.time()-start_time)/60))

        if len(max_r2_nominal_idx.shape)==1:
            max_r2_nominal = max_r2_nominal.reshape(-1,1)
            max_r2_nominal_idx = max_r2_nominal_idx.reshape(-1,1)

        idxmax = np.argmax(max_r2_nominal, 0)
        variant_ix = [max_r2_nominal_idx[i,k] for k,i in enumerate(idxmax)]
        variant_id = genotype_df.index[variant_ix]

        max_r2_nominal = np.max(max_r2_nominal, 0)
        tstat = np.sqrt( dof*max_r2_nominal / (1-max_r2_nominal) )
        minp_nominal = 2*stats.t.cdf(-np.abs(tstat), dof)

        max_r2_empirical = np.max(max_r2_empirical, 0)
        tstat = np.sqrt( dof*max_r2_empirical / (1-max_r2_empirical) )
        minp_empirical = 2*stats.t.cdf(-np.abs(tstat), dof)

        pval_perm = np.array([(np.sum(minp_empirical<=p)+1)/(nperms+1) for p in minp_nominal])
        # calculate beta p-values for each phenotype:
        beta_shape1, beta_shape2, true_dof, minp_vec = fit_beta_parameters(max_r2_empirical, dof, tol=1e-4, return_minp=True)
        pval_true_dof = pval_from_corr(max_r2_nominal, true_dof)
        pval_beta = stats.beta.cdf(pval_true_dof, beta_shape1, beta_shape2)

    if not split_chr or pval_df is not None:
        fit_df = pd.DataFrame(OrderedDict([
            ('beta_shape1', beta_shape1),
            ('beta_shape2', beta_shape2),
            ('true_df', true_dof),
            ('min_pval_true_df', pval_true_dof),
            ('variant_id', variant_id),
            ('min_pval_nominal', minp_nominal),
            ('pval_perm', pval_perm),
            ('pval_beta', pval_beta),
        ]), index=phenotype_df.index)
    else:
        fit_df = pd.DataFrame(OrderedDict([
            ('beta_shape1', beta_shape1),
            ('beta_shape2', beta_shape2),
            ('true_df', true_dof),
        ]), index=chroms)

    if split_chr:
        return fit_df, minp_empirical
    else:
        return fit_df, minp_vec

def map_trans_tfrecord(vcf_tfrecord, phenotype_df, covariates_df, interaction_s=None, return_sparse=True, pval_threshold=1e-5, maf_threshold=0.05, batch_size=50000, logger=None):
    """Run trans-QTL mapping from genotypes in tfrecord"""
    if logger is None:
        logger = SimpleLogger()
    assert np.all(phenotype_df.columns==covariates_df.index)

    with open(vcf_tfrecord+'.samples') as f:
        vcf_sample_ids = f.read().strip().split('\n')
    n_samples_vcf = len(vcf_sample_ids)

    with gzip.open(vcf_tfrecord+'.variants.gz', 'rt') as f:
        variant_ids = f.read().strip().split('\n')
    variant_dict = {i:j for i,j in enumerate(variant_ids)}
    n_variants = len(variant_ids)

    # index of VCF samples corresponding to phenotypes
    ix_t = get_sample_indexes(vcf_sample_ids, phenotype_df)
    n_samples = phenotype_df.shape[1]

    # batched_dataset = dataset.apply(tf.contrib.data.padded_batch(batch_size, padded_shapes=[[batch_size], [None]]))
    # batched_dataset = dataset.padded_batch(batch_size, padded_shapes=(batch_size,n_samples), padding_values=0)
    with tf.device('/cpu:0'):
        batched_dataset = genotypeio.make_batched_dataset(vcf_tfrecord, batch_size, ix_t=ix_t)

        iterator = batched_dataset.make_one_shot_iterator()
        next_element = iterator.get_next()
        next_element = genotypeio.pad_up_to(next_element, [batch_size, n_samples])  # not clear if right/best way to do this

    logger.write('  * {} samples'.format(n_samples))
    logger.write('  * {} phenotypes'.format(phenotype_df.shape[0]))
    logger.write('  * {} covariates'.format(covariates_df.shape[1]))
    logger.write('  * {} variants'.format(n_variants))
    if interaction_s is not None:
        logger.write('  * including interaction term')

    num_batches = int(np.ceil(np.true_divide(n_variants, batch_size)))

    # calculate correlation threshold for sparse output
    if return_sparse:
        dof = n_samples - 2 - covariates_df.shape[1]
        t = stats.t.ppf(pval_threshold/2, dof)**2 / dof
        r2_threshold = t / (1+t)
    else:
        r2_threshold = None

    if interaction_s is None:
        genotypes, phenotypes, covariates = initialize_data(phenotype_df, covariates_df, batch_size=batch_size)
        with tf.device('/gpu:0'):
            p_values, maf = calculate_association(genotypes, phenotypes, covariates, return_sparse=return_sparse, r2_threshold=r2_threshold)
    else:
        genotypes, phenotypes, covariates, interaction = initialize_data(phenotype_df, covariates_df, batch_size=batch_size, interaction_s=interaction_s)
        p_values, maf = calculate_association(genotypes, phenotypes, covariates, interaction_t=interaction, return_sparse=return_sparse, r2_threshold=r2_threshold)

    # g = _parse_function(next_element, batch_size, n_samples, ix_t)
    init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())

    start_time = time.time()
    with tf.Session() as sess:
        # run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
        # run_metadata = tf.RunMetadata()

        # writer = tf.summary.FileWriter('logs', sess.graph, session=sess)
        sess.run(init_op)
        pval_list = []
        maf_list = []
        for i in range(num_batches):
            sys.stdout.write('\rProcessing batch {}/{}'.format(i+1, num_batches))
            sys.stdout.flush()

            g_iter = sess.run(next_element)
            # g_iter = sess.run(g)
            p_ = sess.run([p_values, maf], feed_dict={genotypes:g_iter})#, options=run_options, run_metadata=run_metadata)
            # writer.add_run_metadata(run_metadata, 'batch%d' % i)

            pval_list.append(p_[0])
            maf_list.append(p_[1])

        if return_sparse:
            pval = tf.sparse_concat(0, pval_list).eval()
        else:
            pval = tf.concat(pval_list, 0).eval()
        maf = tf.concat(maf_list, 0).eval()
        print()
        # writer.close()

    logger.write('Time elapsed: {:.2f} min'.format((time.time()-start_time)/60))

    if return_sparse:
        ix = pval.indices[:,0]<n_variants  # truncate last batch
        v = [variant_dict[i] for i in pval.indices[ix,0]]
        pval_df = pd.DataFrame(
            np.array([v, phenotype_df.index[pval.indices[ix,1]], pval.values[ix], maf[ix]]).T,
            columns=['variant_id', 'phenotype_id', 'pval', 'maf']
        )
        pval_df['pval'] = pval_df['pval'].astype(np.float64)
        pval_df['maf'] = pval_df['maf'].astype(np.float32)
    else:
        # truncate last batch
        pval = pval[:n_variants]
        maf = maf[:n_variants]
        # add indices
        pval_df = pd.DataFrame(pval, index=variant_ids, columns=[i for i in phenotype_df.index])
        pval_df['maf'] = maf
        pval_df.index.name = 'variant_id'

    if maf_threshold is not None and maf_threshold>0:
        logger.write('  * filtering output by MAF >= {}'.format(maf_threshold))
        pval_df = pval_df[pval_df['maf']>=maf_threshold]

    return pval_df

def _in_cis(chrom, pos, gene_id, tss_dict, window=1000000):
    """Test if a variant-gene pair is in cis"""
    if chrom==tss_dict[gene_id]['chr']:
        tss = tss_dict[gene_id]['tss']
        if pos>=tss-window and pos<=tss+window:
            return True
        else:
            return False
    else:
        return False

def filter_cis(pval_df, tss_dict, window=1000000):
    """Filter out cis-QTLs"""
    drop_ix = []
    for k,gene_id,variant_id in zip(pval_df['phenotype_id'].index, pval_df['phenotype_id'], pval_df['variant_id']):
        chrom, pos = variant_id.split('_',2)[:2]
        pos = int(pos)
        if _in_cis(chrom, pos, gene_id, tss_dict, window=window):
            drop_ix.append(k)
    return pval_df.drop(drop_ix)

#------------------------------------------------------------------------------
#  Input parsers
#------------------------------------------------------------------------------
def read_phenotype_bed(phenotype_bed):
    """Load phenotype BED file as phenotype and TSS DataFrames"""
    if phenotype_bed.endswith('.bed.gz'):
        phenotype_df = pd.read_csv(phenotype_bed, sep='\t', index_col=3, dtype={'#chr':str, '#Chr':str})
    elif phenotype_bed.endswith('.parquet'):
        phenotype_df = pd.read_parquet(phenotype_bed)
    else:
        raise ValueError('Unsupported file type.')
    phenotype_df = phenotype_df.rename(columns={i:i.lower() for i in phenotype_df.columns[:3]})
    phenotype_pos_df = phenotype_df[['#chr', 'end']].rename(columns={'#chr':'chr', 'end':'tss'})
    phenotype_df = phenotype_df.drop(['#chr', 'start', 'end'], axis=1)
    return phenotype_df, phenotype_pos_df

def main():
    parser = argparse.ArgumentParser(description='tensorQTL: GPU-based QTL mapper')
    parser.add_argument('genotype_path', help='Genotypes/dosages in PLINK or tfrecord format')
    parser.add_argument('phenotype_bed', help='Phenotypes in BED format')
    parser.add_argument('prefix', help='Prefix for output file names')
    parser.add_argument('--mode', default='cis', choices=['cis', 'cis_nominal', 'cis_independent', 'trans'], help='Mapping mode. Default: cis')
    parser.add_argument('--covariates', default=None, help='Covariates file, tab-delimited, covariates x samples')
    parser.add_argument('--permutations', default=10000, help='Number of permutations. Default: 10000')
    parser.add_argument('--interaction', default=None, type=str, help='Interaction term')
    parser.add_argument('--cis_output', default=None, type=str, help="Output from 'cis' mode with q-values. Required for independent cis-QTL mapping.")
    parser.add_argument('--phenotype_groups', default=None, type=str, help='Phenotype groups. Header-less TSV with two columns: phenotype_id, group_id')
    parser.add_argument('--window', default=1000000, type=np.int32, help='Cis-window size, in bases. Default: 1000000.')
    parser.add_argument('--pval_threshold', default=None, type=np.float64, help='Output only significant phenotype-variant pairs with a p-value below threshold. Default: 1e-5 for trans-QTL')
    parser.add_argument('--maf_threshold', default=None, type=np.float64, help='Include only genotypes with minor allele frequency >=maf_threshold. Default: 0')
    parser.add_argument('--maf_threshold_interaction', default=0.05, type=np.float64, help='MAF threshold for interactions, applied to lower and upper half of samples')
    parser.add_argument('--return_dense', action='store_true', help='Return dense output for trans-QTL.')
    parser.add_argument('--return_r2', action='store_true', help='Return r2 (only for sparse trans-QTL output)')
    parser.add_argument('--best_only', action='store_true', help='Produce output only for the top association/phenotype')
    parser.add_argument('--output_text', action='store_true', help='Write output in txt.gz format instead of parquet (trans-QTL mode only)')
    parser.add_argument('--batch_size', type=int, default=50000, help='Batch size. Reduce this if encountering OOM errors.')
    parser.add_argument('--fdr', default=0.05, type=np.float64, help='FDR for cis-QTLs')
    parser.add_argument('--qvalue_lambda', default=None, type=np.float64, help='lambda parameter for pi0est in qvalue.')
    parser.add_argument('-o', '--output_dir', default='.', help='Output directory')
    args = parser.parse_args()

    # os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

    # check inputs
    if args.mode=='cis_independent' and (args.cis_output is None or not os.path.exists(args.cis_output)):
        raise ValueError("Output from 'cis' mode must be provided.")

    logger = SimpleLogger(os.path.join(args.output_dir, args.prefix+'.tensorQTL.{}.log'.format(args.mode)))
    logger.write('[{}] Running TensorQTL: {}-QTL mapping'.format(datetime.now().strftime("%b %d %H:%M:%S"), args.mode.split('_')[0]))

    # load inputs
    logger.write('  * reading phenotypes ({})'.format(args.phenotype_bed))
    phenotype_df, phenotype_pos_df = read_phenotype_bed(args.phenotype_bed)

    tss_dict = phenotype_pos_df.T.to_dict()
    if args.covariates is not None:
        logger.write('  * reading covariates ({})'.format(args.covariates))
        covariates_df = pd.read_csv(args.covariates, sep='\t', index_col=0).T
        assert np.all(phenotype_df.columns==covariates_df.index)
    if args.interaction is not None:
        logger.write('  * reading interaction term ({})'.format(args.interaction))
        interaction_s = pd.read_csv(args.interaction, sep='\t', index_col=0, header=None, squeeze=True)
        assert np.all(interaction_s.index==covariates_df.index)
    else:
        interaction_s = None

    if args.maf_threshold is None:
        if args.mode=='trans':
            maf_threshold = 0.05
        else:
            maf_threshold = 0
    else:
        maf_threshold = args.maf_threshold

    if args.phenotype_groups is not None:
        group_s = pd.read_csv(args.phenotype_groups, sep='\t', index_col=0, header=None, squeeze=True)
        # verify sort order
        group_dict = group_s.to_dict()
        previous_group = ''
        parsed_groups = 0
        for i in phenotype_df.index:
            if group_dict[i]!=previous_group:
                parsed_groups += 1
                previous_group = group_dict[i]
        if not parsed_groups == len(group_s.unique()):
            raise ValueError('Groups defined in input do not match phenotype file (check sort order).')
    else:
        group_s = None

    if args.mode.startswith('cis'):
        pr = genotypeio.PlinkReader(args.genotype_path, select_samples=phenotype_df.columns, dtype=np.int8)
        if args.mode=='cis':
            res_df = map_cis(pr, phenotype_df, phenotype_pos_df, covariates_df, group_s=group_s, nperm=args.permutations, logger=logger)
            logger.write('  * writing output')
            out_file = os.path.join(args.output_dir, args.prefix+'.cis_qtl.txt.gz')
            if has_rpy2 and group_s is None:
                calculate_qvalues(res_df, fdr=args.fdr, qvalue_lambda=args.qvalue_lambda)
            res_df.to_csv(out_file, sep='\t', float_format='%.6g', compression='gzip')
        elif args.mode=='cis_nominal':
            if interaction_s is None:
                map_cis_nominal(pr, phenotype_df, phenotype_pos_df, covariates_df, args.prefix,
                                output_dir=args.output_dir, logger=logger)
            else:
                map_cis_interaction_nominal(pr, phenotype_df, phenotype_pos_df, covariates_df, interaction_s,
                                args.prefix, maf_threshold_interaction=args.maf_threshold_interaction,
                                best_only=args.best_only, output_dir=args.output_dir, logger=logger)
        elif args.mode=='cis_independent':
            summary_df = pd.read_csv(args.cis_output, sep='\t', index_col=0)
            summary_df.rename(columns={'minor_allele_samples':'ma_samples', 'minor_allele_count':'ma_count'}, inplace=True)
            res_df = map_cis_independent(pr, summary_df, phenotype_df, phenotype_pos_df, covariates_df,
                                         fdr=args.fdr, nperm=args.permutations, logger=logger)
            logger.write('  * writing output')
            out_file = os.path.join(args.output_dir, args.prefix+'.cis_independent_qtl.txt.gz')
            res_df.to_csv(out_file, sep='\t', index=False, float_format='%.6g', compression='gzip')
    elif args.mode=='trans':
        return_sparse = not args.return_dense
        pval_threshold = args.pval_threshold
        if pval_threshold is None and return_sparse:
            pval_threshold = 1e-5
            logger.write('  * p-value threshold: {:.2g}'.format(pval_threshold))

        genotype_df = genotypeio.load_genotypes(args.genotype_path)
        pval_df = map_trans(genotype_df, phenotype_df, covariates_df, interaction_s=interaction_s,
                  return_sparse=return_sparse, pval_threshold=pval_threshold,
                  maf_threshold=maf_threshold, batch_size=args.batch_size,
                  return_r2=args.return_r2, logger=logger)

        logger.write('  * filtering out cis-QTLs (within +/-5Mb)')
        pval_df = filter_cis(pval_df, tss_dict, window=5000000)

        logger.write('  * writing output')
        if not args.output_text:
            pval_df.to_parquet(os.path.join(args.output_dir, args.prefix+'.trans_qtl_pairs.parquet'))
        else:
            out_file = os.path.join(args.output_dir, args.prefix+'.trans_qtl_pairs.txt.gz')
            pval_df.to_csv(out_file, sep='\t', index=False, float_format='%.6g', compression='gzip')

    logger.write('[{}] Finished mapping'.format(datetime.now().strftime("%b %d %H:%M:%S")))

if __name__ == '__main__':
    main()

reason: import numpy,import scipy,from scipy


text:
"""this is a test case."""
import time
import path_magic
from mesh import TransfiniteMesh
from function_space import FunctionSpace
from inner_product import MeshFunction, inner
import numpy as np
from scipy.sparse import coo_matrix
from forms import Form
from basis_forms import BasisForm
from coboundaries import d, d_21_lobatto_outer
from assemble import assemble
from scipy import sparse
import matplotlib.pyplot as plt
from matplotlib.pyplot import plot, draw, show

start_time = time.time()

dim = 2
elements_layout = (1, 1)

p = (5, 5)
p2 = (4, 4)

primal_is_inner = False

def gamma1(s):
    """Parameterized BOTTOM boundary curve of the domain."""
    return s, 0 * np.ones(np.shape(s))

def gamma2(t):
    """Parameterized RIGHT boundary curve of the domain."""
    return 1 * np.ones(np.shape(t)), t

def gamma3(s):
    """Parameterized TOP boundary curve of the domain."""
    return s, 1 * np.ones(np.shape(s))

def gamma4(t):
    """Parameterized LEFT boundary curve of the domain."""
    return 0 * np.ones(np.shape(t)), t

def dgamma1(s):
    """Differential of gamma1 wrt to xi."""
    return 1 * np.ones(np.shape(s)), 0 * np.ones(np.shape(s))

def dgamma2(t):
    """Differential of gamma2 wrt to eta."""
    return 0 * np.ones(np.shape(t)), 1 * np.ones(np.shape(t))

def dgamma3(s):
    """Differential of gamma3 wrt to xi."""
    return 1 * np.ones(np.shape(s)), 0 * np.ones(np.shape(s))

def dgamma4(t):
    """Differential of gamma4 wrt to eta."""
    return 0 * np.ones(np.shape(t)), 1 * np.ones(np.shape(t))

def source(x, y):
    """Define the source term."""
    return -8 * np.pi**2 * np.sin(2 * np.pi * x) * np.sin(2 * np.pi * y)

def k_11():
    """Define the k11 component of permeability tensor."""
    return np.ones((1, 1))

def k_12():
    """Define the k12 component of permeability tensor."""
    return np.zeros((1, 1))

def k_22():
    """Define the k22 component of permeability tensor."""
    return k_11()

def manufactured_solution(x, y):
    # return np.sin(np.pi * x / 2)
    # return 1 + x
    return np.sin(2 * np.pi * x) * np.sin(2 * np.pi * y)

gamma = (gamma1, gamma2, gamma3, gamma4)
dgamma = (dgamma1, dgamma2, dgamma3, dgamma4)

ref_mesh = TransfiniteMesh(dim, elements_layout, gamma, dgamma)

func_space_1_lobatto = FunctionSpace(ref_mesh, '1-lobatto', p, primal_is_inner)
func_space_1_lobatto.dof_map.continous_dof = False

func_space_0_ext_gauss = FunctionSpace(ref_mesh, '0-ext_gauss', p2, primal_is_inner)

func_space_2_lobatto = FunctionSpace(ref_mesh, '2-lobatto', p, primal_is_inner)

anisotropic_tensor = MeshFunction(ref_mesh)
anisotropic_tensor.discrete_tensor = [k_11(), k_12(), k_22()]

source_form = Form(func_space_2_lobatto)
source_form.discretize(source)

phi_0_exact = Form(func_space_0_ext_gauss)
phi_0_exact.discretize(manufactured_solution)

print(np.shape(phi_0_exact.cochain))
print(phi_0_exact.cochain)
# define basis forms from function space

basis_0 = BasisForm(func_space_0_ext_gauss)
basis_1 = BasisForm(func_space_1_lobatto)
basis_2 = BasisForm(func_space_2_lobatto)

basis_0.quad_grid = 'gauss'
basis_1.quad_grid = 'lobatto'
basis_2.quad_grid = 'lobatto'

# define forms for soultion using function space

q_1 = Form(func_space_1_lobatto)
# q_1.basis.quad_grid = 'lobatto'
phi_0 = Form(func_space_0_ext_gauss)
# phi_0.basis.quad_grid = 'gauss'

E21 = d_21_lobatto_outer(p)

"""extended E21 with virtual volumes"""

px, py = p

virtual_E21 = np.zeros((px * 4, 2 * px * (px + 1)))

for j in range(px):
    vol_id = px**2 + j
    edge_id = px * (px + 1) + j
    virtual_E21[vol_id - px**2, edge_id] = 1
    virtual_E21[vol_id - px**2 + px, edge_id + px**2] = -1

for i in range(px):
    vol_id = px**2 + 2 * px + i
    edge_id = (px + 1) * i
    virtual_E21[vol_id - px**2, edge_id] = 1
    virtual_E21[vol_id - px**2 + px, edge_id + px] = -1

# define the mass and wedge matrices

M_1 = inner(basis_1, basis_1)
M1_assembled = assemble(M_1, (func_space_1_lobatto, func_space_1_lobatto))

W_i_2 = basis_2.wedged(basis_0)
Wi_2_E_21 = np.dot(W_i_2, E21)

lhs = sparse.bmat([[M1_assembled, Wi_2_E_21.transpose(), virtual_E21.transpose()], [
                  Wi_2_E_21, None, None], [virtual_E21, None, None]]).tolil()

rhs_1 = np.zeros(2 * px * (px + 1))
rhs_2 = np.dot(W_i_2, source_form.cochain)
rhs_3 = np.zeros(4 * px)
rhs = np.hstack((rhs_1, rhs_2, rhs_3))

# implementing BC's

nr_ghost_points = px * 4
dof_ghost_points_boundary = range(2 * px * (px + 1) + px**2, 2 * px * (px + 1) + px**2 + px * 4)

# changing lhs

lhs[-nr_ghost_points:, :] = np.zeros((nr_ghost_points, np.shape(lhs)[1]))
lhs[dof_ghost_points_boundary, dof_ghost_points_boundary] = 1
print(np.linalg.det(lhs.todense()))

# changing rhs

solution = sparse.linalg.spsolve(lhs.tocsc(), rhs)

end_time = time.time()

print("The total time taken by the program is : ", end_time - start_time)

q_1.cochain = solution[:func_space_1_lobatto.num_dof]
phi_0_internal = solution[func_space_1_lobatto.num_dof:func_space_1_lobatto.num_dof +
                          func_space_0_ext_gauss.num_internal_local_dof]

# plot sparsity pattern
plt.figure(1)
plt.spy(lhs)
plt.show(block=False)

xi = eta = np.linspace(-1, 1, 30)

q_1.reconstruct(xi, eta)
(x, y), u_x, u_y = q_1.export_to_plot()
q_x = u_y
q_y = -u_x
print(np.max(q_x), np.min(q_x))

plt.figure(2)
plt.contourf(x, y, q_x)
plt.colorbar()
plt.show(block=False)

q_x_exact = 2 * np.pi * np.cos(2 * np.pi * x) * np.sin(2 * np.pi * y)
print(np.max(q_x_exact), np.min(q_x_exact))

plt.figure(3)
plt.contourf(x, y, q_x_exact)

plt.colorbar()
plt.show(block=False)

# phi_0.reconstruct(xi, eta)
# (x1, y1), phi_plot = phi_0.export_to_plot()

# plt.figure(4)
# plt.contourf(x1, y1, phi_plot)

# plt.show()

ext: py

reason: import numpy,from scipy


text:
import os
import cv2
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import interactive
import scipy

def rgb_min_image(image):
    # extractes the min of the rgb values and outputs
    # a gray scale image
    rgb_image = np.amin(image, axis= 2)
    return rgb_image

def min_filter(image):
    # perfroms the min filter on 15 by 15 area
    for k in range (3):
        i_image = image.copy()
        temp_image = image[:,:,k].copy()
        [row,col] = temp_image.shape
        temp_image = cv2.copyMakeBorder(temp_image, 14, 14, 14, 14, cv2.BORDER_REFLECT)
        for i in range(row):
            for j in range(col):
                i_image[i,j,k] = (temp_image[i:14+i,j:14+j]).min()
    return i_image

def dark_channel(image):
    # output the dark channel as the image
    new_image = image.copy()
    min_image = min_filter(new_image)
    dark_prior = rgb_min_image(min_image)
    return dark_prior

def transmition_map(image,A,w):
    #finds the transmition map for the image
    image_new =  np.divide(image,A).astype(float)
    new_dark = dark_channel(image_new)
    transmition = 1 - w*new_dark
    return transmition

def A_estimator(image,dark_prior):
    #Used the information extracted from the dark prior
    #find a value for A
    image_copy = image.copy()
    [row,col,dem] = image_copy.shape
    dark_copy = dark_prior.copy()
    num = np.round(row*col*0.001).astype(int)
    j = sorted(np.asarray(dark_copy).reshape(-1), reverse=True)[:num]
    ind = np.unravel_index(j[0], dark_copy.shape)
    max_val = image_copy[ind[0],ind[1],:]
    for element in j:
        ind = np.unravel_index(element, dark_copy.shape)
        if (sum(max_val[:]) < sum(image_copy[ind[0],ind[1],:])):
            max_val[:] = image_copy[ind[0],ind[1],:]
    A = image_copy
    A[:,:,:] = max_val[:]
    return A

def Radience_cal(image,A,Transmission_map,t_not):
    #Used information from the transmit map to remove haze from the image.
    image_copy = image.copy()
    A_copy = A.copy()
    Transmission_map_copy = (Transmission_map.copy()).astype(float)
    divisor = np.maximum(Transmission_map_copy,t_not)
    radience = (image.copy()).astype(float)
    for i in range(3):
        radience[:,:,i] = np.divide(((image_copy[:,:,i]).astype(float) - A[0,0,i]),divisor) + A[0,0,i]
    #radience = 255*(radience/np.max(radience))
    radience[radience>255]=255
    radience[radience<0]=0
    return radience.astype('uint8')

def L_calculator(image,Transmission_map):
    #helps fine tune the transmition map for a better result
    epsalon = 10**(-8)
    h,w = image.shape[:2]
    window_area = (2*r + 1)**2
    n_vals = (w - 2*r)*(h - 2*r)*window_area**2
    k = 0
    # data for matting laplacian in coordinate form
    i = np.empty(n_vals, dtype=np.int32)
    j = np.empty(n_vals, dtype=np.int32)
    v = np.empty(n_vals, dtype=np.float64)

    # for each pixel of image
    for y in range(r, h - r):
        for x in range(r, w - r):

            # gather neighbors of current pixel in 3x3 window
            n = image[y-r:y+r+1, x-r:x+r+1]
            u = np.zeros(3)
            for p in range(3):
                u[p] = n[:, :, p].mean()
            c = n - u

            # calculate covariance matrix over color channels
            cov = np.zeros((3, 3))
            for p in range(3):
                for q in range(3):
                    cov[p, q] = np.mean(c[:, :, p]*c[:, :, q])

            # calculate inverse covariance of window
            inv_cov = np.linalg.inv(cov + epsilon/window_area * np.eye(3))

            # for each pair ((xi, yi), (xj, yj)) in a 3x3 window
            for dyi in range(2*r + 1):
                for dxi in range(2*r + 1):
                    for dyj in range(2*r + 1):
                        for dxj in range(2*r + 1):
                            i[k] = (x + dxi - r) + (y + dyi - r)*w
                            j[k] = (x + dxj - r) + (y + dyj - r)*w
                            temp = c[dyi, dxi].dot(inv_cov).dot(c[dyj, dxj])
                            v[k] = (1.0 if (i[k] == j[k]) else 0.0) - (1 + temp)/window_area
                            k += 1
    h,w = Transmission_map.shape
    L = scipy.sparse.csr_matrix((v, (i, j)), shape=(w*h, w*h))
    return L

def soft_matting(L,image,t_map):
    image_copy = image.copy()
    lamda = 10**(-4)
    U = np.identity(L.shape[0])
    t_map_mat = t_map*(L+lamda*U)/lamda
    return t_map_mat

def guided_filter(image,guide,diameter,epsilon):
    w_size = diameter+1
    # Exatrcation the mean of the image by blurring
    meanI=cv2.blur(image,(w_size,w_size))
    mean_Guide=cv2.blur(guide,(w_size,w_size))
    # Extracting the auto correlation
    II=image**2
    corrI=cv2.blur(II,(w_size,w_size))
    # Finding the correlation between image and guide
    I_guide=image*guide
    corrIG=cv2.blur(I_guide,(w_size,w_size))
    # using the mean of the image to find the variance of each point
    varI=corrI-meanI**2
    covIG=corrIG-meanI*mean_Guide
    #covIG normalized with a epsilon factor
    a=covIG/(varI+epsilon)
    #a is used to find the b
    b=mean_Guide-a*meanI
    meanA=cv2.blur(a,(w_size,w_size))
    meanB=cv2.blur(b,(w_size,w_size))
    transmission_rate=meanA*image+meanB
    # normalizaing of the transimational map
    transmission_rate = transmission_rate/np.max(transmission_rate)
    return transmission_rate

#---------------------------------------------------------------------------------------
#---------------------------------------------------------------------------------------

# inputing the information
base_path  =  os.getcwd()
test_Haze = os.listdir(base_path+ '/data_set/Training_Set/hazy')
test_GT = os.listdir(base_path+ '/data_set/Training_Set/GT')
image = cv2.imread( base_path + "/data_set/Test_Set/Bridge.jpg",cv2.IMREAD_COLOR)
# extracting the minmum value from 15 by 15 patch
min_image = min_filter(image)

# perfroming the minmin with 15by15 min filter
dark_prior = rgb_min_image(min_image)
# displaying the results
fig, axes= plt.subplots(nrows=1, ncols=3,figsize=(20,5))
plt.suptitle('Stages of Dark channel')
axes[0].imshow(cv2.cvtColor(image,cv2.COLOR_BGR2RGB))
axes[0].set_title('original image')
axes[1].imshow(cv2.cvtColor(min_image,cv2.COLOR_BGR2RGB))
axes[1].set_title('The min 15 patch image',)
axes[2].imshow(dark_prior,cmap='gray')
axes[2].set_title('The dark prior')
interactive(True)
plt.show()

A = A_estimator(image,dark_prior)
fig, axes= plt.subplots(nrows=1, ncols=3,figsize=(20,5))
axes[0].imshow(cv2.cvtColor(image,cv2.COLOR_BGR2RGB))
axes[0].set_title('original image')
axes[1].imshow(A,cmap='gray')
axes[1].set_title('The Ambiance image')
Transmition_image = transmition_map(image,A,0.95)
axes[2].imshow(Transmition_image,cmap='gray')
axes[2].set_title('The transmitance image')

plt.show()

fig, axes= plt.subplots(nrows=1, ncols=3,figsize=(20,5))
radience_image = Radience_cal(image,A,Transmition_image,0.1)
axes[0].imshow(cv2.cvtColor(image,cv2.COLOR_BGR2RGB))
axes[0].set_title('original image')
axes[1].imshow(Transmition_image,cmap='gray')
axes[1].set_title('The transmitance image')
axes[2].imshow(cv2.cvtColor(radience_image,cv2.COLOR_BGR2RGB))
axes[2].set_title('Haze Free image')

plt.show()

epsilon = 10**-8
img_gray=cv2.cvtColor(image,cv2.COLOR_BGR2GRAY)
# refined the transmition map using the guide filter
refine_Transmission_image=guided_filter(img_gray.astype(np.float32),Transmition_image.astype(np.float32),100,epsilon)
refine_radience_image = Radience_cal(image,A,refine_Transmission_image,0.1)
# diplaying the refined results
fig, axes= plt.subplots(nrows=1, ncols=3,figsize=(20,5))
radience_image = Radience_cal(image,A,Transmition_image,0.1)
axes[0].imshow(cv2.cvtColor(image,cv2.COLOR_BGR2RGB))
axes[0].set_title('original image')
axes[1].imshow(refine_Transmission_image,cmap='gray')
axes[1].set_title('The Refine Transmitance image')
axes[2].imshow(cv2.cvtColor(refine_radience_image,cv2.COLOR_BGR2RGB))
axes[2].set_title('Haze Free image')
interactive(False)
plt.show()

reason: import numpy,import scipy


text:
import numpy as np
from npstructures import RaggedArray, RaggedView, RaggedShape, npdataclass
from .encodings import BaseEncoding, QualityEncoding
from .datatypes import SequenceEntry, SequenceEntryWithQuality
from .sequences import Sequences

NEWLINE = 10

class FileBuffer:
    _buffer_divisor = 1
    COMMENT = 0

    def __init__(self, data, new_lines):
        self._data = np.asanyarray(data)
        self._new_lines = np.asanyarray(new_lines)
        self._is_validated = False
        self.size = self._data.size

    @classmethod
    def from_raw_buffer(cls, raw_buffer) -> "FileBuffer":
        """Create a buffer with full entries

        A raw buffer can end with data that does not represent full entries.
        This method extracts all the full entries, so that the next buffer can
        start from the last incomplete entry.

        Parameters
        ----------
        chunk : np.ndarray
            Raw buffer with data that might end with incomplete entry

        Returns
        -------
        'FileBuffer'
            Buffer with complete entries

        Examples
        --------

        """

        return NotImplemented

    @classmethod
    def from_data(cls, data: npdataclass) -> "FileBuffer":
        """Create FileBuffer from a data set

        Create a FileBuffer that can be written to file

        Parameters
        ----------
        data : npdataclass
            Data set containing the data to be written

        Returns
        -------
        'FileBuffer'
            FileBuffer containing the data
        """
        return NotImplemented

    def validate_if_not(self):
        if not self._is_validated:
            self._validate()

    def get_data(self) -> npdataclass:
        """Extract the data from the buffer

        The default way to extract data from the the buffer

        Returns
        -------
        npdataclass
            Data set containing the data from the buffer
        """
        return NotImplemented

    def _move_intervals_to_2d_array(self, starts, ends, fill_value=0):
        n_intervals = starts.size
        n_chars = ends - starts
        from_indices, _ = RaggedView(starts, n_chars).get_flat_indices()
        max_chars = np.max(n_chars)
        array = np.full(n_intervals * max_chars, fill_value, dtype=np.uint8)
        to_indices, _ = RaggedView(
            max_chars * np.arange(1, n_intervals + 1) - n_chars, n_chars
        ).get_flat_indices()
        array[to_indices] = self._data[from_indices]
        return array.reshape((n_intervals, max_chars))

    def _move_intervals_to_ragged_array(self, starts, ends=None, lens=None):
        if lens is None:
            lens = ends - starts
        indices, shape = RaggedView(starts, lens).get_flat_indices()
        return Sequences(self._data[indices], shape)

class OneLineBuffer(FileBuffer):
    n_lines_per_entry = 2
    _buffer_divisor = 32

    @classmethod
    def from_raw_buffer(cls, chunk) -> "OneLineBuffer":
        """Create a buffer with full entries

        Extract complete entries, i. e. a number of lines that is divisible by lines per entry

        Parameters
        ----------
        chunk : np.ndarray
            Raw buffer with data that might end with incomplete entry

        Returns
        -------
        'OneLineBuffer'
            Buffer with complete entries

        Examples
        --------
        8

        """
        new_lines = np.flatnonzero(chunk == NEWLINE)
        n_lines = new_lines.size
        assert n_lines >= cls.n_lines_per_entry, "No complete entry in buffer"
        new_lines = new_lines[: n_lines - (n_lines % cls.n_lines_per_entry)]
        return cls(chunk[: new_lines[-1] + 1], new_lines)

    def get_sequences(self) -> Sequences:
        self.validate_if_not()
        sequence_starts = self._new_lines[:: self.n_lines_per_entry] + 1
        sequence_lens = self._new_lines[1 :: self.n_lines_per_entry] - sequence_starts
        indices, shape = RaggedView(sequence_starts, sequence_lens).get_flat_indices()
        m = indices.size
        d = m % self._buffer_divisor
        seq = np.empty(m - d + self._buffer_divisor, dtype=self._data.dtype)
        seq[:m] = self._data[indices]
        return Sequences(seq, shape)

    def get_data(self):
        self.validate_if_not()
        starts = np.insert(self._new_lines, 0, -1)
        lengths = np.diff(starts)
        self.lines = Sequences(self._data, RaggedShape(lengths))
        sequences = self.lines[1 :: self.n_lines_per_entry, :-1]
        headers = self.lines[:: self.n_lines_per_entry, 1:-1]
        return SequenceEntry(headers, sequences)

    @classmethod
    def from_data(cls, entries):
        name_lengths = entries.name.shape.lengths
        sequence_lengths = entries.sequence.shape.lengths
        line_lengths = np.hstack(
            (name_lengths[:, None] + 2, sequence_lengths[:, None] + 1)
        ).ravel()
        buf = np.empty(line_lengths.sum(), dtype=np.uint8)
        lines = RaggedArray(buf, line_lengths)
        step = cls.n_lines_per_entry
        lines[0::step, 1:-1] = entries.name
        lines[1::step, :-1] = entries.sequence

        lines[0::step, 0] = ord(">")

        lines[:, -1] = ord("\n")
        return buf

    def _validate(self):
        n_lines = self._new_lines.size
        assert n_lines % self.n_lines_per_entry == 0, "Wrong number of lines in buffer"
        header_idxs = (
            self._new_lines[self.n_lines_per_entry - 1 : -1 : self.n_lines_per_entry]
            + 1
        )
        assert np.all(self._data[header_idxs] == self.HEADER)
        self._is_validated = True

class TwoLineFastaBuffer(OneLineBuffer):
    HEADER = 62
    n_lines_per_entry = 2
    _encoding = BaseEncoding

class FastQBuffer(OneLineBuffer):
    HEADER = 64
    n_lines_per_entry = 4
    _encoding = BaseEncoding
    dataclass = SequenceEntryWithQuality

    def get_data(self):
        seq_entry = super().get_data()
        quality = QualityEncoding.encode(
            self.lines[3 :: self.n_lines_per_entry, :-1]
        )
        return SequenceEntryWithQuality(seq_entry.name, seq_entry.sequence, quality)

    @classmethod
    def _get_line_lens(cls, entries):
        name_lengths = entries.name.shape.lengths[:, None]
        sequence_lengths = entries.sequence.shape.lengths[:, None]
        return (
            np.hstack(
                (
                    name_lengths + 1,
                    sequence_lengths,
                    np.ones_like(sequence_lengths),
                    sequence_lengths,
                )
            ).ravel()
            + 1
        )

    @classmethod
    def from_data(cls, entries):
        line_lengths = cls._get_line_lens(entries)
        buf = np.empty(line_lengths.sum(), dtype=np.uint8)
        lines = RaggedArray(buf, line_lengths)
        step = cls.n_lines_per_entry
        lines[0::step, 1:-1] = entries.name
        lines[1::step, :-1] = entries.sequence
        lines[2::step, 0] = ord("+")
        lines[3::step, :-1] = QualityEncoding.decode(entries.quality)
        lines[0::step, 0] = cls.HEADER
        lines[:, -1] = ord("\n")

        return buf

max_stars_repo_name: knutdrand/bionumpy

max_issues_repo_name: knutdrand/bionumpy

max_forks_repo_name: knutdrand/bionumpy

reason: import numpy


text:
'''
Export script to create tabular dataset from the metacorps web app's
MongoDB database. A mongodump of this database is available at
http://metacorps.io/static/data/nov-15-2017-metacorps-dump.zip (594M)
'''
import numpy as np
import pandas as pd

from nltk.tokenize import RegexpTokenizer
from pymongo import MongoClient

CLIENT = MongoClient()

def make_csv(project_name='Viomet Sep-Nov 2012',
             output_path='viomet-2012.csv',
             random_seed=42):
    '''
    Keep random_seed=42 to reproduce results in paper.
    '''

    project = CLIENT.metacorps.project.find(
        {'name': 'Viomet Sep-Nov 2012'}
    ).next()

    facets = project['facets']

    facet_docs = [
        CLIENT.metacorps.facet.find({'_id': id_}).next()
        # the first three facets are hit, attack, and beat
        for id_ in facets[:3]
    ]

    tokenizer = RegexpTokenizer(r'\w+')

    def _preprocess(text):
        text = text.lower()
        text = ' '.join(tokenizer.tokenize(text))
        return _replacements(text)

    text_metaphor_rows = [
        (
            doc['word'],
            instance['reference_url'],
            _preprocess(instance['text']),
            int(instance['figurative'])
        )
        for doc in facet_docs
        for instance in doc['instances']
    ]

    # create tabular format of data
    data = pd.DataFrame(
        data=text_metaphor_rows,
        columns=['word', 'reference_url', 'text', 'is_metaphor']
    )

    data.to_csv(output_path, index=False)

    return data

    # XXX removing XXX
    # now sample with replacement to get a balanced dataset
    # n_metaphor = data['is_metaphor'].sum()
    # n_rows = len(data)
    # n_to_sample = n_rows - (2 * n_metaphor)
    # print(
    #     (
    #         '{} out of {} are metaphor. '
    #         'Adding {} more rows by sampling with replacement for a total '
    #         'number of {} rows.'
    #     ).format(n_metaphor, n_rows, n_to_sample, n_to_sample + n_rows)
    # )

    # metaphor_rows = data[data.is_metaphor == 1]

    # np.random.seed()
    # indexes_to_sample = np.random.choice(range(n_metaphor), n_to_sample)

    # augmentation = metaphor_rows.iloc[indexes_to_sample]

    # data = data.append(augmentation, ignore_index=True)

    # data.to_csv(output_path, index=False)

    # return data

def _replacements(text):
    text = text.replace('i m ', 'im ')
    text = text.replace('dont t ', 'dont ')
    text = text.replace('u s ', 'U.S. ')
    text = text.replace('it s ', 'it\'s ')
    text = text.replace('we re ', 'we\'re ')
    text = text.replace('they re ', 'they\'re ')
    text = text.replace('can t ', 'can\'t ')
    text = text.replace('palins', 'palin')
    return text

max_stars_repo_name: mt-digital/metacorps-nn

max_issues_repo_name: mt-digital/metacorps-nn

max_forks_repo_name: mt-digital/metacorps-nn


text:
import numpy as np
from numpy import sin
from timeit import default_timer as timer
from numba import vectorize, cuda, types
import math
## from pyculib import rand FIGURE OUT HOW TO GET PYCULIB WORKING

@vectorize(["float32(float32, float32)"], target='cuda')
def add(a, b):
    return a + b

@vectorize(["float32(float32, float32)"], target='cuda')
def testfunc(a, b):
    return (a*b)**(a/100)

@vectorize(["float32(float32, float32)", "float32(float32, float32)"], nopython=True)
def sintest(a, b):
    return (sin(a))+(sin(b))

@vectorize(['f8(f8, f8)', 'u1(u1, u1)'])
def sinc(a, b):
    return math.sin(a) + math.sin(b)

def sintest2(a, b):
    return math.sin(a) + math.sin(b)

N = 10000000

def main(function):
    T = np.ones(N, dtype=np.uint8)
    A = cuda.device_array_like(T)
    B = cuda.device_array_like(T)
    A3 = np.random.poisson(lam=1, size=N)
    B3 = np.random.poisson(lam=1, size=N)
    X = cuda.device_array(N, dtype=np.float16)
    start = timer()
    X = function(A3, B3)
    print(X)
    add_time1 = timer() - start
    free = cuda.current_context().get_memory_info()[0]/1e9
    total = cuda.current_context().get_memory_info()[1]/1e9
    print('----------------------------------------')
    print(" CUDA Addition took {} seconds".format(add_time1))
    print(" Free memory: {} GB, Total memory: {} GB".format(free, total))
    print('----------------------------------------\n')
    return add_time1

def testfunc2(function):
    A2 = np.ones(N, dtype=np.uint8)
    B2 = np.ones(N, dtype=np.uint8)
    X2 = np.zeros(N, dtype=np.float32)
    A4 = np.random.poisson(lam=1, size=N)
    B4 = np.random.poisson(lam=1, size=N)
    start = timer()
    for i in range(0, N):
        X2[i] = function(A4[i], B4[i])
    add_time2 = timer() - start
    print(X2)
    print('----------------------------------------')
    print(" Python Addition took {} seconds ".format(add_time2))
    print('----------------------------------------\n')
    return add_time2

t1 = main(sinc)
t2 = testfunc2(sintest2)

print('Speed ratio: {}'.format(t2/t1))

max_stars_repo_name: JCCPort/Tests

max_issues_repo_name: JCCPort/Tests

max_forks_repo_name: JCCPort/Tests

reason: import numpy,from numpy,from numba