boltz2-python-client 0.3.2__py3-none-any.whl → 0.3.3__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.3.3"
 __author__ = "NVIDIA Corporation"
 __email__ = "bionemo-support@nvidia.com"
 

boltz2_client/a3m_to_csv_converter.py

@@ -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
     )
     

boltz2_client/cli.py

@@ -1534,10 +1534,12 @@ def screen(ctx, target_sequence, compounds_file, target_name, output_dir, no_aff
 @click.option('--pairing-mode', type=click.Choice(['auto', 'taxid', 'uniref']),
               default='auto',
               help='Pairing identifier mode: auto (default, like ColabFold), taxid, or uniref')
+@click.option('--include-unpaired', is_flag=True, default=False,
+              help='Include unpaired sequences in block-diagonal format (maximizes MSA depth)')
 @click.pass_context
 def convert_msa_command(ctx, a3m_files: Tuple[str, ...], chain_ids: str, 
                         output: str, max_pairs: Optional[int], 
-                        pairing_strategy: str, pairing_mode: str):
+                        pairing_strategy: str, pairing_mode: str, include_unpaired: bool):
     """
     Convert ColabFold A3M monomer MSA files to Boltz2 multimer CSV format.
     
@@ -1607,6 +1609,8 @@ def convert_msa_command(ctx, a3m_files: Tuple[str, ...], chain_ids: str,
     print_info(f"Output: {output}")
     print_info(f"Pairing strategy: {pairing_strategy}")
     print_info(f"Pairing mode: {pairing_mode}")
+    if include_unpaired:
+        print_info("Include unpaired: Yes (block-diagonal format)")
     if max_pairs:
         print_info(f"Max pairs: {max_pairs}")
     
@@ -1624,6 +1628,7 @@ def convert_msa_command(ctx, a3m_files: Tuple[str, ...], chain_ids: str,
                 output_path=Path(output),
                 pairing_strategy=pairing_strategy,
                 use_tax_id=use_tax_id,
+                include_unpaired=include_unpaired,
                 max_pairs=max_pairs
             )
             
@@ -1666,6 +1671,8 @@ def convert_msa_command(ctx, a3m_files: Tuple[str, ...], chain_ids: str,
 @click.option('--pairing-mode', type=click.Choice(['auto', 'taxid', 'uniref']),
               default='auto',
               help='Pairing identifier mode: auto (default), taxid, or uniref')
+@click.option('--include-unpaired', is_flag=True, default=False,
+              help='Include unpaired sequences in block-diagonal format (maximizes MSA depth)')
 @click.option('--recycling-steps', type=int, default=3,
               help='Number of recycling steps (default: 3)')
 @click.option('--sampling-steps', type=int, default=200,
@@ -1675,7 +1682,7 @@ def convert_msa_command(ctx, a3m_files: Tuple[str, ...], chain_ids: str,
 @click.pass_context
 def multimer_msa_command(ctx, a3m_files: Tuple[str, ...], chain_ids: str, 
                          output: Optional[str], save_csv: bool, save_all: bool,
-                         max_pairs: Optional[int], pairing_mode: str, 
+                         max_pairs: Optional[int], pairing_mode: str, include_unpaired: bool,
                          recycling_steps: int, sampling_steps: int, 
                          diffusion_samples: int):
     """
@@ -1775,10 +1782,12 @@ def multimer_msa_command(ctx, a3m_files: Tuple[str, ...], chain_ids: str,
                 a3m_files=a3m_file_dict,
                 pairing_strategy='greedy',
                 use_tax_id=use_tax_id,
+                include_unpaired=include_unpaired,
                 max_pairs=max_pairs
             )
             
-            progress.update(task, description=f"✓ Paired {result.num_pairs} sequences")
+            unpaired_msg = " (+ unpaired)" if include_unpaired else ""
+            progress.update(task, description=f"✓ Paired {result.num_pairs} sequences{unpaired_msg}")
         
         print_info(f"Paired sequences: {result.num_pairs}")
         

