Aleaverfay/temp sparse dispatch fa elec #324

aleaverfay · 2025-10-06T18:43:53Z

FA-elec refactored for rotamer pair / sparse block pair evaluation.

Unit tests for sparse evaluation now properly working.

Torch 2.5.0, 2.5.1, and 2.7 also seem to work just fine for tmol; Torch 2.4 does not work. Also updating the code relying on cattrs.structure The newest versions of cattrs (>=24.1.0) throw errors that previously were not errors or that could be worked around or something so that now the code in tmol/chemical/restypes.py produces errors that weren't there before. One small issue: numpy.array cannot be used as a type name for attrs and this produces a weird complaint out of cattrs. numpy.array is for creating ndrarrays; ndarrays are the types, array is the function

…s table

…dd to packer tests

Nicely enough, termini variants are being properly handled already; the tests simply had to be updated.

Now the sampler is responsible for mapping their samples onto the DOFs that define the conformations, allowing samplers to define more than only the chi DOFs for the samples; e.g., if I want to build hydroxyl rotamers onto the heavy-atom coordinates of an existing SER

…sampler

This code is "used" by the on-the-fly RPE calculation code which itself is not really used

…onsidered-bts

This allows me to talk about global-block-type index (not a great name) when building rotamers for disallowed block types as will happen when a block has all of its block-types disallowed.

…tion" error

…mp_jflat_sparse_disp_torch28

…ual root in KTO

…tten PoseStack creator yet

…not matching, yikes!

Pack rotamers is now agreeing with the final energy after scoring

…two poses in a PoseStack

…s call

This includes a bug in neighbor detection that was allowing the first rotamer from pose i+1 to get checked as a neighbor with rotamers of the last residue of pose i!

…h_packing

…ispatch_packing

