Fast batched nonbonded on interaction groups, as fxn of ligand charges and epsilons #685

maxentile · 2022-04-01T03:56:11Z

Improved Jax utilities for processing batches of snapshots, with the aim of computing a batch of nonbonded energies as an efficient differentiable function of ligand nonbonded parameters.

Main feature:

Utilities for precomputing small summary vectors of a stored trajectory, so that ligand-environment nonbonded interaction group energies as a function of ligand charges (or ligand epsilons) can be computed as dot products with the small summary vectors. This is directly inspired by the linear basis function approach developed by Levi Naden and Michael Shirts in https://pubmed.ncbi.nlm.nih.gov/26580188/ .

(Note: These utilities make reweighting virtually a no-op, cost-wise... Rough timings are ~5 microseconds to compute a reweighted hydration free energy in explicit solvent (using Jax in double-precision CPU mode)...)

Additional refactoring:

Refactor jax nonbonded functions to use pairs ([N, 2] array) instead of a pair of inds_l, inds_r (2 [N,] arrays)
Extract function for converting nonbonded exceptions/exclusions to rescale masks from tests.common: ~~nonbonded.convert_exceptions_to_rescale_masks~~ nonbonded.convert_exclusions_to_rescale_masks
Extract function to get pairs_from_interaction_groups
Move unused jax function rescale_coordinates to attic
- (unsure if broken?)
~~jax.ops.index_update(x, idx, y) -> x.at[idx].set(y), x.at[idx].multiply(0.5)~~
- This update was already applied in Replaces index_update with .at().set() #672 , but in resolving merge conflicts I tweaked these slightly to use named slices (params.at[:,0].multiply(0.5) -> params.at[charge_indices].multiply(0.5))
np.indices -> np.meshgrid (following Add NonbondedInteractionGroup potential, rename existing nonbonded potentials #578 (comment) )
Change the return type of reference nonbonded interaction groups (no longer returns pairs)
Add functions to document Lennard-Jones combining rules

Reviewer note: This branch "vmap-jax-nonbonded-ig" is a scope-reduced version of "vmap-jax-nonbonded" (omitting complexity of handling intramolecular exceptions / exclusions, etc., for inclusion in a later PR after deboggling week).

Utility function to find atom pairs that depend on ligand and contribute to nonbonded sum for each conf in a batch of confs...

Addressing #578 (comment)

Following #578 (comment)

Possibly save some debugging time in the future...

Completes 852fe29

* compute_batch_nonbonded_exceptions * get_intramolecular_nb_pairs_and_prefactors * construct_batch_nonbonded_hl_as_fxn_of_ligand_nb_params * construct_batch_nonbonded_ll_as_fxn_of_ligand_nb_params

Addressing #685 (comment)

Applying suggestion from #685 (comment)

Addressing #685 (comment)

Applying suggestion in #685 (comment)

badisa · 2022-04-04T14:11:46Z

tests/common.py

-    for (i, j), exc in zip(exclusion_idxs, scales[:, 1]):
-        lj_rescale_mask[i][j] = 1 - exc
-        lj_rescale_mask[j][i] = 1 - exc
+    charge_rescale_mask, lj_rescale_mask = nonbonded.convert_exceptions_to_rescale_masks(exclusion_idxs, scales, N)


Point of clarification, is an exception an exclusion?

Ahh, this should be "exclusions" -- will rename

(I thought "exceptions" are pairs that need to be handled differently, "exclusions" are pairs that are omitted entirely (special case of exceptions), but this is not actually the convention.)

Renamed in c4077ae

badisa · 2022-04-04T14:14:27Z

timemachine/fe/topology.py

@@ -546,8 +547,12 @@ def parameterize_nonbonded(self, ff_q_params, ff_lj_params):
        qlj_params, nb_potential = super().parameterize_nonbonded(ff_q_params, ff_lj_params)

        # halve the strength of the charge and the epsilon parameters
-        src_qlj_params = jnp.asarray(qlj_params).at[:, 0].multiply(0.5)
-        src_qlj_params = jnp.asarray(src_qlj_params).at[:, 2].multiply(0.5)
+        charge_indices = jax.numpy.index_exp[:, 0]


Is there a reason jax.numpy was used here instead of jnp?

Nope -- good catch!

~Also is there a reason to use the index_exp syntax over the .at syntax? ~

edit: Is there a reason to use index_exp with .at rather than do the indexing directly? Just to reuse indices and make it more explicit?

Just to be more explicit -- should have identical behavior

Partially addressed in 4becd4a

badisa · 2022-04-04T14:18:26Z

timemachine/potentials/jax_utils.py

@@ -1,34 +1,30 @@
-from typing import Tuple
-
-import jax
 import jax.numpy as np
 import numpy as onp
 from jax import vmap

 Array = onp.array


nit: Since you are changing the usage, might be good to switch this to onp.typing.NDArray?

Updated in 78ac174

Addressing #685 (comment)

badisa · 2022-04-04T14:24:16Z

timemachine/potentials/nonbonded.py

+    traj: [T, N, D] array
+    boxes: diagonal [T, D, D] array (or [T] array of Nones)
+    charges: [N] array
+    ligand_indices: [N_lig] array of ints


Are there any limitations we want to place on ligand_indices and env_indices? Should they be disjoint?

They should be disjoint (throw error if not), and ligand_indices should be smaller than env_indices (print warning if not) -- will fix

Partially addressed in 762fc61 (and applied to nonbonded_v3_interaction_groups as well)

Addressing #685 (comment)

* Partially address #685 (comment) * Also validate input for nonbonded_v3_interaction_groups

badisa · 2022-04-04T15:05:50Z

tests/test_jax_nonbonded.py

+    reweight_ref = jit(make_reweighter(u_batch_ref))
+    reweight_test = jit(make_reweighter(u_batch_test))
+
+    for _ in range(50):


This test takes about 50 seconds on CI, does this need 50 iterations to be correct?

Ahh, on Friday it was ~6 seconds -- will track down which change made it more expensive and reduce iterations if needed

Locally ~80-90% of the iteration time is in computing reference gradients

timemachine/tests/test_jax_nonbonded.py

Line 463 in 762fc61

v_ref, g_ref = value_and_grad(reweight_ref, argnums=argnum)(eps_ligand, q_ligand)

, added in response to #685 (comment) .

Will rearrange this slightly and reduce the number of iterations from 50 to 5...

Addressed in 930ae0d

(in initial version of test the iterations were ~ 0.03 seconds each, but in current version the iterations are ~0.3 seconds each, so it's good to reduce the number of iterations!)

Addressing #685 (comment)

maxentile · 2022-04-04T15:27:01Z

Thanks for the reviews! Enabling auto-merge

restores functionality from #685 overwritten in #931

* restore lj_eps_prefactor special case restores functionality from #685 overwritten in #931 * add lj eps branch in test

maxentile added 30 commits February 17, 2022 18:41

Add get_ligand_dependent_indices_batch

0ef46c5

Utility function to find atom pairs that depend on ligand and contribute to nonbonded sum for each conf in a batch of confs...

Categorize jax_utils.py contents

da4770b

Move unused rescale_coordinates to attic

8f225c2

Refactor (inds_i, inds_j) --> pairs, improve docs

213752b

Addressing #578 (comment)

Use np.meshgrid instead of np.indices

1df2908

Following #578 (comment)

Convert np.triu_indices return

be6c17c

Move inverted exclusion format definition from tests to nonbonded

852fe29

Possibly save some debugging time in the future...

Missed file in last commit

9a24c2a

Completes 852fe29

Typo

853c7c6

Add batch nonbonded utilities

00cc7c6

Merge branch 'master' into jax_reference_reweighting

3108f8c

Replace get_group_group_indices with pairs_from_interaction_groups

2150b05

index_update(x, idx, y) -> x.at(idx).set(y)

72e20a1

index_update(x, idx, 0.5 * x[idx]) -> x.at(idx).multiply(0.5)

4e785cc

Rename batched_neighbor_inds -> get_interacting_pair_indices_batch

9f8e673

Reorder input arguments in less surprising way

db95c01

Improve shape annotations, add todo note

1ea420d

Fix typo: x.at(idx) -> x.at[idx]

bf2e9af

Add batch nonbonded HL, LL functions

c1bef0e

* compute_batch_nonbonded_exceptions * get_intramolecular_nb_pairs_and_prefactors * construct_batch_nonbonded_hl_as_fxn_of_ligand_nb_params * construct_batch_nonbonded_ll_as_fxn_of_ligand_nb_params

Typo: repeated ligand_indices

a99b2dd

Typo: missing argument

5df7fbb

Merge branch 'master' into vmap-jax-nonbonded

66359ff

Merge branch 'master' into vmap-jax-nonbonded

bb5c728

Merge branch 'master' into vmap-jax-nonbonded

374bf8a

Merge branch 'master' into vmap-jax-nonbonded

cf02e90

Merge branch 'master' into vmap-jax-nonbonded

02b5945

Revert accidental change to nonbonded_v3_interaction_groups

fce09dd

Use convert_exceptions_to_rescale_masks in enhanced.py

23ad660

Merge branch 'master' into vmap-jax-nonbonded

2b4513c

Remove list of contents

9786f21

maxentile added 6 commits April 1, 2022 15:56

Improve docs, extract repeated distance utility

b940ab9

Addressing #685 (comment)

Axis keyword

3fbc030

Applying suggestion from #685 (comment)

jax.ops.index -> jax.numpy.index_exp

a7f3e2c

Addressing #685 (comment)

Clarify comment about pairs

4cdf070

Addressing #685 (comment)

Clarify index expression

0ed8175

Applying suggestion in #685 (comment)

Merge branch 'master' into vmap-jax-nonbonded-ig

ec937b7

badisa reviewed Apr 4, 2022

View reviewed changes

maxentile added 2 commits April 4, 2022 10:21

jax.numpy -> jnp

4becd4a

Addressing #685 (comment)

Rename exceptions -> exclusions

c4077ae

Addressing #685 (comment)

badisa reviewed Apr 4, 2022

View reviewed changes

maxentile added 2 commits April 4, 2022 10:33

onp.array -> onp.typing.NDArray

78ac174

Addressing #685 (comment)

Add validate_interaction_group_idxs

762fc61

* Partially address #685 (comment) * Also validate input for nonbonded_v3_interaction_groups

badisa reviewed Apr 4, 2022

View reviewed changes

Reduce random test iterations, set random seed

930ae0d

Addressing #685 (comment)

badisa approved these changes Apr 4, 2022

View reviewed changes

maxentile enabled auto-merge (squash) April 4, 2022 15:28

maxentile merged commit 0c4b610 into master Apr 4, 2022

proteneer added the cr_scipy label Apr 11, 2022

maxentile deleted the vmap-jax-nonbonded-ig branch June 15, 2022 13:04

maxentile mentioned this pull request Nov 18, 2022

Image Frames from RBFE #905

Merged

maxentile mentioned this pull request Jan 6, 2023

Expand linear basis to support variable LJ sigma parameters #931

Merged

This was referenced Mar 6, 2023

[WIP] Update reaction field #969

Closed

Constructing potential energy functions for many systems on same device #682

Closed

maxentile mentioned this pull request Dec 4, 2023

[WIP] Prefactor refactor #1210

Closed

6 tasks

maxentile added a commit that referenced this pull request Dec 5, 2023

restore lj_eps_prefactor special case

c5fd3f3

restores functionality from #685 overwritten in #931

maxentile mentioned this pull request Dec 5, 2023

Lj eps prefactor #1212

Merged

maxentile added a commit that referenced this pull request Dec 5, 2023

Lj eps prefactor (#1212)

7fa0ed5

* restore lj_eps_prefactor special case restores functionality from #685 overwritten in #931 * add lj eps branch in test

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fast batched nonbonded on interaction groups, as fxn of ligand charges and epsilons #685

Fast batched nonbonded on interaction groups, as fxn of ligand charges and epsilons #685

maxentile commented Apr 1, 2022 •

edited

Loading

badisa Apr 4, 2022

maxentile Apr 4, 2022

maxentile Apr 4, 2022

badisa Apr 4, 2022

maxentile Apr 4, 2022

badisa Apr 4, 2022 •

edited

Loading

maxentile Apr 4, 2022

maxentile Apr 4, 2022

badisa Apr 4, 2022

maxentile Apr 4, 2022

badisa Apr 4, 2022

maxentile Apr 4, 2022

maxentile Apr 4, 2022

badisa Apr 4, 2022

maxentile Apr 4, 2022

maxentile Apr 4, 2022 •

edited

Loading

maxentile Apr 4, 2022

maxentile commented Apr 4, 2022

Fast batched nonbonded on interaction groups, as fxn of ligand charges and epsilons #685

Fast batched nonbonded on interaction groups, as fxn of ligand charges and epsilons #685

Conversation

maxentile commented Apr 1, 2022 • 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

badisa Apr 4, 2022 • 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

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

maxentile Apr 4, 2022 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

maxentile commented Apr 4, 2022

maxentile commented Apr 1, 2022 •

edited

Loading

badisa Apr 4, 2022 •

edited

Loading

maxentile Apr 4, 2022 •

edited

Loading