boltz2-python-client 0.3.2__py3-none-any.whl → 0.4.0__py3-none-any.whl

boltz2_client/__init__.py

@@ -26,7 +26,7 @@ Example:
     >>> print(f"Confidence: {result.confidence_scores[0]:.3f}")
 """
 
-__version__ = "0.3.2"
+__version__ = "0.4.0"
 __author__ = "NVIDIA Corporation"
 __email__ = "bionemo-support@nvidia.com"
 
@@ -43,6 +43,8 @@ from .models import (
     AlignmentFormat,
     HealthStatus,
     ServiceMetadata,
+    StructuralTemplate,
+    Modification,
 )
 from .models_affinity import AffinityPrediction
 from .exceptions import (
@@ -81,6 +83,11 @@ from .a3m_to_csv_converter import (
     create_paired_msa_per_chain,
     save_prediction_outputs,
     get_prediction_summary,
+    save_pae_matrix,
+    save_pde_matrix,
+    get_pae_summary,
+    convert_cif_to_pdb,
+    convert_pdb_to_cif,
     ConversionResult,
     GreedyPairingStrategy,
     CompletePairingStrategy,
@@ -130,6 +137,8 @@ __all__ = [
     "HealthStatus",
     "ServiceMetadata",
     "AffinityPrediction",
+    "StructuralTemplate",
+    "Modification",
     
     # Exceptions
     "Boltz2Error",
@@ -167,6 +176,11 @@ __all__ = [
     "create_paired_msa_per_chain",
     "save_prediction_outputs",
     "get_prediction_summary",
+    "save_pae_matrix",
+    "save_pde_matrix",
+    "get_pae_summary",
+    "convert_cif_to_pdb",
+    "convert_pdb_to_cif",
     "ConversionResult",
     "GreedyPairingStrategy",
     "CompletePairingStrategy",

boltz2_client/a3m_to_csv_converter.py

@@ -23,7 +23,7 @@ Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
 import re
 import csv
 import logging
-from typing import Dict, List, Optional, Tuple, Set, NamedTuple
+from typing import Dict, List, Optional, Tuple, Set, NamedTuple, Union
 from pathlib import Path
 from dataclasses import dataclass, field
 from collections import defaultdict
@@ -1192,6 +1192,11 @@ class A3MToCSVConverter:
         - Rows with the same 'key' across different chain CSVs are paired
         - This enables the model to identify co-evolved sequences from same organisms
         
+        When include_unpaired=True, also includes:
+        - Unpaired sequences in "block-diagonal" format where only one chain has
+          a sequence and others have gaps. This maximizes MSA depth while still
+          providing pairing information where available.
+        
         Args:
             msas: Dictionary mapping chain IDs to parsed A3MMSA objects
             output_path: Optional path to save the combined CSV file
@@ -1204,35 +1209,79 @@ class A3MToCSVConverter:
         # Find paired sequences
         pairs = self.pairing_strategy.find_pairs(msas)
         
+        # Track which sequences have been paired (to find unpaired ones later)
+        paired_sequence_ids: Dict[str, Set[str]] = {chain_id: set() for chain_id in chain_ids}
+        for pair in pairs:
+            for chain_id, seq in pair.items():
+                paired_sequence_ids[chain_id].add(seq.identifier)
+        
         if self.max_pairs:
             pairs = pairs[:self.max_pairs]
         
         logger.info(f"Found {len(pairs)} paired sequence sets across {len(chain_ids)} chains")
         
-        # Extract query sequences
+        # Extract query sequences and their lengths (for gap filling)
         query_sequences = {}
+        query_lengths = {}
         for chain_id, msa in msas.items():
             query = msa.get_query()
             if query:
-                query_sequences[chain_id] = self._clean_sequence(query.sequence)
+                clean_seq = self._clean_sequence(query.sequence)
+                query_sequences[chain_id] = clean_seq
+                query_lengths[chain_id] = len(clean_seq)
+        
+        # Collect unpaired sequences if requested
+        unpaired_sequences: Dict[str, List[A3MSequence]] = {chain_id: [] for chain_id in chain_ids}
+        num_unpaired = 0
+        
+        if self.include_unpaired:
+            for chain_id, msa in msas.items():
+                for seq in msa.sequences:
+                    if seq.is_query:
+                        continue
+                    if seq.identifier not in paired_sequence_ids[chain_id]:
+                        unpaired_sequences[chain_id].append(seq)
+                        num_unpaired += 1
+            
+            logger.info(f"Found {num_unpaired} unpaired sequences to include in block-diagonal format")
+        
+        # Calculate total rows: paired + unpaired (block-diagonal)
+        current_key = len(pairs)  # Start unpaired keys after paired ones
         
         # Build per-chain CSVs (for Boltz2 NIM API)
         csv_per_chain: Dict[str, str] = {}
         for chain_id in chain_ids:
             chain_lines = ["key,sequence"]
+            
+            # Add paired sequences
             for idx, pair in enumerate(pairs, start=1):
                 if chain_id in pair:
                     seq = self._clean_sequence(pair[chain_id].sequence)
                     chain_lines.append(f"{idx},{seq}")
                 else:
                     # Use gaps if sequence not found for this chain
-                    ref_msa = msas[chain_id]
-                    query = ref_msa.get_query()
-                    if query:
-                        gap_seq = '-' * len(self._clean_sequence(query.sequence))
-                        chain_lines.append(f"{idx},{gap_seq}")
+                    gap_seq = '-' * query_lengths.get(chain_id, 0)
+                    chain_lines.append(f"{idx},{gap_seq}")
+            
             csv_per_chain[chain_id] = '\n'.join(chain_lines)
         
+        # Add unpaired sequences in block-diagonal format
+        if self.include_unpaired:
+            unpaired_key = current_key + 1
+            for source_chain_id in chain_ids:
+                for unpaired_seq in unpaired_sequences[source_chain_id]:
+                    # Add this unpaired sequence to all chain CSVs
+                    for chain_id in chain_ids:
+                        if chain_id == source_chain_id:
+                            # This chain has the sequence
+                            seq = self._clean_sequence(unpaired_seq.sequence)
+                            csv_per_chain[chain_id] += f"\n{unpaired_key},{seq}"
+                        else:
+                            # Other chains get gaps
+                            gap_seq = '-' * query_lengths.get(chain_id, 0)
+                            csv_per_chain[chain_id] += f"\n{unpaired_key},{gap_seq}"
+                    unpaired_key += 1
+        
         # Build combined CSV (for reference/documentation)
         csv_lines = ["key,sequence"]
         for idx, pair in enumerate(pairs, start=1):
@@ -1242,15 +1291,31 @@ class A3MToCSVConverter:
                     seq = self._clean_sequence(pair[chain_id].sequence)
                     sequences.append(seq)
                 else:
-                    ref_msa = msas[chain_id]
-                    query = ref_msa.get_query()
-                    if query:
-                        sequences.append('-' * len(self._clean_sequence(query.sequence)))
+                    sequences.append('-' * query_lengths.get(chain_id, 0))
             concatenated = ':'.join(sequences)
             csv_lines.append(f"{idx},{concatenated}")
         
+        # Add unpaired sequences to combined CSV in block-diagonal format
+        if self.include_unpaired:
+            unpaired_key = current_key + 1
+            for source_chain_id in chain_ids:
+                for unpaired_seq in unpaired_sequences[source_chain_id]:
+                    sequences = []
+                    for chain_id in chain_ids:
+                        if chain_id == source_chain_id:
+                            seq = self._clean_sequence(unpaired_seq.sequence)
+                            sequences.append(seq)
+                        else:
+                            sequences.append('-' * query_lengths.get(chain_id, 0))
+                    concatenated = ':'.join(sequences)
+                    csv_lines.append(f"{unpaired_key},{concatenated}")
+                    unpaired_key += 1
+        
         csv_content = '\n'.join(csv_lines)
         
+        total_rows = len(pairs) + (num_unpaired if self.include_unpaired else 0)
+        logger.info(f"Total MSA rows: {total_rows} ({len(pairs)} paired + {num_unpaired if self.include_unpaired else 0} unpaired)")
+        
         # Save files if path provided
         output_paths_per_chain = None
         if output_path:
@@ -1355,6 +1420,7 @@ def convert_a3m_to_multimer_csv(
     output_path: Optional[Path] = None,
     pairing_strategy: str = "greedy",
     use_tax_id: Optional[bool] = None,  # None = auto-detect (ColabFold default)
+    include_unpaired: bool = False,
     max_pairs: Optional[int] = None
 ) -> ConversionResult:
     """
