text:
import numpy as np
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.feature_selection import VarianceThreshold, SelectKBest, f_regression
from sklearn.impute import SimpleImputer
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler

from models.preprocessors.threshold_selectors import CorThreshold
from models.preprocessors.column_selectors import make_col_selector

class Preprocessor(BaseEstimator, TransformerMixin):
    def __init__(self, desc_col_selector, fgp_col_selector, p=0.9, cor_th=0.9, k='all'):
        self.desc_col_selector = desc_col_selector
        self.fgp_col_selector = fgp_col_selector
        self.p = p
        self.cor_th = cor_th
        self.k = k

    def _init_hidden_models(self):
        # TODO: not included here the non-retained predictor for simplicity

        # Try out MultipleImputer (similar to MICE when sample_posterior = True and random_state = number)
        # Try out PCA and supress SelectKBest
        self._desc_pipeline = Pipeline([
            ('scaler', StandardScaler()),
            ('imputation', SimpleImputer(missing_values=np.nan, strategy='median', add_indicator=True)),
            ('var_threshold', VarianceThreshold()),
            ('cor_selector', CorThreshold(threshold=self.cor_th)),
            ('f_selector', SelectKBest(score_func=f_regression, k=self.k))
        ])
        self._fgp_vs = VarianceThreshold(threshold=self.p * (1 - self.p))

    def fit(self, X, y=None):
        self._init_hidden_models()
        X_desc = self.desc_col_selector(X)
        X_fgp = self.fgp_col_selector(X)

        self._desc_pipeline.fit(X_desc, y)
        self._fgp_vs.fit(X_fgp)
        return self

    def transform(self, X, y=None):
        X_desc = self.desc_col_selector(X)
        X_fgp = self.fgp_col_selector(X)

        X_desc_proc = self._desc_pipeline.transform(X_desc)
        X_fgp_proc = self._fgp_vs.transform(X_fgp)
        new_X = np.concatenate([X_desc_proc, X_fgp_proc], axis=1)

        self.transformed_desc_col_selector = make_col_selector(X_desc_proc.shape[1], return_first_half=True)
        self.transformed_fgp_col_selector = make_col_selector(X_desc_proc.shape[1], return_first_half=False)

        return new_X

ext: py

reason: import numpy


text:
from online_leastsq import OnlineLeastSquares
import numpy as np
import math

class LinearSystemID:
    """System identification for a system y = Ax + Bu + C with m state
    variables and n inputs."""
    def __init__(self,m,n):
        self.m,self.n = m,n
        self.coeffPattern = [None,None,None]
        self.estimators = [OnlineLeastSquares(self.m+self.n+1) for i in xrange(m)]

    def setPattern(self,Apattern,Bpattern,Cpattern):
        """The patterns are list-of-lists of size mxm, mxn, and a list of
        size m, which indicate whether the (i,j)'th entry of the A,B, and C
        matrices are fixed (respectively).  A None entry indicates a free
        coefficient, while a numeric entry indicates a fixed coefficient.
        """
        self.coeffPattern = [Apattern,Bpattern,Cpattern]
        for i in xrange(self.m):
            self._updateEstimatorSize(i)

    def fixA(self,i,j,value):
        """Sets the i,j'th entry of the A matrix to a fixed value"""
        if self.coeffPattern[0] == None:
            m,n=self.m,self.n
            self.coeffPattern[0] = [[None]*m for i in xrange(m)]
        self.coeffPattern[0][i][j]=value
        self._updateEstimatorSize(i)

    def fixB(self,i,j,value):
        """Sets the i,j'th entry of the B matrix to a fixed value"""
        if self.coeffPattern[1] == None:
            m,n=self.m,self.n
            self.coeffPattern = [[None]*n for i in xrange(m)]
        self.coeffPattern[1][i][j]=value
        self._updateEstimatorSize(i)

    def fixC(self,i,value):
        """Sets the i'th entry of the C vector to a fixed value"""
        if self.coeffPattern[2] == None:
            m,n=self.m,self.n
            self.coeffPattern[2] = [None]*m
        self.coeffPattern[2][i]=value
        self._updateEstimatorSize(i)

    def add(self,x,u,y,weight=1.0):
        """Adds a new datapoint to the estimator"""
        assert(len(y)==self.m)
        assert(len(x)==self.m)
        assert(len(u)==self.n)
        if isinstance(x,np.ndarray): x = x.tolist()
        if isinstance(u,np.ndarray): u = u.tolist()
        xu1 = x + u + [1.0]
        if self.coeffPattern == [None,None,None]:
            for yi,e in zip(y,self.estimators):
                e.add(xu1,yi,weight)
        else:
            #each row might have some fixed values
            for i,(yi,e) in enumerate(zip(y,self.estimators)):
                if e == None: continue
                (xuc,constOffset) = self._toEstimator(i,x,u)
                rhs = yi - constOffset
                e.add(xuc,rhs,weight)
        return

    def discount(self,discountFactor,type='geometric'):
        """Reduces the effects of prior readings."""
        for e in self.estimators:
            e.discount(discountFactor,type)
        return

    def setModelPrior(self,A,B,C,priorWeight):
        """Adds in a prior belief for the model.
        Must be called AFTER fixing coefficients and BEFORE adding any
        datapoints."""
        Cpattern = self.coeffPattern[2]
        for i in xrange(self.m):
            ai = A[i,:].tolist()
            bi = B[i,:].tolist()
            (xuc,constant) = self._toEstimator(i,ai,bi)
            if Cpattern == None or Cpattern[i] == None:
                xuc[-1] = C[i]
            self.estimators[i].setPrior(np.array(xuc),priorWeight)
        return

    def getModel(self):
        """Returns the estimated triple (A,B,C) as numpy arrays"""
        m,n = self.m,self.n
        A = np.zeros((m,m))
        B = np.zeros((m,n))
        C = np.zeros(m)
        Apattern,Bpattern,Cpattern = self.coeffPattern
        for i,e in enumerate(self.estimators):
            aofs = 0
            bofs = m
            cofs = m+n
            if Apattern==None:
                ai = e.x[aofs:m+aofs]
            else:
                bofs=aofs
                ai = []
                for j,pj in enumerate(Apattern[i]):
                    if pj == None:
                        ai.append(e.x[bofs])
                        bofs += 1
                    else:
                        ai.append(pj)
            if Bpattern==None:
                bi = e.x[bofs:n+bofs]
            else:
                cofs=bofs
                bi = []
                for j,pj in enumerate(Bpattern[i]):
                    if pj == None:
                        bi.append(e.x[cofs])
                        cofs += 1
                    else:
                        bi.append(pj)
            if Cpattern==None:
                ci = e.x[cofs]
                cofs+=1
            else:
                if Cpattern[i] == None:
                    ci = e.x[cofs]
                    cofs+=1
                else:
                    ci = Cpattern[i]
            assert(cofs == e.n)
            assert len(ai)==m
            assert len(bi)==n
            A[i,:] = ai
            B[i,:] = bi
            C[i] = ci
        return (A,B,C)

    def getOutput(self,x,u):
        """Returns the estimate A*x+B*u+c"""
        assert(len(x)==self.m)
        assert(len(u)==self.n)
        if isinstance(x,np.ndarray): x = x.tolist()
        if isinstance(u,np.ndarray): u = u.tolist()
        dx = []
        if self.coeffPattern == [None,None,None]:
            xuc = np.array(x + u + [1.0])
            for e in self.estimators:
                dx.append(np.dot(e.x,xuc))
        else:
            for i,e in enumerate(self.estimators):
                (xuc,constOffset) = self._toEstimator(i,x,u)
                dx.append(np.dot(e.x,xuc)+constOffset)
        return dx

    def _updateEstimatorSize(self,index):
        """Helper."""
        Apattern,Bpattern,Cpattern = self.coeffPattern
        m,n = self.m,self.n
        numFixed = 0
        if Apattern!=None:
            numFixed += len(v for v in Apattern[index] if v != None)
        if Bpattern!=None:
            numFixed += len(v for v in Bpattern[index] if v != None)
        if Cpattern!=None:
            if Cpattern[index]!=None:
                numFixed += 1
        if numFixed==m+n+1:
            self.estimators[index]=None
        else:
            self.estimators[index]=OnlineLeastSquares(m+n+1-numFixed)
        return

    def _toEstimator(self,index,x,u):
        """Helper: Projects x,u to the pattern xe taken by the index'th
        estimator. Returns the pair (xe,constant offset) where the index'th
        row of Ax+Bu+C is equal to dot(xe,estimator.coeffs) + constOffset."""
        Apattern,Bpattern,Cpattern = self.coeffPattern
        xuc = []
        constOffset = 0
        if Apattern == None:
            xuc += x
        else:
            xuc += [xj for (xj,pj) in zip(x,Apattern[index]) if pj == None]
            constOffset += sum([xj*pj for (xj,pj) in zip(x,Apattern[index]) if pj != None])
        if Bpattern == None:
            xuc += u
        else:
            xuc += [uj for (uj,pj) in zip(u,Bpattern[index]) if pj == None]
            constOffset += sum([uj*pj for (uj,pj) in zip(u,Bpattern[index]) if pj != None])
        if Cpattern == None:
            xuc += [1.0]
        else:
            constOffset = Cpattern[index]
        return (xuc,constOffset)

def testLinearSystemID(A,B,C,N,xScale=1,uScale=1):
    m,n = B.shape
    sysid = LinearSystemID(m,n)
    print "Actual model:",
    print "A:"
    print A
    print "B:"
    print B
    print "C:"
    print C

    X = [np.random.rand(m)*xScale*2-np.ones(m)*xScale for i in range(N)]
    U = [np.random.rand(n)*uScale-np.ones(n)*uScale for i in range(N)]
    for x,u in zip(X,U):
        y = np.dot(A,x)+np.dot(B,u)+C
        sysid.add(x,u,y)
    print "Estimated model:",
    eA,eB,eC=sysid.getModel()
    print "A:"
    print eA
    print "B:"
    print eB
    print "C:"
    print eC
    resid = 0.0
    for x,u in zip(X,U):
        y = np.dot(A,x)+np.dot(B,u)+C
        ey = np.dot(eA,x)+np.dot(eB,u)+C
        #print "Error",np.dot(ey,ey)
        resid += np.dot(ey-y,ey-y)
    print "RMSE",math.sqrt(resid/len(X))

def testNonlinearSystemID(func,m,n,N,discountParams,x0,uScale=1):
    sysid = LinearSystemID(m,n)
    sysid.add(x0,np.zeros(n),x0)

    if not hasattr(discountParams,'__iter__'):
        discountParams = (discountParams,'geometric')

    #simulate a random trace
    X = [x0]
    U = []
    for i in xrange(N):
        u = np.random.rand(n)*uScale*2.0 - np.ones(n)*uScale
        U.append(u)
        X.append(func(X[-1],u))
    resid = 0.0
    discresid = 0.0
    errors = []
    for x,u,xn in zip(X[:-1],U,X[1:]):
        xpred = sysid.getOutput(x,u)
        e2 = np.dot(xn-xpred,xn-xpred)
        errors.append(e2)
        resid += e2
        discount = discountParams[0] if discountParams[1]=='geometric' else (1.0/discountParams[0]-1)/sysid.estimators[0].sumWeight
        discresid = discresid*discount + e2
        sysid.discount(*discountParams)
        sysid.add(x,u,xn)
    print "Nonlinear estimation with discount",discountParams
    print "Final estimated model:",
    eA,eB,eC=sysid.getModel()
    print "A:"
    print eA
    print "B:"
    print eB
    print "C:"
    print eC
    print "RMSE",math.sqrt(resid/(len(X)-1))
    print "Discounted squared errors",discresid
    #print "Errors",errors

