boltzkit.utils.molecular.energy_eval

Functions

create_simulation(topology, system[, ...])

Create a simulation object using the given parameters.

evaluate_energy_single(sim, pos[, ...])

Evaluate energy for a set of atomic positions specified in nanometers (nm).

get_openmm_platform([platform_name])

Classes

EnergyEval

ParallelEnergyEval

SequentialEnergyEval

boltzkit.utils.molecular.energy_eval._test_ABI_set_positions_compatibility(sim: Simulation)[source]

Test whether positions can be safely set in an OpenMM Simulation when using NumPy >= 2.0.0.

OpenMM may have been compiled against NumPy < 2.0.0 while the current environment uses NumPy >= 2.0.0. This can lead to a C-ABI mismatch, which may trigger warnings or errors if NumPy arrays are passed between Python and OpenMM (e.g., using asNumpy=True or passing ndarray inputs).

This function performs a test by attempting to set positions using the simulation’s context and verifying that the positions are correctly applied.

Parameters:

sim (app.Simulation) – The OpenMM Simulation instance to test for safe position setting.

Raises:

RuntimeError – If the test fails, indicating that positions cannot be safely set with the current OpenMM/NumPy combination. This may happen due to a NumPy C-ABI mismatch (OpenMM compiled with NumPy < 2.0.0 while NumPy >= 2.0.0 is installed). In that case, consider: - downgrading NumPy to < 2.0, or - rebuilding OpenMM against NumPy >= 2.0.

boltzkit.utils.molecular.energy_eval.evaluate_energy_single(sim: Simulation, pos: list[Vec3] | ndarray | Quantity, include_energy: bool = True, include_forces: bool = True) ndarray[source]

Evaluate energy for a set of atomic positions specified in nanometers (nm).

Parameters:
  • sim (app.Simulation) – Description

  • pos (list[mm.Vec3] | np.ndarray | unit.Quantity) – Atomic positions as list of vectors or numpy array of shape (n_atoms, 3)

  • include_energy (bool) – Whether to compute energy

  • include_forces (bool) – Whether to compute forces

Returns:

Return optional energy (in eV) and forces in (eV/nm)

Return type:

ndarray[_AnyShape, dtype[Any]]

boltzkit.utils.molecular.energy_eval.get_openmm_platform(platform_name: Literal['CPU', 'CUDA'] | None = None)[source]
boltzkit.utils.molecular.energy_eval.create_simulation(topology: Topology, system: System, platform_name: Literal['CPU', 'CUDA'] | None = None, integrator: Integrator | None = None)[source]

Create a simulation object using the given parameters.

Parameters:
  • topology (openmm.app.Topology) – Topology

  • system (openmm.System) – System

  • platform_name (Literal["CPU", "CUDA"] | None) – Executation platform. None means the platform is determined automatically, using CUDA if available.

  • integrator – Optional integrator (default is None). This is only relevant if the resulting simulation object is used

for something else than static energy evaluation. In this case, a Langevin integrator is used a dummy object. :type integrator: mm.Integrator | None

class boltzkit.utils.molecular.energy_eval.EnergyEval[source]

Bases: ABC

__init__()[source]
abstractmethod evaluate_batch(x: ndarray, include_energy: bool = True, include_forces: bool = True) tuple[ndarray | None, ndarray | None][source]
class boltzkit.utils.molecular.energy_eval.SequentialEnergyEval[source]

Bases: EnergyEval

__init__(topology: Topology, system: System, platform: Literal['CPU', 'CUDA'] | None = 'CPU')[source]
evaluate_batch(x, include_energy=True, include_forces=True)[source]
class boltzkit.utils.molecular.energy_eval.ParallelEnergyEval[source]

Bases: object

__init__(topology: Topology, system: System, platform: Literal['CPU', 'CUDA'] | None = 'CPU', n_workers: int = -1)[source]
evaluate_batch(x: ndarray, include_energy=True, include_forces=True, fragmentation=1)[source]

fragmentation (int, optional, default=1): Controls how samples are divided across workers.

By default, N samples are evenly split among K workers (blocks of size N/K), which minimizes interprocess communication when all workers have similar speeds.

If some workers are slower, they can become a bottleneck. Setting fragmentation > 1 splits the N samples into smaller blocks of size N / (K * fragmentation). This increases communication overhead but allows faster workers to process more blocks than slower ones, helping balance the workload.