|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| import json
|
| import os
|
| import subprocess
|
| import tempfile
|
| from typing import Any, Dict, List
|
|
|
| from pxdbench.globals import MPNN_CKPT_PATH, _require
|
|
|
|
|
| class VanillaMPNNPredictor:
|
| def __init__(self, cfg, device_id: int = 0, seed: int = None, verbose=True):
|
| self.cfg = cfg
|
| self.device_id = device_id
|
| self.seed = seed
|
| self.verbose = verbose
|
| _require(
|
| os.path.join(
|
| MPNN_CKPT_PATH[self.cfg["model_type"]], self.cfg["model_name"] + ".pt"
|
| )
|
| )
|
|
|
| dir_name = os.path.dirname(__file__)
|
| self.prepare_script_path = os.path.join(dir_name, "parse_multiple_chains.py")
|
| self.run_script_path = os.path.join(dir_name, "protein_mpnn_run.py")
|
| self.env = os.environ.copy()
|
| self.env["CUDA_VISIBLE_DEVICES"] = str(device_id)
|
|
|
| def prepare_jsonl(self, input_dir: str, pdb_names: list[str]):
|
| """
|
| Prepare a JSONL file for ProteinMPNN input by parsing multiple chains from PDB files.
|
|
|
| Args:
|
| input_dir (str): Directory containing input PDB files.
|
| pdb_names (list[str]): List of PDB base names (without '.pdb') to process.
|
|
|
| Returns:
|
| str: Path to the output JSONL file.
|
| """
|
|
|
| with tempfile.NamedTemporaryFile(suffix=".jsonl", delete=False) as f:
|
| output_path = f.name
|
|
|
| output_path = os.path.join(input_dir, "parsed_pdbs.jsonl")
|
| cmd = [
|
| "python3",
|
| "-u",
|
| self.prepare_script_path,
|
| "--input_path",
|
| input_dir,
|
| "--output",
|
| output_path,
|
| ]
|
|
|
| try:
|
| process = subprocess.Popen(
|
| cmd,
|
| stdout=subprocess.PIPE,
|
| stderr=subprocess.PIPE,
|
| text=True,
|
| env=self.env,
|
| )
|
| while True:
|
| output = process.stdout.readline()
|
| if output == "" and process.poll() is not None:
|
| break
|
| if output:
|
| print(output.strip())
|
|
|
| while True:
|
| error = process.stderr.readline()
|
| if error == "" and process.poll() is not None:
|
| break
|
| if error:
|
| print(error.strip())
|
|
|
| returncode = process.wait()
|
| print(f"Run subprocess success: {returncode}")
|
|
|
| except Exception as e:
|
| print(f"Run subprocess fail: {str(e)}")
|
| raise
|
|
|
| result = []
|
| with open(output_path, "r") as infile:
|
| for line in infile:
|
| entry = json.loads(line)
|
| if entry.get("name") in pdb_names:
|
| result.append(entry)
|
| with open(output_path, "w") as f:
|
| for entry in result:
|
| f.write(json.dumps(entry) + "\n")
|
| return output_path
|
|
|
| def run_mpnn(self, jsonl_path: str, pdb_names: list[str], num_seqs: int):
|
| """
|
| Run ProteinMPNN to design sequences for the given PDB names.
|
|
|
| Args:
|
| jsonl_path (str): Path to the input JSONL file containing parsed PDB data.
|
| pdb_names (list[str]): List of PDB base names (without '.pdb') to process.
|
| num_seqs (int): Number of sequences to generate per PDB.
|
|
|
| Returns:
|
| None
|
| """
|
|
|
| model_type = self.cfg["model_type"]
|
| model_name = self.cfg["model_name"]
|
| assert model_type in ["ca", "bb", "soluable"]
|
| path_to_model_weights = MPNN_CKPT_PATH[model_type]
|
|
|
| output_dir = os.path.dirname(jsonl_path)
|
| cmd = [
|
| "python3",
|
| "-u",
|
| self.run_script_path,
|
| "--jsonl_path",
|
| jsonl_path,
|
| "--out_folder",
|
| output_dir,
|
| "--num_seq_per_target",
|
| str(num_seqs),
|
| "--sampling_temp",
|
| self.cfg["temperature"],
|
| "--batch_size",
|
| str(num_seqs),
|
| "--model_name",
|
| model_name,
|
| "--path_to_model_weights",
|
| path_to_model_weights,
|
| ]
|
| if model_type == "ca":
|
| cmd.extend(["--ca_only"])
|
| elif model_type == "soluble":
|
| cmd.extend(["--use_soluble_model"])
|
| if self.seed is not None:
|
| cmd.extend(["--seed", str(self.seed)])
|
|
|
| print("Run Protein MPNN with %s" % (" ".join(cmd)))
|
| try:
|
| process = subprocess.Popen(
|
| cmd,
|
| stdout=subprocess.PIPE,
|
| stderr=subprocess.PIPE,
|
| text=True,
|
| env=self.env,
|
| )
|
| while True:
|
| output = process.stdout.readline()
|
| if output == "" and process.poll() is not None:
|
| break
|
| if output:
|
| print(output.strip())
|
|
|
| while True:
|
| error = process.stderr.readline()
|
| if error == "" and process.poll() is not None:
|
| break
|
| if error:
|
| print(error.strip())
|
|
|
| returncode = process.wait()
|
| print(f"Run subprocess success: {returncode}")
|
|
|
| except Exception as e:
|
| print(f"Run subprocess fail: {str(e)}")
|
| raise
|
|
|
|
|
| seq_dir = os.path.join(output_dir, "seqs")
|
| result = []
|
| for name in pdb_names:
|
| fasta_file = name + ".fa"
|
| fasta_path = os.path.join(seq_dir, fasta_file)
|
| if not os.path.exists(fasta_path):
|
| continue
|
|
|
| with open(fasta_path, "r") as f:
|
| lines = f.readlines()
|
|
|
| for i, line_no in enumerate(
|
| range(4, int(num_seqs) * 2 + 3, 2)
|
| ):
|
| if line_no <= len(lines):
|
| sequence = lines[line_no - 1].strip()
|
| result.append({"name": name, "seq_idx": i, "sequence": sequence})
|
| return result
|
|
|
| def design_monomer(
|
| self, pdb_dir: str, pdb_names: List[str], num_samples: int
|
| ) -> List[Dict]:
|
| """
|
| Design sequences for monomer proteins using ProteinMPNN.
|
|
|
| Args:
|
| pdb_dir (str): Directory containing input PDB files.
|
| pdb_names (list[str]): List of PDB base names (without '.pdb') to process.
|
| num_samples (int): Number of sequences to generate per PDB.
|
|
|
| Returns:
|
| list[dict]: List of design results with keys 'name' (PDB name), 'seq_idx' (sequence index),
|
| and 'sequence' (designed amino acid sequence).
|
| """
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| jsonl_path = self.prepare_jsonl(pdb_dir, pdb_names)
|
| output = self.run_mpnn(jsonl_path, pdb_names, num_samples)
|
| os.unlink(jsonl_path)
|
| return output
|
|
|
| def design_binder(
|
| self,
|
| pdb_dir: str,
|
| pdb_names: List[str],
|
| num_samples: int,
|
| binder_chains: List[str],
|
| cond_chains: List[str],
|
| ) -> List[Dict]:
|
| """
|
| Design sequences for binder proteins using ProteinMPNN.
|
|
|
| Args:
|
| pdb_dir (str): Directory containing input PDB files.
|
| pdb_names (list[str]): List of PDB base names (without '.pdb') to process.
|
| num_samples (int): Number of sequences to generate per PDB.
|
| binder_chains (list[str]): List of binder chain IDs.
|
| cond_chains (list[str]): List of conditional chain IDs.
|
|
|
| Returns:
|
| list[dict]: List of design results with keys 'name' (PDB name), 'seq_idx' (sequence index),
|
| and 'sequence' (designed amino acid sequence).
|
| """
|
| input_data = {
|
| "pdb_dir": pdb_dir,
|
| "pdb_names": pdb_names,
|
| "num_samples": num_samples,
|
| "binder_chains": binder_chains,
|
| "cond_chains": cond_chains,
|
| "mpnn_cfg": self.cfg.to_dict(),
|
| "design_type": "binder",
|
| }
|
| output = self.run(input_data)
|
| return output
|
|
|
