|
| 1 | +# SPDX-License-Identifier: LGPL-3.0-or-later |
| 2 | +"""Minimum pairwise distance check for frame validity filtering.""" |
| 3 | + |
| 4 | +from __future__ import ( |
| 5 | + annotations, |
| 6 | +) |
| 7 | + |
| 8 | +import numpy as np |
| 9 | + |
| 10 | +_MIN_PAIR_DIST_BLOCK_PAIRS = 262_144 |
| 11 | + |
| 12 | + |
| 13 | +def compute_min_pair_dist_single( |
| 14 | + coord: np.ndarray, |
| 15 | + box: np.ndarray | None, |
| 16 | + atype: np.ndarray, |
| 17 | + stop_below: float | None = None, |
| 18 | +) -> float: |
| 19 | + """Compute the minimum pairwise atomic distance for a single frame. |
| 20 | +
|
| 21 | + Parameters |
| 22 | + ---------- |
| 23 | + coord : np.ndarray |
| 24 | + Atomic coordinates, flattened with shape (natoms * 3,) |
| 25 | + or reshaped as (natoms, 3). |
| 26 | + box : np.ndarray or None |
| 27 | + Box vectors with shape (9,) for PBC, or None for non-PBC. |
| 28 | + atype : np.ndarray |
| 29 | + Atom types with shape (natoms,). Virtual atoms (type < 0) |
| 30 | + are excluded from the distance check. |
| 31 | + stop_below : float or None |
| 32 | + Optional early-stop threshold. If a block has any pair closer |
| 33 | + than this value, the block minimum is returned immediately. |
| 34 | +
|
| 35 | + Returns |
| 36 | + ------- |
| 37 | + float |
| 38 | + Minimum pairwise distance. Returns inf if fewer than 2 |
| 39 | + real atoms exist. |
| 40 | + """ |
| 41 | + coord = coord.reshape(-1, 3) |
| 42 | + |
| 43 | + # === Step 1. Filter out virtual atoms === |
| 44 | + real_mask = atype.ravel() >= 0 |
| 45 | + real_coord = coord[real_mask] |
| 46 | + n_real = real_coord.shape[0] |
| 47 | + if n_real < 2: |
| 48 | + return float("inf") |
| 49 | + |
| 50 | + # === Step 2. Prepare minimum image convention for PBC === |
| 51 | + if box is not None: |
| 52 | + cell = box.reshape(3, 3) |
| 53 | + inv_cell = np.linalg.inv(cell) |
| 54 | + else: |
| 55 | + cell = None |
| 56 | + inv_cell = None |
| 57 | + |
| 58 | + # === Step 3. Compute distances in bounded row blocks === |
| 59 | + block_size = max(1, min(n_real, _MIN_PAIR_DIST_BLOCK_PAIRS // n_real)) |
| 60 | + min_dist_sq = float("inf") |
| 61 | + stop_dist_sq = ( |
| 62 | + float(stop_below) * float(stop_below) |
| 63 | + if stop_below is not None and stop_below > 0.0 |
| 64 | + else None |
| 65 | + ) |
| 66 | + for start in range(0, n_real, block_size): |
| 67 | + stop = min(start + block_size, n_real) |
| 68 | + diff = real_coord[np.newaxis, :, :] - real_coord[start:stop, np.newaxis, :] |
| 69 | + |
| 70 | + if cell is not None and inv_cell is not None: |
| 71 | + frac_diff = diff @ inv_cell |
| 72 | + frac_diff -= np.round(frac_diff) |
| 73 | + diff = frac_diff @ cell |
| 74 | + |
| 75 | + dist_sq = np.sum(diff * diff, axis=-1) |
| 76 | + rows = np.arange(stop - start, dtype=np.int64) |
| 77 | + dist_sq[rows, start + rows] = np.inf |
| 78 | + min_dist_sq = min(min_dist_sq, float(dist_sq.min())) |
| 79 | + if min_dist_sq == 0.0 or ( |
| 80 | + stop_dist_sq is not None and min_dist_sq < stop_dist_sq |
| 81 | + ): |
| 82 | + break |
| 83 | + |
| 84 | + return float(np.sqrt(min_dist_sq)) |
0 commit comments