run_pretrained_openfold.py 16.7 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
Sam DeLuca's avatar
Sam DeLuca committed
22
from copy import deepcopy
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
23

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
24
import pickle
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
25
26
27
from pytorch_lightning.utilities.deepspeed import (
    convert_zero_checkpoint_to_fp32_state_dict
)
28
29
import random
import sys
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
30
31
32
import time
import torch

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

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
46
from scripts.utils import add_data_args
47

48

49
50
51
52
53
logging.basicConfig()
logger = logging.getLogger(__file__)
logger.setLevel(level=logging.INFO)


Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
54
55
56
57
58
59
60
61
def precompute_alignments(tags, seqs, alignment_dir, args):
    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}")

        local_alignment_dir = os.path.join(alignment_dir, tag)
        if(args.use_precomputed_alignments is None):
62
            logger.info(f"Generating alignments for {tag}...")
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
63
64
            if not os.path.exists(local_alignment_dir):
                os.makedirs(local_alignment_dir)
65

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
66
67
68
69
70
71
72
73
74
75
76
77
            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,
                no_cpus=args.cpus,
            )
            alignment_runner.run(
78
                tmp_fasta_path, local_alignment_dir
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
79
80
81
82
83
84
85
86
87
88
89
90
91
92
            )

        # Remove temporary FASTA file
        os.remove(tmp_fasta_path)


def run_model(model, batch, tag, args):
    with torch.no_grad():
        batch = {
            k:torch.as_tensor(v, device=args.model_device) 
            for k,v in batch.items()
        }
 
        # Disable templates if there aren't any in the batch
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
93
        model.config.template.enabled = model.config.template.enabled and any([
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
94
95
96
            "template_" in k for k in batch
        ])

97
        logger.info(f"Running inference for {tag}...")
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
98
99
        t = time.perf_counter()
        out = model(batch)
100
        logger.info(f"Inference time: {time.perf_counter() - t}")
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
    
    return out


def prep_output(out, batch, feature_dict, feature_processor, args):
    plddt = out["plddt"]
    mean_plddt = np.mean(plddt)
    
    plddt_b_factors = np.repeat(
        plddt[..., None], residue_constants.atom_type_num, axis=-1
    )

    # Prep protein metadata
    template_domain_names = []
    template_chain_index = None
116
    if(feature_processor.config.common.use_templates and "template_domain_names" in feature_dict):
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
        template_domain_names = [
            t.decode("utf-8") for t in feature_dict["template_domain_names"]
        ]

        # This works because templates are not shuffled during inference
        template_domain_names = template_domain_names[
            :feature_processor.config.predict.max_templates
        ]

        if("template_chain_index" in feature_dict):
            template_chain_index = feature_dict["template_chain_index"]
            template_chain_index = template_chain_index[
                :feature_processor.config.predict.max_templates
            ]

    no_recycling = feature_processor.config.common.max_recycling_iters
    remark = ', '.join([
        f"no_recycling={no_recycling}",
        f"max_templates={feature_processor.config.predict.max_templates}",
136
        f"config_preset={args.config_preset}",
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
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
    ])

    # For multi-chain FASTAs
    ri = feature_dict["residue_index"]
    chain_index = (ri - np.arange(ri.shape[0])) / args.multimer_ri_gap
    chain_index = chain_index.astype(np.int64)
    cur_chain = 0
    prev_chain_max = 0
    for i, c in enumerate(chain_index):
        if(c != cur_chain):
            cur_chain = c
            prev_chain_max = i + cur_chain * args.multimer_ri_gap

        batch["residue_index"][i] -= prev_chain_max

    unrelaxed_protein = protein.from_prediction(
        features=batch,
        result=out,
        b_factors=plddt_b_factors,
        chain_index=chain_index,
        remark=remark,
        parents=template_domain_names,
        parents_chain_index=template_chain_index,
    )

    return unrelaxed_protein


165
def generate_batch(fasta_file, fasta_dir, alignment_dir, data_processor, feature_processor, prediction_dir):
166
167
    with open(os.path.join(fasta_dir, fasta_file), "r") as fp:
        data = fp.read()
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
168

169
170
171
172
173
    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
174

175
    tags = [t.split()[0] for t in tags]
176
    # assert len(tags) == len(set(tags)), "All FASTA tags must be unique"
177
178
    tag = '-'.join(tags)

179
180
181
182
183
184
185
186
187
188
189
190
    output_name = f'{tag}_{args.config_preset}'
    if args.output_postfix is not None:
        output_name = f'{output_name}_{args.output_postfix}'

    # Save the unrelaxed PDB.
    unrelaxed_output_path = os.path.join(
        prediction_dir, f'{output_name}_unrelaxed.pdb'
    )

    if os.path.exists(unrelaxed_output_path):
        return

