toy_net

Toy architecture for testing purposes.

class ToyEnergy(max_number_basis: int, hidden_dimensions: list[int])

Toy energy module for testing purposes.

__init__(max_number_basis: int, hidden_dimensions: list[int]) → None

Initializes the ToyEnergy module.

Parameters:

• max_number_basis (int) – Maximal number of basis per atom

• hidden_dimensions (list[int]) – number of dimensions of the hidden layers in the mlp

forward(batch, edge_index: Tensor, padded_coeffs: Tensor) → Tensor

Calculates the forward pass of the ToyEnergy module.

Parameters:

• batch (Batch) – Batch object containing the data

• edge_index (Tensor) – edge index of the graph

• padded_coeffs (Tensor) – zero padded coefficients

Returns:

predicted energy

Return type:

Tensor

message(pos_i: Tensor, pos_j: Tensor, coeffs_j: Tensor) → Tensor

Message function to calculate the message from node j to i.

Parameters:

• pos_i (Tensor) – position of the receiving node

• pos_j (Tensor) – position of the sending node

• coeffs_j (Tensor) – coefficients of the sending node

Returns:

Message

Return type:

Tensor

class ToyInitialGuessDelta(max_number_basis: int, hidden_dimensions: list[int])

MessagePassing Module for calculating the inital_guess_delta.

__init__(max_number_basis: int, hidden_dimensions: list[int]) → None

Initializes the ToyInitialGuessDelta module.

Parameters:

• max_number_basis (int) – maximal number of basis functions per atom

• hidden_dimensions (list[int]) – dimensions of the hidden layers in the mlp

forward(batch, edge_index: Tensor, padded_coeffs: Tensor) → Tensor

Calculates the forward pass of the ToyCoeffs module.

Parameters:

• batch (Batch) – Batch object containing the data

• edge_index (Tensor) – edge index of the graph

• padded_coeffs (Tensor) – zero padded coefficients

Returns:

Coefficient inital_guess_delta

Return type:

Tensor

message(pos_i: Tensor, pos_j: Tensor, coeffs_j: Tensor) → Tensor

Message function to calculate the message from node j to i.

Parameters:

• pos_i (Tensor) – position of the receiving node

• pos_j (Tensor) – position of the sending node

• coeffs_j (Tensor) – coefficients of the sending node

Returns:

Message

Return type:

Tensor

class ToyNet(basis_info: BasisInfo, k_neighbors: int, hidden_dimension_energy: tuple[int, ...] = (128, 128, 128), hidden_dimension_coeffs: tuple[int, ...] = (128, 128, 128))

Toy net module for testing purposes.

__init__(basis_info: BasisInfo, k_neighbors: int, hidden_dimension_energy: tuple[int, ...] = (128, 128, 128), hidden_dimension_coeffs: tuple[int, ...] = (128, 128, 128)) → None

Initializes the ToyNet module.

Parameters:

• basis_info (BasisInfo) – BasisInfo object containing information about the basis for the dataset

• k_neighbors (int) – Number of neighbors to consider in the graph creation

• hidden_dimension_energy (list[int], optional) – list of dimensions of the hidden layer in the mlp calculating the energy. Defaults to [128,128,128].

• hidden_dimension_coeffs (list[int], optional) – list of dimensions of the hidden layer in the mlp calculating the coefficients. Defaults to [128,128,128].

forward(batch) → tuple[Tensor, Tensor]

Calculates the forward pass of the ToyNet module. By independently calculating the energy and the predicted delta to the ground state, to be used for the learned initial guess in OF-DFT.

Parameters:

batch (Batch) – Batch object containing the data

Returns:

The predicted energy and the predicted coefficients (in order)

Return type:

tuple[Tensor, Tensor]