RosettaConstraintGenerate/constraint_from_structure.py at main · cemilcansaylan/RosettaConstraintGenerate · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
#!/usr/bin/env python
# coding: utf-8

"""
Created on Wed Sep 15 22:27:56 2021
@author: cemil can saylan
"""

from Bio.PDB import *
from Bio import SeqIO
import numpy

from absl import app
from absl import flags

flags.DEFINE_string('structure', None, 'PDB Structure')
flags.DEFINE_string('output', 'constraint.cst', 'constraint file name')
flags.DEFINE_boolean('printfasta', False, 'Print sequence of structure')
flags.DEFINE_integer('distance', 20, 'Atompair radius')

FLAGS = flags.FLAGS
flags.mark_flag_as_required("structure")
flags.mark_flag_as_required("distance")


def PrintFasta(structure):
    f = open("input.fasta", "w")
    for record in SeqIO.parse(structure, "pdb-atom"):
        f.write(str(record.seq))
    f.close()

def AtomPairNearest(structure, radius):
    """
    Only for CA
    """

    res_list = Selection.unfold_entities(structure, "R")
    atoms  = Selection.unfold_entities(structure, 'A')

    ns = NeighborSearch(atoms)
    Nearest_Atom = []

    for res in res_list:
        near = ns.search(res["CA"].coord, radius)
        for a in near:
            if a.id == "CA":
                diff = res["CA"].coord - a.coord
                dist = numpy.sqrt(numpy.sum(diff * diff))
                if dist == 0.0:
                    continue
                Nearest_Atom.append([res["CA"].id,
                                     res["CA"].get_parent().id[1],
                                     a.id,
                                     a.get_parent().id[1],
                                     dist])
    return Nearest_Atom

def Dih_Angle(structure):
    """
    Get the phi, psi information from the structure
    """
    dihedrals = []
    angles = []
    structure.atom_to_internal_coordinates()
    for r in structure.get_residues():
        if r.internal_coord:
            angles.append(r.internal_coord.get_angle("N:CA:C"))
            dihedrals.append([r.internal_coord.get_angle("phi"),
                             r.internal_coord.get_angle("psi")])
    return dihedrals, angles

def PrintConstraint(pdb, filename, radius):
    """
    psi N-Cα-C-N
    phi C-N-Cα-C
    """
    res_list = Selection.unfold_entities(pdb, "R")

    Nearest_Atom = AtomPairNearest(pdb, radius)
    Dihedrals, Angles = Dih_Angle(pdb)

    f = open(filename, "w")
    for res_num in range(1,len(res_list)+1):
        if res_num == 1:
            f.write(f"Dihedral N {res_num} CA {res_num} C {res_num} N {res_num+1} LINEAR_PENALTY {Dihedrals[res_num-1][1]*0.01745329252} 0 0 10.0\n")
            f.write(f"Angle N {res_num} CA {res_num} C {res_num} LINEAR_PENALTY {Angles[res_num-1]} 0 0 10.0\n")
        elif res_num == len(res_list):
            f.write(f"Dihedral C {res_num-1} N {res_num} CA {res_num} C {res_num} LINEAR_PENALTY {Dihedrals[res_num-1][0]*0.01745329252} 0 0 10.0\n")
            f.write(f"Angle N {res_num} CA {res_num} C {res_num} LINEAR_PENALTY {Angles[res_num-1]} 0 0 10.0\n")
        else:
            f.write(f"Dihedral N {res_num} CA {res_num} C {res_num} N {res_num+1} LINEAR_PENALTY {Dihedrals[res_num-1][1]*0.01745329252} 0 0 10.0\n")
            f.write(f"Dihedral C {res_num-1} N {res_num} CA {res_num} C {res_num} LINEAR_PENALTY {Dihedrals[res_num-1][0]*0.01745329252} 0 0 10.0\n")
            f.write(f"Angle N {res_num} CA {res_num} C {res_num} LINEAR_PENALTY {Angles[res_num-1]} 0 0 10.0\n")


    for atompair in Nearest_Atom:
        f.write(f"AtomPair CA {atompair[1]} CA {atompair[3]} LINEAR_PENALTY {atompair[4]} 0 0 10.0\n")

    f.close()

def main(argv):
    parser = PDBParser()
    pdb_file = parser.get_structure("demo", FLAGS.structure)

    PrintConstraint(pdb_file, FLAGS.output, FLAGS.distance)

    if FLAGS.printfasta:
        PrintFasta(FLAGS.structure)

if __name__ == '__main__':
    app.run(main)