text:
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import copy
import logging

import nglview as ngl
import networkx as nx

logging.basicConfig(format='%(asctime)s %(message)s')

DEBUG = True
if DEBUG == True:
    logging.basicConfig(level = logging.DEBUG )
    #logging.basicConfig( 'root':{'handlers':('console', 'file'), 'level':'DEBUG'} )

## bead properties

aaList = list( 'ACDEFGHIKLMNPQRSTVWY' )
aaList += ['CA', 'C', 'H', 'N', 'O', 'OH', 'S', 'NT', 'NE', 'NH', 'OT', 'CB', 'CG', 'CD', 'CE', 'CZ', 'CH']
#aaList = ['CA', 'C', 'H', 'N', 'O', 'S']
#print( len(aaList))

beadDictRadius = { a:0.77 for a in aaList }
beadDictProp1 = { a:0 for a in aaList }

ndim = 3

def getDistance( p ):
    return np.sqrt( np.sum( (p[0] - p[1]) ** 2 ) )

def getAngle( p ):
    p -= p[1]
    v1 = p[0] / np.linalg.norm(p[0])
    v2 = p[2] / np.linalg.norm(p[2])
    d = np.dot(v1, v2)
    angle = np.arccos(d)

    if np.isnan( angle ):
        return 0
    else:
        return np.degrees( angle )

def getTorsionAngle(p):
    ## https://stackoverflow.com/questions/20305272/dihedral-torsion-angle-from-four-points-in-cartesian-coordinates-in-python
    """Praxeolitic formula
    1 sqrt, 1 cross product"""

    p = np.array( p )

    b0 = -1.0 * ( p[1] - p[0])
    b1 = p[2] - p[1]
    b2 = p[3] - p[2]

    # normalize b1 so that it does not influence magnitude of vector
    # rejections that come next
    b1 /= np.linalg.norm(b1)

    # vector rejections
    # v = projection of b0 onto plane perpendicular to b1
    #   = b0 minus component that aligns with b1
    # w = projection of b2 onto plane perpendicular to b1
    #   = b2 minus component that aligns with b1
    v = b0 - np.dot(b0, b1)*b1
    w = b2 - np.dot(b2, b1)*b1

    # angle between v and w in a plane is the torsion angle
    # v and w may not be normalized but that's fine since tan is y/x
    x = np.dot(v, w)
    if ndim == 2:
        b1 = b1.tolist() + [0]
        v = v.tolist() + [0]
        y0 = np.cross( b1, v )[:-1]

    else:
        y0 = np.cross( b1, v )
    y = np.dot( y0, w)
    theta = np.degrees(np.arctan2(y, x))
    if np.isnan( theta ):
        return 0
    else:
        return theta

class bead: # residue or atom or any smallest entity in the model
    def __init__( self, code='C', name='C', bno=0, resno=0, chain='A' ):
        self.loc = np.zeros( (ndim) ) # location or coordinates
        self.name = name # one-letter code
        self.code = code
        self.radius = beadDictRadius[ code ] # radius or coarse-grained radius
        self.prop1 = beadDictProp1[ code ] # some property - just an example
        self.interaction = 0 # count or some quantity
        self.no = bno # some interation no for referencing within chain
        self.resno = resno # some residue no for referencing within chain
        self.chain = chain

    def move(self, dx ): ## update location by an increment
        self.loc += dx

    def __str__(self):
        return  'bead {}({}):{} in res {} and chain {} '.format( self.name, self.code, self.no, self.resno, self.chain  )

    def __repr__(self):
        return self.__str__()

class chain: ## a molecule
    def __init__(self, seqCode=[], seq=[], chname='A', loc=np.zeros( (ndim) ), initConfor='zigzag' ):
        self.bindex = 0 # loop over beads like a iterator
        self.name = chname
        self.seqCode = []
        self.seq = []
        self.bbInd = []
        #self.seqCode = seqCode
        if len(seq) == 0:
            seq = seqCode
        else:
            if len(seq)  < len(seqCode):
                logging.warning(  'Warning bead Codes and Names differ in count {}:{}'.format( len(seqCode), len(seq) ) )

        self.residues = []
        self.bondPair = []
        self.initConfor = initConfor
        #self.beads = [bead(s, bi, bi, self.name) for bi, s in enumerate( seq ) ]
        if len( seqCode ) == 0:
            self.nbeads = 0
            self.beads = []
        else:
            self.nbeads = 1
            self.beads = [ bead( seqCode[0], seq[0], 0, 0, self.name) ]
            self.beads[0].loc = loc
        self.update()
        for bi in range( 1, len(seqCode) ):
            self.grow(bcode=seqCode[bi], bname=seq[bi], addtoBead=bi-1, resno=bi, conType=initConfor )
            #self.bondPair.append( (bi-1, bi) )
            #self.setConformation(bi-1, bi, conType=initConfor)
        #self.update()

    ## set loc/orientation second bead with respect previous
    ## based bead radiuss
    def setConformation(self, b1, b2, conType ='zigzag' ):
        if type(conType) == str:
            if  conType == 'linear':
                r1 = self.beads[b1].radius
                r2 = self.beads[b2].radius
                d = r1 + r2
                sign = 1
                self.beads[b2].loc = self.beads[b1].loc + [0, d, 0]

            elif  conType == 'zigzag':
                r1 = self.beads[b1].radius
                r2 = self.beads[b2].radius
                d = np.sqrt( r1 + r2 )
                sign = 1 if b2%2 == 0 else -1
                self.beads[b2].loc = self.beads[b1].loc + [d, d * sign, 0]

        else: ## conType must be a dict
            if 'd' in conType:
                r1 = 0
                r2 = 0
                d = np.sqrt( conType['d'] )
            else:
                r1 = self.beads[b1].radius
                r2 = self.beads[b2].radius
                d = np.sqrt( r1 + r2 )

            sign = 1 if b2%2 == 0 else -1
            self.beads[b2].loc = self.beads[b1].loc + [d, d * sign, 0]

            if False and 'ang' in conType : # and 'angleTriple' in conType:
                ang = conType['ang'][0]
                if b1 - 1 > 0 :
                    self.setAngle( [b1 - 1, b1, b2], ang  )

            '''
            if 'ang' in conType and 'angleTriple' in conType:
                for ang, angleTriple in zip( conType['ang'], conType['angleTriple'] ):
                    self.setAngle( angleTriple, ang  )

            if 'tor' in conType and 'torsionQuad' in conType:
                for tor, torsionQuad in zip( conType['tor'], conType['torsionQuad'] ):
                    self.setTorsionAngle( torsionQuad, tor  )
            '''
        if  True in np.isnan( self.beads[b2].loc ):
            logging.error( ' NaN value in location, something went wrong while assiging coordinates {}{}:{}'.format(r1, r2, d) )
            logging.error(   '{}, {} '.format( d,  sign ) )
        return True

    def get_coords(self):
        return np.asarray( [b.loc for b in self.beads  ])

    ## assumming bondPair is in order
    def get_angleTriple(self):
        angleTriple = []
        n = len(self.bondPair)
        for i in range(n-1):
            a = self.bondPair[i]
            for j in range( i+1, n):
                b = self.bondPair[j]
                if a[1] == b[0]:
                    angleTriple.append( a + ( b[-1], )  )
        return angleTriple

    ## assuming angleTriple is in order
    def get_torsionQuad(self):
        torsionQuad = []
        n = len(self.angleTriple)
        for i in range(n-1):
            a = self.angleTriple[i]
            for j in range( i+1, n):
                b = self.angleTriple[j]
                if a[1:] == b[:-1] :
                    torsionQuad.append( a + ( b[-1], ) )
        return torsionQuad

    def update(self): ## topology change
        self.update0()
        self.update1()

    def update0(self): ## topology change
        self.angleTriple = self.get_angleTriple()
        self.torsionQuad = self.get_torsionQuad()

    def update1(self): ## conformation change
        self.coords = self.get_coords()
        self.torsions = [ getTorsionAngle( self.coords[q,:] ) for q in self.torsionQuad ]
        self.update2()

    def update2(self):
        self.bonds = [ getDistance( self.coords[bp,:]) for bp in self.bondPair ]
        self.angles = [ getAngle( self.coords[t,:]) for t in self.angleTriple ]

    def grow(self, bcode='C', bname='C', addtoBead=-1, resno=0, conType='zigzag' ):
        bi = self.nbeads
        if addtoBead == -1:
            addtoBead = bi-1
        b = bead(bcode, bname, bi, resno, self.name)
        self.beads.append( b )
        self.nbeads += 1
        self.seq.append( bname )
        self.seqCode.append( bcode )
        self.update()
        if bi != 0:
            self.setConformation(addtoBead, bi, conType)
            self.bondPair.append( (addtoBead, bi) )
        self.update()

    def __str__(self):
        return 'chain {} with {} bead(s)'.format(self.name, self.nbeads)

    def __repr__(self):
        return self.__str__()

    def __iter__(self):
        yield from self.beads

    def __next__(self):
        if self.bindex == 0:
            raise StopIteration
        self.bindex = self.bindex - 1
        return self.beads[self.bindex]

    def __getitem__(self, index):
        return self.beads[index]

    def copy(self):
        return copy.deepcopy( self )

    def draw(self, ax):
        pradius = 100
        coords = self.coords
        radiuss = [ b.radius * pradius for b in self ]
        for pair in self.bondPair:
            if pair[0] in self.bbInd and pair[1] in self.bbInd :
                plt.plot( coords[ pair, 0], coords[ pair, 1], 'k-'  )
            else:
                plt.plot( coords[ pair, 0], coords[ pair, 1], 'r-'  )

        ax.scatter( coords[:,0], coords[:, 1], radiuss, marker='o' )

    def rotateAngle(self, angleTriple, theta=10, endIndex=0, beadIndices=[] ): # one entry from angleTriple list
        startIndex = angleTriple[1]
        if endIndex == 0:
            endIndex = self.nbeads - 1
        c = self.coords[ startIndex, :  ] # temp center
        if len(beadIndices) == 0:
            beadIndices = list( range(startIndex, endIndex) )
        coordsSub = self.coords[ beadIndices, :  ].copy() - c ## rest of the chain
        #axis = coordsSub[1]  # rotation axis
        v0 = self.coords[ angleTriple[0], :  ] - c
        v0 /= np.linalg.norm( v0 ) # unit vector
        v1 = self.coords[ angleTriple[2], :  ] - c
        v1 /= np.linalg.norm( v1 ) # unit vector
        axis = np.cross( v0, v1 ) ### considering an axis perpendicular to the two vectors
        axis /= np.linalg.norm( axis )
        ##
        ## alternate axis
        v1s = [ [0, 0, 1], [0, 1, 0], [1, 0, 0], [1, 1, 1] ]
        v1i = 0
        while True in np.isnan(axis): ## two parallel lines in an axis
            axis = np.cross( v0, v1s[v1i] )
            axis /= np.linalg.norm( axis )
            v1i + 1

        rad = np.radians( theta )
        q0 = np.cos( rad/2 )
        q1 = np.sin( rad/2 ) * axis[0]
        q2 = np.sin( rad/2 ) * axis[1]
        q3 = np.sin( rad/2 ) * axis[2]
        Q = [ [(q0**2 + q1**2 - q2**2 - q3**2),      2*(q1*q2 - q0*q3),          2*(q1*q3 + q0*q2)],
             [ 2*(q2*q1 + q0*q3),     (q0**2 - q1**2 + q2**2 - q3**2),      2*(q2*q3 - q0*q1) ],
             [ 2*(q3*q1 - q0*q2),          2*(q3*q2 + q0*q1),     (q0**2 - q1**2 - q2**2 + q3**2) ] ]
        Q = np.array( Q )
        newCoords = np.matmul( coordsSub, Q.T )

        for i, bi in enumerate( beadIndices ):
            self.beads[bi].loc = newCoords[i] + c
        self.update1()

        if True in np.isnan( newCoords ):
            logging.error( 'Nan found in coordinates during rotation. Info below:' )
            logging.error( ['theta, rad, axis, v0, v1', theta, rad, axis, v0, v1 ] )
            logging.error( ['angleTriple, self.nbeads', angleTriple, self.nbeads ])
            logging.error( [ 'self.coords[ angleTriple ]', self.coords[ angleTriple ] ]  )
            logging.error(  ['Q', Q] )
            return False
        return True

    def rotateTorsion(self, torsionQuad, theta=10, endIndex=0, beadIndices=[] ): # one entry from torsionQuad list
        startIndex = torsionQuad[1]
        if endIndex == 0:
            endIndex = self.nbeads - 1
        c = self.coords[ startIndex, :  ] # temp center
        if len(beadIndices) == 0:
            beadIndices = list( range(startIndex, endIndex) )
        coordsSub = self.coords[ beadIndices, :  ].copy() - c ## rest of the chain
        #axis = coordsSub[1]  # rotation axis
        axis = self.coords[ torsionQuad[2], :  ] - c
        axis /= np.linalg.norm( axis ) # unit vector

        rad = np.radians( theta )
        q0 = np.cos( rad/2 )
        q1 = np.sin( rad/2 ) * axis[0]
        q2 = np.sin( rad/2 ) * axis[1]
        q3 = np.sin( rad/2 ) * axis[2]
        Q = [ [(q0**2 + q1**2 - q2**2 - q3**2),      2*(q1*q2 - q0*q3),          2*(q1*q3 + q0*q2)],
             [ 2*(q2*q1 + q0*q3),     (q0**2 - q1**2 + q2**2 - q3**2),      2*(q2*q3 - q0*q1) ],
             [ 2*(q3*q1 - q0*q2),          2*(q3*q2 + q0*q1),     (q0**2 - q1**2 - q2**2 + q3**2) ] ]
        Q = np.array( Q )
        newCoords = np.matmul( coordsSub, Q.T )

        for i, bi in enumerate( beadIndices ):
            self.beads[bi].loc = newCoords[i] + c
        self.update1()

        return True

    def setAngle( self, angleTriple, theta=10, endIndex=0, beadIndices=[]  ):
        theta0 = getAngle( self.coords[angleTriple,:] )
        logging.debug( ['setangle', angleTriple, theta0] )
        s = self.rotateAngle(angleTriple, theta - theta0 )
        if DEBUG:
            theta0 = getAngle( self.coords[angleTriple,:] )
            logging.debug( ['setangle', angleTriple, theta0] )
        return s

    def setTorsionAngle( self, torsionQuad, theta=10, endIndex=0, beadIndices=[]  ):
        theta0 = getTorsionAngle( self.coords[torsionQuad,:] )
        logging.debug( ['settorsion', torsionQuad, theta0] )
        s = self.rotateTorsion(torsionQuad, theta - theta0 )
        if DEBUG:
            theta0 = getTorsionAngle( self.coords[torsionQuad,:] )
            logging.debug( ['settorsion', torsionQuad, theta0] )
        return s

