Constraints #305

jflat06 · 2024-07-23T17:15:42Z

You can create a constraint set and add constraints to it with:

add_constraints(your_constraint_function, atom_indices, parameters)

atom_indices: a tensor that is (constraint_ind x atom_num x (pose_ind, res_ind, atom_ind). You must supply at least 2 and at most 4 atoms in the atom_num dimension.

parameters: a (constraint_ind x 6) float tensor that you can use to store parameters on a per-constraint basis.

your_constraint_function: this must take two args, (atoms, params). atoms will be filled with a (n_constraints x 4 x 3) tensor containing atom coordinates. If you supplied less than 4 atom indices when you created the constraints, you must not use any extra atoms (they will be filled with garbage). parameters will be the parameter tensor you supplied when you added the constraints.

You can call add_constraints multiple times and they will be concatenated.

An example:

    def constfn(atoms, params):
        return params[:, 0]

    cnstr_atoms = torch.full((2, 2, 3), 0, dtype=torch.int32, device=torch_device)
    cnstr_params = torch.full((2, 6), 0, dtype=torch.float32, device=torch_device)

    cnstr_atoms[0,0] = torch.tensor([0,0,0]) # constraint 0, atom 0 - assigned to atom at pose 0, res 0, atom 0
    cnstr_atoms[0,1] = torch.tensor([0,1,1]) # constraint 0, atom 1 - assigned to atom at pose 0, res 1, atom 1
    cnstr_params[0, 0] = 2 # constraint 0, parameter 0 - assigned to 2

    cnstr_atoms[1,0] = torch.tensor([1,0,0])
    cnstr_atoms[1,1] = torch.tensor([1,1,1])
    cnstr_params[1, 0] = 4

    constraints.add_constraints(constfn, cnstr_atoms, cnstr_params)

    cnstr_atoms[0,0] = torch.tensor([0,1,0])
    cnstr_atoms[0,1] = torch.tensor([0,2,1])
    cnstr_params[0, 0] = 20

    cnstr_atoms[1,0] = torch.tensor([1,0,0])
    cnstr_atoms[1,1] = torch.tensor([1,2,1])
    cnstr_params[1, 0] = 40

    constraints.add_constraints(constfn, cnstr_atoms, cnstr_params)

This returns:

[[[[ 0.  1.  0.  0.  0.]
   [ 1.  0. 10.  0.  0.]
   [ 0. 10.  0.  0.  0.]
   [ 0.  0.  0.  0.  0.]
   [ 0.  0.  0.  0.  0.]]

  [[ 0.  2. 20.  0.  0.]
   [ 2.  0.  0.  0.  0.]
   [20.  0.  0.  0.  0.]
   [ 0.  0.  0.  0.  0.]
   [ 0.  0.  0.  0.  0.]]]]

Right now, the score for a constraint is statically split between the residues of the first and second atoms. You could utilize this convention to manipulate where the scores end up by ordering the atoms in your input atom index tensor appropriately. I'm not sure if that's ideal, but may be OK for now.

jflat06 · 2024-09-13T23:40:41Z

tmol/score/constraint/constraint_energy_term.py

+        return tmol.score.terms.constraint_creator.ConstraintTermCreator.score_types()
+
+    def n_bodies(self):
+        return 1


Not totally sure what to do here, since this could involve up to 4 'bodies'. But as I understand it, this is mostly for dispatch reasons, which isn't a concern in this case.

jflat06 · 2024-09-13T23:41:20Z

tmol/score/constraint/constraint_energy_term.py

+    def bounded(cls, atoms, params):
+        lb = params[:, 0]
+        ub = params[:, 1]
+        sd = params[:, 2]
+        rswitch = params[:, 3]
+
+        atoms1 = atoms[:, 0]
+        atoms2 = atoms[:, 1]
+        dist = torch.linalg.norm(atoms1 - atoms2, dim=-1)
+
+        ret = torch.full_like(dist, 0)
+
+        ub2 = ub + 0.5 * sd
+
+        g0 = dist < lb
+        # g1 = torch.logical_and(lb <= dist, dist <= ub) # default 0
+        g2 = torch.logical_and(ub < dist, dist <= ub2)
+        g3 = dist > ub2
+
+        ret[g0] = ((lb - dist) / sd) ** 2
+        # ret[g1] = 0
+        ret[g2] = ((dist - ub) / sd) ** 2
+        ret[g3] = 2 * rswitch * ((dist - ub) / sd) - rswitch**2
+
+        return ret


This logic could use some checking (or a baseline comparison vs Rosetta)

Look at electrostatics for an example of sliding atom for testing.

jflat06 · 2024-09-13T23:41:37Z

tmol/score/constraint/constraint_energy_term.py

+    @classmethod
+    def circularharmonic_angle(cls, atoms, params):
+        return cls.circularharmonic(atoms, params)
+
+    @classmethod
+    def circularharmonic_torsion(cls, atoms, params):
+        return cls.circularharmonic(atoms, params, torsion=True)
+
+    @classmethod
+    def circularharmonic(cls, atoms, params, torsion=False):


Is there a cleaner way of breaking this up?

Remove the non-torsion version.

jflat06 · 2024-09-13T23:45:45Z

tmol/score/constraint/constraint_whole_pose_module.py

+            flattened.index_add_(0, indices1, scores / 2)
+            flattened.index_add_(0, indices2, scores / 2)


Right now this splits the score between the residues of the first and second atoms that were provided.

This makes it a bit awkward because it means that a torsion has to be specified in (1)(4)(2)(3) order rather than (1)(2)(3)(4) order if you want to have the score split between them.

I could change it to be the 1st and 4th atoms instead, but this would mean that someone writing a simple distance constraint would have to provide a tensor with 4 atoms (and then only use the 1st and 4th in their function). I'm not sure what the best solution is to this.

Copy last atom into slot 4.

jflat06 · 2024-09-13T23:46:14Z

tmol/score/score_types.py

@@ -7,6 +7,7 @@ class ScoreType(AutoNumber):
    cart_torsions = ()
    cart_impropers = ()
    cart_hxltorsions = ()
+    constraint = ()


All under a single score term right now.

Save this for later.

jflat06 · 2024-09-13T23:47:17Z

tmol/tests/score/constraint/test_constraint_energy_term.py

+    cnstr_params[0, 0] = 1.47  # lb TODO
+    cnstr_params[0, 1] = 1.47  # ub TODO
+    cnstr_params[0, 2] = 1.47  # sd TODO
+    cnstr_params[0, 3] = 1.47  # rswitch TODO


Some examples of a realistic bounded constraint would be very appreciated.

something like 1,3,1,1
sweep 0-8

jflat06 · 2024-09-13T23:52:27Z

tmol/pose/constraint_set.py

+    def add_constraints_to_all_poses(self, fn, nposes, atom_indices, params=None):
+        self.add_constraints(fn, atom_indices, params, nposes=nposes)
+
+    def add_constraints(self, fn, atom_indices, params=None, nposes=0):


Right now adding constraints to all poses is done with a different function. This was originally done by simply omitting the 'pose' index from the atom specifiers. I'm not sure which is preferable.

CSTs live in posestack
Deep copy on pose copy
None if no csts

jflat06 · 2024-09-13T23:55:10Z

tmol/pose/constraint_set.py

+    def add_constraints_to_all_poses(self, fn, nposes, atom_indices, params=None):
+        self.add_constraints(fn, atom_indices, params, nposes=nposes)
+
+    def add_constraints(self, fn, atom_indices, params=None, nposes=0):


We have to pass in the number of poses in order to replicate a copy of each constraint to all poses, since this constraint_set doesn't have a reference to the pose.

I'm thinking more and more that these functions should be going through pose stack (and then maybe managed internally in the pose_stack by a constraint_set class?). The constraints are fundamentally attached to the pose_stack anyway, especially if you get the atom names from the pose stack's block types like in the test cases.

jflat06 · 2024-09-13T23:57:05Z

@fdimaio

This should be ready for review. I've tagged a number of notes/questions in the code for some feedback.

Also, I took a look at the tests failing and it seems like something might be up with the testing server:

⚠️ Warning: Checkout failed! Error running /usr/bin/git remote set-url origin https://github.com/uw-ipd/tmol: exit status 128 (Attempt 3/3 Retrying in 2s)
🚨 Error: Error running /usr/bin/git remote set-url origin https://github.com/uw-ipd/tmol: exit status 128

codecov · 2024-10-17T23:19:55Z

Codecov Report

Attention: Patch coverage is 97.98271% with 7 lines in your changes missing coverage. Please review.

Project coverage is 94.78%. Comparing base (2cc46cd) to head (de4e99f).
Report is 6 commits behind head on master.

Files with missing lines	Patch %	Lines
tmol/score/constraint/constraint_energy_term.py	95.58%	3 Missing ⚠️
tmol/score/terms/constraint_creator.py	85.71%	2 Missing ⚠️
...l/score/constraint/constraint_whole_pose_module.py	97.91%	1 Missing ⚠️
tmol/tests/score/common/test_energy_term.py	93.33%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #305      +/-   ##
==========================================
+ Coverage   94.71%   94.78%   +0.06%     
==========================================
  Files         299      305       +6     
  Lines       20415    20769     +354     
==========================================
+ Hits        19337    19685     +348     
- Misses       1078     1084       +6

Flag	Coverage Δ
_shrug_Testing_CPU	`90.06% <97.98%> (+0.14%)`	⬆️
_shrug_Testing_CPU_debug_w_o_jit	`92.27% <97.98%> (+0.10%)`	⬆️
_shrug_Testing_CUDA	`92.55% <97.98%> (+0.09%)`	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

fdimaio

lgtm (pending tests)

… constraint functions. Adding functionality to the testing infrastructure to allow for editing the pose_stack from the parent test case (needed for adding the constraint set)

…alk between tests. Switch the score attribution to use the first and last provided atoms for each constraint. Adding test to score a range of distances for the bounded test. Adding functionality to test_block_pair_scoring base function so that it can be re-used by other tests. Updating bounded test parameters. Fixing bug in bounded test that broke when there was more than 1 being evaluated.

* Fix env build issue * take 2 * take 3 * Fix benchmarking bug

jflat06 commented Sep 13, 2024

View reviewed changes

fdimaio approved these changes Oct 31, 2024

View reviewed changes

jflat06 added 19 commits December 11, 2024 15:39

Initial commit for constraints

c419bf2

Use lambdas instead.

93e7531

constraint basics working.

0a290c9

lots of progress

dd6b697

progress

a2cde61

almost done

3e649cb

almost done

49d52f6

Adding constraint energy function with the ability to specify pytorch…

0351971

… constraint functions. Adding functionality to the testing infrastructure to allow for editing the pose_stack from the parent test case (needed for adding the constraint set)

Adding constraint_set class, turning off update_baselines.

0963547

comments

21201eb

removing extraneous includes

5569cc1

Adding private hint to constraint set

b2ba118

linting

42f7197

fixing bug in tests for the constraint energy term.

1b67107

Testing buildkite benchmarking

841ebb4

further buildkite experimentation

42bf53e

more buildkite

1188633

improve constraint testing coverage

ff8a276

jflat06 and others added 3 commits December 11, 2024 15:46

adding more test coverage

8e0e753

remove unused var

218e8f4

Fix env build issue (#307)

de4e99f

* Fix env build issue * take 2 * take 3 * Fix benchmarking bug

jflat06 force-pushed the jflat06/constraints branch from 942cc24 to de4e99f Compare December 11, 2024 23:54

jflat06 merged commit e5e9724 into master Dec 12, 2024
12 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Constraints #305

Constraints #305

jflat06 commented Jul 23, 2024 •

edited

Loading

jflat06 Sep 13, 2024

jflat06 Sep 13, 2024

jflat06 Oct 17, 2024

jflat06 Sep 13, 2024

jflat06 Oct 17, 2024

jflat06 Sep 13, 2024

jflat06 Oct 17, 2024

jflat06 Sep 13, 2024

jflat06 Oct 17, 2024

jflat06 Sep 13, 2024

fdimaio Oct 17, 2024 •

edited

Loading

jflat06 Sep 13, 2024

fdimaio Oct 17, 2024

jflat06 Sep 13, 2024

fdimaio Oct 17, 2024

jflat06 commented Sep 13, 2024

codecov bot commented Oct 17, 2024 •

edited

Loading

fdimaio left a comment

		flattened.index_add_(0, indices1, scores / 2)
		flattened.index_add_(0, indices2, scores / 2)

Constraints #305

Constraints #305

Conversation

jflat06 commented Jul 23, 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

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

fdimaio Oct 17, 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

jflat06 commented Sep 13, 2024

codecov bot commented Oct 17, 2024 • edited Loading

Codecov Report

fdimaio left a comment

Choose a reason for hiding this comment

jflat06 commented Jul 23, 2024 •

edited

Loading

fdimaio Oct 17, 2024 •

edited

Loading

codecov bot commented Oct 17, 2024 •

edited

Loading