From a631d5a14a5a31971d9142c073087b8d1996361d Mon Sep 17 00:00:00 2001 From: Mohammadreza Date: Mon, 26 Jan 2026 16:28:45 -0600 Subject: [PATCH 1/5] Add header for Chebyshev anisotropic pair potential --- src/ChebyshevTensorAnisotropicPairPotential.h | 130 ++++++++++++++++++ 1 file changed, 130 insertions(+) create mode 100644 src/ChebyshevTensorAnisotropicPairPotential.h diff --git a/src/ChebyshevTensorAnisotropicPairPotential.h b/src/ChebyshevTensorAnisotropicPairPotential.h new file mode 100644 index 0000000..626771d --- /dev/null +++ b/src/ChebyshevTensorAnisotropicPairPotential.h @@ -0,0 +1,130 @@ +// Copyright (c) 2018-2020, Michael P. Howard +// Copyright (c) 2021-2025, Auburn University +// Part of azplugins, released under the BSD 3-Clause License. + +/*! + * \file ChebyshevTensorAnisotropicPairPotential.h + * \brief Declaration of ChebyshevTensorAnisotropicPairPotential + */ + +#ifndef AZPLUGINS_CHEBYSHEV_TENSOR_ANISO_PAIR_POTENTIAL_H_ +#define AZPLUGINS_CHEBYSHEV_TENSOR_ANISO_PAIR_POTENTIAL_H_ + +#ifdef NVCC +#error This header cannot be compiled by nvcc +#endif + +#include +#include +#include + +#include "hoomd/ForceCompute.h" +#include "hoomd/GPUArray.h" +#include "hoomd/HOOMDMath.h" +#include "hoomd/Index1D.h" + +#include "hoomd/md/NeighborList.h" + +namespace hoomd + { +namespace azplugins + { + +class R0Interpolator; + +class PYBIND11_EXPORT ChebyshevTensorAnisotropicPairPotential : public ForceCompute + { + public: + //! Constructor + ChebyshevTensorAnisotropicPairPotential(std::shared_ptr sysdef, + std::shared_ptr nlist); + + //! Destructor + virtual ~ChebyshevTensorAnisotropicPairPotential(); + + // Getters + std::shared_ptr getNeighborList() const + { + return m_nlist; + } + + /// 6x2 domain: stored as 6 entries of Scalar2 = (min,max) + const GPUArray& getApproximationDomain() const + { + return m_domain; + } + + /// Term degrees (Nterms x 6) + const GPUArray& getChebyshevTermList() const + { + return m_terms; + } + + /// Coefficients (length = Nterms) + const GPUArray& getCoefficients() const + { + return m_coeffs; + } + + /// Max degree per dimension length (length = 6) + const GPUArray& getMaxDegreePerDim() const + { + return m_max_degree; + } + + /// Number of terms (int) + unsigned int getNTerms() const + { + return m_Nterms; + } + + /// R0 interpolator object + std::shared_ptr getR0Interpolator() const + { + return m_r0_interp; + } + + /// Allocate storage for term list and coefficients. + void resizeTerms(unsigned int Nterms); + + void setR0Interpolator(std::shared_ptr interp) + { + m_r0_interp = interp; + } + + protected: + void computeForces(uint64_t timestep) override; + + private: + // 1) neighbor list object + std::shared_ptr m_nlist; + + // 2) approximation domain (6x2): 6 rows, each is (min,max) + GPUArray m_domain; + + // 3) r0 interpolation object + std::shared_ptr m_r0_interp; + + // 4) Chebyshev term list (Nterms x 6) + GPUArray m_terms; + + // 5) coeffs (Nterms) + GPUArray m_coeffs; + + // 6) max degree per dimension (6) + GPUArray m_max_degree; + + // 7) number of terms + unsigned int m_Nterms; + }; + +namespace detail + { +///! exports to Python +void export_ChebyshevTensorAnisotropicPairPotential(pybind11::module& m); + } // end namespace detail + + } // end namespace azplugins + } // end namespace hoomd + +#endif // AZPLUGINS_CHEBYSHEV_TENSOR_ANISO_PAIR_POTENTIAL_H_ From c0323f1a22c011610a9c8e8d347e6bc96a929e36 Mon Sep 17 00:00:00 2001 From: Mohammadreza Date: Thu, 5 Feb 2026 16:58:58 -0600 Subject: [PATCH 2/5] Update header interface --- src/ChebyshevTensorAnisotropicPairPotential.h | 55 ++++++++----------- 1 file changed, 24 insertions(+), 31 deletions(-) diff --git a/src/ChebyshevTensorAnisotropicPairPotential.h b/src/ChebyshevTensorAnisotropicPairPotential.h index 626771d..7d2374f 100644 --- a/src/ChebyshevTensorAnisotropicPairPotential.h +++ b/src/ChebyshevTensorAnisotropicPairPotential.h @@ -7,13 +7,14 @@ * \brief Declaration of ChebyshevTensorAnisotropicPairPotential */ -#ifndef AZPLUGINS_CHEBYSHEV_TENSOR_ANISO_PAIR_POTENTIAL_H_ -#define AZPLUGINS_CHEBYSHEV_TENSOR_ANISO_PAIR_POTENTIAL_H_ +#ifndef AZPLUGINS_CHEBYSHEV_TENSOR_ANISOTROPIC_PAIR_POTENTIAL_H_ +#define AZPLUGINS_CHEBYSHEV_TENSOR_ANISOTROPIC_PAIR_POTENTIAL_H_ #ifdef NVCC #error This header cannot be compiled by nvcc #endif +#include #include #include #include @@ -30,14 +31,18 @@ namespace hoomd namespace azplugins { -class R0Interpolator; +class LinearInterpolator5D; class PYBIND11_EXPORT ChebyshevTensorAnisotropicPairPotential : public ForceCompute { public: //! Constructor ChebyshevTensorAnisotropicPairPotential(std::shared_ptr sysdef, - std::shared_ptr nlist); + std::shared_ptr nlist, + const std::array& domain, + const std::vector& r0_data, + const std::vector& terms, + const std::vector& coeffs); //! Destructor virtual ~ChebyshevTensorAnisotropicPairPotential(); @@ -54,6 +59,12 @@ class PYBIND11_EXPORT ChebyshevTensorAnisotropicPairPotential : public ForceComp return m_domain; } + /// r0 data (theta, phi, alpha, beta, gamma) (N x 6) + const GPUArray& getR0Data() const + { + return m_r0_data; + } + /// Term degrees (Nterms x 6) const GPUArray& getChebyshevTermList() const { @@ -66,32 +77,14 @@ class PYBIND11_EXPORT ChebyshevTensorAnisotropicPairPotential : public ForceComp return m_coeffs; } - /// Max degree per dimension length (length = 6) - const GPUArray& getMaxDegreePerDim() const - { - return m_max_degree; - } - - /// Number of terms (int) unsigned int getNTerms() const { return m_Nterms; } - /// R0 interpolator object - std::shared_ptr getR0Interpolator() const - { - return m_r0_interp; - } - /// Allocate storage for term list and coefficients. void resizeTerms(unsigned int Nterms); - void setR0Interpolator(std::shared_ptr interp) - { - m_r0_interp = interp; - } - protected: void computeForces(uint64_t timestep) override; @@ -102,20 +95,20 @@ class PYBIND11_EXPORT ChebyshevTensorAnisotropicPairPotential : public ForceComp // 2) approximation domain (6x2): 6 rows, each is (min,max) GPUArray m_domain; - // 3) r0 interpolation object - std::shared_ptr m_r0_interp; + // 3) intenal r0 linear interpolator + std::unique_ptr m_r0_interp; + + // 4) r0_data + GPUArray m_r0_data; - // 4) Chebyshev term list (Nterms x 6) + // 5) Chebyshev term list (Nterms x 6) GPUArray m_terms; - // 5) coeffs (Nterms) + // 6) coeffs (Nterms) GPUArray m_coeffs; - // 6) max degree per dimension (6) - GPUArray m_max_degree; - // 7) number of terms - unsigned int m_Nterms; + unsigned int m_Nterms = 0; }; namespace detail @@ -127,4 +120,4 @@ void export_ChebyshevTensorAnisotropicPairPotential(pybind11::module& m); } // end namespace azplugins } // end namespace hoomd -#endif // AZPLUGINS_CHEBYSHEV_TENSOR_ANISO_PAIR_POTENTIAL_H_ +#endif // AZPLUGINS_CHEBYSHEV_TENSOR_ANISOTROPIC_PAIR_POTENTIAL_H_ From 189f25f2d8fdeff994720bc45605aba81a201fb7 Mon Sep 17 00:00:00 2001 From: Mohammadreza Date: Mon, 16 Feb 2026 09:17:05 -0600 Subject: [PATCH 3/5] Update header interface --- src/ChebyshevAnisotropicPairPotential.h | 100 ++++++++++++++ src/ChebyshevTensorAnisotropicPairPotential.h | 123 ------------------ 2 files changed, 100 insertions(+), 123 deletions(-) create mode 100644 src/ChebyshevAnisotropicPairPotential.h delete mode 100644 src/ChebyshevTensorAnisotropicPairPotential.h diff --git a/src/ChebyshevAnisotropicPairPotential.h b/src/ChebyshevAnisotropicPairPotential.h new file mode 100644 index 0000000..dd14b25 --- /dev/null +++ b/src/ChebyshevAnisotropicPairPotential.h @@ -0,0 +1,100 @@ +// Copyright (c) 2018-2020, Michael P. Howard +// Copyright (c) 2021-2025, Auburn University +// Part of azplugins, released under the BSD 3-Clause License. + +/*! + * \file ChebyshevAnisotropicPairPotential.h + * \brief Declaration of ChebyshevAnisotropicPairPotential + */ + +#ifndef AZPLUGINS_CHEBYSHEV_ANISOTROPIC_PAIR_POTENTIAL_H_ +#define AZPLUGINS_CHEBYSHEV_ANISOTROPIC_PAIR_POTENTIAL_H_ + +#ifdef NVCC +#error This header cannot be compiled by nvcc +#endif + +#include +#include +#include +#include + +#include "hoomd/ForceCompute.h" +#include "hoomd/GPUArray.h" +#include "hoomd/HOOMDMath.h" +#include "hoomd/Index1D.h" + +#include "hoomd/md/NeighborList.h" + +namespace hoomd + { +namespace azplugins + { + +class LinearInterpolator5D; + +class PYBIND11_EXPORT ChebyshevAnisotropicPairPotential : public ForceCompute + { + public: + //! Constructor + ChebyshevAnisotropicPairPotential(std::shared_ptr sysdef, + std::shared_ptr nlist, + const Scalar* domain, + const Scalar* r0_data, + const unsigned int* r0_shape, + unsigned int Nterms, + const unsigned int* terms, + const Scalar* coeffs); + + //! Destructor + virtual ~ChebyshevAnisotropicPairPotential(); + + // Getters + std::shared_ptr getNeighborList() const + { + return m_nlist; + } + + /// 6x2 domain: stored as 6 entries of Scalar2 = (min,max) + const GPUArray& getApproximationDomain() const + { + return m_domain; + } + + protected: + void computeForces(uint64_t timestep) override; + + // neighbor list object + std::shared_ptr m_nlist; + + // approximation domain (6x2): 6 rows, each is (min,max) + GPUArray m_domain; + + // intenal r0 linear interpolator + std::unique_ptr m_r0_interp; + + // r0_data + GPUArray m_r0_data; + + std::array m_r0_shape; + + // Chebyshev term list (Nterms x 6) + GPUArray m_terms; + + // coeffs (Nterms) + GPUArray m_coeffs; + + // number of terms + unsigned int m_Nterms = 0; + }; + +namespace detail + { +///! exports to Python +void export_ChebyshevAnisotropicPairPotential(pybind11::module& m); + } // end namespace detail + + } // end namespace azplugins + } // end namespace hoomd + +#endif // AZPLUGINS_CHEBYSHEV_ANISOTROPIC_PAIR_POTENTIAL_H_ diff --git a/src/ChebyshevTensorAnisotropicPairPotential.h b/src/ChebyshevTensorAnisotropicPairPotential.h deleted file mode 100644 index 7d2374f..0000000 --- a/src/ChebyshevTensorAnisotropicPairPotential.h +++ /dev/null @@ -1,123 +0,0 @@ -// Copyright (c) 2018-2020, Michael P. Howard -// Copyright (c) 2021-2025, Auburn University -// Part of azplugins, released under the BSD 3-Clause License. - -/*! - * \file ChebyshevTensorAnisotropicPairPotential.h - * \brief Declaration of ChebyshevTensorAnisotropicPairPotential - */ - -#ifndef AZPLUGINS_CHEBYSHEV_TENSOR_ANISOTROPIC_PAIR_POTENTIAL_H_ -#define AZPLUGINS_CHEBYSHEV_TENSOR_ANISOTROPIC_PAIR_POTENTIAL_H_ - -#ifdef NVCC -#error This header cannot be compiled by nvcc -#endif - -#include -#include -#include -#include - -#include "hoomd/ForceCompute.h" -#include "hoomd/GPUArray.h" -#include "hoomd/HOOMDMath.h" -#include "hoomd/Index1D.h" - -#include "hoomd/md/NeighborList.h" - -namespace hoomd - { -namespace azplugins - { - -class LinearInterpolator5D; - -class PYBIND11_EXPORT ChebyshevTensorAnisotropicPairPotential : public ForceCompute - { - public: - //! Constructor - ChebyshevTensorAnisotropicPairPotential(std::shared_ptr sysdef, - std::shared_ptr nlist, - const std::array& domain, - const std::vector& r0_data, - const std::vector& terms, - const std::vector& coeffs); - - //! Destructor - virtual ~ChebyshevTensorAnisotropicPairPotential(); - - // Getters - std::shared_ptr getNeighborList() const - { - return m_nlist; - } - - /// 6x2 domain: stored as 6 entries of Scalar2 = (min,max) - const GPUArray& getApproximationDomain() const - { - return m_domain; - } - - /// r0 data (theta, phi, alpha, beta, gamma) (N x 6) - const GPUArray& getR0Data() const - { - return m_r0_data; - } - - /// Term degrees (Nterms x 6) - const GPUArray& getChebyshevTermList() const - { - return m_terms; - } - - /// Coefficients (length = Nterms) - const GPUArray& getCoefficients() const - { - return m_coeffs; - } - - unsigned int getNTerms() const - { - return m_Nterms; - } - - /// Allocate storage for term list and coefficients. - void resizeTerms(unsigned int Nterms); - - protected: - void computeForces(uint64_t timestep) override; - - private: - // 1) neighbor list object - std::shared_ptr m_nlist; - - // 2) approximation domain (6x2): 6 rows, each is (min,max) - GPUArray m_domain; - - // 3) intenal r0 linear interpolator - std::unique_ptr m_r0_interp; - - // 4) r0_data - GPUArray m_r0_data; - - // 5) Chebyshev term list (Nterms x 6) - GPUArray m_terms; - - // 6) coeffs (Nterms) - GPUArray m_coeffs; - - // 7) number of terms - unsigned int m_Nterms = 0; - }; - -namespace detail - { -///! exports to Python -void export_ChebyshevTensorAnisotropicPairPotential(pybind11::module& m); - } // end namespace detail - - } // end namespace azplugins - } // end namespace hoomd - -#endif // AZPLUGINS_CHEBYSHEV_TENSOR_ANISOTROPIC_PAIR_POTENTIAL_H_ From 67e1bdc95acaad92363221be426a1d69c9b004f1 Mon Sep 17 00:00:00 2001 From: Mohammadreza Date: Fri, 27 Feb 2026 10:02:57 -0600 Subject: [PATCH 4/5] Update Chebyshev anisotropic pair potential and add linear interpolant --- src/CMakeLists.txt | 1 + src/ChebyshevAnisotropicPairPotential.cc | 129 ++++++++++++++++ src/ChebyshevAnisotropicPairPotential.h | 48 +++--- src/LinearInterpolator5D.h | 186 +++++++++++++++++++++++ src/module.cc | 2 + src/pair.py | 56 +++++++ src/pytest/CMakeLists.txt | 1 + src/pytest/test_chebyshev.py | 71 +++++++++ 8 files changed, 473 insertions(+), 21 deletions(-) create mode 100644 src/ChebyshevAnisotropicPairPotential.cc create mode 100644 src/LinearInterpolator5D.h create mode 100644 src/pytest/test_chebyshev.py diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index a5529d9..304b83f 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -3,6 +3,7 @@ set(COMPONENT_NAME azplugins) # TODO: List all host C++ source code files in _${COMPONENT_NAME}_sources. set(_${COMPONENT_NAME}_sources + ChebyshevAnisotropicPairPotential.cc ConstantFlow.cc export_ImagePotentialBondHarmonic.cc module.cc diff --git a/src/ChebyshevAnisotropicPairPotential.cc b/src/ChebyshevAnisotropicPairPotential.cc new file mode 100644 index 0000000..b77ae2f --- /dev/null +++ b/src/ChebyshevAnisotropicPairPotential.cc @@ -0,0 +1,129 @@ +// Copyright (c) 2018-2020, Michael P. Howard +// Copyright (c) 2021-2025, Auburn University +// Part of azplugins, released under the BSD 3-Clause License. + +/*! + * \file ChebyshevAnisotropicPairPotential.h + * \brief Definition of ChebyshevAnisotropicPairPotential + */ + +#include "ChebyshevAnisotropicPairPotential.h" +#include "LinearInterpolator5D.h" + +namespace hoomd + { +namespace azplugins + { + +ChebyshevAnisotropicPairPotential::ChebyshevAnisotropicPairPotential( + std::shared_ptr sysdef, + std::shared_ptr nlist, + const Scalar* domain, + const float r_cut, + const unsigned int* terms, + const Scalar* coeffs, + unsigned int Nterms, + const Scalar* r0_data, + const unsigned int* r0_shape) + : ForceCompute(sysdef), m_nlist(nlist), m_r_cut(r_cut), m_Nterms(Nterms) + { + } + +ChebyshevAnisotropicPairPotential::~ChebyshevAnisotropicPairPotential() { } + +void ChebyshevAnisotropicPairPotential::computeForces(uint64_t timestep) + { + if (m_nlist) + { + m_nlist->compute(timestep); + } + + const unsigned int N = m_pdata->getN(); + + ArrayHandle h_force(m_force, access_location::host, access_mode::readwrite); + ArrayHandle h_torque(m_torque, access_location::host, access_mode::readwrite); + + for (unsigned int i = 0; i < N; ++i) + { + h_force.data[i] = make_scalar4(Scalar(0), Scalar(0), Scalar(0), Scalar(0)); + h_torque.data[i] = make_scalar4(Scalar(0), Scalar(0), Scalar(0), Scalar(0)); + } + } + +namespace detail + { + +void export_ChebyshevAnisotropicPairPotential(pybind11::module& m) + { + namespace py = pybind11; + using NL = hoomd::md::NeighborList; + + py::class_>( + m, + "ChebyshevAnisotropicPairPotential", + py::base()) + .def(py::init( + [](std::shared_ptr sysdef, + std::shared_ptr nlist, + py::array_t domain, + float r_cut, + py::array_t terms, + py::array_t coeffs, + py::array_t r0_data) + { + // domain must be (5,2) - rho is always in (0, 1) + if (domain.ndim() != 2 || domain.shape(0) != 5 || domain.shape(1) != 2) + { + throw std::runtime_error("domain must have shape (5,2)."); + } + + // terms must be (Nterms,6) + if (terms.ndim() != 2 || terms.shape(1) != 6) + { + throw std::runtime_error("terms must have shape (Nterms,6)."); + } + + const unsigned int Nterms = static_cast(terms.shape(0)); + + // coeffs must be (Nterms,) + if (coeffs.ndim() != 1 || static_cast(coeffs.shape(0)) != Nterms) + { + throw std::runtime_error("coeffs must have shape (Nterms,)."); + } + + // r0_data must be 5D + if (r0_data.ndim() != 5) + { + throw std::runtime_error("r0_data must be a 5D array."); + } + + // Infer r0_shape from r0_data.shape + std::array r0_shape; + for (unsigned int k = 0; k < 5; ++k) + { + const auto dim = r0_data.shape(k); + if (dim < 2) + { + throw std::runtime_error("r0_data has invalid dimension(s)."); + } + r0_shape[k] = static_cast(dim); + } + + return std::make_shared(sysdef, + nlist, + domain.data(), + r_cut, + terms.data(), + coeffs.data(), + Nterms, + r0_data.data(), + r0_shape.data()); + })) + .def_property_readonly("r_cut", &ChebyshevAnisotropicPairPotential::getRCut) + .def_property_readonly("n_terms", &ChebyshevAnisotropicPairPotential::getNTerms); + } + + } // end namespace detail + } // namespace azplugins + } // namespace hoomd diff --git a/src/ChebyshevAnisotropicPairPotential.h b/src/ChebyshevAnisotropicPairPotential.h index dd14b25..d054907 100644 --- a/src/ChebyshevAnisotropicPairPotential.h +++ b/src/ChebyshevAnisotropicPairPotential.h @@ -31,7 +31,7 @@ namespace hoomd namespace azplugins { -class LinearInterpolator5D; +template class LinearInterpolator5D; class PYBIND11_EXPORT ChebyshevAnisotropicPairPotential : public ForceCompute { @@ -40,11 +40,12 @@ class PYBIND11_EXPORT ChebyshevAnisotropicPairPotential : public ForceCompute ChebyshevAnisotropicPairPotential(std::shared_ptr sysdef, std::shared_ptr nlist, const Scalar* domain, - const Scalar* r0_data, - const unsigned int* r0_shape, - unsigned int Nterms, + const float r_cut, const unsigned int* terms, - const Scalar* coeffs); + const Scalar* coeffs, + unsigned int Nterms, + const Scalar* r0_data, + const unsigned int* r0_shape); //! Destructor virtual ~ChebyshevAnisotropicPairPotential(); @@ -55,37 +56,42 @@ class PYBIND11_EXPORT ChebyshevAnisotropicPairPotential : public ForceCompute return m_nlist; } - /// 6x2 domain: stored as 6 entries of Scalar2 = (min,max) + /// 5x2 domain: stored as 5 entries of Scalar2 = (min,max) const GPUArray& getApproximationDomain() const { return m_domain; } + /// Read-only cutoff radius + const float getRCut() const + { + return m_r_cut; + } + + /// Read-only number of Chebyshev terms + unsigned int getNTerms() const + { + return m_Nterms; + } + protected: void computeForces(uint64_t timestep) override; - // neighbor list object - std::shared_ptr m_nlist; + std::shared_ptr m_nlist; //!< Neighbor list - // approximation domain (6x2): 6 rows, each is (min,max) - GPUArray m_domain; + GPUArray m_domain; //!< Approximation domain (5x2): 5 rows, each is (min, max) - // intenal r0 linear interpolator - std::unique_ptr m_r0_interp; + float m_r_cut; //!< cut-off distance in approximation domain - // r0_data - GPUArray m_r0_data; + GPUArray m_terms; //!< Chebyshev term list (Nterms x 6) - std::array m_r0_shape; + GPUArray m_coeffs; //!< Coefficients corresponding to each term - // Chebyshev term list (Nterms x 6) - GPUArray m_terms; + unsigned int m_Nterms; //!< Number of terms - // coeffs (Nterms) - GPUArray m_coeffs; + GPUArray m_r0_data; //!< R0 data - // number of terms - unsigned int m_Nterms = 0; + GPUArray m_r0_shape; //!< Number of points used along each dimension to sample r0 }; namespace detail diff --git a/src/LinearInterpolator5D.h b/src/LinearInterpolator5D.h new file mode 100644 index 0000000..9e1f292 --- /dev/null +++ b/src/LinearInterpolator5D.h @@ -0,0 +1,186 @@ +// Copyright (c) 2018-2020, Michael P. Howard +// Copyright (c) 2021-2025, Auburn University +// Part of azplugins, released under the BSD 3-Clause License. + +#ifndef AZPLUGINS_LINEAR_INTERPOLATOR_5D_H_ +#define AZPLUGINS_LINEAR_INTERPOLATOR_5D_H_ + +#include +#include +#include + +#include "hoomd/HOOMDMath.h" + +#if defined(__HIPCC__) || defined(__CUDACC__) +#define AZPLUGINS_HOSTDEVICE __host__ __device__ +#define AZPLUGINS_FORCEINLINE __forceinline__ +#else +#define AZPLUGINS_HOSTDEVICE +#define AZPLUGINS_FORCEINLINE inline +#endif + +namespace hoomd + { +namespace azplugins + { + +class FiveDimensionalIndex + { + public: + AZPLUGINS_HOSTDEVICE AZPLUGINS_FORCEINLINE FiveDimensionalIndex() + : n0_(0), n1_(0), n2_(0), n3_(0), n4_(0) + { + } + + AZPLUGINS_HOSTDEVICE AZPLUGINS_FORCEINLINE FiveDimensionalIndex(unsigned int n0, + unsigned int n1, + unsigned int n2, + unsigned int n3, + unsigned int n4) + : n0_(n0), n1_(n1), n2_(n2), n3_(n3), n4_(n4) + { + } + + AZPLUGINS_HOSTDEVICE AZPLUGINS_FORCEINLINE unsigned int operator()(unsigned int i0, + unsigned int i1, + unsigned int i2, + unsigned int i3, + unsigned int i4) const + { + unsigned int idx = i0; + idx = idx * n1_ + i1; + idx = idx * n2_ + i2; + idx = idx * n3_ + i3; + idx = idx * n4_ + i4; + return idx; + } + + AZPLUGINS_HOSTDEVICE AZPLUGINS_FORCEINLINE unsigned int size() const + { + return n0_ * n1_ * n2_ * n3_ * n4_; + } + + private: + unsigned int n0_, n1_, n2_, n3_, n4_; + }; + +// T is the stored data type. +template class LinearInterpolator5D + { + public: + AZPLUGINS_HOSTDEVICE AZPLUGINS_FORCEINLINE LinearInterpolator5D() : data_(nullptr), nindex_() + { + for (int d = 0; d < 5; ++d) + { + n_[d] = 0; + lo_[d] = Scalar(0); + hi_[d] = Scalar(0); + dx_[d] = Scalar(0); + } + } + + AZPLUGINS_HOSTDEVICE AZPLUGINS_FORCEINLINE + LinearInterpolator5D(const T* data, const unsigned int* n, const Scalar* lo, const Scalar* hi) + : data_(data) + { + for (int d = 0; d < 5; ++d) + { + n_[d] = n[d]; + lo_[d] = lo[d]; + hi_[d] = hi[d]; + + assert(n_[d] >= 2); + dx_[d] = (hi_[d] - lo_[d]) / Scalar(n_[d] - 1); + } + + nindex_ = FiveDimensionalIndex(n_[0], n_[1], n_[2], n_[3], n_[4]); + } + + // Interpolate at (x0, x1, x2, x3, x4). + AZPLUGINS_HOSTDEVICE AZPLUGINS_FORCEINLINE Scalar + operator()(Scalar x0, Scalar x1, Scalar x2, Scalar x3, Scalar x4) const + { + const Scalar x[5] = {x0, x1, x2, x3, x4}; + + Scalar f[5]; + for (int d = 0; d < 5; ++d) + { + f[d] = (x[d] - lo_[d]) / dx_[d]; + } + + int bin[5]; + Scalar dloc[5]; + + for (int dim = 0; dim < 5; ++dim) + { + bin[dim] = (int)std::floor((double)f[dim]); + + if (f[dim] == Scalar(n_[dim] - 1) && x[dim] == hi_[dim]) + { + --bin[dim]; + } + + dloc[dim] = f[dim] - Scalar(bin[dim]); + + assert(bin[dim] >= 0); + assert(bin[dim] < (int)n_[dim] - 1); + } + + const Scalar xd = dloc[0]; + const Scalar yd = dloc[1]; + const Scalar zd = dloc[2]; + const Scalar wd = dloc[3]; + const Scalar vd = dloc[4]; + + Scalar c[32]; + for (unsigned int mask = 0; mask < 32; ++mask) + { + const unsigned int i0 = (unsigned int)(bin[0] + ((mask >> 0) & 1u)); + const unsigned int i1 = (unsigned int)(bin[1] + ((mask >> 1) & 1u)); + const unsigned int i2 = (unsigned int)(bin[2] + ((mask >> 2) & 1u)); + const unsigned int i3 = (unsigned int)(bin[3] + ((mask >> 3) & 1u)); + const unsigned int i4 = (unsigned int)(bin[4] + ((mask >> 4) & 1u)); + + c[mask] = Scalar(data_[nindex_(i0, i1, i2, i3, i4)]); + } + + Scalar c0[16]; + for (unsigned int i = 0; i < 16; ++i) + { + c0[i] = c[2 * i] * (Scalar(1) - xd) + c[2 * i + 1] * xd; + } + + Scalar c1[8]; + for (unsigned int i = 0; i < 8; ++i) + { + c1[i] = c0[2 * i] * (Scalar(1) - yd) + c0[2 * i + 1] * yd; + } + + Scalar c2[4]; + for (unsigned int i = 0; i < 4; ++i) + { + c2[i] = c1[2 * i] * (Scalar(1) - zd) + c1[2 * i + 1] * zd; + } + + Scalar c3[2]; + for (unsigned int i = 0; i < 2; ++i) + { + c3[i] = c2[2 * i] * (Scalar(1) - wd) + c2[2 * i + 1] * wd; + } + + return c3[0] * (Scalar(1) - vd) + c3[1] * vd; + } + + private: + const T* data_; + unsigned int n_[5]; + Scalar lo_[5]; + Scalar hi_[5]; + Scalar dx_[5]; + FiveDimensionalIndex nindex_; + }; + + } // namespace azplugins + } // namespace hoomd + +#endif // AZPLUGINS_LINEAR_INTERPOLATOR_5D_H_ diff --git a/src/module.cc b/src/module.cc index 1ab81c7..d22f6dd 100644 --- a/src/module.cc +++ b/src/module.cc @@ -69,6 +69,7 @@ void export_ParabolicFlow(pybind11::module&); // pair void export_AnisoPotentialPairTwoPatchMorse(pybind11::module&); +void export_ChebyshevAnisotropicPairPotential(pybind11::module&); void export_PotentialPairColloid(pybind11::module&); void export_PotentialPairExpandedYukawa(pybind11::module&); void export_PotentialPairHertz(pybind11::module&); @@ -141,6 +142,7 @@ PYBIND11_MODULE(_azplugins, m) // pair export_AnisoPotentialPairTwoPatchMorse(m); + export_ChebyshevAnisotropicPairPotential(m); export_PotentialPairColloid(m); export_PotentialPairExpandedYukawa(m); export_PotentialPairHertz(m); diff --git a/src/pair.py b/src/pair.py index fbefb07..8e982ae 100644 --- a/src/pair.py +++ b/src/pair.py @@ -4,13 +4,69 @@ """Pair potentials.""" +import numpy from hoomd.azplugins import _azplugins from hoomd.data.parameterdicts import ParameterDict, TypeParameterDict from hoomd.data.typeparam import TypeParameter from hoomd.md import pair +from hoomd.md.force import Force from hoomd.variant import Variant +class ChebyshevAnisotropicPairPotential(Force): + """Chebyshev anisotropic pair potential.""" + + def __init__(self, nlist, domain, terms, coeffs, r0_data, r_cut=3.0): + super().__init__() + self._nlist = nlist + + self._domain = numpy.asarray(domain) + self._r_cut = float(r_cut) + self._terms = numpy.asarray(terms, dtype=numpy.uint32) + self._coeffs = numpy.asarray(coeffs) + self._r0_data = numpy.asarray(r0_data) + + if self._domain.shape != (5, 2): + raise ValueError("domain must have shape (5, 2).") + if self._terms.ndim != 2 or self._terms.shape[1] != 6: + raise ValueError("terms must have shape (Nterms, 6).") + nterms = int(self._terms.shape[0]) + if self._coeffs.ndim != 1 or int(self._coeffs.shape[0]) != nterms: + raise ValueError("coeffs must have shape (Nterms,).") + if self._r0_data.ndim != 5: + raise ValueError("r0_data must be a 5D array.") + + @property + def r_cut(self): + """Cut-off distance in approximation domain""" + return self._r_cut + + @property + def n_terms(self): + """Number of terms.""" + return int(self._terms.shape[0]) + + @property + def r0_shape(self): + """r0 table shape.""" + return tuple(int(x) for x in self._r0_data.shape) + + def _attach_hook(self): + self._nlist._attach(self._simulation) + + self._cpp_obj = _azplugins.ChebyshevAnisotropicPairPotential( + self._simulation.state._cpp_sys_def, + self._nlist._cpp_obj, + self._domain, + self._r_cut, + self._terms, + self._coeffs, + self._r0_data, + ) + + super()._attach_hook() + + class Colloid(pair.Pair): r"""Colloid pair potential. diff --git a/src/pytest/CMakeLists.txt b/src/pytest/CMakeLists.txt index 157b326..4676d3a 100644 --- a/src/pytest/CMakeLists.txt +++ b/src/pytest/CMakeLists.txt @@ -3,6 +3,7 @@ set(test_files __init__.py test_bond.py test_compute.py + test_chebyshev.py test_external.py test_flow.py test_pair.py diff --git a/src/pytest/test_chebyshev.py b/src/pytest/test_chebyshev.py new file mode 100644 index 0000000..c0b16cc --- /dev/null +++ b/src/pytest/test_chebyshev.py @@ -0,0 +1,71 @@ +# Copyright (c) 2018-2020, Michael P. Howard +# Copyright (c) 2021-2025, Auburn University +# Part of azplugins, released under the BSD 3-Clause License. + +import numpy +import hoomd +import hoomd.azplugins + + +def test_attach_and_zero_force(simulation_factory, two_particle_snapshot_factory): + """Construct, attach, and check force/torque output.""" + + # Construct the Python object + nlist = hoomd.md.nlist.Cell(buffer=0.4) + + domain = numpy.zeros((5, 2), dtype=numpy.float64) + terms = numpy.zeros((2, 6), dtype=numpy.uint32) + coeffs = numpy.zeros((2,), dtype=numpy.float64) + r0_data = numpy.zeros((2, 2, 2, 2, 2), dtype=numpy.float64) + r_cut = 3.0 + + pot = hoomd.azplugins.pair.ChebyshevAnisotropicPairPotential( + nlist=nlist, + domain=domain, + r_cut=r_cut, + terms=terms, + coeffs=coeffs, + r0_data=r0_data, + ) + + # Pre-attach checks + assert numpy.isclose(pot.r_cut, r_cut) + assert pot.n_terms == 2 + assert pot.r0_shape == (2, 2, 2, 2, 2) + + # Attach via a 0-step simulation + snap = two_particle_snapshot_factory() + if snap.communicator.rank == 0: + snap.particles.position[:] = [[-0.5, 0.0, 0.0], [0.5, 0.0, 0.0]] + snap.particles.orientation[:] = [[1, 0, 0, 0], [1, 0, 0, 0]] + snap.particles.moment_inertia[:] = [0.1, 0.1, 0.1] + + sim = simulation_factory(snap) + + integrator = hoomd.md.Integrator(dt=0.001) + nve = hoomd.md.methods.ConstantVolume(hoomd.filter.All()) + integrator.methods = [nve] + + integrator.forces = [pot] + sim.operations.integrator = integrator + + # Attach all objects + sim.run(0) + + # After attach + assert hasattr(pot, "_cpp_obj") + assert pot._cpp_obj is not None + + # Post-attach checks + assert numpy.isclose(pot.r_cut, r_cut) + assert pot.n_terms == 2 + assert pot.r0_shape == (2, 2, 2, 2, 2) + + assert pot._cpp_obj.n_terms == 2 + assert numpy.isclose(pot._cpp_obj.r_cut, r_cut) + + # Check Force/torque/energy outputs + if sim.device.communicator.rank == 0: + numpy.testing.assert_array_equal(pot.forces, numpy.zeros((2, 3))) + numpy.testing.assert_array_equal(pot.torques, numpy.zeros((2, 3))) + numpy.testing.assert_array_equal(pot.energies, numpy.zeros((2,))) From 45a1740044c627952e645be430bc6b020c7bcdb2 Mon Sep 17 00:00:00 2001 From: Mohammadreza Date: Thu, 19 Mar 2026 11:27:43 -0500 Subject: [PATCH 5/5] Update linear interpolation, and initiate compute force --- src/ChebyshevAnisotropicPairPotential.cc | 322 ++++++++++++++++++++++- src/ChebyshevAnisotropicPairPotential.h | 14 +- src/LinearInterpolator5D.h | 168 ++++++------ src/pair.py | 46 ++-- src/pytest/test_chebyshev.py | 78 +++--- 5 files changed, 472 insertions(+), 156 deletions(-) diff --git a/src/ChebyshevAnisotropicPairPotential.cc b/src/ChebyshevAnisotropicPairPotential.cc index b77ae2f..bc0b04c 100644 --- a/src/ChebyshevAnisotropicPairPotential.cc +++ b/src/ChebyshevAnisotropicPairPotential.cc @@ -10,6 +10,14 @@ #include "ChebyshevAnisotropicPairPotential.h" #include "LinearInterpolator5D.h" +#include "hoomd/BoxDim.h" +#include "hoomd/VectorMath.h" + +#include +#include +#include +#include + namespace hoomd { namespace azplugins @@ -19,7 +27,7 @@ ChebyshevAnisotropicPairPotential::ChebyshevAnisotropicPairPotential( std::shared_ptr sysdef, std::shared_ptr nlist, const Scalar* domain, - const float r_cut, + const Scalar r_cut, const unsigned int* terms, const Scalar* coeffs, unsigned int Nterms, @@ -27,26 +35,324 @@ ChebyshevAnisotropicPairPotential::ChebyshevAnisotropicPairPotential( const unsigned int* r0_shape) : ForceCompute(sysdef), m_nlist(nlist), m_r_cut(r_cut), m_Nterms(Nterms) { + { + GPUArray domain_arr(5, m_exec_conf); + m_domain.swap(domain_arr); + + ArrayHandle h_domain(m_domain, access_location::host, access_mode::readwrite); + for (unsigned int d = 0; d < 5; ++d) + { + h_domain.data[d] = make_scalar2(domain[2 * d], domain[2 * d + 1]); + } + } + + // terms: shape (Nterms, 6), stored flat + { + GPUArray terms_arr(static_cast(Nterms) * 6, m_exec_conf); + m_terms.swap(terms_arr); + + ArrayHandle h_terms(m_terms, access_location::host, access_mode::readwrite); + std::copy(terms, terms + static_cast(Nterms) * 6, h_terms.data); + } + + // coeffs: shape (Nterms,) + { + GPUArray coeffs_arr(Nterms, m_exec_conf); + m_coeffs.swap(coeffs_arr); + + ArrayHandle h_coeffs(m_coeffs, access_location::host, access_mode::readwrite); + std::copy(coeffs, coeffs + Nterms, h_coeffs.data); + } + + // r0_shape: length 5 + { + GPUArray shape_arr(5, m_exec_conf); + m_r0_shape.swap(shape_arr); + + ArrayHandle h_shape(m_r0_shape, + access_location::host, + access_mode::readwrite); + std::copy(r0_shape, r0_shape + 5, h_shape.data); + } + + // r0_data: flat array, length = product(r0_shape) + size_t n_r0 = 1; + for (unsigned int d = 0; d < 5; ++d) + { + n_r0 *= static_cast(r0_shape[d]); + } + + { + GPUArray r0_arr(n_r0, m_exec_conf); + m_r0_data.swap(r0_arr); + + ArrayHandle h_r0(m_r0_data, access_location::host, access_mode::readwrite); + std::copy(r0_data, r0_data + n_r0, h_r0.data); + } } ChebyshevAnisotropicPairPotential::~ChebyshevAnisotropicPairPotential() { } void ChebyshevAnisotropicPairPotential::computeForces(uint64_t timestep) { - if (m_nlist) + // start by updating the neighborlist + m_nlist->compute(timestep); + + // access neighbor list, particle data, and simulation box. + ArrayHandle h_n_neigh(m_nlist->getNNeighArray(), + access_location::host, + access_mode::read); + ArrayHandle h_nlist(m_nlist->getNListArray(), + access_location::host, + access_mode::read); + ArrayHandle h_head_list(m_nlist->getHeadList(), + access_location::host, + access_mode::read); + ArrayHandle h_pos(m_pdata->getPositions(), access_location::host, access_mode::read); + ArrayHandle h_orientation(m_pdata->getOrientationArray(), + access_location::host, + access_mode::read); + ArrayHandle h_domain(m_domain, access_location::host, access_mode::read); + ArrayHandle h_r0_data(m_r0_data, access_location::host, access_mode::read); + ArrayHandle h_r0_shape(m_r0_shape, access_location::host, access_mode::read); + + const BoxDim box = m_pdata->getGlobalBox(); + const Scalar rcutsq = m_r_cut * m_r_cut; + const Scalar h = Scalar(1.0e-6); + + Scalar lo[5]; + Scalar hi[5]; + for (unsigned int d = 0; d < 5; ++d) { - m_nlist->compute(timestep); + lo[d] = h_domain.data[d].x; + hi[d] = h_domain.data[d].y; } - const unsigned int N = m_pdata->getN(); + LinearInterpolator5D interp(h_r0_data.data, h_r0_shape.data, lo, hi); + + // need to start from a zero force and torque + m_force.zeroFill(); + m_torque.zeroFill(); ArrayHandle h_force(m_force, access_location::host, access_mode::readwrite); ArrayHandle h_torque(m_torque, access_location::host, access_mode::readwrite); + const unsigned int N = m_pdata->getN(); + for (unsigned int i = 0; i < N; ++i) { - h_force.data[i] = make_scalar4(Scalar(0), Scalar(0), Scalar(0), Scalar(0)); - h_torque.data[i] = make_scalar4(Scalar(0), Scalar(0), Scalar(0), Scalar(0)); + // particle i position and orientation + const Scalar3 pi = make_scalar3(h_pos.data[i].x, h_pos.data[i].y, h_pos.data[i].z); + const quat q_i(h_orientation.data[i]); + const quat q_i_conj = conj(q_i); + + // initialize current particle force and torque + Scalar3 fi = make_scalar3(0, 0, 0); + Scalar3 ti = make_scalar3(0, 0, 0); + + const size_t myHead = h_head_list.data[i]; + const unsigned int size = (unsigned int)h_n_neigh.data[i]; + + for (unsigned int k = 0; k < size; ++k) + { + // access the index + const unsigned int j = h_nlist.data[myHead + k]; + assert(j < m_pdata->getN() + m_pdata->getNGhosts()); + + const Scalar3 pj = make_scalar3(h_pos.data[j].x, h_pos.data[j].y, h_pos.data[j].z); + Scalar3 dx = pi - pj; + // apply periodic boundary conditions + dx = box.minImage(dx); + + // cut-off check + const Scalar rsq = dot(dx, dx); + if (rsq > rcutsq) + { + continue; + } + + // particle j, orientation quaternion + const quat q_j(h_orientation.data[j]); + // dx is in lab frame, so rotate dx by conj(q_i) + const vec3 dx_lab(dx.x, dx.y, dx.z); + const vec3 dx_body = rotate(q_i_conj, dx_lab); + // relative orientation of j with respect to i + const quat q_rel = q_i_conj * q_j; + + // convert position to spherical coordinates + const Scalar r = fast::sqrt(dot(dx_body, dx_body)); + Scalar theta = Scalar(0); + Scalar phi = Scalar(0); + + if (r > Scalar(0)) + { + theta = std::atan2(dx_body.y, dx_body.x); + if (theta < Scalar(0)) + { + theta += Scalar(2.0) * M_PI; + } + + Scalar cosphi = dx_body.z / r; + if (cosphi < Scalar(-1)) + { + cosphi = Scalar(-1); + } + else if (cosphi > Scalar(1)) + { + cosphi = Scalar(1); + } + + phi = std::acos(cosphi); + } + + // get the columns of an active rotation matrix + const vec3 ex = rotate(q_rel, vec3(1, 0, 0)); + const vec3 ey = rotate(q_rel, vec3(0, 1, 0)); + const vec3 ez = rotate(q_rel, vec3(0, 0, 1)); + + Scalar alpha = Scalar(0); + Scalar beta = Scalar(0); + Scalar gamma = Scalar(0); + + // get the rotation angles by R_ZXZ (body-fixed) = R_q + if (ez.z < Scalar(-1)) + { + beta = Scalar(M_PI); + } + else if (ez.z > Scalar(1)) + { + beta = Scalar(0); + } + else + { + beta = std::acos(ez.z); + } + + if (beta > Scalar(1e-7) && beta < Scalar(M_PI - 1e-7)) + { + alpha = std::atan2(ez.x, -ez.y); + gamma = std::atan2(ex.z, ey.z); + } + else if (beta <= Scalar(1e-7)) + { + alpha = Scalar(0); + gamma = std::atan2(ex.y, ex.x); + } + else + { + alpha = Scalar(0); + gamma = std::atan2(-ex.y, ex.x); + } + + if (alpha < Scalar(0)) + { + alpha += Scalar(2) * M_PI; + } + if (gamma < Scalar(0)) + { + gamma += Scalar(2) * M_PI; + } + + // compute r0 and its derivatives + const Scalar r0 = interp(theta, phi, alpha, beta, gamma); + Scalar dr0_dtheta = Scalar(0); + Scalar dr0_dphi = Scalar(0); + Scalar dr0_dalpha = Scalar(0); + Scalar dr0_dbeta = Scalar(0); + Scalar dr0_dgamma = Scalar(0); + + // d r0 / d theta + if (theta - h < lo[0]) + { + dr0_dtheta = (interp(theta + h, phi, alpha, beta, gamma) - r0) / h; + } + else if (theta + h > hi[0]) + { + dr0_dtheta = (r0 - interp(theta - h, phi, alpha, beta, gamma)) / h; + } + else + { + dr0_dtheta = (interp(theta + h, phi, alpha, beta, gamma) + - interp(theta - h, phi, alpha, beta, gamma)) + / (Scalar(2) * h); + } + + // d r0 / d phi + if (phi - h < lo[1]) + { + dr0_dphi = (interp(theta, phi + h, alpha, beta, gamma) - r0) / h; + } + else if (phi + h > hi[1]) + { + dr0_dphi = (r0 - interp(theta, phi - h, alpha, beta, gamma)) / h; + } + else + { + dr0_dphi = (interp(theta, phi + h, alpha, beta, gamma) + - interp(theta, phi - h, alpha, beta, gamma)) + / (Scalar(2) * h); + } + + // d r0 / d alpha + if (alpha - h < lo[2]) + { + dr0_dalpha = (interp(theta, phi, alpha + h, beta, gamma) - r0) / h; + } + else if (alpha + h > hi[2]) + { + dr0_dalpha = (r0 - interp(theta, phi, alpha - h, beta, gamma)) / h; + } + else + { + dr0_dalpha = (interp(theta, phi, alpha + h, beta, gamma) + - interp(theta, phi, alpha - h, beta, gamma)) + / (Scalar(2) * h); + } + + // d r0 / d beta + if (beta - h < lo[3]) + { + dr0_dbeta = (interp(theta, phi, alpha, beta + h, gamma) - r0) / h; + } + else if (beta + h > hi[3]) + { + dr0_dbeta = (r0 - interp(theta, phi, alpha, beta - h, gamma)) / h; + } + else + { + dr0_dbeta = (interp(theta, phi, alpha, beta + h, gamma) + - interp(theta, phi, alpha, beta - h, gamma)) + / (Scalar(2) * h); + } + + // d r0 / d gamma + if (gamma - h < lo[4]) + { + dr0_dgamma = (interp(theta, phi, alpha, beta, gamma + h) - r0) / h; + } + else if (gamma + h > hi[4]) + { + dr0_dgamma = (r0 - interp(theta, phi, alpha, beta, gamma - h)) / h; + } + else + { + dr0_dgamma = (interp(theta, phi, alpha, beta, gamma + h) + - interp(theta, phi, alpha, beta, gamma - h)) + / (Scalar(2) * h); + } + + // compute J + } + + h_force.data[i].x += fi.x; + h_force.data[i].y += fi.y; + h_force.data[i].z += fi.z; + h_force.data[i].w += Scalar(0.0); + + h_torque.data[i].x += ti.x; + h_torque.data[i].y += ti.y; + h_torque.data[i].z += ti.z; + h_torque.data[i].w += Scalar(0.0); } } @@ -67,7 +373,7 @@ void export_ChebyshevAnisotropicPairPotential(pybind11::module& m) [](std::shared_ptr sysdef, std::shared_ptr nlist, py::array_t domain, - float r_cut, + Scalar r_cut, py::array_t terms, py::array_t coeffs, py::array_t r0_data) @@ -121,7 +427,7 @@ void export_ChebyshevAnisotropicPairPotential(pybind11::module& m) r0_shape.data()); })) .def_property_readonly("r_cut", &ChebyshevAnisotropicPairPotential::getRCut) - .def_property_readonly("n_terms", &ChebyshevAnisotropicPairPotential::getNTerms); + .def_property_readonly("num_terms", &ChebyshevAnisotropicPairPotential::getNTerms); } } // end namespace detail diff --git a/src/ChebyshevAnisotropicPairPotential.h b/src/ChebyshevAnisotropicPairPotential.h index d054907..7fc549b 100644 --- a/src/ChebyshevAnisotropicPairPotential.h +++ b/src/ChebyshevAnisotropicPairPotential.h @@ -31,8 +31,6 @@ namespace hoomd namespace azplugins { -template class LinearInterpolator5D; - class PYBIND11_EXPORT ChebyshevAnisotropicPairPotential : public ForceCompute { public: @@ -40,7 +38,7 @@ class PYBIND11_EXPORT ChebyshevAnisotropicPairPotential : public ForceCompute ChebyshevAnisotropicPairPotential(std::shared_ptr sysdef, std::shared_ptr nlist, const Scalar* domain, - const float r_cut, + const Scalar r_cut, const unsigned int* terms, const Scalar* coeffs, unsigned int Nterms, @@ -63,7 +61,7 @@ class PYBIND11_EXPORT ChebyshevAnisotropicPairPotential : public ForceCompute } /// Read-only cutoff radius - const float getRCut() const + Scalar getRCut() const { return m_r_cut; } @@ -81,7 +79,7 @@ class PYBIND11_EXPORT ChebyshevAnisotropicPairPotential : public ForceCompute GPUArray m_domain; //!< Approximation domain (5x2): 5 rows, each is (min, max) - float m_r_cut; //!< cut-off distance in approximation domain + Scalar m_r_cut; //!< cut-off distance in approximation domain GPUArray m_terms; //!< Chebyshev term list (Nterms x 6) @@ -94,12 +92,6 @@ class PYBIND11_EXPORT ChebyshevAnisotropicPairPotential : public ForceCompute GPUArray m_r0_shape; //!< Number of points used along each dimension to sample r0 }; -namespace detail - { -///! exports to Python -void export_ChebyshevAnisotropicPairPotential(pybind11::module& m); - } // end namespace detail - } // end namespace azplugins } // end namespace hoomd diff --git a/src/LinearInterpolator5D.h b/src/LinearInterpolator5D.h index 9e1f292..7340434 100644 --- a/src/LinearInterpolator5D.h +++ b/src/LinearInterpolator5D.h @@ -27,8 +27,10 @@ namespace azplugins class FiveDimensionalIndex { public: - AZPLUGINS_HOSTDEVICE AZPLUGINS_FORCEINLINE FiveDimensionalIndex() - : n0_(0), n1_(0), n2_(0), n3_(0), n4_(0) + AZPLUGINS_HOSTDEVICE AZPLUGINS_FORCEINLINE FiveDimensionalIndex() : m_n {0, 0, 0, 0, 0} { } + + AZPLUGINS_HOSTDEVICE AZPLUGINS_FORCEINLINE explicit FiveDimensionalIndex(const unsigned int* n) + : m_n {n[0], n[1], n[2], n[3], n[4]} { } @@ -37,7 +39,7 @@ class FiveDimensionalIndex unsigned int n2, unsigned int n3, unsigned int n4) - : n0_(n0), n1_(n1), n2_(n2), n3_(n3), n4_(n4) + : m_n {n0, n1, n2, n3, n4} { } @@ -48,139 +50,153 @@ class FiveDimensionalIndex unsigned int i4) const { unsigned int idx = i0; - idx = idx * n1_ + i1; - idx = idx * n2_ + i2; - idx = idx * n3_ + i3; - idx = idx * n4_ + i4; + idx = idx * m_n[1] + i1; + idx = idx * m_n[2] + i2; + idx = idx * m_n[3] + i3; + idx = idx * m_n[4] + i4; return idx; } AZPLUGINS_HOSTDEVICE AZPLUGINS_FORCEINLINE unsigned int size() const { - return n0_ * n1_ * n2_ * n3_ * n4_; + return m_n[0] * m_n[1] * m_n[2] * m_n[3] * m_n[4]; + } + + AZPLUGINS_HOSTDEVICE AZPLUGINS_FORCEINLINE unsigned int getN(unsigned int dim) const + { + return m_n[dim]; } private: - unsigned int n0_, n1_, n2_, n3_, n4_; + unsigned int m_n[5]; }; -// T is the stored data type. +/*! \brief 5D multilinear interpolation on a uniform rectilinear grid. + + This is an extension of three-dimensional linear interpolation + from (https://github.com/mphowardlab/flyft/blob/main/src/grid_interpolator.cc). + +*/ template class LinearInterpolator5D { public: - AZPLUGINS_HOSTDEVICE AZPLUGINS_FORCEINLINE LinearInterpolator5D() : data_(nullptr), nindex_() + AZPLUGINS_HOSTDEVICE AZPLUGINS_FORCEINLINE LinearInterpolator5D() : m_data(nullptr), m_indexer() { for (int d = 0; d < 5; ++d) { - n_[d] = 0; - lo_[d] = Scalar(0); - hi_[d] = Scalar(0); - dx_[d] = Scalar(0); + m_lo[d] = Scalar(0); + m_hi[d] = Scalar(0); + m_dx[d] = Scalar(0); } } AZPLUGINS_HOSTDEVICE AZPLUGINS_FORCEINLINE LinearInterpolator5D(const T* data, const unsigned int* n, const Scalar* lo, const Scalar* hi) - : data_(data) + : m_data(data), m_indexer(n) { for (int d = 0; d < 5; ++d) { - n_[d] = n[d]; - lo_[d] = lo[d]; - hi_[d] = hi[d]; + const unsigned int nd = n[d]; + assert(nd >= 2); - assert(n_[d] >= 2); - dx_[d] = (hi_[d] - lo_[d]) / Scalar(n_[d] - 1); + m_lo[d] = lo[d]; + m_hi[d] = hi[d]; + m_dx[d] = (m_hi[d] - m_lo[d]) / Scalar(nd - 1); } - nindex_ = FiveDimensionalIndex(n_[0], n_[1], n_[2], n_[3], n_[4]); + assert(m_indexer.size() > 0); } - // Interpolate at (x0, x1, x2, x3, x4). + //! Interpolate at (x0, x1, x2, x3, x4). AZPLUGINS_HOSTDEVICE AZPLUGINS_FORCEINLINE Scalar operator()(Scalar x0, Scalar x1, Scalar x2, Scalar x3, Scalar x4) const { const Scalar x[5] = {x0, x1, x2, x3, x4}; - Scalar f[5]; - for (int d = 0; d < 5; ++d) - { - f[d] = (x[d] - lo_[d]) / dx_[d]; - } - + // Compute the cell bin and fractional coordinate in each dimension. int bin[5]; - Scalar dloc[5]; + Scalar frac[5]; - for (int dim = 0; dim < 5; ++dim) + for (int d = 0; d < 5; ++d) { - bin[dim] = (int)std::floor((double)f[dim]); + const unsigned int nd = m_indexer.getN(static_cast(d)); + const Scalar f = (x[d] - m_lo[d]) / m_dx[d]; + + int b = static_cast(std::floor(static_cast(f))); - if (f[dim] == Scalar(n_[dim] - 1) && x[dim] == hi_[dim]) + // If exactly at the top boundary, shift into the last valid cell so + // that (b+1) remains in bounds. + if (f == Scalar(nd - 1) && x[d] == m_hi[d]) { - --bin[dim]; + --b; } - dloc[dim] = f[dim] - Scalar(bin[dim]); + assert(b >= 0); + assert(b < static_cast(nd) - 1); - assert(bin[dim] >= 0); - assert(bin[dim] < (int)n_[dim] - 1); + bin[d] = b; + frac[d] = f - Scalar(b); } - const Scalar xd = dloc[0]; - const Scalar yd = dloc[1]; - const Scalar zd = dloc[2]; - const Scalar wd = dloc[3]; - const Scalar vd = dloc[4]; - + // Load the 2^5=32 corners of the surrounding 5D cell. + Scalar c0[32]; Scalar c[32]; - for (unsigned int mask = 0; mask < 32; ++mask) - { - const unsigned int i0 = (unsigned int)(bin[0] + ((mask >> 0) & 1u)); - const unsigned int i1 = (unsigned int)(bin[1] + ((mask >> 1) & 1u)); - const unsigned int i2 = (unsigned int)(bin[2] + ((mask >> 2) & 1u)); - const unsigned int i3 = (unsigned int)(bin[3] + ((mask >> 3) & 1u)); - const unsigned int i4 = (unsigned int)(bin[4] + ((mask >> 4) & 1u)); - - c[mask] = Scalar(data_[nindex_(i0, i1, i2, i3, i4)]); - } - Scalar c0[16]; - for (unsigned int i = 0; i < 16; ++i) + for (unsigned int mask = 0; mask < 32; ++mask) { - c0[i] = c[2 * i] * (Scalar(1) - xd) + c[2 * i + 1] * xd; + const unsigned int i0 + = static_cast(bin[0] + static_cast((mask >> 0) & 1u)); + const unsigned int i1 + = static_cast(bin[1] + static_cast((mask >> 1) & 1u)); + const unsigned int i2 + = static_cast(bin[2] + static_cast((mask >> 2) & 1u)); + const unsigned int i3 + = static_cast(bin[3] + static_cast((mask >> 3) & 1u)); + const unsigned int i4 + = static_cast(bin[4] + static_cast((mask >> 4) & 1u)); + + // Implicit conversion from T to Scalar is intended. + c0[mask] = m_data[m_indexer(i0, i1, i2, i3, i4)]; } - Scalar c1[8]; - for (unsigned int i = 0; i < 8; ++i) - { - c1[i] = c0[2 * i] * (Scalar(1) - yd) + c0[2 * i + 1] * yd; - } + // For each dimension d, collapse pairs of points that differ in bit d. + Scalar* in = c0; + Scalar* out = c; + unsigned int len = 32; - Scalar c2[4]; - for (unsigned int i = 0; i < 4; ++i) + for (int d = 0; d < 5; ++d) { - c2[i] = c1[2 * i] * (Scalar(1) - zd) + c1[2 * i + 1] * zd; - } + const Scalar t = frac[d]; + const Scalar omt = Scalar(1) - t; + const unsigned int out_len = len / 2; - Scalar c3[2]; - for (unsigned int i = 0; i < 2; ++i) - { - c3[i] = c2[2 * i] * (Scalar(1) - wd) + c2[2 * i + 1] * wd; + for (unsigned int i = 0; i < out_len; ++i) + { + out[i] = in[2 * i] * omt + in[2 * i + 1] * t; + } + // Swap input/output + Scalar* tmp = in; + in = out; + out = tmp; + len = out_len; } - return c3[0] * (Scalar(1) - vd) + c3[1] * vd; + // After 5 reductions, len==1 and in[0] holds the interpolated value. + return in[0]; } private: - const T* data_; - unsigned int n_[5]; - Scalar lo_[5]; - Scalar hi_[5]; - Scalar dx_[5]; - FiveDimensionalIndex nindex_; + const T* m_data; + Scalar m_lo[5]; + Scalar m_hi[5]; + Scalar m_dx[5]; + FiveDimensionalIndex m_indexer; }; } // namespace azplugins } // namespace hoomd +#undef AZPLUGINS_HOSTDEVICE +#undef AZPLUGINS_FORCEINLINE + #endif // AZPLUGINS_LINEAR_INTERPOLATOR_5D_H_ diff --git a/src/pair.py b/src/pair.py index 8e982ae..56319f1 100644 --- a/src/pair.py +++ b/src/pair.py @@ -16,40 +16,32 @@ class ChebyshevAnisotropicPairPotential(Force): """Chebyshev anisotropic pair potential.""" - def __init__(self, nlist, domain, terms, coeffs, r0_data, r_cut=3.0): + def __init__(self, nlist, domain, terms, coeffs, r0, r_cut): super().__init__() + self._nlist = nlist - self._domain = numpy.asarray(domain) - self._r_cut = float(r_cut) + param_dict = ParameterDict(r_cut=float) + param_dict["r_cut"] = float(r_cut) + self._param_dict.update(param_dict) + + self._domain = numpy.asarray(domain, dtype=numpy.float64) self._terms = numpy.asarray(terms, dtype=numpy.uint32) - self._coeffs = numpy.asarray(coeffs) - self._r0_data = numpy.asarray(r0_data) + self._coeffs = numpy.asarray(coeffs, dtype=numpy.float64) + + self.r0 = numpy.asarray(r0, dtype=numpy.float64) if self._domain.shape != (5, 2): - raise ValueError("domain must have shape (5, 2).") + raise ValueError("domain must have shape (5,2).") if self._terms.ndim != 2 or self._terms.shape[1] != 6: - raise ValueError("terms must have shape (Nterms, 6).") - nterms = int(self._terms.shape[0]) - if self._coeffs.ndim != 1 or int(self._coeffs.shape[0]) != nterms: - raise ValueError("coeffs must have shape (Nterms,).") - if self._r0_data.ndim != 5: - raise ValueError("r0_data must be a 5D array.") - - @property - def r_cut(self): - """Cut-off distance in approximation domain""" - return self._r_cut + raise ValueError("terms must have shape (Nterms,6).") - @property - def n_terms(self): - """Number of terms.""" - return int(self._terms.shape[0]) + n_terms = int(self._terms.shape[0]) + if self._coeffs.ndim != 1 or int(self._coeffs.shape[0]) != n_terms: + raise ValueError("coeffs must have shape (Nterms,).") - @property - def r0_shape(self): - """r0 table shape.""" - return tuple(int(x) for x in self._r0_data.shape) + if self.r0.ndim != 5: + raise ValueError("r0 must be a 5D array.") def _attach_hook(self): self._nlist._attach(self._simulation) @@ -58,10 +50,10 @@ def _attach_hook(self): self._simulation.state._cpp_sys_def, self._nlist._cpp_obj, self._domain, - self._r_cut, + self.r_cut, self._terms, self._coeffs, - self._r0_data, + self.r0, ) super()._attach_hook() diff --git a/src/pytest/test_chebyshev.py b/src/pytest/test_chebyshev.py index c0b16cc..30a3f51 100644 --- a/src/pytest/test_chebyshev.py +++ b/src/pytest/test_chebyshev.py @@ -7,33 +7,11 @@ import hoomd.azplugins -def test_attach_and_zero_force(simulation_factory, two_particle_snapshot_factory): +def test_chebyshev_construct_attach_zero( + simulation_factory, two_particle_snapshot_factory +): """Construct, attach, and check force/torque output.""" - # Construct the Python object - nlist = hoomd.md.nlist.Cell(buffer=0.4) - - domain = numpy.zeros((5, 2), dtype=numpy.float64) - terms = numpy.zeros((2, 6), dtype=numpy.uint32) - coeffs = numpy.zeros((2,), dtype=numpy.float64) - r0_data = numpy.zeros((2, 2, 2, 2, 2), dtype=numpy.float64) - r_cut = 3.0 - - pot = hoomd.azplugins.pair.ChebyshevAnisotropicPairPotential( - nlist=nlist, - domain=domain, - r_cut=r_cut, - terms=terms, - coeffs=coeffs, - r0_data=r0_data, - ) - - # Pre-attach checks - assert numpy.isclose(pot.r_cut, r_cut) - assert pot.n_terms == 2 - assert pot.r0_shape == (2, 2, 2, 2, 2) - - # Attach via a 0-step simulation snap = two_particle_snapshot_factory() if snap.communicator.rank == 0: snap.particles.position[:] = [[-0.5, 0.0, 0.0], [0.5, 0.0, 0.0]] @@ -46,25 +24,57 @@ def test_attach_and_zero_force(simulation_factory, two_particle_snapshot_factory nve = hoomd.md.methods.ConstantVolume(hoomd.filter.All()) integrator.methods = [nve] + nlist = hoomd.md.nlist.Cell(buffer=0.4) + + domain = numpy.asarray( + [ + [0.0, 2.0 * numpy.pi], # theta + [0.0, numpy.pi], # phi + [0.0, 2.0 * numpy.pi], # alpha + [0.0, numpy.pi], # beta + [0.0, 2.0 * numpy.pi], # gamma + ], + dtype=numpy.float64, + ) + + terms = numpy.asarray( + [ + [0, 0, 0, 0, 0, 0], + [1, 0, 2, 0, 1, 3], + ], + dtype=numpy.uint32, + ) + + coeffs = numpy.asarray([1.0, -0.25], dtype=numpy.float64) + + # r0 must be 5D (and each dimension >= 2) + r0 = (numpy.arange(32, dtype=numpy.float64).reshape((2, 2, 2, 2, 2))) * 0.01 + + r_cut = 3.0 + + pot = hoomd.azplugins.pair.ChebyshevAnisotropicPairPotential( + nlist=nlist, domain=domain, terms=terms, coeffs=coeffs, r0=r0, r_cut=r_cut + ) + + assert numpy.isclose(pot.r_cut, r_cut) + assert isinstance(pot.r0, numpy.ndarray) + assert pot.r0.ndim == 5 + assert pot.r0.shape == (2, 2, 2, 2, 2) + integrator.forces = [pot] sim.operations.integrator = integrator - # Attach all objects + # attach sim.run(0) - # After attach + # check if attach happened assert hasattr(pot, "_cpp_obj") assert pot._cpp_obj is not None - # Post-attach checks + # recheck key properties after attach assert numpy.isclose(pot.r_cut, r_cut) - assert pot.n_terms == 2 - assert pot.r0_shape == (2, 2, 2, 2, 2) - - assert pot._cpp_obj.n_terms == 2 - assert numpy.isclose(pot._cpp_obj.r_cut, r_cut) + assert pot.r0.shape == (2, 2, 2, 2, 2) - # Check Force/torque/energy outputs if sim.device.communicator.rank == 0: numpy.testing.assert_array_equal(pot.forces, numpy.zeros((2, 3))) numpy.testing.assert_array_equal(pot.torques, numpy.zeros((2, 3)))