add option to return ADPs as an array or dictionary

Author: cadenmyers13

src/diffpy/structure/structure.py

@@ -213,29 +213,61 @@ def get_cartesian_coordinates(self):
         cartn_coords = numpy.array([a.xyz_cartn for a in self])
         return cartn_coords
 
-    def get_anisotropic_displacement_parameters(self):
-        """Return array of anisotropic displacement parameters of all
-        `Atoms` in this structure.
+    def get_anisotropic_displacement_parameters(self, return_array=False):
+        """Return a dictionary of anisotropic displacement parameters
+        for all atoms.
+
+        Parameters
+        ----------
+        return_array : bool, optional
+            If True, return anisotropic displacement parameters as a numpy array instead of a dictionary.
 
         Returns
         -------
-        numpy.ndarray
-            The array of anisotropic displacement parameters of all `Atoms` in this structure.
+        dict
+            The dictionary of anisotropic displacement parameters for all atoms in this structure.
+            Keys are of the form 'Element_i_Ujk', e.g. 'C_0_11', 'C_0_12'.
         """
-        aniso_adps = numpy.array([a.U for a in self])
-        return aniso_adps
+        if return_array:
+            aniso_adps = numpy.array([a.U for a in self])
+            return aniso_adps
+        else:
+            adp_dict = {}
+            for i, atom in enumerate(self):
+                element = atomBareSymbol(atom.element)
+                adp_dict[f"{element}_{i}_11"] = self.U11[i]
+                adp_dict[f"{element}_{i}_22"] = self.U22[i]
+                adp_dict[f"{element}_{i}_33"] = self.U33[i]
+                adp_dict[f"{element}_{i}_12"] = self.U12[i]
+                adp_dict[f"{element}_{i}_13"] = self.U13[i]
+                adp_dict[f"{element}_{i}_23"] = self.U23[i]
+            return adp_dict
+
+    def get_isotropic_displacement_parameters(self, return_array=False):
+        """Return a dictionary of isotropic displacement parameters for
+        all atoms.
 
-    def get_isotropic_displacement_parameters(self):
-        """Return array of isotropic displacement parameters of all
-        `Atoms` in this structure.
+        Parameters
+        ----------
+        return_array : bool, optional
+            If True, return isotropic displacement parameters as a numpy array instead of a dictionary.
+            Default is False.
 
         Returns
         -------
-        numpy.ndarray
-            The array of isotropic displacement parameters of all `Atoms` in this structure.
+        dict
+            The dictionary of isotropic displacement parameters for all atoms in this structure.
+            Keys are of the form 'Element_i_Uiso', e.g. 'C_0_Uiso'.
         """
-        iso_adps = numpy.array([a.Uisoequiv for a in self])
-        return iso_adps
+        if return_array:
+            iso_adps = numpy.array([a.Uisoequiv for a in self])
+            return iso_adps
+        else:
+            iso_dict = {}
+            for i, atom in enumerate(self):
+                element = atomBareSymbol(atom.element)
+                iso_dict[f"{element}_{i+1}_Uiso"] = self.Uisoequiv[i]
+            return iso_dict
 
     def get_occupancies(self):
         """Return array of occupancies of all `Atoms` in this structure.