{boltz2_python_client-0.3.2.dist-info → boltz2_python_client-0.3.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: boltz2-python-client
-Version: 0.3.2
+Version: 0.3.3
 Summary: Python client for Boltz-2 protein structure prediction API with covalent complex and multi-endpoint support
 Author-email: NVIDIA Corporation <bionemo-support@nvidia.com>
 Maintainer-email: NVIDIA Corporation <bionemo-support@nvidia.com>

{boltz2_python_client-0.3.2.dist-info → boltz2_python_client-0.3.3.dist-info}/RECORD

@@ -1,7 +1,7 @@
-boltz2_client/__init__.py,sha256=gQfHL4fPYBPoaGYpCGuy-WAZClzYZsSiRqDNEuEN2IA,5321
+boltz2_client/__init__.py,sha256=YAjmLUxO4KE2a8wVJ4WA1tXFEDRnNalN_yOf6gNb7rk,5321
 boltz2_client/__main__.py,sha256=LamY2ky-DaWf9mGqDZljQwotWICRNeSwq9Y__erB3bQ,140
-boltz2_client/a3m_to_csv_converter.py,sha256=_7b7XHcfu-RttrA40SzsqBFlVw8-KHvEUkryQtgEpkk,60644
-boltz2_client/cli.py,sha256=nnJMv0aK_LmiD0vxbZKqZiG64jci5wlueoPD1BzuI2A,87321
+boltz2_client/a3m_to_csv_converter.py,sha256=4ba0P_7Aqg80Hc4rKZZ_zegpIF-eyjmXT6c_3yo8CdM,64705
+boltz2_client/cli.py,sha256=bo6aYI0pEtoeRPW6zBpeymxum4vxdJJBE5RSps042Rg,87970
 boltz2_client/client.py,sha256=FDz1X0lnB5B0Oi0anHVFYY-VZpVWjyieyX9APmapKOg,57125
 boltz2_client/exceptions.py,sha256=x-5MvZfsN1sLVJjiy8DMNfdmzJDxTosl2QGYKnrhz5k,6725
 boltz2_client/models.py,sha256=XHFTq8h0pUzK5FmOxvBm9aqZNosF75r7CMdEvtV2A50,20171
@@ -12,9 +12,9 @@ boltz2_client/utils.py,sha256=myk4O1ZiJmUvQ-1gQfj_SiDqh3wBINAsLFxH5P5X-cQ,12193
 boltz2_client/virtual_screening.py,sha256=_ykwsjg8to8MAGZGrsaOmiS3Bt9mH5DF_RUxl2Dy5ZU,23857
 boltz2_client/data/__init__.py,sha256=Qy01JhGtCtEJVIMA16fks-ZbL5r0z5JOqHri9bjkWVY,112
 boltz2_client/data/speclist.txt,sha256=5e0no6pTa5HsgErt2fMavdzZs8NZ_iZzBHEfiWKQ1lA,2438338
-boltz2_python_client-0.3.2.dist-info/licenses/LICENSE,sha256=7sYV_pvGtWzybTdOcFXLDPRZL2VolpIl6nt8akVEO94,1111
-boltz2_python_client-0.3.2.dist-info/METADATA,sha256=zdIzySJt-zNVLiEWf-yTeSOtkyEdjjnuc4bqnixpu9M,23920
-boltz2_python_client-0.3.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-boltz2_python_client-0.3.2.dist-info/entry_points.txt,sha256=3whjBtXqltfepvt46DSJNpeUGCZGLRFYnccMEVXMCes,49
-boltz2_python_client-0.3.2.dist-info/top_level.txt,sha256=INM1MYyL_h2nQkTFjZE7AmrgX4ZePTr4uw4n7Sag3ao,14
-boltz2_python_client-0.3.2.dist-info/RECORD,,
+boltz2_python_client-0.3.3.dist-info/licenses/LICENSE,sha256=7sYV_pvGtWzybTdOcFXLDPRZL2VolpIl6nt8akVEO94,1111
+boltz2_python_client-0.3.3.dist-info/METADATA,sha256=dTP4OStXEXbsiZNDbNN7ZA2fuebBC_QyR9g4-E75ZT8,23920
+boltz2_python_client-0.3.3.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+boltz2_python_client-0.3.3.dist-info/entry_points.txt,sha256=3whjBtXqltfepvt46DSJNpeUGCZGLRFYnccMEVXMCes,49
+boltz2_python_client-0.3.3.dist-info/top_level.txt,sha256=INM1MYyL_h2nQkTFjZE7AmrgX4ZePTr4uw4n7Sag3ao,14
+boltz2_python_client-0.3.3.dist-info/RECORD,,

{boltz2_python_client-0.3.2.dist-info → boltz2_python_client-0.3.3.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.9.0)
+Generator: setuptools (80.10.2)
 Root-Is-Purelib: true
 Tag: py3-none-any