search

coarse-graining / cgnet

learning coarse-grained force fields
BSD 3-Clause "New" or "Revised" License
57 stars 26 forks source link

Avg model sim #179

Closed nec4 closed 4 years ago

nec4 commented 4 years ago

Development:

Heyo! Here is the idea I had for averaging models for simulations. The basic idea is based on the following:

  1. Model averaging is done by averaging the predicted forces/potential energies over a series of input models
  2. Simulation steps (either overdamped or normal Langevin) are then carried out using these averaged forces.

To accomplish this, I made the following changes/implementations to the Simulation class:

  1. Extraction of potential, forces from the model at each time step now has a dedicated method called calculate_potential_and_forces(). This is for ease in overriding with a method that can process (eg, average) the potential energies and forces before use in the calculation of new positions/velocities in the next simulation step.
  2. The checks for models in training/eval mode has been moved to a specific method called _input_model_check(). This method also checks for embeddings if the model has a SchnetFeature in its architecture.

After these changes, I made a new class, MultiModelSimulation, that inherits from the Simulation class:

  1. This class is instanced with a list of models, rather than just a single model. This list is stored as the attribute self.models.
  2. This class overrides the _check_input_model which individually performs the same model checks as in the Simulation class, but instead carries these out individually over each model
  3. This class overrides the calculate_potential_and_forces() method such that the potential and forces are averaged over all the models provided.

I have also made a changea to the HarmonicPotential class in test_simulation.py: Previously, the model output a zero/meaningless potential, but for the purpose of making sure that the potential and forces are averaged correctly in MultiModelSimulation, I have given the output potential a physical value. I have written one test (open to writing more!) that tests the ability to average the forces and potentials in MultiModelSimulation.

Some further ideas/potential problems (not an exhaustive list - feel free to add/discuss!):

  1. Is it important to serially mount the models on the GPU versus simultaneously mounting them? This may be important to consider if all the models may not simultaneously fit on the GPU.
nec4 commented 4 years ago

I was thinking that it would be less verbose to simply point the users to the original documentation of simulation kwargs. Although, on a second look, this would may not be as useful if a user is calling help(MultiModelSimulation). Therefore, I have copied the full documentation over - see 87810fa. I have left out the information on the integration scheme, but I have put in a note at the top of the MultiModelSimulation docs that points users to help(cgnet.network.Simulation) if they wish to see this information. Let me know what you think!

nec4 commented 4 years ago

I added some new documentation/comments/notes in light of your suggestions - let me know what you think, and if there is anything to improve/change!

brookehus commented 4 years ago

more comments - let me know if they're unclear. can you write a test that compares a Simulation with one model and a MultiModelSimulation with a list of one model? with a random seed they should come out the same, right?

nec4 commented 4 years ago

Thanks for the suggestions! I added a new test in ff1a437 that tests the equality of simulation outputs for a single model used in both Simulation and MultiModelSimulation. The test only passes for me on my local machine if n_sims=1, but it passes every time for any random for n_sims on Clementi machines. Not sure what's going on - I'm gonna try a clean miniconda install and see if it helps.

nec4 commented 4 years ago

I found a mistake I made in the averaging dimension for MultiModelSImulation (it should have been 0 instead of 1) - this was only due to writing the second test - Thank you for the suggestion! See f12a115 for details.

nec4 commented 4 years ago

Great suggestions and great ideas for input_model_checks. I have addressed the comments above, and will now work on pep8. Let me know if you see anything problematic or something that can be improved!

nec4 commented 4 years ago

Thanks for the additional suggestions! I think I have addressed all of the latest comments. Let me know if you see anything else that can be improved/changed!