run_pretrained_openfold.py 18.6 KB
Newer Older
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
1
2
# Copyright 2021 AlQuraishi Laboratory
# Copyright 2021 DeepMind Technologies Limited
3
#
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
4
5
6
7
8
9
10
11
12
13
14
# 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.
15
import argparse
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
16
import logging
17
import math
18
import numpy as np
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
19
import os
20
21
22
import pickle
import random
import time
23
import json
24

25
26
27
logging.basicConfig()
logger = logging.getLogger(__file__)
logger.setLevel(level=logging.INFO)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
28

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
29
import torch
30
31
32
torch_versions = torch.__version__.split(".")
torch_major_version = int(torch_versions[0])
torch_minor_version = int(torch_versions[1])
33
if (
34
    torch_major_version > 1 or
35
36
37
38
39
40
41
    (torch_major_version == 1 and torch_minor_version >= 12)
):
    # Gives a large speedup on Ampere-class GPUs
    torch.set_float32_matmul_precision("high")

torch.set_grad_enabled(False)

42
from openfold.config import model_config
43
from openfold.data import templates, feature_pipeline, data_pipeline
44
45
46
47
48
from openfold.data.tools import hhsearch, hmmsearch
from openfold.np import protein
from openfold.utils.script_utils import (load_models_from_command_line, parse_fasta, run_model,
                                         prep_output, relax_protein)
from openfold.utils.tensor_utils import tensor_tree_map
49
50
51
52
from openfold.utils.trace_utils import (
    pad_feature_dict_seq,
    trace_model_,
)
53

54
from scripts.precompute_embeddings import EmbeddingGenerator
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
55
from scripts.utils import add_data_args
56

57

58
TRACING_INTERVAL = 50
59
60


61
def precompute_alignments(tags, seqs, alignment_dir, args):
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
62
63
64
65
66
    for tag, seq in zip(tags, seqs):
        tmp_fasta_path = os.path.join(args.output_dir, f"tmp_{os.getpid()}.fasta")
        with open(tmp_fasta_path, "w") as fp:
            fp.write(f">{tag}\n{seq}")

67
        local_alignment_dir = os.path.join(alignment_dir, tag)
68

69
        if args.use_precomputed_alignments is None:
70
            logger.info(f"Generating alignments for {tag}...")
71

72
73
74
75
76
77
78
79
80
81
82
83
84
            os.makedirs(local_alignment_dir, exist_ok=True)

            if "multimer" in args.config_preset:
                template_searcher = hmmsearch.Hmmsearch(
                    binary_path=args.hmmsearch_binary_path,
                    hmmbuild_binary_path=args.hmmbuild_binary_path,
                    database_path=args.pdb_seqres_database_path,
                )
            else:
                template_searcher = hhsearch.HHSearch(
                    binary_path=args.hhsearch_binary_path,
                    databases=[args.pdb70_database_path],
                )
85

86
            # In seqemb mode, use AlignmentRunner only to generate templates
87
88
89
90
            if args.use_single_seq_mode:
                alignment_runner = data_pipeline.AlignmentRunner(
                    jackhmmer_binary_path=args.jackhmmer_binary_path,
                    uniref90_database_path=args.uniref90_database_path,
91
                    template_searcher=template_searcher,
92
93
                    no_cpus=args.cpus,
                )
94
                embedding_generator = EmbeddingGenerator()
95
                embedding_generator.run(tmp_fasta_path, alignment_dir)
96
97
98
99
100
101
102
            else:
                alignment_runner = data_pipeline.AlignmentRunner(
                    jackhmmer_binary_path=args.jackhmmer_binary_path,
                    hhblits_binary_path=args.hhblits_binary_path,
                    uniref90_database_path=args.uniref90_database_path,
                    mgnify_database_path=args.mgnify_database_path,
                    bfd_database_path=args.bfd_database_path,
103
                    uniref30_database_path=args.uniref30_database_path,
104
                    uniclust30_database_path=args.uniclust30_database_path,
105
                    uniprot_database_path=args.uniprot_database_path,
106
107
                    template_searcher=template_searcher,
                    use_small_bfd=args.bfd_database_path is None,
108
                    no_cpus=args.cpus
109
                )
