Skip to content

Commit

Permalink
calcanglesidx implementation and tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@richardjgowers
richardjgowers committed Aug 2, 2024
1 parent e725f8c commit ae6817d
Show file tree
Hide file tree
Showing 3 changed files with 138 additions and 2 deletions.
57 changes: 57 additions & 0 deletions libdistopia/src/distopia.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -788,7 +788,64 @@ namespace distopia {
template <typename T, typename B>
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<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>)];
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);
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) {
Expand Down
27 changes: 26 additions & 1 deletion libdistopia/test/test_fixtures.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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];
Expand All @@ -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];

Expand All @@ -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; i<nidx; i++) {
b_idx[i] = big_idx[i];
b_coords_contig[i*3 + 0] = coords[b_idx[i] * 3];
b_coords_contig[i*3 + 1] = coords[b_idx[i] * 3 + 1];
b_coords_contig[i*3 + 2] = coords[b_idx[i] * 3 + 2];
}
RandomInt(big_idx, nidx, 0, ncoords);
for (size_t i=0; i<nidx; i++) {
c_idx[i] = big_idx[i];
c_coords_contig[i*3 + 0] = coords[c_idx[i] * 3];
c_coords_contig[i*3 + 1] = coords[c_idx[i] * 3 + 1];
c_coords_contig[i*3 + 2] = coords[c_idx[i] * 3 + 2];
}
RandomInt(big_idx, nidx, 0, ncoords);
for (size_t i=0; i<nidx; i++) {
d_idx[i] = big_idx[i];
d_coords_contig[i*3 + 0] = coords[d_idx[i] * 3];
d_coords_contig[i*3 + 1] = coords[d_idx[i] * 3 + 1];
d_coords_contig[i*3 + 2] = coords[d_idx[i] * 3 + 2];
}

box[0] = boxsize;
box[1] = boxsize;
Expand All @@ -190,9 +211,13 @@ class CoordinatesIdx : public ::testing::Test {
delete[] coords;
delete[] a_coords_contig;
delete[] b_coords_contig;
delete[] c_coords_contig;
delete[] d_coords_contig;
delete[] big_idx;
delete[] a_idx;
delete[] b_idx;
delete[] c_idx;
delete[] d_idx;
delete[] ref_results;
delete[] results;
}
Expand Down
56 changes: 55 additions & 1 deletion libdistopia/test/test_mda_match.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -536,4 +536,58 @@ TYPED_TEST(CoordinatesIdx, CalcBondsTriclinicIdx) {
for (std::size_t i=0; i<this->nidx; i++) {
EXPECT_NEAR(this->ref_results[i], this->results[i], abs_err);
}
}
}


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; i<this->nidx; 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; i<this->nidx; 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; i<this->nidx; i++) {
EXPECT_NEAR(this->ref_results[i], this->results[i], abs_err);
}
}

0 comments on commit ae6817d

Please sign in to comment.