Move some of the structure plotting functions from contrib

CHANGELOG.md

@@ -1,5 +1,7 @@
 # Unreleased
 
+- Add ploting functions from pyiron_contrib to StructureContainer ([#710](https://github.com/pyiron/pyiron_atomistics/pull/710))
+
 # pyiron_atomistics-0.2.41
 
 - Symmetrize vectors ([#576](https://github.com/pyiron/pyiron_atomistics/pull/576))

pyiron_atomistics/atomistics/job/structurecontainer.py

@@ -154,3 +154,10 @@ def from_hdf(self, hdf=None, group_name=None):
         else:
             super().from_hdf(hdf=hdf, group_name=group_name)
             self._container.from_hdf(hdf=self.project_hdf5, group_name="structures")
+
+    @property
+    def plot(self):
+        """
+        Accessor for :class:`~.StructurePlots` instance using these structures.
+        """
+        return self._container.plot

pyiron_atomistics/atomistics/structure/structurestorage.py

@@ -8,13 +8,21 @@
 
 from typing import List
 
+import matplotlib.pyplot as plt
 import numpy as np
+import pandas as pd
 
-from pyiron_base import FlattenedStorage
+from pyiron_base import FlattenedStorage, ImportAlarm
 from pyiron_atomistics.atomistics.structure.atom import Atom
 from pyiron_atomistics.atomistics.structure.atoms import Atoms
+from pyiron_atomistics.atomistics.structure.neighbors import NeighborsTrajectory
 from pyiron_atomistics.atomistics.structure.has_structure import HasStructure
 
+with ImportAlarm(
+    "Some plotting functionality requires the seaborn library."
+) as seaborn_alarm:
+    import seaborn as sns
+
 
 class StructureStorage(FlattenedStorage, HasStructure):
     """
@@ -86,6 +94,7 @@ def __init__(self, num_atoms=1, num_structures=1):
         """
         super().__init__(num_elements=num_atoms, num_chunks=num_structures)
         self._element_cache = None
+        self._plots = None
 
     def _init_arrays(self):
         super()._init_arrays()
@@ -241,3 +250,283 @@ def _number_of_structures(self):
 
     def _get_hdf_group_name(self):
         return "structures"
+
+    @property
+    def plot(self):
+        """
+        Accessor for :class:`.StructurePlots` instance using these structures.
+        """
+        if self._plots is None:
+            self._plots = StructurePlots(self)
+        return self._plots
+
+
+class StructurePlots:
+    """
+    Simple interface to plot various properties of structures.
+    """
+
+    @seaborn_alarm
+    def __init__(self, store: StructureStorage):
+        self._store = store
+        self._neigh = None
+
+    def cell(self, angle_in_degrees=True):
+        """
+        Plot histograms of cell parameters.
+
+        Plotted are atomic volume, density, cell vector lengths and cell vector angles in separate subplots all on a
+        log-scale.
+
+        Args:
+            angle_in_degrees (bool): whether unit for angles is degree or radians
+
+        Returns:
+            `DataFrame`: contains the plotted information in the columns:
+                            - a: length of first vector
+                            - b: length of second vector
+                            - c: length of third vector
+                            - alpha: angle between first and second vector
+                            - beta: angle between second and third vector
+                            - gamma: angle between third and first vector
+                            - V: volume of the cell
+                            - N: number of atoms in the cell
+        """
+        N = self._store.get_array("length")
+        C = self._store.get_array("cell")
+
+        def get_angle(cell, idx=0):
+            return np.arccos(
+                np.dot(cell[idx], cell[(idx + 1) % 3])
+                / np.linalg.norm(cell[idx])
+                / np.linalg.norm(cell[(idx + 1) % 3])
+            )
+
+        def extract(n, c):
+            return {
+                "a": np.linalg.norm(c[0]),
+                "b": np.linalg.norm(c[1]),
+                "c": np.linalg.norm(c[2]),
+                "alpha": get_angle(c, 0),
+                "beta": get_angle(c, 1),
+                "gamma": get_angle(c, 2),
+            }
+
+        df = pd.DataFrame([extract(n, c) for n, c in zip(N, C)])
+        df["V"] = np.linalg.det(C)
+        df["N"] = N
+        if angle_in_degrees:
+            df["alpha"] = np.rad2deg(df["alpha"])
+            df["beta"] = np.rad2deg(df["beta"])
+            df["gamma"] = np.rad2deg(df["gamma"])
+
+        plt.subplot(1, 4, 1)
+        plt.title("Atomic Volume")
+        plt.hist(df.V / df.N, bins=20, log=True)
+        plt.xlabel(r"$V$ [$\AA^3$]")
+
+        plt.subplot(1, 4, 2)
+        plt.title("Density")
+        plt.hist(df.N / df.V, bins=20, log=True)
+        plt.xlabel(r"$\rho$ [$\AA^{-3}$]")
+
+        plt.subplot(1, 4, 3)
+        plt.title("Lattice Vector Lengths")
+        plt.hist([df.a, df.b, df.c], log=True)
+        plt.xlabel(r"$a,b,c$ [$\AA$]")
+
+        plt.subplot(1, 4, 4)
+        plt.title("Lattice Vector Angles")
+        plt.hist([df.alpha, df.beta, df.gamma], log=True)
+        if angle_in_degrees:
+            label = r"$\alpha,\beta,\gamma$ [°]"
+        else:
+            label = r"$\alpha,\beta,\gamma$ [rad]"
+        plt.xlabel(label)
+
+        return df
+
+    def _calc_spacegroups(self, symprec=1e-3):
+        """
+        Calculate space groups of all structures.
+
+        Args:
+            symprec (float): symmetry precision given to spglib
+
+        Returns:
+            DataFrame: contains columns 'crystal_system' (str) and 'space_group' (int) for each structure
+        """
+
+        def get_crystal_system(num):
+            if num in range(1, 3):
+                return "triclinic"
+            elif num in range(3, 16):
+                return "monoclinic"
+            elif num in range(16, 75):
+                return "orthorombic"
+            elif num in range(75, 143):
+                return "trigonal"
+            elif num in range(143, 168):
+                return "tetragonal"
+            elif num in range(168, 195):
+                return "hexagonal"
+            elif num in range(195, 230):
+                return "cubic"
+
+        def extract(s):
+            spg = s.get_symmetry(symprec=symprec).spacegroup["Number"]
+            return {"space_group": spg, "crystal_system": get_crystal_system(spg)}
+
+        return pd.DataFrame(map(extract, self._store.iter_structures()))
+
+    def spacegroups(self, symprec=1e-3):
+        """
+        Plot histograms of space groups and crystal systems.
+
+        Spacegroups and crystal systems are plotted in separate subplots.
+
+        Args:
+            symprec (float): precision of the symmetry search (passed to spglib)
+
+        Returns:
+            DataFrame: contains two columns "space_group", "crystal_system"
+                       for each structure
+        """
+
+        df = self._calc_spacegroups(symprec=symprec)
+        plt.subplot(1, 2, 1)
+        plt.hist(df.space_group, bins=230)
+        plt.xlabel("Space Group")
+
+        plt.subplot(1, 2, 2)
+        l, h = np.unique(df.crystal_system, return_counts=True)
+        sort_key = {
+            "triclinic": 1,
+            "monoclinic": 3,
+            "orthorombic": 16,
+            "trigonal": 75,
+            "tetragonal": 143,
+            "hexagonal": 168,
+            "cubic": 195,
+        }
+        I = np.argsort([sort_key[ll] for ll in l])
+        plt.bar(l[I], h[I])
+        plt.xlabel("Crystal System")
+        plt.xticks(rotation=35)
+        return df
+
+    def _calc_neighbors(self, num_neighbors):
+        """
+        Calculate the neighbor information with additional caching.
+
+        If 'distances' and 'shells' are provided in the underlying store, they are returned directly without checking
+        `num_neighbors`.
+
+        If they are not provided there, they are calculated here and cached.
+        Recalculation happens when a different `num_neighbors` is provided than in a previous call or the underlying
+        store changes.
+
+        If `num_neighbors` is `None` on the first call, the default is 36.
+
+        Returns:
+            dict: with keys 'distances' and 'shells' containing the respective flattened arrays from
+            :meth:`.Atoms.get_neighbors`.
+        """
+        if self._store.has_array("distances") and self._store.has_array("shells"):
+            return {
+                "distances": self._store["distances"],
+                "shells": self._store["shells"],
+            }
+        # check that _store and _neigh are still consistent
+        cur_neighbors = self._neigh.has_array("distances")["shape"][0]
+        if (
+            self._neigh is None
+            or len(self._store) != len(self._neigh)
+            or (num_neighbors is None or cur_neighbors != num_neighbors)
+        ):
+            if num_neighbors is None:
+                num_neighbors = 36
+            self._neigh = FlattenedStorage()
+            neigh_traj = NeighborsTrajectory(
+                has_structure=self._store,
+                num_neighbors=num_neighbors,
+                store=self._neigh,
+            )
+        return {
+            "distances": self._neigh["distances"],
+            "shells": self._neigh["shells"],
+        }
+
+    def coordination(self, num_shells=4, log=True, num_neighbors=None):
+        """
+        Plot histogram of coordination in neighbor shells.
+
+        Computes one histogram of the number of neighbors in each neighbor shell up to `num_shells` and then plots them
+        together.
+
+        If the underlying :class:`.StructureStorage` has a 'shells' array defined it is used, if not it is calculated on
+        the fly.
+
+        Args:
+            num_shells (int): maximum shell to plot
+            num_neighbors (int): maximum number of neighbors to calculate, when 'shells' is not defined in storage,
+                                 default is the value from the previous call or 36
+            log (float): plot histogram values on a log scale
+        """
+        neigh = self._calc_neighbors(num_neighbors=num_neighbors)
+        shells = neigh["shells"]
+
+        shell_index = (
+            shells[np.newaxis, :, :]
+            == np.arange(1, num_shells + 1)[:, np.newaxis, np.newaxis]
+        )
+        neigh_count = shell_index.sum(axis=-1)
+        ticks = np.arange(neigh_count.min(), neigh_count.max() + 1)
+        plt.hist(
+            neigh_count.T,
+            bins=ticks - 0.5,
+            log=True,
+            label=[f"{i}." for i in range(1, num_shells + 1)],
+        )
+        plt.xticks(ticks)
+        plt.xlabel("Number of Neighbors")
+        plt.legend(title="Shell")
+        plt.title("Neighbor Coordination in Shells")
+
+    def distances(self, bins=50, num_neighbors=None):
+        """
+        Plot a histogram of the neighbor distances.
+
+        Args:
+            bins (int): number of bins
+            num_neighbors (int): maximum number of neighbors to calculate, when 'shells' or 'distances' are not defined in storage
+                                 default is the value from the previous call or 36
+        """
+        neigh = self._calc_neighbors(num_neighbors=num_neighbors)
+        distances = neigh["distances"]
+
+        plt.hist(distances.flatten(), bins=bins)
+        plt.xlabel(r"Distance [$\AA$]")
+        plt.ylabel("Neighbor count")
+
+    def shell_distances(self, num_shells=4, num_neighbors=None):
+        """
+        Plot a violin plot of the neighbor distances in shells up to `num_shells`.
+
+        Args:
+            num_shells (int): maximum shell to plot
+            num_neighbors (int): maximum number of neighbors to calculate, when 'shells' or 'distances' are not defined in storage
+                                 default is the value from the previous call or 36
+        """
+        neigh = self._calc_neighbors(num_neighbors=num_neighbors)
+        shells = neigh["shells"]
+        distances = neigh["distances"]
+
+        R = distances.flatten()
+        S = shells.ravel()
+        d = pd.DataFrame(
+            {"distance": R[S < num_shells + 1], "shells": S[S < num_shells + 1]}
+        )
+        sns.violinplot(y=d.shells, x=d.distance, scale="width", orient="h")
+        plt.xlabel(r"Distance [$\AA$]")
+        plt.ylabel("Shell")