def write_pdb(chain, filename, rows=[], mode='a'): #writes pdb based on atom records
    #	try:
        with open(filename, mode) as f:
            if len(rows) == 0:
                rows = np.arange( chain.nbeads)
            if len( chain.residues ) == 0:
                residues = [ 'UNK' ] * chain.nbeads
            else:
                residues = chain.residues
            lineDict  = { }
            for i in rows:

                a = chain[i] #atom
                atomname = '{0:<4s}'.format( a.name )
                line = 'ATOM  {0:>5d} {1:<4s}{2:1s}{3:>3s} {4:1s}{5:>4d}{6:1s}   {7:>8.3f}{8:>8.3f}{9:>8.3f}{10:>6.2f}{11:>6.2f}{12:<10s}{13:>2s}{14:>2.1s}'
                line = line.format(a.no, atomname, '', residues[a.resno], a.chain, a.resno, '', a.loc[0], a.loc[1], a.loc[2], 1, 0, ' ', a.code[0], '')
                line += '\n'

                if a.resno in lineDict.keys():
                    lineDict[a.resno].append( line )
                else:
                    lineDict[a.resno] = [line]

            resnos = list( set( list( lineDict.keys() ) ) )
            for i in resnos:
                for line in lineDict[i]:
                    f.write(line)
        return None

CT_N_Ca = { 'd': 1.47,  'bb' : True }
CT_Ca_C = { 'd': 1.53,  'bb' : True, 'ang' : [ 121 ] } # CA-C-N or CA-C-O
CT_C_N = { 'd': 1.32,  'bb' : True,  'ang' : [ 123 ] } # C-N-C
CT_C_O = {'d': 1.24,  'bb' : True, }

CT_C_Na = { 'd': 1.352,  'bb' : True, }
CT_C_Na_His = { 'd': 1.37,  'bb' : True, }
CT_C_Nd = { 'd': 1.32,  'bb' : True, }
CT_C_OH = {'d': 1.43,  'bb' : True, }
CT_C_Od = {'d': 1.23,  'bb' : True, }

CT_C_Na_Trp = { 'd': 1.384,  'bb' : True, }
CT_Car_Car_Trp = {'d':1.384, 'bb' : False, 'ang' : [121.3] }

CT_Car_Car = {'d':1.41, 'bb' : False, 'ang' : [120] }
CT_C_C = {'d':1.535, 'bb' : False, 'ang' : [111.17] }
CT_Cd_Cd = {'d':1.339, 'bb' : False, 'ang' : [121.3] }
CT_C_S = {'d':1.83, 'bb' : False, 'ang' : [111.17] }
#if initConfor == 'alpha':

def buildPeptide( resList, chname='A', loc=np.zeros( (ndim) ), initConfor0='alpha', Ncap='' ):
    ch1 = chain( [], [], chname, loc, initConfor0)
    ch1.residues = resList
    ch1.bbInd = []
    ch1.bbCode = []

    initConfor = 'linear' # sc
    initConfor1 = 'linear' # mc
    phi = -140; psi = 130

    phi = -57; psi=-47
    #phi = -80; psi=150
    ## add backbone first
    ## add first residue with N terminal atom/ Ncap
    resno = -1

    ## declare conType
    if Ncap == '':
        ch1.grow(bcode='NT', bname='NT', addtoBead=-1, resno=0, conType=CT_C_N )
        biN = 0 ## reference for backbone N atom
        #ch1.bbInd.extend( []  ) ## add to backbone except the last atom - which will be added at the end
    else:
        ch1.grow(bcode='NT', bname='NT', addtoBead=-1, resno=0, conType=CT_C_N )
        biN = 0
        #ch1.bbInd.extend( []  )

    biCAo = '' ## old residue
    biCo = ''  ## old residue
    for resno, resName in enumerate( resList ):
        if resno > 0:
            biCAo = biCA ## old residue
            biCo = biC  ## old residue

        ## add mainchain
        if resno != 0: ## add spl Ncaps
            ch1.grow(bcode='N', bname='N', addtoBead=biC, resno=resno, conType=CT_N_Ca )
            biN = ch1[-1].no
        ch1.grow(bcode='CA', bname='CA', addtoBead=biN, resno=resno, conType=CT_N_Ca )
        biCA = biN + 1 ## index for CA atom too grow sidechain
        ### add sidechain - ACDEFGHIKLMNPQRSTVWY
        initConfor = 'linear'
        bi = biCA ## both same in this version
        if resName == 'GLY':
            None
        elif resName == 'ALA':
            ch1.grow(bcode='CB', bname='CB', addtoBead=biCA, resno=resno, conType=CT_C_C )
        elif resName == 'VAL':
            ch1.grow(bcode='CB', bname='CB', addtoBead=biCA, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='CG', bname='CG1', addtoBead=bi+1, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='CG', bname='CG2', addtoBead=bi+1, resno=resno, conType=CT_C_C )
        elif resName == 'LEU':
            ch1.grow(bcode='CB', bname='CB', addtoBead=biCA, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='CG', bname='CG', addtoBead=bi+1, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='CD', bname='CD1', addtoBead=bi+2, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='CD', bname='CD2', addtoBead=bi+2, resno=resno, conType=CT_C_C )
        elif resName == 'ILE':
            ch1.grow(bcode='C', bname='CB', addtoBead=biCA, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='C', bname='CG', addtoBead=bi+1, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='C', bname='CG', addtoBead=bi+1, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='C', bname='CD', addtoBead=bi+2, resno=resno, conType=CT_C_C )

        elif resName == 'PRO':
            ch1.grow(bcode='C', bname='CB', addtoBead=biCA, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='C', bname='CG', addtoBead=bi+1, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='C', bname='CD', addtoBead=bi+2, resno=resno, conType=CT_C_C )
            ch1.bondPair.append( [biCA-1, bi+3] ) ## add bond between N to Cdelta
        elif resName == 'CYS':
            ch1.grow(bcode='C', bname='CB', addtoBead=biCA, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='C', bname='S', addtoBead=bi+1, resno=resno, conType=CT_C_C )
        elif resName == 'MET':
            ch1.grow(bcode='C', bname='CB', addtoBead=biCA, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='C', bname='CG', addtoBead=bi+1, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='S', bname='S', addtoBead=bi+2, resno=resno, conType=CT_C_S )
            ch1.grow(bcode='C', bname='CE', addtoBead=bi+3, resno=resno, conType=CT_C_C )

        elif resName == 'THR':
            ch1.grow(bcode='CB', bname='CB', addtoBead=biCA, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='CG', bname='CG', addtoBead=bi+1, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='OH', bname='OH', addtoBead=bi+1, resno=resno, conType=CT_C_OH )

        elif resName == 'SER':
            ch1.grow(bcode='C', bname='CB', addtoBead=biCA, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='OH', bname='OH', addtoBead=bi+1, resno=resno, conType=CT_C_OH )
        elif resName == 'ASN':
            ch1.grow(bcode='C', bname='CB', addtoBead=biCA, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='C', bname='C', addtoBead=bi+1, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='C', bname='O', addtoBead=bi+2, resno=resno, conType=CT_C_Od )
            ch1.grow(bcode='C', bname='N', addtoBead=bi+2, resno=resno, conType=CT_C_Na )
        elif resName == 'GLN':
            ch1.grow(bcode='C', bname='CG', addtoBead=bi+1, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='C', bname='C', addtoBead=bi+2, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='C', bname='O', addtoBead=bi+3, resno=resno, conType=CT_C_Od )
            ch1.grow(bcode='C', bname='N', addtoBead=bi+3, resno=resno, conType=CT_C_Na )
            ch1.grow(bcode='C', bname='CB', addtoBead=biCA, resno=resno, conType=CT_C_C )
        elif resName == 'ASP':
            ch1.grow(bcode='C', bname='CB', addtoBead=biCA, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='C', bname='C', addtoBead=bi+1, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='O', bname='O', addtoBead=bi+2, resno=resno, conType=CT_C_O )
            ch1.grow(bcode='O', bname='O', addtoBead=bi+2, resno=resno, conType=CT_C_O )
        elif resName == 'GLU':
            ch1.grow(bcode='C', bname='CB', addtoBead=biCA, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='C', bname='CG', addtoBead=bi+1, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='C', bname='C', addtoBead=bi+2, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='O', bname='O', addtoBead=bi+3, resno=resno, conType=CT_C_O )
            ch1.grow(bcode='O', bname='O', addtoBead=bi+3, resno=resno, conType=CT_C_O )
        elif resName == 'HIS':
            ch1.grow(bcode='C', bname='CB', addtoBead=biCA, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='C', bname='CG', addtoBead=bi+1, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='C', bname='CD2', addtoBead=bi+2, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='C', bname='NE2', addtoBead=bi+3, resno=resno, conType=CT_C_Na_His )
            ch1.grow(bcode='C', bname='CE1', addtoBead=bi+4, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='C', bname='ND1', addtoBead=bi+5, resno=resno, conType=CT_C_Nd )
            ch1.bondPair.append( [bi+2, bi+6] ) ## add bond between NH to Cgamma
        elif resName == 'LYS':
            ch1.grow(bcode='C', bname='CB', addtoBead=biCA, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='C', bname='CG', addtoBead=bi+1, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='C', bname='CD', addtoBead=bi+2, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='C', bname='CE', addtoBead=bi+3, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='N', bname='N', addtoBead=bi+4, resno=resno, conType=CT_C_Na )
        elif resName == 'ARG':
            ch1.grow(bcode='C', bname='CB', addtoBead=biCA, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='C', bname='CG', addtoBead=bi+1, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='C', bname='CD', addtoBead=bi+2, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='C', bname='NE', addtoBead=bi+3, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='C', bname='CZ', addtoBead=bi+4, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='N', bname='NH', addtoBead=bi+5, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='N', bname='NH', addtoBead=bi+5, resno=resno, conType=CT_C_C )

        elif resName == 'TRP':
            ch1.grow(bcode='C', bname='CB', addtoBead=biCA, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='C', bname='CG', addtoBead=bi+1, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='C', bname='CD', addtoBead=bi+2, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='C', bname='CE', addtoBead=bi+3, resno=resno, conType=CT_Car_Car_Trp )
            ch1.grow(bcode='N', bname='NE', addtoBead=bi+4, resno=resno, conType=CT_C_Na_Trp )
            ch1.grow(bcode='C', bname='CD', addtoBead=bi+5, resno=resno, conType=CT_C_C )
            ch1.bondPair.append( [bi+1, bi+6] )

            ch1.grow(bcode='C', bname='CZ', addtoBead=bi+4, resno=resno, conType=CT_Car_Car_Trp )
            ch1.grow(bcode='C', bname='CH', addtoBead=bi+7, resno=resno, conType=CT_Car_Car_Trp )
            ch1.grow(bcode='C', bname='CZ', addtoBead=bi+8, resno=resno, conType=CT_Car_Car_Trp )
            ch1.grow(bcode='C', bname='CE', addtoBead=bi+9, resno=resno, conType=CT_Car_Car_Trp)
            ch1.bondPair.append( [bi+6, bi+10] ) ## complete the ring
        elif resName == 'TYR':
            ch1.grow(bcode='C', bname='CB', addtoBead=biCA, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='C', bname='CG', addtoBead=bi+1, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='C', bname='CD', addtoBead=bi+2, resno=resno, conType=CT_Car_Car )
            ch1.grow(bcode='C', bname='CE', addtoBead=bi+3, resno=resno, conType=CT_Car_Car )
            ch1.grow(bcode='C', bname='CZ', addtoBead=bi+4, resno=resno, conType=CT_Car_Car )
            ch1.grow(bcode='C', bname='CE', addtoBead=bi+5, resno=resno, conType=CT_Car_Car )
            ch1.grow(bcode='C', bname='CD', addtoBead=bi+6, resno=resno, conType=CT_Car_Car )
            ch1.grow(bcode='C', bname='OH', addtoBead=bi+5, resno=resno, conType=CT_C_OH )
            ch1.bondPair.append( [bi+2, bi+7] ) ## complete the ring
        elif resName == 'PHE':
            ch1.grow(bcode='C', bname='CB', addtoBead=biCA, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='C', bname='CG', addtoBead=bi+1, resno=resno, conType=CT_C_C )
            ch1.grow(bcode='C', bname='CD', addtoBead=bi+2, resno=resno, conType=CT_Car_Car )
            ch1.grow(bcode='C', bname='CE', addtoBead=bi+3, resno=resno, conType=CT_Car_Car )
            ch1.grow(bcode='C', bname='CZ', addtoBead=bi+4, resno=resno, conType=CT_Car_Car )
            ch1.grow(bcode='C', bname='CE', addtoBead=bi+5, resno=resno, conType=CT_Car_Car )
            ch1.grow(bcode='C', bname='CD', addtoBead=bi+6, resno=resno, conType=CT_Car_Car )
            ch1.bondPair.append( [bi+2, bi+7] ) ## complete the ring
            print( bi )
            tsets = [ [2, 3, 4], [3, 4, 5], [4, 5, 6], [5, 6, 7], [6, 7, 2], [7, 2, 3] ]
            tall = [ bi+x for x in range(2,7) ]
            for t in tsets:
                a = [ bi + x for x in t ]
                ch1.setAngle( a, 120, beadIndices= a )

        ### added side chain - add rest of backbone
        ch1.grow(bcode='C', bname='C', addtoBead=biCA, resno=resno, conType=CT_Ca_C )
        biC = ch1[-1].no
        if resno < len(resList)-1:
            ch1.grow(bcode='O', bname='O', addtoBead=biC, resno=resno, conType=CT_C_O )
        else:
         ## add C-terminal oxygen
            ch1.grow(bcode='OT', bname='OT', addtoBead=biC, resno=resno, conType=CT_C_O )
        ch1.bbInd.extend( [ biN, biCA, biC, biC+1 ] )
         # [ 114, 121 ] } # CA-C-N or CA-C-O

        ### adjust bond angle and torsion angles
        ## adjust bond angle
        sign = 1 #  (-1 if resno%2 != 0 else 1) * -1

        if False:
            if resno > 0:
                ch1.setAngle( [biCo, biCA-1, biCA ], -123 * sign ) # C-N-Ca
                ch1.setAngle( [ biCAo, biCo, biCA-1 ], 114 * sign )  #Ca-C-N
                ch1.setAngle( [ biCo+1, biCo, biCA-1 ], -125 * sign ) # O-C-N
            ch1.setAngle( [ biCA-1, biCA, biC ], 110 * sign )  # N-Ca-C
            ch1.setAngle( [ biCA, biC, biC+1 ], 121 * sign )  # Ca-C-O

            if resName != 'GLY':
                ch1.setAngle( [biCA-1, biCA, biCA+1 ], -111.17 * sign, beadIndices= list(range( biCA, biC )) )

            ## check
            if False and resno > 0:
                for a in [[biCo, biCA-1, biCA ], [ biCo+1, biCo, biCA-1 ], [ biCAo, biCo, biCA-1 ],
                          [ biCA-1, biCA, biC ], [ biCA, biC, biC+1 ] , [biCA-1, biCA, biCA+1 ] ]:
                    print( [resno, [-123, -125, -114, -119, -125, 120], getAngle( ch1.coords[ a ] ) ] )

        ## adust torsion angles -- bb
        if True and resno > 0:
            if resName != 'PRO':
                ch1.setTorsionAngle( [biCAo, biCo, biCA-1, biCA ], 180 ) # Ca-C-N-Ca
            else:
                ch1.setTorsionAngle( [biCAo, biCo, biCA-1, biCA ], 0 )

            ch1.setTorsionAngle( [biCo, biCA-1, biCA, biC ], phi * sign ) # C-N-Ca-C
            ch1.setTorsionAngle( [biCAo-1, biCAo, biCo, biCA-1 ], psi * sign) # N-Ca-C-N

        ## check
        if False and resno > 0:
            for a in [[biCo, biCA-1, biCA ], [ biCo+1, biCo, biCA-1 ], [ biCAo, biCo, biCA-1 ],
                      [ biCA-1, biCA, biC ], [ biCA, biC, biC+1 ] , [biCA-1, biCA, biCA+1 ] ]:
                print( [resno, [-123, -125, -114, -119, -125, 120], getAngle( ch1.coords[ a ] ) ] )

        if True and resno > 0:
            for a in [ [biCAo, biCo, biCA-1, biCA ], [biCo, biCA-1, biCA, biC ], [biCAo-1, biCAo, biCo, biCA-1 ] ]:
                print( [resno, ['omega, phi, psi'], getTorsionAngle( ch1.coords[ a ] ) ] )

    for resno, res in enumerate( resList):
        if True and resno > 0:
            for a in [ [biCAo, biCo, biCA-1, biCA ], [biCo, biCA-1, biCA, biC ], [biCAo-1, biCAo, biCo, biCA-1 ] ]:
                print( [resno, ['omega, phi, psi'], getTorsionAngle( ch1.coords[ a ] ) ] )

    ####
    ch1.bbCode = [ ch1[b].code for b in ch1.bbInd ]
    return ch1

