Source code for IMDgroup.pymatgen.diffusion.neb

# MIT License
#
# Copyright (c) 2024-2025 Inverse Materials Design Group
#
# Author: Ihor Radchenko <yantar92@posteo.net>
#
# This file is a part of IMDgroup-pymatgen package
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.


"""NEB pair generator for diffusion paths."""

import logging
import warnings
from multiprocessing import Pool
from collections.abc import Sequence
from typing import cast
from alive_progress import alive_bar
import numpy as np
import networkx as nx
from networkx import MultiDiGraph
from networkx.algorithms.cycles import (
    # `_johnson_cycle_search` is a private symbol, so Pyright cannot see it;
    # the public `nx.simple_cycles` alternative is far slower.
    _johnson_cycle_search as johnson_cycle_search,  # pyright: ignore[reportAttributeAccessIssue]
)
from pymatgen.core import Structure, PeriodicSite
from IMDgroup.pymatgen.core.structure import\
    merge_structures, structure_matches
from IMDgroup.pymatgen.transformations.symmetry_clone\
    import SymmetryCloneTransformation
from IMDgroup.pymatgen.core.structure import\
    structure_diff, get_matched_structure

logger = logging.getLogger(__name__)


[docs] class get_neb_pairs_warning(UserWarning): """Warning emitted during NEB pair generation."""
[docs] class NEB_Graph(MultiDiGraph): """Graph representing diffusion paths between structures. Nodes are structure indices. Edges are diffusion paths with attributes ``distance``, ``vector``, and ``energy_barrier``. Attributes: structures: List of Structure objects forming the graph nodes. multithread: Whether to use multithreading for distance matrix. jimage_idxs: Indices of sites for which periodic images are considered when computing displacement vectors. """
[docs] def __init__( self, structures: list[Structure] | None = None, jimage_idxs: list[int] | None = None, multithread: bool = False): """Build a complete NEB diffusion graph. All structures must share the same lattice and have one-to-one site correspondence. When ``jimage_idxs`` is None, edges use the shortest distances between structures under periodic boundary conditions (self-self paths are ignored). When ``jimage_idxs`` is provided, the specified site indices are used for generating multiple periodic images (range -1..1 in each direction), enabling self-self diffusion path discovery. Args: structures: List of structures forming the graph nodes. jimage_idxs: Site indices for periodic image enumeration, or None for shortest-path only. multithread: Whether to use multithreading for distance matrix computation. """ super().__init__() if structures is None: self.structures = [] else: self.structures = structures self.multithread = multithread self.jimage_idxs = jimage_idxs self.__cycle_cache = [] if structures is None: return all_vecs = self.__get_all_vecs() with alive_bar( len(all_vecs), title="Adding diffusion paths to graph") as progress_bar: for from_idx, to_idx, vec in all_vecs: if jimage_idxs is None: self.__add_edge(from_idx, to_idx, vec) else: for jimage in [[i, j, k] for i in range(-1, 2) for j in range(-1, 2) for k in range(-1, 2)]: vec2 = vec.copy() for idx in jimage_idxs: # FIXME: Pymatgen's idx does not do the type properly site_from =\ cast(PeriodicSite, structures[from_idx][idx]).to_unit_cell() site_to = cast(PeriodicSite, structures[to_idx][idx]).to_unit_cell() assert site_from is not None assert site_to is not None lattice = structures[from_idx].lattice vec2[idx] = lattice.get_cartesian_coords( site_to.frac_coords + jimage - site_from.frac_coords) self.__add_edge(from_idx, to_idx, vec2) progress_bar() # pylint: disable=not-callable
def __add_edge(self, from_idx: int, to_idx: int, vector): def get_distance(vec): distance = 0 for idx, v in enumerate(vec): # When we consider self-self paths, there is no way # computing real side displacemnts along the way and # we can only know how much jimage_idxs sites move. # But then such self-self paths cannot be compared to # full distance between distint structures as we do # have other site displacement then. # To be consistent, only consider displacements of # jimage_idxs, when they are provided. if self.jimage_idxs is not None and\ idx not in self.jimage_idxs: continue d = np.linalg.norm(v) if d > 0.5: distance += d return distance distance = get_distance(vector) from_energy = self.structures[from_idx].properties['final_energy'] to_energy = self.structures[to_idx].properties['final_energy'] energy_barrier = to_energy - from_energy self.add_edge( from_idx, to_idx, distance=distance, vector=np.array(vector), energy_barrier=energy_barrier ) if to_idx != from_idx: self.add_edge( to_idx, from_idx, distance=distance, vector=-np.array(vector), energy_barrier=-energy_barrier ) @staticmethod def _get_vec( from_idx: int, to_idx: int, from_struct: Structure, to_struct: Structure, progress_bar=None): """Compute minimal vector connecting from_struct and to_struct. Return (from_idx, to_idx, vec). """ vec = structure_diff(from_struct, to_struct, tol=0, match_first=False) if progress_bar is not None: progress_bar() # pylint: disable=not-callable return (from_idx, to_idx, vec) def __get_all_vecs(self): """Compute all structure diff vectors for the NEB graph.""" if self.multithread: return self._compute_vecs_multithreaded() return self._compute_vecs_singlethreaded() def _compute_vecs_multithreaded(self): """Compute vectors using multithreading.""" with alive_bar( None, title='Computing distance matrix') as progress_bar: with Pool() as pool: # pylint: disable=not-callable progress_bar(len(self.structures) ** 2 / 2) return pool.starmap( self._get_vec, [(from_idx, to_idx, from_struct, to_struct) for from_idx, from_struct in enumerate(self.structures) for to_idx, to_struct in enumerate(self.structures) if from_idx <= to_idx] ) def _compute_vecs_singlethreaded(self): """Compute vectors without multithreading.""" with alive_bar( int(len(self.structures) ** 2 / 2), title='Computing distance matrix') as progress_bar: return [ self._get_vec( from_idx, to_idx, from_struct, to_struct, progress_bar) for from_idx, from_struct in enumerate(self.structures) for to_idx, to_struct in enumerate(self.structures) if from_idx <= to_idx ]
[docs] def connected(self, idxs: list[int] | None = None): """Check whether the graph is connected. Args: idxs: When provided, only check connectivity of these vertex indices. Returns: bool: True if all relevant vertices are reachable. """ n_vertices = len(self.structures) # Visit matrix: False when we cannot reach a site; True - we can. visited = [False for _ in range(n_vertices)] if idxs is None: idxs = [idx for idx, _ in enumerate(self.structures)] queue = [idxs[0]] while len(queue) > 0: from_idx = queue.pop(0) visited[from_idx] = True for _, to_idx in self.edges(from_idx): if not visited[to_idx] and to_idx not in queue: queue.append(to_idx) return all(visited[idx] for idx in idxs)
[docs] def diffusion_path_infinite(self, start_idx): """Check whether a vertex lies on an infinite diffusion path. An infinite path is a cycle whose sum of displacement vectors is non-zero, meaning the diffusing atom can move without bound through the material. Args: start_idx: Vertex index to check. Returns: bool: True if the vertex is on an infinite diffusion path. """ # Note: It is sufficient to check one START_IDX and skip all the # symmetrically equivalent clones - by definition they should # yield exactly the same results. # Of course, such simplification won't work when comparing # non-equivalent starting positions for the diffusion. # visited = np.full(len(self.structures), False) logger.debug( "Searching infinite diffusion paths including %d", start_idx) def _check_cycle(cycle, acc=0, full_cycle=None): if full_cycle is None: full_cycle = cycle # logger.debug( # "tot=%f (%s -- %s)", np.linalg.norm(acc), full_cycle, cycle) if len(cycle) <= 1: assert full_cycle is not None if np.isclose(np.linalg.norm(acc), 0): return False return True for _, to, vec in self.edges(cycle[0], data='vector'): if to == cycle[1] and _check_cycle( cycle[1:], acc + vec, full_cycle): return True return False def _debug_found(cycle, cached=False): logger.debug( "%sFound infinite diffusion path for %d: %s", "(cached) " if cached else "", start_idx, " -> ".join([str(i) for i in cycle]) ) self.__cycle_cache = [ cycle for cycle in self.__cycle_cache if start_idx not in cycle or _check_cycle(cycle) ] # Now, all the cached cycles containing start_idx are valid. # Check if there are any left. for cycle in self.__cycle_cache: if start_idx in cycle: _debug_found(cycle, cached=True) return True components = nx.strongly_connected_components(self) subgraph = None for c in components: if start_idx in c: subgraph = self.subgraph(c) break assert subgraph is not None assert start_idx in subgraph.nodes() n_skipped = 0 max_skipped = int(1E6) for cycle in johnson_cycle_search(subgraph, [start_idx]): if start_idx not in cycle: continue cycle = cycle + [cycle[0]] if _check_cycle(cycle): self.__cycle_cache.append(cycle) _debug_found(cycle, cached=False) return True n_skipped += 1 # logger.debug("skipped closed cycle: %d", n_skipped) if n_skipped > max_skipped: warnings.warn( "Unable to find infinite diffusion path" f" for {start_idx} after attempting" f" {max_skipped} paths" ) break return False
[docs] def all_diffusion_paths_infinite( self, idxs: list[int] | None = None) -> bool: """Check whether all given vertices are on infinite diffusion paths. Structures with the same ``_orig_idx`` property are assumed symmetrically equivalent and checked only once. Args: idxs: Indices to check. Defaults to all vertices. Returns: bool: True if all checked vertices are on infinite paths. """ _checked = [] if idxs is None: idxs = [idx for idx, _ in enumerate(self.structures)] for idx in idxs: struct = self.structures[idx] if struct.properties['_orig_idx'] in _checked: continue if not self.diffusion_path_infinite(idx): return False _checked.append(struct.properties['_orig_idx']) return True
[docs] def get_min_cutoff(self, idx_connected: list[int] | None = None): """Find the smallest edge distance cutoff that keeps the graph connected. Also requires that all vertices be on infinite diffusion paths after applying the cutoff. Args: idx_connected: Indices that must remain connected. Defaults to all vertices. Returns: float: Minimum distance cutoff in Angstrom that satisfies both connectivity and infinite-path constraints. """ # Visit matrix: False when we cannot reach a site; True - we can. visited = [False] * len(self.structures) if idx_connected is None: start_idx = 0 must_visit = [True] * len(self.structures) else: start_idx = idx_connected[0] must_visit = [False] * len(self.structures) for idx in idx_connected: must_visit[idx] = True with alive_bar( None, title='Auto-detecting cutoff'): def bfs1(queue, dist_cutoff): while len(queue) > 0: from_idx = queue.pop(0) if not visited[from_idx]: visited[from_idx] = True distances = [ (dist, to_idx) for _, to_idx, dist in self.edges(from_idx, data='distance') if not visited[to_idx]] for distance, to_idx in sorted(distances): if not visited[to_idx] and not distance > dist_cutoff\ and to_idx not in queue: logger.debug( "coverage: %s -> %s (%fÅ)", from_idx, to_idx, distance ) queue.append(to_idx) visited[start_idx] = True all_distances_sorted = np.unique( [dist for _, _, dist in self.edges(data='distance')]) prev_d = all_distances_sorted[0] tmp = [prev_d] # Avoid iterating over distances that are very close. # Ensure at least 1% increment. for d in all_distances_sorted[1:]: if prev_d * 1.01 < d: prev_d = d tmp.append(d) all_distances_sorted = tmp for max_dist in all_distances_sorted: logger.debug( "Trying to reach all the sites via <=%.2fÅ long paths", max_dist ) queue = [idx for idx, v in enumerate(visited) if v] bfs1(queue, max_dist) all_visited = all(did for must, did in zip(must_visit, visited) if must) all_infinite = True if all_visited: data = [ (from_idx, to_idx, key, data) for from_idx, to_idx, key, data in self.edges(keys=True, data=True) if data['distance'] > max_dist ] for from_idx, to_idx, key, _ in data: self.remove_edge(from_idx, to_idx, key) all_infinite =\ self.all_diffusion_paths_infinite(idx_connected) for from_idx, to_idx, _, d in data: self.add_edges_from([(from_idx, to_idx, d)]) if all_visited and all_infinite: for distance in all_distances_sorted: # Return _larger_ distance to avoid float # comparison precision problems if distance != np.inf and distance > max_dist: logger.debug( "Max required distance %f. Next: %f", max_dist, distance) # Use slightly larger value to avoid # comparison errors. return (max_dist + distance) / 2.0 * 1.01 # cutoff is the largest distance, return something # slightly higher logger.debug( "Cutoff is the largest diffusion path length" ", returning x2") return max_dist * 2 raise AssertionError( f"bfs: This must not happen (visited: {visited})")
def _remove_duplicates( structures: Sequence[Structure | None], idxs: list[int] | None = None, multithread: bool = False, ) -> list[Structure | None]: """Remove duplicates from a list of structures. The first occurrence is kept; subsequent matches are replaced with None. The comparison can be restricted to a subset of site indices. Args: structures: Sequence of structures. ``None`` entries are skipped and preserved in the output. idxs: When provided, only compare proximity of sites at these indices. multithread: Whether to use multithreading for duplicate detection. Returns: list[Structure | None]: Input list with duplicates replaced by None, preserving order. """ uniq_structures = [] logger.info( "gen_neb_pairs: Checking duplicates among %d structures...", len(structures)) if idxs is not None: def __check_idxs(struct1, struct2): for idx in idxs: if struct1[idx].distance(struct2[idx]) > 0.5: return False return True cmp_fun = __check_idxs # Can't pickle local function multithread = False else: cmp_fun = None with alive_bar(len(structures), title='Checking duplicates')\ as progress_bar: for struct in structures: if struct is None: uniq_structures.append(None) progress_bar() # pylint: disable=not-callable continue if structure_matches( struct, uniq_structures, warn=True, cmp_fun=cmp_fun, multithread=multithread): uniq_structures.append(None) else: uniq_structures.append(struct) progress_bar() # pylint: disable=not-callable logger.info( "gen_neb_pairs: Checking duplicates... done (removed %d)", len(structures) - len([x for x in uniq_structures if x is not None])) if len(structures) - len(uniq_structures) > 0: logger.info( "gen_neb_pairs: Using structure enumeration" " preserving the original order (including duplicates)") return uniq_structures
[docs] def get_neb_pairs( structures: list[Structure], prototype: Structure, cutoff: float | None | str = None, remove_compound: bool = False, multithread: bool = False, limit: None | int = None, return_unfiltered: bool = False ) -> list[tuple[Structure, Structure]] | tuple[list[tuple[Structure, Structure]], list[tuple[Structure, Structure]]]: """Construct all unique diffusion NEB pairs from a set of structures. The structures must share the same lattice and be derived from a common ``prototype`` structure. Symmetry operations of the prototype are applied to enumerate equivalent diffusion paths. Args: structures: Candidate structures (e.g. with interstitial atoms at various positions). prototype: Reference structure defining the host lattice and symmetry. cutoff: Maximum allowed displacement distance (Angstrom). When ``'auto'``, the smallest cutoff that keeps all lowest-energy structures connected is determined automatically. remove_compound: When True, remove paths that can be composed from shorter paths (heuristic simplification). multithread: Whether to use multithreading. limit: Maximum number of unique NEB pairs. return_unfiltered: When True, return both the filtered unique pairs and all (symmetry-unfiltered) pairs. Returns: list[tuple[Structure, Structure]]: Pairs of (start, end) structures for NEB calculations. When ``return_unfiltered`` is True, returns ``(unique_pairs, all_pairs)``. """ # Arrange 1-to-1 site matching in all the provided structures # including prototype (needed to detect insertion sites) for idx, struct in enumerate(structures): if struct is not None: structures[idx] =\ get_matched_structure(prototype, struct) uniq_structures = structures # Remove duplicates in the inserted sites. if len(structures[0]) > len(prototype): idxs = list(range(len(prototype), len(structures[0]))) logger.info( "Removing duplicate positions for insertions into prototype") uniq_structures = _remove_duplicates(structures, idxs, multithread) # Remove duplicates. logger.info("Removing symmetry duplicates") uniq_structures = _remove_duplicates( uniq_structures, multithread=multithread) # Inform user about structure numbers logger.info("Assigning indices") for idx, struct in enumerate(uniq_structures): logger.info( "#%d: %s", idx, struct.properties['origin_path'] if struct is not None else "ignore (duplicate)" ) # Compute all the symmetrically equivalent structure clones all_clones = [] for idx, struct in enumerate(uniq_structures): if struct is None: continue logger.info("Enumerating clones in structure #%d", idx) trans = SymmetryCloneTransformation(prototype) clones = trans.get_all_clones(struct, multithread=multithread) for clone_idx, clone in enumerate(clones): clone.properties['_orig_idx'] = idx clone.properties['_clone_idx'] =\ len(all_clones) + clone_idx logger.info( "#%d clones assigned indices #%d..#%d", idx, len(all_clones), len(all_clones) + len(clones) - 1 ) all_clones += clones logger.info("Found %d clones", len(all_clones)) # Build diffusion graph connecting all the clones neb_graph = NEB_Graph( all_clones, jimage_idxs=list(range(len(prototype), len(structures[0]))) if len(structures[0]) > len(prototype) else None, multithread=multithread) # Select structures that must remain connected in the graph # Currently, we simply maintain connectivity of the lowest-energy # structures (maybe through higher-energy intermediates). energies = [] assert neb_graph.structures is not None for idx, struct in enumerate(neb_graph.structures): energy = struct.properties.get('final_energy') if energy is None: origin_path = struct.properties.get('origin_path') raise ValueError(f"Energy data is missing for {origin_path}") logger.info("Energy %d: %f", idx, energy) energies.append(energy) # Only take the lowest-energy structure # +1E-9 is to counter floating point error where structure clones # may have slight variation of energies energy_threshold = sorted(np.unique(energies))[0] + 1E-9 logger.info("Energy theshold: %f", energy_threshold) low_en_idxs = [ idx for idx, s in enumerate(neb_graph.structures) if not s.properties['final_energy'] > energy_threshold] logger.info( "Diffusion graph connectivity will be limited" " to lowest-energy configurations: %s", low_en_idxs ) # If diffusion distance cutoff is not provided, detect it # automatically, choosing the minimal possible cutoff that # maintains structure connectivity. if cutoff == 'auto': logger.info("Determining minimal possible diffusion distance cutoff") cutoff = neb_graph.get_min_cutoff(low_en_idxs) logger.info('Found optimal cutoff to cover all sites: %f', cutoff) # Only keep graph egdes shorter than cutoff assert isinstance(cutoff, float) n_edges = 0 # Cast to list because MultiDiGraph.edges cannot handle # changing edges on the fly. for from_idx, to_idx, key, edge_len in\ list(neb_graph.edges(data='distance', keys=True)): if not edge_len < cutoff: neb_graph.remove_edge(from_idx, to_idx, key) else: n_edges += 1 logger.info("Found %d paths shorter than cutoff (%f)", n_edges, cutoff) def _connected_and_infinite(): is_connected = neb_graph.connected(low_en_idxs) is_infinite = False if is_connected: is_infinite =\ neb_graph.all_diffusion_paths_infinite(low_en_idxs) return is_connected and is_infinite # Loop over largest known (or estimated as energy difference) # barriers and remove as many as possible. Always keep <=0 barriers. n_removed = 0 # We use rounding here because exact values will have floating # point errors. barriers = [(np.round(data['energy_barrier'], 5), np.round(data['distance'], 2), from_idx, to_idx, key) for from_idx, to_idx, key, data in neb_graph.edges(keys=True, data=True) if data['energy_barrier'] >= 1E-9] logger.info('Removing high-energy barriers') with alive_bar( len(barriers), title='Removing high-energy barriers', ) as progress_bar: # Try removing one by one, starting from the highest. # For the same energy barrier, remove longer paths first. for en, _, from_idx, to_idx, key in sorted(barriers, reverse=True): data = neb_graph.edges[from_idx, to_idx, key] neb_graph.remove_edge(from_idx, to_idx, key) if _connected_and_infinite(): logger.debug( "%d -> %d (%d): removed (%feV)", from_idx, to_idx, key, en) n_removed += 1 else: logger.info( "%d -> %d (%d): kept (%feV)", from_idx, to_idx, key, en) neb_graph.add_edges_from([(from_idx, to_idx, data)]) progress_bar() # pylint: disable=not-callable logger.info("Removed %d high-energy barriers", n_removed) # If requested, remove as many as possible diffusion paths, # starting from the longest, while keeping graph connectivity. if remove_compound: logger.info("Removing compound paths") edges = [(distance, from_idx, to_idx, key) for from_idx, to_idx, key, distance in neb_graph.edges(keys=True, data='distance')] n_edges = len(edges) with alive_bar( n_edges, title='Removing compound paths') as progress_bar: for edge_len, from_idx, to_idx, key in sorted(edges, reverse=True): data = neb_graph.edges[from_idx, to_idx, key] neb_graph.remove_edge(from_idx, to_idx, key) if _connected_and_infinite(): logger.debug( "%d -> %d (%d): removed (%fÅ)", from_idx, to_idx, key, edge_len ) n_edges -= 1 else: logger.info( "%d -> %d (%d): kept (%fÅ)", from_idx, to_idx, key, edge_len ) neb_graph.add_edges_from([(from_idx, to_idx, data)]) progress_bar() # pylint: disable=not-callable logger.info("Found %d non-compound paths", n_edges) logger.info( "Final diffusion graph: %s", [(from_idx, to_idx) for from_idx, to_idx, _ in neb_graph.edges] ) for from_idx, to_idx, key, data in neb_graph.edges(data=True, keys=True): max_v = [0, 0, 0] for v in data['vector']: if np.linalg.norm(v) > np.linalg.norm(max_v): max_v = v with np.printoptions(precision=2, suppress=True): logger.info( "%d -> %d (%d): %fÅ; %feV; %s", from_idx, to_idx, key, data['distance'], data['energy_barrier'], max_v ) # Get rid of symmetrically equivalent diffusion paths. logger.info("Searching unique diffusion paths") unique_edges = [] merged_pairs = [] _known_dists = [] edges = [] dists = [] for from_idx, to_idx, key, dist in\ neb_graph.edges(data='distance', keys=True): edges.append((from_idx, to_idx, key)) dists.append(dist) n_edges = len(edges) with alive_bar( n_edges, title='Removing equivalent paths') as progress_bar: for dist, (from_idx, to_idx, key) in sorted(zip(dists, edges)): logger.debug('Processing path %d -> %d: %fÅ', from_idx, to_idx, dist) if np.isclose(dist, 0): logger.info('Skipping too short path') continue # Two paths are equivalent when they (1) form a # symmetrically unique structure when combined; # (2) when they length/vector is the same. merged = merge_structures( [all_clones[from_idx], all_clones[to_idx]]) # Equivalent paths must have the same length. close_pair_idxs = [ idx for idx, d in enumerate(_known_dists) if np.isclose(dist, d)] if len(close_pair_idxs) > 0 and structure_matches( merged, [merged_pairs[idx] for idx in close_pair_idxs], multithread=multithread): logger.debug("%s path is non-unique", (from_idx, to_idx)) progress_bar() # pylint: disable=not-callable continue unique_edges.append((dist, from_idx, to_idx, key)) merged_pairs.append(merged) _known_dists.append(dist) progress_bar() # pylint: disable=not-callable if limit and len(unique_edges) >= limit: break logger.info("Found %d unique paths", len(unique_edges)) # Sort by lentgh unique_edges = sorted(unique_edges) def get_pair(from_idx, to_idx, key): origin_struct = all_clones[from_idx] target_struct = origin_struct.copy() for idx, v in enumerate( neb_graph.edges[from_idx, to_idx, key]['vector']): target_struct.translate_sites( [idx], v, frac_coords=False, to_unit_cell=False) return (origin_struct, target_struct) # Use full displacement vector to produce final diffusion point. pairs = [] for _, from_idx, to_idx, key in unique_edges: pairs.append(get_pair(from_idx, to_idx, key)) if return_unfiltered: unfiltered_pairs = [ get_pair(from_idx, to_idx, key) for from_idx, to_idx, key in neb_graph.edges(keys=True) ] return pairs, unfiltered_pairs return pairs