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..bc0b04c --- /dev/null +++ b/src/ChebyshevAnisotropicPairPotential.cc @@ -0,0 +1,435 @@ +// 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" + +#include "hoomd/BoxDim.h" +#include "hoomd/VectorMath.h" + +#include +#include +#include +#include + +namespace hoomd + { +namespace azplugins + { + +ChebyshevAnisotropicPairPotential::ChebyshevAnisotropicPairPotential( + std::shared_ptr sysdef, + std::shared_ptr nlist, + const Scalar* domain, + const Scalar 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) + { + { + 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) + { + // 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) + { + lo[d] = h_domain.data[d].x; + hi[d] = h_domain.data[d].y; + } + + 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) + { + // 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); + } + } + +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, + Scalar 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("num_terms", &ChebyshevAnisotropicPairPotential::getNTerms); + } + + } // end namespace detail + } // namespace azplugins + } // namespace hoomd diff --git a/src/ChebyshevAnisotropicPairPotential.h b/src/ChebyshevAnisotropicPairPotential.h new file mode 100644 index 0000000..7fc549b --- /dev/null +++ b/src/ChebyshevAnisotropicPairPotential.h @@ -0,0 +1,98 @@ +// 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 PYBIND11_EXPORT ChebyshevAnisotropicPairPotential : public ForceCompute + { + public: + //! Constructor + ChebyshevAnisotropicPairPotential(std::shared_ptr sysdef, + std::shared_ptr nlist, + const Scalar* domain, + const Scalar r_cut, + const unsigned int* terms, + const Scalar* coeffs, + unsigned int Nterms, + const Scalar* r0_data, + const unsigned int* r0_shape); + + //! Destructor + virtual ~ChebyshevAnisotropicPairPotential(); + + // Getters + std::shared_ptr getNeighborList() const + { + return m_nlist; + } + + /// 5x2 domain: stored as 5 entries of Scalar2 = (min,max) + const GPUArray& getApproximationDomain() const + { + return m_domain; + } + + /// Read-only cutoff radius + Scalar 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; + + std::shared_ptr m_nlist; //!< Neighbor list + + GPUArray m_domain; //!< Approximation domain (5x2): 5 rows, each is (min, max) + + Scalar m_r_cut; //!< cut-off distance in approximation domain + + GPUArray m_terms; //!< Chebyshev term list (Nterms x 6) + + GPUArray m_coeffs; //!< Coefficients corresponding to each term + + unsigned int m_Nterms; //!< Number of terms + + GPUArray m_r0_data; //!< R0 data + + GPUArray m_r0_shape; //!< Number of points used along each dimension to sample r0 + }; + + } // end namespace azplugins + } // end namespace hoomd + +#endif // AZPLUGINS_CHEBYSHEV_ANISOTROPIC_PAIR_POTENTIAL_H_ diff --git a/src/LinearInterpolator5D.h b/src/LinearInterpolator5D.h new file mode 100644 index 0000000..7340434 --- /dev/null +++ b/src/LinearInterpolator5D.h @@ -0,0 +1,202 @@ +// 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() : 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]} + { + } + + AZPLUGINS_HOSTDEVICE AZPLUGINS_FORCEINLINE FiveDimensionalIndex(unsigned int n0, + unsigned int n1, + unsigned int n2, + unsigned int n3, + unsigned int n4) + : m_n {n0, n1, n2, n3, 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 * 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 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 m_n[5]; + }; + +/*! \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() : m_data(nullptr), m_indexer() + { + for (int d = 0; d < 5; ++d) + { + 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) + : m_data(data), m_indexer(n) + { + for (int d = 0; d < 5; ++d) + { + const unsigned int nd = n[d]; + assert(nd >= 2); + + m_lo[d] = lo[d]; + m_hi[d] = hi[d]; + m_dx[d] = (m_hi[d] - m_lo[d]) / Scalar(nd - 1); + } + + assert(m_indexer.size() > 0); + } + + //! 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}; + + // Compute the cell bin and fractional coordinate in each dimension. + int bin[5]; + Scalar frac[5]; + + for (int d = 0; d < 5; ++d) + { + 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 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]) + { + --b; + } + + assert(b >= 0); + assert(b < static_cast(nd) - 1); + + bin[d] = b; + frac[d] = f - Scalar(b); + } + + // 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 + = 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)]; + } + + // For each dimension d, collapse pairs of points that differ in bit d. + Scalar* in = c0; + Scalar* out = c; + unsigned int len = 32; + + for (int d = 0; d < 5; ++d) + { + const Scalar t = frac[d]; + const Scalar omt = Scalar(1) - t; + const unsigned int out_len = len / 2; + + 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; + } + + // After 5 reductions, len==1 and in[0] holds the interpolated value. + return in[0]; + } + + private: + 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/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..56319f1 100644 --- a/src/pair.py +++ b/src/pair.py @@ -4,13 +4,61 @@ """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, r_cut): + super().__init__() + + self._nlist = nlist + + 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, 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).") + if self._terms.ndim != 2 or self._terms.shape[1] != 6: + raise ValueError("terms must have shape (Nterms,6).") + + 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,).") + + if self.r0.ndim != 5: + raise ValueError("r0 must be a 5D array.") + + 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, + ) + + 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..30a3f51 --- /dev/null +++ b/src/pytest/test_chebyshev.py @@ -0,0 +1,81 @@ +# 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_chebyshev_construct_attach_zero( + simulation_factory, two_particle_snapshot_factory +): + """Construct, attach, and check force/torque output.""" + + 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] + + 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 + sim.run(0) + + # check if attach happened + assert hasattr(pot, "_cpp_obj") + assert pot._cpp_obj is not None + + # recheck key properties after attach + assert numpy.isclose(pot.r_cut, r_cut) + assert pot.r0.shape == (2, 2, 2, 2, 2) + + 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,)))