import barnaba as bb
import pickle

top = "topology.pdb"
traj = "trajectory.dcd" 
native = "2KOC.pdb" 

"""
Explanation: Calculate quantities for analysis
These notebooks describe how to calculate the data and how to produce figures in the manuscript 
"Barnaba: Software for Analysis of Nucleic Acids Structures and Trajectories". 
Here, we calculate different quantities over the entire trajectory:
- eRMSD from reference structure
- RMSD from reference structure
- Base-pairing and base-stacking detection (annotation)
- Dot-bracket annotation
- Backbone torsion angles
- $^3$J scalar couplings
- Relative position and orientation between nucleobases as G-vectors, as defined in Bottaro, Di Palma, Bussi. NAR 2014.
All the data is saved to pickle files for later analysis. The MD trajectory is taken from the paper "RNA force field with accuracy comparable to state-of- the-art protein force fields", PNAS, 2017.
First, we import the modules barnaba and pickles, and define the location of the topology/trajectory file, as well as the location of the native, reference structure.
End of explanation
"""


# calculate ermsd and store in a pickle file 
fname = "ermsd.p"
ermsd = bb.ermsd(native,traj,topology=top)
pickle.dump(ermsd[1:],open(fname, "w"))

# calculate rmsd and store in a pickle file 
fname = "rmsd.p"
print "# calculate %s" % fname
rmsd = bb.rmsd(native,traj,topology=top,heavy_atom=True)
pickle.dump(rmsd[1:],open(fname, "w"))

# calculate annotation and store in pickle file
fname = "pairs.p"
print "# calculate %s" % fname
stackings, pairings, res = bb.annotate(traj,topology=top)
pickle.dump([pairings[1:], res],open(fname, "w"))

# calculate dot-bracket annotation and store in pickle file
fname = "dotbracket.p"
dotbr,ss = bb.dot_bracket(pairings,res)
pickle.dump([dotbr[1:], res],open(fname, "w"))

# Calculate torsion angles
fname = "angles.p"
print "# calculate %s" % fname
angles,res = bb.backbone_angles(traj,topology=top)
pickle.dump([angles[1:],res],open(fname, "w"))

# calculate couplings and save to pickle
fname = "couplings.p"
print "# calculate %s" % fname
couplings,res = bb.jcouplings(traj,topology=top)
pickle.dump([couplings[1:],res],open(fname, "w"))

# calculate couplings and save to pickle
fname = "gvec.p"
print "# calculate %s" % fname
gvec,seq = bb.dump_gvec(traj,topology=top)
pickle.dump([gvec[1:],seq],open(fname, "w"))




"""
Explanation: Now we calculate all the quantitites described above and save the data to pickle file. Note that heavy_atoms=True in RMSD calculation. Default mode is backbone-only.
End of explanation
"""


