pantea.atoms package

Submodules

pantea.atoms.box module

class pantea.atoms.box.Box(lattice)

Bases: BaseJaxPytreeDataClass

Simulation box is used to apply periodic boundary conditions (PBC) in the presence of a lattice info.

Warning

The current implementation only works for orthogonal cells and does not support triclinic cells.

apply_pbc(dx)

Apply periodic boundary condition (PBC) on the atom positions.

For this method to function correctly, it is essential that all atoms are initially positioned within the boundaries of the box. Otherwise, the results may not be as anticipated. This could happen for when example time step is too large.

Parameters

dx (Array) – positional differences

Returns

PBC applied position differences

Return type

Array

property dtype: _ScalarMeta

Return type

_ScalarMeta

classmethod from_list(data, dtype=None)

Return type

Box

lattice: ArrayImpl

property length: ArrayImpl

Return length of cell in x, y, and z-directions.

Return type

ArrayImpl

property lx: ArrayImpl

Return length of cell in x-direction.

Return type

ArrayImpl

property ly: ArrayImpl

Return length of cell in y-direction.

Return type

ArrayImpl

property lz: ArrayImpl

Return length of cell in z-direction.

Return type

ArrayImpl

property volume: ArrayImpl

Return calculated volume of the box.

Return type

ArrayImpl

wrap_into_box(positions)

Adjust the coordinates of the atoms to ensure they fall within the boundaries of the simulation box that has periodic boundary conditions (PBC).

Parameters

positions (Array) – atom positions

Returns

wrapped atom positions using the PBC.

Return type

Array

pantea.atoms.distance module

class pantea.atoms.distance.StructureInterface(*args, **kwargs)

Bases: Protocol

lattice: ArrayImpl

positions: ArrayImpl

pantea.atoms.distance.calculate_distances(structure, atom_index=None, neighbor_atom_index=None, with_aux=False)

Calculate distances between specific atoms and the neighboring atoms in the structure. This function optionally also returns the corresponding position differences.

If atom indices are not specified, all atoms in the structure will be taken into account. Similarly, if neighbor indices are not provided, all neighboring atoms will be considered.

Parameters

• structure (StructureInterface) – input structure

• atom_index (Optional[Array], optional) – array of atom indices (zero-based)

• neighbor_atom_index (Optional[Array], optional) – array of indices of neighbor atoms (zero-based), defaults to None

• with_aux – whether returning position differences, defaults to False

Type

bool, optional

Returns

either an array of distances or tuple of distances together and position differences

Return type

Union[Array, tuple[Array, Array]]

pantea.atoms.element module

class pantea.atoms.element.ElementMap(unique_elements, element_to_atomic_number, element_to_atom_type, atom_type_to_element)

Bases: object

Mapping element name to atom type and more.

It assigns an atomic type to each element which allows efficient array processing (e.g. applying conditions) using the arrays of integer (i.e. atom types) instead of strings.

atom_type_to_element: Dict[int, str]

element_to_atom_type: Dict[str, ArrayImpl]

element_to_atomic_number: Dict[str, int]

classmethod from_list(elements)

Create dictionary to map elements, atom types, and atomic numbers. The atom types are sorted based on elements’ atomic number.

Return type

ElementMap

get_atom_type_from_element(name)

Map element name to atom type.

Return type

int

classmethod get_atomic_mass_from_element(name)

Return type

float

classmethod get_atomic_number_from_element(name)

Return type

int

get_element_from_atom_type(value)

Map atom type to element name.

Return type

str

classmethod get_element_from_atomic_number(value)

Return type

str

classmethod get_masses_from_structure(structure)

Get array of atomic masses.

Return type

ArrayImpl

unique_elements: Tuple[str, ...]

class pantea.atoms.element.StructureInterface(*args, **kwargs)

Bases: Protocol

atom_types: ArrayImpl

element_map: ElementMap

pantea.atoms.neighbor module

class pantea.atoms.neighbor.Neighbor(r_cutoff, masks)

Bases: BaseJaxPytreeDataClass

Finding neighboring atoms.

This is useful for efficiently determining the neighboring atoms within a specified cutoff radius. The neighbor list allows for faster calculations properties that depend on nearby atoms, such as computing forces, energies, or evaluating interatomic distances.

