diff --git a/dpdata/deepmd/comp.py b/dpdata/deepmd/comp.py index 1cd2b055a..3ac239b05 100644 --- a/dpdata/deepmd/comp.py +++ b/dpdata/deepmd/comp.py @@ -8,9 +8,12 @@ def _cond_load_data(fname) : tmp = np.load(fname) return tmp -def _load_set(folder) : - cells = np.load(os.path.join(folder, 'box.npy')) +def _load_set(folder, nopbc: bool) : coords = np.load(os.path.join(folder, 'coord.npy')) + if nopbc: + cells = np.zeros((coords.shape[0], 3,3)) + else: + cells = np.load(os.path.join(folder, 'box.npy')) eners = _cond_load_data(os.path.join(folder, 'energy.npy')) forces = _cond_load_data(os.path.join(folder, 'force.npy')) virs = _cond_load_data(os.path.join(folder, 'virial.npy')) @@ -22,6 +25,8 @@ def to_system_data(folder, # data is empty data = load_type(folder, type_map = type_map) data['orig'] = np.zeros([3]) + if os.path.isfile(os.path.join(folder, "nopbc")): + data['nopbc'] = True sets = sorted(glob.glob(os.path.join(folder, 'set.*'))) all_cells = [] all_coords = [] @@ -29,7 +34,7 @@ def to_system_data(folder, all_forces = [] all_virs = [] for ii in sets : - cells, coords, eners, forces, virs = _load_set(ii) + cells, coords, eners, forces, virs = _load_set(ii, data.get('nopbc', False)) nframes = np.reshape(cells, [-1,3,3]).shape[0] all_cells.append(np.reshape(cells, [nframes,3,3])) all_coords.append(np.reshape(coords, [nframes,-1,3])) @@ -50,8 +55,6 @@ def to_system_data(folder, data['forces'] = np.concatenate(all_forces, axis = 0) if len(all_virs) > 0: data['virials'] = np.concatenate(all_virs, axis = 0) - if os.path.isfile(os.path.join(folder, "nopbc")): - data['nopbc'] = True return data diff --git a/dpdata/deepmd/raw.py b/dpdata/deepmd/raw.py index ab3290f61..49744d0e5 100644 --- a/dpdata/deepmd/raw.py +++ b/dpdata/deepmd/raw.py @@ -33,10 +33,13 @@ def to_system_data(folder, type_map = None, labels = True) : if os.path.isdir(folder) : data = load_type(folder, type_map = type_map) data['orig'] = np.zeros([3]) - data['cells'] = np.loadtxt(os.path.join(folder, 'box.raw')) - data['coords'] = np.loadtxt(os.path.join(folder, 'coord.raw')) - data['cells'] = np.reshape(data['cells'], [-1, 3, 3]) - nframes = data['cells'].shape[0] + data['coords'] = np.loadtxt(os.path.join(folder, 'coord.raw'), ndmin=2) + nframes = data['coords'].shape[0] + if os.path.isfile(os.path.join(folder, "nopbc")): + data['nopbc'] = True + data['cells'] = np.zeros((nframes, 3,3)) + else: + data['cells'] = np.loadtxt(os.path.join(folder, 'box.raw'), ndmin=2) data['cells'] = np.reshape(data['cells'], [nframes, 3, 3]) data['coords'] = np.reshape(data['coords'], [nframes, -1, 3]) if labels : diff --git a/tests/comp_sys.py b/tests/comp_sys.py index d4934c1cb..2b87828c5 100644 --- a/tests/comp_sys.py +++ b/tests/comp_sys.py @@ -35,10 +35,11 @@ def test_nframs(self): def test_cell(self): self.assertEqual(self.system_1.get_nframes(), self.system_2.get_nframes()) - np.testing.assert_almost_equal(self.system_1.data['cells'], - self.system_2.data['cells'], - decimal = self.places, - err_msg = 'cell failed') + if not self.system_1.nopbc and not self.system_2.nopbc: + np.testing.assert_almost_equal(self.system_1.data['cells'], + self.system_2.data['cells'], + decimal = self.places, + err_msg = 'cell failed') def test_coord(self): self.assertEqual(self.system_1.get_nframes(),