Anisotropic PIPs and Models

Functions

get_mask	Generates a mask matrix identifying unique atom pairs from a list of atom types.
get_f_mask	Creates a JIT-compiled function to compute anisotropic masked distances.
lambda_random_init	Randomizes the length scale parameters in a PIP model parameter tree.

Classes

PIPAniso	Permutationally Invariant Polynomials layer (Anisotropic).
LayerPIPAniso	Energy model combining LayerPIPAniso and a dense output layer.
EnergyPIPAniso	Energy model using PIP.

class molpipx.pip_anisotropic_flax.EnergyPIPAniso(*args, **kwargs)

Bases: Module

Energy model using PIP.

Variables:

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

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

• f_mask (Callable) – callable function that returns the masked distances

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

• n_pairs (int) – Number of unique atom pairs.

__call__(inputs)

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

Parameters:

• inputs – geometries, (batch,number of atoms, 3)

• mask – array with the mask for each atom-atom type length scale parameter

Returns:

energy of each geometry, (batch, 1)

Return type:

energy

class molpipx.pip_anisotropic_flax.LayerPIPAniso(*args, **kwargs)

Bases: Module

Energy model combining LayerPIPAniso and a dense output layer.

Variables:

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

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

• f_mask (Callable) – Function that returns the masked distances.

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

• n_pairs (int) – Number of unique atom pairs.

__call__(inputs)

Applies a vectorize map to PIP.apply.

Parameters:

• inputs – geometries, (batch,number of atoms, 3)

• mask – mask parameters for each atom-atom type length scale parameter

Returns:

pip vector for each geometry, (batch, number of pips)

Return type:

pip_vector

class molpipx.pip_anisotropic_flax.PIPAniso(*args, **kwargs)

Bases: Module

Permutationally Invariant Polynomials layer (Anisotropic).

Variables:

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

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

• f_mask (Callable) – Function that returns the masked distances.

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

• bias_init (Callable) – Initializer function for the l parameter.

• n_pairs (int) – Number of unique atom pairs.

__call__(input)

Applies the PIP transformation to the inputs. Anisotropic version, where we considered different atom-atom type length scale parameter (warning: this function only works for a single input)

Parameters:

• inputs – geometry, array of number of atoms by 3 (Na x 3)

• mask – array with the mask for each atom-atom type length scale parameter

Returns:

pip vector

Return type:

pip_vector

molpipx.pip_anisotropic_flax.get_f_mask(mask: jaxtyping.Array) → Callable

Creates a JIT-compiled function to compute anisotropic masked distances.

Parameters:

mask (Array) – The binary mask array generated by get_mask.

Returns:

A function f_mask(l, inputs) that returns weighted distances.

Return type:

Callable

molpipx.pip_anisotropic_flax.get_mask(atom_types: list) → Tuple

Generates a mask matrix identifying unique atom pairs from a list of atom types.

Parameters:

atom_types (list) – A list of atomic symbols (e.g., ['H', 'H', 'O']).

Returns:

A tuple containing the mask array and a list of unique pair strings.

Return type:

Tuple

molpipx.pip_anisotropic_flax.lambda_random_init(params_pip: Any, key: jaxtyping.Array) → Any

Randomizes the length scale parameters in a PIP model parameter tree.

Parameters:

• params_pip (Any) – The Flax parameter PyTree.

• key (KeyArray) – A JAX random key.

Returns:

The updated parameter PyTree with randomized lambda values.

Return type:

Any