if __name__=='__main__':
    print "Testing standard double integrator with offset"""
    mass = 5.0
    f = 1.0
    A = np.array([[0,1],[0,0]])
    B = np.array([[0],[1/mass]])
    C = np.array([0,f/mass])
    testLinearSystemID(A,B,C,100)

    def ABCFunc(x,u):
        dt = 0.1
        return x + dt*(np.dot(A,x)+np.dot(B,u)+C)
    x0 = np.array([0,0])
    """
    testNonlinearSystemID(ABCFunc,2,1,100,1.0,x0)
    testNonlinearSystemID(ABCFunc,2,1,100,0.95,x0)
    testNonlinearSystemID(ABCFunc,2,1,100,0.9,x0)
    testNonlinearSystemID(ABCFunc,2,1,100,0.8,x0)
    exit(1)
    """

    print "Testing damped pendulum with gravity"""
    def pendulumFunc(x,u):
        theta,dtheta = x
        torque = u[0]
        g = -9.8
        dt = 0.1
        mass = 1
        kD = 0.5
        ddtheta = torque/mass + math.sin(theta)*g - kD*dtheta
        return np.array([theta+dt*dtheta,dtheta+dt*ddtheta])
    x0 = np.array([math.pi*0.5,0])
    testNonlinearSystemID(pendulumFunc,2,1,1000,1.0,x0)
    testNonlinearSystemID(pendulumFunc,2,1,1000,0.99,x0)
    testNonlinearSystemID(pendulumFunc,2,1,1000,0.95,x0)
    testNonlinearSystemID(pendulumFunc,2,1,1000,(0.01,'hyperbolic'),x0)
    testNonlinearSystemID(pendulumFunc,2,1,1000,(0.05,'hyperbolic'),x0)
    testNonlinearSystemID(pendulumFunc,2,1,1000,(0.1,'hyperbolic'),x0)
    testNonlinearSystemID(pendulumFunc,2,1,1000,(0.2,'hyperbolic'),x0)
    testNonlinearSystemID(pendulumFunc,2,1,1000,(0.5,'hyperbolic'),x0)
    testNonlinearSystemID(pendulumFunc,2,1,1000,(0.75,'hyperbolic'),x0)
    """
    testNonlinearSystemID(pendulumFunc,2,1,1000,0.9,x0)
    testNonlinearSystemID(pendulumFunc,2,1,1000,0.8,x0)
    testNonlinearSystemID(pendulumFunc,2,1,1000,0.7,x0)
    testNonlinearSystemID(pendulumFunc,2,1,1000,0.5,x0)
    testNonlinearSystemID(pendulumFunc,2,1,1000,0.25,x0)
    """
    print "Actual:"
    print "[q' ] = [ 1, 0.1 ]*[q ] + [0  ]*u + [0                    ]"
    print "[dq']   [ 0,   1 ] [dq]   [0.1]     [-.98 sin(q) - 0.05 dq]"


text:
"""Compiled and parallelized functions."""

import warnings
import numpy as np
from numba import jit

# Disable bad division warning when summing up squares
warnings.filterwarnings("ignore", category=RuntimeWarning)

@jit(nopython=True, parallel=True)
def loop_u(u, n_h, X_v, U, U_no_noise, X, mu_lhs, u_noise=0., x_noise=0.):
    """Return the inputs/snapshots matrices from parallel computation."""
    # pylint: disable=not-an-iterable

    for i in range(mu_lhs.shape[0]):
        X_v[i, :] = mu_lhs[i]
        U_i_no_noise = u(X, 0, X_v[i, :]).reshape((n_h,))
        if x_noise > 0.:
            # X_v[i, :] += x_noise*np.std(X_v[i, :])*np.random.randn(X_v[i, :].shape[0])
            X_v[i, :] += np.random.normal(0., x_noise*np.std(X_v[i, :]), X_v.shape[1])
        U_i = u(X, 0, X_v[i, :]).reshape((n_h,))
        if u_noise > 0.:
            U_i += u_noise*np.std(U_i)*np.random.randn(U_i.shape[0])
        U[:, i] = U_i
        U_no_noise[:, i] = U_i_no_noise
    U_struct = U
    return X_v, U, U_struct, U_no_noise

@jit(nopython=True, parallel=True)
def loop_u_t(u, n_t, n_v, n_xyz, n_h,
             X_v, U, U_no_noise, U_struct, X, mu_lhs, t_min, t_max, u_noise=0., x_noise=0.):
    """Return the inputs/snapshots matrices from parallel computation (w/ t)."""
    # Creating the time steps
    t = np.linspace(t_min, t_max, n_t)
    tT = t.reshape((n_t, 1))
    # pylint: disable=not-an-iterable
    for i in range(mu_lhs.shape[0]):
        # Getting the snapshot times indices
        s = n_t * i
        e = n_t * (i + 1)

        # Setting the regression inputs (t, mu)
        mu_i_no_noise = mu_lhs[i, :]
        mu_i = mu_lhs[i, :]
        dev = np.std(mu_i)
        if dev == 0.:
            dev = mu_i[0]
        mu_i = mu_i_no_noise + \
               x_noise*dev*np.random.randn(mu_i.shape[0])
        X_v[s:e, :] = np.hstack((tT, np.ones_like(tT)*mu_i))

        # Calling the analytical solution function
        Ui = np.zeros((n_v, n_xyz, n_t))
        Ui_no_noise = np.zeros((n_v, n_xyz, n_t))
        for j in range(n_t):
            Uij = u(X, t[j], mu_i)
            Uij_no_noise = u(X, t[j], mu_i_no_noise)
            if u_noise > 0.:
                Uij += u_noise*np.std(Uij)*np.random.randn(Uij.shape[0], Uij.shape[1])
            Ui[:, :, j] = Uij
            Ui_no_noise[:, :, j] = Uij_no_noise

        U[:, s:e] = Ui.reshape((n_h, n_t))
        U_no_noise[:, s:e] = Ui_no_noise.reshape((n_h, n_t))
        U_struct[:, :, i] = U[:, s:e]
    return X_v, U, U_struct, U_no_noise

@jit(nopython=True, parallel=True)
def lhs(n, samples):
    """Borrowed, parallelized __lhscentered() from pyDOE."""

    # Generate the intervals
    cut = np.linspace(0, 1, samples + 1)

    # Fill points uniformly in each interval
    u = np.random.rand(samples, n)
    a = cut[:samples]
    b = cut[1:samples + 1]
    rdpoints = np.zeros(u.shape)
    for j in range(n):
        rdpoints[:, j] = u[:, j]*(b-a) + a

    # Make the random pairings
    H = np.zeros(rdpoints.shape)
    for j in range(n):
        order = np.random.permutation(np.arange(samples))
        H[:, j] = rdpoints[order, j]

    return H

max_stars_repo_name: flabowski/POD-UQNN

max_issues_repo_name: flabowski/POD-UQNN

max_forks_repo_name: flabowski/POD-UQNN

reason: import numpy,from numba


text:
"""
Copyright © 2017 United States Government as represented by the Administrator, National Aeronautics and Space Administration.
No Copyright is claimed in the United States under Title 17, U.S. Code.  All Other Rights Reserved.
"""
import plotly.graph_objs as go
import numpy as np
from util import arg_to_latex, beautify_latex, cast_0_dim
from plotly import figure_factory as ff
import pandas as pd
from collections import defaultdict

def line_plot(result, titles, verbose = False, **kwargs):
	'''N-d line plot f(t)'''
	f = result[titles['variable']]
	t_name, t = list(result.items())[0]
	if len(result) == 2:
		if len(f.shape) == 1:
			if verbose:
				print('\t1-d output', f.shape)
			assert f.shape == t.shape
			y = f
			layout = go.Layout(
				title = titles['title'],
				xaxis = dict(title = '${}$'.format(t_name)),
				yaxis = dict(title = titles['title_lhs']))
			trace = go.Scatter(x = t, y = y, name = titles['title_short'])
			chart_type = 'line'

		elif f.shape[1] == 2:
			if verbose:
				print('\t2-d output', f.shape)
			x = f[:,0]
			y = f[:,1]
			# t_name, t = result.items()[0]
			text = ["{}:{}".format(t_name, v) for v in t]
			layout = go.Layout(title = titles['title'],
							   xaxis = dict(title = 'x' + titles['units']),
							  yaxis = dict(title = 'y' + titles['units']))
			trace = go.Scatter(x = x, y = y, text = text)
			chart_type = '2d-line'

		elif f.shape[1] == 3:
			if verbose:
				print('\t3-d output', f.shape)
			x = f[:,0]
			y = f[:,1]
			z = f[:,2]
			# t_name, t = result.items()[0]
			text = ["{}:{}".format(t_name, v) for v in t]
			layout = go.Layout(
				title = titles['title'],
				scene = dict(
					xaxis = dict(title = 'x' + titles['units']),
					yaxis = dict(title = 'y' + titles['units']),
					zaxis = dict(title = 'z' + titles['units']),
					))
			trace = go.Scatter3d(x = x, y = y, z = z, text = text, mode = 'lines')
			chart_type = '3d-line'

	elif len(result) == 4:
		arg0, val0 = t_name, t
		arg1, val1 = list(result.items())[1]
		arg2, val2 = list(result.items())[2]

		assert(f.shape == t.shape)

		text = ["{}:{}".format(titles['variable'], v) for v in f]

		trace = go.Scatter3d(
			x = val0,
			y = val1,
			z = val2,
			text = text,
			marker = dict(
				color = result[titles['variable']],
				colorscale = 'Viridis',
				colorbar = dict(title = titles['title_short'])),
			line = dict(
				color = result[titles['variable']],
				colorscale='Viridis'),
			mode = 'markers+lines')
		layout = go.Layout(
			title = titles['title'],
			scene = dict(
				xaxis = dict(title = arg0),
				yaxis = dict(title = arg1),
				zaxis = dict(title = arg2),))
		chart_type = '4d-line'
	else:
		raise NotImplementedError('shape not supported: {}'.format(f.shape))

	return [trace], chart_type, layout

def vector_plot(result, titles, verbose = False, **kwargs):
	variable = titles['variable']
	val0 = list(result.values())[0]

	if (result[variable].shape == val0.shape) & (val0.shape[1] == 2):
		if verbose:
			print('\t 2-d output', result[variable].shape)
		u = result[variable][:,0]
		v = result[variable][:,1]
		x = val0[:,0]
		y = val0[:,1]

		trace = ff.create_quiver(
			x, y, u, v,
			# scale=.25,
			# arrow_scale=.4,
			name='quiver',
			line=dict(width=1))['data'][0] # extracts quiver trace

		layout = go.Layout(title = titles['title'],
							  xaxis = dict(title = 'x'),
							  yaxis = dict(title = 'y'))
		chart_type = '2d-vector'

	elif (result[variable].shape == val0.shape) & (val0.shape[1] == 3):
		if verbose:
			print('\t 3-d output', result[variable].shape)
		u = result[variable][:,0].tolist()
		v = result[variable][:,1].tolist()
		w = result[variable][:,2].tolist()
		x = val0[:,0].tolist()
		y = val0[:,1].tolist()
		z = val0[:,2].tolist()
		trace = go.Cone(x = x, y = y, z = z, u = u, v = v, w = w)

		layout = go.Layout(
			title = titles['title'],
			scene = dict(
				xaxis = dict(title = 'x'),
				yaxis = dict(title = 'y'),
				zaxis = dict(title = 'z'),
				))
		chart_type = '3d-vector'
	else:
		raise NotImplementedError('plot type not supported yet')

	return [trace], chart_type, layout

