grab forces by name instead of hard-coded index

[mikemhenry]

Description

While working on #949 I found that we were grabbing forces by index, and it seems some changes to openmm changed the order of these forces, causing errors. See https://github.com/choderalab/perses/pull/953/files#issuecomment-1083962468

Motivation and context

Resolves #???

How has this been tested?

Change log

[codecov]

Codecov Report

Merging #976 (b1e9eb4) into main (f80b0d6) will decrease coverage by 0.00%.
The diff coverage is 63.15%.

❗ Current head b1e9eb4 differs from pull request most recent head c43dd06. Consider uploading reports for the commit c43dd06 to get more accurate results

[jchodera]

Looks good, thanks!

[jchodera]

FYI: There's now a force.getName() method that can be used instead, though this is a recent addition.
The name can even be set by users if nonstandard names are desired.

[zhang-ivy]

@mikemhenry : Can you clarify what changes in openmm caused the change in order of the forces?

[ijpulidos]

This looks good, I'm just worried that there's no guarantee that the forces names will be different. According to official doc they would default to using the class name if the name is not specified. When was this changed upstream? I tried looking but couldn't find a direct cause for it.

[mikemhenry]

@mikemhenry : Can you clarify what changes in openmm caused the change in order of the forces?

No idea what changed, but it seems clear the order can change either from where we set them or grab them so getting by name seems safer

This looks good, I'm just worried that there's no guarantee that the forces names will be different. According to official doc they would default to using the class name if the name is not specified. When was this changed upstream? I tried looking but couldn't find a direct cause for it.

I did print all the names and there are no duplicates, given this is a test, it seems it would be up to whoever is re-writing the test and adding duplicate forces names to not do that.

[mikemhenry]

RE: Our dev meeting, I'm going to see what prints out on CI to see what the order is for OpenMM 7.7 and Nightly

[mikemhenry]

Nightly the order is:

{'HarmonicBondForce': 0, 'PeriodicTorsionForce': 1, 'NonbondedForce': 2, 'HarmonicAngleForce': 3}

On openMM 7.7.0

{'HarmonicBondForce': 0, 'HarmonicAngleForce': 1, 'PeriodicTorsionForce': 2, 'NonbondedForce': 3}

I have no idea why it changes.

[zhang-ivy]

@mikemhenry : Hmm very strange. Is it worth opening an issue in openmm about this to see what may have changed in the nightly?
Either way, I think you should be good to merge this PR, since the test is written to check the vanilla hybrid system, which as its written now, only has one PeriodicTorsionForce and one NonbondedForce, so I think this PR is fine for the purposes of fixing this test.

[mikemhenry]

Looks like I had to fix another hard-coded force index. Not sure if I should fix all of them or just the ones that are giving us problems.

[zhang-ivy]

@mikemhenry : Maybe would be best to fix all of them?

[jchodera]

Let's fix anything that used a hard-coded force index.

[zhang-ivy]

The nightlys are still failing here.

________________ test_PointMutationExecutor_endstate_validation ________________
[gw1] linux -- Python 3.9.12 /usr/share/miniconda/envs/test/bin/python

    def test_PointMutationExecutor_endstate_validation():
        from pkg_resources import resource_filename
        from simtk import unit
    
        from perses.app.relative_point_mutation_setup import PointMutationExecutor
    
        pdb_filename = resource_filename("perses", "data/ala_vacuum.pdb")
>       PointMutationExecutor(
            pdb_filename,
            "1",
            "2",
            "ASP",
            ionic_strength=0.15 * unit.molar,
            flatten_torsions=False,
            flatten_exceptions=False,
            conduct_endstate_validation=True,
            generate_unmodified_hybrid_topology_factory=True,
            generate_repartitioned_hybrid_topology_factory=True,
            generate_rest_capable_hybrid_topology_factory=True
        )
perses/tests/test_relative_point_mutation_setup.py:27: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
perses/app/relative_point_mutation_setup.py:440: in __init__
    validate_endstate_energies_point(htf, endstate=endstate, minimize=True)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