max_stars_repo_name: rpkarandev/propylib

max_issues_repo_name: rpkarandev/propylib

max_forks_repo_name: rpkarandev/propylib

reason: import numpy,import networkx


text:
import pandas as pd
import numpy as np
import sys
import argparse
import os
from utils import *

__version__ = '1.0.0'
SOFTWARE_CORRESPONDENCE_EMAIL1 = 'jxiaoae@connect.ust.hk'
SOFTWARE_CORRESPONDENCE_EMAIL2 = 'mcaiad@connect.ust.hk'

HEADER = """
<><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><>
<>
<> LOG-TRAM: Leveraging the local genetic structure for trans-ancestry association mapping
<> Version: %s
<> MIT License
<>
<><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><>
<> Software-related correspondence: %s or %s
<><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><>
<> example:
    python <install path>/src/LOG-TRAM.py \\
        --out test \\
        --sumstats-popu1 ss_file1,ss_name1 \\
        --sumstats-popu2 ss_file2,ss_name2 \\
        --ldscores ldsc_annot_EUR_EAS_1mb_TGP_hm3_chr@_std \\
        --out-harmonized

<><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><>
""" % (__version__, SOFTWARE_CORRESPONDENCE_EMAIL1, SOFTWARE_CORRESPONDENCE_EMAIL2)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Leverage local genetic architecture for trans-ancestry association mapping')
    # Input / output
    parser.add_argument('--out', type=str, help='output file path', required=True)
    parser.add_argument('--sumstats-popu1', type=str, nargs="+", metavar="FILE,PHENOTYPE",
        help='summary statisitcs F(file path),P(phenotype) of population 1, separated by whitespace',required=True)
    parser.add_argument('--sumstats-popu2', type=str, nargs="+", metavar="FILE,PHENOTYPE",
        help='summary statisitcs F(file path),P(phenotype) of population 2, separated by whitespace',required=True)
    parser.add_argument('--ldscores', type=str,
        help='specifies prefix of the LD score files computed by S-LDXR \
            (popu1 <corresponding to population of --sumstats-popu1>, \
            popu2 <corresponding to population of --sumstats-popu2>, trans-ethnic), \
            If the filename prefix contains the symbol @, LOG-TRAM will replace the @ symbol \
            with chromosome number, then add the suffix _pop1.gz/_pop2.gz/_te.gz',required=True)
    parser.add_argument('--use_snps', type=str,
        help='SNPs list file (one rsID per line), If specified, this list will be used to restrict the final list of SNPs reported')
    parser.add_argument('--out-harmonized', action="store_true",
        help='If specified, LOG-TRAM will output harmonized summary statistics to disk')
    parser.add_argument('--out-reg-coef', action="store_true",
        help='If specified, LOG-TRAM will output LD score regression coeficients to disk')
    # Regression option
    parser.add_argument("--reg-int-ident", action="store_true",
        help="Optional argument indicating that the LDscore regression intercept matrix should be set to be the identity matrix.")
    parser.add_argument("--reg-int-diag", action="store_true",
        help="Optional argument indicating that the LDscore regression intercept matrix should have off-diagonal elements set to zero")
    # Summary Statistics Filtering Options
    parser.add_argument("--allowed-chr-values", type=str.upper, nargs="+",
        help='specify the allowed values for the chromosome')
    parser.add_argument("--remove-palindromic-snps", action="store_true",
        help='This option removes the SNPs whose major and minor alleles form a base pair')
    parser.add_argument('--num_threads', type=str, help='number of threads', default="32")
    args = parser.parse_args()

    os.environ["OMP_NUM_THREADS"] = args.num_threads
    os.environ["OPENBLAS_NUM_THREADS"] = args.num_threads
    os.environ["MKL_NUM_THREADS"] = args.num_threads
    os.environ["VECLIB_MAXIMUM_THREADS"] = args.num_threads
    os.environ["NUMEXPR_NUM_THREADS"] = args.num_threads

    logger = configure_logging(args.out)

    logger.info(HEADER)
    logger.info("See full log at: %s\n", os.path.abspath(args.out+'.log'))
    logger.info("\nProgram executed via:\n%s\n", ' '.join(sys.argv).replace("--", " \\ \n\t--"))

    logger.info("Reading in the base LD Scores")
    df_ldsc_pop1 = pd.concat(
        (pd.read_csv(args.ldscores.replace('@',str(c))+'_pop1.gz',
        compression='infer',sep='\t',dtype={'CHR':str})[['CHR','SNP','BP','base']] for c in range(1,23)),
        ignore_index=True)
    df_ldsc_pop1 = df_ldsc_pop1.dropna()
    df_ldsc_pop2 = pd.concat(
        (pd.read_csv(args.ldscores.replace('@',str(c))+'_pop2.gz',
        compression='infer',sep='\t',dtype={'CHR':str})[['CHR','SNP','BP','base']] for c in range(1,23)),
        ignore_index=True)
    df_ldsc_pop2 = df_ldsc_pop2.dropna()
    df_ldsc_te = pd.concat(
        (pd.read_csv(args.ldscores.replace('@',str(c))+'_te.gz',
        compression='infer',sep='\t',dtype={'CHR':str})[['CHR','SNP','BP','base']] for c in range(1,23)),
        ignore_index=True)
    df_ldsc_te = df_ldsc_te.dropna()
    logger.info("LD score matrix shape: {}x{}".format(df_ldsc_pop1.shape[0],df_ldsc_pop1.shape[1]))

    # Check / read in summary stats and then QC
    logger.info("Reading in summary statistics")
    sumstats_pop1 = {}
    sumstats_pop1_names = []
    for ssfile in args.sumstats_popu1:
        sumstats_pop1_names.append(ssfile.split(',')[1])
        sumstats_pop1[ssfile.split(',')[1]] = pd.read_csv(ssfile.split(',')[0],
            sep='\t',compression='infer',dtype={'CHR':str}).drop_duplicates('SNP').dropna()
        logger.info("Sumstats {} shape: {}x{}".format(ssfile.split(',')[1],
            sumstats_pop1[ssfile.split(',')[1]].shape[0],sumstats_pop1[ssfile.split(',')[1]].shape[1]))
    sumstats_pop2 = {}
    sumstats_pop2_names = []
    for ssfile in args.sumstats_popu2:
        sumstats_pop2_names.append(ssfile.split(',')[1])
        sumstats_pop2[ssfile.split(',')[1]] = pd.read_csv(ssfile.split(',')[0],
            sep='\t',compression='infer',dtype={'CHR':str}).drop_duplicates('SNP').dropna()
        logger.info("Sumstats {} shape: {}x{}".format(ssfile.split(',')[1],
            sumstats_pop2[ssfile.split(',')[1]].shape[0],sumstats_pop2[ssfile.split(',')[1]].shape[1]))
    P1n, P2n = len(sumstats_pop1_names), len(sumstats_pop2_names)

    # matching SNPs
    logger.info('Matching SNPs from LD scores and GWAS Sumstats')
    snps_lst = [set(df_ldsc_pop1['SNP'].values),set(df_ldsc_pop2['SNP'].values),set(df_ldsc_te['SNP'].values)]+\
        [set(sumstats_pop1[_]['SNP'].values) for _ in sumstats_pop1.keys()]+\
            [set(sumstats_pop2[_]['SNP'].values) for _ in sumstats_pop2.keys()]
    snps_common = snps_lst[0].intersection(*snps_lst)
    if args.use_snps is not None:
        logger.info('loading user-supplied SNPs from {}'.format(args.use_snps))
        snps_common_inp = np.loadtxt(args.use_snps,dtype=str)
        snps_common = list(snps_common & set(snps_common_inp))
    logger.info('Number of SNPS in initial intersection of all sources {}'.format(len(snps_common)))
    df_ldsc_pop1 = df_ldsc_pop1.loc[df_ldsc_pop1['SNP'].isin(snps_common)].reset_index(drop=True)
    df_ldsc_pop2 = df_ldsc_pop2.loc[df_ldsc_pop2['SNP'].isin(snps_common)].reset_index(drop=True)
    df_ldsc_te = df_ldsc_te.loc[df_ldsc_te['SNP'].isin(snps_common)].reset_index(drop=True)
    for k,v in sumstats_pop1.items():
        if 'chr' in v['CHR'].values[0]:
            v['CHR'] = v['CHR'].map(lambda x:x.replace('chr',''))
        if 'Chr' in v['CHR'].values[0]:
            v['CHR'] = v['CHR'].map(lambda x:x.replace('Chr',''))
        if 'CHR' in v['CHR'].values[0]:
            v['CHR'] = v['CHR'].map(lambda x:x.replace('CHR',''))
        v = v.loc[v['CHR'].isin([str(_) for _ in range(1,23)])]
        v['CHR'] = v['CHR'].astype(int)
        sumstats_pop1[k] = v.loc[v['SNP'].isin(snps_common)].sort_values(['CHR','BP']).reset_index(drop=True)
    for k,v in sumstats_pop2.items():
        if 'chr' in v['CHR'].values[0]:
            v['CHR'] = v['CHR'].map(lambda x:x.replace('chr',''))
        if 'Chr' in v['CHR'].values[0]:
            v['CHR'] = v['CHR'].map(lambda x:x.replace('Chr',''))
        if 'CHR' in v['CHR'].values[0]:
            v['CHR'] = v['CHR'].map(lambda x:x.replace('CHR',''))
        v = v.loc[v['CHR'].isin([str(_) for _ in range(1,23)])]
        v['CHR'] = v['CHR'].astype(int)
        sumstats_pop2[k] = v.loc[v['SNP'].isin(snps_common)].sort_values(['CHR','BP']).reset_index(drop=True)

    # Removing ambiguous SNPs
    logger.info('Removing ambiguous SNPs...')
    for k,v in sumstats_pop1.items():
        v['A1_int'] = v['A1'].map(lambda x: 1 if x in ['A','T'] else 2)
        v['A2_int'] = v['A2'].map(lambda x: 1 if x in ['A','T'] else 2)
    for k,v in sumstats_pop2.items():
        v['A1_int'] = v['A1'].map(lambda x: 1 if x in ['A','T'] else 2)
        v['A2_int'] = v['A2'].map(lambda x: 1 if x in ['A','T'] else 2)
    sumstat_ref = sumstats_pop1[sumstats_pop1_names[0]]
    sumstat_ref_allelesum = sumstat_ref['A1_int'].values+sumstat_ref['A2_int'].values
    if args.remove_palindromic_snps:
        snps_rm_idx = (np.array([(sumstat_ref_allelesum % 2 == 0)]+\
            [(v['A1_int'].values+v['A2_int'].values != sumstat_ref_allelesum) for v in sumstats_pop1.values()]+\
            [(v['A1_int'].values+v['A2_int'].values != sumstat_ref_allelesum) for v in sumstats_pop2.values()])).sum(axis=0)>0
    else:
        snps_rm_idx = (np.array([(v['A1_int'].values+v['A2_int'].values != sumstat_ref_allelesum) for v in sumstats_pop1.values()]+\
            [(v['A1_int'].values+v['A2_int'].values != sumstat_ref_allelesum) for v in sumstats_pop2.values()])).sum(axis=0)>0
    logger.info('Number of ambiguous SNPS have been removed: {}'.format(snps_rm_idx.sum()))
    df_ldsc_pop1 = df_ldsc_pop1.loc[~snps_rm_idx,:].reset_index(drop=True)
    df_ldsc_pop2 = df_ldsc_pop2.loc[~snps_rm_idx,:].reset_index(drop=True)
    df_ldsc_te = df_ldsc_te.loc[~snps_rm_idx,:].reset_index(drop=True)
    for k,v in sumstats_pop1.items():
        sumstats_pop1[k] = v.loc[~snps_rm_idx,:].reset_index(drop=True)
    for k,v in sumstats_pop2.items():
        sumstats_pop2[k] = v.loc[~snps_rm_idx,:].reset_index(drop=True)

    logger.info('Aligning the minor allele according to a designated reference summary statistics DataFrame: {}'.format(sumstats_pop1_names[0]))
    allele_ref = sumstats_pop1[sumstats_pop1_names[0]][['SNP','A1','A2']].copy()
    allele_ref.columns = ['SNP','A1_ref','A2_ref']
    for k,v in sumstats_pop1.items():
        v = v.merge(allele_ref,left_on='SNP',right_on='SNP',how='inner')
        index_flip = v['A1'].values!=v['A1_ref'].values
        logger.info('{} SNPs need to be flipped for sumstats: {}'.format(index_flip.sum(),k))
        v.loc[index_flip,'BETA'] = -v.loc[index_flip,'BETA']
        v.loc[index_flip,'Z'] = -v.loc[index_flip,'Z']
        v['CHR'] = v['CHR'].astype(str)
        sumstats_pop1[k] = v
        if args.out_harmonized:
            v = v[['CHR','BP','SNP','A1_ref','A2_ref','FRQ','BETA','SE','Z','P','N']]
            v.columns = ['CHR','BP','SNP','A1','A2','FRQ','BETA','SE','Z','P','N']
            v.to_csv(args.out+'_harmonized_pop1_'+k+'.txt',sep='\t',index=None)
    for k,v in sumstats_pop2.items():
        v = v.merge(allele_ref,left_on='SNP',right_on='SNP',how='inner')
        index_flip = v['A1'].values!=v['A1_ref'].values
        logger.info('{} SNPs need to be flipped for sumstats: {}'.format(index_flip.sum(),k))
        v.loc[index_flip,'BETA'] = -v.loc[index_flip,'BETA']
        v.loc[index_flip,'Z'] = -v.loc[index_flip,'Z']
        v['CHR'] = v['CHR'].astype(str)
        sumstats_pop2[k] = v
        if args.out_harmonized:
            v = v[['CHR','BP','SNP','A1_ref','A2_ref','FRQ','BETA','SE','Z','P','N']]
            v.columns = ['CHR','BP','SNP','A1','A2','FRQ','BETA','SE','Z','P','N']
            v.to_csv(args.out+'_harmonized_pop2_'+k+'.txt',sep='\t',index=None)

    # convert to stdized units
    for k,v in sumstats_pop1.items():
        #conversion_factor_col = np.sqrt(2.0 * v['FRQ'] * (1.0 - v['FRQ']))
        #v['BETA'] = v['BETA'] * conversion_factor_col
        #v['SE'] = v['SE'] * conversion_factor_col
        v['BETA'] = v.apply(lambda x:x['Z']/np.sqrt(x['N']),axis=1)
        v['SE'] = v['N'].map(lambda x:1/np.sqrt(x))
    for k,v in sumstats_pop2.items():
        #conversion_factor_col = np.sqrt(2.0 * v['FRQ'] * (1.0 - v['FRQ']))
        #v['BETA'] = v['BETA'] * conversion_factor_col
        #v['SE'] = v['SE'] * conversion_factor_col
        v['BETA'] = v.apply(lambda x:x['Z']/np.sqrt(x['N']),axis=1)
        v['SE'] = v['N'].map(lambda x:1/np.sqrt(x))

    # Run LD score regressions
    logger.info("Running LD Score regression for the whole genome.")
    Omega_global, Sigma_LD = run_ldscore_regressions(sumstats_pop1_names,sumstats_pop1,sumstats_pop2_names,\
        sumstats_pop2,df_ldsc_pop1[['base']].values,df_ldsc_pop2[['base']].values,df_ldsc_te[['base']].values)
    Omega_global = Omega_global.squeeze()
    logger.info("Regression coefficients (LD):\n%s", Omega_global)
    logger.info("Regression coefficients (Intercept):\n%s", Sigma_LD)

    if args.reg_int_diag:
        Sigma_LD = np.diag(np.diag(Sigma_LD))
        logger.info("Modified regression coefficients (Intercept):\n%s", Sigma_LD)
    if args.reg_int_ident:
        Sigma_LD = np.eye(Sigma_LD.shape[0])
        logger.info("Modified regression coefficients (Intercept):\n%s", Sigma_LD)

    # beta, se, N_eff
    tram_res_beta = np.zeros((df_ldsc_pop1.shape[0],P1n+P2n))
    tram_res_beta_se = np.zeros((df_ldsc_pop1.shape[0],P1n+P2n))

    if args.allowed_chr_values is None:
        chrs = [str(_) for _ in range(1,23)]
    else:
        chrs = args.allowed_chr_values

    res_omegas_local = {}
    # Run LD score regressions for annotation
    logger.info("\nRunning LD Score regression and LOG-TRAM for local regions chromosome by chromosome.\n")
    for c in chrs:
        ldscore1_c = pd.read_csv(args.ldscores.replace('@',c)+'_pop1.gz',compression='infer',sep='\t')
        ldscore2_c = pd.read_csv(args.ldscores.replace('@',c)+'_pop2.gz',compression='infer',sep='\t')
        ldscorete_c = pd.read_csv(args.ldscores.replace('@',c)+'_te.gz',compression='infer',sep='\t')
        ldscore1_c = ldscore1_c.loc[ldscore1_c['SNP'].isin(df_ldsc_pop1['SNP'])].reset_index(drop=True)
        ldscore2_c = ldscore2_c.loc[ldscore2_c['SNP'].isin(df_ldsc_pop1['SNP'])].reset_index(drop=True)
        ldscorete_c = ldscorete_c.loc[ldscorete_c['SNP'].isin(df_ldsc_pop1['SNP'])].reset_index(drop=True)
        chr_snp_idx = df_ldsc_pop1.loc[df_ldsc_pop1['CHR']==c].index.values
        snp_num_c = chr_snp_idx.shape[0]
        logger.info("Processing chromosome {}, SNPs number: {}".format(c,chr_snp_idx.shape[0]))

        sumstats_pop1_c = {}
        beta_array = np.zeros((snp_num_c,P1n+P2n))
        se_array = np.zeros((snp_num_c,P1n+P2n))
        for i,ss_name in enumerate(sumstats_pop1_names):
            v = sumstats_pop1[ss_name].copy()
            v = v.loc[v['CHR']==c].reset_index(drop=True)
            sumstats_pop1_c[ss_name] = v
            beta_array[:,i] = v['BETA'].values
            se_array[:,i] = v['SE'].values
        sumstats_pop2_c = {}
        for i,ss_name in enumerate(sumstats_pop2_names):
            v = sumstats_pop2[ss_name].copy()
            v = v.loc[v['CHR']==c].reset_index(drop=True)
            sumstats_pop2_c[ss_name] = v
            beta_array[:,i+P1n] = v['BETA'].values
            se_array[:,i+P1n] = v['SE'].values

        name_windows = list(ldscore1_c.columns.values[4:])
        n_window = len(name_windows)
        ldscore1_cbase = ldscore1_c['base'].values
        ldscore2_cbase = ldscore2_c['base'].values
        ldscorete_cbase = ldscorete_c['base'].values
        logger.info("find {} local regions in chromosome {}".format(len(name_windows),c))
        omega_windows = np.zeros((len(name_windows),P1n+P2n,P1n+P2n))

        omega_ldscore_snps = np.zeros((snp_num_c,P1n+P2n,P1n+P2n))
        sigma_snps = np.zeros((snp_num_c,P1n+P2n,P1n+P2n))

        ldscore1_jk = np.zeros((snp_num_c,2))
        ldscore1_jk[:,0] = ldscore1_cbase
        ldscore2_jk = np.zeros((snp_num_c,2))
        ldscore2_jk[:,0] = ldscore2_cbase
        ldscorex_jk = np.zeros((snp_num_c,2))
        ldscorex_jk[:,0] = ldscorete_cbase

        for k,name_window in enumerate(name_windows):
            start,end = int(name_window.split('_')[1]),int(name_window.split('_')[2])
            idx_snp_window_chr = ldscore1_c.loc[(ldscore1_c['BP']>=start)&(ldscore1_c['BP']<end)].index.values
            p_window = idx_snp_window_chr.shape[0]
            if p_window==0:
                continue

            ldscore1_jk[:,1] = ldscore1_c[name_window].values
            ldscore2_jk[:,1] = ldscore2_c[name_window].values
            ldscorex_jk[:,1] = ldscorete_c[name_window].values

            omega_window, _ = run_ldscore_regressions(sumstats_pop1_names,sumstats_pop1_c,sumstats_pop2_names,\
                sumstats_pop2_c,ldscore1_jk,ldscore2_jk,ldscorex_jk,Sigma_LD)
            omega_window = make_positive_definite(omega_window.sum(axis=0))
            omega_windows[k] = omega_window

            for j in range(p_window):
                sigma_snps[idx_snp_window_chr[j]] = np.diag(se_array[idx_snp_window_chr[j]]) @ Sigma_LD @ np.diag(se_array[idx_snp_window_chr[j]])
                ld_matrx_j = np.zeros_like(omega_window)
                ld_matrx_j[:P1n,:P1n] = ldscore1_cbase[idx_snp_window_chr[j]]
                ld_matrx_j[:P1n,P1n:] = ld_matrx_j[P1n:,:P1n] = ldscorete_cbase[idx_snp_window_chr[j]]
                ld_matrx_j[P1n:,P1n:] = ldscore2_cbase[idx_snp_window_chr[j]]
                omega_ldscore_snps[idx_snp_window_chr[j]] = omega_window*ld_matrx_j

        # Run the LOG-TRAM method
        new_betas, new_beta_ses = run_tram_method(beta_array,omega_ldscore_snps,sigma_snps)
        tram_res_beta[chr_snp_idx] = new_betas
        tram_res_beta_se[chr_snp_idx] = new_beta_ses
        res_omegas_local[c] = dict(zip(name_windows,omega_windows))
    if args.out_reg_coef:
        logger.info("Saving local regions regression coefficients to file %s \n",args.out+'_TRAM_reg_coefs.npy')
        np.save(args.out+'_TRAM_reg_coefs.npy',res_omegas_local)

    logger.info("Preparing results")
    sumstats_pop1_new = {}
    for i,ss_name in enumerate(sumstats_pop1_names):
        ss_df = sumstats_pop1[ss_name].copy()
        ss_df['BETA_tram'] = tram_res_beta[:,i]
        ss_df['SE_tram'] = tram_res_beta_se[:,i]

        ss_df.loc[ss_df['BETA_tram']==0,'BETA_tram'] = ss_df.loc[ss_df['BETA_tram']==0,'BETA']
        ss_df.loc[ss_df['SE_tram']==0,'SE_tram'] = ss_df.loc[ss_df['SE_tram']==0,'SE']
        ss_df.loc[ss_df['BETA_tram'].isnull(),'BETA_tram'] = ss_df.loc[ss_df['BETA_tram'].isnull(),'BETA']
        ss_df.loc[ss_df['SE_tram'].isnull(),'SE_tram'] = ss_df.loc[ss_df['SE_tram'].isnull(),'SE']

        ss_df['Z_tram'] = ss_df['BETA_tram']/ss_df['SE_tram']
        ss_df = ss_df.loc[ss_df['CHR'].isin(chrs)]
        ss_df['P_tram'] = calculate_p(ss_df['Z_tram'].values)
        conversion_factor_col = np.reciprocal(np.sqrt(2.0 * ss_df['FRQ'] * (1.0 - ss_df['FRQ'])))
        ss_df['BETA_tram'] = ss_df['BETA_tram']*conversion_factor_col
        ss_df['SE_tram'] = ss_df['SE_tram']*conversion_factor_col
        ss_df = ss_df[['CHR','BP','SNP','A1_ref','A2_ref','FRQ','BETA_tram','SE_tram','Z_tram','P_tram','N']]
        ss_df.columns = ['CHR','BP','SNP','A1','A2','FRQ','BETA','SE','Z','P','N']
        sumstats_pop1_new[ss_name] = ss_df
    sumstats_pop2_new = {}
    for i,ss_name in enumerate(sumstats_pop2_names):
        ss_df = sumstats_pop2[ss_name].copy()
        ss_df['BETA_tram'] = tram_res_beta[:,i+P1n]
        ss_df['SE_tram'] = tram_res_beta_se[:,i+P1n]

        ss_df.loc[ss_df['BETA_tram']==0,'BETA_tram'] = ss_df.loc[ss_df['BETA_tram']==0,'BETA']
        ss_df.loc[ss_df['SE_tram']==0,'SE_tram'] = ss_df.loc[ss_df['SE_tram']==0,'SE']
        ss_df.loc[ss_df['BETA_tram'].isnull(),'BETA_tram'] = ss_df.loc[ss_df['BETA_tram'].isnull(),'BETA']
        ss_df.loc[ss_df['SE_tram'].isnull(),'SE_tram'] = ss_df.loc[ss_df['SE_tram'].isnull(),'SE']

        ss_df['Z_tram'] = ss_df['BETA_tram']/ss_df['SE_tram']
        ss_df = ss_df.loc[ss_df['CHR'].isin(chrs)]
        ss_df['P_tram'] = calculate_p(ss_df['Z_tram'].values)
        conversion_factor_col = np.reciprocal(np.sqrt(2.0 * ss_df['FRQ'] * (1.0 - ss_df['FRQ'])))
        ss_df['BETA_tram'] = ss_df['BETA_tram']*conversion_factor_col
        ss_df['SE_tram'] = ss_df['SE_tram']*conversion_factor_col
        ss_df = ss_df[['CHR','BP','SNP','A1_ref','A2_ref','FRQ','BETA_tram','SE_tram','Z_tram','P_tram','N']]
        ss_df.columns = ['CHR','BP','SNP','A1','A2','FRQ','BETA','SE','Z','P','N']
        sumstats_pop2_new[ss_name] = ss_df

    # Run LD score regressions
    logger.info("Running LD Score regression for the whole genome after Meta-GWAS.")
    Omega_global_new, Sigma_LD_new = run_ldscore_regressions(sumstats_pop1_names,sumstats_pop1_new,sumstats_pop2_names,\
        sumstats_pop2_new,df_ldsc_pop1.loc[df_ldsc_pop1['CHR'].isin(chrs),['base']].values,\
            df_ldsc_pop2.loc[df_ldsc_pop2['CHR'].isin(chrs),['base']].values,\
                df_ldsc_te.loc[df_ldsc_te['CHR'].isin(chrs),['base']].values)
    Omega_global_new = Omega_global_new.squeeze()
    logger.info("Regression coefficients (LD):\n%s", Omega_global_new)
    logger.info("Regression coefficients (Intercept):\n%s", Sigma_LD_new)

    logger.info("Estimating the effective sample size and save output to disk")
    for i,ss_name in enumerate(sumstats_pop1_names):
        logger.info("Writing %s (Population 1) to disk", ss_name)
        ss_df_new = sumstats_pop1_new[ss_name]
        ss_df = sumstats_pop1[ss_name]
        Neff_f = ((ss_df_new['Z']**2).mean() - Sigma_LD_new[i,i])/((ss_df['Z']**2).mean() - Sigma_LD[i,i])
        if Neff_f<1:
            Neff_f = ((ss_df_new['Z']**2).mean() - 1)/((ss_df['Z']**2).mean() - Sigma_LD[i,i])
        ss_df_new['N_eff'] = Neff_f*ss_df_new['N'].values
        ss_df_new.to_csv(args.out+'_TRAM_pop1_'+ss_name+'.txt',sep='\t',index=None)

    for i,ss_name in enumerate(sumstats_pop2_names):
        logger.info("Writing %s (Population 2) to disk", ss_name)
        ss_df_new = sumstats_pop2_new[ss_name]
        ss_df = sumstats_pop2[ss_name]
        Neff_f = ((ss_df_new['Z']**2).mean() - Sigma_LD_new[i+P1n,i+P1n])/((ss_df['Z']**2).mean() - Sigma_LD[i+P1n,i+P1n])
        if Neff_f<1:
            Neff_f = ((ss_df_new['Z']**2).mean() - 1)/((ss_df['Z']**2).mean() - Sigma_LD[i+P1n,i+P1n])
        ss_df_new['N_eff'] = Neff_f*ss_df_new['N'].values
        ss_df_new.to_csv(args.out+'_TRAM_pop2_'+ss_name+'.txt',sep='\t',index=None)
    logger.info("\nExecution complete")

