ramannoodle.pmodel.torch

Modules for polarizability models implemented with PyTorch.

class ramannoodle.pmodel.torch.PotGNN(ref_structure, cutoff, size_node_embedding, size_edge_embedding, num_message_passes, gaussian_filter_start, gaussian_filter_end, mean_polarizability, stddev_polarizability)

Bases: Module, PolarizabilityModel

POlarizability Tensor Graph Neural Network (PotGNN).

The architecture was inspired by the “direct force architecture” developed in Park et al.; npj Computational Materials (2021)7:73; doi:10.1038/s41524-021-00543-3. Implementation adapted from torch_geometric.nn.models.GNNFF authored by @ken2403 and merged by @rusty1s.

The architecture of this model is still somewhat in flux. More complete documentation for this model, including a description of the architecture and discussion of design choices, will be available at a later date.

Parameters:

  • ref_structure (ReferenceStructure) – Reference structure from which nodes/edges are determined.

  • cutoff (float) – (Å) Cutoff distance for edges.

  • size_node_embedding (int)

  • size_edge_embedding (int)

  • num_message_passes (int)

  • gaussian_filter_start (float) – (Å) Lower bound of the Gaussian filter used in initial edge embedding.

  • gaussian_filter_end (float) – (Å) Upper bound of the Gaussian filter used in initial edge embedding.

  • mean_polarizability (ndarray[Any, dtype[float64]]) – Array with shape (3,3).

  • stddev_polarizability (ndarray[Any, dtype[float64]]) – Array with shape (3,3).

calc_polarizabilities(positions_batch)

Return estimated polarizabilities for a batch of fractional positions.

Parameters:

positions_batch (ndarray[Any, dtype[float64]]) – (fractional) Array with shape (S,N,3) where S is the number of samples and N is the number of atoms.

Return type:

ndarray[Any, dtype[float64]]

Returns:

Array with shape (S,3,3).

forward(lattice, atomic_numbers, positions)

Forward pass.

Parameters:

  • lattice (Tensor) – (Å) Tensor with size [S,3,3] where S is the number of samples.

  • atomic_numbers (Tensor) – Tensor with size [S,N] where N is the number of atoms.

  • positions (Tensor) – (fractional) Tensor with size [S,N,3].

Return type:

Tensor

Returns:

Polarizability vectors with size [S,6]. To convert into tensor form, see polarizability_vectors_to_tensors().

reset_parameters()

Reset model parameters.

Return type:

None

ramannoodle.pmodel.torch.train_single_epoch(model, training_set, validation_set, batch_size, optimizer, loss_function)

Train PotGNN model for a single epoch on the default device.

Parameters:
Return type:

tuple[float, float, ndarray[Any, dtype[float64]]]

Returns:

  1. mean training loss

  2. mean validation loss

  3. mean variance of predictions on validation set – Array with shape [6,]