aimnet 0.0.1__py3-none-any.whl

aimnet/calculators/nb_kernel_cpu.py

@@ -0,0 +1,222 @@
+import numba
+import numpy as np
+
+
+@numba.njit(cache=True, parallel=False, fastmath=True)
+def _nbmat_dual_cpu(
+    coord: np.ndarray,  # float, (N, 3)
+    cutoff1_squared: float,
+    cutoff2_squared: float,
+    mol_idx: np.ndarray,  # int, (N,)
+    mol_end_idx: np.ndarray,  # int, (M,)
+    nbmat1: np.ndarray,  # int, (N, maxnb1)
+    nbmat2: np.ndarray,  # int, (N, maxnb2)
+    nnb1: np.ndarray,  # int, zeros, (N,)
+    nnb2: np.ndarray,  # int, zeros, (N,)
+):
+    maxnb1 = nbmat1.shape[1]
+    maxnb2 = nbmat2.shape[1]
+    N = coord.shape[0]
+    for i in range(N):
+        c_i = coord[i]
+        _mol_idx = mol_idx[i]
+        _j_start = i + 1
+        _j_end = mol_end_idx[_mol_idx]
+        for j in range(_j_start, _j_end):
+            diff = c_i - coord[j]
+            dx, dy, dz = diff[0], diff[1], diff[2]
+            dist2 = dx * dx + dy * dy + dz * dz
+            if dist2 < cutoff1_squared:
+                pos = nnb1[i]
+                nnb1[i] += 1
+                if pos < maxnb1:
+                    nbmat1[i, pos] = j
+            if dist2 < cutoff2_squared:
+                pos = nnb2[i]
+                nnb2[i] += 1
+                if pos < maxnb2:
+                    nbmat2[i, pos] = j
+    _expand_nb(nnb1, nbmat1)
+    _expand_nb(nnb2, nbmat2)
+
+
+@numba.njit(cache=True, parallel=False, fastmath=True)
+def _nbmat_cpu(
+    coord: np.ndarray,  # float, (N, 3)
+    cutoff1_squared: float,
+    mol_idx: np.ndarray,  # int, (N,)
+    mol_end_idx: np.ndarray,  # int, (M,)
+    nbmat1: np.ndarray,  # int, (N, maxnb1)
+    nnb1: np.ndarray,  # int, zeros, (N,)
+):
+    maxnb1 = nbmat1.shape[1]
+    N = coord.shape[0]
+    for i in range(N):
+        c_i = coord[i]
+        _mol_idx = mol_idx[i]
+        _j_start = i + 1
+        _j_end = mol_end_idx[_mol_idx]
+        for j in range(_j_start, _j_end):
+            diff = c_i - coord[j]
+            dx, dy, dz = diff[0], diff[1], diff[2]
+            dist2 = dx * dx + dy * dy + dz * dz
+            if dist2 < cutoff1_squared:
+                pos = nnb1[i]
+                nnb1[i] += 1
+                if pos < maxnb1:
+                    nbmat1[i, pos] = j
+    _expand_nb(nnb1, nbmat1)
+
+
+@numba.njit(cache=True, inline="always")
+def _expand_nb(nnb, nbmat):
+    nnb_copy = nnb.copy()
+    N = nnb.shape[0]
+    for i in range(N):
+        for m in range(nnb_copy[i]):
+            if m >= nbmat.shape[1]:
+                continue
+            j = nbmat[i, m]
+            if j < N:
+                pos = nnb[j]
+                nnb[j] += 1
+                if pos < nbmat.shape[1]:
+                    nbmat[j, pos] = i
+
+
+@numba.njit(cache=True, inline="always")
+def _expand_nb_pbc(nnb, nbmat, shifts):
+    nnb_copy = nnb.copy()
+    N = nnb.shape[0]
+    for i in range(N):
+        for m in range(nnb_copy[i]):
+            if m >= nbmat.shape[1]:
+                continue
+            j = nbmat[i, m]
+            if j < N:
+                pos = nnb[j]
+                nnb[j] += 1
+                if pos < nbmat.shape[1]:
+                    nbmat[j, pos] = i
+                    shift = shifts[i, m]
+                    shifts[j, pos] = -shift
+
+
+@numba.njit(cache=True)
+def _expand_shifts(nshift):
+    tot_shifts = (nshift[0] + 1) * (2 * nshift[1] + 1) * (2 * nshift[2] + 1)
+    shifts = np.zeros((tot_shifts, 3), dtype=np.float32)
+    i = 0
+    for k1 in range(-nshift[0], nshift[0] + 1):
+        for k2 in range(-nshift[1], nshift[1] + 1):
+            for k3 in range(-nshift[2], nshift[2] + 1):
+                if k1 > 0 or (k1 == 0 and k2 > 0) or (k1 == 0 and k2 == 0 and k3 >= 0):
+                    shifts[i, 0] = k1
+                    shifts[i, 1] = k2
+                    shifts[i, 2] = k3
+                    i += 1
+    shifts = shifts[:i]
+    return shifts
+
+
+@numba.njit(cache=True, parallel=False, fastmath=True)
+def shift_coords(coord, cell, shifts):
+    N = coord.shape[0]
+    S = shifts.shape[0]
+    # pre-compute shifted coords
+    coord_shifted = np.empty((N, S, 3), dtype=coord.dtype)
+    for i in range(N):
+        for s in range(S):
+            shift = shifts[s]
+            c_x = coord[i, 0] + shift[0] * cell[0, 0] + shift[1] * cell[1, 0] + shift[2] * cell[2, 0]
+            c_y = coord[i, 1] + shift[0] * cell[0, 1] + shift[1] * cell[1, 1] + shift[2] * cell[2, 1]
+            c_z = coord[i, 2] + shift[0] * cell[0, 2] + shift[1] * cell[1, 2] + shift[2] * cell[2, 2]
+            coord_shifted[i, s] = c_x, c_y, c_z
+    return coord_shifted
+
+
+@numba.njit(cache=True, parallel=False, fastmath=True)
+def _nbmat_pbc_cpu(
+    coord: np.ndarray,  # float, (N, 3)
+    cell: np.ndarray,  # float, (3, 3)
+    cutoff1_squared: float,
+    shifts: np.ndarray,  # float, (S, 3)
+    nnb1: np.ndarray,  # int, zeros, (N,)
+    nbmat1: np.ndarray,  # int, (N, M)
+    shifts1: np.ndarray,  # int, (N, M, 3)
+):
+    maxnb1 = nbmat1.shape[1]
+    N = coord.shape[0]
+    S = shifts.shape[0]
+
+    coord_shifted = shift_coords(coord, cell, shifts)
+
+    for i in range(N):
+        c_i = coord[i]
+        for s in range(S):
+            shift = shifts[s]
+            zero_shift = shift[0] == 0 and shift[1] == 0 and shift[2] == 0
+            _j_end = i if zero_shift else N
+            for j in range(_j_end):
+                c_j = coord_shifted[j, s]
+                dx = c_i[0] - c_j[0]
+                dy = c_i[1] - c_j[1]
+                dz = c_i[2] - c_j[2]
+                r2 = dx * dx + dy * dy + dz * dz
+                if r2 < cutoff1_squared:
+                    pos = nnb1[i]
+                    nnb1[i] += 1
+                    if pos < maxnb1:
+                        nbmat1[i, pos] = j
+                        shifts1[i, pos] = shift
+    _expand_nb_pbc(nnb1, nbmat1, shifts1)
+
+
+@numba.njit(cache=True, parallel=False, fastmath=True)
+def _nbmat_dual_pbc_cpu(
+    coord: np.ndarray,  # float, (N, 3)
+    cell: np.ndarray,  # float, (3, 3)
+    cutoff1_squared: float,
+    cutoff2_squared: float,
+    shifts: np.ndarray,  # float, (S, 3)
+    nnb1: np.ndarray,  # int, zeros, (N,)
+    nnb2: np.ndarray,  # int, zeros, (N,)
+    nbmat1: np.ndarray,  # int, (N, M)
+    nbmat2: np.ndarray,  # int, (N, M)
+    shifts1: np.ndarray,  # int, (N, M, 3)
+    shifts2: np.ndarray,  # int, (N, M, 3)
+):
+    maxnb1 = nbmat1.shape[1]
+    maxnb2 = nbmat2.shape[1]
+    N = coord.shape[0]
+    S = shifts.shape[0]
+
+    coord_shifted = shift_coords(coord, cell, shifts)
+
+    for i in range(N):
+        c_i = coord[i]
+        for s in range(S):
+            shift = shifts[s]
+            zero_shift = shift[0] == 0 and shift[1] == 0 and shift[2] == 0
+            _j_end = i if zero_shift else N
+            for j in range(_j_end):
+                c_j = coord_shifted[j, s]
+                dx = c_i[0] - c_j[0]
+                dy = c_i[1] - c_j[1]
+                dz = c_i[2] - c_j[2]
+                r2 = dx * dx + dy * dy + dz * dz
+                if r2 < cutoff1_squared:
+                    pos = nnb1[i]
+                    nnb1[i] += 1
+                    if pos < maxnb1:
+                        nbmat1[i, pos] = j
+                        shifts1[i, pos] = shift
+                if r2 < cutoff2_squared:
+                    pos = nnb2[i]
+                    nnb2[i] += 1
+                    if pos < maxnb2:
+                        nbmat2[i, pos] = j
+                        shifts2[i, pos] = shift
+
+    _expand_nb_pbc(nnb1, nbmat1, shifts1)
+    _expand_nb_pbc(nnb2, nbmat2, shifts2)

