
xplor.requireVersion("2.34")

#
# slow cooling protocol in torsion angle space for protein G. Uses 
# NOE, RDC, J-coupling restraints.
#
# this version refines from a reasonable model structure.
#
# CDS 2005/05/10
#

xplor.parseArguments()
outFilename = "SCRIPT_STRUCTURE.sa"
numberOfStructures=100   #usually you want to create at least 20 

# protocol module has many high-level helper functions.
#
import protocol
protocol.initRandomSeed()   #explicitly set random seed
protocol.initTopology("cyc")
protocol.initParams("cyc")
protocol.initParams("protein")

import psfGen
psfGen.residueTypes['protein'].append('CYC')

# generate PSF data from sequence and initialize the correct parameters.
#
#from psfGen import seqToPSF
#seqToPSF('protG.seq')
#protocol.initStruct("g_new.psf") # - or from file

# generate a random extended structure with correct covalent geometry
#  saves the generated structure in the indicated file for faster startup
#  next time.
#
#protocol.genExtendedStructure("gb1_extended_%d.pdb" %
#                              protocol.initialRandomSeed())

# or read an existing model
#
protocol.loadPDB("anneal_55.sa")
xplor.simulation.deleteAtoms("not known")

#
# annealing settings
#

command = xplor.command
command("""param
BOND	(name	C1C)	(name	C2C)	1000	1.5118
BOND	(name	C2C)	(name	C3C)	1000	1.5418
BOND	(name	C3C)	(name	C4C)	1000	1.5124
BOND	(name	C4C)	(name	N_C)	1000	1.3697
BOND	(name	C1C)	(name	N_C)	1000	1.3386
BOND	(name	C3C)	(name	CAC)	1000	1.5292
BOND	(name	C4B)	(name	N_B)	1000	1.3465
BOND	(name	C1B)	(name	N_B)	1000	1.3923
BOND	(name	C1D)	(name	C2D)	1000	1.3967
BOND	(name	C3D)	(name	C4D)	1000	1.4031
BOND	(name	C3A)	(name	C4A)	1000	1.3967
BOND	(name	N_D)	(name	C1D)	1000	1.3675
BOND	(name	N_D)	(name	C4D)	1000	1.3863
BOND	(name	N_A)	(name	C1A)	1000	1.3863
BOND	(name	N_A)	(name	C4A)	1000	1.3675

ANGLE	(name	H2C)	(name	C2C)	(name	C1C)	500	110.5
ANGLE	(name	H2C)	(name	C2C)	(name	C3C)	500	110.5
ANGLE	(name	CMC)	(name	C2C)	(name	C1C)	500	110.5
ANGLE	(name	CMC)	(name	C2C)	(name	C3C)	500	110.5
ANGLE	(name	H3C)	(name	C3C)	(name	C2C)	500	110.5
ANGLE	(name	H3C)	(name	C3C)	(name	C4C)	500	110.5
ANGLE	(name	H3C)	(name	C3C)	(name	CAC)	500	110.6
ANGLE	(name	CAC)	(name	C3C)	(name	C2C)	500	110.6
ANGLE	(name	CAC)	(name	C3C)	(name	C4C)	500	110.6

ANGLE	(name	C1C)	(name	C2C)	(name	C3C)	500	104.1
ANGLE	(name	C2C)	(name	C3C)	(name	C4C)	500	103.8
ANGLE	(name	C3C)	(name	C4C)	(name	N_C)	500	108.8
ANGLE	(name	C1C)	(name	N_C)	(name	C4C)	500	113.8
ANGLE	(name	C2C)	(name	C1C)	(name	N_C)	500	109.5

ANGLE	(name	C4B)	(name	N_B)	(name	C1B)	500	108.3
ANGLE	(name	N_D)	(name	C1D)	(name	C2D)	500	108.4
ANGLE	(name	N_D)	(name	C4D)	(name	C3D)	500	107.8
ANGLE	(name	N_A)	(name	C4A)	(name	C3A)	500	108.4
ANGLE	(name	C1D)	(name	C2D)	(name	C3D)	500	107.9
ANGLE	(name	C4D)	(name	C3D)	(name	C2D)	500	107.6
ANGLE	(name	H_C)	(name	N_C)	(name	C1C)	500	123.1
ANGLE	(name	H_C)	(name	N_C)	(name	C4C)	500	123.1
ANGLE	(name	H_B)	(name	N_B)	(name	C4B)	500	125.85
ANGLE	(name	H_B)	(name	N_B)	(name	C1B)	500	125.85

IMPROPER (name  SG) (name C3C) (name CBC) (name HAC) 500.0 0 -65.5000
IMPROPER (name CAC) (name C2C) (name C4C) (name H3C) 500.0 0 -64.6000
IMPROPER (name CMC) (name C1C) (name C3C) (name H2C) 500.0 0  64.6000
end""")
protocol.initParams("protein")
#protocol.fixupCovalentGeom(maxIters=100,useVDW=1)

#
# 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
#
from posDiffPotTools import create_PosDiffPot
refRMSD = create_PosDiffPot("refRMSD","resid 599:754 and (name CA or name C or name N)",
                            pdbFile='gaf6012model+582.pdb')

# orientation Tensor - used with the dipolar coupling term
#  one for each medium
#   For each medium, specify a name, and initial values of Da, Rh.
#
from varTensorTools import create_VarTensor
media={}
#                        medium  Da   rhombicity
for (medium,Da,Rh) in [ ('phage',   17.98, 0.37) ]:
    oTensor = create_VarTensor(medium)
    oTensor.setDa(Da)
    oTensor.setRh(Rh)
    media[medium] = oTensor
    pass
    

