Chemical Shift Differences

#A macro to perform 1H 15N chemical shift perturbation analysis comparing peaks in different spectra eg. a wild type and a mutant
#This macro compares the differences between the ppmPositions of identically assigned peaks, based on SeqCode and 
#plots them as a barchart as function of residue number.
#It further generates a script file which can be run in PyMOL, showing the residues affected coloured from least to most affected, 
#and unaffected or unknown in white. 
# 
# Uses "data2bfactor.py" which can be found at "https://gist.github.com/Croydon-Brixton/ca916769c285faed69740d384b5daf23
# note: Alternative is to read pdb file and change the b-factors (e.g. using BioPython, already embedded in CCPN) and use modified pdb file loading into PyMol
# To run, change the following parameters:
#     Chemical shift pertubation function: (((H_1 - H_2)^2+(WF*(N_1 - N_2))^2))^(1/2) WF=weighting factor
#     origPeakList (peaklist containing peaks from one spectrum),
#     newPeakList (containing peaks in the other spectrum)
#     xlims (residue range to plot)
#     Figure_Filepath (filepath to put the generated figure),
#     delta_filepath (filepath for generated deltas),
#     pdbPath (filepath for pdb file containing the structure to map onto),
#     filePath (for PyMol file)
#     d2bPath (the location of your data2bfactor.py)
#     objName (the name of your object in PyMOL)
#Copyright 2024 Pernille Vosbein and Brian Smith
###############################################################################################
import pandas as pd
import matplotlib.pyplot as plt
import csv
import os
from ccpn.ui.gui.modules.PyMolUtil import CodeBlock
import numpy as np

###Change appropriately
WF=0.2 #Weighting factor for Nitrogen
origPeakList=get('PL:Ubi_free_0.2')
newPeakList=get('PL:Ubi_Uba_500.1')
xlims=[0,77] #one before to one after the residue range to plot
Figure_Filepath='/Users/vad5/Desktop/ChemShiftDiffs.png'
delta_filepath= '/Users/vad5/Desktop/Deltastab'
csv_filename ='DeltaDeltas_Ubi.csv'
#For plotting in PyMOL. Input structure to map onto as string, and string for the filepath to the file it outputs. 
d2bPath = '/Users/vad5/Documents/V3Projects/Macros/data2bfactor.py' # location of the data2bfactor script
pdbPath = '/Users/vad5/Documents/Tutorials/NewTitrationTutorial/CcpnUbiTitrationTutorial/1ubq.pdb' # the file containing the structure to map onto
filePath = '/Users/vad5/Desktop/Ubi_mapped.pml' # the filepath for the PyMOL file storing the mapped residues
objName = '1ubq' # name of the object in PyMOL (usually the PDB code)

###Script starts here
def shiftsDfFromPeakList(pl):
   '''harvest H and N shifts from assigned peaks in a peakList and return a dataFrame'''
   
   #should check input is valid
   dataArray= []
   
   Hidx, Nidx = pl.spectrum.isotopeCodes.index('1H'), pl.spectrum.isotopeCodes.index('15N')
   #assuming that peaks are assigned to single NmrResidue in both dimensions
   for peak in pl.peaks[:]:
     if peak.assignments:
       if peak.assignments[0][Nidx].name == 'N':
         row_dict={}
         seqCode=peak.assignments[0][0].nmrResidue.sequenceCode
         nmrChain=peak.assignments[0][0].nmrResidue.nmrChain.pid
         peakPid=peak.pid
         nmrAtomH=peak.assignments[0][Hidx]
         nmrAtomN=peak.assignments[0][Nidx]
         shiftH=peak.ppmPositions[Hidx]
         shiftN=peak.ppmPositions[Nidx]
         row_dict.update(seqCode=seqCode, nmrChain=nmrChain, peakPid=peakPid, nmrAtomH=nmrAtomH, nmrAtomN=nmrAtomN, shiftH=shiftH, shiftN=shiftN)
         dataArray.append(row_dict)
         
   dataFrame=pd.DataFrame(dataArray)
    
   return dataFrame

myDataFrameA = shiftsDfFromPeakList(origPeakList)
myDataFrameB = shiftsDfFromPeakList(newPeakList)
# Creates a Dataframe with layout:
#seqCode  nmrChain  peakPid  nmrAtomH  nmrAtomN  shiftH  shiftN  
# merging the DataFrames on SeqCode and giving each column suffix either o (orig) or n (new)
mergedFrame=myDataFrameA.merge(myDataFrameB, on=['seqCode'], suffixes=('_o', '_n')) 
print(mergedFrame)
 