htf = <perses.annihilation.relative.RESTCapableHybridTopologyFactory object at 0x7f35103c5bb0>
endstate = 0, minimize = True

    def validate_endstate_energies_point(htf, endstate=0, minimize=False):
        """
        ** Used for validating endstate energies for RESTCapableHybridTopologyFactory **
    
        Check that the hybrid system's energy (without unique old/new valence energy) matches the original system's energy for the positions in the htf.
    
        E.g. at endstate=0, the hybrid system's energy (with unique new valence terms zeroed) should match the old system's energy.
    
        Parameters
        ----------
        htf : RESTCapableHybridTopologyFactory
            the RESTCapableHybridTopologyFactory to test
        endstate : int, default=0
            the endstate to test (0 or 1)
        minimize : bool, default=False
            whether to minimize the positions before testing that the energies match
    
        """
        from perses.dispersed import feptasks
    
        # Check that endstate is 0 or 1
        assert endstate in [0, 1], "Endstate must be 0 or 1"
    
        # Get original system
        system = htf._topology_proposal.old_system if endstate == 0 else htf._topology_proposal.new_system
    
        # Get hybrid system, positions, and forces
        hybrid_system = htf.hybrid_system
        hybrid_positions = htf.hybrid_positions
    
        force_dict = {force.getName(): force for force in hybrid_system.getForces()}
        bond_force = force_dict['CustomBondForce']
        angle_force = force_dict['CustomAngleForce']
        torsion_force = force_dict['CustomTorsionForce']
        electrostatics_force = force_dict['CustomNonbondedForce_electrostatics']
        scaled_sterics_force = force_dict['CustomNonbondedForce_sterics']
        exceptions_force = force_dict['CustomBondForce_exceptions']
        reciprocal_force = force_dict['NonbondedForce_reciprocal']
        nonscaled_sterics_force = force_dict['NonbondedForce_sterics']
    
        forces = [bond_force, angle_force, torsion_force, electrostatics_force, scaled_sterics_force]
        force_names = ['bonds', 'angles', 'torsions', 'electrostatics', 'sterics']
    
        # Set global parameters for valence + electrostatics/scaled_sterics forces
        lambda_old = 1 if endstate == 0 else 0
        lambda_new = 0 if endstate == 0 else 1
        for force, name in zip(forces, force_names):
            for i in range(force.getNumGlobalParameters()):
                if force.getGlobalParameterName(i) == f'lambda_alchemical_{name}_old':
                    force.setGlobalParameterDefaultValue(i, lambda_old)
                if force.getGlobalParameterName(i) == f'lambda_alchemical_{name}_new':
                    force.setGlobalParameterDefaultValue(i, lambda_new)
    
        # Set global parameters for exceptions force
        old_parameter_names = ['lambda_alchemical_electrostatics_exceptions_old',
                               'lambda_alchemical_sterics_exceptions_old']
        new_parameter_names = ['lambda_alchemical_electrostatics_exceptions_new',
                               'lambda_alchemical_sterics_exceptions_new']
        for i in range(exceptions_force.getNumGlobalParameters()):
            if exceptions_force.getGlobalParameterName(i) in old_parameter_names:
                exceptions_force.setGlobalParameterDefaultValue(i, lambda_old)
            elif exceptions_force.getGlobalParameterName(i) in new_parameter_names:
                exceptions_force.setGlobalParameterDefaultValue(i, lambda_new)
    
        # Set global parameters for reciprocal force
        for i in range(reciprocal_force.getNumGlobalParameters()):
            if reciprocal_force.getGlobalParameterName(i) == 'lambda_alchemical_electrostatics_reciprocal':
                reciprocal_force.setGlobalParameterDefaultValue(i, lambda_new)
    
        # Zero the unique old/new valence terms at lambda = 1/0
        hybrid_to_bond_indices = htf._hybrid_to_new_bond_indices if endstate == 0 else htf._hybrid_to_old_bond_indices
        hybrid_to_angle_indices = htf._hybrid_to_new_angle_indices if endstate == 0 else htf._hybrid_to_old_angle_indices
        hybrid_to_torsion_indices = htf._hybrid_to_new_torsion_indices if endstate == 0 else htf._hybrid_to_old_torsion_indices
        for hybrid_idx, idx in hybrid_to_bond_indices.items():
            p1, p2, hybrid_params = bond_force.getBondParameters(hybrid_idx)
            hybrid_params = list(hybrid_params)
            hybrid_params[-2] *= 0  # zero K_old
            hybrid_params[-1] *= 0  # zero K_new
            bond_force.setBondParameters(hybrid_idx, p1, p2, hybrid_params)
        for hybrid_idx, idx in hybrid_to_angle_indices.items():
            p1, p2, p3, hybrid_params = angle_force.getAngleParameters(hybrid_idx)
            hybrid_params = list(hybrid_params)
            hybrid_params[-2] *= 0
            hybrid_params[-1] *= 0
            angle_force.setAngleParameters(hybrid_idx, p1, p2, p3, hybrid_params)
        for hybrid_idx, idx in hybrid_to_torsion_indices.items():
            p1, p2, p3, p4, hybrid_params = torsion_force.getTorsionParameters(hybrid_idx)
            hybrid_params = list(hybrid_params)
            hybrid_params[-2] *= 0
            hybrid_params[-1] *= 0
            torsion_force.setTorsionParameters(hybrid_idx, p1, p2, p3, p4, hybrid_params)
    
        # Get energy components of hybrid system
        thermostate_hybrid = states.ThermodynamicState(system=hybrid_system, temperature=temperature)
        integrator_hybrid = openmm.VerletIntegrator(1.0 * unit.femtosecond)
        context_hybrid = thermostate_hybrid.create_context(integrator_hybrid)
        if minimize:
            sampler_state = states.SamplerState(hybrid_positions)
            feptasks.minimize(thermostate_hybrid, sampler_state)
            hybrid_positions = sampler_state.positions
        context_hybrid.setPositions(hybrid_positions)
        components_hybrid = compute_potential_components(context_hybrid, beta=beta)
    
        # Get energy components of original system
        thermostate_other = states.ThermodynamicState(system=system, temperature=temperature)
        integrator_other = openmm.VerletIntegrator(1.0 * unit.femtosecond)
        context_other = thermostate_other.create_context(integrator_other)
        positions = htf.old_positions(hybrid_positions) if endstate == 0 else htf.new_positions(hybrid_positions)
        context_other.setPositions(positions)
        components_other = compute_potential_components(context_other, beta=beta)
    
        # Check that each of the valence force energies are concordant
        # TODO: Instead of checking with np.isclose(), check whether the ratio of differences is less than a specified energy threshold (like in validate_endstate_energies())
        for i in range(3):
            print(f"{components_other[i][0]} -- og: {components_other[i][1]}, hybrid: {components_hybrid[i][1]}")
>           assert np.isclose(components_other[i][1], components_hybrid[i][1])
E           AssertionError

[mikemhenry]

@zhang-ivy That is a different issue #993 (which I would love ideas on how to fix/what you think might be causing it, there is a TODO that looks like it could be the solution, but I wonder if there is some regression causing us to get different energies now)
So I think this one is good to go? @ijpulidos @zhang-ivy

[ijpulidos]

@mikemhenry I agree, it is a different issue. This one should be good to go.

[ijpulidos]

Looks good to me.

[zhang-ivy]

Ah thanks! I didn’t notice that issue before On May 12, 2022, at 15:51, Mike Henry ***@***.***> wrote:  Merged #976<#976> into main. — Reply to this email directly, view it on GitHub<#976 (comment)>, or unsubscribe<https://github.com/notifications/unsubscribe-auth/AIPGJCTMKQBPB7XMGPL66VLVJVOLFANCNFSM5SD4A7AQ>. You are receiving this because you were mentioned.Message ID: ***@***.***>