def contour_plot(result, titles, indexing, verbose = False, **kwargs):
	variable = titles['variable']
	if verbose:
		print('\t2-d output', result[variable].shape)
	z = result[variable]
	title = titles['title']
	arg0, val0 = list(result.items())[0]
	arg1, val1 = list(result.items())[1]

	if verbose:
		print('2 inputs: {} {}, {} {}'.format(arg0, val0.shape, arg1, val1.shape))
	traces = []
	if len(val0.shape) == 1:
		if verbose:
			print('\t\t1-d args', end=' ')
		if result[variable].shape[0] == val0.shape[0]:
			if indexing == 'ij':
				if verbose:
					print('{} indexing'.format(indexing))
				trace = go.Contour(x = val0,
								   y = val1,
								   z = z.T)
			else:
				if verbose:
					print('{} indexing'.format(indexing))
				trace = go.Contour(x = val0,
								   y = val1,
								   z = z)

		else:
			if verbose:
				print('xy indexing')
			trace = go.Contour(x = val0,
							   y = val1,
							   z = z)
		layout = go.Layout(
			title = title,
			xaxis = dict(title = '${}$'.format(arg0)),
			yaxis = dict(title = '${}$'.format(arg1)))
		traces.append(trace)
		chart_type = '2d-grid'
	else:
		if verbose:
			print('\t\t2-d args', val0.shape, val1.shape)
		assert val0.shape == val1.shape
		assert val0.shape == z.shape
		xaxis = dict(title = beautify_latex('${}$'.format(arg_to_latex(arg0))))
		yaxis = dict(title = beautify_latex('${}$'.format(arg_to_latex(arg1))))
		carpet_traces, layout = carpet_plot(result, title, xaxis, yaxis)
		traces.extend(carpet_traces)
		chart_type = '2d-skew'

	return traces, chart_type, layout

def carpet_plot(results, title, xaxis, yaxis, indexing = 'xy', **kwargs):
	'''Assumes ordered dict where values have the same shape'''
	# print title, xaxis, yaxis
	arg0, val0 = list(results.items())[0]
	arg1, val1 = list(results.items())[1]
	arg2, val2 = list(results.items())[2]

	a = list(range(val0.shape[0]))
	b = list(range(val0.shape[1]))

	aa, bb = np.meshgrid(b,a, indexing = indexing)

	trace1 = go.Contourcarpet(
		a = aa.ravel(),
		b = bb.ravel(),
		z = val2.ravel(),
	)

	trace2 = go.Carpet(
		a = aa.ravel(),
		b = bb.ravel(),
		x = val0.ravel(),
		y = val1.ravel(),
	)
	traces = [trace1, trace2]
	layout = go.Layout(
		title = title,
		xaxis = xaxis,
		yaxis =  yaxis)
	return traces, layout

def plane(result, titles, indexing = 'xy', verbose = False, **kwargs):
	variable = titles['variable']
	arg0, val0 = list(result.items())[0]
	arg1, val1 = list(result.items())[1]
	arg2, val2 = list(result.items())[2]
	traces = []
	if verbose:
		print('\t3-d output', result[variable].shape)
	surfacecolor = result[variable].squeeze()
	create_meshgrid = True
	for x in val0, val1, val2:
		if len(get_arg_shapes(x)[0]) != 1:
			create_meshgrid = False
			ones = np.ones_like(x)
	if create_meshgrid:
		xx, yy, zz = np.meshgrid(val0, val1, val2)
	else:
		xx, yy, zz = [v*ones for v in [val0, val1, val2]]
	traces.append(go.Surface(
		x = xx.squeeze(),
		y = yy.squeeze(),
		z = zz.squeeze(),
		surfacecolor = surfacecolor))
	layout = go.Layout(
		title = titles['title'],
		scene = dict(
			xaxis = dict(title = arg0),
			yaxis = dict(title = arg1),
			zaxis = dict(title = arg2),
		)
	)
	chart_type = '3d-plane'

	return traces, chart_type, layout

def surface(result, titles, verbose = False, **kwargs):
	variable = titles['variable']
	title = titles['title']
	arg0, val0 = list(result.items())[0]
	arg1, val1 = list(result.items())[1]
	arg2, val2 = list(result.items())[2]

	if verbose:
		print('\t2-d output', result[variable].shape)

	traces = []

	if len(result[variable].shape) == 2:
		surfacecolor = result[variable]
		xx = cast_0_dim(val0, result[variable])
		yy = cast_0_dim(val1, result[variable])
		zz = cast_0_dim(val2, result[variable])
		traces.append(go.Surface(
			x = xx,
			y = yy,
			z = zz,
			surfacecolor = surfacecolor))
		layout = go.Layout(
			title = title,
			scene = dict(
				xaxis = dict(title = arg0),
				yaxis = dict(title = arg1),
				zaxis = dict(title = arg2),))
		chart_type = '3d-surface-scalar'

	elif len(result[variable].shape) == 1:
		if verbose:
			print('\tscalar output', result[variable].shape)

		elif result[variable].shape[0] == 1:
			traces.append(go.Surface(
				x = val0,
				y = val1,
				z = val2,
				))
			layout = go.Layout(
				title = title,
				scene = dict(
					xaxis = dict(title = arg0),
					yaxis = dict(title = arg1),
					zaxis = dict(title = arg2),
					))
			chart_type = '3d-surface'

	elif len(result[variable]) == 1: #scalar or color value?
		if verbose:
			print('\tsingle valued output', result[variable].shape)
		surfacecolor = result[variable]
		arg0, val0 = list(result.items())[0]
		arg1, val1 = list(result.items())[1]
		arg2, val2 = list(result.items())[2]
		xx = val0
		yy = val1
		zz = val2
		traces.append(go.Surface(
			x = xx,
			y = yy,
			z = zz,
			))
		layout = go.Layout(
			title = title,
			scene = dict(
				xaxis = dict(title = arg0),
				yaxis = dict(title = arg1),
				zaxis = dict(title = arg2),
			)
		)

	return traces, chart_type, layout

plot_dict = {
	(1,)	:	{(('N','M'), ('N','M'), ('N','M')): {'name': '3d-parametric', 'func': surface}},
	('N',)	:	{
		(('N',),) : {'name': '1d-line', 'func': line_plot},
		(('N',),('N',),('N',)): {'name': '3d-line-scalar', 'func': line_plot},},
	('N',2) :	{
		(('N',),) :{'name': '2d-line', 'func': line_plot},
		(('N',2),):{'name': '2d-vector', 'func': vector_plot},
	},
	('N',3) : 	{
		(('N',),) :{'name': '3d-line', 'func': line_plot},
		(('N',3),):{'name': '3d-vector', 'func': vector_plot},
	},
	('N','M'):	{
		(('N',),('M',)) :{'name': '2d-contour', 'func': contour_plot},
		(('N','M'),('N','M')):{'name': '2d-contour-skew', 'func': contour_plot},
		(('N','M'), ('N','M'), ('N','M')): {'name': '3d-parametric-scalar', 'func': surface},
		((1,),('N','M'),('N','M')):{'name': '3d-plane', 'func': plane},
		(('N','M'),(1,),('N','M')):{'name': '3d-plane', 'func': plane},
		(('N','M'),('N','M'),(1,)):{'name': '3d-plane', 'func': plane},
	},
	('N','M',1): {
		((1,),('N',),('M',)):{'name': '3d-plane', 'func': plane},
		(('N',),(1,),('M',)):{'name': '3d-plane', 'func': plane},
		(('N',),('M',),(1,)):{'name': '3d-plane', 'func': plane},
	},
}

def get_arg_shapes(*args):
	shapes = []
	for a in args:
		if type(a) == np.ndarray:
			shape = a.shape
		else:
			try:
				shape = (len(a),)
			except:
				shape = (1,)
		shapes.append(shape)
	return shapes

def get_plot_key(out_shape, *arg_shapes):
	shapes_match = all([(a == out_shape) for a in list(arg_shapes)])
	nargs = len(arg_shapes)
	arg_dims = ''
	if nargs == 1:
		if len(out_shape) == 1:
			out_dim = 'N'
		else:
			out_dim = ('N', out_shape[-1])
	else:
		if len(out_shape) == 1:
			if out_shape[0] == 1:
				out_dim = 1,
			else:
				out_dim = 'N',
		elif len(out_shape) == 2:
			out_dim = 'N','M'
		elif len(out_shape) == 3:
			if 1 in out_shape:
				out_dim = 'N', 'M', 1
			else:
				out_dim = 'N','M','L'

	out_dim = tuple(out_dim)
	if shapes_match:
		arg_dims = tuple(nargs*[out_dim])
	else:
		if nargs == 1:
			if out_dim == ('N',2):
				if arg_shapes[0][0] == out_shape[0]:
					arg_dims = (('N',),)
			elif out_dim == ('N',3):
				if arg_shapes[0][0] == out_shape[0]:
					arg_dims = (('N',),)
		elif nargs == 2:
			if out_dim == ('N','M'):
				if (arg_shapes[0][0] in out_shape) & (arg_shapes[1][0] in out_shape):
					arg_dims = (('N',),('M',))
		elif nargs == 3:
			if out_dim == (1,):
				if len(set(arg_shapes)) == 1:
					if len(arg_shapes[0]) == 2:
						arg_dims = tuple(3*[('N','M')])
			elif out_dim == ('N','M'):
				arg_set = set([1])
				for arg in arg_shapes:
					arg_set.update(set(arg))
				if arg_set - set(out_shape) == set([1]):
					if arg_shapes[0] == (1,):
						arg_dims = ((1,), ('N','M'), ('N','M'))
					elif arg_shapes[1] == (1,):
						arg_dims = (('N','M'),(1,),('N','M',))
					elif arg_shapes[2] == (1,):
						arg_dims = (('N','M'),('N','M'),(1,))
			elif out_dim == ('N','M',1):
				arg_set = set([1])
				for arg in arg_shapes:
					arg_set.update(set(arg))
				if arg_shapes[0] == (1,):
					arg_dims = ((1,), ('N',), ('M',))
				elif arg_shapes[1] == (1,):
					arg_dims = (('N',),(1,),('M',))
				elif arg_shapes[2] == (1,):
					arg_dims = (('N',),('M',),(1,))
	if arg_dims == '':
		raise NotImplementedError('No way to handle out_shape {}, with arg shapes:{}'.format(out_shape, arg_shapes))
	return out_dim, arg_dims

plot_types = dict()
for out_shape, v in list(plot_dict.items()):
	for arg_shapes, v_ in list(v.items()):
		plot_key = get_plot_key(out_shape,*arg_shapes)
		if plot_key in plot_types:
			raise KeyError('plot_key already present {}'.format(plot_key))
		plot_types[plot_key] = [v_['name'], v_['func']]

plot_types = pd.DataFrame(plot_types).T
plot_types.index.set_names(['out_shape', 'arg_shapes'], inplace = True)
plot_types.columns = ['plot_type', 'function']

def test_plot_keys():
	for k in plot_types.to_dict(orient = 'index'):
		try:
			plot_dict[k[0]][k[1]]['name']
		except KeyError:
			print('could not find', k[0], k[1])
			raise

max_stars_repo_name: scivision/Kamodo

max_issues_repo_name: scivision/Kamodo

max_forks_repo_name: scivision/Kamodo


text:
'''
Queue implementation using a circular numpy array.
'''

import numpy as np
from stack import Stack

class QueueNP:
    def __init__(self, capacity=3):
        self.array = np.empty(capacity, dtype=np.uint8)
        self.length = capacity
        self.count = 0
        self.front = 0
        self.rear = 0

    def is_empty(self):
        return self.count == 0

    def is_full(self):
        return self.count == self.length

    def peek(self):
        return self.array[self.front]

    def increment_front(self):
        self.front = (self.front + 1) % self.length

    def increment_rear(self):
        self.rear = (self.rear + 1) % self.length

    def enqueue(self, item):
        self.array[self.rear] = item
        self.increment_rear()
        if self.is_full():
            self.increment_front()
        else:
            self.count += 1

    def dequeue(self):
        if self.is_empty(): raise IndexError()
        item =  self.array[self.front]
        self.array[self.front] = 0
        self.increment_front()
        self.count -= 1
        return item

    def reverse(self):
        stack = Stack()

        while not self.is_empty():
            stack.push(self.dequeue())

        while not stack.is_empty():
            self.enqueue(stack.pop())

    def __str__(self):
        lst = list(self.array)
        front, length  = self.front, self.length
        ans = []
        for _ in range(self.count):
            ans.append(lst[front])
            front = (front + 1) % length
        return f'{ans}'