#calculate the ShiftDistance and add to new column
mergedFrame['ShiftDist']=((mergedFrame['shiftH_n']-mergedFrame['shiftH_o'])**2+(WF*(mergedFrame['shiftN_n']-mergedFrame['shiftN_o']))**2)**(1/2)
dataToPlot = mergedFrame[['seqCode','ShiftDist']]
#dropping incompletely assigned '@' NmrResidues
dataToPlot = dataToPlot.drop(dataToPlot[(dataToPlot['seqCode'].str.contains('@') == True)].index)
dataToPlot=dataToPlot.astype('float')
#create a new df with the data to plot and sort them by the seqCode column
dataToPlot = dataToPlot.sort_values(by=['seqCode']) 
#make a complete series so that missing residues can be plotted and coloured
# uses lowest to highest assigned seqCode
# use xlims to create filled if you need a wider range of residue numbers
filled = pd.Series(np.arange(dataToPlot.iloc[0,0], dataToPlot.iloc[-1, 0]+1))
filled = filled.to_frame()
filled.columns = ['seqCode']
merged = pd.merge(filled, dataToPlot, on='seqCode', how ='left')

#Save dataframe as CSV file
merged.to_csv(csv_filename, sep=',', index=True, na_rep='0')
merged['no_residue'] = [-0.005 if x == True else 1 for x in merged['ShiftDist'].isna()]
noRes = merged[merged['no_residue'] == -0.005]
#plt.plot(dataToPlot['seqCode'],dataToPlot['ShiftDist']) # plot the data
#fig, plt = plt.subplots(figsize=(12,12))
fig = plt.figure(figsize=(10,6))
ax = plt.subplot()
ax.bar(merged['seqCode'],merged['ShiftDist']) # plot the data as barchart
ax.scatter(noRes['seqCode'],noRes['no_residue'], marker='v', color='black', s=10, linewidths=0.7, label ='Residue not assigned') # plot the data as barchart
#ax.axhline(y=0.01, linestyle='--', color='grey') # setting a horizontal line  
#ax.xticks(rotation = 25, fontsize=12) #set rotation of ticks on the x-axis
#ax.yticks(fontsize=12) #set rotation of ticks on the x-axis
ax.set_xlabel('Residue number', fontsize=14)
ax.set_ylabel('Chemical shift perturbation', fontsize=14)
ax.set_xlim(xlims) # setting the x-axis limitation on the plot
ax.legend()
plt.savefig(Figure_Filepath)   # save as a figure in the specified path
plt.show()

#This part is for mapping onto the PyMOL structure part is for mapping onto the PyMOL structure
Deltas = dataToPlot
Deltas['seqCode']=Deltas['seqCode'].astype(int) 

max_bfactor= Deltas['ShiftDist'].max()
min_bfactor=Deltas['ShiftDist'].min()
bfactor_range='[' + str(min_bfactor) + ',' + str(max_bfactor) + ']'
def MapDeltaToPyMol(filePath, pdbPath, Deltas):
    # adjust the following to reflect the location of your data2bfactor.py and the name of your object in PyMOL
    #writing the Deltas DataFrame to a csv file.
    Deltas.to_csv(delta_filepath, index=False, sep='\t', header=False)
    #writing the PyMOL script
    if not os.path.exists(pdbPath):
       print("Warning: PDB file not found at " + pdbPath + ". You will need to edit the PyMOL script.")
    codeBlock = CodeBlock()
    imp = codeBlock.addImport
    cmd = codeBlock.addCmd
    imp('from pymol import cmd')
    cmd('cmd.reinitialize', )
    cmd('cmd.run', d2bPath)
    cmd('cmd.load', pdbPath)
    cmd('cmd.hide', 'lines')
    cmd('cmd.show', 'cartoon')
    cmd('cmd.color', 'white')
    cmd('cmd.alter', objName, 'b=99.9')      #set all B-factors to 99.9 in PyMOL
    cmd('data2b_res', objName, delta_filepath)       #use the data2b_res function
    cmd('cmd.spectrum', "b", "yellow_red", objName , min_bfactor, max_bfactor) #define the plotting
    cmd('cmd.ramp_new', "count",  objName, bfactor_range, "[yellow, red ]") # optionalm, creates a ramp for the plot
    cmd('cmd.color',  'white', '(b=99.9)') # set the colour of all the residues with b-factor 99.9 to white
    cmd('cmd.recolor')
    codeBlock.toFile(filePath) # makes a PyMoL file that can then be loaded in PyMOL
    return filePath
			    
MapDeltaToPyMol(filePath,pdbPath,Deltas)