Skip to content

Commit

Permalink
chore(pt): make comm_dict for dpa2 noncompulsory when nghost is 0 (#…
Browse files Browse the repository at this point in the history 
…4144)

Construct and pass a mapping instead when there is no ghost atom.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## 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.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
  • Loading branch information
@njzjz
njzjz committed Sep 19, 2024
1 parent a9aad68 commit e1b6aec
Show file tree
Hide file tree
Showing 4 changed files with 243 additions and 9 deletions.
16 changes: 11 additions & 5 deletions source/api_cc/src/DeepPotPT.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -164,11 +164,12 @@ void DeepPotPT::compute(ENERGYVTYPE& ener,
std::vector<std::int64_t> 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<torch::Tensor> 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);
Expand Down Expand Up @@ -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<torch::Tensor> optional_tensor;
c10::optional<torch::Tensor> fparam_tensor;
if (!fparam.empty()) {
fparam_tensor =
Expand All @@ -219,15 +225,15 @@ void DeepPotPT::compute(ENERGYVTYPE& ener,
.to(device);
}
c10::Dict<c10::IValue, c10::IValue> 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");
Expand Down
2 changes: 1 addition & 1 deletion source/api_cc/tests/test_deeppot_a_fparam_aparam_pt.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -95,7 +95,7 @@ class TestInferDeepPotAFParamAParamPt : public ::testing::Test {
}
};

void TearDown() override { remove("fparam_aparam.pb"); };
void TearDown() override {};
};

TYPED_TEST_SUITE(TestInferDeepPotAFParamAParamPt, ValueTypes);
Expand Down
232 changes: 230 additions & 2 deletions source/api_cc/tests/test_deeppot_dpa_pt.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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()=}")
Expand Down Expand Up @@ -95,7 +95,7 @@ class TestInferDeepPotDpaPt : public ::testing::Test {
}
};

void TearDown() override { remove("deeppot.pb"); };
void TearDown() override {};
};

TYPED_TEST_SUITE(TestInferDeepPotDpaPt, ValueTypes);
Expand Down Expand Up @@ -163,3 +163,231 @@ TYPED_TEST(TestInferDeepPotDpaPt, cpu_build_nlist_atomic) {
EXPECT_LT(fabs(atom_vir[ii] - expected_v[ii]), EPSILON);
}
}

template <class VALUETYPE>
class TestInferDeepPotDpaPtNopbc : public ::testing::Test {
protected:
std::vector<VALUETYPE> 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<int> atype = {0, 1, 1, 0, 1, 1};
std::vector<VALUETYPE> 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<VALUETYPE> expected_e = {
-95.13216447995296, -188.10146505781867, -187.74742451023172,
-94.73864717001219, -187.76956603003393, -187.76904550434332};
std::vector<VALUETYPE> 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<VALUETYPE> 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<VALUETYPE> 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<VALUETYPE>& coord = this->coord;
std::vector<int>& atype = this->atype;
std::vector<VALUETYPE>& box = this->box;
std::vector<VALUETYPE>& expected_e = this->expected_e;
std::vector<VALUETYPE>& expected_f = this->expected_f;
std::vector<VALUETYPE>& expected_v = this->expected_v;
int& natoms = this->natoms;
double& expected_tot_e = this->expected_tot_e;
std::vector<VALUETYPE>& expected_tot_v = this->expected_tot_v;
deepmd::DeepPot& dp = this->dp;
double ener;
std::vector<VALUETYPE> 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<VALUETYPE>& coord = this->coord;
std::vector<int>& atype = this->atype;
std::vector<VALUETYPE>& box = this->box;
std::vector<VALUETYPE>& expected_e = this->expected_e;
std::vector<VALUETYPE>& expected_f = this->expected_f;
std::vector<VALUETYPE>& expected_v = this->expected_v;
int& natoms = this->natoms;
double& expected_tot_e = this->expected_tot_e;
std::vector<VALUETYPE>& expected_tot_v = this->expected_tot_v;
deepmd::DeepPot& dp = this->dp;
double ener;
std::vector<VALUETYPE> 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<VALUETYPE>& coord = this->coord;
std::vector<int>& atype = this->atype;
std::vector<VALUETYPE>& box = this->box;
std::vector<VALUETYPE>& expected_e = this->expected_e;
std::vector<VALUETYPE>& expected_f = this->expected_f;
std::vector<VALUETYPE>& expected_v = this->expected_v;
int& natoms = this->natoms;
double& expected_tot_e = this->expected_tot_e;
std::vector<VALUETYPE>& expected_tot_v = this->expected_tot_v;
deepmd::DeepPot& dp = this->dp;
double ener;
std::vector<VALUETYPE> force, virial;

std::vector<std::vector<int> > 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<int> ilist(natoms), numneigh(natoms);
std::vector<int*> 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<VALUETYPE>& coord = this->coord;
std::vector<int>& atype = this->atype;
std::vector<VALUETYPE>& box = this->box;
std::vector<VALUETYPE>& expected_e = this->expected_e;
std::vector<VALUETYPE>& expected_f = this->expected_f;
std::vector<VALUETYPE>& expected_v = this->expected_v;
int& natoms = this->natoms;
double& expected_tot_e = this->expected_tot_e;
std::vector<VALUETYPE>& expected_tot_v = this->expected_tot_v;
deepmd::DeepPot& dp = this->dp;
double ener;
std::vector<VALUETYPE> force, virial, atom_ener, atom_vir;

std::vector<std::vector<int> > 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<int> ilist(natoms), numneigh(natoms);
std::vector<int*> 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);
}
}
2 changes: 1 addition & 1 deletion source/api_cc/tests/test_deeppot_pt.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,7 @@ class TestInferDeepPotAPt : public ::testing::Test {
}
};

void TearDown() override { remove("deeppot.pb"); };
void TearDown() override {};
};

TYPED_TEST_SUITE(TestInferDeepPotAPt, ValueTypes);
Expand Down

0 comments on commit e1b6aec

Please sign in to comment.