(Do not merge) Reference Smoothcore Implementation. #1380

proteneer · 2024-09-12T22:33:47Z

This PR designs and implements a new softcore potential, the SmoothcoreNonbondedInteractionGroup, that attempts to overcome two major failings of our existing 4D decoupling regime:

The net-charge and polarization can vary drastically along the lambda schedule. In fact even computing an effective net charge can be quite difficult due to the non-linear response of the 4D coupling parameter. This can result in rapid oscillations that require a concomitant re-orientation of the environment (eg. water inversions).
Second-order phase transition effects with large dummy groups when they're lifted "all-at-once" in 4D. Note that a sequential approach can mitigate some of this, but may still suffer from the large oscillations in net charge.

Instead, SmoothcoreNonbondedInteractionGroup takes in a list of anchoring_idxs of size num_atoms, specifying each atom's "anchoring" point. Each atom still retains the w component in the parameters, whose values are now strictly in the closed interval [0, 1] (as opposed to the old [0, cutoff]) and monotone along lambda. Here, w=0 still denotes "fully-interacting"/"fully-extended", and w=1 denotes "fully-non-interacting"/"fully-absorbed". Using the anchor_idx information in conjunction with w, a vector is drawn and used for scaling purposes to determine the virtual site location:

anchor_confs = old_conf[anchor_idxs]
rescale_factors = 1 - w_coords
new_conf = anchor_confs + rescale_factors * (old_conf - anchor_confs)

For example:

above: O is the anchoring atom for the dummy hydrogen. A scaled vector r_OH is used to position the location of the hydrogen.

This approach also makes it much easier to balance the net-charge, the all the charge sites are at full strength (but may be smushed/averaged out). At the end-states, lambda=0, and lambda=1, the resulting energies should be identical to the current SingleTopology end-states. A test is added to verify this behavior.

Todos

maxentile

Nice! Useful generalization and extension of #415

Makes sense, but a couple aspects are outside my current intuition:

[stability worry] Do any forces become unstable as the virtual sites become very close to the anchor site?
- This appears fine, but some randomized checks (e.g. running MD on random RHFE/RBFE edges with random values of w_i) would provide additional peace of mind before proceeding to the GPU implementation.
[possible efficiency enhancements] When the group's w_i's approach 1, the collapsed anchor group will look like a single particle that might make stronger-than usual interactions with environment.
- (consider if the group's |\sum_i q_i| is large..., or, even if |\sum_i q_i|== 0, but there's a large number of atoms in the group, then the collapsed site will appear sort of like a single LJ particle with a large epsilon parameter / a very deep well)
- In both cases, I suspect it may help with efficiency if an additional parameter is exposed, which scales the interaction energy by some factor <= 1, factor * U_interaction_group(x)... (Or, this could be emulated by users of this function, by scaling down the LJ epsilons by some factor like mean(epsilons) / sum(epsilons), etc.)

In the description of #415 , you mentioned this implementation is related to a suggestion from Gilson in the context of docking -- if you happen to have a specific work in mind, capturing that reference in the docstring or PR description would be helpful.

proteneer · 2024-09-16T22:12:09Z

Improved API a bit so that each row atom takes in an anchoring atom (will expand later on to groups)
Added tests for interpolation of charges and lennard jones parameters
Added test for solvated hif2a system with standard pair of hif2a mols

proteneer · 2024-09-17T02:54:47Z

Self-consistency w.r.t. Single Topology code at the end-states achieved on a real system (from hif2a) - GPU implementation tomorrow!

mcwitt · 2024-09-17T17:14:22Z

timemachine/cpp/src/kernels/k_nonbonded.cuh

+    for (int i = 0; i < COORDS_DIM; i++) {
+        RealType atom_pos = coords[atom_idx * COORDS_DIM + i];
+        RealType anchor_pos = coords[anchor_idx * COORDS_DIM + i];
+        coords_out[atom_idx * COORDS_DIM + i] = anchor_pos + rescale * (atom_pos - anchor_pos);


nit: could this be simplified to

Suggested change

coords_out[atom_idx * COORDS_DIM + i] = anchor_pos + rescale * (atom_pos - anchor_pos);

coords_out[atom_idx * COORDS_DIM + i] = atom_pos + params[atom_idx * PARAMS_DIM + 3] * (anchor_pos - atom_pos);

(removing rescale definition)?

rescale is removed by the compiler (would prefer to make the rescaling explicit).

timemachine/cpp/src/kernels/k_nonbonded.cuh

mcwitt · 2024-09-17T17:17:36Z

timemachine/cpp/src/smoothcore_nonbonded_interaction_group.cu

+template class SmoothcoreNonbondedInteractionGroup<double>;
+template class SmoothcoreNonbondedInteractionGroup<float>;
+
+} // namespace timemachine


Request: briefly describe what changed here relative to nonbonded_interaction_group.cu? Ideally, could create one commit that copies from nonbonded_interaction_group.cu and a second commit with the changes.

Best way I can think of is to do a manual diff between the two files. This turned out to be much trickier than I initially expected, so for safety reasons I decided to just make a new file for now (and eventually deprecate the old NonbondedInteractionGroup).

mcwitt · 2024-09-17T17:21:28Z

timemachine/cpp/src/smoothcore_nonbonded_interaction_group.hpp

@@ -0,0 +1,103 @@
+#pragma once


Would it be desirable to also have a smoothcore implementation of NonbondedPairListPrecomputed for intramolecular nonbonded interactions?

yes in theory - but in practice this is tricky since we need to deal with exclusions (1-2/1-3/1-4)

mcwitt · 2024-09-17T17:25:45Z

timemachine/potentials/nonbonded.py

