diff --git a/source/api_cc/include/DeepSpinTF.h b/source/api_cc/include/DeepSpinTF.h index 547f8140cf..e628c565c5 100644 --- a/source/api_cc/include/DeepSpinTF.h +++ b/source/api_cc/include/DeepSpinTF.h @@ -259,14 +259,16 @@ class DeepSpinTF : public DeepSpinBackend { const int nghost, const std::vector& spin, const int numb_types, - const int numb_types_spin); + const int numb_types_spin, + const int nframes); template void extend_nlist(std::vector& extend_dcoord, std::vector& extend_atype, const std::vector& dcoord_, const std::vector& dspin_, - const std::vector& datype_); + const std::vector& datype_, + const int nframes); void cum_sum(std::map&, std::map&); diff --git a/source/api_cc/src/DeepSpinTF.cc b/source/api_cc/src/DeepSpinTF.cc index 1981f025e3..ed9ca849b9 100644 --- a/source/api_cc/src/DeepSpinTF.cc +++ b/source/api_cc/src/DeepSpinTF.cc @@ -176,9 +176,23 @@ static void run_model( for (size_t ii = 0; ii < static_cast(nframes) * nall * 3; ++ii) { dforce[ii] = of(ii); } + if (output_ae.NumElements() % nframes != 0) { + throw deepmd::deepmd_exception( + "TensorFlow atomic-energy output is not divisible by nframes."); + } + const size_t nloc_energy = output_ae.NumElements() / nframes; + if (nloc_energy > nall) { + throw deepmd::deepmd_exception( + "TensorFlow atomic-energy output has more atoms than the extended " + "DeepSpin system."); + } + // Spin models emit atomic energies for physical atoms only. The extended + // virtual atoms sort after all physical types, so keep their slots zero and + // use the actual output width as the per-frame source stride. for (int ii = 0; ii < nframes; ++ii) { - for (int jj = 0; jj < nloc; ++jj) { - datom_energy[ii * nall + jj] = oae(ii * nloc + jj); + for (size_t jj = 0; jj < nloc_energy; ++jj) { + datom_energy[static_cast(ii) * nall + jj] = + oae(static_cast(ii) * nloc_energy + jj); } } for (size_t ii = 0; ii < static_cast(nframes) * nall * 9; ++ii) { @@ -602,7 +616,7 @@ void DeepSpinTF::compute(ENERGYVTYPE& dener, std::vector extend_dcoord; std::vector extend_atype; - extend_nlist(extend_dcoord, extend_atype, dcoord_, dspin_, datype_); + extend_nlist(extend_dcoord, extend_atype, dcoord_, dspin_, datype_, nframes); atommap = deepmd::AtomMap(extend_atype.begin(), extend_atype.end()); @@ -636,13 +650,20 @@ void DeepSpinTF::compute(ENERGYVTYPE& dener, // backward force and mag. dforce_.resize(static_cast(nframes) * nloc * 3); dforce_mag_.resize(static_cast(nframes) * nloc * 3); - for (int ii = 0; ii < nloc; ++ii) { - for (int dd = 0; dd < 3; ++dd) { - dforce_[3 * ii + dd] = dforce_tmp[3 * ii + dd]; - if (datype_[ii] < ntypes_spin) { - dforce_mag_[3 * ii + dd] = dforce_tmp[3 * (ii + nloc) + dd]; - } else { - dforce_mag_[3 * ii + dd] = 0.0; + const size_t extend_nall = extend_atype.size(); + for (int ff = 0; ff < nframes; ++ff) { + for (int ii = 0; ii < nloc; ++ii) { + const size_t output_atom = static_cast(ff) * nloc + ii; + const size_t extended_atom = static_cast(ff) * extend_nall + ii; + for (int dd = 0; dd < 3; ++dd) { + dforce_[output_atom * 3 + dd] = dforce_tmp[extended_atom * 3 + dd]; + if (datype_[ii] < ntypes_spin) { + const size_t virtual_atom = + static_cast(ff) * extend_nall + ii + nloc; + dforce_mag_[output_atom * 3 + dd] = dforce_tmp[virtual_atom * 3 + dd]; + } else { + dforce_mag_[output_atom * 3 + dd] = 0.0; + } } } } @@ -735,7 +756,7 @@ void DeepSpinTF::compute(ENERGYVTYPE& dener, extend(extend_inum, extend_ilist, extend_numneigh, extend_neigh, extend_firstneigh, extend_dcoord, extend_dtype, extend_nghost, new_idx_map, old_idx_map, lmp_list, dcoord_, datype_, nghost, dspin_, - ntypes, ntypes_spin); + ntypes, ntypes_spin, nframes); InputNlist extend_lmp_list(extend_inum, &extend_ilist[0], &extend_numneigh[0], &extend_firstneigh[0]); extend_lmp_list.set_mask(lmp_list.mask); @@ -808,22 +829,30 @@ void DeepSpinTF::compute(ENERGYVTYPE& dener, dforce_mag_.resize(static_cast(nframes) * nall * 3); datom_energy_.resize(static_cast(nframes) * nall); datom_virial_.resize(static_cast(nframes) * nall * 9); - for (int ii = 0; ii < nall; ++ii) { - int new_idx = new_idx_map[ii]; - for (int dd = 0; dd < 3; ++dd) { - dforce_[3 * ii + dd] = dforce_tmp[3 * new_idx + dd]; - datom_energy_[ii] = datom_energy_tmp[new_idx]; - - if (datype_[ii] < ntypes_spin && ii < nloc) { - dforce_mag_[3 * ii + dd] = dforce_tmp[3 * (new_idx + nloc) + dd]; - } else if (datype_[ii] < ntypes_spin) { - dforce_mag_[3 * ii + dd] = dforce_tmp[3 * (new_idx + nghost) + dd]; - } else { - dforce_mag_[3 * ii + dd] = 0.0; + const size_t extended_nall = fwd_map.size(); + for (int ff = 0; ff < nframes; ++ff) { + for (int ii = 0; ii < nall; ++ii) { + const int new_idx = new_idx_map[ii]; + const size_t output_atom = static_cast(ff) * nall + ii; + const size_t extended_atom = + static_cast(ff) * extended_nall + new_idx; + datom_energy_[output_atom] = datom_energy_tmp[extended_atom]; + for (int dd = 0; dd < 3; ++dd) { + dforce_[output_atom * 3 + dd] = dforce_tmp[extended_atom * 3 + dd]; + + if (datype_[ii] < ntypes_spin) { + const int virtual_idx = new_idx + (ii < nloc ? nloc : nghost); + const size_t virtual_atom = + static_cast(ff) * extended_nall + virtual_idx; + dforce_mag_[output_atom * 3 + dd] = dforce_tmp[virtual_atom * 3 + dd]; + } else { + dforce_mag_[output_atom * 3 + dd] = 0.0; + } + } + for (int dd = 0; dd < 9; ++dd) { + datom_virial_[output_atom * 9 + dd] = + datom_virial_tmp[extended_atom * 9 + dd]; } - } - for (int dd = 0; dd < 9; ++dd) { - datom_virial_[ii * 9 + dd] = datom_virial_tmp[new_idx * 9 + dd]; } } } @@ -1004,7 +1033,8 @@ void DeepSpinTF::extend(int& extend_inum, const int nghost, const std::vector& spin, const int numb_types, - const int numb_types_spin) { + const int numb_types_spin, + const int nframes) { extend_ilist.clear(); extend_numneigh.clear(); extend_neigh.clear(); @@ -1021,8 +1051,10 @@ void DeepSpinTF::extend(int& extend_inum, get_vector(spin_norm, "spin_attr/spin_norm"); } - int nall = dcoord.size() / 3; + int nall = atype.size(); int nloc = nall - nghost; + assert(static_cast(nframes) * nall * 3 == dcoord.size()); + assert(dcoord.size() == spin.size()); assert(nloc == lmp_list.inum); // record numb_types_real and nloc_virt @@ -1128,26 +1160,36 @@ void DeepSpinTF::extend(int& extend_inum, } // extend coord - extend_dcoord.resize(static_cast(extend_nall) * 3); - for (int ii = 0; ii < nloc; ii++) { - for (int jj = 0; jj < 3; jj++) { - extend_dcoord[new_idx_map[ii] * 3 + jj] = dcoord[ii * 3 + jj]; - if (atype[ii] < numb_types_spin) { - double temp_dcoord = dcoord[ii * 3 + jj] + spin[ii * 3 + jj] / - spin_norm[atype[ii]] * - virtual_len[atype[ii]]; - extend_dcoord[(new_idx_map[ii] + nloc) * 3 + jj] = temp_dcoord; + extend_dcoord.resize(static_cast(nframes) * extend_nall * 3); + for (int ff = 0; ff < nframes; ++ff) { + const size_t input_offset = static_cast(ff) * nall * 3; + const size_t output_offset = static_cast(ff) * extend_nall * 3; + for (int ii = 0; ii < nloc; ii++) { + for (int jj = 0; jj < 3; jj++) { + extend_dcoord[output_offset + new_idx_map[ii] * 3 + jj] = + dcoord[input_offset + ii * 3 + jj]; + if (atype[ii] < numb_types_spin) { + const VALUETYPE temp_dcoord = dcoord[input_offset + ii * 3 + jj] + + spin[input_offset + ii * 3 + jj] / + spin_norm[atype[ii]] * + virtual_len[atype[ii]]; + extend_dcoord[output_offset + (new_idx_map[ii] + nloc) * 3 + jj] = + temp_dcoord; + } } } - } - for (int ii = nloc; ii < nall; ii++) { - for (int jj = 0; jj < 3; jj++) { - extend_dcoord[new_idx_map[ii] * 3 + jj] = dcoord[ii * 3 + jj]; - if (atype[ii] < numb_types_spin) { - double temp_dcoord = dcoord[ii * 3 + jj] + spin[ii * 3 + jj] / - spin_norm[atype[ii]] * - virtual_len[atype[ii]]; - extend_dcoord[(new_idx_map[ii] + nghost) * 3 + jj] = temp_dcoord; + for (int ii = nloc; ii < nall; ii++) { + for (int jj = 0; jj < 3; jj++) { + extend_dcoord[output_offset + new_idx_map[ii] * 3 + jj] = + dcoord[input_offset + ii * 3 + jj]; + if (atype[ii] < numb_types_spin) { + const VALUETYPE temp_dcoord = dcoord[input_offset + ii * 3 + jj] + + spin[input_offset + ii * 3 + jj] / + spin_norm[atype[ii]] * + virtual_len[atype[ii]]; + extend_dcoord[output_offset + (new_idx_map[ii] + nghost) * 3 + jj] = + temp_dcoord; + } } } } @@ -1183,7 +1225,8 @@ template void DeepSpinTF::extend( const int nghost, const std::vector& spin, const int numb_types, - const int numb_types_spin); + const int numb_types_spin, + const int nframes); template void DeepSpinTF::extend( int& extend_inum, @@ -1202,14 +1245,16 @@ template void DeepSpinTF::extend( const int nghost, const std::vector& spin, const int numb_types, - const int numb_types_spin); + const int numb_types_spin, + const int nframes); template void DeepSpinTF::extend_nlist(std::vector& extend_dcoord, std::vector& extend_atype, const std::vector& dcoord_, const std::vector& dspin_, - const std::vector& datype_) { + const std::vector& datype_, + const int nframes) { if (dtype == tensorflow::DT_DOUBLE) { get_vector(virtual_len, "spin_attr/virtual_len"); get_vector(spin_norm, "spin_attr/spin_norm"); @@ -1228,20 +1273,27 @@ void DeepSpinTF::extend_nlist(std::vector& extend_dcoord, } } int extend_nall = nloc + nloc_spin; - extend_dcoord.resize(static_cast(extend_nall) * 3); + assert(static_cast(nframes) * nloc * 3 == dcoord_.size()); + assert(dcoord_.size() == dspin_.size()); + extend_dcoord.resize(static_cast(nframes) * extend_nall * 3); extend_atype.resize(extend_nall); for (int ii = 0; ii < nloc; ii++) { extend_atype[ii] = datype_[ii]; if (datype_[ii] < ntypes_spin) { extend_atype[ii + nloc] = datype_[ii] + ntypes - ntypes_spin; } - for (int jj = 0; jj < 3; jj++) { - extend_dcoord[ii * 3 + jj] = dcoord_[ii * 3 + jj]; - if (datype_[ii] < ntypes_spin) { - extend_dcoord[(ii + nloc) * 3 + jj] = - dcoord_[ii * 3 + jj] + dspin_[ii * 3 + jj] / - spin_norm[datype_[ii]] * - virtual_len[datype_[ii]]; + for (int ff = 0; ff < nframes; ++ff) { + const size_t input_offset = static_cast(ff) * nloc * 3; + const size_t output_offset = static_cast(ff) * extend_nall * 3; + for (int jj = 0; jj < 3; jj++) { + extend_dcoord[output_offset + ii * 3 + jj] = + dcoord_[input_offset + ii * 3 + jj]; + if (datype_[ii] < ntypes_spin) { + extend_dcoord[output_offset + (ii + nloc) * 3 + jj] = + dcoord_[input_offset + ii * 3 + jj] + + dspin_[input_offset + ii * 3 + jj] / spin_norm[datype_[ii]] * + virtual_len[datype_[ii]]; + } } } } @@ -1252,11 +1304,13 @@ template void DeepSpinTF::extend_nlist( std::vector& extend_atype, const std::vector& dcoord_, const std::vector& dspin_, - const std::vector& datype_); + const std::vector& datype_, + const int nframes); template void DeepSpinTF::extend_nlist(std::vector& extend_dcoord, std::vector& extend_atype, const std::vector& dcoord_, const std::vector& dspin_, - const std::vector& datype_); + const std::vector& datype_, + const int nframes); #endif diff --git a/source/api_cc/src/common.cc b/source/api_cc/src/common.cc index dc604b32e9..181d09e376 100644 --- a/source/api_cc/src/common.cc +++ b/source/api_cc/src/common.cc @@ -184,8 +184,14 @@ void deepmd::select_real_atoms_coord(std::vector& dcoord, nloc_real = nall_real - nghost_real; dcoord.resize(static_cast(nframes) * nall_real * 3); datype.resize(nall_real); - // fwd map - select_map(dcoord, dcoord_, fwd_map, 3, nframes, nall_real, nall); + // Coordinate buffers can contain an extended atom set while aparam keeps + // the caller's original atom stride (for example DeepSpin virtual atoms). + // Infer the coordinate stride from its own frame-major buffer instead of + // reusing the aparam atom count supplied through ``nall``. + const int coord_nall = dcoord_.size() / static_cast(nframes) / 3; + assert(static_cast(nframes) * coord_nall * 3 == dcoord_.size()); + select_map(dcoord, dcoord_, fwd_map, 3, nframes, nall_real, + coord_nall); select_map(datype, datype_, fwd_map, 1); // aparam if (daparam > 0) { diff --git a/source/api_cc/tests/test_deeppot_tf_spin.cc b/source/api_cc/tests/test_deeppot_tf_spin.cc index 9c32e5c3d5..c12252a4f3 100644 --- a/source/api_cc/tests/test_deeppot_tf_spin.cc +++ b/source/api_cc/tests/test_deeppot_tf_spin.cc @@ -13,6 +13,30 @@ #include "neighbor_list.h" #include "test_utils.h" +namespace { +template +std::vector concatenate_frames(const std::vector& first, + const std::vector& second) { + std::vector result = first; + result.insert(result.end(), second.begin(), second.end()); + return result; +} + +template +void expect_two_frame_values(const std::vector& actual, + const std::vector& first, + const std::vector& second, + const double tolerance) { + ASSERT_EQ(actual.size(), first.size() + second.size()); + for (size_t ii = 0; ii < first.size(); ++ii) { + EXPECT_NEAR(actual[ii], first[ii], tolerance); + } + for (size_t ii = 0; ii < second.size(); ++ii) { + EXPECT_NEAR(actual[first.size() + ii], second[ii], tolerance); + } +} +} // namespace + template class TestInferDeepSpin : public ::testing::Test { protected: @@ -161,6 +185,64 @@ TYPED_TEST(TestInferDeepSpin, cpu_build_nlist_atomic) { } } +TYPED_TEST(TestInferDeepSpin, cpu_build_nlist_atomic_two_frames) { + using VALUETYPE = TypeParam; + auto coord_second = this->coord; + coord_second[0] += 0.07; + coord_second[3] -= 0.04; + auto spin_second = this->spin; + spin_second[2] *= -0.5; + spin_second[5] *= 0.25; + + double energy_first, energy_second; + std::vector force_first, force_mag_first, virial_first; + std::vector atom_energy_first, atom_virial_first; + std::vector force_second, force_mag_second, virial_second; + std::vector atom_energy_second, atom_virial_second; + this->dp.compute(energy_first, force_first, force_mag_first, virial_first, + atom_energy_first, atom_virial_first, this->coord, + this->spin, this->atype, this->box); + this->dp.compute(energy_second, force_second, force_mag_second, virial_second, + atom_energy_second, atom_virial_second, coord_second, + spin_second, this->atype, this->box); + EXPECT_GT(std::fabs(energy_first - energy_second), EPSILON); + + const auto coord = concatenate_frames(this->coord, coord_second); + const auto spin = concatenate_frames(this->spin, spin_second); + const auto box = concatenate_frames(this->box, this->box); + std::vector energy; + std::vector force, force_mag, virial, atom_energy, atom_virial; + this->dp.compute(energy, force, force_mag, virial, atom_energy, atom_virial, + coord, spin, this->atype, box); + + ASSERT_EQ(energy.size(), 2U); + EXPECT_NEAR(energy[0], energy_first, EPSILON); + EXPECT_NEAR(energy[1], energy_second, EPSILON); + { + SCOPED_TRACE("force"); + expect_two_frame_values(force, force_first, force_second, EPSILON); + } + { + SCOPED_TRACE("magnetic force"); + expect_two_frame_values(force_mag, force_mag_first, force_mag_second, + EPSILON); + } + { + SCOPED_TRACE("virial"); + expect_two_frame_values(virial, virial_first, virial_second, EPSILON); + } + { + SCOPED_TRACE("atomic energy"); + expect_two_frame_values(atom_energy, atom_energy_first, atom_energy_second, + EPSILON); + } + { + SCOPED_TRACE("atomic virial"); + expect_two_frame_values(atom_virial, atom_virial_first, atom_virial_second, + EPSILON); + } +} + template class TestInferDeepSpinNopbc : public ::testing::Test { protected: @@ -330,7 +412,7 @@ TYPED_TEST(TestInferDeepSpinNopbc, cpu_lmp_nlist) { double ener; std::vector force, force_mag, virial; - std::vector > nlist_data = {{1}, {0}, {3}, {2}}; + std::vector> nlist_data = {{1}, {0}, {3}, {2}}; std::vector ilist(natoms), numneigh(natoms); std::vector firstneigh(natoms); deepmd::InputNlist inlist(natoms, &ilist[0], &numneigh[0], &firstneigh[0]); @@ -369,7 +451,7 @@ TYPED_TEST(TestInferDeepSpinNopbc, cpu_lmp_nlist_atomic) { double ener; std::vector force, force_mag, virial, atom_ener, atom_vir; - std::vector > nlist_data = {{1}, {0}, {3}, {2}}; + std::vector> nlist_data = {{1}, {0}, {3}, {2}}; std::vector ilist(natoms), numneigh(natoms); std::vector firstneigh(natoms); deepmd::InputNlist inlist(natoms, &ilist[0], &numneigh[0], &firstneigh[0]); @@ -398,3 +480,69 @@ TYPED_TEST(TestInferDeepSpinNopbc, cpu_lmp_nlist_atomic) { EXPECT_LT(fabs(atom_vir[ii] - expected_v[ii]), EPSILON); } } + +TYPED_TEST(TestInferDeepSpinNopbc, cpu_lmp_nlist_atomic_two_frames) { + using VALUETYPE = TypeParam; + auto coord_second = this->coord; + coord_second[0] += 0.07; + coord_second[3] -= 0.04; + auto spin_second = this->spin; + spin_second[2] *= -0.5; + spin_second[5] *= 0.25; + + const int natoms = static_cast(this->atype.size()); + std::vector> nlist_data = {{1}, {0}, {3}, {2}}; + std::vector ilist(natoms), numneigh(natoms); + std::vector firstneigh(natoms); + deepmd::InputNlist inlist(natoms, ilist.data(), numneigh.data(), + firstneigh.data()); + convert_nlist(inlist, nlist_data); + + double energy_first, energy_second; + std::vector force_first, force_mag_first, virial_first; + std::vector atom_energy_first, atom_virial_first; + std::vector force_second, force_mag_second, virial_second; + std::vector atom_energy_second, atom_virial_second; + this->dp.compute(energy_first, force_first, force_mag_first, virial_first, + atom_energy_first, atom_virial_first, this->coord, + this->spin, this->atype, this->box, 0, inlist, 0); + this->dp.compute(energy_second, force_second, force_mag_second, virial_second, + atom_energy_second, atom_virial_second, coord_second, + spin_second, this->atype, this->box, 0, inlist, 0); + EXPECT_GT(std::fabs(energy_first - energy_second), EPSILON); + + const auto coord = concatenate_frames(this->coord, coord_second); + const auto spin = concatenate_frames(this->spin, spin_second); + const auto box = concatenate_frames(this->box, this->box); + std::vector energy; + std::vector force, force_mag, virial, atom_energy, atom_virial; + this->dp.compute(energy, force, force_mag, virial, atom_energy, atom_virial, + coord, spin, this->atype, box, 0, inlist, 0); + + ASSERT_EQ(energy.size(), 2U); + EXPECT_NEAR(energy[0], energy_first, EPSILON); + EXPECT_NEAR(energy[1], energy_second, EPSILON); + { + SCOPED_TRACE("force"); + expect_two_frame_values(force, force_first, force_second, EPSILON); + } + { + SCOPED_TRACE("magnetic force"); + expect_two_frame_values(force_mag, force_mag_first, force_mag_second, + EPSILON); + } + { + SCOPED_TRACE("virial"); + expect_two_frame_values(virial, virial_first, virial_second, EPSILON); + } + { + SCOPED_TRACE("atomic energy"); + expect_two_frame_values(atom_energy, atom_energy_first, atom_energy_second, + EPSILON); + } + { + SCOPED_TRACE("atomic virial"); + expect_two_frame_values(atom_virial, atom_virial_first, atom_virial_second, + EPSILON); + } +}