The current implementation relies on cutoff masks, which is different from conventional methods used to update the neighbor list (such as defining neighbor indices). The rationale behind this approach is that JAX executes efficiently on vectorized variables, offering faster performance compared to simple Python loops.

Note

For MD simulations, re-neighboring the list is required every few steps. This is usually implemented together with defining a skin radius.

classmethod from_structure(structure, r_cutoff, with_aux=False)

Return type

Union[Neighbor, Tuple[ArrayImpl, ArrayImpl]]

masks: ArrayImpl

r_cutoff: ArrayImpl

class pantea.atoms.neighbor.StructureInterface(*args, **kwargs)

Bases: Protocol

lattice: ArrayImpl

positions: ArrayImpl

pantea.atoms.structure module

class pantea.atoms.structure.Structure(positions, forces, energies, charges, total_energy, total_charge, atom_types, element_map, box=None)

Bases: BaseJaxPytreeDataClass

A Structure contains arrays that store atomic attributes for a collection of atoms in a simulation box and at any time.

The currently existing attributes of atoms within a Structure are as follows:

• positions: positions of atoms, an array of (natoms, 3)

• forces: force components, an array of (natoms, 3)

• energies: associated atom potential energies, an array of (natoms,)

• charges: charge of atoms, an array of (natoms,)

• total_energy: total potential energy, scalar value

• total_charge: total charge, scalar value

Structure serves as a fundamental data container for atoms. Multiple structures can be used to train a potential, or alternatively, the total energy and force components can be computed for a specific structure.

Each structure has additional attributes as:

• Box: applying periodic boundary condition (PBC) along x, y, and z directions

• ElementMap: determines how to extract assigned atom types from the element and vice versa

Note

The structure can be viewed as a separate domain for implementing MPI in large-scale molecular dynamics (MD) simulations, as demonstrated in the miniMD code.

add_energy_offset(atom_energy)

Add the input reference energies to individual atoms and the total energy.

Parameters

atom_energy (Dict[Element, float]]) – atom reference energy

Return type

None

as_kernel_args()

Return type

StructureAsKernelArgs

atom_types: ArrayImpl

box: Optional[Box] = None

charges: ArrayImpl

property dtype: _ScalarMeta

Return data type of the arrays in the structure (default is float64).

Return type

_ScalarMeta

element_map: ElementMap

energies: ArrayImpl

forces: ArrayImpl

classmethod from_ase(atoms, dtype=None)

Create an instance of the structure from ASE atoms input.

Parameters

• atoms (Atoms) – input ASE atoms instance

• dtype (Optional[_ScalarMeta]) – data type for arrays, defaults to None

Return type

Structure

Returns

initialized structure

classmethod from_dict(data, dtype=None)

Create a structure using an input data dictionary that contains distinct lists of positions, forces, elements, lattice, etc.

Parameters

• data (Dict[str, Any]) – input data

• dtype (Optional[_ScalarMeta]) – data type for arrays, defaults to None

Return type

Structure

Returns

the initialized Structure

get_elements()

Get array of elements.

Return type

Tuple[str, ...]

get_forces_per_element()

Get force components per element.

Return type

Dict[str, ArrayImpl]

get_positions_per_element()

Get position of atoms per element.

Return type

Dict[str, ArrayImpl]

get_unique_elements()

Return type

Tuple[str, ...]

property lattice: Optional[ArrayImpl]

Return the cell matrix (3x3).

Return type

Optional[ArrayImpl]

property natoms: int

Return number of atoms in the structure

Return type

int

positions: ArrayImpl

remove_energy_offset(atom_energy)

Remove the input reference energies from individual atoms and the total energy.

Parameters

atom_energy (Dict[Element, float]]) – atom reference energy

Return type

None

select(element)

Retrieve the indices of all atoms that correspond to the given element.

Parameters

element (str) – element name (e.g. H for hydrogen)

Return type

ArrayImpl

Returns

atom indices

to_ase()

Represent the structure as ASE atoms.

The returned object can be utilized with the ASE package for visualization or modification of the structure.

Return type

Atoms

Returns

ASE representation of the structure

to_dict()

The atomic attributes are represented as a dictionary of NumPy arrays. This format can be employed, for instance, when saving the structure data into a file.

Returns

dictionary of atom attributes.

Return type

Dict[str, np.ndarray]

total_charge: ArrayImpl

total_energy: ArrayImpl

