From fe0fe9bfceea9f0f7bbcd15eaed3b59f28394c1a Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Mon, 24 Jan 2022 14:10:45 -0800 Subject: [PATCH 01/51] Rename NonbondedDense -> NonbondedAllPairs --- timemachine/cpp/CMakeLists.txt | 2 +- ...ed_dense.cuh => k_nonbonded_all_pairs.cuh} | 0 timemachine/cpp/src/nonbonded.hpp | 4 +-- ...bonded_dense.cu => nonbonded_all_pairs.cu} | 30 +++++++++---------- ...nded_dense.hpp => nonbonded_all_pairs.hpp} | 6 ++-- timemachine/cpp/src/nonbonded_pairs.cu | 2 +- timemachine/cpp/src/wrap_kernels.cpp | 2 -- 7 files changed, 22 insertions(+), 24 deletions(-) rename timemachine/cpp/src/kernels/{k_nonbonded_dense.cuh => k_nonbonded_all_pairs.cuh} (100%) rename timemachine/cpp/src/{nonbonded_dense.cu => nonbonded_all_pairs.cu} (94%) rename timemachine/cpp/src/{nonbonded_dense.hpp => nonbonded_all_pairs.hpp} (95%) diff --git a/timemachine/cpp/CMakeLists.txt b/timemachine/cpp/CMakeLists.txt index 016b3301f..0bacf307c 100644 --- a/timemachine/cpp/CMakeLists.txt +++ b/timemachine/cpp/CMakeLists.txt @@ -42,7 +42,7 @@ pybind11_add_module(${LIBRARY_NAME} SHARED NO_EXTRAS src/gpu_utils.cu src/vendored/hilbert.cpp src/nonbonded.cu - src/nonbonded_dense.cu + src/nonbonded_all_pairs.cu src/nonbonded_pairs.cu src/neighborlist.cu src/harmonic_bond.cu diff --git a/timemachine/cpp/src/kernels/k_nonbonded_dense.cuh b/timemachine/cpp/src/kernels/k_nonbonded_all_pairs.cuh similarity index 100% rename from timemachine/cpp/src/kernels/k_nonbonded_dense.cuh rename to timemachine/cpp/src/kernels/k_nonbonded_all_pairs.cuh diff --git a/timemachine/cpp/src/nonbonded.hpp b/timemachine/cpp/src/nonbonded.hpp index 663c13a52..68a8ef7fc 100644 --- a/timemachine/cpp/src/nonbonded.hpp +++ b/timemachine/cpp/src/nonbonded.hpp @@ -1,6 +1,6 @@ #pragma once -#include "nonbonded_dense.hpp" +#include "nonbonded_all_pairs.hpp" #include "nonbonded_pairs.hpp" #include "potential.hpp" #include @@ -10,7 +10,7 @@ namespace timemachine { template class Nonbonded : public Potential { private: - NonbondedDense dense_; + NonbondedAllPairs dense_; static const bool Negated = true; NonbondedPairs exclusions_; // implement exclusions as negated NonbondedPairs diff --git a/timemachine/cpp/src/nonbonded_dense.cu b/timemachine/cpp/src/nonbonded_all_pairs.cu similarity index 94% rename from timemachine/cpp/src/nonbonded_dense.cu rename to timemachine/cpp/src/nonbonded_all_pairs.cu index 5a66343c5..95bd2df69 100644 --- a/timemachine/cpp/src/nonbonded_dense.cu +++ b/timemachine/cpp/src/nonbonded_all_pairs.cu @@ -10,10 +10,10 @@ #include "fixed_point.hpp" #include "gpu_utils.cuh" -#include "nonbonded_dense.hpp" +#include "nonbonded_all_pairs.hpp" #include "vendored/hilbert.h" -#include "k_nonbonded_dense.cuh" +#include "k_nonbonded_all_pairs.cuh" #include #include @@ -22,7 +22,7 @@ namespace timemachine { template -NonbondedDense::NonbondedDense( +NonbondedAllPairs::NonbondedAllPairs( const std::vector &lambda_plane_idxs, // [N] const std::vector &lambda_offset_idxs, // [N] const double beta, @@ -138,7 +138,7 @@ NonbondedDense::NonbondedDense( gpuErrchk(cudaMalloc(&d_sort_storage_, d_sort_storage_bytes_)); }; -template NonbondedDense::~NonbondedDense() { +template NonbondedAllPairs::~NonbondedAllPairs() { gpuErrchk(cudaFree(d_lambda_plane_idxs_)); gpuErrchk(cudaFree(d_lambda_offset_idxs_)); @@ -173,16 +173,16 @@ template NonbondedDense -void NonbondedDense::set_nblist_padding(double val) { +void NonbondedAllPairs::set_nblist_padding(double val) { nblist_padding_ = val; } -template void NonbondedDense::disable_hilbert_sort() { +template void NonbondedAllPairs::disable_hilbert_sort() { disable_hilbert_ = true; } template -void NonbondedDense::hilbert_sort( +void NonbondedAllPairs::hilbert_sort( const double *d_coords, const double *d_box, cudaStream_t stream) { const int tpb = 32; @@ -217,7 +217,7 @@ void __global__ k_arange(int N, unsigned int *arr) { } template -void NonbondedDense::execute_device( +void NonbondedAllPairs::execute_device( const int N, const int P, const double *d_x, @@ -248,14 +248,14 @@ void NonbondedDense::execute_device( if (N != N_) { std::cout << N << " " << N_ << std::endl; - throw std::runtime_error("NonbondedDense::execute_device() N != N_"); + throw std::runtime_error("NonbondedAllPairs::execute_device() N != N_"); } const int M = Interpolated ? 2 : 1; if (P != M * N_ * 3) { std::cout << P << " " << N_ << std::endl; - throw std::runtime_error("NonbondedDense::execute_device() P != M*N_*3"); + throw std::runtime_error("NonbondedAllPairs::execute_device() P != M*N_*3"); } // identify which tiles contain interpolated parameters @@ -410,7 +410,7 @@ void NonbondedDense::execute_device( } template -void NonbondedDense::du_dp_fixed_to_float( +void NonbondedAllPairs::du_dp_fixed_to_float( const int N, const int P, const unsigned long long *du_dp, double *du_dp_float) { // In the interpolated case we have derivatives for the initial and final parameters @@ -426,9 +426,9 @@ void NonbondedDense::du_dp_fixed_to_float( } } -template class NonbondedDense; -template class NonbondedDense; -template class NonbondedDense; -template class NonbondedDense; +template class NonbondedAllPairs; +template class NonbondedAllPairs; +template class NonbondedAllPairs; +template class NonbondedAllPairs; } // namespace timemachine diff --git a/timemachine/cpp/src/nonbonded_dense.hpp b/timemachine/cpp/src/nonbonded_all_pairs.hpp similarity index 95% rename from timemachine/cpp/src/nonbonded_dense.hpp rename to timemachine/cpp/src/nonbonded_all_pairs.hpp index b0e82b00f..af56dcb61 100644 --- a/timemachine/cpp/src/nonbonded_dense.hpp +++ b/timemachine/cpp/src/nonbonded_all_pairs.hpp @@ -25,7 +25,7 @@ typedef void (*k_nonbonded_fn)( unsigned long long *__restrict__ du_dl_buffer, unsigned long long *__restrict__ u_buffer); -template class NonbondedDense : public Potential { +template class NonbondedAllPairs : public Potential { private: std::array kernel_ptrs_; @@ -83,14 +83,14 @@ template class NonbondedDense : public Po void set_nblist_padding(double val); void disable_hilbert_sort(); - NonbondedDense( + NonbondedAllPairs( const std::vector &lambda_plane_idxs, // N const std::vector &lambda_offset_idxs, // N const double beta, const double cutoff, const std::string &kernel_src); - ~NonbondedDense(); + ~NonbondedAllPairs(); virtual void execute_device( const int N, diff --git a/timemachine/cpp/src/nonbonded_pairs.cu b/timemachine/cpp/src/nonbonded_pairs.cu index f6b7c41ef..ed43badde 100644 --- a/timemachine/cpp/src/nonbonded_pairs.cu +++ b/timemachine/cpp/src/nonbonded_pairs.cu @@ -162,7 +162,7 @@ void NonbondedPairs::execute_device( } } -// TODO: this implementation is duplicated from NonbondedDense. Worth adding NonbondedBase? +// TODO: this implementation is duplicated from NonbondedAllPairs template void NonbondedPairs::du_dp_fixed_to_float( const int N, const int P, const unsigned long long *du_dp, double *du_dp_float) { diff --git a/timemachine/cpp/src/wrap_kernels.cpp b/timemachine/cpp/src/wrap_kernels.cpp index a0a5549ef..907cf4a4e 100644 --- a/timemachine/cpp/src/wrap_kernels.cpp +++ b/timemachine/cpp/src/wrap_kernels.cpp @@ -14,8 +14,6 @@ #include "integrator.hpp" #include "neighborlist.hpp" #include "nonbonded.hpp" -#include "nonbonded_dense.hpp" -#include "nonbonded_pairs.hpp" #include "periodic_torsion.hpp" #include "potential.hpp" #include "rmsd_align.hpp" From 5848e880526b1332eb85d0b49af56649a38ec6f7 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Mon, 24 Jan 2022 14:14:53 -0800 Subject: [PATCH 02/51] Rename NonbondedPairs -> NonbondedPairList --- timemachine/cpp/CMakeLists.txt | 2 +- ...ed_pairs.cuh => k_nonbonded_pair_list.cuh} | 2 +- timemachine/cpp/src/nonbonded.hpp | 4 +-- ...bonded_pairs.cu => nonbonded_pair_list.cu} | 30 +++++++++---------- ...nded_pairs.hpp => nonbonded_pair_list.hpp} | 6 ++-- 5 files changed, 22 insertions(+), 22 deletions(-) rename timemachine/cpp/src/kernels/{k_nonbonded_pairs.cuh => k_nonbonded_pair_list.cuh} (99%) rename timemachine/cpp/src/{nonbonded_pairs.cu => nonbonded_pair_list.cu} (87%) rename timemachine/cpp/src/{nonbonded_pairs.hpp => nonbonded_pair_list.hpp} (94%) diff --git a/timemachine/cpp/CMakeLists.txt b/timemachine/cpp/CMakeLists.txt index 0bacf307c..9e2c6a76b 100644 --- a/timemachine/cpp/CMakeLists.txt +++ b/timemachine/cpp/CMakeLists.txt @@ -43,7 +43,7 @@ pybind11_add_module(${LIBRARY_NAME} SHARED NO_EXTRAS src/vendored/hilbert.cpp src/nonbonded.cu src/nonbonded_all_pairs.cu - src/nonbonded_pairs.cu + src/nonbonded_pair_list.cu src/neighborlist.cu src/harmonic_bond.cu src/harmonic_angle.cu diff --git a/timemachine/cpp/src/kernels/k_nonbonded_pairs.cuh b/timemachine/cpp/src/kernels/k_nonbonded_pair_list.cuh similarity index 99% rename from timemachine/cpp/src/kernels/k_nonbonded_pairs.cuh rename to timemachine/cpp/src/kernels/k_nonbonded_pair_list.cuh index 7d95f18a2..0c6c9290b 100644 --- a/timemachine/cpp/src/kernels/k_nonbonded_pairs.cuh +++ b/timemachine/cpp/src/kernels/k_nonbonded_pair_list.cuh @@ -14,7 +14,7 @@ void __device__ __forceinline__ accumulate(unsigned long long *__restrict acc, u } template -void __global__ k_nonbonded_pairs( +void __global__ k_nonbonded_pair_list( const int M, // number of pairs const double *__restrict__ coords, const double *__restrict__ params, diff --git a/timemachine/cpp/src/nonbonded.hpp b/timemachine/cpp/src/nonbonded.hpp index 68a8ef7fc..5c8d06e71 100644 --- a/timemachine/cpp/src/nonbonded.hpp +++ b/timemachine/cpp/src/nonbonded.hpp @@ -1,7 +1,7 @@ #pragma once #include "nonbonded_all_pairs.hpp" -#include "nonbonded_pairs.hpp" +#include "nonbonded_pair_list.hpp" #include "potential.hpp" #include @@ -13,7 +13,7 @@ template class Nonbonded : public Potenti NonbondedAllPairs dense_; static const bool Negated = true; - NonbondedPairs exclusions_; // implement exclusions as negated NonbondedPairs + NonbondedPairList exclusions_; // implement exclusions as negated NonbondedPairList public: // these are marked public but really only intended for testing. diff --git a/timemachine/cpp/src/nonbonded_pairs.cu b/timemachine/cpp/src/nonbonded_pair_list.cu similarity index 87% rename from timemachine/cpp/src/nonbonded_pairs.cu rename to timemachine/cpp/src/nonbonded_pair_list.cu index ed43badde..f08cc7eca 100644 --- a/timemachine/cpp/src/nonbonded_pairs.cu +++ b/timemachine/cpp/src/nonbonded_pair_list.cu @@ -1,14 +1,14 @@ #include "gpu_utils.cuh" -#include "k_nonbonded_pairs.cuh" +#include "k_nonbonded_pair_list.cuh" #include "math_utils.cuh" -#include "nonbonded_pairs.hpp" +#include "nonbonded_pair_list.hpp" #include #include namespace timemachine { template -NonbondedPairs::NonbondedPairs( +NonbondedPairList::NonbondedPairList( const std::vector &pair_idxs, // [M, 2] const std::vector &scales, // [M, 2] const std::vector &lambda_plane_idxs, // [N] @@ -72,7 +72,7 @@ NonbondedPairs::NonbondedPairs( }; template -NonbondedPairs::~NonbondedPairs() { +NonbondedPairList::~NonbondedPairList() { gpuErrchk(cudaFree(d_pair_idxs_)); gpuErrchk(cudaFree(d_scales_)); gpuErrchk(cudaFree(d_lambda_plane_idxs_)); @@ -89,7 +89,7 @@ NonbondedPairs::~NonbondedPairs() { }; template -void NonbondedPairs::execute_device( +void NonbondedPairList::execute_device( const int N, const int P, const double *d_x, @@ -129,7 +129,7 @@ void NonbondedPairs::execute_device( int num_blocks_pairs = ceil_divide(M_, tpb); - k_nonbonded_pairs<<>>( + k_nonbonded_pair_list<<>>( M_, d_x, Interpolated ? d_p_interp_ : d_p, @@ -164,7 +164,7 @@ void NonbondedPairs::execute_device( // TODO: this implementation is duplicated from NonbondedAllPairs template -void NonbondedPairs::du_dp_fixed_to_float( +void NonbondedPairList::du_dp_fixed_to_float( const int N, const int P, const unsigned long long *du_dp, double *du_dp_float) { // In the interpolated case we have derivatives for the initial and final parameters @@ -180,16 +180,16 @@ void NonbondedPairs::du_dp_fixed_to_float( } } -template class NonbondedPairs; -template class NonbondedPairs; +template class NonbondedPairList; +template class NonbondedPairList; -template class NonbondedPairs; -template class NonbondedPairs; +template class NonbondedPairList; +template class NonbondedPairList; -template class NonbondedPairs; -template class NonbondedPairs; +template class NonbondedPairList; +template class NonbondedPairList; -template class NonbondedPairs; -template class NonbondedPairs; +template class NonbondedPairList; +template class NonbondedPairList; } // namespace timemachine diff --git a/timemachine/cpp/src/nonbonded_pairs.hpp b/timemachine/cpp/src/nonbonded_pair_list.hpp similarity index 94% rename from timemachine/cpp/src/nonbonded_pairs.hpp rename to timemachine/cpp/src/nonbonded_pair_list.hpp index dac1c6e6d..a2770bd36 100644 --- a/timemachine/cpp/src/nonbonded_pairs.hpp +++ b/timemachine/cpp/src/nonbonded_pair_list.hpp @@ -6,7 +6,7 @@ namespace timemachine { -template class NonbondedPairs : public Potential { +template class NonbondedPairList : public Potential { private: int *d_pair_idxs_; // [M, 2] @@ -36,7 +36,7 @@ template class NonbondedPai jitify::KernelInstantiation compute_add_du_dp_interpolated_; public: - NonbondedPairs( + NonbondedPairList( const std::vector &pair_idxs, // [M, 2] const std::vector &scales, // [M, 2] const std::vector &lambda_plane_idxs, // [N] @@ -45,7 +45,7 @@ template class NonbondedPai const double cutoff, const std::string &kernel_src); - ~NonbondedPairs(); + ~NonbondedPairList(); virtual void execute_device( const int N, From 5bf391ef6690caa8bd877483d32b93be20a88695 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Mon, 24 Jan 2022 14:54:05 -0800 Subject: [PATCH 03/51] Copy NonbondedInteractionGroup from NonbondedAllPairs Rename `k_nonbonded_all_pairs.cuh` to (shared) `k_nonbonded.cuh` Move concrete implementations in k_nonbonded.cuh header to k_nonbonded.cu to prevent "multiple definitions" at linking stage --- timemachine/cpp/CMakeLists.txt | 2 + timemachine/cpp/src/kernels/k_nonbonded.cu | 48 ++ ...onbonded_all_pairs.cuh => k_nonbonded.cuh} | 35 +- timemachine/cpp/src/nonbonded_all_pairs.cu | 10 +- .../cpp/src/nonbonded_interaction_group.cu | 423 ++++++++++++++++++ .../cpp/src/nonbonded_interaction_group.hpp | 111 +++++ 6 files changed, 588 insertions(+), 41 deletions(-) create mode 100644 timemachine/cpp/src/kernels/k_nonbonded.cu rename timemachine/cpp/src/kernels/{k_nonbonded_all_pairs.cuh => k_nonbonded.cuh} (95%) create mode 100644 timemachine/cpp/src/nonbonded_interaction_group.cu create mode 100644 timemachine/cpp/src/nonbonded_interaction_group.hpp diff --git a/timemachine/cpp/CMakeLists.txt b/timemachine/cpp/CMakeLists.txt index 9e2c6a76b..9a227fd29 100644 --- a/timemachine/cpp/CMakeLists.txt +++ b/timemachine/cpp/CMakeLists.txt @@ -44,6 +44,7 @@ pybind11_add_module(${LIBRARY_NAME} SHARED NO_EXTRAS src/nonbonded.cu src/nonbonded_all_pairs.cu src/nonbonded_pair_list.cu + src/nonbonded_interaction_group.cu src/neighborlist.cu src/harmonic_bond.cu src/harmonic_angle.cu @@ -54,6 +55,7 @@ pybind11_add_module(${LIBRARY_NAME} SHARED NO_EXTRAS src/rmsd_align.cpp src/summed_potential.cu src/device_buffer.cu + src/kernels/k_nonbonded.cu src/kernels/nonbonded_common.cu ) diff --git a/timemachine/cpp/src/kernels/k_nonbonded.cu b/timemachine/cpp/src/kernels/k_nonbonded.cu new file mode 100644 index 000000000..53f66a3d4 --- /dev/null +++ b/timemachine/cpp/src/kernels/k_nonbonded.cu @@ -0,0 +1,48 @@ +#include "k_nonbonded.cuh" + +void __global__ k_coords_to_kv( + const int N, + const double *coords, + const double *box, + const unsigned int *bin_to_idx, + unsigned int *keys, + unsigned int *vals) { + + const int atom_idx = blockIdx.x * blockDim.x + threadIdx.x; + + if (atom_idx >= N) { + return; + } + + // these coords have to be centered + double bx = box[0 * 3 + 0]; + double by = box[1 * 3 + 1]; + double bz = box[2 * 3 + 2]; + + double binWidth = max(max(bx, by), bz) / 255.0; + + double x = coords[atom_idx * 3 + 0]; + double y = coords[atom_idx * 3 + 1]; + double z = coords[atom_idx * 3 + 2]; + + x -= bx * floor(x / bx); + y -= by * floor(y / by); + z -= bz * floor(z / bz); + + unsigned int bin_x = x / binWidth; + unsigned int bin_y = y / binWidth; + unsigned int bin_z = z / binWidth; + + keys[atom_idx] = bin_to_idx[bin_x * 256 * 256 + bin_y * 256 + bin_z]; + // uncomment below if you want to preserve the atom ordering + // keys[atom_idx] = atom_idx; + vals[atom_idx] = atom_idx; +} + +void __global__ k_arange(int N, unsigned int *arr) { + const int atom_idx = blockIdx.x * blockDim.x + threadIdx.x; + if (atom_idx >= N) { + return; + } + arr[atom_idx] = atom_idx; +} diff --git a/timemachine/cpp/src/kernels/k_nonbonded_all_pairs.cuh b/timemachine/cpp/src/kernels/k_nonbonded.cuh similarity index 95% rename from timemachine/cpp/src/kernels/k_nonbonded_all_pairs.cuh rename to timemachine/cpp/src/kernels/k_nonbonded.cuh index c2cab8c16..a14b5d3e3 100644 --- a/timemachine/cpp/src/kernels/k_nonbonded_all_pairs.cuh +++ b/timemachine/cpp/src/kernels/k_nonbonded.cuh @@ -5,6 +5,8 @@ #include "nonbonded_common.cuh" #include "surreal.cuh" +void __global__ k_arange(int N, unsigned int *arr); + // generate kv values from coordinates to be radix sorted void __global__ k_coords_to_kv( const int N, @@ -12,38 +14,7 @@ void __global__ k_coords_to_kv( const double *box, const unsigned int *bin_to_idx, unsigned int *keys, - unsigned int *vals) { - - const int atom_idx = blockIdx.x * blockDim.x + threadIdx.x; - - if (atom_idx >= N) { - return; - } - - // these coords have to be centered - double bx = box[0 * 3 + 0]; - double by = box[1 * 3 + 1]; - double bz = box[2 * 3 + 2]; - - double binWidth = max(max(bx, by), bz) / 255.0; - - double x = coords[atom_idx * 3 + 0]; - double y = coords[atom_idx * 3 + 1]; - double z = coords[atom_idx * 3 + 2]; - - x -= bx * floor(x / bx); - y -= by * floor(y / by); - z -= bz * floor(z / bz); - - unsigned int bin_x = x / binWidth; - unsigned int bin_y = y / binWidth; - unsigned int bin_z = z / binWidth; - - keys[atom_idx] = bin_to_idx[bin_x * 256 * 256 + bin_y * 256 + bin_z]; - // uncomment below if you want to preserve the atom ordering - // keys[atom_idx] = atom_idx; - vals[atom_idx] = atom_idx; -} + unsigned int *vals); template void __global__ k_check_rebuild_box(const int N, const double *new_box, const double *old_box, int *rebuild) { diff --git a/timemachine/cpp/src/nonbonded_all_pairs.cu b/timemachine/cpp/src/nonbonded_all_pairs.cu index 95bd2df69..f874073a9 100644 --- a/timemachine/cpp/src/nonbonded_all_pairs.cu +++ b/timemachine/cpp/src/nonbonded_all_pairs.cu @@ -13,7 +13,7 @@ #include "nonbonded_all_pairs.hpp" #include "vendored/hilbert.h" -#include "k_nonbonded_all_pairs.cuh" +#include "k_nonbonded.cuh" #include #include @@ -208,14 +208,6 @@ void NonbondedAllPairs::hilbert_sort( gpuErrchk(cudaPeekAtLastError()); } -void __global__ k_arange(int N, unsigned int *arr) { - const int atom_idx = blockIdx.x * blockDim.x + threadIdx.x; - if (atom_idx >= N) { - return; - } - arr[atom_idx] = atom_idx; -} - template void NonbondedAllPairs::execute_device( const int N, diff --git a/timemachine/cpp/src/nonbonded_interaction_group.cu b/timemachine/cpp/src/nonbonded_interaction_group.cu new file mode 100644 index 000000000..217b1a6e5 --- /dev/null +++ b/timemachine/cpp/src/nonbonded_interaction_group.cu @@ -0,0 +1,423 @@ +#include "vendored/jitify.hpp" +#include +#include +#include +#include +#include +#include +#include +#include + +#include "fixed_point.hpp" +#include "gpu_utils.cuh" +#include "nonbonded_interaction_group.hpp" +#include "vendored/hilbert.h" + +#include "k_nonbonded.cuh" + +#include +#include +#include + +namespace timemachine { + +template +NonbondedInteractionGroup::NonbondedInteractionGroup( + const std::vector &lambda_plane_idxs, // [N] + const std::vector &lambda_offset_idxs, // [N] + const double beta, + const double cutoff, + const std::string &kernel_src) + : N_(lambda_offset_idxs.size()), cutoff_(cutoff), nblist_(lambda_offset_idxs.size()), beta_(beta), + d_sort_storage_(nullptr), d_sort_storage_bytes_(0), nblist_padding_(0.1), disable_hilbert_(false), + kernel_ptrs_({// enumerate over every possible kernel combination + // U: Compute U + // X: Compute DU_DL + // L: Compute DU_DX + // P: Compute DU_DP + // U X L P + &k_nonbonded_unified, + &k_nonbonded_unified, + &k_nonbonded_unified, + &k_nonbonded_unified, + &k_nonbonded_unified, + &k_nonbonded_unified, + &k_nonbonded_unified, + &k_nonbonded_unified, + &k_nonbonded_unified, + &k_nonbonded_unified, + &k_nonbonded_unified, + &k_nonbonded_unified, + &k_nonbonded_unified, + &k_nonbonded_unified, + &k_nonbonded_unified, + &k_nonbonded_unified}), + compute_w_coords_instance_(kernel_cache_.program(kernel_src.c_str()).kernel("k_compute_w_coords").instantiate()), + compute_permute_interpolated_( + kernel_cache_.program(kernel_src.c_str()).kernel("k_permute_interpolated").instantiate()), + compute_add_du_dp_interpolated_( + kernel_cache_.program(kernel_src.c_str()).kernel("k_add_du_dp_interpolated").instantiate()) { + + if (lambda_offset_idxs.size() != N_) { + throw std::runtime_error("lambda offset idxs need to have size N"); + } + + if (lambda_offset_idxs.size() != lambda_plane_idxs.size()) { + throw std::runtime_error("lambda offset idxs and plane idxs need to be equivalent"); + } + + gpuErrchk(cudaMalloc(&d_lambda_plane_idxs_, N_ * sizeof(*d_lambda_plane_idxs_))); + gpuErrchk(cudaMemcpy( + d_lambda_plane_idxs_, &lambda_plane_idxs[0], N_ * sizeof(*d_lambda_plane_idxs_), cudaMemcpyHostToDevice)); + + gpuErrchk(cudaMalloc(&d_lambda_offset_idxs_, N_ * sizeof(*d_lambda_offset_idxs_))); + gpuErrchk(cudaMemcpy( + d_lambda_offset_idxs_, &lambda_offset_idxs[0], N_ * sizeof(*d_lambda_offset_idxs_), cudaMemcpyHostToDevice)); + + gpuErrchk(cudaMalloc(&d_perm_, N_ * sizeof(*d_perm_))); + + gpuErrchk(cudaMalloc(&d_sorted_x_, N_ * 3 * sizeof(*d_sorted_x_))); + + gpuErrchk(cudaMalloc(&d_w_, N_ * sizeof(*d_w_))); + gpuErrchk(cudaMalloc(&d_dw_dl_, N_ * sizeof(*d_dw_dl_))); + gpuErrchk(cudaMalloc(&d_sorted_w_, N_ * sizeof(*d_sorted_w_))); + gpuErrchk(cudaMalloc(&d_sorted_dw_dl_, N_ * sizeof(*d_sorted_dw_dl_))); + + gpuErrchk(cudaMalloc(&d_unsorted_p_, N_ * 3 * sizeof(*d_unsorted_p_))); // interpolated + gpuErrchk(cudaMalloc(&d_sorted_p_, N_ * 3 * sizeof(*d_sorted_p_))); // interpolated + gpuErrchk(cudaMalloc(&d_unsorted_dp_dl_, N_ * 3 * sizeof(*d_unsorted_dp_dl_))); // interpolated + gpuErrchk(cudaMalloc(&d_sorted_dp_dl_, N_ * 3 * sizeof(*d_sorted_dp_dl_))); // interpolated + gpuErrchk(cudaMalloc(&d_sorted_du_dx_, N_ * 3 * sizeof(*d_sorted_du_dx_))); + gpuErrchk(cudaMalloc(&d_sorted_du_dp_, N_ * 3 * sizeof(*d_sorted_du_dp_))); + gpuErrchk(cudaMalloc(&d_du_dp_buffer_, N_ * 3 * sizeof(*d_du_dp_buffer_))); + + gpuErrchk(cudaMallocHost(&p_ixn_count_, 1 * sizeof(*p_ixn_count_))); + + gpuErrchk(cudaMalloc(&d_nblist_x_, N_ * 3 * sizeof(*d_nblist_x_))); + gpuErrchk(cudaMemset(d_nblist_x_, 0, N_ * 3 * sizeof(*d_nblist_x_))); // set non-sensical positions + gpuErrchk(cudaMalloc(&d_nblist_box_, 3 * 3 * sizeof(*d_nblist_x_))); + gpuErrchk(cudaMemset(d_nblist_box_, 0, 3 * 3 * sizeof(*d_nblist_x_))); + gpuErrchk(cudaMalloc(&d_rebuild_nblist_, 1 * sizeof(*d_rebuild_nblist_))); + gpuErrchk(cudaMallocHost(&p_rebuild_nblist_, 1 * sizeof(*p_rebuild_nblist_))); + + gpuErrchk(cudaMalloc(&d_sort_keys_in_, N_ * sizeof(d_sort_keys_in_))); + gpuErrchk(cudaMalloc(&d_sort_keys_out_, N_ * sizeof(d_sort_keys_out_))); + gpuErrchk(cudaMalloc(&d_sort_vals_in_, N_ * sizeof(d_sort_vals_in_))); + + // initialize hilbert curve + std::vector bin_to_idx(256 * 256 * 256); + for (int i = 0; i < 256; i++) { + for (int j = 0; j < 256; j++) { + for (int k = 0; k < 256; k++) { + + bitmask_t hilbert_coords[3]; + hilbert_coords[0] = i; + hilbert_coords[1] = j; + hilbert_coords[2] = k; + + unsigned int bin = static_cast(hilbert_c2i(3, 8, hilbert_coords)); + bin_to_idx[i * 256 * 256 + j * 256 + k] = bin; + } + } + } + + gpuErrchk(cudaMalloc(&d_bin_to_idx_, 256 * 256 * 256 * sizeof(*d_bin_to_idx_))); + gpuErrchk( + cudaMemcpy(d_bin_to_idx_, &bin_to_idx[0], 256 * 256 * 256 * sizeof(*d_bin_to_idx_), cudaMemcpyHostToDevice)); + + // estimate size needed to do radix sorting, this can use uninitialized data. + cub::DeviceRadixSort::SortPairs( + d_sort_storage_, d_sort_storage_bytes_, d_sort_keys_in_, d_sort_keys_out_, d_sort_vals_in_, d_perm_, N_); + + gpuErrchk(cudaPeekAtLastError()); + gpuErrchk(cudaMalloc(&d_sort_storage_, d_sort_storage_bytes_)); +}; + +template +NonbondedInteractionGroup::~NonbondedInteractionGroup() { + + gpuErrchk(cudaFree(d_lambda_plane_idxs_)); + gpuErrchk(cudaFree(d_lambda_offset_idxs_)); + gpuErrchk(cudaFree(d_du_dp_buffer_)); + gpuErrchk(cudaFree(d_perm_)); // nullptr if we never built nblist + + gpuErrchk(cudaFree(d_bin_to_idx_)); + gpuErrchk(cudaFree(d_sorted_x_)); + + gpuErrchk(cudaFree(d_w_)); + gpuErrchk(cudaFree(d_dw_dl_)); + gpuErrchk(cudaFree(d_sorted_w_)); + gpuErrchk(cudaFree(d_sorted_dw_dl_)); + gpuErrchk(cudaFree(d_unsorted_p_)); + gpuErrchk(cudaFree(d_sorted_p_)); + gpuErrchk(cudaFree(d_unsorted_dp_dl_)); + gpuErrchk(cudaFree(d_sorted_dp_dl_)); + gpuErrchk(cudaFree(d_sorted_du_dx_)); + gpuErrchk(cudaFree(d_sorted_du_dp_)); + + gpuErrchk(cudaFree(d_sort_keys_in_)); + gpuErrchk(cudaFree(d_sort_keys_out_)); + gpuErrchk(cudaFree(d_sort_vals_in_)); + gpuErrchk(cudaFree(d_sort_storage_)); + + gpuErrchk(cudaFreeHost(p_ixn_count_)); + + gpuErrchk(cudaFree(d_nblist_x_)); + gpuErrchk(cudaFree(d_nblist_box_)); + gpuErrchk(cudaFree(d_rebuild_nblist_)); + gpuErrchk(cudaFreeHost(p_rebuild_nblist_)); +}; + +template +void NonbondedInteractionGroup::set_nblist_padding(double val) { + nblist_padding_ = val; +} + +template +void NonbondedInteractionGroup::disable_hilbert_sort() { + disable_hilbert_ = true; +} + +template +void NonbondedInteractionGroup::hilbert_sort( + const double *d_coords, const double *d_box, cudaStream_t stream) { + + const int tpb = 32; + const int B = (N_ + tpb - 1) / tpb; + + k_coords_to_kv<<>>(N_, d_coords, d_box, d_bin_to_idx_, d_sort_keys_in_, d_sort_vals_in_); + + gpuErrchk(cudaPeekAtLastError()); + + cub::DeviceRadixSort::SortPairs( + d_sort_storage_, + d_sort_storage_bytes_, + d_sort_keys_in_, + d_sort_keys_out_, + d_sort_vals_in_, + d_perm_, + N_, + 0, // begin bit + sizeof(*d_sort_keys_in_) * 8, // end bit + stream // cudaStream + ); + + gpuErrchk(cudaPeekAtLastError()); +} + +template +void NonbondedInteractionGroup::execute_device( + const int N, + const int P, + const double *d_x, + const double *d_p, // 2 * N * 3 + const double *d_box, // 3 * 3 + const double lambda, + unsigned long long *d_du_dx, + unsigned long long *d_du_dp, + unsigned long long *d_du_dl, + unsigned long long *d_u, + cudaStream_t stream) { + + // (ytz) the nonbonded algorithm proceeds as follows: + + // (done in constructor), construct a hilbert curve mapping each of the 256x256x256 cells into an index. + // a. decide if we need to rebuild the neighborlist, if so: + // - look up which cell each particle belongs to, and its linear index along the hilbert curve. + // - use radix pair sort keyed on the hilbert index with values equal to the atomic index + // - resulting sorted values is the permutation array. + // - permute lambda plane/offsets, coords + // b. else: + // - permute new coords + // c. permute parameters + // d. compute the nonbonded interactions using the neighborlist + // e. inverse permute the forces, du/dps into the original index. + // f. u and du/dl is buffered into a per-particle array, and then reduced. + // g. note that du/dl is not an exact per-particle du/dl - it is only used for reduction purposes. + + if (N != N_) { + std::cout << N << " " << N_ << std::endl; + throw std::runtime_error("NonbondedInteractionGroup::execute_device() N != N_"); + } + + const int M = Interpolated ? 2 : 1; + + if (P != M * N_ * 3) { + std::cout << P << " " << N_ << std::endl; + throw std::runtime_error("NonbondedInteractionGroup::execute_device() P != M*N_*3"); + } + + // identify which tiles contain interpolated parameters + + const int tpb = 32; + const int B = (N + tpb - 1) / tpb; + + dim3 dimGrid(B, 3, 1); + + // (ytz) see if we need to rebuild the neighborlist. + // (ytz + jfass): note that this logic needs to change if we use NPT later on since a resize in the box + // can introduce new interactions. + k_check_rebuild_coords_and_box + <<>>(N, d_x, d_nblist_x_, d_box, d_nblist_box_, nblist_padding_, d_rebuild_nblist_); + gpuErrchk(cudaPeekAtLastError()); + + // we can optimize this away by doing the check on the GPU directly. + gpuErrchk(cudaMemcpyAsync( + p_rebuild_nblist_, d_rebuild_nblist_, 1 * sizeof(*p_rebuild_nblist_), cudaMemcpyDeviceToHost, stream)); + gpuErrchk(cudaStreamSynchronize(stream)); // slow! + + if (p_rebuild_nblist_[0] > 0) { + + // (ytz): update the permutation index before building neighborlist, as the neighborlist is tied + // to a particular sort order + if (!disable_hilbert_) { + this->hilbert_sort(d_x, d_box, stream); + } else { + k_arange<<>>(N, d_perm_); + gpuErrchk(cudaPeekAtLastError()); + } + + // compute new coordinates, new lambda_idxs, new_plane_idxs + k_permute<<>>(N, d_perm_, d_x, d_sorted_x_); + gpuErrchk(cudaPeekAtLastError()); + nblist_.build_nblist_device(N, d_sorted_x_, d_box, cutoff_ + nblist_padding_, stream); + gpuErrchk(cudaMemcpyAsync( + p_ixn_count_, nblist_.get_ixn_count(), 1 * sizeof(*p_ixn_count_), cudaMemcpyDeviceToHost, stream)); + + std::vector h_box(9); + gpuErrchk(cudaMemcpyAsync(&h_box[0], d_box, 3 * 3 * sizeof(*d_box), cudaMemcpyDeviceToHost, stream)); + // Verify that the cutoff and box size are valid together. If cutoff is greater than half the box + // then a particle can interact with multiple periodic copies. + const double db_cutoff = (cutoff_ + nblist_padding_) * 2; + cudaStreamSynchronize(stream); + // Verify that box is orthogonal and the width of the box in all dimensions is greater than twice the cutoff + for (int i = 0; i < 9; i++) { + if (i == 0 || i == 4 || i == 8) { + if (h_box[i] < db_cutoff) { + throw std::runtime_error( + "Cutoff with padding is more than half of the box width, neighborlist is no longer reliable"); + } + } else if (h_box[i] != 0.0) { + throw std::runtime_error("Provided non-ortholinear box, unable to compute nonbonded energy"); + } + } + + gpuErrchk(cudaMemsetAsync(d_rebuild_nblist_, 0, sizeof(*d_rebuild_nblist_), stream)); + gpuErrchk(cudaMemcpyAsync(d_nblist_x_, d_x, N * 3 * sizeof(*d_x), cudaMemcpyDeviceToDevice, stream)); + gpuErrchk(cudaMemcpyAsync(d_nblist_box_, d_box, 3 * 3 * sizeof(*d_box), cudaMemcpyDeviceToDevice, stream)); + } else { + k_permute<<>>(N, d_perm_, d_x, d_sorted_x_); + gpuErrchk(cudaPeekAtLastError()); + } + + // do parameter interpolation here + if (Interpolated) { + CUresult result = compute_permute_interpolated_.configure(dimGrid, tpb, 0, stream) + .launch(lambda, N, d_perm_, d_p, d_sorted_p_, d_sorted_dp_dl_); + if (result != 0) { + throw std::runtime_error("Driver call to k_permute_interpolated failed"); + } + } else { + k_permute<<>>(N, d_perm_, d_p, d_sorted_p_); + gpuErrchk(cudaPeekAtLastError()); + gpuErrchk(cudaMemsetAsync(d_sorted_dp_dl_, 0, N * 3 * sizeof(*d_sorted_dp_dl_), stream)) + } + + // this stream needs to be synchronized so we can be sure that p_ixn_count_ is properly set. + // reset buffers and sorted accumulators + if (d_du_dx) { + gpuErrchk(cudaMemsetAsync(d_sorted_du_dx_, 0, N * 3 * sizeof(*d_sorted_du_dx_), stream)) + } + if (d_du_dp) { + gpuErrchk(cudaMemsetAsync(d_sorted_du_dp_, 0, N * 3 * sizeof(*d_sorted_du_dp_), stream)) + } + + // update new w coordinates + // (tbd): cache lambda value for equilibrium calculations + CUresult result = compute_w_coords_instance_.configure(B, tpb, 0, stream) + .launch(N, lambda, cutoff_, d_lambda_plane_idxs_, d_lambda_offset_idxs_, d_w_, d_dw_dl_); + if (result != 0) { + throw std::runtime_error("Driver call to k_compute_w_coords"); + } + + gpuErrchk(cudaPeekAtLastError()); + k_permute_2x<<>>(N, d_perm_, d_w_, d_dw_dl_, d_sorted_w_, d_sorted_dw_dl_); + gpuErrchk(cudaPeekAtLastError()); + + // look up which kernel we need for this computation + int kernel_idx = 0; + kernel_idx |= d_du_dp ? 1 << 0 : 0; + kernel_idx |= d_du_dl ? 1 << 1 : 0; + kernel_idx |= d_du_dx ? 1 << 2 : 0; + kernel_idx |= d_u ? 1 << 3 : 0; + + kernel_ptrs_[kernel_idx]<<>>( + N, + d_sorted_x_, + d_sorted_p_, + d_box, + d_sorted_dp_dl_, + d_sorted_w_, + d_sorted_dw_dl_, + beta_, + cutoff_, + nblist_.get_ixn_tiles(), + nblist_.get_ixn_atoms(), + d_sorted_du_dx_, + d_sorted_du_dp_, + d_du_dl, // switch to nullptr if we don't request du_dl + d_u // switch to nullptr if we don't request energies + ); + + gpuErrchk(cudaPeekAtLastError()); + + // coords are N,3 + if (d_du_dx) { + k_inv_permute_accum<<>>(N, d_perm_, d_sorted_du_dx_, d_du_dx); + gpuErrchk(cudaPeekAtLastError()); + } + + // params are N,3 + // this needs to be an accumulated permute + if (d_du_dp) { + k_inv_permute_assign<<>>(N, d_perm_, d_sorted_du_dp_, d_du_dp_buffer_); + gpuErrchk(cudaPeekAtLastError()); + } + + if (d_du_dp) { + if (Interpolated) { + CUresult result = compute_add_du_dp_interpolated_.configure(dimGrid, tpb, 0, stream) + .launch(lambda, N, d_du_dp_buffer_, d_du_dp); + if (result != 0) { + throw std::runtime_error("Driver call to k_add_du_dp_interpolated failed"); + } + } else { + k_add_ull_to_ull<<>>(N, d_du_dp_buffer_, d_du_dp); + } + gpuErrchk(cudaPeekAtLastError()); + } +} + +template +void NonbondedInteractionGroup::du_dp_fixed_to_float( + const int N, const int P, const unsigned long long *du_dp, double *du_dp_float) { + + // In the interpolated case we have derivatives for the initial and final parameters + const int num_tuples = Interpolated ? N * 2 : N; + + for (int i = 0; i < num_tuples; i++) { + const int idx_charge = i * 3 + 0; + const int idx_sig = i * 3 + 1; + const int idx_eps = i * 3 + 2; + du_dp_float[idx_charge] = FIXED_TO_FLOAT_DU_DP(du_dp[idx_charge]); + du_dp_float[idx_sig] = FIXED_TO_FLOAT_DU_DP(du_dp[idx_sig]); + du_dp_float[idx_eps] = FIXED_TO_FLOAT_DU_DP(du_dp[idx_eps]); + } +} + +template class NonbondedInteractionGroup; +template class NonbondedInteractionGroup; +template class NonbondedInteractionGroup; +template class NonbondedInteractionGroup; + +} // namespace timemachine diff --git a/timemachine/cpp/src/nonbonded_interaction_group.hpp b/timemachine/cpp/src/nonbonded_interaction_group.hpp new file mode 100644 index 000000000..c5433d73f --- /dev/null +++ b/timemachine/cpp/src/nonbonded_interaction_group.hpp @@ -0,0 +1,111 @@ +#pragma once + +#include "neighborlist.hpp" +#include "potential.hpp" +#include "vendored/jitify.hpp" +#include +#include + +namespace timemachine { + +typedef void (*k_nonbonded_fn)( + const int N, + const double *__restrict__ coords, + const double *__restrict__ params, // [N] + const double *__restrict__ box, + const double *__restrict__ dl_dp, + const double *__restrict__ coords_w, // 4D coords + const double *__restrict__ dw_dl, // 4D derivatives + const double beta, + const double cutoff, + const int *__restrict__ ixn_tiles, + const unsigned int *__restrict__ ixn_atoms, + unsigned long long *__restrict__ du_dx, + unsigned long long *__restrict__ du_dp, + unsigned long long *__restrict__ du_dl_buffer, + unsigned long long *__restrict__ u_buffer); + +template class NonbondedInteractionGroup : public Potential { + +private: + std::array kernel_ptrs_; + + int *d_lambda_plane_idxs_; + int *d_lambda_offset_idxs_; + int *p_ixn_count_; // pinned memory + + double beta_; + double cutoff_; + Neighborlist nblist_; + + const int N_; + + double nblist_padding_; + double *d_nblist_x_; // coords which were used to compute the nblist + double *d_nblist_box_; // box which was used to rebuild the nblist + int *d_rebuild_nblist_; // whether or not we have to rebuild the nblist + int *p_rebuild_nblist_; // pinned + + unsigned int *d_perm_; // hilbert curve permutation + + double *d_w_; // + double *d_dw_dl_; // + + double *d_sorted_x_; // + double *d_sorted_w_; // + double *d_sorted_dw_dl_; // + double *d_sorted_p_; // + double *d_unsorted_p_; // + double *d_sorted_dp_dl_; + double *d_unsorted_dp_dl_; + unsigned long long *d_sorted_du_dx_; // + unsigned long long *d_sorted_du_dp_; // + unsigned long long *d_du_dp_buffer_; // + + unsigned int *d_bin_to_idx_; + unsigned int *d_sort_keys_in_; + unsigned int *d_sort_keys_out_; + unsigned int *d_sort_vals_in_; + unsigned int *d_sort_storage_; + size_t d_sort_storage_bytes_; + + bool disable_hilbert_; + + void hilbert_sort(const double *d_x, const double *d_box, cudaStream_t stream); + + jitify::JitCache kernel_cache_; + jitify::KernelInstantiation compute_w_coords_instance_; + jitify::KernelInstantiation compute_permute_interpolated_; + jitify::KernelInstantiation compute_add_du_dp_interpolated_; + +public: + // these are marked public but really only intended for testing. + void set_nblist_padding(double val); + void disable_hilbert_sort(); + + NonbondedInteractionGroup( + const std::vector &lambda_plane_idxs, // N + const std::vector &lambda_offset_idxs, // N + const double beta, + const double cutoff, + const std::string &kernel_src); + + ~NonbondedInteractionGroup(); + + virtual void execute_device( + const int N, + const int P, + const double *d_x, + const double *d_p, + const double *d_box, + const double lambda, + unsigned long long *d_du_dx, + unsigned long long *d_du_dp, + unsigned long long *d_du_dl, + unsigned long long *d_u, + cudaStream_t stream) override; + + void du_dp_fixed_to_float(const int N, const int P, const unsigned long long *du_dp, double *du_dp_float) override; +}; + +} // namespace timemachine From b915a4ea7bebf6aea96137833e290ace9549fcb3 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Tue, 15 Feb 2022 12:03:00 -0800 Subject: [PATCH 04/51] Update interface --- timemachine/cpp/src/nonbonded_interaction_group.cu | 6 ++++-- timemachine/cpp/src/nonbonded_interaction_group.hpp | 8 ++++++-- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/timemachine/cpp/src/nonbonded_interaction_group.cu b/timemachine/cpp/src/nonbonded_interaction_group.cu index 217b1a6e5..e57e0c7a5 100644 --- a/timemachine/cpp/src/nonbonded_interaction_group.cu +++ b/timemachine/cpp/src/nonbonded_interaction_group.cu @@ -6,6 +6,7 @@ #include #include #include +#include #include #include "fixed_point.hpp" @@ -23,13 +24,14 @@ namespace timemachine { template NonbondedInteractionGroup::NonbondedInteractionGroup( + const std::set &row_atom_idxs, const std::vector &lambda_plane_idxs, // [N] const std::vector &lambda_offset_idxs, // [N] const double beta, const double cutoff, const std::string &kernel_src) - : N_(lambda_offset_idxs.size()), cutoff_(cutoff), nblist_(lambda_offset_idxs.size()), beta_(beta), - d_sort_storage_(nullptr), d_sort_storage_bytes_(0), nblist_padding_(0.1), disable_hilbert_(false), + : N_(lambda_offset_idxs.size()), NR_(row_atom_idxs.size()), NC_(N_ - NR_), cutoff_(cutoff), nblist_(N_), + beta_(beta), d_sort_storage_(nullptr), d_sort_storage_bytes_(0), nblist_padding_(0.1), disable_hilbert_(false), kernel_ptrs_({// enumerate over every possible kernel combination // U: Compute U // X: Compute DU_DL diff --git a/timemachine/cpp/src/nonbonded_interaction_group.hpp b/timemachine/cpp/src/nonbonded_interaction_group.hpp index c5433d73f..266fb06ae 100644 --- a/timemachine/cpp/src/nonbonded_interaction_group.hpp +++ b/timemachine/cpp/src/nonbonded_interaction_group.hpp @@ -4,6 +4,7 @@ #include "potential.hpp" #include "vendored/jitify.hpp" #include +#include #include namespace timemachine { @@ -28,6 +29,10 @@ typedef void (*k_nonbonded_fn)( template class NonbondedInteractionGroup : public Potential { private: + const int N_; // N_ = NC_ + NR_ + const int NR_; // number of row atoms + const int NC_; // number of column atoms + std::array kernel_ptrs_; int *d_lambda_plane_idxs_; @@ -38,8 +43,6 @@ template class NonbondedInteractionGroup double cutoff_; Neighborlist nblist_; - const int N_; - double nblist_padding_; double *d_nblist_x_; // coords which were used to compute the nblist double *d_nblist_box_; // box which was used to rebuild the nblist @@ -84,6 +87,7 @@ template class NonbondedInteractionGroup void disable_hilbert_sort(); NonbondedInteractionGroup( + const std::set &row_atom_idxs, const std::vector &lambda_plane_idxs, // N const std::vector &lambda_offset_idxs, // N const double beta, From 3b93f01bdcfb59dd9ef813eda96a7b4d7e5f2a1f Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Tue, 15 Feb 2022 12:17:57 -0800 Subject: [PATCH 05/51] Copy indices to device --- .../cpp/src/nonbonded_interaction_group.cu | 29 +++++++++++++++++++ .../cpp/src/nonbonded_interaction_group.hpp | 3 ++ 2 files changed, 32 insertions(+) diff --git a/timemachine/cpp/src/nonbonded_interaction_group.cu b/timemachine/cpp/src/nonbonded_interaction_group.cu index e57e0c7a5..d33103c4e 100644 --- a/timemachine/cpp/src/nonbonded_interaction_group.cu +++ b/timemachine/cpp/src/nonbonded_interaction_group.cu @@ -6,6 +6,7 @@ #include #include #include +#include #include #include @@ -22,6 +23,11 @@ namespace timemachine { +std::vector set_to_vector(const std::set &s) { + std::vector v(s.begin(), s.end()); + return v; +} + template NonbondedInteractionGroup::NonbondedInteractionGroup( const std::set &row_atom_idxs, @@ -60,6 +66,17 @@ NonbondedInteractionGroup::NonbondedInteractionGroup( compute_add_du_dp_interpolated_( kernel_cache_.program(kernel_src.c_str()).kernel("k_add_du_dp_interpolated").instantiate()) { + // compute set of column atoms as set difference + std::vector all_atom_idxs(N_); + std::iota(all_atom_idxs.begin(), all_atom_idxs.end(), 0); + std::set col_atom_idxs; + std::set_difference( + all_atom_idxs.begin(), + all_atom_idxs.end(), + row_atom_idxs.begin(), + row_atom_idxs.end(), + std::inserter(col_atom_idxs, col_atom_idxs.end())); + if (lambda_offset_idxs.size() != N_) { throw std::runtime_error("lambda offset idxs need to have size N"); } @@ -68,6 +85,16 @@ NonbondedInteractionGroup::NonbondedInteractionGroup( throw std::runtime_error("lambda offset idxs and plane idxs need to be equivalent"); } + std::vector col_atom_idxs_v(set_to_vector(col_atom_idxs)); + gpuErrchk(cudaMalloc(&d_col_atom_idxs_, NC_ * sizeof(*d_col_atom_idxs_))); + gpuErrchk( + cudaMemcpy(d_col_atom_idxs_, &col_atom_idxs_v[0], NC_ * sizeof(*d_col_atom_idxs_), cudaMemcpyHostToDevice)); + + std::vector row_atom_idxs_v(set_to_vector(row_atom_idxs)); + gpuErrchk(cudaMalloc(&d_row_atom_idxs_, NR_ * sizeof(*d_row_atom_idxs_))); + gpuErrchk( + cudaMemcpy(d_row_atom_idxs_, &row_atom_idxs_v[0], NR_ * sizeof(*d_row_atom_idxs_), cudaMemcpyHostToDevice)); + gpuErrchk(cudaMalloc(&d_lambda_plane_idxs_, N_ * sizeof(*d_lambda_plane_idxs_))); gpuErrchk(cudaMemcpy( d_lambda_plane_idxs_, &lambda_plane_idxs[0], N_ * sizeof(*d_lambda_plane_idxs_), cudaMemcpyHostToDevice)); @@ -137,6 +164,8 @@ NonbondedInteractionGroup::NonbondedInteractionGroup( template NonbondedInteractionGroup::~NonbondedInteractionGroup() { + gpuErrchk(cudaFree(d_col_atom_idxs_)); + gpuErrchk(cudaFree(d_row_atom_idxs_)); gpuErrchk(cudaFree(d_lambda_plane_idxs_)); gpuErrchk(cudaFree(d_lambda_offset_idxs_)); diff --git a/timemachine/cpp/src/nonbonded_interaction_group.hpp b/timemachine/cpp/src/nonbonded_interaction_group.hpp index 266fb06ae..be942add2 100644 --- a/timemachine/cpp/src/nonbonded_interaction_group.hpp +++ b/timemachine/cpp/src/nonbonded_interaction_group.hpp @@ -35,6 +35,9 @@ template class NonbondedInteractionGroup std::array kernel_ptrs_; + unsigned int *d_col_atom_idxs_; + unsigned int *d_row_atom_idxs_; + int *d_lambda_plane_idxs_; int *d_lambda_offset_idxs_; int *p_ixn_count_; // pinned memory From 0c06115a53b8a2b1c187c40516e557e7cdeb65e1 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Wed, 9 Feb 2022 12:49:04 -0800 Subject: [PATCH 06/51] Add wrapper for NonbondedInteractionGroup --- timemachine/cpp/src/wrap_kernels.cpp | 75 ++++++++++++++++++++++++++++ timemachine/lib/potentials.py | 4 ++ 2 files changed, 79 insertions(+) diff --git a/timemachine/cpp/src/wrap_kernels.cpp b/timemachine/cpp/src/wrap_kernels.cpp index 907cf4a4e..d695bb14c 100644 --- a/timemachine/cpp/src/wrap_kernels.cpp +++ b/timemachine/cpp/src/wrap_kernels.cpp @@ -14,6 +14,7 @@ #include "integrator.hpp" #include "neighborlist.hpp" #include "nonbonded.hpp" +#include "nonbonded_interaction_group.hpp" #include "periodic_torsion.hpp" #include "potential.hpp" #include "rmsd_align.hpp" @@ -705,6 +706,74 @@ template void declare_nonbonded(py::modul py::arg("transform_lambda_w") = "lambda"); } +std::set unique_idxs(const std::vector &idxs) { + std::set unique_idxs(idxs.begin(), idxs.end()); + if (unique_idxs.size() < idxs.size()) { + throw std::runtime_error("atom indices must be unique"); + } + return unique_idxs; +} + +template +void declare_nonbonded_interaction_group(py::module &m, const char *typestr) { + using Class = timemachine::NonbondedInteractionGroup; + std::string pyclass_name = std::string("NonbondedInteractionGroup_") + typestr; + py::class_, timemachine::Potential>( + m, pyclass_name.c_str(), py::buffer_protocol(), py::dynamic_attr()) + .def("set_nblist_padding", &timemachine::NonbondedInteractionGroup::set_nblist_padding) + .def( + "disable_hilbert_sort", + &timemachine::NonbondedInteractionGroup::disable_hilbert_sort) + .def( + py::init([](const py::array_t &row_atom_idxs_i, + const py::array_t &lambda_plane_idxs_i, + const py::array_t &lambda_offset_idxs_i, + const double beta, + const double cutoff, + const std::string &transform_lambda_charge = "lambda", + const std::string &transform_lambda_sigma = "lambda", + const std::string &transform_lambda_epsilon = "lambda", + const std::string &transform_lambda_w = "lambda") { + std::vector row_atom_idxs(row_atom_idxs_i.size()); + std::memcpy(row_atom_idxs.data(), row_atom_idxs_i.data(), row_atom_idxs_i.size() * sizeof(int)); + std::set unique_row_atom_idxs(unique_idxs(row_atom_idxs)); + + std::vector lambda_plane_idxs(lambda_plane_idxs_i.size()); + std::memcpy( + lambda_plane_idxs.data(), lambda_plane_idxs_i.data(), lambda_plane_idxs_i.size() * sizeof(int)); + + std::vector lambda_offset_idxs(lambda_offset_idxs_i.size()); + std::memcpy( + lambda_offset_idxs.data(), lambda_offset_idxs_i.data(), lambda_offset_idxs_i.size() * sizeof(int)); + + std::string dir_path = dirname(__FILE__); + std::string kernel_dir = dir_path + "/kernels"; + std::string src_path = kernel_dir + "/k_lambda_transformer_jit.cuh"; + std::ifstream t(src_path); + std::string source_str((std::istreambuf_iterator(t)), std::istreambuf_iterator()); + source_str = std::regex_replace(source_str, std::regex("KERNEL_DIR"), kernel_dir); + source_str = + std::regex_replace(source_str, std::regex("CUSTOM_EXPRESSION_CHARGE"), transform_lambda_charge); + source_str = + std::regex_replace(source_str, std::regex("CUSTOM_EXPRESSION_SIGMA"), transform_lambda_sigma); + source_str = + std::regex_replace(source_str, std::regex("CUSTOM_EXPRESSION_EPSILON"), transform_lambda_epsilon); + source_str = std::regex_replace(source_str, std::regex("CUSTOM_EXPRESSION_W"), transform_lambda_w); + + return new timemachine::NonbondedInteractionGroup( + unique_row_atom_idxs, lambda_plane_idxs, lambda_offset_idxs, beta, cutoff, source_str); + }), + py::arg("row_atom_idxs_i"), + py::arg("lambda_plane_idxs_i"), + py::arg("lambda_offset_idxs_i"), + py::arg("beta"), + py::arg("cutoff"), + py::arg("transform_lambda_charge") = "lambda", + py::arg("transform_lambda_sigma") = "lambda", + py::arg("transform_lambda_epsilon") = "lambda", + py::arg("transform_lambda_w") = "lambda"); +} + void declare_barostat(py::module &m) { using Class = timemachine::MonteCarloBarostat; @@ -810,5 +879,11 @@ PYBIND11_MODULE(custom_ops, m) { declare_nonbonded(m, "f64"); declare_nonbonded(m, "f32"); + declare_nonbonded_interaction_group(m, "f64_interpolated"); + declare_nonbonded_interaction_group(m, "f32_interpolated"); + + declare_nonbonded_interaction_group(m, "f64"); + declare_nonbonded_interaction_group(m, "f32"); + declare_context(m); } diff --git a/timemachine/lib/potentials.py b/timemachine/lib/potentials.py index 7d64b8ba1..23720f3fb 100644 --- a/timemachine/lib/potentials.py +++ b/timemachine/lib/potentials.py @@ -278,3 +278,7 @@ def unbound_impl(self, precision): custom_ctor = getattr(custom_ops, cls_name_base) return custom_ctor(*self.args) + + +class NonbondedInteractionGroup(CustomOpWrapper): + pass From 6774f285764bed2233d1473742a5fd2c31933fa2 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Wed, 9 Feb 2022 12:49:52 -0800 Subject: [PATCH 07/51] Add test for raising on invalid indices --- tests/test_nonbonded_interaction_group.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) create mode 100644 tests/test_nonbonded_interaction_group.py diff --git a/tests/test_nonbonded_interaction_group.py b/tests/test_nonbonded_interaction_group.py new file mode 100644 index 000000000..3ac12314b --- /dev/null +++ b/tests/test_nonbonded_interaction_group.py @@ -0,0 +1,21 @@ +import numpy as np +import pytest + +from timemachine.lib.potentials import NonbondedInteractionGroup + + +def test_nonbonded_interaction_group_invalid_indices(): + def make_potential(row_atom_idxs, num_atoms): + lambda_plane_idxs = [0] * num_atoms + lambda_offset_idxs = [0] * num_atoms + return NonbondedInteractionGroup(row_atom_idxs, lambda_plane_idxs, lambda_offset_idxs, 1.0, 1.0).unbound_impl( + np.float64 + ) + + with pytest.raises(RuntimeError) as e: + make_potential([], 1) + assert "row_atom_idxs must be nonempty" in str(e) + + with pytest.raises(RuntimeError) as e: + make_potential([1, 1], 3) + assert "atom indices must be unique" in str(e) From 71c96e3dfd7278ff6dd5d77748fe64cd3f262795 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Tue, 15 Feb 2022 12:18:05 -0800 Subject: [PATCH 08/51] Check for empty index set --- timemachine/cpp/src/nonbonded_interaction_group.cu | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/timemachine/cpp/src/nonbonded_interaction_group.cu b/timemachine/cpp/src/nonbonded_interaction_group.cu index d33103c4e..b89249602 100644 --- a/timemachine/cpp/src/nonbonded_interaction_group.cu +++ b/timemachine/cpp/src/nonbonded_interaction_group.cu @@ -66,6 +66,10 @@ NonbondedInteractionGroup::NonbondedInteractionGroup( compute_add_du_dp_interpolated_( kernel_cache_.program(kernel_src.c_str()).kernel("k_add_du_dp_interpolated").instantiate()) { + if (NR_ == 0) { + throw std::runtime_error("row_atom_idxs must be nonempty"); + } + // compute set of column atoms as set difference std::vector all_atom_idxs(N_); std::iota(all_atom_idxs.begin(), all_atom_idxs.end(), 0); From 980185e9be1a3cc6e1533bc9d926ab3b9cae24db Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Wed, 9 Feb 2022 15:06:38 -0800 Subject: [PATCH 09/51] Add basic correctness test --- tests/test_nonbonded_interaction_group.py | 108 ++++++++++++++++++++++ 1 file changed, 108 insertions(+) diff --git a/tests/test_nonbonded_interaction_group.py b/tests/test_nonbonded_interaction_group.py index 3ac12314b..19692a63a 100644 --- a/tests/test_nonbonded_interaction_group.py +++ b/tests/test_nonbonded_interaction_group.py @@ -1,9 +1,117 @@ +import jax + +jax.config.update("jax_enable_x64", True) + import numpy as np import pytest +from common import GradientTest, prepare_reference_nonbonded +from simtk.openmm import app +from timemachine.fe.utils import to_md_units +from timemachine.ff.handlers import openmm_deserializer +from timemachine.lib import potentials from timemachine.lib.potentials import NonbondedInteractionGroup +@pytest.fixture(autouse=True) +def set_random_seed(): + np.random.seed(2021) + yield + + +@pytest.fixture +def test_system(): + pdb_path = "tests/data/5dfr_solv_equil.pdb" + host_pdb = app.PDBFile(pdb_path) + ff = app.ForceField("amber99sbildn.xml", "tip3p.xml") + return ( + ff.createSystem(host_pdb.topology, nonbondedMethod=app.NoCutoff, constraints=None, rigidWater=False), + host_pdb.positions, + host_pdb.topology.getPeriodicBoxVectors(), + ) + + +@pytest.fixture +def ref_nonbonded_potential(test_system): + host_system, _, _ = test_system + host_fns, _ = openmm_deserializer.deserialize_system(host_system, cutoff=1.0) + + nonbonded_fn = None + for f in host_fns: + if isinstance(f, potentials.Nonbonded): + nonbonded_fn = f + + return nonbonded_fn + + +@pytest.fixture +def test_conf(test_system): + _, host_coords, _ = test_system + return np.array([[to_md_units(x), to_md_units(y), to_md_units(z)] for x, y, z in host_coords]) + + +@pytest.fixture +def test_box(test_system): + _, _, box = test_system + return np.asarray(box / box.unit) + + +@pytest.mark.parametrize("beta", [2.0]) +@pytest.mark.parametrize("cutoff", [1.1]) +@pytest.mark.parametrize("precision,rtol,atol", [(np.float64, 1e-8, 1e-8), (np.float32, 1e-4, 5e-4)]) +@pytest.mark.parametrize("num_row_atoms", [1, 15]) +@pytest.mark.parametrize("num_atoms", [33, 231, 1050]) +def test_nonbonded_interaction_group_correctness( + num_atoms, num_row_atoms, precision, rtol, atol, cutoff, beta, ref_nonbonded_potential, test_conf, test_box +): + + test_conf = test_conf[:num_atoms] + test_params = ref_nonbonded_potential.params[:num_atoms, :] + test_lambda_offset_idxs = np.zeros(num_atoms, dtype=np.int32) + + def make_reference_nonbonded(lambda_plane_idxs): + return prepare_reference_nonbonded( + params=test_params, + exclusion_idxs=np.array([], dtype=np.int32), + scales=np.zeros((0, 2), dtype=np.float64), + lambda_plane_idxs=lambda_plane_idxs, + lambda_offset_idxs=test_lambda_offset_idxs, + beta=beta, + cutoff=cutoff, + ) + + ref_allpairs = make_reference_nonbonded(np.zeros(num_atoms, dtype=np.int32)) + + num_col_atoms = num_atoms - num_row_atoms + + ref_allpairs_minus_ixngroups = make_reference_nonbonded( + np.concatenate((np.zeros(num_row_atoms, dtype=np.int32), np.ones(num_col_atoms, dtype=np.int32))) + ) + + def ref_ixngroups(*args): + return ref_allpairs(*args) - ref_allpairs_minus_ixngroups(*args) + + test_ixngroups = NonbondedInteractionGroup( + np.arange(0, num_row_atoms, dtype=np.int32), + np.zeros(num_atoms, dtype=np.int32), # lambda plane indices + test_lambda_offset_idxs, + beta, + cutoff, + ) + + GradientTest().compare_forces( + test_conf, + test_params, + test_box, + lamb=0.1, + ref_potential=ref_ixngroups, + test_potential=test_ixngroups, + rtol=rtol, + atol=atol, + precision=precision, + ) + + def test_nonbonded_interaction_group_invalid_indices(): def make_potential(row_atom_idxs, num_atoms): lambda_plane_idxs = [0] * num_atoms From 945e62a015ce6f5e17b1a600ed53d066d4f91bce Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Tue, 15 Feb 2022 12:36:22 -0800 Subject: [PATCH 10/51] Implement potential General strategy: modify permutation to 1. sort each group of atoms into a contiguous block 2. hilbert sort each group separately --- timemachine/cpp/src/kernels/k_nonbonded.cu | 43 +++++++++++++++ timemachine/cpp/src/kernels/k_nonbonded.cuh | 10 ++++ .../cpp/src/nonbonded_interaction_group.cu | 52 +++++++++++++------ .../cpp/src/nonbonded_interaction_group.hpp | 8 ++- 4 files changed, 96 insertions(+), 17 deletions(-) diff --git a/timemachine/cpp/src/kernels/k_nonbonded.cu b/timemachine/cpp/src/kernels/k_nonbonded.cu index 53f66a3d4..719466555 100644 --- a/timemachine/cpp/src/kernels/k_nonbonded.cu +++ b/timemachine/cpp/src/kernels/k_nonbonded.cu @@ -39,6 +39,49 @@ void __global__ k_coords_to_kv( vals[atom_idx] = atom_idx; } +// TODO: DRY with k_coords_to_kv +void __global__ k_coords_to_kv_gather( + const int N, + const unsigned int *atom_idxs, + const double *coords, + const double *box, + const unsigned int *bin_to_idx, + unsigned int *keys, + unsigned int *vals) { + + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + + if (idx >= N) { + return; + } + + const int atom_idx = atom_idxs[idx]; + + // these coords have to be centered + double bx = box[0 * 3 + 0]; + double by = box[1 * 3 + 1]; + double bz = box[2 * 3 + 2]; + + double binWidth = max(max(bx, by), bz) / 255.0; + + double x = coords[atom_idx * 3 + 0]; + double y = coords[atom_idx * 3 + 1]; + double z = coords[atom_idx * 3 + 2]; + + x -= bx * floor(x / bx); + y -= by * floor(y / by); + z -= bz * floor(z / bz); + + unsigned int bin_x = x / binWidth; + unsigned int bin_y = y / binWidth; + unsigned int bin_z = z / binWidth; + + keys[idx] = bin_to_idx[bin_x * 256 * 256 + bin_y * 256 + bin_z]; + // uncomment below if you want to preserve the atom ordering + // keys[idx] = atom_idx; + vals[idx] = atom_idx; +} + void __global__ k_arange(int N, unsigned int *arr) { const int atom_idx = blockIdx.x * blockDim.x + threadIdx.x; if (atom_idx >= N) { diff --git a/timemachine/cpp/src/kernels/k_nonbonded.cuh b/timemachine/cpp/src/kernels/k_nonbonded.cuh index a14b5d3e3..eefe68a9b 100644 --- a/timemachine/cpp/src/kernels/k_nonbonded.cuh +++ b/timemachine/cpp/src/kernels/k_nonbonded.cuh @@ -16,6 +16,16 @@ void __global__ k_coords_to_kv( unsigned int *keys, unsigned int *vals); +// variant of k_coords_to_kv allowing the selection of a subset of coordinates +void __global__ k_coords_to_kv_gather( + const int N, // number of atoms in selection + const unsigned int *atom_idxs, // [N] indices of atoms to select + const double *coords, + const double *box, + const unsigned int *bin_to_idx, + unsigned int *keys, + unsigned int *vals); + template void __global__ k_check_rebuild_box(const int N, const double *new_box, const double *old_box, int *rebuild) { diff --git a/timemachine/cpp/src/nonbonded_interaction_group.cu b/timemachine/cpp/src/nonbonded_interaction_group.cu index b89249602..d730e7842 100644 --- a/timemachine/cpp/src/nonbonded_interaction_group.cu +++ b/timemachine/cpp/src/nonbonded_interaction_group.cu @@ -36,7 +36,7 @@ NonbondedInteractionGroup::NonbondedInteractionGroup( const double beta, const double cutoff, const std::string &kernel_src) - : N_(lambda_offset_idxs.size()), NR_(row_atom_idxs.size()), NC_(N_ - NR_), cutoff_(cutoff), nblist_(N_), + : N_(lambda_offset_idxs.size()), NR_(row_atom_idxs.size()), NC_(N_ - NR_), cutoff_(cutoff), nblist_(NC_, NR_), beta_(beta), d_sort_storage_(nullptr), d_sort_storage_bytes_(0), nblist_padding_(0.1), disable_hilbert_(false), kernel_ptrs_({// enumerate over every possible kernel combination // U: Compute U @@ -160,7 +160,13 @@ NonbondedInteractionGroup::NonbondedInteractionGroup( // estimate size needed to do radix sorting, this can use uninitialized data. cub::DeviceRadixSort::SortPairs( - d_sort_storage_, d_sort_storage_bytes_, d_sort_keys_in_, d_sort_keys_out_, d_sort_vals_in_, d_perm_, N_); + d_sort_storage_, + d_sort_storage_bytes_, + d_sort_keys_in_, + d_sort_keys_out_, + d_sort_vals_in_, + d_perm_, + std::max(NC_, NR_)); gpuErrchk(cudaPeekAtLastError()); gpuErrchk(cudaMalloc(&d_sort_storage_, d_sort_storage_bytes_)); @@ -215,12 +221,18 @@ void NonbondedInteractionGroup::disable_hilbert_sort() { template void NonbondedInteractionGroup::hilbert_sort( - const double *d_coords, const double *d_box, cudaStream_t stream) { + const int N, + const unsigned int *d_atom_idxs, + const double *d_coords, + const double *d_box, + unsigned int *d_perm, + cudaStream_t stream) { const int tpb = 32; - const int B = (N_ + tpb - 1) / tpb; + const int B = ceil_divide(N, tpb); - k_coords_to_kv<<>>(N_, d_coords, d_box, d_bin_to_idx_, d_sort_keys_in_, d_sort_vals_in_); + k_coords_to_kv_gather<<>>( + N, d_atom_idxs, d_coords, d_box, d_bin_to_idx_, d_sort_keys_in_, d_sort_vals_in_); gpuErrchk(cudaPeekAtLastError()); @@ -230,8 +242,8 @@ void NonbondedInteractionGroup::hilbert_sort( d_sort_keys_in_, d_sort_keys_out_, d_sort_vals_in_, - d_perm_, - N_, + d_perm, + N, 0, // begin bit sizeof(*d_sort_keys_in_) * 8, // end bit stream // cudaStream @@ -285,15 +297,14 @@ void NonbondedInteractionGroup::execute_device( // identify which tiles contain interpolated parameters const int tpb = 32; - const int B = (N + tpb - 1) / tpb; - - dim3 dimGrid(B, 3, 1); + const int B = ceil_divide(N_, tpb); // (ytz) see if we need to rebuild the neighborlist. // (ytz + jfass): note that this logic needs to change if we use NPT later on since a resize in the box // can introduce new interactions. k_check_rebuild_coords_and_box - <<>>(N, d_x, d_nblist_x_, d_box, d_nblist_box_, nblist_padding_, d_rebuild_nblist_); + <<>>(N_, d_x, d_nblist_x_, d_box, d_nblist_box_, nblist_padding_, d_rebuild_nblist_); + gpuErrchk(cudaPeekAtLastError()); // we can optimize this away by doing the check on the GPU directly. @@ -301,21 +312,28 @@ void NonbondedInteractionGroup::execute_device( p_rebuild_nblist_, d_rebuild_nblist_, 1 * sizeof(*p_rebuild_nblist_), cudaMemcpyDeviceToHost, stream)); gpuErrchk(cudaStreamSynchronize(stream)); // slow! + dim3 dimGrid(B, 3, 1); + if (p_rebuild_nblist_[0] > 0) { // (ytz): update the permutation index before building neighborlist, as the neighborlist is tied // to a particular sort order if (!disable_hilbert_) { - this->hilbert_sort(d_x, d_box, stream); + this->hilbert_sort(NC_, d_col_atom_idxs_, d_x, d_box, d_perm_, stream); + this->hilbert_sort(NR_, d_row_atom_idxs_, d_x, d_box, d_perm_ + NC_, stream); } else { - k_arange<<>>(N, d_perm_); - gpuErrchk(cudaPeekAtLastError()); + gpuErrchk(cudaMemcpyAsync( + d_perm_, d_col_atom_idxs_, NC_ * sizeof(*d_col_atom_idxs_), cudaMemcpyDeviceToDevice, stream)); + gpuErrchk(cudaMemcpyAsync( + d_perm_ + NC_, d_row_atom_idxs_, NR_ * sizeof(*d_row_atom_idxs_), cudaMemcpyDeviceToDevice, stream)); } // compute new coordinates, new lambda_idxs, new_plane_idxs - k_permute<<>>(N, d_perm_, d_x, d_sorted_x_); + k_permute<<>>(N_, d_perm_, d_x, d_sorted_x_); gpuErrchk(cudaPeekAtLastError()); - nblist_.build_nblist_device(N, d_sorted_x_, d_box, cutoff_ + nblist_padding_, stream); + + nblist_.build_nblist_device( + NC_, NR_, d_sorted_x_, d_sorted_x_ + 3 * NC_, d_box, cutoff_ + nblist_padding_, stream); gpuErrchk(cudaMemcpyAsync( p_ixn_count_, nblist_.get_ixn_count(), 1 * sizeof(*p_ixn_count_), cudaMemcpyDeviceToHost, stream)); @@ -325,6 +343,7 @@ void NonbondedInteractionGroup::execute_device( // then a particle can interact with multiple periodic copies. const double db_cutoff = (cutoff_ + nblist_padding_) * 2; cudaStreamSynchronize(stream); + // Verify that box is orthogonal and the width of the box in all dimensions is greater than twice the cutoff for (int i = 0; i < 9; i++) { if (i == 0 || i == 4 || i == 8) { @@ -369,6 +388,7 @@ void NonbondedInteractionGroup::execute_device( // update new w coordinates // (tbd): cache lambda value for equilibrium calculations + // TODO: skip computing w coords for non-interacting atoms? CUresult result = compute_w_coords_instance_.configure(B, tpb, 0, stream) .launch(N, lambda, cutoff_, d_lambda_plane_idxs_, d_lambda_offset_idxs_, d_w_, d_dw_dl_); if (result != 0) { diff --git a/timemachine/cpp/src/nonbonded_interaction_group.hpp b/timemachine/cpp/src/nonbonded_interaction_group.hpp index be942add2..0d3bd739d 100644 --- a/timemachine/cpp/src/nonbonded_interaction_group.hpp +++ b/timemachine/cpp/src/nonbonded_interaction_group.hpp @@ -77,7 +77,13 @@ template class NonbondedInteractionGroup bool disable_hilbert_; - void hilbert_sort(const double *d_x, const double *d_box, cudaStream_t stream); + void hilbert_sort( + const int N, + const unsigned int *d_atom_idxs, + const double *d_x, + const double *d_box, + unsigned int *d_perm, + cudaStream_t stream); jitify::JitCache kernel_cache_; jitify::KernelInstantiation compute_w_coords_instance_; From b532e89675875d6f9cae3d82edcc3a5cf0367b2b Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Sun, 13 Feb 2022 08:53:31 -0800 Subject: [PATCH 11/51] Skip kernel invocation if there are no interactions --- .../cpp/src/nonbonded_interaction_group.cu | 39 ++++++++++--------- 1 file changed, 21 insertions(+), 18 deletions(-) diff --git a/timemachine/cpp/src/nonbonded_interaction_group.cu b/timemachine/cpp/src/nonbonded_interaction_group.cu index d730e7842..254a6a3a7 100644 --- a/timemachine/cpp/src/nonbonded_interaction_group.cu +++ b/timemachine/cpp/src/nonbonded_interaction_group.cu @@ -406,25 +406,28 @@ void NonbondedInteractionGroup::execute_device( kernel_idx |= d_du_dx ? 1 << 2 : 0; kernel_idx |= d_u ? 1 << 3 : 0; - kernel_ptrs_[kernel_idx]<<>>( - N, - d_sorted_x_, - d_sorted_p_, - d_box, - d_sorted_dp_dl_, - d_sorted_w_, - d_sorted_dw_dl_, - beta_, - cutoff_, - nblist_.get_ixn_tiles(), - nblist_.get_ixn_atoms(), - d_sorted_du_dx_, - d_sorted_du_dp_, - d_du_dl, // switch to nullptr if we don't request du_dl - d_u // switch to nullptr if we don't request energies - ); + if (p_ixn_count_[0] > 0) { + + kernel_ptrs_[kernel_idx]<<>>( + N, + d_sorted_x_, + d_sorted_p_, + d_box, + d_sorted_dp_dl_, + d_sorted_w_, + d_sorted_dw_dl_, + beta_, + cutoff_, + nblist_.get_ixn_tiles(), + nblist_.get_ixn_atoms(), + d_sorted_du_dx_, + d_sorted_du_dp_, + d_du_dl, // switch to nullptr if we don't request du_dl + d_u // switch to nullptr if we don't request energies + ); - gpuErrchk(cudaPeekAtLastError()); + gpuErrchk(cudaPeekAtLastError()); + } // coords are N,3 if (d_du_dx) { From 0f247852ddfd62b0e353709d717a53f865fc06e8 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Sun, 13 Feb 2022 23:07:41 -0800 Subject: [PATCH 12/51] Add offset to row atom indices in case NR != 0 In the interacting groups case (NR != 0), the row atom indices are offset in the sorted arrays (e.g. d_sorted_x) by the number of column atoms --- timemachine/cpp/src/kernels/k_nonbonded.cuh | 49 +++++++++++++------ timemachine/cpp/src/nonbonded_all_pairs.cu | 1 + timemachine/cpp/src/nonbonded_all_pairs.hpp | 3 +- .../cpp/src/nonbonded_interaction_group.cu | 3 +- .../cpp/src/nonbonded_interaction_group.hpp | 3 +- 5 files changed, 40 insertions(+), 19 deletions(-) diff --git a/timemachine/cpp/src/kernels/k_nonbonded.cuh b/timemachine/cpp/src/kernels/k_nonbonded.cuh index eefe68a9b..a542795e6 100644 --- a/timemachine/cpp/src/kernels/k_nonbonded.cuh +++ b/timemachine/cpp/src/kernels/k_nonbonded.cuh @@ -263,7 +263,8 @@ template < bool COMPUTE_DU_DP> // void __device__ __forceinline__ v_nonbonded_unified( void __device__ v_nonbonded_unified( - const int N, + const int NC, + const int NR, const double *__restrict__ coords, const double *__restrict__ params, // [N] const double *__restrict__ box, @@ -297,6 +298,12 @@ void __device__ v_nonbonded_unified( // int lambda_offset_i = atom_i_idx < N ? lambda_offset_idxs[atom_i_idx] : 0; // int lambda_plane_i = atom_i_idx < N ? lambda_plane_idxs[atom_i_idx] : 0; + int N = NC + NR; + + if (NR != 0) { + atom_i_idx += NC; + } + RealType ci_x = atom_i_idx < N ? coords[atom_i_idx * 3 + 0] : 0; RealType ci_y = atom_i_idx < N ? coords[atom_i_idx * 3 + 1] : 0; RealType ci_z = atom_i_idx < N ? coords[atom_i_idx * 3 + 2] : 0; @@ -329,15 +336,15 @@ void __device__ v_nonbonded_unified( // int lambda_offset_j = atom_j_idx < N ? lambda_offset_idxs[atom_j_idx] : 0; // int lambda_plane_j = atom_j_idx < N ? lambda_plane_idxs[atom_j_idx] : 0; - RealType cj_x = atom_j_idx < N ? coords[atom_j_idx * 3 + 0] : 0; - RealType cj_y = atom_j_idx < N ? coords[atom_j_idx * 3 + 1] : 0; - RealType cj_z = atom_j_idx < N ? coords[atom_j_idx * 3 + 2] : 0; - RealType cj_w = atom_j_idx < N ? coords_w[atom_j_idx] : 0; + RealType cj_x = atom_j_idx < NC ? coords[atom_j_idx * 3 + 0] : 0; + RealType cj_y = atom_j_idx < NC ? coords[atom_j_idx * 3 + 1] : 0; + RealType cj_z = atom_j_idx < NC ? coords[atom_j_idx * 3 + 2] : 0; + RealType cj_w = atom_j_idx < NC ? coords_w[atom_j_idx] : 0; - RealType dq_dl_j = atom_j_idx < N ? dp_dl[atom_j_idx * 3 + 0] : 0; - RealType dsig_dl_j = atom_j_idx < N ? dp_dl[atom_j_idx * 3 + 1] : 0; - RealType deps_dl_j = atom_j_idx < N ? dp_dl[atom_j_idx * 3 + 2] : 0; - RealType dw_dl_j = atom_j_idx < N ? dw_dl[atom_j_idx] : 0; + RealType dq_dl_j = atom_j_idx < NC ? dp_dl[atom_j_idx * 3 + 0] : 0; + RealType dsig_dl_j = atom_j_idx < NC ? dp_dl[atom_j_idx * 3 + 1] : 0; + RealType deps_dl_j = atom_j_idx < NC ? dp_dl[atom_j_idx * 3 + 2] : 0; + RealType dw_dl_j = atom_j_idx < NC ? dw_dl[atom_j_idx] : 0; unsigned long long gj_x = 0; unsigned long long gj_y = 0; @@ -347,9 +354,9 @@ void __device__ v_nonbonded_unified( int lj_param_idx_sig_j = atom_j_idx * 3 + 1; int lj_param_idx_eps_j = atom_j_idx * 3 + 2; - RealType qj = atom_j_idx < N ? params[charge_param_idx_j] : 0; - RealType sig_j = atom_j_idx < N ? params[lj_param_idx_sig_j] : 0; - RealType eps_j = atom_j_idx < N ? params[lj_param_idx_eps_j] : 0; + RealType qj = atom_j_idx < NC ? params[charge_param_idx_j] : 0; + RealType sig_j = atom_j_idx < NC ? params[lj_param_idx_sig_j] : 0; + RealType eps_j = atom_j_idx < NC ? params[lj_param_idx_eps_j] : 0; unsigned long long g_qj = 0; unsigned long long g_sigj = 0; @@ -387,7 +394,8 @@ void __device__ v_nonbonded_unified( // (ytz): note that d2ij must be *strictly* less than cutoff_squared. This is because we set the // non-interacting atoms to exactly real_cutoff*real_cutoff. This ensures that atoms who's 4th dimension // is set to cutoff are non-interacting. - if (d2ij < cutoff_squared && atom_j_idx > atom_i_idx && atom_j_idx < N && atom_i_idx < N) { + if (d2ij < cutoff_squared && atom_i_idx < N && + (NR == 0 && atom_j_idx > atom_i_idx && atom_j_idx < N || NR != 0 && atom_j_idx < NC)) { // electrostatics RealType u; @@ -528,7 +536,8 @@ void __device__ v_nonbonded_unified( template void __global__ k_nonbonded_unified( - const int N, + const int NC, + const int NR, const double *__restrict__ coords, const double *__restrict__ params, // [N] const double *__restrict__ box, @@ -548,6 +557,12 @@ void __global__ k_nonbonded_unified( int row_block_idx = ixn_tiles[tile_idx]; int atom_i_idx = row_block_idx * 32 + threadIdx.x; + int N = NC + NR; + + if (NR != 0) { + atom_i_idx += NC; + } + RealType dq_dl_i = atom_i_idx < N ? dp_dl[atom_i_idx * 3 + 0] : 0; RealType dsig_dl_i = atom_i_idx < N ? dp_dl[atom_i_idx * 3 + 1] : 0; RealType deps_dl_i = atom_i_idx < N ? dp_dl[atom_i_idx * 3 + 2] : 0; @@ -568,7 +583,8 @@ void __global__ k_nonbonded_unified( if (tile_is_vanilla) { v_nonbonded_unified( - N, + NC, + NR, coords, params, box, @@ -585,7 +601,8 @@ void __global__ k_nonbonded_unified( u_buffer); } else { v_nonbonded_unified( - N, + NC, + NR, coords, params, box, diff --git a/timemachine/cpp/src/nonbonded_all_pairs.cu b/timemachine/cpp/src/nonbonded_all_pairs.cu index f874073a9..6ac40b3c1 100644 --- a/timemachine/cpp/src/nonbonded_all_pairs.cu +++ b/timemachine/cpp/src/nonbonded_all_pairs.cu @@ -356,6 +356,7 @@ void NonbondedAllPairs::execute_device( kernel_ptrs_[kernel_idx]<<>>( N, + 0, d_sorted_x_, d_sorted_p_, d_box, diff --git a/timemachine/cpp/src/nonbonded_all_pairs.hpp b/timemachine/cpp/src/nonbonded_all_pairs.hpp index af56dcb61..072eafb38 100644 --- a/timemachine/cpp/src/nonbonded_all_pairs.hpp +++ b/timemachine/cpp/src/nonbonded_all_pairs.hpp @@ -9,7 +9,8 @@ namespace timemachine { typedef void (*k_nonbonded_fn)( - const int N, + const int NC, + const int NR, const double *__restrict__ coords, const double *__restrict__ params, // [N] const double *__restrict__ box, diff --git a/timemachine/cpp/src/nonbonded_interaction_group.cu b/timemachine/cpp/src/nonbonded_interaction_group.cu index 254a6a3a7..9b55d9ba3 100644 --- a/timemachine/cpp/src/nonbonded_interaction_group.cu +++ b/timemachine/cpp/src/nonbonded_interaction_group.cu @@ -409,7 +409,8 @@ void NonbondedInteractionGroup::execute_device( if (p_ixn_count_[0] > 0) { kernel_ptrs_[kernel_idx]<<>>( - N, + NC_, + NR_, d_sorted_x_, d_sorted_p_, d_box, diff --git a/timemachine/cpp/src/nonbonded_interaction_group.hpp b/timemachine/cpp/src/nonbonded_interaction_group.hpp index 0d3bd739d..394bc1bb2 100644 --- a/timemachine/cpp/src/nonbonded_interaction_group.hpp +++ b/timemachine/cpp/src/nonbonded_interaction_group.hpp @@ -10,7 +10,8 @@ namespace timemachine { typedef void (*k_nonbonded_fn)( - const int N, + const int NC, + const int NR, const double *__restrict__ coords, const double *__restrict__ params, // [N] const double *__restrict__ box, From ecf76d848fd140cca020c05ec340243f1901c623 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Tue, 15 Feb 2022 17:40:57 -0800 Subject: [PATCH 13/51] Avoid "test_" prefix in fixtures --- tests/test_nonbonded_interaction_group.py | 43 ++++++++++++++--------- 1 file changed, 26 insertions(+), 17 deletions(-) diff --git a/tests/test_nonbonded_interaction_group.py b/tests/test_nonbonded_interaction_group.py index 19692a63a..ad4b39c8e 100644 --- a/tests/test_nonbonded_interaction_group.py +++ b/tests/test_nonbonded_interaction_group.py @@ -20,7 +20,7 @@ def set_random_seed(): @pytest.fixture -def test_system(): +def example_system(): pdb_path = "tests/data/5dfr_solv_equil.pdb" host_pdb = app.PDBFile(pdb_path) ff = app.ForceField("amber99sbildn.xml", "tip3p.xml") @@ -32,8 +32,8 @@ def test_system(): @pytest.fixture -def ref_nonbonded_potential(test_system): - host_system, _, _ = test_system +def example_nonbonded_potential(example_system): + host_system, _, _ = example_system host_fns, _ = openmm_deserializer.deserialize_system(host_system, cutoff=1.0) nonbonded_fn = None @@ -45,14 +45,14 @@ def ref_nonbonded_potential(test_system): @pytest.fixture -def test_conf(test_system): - _, host_coords, _ = test_system +def example_coords(example_system): + _, host_coords, _ = example_system return np.array([[to_md_units(x), to_md_units(y), to_md_units(z)] for x, y, z in host_coords]) @pytest.fixture -def test_box(test_system): - _, _, box = test_system +def example_box(example_system): + _, _, box = example_system return np.asarray(box / box.unit) @@ -62,20 +62,29 @@ def test_box(test_system): @pytest.mark.parametrize("num_row_atoms", [1, 15]) @pytest.mark.parametrize("num_atoms", [33, 231, 1050]) def test_nonbonded_interaction_group_correctness( - num_atoms, num_row_atoms, precision, rtol, atol, cutoff, beta, ref_nonbonded_potential, test_conf, test_box + num_atoms, + num_row_atoms, + precision, + rtol, + atol, + cutoff, + beta, + example_nonbonded_potential, + example_coords, + example_box, ): - test_conf = test_conf[:num_atoms] - test_params = ref_nonbonded_potential.params[:num_atoms, :] - test_lambda_offset_idxs = np.zeros(num_atoms, dtype=np.int32) + coords = example_coords[:num_atoms] + params = example_nonbonded_potential.params[:num_atoms, :] + lambda_offset_idxs = np.zeros(num_atoms, dtype=np.int32) def make_reference_nonbonded(lambda_plane_idxs): return prepare_reference_nonbonded( - params=test_params, + params=params, exclusion_idxs=np.array([], dtype=np.int32), scales=np.zeros((0, 2), dtype=np.float64), lambda_plane_idxs=lambda_plane_idxs, - lambda_offset_idxs=test_lambda_offset_idxs, + lambda_offset_idxs=lambda_offset_idxs, beta=beta, cutoff=cutoff, ) @@ -94,15 +103,15 @@ def ref_ixngroups(*args): test_ixngroups = NonbondedInteractionGroup( np.arange(0, num_row_atoms, dtype=np.int32), np.zeros(num_atoms, dtype=np.int32), # lambda plane indices - test_lambda_offset_idxs, + lambda_offset_idxs, beta, cutoff, ) GradientTest().compare_forces( - test_conf, - test_params, - test_box, + coords, + params, + example_box, lamb=0.1, ref_potential=ref_ixngroups, test_potential=test_ixngroups, From fcd419922d79f35514996bda6b8ee09ecf348f71 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Tue, 15 Feb 2022 18:18:25 -0800 Subject: [PATCH 14/51] Test with random non-contiguous row atom indices --- tests/test_nonbonded_interaction_group.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/tests/test_nonbonded_interaction_group.py b/tests/test_nonbonded_interaction_group.py index ad4b39c8e..bef5e4eb2 100644 --- a/tests/test_nonbonded_interaction_group.py +++ b/tests/test_nonbonded_interaction_group.py @@ -15,10 +15,15 @@ @pytest.fixture(autouse=True) def set_random_seed(): - np.random.seed(2021) + np.random.seed(2022) yield +@pytest.fixture() +def rng(): + return np.random.default_rng(2022) + + @pytest.fixture def example_system(): pdb_path = "tests/data/5dfr_solv_equil.pdb" @@ -72,6 +77,7 @@ def test_nonbonded_interaction_group_correctness( example_nonbonded_potential, example_coords, example_box, + rng, ): coords = example_coords[:num_atoms] @@ -91,17 +97,17 @@ def make_reference_nonbonded(lambda_plane_idxs): ref_allpairs = make_reference_nonbonded(np.zeros(num_atoms, dtype=np.int32)) - num_col_atoms = num_atoms - num_row_atoms + row_atom_idxs = rng.choice(num_atoms, size=num_row_atoms, replace=False).astype(np.int32) + lambda_plane_idxs = np.zeros(num_atoms, dtype=np.int32) + lambda_plane_idxs[row_atom_idxs] = 1 - ref_allpairs_minus_ixngroups = make_reference_nonbonded( - np.concatenate((np.zeros(num_row_atoms, dtype=np.int32), np.ones(num_col_atoms, dtype=np.int32))) - ) + ref_allpairs_minus_ixngroups = make_reference_nonbonded(lambda_plane_idxs) def ref_ixngroups(*args): return ref_allpairs(*args) - ref_allpairs_minus_ixngroups(*args) test_ixngroups = NonbondedInteractionGroup( - np.arange(0, num_row_atoms, dtype=np.int32), + row_atom_idxs, np.zeros(num_atoms, dtype=np.int32), # lambda plane indices lambda_offset_idxs, beta, From 4f4bc51cade089c27f804519ef0e8fa04d6a267e Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Tue, 15 Feb 2022 18:46:10 -0800 Subject: [PATCH 15/51] Remove obsolete comments We now rebuild the neighborlist whenever the box changes at all --- timemachine/cpp/src/nonbonded_all_pairs.cu | 2 -- timemachine/cpp/src/nonbonded_interaction_group.cu | 2 -- 2 files changed, 4 deletions(-) diff --git a/timemachine/cpp/src/nonbonded_all_pairs.cu b/timemachine/cpp/src/nonbonded_all_pairs.cu index 6ac40b3c1..2fe773cf1 100644 --- a/timemachine/cpp/src/nonbonded_all_pairs.cu +++ b/timemachine/cpp/src/nonbonded_all_pairs.cu @@ -258,8 +258,6 @@ void NonbondedAllPairs::execute_device( dim3 dimGrid(B, 3, 1); // (ytz) see if we need to rebuild the neighborlist. - // (ytz + jfass): note that this logic needs to change if we use NPT later on since a resize in the box - // can introduce new interactions. k_check_rebuild_coords_and_box <<>>(N, d_x, d_nblist_x_, d_box, d_nblist_box_, nblist_padding_, d_rebuild_nblist_); gpuErrchk(cudaPeekAtLastError()); diff --git a/timemachine/cpp/src/nonbonded_interaction_group.cu b/timemachine/cpp/src/nonbonded_interaction_group.cu index 9b55d9ba3..64014e406 100644 --- a/timemachine/cpp/src/nonbonded_interaction_group.cu +++ b/timemachine/cpp/src/nonbonded_interaction_group.cu @@ -300,8 +300,6 @@ void NonbondedInteractionGroup::execute_device( const int B = ceil_divide(N_, tpb); // (ytz) see if we need to rebuild the neighborlist. - // (ytz + jfass): note that this logic needs to change if we use NPT later on since a resize in the box - // can introduce new interactions. k_check_rebuild_coords_and_box <<>>(N_, d_x, d_nblist_x_, d_box, d_nblist_box_, nblist_padding_, d_rebuild_nblist_); From 6185c4b367ade5964881d92e6bc1ee2071f13181 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Wed, 16 Feb 2022 09:53:22 -0800 Subject: [PATCH 16/51] Define N as const --- timemachine/cpp/src/kernels/k_nonbonded.cuh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/timemachine/cpp/src/kernels/k_nonbonded.cuh b/timemachine/cpp/src/kernels/k_nonbonded.cuh index a542795e6..dc20825d9 100644 --- a/timemachine/cpp/src/kernels/k_nonbonded.cuh +++ b/timemachine/cpp/src/kernels/k_nonbonded.cuh @@ -298,7 +298,7 @@ void __device__ v_nonbonded_unified( // int lambda_offset_i = atom_i_idx < N ? lambda_offset_idxs[atom_i_idx] : 0; // int lambda_plane_i = atom_i_idx < N ? lambda_plane_idxs[atom_i_idx] : 0; - int N = NC + NR; + const int N = NC + NR; if (NR != 0) { atom_i_idx += NC; @@ -557,7 +557,7 @@ void __global__ k_nonbonded_unified( int row_block_idx = ixn_tiles[tile_idx]; int atom_i_idx = row_block_idx * 32 + threadIdx.x; - int N = NC + NR; + const int N = NC + NR; if (NR != 0) { atom_i_idx += NC; From efa4828f649dc93bf2aa345cfe007ec1266531ee Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Wed, 16 Feb 2022 09:57:20 -0800 Subject: [PATCH 17/51] Move input validation to top --- .../cpp/src/nonbonded_interaction_group.cu | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/timemachine/cpp/src/nonbonded_interaction_group.cu b/timemachine/cpp/src/nonbonded_interaction_group.cu index 64014e406..b4b7863ae 100644 --- a/timemachine/cpp/src/nonbonded_interaction_group.cu +++ b/timemachine/cpp/src/nonbonded_interaction_group.cu @@ -70,6 +70,14 @@ NonbondedInteractionGroup::NonbondedInteractionGroup( throw std::runtime_error("row_atom_idxs must be nonempty"); } + if (lambda_offset_idxs.size() != N_) { + throw std::runtime_error("lambda offset idxs need to have size N"); + } + + if (lambda_offset_idxs.size() != lambda_plane_idxs.size()) { + throw std::runtime_error("lambda offset idxs and plane idxs need to be equivalent"); + } + // compute set of column atoms as set difference std::vector all_atom_idxs(N_); std::iota(all_atom_idxs.begin(), all_atom_idxs.end(), 0); @@ -81,14 +89,6 @@ NonbondedInteractionGroup::NonbondedInteractionGroup( row_atom_idxs.end(), std::inserter(col_atom_idxs, col_atom_idxs.end())); - if (lambda_offset_idxs.size() != N_) { - throw std::runtime_error("lambda offset idxs need to have size N"); - } - - if (lambda_offset_idxs.size() != lambda_plane_idxs.size()) { - throw std::runtime_error("lambda offset idxs and plane idxs need to be equivalent"); - } - std::vector col_atom_idxs_v(set_to_vector(col_atom_idxs)); gpuErrchk(cudaMalloc(&d_col_atom_idxs_, NC_ * sizeof(*d_col_atom_idxs_))); gpuErrchk( From 3b3ea55c89d5b03ad25bf10cc8f07f92135f5b04 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Wed, 16 Feb 2022 10:29:23 -0800 Subject: [PATCH 18/51] Remove guard We should return early if ixn count is zero. Will add a failing test, then implement as an early return --- .../cpp/src/nonbonded_interaction_group.cu | 41 +++++++++---------- 1 file changed, 19 insertions(+), 22 deletions(-) diff --git a/timemachine/cpp/src/nonbonded_interaction_group.cu b/timemachine/cpp/src/nonbonded_interaction_group.cu index b4b7863ae..0de1ee305 100644 --- a/timemachine/cpp/src/nonbonded_interaction_group.cu +++ b/timemachine/cpp/src/nonbonded_interaction_group.cu @@ -404,29 +404,26 @@ void NonbondedInteractionGroup::execute_device( kernel_idx |= d_du_dx ? 1 << 2 : 0; kernel_idx |= d_u ? 1 << 3 : 0; - if (p_ixn_count_[0] > 0) { - - kernel_ptrs_[kernel_idx]<<>>( - NC_, - NR_, - d_sorted_x_, - d_sorted_p_, - d_box, - d_sorted_dp_dl_, - d_sorted_w_, - d_sorted_dw_dl_, - beta_, - cutoff_, - nblist_.get_ixn_tiles(), - nblist_.get_ixn_atoms(), - d_sorted_du_dx_, - d_sorted_du_dp_, - d_du_dl, // switch to nullptr if we don't request du_dl - d_u // switch to nullptr if we don't request energies - ); + kernel_ptrs_[kernel_idx]<<>>( + NC_, + NR_, + d_sorted_x_, + d_sorted_p_, + d_box, + d_sorted_dp_dl_, + d_sorted_w_, + d_sorted_dw_dl_, + beta_, + cutoff_, + nblist_.get_ixn_tiles(), + nblist_.get_ixn_atoms(), + d_sorted_du_dx_, + d_sorted_du_dp_, + d_du_dl, // switch to nullptr if we don't request du_dl + d_u // switch to nullptr if we don't request energies + ); - gpuErrchk(cudaPeekAtLastError()); - } + gpuErrchk(cudaPeekAtLastError()); // coords are N,3 if (d_du_dx) { From d0f154a6c65a8f2e4fbf7e6f9610c6f17e295f70 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Wed, 16 Feb 2022 10:27:28 -0800 Subject: [PATCH 19/51] Add test for case with zero interactions --- tests/test_nonbonded_interaction_group.py | 31 +++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/tests/test_nonbonded_interaction_group.py b/tests/test_nonbonded_interaction_group.py index bef5e4eb2..a2eb6ffee 100644 --- a/tests/test_nonbonded_interaction_group.py +++ b/tests/test_nonbonded_interaction_group.py @@ -61,6 +61,37 @@ def example_box(example_system): return np.asarray(box / box.unit) +def test_nonbonded_interaction_group_zero_interactions(num_row_atoms, rng: np.random.Generator): + num_atoms = 33 + num_row_atoms = 15 + beta = 2.0 + lamb = 0.1 + cutoff = 1.1 + box = 10.0 * np.eye(3) + coords = rng.uniform(0, 1, size=(num_atoms, 3)) + row_atom_idxs = rng.choice(num_atoms, size=num_row_atoms, replace=False).astype(np.int32) + + # shift row atoms in x by twice the cutoff + coords[row_atom_idxs, 0] += 2 * cutoff + + params = rng.uniform(0, 1, size=(num_atoms, 3)) + + potential = NonbondedInteractionGroup( + row_atom_idxs, + np.zeros(num_atoms, dtype=np.int32), + np.zeros(num_atoms, dtype=np.int32), + beta, + cutoff, + ) + + du_dx, du_dp, du_dl, u = potential.unbound_impl(np.float64).execute(coords, params, box, lamb) + + assert (du_dx == 0).all() + assert (du_dp == 0).all() + assert du_dl == 0 + assert u == 0 + + @pytest.mark.parametrize("beta", [2.0]) @pytest.mark.parametrize("cutoff", [1.1]) @pytest.mark.parametrize("precision,rtol,atol", [(np.float64, 1e-8, 1e-8), (np.float32, 1e-4, 5e-4)]) From 3eb51ceb7041cbc6471e7bb187f4bb031b47f62d Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Wed, 16 Feb 2022 10:44:12 -0800 Subject: [PATCH 20/51] Fix typo --- tests/test_nonbonded_interaction_group.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_nonbonded_interaction_group.py b/tests/test_nonbonded_interaction_group.py index a2eb6ffee..75a6fd954 100644 --- a/tests/test_nonbonded_interaction_group.py +++ b/tests/test_nonbonded_interaction_group.py @@ -61,7 +61,7 @@ def example_box(example_system): return np.asarray(box / box.unit) -def test_nonbonded_interaction_group_zero_interactions(num_row_atoms, rng: np.random.Generator): +def test_nonbonded_interaction_group_zero_interactions(rng: np.random.Generator): num_atoms = 33 num_row_atoms = 15 beta = 2.0 From 0cd9168ee8e7cea80b827b09b950ada793de4e37 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Wed, 16 Feb 2022 10:46:53 -0800 Subject: [PATCH 21/51] Return early if neighborlist is empty --- timemachine/cpp/src/nonbonded_interaction_group.cu | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/timemachine/cpp/src/nonbonded_interaction_group.cu b/timemachine/cpp/src/nonbonded_interaction_group.cu index 0de1ee305..ec482ba77 100644 --- a/timemachine/cpp/src/nonbonded_interaction_group.cu +++ b/timemachine/cpp/src/nonbonded_interaction_group.cu @@ -362,6 +362,11 @@ void NonbondedInteractionGroup::execute_device( gpuErrchk(cudaPeekAtLastError()); } + // if the neighborlist is empty, we can return early + if (p_ixn_count_[0] == 0) { + return; + } + // do parameter interpolation here if (Interpolated) { CUresult result = compute_permute_interpolated_.configure(dimGrid, tpb, 0, stream) From 05b1b35a1f6600973025b5944f88000b3f0ba05c Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Wed, 16 Feb 2022 11:11:25 -0800 Subject: [PATCH 22/51] Attempt to clarify condition --- timemachine/cpp/src/kernels/k_nonbonded.cuh | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/timemachine/cpp/src/kernels/k_nonbonded.cuh b/timemachine/cpp/src/kernels/k_nonbonded.cuh index dc20825d9..00e1e7f5a 100644 --- a/timemachine/cpp/src/kernels/k_nonbonded.cuh +++ b/timemachine/cpp/src/kernels/k_nonbonded.cuh @@ -391,12 +391,17 @@ void __device__ v_nonbonded_unified( d2ij += delta_w * delta_w; } + const bool valid_ij = + atom_i_idx < N && + ((NR == 0) ? atom_i_idx < atom_j_idx && atom_j_idx < N // all-pairs case, only compute the upper tri + // 0 <= i < N, i < j < N + : atom_j_idx < NC); // ixn groups case, compute all pairwise ixns + // NC <= i < N, 0 <= j < NC + // (ytz): note that d2ij must be *strictly* less than cutoff_squared. This is because we set the // non-interacting atoms to exactly real_cutoff*real_cutoff. This ensures that atoms who's 4th dimension // is set to cutoff are non-interacting. - if (d2ij < cutoff_squared && atom_i_idx < N && - (NR == 0 && atom_j_idx > atom_i_idx && atom_j_idx < N || NR != 0 && atom_j_idx < NC)) { - + if (d2ij < cutoff_squared && valid_ij) { // electrostatics RealType u; RealType es_prefactor; From 43fd5223e6432b765c741d9d618740102998a8f0 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Wed, 16 Feb 2022 11:17:13 -0800 Subject: [PATCH 23/51] Move stray comment, stream synchronization --- timemachine/cpp/src/nonbonded_all_pairs.cu | 5 +++-- timemachine/cpp/src/nonbonded_interaction_group.cu | 5 +++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/timemachine/cpp/src/nonbonded_all_pairs.cu b/timemachine/cpp/src/nonbonded_all_pairs.cu index 2fe773cf1..64cfafa52 100644 --- a/timemachine/cpp/src/nonbonded_all_pairs.cu +++ b/timemachine/cpp/src/nonbonded_all_pairs.cu @@ -285,12 +285,14 @@ void NonbondedAllPairs::execute_device( gpuErrchk(cudaMemcpyAsync( p_ixn_count_, nblist_.get_ixn_count(), 1 * sizeof(*p_ixn_count_), cudaMemcpyDeviceToHost, stream)); + // this stream needs to be synchronized so we can be sure that p_ixn_count_ is properly set. + cudaStreamSynchronize(stream); + std::vector h_box(9); gpuErrchk(cudaMemcpyAsync(&h_box[0], d_box, 3 * 3 * sizeof(*d_box), cudaMemcpyDeviceToHost, stream)); // Verify that the cutoff and box size are valid together. If cutoff is greater than half the box // then a particle can interact with multiple periodic copies. const double db_cutoff = (cutoff_ + nblist_padding_) * 2; - cudaStreamSynchronize(stream); // Verify that box is orthogonal and the width of the box in all dimensions is greater than twice the cutoff for (int i = 0; i < 9; i++) { if (i == 0 || i == 4 || i == 8) { @@ -324,7 +326,6 @@ void NonbondedAllPairs::execute_device( gpuErrchk(cudaMemsetAsync(d_sorted_dp_dl_, 0, N * 3 * sizeof(*d_sorted_dp_dl_), stream)) } - // this stream needs to be synchronized so we can be sure that p_ixn_count_ is properly set. // reset buffers and sorted accumulators if (d_du_dx) { gpuErrchk(cudaMemsetAsync(d_sorted_du_dx_, 0, N * 3 * sizeof(*d_sorted_du_dx_), stream)) diff --git a/timemachine/cpp/src/nonbonded_interaction_group.cu b/timemachine/cpp/src/nonbonded_interaction_group.cu index ec482ba77..2ab29cb93 100644 --- a/timemachine/cpp/src/nonbonded_interaction_group.cu +++ b/timemachine/cpp/src/nonbonded_interaction_group.cu @@ -335,12 +335,14 @@ void NonbondedInteractionGroup::execute_device( gpuErrchk(cudaMemcpyAsync( p_ixn_count_, nblist_.get_ixn_count(), 1 * sizeof(*p_ixn_count_), cudaMemcpyDeviceToHost, stream)); + // this stream needs to be synchronized so we can be sure that p_ixn_count_ is properly set. + cudaStreamSynchronize(stream); + std::vector h_box(9); gpuErrchk(cudaMemcpyAsync(&h_box[0], d_box, 3 * 3 * sizeof(*d_box), cudaMemcpyDeviceToHost, stream)); // Verify that the cutoff and box size are valid together. If cutoff is greater than half the box // then a particle can interact with multiple periodic copies. const double db_cutoff = (cutoff_ + nblist_padding_) * 2; - cudaStreamSynchronize(stream); // Verify that box is orthogonal and the width of the box in all dimensions is greater than twice the cutoff for (int i = 0; i < 9; i++) { @@ -380,7 +382,6 @@ void NonbondedInteractionGroup::execute_device( gpuErrchk(cudaMemsetAsync(d_sorted_dp_dl_, 0, N * 3 * sizeof(*d_sorted_dp_dl_), stream)) } - // this stream needs to be synchronized so we can be sure that p_ixn_count_ is properly set. // reset buffers and sorted accumulators if (d_du_dx) { gpuErrchk(cudaMemsetAsync(d_sorted_du_dx_, 0, N * 3 * sizeof(*d_sorted_du_dx_), stream)) From 653fe4d351abdd04a6254b88e05614c2ef25b4d6 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Wed, 16 Feb 2022 11:30:13 -0800 Subject: [PATCH 24/51] Remove print statements, make errors more informative --- timemachine/cpp/src/nonbonded_all_pairs.cu | 10 ++++++---- timemachine/cpp/src/nonbonded_interaction_group.cu | 10 ++++++---- 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/timemachine/cpp/src/nonbonded_all_pairs.cu b/timemachine/cpp/src/nonbonded_all_pairs.cu index 64cfafa52..b90956b2f 100644 --- a/timemachine/cpp/src/nonbonded_all_pairs.cu +++ b/timemachine/cpp/src/nonbonded_all_pairs.cu @@ -239,15 +239,17 @@ void NonbondedAllPairs::execute_device( // g. note that du/dl is not an exact per-particle du/dl - it is only used for reduction purposes. if (N != N_) { - std::cout << N << " " << N_ << std::endl; - throw std::runtime_error("NonbondedAllPairs::execute_device() N != N_"); + throw std::runtime_error( + "NonbondedAllPairs::execute_device(): expected N == N_, got N=" + std::to_string(N) + + ", N_=" + std::to_string(N_)); } const int M = Interpolated ? 2 : 1; if (P != M * N_ * 3) { - std::cout << P << " " << N_ << std::endl; - throw std::runtime_error("NonbondedAllPairs::execute_device() P != M*N_*3"); + throw std::runtime_error( + "NonbondedAllPairs::execute_device(): expected P == M*N_*3, got P=" + std::to_string(P) + + ", M*N_*3=" + std::to_string(M * N_ * 3)); } // identify which tiles contain interpolated parameters diff --git a/timemachine/cpp/src/nonbonded_interaction_group.cu b/timemachine/cpp/src/nonbonded_interaction_group.cu index 2ab29cb93..8919d1966 100644 --- a/timemachine/cpp/src/nonbonded_interaction_group.cu +++ b/timemachine/cpp/src/nonbonded_interaction_group.cu @@ -283,15 +283,17 @@ void NonbondedInteractionGroup::execute_device( // g. note that du/dl is not an exact per-particle du/dl - it is only used for reduction purposes. if (N != N_) { - std::cout << N << " " << N_ << std::endl; - throw std::runtime_error("NonbondedInteractionGroup::execute_device() N != N_"); + throw std::runtime_error( + "NonbondedAllPairs::execute_device(): expected N == N_, got N=" + std::to_string(N) + + ", N_=" + std::to_string(N_)); } const int M = Interpolated ? 2 : 1; if (P != M * N_ * 3) { - std::cout << P << " " << N_ << std::endl; - throw std::runtime_error("NonbondedInteractionGroup::execute_device() P != M*N_*3"); + throw std::runtime_error( + "NonbondedAllPairs::execute_device(): expected P == M*N_*3, got P=" + std::to_string(P) + + ", M*N_*3=" + std::to_string(M * N_ * 3)); } // identify which tiles contain interpolated parameters From a70a3309fce2c2edc9ab93322e89a0b7119af848 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Wed, 16 Feb 2022 11:40:43 -0800 Subject: [PATCH 25/51] Add comments for clarity, remove extra placeholders --- timemachine/cpp/src/nonbonded_all_pairs.hpp | 27 ++++++++++--------- .../cpp/src/nonbonded_interaction_group.hpp | 27 ++++++++++--------- 2 files changed, 30 insertions(+), 24 deletions(-) diff --git a/timemachine/cpp/src/nonbonded_all_pairs.hpp b/timemachine/cpp/src/nonbonded_all_pairs.hpp index 072eafb38..b1f0ee689 100644 --- a/timemachine/cpp/src/nonbonded_all_pairs.hpp +++ b/timemachine/cpp/src/nonbonded_all_pairs.hpp @@ -49,21 +49,24 @@ template class NonbondedAllPairs : public unsigned int *d_perm_; // hilbert curve permutation - double *d_w_; // - double *d_dw_dl_; // - - double *d_sorted_x_; // - double *d_sorted_w_; // - double *d_sorted_dw_dl_; // - double *d_sorted_p_; // - double *d_unsorted_p_; // + double *d_w_; // 4D coordinates + double *d_dw_dl_; + + // "sorted" means sorted according to the hilbert curve index + // "unsorted" means in the original order of the input + double *d_sorted_x_; // sorted coordinates + double *d_sorted_w_; // sorted 4D coordinates + double *d_sorted_dw_dl_; + double *d_sorted_p_; // sorted parameters + double *d_unsorted_p_; // unsorted parameters double *d_sorted_dp_dl_; double *d_unsorted_dp_dl_; - unsigned long long *d_sorted_du_dx_; // - unsigned long long *d_sorted_du_dp_; // - unsigned long long *d_du_dp_buffer_; // + unsigned long long *d_sorted_du_dx_; + unsigned long long *d_sorted_du_dp_; + unsigned long long *d_du_dp_buffer_; - unsigned int *d_bin_to_idx_; + // used for hilbert sorting + unsigned int *d_bin_to_idx_; // mapping from 256x256x256 grid to hilbert curve index unsigned int *d_sort_keys_in_; unsigned int *d_sort_keys_out_; unsigned int *d_sort_vals_in_; diff --git a/timemachine/cpp/src/nonbonded_interaction_group.hpp b/timemachine/cpp/src/nonbonded_interaction_group.hpp index 394bc1bb2..50f7a3b0e 100644 --- a/timemachine/cpp/src/nonbonded_interaction_group.hpp +++ b/timemachine/cpp/src/nonbonded_interaction_group.hpp @@ -55,21 +55,24 @@ template class NonbondedInteractionGroup unsigned int *d_perm_; // hilbert curve permutation - double *d_w_; // - double *d_dw_dl_; // - - double *d_sorted_x_; // - double *d_sorted_w_; // - double *d_sorted_dw_dl_; // - double *d_sorted_p_; // - double *d_unsorted_p_; // + double *d_w_; // 4D coordinates + double *d_dw_dl_; + + // "sorted" means sorted according to the hilbert curve index + // "unsorted" means in the original order of the input + double *d_sorted_x_; // sorted coordinates + double *d_sorted_w_; // sorted 4D coordinates + double *d_sorted_dw_dl_; + double *d_sorted_p_; // sorted parameters + double *d_unsorted_p_; // unsorted parameters double *d_sorted_dp_dl_; double *d_unsorted_dp_dl_; - unsigned long long *d_sorted_du_dx_; // - unsigned long long *d_sorted_du_dp_; // - unsigned long long *d_du_dp_buffer_; // + unsigned long long *d_sorted_du_dx_; + unsigned long long *d_sorted_du_dp_; + unsigned long long *d_du_dp_buffer_; - unsigned int *d_bin_to_idx_; + // used for hilbert sorting + unsigned int *d_bin_to_idx_; // mapping from 256x256x256 grid to hilbert curve index unsigned int *d_sort_keys_in_; unsigned int *d_sort_keys_out_; unsigned int *d_sort_vals_in_; From 690a57cf4f00cf55e98abfb14e51f5fb34baa471 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Wed, 16 Feb 2022 11:44:15 -0800 Subject: [PATCH 26/51] Remove obsolete comment --- timemachine/cpp/src/nonbonded_interaction_group.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/timemachine/cpp/src/nonbonded_interaction_group.cu b/timemachine/cpp/src/nonbonded_interaction_group.cu index 8919d1966..85a2617a8 100644 --- a/timemachine/cpp/src/nonbonded_interaction_group.cu +++ b/timemachine/cpp/src/nonbonded_interaction_group.cu @@ -180,7 +180,7 @@ NonbondedInteractionGroup::~NonbondedInteractionGroup() gpuErrchk(cudaFree(d_lambda_plane_idxs_)); gpuErrchk(cudaFree(d_lambda_offset_idxs_)); gpuErrchk(cudaFree(d_du_dp_buffer_)); - gpuErrchk(cudaFree(d_perm_)); // nullptr if we never built nblist + gpuErrchk(cudaFree(d_perm_)); gpuErrchk(cudaFree(d_bin_to_idx_)); gpuErrchk(cudaFree(d_sorted_x_)); From 1097ca496d5e8cbff6205b2d090827d19a8c3ff6 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Wed, 16 Feb 2022 11:46:30 -0800 Subject: [PATCH 27/51] Remove obsolete comment --- timemachine/cpp/src/nonbonded_all_pairs.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/timemachine/cpp/src/nonbonded_all_pairs.cu b/timemachine/cpp/src/nonbonded_all_pairs.cu index b90956b2f..5362d6503 100644 --- a/timemachine/cpp/src/nonbonded_all_pairs.cu +++ b/timemachine/cpp/src/nonbonded_all_pairs.cu @@ -143,7 +143,7 @@ template NonbondedAllPairs Date: Wed, 16 Feb 2022 12:01:37 -0800 Subject: [PATCH 28/51] Add description to test --- tests/test_nonbonded_interaction_group.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/tests/test_nonbonded_interaction_group.py b/tests/test_nonbonded_interaction_group.py index 75a6fd954..f05ab02fd 100644 --- a/tests/test_nonbonded_interaction_group.py +++ b/tests/test_nonbonded_interaction_group.py @@ -110,6 +110,24 @@ def test_nonbonded_interaction_group_correctness( example_box, rng, ): + """Compares with reference nonbonded_v3 potential, which computes + the sum of all pairwise interactions. This uses the identity + + U = U_A + U_B + U_AB + + where + - `U` is the all-pairs potential over all atoms + - `U_A`, `U_B` are all-pairs potentials for interacting groups A + and B, respectively + - `U_AB` is the "interaction group" potential, i.e. the sum of + pairwise interactions `(a, b)` where `a` is in `A` and `b` is in + `B` + + The quantity `U` is computed using the reference potential over + all atoms, and `U_A + U_B` computed using the reference potential + over all atoms separated into 2 lambda planes according to which + interacting group they belong + """ coords = example_coords[:num_atoms] params = example_nonbonded_potential.params[:num_atoms, :] From eb07db3088a77625e177e4c27f005ff4ac360364 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Wed, 16 Feb 2022 13:08:43 -0800 Subject: [PATCH 29/51] Add test comparing with nonbonded_v3_on_specific_pairs --- tests/test_nonbonded_interaction_group.py | 72 +++++++++++++++++++++-- 1 file changed, 67 insertions(+), 5 deletions(-) diff --git a/tests/test_nonbonded_interaction_group.py b/tests/test_nonbonded_interaction_group.py index f05ab02fd..c48f55b1a 100644 --- a/tests/test_nonbonded_interaction_group.py +++ b/tests/test_nonbonded_interaction_group.py @@ -2,6 +2,8 @@ jax.config.update("jax_enable_x64", True) +import itertools + import numpy as np import pytest from common import GradientTest, prepare_reference_nonbonded @@ -11,6 +13,7 @@ from timemachine.ff.handlers import openmm_deserializer from timemachine.lib import potentials from timemachine.lib.potentials import NonbondedInteractionGroup +from timemachine.potentials import nonbonded @pytest.fixture(autouse=True) @@ -37,7 +40,7 @@ def example_system(): @pytest.fixture -def example_nonbonded_potential(example_system): +def example_nonbonded_params(example_system): host_system, _, _ = example_system host_fns, _ = openmm_deserializer.deserialize_system(host_system, cutoff=1.0) @@ -46,7 +49,8 @@ def example_nonbonded_potential(example_system): if isinstance(f, potentials.Nonbonded): nonbonded_fn = f - return nonbonded_fn + assert nonbonded_fn is not None + return nonbonded_fn.params @pytest.fixture @@ -97,7 +101,7 @@ def test_nonbonded_interaction_group_zero_interactions(rng: np.random.Generator) @pytest.mark.parametrize("precision,rtol,atol", [(np.float64, 1e-8, 1e-8), (np.float32, 1e-4, 5e-4)]) @pytest.mark.parametrize("num_row_atoms", [1, 15]) @pytest.mark.parametrize("num_atoms", [33, 231, 1050]) -def test_nonbonded_interaction_group_correctness( +def test_nonbonded_interaction_group_consistency_allpairs( num_atoms, num_row_atoms, precision, @@ -105,7 +109,7 @@ def test_nonbonded_interaction_group_correctness( atol, cutoff, beta, - example_nonbonded_potential, + example_nonbonded_params, example_coords, example_box, rng, @@ -130,7 +134,7 @@ def test_nonbonded_interaction_group_correctness( """ coords = example_coords[:num_atoms] - params = example_nonbonded_potential.params[:num_atoms, :] + params = example_nonbonded_params[:num_atoms, :] lambda_offset_idxs = np.zeros(num_atoms, dtype=np.int32) def make_reference_nonbonded(lambda_plane_idxs): @@ -176,6 +180,64 @@ def ref_ixngroups(*args): ) +@pytest.mark.parametrize("beta", [2.0]) +@pytest.mark.parametrize("cutoff", [1.1]) +@pytest.mark.parametrize("precision,rtol,atol", [(np.float64, 1e-8, 1e-8), (np.float32, 1e-4, 5e-4)]) +@pytest.mark.parametrize("num_row_atoms", [1, 15]) +@pytest.mark.parametrize("num_atoms", [33, 231]) +def test_nonbonded_interaction_group_consistency_pairwise( + num_atoms, + num_row_atoms, + precision, + rtol, + atol, + cutoff, + beta, + example_nonbonded_params, + example_coords, + example_box, + rng, +): + """Compares with reference nonbonded_v3_on_specific_pairs potential, which computes + the sum of interactions for a list of pairs. + """ + + coords = example_coords[:num_atoms] + params = example_nonbonded_params[:num_atoms, :] + + row_atom_idxs = rng.choice(num_atoms, size=num_row_atoms, replace=False).astype(np.int32) + col_atom_idxs = np.setdiff1d(np.arange(num_atoms), row_atom_idxs) + pairs = itertools.product(row_atom_idxs, col_atom_idxs) + inds_l, inds_r = (np.array(ps) for ps in zip(*pairs)) + + def ref_ixngroups(coords, params, box, _): + vdW, electrostatics = nonbonded.nonbonded_v3_on_specific_pairs( + coords, params, box, inds_l, inds_r, beta, cutoff + ) + total_energy = jax.numpy.sum(vdW + electrostatics) + return total_energy + + test_ixngroups = NonbondedInteractionGroup( + row_atom_idxs, + np.zeros(num_atoms, dtype=np.int32), + np.zeros(num_atoms, dtype=np.int32), + beta, + cutoff, + ) + + GradientTest().compare_forces( + coords, + params, + example_box, + lamb=0.1, + ref_potential=ref_ixngroups, + test_potential=test_ixngroups, + rtol=rtol, + atol=atol, + precision=precision, + ) + + def test_nonbonded_interaction_group_invalid_indices(): def make_potential(row_atom_idxs, num_atoms): lambda_plane_idxs = [0] * num_atoms From 3e36b51ed4c1062730183f88732981a4281dfdcd Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Wed, 16 Feb 2022 14:24:09 -0800 Subject: [PATCH 30/51] Improve accuracy of comments --- timemachine/cpp/src/nonbonded_all_pairs.hpp | 8 ++++++-- timemachine/cpp/src/nonbonded_interaction_group.hpp | 10 ++++++++-- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/timemachine/cpp/src/nonbonded_all_pairs.hpp b/timemachine/cpp/src/nonbonded_all_pairs.hpp index b1f0ee689..1ad6bbdec 100644 --- a/timemachine/cpp/src/nonbonded_all_pairs.hpp +++ b/timemachine/cpp/src/nonbonded_all_pairs.hpp @@ -52,8 +52,12 @@ template class NonbondedAllPairs : public double *d_w_; // 4D coordinates double *d_dw_dl_; - // "sorted" means sorted according to the hilbert curve index - // "unsorted" means in the original order of the input + // "sorted" means + // - if hilbert sorting enabled, atoms are sorted according to the + // hilbert curve index + // - otherwise, atom ordering is preserved with respect to input + // + // "unsorted" means the atom ordering is preserved with respect to input double *d_sorted_x_; // sorted coordinates double *d_sorted_w_; // sorted 4D coordinates double *d_sorted_dw_dl_; diff --git a/timemachine/cpp/src/nonbonded_interaction_group.hpp b/timemachine/cpp/src/nonbonded_interaction_group.hpp index 50f7a3b0e..b3ab817bb 100644 --- a/timemachine/cpp/src/nonbonded_interaction_group.hpp +++ b/timemachine/cpp/src/nonbonded_interaction_group.hpp @@ -58,8 +58,14 @@ template class NonbondedInteractionGroup double *d_w_; // 4D coordinates double *d_dw_dl_; - // "sorted" means sorted according to the hilbert curve index - // "unsorted" means in the original order of the input + // "sorted" means + // - if hilbert sorting enabled, atoms are sorted into contiguous + // blocks by interaction group, and each block is hilbert-sorted + // independently + // - otherwise, atoms are sorted into contiguous blocks by + // interaction group, with arbitrary ordering within each block + // + // "unsorted" means the atom ordering is preserved with respect to input double *d_sorted_x_; // sorted coordinates double *d_sorted_w_; // sorted 4D coordinates double *d_sorted_dw_dl_; From 2dfb18d5312df824c2ee70452ee2b81e99053475 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Wed, 16 Feb 2022 15:59:13 -0800 Subject: [PATCH 31/51] Rename test to something less confusing, add type annotation --- tests/test_nonbonded_interaction_group.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_nonbonded_interaction_group.py b/tests/test_nonbonded_interaction_group.py index c48f55b1a..d588b827a 100644 --- a/tests/test_nonbonded_interaction_group.py +++ b/tests/test_nonbonded_interaction_group.py @@ -112,7 +112,7 @@ def test_nonbonded_interaction_group_consistency_allpairs( example_nonbonded_params, example_coords, example_box, - rng, + rng: np.random.Generator, ): """Compares with reference nonbonded_v3 potential, which computes the sum of all pairwise interactions. This uses the identity @@ -185,7 +185,7 @@ def ref_ixngroups(*args): @pytest.mark.parametrize("precision,rtol,atol", [(np.float64, 1e-8, 1e-8), (np.float32, 1e-4, 5e-4)]) @pytest.mark.parametrize("num_row_atoms", [1, 15]) @pytest.mark.parametrize("num_atoms", [33, 231]) -def test_nonbonded_interaction_group_consistency_pairwise( +def test_nonbonded_interaction_group_consistency_specific_pairs( num_atoms, num_row_atoms, precision, From 39907d6cc2bfe9ba36d4fffdecdeb93889507103 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Wed, 16 Feb 2022 16:59:32 -0800 Subject: [PATCH 32/51] Account for difference with zero LJ parameters in test --- tests/test_nonbonded_interaction_group.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/test_nonbonded_interaction_group.py b/tests/test_nonbonded_interaction_group.py index d588b827a..3fe40095b 100644 --- a/tests/test_nonbonded_interaction_group.py +++ b/tests/test_nonbonded_interaction_group.py @@ -210,11 +210,18 @@ def test_nonbonded_interaction_group_consistency_specific_pairs( pairs = itertools.product(row_atom_idxs, col_atom_idxs) inds_l, inds_r = (np.array(ps) for ps in zip(*pairs)) + atom_is_hydrogen = (params[:, 1] == 0) & (params[:, 2] == 0) + pair_has_hydrogen = atom_is_hydrogen[inds_l] | atom_is_hydrogen[inds_r] + def ref_ixngroups(coords, params, box, _): vdW, electrostatics = nonbonded.nonbonded_v3_on_specific_pairs( coords, params, box, inds_l, inds_r, beta, cutoff ) - total_energy = jax.numpy.sum(vdW + electrostatics) + + # custom ops implementation returns du/dp = 0 for LJ params at zero + vdW_pinned = vdW.at[pair_has_hydrogen].set(0.0) + + total_energy = jax.numpy.sum(vdW_pinned + electrostatics) return total_energy test_ixngroups = NonbondedInteractionGroup( From 561ce2ffd42b75aeeb047c5d0df496768602308c Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Thu, 17 Feb 2022 08:25:37 -0800 Subject: [PATCH 33/51] Move jax reference potential out of test --- tests/test_nonbonded_interaction_group.py | 19 +++++++------------ timemachine/potentials/nonbonded.py | 13 +++++++++++++ 2 files changed, 20 insertions(+), 12 deletions(-) diff --git a/tests/test_nonbonded_interaction_group.py b/tests/test_nonbonded_interaction_group.py index 3fe40095b..f3e35e783 100644 --- a/tests/test_nonbonded_interaction_group.py +++ b/tests/test_nonbonded_interaction_group.py @@ -2,8 +2,6 @@ jax.config.update("jax_enable_x64", True) -import itertools - import numpy as np import pytest from common import GradientTest, prepare_reference_nonbonded @@ -185,7 +183,7 @@ def ref_ixngroups(*args): @pytest.mark.parametrize("precision,rtol,atol", [(np.float64, 1e-8, 1e-8), (np.float32, 1e-4, 5e-4)]) @pytest.mark.parametrize("num_row_atoms", [1, 15]) @pytest.mark.parametrize("num_atoms", [33, 231]) -def test_nonbonded_interaction_group_consistency_specific_pairs( +def test_nonbonded_interaction_group_correctness( num_atoms, num_row_atoms, precision, @@ -198,27 +196,24 @@ def test_nonbonded_interaction_group_consistency_specific_pairs( example_box, rng, ): - """Compares with reference nonbonded_v3_on_specific_pairs potential, which computes - the sum of interactions for a list of pairs. - """ + "Compares with jax reference implementation." coords = example_coords[:num_atoms] params = example_nonbonded_params[:num_atoms, :] row_atom_idxs = rng.choice(num_atoms, size=num_row_atoms, replace=False).astype(np.int32) col_atom_idxs = np.setdiff1d(np.arange(num_atoms), row_atom_idxs) - pairs = itertools.product(row_atom_idxs, col_atom_idxs) - inds_l, inds_r = (np.array(ps) for ps in zip(*pairs)) + # hydrogen atoms have lennard-jones parameters set to 0 atom_is_hydrogen = (params[:, 1] == 0) & (params[:, 2] == 0) - pair_has_hydrogen = atom_is_hydrogen[inds_l] | atom_is_hydrogen[inds_r] def ref_ixngroups(coords, params, box, _): - vdW, electrostatics = nonbonded.nonbonded_v3_on_specific_pairs( - coords, params, box, inds_l, inds_r, beta, cutoff + vdW, electrostatics, pairs = nonbonded.nonbonded_v3_interaction_groups( + coords, params, box, row_atom_idxs, col_atom_idxs, beta, cutoff ) - # custom ops implementation returns du/dp = 0 for LJ params at zero + # custom ops implementation returns du/dp = 0 for lennard-jones params at zero + pair_has_hydrogen = atom_is_hydrogen[pairs[:, 0]] | atom_is_hydrogen[pairs[:, 1]] vdW_pinned = vdW.at[pair_has_hydrogen].set(0.0) total_energy = jax.numpy.sum(vdW_pinned + electrostatics) diff --git a/timemachine/potentials/nonbonded.py b/timemachine/potentials/nonbonded.py index 0cf224eb5..f1381aea9 100644 --- a/timemachine/potentials/nonbonded.py +++ b/timemachine/potentials/nonbonded.py @@ -1,3 +1,5 @@ +import itertools + import jax.numpy as np from jax.ops import index, index_update from jax.scipy.special import erfc @@ -251,6 +253,17 @@ def apply_cutoff(x): return vdW, electrostatics +def nonbonded_v3_interaction_groups(conf, params, box, inds_l, inds_r, beta: float, cutoff: Optional[float] = None): + """Nonbonded interactions between all pairs of atoms $(i, j)$ + where $i$ is in the first set and $j$ in the second. + + See nonbonded_v3 docstring for more details + """ + pairs = np.array(list(itertools.product(inds_l, inds_r))) + vdW, electrostatics = nonbonded_v3_on_specific_pairs(conf, params, box, pairs[:, 0], pairs[:, 1], beta, cutoff) + return vdW, electrostatics, pairs + + def validate_coulomb_cutoff(cutoff=1.0, beta=2.0, threshold=1e-2): """check whether f(r) = erfc(beta * r) <= threshold at r = cutoff following https://github.com/proteneer/timemachine/pull/424#discussion_r629678467""" From dd543eadce68c14717bb2272a7954de1a38e330a Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Thu, 17 Feb 2022 08:54:14 -0800 Subject: [PATCH 34/51] Test with multiple lambda values --- tests/test_nonbonded_interaction_group.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/test_nonbonded_interaction_group.py b/tests/test_nonbonded_interaction_group.py index f3e35e783..6fa657e93 100644 --- a/tests/test_nonbonded_interaction_group.py +++ b/tests/test_nonbonded_interaction_group.py @@ -94,6 +94,7 @@ def test_nonbonded_interaction_group_zero_interactions(rng: np.random.Generator) assert u == 0 +@pytest.mark.parametrize("lamb", [0.0, 0.1]) @pytest.mark.parametrize("beta", [2.0]) @pytest.mark.parametrize("cutoff", [1.1]) @pytest.mark.parametrize("precision,rtol,atol", [(np.float64, 1e-8, 1e-8), (np.float32, 1e-4, 5e-4)]) @@ -107,6 +108,7 @@ def test_nonbonded_interaction_group_consistency_allpairs( atol, cutoff, beta, + lamb, example_nonbonded_params, example_coords, example_box, @@ -169,7 +171,7 @@ def ref_ixngroups(*args): coords, params, example_box, - lamb=0.1, + lamb=lamb, ref_potential=ref_ixngroups, test_potential=test_ixngroups, rtol=rtol, @@ -178,6 +180,7 @@ def ref_ixngroups(*args): ) +@pytest.mark.parametrize("lamb", [0.0, 0.1]) @pytest.mark.parametrize("beta", [2.0]) @pytest.mark.parametrize("cutoff", [1.1]) @pytest.mark.parametrize("precision,rtol,atol", [(np.float64, 1e-8, 1e-8), (np.float32, 1e-4, 5e-4)]) @@ -191,6 +194,7 @@ def test_nonbonded_interaction_group_correctness( atol, cutoff, beta, + lamb, example_nonbonded_params, example_coords, example_box, @@ -231,7 +235,7 @@ def ref_ixngroups(coords, params, box, _): coords, params, example_box, - lamb=0.1, + lamb=lamb, ref_potential=ref_ixngroups, test_potential=test_ixngroups, rtol=rtol, From 25fd4f6561002d4c25e5f5125783ed33eec63359 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Thu, 17 Feb 2022 13:13:00 -0800 Subject: [PATCH 35/51] Move cudaStreamSynchronize after cudaMemcpyAsync This restores the original position of stream synchronization. The CUDA docs indicate that this shouldn't matter, since transfers from device to (pageable) host memory only return once the copy has completed. But, in the interest of keeping the diff minimal, it seems best to restore the original order. https://docs.nvidia.com/cuda/cuda-runtime-api/api-sync-behavior.html --- timemachine/cpp/src/nonbonded_all_pairs.cu | 6 ++++-- timemachine/cpp/src/nonbonded_interaction_group.cu | 5 +++-- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/timemachine/cpp/src/nonbonded_all_pairs.cu b/timemachine/cpp/src/nonbonded_all_pairs.cu index 5362d6503..66235bd30 100644 --- a/timemachine/cpp/src/nonbonded_all_pairs.cu +++ b/timemachine/cpp/src/nonbonded_all_pairs.cu @@ -287,14 +287,16 @@ void NonbondedAllPairs::execute_device( gpuErrchk(cudaMemcpyAsync( p_ixn_count_, nblist_.get_ixn_count(), 1 * sizeof(*p_ixn_count_), cudaMemcpyDeviceToHost, stream)); + std::vector h_box(9); + gpuErrchk(cudaMemcpyAsync(&h_box[0], d_box, 3 * 3 * sizeof(*d_box), cudaMemcpyDeviceToHost, stream)); + // this stream needs to be synchronized so we can be sure that p_ixn_count_ is properly set. cudaStreamSynchronize(stream); - std::vector h_box(9); - gpuErrchk(cudaMemcpyAsync(&h_box[0], d_box, 3 * 3 * sizeof(*d_box), cudaMemcpyDeviceToHost, stream)); // Verify that the cutoff and box size are valid together. If cutoff is greater than half the box // then a particle can interact with multiple periodic copies. const double db_cutoff = (cutoff_ + nblist_padding_) * 2; + // Verify that box is orthogonal and the width of the box in all dimensions is greater than twice the cutoff for (int i = 0; i < 9; i++) { if (i == 0 || i == 4 || i == 8) { diff --git a/timemachine/cpp/src/nonbonded_interaction_group.cu b/timemachine/cpp/src/nonbonded_interaction_group.cu index 85a2617a8..e72369834 100644 --- a/timemachine/cpp/src/nonbonded_interaction_group.cu +++ b/timemachine/cpp/src/nonbonded_interaction_group.cu @@ -337,11 +337,12 @@ void NonbondedInteractionGroup::execute_device( gpuErrchk(cudaMemcpyAsync( p_ixn_count_, nblist_.get_ixn_count(), 1 * sizeof(*p_ixn_count_), cudaMemcpyDeviceToHost, stream)); + std::vector h_box(9); + gpuErrchk(cudaMemcpyAsync(&h_box[0], d_box, 3 * 3 * sizeof(*d_box), cudaMemcpyDeviceToHost, stream)); + // this stream needs to be synchronized so we can be sure that p_ixn_count_ is properly set. cudaStreamSynchronize(stream); - std::vector h_box(9); - gpuErrchk(cudaMemcpyAsync(&h_box[0], d_box, 3 * 3 * sizeof(*d_box), cudaMemcpyDeviceToHost, stream)); // Verify that the cutoff and box size are valid together. If cutoff is greater than half the box // then a particle can interact with multiple periodic copies. const double db_cutoff = (cutoff_ + nblist_padding_) * 2; From 860dc9df70dc2489861f0e995cc7310a60fc81a6 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Thu, 17 Feb 2022 13:54:32 -0800 Subject: [PATCH 36/51] Update reference to only compute vdW force when eps_ij != 0 Remove now-obsolete workaround --- tests/test_nonbonded_interaction_group.py | 13 ++----------- timemachine/potentials/nonbonded.py | 2 +- 2 files changed, 3 insertions(+), 12 deletions(-) diff --git a/tests/test_nonbonded_interaction_group.py b/tests/test_nonbonded_interaction_group.py index 6fa657e93..4893ab9e4 100644 --- a/tests/test_nonbonded_interaction_group.py +++ b/tests/test_nonbonded_interaction_group.py @@ -208,20 +208,11 @@ def test_nonbonded_interaction_group_correctness( row_atom_idxs = rng.choice(num_atoms, size=num_row_atoms, replace=False).astype(np.int32) col_atom_idxs = np.setdiff1d(np.arange(num_atoms), row_atom_idxs) - # hydrogen atoms have lennard-jones parameters set to 0 - atom_is_hydrogen = (params[:, 1] == 0) & (params[:, 2] == 0) - def ref_ixngroups(coords, params, box, _): - vdW, electrostatics, pairs = nonbonded.nonbonded_v3_interaction_groups( + vdW, electrostatics, _ = nonbonded.nonbonded_v3_interaction_groups( coords, params, box, row_atom_idxs, col_atom_idxs, beta, cutoff ) - - # custom ops implementation returns du/dp = 0 for lennard-jones params at zero - pair_has_hydrogen = atom_is_hydrogen[pairs[:, 0]] | atom_is_hydrogen[pairs[:, 1]] - vdW_pinned = vdW.at[pair_has_hydrogen].set(0.0) - - total_energy = jax.numpy.sum(vdW_pinned + electrostatics) - return total_energy + return jax.numpy.sum(vdW + electrostatics) test_ixngroups = NonbondedInteractionGroup( row_atom_idxs, diff --git a/timemachine/potentials/nonbonded.py b/timemachine/potentials/nonbonded.py index f1381aea9..557da0a0e 100644 --- a/timemachine/potentials/nonbonded.py +++ b/timemachine/potentials/nonbonded.py @@ -244,7 +244,7 @@ def apply_cutoff(x): # vdW by Lennard-Jones sig_ij = apply_cutoff(sig[inds_l] + sig[inds_r]) eps_ij = apply_cutoff(eps[inds_l] * eps[inds_r]) - vdW = lennard_jones(dij, sig_ij, eps_ij) + vdW = np.where(eps_ij != 0, lennard_jones(dij, sig_ij, eps_ij), 0) # Electrostatics by direct-space part of PME qij = apply_cutoff(charges[inds_l] * charges[inds_r]) From b7ebc26ba9240b54380470487307193b81438407 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Thu, 17 Feb 2022 14:28:07 -0800 Subject: [PATCH 37/51] Rename coords -> conf --- tests/test_nonbonded_interaction_group.py | 28 +++++++++++------------ 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/tests/test_nonbonded_interaction_group.py b/tests/test_nonbonded_interaction_group.py index 4893ab9e4..110de07d3 100644 --- a/tests/test_nonbonded_interaction_group.py +++ b/tests/test_nonbonded_interaction_group.py @@ -52,9 +52,9 @@ def example_nonbonded_params(example_system): @pytest.fixture -def example_coords(example_system): - _, host_coords, _ = example_system - return np.array([[to_md_units(x), to_md_units(y), to_md_units(z)] for x, y, z in host_coords]) +def example_conf(example_system): + _, host_conf, _ = example_system + return np.array([[to_md_units(x), to_md_units(y), to_md_units(z)] for x, y, z in host_conf]) @pytest.fixture @@ -70,11 +70,11 @@ def test_nonbonded_interaction_group_zero_interactions(rng: np.random.Generator) lamb = 0.1 cutoff = 1.1 box = 10.0 * np.eye(3) - coords = rng.uniform(0, 1, size=(num_atoms, 3)) + conf = rng.uniform(0, 1, size=(num_atoms, 3)) row_atom_idxs = rng.choice(num_atoms, size=num_row_atoms, replace=False).astype(np.int32) # shift row atoms in x by twice the cutoff - coords[row_atom_idxs, 0] += 2 * cutoff + conf[row_atom_idxs, 0] += 2 * cutoff params = rng.uniform(0, 1, size=(num_atoms, 3)) @@ -86,7 +86,7 @@ def test_nonbonded_interaction_group_zero_interactions(rng: np.random.Generator) cutoff, ) - du_dx, du_dp, du_dl, u = potential.unbound_impl(np.float64).execute(coords, params, box, lamb) + du_dx, du_dp, du_dl, u = potential.unbound_impl(np.float64).execute(conf, params, box, lamb) assert (du_dx == 0).all() assert (du_dp == 0).all() @@ -110,7 +110,7 @@ def test_nonbonded_interaction_group_consistency_allpairs( beta, lamb, example_nonbonded_params, - example_coords, + example_conf, example_box, rng: np.random.Generator, ): @@ -133,7 +133,7 @@ def test_nonbonded_interaction_group_consistency_allpairs( interacting group they belong """ - coords = example_coords[:num_atoms] + conf = example_conf[:num_atoms] params = example_nonbonded_params[:num_atoms, :] lambda_offset_idxs = np.zeros(num_atoms, dtype=np.int32) @@ -168,7 +168,7 @@ def ref_ixngroups(*args): ) GradientTest().compare_forces( - coords, + conf, params, example_box, lamb=lamb, @@ -196,21 +196,21 @@ def test_nonbonded_interaction_group_correctness( beta, lamb, example_nonbonded_params, - example_coords, + example_conf, example_box, rng, ): "Compares with jax reference implementation." - coords = example_coords[:num_atoms] + conf = example_conf[:num_atoms] params = example_nonbonded_params[:num_atoms, :] row_atom_idxs = rng.choice(num_atoms, size=num_row_atoms, replace=False).astype(np.int32) col_atom_idxs = np.setdiff1d(np.arange(num_atoms), row_atom_idxs) - def ref_ixngroups(coords, params, box, _): + def ref_ixngroups(conf, params, box, _): vdW, electrostatics, _ = nonbonded.nonbonded_v3_interaction_groups( - coords, params, box, row_atom_idxs, col_atom_idxs, beta, cutoff + conf, params, box, row_atom_idxs, col_atom_idxs, beta, cutoff ) return jax.numpy.sum(vdW + electrostatics) @@ -223,7 +223,7 @@ def ref_ixngroups(coords, params, box, _): ) GradientTest().compare_forces( - coords, + conf, params, example_box, lamb=lamb, From 89ee22f452c7c21f35ed559e17930fecb2e8bfa6 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Thu, 17 Feb 2022 14:29:00 -0800 Subject: [PATCH 38/51] Reorder tests --- tests/test_nonbonded_interaction_group.py | 144 +++++++++++----------- 1 file changed, 72 insertions(+), 72 deletions(-) diff --git a/tests/test_nonbonded_interaction_group.py b/tests/test_nonbonded_interaction_group.py index 110de07d3..4ab351d32 100644 --- a/tests/test_nonbonded_interaction_group.py +++ b/tests/test_nonbonded_interaction_group.py @@ -63,6 +63,23 @@ def example_box(example_system): return np.asarray(box / box.unit) +def test_nonbonded_interaction_group_invalid_indices(): + def make_potential(row_atom_idxs, num_atoms): + lambda_plane_idxs = [0] * num_atoms + lambda_offset_idxs = [0] * num_atoms + return NonbondedInteractionGroup(row_atom_idxs, lambda_plane_idxs, lambda_offset_idxs, 1.0, 1.0).unbound_impl( + np.float64 + ) + + with pytest.raises(RuntimeError) as e: + make_potential([], 1) + assert "row_atom_idxs must be nonempty" in str(e) + + with pytest.raises(RuntimeError) as e: + make_potential([1, 1], 3) + assert "atom indices must be unique" in str(e) + + def test_nonbonded_interaction_group_zero_interactions(rng: np.random.Generator): num_atoms = 33 num_row_atoms = 15 @@ -99,8 +116,8 @@ def test_nonbonded_interaction_group_zero_interactions(rng: np.random.Generator) @pytest.mark.parametrize("cutoff", [1.1]) @pytest.mark.parametrize("precision,rtol,atol", [(np.float64, 1e-8, 1e-8), (np.float32, 1e-4, 5e-4)]) @pytest.mark.parametrize("num_row_atoms", [1, 15]) -@pytest.mark.parametrize("num_atoms", [33, 231, 1050]) -def test_nonbonded_interaction_group_consistency_allpairs( +@pytest.mark.parametrize("num_atoms", [33, 231]) +def test_nonbonded_interaction_group_correctness( num_atoms, num_row_atoms, precision, @@ -112,57 +129,26 @@ def test_nonbonded_interaction_group_consistency_allpairs( example_nonbonded_params, example_conf, example_box, - rng: np.random.Generator, + rng, ): - """Compares with reference nonbonded_v3 potential, which computes - the sum of all pairwise interactions. This uses the identity - - U = U_A + U_B + U_AB - - where - - `U` is the all-pairs potential over all atoms - - `U_A`, `U_B` are all-pairs potentials for interacting groups A - and B, respectively - - `U_AB` is the "interaction group" potential, i.e. the sum of - pairwise interactions `(a, b)` where `a` is in `A` and `b` is in - `B` - - The quantity `U` is computed using the reference potential over - all atoms, and `U_A + U_B` computed using the reference potential - over all atoms separated into 2 lambda planes according to which - interacting group they belong - """ + "Compares with jax reference implementation." conf = example_conf[:num_atoms] params = example_nonbonded_params[:num_atoms, :] - lambda_offset_idxs = np.zeros(num_atoms, dtype=np.int32) - - def make_reference_nonbonded(lambda_plane_idxs): - return prepare_reference_nonbonded( - params=params, - exclusion_idxs=np.array([], dtype=np.int32), - scales=np.zeros((0, 2), dtype=np.float64), - lambda_plane_idxs=lambda_plane_idxs, - lambda_offset_idxs=lambda_offset_idxs, - beta=beta, - cutoff=cutoff, - ) - - ref_allpairs = make_reference_nonbonded(np.zeros(num_atoms, dtype=np.int32)) row_atom_idxs = rng.choice(num_atoms, size=num_row_atoms, replace=False).astype(np.int32) - lambda_plane_idxs = np.zeros(num_atoms, dtype=np.int32) - lambda_plane_idxs[row_atom_idxs] = 1 - - ref_allpairs_minus_ixngroups = make_reference_nonbonded(lambda_plane_idxs) + col_atom_idxs = np.setdiff1d(np.arange(num_atoms), row_atom_idxs) - def ref_ixngroups(*args): - return ref_allpairs(*args) - ref_allpairs_minus_ixngroups(*args) + def ref_ixngroups(conf, params, box, _): + vdW, electrostatics, _ = nonbonded.nonbonded_v3_interaction_groups( + conf, params, box, row_atom_idxs, col_atom_idxs, beta, cutoff + ) + return jax.numpy.sum(vdW + electrostatics) test_ixngroups = NonbondedInteractionGroup( row_atom_idxs, - np.zeros(num_atoms, dtype=np.int32), # lambda plane indices - lambda_offset_idxs, + np.zeros(num_atoms, dtype=np.int32), + np.zeros(num_atoms, dtype=np.int32), beta, cutoff, ) @@ -185,8 +171,8 @@ def ref_ixngroups(*args): @pytest.mark.parametrize("cutoff", [1.1]) @pytest.mark.parametrize("precision,rtol,atol", [(np.float64, 1e-8, 1e-8), (np.float32, 1e-4, 5e-4)]) @pytest.mark.parametrize("num_row_atoms", [1, 15]) -@pytest.mark.parametrize("num_atoms", [33, 231]) -def test_nonbonded_interaction_group_correctness( +@pytest.mark.parametrize("num_atoms", [33, 231, 1050]) +def test_nonbonded_interaction_group_consistency_allpairs( num_atoms, num_row_atoms, precision, @@ -198,26 +184,57 @@ def test_nonbonded_interaction_group_correctness( example_nonbonded_params, example_conf, example_box, - rng, + rng: np.random.Generator, ): - "Compares with jax reference implementation." + """Compares with reference nonbonded_v3 potential, which computes + the sum of all pairwise interactions. This uses the identity + + U = U_A + U_B + U_AB + + where + - `U` is the all-pairs potential over all atoms + - `U_A`, `U_B` are all-pairs potentials for interacting groups A + and B, respectively + - `U_AB` is the "interaction group" potential, i.e. the sum of + pairwise interactions `(a, b)` where `a` is in `A` and `b` is in + `B` + + The quantity `U` is computed using the reference potential over + all atoms, and `U_A + U_B` computed using the reference potential + over all atoms separated into 2 lambda planes according to which + interacting group they belong + """ conf = example_conf[:num_atoms] params = example_nonbonded_params[:num_atoms, :] + lambda_offset_idxs = np.zeros(num_atoms, dtype=np.int32) + + def make_reference_nonbonded(lambda_plane_idxs): + return prepare_reference_nonbonded( + params=params, + exclusion_idxs=np.array([], dtype=np.int32), + scales=np.zeros((0, 2), dtype=np.float64), + lambda_plane_idxs=lambda_plane_idxs, + lambda_offset_idxs=lambda_offset_idxs, + beta=beta, + cutoff=cutoff, + ) + + ref_allpairs = make_reference_nonbonded(np.zeros(num_atoms, dtype=np.int32)) row_atom_idxs = rng.choice(num_atoms, size=num_row_atoms, replace=False).astype(np.int32) - col_atom_idxs = np.setdiff1d(np.arange(num_atoms), row_atom_idxs) + lambda_plane_idxs = np.zeros(num_atoms, dtype=np.int32) + lambda_plane_idxs[row_atom_idxs] = 1 - def ref_ixngroups(conf, params, box, _): - vdW, electrostatics, _ = nonbonded.nonbonded_v3_interaction_groups( - conf, params, box, row_atom_idxs, col_atom_idxs, beta, cutoff - ) - return jax.numpy.sum(vdW + electrostatics) + ref_allpairs_minus_ixngroups = make_reference_nonbonded(lambda_plane_idxs) + + def ref_ixngroups(*args): + return ref_allpairs(*args) - ref_allpairs_minus_ixngroups(*args) test_ixngroups = NonbondedInteractionGroup( row_atom_idxs, - np.zeros(num_atoms, dtype=np.int32), - np.zeros(num_atoms, dtype=np.int32), + np.zeros(num_atoms, dtype=np.int32), # lambda plane indices + lambda_offset_idxs, beta, cutoff, ) @@ -233,20 +250,3 @@ def ref_ixngroups(conf, params, box, _): atol=atol, precision=precision, ) - - -def test_nonbonded_interaction_group_invalid_indices(): - def make_potential(row_atom_idxs, num_atoms): - lambda_plane_idxs = [0] * num_atoms - lambda_offset_idxs = [0] * num_atoms - return NonbondedInteractionGroup(row_atom_idxs, lambda_plane_idxs, lambda_offset_idxs, 1.0, 1.0).unbound_impl( - np.float64 - ) - - with pytest.raises(RuntimeError) as e: - make_potential([], 1) - assert "row_atom_idxs must be nonempty" in str(e) - - with pytest.raises(RuntimeError) as e: - make_potential([1, 1], 3) - assert "atom indices must be unique" in str(e) From a6486c198e24755f10159460d4770e4648109f9c Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Thu, 17 Feb 2022 15:08:45 -0800 Subject: [PATCH 39/51] Add consistency test applying constant shift to one group --- tests/test_nonbonded_interaction_group.py | 77 ++++++++++++++++++++++- 1 file changed, 76 insertions(+), 1 deletion(-) diff --git a/tests/test_nonbonded_interaction_group.py b/tests/test_nonbonded_interaction_group.py index 4ab351d32..a1735473d 100644 --- a/tests/test_nonbonded_interaction_group.py +++ b/tests/test_nonbonded_interaction_group.py @@ -172,7 +172,7 @@ def ref_ixngroups(conf, params, box, _): @pytest.mark.parametrize("precision,rtol,atol", [(np.float64, 1e-8, 1e-8), (np.float32, 1e-4, 5e-4)]) @pytest.mark.parametrize("num_row_atoms", [1, 15]) @pytest.mark.parametrize("num_atoms", [33, 231, 1050]) -def test_nonbonded_interaction_group_consistency_allpairs( +def test_nonbonded_interaction_group_consistency_allpairs_lambda_planes( num_atoms, num_row_atoms, precision, @@ -250,3 +250,78 @@ def ref_ixngroups(*args): atol=atol, precision=precision, ) + + +@pytest.mark.parametrize("lamb", [0.0, 0.1]) +@pytest.mark.parametrize("beta", [2.0]) +@pytest.mark.parametrize("cutoff", [1.1]) +@pytest.mark.parametrize("precision,rtol,atol", [(np.float64, 1e-12, 0), (np.float32, 1e-5, 0)]) +@pytest.mark.parametrize("num_row_atoms", [1, 15]) +@pytest.mark.parametrize("num_atoms", [33, 231]) +def test_nonbonded_interaction_group_consistency_allpairs_constant_shift( + num_atoms, + num_row_atoms, + precision, + rtol, + atol, + cutoff, + beta, + lamb, + example_nonbonded_params, + example_conf, + example_box, + rng: np.random.Generator, +): + """Compares with reference nonbonded_v3 potential, which computes + the sum of all pairwise interactions. This uses the identity + + U(x') - U(x) = U_AB(x') - U_AB(x) + + where + - `U` is the all-pairs potential over all atoms + - `U_A`, `U_B` are all-pairs potentials for interacting groups A + and B, respectively + - `U_AB` is the "interaction group" potential, i.e. the sum of + pairwise interactions `(a, b)` where `a` is in `A` and `b` is in + `B` + - the transformation x -> x' does not affect `U_A` or `U_B` (e.g. + a constant translation applied to each atom in one group) + """ + + conf = example_conf[:num_atoms] + params = example_nonbonded_params[:num_atoms, :] + + def ref_allpairs(conf): + return prepare_reference_nonbonded( + params=params, + exclusion_idxs=np.array([], dtype=np.int32), + scales=np.zeros((0, 2), dtype=np.float64), + lambda_plane_idxs=np.zeros(num_atoms, dtype=np.int32), + lambda_offset_idxs=np.zeros(num_atoms, dtype=np.int32), + beta=beta, + cutoff=cutoff, + )(conf, params, example_box, lamb) + + row_atom_idxs = rng.choice(num_atoms, size=num_row_atoms, replace=False).astype(np.int32) + + def test_ixngroups(conf): + _, _, _, u = ( + NonbondedInteractionGroup( + row_atom_idxs, + np.zeros(num_atoms, dtype=np.int32), + np.zeros(num_atoms, dtype=np.int32), + beta, + cutoff, + ) + .unbound_impl(precision) + .execute(conf, params, example_box, lamb) + ) + return u + + conf_prime = np.array(conf) + conf_prime[row_atom_idxs] += rng.normal(0, 0.01, size=(3,)) + + ref_delta = ref_allpairs(conf_prime) - ref_allpairs(conf) + test_delta = test_ixngroups(conf_prime) - test_ixngroups(conf) + + np.testing.assert_allclose(ref_delta, test_delta, rtol=rtol, atol=atol) From 6163d3cb8ccda9ade69c456510dc881299bf0ff5 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Thu, 17 Feb 2022 15:15:46 -0800 Subject: [PATCH 40/51] Rename row/col -> ligand/host in test Row/col is an implementation detail; ligand/host reflects a common usage --- tests/test_nonbonded_interaction_group.py | 46 +++++++++++------------ 1 file changed, 23 insertions(+), 23 deletions(-) diff --git a/tests/test_nonbonded_interaction_group.py b/tests/test_nonbonded_interaction_group.py index a1735473d..571594fb7 100644 --- a/tests/test_nonbonded_interaction_group.py +++ b/tests/test_nonbonded_interaction_group.py @@ -64,10 +64,10 @@ def example_box(example_system): def test_nonbonded_interaction_group_invalid_indices(): - def make_potential(row_atom_idxs, num_atoms): + def make_potential(ligand_idxs, num_atoms): lambda_plane_idxs = [0] * num_atoms lambda_offset_idxs = [0] * num_atoms - return NonbondedInteractionGroup(row_atom_idxs, lambda_plane_idxs, lambda_offset_idxs, 1.0, 1.0).unbound_impl( + return NonbondedInteractionGroup(ligand_idxs, lambda_plane_idxs, lambda_offset_idxs, 1.0, 1.0).unbound_impl( np.float64 ) @@ -82,21 +82,21 @@ def make_potential(row_atom_idxs, num_atoms): def test_nonbonded_interaction_group_zero_interactions(rng: np.random.Generator): num_atoms = 33 - num_row_atoms = 15 + num_atoms_ligand = 15 beta = 2.0 lamb = 0.1 cutoff = 1.1 box = 10.0 * np.eye(3) conf = rng.uniform(0, 1, size=(num_atoms, 3)) - row_atom_idxs = rng.choice(num_atoms, size=num_row_atoms, replace=False).astype(np.int32) + ligand_idxs = rng.choice(num_atoms, size=num_atoms_ligand, replace=False).astype(np.int32) - # shift row atoms in x by twice the cutoff - conf[row_atom_idxs, 0] += 2 * cutoff + # shift ligand atoms in x by twice the cutoff + conf[ligand_idxs, 0] += 2 * cutoff params = rng.uniform(0, 1, size=(num_atoms, 3)) potential = NonbondedInteractionGroup( - row_atom_idxs, + ligand_idxs, np.zeros(num_atoms, dtype=np.int32), np.zeros(num_atoms, dtype=np.int32), beta, @@ -115,11 +115,11 @@ def test_nonbonded_interaction_group_zero_interactions(rng: np.random.Generator) @pytest.mark.parametrize("beta", [2.0]) @pytest.mark.parametrize("cutoff", [1.1]) @pytest.mark.parametrize("precision,rtol,atol", [(np.float64, 1e-8, 1e-8), (np.float32, 1e-4, 5e-4)]) -@pytest.mark.parametrize("num_row_atoms", [1, 15]) +@pytest.mark.parametrize("num_atoms_ligand", [1, 15]) @pytest.mark.parametrize("num_atoms", [33, 231]) def test_nonbonded_interaction_group_correctness( num_atoms, - num_row_atoms, + num_atoms_ligand, precision, rtol, atol, @@ -136,17 +136,17 @@ def test_nonbonded_interaction_group_correctness( conf = example_conf[:num_atoms] params = example_nonbonded_params[:num_atoms, :] - row_atom_idxs = rng.choice(num_atoms, size=num_row_atoms, replace=False).astype(np.int32) - col_atom_idxs = np.setdiff1d(np.arange(num_atoms), row_atom_idxs) + ligand_idxs = rng.choice(num_atoms, size=num_atoms_ligand, replace=False).astype(np.int32) + host_idxs = np.setdiff1d(np.arange(num_atoms), ligand_idxs) def ref_ixngroups(conf, params, box, _): vdW, electrostatics, _ = nonbonded.nonbonded_v3_interaction_groups( - conf, params, box, row_atom_idxs, col_atom_idxs, beta, cutoff + conf, params, box, ligand_idxs, host_idxs, beta, cutoff ) return jax.numpy.sum(vdW + electrostatics) test_ixngroups = NonbondedInteractionGroup( - row_atom_idxs, + ligand_idxs, np.zeros(num_atoms, dtype=np.int32), np.zeros(num_atoms, dtype=np.int32), beta, @@ -170,11 +170,11 @@ def ref_ixngroups(conf, params, box, _): @pytest.mark.parametrize("beta", [2.0]) @pytest.mark.parametrize("cutoff", [1.1]) @pytest.mark.parametrize("precision,rtol,atol", [(np.float64, 1e-8, 1e-8), (np.float32, 1e-4, 5e-4)]) -@pytest.mark.parametrize("num_row_atoms", [1, 15]) +@pytest.mark.parametrize("num_atoms_ligand", [1, 15]) @pytest.mark.parametrize("num_atoms", [33, 231, 1050]) def test_nonbonded_interaction_group_consistency_allpairs_lambda_planes( num_atoms, - num_row_atoms, + num_atoms_ligand, precision, rtol, atol, @@ -222,9 +222,9 @@ def make_reference_nonbonded(lambda_plane_idxs): ref_allpairs = make_reference_nonbonded(np.zeros(num_atoms, dtype=np.int32)) - row_atom_idxs = rng.choice(num_atoms, size=num_row_atoms, replace=False).astype(np.int32) + ligand_idxs = rng.choice(num_atoms, size=num_atoms_ligand, replace=False).astype(np.int32) lambda_plane_idxs = np.zeros(num_atoms, dtype=np.int32) - lambda_plane_idxs[row_atom_idxs] = 1 + lambda_plane_idxs[ligand_idxs] = 1 ref_allpairs_minus_ixngroups = make_reference_nonbonded(lambda_plane_idxs) @@ -232,7 +232,7 @@ def ref_ixngroups(*args): return ref_allpairs(*args) - ref_allpairs_minus_ixngroups(*args) test_ixngroups = NonbondedInteractionGroup( - row_atom_idxs, + ligand_idxs, np.zeros(num_atoms, dtype=np.int32), # lambda plane indices lambda_offset_idxs, beta, @@ -256,11 +256,11 @@ def ref_ixngroups(*args): @pytest.mark.parametrize("beta", [2.0]) @pytest.mark.parametrize("cutoff", [1.1]) @pytest.mark.parametrize("precision,rtol,atol", [(np.float64, 1e-12, 0), (np.float32, 1e-5, 0)]) -@pytest.mark.parametrize("num_row_atoms", [1, 15]) +@pytest.mark.parametrize("num_atoms_ligand", [1, 15]) @pytest.mark.parametrize("num_atoms", [33, 231]) def test_nonbonded_interaction_group_consistency_allpairs_constant_shift( num_atoms, - num_row_atoms, + num_atoms_ligand, precision, rtol, atol, @@ -302,12 +302,12 @@ def ref_allpairs(conf): cutoff=cutoff, )(conf, params, example_box, lamb) - row_atom_idxs = rng.choice(num_atoms, size=num_row_atoms, replace=False).astype(np.int32) + ligand_idxs = rng.choice(num_atoms, size=num_atoms_ligand, replace=False).astype(np.int32) def test_ixngroups(conf): _, _, _, u = ( NonbondedInteractionGroup( - row_atom_idxs, + ligand_idxs, np.zeros(num_atoms, dtype=np.int32), np.zeros(num_atoms, dtype=np.int32), beta, @@ -319,7 +319,7 @@ def test_ixngroups(conf): return u conf_prime = np.array(conf) - conf_prime[row_atom_idxs] += rng.normal(0, 0.01, size=(3,)) + conf_prime[ligand_idxs] += rng.normal(0, 0.01, size=(3,)) ref_delta = ref_allpairs(conf_prime) - ref_allpairs(conf) test_delta = test_ixngroups(conf_prime) - test_ixngroups(conf) From 824342dfdad206762f29cf51af4262513afeaa50 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Thu, 17 Feb 2022 15:26:10 -0800 Subject: [PATCH 41/51] Clean up docstrings --- tests/test_nonbonded_interaction_group.py | 36 +++++++++++------------ 1 file changed, 17 insertions(+), 19 deletions(-) diff --git a/tests/test_nonbonded_interaction_group.py b/tests/test_nonbonded_interaction_group.py index 571594fb7..dd9efe70a 100644 --- a/tests/test_nonbonded_interaction_group.py +++ b/tests/test_nonbonded_interaction_group.py @@ -192,17 +192,16 @@ def test_nonbonded_interaction_group_consistency_allpairs_lambda_planes( U = U_A + U_B + U_AB where - - `U` is the all-pairs potential over all atoms - - `U_A`, `U_B` are all-pairs potentials for interacting groups A - and B, respectively - - `U_AB` is the "interaction group" potential, i.e. the sum of - pairwise interactions `(a, b)` where `a` is in `A` and `b` is in - `B` - - The quantity `U` is computed using the reference potential over - all atoms, and `U_A + U_B` computed using the reference potential - over all atoms separated into 2 lambda planes according to which - interacting group they belong + - U is the all-pairs potential over all atoms + - U_A, U_B are all-pairs potentials for interacting groups A and + B, respectively + - U_AB is the "interaction group" potential, i.e. the sum of + pairwise interactions (a, b) where "a" is in A and "b" is in B + + U is computed using the reference potential over all atoms, and + U_A + U_B computed using the reference potential over all atoms + separated into 2 lambda planes according to which interacting + group they belong """ conf = example_conf[:num_atoms] @@ -278,14 +277,13 @@ def test_nonbonded_interaction_group_consistency_allpairs_constant_shift( U(x') - U(x) = U_AB(x') - U_AB(x) where - - `U` is the all-pairs potential over all atoms - - `U_A`, `U_B` are all-pairs potentials for interacting groups A - and B, respectively - - `U_AB` is the "interaction group" potential, i.e. the sum of - pairwise interactions `(a, b)` where `a` is in `A` and `b` is in - `B` - - the transformation x -> x' does not affect `U_A` or `U_B` (e.g. - a constant translation applied to each atom in one group) + - U is the all-pairs potential over all atoms + - U_A, U_B are all-pairs potentials for interacting groups A and + B, respectively + - U_AB is the "interaction group" potential, i.e. the sum of + pairwise interactions (a, b) where "a" is in A and "b" is in B + - the transformation x -> x' does not affect U_A or U_B (e.g. a + constant translation applied to each atom in one group) """ conf = example_conf[:num_atoms] From a986a342dd9971b79cc82775acef99d056278437 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Thu, 17 Feb 2022 15:32:21 -0800 Subject: [PATCH 42/51] Adjust test tolerances --- tests/test_nonbonded_interaction_group.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_nonbonded_interaction_group.py b/tests/test_nonbonded_interaction_group.py index dd9efe70a..18d07e4b0 100644 --- a/tests/test_nonbonded_interaction_group.py +++ b/tests/test_nonbonded_interaction_group.py @@ -254,7 +254,7 @@ def ref_ixngroups(*args): @pytest.mark.parametrize("lamb", [0.0, 0.1]) @pytest.mark.parametrize("beta", [2.0]) @pytest.mark.parametrize("cutoff", [1.1]) -@pytest.mark.parametrize("precision,rtol,atol", [(np.float64, 1e-12, 0), (np.float32, 1e-5, 0)]) +@pytest.mark.parametrize("precision,rtol,atol", [(np.float64, 1e-8, 1e-8), (np.float32, 1e-4, 5e-4)]) @pytest.mark.parametrize("num_atoms_ligand", [1, 15]) @pytest.mark.parametrize("num_atoms", [33, 231]) def test_nonbonded_interaction_group_consistency_allpairs_constant_shift( From 91d1722b40dfae61d46239686ff021c44aa5e7bd Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Thu, 17 Feb 2022 17:46:19 -0800 Subject: [PATCH 43/51] Replace itertools product with meshgrid Roughly 5000x speedup for arrays of length 100 --- timemachine/potentials/nonbonded.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/timemachine/potentials/nonbonded.py b/timemachine/potentials/nonbonded.py index 557da0a0e..4082c83ac 100644 --- a/timemachine/potentials/nonbonded.py +++ b/timemachine/potentials/nonbonded.py @@ -259,7 +259,7 @@ def nonbonded_v3_interaction_groups(conf, params, box, inds_l, inds_r, beta: flo See nonbonded_v3 docstring for more details """ - pairs = np.array(list(itertools.product(inds_l, inds_r))) + pairs = np.stack(np.meshgrid(inds_l, inds_r)).reshape(2, -1).T vdW, electrostatics = nonbonded_v3_on_specific_pairs(conf, params, box, pairs[:, 0], pairs[:, 1], beta, cutoff) return vdW, electrostatics, pairs From efb100d8b6aa7ce3ed146952b5aa0040237ee5b2 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Thu, 17 Feb 2022 17:57:42 -0800 Subject: [PATCH 44/51] Remove unused import --- timemachine/potentials/nonbonded.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/timemachine/potentials/nonbonded.py b/timemachine/potentials/nonbonded.py index 4082c83ac..006351a52 100644 --- a/timemachine/potentials/nonbonded.py +++ b/timemachine/potentials/nonbonded.py @@ -1,5 +1,3 @@ -import itertools - import jax.numpy as np from jax.ops import index, index_update from jax.scipy.special import erfc From 90d0cd588762edd30b402566c84a34c2e0205a76 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Fri, 18 Feb 2022 09:12:28 -0800 Subject: [PATCH 45/51] Remove obsolete comment --- timemachine/cpp/src/nonbonded_interaction_group.cu | 1 - 1 file changed, 1 deletion(-) diff --git a/timemachine/cpp/src/nonbonded_interaction_group.cu b/timemachine/cpp/src/nonbonded_interaction_group.cu index e72369834..b33594870 100644 --- a/timemachine/cpp/src/nonbonded_interaction_group.cu +++ b/timemachine/cpp/src/nonbonded_interaction_group.cu @@ -395,7 +395,6 @@ void NonbondedInteractionGroup::execute_device( // update new w coordinates // (tbd): cache lambda value for equilibrium calculations - // TODO: skip computing w coords for non-interacting atoms? CUresult result = compute_w_coords_instance_.configure(B, tpb, 0, stream) .launch(N, lambda, cutoff_, d_lambda_plane_idxs_, d_lambda_offset_idxs_, d_w_, d_dw_dl_); if (result != 0) { From 695f2c179d62588ecae09037a820c500c6234156 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Fri, 18 Feb 2022 09:59:49 -0800 Subject: [PATCH 46/51] Remove no-op checks --- timemachine/cpp/src/nonbonded_all_pairs.cu | 4 ---- timemachine/cpp/src/nonbonded_interaction_group.cu | 4 ---- 2 files changed, 8 deletions(-) diff --git a/timemachine/cpp/src/nonbonded_all_pairs.cu b/timemachine/cpp/src/nonbonded_all_pairs.cu index 66235bd30..b618e06dc 100644 --- a/timemachine/cpp/src/nonbonded_all_pairs.cu +++ b/timemachine/cpp/src/nonbonded_all_pairs.cu @@ -63,10 +63,6 @@ NonbondedAllPairs::NonbondedAllPairs( compute_add_du_dp_interpolated_( kernel_cache_.program(kernel_src.c_str()).kernel("k_add_du_dp_interpolated").instantiate()) { - if (lambda_offset_idxs.size() != N_) { - throw std::runtime_error("lambda offset idxs need to have size N"); - } - if (lambda_offset_idxs.size() != lambda_plane_idxs.size()) { throw std::runtime_error("lambda offset idxs and plane idxs need to be equivalent"); } diff --git a/timemachine/cpp/src/nonbonded_interaction_group.cu b/timemachine/cpp/src/nonbonded_interaction_group.cu index b33594870..022f0ff3d 100644 --- a/timemachine/cpp/src/nonbonded_interaction_group.cu +++ b/timemachine/cpp/src/nonbonded_interaction_group.cu @@ -70,10 +70,6 @@ NonbondedInteractionGroup::NonbondedInteractionGroup( throw std::runtime_error("row_atom_idxs must be nonempty"); } - if (lambda_offset_idxs.size() != N_) { - throw std::runtime_error("lambda offset idxs need to have size N"); - } - if (lambda_offset_idxs.size() != lambda_plane_idxs.size()) { throw std::runtime_error("lambda offset idxs and plane idxs need to be equivalent"); } From 6c5598bb382df0f5b071d7a8028c387bb1b6ad9a Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Fri, 18 Feb 2022 10:36:58 -0800 Subject: [PATCH 47/51] Use nonzero offset indices in consistency checks --- tests/test_nonbonded_interaction_group.py | 21 +++++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/tests/test_nonbonded_interaction_group.py b/tests/test_nonbonded_interaction_group.py index 18d07e4b0..e61f71b80 100644 --- a/tests/test_nonbonded_interaction_group.py +++ b/tests/test_nonbonded_interaction_group.py @@ -206,7 +206,9 @@ def test_nonbonded_interaction_group_consistency_allpairs_lambda_planes( conf = example_conf[:num_atoms] params = example_nonbonded_params[:num_atoms, :] - lambda_offset_idxs = np.zeros(num_atoms, dtype=np.int32) + + max_abs_offset_idx = 2 # i.e., lambda_offset_idxs in {-2, -1, 0, 1, 2} + lambda_offset_idxs = rng.integers(-max_abs_offset_idx, max_abs_offset_idx + 1, size=num_atoms, dtype=np.int32) def make_reference_nonbonded(lambda_plane_idxs): return prepare_reference_nonbonded( @@ -222,8 +224,12 @@ def make_reference_nonbonded(lambda_plane_idxs): ref_allpairs = make_reference_nonbonded(np.zeros(num_atoms, dtype=np.int32)) ligand_idxs = rng.choice(num_atoms, size=num_atoms_ligand, replace=False).astype(np.int32) + + # for reference U_A + U_B computation, ensure minimum distance + # between a host and ligand atom is at least one cutoff distance + # when lambda = 1 lambda_plane_idxs = np.zeros(num_atoms, dtype=np.int32) - lambda_plane_idxs[ligand_idxs] = 1 + lambda_plane_idxs[ligand_idxs] = 2 * max_abs_offset_idx + 1 ref_allpairs_minus_ixngroups = make_reference_nonbonded(lambda_plane_idxs) @@ -289,13 +295,16 @@ def test_nonbonded_interaction_group_consistency_allpairs_constant_shift( conf = example_conf[:num_atoms] params = example_nonbonded_params[:num_atoms, :] + lambda_plane_idxs = rng.integers(-2, 3, size=num_atoms, dtype=np.int32) + lambda_offset_idxs = rng.integers(-2, 3, size=num_atoms, dtype=np.int32) + def ref_allpairs(conf): return prepare_reference_nonbonded( params=params, exclusion_idxs=np.array([], dtype=np.int32), scales=np.zeros((0, 2), dtype=np.float64), - lambda_plane_idxs=np.zeros(num_atoms, dtype=np.int32), - lambda_offset_idxs=np.zeros(num_atoms, dtype=np.int32), + lambda_plane_idxs=lambda_plane_idxs, + lambda_offset_idxs=lambda_offset_idxs, beta=beta, cutoff=cutoff, )(conf, params, example_box, lamb) @@ -306,8 +315,8 @@ def test_ixngroups(conf): _, _, _, u = ( NonbondedInteractionGroup( ligand_idxs, - np.zeros(num_atoms, dtype=np.int32), - np.zeros(num_atoms, dtype=np.int32), + lambda_plane_idxs, + lambda_offset_idxs, beta, cutoff, ) From dbf4fbe91da42b13aee01e59e7f69367a3807709 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Fri, 18 Feb 2022 11:13:07 -0800 Subject: [PATCH 48/51] Use nonzero planes, offsets in correctness test --- tests/test_nonbonded_interaction_group.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/tests/test_nonbonded_interaction_group.py b/tests/test_nonbonded_interaction_group.py index e61f71b80..0a6d5d6be 100644 --- a/tests/test_nonbonded_interaction_group.py +++ b/tests/test_nonbonded_interaction_group.py @@ -11,7 +11,7 @@ from timemachine.ff.handlers import openmm_deserializer from timemachine.lib import potentials from timemachine.lib.potentials import NonbondedInteractionGroup -from timemachine.potentials import nonbonded +from timemachine.potentials import jax_utils, nonbonded @pytest.fixture(autouse=True) @@ -136,19 +136,28 @@ def test_nonbonded_interaction_group_correctness( conf = example_conf[:num_atoms] params = example_nonbonded_params[:num_atoms, :] + lambda_plane_idxs = rng.integers(-2, 3, size=num_atoms, dtype=np.int32) + lambda_offset_idxs = rng.integers(-2, 3, size=num_atoms, dtype=np.int32) + ligand_idxs = rng.choice(num_atoms, size=num_atoms_ligand, replace=False).astype(np.int32) host_idxs = np.setdiff1d(np.arange(num_atoms), ligand_idxs) - def ref_ixngroups(conf, params, box, _): + def ref_ixngroups(conf, params, box, lamb): + + # compute 4d coordinates + w = jax_utils.compute_lifting_parameter(lamb, lambda_plane_idxs, lambda_offset_idxs, cutoff) + conf_4d = jax_utils.augment_dim(conf, w) + box_4d = (1000 * jax.numpy.eye(4)).at[:3, :3].set(box) + vdW, electrostatics, _ = nonbonded.nonbonded_v3_interaction_groups( - conf, params, box, ligand_idxs, host_idxs, beta, cutoff + conf_4d, params, box_4d, ligand_idxs, host_idxs, beta, cutoff ) return jax.numpy.sum(vdW + electrostatics) test_ixngroups = NonbondedInteractionGroup( ligand_idxs, - np.zeros(num_atoms, dtype=np.int32), - np.zeros(num_atoms, dtype=np.int32), + lambda_plane_idxs, + lambda_offset_idxs, beta, cutoff, ) From 6d199ca4e21366d4c84d0ea81f87f668368efc94 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Fri, 18 Feb 2022 11:39:52 -0800 Subject: [PATCH 49/51] Expose interpolated potential to Python API, add test --- tests/test_nonbonded_interaction_group.py | 69 ++++++++++++++++++++++- tests/test_parameter_interpolation.py | 12 +--- timemachine/lib/potentials.py | 13 +++++ timemachine/potentials/nonbonded.py | 18 ++++++ 4 files changed, 101 insertions(+), 11 deletions(-) diff --git a/tests/test_nonbonded_interaction_group.py b/tests/test_nonbonded_interaction_group.py index 0a6d5d6be..18574c1b6 100644 --- a/tests/test_nonbonded_interaction_group.py +++ b/tests/test_nonbonded_interaction_group.py @@ -10,7 +10,7 @@ from timemachine.fe.utils import to_md_units from timemachine.ff.handlers import openmm_deserializer from timemachine.lib import potentials -from timemachine.lib.potentials import NonbondedInteractionGroup +from timemachine.lib.potentials import NonbondedInteractionGroup, NonbondedInteractionGroupInterpolated from timemachine.potentials import jax_utils, nonbonded @@ -175,6 +175,73 @@ def ref_ixngroups(conf, params, box, lamb): ) +@pytest.mark.parametrize("lamb", [0.0, 0.1, 0.9, 1.0]) +@pytest.mark.parametrize("beta", [2.0]) +@pytest.mark.parametrize("cutoff", [1.1]) +@pytest.mark.parametrize("precision,rtol,atol", [(np.float64, 1e-8, 1e-8), (np.float32, 1e-4, 5e-4)]) +@pytest.mark.parametrize("num_atoms_ligand", [1, 15]) +@pytest.mark.parametrize("num_atoms", [33]) +def test_nonbonded_interaction_group_interpolated_correctness( + num_atoms, + num_atoms_ligand, + precision, + rtol, + atol, + cutoff, + beta, + lamb, + example_nonbonded_params, + example_conf, + example_box, + rng, +): + "Compares with jax reference implementation, with parameter interpolation." + + conf = example_conf[:num_atoms] + params_initial = example_nonbonded_params[:num_atoms, :] + params_final = params_initial + rng.normal(0, 0.01, size=params_initial.shape) + params = np.concatenate((params_initial, params_final)) + + lambda_plane_idxs = rng.integers(-2, 3, size=num_atoms, dtype=np.int32) + lambda_offset_idxs = rng.integers(-2, 3, size=num_atoms, dtype=np.int32) + + ligand_idxs = rng.choice(num_atoms, size=num_atoms_ligand, replace=False).astype(np.int32) + host_idxs = np.setdiff1d(np.arange(num_atoms), ligand_idxs) + + @nonbonded.interpolated + def ref_ixngroups(conf, params, box, lamb): + + # compute 4d coordinates + w = jax_utils.compute_lifting_parameter(lamb, lambda_plane_idxs, lambda_offset_idxs, cutoff) + conf_4d = jax_utils.augment_dim(conf, w) + box_4d = (1000 * jax.numpy.eye(4)).at[:3, :3].set(box) + + vdW, electrostatics, _ = nonbonded.nonbonded_v3_interaction_groups( + conf_4d, params, box_4d, ligand_idxs, host_idxs, beta, cutoff + ) + return jax.numpy.sum(vdW + electrostatics) + + test_ixngroups = NonbondedInteractionGroupInterpolated( + ligand_idxs, + lambda_plane_idxs, + lambda_offset_idxs, + beta, + cutoff, + ) + + GradientTest().compare_forces( + conf, + params, + example_box, + lamb=lamb, + ref_potential=ref_ixngroups, + test_potential=test_ixngroups, + rtol=rtol, + atol=atol, + precision=precision, + ) + + @pytest.mark.parametrize("lamb", [0.0, 0.1]) @pytest.mark.parametrize("beta", [2.0]) @pytest.mark.parametrize("cutoff", [1.1]) diff --git a/tests/test_parameter_interpolation.py b/tests/test_parameter_interpolation.py index 099a5fd5d..9299b05da 100644 --- a/tests/test_parameter_interpolation.py +++ b/tests/test_parameter_interpolation.py @@ -9,15 +9,7 @@ from common import GradientTest, prepare_water_system from timemachine.lib import potentials - - -def interpolated_potential(conf, params, box, lamb, u_fn): - assert params.size % 2 == 0 - - CP = params.shape[0] // 2 - new_params = (1 - lamb) * params[:CP] + lamb * params[CP:] - - return u_fn(conf, new_params, box, lamb) +from timemachine.potentials import nonbonded class TestInterpolatedPotential(GradientTest): @@ -57,7 +49,7 @@ def test_nonbonded(self): print("lambda", lamb, "cutoff", cutoff, "precision", precision, "xshape", coords.shape) - ref_interpolated_potential = functools.partial(interpolated_potential, u_fn=ref_potential) + ref_interpolated_potential = nonbonded.interpolated(ref_potential) test_interpolated_potential = potentials.NonbondedInterpolated(*test_potential.args) diff --git a/timemachine/lib/potentials.py b/timemachine/lib/potentials.py index 23720f3fb..26346a5d8 100644 --- a/timemachine/lib/potentials.py +++ b/timemachine/lib/potentials.py @@ -282,3 +282,16 @@ def unbound_impl(self, precision): class NonbondedInteractionGroup(CustomOpWrapper): pass + + +class NonbondedInteractionGroupInterpolated(NonbondedInteractionGroup): + def unbound_impl(self, precision): + cls_name_base = "NonbondedInteractionGroup" + if precision == np.float64: + cls_name_base += "_f64_interpolated" + else: + cls_name_base += "_f32_interpolated" + + custom_ctor = getattr(custom_ops, cls_name_base) + + return custom_ctor(*self.args) diff --git a/timemachine/potentials/nonbonded.py b/timemachine/potentials/nonbonded.py index 006351a52..14de0136d 100644 --- a/timemachine/potentials/nonbonded.py +++ b/timemachine/potentials/nonbonded.py @@ -1,3 +1,5 @@ +import functools + import jax.numpy as np from jax.ops import index, index_update from jax.scipy.special import erfc @@ -262,6 +264,22 @@ def nonbonded_v3_interaction_groups(conf, params, box, inds_l, inds_r, beta: flo return vdW, electrostatics, pairs +def interpolated(u_fn): + @functools.wraps(u_fn) + def wrapper(conf, params, box, lamb): + + # params is expected to be the concatenation of initial + # (lambda = 0) and final (lamda = 1) parameters, each of + # length num_atoms + assert params.size % 2 == 0 + num_atoms = params.shape[0] // 2 + + new_params = (1 - lamb) * params[:num_atoms] + lamb * params[num_atoms:] + return u_fn(conf, new_params, box, lamb) + + return wrapper + + def validate_coulomb_cutoff(cutoff=1.0, beta=2.0, threshold=1e-2): """check whether f(r) = erfc(beta * r) <= threshold at r = cutoff following https://github.com/proteneer/timemachine/pull/424#discussion_r629678467""" From 22f9dc9b358490867e41469dfd5a88cf3a1247cd Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Fri, 18 Feb 2022 13:16:04 -0800 Subject: [PATCH 50/51] Remove unused import --- tests/test_parameter_interpolation.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_parameter_interpolation.py b/tests/test_parameter_interpolation.py index 9299b05da..4605b2f01 100644 --- a/tests/test_parameter_interpolation.py +++ b/tests/test_parameter_interpolation.py @@ -2,7 +2,6 @@ config.update("jax_enable_x64", True) import copy -import functools import jax.numpy as jnp import numpy as np From cdd37ff1794adda24d98bddda860c241b8a65328 Mon Sep 17 00:00:00 2001 From: Matt Wittmann Date: Fri, 18 Feb 2022 12:39:06 -0800 Subject: [PATCH 51/51] Pass tuples to silence type warnings --- tests/test_nonbonded_interaction_group.py | 24 +++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/tests/test_nonbonded_interaction_group.py b/tests/test_nonbonded_interaction_group.py index 18574c1b6..a786c2c20 100644 --- a/tests/test_nonbonded_interaction_group.py +++ b/tests/test_nonbonded_interaction_group.py @@ -88,7 +88,7 @@ def test_nonbonded_interaction_group_zero_interactions(rng: np.random.Generator) cutoff = 1.1 box = 10.0 * np.eye(3) conf = rng.uniform(0, 1, size=(num_atoms, 3)) - ligand_idxs = rng.choice(num_atoms, size=num_atoms_ligand, replace=False).astype(np.int32) + ligand_idxs = rng.choice(num_atoms, size=(num_atoms_ligand,), replace=False).astype(np.int32) # shift ligand atoms in x by twice the cutoff conf[ligand_idxs, 0] += 2 * cutoff @@ -136,10 +136,10 @@ def test_nonbonded_interaction_group_correctness( conf = example_conf[:num_atoms] params = example_nonbonded_params[:num_atoms, :] - lambda_plane_idxs = rng.integers(-2, 3, size=num_atoms, dtype=np.int32) - lambda_offset_idxs = rng.integers(-2, 3, size=num_atoms, dtype=np.int32) + lambda_plane_idxs = rng.integers(-2, 3, size=(num_atoms,), dtype=np.int32) + lambda_offset_idxs = rng.integers(-2, 3, size=(num_atoms,), dtype=np.int32) - ligand_idxs = rng.choice(num_atoms, size=num_atoms_ligand, replace=False).astype(np.int32) + ligand_idxs = rng.choice(num_atoms, size=(num_atoms_ligand,), replace=False).astype(np.int32) host_idxs = np.setdiff1d(np.arange(num_atoms), ligand_idxs) def ref_ixngroups(conf, params, box, lamb): @@ -202,10 +202,10 @@ def test_nonbonded_interaction_group_interpolated_correctness( params_final = params_initial + rng.normal(0, 0.01, size=params_initial.shape) params = np.concatenate((params_initial, params_final)) - lambda_plane_idxs = rng.integers(-2, 3, size=num_atoms, dtype=np.int32) - lambda_offset_idxs = rng.integers(-2, 3, size=num_atoms, dtype=np.int32) + lambda_plane_idxs = rng.integers(-2, 3, size=(num_atoms,), dtype=np.int32) + lambda_offset_idxs = rng.integers(-2, 3, size=(num_atoms,), dtype=np.int32) - ligand_idxs = rng.choice(num_atoms, size=num_atoms_ligand, replace=False).astype(np.int32) + ligand_idxs = rng.choice(num_atoms, size=(num_atoms_ligand,), replace=False).astype(np.int32) host_idxs = np.setdiff1d(np.arange(num_atoms), ligand_idxs) @nonbonded.interpolated @@ -284,7 +284,7 @@ def test_nonbonded_interaction_group_consistency_allpairs_lambda_planes( params = example_nonbonded_params[:num_atoms, :] max_abs_offset_idx = 2 # i.e., lambda_offset_idxs in {-2, -1, 0, 1, 2} - lambda_offset_idxs = rng.integers(-max_abs_offset_idx, max_abs_offset_idx + 1, size=num_atoms, dtype=np.int32) + lambda_offset_idxs = rng.integers(-max_abs_offset_idx, max_abs_offset_idx + 1, size=(num_atoms,), dtype=np.int32) def make_reference_nonbonded(lambda_plane_idxs): return prepare_reference_nonbonded( @@ -299,7 +299,7 @@ def make_reference_nonbonded(lambda_plane_idxs): ref_allpairs = make_reference_nonbonded(np.zeros(num_atoms, dtype=np.int32)) - ligand_idxs = rng.choice(num_atoms, size=num_atoms_ligand, replace=False).astype(np.int32) + ligand_idxs = rng.choice(num_atoms, size=(num_atoms_ligand,), replace=False).astype(np.int32) # for reference U_A + U_B computation, ensure minimum distance # between a host and ligand atom is at least one cutoff distance @@ -371,8 +371,8 @@ def test_nonbonded_interaction_group_consistency_allpairs_constant_shift( conf = example_conf[:num_atoms] params = example_nonbonded_params[:num_atoms, :] - lambda_plane_idxs = rng.integers(-2, 3, size=num_atoms, dtype=np.int32) - lambda_offset_idxs = rng.integers(-2, 3, size=num_atoms, dtype=np.int32) + lambda_plane_idxs = rng.integers(-2, 3, size=(num_atoms,), dtype=np.int32) + lambda_offset_idxs = rng.integers(-2, 3, size=(num_atoms,), dtype=np.int32) def ref_allpairs(conf): return prepare_reference_nonbonded( @@ -385,7 +385,7 @@ def ref_allpairs(conf): cutoff=cutoff, )(conf, params, example_box, lamb) - ligand_idxs = rng.choice(num_atoms, size=num_atoms_ligand, replace=False).astype(np.int32) + ligand_idxs = rng.choice(num_atoms, size=(num_atoms_ligand,), replace=False).astype(np.int32) def test_ixngroups(conf): _, _, _, u = (