Skip to content

Commit 945b670

Browse files
authored
refactor: tests for structure.template (#309)
* refactor: tests for structure.template lock in the behavior seen in #306 * run linter * missing some __init__.py files - otherwise it's ambigous to pytest * more comprehensive tests * remove crappy comments * remove test_template_featurization.py * e2e test: template save single-sequence prediction * review comments from Jennifer
1 parent 2a132de commit 945b670

6 files changed

Lines changed: 5149 additions & 44 deletions

File tree

openfold3/tests/core/data/pipelines/preprocessing/__init__.py

Whitespace-only changes.

openfold3/tests/core/data/primitives/structure/__init__.py

Whitespace-only changes.
Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
# Copyright 2026 AlQuraishi Laboratory
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import dataclasses
16+
from pathlib import Path
17+
18+
import numpy as np
19+
import pytest
20+
21+
from openfold3.core.data.primitives.structure.template import sample_templates
22+
from openfold3.tests.utils.template_helpers import (
23+
TEMPLATE_ID,
24+
make_cache_entry,
25+
template_structure_array_path,
26+
write_cache_npz,
27+
)
28+
29+
30+
def _cache_entry():
31+
"""The single template both written to the cache and expected back out."""
32+
return make_cache_entry([[1, 1], [2, 2]])
33+
34+
35+
def _write_cache_npz(path: Path) -> Path:
36+
return write_cache_npz(path, {TEMPLATE_ID: _cache_entry()})
37+
38+
39+
def _assembly_data(cache_npz: Path) -> dict:
40+
return {
41+
"A": {
42+
"template_ids": [TEMPLATE_ID],
43+
"cache_entry_file_path": cache_npz,
44+
}
45+
}
46+
47+
48+
def _cache_none(tmp_path: Path) -> None:
49+
return None
50+
51+
52+
def _cache_tmp_path(tmp_path: Path) -> Path:
53+
return tmp_path
54+
55+
56+
def _structure_arrays_none(tmp_path: Path) -> None:
57+
"""No preparsed structure arrays -> the existence filter is skipped."""
58+
return None
59+
60+
61+
def _structure_arrays_dummy(tmp_path: Path) -> Path:
62+
"""A dummy structure erray (empty inside)"""
63+
array_dir = tmp_path / "arrays"
64+
struct_path = template_structure_array_path(array_dir)
65+
struct_path.parent.mkdir(parents=True, exist_ok=True)
66+
struct_path.touch()
67+
return array_dir
68+
69+
70+
@pytest.mark.parametrize(
71+
"make_cache_directory, make_structure_array_directory, expected",
72+
[
73+
pytest.param(
74+
_cache_none, _structure_arrays_none, {}, id="no_cache__no_arrays__drops"
75+
),
76+
pytest.param(
77+
_cache_none, _structure_arrays_dummy, {}, id="no_cache__with_arrays__drops"
78+
),
79+
pytest.param(
80+
_cache_tmp_path,
81+
_structure_arrays_none,
82+
{TEMPLATE_ID: _cache_entry()},
83+
id="cache__no_arrays__loads",
84+
),
85+
pytest.param(
86+
_cache_tmp_path,
87+
_structure_arrays_dummy,
88+
{TEMPLATE_ID: _cache_entry()},
89+
id="cache__with_arrays__loads",
90+
),
91+
],
92+
)
93+
def test_sample_templates_cache_directory_gate(
94+
tmp_path, make_cache_directory, make_structure_array_directory, expected
95+
):
96+
cache_npz = _write_cache_npz(tmp_path / "chainA.npz")
97+
98+
actual = sample_templates(
99+
assembly_data=_assembly_data(cache_npz),
100+
template_cache_directory=make_cache_directory(tmp_path),
101+
n_templates=4,
102+
take_top_k=True, # deterministic: k = min(len(ids), n_templates)
103+
chain_id="A",
104+
template_structure_array_directory=make_structure_array_directory(tmp_path),
105+
template_file_format="npz",
106+
)
107+
108+
np.testing.assert_equal(
109+
{k: dataclasses.asdict(v) for k, v in actual.items()},
110+
{k: dataclasses.asdict(v) for k, v in expected.items()},
111+
)

0 commit comments

Comments
 (0)