run_pretrained_openfold.py 9.39 KB
Newer Older
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# Copyright 2021 AlQuraishi Laboratory
# Copyright 2021 DeepMind Technologies Limited
# 
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

16
import argparse
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
17
from datetime import date
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
18
import gc
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
19
import logging
20
import numpy as np
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
21
import os
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
22

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
23
import pickle
24
25
import random
import sys
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
26
27
28
import time
import torch

29
from openfold.config import model_config
30
from openfold.data import templates, feature_pipeline, data_pipeline
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
31
from openfold.model.model import AlphaFold
32
from openfold.model.torchscript import script_preset_
33
from openfold.np import residue_constants, protein
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
34
35
import openfold.np.relax.relax as relax
from openfold.utils.import_weights import (
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
36
37
    import_jax_weights_,
)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
38
from openfold.utils.tensor_utils import (
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
39
40
41
    tensor_tree_map,
)

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
42
from scripts.utils import add_data_args
43

44

45
46
def main(args):
    config = model_config(args.model_name)
47
    model = AlphaFold(config)
48
    model = model.eval()
49
    import_jax_weights_(model, args.param_path, version=args.model_name)
50
    #script_preset_(model)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
51
    model = model.to(args.model_device)
52
 
53
54
55
    template_featurizer = templates.TemplateHitFeaturizer(
        mmcif_dir=args.template_mmcif_dir,
        max_template_date=args.max_template_date,
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
56
        max_hits=config.data.predict.max_templates,
57
        kalign_binary_path=args.kalign_binary_path,
58
        release_dates_path=args.release_dates_path,
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
59
60
        obsolete_pdbs_path=args.obsolete_pdbs_path
    )
61

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
62
63
    use_small_bfd=(args.bfd_database_path is None)

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
64
    data_processor = data_pipeline.DataPipeline(
65
66
67
68
        template_featurizer=template_featurizer,
    )

    output_dir_base = args.output_dir
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
69
    random_seed = args.data_random_seed
70
71
    if random_seed is None:
        random_seed = random.randrange(sys.maxsize)
72
    feature_processor = feature_pipeline.FeaturePipeline(config.data)
73
74
    if not os.path.exists(output_dir_base):
        os.makedirs(output_dir_base)
Gustaf's avatar
Gustaf committed
75
    if(args.use_precomputed_alignments is None):
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
76
        alignment_dir = os.path.join(output_dir_base, "alignments")
Gustaf's avatar
Gustaf committed
77
78
    else:
        alignment_dir = args.use_precomputed_alignments
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
79

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
80
81
82
83
    for fasta_file in os.listdir(args.fasta_dir):
        # Gather input sequences
        with open(os.path.join(args.fasta_dir, fasta_file), "r") as fp:
            data = fp.read()
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
84

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
85
86
87
88
89
        lines = [
            l.replace('\n', '') 
            for prot in data.split('>') for l in prot.strip().split('\n', 1)
        ][1:]
        tags, seqs = lines[::2], lines[1::2]
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
90

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
91
92
93
94
        assert len(seqs) == 1, "Input FASTAs may only contain one sequence"
        tag, seq = tags[0], seqs[0]

        fasta_path = os.path.join(args.output_dir, f"tmp_{os.getpid()}.fasta")
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
        with open(fasta_path, "w") as fp:
            fp.write(f">{tag}\n{seq}")

        logging.info("Generating features...") 
        local_alignment_dir = os.path.join(alignment_dir, tag)
        if(args.use_precomputed_alignments is None):
            if not os.path.exists(local_alignment_dir):
                os.makedirs(local_alignment_dir)
            
            alignment_runner = data_pipeline.AlignmentRunner(
                jackhmmer_binary_path=args.jackhmmer_binary_path,
                hhblits_binary_path=args.hhblits_binary_path,
                hhsearch_binary_path=args.hhsearch_binary_path,
                uniref90_database_path=args.uniref90_database_path,
                mgnify_database_path=args.mgnify_database_path,
                bfd_database_path=args.bfd_database_path,
                uniclust30_database_path=args.uniclust30_database_path,
                pdb70_database_path=args.pdb70_database_path,
                use_small_bfd=use_small_bfd,
                no_cpus=args.cpus,
            )
            alignment_runner.run(
                fasta_path, local_alignment_dir
            )
119
    
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
120
121
122
        feature_dict = data_processor.process_fasta(
            fasta_path=fasta_path, alignment_dir=local_alignment_dir
        )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
123

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
124
125
        # Remove temporary FASTA file
        os.remove(fasta_path)
126
    
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
127
128
129
        processed_feature_dict = feature_processor.process_features(
            feature_dict, mode='predict',
        )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
130

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
131
132
133
134
135
136
137
138
        logging.info("Executing model...")
        batch = processed_feature_dict
        with torch.no_grad():
            batch = {
                k:torch.as_tensor(v, device=args.model_device) 
                for k,v in batch.items()
            }
        
