Repository: GEMDAT-repos/GEMDAT (language: python)

Issue:
Radial distribution functions
These are calculated in `calc_rdfs.m`. Data can be verified against `rdf.mat`.

- [ ] `rdf.distributions`
- [ ] `rdf.integrated`
- [ ] `rdf.rdf_names`
- [ ] `rdf.elements`
- [ ] `rdf.max_dist`
- [ ] `rdf.resolution`
- [ ] `rdf.total`


Discussion / hints:

Fix patch:
diff --git a/pyproject.toml b/pyproject.toml
index 54c1d1cc..3d5114a4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -34,6 +34,7 @@ classifiers = [
 dependencies = [
   'pymatgen >= 2023.7.14',
   'numpy',
+  'rich',
   'scipy',
   'matplotlib',
   'MDAnalysis',
diff --git a/src/rdf.py b/src/rdf.py
new file mode 100644
index 00000000..d36fa80b
--- /dev/null
+++ b/src/rdf.py
@@ -0,0 +1,163 @@
+from collections import defaultdict
+
+import numpy as np
+from gemdat import SimulationData, SitesData
+from pymatgen.core import Structure
+from rich.progress import track
+
+
+def _uniqify_labels(arr, labels: list[str]) -> np.ndarray:
+    """Helper function to uniqify labels."""
+    unique_labels = list(set(labels))
+    mapping = np.array([-1] + [unique_labels.index(label) for label in labels])
+
+    palette = np.arange(len(labels), dtype=int)
+
+    index = np.digitize(arr, palette, right=True)
+    return mapping[index]
+
+
+def _get_states(labels: list[str]) -> dict[int, str]:
+    """Helper function to generate a list of states from the labels."""
+    unique_labels = list(set(labels))
+
+    states = {}
+
+    r = [-1] + list(range(len(unique_labels)))
+
+    for i in r:
+        for j in r:
+            for k in r:
+                if i != -1:
+                    state = '@' + unique_labels[i]
+                elif j == -1 or k == -1:
+                    state = '~>' + unique_labels[j]
+                else:
+                    state = unique_labels[j] + '->' + unique_labels[k]
+
+                states[int(i * 1e6 + j * 1e3 + k)] = state
+
+    return states
+
+
+def _get_states_array(sites: SimulationData, labels: list[str]) -> np.ndarray:
+    """Helper function to generate integer array of transition states."""
+    atom_sites = _uniqify_labels(sites.atom_sites, labels)
+    atom_sites_from = _uniqify_labels(sites.atom_sites_from, labels)
+    atom_sites_to = _uniqify_labels(sites.atom_sites_to, labels)
+
+    states_array = (atom_sites * 1e6 + atom_sites_from * 1e3 +
+                    atom_sites_to).astype(int)
+
+    return states_array
+
+
+def _get_symbol_indices(structure: Structure) -> dict[str, np.ndarray]:
+    """Helper function to generate symbol indices."""
+    symbols = structure.symbol_set
+    return {
+        symbol:
+        np.argwhere([e.name == symbol for e in structure.species]).flatten()
+        for symbol in symbols
+    }
+
+
+def calculate_rdfs(
+        *,
+        data: SimulationData,
+        sites: SitesData,
+        diff_coords: np.ndarray,
+        n_steps: int,
+        equilibration_steps: int,
+        max_dist: float = 5.0,
+        resolution: float = 0.1) -> dict[str, dict[str, np.ndarray]]:
+    """
+    Parameters
+    ----------
+    data : SimulationData
+        Input simulation data
+    sites : SitesData
+        Input sites data
+    diff_coords : np.ndarray
+        Input coordinates for diffusing element (extras)
+    n_steps : int
+        Total number of simulation steps (extras)
+    equilibration_steps : int
+        Number of equilibration steps (extras)
+    max_dist : float, optional
+        Max distance for rdf calculation
+    resolution : float, optional
+        Width of the bins
+
+    Returns
+    -------
+    rdfs : dict[str, np.ndarray]
+        Dictionary with rdf arrays per symbol
+    """
+    structure = data.structure
+    lattice = structure.lattice
+
+    coords = data.trajectory_coords[equilibration_steps:]
+
+    states2str = _get_states(sites.structure.labels)
+    states_array = _get_states_array(sites, sites.structure.labels)
+    symbol_indices = _get_symbol_indices(structure)
+
+    bins = np.arange(0, max_dist + resolution, resolution)
+    length = len(bins) + 1
+
+    rdfs: dict[tuple[str, str],
+               np.ndarray] = defaultdict(lambda: np.zeros(length, dtype=int))
+
+    for i in track(range(n_steps), transient=True):
+
+        t_coords = coords[i]
+        t_diff_coords = diff_coords[i]
+
+        dists = lattice.get_all_distances(t_diff_coords, t_coords)
+
+        rdf = np.digitize(dists, bins, right=True)
+
+        states = np.unique(states_array[i], axis=0)
+
+        t_states = states_array[i]
+
+        for state in states:
+            k_idx = np.argwhere(t_states == state)
+            state_str = states2str[state]
+
+            for symbol, symbol_idx in symbol_indices.items():
+                rdf_state = rdf[k_idx, symbol_idx].flatten()
+                rdfs[state_str, symbol] += np.bincount(rdf_state,
+                                                       minlength=length)
+
+    new_rdf: dict[str, dict[str, np.ndarray]] = defaultdict(dict)
+
+    for (k0, k1), v in rdfs.items():
+        # Drop last element with distance > max_dist
+        new_rdf[k0][k1] = v[:-1]
+
+    return new_rdf
+
+
+def plot_rdf(rdfs: dict[str, np.ndarray], name: str | None = None):
+    """Plot radial distribution function.
+
+    Parameters
+    ----------
+    rdfs : dict[str, np.ndarray]
+        Dictionary with rdf array per symbol
+    """
+    import matplotlib.pyplot as plt
+
+    fig, ax = plt.subplots()
+
+    for symbol, rdf in rdfs.items():
+        ax.plot(rdf[:-1], label=symbol)
+
+    suffix = f' ({name})' if name else ''
+
+    ax.legend()
+    ax.set(title=f'Radial distribution function{suffix}',
+           xlabel='Distance (Ang)',
+           ylabel='Counts')


Test patch:
diff --git a/tests/integration_test.py b/tests/integration_test.py
index 7bd62bb0..e1c8fa24 100644
--- a/tests/integration_test.py
+++ b/tests/integration_test.py
@@ -1,7 +1,7 @@
 """
 Run integration test with:
 
-VASP_XML=/run/media/stef/Scratch/md-analysis-matlab-example/vasprun.xml pytest
+VASP_XML=/home/stef/md-analysis-matlab-example/vasprun.xml pytest
 """
 
 import os
@@ -9,14 +9,17 @@
 
 import numpy as np
 import pytest
-from gemdat import SimulationData
+from gemdat import SimulationData, SitesData
+from gemdat.io import load_known_material
+from gemdat.rdf import calculate_rdfs
 from gemdat.volume import trajectory_to_volume
 
 VASP_XML = os.environ.get('VASP_XML')
 
 vaspxml_available = pytest.mark.skipif(
     VASP_XML is None,
-    reason='vasprun.xml test sim_data is required for this test.')
+    reason='Simulation data from vasprun.xml example is required for this test.'
+)
 
 
 @pytest.fixture
@@ -38,6 +41,31 @@ def gemdat_results():
     return (data, extras)
 
 
+@pytest.fixture
+def gemdat_results_subset():
+    # Reduced number of time steps for slow calculations
+    equilibration_steps = 74000
+    diffusing_element = 'Li'
+    diffusion_dimensions = 3
+    z_ion = 1
+
+    data = SimulationData.from_vasprun(VASP_XML)
+
+    extras = data.calculate_all(
+        equilibration_steps=equilibration_steps,
+        diffusing_element=diffusing_element,
+        z_ion=z_ion,
+        diffusion_dimensions=diffusion_dimensions,
+    )
+
+    return (data, extras)
+
+
+@pytest.fixture
+def structure():
+    return load_known_material('argyrodite')
+
+
 @vaspxml_available
 def test_volume(gemdat_results):
     data, extras = gemdat_results
@@ -76,14 +104,9 @@ def test_tracer(gemdat_results):
 
 
 @vaspxml_available
-def test_sites(gemdat_results):
-    from gemdat.io import load_known_material
-    from gemdat.sites import SitesData
-
+def test_sites(gemdat_results, structure):
     data, extras = gemdat_results
 
-    structure = load_known_material('argyrodite')
-
     sites = SitesData(structure)
     sites.calculate_all(data=data, extras=extras)
 
@@ -160,3 +183,32 @@ def test_sites(gemdat_results):
     assert sites.coll_matrix[0, 0] == 437
 
     assert sites.multi_coll.sum() == 174380
+
+
+@vaspxml_available
+def test_rdf(gemdat_results_subset, structure):
+    data, extras = gemdat_results_subset
+
+    structure = load_known_material('argyrodite')
+
+    sites = SitesData(structure)
+    sites.calculate_all(data=data, extras=extras)
+
+    rdfs = calculate_rdfs(
+        data=data,
+        sites=sites,
+        diff_coords=extras.diff_coords,
+        n_steps=extras.n_steps,
+        equilibration_steps=extras.equilibration_steps,
+        max_dist=5,
+    )
+
+    expected_states = {'~>Li48h', '@Li48h', 'Li48h->Li48h'}
+    expected_symbols = set(data.structure.symbol_set)
+
+    assert isinstance(rdfs, dict)
+
+    for state, rdf in rdfs.items():
+        assert state in expected_states
+        assert set(rdf.keys()) == expected_symbols
+        assert all(len(arr) == 51 for arr in rdf.values())