boltz2-python-client 0.3.2__tar.gz → 0.3.3__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (92) hide show
  1. {boltz2_python_client-0.3.2/boltz2_python_client.egg-info → boltz2_python_client-0.3.3}/PKG-INFO +1 -1
  2. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_client/__init__.py +1 -1
  3. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_client/a3m_to_csv_converter.py +88 -11
  4. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_client/cli.py +12 -3
  5. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3/boltz2_python_client.egg-info}/PKG-INFO +1 -1
  6. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_python_client.egg-info/SOURCES.txt +1 -0
  7. boltz2_python_client-0.3.3/tests/test_a3m_to_csv_converter.py +579 -0
  8. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/ASYNC_GUIDE.md +0 -0
  9. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/COVALENT_COMPLEX_GUIDE.md +0 -0
  10. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/LICENSE +0 -0
  11. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/MANIFEST.in +0 -0
  12. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/PARAMETERS.md +0 -0
  13. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/README.md +0 -0
  14. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/YAML_GUIDE.md +0 -0
  15. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_client/__main__.py +0 -0
  16. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_client/client.py +0 -0
  17. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_client/data/__init__.py +0 -0
  18. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_client/data/speclist.txt +0 -0
  19. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_client/exceptions.py +0 -0
  20. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_client/models.py +0 -0
  21. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_client/models_affinity.py +0 -0
  22. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_client/msa_search.py +0 -0
  23. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_client/multi_endpoint_client.py +0 -0
  24. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_client/utils.py +0 -0
  25. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_client/virtual_screening.py +0 -0
  26. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_python_client.egg-info/dependency_links.txt +0 -0
  27. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_python_client.egg-info/entry_points.txt +0 -0
  28. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_python_client.egg-info/requires.txt +0 -0
  29. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_python_client.egg-info/top_level.txt +0 -0
  30. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/01_basic_protein_folding.py +0 -0
  31. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/02_protein_structure_prediction_with_msa.py +0 -0
  32. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/03_protein_ligand_complex.py +0 -0
  33. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/04_covalent_bonding.py +0 -0
  34. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/05_dna_protein_complex.py +0 -0
  35. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/06_yaml_configurations.py +0 -0
  36. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/07_advanced_parameters.py +0 -0
  37. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/08_affinity_prediction_simple.py +0 -0
  38. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/09_virtual_screening.py +0 -0
  39. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/10_msa_search_integration.py +0 -0
  40. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/11_msa_search_large_protein.py +0 -0
  41. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/12_msa_affinity_prediction.py +0 -0
  42. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/15_a3m_to_multimer_csv.py +0 -0
  43. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/barnase_barstar_with_msa.py +0 -0
  44. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/cdk4_msa_affinity/cdk4_direct.a3m +0 -0
  45. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/cdk4_msa_affinity/cdk4_msa.a3m +0 -0
  46. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/cdk4_msa_affinity/cdk4_simple.a3m +0 -0
  47. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/cdk4_msa_affinity_example.py +0 -0
  48. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/comprehensive_multi_endpoint_demo.py +0 -0
  49. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/debug_msa_conversion.py +0 -0
  50. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/dimer_examples.py +0 -0
  51. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/fix_cdk4_msa_search.py +0 -0
  52. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/msa-kras-g12c_combined.a3m +0 -0
  53. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/msa_search_simple_demo.py +0 -0
  54. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/multi_endpoint_screening.py +0 -0
  55. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/multi_protein_complex.yaml +0 -0
  56. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/multimer_file_output_example.py +0 -0
  57. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/multimer_msa_example.py +0 -0
  58. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/protein_ligand.yaml +0 -0
  59. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/quick_barnase_barstar.py +0 -0
  60. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/quick_barnase_barstar_msa.py +0 -0
  61. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/quick_gcn4_homodimer.py +0 -0
  62. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/quick_heterodimer_example.py +0 -0
  63. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/quick_msa_test.py +0 -0
  64. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/sars_cov2_mpro_nirmatrelvir.yaml +0 -0
  65. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/simple_dimer_examples.py +0 -0
  66. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/test_boltz2_server.py +0 -0
  67. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/test_cdk4_msa_fixed.py +0 -0
  68. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/test_cdk4_msa_simple.py +0 -0
  69. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/test_msa_a3m_simple.py +0 -0
  70. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/test_msa_endpoint.py +0 -0
  71. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/test_notebook_fixes.py +0 -0
  72. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/licenses/PyYAML-LICENSE +0 -0
  73. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/licenses/README.md +0 -0
  74. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/licenses/aiofiles-LICENSE +0 -0
  75. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/licenses/click-LICENSE +0 -0
  76. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/licenses/httpx-LICENSE +0 -0
  77. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/licenses/py3Dmol-LICENSE +0 -0
  78. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/licenses/pydantic-LICENSE +0 -0
  79. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/licenses/rich-LICENSE +0 -0
  80. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/licenses/typing-extensions-LICENSE +0 -0
  81. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/pyproject.toml +0 -0
  82. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/setup.cfg +0 -0
  83. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/tests/__init__.py +0 -0
  84. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/tests/conftest.py +0 -0
  85. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/tests/test_basic.py +0 -0
  86. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/tests/test_cli_multi_endpoint.py +0 -0
  87. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/tests/test_comprehensive_stress.py +0 -0
  88. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/tests/test_examples_syntax.py +0 -0
  89. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/tests/test_integration_scenarios.py +0 -0
  90. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/tests/test_msa_search.py +0 -0
  91. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/tests/test_multi_endpoint_functionality.py +0 -0
  92. {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/tests/test_real_endpoints.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: boltz2-python-client
3
- Version: 0.3.2
3
+ Version: 0.3.3
4
4
  Summary: Python client for Boltz-2 protein structure prediction API with covalent complex and multi-endpoint support
5
5
  Author-email: NVIDIA Corporation <bionemo-support@nvidia.com>
6
6
  Maintainer-email: NVIDIA Corporation <bionemo-support@nvidia.com>
@@ -26,7 +26,7 @@ Example:
26
26
  >>> print(f"Confidence: {result.confidence_scores[0]:.3f}")
27
27
  """
28
28
 
29
- __version__ = "0.3.2"
29
+ __version__ = "0.3.3"
30
30
  __author__ = "NVIDIA Corporation"
31
31
  __email__ = "bionemo-support@nvidia.com"
32
32
 
@@ -1192,6 +1192,11 @@ class A3MToCSVConverter:
1192
1192
  - Rows with the same 'key' across different chain CSVs are paired
1193
1193
  - This enables the model to identify co-evolved sequences from same organisms
1194
1194
 
1195
+ When include_unpaired=True, also includes:
1196
+ - Unpaired sequences in "block-diagonal" format where only one chain has
1197
+ a sequence and others have gaps. This maximizes MSA depth while still
1198
+ providing pairing information where available.
1199
+
1195
1200
  Args:
1196
1201
  msas: Dictionary mapping chain IDs to parsed A3MMSA objects
1197
1202
  output_path: Optional path to save the combined CSV file
@@ -1204,35 +1209,79 @@ class A3MToCSVConverter:
1204
1209
  # Find paired sequences
1205
1210
  pairs = self.pairing_strategy.find_pairs(msas)
1206
1211
 
1212
+ # Track which sequences have been paired (to find unpaired ones later)
1213
+ paired_sequence_ids: Dict[str, Set[str]] = {chain_id: set() for chain_id in chain_ids}
1214
+ for pair in pairs:
1215
+ for chain_id, seq in pair.items():
1216
+ paired_sequence_ids[chain_id].add(seq.identifier)
1217
+
1207
1218
  if self.max_pairs:
1208
1219
  pairs = pairs[:self.max_pairs]
1209
1220
 
1210
1221
  logger.info(f"Found {len(pairs)} paired sequence sets across {len(chain_ids)} chains")
1211
1222
 
1212
- # Extract query sequences
1223
+ # Extract query sequences and their lengths (for gap filling)
1213
1224
  query_sequences = {}
1225
+ query_lengths = {}
1214
1226
  for chain_id, msa in msas.items():
1215
1227
  query = msa.get_query()
1216
1228
  if query:
1217
- query_sequences[chain_id] = self._clean_sequence(query.sequence)
1229
+ clean_seq = self._clean_sequence(query.sequence)
1230
+ query_sequences[chain_id] = clean_seq
1231
+ query_lengths[chain_id] = len(clean_seq)
1232
+
1233
+ # Collect unpaired sequences if requested
1234
+ unpaired_sequences: Dict[str, List[A3MSequence]] = {chain_id: [] for chain_id in chain_ids}
1235
+ num_unpaired = 0
1236
+
1237
+ if self.include_unpaired:
1238
+ for chain_id, msa in msas.items():
1239
+ for seq in msa.sequences:
1240
+ if seq.is_query:
1241
+ continue
1242
+ if seq.identifier not in paired_sequence_ids[chain_id]:
1243
+ unpaired_sequences[chain_id].append(seq)
1244
+ num_unpaired += 1
1245
+
1246
+ logger.info(f"Found {num_unpaired} unpaired sequences to include in block-diagonal format")
1247
+
1248
+ # Calculate total rows: paired + unpaired (block-diagonal)
1249
+ current_key = len(pairs) # Start unpaired keys after paired ones
1218
1250
 
1219
1251
  # Build per-chain CSVs (for Boltz2 NIM API)
1220
1252
  csv_per_chain: Dict[str, str] = {}
1221
1253
  for chain_id in chain_ids:
1222
1254
  chain_lines = ["key,sequence"]
1255
+
1256
+ # Add paired sequences
1223
1257
  for idx, pair in enumerate(pairs, start=1):
1224
1258
  if chain_id in pair:
1225
1259
  seq = self._clean_sequence(pair[chain_id].sequence)
1226
1260
  chain_lines.append(f"{idx},{seq}")
1227
1261
  else:
1228
1262
  # Use gaps if sequence not found for this chain
1229
- ref_msa = msas[chain_id]
1230
- query = ref_msa.get_query()
1231
- if query:
1232
- gap_seq = '-' * len(self._clean_sequence(query.sequence))
1233
- chain_lines.append(f"{idx},{gap_seq}")
1263
+ gap_seq = '-' * query_lengths.get(chain_id, 0)
1264
+ chain_lines.append(f"{idx},{gap_seq}")
1265
+
1234
1266
  csv_per_chain[chain_id] = '\n'.join(chain_lines)
1235
1267
 
1268
+ # Add unpaired sequences in block-diagonal format
1269
+ if self.include_unpaired:
1270
+ unpaired_key = current_key + 1
1271
+ for source_chain_id in chain_ids:
1272
+ for unpaired_seq in unpaired_sequences[source_chain_id]:
1273
+ # Add this unpaired sequence to all chain CSVs
1274
+ for chain_id in chain_ids:
1275
+ if chain_id == source_chain_id:
1276
+ # This chain has the sequence
1277
+ seq = self._clean_sequence(unpaired_seq.sequence)
1278
+ csv_per_chain[chain_id] += f"\n{unpaired_key},{seq}"
1279
+ else:
1280
+ # Other chains get gaps
1281
+ gap_seq = '-' * query_lengths.get(chain_id, 0)
1282
+ csv_per_chain[chain_id] += f"\n{unpaired_key},{gap_seq}"
1283
+ unpaired_key += 1
1284
+
1236
1285
  # Build combined CSV (for reference/documentation)
1237
1286
  csv_lines = ["key,sequence"]
1238
1287
  for idx, pair in enumerate(pairs, start=1):
@@ -1242,15 +1291,31 @@ class A3MToCSVConverter:
1242
1291
  seq = self._clean_sequence(pair[chain_id].sequence)
1243
1292
  sequences.append(seq)
1244
1293
  else:
1245
- ref_msa = msas[chain_id]
1246
- query = ref_msa.get_query()
1247
- if query:
1248
- sequences.append('-' * len(self._clean_sequence(query.sequence)))
1294
+ sequences.append('-' * query_lengths.get(chain_id, 0))
1249
1295
  concatenated = ':'.join(sequences)
1250
1296
  csv_lines.append(f"{idx},{concatenated}")
1251
1297
 
1298
+ # Add unpaired sequences to combined CSV in block-diagonal format
1299
+ if self.include_unpaired:
1300
+ unpaired_key = current_key + 1
1301
+ for source_chain_id in chain_ids:
1302
+ for unpaired_seq in unpaired_sequences[source_chain_id]:
1303
+ sequences = []
1304
+ for chain_id in chain_ids:
1305
+ if chain_id == source_chain_id:
1306
+ seq = self._clean_sequence(unpaired_seq.sequence)
1307
+ sequences.append(seq)
1308
+ else:
1309
+ sequences.append('-' * query_lengths.get(chain_id, 0))
1310
+ concatenated = ':'.join(sequences)
1311
+ csv_lines.append(f"{unpaired_key},{concatenated}")
1312
+ unpaired_key += 1
1313
+
1252
1314
  csv_content = '\n'.join(csv_lines)
1253
1315
 
1316
+ total_rows = len(pairs) + (num_unpaired if self.include_unpaired else 0)
1317
+ logger.info(f"Total MSA rows: {total_rows} ({len(pairs)} paired + {num_unpaired if self.include_unpaired else 0} unpaired)")
1318
+
1254
1319
  # Save files if path provided
1255
1320
  output_paths_per_chain = None
1256
1321
  if output_path:
@@ -1355,6 +1420,7 @@ def convert_a3m_to_multimer_csv(
1355
1420
  output_path: Optional[Path] = None,
1356
1421
  pairing_strategy: str = "greedy",
1357
1422
  use_tax_id: Optional[bool] = None, # None = auto-detect (ColabFold default)
1423
+ include_unpaired: bool = False,
1358
1424
  max_pairs: Optional[int] = None
1359
1425
  ) -> ConversionResult:
1360
1426
  """
@@ -1374,6 +1440,10 @@ def convert_a3m_to_multimer_csv(
1374
1440
  This is how ColabFold behaves.
1375
1441
  - True: Force TaxID pairing (requires OX= or species codes in headers)
1376
1442
  - False: Force UniRef/organism ID pairing (works with standard ColabFold output)
1443
+ include_unpaired: If True, include sequences without cross-chain matches in
1444
+ "block-diagonal" format (one chain has sequence, others have gaps).
1445
+ This maximizes MSA depth while still providing pairing where available.
1446
+ Default: False (only paired sequences are included)
1377
1447
  max_pairs: Maximum number of pairs to include
1378
1448
 
1379
1449
  Returns:
@@ -1388,6 +1458,12 @@ def convert_a3m_to_multimer_csv(
1388
1458
  ... # use_tax_id=None means auto-detect (default)
1389
1459
  ... )
1390
1460
  >>> print(f"Created {result.num_pairs} paired sequences")
1461
+
1462
+ >>> # Include unpaired sequences for maximum MSA depth
1463
+ >>> result = convert_a3m_to_multimer_csv(
1464
+ ... a3m_files={'A': Path('chain_A.a3m'), 'B': Path('chain_B.a3m')},
1465
+ ... include_unpaired=True # Block-diagonal format
1466
+ ... )
1391
1467
 
1392
1468
  ColabFold Compatibility:
1393
1469
  Standard ColabFold A3M files use UniRef100 cluster IDs without TaxID information.
@@ -1417,6 +1493,7 @@ def convert_a3m_to_multimer_csv(
1417
1493
 
1418
1494
  converter = A3MToCSVConverter(
1419
1495
  pairing_strategy=strategy,
1496
+ include_unpaired=include_unpaired,
1420
1497
  max_pairs=max_pairs
1421
1498
  )
1422
1499
 
@@ -1534,10 +1534,12 @@ def screen(ctx, target_sequence, compounds_file, target_name, output_dir, no_aff
1534
1534
  @click.option('--pairing-mode', type=click.Choice(['auto', 'taxid', 'uniref']),
1535
1535
  default='auto',
1536
1536
  help='Pairing identifier mode: auto (default, like ColabFold), taxid, or uniref')
1537
+ @click.option('--include-unpaired', is_flag=True, default=False,
1538
+ help='Include unpaired sequences in block-diagonal format (maximizes MSA depth)')
1537
1539
  @click.pass_context
1538
1540
  def convert_msa_command(ctx, a3m_files: Tuple[str, ...], chain_ids: str,
1539
1541
  output: str, max_pairs: Optional[int],
1540
- pairing_strategy: str, pairing_mode: str):
1542
+ pairing_strategy: str, pairing_mode: str, include_unpaired: bool):
1541
1543
  """
1542
1544
  Convert ColabFold A3M monomer MSA files to Boltz2 multimer CSV format.
1543
1545
 
@@ -1607,6 +1609,8 @@ def convert_msa_command(ctx, a3m_files: Tuple[str, ...], chain_ids: str,
1607
1609
  print_info(f"Output: {output}")
1608
1610
  print_info(f"Pairing strategy: {pairing_strategy}")
1609
1611
  print_info(f"Pairing mode: {pairing_mode}")
1612
+ if include_unpaired:
1613
+ print_info("Include unpaired: Yes (block-diagonal format)")
1610
1614
  if max_pairs:
1611
1615
  print_info(f"Max pairs: {max_pairs}")
1612
1616
 
@@ -1624,6 +1628,7 @@ def convert_msa_command(ctx, a3m_files: Tuple[str, ...], chain_ids: str,
1624
1628
  output_path=Path(output),
1625
1629
  pairing_strategy=pairing_strategy,
1626
1630
  use_tax_id=use_tax_id,
1631
+ include_unpaired=include_unpaired,
1627
1632
  max_pairs=max_pairs
1628
1633
  )
1629
1634
 
@@ -1666,6 +1671,8 @@ def convert_msa_command(ctx, a3m_files: Tuple[str, ...], chain_ids: str,
1666
1671
  @click.option('--pairing-mode', type=click.Choice(['auto', 'taxid', 'uniref']),
1667
1672
  default='auto',
1668
1673
  help='Pairing identifier mode: auto (default), taxid, or uniref')
1674
+ @click.option('--include-unpaired', is_flag=True, default=False,
1675
+ help='Include unpaired sequences in block-diagonal format (maximizes MSA depth)')
1669
1676
  @click.option('--recycling-steps', type=int, default=3,
1670
1677
  help='Number of recycling steps (default: 3)')
1671
1678
  @click.option('--sampling-steps', type=int, default=200,
@@ -1675,7 +1682,7 @@ def convert_msa_command(ctx, a3m_files: Tuple[str, ...], chain_ids: str,
1675
1682
  @click.pass_context
1676
1683
  def multimer_msa_command(ctx, a3m_files: Tuple[str, ...], chain_ids: str,
1677
1684
  output: Optional[str], save_csv: bool, save_all: bool,
1678
- max_pairs: Optional[int], pairing_mode: str,
1685
+ max_pairs: Optional[int], pairing_mode: str, include_unpaired: bool,
1679
1686
  recycling_steps: int, sampling_steps: int,
1680
1687
  diffusion_samples: int):
1681
1688
  """
@@ -1775,10 +1782,12 @@ def multimer_msa_command(ctx, a3m_files: Tuple[str, ...], chain_ids: str,
1775
1782
  a3m_files=a3m_file_dict,
1776
1783
  pairing_strategy='greedy',
1777
1784
  use_tax_id=use_tax_id,
1785
+ include_unpaired=include_unpaired,
1778
1786
  max_pairs=max_pairs
1779
1787
  )
1780
1788
 
1781
- progress.update(task, description=f" Paired {result.num_pairs} sequences")
1789
+ unpaired_msg = " (+ unpaired)" if include_unpaired else ""
1790
+ progress.update(task, description=f"✓ Paired {result.num_pairs} sequences{unpaired_msg}")
1782
1791
 
1783
1792
  print_info(f"Paired sequences: {result.num_pairs}")
1784
1793
 
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: boltz2-python-client
3
- Version: 0.3.2
3
+ Version: 0.3.3
4
4
  Summary: Python client for Boltz-2 protein structure prediction API with covalent complex and multi-endpoint support
5
5
  Author-email: NVIDIA Corporation <bionemo-support@nvidia.com>
6
6
  Maintainer-email: NVIDIA Corporation <bionemo-support@nvidia.com>
@@ -79,6 +79,7 @@ licenses/rich-LICENSE
79
79
  licenses/typing-extensions-LICENSE
80
80
  tests/__init__.py
81
81
  tests/conftest.py
82
+ tests/test_a3m_to_csv_converter.py
82
83
  tests/test_basic.py
83
84
  tests/test_cli_multi_endpoint.py
84
85
  tests/test_comprehensive_stress.py
@@ -0,0 +1,579 @@
1
+ #!/usr/bin/env python3
2
+ """
3
+ Tests for A3M to CSV Multimer MSA Converter.
4
+
5
+ Tests cover:
6
+ 1. A3M parsing (headers, sequences, various formats)
7
+ 2. Pairing strategies (greedy, complete)
8
+ 3. TaxID vs UniRef ID pairing
9
+ 4. CSV generation
10
+ 5. Block-diagonal unpaired sequences (new feature)
11
+ """
12
+
13
+ import pytest
14
+ from pathlib import Path
15
+ from boltz2_client.a3m_to_csv_converter import (
16
+ A3MParser,
17
+ A3MMSA,
18
+ A3MSequence,
19
+ A3MToCSVConverter,
20
+ GreedyPairingStrategy,
21
+ CompletePairingStrategy,
22
+ TaxonomyPairingStrategy,
23
+ SpeciesMapper,
24
+ convert_a3m_to_multimer_csv,
25
+ create_paired_msa_per_chain,
26
+ ConversionResult,
27
+ SPECIES_TO_TAXID,
28
+ )
29
+
30
+
31
+ # Test data
32
+ CHAIN_A_A3M = """>Query|-|Query
33
+ MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG
34
+ >UniRef100_A0A2N5EEG3 340 0.994
35
+ MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG
36
+ >UniRef100_UPI000755DEB5 340 0.994
37
+ MKTVRQERLKSIVRILERSKDPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG
38
+ >UniRef100_Q54321 320 0.95
39
+ MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLG
40
+ >UniRef100_B67890 315 0.90
41
+ MKTVRQERLKSIVRILERSKEPVSGAQLAEDLSVSRQVIVQDIAYLRSLG
42
+ """
43
+
44
+ CHAIN_B_A3M = """>Query|-|Query
45
+ MVTPEGNVSLVDESLLVGVTDEDRAVRSAHQFYERLIGLWAPAVMEAAHELGVFAALAEAPAD
46
+ >UniRef100_A0A2N5EEG3 340 0.994
47
+ MVTPEGNVSLVDESLLVGVTDEDRAVRSAHQFYERLIGLWAPAVMEAAHELGVFAALAEAPAD
48
+ >UniRef100_UPI000755DEB5 340 0.994
49
+ MVTPEGNVSLVDESLLVGVTDEDRAVRSAHQFYERLIGLWAPAVMEAAHELGVFAALAERPAD
50
+ >UniRef100_X99999 300 0.85
51
+ MVTPEGNVSLVDESLLVGVTDEDRAVRSAHQFYERLIGLWAPAVMEAAHELGVFAALAE
52
+ >UniRef100_B67890 315 0.90
53
+ MVTPEGNVSLVDESLLVGVTDEDRAVRSAHQFYERLIGLWAPAVMEAAHELGVFAALAEAPAD
54
+ """
55
+
56
+ # A3M with TaxID annotations
57
+ CHAIN_A_TAXID_A3M = """>Query|-|Query
58
+ MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG
59
+ >tr|A0A0B4J2F2|A0A0B4J2F2_HUMAN
60
+ MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG
61
+ >tr|P12345|P12345_MOUSE
62
+ MKTVRQERLKSIVRILERSKDPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG
63
+ >UniRef100_XYZ123 OX=7955
64
+ MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLG
65
+ """
66
+
67
+ CHAIN_B_TAXID_A3M = """>Query|-|Query
68
+ MVTPEGNVSLVDESLLVGVTDEDRAVRSAHQFYERLIGLWAPAVMEAAHELGVFAALAEAPAD
69
+ >tr|Q99999|Q99999_HUMAN
70
+ MVTPEGNVSLVDESLLVGVTDEDRAVRSAHQFYERLIGLWAPAVMEAAHELGVFAALAEAPAD
71
+ >tr|Q54321|Q54321_MOUSE
72
+ MVTPEGNVSLVDESLLVGVTDEDRAVRSAHQFYERLIGLWAPAVMEAAHELGVFAALAERPAD
73
+ >UniRef100_ABC456 OX=7955
74
+ MVTPEGNVSLVDESLLVGVTDEDRAVRSAHQFYERLIGLWAPAVMEAAHELGVFAALAE
75
+ """
76
+
77
+
78
+ class TestA3MParser:
79
+ """Tests for A3M file parsing."""
80
+
81
+ def test_parse_basic(self):
82
+ """Test basic A3M parsing."""
83
+ msa = A3MParser.parse(CHAIN_A_A3M)
84
+
85
+ assert len(msa.sequences) == 5
86
+ assert msa.sequences[0].is_query == True
87
+
88
+ def test_parse_query_sequence(self):
89
+ """Test query sequence identification."""
90
+ msa = A3MParser.parse(CHAIN_A_A3M)
91
+
92
+ query = msa.get_query()
93
+ assert query is not None
94
+ assert query.is_query == True
95
+ assert len(query.sequence) == 65 # Actual length of test sequence
96
+
97
+ def test_parse_uniref_headers(self):
98
+ """Test parsing UniRef-style headers."""
99
+ msa = A3MParser.parse(CHAIN_A_A3M)
100
+
101
+ # Skip query, check first UniRef entry
102
+ seq = msa.sequences[1]
103
+ assert seq.identifier == "UniRef100_A0A2N5EEG3"
104
+ assert seq.organism_id == "A0A2N5EEG3"
105
+
106
+ def test_parse_uniprot_headers(self):
107
+ """Test parsing UniProt-style headers (tr|...|..._SPECIES)."""
108
+ a3m_content = """>Query
109
+ MKTVRQERL
110
+ >tr|A0A0B4J2F2|A0A0B4J2F2_HUMAN
111
+ MKTVRQERL
112
+ >sp|P04637|P53_HUMAN
113
+ MKTVRQERL
114
+ """
115
+ msa = A3MParser.parse(a3m_content)
116
+
117
+ assert len(msa.sequences) == 3
118
+
119
+ # Check UniProt entry parsing
120
+ seq1 = msa.sequences[1]
121
+ assert seq1.identifier == "A0A0B4J2F2"
122
+ assert seq1.species == "HUMAN"
123
+ assert seq1.tax_id == "9606" # Should map HUMAN -> 9606
124
+
125
+ def test_parse_ox_field(self):
126
+ """Test parsing explicit OX= taxonomic ID."""
127
+ a3m_content = """>Query
128
+ MKTVRQERL
129
+ >UniRef100_ABC123 OX=9606 n=5
130
+ MKTVRQERL
131
+ """
132
+ msa = A3MParser.parse(a3m_content)
133
+
134
+ seq = msa.sequences[1]
135
+ assert seq.tax_id == "9606"
136
+
137
+ def test_get_organism_ids(self):
138
+ """Test getting unique organism IDs."""
139
+ msa = A3MParser.parse(CHAIN_A_A3M)
140
+
141
+ org_ids = msa.get_organism_ids()
142
+ assert len(org_ids) == 4 # 4 non-query sequences
143
+ assert "A0A2N5EEG3" in org_ids
144
+ assert "UPI000755DEB5" in org_ids
145
+
146
+ def test_get_tax_ids(self):
147
+ """Test getting unique TaxIDs."""
148
+ msa = A3MParser.parse(CHAIN_A_TAXID_A3M)
149
+
150
+ tax_ids = msa.get_tax_ids()
151
+ assert "9606" in tax_ids # HUMAN
152
+ assert "10090" in tax_ids # MOUSE
153
+ assert "7955" in tax_ids # From OX= field
154
+
155
+
156
+ class TestA3MSequence:
157
+ """Tests for A3MSequence header parsing."""
158
+
159
+ def test_uniprot_format(self):
160
+ """Test UniProt format parsing."""
161
+ seq = A3MSequence(
162
+ header=">tr|A0A0B4J2F2|A0A0B4J2F2_HUMAN",
163
+ sequence="MKTVRQERL"
164
+ )
165
+ assert seq.identifier == "A0A0B4J2F2"
166
+ assert seq.species == "HUMAN"
167
+ assert seq.tax_id == "9606"
168
+
169
+ def test_uniref_format(self):
170
+ """Test UniRef format parsing."""
171
+ seq = A3MSequence(
172
+ header=">UniRef100_A0A2N5EEG3 340 0.994",
173
+ sequence="MKTVRQERL"
174
+ )
175
+ assert seq.identifier == "UniRef100_A0A2N5EEG3"
176
+ assert seq.organism_id == "A0A2N5EEG3"
177
+
178
+ def test_ox_field_parsing(self):
179
+ """Test OX= field extraction."""
180
+ seq = A3MSequence(
181
+ header=">UniRef100_ABC123 OX=9606 n=5 Tax=Homo sapiens",
182
+ sequence="MKTVRQERL"
183
+ )
184
+ assert seq.tax_id == "9606"
185
+
186
+ def test_ncbi_format(self):
187
+ """Test NCBI format with [Species name]."""
188
+ seq = A3MSequence(
189
+ header=">gi|123|ref|NP_001.1| protein [Homo sapiens]",
190
+ sequence="MKTVRQERL"
191
+ )
192
+ assert seq.species == "Homo sapiens"
193
+ assert seq.tax_id == "9606"
194
+
195
+
196
+ class TestSpeciesMapper:
197
+ """Tests for species code to TaxID mapping."""
198
+
199
+ def test_builtin_mapping(self):
200
+ """Test built-in species mapping."""
201
+ assert SpeciesMapper.get_tax_id("HUMAN") == "9606"
202
+ assert SpeciesMapper.get_tax_id("MOUSE") == "10090"
203
+ assert SpeciesMapper.get_tax_id("ECOLI") == "562"
204
+ assert SpeciesMapper.get_tax_id("YEAST") == "559292"
205
+
206
+ def test_unknown_species(self):
207
+ """Test unknown species returns None."""
208
+ assert SpeciesMapper.get_tax_id("UNKNOWN_SPECIES_XYZ") is None
209
+
210
+ def test_builtin_mapping_coverage(self):
211
+ """Test that built-in mapping has common organisms."""
212
+ assert "HUMAN" in SPECIES_TO_TAXID
213
+ assert "MOUSE" in SPECIES_TO_TAXID
214
+ assert "ECOLI" in SPECIES_TO_TAXID
215
+ assert "DROME" in SPECIES_TO_TAXID # Drosophila
216
+ assert "ARATH" in SPECIES_TO_TAXID # Arabidopsis
217
+
218
+
219
+ class TestPairingStrategies:
220
+ """Tests for sequence pairing strategies."""
221
+
222
+ def test_greedy_pairing_uniref(self):
223
+ """Test greedy pairing with UniRef IDs."""
224
+ msa_a = A3MParser.parse(CHAIN_A_A3M)
225
+ msa_b = A3MParser.parse(CHAIN_B_A3M)
226
+
227
+ msa_a.chain_id = 'A'
228
+ msa_b.chain_id = 'B'
229
+
230
+ strategy = GreedyPairingStrategy(use_tax_id=False)
231
+ pairs = strategy.find_pairs({'A': msa_a, 'B': msa_b})
232
+
233
+ # Query + 3 common organisms (A0A2N5EEG3, UPI000755DEB5, B67890)
234
+ assert len(pairs) == 4
235
+
236
+ def test_greedy_pairing_taxid(self):
237
+ """Test greedy pairing with TaxIDs."""
238
+ msa_a = A3MParser.parse(CHAIN_A_TAXID_A3M)
239
+ msa_b = A3MParser.parse(CHAIN_B_TAXID_A3M)
240
+
241
+ msa_a.chain_id = 'A'
242
+ msa_b.chain_id = 'B'
243
+
244
+ strategy = GreedyPairingStrategy(use_tax_id=True)
245
+ pairs = strategy.find_pairs({'A': msa_a, 'B': msa_b})
246
+
247
+ # Query + 3 common TaxIDs (9606/HUMAN, 10090/MOUSE, 7955/zebrafish)
248
+ assert len(pairs) == 4
249
+
250
+ def test_taxonomy_strategy(self):
251
+ """Test taxonomy-based pairing strategy."""
252
+ msa_a = A3MParser.parse(CHAIN_A_TAXID_A3M)
253
+ msa_b = A3MParser.parse(CHAIN_B_TAXID_A3M)
254
+
255
+ msa_a.chain_id = 'A'
256
+ msa_b.chain_id = 'B'
257
+
258
+ strategy = TaxonomyPairingStrategy(strategy="greedy")
259
+ pairs = strategy.find_pairs({'A': msa_a, 'B': msa_b})
260
+
261
+ assert len(pairs) >= 1 # At least query
262
+
263
+
264
+ class TestA3MToCSVConverter:
265
+ """Tests for the main converter class."""
266
+
267
+ def test_basic_conversion_with_uniref(self):
268
+ """Test basic conversion with UniRef IDs (use_tax_id=False)."""
269
+ # Use UniRef-based pairing for test data without TaxIDs
270
+ strategy = GreedyPairingStrategy(use_tax_id=False)
271
+ converter = A3MToCSVConverter(pairing_strategy=strategy)
272
+ result = converter.convert_content(
273
+ a3m_contents={'A': CHAIN_A_A3M, 'B': CHAIN_B_A3M}
274
+ )
275
+
276
+ assert result.num_pairs == 4 # Query + 3 common organisms
277
+ assert len(result.chain_ids) == 2
278
+ assert 'A' in result.csv_per_chain
279
+ assert 'B' in result.csv_per_chain
280
+
281
+ def test_basic_conversion_with_taxid(self):
282
+ """Test basic conversion with TaxID-based pairing."""
283
+ # Use TaxID-based pairing for data with TaxIDs
284
+ strategy = GreedyPairingStrategy(use_tax_id=True)
285
+ converter = A3MToCSVConverter(pairing_strategy=strategy)
286
+ result = converter.convert_content(
287
+ a3m_contents={'A': CHAIN_A_TAXID_A3M, 'B': CHAIN_B_TAXID_A3M}
288
+ )
289
+
290
+ # Query + 3 common TaxIDs (9606/HUMAN, 10090/MOUSE, 7955/zebrafish)
291
+ assert result.num_pairs == 4
292
+
293
+ def test_csv_format(self):
294
+ """Test CSV output format."""
295
+ strategy = GreedyPairingStrategy(use_tax_id=False)
296
+ converter = A3MToCSVConverter(pairing_strategy=strategy)
297
+ result = converter.convert_content(
298
+ a3m_contents={'A': CHAIN_A_A3M, 'B': CHAIN_B_A3M}
299
+ )
300
+
301
+ # Check CSV header
302
+ assert result.csv_per_chain['A'].startswith('key,sequence')
303
+ assert result.csv_per_chain['B'].startswith('key,sequence')
304
+
305
+ # Check that keys match across chains
306
+ lines_a = result.csv_per_chain['A'].split('\n')
307
+ lines_b = result.csv_per_chain['B'].split('\n')
308
+
309
+ keys_a = [line.split(',')[0] for line in lines_a[1:] if line]
310
+ keys_b = [line.split(',')[0] for line in lines_b[1:] if line]
311
+
312
+ assert keys_a == keys_b # Keys should match
313
+
314
+ def test_include_unpaired_false(self):
315
+ """Test that unpaired sequences are excluded by default."""
316
+ strategy = GreedyPairingStrategy(use_tax_id=False)
317
+ converter = A3MToCSVConverter(pairing_strategy=strategy, include_unpaired=False)
318
+ result = converter.convert_content(
319
+ a3m_contents={'A': CHAIN_A_A3M, 'B': CHAIN_B_A3M}
320
+ )
321
+
322
+ # Count rows (excluding header)
323
+ lines_a = result.csv_per_chain['A'].strip().split('\n')
324
+ assert len(lines_a) == 5 # header + 4 pairs (query + 3 common)
325
+
326
+ def test_include_unpaired_true(self):
327
+ """Test that unpaired sequences are included in block-diagonal format."""
328
+ strategy = GreedyPairingStrategy(use_tax_id=False)
329
+ converter = A3MToCSVConverter(pairing_strategy=strategy, include_unpaired=True)
330
+ result = converter.convert_content(
331
+ a3m_contents={'A': CHAIN_A_A3M, 'B': CHAIN_B_A3M}
332
+ )
333
+
334
+ # Count rows (excluding header)
335
+ lines_a = result.csv_per_chain['A'].strip().split('\n')
336
+
337
+ # Should have: header + 4 paired + unpaired from A (Q54321) + unpaired from B (X99999)
338
+ # Chain A has Q54321 that's not in B
339
+ # Chain B has X99999 that's not in A
340
+ assert len(lines_a) > 5 # More than just paired
341
+
342
+ def test_unpaired_block_diagonal_format(self):
343
+ """Test that unpaired sequences use gaps for other chains."""
344
+ strategy = GreedyPairingStrategy(use_tax_id=False)
345
+ converter = A3MToCSVConverter(pairing_strategy=strategy, include_unpaired=True)
346
+ result = converter.convert_content(
347
+ a3m_contents={'A': CHAIN_A_A3M, 'B': CHAIN_B_A3M}
348
+ )
349
+
350
+ # Parse CSVs
351
+ lines_a = result.csv_per_chain['A'].strip().split('\n')
352
+ lines_b = result.csv_per_chain['B'].strip().split('\n')
353
+
354
+ # Find rows where one chain has gaps
355
+ gap_rows_a = [line for line in lines_a[1:] if ',' in line and line.split(',')[1].startswith('-')]
356
+ gap_rows_b = [line for line in lines_b[1:] if ',' in line and line.split(',')[1].startswith('-')]
357
+
358
+ # Chain A should have gap rows (for B's unpaired sequences)
359
+ # Chain B should have gap rows (for A's unpaired sequences)
360
+ # At least one chain should have gaps
361
+ assert len(gap_rows_a) > 0 or len(gap_rows_b) > 0
362
+
363
+ def test_max_pairs_limit(self):
364
+ """Test max_pairs parameter."""
365
+ strategy = GreedyPairingStrategy(use_tax_id=False)
366
+ converter = A3MToCSVConverter(pairing_strategy=strategy, max_pairs=2)
367
+ result = converter.convert_content(
368
+ a3m_contents={'A': CHAIN_A_A3M, 'B': CHAIN_B_A3M}
369
+ )
370
+
371
+ assert result.num_pairs == 2
372
+
373
+ def test_query_sequences_extracted(self):
374
+ """Test that query sequences are correctly extracted."""
375
+ converter = A3MToCSVConverter()
376
+ result = converter.convert_content(
377
+ a3m_contents={'A': CHAIN_A_A3M, 'B': CHAIN_B_A3M}
378
+ )
379
+
380
+ assert 'A' in result.query_sequences
381
+ assert 'B' in result.query_sequences
382
+ assert len(result.query_sequences['A']) == 65 # Actual length
383
+ assert len(result.query_sequences['B']) == 63
384
+
385
+
386
+ class TestConvenienceFunctions:
387
+ """Tests for convenience functions."""
388
+
389
+ def test_convert_a3m_to_multimer_csv(self, tmp_path):
390
+ """Test the main convenience function."""
391
+ # Create temp A3M files
392
+ chain_a_file = tmp_path / "chain_A.a3m"
393
+ chain_b_file = tmp_path / "chain_B.a3m"
394
+ chain_a_file.write_text(CHAIN_A_A3M)
395
+ chain_b_file.write_text(CHAIN_B_A3M)
396
+
397
+ output_file = tmp_path / "paired.csv"
398
+
399
+ result = convert_a3m_to_multimer_csv(
400
+ a3m_files={'A': chain_a_file, 'B': chain_b_file},
401
+ output_path=output_file
402
+ )
403
+
404
+ assert result.num_pairs == 4
405
+ assert output_file.exists()
406
+
407
+ def test_convert_with_include_unpaired(self, tmp_path):
408
+ """Test conversion with include_unpaired=True."""
409
+ chain_a_file = tmp_path / "chain_A.a3m"
410
+ chain_b_file = tmp_path / "chain_B.a3m"
411
+ chain_a_file.write_text(CHAIN_A_A3M)
412
+ chain_b_file.write_text(CHAIN_B_A3M)
413
+
414
+ result = convert_a3m_to_multimer_csv(
415
+ a3m_files={'A': chain_a_file, 'B': chain_b_file},
416
+ include_unpaired=True
417
+ )
418
+
419
+ # With unpaired, we should have more rows
420
+ lines = result.csv_per_chain['A'].strip().split('\n')
421
+ assert len(lines) > 5 # More than header + 4 paired
422
+
423
+ def test_create_paired_msa_per_chain(self, tmp_path):
424
+ """Test creating MSA data structures for Boltz2."""
425
+ chain_a_file = tmp_path / "chain_A.a3m"
426
+ chain_b_file = tmp_path / "chain_B.a3m"
427
+ chain_a_file.write_text(CHAIN_A_A3M)
428
+ chain_b_file.write_text(CHAIN_B_A3M)
429
+
430
+ result = convert_a3m_to_multimer_csv(
431
+ a3m_files={'A': chain_a_file, 'B': chain_b_file}
432
+ )
433
+
434
+ msa_per_chain = create_paired_msa_per_chain(result)
435
+
436
+ assert 'A' in msa_per_chain
437
+ assert 'B' in msa_per_chain
438
+ assert 'paired' in msa_per_chain['A']
439
+ assert 'csv' in msa_per_chain['A']['paired']
440
+
441
+
442
+ class TestAutoDetection:
443
+ """Tests for auto-detection of pairing mode."""
444
+
445
+ def test_auto_detect_taxid_mode(self, tmp_path):
446
+ """Test auto-detection when TaxIDs are present."""
447
+ chain_a_file = tmp_path / "chain_A.a3m"
448
+ chain_b_file = tmp_path / "chain_B.a3m"
449
+ chain_a_file.write_text(CHAIN_A_TAXID_A3M)
450
+ chain_b_file.write_text(CHAIN_B_TAXID_A3M)
451
+
452
+ # With use_tax_id=None, should auto-detect
453
+ result = convert_a3m_to_multimer_csv(
454
+ a3m_files={'A': chain_a_file, 'B': chain_b_file},
455
+ use_tax_id=None # Auto-detect
456
+ )
457
+
458
+ # Should successfully create pairs
459
+ assert result.num_pairs >= 1
460
+
461
+ def test_force_taxid_mode(self, tmp_path):
462
+ """Test forcing TaxID mode."""
463
+ chain_a_file = tmp_path / "chain_A.a3m"
464
+ chain_b_file = tmp_path / "chain_B.a3m"
465
+ chain_a_file.write_text(CHAIN_A_TAXID_A3M)
466
+ chain_b_file.write_text(CHAIN_B_TAXID_A3M)
467
+
468
+ result = convert_a3m_to_multimer_csv(
469
+ a3m_files={'A': chain_a_file, 'B': chain_b_file},
470
+ use_tax_id=True
471
+ )
472
+
473
+ assert result.num_pairs >= 1
474
+
475
+ def test_force_uniref_mode(self, tmp_path):
476
+ """Test forcing UniRef mode."""
477
+ chain_a_file = tmp_path / "chain_A.a3m"
478
+ chain_b_file = tmp_path / "chain_B.a3m"
479
+ chain_a_file.write_text(CHAIN_A_A3M)
480
+ chain_b_file.write_text(CHAIN_B_A3M)
481
+
482
+ result = convert_a3m_to_multimer_csv(
483
+ a3m_files={'A': chain_a_file, 'B': chain_b_file},
484
+ use_tax_id=False
485
+ )
486
+
487
+ assert result.num_pairs == 4
488
+
489
+
490
+ class TestEdgeCases:
491
+ """Tests for edge cases and error handling."""
492
+
493
+ def test_empty_a3m(self):
494
+ """Test handling of empty A3M content."""
495
+ converter = A3MToCSVConverter()
496
+
497
+ empty_a3m = ""
498
+ msa = A3MParser.parse(empty_a3m)
499
+ assert len(msa.sequences) == 0
500
+
501
+ def test_single_sequence_a3m(self):
502
+ """Test A3M with only query sequence."""
503
+ single_a3m = """>Query
504
+ MKTVRQERL
505
+ """
506
+ msa = A3MParser.parse(single_a3m)
507
+ assert len(msa.sequences) == 1
508
+
509
+ def test_no_common_organisms(self):
510
+ """Test when chains have no common organisms."""
511
+ chain_a = """>Query
512
+ MKTVRQERL
513
+ >UniRef100_OnlyInA
514
+ MKTVRQERL
515
+ """
516
+ chain_b = """>Query
517
+ MVTPEGNVS
518
+ >UniRef100_OnlyInB
519
+ MVTPEGNVS
520
+ """
521
+ converter = A3MToCSVConverter()
522
+ result = converter.convert_content(
523
+ a3m_contents={'A': chain_a, 'B': chain_b}
524
+ )
525
+
526
+ # Should only have query pair
527
+ assert result.num_pairs == 1
528
+
529
+ def test_no_common_with_include_unpaired(self):
530
+ """Test including unpaired when there are no common organisms."""
531
+ chain_a = """>Query
532
+ MKTVRQERL
533
+ >UniRef100_OnlyInA
534
+ MKTVRQERL
535
+ """
536
+ chain_b = """>Query
537
+ MVTPEGNVS
538
+ >UniRef100_OnlyInB
539
+ MVTPEGNVS
540
+ """
541
+ converter = A3MToCSVConverter(include_unpaired=True)
542
+ result = converter.convert_content(
543
+ a3m_contents={'A': chain_a, 'B': chain_b}
544
+ )
545
+
546
+ # Should have query pair + 2 unpaired
547
+ lines_a = result.csv_per_chain['A'].strip().split('\n')
548
+ assert len(lines_a) == 4 # header + query + 2 unpaired
549
+
550
+ def test_three_chains(self):
551
+ """Test with three chains."""
552
+ chain_a = """>Query
553
+ MKTVRQERL
554
+ >UniRef100_Common
555
+ MKTVRQERL
556
+ """
557
+ chain_b = """>Query
558
+ MVTPEGNVS
559
+ >UniRef100_Common
560
+ MVTPEGNVS
561
+ """
562
+ chain_c = """>Query
563
+ MAAAAEEEE
564
+ >UniRef100_Common
565
+ MAAAAEEEE
566
+ """
567
+ # Use UniRef-based pairing for test data without TaxIDs
568
+ strategy = GreedyPairingStrategy(use_tax_id=False)
569
+ converter = A3MToCSVConverter(pairing_strategy=strategy)
570
+ result = converter.convert_content(
571
+ a3m_contents={'A': chain_a, 'B': chain_b, 'C': chain_c}
572
+ )
573
+
574
+ assert len(result.chain_ids) == 3
575
+ assert result.num_pairs == 2 # Query + Common
576
+
577
+
578
+ if __name__ == "__main__":
579
+ pytest.main([__file__, "-v"])