From 9f6e9e87a75161ddb7e87feb57c689ded927e520 Mon Sep 17 00:00:00 2001 From: Jeff Flatten Date: Wed, 17 Apr 2024 15:30:16 -0700 Subject: [PATCH 1/6] Pose_stack->RF2 function, tests, init.py hooks. --- tmol/__init__.py | 4 ++ tmol/io/pose_stack_from_rosettafold2.py | 50 +++++++++++++++++++ .../io/test_pose_stack_from_rosettafold2.py | 12 +++++ 3 files changed, 66 insertions(+) diff --git a/tmol/__init__.py b/tmol/__init__.py index c44027500..d1fa1131e 100644 --- a/tmol/__init__.py +++ b/tmol/__init__.py @@ -21,6 +21,7 @@ ) from tmol.io.pose_stack_from_rosettafold2 import ( # noqa: F401 pose_stack_from_rosettafold2, + pose_stack_to_rosettafold2, canonical_form_from_rosettafold2, canonical_ordering_for_rosettafold2, packed_block_types_for_rosettafold2, @@ -32,6 +33,9 @@ from tmol.score import beta2016_score_function # noqa: F401 from tmol.score.score_function import ScoreFunction # noqa: F401 +from tmol.optimization.sfxn_modules import CartesianSfxnNetwork as cart_sfxn_network + +from tmol.optimization.lbfgs_armijo import LBFGS_Armijo as lbfgs_armijo try: __version__ = version("tmol") diff --git a/tmol/io/pose_stack_from_rosettafold2.py b/tmol/io/pose_stack_from_rosettafold2.py index 8878472b2..53c2e81c3 100644 --- a/tmol/io/pose_stack_from_rosettafold2.py +++ b/tmol/io/pose_stack_from_rosettafold2.py @@ -143,6 +143,56 @@ def canonical_form_from_rosettafold2( ) +def pose_stack_to_rosettafold2(pose_stack, chainlens): + from tmol.io.pose_stack_deconstruction import canonical_form_from_pose_stack + + device = pose_stack.device + n_poses = 1 # RF2 does not presently do batch processing + max_n_res = sum(chainlens) + max_n_ats = 27 + + rf2_pose_ind_for_atom = ( + torch.arange(n_poses, dtype=torch.int64, device=device) + .reshape(-1, 1, 1) + .expand(-1, max_n_res, max_n_ats) + ) + rf2_res_ind_for_atom = ( + torch.arange(max_n_res, dtype=torch.int64, device=device) + .reshape(1, -1, 1) + .expand(n_poses, -1, max_n_ats) + ) + + co = canonical_ordering_for_rosettafold2() + ( + _, # rf22t_rtmap, + rf22t_atmap, + rf2_at_is_real_map, + _, # supress_atom_for_nterm, + ) = _get_rf2_2_tmol_mappings(device) + + canonical_form = canonical_form_from_pose_stack(co, pose_stack) + + seq = canonical_form[1] + # tmol_restypes = rf22t_rtmap[seq] + atom_mapping = rf22t_atmap[seq] + rf2_at_is_real = rf2_at_is_real_map[seq] + + rf2_coords = torch.full( + (n_poses, max_n_res, max_n_ats, 3), + numpy.NaN, + dtype=torch.float32, + device=device, + ) + + rf2_coords[rf2_at_is_real] = canonical_form[2][ + rf2_pose_ind_for_atom[rf2_at_is_real], + rf2_res_ind_for_atom[rf2_at_is_real], + atom_mapping[rf2_at_is_real], + ] + + return rf2_coords + + @toolz.functoolz.memoize def _paramdb_for_rosettafold2() -> ParameterDatabase: """Construct the paramdb representing the subset of residues that diff --git a/tmol/tests/io/test_pose_stack_from_rosettafold2.py b/tmol/tests/io/test_pose_stack_from_rosettafold2.py index 426ad6e18..026e409f6 100644 --- a/tmol/tests/io/test_pose_stack_from_rosettafold2.py +++ b/tmol/tests/io/test_pose_stack_from_rosettafold2.py @@ -4,6 +4,7 @@ from tmol.io.pose_stack_from_rosettafold2 import ( pose_stack_from_rosettafold2, + pose_stack_to_rosettafold2, canonical_form_from_rosettafold2, _paramdb_for_rosettafold2, canonical_ordering_for_rosettafold2, @@ -48,6 +49,17 @@ def test_multi_chain_rosettafold2_pose_stack_construction( assert ps.packed_block_types is pbt +def test_from_to_rosettafold2(rosettafold2_ubq_pred, torch_device): + rosettafold2_ubq_pred["chainlens"] = [76] + ps = pose_stack_from_rosettafold2(**rosettafold2_ubq_pred) + rf2ubq = pose_stack_to_rosettafold2(ps, rosettafold2_ubq_pred["chainlens"]) + print(rosettafold2_ubq_pred["xyz"].unsqueeze(0).shape) + print(rf2ubq[0, 1:73]) + assert torch.allclose( + rosettafold2_ubq_pred["xyz"].unsqueeze(0)[0, 1:73], rf2ubq[0, 1:73] + ) + + def test_create_canonical_form_from_rosettafold2_ubq_stability( rosettafold2_ubq_pred, torch_device ): From 7b831df6e2fa727107a3c2426c2e898cb5c59737 Mon Sep 17 00:00:00 2001 From: Jeff Flatten Date: Fri, 19 Apr 2024 13:58:25 -0700 Subject: [PATCH 2/6] Accounting for suppressed ats in the tmol->rf2 function and test. --- tmol/io/pose_stack_from_rosettafold2.py | 18 ++++++++++++++++-- .../io/test_pose_stack_from_rosettafold2.py | 10 +++++----- 2 files changed, 21 insertions(+), 7 deletions(-) diff --git a/tmol/io/pose_stack_from_rosettafold2.py b/tmol/io/pose_stack_from_rosettafold2.py index 53c2e81c3..ed0e19800 100644 --- a/tmol/io/pose_stack_from_rosettafold2.py +++ b/tmol/io/pose_stack_from_rosettafold2.py @@ -167,7 +167,7 @@ def pose_stack_to_rosettafold2(pose_stack, chainlens): _, # rf22t_rtmap, rf22t_atmap, rf2_at_is_real_map, - _, # supress_atom_for_nterm, + supress_atom_for_nterm, ) = _get_rf2_2_tmol_mappings(device) canonical_form = canonical_form_from_pose_stack(co, pose_stack) @@ -175,6 +175,7 @@ def pose_stack_to_rosettafold2(pose_stack, chainlens): seq = canonical_form[1] # tmol_restypes = rf22t_rtmap[seq] atom_mapping = rf22t_atmap[seq] + supressed = supress_atom_for_nterm[seq] rf2_at_is_real = rf2_at_is_real_map[seq] rf2_coords = torch.full( @@ -190,7 +191,20 @@ def pose_stack_to_rosettafold2(pose_stack, chainlens): atom_mapping[rf2_at_is_real], ] - return rf2_coords + suppressed_mapped = torch.full( + (n_poses, max_n_res, max_n_ats), + False, + dtype=torch.bool, + device=device, + ) + suppressed_mapped[rf2_at_is_real] = supressed[ + rf2_pose_ind_for_atom[rf2_at_is_real], + rf2_res_ind_for_atom[rf2_at_is_real], + atom_mapping[rf2_at_is_real], + ] + return rf2_coords, torch.logical_and( + rf2_at_is_real, torch.logical_not(suppressed_mapped) + ) @toolz.functoolz.memoize diff --git a/tmol/tests/io/test_pose_stack_from_rosettafold2.py b/tmol/tests/io/test_pose_stack_from_rosettafold2.py index 026e409f6..404b10472 100644 --- a/tmol/tests/io/test_pose_stack_from_rosettafold2.py +++ b/tmol/tests/io/test_pose_stack_from_rosettafold2.py @@ -1,6 +1,7 @@ import os import torch +from tmol.tests.score.common.test_energy_term import assert_allclose from tmol.io.pose_stack_from_rosettafold2 import ( pose_stack_from_rosettafold2, @@ -52,11 +53,10 @@ def test_multi_chain_rosettafold2_pose_stack_construction( def test_from_to_rosettafold2(rosettafold2_ubq_pred, torch_device): rosettafold2_ubq_pred["chainlens"] = [76] ps = pose_stack_from_rosettafold2(**rosettafold2_ubq_pred) - rf2ubq = pose_stack_to_rosettafold2(ps, rosettafold2_ubq_pred["chainlens"]) - print(rosettafold2_ubq_pred["xyz"].unsqueeze(0).shape) - print(rf2ubq[0, 1:73]) - assert torch.allclose( - rosettafold2_ubq_pred["xyz"].unsqueeze(0)[0, 1:73], rf2ubq[0, 1:73] + + rf2ubq, rf2_ats = pose_stack_to_rosettafold2(ps, rosettafold2_ubq_pred["chainlens"]) + assert_allclose( + rosettafold2_ubq_pred["xyz"].unsqueeze(0)[rf2_ats], rf2ubq[rf2_ats], 1e-5, 1e-3 ) From 4a3bf7792eb2323949a171b43e3a110f0bdf0a32 Mon Sep 17 00:00:00 2001 From: Jeff Flatten Date: Wed, 29 May 2024 18:25:49 -0700 Subject: [PATCH 3/6] Map H1s back into H for terminal variants on tmol->RF2. --- tmol/io/pose_stack_from_rosettafold2.py | 80 ++++++++++++++++++- .../io/test_pose_stack_from_rosettafold2.py | 4 + tmol/tests/score/common/test_energy_term.py | 4 +- 3 files changed, 85 insertions(+), 3 deletions(-) diff --git a/tmol/io/pose_stack_from_rosettafold2.py b/tmol/io/pose_stack_from_rosettafold2.py index ed0e19800..662da7c37 100644 --- a/tmol/io/pose_stack_from_rosettafold2.py +++ b/tmol/io/pose_stack_from_rosettafold2.py @@ -168,13 +168,18 @@ def pose_stack_to_rosettafold2(pose_stack, chainlens): rf22t_atmap, rf2_at_is_real_map, supress_atom_for_nterm, - ) = _get_rf2_2_tmol_mappings(device) + hydrogens, + h_to_h1, + ) = _get_tmol_2_rf2_mappings(device) + + # torch.set_printoptions(threshold=10_000, linewidth=256) canonical_form = canonical_form_from_pose_stack(co, pose_stack) seq = canonical_form[1] - # tmol_restypes = rf22t_rtmap[seq] atom_mapping = rf22t_atmap[seq] + hydro = hydrogens[seq] + h1s = h_to_h1[seq] supressed = supress_atom_for_nterm[seq] rf2_at_is_real = rf2_at_is_real_map[seq] @@ -191,6 +196,15 @@ def pose_stack_to_rosettafold2(pose_stack, chainlens): atom_mapping[rf2_at_is_real], ] + # Find any NaN hydrogens and copy from the H1 instead + nans = torch.isnan(rf2_coords).any(-1) + terminal_hs = torch.logical_and(nans, hydro) + rf2_coords[terminal_hs] = canonical_form[2][ + rf2_pose_ind_for_atom[terminal_hs], + rf2_res_ind_for_atom[terminal_hs], + h1s[terminal_hs.any(-1)], + ] + suppressed_mapped = torch.full( (n_poses, max_n_res, max_n_ats), False, @@ -309,3 +323,65 @@ def _get_rf2_2_tmol_mappings(device: torch.device): tmol_ind = co.restypes_atom_index_mapping[i_3lc][atname.strip()] supress_atom_at_nterm[i, tmol_ind] = True return rt_map, atname_map, at_is_real, supress_atom_at_nterm + + +@toolz.functoolz.memoize +def _get_tmol_2_rf2_mappings(device: torch.device): + """Same logic is the RF2->tmol function, but additionally provides + a tensor marking the hydrogens in the RF2 index space, as well as + a tensor giving the tmol 1H index for a residue when indexed by + the RF2 residue type index + """ + + co = canonical_ordering_for_rosettafold2() + from tmol.extern.rosettafold2.chemical import ( + num2aa, + aa2long, + ) + + rf2_atom_names_for_name3s = { + x: [at.strip() if at is not None else "" for at in y] + for x, y in zip(num2aa, aa2long) + } + + (rt_map, atname_map, at_is_real) = co.create_src_2_tmol_mappings( + num2aa, rf2_atom_names_for_name3s, device + ) + + src_max_n_ats = len(rf2_atom_names_for_name3s[num2aa[0]]) + hydrogens = torch.zeros( + ( + len(num2aa), + src_max_n_ats, + ), + dtype=torch.bool, + device=device, + ) + h_to_h1 = torch.full( + (len(num2aa),), + -1, + dtype=torch.int64, + device=device, + ) + + # also want to turn off n-term "H" atoms + supress_atom_at_nterm = torch.zeros( + ( + len(num2aa), + co.max_n_canonical_atoms, + ), + dtype=torch.bool, + device=device, + ) + + for i, i_3lc in enumerate(num2aa): + if i_3lc not in co.restype_io_equiv_classes: + continue + for j, atname in enumerate(rf2_atom_names_for_name3s[i_3lc]): + if atname.strip() == "H": + tmol_ind = co.restypes_atom_index_mapping[i_3lc][atname.strip()] + supress_atom_at_nterm[i, tmol_ind] = True + + hydrogens[i, j] = True + h_to_h1[i] = co.restypes_atom_index_mapping[i_3lc]["1H"] + return rt_map, atname_map, at_is_real, supress_atom_at_nterm, hydrogens, h_to_h1 diff --git a/tmol/tests/io/test_pose_stack_from_rosettafold2.py b/tmol/tests/io/test_pose_stack_from_rosettafold2.py index 404b10472..10959eb89 100644 --- a/tmol/tests/io/test_pose_stack_from_rosettafold2.py +++ b/tmol/tests/io/test_pose_stack_from_rosettafold2.py @@ -52,9 +52,13 @@ def test_multi_chain_rosettafold2_pose_stack_construction( def test_from_to_rosettafold2(rosettafold2_ubq_pred, torch_device): rosettafold2_ubq_pred["chainlens"] = [76] + + # RF2->tmol ps = pose_stack_from_rosettafold2(**rosettafold2_ubq_pred) + # tmol->RF2 rf2ubq, rf2_ats = pose_stack_to_rosettafold2(ps, rosettafold2_ubq_pred["chainlens"]) + assert_allclose( rosettafold2_ubq_pred["xyz"].unsqueeze(0)[rf2_ats], rf2ubq[rf2_ats], 1e-5, 1e-3 ) diff --git a/tmol/tests/score/common/test_energy_term.py b/tmol/tests/score/common/test_energy_term.py index a4cb5131c..cd980da4b 100644 --- a/tmol/tests/score/common/test_energy_term.py +++ b/tmol/tests/score/common/test_energy_term.py @@ -83,7 +83,9 @@ def get_notallclose_msg(analytical, numerical, atol, rtol): def assert_allclose(baseline, measured, atol, rtol): try: - numpy.testing.assert_allclose(baseline, measured, atol=atol, rtol=rtol) + numpy.testing.assert_allclose( + baseline.cpu(), measured.cpu(), atol=atol, rtol=rtol + ) except AssertionError: raise AssertionError(get_notallclose_msg(measured, baseline, atol, rtol)) From db5f71aa19dad9970548279038f81c6043a07a59 Mon Sep 17 00:00:00 2001 From: Jeff Flatten Date: Mon, 3 Jun 2024 09:39:01 -0700 Subject: [PATCH 4/6] Moving the extra info on the tmol->rf2 function to a helper function to avoid complicating the interface for rf2 users. --- tmol/io/pose_stack_from_rosettafold2.py | 6 +++++- tmol/tests/io/test_pose_stack_from_rosettafold2.py | 6 ++++-- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/tmol/io/pose_stack_from_rosettafold2.py b/tmol/io/pose_stack_from_rosettafold2.py index 662da7c37..51e98c841 100644 --- a/tmol/io/pose_stack_from_rosettafold2.py +++ b/tmol/io/pose_stack_from_rosettafold2.py @@ -143,7 +143,7 @@ def canonical_form_from_rosettafold2( ) -def pose_stack_to_rosettafold2(pose_stack, chainlens): +def pose_stack_to_rosettafold2_with_suppressed(pose_stack, chainlens): from tmol.io.pose_stack_deconstruction import canonical_form_from_pose_stack device = pose_stack.device @@ -221,6 +221,10 @@ def pose_stack_to_rosettafold2(pose_stack, chainlens): ) +def pose_stack_to_rosettafold2(pose_stack, chainlens): + return pose_stack_to_rosettafold2_with_suppressed(pose_stack, chainlens)[0] + + @toolz.functoolz.memoize def _paramdb_for_rosettafold2() -> ParameterDatabase: """Construct the paramdb representing the subset of residues that diff --git a/tmol/tests/io/test_pose_stack_from_rosettafold2.py b/tmol/tests/io/test_pose_stack_from_rosettafold2.py index 10959eb89..4be24429c 100644 --- a/tmol/tests/io/test_pose_stack_from_rosettafold2.py +++ b/tmol/tests/io/test_pose_stack_from_rosettafold2.py @@ -5,7 +5,7 @@ from tmol.io.pose_stack_from_rosettafold2 import ( pose_stack_from_rosettafold2, - pose_stack_to_rosettafold2, + pose_stack_to_rosettafold2_with_suppressed, canonical_form_from_rosettafold2, _paramdb_for_rosettafold2, canonical_ordering_for_rosettafold2, @@ -57,7 +57,9 @@ def test_from_to_rosettafold2(rosettafold2_ubq_pred, torch_device): ps = pose_stack_from_rosettafold2(**rosettafold2_ubq_pred) # tmol->RF2 - rf2ubq, rf2_ats = pose_stack_to_rosettafold2(ps, rosettafold2_ubq_pred["chainlens"]) + rf2ubq, rf2_ats = pose_stack_to_rosettafold2_with_suppressed( + ps, rosettafold2_ubq_pred["chainlens"] + ) assert_allclose( rosettafold2_ubq_pred["xyz"].unsqueeze(0)[rf2_ats], rf2ubq[rf2_ats], 1e-5, 1e-3 From 23794cce0abb94953ca86b9fd757dbe86ba7e6c1 Mon Sep 17 00:00:00 2001 From: Jeff Flatten Date: Mon, 3 Jun 2024 12:54:57 -0700 Subject: [PATCH 5/6] Fix test breaking issue. --- tmol/tests/io/test_pose_stack_from_rosettafold2.py | 5 ++++- tmol/tests/score/common/test_energy_term.py | 5 ++--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/tmol/tests/io/test_pose_stack_from_rosettafold2.py b/tmol/tests/io/test_pose_stack_from_rosettafold2.py index 4be24429c..163ece971 100644 --- a/tmol/tests/io/test_pose_stack_from_rosettafold2.py +++ b/tmol/tests/io/test_pose_stack_from_rosettafold2.py @@ -62,7 +62,10 @@ def test_from_to_rosettafold2(rosettafold2_ubq_pred, torch_device): ) assert_allclose( - rosettafold2_ubq_pred["xyz"].unsqueeze(0)[rf2_ats], rf2ubq[rf2_ats], 1e-5, 1e-3 + rosettafold2_ubq_pred["xyz"].unsqueeze(0)[rf2_ats].cpu(), + rf2ubq[rf2_ats].cpu(), + 1e-5, + 1e-3, ) diff --git a/tmol/tests/score/common/test_energy_term.py b/tmol/tests/score/common/test_energy_term.py index cd980da4b..9ec4563f1 100644 --- a/tmol/tests/score/common/test_energy_term.py +++ b/tmol/tests/score/common/test_energy_term.py @@ -83,9 +83,7 @@ def get_notallclose_msg(analytical, numerical, atol, rtol): def assert_allclose(baseline, measured, atol, rtol): try: - numpy.testing.assert_allclose( - baseline.cpu(), measured.cpu(), atol=atol, rtol=rtol - ) + numpy.testing.assert_allclose(baseline, measured, atol=atol, rtol=rtol) except AssertionError: raise AssertionError(get_notallclose_msg(measured, baseline, atol, rtol)) @@ -378,6 +376,7 @@ def score(coords): scale = 0.01 * torch.arange( torch.numel(scores), device=scores.device ).reshape(scores.shape) + print(scale) return torch.sum(scale * scores) # monkeypatch more sane error reporting From 79038b5c8250b7a08e3bc6a154125bb49bad9d75 Mon Sep 17 00:00:00 2001 From: Jeff Flatten Date: Tue, 25 Jun 2024 17:18:17 -0700 Subject: [PATCH 6/6] Linting exception for api funcs, removing print statement. --- tmol/__init__.py | 6 ++++-- tmol/tests/score/common/test_energy_term.py | 1 - 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/tmol/__init__.py b/tmol/__init__.py index d1fa1131e..c4fe89a29 100644 --- a/tmol/__init__.py +++ b/tmol/__init__.py @@ -33,9 +33,11 @@ from tmol.score import beta2016_score_function # noqa: F401 from tmol.score.score_function import ScoreFunction # noqa: F401 -from tmol.optimization.sfxn_modules import CartesianSfxnNetwork as cart_sfxn_network +from tmol.optimization.sfxn_modules import ( + CartesianSfxnNetwork as cart_sfxn_network, +) # noqa: F401 -from tmol.optimization.lbfgs_armijo import LBFGS_Armijo as lbfgs_armijo +from tmol.optimization.lbfgs_armijo import LBFGS_Armijo as lbfgs_armijo # noqa: F401 try: __version__ = version("tmol") diff --git a/tmol/tests/score/common/test_energy_term.py b/tmol/tests/score/common/test_energy_term.py index 9ec4563f1..a4cb5131c 100644 --- a/tmol/tests/score/common/test_energy_term.py +++ b/tmol/tests/score/common/test_energy_term.py @@ -376,7 +376,6 @@ def score(coords): scale = 0.01 * torch.arange( torch.numel(scores), device=scores.device ).reshape(scores.shape) - print(scale) return torch.sum(scale * scores) # monkeypatch more sane error reporting