Olamide Isreal
f971fd0e21
- Nextflow pipeline using chai1 Docker image from Harbor - S3-based input/output paths (s3://omic/eureka/chai-lab/) - GPU-accelerated protein folding with MSA support Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
107 lines
4.4 KiB
Python
Executable File
107 lines
4.4 KiB
Python
Executable File
# Copyright (c) 2024 Chai Discovery, Inc.
|
|
# Licensed under the Apache License, Version 2.0.
|
|
# See the LICENSE file for details.
|
|
|
|
"""
|
|
Tests for inference dataset.
|
|
"""
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from chai_lab.data.dataset.inference_dataset import Input, load_chains_from_raw
|
|
from chai_lab.data.dataset.structure.all_atom_residue_tokenizer import (
|
|
AllAtomResidueTokenizer,
|
|
)
|
|
from chai_lab.data.dataset.structure.all_atom_structure_context import (
|
|
AllAtomStructureContext,
|
|
)
|
|
from chai_lab.data.dataset.structure.chain import Chain
|
|
from chai_lab.data.parsing.structure.entity_type import EntityType
|
|
from chai_lab.data.sources.rdkit import RefConformerGenerator
|
|
|
|
|
|
@pytest.fixture
|
|
def tokenizer() -> AllAtomResidueTokenizer:
|
|
return AllAtomResidueTokenizer(RefConformerGenerator())
|
|
|
|
|
|
def test_malformed_smiles(tokenizer: AllAtomResidueTokenizer):
|
|
"""Malformed SMILES should be dropped."""
|
|
# Zn ligand is malformed (should be [Zn+2])
|
|
inputs = [
|
|
Input("RKDESES", entity_type=EntityType.PROTEIN.value, entity_name="foo"),
|
|
Input("Zn", entity_type=EntityType.LIGAND.value, entity_name="bar"),
|
|
Input("RKEEE", entity_type=EntityType.PROTEIN.value, entity_name="baz"),
|
|
Input("EEEEEEEEEEEE", entity_type=EntityType.PROTEIN.value, entity_name="boz"),
|
|
]
|
|
chains = load_chains_from_raw(
|
|
inputs,
|
|
identifier="test",
|
|
tokenizer=tokenizer,
|
|
)
|
|
assert len(chains) == 3
|
|
for chain in chains:
|
|
# NOTE this check is only valid because there are no residues that are tokenized per-atom
|
|
# Ensures that the entity data and the structure context in each chain are paired correctly
|
|
assert chain.structure_context.num_tokens == len(
|
|
chain.entity_data.full_sequence
|
|
)
|
|
|
|
|
|
def test_ions_parsing(tokenizer: AllAtomResidueTokenizer):
|
|
"""Ions as SMILES strings should carry the correct charge."""
|
|
inputs = [Input("[Mg+2]", entity_type=EntityType.LIGAND.value, entity_name="foo")]
|
|
chains = load_chains_from_raw(inputs, identifier="foo", tokenizer=tokenizer)
|
|
assert len(chains) == 1
|
|
chain = chains[0]
|
|
assert chain.structure_context.num_atoms == 1
|
|
assert chain.structure_context.atom_ref_charge == 2
|
|
assert chain.structure_context.atom_ref_element.item() == 12
|
|
|
|
|
|
def test_protein_with_smiles(tokenizer: AllAtomResidueTokenizer):
|
|
"""Complex with multiple duplicated protein chains and SMILES ligands."""
|
|
# Based on https://www.rcsb.org/structure/1AFS
|
|
seq = "MDSISLRVALNDGNFIPVLGFGTTVPEKVAKDEVIKATKIAIDNGFRHFDSAYLYEVEEEVGQAIRSKIEDGTVKREDIFYTSKLWSTFHRPELVRTCLEKTLKSTQLDYVDLYIIHFPMALQPGDIFFPRDEHGKLLFETVDICDTWEAMEKCKDAGLAKSIGVSNFNCRQLERILNKPGLKYKPVCNQVECHLYLNQSKMLDYCKSKDIILVSYCTLGSSRDKTWVDQKSPVLLDDPVLCAIAKKYKQTPALVALRYQLQRGVVPLIRSFNAKRIKELTQVFEFQLASEDMKALDGLNRNFRYNNAKYFDDHPNHPFTDEN"
|
|
nap = "NC(=O)c1ccc[n+](c1)[CH]2O[CH](CO[P]([O-])(=O)O[P](O)(=O)OC[CH]3O[CH]([CH](O[P](O)(O)=O)[CH]3O)n4cnc5c(N)ncnc45)[CH](O)[CH]2O"
|
|
tes = "O=C4C=C3C(C2CCC1(C(CCC1O)C2CC3)C)(C)CC4"
|
|
inputs = [
|
|
Input(seq, EntityType.PROTEIN.value, entity_name="A"),
|
|
Input(seq, EntityType.PROTEIN.value, entity_name="B"),
|
|
Input(nap, EntityType.LIGAND.value, entity_name="C"),
|
|
Input(nap, EntityType.LIGAND.value, entity_name="D"),
|
|
Input(tes, EntityType.LIGAND.value, entity_name="E"),
|
|
Input(tes, EntityType.LIGAND.value, entity_name="F"),
|
|
]
|
|
chains: list[Chain] = load_chains_from_raw(inputs, tokenizer=tokenizer)
|
|
assert len(chains) == len(inputs)
|
|
|
|
example = AllAtomStructureContext.merge(
|
|
[chain.structure_context for chain in chains]
|
|
)
|
|
|
|
# Should be 1 protein chain, 2 ligand chains
|
|
assert example.token_entity_id.unique().numel() == 3
|
|
assert example.token_asym_id.unique().numel() == 6
|
|
|
|
# Check protein chains
|
|
prot_entity_ids = example.token_entity_id[
|
|
example.token_entity_type == EntityType.PROTEIN.value
|
|
]
|
|
assert torch.unique(prot_entity_ids).numel() == 1
|
|
prot_sym_ids = example.token_sym_id[
|
|
example.token_entity_type == EntityType.PROTEIN.value
|
|
]
|
|
assert torch.unique(prot_sym_ids).numel() == 2 # Two copies of this chain
|
|
|
|
# Check ligand chains
|
|
lig_entity_ids = example.token_entity_id[
|
|
example.token_entity_type == EntityType.LIGAND.value
|
|
]
|
|
assert torch.unique(lig_entity_ids).numel() == 2
|
|
lig_sym_ids = example.token_sym_id[
|
|
example.token_entity_type == EntityType.LIGAND.value
|
|
]
|
|
assert torch.unique(lig_sym_ids).numel() == 2 # Two copies of each ligand
|