def pool_threading(nthreads=None):
    if nthreads is None:
        nthreads = omp_num_threads()
    try:
        import mkl
        old_mkl_num_threads = mkl.get_max_threads()
        mkl.set_num_threads(1)
    except ImportError:
        pass
    old_omp_num_threads = os.getenv('OMP_NUM_THREADS')
    os.environ['OMP_NUM_THREADS'] = '1'

    pool = multiprocessing.dummy.Pool(nthreads)
    yield pool

    pool.close()
    pool.join()
    try:
        mkl.set_num_threads(old_mkl_num_threads)
    except NameError:
        pass
    if old_omp_num_threads is not None:
        os.environ['OMP_NUM_THREADS'] = old_omp_num_threads
    else:
        del os.environ['OMP_NUM_THREADS']

def handle_mkl(max_threads):
    """Set max threads if mkl is availavle"""
    try:
        import mkl
        mkl.set_num_threads(max_threads)
    except ImportError:
        pass

def run(expt, display=False):
    if isinstance(expt, str):
        expt = get_experiment(expt)

    storage.ensure_directory(expt.figures_dir())
    storage.ensure_directory(expt.outputs_dir())

    mkl.set_num_threads(1)
        
    v = gnp.garray(expt.dataset.load().as_matrix())
    v = 0.999 * v + 0.001 * 0.5

    if expt.permute:
        idxs = np.random.permutation(v.shape[0])
        v = v[idxs]

    if display:
        expt.diagnostics += ['objective']
        expt.show_after = expt.save_after
    visuals = Visuals(expt, v)


    if expt.init_rbm == 'base_rates':
        init_rbm = None
    elif isinstance(expt.init_rbm, TrainedRBM):
        init_rbm = load_trained_rbm(expt.init_rbm.location).convert_to_garrays()
    else:
        raise RuntimeError('Unknown init_rbm')

    assert isinstance(expt.training, rbm_training.TrainingParams)
    rbm_training.train_rbm(v, expt.nhid, expt.training, after_step=visuals.after_step, init_rbm=init_rbm)

def configure(num_jobs=8, TEST=False, subtract=0, num_proc=None, num_thread_per_proc=None):
    '''
    num_jobs is typically the # of genes we are parallelizing over
    '''
    if num_proc is None:
        num_proc = multiprocessing.cpu_count() - subtract

    if num_jobs > num_proc:
        num_jobs = num_proc

    if num_thread_per_proc is None:
        num_thread_per_proc = int(np.floor(num_proc/num_jobs))

    if TEST:
        num_jobs = 1
        num_thread_per_proc = 1

    try:
        import mkl
        mkl.set_num_threads(num_thread_per_proc)    
    except ImportError:
        print "MKL not available, so I'm not adjusting the number of threads"

    print "Launching %d jobs with %d MKL threads each" % (num_jobs, num_thread_per_proc)

    return num_jobs

def run_all(nsubs, nrois, nthreads=2):
    import mkl
    mkl.set_num_threads(nthreads)
    
    print "setup"
    cmats, design, grp = gen_data(nsubs=nsubs, nrois=nrois, nvoxs=nrois)
    
    print "degree"
    dall = time_fun(local_degree, cmats, design)
    
    print "glm"
    gall = time_fun(local_glm, cmats, design)
    
    print "mdmr"
    mall = time_fun(local_mdmr, cmats, design)
    
    print "svm"
    sall = time_fun(local_svm, cmats, grp)
    
    print "kmeans"
    kall = time_fun(local_kmeans, cmats, grp)
    
    print "end"
    times = np.vstack((dall, gall, mall, sall, kall))
    
    return times