def main():
    def prnq(q):
        print('\n\t q: ', q)

    def prndq(q):
        print('\n\t dequeue: ', q.dequeue())

    q = QueueNP()
    for x in range(1, 5):
        q.enqueue(x)
        prnq(q)

    prndq(q)
    prnq(q)
    print()
    quit()

    for y in range(55, 66):
        q.enqueue(y)
        print('\t q.enqueue', y, ':  ', q)
        print('\t dequeue: ', q.dequeue(), ':  ', q)

main()

max_stars_repo_name: jwbat/python-datastructures

max_issues_repo_name: jwbat/python-datastructures

max_forks_repo_name: jwbat/python-datastructures


text:
'''
Created on 2009-08-11

@author: malem303
'''
import unittest

from imugrabber.algorithms import fong_accelero, utils, statistics
from imugrabber.algorithms import io
import os

import scipy as sp
from numpy import testing

class FongTests(unittest.TestCase):

    def setUp(self):

        self.misalignmentsAndScales = sp.array([[  4.08269136e-03,  -1.96002082e-05,   1.16692771e-04],
                                                 [ -6.73123099e-06,   3.86658837e-03,  -2.77361987e-04],
                                                 [ -6.43895175e-05,   2.91260930e-04,  3.93614477e-03]])
        self.biases = sp.array([[ 604.00283039],
                             [ 480.33539568],
                             [ 522.23054001]])

    def test_fit_parameters(self):

        dataSet = io.float_columns_from_CSV(csvFileName = "FongTests.csv", path = io._test_path)

        x, y, z = dataSet['avgX'], dataSet['avgY'], dataSet['avgZ']

        measures = utils.build_measures_matrix(x, y, z)

        misalignmentsAndScales, biases  = fong_accelero.fit(measures)

        testing.assert_almost_equal(misalignmentsAndScales, self.misalignmentsAndScales)

        testing.assert_almost_equal(biases, self.biases)

if __name__ == "__main__":
    #import sys;sys.argv = ['', 'Test.testName']
    unittest.main()

max_stars_repo_name: maxlem/AVRCpp

max_issues_repo_name: maxlem/AVRCpp

max_forks_repo_name: maxlem/AVRCpp

reason: from numpy,import scipy


text:
import numpy as np
from numba import jit
import sample_households
import sample_comorbidities
import functions
import pickle

"""3. RUN SIMULATION."""

def run_complete_simulation(seed, country, contact_matrix, p_mild_severe, p_severe_critical, p_critical_death, mean_time_to_isolate_factor, lockdown_factor_age,
    p_infect_household, fraction_stay_home, params, load_population=False):
    '''
    Runs simulation with given parameters (see run_simulation for details).
    If load_population is true, reads in the simulated population for the given country and size.
    Otherwise, creates and saves a new population.
    '''
    n = int(params['n'])
    n_ages = int(params['n_ages'])
    np.random.seed(seed)
    if load_population:
        print('loading')
        age, households, diabetes, hypertension, age_groups = pickle.load(open('{}_population_{}.pickle'.format(country, n), 'rb'))
    else:
        if country == "Italy":
          households, age = sample_households.sample_households_italy(n)
        elif country == "Germany":
          households, age = sample_households.sample_households_germany(n)
        elif country == "UK":
          households, age = sample_households.sample_households_uk(n)
        elif country == "Spain":
          households, age = sample_households.sample_households_spain(n)
        elif country == "China":
          households, age = sample_households.sample_households_china(n)
        else:
          households, age = sample_households.sample_households_un(n, country)
        age_groups = tuple([np.where(age == i)[0] for i in range(0, n_ages)])
        diabetes, hypertension = sample_comorbidities.sample_joint_comorbidities(age, country)

        data = (age, households, diabetes, hypertension, age_groups)
        with open('{}_population_{}.pickle'.format(country, n), 'wb') as f:
            pickle.dump(data, f)

    print('starting simulation')
    return run_model(seed, households, age, age_groups, diabetes, hypertension, contact_matrix, p_mild_severe, p_severe_critical, p_critical_death, mean_time_to_isolate_factor, lockdown_factor_age, p_infect_household, fraction_stay_home, params)

@jit(nopython=True)
def get_isolation_factor(age, mean_time_to_isolate_factor):
    for i in range(len(mean_time_to_isolate_factor)):
        if age >= mean_time_to_isolate_factor[i, 0] and age <= mean_time_to_isolate_factor[i, 1]:
            return mean_time_to_isolate_factor[i, 2]
    return 1

@jit(nopython=True)
def get_lockdown_factor_age(age, lockdown_factor_age):
    for i in range(len(lockdown_factor_age)):
        if age >= lockdown_factor_age[i, 0] and age <= lockdown_factor_age[i, 1]:
            return lockdown_factor_age[i, 2]
    return 1

@jit(nopython=True)
def do_contact_tracing(i, infected_by, p_trace_outside, Q, S, t, households, p_trace_household, Documented, time_documented, traced):
    '''
    Recursively implements tracing the tree of infectious defined by infected_by.
    NOT CURRENTLY USED in the preprint. There is no validation of this procedure
    for any particular example.
    '''
    #trace contacts within the household
    for j in range(households.shape[1]):
        contact = households[i, j]
        if contact == -1:
            break
        if not S[t-1, contact] and not traced[contact] and np.random.rand() < p_trace_household:
            Q[t, contact] = True
            Documented[t, contact] = True
            traced[contact] = True
            time_documented[contact] = t
            do_contact_tracing(contact, infected_by, p_trace_outside, Q, S, t, households, p_trace_household, Documented, time_documented, traced)
    #trace outside of household contacts
    for j in range(infected_by.shape[1]):
        contact = infected_by[i, j]
        if contact == -1:
            break
        if np.random.rand() < p_trace_outside:
            Q[t, contact] = True
            Documented[t, contact] = True
            time_documented[contact] = t
            traced[contact] = True
            do_contact_tracing(contact, infected_by, p_trace_outside, Q, S, t, households, p_trace_household, Documented, time_documented, traced)