* Update modern GPU to CUDA13 compatibility (See moderngpu/moderngpu#56) * Add conda/mamba-installed nvtx to search path * Update NVidia conda channel to torch 2.8.0 for CUDA 13 compatibility * Add nvtx to conda packages * Fix broken unit test enable_tensor_view that lead to a "Illegal Operation" error * Change error-propogation command location in "linting" script until after conda activation

…spatch_fa_elec

jflat06 · 2025-10-06T19:12:40Z

Have you had a chance to benchmark this vs the old code?

jflat06

Seems like there's stuff from the packing branch and master mixed in with this, so I did not review that.

The elec code seems fine I think. I had some comments on the testing code.

jflat06 · 2025-10-06T23:40:06Z

tmol/kinematics/compiled/compiled.cuda.cu

      // reindexing function
      nvtx_range_push("dispatch::segscan");
      auto k_reindex = [=] MGPU_DEVICE(int index, int seg, int rank) {
+        if (nodestart + index >= nodes.size(0) || nodestart + index < 0) {


I assume this and the other unrelated stuff is just coming from master?

jflat06 · 2025-10-06T23:46:17Z

tmol/score/elec/potentials/elec_pose_score.impl.hh

+  assert(first_rot_block_type.size(0) == n_poses);
+  assert(first_rot_block_type.size(1) == max_n_blocks);
+
+  assert(block_ind_for_rot.size(0) == n_rots);
+  assert(pose_ind_for_rot.size(0) == n_rots);
+  assert(block_type_ind_for_rot.size(0) == n_rots);
+
+  assert(n_rots_for_pose.size(0) == n_poses);
+  assert(rot_offset_for_pose.size(0) == n_poses);
+
+  assert(n_rots_for_block.size(0) == n_poses);
+  assert(n_rots_for_block.size(1) == max_n_blocks);
+
+  assert(rot_offset_for_block.size(0) == n_poses);
+  assert(rot_offset_for_block.size(1) == max_n_blocks);


Since we've established a common set of parameters, I was thinking it might be good to make a function for sanity checks like this to cover checking the common parameters.

Just an idea, not something for this PR.

jflat06 · 2025-10-06T23:46:58Z

tmol/score/elec/potentials/elec_pose_score.impl.hh

    if (block_type1 < 0 || block_type2 < 0) {
      return;


We decided we didn't need this, right?

jflat06 · 2025-10-06T23:54:17Z

tmol/tests/score/common/test_energy_term.py

+    @classmethod
+    def get_block_pair_dispatch_indices(cls, nres, device):
+        i_p = []
+        i_r1 = []
+        i_r2 = []
+        for i in range(nres):
+            for j in range(nres):
+                if i <= j:
+                    i_p += [0]
+                    i_r1 += [i]
+                    i_r2 += [j]
+        return torch.tensor([i_p, i_r1, i_r2], dtype=torch.int32, device=device)


This should be unused and deleted, not sure why it's added back.

jflat06 · 2025-10-06T23:54:57Z

tmol/tests/score/common/test_energy_term.py

+    def get_whole_pose_scorer(cls, pose_stack, param_db, device):
        energy_term = cls.energy_term_class(param_db=param_db, device=device)

        for bt in pose_stack.packed_block_types.active_block_types:
            energy_term.setup_block_type(bt)
        energy_term.setup_packed_block_types(pose_stack.packed_block_types)
        energy_term.setup_poses(pose_stack)

-        return (
-            energy_term.render_whole_pose_scoring_module(pose_stack)
-            if block_pair_scoring is False
-            else energy_term.render_block_pair_scoring_module(pose_stack)
-        )
+        return energy_term.render_whole_pose_scoring_module(pose_stack)
+
+    @classmethod
+    def get_block_pair_scorer(cls, pose_stack, param_db, device):
+        energy_term = cls.energy_term_class(param_db=param_db, device=device)
+
+        for bt in pose_stack.packed_block_types.active_block_types:
+            energy_term.setup_block_type(bt)
+        energy_term.setup_packed_block_types(pose_stack.packed_block_types)
+        energy_term.setup_poses(pose_stack)
+
+        return energy_term.render_block_pair_scoring_module(pose_stack)


Do we want 2 completely separate functions for this? Seems like a bit of a waste to duplicate the code for essentially a 1 line difference.

jflat06 · 2025-10-06T23:56:29Z

tmol/tests/score/common/test_energy_term.py

+        # reshape these values for the new block-pair scoring scheme
+        gold_vals_upper_triangle = numpy.zeros_like(gold_vals)
+        assert gold_vals.shape[2] == gold_vals.shape[3]
+        n_terms = gold_vals.shape[0]
+        n_poses = gold_vals.shape[1]
+        n_res = gold_vals.shape[2]
+        assert n_poses == 1
+        inds_arange = numpy.arange(n_res, dtype=int)
+        inds_arange_i = inds_arange.reshape(-1, 1)
+        inds_arange_j = inds_arange.reshape(1, -1)
+        ij_is_upper_triangle = inds_arange_i < inds_arange_j
+        ij_is_lower_triangle = inds_arange_i > inds_arange_j
+        ij_is_diagonal = inds_arange_i == inds_arange_j
+        gold_vals_upper_triangle[:, :, ij_is_diagonal] = gold_vals[:, :, ij_is_diagonal]
+        gold_vals_upper_triangle[:, :, ij_is_upper_triangle] = 2 * gold_vals[:, :, ij_is_upper_triangle]


Mind elaborating a bit on this code? I'm not sure I follow on why this is necessary.

aleaverfay added 30 commits June 26, 2025 14:35

Progress in getting interaction graph constructed from sparse energie…

490d289

…s table

Fix multi-pose IG test case

e6d0018

Beginning the hacking on simA stuff

fd10964

Very close to getting packing compiled

89d1d36

SimA now compiles; unit tests are not working yet.

52f5d7b

First version of running (but crashing!) SimA

6d00fd0

Single-pose packing works properly on the GPU

712d8eb

Working version of multi-pose simulated annealing

3fe5a80

Create a fixture of pre-computed rotamer pair energies for 1ubq and a…

9fa7ddd

…dd to packer tests

Trim comments from the ig tests

b4f1f1c

Updating rotamer-building tests to for termini variants

3a4e705

Nicely enough, termini variants are being properly handled already; the tests simply had to be updated.

Fixing issue with offsets not reading from correct block type in chi_…

e2b56bc

…sampler

Fixing last of the rotamer-building tests

5015725

Fix torch/numpy indexing issue on device

6463404

Removing unused rotamer-bounding-sphere caclulation code

1f19f47

This code is "used" by the on-the-fly RPE calculation code which itself is not really used

Saving progress on (currently buggy) refactoring of allowed-bts --> c…

1a781e5

…onsidered-bts

Complete refactoring to track considered but disallowed block types

f3b0e0c

This allows me to talk about global-block-type index (not a great name) when building rotamers for disallowed block types as will happen when a block has all of its block-types disallowed.

Fix a unit test

2a50d24

Add IncludeCurrentRotamer sampler and tests

e118cf5

Fix PackerTask unit test

a7d85df

Save progress w/ fixing bugs in IncludeCurrentSampler

5935024

Update modern GPU to CUDA13 compatibility

8061fae

Add conda/mamba-installed nvtx to search path

1502d77

Update to torch 2.8 / CUDA 13 compatibility

6d6116b

Small adjustment to env.yaml

80372ec

Fix broken unit test enable_tensor_view that lead to a "Illegal Opera…

c42b840

…tion" error

Merge aleaverfay/update_to_torch2p8 into jflat's sparse_dispatch branch

eb553d9

Merge branch 'aleaverfay/multi_pose_stack_packing' into aleaverfay/te…

943e395

…mp_jflat_sparse_disp_torch28

aleaverfay added 21 commits September 29, 2025 14:53

Fix off-by-one error in IncludeCurrentSampler due to presence of virt…

894fc24

…ual root in KTO

Running version of pack_rotamers with LJ + Hbond energy; have not wri…

5ff5e36

…tten PoseStack creator yet

Add utility to create PoseStack for final rotamer assignment; scores …

09e3655

…not matching, yikes!

One-body energies have to be segregated from two-body energies;

2e3908c

Pack rotamers is now agreeing with the final energy after scoring

Fix bug in chain deduction that shows up when dealing with more than …

ca9e1aa

…two poses in a PoseStack

Add extra-chi-samples option to PackerTask; add test for pack_rotamer…

0c3191f

…s call

linting

b59a4aa

linting

92a9fdd

Reverting change to linux-cuda pinned versions for testing server

ff63e02

Reverting change to linux-cuda pinned versions for testing server

8e3b365

Try to install on nvidia/label/cuda-12.8.0 instead of cuda-12.8

ee602f9

Fix several errors in handling jagged PoseStack packing

f0aa5cd

This includes a bug in neighbor detection that was allowing the first rotamer from pose i+1 to get checked as a neighbor with rotamers of the last residue of pose i!

Beef up the pack_rotamers unit test

8b90004

Merge branch 'jflat06/sparse-dispatch' into aleaverfay/sparse_dispatc…

c7b2659

…h_packing

Merge branch 'aleaverfay/update_to_torch2p8' into aleaverfay/sparse_d…

de7dd36

…ispatch_packing

Merge branch 'master' into aleaverfay/sparse_dispatch_packing

3524f42

Add macos build

0a6f659

Significant progress in rotamre energy evaluation for fa_elec

4709c74

Fix base class EnergyTermTestBase for new scoring scheme; works for LJLK

8fb1cd4

FAElec now working in all unit tests

4542da5

aleaverfay requested a review from jflat06 October 6, 2025 18:43

Merge branch 'jflat06/sparse-dispatch' into aleaverfay/temp_sparse_di…

77ca4b6

…spatch_fa_elec

Turn on fa_elec in the packer unit test

0576fe3

jflat06 reviewed Oct 6, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Aleaverfay/temp sparse dispatch fa elec #324

Aleaverfay/temp sparse dispatch fa elec #324

Uh oh!

aleaverfay commented Oct 6, 2025

Uh oh!

jflat06 commented Oct 6, 2025

Uh oh!

jflat06 left a comment

Uh oh!

jflat06 Oct 6, 2025

Uh oh!

jflat06 Oct 6, 2025

Uh oh!

jflat06 Oct 6, 2025

Uh oh!

jflat06 Oct 6, 2025

Uh oh!

jflat06 Oct 6, 2025

Uh oh!

jflat06 Oct 6, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Aleaverfay/temp sparse dispatch fa elec #324

Are you sure you want to change the base?

Aleaverfay/temp sparse dispatch fa elec #324

Uh oh!

Conversation

aleaverfay commented Oct 6, 2025

Uh oh!

jflat06 commented Oct 6, 2025

Uh oh!

jflat06 left a comment

Choose a reason for hiding this comment

Uh oh!

jflat06 Oct 6, 2025

Choose a reason for hiding this comment

Uh oh!

jflat06 Oct 6, 2025

Choose a reason for hiding this comment

Uh oh!

jflat06 Oct 6, 2025

Choose a reason for hiding this comment

Uh oh!

jflat06 Oct 6, 2025

Choose a reason for hiding this comment

Uh oh!

jflat06 Oct 6, 2025

Choose a reason for hiding this comment

Uh oh!

jflat06 Oct 6, 2025

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants