Implement energy landscapes and pathways from simulation

mkdocs.yml

@@ -12,6 +12,7 @@ nav:
     - Simple Example: notebooks/simple.ipynb
     - Loading a CIF File: notebooks/CIF.ipynb
     - Autogenerated Sites: notebooks/auto.ipynb
+    - Pathways: notebooks/paths.ipynb
   - Python API:
     - gemdat: api/gemdat.md
     - gemdat.collective: api/gemdat_collective.md

src/gemdat/path.py

@@ -0,0 +1,211 @@
+"""This module contains classes for computing optimal and percolating paths
+between sites."""
+
+from __future__ import annotations
+
+import networkx as nx
+import numpy as np
+
+
+class Pathway:
+    """Container class for paths between sites.
+
+    Attributes
+    ----------
+    path: list[tuple]
+        List of coordinates of the path
+    path_energy: list[float]
+        Energy along the path
+    """
+
+    def __init__(self, sites: list[tuple], energy: list[float]):
+        """Store event data for jumps and transitions between sites.
+
+        Parameters
+        ----------
+        sites: list[tuple]
+            List of coordinates of the sites definint the path
+        energy: list[float]
+            List of the energy along the path
+        """
+        self.sites = sites
+        self.energy = energy
+
+    def wrap(self, F: np.ndarray) -> Pathway:
+        """Wrap path in periodic boundary conditions.
+
+        Parameters
+        ----------
+        F: np.ndarray
+            Grid in which the path sites will be wrapped
+        """
+        X, Y, Z = F.shape
+        self.sites = [(x % X, y % Y, z % Z) for x, y, z in self.sites]
+
+        return self
+
+
+def _wrap_pbc(coord: np.ndarray, size: tuple) -> np.ndarray:
+    """Wraps coordinates in periodic boundaries.
+
+    Parameters
+    ----------
+    coord: np.ndarray
+        Coordinates of one or multiple sites
+    size: tuple
+        Boundary size
+
+    Returns:
+    -------
+    wrapped_coord: np.ndarray
+        Coordinates inside the PBC
+    """
+    wrapped_coord = coord % size
+    return wrapped_coord
+
+
+def free_energy_graph(F: np.ndarray,
+                      max_energy_threshold: float = 1e20,
+                      diagonal: bool = True) -> nx.Graph:
+    """Compute the graph of the free energy for networkx functions.
+
+    Parameters
+    ----------
+    F : np.ndarray
+        Free energy on the 3d grid
+    max_energy_threshold : float, optional
+        Maximum energy threshold for the path to be considered valid
+    diagonal : bool
+        If True, allows diagonal grid moves
+
+    Returns
+    -------
+    G : nx.Graph
+        Graph of free energy
+    """
+
+    # Define possible movements in 3D space
+    movements = [(1, 0, 0), (-1, 0, 0), (0, 1, 0), (0, -1, 0), (0, 0, 1),
+                 (0, 0, -1)]
+    if diagonal:
+        movements.extend([(1, 1, 0), (-1, -1, 0), (1, -1, 0), (-1, 1, 0),
+                          (1, 0, 1), (-1, 0, -1), (1, 0, -1), (-1, 0, 1),
+                          (0, 1, 1), (0, -1, -1), (0, 1, -1), (0, -1, 1),
+                          (1, 1, 1), (-1, -1, -1), (1, -1, -1), (-1, 1, 1)])
+
+    G = nx.Graph()
+    for index, Fi in np.ndenumerate(F):
+        if 0 <= Fi < max_energy_threshold:
+            G.add_node(index, energy=Fi)
+    for node in G.nodes:
+        for move in movements:
+            neighbor = tuple(_wrap_pbc(node + np.array(move), F.shape))
+            if neighbor in G.nodes:
+                G.add_edge(node, neighbor, weight=F[neighbor])
+
+    return G
+
+
+def optimal_path(
+    F_graph: nx.Graph,
+    start: tuple,
+    end: tuple,
+) -> Pathway:
+    """Calculate the shortest cost-effective path from start to end using
+    Networkx library.
+
+    Parameters
+    ----------
+    F_graph : nx.Graph
+        Graph of the free energy
+    start : tuple
+        Coordinates of the starting point
+    end: tuple
+        Coordinates of the ending point
+
+    Returns
+    -------
+    path: Pathway
+        Optimal path on the graph between start and end
+    """
+
+    optimal_path = nx.shortest_path(F_graph,
+                                    source=start,
+                                    target=end,
+                                    weight='weight')
+    path_energy = [F_graph.nodes[node]['energy'] for node in optimal_path]
+
+    path = Pathway(optimal_path, path_energy)
+
+    return path
+
+
+def find_best_perc_path(F: np.ndarray,
+                        peaks: np.ndarray,
+                        percolate_x: bool = True,
+                        percolate_y: bool = False,
+                        percolate_z: bool = False) -> Pathway:
+    """Calculate the best percolating path.
+
+    Parameters
+    ----------
+    F : np.ndarray
+        Energy grid that will be used to calculate the shortest path
+    peaks : np.ndarray
+        List of the peaks that correspond to high probability regions
+    percolate_x : bool
+        If True, consider paths that percolate along the x dimension
+    percolate_y : bool
+        If True, consider paths that percolate along the y dimension
+    percolate_z : bool
+        If True, consider paths that percolate along the z dimension
+
+    Returns
+    -------
+    best_percolating_path: Pathway
+        Optimal path that percolates the graph in the specified directions
+    """
+    xyz_real = F.shape
+
+    # 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([], [])
+
+    # Tile the grind in the percolation directions
+    F = np.tile(F, (1 + percolate_x, 1 + percolate_y, 1 + percolate_z))
+    X, Y, Z = F.shape
+
+    # Get F on a graph
+    F_graph = free_energy_graph(F, max_energy_threshold=1e7, diagonal=True)
+
+    # Find the lowest cost path that percolates along the x dimension
+    total_energy_cost = float('inf')
+    # reaching the percolating image
+    image = tuple(
+        x * px
+        for x, px in zip(xyz_real, (percolate_x, percolate_y, percolate_z)))
+
+    for starting_point in peaks:
+
+        # Get the end point which is a periodic image of the peak
+        end_point = starting_point + image
+
+        # Find the shortest percolating path through this peak
+        try:
+            path = optimal_path(
+                F_graph,
+                tuple(starting_point),
+                tuple(end_point),
+            )
+        except nx.NetworkXNoPath:
+            continue
+
+        # Calculate the path cost
+        path_cost = np.sum(path.energy)
+
+        if path_cost < total_energy_cost:
+            total_energy_cost = path_cost
+            best_percolating_path = path
+
+    return best_percolating_path

