Utilities for atom manipulation

cryojax.atom_util is a submodule that contains functions for manipulating arrays of atoms and their properties.

Chemistry

cryojax.atom_util.split_atoms_by_element(atomic_numbers: Int[ndarray, '_'], atom_pytree: PyTree[numpy.ndarray, 'T'], atom_axis: int = 0) -> tuple[PyTree[tuple[numpy.ndarray, ...], 'T'], Int[ndarray, '_']]

Given atomic numbers, split a pytree of numpy arrays representing atom properties into a tuple where each element is the pytree of tuples where each element is the array for a given atomic number.

Arguments:

• atomic_numbers: Atomic numbers as a numpy arrays.
• atom_pytree: A pytree of numpy arrays,
• atom_axis: The axis representing the atom index. Leading axes are assumed to be batched axes.

Returns:

A pytree with tree structure matching atom_pytree, where arrays have been replaced with tuples of arrays.

Geometry

cryojax.atom_util.center_atom_positions(atom_positions: Float[ndarray, '... n_atoms 3'], atom_masses: Float[ndarray, '... n_atoms'] | None = None, atom_axis: int = 0) -> Float[ndarray, '... n_atoms 3']

Center positions of a numpy array of atoms.

Arguments:

• atom_positions: The array of atom positions.
• atom_masses: Optionally, provide the masses of the atoms to center by mass.
• atom_axis: The axis representing the atom index. Leading axes are assumed to be batched axes. atom_positions and atom_masses are assumed to have the same batched axes.

Returns:

The centered atom_positions. This is centered by geometry if atom_masses = None and by mass otherwise.