                         volatile float * __restrict pts, // (nnodes + 1) x 2
                         volatile float * __restrict attr_forces, // (N x 2)
                         volatile float * __restrict rep_forces, // (nnodes + 1) x 2
+                        volatile float * __restrict gains,
                         volatile float * __restrict old_forces) // (N x 2)
 {
   register int i, inc;
-  register float tmpx, tmpy;
+  register float dx, dy, ux, uy, gx, xy;
 
   // iterate over all bodies assigned to thread
   // TODO: fix momentum at step 0
   inc = blockDim.x * gridDim.x;
   for (i = threadIdx.x + blockIdx.x * blockDim.x; i < N; i += inc) {
-      tmpx = 4.0f * (exaggeration*attr_forces[i] - (rep_forces[i] / norm));
-      tmpy = 4.0f * (exaggeration*attr_forces[i + N] - (rep_forces[nnodes + 1 + i] / norm));
-      tmpx = momentum * tmpx + (1 - momentum) * old_forces[i];
-      tmpy = momentum * tmpy + (1 - momentum) * old_forces[i + N];
-
-      // If the point is screwed up, just put it in the middle of the screen
-      if (pts[i] != pts[i] || isinf(pts[i])) 
-        pts[i] = 0;
-      if (pts[i + nnodes + 1] != pts[i + nnodes + 1] || isinf(pts[i + nnodes + 1])) 
-        pts[i + nnodes + 1] = 0;
+        ux = old_forces[i];
+        uy = old_forces[N + i];
+        gx = gains[i];
+        gy = gains[N + i];
+        dx = 4.0f * (exaggeration*attr_forces[i] - (rep_forces[i] / norm));
+        dy = 4.0f * (exaggeration*attr_forces[i + N] - (rep_forces[nnodes + 1 + i] / norm));
+
+        gx = (signbit(dx) != signbit(ux)) ? gx + 0.2 : gx * 0.8;
+        gy = (signbit(dy) != signbit(uy)) ? gy + 0.2 : gy * 0.8;
+        gx = (gx < 0.01) ? 0.01 : gx;
+        gy = (gy < 0.01) ? 0.01 : gy;
+
+        ux = momentum * ux - eta * gx * dx;
+        uy = momentum * uy - eta * gy * dy;
+
+        pts[i] += ux;
+        pts[i + nnodes + 1] += uy;
+
+        old_forces[i] = ux;
+        old_forces[N + i] = uy;
+        gains[i] = gx;
+        gains[N + i] = gy;
+
         
-      pts[i] -= eta * tmpx;
-      pts[i + nnodes + 1] -= eta * tmpy;
-      old_forces[i] = tmpx;
-      old_forces[i + N] = tmpy;
+
+        // tmpx = momentum * tmpx + (1 - momentum) * old_forces[i];
+        // tmpy = momentum * tmpy + (1 - momentum) * old_forces[i + N];
+
+        // // If the point is screwed up, just put it in the middle of the screen
+        // if (pts[i] != pts[i] || isinf(pts[i])) 
+        // pts[i] = 0;
+        // if (pts[i + nnodes + 1] != pts[i + nnodes + 1] || isinf(pts[i + nnodes + 1])) 
+        // pts[i + nnodes + 1] = 0;
+
+        // pts[i] -= eta * tmpx;
+        // pts[i + nnodes + 1] -= eta * tmpy;
+        // old_forces[i] = tmpx;
+        // old_forces[i + N] = tmpy;
    }
 }
 

                                                                      thrust::raw_pointer_cast(pts.data()),
                                                                      thrust::raw_pointer_cast(attr_forces.data()),
                                                                      thrust::raw_pointer_cast(rep_forces.data()),
                                                                      thrust::raw_pointer_cast(old_forces.data()));
          gpuErrchk(cudaDeviceSynchronize());
          thrust::transform(pts.begin(), pts.end(), random_vec.begin(), pts.begin(), thrust::plus<float>());

        thrust::device_vector<float> forceProd(sparsePij.size());
        thrust::device_vector<float> rep_forces((nnodes + 1) * 2, 0);
        thrust::device_vector<float> attr_forces(N_POINTS * 2, 0);
        thrust::device_vector<float> gains(N_POINTS * 2, 1);
        thrust::device_vector<float> old_forces(N_POINTS * 2, 0); // for momentum
        thrust::device_vector<int> errl(1);
        thrust::device_vector<int> startl(nnodes + 1);
