Add peak labels to free energy plots

src/gemdat/path.py

@@ -7,6 +7,9 @@
 
 import networkx as nx
 import numpy as np
+from pymatgen.core import Structure
+
+from gemdat.volume import Volume
 
 
 @dataclass
@@ -15,13 +18,34 @@ class Pathway:
 
     Attributes
     ----------
-    path: list[tuple]
-        List of coordinates of the sites definint the path
-    path_energy: list[float]
+    sites: list[tuple]
+        List of voxel coordinates of the sites defining the path
+    cart_sites: list[tuple]
+        List of cartesian coordinates of the sites defining the path
+    energy: list[float]
         List of the energy along the path
+    nearest_structure_coord: list[np.ndarray]
+        List of cartesian coordinates of the closest site of the reference structure
+    nearest_structure_label: list[str]
+        List of the label of the closest site of the reference structure
     """
     sites: list[tuple[int, int, int]]
+    cart_sites: list[tuple[float, float, float]]
     energy: list[float]
+    nearest_structure_coord: list[np.ndarray]
+    nearest_structure_label: list[str]
+
+    def cartesian_path(self, vol: Volume):
+        """Convert voxel coordinates to cartesian coordinates.
+
+        Parameters
+        ----------
+        vol : Volume
+            Volume object containing the grid information
+        """
+        self.cart_sites = [
+            tuple(np.asarray(site) * vol.resolution) for site in self.sites
+        ]
 
     def wrap(self, F: np.ndarray):
         """Wrap path in periodic boundary conditions in-place.
@@ -34,6 +58,29 @@ def wrap(self, F: np.ndarray):
         X, Y, Z = F.shape
         self.sites = [(x % X, y % Y, z % Z) for x, y, z in self.sites]
 
+    def path_over_structure(self, structure: Structure, structure_map: dict):
+        """Find the nearest site of the structure to the path sites.
+
+        Parameters
+        ----------
+        structure : Structure
+            Reference structure
+        structure_map : dict
+            Dictionary of the closest site of the structure to each coordinate point
+        """
+        if not self.cart_sites:
+            raise ValueError(
+                'Cartesian coordinates of the path sites are not defined')
+
+        self.nearest_structure_label = [
+            structure.labels[structure_map.get(site, None)]
+            for site in self.cart_sites
+        ]
+        self.nearest_structure_coord = [
+            structure.cart_coords[structure_map.get(site, None)]
+            for site in self.cart_sites
+        ]
+
     @property
     def cost(self) -> float:
         """Calculate the path cost."""
@@ -124,7 +171,7 @@ def optimal_path(
                                     weight='weight')
     path_energy = [F_graph.nodes[node]['energy'] for node in optimal_path]
 
-    path = Pathway(optimal_path, path_energy)
+    path = Pathway(optimal_path, [], path_energy, [], [])
 
     return path
 
