Code cleaning and minor updates

brennanaba · Nov 10, 2022 · 08143a5 · 08143a5
1 parent 07f4bf5
commit 08143a5
Show file tree

Hide file tree

Showing 7 changed files with 14 additions and 11 deletions.
diff --git a/ImmuneBuilder/ABodyBuilder2.py b/ImmuneBuilder/ABodyBuilder2.py
@@ -62,7 +62,7 @@ def save_all(self, dirname=None, filename=None):
         np.save(os.path.join(dirname,"error_estimates"), self.error_estimates.mean(0).cpu().numpy())
         final_filename = os.path.join(dirname, filename)
         refine(os.path.join(dirname,"rank0_unrefined.pdb"), final_filename)
-        add_errors_as_bfactors(final_filename, self.error_estimates.mean(0).sqrt().cpu().numpy(), new_txt=[header])
+        add_errors_as_bfactors(final_filename, self.error_estimates.mean(0).sqrt().cpu().numpy(), header=[header])
 
 
     def save(self, filename=None):
@@ -82,7 +82,7 @@ def save(self, filename=None):
 
         if not success:
             print(f"FAILED TO REFINE {filename}.\nSaving anyways.", flush=True)
-        add_errors_as_bfactors(filename, self.error_estimates.mean(0).sqrt().cpu().numpy(), new_txt=[header])  
+        add_errors_as_bfactors(filename, self.error_estimates.mean(0).sqrt().cpu().numpy(), header=[header])  
 
 
 class ABodyBuilder2:

diff --git a/ImmuneBuilder/NanoBodyBuilder2.py b/ImmuneBuilder/NanoBodyBuilder2.py
@@ -62,7 +62,7 @@ def save_all(self, dirname=None, filename=None):
         np.save(os.path.join(dirname,"error_estimates"), self.error_estimates.mean(0).cpu().numpy())
         final_filename = os.path.join(dirname, filename)
         refine(os.path.join(dirname,"rank0_unrefined.pdb"), final_filename)
-        add_errors_as_bfactors(final_filename, self.error_estimates.mean(0).sqrt().cpu().numpy(), new_txt=[header])
+        add_errors_as_bfactors(final_filename, self.error_estimates.mean(0).sqrt().cpu().numpy(), header=[header])
 
 
     def save(self, filename=None):
@@ -82,7 +82,7 @@ def save(self, filename=None):
 
         if not success:
             print(f"FAILED TO REFINE {filename}.\nSaving anyways.", flush=True)
-        add_errors_as_bfactors(filename, self.error_estimates.mean(0).sqrt().cpu().numpy(), new_txt=[header])  
+        add_errors_as_bfactors(filename, self.error_estimates.mean(0).sqrt().cpu().numpy(), header=[header])  
 
 
 class NanoBodyBuilder2:

diff --git a/ImmuneBuilder/TCRBuilder2.py b/ImmuneBuilder/TCRBuilder2.py
@@ -62,7 +62,7 @@ def save_all(self, dirname=None, filename=None):
         np.save(os.path.join(dirname,"error_estimates"), self.error_estimates.mean(0).cpu().numpy())
         final_filename = os.path.join(dirname, filename)
         refine(os.path.join(dirname,"rank0_unrefined.pdb"), final_filename)
-        add_errors_as_bfactors(final_filename, self.error_estimates.mean(0).sqrt().cpu().numpy(), new_txt=[header])
+        add_errors_as_bfactors(final_filename, self.error_estimates.mean(0).sqrt().cpu().numpy(), header=[header])
 
 
     def save(self, filename=None):
@@ -82,7 +82,7 @@ def save(self, filename=None):
 
         if not success:
             print(f"FAILED TO REFINE {filename}.\nSaving anyways.", flush=True)
-        add_errors_as_bfactors(filename, self.error_estimates.mean(0).sqrt().cpu().numpy(), new_txt=[header])  
+        add_errors_as_bfactors(filename, self.error_estimates.mean(0).sqrt().cpu().numpy(), header=[header])  
 
 
 class TCRBuilder2:
@@ -164,7 +164,7 @@ def command_line_interface():
         print("Running sequences through deep learning model...", flush=True)
 
     try:
-        antibody = tcr = TCRBuilder2().predict(seqs)
+        tcr = TCRBuilder2().predict(seqs)
     except AssertionError as e:
         print(e, flush=True)
         sys.exit(1)

diff --git a/ImmuneBuilder/models.py b/ImmuneBuilder/models.py
@@ -225,7 +225,7 @@ def forward(self, node_features, sequence):
 
         # Remove atoms of side chains with outrageous clashes
         ds = torch.linalg.norm(all_atoms[None,:,None] - all_atoms[:,None,:,None], axis = -1)
-        ds[(ds!=ds) | (ds==0.0)] = 10
+        ds[torch.isnan(ds!=ds) | (ds==0.0)] = 10
         min_ds = ds.min(dim=-1)[0].min(dim=-1)[0].min(dim=-1)[0]
         all_atoms[min_ds < 0.2, 5:, :] = float("Nan")
 

diff --git a/ImmuneBuilder/rigids.py b/ImmuneBuilder/rigids.py
@@ -157,8 +157,10 @@ def __init__(self, origin, rot):
     def __matmul__(self, other):
         if isinstance(other, Vector):
             return self.rot @ other + self.origin
-        if isinstance(other, Rigid):
+        elif isinstance(other, Rigid):
             return Rigid(self.rot @ other.origin + self.origin, self.rot @ other.rot)
+        else:
+            raise TypeError(f"can't multiply rigid by object of type {type(other)}")
 
     def inv(self):
         inv_rot = self.rot.inv()

diff --git a/ImmuneBuilder/util.py b/ImmuneBuilder/util.py
@@ -103,11 +103,12 @@ def sequence_dict_from_fasta(fasta_file):
     return out
 
 
-def add_errors_as_bfactors(filename, errors, new_txt=[]):
+def add_errors_as_bfactors(filename, errors, header=[]):
 
     with open(filename) as file:
         txt = file.readlines()
 
+    new_txt = header.copy()
     residue_index = -1
     position = "  "
 

diff --git a/setup.py b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name='ImmuneBuilder',
-    version='0.0.3',
+    version='0.0.4',
     description='Set of functions to predict the structure of immune receptor proteins',
     license='BSD 3-clause license',
     maintainer='Brennan Abanades',