diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index a5529d9..c2c4499 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -65,6 +65,12 @@ set(_aniso_pair_evaluators TwoPatchMorse ) +set(_chebyshev_symmetries + Null + Cube + Tetrahedron + ) + # process bond potentials foreach(_evaluator ${_bond_evaluators}) configure_file(export_PotentialBond.cc.inc @@ -160,6 +166,23 @@ foreach(_evaluator ${_aniso_pair_evaluators}) endif() endforeach() +# process Chebyshev anisotropic pair potential +foreach(_symmetry ${_chebyshev_symmetries}) + configure_file(export_ChebyshevAnisotropicPairPotential.cc.inc + export_ChebyshevAnisotropicPairPotential${_symmetry}.cc + @ONLY) + set(_${COMPONENT_NAME}_sources ${_${COMPONENT_NAME}_sources} + export_ChebyshevAnisotropicPairPotential${_symmetry}.cc) + + if (ENABLE_HIP) + configure_file(export_ChebyshevAnisotropicPairPotentialGPU.cc.inc + export_ChebyshevAnisotropicPairPotential${_symmetry}GPU.cc + @ONLY) + set(_${COMPONENT_NAME}_sources ${_${COMPONENT_NAME}_sources} + export_ChebyshevAnisotropicPairPotential${_symmetry}GPU.cc) + endif() +endforeach() + # process velocity field geometries set(_binning_geometries Cartesian diff --git a/src/ChebyshevAnisotropicPairPotential.h b/src/ChebyshevAnisotropicPairPotential.h new file mode 100644 index 0000000..fbf69eb --- /dev/null +++ b/src/ChebyshevAnisotropicPairPotential.h @@ -0,0 +1,743 @@ +// 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 Templated class for the Chebyshev anisotropic pair potential. + * + * The class is templated on a \c ShapeSymmetryT parameter that provides + * a symmetry reduction of the angular coordinates. + */ + +#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 +#include + +#include "hoomd/BoxDim.h" +#include "hoomd/ForceCompute.h" +#include "hoomd/GPUArray.h" +#include "hoomd/HOOMDMath.h" +#include "hoomd/Index1D.h" +#include "hoomd/VectorMath.h" +#include "hoomd/md/NeighborList.h" + +#include "LinearInterpolator5D.h" +#include "ShapeSymmetry.h" + +namespace hoomd + { +namespace azplugins + { + +//! Scale a coordinate from [lo, hi] to the Chebyshev domain [-1, 1]. +inline Scalar chebScale(Scalar x, Scalar lo, Scalar hi) + { + return (Scalar(2) * (x - lo) / (hi - lo)) - Scalar(1); + } + +//! Evaluate Chebyshev polynomials of the first kind and their derivatives +//! from degree 0 up to max_deg, using the three-term recurrence relation. +/*! + T_0(x) = 1 T'_0(x) = 0 + T_1(x) = x T'_1(x) = 1 + T_{n+1}(x) = 2x T_n - T_{n-1} T'_{n+1}(x) = 2 T_n + 2x T'_n - T'_{n-1} + + + \param x Evaluation point in [-1, 1] + \param max_deg Highest polynomial degree to compute + \param T Output: T[n] = T_n(x) for n = 0 .. max_deg (size >= max_deg+1) + \param dT Output: dT[n] = T'_n(x) for n = 0 .. max_deg (size >= max_deg+1) +*/ +inline void chebEvaluate(Scalar x, unsigned int max_deg, Scalar* T, Scalar* dT) + { + T[0] = Scalar(1); + dT[0] = Scalar(0); + + if (max_deg == 0) + return; + + T[1] = x; + dT[1] = Scalar(1); + + const Scalar two_x = Scalar(2) * x; + for (unsigned int n = 1; n < max_deg; ++n) + { + T[n + 1] = two_x * T[n] - T[n - 1]; + dT[n + 1] = Scalar(2) * T[n] + two_x * dT[n] - dT[n - 1]; + } + } + +//! Chebyshev anisotropic pair potential, templated on a symmetry reducer. +/*! + * \tparam ShapeSymmetryT A class providing a static \c domain_upper[5] array + * and a static \c reduce(theta, phi, alpha, beta, gamma) method that + * maps the angles into a fundamental domain and returns the applied + * rotation as a quaternion. See \c ShapeSymmetry.h. + */ +template +class PYBIND11_EXPORT ChebyshevAnisotropicPairPotential : public ForceCompute + { + public: + static constexpr unsigned int num_coordinates = 6; + static constexpr unsigned int num_angle_coordinates = num_coordinates - 1; + + ChebyshevAnisotropicPairPotential(std::shared_ptr sysdef, + std::shared_ptr nlist, + 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(); + + //! Detach from the neighbor list (called when removing from simulation) + virtual void notifyDetach(); + + // Getters + std::shared_ptr getNeighborList() const + { + return m_nlist; + } + + /// 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: + // member variables + + std::shared_ptr m_nlist; //!< Neighbor list + + Scalar m_r_cut; //!< Cut-off distance in approximation domain + Scalar m_nlist_r_cut; //!< Neighbor-list cutoff = ceil(max(r0) + r_cut) + + std::shared_ptr> m_r_cut_nlist; //!< r_cut matrix shared with nlist + bool m_attached = true; //!< Whether attached to the simulation + + GPUArray m_terms; //!< Chebyshev term list (Nterms x num_coordinates) + 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; //!< Points per dimension to sample r0 + + std::array + m_max_deg; //!< Maximum Chebyshev degree per coordinate + std::array + m_cheb_scale; //!< Chain-rule scale factors for each coordinate + unsigned int m_max_deg_global; //!< Maximum Chebyshev degree over all coordinates + Index2D m_cheb_idx; //!< Indexer for flat Chebyshev scratch storage + std::vector m_cheb_T_flat; //!< Chebyshev polynomial scratch storage + std::vector m_cheb_dT_flat; //!< Chebyshev derivative scratch storage + + void initializeChebyshevData(); + void computeForces(uint64_t timestep) override; + }; + +// Constructor +template +ChebyshevAnisotropicPairPotential::ChebyshevAnisotropicPairPotential( + std::shared_ptr sysdef, + std::shared_ptr nlist, + 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), m_max_deg_global(0), + m_cheb_idx() + { + // terms: shape (Nterms, num_coordinates), stored flat + { + GPUArray terms_arr(static_cast(Nterms) * num_coordinates, + m_exec_conf); + m_terms.swap(terms_arr); + + ArrayHandle h_terms(m_terms, access_location::host, access_mode::overwrite); + std::copy(terms, terms + static_cast(Nterms) * num_coordinates, 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::overwrite); + std::copy(coeffs, coeffs + Nterms, h_coeffs.data); + } + + // r0_shape: length num_angle_coordinates + { + GPUArray shape_arr(num_angle_coordinates, m_exec_conf); + m_r0_shape.swap(shape_arr); + + ArrayHandle h_shape(m_r0_shape, + access_location::host, + access_mode::overwrite); + std::copy(r0_shape, r0_shape + num_angle_coordinates, h_shape.data); + } + + // r0_data: flat array, length = product(r0_shape) + unsigned int n_r0 = 1; + for (unsigned int d = 0; d < num_angle_coordinates; ++d) + { + n_r0 *= 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::overwrite); + std::copy(r0_data, r0_data + n_r0, h_r0.data); + } + + // neighbor list subscriber + Scalar max_r0 = *std::max_element(r0_data, r0_data + n_r0); + m_nlist_r_cut = std::ceil(max_r0 + m_r_cut); + + const Index2D typpair_idx(m_pdata->getNTypes()); + m_r_cut_nlist = std::make_shared>(typpair_idx.getNumElements(), m_exec_conf); + { + ArrayHandle h_r_cut_nlist(*m_r_cut_nlist, + access_location::host, + access_mode::overwrite); + std::fill(h_r_cut_nlist.data, + h_r_cut_nlist.data + typpair_idx.getNumElements(), + m_nlist_r_cut); + } + m_nlist->addRCutMatrix(m_r_cut_nlist); + m_nlist->notifyRCutMatrixChange(); + + initializeChebyshevData(); + } + +// Destructor +template +ChebyshevAnisotropicPairPotential::~ChebyshevAnisotropicPairPotential() + { + if (m_attached) + { + m_nlist->removeRCutMatrix(m_r_cut_nlist); + } + } + +// notifyDetach +template +void ChebyshevAnisotropicPairPotential::notifyDetach() + { + if (m_attached) + { + m_nlist->removeRCutMatrix(m_r_cut_nlist); + } + m_attached = false; + } + +// initializeChebyshevData +template +void ChebyshevAnisotropicPairPotential::initializeChebyshevData() + { + m_max_deg.fill(0); + + ArrayHandle h_terms(m_terms, access_location::host, access_mode::read); + for (unsigned int t = 0; t < m_Nterms; ++t) + { + for (unsigned int c = 0; c < num_coordinates; ++c) + { + const unsigned int deg = h_terms.data[t * num_coordinates + c]; + if (deg > m_max_deg[c]) + { + m_max_deg[c] = deg; + } + } + } + + m_cheb_scale[0] = Scalar(2); + for (unsigned int d = 0; d < num_angle_coordinates; ++d) + { + m_cheb_scale[d + 1] + = Scalar(2) / (ShapeSymmetryT::domain_upper[d] - ShapeSymmetryT::domain_lower[d]); + } + + m_max_deg_global = 0; + for (unsigned int c = 0; c < num_coordinates; ++c) + { + if (m_max_deg[c] > m_max_deg_global) + { + m_max_deg_global = m_max_deg[c]; + } + } + + m_cheb_idx = Index2D(m_max_deg_global + 1, num_coordinates); + m_cheb_T_flat.resize(m_cheb_idx.getNumElements()); + m_cheb_dT_flat.resize(m_cheb_idx.getNumElements()); + } + +// computeForces +template +void ChebyshevAnisotropicPairPotential::computeForces(uint64_t timestep) + { + // start by updating the neighborlist + m_nlist->compute(timestep); + + // check neighbor list storage mode + const bool use_third_law = (m_nlist->getStorageMode() == hoomd::md::NeighborList::half); + const unsigned int N_local = m_pdata->getN(); + // 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_tag(m_pdata->getTags(), 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); + ArrayHandle h_terms(m_terms, access_location::host, access_mode::read); + ArrayHandle h_coeffs(m_coeffs, access_location::host, access_mode::read); + + const BoxDim box = m_pdata->getGlobalBox(); + const Scalar nlist_rcutsq = m_nlist_r_cut * m_nlist_r_cut; + const Scalar fd_step = Scalar(1.0e-6); + + LinearInterpolator5D interp(h_r0_data.data, + h_r0_shape.data, + ShapeSymmetryT::domain_lower, + ShapeSymmetryT::domain_upper); + + 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(); + + //! Euler-angle singularity tolerance for the alpha/gamma extraction. + const Scalar euler_singularity_tol = Scalar(1e-7); + + const Scalar phi_eval_min = ShapeSymmetryT::domain_lower[1]; + const Scalar phi_eval_max = ShapeSymmetryT::domain_upper[1]; + const Scalar beta_eval_min = ShapeSymmetryT::domain_lower[3]; + const Scalar beta_eval_max = ShapeSymmetryT::domain_upper[3]; + + for (unsigned int i = 0; i < N; ++i) + { + // Per-pair position and orientation are loaded inside the loop after + // sorting by tag (on full lists) + + // Initialize particle force, torque, and energy + Scalar3 fi = make_scalar3(0, 0, 0); + Scalar3 ti = make_scalar3(0, 0, 0); + Scalar pei = Scalar(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()); + + // Sort the pair by tag + unsigned int eval_a = i; + unsigned int eval_b = j; + bool i_is_eval_a = true; + { + const unsigned int tag_i = h_tag.data[i]; + const unsigned int tag_j = h_tag.data[j]; + if (tag_j < tag_i) + { + eval_a = j; + eval_b = i; + i_is_eval_a = false; + } + } + + const Scalar3 pos_a + = make_scalar3(h_pos.data[eval_a].x, h_pos.data[eval_a].y, h_pos.data[eval_a].z); + const Scalar3 pos_b + = make_scalar3(h_pos.data[eval_b].x, h_pos.data[eval_b].y, h_pos.data[eval_b].z); + Scalar3 dx = pos_a - pos_b; + // Apply periodic boundary conditions + dx = box.minImage(dx); + // Neighbor-list cutoff check (center-center distance). + const Scalar rsq = dot(dx, dx); + if (rsq > nlist_rcutsq) + { + continue; + } + + const quat q_a(h_orientation.data[eval_a]); + const quat q_b(h_orientation.data[eval_b]); + const quat q_a_conj = conj(q_a); + // dx is in lab frame, so rotate dx by conj(q_a) + const vec3 dx_body = rotate(q_a_conj, vec3(dx)); + // Relative orientation of eval_b with respect to eval_a: + // q_rel = conj(q_a) * q_b + // ref: + // https://www.mathworks.com/help/fusion/ug/rotations-orientation-and-quaternions.html + const quat q_rel = q_a_conj * q_b; + + // Convert position to spherical coordinates + // Skip overlapping particles + const Scalar r = fast::sqrt(dot(dx_body, dx_body)); + if (r < Scalar(1e-12)) + { + continue; + } + + Scalar 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); + Scalar phi = slow::acos(cosphi); + + // Build rotation matrix from the relative quaternion and extract + // ZXZ Euler angles (alpha, beta, gamma) + const rotmat3 R(q_rel); + + Scalar alpha = Scalar(0); + Scalar beta = Scalar(0); + Scalar gamma = Scalar(0); + + Scalar clamped_r22 = R.row2.z; + if (clamped_r22 < Scalar(-1)) + clamped_r22 = Scalar(-1); + else if (clamped_r22 > Scalar(1)) + clamped_r22 = Scalar(1); + beta = slow::acos(clamped_r22); + + if (beta > euler_singularity_tol && beta < Scalar(M_PI) - euler_singularity_tol) + { + alpha = std::atan2(R.row0.z, -R.row1.z); + gamma = std::atan2(R.row2.x, R.row2.y); + if (alpha < Scalar(0)) + alpha += Scalar(2) * M_PI; + } + else + { + alpha = Scalar(0); + gamma + = std::atan2((beta <= euler_singularity_tol) ? R.row0.y : -R.row0.y, R.row0.x); + } + + if (gamma < Scalar(0)) + gamma += Scalar(2) * M_PI; + + // Symmetry reduction + // The symmetry evaluator maps (theta, phi, alpha, beta, gamma) + // into the reduced domain and returns the cumulative + // quaternion rotation that was applied. We keep the transformation so + // that forces and torques can be rotated back to the + // original frame at the end. + const quat sym_transformation + = ShapeSymmetryT::reduce(theta, phi, alpha, beta, gamma); + + // Move phi and beta away from 0 and pi to avoid 1/sin(beta or phi) + // singularity in the Jacobian (used the same threshold as beta). + if (beta < beta_eval_min) + beta = beta_eval_min; + else if (beta > beta_eval_max) + beta = beta_eval_max; + + if (phi < phi_eval_min) + phi = phi_eval_min; + else if (phi > phi_eval_max) + phi = phi_eval_max; + + // Compute r0 and all 5 derivatives + Scalar r0; + Scalar dr0[num_angle_coordinates]; + interp.valueAndDerivatives(theta, phi, alpha, beta, gamma, fd_step, r0, dr0); + const Scalar dr0_dtheta = dr0[0]; + const Scalar dr0_dphi = dr0[1]; + const Scalar dr0_dalpha = dr0[2]; + const Scalar dr0_dbeta = dr0[3]; + const Scalar dr0_dgamma = dr0[4]; + + // Compute rho + const Scalar inv_r = Scalar(1) / r; + const Scalar inv_r0 = Scalar(1) / r0; + const Scalar inv_r0_rcut = Scalar(1) / (r0 + m_r_cut); + const Scalar rho_denom = inv_r0_rcut - inv_r0; + const Scalar rho_num = inv_r - inv_r0; + Scalar rho = rho_num / rho_denom; + + if (rho > Scalar(1)) + continue; + + // Save raw rho for energy extrapolation if rho < 0 + const Scalar rho_energy = rho; + if (rho < Scalar(0)) + rho = Scalar(0); + + // Compute drho/dr and drho/dr0 + const Scalar inv_r_sq = inv_r * inv_r; + const Scalar inv_r0_sq = inv_r0 * inv_r0; + const Scalar inv_r0_rcut_sq = inv_r0_rcut * inv_r0_rcut; + const Scalar rho_denom_sq = rho_denom * rho_denom; + + const Scalar drho_dr = -inv_r_sq / rho_denom; + const Scalar drho_dr0 + = (inv_r0_sq * rho_denom - rho_num * (inv_r0_sq - inv_r0_rcut_sq)) / rho_denom_sq; + + // Chebyshev evaluation: scale each coordinate to [-1,1] + // and evaluate polynomials + derivatives up to max degree. + chebEvaluate(chebScale(rho, Scalar(0), Scalar(1)), + m_max_deg[0], + m_cheb_T_flat.data() + m_cheb_idx(0, 0), + m_cheb_dT_flat.data() + m_cheb_idx(0, 0)); + + const Scalar ang_coords[num_angle_coordinates] = {theta, phi, alpha, beta, gamma}; + for (unsigned int c = 0; c < 5; ++c) + { + chebEvaluate(chebScale(ang_coords[c], + ShapeSymmetryT::domain_lower[c], + ShapeSymmetryT::domain_upper[c]), + m_max_deg[c + 1], + m_cheb_T_flat.data() + m_cheb_idx(0, c + 1), + m_cheb_dT_flat.data() + m_cheb_idx(0, c + 1)); + } + + // Evaluate u and du/d(coord_k) + Scalar u = Scalar(0); + Scalar du[num_coordinates] + = {Scalar(0), Scalar(0), Scalar(0), Scalar(0), Scalar(0), Scalar(0)}; + + for (unsigned int t = 0; t < m_Nterms; ++t) + { + const unsigned int* degs = h_terms.data + num_coordinates * t; + const Scalar coeff = h_coeffs.data[t]; + + Scalar T_vals[num_coordinates]; + Scalar dT_vals[num_coordinates]; + for (unsigned int c = 0; c < num_coordinates; ++c) + { + T_vals[c] = m_cheb_T_flat[m_cheb_idx(degs[c], c)]; + dT_vals[c] = m_cheb_dT_flat[m_cheb_idx(degs[c], c)]; + } + + Scalar prefix[num_coordinates + 1]; + prefix[0] = Scalar(1); + for (unsigned int c = 0; c < num_coordinates; ++c) + prefix[c + 1] = prefix[c] * T_vals[c]; + + Scalar suffix[num_coordinates + 1]; + suffix[num_coordinates] = Scalar(1); + for (int c = static_cast(num_coordinates) - 1; c >= 0; --c) + suffix[c] = suffix[c + 1] * T_vals[c]; + + u += coeff * prefix[num_coordinates]; + + for (unsigned int c = 0; c < num_coordinates; ++c) + du[c] += coeff * dT_vals[c] * m_cheb_scale[c] * prefix[c] * suffix[c + 1]; + } + + // Linear extrapolation for energy when rho < 0 + const Scalar u_energy = (rho_energy < Scalar(0)) ? (u + rho_energy * du[0]) : u; + + // Jacobian matrix J (6x6). + // J maps the potential-derivative vector + // [du/drho, du/dtheta, du/dphi, du/dalpha, du/dbeta, du/dgamma] + // to the lab-frame force and torque: + // [F_x, F_y, F_z, tau_x, tau_y, tau_z] + Scalar s_th, c_th; + fast::sincos(theta, s_th, c_th); + Scalar s_ph, c_ph; + fast::sincos(phi, s_ph, c_ph); + Scalar s_b, c_b; + fast::sincos(beta, s_b, c_b); + Scalar s_a, c_a; + fast::sincos(alpha, s_a, c_a); + + const Scalar inv_r_s_ph = inv_r / s_ph; + const Scalar inv_s_b = Scalar(1) / s_b; + + const Scalar A = drho_dr0 * dr0_dtheta * inv_r_s_ph; + const Scalar B = drho_dr0 * dr0_dphi * inv_r; + const Scalar C = drho_dr0 * dr0_dalpha * inv_s_b; + const Scalar D = drho_dr0 * dr0_dgamma * inv_s_b; + + const Scalar f_x_red = (-c_th * s_ph * drho_dr + s_th * A - c_th * c_ph * B) * du[0] + + (s_th * inv_r_s_ph) * du[1] + (-c_th * c_ph * inv_r) * du[2]; + + const Scalar f_y_red = (-s_th * s_ph * drho_dr - c_th * A - s_th * c_ph * B) * du[0] + + (-c_th * inv_r_s_ph) * du[1] + (-s_th * c_ph * inv_r) * du[2]; + + const Scalar f_z_red = (-c_ph * drho_dr + s_ph * B) * du[0] + (s_ph * inv_r) * du[2]; + + const Scalar tau_x_red = (c_b * s_a * C - c_a * drho_dr0 * dr0_dbeta - s_a * D) * du[0] + + (c_b * s_a * inv_s_b) * du[3] + (-c_a) * du[4] + + (-s_a * inv_s_b) * du[5]; + + const Scalar tau_y_red = (-c_b * c_a * C - s_a * drho_dr0 * dr0_dbeta + c_a * D) * du[0] + + (-c_a * c_b * inv_s_b) * du[3] + (-s_a) * du[4] + + (c_a * inv_s_b) * du[5]; + + const Scalar tau_z_red = (-drho_dr0 * dr0_dalpha) * du[0] + (-Scalar(1)) * du[3]; + + // Rotate back to original frame + const quat sym_inv = conj(sym_transformation); + const vec3 f_red(f_x_red, f_y_red, f_z_red); + const vec3 tau_red(tau_x_red, tau_y_red, tau_z_red); + const vec3 f_lab = rotate(sym_inv, f_red); + const vec3 tau_lab = rotate(sym_inv, tau_red); + + const Scalar f_x = f_lab.x; + const Scalar f_y = f_lab.y; + const Scalar f_z = f_lab.z; + const Scalar tau_x = tau_lab.x; + const Scalar tau_y = tau_lab.y; + const Scalar tau_z = tau_lab.z; + + // Writeback + const Scalar sign = i_is_eval_a ? Scalar(1) : Scalar(-1); + + fi.x += sign * f_x; + fi.y += sign * f_y; + fi.z += sign * f_z; + + ti.x += sign * tau_x; + ti.y += sign * tau_y; + ti.z += sign * tau_z; + + pei += u_energy; + + // Newton's third law for half neighbor list + if (use_third_law && j < N_local) + { + h_force.data[j].x -= sign * f_x; + h_force.data[j].y -= sign * f_y; + h_force.data[j].z -= sign * f_z; + h_force.data[j].w += Scalar(0.5) * u_energy; + + h_torque.data[j].x -= sign * tau_x; + h_torque.data[j].y -= sign * tau_y; + h_torque.data[j].z -= sign * tau_z; + } + } + + 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.5) * pei; + + 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 + { + +//! Export one subclass of ChebyshevAnisotropicPairPotential to python. +/*! + * \param m pybind11 module. + * \param name Name the class should have in the python module (must be + * unique per symmetry). + * \tparam ShapeSymmetryT Symmetry evaluator type. + */ +template +void export_ChebyshevAnisotropicPairPotential(pybind11::module& m, const std::string& name) + { + namespace py = pybind11; + using NL = hoomd::md::NeighborList; + using Pot = ChebyshevAnisotropicPairPotential; + + py::class_>(m, name.c_str()) + .def(py::init( + [](std::shared_ptr sysdef, + std::shared_ptr nlist, + Scalar r_cut, + py::array_t terms, + py::array_t coeffs, + py::array_t r0_data) + { + // 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, + r_cut, + terms.data(), + coeffs.data(), + Nterms, + r0_data.data(), + r0_shape.data()); + })) + .def_property_readonly("r_cut", &Pot::getRCut) + .def_property_readonly("num_terms", &Pot::getNTerms); + } + + } // end namespace detail + } // namespace azplugins + } // namespace hoomd + +#endif // AZPLUGINS_CHEBYSHEV_ANISOTROPIC_PAIR_POTENTIAL_H_ diff --git a/src/ChebyshevAnisotropicPairPotentialGPU.h b/src/ChebyshevAnisotropicPairPotentialGPU.h new file mode 100644 index 0000000..17b9444 --- /dev/null +++ b/src/ChebyshevAnisotropicPairPotentialGPU.h @@ -0,0 +1,146 @@ +// 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_CHEBYSHEV_ANISOTROPIC_PAIR_POTENTIAL_GPU_H_ +#define AZPLUGINS_CHEBYSHEV_ANISOTROPIC_PAIR_POTENTIAL_GPU_H_ + +#include +#ifdef ENABLE_HIP +#include "ChebyshevAnisotropicPairPotential.h" + +/*!\file ChebyshevAnisotropicPairPotentialGPU.h + \brief Defines a GPU shell for the Chebyshev anisotropic pair potential. + \note This header cannot be compiled by nvcc +*/ + +#ifdef __HIPCC__ +#error This header cannot be compiled by nvcc +#endif + +#include + +namespace hoomd + { +namespace azplugins + { + +//! Chebyshev anisotropic pair potential, templated on a symmetry reducer. +/*! + * \tparam ShapeSymmetryT A class providing a static \c domain_upper[5] and domain_lower[5] array + * and a static \c reduce(theta, phi, alpha, beta, gamma) method that + * maps the angles into a fundamental domain and returns the applied + * rotation as a quaternion. See \c ShapeSymmetry.h. + */ +template +class PYBIND11_EXPORT ChebyshevAnisotropicPairPotentialGPU + : public ChebyshevAnisotropicPairPotential + { + public: + using Base = ChebyshevAnisotropicPairPotential; + + ChebyshevAnisotropicPairPotentialGPU(std::shared_ptr sysdef, + std::shared_ptr nlist, + const Scalar r_cut, + const unsigned int* terms, + const Scalar* coeffs, + unsigned int Nterms, + const Scalar* r0_data, + const unsigned int* r0_shape); + + virtual ~ChebyshevAnisotropicPairPotentialGPU() { } + }; + +template +ChebyshevAnisotropicPairPotentialGPU::ChebyshevAnisotropicPairPotentialGPU( + std::shared_ptr sysdef, + std::shared_ptr nlist, + const Scalar r_cut, + const unsigned int* terms, + const Scalar* coeffs, + unsigned int Nterms, + const Scalar* r0_data, + const unsigned int* r0_shape) + : Base(sysdef, nlist, r_cut, terms, coeffs, Nterms, r0_data, r0_shape) + { + if (!this->m_exec_conf->isCUDAEnabled()) + { + this->m_exec_conf->msg->error() + << "Creating a ChebyshevAnisotropicPairPotentialGPU with no GPU in the " + << "execution configuration" << std::endl; + throw std::runtime_error("Error initializing ChebyshevAnisotropicPairPotentialGPU"); + } + } + +namespace detail + { + +//! Export one GPU subclass of ChebyshevAnisotropicPairPotential to python. +/*! + * \param m pybind11 module. + * \param name Name the class should have in the python module (must be + * unique per symmetry). + * \tparam ShapeSymmetryT Symmetry evaluator type. + */ +template +void export_ChebyshevAnisotropicPairPotentialGPU(pybind11::module& m, const std::string& name) + { + namespace py = pybind11; + using NL = hoomd::md::NeighborList; + using Pot = ChebyshevAnisotropicPairPotentialGPU; + using Base = ChebyshevAnisotropicPairPotential; + + py::class_>(m, name.c_str()) + .def(py::init( + [](std::shared_ptr sysdef, + std::shared_ptr nlist, + Scalar r_cut, + py::array_t terms, + py::array_t coeffs, + py::array_t r0_data) + { + if (terms.ndim() != 2 || terms.shape(1) != Base::num_coordinates) + { + throw std::runtime_error("terms must have shape (Nterms,6)."); + } + + const unsigned int Nterms = static_cast(terms.shape(0)); + + if (coeffs.ndim() != 1 || static_cast(coeffs.shape(0)) != Nterms) + { + throw std::runtime_error("coeffs must have shape (Nterms,)."); + } + + if (r0_data.ndim() != Base::num_angle_coordinates) + { + throw std::runtime_error("r0_data must be a 5D array."); + } + + std::array r0_shape; + for (unsigned int k = 0; k < Base::num_angle_coordinates; ++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, + r_cut, + terms.data(), + coeffs.data(), + Nterms, + r0_data.data(), + r0_shape.data()); + })); + } + + } // end namespace detail + } // end namespace azplugins + } // end namespace hoomd + +#endif // ENABLE_HIP +#endif // AZPLUGINS_CHEBYSHEV_ANISOTROPIC_PAIR_POTENTIAL_GPU_H_ diff --git a/src/LinearInterpolator5D.h b/src/LinearInterpolator5D.h new file mode 100644 index 0000000..887a801 --- /dev/null +++ b/src/LinearInterpolator5D.h @@ -0,0 +1,330 @@ +// 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 Index5D + { + public: + AZPLUGINS_HOSTDEVICE AZPLUGINS_FORCEINLINE Index5D() : m_n {0, 0, 0, 0, 0} { } + + AZPLUGINS_HOSTDEVICE AZPLUGINS_FORCEINLINE explicit Index5D(const unsigned int* n) + : m_n {n[0], n[1], n[2], n[3], n[4]} + { + } + + AZPLUGINS_HOSTDEVICE AZPLUGINS_FORCEINLINE + Index5D(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); + } + + //! Constructor accepting the domain as a Scalar2 array + AZPLUGINS_HOSTDEVICE AZPLUGINS_FORCEINLINE LinearInterpolator5D(const T* data, + const unsigned int* n, + const Scalar2* domain_s2) + : 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] = domain_s2[d].x; + m_hi[d] = domain_s2[d].y; + 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(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 corners[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 + corners[mask] = m_data[m_indexer(i0, i1, i2, i3, i4)]; + } + + // For each dimension d, collapse pairs of points that differ in bit d + Scalar scratch[16]; + Scalar* in = corners; + Scalar* out = scratch; + 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 + std::swap(in, out); + len = out_len; + } + + // After 5 reductions, len==1 and in[0] holds the interpolated value + return in[0]; + } + + //! Compute the finite-difference derivative with respect to a single dimension + /*! Uses central differences when possible, falling back to forward or backward + differences at the domain boundaries. + + \param x0 Coordinate along dimension 0 + \param x1 Coordinate along dimension 1 + \param x2 Coordinate along dimension 2 + \param x3 Coordinate along dimension 3 + \param x4 Coordinate along dimension 4 + \param dim Which dimension (0-4) to differentiate with respect to + \param h Finite-difference step size + */ + AZPLUGINS_HOSTDEVICE AZPLUGINS_FORCEINLINE Scalar + derivative(Scalar x0, Scalar x1, Scalar x2, Scalar x3, Scalar x4, int dim, Scalar h) const + { + Scalar x[5] = {x0, x1, x2, x3, x4}; + const Scalar val = (*this)(x[0], x[1], x[2], x[3], x[4]); + + const bool at_lo = (x[dim] - h < m_lo[dim]); + const bool at_hi = (x[dim] + h > m_hi[dim]); + + if (!at_lo && !at_hi) + { + x[dim] += h; + const Scalar f_plus = (*this)(x[0], x[1], x[2], x[3], x[4]); + x[dim] -= Scalar(2) * h; + const Scalar f_minus = (*this)(x[0], x[1], x[2], x[3], x[4]); + return (f_plus - f_minus) / (Scalar(2) * h); + } + else if (at_lo) + { + // forward difference + x[dim] += h; + const Scalar f_plus = (*this)(x[0], x[1], x[2], x[3], x[4]); + return (f_plus - val) / h; + } + else + { + // backward difference + x[dim] -= h; + const Scalar f_minus = (*this)(x[0], x[1], x[2], x[3], x[4]); + return (val - f_minus) / h; + } + } + + //! Compute the interpolated value and all 5 partial derivatives in one call. + /*! Compute the finite-difference derivative with respect to dimensions. + Uses central differences when possible, falling back to forward or backward + differences at the domain boundaries. + + \param x0 Coordinate along dimension 0 + \param x1 Coordinate along dimension 1 + \param x2 Coordinate along dimension 2 + \param x3 Coordinate along dimension 3 + \param x4 Coordinate along dimension 4 + \param h Finite-difference step size + \param value Interpolated value at (x0..x4) + \param deriv Array of 5 partial derivatives (one per dimension) + */ + AZPLUGINS_HOSTDEVICE AZPLUGINS_FORCEINLINE void valueAndDerivatives(Scalar x0, + Scalar x1, + Scalar x2, + Scalar x3, + Scalar x4, + Scalar h, + Scalar& value, + Scalar* deriv) const + { + Scalar x[5] = {x0, x1, x2, x3, x4}; + value = (*this)(x[0], x[1], x[2], x[3], x[4]); + + for (int dim = 0; dim < 5; ++dim) + { + const Scalar x_orig = x[dim]; + const bool at_lo = (x_orig - h < m_lo[dim]); + const bool at_hi = (x_orig + h > m_hi[dim]); + + if (!at_lo && !at_hi) + { + // central difference + x[dim] = x_orig + h; + const Scalar f_plus = (*this)(x[0], x[1], x[2], x[3], x[4]); + x[dim] = x_orig - h; + const Scalar f_minus = (*this)(x[0], x[1], x[2], x[3], x[4]); + deriv[dim] = (f_plus - f_minus) / (Scalar(2) * h); + } + else if (at_lo) + { + // forward difference + x[dim] = x_orig + h; + const Scalar f_plus = (*this)(x[0], x[1], x[2], x[3], x[4]); + deriv[dim] = (f_plus - value) / h; + } + else + { + // backward difference + x[dim] = x_orig - h; + const Scalar f_minus = (*this)(x[0], x[1], x[2], x[3], x[4]); + deriv[dim] = (value - f_minus) / h; + } + x[dim] = x_orig; + } + } + //! Return the lower bound for a given dimension. + AZPLUGINS_HOSTDEVICE AZPLUGINS_FORCEINLINE Scalar getLo(int dim) const + { + return m_lo[dim]; + } + //! Return the upper bound for a given dimension. + AZPLUGINS_HOSTDEVICE AZPLUGINS_FORCEINLINE Scalar getHi(int dim) const + { + return m_hi[dim]; + } + + private: + const T* m_data; + Scalar m_lo[5]; + Scalar m_hi[5]; + Scalar m_dx[5]; + Index5D m_indexer; + }; + + } // namespace azplugins + } // namespace hoomd + +#undef AZPLUGINS_HOSTDEVICE +#undef AZPLUGINS_FORCEINLINE + +#endif // AZPLUGINS_LINEAR_INTERPOLATOR_5D_H_ diff --git a/src/ShapeSymmetry.h b/src/ShapeSymmetry.h new file mode 100644 index 0000000..8c1158f --- /dev/null +++ b/src/ShapeSymmetry.h @@ -0,0 +1,351 @@ +// Copyright (c) 2018-2020, Michael P. Howard +// Copyright (c) 2021-2025, Auburn University +// Part of azplugins, released under the BSD 3-Clause License. + +/*! + * \file ShapeSymmetry.h + * \brief Symmetry evaluators that reduce angular coordinates to a fundamental domain. + */ + +#ifndef AZPLUGINS_SHAPE_SYMMETRY_H_ +#define AZPLUGINS_SHAPE_SYMMETRY_H_ + +#include "hoomd/HOOMDMath.h" +#include "hoomd/VectorMath.h" + +#include + +#ifndef __HIPCC__ +#include +#endif + +#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 + { +namespace detail + { + +//! Convert spherical coordinates to Cartesian coordinates. +AZPLUGINS_HOSTDEVICE inline vec3 sphericalToCartesian(Scalar r, Scalar theta, Scalar phi) + { + Scalar s_th, c_th, s_ph, c_ph; + fast::sincos(theta, s_th, c_th); + fast::sincos(phi, s_ph, c_ph); + return vec3(r * s_ph * c_th, r * s_ph * s_th, r * c_ph); + } + +//! Convert Cartesian coordinates to spherical coordinates. +AZPLUGINS_HOSTDEVICE inline void +cartesianToSpherical(const vec3& v, Scalar& r, Scalar& theta, Scalar& phi) + { + r = fast::sqrt(dot(v, v)); + if (r > Scalar(0)) + { + theta = std::atan2(v.y, v.x); + if (theta < Scalar(0)) + theta += Scalar(2) * Scalar(M_PI); + Scalar cp = v.z / r; + if (cp < Scalar(-1)) + cp = Scalar(-1); + else if (cp > Scalar(1)) + cp = Scalar(1); + phi = slow::acos(cp); + } + else + { + theta = Scalar(0); + phi = Scalar(0); + } + } + +//! Build a quaternion from an axis and an angle. +AZPLUGINS_HOSTDEVICE inline quat quatFromAxisAngle(const vec3& axis, Scalar angle) + { + Scalar s, c; + fast::sincos(Scalar(0.5) * angle, s, c); + return quat(c, s * axis); + } + +//! Build a quaternion from intrinsic ZXZ Euler angles. +AZPLUGINS_HOSTDEVICE inline quat quatFromEulerZXZ(Scalar alpha, Scalar beta, Scalar gamma) + { + const quat qz_a = quatFromAxisAngle(vec3(0, 0, 1), alpha); + const quat qx_b = quatFromAxisAngle(vec3(1, 0, 0), beta); + const quat qz_g = quatFromAxisAngle(vec3(0, 0, 1), gamma); + return qz_a * qx_b * qz_g; + } + +//! Extract intrinsic ZXZ Euler angles from a quaternion. +AZPLUGINS_HOSTDEVICE inline void +eulerFromQuat(const quat& q, Scalar& alpha, Scalar& beta, Scalar& gamma) + { + const rotmat3 R(q); + const Scalar tol = Scalar(1e-7); + + if (R.row2.z < Scalar(-1)) + { + beta = Scalar(M_PI); + } + else if (R.row2.z > Scalar(1)) + { + beta = Scalar(0); + } + else + { + beta = slow::acos(R.row2.z); + } + + if (beta > tol && beta < Scalar(M_PI) - tol) + { + alpha = std::atan2(R.row0.z, -R.row1.z); + gamma = std::atan2(R.row2.x, R.row2.y); + } + else if (beta <= tol) + { + alpha = Scalar(0); + gamma = std::atan2(R.row1.x, R.row0.x); + } + else + { + alpha = Scalar(0); + gamma = std::atan2(-R.row1.x, R.row0.x); + } + + if (alpha < Scalar(0)) + alpha += Scalar(2) * Scalar(M_PI); + if (gamma < Scalar(0)) + gamma += Scalar(2) * Scalar(M_PI); + } + + } // namespace detail + +//! Base class for shared reduced-domain bounds. +class ShapeSymmetry + { + public: + static constexpr Scalar domain_lower[5] + = {Scalar(0.0), Scalar(1.0e-5), Scalar(0.0), Scalar(1.0e-5), Scalar(0.0)}; + }; + +//! Null symmetry: no reduction. +/*! Full natural domain: + theta in [0, 2 pi], phi in [1e-5, pi-1e-5], alpha in [0, 2 pi], + beta in [1e-5, pi-1e-5], gamma in [0, 2 pi]. +*/ +class ShapeSymmetryNull : public ShapeSymmetry + { + public: + //! Upper bounds of the reduced domain. + static constexpr Scalar domain_upper[5] = {Scalar(2.0 * M_PI), + Scalar(M_PI) - Scalar(1e-5), + Scalar(2.0 * M_PI), + Scalar(M_PI) - Scalar(1e-5), + Scalar(2.0 * M_PI)}; + +#ifndef __HIPCC__ + static std::string getName() + { + return "Null"; + } +#endif + + AZPLUGINS_HOSTDEVICE static quat reduce(Scalar& /*theta*/, + Scalar& /*phi*/, + Scalar& /*alpha*/, + Scalar& /*beta*/, + Scalar& /*gamma*/) + { + return quat(Scalar(1), vec3(0, 0, 0)); + } + }; + +//! Cube symmetry evaluator. +/*! Reduced domain: + theta in [0, pi/4], phi in [1e-5, pi/2], alpha in [0, 2 pi], + beta in [1e-5, arccos(1/sqrt(3))], gamma in [0, pi/2]. +*/ +class ShapeSymmetryCube : public ShapeSymmetry + { + public: + //! Upper bounds of the reduced domain. + static constexpr Scalar domain_upper[5] = {Scalar(M_PI / 4.0), + Scalar(M_PI / 2.0), + Scalar(2.0 * M_PI), + Scalar(0.9553166181245093), + Scalar(M_PI / 2.0)}; + +#ifndef __HIPCC__ + static std::string getName() + { + return "Cube"; + } +#endif + + AZPLUGINS_HOSTDEVICE static quat + reduce(Scalar& theta, Scalar& phi, Scalar& alpha, Scalar& beta, Scalar& gamma) + { + // Rotation by pi around x + const quat rot_x_pi + = detail::quatFromAxisAngle(vec3(1, 0, 0), Scalar(M_PI)); + // Rotation by 2 pi/3 around [1,1,1]/sqrt(3) + const quat rot_111(Scalar(0.5), + vec3(Scalar(0.5), Scalar(0.5), Scalar(0.5))); + + quat transformation(Scalar(1), vec3(0, 0, 0)); + + vec3 pos = detail::sphericalToCartesian(Scalar(1), theta, phi); + + // If phi > pi/2, rotate by pi around x to flip z. + if (phi > Scalar(M_PI) / Scalar(2)) + { + pos = rotate(rot_x_pi, pos); + transformation = rot_x_pi * transformation; + } + + // Fold theta into [0, pi/2] by rotating around z. + Scalar r_tmp, th_tmp, ph_tmp; + detail::cartesianToSpherical(pos, r_tmp, th_tmp, ph_tmp); + + const Scalar theta_fold = Scalar(M_PI) / Scalar(2); + if (th_tmp > theta_fold) + { + const Scalar angle = -slow::floor(th_tmp / theta_fold) * theta_fold; + const quat rot_z = detail::quatFromAxisAngle(vec3(0, 0, 1), angle); + pos = rotate(rot_z, pos); + transformation = rot_z * transformation; + } + + // Fold theta into [0, pi/4] using 120-degree rotations around the + // [111] body diagonal (up to 3 attempts). + const Scalar theta_max = Scalar(M_PI) / Scalar(4); + unsigned int n_rot = 0; + detail::cartesianToSpherical(pos, r_tmp, th_tmp, ph_tmp); + while (n_rot < 3 && th_tmp > theta_max) + { + pos = rotate(rot_111, pos); + transformation = rot_111 * transformation; + detail::cartesianToSpherical(pos, r_tmp, th_tmp, ph_tmp); + ++n_rot; + } + + // Write back reduced position angles. + detail::cartesianToSpherical(pos, r_tmp, theta, phi); + + // Apply cumulative transformation to orientation and + // select the best candidate from 3 rotations around [111]. + quat q_orient = detail::quatFromEulerZXZ(alpha, beta, gamma); + quat q_cand = transformation * q_orient; + + Scalar best_a = Scalar(0); + Scalar best_b = Scalar(1e30); + Scalar best_g = Scalar(1e30); + + for (unsigned int i = 0; i < 3; ++i) + { + Scalar a, b, g; + detail::eulerFromQuat(q_cand, a, b, g); + + // If beta > pi/2, reflect. + if (b > Scalar(M_PI) / Scalar(2)) + { + a = std::fmod(a + Scalar(M_PI), Scalar(2) * Scalar(M_PI)); + b = Scalar(M_PI) - b; + g = Scalar(2) * Scalar(M_PI) - g; + if (g < Scalar(0)) + g += Scalar(2) * Scalar(M_PI); + } + + // 4-fold gamma symmetry. + g = std::fmod(g, Scalar(M_PI) / Scalar(2)); + + // Lexicographic pick on (beta, gamma). + if (b < best_b - Scalar(1e-10) || (std::fabs(b - best_b) < Scalar(1e-10) && g < best_g)) + { + best_a = a; + best_b = b; + best_g = g; + } + + q_cand = q_cand * rot_111; + } + + alpha = best_a; + beta = best_b; + gamma = best_g; + + return transformation; + } + }; + +//! Tetrahedron symmetry evaluator. +/*! Reduced domain: + theta in [0, 2 pi/3], phi in [1e-5, pi], alpha in [0, 2 pi], + beta in [1e-5, pi-1e-5], gamma in [0, 2 pi/3]. +*/ +class ShapeSymmetryTetrahedron : public ShapeSymmetry + { + public: + //! Upper bounds of the reduced domain. + static constexpr Scalar domain_upper[5] = {Scalar(2.0 * M_PI / 3.0), + Scalar(M_PI) - Scalar(1e-5), + Scalar(2.0 * M_PI), + Scalar(M_PI - Scalar(1e-5)), + Scalar(2.0 * M_PI / 3.0)}; + +#ifndef __HIPCC__ + static std::string getName() + { + return "Tetrahedron"; + } +#endif + + AZPLUGINS_HOSTDEVICE static quat + reduce(Scalar& theta, Scalar& phi, Scalar& alpha, Scalar& beta, Scalar& gamma) + { + quat transformation(Scalar(1), vec3(0, 0, 0)); + + const Scalar theta_fold = Scalar(2) * Scalar(M_PI) / Scalar(3); + + // Fold theta into [0, 2 pi/3] by rotating around z. + if (theta > theta_fold) + { + vec3 pos = detail::sphericalToCartesian(Scalar(1), theta, phi); + + const Scalar n = slow::floor(theta / theta_fold); + const Scalar angle = -n * theta_fold; + const quat rot_z = detail::quatFromAxisAngle(vec3(0, 0, 1), angle); + pos = rotate(rot_z, pos); + transformation = rot_z * transformation; + + Scalar r_tmp; + detail::cartesianToSpherical(pos, r_tmp, theta, phi); + } + + // Apply transformation to orientation, extract Euler angles. + quat q_orient = detail::quatFromEulerZXZ(alpha, beta, gamma); + quat q_transformed = transformation * q_orient; + detail::eulerFromQuat(q_transformed, alpha, beta, gamma); + + // 3-fold gamma symmetry. + gamma = std::fmod(gamma, theta_fold); + + return transformation; + } + }; + + } // namespace azplugins + } // namespace hoomd + +#undef AZPLUGINS_HOSTDEVICE +#undef AZPLUGINS_FORCEINLINE + +#endif // AZPLUGINS_SHAPE_SYMMETRY_H_ diff --git a/src/export_ChebyshevAnisotropicPairPotential.cc.inc b/src/export_ChebyshevAnisotropicPairPotential.cc.inc new file mode 100644 index 0000000..4cc6507 --- /dev/null +++ b/src/export_ChebyshevAnisotropicPairPotential.cc.inc @@ -0,0 +1,33 @@ +// Copyright (c) 2018-2020, Michael P. Howard +// Copyright (c) 2021-2025, Auburn University +// Part of azplugins, released under the BSD 3-Clause License. + +// Adapted from hoomd/md/export_PotentialPair.cc.inc of HOOMD-blue. +// Copyright (c) 2009-2024 The Regents of the University of Michigan. +// Part of HOOMD-blue, released under the BSD 3-Clause License. + +// See CMakeLists.txt for the source of these variables to be processed by CMake's +// configure_file(). + +// clang-format off +#include "ChebyshevAnisotropicPairPotential.h" + +#define SYMMETRY_CLASS ShapeSymmetry@_symmetry@ +#define EXPORT_FUNCTION export_ChebyshevAnisotropicPairPotential@_symmetry@ +// clang-format on + +namespace hoomd + { +namespace azplugins + { +namespace detail + { +void EXPORT_FUNCTION(pybind11::module& m) + { + export_ChebyshevAnisotropicPairPotential( + m, + "ChebyshevAnisotropicPairPotential@_symmetry@"); + } + } // end namespace detail + } // end namespace azplugins + } // end namespace hoomd diff --git a/src/export_ChebyshevAnisotropicPairPotentialGPU.cc.inc b/src/export_ChebyshevAnisotropicPairPotentialGPU.cc.inc new file mode 100644 index 0000000..2b72073 --- /dev/null +++ b/src/export_ChebyshevAnisotropicPairPotentialGPU.cc.inc @@ -0,0 +1,39 @@ +// Copyright (c) 2018-2020, Michael P. Howard +// Copyright (c) 2021-2025, Auburn University +// Part of azplugins, released under the BSD 3-Clause License. + +// Adapted from hoomd/md/export_PotentialPairGPU.cc.inc of HOOMD-blue. +// Copyright (c) 2009-2026 The Regents of the University of Michigan. +// Part of HOOMD-blue, released under the BSD 3-Clause License. + +// See CMakeLists.txt for the source of these variables to be processed by CMake's +// configure_file(). + +#ifdef ENABLE_HIP + +// clang-format off +#include "ChebyshevAnisotropicPairPotentialGPU.h" + +#define SYMMETRY_CLASS ShapeSymmetry@_symmetry@ +#define EXPORT_FUNCTION export_ChebyshevAnisotropicPairPotential@_symmetry@GPU +// clang-format on + +namespace hoomd + { +namespace azplugins + { +namespace detail + { + +void EXPORT_FUNCTION(pybind11::module& m) + { + export_ChebyshevAnisotropicPairPotentialGPU( + m, + "ChebyshevAnisotropicPairPotential@_symmetry@GPU"); + } + + } // end namespace detail + } // end namespace azplugins + } // end namespace hoomd + +#endif // ENABLE_HIP diff --git a/src/module.cc b/src/module.cc index 1ab81c7..4e75cd7 100644 --- a/src/module.cc +++ b/src/module.cc @@ -69,6 +69,9 @@ void export_ParabolicFlow(pybind11::module&); // pair void export_AnisoPotentialPairTwoPatchMorse(pybind11::module&); +void export_ChebyshevAnisotropicPairPotentialNull(pybind11::module&); +void export_ChebyshevAnisotropicPairPotentialCube(pybind11::module&); +void export_ChebyshevAnisotropicPairPotentialTetrahedron(pybind11::module&); void export_PotentialPairColloid(pybind11::module&); void export_PotentialPairExpandedYukawa(pybind11::module&); void export_PotentialPairHertz(pybind11::module&); @@ -98,6 +101,9 @@ void export_SphericalHarmonicBarrierGPU(pybind11::module&); // pair void export_AnisoPotentialPairTwoPatchMorseGPU(pybind11::module&); +void export_ChebyshevAnisotropicPairPotentialNullGPU(pybind11::module&); +void export_ChebyshevAnisotropicPairPotentialCubeGPU(pybind11::module&); +void export_ChebyshevAnisotropicPairPotentialTetrahedronGPU(pybind11::module&); void export_PotentialPairColloidGPU(pybind11::module&); void export_PotentialPairExpandedYukawaGPU(pybind11::module&); void export_PotentialPairHertzGPU(pybind11::module&); @@ -141,6 +147,9 @@ PYBIND11_MODULE(_azplugins, m) // pair export_AnisoPotentialPairTwoPatchMorse(m); + export_ChebyshevAnisotropicPairPotentialNull(m); + export_ChebyshevAnisotropicPairPotentialCube(m); + export_ChebyshevAnisotropicPairPotentialTetrahedron(m); export_PotentialPairColloid(m); export_PotentialPairExpandedYukawa(m); export_PotentialPairHertz(m); @@ -170,6 +179,9 @@ PYBIND11_MODULE(_azplugins, m) // pair export_AnisoPotentialPairTwoPatchMorseGPU(m); + export_ChebyshevAnisotropicPairPotentialNullGPU(m); + export_ChebyshevAnisotropicPairPotentialCubeGPU(m); + export_ChebyshevAnisotropicPairPotentialTetrahedronGPU(m); export_PotentialPairColloidGPU(m); export_PotentialPairExpandedYukawaGPU(m); export_PotentialPairHertzGPU(m); diff --git a/src/pair.py b/src/pair.py index fbefb07..10884ed 100644 --- a/src/pair.py +++ b/src/pair.py @@ -4,13 +4,106 @@ """Pair potentials.""" +import numpy +import hoomd 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.md import _md from hoomd.variant import Variant +class ChebyshevAnisotropicPairPotential(Force): + """Chebyshev anisotropic pair potential. + + Base class for Chebyshev anisotropic pair potentials. It corresponds to + the ``ShapeSymmetryNull`` (no symmetry reduction), so it + can be instantiated directly for shapes with no symmetry. Shape-specific + subclasses (e.g. ``ChebyshevAnisotropicPairPotentialCube``) override + ``_cpp_class_name`` to select a different compiled C++ symmetry. + """ + + _ext_module = _azplugins + _cpp_class_name = "ChebyshevAnisotropicPairPotentialNull" + + def __init__(self, nlist, 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._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._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.") + + if any(dim < 2 for dim in self.r0.shape): + raise ValueError( + "r0 must have at least 2 grid points along each of its 5 dimensions." + ) + + def _attach_hook(self): + self._nlist._attach(self._simulation) + + if isinstance(self._simulation.device, hoomd.device.CPU): + cls = getattr(self._ext_module, self._cpp_class_name) + self._nlist._cpp_obj.setStorageMode(_md.NeighborList.storageMode.half) + else: + cls = getattr(self._ext_module, self._cpp_class_name + "GPU") + self._nlist._cpp_obj.setStorageMode(_md.NeighborList.storageMode.full) + + self._cpp_obj = cls( + self._simulation.state._cpp_sys_def, + self._nlist._cpp_obj, + self.r_cut, + self._terms, + self._coeffs, + self.r0, + ) + + super()._attach_hook() + + def _detach_hook(self): + self._nlist._detach() + + +class ChebyshevAnisotropicPairPotentialCube(ChebyshevAnisotropicPairPotential): + """Chebyshev anisotropic pair potential with cube symmetry reduction. + + Reduced domain: + theta in [0, pi/4], phi in [1e-5, pi/2], alpha in [0, 2 pi], + beta in [1e-5, arccos(1/sqrt(3))], gamma in [0, pi/2]. + """ + + _cpp_class_name = "ChebyshevAnisotropicPairPotentialCube" + + +class ChebyshevAnisotropicPairPotentialTetrahedron(ChebyshevAnisotropicPairPotential): + """Chebyshev anisotropic pair potential with tetrahedron symmetry reduction. + + Reduced domain: + theta in [0, 2 pi/3], phi in [1e-5, pi-1e-5], alpha in [0, 2 pi], + beta in [1e-5, pi-1e-5], gamma in [0, 2 pi/3]. + """ + + _cpp_class_name = "ChebyshevAnisotropicPairPotentialTetrahedron" + + 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..c6aa277 --- /dev/null +++ b/src/pytest/test_chebyshev.py @@ -0,0 +1,636 @@ +# Copyright (c) 2018-2020, Michael P. Howard +# Copyright (c) 2021-2025, Auburn University +# Part of azplugins, released under the BSD 3-Clause License. + +"""Chebyshev anisotropic pair potential unit tests.""" + +import collections + +import numpy +from scipy.spatial.transform import Rotation + +import hoomd +import hoomd.azplugins + +import pytest + + +_DEVICE_PARAMS = ["cpu"] + +if hoomd.version.gpu_enabled: + try: + if len(hoomd.device.GPU.get_available_devices()) > 0: + _DEVICE_PARAMS.append("gpu") + except Exception: + pass + + +@pytest.fixture(params=_DEVICE_PARAMS) +def simulation_factory(request): + """Create a Simulation on CPU, and on GPU when available.""" + + def make_simulation(snapshot): + if request.param == "cpu": + device = hoomd.device.CPU() + else: + device = hoomd.device.GPU() + + sim = hoomd.Simulation(device=device, seed=1) + sim.create_state_from_snapshot(snapshot) + return sim + + return make_simulation + + +@pytest.fixture +def two_particle_snapshot_factory(): + """Create a basic 2-particle snapshot for pair-potential tests.""" + + def make_snapshot(): + snap = hoomd.Snapshot() + + if snap.communicator.rank == 0: + snap.configuration.box = [20, 20, 20, 0, 0, 0] + snap.particles.N = 2 + snap.particles.types = ["A"] + snap.particles.typeid[:] = [0, 0] + snap.particles.position[:] = [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]] + snap.particles.orientation[:] = [[1.0, 0.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0]] + snap.particles.moment_inertia[:] = [0.0, 0.0, 0.0] + + return snap + + return make_snapshot + + +def good_kwargs(): + """A set of constructor kwargs known to be valid.""" + return dict( + nlist=hoomd.md.nlist.Cell(buffer=0.4), + terms=numpy.zeros((1, 6), dtype=numpy.uint32), + coeffs=numpy.zeros((1,), dtype=numpy.float64), + r0=numpy.zeros((2, 2, 2, 2, 2), dtype=numpy.float64), + r_cut=3.0, + ) + + +def test_chebyshev_invalid_terms_shape(): + """Raise ValueError when ``terms`` is not (Nterms, 6).""" + kwargs = good_kwargs() + kwargs["terms"] = numpy.zeros((1, 5), dtype=numpy.uint32) + kwargs["coeffs"] = numpy.zeros((1,), dtype=numpy.float64) + with pytest.raises(ValueError, match=r"terms must have shape \(Nterms, 6\)\."): + hoomd.azplugins.pair.ChebyshevAnisotropicPairPotential(**kwargs) + + +def test_chebyshev_invalid_coeffs_shape(): + """Raise ValueError when ``coeffs`` length does not match Nterms.""" + kwargs = good_kwargs() + kwargs["terms"] = numpy.zeros((2, 6), dtype=numpy.uint32) + kwargs["coeffs"] = numpy.zeros((1,), dtype=numpy.float64) + with pytest.raises(ValueError, match=r"coeffs must have shape \(Nterms,\)"): + hoomd.azplugins.pair.ChebyshevAnisotropicPairPotential(**kwargs) + + +def test_chebyshev_invalid_r0_ndim(): + """Raise ValueError when ``r0`` is not a 5D array.""" + kwargs = good_kwargs() + kwargs["r0"] = numpy.zeros((2, 2, 2, 2), dtype=numpy.float64) + with pytest.raises(ValueError, match=r"r0 must be a 5D array\."): + hoomd.azplugins.pair.ChebyshevAnisotropicPairPotential(**kwargs) + + +def test_chebyshev_invalid_r0_shape(): + """Raise ValueError when ``r0`` has a dimension with less than 2 points.""" + kwargs = good_kwargs() + kwargs["r0"] = numpy.zeros((2, 2, 1, 2, 2), dtype=numpy.float64) + with pytest.raises( + ValueError, + match=r"r0 must have at least 2 grid points along each of its 5 dimensions\.", + ): + hoomd.azplugins.pair.ChebyshevAnisotropicPairPotential(**kwargs) + + +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) + integrator.methods = [hoomd.md.methods.ConstantVolume(hoomd.filter.All())] + + nlist = hoomd.md.nlist.Cell(buffer=0.4) + + terms = numpy.asarray( + [ + [0, 0, 0, 0, 0, 0], + [1, 0, 2, 0, 1, 3], + ], + dtype=numpy.uint32, + ) + coeffs = numpy.asarray([0.0, 0.0], dtype=numpy.float64) + + # r0 must be 5D (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, 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,))) + + +# Parameters that are identical across every test. +rc = 3.0 +phi_min = 1e-5 +beta_min = 1e-5 + +terms = numpy.array( + [ + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 1, 0, 0], + [1, 0, 0, 0, 0, 0], + [1, 0, 0, 1, 0, 0], + ], + dtype=numpy.uint32, +) +coeffs = numpy.array([1.0, 0.25, 1.5, -1.0], dtype=numpy.float64) + +# r0 data: shape (3, 2, 3, 2, 3) = 108 values. +r0_data = numpy.array([1, 2.1, 3.2] * 36, dtype=numpy.float64).reshape(3, 2, 3, 2, 3) + +PotentialTestCase = collections.namedtuple( + "PotentialTestCase", + [ + "name", + "potential", + "r0", + "rho", + "theta", + "phi", + "alpha", + "beta", + "gamma", + "energy", + "force", + "torque", + "zero_output", + ], +) + +potential_tests = [] + +# Null symmetry +potential_tests += [ + PotentialTestCase( + "null_point_1", + hoomd.azplugins.pair.ChebyshevAnisotropicPairPotential, + 2.1, + 0.2, + numpy.pi / 4, + numpy.pi / 4, + 2 * numpy.pi / 5, + numpy.pi / 2, + numpy.pi, + -0.41, + (-1.324, -1.324, -1.872), + (0.944, -0.307, -0.271), + False, + ), + PotentialTestCase( + "null_point_2", + hoomd.azplugins.pair.ChebyshevAnisotropicPairPotential, + 2.1, + -0.1, + numpy.pi / 4, + numpy.pi / 4, + 2 * numpy.pi / 5, + numpy.pi / 2, + numpy.pi, + -1.67, + (-1.906, -1.906, -2.695), + (1.226, -0.398, -0.398), + False, + ), + PotentialTestCase( + "null_point_3", + hoomd.azplugins.pair.ChebyshevAnisotropicPairPotential, + 2.1, + 0.0, + numpy.pi / 4, + numpy.pi - phi_min, + 2 * numpy.pi / 15, + numpy.pi / 2, + numpy.pi, + -1.583, + (0.0, 0.0, 3.832), + (0.546, -1.226, -0.398), + False, + ), + PotentialTestCase( + "null_point_4", + hoomd.azplugins.pair.ChebyshevAnisotropicPairPotential, + 2.1, + 0.2, + numpy.pi / 4, + numpy.pi / 4, + 2 * numpy.pi / 5, + beta_min, + numpy.pi, + -0.41, + (-1.324, -1.324, -1.872), + (120148.0, -39038.6, -0.271), + False, + ), + PotentialTestCase( + "null_point_5", + hoomd.azplugins.pair.ChebyshevAnisotropicPairPotential, + 1.1375, + 0.95, + numpy.pi / 4, + numpy.pi / 6, + 2 * numpy.pi / 5, + numpy.pi / 2, + numpy.pi / 8, + 2.74, + (-0.174, -0.174, -0.427), + (0.207, -0.067, 0.207), + False, + ), + PotentialTestCase( + "null_point_6", + hoomd.azplugins.pair.ChebyshevAnisotropicPairPotential, + 1.1375, + 1.05, + numpy.pi / 4, + numpy.pi / 6, + 2 * numpy.pi / 5, + numpy.pi / 2, + numpy.pi / 8, + 0.0, + None, + None, + True, + ), +] + +# Cube symmetry +potential_tests += [ + PotentialTestCase( + "cube_point_1", + hoomd.azplugins.pair.ChebyshevAnisotropicPairPotentialCube, + 2.46666667, + 0.2, + numpy.pi / 8, + numpy.pi / 5, + 2 * numpy.pi / 5, + numpy.pi / 6, + numpy.pi / 3, + -0.41, + (-1.335, -0.553, -1.989), + (7.395, -2.403, -0.271), + False, + ), + PotentialTestCase( + "cube_point_2", + hoomd.azplugins.pair.ChebyshevAnisotropicPairPotentialCube, + 2.46666667, + -0.1, + numpy.pi / 8, + numpy.pi / 5, + 2 * numpy.pi / 5, + numpy.pi / 6, + numpy.pi / 3, + -1.67, + (-1.875, -0.777, -2.793), + (9.579, -3.113, -0.398), + False, + ), + PotentialTestCase( + "cube_point_3", + hoomd.azplugins.pair.ChebyshevAnisotropicPairPotentialCube, + 2.62072583, + 0.0, + 2 * numpy.pi / 7, + numpy.pi / 9, + 2 * numpy.pi / 15, + numpy.pi / 4, + numpy.pi / 5, + -0.751, + (-0.518, -0.65, -2.285), + (4.223, -0.398, 3.08), + False, + ), + PotentialTestCase( + "cube_point_4", + hoomd.azplugins.pair.ChebyshevAnisotropicPairPotentialCube, + 2.11254315, + 0.0, + 2 * numpy.pi / 7, + 2 * numpy.pi / 3, + 2 * numpy.pi / 15, + numpy.pi / 3, + numpy.pi / 5, + -0.427, + (-1.256, -1.575, 1.163), + (4.872, -0.906, 0.398), + False, + ), + PotentialTestCase( + "cube_point_5", + hoomd.azplugins.pair.ChebyshevAnisotropicPairPotentialCube, + 1.0, + 0.2, + numpy.pi / 4, + numpy.pi / 4, + 2 * numpy.pi / 5, + beta_min, + numpy.pi, + -0.41, + (-2.023, -2.023, -2.861), + (6.31798953e05, -2.05283924e05, -0.271), + False, + ), + PotentialTestCase( + "cube_point_6", + hoomd.azplugins.pair.ChebyshevAnisotropicPairPotentialCube, + 1.0, + 0.95, + numpy.pi, + 2 * numpy.pi / 3, + 2 * numpy.pi / 5, + 2 * numpy.pi / 3, + 2 * numpy.pi, + 2.61, + (0.363, 0.0, 0.209), + (-1.135, 0.369, -0.207), + False, + ), + PotentialTestCase( + "cube_point_7", + hoomd.azplugins.pair.ChebyshevAnisotropicPairPotentialCube, + 1.0, + 0.95, + 0.0, + 1.0471975511965979, + 1.8849555921538759, + 0.5235987755982987, + 0.0, + 2.61, + (-0.363, -0.0, -0.209), + (1.135, 0.369, 0.207), + False, + ), + PotentialTestCase( + "cube_point_8", + hoomd.azplugins.pair.ChebyshevAnisotropicPairPotentialCube, + 1.0, + 1.05, + numpy.pi / 4, + numpy.pi / 6, + 2 * numpy.pi / 5, + numpy.pi / 2, + numpy.pi / 8, + 0.0, + None, + None, + True, + ), +] + +# Tetrahedron symmetry +potential_tests += [ + PotentialTestCase( + "tetrahedron_point_1", + hoomd.azplugins.pair.ChebyshevAnisotropicPairPotentialTetrahedron, + 2.1, + 0.2, + numpy.pi / 8, + numpy.pi / 5, + 2 * numpy.pi / 5, + numpy.pi / 6, + numpy.pi / 3, + -0.41, + (-1.437, -0.595, -2.142), + (6.111, -1.985, -0.271), + False, + ), + PotentialTestCase( + "tetrahedron_point_2", + hoomd.azplugins.pair.ChebyshevAnisotropicPairPotentialTetrahedron, + 2.1, + -0.1, + numpy.pi / 8, + numpy.pi / 5, + 2 * numpy.pi / 5, + numpy.pi / 6, + numpy.pi / 3, + -1.67, + (-2.07, -0.857, -3.084), + (8.013, -2.604, -0.398), + False, + ), + PotentialTestCase( + "tetrahedron_point_3", + hoomd.azplugins.pair.ChebyshevAnisotropicPairPotentialTetrahedron, + 1.66, + 0.0, + 3 * numpy.pi / 2, + phi_min, + 2 * numpy.pi / 15, + numpy.pi / 4, + numpy.pi / 5, + -0.75, + (0.0, 0.0, -3.182), + (2.084, -4.680, -0.398), + False, + ), + PotentialTestCase( + "tetrahedron_point_4", + hoomd.azplugins.pair.ChebyshevAnisotropicPairPotentialTetrahedron, + 2.1, + 0.65, + numpy.pi / 4, + numpy.pi / 4, + 2 * numpy.pi / 5, + beta_min, + numpy.pi, + 1.48, + (-0.649, -0.649, -0.917), + (1.54802229e05, -5.02982930e04, 0.016), + False, + ), + PotentialTestCase( + "tetrahedron_point_5", + hoomd.azplugins.pair.ChebyshevAnisotropicPairPotentialTetrahedron, + 2.1, + 0.2, + numpy.pi / 4, + numpy.pi - phi_min, + 2 * numpy.pi / 5, + numpy.pi - beta_min, + numpy.pi, + -0.41, + (0.0, 0.0, 2.647), + (2.57516972e05, -8.36723360e04, -0.271), + False, + ), + PotentialTestCase( + "tetrahedron_point_6", + hoomd.azplugins.pair.ChebyshevAnisotropicPairPotentialTetrahedron, + 1.0, + 0.95, + numpy.pi, + 2 * numpy.pi / 3, + 2 * numpy.pi / 5, + 2 * numpy.pi / 3, + 2 * numpy.pi, + 1.873, + (0.146, 0.0, 0.084), + (0.371, -0.120, 0.207), + False, + ), + PotentialTestCase( + "tetrahedron_point_7", + hoomd.azplugins.pair.ChebyshevAnisotropicPairPotentialTetrahedron, + 1.0, + 0.95, + 1.0471975511965979, + 2.0943951023931953, + 5.445427266222309, + 2.0943951023931953, + 0.0, + 1.873, + (-0.073, -0.127, 0.084), + (-0.29, -0.261, 0.207), + False, + ), + PotentialTestCase( + "tetrahedron_point_8", + hoomd.azplugins.pair.ChebyshevAnisotropicPairPotentialTetrahedron, + 1.0, + 1.05, + numpy.pi / 4, + numpy.pi / 6, + 2 * numpy.pi / 5, + numpy.pi / 2, + numpy.pi / 8, + 0.0, + None, + None, + True, + ), +] + + +def rho_to_r(rho, r0, rc): + """Invert rho = (1/r - 1/r0) / (1/(r0+rc) - 1/r0).""" + inv_r0 = 1.0 / r0 + inv_r0_rc = 1.0 / (r0 + rc) + inv_r = rho * (inv_r0_rc - inv_r0) + inv_r0 + return 1.0 / inv_r + + +@pytest.mark.parametrize("potential_test", potential_tests, ids=lambda x: x.name) +def test_energy_force_and_torque( + simulation_factory, two_particle_snapshot_factory, potential_test +): + """Test energy, force, and torque evaluation.""" + snap = two_particle_snapshot_factory() + if snap.communicator.rank == 0: + r = rho_to_r(potential_test.rho, potential_test.r0, rc) + + dx = r * numpy.sin(potential_test.phi) * numpy.cos(potential_test.theta) + dy = r * numpy.sin(potential_test.phi) * numpy.sin(potential_test.theta) + dz = r * numpy.cos(potential_test.phi) + + q_j = Rotation.from_euler( + "ZXZ", + [potential_test.alpha, potential_test.beta, potential_test.gamma], + ).as_quat(scalar_first=True) + + snap.particles.position[:] = [[0.0, 0.0, 0.0], [-dx, -dy, -dz]] + snap.particles.orientation[:] = [[1, 0, 0, 0], q_j] + 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] + + potential = potential_test.potential( + nlist=hoomd.md.nlist.Cell(buffer=1), + terms=terms, + coeffs=coeffs, + r0=r0_data, + r_cut=rc, + ) + integrator.forces = [potential] + + sim.operations.integrator = integrator + sim.run(0) + if sim.device.communicator.rank == 0: + if potential_test.zero_output: + numpy.testing.assert_allclose(potential.energies, [0.0, 0.0], atol=1e-10) + numpy.testing.assert_allclose( + potential.forces, + [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], + atol=1e-10, + ) + numpy.testing.assert_allclose( + potential.torques, + [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], + atol=1e-10, + ) + else: + e = potential_test.energy + f = numpy.array(potential_test.force) + T = numpy.array(potential_test.torque) + + numpy.testing.assert_allclose( + potential.energies, + [0.5 * e, 0.5 * e], + atol=1e-3, + rtol=1e-3, + ) + numpy.testing.assert_allclose( + potential.forces, + [f, -f], + atol=1e-3, + rtol=1e-3, + ) + numpy.testing.assert_allclose( + potential.torques, + [T, -T], + atol=1e-3, + rtol=1e-3, + )