From 54ef185b0c6b8a4793915ec5399b0bdb7b2a984f Mon Sep 17 00:00:00 2001 From: lucidrains Date: Fri, 27 Sep 2024 09:05:55 -0700 Subject: [PATCH] fibonnaci sphere --- alphafold3_pytorch/alphafold3.py | 25 ++++++++++++++++++++++--- 1 file changed, 22 insertions(+), 3 deletions(-) diff --git a/alphafold3_pytorch/alphafold3.py b/alphafold3_pytorch/alphafold3.py index ff8ffdcc..ff69e41f 100644 --- a/alphafold3_pytorch/alphafold3.py +++ b/alphafold3_pytorch/alphafold3.py @@ -5681,7 +5681,7 @@ def _compute_unresolved_rasa( return unresolved_rasa.mean() @typecheck - def _compute_unresolved_rasa( + def _inhouse_compute_unresolved_rasa( self, unresolved_cid: int, unresolved_residue_mask: Bool[" n"], @@ -5689,7 +5689,8 @@ def _compute_unresolved_rasa( molecule_ids: Int[" n"], molecule_atom_lens: Int[" n"], atom_pos: Float["m 3"], - atom_mask: Bool[" m"], + atom_mask: Bool[" m"], + fibonacci_sphere_n = 200 # they use 200 in mkdssp ) -> Float[""]: """Compute the unresolved relative solvent accessible surface area (RASA) for proteins. @@ -5737,7 +5738,25 @@ def _compute_unresolved_rasa( # write custom RSA function here - raise NotImplementedError + # first constitute the fibonacci sphere + + num_surface_dots = fibonacci_sphere_n * 2. + 1 + golden_ratio = 1. + math.sqrt(5.) / 2 + + arange = torch.arange(-fibonacci_sphere_n, fibonacci_sphere_n + 1) # for example, N = 3 -> [-3, -2, -1, 0, 1, 2, 3] + + lat = torch.asin((2. * arange) / num_surface_dots) + lon = torch.fmod(arange, golden_ratio) * 2 * math.pi / golden_ratio + + surface_dots = torch.stack(( + lon.sin() * lat.cos(), + lon.cos() * lat.cos(), + lat.sin() + ), dim = -1) + + weight = (4. * math.pi) / num_surface_dots + + # rest of logic written by @xluo rasa = [] aatypes = []