mor einference fixes

This commit is contained in:
Raktim Mitra
2026-02-05 12:40:50 -08:00
parent 4caba95d32
commit 45e8f837ce
2 changed files with 19 additions and 8 deletions

View File

@@ -117,7 +117,6 @@ def fetch_motif_residue_(
NB: For glycines, we extend the array with a CB position so as to not leak whether
the original residue is a glycine if sequence is masked during inference.
"""
assert (
src_atom_array is not None
), "Motif provided in contigs, but no input provided. input={} contig={}".format(
@@ -227,9 +226,15 @@ def fetch_motif_residue_(
"is_motif_atom_with_fixed_coord", np.zeros(subarray.shape[0], dtype=int)
)
else:
subarray.set_annotation(
"is_motif_atom_with_fixed_coord", subarray.is_motif_atom.copy()
)
if "is_motif_atom" in subarray.get_annotation_categories():
subarray.set_annotation(
"is_motif_atom_with_fixed_coord", subarray.is_motif_atom.copy()
)
else:
subarray.set_annotation(
"is_motif_atom_with_fixed_coord", np.array([True]*len(subarray))
)
if flexible_backbone:
backbone_atoms = ["N", "CA", "C", "O"]
is_flexible_motif_atom = np.isin(subarray.atom_name, backbone_atoms)
@@ -242,12 +247,17 @@ def fetch_motif_residue_(
"is_flexible_motif_atom", np.zeros(subarray.shape[0], dtype=bool)
)
if to_unindex:
subarray.set_annotation(
"is_motif_atom_unindexed", subarray.is_motif_atom.copy()
)
if "is_motif_atom" in subarray.get_annotation_categories():
subarray.set_annotation(
"is_motif_atom_unindexed", subarray.is_motif_atom.copy()
)
else:
subarray.set_annotation(
"is_motif_atom_unindexed", np.array([True]*len(subarray))
)
# Subset to desired motif atoms
subarray = subarray[subarray.is_motif_atom.astype(bool)]
# ... Relax sequence constraint if provided
if (
exists(unfixed_sequence_components)

View File

@@ -237,6 +237,7 @@ class PadTokensWithVirtualAtoms(Transform):
mask = get_af3_token_representative_masks(
token, central_atom=central_atom
)
assert_single_representative(token, central_atom=central_atom)
# ... Create virtual atoms