CustomNonbondedForce supports computed values

[peastman]

Implements #3368. You can call addComputedValue() on a CustomNonbondedForce to add a computed value. It is computed for each particle based on the global and per-particle parameters, and can appear in the energy expression.

[zhang-ivy]

Thanks to @jchodera and @ijpulidos , I was able to build OpenMM from source! I will test this feature out next week.

[zhang-ivy]

I've successfully added computed values to my rest factory implementation and modified my custom expressions accordingly. The energy validation tests are still passing with these changes, so I think this feature is working as expected. I ran out of time today to actually test how much this feature speeds up my rest factory for context.getState() and integrator.step() (doing the latter after updating the global param every step) calls, but will do that next week and report back!

[peastman]

Thanks!

[zhang-ivy]

@peastman : Sorry for the delay -- it took me a couple iterations to make sure I was running the right tests.

Below is my code for timing getState() calls after changing a global parameter:

def run_experiment(system, state, is_rest=False):
    # Set up context
    integrator = openmm.LangevinMiddleIntegrator(300, 1.0, 0.004)
    platform = openmm.Platform.getPlatformByName("CUDA")
    platform.setPropertyDefaultValue('Precision', 'mixed')
    platform.setPropertyDefaultValue('DeterministicForces', 'true')
    context = openmm.Context(system, integrator, platform)
    context.setPeriodicBoxVectors(*state.getPeriodicBoxVectors())
    context.setPositions(state.getPositions())
    
    # Time getState()
    force_group_idx = 0
    context.getState(getEnergy=True, groups=2**force_group_idx)

    global_param = 'lambda_alchemical_sterics_new' if is_rest else 'lambda_sterics_insert'
    
    nsamples = 10
    elapsed_time = 0.0
    for i in range(nsamples):
        initial_time = time.time()
        context.setParameter(global_param, 1 - i*0.1)
        context.getState(getEnergy=True, groups=2**force_group_idx)
        elapsed_time += (time.time() - initial_time)
        time.sleep(0.1)
    print("mean: ", elapsed_time / nsamples)

    del context, integrator

And here is my code for timing step() calls after changing a global parameter:

def run_experiment_step(system, state, is_rest=False):
    # Set up context
    integrator = openmm.LangevinMiddleIntegrator(300, 1.0, 0.004)
    platform = openmm.Platform.getPlatformByName("CUDA")
    platform.setPropertyDefaultValue('Precision', 'mixed')
    platform.setPropertyDefaultValue('DeterministicForces', 'true')
    context = openmm.Context(system, integrator, platform)
    context.setPeriodicBoxVectors(*state.getPeriodicBoxVectors())
    context.setPositions(state.getPositions())
    
    global_param = 'lambda_alchemical_sterics_new' if is_rest else 'lambda_sterics_insert'
    
    integrator.step(10)
    context.getState(getPositions=True)
    elapsed_time = 0.0
    initial_time = time.time()
    for i in range(100):
        context.setParameter(global_param, 0.1*(i%2))
        integrator.step(1)
    context.getState(getPositions=True)
    elapsed_time += (time.time() - initial_time)  
    print("mean: ", elapsed_time / 100)
    
    del context, integrator

I ran these two code snippets on 5 different systems:
a. Vanilla hybrid system (generated by the default factory in perses)
b. REST hybrid system (generated by the new REST factory I've been working on)
c. REST hybrid system with real-real (non-scaled) interactions in a NonbondedForce and scaled interactions in a CustomNonbondedForce
d. REST hybrid system with custom expressions re-written to use new computed values feature
e. REST hybrid system with both c and d

Each system was tested with 3 different LRC settings:

• LRC completely off
• LRC on for real-real interactions (NonbondedForce), off for scaled interactions (CustomNonbondedForce)
• LRC completely on

As a control, I first ran these experiments WITHOUT changing the global parameter (units = seconds):

Here are the results when I do change the global parameter (units = seconds):

Takeaways:

• Moving the real-real interactions + using computed values (with LRC on for real-real and and off for scaled) is a workable solution for now
• The computed values feature offers a 2x speed up when LRC is on for getState() calls, but it is still very slow compared to when LRC is off, so we still need to figure out how to speed up LRC correction in CustomNonbondedForce (perhaps via moving LRC computation to gpu kernels).

Note that I am running these experiments with 1 thread on a GeForce GTX 1080Ti.

[peastman]

Thanks!