core.models.gemnet_oc

core.models.gemnet_oc#

Subpackages#

Submodules#

Classes#

GemNetOC

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)#
init_shared_basis_layers(num_radial, num_spherical, emb_size_rbf, emb_size_cbf, emb_size_sbf)#
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#
Contents