Virtual sites can depend on other virtual sites (#4348)

* Reference platform supports nested virtual sites

* Common platform supports nested virtual sites

* Fixed force distribution from nested virtual sites

* Fixed test failures

openmmapi/src/ContextImpl.cpp

@@ -62,16 +62,9 @@ ContextImpl::ContextImpl(Context& owner, const System& system, Integrator& integ
     
     // Check for errors in virtual sites and massless particles.
     
-    for (int i = 0; i < numParticles; i++) {
-        if (system.isVirtualSite(i)) {
-            if (system.getParticleMass(i) != 0.0)
+    for (int i = 0; i < numParticles; i++)
+        if (system.isVirtualSite(i) && system.getParticleMass(i) != 0.0)
             throw OpenMMException("Virtual site has nonzero mass");
-            const VirtualSite& site = system.getVirtualSite(i);
-            for (int j = 0; j < site.getNumParticles(); j++)
-                if (system.isVirtualSite(site.getParticle(j)))
-                    throw OpenMMException("A virtual site cannot depend on another virtual site");
-        }
-    }
     set<pair<int, int> > constraintAtoms;
     for (int i = 0; i < system.getNumConstraints(); i++) {
         int particle1, particle2;

platforms/common/include/openmm/common/IntegrationUtilities.h

@@ -176,7 +176,8 @@ protected:
     ComputeArray vsiteLocalCoordsWeights;
     ComputeArray vsiteLocalCoordsPos;
     ComputeArray vsiteLocalCoordsStartIndex;
-    int randomPos, lastSeed, numVsites;
+    ComputeArray vsiteStage;
+    int randomPos, lastSeed, numVsites, numVsiteStages;
     bool hasOverlappingVsites;
     mm_double2 lastStepSize;
     struct ShakeCluster;

platforms/common/src/IntegrationUtilities.cpp

@@ -529,6 +529,38 @@ IntegrationUtilities::IntegrationUtilities(ComputeContext& context, const System
         if (atomCounts[i] > 1)
             hasOverlappingVsites = true;
 
+    // Divide virtual sites into stages to resolve dependencies between them.
+
+    set<int> sites;
+    vector<int> vsiteStageVec(numAtoms, -1);
+    for (int i = 0; i < numAtoms; i++)
+        if (system.isVirtualSite(i)) {
+            sites.insert(i);
+            vsiteStageVec[i] = numAtoms;
+        }
+    numVsiteStages = 0;
+    int remainingSites = 0;
+    while (sites.size() > 0) {
+        if (sites.size() == remainingSites)
+            throw OpenMMException("Virtual site definitions are circular");
+        for (auto index = sites.begin(); index != sites.end();) {
+            const VirtualSite& site = system.getVirtualSite(*index);
+            bool canCompute = true;
+            for (int i = 0; i < site.getNumParticles(); i++)
+                if (vsiteStageVec[site.getParticle(i)] >= numVsiteStages)
+                    canCompute = false;
+            if (canCompute) {
+                vsiteStageVec[*index] = numVsiteStages;
+                index = sites.erase(index);
+            }
+            else
+                ++index;
+        }
+        numVsiteStages++;
+    }
+    vsiteStage.initialize<int>(context, vsiteStageVec.size(), "vsiteStages");
+    vsiteStage.upload(vsiteStageVec);
+
     // Create the kernels used by this class.
 
     map<string, string> defines;
@@ -542,6 +574,8 @@ IntegrationUtilities::IntegrationUtilities(ComputeContext& context, const System
     defines["PADDED_NUM_ATOMS"] = context.intToString(context.getPaddedNumAtoms());
     if (hasOverlappingVsites)
         defines["HAS_OVERLAPPING_VSITES"] = "1";
+    if (numVsiteStages > 1)
+        defines["MULTIPLE_VSITE_STAGES"] = "1";
     ComputeProgram program = context.compileProgram(CommonKernelSources::integrationUtilities, defines);
     settlePosKernel = program->createKernel("applySettleToPositions");
     settleVelKernel = program->createKernel("applySettleToVelocities");
@@ -577,6 +611,8 @@ IntegrationUtilities::IntegrationUtilities(ComputeContext& context, const System
     vsitePositionKernel->addArg(vsiteLocalCoordsWeights);
     vsitePositionKernel->addArg(vsiteLocalCoordsPos);
     vsitePositionKernel->addArg(vsiteLocalCoordsStartIndex);
+    vsitePositionKernel->addArg(vsiteStage);
+    vsitePositionKernel->addArg();
     vsiteForceKernel->addArg(context.getPosq());
     if (context.getUseMixedPrecision())
         vsiteForceKernel->addArg(context.getPosqCorrection());
@@ -594,6 +630,8 @@ IntegrationUtilities::IntegrationUtilities(ComputeContext& context, const System
     vsiteForceKernel->addArg(vsiteLocalCoordsWeights);
     vsiteForceKernel->addArg(vsiteLocalCoordsPos);
     vsiteForceKernel->addArg(vsiteLocalCoordsStartIndex);
+    vsiteForceKernel->addArg(vsiteStage);
+    vsiteForceKernel->addArg();
     for (int i = 0; i < 3; i++)
         vsiteSaveForcesKernel->addArg();
 
@@ -736,8 +774,10 @@ void IntegrationUtilities::applyVelocityConstraints(double tol) {
 
 void IntegrationUtilities::computeVirtualSites() {
     ContextSelector selector(context);
-    if (numVsites > 0)
+    for (int i = 0; i < numVsiteStages; i++) {
+        vsitePositionKernel->setArg(14, i);
         vsitePositionKernel->execute(numVsites);
+    }
 }
 
 void IntegrationUtilities::initRandomNumberGenerator(unsigned int randomNumberSeed) {

platforms/common/src/kernels/integrationUtilities.cc

@@ -808,12 +808,17 @@ KERNEL void computeVirtualSites(GLOBAL real4* RESTRICT posq, GLOBAL real4* RESTR
         GLOBAL const int4* RESTRICT outOfPlaneAtoms, GLOBAL const real4* RESTRICT outOfPlaneWeights,
         GLOBAL const int* RESTRICT localCoordsIndex, GLOBAL const int* RESTRICT localCoordsAtoms,
         GLOBAL const real* RESTRICT localCoordsWeights, GLOBAL const real4* RESTRICT localCoordsPos,
-        GLOBAL const int* RESTRICT localCoordsStartIndex) {
+        GLOBAL const int* RESTRICT localCoordsStartIndex, GLOBAL const int* RESTRICT vsiteStage,
+        int currentStage) {
     
     // Two particle average sites.
     
     for (int index = GLOBAL_ID; index < NUM_2_AVERAGE; index += GLOBAL_SIZE) {
         int4 atoms = avg2Atoms[index];
+#ifdef MULTIPLE_VSITE_STAGES
+        if (vsiteStage[atoms.x] != currentStage)
+            continue;
+#endif
         real2 weights = avg2Weights[index];
         mixed4 pos = loadPos(posq, posqCorrection, atoms.x);
         mixed4 pos1 = loadPos(posq, posqCorrection, atoms.y);
@@ -828,6 +833,10 @@ KERNEL void computeVirtualSites(GLOBAL real4* RESTRICT posq, GLOBAL real4* RESTR
     
     for (int index = GLOBAL_ID; index < NUM_3_AVERAGE; index += GLOBAL_SIZE) {
         int4 atoms = avg3Atoms[index];
+#ifdef MULTIPLE_VSITE_STAGES
+        if (vsiteStage[atoms.x] != currentStage)
+            continue;
+#endif
         real4 weights = avg3Weights[index];
         mixed4 pos = loadPos(posq, posqCorrection, atoms.x);
         mixed4 pos1 = loadPos(posq, posqCorrection, atoms.y);
@@ -843,6 +852,10 @@ KERNEL void computeVirtualSites(GLOBAL real4* RESTRICT posq, GLOBAL real4* RESTR
     
     for (int index = GLOBAL_ID; index < NUM_OUT_OF_PLANE; index += GLOBAL_SIZE) {
         int4 atoms = outOfPlaneAtoms[index];
+#ifdef MULTIPLE_VSITE_STAGES
+        if (vsiteStage[atoms.x] != currentStage)
+            continue;
+#endif
         real4 weights = outOfPlaneWeights[index];
         mixed4 pos = loadPos(posq, posqCorrection, atoms.x);
         mixed4 pos1 = loadPos(posq, posqCorrection, atoms.y);
@@ -861,6 +874,10 @@ KERNEL void computeVirtualSites(GLOBAL real4* RESTRICT posq, GLOBAL real4* RESTR
     
     for (int index = GLOBAL_ID; index < NUM_LOCAL_COORDS; index += GLOBAL_SIZE) {
         int siteAtomIndex = localCoordsIndex[index];
+#ifdef MULTIPLE_VSITE_STAGES
+        if (vsiteStage[siteAtomIndex] != currentStage)
+            continue;
+#endif
         int start = localCoordsStartIndex[index];
         int end = localCoordsStartIndex[index+1];
         mixed3 origin = make_mixed3(0), xdir = make_mixed3(0), ydir = make_mixed3(0);
@@ -915,12 +932,17 @@ KERNEL void distributeVirtualSiteForces(GLOBAL const real4* RESTRICT posq, GLOBA
         GLOBAL const int4* RESTRICT outOfPlaneAtoms, GLOBAL const real4* RESTRICT outOfPlaneWeights,
         GLOBAL const int* RESTRICT localCoordsIndex, GLOBAL const int* RESTRICT localCoordsAtoms,
         GLOBAL const real* RESTRICT localCoordsWeights, GLOBAL const real4* RESTRICT localCoordsPos,
-        GLOBAL const int* RESTRICT localCoordsStartIndex) {
+        GLOBAL const int* RESTRICT localCoordsStartIndex, GLOBAL const int* RESTRICT vsiteStage,
+        int currentStage) {
     
     // Two particle average sites.
     
     for (int index = GLOBAL_ID; index < NUM_2_AVERAGE; index += GLOBAL_SIZE) {
         int4 atoms = avg2Atoms[index];
+#ifdef MULTIPLE_VSITE_STAGES
+        if (vsiteStage[atoms.x] != currentStage)
+            continue;
+#endif
         real2 weights = avg2Weights[index];
         real3 f = loadForce(atoms.x, force);
         addForce(atoms.y, force, f*weights.x);
@@ -931,6 +953,10 @@ KERNEL void distributeVirtualSiteForces(GLOBAL const real4* RESTRICT posq, GLOBA
     
     for (int index = GLOBAL_ID; index < NUM_3_AVERAGE; index += GLOBAL_SIZE) {
         int4 atoms = avg3Atoms[index];
+#ifdef MULTIPLE_VSITE_STAGES
+        if (vsiteStage[atoms.x] != currentStage)
+            continue;
+#endif
         real4 weights = avg3Weights[index];
         real3 f = loadForce(atoms.x, force);
         addForce(atoms.y, force, f*weights.x);
@@ -942,6 +968,10 @@ KERNEL void distributeVirtualSiteForces(GLOBAL const real4* RESTRICT posq, GLOBA
     
     for (int index = GLOBAL_ID; index < NUM_OUT_OF_PLANE; index += GLOBAL_SIZE) {
         int4 atoms = outOfPlaneAtoms[index];
+#ifdef MULTIPLE_VSITE_STAGES
+        if (vsiteStage[atoms.x] != currentStage)
+            continue;
+#endif
         real4 weights = outOfPlaneWeights[index];
         mixed4 pos1 = loadPos(posq, posqCorrection, atoms.y);
         mixed4 pos2 = loadPos(posq, posqCorrection, atoms.z);
@@ -964,6 +994,10 @@ KERNEL void distributeVirtualSiteForces(GLOBAL const real4* RESTRICT posq, GLOBA
     
     for (int index = GLOBAL_ID; index < NUM_LOCAL_COORDS; index += GLOBAL_SIZE) {
         int siteAtomIndex = localCoordsIndex[index];
+#ifdef MULTIPLE_VSITE_STAGES
+        if (vsiteStage[siteAtomIndex] != currentStage)
+            continue;
+#endif
         int start = localCoordsStartIndex[index];
         int end = localCoordsStartIndex[index+1];
         mixed3 origin = make_mixed3(0), xdir = make_mixed3(0), ydir = make_mixed3(0);

platforms/cpu/src/CpuKernels.cpp

@@ -6,7 +6,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2013-2022 Stanford University and the Authors.      *
+ * Portions copyright (c) 2013-2023 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -85,6 +85,11 @@ static ReferenceConstraints& extractConstraints(ContextImpl& context) {
     return *data->constraints;
 }
 
+static const ReferenceVirtualSites& extractVirtualSites(ContextImpl& context) {
+    ReferencePlatform::PlatformData* data = reinterpret_cast<ReferencePlatform::PlatformData*>(context.getPlatformData());
+    return *data->virtualSites;
+}
+
 static map<string, double>& extractEnergyParameterDerivatives(ContextImpl& context) {
     ReferencePlatform::PlatformData* data = reinterpret_cast<ReferencePlatform::PlatformData*>(context.getPlatformData());
     return *data->energyParameterDerivatives;
@@ -1461,6 +1466,7 @@ void CpuIntegrateLangevinStepKernel::execute(ContextImpl& context, const Langevi
             delete dynamics;
         dynamics = new CpuLangevinDynamics(context.getSystem().getNumParticles(), stepSize, friction, temperature, data.threads, data.random);
         dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
+        dynamics->setVirtualSites(extractVirtualSites(context));
         prevTemp = temperature;
         prevFriction = friction;
         prevStepSize = stepSize;
@@ -1501,6 +1507,7 @@ void CpuIntegrateLangevinMiddleStepKernel::execute(ContextImpl& context, const L
             delete dynamics;
         dynamics = new CpuLangevinMiddleDynamics(context.getSystem().getNumParticles(), stepSize, friction, temperature, data.threads, data.random);
         dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
+        dynamics->setVirtualSites(extractVirtualSites(context));
         prevTemp = temperature;
         prevFriction = friction;
         prevStepSize = stepSize;

platforms/cuda/src/CudaIntegrationUtilities.cpp

@@ -134,8 +134,9 @@ void CudaIntegrationUtilities::applyConstraintsImpl(bool constrainVelocities, do
 
 void CudaIntegrationUtilities::distributeForcesFromVirtualSites() {
     ContextSelector selector(context);
-    if (numVsites > 0) {
+    for (int i = numVsiteStages-1; i >= 0; i--) {
         vsiteForceKernel->setArg(2, context.getLongForceBuffer());
+        vsiteForceKernel->setArg(15, i);
         vsiteForceKernel->execute(numVsites);
     }
 }

platforms/opencl/src/OpenCLIntegrationUtilities.cpp

@@ -131,8 +131,9 @@ void OpenCLIntegrationUtilities::applyConstraintsImpl(bool constrainVelocities,
 }
 
 void OpenCLIntegrationUtilities::distributeForcesFromVirtualSites() {
-    if (numVsites > 0) {
+    for (int i = numVsiteStages-1; i >= 0; i--) {
         vsiteForceKernel->setArg(2, context.getLongForceBuffer());
+        vsiteForceKernel->setArg(15, i);
         vsiteForceKernel->execute(numVsites);
         vsiteSaveForcesKernel->setArg(0, context.getLongForceBuffer());
         vsiteSaveForcesKernel->setArg(1, context.getForceBuffers());

platforms/reference/include/ReferenceDynamics.h

@@ -26,6 +26,7 @@
 #define __ReferenceDynamics_H__
 
 #include "ReferenceConstraintAlgorithm.h"
+#include "ReferenceVirtualSites.h"
 #include "openmm/System.h"
 #include <cstddef>
 #include <vector>
@@ -52,8 +53,8 @@ class OPENMM_EXPORT ReferenceDynamics {
       double _deltaT;
       double _temperature;
 
-      int _ownReferenceConstraint;
       ReferenceConstraintAlgorithm* _referenceConstraint;
+      const ReferenceVirtualSites* virtualSites;
       
    public:
 
@@ -170,6 +171,16 @@ class OPENMM_EXPORT ReferenceDynamics {
          --------------------------------------------------------------------------------------- */
       
       void setReferenceConstraintAlgorithm(ReferenceConstraintAlgorithm* referenceConstraint);
+      
+      /**
+       * Get the object used to compute virtual sites.
+       */
+      const ReferenceVirtualSites& getVirtualSites() const;
+
+      /**
+       * Set the object used to compute virtual sites.
+       */
+      void setVirtualSites(const ReferenceVirtualSites& sites);
 };
 
 } // namespace OpenMM

platforms/reference/include/ReferencePlatform.h

@@ -9,7 +9,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2008-2019 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2023 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -36,6 +36,7 @@
 #include "openmm/System.h"
 #include "openmm/internal/windowsExport.h"
 #include "ReferenceConstraints.h"
+#include "ReferenceVirtualSites.h"
 #include <map>
 #include <vector>
 
@@ -72,6 +73,7 @@ public:
     Vec3* periodicBoxSize;
     Vec3* periodicBoxVectors;
     ReferenceConstraints* constraints;
+    ReferenceVirtualSites* virtualSites;
     std::map<std::string, double>* energyParameterDerivatives;
 };
 } // namespace OpenMM

platforms/reference/include/ReferenceVirtualSites.h

@@ -6,7 +6,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2012 Stanford University and the Authors.           *
+ * Portions copyright (c) 2012-2023 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -40,14 +40,17 @@ namespace OpenMM {
 
 class OPENMM_EXPORT ReferenceVirtualSites {
 public:
+    ReferenceVirtualSites(const System& system);
     /**
      * Compute the positions of all virtual sites.
      */
-    static void computePositions(const OpenMM::System& system, std::vector<OpenMM::Vec3>& atomCoordinates);
+    void computePositions(const System& system, std::vector<Vec3>& atomCoordinates) const;
     /**
      * Distribute forces from virtual sites to the atoms they are based on.
      */
-    static void distributeForces(const OpenMM::System& system, const std::vector<OpenMM::Vec3>& atomCoordinates, std::vector<OpenMM::Vec3>& forces);
+    void distributeForces(const System& system, const std::vector<Vec3>& atomCoordinates, std::vector<Vec3>& forces) const;
+private:
+    std::vector<int> order;
 };
 
 } // namespace OpenMM

platforms/reference/src/ReferenceKernels.cpp

@@ -121,6 +121,11 @@ static ReferenceConstraints& extractConstraints(ContextImpl& context) {
     return *data->constraints;
 }
 
+static const ReferenceVirtualSites& extractVirtualSites(ContextImpl& context) {
+    ReferencePlatform::PlatformData* data = reinterpret_cast<ReferencePlatform::PlatformData*>(context.getPlatformData());
+    return *data->virtualSites;
+}
+
 static map<string, double>& extractEnergyParameterDerivatives(ContextImpl& context) {
     ReferencePlatform::PlatformData* data = reinterpret_cast<ReferencePlatform::PlatformData*>(context.getPlatformData());
     return *data->energyParameterDerivatives;
@@ -175,7 +180,7 @@ double ReferenceCalcForcesAndEnergyKernel::finishComputation(ContextImpl& contex
     if (!includeForces)
         extractForces(context) = savedForces; // Restore the forces so computing the energy doesn't overwrite the forces with incorrect values.
     else
-        ReferenceVirtualSites::distributeForces(context.getSystem(), extractPositions(context), extractForces(context));
+        extractVirtualSites(context).distributeForces(context.getSystem(), extractPositions(context), extractForces(context));
     return 0.0;
 }
 
@@ -341,7 +346,7 @@ ReferenceApplyConstraintsKernel::~ReferenceApplyConstraintsKernel() {
 void ReferenceApplyConstraintsKernel::apply(ContextImpl& context, double tol) {
     vector<Vec3>& positions = extractPositions(context);
     extractConstraints(context).apply(positions, positions, inverseMasses, tol);
-    ReferenceVirtualSites::computePositions(context.getSystem(), positions);
+    extractVirtualSites(context).computePositions(context.getSystem(), positions);
 }
 
 void ReferenceApplyConstraintsKernel::applyToVelocities(ContextImpl& context, double tol) {
@@ -355,7 +360,7 @@ void ReferenceVirtualSitesKernel::initialize(const System& system) {
 
 void ReferenceVirtualSitesKernel::computePositions(ContextImpl& context) {
     vector<Vec3>& positions = extractPositions(context);
-    ReferenceVirtualSites::computePositions(context.getSystem(), positions);
+    extractVirtualSites(context).computePositions(context.getSystem(), positions);
 }
 
 void ReferenceCalcHarmonicBondForceKernel::initialize(const System& system, const HarmonicBondForce& force) {
@@ -2290,6 +2295,7 @@ void ReferenceIntegrateVerletStepKernel::execute(ContextImpl& context, const Ver
             delete dynamics;
         dynamics = new ReferenceVerletDynamics(context.getSystem().getNumParticles(), stepSize);
         dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
+        dynamics->setVirtualSites(extractVirtualSites(context));
         prevStepSize = stepSize;
     }
     dynamics->update(context.getSystem(), posData, velData, forceData, masses, integrator.getConstraintTolerance());
@@ -2327,6 +2333,7 @@ void ReferenceIntegrateNoseHooverStepKernel::execute(ContextImpl& context, const
             delete dynamics;
         dynamics = new ReferenceNoseHooverDynamics(context.getSystem().getNumParticles(), stepSize);
         dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
+        dynamics->setVirtualSites(extractVirtualSites(context));
         prevStepSize = stepSize;
     }
     dynamics->step1(context, context.getSystem(), posData, velData, forceData, masses, integrator.getConstraintTolerance(), forcesAreValid,
@@ -2594,6 +2601,7 @@ void ReferenceIntegrateLangevinStepKernel::execute(ContextImpl& context, const L
                 friction, 
                 temperature);
         dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
+        dynamics->setVirtualSites(extractVirtualSites(context));
         prevTemp = temperature;
         prevFriction = friction;
         prevStepSize = stepSize;
@@ -2637,6 +2645,7 @@ void ReferenceIntegrateLangevinMiddleStepKernel::execute(ContextImpl& context, c
                 friction, 
                 temperature);
         dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
+        dynamics->setVirtualSites(extractVirtualSites(context));
         prevTemp = temperature;
         prevFriction = friction;
         prevStepSize = stepSize;
@@ -2681,6 +2690,7 @@ void ReferenceIntegrateBrownianStepKernel::execute(ContextImpl& context, const B
                 friction, 
                 temperature);
         dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
+        dynamics->setVirtualSites(extractVirtualSites(context));
         prevTemp = temperature;
         prevFriction = friction;
         prevStepSize = stepSize;
@@ -2721,6 +2731,7 @@ double ReferenceIntegrateVariableLangevinStepKernel::execute(ContextImpl& contex
             delete dynamics;
         dynamics = new ReferenceVariableStochasticDynamics(context.getSystem().getNumParticles(), friction, temperature, errorTol);
         dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
+        dynamics->setVirtualSites(extractVirtualSites(context));
         prevTemp = temperature;
         prevFriction = friction;
         prevErrorTol = errorTol;
@@ -2764,6 +2775,7 @@ double ReferenceIntegrateVariableVerletStepKernel::execute(ContextImpl& context,
             delete dynamics;
         dynamics = new ReferenceVariableVerletDynamics(context.getSystem().getNumParticles(), errorTol);
         dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
+        dynamics->setVirtualSites(extractVirtualSites(context));
         prevErrorTol = errorTol;
     }
     double maxStepSize = maxTime-data.time;
@@ -2816,6 +2828,7 @@ void ReferenceIntegrateCustomStepKernel::execute(ContextImpl& context, CustomInt
     // Execute the step.
     
     dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
+    dynamics->setVirtualSites(extractVirtualSites(context));
     dynamics->update(context, context.getSystem().getNumParticles(), posData, velData, forceData, masses, globals, perDofValues, forcesAreValid, integrator.getConstraintTolerance());
     
     // Record changed global variables.

platforms/reference/src/ReferencePlatform.cpp

@@ -6,7 +6,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2008-2019 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2023 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -104,6 +104,7 @@ ReferencePlatform::PlatformData::PlatformData(const System& system) : time(0.0),
     periodicBoxSize = new Vec3();
     periodicBoxVectors = new Vec3[3];
     constraints = new ReferenceConstraints(system);
+    virtualSites = new ReferenceVirtualSites(system);
     energyParameterDerivatives = new map<string, double>();
 }
 
@@ -114,5 +115,6 @@ ReferencePlatform::PlatformData::~PlatformData() {
     delete periodicBoxSize;
     delete[] periodicBoxVectors;
     delete constraints;
+    delete virtualSites;
     delete energyParameterDerivatives;
 }

platforms/reference/src/SimTKReference/ReferenceBrownianDynamics.cpp

@@ -27,7 +27,6 @@
 
 #include "SimTKOpenMMUtilities.h"
 #include "ReferenceBrownianDynamics.h"
-#include "ReferenceVirtualSites.h"
 #include "openmm/OpenMMException.h"
 
 #include <cstdio>
@@ -136,6 +135,6 @@ void ReferenceBrownianDynamics::update(const OpenMM::System& system, vector<Vec3
                atomCoordinates[i][j] = xPrime[i][j];
            }
    }
-   ReferenceVirtualSites::computePositions(system, atomCoordinates);
+   getVirtualSites().computePositions(system, atomCoordinates);
    incrementTimeStep();
 }

platforms/reference/src/SimTKReference/ReferenceCustomDynamics.cpp

@@ -351,7 +351,7 @@ void ReferenceCustomDynamics::update(ContextImpl& context, int numberOfAtoms, ve
         }
         step = nextStep;
     }
-    ReferenceVirtualSites::computePositions(context.getSystem(), atomCoordinates);
+    getVirtualSites().computePositions(context.getSystem(), atomCoordinates);
     incrementTimeStep();
     recordChangedParameters(context, globals);
 }

platforms/reference/src/SimTKReference/ReferenceDynamics.cpp

@@ -27,6 +27,7 @@
 
 #include "SimTKOpenMMUtilities.h"
 #include "ReferenceDynamics.h"
+#include "openmm/OpenMMException.h"
 
 #include <cstdio>
 
@@ -45,10 +46,9 @@ using namespace OpenMM;
    --------------------------------------------------------------------------------------- */
 
 ReferenceDynamics::ReferenceDynamics(int numberOfAtoms,  double deltaT, double temperature) : 
-                  _numberOfAtoms(numberOfAtoms), _deltaT(deltaT), _temperature(temperature) {
+                  _numberOfAtoms(numberOfAtoms), _deltaT(deltaT), _temperature(temperature), virtualSites(NULL) {
 
    _timeStep = 0;
-   _ownReferenceConstraint = false;
    _referenceConstraint = NULL;
 }
 
@@ -59,9 +59,6 @@ ReferenceDynamics::ReferenceDynamics(int numberOfAtoms,  double deltaT, double t
    --------------------------------------------------------------------------------------- */
 
 ReferenceDynamics::~ReferenceDynamics() {
-   if (_ownReferenceConstraint) {
-      delete _referenceConstraint;
-   }
 }
 
 /**---------------------------------------------------------------------------------------
@@ -155,14 +152,7 @@ ReferenceConstraintAlgorithm* ReferenceDynamics::getReferenceConstraintAlgorithm
    --------------------------------------------------------------------------------------- */
 
 void ReferenceDynamics::setReferenceConstraintAlgorithm(ReferenceConstraintAlgorithm* referenceConstraint) {
-   // delete if own
-
-   if (_referenceConstraint && _ownReferenceConstraint) {
-      delete _referenceConstraint;
-   }
-
    _referenceConstraint = referenceConstraint;
-   _ownReferenceConstraint = 0;
 }
 
 /**---------------------------------------------------------------------------------------
@@ -182,3 +172,13 @@ void ReferenceDynamics::setReferenceConstraintAlgorithm(ReferenceConstraintAlgor
 void ReferenceDynamics::update(const OpenMM::System& system, vector<Vec3>& atomCoordinates,
                                vector<Vec3>& velocities, vector<Vec3>& forces, vector<double>& masses, double tolerance) {
 }
+
+const ReferenceVirtualSites& ReferenceDynamics::getVirtualSites() const {
+    if (virtualSites == NULL)
+        throw OpenMMException("ReferenceVirtualSites has not been set");
+    return *virtualSites;
+}
+
+void ReferenceDynamics::setVirtualSites(const ReferenceVirtualSites& sites) {
+    virtualSites = &sites;
+}

platforms/reference/src/SimTKReference/ReferenceLangevinMiddleDynamics.cpp

@@ -122,6 +122,6 @@ void ReferenceLangevinMiddleDynamics::update(ContextImpl& context, vector<Vec3>&
 
     updatePart3(context, numberOfAtoms, atomCoordinates, velocities, inverseMasses, xPrime);
 
-    ReferenceVirtualSites::computePositions(context.getSystem(), atomCoordinates);
+    getVirtualSites().computePositions(context.getSystem(), atomCoordinates);
     incrementTimeStep();
 }

platforms/reference/src/SimTKReference/ReferenceNoseHooverDynamics.cpp

@@ -202,7 +202,7 @@ void ReferenceNoseHooverDynamics::step2(OpenMM::ContextImpl &context, const Open
         }
     } /* end of hard wall constraint part */
 
-    ReferenceVirtualSites::computePositions(context.getSystem(), atomCoordinates);
+    getVirtualSites().computePositions(context.getSystem(), atomCoordinates);
 
     incrementTimeStep();
 }

platforms/reference/src/SimTKReference/ReferenceStochasticDynamics.cpp

@@ -189,6 +189,6 @@ void ReferenceStochasticDynamics::update(const OpenMM::System& system, vector<Ve
 
    updatePart3(numberOfAtoms, atomCoordinates, velocities, inverseMasses, xPrime);
 
-   ReferenceVirtualSites::computePositions(system, atomCoordinates);
+   getVirtualSites().computePositions(system, atomCoordinates);
    incrementTimeStep();
 }

platforms/reference/src/SimTKReference/ReferenceVariableStochasticDynamics.cpp

@@ -233,6 +233,6 @@ void ReferenceVariableStochasticDynamics::update(const OpenMM::System& system, v
        }
    }
 
-   ReferenceVirtualSites::computePositions(system, atomCoordinates);
+   getVirtualSites().computePositions(system, atomCoordinates);
    incrementTimeStep();
 }

platforms/reference/src/SimTKReference/ReferenceVariableVerletDynamics.cpp

@@ -149,7 +149,7 @@ void ReferenceVariableVerletDynamics::update(const OpenMM::System& system, vecto
                atomCoordinates[i][j] = xPrime[i][j];
            }
    }
-   ReferenceVirtualSites::computePositions(system, atomCoordinates);
+   getVirtualSites().computePositions(system, atomCoordinates);
    incrementTimeStep();
 }