191
192
193
194
195
196
197
198
199
200
201
    precompute_alignments(tags, seqs, alignment_dir, args)

    tmp_fasta_path = os.path.join(args.output_dir, f"tmp_{os.getpid()}.fasta")
    if len(seqs) == 1:
        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(
            fasta_path=tmp_fasta_path, alignment_dir=local_alignment_dir
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
202
        )
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
    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)

    processed_feature_dict = feature_processor.process_features(
        feature_dict, mode='predict',
    )
    return processed_feature_dict, tag, feature_dict


def load_models_from_command_line(args, config):
    # Create the output directory
    os.makedirs(args.output_dir, exist_ok=True)
    if args.jax_param_path:
        for path in args.jax_param_path.split(","):
            model = AlphaFold(config)
            model = model.eval()
            import_jax_weights_(
                model, path, version=args.model_name
            )
            model = model.to(args.model_device)
232
233
234
            logger.info(
                f"Successfully loaded JAX parameters at {args.jax_param_path}..."
            )
235
236
            yield model, None
    if args.openfold_checkpoint_path:
Sam DeLuca's avatar
wip  
Sam DeLuca committed
237
        for path in args.openfold_checkpoint_path.split(","):
238
239
            model = AlphaFold(config)
            model = model.eval()
Sam DeLuca's avatar
Sam DeLuca committed
240
241
242
243
244
            checkpoint_basename = os.path.splitext(
                os.path.basename(
                    os.path.normpath(path)
                )
            )[0]
245
            if os.path.isdir(path):
246
                # A DeepSpeed checkpoint
247
248
249
                ckpt_path = os.path.join(
                    args.output_dir,
                    checkpoint_basename + ".pt",
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
250
251
                )

252
253
                if not os.path.isfile(ckpt_path):
                    convert_zero_checkpoint_to_fp32_state_dict(
Sam DeLuca's avatar
wip  
Sam DeLuca committed
254
                        path,
255
256
                        ckpt_path,
                    )
257
258
                d = torch.load(ckpt_path)
                model.load_state_dict(d["ema"]["params"])
259
260
            else:
                ckpt_path = path
261
                d = torch.load(ckpt_path)
262

263
264
265
266
                if ("ema" in d):
                    # The public weights have had this done to them already
                    d = d["ema"]["params"]
                model.load_state_dict(d)
267
            model = model.to(args.model_device)
268
269
270
            logger.info(
                f"Loaded OpenFold parameters at {args.openfold_checkpoint_path}..."
            )
271
            yield model, checkpoint_basename
Sam DeLuca's avatar
wip  
Sam DeLuca committed
272
    if not args.jax_param_path and not args.openfold_checkpoint_path:
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
273
274
275
276
277
        raise ValueError(
            "At least one of jax_param_path or openfold_checkpoint_path must "
            "be specified."
        )

278
279
280
281
282

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

283
    config = model_config(args.config_preset)
284
285
286
    template_featurizer = templates.TemplateHitFeaturizer(
        mmcif_dir=args.template_mmcif_dir,
        max_template_date=args.max_template_date,
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
287
        max_hits=config.data.predict.max_templates,
288
        kalign_binary_path=args.kalign_binary_path,
289
        release_dates_path=args.release_dates_path,
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
290
291
        obsolete_pdbs_path=args.obsolete_pdbs_path
    )
292

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
293
    data_processor = data_pipeline.DataPipeline(
294
295
296
297
        template_featurizer=template_featurizer,
    )

    output_dir_base = args.output_dir
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
298
    random_seed = args.data_random_seed
299
300
    if random_seed is None:
        random_seed = random.randrange(sys.maxsize)
301
    feature_processor = feature_pipeline.FeaturePipeline(config.data)
302
303
    if not os.path.exists(output_dir_base):
        os.makedirs(output_dir_base)
304
    if args.use_precomputed_alignments is None:
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
305
        alignment_dir = os.path.join(output_dir_base, "alignments")
Gustaf's avatar
Gustaf committed
306
307
    else:
        alignment_dir = args.use_precomputed_alignments
308
        logger.info(f"Using precomputed alignments at {alignment_dir}...")
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
309

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
310
311
312
    prediction_dir = os.path.join(args.output_dir, "predictions")
    os.makedirs(prediction_dir, exist_ok=True)

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
313
    for fasta_file in os.listdir(args.fasta_dir):
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
314

Sam DeLuca's avatar
Sam DeLuca committed
315
        batch_data = generate_batch(
316
317
318
319
320
321
            fasta_file,
            args.fasta_dir,
            alignment_dir,
            data_processor,
            feature_processor,
            prediction_dir)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
322

Sam DeLuca's avatar
Sam DeLuca committed
323
324
325
326
327
328
        if batch_data is None:
            # this file has already been processed
            continue

        batch, tag, feature_dict = batch_data

329
        for model, model_version in load_models_from_command_line(args, config):
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
330

Sam DeLuca's avatar
Sam DeLuca committed
331
332
            working_batch = deepcopy(batch)
            out = run_model(model, working_batch, tag, args)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
333

334
            # Toss out the recycling dimensions --- we don't need them anymore
