xplor.requireVersion("2.18")

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


(opts,args) = xplor.parseArguments(["quick"]) # check for command-line typos

quick=False
for opt in opts:
    if opt[0]=="quick":  #specify -quick to just test that the script runs
        quick=True
        pass
    pass


outFilename = "SCRIPT_STRUCTURE.sa"
numberOfStructures=20

if quick:
    numberOfStructures=3
    pass

# protocol module has many high-level helper functions.
#
import protocol
protocol.initRandomSeed(3421)   #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('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("model.pdb")
xplor.simulation.deleteAtoms("not known")

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","name CA or name C or name N",
                            pdbFile='g_xray.pdb',
                            cmpSel="not name H*")

# 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 [ ('t',   -6.5, 0.62),
                        ('b',   -9.9, 0.23) ]:
    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)
     ]:
    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
#
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_sum_3.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_g_all.tbl",
                       useDefaults=0)
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 1:56") # selection should exclude disordered tails
potList.append(gyr)
rampedParams.append( MultRamp(.002,1,"gyr.setScale(VALUE)") )

# hbda - distance/angle bb hbond term
#
protocol.initHBDA('hbda.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
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, 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()

# 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
    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']],
              averageCrossTerms=refRMSD,
              averageTopFraction=0.5, #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()