virtual_sites

Virtual-site geometry definitions for molecular mechanics.

import mlx_atomistic.virtual_sites

Classes

LocalCoordinates

class LocalCoordinates
    def __init__(particle1: int, particle2: int, particle3: int, weight1: float, weight2: float, weight3: float, local_x: float, local_y: float, local_z: float)

Virtual site defined by local coordinate offsets from three parent atoms.

The virtual site position is: p = p1w1 + p2w2 + p3w3 + xu1 + yu2 + zu3

Parameters

Name	Type	Default	Description
particle1	int
particle2	int
particle3	int
weight1	float
weight2	float
weight3	float
local_x	float
local_y	float
local_z	float

Properties

• parent_atoms tuple[int, int, int] — Indices of the three parent atoms.

Methods

compute_position

def compute_position(positions: mx.array) -> mx.array

Return the virtual-site position from the parent-atom positions.

Parameters

Name	Type	Default	Description
positions	mx.array		Cartesian positions of all atoms, shape (n_atoms, 3).

Returns

• mx.array — The virtual-site position, shape (3,).

OutOfPlane

class OutOfPlane
    def __init__(particle1: int, particle2: int, particle3: int, weight1: float, weight2: float, weight3: float, distance: float)

Virtual site computed from three parent atoms using cross-product offset.

Following the OpenMM convention, the position is computed as: p = p1w1 + p2w2 + p3w3 + d(cross(p2-p1, p3-p1) / |cross(p2-p1, p3-p1)|) where d is the out-of-plane distance.

Parameters

Name	Type	Default	Description
particle1	int
particle2	int
particle3	int
weight1	float
weight2	float
weight3	float
distance	float

Properties

• parent_atoms tuple[int, int, int] — Indices of the three parent atoms.

Methods

compute_position

def compute_position(positions: mx.array) -> mx.array

Return the virtual-site position from the parent-atom positions.

Parameters

Name	Type	Default	Description
positions	mx.array		Cartesian positions of all atoms, shape (n_atoms, 3).

Returns

• mx.array — The virtual-site position, shape (3,).

ThreeParticleAverage

class ThreeParticleAverage
    def __init__(particle1: int, particle2: int, particle3: int, weight1: float, weight2: float, weight3: float)

Virtual site at a weighted average of three parent atom positions.

Parameters

Name	Type	Default	Description
particle1	int
particle2	int
particle3	int
weight1	float
weight2	float
weight3	float

Properties

• parent_atoms tuple[int, int, int] — Indices of the three parent atoms.

Methods

compute_position

def compute_position(positions: mx.array) -> mx.array

Return the virtual-site position from the parent-atom positions.

Parameters

Name	Type	Default	Description
positions	mx.array		Cartesian positions of all atoms, shape (n_atoms, 3).

Returns

• mx.array — The virtual-site position, shape (3,).

TwoParticleAverage

class TwoParticleAverage
    def __init__(particle1: int, particle2: int, weight1: float, weight2: float)

Virtual site at a weighted average of two parent atom positions.

Parameters

Name	Type	Default	Description
particle1	int
particle2	int
weight1	float
weight2	float

Properties

• parent_atoms tuple[int, int] — Indices of the two parent atoms.

Methods

compute_position

def compute_position(positions: mx.array) -> mx.array

Return the virtual-site position from the parent-atom positions.

Parameters

Name	Type	Default	Description
positions	mx.array		Cartesian positions of all atoms, shape (n_atoms, 3).

Returns

• mx.array — The virtual-site position, shape (3,).

VirtualSiteManager

class VirtualSiteManager
    def __init__(virtual_sites: tuple, n_real_atoms: int)

Manages virtual site position reconstruction and force redistribution.

Real atom indices occupy positions 0 .. n_real_atoms - 1. Virtual site indices follow at n_real_atoms .. n_real_atoms + n_virtual_sites - 1 in the same order as virtual_sites.

Parameters

Name	Type	Default	Description
virtual_sites	tuple
n_real_atoms	int

Properties

• n_total_atoms int — Total atom count: real atoms plus virtual sites.
• n_virtual_sites int — Number of managed virtual sites.

Methods

extend_positions

