optimizeCentre_oop.py

############################################################################
##This is used to find the waters around the high density centre. This script
##use two outputs, which were generated by reasDCD File.py and calculatCentre
##_oop.py. We used several steps to optimize the founded coordidates centre
##and the detected centre to make them consistent with each other.
##The outputfile is a set of water coordidates and according centre index.
##Moreover, the waterInformation arround the centre were write to an outputFile.
##
##Xianqiang Sun
##TheoChem&Bio
##KTH
##2012-05-14
###########################################################################

import numpy
import sys
sys.path.append('/home/x/xiansu/pfs/program/numpy/lib/python2.6/site-packages')
from Numeric import *
from datetime import datetime

waters=open('distance.txt','r')
centreFile=open('waterCentre.txt','r')
optimizedCentreOutput=open('optimizedCentre.txt','w')
waterInforByCentreOutput=open('waterInforByCentre.txt','w')
waterInforByCentreOutput.write('resname    atomid    resnumber    X    Y     Z      segname    frameNo     centreNo'+ '\n')


def formatWaterInfor(self):
    waterInfor=[]
    for line in self:
        line=line.split()
        if line[0]=='TIP3':
            line[1]=line[1][:-1]
            line[-1]=line[-1][1:-2]
        elif line[0]=='frame':
            frameNo=line[1]
        if line[0]=='TIP3':
            line.append(str(frameNo))
            waterInfor.append(line)
    return waterInfor

##formatWaterInfor is used to read the waters information from the input
##file and translate it into list format. The frame No. information were
##also added to the list. the format is [[waterInfor]....]

def getWaterCoord(self,firstFrame,frequency):
    waterCoord=[]
    for waterAtom in self:
        frameNo=waterAtom[-1]
##        print frameNo
        if int(frameNo)>=firstFrame:
            if (int(frameNo)/frequency)-(int(frameNo)/float(frequency))==0:
                atomCoord=waterAtom[3:-2]
                atomCoordFloat=[]
                for coord in atomCoord:
                    coord=float(coord)
                    atomCoordFloat.append(coord)
                waterCoord.append(atomCoordFloat)
    return waterCoord
##getWaterCoord use the output of 'formatWaterInfor' as the input. It will
##read the atom Coordiates and extract them ad the float point numbers. Moreover,
##the firstFrame number and frequency can be seted to save computational time. Each
##coordidates can be save as one element in one list, the format is:[[x1,y1,z2]
##,[x2,y2,z2]...]


def getCentreCoord(self):
    centreCoord=[]
    for centre in self:
        centre=centre.split()
        centreFloat=[]
        for coord in centre:
            coord =float(coord)
            centreFloat.append(coord)
            
        centreCoord.append(centreFloat)
    return centreCoord

##getCentreCoord read the output file from the 'calculatCentre_oop.py'. Read the
##coorditates and save them as list with the float format. The output format is
##the format is:[[x1,y1,z2],[x2,y2,z2]...]

def compareCoord(self,waterCoord,threshould):
    waterByCentre=[]
    for centre in range(len(self)):
        a=numpy.array(self[centre])
        waterByEachCentre=[]
        for i in range(len(waterCoord)):
            b=numpy.array(waterCoord[i])
            dist = numpy.linalg.norm(a-b)
            if dist<=threshould:
                waterByEachCentre.append(waterCoord[i])
        waterByCentre.append(waterByEachCentre)
    return waterByCentre

##compareCoord is used to compare two sets of coordidates with certain threshould.
##the output is a set of coorditates set with ranked by each coordidate in self.
##the format of output is [[[x1,y1,z1]...],[...]].

def compareOneCoorWithSet(self,waterCoord,threshould):
    a=numpy.array(self)
    waterByEachCentre=[]
    for i in range(len(waterCoord)):
        b=numpy.array(waterCoord[i])
        dist = numpy.linalg.norm(a-b)
        if dist<=threshould:
            waterByEachCentre.append(waterCoord[i])
    return waterByEachCentre

##compareCoord is used to compare one coordidate with a set of coordidates with certain threshould.
##the output is a set of coorditates. The format of output is [[x1,y1,z1]...].
    

def getCentre(self):
    centreAll=[]
    for coordSet in self:
        coordNo=0
        xsum=0
        ysum=0
        zsum=0
        for coord in coordSet:
            coordNo+=1
            xsum=xsum+coord[0]
            ysum=ysum+coord[1]
            zsum=zsum+coord[2]
        centre=[]
        centre.append(xsum/coordNo)
        centre.append(ysum/coordNo)
        centre.append(zsum/coordNo)
        centreAll.append(centre)
    return centreAll

