From e1b6aec8c5c576bb154be4177e566dde239d6707 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Thu, 19 Sep 2024 08:38:11 -0400 Subject: [PATCH] chore(pt): make `comm_dict` for dpa2 noncompulsory when nghost is 0 (#4144) Construct and pass a mapping instead when there is no ghost atom. ## Summary by CodeRabbit - **New Features** - Introduced a new test class to evaluate the `DeepPot` model without periodic boundary conditions, enhancing testing coverage. - **Bug Fixes** - Corrected the order of returned values in existing tests for clarity and accuracy. - **Documentation** - Updated the cleanup process in test classes to improve test execution consistency. - **Chores** - Simplified the `TearDown` methods in multiple test classes, removing unnecessary file deletion steps. --------- Signed-off-by: Jinzhe Zeng --- source/api_cc/src/DeepPotPT.cc | 16 +- .../tests/test_deeppot_a_fparam_aparam_pt.cc | 2 +- source/api_cc/tests/test_deeppot_dpa_pt.cc | 232 +++++++++++++++++- source/api_cc/tests/test_deeppot_pt.cc | 2 +- 4 files changed, 243 insertions(+), 9 deletions(-) diff --git a/source/api_cc/src/DeepPotPT.cc b/source/api_cc/src/DeepPotPT.cc index 3d406879de..c03576635a 100644 --- a/source/api_cc/src/DeepPotPT.cc +++ b/source/api_cc/src/DeepPotPT.cc @@ -164,11 +164,12 @@ void DeepPotPT::compute(ENERGYVTYPE& ener, std::vector atype_64(datype.begin(), datype.end()); at::Tensor atype_Tensor = torch::from_blob(atype_64.data(), {1, nall_real}, int_option).to(device); + c10::optional mapping_tensor; if (ago == 0) { nlist_data.copy_from_nlist(lmp_list); nlist_data.shuffle_exclude_empty(fwd_map); nlist_data.padding(); - if (do_message_passing == 1) { + if (do_message_passing == 1 && nghost > 0) { int nswap = lmp_list.nswap; torch::Tensor sendproc_tensor = torch::from_blob(lmp_list.sendproc, {nswap}, int32_option); @@ -196,11 +197,16 @@ void DeepPotPT::compute(ENERGYVTYPE& ener, comm_dict.insert("recv_num", recvnum_tensor); comm_dict.insert("communicator", communicator_tensor); } + if (do_message_passing == 1 && nghost == 0) { + // for the situation that no ghost atoms (e.g. serial nopbc) + // set the mapping arange(nloc) is enough + auto option = torch::TensorOptions().device(device).dtype(torch::kInt64); + mapping_tensor = at::arange(nloc_real, option).unsqueeze(0); + } } at::Tensor firstneigh = createNlistTensor(nlist_data.jlist); firstneigh_tensor = firstneigh.to(torch::kInt64).to(device); bool do_atom_virial_tensor = atomic; - c10::optional optional_tensor; c10::optional fparam_tensor; if (!fparam.empty()) { fparam_tensor = @@ -219,15 +225,15 @@ void DeepPotPT::compute(ENERGYVTYPE& ener, .to(device); } c10::Dict outputs = - (do_message_passing == 1) + (do_message_passing == 1 && nghost > 0) ? module .run_method("forward_lower", coord_wrapped_Tensor, atype_Tensor, - firstneigh_tensor, optional_tensor, fparam_tensor, + firstneigh_tensor, mapping_tensor, fparam_tensor, aparam_tensor, do_atom_virial_tensor, comm_dict) .toGenericDict() : module .run_method("forward_lower", coord_wrapped_Tensor, atype_Tensor, - firstneigh_tensor, optional_tensor, fparam_tensor, + firstneigh_tensor, mapping_tensor, fparam_tensor, aparam_tensor, do_atom_virial_tensor) .toGenericDict(); c10::IValue energy_ = outputs.at("energy"); diff --git a/source/api_cc/tests/test_deeppot_a_fparam_aparam_pt.cc b/source/api_cc/tests/test_deeppot_a_fparam_aparam_pt.cc index dfaf0abc06..e394f86b2a 100644 --- a/source/api_cc/tests/test_deeppot_a_fparam_aparam_pt.cc +++ b/source/api_cc/tests/test_deeppot_a_fparam_aparam_pt.cc @@ -95,7 +95,7 @@ class TestInferDeepPotAFParamAParamPt : public ::testing::Test { } }; - void TearDown() override { remove("fparam_aparam.pb"); }; + void TearDown() override {}; }; TYPED_TEST_SUITE(TestInferDeepPotAFParamAParamPt, ValueTypes); diff --git a/source/api_cc/tests/test_deeppot_dpa_pt.cc b/source/api_cc/tests/test_deeppot_dpa_pt.cc index 94cadeaf19..7bf4cbf376 100644 --- a/source/api_cc/tests/test_deeppot_dpa_pt.cc +++ b/source/api_cc/tests/test_deeppot_dpa_pt.cc @@ -36,7 +36,7 @@ class TestInferDeepPotDpaPt : public ::testing::Test { // atype = np.array([0, 1, 1, 0, 1, 1]) // box = np.array([13., 0., 0., 0., 13., 0., 0., 0., 13.]).reshape(1, -1) // dp = DeepPot("deeppot_dpa.pth") - // e, v, f, ae, av = dp.eval(coord, box, atype, atomic=True) + // e, f, v, ae, av = dp.eval(coord, box, atype, atomic=True) // np.set_printoptions(precision=16) // print(f"{e.ravel()=} {v.ravel()=} {f.ravel()=} {ae.ravel()=} // {av.ravel()=}") @@ -95,7 +95,7 @@ class TestInferDeepPotDpaPt : public ::testing::Test { } }; - void TearDown() override { remove("deeppot.pb"); }; + void TearDown() override {}; }; TYPED_TEST_SUITE(TestInferDeepPotDpaPt, ValueTypes); @@ -163,3 +163,231 @@ TYPED_TEST(TestInferDeepPotDpaPt, cpu_build_nlist_atomic) { EXPECT_LT(fabs(atom_vir[ii] - expected_v[ii]), EPSILON); } } + +template +class TestInferDeepPotDpaPtNopbc : public ::testing::Test { + protected: + std::vector coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74, + 00.25, 3.32, 1.68, 3.36, 3.00, 1.81, + 3.51, 2.51, 2.60, 4.27, 3.22, 1.56}; + std::vector atype = {0, 1, 1, 0, 1, 1}; + std::vector box = {}; + // Generated by the following Python code: + // import numpy as np + // from deepmd.infer import DeepPot + // coord = np.array([ + // 12.83, 2.56, 2.18, 12.09, 2.87, 2.74, + // 00.25, 3.32, 1.68, 3.36, 3.00, 1.81, + // 3.51, 2.51, 2.60, 4.27, 3.22, 1.56 + // ]).reshape(1, -1) + // atype = np.array([0, 1, 1, 0, 1, 1]) + // box = None + // dp = DeepPot("deeppot_dpa.pth") + // e, f, v, ae, av = dp.eval(coord, box, atype, atomic=True) + // np.set_printoptions(precision=16) + // print(f"{e.ravel()=} {v.ravel()=} {f.ravel()=} {ae.ravel()=} + // {av.ravel()=}") + + std::vector expected_e = { + -95.13216447995296, -188.10146505781867, -187.74742451023172, + -94.73864717001219, -187.76956603003393, -187.76904550434332}; + std::vector expected_f = { + 0.7486830600282869, -0.240322915088127, -0.3943366458127905, + -0.1776248813665344, 0.2359143394202788, 0.4210018319063822, + -0.2368532809002255, 0.0291156803500336, -0.0219651427265617, + -1.407280069394403, 0.4932116549421467, -0.9482072853582465, + -0.1501958909452974, -0.9720722611839484, 1.5128172910814666, + 1.2232710625781733, 0.4541535015596165, -0.569310049090249}; + std::vector expected_v = { + 1.4724482801774368e+00, -1.8952544175284314e-01, -2.0502896614522359e-01, + -2.0361724110178425e-01, 5.4221646102123211e-02, 8.7963957026666373e-02, + -1.3233356224791937e-01, 8.3907068051133571e-02, 1.6072164570432412e-01, + 2.2913216241740741e+00, -6.0712170533586352e-02, 1.2802395909429765e-01, + 6.9581050483420448e-03, 2.0894022035588655e-02, 4.3408316864598340e-02, + -1.4144392402206662e-03, 3.6852652738654124e-02, 7.7149761552687490e-02, + 5.6814285976509526e-01, -7.0738211182030164e-02, 5.4514470128648518e-02, + -7.1339324275474125e-02, 9.8158535704203354e-03, -8.3431069537701560e-03, + 5.4072790262097083e-02, -8.1976736911977682e-03, 7.6505804915597275e-03, + 1.6869950835783332e-01, 2.1880432930426963e-02, 1.0308234746703970e-01, + 9.1015395953307099e-02, 7.1788910181538768e-02, -1.4119552688428305e-01, + -1.4977320631771729e-01, -1.0982955047012899e-01, 2.3324521962640055e-01, + 8.1569862372597679e-01, 6.2848559999917952e-02, -4.5341405643671506e-02, + -3.9134119664198064e-01, 4.1651372430088562e-01, -5.8173709994663803e-01, + 6.6155672230934037e-01, -6.4774042800560672e-01, 9.0924772156749301e-01, + 2.0503134548416586e+00, 1.9684008914564011e-01, -3.1711040533580070e-01, + 5.2891751962511613e-01, 8.7385258358844808e-02, -1.5487618319904839e-01, + -7.1396830520028809e-01, -1.0977171171532918e-01, 1.9792085656111236e-01}; + int natoms; + double expected_tot_e; + std::vector expected_tot_v; + + deepmd::DeepPot dp; + + void SetUp() override { + dp.init("../../tests/infer/deeppot_dpa.pth"); + + natoms = expected_e.size(); + EXPECT_EQ(natoms * 3, expected_f.size()); + EXPECT_EQ(natoms * 9, expected_v.size()); + expected_tot_e = 0.; + expected_tot_v.resize(9); + std::fill(expected_tot_v.begin(), expected_tot_v.end(), 0.); + for (int ii = 0; ii < natoms; ++ii) { + expected_tot_e += expected_e[ii]; + } + for (int ii = 0; ii < natoms; ++ii) { + for (int dd = 0; dd < 9; ++dd) { + expected_tot_v[dd] += expected_v[ii * 9 + dd]; + } + } + }; + + void TearDown() override {}; +}; + +TYPED_TEST_SUITE(TestInferDeepPotDpaPtNopbc, ValueTypes); + +TYPED_TEST(TestInferDeepPotDpaPtNopbc, cpu_build_nlist) { + using VALUETYPE = TypeParam; + std::vector& coord = this->coord; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_e = this->expected_e; + std::vector& expected_f = this->expected_f; + std::vector& expected_v = this->expected_v; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e; + std::vector& expected_tot_v = this->expected_tot_v; + deepmd::DeepPot& dp = this->dp; + double ener; + std::vector force, virial; + dp.compute(ener, force, virial, coord, atype, box); + + EXPECT_EQ(force.size(), natoms * 3); + EXPECT_EQ(virial.size(), 9); + + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + } + for (int ii = 0; ii < 3 * 3; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } +} + +TYPED_TEST(TestInferDeepPotDpaPtNopbc, cpu_build_nlist_atomic) { + using VALUETYPE = TypeParam; + std::vector& coord = this->coord; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_e = this->expected_e; + std::vector& expected_f = this->expected_f; + std::vector& expected_v = this->expected_v; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e; + std::vector& expected_tot_v = this->expected_tot_v; + deepmd::DeepPot& dp = this->dp; + double ener; + std::vector force, virial, atom_ener, atom_vir; + dp.compute(ener, force, virial, atom_ener, atom_vir, coord, atype, box); + + EXPECT_EQ(force.size(), natoms * 3); + EXPECT_EQ(virial.size(), 9); + EXPECT_EQ(atom_ener.size(), natoms); + EXPECT_EQ(atom_vir.size(), natoms * 9); + + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + } + for (int ii = 0; ii < 3 * 3; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } + for (int ii = 0; ii < natoms; ++ii) { + EXPECT_LT(fabs(atom_ener[ii] - expected_e[ii]), EPSILON); + } + for (int ii = 0; ii < natoms * 9; ++ii) { + EXPECT_LT(fabs(atom_vir[ii] - expected_v[ii]), EPSILON); + } +} + +TYPED_TEST(TestInferDeepPotDpaPtNopbc, cpu_lmp_nlist) { + using VALUETYPE = TypeParam; + std::vector& coord = this->coord; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_e = this->expected_e; + std::vector& expected_f = this->expected_f; + std::vector& expected_v = this->expected_v; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e; + std::vector& expected_tot_v = this->expected_tot_v; + deepmd::DeepPot& dp = this->dp; + double ener; + std::vector force, virial; + + std::vector > nlist_data = { + {1, 2, 3, 4, 5}, {0, 2, 3, 4, 5}, {0, 1, 3, 4, 5}, + {0, 1, 2, 4, 5}, {0, 1, 2, 3, 5}, {0, 1, 2, 3, 4}}; + std::vector ilist(natoms), numneigh(natoms); + std::vector firstneigh(natoms); + deepmd::InputNlist inlist(natoms, &ilist[0], &numneigh[0], &firstneigh[0]); + convert_nlist(inlist, nlist_data); + dp.compute(ener, force, virial, coord, atype, box, 0, inlist, 0); + + EXPECT_EQ(force.size(), natoms * 3); + EXPECT_EQ(virial.size(), 9); + + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + } + for (int ii = 0; ii < 3 * 3; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } +} + +TYPED_TEST(TestInferDeepPotDpaPtNopbc, cpu_lmp_nlist_atomic) { + using VALUETYPE = TypeParam; + std::vector& coord = this->coord; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_e = this->expected_e; + std::vector& expected_f = this->expected_f; + std::vector& expected_v = this->expected_v; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e; + std::vector& expected_tot_v = this->expected_tot_v; + deepmd::DeepPot& dp = this->dp; + double ener; + std::vector force, virial, atom_ener, atom_vir; + + std::vector > nlist_data = { + {1, 2, 3, 4, 5}, {0, 2, 3, 4, 5}, {0, 1, 3, 4, 5}, + {0, 1, 2, 4, 5}, {0, 1, 2, 3, 5}, {0, 1, 2, 3, 4}}; + std::vector ilist(natoms), numneigh(natoms); + std::vector firstneigh(natoms); + deepmd::InputNlist inlist(natoms, &ilist[0], &numneigh[0], &firstneigh[0]); + convert_nlist(inlist, nlist_data); + dp.compute(ener, force, virial, atom_ener, atom_vir, coord, atype, box, 0, + inlist, 0); + + EXPECT_EQ(force.size(), natoms * 3); + EXPECT_EQ(virial.size(), 9); + EXPECT_EQ(atom_ener.size(), natoms); + EXPECT_EQ(atom_vir.size(), natoms * 9); + + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + } + for (int ii = 0; ii < 3 * 3; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } + for (int ii = 0; ii < natoms; ++ii) { + EXPECT_LT(fabs(atom_ener[ii] - expected_e[ii]), EPSILON); + } + for (int ii = 0; ii < natoms * 9; ++ii) { + EXPECT_LT(fabs(atom_vir[ii] - expected_v[ii]), EPSILON); + } +} diff --git a/source/api_cc/tests/test_deeppot_pt.cc b/source/api_cc/tests/test_deeppot_pt.cc index 852aef611e..af5b533a5f 100644 --- a/source/api_cc/tests/test_deeppot_pt.cc +++ b/source/api_cc/tests/test_deeppot_pt.cc @@ -74,7 +74,7 @@ class TestInferDeepPotAPt : public ::testing::Test { } }; - void TearDown() override { remove("deeppot.pb"); }; + void TearDown() override {}; }; TYPED_TEST_SUITE(TestInferDeepPotAPt, ValueTypes);