def extend_positions(positions: mx.array) -> mx.array

Append virtual site positions to real-atom positions.

Returns an array of shape (n_total_atoms, 3) whose first n_real_atoms rows equal positions and whose remaining rows are reconstructed from the parent atoms.

Parameters

Name	Type	Default	Description
positions	mx.array

Returns

• mx.array

reconstruct_positions

def reconstruct_positions(positions: mx.array) -> mx.array

Recompute virtual site rows from current parent atom positions.

positions must have shape (n_total_atoms, 3). The first n_real_atoms rows are taken as-is; the remaining rows are overwritten with freshly reconstructed virtual site positions.

Parameters

Name	Type	Default	Description
positions	mx.array

Returns

• mx.array

redistribute_forces

def redistribute_forces(forces: mx.array, positions: mx.array) -> mx.array

Redistribute virtual site forces to parent atoms.

forces has shape (n_total_atoms, 3) and positions has shape (n_total_atoms, 3). Returns a (n_real_atoms, 3) array where each virtual site’s force has been distributed to its parents.

Parameters

Name	Type	Default	Description
forces	mx.array
positions	mx.array

Returns

• mx.array

Functions

compute_virtual_site_positions

def compute_virtual_site_positions(virtual_sites: tuple, positions: mx.array) -> mx.array

Compute positions for a sequence of virtual sites given parent atom positions.

Parameters

Name	Type	Default	Description
virtual_sites	tuple
positions	mx.array

Returns

• mx.array

redistribute_virtual_site_force

def redistribute_virtual_site_force(vs: TwoParticleAverage | ThreeParticleAverage | OutOfPlane | LocalCoordinates, positions: mx.array, vs_force: mx.array) -> dict[int, mx.array]

Distribute the force on a single virtual site to its parent atoms.

For linear-combination types (TwoParticleAverage, ThreeParticleAverage) the redistribution uses the constant weights. For position-dependent types (OutOfPlane, LocalCoordinates) a vector-Jacobian product (VJP) is computed via mx.grad.

Parameters

Name	Type	Default	Description
vs	TwoParticleAverage | ThreeParticleAverage | OutOfPlane | LocalCoordinates
positions	mx.array
vs_force	mx.array

Returns

• dict[int, mx.array]

tip4p_ew_m_site_weights

def tip4p_ew_m_site_weights(*, oh_distance: float = TIP4P_EW_OH_DISTANCE_ANGSTROM, hoh_angle_degrees: float = TIP4P_EW_HOH_ANGLE_DEGREES, om_distance: float = TIP4P_EW_OM_DISTANCE_ANGSTROM) -> tuple[float, float, float]

Return OpenMM-compatible O/H/H weights for the TIP4P-Ew M site.

Parameters

Name	Type	Default	Description
oh_distance	float	TIP4P_EW_OH_DISTANCE_ANGSTROM
hoh_angle_degrees	float	TIP4P_EW_HOH_ANGLE_DEGREES
om_distance	float	TIP4P_EW_OM_DISTANCE_ANGSTROM

Returns

• tuple[float, float, float]

tip4p_ew_reference_positions

def tip4p_ew_reference_positions() -> np.ndarray

Return O, H1, H2, M coordinates for an ideal TIP4P-Ew water in angstrom.

Returns

• np.ndarray

tip4p_ew_virtual_site

def tip4p_ew_virtual_site(oxygen: int, hydrogen1: int, hydrogen2: int, *, oh_distance: float = TIP4P_EW_OH_DISTANCE_ANGSTROM, hoh_angle_degrees: float = TIP4P_EW_HOH_ANGLE_DEGREES, om_distance: float = TIP4P_EW_OM_DISTANCE_ANGSTROM) -> ThreeParticleAverage

Create the TIP4P-Ew charge-site geometry for one water molecule.

Parameters

Name	Type	Default	Description
oxygen	int
hydrogen1	int
hydrogen2	int
oh_distance	float	TIP4P_EW_OH_DISTANCE_ANGSTROM
hoh_angle_degrees	float	TIP4P_EW_HOH_ANGLE_DEGREES
om_distance	float	TIP4P_EW_OM_DISTANCE_ANGSTROM

Returns

• ThreeParticleAverage