ext: py

max_stars_repo_name: YangLabHKUST/LOG-TRAM

max_issues_repo_name: YangLabHKUST/LOG-TRAM

max_forks_repo_name: YangLabHKUST/LOG-TRAM

reason: import numpy


text:
#!/usr/bin/env python

import cv2
import matplotlib.pyplot as plt
import numpy as np
import torch
import torchvision.transforms as transforms
import torchvision.transforms.functional as VF
import torchvision.utils as vutils
from PIL import Image

def load_image(image_path):
    """Loads image as a PIL RGB image.

        Args:
            - **image_path (str) - **: A path to the image

        Returns:
            An instance of PIL.Image.Image in RGB

    """

    return Image.open(image_path).convert('RGB')

def apply_transforms(image, size=224):
    """Transforms a PIL image to torch.Tensor.

    Applies a series of tranformations on PIL image including a conversion
    to a tensor. The returned tensor has a shape of :math:`(N, C, H, W)` and
    is ready to be used as an input to neural networks.

    First the image is resized to 256, then cropped to 224. The `means` and
    `stds` for normalisation are taken from numbers used in ImageNet, as
    currently developing the package for visualizing pre-trained models.

    The plan is to to expand this to handle custom size/mean/std.

    Args:
        image (PIL.Image.Image or numpy array)
        size (int, optional, default=224): Desired size (width/height) of the
            output tensor

    Shape:
        Input: :math:`(C, H, W)` for numpy array
        Output: :math:`(N, C, H, W)`

    Returns:
        torch.Tensor (torch.float32): Transformed image tensor

    Note:
        Symbols used to describe dimensions:
            - N: number of Images in a batch
            - C: number of channels
            - H: height of the image
            - W: width of the image

    """

    if not isinstance(image, Image.Image):
        image = VF.to_pil_image(image)

    means = [0.485, 0.456, 0.406]
    stds = [0.229, 0.224, 0.225]

    transform = transforms.Compose([
        transforms.Resize(size),
        transforms.CenterCrop(size),
        transforms.ToTensor(),
        transforms.Normalize(means, stds)
    ])

    tensor = transform(image).unsqueeze(0)

    tensor.requires_grad = True

    return tensor

