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.
- {boltz2_python_client-0.3.2/boltz2_python_client.egg-info → boltz2_python_client-0.3.3}/PKG-INFO +1 -1
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_client/__init__.py +1 -1
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_client/a3m_to_csv_converter.py +88 -11
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_client/cli.py +12 -3
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3/boltz2_python_client.egg-info}/PKG-INFO +1 -1
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_python_client.egg-info/SOURCES.txt +1 -0
- boltz2_python_client-0.3.3/tests/test_a3m_to_csv_converter.py +579 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/ASYNC_GUIDE.md +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/COVALENT_COMPLEX_GUIDE.md +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/LICENSE +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/MANIFEST.in +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/PARAMETERS.md +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/README.md +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/YAML_GUIDE.md +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_client/__main__.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_client/client.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_client/data/__init__.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_client/data/speclist.txt +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_client/exceptions.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_client/models.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_client/models_affinity.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_client/msa_search.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_client/multi_endpoint_client.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_client/utils.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_client/virtual_screening.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_python_client.egg-info/dependency_links.txt +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_python_client.egg-info/entry_points.txt +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_python_client.egg-info/requires.txt +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_python_client.egg-info/top_level.txt +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/01_basic_protein_folding.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/02_protein_structure_prediction_with_msa.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/03_protein_ligand_complex.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/04_covalent_bonding.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/05_dna_protein_complex.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/06_yaml_configurations.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/07_advanced_parameters.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/08_affinity_prediction_simple.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/09_virtual_screening.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/10_msa_search_integration.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/11_msa_search_large_protein.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/12_msa_affinity_prediction.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/15_a3m_to_multimer_csv.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/barnase_barstar_with_msa.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/cdk4_msa_affinity/cdk4_direct.a3m +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/cdk4_msa_affinity/cdk4_msa.a3m +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/cdk4_msa_affinity/cdk4_simple.a3m +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/cdk4_msa_affinity_example.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/comprehensive_multi_endpoint_demo.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/debug_msa_conversion.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/dimer_examples.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/fix_cdk4_msa_search.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/msa-kras-g12c_combined.a3m +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/msa_search_simple_demo.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/multi_endpoint_screening.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/multi_protein_complex.yaml +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/multimer_file_output_example.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/multimer_msa_example.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/protein_ligand.yaml +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/quick_barnase_barstar.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/quick_barnase_barstar_msa.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/quick_gcn4_homodimer.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/quick_heterodimer_example.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/quick_msa_test.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/sars_cov2_mpro_nirmatrelvir.yaml +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/simple_dimer_examples.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/test_boltz2_server.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/test_cdk4_msa_fixed.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/test_cdk4_msa_simple.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/test_msa_a3m_simple.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/test_msa_endpoint.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/test_notebook_fixes.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/licenses/PyYAML-LICENSE +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/licenses/README.md +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/licenses/aiofiles-LICENSE +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/licenses/click-LICENSE +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/licenses/httpx-LICENSE +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/licenses/py3Dmol-LICENSE +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/licenses/pydantic-LICENSE +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/licenses/rich-LICENSE +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/licenses/typing-extensions-LICENSE +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/pyproject.toml +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/setup.cfg +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/tests/__init__.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/tests/conftest.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/tests/test_basic.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/tests/test_cli_multi_endpoint.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/tests/test_comprehensive_stress.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/tests/test_examples_syntax.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/tests/test_integration_scenarios.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/tests/test_msa_search.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/tests/test_multi_endpoint_functionality.py +0 -0
- {boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/tests/test_real_endpoints.py +0 -0
{boltz2_python_client-0.3.2/boltz2_python_client.egg-info → boltz2_python_client-0.3.3}/PKG-INFO
RENAMED
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: boltz2-python-client
|
|
3
|
-
Version: 0.3.
|
|
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>
|
{boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_client/a3m_to_csv_converter.py
RENAMED
|
@@ -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
|
-
|
|
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
|
-
|
|
1230
|
-
|
|
1231
|
-
|
|
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
|
-
|
|
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
|
-
|
|
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
|
|
{boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3/boltz2_python_client.egg-info}/PKG-INFO
RENAMED
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: boltz2-python-client
|
|
3
|
-
Version: 0.3.
|
|
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>
|
|
@@ -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"])
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_client/multi_endpoint_client.py
RENAMED
|
File without changes
|
|
File without changes
|
{boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_client/virtual_screening.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
{boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/boltz2_python_client.egg-info/requires.txt
RENAMED
|
File without changes
|
|
File without changes
|
{boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/01_basic_protein_folding.py
RENAMED
|
File without changes
|
|
File without changes
|
{boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/03_protein_ligand_complex.py
RENAMED
|
File without changes
|
|
File without changes
|
{boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/05_dna_protein_complex.py
RENAMED
|
File without changes
|
{boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/06_yaml_configurations.py
RENAMED
|
File without changes
|
{boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/07_advanced_parameters.py
RENAMED
|
File without changes
|
{boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/08_affinity_prediction_simple.py
RENAMED
|
File without changes
|
|
File without changes
|
{boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/10_msa_search_integration.py
RENAMED
|
File without changes
|
{boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/11_msa_search_large_protein.py
RENAMED
|
File without changes
|
{boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/12_msa_affinity_prediction.py
RENAMED
|
File without changes
|
{boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/15_a3m_to_multimer_csv.py
RENAMED
|
File without changes
|
{boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/barnase_barstar_with_msa.py
RENAMED
|
File without changes
|
{boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/cdk4_msa_affinity/cdk4_direct.a3m
RENAMED
|
File without changes
|
{boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/cdk4_msa_affinity/cdk4_msa.a3m
RENAMED
|
File without changes
|
{boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/cdk4_msa_affinity/cdk4_simple.a3m
RENAMED
|
File without changes
|
{boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/cdk4_msa_affinity_example.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/msa-kras-g12c_combined.a3m
RENAMED
|
File without changes
|
{boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/msa_search_simple_demo.py
RENAMED
|
File without changes
|
{boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/multi_endpoint_screening.py
RENAMED
|
File without changes
|
{boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/multi_protein_complex.yaml
RENAMED
|
File without changes
|
{boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/multimer_file_output_example.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/quick_barnase_barstar_msa.py
RENAMED
|
File without changes
|
|
File without changes
|
{boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/quick_heterodimer_example.py
RENAMED
|
File without changes
|
|
File without changes
|
{boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/examples/sars_cov2_mpro_nirmatrelvir.yaml
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/licenses/typing-extensions-LICENSE
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/tests/test_comprehensive_stress.py
RENAMED
|
File without changes
|
|
File without changes
|
{boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/tests/test_integration_scenarios.py
RENAMED
|
File without changes
|
|
File without changes
|
{boltz2_python_client-0.3.2 → boltz2_python_client-0.3.3}/tests/test_multi_endpoint_functionality.py
RENAMED
|
File without changes
|
|
File without changes
|