aimnet/calculators/nb_kernel_cuda.py

@@ -0,0 +1,217 @@
+# type: ignore
+import numba
+import numba.cuda
+from numba.core import config
+
+config.CUDA_LOW_OCCUPANCY_WARNINGS = 0
+
+
+@numba.cuda.jit(fastmath=True, cache=True)
+def _nbmat_dual_cuda(coord, cutoff1_squared, cutoff2_squared, mol_idx, mol_end_idx, nbmat1, nbmat2, nnb1, nnb2):
+    N = coord.shape[0]
+    i = numba.cuda.grid(1)
+
+    if i >= N:
+        return
+
+    c0 = coord[i, 0]
+    c1 = coord[i, 1]
+    c2 = coord[i, 2]
+
+    maxnb1 = nbmat1.shape[1]
+    maxnb2 = nbmat2.shape[1]
+
+    _mol_idx = mol_idx[i]
+    _j_start = i + 1
+    _j_end = mol_end_idx[_mol_idx]
+
+    for j in range(_j_start, _j_end):
+        d0 = c0 - coord[j, 0]
+        d1 = c1 - coord[j, 1]
+        d2 = c2 - coord[j, 2]
+        dist_squared = d0 * d0 + d1 * d1 + d2 * d2
+        if dist_squared < cutoff1_squared:
+            pos = numba.cuda.atomic.add(nnb1, i, 1)
+            if pos < maxnb1:
+                nbmat1[i, pos] = j
+            pos = numba.cuda.atomic.add(nnb1, j, 1)
+            if pos < maxnb1:
+                nbmat1[j, pos] = i
+        if dist_squared < cutoff2_squared:
+            pos = numba.cuda.atomic.add(nnb2, i, 1)
+            if pos < maxnb2:
+                nbmat2[i, pos] = j
+            pos = numba.cuda.atomic.add(nnb2, j, 1)
+            if pos < maxnb2:
+                nbmat2[j, pos] = i
+
+
+@numba.cuda.jit(fastmath=True, cache=True)
+def _nbmat_cuda(coord, cutoff1_squared, mol_idx, mol_end_idx, nbmat1, nnb1):
+    N = coord.shape[0]
+    i = numba.cuda.grid(1)
+
+    if i >= N:
+        return
+
+    c0 = coord[i, 0]
+    c1 = coord[i, 1]
+    c2 = coord[i, 2]
+
+    maxnb1 = nbmat1.shape[1]
+
+    _mol_idx = mol_idx[i]
+    _j_start = i + 1
+    _j_end = mol_end_idx[_mol_idx]
+
+    for j in range(_j_start, _j_end):
+        d0 = c0 - coord[j, 0]
+        d1 = c1 - coord[j, 1]
+        d2 = c2 - coord[j, 2]
+        dist_squared = d0 * d0 + d1 * d1 + d2 * d2
+        if dist_squared < cutoff1_squared:
+            pos = numba.cuda.atomic.add(nnb1, i, 1)
+            if pos < maxnb1:
+                nbmat1[i, pos] = j
+            pos = numba.cuda.atomic.add(nnb1, j, 1)
+            if pos < maxnb1:
+                nbmat1[j, pos] = i
+
+
+@numba.cuda.jit(cache=True, fastmath=True)
+def _nbmat_pbc_dual_cuda(
+    coord,  # N, 3
+    cell,  # 3, 3
+    cutoff1_squared: float,
+    cutoff2_squared: float,
+    shifts,  # S, 3
+    nnb1,  # N
+    nnb2,  # N
+    nbmat1,  # N, M
+    nbmat2,  # N, K
+    shifts1,  # N, M, 3
+    shifts2,  # N, K, 3
+):
+    idx = numba.cuda.grid(1)
+
+    _n = coord.shape[0]
+    _s = shifts.shape[0]
+
+    shift_idx, atom_idx = idx // _n, idx % _n
+    if shift_idx >= _s:
+        return
+
+    maxnb1 = nbmat1.shape[1]
+    maxnb2 = nbmat2.shape[1]
+
+    shift_x = shifts[shift_idx, 0]
+    shift_y = shifts[shift_idx, 1]
+    shift_z = shifts[shift_idx, 2]
+
+    zero_shift = shift_x == 0 and shift_y == 0 and shift_z == 0
+
+    shift_x = numba.float32(shift_x)
+    shift_y = numba.float32(shift_y)
+    shift_z = numba.float32(shift_z)
+
+    coord_shifted_x = coord[atom_idx, 0] + shift_x * cell[0, 0] + shift_y * cell[1, 0] + shift_z * cell[2, 0]
+    coord_shifted_y = coord[atom_idx, 1] + shift_x * cell[0, 1] + shift_y * cell[1, 1] + shift_z * cell[2, 1]
+    coord_shifted_z = coord[atom_idx, 2] + shift_x * cell[0, 2] + shift_y * cell[1, 2] + shift_z * cell[2, 2]
+
+    for i in range(_n):
+        if zero_shift and i >= atom_idx:
+            continue
+
+        dx = coord_shifted_x - coord[i, 0]
+        dy = coord_shifted_y - coord[i, 1]
+        dz = coord_shifted_z - coord[i, 2]
+
+        r2 = dx * dx + dy * dy + dz * dz
+
+        if r2 < cutoff1_squared:
+            pos = numba.cuda.atomic.add(nnb1, i, 1)
+            if pos < maxnb1:
+                nbmat1[i, pos] = atom_idx
+                shifts1[i, pos, 0] = shift_x
+                shifts1[i, pos, 1] = shift_y
+                shifts1[i, pos, 2] = shift_z
+            pos = numba.cuda.atomic.add(nnb1, atom_idx, 1)
+            if pos < maxnb1:
+                nbmat1[atom_idx, pos] = i
+                shifts1[atom_idx, pos, 0] = -shift_x
+                shifts1[atom_idx, pos, 1] = -shift_y
+                shifts1[atom_idx, pos, 2] = -shift_z
+
+        if r2 < cutoff2_squared:
+            pos = numba.cuda.atomic.add(nnb2, i, 1)
+            if pos < maxnb2:
+                nbmat2[i, pos] = atom_idx
+                shifts2[i, pos, 0] = shift_x
+                shifts2[i, pos, 1] = shift_y
+                shifts2[i, pos, 2] = shift_z
+            pos = numba.cuda.atomic.add(nnb2, atom_idx, 1)
+            if pos < maxnb2:
+                nbmat2[atom_idx, pos] = i
+                shifts2[atom_idx, pos, 0] = -shift_x
+                shifts2[atom_idx, pos, 1] = -shift_y
+                shifts2[atom_idx, pos, 2] = -shift_z
+
+
+@numba.cuda.jit(cache=True, fastmath=True)
+def _nbmat_pbc_cuda(
+    coord,  # N, 3
+    cell,  # 3, 3
+    cutoff1_squared: float,
+    shifts,  # S, 3
+    nnb1,  # N
+    nbmat1,  # N, M
+    shifts1,  # N, M, 3
+):
+    idx = numba.cuda.grid(1)
+
+    _n = coord.shape[0]
+    _s = shifts.shape[0]
+
+    shift_idx, atom_idx = idx // _n, idx % _n
+    if shift_idx >= _s:
+        return
+
+    maxnb1 = nbmat1.shape[1]
+
+    shift_x = shifts[shift_idx, 0]
+    shift_y = shifts[shift_idx, 1]
+    shift_z = shifts[shift_idx, 2]
+
+    zero_shift = shift_x == 0 and shift_y == 0 and shift_z == 0
+
+    shift_x = numba.float32(shift_x)
+    shift_y = numba.float32(shift_y)
+    shift_z = numba.float32(shift_z)
+
+    coord_shifted_x = coord[atom_idx, 0] + shift_x * cell[0, 0] + shift_y * cell[1, 0] + shift_z * cell[2, 0]
+    coord_shifted_y = coord[atom_idx, 1] + shift_x * cell[0, 1] + shift_y * cell[1, 1] + shift_z * cell[2, 1]
+    coord_shifted_z = coord[atom_idx, 2] + shift_x * cell[0, 2] + shift_y * cell[1, 2] + shift_z * cell[2, 2]
+
+    for i in range(_n):
+        if zero_shift and i >= atom_idx:
+            continue
+
+        dx = coord_shifted_x - coord[i, 0]
+        dy = coord_shifted_y - coord[i, 1]
+        dz = coord_shifted_z - coord[i, 2]
+
+        r2 = dx * dx + dy * dy + dz * dz
+
+        if r2 < cutoff1_squared:
+            pos = numba.cuda.atomic.add(nnb1, i, 1)
+            if pos < maxnb1:
+                nbmat1[i, pos] = atom_idx
+                shifts1[i, pos, 0] = shift_x
+                shifts1[i, pos, 1] = shift_y
+                shifts1[i, pos, 2] = shift_z
+            pos = numba.cuda.atomic.add(nnb1, atom_idx, 1)
+            if pos < maxnb1:
+                nbmat1[atom_idx, pos] = i
+                shifts1[atom_idx, pos, 0] = -shift_x
+                shifts1[atom_idx, pos, 1] = -shift_y
+                shifts1[atom_idx, pos, 2] = -shift_z

