lint and cleanup

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -1,43 +1,5 @@
 from __future__ import annotations
 
-"""
-global ein notation:
-
-b - batch
-ba - batch with augmentation
-h - heads
-n - residue sequence length
-i - residue sequence length (source)
-j - residue sequence length (target)
-m - atom sequence length
-nw - windowed sequence length
-d - feature dimension
-ds - feature dimension (single)
-dp - feature dimension (pairwise)
-dap - feature dimension (atompair)
-dapi - feature dimension (atompair input)
-da - feature dimension (atom)
-dai - feature dimension (atom input)
-t - templates
-s - msa
-r - registers
-"""
-
-"""
-additional_residue_feats: [*, 10]:
-
-0: residue_index
-1: token_index
-2: asym_id
-3: entity_id
-4: sym_id
-5: restype (must be one hot encoded to 32)
-6: is_protein
-7: is_rna
-8: is_dna
-9: is_ligand
-"""
-
 from math import pi, sqrt
 from pathlib import Path
 from functools import partial, wraps
@@ -85,6 +47,44 @@
 
 from importlib.metadata import version
 
+"""
+global ein notation:
+
+b - batch
+ba - batch with augmentation
+h - heads
+n - residue sequence length
+i - residue sequence length (source)
+j - residue sequence length (target)
+m - atom sequence length
+nw - windowed sequence length
+d - feature dimension
+ds - feature dimension (single)
+dp - feature dimension (pairwise)
+dap - feature dimension (atompair)
+dapi - feature dimension (atompair input)
+da - feature dimension (atom)
+dai - feature dimension (atom input)
+t - templates
+s - msa
+r - registers
+"""
+
+"""
+additional_residue_feats: [*, 10]:
+
+0: residue_index
+1: token_index
+2: asym_id
+3: entity_id
+4: sym_id
+5: restype (must be one hot encoded to 32)
+6: is_protein
+7: is_rna
+8: is_dna
+9: is_ligand
+"""
+
 # constants
 
 ADDITIONAL_RESIDUE_FEATS = 10
@@ -241,24 +241,6 @@ def repeat_consecutive_with_lens(
 
     return output
 
-def repeat_pairwise_consecutive_with_lens(
-    feats: Float['b n n dp'],
-    lens: Int['b n']
-) -> Float['b m m dp']:
-
-    repeated_lens = repeat(lens, 'b ... -> (b repeat) ...', repeat = feats.shape[1])
-    feats, ps = pack_one(feats, '* n dp')
-    feats = repeat_consecutive_with_lens(feats, repeated_lens)
-    feats = unpack_one(feats, ps, '* n dp')
-
-    feats = rearrange(feats, 'b i j dp -> b j i dp')
-    repeated_lens = repeat(lens, 'b ... -> (b repeat) ...', repeat = feats.shape[1])
-    feats, ps = pack_one(feats, '* n dp')
-    feats = repeat_consecutive_with_lens(feats, repeated_lens)
-    feats = unpack_one(feats, ps, '* n dp')
-    feats = rearrange(feats, 'b j i dp -> b i j dp')
-    return feats
-
 # linear and outer sum
 # for single repr -> pairwise pattern throughout this architecture
 
@@ -2105,8 +2087,6 @@ def forward(
         align_weights = atom_pos_ground_truth.new_ones(atom_pos_ground_truth.shape[:2])
 
         if exists(additional_residue_feats):
-            w = self.net.atoms_per_window
-
             is_nucleotide_or_ligand_fields = (additional_residue_feats[..., 7:] != 0.).unbind(dim = -1)
 
             is_nucleotide_or_ligand_fields = tuple(repeat_consecutive_with_lens(t, residue_atom_lens) for t in is_nucleotide_or_ligand_fields)

alphafold3_pytorch/attention.py

@@ -2,7 +2,7 @@
 from typing import NamedTuple, Tuple
 
 import torch
-from torch import nn
+from torch import nn, Tensor
 import torch.nn.functional as F
 from torch.nn import Module
 

alphafold3_pytorch/trainer.py

@@ -29,14 +29,14 @@ class Alphafold3Input(TypedDict):
     additional_residue_feats:   Float['n 10']
     templates:                  Float['t n n dt']
     msa:                        Float['s n dm']
-    template_mask:              Bool['t'] | None
-    msa_mask:                   Bool['s'] | None
+    template_mask:              Bool[' t'] | None
+    msa_mask:                   Bool[' s'] | None
     atom_pos:                   Float['m 3'] | None
-    residue_atom_indices:       Int['n'] | None
+    residue_atom_indices:       Int[' n'] | None
     distance_labels:            Int['n n'] | None
     pae_labels:                 Int['n n'] | None
-    pde_labels:                 Int['n'] | None
-    resolved_labels:            Int['n'] | None
+    pde_labels:                 Int[' n'] | None
+    resolved_labels:            Int[' n'] | None
 
 # helpers
 
@@ -406,4 +406,4 @@ def __call__(
 
             self.log(**test_loss_breakdown)
 
-        print(f'training complete')
+        print('training complete')

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphafold3-pytorch"
-version = "0.1.1"
+version = "0.1.2"
 description = "Alphafold 3 - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }