cleanup

Author: lucidrains

alphafold3_pytorch/attention.py

@@ -427,7 +427,7 @@ def local_attn(
         # maybe softclamp
 
         if self.enable_attn_softclamp:
-            sim = softclamp_value(sim, self.attn_softclamp_value)
+            sim = softclamp(sim, self.attn_softclamp_value)
 
         # windowed masking - for masking out atoms not belonging to the same molecule / polypeptide / nucleic acid in sequence-local attention
 
@@ -528,7 +528,7 @@ def forward(
         # maybe softclamp
 
         if self.enable_attn_softclamp:
-            sim = softclamp_value(sim, self.attn_softclamp_value)
+            sim = softclamp(sim, self.attn_softclamp_value)
 
         # masking
 

alphafold3_pytorch/inputs.py

@@ -2,7 +2,7 @@
 
 from functools import wraps, partial
 from dataclasses import dataclass, asdict, field
-from typing import Type, Literal, Callable, List, Any
+from typing import Type, Literal, Callable, List, Any, Tuple
 
 import torch
 from torch import tensor
@@ -350,7 +350,7 @@ class Alphafold3Input:
     ligands:                    List[Mol | str] = imm_list() # can be given as smiles
     ds_dna:                     List[Int[' _'] | str] = imm_list()
     ds_rna:                     List[Int[' _'] | str] = imm_list()
-    atom_parent_ids:            Int['m'] | None = None
+    atom_parent_ids:            Int[' m'] | None = None
     additional_token_feats:     Float[f'n dtf'] | None = None
     templates:                  Float['t n n dt'] | None = None
     msa:                        Float['s n dm'] | None = None
@@ -508,7 +508,7 @@ def alphafold3_input_to_molecule_input(
     # convert ligands to rdchem.Mol
 
     ligands = list(alphafold3_input.ligands)
-    mol_ligands = [(mol_from_smile(l) if isinstance(l, str) else l) for l in ligands]
+    mol_ligands = [(mol_from_smile(ligand) if isinstance(ligand, str) else ligand) for ligand in ligands]
 
     # create the molecule input
 
@@ -643,7 +643,7 @@ def get_num_atoms_per_chain(chains: List[List[Mol]]) -> List[int]:
     num_protein_atoms = get_num_atoms_per_chain(mol_proteins)
     num_ss_rna_atoms = get_num_atoms_per_chain(mol_ss_rnas)
     num_ss_dna_atoms = get_num_atoms_per_chain(mol_ss_dnas)
-    num_ligand_atoms = [l.GetNumAtoms() for l in mol_ligands]
+    num_ligand_atoms = [ligand.GetNumAtoms() for ligand in mol_ligands]
 
     atom_counts = [*num_protein_atoms, *num_ss_rna_atoms, *num_ss_dna_atoms, *num_ligand_atoms, num_metal_ions]
 
@@ -665,7 +665,7 @@ def get_num_atoms_per_chain(chains: List[List[Mol]]) -> List[int]:
     num_protein_tokens = [len(protein) for protein in proteins]
     num_ss_rna_tokens = [len(rna) for rna in ss_rnas]
     num_ss_dna_tokens = [len(dna) for dna in ss_dnas]
-    num_ligand_tokens = [l.GetNumAtoms() for l in mol_ligands]
+    num_ligand_tokens = [ligand.GetNumAtoms() for ligand in mol_ligands]
 
     token_repeats = tensor([*num_protein_tokens, *num_ss_rna_tokens, *num_ss_dna_tokens, *num_ligand_tokens, num_metal_ions])
 

alphafold3_pytorch/life.py

@@ -246,9 +246,9 @@ def reverse_complement(
     return ''.join(complement[::-1])
 
 @typecheck
-def reverse_complement_tensor(t: Int['n']):
+def reverse_complement_tensor(t: Int[' n']):
     complement = NUCLEIC_ACID_COMPLEMENT_TENSOR[t]
-    reverse_complement = t.flip(dims = (-1,))
+    reverse_complement = complement.flip(dims = (-1,))
     return reverse_complement
 
 # metal ions

pyproject.toml

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