File size: 5,380 Bytes
0605e17
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
from pymol import cmd
import requests


# from gradio_client import Client


def color_plddt(selection="all"):
    """
    AUTHOR
    Jinyuan Sun
    https://github.com/JinyuanSun/PymolFold/tree/main
    MIT License

    DESCRIPTION
    Colors Predicted Structures by pLDDT

    USAGE
    color_plddt sele

    PARAMETERS

    sele (string)
    The name of the selection/object to color by pLDDT. Default: all
    """
    # Alphafold color scheme for plddt
    cmd.set_color("high_lddt_c", [0, 0.325490196078431, 0.843137254901961])
    cmd.set_color(
        "normal_lddt_c", [0.341176470588235, 0.792156862745098, 0.976470588235294]
    )
    cmd.set_color("medium_lddt_c", [1, 0.858823529411765, 0.070588235294118])
    cmd.set_color("low_lddt_c", [1, 0.494117647058824, 0.270588235294118])

    # test the scale of predicted_lddt (0~1 or 0~100 ) as b-factors
    cmd.select("test_b_scale", f"b>1 and ({selection})")
    b_scale = cmd.count_atoms("test_b_scale")

    if b_scale > 0:
        cmd.select("high_lddt", f"({selection}) and (b >90 or b =90)")
        cmd.select("normal_lddt", f"({selection}) and ((b <90 and b >70) or (b =70))")
        cmd.select("medium_lddt", f"({selection}) and ((b <70 and b >50) or (b=50))")
        cmd.select("low_lddt", f"({selection}) and ((b <50 and b >0 ) or (b=0))")
    else:
        cmd.select("high_lddt", f"({selection}) and (b >.90 or b =.90)")
        cmd.select(
            "normal_lddt", f"({selection}) and ((b <.90 and b >.70) or (b =.70))"
        )
        cmd.select("medium_lddt", f"({selection}) and ((b <.70 and b >.50) or (b=.50))")
        cmd.select("low_lddt", f"({selection}) and ((b <.50 and b >0 ) or (b=0))")

    cmd.delete("test_b_scale")

    # set color based on plddt values
    cmd.color("high_lddt_c", "high_lddt")
    cmd.color("normal_lddt_c", "normal_lddt")
    cmd.color("medium_lddt_c", "medium_lddt")
    cmd.color("low_lddt_c", "low_lddt")

    # set background color
    cmd.bg_color("white")


def query_rosettafold2(
    sequence: str,
    jobname: str,
    sym: str = "X",
    order: int = 1,
    msa_concat_mode: str = "diag",
    msa_method: str = "single_sequence",
    pair_mode: str = "unpaired_paired",
    collapse_identical: bool = True,
    num_recycles: int = 0,
    use_mlm: bool = True,
    use_dropout: bool = True,
    max_msa: int = 16,
    random_seed: int = 0,
    num_models: int = 0,
):
    """
    AUTHOR
    Simon Duerr
    https://twitter.com/simonduerr


    DESCRIPTION
    Predict a structure using rosettafold2

    USAGE
    rosettafold2 sequence, jobname, [sym, order, msa_concat_mode, msa_method, pair_mode, collapse_identical, num_recycles, use_mlm, use_dropout, max_msa, random_seed, num_models]

    PARAMETERS

    sequence: (string)
    one letter amino acid codes that you want to predict

    jobname: string
    name of the pdbfile that will be outputted

    sym: string
    symmetry Default: X

    order:
    Default 1,

    msa_concat_mode:
    MSA concatenation mode Default:"diag" Options: "diag", "repeat", "default"

    msa_method:
    MSA method Default:"single_sequence" Options: "mmseqs2", "single_sequence"

    pair_mode:
    Pair mode Default:"unpaired_paired" Options: "unpaired_paired", "paired", "unpaired"

    collapse_identical:
    Collapse identical sequences Default:True

    num_recycles:
    Number of recycles Default:0 Options: 0, 1, 3, 6, 12, 24

    use_mlm:
    Use MLM Default:True

    use_dropout:
    Use dropout Default:True

    max_msa:
    Max MSA Default:16

    random_seed:
    Random seed Default:0

    num_models:
    Number of models Default:0
    """
    response = requests.post(
        "http://localhost:7860/run/rosettafold2/",
        json={
            "data": [
                sequence,  # str  in 'sequence' Textbox component
                jobname,  # str  in 'jobname' Textbox component
                sym,  # str  in 'sym' Textbox component
                order,  # int | float (numeric value between 1 and 12) in 'order' Slider component
                "diag",  # str (Option from: ['diag', 'repeat', 'default']) in 'msa_concat_mode' Dropdown component
                "single_sequence",  # str (Option from: ['mmseqs2', 'single_sequence', 'custom_a3m']) in 'msa_method' Dropdown component
                "unpaired_paired",  # str (Option from: ['unpaired_paired', 'paired', 'unpaired']) in 'pair_mode' Dropdown component
                True,  # bool  in 'collapse_identical' Checkbox component
                0,  # int (Option from: ['0', '1', '3', '6', '12', '24']) in 'num_recycles' Dropdown component
                True,  # bool  in 'use_mlm' Checkbox component
                True,  # bool  in 'use_dropout' Checkbox component
                16,  # int (Option from: ['16', '32', '64', '128', '256', '512']) in 'max_msa' Dropdown component
                0,  # int  in 'random_seed' Textbox component
                1,  # int (Option from: ['1', '2', '4', '8', '16', '32']) in 'num_models' Dropdown component
            ]
        },
    ).json()
    print(response)
    try:
        data = response["data"]
    except KeyError:
        print(response["error"])
        return None
    with open(f"{jobname}.pdb", "w") as out:
        out.writelines(data)
    cmd.load(f"{jobname}.pdb")


cmd.extend("rosettafold2", query_rosettafold2)
cmd.extend("color_plddt", color_plddt)