reaxkit 1.0.0__py3-none-any.whl

reaxkit/workflows/composed/coordination_workflow.py

@@ -0,0 +1,393 @@
+"""
+Coordination workflow for ReaxKit.
+
+This workflow analyzes atomic coordination using ReaxFF bond-order data
+from `fort.7` and structural information from `xmolout`.
+
+It provides two main capabilities:
+
+- Analyze coordination status per atom per frame (under-, correctly-, or over-coordinated)
+  based on atom valence and total bond order.
+- Relabel atom types according to coordination status and write a new `xmolout` file
+  for visualization or downstream analysis.
+
+The workflow supports flexible frame selection, user-defined or ffield-derived valences,
+and multiple labeling modes.
+"""
+
+
+from __future__ import annotations
+import argparse
+from typing import Dict, Any, Optional, Literal, List
+import pandas as pd
+
+from reaxkit.io.handlers.xmolout_handler import XmoloutHandler
+from reaxkit.io.generators.xmolout_generator import write_xmolout_from_frames
+from reaxkit.analysis.per_file.fort7_analyzer import per_atom_coordination_status_over_frames
+from reaxkit.utils.path import resolve_output_path
+from reaxkit.io.handlers.ffield_handler import FFieldHandler
+
+# -----------------------------
+# Helpers
+# -----------------------------
+def _normalize_frames(xh: XmoloutHandler, frames: Optional[str]) -> list[int]:
+    """
+    Normalize a frame-selection string into a sorted list of frame indices.
+
+    Supported formats:
+      - 'start:end:step'  (any of the 3 can be omitted, e.g. '::5', '10:', ':100')
+      - 'i,j,k' or 'i j k' (explicit indices)
+    If frames is None, all frames [0, n_frames) are used.
+    """
+    n = xh.n_frames()
+
+    if not frames:
+        return list(range(n))
+
+    s = frames.strip()
+
+    # Slice-like 'start:end:step'
+    if ":" in s:
+        parts = s.split(":")
+        if len(parts) > 3:
+            raise ValueError(f"Invalid --frames slice specification: {frames!r}")
+        start_s = parts[0].strip() if len(parts) >= 1 else ""
+        end_s   = parts[1].strip() if len(parts) >= 2 else ""
+        step_s  = parts[2].strip() if len(parts) == 3 else ""
+
+        start = int(start_s) if start_s else 0
+        end   = int(end_s)   if end_s   else n
+        step  = int(step_s)  if step_s  else 1
+        if step == 0:
+            raise ValueError("Step in --frames slice cannot be 0.")
+
+        start = max(0, start)
+        end   = min(n, end)
+        return list(range(start, end, step))
+
+    # Comma/space-separated explicit indices
+    idx: List[int] = []
+    for tok in s.replace(",", " ").split():
+        i = int(tok)
+        if 0 <= i < n:
+            idx.append(i)
+    return sorted(set(idx))
+
+
+def _valences_from_ffield(ffield_path: str) -> Dict[str, float]:
+    """
+    Read atom valencies from ffield atom section.
+
+    Returns a dict that includes BOTH:
+      - symbol -> valency  (e.g., "O": 2)
+      - atom_index -> valency as string key (e.g., "2": 2)
+    This makes it robust whether xmolout uses symbols or numeric atom types.
+    """
+    fh = FFieldHandler(ffield_path)
+    atom_df = fh.section_df(FFieldHandler.SECTION_ATOM)
+
+    if "valency" not in atom_df.columns:
+        raise SystemExit("❌ ffield atom section missing 'valency' column.")
+
+    out: Dict[str, float] = {}
+    used_pairs = []  # for printing
+
+    for atom_idx, row in atom_df.iterrows():
+        # atom_idx is the ffield atom index (usually 1-based)
+        sym = row.get("symbol")
+        val = row.get("valency")
+
+        if val is None or (isinstance(val, float) and pd.isna(val)):
+            continue
+
+        try:
+            v = float(val)
+        except Exception:
+            continue
+
+        # key by atom index (string) too, for numeric xmolout types
+        out[str(int(atom_idx))] = v
+
+        # key by symbol if present
+        if sym is not None and not (isinstance(sym, float) and pd.isna(sym)):
+            s = str(sym).strip()
+            if s:
+                out[s] = v
+                used_pairs.append(f"{s}={v:g}")
+            else:
+                used_pairs.append(f"{int(atom_idx)}={v:g}")
+        else:
+            used_pairs.append(f"{int(atom_idx)}={v:g}")
+
+    if not out:
+        raise SystemExit("❌ No usable valencies found in ffield atom section.")
+
+    # Print one clean line showing what you used (symbol-first when available)
+    # (You can keep it in this helper so both tasks benefit.)
+    print("[Valences] " + ", ".join(used_pairs))
+    return out
+
+
+def _parse_kv_map(s: Optional[str], value_cast=float) -> Dict[str, float]:
+    if not s:
+        return {}
+    out: Dict[str, float] = {}
+    for item in s.split(","):
+        if not item.strip():
+            continue
+        if "=" not in item:
+            raise ValueError(f"Invalid mapping entry {item!r}; use key=value,comma-separated.")
+        k, v = item.split("=", 1)
+        out[k.strip()] = value_cast(v.strip())
+    return out
+
+
+def _parse_status_labels(s: Optional[str]) -> Dict[int, str]:
+    if not s:
+        return {-1: "U", 0: "C", 1: "O"}
+    out: Dict[int, str] = {}
+    for item in s.split(","):
+        if not item.strip():
+            continue
+        if "=" not in item:
+            raise ValueError(f"Invalid label entry {item!r}; use -1=U,0=C,1=O")
+        k, v = item.split("=", 1)
+        ki = int(k.strip())
+        if ki not in (-1, 0, 1):
+            raise ValueError("Status keys must be -1, 0, or 1")
+        out[ki] = v.strip()
+    for k in (-1, 0, 1):
+        out.setdefault(k, {-1: "U", 0: "C", 1: "O"}[k])
+    return out
+
+
+def _frame_record_from_handler(xh: XmoloutHandler, i: int) -> Dict[str, Any]:
+    fr = xh.frame(i)
+    df = xh.dataframe()
+    row = df.iloc[i] if i < len(df) else pd.Series()
+    return {
+        "iter": int(fr.get("iter", int(row["iter"]) if "iter" in row else i)),
+        "coords": fr["coords"],
+        "atom_types": fr["atom_types"],
+        "energy": float(row["energy"]) if "energy" in row else 0.0,
+        "a": float(row["a"]) if "a" in row else 1.0,
+        "b": float(row["b"]) if "b" in row else 1.0,
+        "c": float(row["c"]) if "c" in row else 1.0,
+        "alpha": float(row["alpha"]) if "alpha" in row else 90.0,
+        "beta": float(row["beta"]) if "beta" in row else 90.0,
+        "gamma": float(row["gamma"]) if "gamma" in row else 90.0,
+    }
+
+
+# -----------------------------
+# Action implementations
+# -----------------------------
+def _task_analyze(args: argparse.Namespace) -> int:
+    xh = XmoloutHandler(args.xmolout)
+    from reaxkit.io.handlers.fort7_handler import Fort7Handler
+    f7 = Fort7Handler(args.fort7)
+
+    valences = _parse_kv_map(args.valences, float)
+    if valences:
+        # User explicitly provided valences; print what will be used
+        print("[Valences] " + ", ".join(f"{k}={v:g}" for k, v in valences.items()))
+    else:
+        valences = _valences_from_ffield(args.ffield)
+    if not valences:
+        raise SystemExit("❌ Provide --valences like 'Mg=2,O=2'.")
+
+    frames = _normalize_frames(xh, args.frames)
+
+    df = per_atom_coordination_status_over_frames(
+        f7, xh,
+        valences=valences,
+        threshold=args.threshold,
+        frames=frames,
+        iterations=None,
+        require_all_valences=not args.allow_missing_valences,
+    )
+
+    workflow_name = args.kind
+    if args.export:
+        export_path = resolve_output_path(args.export, workflow_name)
+        export_path.parent.mkdir(parents=True, exist_ok=True)
+        df.to_csv(export_path, index=False)
+        print(f"[Done] Exported coordination table to {export_path}")
+    else:
+        summary = df.groupby(["frame_index", "status"], dropna=False) \
+            .size().unstack(fill_value=0)
+        print(summary)
+
+    return 0
+
+
+def _task_relabel(args: argparse.Namespace) -> int:
+    xh = XmoloutHandler(args.xmolout)
+    from reaxkit.io.handlers.fort7_handler import Fort7Handler
+    f7 = Fort7Handler(args.fort7)
+
+    valences = _parse_kv_map(args.valences, float)
+    if valences:
+        # User explicitly provided valences; print what will be used
+        print("[Valences] " + ", ".join(f"{k}={v:g}" for k, v in valences.items()))
+    else:
+        valences = _valences_from_ffield(args.ffield)
+    if not valences:
+        raise SystemExit("❌ Provide --valences like 'Mg=2,O=2'.")
+
+    labels = _parse_status_labels(args.labels)
+    mode: Literal["global", "by_type"] = args.mode
+    frames = _normalize_frames(xh, args.frames)
+
+    status_df = per_atom_coordination_status_over_frames(
+        f7, xh,
+        valences=valences,
+        threshold=args.threshold,
+        frames=frames,
+        iterations=None,
+        require_all_valences=not args.allow_missing_valences,
+    )
+
+    if status_df.empty:
+        blocks = (_frame_record_from_handler(xh, i) for i in frames)
+        write_xmolout_from_frames(
+            blocks, args.output,
+            simulation_name=args.simulation or getattr(xh, "simulation_name", None) or "MD",
+            precision=args.precision,
+        )
+        print(f"⚠️ No frames selected or empty status; wrote pass-through xmolout to {args.output}")
+        return 0
+
+    sim_df = xh.dataframe()
+    frame_blocks: List[Dict[str, Any]] = []
+    for fi, g in status_df.groupby("frame_index", sort=True):
+        fi = int(fi)
+        fr = xh.frame(fi)
+        coords = fr["coords"]
+        orig_types = fr["atom_types"]
+        n = len(orig_types)
+
+        g_sorted = g.sort_values("atom_id")  # 1-based ids
+        if len(g_sorted) != n:
+            raise SystemExit(f"❌ Frame {fi}: mismatch atoms between xmolout ({n}) and status rows ({len(g_sorted)}).")
+
+        new_types: List[str] = []
+        for k in range(n):
+            t0 = str(orig_types[k])
+            st_val = g_sorted.iloc[k]["status"]
+            if pd.isna(st_val):
+                new_types.append(t0)
+                continue
+            st = int(st_val)
+            tag = labels.get(st, {-1: "U", 0: "C", 1: "O"}[st])
+            if mode == "global":
+                new_types.append(str(tag))
+            else:
+                if st == 0 and args.keep_coord_original:
+                    new_types.append(t0)
+                else:
+                    new_types.append(f"{t0}{tag}")
+
+        row = sim_df.iloc[fi] if fi < len(sim_df) else pd.Series()
+        frame_blocks.append({
+            "iter": int(g_sorted.iloc[0]["iter"]) if "iter" in g_sorted.columns else int(fr["iter"]),
+            "coords": coords,
+            "atom_types": new_types,
+            "E_pot": float(row["E_pot"]) if "E_pot" in row else 0.0,
+            "a": float(row["a"]) if "a" in row else 1.0,
+            "b": float(row["b"]) if "b" in row else 1.0,
+            "c": float(row["c"]) if "c" in row else 1.0,
+            "alpha": float(row["alpha"]) if "alpha" in row else 90.0,
+            "beta":  float(row["beta"])  if "beta"  in row else 90.0,
+            "gamma": float(row["gamma"]) if "gamma" in row else 90.0,
+        })
+
+    workflow_name = args.kind
+    output_path = resolve_output_path(args.output, workflow_name)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    write_xmolout_from_frames(
+        frame_blocks,
+        output_path,
+        simulation_name=args.simulation or getattr(xh, "simulation_name", None) or "MD",
+        precision=args.precision,
+    )
+
+    print(f"[Done] Wrote relabeled xmolout to {output_path}")
+    return 0
+
+
+# -----------------------------
+# Template-style registration
+# -----------------------------
+def _add_common_coord_args(p: argparse.ArgumentParser) -> None:
+    """Common args shared by coord analyze / relabel."""
+    p.add_argument("--xmolout", default="xmolout", help="Path to xmolout.")
+    p.add_argument("--fort7", default="fort.7", help="Path to fort.7 file.")
+    p.add_argument("--valences", default=None,
+                   help="Optional: override type valences, e.g. 'Mg=2,O=2,H=1'. "
+                        "If omitted, reads from ffield atom section ('valency').")
+    p.add_argument("--ffield", default="ffield", help="Path to ffield (used if --valences not given).")
+    p.add_argument("--threshold", type=float, default=0.3, help="Tolerance around valence.")
+    p.add_argument("--frames", default=None,
+                   help="Frame selection, e.g. '0:100:5' or '0,5,10'.")
+    p.add_argument("--allow-missing-valences", action="store_true",
+                   help="Keep atoms with unknown valence (status=NaN).")
+
+
+def register_tasks(subparsers: argparse._SubParsersAction) -> None:
+    """
+    Register coordination-related CLI subcommands.
+
+    This function defines the task-level interface for the
+    `reaxkit file` workflow and attaches its tasks (i.e., subcommands).
+
+    Each task may share common input arguments and defines task-specific options as needed.
+    """
+
+    # -------------------- analyze --------------------
+    p1 = subparsers.add_parser(
+        "analyze",
+        help="Analyze coordination per atom per frame.",
+        description=(
+            "Analyzes coordination per atom per frame.\n"
+            "It determines if an atom is under-coordinated, over-coordinated, or coordinates, based on its valence and "
+            "total bond order.\n"
+            "Examples:\n"
+            "  reaxkit coord analyze --export coord_analysis.csv\n"
+            "  reaxkit coord analyze --valences 'Mg=2,O=2' --frames 0:200:2 --export coord_0_200_2.csv\n"
+        ),
+        formatter_class=argparse.RawTextHelpFormatter,
+    )
+
+    _add_common_coord_args(p1)
+    p1.add_argument("--export", default=None, help="Path to export coordination CSV.")
+    p1.set_defaults(_run=_task_analyze)
+
+
+    # -------------------- relabel --------------------
+    p2 = subparsers.add_parser(
+        "relabel",
+        help="Relabel atom types by coordination and write a new xmolout.",
+        description=(
+            "Relabel atom types by coordination and write a new xmolout.\n"
+            "Examples:\n"
+            "  reaxkit coord relabel --output xmolout_relabeled --mode global --labels=-1=U,0=C,1=O\n"
+            "  reaxkit coord relabel --output xmolout_type --mode by_type --keep-coord-original\n"
+            "  reaxkit coord relabel --valences 'Mg=2,O=2,Zn=2' --frames 0:400:5 "
+            "--output xmolout_relabeled --mode global\n"
+        ),
+        formatter_class=argparse.RawTextHelpFormatter,
+    )
+
+    _add_common_coord_args(p2)
+    p2.add_argument("--output", required=True, help="Output xmolout path.")
+    p2.add_argument("--mode", choices=["global", "by_type"], default="global",
+                    help="Relabeling mode.")
+    p2.add_argument("--labels", default=None,
+                    help="Status→tag map, e.g. '-1=U,0=C,1=O'.")
+    p2.add_argument("--keep-coord-original", action="store_true",
+                    help="In by_type mode, keep original label when status==0.")
+    p2.add_argument("--simulation", default=None, help="Override header simulation name.")
+    p2.add_argument("--precision", type=int, default=6, help="Float precision.")
+    p2.set_defaults(_run=_task_relabel)