class pantea.atoms.structure.StructureAsKernelArgs(positions: Array, atom_types: Array, lattice: Array, total_energy: Array, element_map: Dict[Element, Array])

Bases: tuple

This is a jit-complied compatible representation of structure to be used for for energy and force computing kernels.

Create new instance of StructureAsKernelArgs(positions, atom_types, lattice, total_energy, element_map)

atom_types: ArrayImpl

Alias for field number 1

element_map: Dict[str, ArrayImpl]

Alias for field number 4

lattice: ArrayImpl

Alias for field number 2

positions: ArrayImpl

Alias for field number 0

total_energy: ArrayImpl

Alias for field number 3

Module contents

class pantea.atoms.Box(lattice)

Bases: BaseJaxPytreeDataClass

Simulation box is used to apply periodic boundary conditions (PBC) in the presence of a lattice info.

Warning

The current implementation only works for orthogonal cells and does not support triclinic cells.

apply_pbc(dx)

Apply periodic boundary condition (PBC) on the atom positions.

For this method to function correctly, it is essential that all atoms are initially positioned within the boundaries of the box. Otherwise, the results may not be as anticipated. This could happen for when example time step is too large.

Parameters

dx (Array) – positional differences

Returns

PBC applied position differences

Return type

Array

property dtype: _ScalarMeta

Return type

_ScalarMeta

classmethod from_list(data, dtype=None)

Return type

Box

lattice: ArrayImpl

property length: ArrayImpl

Return length of cell in x, y, and z-directions.

Return type

ArrayImpl

property lx: ArrayImpl

Return length of cell in x-direction.

Return type

ArrayImpl

property ly: ArrayImpl

Return length of cell in y-direction.

Return type

ArrayImpl

property lz: ArrayImpl

Return length of cell in z-direction.

Return type

ArrayImpl

property volume: ArrayImpl

Return calculated volume of the box.

Return type

ArrayImpl

wrap_into_box(positions)

Adjust the coordinates of the atoms to ensure they fall within the boundaries of the simulation box that has periodic boundary conditions (PBC).

Parameters

positions (Array) – atom positions

Returns

wrapped atom positions using the PBC.

Return type

Array

class pantea.atoms.ElementMap(unique_elements, element_to_atomic_number, element_to_atom_type, atom_type_to_element)

Bases: object

Mapping element name to atom type and more.

It assigns an atomic type to each element which allows efficient array processing (e.g. applying conditions) using the arrays of integer (i.e. atom types) instead of strings.

atom_type_to_element: Dict[int, str]

element_to_atom_type: Dict[str, ArrayImpl]

element_to_atomic_number: Dict[str, int]

classmethod from_list(elements)

Create dictionary to map elements, atom types, and atomic numbers. The atom types are sorted based on elements’ atomic number.

Return type

ElementMap

get_atom_type_from_element(name)

Map element name to atom type.

Return type

int

classmethod get_atomic_mass_from_element(name)

Return type

float

classmethod get_atomic_number_from_element(name)

Return type

int

get_element_from_atom_type(value)

Map atom type to element name.

Return type

str

classmethod get_element_from_atomic_number(value)

Return type

str

classmethod get_masses_from_structure(structure)

Get array of atomic masses.

Return type

ArrayImpl

unique_elements: Tuple[str, ...]

class pantea.atoms.Neighbor(r_cutoff, masks)

Bases: BaseJaxPytreeDataClass

Finding neighboring atoms.

This is useful for efficiently determining the neighboring atoms within a specified cutoff radius. The neighbor list allows for faster calculations properties that depend on nearby atoms, such as computing forces, energies, or evaluating interatomic distances.

The current implementation relies on cutoff masks, which is different from conventional methods used to update the neighbor list (such as defining neighbor indices). The rationale behind this approach is that JAX executes efficiently on vectorized variables, offering faster performance compared to simple Python loops.

Note

For MD simulations, re-neighboring the list is required every few steps. This is usually implemented together with defining a skin radius.

classmethod from_structure(structure, r_cutoff, with_aux=False)

Return type

Union[Neighbor, Tuple[ArrayImpl, ArrayImpl]]

masks: ArrayImpl

r_cutoff: ArrayImpl

class pantea.atoms.Structure(positions, forces, energies, charges, total_energy, total_charge, atom_types, element_map, box=None)

Bases: BaseJaxPytreeDataClass