@jit(nopython=True)
def run_model(seed, households, age, age_groups, diabetes, hypertension, contact_matrix, p_mild_severe, p_severe_critical, p_critical_death, mean_time_to_isolate_factor, lockdown_factor_age, p_infect_household, fraction_stay_home, params):
    print('run_model')
    """Run the SEIR model to completion.

    Args:
        seed (int): Random seed.
        households (int n x max_household_size matrix): Household structure (adjacency list format, each row terminated with -1s)
        age (int vector of length n): Age of each individual.
        diabetes (bool vector of length n): Diabetes state of each individual.
        hypertension (bool vector of length n): Hypertension state of each
            individual.
        contact_matrix (float matrix n_ages x n_ages): expected number of daily contacts between each pair of age groups
        p_mild_severe (float matrix n_ages x 2 x 2): probability of mild->severe transition for each age/diabetes/hypertension status
        p_severe_critical (float matrix n_ages x 2 x 2): as above but for severe->critical
        p_critical_death (float matrix n_ages x 2 x 2): as above but for critical->death
        mean_time_to_isolate_factor (float vector n_ages): scaling applied to the mean time to isolate for mild cases per age group
        lockdown_factor_age (float vector n_ages): per-age reduction in contact during lockdown. Not currently used.
        p_infect_household (float vector n): probability for each individual to infect each household member per day.
        fraction_stay_home (float vector n_ages): fraction of each age group assigned to shelter in place
        params: dict with remaining scalar parameters

    Returns:
        S (bool T x n matrix): Matrix where S[i][j] represents
            whether individual i was in the Susceptible state at time j.
        E (bool T x n matrix): same for Exposed state.
        Mild (bool T x n matrix): same for Mild state.
        Severe (bool T x n matrix): same for Severe state.
        Critical (bool T x n matrix): same for Critical state.
        R (bool T x n matrix): same for Recovered state.
        D (bool T x n matrix): same for Dead state.
        Q (bool T x n matrix): same for Quarantined state.
        num_infected_by (n vector): num_infected_by[i] is the number of individuals
            infected by individual i. -1 if they never became infectious
        time_documented (n vector): time step when each individual became a documented case, 0 if never documented
        time_to_activation (n vector): incubation time drawn for this individual, 0 if no time drawn
        time_to_death (n vector):  time-to-event for critical -> death for this individual, 0 if no time drawn
        time_to_recovery (n vector): time-to-event for infectious -> recovery for this individual, 0 if no time drawn
        time_critical (n vector): time step that this individual entered the critical state, 0 if they never entered
        time_exposed (n vector): time step that this individual became asymptomatically infected, -1 if this never happend
        num_infected_asympt (n vector): number of others infected by this individual while asymptomatic, -1 if never became infectious
        age (n vector): age of each individual
        time_infected (n vector): time step that this individual became mildly infectious, 0 if they never became infectious
        time_to_severe (n vector): time-to-event between mild and severe cases for this individual, 0 if never drawn
    """
    time_to_activation_mean = params['time_to_activation_mean']
    time_to_activation_std = params['time_to_activation_std']
    mean_time_to_death = params['mean_time_to_death']
    mean_time_critical_recovery = params['mean_time_critical_recovery']
    mean_time_severe_recovery = params['mean_time_severe_recovery']
    mean_time_to_severe = params['mean_time_to_severe']
    mean_time_mild_recovery = params['mean_time_mild_recovery']
    mean_time_to_critical = params['mean_time_to_critical']
    p_documented_in_mild = params['p_documented_in_mild']
    mean_time_to_isolate_asympt = params['mean_time_to_isolate_asympt']
    asymptomatic_transmissibility = params['asymptomatic_transmissibility']
    p_infect_given_contact = params['p_infect_given_contact']
    T = int(params['T'])
    initial_infected_fraction = params['initial_infected_fraction']
    t_lockdown = int(params['t_lockdown'])
    lockdown_factor = params['lockdown_factor']
    mean_time_to_isolate = params['mean_time_to_isolate']
    n = int(params['n'])
    n_ages = int(params['n_ages'])
    contact_tracing = bool(params['contact_tracing'])
    p_trace_outside = params['p_trace_outside']
    p_trace_household = params['p_trace_household']
    t_tracing_start = int(params['t_tracing_start'])
    t_stayinghome_start = int(params['t_stayhome_start'])
    if contact_tracing:
        tracing_enabled = True
        contact_tracing = False

    np.random.seed(seed)
    max_household_size = households.shape[1]
    S = np.zeros((T, n), dtype=np.bool8)
    E = np.zeros((T, n), dtype=np.bool8)
    Mild = np.zeros((T, n), dtype=np.bool8)
    Documented = np.zeros((T, n), dtype=np.bool8)
    Severe = np.zeros((T, n), dtype=np.bool8)
    Critical = np.zeros((T, n), dtype=np.bool8)
    R = np.zeros((T, n), dtype=np.bool8)
    D = np.zeros((T, n), dtype=np.bool8)
    Q = np.zeros((T, n), dtype=np.bool8)
    traced = np.zeros((n), dtype=np.bool8)
    #whether each individual is assigned to shelter in place
    Home_real = np.zeros(n, dtype=np.bool8)
    Home_real[:] = False
    for i in range(n_ages):
        matches = np.where(age == i)[0]
        if matches.shape[0] > 0:
            to_stay_home = np.random.choice(matches, int(fraction_stay_home[i]*matches.shape[0]), replace=False)
            Home_real[to_stay_home] = True
    #no one shelters in place until t_stayhome_start
    dummy_Home = np.zeros(n, dtype=np.bool8)
    dummy_Home[:] = False
    Home = dummy_Home
    initial_infected = np.random.choice(n, int(initial_infected_fraction*n), replace=False)
    S[0] = True
    E[0] = False
    R[0] = False
    D[0] = False
    Mild[0] = False
    Documented[0]=False
    Severe[0] = False
    Critical[0] = False

    infected_by = np.zeros((n, 100), dtype=np.int32)
    infected_by[:] = -1

    time_exposed = np.zeros(n)
    time_infected = np.zeros(n)
    time_severe = np.zeros(n)
    time_critical = np.zeros(n)
    time_documented=np.zeros(n)
    time_exposed[:] = -1
    #total number of infections caused by every individual, -1 if never become infectious
    num_infected_by = np.zeros(n)
    num_infected_by_outside = np.zeros(n, dtype=np.int32)
    num_infected_asympt = np.zeros(n)
    num_infected_by[:] = -1
    num_infected_by_outside[:] = -1
    num_infected_asympt[:] = -1
    time_to_severe = np.zeros(n)
    time_to_recovery = np.zeros(n)
    time_to_critical = np.zeros(n)
    time_to_death = np.zeros(n)
    time_to_isolate = np.zeros(n)
    time_to_activation = np.zeros(n)
    #initialize values for individuals infected at the starting step
    for i in range(initial_infected.shape[0]):
        E[0, initial_infected[i]] = True
        S[0, initial_infected[i]] = False
        time_exposed[initial_infected[i]] = 0
        num_infected_by[initial_infected[i]] = 0
        num_infected_by_outside[initial_infected[i]] = 0
        num_infected_asympt[initial_infected[i]] = 0
        time_to_activation[initial_infected[i]] = functions.threshold_log_normal(time_to_activation_mean, time_to_activation_std)

    print('Initialized finished')
    print('mean_time_to_isolate',mean_time_to_isolate)
    for t in range(1, T):
        if t % 10 == 0:
            print(t,"/",T)
        if t == t_lockdown:
        #   implements a reduction in contact per age group (instead of a single factor)
        #   applied to the whole population. Currently not used.
        #    for i in range(contact_matrix.shape[0]):
        #        for j in range(contact_matrix.shape[1]):
        #            contact_matrix[i,j] = contact_matrix[i,j]/(get_lockdown_factor_age(i,lockdown_factor_age)*get_lockdown_factor_age(j,lockdown_factor_age))
            contact_matrix = contact_matrix/lockdown_factor
        if t == t_tracing_start and tracing_enabled:
            contact_tracing = True
        if t == t_stayinghome_start:
            Home = Home_real
        S[t] = S[t-1]
        E[t] = E[t-1]
        Mild[t] = Mild[t-1]
        Documented[t]=Documented[t-1]
        Severe[t] = Severe[t-1]
        Critical[t] = Critical[t-1]
        R[t] = R[t-1]
        D[t] = D[t-1]
        Q[t] = Q[t-1]
        for i in range(n):
            #exposed -> (mildly) infected
            if E[t-1, i]:
                if t - time_exposed[i] == time_to_activation[i]:
                    Mild[t, i] = True
                    time_infected[i] = t
                    E[t, i] = False
                    #draw whether they will progress to severe illness
                    if np.random.rand() < p_mild_severe[age[i], diabetes[i], hypertension[i]]:
                        time_to_severe[i] = functions.threshold_exponential(mean_time_to_severe)
                        time_to_recovery[i] = np.inf
                    #draw time to recovery
                    else:
                        time_to_recovery[i] = functions.threshold_exponential(mean_time_mild_recovery)
                        time_to_severe[i] = np.inf
                    #draw time to isolation
                    time_to_isolate[i] = functions.threshold_exponential(mean_time_to_isolate*get_isolation_factor(age[i], mean_time_to_isolate_factor))
                    if time_to_isolate[i] == 0:
                        Q[t, i] = True
            #symptomatic individuals
            if (Mild[t-1, i] or Severe[t-1, i] or Critical[t-1, i]):
                #recovery
                if t - time_infected[i] == time_to_recovery[i]:
                    R[t, i] = True
                    Mild[t, i] = Severe[t, i] = Critical[t, i] = Q[t, i] = False
                    continue
                if Mild[t-1, i] and not Documented[t-1, i]:
                    #mild cases are documented with some probability each day
                    if np.random.rand() < p_documented_in_mild:
                        Documented[t, i] = True
                        time_documented[i] = t
                        traced[i] = True
                        #do contact tracing (if enabled) whenever a case becomes documented
                        if contact_tracing:
                            Q[t, i] = True
                            do_contact_tracing(i, infected_by, p_trace_outside, Q, S, t, households, p_trace_household, Documented, time_documented, traced)
                #progression between infection states
                if Mild[t-1, i] and t - time_infected[i] == time_to_severe[i]:
                    Mild[t, i] = False
                    Severe[t, i] = True
                    #assume that severe cases are always documented
                    if not Documented[t-1, i]:
                        Documented[t, i] = True
                        time_documented[i] = t
                        traced[i] = True
                        #assume that contact tracing is always started for severe cases
                        #(if enabled)
                        if contact_tracing:
                            Q[t, i] = True
                            do_contact_tracing(i, infected_by, p_trace_outside, Q, S, t, households, p_trace_household, Documented, time_documented, traced)
                    Q[t, i] = True
                    time_severe[i] = t
                    if np.random.rand() < p_severe_critical[age[i], diabetes[i], hypertension[i]]:
                        time_to_critical[i] = functions.threshold_exponential(mean_time_to_critical)
                        time_to_recovery[i] = np.inf
                    else:
                        time_to_recovery[i] = functions.threshold_exponential(mean_time_severe_recovery) + time_to_severe[i]
                        time_to_critical[i] = np.inf
                elif Severe[t-1, i] and t - time_severe[i] == time_to_critical[i]:
                    Severe[t, i] = False
                    Critical[t, i] = True
                    time_critical[i] = t
                    if np.random.rand() < p_critical_death[age[i], diabetes[i], hypertension[i]]:
                        time_to_death[i] = functions.threshold_exponential(mean_time_to_death)
                        time_to_recovery[i] = np.inf
                    else:
                        time_to_recovery[i] = functions.threshold_exponential(mean_time_critical_recovery) + time_to_severe[i] + time_to_critical[i]
                        time_to_death[i] = np.inf
                #risk of mortality for critically ill patients
                elif Critical[t-1, i]:
                    if t - time_critical[i] == time_to_death[i]:
                        Critical[t, i] = False
                        Q[t, i] = False
                        D[t, i] = True
            if E[t-1, i] or Mild[t-1, i] or Severe[t-1, i] or Critical[t-1, i]:
                #not isolated: either enter isolation or infect others
                if not Q[t-1, i]:
                    #isolation
                    if not E[t-1, i] and t - time_infected[i] == time_to_isolate[i]:
                        Q[t, i] = True
                        continue
                    if E[t-1, i] and t - time_exposed[i] == time_to_isolate[i]:
                        Q[t, i] = True
                        continue
                    #infect within family
                    for j in range(max_household_size):
                        if households[i,j] == -1:
                            break
                        contact = households[i,j]
                        infectiousness = p_infect_household[i]
                        if E[t-1, i]:
                            infectiousness *= asymptomatic_transmissibility
                        if S[t-1, contact] and np.random.rand() < infectiousness:
                                E[t, contact] = True
                                num_infected_by[contact] = 0
                                num_infected_by_outside[contact] = 0
                                num_infected_asympt[contact] = 0
                                S[t, contact] = False
                                time_to_isolate[contact] = functions.threshold_exponential(mean_time_to_isolate_asympt*get_isolation_factor(age[contact], mean_time_to_isolate_factor))
                                if time_to_isolate[contact] == 0:
                                    Q[t, contact] = True
                                time_exposed[contact] = t
                                time_to_activation[contact] = functions.threshold_log_normal(time_to_activation_mean, time_to_activation_std)
                                num_infected_by[i] += 1
                                if E[t-1, i]:
                                    num_infected_asympt[i] += 1
                    #infect across families
                    if not Home[i]:
                        infectiousness = p_infect_given_contact
                        #lower infectiousness for asymptomatic individuals
                        if E[t-1, i]:
                            infectiousness *= asymptomatic_transmissibility
                        #draw a Poisson-distributed number of contacts for each age group
                        for contact_age in range(n_ages):
                            if age_groups[contact_age].shape[0] == 0:
                                continue
                            num_contacts = np.random.poisson(contact_matrix[age[i], contact_age])
                            for j in range(num_contacts):
                                #if the contact becomes infected, handle bookkeeping
                                if np.random.rand() < infectiousness:
                                    contact = np.random.choice(age_groups[contact_age])
                                    if S[t-1, contact] and not Home[contact]:
                                        E[t, contact] = True
                                        num_infected_by[contact] = 0
                                        num_infected_by_outside[contact] = 0
                                        num_infected_asympt[contact] = 0
                                        S[t, contact] = False
                                        time_to_isolate[contact] = functions.threshold_exponential(mean_time_to_isolate_asympt*get_isolation_factor(age[contact], mean_time_to_isolate_factor))
                                        if time_to_isolate[contact] == 0:
                                            Q[t, contact] = True
                                        time_exposed[contact] = t
                                        time_to_activation[contact] = functions.threshold_log_normal(time_to_activation_mean, time_to_activation_std)
                                        num_infected_by[i] += 1
                                        infected_by[i, num_infected_by_outside[i]] = contact
                                        num_infected_by_outside[i] += 1
                                        if E[t-1, i]:
                                            num_infected_asympt[i] += 1

    return S, E, Mild, Documented, Severe, Critical, R, D, Q, num_infected_by,time_documented, time_to_activation, time_to_death, time_to_recovery, time_critical, time_exposed, num_infected_asympt, age, time_infected, time_to_severe


text:
"""
:Author: Arthur Goldberg <Arthur.Goldberg@mssm.edu>
:Date: 2019-11-21
:Copyright: 2019, Karr Lab
:License: MIT
"""

import math
import numpy
import random
import timeit
import unicodedata
import unittest

from wc_sim.utils.time_ordered_list import DoublyLinkedNode, TimeOrderedList, LinearInterpolationList

class TestDoublyLinkedNode(unittest.TestCase):

    def test_node(self):

        def is_connected(test_case, left, right):
            test_case.assertEqual(left._next._prev, left)
            test_case.assertEqual(right._prev._next, right)
            test_case.assertEqual(left._next, right)
            test_case.assertEqual(left, right._prev)

        n = DoublyLinkedNode(3, 'x')
        self.assertIn('_time', str(n))
        self.assertIn('_data', str(n))
        self.assertIn('is None', str(n))
        o = DoublyLinkedNode(4, 'y')

        # test link_on_left
        n.link_on_left(o)
        is_connected(self, n, o)
        self.assertIn('is not None', str(n))
        self.assertIn('is not None', str(o))

        # test is_equal
        p = DoublyLinkedNode(4, 'y')
        self.assertTrue(p.is_equal(p))
        self.assertTrue(p.is_equal(o))
        self.assertTrue(o.is_equal(p))
        p._time += 1
        self.assertTrue(p.is_equal(p))
        self.assertFalse(p.is_equal(o))
        self.assertFalse(p.is_equal('hi'))

