core.models.gemnet_oc#
Subpackages#
Submodules#
Classes#
Package Contents#
- class core.models.gemnet_oc.GemNetOC(num_spherical: int, num_radial: int, num_blocks: int, emb_size_atom: int, emb_size_edge: int, emb_size_trip_in: int, emb_size_trip_out: int, emb_size_quad_in: int, emb_size_quad_out: int, emb_size_aint_in: int, emb_size_aint_out: int, emb_size_rbf: int, emb_size_cbf: int, emb_size_sbf: int, num_before_skip: int, num_after_skip: int, num_concat: int, num_atom: int, num_output_afteratom: int, num_atom_emb_layers: int = 0, num_global_out_layers: int = 2, regress_forces: bool = True, direct_forces: bool = False, use_pbc: bool = True, use_pbc_single: bool = False, scale_backprop_forces: bool = False, cutoff: float = 6.0, cutoff_qint: float | None = None, cutoff_aeaint: float | None = None, cutoff_aint: float | None = None, max_neighbors: int = 50, max_neighbors_qint: int | None = None, max_neighbors_aeaint: int | None = None, max_neighbors_aint: int | None = None, enforce_max_neighbors_strictly: bool = True, rbf: dict[str, str] | None = None, rbf_spherical: dict | None = None, envelope: dict[str, str | int] | None = None, cbf: dict[str, str] | None = None, sbf: dict[str, str] | None = None, extensive: bool = True, forces_coupled: bool = False, output_init: str = 'HeOrthogonal', activation: str = 'silu', quad_interaction: bool = False, atom_edge_interaction: bool = False, edge_atom_interaction: bool = False, atom_interaction: bool = False, scale_basis: bool = False, qint_tags: list | None = None, num_elements: int = 83, otf_graph: bool = False, scale_file: str | None = None, **kwargs)#
Bases:
torch.nn.Module
,fairchem.core.models.base.GraphModelMixin
- Parameters:
num_spherical (int) – Controls maximum frequency.
num_radial (int) – Controls maximum frequency.
num_blocks (int) – Number of building blocks to be stacked.
emb_size_atom (int) – Embedding size of the atoms.
emb_size_edge (int) – Embedding size of the edges.
emb_size_trip_in (int) – (Down-projected) embedding size of the quadruplet edge embeddings before the bilinear layer.
emb_size_trip_out (int) – (Down-projected) embedding size of the quadruplet edge embeddings after the bilinear layer.
emb_size_quad_in (int) – (Down-projected) embedding size of the quadruplet edge embeddings before the bilinear layer.
emb_size_quad_out (int) – (Down-projected) embedding size of the quadruplet edge embeddings after the bilinear layer.
emb_size_aint_in (int) – Embedding size in the atom interaction before the bilinear layer.
emb_size_aint_out (int) – Embedding size in the atom interaction after the bilinear layer.
emb_size_rbf (int) – Embedding size of the radial basis transformation.
emb_size_cbf (int) – Embedding size of the circular basis transformation (one angle).
emb_size_sbf (int) – Embedding size of the spherical basis transformation (two angles).
num_before_skip (int) – Number of residual blocks before the first skip connection.
num_after_skip (int) – Number of residual blocks after the first skip connection.
num_concat (int) – Number of residual blocks after the concatenation.
num_atom (int) – Number of residual blocks in the atom embedding blocks.
num_output_afteratom (int) – Number of residual blocks in the output blocks after adding the atom embedding.
num_atom_emb_layers (int) – Number of residual blocks for transforming atom embeddings.
num_global_out_layers (int) – Number of final residual blocks before the output.
regress_forces (bool) – Whether to predict forces. Default: True
direct_forces (bool) – If True predict forces based on aggregation of interatomic directions. If False predict forces based on negative gradient of energy potential.
use_pbc (bool) – Whether to use periodic boundary conditions.
use_pbc_single – Process batch PBC graphs one at a time
scale_backprop_forces (bool) – Whether to scale up the energy and then scales down the forces to prevent NaNs and infs in backpropagated forces.
cutoff (float) – Embedding cutoff for interatomic connections and embeddings in Angstrom.
cutoff_qint (float) – Quadruplet interaction cutoff in Angstrom. Optional. Uses cutoff per default.
cutoff_aeaint (float) – Edge-to-atom and atom-to-edge interaction cutoff in Angstrom. Optional. Uses cutoff per default.
cutoff_aint (float) – Atom-to-atom interaction cutoff in Angstrom. Optional. Uses maximum of all other cutoffs per default.
max_neighbors (int) – Maximum number of neighbors for interatomic connections and embeddings.
max_neighbors_qint (int) – Maximum number of quadruplet interactions per embedding. Optional. Uses max_neighbors per default.
max_neighbors_aeaint (int) – Maximum number of edge-to-atom and atom-to-edge interactions per embedding. Optional. Uses max_neighbors per default.
max_neighbors_aint (int) – Maximum number of atom-to-atom interactions per atom. Optional. Uses maximum of all other neighbors per default.
enforce_max_neighbors_strictly (bool) – When subselected edges based on max_neighbors args, arbitrarily select amongst degenerate edges to have exactly the correct number.
rbf (dict) – Name and hyperparameters of the radial basis function.
rbf_spherical (dict) – Name and hyperparameters of the radial basis function used as part of the circular and spherical bases. Optional. Uses rbf per default.
envelope (dict) – Name and hyperparameters of the envelope function.
cbf (dict) – Name and hyperparameters of the circular basis function.
sbf (dict) – Name and hyperparameters of the spherical basis function.
extensive (bool) – Whether the output should be extensive (proportional to the number of atoms)
forces_coupled (bool) – If True, enforce that |F_st| = |F_ts|. No effect if direct_forces is False.
output_init (str) – Initialization method for the final dense layer.
activation (str) – Name of the activation function.
scale_file (str) – Path to the pytorch file containing the scaling factors.
quad_interaction (bool) – Whether to use quadruplet interactions (with dihedral angles)
atom_edge_interaction (bool) – Whether to use atom-to-edge interactions
edge_atom_interaction (bool) – Whether to use edge-to-atom interactions
atom_interaction (bool) – Whether to use atom-to-atom interactions
scale_basis (bool) – Whether to use a scaling layer in the raw basis function for better numerical stability.
qint_tags (list) – Which atom tags to use quadruplet interactions for. 0=sub-surface bulk, 1=surface, 2=adsorbate atoms.
- num_blocks#
- extensive#
- activation#
- atom_edge_interaction#
- edge_atom_interaction#
- atom_interaction#
- quad_interaction#
- qint_tags#
- otf_graph#
- enforce_max_neighbors_strictly#
- use_pbc#
- use_pbc_single#
- direct_forces#
- forces_coupled#
- regress_forces#
- force_scaler#
- atom_emb#
- edge_emb#
- int_blocks = []#
- out_blocks = []#
- out_mlp_E#
- out_energy#
- out_initializer#
- set_cutoffs(cutoff, cutoff_qint, cutoff_aeaint, cutoff_aint)#
- set_max_neighbors(max_neighbors, max_neighbors_qint, max_neighbors_aeaint, max_neighbors_aint)#
- init_basis_functions(num_radial, num_spherical, rbf, rbf_spherical, envelope, cbf, sbf, scale_basis)#
- calculate_quad_angles(V_st, V_qint_st, quad_idx)#
Calculate angles for quadruplet-based message passing.
- Parameters:
V_st (Tensor, shape = (nAtoms, 3)) – Normalized directions from s to t
V_qint_st (Tensor, shape = (nAtoms, 3)) – Normalized directions from s to t for the quadruplet interaction graph
quad_idx (dict of torch.Tensor) – Indices relevant for quadruplet interactions.
- Returns:
cosφ_cab (Tensor, shape = (num_triplets_inint,)) – Cosine of angle between atoms c -> a <- b.
cosφ_abd (Tensor, shape = (num_triplets_qint,)) – Cosine of angle between atoms a -> b -> d.
angle_cabd (Tensor, shape = (num_quadruplets,)) – Dihedral angle between atoms c <- a-b -> d.
- select_symmetric_edges(tensor: torch.Tensor, mask: torch.Tensor, reorder_idx: torch.Tensor, opposite_neg) torch.Tensor #
Use a mask to remove values of removed edges and then duplicate the values for the correct edge direction.
- Parameters:
tensor (torch.Tensor) – Values to symmetrize for the new tensor.
mask (torch.Tensor) – Mask defining which edges go in the correct direction.
reorder_idx (torch.Tensor) – Indices defining how to reorder the tensor values after concatenating the edge values of both directions.
opposite_neg (bool) – Whether the edge in the opposite direction should use the negative tensor value.
- Returns:
tensor_ordered – A tensor with symmetrized values.
- Return type:
torch.Tensor
- symmetrize_edges(graph, batch_idx)#
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. We only use i->j edges here. So we lose some j->i edges and add others by making it symmetric.
- subselect_edges(data, graph, cutoff=None, max_neighbors=None)#
Subselect edges using a stricter cutoff and max_neighbors.
- generate_graph_dict(data, cutoff, max_neighbors)#
Generate a radius/nearest neighbor graph.
- subselect_graph(data, graph, cutoff, max_neighbors, cutoff_orig, max_neighbors_orig)#
If the new cutoff and max_neighbors is different from the original, subselect the edges of a given graph.
- get_graphs_and_indices(data)#
“Generate embedding and interaction graphs and indices.
- get_bases(main_graph, a2a_graph, a2ee2a_graph, qint_graph, trip_idx_e2e, trip_idx_a2e, trip_idx_e2a, quad_idx, num_atoms)#
Calculate and transform basis functions.
- forward(data)#
- property num_params: int#