def apply_transforms_v0(image, size=224):
    """Transforms a PIL image to torch.Tensor.

    Applies a series of tranformations on PIL image including a conversion
    to a tensor. The returned tensor has a shape of :math:`(N, C, H, W)` and
    is ready to be used as an input to neural networks.

    First the image is resized to 256, then cropped to 224. The `means` and
    `stds` for normalisation are taken from numbers used in ImageNet, as
    currently developing the package for visualizing pre-trained models.

    The plan is to to expand this to handle custom size/mean/std.

    Args:
        image (PIL.Image.Image or numpy array)
        size (int, optional, default=224): Desired size (width/height) of the
            output tensor

    Shape:
        Input: :math:`(C, H, W)` for numpy array
        Output: :math:`(N, C, H, W)`

    Returns:
        torch.Tensor (torch.float32): Transformed image tensor

    Note:
        Symbols used to describe dimensions:
            - N: number of Images in a batch
            - C: number of channels
            - H: height of the image
            - W: width of the image

    """

    if not isinstance(image, Image.Image):
        image = VF.to_pil_image(image)

    means = [0.485, 0.456, 0.406]
    stds = [0.229, 0.224, 0.225]

    transform = transforms.Compose([
        transforms.Resize(size),
        transforms.CenterCrop(size),
        transforms.ToTensor()
    ])

    tensor = transform(image).unsqueeze(0)

    tensor.requires_grad = True

    return tensor

def denormalize(tensor):
    """Reverses the normalisation on a tensor.

    Performs a reverse operation on a tensor, so the pixel value range is
    between 0 and 1. Useful for when plotting a tensor into an image.

    Normalisation: (image - mean) / std
    Denormalisation: image * std + mean

    Args:
        tensor (torch.Tensor, dtype=torch.float32): Normalized image tensor

    Shape:
        Input: :math:`(N, C, H, W)`
        Output: :math:`(N, C, H, W)` (same shape as input)

    Return:
        torch.Tensor (torch.float32): Demornalised image tensor with pixel
            values between [0, 1]

    Note:
        Symbols used to describe dimensions:
            - N: number of Images in a batch
            - C: number of channels
            - H: height of the image
            - W: width of the image

    """

    means = [0.485, 0.456, 0.406]
    stds = [0.229, 0.224, 0.225]

    denormalized = tensor.clone()

    for channel, mean, std in zip(denormalized[0], means, stds):
        channel.mul_(std).add_(mean)

    return denormalized

def standardize_and_clip(tensor, min_value=0.0, max_value=1.0):
    """Standardizes and clips input tensor.

    Standardize the input tensor (mean = 0.0, std = 1.0), ensures std is 0.1
    and clips it to values between min/max (default: 0.0/1.0).

    Args:
        tensor (torch.Tensor):
        min_value (float, optional, default=0.0)
        max_value (float, optional, default=1.0)

    Shape:
        Input: :math:`(C, H, W)`
        Output: Same as the input

    Return:
        torch.Tensor (torch.float32): Normalised tensor with values between
            [min_value, max_value]

    """

    tensor = tensor.detach().cpu()

    mean = tensor.mean()
    std = tensor.std()
    if std == 0:
        std += 1e-7

    standardized = tensor.sub(mean).div(std).mul(0.1)
    clipped = standardized.add(0.5).clamp(min_value, max_value)

    return clipped

