From 08143a5430ffbcbf0fd7ddf5bcff6831ad687eb1 Mon Sep 17 00:00:00 2001 From: Brennan Abanades Date: Thu, 10 Nov 2022 15:10:36 +0000 Subject: [PATCH] Code cleaning and minor updates --- ImmuneBuilder/ABodyBuilder2.py | 4 ++-- ImmuneBuilder/NanoBodyBuilder2.py | 4 ++-- ImmuneBuilder/TCRBuilder2.py | 6 +++--- ImmuneBuilder/models.py | 2 +- ImmuneBuilder/rigids.py | 4 +++- ImmuneBuilder/util.py | 3 ++- setup.py | 2 +- 7 files changed, 14 insertions(+), 11 deletions(-) diff --git a/ImmuneBuilder/ABodyBuilder2.py b/ImmuneBuilder/ABodyBuilder2.py index a7ddb53..2f74080 100644 --- 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 index 881e02b..90c6b9a 100644 --- 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 index 528d927..0109680 100644 --- 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 index 18c7c27..9f58050 100644 --- 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 index 54dddbe..63736d9 100644 --- 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 index e2ccf04..9dda49d 100644 --- 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 index 1e570a9..0bebf8f 100644 --- 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',