deepmd.dpmodel.utils.dist_check

deepmd.dpmodel.utils.dist_check#

Pair-clearance check for frame validity filtering.

A frame is admissible when every atom pair keeps the clearance the training-frame filter prescribes. That clearance is a fraction of each pair’s own length scale, so the check is scale-free: the quantity compared against one is the margin

\[m = \min_{i<j} \frac{d_{ij}}{h_{t_i} + h_{t_j}},\]

with \(h_t\) the half-threshold of type \(t\). A radius given directly in Å is the same rule measured against a unit length scale, where every type carries the same half-threshold and the margin reduces to the minimum pair distance divided by that radius.

Attributes#

`_MIN_PAIR_DIST_BLOCK_PAIRS`
`_MIN_PAIR_DIST_BATCH_FRAMES`

Functions#

`pair_half_thresholds`(→ numpy.ndarray)	Tabulate the half-threshold each type contributes to a pair.
`requirement_half_thresholds`(→ numpy.ndarray)	Resolve and retain the half-thresholds a derived margin measures against.
`compute_min_pair_margin_single`(→ float)	Compute the minimum pair margin of a single frame.
`compute_min_pair_margin_batch`(→ numpy.ndarray)	Compute the minimum pair margin of every frame of a batch.
`_scan_by_atom_count`(→ numpy.ndarray)	Scan frames of real atoms laid out on one flat node axis.
`_scan_uniform`(→ numpy.ndarray)	Scan frames of one atom count, all of whose atoms are real.
`_min_pair_margin_block`(→ numpy.ndarray)	Minimum pair margin of frames of equal atom count.
`pair_margin_frame_mask`(→ numpy.ndarray \| None)	Mark the frames of a batch whose atom pairs all clear their window.
`select_frames`(→ dict[str, Any])	Keep the masked frames of a normalized NumPy batch.

Module Contents#

deepmd.dpmodel.utils.dist_check._MIN_PAIR_DIST_BLOCK_PAIRS = 262144[source]#

deepmd.dpmodel.utils.dist_check._MIN_PAIR_DIST_BATCH_FRAMES = 8[source]#

deepmd.dpmodel.utils.dist_check.pair_half_thresholds(fraction: float, scale: str, type_map: list[str] | None = None, ntypes: int | None = None) → numpy.ndarray[source]#

Tabulate the half-threshold each type contributes to a pair.

A pair’s threshold is half[a] + half[b]. Under the "covalent" scale that sum is the given fraction of the pair’s covalent bond length, so one fraction describes every element combination; under "absolute" every type carries the unit radius, the sum is the fraction itself, and the margin reduces to the pair distance divided by a threshold in Å.

Parameters:

fractionfloat: Filter radius as a fraction of the pair’s length scale, or directly in Å under the "absolute" scale.
scalestr: Either "covalent" or "absolute".
type_maplist[str], optional: Element symbols, required by the "covalent" scale.
ntypesint, optional: Number of types; taken from type_map when omitted.

Returns:

np.ndarray: Half-thresholds in Å with shape (ntypes,).

Raises:

ValueError: If the fraction is not positive, if the scale is unknown, if the covalent scale is requested without a type map, or if the map does not cover every type.

deepmd.dpmodel.utils.dist_check.requirement_half_thresholds(fraction: float, scale: str, type_map: tuple[str, ...] | None, ntypes: int) → numpy.ndarray[source]#

Resolve and retain the half-thresholds a derived margin measures against.

A reader scans every frame and every batch against one window and one type map, so the table is resolved once per combination instead of once per scan. The result is shared and therefore read-only.

Parameters:

fractionfloat: Filter radius, as a fraction of the pair length scale.
scalestr: Either "covalent" or "absolute".
type_maptuple[str, …], optional: Type names, required by the "covalent" scale.
ntypesint: Number of model types.

Returns:

np.ndarray: Half-thresholds in Å with shape (ntypes,).

