Init Drude velocities (relative temp is bad currently)

plugins/drude/openmmapi/include/openmm/DrudeLangevinIntegrator.h

@@ -194,6 +194,15 @@ protected:
      * Compute the kinetic energy of the system at the current time.
      */
     double computeKineticEnergy();
+    /**
+     * Return a list of velocities normally distributed around a target temperature, correctly
+     * handling center of mass velocities for the Drude pairs.
+     *
+     * @param system the system whose velocities are to be initialized.
+     * @param temperature the target temperature in Kelvin.
+     * @param randomSeed the random number seed to use when selecting velocities 
+     */
+    virtual std::vector<Vec3> getVelocitiesForTemperature(const System &system, double temperature, int randomSeed) const override;
 private:
     double temperature, friction, drudeTemperature, drudeFriction, maxDrudeDistance;
     int randomNumberSeed;

plugins/drude/openmmapi/src/DrudeLangevinIntegrator.cpp

@@ -29,12 +29,16 @@
  * USE OR OTHER DEALINGS IN THE SOFTWARE.                                     *
  * -------------------------------------------------------------------------- */
 
+#include "sfmt/SFMT.h"
+#include "SimTKOpenMMRealType.h"
 #include "openmm/DrudeLangevinIntegrator.h"
 #include "openmm/Context.h"
 #include "openmm/OpenMMException.h"
 #include "openmm/internal/ContextImpl.h"
 #include "openmm/DrudeKernels.h"
+#include <cmath>
 #include <ctime>
+#include <set>
 #include <string>
 
 using namespace OpenMM;
@@ -99,6 +103,82 @@ double DrudeLangevinIntegrator::computeKineticEnergy() {
     return kernel.getAs<IntegrateDrudeLangevinStepKernel>().computeKineticEnergy(*context, *this);
 }
 
+std::vector<Vec3> DrudeLangevinIntegrator::getVelocitiesForTemperature(const System &system, double temperature, int randomSeed) const {
+    // Find the underlying Drude force object
+    const DrudeForce* drudeForce = NULL;
+    for (int i = 0; i < system.getNumForces(); i++)
+        if (dynamic_cast<const DrudeForce*>(&system.getForce(i)) != NULL) {
+            if (drudeForce == NULL)
+                drudeForce = dynamic_cast<const DrudeForce*>(&system.getForce(i));
+            else
+                throw OpenMMException("The System contains multiple DrudeForces");
+        }
+    if (drudeForce == NULL)
+        throw OpenMMException("The System does not contain a DrudeForce");
+
+    // Figure out which particles are individual and which are Drude pairs
+    std::set<int> particles;
+    std::vector<std::pair<int, int>> pairParticles;
+    for (int i = 0; i < system.getNumParticles(); i++) {
+        particles.insert(i);
+    }
+    for (int i = 0; i < drudeForce->getNumParticles(); i++) {
+        int p, p1, p2, p3, p4;
+        double charge, polarizability, aniso12, aniso34;
+        drudeForce->getParticleParameters(i, p, p1, p2, p3, p4, charge, polarizability, aniso12, aniso34);
+        particles.erase(p);
+        particles.erase(p1);
+        pairParticles.emplace_back(p, p1);
+    }
+    std::vector<int> normalParticles(particles.begin(), particles.end());
+
+    // Generate the list of Gaussian random numbers.
+    OpenMM_SFMT::SFMT sfmt;
+    init_gen_rand(randomSeed, sfmt);
+    std::vector<double> randoms;
+    while (randoms.size() < system.getNumParticles()*3) {
+        double x, y, r2;
+        do {
+            x = 2.0*genrand_real2(sfmt)-1.0;
+            y = 2.0*genrand_real2(sfmt)-1.0;
+            r2 = x*x + y*y;
+        } while (r2 >= 1.0 || r2 == 0.0);
+        double multiplier = sqrt((-2.0*std::log(r2))/r2);
+        randoms.push_back(x*multiplier);
+        randoms.push_back(y*multiplier);
+    }
+
+    // Assign the velocities.
+    std::vector<Vec3> velocities(system.getNumParticles(), Vec3());
+    int nextRandom = 0;
+    // First the indivitual atoms
+    for (const auto &atom : normalParticles ) {
+        double mass = system.getParticleMass(atom);
+        if (mass != 0) {
+            double velocityScale = sqrt(BOLTZ*temperature/mass);
+            velocities[atom] = Vec3(randoms[nextRandom++], randoms[nextRandom++], randoms[nextRandom++])*velocityScale;
+        }
+    }
+    // Now the particle-Drude pairs
+    for (const auto &pair : pairParticles ) {
+        const auto atom1 = pair.first;
+        const auto atom2 = pair.second;
+        double mass1 = system.getParticleMass(atom1);
+        double mass2 = system.getParticleMass(atom2);
+        if (mass1 != 0 && mass2 != 0) {
+            double invMass = 1.0 / (mass1 + mass2);
+            double redMass = mass1 * mass2 * invMass;
+            double fracM1 = mass1 * invMass;
+            double fracM2 = mass2 * invMass;
+            Vec3 comVelocity = Vec3(randoms[nextRandom++], randoms[nextRandom++], randoms[nextRandom++])*sqrt(BOLTZ*temperature*invMass);
+            Vec3 relVelocity = Vec3(randoms[nextRandom++], randoms[nextRandom++], randoms[nextRandom++])*sqrt(BOLTZ*drudeTemperature*redMass);
+            velocities[atom1] = comVelocity - fracM2 * relVelocity;
+            velocities[atom2] = comVelocity + fracM1 * relVelocity;
+        }
+    }
+    return velocities;
+}
+
 void DrudeLangevinIntegrator::step(int steps) {
     if (context == NULL)
         throw OpenMMException("This Integrator is not bound to a context!");    

plugins/drude/platforms/reference/tests/TestReferenceDrudeLangevinIntegrator.cpp

@@ -239,9 +239,68 @@ void testForceEnergyConsistency() {
     }
 }
 
