Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add test for ion charge #214

Merged
merged 2 commits into from
Dec 4, 2024
Merged

Add test for ion charge #214

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3b43b2a
Add test for ion charge
wukevin Dec 4, 2024
0ddfc90
Check that Mg has the right element number
wukevin Dec 4, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 20 additions & 2 deletions tests/test_inference_dataset.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,8 @@
Tests for inference dataset.
"""

import pytest

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,
Expand All @@ -14,7 +16,12 @@
from chai_lab.data.sources.rdkit import RefConformerGenerator


def test_malformed_smiles():
@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 = [
Expand All @@ -26,7 +33,7 @@ def test_malformed_smiles():
chains = load_chains_from_raw(
inputs,
identifier="test",
tokenizer=AllAtomResidueTokenizer(RefConformerGenerator()),
tokenizer=tokenizer,
)
assert len(chains) == 3
for chain in chains:
Expand All @@ -35,3 +42,14 @@ def test_malformed_smiles():
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
Loading