
xplor.requireVersion("2.24")


#
# slow cooling protocol in torsion angle space for protein G. Uses 
# NOE, J-coupling restraints.
#
# this script performs annealing from an extended structure.
# It is faster than the original anneal.py
#
# CDS 2009/07/24
#

# this checks for typos on the command-line. User-customized arguments can
# also be specified.
#
xplor.parseArguments()


# filename for output structures. This string must contain the STRUCTURE
# literal so that each calculated structure has a unique name. The SCRIPT
# literal is replaced by this filename (or stdin if redirected using <),
# but it is optional.
#
outFilename = "Cks1_STRUCTURE.sa"
numberOfStructures=20   #usually you want to create at least 20 

# protocol module has many high-level helper functions.
#
import protocol

protocol.initRandomSeed()   #set random seed - by time

command = xplor.command

# Load pregenerated .psf and .pdb files.
#
protocol.initParams('parallh22x.pro')
protocol.initParams('parallh22x_ole.lip')
protocol.initStruct('Cks1_psfgen.psf')
protocol.initCoords('Cks1_psfgen_min.pdb')

# generate PSF data from sequence and initialize the correct parameters.
#
#from psfGen import seqToPSF
#seqToPSF('aS_wt.seq', startResid=0, psfFilename='aS_wt.psf', amidate_cterm=True)
#protocol.initCoords(files='aS_wt2.pdb')

# generate random extended initial structure with correct covalent geometry
#
try:
    protocol.genExtendedStructure(maxFixupIters=100)
except protocol.CovalentViolation:
    pass


#Write a pdb to match the generated psf
#from pdbTool import PDBTool
#pdb1=PDBTool("out.pdb")
#pdb1.write()

#
# a PotList contains a list of potential terms. This is used to specify which
# terms are active during refinement.
#
from potList import PotList
potList = PotList()

# parameters to ramp up during the simulated annealing protocol
#
from simulationTools import MultRamp, StaticRamp, InitialParams

rampedParams=[]
highTempParams=[]

# compare atomic Cartesian rmsd with a reference structure
#  backbone and heavy atom RMSDs will be printed in the output
#  structure files
# set to reference input structure
from posDiffPotTools import create_PosDiffPot
refRMSD = create_PosDiffPot("refRMSD","name CA or name C or name N",
                            pdbFile='Cks1_psfgen_min.pdb',
                            psf='Cks1_psfgen.psf',
                            cmpSel="name CA or name C or name N or name O")


# set up NOE potential
noe=PotList('noe')
potList.append(noe)
from noePotTools import create_NOEPot
for (name,scale,file) in [('all',1,"Cks1_nOe.tbl"),
                          #add entries for additional tables
                          ]:
    pot = create_NOEPot(name,file)
    # pot.setPotType("soft") # if you think there may be bad NOEs
    pot.setScale(scale)
    noe.append(pot)
rampedParams.append( MultRamp(2,30, "noe.setScale( VALUE )") )

# set up J coupling - with Karplus coefficients
from jCoupPotTools import create_JCoupPot
jCoup = create_JCoupPot("jcoup","Cks1_jna_coup.tbl",
                        A=6.98,B=-1.38,C=1.72,phase=-60.0)
potList.append(jCoup)

# Set up dihedral angles
from xplorPot import XplorPot
dihedralRestraintFilename="Cks1_angles.tbl"
protocol.initDihedrals(dihedralRestraintFilename,
                      #useDefaults=False  # by default, symmetric sidechain
                                           # restraints are included
                       )
potList.append( XplorPot('CDIH') )
highTempParams.append( StaticRamp("potList['CDIH'].setScale(10)") )
rampedParams.append( StaticRamp("potList['CDIH'].setScale(200)") )
# set custom values of threshold values for violation calculation
#
potList['CDIH'].setThreshold( 5 )



# gyration volume term 
#
from gyrPotTools import create_GyrPot
gyr = create_GyrPot("Vgyr",
                    "resid 1:14") # selection should exclude disordered tails
potList.append(gyr)
rampedParams.append( MultRamp(.002,1,"gyr.setScale(VALUE)") )

# hbdb - hbond database-based term
#
protocol.initHBDB()
potList.append( XplorPot('HBDB') )

#New torsion angle database potential
#
from torsionDBPotTools import create_TorsionDBPot
torsionDB = create_TorsionDBPot('torsionDB')
potList.append( torsionDB )
rampedParams.append( MultRamp(.002,2,"torsionDB.setScale(VALUE)") )