+void testInitialTemperature() {
+    // Check temperature initialization for a collection of randomly placed particles
+    const int numRealParticles = 500000;
+    const int numParticles = 2 * numRealParticles;
+    const int nDoF = 3 * numRealParticles;
+    const double targetTemperature = 300;
+    const double drudeTemperature = 1;
+    const double realMass = 10;
+    const double drudeMass = 1;
+    System system;
+    OpenMM_SFMT::SFMT sfmt;
+    init_gen_rand(0, sfmt);
+    std::vector<Vec3> positions(numParticles);
+    DrudeForce* drude = new DrudeForce();
+
+    for (int i = 0; i < numRealParticles; i++) {
+        system.addParticle(realMass);
+        system.addParticle(drudeMass);
+        positions[2*i][0] = genrand_real2(sfmt);
+        positions[2*i][1] = genrand_real2(sfmt);
+        positions[2*i][2] = genrand_real2(sfmt);
+        positions[2*i+1][0] = positions[2*i][0] + 0.01*genrand_real2(sfmt);
+        positions[2*i+1][1] = positions[2*i][1] + 0.01*genrand_real2(sfmt);
+        positions[2*i+1][2] = positions[2*i][2] + 0.01*genrand_real2(sfmt);
+        drude->addParticle(2*i+1, 2*i, -1, -1, -1, -1.0, 0.001, 1, 1);
+    }
+    system.addForce(drude);
+
+    DrudeLangevinIntegrator integrator(targetTemperature, 25, drudeTemperature, 25, 0.001);
+    Platform& platform = Platform::getPlatformByName("Reference");
+    Context context(system, integrator, platform);
+    context.setPositions(positions);
+    context.setVelocitiesToTemperature(targetTemperature);
+    auto velocities = context.getState(State::Velocities).getVelocities();
+    double comKineticEnergy = 0;
+    double relKineticEnergy = 0;
+    for (int i = 0; i < numRealParticles; i++) {
+        int m1 = realMass;
+        int m2 = drudeMass;
+        Vec3 v1 = velocities[2*i];
+        Vec3 v2 = velocities[2*i + 1];
+        double invMass = 1.0 / (m1 + m2);
+        double redMass = m1 * m2 * invMass;
+        double fracM1 = m1 * invMass;
+        double fracM2 = m2 * invMass;
+        Vec3 comVelocity = fracM1 * v1 + fracM2 * v2;
+        Vec3 relVelocity = v2 - v1;
+
+        comKineticEnergy += 0.5 * (m1 + m2) * comVelocity.dot(comVelocity);
+        relKineticEnergy += 0.5 * redMass * relVelocity.dot(relVelocity);
+    }
+    double comTemperature = (2*comKineticEnergy / (nDoF*BOLTZ));
+    double relTemperature = (2*relKineticEnergy / (nDoF*BOLTZ));
+    std::cout << comTemperature << std::endl;
+    std::cout << relTemperature << std::endl;
+    //ASSERT_USUALLY_EQUAL_TOL(targetTemperature, temperature, 0.01);
+}
+
 int main() {
     try {
         registerDrudeReferenceKernelFactories();
+        testInitialTemperature();
         testSinglePair();
         testWater();
         testForceEnergyConsistency();