PIP Neural Networks

Classes

MLP	A standard Multi-Layer Perceptron (MLP) module.
PIPNN	Neural Network model that uses PIPs as input features.

class molpipx.pipnn_flax.MLP(*args, **kwargs)

Bases: Module

A standard Multi-Layer Perceptron (MLP) module.

Variables:

• features (Tuple[int]) – A tuple defining the number of neurons in each hidden layer.

• act_fun (Callable) – The activation function applied after each hidden layer. Defaults to nn.tanh.

__call__(x)

Applies the MLP to the input.

Parameters:

x (Array) – Input tensor.

Returns:

The output of the network (scalar per batch item).

Return type:

Array

setup()

Initializes the dense layers based on the features tuple.

class molpipx.pipnn_flax.PIPNN(*args, **kwargs)

Bases: Module

Neural Network model that uses PIPs as input features.

Variables:

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

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

• features (Tuple[int]) – A tuple defining the number of neurons in each hidden layer of the MLP.

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

• act_fun (Callable) – The activation function used in the MLP. Defaults to nn.tanh.

__call__(x)

Computes the energy for the input geometries.

Parameters:

x (Array) – Batch of geometries (Batch, Na, 3).

Returns:

Predicted energy values (Batch, 1).

Return type:

Array

setup()

Initializes the PIP layer and the MLP layers.