reaxkit 1.0.0__py3-none-any.whl

reaxkit/analysis/composed/RDF_analyzer.py

@@ -0,0 +1,560 @@
+"""
+Radial Distribution Function (RDF) analysis utilities.
+
+This module provides single-frame and multi-frame RDF calculations for
+ReaxFF trajectories using either the FREUD or OVITO backends, as well as
+helper utilities for extracting RDF-derived structural descriptors.
+
+Typical use cases include:
+
+- computing total or partial RDFs for selected frames
+- averaging RDFs across multiple frames
+- extracting RDF-based properties such as peak positions or integrated areas
+"""
+
+
+from __future__ import annotations
+from typing import Iterable, Optional, Sequence, Tuple, List, Dict, Union
+import numpy as np
+import pandas as pd
+
+# FREUD backend (pip install freud-analysis)
+import freud
+
+from reaxkit.io.handlers.xmolout_handler import XmoloutHandler
+
+
+# ==========================================================
+# =============== Single-frame RDF backends ================
+# ==========================================================
+
+def _rdf_freud_frame(
+    handler: XmoloutHandler,
+    frame_index: int,
+    *,
+    types_a: Optional[Iterable[str]] = None,
+    types_b: Optional[Iterable[str]] = None,
+    r_max: Optional[float] = None,
+    bins: int = 200,
+    drop_first_bin: bool = True,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Compute the RDF for a single frame using the FREUD backend.
+
+    Works on
+    --------
+    XmoloutHandler — ``xmolout``
+
+    Parameters
+    ----------
+    handler : XmoloutHandler
+        Parsed trajectory handler.
+    frame_index : int
+        Frame index to analyze.
+    types_a, types_b : iterable of str, optional
+        Atom types defining A–B RDF. If None, all atoms are used.
+    r_max : float, optional
+        Maximum radius. Defaults to half the smallest box length.
+    bins : int, default=200
+        Number of RDF bins.
+    drop_first_bin : bool, default=True
+        Drop the first bin to avoid self-counting artifacts.
+
+    Returns
+    -------
+    r : np.ndarray
+        Radial distances.
+    g : np.ndarray
+        RDF values g(r).
+
+    Examples
+    --------
+    >>> r, g = _rdf_freud_frame(xh, frame_index=0, types_a=["Al"], types_b=["N"])
+    """
+    df = handler.dataframe()
+    fr = handler.frame(int(frame_index))
+
+    coords = fr["coords"]
+    atom_types = np.asarray(fr["atom_types"], dtype=str)
+
+    # Box from sim DF (supports non-orthogonal if available)
+    row = df.iloc[int(frame_index)]
+    a, b, c = float(row["a"]), float(row["b"]), float(row["c"])
+    alpha = float(row["alpha"]) if "alpha" in df.columns else 90.0
+    beta  = float(row["beta"])  if "beta"  in df.columns else 90.0
+    gamma = float(row["gamma"]) if "gamma" in df.columns else 90.0
+
+    box = freud.box.Box.from_box_lengths_and_angles(
+        a, b, c, np.radians(alpha), np.radians(beta), np.radians(gamma)
+    )
+
+    # A/B masks
+    A_mask = np.ones(len(atom_types), dtype=bool) if types_a is None else np.isin(atom_types, list(types_a))
+    B_mask = np.ones(len(atom_types), dtype=bool) if types_b is None else np.isin(atom_types, list(types_b))
+    A = coords[A_mask]
+    B = coords[B_mask]
+
+    if r_max is None:
+        r_max = 0.5 * float(min(a, b, c))
+
+    rdf = freud.density.RDF(bins=int(bins), r_max=float(r_max))
+    rdf.compute((box, B), query_points=A)
+
+    r = rdf.bin_centers.copy()
+    g = rdf.rdf.copy()
+    if drop_first_bin and len(r) > 1:
+        r, g = r[1:], g[1:]
+    return r, g
+
+
+def _write_xyz_temp(coords: np.ndarray, types: Sequence[str]) -> str:
+    import tempfile, os
+    def _norm(sym: str) -> str:
+        s = str(sym).strip()
+        return s[:1].upper() + s[1:].lower()
+    n = coords.shape[0]
+    fd, path = tempfile.mkstemp(suffix=".xyz", prefix="reaxkit_rdf_")
+    os.close(fd)
+    with open(path, "w") as f:
+        f.write(f"{n}\n")
+        f.write("generated by reaxkit RDF (single-frame)\n")
+        for t, (x, y, z) in zip(types, coords):
+            f.write(f"{_norm(t)} {x:.9f} {y:.9f} {z:.9f}\n")
+    return path
+
+
+def _rdf_ovito_total_frame(
+    handler: XmoloutHandler,
+    frame_index: int,
+    *,
+    r_max: float = 4.0,
+    bins: int = 200,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Compute the total RDF for a single frame using the OVITO backend.
+
+    Works on
+    --------
+    XmoloutHandler — ``xmolout``
+
+    Parameters
+    ----------
+    handler : XmoloutHandler
+        Parsed trajectory handler.
+    frame_index : int
+        Frame index to analyze.
+    r_max : float, default=4.0
+        RDF cutoff radius.
+    bins : int, default=200
+        Number of RDF bins.
+
+    Returns
+    -------
+    r : np.ndarray
+        Radial distances.
+    g : np.ndarray
+        RDF values g(r).
+
+    Examples
+    --------
+    >>> r, g = _rdf_ovito_total_frame(xh, frame_index=10)
+    """
+    fr = handler.frame(int(frame_index))
+    coords = fr["coords"]
+    types = fr["atom_types"]
+
+    import os
+    xyz = _write_xyz_temp(coords, types)
+
+    try:
+        from ovito.io import import_file
+        from ovito.modifiers import CoordinationAnalysisModifier
+    except Exception as e:
+        raise ImportError(
+            "OVITO is required for RDF analysis. "
+            "Install reaxkit[viz] and run in an environment with EGL/Qt."
+        ) from e
+
+    try:
+        pipe = import_file(xyz)
+        pipe.modifiers.append(CoordinationAnalysisModifier(cutoff=float(r_max), number_of_bins=int(bins), partial=False))
+        data = pipe.compute()
+        table = data.tables["coordination-rdf"]
+        arr = table.xy()
+        r = np.asarray(arr[:, 0])
+        g = np.asarray(arr[:, 1])
+        return r, g
+    finally:
+        try: os.remove(xyz)
+        except OSError: pass
+
+
+def _rdf_ovito_partial_frame(
+    handler: XmoloutHandler,
+    frame_index: int,
+    *,
+    r_max: float = 4.0,
+    bins: int = 200,
+    type_a: str,
+    type_b: str,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Compute a partial A–B RDF for a single frame using the OVITO backend.
+
+    Works on
+    --------
+    XmoloutHandler — ``xmolout``
+
+    Parameters
+    ----------
+    handler : XmoloutHandler
+        Parsed trajectory handler.
+    frame_index : int
+        Frame index to analyze.
+    r_max : float, default=4.0
+        RDF cutoff radius.
+    bins : int, default=200
+        Number of RDF bins.
+    type_a, type_b : str
+        Atom types defining the partial RDF (e.g. ``Al``–``N``).
+
+    Returns
+    -------
+    r : np.ndarray
+        Radial distances.
+    g : np.ndarray
+        RDF values g(r).
+
+    Examples
+    --------
+    >>> r, g = _rdf_ovito_partial_frame(xh, 0, type_a="Al", type_b="N")
+    """
+    fr = handler.frame(int(frame_index))
+    coords = fr["coords"]
+    types = fr["atom_types"]
+
+    import os
+    xyz = _write_xyz_temp(coords, types)
+
+    try:
+        from ovito.io import import_file
+        from ovito.modifiers import CoordinationAnalysisModifier
+    except Exception as e:
+        raise ImportError(
+            "OVITO is required for RDF analysis. "
+            "Install reaxkit[viz] and run in an environment with EGL/Qt."
+        ) from e
+
+    try:
+        pipe = import_file(xyz)
+        pipe.modifiers.append(CoordinationAnalysisModifier(cutoff=float(r_max), number_of_bins=int(bins), partial=True))
+        data = pipe.compute()
+        base = "coordination-rdf"
+        if base not in data.tables:
+            raise KeyError(f"'{base}' table not found. Available: {list(data.tables.keys())}")
+        table = data.tables[base]
+
+        # locate the pair component
+        try:
+            names = list(table.y.component_names)
+        except Exception:
+            names = list(getattr(table.y, "components", []))
+        want = f"{type_a}-{type_b}"
+        if want not in names and f"{type_b}-{type_a}" in names:
+            want = f"{type_b}-{type_a}"
+        idx = names.index(want)
+
+        arr = table.xy()
+        r = np.asarray(arr[:, 0])
+        g = np.asarray(arr[:, 1 + idx])
+        return r, g
+    finally:
+        try: os.remove(xyz)
+        except OSError: pass
+
+
+# ==========================================================
+# ================= RDF utilities (no smoothing) ===========
+# ==========================================================
+
+def _dominant_peak(r: np.ndarray, g: np.ndarray) -> tuple[float, float]:
+    """Return the position and height of the dominant (global) RDF peak.
+
+    Parameters
+    ----------
+    r : np.ndarray
+        Radial distances.
+    g : np.ndarray
+        RDF values.
+
+    Returns
+    -------
+    r_peak : float
+        Radius at the global maximum of g(r).
+    g_peak : float
+        Height of the global maximum.
+
+    Examples
+    --------
+    >>> r_peak, g_peak = _dominant_peak(r, g)
+    """
+    if len(g) == 0:
+        return float("nan"), float("nan")
+    k = int(np.argmax(g))
+    return float(r[k]), float(g[k])
+
+
+def _first_local_max(r: np.ndarray, g: np.ndarray) -> tuple[float, float]:
+    """Return the first local maximum of an RDF curve.
+
+    Parameters
+    ----------
+    r : np.ndarray
+        Radial distances.
+    g : np.ndarray
+        RDF values.
+
+    Returns
+    -------
+    r_peak : float
+        Radius of the first local maximum.
+    g_peak : float
+        Height of the first local maximum.
+
+    Notes
+    -----
+    Falls back to the dominant peak if no local maximum is found.
+
+    Examples
+    --------
+    >>> r_first, g_first = _first_local_max(r, g)
+    """
+    n = len(g)
+    for i in range(1, n - 1):
+        if g[i] > g[i - 1] and g[i] > g[i + 1]:
+            return float(r[i]), float(g[i])
+    return _dominant_peak(r, g)
+
+
+# ==========================================================
+# ============== Multi-frame RDF + properties ==============
+# ==========================================================
+
+def rdf_using_freud(
+    handler: XmoloutHandler,
+    *,
+    frames: Optional[Iterable[int]] = None,
+    types_a: Optional[Iterable[str]] = None,
+    types_b: Optional[Iterable[str]] = None,
+    r_max: Optional[float] = None,
+    bins: int = 200,
+    average: bool = True,
+    return_stack: bool = False,
+) -> Tuple[np.ndarray, Union[np.ndarray, List[np.ndarray]]]:
+    """Compute RDFs across multiple frames using the FREUD backend.
+
+    Works on
+    --------
+    XmoloutHandler — ``xmolout``
+
+    Parameters
+    ----------
+    frames : iterable of int, optional
+        Frame indices to include.
+    types_a, types_b : iterable of str, optional
+        Atom types defining A–B RDF.
+    average : bool, default=True
+        If True, return the average RDF across frames.
+    return_stack : bool, default=False
+        If True and ``average=False``, return all per-frame RDFs.
+
+    Returns
+    -------
+    r : np.ndarray
+        Radial distances.
+    g : np.ndarray or list[np.ndarray]
+        Averaged RDF, last RDF, or stack of RDFs.
+
+    Examples
+    --------
+    >>> r, g = rdf_using_freud(xh, frames=range(100), average=True)
+    """
+    df = handler.dataframe()
+    if frames is None:
+        frames = range(len(df))
+    frames = list(frames)
+
+    r_ref: np.ndarray | None = None
+    stack: List[np.ndarray] = []
+
+    for i in frames:
+        r, g = _rdf_freud_frame(
+            handler, frame_index=int(i),
+            types_a=types_a, types_b=types_b, r_max=r_max, bins=bins, drop_first_bin=True
+        )
+        if r_ref is None:
+            r_ref = r
+        else:
+            if len(r) != len(r_ref) or np.max(np.abs(r - r_ref)) > 1e-10:
+                raise ValueError("R grids differ between frames; fix bins/r_max.")
+        stack.append(g)
+
+    if r_ref is None:
+        raise ValueError("No frames selected.")
+
+    if average:
+        return r_ref, np.mean(np.vstack(stack), axis=0)
+    if return_stack:
+        return r_ref, stack
+    return r_ref, stack[-1]
+
+
+def rdf_using_ovito(
+    handler: XmoloutHandler,
+    *,
+    frames: Optional[Iterable[int]] = None,
+    r_max: float = 4.0,
+    bins: int = 200,
+    types_a: Optional[Iterable[str]] = None,  # if both provided => partial
+    types_b: Optional[Iterable[str]] = None,
+    average: bool = True,
+    return_stack: bool = False,
+) -> Tuple[np.ndarray, Union[np.ndarray, List[np.ndarray]]]:
+    """Compute RDFs across multiple frames using the OVITO backend.
+
+    Works on
+    --------
+    XmoloutHandler — ``xmolout``
+
+    Parameters
+    ----------
+    frames : iterable of int, optional
+        Frame indices to include.
+    types_a, types_b : iterable of str, optional
+        If provided, compute partial RDF.
+    average : bool, default=True
+        Return the average RDF across frames.
+
+    Returns
+    -------
+    r : np.ndarray
+        Radial distances.
+    g : np.ndarray or list[np.ndarray]
+        RDF data.
+
+    Examples
+    --------
+    >>> r, g = rdf_using_ovito(xh, types_a=["Al"], types_b=["N"])
+    """
+    df = handler.dataframe()
+    if frames is None:
+        frames = range(len(df))
+    frames = list(frames)
+
+    do_partial = (types_a is not None) and (types_b is not None)
+    if do_partial:
+        ta = next(iter(types_a))
+        tb = next(iter(types_b))
+
+    r_ref: np.ndarray | None = None
+    stack: List[np.ndarray] = []
+
+    for i in frames:
+        if do_partial:
+            r, g = _rdf_ovito_partial_frame(handler, int(i), r_max=r_max, bins=bins, type_a=str(ta), type_b=str(tb))
+        else:
+            r, g = _rdf_ovito_total_frame(handler, int(i), r_max=r_max, bins=bins)
+
+        if r_ref is None:
+            r_ref = r
+        else:
+            if len(r) != len(r_ref) or np.max(np.abs(r - r_ref)) > 1e-10:
+                raise ValueError("R grids differ between frames; fix bins/r_max.")
+        stack.append(g)
+
+    if r_ref is None:
+        raise ValueError("No frames selected.")
+
+    if average:
+        return r_ref, np.mean(np.vstack(stack), axis=0)
+    if return_stack:
+        return r_ref, stack
+    return r_ref, stack[-1]
+
+
+def rdf_property_over_frames(
+    handler: XmoloutHandler,
+    *,
+    backend: str = "ovito",
+    frames: Optional[Iterable[int]] = None,
+    property: str = "first_peak",  # 'first_peak' | 'dominant_peak' | 'area' | 'excess_area'
+    r_max: Optional[float] = None,  # for OVITO total/partial also accepts float
+    bins: int = 200,
+    types_a: Optional[Iterable[str]] = None,
+    types_b: Optional[Iterable[str]] = None,
+) -> pd.DataFrame:
+    """Compute a single RDF-derived property for each simulation frame.
+
+    Works on
+    --------
+    XmoloutHandler — ``xmolout``
+
+    Parameters
+    ----------
+    backend : {"ovito", "freud"}, default="ovito"
+        RDF backend to use.
+    property : {"first_peak", "dominant_peak", "area", "excess_area"}
+        RDF-based property to compute.
+    frames : iterable of int, optional
+        Frame indices to include.
+    types_a, types_b : iterable of str, optional
+        Atom types defining partial RDF.
+
+    Returns
+    -------
+    pandas.DataFrame
+        One row per frame with frame index, iteration, and property values.
+
+    Examples
+    --------
+    >>> df = rdf_property_over_frames(xh, property="first_peak")
+    """
+    df_sim = handler.dataframe()
+    if frames is None:
+        frames = range(len(df_sim))
+    frames = list(frames)
+
+    prop = property.strip().lower()
+    allowed = {"first_peak", "dominant_peak", "area", "excess_area"}
+    if prop not in allowed:
+        raise ValueError(f"property must be one of {allowed}")
+
+    rows: List[Dict[str, float]] = []
+
+    for i in frames:
+        if backend.lower() == "freud":
+            r, g = _rdf_freud_frame(handler, int(i), types_a=types_a, types_b=types_b, r_max=r_max, bins=bins)
+        elif backend.lower() == "ovito":
+            if (types_a is not None) and (types_b is not None):
+                ta = next(iter(types_a)); tb = next(iter(types_b))
+                r, g = _rdf_ovito_partial_frame(handler, int(i), r_max=float(r_max or 4.0), bins=bins, type_a=str(ta), type_b=str(tb))
+            else:
+                r, g = _rdf_ovito_total_frame(handler, int(i), r_max=float(r_max or 4.0), bins=bins)
+        else:
+            raise ValueError("backend must be 'freud' or 'ovito'")
+
+        if prop == "first_peak":
+            rp, gp = _first_local_max(r, g)
+            out = {"r_first_peak": rp, "g_first_peak": gp}
+        elif prop == "dominant_peak":
+            rp, gp = _dominant_peak(r, g)
+            out = {"r_peak": rp, "g_peak": gp}
+        elif prop == "area":
+            out = {"area": float(np.trapezoid(g, r)) if len(r) else np.nan}
+        else:  # "excess_area"
+            out = {"excess_area": float(np.trapezoid(g - 1.0, r)) if len(r) else np.nan}
+
+        rows.append({
+            "frame_index": int(i),
+            "iter": int(df_sim.iloc[i]["iter"]) if "iter" in df_sim.columns else int(i),
+            **out,
+        })
+
+    return pd.DataFrame(rows).sort_values("frame_index").reset_index(drop=True)