custom_force

Custom force terms with symbolic expression evaluation.

import mlx_atomistic.custom_force

Classes

CustomForcePotential

class CustomForcePotential
    def __init__(indices: object, expression: str, parameters: object = (), global_parameters: object = (), name: str = 'custom_force', term_type: str = 'bond', supports_virial: bool = True)

Custom force term defined by a symbolic expression.

The expression string is parsed at construction time into an AST of mlx.core operations. Supported syntax:

• arithmetic: +, -, *, /, **
• functions: cos, sin, exp, log, sqrt, abs
• variables: geometric quantities (r, theta) and named per-term parameters supplied via parameters or global_parameters.

Parameters

indices : array-like (n_terms, W) index array where W is the term width (2 for bond/pair, 3 for angle). Each row lists the atom indices participating in that term. expression : str Symbolic energy-per-term expression. parameters : dict[str, array-like], optional Named per-term parameter arrays, each with shape (n_terms,). global_parameters : dict[str, float], optional Named scalar parameters available in the expression. name : str Descriptive name for diagnostics. term_type : str One of "bond" (2-body distance), "angle" (3-body angle), or "pair" (nonbonded-like 2-body). supports_virial : bool Whether the term advertises virial support.

Parameters

Name	Type	Default	Description
indices	object
expression	str
parameters	object	()
global_parameters	object	()
name	str	'custom_force'
term_type	str	'bond'
supports_virial	bool	True

Methods

energy_forces

def energy_forces(positions: mx.array, cell: Cell | None = None, pairs: mx.array | None = None) -> tuple[mx.array, mx.array]

Return the custom force-term energy and per-atom forces.

Parameters

Name	Type	Default	Description
positions	mx.array		Atomic coordinates, shape (n_atoms, 3).
cell	Cell | None	None	Optional periodic cell; when given, distances use the minimum-image convention. Defaults to None.
pairs	mx.array | None	None	Accepted for interface uniformity and ignored; the term uses its stored index list. Defaults to None.

Returns

• tuple[mx.array, mx.array] — An (energy, forces) tuple: scalar energy and per-atom forces of shape (n_atoms, 3).

potential_energy

def potential_energy(positions: mx.array, cell: Cell | None = None) -> mx.array

Return the custom force-term energy from the symbolic expression.

Parameters

Name	Type	Default	Description
positions	mx.array		Atomic coordinates, shape (n_atoms, 3).
cell	Cell | None	None	Optional periodic cell; when given, distances use the minimum-image convention. Defaults to None.

Returns

• mx.array — Total custom force-term energy as a scalar array.

Functions

evaluate_expression

def evaluate_expression(node: _ExprNode, env: dict[str, mx.array]) -> mx.array

Evaluate a parsed expression tree against a variable environment.

Parameters

Name	Type	Default	Description
node	_ExprNode		A parsed expression AST node.
env	dict[str, mx.array]		Mapping from variable name to its mx.array value.

Returns

• mx.array — The evaluated result as an mx.array.

parse_expression

def parse_expression(expr: str) -> _ExprNode

Parse a force-expression string into an evaluable expression tree.

Parameters

Name	Type	Default	Description
expr	str		The expression source string.

Returns

• _ExprNode — The parsed expression AST node.