+import itertools
+
+
+def nonbonded_block_rescaled(xi, xj, box, params_i, params_j, beta, cutoff, anchors_i, groups_i):


nit: is this currently a different interface than is exposed by the GPU potential? (don't immediately see the concept of "groups" there)

stale - added a standard API and wrapped inside potentials.py

mcwitt · 2024-09-17T17:30:17Z

timemachine/potentials/nonbonded.py

+    xi_a = xi[np.repeat(anchors_i, [len(x) for x in groups_i])]
+
+    group_idxs = jnp.array(list(itertools.chain(*groups_i)))
+
+    excluded_idxs = list(set(range(len(xi))).difference(group_idxs.tolist()))


nit: can be optimized to avoid Python loops (OK to leave for later)

proteneer · 2024-09-18T00:47:37Z

[stability worry] Do any forces become unstable as the virtual sites become very close to the anchor site?

Checking the gradients for some of these computations - I don't think so. The exact expression for the gradient computation would only ever scale down the forces, not scale up.

(consider if the group's |\sum_i q_i| is large..., or, even if |\sum_i q_i|== 0, but there's a large number of atoms in the group, then the collapsed site will appear sort of like a single LJ particle with a large epsilon parameter / a very deep well)

Correct, there's two to-be-tuned parameters describing the strength near the (0, 1) lambda boundary.

MIN_SMOOTHCORE_SIG_HALF = 0.05
MIN_SMOOTHCORE_EPS_SQRT = 0.10

edit: to elaborate a bit more, it's possible that a super-collapsed site may not be 1-step removable. Here, 1-step removable means when going from lambda=0.999 to lambda=1.0, the eps/sig parameters go from (0.05, 0.10) to the non-interacting state (0, 0) instantly. For a single LJ particle, parameterized by (0.05, 0.10), it should be 1-step removable, but when many dummy particles collapse, the net effect might be something like (0.05*num_dummy_particles, 0.10) which may no longer be one-step removal. An additional staging step where they're turned off sequentially may be necessary.

In both cases, I suspect it may help with efficiency if an additional parameter is exposed, which scales the interaction energy by some factor <= 1, factor * U_interaction_group(x)... (Or, this could be emulated by users of this function, by scaling down the LJ epsilons by some factor like mean(epsilons) / sum(epsilons), etc.)

As noted, same effect can be achieved by scaling down eps/qs via parameter interpolation

In the description of #415 , you mentioned this implementation is related to a suggestion from Gilson in the context of docking -- if you happen to have a specific work in mind, capturing that reference in the docstring or PR description would be helpful.

I need to dig this up - having a tough time finding it right now.

proteneer · 2024-09-18T02:04:19Z

Updated description.

proteneer · 2024-09-19T19:23:17Z

Rebased onto master.

Fix mypy bug WIP

proteneer requested a review from maxentile September 12, 2024 22:33

maxentile approved these changes Sep 13, 2024

View reviewed changes

mcwitt reviewed Sep 17, 2024

View reviewed changes

timemachine/cpp/src/kernels/k_nonbonded.cuh Outdated Show resolved Hide resolved

mcwitt reviewed Sep 17, 2024

View reviewed changes

proteneer mentioned this pull request Sep 18, 2024

remove extraneous summed potential now that the L-E ixn groups are merged #1384

Merged

proteneer force-pushed the smoothcores branch from cb3261b to 3d1dfe6 Compare September 19, 2024 19:22

Rebased

41a17e9

Fix mypy bug WIP

proteneer force-pushed the smoothcores branch from 99dfbf3 to 41a17e9 Compare September 20, 2024 16:56

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

(Do not merge) Reference Smoothcore Implementation. #1380

(Do not merge) Reference Smoothcore Implementation. #1380

proteneer commented Sep 12, 2024 •

edited

Loading

maxentile left a comment

proteneer commented Sep 16, 2024

proteneer commented Sep 17, 2024

mcwitt Sep 17, 2024

proteneer Sep 18, 2024

mcwitt Sep 17, 2024

proteneer Sep 18, 2024 •

edited

Loading

mcwitt Sep 17, 2024

proteneer Sep 18, 2024 •

edited

Loading

mcwitt Sep 17, 2024

proteneer Sep 18, 2024

mcwitt Sep 17, 2024

proteneer Sep 18, 2024

proteneer commented Sep 18, 2024 •

edited

Loading

proteneer commented Sep 18, 2024

proteneer commented Sep 19, 2024

	coords_out[atom_idx * COORDS_DIM + i] = anchor_pos + rescale * (atom_pos - anchor_pos);
	coords_out[atom_idx * COORDS_DIM + i] = atom_pos + params[atom_idx * PARAMS_DIM + 3] * (anchor_pos - atom_pos);

		import itertools


		def nonbonded_block_rescaled(xi, xj, box, params_i, params_j, beta, cutoff, anchors_i, groups_i):

(Do not merge) Reference Smoothcore Implementation. #1380

Are you sure you want to change the base?

(Do not merge) Reference Smoothcore Implementation. #1380

Conversation

proteneer commented Sep 12, 2024 • edited Loading

Todos

maxentile left a comment

Choose a reason for hiding this comment

proteneer commented Sep 16, 2024

proteneer commented Sep 17, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

proteneer Sep 18, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

proteneer Sep 18, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

proteneer commented Sep 18, 2024 • edited Loading

proteneer commented Sep 18, 2024

proteneer commented Sep 19, 2024

proteneer commented Sep 12, 2024 •

edited

Loading

proteneer Sep 18, 2024 •

edited

Loading

proteneer Sep 18, 2024 •

edited

Loading

proteneer commented Sep 18, 2024 •

edited

Loading