Spaces:
Build error
Build error
File size: 10,242 Bytes
8ef403e |
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 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 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 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 |
#!/usr/bin/env python3
"""
Structure preparation module for AbMelt inference pipeline.
Handles antibody structure generation and preprocessing.
"""
import os
import sys
import logging
from pathlib import Path
import numpy as np
from typing import Dict, List, Optional, Tuple
# Add the original AbMelt src to path for imports
sys.path.append(str(Path(__file__).parent.parent.parent / "AbMelt" / "src"))
try:
from structure import immune_builder
from Bio.PDB import PDBParser
from Bio.SeqUtils import seq1
except ImportError as e:
logging.error(f"Failed to import required modules: {e}")
raise
logger = logging.getLogger(__name__)
def prepare_structure(antibody: Dict, config: Dict) -> Dict[str, str]:
"""
Prepare antibody structure for MD simulation.
Args:
antibody: Dictionary containing antibody information
config: Configuration dictionary
Returns:
Dictionary with paths to prepared structure files
"""
logger.info(f"Preparing structure for antibody: {antibody['name']}")
if antibody["type"] == "pdb":
return _prepare_from_pdb(antibody, config)
elif antibody["type"] == "sequences":
return _prepare_from_sequences(antibody, config)
else:
raise ValueError(f"Unsupported antibody type: {antibody['type']}")
def _prepare_from_pdb(antibody: Dict, config: Dict) -> Dict[str, str]:
"""Prepare structure from existing PDB file."""
pdb_file = antibody["pdb_file"]
antibody_name = antibody["name"]
# Create working directory
work_dir = Path(config["paths"]["temp_dir"])
work_dir.mkdir(parents=True, exist_ok=True)
# Copy PDB to working directory
import shutil
target_pdb = work_dir / f"{antibody_name}.pdb"
shutil.copy2(pdb_file, target_pdb)
# Preprocess the structure
return _preprocess_structure(str(target_pdb), antibody_name, work_dir, config)
def _prepare_from_sequences(antibody: Dict, config: Dict) -> Dict[str, str]:
"""Generate structure from heavy and light chain sequences."""
heavy_chain = antibody["heavy_chain"]
light_chain = antibody["light_chain"]
antibody_name = antibody["name"]
# Create working directory
work_dir = Path(config["paths"]["temp_dir"])
work_dir.mkdir(parents=True, exist_ok=True)
# Generate structure using ImmuneBuilder
logger.info("Generating structure using ImmuneBuilder...")
sequence = {'H': heavy_chain, 'L': light_chain}
pdb_file = work_dir / f"{antibody_name}.pdb"
try:
immune_builder(sequence, str(pdb_file))
logger.info(f"Structure generated: {pdb_file}")
except Exception as e:
logger.error(f"Failed to generate structure: {e}")
raise
# Preprocess the structure
return _preprocess_structure(str(pdb_file), antibody_name, work_dir, config)
def _preprocess_structure(pdb_file: str, antibody_name: str, work_dir: Path, config: Dict) -> Dict[str, str]:
"""Basic structure validation and organization."""
logger.info("Validating structure...")
# Validate the PDB structure
if not validate_structure(pdb_file):
logger.error("Structure validation failed")
raise Exception("Structure validation failed")
# Rename chains to A/B convention for compatibility with downstream code
logger.info("Renaming chains to A/B convention...")
rename_chains_to_ab(pdb_file)
# Extract chain sequences for reference
chains = get_chain_sequences(pdb_file)
logger.info(f"Found chains: {list(chains.keys())}")
# Return basic structure information
structure_files = {
"pdb_file": str(work_dir / f"{antibody_name}.pdb"),
"work_dir": str(work_dir),
"chains": chains
}
logger.info("Structure preparation completed successfully")
return structure_files
def validate_structure(pdb_file: str) -> bool:
"""Validate that the PDB file contains a proper antibody structure."""
try:
parser = PDBParser(QUIET=True)
structure = parser.get_structure("antibody", pdb_file)
# Check for heavy and light chains
chains = list(structure.get_chains())
if len(chains) < 2:
logger.warning("Structure should contain at least 2 chains (heavy and light)")
return False
# Basic validation passed
return True
except Exception as e:
logger.error(f"Structure validation failed: {e}")
return False
def rename_chains_to_ab(pdb_file: str) -> None:
"""
Rename chains in PDB file to use A/B convention instead of H/L.
This ensures compatibility with downstream code that expects A=light, B=heavy.
Args:
pdb_file: Path to PDB file to modify in-place
"""
from Bio.PDB import PDBIO
from anarci import anarci
try:
parser = PDBParser(QUIET=True)
structure = parser.get_structure("antibody", pdb_file)
# Identify which chain is heavy vs light using ANARCI
chains = list(structure.get_chains())
chain_mapping = {} # old_id -> new_id
for chain in chains:
chain_id = chain.id
sequence = seq1(''.join(residue.resname for residue in chain))
# Use ANARCI to identify chain type
results = anarci([(chain_id, sequence)], scheme="imgt", output=False)
numbering, alignment_details, hit_tables = results
if numbering[0] and numbering[0][0] and hit_tables and len(hit_tables[0]) > 1:
best_hit = hit_tables[0][1]
hit_id = best_hit[0]
if '_' in hit_id:
chain_type = hit_id.split('_')[1]
if chain_type in ['K', 'L']: # Light chain
chain_mapping[chain_id] = 'A'
elif chain_type == 'H': # Heavy chain
chain_mapping[chain_id] = 'B'
# Rename chains if needed
if chain_mapping:
logger.info(f"Renaming chains: {chain_mapping}")
for chain in structure.get_chains():
if chain.id in chain_mapping:
chain.id = chain_mapping[chain.id]
# Save modified structure
io = PDBIO()
io.set_structure(structure)
io.save(pdb_file)
logger.info(f"Chains renamed successfully. New chain IDs: {[c.id for c in structure.get_chains()]}")
else:
logger.warning("Could not identify chains for renaming")
except Exception as e:
logger.error(f"Failed to rename chains: {e}")
raise
def get_chain_sequences(pdb_file: str) -> Dict[str, str]:
"""Extract heavy and light chain sequences from PDB file."""
try:
parser = PDBParser(QUIET=True)
structure = parser.get_structure("antibody", pdb_file)
# print unique chain ids
print(f"Unique chain ids: {list(set(chain.id for chain in structure.get_chains()))}")
chains = {chain.id: seq1(''.join(residue.resname for residue in chain))
for chain in structure.get_chains()}
return chains
except Exception as e:
logger.error(f"Failed to extract sequences: {e}")
return {}
# Convenience functions for direct usage
def generate_structure_from_sequences(heavy_chain: str, light_chain: str,
output_file: str = "antibody.pdb") -> str:
"""Generate antibody structure from sequences using ImmuneBuilder."""
sequence = {'H': heavy_chain, 'L': light_chain}
immune_builder(sequence, output_file)
return output_file
def load_existing_structure_files(antibody: Dict, config: Dict) -> Dict[str, str]:
"""
Load existing structure files and validate they exist.
Args:
antibody: Dictionary containing antibody information
config: Configuration dictionary
Returns:
Dictionary with paths to structure files matching format from prepare_structure
Raises:
FileNotFoundError: If required files are missing
"""
logger.info("Loading existing structure files...")
antibody_name = antibody['name']
work_dir = Path(config["paths"]["temp_dir"]).resolve()
# Required file for structure step (processed files are created during MD preprocessing)
pdb_file = work_dir / f"{antibody_name}.pdb"
# Validate required file exists
if not pdb_file.exists():
error_msg = f"Required structure file not found when skipping structure preparation:\n"
error_msg += f" - {pdb_file}\n"
error_msg += f"\nWork directory: {work_dir}"
raise FileNotFoundError(error_msg)
# Extract chain sequences from PDB
try:
parser = PDBParser(QUIET=True)
structure = parser.get_structure("antibody", str(pdb_file))
chains = {chain.id: seq1(''.join(residue.resname for residue in chain))
for chain in structure.get_chains()}
except Exception as e:
logger.warning(f"Failed to extract chain sequences from PDB: {e}")
chains = {}
structure_files = {
"pdb_file": str(pdb_file),
"work_dir": str(work_dir),
"chains": chains
}
logger.info(f"Successfully loaded structure files from {work_dir}")
logger.info(f"Found chains: {list(chains.keys())}")
return structure_files
def prepare_pdb_for_analysis(pdb_file: str, output_dir: str) -> Dict[str, str]:
"""Prepare existing PDB file for analysis."""
antibody = {
"name": Path(pdb_file).stem,
"pdb_file": pdb_file,
"type": "pdb"
}
config = {
"paths": {"temp_dir": output_dir}
}
return prepare_structure(antibody, config)
|