@@ -1374,6 +1440,10 @@ def convert_a3m_to_multimer_csv(
                              This is how ColabFold behaves.
             - True: Force TaxID pairing (requires OX= or species codes in headers)
             - False: Force UniRef/organism ID pairing (works with standard ColabFold output)
+        include_unpaired: If True, include sequences without cross-chain matches in 
+                         "block-diagonal" format (one chain has sequence, others have gaps).
+                         This maximizes MSA depth while still providing pairing where available.
+                         Default: False (only paired sequences are included)
         max_pairs: Maximum number of pairs to include
         
     Returns:
@@ -1388,6 +1458,12 @@ def convert_a3m_to_multimer_csv(
         ...     # use_tax_id=None means auto-detect (default)
         ... )
         >>> print(f"Created {result.num_pairs} paired sequences")
+        
+        >>> # Include unpaired sequences for maximum MSA depth
+        >>> result = convert_a3m_to_multimer_csv(
+        ...     a3m_files={'A': Path('chain_A.a3m'), 'B': Path('chain_B.a3m')},
+        ...     include_unpaired=True  # Block-diagonal format
+        ... )
     
     ColabFold Compatibility:
         Standard ColabFold A3M files use UniRef100 cluster IDs without TaxID information.
@@ -1417,6 +1493,7 @@ def convert_a3m_to_multimer_csv(
     
     converter = A3MToCSVConverter(
         pairing_strategy=strategy,
+        include_unpaired=include_unpaired,
         max_pairs=max_pairs
     )
     
@@ -1659,3 +1736,277 @@ def get_prediction_summary(response) -> Dict:
     
     return summary
 
+
+def save_pae_matrix(
+    pae: List[List[List[float]]],
+    output_path: Union[str, Path],
+    format: str = "json"
+) -> Path:
+    """
+    Save PAE (Predicted Aligned Error) matrix to file.
+    
+    Args:
+        pae: PAE matrix from PredictionResponse.pae
+             Shape: [num_models, num_residues, num_residues]
+        output_path: Path to save the file
+        format: Output format - "json" or "npy" (numpy)
+        
+    Returns:
+        Path to saved file
+        
+    Example:
+        >>> if response.pae:
+        ...     save_pae_matrix(response.pae, "structure.pae.json")
+    """
+    import json
+    
+    output_path = Path(output_path)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    
+    if format == "json":
+        with open(output_path, "w") as f:
+            json.dump({"pae": pae}, f, indent=2)
+    elif format == "npy":
+        try:
+            import numpy as np
+            np.save(output_path, np.array(pae))
+        except ImportError:
+            raise ImportError("numpy is required for .npy format. Install with: pip install numpy")
+    else:
+        raise ValueError(f"Unsupported format: {format}. Use 'json' or 'npy'")
+    
+    return output_path
+
+
+def save_pde_matrix(
+    pde: List[List[List[float]]],
+    output_path: Union[str, Path],
+    format: str = "json"
+) -> Path:
+    """
+    Save PDE (Predicted Distance Error) matrix to file.
+    
+    Args:
+        pde: PDE matrix from PredictionResponse.pde
+             Shape: [num_models, num_residues, num_residues]
+        output_path: Path to save the file
+        format: Output format - "json" or "npy" (numpy)
+        
+    Returns:
+        Path to saved file
+        
+    Example:
+        >>> if response.pde:
+        ...     save_pde_matrix(response.pde, "structure.pde.json")
+    """
+    import json
+    
+    output_path = Path(output_path)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    
+    if format == "json":
+        with open(output_path, "w") as f:
+            json.dump({"pde": pde}, f, indent=2)
+    elif format == "npy":
+        try:
+            import numpy as np
+            np.save(output_path, np.array(pde))
+        except ImportError:
+            raise ImportError("numpy is required for .npy format. Install with: pip install numpy")
+    else:
+        raise ValueError(f"Unsupported format: {format}. Use 'json' or 'npy'")
+    
+    return output_path
+
+
+def get_pae_summary(pae: List[List[List[float]]]) -> Dict:
+    """
+    Get summary statistics for PAE matrix.
+    
+    Args:
+        pae: PAE matrix from PredictionResponse.pae
+        
+    Returns:
+        Dictionary with shape, mean, min, max values
+        
+    Example:
+        >>> summary = get_pae_summary(response.pae)
+        >>> print(f"Mean PAE: {summary['mean']:.2f} Å")
+    """
+    if not pae:
+        return {}
+    
+    # Flatten to get statistics
+    all_values = []
+    for model_pae in pae:
+        for row in model_pae:
+            all_values.extend(row)
+    
+    return {
+        "shape": f"{len(pae)}x{len(pae[0])}x{len(pae[0][0])}",
+        "num_models": len(pae),
+        "num_residues": len(pae[0]),
+        "mean": sum(all_values) / len(all_values),
+        "min": min(all_values),
+        "max": max(all_values),
+    }
+
+
+def convert_cif_to_pdb(
+    cif_input: Union[str, Path],
+    pdb_output: Union[str, Path],
+    structure_id: str = "PRED"
+) -> Path:
+    """
+    Convert mmCIF file to PDB format.
+    
+    Uses gemmi (preferred, faster) or BioPython as fallback.
+    
+    Args:
+        cif_input: Path to input CIF file or CIF content as string
+        pdb_output: Path for output PDB file
+        structure_id: Structure ID to use in PDB file (default: "PRED")
+        
+    Returns:
+        Path to the saved PDB file
+        
+    Example:
+        >>> convert_cif_to_pdb("structure.cif", "structure.pdb")
+        >>> # Or from string content
+        >>> convert_cif_to_pdb(cif_content, "structure.pdb")
+        
+    Note:
+        Requires either gemmi or BioPython to be installed:
+        - pip install gemmi  (recommended, faster)
+        - pip install biopython
+    """
+    import tempfile
+    
+    pdb_output = Path(pdb_output)
+    pdb_output.parent.mkdir(parents=True, exist_ok=True)
+    
+    # Check if input is a file path or string content
+    cif_path = Path(cif_input) if isinstance(cif_input, (str, Path)) and Path(cif_input).exists() else None
+    
+    # If input is string content, write to temp file
+    if cif_path is None:
+        with tempfile.NamedTemporaryFile(mode='w', suffix='.cif', delete=False) as f:
+            f.write(str(cif_input))
+            cif_path = Path(f.name)
+        cleanup_temp = True
+    else:
+        cif_path = Path(cif_input)
+        cleanup_temp = False
+    
+    try:
+        # Try gemmi first (faster)
+        try:
+            import gemmi
+            structure = gemmi.read_structure(str(cif_path))
+            structure.write_pdb(str(pdb_output))
+            return pdb_output
+        except ImportError:
+            pass
+        
+        # Fallback to BioPython
+        try:
+            from Bio.PDB import MMCIFParser, PDBIO
+            import warnings
+            
+            # Suppress BioPython warnings about incomplete structures
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore")
+                
+                parser = MMCIFParser(QUIET=True)
+                structure = parser.get_structure(structure_id, str(cif_path))
+                
+                io = PDBIO()
+                io.set_structure(structure)
+                io.save(str(pdb_output))
+                
+            return pdb_output
+        except ImportError:
+            raise ImportError(
+                "CIF to PDB conversion requires either gemmi or BioPython. "
+                "Install with: pip install gemmi  (recommended) or pip install biopython"
+            )
+    finally:
+        # Clean up temp file if created
+        if cleanup_temp and cif_path.exists():
+            cif_path.unlink()
+
+
+def convert_pdb_to_cif(
+    pdb_input: Union[str, Path],
+    cif_output: Union[str, Path],
+    structure_id: str = "PRED"
+) -> Path:
+    """
+    Convert PDB file to mmCIF format.
+    
+    Uses gemmi (preferred, faster) or BioPython as fallback.
+    
+    Args:
+        pdb_input: Path to input PDB file or PDB content as string
+        cif_output: Path for output CIF file
+        structure_id: Structure ID to use (default: "PRED")
+        
+    Returns:
+        Path to the saved CIF file
+        
+    Example:
+        >>> convert_pdb_to_cif("structure.pdb", "structure.cif")
+    """
+    import tempfile
+    
+    cif_output = Path(cif_output)
+    cif_output.parent.mkdir(parents=True, exist_ok=True)
+    
+    # Check if input is a file path or string content
+    pdb_path = Path(pdb_input) if isinstance(pdb_input, (str, Path)) and Path(pdb_input).exists() else None
+    
+    # If input is string content, write to temp file
+    if pdb_path is None:
+        with tempfile.NamedTemporaryFile(mode='w', suffix='.pdb', delete=False) as f:
+            f.write(str(pdb_input))
+            pdb_path = Path(f.name)
+        cleanup_temp = True
+    else:
+        pdb_path = Path(pdb_input)
+        cleanup_temp = False
+    
+    try:
+        # Try gemmi first (faster)
+        try:
+            import gemmi
+            structure = gemmi.read_structure(str(pdb_path))
+            structure.write_minimal_cif(str(cif_output))
+            return cif_output
+        except ImportError:
+            pass
+        
+        # Fallback to BioPython
+        try:
+            from Bio.PDB import PDBParser, MMCIFIO
+            import warnings
+            
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore")
+                
+                parser = PDBParser(QUIET=True)
+                structure = parser.get_structure(structure_id, str(pdb_path))
+                
+                io = MMCIFIO()
+                io.set_structure(structure)
+                io.save(str(cif_output))
+                
+            return cif_output
+        except ImportError:
+            raise ImportError(
+                "PDB to CIF conversion requires either gemmi or BioPython. "
+                "Install with: pip install gemmi  (recommended) or pip install biopython"
+            )
+    finally:
+        if cleanup_temp and pdb_path.exists():
+            pdb_path.unlink()
+