diff --git a/esm/utils/structure/protein_chain.py b/esm/utils/structure/protein_chain.py index 4e433bd4..d994a1aa 100644 --- a/esm/utils/structure/protein_chain.py +++ b/esm/utils/structure/protein_chain.py @@ -1244,7 +1244,9 @@ def infer_cbeta(self, infer_cbeta_for_glycine: bool = False) -> ProteinChain: atom37_positions = self.atom37_positions.copy() atom37_mask = self.atom37_mask.copy() - inferred_cbeta_positions = self.inferred_cbeta + # Copy so the in-place glycine NaN-fill below does not mutate the cached + # `inferred_cbeta` array (and everything derived from it, e.g. pdist_CB). + inferred_cbeta_positions = self.inferred_cbeta.copy() if not infer_cbeta_for_glycine: inferred_cbeta_positions[np.array(list(self.sequence)) == "G", :] = np.nan diff --git a/esm/utils/structure/protein_chain_test.py b/esm/utils/structure/protein_chain_test.py new file mode 100644 index 00000000..5409f4b2 --- /dev/null +++ b/esm/utils/structure/protein_chain_test.py @@ -0,0 +1,37 @@ +"""Tests for protein_chain.py""" + +import numpy as np + +from esm.utils.structure.protein_chain import ProteinChain + + +def _make_chain(sequence: str) -> ProteinChain: + n = len(sequence) + rng = np.random.default_rng(0) + coords = rng.uniform(-10, 10, (n, 37, 3)).astype(np.float32) + return ProteinChain( + id="test", + sequence=sequence, + chain_id="A", + entity_id=1, + residue_index=np.arange(n), + insertion_code=np.full(n, "", dtype="