def invariants(test_case, time_ordered_list):
    # head tail invariant
    test_case.assertEqual(time_ordered_list._head._next._prev, time_ordered_list._head)
    test_case.assertEqual(time_ordered_list._tail._prev._next, time_ordered_list._tail)

    # ordered invariant
    node = time_ordered_list._head
    while node._next is not None:
        test_case.assertTrue(node._time < node._next._time)
        node = node._next

class TestTimeOrderedList(unittest.TestCase):

    def test_time_ordered_list(self):

        ### test internal methods ###
        tol = TimeOrderedList()
        invariants(self, tol)

        # test _insert_after
        value = 'x'
        n = DoublyLinkedNode(1, value)
        tol._insert_after(tol._head, n)
        self.assertEqual(tol.len(), 1)
        invariants(self, tol)
        self.assertEqual(tol._head._next._data, value)
        self.assertEqual(tol._tail._prev._data, value)

        self.assertIn('nodes:', str(tol))
        self.assertIn(value, str(tol))

        # test _delete
        self.assertEqual(tol._delete(n), n)
        invariants(self, tol)
        self.assertTrue(tol.is_empty())

        self.assertEqual(tol._delete(n), None)
        invariants(self, tol)

        # test _check_time
        bad_time = 'x'
        with self.assertRaisesRegex(ValueError, "time '.*' isn't a number"):
            tol._check_time('method', bad_time)
        bad_time = float('inf')
        with self.assertRaisesRegex(ValueError, 'cannot .* at time'):
            tol._check_time('method', bad_time)
        bad_time = float('NaN')
        with self.assertRaisesRegex(ValueError, 'cannot .* at time'):
            tol._check_time('method', bad_time)
        good_time = 3
        self.assertEqual(tol._check_time('method', good_time), None)
        good_time = 1.34e4
        self.assertEqual(tol._check_time('method', good_time), None)

        ### test external methods ###
        tol = TimeOrderedList()
        self.assertTrue(tol.is_empty())
        self.assertEqual(tol.len(), 0)
        time = 1
        tol.insert(time, 3)
        invariants(self, tol)

        tol.clear()
        self.assertTrue(tol.is_empty())
        invariants(self, tol)

        # test insert
        value = 'z'
        tol.insert(time, value)
        invariants(self, tol)

        with self.assertRaisesRegex(ValueError, 'cannot insert at time'):
            tol.insert(float("inf"), value)
        with self.assertRaisesRegex(ValueError, 'cannot insert at time'):
            tol.insert(float("NaN"), value)
        time = 3
        tol.insert(time, value)
        with self.assertRaisesRegex(ValueError, "time .* already in queue"):
            tol.insert(time, value)

        times = [4, -3, 2, 10, 6]
        for time in times:
            tol.insert(time, time)
        invariants(self, tol)

        # test find
        for time in times:
            n = tol.find(time)
            self.assertEqual(n._data, time)
        n = tol.find(max(times) + 1)
        self.assertEqual(n._time, max(times))
        n = tol.find(min(times) - 1)
        self.assertEqual(n._time, -float('inf'))
        with self.assertRaisesRegex(ValueError, "time .* isn't a number"):
            tol.find('h')

        # test delete
        len = tol.len()
        for time in times:
            n = tol.delete(time)
            self.assertTrue(isinstance(n, DoublyLinkedNode))
            self.assertEqual(tol.len(), len - 1)
            len -= 1
        with self.assertRaisesRegex(ValueError, "no node with time .* in queue"):
            tol.delete(times[0])

        # test gc_old_data
        # test empty list
        tol = TimeOrderedList()
        self.assertEqual(tol.gc_old_data(1), 0)
        invariants(self, tol)

        maximum = 10
        for t in reversed(range(maximum)):
            tol.insert(t, t)
        invariants(self, tol)

        # try to keep more than available
        self.assertEqual(tol.gc_old_data(1, min_to_keep=maximum + 1), 0)
        self.assertEqual(tol.len(), maximum)
        invariants(self, tol)

        # delete some
        time_cutoff = 4
        self.assertEqual(tol.gc_old_data(time_cutoff), time_cutoff)
        invariants(self, tol)
        self.assertEqual(tol.len(), maximum - time_cutoff)

        # delete none because time cutoff is too early
        self.assertEqual(tol.gc_old_data(time_cutoff), 0)
        invariants(self, tol)

        # limit deletions with min_to_keep
        self.assertEqual(tol.gc_old_data(maximum, min_to_keep=tol.len()-1), 1)
        invariants(self, tol)

        # delete everything
        self.assertEqual(tol.gc_old_data(maximum, min_to_keep=0), 5)
        self.assertEqual(tol.len(), 0)
        invariants(self, tol)

    def test_time_ordered_list_performance(self):
        # measure the performance of TimeOrderedList.find()
        print()
        def measure_time_ordered_list_performance(num_items, timeit_num=100, num_measurements=100):
            tol = TimeOrderedList()
            value = 'x'
            # reverse to build rapidly
            for time in reversed(range(num_items)):
                tol.insert(time, value)
            find_times = []
            for _ in range(num_measurements):
                time = num_items * random.random()
                # find is the expensive method
                find_expr = f'tol.find({time})'
                m = timeit.timeit(find_expr, number=timeit_num, globals=locals())
                find_times.append(m)
            factor = 1_000_000 / timeit_num
            # time in microsec
            mean_time = factor * numpy.mean(find_times)
            print(f"{num_items}\t{mean_time:.2f}")

        mu = unicodedata.lookup('GREEK SMALL LETTER MU')
        print(f"# TOL items\tE[time(find)] ({mu}sec)")
        factors_of_ten = 4
        all_num_items = [10 ** factor_of_ten for factor_of_ten in range(factors_of_ten)]
        for num_items in all_num_items:
            measure_time_ordered_list_performance(num_items)

class TestLinearInterpolationList(unittest.TestCase):

    def test_read(self):

        # test read
        tol = LinearInterpolationList()
        first_2_values = 5
        last_value = 10
        time_n_values = [(1, first_2_values), (2, first_2_values), (3, last_value), ]
        for time, value in time_n_values:
            tol.insert(time, value)
        invariants(self, tol)
        for time, value in time_n_values:
            self.assertEqual(tol.read(time), value)
        self.assertTrue(math.isnan(tol.read(0.5)))
        self.assertEqual(tol.read(1.01), first_2_values)
        self.assertEqual(tol.read(1.5), first_2_values)
        self.assertEqual(tol.read(2.5), 0.5*(first_2_values + last_value))
        self.assertEqual(tol.read(3.5), last_value)


text:
import os
import glob
import sys
import time
from optparse import OptionParser

import dask.dataframe as dd
import h5py
import pandas as pd
import numpy as np
import scipy as sp
import scipy.stats as scst

from struct_lmm import StructLMM
import limix
import qtl_output
import qtl_loader_utils
import qtl_parse_args
import qtl_utilities as utils
from qtl_snp_qc import do_snp_qc
from sklearn.preprocessing import Imputer

def struct_lmm(geno_df,
                pheno,
                env,
                covs = None,
                rhos = None,
                no_association_test = False,
                no_interaction_test = False):
    """
    Utility function to run StructLMM

    Parameters
    ----------
    geno_df: (`N`, `S`) pandas.DataFrame
        genotype matrix for `N` samples, `S` SNPis
    pheno : (`N`, 1) ndarray
        phenotype vector
    env : (`N`, `K`)
          Environmental matrix (indviduals by number of environments)
    covs : (`N`, L) ndarray
        fixed effect design for covariates `N` samples and `L` covariates.
    rhos : list
        list of ``rho`` values.
        ``rho=0`` correspond to no persistent effect (only GxE);
        ``rho=1`` corresponds to only persitent effect (no GxE);
        By default, ``rho=[0, 0.2, 0.4, 0.6, 0.8, 1.]``
    no_association_test : bool
        if True the association test is not considered.
        The default value is False.
    no_interaction_test : bool
        if True the interaction test is not considered.
        The default value is False.

    Returns
    -------
    res : *:class:`pandas.DataFrame`*
        contains pv of joint test (if no_association_test is False),
        pv of interaction test (if no_interaction_test is False).
    """
    #import pdb; pdb.set_trace()
    if covs is None:
        covs = sp.ones((env.shape[0], 1))
    if rhos is None:
        rhos = [0, .2, .4, .6, .8, 1.]
    if not no_association_test:
        # slmm fit null
        slmm = StructLMM(pheno, env, W=env, rho_list=rhos)
        null = slmm.fit_null(F=covs, verbose=False)
    if not no_interaction_test:
        # slmm int
        slmm_int = StructLMM(pheno, env, W=env, rho_list=[0])
    t0 = time.time()
    #import pdb; pdb.set_trace()
    res = pd.DataFrame(data = [], index = geno_df.columns.values)
    pv = sp.zeros(geno_df.shape[1])
    pv_int = sp.zeros(geno_df.shape[1])

    for snp in xrange(geno_df.shape[1]):
        x = geno_df.values[:, [snp]]
        if not no_association_test:
            # association test
            _p, _rho = slmm.score_2_dof(x) #second arg: optimal rho
            pv[snp] = _p
        if not no_interaction_test:
            # interaction test
            covs1 = sp.hstack((covs, x))
            null = slmm_int.fit_null(F=covs1, verbose=False)
            _p, _rho = slmm_int.score_2_dof(x)
            pv_int[snp] = _p
    # add pvalues to res
    if not no_association_test:
        res['pv'] = pv
    if not no_interaction_test:
        res['pv_int'] = pv_int
    t = time.time() - t0
    #print '%.2f s elapsed' % t
    return res