src/gemdat/plots/__init__.py

@@ -8,6 +8,10 @@
     jumps_3d,
     jumps_3d_animation,
 )
+from .matplotlib._paths import (
+    energy_along_path,
+    path_on_grid,
+)
 from .matplotlib._rdf import radial_distribution
 from .matplotlib._vibration import frequency_vs_occurence
 
@@ -21,6 +25,7 @@
     jumps_vs_distance,
     jumps_vs_time,
 )
+from .plotly._paths import path_on_landscape
 from .plotly._vibration import vibrational_amplitudes
 
 __all__ = [
@@ -36,4 +41,7 @@
     'jumps_3d',
     'jumps_3d_animation',
     'radial_distribution',
+    'energy_along_path',
+    'path_on_grid',
+    'path_on_landscape',
 ]

src/gemdat/plots/matplotlib/_paths.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+import matplotlib.pyplot as plt
+
+from gemdat.path import Pathway
+
+
+def energy_along_path(*, path: Pathway) -> plt.Figure:
+    """Plot energy along specified path.
+
+    Parameters
+    ----------
+    path : Pathway
+        Pathway object containing the energy along the path
+
+    Returns
+    -------
+    fig : matplotlib.figure.Figure
+        Output figure
+    """
+    fig, ax = plt.subplots()
+
+    ax.plot(range(len(path.energy)), path.energy, marker='o', color='r')
+    ax.set(xlabel='Step', ylabel='Energy')
+
+    return fig
+
+
+def path_on_grid(*, path: Pathway) -> plt.Figure:
+    """Plot the 3d coordinates of the points that define a path.
+
+    Parameters
+    ----------
+    path : Pathway
+        Pathway to plot
+
+    Returns
+    -------
+    fig : matplotlib.figure.Figure
+        Output figure
+    """
+
+    # Create a colormap to visualize the path
+    colormap = plt.get_cmap()
+    normalize = plt.Normalize(0, len(path.energy))
+
+    fig, ax = plt.subplots()
+    ax = fig.add_subplot(111, projection='3d')
+
+    path_x, path_y, path_z = zip(*path.sites)
+
+    for i in range(len(path.energy) - 1):
+        ax.plot(path_x[i:i + 1],
+                path_y[i:i + 1],
+                path_z[i:i + 1],
+                color=colormap(normalize(i)),
+                marker='o',
+                linestyle='-')
+
+    ax.set_xlabel('X')
+    ax.set_ylabel('Y')
+    sm = plt.cm.ScalarMappable(cmap=colormap, norm=normalize)
+    sm.set_array([])
+    cbar = plt.colorbar(sm, ax=ax)
+    cbar.set_label('Steps')
+
+    return fig

src/gemdat/plots/plotly/_paths.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+import numpy as np
+import plotly.graph_objects as go
+
+from gemdat.path import Pathway
+from gemdat.volume import Structure, Volume
+
+from ._density import density
+
+
+def path_on_landscape(
+    vol: Volume,
+    path: Pathway,
+    structure: Structure,
+) -> go.Figure:
+    """Ploth path over free energy.
+
+    Uses plotly as plotting backend.
+
+    Arguments
+    ---------
+    vol : Volume
+        Input volume to create the landscape
+    path : Pathway
+        Pathway to plot
+    structure : Structure
+        Input structure
+
+    Returns
+    -------
+    fig : go.Figure
+        Output as plotly figure
+    """
+
+    fig = density(vol.to_probability(), structure)
+
+    path = np.array(path.sites) * vol.resolution
+    x_path, y_path, z_path = path.T
+
+    # Color path and endpoints differently
+    color = ['red' for _ in x_path]
+    color[0] = 'blue'
+    color[-1] = 'blue'
+
+    fig.add_trace(
+        go.Scatter3d(
+            x=x_path,
+            y=y_path,
+            z=z_path,
+            mode='markers+lines',
+            line={'width': 3},
+            marker={
+                'size': 6,
+                'color': color,
+                'symbol': 'circle',
+                'opacity': 0.9
+            },
+            name='Path',
+        ))
+
+    return fig