nvalchemi.dynamics.hooks.FreezeAtomsHook

class nvalchemi.dynamics.hooks.FreezeAtomsHook(frequency=1, freeze_category=AtomCategory.SPECIAL.value, zero_forces=True)

Freeze selected atoms during molecular dynamics simulation.

During dynamics, certain atoms may need to remain fixed in place, such as substrate atoms in surface simulations, boundary atoms in slab models, or anchor atoms in constrained optimization. This hook identifies atoms by their atom_categories field and constrains them by:

1. Snapshotting positions — At BEFORE_PRE_UPDATE, the hook snapshots all atomic positions.

2. Restoring positions — At AFTER_POST_UPDATE, the hook restores positions of frozen atoms using torch.where, effectively undoing any displacement applied by the integrator.

3. Zeroing velocities — Velocities of frozen atoms are set to zero to prevent momentum accumulation.

4. Optionally zeroing forces — By default, forces on frozen atoms are also zeroed. This prevents force contributions from propagating through the integrator and ensures clean energy conservation diagnostics.

The hook fires at two stages: BEFORE_PRE_UPDATE (to snapshot positions) and AFTER_POST_UPDATE (to restore frozen positions and zero velocities/forces). This two-stage design enables torch.compile(fullgraph=True) compatibility by avoiding data-dependent branching.

Parameters:

• frequency (int, optional) – Apply constraints every frequency steps. Default 1 (every step). Setting this higher than 1 is not recommended as frozen atoms will drift between constraint applications.

• freeze_category (int, optional) – The atom_categories value that identifies frozen atoms. Default is AtomCategory.SPECIAL.value (-1). Atoms with batch.atom_categories == freeze_category will be frozen.

• zero_forces (bool, optional) – Whether to zero forces on frozen atoms. Default True. Set to False if you need to measure forces on frozen atoms for analysis purposes.

frequency

Constraint application frequency in steps.

Type:

int

freeze_category

Category value identifying frozen atoms.

Type:

int

zero_forces

Whether forces are zeroed on frozen atoms.

Type:

bool

stage

Primary stage, set to BEFORE_PRE_UPDATE for protocol compliance.

Type:

HookStageEnum

stages

Tuple of stages at which this hook fires: BEFORE_PRE_UPDATE and AFTER_POST_UPDATE.

Type:

tuple[HookStageEnum, …]

Examples

Freeze atoms marked as SPECIAL (default):

>>> from nvalchemi.dynamics.hooks import FreezeAtomsHook
>>> hook = FreezeAtomsHook()
>>> dynamics = DemoDynamics(model=model, n_steps=1000, hooks=[hook])
>>> dynamics.run(batch)

Freeze bulk atoms instead:

>>> from nvalchemi._typing import AtomCategory
>>> hook = FreezeAtomsHook(freeze_category=AtomCategory.BULK.value)

Keep forces for analysis:

>>> hook = FreezeAtomsHook(zero_forces=False)

Notes

• Fires at two stages: BEFORE_PRE_UPDATE (snapshot all positions) and AFTER_POST_UPDATE (restore frozen positions via torch.where).

• Uses torch.where for branchless GPU-vectorized restore, enabling torch.compile(fullgraph=True) compatibility.

• All positions are snapshotted each step (not just frozen ones) to avoid shape-dependent logic.

• When using with WrapPeriodicHook, both hooks fire at AFTER_POST_UPDATE. Registration order determines execution order; register this hook before the periodic wrapping hook to ensure frozen positions are restored before wrapping is applied.

__init__(frequency=1, freeze_category=AtomCategory.SPECIAL.value, zero_forces=True)

Parameters:

• frequency (int)

• freeze_category (int)

• zero_forces (bool)

Return type:

None

Methods

__init__([frequency, freeze_category, ...])

Attributes

stage
stages