diff --git a/libdistopia/src/distopia.cpp b/libdistopia/src/distopia.cpp index 180d8fc2..06198b95 100644 --- a/libdistopia/src/distopia.cpp +++ b/libdistopia/src/distopia.cpp @@ -788,7 +788,64 @@ namespace distopia { template void CalcAnglesIdx(const T *coords, const unsigned int *a_idx, const unsigned int *b_idx, const unsigned int *c_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)]; + 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); + memcpy(a_sub, a_idx, sizeof(int) * n); + memcpy(b_sub, b_idx, sizeof(int) * n); + memcpy(c_sub, c_idx, sizeof(int) * n); + a_idx = a_sub; + b_idx = b_sub; + c_idx = c_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); + + 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); + + auto result = Angle(a_x, a_y, a_z, b_x, b_y, b_z, c_x, c_y, c_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); + + auto result = Angle(a_x, a_y, a_z, b_x, b_y, b_z, c_x, c_y, c_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_fixtures.h b/libdistopia/test/test_fixtures.h index 5939c87e..bc2db44b 100644 --- a/libdistopia/test/test_fixtures.h +++ b/libdistopia/test/test_fixtures.h @@ -134,9 +134,13 @@ class CoordinatesIdx : public ::testing::Test { T *coords = nullptr; T *a_coords_contig = nullptr; T *b_coords_contig = nullptr; + T *c_coords_contig = nullptr; + T *d_coords_contig = nullptr; size_t *big_idx; unsigned int *a_idx = nullptr; unsigned int *b_idx = nullptr; + unsigned int *c_idx = nullptr; + unsigned int *d_idx = nullptr; T *ref_results = nullptr; T *results = nullptr; T box[3]; @@ -149,9 +153,13 @@ class CoordinatesIdx : public ::testing::Test { coords = new T[ncoords * 3]; a_coords_contig = new T[nidx * 3]; b_coords_contig = new T[nidx * 3]; + c_coords_contig = new T[nidx * 3]; + d_coords_contig = new T[nidx * 3]; big_idx = new size_t[nidx]; a_idx = new unsigned int[nidx]; b_idx = new unsigned int[nidx]; + c_idx = new unsigned int[nidx]; + d_idx = new unsigned int[nidx]; ref_results = new T[nidx]; results = new T[nidx]; @@ -166,13 +174,26 @@ class CoordinatesIdx : public ::testing::Test { a_coords_contig[i*3 + 2] = coords[a_idx[i] * 3 + 2]; } RandomInt(big_idx, nidx, 0, ncoords); - // copy bigidx into smaller, and also make contig coords array for (size_t i=0; inidx; i++) { EXPECT_NEAR(this->ref_results[i], this->results[i], abs_err); } -} \ No newline at end of file +} + + +TYPED_TEST(CoordinatesIdx, CalcAnglesNoBoxIdx) { + int ncoords = 250; + int nidx = 100; + + this->SetUp(ncoords, nidx, BOXSIZE, 3 * BOXSIZE); + + // reference result + distopia::CalcAnglesNoBox(this->a_coords_contig, this->b_coords_contig, this->c_coords_contig, this->nidx, this->ref_results); + + // test the idx + distopia::CalcAnglesNoBoxIdx(this->coords, this->a_idx, this->b_idx, this->c_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, CalcAnglesOrthoIdx) { + int ncoords = 250; + int nidx = 100; + + this->SetUp(ncoords, nidx, BOXSIZE, 3 * BOXSIZE); + + // reference result + distopia::CalcAnglesOrtho(this->a_coords_contig, this->b_coords_contig, this->c_coords_contig, this->nidx, this->box, this->ref_results); + + // test the idx + distopia::CalcAnglesOrthoIdx(this->coords, this->a_idx, this->b_idx, this->c_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, CalcAnglesTriclinicIdx) { + int ncoords = 250; + int nidx = 100; + + this->SetUp(ncoords, nidx, BOXSIZE, 3 * BOXSIZE); + + // reference result + distopia::CalcAnglesTriclinic(this->a_coords_contig, this->b_coords_contig, this->c_coords_contig, this->nidx, this->triclinic_box, this->ref_results); + + // test the idx + distopia::CalcAnglesTriclinicIdx(this->coords, this->a_idx, this->b_idx, this->c_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); + } +}