Replace deepcopy with coord save/restore in CIF writing

pxdesign/runner/dumper.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import copy
 import json
 import os
 from pathlib import Path
@@ -195,12 +194,12 @@ def save_structure_cif(
         entity_poly_type (dict[str, str]): The entity poly type information.
         pdb_id (str): The PDB ID for the entry.
     """
-    pred_atom_array = copy.deepcopy(atom_array)
-    pred_pose = pred_coordinate.cpu().numpy()
-    pred_atom_array.coord = pred_pose
+    original_coord = atom_array.coord
+    atom_array.coord = pred_coordinate.cpu().numpy()
     save_atoms_to_cif(
         output_fpath,
-        pred_atom_array,
+        atom_array,
         entity_poly_type,
         pdb_id,
     )
+    atom_array.coord = original_coord

tests/test_avoid_deepcopy_cif.py

@@ -0,0 +1,44 @@
+"""Test that CIF writing uses coord swap instead of deepcopy."""
+
+import ast
+from pathlib import Path
+
+import numpy as np
+
+
+def test_no_deepcopy_in_dumper():
+    """Verify that dumper.py no longer uses copy.deepcopy."""
+    filepath = Path(__file__).parent.parent / "pxdesign" / "runner" / "dumper.py"
+    tree = ast.parse(filepath.read_text())
+
+    for node in ast.walk(tree):
+        if isinstance(node, ast.Call):
+            func = node.func
+            if (
+                isinstance(func, ast.Attribute)
+                and func.attr == "deepcopy"
+                and isinstance(func.value, ast.Name)
+                and func.value.id == "copy"
+            ):
+                raise AssertionError(
+                    "Found copy.deepcopy() in dumper.py. "
+                    "Use coord save/restore instead to avoid N_sample full copies."
+                )
+
+
+def test_coord_restore_correctness():
+    """Verify the coord swap pattern restores original coordinates."""
+    original_coords = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+
+    class MockArray:
+        def __init__(self):
+            self.coord = original_coords.copy()
+
+    mock = MockArray()
+    new_coords = np.array([[10.0, 20.0, 30.0], [40.0, 50.0, 60.0]])
+
+    saved = mock.coord
+    mock.coord = new_coords
+    assert np.array_equal(mock.coord, new_coords), "Should have new coords during write"
+    mock.coord = saved
+    assert np.array_equal(mock.coord, original_coords), "Should restore original coords after write"