# dipolar coupling restraints for protein amide NH.  
#
# collect all RDCs in the rdcs PotList
#
# RDC scaling. Three possible contributions.
#   1) gamma_A * gamma_B / r_AB^3 prefactor. So that the same Da can be used
#      for different expts. in the same medium. Sometimes the data is
#      prescaled so that this is not needed. scale_toNH() is used for this.
#      Note that if the expt. data has been prescaled, the values for rdc rmsd
#      reported in the output will relative to the scaled values- not the expt.
#      values.
#   2) expt. error scaling. Used here. A scale factor equal to 1/err^2
#      (relative to that for NH) is used.
#   3) sometimes the reciprocal of the Da^2 is used if there is a large
#      spread in Da values. Not used here.
#
from rdcPotTools import create_RDCPot, scale_toNH
rdcs = PotList('rdc') 
for (medium,expt,file,                 scale) in \
    [('phage','NH' ,'RDC.tbl'       ,1)
     ]:
    rdc = create_RDCPot("%s_%s"%(medium,expt),file,media[medium])

    #1) scale prefactor relative to NH
    #   see python/rdcPotTools.py for exact calculation
    # scale_toNH(rdc) - not needed for these datasets -
    #                        but non-NH reported rmsd values will be wrong.

    #3) Da rescaling factor (separate multiplicative factor)
    # scale *= ( 1. / rdc.oTensor.Da(0) )**2
    rdc.setScale(scale)
    rdc.setShowAllRestraints(1) #all restraints are printed during analysis
    rdc.setThreshold(1.5)       # in Hz
    rdcs.append(rdc)
    pass
potList.append(rdcs)
rampedParams.append( MultRamp(0.05,5.0, "rdcs.setScale( VALUE )") )

# calc. initial tensor orientation
# and setup tensor calculation during simulated annealing
#
from varTensorTools import calcTensorOrientation, calcTensor
for medium in media.keys():
    calcTensorOrientation(media[medium])
    rampedParams.append( StaticRamp("calcTensor(media['%s'])" % medium) )
    pass

# set up NOE potential
noe=PotList('noe')
potList.append(noe)
from noePotTools import create_NOEPot
for (name,scale,file) in [('all',1,"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","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
protocol.initDihedrals("dihed.tbl",
                       #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 ) #5 degrees is the default value, though



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

# hbda - distance/angle bb hbond 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)") )
potList.append( XplorPot('DIHE') )
potList['DIHE'].setThreshold( 10 )


# Give atoms uniform weights, except for the anisotropy axis
#
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.group("resname CYC and (name C1C or name C2C or name C3C or name C4C or name C1D or name C2D or name C3D or name C4D or name C1A or name C2A or name C3A or name C4A or name CHA or name CHD or name N_A or name N_C or name N_D)")

# initially minimize in Cartesian space with only the covalent constraints.
#   Note that bonds, angles and many impropers can't change with the 
#   internal torsion-angle dynamics
#   breaks bonds topologically - doesn't change force field
#
#dyn.potList().add( XplorPot("BOND") )
#dyn.potList().add( XplorPot("ANGL") )
#dyn.potList().add( XplorPot("IMPR") )
#
#dyn.breakAllBondsIn("not resname ANI")
#import varTensorTools
#for m in media.values():
#    m.setFreedom("fix")                 #fix tensor parameters
#    varTensorTools.topologySetup(dyn,m) #setup tensor topology
#
#protocol.initMinimize(dyn,numSteps=1000)
#dyn.run()

# reset ivm topology for torsion-angle dynamics
#
#dyn.reset()

for m in media.values():
#    m.setFreedom("fixDa, fixRh")        #fix tensor Rh, Da, vary orientation
    m.setFreedom("varyDa, varyRh")      #vary tensor Rh, Da, vary orientation
protocol.torsionTopology(dyn)

# minc used for final cartesian minimization
#
minc = IVM()
protocol.initMinimize(minc)

for m in media.values():
    m.setFreedom("varyDa, varyRh")    #allow all tensor parameters float here
    pass
protocol.cartesianTopology(minc)



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

def accept(potList):
    """
    return True if current structure meets acceptance criteria
    """
    if potList['noe'].violations()>0:
        return False
    if potList['rdc'].rms()>1.2: #this might be tightened some
        return False
    if potList['CDIH'].violations()>0:
        return False
    if potList['BOND'].violations()>0:
        return False
    if potList['ANGL'].violations()>0:
        return False
    if potList['IMPR'].violations()>1:
        return False
    
    return True
    
def calcOneStructure(loopInfo):
    """ this function calculates a single structure, performs analysis on the
    structure, and then writes out a pdb file, with remarks.
    """

    # 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=10,    # stops at 10ps or 5000 steps
                          numSteps=5000,   # 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=100,       #at each temp: 100 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- atom minimization
    #
    protocol.initMinimize(minc,
                          potList=potList,
                          dEPred=10)
    minc.run()

    #do analysis and write structure when this function returns
    pass

from simulationTools import StructureLoop, FinalParams
StructureLoop(numStructures=numberOfStructures,
              doWriteStructures=True,
              pdbTemplate=outFilename,
              structLoopAction=calcOneStructure,
              genViolationStats=True,
              averagePotList=potList,
              averageSortPots=[potList['BOND'],potList['ANGL'],potList['IMPR'],
                               noe,rdcs,potList['CDIH'],potList['DIHE']],
              averageCrossTerms=refRMSD,
              averageTopFraction=0.2, #report only on best 50% of structs
#              averageAccept=accept,   #only use structures which pass accept()
              averageContext=FinalParams(rampedParams),
              averageFilename="SCRIPT_ave.pdb",    #generate regularized ave structure
              averageFitSel="name CA",
              averageCompSel="not resname ANI and not name H*"     ).run()