def silhouette_original_clusterings(dataset='CB1', neuropil='Antennal_lobe', clusterer_or_k=60):
    """Returns a pandas dataframe with the silhouette index of each cluster member.
    The dataframe have columns (cluster_id, member_id, silhouette).
    """

    # Read the expression matrix
    print('Reading expression matrix')
    Xdf = ExpressionDataset.dataset(dset=dataset, neuropil=neuropil).Xdf(index_type='string')

    # Generate a flat map cluster_id -> members
    print('Finding cluster assignments')
    clusters_df, _ = get_original_clustering(dataset=dataset, neuropil=neuropil,
                                             clusterer_or_k=clusterer_or_k)
    dfs = []
    for cluster_id, members in zip(clusters_df.cluster_id,
                                   clusters_df.original_voxels_in_cluster):
        dfs.append(pd.DataFrame({'cluster_id': cluster_id, 'member_id': members}))
    members_df = pd.concat(dfs).set_index('member_id').loc[Xdf.index]

    # Compute the distance matrix - this must be parameterised
    print('Computing distance')
    import mkl
    mkl.set_num_threads(6)
    D = dicedist_metric(Xdf)

    # Compute silhouette
    # Here we could go for the faster implementation in third_party, if needed
    print('Computing silhouette index')
    members_df['silhouette'] = silhouette_samples(D.values,
                                                  members_df.cluster_id.values,
                                                  metric='precomputed')
    return (members_df.
            reset_index().
            rename(columns=lambda col: {'index': 'member_id'}.get(col, col))
            [['cluster_id', 'member_id', 'silhouette']])

def fftvec(vec):
    """
    performs a fft on a vector with 3 components in the first index position
    This is really just a wrapper for fft, fftn and their inverses
    """
    try:
        from anfft import fft, fftn
        fft_type = 1
    except:
#        print "Could not import anfft, importing scipy instead."
#Update 9/18/2013 -- numpy with mkl is way faster than scipy
        import mkl
        mkl.set_num_threads(8)
        from numpy.fft import fft, fftn
        fft_type = 0
        
    if force_gpu:
        fft_type = 2 #set gpu fft's manually -- not sure what a automatic way would look like
    
    from numpy import complex64, shape, array, empty
    if vec.ndim > 2:
        if vec.shape[0] == 3:
            # "Vector": first index has size 3 so fft the other columns
            if fft_type==1:
                return array([fftn(i,measure=True) for i in vec]).astype(complex64)
#                result = empty(vec.shape, dtype=complex64)
#                result[0] = fftn(vec[0], measure=True).astype(complex64)
#                result[1] = fftn(vec[1], measure=True).astype(complex64)
#                result[2] = fftn(vec[2], measure=True).astype(complex64)
#                return result
                
            elif fft_type==0:
                return fftn(vec, axes=range(1,vec.ndim)).astype(complex64)
            elif fft_type==2:
#                return array([gpu_fft(i) for i in vec.astype(complex64)])
                result = empty(vec.shape, dtype=complex64)
                result[0] = gpu_fft(vec[0].copy())
                result[1] = gpu_fft(vec[1].copy())
                result[2] = gpu_fft(vec[2].copy())
                return result
        else: # "Scalar", fft the whole thing
            if fft_type==1:
                return fftn(vec,measure=True).astype(complex64)
            elif fft_type==0:
                return fftn(vec).astype(complex64)
            elif fft_type==2:
                return gpu_fft(vec.copy())
    elif vec.ndim == 1: #Not a vector, so use fft
        if fft_type==1:
            return fft(vec,measure = True).astype(complex64)
        elif fft_type==0:
            return fft(vec).astype(complex64)
        elif fft_type==2:
            return gpu_fft(vec.astype(complex64))
    else:
        #0th index is 3, so its a vector
        #return fft(vec, axis=1).astype(complex64)
        return array([fft(i) for i in vec])

