
xplor.requireVersion("2.14.4")

#
# 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() # check for typos on the command-line

outFilename = "SCRIPT_STRUCTURE.sa"
numberOfStructures=100

# protocol module has many high-level helper functions.
#
import protocol
protocol.initRandomSeed(38755341)   #explicitly set random seed

#
# annealing settings
#

command = xplor.command

protocol.initParams("protein")

# generate PSF data from sequence and initialize the correct parameters.
#
from psfGen import seqToPSF
seqToPSF('protein.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.initCoords("model.pdb")
protocol.addUnknownAtoms()
protocol.fixupCovalentGeom(maxIters=100,useVDW=1)

#
# a PotList conatins 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=[]

# 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 [('p',   -16, 0.6)
                       ]:
    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 \
#    [('t','NH' ,'tmv107_nh.tbl'       ,1),
#     ('t','NCO','tmv107_nc.tbl'       ,.05),
#     ('t','HNC','tmv107_hnc.tbl'      ,.108),
#     ('b','NH' ,'bicelles_new_nh.tbl' ,1),
#     ('b','NCO','bicelles_new_nc.tbl' ,.05),
#     ('b','HNC','bicelles_new_hnc.tbl',.108)
#     ]:
for (medium,expt,file,                 scale) in \
    [('p','NH' ,'protein_pf1_jnh.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 *= ( 9.9 / rdcs[name].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)
# MEH: originally at 0.05,5.0
rampedParams.append( MultRamp(0.05,20.0, "rdcs.setScale( VALUE )") ) 

# calc. initial tensor orientation
#
from varTensorTools import calcTensorOrientation
for medium in media.values():
    calcTensorOrientation(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)
    
# changed from 2, 30 to 2, 50 by MEH
rampedParams.append( MultRamp(2, 50, "noe.setScale( VALUE )") ) 

# set up J coupling - with Karplus coefficients
from jCoupPotTools import create_JCoupPot
jCoup = create_JCoupPot("jcoup","empty.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("empty.tbl",
                       useDefaults=0)
potList.append( XplorPot('CDIH') )
highTempParams.append( StaticRamp("potList['CDIH'].setScale(50)") )
rampedParams.append( StaticRamp("potList['CDIH'].setScale(500)") )
# set custom values of threshold values for violation calculation
#
potList['CDIH'].setThreshold( 1 )



# radius of gyration term 
#
#protocol.initCollapse("resid 11:140",
#                      Rtarget=13.0)  # Rgyr = (2.2 * Nresidues^0.38) - 1
#potList.append( XplorPot('COLL') )

# hbda - distance/angle bb hbond term
#
protocol.initHBDA('empty.tbl')
potList.append( XplorPot('HBDA') )

#Rama torsion angle database
#
protocol.initRamaDatabase()
potList.append( XplorPot('RAMA') )
rampedParams.append( MultRamp(.002,1,"potList['RAMA'].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 for high Temp MD
highTempParams.append( StaticRamp("""protocol.initNBond(cutnb=100,
                                                        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, except for the anisotropy axis
#
from atomAction import SetProperty
import varTensorTools
AtomSel("not resname ANI").apply( SetProperty("mass",100.) )
varTensorTools.massSetup(media.values(),300)
AtomSel("all            ").apply( SetProperty("fric",10.) )


# 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()

# 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,oTensors=media.values())

# 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,oTensors=media.values())



# 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()>2:
        return False
    if potList['rdc'].rms()>0.8: #this might be tightened some
        return False
    if potList['CDIH'].violations()>2:
        return False
    if potList['BOND'].rms()>0.1:
        return False
    if potList['ANGL'].rms()>2:
        return False
    if potList['IMPR'].rms()>4:
        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=200,       #at each temp: 200 steps or
                          finalTime=.4 ,       # .4ps, 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
    loopInfo.writeStructure(potList)
    pass



from simulationTools import StructureLoop, FinalParams
StructureLoop(numStructures=numberOfStructures,
              pdbTemplate=outFilename,
              structLoopAction=calcOneStructure,
              genViolationStats=1,
              averagePotList=potList,
#              averageSortPots=[potList['BOND'],potList['ANGL'],potList['IMPR'],
#                               noe,rdcs,potList['CDIH']],
              averageSortPots=[potList['BOND'],potList['ANGL'],potList['IMPR'],
                               noe,rdcs,potList['CDIH']],
              averageTopFraction=0.4, #report only on top fraction of structs
#              averageAccept=accept,   #only use structures which pass accept()
              averageContext=FinalParams(rampedParams),
              averageFilename="ave.pdb",    #generate average structure
              averageFitSel="name CA",
              averageCompSel="not hydro"
              ).run()

