Merge 5d3b5b4 into 00ad9cf

diffsims/libraries/structure_library.py

@@ -56,6 +56,36 @@ def __init__(self,
         for ident, struct, ori in zip(identifiers, structures, orientations):
             self.struct_lib[ident] = (struct, ori)
 
+    def get_library_size(cls, 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 parameter to_print is set to
+        True. Works with both cls.orientations on the form [1,2]
+        and [[1,2], [2,3]](rotation lists).
+
+        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(cls.orientations)):
+            if type(cls.orientations[i]) != list:
+                size_library += 1
+            else:
+                size_library += len(cls.orientations[i])
+            if to_print == True and type(cls.orientations[i]) == list:
+                print(cls.identifiers[i], "has", \
+                len(cls.orientations[i]), "number of entries.")
+                if to_print == True:
+                    print("\nIn total:", size_library, "number of entries")
+        return size_library
+
     @classmethod
     def from_orientation_lists(cls, identifiers, structures, orientations):
         """

diffsims/tests/test_library/test_structure_library.py

@@ -21,6 +21,17 @@
 
 from diffsims.libraries.structure_library import StructureLibrary
 
+def test_get_library_size():
+    identifiers = ['a', 'b']
+    structures = [1, 2]
+    rotation_list = [[3,1], [4,5]]
+    orientations = [1, 2]
+    first_library = StructureLibrary(identifiers, structures, rotation_list)
+    second_library = StructureLibrary(identifiers, structures, orientations)
+    # Test for rotation list
+    assert first_library.get_library_size() == 4
+    # Test for single orientation
+    assert second_library.get_library_size() == 2
 
 def test_from_orientations_method():
     identifiers = ['a', 'b']