110

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
111
            alignment_runner.run(
112
                tmp_fasta_path, local_alignment_dir
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
113
            )
114
115
116
117
        else:
            logger.info(
                f"Using precomputed alignments for {tag} at {alignment_dir}..."
            )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
118
119
120
121
122

        # Remove temporary FASTA file
        os.remove(tmp_fasta_path)


123
124
125
126
def round_up_seqlen(seqlen):
    return int(math.ceil(seqlen / TRACING_INTERVAL)) * TRACING_INTERVAL


127
128
129
130
131
132
133
def generate_feature_dict(
    tags,
    seqs,
    alignment_dir,
    data_processor,
    args,
):
134
    tmp_fasta_path = os.path.join(args.output_dir, f"tmp_{os.getpid()}.fasta")
135
136
137
138
139
140
141
142
143
144

    if "multimer" in args.config_preset:
        with open(tmp_fasta_path, "w") as fp:
            fp.write(
                '\n'.join([f">{tag}\n{seq}" for tag, seq in zip(tags, seqs)])
            )
        feature_dict = data_processor.process_fasta(
            fasta_path=tmp_fasta_path, alignment_dir=alignment_dir,
        )
    elif len(seqs) == 1:
145
        tag = tags[0]
146
147
148
149
150
151
        seq = seqs[0]
        with open(tmp_fasta_path, "w") as fp:
            fp.write(f">{tag}\n{seq}")

        local_alignment_dir = os.path.join(alignment_dir, tag)
        feature_dict = data_processor.process_fasta(
152
153
154
            fasta_path=tmp_fasta_path,
            alignment_dir=local_alignment_dir,
            seqemb_mode=args.use_single_seq_mode,
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
155
        )
156
157
158
159
160
161
162
163
164
165
166
167
    else:
        with open(tmp_fasta_path, "w") as fp:
            fp.write(
                '\n'.join([f">{tag}\n{seq}" for tag, seq in zip(tags, seqs)])
            )
        feature_dict = data_processor.process_multiseq_fasta(
            fasta_path=tmp_fasta_path, super_alignment_dir=alignment_dir,
        )

    # Remove temporary FASTA file
    os.remove(tmp_fasta_path)

168
    return feature_dict
169

170

171
172
def list_files_with_extensions(dir, extensions):
    return [f for f in os.listdir(dir) if f.endswith(extensions)]
173

174

175
def main(args):
176
    # Create the output directory
177
178
    os.makedirs(args.output_dir, exist_ok=True)

179
180
    if args.config_preset.startswith("seq"):
        args.use_single_seq_mode = True
181

182
183
184
185
186
    config = model_config(
        args.config_preset, 
        long_sequence_inference=args.long_sequence_inference,
        use_deepspeed_evoformer_attention=args.use_deepspeed_evoformer_attention,
        )
187

188
    if args.experiment_config_json:
189
190
191
        with open(args.experiment_config_json, 'r') as f:
            custom_config_dict = json.load(f)
        config.update_from_flattened_dict(custom_config_dict)
192

193
194
    if args.trace_model:
        if not config.data.predict.fixed_size:
195
196
197
            raise ValueError(
                "Tracing requires that fixed_size mode be enabled in the config"
            )
Christina Floristean's avatar
Christina Floristean committed
198
199

    is_multimer = "multimer" in args.config_preset
rostro36's avatar
rostro36 committed
200
    is_custom_template = "use_custom_template" in args and args.use_custom_template
