Core PIPs

Classes

PIP	Permutationally Invariant Polynomials layer.
PIPlayer	Vectorized wrapper for PIP to handle batches of geometries.
EnergyPIP	End-to-end energy model combining PIPlayer and a linear output layer.
PIPlayerGP	Wrapper for PIPlayer to reshape the inputs before passing to PIPlayer.

class molpipx.pip_flax.EnergyPIP(*args, **kwargs)

Bases: Module

End-to-end energy model combining PIPlayer and a linear output layer.

Variables:

• f_mono (Callable) – Function that returns the monomials.

• f_poly (Callable) – Function that returns the polynomials.

• l (float) – Initial value of the Morse variables length scale parameter.

• trainable_l (bool) – If True, the length scale l is optimized.

__call__(inputs)

Applies the PIPLayer and a nn.Dense modules to compute the energy.

Parameters:

inputs (Array) – Batch of geometries with shape (Batch, Number of atoms, 3).

Returns:

Energy values for each geometry, shape (Batch, 1).

Return type:

Array

class molpipx.pip_flax.PIP(*args, **kwargs)

Bases: Module

Permutationally Invariant Polynomials layer.

Variables:

• f_mono (Callable) – callable function that returns the monomials

• f_poly (Callable) – callable function that returns the polynomials

• l (float) – initial value of the morse variables length scale parameter

• bias_init (Callable) – initializer function for the the l parameter

• trainable_l (bool) – If True, the length scale l is optimized during training.

__call__(input)

Applies the PIP transformation to a single input geometry.

Warning

This function works on a single geometry (Na x 3). Use PIPlayer for batches.

Parameters:

input (Array) – Geometry array of shape (Number of atoms, 3).

Returns:

The PIP vector.

Return type:

Array

class molpipx.pip_flax.PIPlayer(*args, **kwargs)

Bases: Module

Vectorized wrapper for PIP to handle batches of geometries.

See PIP for more information.

Variables:

• f_mono (Callable) – Function that returns the monomials.

• f_poly (Callable) – Function that returns the polynomials.

• l (float) – Initial value of the Morse variables length scale parameter.

• trainable_l (bool) – If True, the length scale l is optimized.

__call__(inputs)

Applies a vectorized map of PIP to a batch of inputs.

Parameters:

inputs (Array) – Batch of geometries with shape (Batch, Number of atoms, 3).

Returns:

PIP vectors for each geometry, shape (Batch, Number of PIPs).

Return type:

Array

class molpipx.pip_flax.PIPlayerGP(*args, **kwargs)

Bases: Module

Wrapper for PIPlayer to reshape the inputs before passing to PIPlayer. :ivar f_mono: Function that returns the monomials. :vartype f_mono: Callable :ivar f_poly: Function that returns the polynomials. :vartype f_poly: Callable :ivar l: Initial value of the Morse variables length scale parameter. :vartype l: float :ivar trainable_l: If True, the length scale l is optimized.

__call__(inputs)

Reshapes flattened inputs and applies the PIP transformation.

Parameters:

inputs (Array) – Geometries. Can be (Batch, Na, 3) or flattened (Batch * Na * 3).

Returns:

PIP vectors for each geometry, shape (Batch, Number of PIPs).

Return type:

Array