A Structure contains arrays that store atomic attributes for a collection of atoms in a simulation box and at any time.

The currently existing attributes of atoms within a Structure are as follows:

• positions: positions of atoms, an array of (natoms, 3)

• forces: force components, an array of (natoms, 3)

• energies: associated atom potential energies, an array of (natoms,)

• charges: charge of atoms, an array of (natoms,)

• total_energy: total potential energy, scalar value

• total_charge: total charge, scalar value

Structure serves as a fundamental data container for atoms. Multiple structures can be used to train a potential, or alternatively, the total energy and force components can be computed for a specific structure.

Each structure has additional attributes as:

• Box: applying periodic boundary condition (PBC) along x, y, and z directions

• ElementMap: determines how to extract assigned atom types from the element and vice versa

Note

The structure can be viewed as a separate domain for implementing MPI in large-scale molecular dynamics (MD) simulations, as demonstrated in the miniMD code.

add_energy_offset(atom_energy)

Add the input reference energies to individual atoms and the total energy.

Parameters

atom_energy (Dict[Element, float]]) – atom reference energy

Return type

None

as_kernel_args()

Return type

StructureAsKernelArgs

atom_types: ArrayImpl

box: Optional[Box] = None

charges: ArrayImpl

property dtype: _ScalarMeta

Return data type of the arrays in the structure (default is float64).

Return type

_ScalarMeta

element_map: ElementMap

energies: ArrayImpl

forces: ArrayImpl

classmethod from_ase(atoms, dtype=None)

Create an instance of the structure from ASE atoms input.

Parameters

• atoms (Atoms) – input ASE atoms instance

• dtype (Optional[_ScalarMeta]) – data type for arrays, defaults to None

Return type

Structure

Returns

initialized structure

classmethod from_dict(data, dtype=None)

Create a structure using an input data dictionary that contains distinct lists of positions, forces, elements, lattice, etc.

Parameters

• data (Dict[str, Any]) – input data

• dtype (Optional[_ScalarMeta]) – data type for arrays, defaults to None

Return type

Structure

Returns

the initialized Structure

get_elements()

Get array of elements.

Return type

Tuple[str, ...]

get_forces_per_element()

Get force components per element.

Return type

Dict[str, ArrayImpl]

get_positions_per_element()

Get position of atoms per element.

Return type

Dict[str, ArrayImpl]

get_unique_elements()

Return type

Tuple[str, ...]

property lattice: Optional[ArrayImpl]

Return the cell matrix (3x3).

Return type

Optional[ArrayImpl]

property natoms: int

Return number of atoms in the structure

Return type

int

positions: ArrayImpl

remove_energy_offset(atom_energy)

Remove the input reference energies from individual atoms and the total energy.

Parameters

atom_energy (Dict[Element, float]]) – atom reference energy

Return type

None

select(element)

Retrieve the indices of all atoms that correspond to the given element.

Parameters

element (str) – element name (e.g. H for hydrogen)

Return type

ArrayImpl

Returns

atom indices

to_ase()

Represent the structure as ASE atoms.

The returned object can be utilized with the ASE package for visualization or modification of the structure.

Return type

Atoms

Returns

ASE representation of the structure

to_dict()

The atomic attributes are represented as a dictionary of NumPy arrays. This format can be employed, for instance, when saving the structure data into a file.

Returns

dictionary of atom attributes.

Return type

Dict[str, np.ndarray]

total_charge: ArrayImpl

total_energy: ArrayImpl

pantea.atoms.calculate_distances(structure: StructureInterface, atom_index: Optional[Array] = None, neighbor_atom_index: Optional[Array] = None, with_aux: bool = False) → Union[Array, tuple[Array, Array]]

Calculate distances between specific atoms and the neighboring atoms in the structure. This function optionally also returns the corresponding position differences.

If atom indices are not specified, all atoms in the structure will be taken into account. Similarly, if neighbor indices are not provided, all neighboring atoms will be considered.

Parameters

• structure (StructureInterface) – input structure

• atom_index (Optional[Array], optional) – array of atom indices (zero-based)

• neighbor_atom_index (Optional[Array], optional) – array of indices of neighbor atoms (zero-based), defaults to None

• with_aux – whether returning position differences, defaults to False

Type

bool, optional

Returns

either an array of distances or tuple of distances together and position differences

Return type

Union[Array, tuple[Array, Array]]