rostro36's avatar
rostro36 committed
201
202
203
204
205
206
207
208
    if is_custom_template:
        template_featurizer = templates.CustomHitFeaturizer(
            mmcif_dir=args.template_mmcif_dir,
            max_template_date="9999-12-31", # just dummy, not used
            max_hits=-1, # just dummy, not used
            kalign_binary_path=args.kalign_binary_path
            )
    elif is_multimer:
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
        template_featurizer = templates.HmmsearchHitFeaturizer(
            mmcif_dir=args.template_mmcif_dir,
            max_template_date=args.max_template_date,
            max_hits=config.data.predict.max_templates,
            kalign_binary_path=args.kalign_binary_path,
            release_dates_path=args.release_dates_path,
            obsolete_pdbs_path=args.obsolete_pdbs_path
        )
    else:
        template_featurizer = templates.HhsearchHitFeaturizer(
            mmcif_dir=args.template_mmcif_dir,
            max_template_date=args.max_template_date,
            max_hits=config.data.predict.max_templates,
            kalign_binary_path=args.kalign_binary_path,
            release_dates_path=args.release_dates_path,
            obsolete_pdbs_path=args.obsolete_pdbs_path
        )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
226
    data_processor = data_pipeline.DataPipeline(
227
228
        template_featurizer=template_featurizer,
    )
229
    if is_multimer:
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
230
231
232
233
        data_processor = data_pipeline.DataPipelineMultimer(
            monomer_data_pipeline=data_processor,
        )

234
    output_dir_base = args.output_dir
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
235
    random_seed = args.data_random_seed
236
    if random_seed is None:
237
        random_seed = random.randrange(2 ** 32)
238

239
240
    np.random.seed(random_seed)
    torch.manual_seed(random_seed + 1)
241
    feature_processor = feature_pipeline.FeaturePipeline(config.data)
242
243
    if not os.path.exists(output_dir_base):
        os.makedirs(output_dir_base)
244
    if args.use_precomputed_alignments is None:
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
245
        alignment_dir = os.path.join(output_dir_base, "alignments")
Gustaf's avatar
Gustaf committed
246
247
    else:
        alignment_dir = args.use_precomputed_alignments
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
248

249
250
    tag_list = []
    seq_list = []
251
    for fasta_file in list_files_with_extensions(args.fasta_dir, (".fasta", ".fa")):
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
252
        # Gather input sequences
Christina Floristean's avatar
Christina Floristean committed
253
        fasta_path = os.path.join(args.fasta_dir, fasta_file)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
254
255
        with open(fasta_path, "r") as fp:
            data = fp.read()
256

257
        tags, seqs = parse_fasta(data)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
258

259
        if not is_multimer and len(tags) != 1:
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
260
261
262
            print(
                f"{fasta_path} contains more than one sequence but "
                f"multimer mode is not enabled. Skipping..."
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
263
            )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
264
            continue
Christina Floristean's avatar
Christina Floristean committed
265

266
        # assert len(tags) == len(set(tags)), "All FASTA tags must be unique"
267
        tag = '-'.join(tags)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
268

269
        tag_list.append((tag, tags))
270
271
272
273
274
        seq_list.append(seqs)

    seq_sort_fn = lambda target: sum([len(s) for s in target[1]])
    sorted_targets = sorted(zip(tag_list, seq_list), key=seq_sort_fn)
    feature_dicts = {}
275
276
277
278
279

    if is_multimer and args.openfold_checkpoint_path:
        raise ValueError(
            '`openfold_checkpoint_path` was specified, but no OpenFold checkpoints are available for multimer mode')

280
281
282
283
284
285
    model_generator = load_models_from_command_line(
        config,
        args.model_device,
        args.openfold_checkpoint_path,
        args.jax_param_path,
        args.output_dir)
286

287
    for model, output_directory in model_generator:
288
        cur_tracing_interval = 0
289
        for (tag, tags), seqs in sorted_targets:
290
291
292
            output_name = f'{tag}_{args.config_preset}'
            if args.output_postfix is not None:
                output_name = f'{output_name}_{args.output_postfix}'
293

294
            # Does nothing if the alignments have already been computed
295
            precompute_alignments(tags, seqs, alignment_dir, args)
296

297
            feature_dict = feature_dicts.get(tag, None)
298
            if feature_dict is None:
