diff --git a/dpdata/qe/scf.py b/dpdata/qe/scf.py index f78be7cb3..72202358d 100755 --- a/dpdata/qe/scf.py +++ b/dpdata/qe/scf.py @@ -107,9 +107,10 @@ def get_energy(lines): return energy -def get_force(lines): +def get_force(lines, natoms): blk = get_block(lines, "Forces acting on atoms", skip=1) ret = [] + blk = blk[0 : sum(natoms)] for ii in blk: ret.append([float(jj) for jj in ii.split("=")[1].split()]) ret = np.array(ret) @@ -146,7 +147,7 @@ def get_frame(fname): cell = get_cell(inlines) atom_names, natoms, types, coords = get_coords(inlines, cell) energy = get_energy(outlines) - force = get_force(outlines) + force = get_force(outlines, natoms) stress = get_stress(outlines) * np.linalg.det(cell) return ( atom_names,