Exact PME electrostatics treatment and fast computation of reduced potentials

[andrrizzi]

Fix #319.

I've run some timeit tests, and reading the charges from the CustomNonbondedForce is actually pretty fast so I don't think implementing the caching system for the charges is worth it. With

import timeit

def read_charges(system):
    original_charges_force = None
    for force in system.getForces():
        if (isinstance(force, openmm.CustomNonbondedForce) and
                    force.getEnergyFunction() == '0.0;'):
            original_charges_force = force

    # Set alchemical atoms charges.
    _, alchemical_atoms = original_charges_force.getInteractionGroupParameters(0)
    for atom_idx in alchemical_atoms:
        original_charge = original_charges_force.getParticleParameters(atom_idx)[0]

test_system = testsystems.SrcExplicit()
factory = AbsoluteAlchemicalFactory(alchemical_pme_treatment='exact')

# Alchemically modify first 150 atoms of Src.
alchemical_region = AlchemicalRegion(alchemical_atoms=range(150))
alchemical_system = factory.create_alchemical_system(test_system.system, alchemical_region)
globals = {
    'read_charges': read_charges,
    'alchemical_system': alchemical_system,
}
number = 1000
tot_time = timeit.timeit('read_charges(alchemical_system)', globals=globals, number=number)
print('read_charges: time for {} executions: {} s (time per execution {} s)'.format(number, tot_time, tot_time / number))

I got on the cluster (CUDA platform)

read_charges: time for 1000 executions: 0.13602294865995646 s (time per execution 0.00013602294865995646 s)

[andrrizzi]

Tests pass finally.

[andrrizzi]

I'm going to push here also the implementation of #321 while this is waiting for review because I need everything in one place.

[andrrizzi]

I've added another couple of features/optimizations.

• There's a new option AlchemicalState.update_alchemical_charges that makes changes in electrostatics with exact PME treatment incompatible. This will allow us to simulate electrostatics on multiple contexts and avoid calling updateParametersInContext.
• SamplerState now strips the unit only once to positions and velocities on apply_to_context and cache them internally. The cache is invalidated if positions and velocities are changed. I've added the utility class TrackedQuantity for this purpose.

[Lnaden]

This looks good to me, I have a couple places where I thought of a "what if X happens," but they are corner cases. Do you know if the tests should pass at this point, or are they broken pending other fixes?

[Lnaden]

What happens here if _UPDATE_ALCHEMICAL_CHARGES_PARAMETER is False?

[andrrizzi]

_UPDATE_ALCHEMICAL_CHARGES_PARAMETER is a module-level constant so it should be always a string.

[Lnaden]

Is it possible to fall back on not using force groups if there are not enough available?

[andrrizzi]

You can set split_alchemical_forces to False if there are not enough force groups available. Otherwise, RuntimeError is raised.

[Lnaden]

Right, but I mean can we just default to fall back to non-split_alchemical_forces if there are too many forces to split into the 32 groups?

[andrrizzi]

Ah! I see, sorry. I went for the exception because in general, I would prefer to know if there's a performance penalty instead of proceeding silently, but I can change this behavior if you think we should automatically fall back to split_alchemical_forces=False.

I think it'd be very easy for an application to go with something like this anyway:

try:
    factory.create_alchemical_system(...)
except RuntimeError:
    factory.split_alchemical_forces = False
    factory.create_alchemical_system(...)

[Lnaden]

I don't understand the need for the change do key-less dict here?

[andrrizzi]

Each alchemically_modify_X() method now has a different signature. The key of the returned dict is the parameter that controls that force so that we know upstream in which force group to put it.

[Lnaden]

What happens if someone specifies an exact force and not PME?

[andrrizzi]

As long as reference_force is a NonbondedForce, the method should behave correctly. We call this only inside _alchemically_modify_NonbondedForce.

[Lnaden]

What happened to the docstring here?

[andrrizzi]

potential_energy and kinetic_energy are properties now so I moved their docstrings in the property definitions.

[andrrizzi]

Thanks for the review!

I just noticed that in AlchemicalState.__setstate__ we read the update_alchemical_charges attribute from the dict containing the serialization. I'll make that optional (that defaults to True) or it'll break backwards compatibility of our YANK datasets.