Sam DeLuca's avatar
Sam DeLuca committed
335
            working_batch = tensor_tree_map(lambda x: np.array(x[..., -1].cpu()), working_batch)
336
337
338
            out = tensor_tree_map(lambda x: np.array(x.cpu()), out)

            unrelaxed_protein = prep_output(
Sam DeLuca's avatar
Sam DeLuca committed
339
                out, working_batch, feature_dict, feature_processor, args
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
340
            )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
341

Sam DeLuca's avatar
bugfix  
Sam DeLuca committed
342
            output_name = f'{tag}_{args.config_preset}'
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
343

344
345
346
347
            if model_version is not None:
                output_name = f'{output_name}_{model_version}'
            if args.output_postfix is not None:
                output_name = f'{output_name}_{args.output_postfix}'
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
348

349
350
351
352
353
354
            # Save the unrelaxed PDB.
            unrelaxed_output_path = os.path.join(
                prediction_dir, f'{output_name}_unrelaxed.pdb'
            )
            with open(unrelaxed_output_path, 'w') as fp:
                fp.write(protein.to_pdb(unrelaxed_protein))
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
355

356
            logger.info(f"Output written to {unrelaxed_output_path}...")
357
358
359
360
361
            if not args.skip_relaxation:
                amber_relaxer = relax.AmberRelaxation(
                    use_gpu=(args.model_device != "cpu"),
                    **config.relax,
                )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
362

363
                # Relax the prediction.
364
                logger.info(f"Running relaxation on {unrelaxed_output_path}...")
365
366
367
368
369
370
371
                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
372
                logger.info(f"Relaxation time: {time.perf_counter() - t}")
373
374
375
376
377
378
379

                # Save the relaxed PDB.
                relaxed_output_path = os.path.join(
                    prediction_dir, f'{output_name}_relaxed.pdb'
                )
                with open(relaxed_output_path, 'w') as fp:
                    fp.write(relaxed_pdb_str)
380
                logger.info(f"Relaxed output written to {relaxed_output_path}...")
381

382
383
384
385
386
387
            if args.save_outputs:
                output_dict_path = os.path.join(
                    args.output_dir, f'{output_name}_output_dict.pkl'
                )
                with open(output_dict_path, "wb") as fp:
                    pickle.dump(out, fp, protocol=pickle.HIGHEST_PROTOCOL)
388

Sam DeLuca's avatar
Sam DeLuca committed
389
                logger.info(f"Model output written to {output_dict_path}...")
390
391
392

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
393
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
394
395
        "fasta_dir", type=str,
        help="Path to directory containing FASTA files, one sequence per file"
396
    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
397
398
399
    parser.add_argument(
        "template_mmcif_dir", type=str,
    )
Gustaf's avatar
Gustaf committed
400
401
402
403
    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
404
    )
405
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
406
407
        "--output_dir", type=str, default=os.getcwd(),
        help="""Name of the directory in which to output the prediction""",
408
409
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
410
        "--model_device", type=str, default="cpu",
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
411
412
        help="""Name of the device on which to run the model. Any valid torch
             device name is accepted (e.g. "cpu", "cuda:0")"""
413
414
    )
    parser.add_argument(
415
        "--config_preset", type=str, default="model_1",
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
416
417
        help="""Name of a model config. Choose one of model_{1-5} or 
             model_{1-5}_ptm, as defined on the AlphaFold GitHub."""
418
419
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
420
421
422
423
424
425
426
427
428
        "--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"""
    )
    parser.add_argument(
        "--openfold_checkpoint_path", type=str, default=None,
        help="""Path to OpenFold checkpoint. Can be either a DeepSpeed 
             checkpoint directory or a .pt file"""
429
    )
430
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
431
        "--save_outputs", action="store_true", default=False,
432
433
        help="Whether to save all model outputs, including embeddings, etc."
    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
434
435
    parser.add_argument(
        "--cpus", type=int, default=4,
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
436
        help="""Number of CPUs with which to run alignment tools"""
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
437
    )
438
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
439
        "--preset", type=str, default='full_dbs',
440
441
        choices=('reduced_dbs', 'full_dbs')
    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
442
443
444
445
    parser.add_argument(
        "--output_postfix", type=str, default=None,
        help="""Postfix for output prediction filenames"""
    )
446
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
447
448
449
450
        "--data_random_seed", type=str, default=None
    )
    parser.add_argument(
        "--skip_relaxation", action="store_true", default=False,
451
    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
452
453
454
455
    parser.add_argument(
        "--multimer_ri_gap", type=int, default=200,
        help="""Residue index offset between multiple sequences, if provided"""
    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
456
    add_data_args(parser)
457
458
    args = parser.parse_args()

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
459
460
    if(args.jax_param_path is None and args.openfold_checkpoint_path is None):
        args.jax_param_path = os.path.join(
461
            "openfold", "resources", "params", 
462
            "params_" + args.config_preset + ".npz"
463
464
        )

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
465
466
467
468
469
470
    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"""
        )

471
    main(args)