diff --git a/deepmd/entrypoints/test.py b/deepmd/entrypoints/test.py index d285c996e8..182b5a6eba 100644 --- a/deepmd/entrypoints/test.py +++ b/deepmd/entrypoints/test.py @@ -934,18 +934,40 @@ def test_dipole( if detail_file is not None: detail_path = Path(detail_file) + if not atomic: + pe = np.concatenate( + ( + np.reshape(test_data["dipole"][:numb_test], [-1, 3]), + np.reshape(dipole, [-1, 3]), + ), + axis=1, + ) + header_text = "data_x data_y data_z pred_x pred_y pred_z" + else: + pe = np.concatenate( + ( + np.reshape( + test_data["atomic_dipole"][:numb_test], [-1, 3 * sel_natoms] + ), + np.reshape(dipole, [-1, 3 * sel_natoms]), + ), + axis=1, + ) + header_text = [ + f"{letter}{number}" + for number in range(1, sel_natoms + 1) + for letter in ["data_x", "data_y", "data_z"] + ] + [ + f"{letter}{number}" + for number in range(1, sel_natoms + 1) + for letter in ["pred_x", "pred_y", "pred_z"] + ] + header_text = " ".join(header_text) - pe = np.concatenate( - ( - np.reshape(test_data["dipole"][:numb_test], [-1, 3]), - np.reshape(dipole, [-1, 3]), - ), - axis=1, - ) np.savetxt( detail_path.with_suffix(".out"), pe, - header="data_x data_y data_z pred_x pred_y pred_z", + header=header_text, ) return {"rmse": (rmse_f, dipole.size)}