From 1185bc1e1eaabec040a8f5b89446e3f77bb047ee Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Tue, 27 Feb 2024 02:10:53 +0800 Subject: [PATCH] Fix uts; change force_real to force --- deepmd/pt/infer/deep_eval.py | 23 ++----------------- deepmd/pt/loss/ener_spin.py | 6 ++--- deepmd/pt/model/model/spin_model.py | 10 ++++---- deepmd/tf/descriptor/se_a.py | 6 ++--- source/tests/pt/model/test_autodiff.py | 6 ++--- source/tests/pt/model/test_ener_spin_model.py | 2 +- source/tests/pt/model/test_forward_lower.py | 8 +++---- source/tests/pt/model/test_permutation.py | 4 ++-- source/tests/pt/model/test_rot.py | 6 ++--- source/tests/pt/model/test_smooth.py | 4 ++-- source/tests/pt/model/test_trans.py | 4 ++-- source/tests/pt/model/test_unused_params.py | 11 +++------ 12 files changed, 32 insertions(+), 58 deletions(-) diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index c693586aa9..9a8af68aa2 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -414,7 +414,6 @@ def eval_model( energy_out = [] atomic_energy_out = [] force_out = [] - force_real_out = [] force_mag_out = [] virial_out = [] atomic_virial_out = [] @@ -502,8 +501,6 @@ def eval_model( ) if "force" in batch_output: force_out.append(batch_output["force"].detach().cpu().numpy()) - if "force_real" in batch_output: - force_real_out.append(batch_output["force_real"].detach().cpu().numpy()) if "force_mag" in batch_output: force_mag_out.append(batch_output["force_mag"].detach().cpu().numpy()) if "virial" in batch_output: @@ -525,8 +522,6 @@ def eval_model( atomic_energy_out.append(batch_output["atom_energy"]) if "force" in batch_output: force_out.append(batch_output["force"]) - if "force_real" in batch_output: - force_real_out.append(batch_output["force_real"]) if "force_mag" in batch_output: force_mag_out.append(batch_output["force_mag"]) if "virial" in batch_output: @@ -549,11 +544,6 @@ def eval_model( force_out = ( np.concatenate(force_out) if force_out else np.zeros([nframes, natoms, 3]) ) - force_real_out = ( - np.concatenate(force_real_out) - if force_real_out - else np.zeros([nframes, natoms, 3]) - ) force_mag_out = ( np.concatenate(force_mag_out) if force_mag_out @@ -593,13 +583,6 @@ def eval_model( [nframes, natoms, 3], dtype=GLOBAL_PT_FLOAT_PRECISION, device=DEVICE ) ) - force_real_out = ( - torch.cat(force_real_out) - if force_real_out - else torch.zeros( - [nframes, natoms, 3], dtype=GLOBAL_PT_FLOAT_PRECISION, device=DEVICE - ) - ) force_mag_out = ( torch.cat(force_mag_out) if force_mag_out @@ -628,12 +611,10 @@ def eval_model( else: results_dict = { "energy": energy_out, + "force": force_out, "virial": virial_out, } - if not getattr(model, "__USE_SPIN_INPUT__", False): - results_dict["force"] = force_out - else: - results_dict["force_real"] = force_real_out + if getattr(model, "__USE_SPIN_INPUT__", False): results_dict["force_mag"] = force_mag_out if atomic: results_dict["atom_energy"] = atomic_energy_out diff --git a/deepmd/pt/loss/ener_spin.py b/deepmd/pt/loss/ener_spin.py index ef0a4d9c1e..5a9e5bc528 100644 --- a/deepmd/pt/loss/ener_spin.py +++ b/deepmd/pt/loss/ener_spin.py @@ -106,9 +106,9 @@ def forward(self, model_pred, label, natoms, learning_rate, mae=False): ) more_loss["mae_e_all"] = mae_e_all.detach() - if self.has_fr and "force_real" in model_pred and "force" in label: + if self.has_fr and "force" in model_pred and "force" in label: if not self.use_l1_all: - diff_fr = label["force"] - model_pred["force_real"] + diff_fr = label["force"] - model_pred["force"] l2_force_real_loss = torch.mean(torch.square(diff_fr)) if not self.inference: more_loss["l2_force_r_loss"] = l2_force_real_loss.detach() @@ -120,7 +120,7 @@ def forward(self, model_pred, label, natoms, learning_rate, mae=False): more_loss["mae_fr"] = mae_fr.detach() else: l1_force_real_loss = F.l1_loss( - label["force"], model_pred["force_real"], reduction="none" + label["force"], model_pred["force"], reduction="none" ) more_loss["mae_fr"] = l1_force_real_loss.mean().detach() l1_force_real_loss = l1_force_real_loss.sum(-1).mean(-1).sum() diff --git a/deepmd/pt/model/model/spin_model.py b/deepmd/pt/model/model/spin_model.py index f1c7583706..1cd9b132c9 100644 --- a/deepmd/pt/model/model/spin_model.py +++ b/deepmd/pt/model/model/spin_model.py @@ -489,7 +489,7 @@ def forward( model_predict["energy"] = model_ret["energy_redu"] model_predict["mask_mag"] = model_ret["mask_mag"] if self.backbone_model.do_grad_r("energy"): - model_predict["force_real"] = model_ret["energy_derv_r"].squeeze(-2) + model_predict["force"] = model_ret["energy_derv_r"].squeeze(-2) model_predict["force_mag"] = model_ret["energy_derv_r_mag"].squeeze(-2) if self.backbone_model.do_grad_c("energy"): model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2) @@ -498,7 +498,7 @@ def forward( else: assert model_ret["dforce_real"] is not None assert model_ret["dforce_mag"] is not None - model_predict["force_real"] = model_ret["dforce_real"] + model_predict["force"] = model_ret["dforce_real"] model_predict["force_mag"] = model_ret["dforce_mag"] return model_predict @@ -530,9 +530,7 @@ def forward_lower( model_predict["energy"] = model_ret["energy_redu"] model_predict["mask_mag"] = model_ret["mask_mag"] if self.backbone_model.do_grad_r("energy"): - model_predict["extended_force_real"] = model_ret[ - "energy_derv_r" - ].squeeze(-2) + model_predict["extended_force"] = model_ret["energy_derv_r"].squeeze(-2) model_predict["extended_force_mag"] = model_ret[ "energy_derv_r_mag" ].squeeze(-2) @@ -545,7 +543,7 @@ def forward_lower( else: assert model_ret["dforce_real"] is not None assert model_ret["dforce_mag"] is not None - model_predict["extended_force_real"] = model_ret["dforce_real"] + model_predict["extended_force"] = model_ret["dforce_real"] model_predict["extended_force_mag"] = model_ret["dforce_mag"] else: model_predict = model_ret diff --git a/deepmd/tf/descriptor/se_a.py b/deepmd/tf/descriptor/se_a.py index 3dc0d2257d..38e13894df 100644 --- a/deepmd/tf/descriptor/se_a.py +++ b/deepmd/tf/descriptor/se_a.py @@ -141,8 +141,6 @@ class DescrptSeA(DescrptSe): exclude_types : List[List[int]] The excluded pairs of types which have no interaction with each other. For example, `[[0, 1]]` means no interaction between type 0 and type 1. - env_protection: float - Protection parameter to prevent division by zero errors during environment matrix calculations. set_davg_zero Set the shift of embedding net input to zero. activation_function @@ -153,6 +151,8 @@ class DescrptSeA(DescrptSe): Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed multi_task If the model has multi fitting nets to train. + env_protection: float + Protection parameter to prevent division by zero errors during environment matrix calculations. References ---------- @@ -174,7 +174,6 @@ def __init__( seed: Optional[int] = None, type_one_side: bool = True, exclude_types: List[List[int]] = [], - env_protection: float = 0.0, # not implement!! set_davg_zero: bool = False, activation_function: str = "tanh", precision: str = "default", @@ -182,6 +181,7 @@ def __init__( multi_task: bool = False, spin: Optional[Spin] = None, stripped_type_embedding: bool = False, + env_protection: float = 0.0, # not implement!! **kwargs, ) -> None: """Constructor.""" diff --git a/source/tests/pt/model/test_autodiff.py b/source/tests/pt/model/test_autodiff.py index 8631fb074f..d8f4e51dba 100644 --- a/source/tests/pt/model/test_autodiff.py +++ b/source/tests/pt/model/test_autodiff.py @@ -67,7 +67,7 @@ def test( if not test_spin: test_keys = ["energy", "force", "virial"] else: - test_keys = ["energy", "force_real", "force_mag", "virial"] + test_keys = ["energy", "force", "force_mag", "virial"] def np_infer_coord( coord, @@ -114,7 +114,7 @@ def ff_spin(_spin): else: # real force fdf = -finite_difference(ff_coord, coord, delta=delta).squeeze() - rff = np_infer_coord(coord)["force_real"] + rff = np_infer_coord(coord)["force"] np.testing.assert_almost_equal(fdf, rff, decimal=places) # magnetic force fdf = -finite_difference(ff_spin, spin, delta=delta).squeeze() @@ -143,7 +143,7 @@ def test( if not test_spin: test_keys = ["energy", "force", "virial"] else: - test_keys = ["energy", "force_real", "force_mag", "virial"] + test_keys = ["energy", "force", "force_mag", "virial"] def np_infer( new_cell, diff --git a/source/tests/pt/model/test_ener_spin_model.py b/source/tests/pt/model/test_ener_spin_model.py index 10bf327c14..227f5f402c 100644 --- a/source/tests/pt/model/test_ener_spin_model.py +++ b/source/tests/pt/model/test_ener_spin_model.py @@ -101,7 +101,7 @@ def test_output_shape( nframes, nloc = self.coord.shape[:2] torch.testing.assert_close(result["energy"].shape, [nframes, 1]) torch.testing.assert_close(result["atom_energy"].shape, [nframes, nloc, 1]) - torch.testing.assert_close(result["force_real"].shape, [nframes, nloc, 3]) + torch.testing.assert_close(result["force"].shape, [nframes, nloc, 3]) torch.testing.assert_close(result["force_mag"].shape, [nframes, nloc, 3]) torch.testing.assert_close(result["virial"].shape, [nframes, 9]) diff --git a/source/tests/pt/model/test_forward_lower.py b/source/tests/pt/model/test_forward_lower.py index a5c9101e34..553649d4f6 100644 --- a/source/tests/pt/model/test_forward_lower.py +++ b/source/tests/pt/model/test_forward_lower.py @@ -65,7 +65,7 @@ def test( if not test_spin: test_keys = ["energy", "force", "virial"] else: - test_keys = ["energy", "force_real", "force_mag", "virial"] + test_keys = ["energy", "force", "force_mag", "virial"] result_forward = eval_model( self.model, @@ -105,7 +105,7 @@ def test( torch.testing.assert_close( result_forward_lower[key], result_forward[key], rtol=prec, atol=prec ) - elif key in ["force", "force_real", "force_mag"]: + elif key in ["force", "force_mag"]: reduced_vv = reduce_tensor( result_forward_lower[f"extended_{key}"], mapping, natoms ) @@ -165,8 +165,8 @@ def setUp(self): class TestEnergyModelSpinSeA(unittest.TestCase, ForwardLowerTest): def setUp(self): - # still need to figure out why only 1e-6 rtol and atol - self.prec = 1e-6 + # still need to figure out why only 1e-5 rtol and atol + self.prec = 1e-5 model_params = copy.deepcopy(model_spin) self.type_split = False self.test_spin = True diff --git a/source/tests/pt/model/test_permutation.py b/source/tests/pt/model/test_permutation.py index b843411191..1f6d054034 100644 --- a/source/tests/pt/model/test_permutation.py +++ b/source/tests/pt/model/test_permutation.py @@ -238,7 +238,7 @@ def test( if not test_spin: test_keys = ["energy", "force", "virial"] else: - test_keys = ["energy", "force_real", "force_mag", "virial"] + test_keys = ["energy", "force", "force_mag", "virial"] result_0 = eval_model( self.model, coord.unsqueeze(0), @@ -259,7 +259,7 @@ def test( for key in test_keys: if key in ["energy"]: torch.testing.assert_close(ret0[key], ret1[key], rtol=prec, atol=prec) - elif key in ["force", "force_real", "force_mag"]: + elif key in ["force", "force_mag"]: torch.testing.assert_close( ret0[key][idx_perm], ret1[key], rtol=prec, atol=prec ) diff --git a/source/tests/pt/model/test_rot.py b/source/tests/pt/model/test_rot.py index 53ea88bb7b..39818849bf 100644 --- a/source/tests/pt/model/test_rot.py +++ b/source/tests/pt/model/test_rot.py @@ -43,7 +43,7 @@ def test( if not test_spin: test_keys = ["energy", "force", "virial"] else: - test_keys = ["energy", "force_real", "force_mag", "virial"] + test_keys = ["energy", "force", "force_mag", "virial"] rmat = torch.tensor(special_ortho_group.rvs(3), dtype=dtype, device=env.DEVICE) # rotate only coord and shift to the center of cell @@ -68,7 +68,7 @@ def test( for key in test_keys: if key in ["energy"]: torch.testing.assert_close(ret0[key], ret1[key], rtol=prec, atol=prec) - elif key in ["force", "force_real", "force_mag"]: + elif key in ["force", "force_mag"]: torch.testing.assert_close( torch.matmul(ret0[key], rmat), ret1[key], rtol=prec, atol=prec ) @@ -114,7 +114,7 @@ def test( for key in test_keys: if key in ["energy"]: torch.testing.assert_close(ret0[key], ret1[key], rtol=prec, atol=prec) - elif key in ["force", "force_real", "force_mag"]: + elif key in ["force", "force_mag"]: torch.testing.assert_close( torch.matmul(ret0[key], rmat), ret1[key], rtol=prec, atol=prec ) diff --git a/source/tests/pt/model/test_smooth.py b/source/tests/pt/model/test_smooth.py index be6e3368be..f40a0a29b9 100644 --- a/source/tests/pt/model/test_smooth.py +++ b/source/tests/pt/model/test_smooth.py @@ -73,7 +73,7 @@ def test( if not test_spin: test_keys = ["energy", "force", "virial"] else: - test_keys = ["energy", "force_real", "force_mag", "virial"] + test_keys = ["energy", "force", "force_mag", "virial"] result_0 = eval_model( self.model, @@ -114,7 +114,7 @@ def compare(ret0, ret1): torch.testing.assert_close( ret0[key], ret1[key], rtol=rprec, atol=aprec ) - elif key in ["force", "force_real", "force_mag"]: + elif key in ["force", "force_mag"]: # plus 1. to avoid the divided-by-zero issue torch.testing.assert_close( 1.0 + ret0[key], 1.0 + ret1[key], rtol=rprec, atol=aprec diff --git a/source/tests/pt/model/test_trans.py b/source/tests/pt/model/test_trans.py index 656f9f7380..475cc301c7 100644 --- a/source/tests/pt/model/test_trans.py +++ b/source/tests/pt/model/test_trans.py @@ -46,7 +46,7 @@ def test( if not test_spin: test_keys = ["energy", "force", "virial"] else: - test_keys = ["energy", "force_real", "force_mag", "virial"] + test_keys = ["energy", "force", "force_mag", "virial"] result_0 = eval_model( self.model, coord.unsqueeze(0), @@ -65,7 +65,7 @@ def test( ret1 = {key: result_1[key].squeeze(0) for key in test_keys} prec = 1e-10 for key in test_keys: - if key in ["energy", "force", "force_real", "force_mag"]: + if key in ["energy", "force", "force_mag"]: torch.testing.assert_close(ret0[key], ret1[key], rtol=prec, atol=prec) elif key == "virial": if not hasattr(self, "test_virial") or self.test_virial: diff --git a/source/tests/pt/model/test_unused_params.py b/source/tests/pt/model/test_unused_params.py index c20a5f1dc5..f1d4591798 100644 --- a/source/tests/pt/model/test_unused_params.py +++ b/source/tests/pt/model/test_unused_params.py @@ -64,14 +64,9 @@ def _test_unused(self, model_params): coord = torch.matmul(coord, cell) atype = torch.IntTensor([0, 0, 0, 1, 1]).to(env.DEVICE) idx_perm = [1, 0, 4, 3, 2] - e0, f0, v0 = eval_model( - self.model, coord.unsqueeze(0), cell.unsqueeze(0), atype - ) - ret0 = { - "energy": e0.squeeze(0), - "force": f0.squeeze(0), - "virial": v0.squeeze(0), - } + result_0 = eval_model(self.model, coord.unsqueeze(0), cell.unsqueeze(0), atype) + test_keys = ["energy", "force", "virial"] + ret0 = {key: result_0[key].squeeze(0) for key in test_keys} # use computation graph to find all contributing tensors def get_contributing_params(y, top_level=True):