deepmd.dpmodel.utils.dist_check.compute_min_pair_margin_single(coord: numpy.ndarray, box: numpy.ndarray | None, atype: numpy.ndarray, half: numpy.ndarray, screened: bool = False) → float[source]#

Compute the minimum pair margin of a single frame.

Pairs are visited in bounded row blocks, every one of them unless the scan is screened.

Parameters:

coordnp.ndarray: Atomic coordinates, flattened with shape (natoms * 3,) or reshaped as (natoms, 3), in Å.
boxnp.ndarray or None: Box vectors with shape (9,) for PBC, or None for non-PBC.
atypenp.ndarray: Atom types with shape (natoms,). Virtual atoms (type < 0) are excluded from the check.
halfnp.ndarray: Half-threshold of every type with shape (ntypes,) in Å.
screenedbool: Whether the scan may stop at a margin of one. A caller that only asks whether the frame clears its window learns that from the first pair inside it, and the remaining pairs of a failing frame cost the scan its whole cost on the largest frames.

Returns:

float: Minimum pair margin, or inf if fewer than 2 real atoms exist. A screened scan returns the exact margin for a frame that clears its window and some margin below one for a frame that does not.

deepmd.dpmodel.utils.dist_check.compute_min_pair_margin_batch(coord: numpy.ndarray, box: numpy.ndarray | None, atype: numpy.ndarray, half: numpy.ndarray, n_node: numpy.ndarray | None = None, screened: bool = False) → numpy.ndarray[source]#

Compute the minimum pair margin of every frame of a batch.

A batch is scanned in one pass over the frames that share an atom count rather than one pass per frame, and the scan of a frame visits the atom pairs that lie close along one lattice direction instead of every pair (see _min_pair_margin_block()).

Parameters:

coordnp.ndarray: Atomic coordinates, either (nframes, natoms, 3) (equivalently (nframes, natoms * 3)) for a rectangular batch, or (nnodes, 3) (equivalently (nnodes * 3,)) on the flat node axis of a ragged one, in Å.
boxnp.ndarray or None: Box vectors with shape (nframes, 9) for periodic frames, or None.
atypenp.ndarray: Atom types, (nframes, natoms) or (nnodes,). Virtual atoms (type < 0) are excluded from the check.
halfnp.ndarray: Half-threshold of every type with shape (ntypes,) in Å.
n_nodenp.ndarray or None: Per-frame atom counts of a ragged batch, with shape (nframes,). None marks a rectangular batch.
screenedbool: Bound the scan of a frame at a margin of one, which is all the frame filter reads. A frame holding a pair inside its window then keeps its exact margin, and a frame clearing it carries a margin above one that it reaches.

Returns:

np.ndarray: Margin of every frame, with shape (nframes,). A frame without a pair of real atoms carries inf.

deepmd.dpmodel.utils.dist_check._scan_by_atom_count(positions: numpy.ndarray, halves: numpy.ndarray, boxes: numpy.ndarray | None, counts: numpy.ndarray, screened: bool) → numpy.ndarray[source]#

Scan frames of real atoms laid out on one flat node axis.

Frames of equal atom count form the rectangular blocks the scan works on. A batch of the LMDB reader normally holds one count, and a padded rectangular batch reduces to one as soon as its filler atoms are left out, so this usually yields a single block; a block too small to pay for the batched scan is served frame by frame instead.

Parameters:

positionsnp.ndarray: Coordinates of the real atoms of every frame in order, (nnodes, 3).
halvesnp.ndarray: Half-threshold of each of those atoms, (nnodes,) in Å.
boxesnp.ndarray or None: Box vectors with shape (nframes, 9), or None for open boundaries.
countsnp.ndarray: Real atom count of every frame, with shape (nframes,).
screenedbool: Whether the scan is bounded at a margin of one.

Returns:

np.ndarray: Margin of every frame, with shape (nframes,).

deepmd.dpmodel.utils.dist_check._scan_uniform(coord: numpy.ndarray, halves: numpy.ndarray, boxes: numpy.ndarray | None, screened: bool) → numpy.ndarray[source]#