def parallel_loop(args):

    import numpy as np
    import time

    import pysparsefht
    from utils import random_k_sparse

    try:
        import mkl as mkl_service
        # for such parallel processing, it is better 
        # to deactivate multithreading in mkl
        mkl_service.set_num_threads(1)
    except ImportError:
        pass

    n = args[0]

    b = np.arange(1, n-1)
    K = 2**b
    B = 2**b

    # compute value of C
    C = np.empty_like(b)
    C[:np.floor(n/3)] = n/b[:np.floor(n/3)]
    C[np.floor(n/3):np.floor(2*n/3)] = 3
    C[np.floor(2*n/3):] = n / (n - b[np.floor(2*n/3):])

    algo_name = params['algo_name']
    seed = args[1]

    if algo_name == 'RANDOM':
        algo = pysparsefht.ALGO_RANDOM
    elif algo_name == 'DETERMINISTIC':
        algo = pysparsefht.ALGO_OPTIMIZED
    else:
        ValueError('No such algorithm.')

    # initialize rng
    np.random.seed(seed)

    # a list for return values
    ret = []

    # generate a seed for the C RNG
    C_seed = np.random.randint(4294967295, dtype=np.uint32)

    # create sparse vector
    Tsfht, Tfht = pysparsefht.benchmark(K, B, C, 2**n, 
            loops=params['inner_loops'], warm=params['warm'], body=params['body'], max_mag=params['max_mag'],
            sfht_max_iter=params['max_iter'], seed=C_seed)

    return [Tsfht, Tfht, b, C]

    def _initialize_fft(self):

        """ Define the two-dimensional FFT methods.
        """

        if self.use_mkl:
            import mkl
            mkl.set_num_threads(self.nthreads)
            import mkl_fft
            self.fft =  (lambda x : mkl_fft.fft2(x))
            self.ifft = (lambda x : mkl_fft.ifft2(x))
        else:
            self.fft =  (lambda x : np.fft.fft2(x))
            self.ifft = (lambda x : np.fft.ifft2(x))

def dtest(n=50, d=0.0, r=0.0, model_covariate=True, niters=100, nperms=4999):
    import mkl
    mkl.set_num_threads(2)
    
    d   = float(d)
    r   = float(r)
        
    # Data/Distances
    pvals = np.zeros(niters)
    Fvals = np.zeros(niters)
    for i in xrange(niters):
        # Design
        
        ## Categorical
        gp  = np.repeat([0, 1], n/2)
        np.random.shuffle(gp)
        x   = gp*d + np.random.standard_normal(n)
        
        ## Continuous
        # see http://stackoverflow.com/questions/16024677/generate-correlated-data-in-python-3-3
        # and http://stats.stackexchange.com/questions/19367/creating-two-random-sequences-with-50-correlation?lq=1
        uncorrelated    = np.random.standard_normal((2,n))
        motion          = uncorrelated[0]
        y               = r*motion + np.sqrt(1-r**2)*uncorrelated[1]
        
        ## Design Matrix
        if model_covariate:
            design = np.vstack((np.ones(n), gp, motion)).T
        else:
            design = np.vstack((np.ones(n), gp)).T
                
        # Date
        points = np.vstack((x,y)).T
        
        # Distances
        dmat  = euclidean_distances(points)
        dmats = dmat[np.newaxis,:,:]
        
        # Only the group effect is the variable of interest
        cols = [1]
        
        # Call MDMR
        pval, Fval, _, _ = mdmr(dmats, design, cols, nperms)
        
        pvals[i] = pval
        Fvals[i] = Fval
    
    return pvals, Fvals

def prep_cwas_workflow(c, subject_infos):
    from CPAC.cwas import create_cwas
    import numpy as np
    
    try:
        import mkl
        mkl.set_num_threads(c.cwasThreads)
    except ImportError:
        pass
    
    print 'Preparing CWAS workflow'
    p_id, s_ids, scan_ids, s_paths = (list(tup) for tup in zip(*subject_infos))
    print 'Subjects', s_ids
    
    # Read in list of subject functionals
    lines   = open(c.cwasFuncFiles).readlines()
    spaths  = [ l.strip().strip('"') for l in lines ]
    
    # Read in design/regressor file
    regressor = np.loadtxt(c.cwasRegressorFile)

    wf = pe.Workflow(name='cwas_workflow')
    wf.base_dir = c.workingDirectory
    
    cw = create_cwas()
    cw.inputs.inputspec.roi         = c.cwasROIFile
    cw.inputs.inputspec.subjects    = spaths
    cw.inputs.inputspec.regressor   = regressor
    cw.inputs.inputspec.cols        = c.cwasRegressorCols
    cw.inputs.inputspec.f_samples   = c.cwasFSamples
    cw.inputs.inputspec.strata      = c.cwasRegressorStrata # will stay None?
    cw.inputs.inputspec.parallel_nodes = c.cwasParallelNodes
    cw.inputs.inputspec.memory_limit = c.cwasMemory
    cw.inputs.inputspec.dtype       = c.cwasDtype
    
    ds = pe.Node(nio.DataSink(), name='cwas_sink')
    out_dir = os.path.dirname(s_paths[0]).replace(s_ids[0], 'cwas_results')
    ds.inputs.base_directory = out_dir
    ds.inputs.container = ''

    wf.connect(cw, 'outputspec.F_map',
               ds, 'F_map')
    wf.connect(cw, 'outputspec.p_map',
               ds, 'p_map')

    wf.run(plugin='MultiProc',
                         plugin_args={'n_procs': c.numCoresPerSubject})

