Implemented tabulated functions in CUDA platform

libraries/lepton/include/lepton/Operation.h

@@ -211,6 +211,9 @@ public:
         return function->evaluate(args);
     }
     ExpressionTreeNode differentiate(const std::vector<ExpressionTreeNode>& children, const std::vector<ExpressionTreeNode>& childDerivs, const std::string& variable) const;
+    const std::vector<int>& getDerivOrder() const {
+        return derivOrder;
+    }
 private:
     std::string name;
     CustomFunction* function;

platforms/cuda/src/CudaKernels.cpp

@@ -423,6 +423,17 @@ void CudaCalcCustomNonbondedForceKernel::initialize(const System& system, const
     for (int i = 0; i < numExceptions; i++)
         force.getExceptionParameters(exceptions[i], exceptionParticle1[i], exceptionParticle2[i], exceptionParams[i]);
 
+    // Record the tabulated functions.
+
+    for (int i = 0; i < force.getNumFunctions(); i++) {
+        string name;
+        vector<double> values;
+        double min, max;
+        bool interpolating;
+        force.getFunctionParameters(i, name, values, min, max, interpolating);
+        gpuSetTabulatedFunction(gpu, i, name, values, min, max, interpolating);
+    }
+
     // Record information for the expressions.
 
     vector<string> paramNames;

platforms/cuda/src/kernels/cudatypes.h

@@ -232,6 +232,7 @@ static const int GT2XX_RANDOM_THREADS_PER_BLOCK         = 384;
 static const int G8X_NONBOND_WORKUNITS_PER_SM           = 220;
 static const int GT2XX_NONBOND_WORKUNITS_PER_SM         = 256;
 static const unsigned int MAX_STACK_SIZE = 8;
+static const unsigned int MAX_TABULATED_FUNCTIONS = 8;
 
 static const float PI = 3.14159265358979323846f;
 
@@ -244,7 +245,7 @@ enum CudaNonbondedMethod
 };
 
 enum ExpressionOp {
-    CONSTANT = 0, VARIABLE0, VARIABLE1, VARIABLE2, VARIABLE3, VARIABLE4, VARIABLE5, VARIABLE6, VARIABLE7, VARIABLE8, GLOBAL, CUSTOM, ADD, SUBTRACT, MULTIPLY, DIVIDE,
+    CONSTANT = 0, VARIABLE0, VARIABLE1, VARIABLE2, VARIABLE3, VARIABLE4, VARIABLE5, VARIABLE6, VARIABLE7, VARIABLE8, GLOBAL, CUSTOM, CUSTOM_DERIV, ADD, SUBTRACT, MULTIPLY, DIVIDE,
         POWER, NEGATE, SQRT, EXP, LOG, SIN, COS, SEC, CSC, TAN, COT, ASIN, ACOS, ATAN, SQUARE, CUBE, RECIPROCAL, ADD_CONSTANT, MULTIPLY_CONSTANT, POWER_CONSTANT
 };
 