##getCentre use a set of coordidate sets to find the centre of these sets. The input can be [[[x,y,z],[x,y,z]..]...]
##the output of this function is centre for each coordidate set [[xcen,ycen,zcen]...]
def getOneCentre(self):

    xsum=0
    ysum=0
    zsum=0
    coordNo=0
    for coord in self:
        coordNo+=1
        xsum=xsum+coord[0]
        ysum=ysum+coord[1]
        zsum=zsum+coord[2]
        centre=[]
        centre.append(xsum/coordNo)
        centre.append(ysum/coordNo)
        centre.append(zsum/coordNo)
    return centre
##getCentre use a set of coordidates to find the centre of these coordidates. The input can be [[x,y,z],[x,y,z]..]
##the output of this function is centre for these coordidates [xcen,ycen,zcen]        
        

def optimizeCentre(self,oldCentre,waterCoord,threshould):
    for i in range(len(self)):
        a=numpy.array(self[i])
        b=numpy.array(oldCentre[i])
        dist=numpy.linalg.norm(a-b)
        if dist>=0.1:
            oldCentre[i]=self[i]
        for number in range(100):
            if dist>=0.01:
                print dist
                print 'optimizing centre ',oldCentre[i],' No ', i
##                oldCentre[i]=self[i]
                c=numpy.array(oldCentre[i])
                optimizedWater=compareOneCoorWithSet(oldCentre[i],waterCoord,threshould)
                optimizedCentre=getOneCentre(optimizedWater)
                oldCentre[i]=optimizedCentre
                d=numpy.array(optimizedCentre)
                dist=numpy.linalg.norm(c-d)

            else:
                print 'the distance is', dist
                break

    return oldCentre



##optimize caompare the data in two list of coordidats(newcentre generated by getCentre, oldcenter by
##compareCentre, and waterfoord for all the water coordidated. the threshould were used in 'compraeOneCoorWithSet'
##One should mention that this function use two other function. therefore, the ranking of this function should be
##take care). At last we obtained several coordidates with optimized centre. Nevermind about the name oldCentre
##at last return. I have updated it.

def getWaterInforFromCentre(self,waterCoord,waterInfor,threshould):
    waterByCentre=[]
    waterInforByCentre=[]
    for centre in range(len(self)):
        a=numpy.array(self[centre])
        waterByEachCentre=[]
        waterInforByEachCentre=[]
        for i in range(len(waterCoord)):
            b=numpy.array(waterCoord[i])
            dist = numpy.linalg.norm(a-b)
            if dist<=threshould:
                waterByEachCentre.append(waterCoord[i])
                waterInfor[i].append(centre)
                waterInforByEachCentre.append(waterInfor[i])
        waterInforByCentre.append(waterInforByEachCentre)
        waterByCentre.append(waterByEachCentre)
    return waterInforByCentre



##getWaterInforFromCentre is used to compare each coordidates in self and waterCoord with certain threshould. Then
##water information will be obtianed with waterInfor. The waterCoord should always matched with waterinfor.
##.in with certain threshould.
##the output is a set of water Information set according to self.
##the format of output is [[[x1,y1,z1]...],[...]].

def writeListSets(self,outputFileName):
    
    for i in self:
##        eachWater=''
        for j in i:
            eachWater=''
            for item in j:
                eachWater=eachWater+str(item)+'    '
            print eachWater
            eachWater=eachWater+'\n'
            outputFileName.write(eachWater)
    outputFileName.close()
##writeListSets use a set of water information list and write it to file. The input of the file looks like
##[[[waterInfor],...]...].
        

def writeCentre(self,outPutFileName):
    for i in self:
        outPutFileName.write(str(i[0])+'    '+str(i[1])+'    '+str(i[2])+'\n')
    outPutFileName.close()

##This function read a set of coorditates and write it to file. The input is:[[x1,y1,z2]
##,[x2,y2,z2]...],outPutFileName is the output file name.

waterInfor=formatWaterInfor(waters)
print len(waterInfor)
waterCoords=getWaterCoord(waterInfor,1,1)
print len(waterCoords)

waterCentre=getCentreCoord(centreFile)
print len(waterCentre)

waterAroundCentre=compareCoord(waterCentre,waterCoords,1.2)

print len(waterAroundCentre)

newWaterCentre=getCentre(waterAroundCentre)

print len(newWaterCentre)

optimizedCentre=optimizeCentre(newWaterCentre,waterCentre,waterCoords,1.2)
print len(optimizedCentre)

waterInforFromCentre=getWaterInforFromCentre(optimizedCentre,waterCoords,waterInfor,1.2)
print len(waterInforFromCentre)
writeListSets(waterInforFromCentre,waterInforByCentreOutput)


writeCentre(optimizedCentre,optimizedCentreOutput)