Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
16 changes: 13 additions & 3 deletions dpdata/system.py
Original file line number Diff line number Diff line change
Expand Up @@ -461,10 +461,20 @@ def sub_system(self, f_idx: int | slice | list | np.ndarray):
slice(None) for _ in self.data[tt.name].shape
]
new_shape[axis_nframes] = f_idx
tmp.data[tt.name] = self.data[tt.name][tuple(new_shape)]
# Advanced indexing already produces an independent array.
# Copy only slice results that still share storage, avoiding a
# second full memcpy for permutations and index lists.
source = self.data[tt.name]
selected = source[tuple(new_shape)]
if np.shares_memory(selected, source):
selected = selected.copy()
tmp.data[tt.name] = selected
else:
# keep the original data
tmp.data[tt.name] = self.data[tt.name]
# Frame-independent values (atom names/types, ``orig``, and
# optional metadata) are usually lists or arrays. Copy them
# as well as frame data so editing a slice can never mutate
# the source system through a shared mutable object.
tmp.data[tt.name] = deepcopy(self.data[tt.name])
return tmp

def append(self, system: System) -> bool:
Expand Down
35 changes: 35 additions & 0 deletions tests/test_system_append.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,41 @@ def test_failed_append(self):
)


class TestAppendOwnership(unittest.TestCase):
"""Regression tests for copy-on-slice and copy-on-first-append semantics."""

def test_sub_system_does_not_alias_metadata(self):
system = dpdata.System(
data={
"atom_names": ["H"],
"atom_numbs": [1],
"atom_types": np.array([0]),
"orig": np.zeros(3),
"cells": np.eye(3).reshape(1, 3, 3),
"coords": np.zeros((1, 1, 3)),
}
)
sub = system[0:1]
sub.data["atom_names"][0] = "X"
sub.data["atom_types"][0] = 1
sub.data["coords"][0, 0, 0] = 123.0
sub.data["cells"][0, 0, 0] = 456.0
self.assertEqual(system.data["atom_names"], ["H"])
np.testing.assert_array_equal(system.data["atom_types"], [0])
self.assertEqual(system.data["coords"][0, 0, 0], 0.0)
self.assertEqual(system.data["cells"][0, 0, 0], 1.0)

def test_first_append_does_not_alias_source(self):
source = dpdata.System("poscars/POSCAR.oh.d", fmt="vasp/poscar")
target = dpdata.System()
target.append(source)

source.data["atom_names"][0] = "X"
source.data["coords"][0, 0, 0] = 123.0
self.assertEqual(target.data["atom_names"][0], "O")
Comment thread
njzjz-bot marked this conversation as resolved.
self.assertNotEqual(target.data["coords"][0, 0, 0], 123.0)

Comment thread
coderabbitai[bot] marked this conversation as resolved.

class TestVaspXmlAppend(unittest.TestCase, CompLabeledSys, IsPBC):
def setUp(self):
self.places = 6
Expand Down
Loading