diff --git a/psiflow/geometry.py b/psiflow/geometry.py index 9e5f8f3..9654ec2 100644 --- a/psiflow/geometry.py +++ b/psiflow/geometry.py @@ -273,14 +273,27 @@ def from_string(cls, s: str, natoms: Optional[int] = None) -> Optional[Geometry] comment_dict = key_val_str_to_dict_regex(comment) # read and format per_atom data + column_indices = {} + if 'Properties' in comment_dict: + properties = comment_dict['Properties'].split(':') + count = 0 + for i in range(len(properties) // 3): + name = properties[3 * i] + ncolumns = int(properties[3 * i + 2]) + column_indices[name] = count + count += ncolumns + assert 'pos' in column_indices # positions need to be there + per_atom = np.recarray(natoms, dtype=per_atom_dtype) per_atom.forces[:] = np.nan + POS_INDEX = column_indices.get('pos', 1) + FORCES_INDEX = column_indices.get('forces', None) for i in range(natoms): values = lines[i + 1].split() per_atom.numbers[i] = chemical_symbols.index(values[0]) - per_atom.positions[i, :] = [float(_) for _ in values[1:4]] - if len(values) > 4: - per_atom.forces[i, :] = [float(_) for _ in values[4:7]] + per_atom.positions[i, :] = [float(_) for _ in values[POS_INDEX:POS_INDEX + 3]] + if FORCES_INDEX is not None: + per_atom.forces[i, :] = [float(_) for _ in values[FORCES_INDEX:FORCES_INDEX + 3]] order = {} for key, value in comment_dict.items(): diff --git a/tests/test_data.py b/tests/test_data.py index 294ed33..738993f 100644 --- a/tests/test_data.py +++ b/tests/test_data.py @@ -150,21 +150,25 @@ def test_readwrite_cycle(dataset, tmp_path): assert "test" in states[2].order s = """3 -energy=3.0 phase=c7eq -C 0 1 2 -O 1 2 3 -F 4 5 6 +energy=3.0 phase=c7eq Properties=species:S:1:pos:R:3:momenta:R:3:forces:R:3 +C 0 1 2 3 4 5 6 7 8 +O 1 2 3 4 5 6 7 8 9 +F 2 3 4 5 6 7 8 9 10 """ geometry = Geometry.from_string(s, natoms=None) assert len(geometry) == 3 assert geometry.energy == 3.0 assert geometry.phase == "c7eq" assert not geometry.periodic - assert np.all(np.isnan(geometry.per_atom.forces)) + assert np.all(np.logical_not(np.isnan(geometry.per_atom.forces))) assert np.allclose( geometry.per_atom.numbers, np.array([6, 8, 9]), ) + assert np.allclose( + geometry.per_atom.forces, + np.array([[6,7,8], [7,8,9], [8,9,10]]), + ) s = """7 O 0.269073490000000 0.952731530000000 0.639899630000000