Specifying hybrid topology factory (htf) from input file

[ijpulidos]

Description

Created a common API point to handle the Hybrid Topology Factory creation. It can handle the 'vanilla' HybridTopologyFactory and the new RESTCapableHybridTopologyFactory.

The new parameter for the input yaml file is hybrid_topology_factory and it expects the full name of the class as a value. It raises an error if you provide an invalid name.

Motivation and context

Resolves #899

How has this been tested?

Tested locally by running the protein-ligand example from cli.

Change log

[ijpulidos]

@zhang-ivy Can you help me figure out why I am getting the following error? This happens when I try to run the protein-ligand benchmark code specifying the RESTCapableHybridTopologyFactory.

Traceback (most recent call last):
  File "/home/user/miniconda3/envs/perses-dev/bin/perses-relative", line 33, in <module>
    sys.exit(load_entry_point('perses', 'console_scripts', 'perses-relative')())
  File "/home/user/workdir/repos/perses/perses/app/setup_relative_calculation.py", line 800, in run
    setup_dict = run_setup(setup_options)
  File "/home/user/workdir/repos/perses/perses/app/setup_relative_calculation.py", line 614, in run_setup
    zero_state_error, one_state_error = validate_endstate_energies(_top_prop, _htf, _forward_added_valence_energy, _reverse_subtracted_valence_energy, beta = 1.0/(kB*temperature), ENERGY_THRESHOLD = ENERGY_THRESHOLD)#, trajectory_directory=f'{xml_directory}{phase}')
  File "/home/user/workdir/repos/perses/perses/tests/utils.py", line 796, in validate_endstate_energies
    nonalch_zero, nonalch_one, alch_zero, alch_one = generate_endpoint_thermodynamic_states(hybrid_system, top_proposal, repartitioned_endstate)
  File "/home/user/workdir/repos/perses/perses/tests/utils.py", line 387, in generate_endpoint_thermodynamic_states
    check_system(system)
  File "/home/user/workdir/repos/perses/perses/tests/utils.py", line 346, in check_system
    force = forces['PeriodicTorsionForce']
KeyError: 'PeriodicTorsionForce'

As the error says it is expecting that key here, but the forces dictionary doesn't have it.

[zhang-ivy]

@ijpulidos : The error arises because validate_endstate_energies should not be called on the RESTCapableHybridTopologyFactory, it was written for the HybridTopologyFactory. Here is the logic for how to run energy validation depending on the factory selected (that I use for protein mutations): https://github.com/choderalab/perses/blob/main/perses/app/relative_point_mutation_setup.py#L401-L440
You'll likely need to adapt the small molecule setup pipeline to use something similar. Note that energy validation should only be done if conduct_endstate_validation is True and flatten_torsions/flatten_exceptions are both False

[zhang-ivy]

@ijpulidos : I think you are missing a bunch of HybridTopologyFactory and RESTCapableHybridTopologyFactory arguments: https://github.com/choderalab/perses/blob/main/perses/app/relative_point_mutation_setup.py#L444-L460.

[zhang-ivy]

Most of the ones in the code I linked are for the HybridTopologyFactory. The only ones that are used for the RESTCapableHybridTopologyFactory are: rest_radius=rest_radius, w_scale=w_scale -- you'll want to make sure the yaml has options for specifying these.

[ijpulidos]

Yes, you are correct.

[ijpulidos]

@zhang-ivy You are correct, but I'm using what used to be in the setup_relative_calculation.py script originally here and here. I agree, we are not specifying all of the parameters, using many of the defaults. But I don't want to change the logic here, unless this has been missing since the beginning and it's something we want to change.

[ijpulidos]

I had to manually add the new option/key to the fah_generator.py in order for tests to pass. This is less than ideal but I think we want this if we want to move forward with this PR without having to do a big refactor that will eventually be overridden when we change things to the new Settings object from the OpenFE folks.

[codecov]

Codecov Report

Merging #988 (08bc9d3) into main (6aedcd0) will decrease coverage by 0.01%.
The diff coverage is 67.21%.

[ijpulidos]

I had to manually add the new option/key to the fah_generator.py in order for tests to pass. This is less than ideal but I think we want this if we want to move forward with this PR without having to do a big refactor that will eventually be overridden when we change things to the new Settings object from the OpenFE folks.

Related #995

[zhang-ivy]

When validating the RESTCapableHybridTopologyFactory, we need to run validate_endstate_energies_point twice -- once at each endstate. We don't want to run the test twice on the same htf though, since the test involves changing hybrid system parameters. I recommend making a copy of the htf within this for loop, like I do here

[ijpulidos]

As discussed, it's probably better to make the copy inside the validate_endstate_energies_point. This will help with not having to remember the function alters the htf and also should be better in terms of memory management. We would need to change this as well for validate_endstate_energies_md and make the corresponding changes in relative_point_mutation_setup.py and test_relative.py. I'll do those, it should be an easy fix.

[zhang-ivy]

I think a more useful error message here would be something like f"You specified an unsupported factory type: {factory_name}. Currently, the supported factories are: HybridTopologyFactory and RESTCapableHybridTopologyFactory."

[zhang-ivy]

One more thing that is no longer necessary due to your suggestion to make a copy of htf in the validate_endstate_energies_point/md() functions is saving the htf as a pickle here

All the pickle code can be removed and we can just call:

for endstate in [0, 1]:
    validate_endstate_energies_point(htf, endstate=endstate, minimize=True)

The pickle code is only in there because I wanted to run the validate function on fresh htfs each time (aka the same reason why I was making copies of the htfs).