UMA¶

class dipm.models.uma.models.UMA(*args: Any, **kwargs: Any)¶

The UMA model flax module. It is derived from the ForceModel class.

References

Brandon M. Wood, Misko Dzamba, Xiang Fu, Meng Gao, Muhammed Shuaibi, Luis Barroso-Luque, Kareem Abdelmaqsoud, Vahe Gharakhanyan, John R. Kitchin, Daniel S. Levine, etc. UMA: A Family of Universal Models for Atoms. arXiv, Jun 2025. URL: https://arxiv.org/abs/2506.23971.

config¶

Hyperparameters / configuration for the UMA model, see UMAConfig.

Type:: dipm.models.uma.config.UMAConfig

dataset_info¶: Hyperparameters dictated by the dataset (e.g., cutoff radius or average number of neighbors).

__call__(edge_vectors: Array, node_species: Array, senders: Array, receivers: Array, n_node: Array, rngs: Rngs | None = None, *, csd_mixed_emb: Array, ctx: dict) → Array¶: Compute node-wise energy summands. This function must be overridden by the implementation of ForceModel.

class dipm.models.uma.config.UMAConfig(*, force_head: bool = False, param_dtype: DtypeEnum = DtypeEnum.F32, task_list: list[str] | None = None, num_layers: Annotated[int, Gt(gt=0)] = 12, lmax: Annotated[int, Gt(gt=0)] = 6, mmax: Annotated[int, Ge(ge=0)] = 2, sphere_channels: Annotated[int, Gt(gt=0)] = 128, edge_channels: Annotated[int, Gt(gt=0)] = 128, hidden_channels: Annotated[int, Gt(gt=0)] = 128, num_rbf: Annotated[int, Gt(gt=0)] = 600, grid_resolution: Annotated[int, Gt(gt=0)] = None, norm_type: LayerNormType = LayerNormType.LAYER_NORM_SH, act_type: ActivationType = ActivationType.GATE, ff_type: FeedForwardType = FeedForwardType.GRID, atomic_energies: str | dict[int, float] | None = None, num_experts: int = 8, mole_dropout: float = 0.0, use_composition_embedding: bool = False)¶

The configuration / hyperparameters of the UMA model.

num_layers¶

Number of UMA layers. Default is 12.

Type:: int

lmax¶

Maximum degree of the spherical harmonics (1 to 10).

Type:: int

mmax¶

Maximum order of the spherical harmonics (0 to lmax).

Type:: int

sphere_channels¶

Number of spherical channels. Default is 128.

Type:: int

edge_channels¶

Number of channels for the edge invariant features. Default is 128.

Type:: int

hidden_channels¶

Number of hidden channels in the UMA layer. Default is 128.

Type:: int

num_rbf¶

Number of basis functions used in the embedding block. Default is 600.

Type:: int

grid_resolution¶

Resolution of SO3Grid used in Activation. Examples are 18, 16, 14, None (default, decided automatically).

Type:: int

norm_type¶

Type of normalization layer. Options are “layer_norm”, “layer_norm_sh” (default) and “rms_norm_sh”.

Type:: dipm.layers.escn.layernorm.LayerNormType

act_type¶

Type of activation function. Options are “gate” (default) and “s2_sep”.

Type:: dipm.models.uma.blocks.ActivationType

ff_type¶

Type of feed-forward function. Options are “spectral” (default) and “grid”.

Type:: dipm.models.uma.blocks.FeedForwardType

atomic_energies¶

How to treat the atomic energies. If set to None (default) or the string "average", then the average atomic energies stored in the dataset info are used. It can also be set to the string "zero" which means not to use any atomic energies in the model. Lastly, one can also pass an atomic energies dictionary via this parameter different from the one in the dataset info, that is used.

Type:: str | dict[int, float] | None

num_experts¶

Number of experts in the MoLE block. Default is 8.

Type:: int

mole_dropout¶

Dropout rate for MoLE router. Default is 0.0.

Type:: float

use_composition_embedding¶

Whether to use composition embedding in MoLE router.

Type:: bool

class dipm.models.uma.blocks.ActivationType(*values)¶: Type of attention activation.

class dipm.models.uma.blocks.FeedForwardType(*values)¶: Type of feed-forward layer.