From 2ca3adfb96f9c714a14103076f577c2c21236ec8 Mon Sep 17 00:00:00 2001 From: Richard Gowers Date: Tue, 6 Aug 2024 11:07:44 +0100 Subject: [PATCH] adds dihedrals idx implementation and tests --- libdistopia/src/distopia.cpp | 75 +++++++++++++++++++++++++++++ libdistopia/test/test_mda_match.cpp | 54 +++++++++++++++++++++ 2 files changed, 129 insertions(+) diff --git a/libdistopia/src/distopia.cpp b/libdistopia/src/distopia.cpp index 61ee43fb..944d23b2 100644 --- a/libdistopia/src/distopia.cpp +++ b/libdistopia/src/distopia.cpp @@ -851,6 +851,81 @@ namespace distopia { template void CalcDihedralsIdx(const T *coords, const unsigned int *a_idx, const unsigned int *b_idx, const unsigned int *c_idx, const unsigned int *d_idx, int n, T *out, const B &box) { + // Logic here closely follows BondsIdx, except with an additional coordinate + const hn::ScalableTag d; + int nlanes = hn::Lanes(d); + + unsigned int a_sub[3 * HWY_MAX_LANES_D(hn::ScalableTag)]; + unsigned int b_sub[3 * HWY_MAX_LANES_D(hn::ScalableTag)]; + unsigned int c_sub[3 * HWY_MAX_LANES_D(hn::ScalableTag)]; + unsigned int d_sub[3 * HWY_MAX_LANES_D(hn::ScalableTag)]; + T out_sub[HWY_MAX_LANES_D(hn::ScalableTag)]; + T *dst; + + if (HWY_UNLIKELY(n < nlanes)) { + memset(a_sub, 0, sizeof(int) * nlanes); + memset(b_sub, 0, sizeof(int) * nlanes); + memset(c_sub, 0, sizeof(int) * nlanes); + memset(d_sub, 0, sizeof(int) * nlanes); + memcpy(a_sub, a_idx, sizeof(int) * n); + memcpy(b_sub, b_idx, sizeof(int) * n); + memcpy(c_sub, c_idx, sizeof(int) * n); + memcpy(d_sub, d_idx, sizeof(int) * n); + a_idx = a_sub; + b_idx = b_sub; + c_idx = c_sub; + d_idx = d_sub; + dst = out_sub; + } else { + dst = out; + } + + auto a_x = hn::Undefined(d); + auto a_y = hn::Undefined(d); + auto a_z = hn::Undefined(d); + auto b_x = hn::Undefined(d); + auto b_y = hn::Undefined(d); + auto b_z = hn::Undefined(d); + auto c_x = hn::Undefined(d); + auto c_y = hn::Undefined(d); + auto c_z = hn::Undefined(d); + auto d_x = hn::Undefined(d); + auto d_y = hn::Undefined(d); + auto d_z = hn::Undefined(d); + + for (size_t i=0; i <= n - nlanes; i += nlanes) { + LoadInterleavedIdx(a_idx + i, coords, a_x, a_y, a_z); + LoadInterleavedIdx(b_idx + i, coords, b_x, b_y, b_z); + LoadInterleavedIdx(c_idx + i, coords, c_x, c_y, c_z); + LoadInterleavedIdx(d_idx + i, coords, d_x, d_y, d_z); + + auto result = Dihedral(a_x, a_y, a_z, + b_x, b_y, b_z, + c_x, c_y, c_z, + d_x, d_y, d_z, + box); + + hn::StoreU(result, d, dst + i); + } + size_t rem = n % nlanes; + if (rem) { + LoadInterleavedIdx(a_idx + n - nlanes, coords, a_x, a_y, a_z); + LoadInterleavedIdx(b_idx + n - nlanes, coords, b_x, b_y, b_z); + LoadInterleavedIdx(c_idx + n - nlanes, coords, c_x, c_y, c_z); + LoadInterleavedIdx(d_idx + n - nlanes, coords, d_x, d_y, d_z); + + auto result = Dihedral(a_x, a_y, a_z, + b_x, b_y, b_z, + c_x, c_y, c_z, + d_x, d_y, d_z, + box); + + hn::StoreU(result, d, dst + n - nlanes); + } + + if (HWY_UNLIKELY(n < nlanes)) { + memcpy(out, out_sub, n * sizeof(T)); + } } void CalcBondsNoBoxDouble(const double *a, const double *b, int n, double *out) { diff --git a/libdistopia/test/test_mda_match.cpp b/libdistopia/test/test_mda_match.cpp index 4fe1bdf2..637bba0b 100644 --- a/libdistopia/test/test_mda_match.cpp +++ b/libdistopia/test/test_mda_match.cpp @@ -591,3 +591,57 @@ TYPED_TEST(CoordinatesIdx, CalcAnglesTriclinicIdx) { EXPECT_NEAR(this->ref_results[i], this->results[i], abs_err); } } + + +TYPED_TEST(CoordinatesIdx, CalcDihedralsNoBoxIdx) { + int ncoords = 250; + int nidx = 100; + + this->SetUp(ncoords, nidx, BOXSIZE, 3 * BOXSIZE); + + // reference result + distopia::CalcDihedralsNoBox(this->a_coords_contig, this->b_coords_contig, this->c_coords_contig, this->d_coords_contig, this->nidx, this->ref_results); + + // test the idx + distopia::CalcDihedralsNoBoxIdx(this->coords, this->a_idx, this->b_idx, this->c_idx, this->d_idx, this->nidx, this->results); + + for (std::size_t i=0; inidx; i++) { + EXPECT_NEAR(this->ref_results[i], this->results[i], abs_err); + } +} + + +TYPED_TEST(CoordinatesIdx, CalcDihedralsOrthoIdx) { + int ncoords = 250; + int nidx = 100; + + this->SetUp(ncoords, nidx, BOXSIZE, 3 * BOXSIZE); + + // reference result + distopia::CalcDihedralsOrtho(this->a_coords_contig, this->b_coords_contig, this->c_coords_contig, this->d_coords_contig, this->nidx, this->box, this->ref_results); + + // test the idx + distopia::CalcDihedralsOrthoIdx(this->coords, this->a_idx, this->b_idx, this->c_idx, this->d_idx, this->nidx, this->box, this->results); + + for (std::size_t i=0; inidx; i++) { + EXPECT_NEAR(this->ref_results[i], this->results[i], abs_err); + } +} + + +TYPED_TEST(CoordinatesIdx, CalcDihedralsTriclinicIdx) { + int ncoords = 250; + int nidx = 100; + + this->SetUp(ncoords, nidx, BOXSIZE, 3 * BOXSIZE); + + // reference result + distopia::CalcDihedralsTriclinic(this->a_coords_contig, this->b_coords_contig, this->c_coords_contig, this->d_coords_contig, this->nidx, this->triclinic_box, this->ref_results); + + // test the idx + distopia::CalcDihedralsTriclinicIdx(this->coords, this->a_idx, this->b_idx, this->c_idx, this->d_idx, this->nidx, this->triclinic_box, this->results); + + for (std::size_t i=0; inidx; i++) { + EXPECT_NEAR(this->ref_results[i], this->results[i], abs_err); + } +} \ No newline at end of file