Scan frames of one atom count, all of whose atoms are real.

A cell of all zeros marks a frame without periodicity, which the readers may mix with periodic ones; the two kinds are scanned apart because the minimum image of a frame only exists with a cell.

Parameters:

coordnp.ndarray: Coordinates with shape (nframes, natoms, 3).
halvesnp.ndarray: Half-threshold of every atom, (nframes, natoms) in Å.
boxesnp.ndarray or None: Box vectors with shape (nframes, 9), or None for open boundaries.
screenedbool: Whether the scan is bounded at a margin of one.

Returns:

np.ndarray: Margin of every frame, with shape (nframes,).

deepmd.dpmodel.utils.dist_check._min_pair_margin_block(coord: numpy.ndarray, halves: numpy.ndarray, box: numpy.ndarray | None, screened: bool = False) → numpy.ndarray[source]#

Minimum pair margin of frames of equal atom count.

The atoms of every frame are ordered along one direction – the lattice planes that lie furthest apart, or the widest Cartesian spread without a cell – and each atom is paired with its successors in that order. A pair can only improve a frame’s margin m when its distance falls below m times its own threshold, hence below m times twice the largest half-threshold the frame carries; the separation along the ordering direction only grows with the distance in the order, so a frame is finished as soon as every atom’s window reaches that far. The scan therefore visits the successors that lie within one slab instead of every pair, and its bound is a bound on the margin rather than on the distance: the pair that minimizes the margin need not be the pair that minimizes the distance, and the scan never assumes it is.

Under periodicity the order is cyclic and every unordered pair is met from both of its atoms, at the windows shift and natoms - shift. That is what lets the advance be measured in the direction of the order: the pair whose periodic image is close the other way around is met from its other atom, where the advance is short.

Parameters:

coordnp.ndarray: Coordinates with shape (nframes, natoms, 3). Every atom counts.
halvesnp.ndarray: Half-threshold of every atom, (nframes, natoms) in Å.
boxnp.ndarray or None: Box vectors with shape (nframes, 9), or None for open boundaries.
screenedbool: Whether the scan is bounded at a margin of one.

Returns:

np.ndarray: Margin of every frame, with shape (nframes,).

deepmd.dpmodel.utils.dist_check.pair_margin_frame_mask(batch: dict[str, Any], enabled: bool) → numpy.ndarray | None[source]#

Mark the frames of a batch whose atom pairs all clear their window.

The filter radius the requirement was registered with is already spent: it sized the per-type half-thresholds the margin was measured against. What is left here is the same comparison against one for every length scale, so the only decision this function still takes from the configuration is whether the filter runs at all.

Parameters:

batchdict[str, Any]: Normalized NumPy batch carrying the derived pair_margin field, which holds each frame’s margin against its own pair thresholds.
enabledbool: Whether the pair-clearance filter is configured for this data.

Returns:

np.ndarray or None: Boolean mask with shape (nframes,), or None when the filter is disabled.

Raises:

RuntimeError: If the batch carries a placeholder instead of a derived margin, which would otherwise accept every frame silently.

deepmd.dpmodel.utils.dist_check.select_frames(batch: dict[str, Any], frame_mask: numpy.ndarray) → dict[str, Any][source]#

Keep the masked frames of a normalized NumPy batch.

A rectangular batch leads every per-frame and per-atom array with the frame axis. A ragged batch, which carries n_node, leads its per-frame arrays with the frame axis as well but concatenates its per-atom arrays frame by frame on one flat node axis, with a whole number of rows per atom; those are sliced by the node mask that the frame mask expands to.

Parameters:

batchdict[str, Any]: Normalized NumPy batch, rectangular or ragged.
frame_masknp.ndarray: Boolean mask with shape (nframes,).

Returns:

dict[str, Any]: New batch in which every array on the frame axis or on the flat node axis is sliced; all other entries pass through.