def dtest(pos_nodes=0, effect=0.0, dist="euclidean", n=100, nodes=400, nperms=4999, iters=100):
    import mkl
    mkl.set_num_threads(2)
    
    print "Start"

    #print "Categorical Effect"
    grp  = np.repeat([0, 1], n/2)
    np.random.shuffle(grp)

    #print "Design Matrix"
    design = np.vstack((np.ones(n), grp)).T
    cols = [1]

    #print "Distance Matrices"
    dmats = np.zeros((iters,n,n))
    for i in xrange(iters):
        #if (i % 10) == 0:
        #    print i,
        # Data
        ## Fist, I created the matrix with the random data
        points = np.random.standard_normal((n,nodes))
        ## Second, I select a random selection of nodes to add the effect
        neg_nodes = nodes - pos_nodes
        change_nodes = np.repeat([0,1], [neg_nodes, pos_nodes])
        np.random.shuffle(change_nodes)
        ## Finally, I add the effect to a select subjects and nodes
        for i in (change_nodes==1).nonzero()[0]:
            points[grp==1,i] += effect
        
        # Compute Distances
        if dist == "euclidean":
            dmat    = euclidean_distances(points)
        elif dist == "pearson":
            dmat    = compute_distances(points)
        else:
            raise Exception("Unknown distance measure %s" % dist)
        dmats[i]    = dmat
    #print ""

    #print "MDMR"
    pvals = []; Fvals = [];
    pvals, Fvals, _, _ = mdmr(dmats, design, cols, nperms)
    
    #print "Done"
    return pvals, Fvals

    def __init__(self, *args, **kwargs):
        super(Worker, self).__init__()

        self.nthreads = pyrat._nthreads  # number of threads for processing
        if pyrat._debug is True:
            self.nthreads = 1

        try:
            import mkl
            if self.nthreads > 1:  # switch of mkl multithreading
                mkl.set_num_threads(1)  # because we do it ourself
            else:
                mkl.set_num_threads(999)
        except ImportError:
            pass

        # self.blockprocess = True                                               # blockprocessing on/off
        # self.blocksize = 128                                                   # size of single block

        for para in self.para:  # copy defaults to self
            setattr(self, para['var'], para['value'])
        for (k, v) in kwargs.items():  # copy keywords to self
            setattr(self, k, v)  # eventually overwriting defaults
        if not hasattr(self, 'layer'):  # if no keyword was used
            self.layer = pyrat.data.active  # use active layer
        # --------------------------------------------------

        self.name = self.__class__.__name__  # name of worker class (string)
        self.input = ''  # input layer(s)
        self.output = ''  # output layer(s)
        self.blockoverlap = 0  # block overlap
        self.vblock = False  # vertical blocks on/off
        self.blocks = []  # list of block boundaries
        self.valid = []  # valid part of each block
        # self.block = False                                                     # actual block range / validity
        self.allowed_ndim = False
        self.require_para = False
        self.allowed_dtype = False

