From e2b61059aee0a5bc974b27fe81ef2943ab08f921 Mon Sep 17 00:00:00 2001
From: winklemad <winklemad@outlook.com>
Date: Mon, 13 Jul 2026 21:45:59 +0530
Subject: [PATCH] Fix infer_cbeta corrupting the cached inferred_cbeta array

ProteinChain.infer_cbeta() takes the cached `inferred_cbeta` array by reference
and NaN-fills the glycine rows in place to build a new chain. Because
`inferred_cbeta` is a cached_property, this permanently corrupts the original
chain's cached array -- and everything derived from it, e.g. `pdist_CB` -- so the
value of `inferred_cbeta` / `pdist_CB` depends on whether `infer_cbeta()` was ever
called on the chain.

The same method already copies `atom37_positions` / `atom37_mask` before mutating
them, and the sibling `ProteinComplex.infer_cbeta` computes a fresh array. Copy
the array before the in-place NaN-fill. Added a regression test.
---
 esm/utils/structure/protein_chain.py      |  4 ++-
 esm/utils/structure/protein_chain_test.py | 37 +++++++++++++++++++++++
 2 files changed, 40 insertions(+), 1 deletion(-)
 create mode 100644 esm/utils/structure/protein_chain_test.py

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="<U4"),
+        atom37_positions=coords,
+        atom37_mask=np.ones((n, 37), dtype=bool),
+        confidence=np.ones(n),
+    )
+
+
+def test_infer_cbeta_does_not_mutate_cached_inferred_cbeta():
+    # infer_cbeta() NaN-fills the glycine CB positions to build a *new* chain; it
+    # must not mutate the original chain's cached `inferred_cbeta` array (nor
+    # anything derived from it, such as pdist_CB).
+    chain = _make_chain("AGA")  # glycine at index 1
+
+    original = chain.inferred_cbeta.copy()
+    original_pdist = chain.pdist_CB.copy()
+
+    chain.infer_cbeta()
+
+    np.testing.assert_array_equal(chain.inferred_cbeta, original)
+    np.testing.assert_array_equal(chain.pdist_CB, original_pdist)