coarse-graining / cgnet

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

PhysNet style RBF #128

Closed nec4 closed 5 years ago

nec4 commented 5 years ago

Development:

Hey guys - here is an implementation of the PhysNet telescoping radial basis functions. I can't think of a good name for the class, so I called it TelescopingRBF() (definitely open to better names). Basically, the centers are clustered uniformly but a modulation envelope has the effect of creating more higher resolution functions at smaller distances and lower resolution functions at larger distances. As an aside, acceptable training losses could only be reached after I extended the distance cutoff to something larger than 5.0 angstroms. Currently, the default values for n_gaussians and cutoff are the same as those from the PhysNet paper (64, and 10 angstroms respectively). In addition, I made some slight modifications to SchnetFeature to include basis function type logic and device passing to TelescopingRBF if used. Tests are still needed.

Feel free to commit changes directly here. I am training an ALA2 CB model using this expansion basis overnight - last I looked the test loss was pretty good compared to our last results.

Screen Shot 2019-11-15 at 2 23 57 AM

.

Dom1L commented 5 years ago

Thanks for the great work! I will take a look :)

nec4 commented 5 years ago

Ok guys - I added a test based on @Dom1L 's test_radial_basis_function() in test_feature_utils.py and addressed some of the suggestions above. The test is failing, but the values between the manual and TelescopingRBF computations are not that far off (?). I'm not sure what is going on to cause that. Other than that, I fixed a minus sign error in TelescopingRBF.forward(), where I was originally missing a minus sign in the exponential expansion of the distances inside the argument of the gaussian function.

nec4 commented 5 years ago

Oops - forgot to push the last commit - that one changes mod to modulation_envelope in TelescopingRBF.modulation().

brookehus commented 5 years ago

I addressed your test failure in #129. I'm a little sketched out by the test, because it's basically the exact reimplementation of the method, but I'm also not sure how to do it any differently. Are there any end cases, like a cutoff of zero or a huge number, that can be tested for certain properties?

nec4 commented 5 years ago

Ok - I merged @brookehus 's changes and implementation of a tolerance parameter to get ride of very small values beyond the tail end of the telescoping RBFs. If I set the tolerance=1e-8 in the telescoping rbf instanced in test_telescoping_rbf(), the test passes most of the time. What do you guys think?

nec4 commented 5 years ago

A cutoff of zero would give n_gaussian copies of the same basis function, each centered at the same spot. Not sure if there is meaningful info in having the cutoff be a huge number - the range of basis functions will just be scaled accordingly and depending on n_gaussian there may be some appreciable gaps between sequential basis functions.

nec4 commented 5 years ago

Update: added a small test for the centers using a zero cutoff.

brookehus commented 5 years ago

@coarse-graining/developers team, we gotta move on this. What's the status? Review ready?

nec4 commented 5 years ago

Review is ready.

brookehus commented 5 years ago

More comments. -1 for giving this a merge candidate label before pep8 compliance! But great work. :)

nec4 commented 5 years ago

Okay - I think I addressed the above comments. Let me know if you think of or see any more!

nec4 commented 5 years ago

Thanks for the review - I have no problem changing the name (in fact I think its better). We just need to keep this in mind for models that have been trained using the old name. Because there are no trainable parameters in this module, we can always load the old models and replace instances of TelescopingRBF with ModulatedRBF and save them as new models.

brookehus commented 5 years ago

There is a meta-issue here which is that we are doing "production" with a branch not merged to master! And as discussed on slack, that means that we have some stuff running with the previous name TelescopingRBF. For now, decide if you want to:

(1) deal with it by hand. In this case, keep the physnet_rbf branch alive. This will get messy when it needs to start importing from other branches if we add more stuff, but I think we can handle it.

or

(2) keep the telescoping option for SchnetFeature's basis_function_type, and add back in a TelescopingRBF method that inherits everything from ModulatedRBF and prints a DeprecationWarning. At some point in the future, when everything is free from calling/using TelescopingRBF, we remove it.

In this case I don't mind either way, and it's up to you if you want to try your hand at deprecation infrastructure or if you want to just expedite merging this in.

What I do mind is running production before a merge to master! It's true that a lot of our comments here are nit-picking on the documentation, but obviously reviews can find logical errors in the actual code too! if someone is about to run production on a branch that's not master, we need to communicate more clearly about what PR's are holding it back, and prioritize getting them done.

nec4 commented 5 years ago

We can always rerun things because we are not deep into production - no worries. I think it's easier to deal with this issue by hand in the limited number of models affected.