aimnet/calculators/nbmat.py

@@ -0,0 +1,220 @@
+from typing import Optional, Tuple
+
+import torch
+from torch import Tensor
+
+from .nb_kernel_cpu import _expand_shifts
+
+
+class TooManyNeighborsError(Exception):
+    pass
+
+
+if torch.cuda.is_available():
+    import numba.cuda
+
+    if not numba.cuda.is_available():
+        raise ImportError("PyTorch CUDA is available, but Numba CUDA is not available.")
+    _numba_cuda_available = True
+    from .nb_kernel_cuda import _nbmat_cuda, _nbmat_dual_cuda, _nbmat_pbc_cuda, _nbmat_pbc_dual_cuda
+
+    _kernel_nbmat = _nbmat_cuda
+    _kernel_nbmat_dual = _nbmat_dual_cuda
+    _kernel_nbmat_pbc = _nbmat_pbc_cuda
+    _kernel_nbmat_pbc_dual = _nbmat_pbc_dual_cuda
+else:
+    _numba_cuda_available = False
+    from .nb_kernel_cpu import _nbmat_cpu, _nbmat_dual_cpu, _nbmat_dual_pbc_cpu, _nbmat_pbc_cpu
+
+    _kernel_nbmat = _nbmat_cpu
+    _kernel_nbmat_dual = _nbmat_dual_cpu
+    _kernel_nbmat_pbc = _nbmat_pbc_cpu
+    _kernel_nbmat_pbc_dual = _nbmat_dual_pbc_cpu
+
+
+def calc_nbmat(
+    coord: Tensor,
+    cutoffs: Tuple[float, Optional[float]],
+    maxnb: Tuple[int, Optional[int]],
+    cell: Optional[Tensor] = None,
+    mol_idx: Optional[Tensor] = None,
+):
+    device = coord.device
+    N = coord.shape[0]
+
+    _pbc = cell is not None
+    if _pbc and mol_idx is not None and mol_idx[-1] > 0:
+        raise ValueError("Multiple molecules are not supported with PBC.")
+
+    if mol_idx is None:
+        mol_idx = torch.zeros(N, dtype=torch.long, device=device)
+        mol_end_idx = torch.tensor([N], dtype=torch.long, device=device)
+    else:
+        _, mol_size = torch.unique(mol_idx, return_counts=True)
+        mol_end_idx = mol_size.cumsum(0)
+
+    if _numba_cuda_available and device.type != "cuda":
+        raise ValueError("Numba CUDA is available, but the input tensors are not on CUDA.")
+
+    _cuda = device.type == "cuda" and _numba_cuda_available
+    _dual_cutoff = cutoffs[1] is not None
+    if _dual_cutoff and maxnb[1] is None:
+        raise ValueError("maxnb[1] must be specified for dual cutoff.")
+
+    nnb1 = torch.zeros(N, dtype=torch.long, device=device)
+    nbmat1 = torch.full((N + 1, maxnb[0]), N, dtype=torch.long, device=device)
+
+    if _dual_cutoff:
+        nnb2 = torch.zeros(N, dtype=torch.long, device=device)
+        nbmat2 = torch.full((N + 1, maxnb[1]), N, dtype=torch.long, device=device)  # type: ignore
+
+    if _pbc:
+        cell_inv = torch.inverse(cell)  # type: ignore[arg-type]
+        cutoff = max(cutoffs) if _dual_cutoff else cutoffs[0]  # type: ignore
+        nshift = torch.ceil(cutoff * cell_inv.norm(dim=-1)).to(torch.long).cpu().numpy()
+        shifts = _expand_shifts(nshift)
+        S = shifts.shape[0]
+        shifts = torch.from_numpy(shifts).to(device)
+        shifts1 = torch.zeros(N + 1, maxnb[0], 3, dtype=torch.long, device=device)
+        if _dual_cutoff:
+            shifts2 = torch.zeros(N + 1, maxnb[1], 3, dtype=torch.long, device=device)  # type: ignore
+    else:
+        S = 1
+
+    # convert tensors and launch the kernel
+    if _cuda:
+        _coord = numba.cuda.as_cuda_array(coord)
+        _mol_idx = numba.cuda.as_cuda_array(mol_idx)
+        _mol_end_idx = numba.cuda.as_cuda_array(mol_end_idx)
+        _nnb1 = numba.cuda.as_cuda_array(nnb1)
+        _nbmat1 = numba.cuda.as_cuda_array(nbmat1)
+        if _dual_cutoff:
+            _nnb2 = numba.cuda.as_cuda_array(nnb2)
+            _nbmat2 = numba.cuda.as_cuda_array(nbmat2)
+        if _pbc:
+            _cell = numba.cuda.as_cuda_array(cell)
+            _shifts = numba.cuda.as_cuda_array(shifts)
+            _shifts1 = numba.cuda.as_cuda_array(shifts1)
+            if _dual_cutoff:
+                _shifts2 = numba.cuda.as_cuda_array(shifts2)
+        threads_per_block = 32
+        blocks_per_grid = (N * S + (threads_per_block - 1)) // threads_per_block
+
+        if _pbc:
+            if _dual_cutoff:
+                _kernel_nbmat_pbc_dual[blocks_per_grid, threads_per_block](  # type: ignore
+                    _coord,
+                    _cell,
+                    cutoffs[0] ** 2,
+                    cutoffs[1] ** 2,  # type: ignore
+                    _shifts,
+                    _nnb1,
+                    _nnb2,
+                    _nbmat1,
+                    _nbmat2,
+                    _shifts1,
+                    _shifts2,
+                )
+            else:
+                _kernel_nbmat_pbc[blocks_per_grid, threads_per_block](  # type: ignore
+                    _coord,
+                    _cell,
+                    cutoffs[0] ** 2,
+                    _shifts,
+                    _nnb1,
+                    _nbmat1,
+                    _shifts1,
+                )
+        else:
+            if _dual_cutoff:
+                _kernel_nbmat_dual[blocks_per_grid, threads_per_block](  # type: ignore
+                    _coord,
+                    cutoffs[0] ** 2,
+                    cutoffs[1] ** 2,  # type: ignore
+                    _mol_idx,
+                    _mol_end_idx,
+                    _nbmat1,
+                    _nbmat2,
+                    _nnb1,
+                    _nnb2,
+                )
+            else:
+                _kernel_nbmat[blocks_per_grid, threads_per_block](  # type: ignore
+                    _coord,
+                    cutoffs[0] ** 2,
+                    _mol_idx,
+                    _mol_end_idx,
+                    _nbmat1,
+                    _nnb1,
+                )
+
+    else:
+        _coord = coord.numpy()
+        _mol_idx = mol_idx.numpy()
+        _mol_end_idx = mol_end_idx.numpy()
+        _nnb1 = nnb1.numpy()
+        _nbmat1 = nbmat1.numpy()
+        if _dual_cutoff:
+            _nnb2 = nnb2.numpy()
+            _nbmat2 = nbmat2.numpy()
+        if _pbc:
+            _cell = cell.numpy()  # type: ignore[union-attr]
+            _shifts = shifts.numpy()
+
+        if _pbc:
+            _shifts1 = shifts1.numpy()
+            if _dual_cutoff:
+                _shifts2 = shifts2.numpy()
+                _kernel_nbmat_pbc_dual(
+                    _coord,
+                    _cell,
+                    cutoffs[0] ** 2,
+                    cutoffs[1] ** 2,  # type: ignore
+                    _shifts,
+                    _nnb1,
+                    _nnb2,
+                    _nbmat1,
+                    _nbmat2,
+                    _shifts1,  # type: ignore
+                    _shifts2,  # type: ignore
+                )  # type: ignore
+            else:
+                _kernel_nbmat_pbc(_coord, _cell, cutoffs[0] ** 2, _shifts, _nnb1, _nbmat1, _shifts1)  # type: ignore
+        else:
+            if _dual_cutoff:
+                _kernel_nbmat_dual(
+                    _coord,
+                    cutoffs[0] ** 2,
+                    cutoffs[1] ** 2,  # type: ignore
+                    _mol_idx,
+                    _mol_end_idx,
+                    _nbmat1,
+                    _nbmat2,
+                    _nnb1,
+                    _nnb2,
+                )  # type: ignore
+            else:
+                _kernel_nbmat(_coord, cutoffs[0] ** 2, _mol_idx, _mol_end_idx, _nbmat1, _nnb1)
+
+    if not _pbc:
+        shifts1 = None  # type: ignore[assignment]
+        shifts2 = None  # type: ignore[assignment]
+
+    nnb1_max = nnb1.max().item()
+    if nnb1_max > maxnb[0]:
+        raise TooManyNeighborsError(f"maxnb is too small: {nnb1_max=}, {maxnb=}")
+    nbmat1 = nbmat1[:, :nnb1_max]  # type: ignore
+    if _pbc:
+        shifts1 = shifts1[:, :nnb1_max]  # type: ignore
+    if _dual_cutoff:
+        nnb2_max = nnb2.max().item()
+        if nnb2_max > maxnb[1]:  # type: ignore
+            raise TooManyNeighborsError(f"maxnb is too small: {nnb1_max=}, {nnb2_max=}, {maxnb=}")
+        nbmat2 = nbmat2[:, :nnb2_max]
+        if _pbc:
+            shifts2 = shifts2[:, :nnb2_max]  # type: ignore
+    else:
+        nbmat2 = None
+        if _pbc:
+            shifts2 = None
+    return nbmat1, nbmat2, shifts1, shifts2

aimnet/cli.py

@@ -0,0 +1,22 @@
+import click
+
+from .train.calc_sae import calc_sae
+from .train.pt2jpt import jitcompile
+from .train.train import train
+
+
+@click.group()
+def cli():
+    """AIMNet2 command line tool"""
+
+
+cli.add_command(train, name="train")
+cli.add_command(jitcompile, name="jitcompile")
+cli.add_command(calc_sae, name="calc_sae")
+
+
+if __name__ == "__main__":
+    import logging
+
+    logging.basicConfig(level=logging.INFO)
+    cli()