def format_for_plotting(tensor):
    """Formats the shape of tensor for plotting.

    Tensors typically have a shape of :math:`(N, C, H, W)` or :math:`(C, H, W)`
    which is not suitable for plotting as Images. This function formats an
    input tensor :math:`(H, W, C)` for RGB and :math:`(H, W)` for mono-channel
    data.

    Args:
        tensor (torch.Tensor, torch.float32): Image tensor

    Shape:
        Input: :math:`(N, C, H, W)` or :math:`(C, H, W)`
        Output: :math:`(H, W, C)` or :math:`(H, W)`, respectively

    Return:
        torch.Tensor (torch.float32): Formatted image tensor (detached)

    Note:
        Symbols used to describe dimensions:
            - N: number of Images in a batch
            - C: number of channels
            - H: height of the image
            - W: width of the image

    """

    has_batch_dimension = len(tensor.shape) == 4
    formatted = tensor.clone()

    if has_batch_dimension:
        formatted = tensor.squeeze(0)

    if formatted.shape[0] == 1:
        return formatted.squeeze(0).detach()
    else:
        return formatted.permute(1, 2, 0).detach()

def visualize(input_, gradients, save_path=None, cmap='viridis', alpha=0.7):
    """ Method to plot the explanation.

        # Arguments
            input_: Tensor. Original image.
            gradients: Tensor. Saliency map result.
            save_path: String. Defaults to None.
            cmap: Defaults to be 'viridis'.
            alpha: Defaults to be 0.7.

    """

    input_ = format_for_plotting(denormalize(input_))
    gradients = format_for_plotting(standardize_and_clip(gradients))

    subplots = [
        ('Input image', [(input_, None, None)]),
        ('Saliency map across RGB channels', [(gradients, None, None)]),
        ('Overlay', [(input_, None, None), (gradients, cmap, alpha)])
    ]

    num_subplots = len(subplots)

    fig = plt.figure(figsize=(16, 3))

    for i, (title, images) in enumerate(subplots):
        ax = fig.add_subplot(1, num_subplots, i + 1)
        ax.set_axis_off()

        for image, cmap, alpha in images:
            ax.imshow(image, cmap=cmap, alpha=alpha)

        ax.set_title(title)
    if save_path is not None:
        plt.savefig(save_path)

def basic_visualize(input_, gradients, save_path=None, weight=None, cmap='viridis', alpha=0.7):
    """ Method to plot the explanation.

        # Arguments
            input_: Tensor. Original image.
            gradients: Tensor. Saliency map result.
            save_path: String. Defaults to None.
            cmap: Defaults to be 'viridis'.
            alpha: Defaults to be 0.7.

    """
    input_ = format_for_plotting(denormalize(input_))
    gradients = format_for_plotting(standardize_and_clip(gradients))

    subplots = [
        ('Saliency map across RGB channels', [(gradients, None, None)]),
        ('Overlay', [(input_, None, None), (gradients, cmap, alpha)])
    ]

    num_subplots = len(subplots)

    fig = plt.figure(figsize=(4, 4))

    for i, (title, images) in enumerate(subplots):
        ax = fig.add_subplot(1, num_subplots, i + 1)
        ax.set_axis_off()

        for image, cmap, alpha in images:
            ax.imshow(image, cmap=cmap, alpha=alpha)

    if save_path is not None:
        plt.savefig(save_path)

def show_heatmap(mask, title=None):
    """Make heatmap from mask and synthesize GradCAM result image using heatmap and img.
    Args:
        mask (Tensor): shape (1, 1, H, W)
    Return:
        heatmap (Tensor): shape (1, 3, H, W)
        :param title:
    """
    mask = (mask - mask.min()).div(mask.max() - mask.min()).data
    heatmap = cv2.applyColorMap(np.uint8(255 * mask.squeeze().float()), cv2.COLORMAP_JET)  # [H, W, 3]
    heatmap = torch.from_numpy(heatmap).permute(2, 0, 1).float().div(255)

    b, g, r = heatmap.split(1)
    heatmap = torch.cat([r, g, b])
    heatmap = (heatmap - heatmap.min()) / (heatmap.max() - heatmap.min())
    if title is not None:
        vutils.save_image(heatmap, title)

    return heatmap

def show_cam(img, mask, title=None):
    """Make heatmap from mask and synthesize GradCAM result image using heatmap and img.
    Args:
        img (Tensor): shape (1, 3, H, W)
        mask (Tensor): shape (1, 1, H, W)
    Return:
        heatmap (Tensor): shape (3, H, W)
        cam (Tensor): synthesized GradCAM cam of the same shape with heatmap.
        :param title:
    """
    mask = (mask - mask.min()).div(mask.max() - mask.min()).data
    heatmap = cv2.applyColorMap(np.uint8(255 * mask.squeeze().float()), cv2.COLORMAP_JET)  # [H, W, 3]
    heatmap = torch.from_numpy(heatmap).permute(2, 0, 1).float().div(255)

    b, g, r = heatmap.split(1)
    heatmap = torch.cat([r, g, b])
    heatmap = (heatmap - heatmap.min()) / (heatmap.max() - heatmap.min())

    # cam = heatmap + img.cpu()
    cam = 1 * (1 - mask ** 0.8) * img + (mask ** 0.8) * heatmap
    # cam = (cam - cam.min()) / (cam.max() - cam.min())
    if title is not None:
        vutils.save_image(cam, title)

    return cam

def preprocess_img(cv_img):
    """Turn a opencv image into tensor and normalize it"""
    # revert the channels from BGR to RGB
    img = cv_img.copy()[:, :, ::-1]
    # convert tor tensor
    img = torch.from_numpy(np.ascontiguousarray(np.transpose(img, (2, 0, 1))))
    # Normalize
    transform_norm = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    norm_img = transform_norm(img).unsqueeze(0)

    return img, norm_img

def convert_to_gray(x, percentile=99):
    """
    Args:
        x: torch tensor with shape of (1, 3, H, W)
        percentile: int
    Return:
        result: shape of (1, 1, H, W)
    """
    x_2d = torch.abs(x).sum(dim=1).squeeze(0)
    v_max = np.percentile(x_2d, percentile)
    v_min = torch.min(x_2d)
    torch.clamp_((x_2d - v_min) / (v_max - v_min), 0, 1)
    return x_2d.unsqueeze(0).unsqueeze(0)

def write_video(inputpath, outputname, img_num, fps=10):
    """Generate videos
    Args:
        input_path: the path for input Images
        output_name: the output name for the video
        img_num: the number of the input Images
        fps: frames per second
    """

    fourcc = cv2.VideoWriter_fourcc('M', 'J', 'P', 'G')
    videoWriter = cv2.VideoWriter(outputname, fourcc, fps, (1000, 1000))
    for i in range(img_num):
        img_no = i + 1
        print(inputpath + 'video' + str(img_no) + '.jpg')
        cv_img = cv2.imread(inputpath + 'video' + str(img_no) + '.jpg', 1)
        videoWriter.write(cv_img)
    videoWriter.release()

max_stars_repo_name: wofmanaf/Group-CAM

max_issues_repo_name: wofmanaf/Group-CAM

max_forks_repo_name: wofmanaf/Group-CAM


text:
from sympy import S, pi, GoldenRatio as phi, sqrt
from sympy.ntheory.continued_fraction import \
    (continued_fraction_periodic as cf_p,
     continued_fraction_iterator as cf_i,
     continued_fraction_convergents as cf_c,
     continued_fraction_reduce as cf_r,
     continued_fraction as cf)
from sympy.utilities.pytest import raises

def test_continued_fraction():
    assert cf_p(1, 1, 10, 0) == cf_p(1, 1, 0, 1)
    assert cf_p(1, -1, 10, 1) == cf_p(-1, 1, 10, -1)
    t = sqrt(2)
    assert cf((1 + t)*(1 - t)) == cf(-1)
    for n in [0, 2, S(2)/3, sqrt(2), 3*sqrt(2), 1 + 2*sqrt(3)/5,
            (2 - 3*sqrt(5))/7, 1 + sqrt(2), (-5 + sqrt(17))/4]:
        assert (cf_r(cf(n)) - n).expand() == 0
        assert (cf_r(cf(-n)) + n).expand() == 0
    raises(ValueError, lambda: cf(sqrt(2 + sqrt(3))))
    raises(ValueError, lambda: cf(sqrt(2) + sqrt(3)))
    raises(ValueError, lambda: cf(pi))
    raises(ValueError, lambda: cf(.1))

    raises(ValueError, lambda: cf_p(1, 0, 0))
    raises(ValueError, lambda: cf_p(1, 1, -1))
    assert cf_p(4, 3, 0) == [1, 3]
    assert cf_p(0, 3, 5) == [0, 1, [2, 1, 12, 1, 2, 2]]
    assert cf_p(1, 1, 0) == [1]
    assert cf_p(3, 4, 0) == [0, 1, 3]
    assert cf_p(4, 5, 0) == [0, 1, 4]
    assert cf_p(5, 6, 0) == [0, 1, 5]
    assert cf_p(11, 13, 0) == [0, 1, 5, 2]
    assert cf_p(16, 19, 0) == [0, 1, 5, 3]
    assert cf_p(27, 32, 0) == [0, 1, 5, 2, 2]
    assert cf_p(1, 2, 5) == [[1]]
    assert cf_p(0, 1, 2) == [1, [2]]
    assert cf_p(6, 7, 49) == [1, 1, 6]
    assert cf_p(3796, 1387, 0) == [2, 1, 2, 1, 4]
    assert cf_p(3245, 10000) == [0, 3, 12, 4, 13]
    assert cf_p(1932, 2568) == [0, 1, 3, 26, 2]
    assert cf_p(6589, 2569) == [2, 1, 1, 3, 2, 1, 3, 1, 23]

    def take(iterator, n=7):
        res = []
        for i, t in enumerate(cf_i(iterator)):
            if i >= n:
                break
            res.append(t)
        return res

    assert take(phi) == [1, 1, 1, 1, 1, 1, 1]
    assert take(pi) == [3, 7, 15, 1, 292, 1, 1]

    assert list(cf_i(S(17)/12)) == [1, 2, 2, 2]
    assert list(cf_i(S(-17)/12)) == [-2, 1, 1, 2, 2]

    assert list(cf_c([1, 6, 1, 8])) == [S(1), S(7)/6, S(8)/7, S(71)/62]
    assert list(cf_c([2])) == [S(2)]
    assert list(cf_c([1, 1, 1, 1, 1, 1, 1])) == [S.One, S(2), S(3)/2, S(5)/3,
                                                 S(8)/5, S(13)/8, S(21)/13]
    assert list(cf_c([1, 6, S(-1)/2, 4])) == [S.One, S(7)/6, S(5)/4, S(3)/2]

    assert cf_r([1, 6, 1, 8]) == S(71)/62
    assert cf_r([3]) == S(3)
    assert cf_r([-1, 5, 1, 4]) == S(-24)/29
    assert (cf_r([0, 1, 1, 7, [24, 8]]) - (sqrt(3) + 2)/7).expand() == 0
    assert cf_r([1, 5, 9]) == S(55)/46
    assert (cf_r([[1]]) - (sqrt(5) + 1)/2).expand() == 0
    assert cf_r([-3, 1, 1, [2]]) == -1 - sqrt(2)

max_stars_repo_name: Geektrovert/sympy

reason: from sympy


text:
# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:
### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
#
#   See COPYING file distributed along with the PyMVPA package for the
#   copyright and license terms.
#
### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
"""NeuroImaging Markup Language (NIML) support.

Supports storing most typical values (samples, feature attributes, sample
attributes, dataset attributes) that are in a dataset in NIML format, as
long as these values are array-like.

Notes
-----
No support for 'sophisticated' values such as Mappers

.. versionadded:: 2.3.0

"""

__docformat__ = 'restructuredtext'

import numpy as np

import os

from mvpa2.support.nibabel import afni_niml_dset as niml_dset

from mvpa2.base.collections import SampleAttributesCollection, \
        FeatureAttributesCollection, DatasetAttributesCollection, \
        ArrayCollectable

from mvpa2.base import warning, debug, externals
from mvpa2.datasets.base import Dataset
from mvpa2.base import dataset

if externals.exists('h5py'):
    from mvpa2.base.hdf5 import h5save, h5load

_PYMVPA_PREFIX = 'PYMVPA'
_PYMVPA_SEP = '_'

