adds dihedrals idx implementation and tests

MDAnalysis · Aug 6, 2024 · 2ca3adf · 2ca3adf
1 parent 0cfe9e6
commit 2ca3adf
Show file tree

Hide file tree

Showing 2 changed files with 129 additions and 0 deletions.
diff --git a/libdistopia/src/distopia.cpp b/libdistopia/src/distopia.cpp
@@ -851,6 +851,81 @@ namespace distopia {
         template <typename T, typename B>
         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<T> d;
+            int nlanes = hn::Lanes(d);
+
+            unsigned int a_sub[3 * HWY_MAX_LANES_D(hn::ScalableTag<T>)];
+            unsigned int b_sub[3 * HWY_MAX_LANES_D(hn::ScalableTag<T>)];
+            unsigned int c_sub[3 * HWY_MAX_LANES_D(hn::ScalableTag<T>)];
+            unsigned int d_sub[3 * HWY_MAX_LANES_D(hn::ScalableTag<T>)];
+            T out_sub[HWY_MAX_LANES_D(hn::ScalableTag<T>)];
+            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
@@ -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; i<this->nidx; 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; i<this->nidx; 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; i<this->nidx; i++) {
+        EXPECT_NEAR(this->ref_results[i], this->results[i], abs_err);
+    }
+}