Fixed bugs in multipoles with ZOnly axis type

plugins/amoeba/platforms/cuda/src/kernels/multipoles.cu

@@ -27,12 +27,16 @@ extern "C" __global__ void computeLabFrameMoments(real4* __restrict__ posq, int4
         if (particles.z >= 0) {
             real4 thisParticlePos = posq[atom];
             real4 posZ = posq[particles.z];
-            real3 vectorZ = make_real3(posZ.x-thisParticlePos.x, posZ.y-thisParticlePos.y, posZ.z-thisParticlePos.z);
+            real3 vectorZ = normalize(make_real3(posZ.x-thisParticlePos.x, posZ.y-thisParticlePos.y, posZ.z-thisParticlePos.z));
             int axisType = particles.w; 
             real4 posX;
             real3 vectorX;
-            if (axisType >= 4)
-                vectorX = make_real3((real) 0.1f);
+            if (axisType >= 4) {
+                if (fabs(vectorZ.x) < 0.866)
+                    vectorX = make_real3(1, 0, 0);
+                else
+                    vectorX = make_real3(0, 1, 0);
+            }
             else {
                 posX = posq[particles.x];
                 vectorX = make_real3(posX.x-thisParticlePos.x, posX.y-thisParticlePos.y, posX.z-thisParticlePos.z);
@@ -81,8 +85,6 @@ extern "C" __global__ void computeLabFrameMoments(real4* __restrict__ posq, int4
         
             // branch based on axis type
                     
-            vectorZ = normalize(vectorZ);
-        
             if (axisType == 1) {
         
                 // bisector
@@ -362,8 +364,12 @@ extern "C" __global__ void mapTorqueToForce(unsigned long long* __restrict__ for
             norms[U] = normVector(vector[U]);
             if (axisType != 4 && particles.x >= 0)
                 vector[V] = atomPos - trimTo3(posq[particles.x]);
+            else {
+                if (fabs(vector[U].x/norms[U]) < 0.866)
+                    vector[V] = make_real3(1, 0, 0);
                 else
-                vector[V] = make_real3(0.1f);
+                    vector[V] = make_real3(0, 1, 0);
+            }
             norms[V] = normVector(vector[V]);
         
             // W = UxV
@@ -488,7 +494,7 @@ extern "C" __global__ void mapTorqueToForce(unsigned long long* __restrict__ for
             else if (axisType == 4) {
                 // z-only
         
-                forces[Z] = vector[UV]*dphi[V]/(norms[U]*angles[UV][1]);
+                forces[Z] = vector[UV]*dphi[V]/(norms[U]*angles[UV][1]) + vector[UW]*dphi[W]/norms[U];
                 forces[X] = make_real3(0);
                 forces[Y] = make_real3(0);
                 forces[I] = -forces[Z];

plugins/amoeba/platforms/cuda/tests/TestCudaAmoebaMultipoleForce.cpp

@@ -3223,6 +3223,55 @@ void testZBisect() {
     ASSERT_EQUAL_TOL(-84.1532, state.getPotentialEnergy(), 0.01);
 }
 
+void testZOnly() {
+    int numParticles = 3;
+    System system;
+    for (int i = 0; i < numParticles; i++)
+        system.addParticle(1.0);
+    AmoebaMultipoleForce* force = new AmoebaMultipoleForce();
+    system.addForce(force);
+    vector<double> d(3), q(9, 0.0);
+    d[0] = 0.05;
+    d[1] = -0.05;
+    d[2] = 0.1;
+    force->addMultipole(0.0, d, q, AmoebaMultipoleForce::ZOnly, 1, 0, 0, 0.39, 0.33, 0.001);
+    force->addMultipole(0.0, d, q, AmoebaMultipoleForce::Bisector, 0, 2, 0, 0.39, 0.33, 0.001);
+    force->addMultipole(0.0, d, q, AmoebaMultipoleForce::ZOnly, 1, 0, 0, 0.39, 0.33, 0.001);
+    vector<Vec3> positions(3);
+    positions[0] = Vec3(0, 0, 0);
+    positions[1] = Vec3(0.2, 0, 0);
+    positions[2] = Vec3(0.2, 0.2, -0.05);
+    
+    // Evaluate the forces.
+    
+    LangevinIntegrator integrator(0.0, 0.1, 0.01);
+    Context context(system, integrator, Platform::getPlatformByName("CUDA"));
+    context.setPositions(positions);
+    State state = context.getState(State::Forces);
+    double norm = 0.0;
+    for (Vec3 f : state.getForces())
+        norm += f[0]*f[0] + f[1]*f[1] + f[2]*f[2];
+    norm = std::sqrt(norm);
+
+    // Take a small step in the direction of the energy gradient and see whether the potential energy changes by the expected amount.
+
+    const double delta = 1e-3;
+    double step = 0.5*delta/norm;
+    vector<Vec3> positions2(numParticles), positions3(numParticles);
+    for (int i = 0; i < numParticles; ++i) {
+        Vec3 p = positions[i];
+        Vec3 f = state.getForces()[i];
+        positions2[i] = Vec3(p[0]-f[0]*step, p[1]-f[1]*step, p[2]-f[2]*step);
+        positions3[i] = Vec3(p[0]+f[0]*step, p[1]+f[1]*step, p[2]+f[2]*step);
+    }
+    context.setPositions(positions2);
+    State state2 = context.getState(State::Energy);
+    context.setPositions(positions3);
+    State state3 = context.getState(State::Energy);
+    ASSERT_EQUAL_TOL(state2.getPotentialEnergy(), state3.getPotentialEnergy()+norm*delta, 1e-3)
+}
+
+
 int main(int argc, char* argv[]) {
     try {
         std::cout << "TestCudaAmoebaMultipoleForce running test..." << std::endl;
@@ -3280,6 +3329,10 @@ int main(int argc, char* argv[]) {
         
         testZBisect();
         
+        // test the ZOnly axis type.
+        
+        testZOnly();
+
     } catch(const std::exception& e) {
         std::cout << "exception: " << e.what() << std::endl;
         std::cout << "FAIL - ERROR.  Test failed." << std::endl;

plugins/amoeba/platforms/reference/src/SimTKReference/AmoebaReferenceMultipoleForce.cpp

@@ -420,7 +420,10 @@ void AmoebaReferenceMultipoleForce::applyRotationMatrixToParticle(      Multipol
 
         // z-only
 
-        vectorX = Vec3(0.1, 0.1, 0.1);
+        if (fabs(vectorZ[0]) < 0.866)
+            vectorX = Vec3(1.0, 0.0, 0.0);
+        else
+            vectorX = Vec3(0.0, 1.0, 0.0);
     }
     else {
         vectorX = particleX->position - particleI.position;
@@ -1755,7 +1758,7 @@ void AmoebaReferenceMultipoleForce::mapTorqueToForceForParticle(const MultipoleP
         // z-only
 
         for (int ii = 0; ii < 3; ii++) {
-            double du                            = vectorUV[ii]*dphi[V]/(norms[U]*angles[UV][1]);
+            double du                            = vectorUV[ii]*dphi[V]/(norms[U]*angles[UV][1]) + vectorUW[ii]*dphi[W]/norms[U];
             forces[particleU.particleIndex][ii] -= du;
             forces[particleI.particleIndex][ii] += du;
         }

plugins/amoeba/platforms/reference/tests/TestReferenceAmoebaMultipoleForce.cpp

@@ -3030,6 +3030,54 @@ void testZBisect() {
     ASSERT_EQUAL_TOL(-84.1532, state.getPotentialEnergy(), 0.01);
 }
 
+void testZOnly() {
+    int numParticles = 3;
+    System system;
+    for (int i = 0; i < numParticles; i++)
+        system.addParticle(1.0);
+    AmoebaMultipoleForce* force = new AmoebaMultipoleForce();
+    system.addForce(force);
+    vector<double> d(3), q(9, 0.0);
+    d[0] = 0.05;
+    d[1] = -0.05;
+    d[2] = 0.1;
+    force->addMultipole(0.0, d, q, AmoebaMultipoleForce::ZOnly, 1, 0, 0, 0.39, 0.33, 0.001);
+    force->addMultipole(0.0, d, q, AmoebaMultipoleForce::Bisector, 0, 2, 0, 0.39, 0.33, 0.001);
+    force->addMultipole(0.0, d, q, AmoebaMultipoleForce::ZOnly, 1, 0, 0, 0.39, 0.33, 0.001);
+    vector<Vec3> positions(3);
+    positions[0] = Vec3(0, 0, 0);
+    positions[1] = Vec3(0.2, 0, 0);
+    positions[2] = Vec3(0.2, 0.2, -0.05);
+    
+    // Evaluate the forces.
+    
+    LangevinIntegrator integrator(0.0, 0.1, 0.01);
+    Context context(system, integrator, Platform::getPlatformByName("Reference"));
+    context.setPositions(positions);
+    State state = context.getState(State::Forces);
+    double norm = 0.0;
+    for (Vec3 f : state.getForces())
+        norm += f[0]*f[0] + f[1]*f[1] + f[2]*f[2];
+    norm = std::sqrt(norm);
+
+    // Take a small step in the direction of the energy gradient and see whether the potential energy changes by the expected amount.
+
+    const double delta = 1e-3;
+    double step = 0.5*delta/norm;
+    vector<Vec3> positions2(numParticles), positions3(numParticles);
+    for (int i = 0; i < numParticles; ++i) {
+        Vec3 p = positions[i];
+        Vec3 f = state.getForces()[i];
+        positions2[i] = Vec3(p[0]-f[0]*step, p[1]-f[1]*step, p[2]-f[2]*step);
+        positions3[i] = Vec3(p[0]+f[0]*step, p[1]+f[1]*step, p[2]+f[2]*step);
+    }
+    context.setPositions(positions2);
+    State state2 = context.getState(State::Energy);
+    context.setPositions(positions3);
+    State state3 = context.getState(State::Energy);
+    ASSERT_EQUAL_TOL(state2.getPotentialEnergy(), state3.getPotentialEnergy()+norm*delta, 1e-3)
+}
+
 int main(int numberOfArguments, char* argv[]) {
 
     try {
@@ -3083,6 +3131,10 @@ int main(int numberOfArguments, char* argv[]) {
         // test the ZBisect axis type.
         
         testZBisect();
+        
+        // test the ZOnly axis type.
+        
+        testZOnly();
     }
     catch(const std::exception& e) {
         std::cout << "exception: " << e.what() << std::endl;