add docstring and rename get_assemblies

BradyAJohnston · Aug 2, 2023 · 84155af · 84155af
1 parent 73087bc
commit 84155af
Show file tree

Hide file tree

Showing 6 changed files with 21 additions and 6 deletions.
diff --git a/MolecularNodes/assembly/__init__.py b/MolecularNodes/assembly/__init__.py
@@ -32,4 +32,19 @@ def get_transformations(self, assembly_id):
         |          |        |        translation vector 
         |          |        |        |
         list[tuple[ndarray, ndarray, ndarray]]]
+        """
+
+    @abstractmethod
+    def get_assemblies(self):
+        """
+        Parse all the transformations for each assembly, returning a dictionary of 
+        key:value pairs of assembly_id:transformations. The transformations list
+        comes from the `get_transformations(assembly_id)` method.
+        
+        Dictionary of all assemblies
+        |     Assembly ID
+        |     |   List of transformations to create biological assembly.
+        |     |   | 
+        dict{'1', list[transformations]}
+        
         """
diff --git a/MolecularNodes/assembly/cif.py b/MolecularNodes/assembly/cif.py
@@ -73,7 +73,7 @@ def get_transformations(self, assembly_id):
 
         return transformations
 
-    def get_all_transformations(self):
+    def get_assemblies(self):
         assembly_dict = {}
         for assembly_id in self.list_assemblies():
             assembly_dict[assembly_id] = self.get_transformations(assembly_id)

diff --git a/MolecularNodes/assembly/mmtf.py b/MolecularNodes/assembly/mmtf.py
@@ -48,7 +48,7 @@ def get_transformations(self, assembly_id):
 
         return transformations
 
-    def get_all_transformations(self):
+    def get_assemblies(self):
         assembly_dict = {}
         for assembly_id in self.list_assemblies():
             assembly_dict[assembly_id] = self.get_transformations(assembly_id)

diff --git a/MolecularNodes/assembly/pdb.py b/MolecularNodes/assembly/pdb.py
@@ -88,7 +88,7 @@ def get_transformations(self, assembly_id):
 
         return transformations
 
-    def get_all_transformations(self):
+    def get_assemblies(self):
         assembly_dict = {}
         for assembly_id in self.list_assemblies():
             assembly_dict[assembly_id] = self.get_transformations(assembly_id)

diff --git a/MolecularNodes/load.py b/MolecularNodes/load.py
@@ -116,7 +116,7 @@ def molecule_rcsb(
 
     try:
         parsed_assembly_file = assembly.mmtf.MMTFAssemblyParser(file)
-        mol_object['biological_assemblies'] = parsed_assembly_file.get_all_transformations()
+        mol_object['biological_assemblies'] = parsed_assembly_file.get_assemblies()
     except InvalidFileError:
         pass
 
@@ -143,14 +143,14 @@ def molecule_local(
     if file_ext == '.pdb':
         mol, file = open_structure_local_pdb(file_path, include_bonds)
         try:
-            transforms = assembly.pdb.PDBAssemblyParser(file).get_all_transformations()
+            transforms = assembly.pdb.PDBAssemblyParser(file).get_assemblies()
         except InvalidFileError:
             transforms = None
 
     elif file_ext == '.pdbx' or file_ext == '.cif':
         mol, file = open_structure_local_pdbx(file_path, include_bonds)
         try:
-            transforms = assembly.cif.CIFAssemblyParser(file).get_all_transformations()
+            transforms = assembly.cif.CIFAssemblyParser(file).get_assemblies()
         except InvalidFileError:
             transforms = None
 

diff --git a/tests/data/md_ppr/.first_5_frames.xtc_offsets.npz b/tests/data/md_ppr/.first_5_frames.xtc_offsets.npz