From 61be15833235eb2f41378ef3594b51196a99d439 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 20 Feb 2024 02:48:06 -0500 Subject: [PATCH 01/10] checkpoint Signed-off-by: Jinzhe Zeng --- deepmd/tf/fit/dos.py | 9 +++-- deepmd/tf/fit/ener.py | 23 ++++++------ deepmd/utils/data_system.py | 10 ++++-- deepmd/utils/out_stat.py | 71 +++++++++++++++++++++++++++++++++++++ source/tests/tf/common.py | 11 ++++-- 5 files changed, 104 insertions(+), 20 deletions(-) create mode 100644 deepmd/utils/out_stat.py diff --git a/deepmd/tf/fit/dos.py b/deepmd/tf/fit/dos.py index e8681f47ea..3e2d63e184 100644 --- a/deepmd/tf/fit/dos.py +++ b/deepmd/tf/fit/dos.py @@ -43,6 +43,9 @@ from deepmd.tf.utils.network import ( one_layer_rand_seed_shift, ) +from deepmd.utils.out_stat import ( + compute_output_stat, +) log = logging.getLogger(__name__) @@ -225,8 +228,10 @@ def _compute_output_stats(self, all_stat, rcond=1e-3, mixed_type=False): sys_tynatom = np.reshape(sys_tynatom, [nsys, -1]) sys_tynatom = sys_tynatom[:, 2:] - dos_shift, resd, rank, s_value = np.linalg.lstsq( - sys_tynatom, sys_dos, rcond=rcond + dos_shift = compute_output_stat( + sys_dos, + sys_tynatom, + rcond=rcond, ) return dos_shift diff --git a/deepmd/tf/fit/ener.py b/deepmd/tf/fit/ener.py index 074856ea6c..3980bd0d46 100644 --- a/deepmd/tf/fit/ener.py +++ b/deepmd/tf/fit/ener.py @@ -53,6 +53,9 @@ from deepmd.tf.utils.spin import ( Spin, ) +from deepmd.utils.out_stat import ( + compute_output_stat, +) if TYPE_CHECKING: pass @@ -290,20 +293,16 @@ def _compute_output_stats(self, all_stat, rcond=1e-3, mixed_type=False): # In this situation, we directly use these assigned energies instead of computing stats. # This will make the loss decrease quickly assigned_atom_ener = np.array( - [ee for ee in self.atom_ener_v if ee is not None] + [ee if ee is not None else np.nan for ee in self.atom_ener_v] ) - assigned_ener_idx = [ - ii for ii, ee in enumerate(self.atom_ener_v) if ee is not None - ] - # np.dot out size: nframe - sys_ener -= np.dot(sys_tynatom[:, assigned_ener_idx], assigned_atom_ener) - sys_tynatom[:, assigned_ener_idx] = 0.0 - energy_shift, resd, rank, s_value = np.linalg.lstsq( - sys_tynatom, sys_ener, rcond=rcond + else: + assigned_atom_ener = None + energy_shift = compute_output_stat( + sys_ener, + sys_tynatom, + assigned_bias=assigned_atom_ener, + rcond=rcond, ) - if len(self.atom_ener) > 0: - for ii in assigned_ener_idx: - energy_shift[ii] = self.atom_ener_v[ii] return energy_shift def compute_input_stats(self, all_stat: dict, protection: float = 1e-2) -> None: diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index 20111558cf..90a45c7f44 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -22,6 +22,9 @@ from deepmd.utils.data import ( DeepmdData, ) +from deepmd.utils.out_stat import ( + compute_output_stat, +) log = logging.getLogger(__name__) @@ -248,10 +251,11 @@ def compute_energy_shift(self, rcond=None, key="energy"): sys_tynatom = np.array(self.natoms_vec, dtype=GLOBAL_NP_FLOAT_PRECISION) sys_tynatom = np.reshape(sys_tynatom, [self.nsystems, -1]) sys_tynatom = sys_tynatom[:, 2:] - energy_shift, resd, rank, s_value = np.linalg.lstsq( - sys_tynatom, sys_ener, rcond=rcond + energy_shift = compute_output_stat( + sys_ener.reshape(-1, 1), + sys_tynatom, ) - return energy_shift + return energy_shift.ravel() def add_dict(self, adict: dict) -> None: """Add items to the data system by a `dict`. diff --git a/deepmd/utils/out_stat.py b/deepmd/utils/out_stat.py new file mode 100644 index 0000000000..66d195afe8 --- /dev/null +++ b/deepmd/utils/out_stat.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Output statistics.""" +from typing import ( + Optional, +) + +import numpy as np + + +def compute_output_stat( + output_redu: np.ndarray, + natoms: np.ndarray, + assigned_bias: Optional[np.ndarray] = None, + rcond: Optional[float] = None, +) -> np.ndarray: + """Compute the output statistics. + + Given the reduced output value and the number of atoms for each atom, + compute the least-squares solution as the atomic output bais. + + Parameters + ---------- + output_redu + The reduced output value, shape is [nframes, ndim]. + natoms + The number of atoms for each atom, shape is [nframes, ntypes]. + assigned_bias + The assigned output bias, shape is [ntypes, ndim]. Set to nan + if not assigned. + rcond + Cut-off ratio for small singular values of a. + + Returns + ------- + np.ndarray + The computed output bias, shape is [ntypes, ndim]. + """ + output_redu = np.array(output_redu) + natoms = np.array(natoms) + # check shape + assert output_redu.ndim == 2 + assert natoms.ndim == 2 + assert output_redu.shape[0] == natoms.shape[0] # nframes + if assigned_bias is not None: + assigned_bias = np.array(assigned_bias) + assert assigned_bias.ndim == 2 + assert assigned_bias.shape[0] == natoms.shape[1] # ntypes + assert assigned_bias.shape[1] == output_redu.shape[1] # ndim + # compute output bias + if assigned_bias is not None: + # Atomic energies stats are incorrect if atomic energies are assigned. + # In this situation, we directly use these assigned energies instead of computing stats. + # This will make the loss decrease quickly + assigned_bias_atom_mask = ~np.isnan(assigned_bias).any(axis=1) + # assigned_bias_masked: nmask, ndim + assigned_bias_masked = assigned_bias[assigned_bias_atom_mask] + # assigned_bias_natoms: nframes, nmask + assigned_bias_natoms = natoms[:, assigned_bias_atom_mask] + # output_redu: nframes, ndim + output_redu -= np.einsum( + "ij,jk->ik", assigned_bias_natoms, assigned_bias_masked + ) + # remove assigned atom + natoms[:, assigned_bias_atom_mask] = 0 + + # computed_output_bias: ntypes, ndim + computed_output_bias, _, _, _ = np.linalg.lstsq(natoms, output_redu, rcond=rcond) + if assigned_bias is not None: + # add back assigned atom; this might not be required + computed_output_bias[assigned_bias_atom_mask] = assigned_bias_masked + return computed_output_bias diff --git a/source/tests/tf/common.py b/source/tests/tf/common.py index a83397c11c..a9abcbe143 100644 --- a/source/tests/tf/common.py +++ b/source/tests/tf/common.py @@ -17,6 +17,9 @@ tf, ) from deepmd.tf.utils import random as dp_random +from deepmd.utils.out_stat import ( + compute_output_stat, +) if GLOBAL_NP_FLOAT_PRECISION == np.float32: global_default_fv_hh = 1e-2 @@ -1041,10 +1044,12 @@ def compute_energy_shift(self): sys_tynatom = np.array(self.natoms_vec, dtype=GLOBAL_NP_FLOAT_PRECISION) sys_tynatom = np.reshape(sys_tynatom, [self.nsystems, -1]) sys_tynatom = sys_tynatom[:, 2:] - energy_shift, resd, rank, s_value = np.linalg.lstsq( - sys_tynatom, sys_ener, rcond=None + energy_shift = compute_output_stat( + sys_ener.reshape(-1, 1), + sys_tynatom, + rcond=None, ) - return energy_shift + return energy_shift.ravel() def process_sys_weights(self, sys_weights): sys_weights = np.array(sys_weights) From fc9ae4ed64eae3a607fca78f940f51fa3dd87277 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 20 Feb 2024 04:57:34 -0500 Subject: [PATCH 02/10] merge Signed-off-by: Jinzhe Zeng --- deepmd/pt/model/descriptor/descriptor.py | 8 ------ deepmd/pt/model/descriptor/dpa1.py | 8 ------ deepmd/pt/model/descriptor/dpa2.py | 8 ------ deepmd/pt/model/descriptor/se_a.py | 8 ------ deepmd/pt/model/task/dipole.py | 8 ------ deepmd/pt/model/task/ener.py | 33 ++++++++++++++--------- deepmd/pt/model/task/fitting.py | 8 ------ deepmd/pt/model/task/polarizability.py | 8 ------ deepmd/pt/utils/stat.py | 21 --------------- deepmd/tf/fit/ener.py | 4 +-- source/tests/pt/test_stat.py | 34 ++++++++---------------- 11 files changed, 34 insertions(+), 114 deletions(-) diff --git a/deepmd/pt/model/descriptor/descriptor.py b/deepmd/pt/model/descriptor/descriptor.py index 16659e444d..4db88e237f 100644 --- a/deepmd/pt/model/descriptor/descriptor.py +++ b/deepmd/pt/model/descriptor/descriptor.py @@ -82,14 +82,6 @@ def get_data_process_key(cls, config): descrpt_type = config["type"] return Descriptor.__plugins.plugins[descrpt_type].get_data_process_key(config) - @property - def data_stat_key(self): - """ - Get the keys for the data statistic of the descriptor. - Return a list of statistic names needed, such as "sumr", "suma" or "sumn". - """ - raise NotImplementedError("data_stat_key is not implemented!") - def __new__(cls, *args, **kwargs): if cls is Descriptor: try: diff --git a/deepmd/pt/model/descriptor/dpa1.py b/deepmd/pt/model/descriptor/dpa1.py index d948c7abf7..8988f3279f 100644 --- a/deepmd/pt/model/descriptor/dpa1.py +++ b/deepmd/pt/model/descriptor/dpa1.py @@ -142,14 +142,6 @@ def get_data_process_key(cls, config): assert descrpt_type in ["dpa1", "se_atten"] return {"sel": config["sel"], "rcut": config["rcut"]} - @property - def data_stat_key(self): - """ - Get the keys for the data statistic of the descriptor. - Return a list of statistic names needed, such as "sumr", "suma" or "sumn". - """ - return ["sumr", "suma", "sumn", "sumr2", "suma2"] - def serialize(self) -> dict: """Serialize the obj to dict.""" raise NotImplementedError diff --git a/deepmd/pt/model/descriptor/dpa2.py b/deepmd/pt/model/descriptor/dpa2.py index 90ec56e0bf..6fa0f36e69 100644 --- a/deepmd/pt/model/descriptor/dpa2.py +++ b/deepmd/pt/model/descriptor/dpa2.py @@ -320,14 +320,6 @@ def get_data_process_key(cls, config): "rcut": [config["repinit_rcut"], config["repformer_rcut"]], } - @property - def data_stat_key(self): - """ - Get the keys for the data statistic of the descriptor. - Return a list of statistic names needed, such as "sumr", "suma" or "sumn". - """ - return ["sumr", "suma", "sumn", "sumr2", "suma2"] - def serialize(self) -> dict: """Serialize the obj to dict.""" raise NotImplementedError diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py index 0f8add1d8d..58d50b09aa 100644 --- a/deepmd/pt/model/descriptor/se_a.py +++ b/deepmd/pt/model/descriptor/se_a.py @@ -135,14 +135,6 @@ def get_data_process_key(cls, config): assert descrpt_type in ["se_e2_a"] return {"sel": config["sel"], "rcut": config["rcut"]} - @property - def data_stat_key(self): - """ - Get the keys for the data statistic of the descriptor. - Return a list of statistic names needed, such as "sumr", "suma" or "sumn". - """ - return ["sumr", "suma", "sumn", "sumr2", "suma2"] - def forward( self, coord_ext: torch.Tensor, diff --git a/deepmd/pt/model/task/dipole.py b/deepmd/pt/model/task/dipole.py index 4ea66e2636..116074b0a2 100644 --- a/deepmd/pt/model/task/dipole.py +++ b/deepmd/pt/model/task/dipole.py @@ -118,14 +118,6 @@ def output_def(self) -> FittingOutputDef: ] ) - @property - def data_stat_key(self): - """ - Get the keys for the data statistic of the fitting. - Return a list of statistic names needed, such as "bias_atom_e". - """ - return [] - def forward( self, descriptor: torch.Tensor, diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index d57d460f6d..5d9fdec49f 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -28,8 +28,8 @@ from deepmd.pt.utils.env import ( DEFAULT_PRECISION, ) -from deepmd.pt.utils.stat import ( - compute_output_bias, +from deepmd.utils.out_stat import ( + compute_output_stat, ) from deepmd.utils.path import ( DPPath, @@ -101,6 +101,8 @@ def __init__( **kwargs, ): self.dim_out = dim_out + # TODO: atom_ener + self.atom_ener = None super().__init__( var_name=var_name, ntypes=ntypes, @@ -128,16 +130,8 @@ def serialize(self) -> dict: data["dim_out"] = self.dim_out return data - @property - def data_stat_key(self): - """ - Get the keys for the data statistic of the fitting. - Return a list of statistic names needed, such as "bias_atom_e". - """ - return ["bias_atom_e"] - def compute_output_stats(self, merged, stat_file_path: Optional[DPPath] = None): - energy = [item["energy"] for item in merged] + energy = [item[self.var_name] for item in merged] data_mixed_type = "real_natoms_vec" in merged[0] if data_mixed_type: input_natoms = [item["real_natoms_vec"] for item in merged] @@ -148,7 +142,22 @@ def compute_output_stats(self, merged, stat_file_path: Optional[DPPath] = None): if stat_file_path is not None and stat_file_path.is_file(): bias_atom_e = stat_file_path.load_numpy() else: - bias_atom_e = compute_output_bias(energy, input_natoms, rcond=self.rcond) + # shape: (nframes, ndim) + merged_energy = torch.cat(energy).detach().cpu().numpy() + # shape: (nframes, ntypes) + merged_natoms = torch.cat(input_natoms)[:, 2:].detach().cpu().numpy() + if self.atom_ener is not None and len(self.atom_ener) > 0: + assigned_atom_ener = np.array( + [ee if ee is not None else np.nan for ee in self.atom_ener] + ) + else: + assigned_atom_ener = None + bias_atom_e = compute_output_stat( + merged_energy, + merged_natoms, + assigned_bias=assigned_atom_ener, + rcond=self.rcond, + ) if stat_file_path is not None: stat_file_path.save_numpy(bias_atom_e) assert all(x is not None for x in [bias_atom_e]) diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index cade533f1a..0cfde97251 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -122,14 +122,6 @@ def share_params(self, base_class, shared_level, resume=False): else: raise NotImplementedError - @property - def data_stat_key(self): - """ - Get the keys for the data statistic of the fitting. - Return a list of statistic names needed, such as "bias_atom_e". - """ - raise NotImplementedError("data_stat_key is not implemented!") - def change_energy_bias( self, config, model, old_type_map, new_type_map, bias_shift="delta", ntest=10 ): diff --git a/deepmd/pt/model/task/polarizability.py b/deepmd/pt/model/task/polarizability.py index dc8d13ee84..1ed5edf872 100644 --- a/deepmd/pt/model/task/polarizability.py +++ b/deepmd/pt/model/task/polarizability.py @@ -150,14 +150,6 @@ def output_def(self) -> FittingOutputDef: ] ) - @property - def data_stat_key(self): - """ - Get the keys for the data statistic of the fitting. - Return a list of statistic names needed, such as "bias_atom_e". - """ - return [] - def forward( self, descriptor: torch.Tensor, diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py index 051fddd14b..0d4847a798 100644 --- a/deepmd/pt/utils/stat.py +++ b/deepmd/pt/utils/stat.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging -import numpy as np import torch log = logging.getLogger(__name__) @@ -56,23 +55,3 @@ def make_stat_input(datasets, dataloaders, nbatches): sys_stat[key] = sys_stat_list lst.append(sys_stat) return lst - - -def compute_output_bias(energy, natoms, rcond=None): - """Update output bias for fitting net. - - Args: - - energy: Batched energy with shape [nframes, 1]. - - natoms: Batched atom statisics with shape [self.ntypes+2]. - - Returns - ------- - - energy_coef: Average enery per atom for each element. - """ - for i in range(len(energy)): - energy[i] = energy[i].mean(dim=0, keepdim=True) - natoms[i] = natoms[i].double().mean(dim=0, keepdim=True) - sys_ener = torch.cat(energy).cpu() - sys_tynatom = torch.cat(natoms)[:, 2:].cpu() - energy_coef, _, _, _ = np.linalg.lstsq(sys_tynatom, sys_ener, rcond) - return energy_coef diff --git a/deepmd/tf/fit/ener.py b/deepmd/tf/fit/ener.py index 3980bd0d46..c18b587180 100644 --- a/deepmd/tf/fit/ener.py +++ b/deepmd/tf/fit/ener.py @@ -298,11 +298,11 @@ def _compute_output_stats(self, all_stat, rcond=1e-3, mixed_type=False): else: assigned_atom_ener = None energy_shift = compute_output_stat( - sys_ener, + sys_ener.reshape(-1, 1), sys_tynatom, assigned_bias=assigned_atom_ener, rcond=rcond, - ) + ).ravel() return energy_shift def compute_input_stats(self, all_stat: dict, protection: float = 1e-2) -> None: diff --git a/source/tests/pt/test_stat.py b/source/tests/pt/test_stat.py index 5cf6a953cc..3b10cf63b2 100644 --- a/source/tests/pt/test_stat.py +++ b/source/tests/pt/test_stat.py @@ -12,7 +12,6 @@ import dpdata import numpy as np -import torch from deepmd.pt.model.descriptor import ( DescrptSeA, @@ -20,15 +19,15 @@ from deepmd.pt.model.descriptor.dpa1 import ( DescrptDPA1, ) +from deepmd.pt.model.task.ener import ( + EnergyFittingNet, +) from deepmd.pt.utils import ( env, ) from deepmd.pt.utils.dataloader import ( DpLoaderSet, ) -from deepmd.pt.utils.stat import ( - compute_output_bias, -) from deepmd.pt.utils.stat import make_stat_input as my_make from deepmd.tf.common import ( expand_sys_str, @@ -126,28 +125,17 @@ def setUp(self): self.dp_d = self.setup_tf() def test_stat_output(self): - def my_merge(energy, natoms): - energy_lst = [] - natoms_lst = [] - for i in range(len(energy)): - for j in range(len(energy[i])): - energy_lst.append(torch.tensor(energy[i][j])) - natoms_lst.append( - torch.tensor(natoms[i][j]) - .unsqueeze(0) - .expand(energy[i][j].shape[0], -1) - ) - return energy_lst, natoms_lst - - energy = self.dp_sampled["energy"] - natoms = self.dp_sampled["natoms_vec"] - energy, natoms = my_merge(energy, natoms) dp_fn = EnerFitting( self.dp_d.get_ntypes(), self.dp_d.get_dim_out(), self.n_neuron ) - dp_fn.compute_output_stats(self.dp_sampled) - bias_atom_e = compute_output_bias(energy, natoms) - self.assertTrue(np.allclose(dp_fn.bias_atom_e, bias_atom_e[:, 0])) + dp_fn.compute_output_stats(self.dp_sampled, mixed_type=self.mixed_type) + pt_fn = EnergyFittingNet( + self.dp_d.get_ntypes(), self.dp_d.get_dim_out(), self.n_neuron + ) + pt_fn.compute_output_stats(self.my_sampled) + np.testing.assert_allclose( + dp_fn.bias_atom_e, pt_fn.bias_atom_e.detach().cpu().numpy().ravel() + ) # temporarily delete this function for performance of seeds in tf and pytorch may be different """ From 4ba3ae76195eade16d6e2861eef7dbd7cd6b6a64 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 20 Feb 2024 05:01:19 -0500 Subject: [PATCH 03/10] loose the shape of assigned_bias Signed-off-by: Jinzhe Zeng --- deepmd/utils/out_stat.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/deepmd/utils/out_stat.py b/deepmd/utils/out_stat.py index 66d195afe8..90e1649426 100644 --- a/deepmd/utils/out_stat.py +++ b/deepmd/utils/out_stat.py @@ -42,10 +42,9 @@ def compute_output_stat( assert natoms.ndim == 2 assert output_redu.shape[0] == natoms.shape[0] # nframes if assigned_bias is not None: - assigned_bias = np.array(assigned_bias) - assert assigned_bias.ndim == 2 - assert assigned_bias.shape[0] == natoms.shape[1] # ntypes - assert assigned_bias.shape[1] == output_redu.shape[1] # ndim + assigned_bias = np.array(assigned_bias).reshape( + natoms.shape[1], output_redu.shape[1] + ) # compute output bias if assigned_bias is not None: # Atomic energies stats are incorrect if atomic energies are assigned. From abeb202fd0923895be6030eb2a47161484267f2a Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 27 Feb 2024 22:27:53 -0500 Subject: [PATCH 04/10] rcond Signed-off-by: Jinzhe Zeng --- deepmd/utils/data_system.py | 1 + 1 file changed, 1 insertion(+) diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index 90a45c7f44..093cd492e9 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -254,6 +254,7 @@ def compute_energy_shift(self, rcond=None, key="energy"): energy_shift = compute_output_stat( sys_ener.reshape(-1, 1), sys_tynatom, + rcond=rcond, ) return energy_shift.ravel() From 6501b4c209963390716db8a5dd45078ce67493a5 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 27 Feb 2024 22:29:50 -0500 Subject: [PATCH 05/10] to_numpy_array Signed-off-by: Jinzhe Zeng --- deepmd/pt/model/task/ener.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 9623b4d838..0259be87f3 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -28,6 +28,9 @@ from deepmd.pt.utils.env import ( DEFAULT_PRECISION, ) +from deepmd.pt.utils.utils import ( + to_numpy_array, +) from deepmd.utils.out_stat import ( compute_output_stat, ) @@ -148,9 +151,9 @@ def compute_output_stats(self, merged, stat_file_path: Optional[DPPath] = None): bias_atom_e = stat_file_path.load_numpy() else: # shape: (nframes, ndim) - merged_energy = torch.cat(energy).detach().cpu().numpy() + merged_energy = to_numpy_array(torch.cat(energy)) # shape: (nframes, ntypes) - merged_natoms = torch.cat(input_natoms)[:, 2:].detach().cpu().numpy() + merged_natoms = to_numpy_array(torch.cat(input_natoms)[:, 2:]) if self.atom_ener is not None and len(self.atom_ener) > 0: assigned_atom_ener = np.array( [ee if ee is not None else np.nan for ee in self.atom_ener] From bf323039f0d21ccc425aad7a4cb79480e91cc84c Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 27 Feb 2024 23:44:33 -0500 Subject: [PATCH 06/10] support atomic; support std Signed-off-by: Jinzhe Zeng --- deepmd/pt/model/task/ener.py | 4 +-- deepmd/tf/fit/dos.py | 4 +-- deepmd/tf/fit/ener.py | 4 +-- deepmd/utils/data_system.py | 4 +-- deepmd/utils/out_stat.py | 55 +++++++++++++++++++++++++++++++++--- source/tests/tf/common.py | 4 +-- 6 files changed, 61 insertions(+), 14 deletions(-) diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 0259be87f3..d236a0d42f 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -32,7 +32,7 @@ to_numpy_array, ) from deepmd.utils.out_stat import ( - compute_output_stat, + compute_bias_from_redu, ) from deepmd.utils.path import ( DPPath, @@ -160,7 +160,7 @@ def compute_output_stats(self, merged, stat_file_path: Optional[DPPath] = None): ) else: assigned_atom_ener = None - bias_atom_e = compute_output_stat( + bias_atom_e, _ = compute_bias_from_redu( merged_energy, merged_natoms, assigned_bias=assigned_atom_ener, diff --git a/deepmd/tf/fit/dos.py b/deepmd/tf/fit/dos.py index 3e2d63e184..a09878f56b 100644 --- a/deepmd/tf/fit/dos.py +++ b/deepmd/tf/fit/dos.py @@ -44,7 +44,7 @@ one_layer_rand_seed_shift, ) from deepmd.utils.out_stat import ( - compute_output_stat, + compute_bias_from_redu, ) log = logging.getLogger(__name__) @@ -228,7 +228,7 @@ def _compute_output_stats(self, all_stat, rcond=1e-3, mixed_type=False): sys_tynatom = np.reshape(sys_tynatom, [nsys, -1]) sys_tynatom = sys_tynatom[:, 2:] - dos_shift = compute_output_stat( + dos_shift, _ = compute_bias_from_redu( sys_dos, sys_tynatom, rcond=rcond, diff --git a/deepmd/tf/fit/ener.py b/deepmd/tf/fit/ener.py index 72565dc2e8..d5b22ecc1f 100644 --- a/deepmd/tf/fit/ener.py +++ b/deepmd/tf/fit/ener.py @@ -54,7 +54,7 @@ Spin, ) from deepmd.utils.out_stat import ( - compute_output_stat, + compute_bias_from_redu, ) if TYPE_CHECKING: @@ -299,7 +299,7 @@ def _compute_output_stats(self, all_stat, rcond=1e-3, mixed_type=False): ) else: assigned_atom_ener = None - energy_shift = compute_output_stat( + energy_shift, _ = compute_bias_from_redu( sys_ener.reshape(-1, 1), sys_tynatom, assigned_bias=assigned_atom_ener, diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index 093cd492e9..a4272c7de8 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -23,7 +23,7 @@ DeepmdData, ) from deepmd.utils.out_stat import ( - compute_output_stat, + compute_bias_from_redu, ) log = logging.getLogger(__name__) @@ -251,7 +251,7 @@ def compute_energy_shift(self, rcond=None, key="energy"): sys_tynatom = np.array(self.natoms_vec, dtype=GLOBAL_NP_FLOAT_PRECISION) sys_tynatom = np.reshape(sys_tynatom, [self.nsystems, -1]) sys_tynatom = sys_tynatom[:, 2:] - energy_shift = compute_output_stat( + energy_shift, _ = compute_bias_from_redu( sys_ener.reshape(-1, 1), sys_tynatom, rcond=rcond, diff --git a/deepmd/utils/out_stat.py b/deepmd/utils/out_stat.py index 90e1649426..b76cddc954 100644 --- a/deepmd/utils/out_stat.py +++ b/deepmd/utils/out_stat.py @@ -2,21 +2,22 @@ """Output statistics.""" from typing import ( Optional, + Tuple, ) import numpy as np -def compute_output_stat( +def compute_bias_from_redu( output_redu: np.ndarray, natoms: np.ndarray, assigned_bias: Optional[np.ndarray] = None, rcond: Optional[float] = None, -) -> np.ndarray: +) -> Tuple[np.ndarray, np.ndarray]: """Compute the output statistics. Given the reduced output value and the number of atoms for each atom, - compute the least-squares solution as the atomic output bais. + compute the least-squares solution as the atomic output bais and std. Parameters ---------- @@ -34,6 +35,8 @@ def compute_output_stat( ------- np.ndarray The computed output bias, shape is [ntypes, ndim]. + np.ndarray + The computed output std, shape is [ntypes, ndim]. """ output_redu = np.array(output_redu) natoms = np.array(natoms) @@ -67,4 +70,48 @@ def compute_output_stat( if assigned_bias is not None: # add back assigned atom; this might not be required computed_output_bias[assigned_bias_atom_mask] = assigned_bias_masked - return computed_output_bias + # rest_redu: nframes, ndim + rest_redu = output_redu - np.einsum("ij,jk->ik", natoms, computed_output_bias) + output_std = rest_redu.std(axis=0) + return computed_output_bias, output_std + + +def compute_stats_from_atomic( + output: np.ndarray, + atype: np.ndarray, +) -> Tuple[np.ndarray, np.ndarray]: + """Compute the output statistics. + + Given the output value and the type of atoms, + compute the atomic output bais and std. + + Parameters + ---------- + output + The output value, shape is [nframes, nloc, ndim]. + atype + The type of atoms, shape is [nframes, nloc]. + + Returns + ------- + np.ndarray + The computed output bias, shape is [ntypes, ndim]. + np.ndarray + The computed output std, shape is [ntypes, ndim]. + """ + output = np.array(output) + atype = np.array(atype) + # check shape + assert output.ndim == 3 + assert atype.ndim == 2 + assert output.shape[:2] == atype.shape + # compute output bias + nframes, nloc, ndim = output.shape + ntypes = atype.max() + 1 + output_bias = np.zeros((ntypes, ndim)) + output_std = np.zeros((ntypes, ndim)) + for type_i in range(ntypes): + mask = atype == type_i + output_bias[type_i] = output[mask].mean(axis=0) + output_std[type_i] = output[mask].std(axis=0) + return output_bias, output_std diff --git a/source/tests/tf/common.py b/source/tests/tf/common.py index a9abcbe143..5c6ced4a77 100644 --- a/source/tests/tf/common.py +++ b/source/tests/tf/common.py @@ -18,7 +18,7 @@ ) from deepmd.tf.utils import random as dp_random from deepmd.utils.out_stat import ( - compute_output_stat, + compute_bias_from_redu, ) if GLOBAL_NP_FLOAT_PRECISION == np.float32: @@ -1044,7 +1044,7 @@ def compute_energy_shift(self): sys_tynatom = np.array(self.natoms_vec, dtype=GLOBAL_NP_FLOAT_PRECISION) sys_tynatom = np.reshape(sys_tynatom, [self.nsystems, -1]) sys_tynatom = sys_tynatom[:, 2:] - energy_shift = compute_output_stat( + energy_shift, _ = compute_bias_from_redu( sys_ener.reshape(-1, 1), sys_tynatom, rcond=None, From 9a7ff32a423d2830ac8787f27a2e1433befd656d Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 27 Feb 2024 23:46:31 -0500 Subject: [PATCH 07/10] rename to compute_output_stats Signed-off-by: Jinzhe Zeng --- deepmd/tf/fit/polar.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/deepmd/tf/fit/polar.py b/deepmd/tf/fit/polar.py index f5cebf9a39..1f28bf394d 100644 --- a/deepmd/tf/fit/polar.py +++ b/deepmd/tf/fit/polar.py @@ -148,16 +148,14 @@ def get_out_size(self) -> int: """Get the output size. Should be 9.""" return 9 - def compute_input_stats(self, all_stat, protection=1e-2): - """Compute the input statistics. + def compute_output_stats(self, all_stat): + """Compute the output statistics. Parameters ---------- all_stat Dictionary of inputs. can be prepared by model.make_stat_input - protection - Divided-by-zero protection """ if "polarizability" not in all_stat.keys(): self.avgeig = np.zeros([9]) From ced2a3ab2b113cfd89183ba4a1ebb1be3d039f71 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Wed, 28 Feb 2024 00:33:18 -0500 Subject: [PATCH 08/10] add tests Signed-off-by: Jinzhe Zeng --- deepmd/pt/model/task/ener.py | 4 +- deepmd/tf/fit/dos.py | 4 +- deepmd/tf/fit/ener.py | 4 +- deepmd/utils/data_system.py | 4 +- deepmd/utils/out_stat.py | 2 +- source/tests/common/test_out_stat.py | 81 ++++++++++++++++++++++++++++ source/tests/tf/common.py | 4 +- 7 files changed, 92 insertions(+), 11 deletions(-) create mode 100644 source/tests/common/test_out_stat.py diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index d236a0d42f..ff7ae6f8ec 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -32,7 +32,7 @@ to_numpy_array, ) from deepmd.utils.out_stat import ( - compute_bias_from_redu, + compute_stats_from_redu, ) from deepmd.utils.path import ( DPPath, @@ -160,7 +160,7 @@ def compute_output_stats(self, merged, stat_file_path: Optional[DPPath] = None): ) else: assigned_atom_ener = None - bias_atom_e, _ = compute_bias_from_redu( + bias_atom_e, _ = compute_stats_from_redu( merged_energy, merged_natoms, assigned_bias=assigned_atom_ener, diff --git a/deepmd/tf/fit/dos.py b/deepmd/tf/fit/dos.py index a09878f56b..0cc5a7df62 100644 --- a/deepmd/tf/fit/dos.py +++ b/deepmd/tf/fit/dos.py @@ -44,7 +44,7 @@ one_layer_rand_seed_shift, ) from deepmd.utils.out_stat import ( - compute_bias_from_redu, + compute_stats_from_redu, ) log = logging.getLogger(__name__) @@ -228,7 +228,7 @@ def _compute_output_stats(self, all_stat, rcond=1e-3, mixed_type=False): sys_tynatom = np.reshape(sys_tynatom, [nsys, -1]) sys_tynatom = sys_tynatom[:, 2:] - dos_shift, _ = compute_bias_from_redu( + dos_shift, _ = compute_stats_from_redu( sys_dos, sys_tynatom, rcond=rcond, diff --git a/deepmd/tf/fit/ener.py b/deepmd/tf/fit/ener.py index d5b22ecc1f..5f17f95e6d 100644 --- a/deepmd/tf/fit/ener.py +++ b/deepmd/tf/fit/ener.py @@ -54,7 +54,7 @@ Spin, ) from deepmd.utils.out_stat import ( - compute_bias_from_redu, + compute_stats_from_redu, ) if TYPE_CHECKING: @@ -299,7 +299,7 @@ def _compute_output_stats(self, all_stat, rcond=1e-3, mixed_type=False): ) else: assigned_atom_ener = None - energy_shift, _ = compute_bias_from_redu( + energy_shift, _ = compute_stats_from_redu( sys_ener.reshape(-1, 1), sys_tynatom, assigned_bias=assigned_atom_ener, diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index a4272c7de8..592b1f9748 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -23,7 +23,7 @@ DeepmdData, ) from deepmd.utils.out_stat import ( - compute_bias_from_redu, + compute_stats_from_redu, ) log = logging.getLogger(__name__) @@ -251,7 +251,7 @@ def compute_energy_shift(self, rcond=None, key="energy"): sys_tynatom = np.array(self.natoms_vec, dtype=GLOBAL_NP_FLOAT_PRECISION) sys_tynatom = np.reshape(sys_tynatom, [self.nsystems, -1]) sys_tynatom = sys_tynatom[:, 2:] - energy_shift, _ = compute_bias_from_redu( + energy_shift, _ = compute_stats_from_redu( sys_ener.reshape(-1, 1), sys_tynatom, rcond=rcond, diff --git a/deepmd/utils/out_stat.py b/deepmd/utils/out_stat.py index b76cddc954..8f68e32417 100644 --- a/deepmd/utils/out_stat.py +++ b/deepmd/utils/out_stat.py @@ -8,7 +8,7 @@ import numpy as np -def compute_bias_from_redu( +def compute_stats_from_redu( output_redu: np.ndarray, natoms: np.ndarray, assigned_bias: Optional[np.ndarray] = None, diff --git a/source/tests/common/test_out_stat.py b/source/tests/common/test_out_stat.py new file mode 100644 index 0000000000..df8bfdf4d8 --- /dev/null +++ b/source/tests/common/test_out_stat.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest + +import numpy as np + +from deepmd.utils.out_stat import ( + compute_stats_from_atomic, + compute_stats_from_redu, +) + + +class TestOutStat(unittest.TestCase): + def setUp(self) -> None: + rng = np.random.default_rng(20240227) + ndim = 5 + nframes = 1000 + ntypes = 3 + nloc = 1000 + self.atype = rng.integers(0, ntypes, size=(nframes, nloc)) + # compute the number of atoms for each type in each frame + self.natoms = np.zeros((nframes, ntypes), dtype=np.int64) + for i in range(ntypes): + self.natoms[:, i] = (self.atype == i).sum(axis=1) + self.mean = rng.random((ntypes, ndim)) * 1e4 + self.std = rng.random((ntypes, ndim)) * 1e-3 + + # generate random output + self.output = rng.normal( + loc=self.mean[self.atype, :], + scale=self.std[self.atype, :], + size=(nframes, nloc, ndim), + ) + self.output_redu = self.output.sum(axis=1) + + return super().setUp() + + def test_compute_stats_from_redu(self): + bias, std = compute_stats_from_redu(self.output_redu, self.natoms) + np.testing.assert_allclose(bias, self.mean, rtol=1e-7) + np.testing.assert_allclose( + std, + np.sqrt(self.natoms.mean(axis=0) @ np.square(self.std)), + rtol=1e-1, + ) + # ensure the sum is close + np.testing.assert_allclose( + self.output_redu, + self.natoms @ bias, + rtol=1e-7, + ) + + def test_compute_stats_from_redu_with_assigned_bias(self): + assigned_bias = np.full_like(self.mean, np.nan) + assigned_bias[0] = self.mean[0] + bias, std = compute_stats_from_redu( + self.output_redu, + self.natoms, + assigned_bias=assigned_bias, + ) + np.testing.assert_allclose(bias, self.mean, rtol=1e-7) + np.testing.assert_allclose(bias[0], self.mean[0], rtol=1e-14) + np.testing.assert_allclose( + std, + np.sqrt(self.natoms.mean(axis=0) @ np.square(self.std)), + rtol=1e-1, + ) + # ensure the sum is close + np.testing.assert_allclose( + self.output_redu, + self.natoms @ bias, + rtol=1e-7, + ) + + def test_compute_stats_from_atomic(self): + bias, std = compute_stats_from_atomic(self.output, self.atype) + np.testing.assert_allclose(bias, self.mean) + np.testing.assert_allclose( + std, + self.std, + rtol=1e-2, + ) diff --git a/source/tests/tf/common.py b/source/tests/tf/common.py index 5c6ced4a77..0bcb29b4b5 100644 --- a/source/tests/tf/common.py +++ b/source/tests/tf/common.py @@ -18,7 +18,7 @@ ) from deepmd.tf.utils import random as dp_random from deepmd.utils.out_stat import ( - compute_bias_from_redu, + compute_stats_from_redu, ) if GLOBAL_NP_FLOAT_PRECISION == np.float32: @@ -1044,7 +1044,7 @@ def compute_energy_shift(self): sys_tynatom = np.array(self.natoms_vec, dtype=GLOBAL_NP_FLOAT_PRECISION) sys_tynatom = np.reshape(sys_tynatom, [self.nsystems, -1]) sys_tynatom = sys_tynatom[:, 2:] - energy_shift, _ = compute_bias_from_redu( + energy_shift, _ = compute_stats_from_redu( sys_ener.reshape(-1, 1), sys_tynatom, rcond=None, From 9f0b9b7e11bde1298da52e454c1407a83babf223 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Wed, 28 Feb 2024 00:43:32 -0500 Subject: [PATCH 09/10] fix typo Signed-off-by: Jinzhe Zeng --- deepmd/tf/fit/ener.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/deepmd/tf/fit/ener.py b/deepmd/tf/fit/ener.py index 5f17f95e6d..5bad9fcc5c 100644 --- a/deepmd/tf/fit/ener.py +++ b/deepmd/tf/fit/ener.py @@ -304,8 +304,8 @@ def _compute_output_stats(self, all_stat, rcond=1e-3, mixed_type=False): sys_tynatom, assigned_bias=assigned_atom_ener, rcond=rcond, - ).ravel() - return energy_shift + ) + return energy_shift.ravel() def compute_input_stats(self, all_stat: dict, protection: float = 1e-2) -> None: """Compute the input statistics. From 7c12ad09422633d4974c84b32b0469a185509cac Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Wed, 28 Feb 2024 02:20:50 -0500 Subject: [PATCH 10/10] store reference std Signed-off-by: Jinzhe Zeng --- source/tests/common/test_out_stat.py | 55 +++++++++++++++++++++++++--- 1 file changed, 49 insertions(+), 6 deletions(-) diff --git a/source/tests/common/test_out_stat.py b/source/tests/common/test_out_stat.py index df8bfdf4d8..c0cfc25071 100644 --- a/source/tests/common/test_out_stat.py +++ b/source/tests/common/test_out_stat.py @@ -37,10 +37,19 @@ def setUp(self) -> None: def test_compute_stats_from_redu(self): bias, std = compute_stats_from_redu(self.output_redu, self.natoms) np.testing.assert_allclose(bias, self.mean, rtol=1e-7) + reference_std = np.array( + [ + 0.01700638138272794, + 0.01954897296228177, + 0.020281857747683162, + 0.010741237959989648, + 0.020258211828681347, + ] + ) np.testing.assert_allclose( std, - np.sqrt(self.natoms.mean(axis=0) @ np.square(self.std)), - rtol=1e-1, + reference_std, + rtol=1e-7, ) # ensure the sum is close np.testing.assert_allclose( @@ -59,10 +68,19 @@ def test_compute_stats_from_redu_with_assigned_bias(self): ) np.testing.assert_allclose(bias, self.mean, rtol=1e-7) np.testing.assert_allclose(bias[0], self.mean[0], rtol=1e-14) + reference_std = np.array( + [ + 0.017015794087883902, + 0.019549011723239484, + 0.020285565914828625, + 0.01074124012073672, + 0.020283557003416414, + ] + ) np.testing.assert_allclose( std, - np.sqrt(self.natoms.mean(axis=0) @ np.square(self.std)), - rtol=1e-1, + reference_std, + rtol=1e-7, ) # ensure the sum is close np.testing.assert_allclose( @@ -74,8 +92,33 @@ def test_compute_stats_from_redu_with_assigned_bias(self): def test_compute_stats_from_atomic(self): bias, std = compute_stats_from_atomic(self.output, self.atype) np.testing.assert_allclose(bias, self.mean) + reference_std = np.array( + [ + [ + 0.0005452949516910239, + 0.000686732800598535, + 0.00089423457667224, + 7.818017989121455e-05, + 0.0004758637035637342, + ], + [ + 2.0610161678825724e-05, + 0.0007728218734771541, + 0.0004754659308165858, + 0.0001809007655290948, + 0.0008187364708029638, + ], + [ + 0.0007935836092665254, + 0.00031176505013516624, + 0.0005469653430009186, + 0.0005652240916389281, + 0.0006087722080071852, + ], + ] + ) np.testing.assert_allclose( std, - self.std, - rtol=1e-2, + reference_std, + rtol=1e-7, )