def run_structLMM_QTL_analysis(pheno_filename, anno_filename, env_filename, geno_prefix, plinkGenotype, output_dir, window_size = 250000,
                        min_maf = 0.05, min_hwe_P = 0.001, min_call_rate = 0.95, association_mode = True, cis_mode = True,
                        blocksize = 1000, skipAutosomeFiltering = False, gaussianize_method = None, minimum_test_samples = 10,
                        seed = np.random.randint(40000), n_perm = 0, write_permutations = False, relatedness_score = 0.95,
                        feature_variant_covariate_filename = None, snps_filename = None, feature_filename = None, snp_feature_filename = None,
                        genetic_range = 'all', covariates_filename = None, kinship_filename = None, sample_mapping_filename = None,
                        extended_anno_filename = None):
    fill_NaN = Imputer(missing_values = np.nan, strategy = 'mean', axis = 0)
    print('Running structLMM QTL analysis.')
    '''Core function to take input and run QTL tests on a given chromosome.'''

    [phenotype_df, kinship_df, covariate_df, environment_df, sample2individual_df, complete_annotation_df, annotation_df, snp_filter_df, snp_feature_filter_df,
    genetically_unique_individuals, minimum_test_samples, feature_list,bim,fam,bed, chromosome, selectionStart, selectionEnd, feature_variant_covariate_df]=\
    utils.run_structLMM_QTL_analysis_load_intersect_phenotype_environments_covariates_kinship_sample_mapping(pheno_filename = pheno_filename, anno_filename = anno_filename,
                env_filename = env_filename, geno_prefix = geno_prefix, plinkGenotype = plinkGenotype, cis_mode = cis_mode, association_mode = association_mode, skipAutosomeFiltering = skipAutosomeFiltering,
                minimum_test_samples = minimum_test_samples, relatedness_score = relatedness_score, snps_filename = snps_filename,
                feature_filename = feature_filename, snp_feature_filename = snp_feature_filename, selection = genetic_range, covariates_filename = covariates_filename,
                kinship_filename = kinship_filename, sample_mapping_filename = sample_mapping_filename, extended_anno_filename = extended_anno_filename,
                feature_variant_covariate_filename = feature_variant_covariate_filename)
    #import pdb; pdb.set_trace()
    if(feature_list == None or len(feature_list) == 0):
        print ('No features to be tested.')
        sys.exit()

    #Open output files
    qtl_loader_utils.ensure_dir(output_dir)
    if not selectionStart is None :
        output_writer = qtl_output.hdf5_writer(output_dir + 'qtl_results_{}_{}_{}.h5'.format(chromosome, selectionStart, selectionEnd))
    else :
        output_writer = qtl_output.hdf5_writer(output_dir + 'qtl_results_{}.h5'.format(chromosome))
    if(write_permutations):
        if not selectionStart is None :
            permutation_writer = qtl_output.hdf5_permutations_writer(output_dir + 'perm_results_{}_{}_{}.h5'.format(chromosome, selectionStart, selectionEnd), n_perm)
        else :
            permutation_writer = qtl_output.hdf5_permutations_writer(output_dir + 'perm_results_{}.h5'.format(chromosome), n_perm)

    #Arrays to store indices of snps tested and pass and fail QC SNPs for features without missingness.
    tested_snp_idxs = []
    pass_qc_snps_all = []
    fail_qc_snps_all = []
    fail_qc_features = []
    alpha_params = []
    beta_params = []
    n_samples = []
    n_e_samples = []
    currentFeatureNumber = 0
    snpQcInfoMain = None
    for feature_id in feature_list:
        snpQcInfo = None
        currentFeatureNumber+= 1
        if (len(phenotype_df.loc[feature_id,:]) < minimum_test_samples):
            print("Feature: "+feature_id+" not tested: not enough samples do QTL test.")
            continue
        data_written = False
        snpQuery = utils.do_snp_selection(feature_id, complete_annotation_df, bim, cis_mode, window_size, skipAutosomeFiltering)
        snp_cov_df = None
        if(feature_variant_covariate_df is not None):
            if(feature_id in feature_variant_covariate_df['feature'].values):
                covariateSnp = feature_variant_covariate_df['snp_id'].values[feature_variant_covariate_df['feature'] == feature_id]
                if(any(i in bim['snp'].values for i in covariateSnp)):
                    snpQuery_cov = bim.loc[bim['snp'].map(lambda x: x in list(covariateSnp)),:]
                    snp_cov_df_t = pd.DataFrame(data = bed[snpQuery_cov['i'].values,:].compute().transpose(), index = fam.index, columns = snpQuery_cov['snp'],)
                    snp_cov_df = pd.DataFrame(fill_NaN.fit_transform(snp_cov_df_t))
                    snp_cov_df.index = snp_cov_df_t.index
                    snp_cov_df.columns = snp_cov_df_t.columns
                    snp_cov_df_t = None

        if (len(snpQuery) != 0) and (snp_filter_df is not None):
            snpQuery = snpQuery.loc[snpQuery['snp'].map(lambda x: x in list(map(str, snp_filter_df.index)))]

        if (len(snpQuery) != 0) and (snp_feature_filter_df is not None):
            snpQuery = snpQuery.loc[snpQuery['snp'].map(lambda x: x in list(snp_feature_filter_df['snp_id'].loc[snp_feature_filter_df['feature'] == feature_id]))]

        if len(snpQuery) != 0:
            phenotype_ds = phenotype_df.loc[feature_id]
            contains_missing_samples = any(~np.isfinite(phenotype_ds))

            if(contains_missing_samples):
                print ('Feature: ' + feature_id + ' contains missing data.')
                phenotype_ds.dropna(inplace = True)

            '''select indices for relevant individuals in genotype matrix
            These are not unique. NOT to be used to access phenotype/covariates data
            '''
            individual_ids = sample2individual_df.loc[phenotype_ds.index, 'iid'].values
            sample2individual_feature = sample2individual_df.loc[phenotype_ds.index]

            if(contains_missing_samples):
                tmp_unique_individuals = genetically_unique_individuals
                genetically_unique_individuals = utils.get_unique_genetic_samples(kinship_df.loc[individual_ids, individual_ids], relatedness_score);

            if phenotype_ds.empty or (len(genetically_unique_individuals) < minimum_test_samples) :
                print("Feature: " + feature_id + " not tested: not enough samples do QTL test.")
                fail_qc_features.append(feature_id)
                if contains_missing_samples:
                    genetically_unique_individuals = tmp_unique_individuals
                continue
            elif np.var(phenotype_ds.values) == 0:
                print("Feature: " + feature_id + " has no variance in selected individuals.")
                fail_qc_features.append(feature_id)
                if contains_missing_samples:
                    genetically_unique_individuals = tmp_unique_individuals
                continue

            #If no missing samples we can use the previous SNP Qc information before actually loading data.
            #This allowes for more efficient blocking and retreaving of data
            if not contains_missing_samples:
               snpQuery = snpQuery.loc[snpQuery['snp'].map(lambda x: x not in list(map(str, fail_qc_snps_all)))]
            #import pdb; pdb.set_trace()
            n_sample_t = (phenotype_ds.size)
            n_e_samples_t = (len(genetically_unique_individuals))
            print ('For feature: ' + str(currentFeatureNumber) + '/' +str(len(feature_list)) + ' (' + feature_id + '): ' + str(snpQuery.shape[0]) + ' SNPs need to be tested.\n Please stand by.')
            if(n_perm != 0):
                bestPermutationPval = np.ones((n_perm), dtype = np.float)
            for snpGroup in utils.chunker(snpQuery, blocksize):
                snp_idxs = snpGroup['i'].values
                snp_names = snpGroup['snp'].values

                tested_snp_idxs.extend(snp_idxs)

                #subset genotype matrix, we cannot subselect at the same time, do in two steps.
                snp_df = pd.DataFrame(data = bed[snp_idxs,:].compute().transpose(), index = fam.index, columns = snp_names)
                snp_df = snp_df.loc[individual_ids,:]

                #SNP QC.
                #Now we do more proper QC on non-identical samples.
                #However, we do not use it when checking for missingness.
                #That could be extended but gives alot of overhead.
                if not contains_missing_samples:
                    #remove snps from snp_df if they fail QC
                    snp_df = snp_df.loc[:,snp_df.columns[~snp_df.columns.isin(fail_qc_snps_all)]]
                    if snp_df.shape[1] == 0:
                        continue
                    snps_to_test_df = snp_df.loc[:,snp_df.columns[~snp_df.columns.isin(pass_qc_snps_all)]]
                    if snps_to_test_df.shape[1] > 0:
                        #Only do QC on relevant SNPs. join pre-QCed list and new QCed list.
                        if kinship_df is not None:
                            passed_snp_names, failed_snp_names, call_rate, maf, hweP = do_snp_qc(snps_to_test_df.iloc[np.unique(snps_to_test_df.index, return_index = 1)[1]].loc[genetically_unique_individuals,:],
                                                                                min_call_rate, min_maf, min_hwe_P)
                        else:
                            passed_snp_names, failed_snp_names, call_rate, maf, hweP = do_snp_qc(snps_to_test_df, min_call_rate, min_maf, min_hwe_P)
                        snps_to_test_df = None
                        #append snp_names and failed_snp_names
                        pass_qc_snps_all.extend(passed_snp_names)
                        fail_qc_snps_all.extend(failed_snp_names)
                    snp_df = snp_df.loc[:,snp_df.columns[snp_df.columns.isin(pass_qc_snps_all)]]
                else:
                    #Do snp QC for relevant section.
                    if kinship_df is not None:
                        passed_snp_names, failed_snp_names, call_rate, maf, hweP = do_snp_qc(snp_df.iloc[np.unique(snp_df.index, return_index = 1)[1]].loc[genetically_unique_individuals,:],
                                                                                min_call_rate, min_maf, min_hwe_P)
                    else:
                        passed_snp_names, failed_snp_names, call_rate, maf, hweP = do_snp_qc(snp_df, min_call_rate, min_maf, min_hwe_P)
                    snp_df = snp_df.loc[:,snp_df.columns[snp_df.columns.isin(passed_snp_names)]]
                #print('step 0')
                if len(snp_df.columns) == 0:
                    continue

                snpQcInfo_t = None
                if call_rate is not None:
                    snpQcInfo_t = call_rate.transpose()
                    if maf is not None:
                        snpQcInfo_t = snpQcInfo_t.merge(maf.transpose(), how = 'outer')
                        if hweP is not None:
                            snpQcInfo_t = snpQcInfo_t.merge(hweP.transpose(), how = 'outer')
                call_rate = None
                maf = None
                hweP = None
                if snpQcInfo is None and snpQcInfo_t is not None:
                    snpQcInfo = snpQcInfo_t
                elif snpQcInfo_t is not None:
                    snpQcInfo = pd.concat([snpQcInfo, snpQcInfo_t], axis = 1)

                #We could make use of relatedness when imputing.
                snp_matrix_DF = pd.DataFrame(fill_NaN.fit_transform(snp_df), index = snp_df.index, columns = snp_df.columns)
                snp_df = None

#                test if the environments, covariates, kinship, snp and phenotype are in the same order
                if ((all(snp_matrix_DF.index == kinship_df.loc[individual_ids,individual_ids].index) if kinship_df is not None else True) &\
                     (all(phenotype_ds.index == covariate_df.loc[sample2individual_feature['sample'],:].index)if covariate_df is not None else True)&\
                     all(phenotype_ds.index == environment_df.loc[sample2individual_feature['sample'],:].index)&\
                     all(snp_matrix_DF.index == sample2individual_feature.loc[phenotype_ds.index]['iid'])):
                    '''
                    if all lines are in order put in arrays the correct genotype and phenotype
                    x=a if cond1 else b <---> equivalent to if cond1: x=a else x=b;                 better readability of the code
                    '''
                    kinship_mat = kinship_df.loc[individual_ids,individual_ids].values if kinship_df is not None else None
                    env =  environment_df.loc[sample2individual_feature['sample'],:].values
                    if covariate_df is not None:
                        cov_matrix =  covariate_df.loc[sample2individual_feature['sample'],:].values
                    else:
                        cov_matrix = None
                    if(snp_cov_df is not None and cov_matrix is not None):
                        snp_cov_df_tmp = snp_cov_df.loc[individual_ids,:]
                        snp_cov_df_tmp.index = sample2individual_feature['sample']
                        cov_matrix = np.concatenate((cov_matrix,snp_cov_df_tmp.values), 1)
                    elif snp_cov_df is not None :
                        snp_cov_df_tmp = snp_cov_df.loc[individual_ids,:]
                        snp_cov_df_tmp.index = sample2individual_feature['sample']
                        cov_matrix = snp_cov_df_tmp.values
                        #cov_matrix = np.concatenate((np.ones(snp_cov_df_tmp.shape[0]).reshape(np.ones(snp_cov_df_tmp.shape[0]).shape[0],1),snp_cov_df_tmp.values),1)

                    phenotype = utils.force_normal_distribution(phenotype_ds.values, method = gaussianize_method) if gaussianize_method is not None else phenotype_ds.values
                else:
                    print ('There is an issue in mapping phenotypes and genotypes')
                    sys.exit()
                #For limix 1.1 we need to switch to lm our selfs if there is no K.
