aimnet 0.0.1__py3-none-any.whl → 0.1.0__py3-none-any.whl

aimnet/calculators/nb_kernel_cpu.py

@@ -1,222 +0,0 @@
-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

@@ -1,217 +0,0 @@
-# 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

@@ -1,220 +0,0 @@
-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