299
300
301
302
303
304
                feature_dict = generate_feature_dict(
                    tags,
                    seqs,
                    alignment_dir,
                    data_processor,
                    args,
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
305
306
                )

307
                if args.trace_model:
308
309
310
311
312
                    n = feature_dict["aatype"].shape[-2]
                    rounded_seqlen = round_up_seqlen(n)
                    feature_dict = pad_feature_dict_seq(
                        feature_dict, rounded_seqlen,
                    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
313

314
315
                feature_dicts[tag] = feature_dict
            processed_feature_dict = feature_processor.process_features(
Christina Floristean's avatar
Christina Floristean committed
316
                feature_dict, mode='predict', is_multimer=is_multimer
317
318
319
            )

            processed_feature_dict = {
320
321
                k: torch.as_tensor(v, device=args.model_device)
                for k, v in processed_feature_dict.items()
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
322
            }
323

324
325
            if args.trace_model:
                if rounded_seqlen > cur_tracing_interval:
326
327
328
329
330
                    logger.info(
                        f"Tracing model at {rounded_seqlen} residues..."
                    )
                    t = time.perf_counter()
                    trace_model_(model, processed_feature_dict)
331
                    tracing_time = time.perf_counter() - t
332
                    logger.info(
333
                        f"Tracing time: {tracing_time}"
334
335
                    )
                    cur_tracing_interval = rounded_seqlen
Sam DeLuca's avatar
Sam DeLuca committed
336

337
            out = run_model(model, processed_feature_dict, tag, args.output_dir)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
338

339
            # Toss out the recycling dimensions --- we don't need them anymore
340
            processed_feature_dict = tensor_tree_map(
341
                lambda x: np.array(x[..., -1].cpu()),
342
343
                processed_feature_dict
            )
344
345
346
            out = tensor_tree_map(lambda x: np.array(x.cpu()), out)

            unrelaxed_protein = prep_output(
347
348
349
350
                out,
                processed_feature_dict,
                feature_dict,
                feature_processor,
351
                args.config_preset,
352
353
                args.multimer_ri_gap,
                args.subtract_plddt
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
354
            )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
355

356
357
358
            unrelaxed_file_suffix = "_unrelaxed.pdb"
            if args.cif_output:
                unrelaxed_file_suffix = "_unrelaxed.cif"
359
            unrelaxed_output_path = os.path.join(
360
                output_directory, f'{output_name}{unrelaxed_file_suffix}'
361
            )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
362

363
            with open(unrelaxed_output_path, 'w') as fp:
364
365
366
367
                if args.cif_output:
                    fp.write(protein.to_modelcif(unrelaxed_protein))
                else:
                    fp.write(protein.to_pdb(unrelaxed_protein))
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
368

369
            logger.info(f"Output written to {unrelaxed_output_path}...")
370

371
372
            if not args.skip_relaxation:
                # Relax the prediction.
373
                logger.info(f"Running relaxation on {unrelaxed_output_path}...")
374
375
                relax_protein(config, args.model_device, unrelaxed_protein, output_directory, output_name,
                              args.cif_output)
376

377
378
            if args.save_outputs:
                output_dict_path = os.path.join(
379
                    output_directory, f'{output_name}_output_dict.pkl'
380
381
382
                )
                with open(output_dict_path, "wb") as fp:
                    pickle.dump(out, fp, protocol=pickle.HIGHEST_PROTOCOL)
383

Sam DeLuca's avatar
Sam DeLuca committed
384
                logger.info(f"Model output written to {output_dict_path}...")
385
386
387
388


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
389
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
390
        "fasta_dir", type=str,
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
391
        help="Path to directory containing FASTA files, one sequence per file"
392
    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
393
394
395
    parser.add_argument(
        "template_mmcif_dir", type=str,
    )
Gustaf's avatar
Gustaf committed
396
397
398
399
    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
400
    )
rostro36's avatar
rostro36 committed
401
402
403
404
    parser.add_argument(
        "--use_custom_template", action="store_true", default=False,
        help="""Use mmcif given with "template_mmcif_dir" argument as template input."""
    )