#                return[snp_matrix_DF,phenotype, kinship_mat,cov_matrix]
                #import pdb; pdb.set_trace()
                try:
                    if(association_mode == True):
                        res_structlmm = struct_lmm(snp_matrix_DF, np.atleast_2d(phenotype.T).T, np.atleast_2d(env.T).T, covs = np.asarray(cov_matrix, float), no_interaction_test = True)
                    else: # i.e. interaction mode
                        res_structlmm = struct_lmm(snp_matrix_DF, np.atleast_2d(phenotype.T).T, np.atleast_2d(env.T).T, covs = np.asarray(cov_matrix, float), no_association_test = True)
                        # only one of the two analyses can be performed, so copy pvalues from interaction run as simply 'pv'
                        res_structlmm['pv'] = res_structlmm['pv_int']
                except:
                    print (feature_id)
                    print ('Struct-LMM failed')
                    sys.exit()
                #import pdb; pdb.set_trace()
                if(n_perm!=0):
                    pValueBuffer = []
                    totalSnpsToBeTested = (snp_matrix_DF.shape[1]*n_perm)
                    permutationStepSize = np.floor(totalSnpsToBeTested/blocksize)
                    if(permutationStepSize == 0):
                        permutationStepSize = 1
                    perm = 0;
                    if(write_permutations):
                        perm_df = pd.DataFrame(index = range(len(snp_matrix_DF.columns)), columns = ['snp_id'] + ['permutation_' + str(x) for x in range(n_perm)])
                        perm_df['snp_id'] = snp_matrix_DF.columns
                    for currentNperm in utils.chunker(list(range(1, n_perm+1)), permutationStepSize):
                        if kinship_df is not None:
                            temp = utils.get_shuffeld_genotypes_preserving_kinship(genetically_unique_individuals, relatedness_score, snp_matrix_DF,kinship_df.loc[individual_ids,individual_ids], len(currentNperm))
                        else :
                            temp = utils.get_shuffeld_genotypes(snp_matrix_DF, len(currentNperm))
                        #import pdb; pdb.set_trace()
                        # make into "geno_" dataframe for running structLMM
                        temp = pd.DataFrame(data = temp)
                        if(association_mode == True):
                            res_structlmm_perms = struct_lmm(temp, np.atleast_2d(phenotype.T).T, np.atleast_2d(env.T).T, covs = np.asarray(cov_matrix, float), no_interaction_test = True)
                        else:
                            res_structlmm_perms = struct_lmm(temp, np.atleast_2d(phenotype.T).T, np.atleast_2d(env.T).T, covs = np.asarray(cov_matrix, float), no_association_test = True)
                            res_structlmm_perms['pv'] = res_structlmm_perms['pv_int']
                        pValueBuffer.extend(np.asarray(res_structlmm_perms['pv']))
                    if(not(len(pValueBuffer) == totalSnpsToBeTested)):
                        #print(len(pValueBuffer))
                        #print(pValueBuffer)
                        #print(totalSnpsToBeTested)
                        print('Error in blocking logic for permutations.')
                        sys.exit()
                    for relevantOutput in utils.chunker(pValueBuffer, snp_matrix_DF.shape[1]) :
                        if(write_permutations):
                            perm_df['permutation_' + str(perm)] = relevantOutput
                        if(bestPermutationPval[perm] > min(relevantOutput)):
                            bestPermutationPval[perm] = min(relevantOutput)
                        perm+=1
                #print('step 3')
                #add these results to qtl_results
                temp_df = pd.DataFrame(index = range(len(snp_matrix_DF.columns)), columns = ['feature_id','snp_id','p_value','beta','beta_se','n_samples','empirical_feature_p_value'])
                temp_df['snp_id'] = snp_matrix_DF.columns
                temp_df['feature_id'] = feature_id
                temp_df['beta'] = 0
                temp_df['p_value'] = np.asarray(res_structlmm['pv'])
                temp_df['beta_se'] = -1
                #insert default dummy value
                temp_df['empirical_feature_p_value'] = -1.0
                if not temp_df.empty :
                    data_written = True
                    output_writer.add_result_df(temp_df)
                    if(write_permutations):
                        permutation_writer.add_permutation_results_df(perm_df, feature_id)
            if contains_missing_samples:
                snpQcInfo = snpQcInfo.transpose()
                snpQcInfo.to_csv(output_dir + '/snp_qc_metrics_naContaining_feature_{}.txt'.format(feature_id), sep = '\t')
                genetically_unique_individuals = tmp_unique_individuals
            else:
                if (snpQcInfo is not None and snpQcInfoMain is not None):
                    cols_to_use = snpQcInfo.columns.difference(snpQcInfoMain.columns)
                    snpQcInfoMain = pd.concat([snpQcInfoMain, snpQcInfo[cols_to_use]], axis = 1)
                elif snpQcInfo is not None :
                    snpQcInfoMain = snpQcInfo.copy(deep = True)
                #print('step 4')
            #This we need to change in the written file.
        if(n_perm > 1 and data_written):
            #updated_permuted_p_in_hdf5(bestPermutationPval, feature_id);
            alpha_para, beta_para = output_writer.apply_pval_correction(feature_id, bestPermutationPval, cis_mode)
            alpha_params.append(alpha_para)
            beta_params.append(beta_para)
        if(data_written) :
            n_samples.append(phenotype_ds.size)
            n_e_samples.append(len(genetically_unique_individuals))
        if not data_written :
            fail_qc_features.append(feature_id)
        #print('step 5')
    output_writer.close()
    if(write_permutations):
        permutation_writer.close()

    #gather unique indexes of tested snps
    tested_snp_idxs = list(set(tested_snp_idxs))
    #write annotation and snp data to file
    snp_df = pd.DataFrame()
    snp_df['snp_id'] = bim['snp']
    snp_df['chromosome'] = bim['chrom']
    snp_df['position'] = bim['pos']
    snp_df['assessed_allele'] = bim['a1']
    snp_df = snp_df.ix[tested_snp_idxs,:]
    snp_df = snp_df.drop_duplicates()
    snp_df.index = snp_df['snp_id']
    if snpQcInfoMain is not None :
        snp_df = snp_df.transpose()
        snpQcInfoMain = snpQcInfoMain.transpose()
        snpQcInfoMain['snp_id'] = snpQcInfoMain.index
        snpQcInfoMain =  snpQcInfoMain.drop_duplicates().transpose()
        #print(snp_df)
        #print(snpQcInfoMain)
        snp_df = snp_df.merge(snpQcInfoMain, how = 'outer')
        snp_df = snp_df.transpose()
        snp_df.columns = ['snp_id', 'chromosome', 'position', 'assessed_allele', 'call_rate', 'maf', 'hwe_pvalue']
    feature_list = [x for x in feature_list if x not in fail_qc_features]

    annotation_df = annotation_df.loc[feature_list,:]
    annotation_df['n_samples'] = n_samples
    annotation_df['n_e_samples'] = n_e_samples
    if(n_perm>1):
        annotation_df['alpha_param'] = alpha_params
        annotation_df['beta_param'] = beta_params
    if not selectionStart is None :
        snp_df.to_csv(output_dir + '/snp_metadata_{}_{}_{}.txt'.format(chromosome, selectionStart, selectionEnd), sep = '\t', index = False)
        annotation_df.to_csv(output_dir + '/feature_metadata_{}_{}_{}.txt'.format(chromosome, selectionStart, selectionEnd), sep = '\t')
    else :
        snp_df.to_csv(output_dir + '/snp_metadata_{}.txt'.format(chromosome), sep = '\t', index = False)
        annotation_df.to_csv(output_dir + '/feature_metadata_{}.txt'.format(chromosome), sep = '\t')

if __name__=='__main__':
    args = qtl_parse_args.get_struct_args()
    plink  = args.plink
    bgen = args.bgen
    anno_file = args.annotation_file
    extended_anno_file = args.extended_annotation_file
    pheno_file = args.phenotype_file
    env_file = args.environment_file #structLMM envs
    output_dir = args.output_directory
    window_size = args.window
    genetic_range = args.genomic_range
    covariates_file = args.covariates_file
    kinship_file = args.kinship_file
    samplemap_file = args.sample_mapping_file
    min_maf = args.minor_allel_frequency
    min_hwe_P = args.hardy_weinberg_equilibrium
    min_call_rate = args.call_rate
    block_size = args.block_size
    n_perm = int(args.number_of_permutations)
    snps_filename = args.variant_filter
    snp_feature_filename = args.feature_variant_filter
    feature_variant_covariate_filename = args.feature_variant_covariate
    random_seed = args.seed
    feature_filename = args.feature_filter
    relatedness_score = args.relatedness_score
    minimum_test_samples = args.minimum_test_samples
    gaussianize = args.gaussianize_method
    # can only perform cis or trans analysis
    cis = args.cis
    trans = args.trans
    # can only perform association or interaction structLMM
    association_mode = args.association_mode
    interaction_mode = args.interaction_mode

    write_permutations = args.write_permutations
    includeAllChromsomes = args.no_chromosome_filter

    if ((plink is None) and (bgen is None)):
        raise ValueError("No genotypes provided. Either specify a path to a binary plink genotype file or a bgen file.")
    if ((plink is not None) and (bgen is not None)):
        raise ValueError("Only one genotype file can be provided at once, not both plink and bgen")

    if (bgen is not None) :
        plinkGenotype = False
        geno_prefix = bgen
        raise ValueError("Not supported")
    else:
        plinkGenotype = True
        geno_prefix = plink

    if (cis and trans):
        raise ValueError("cis and trans cannot be specified simultaneously")
    elif (not cis and not trans):
        raise ValueError("At least one run mode (-c / -t) is needed.")

    if (association_mode and interaction_mode):
        raise ValueError("association and interaction analyses cannot be specified simultaneously")
    elif(not association_mode and not interaction_mode):
        raise ValueError("At least one run mode (-am / -im) is needed.")

    if (random_seed is None):
        random_seed = np.random.randint(40000)

    if(n_perm == 0 and write_permutations):
        write_permutations = False

    if(n_perm > 1 and n_perm < 10):
        n_perm = 10
        print("Defaults to 10 permutations, if permutations are only used for calibration please give in 1.")
    run_struct_QTL_analysis(pheno_file, anno_file, env_file, geno_prefix, plinkGenotype, output_dir, int(window_size),
                     min_maf = float(min_maf), min_hwe_P = float(min_hwe_P), min_call_rate = float(min_call_rate), association_mode = association_mode,
                     cis_mode = cis, blocksize = int(block_size), skipAutosomeFiltering = includeAllChromsomes, gaussianize_method = gaussianize,
                     seed = int(random_seed), minimum_test_samples = int(minimum_test_samples),  n_perm = int(n_perm), write_permutations = write_permutations,
                     relatedness_score = float(relatedness_score), feature_variant_covariate_filename = feature_variant_covariate_filename,
                     snps_filename = snps_filename, feature_filename = feature_filename, snp_feature_filename = snp_feature_filename,
                     genetic_range = genetic_range, covariates_filename = covariates_file, kinship_filename = kinship_file, sample_mapping_filename = samplemap_file,
                     extended_anno_filename = extended_anno_file)

reason: import numpy,import scipy


text:
from __future__ import print_function
from __future__ import absolute_import
# ============================================================================
# Copyright 2017 BRAIN Corporation. All rights reserved. This software is
# provided to you under BRAIN Corporation's Beta License Agreement and
# your use of the software is governed by the terms of that Beta License
# Agreement, found at http://www.braincorporation.com/betalicense.
# ============================================================================

import numpy as np
import pytest

from bc_gym_planning_env.utilities.frozenarray import freeze_array

def test_frozen_array():
    a = freeze_array(np.array([0, 1, 2]))
    np.testing.assert_array_equal(a, [0, 1, 2])

    assert a[0] == 0

    with pytest.raises(ValueError):
        a[0] = 3

    b = a[:1]
    assert len(b) == 1
    assert b[0] == 0
    with pytest.raises(ValueError):
        b[0] = 4

    # some slicing operations create a copy that is mutable again,
    # however original array is still immutable
    c = a[[0, 2]]
    np.testing.assert_array_equal(c, [0, 2])
    c[1] = 10
    np.testing.assert_array_equal(a, [0, 1, 2])
    with pytest.raises(ValueError):
        a[0] = 3

if __name__ == '__main__':
    test_frozen_array()

max_stars_repo_name: ghostFaceKillah/bc-gym-planning-env

max_issues_repo_name: ghostFaceKillah/bc-gym-planning-env

max_forks_repo_name: ghostFaceKillah/bc-gym-planning-env