Merge #77 from norskie/get_struclib_size

Function: Get StructureLibrary size (total # of orientations)

diffsims/libraries/structure_library.py

@@ -56,6 +56,7 @@ def __init__(self,
         for ident, struct, ori in zip(identifiers, structures, orientations):
             self.struct_lib[ident] = (struct, ori)
 
+
     @classmethod
     def from_orientation_lists(cls, identifiers, structures, orientations):
         """
@@ -103,3 +104,31 @@ def from_crystal_systems(cls, identifiers, structures, systems, resolution, equa
         for system in systems:
             orientations.append(get_grid_stereographic(system, resolution, equal))
         return cls(identifiers, structures, orientations)
+
+    def get_library_size(self, to_print = False):
+        """
+        Returns the the total number of orientations in the
+        current StructureLibrary object. Will also print the number of orientations
+        for each identifier in the library if the to_print==True
+        
+        Parameters
+        ----------
+        to_print : bool
+            Default is 'False'
+        Returns
+        -------
+        size_library : int
+            Total number of entries in the current StructureLibrary object.
+        """
+        size_library = 0
+        for i in range (len(self.orientations)):
+            if len(self.orientations[i]) == 1:
+                size_library += 1
+            else:
+                size_library += len(self.orientations[i])
+            if to_print == True:
+                print(self.identifiers[i], "has", \
+                len(self.orientations[i]), "number of entries.")
+        if to_print == True:
+            print("\nIn total:", size_library, "number of entries")
+        return size_library

diffsims/tests/test_library/test_structure_library.py

@@ -21,6 +21,12 @@
 
 from diffsims.libraries.structure_library import StructureLibrary
 
+def test_get_library_size():
+    identifiers = ['a', 'b']
+    structures = [1, 2]
+    orientations = [[(0, 0, 0), (0.0, 90.0, -180.0)], [(0, 0, 0)]]
+    library = StructureLibrary(identifiers, structures, orientations)
+    assert library.get_library_size(to_print = True) == 3
 
 def test_from_orientations_method():
     identifiers = ['a', 'b']