405
406
407
408
    parser.add_argument(
        "--use_single_seq_mode", action="store_true", default=False,
        help="""Use single sequence embeddings instead of MSAs."""
    )
409
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
410
411
        "--output_dir", type=str, default=os.getcwd(),
        help="""Name of the directory in which to output the prediction""",
412
413
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
414
        "--model_device", type=str, default="cpu",
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
415
416
        help="""Name of the device on which to run the model. Any valid torch
             device name is accepted (e.g. "cpu", "cuda:0")"""
417
418
    )
    parser.add_argument(
419
        "--config_preset", type=str, default="model_1",
420
        help="""Name of a model config preset defined in openfold/config.py"""
421
422
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
423
424
425
426
        "--jax_param_path", type=str, default=None,
        help="""Path to JAX model parameters. If None, and openfold_checkpoint_path
             is also None, parameters are selected automatically according to 
             the model name from openfold/resources/params"""
427
428
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
429
430
431
        "--openfold_checkpoint_path", type=str, default=None,
        help="""Path to OpenFold checkpoint. Can be either a DeepSpeed 
             checkpoint directory or a .pt file"""
432
    )
433
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
434
        "--save_outputs", action="store_true", default=False,
435
436
        help="Whether to save all model outputs, including embeddings, etc."
    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
437
438
    parser.add_argument(
        "--cpus", type=int, default=4,
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
439
        help="""Number of CPUs with which to run alignment tools"""
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
440
    )
441
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
442
        "--preset", type=str, default='full_dbs',
443
444
445
        choices=('reduced_dbs', 'full_dbs')
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
446
447
448
        "--output_postfix", type=str, default=None,
        help="""Postfix for output prediction filenames"""
    )
449
    parser.add_argument(
450
        "--data_random_seed", type=int, default=None
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
451
452
453
    )
    parser.add_argument(
        "--skip_relaxation", action="store_true", default=False,
454
    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
455
456
457
458
    parser.add_argument(
        "--multimer_ri_gap", type=int, default=200,
        help="""Residue index offset between multiple sequences, if provided"""
    )
459
460
461
462
463
464
    parser.add_argument(
        "--trace_model", action="store_true", default=False,
        help="""Whether to convert parts of each model to TorchScript.
                Significantly improves runtime at the cost of lengthy
                'compilation.' Useful for large batch jobs."""
    )
465
466
467
468
469
    parser.add_argument(
        "--subtract_plddt", action="store_true", default=False,
        help=""""Whether to output (100 - pLDDT) in the B-factor column instead
                 of the pLDDT itself"""
    )
470
471
472
    parser.add_argument(
        "--long_sequence_inference", action="store_true", default=False,
        help="""enable options to reduce memory usage at the cost of speed, helps longer sequences fit into GPU memory, see the README for details"""
473
    )
474
475
476
477
    parser.add_argument(
        "--cif_output", action="store_true", default=False,
        help="Output predicted models in ModelCIF format instead of PDB format (default)"
    )
478
479
480
    parser.add_argument(
        "--experiment_config_json", default="", help="Path to a json file with custom config values to overwrite config setting",
    )
481
482
483
484
    parser.add_argument(
        "--use_deepspeed_evoformer_attention", action="store_true", default=False, 
        help="Whether to use the DeepSpeed evoformer attention layer. Must have deepspeed installed in the environment.",
    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
485
    add_data_args(parser)
486
487
    args = parser.parse_args()

488
    if args.jax_param_path is None and args.openfold_checkpoint_path is None:
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
489
        args.jax_param_path = os.path.join(
490
            "openfold", "resources", "params",
491
            "params_" + args.config_preset + ".npz"
492
493
        )

494
    if args.model_device == "cpu" and torch.cuda.is_available():
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
495
496
497
498
        logging.warning(
            """The model is being run on CPU. Consider specifying 
            --model_device for better performance"""
        )
499
    main(args)