@@ -350,6 +351,8 @@ struct cudaGmxSimulation {
     float4*         pCustomExceptionParams;         // Parameters for custom nonbonded exceptions
     unsigned int    customExceptions;               // Number of custom nonbonded exceptions
     unsigned int    customParameters;               // Number of parameters for custom nonbonded interactions
+    float4*         pTabulatedFunctionCoefficients[MAX_TABULATED_FUNCTIONS]; // The spline coefficients for each tabulated function
+    float4*         pTabulatedFunctionParams;       // The min, max, and spacing for each tabulated function
     float2*         pEwaldCosSinSum;                // Pointer to the cos/sin sums (ewald)
     unsigned int    bonds;                          // Number of bonds
     int4*           pBondID;                        // Bond atom and output buffer IDs

platforms/cuda/src/kernels/gpu.cpp

@@ -133,7 +133,7 @@ static const float BOLTZ                    =    (RGAS / KILO);            // (k
 #define DUMP_PARAMETERS 0
 
 template <int SIZE>
-static Expression<SIZE> createExpression(const string& expression, const Lepton::ExpressionProgram& program, const vector<string>& variables,
+static Expression<SIZE> createExpression(gpuContext gpu, const string& expression, const Lepton::ExpressionProgram& program, const vector<string>& variables,
         const vector<string>& globalParamNames, unsigned int& maxStackSize) {
     Expression<SIZE> exp;
     if (program.getNumOperations() > SIZE)
@@ -177,6 +177,14 @@ static Expression<SIZE> createExpression(const string& expression, const Lepton:
                     exp.arg[i] = j;
                 }
                 break;
+            case Operation::CUSTOM:
+                exp.op[i] = dynamic_cast<const Operation::Custom*>(&op)->getDerivOrder()[0] == 0 ? CUSTOM : CUSTOM_DERIV;
+                for (int j = 0; j < MAX_TABULATED_FUNCTIONS; j++)
+                    if (op.getName() == gpu->tabulatedFunctions[j].name) {
+                        exp.arg[i] = j;
+                        break;
+                    }
+                break;
             case Operation::ADD:
                 exp.op[i] = ADD;
                 break;
@@ -573,6 +581,52 @@ void gpuSetNonbondedCutoff(gpuContext gpu, float cutoffDistance, float solventDi
     gpu->sim.reactionFieldC = (1.0f / cutoffDistance)*(3.0f*solventDielectric)/(2.0f*solventDielectric+1.0f);
 }
 
+extern "C"
+void gpuSetTabulatedFunction(gpuContext gpu, int index, const string& name, const vector<double>& values, double min, double max, bool interpolating)
+{
+    if (index < 0 || index >= MAX_TABULATED_FUNCTIONS) {
+        stringstream str;
+        str << "Only " << MAX_TABULATED_FUNCTIONS << " tabulated functions are supported";
+        throw OpenMMException(str.str());
+    }
+    if (gpu->tabulatedFunctions[index].coefficients != NULL)
+        delete gpu->tabulatedFunctions[index].coefficients;
+    CUDAStream<float4>* coeff = new CUDAStream<float4>((int) values.size()-1, 1, "TabulatedFunction");
+    gpu->tabulatedFunctions[index].coefficients = coeff;
+    gpu->sim.pTabulatedFunctionCoefficients[index] = coeff->_pDevData;
+    gpu->tabulatedFunctions[index].name = name;
+    gpu->tabulatedFunctions[index].min = min;
+    gpu->tabulatedFunctions[index].max = max;
+
+    // First create a padded set of function values.
+
+    vector<double> padded(values.size()+2);
+    padded[0] = 2*values[0]-values[1];
+    for (int i = 0; i < (int) values.size(); i++)
+        padded[i+1] = values[i];
+    padded[padded.size()-1] = 2*values[values.size()-1]-values[values.size()-2];
+
+    // Now compute the spline coefficients.
+
+    for (int i = 0; i < (int) values.size()-1; i++) {
+        float4 c;
+        if (interpolating) {
+            c.x = padded[i+1];
+            c.y = 0.5*(-padded[i]+padded[i+2]);
+            c.z = 0.5*(2.0*padded[i]-5.0*padded[i+1]+4.0*padded[i+2]-padded[i+3]);
+            c.w = 0.5*(-padded[i]+3.0*padded[i+1]-3.0*padded[i+2]+padded[i+3]);
+        }
+        else {
+            c.x = (padded[i]+4.0*padded[i+1]+padded[i+2])/6.0;
+            c.y = (-3.0*padded[i]+3.0*padded[i+2])/6.0;
+            c.z = (3.0*padded[i]-6.0*padded[i+1]+3.0*padded[i+2])/6.0;
+            c.w = (-padded[i]+3.0*padded[i+1]-3.0*padded[i+2]+padded[i+3])/6.0;
+        }
+        (*coeff)[i] = c;
+    }
+    coeff->Upload();
+}
+
 extern "C"
 void gpuSetCustomNonbondedParameters(gpuContext gpu, const vector<vector<double> >& parameters, const vector<vector<int> >& exclusions,
             const vector<int>& exceptionAtom1, const vector<int>& exceptionAtom2, const vector<vector<double> >& exceptionParams,
@@ -630,6 +684,39 @@ void gpuSetCustomNonbondedParameters(gpuContext gpu, const vector<vector<double>
     gpu->psCustomExceptionID->Upload();
     gpu->psCustomExceptionParams->Upload();
 
+    // This class serves as a placeholder for custom functions in expressions.
+
+    class FunctionPlaceholder : public Lepton::CustomFunction {
+    public:
+        int getNumArguments() const {
+            return 1;
+        }
+        double evaluate(const double* arguments) const {
+            return 0.0;
+        }
+        double evaluateDerivative(const double* arguments, const int* derivOrder) const {
+            return 0.0;
+        }
+        CustomFunction* clone() const {
+            return new FunctionPlaceholder();
+        }
+    };
+
+    // Record the tabulated functions, which were previously set with calls to gpuSetTabulatedFunction().
+
+    FunctionPlaceholder* fp = new FunctionPlaceholder();
+    map<string, Lepton::CustomFunction*> functions;
+    gpu->psTabulatedFunctionParams = new CUDAStream<float4>(MAX_TABULATED_FUNCTIONS, 1, "TabulatedFunctionRange");
+    gpu->sim.pTabulatedFunctionParams = gpu->psTabulatedFunctionParams->_pDevData;
+    for (int i = 0; i < MAX_TABULATED_FUNCTIONS; i++) {
+        gpuTabulatedFunction& func = gpu->tabulatedFunctions[i];
+        if (func.coefficients != NULL) {
+            (*gpu->psTabulatedFunctionParams)[i] = make_float4(func.min, func.max, func.coefficients->_length/(func.max-func.min), 0.0f);
+            functions[func.name] = fp;
+        }
+    }
+    gpu->psTabulatedFunctionParams->Upload();
+
     // Create the Expressions.
 
     vector<string> variables;
@@ -637,8 +724,8 @@ void gpuSetCustomNonbondedParameters(gpuContext gpu, const vector<vector<double>
     for (int i = 0; i < paramNames.size(); i++)
         variables.push_back(paramNames[i]);
     gpu->sim.customExpressionStackSize = 0;
-    SetCustomNonbondedEnergyExpression(createExpression<128>(energyExp, Lepton::Parser::parse(energyExp).optimize().createProgram(), variables, globalParamNames, gpu->sim.customExpressionStackSize));
-    SetCustomNonbondedForceExpression(createExpression<128>(energyExp, Lepton::Parser::parse(energyExp).differentiate("r").optimize().createProgram(), variables, globalParamNames, gpu->sim.customExpressionStackSize));
+    SetCustomNonbondedEnergyExpression(createExpression<128>(gpu, energyExp, Lepton::Parser::parse(energyExp, functions).optimize().createProgram(), variables, globalParamNames, gpu->sim.customExpressionStackSize));
+    SetCustomNonbondedForceExpression(createExpression<128>(gpu, energyExp, Lepton::Parser::parse(energyExp, functions).differentiate("r").optimize().createProgram(), variables, globalParamNames, gpu->sim.customExpressionStackSize));
     Expression<64> paramExpressions[4];
     vector<string> combiningRuleParams;
     combiningRuleParams.push_back("");
@@ -652,8 +739,9 @@ void gpuSetCustomNonbondedParameters(gpuContext gpu, const vector<vector<double>
             combiningRuleParams.push_back("");
     }
     for (int i = 0; i < paramNames.size(); i++)
-        paramExpressions[i] = createExpression<64>(combiningRules[i], Lepton::Parser::parse(combiningRules[i]).optimize().createProgram(), combiningRuleParams, globalParamNames, gpu->sim.customExpressionStackSize);
+        paramExpressions[i] = createExpression<64>(gpu, combiningRules[i], Lepton::Parser::parse(combiningRules[i], functions).optimize().createProgram(), combiningRuleParams, globalParamNames, gpu->sim.customExpressionStackSize);
     SetCustomNonbondedCombiningRules(paramExpressions);
+    delete fp;
 }
 
 extern "C"
@@ -1522,6 +1610,9 @@ void* gpuInit(int numAtoms, unsigned int device, bool useBlockingSync)
     gpu->psCcmaReducedMass          = NULL;
     gpu->psConstraintMatrixColumn   = NULL;
     gpu->psConstraintMatrixValue    = NULL;
+    gpu->psTabulatedFunctionParams  = NULL;
+    for (int i = 0; i < MAX_TABULATED_FUNCTIONS; i++)
+        gpu->tabulatedFunctions[i].coefficients = NULL;
 
     // Initialize output buffer before reading parameters
     gpu->pOutputBufferCounter       = new unsigned int[gpu->sim.paddedNumberOfAtoms];
@@ -1705,6 +1796,10 @@ void gpuShutDown(gpuContext gpu)
     delete gpu->psCcmaReducedMass;
     delete gpu->psConstraintMatrixColumn;
     delete gpu->psConstraintMatrixValue;
+    delete gpu->psTabulatedFunctionParams;
+    for (int i = 0; i < MAX_TABULATED_FUNCTIONS; i++)
+        if (gpu->tabulatedFunctions[i].coefficients != NULL)
+            delete gpu->tabulatedFunctions[i].coefficients;
     if (gpu->cudpp != 0)
         cudppDestroyPlan(gpu->cudpp);
 

platforms/cuda/src/kernels/gputypes.h

@@ -42,6 +42,12 @@ struct gpuMoleculeGroup {
     std::vector<int> instances;
 };
 
+struct gpuTabulatedFunction {
+    std::string name;
+    double min, max;
+    CUDAStream<float4>* coefficients;
+};
+
 enum SM_VERSION
 {
     SM_10,
@@ -67,6 +73,7 @@ struct _gpuContext {
     std::vector<std::vector<int> > exclusions;
     unsigned char* pAtomSymbol;
     std::vector<gpuMoleculeGroup> moleculeGroups;
+    gpuTabulatedFunction tabulatedFunctions[MAX_TABULATED_FUNCTIONS];
     float iterations;
     float epsfac;
     float solventDielectric;
@@ -92,6 +99,7 @@ struct _gpuContext {
     CUDAStream<float4>* psCustomParams;     // Atom parameters for custom nonbonded force
     CUDAStream<int4>* psCustomExceptionID;  // Atom indices for custom nonbonded exceptions
     CUDAStream<float4>* psCustomExceptionParams; // Parameters for custom nonbonded exceptions
+    CUDAStream<float4>* psTabulatedFunctionParams; // The min, max, and spacing for each tabulated function
     CUDAStream<float2>* psEwaldCosSinSum;
     CUDAStream<float2>* psObcData; 
     CUDAStream<float>* psObcChain;
@@ -182,6 +190,9 @@ void gpuSetCoulombParameters(gpuContext gpu, float epsfac, const std::vector<int
 extern "C"
 void gpuSetNonbondedCutoff(gpuContext gpu, float cutoffDistance, float solventDielectric);
 
+extern "C"
+void gpuSetTabulatedFunction(gpuContext gpu, int index, const std::string& name, const std::vector<double>& values, double min, double max, bool interpolating);
+
 extern "C"
 void gpuSetCustomNonbondedParameters(gpuContext gpu, const std::vector<std::vector<double> >& parameters, const std::vector<std::vector<int> >& exclusions,
             const std::vector<int>& exceptionAtom1, const std::vector<int>& exceptionAtom2, const std::vector<std::vector<double> >& exceptionParams,

platforms/cuda/src/kernels/kCalculateCustomNonbondedForces.cu

@@ -148,6 +148,23 @@ __device__ float kEvaluateExpression_kernel(Expression<SIZE>* expression, float*
                     else if (op == GLOBAL) {
                         STACK(++stackPointer) = globalParams[(int) expression->arg[i]];
                     }
+                    else if (op == CUSTOM || op == CUSTOM_DERIV) {
+                        int function = (int) expression->arg[i];
+                        float x = STACK(stackPointer);
+                        float4 params = cSim.pTabulatedFunctionParams[function];
+                        if (x < params.x || x > params.y)
+                            STACK(stackPointer) = 0.0f;
+                        else
+                        {
+                            int index = floor((x-params.x)*params.z);
+                            float4 coeff = cSim.pTabulatedFunctionCoefficients[function][index];
+                            x = (x-params.x)*params.z-index;
+                            if (op == CUSTOM)
+                                STACK(stackPointer) = coeff.x+x*(coeff.y+x*(coeff.z+x*coeff.w));
+                            else
+                                STACK(stackPointer) = (coeff.y+x*(2.0f*coeff.z+x*3.0f*coeff.w))*params.z;
+                        }
+                    }
                     else if (op == ADD) {
                         float temp = STACK(stackPointer);
                         STACK(--stackPointer) += temp;

platforms/cuda/tests/TestCudaCustomNonbondedForce.cpp

@@ -202,6 +202,38 @@ void testPeriodic() {
     ASSERT_EQUAL_TOL(1.9+1+0.9, state.getPotentialEnergy(), TOL);
 }
 
+void testTabulatedFunction(bool interpolating) {
+    CudaPlatform platform;
+    System system;
+    system.addParticle(1.0);
+    system.addParticle(1.0);
+    VerletIntegrator integrator(0.01);
+    CustomNonbondedForce* forceField = new CustomNonbondedForce("fn(r)+1");
+    forceField->addParticle(vector<double>());
+    forceField->addParticle(vector<double>());
+    vector<double> table;
+    for (int i = 0; i < 21; i++)
+        table.push_back(std::sin(0.25*i));
+    forceField->addFunction("fn", table, 1.0, 6.0, interpolating);
+    system.addForce(forceField);
+    Context context(system, integrator, platform);
+    vector<Vec3> positions(2);
+    positions[0] = Vec3(0, 0, 0);
+    double tol = 0.01;
+    for (int i = 1; i < 30; i++) {
+        double x = (7.0/30.0)*i;
+        positions[1] = Vec3(x, 0, 0);
+        context.setPositions(positions);
+        State state = context.getState(State::Forces | State::Energy);
+        const vector<Vec3>& forces = state.getForces();
+        double force = (x < 1.0 || x > 6.0 ? 0.0 : -std::cos(x-1.0));
+        double energy = (x < 1.0 || x > 6.0 ? 0.0 : std::sin(x-1.0))+1.0;
+        ASSERT_EQUAL_VEC(Vec3(-force, 0, 0), forces[0], 0.1);
+        ASSERT_EQUAL_VEC(Vec3(force, 0, 0), forces[1], 0.1);
+        ASSERT_EQUAL_TOL(energy, state.getPotentialEnergy(), 0.02);
+    }
+}
+
 int main() {
     try {
         testSimpleExpression();
@@ -209,6 +241,8 @@ int main() {
         testExceptions();
         testCutoff();
         testPeriodic();
+        testTabulatedFunction(true);
+        testTabulatedFunction(false);
     }
     catch(const exception& e) {
         cout << "exception: " << e.what() << endl;

platforms/reference/src/ReferenceKernels.cpp

@@ -487,8 +487,8 @@ public:
      */
     void findCoefficients(double& x, double* coeff) const {
         int length = values.size();
-        double spacing = (max-min)/(length-1);
-        int index = std::floor((x-min)/spacing);
+        double scale = (length-1)/(max-min);
+        int index = std::floor((x-min)*scale);
         double points[4];
         points[0] = (index == 0 ? 2*values[0]-values[1] : values[index-1]);
         points[1] = values[index];
@@ -506,7 +506,7 @@ public:
             coeff[2] = (3.0*points[0]-6.0*points[1]+3.0*points[2])/6.0;
             coeff[3] = (-points[0]+3.0*points[1]-3.0*points[2]+points[3])/6.0;
         }
-        x = (x-min)/spacing-index;
+        x = (x-min)*scale-index;
     }
     double evaluate(const double* arguments) const {
         double x = arguments[0];
@@ -523,7 +523,7 @@ public:
         double coeff[4];
         findCoefficients(x, coeff);
         double scale = (values.size()-1)/(max-min);
-        return scale*(coeff[1]+x*(2.0*coeff[2]+x*3.0*coeff[3]));
+        return scale*(coeff[1]+x*(2.0*coeff[2]+x*3.0*coeff[3])); // We assume a first derivative, because that's the only order ever used by CustomNonbondedForce.
     }
     CustomFunction* clone() const {
         return new TabulatedFunction(min, max, values, interpolating);