def from_niml(dset, fa_labels=[], sa_labels=[], a_labels=[]):
    '''Convert a NIML dataset to a Dataset

    Parameters
    ----------
    dset: dict
        Dictionary with NIML key-value pairs, such as obtained from
        mvpa2.support.nibabel.afni_niml_dset.read()
    fa_labels: list
        Keys in dset that are enforced to be feature attributes
    sa_labels: list
        Keys in dset that are enforced to be sample attributes
    a_labels: list
        Keys in dset that are enforced to be dataset attributes

    Returns
    -------
    dataset: mvpa2.base.Dataset
        a PyMVPA Dataset
    '''

    # check for singleton element
    if type(dset) is list and len(dset) == 1:
        # recursive call
        return from_niml(dset[0])

    if not type(dset) is dict:
        raise ValueError("Expected a dict")

    if not 'data' in dset:
        raise ValueError("dset with no data?")

    data = dset['data']
    if len(data.shape) == 1:
        nfeatures = data.shape[0]
        nsamples = 1
    else:
        nfeatures, nsamples = data.shape

    # some labels have predefined destinations
    sa_labels_ = ['labels', 'stats', 'chunks', 'targets'] + sa_labels
    fa_labels_ = ['node_indices', 'center_ids'] + fa_labels
    a_labels_ = ['history'] + a_labels
    ignore_labels = ('data', 'dset_type')

    sa = SampleAttributesCollection(length=nsamples)
    fa = FeatureAttributesCollection(length=nfeatures)
    a = DatasetAttributesCollection()

    labels_collections = [(sa_labels_, sa),
                          (fa_labels_, fa),
                          (a_labels_, a)]

    infix2collection = {'sa':sa,
                      'fa':fa,
                      'a':a}

    infix2length = {'sa':nsamples, 'fa':nfeatures}

    for k, v in dset.iteritems():
        if k in ignore_labels:
            continue

        if k.startswith(_PYMVPA_PREFIX + _PYMVPA_SEP):
            # special PYVMPA field - do the proper conversion
            k_split = k.split(_PYMVPA_SEP)
            if len(k_split) > 2:
                infix = k_split[1].lower()
                collection = infix2collection.get(infix, None)
                if not collection is None:
                    short_k = _PYMVPA_SEP.join(k_split[2:])
                    expected_length = infix2length.get(infix, None)
                    if expected_length:
                        if isinstance(v, np.ndarray) and np.dtype == np.str_:
                            v = str(v)

                        while isinstance(v, basestring):
                            # strings are seperated by ';'
                            # XXX what if this is part of the value
                            # intended by the user?
                            v = v.split(';')

                        if expected_length != len(v):
                            raise ValueError("Unexpected length: %d != %d" %
                                                (expected_length, len(v)))

                        v = ArrayCollectable(v, length=expected_length)

                    collection[short_k] = v
                    continue

        found_label = False

        for label, collection in labels_collections:
            if k in label:
                collection[k] = v
                found_label = True
                break

        if found_label:
            continue

        # try to be smart and deduce this from dimensions.
        # this only works if nfeatures!=nsamples otherwise it would be
        # ambiguous
        # XXX is this ugly?
        if nfeatures != nsamples:
            try:
                n = len(v)
                if n == nfeatures:
                    fa[k] = v
                    continue
                elif n == nsamples:
                    sa[k] = v
                    continue
            except:
                pass

        # don't know what this is - make it a general attribute
        a[k] = v

    ds = Dataset(np.transpose(data), sa=sa, fa=fa, a=a)

    return ds

def to_niml(ds):
    '''Convert a Dataset to a NIML dataset

    Parameters
    ----------
    dataset: mvpa2.base.Dataset
        A PyMVPA Dataset

    Returns
    -------
    dset: dict
        Dictionary with NIML key-value pairs, such as obtained from
        mvpa2.support.nibabel.afni_niml_dset.read()
     '''
    if isinstance(ds, np.ndarray):
        ds = Dataset(ds)

    dset = dict(data=np.transpose(ds.samples))

    node_indices_labels = ('node_indices', 'center_ids', 'ids', 'roi_ids')
    node_indices = _find_node_indices(ds, node_indices_labels)
    if not node_indices is None:
        dset['node_indices'] = node_indices

    sample_labels = ('labels', 'targets')
    labels = _find_sample_labels(ds, sample_labels)
    if not labels is None:
        dset['labels'] = labels

    attr_labels = ('a', 'fa', 'sa')

    # a few labels are directly used in NIML dsets
    # without prefixing it with a pyMVPA string
    # for (dataset, feature, sample) attributes
    # here we define two for sample attributes
    attr_special_labels = ([], [], ['labels', 'stats'])

    for i, attr_label in enumerate(attr_labels):
        attr = getattr(ds, attr_label)
        special_labels = attr_special_labels[i]
        for k in attr.keys():
            v = attr[k]
            if hasattr(v, 'value'):
                v = v.value

            if k in special_labels:
                long_key = k
            else:
                long_key = _PYMVPA_SEP.join((_PYMVPA_PREFIX,
                                             attr_label.upper(), k))

            dset[long_key] = v

    return dset

def hstack(dsets, pad_to_feature_index=None, hstack_method='drop_nonunique',
                set_empty_value=0.):
    '''Stacks NIML datasets while considering node indices

    Parameters
    ----------
    dsets: list
        datasets to be stacked
    pad_to_feature_index: list or int or None
        If a list then it should be of the same length as dsets and indicates
        to which node index the input should be padded. A single int means
        that the same value is used for all dset in dsets. None means
        no padding, and is only allowed for non-sparse datasets.
    hstack_method: str:
        How datasets are stacked; see dataset.hstack.
    set_empty_value: float
        Value to which empty (padded) dataset values are set.

    Returns
    dset: Dataset
        Data combined from all dset in dsets.
    '''

    n = len(dsets)

    # make sure pad_to_feature_index has n values
    if pad_to_feature_index is None or type(pad_to_feature_index) is int:
        pad_to_feature_index = [pad_to_feature_index] * n
    elif len(pad_to_feature_index) != n:
        raise ValueError("illegal pad_to_feature_index: expected list or int")

    # labels that can contain node indices
    node_indices_labels = ('node_indices', 'center_ids', 'ids', 'roi_ids')
    node_indices = []

    # allocate space for output
    padded_dsets = []
    hstack_indices = []
    first_node_index = 0
    for i, (dset, pad_to) in enumerate(zip(dsets, pad_to_feature_index)):
        # get node indices in this dataset
        node_index = _find_node_indices(dset, node_indices_labels)
        if node_index is None:
            node_index = np.arange(dset.nfeatures)
        max_node_index = np.max(node_index)

        # make a stripped version - without node index labels
        stripped_dset = dset.copy()
        for label in node_indices_labels:
            if label in stripped_dset.fa:
                stripped_dset.fa.pop(label)

        # see if padding is needed
        if pad_to is None or pad_to == max_node_index + 1:
            if not np.array_equal(np.arange(max_node_index + 1), np.sort(node_index)):
                raise ValueError("Sparse input %d: need pad_to input" % (i + 1))
            padded_dset = stripped_dset
            other_index = np.arange(0)
        else:
            # have to use empty values
            nfeatures_empty = pad_to - dset.nfeatures
            if nfeatures_empty < 0:
                raise ValueError("Dataset has %d features, cannot pad "
                                    "to %d" % (dset.nfeatures, pad_to))

            # make empty array
            empty_arr = np.zeros((dset.nsamples, nfeatures_empty),
                                    dtype=dset.samples.dtype) + set_empty_value
            empty_dset = Dataset(empty_arr, sa=stripped_dset.sa.copy(deep=True))

            # combine current dset and empty array
            padded_dset = dataset.hstack((stripped_dset, empty_dset), hstack_method)

            # set the proper node indices
            other_index = np.setdiff1d(np.arange(pad_to), node_index)

        # sanity check to make sure that indices are ok
        # XXX could be more informative
        if len(np.setdiff1d(node_index, np.arange(pad_to or max_node_index + 1))):
            raise ValueError("Illegal indices")

        hstack_index = node_index + first_node_index
        hstack_other_index = other_index + first_node_index
        first_node_index += pad_to or (max_node_index + 1) # prepare for next iteration

        padded_dsets.append(padded_dset)
        hstack_indices.append(hstack_index)
        if len(other_index):
            hstack_indices.append(hstack_other_index)

    hstack_dset = dataset.hstack(padded_dsets, hstack_method)
    hstack_indices = np.hstack(hstack_indices)

    hstack_dset.fa[node_indices_labels[0]] = hstack_indices

    return hstack_dset

def _find_sample_labels(dset, sample_labels):
    '''Helper function to find labels in this dataset.
    Looks for any in sample_labels and returns the first one
    that matches '''
    use_label = None

    dset_keys = dset.sa.keys()
    for label in sample_labels:
        if label in dset_keys:
            sample_label = dset.sa[label].value
            if isinstance(sample_label, basestring):
                # split using
                sample_label = sample_label.split(';')

            # they can be of any type so ensure they are strings
            sample_label_list = [str(i) for i in sample_label]
            if len(sample_label_list) != dset.nsamples:
                # unlike node indices here we are more lenient
                # so not throw an exception but just continue
                continue

            use_label = label

            # do not look for any other labels
            break

    return None if use_label is None else sample_label_list

def _find_node_indices(dset, node_indices_labels):
    '''Helper function to find node indices in this dataset
    Sees if any of the node_indices_labels is a feature attribute
    in the dataset and returns it. If they are multiple matches
    ensure they are identical, otherwise raise an error.
    A use case is searchlight results that assignes center_ids as
    a feature attributes, but it should be named node_indices
    before conversion to NIML format'''

    use_label = None

    dset_keys = dset.fa.keys()
    for label in node_indices_labels:
        if label in dset_keys:
            if use_label is None:
                # make vector and ensure all integer values
                node_indices = dset.fa[label].value
                node_indices = np.asarray(node_indices).ravel()
                if len(node_indices) != dset.nfeatures:
                    raise ValueError("Node indices mismatch: found %d values "
                                     " but dataset has %d features" %
                                     (len(node_indices, dset.nfeatures)))
                node_indices_int = np.asarray(node_indices, dtype=np.int)
                if not np.array_equal(node_indices_int, node_indices):
                    raise ValueError("Node indices should have integer values")
                use_label = label

            else:
                if not np.array_equal(dset.fa[label].value, node_indices_int):
                    raise ValueError("Different indices for feature attributes"
                                     " %s and %s" % (use_key, label))

    return None if use_label is None else node_indices_int

def write(fn, ds, form='binary'):
    '''Write a Dataset to a file in NIML format

    Parameters
    ----------
    fn: str
        Filename
    ds: mvpa2.base.Dataset
        Dataset to be stored
    form: str
        Data format: 'binary' or 'text' or 'base64'
    '''
    niml_ds = to_niml(ds)
    niml_dset.write(fn, niml_ds, form=form)

def read(fn):
    '''Read a Dataset from a file in NIML format

    Parameters
    ----------
    fn: str
        Filename
    '''

    readers_converters = [(niml_dset.read, from_niml)]
    if externals.exists('h5py'):
        readers_converters.append((h5load, None))

    for reader, converter in readers_converters:
        try:
            r = reader(fn)
            if converter:
                r = converter(r)
            return r

        except:
            pass

    raise ValueError("Unable to read %s" % fn)

def from_any(x):
    '''Get a Dataset from the input

    Parameters
    ----------
    x: str or dict or Dataset
        Filename, or NIML-dictionary, or a Dataset itself

    Returns
    -------
    ds: mvpa2.base.Dataset
        Dataset instance
    '''
    if isinstance(x, basestring):
        return read(x)
    elif isinstance(x, dict):
        return from_niml(x)
    elif isinstance(x, Dataset):
        return x

    raise ValueError("Not supported: %r" % (x,))

max_stars_repo_name: thomastweets/PyMVPA

max_issues_repo_name: thomastweets/PyMVPA

max_forks_repo_name: thomastweets/PyMVPA


text:
# coding=utf-8
# Copyright 2022 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Create pose regression dataset.

Given a CAD directory structured like:
CAD
  Object Class
    Object files

Render objects at different rotations and save the image and angle of rotation.