139
            t = time.perf_counter()
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
140
            out = model(batch)
141
            logging.info(f"Inference time: {time.perf_counter() - t}")
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
       
        # Toss out the recycling dimensions --- we don't need them anymore
        batch = tensor_tree_map(lambda x: np.array(x[..., -1].cpu()), batch)
        out = tensor_tree_map(lambda x: np.array(x.cpu()), out)
        
        plddt = out["plddt"]
        mean_plddt = np.mean(plddt)
        
        plddt_b_factors = np.repeat(
            plddt[..., None], residue_constants.atom_type_num, axis=-1
        )
    
        unrelaxed_protein = protein.from_prediction(
            features=batch,
            result=out,
            b_factors=plddt_b_factors
        )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
159
160
161
162
163
164
165
166

        # Save the unrelaxed PDB.
        unrelaxed_output_path = os.path.join(
            args.output_dir, f'{tag}_{args.model_name}_unrelaxed.pdb'
        )
        with open(unrelaxed_output_path, 'w') as f:
            f.write(protein.to_pdb(unrelaxed_protein))

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
        if(not args.skip_relaxation):
            amber_relaxer = relax.AmberRelaxation(
                use_gpu=(args.model_device != "cpu"),
                **config.relax,
            )
            
            # Relax the prediction.
            t = time.perf_counter()
            visible_devices = os.getenv("CUDA_VISIBLE_DEVICES", default="")
            if("cuda" in args.model_device):
                device_no = args.model_device.split(":")[-1]
                os.environ["CUDA_VISIBLE_DEVICES"] = device_no
            relaxed_pdb_str, _, _ = amber_relaxer.process(prot=unrelaxed_protein)
            os.environ["CUDA_VISIBLE_DEVICES"] = visible_devices
            logging.info(f"Relaxation time: {time.perf_counter() - t}")
            
            # Save the relaxed PDB.
            relaxed_output_path = os.path.join(
                args.output_dir, f'{tag}_{args.model_name}_relaxed.pdb'
            )
            with open(relaxed_output_path, 'w') as f:
                f.write(relaxed_pdb_str)
189

190
191
192
193
194
195
196
        if(args.save_outputs):
            output_dict_path = os.path.join(
                args.output_dir, f'{tag}_{args.model_name}_output_dict.pkl'
            )
            with open(output_dict_path, "wb") as fp:
                pickle.dump(out, fp, protocol=pickle.HIGHEST_PROTOCOL)

197
198
199

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
200
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
201
202
        "fasta_dir", type=str,
        help="Path to directory containing FASTA files, one sequence per file"
203
    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
204
205
206
    parser.add_argument(
        "template_mmcif_dir", type=str,
    )
Gustaf's avatar
Gustaf committed
207
208
209
210
    parser.add_argument(
        "--use_precomputed_alignments", type=str, default=None,
        help="""Path to alignment directory. If provided, alignment computation 
                is skipped and database path arguments are ignored."""
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
211
    )
212
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
213
214
        "--output_dir", type=str, default=os.getcwd(),
        help="""Name of the directory in which to output the prediction""",
215
216
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
217
        "--model_device", type=str, default="cpu",
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
218
219
        help="""Name of the device on which to run the model. Any valid torch
             device name is accepted (e.g. "cpu", "cuda:0")"""
220
221
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
222
223
224
        "--model_name", type=str, default="model_1",
        help="""Name of a model config. Choose one of model_{1-5} or 
             model_{1-5}_ptm, as defined on the AlphaFold GitHub."""
225
226
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
227
228
229
230
        "--param_path", type=str, default=None,
        help="""Path to model parameters. If None, parameters are selected
             automatically according to the model name from 
             openfold/resources/params"""
231
    )
232
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
233
        "--save_outputs", action="store_true", default=False,
234
235
        help="Whether to save all model outputs, including embeddings, etc."
    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
236
237
    parser.add_argument(
        "--cpus", type=int, default=4,
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
238
        help="""Number of CPUs with which to run alignment tools"""
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
239
    )
240
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
241
        "--preset", type=str, default='full_dbs',
242
243
244
        choices=('reduced_dbs', 'full_dbs')
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
245
246
247
248
        "--data_random_seed", type=str, default=None
    )
    parser.add_argument(
        "--skip_relaxation", action="store_true", default=False,
249
    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
250
    add_data_args(parser)
251
252
253
254
255
256
257
258
    args = parser.parse_args()

    if(args.param_path is None):
        args.param_path = os.path.join(
            "openfold", "resources", "params", 
            "params_" + args.model_name + ".npz"
        )

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
259
260
261
262
263
264
    if(args.model_device == "cpu" and torch.cuda.is_available()):
        logging.warning(
            """The model is being run on CPU. Consider specifying 
            --model_device for better performance"""
        )

265
    main(args)