@@ -159,7 +206,7 @@ def find_best_perc_path(F: np.ndarray,
     # Find percolation using virtual images along the required dimensions
     if not any([percolate_x, percolate_y, percolate_z]):
         print('Warning: percolation is not defined')
-        return Pathway([], [])
+        return Pathway([], [], [], [], [])
 
     # Tile the grind in the percolation directions
     F = np.tile(F, (1 + percolate_x, 1 + percolate_y, 1 + percolate_z))
@@ -174,7 +221,7 @@ def find_best_perc_path(F: np.ndarray,
 
     # Find the lowest cost path that percolates along the x dimension
     best_cost = float('inf')
-    best_path = Pathway([], [])
+    best_path = Pathway([], [], [], [], [])
 
     for starting_point in peaks:
 

src/gemdat/plots/matplotlib/_paths.py

@@ -18,10 +18,37 @@ def energy_along_path(*, path: Pathway) -> plt.Figure:
     fig : matplotlib.figure.Figure
         Output figure
     """
-    fig, ax = plt.subplots()
+    fig, ax = plt.subplots(figsize=(10, 5))
 
     ax.plot(range(len(path.energy)), path.energy, marker='o', color='r')
-    ax.set(xlabel='Step', ylabel='Free energy [eV]')
+    ax.set(ylabel='Free energy [eV]')
+    if path.nearest_structure_label:
+        # Remove repeated labels that would bloat the figure
+        clean_xlabel = [
+            f'{", ".join([f"{coord:.1f}" for coord in path.nearest_structure_coord[i]])}'
+            if (path.nearest_structure_coord[i]
+                != path.nearest_structure_coord[i - 1]).any() else ''
+            for i in range(len(path.sites))
+        ]
+        non_empty_ticks = [
+            i for i, label in enumerate(clean_xlabel) if label != ''
+        ]
+        ax.set_xticks(non_empty_ticks)
+        ax.set_xticklabels([clean_xlabel[i] for i in non_empty_ticks],
+                           rotation=45)
+        # and add on top the site labels
+        clean_xlabel_up = [
+            f'{path.nearest_structure_label[i]}' if
+            (path.nearest_structure_coord[i]
+             != path.nearest_structure_coord[i - 1]).any() else ''
+            for i in range(len(path.sites))
+        ]
+        ax_up = ax.twiny()
+        ax_up.set_xlim(ax.get_xlim())
+        ax_up.set_xticks(non_empty_ticks)
+        ax_up.set_xticklabels([clean_xlabel_up[i] for i in non_empty_ticks],
+                              rotation=45)
+        ax_up.get_yaxis().set_visible(False)
 
     return fig
 

src/gemdat/volume.py

@@ -12,6 +12,7 @@
 from pymatgen.io.vasp import VolumetricData
 from rich.progress import track
 from scipy.constants import physical_constants
+from scipy.spatial import cKDTree
 from skimage.feature import blob_dog
 from skimage.measure import regionprops
 
@@ -122,6 +123,51 @@ def find_peaks(
 
         return coords[:, 0:3].astype(int)
 
+    def nearest_structure_reference(self, structure: Structure):
+        """Find distance and index of the nearest site of the structure for
+        each voxel using a KD-tree.
+
+        Parameters
+        ----------
+        structure : pymatgen.core.structure.Structure
+            Structure of the material to use as reference for nearest site
+
+        Returns
+        -------
+        nearest_structure_map : dict
+            Dictionary that maps each voxel to the index of the nearest site of the structure
+        """
+        # In order to accomodate the periodicity, include the images of the structure sites
+        periodic_structure = []
+        periodic_ids = []
+        for dx in [-1, 0, 1]:
+            for dy in [-1, 0, 1]:
+                for dz in [-1, 0, 1]:
+                    periodic_structure.extend(
+                        structure.cart_coords + np.array([
+                            self.lattice.a * dx, self.lattice.b *
+                            dy, self.lattice.c * dz
+                        ]) * self.resolution)
+                    # store the id of the site in the original structure
+                    periodic_ids.extend(
+                        [i for i in range(len(structure.cart_coords))])
+        # Create a KD-tree from the structure
+        kd_tree = cKDTree(periodic_structure)
+
+        nearest_structure_map = {}
+        possible_sites = np.array(
+            np.meshgrid(*[range(dim) for dim in self.data.shape])).T.reshape(
+                -1, len(self.data.shape)) * self.resolution
+
+        for site in possible_sites:
+            # Query the KD-tree to find the index of the nearest periodic site
+            nearest_structure_index = kd_tree.query(site)[1]
+            # store the index of the corrisponding site in the original structure
+            nearest_structure_map[tuple(
+                site)] = periodic_ids[nearest_structure_index]
+
+        return nearest_structure_map
+
     def to_vasp_volume(self, structure: Structure, *,
                        filename: Optional[str]) -> VolumetricData:
         """Convert to vasp volume.

tests/integration/path_test.py

@@ -5,6 +5,7 @@
 import numpy as np
 import pytest
 
+from gemdat.io import load_known_material
 from gemdat.path import find_best_perc_path
 from gemdat.volume import trajectory_to_volume
 
@@ -29,3 +30,17 @@ def test_find_best_perc_path(vasp_full_vol):
 
     assert isclose(path.cost, 36.39301483423)
     assert path.start_site == (30, 23, 14)
+
+
+@pytest.vaspxml_available  # type: ignore
+def test_nearest_structure_reference(vasp_full_vol):
+    structure = load_known_material('argyrodite')
+    nearest_structure_map = vasp_full_vol.nearest_structure_reference(
+        structure)
+
+    assert nearest_structure_map[(0, 0, 0)] == 21
+    assert nearest_structure_map[(14.7, 0.8999999999999999,
+                                  0.8999999999999999)] == 9
+    assert nearest_structure_map[(18.3, 0.8999999999999999, 9.6)] == 37
+    assert (18.3, 0.899999999999999, 9.6) not in nearest_structure_map
+    assert len(nearest_structure_map) == 71874