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>
26 lines
982 B
Python
Executable File
26 lines
982 B
Python
Executable File
# Copyright (c) 2024 Chai Discovery, Inc.
|
|
# Licensed under the Apache License, Version 2.0.
|
|
# See the LICENSE file for details.
|
|
import torch
|
|
|
|
from chai_lab.data.dataset.msas.msa_context import NO_PAIRING_KEY
|
|
from chai_lab.data.dataset.msas.preprocess import _UKEY_FOR_QUERY, prepair_ukey
|
|
|
|
|
|
def test_prepair_ukey():
|
|
keys = torch.tensor([1, 1, 2, 1, NO_PAIRING_KEY, 2, 3])
|
|
edit_dists = torch.arange(len(keys))
|
|
|
|
paired = prepair_ukey(keys, edit_dists)
|
|
assert list(paired) == [_UKEY_FOR_QUERY, (1, 0), (2, 0), (1, 1), (2, 1), (3, 0)]
|
|
assert set(paired.values()) == set(
|
|
[i for i, val in enumerate(keys.tolist()) if val != NO_PAIRING_KEY]
|
|
)
|
|
|
|
# Reverse the edit distances
|
|
paired = prepair_ukey(keys, torch.tensor(edit_dists.tolist()[::-1]))
|
|
assert list(paired) == [_UKEY_FOR_QUERY, (1, 1), (2, 1), (1, 0), (2, 0), (3, 0)]
|
|
assert set(paired.values()) == set(
|
|
[i for i, val in enumerate(keys.tolist()) if val != NO_PAIRING_KEY]
|
|
)
|