core.models.painn#
Submodules#
Classes#
PaiNN model based on the description in Schütt et al. (2021): |
Package Contents#
- class core.models.painn.PaiNN(hidden_channels: int = 512, num_layers: int = 6, num_rbf: int = 128, cutoff: float = 12.0, max_neighbors: int = 50, rbf: dict[str, str] | None = None, envelope: dict[str, str | int] | None = None, regress_forces: bool = True, direct_forces: bool = True, use_pbc: bool = True, use_pbc_single: bool = False, otf_graph: bool = True, num_elements: int = 83, scale_file: str | None = None)#
Bases:
torch.nn.Module,fairchem.core.models.base.GraphModelMixinPaiNN model based on the description in Schütt et al. (2021): Equivariant message passing for the prediction of tensorial properties and molecular spectra, https://arxiv.org/abs/2102.03150.
- num_layers#
- num_rbf#
- cutoff#
- max_neighbors#
- regress_forces#
- direct_forces#
- otf_graph#
- use_pbc#
- use_pbc_single#
- symmetric_edge_symmetrization = False#
- atom_emb#
- radial_basis#
- message_layers#
- update_layers#
- out_energy#
- inv_sqrt_2#
- reset_parameters() None#
- select_symmetric_edges(tensor, mask, reorder_idx, inverse_neg) torch.Tensor#
- symmetrize_edges(edge_index, cell_offsets, neighbors, batch_idx, reorder_tensors, reorder_tensors_invneg)#
Symmetrize edges to ensure existence of counter-directional edges.
Some edges are only present in one direction in the data, since every atom has a maximum number of neighbors. If symmetric_edge_symmetrization is False, we only use i->j edges here. So we lose some j->i edges and add others by making it symmetric. If symmetric_edge_symmetrization is True, we always use both directions.
- generate_graph_values(data)#
- forward(data)#
- property num_params: int#
- __repr__() str#