File contents

boundHeavyAtomDict = {}

def structurePredictNoePeakList(argServer):

  from ccpnmr.analysis.StructureBasic  import getAtomSetsDistance
  from ccpnmr.analysis.AssignmentBasic import getAtomSetShifts, assignResToDim
  from ccpnmr.analysis.ExperimentBasic import getSpectrumIsotopes, getNoesyPeakLists
  from ccpnmr.analysis.PeakBasic       import pickPeak
  from ccpnmr.analysis.Util            import setPeakListColor, getPeakListColor, setPeakListSymbol

  boundHeavyAtomDict = {} 
  structure = argServer.getStructure()
  
  if not structure:
    argServer.showWarning('No structures in project to work with')
    return

  project   = structure.root
  threshold = argServer.askFloat('Threshold?',5.5)

  argServer.showWarning('Select spectrum to contain mock peak list.')
  
  specDict  = {}
  peakLists = getNoesyPeakLists(structure.root)
  for peakList in peakLists:
    specDict[peakList.dataSource] = None
  
  spectrum = argServer.getSpectrum(specDict.keys())
  if not spectrum:
    argServer.showWarning('No spectrum selected to hold peak list')
    return
    
  isotopes = getSpectrumIsotopes(spectrum)
  if isotopes.count('1H') < 2:
    argServer.showWarning('Spectrum nust have at least two 1H dimensions')
    return
  
  indirectDim = -1
  N = len(isotopes)
  if N > 3:
    argServer.showWarning('Sorry. Macro only works for 2D or 3D NOESYs')
    return
    
  elif N > 2:
    j = 0
    for i in range(N):
      if isotopes[i] == '1H':
        j =  i
    
    while (indirectDim < 0) or  (indirectDim >= N):
      argServer.parent.update_idletasks()
      indirectDim = argServer.askInteger('Indiect 1H dimension?', j+1) -1
  
  directDim    = 0
  heavyDim     = N-1
  heavyIsotope = '15N'
  for i in range(N):
    if isotopes[i] == '1H' and i != indirectDim:
      directDim = i
    elif isotopes[i] != '1H':
      heavyDim = i
      heavyIsotope = isotopes[i]
    
  shiftLists = project.findAllNmrMeasurementLists(className='ShiftList')
  if not shiftLists:
    argServer.showWarning('No shift lists in project')
    return
    
  if len(shiftLists) > 1:
    argServer.showWarning('Now select shift list to predict peak positions.')
    shiftList = argServer.getShiftList()
  else:
    shiftList = shiftLists[0]

  color = 'red'
  if spectrum.peakLists and (getPeakListColor(spectrum.peakLists[0]) == 'red'):
    color = 'blue'
      
  peakList = spectrum.newPeakList(isSimulated=1)
  setPeakListColor(peakList,color)
  setPeakListSymbol(peakList, '+')

  print "Finding hydrogen atoms with coordinates"
  atomSetsDict = {}
  for coordChain in structure.coordChains:
    for coordResidue in coordChain.residues:
      for coordAtom in coordResidue.atoms:
        if coordAtom.name[0] == 'H':
          atom = coordAtom.atom
          if atom and atom.atomSet:
            atomSetsDict[atom.atomSet] = 1
  
  
  print "Getting close atomic pairs"
  atomSetPairs = []        
  atomSets = atomSetsDict.keys()
  for i in range(len(atomSets)-1):
    if i and i % 10 == 0:
      print "  Done %d of %d" % (i, len(atomSets))
    
    for j in range(i+1, len(atomSets)):
      dist = getAtomSetsDistance([atomSets[i],],[atomSets[j],],structure)
      if dist <= threshold:
         atomSetPairs.append( (atomSets[i], atomSets[j], dist) )

  print "Making mock peaks"
  i = 0
  for atomSet1, atomSet2, dist in atomSetPairs:
    if i and i % 100 == 0:
      print "  Done %d of %d" % (i, len(atomSetPairs))
    
    shifts1 = getAtomSetShifts(atomSet1)
    shifts2 = getAtomSetShifts(atomSet2)
    for shift1 in shifts1:
      for shift2 in shifts2:
        if N == 2:
          position = [shift1.value, shift2.value]
          peak = pickPeak(peakList, position, unit=shiftList.unit)
          assignResToDim(peak.peakDims[0],shift1.resonance)
          assignResToDim(peak.peakDims[1],shift2.resonance)
          if shift1 is not shift2:
            position.reverse()
            peak = pickPeak(peakList, position, unit=shiftList.unit)
            assignResToDim(peak.peakDims[0],shift2.resonance)
            assignResToDim(peak.peakDims[1],shift1.resonance)
        else:
          shiftX = getBoundHeavyAtomShift(atomSet1, heavyIsotope, shiftList)
          if shiftX:
            position = [None] * N
            position[directDim]   = shift1.value
            position[indirectDim] = shift2.value
            position[heavyDim]    = shiftX.value
            peak = pickPeak(peakList, position, unit=shiftList.unit)
            assignResToDim(peak.peakDims[directDim],shift1.resonance)
            assignResToDim(peak.peakDims[indirectDim],shift2.resonance)
            assignResToDim(peak.peakDims[heavyDim],shiftX.resonance)
            peak.setAnnotation('%.2f - ' % dist)


          shiftX = getBoundHeavyAtomShift(atomSet2, heavyIsotope, shiftList)
          if shiftX:
            position = [None] * N
            position[directDim]   = shift2.value
            position[indirectDim] = shift1.value
            position[heavyDim]    = shiftX.value
            peak = pickPeak(peakList, position, unit=shiftList.unit)
            assignResToDim(peak.peakDims[directDim],shift2.resonance)
            assignResToDim(peak.peakDims[indirectDim],shift1.resonance)
            assignResToDim(peak.peakDims[heavyDim],shiftX.resonance)
            peak.setAnnotation('%.2f - ' % dist)
    i += 1
    
  print "Done"
          
def getBoundHeavyAtomShift(atomSet, isotope, shiftList):

  if boundHeavyAtomDict.get(atomSet):
    return boundHeavyAtomDict[atomSet]

  shift   = None
  atomH   = atomSet.atoms[0]
  residue = atomH.residue
  
  resonance = None
  
  for chemBond in atomH.chemAtom.chemBonds:
    for chemAtom in chemBond.chemAtoms:
      if chemAtom is not atomH.chemAtom:
        if chemAtom.elementSymbol and chemAtom.elementSymbol in isotope:
          for atom in residue.atoms:
            if atom.chemAtom is chemAtom:
              if atom.atomSet and atom.atomSet.resonanceSets:
                resonance = atom.atomSet.resonanceSets[0].resonances[0]
              break

  if resonance:
    shift = resonance.findFirstShift(parentList = shiftList)

  boundHeavyAtomDict[isotope] = shift
  return shift