Warning: Overwrites pusher.xml in gym package.
"""
from __future__ import print_function
import contextlib
import logging
import os
import pickle
import random
import tempfile

from absl import app
from absl import flags
import gym
import matplotlib
matplotlib.use("Agg")
from matplotlib import pyplot as plt  # pylint: disable=g-import-not-at-top
import numpy as np
import stl
from stl import mesh
import tensorflow.compat.v1 as tf

logging.getLogger("stl").setLevel(logging.ERROR)

flags.DEFINE_string("CAD_dir", None,
                    "Directory of CAD models.")
flags.DEFINE_string("data_dir", None,
                    "Directory where generated data is stored.")
FLAGS = flags.FLAGS

def find_mins_maxs(obj):
  """Find the max dimensions.

  So we can know the bounding box, getting the height, width, length
  (because these are the step size)...

  Args:
    obj: Mesh object.

  Returns:
    minx: Extreme dimension.
    maxx: Extreme dimension.
    miny: Extreme dimension.
    maxy: Extreme dimension.
    minz: Extreme dimension.
    maxz: Extreme dimension.
  """
  minx = maxx = miny = maxy = minz = maxz = None
  for p in obj.points:
    # p contains (x, y, z)
    if minx is None:
      minx = p[stl.Dimension.X]
      maxx = p[stl.Dimension.X]
      miny = p[stl.Dimension.Y]
      maxy = p[stl.Dimension.Y]
      minz = p[stl.Dimension.Z]
      maxz = p[stl.Dimension.Z]
    else:
      maxx = max(p[stl.Dimension.X], maxx)
      minx = min(p[stl.Dimension.X], minx)
      maxy = max(p[stl.Dimension.Y], maxy)
      miny = min(p[stl.Dimension.Y], miny)
      maxz = max(p[stl.Dimension.Z], maxz)
      minz = min(p[stl.Dimension.Z], minz)
  return minx, maxx, miny, maxy, minz, maxz

class MJCModel(object):
  """Mujoco Model."""

  def __init__(self, name):
    self.name = name
    self.root = MJCTreeNode("mujoco").add_attr("model", name)

  @contextlib.contextmanager
  def asfile(self):
    """Usage information.

    Usage:
    model = MJCModel('reacher')
    with model.asfile() as f:
        print f.read()  # prints a dump of the model

    Yields:
      f: File.
    """
    with tempfile.NamedTemporaryFile(
        mode="w+b", suffix=".xml", delete=True) as f:
      self.root.write(f)
      f.seek(0)
      yield f

  def open(self):
    self.file = tempfile.NamedTemporaryFile(
        mode="w+b", suffix=".xml", delete=True)
    self.root.write(self.file)
    self.file.seek(0)
    return self.file

  def save(self, path):
    with open(path, "w") as f:
      self.root.write(f)

  def close(self):
    self.file.close()

class MJCModelRegen(MJCModel):

  def __init__(self, name, regen_fn):
    super(MJCModelRegen, self).__init__(name)
    self.regen_fn = regen_fn

  def regenerate(self):
    self.root = self.regen_fn().root

class MJCTreeNode(object):
  """Mujoco Tree Node."""

  def __init__(self, name):
    self.name = name
    self.attrs = {}
    self.children = []

  def add_attr(self, key, value):
    if isinstance(value, str):
      pass
    elif isinstance(value, list) or isinstance(value, np.ndarray):
      value = " ".join([str(val) for val in value])

    self.attrs[key] = value
    return self

  def __getattr__(self, name):

    def wrapper(**kwargs):
      newnode = MJCTreeNode(name)
      for (k, v) in kwargs.items():  # iteritems in python2
        newnode.add_attr(k, v)
      self.children.append(newnode)
      return newnode

    return wrapper

  def dfs(self):
    yield self
    if self.children:
      for child in self.children:
        for node in child.dfs():
          yield node

  def write(self, ostream, tabs=0):
    """Write out the object as a string."""
    contents = " ".join(['%s="%s"' % (k, v) for (k, v) in self.attrs.items()])
    if self.children:

      ostream.write("\t" * tabs)
      ostream.write("<%s %s>\n" % (self.name, contents))
      for child in self.children:
        child.write(ostream, tabs=tabs + 1)
      ostream.write("\t" * tabs)
      ostream.write("</%s>\n" % self.name)
    else:
      ostream.write("\t" * tabs)
      ostream.write("<%s %s/>\n" % (self.name, contents))

  def __str__(self):
    s = "<" + self.name
    s += " ".join(['%s="%s"' % (k, v) for (k, v) in self.attrs.items()])
    return s + ">"

def pusher(obj_scale=None,
           obj_mass=None,
           obj_damping=None,
           object_pos=(0.45, -0.05, -0.275),
           distr_scale=None,
           axisangle=(0, 0, 1, 1.5),
           distr_mass=None,
           distr_damping=None,
           goal_pos=(0.45, -0.05, -0.3230),
           distractor_pos=(0.45, -0.05, -0.275),
           mesh_file=None,
           mesh_file_path=None,
           distractor_mesh_file=None,
           friction=(.8, .1, .1),
           table_texture=None,
           distractor_texture=None,
           obj_texture=None,
           table_pos="0 0.5 -0.325",
           table_size="1 1 0.1"):
  """Create the pusher Mujoco object from the mesh file."""
  object_pos, goal_pos, distractor_pos, friction = \
      list(object_pos), list(goal_pos), list(distractor_pos), list(friction)
  # For now, only supports one distractor

  #     if obj_scale is None:
  #         #obj_scale = random.uniform(0.5, 1.0)
  #         obj_scale = 4.0
  if obj_mass is None:
    obj_mass = random.uniform(0.2, 2.0)
  if obj_damping is None:
    obj_damping = random.uniform(0.2, 5.0)
  obj_damping = str(obj_damping)

  if distractor_mesh_file:
    if distr_scale is None:
      distr_scale = random.uniform(0.5, 1.0)
    if distr_mass is None:
      distr_mass = random.uniform(0.2, 2.0)
    if distr_damping is None:
      distr_damping = random.uniform(0.2, 5.0)
    distr_damping = str(distr_damping)

  mjcmodel = MJCModel("arm3d")
  mjcmodel.root.compiler(
      inertiafromgeom="true", angle="radian", coordinate="local")
  mjcmodel.root.option(
      timestep="0.01", gravity="0 0 0", iterations="20", integrator="Euler")
  default = mjcmodel.root.default()
  default.joint(armature="0.04", damping=1, limited="true")
  default.geom(
      friction=friction,
      density="300",
      margin="0.002",
      condim="1",
      contype="0",
      conaffinity="0")

  # Make table
  worldbody = mjcmodel.root.worldbody()
  worldbody.light(diffuse=".5 .5 .5", pos="0 0 5", dir="0 0 -1")
  if table_texture:
    worldbody.geom(
        name="table",
        material="table",
        type="plane",
        pos="0 0.5 -0.325",
        size="1 1 0.1",
        contype="1",
        conaffinity="1")
  else:
    worldbody.geom(
        name="table",
        type="plane",
        pos=table_pos,
        size=table_size,
        contype="1",
        conaffinity="1")

  # Process object physical properties
  if mesh_file is not None:
    mesh_object = mesh.Mesh.from_file(mesh_file)

    vol, cog, inertia = mesh_object.get_mass_properties()
    inertia = np.abs(inertia)
    vol = np.abs(vol)
    cog = np.abs(cog)

    minx, maxx, miny, maxy, minz, maxz = find_mins_maxs(mesh_object)
    max_length = max((maxx - minx), max((maxy - miny), (maxz - minz)))

    if max_length > 0.5:
      obj_scale = 5
    else:
      obj_scale = 7

    scale = obj_scale * 0.0012 * (200.0 / max_length)
    # print('max_length=', max_length)
    object_density = np.abs(obj_mass / (vol * scale * scale * scale))
    object_pos[0] -= scale * (minx + maxx) / 2.0
    object_pos[1] -= scale * (miny + maxy) / 2.0
    object_pos[2] = -0.324 - scale * minz
    object_scale = scale
  if distractor_mesh_file is not None:
    distr_mesh_object = mesh.Mesh.from_file(distractor_mesh_file)
    vol, cog, inertia = distr_mesh_object.get_mass_properties()

    inertia = np.abs(inertia)

    minx, maxx, miny, maxy, minz, maxz = find_mins_maxs(distr_mesh_object)
    max_length = max((maxx - minx), max((maxy - miny), (maxz - minz)))
    distr_scale = distr_scale * 0.0012 * (200.0 / max_length)
    distr_density = distr_mass / (vol * distr_scale * distr_scale * distr_scale)
    distractor_pos[0] -= distr_scale * (minx + maxx) / 2.0
    distractor_pos[1] -= distr_scale * (miny + maxy) / 2.0
    distractor_pos[2] = -0.324 - distr_scale * minz

  ## MAKE DISTRACTOR
  if distractor_mesh_file:
    distractor = worldbody.body(name="distractor", pos=distractor_pos)
    if distractor_mesh_file is None:
      distractor.geom(
          rgba="1 1 1 1",
          type="cylinder",
          size="0.05 0.05 0.05",
          density="0.00001",
          contype="1",
          conaffinity="0")
    else:
      if distractor_texture:
        distractor.geom(
            material="distractor",
            conaffinity="0",
            contype="1",
            density=str(distr_density),
            mesh="distractor_mesh",
            rgba="1 1 1 1",
            type="mesh")
      else:
        distractor.geom(
            conaffinity="0",
            contype="1",
            density=str(distr_density),
            mesh="distractor_mesh",
            rgba="1 1 1 1",
            type="mesh")
    distractor.joint(
        name="distractor_slidey",
        type="slide",
        pos="0 0 0",
        axis="0 1 0",
        range="-10.3213 10.3",
        damping=distr_damping)
    distractor.joint(
        name="distractor_slidex",
        type="slide",
        pos="0 0 0",
        axis="1 0 0",
        range="-10.3213 10.3",
        damping=distr_damping)

  # MAKE TARGET OBJECT
  obj = worldbody.body(name="object", pos=object_pos, axisangle=axisangle)
  if mesh_file is None:
    obj.geom(
        rgba="1 1 1 1",
        type="cylinder",
        size="0.05 0.05 0.05",
        density="0.00001",
        contype="1",
        conaffinity="0")
  else:
    if obj_texture:
      obj.geom(
          material="object",
          conaffinity="0",
          contype="1",
          density=str(object_density),
          mesh="object_mesh",
          rgba="1 1 1 1",
          type="mesh")
    else:
      obj.geom(
          conaffinity="0",
          contype="1",
          density=str(object_density),
          mesh="object_mesh",
          rgba="1 1 1 1",
          type="mesh")
  obj.joint(
      name="obj_slidey",
      type="slide",
      pos="0 0 0",
      axis="0 1 0",
      range="-10.3213 10.3",
      damping=obj_damping)
  obj.joint(
      name="obj_slidex",
      type="slide",
      pos="0 0 0",
      axis="1 0 0",
      range="-10.3213 10.3",
      damping=obj_damping)

  goal = worldbody.body(name="goal", pos=goal_pos)
  goal.geom(
      rgba="1 0 0 0",
      type="cylinder",
      size="0.08 0.001 0.1",
      density="0.00001",
      contype="0",
      conaffinity="0")
  goal.joint(
      name="goal_slidey",
      type="slide",
      pos="0 0 0",
      axis="0 1 0",
      range="-10.3213 10.3",
      damping="0.5")
  goal.joint(
      name="goal_slidex",
      type="slide",
      pos="0 0 0",
      axis="1 0 0",
      range="-10.3213 10.3",
      damping="0.5")

  asset = mjcmodel.root.asset()
  if table_texture:
    asset.texture(name="table", file=table_texture, type="2d")
    asset.material(
        shininess="0.3",
        specular="1",
        name="table",
        rgba="0.9 0.9 0.9 1",
        texture="table")
  asset.mesh(
      file=mesh_file_path, name="object_mesh",
      scale=[object_scale] * 3)  # figure out the proper scale
  if obj_texture:
    asset.texture(name="object", file=obj_texture)
    asset.material(
        shininess="0.3",
        specular="1",
        name="object",
        rgba="0.9 0.9 0.9 1",
        texture="object")

  actuator = mjcmodel.root.actuator()
  actuator.motor(joint="goal_slidex", ctrlrange="-2.0 2.0", ctrllimited="true")
  tips_arm = worldbody.body(name="tips_arm", pos=goal_pos)
  tips_arm.geom(
      rgba="1 0 0 0",
      type="cylinder",
      size="0.08 0.001 0.1",
      density="0.00001",
      contype="0",
      conaffinity="0")
  return mjcmodel

def render_images(stl_file, dest_dir):
  """Render images of rotated object."""
  (sub_path, obj) = os.path.split(stl_file)
  obj = obj[:-4]  # Strip extension
  sub_name = os.path.split(sub_path)[-1]
  des_path = os.path.join(dest_dir, "rotate", sub_name, obj)
  if not os.path.exists(des_path):
    os.makedirs(des_path)

  ylabels = []
  for i in range(100):
    # np.random.seed(i**2)
    # random.seed(i**2)
    x_pos = 0.5  # random.uniform(0,1)
    y_pos = 0  # random.uniform(-1,1)
    angle_pos = random.uniform(0, 2 * np.pi)
    model = pusher(
        mesh_file=stl_file,
        mesh_file_path=stl_file,
        object_pos=(x_pos, y_pos, 0.3),
        axisangle=[0, 0, 1, angle_pos],
        table_pos="0.5 0.2 -0.4",
        table_size="1.0 1.6 1")
    model.save(os.path.join(gym.__path__[0][:-4],
                            "gym/envs/mujoco/assets/pusher.xml"))
    # print(GYM_PATH + '/gym/envs/mujoco/assets/pusher.xml')
    # copy2(args.obj_filepath, GYM_PATH+'/gym/envs/mujoco/assets')

    env = gym.envs.make("Pusher-v2")

    screen = env.render(mode="rgb_array")
    # res = cv2.resize(screen, dsize=(128,128), interpolation=cv2.INTER_AREA)
    res = screen
    fig = plt.figure(figsize=(6, 6))
    ax = fig.add_subplot(111)
    ax.axis("off")
    ax.imshow(res)
    ax.set_position([0, 0, 1, 1])
    plt.savefig(des_path + "/" + str(i) + ".png")
    plt.close()

    ylabel = [i, x_pos, y_pos, angle_pos]
    ylabels.append(ylabel)

  pickle.dump(
      ylabels,
      open(des_path + "/" + "label_" + sub_name + obj, "wb"),
      protocol=2)

def main(_):
  CAD_dir = FLAGS.CAD_dir  # pylint: disable=invalid-name

  # Skip these because they are symmetric
  classes_to_skip = ["bottle", "train"]

  for obj_class in tf.io.gfile.listdir(CAD_dir):
    if obj_class in classes_to_skip:
      continue
    for obj in tf.io.gfile.listdir(os.path.join(CAD_dir, obj_class)):
      if obj.endswith(".stl"):
        print("Rendering %s from object class %s" % (obj, obj_class))
        try:
          render_images(os.path.join(CAD_dir, obj_class, obj), FLAGS.data_dir)
        except Exception:  # pylint: disable=broad-except
          print("Failed to render %s from object class %s" % (obj, obj_class))

if __name__ == "__main__":
  app.run(main)