initial_guess_delta_module

Module for the projMINAO guess from [M-OFDFT].

class InitialGuessDeltaModule(input_size: int, basis_info: ~mldft.ml.data.components.basis_info.BasisInfo, hidden_layers: ~typing.List[int], activation_function: ~torch.nn.modules.module.Module | None = <class 'torch.nn.modules.activation.SiLU'>, dataset_statistics: ~mldft.ml.preprocess.dataset_statistics.DatasetStatistics | None = None, weigher_key: str = 'initial_guess_only', dropout: float = 0.0)[source]

Module to calculate the initial guess delta based on the node features computed by previous layers.

__init__(input_size: int, basis_info: ~mldft.ml.data.components.basis_info.BasisInfo, hidden_layers: ~typing.List[int], activation_function: ~torch.nn.modules.module.Module | None = <class 'torch.nn.modules.activation.SiLU'>, dataset_statistics: ~mldft.ml.preprocess.dataset_statistics.DatasetStatistics | None = None, weigher_key: str = 'initial_guess_only', dropout: float = 0.0) → None[source]

Initializes the class. calculates the atom_ptr object, which is needed for reverting the atom hot encoding. Also initializes the mlp with the given parameters.

Parameters:

  • input_size (int) – Input size of the mlp

  • basis_info (BasisInfo) – BasisInfo object

  • activation_function (Module, optional) – Activation function for the MLP. Defaults to SiLU.

  • hidden_layers (Tuple[int, ], optional) – Configurations of the hidden layers in the MLP. Defaults to (32, 32, 32).

  • dataset_statistics (DatasetStatistics, optional) – DatasetStatistics object. Defaults to None. If not None, this will be used to offset and scale the predicted delta according to per-coefficient mean and standard deviation as given in the DatasetStatistics object.

  • weigher_key (str) – Selects the sample weigher for the dataset statistics.

  • dropout (float, optional) – Dropout probability of the MLP. Defaults to 0.0.

forward(x: Tensor, basis_function_ind: Tensor, coeff_ind_to_node_ind: Tensor, batch: Tensor = None) → Tensor[source]

Calculates the forward pass of the module. The coefficients are passed through an MLP and then the atom-hot-encoding is reverted.

Parameters:

  • x (Tensor) – Activations of the previous layer. Shape: (n_atoms, input_size).

  • basis_function_ind (Tensor) – Indices of the basis functions in the molecule. Shape: (n_atoms,).

  • coeff_ind_to_node_ind (Tensor[int]) – Tensor mapping coefficient indices to node (=atom) indices.

  • batch (Tensor, optional) – Batch tensor for LayerNorm inside the mlp.

Returns:

the initial_guess_delta in one vector. Shape: (n_basis,)

Return type:

Tensor