#
# setup parameters for atom-atom repulsive term. (van der Waals-like term)
#
potList.append( XplorPot('VDW') )
rampedParams.append( StaticRamp("protocol.initNBond()") )
rampedParams.append( MultRamp(0.9,0.8,
                              "command('param nbonds repel VALUE end end')") )
rampedParams.append( MultRamp(.004,4,
                              "command('param nbonds rcon VALUE end end')") )
# nonbonded interaction only between CA atoms
highTempParams.append( StaticRamp("""protocol.initNBond(cutnb=100,
                                                        rcon=0.004,
                                                        tolerance=45,
                                                        repel=1.2,
                                                        onlyCA=1)""") )


potList.append( XplorPot("BOND") )
potList.append( XplorPot("ANGL") )
potList['ANGL'].setThreshold( 5 )
rampedParams.append( MultRamp(0.4,1,"potList['ANGL'].setScale(VALUE)") )
potList.append( XplorPot("IMPR") )
potList['IMPR'].setThreshold( 5 )
rampedParams.append( MultRamp(0.1,1,"potList['IMPR'].setScale(VALUE)") )
      


# Give atoms uniform weights, configure bath/molecule friction coeff.
#
protocol.massSetup()


# IVM setup
#   the IVM is used for performing dynamics and minimization in torsion-angle
#   space, and in Cartesian space.
#
from ivm import IVM

dyn = IVM()



#dyn.breakBond("resid 7 and (name C or name CA)")

# configure ivm topology for torsion-angle dynamics

#

protocol.torsionTopology(dyn)



# minc used for final cartesian minimization

#

minc = IVM()

#minc.breakBond("resid 7 and (name C or name CA)")

protocol.initMinimize(minc)



protocol.cartesianTopology(minc)




# object which performs simulated annealing
#
from simulationTools import AnnealIVM
init_t  = 3500.     # Need high temp and slow annealing to converge
cool = AnnealIVM(initTemp =init_t,
                 finalTemp=25,
                 tempStep =12.5,
                 ivm=dyn,
                 rampedParams = rampedParams)


def calcOneStructure(loopInfo):
    """ this function calculates a single structure, performs analysis on the
    structure, and then writes out a pdb file, with remarks.
    """

    # generate a new structure with randomized torsion angles
    #
    from monteCarlo import randomizeTorsions
    randomizeTorsions(dyn)
    protocol.fixupCovalentGeom(maxIters=100,useVDW=1,angleTol=3.5)

    # set torsion angles from restraints
    #
    from torsionTools import setTorsionsFromTable
    setTorsionsFromTable(dihedralRestraintFilename)
    protocol.writePDB(loopInfo.filename()+".init")


    # initialize parameters for high temp dynamics.
    InitialParams( rampedParams )
    # high-temp dynamics setup - only need to specify parameters which
    #   differfrom initial values in rampedParams
    InitialParams( highTempParams )

    # high temp dynamics
    #
    protocol.initDynamics(dyn,
                          potList=potList, # potential terms to use
                          bathTemp=init_t,
                          initVelocities=1,
                          finalTime=100,   # stops at 1000ps or 10000 steps
                          numSteps=1000,   # whichever comes first
                          printInterval=100)

    dyn.setETolerance( init_t/100 )  #used to det. stepsize. default: t/1000 
    dyn.run()

    # initialize parameters for cooling loop
    InitialParams( rampedParams )


    # initialize integrator for simulated annealing
    #
    protocol.initDynamics(dyn,
                          potList=potList,
                          numSteps=1000,       #at each temp: 1000 steps or
                          finalTime=2 ,       # 2ps, whichever is less
                          printInterval=100)

    # perform simulated annealing
    #
    cool.run()
              
              
    # final torsion angle minimization
    #
    protocol.initMinimize(dyn,
                          printInterval=50)
    dyn.run()

    # final all- atomic degrees of freedom minimization
    #
    protocol.initMinimize(minc,
                          potList=potList,
                          dEPred=10)
    minc.run()

    #do analysis and write structure when this function retunrs
    pass



from simulationTools import StructureLoop, FinalParams
StructureLoop(numStructures=numberOfStructures,
              doWriteStructures=True,
              pdbTemplate=outFilename,
              structLoopAction=calcOneStructure,
              genViolationStats=True,
              averageTopFraction=0.5, #report stats on best 50% of structs
              averageSortPots=[potList['BOND'],potList['ANGL'],potList['IMPR'],
                               noe,potList['CDIH']],
              averageContext=FinalParams(rampedParams),
              averageCrossTerms=refRMSD,
              averageFilename="SCRIPT_ave.pdb",
              averagePotList=potList).run()
