Distance scaling in ExpNorm RBFs

[nec4]

I have noticed that the ExpNorm basis functions might benefit from a new parameter that rescales the distances inside the inner expoential. Consider the default setup where the distances are not rescaled and an upper cutoff of 35 angstroms is used (the first figure). While there is a good density of basis functions at small distances, there are just a few at longer distances. By contrast, if the distances are rescaled by some factor, say alpha=0.2. the second figure results in a more even distribution of basis functions (without changing anything else). I think this might be at odds with how the distance expansions are implemented in the TorchMD_GN class, because it always expects a fixed number of parameters common to all distance expansion types. If this is desirable, I would be happy to make a PR and implement tests.

[PhilippThoelke]

This would be a nice addition! Is there a way to compute a good value for alpha from the cutoff range? Do you have an idea how to cleanly implement this additional argument for ExpNorm? It might make sense then to create the RBF module inside create_model then instead of the model's __init__ to avoid redundancy.

[nec4]

I think that is actually a problem. I'm not sure of an intuitive way to choose alpha without visually checking the plots of the basis functions over distances. To me, it makes sense to leave the default value as 1 for consistency with respect to previous works like PhysNet, but I would have to think a bit about how to choose alpha systematically based on the choice of the cutoffs.

[PhilippThoelke]

For now we could also add a special rbf_arg in the argparse of train.py and pass that to the different RBF expansion classes in the constructor. Then each class can handle it individually or just ignore it if it is not needed. This would be similar to how we are passing arguments to the dataset classes right now.

We have this argument:
https://github.com/compsciencelab/torchmd-net/blob/1759652a89c82dc90671a028316b015e1a27a2d5/scripts/train.py#L60
And it is passed to all dataset classes by default:
https://github.com/compsciencelab/torchmd-net/blob/1759652a89c82dc90671a028316b015e1a27a2d5/torchmdnet/data.py#L32-L35

[nec4]

Hello - sorry for the delay on this. Turns out there is a very simple way to tune $\alpha$: It is actually just $\alpha=5/(cutoff_upper-cutoff_lower)$. The 5 is chosen to reproduce the default behavior from PhysNet. Here are a few examples where only the cutoff distances are varied:

I will make a branch and then a PR and if you guys like it we can discuss further.

[giadefa]

Nice

…

On Sun, Jul 18, 2021, 21:02 nec4 ***@***.***> wrote: Hello - sorry for the delay on this. Turns out there is a very simple way to tune $\alpha$: It is actually just $\alpha=5/(cutoff_upper-cutoff_lower)$. The 5 is chosen to reproduce the default behavior from PhysNet. Here are a few examples where only the upper cutoff distance is varied: [image: Screen Shot 2021-07-18 at 9 00 57 PM] <https://user-images.githubusercontent.com/42926839/126079207-b6b7c3e3-bd80-4147-b91e-665f8988bca1.png> I will make a branch and then a PR and if you guys like it we can discuss further. — You are receiving this because you are subscribed to this thread. Reply to this email directly, view it on GitHub <#28 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AB3KUOUWMFRA7PMFLRGJOSDTYMQK3ANCNFSM47ZLL2BQ> .

[PhilippThoelke]

Looks good, thanks!