def mpirun(f,arguments,comm=MPI.COMM_WORLD,bcast=True):
    '''
    Wrapper for the parallel running of f using the mpi4py.

    Parameters
    ----------
    f : callable
        The function to be parallelly run using the mpi4py.
    arguments : list of tuple
        The list of arguments passed to the function f.
    comm : MPI.Comm, optional
        The MPI communicator.
    bcast : True or False
        When True, broadcast the result for all processes;
        Otherwise only the rank 0 process hold the result.

    Returns
    -------
    list
        The returned values of f with respect to the arguments.
    '''
    size=comm.Get_size()
    rank=comm.Get_rank()
    if size>1:
        import mkl
        mkl.set_num_threads(1)
    temp=[]
    for i,argument in enumerate(arguments):
        if i%size==rank:
            temp.append(f(*argument))
    temp=comm.gather(temp,root=0)
    result=[]
    if rank==0:
        for i in range(len(arguments)):
            result.append(temp[i%size][i//size])
    if bcast:
        result=comm.bcast(result,root=0)
    return result

def start_benchmark():
    print "Benchmark starting timing with numpy %s\nVersion: %s" % (numpy.__version__, sys.version)
    print ("-" * 80)

    for cur_threads in threads_range:
        header_set = False

        # This doesn't work: os.environ is not adjusted        
        #os.environ[THREADS_LIMIT_ENV] = '%d' % cur_threads
        
        mkl.set_num_threads(cur_threads)
        print "Maximum number of threads used for computation is : %d" % cur_threads

        header_str = "%20s" % "Function"
        header_str += ' - %9s - Speedup' % 'Time [ms]'

        if cur_threads == 1:        
            timings_single = []

        for ii,fun in enumerate(tests):

            result_str = "%20s" % fun.__name__
            t = timeit.Timer(stmt="%s()" % fun.__name__, setup="from __main__ import %s" % fun.__name__)
            res = t.repeat(repeat=3, number=1)
            timing =  1000.0 * sum(res)/len(res)
            
            if cur_threads == 1:        
                timings_single.append(timing)
        
            result_str += ' - %9.1f - %5.1f' % (timing, timings_single[ii]/timing)
                 
            if not header_set is True:
                print header_str
                header_set = True
        
            print result_str    

#!/usr/bin/env python

import os, sys
from os import path
sys.path.append("/home2/data/Projects/CWAS/pyClusterROI")

if len(sys.argv) != 2:
    sys.exit("Usage: %s num-threads" % sys.argv[0])

# control mkl
import mkl
mkl.set_num_threads(int(sys.argv[1]))


###
# 1. SETUP
###

print "1. Setup"

obase = "/home2/data/Projects/CWAS/development+motion/spatial_cluster"

# functions for connectivity metric
from make_local_connectivity_ones import *

# name of the maskfile that we will be using
roidir = "/home2/data/Projects/CWAS/share/development+motion/rois"
maskfile = path.join(roidir, "mask_gray_4mm.nii.gz")


###

 def __enter__(self):
     mkl.set_num_threads(self.num_threads)

	def __init__(self): 

		import os, sys
		MoleculeName = ''
		if len(sys.argv) < 2:
			print "Ideally you give me a molecule name, Defaulting to /Integrals"
		else: 
			MoleculeName = sys.argv[1]

		self.MoleculeName = MoleculeName # Defined Globally in Main.py 
		self.nocc = 4
		self.nvirt = 4
		self.nmo = 8    # These will be determined on-the-fly reading from disk anyways. 
		self.occ = []
		self.virt = []
		self.all = []
		self.alpha = []
		self.beta = []

		# Try to get a faster timestep by freezing the CoreOrbitals. 
		# if it's not None, then freeze that many orbitals (evenly alpha and beta.) 
		self.FreezeCore = 8 

		self.AvailablePropagators = ["phTDA2TCL","Whole2TCL", "AllPH"]
		self.Propagator = "phTDA2TCL"
		self.Correlated = True
		self.SecularApproximation = 0 # 0 => Nonsecular 1=> Secular 
		self.ImaginaryTime = True
		
		self.Temperature = 303.15*2
		self.TMax =  250.0 
		self.TStep = 0.01
		self.tol = 1e-9
		self.Safety = .9 #Ensures that if Error = Emax, the step size decreases slightly
		self.RK45 = True
		self.LoadPrev = False
#		self.MarkovEta = 0.05
		self.MarkovEta = 0.001
		self.Tc = 3000.0
		self.Compiler = "gcc"
		self.Flags = ["-mtune=native", "-O3"]
		self.latex = True
		self.fluorescence = True # If true, this performs the fluorescence calculations in SpectralAnalysis
		# above a certain energy 
		# no relaxation is included and the denominator is the bare electronic 
		# denominator to avoid integration issues.
		self.UVCutoff = 300.0/27.2113 # No Relaxation above 1eV		

		try:
			mkl.set_num_threads(8)
		except NameError:
			print "No MKL so I can't set the number of threads"
		
		# This is a hack switch to shut off the Boson Correlation Function 
		# ------------------------------------------------------------------------
		# Adiabatic = 0 # Implies numerical integration of expdeltaS and bosons. 
		# Adiabatic = 1 # Implies numerical integration of expdeltaS no bosons. 
		# Adiabatic = 2 # Implies analytical integration of expdeltaS, no bosons.  
		# Adiabatic = 3 # Implies analytical integration of expdeltaS, no bosons, and the perturbative terms are forced to be anti-hermitian. 		
		# Adiabatic = 4 # Implies Markov Approximation. 
		self.Adiabatic = 0

		self.Inhomogeneous = False
		self.InhomogeneousTerms = self.Inhomogeneous
		self.Inhomogenous = False
		self.Undressed = True
		self.ContBath = True # Whether the bath is continuous/There is a continuous bath. 
		self.ReOrg = True 
		self.FiniteDifferenceBosons = False

		self.DipoleGuess = True # if False, a superposition of bright states will be used. 
		self.AllDirections = True # Will initalize three concurrent propagations. 
		self.DirectionSpecific = True		
		
		self.InitialDirection = -1 # 0 = x etc. Only excites in the x direction -1=isotropic
		self.BeginWithStatesOfEnergy = None 
#		self.BeginWithStatesOfEnergy = 18.7/27.2113
#		self.BeginWithStatesOfEnergy = 18.2288355687/27.2113
		self.PulseWidth = 1.7/27.2113
		
		self.Plotting = True
		self.DoCisDecomposition = True
		self.DoBCT = True # plot and fourier transform the bath correlation tensor. 
		self.DoEntropies = True
		if (self.Undressed): 
			self.DoEntropies = False
		self.FieldThreshold = pow(10.0,-7.0)
		self.ExponentialStep = False  #This will be set automatically if a matrix is made. 

		self.LegendFontSize = 14 
		self.LabelFontSize = 16
		
		print "--------------------------------------------"		
		print "Running With Overall Parameters: "
		print "--------------------------------------------"
		print "self.MoleculeName", self.MoleculeName 
		print "self.AllDirections", self.AllDirections 
		print "self.Propagator", self.Propagator 
		print "self.Temperature", self.Temperature 
		print "self.TMax", self.TMax 
#.........这里部分代码省略.........

import numpy
import numpy.fft as fft
numpy.use_fastnumpy = True
import time
#from scipy.fftpack import fft
import mkl

print 'Intel MKL version:', mkl.get_version_string()
print 'Intel cpu_clocks:', mkl.get_cpu_clocks()
print 'Intel cpu_frequency:', mkl.get_cpu_frequency()
#print 'Intel MKL, freeing buffer memory:', mkl.thread_free_buffers()

print 'max Intel threads:', mkl.get_max_threads()

mkl.set_num_threads(2)

N = 2**16

print 'using numpy', numpy.__version__
a = numpy.random.rand(2, N)
print a.shape, 'items'
t0 = time.clock()
for i in range(100):
    continue
base = time.clock()-t0
fftn = fft.fftn
t0 = time.clock()
for i in range(100):
    r = fftn(a, (N,), (1,))
print 'simple loop', time.clock()-t0-base

import os, sys
from os import path as op

import scipy
import nibabel as nb
import numpy as np
from pandas import read_table, read_csv
from patsy import dmatrices, dmatrix

from CPAC.cwas import cwas
from CPAC.cwas.utils import calc_subdists, calc_mdmrs
from CPAC.cwas.subdist import *
from CPAC.cwas.mdmr import mdmr, gen_perms

import mkl
mkl.set_num_threads(8)


####
# Cool Functions
####

def load_subject(filepath, dtype='float64'):
    return nb.load(filepath).get_data().astype(dtype)

def load_subjects(filepaths, dtype='float64'):
    print "Loading Subjects"
    funcs = [ load_subject(fp, dtype) for fp in filepaths ]
    return funcs

def rois2voxels(dat, rois):