# YASARA MACRO
# TOPIC:       5. Structure Determination
# TITLE:       Refine an NMR ensemble in explicit solvent
# REQUIRES:    YASARA Structure and NMR Structure Determination Module
# AUTHOR:      Elmar Krieger, Sander Nabuurs, Chris Spronk
# LICENSE:     GPL
# DESCRIPTION: This macro refines an NMR ensemble in explicit solvent (HOH or DMSO) to increase the accuracy

# To refine an ensemble in water, you need:
# - An ensemble of correctly folded structures in PDB format
#   default filenames are 'ensemble001.pdb', 'ensemble002.pdb' etc.
# - A file with distance and dihedral angle restraints,
#   default filename is 'restraints.tbl'
#
# You can either set the target by clicking on Options > Macro > Set target,
# by providing it as command line argument (see docs at Essentials > The command line),
# or by uncommenting the line below and specifying it directly.
#macrotarget = 'c:\MyProject\1crn'

# If you want to change the defaults, do it here, before the defaults are included.
# (look in nmr_setdefaults to see what you can change)

# Set those defaults that have not been set already
include nmr_setdefaults

# Step 3: Trade speed for accuracy and use the best YASARA got in the final stage
# ===============================================================================
# This means explicit solvent, Particle Mesh Ewald and the YASARA force field.
RestrainPar (defaultpar)
SimSteps 5
Console Off
LogAs (logfile),append=Yes, print 'Explicit solvent refinement:'
for i=start to start+structures-1
  Clear
  # Load the ensemble member, maybe the user just started here 
  exists = FileSize (ensemblefile)(i).pdb
  if exists
    LoadPDB (ensemblefile)(i),Correct=No
  else
    LoadPDB (ensemblefile)(start),Correct=No
  # Set the BFactors for floating assignments in case this macro was run directly, not from nmr_solve
  # Don't affect dummy atoms like the RDC coordinate system
  BFactorAtom Element !Dummy,0
  BFactorAtom (floating),25
  # Correction of naming deviations (Val/Leu) could be useful, but Clean may do too much
  #Clean
  LogAs (logfile),append=Yes, print '  Ensemble structure (i):'
  # Choose a nice orientation
  # (this aligns the major protein axes with the global coordinate system
  #  and makes sure the simulation cell gets as small as possible).
  NiceOriObj 1
  # Select force field for explicit solvent
  ForceField (fofsolvent)
  ScaleForce Angle,2
  ScaleForce Planarity,5
  ScaleForce Dihedral1D,7
  ScaleForce Dihedral2D,7
  ScaleForce Dihedral3D,7
  Cutoff 8
  # Refinement at room temperature and somewhat above, depending on the structure counter
  # This ensures that larger ensembles exploit a larger temperature range
  Temp (298+i%100)K
  # Create the cell
  Cell Auto,Extension=5
  Boundary periodic
  Longrange Coulomb
  
  # Check if there is a maximum allowed violation
  if violdismax!=9999 or violdihmax!=9999 or violrdcmax!=9999
    # Iteratively pull stronger and stronger until requirements are met
    RestrainPot (defaultpot)
    Load(restrainformat) (restrainfile),1,NameFormat=(nameformat)
    EnforceRest violdismax,violdihmax,violrdcmax,defaultscale
    DelRest
  # Fill with solvent
  if solvent=='HOH'
    # Fill with water, predict pKas, place counter ions
    Experiment Neutralization
      WaterDensity 0.997
      pH (ph)
      Speed Fast
    Experiment On
    Wait ExpEnd
    # At this point, NMR structures are still sometimes loosely packed, and water
    # molecules may have ended up in the core. So we delete all waters that are
    # relatively lonely, because nothing is worse than water refinement with
    # water inside the protein.
    DelRes HOH with <7 atoms closer 5 among HOH Atom O
  else
    if solvent=='DMSO'
      # Build a DMSO molecule: First an S, add 3 Hs, and turn the Hs into O and CH3
      dmso = BuildAtom S
      h() = AddHydObj (dmso),3
      SwapBond (h1),all,2
      SwapAtom (h1),O,UpdateBonds=Yes
      SwapAtom (h2) (h3),C,UpdateBonds=Yes,UpdateHyd=Yes
      NameObj (dmso),Solvent
    else
      RaiseError 'The requested solvent "(solvent)" is currently not supported'
    # Fill the cell with the solvent
    FillCellObj Solvent,Density=1.1,RandomOri=100%
    FixObj !Solvent
    # Energy minimize the solvent
    Experiment Minimization
    Experiment On
    Wait ExpEnd
    FreeAll
  # Load the restraints
  RestrainPot (defaultpot)
  Load(restrainformat) (restrainfile),1,NameFormat=(nameformat)
  ScaleRest Distance=(defaultscale1),Dihedral=(defaultscale2),RDC=(defaultscale3)
  # Start the Simulation
  TimeStep 2,1.0
  TempCtrl SteepDes
  Sim On
  ShowMessage 'Starting explicit solvent refinement of structure (i) with a steepest descent minimization...'
  for j=1 to 1000
    Wait 1
    if SpeedMax<3000
      break
  # Correct cis-peptide bonds and wrong isomers also by remodeling
  CorrectCis On,Old=(correctoldcisiso),Proline=(correctcispro)
  CorrectIso On,Old=(correctoldcisiso)
  CorrectConv On
  bestenergy=1e100
  bestenergyshown='not yet determined'
  # Store current SqOffset in 'offset'
  x,x,offset = RestrainPot
  failures=0
  failuresmax=4
  while failures<failuresmax
    Style Ribbon
    # Set strong restraining forces
    RestrainPot SqOffset=(0.5+offset)
    ScaleRest Distance=(strongscale1),Dihedral=(strongscale2),RDC=(strongscale3)
    # Heat
    ShowMessage 'Simulated annealing refinement in solvent, cycle (failures)/(failuresmax). Minimum violation energy is (bestenergyshown). Heating...'
    Temp (298+i%100)K
    TempCtrl Rescale
    Wait (solventsimsteps)
    # And cool
    ShowMessage 'Simulated annealing refinement in solvent, cycle (failures)/(failuresmax). Minimum violation energy is (bestenergyshown). Cooling...'
    TempCtrl Anneal
    Wait 20
    # Back to default restraining forces
    RestrainPot (defaultpot)
    ScaleRest Distance=(defaultscale1),Dihedral=(defaultscale2),RDC=(defaultscale3)
    Wait 20
    # In case something went seriously wrong (explosions), make sure not to save a
    # PDB file with long bonds, which may be skipped by LoadPDB in the next step.
    ScaleRest Distance=0,Dihedral=0,RDC=0
    TempCtrl SteepDes
    for j=1 to 100
      Wait 1
      if SpeedMax<4000
        break
    ScaleRest Distance=(defaultscale1),Dihedral=(defaultscale2),RDC=(defaultscale3)
    # See if the restraint violations are smaller
    violenergy = RestEnergy
    LogAs (logfile),append=Yes, print '    Cycle (failures)/(failuresmax): Violation energy is (0.00+violenergy) (EnergyUnit)'
    if violenergy<bestenergy
      # Yes, replace the structure
      bestenergy=violenergy
      SavePDB 1,(ensemblefile)(i)
      LogAs (logfile),append=Yes, print '      Best energy so far, this is the new member of the final ensemble.'
      failures=0
    else
      failures=failures+1
    bestenergyshown=0.00+bestenergy
  # Finally, force restraints again
  if violdismax!=9999 or violdihmax!=9999 or violrdcmax!=9999
    Clear
    LoadPDB (ensemblefile)(i),Correct=No
    Cell Auto,Extension=5
    RestrainPot (defaultpot)
    Load(restrainformat) (restrainfile),1,NameFormat=(nameformat)
    EnforceRest violdismax,violdihmax,violrdcmax,defaultscale
    SavePDB 1,(ensemblefile)(i)
ShowMessage 'Explicit solvent refinement completed.'
Sim off


# ENFORCE RESTRAINTS UNTIL MAXIMUM ALLOWED VIOLATIONS ARE MET
# ===========================================================
def EnforceRest violdismax,violdihmax,violrdcmax,defaultscale()
  TempCtrl Anneal
  Timestep 2,0.25
  scale()=defaultscale
  Sim On
  do
    # Force as much as necessary
    ScaleRest Distance=(scale1),Dihedral=(scale2),RDC=(scale3)
    Wait 10
    scale()=scale+defaultscale
    dissum,dismax,dihsum,dihmax,rdcsum,rdcmax = RestViol
    ShowMessage 'Current Scale/MaxViolation: Distance=(0.00+scale1)/(0.00+dismax), Dihedral=(0.00+scale2)/(0.00+dihmax), RDC=(0.00+scale3)/(0.00+rdcmax)'
  while scale1<1000 and (dismax>violdismax or dihmax>violdihmax or rdcmax>violrdcmax)