cuRDF 0.1.0__tar.gz

curdf-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Joseph Hart
+
+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.

curdf-0.1.0/PKG-INFO

@@ -0,0 +1,20 @@
+Metadata-Version: 2.4
+Name: cuRDF
+Version: 0.1.0
+Summary: GPU-accelerated radial distribution functions using Toolkit-Ops + PyTorch with MDAnalysis and ASE adapters.
+Requires-Python: >=3.10
+License-File: LICENSE
+Requires-Dist: torch
+Requires-Dist: numpy
+Requires-Dist: nvalchemiops
+Provides-Extra: analysis
+Requires-Dist: MDAnalysis; extra == "analysis"
+Requires-Dist: ase; extra == "analysis"
+Requires-Dist: matplotlib; extra == "analysis"
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Provides-Extra: docs
+Requires-Dist: sphinx; extra == "docs"
+Requires-Dist: sphinx-rtd-theme; extra == "docs"
+Requires-Dist: myst-parser; extra == "docs"
+Dynamic: license-file

curdf-0.1.0/README.md

@@ -0,0 +1,83 @@
+# cuRDF
+
+[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.1085332119.svg)](https://doi.org/10.5281/zenodo.1085332119)
+
+
+CUDA-accelerated radial distribution functions using NVIDIA ALCHEMI Toolkit-Ops O(N) neighbor lists and PyTorch. Compatible with ASE and MDAnalysis.
+
+[![PyPI version](https://badge.fury.io/py/curdf.svg)](https://badge.fury.io/py/curdf)
+[![Tests](https://github.com/josephhart/amorphous-carbon/actions/workflows/tests.yml/badge.svg)](https://github.com/josephhart/amorphous-carbon/actions/workflows/tests.yml)
+
+## Install (editable)
+```
+pip install -e .
+```
+Add `[analysis]` extras if you want MDAnalysis/ASE/matplotlib:
+```
+pip install -e .[analysis]
+```
+
+## Library usage
+```python
+import curdf
+import MDAnalysis as mda
+
+u = mda.Universe("top.data", "traj.dcd")
+bins, gr = curdf.rdf_from_mdanalysis(u, selection="name C", r_min=1.0, r_max=8.0, nbins=200)
+```
+
+ASE first (XYZ/extxyz/ASE .traj):
+```python
+from ase.io import read
+from curdf import rdf_from_ase
+
+atoms = read("structure.xyz")
+bins, gr = rdf_from_ase(atoms, selection=None, r_min=1.0, r_max=8.0, nbins=200)  # selection=None -> all atoms
+```
+
+Cross-species (ASE): provide group A/B indices
+```python
+bins, gr = rdf_from_ase(atoms, selection=[0,1,2], selection_b=[3,4,5], r_min=1.0, r_max=8.0, nbins=200, half_fill=False)
+```
+
+MDAnalysis (explicit dependency required; also supports LAMMPS dump):
+```python
+import MDAnalysis as mda
+from curdf import rdf_from_mdanalysis
+
+u = mda.Universe("top.data", "traj.dcd")
+bins, gr = curdf.rdf_from_mdanalysis(u, selection="name C", r_min=1.0, r_max=8.0, nbins=200)
+```
+
+## CLI
+ASE (XYZ/extxyz/ASE .traj):
+```
+rdf-gpu --format ase --ase-file structure.xyz --selection 0,1,2 --r-max 8 --nbins 200 --device cuda
+```
+
+Cross-species via CLI (ASE indices or MDAnalysis selections):
+```
+rdf-gpu --format ase --ase-file structure.xyz --selection-a 0,1,2 --selection-b 3,4,5 --r-max 8 --nbins 200 --device cuda --ordered-pairs
+```
+(`--selection-b` automatically disables half-fill so pairs are ordered.)
+
+LAMMPS dump (lammpstrj) via MDAnalysis:
+```
+rdf-gpu --format lammps-dump --trajectory dump.lammpstrj --selection "all" --r-max 8 --nbins 200 --device cuda
+```
+
+MDAnalysis:
+```
+rdf-gpu --format mdanalysis --topology top.data --trajectory traj.dcd --selection "name C" --r-max 8 --nbins 200 --device cuda --out results/rdf.npz --plot results/rdf.png
+```
+
+`--ordered-pairs` switches to counting ordered pairs (disable half-fill). `--no-wrap` leaves coordinates unwrapped if you already wrapped them upstream.
+
+## Docs / examples / tests
+- Docs in `docs/` (index, quickstart, api).
+- Examples in `examples/` for basic, ASE, and MDAnalysis workflows.
+- Tests in `tests/` (run with `pytest` or `pip install -e .[dev]` first).
+- Build Sphinx docs with `pip install -e .[docs]` then `sphinx-build -b html docs/source docs/build/html` (footer: "Built with Sphinx using a theme provided by Read the Docs.").
+
+## Citation
+See `CITATION.cff` for how to cite cuRDF in your work.

curdf-0.1.0/cuRDF.egg-info/PKG-INFO

@@ -0,0 +1,20 @@
+Metadata-Version: 2.4
+Name: cuRDF
+Version: 0.1.0
+Summary: GPU-accelerated radial distribution functions using Toolkit-Ops + PyTorch with MDAnalysis and ASE adapters.
+Requires-Python: >=3.10
+License-File: LICENSE
+Requires-Dist: torch
+Requires-Dist: numpy
+Requires-Dist: nvalchemiops
+Provides-Extra: analysis
+Requires-Dist: MDAnalysis; extra == "analysis"
+Requires-Dist: ase; extra == "analysis"
+Requires-Dist: matplotlib; extra == "analysis"
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Provides-Extra: docs
+Requires-Dist: sphinx; extra == "docs"
+Requires-Dist: sphinx-rtd-theme; extra == "docs"
+Requires-Dist: myst-parser; extra == "docs"
+Dynamic: license-file

curdf-0.1.0/cuRDF.egg-info/SOURCES.txt

@@ -0,0 +1,16 @@
+LICENSE
+README.md
+pyproject.toml
+cuRDF.egg-info/PKG-INFO
+cuRDF.egg-info/SOURCES.txt
+cuRDF.egg-info/dependency_links.txt
+cuRDF.egg-info/entry_points.txt
+cuRDF.egg-info/requires.txt
+cuRDF.egg-info/top_level.txt
+curdf/__init__.py
+curdf/adapters.py
+curdf/cell.py
+curdf/cli.py
+curdf/neighbor.py
+curdf/rdf.py
+tests/test_rdf.py

curdf-0.1.0/cuRDF.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

curdf-0.1.0/cuRDF.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+rdf-gpu = curdf.cli:main

curdf-0.1.0/cuRDF.egg-info/requires.txt

@@ -0,0 +1,16 @@
+torch
+numpy
+nvalchemiops
+
+[analysis]
+MDAnalysis
+ase
+matplotlib
+
+[dev]
+pytest
+
+[docs]
+sphinx
+sphinx-rtd-theme
+myst-parser

curdf-0.1.0/cuRDF.egg-info/top_level.txt

@@ -0,0 +1 @@
+curdf

curdf-0.1.0/curdf/__init__.py

@@ -0,0 +1,12 @@
+"""
+cuRDF: GPU-accelerated radial distribution functions with MDAnalysis/ASE adapters.
+"""
+
+from .rdf import compute_rdf
+from .adapters import rdf_from_mdanalysis, rdf_from_ase
+
+__all__ = [
+    "compute_rdf",
+    "rdf_from_mdanalysis",
+    "rdf_from_ase",
+]

curdf-0.1.0/curdf/adapters.py

@@ -0,0 +1,180 @@
+from collections.abc import Iterable
+from typing import Sequence
+
+import numpy as np
+
+try:
+    import MDAnalysis as mda
+    from MDAnalysis.lib.mdamath import triclinic_vectors
+    from MDAnalysis.transformations import wrap as mda_wrap
+except ImportError:
+    mda = None
+    triclinic_vectors = None
+    mda_wrap = None
+
+try:
+    from ase import Atoms
+except ImportError:
+    Atoms = None
+
+from .rdf import accumulate_rdf
+
+
+def _mdanalysis_cell_matrix(dimensions):
+    """
+    MDAnalysis gives [a, b, c, alpha, beta, gamma]; convert to 3x3.
+    """
+    if triclinic_vectors is None:
+        raise ImportError("MDAnalysis not available")
+    return np.array(triclinic_vectors(dimensions), dtype=np.float32)
+
+
+def rdf_from_mdanalysis(
+    universe,
+    selection: str = "all",
+    selection_b: str | None = None,
+    r_min: float = 1.0,
+    r_max: float = 6.0,
+    nbins: int = 100,
+    device="cuda",
+    torch_dtype=None,
+    half_fill: bool = True,
+    max_neighbors: int = 2048,
+    wrap_positions: bool = True,
+):
+    """
+    Compute g(r) from an MDAnalysis Universe across all trajectory frames.
+    selection_b: optional second selection for cross-species RDF (A in selection, B in selection_b).
+    """
+    if mda is None:
+        raise ImportError("MDAnalysis must be installed for rdf_from_mdanalysis")
+    if torch_dtype is None:
+        import torch
+        torch_dtype = torch.float32
+
+    ag_a = universe.select_atoms(selection)
+    ag_b = universe.select_atoms(selection_b) if selection_b is not None else ag_a
+    if wrap_positions and mda_wrap is not None:
+        ag_wrap = ag_a if selection_b is None else (ag_a | ag_b)
+        universe.trajectory.add_transformations(mda_wrap(ag_wrap, compound="atoms"))
+
+    def frames():
+        for ts in universe.trajectory:
+            cell = _mdanalysis_cell_matrix(ts.dimensions)
+            if selection_b is None:
+                yield {
+                    "positions": ag_a.positions.astype(np.float32, copy=False),
+                    "cell": cell,
+                    "pbc": (True, True, True),
+                }
+            else:
+                pos_a = ag_a.positions.astype(np.float32, copy=False)
+                pos_b = ag_b.positions.astype(np.float32, copy=False)
+                pos = np.concatenate([pos_a, pos_b], axis=0)
+                group_a_mask = np.zeros(len(pos), dtype=bool)
+                group_b_mask = np.zeros(len(pos), dtype=bool)
+                group_a_mask[: len(pos_a)] = True
+                group_b_mask[len(pos_a) :] = True
+                yield {
+                    "positions": pos,
+                    "cell": cell,
+                    "pbc": (True, True, True),
+                    "group_a_mask": group_a_mask,
+                    "group_b_mask": group_b_mask,
+                }
+
+    if selection_b is not None and half_fill:
+        half_fill = False  # cross-species -> ordered pairs
+
+    return accumulate_rdf(
+        frames(),
+        r_min=r_min,
+        r_max=r_max,
+        nbins=nbins,
+        device=device,
+        torch_dtype=torch_dtype,
+        half_fill=half_fill,
+        max_neighbors=max_neighbors,
+    )
+
+
+def _extract_selection_indices(selection: Sequence[int] | None, n_atoms: int):
+    if selection is None:
+        return np.arange(n_atoms)
+    idx = np.asarray(selection, dtype=int)
+    if idx.ndim != 1:
+        raise ValueError("selection indices must be 1D")
+    if idx.min(initial=0) < 0 or idx.max(initial=0) >= n_atoms:
+        raise ValueError("selection indices out of bounds")
+    return idx
+
+
+def rdf_from_ase(
+    atoms_or_trajectory,
+    selection: Sequence[int] | None = None,
+    selection_b: Sequence[int] | None = None,
+    r_min: float = 1.0,
+    r_max: float = 6.0,
+    nbins: int = 100,
+    device="cuda",
+    torch_dtype=None,
+    half_fill: bool = True,
+    max_neighbors: int = 2048,
+    wrap_positions: bool = True,
+):
+    """
+    Compute g(r) from an ASE Atoms or iterable of Atoms (trajectory).
+    selection/selection_b: index lists for group A and group B (cross-species). With only selection provided, computes A–A.
+    """
+    if Atoms is None:
+        raise ImportError("ASE must be installed for rdf_from_ase")
+    if torch_dtype is None:
+        import torch
+        torch_dtype = torch.float32
+
+    def _frames_iter():
+        if hasattr(atoms_or_trajectory, "get_positions"):
+            iterable = (atoms_or_trajectory,)
+        elif isinstance(atoms_or_trajectory, Iterable):
+            iterable = atoms_or_trajectory
+        else:
+            raise TypeError("atoms_or_trajectory must be ASE Atoms or iterable of Atoms")
+
+        for frame in iterable:
+            if not hasattr(frame, "get_positions"):
+                raise TypeError("Each frame must be ASE Atoms")
+            n_atoms = len(frame)
+            idx_a = _extract_selection_indices(selection, n_atoms)
+            idx_b = _extract_selection_indices(selection_b, n_atoms) if selection_b is not None else idx_a
+            pos_all = frame.get_positions(wrap=wrap_positions)
+            pos_a = pos_all[idx_a]
+            pos_b = pos_all[idx_b]
+            pos = np.concatenate([pos_a, pos_b], axis=0)
+            cell = np.array(frame.get_cell().array, dtype=np.float32)
+            pbc = tuple(bool(x) for x in frame.get_pbc())
+            group_a_mask = np.zeros(len(pos), dtype=bool)
+            group_b_mask = np.zeros(len(pos), dtype=bool)
+            group_a_mask[: len(pos_a)] = True
+            group_b_mask[len(pos_a) :] = True
+
+            yield {
+                "positions": pos.astype(np.float32, copy=False),
+                "cell": cell,
+                "pbc": pbc,
+                "group_a_mask": group_a_mask,
+                "group_b_mask": group_b_mask,
+            }
+
+    if selection_b is not None and half_fill:
+        half_fill = False  # cross-species -> ordered pairs
+
+    return accumulate_rdf(
+        _frames_iter(),
+        r_min=r_min,
+        r_max=r_max,
+        nbins=nbins,
+        device=device,
+        torch_dtype=torch_dtype,
+        half_fill=half_fill,
+        max_neighbors=max_neighbors,
+    )

curdf-0.1.0/curdf/cell.py

@@ -0,0 +1,32 @@
+import torch
+from torch import Tensor
+
+
+def cell_tensor(cell: Tensor | list | tuple, device, dtype) -> Tensor:
+    """
+    Convert an input (3,3) cell array into a torch tensor shaped (1,3,3).
+    Accepts triclinic cells. Does not assume orthorhombic.
+    """
+    cell_t = torch.as_tensor(cell, device=device, dtype=dtype)
+    if cell_t.shape != (3, 3):
+        raise ValueError(f"cell must be (3,3); got {tuple(cell_t.shape)}")
+    return cell_t.unsqueeze(0)
+
+
+def pbc_tensor(pbc: Tensor | list | tuple, device) -> Tensor:
+    """
+    Convert PBC flags to shape (1,3) boolean tensor.
+    """
+    pbc_t = torch.as_tensor(pbc, device=device, dtype=torch.bool)
+    if pbc_t.shape != (3,):
+        raise ValueError(f"pbc must be (3,); got {tuple(pbc_t.shape)}")
+    return pbc_t.unsqueeze(0)
+
+
+def cell_volume(cell: Tensor) -> float:
+    """
+    Compute volume from a (1,3,3) cell tensor (triclinic supported).
+    """
+    if cell.shape != (1, 3, 3):
+        raise ValueError(f"cell must be (1,3,3); got {tuple(cell.shape)}")
+    return torch.det(cell[0]).abs().item()

curdf-0.1.0/curdf/cli.py

@@ -0,0 +1,166 @@
+import argparse
+import sys
+from pathlib import Path
+
+import numpy as np
+
+
+def _parse_args():
+    p = argparse.ArgumentParser(description="cuRDF: GPU RDF using Toolkit-Ops + PyTorch")
+    p.add_argument(
+        "--format",
+        choices=["mdanalysis", "ase", "lammps-dump"],
+        required=True,
+        help="Input backend",
+    )
+    p.add_argument("--topology", help="Topology file (MDAnalysis)")
+    p.add_argument("--trajectory", nargs="+", help="Trajectory file(s) (MDAnalysis)")
+    p.add_argument("--ase-file", help="Structure/trajectory file readable by ASE")
+    p.add_argument("--ase-index", default=":", help="ASE index (default all frames)")
+    p.add_argument("--selection", default=None, help="(Deprecated) alias for --selection-a")
+    p.add_argument("--selection-a", default=None, help="MDAnalysis selection or ASE comma-separated indices for group A")
+    p.add_argument("--selection-b", default=None, help="MDAnalysis selection or ASE comma-separated indices for group B")
+    p.add_argument("--r-min", type=float, default=1.0)
+    p.add_argument("--r-max", type=float, default=6.0)
+    p.add_argument("--nbins", type=int, default=100)
+    p.add_argument("--device", default="cuda")
+    p.add_argument("--dtype", choices=["float32", "float64"], default="float32")
+    p.add_argument("--half-fill", action="store_true", default=True, help="Use unique pairs (identical species)")
+    p.add_argument("--ordered-pairs", action="store_true", help="Disable half-fill to count ordered pairs")
+    p.add_argument("--max-neighbors", type=int, default=2048)
+    p.add_argument("--no-wrap", action="store_true", help="Skip wrapping positions into the cell")
+    p.add_argument("--plot", type=Path, help="Optional PNG plot output")
+    p.add_argument("--out", type=Path, default=Path("rdf.npz"), help="NPZ output path")
+    return p.parse_args()
+
+
+def main():
+    args = _parse_args()
+    torch_dtype = {"float32": "float32", "float64": "float64"}[args.dtype]
+    half_fill = False if args.ordered_pairs else args.half_fill
+    # Cross-species (selection-b) should use ordered pairs
+    if args.selection_b and half_fill:
+        half_fill = False
+
+    if args.format == "mdanalysis":
+        if args.topology is None or args.trajectory is None:
+            sys.exit("For mdanalysis format, provide --topology and --trajectory")
+        try:
+            import MDAnalysis as mda
+        except ImportError:
+            sys.exit("MDAnalysis is required for --format mdanalysis")
+
+        u = mda.Universe(args.topology, *args.trajectory)
+        selection_a = args.selection_a or args.selection
+        selection_b = args.selection_b
+        if selection_a is None:
+            selection_a = "all"
+        from .adapters import rdf_from_mdanalysis
+
+        bins, gr = rdf_from_mdanalysis(
+            u,
+            selection=selection_a,
+            selection_b=selection_b,
+            r_min=args.r_min,
+            r_max=args.r_max,
+            nbins=args.nbins,
+            device=args.device,
+            torch_dtype=getattr(__import__("torch"), torch_dtype),
+            half_fill=half_fill,
+            max_neighbors=args.max_neighbors,
+            wrap_positions=not args.no_wrap,
+        )
+    elif args.format == "lammps-dump":
+        if args.trajectory is None:
+            sys.exit("For lammps-dump format, provide --trajectory (LAMMPS dump / lammpstrj)")
+        try:
+            import MDAnalysis as mda
+        except ImportError:
+            sys.exit("MDAnalysis is required for --format lammps-dump")
+
+        try:
+            u = mda.Universe(args.trajectory[0], format="LAMMPSDUMP")
+        except Exception as exc:
+            sys.exit(f"Failed to load LAMMPS dump: {exc}")
+
+        selection_a = args.selection_a or args.selection
+        selection_b = args.selection_b
+        if selection_a is None:
+            selection_a = "all"
+        from .adapters import rdf_from_mdanalysis
+
+        bins, gr = rdf_from_mdanalysis(
+            u,
+            selection=selection_a,
+            selection_b=selection_b,
+            r_min=args.r_min,
+            r_max=args.r_max,
+            nbins=args.nbins,
+            device=args.device,
+            torch_dtype=getattr(__import__("torch"), torch_dtype),
+            half_fill=half_fill,
+            max_neighbors=args.max_neighbors,
+            wrap_positions=not args.no_wrap,
+        )
+    else:
+        if args.ase_file is None:
+            sys.exit("For ase format, provide --ase-file")
+        try:
+            import ase.io
+        except ImportError:
+            sys.exit("ASE is required for --format ase")
+
+        allowed_ext = {".xyz", ".extxyz", ".traj"}
+        if Path(args.ase_file).suffix.lower() not in allowed_ext:
+            sys.exit(f"ASE mode supports {sorted(allowed_ext)}; got {args.ase_file}")
+
+        frames = ase.io.read(args.ase_file, index=args.ase_index)
+        if isinstance(frames, list):
+            atoms_or_traj = frames
+        else:
+            atoms_or_traj = frames
+
+        sel_a = None
+        sel_b = None
+        selection_a = args.selection_a or args.selection
+        if selection_a:
+            sel_a = [int(x) for x in selection_a.split(",") if x.strip()]
+        if args.selection_b:
+            sel_b = [int(x) for x in args.selection_b.split(",") if x.strip()]
+
+        from .adapters import rdf_from_ase
+
+        bins, gr = rdf_from_ase(
+            atoms_or_traj,
+            selection=sel_a,
+            selection_b=sel_b,
+            r_min=args.r_min,
+            r_max=args.r_max,
+            nbins=args.nbins,
+            device=args.device,
+            torch_dtype=getattr(__import__("torch"), torch_dtype),
+            half_fill=half_fill,
+            max_neighbors=args.max_neighbors,
+            wrap_positions=not args.no_wrap,
+        )
+
+    args.out.parent.mkdir(parents=True, exist_ok=True)
+    np.savez(args.out, bins=bins, gr=gr)
+
+    if args.plot:
+        try:
+            import matplotlib
+            matplotlib.use("Agg")
+            import matplotlib.pyplot as plt
+        except ImportError:
+            sys.exit("matplotlib required for plotting")
+        plt.plot(bins, gr)
+        plt.xlabel("r (A)")
+        plt.ylabel("g(r)")
+        plt.hlines(1.0, xmin=args.r_min, xmax=args.r_max, colors="k", linestyles="dashed")
+        plt.savefig(args.plot, dpi=300)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

curdf-0.1.0/curdf/neighbor.py

@@ -0,0 +1,29 @@
+from nvalchemiops.neighborlist import neighbor_list
+import torch
+from torch import Tensor
+
+
+def build_neighbor_list(
+    positions: Tensor,
+    r_max: float,
+    cell: Tensor,
+    pbc: Tensor,
+    half_fill: bool = True,
+    max_neighbors: int = 2048,
+    method: str = "cell_list",
+):
+    """
+    Wrap Toolkit-Ops neighbor list to keep a single import site.
+    Returns nlist (2, num_pairs), shifts (num_pairs,3).
+    """
+    nlist, _, shifts = neighbor_list(
+        positions,
+        float(r_max),
+        cell=cell,
+        pbc=pbc,
+        return_neighbor_list=True,
+        half_fill=half_fill,
+        max_neighbors=max_neighbors,
+        method=method,
+    )
+    return nlist, shifts

curdf-0.1.0/curdf/rdf.py

@@ -0,0 +1,221 @@
+import math
+from typing import Iterable
+
+import torch
+from torch import Tensor
+
+from .cell import cell_tensor, cell_volume, pbc_tensor
+from .neighbor import build_neighbor_list
+
+
+def _update_counts(
+    counts: Tensor,
+    positions: Tensor,
+    cell: Tensor,
+    pbc: Tensor,
+    edges: Tensor,
+    r_min: float,
+    r_max: float,
+    half_fill: bool,
+    max_neighbors: int,
+    group_a_mask: Tensor | None = None,
+    group_b_mask: Tensor | None = None,
+) -> float:
+    """
+    Accumulate pair counts for one frame.
+    Returns normalization factor (n_group_a * rho_group_b) so the caller can normalize after multiple frames.
+    """
+    dr = (r_max - r_min) / (len(edges) - 1)
+
+    nlist, shifts = build_neighbor_list(
+        positions,
+        r_max,
+        cell=cell,
+        pbc=pbc,
+        half_fill=half_fill,
+        max_neighbors=max_neighbors,
+    )
+
+    src = nlist[0].to(torch.int64)
+    tgt = nlist[1].to(torch.int64)
+
+    shift_cart = (shifts.to(positions.dtype) @ cell[0])
+    dr_vec = (positions[tgt] + shift_cart) - positions[src]
+    dist = torch.linalg.norm(dr_vec, dim=1)
+
+    valid = (dist >= r_min) & (dist < r_max)
+    if group_a_mask is not None and group_b_mask is not None:
+        src_mask = group_a_mask[src]
+        tgt_mask = group_b_mask[tgt]
+        valid = valid & src_mask & tgt_mask
+    dist = dist[valid]
+
+    bin_idx = torch.floor((dist - r_min) / dr).to(torch.int64)
+    bin_idx = torch.clamp(bin_idx, 0, counts.numel() - 1)
+    counts.scatter_add_(0, bin_idx, torch.ones_like(bin_idx, dtype=torch.int64))
+
+    volume = cell_volume(cell)
+    if group_a_mask is not None and group_b_mask is not None:
+        n_a = group_a_mask.sum().item()
+        n_b = group_b_mask.sum().item()
+        norm_factor = n_a * (n_b / volume)
+    else:
+        n_atoms = positions.shape[0]
+        norm_factor = n_atoms * (n_atoms / volume)  # n_atoms * rho
+    return norm_factor
+
+
+def _finalize_gr(
+    counts: Tensor,
+    edges: Tensor,
+    total_norm: float,
+    half_fill: bool,
+    cross_mode: bool,
+) -> tuple[Tensor, Tensor]:
+    r1 = edges[:-1]
+    r2 = edges[1:]
+    shell_vol = (4.0 / 3.0) * math.pi * (r2**3 - r1**3)
+    pair_factor = 1.0 if cross_mode else (2.0 if half_fill else 1.0)
+
+    if total_norm == 0:
+        raise ValueError("Total normalization is zero; no frames processed?")
+    g_r = (pair_factor * counts.to(r1.dtype)) / (shell_vol * total_norm)
+    centers = (edges[:-1] + edges[1:]) * 0.5
+    return centers, g_r
+
+
+@torch.no_grad()
+def compute_rdf(
+    positions,
+    cell,
+    pbc=(True, True, True),
+    r_min: float = 1.0,
+    r_max: float = 6.0,
+    nbins: int = 100,
+    device: str | torch.device = "cuda",
+    torch_dtype: torch.dtype = torch.float32,
+    half_fill: bool = True,
+    max_neighbors: int = 2048,
+    group_a_indices=None,
+    group_b_indices=None,
+):
+    """
+    Compute g(r) for a single frame of positions.
+
+    Args:
+        positions: array-like (N,3)
+        cell: (3,3) cell matrix (triclinic allowed)
+        pbc: iterable of 3 booleans
+        r_min/r_max/nbins: histogram parameters
+        half_fill: True for identical species (unique pairs); False for ordered pairs
+        max_neighbors: passed to Toolkit-Ops neighbor list
+        group_a_indices/group_b_indices: optional index lists for cross-species RDF.
+            If provided, counts pairs with src in A and tgt in B. When both are None,
+            uses all atoms (identical-species mode).
+    """
+    device = torch.device(device)
+    pos_t = torch.as_tensor(positions, device=device, dtype=torch_dtype)
+    if pos_t.ndim != 2 or pos_t.shape[1] != 3:
+        raise ValueError(f"positions must be (N,3); got {tuple(pos_t.shape)}")
+
+    cell_t = cell_tensor(cell, device=device, dtype=torch_dtype)
+    pbc_t = pbc_tensor(pbc, device=device)
+
+    edges = torch.linspace(r_min, r_max, nbins + 1, device=device, dtype=torch_dtype)
+    counts = torch.zeros(nbins, device=device, dtype=torch.int64)
+
+    group_a_mask = group_b_mask = None
+    if group_a_indices is not None:
+        group_a_mask = torch.zeros(pos_t.shape[0], device=device, dtype=torch.bool)
+        group_a_mask[torch.as_tensor(group_a_indices, device=device, dtype=torch.long)] = True
+    if group_b_indices is not None:
+        group_b_mask = torch.zeros(pos_t.shape[0], device=device, dtype=torch.bool)
+        group_b_mask[torch.as_tensor(group_b_indices, device=device, dtype=torch.long)] = True
+    elif group_a_mask is not None:
+        group_b_mask = group_a_mask
+
+    total_norm = _update_counts(
+        counts,
+        pos_t,
+        cell=cell_t,
+        pbc=pbc_t,
+        edges=edges,
+        r_min=r_min,
+        r_max=r_max,
+        half_fill=half_fill,
+        max_neighbors=max_neighbors,
+        group_a_mask=group_a_mask,
+        group_b_mask=group_b_mask,
+    )
+
+    cross_mode = group_a_mask is not None and group_b_mask is not None and not torch.equal(
+        group_a_mask, group_b_mask
+    )
+    centers, g_r = _finalize_gr(
+        counts, edges, total_norm, half_fill=half_fill, cross_mode=cross_mode
+    )
+    return centers.cpu().numpy(), g_r.cpu().numpy()
+
+
+@torch.no_grad()
+def accumulate_rdf(
+    frames: Iterable[dict],
+    r_min: float,
+    r_max: float,
+    nbins: int,
+    device: str | torch.device,
+    torch_dtype: torch.dtype,
+    half_fill: bool,
+    max_neighbors: int,
+):
+    """
+    General accumulator for multiple frames.
+    frames: iterable yielding dicts with keys positions, cell, pbc, and optional group_a_mask/group_b_mask
+    """
+    device = torch.device(device)
+    edges = torch.linspace(r_min, r_max, nbins + 1, device=device, dtype=torch_dtype)
+    counts = torch.zeros(nbins, device=device, dtype=torch.int64)
+    total_norm = 0.0
+
+    cross_flag = False
+
+    for frame in frames:
+        pos_t = torch.as_tensor(frame["positions"], device=device, dtype=torch_dtype)
+        cell_t = cell_tensor(frame["cell"], device=device, dtype=torch_dtype)
+        pbc_t = pbc_tensor(frame["pbc"], device=device)
+        group_a_mask = frame.get("group_a_mask")
+        group_b_mask = frame.get("group_b_mask")
+        if group_a_mask is not None:
+            group_a_mask = torch.as_tensor(group_a_mask, device=device, dtype=torch.bool)
+        if group_b_mask is not None:
+            group_b_mask = torch.as_tensor(group_b_mask, device=device, dtype=torch.bool)
+        elif group_a_mask is not None:
+            group_b_mask = group_a_mask
+        if group_a_mask is not None and group_b_mask is not None and not torch.equal(
+            group_a_mask, group_b_mask
+        ):
+            cross_flag = True
+
+        norm = _update_counts(
+            counts,
+            pos_t,
+            cell=cell_t,
+            pbc=pbc_t,
+            edges=edges,
+            r_min=r_min,
+            r_max=r_max,
+            half_fill=half_fill,
+            max_neighbors=max_neighbors,
+            group_a_mask=group_a_mask,
+            group_b_mask=group_b_mask,
+        )
+        total_norm += norm
+
+    centers, g_r = _finalize_gr(
+        counts,
+        edges,
+        total_norm,
+        half_fill=half_fill,
+        cross_mode=cross_flag,
+    )
+    return centers.cpu().numpy(), g_r.cpu().numpy()

curdf-0.1.0/pyproject.toml

@@ -0,0 +1,32 @@
+[build-system]
+requires = ["setuptools>=64", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "cuRDF"
+version = "0.1.0"
+description = "GPU-accelerated radial distribution functions using Toolkit-Ops + PyTorch with MDAnalysis and ASE adapters."
+requires-python = ">=3.10"
+dependencies = [
+  "torch",
+  "numpy",
+  "nvalchemiops",
+]
+
+[project.optional-dependencies]
+analysis = [
+  "MDAnalysis",
+  "ase",
+  "matplotlib",
+]
+dev = [
+  "pytest",
+]
+docs = [
+  "sphinx",
+  "sphinx-rtd-theme",
+  "myst-parser",
+]
+
+[project.scripts]
+rdf-gpu = "curdf.cli:main"

curdf-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

curdf-0.1.0/tests/test_rdf.py

@@ -0,0 +1,78 @@
+import math
+import numpy as np
+import torch
+
+from curdf.rdf import compute_rdf, accumulate_rdf
+
+
+def test_compute_rdf_single_frame_two_atoms():
+    positions = np.array([[0.0, 0.0, 0.0], [1.5, 0.0, 0.0]], dtype=np.float32)
+    cell = np.diag([10.0, 10.0, 10.0])
+    bins, gr = compute_rdf(
+        positions,
+        cell,
+        pbc=(True, True, True),
+        r_min=0.0,
+        r_max=5.0,
+        nbins=5,
+        device="cpu",
+        torch_dtype=torch.float32,
+        half_fill=True,
+    )
+    # Distance 1.5 should fall into bin index 1 (edges 0,1,2,...)
+    assert math.isclose(bins[1], 1.5)
+    # Expected g(r) from analytic normalization
+    shell_vol = (4.0 / 3.0) * math.pi * (2.0**3 - 1.0**3)
+    volume = 10.0**3
+    n_atoms = 2
+    rho = n_atoms / volume
+    expected = (2.0 * 1.0) / (shell_vol * (n_atoms * rho))
+    assert math.isclose(gr[1], expected, rel_tol=1e-5)
+
+
+def test_accumulate_rdf_multiple_frames():
+    cell = np.diag([8.0, 8.0, 8.0])
+    positions1 = np.array([[0.0, 0.0, 0.0], [2.0, 0.0, 0.0]], dtype=np.float32)
+    positions2 = np.array([[0.0, 0.0, 0.0], [3.0, 0.0, 0.0]], dtype=np.float32)
+    frames = [
+        {"positions": positions1, "cell": cell, "pbc": (True, True, True)},
+        {"positions": positions2, "cell": cell, "pbc": (True, True, True)},
+    ]
+    bins, gr = accumulate_rdf(
+        frames,
+        r_min=0.0,
+        r_max=5.0,
+        nbins=5,
+        device="cpu",
+        torch_dtype=torch.float32,
+        half_fill=True,
+        max_neighbors=1024,
+    )
+    # Distances fall into bins near 1 and 3; check both nonzero
+    assert gr[1] > 0
+    assert gr[3] > 0
+
+
+def test_compute_rdf_cross_species():
+    # A at origin, B at 1.5
+    positions = np.array([[0.0, 0.0, 0.0], [1.5, 0.0, 0.0]], dtype=np.float32)
+    cell = np.diag([10.0, 10.0, 10.0])
+    bins, gr = compute_rdf(
+        positions,
+        cell,
+        pbc=(True, True, True),
+        r_min=0.0,
+        r_max=5.0,
+        nbins=5,
+        device="cpu",
+        torch_dtype=torch.float32,
+        half_fill=False,  # cross-species -> no pair doubling
+        group_a_indices=[0],
+        group_b_indices=[1],
+    )
+    shell_vol = (4.0 / 3.0) * math.pi * (2.0**3 - 1.0**3)
+    volume = 10.0**3
+    n_a = 1
+    n_b = 1
+    expected = 1.0 / (shell_vol * (n_a * (n_b / volume)))
+    assert math.isclose(gr[1], expected, rel_tol=1e-5)