Skip to content

Hoist b_factor and occupancy annotations out of per-sample loop #18

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

Open
longleo17 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Hoist b_factor and occupancy annotations out of per-sample loop #18

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
18 changes: 8 additions & 10 deletions pxdesign/runner/dumper.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -136,26 +136,24 @@ def _save_structure(
entity_poly_type=None,
):
N_sample = pred_coordinates.shape[0]
# Set annotations once before the loop
atom_array.set_annotation(
"b_factor", np.round(np.zeros(len(atom_array)).astype(float), 2)
)
if "occupancy" not in atom_array._annot:
atom_array.set_annotation(
"occupancy", np.round(np.ones(len(atom_array)), 2)
)
for sample_idx in range(N_sample):
output_fpath = os.path.join(
prediction_save_dir, f"{sample_name}_sample_{sample_idx}.cif"
)
# fake b_factor
atom_array.set_annotation(
"b_factor", np.round(np.zeros(len(atom_array)).astype(float), 2)
)
if "occupancy" not in atom_array._annot:
# fake occupancy
atom_array.set_annotation(
"occupancy", np.round(np.ones(len(atom_array)), 2)
)
save_structure_cif(
atom_array,
pred_coordinates[sample_idx],
output_fpath,
entity_poly_type,
sample_name,
# save_wounresol=False,
)

def _save_confidence(
Expand Down
Empty file added tests/__init__.py
View file
Open in desktop
Empty file.
17 changes: 17 additions & 0 deletions tests/test_parallel_cif_writing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
"""Test that annotations are set once before the CIF writing loop."""

from pathlib import Path


def test_annotations_set_before_loop():
"""Verify b_factor and occupancy are set before the per-sample loop."""
source = (
Path(__file__).parent.parent / "pxdesign" / "runner" / "dumper.py"
).read_text()
method_start = source.find("def _save_structure")
method_source = source[method_start:]
b_factor_pos = method_source.find("b_factor")
loop_pos = method_source.find("for sample_idx")
assert b_factor_pos < loop_pos, (
Comment on lines +12 to +15
Copy link

Copilot AI Mar 3, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This test can produce false positives because str.find() returns -1 when the substring is missing (e.g., if b_factor is renamed/removed). In that case b_factor_pos < loop_pos can still pass, so the test would not actually enforce the intended ordering. Consider asserting that method_start, b_factor_pos, and loop_pos are all != -1 and/or searching for the more specific pattern set_annotation("b_factor" (and similarly for occupancy) or using ast parsing to validate the statement order reliably.

Suggested change
method_source = source[method_start:]
b_factor_pos = method_source.find("b_factor")
loop_pos = method_source.find("for sample_idx")
assert b_factor_pos < loop_pos, (
assert method_start != -1, "Expected to find '_save_structure' in dumper.py"
method_source = source[method_start:]
b_factor_pos = method_source.find('set_annotation("b_factor"')
occupancy_pos = method_source.find('set_annotation("occupancy"')
loop_pos = method_source.find("for sample_idx")
assert b_factor_pos != -1, "Expected b_factor annotation to be set in _save_structure"
assert occupancy_pos != -1, "Expected occupancy annotation to be set in _save_structure"
assert loop_pos != -1, "Expected per-sample loop ('for sample_idx') in _save_structure"
assert b_factor_pos < loop_pos and occupancy_pos < loop_pos, (

Copilot uses AI. Check for mistakes.
"Annotations (b_factor, occupancy) should be set once before the per-sample loop"
)