diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py index 2bcc87ce42..1c5e3f1c52 100644 --- a/deepmd/pt/utils/stat.py +++ b/deepmd/pt/utils/stat.py @@ -546,7 +546,17 @@ def compute_output_stats_atomic( ] for kk in keys } - # shape: (nframes, nloc, ndim) + # reshape outputs [nframes, nloc * ndim] --> reshape to [nframes * nloc, 1, ndim] for concatenation + # reshape natoms [nframes, nloc] --> reshape to [nframes * nolc, 1] for concatenation + natoms = {k: [sys_v.reshape(-1, 1) for sys_v in v] for k, v in natoms.items()} + outputs = { + k: [ + sys.reshape(natoms[k][sys_idx].shape[0], 1, -1) + for sys_idx, sys in enumerate(v) + ] + for k, v in outputs.items() + } + merged_output = { kk: to_numpy_array(torch.cat(outputs[kk])) for kk in keys diff --git a/source/tests/pt/water_tensor/dipole/O78H156/set.000/atomic_dipole.npy b/source/tests/pt/water_tensor/dipole/O78H156/set.000/atomic_dipole.npy new file mode 100644 index 0000000000..3468b4548b Binary files /dev/null and b/source/tests/pt/water_tensor/dipole/O78H156/set.000/atomic_dipole.npy differ diff --git a/source/tests/pt/water_tensor/dipole/O78H156/set.000/box.npy b/source/tests/pt/water_tensor/dipole/O78H156/set.000/box.npy new file mode 100644 index 0000000000..c8e3a99705 Binary files /dev/null and b/source/tests/pt/water_tensor/dipole/O78H156/set.000/box.npy differ diff --git a/source/tests/pt/water_tensor/dipole/O78H156/set.000/coord.npy b/source/tests/pt/water_tensor/dipole/O78H156/set.000/coord.npy new file mode 100644 index 0000000000..695ea58f8c Binary files /dev/null and b/source/tests/pt/water_tensor/dipole/O78H156/set.000/coord.npy differ diff --git a/source/tests/pt/water_tensor/dipole/O78H156/set.000/energy.npy b/source/tests/pt/water_tensor/dipole/O78H156/set.000/energy.npy new file mode 100644 index 0000000000..06fa030321 Binary files /dev/null and b/source/tests/pt/water_tensor/dipole/O78H156/set.000/energy.npy differ diff --git a/source/tests/pt/water_tensor/dipole/O78H156/set.000/force.npy b/source/tests/pt/water_tensor/dipole/O78H156/set.000/force.npy new file mode 100644 index 0000000000..5fc063cb27 Binary files /dev/null and b/source/tests/pt/water_tensor/dipole/O78H156/set.000/force.npy differ diff --git a/source/tests/pt/water_tensor/dipole/O78H156/type.raw b/source/tests/pt/water_tensor/dipole/O78H156/type.raw new file mode 100644 index 0000000000..720f740b1a --- /dev/null +++ b/source/tests/pt/water_tensor/dipole/O78H156/type.raw @@ -0,0 +1,234 @@ +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 diff --git a/source/tests/pt/water_tensor/dipole/O78H156/type_map.raw b/source/tests/pt/water_tensor/dipole/O78H156/type_map.raw new file mode 100644 index 0000000000..e900768b1d --- /dev/null +++ b/source/tests/pt/water_tensor/dipole/O78H156/type_map.raw @@ -0,0 +1,2 @@ +O +H diff --git a/source/tests/pt/water_tensor/dipole/O96H192/set.000/atomic_dipole.npy b/source/tests/pt/water_tensor/dipole/O96H192/set.000/atomic_dipole.npy new file mode 100644 index 0000000000..6284159f35 Binary files /dev/null and b/source/tests/pt/water_tensor/dipole/O96H192/set.000/atomic_dipole.npy differ diff --git a/source/tests/pt/water_tensor/dipole/O96H192/set.000/box.npy b/source/tests/pt/water_tensor/dipole/O96H192/set.000/box.npy new file mode 100644 index 0000000000..c8e3a99705 Binary files /dev/null and b/source/tests/pt/water_tensor/dipole/O96H192/set.000/box.npy differ diff --git a/source/tests/pt/water_tensor/dipole/O96H192/set.000/coord.npy b/source/tests/pt/water_tensor/dipole/O96H192/set.000/coord.npy new file mode 100644 index 0000000000..06a21871a5 Binary files /dev/null and b/source/tests/pt/water_tensor/dipole/O96H192/set.000/coord.npy differ diff --git a/source/tests/pt/water_tensor/dipole/O96H192/set.000/energy.npy b/source/tests/pt/water_tensor/dipole/O96H192/set.000/energy.npy new file mode 100644 index 0000000000..8ad89998d1 Binary files /dev/null and b/source/tests/pt/water_tensor/dipole/O96H192/set.000/energy.npy differ diff --git a/source/tests/pt/water_tensor/dipole/O96H192/set.000/force.npy b/source/tests/pt/water_tensor/dipole/O96H192/set.000/force.npy new file mode 100644 index 0000000000..ec101f7c0a Binary files /dev/null and b/source/tests/pt/water_tensor/dipole/O96H192/set.000/force.npy differ diff --git a/source/tests/pt/water_tensor/dipole/O96H192/type.raw b/source/tests/pt/water_tensor/dipole/O96H192/type.raw new file mode 100644 index 0000000000..640bf07ba6 --- /dev/null +++ b/source/tests/pt/water_tensor/dipole/O96H192/type.raw @@ -0,0 +1,288 @@ +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 +0 +1 +1 diff --git a/source/tests/pt/water_tensor/dipole/O96H192/type_map.raw b/source/tests/pt/water_tensor/dipole/O96H192/type_map.raw new file mode 100644 index 0000000000..e900768b1d --- /dev/null +++ b/source/tests/pt/water_tensor/dipole/O96H192/type_map.raw @@ -0,0 +1,2 @@ +O +H diff --git a/source/tests/pt/water_tensor/se_e2_a.json b/source/tests/pt/water_tensor/se_e2_a.json index e53caafc96..fc01107895 100644 --- a/source/tests/pt/water_tensor/se_e2_a.json +++ b/source/tests/pt/water_tensor/se_e2_a.json @@ -60,7 +60,9 @@ "training_data": { "systems": [ "pt/water_tensor/dipole/atomic_system", - "pt/water_tensor/dipole/global_system" + "pt/water_tensor/dipole/global_system", + "pt/water_tensor/dipole/O78H156", + "pt/water_tensor/dipole/O96H192" ], "batch_size": "auto", "_comment8": "that's all"