mattertune.configs.backbones.jmp.model

class mattertune.configs.backbones.jmp.model.CutoffsConfig(*, main, aeaint, qint, aint)[source]
Parameters:

  • main (float)

  • aeaint (float)

  • qint (float)

  • aint (float)

main: float
aeaint: float
qint: float
aint: float
classmethod from_constant(value)[source]
Parameters:

value (float)

class mattertune.configs.backbones.jmp.model.FinetuneModuleBaseConfig(*, properties, optimizer, lr_scheduler=None, ignore_gpu_batch_transform_error=True, normalizers={})[source]
Parameters:

  • properties (Sequence[PropertyConfig])

  • optimizer (OptimizerConfig)

  • lr_scheduler (LRSchedulerConfig | None)

  • ignore_gpu_batch_transform_error (bool)

  • normalizers (Mapping[str, Sequence[NormalizerConfig]])

properties: Sequence[PropertyConfig]

Properties to predict.

optimizer: OptimizerConfig

Optimizer.

lr_scheduler: LRSchedulerConfig | None

Learning Rate Scheduler

ignore_gpu_batch_transform_error: bool

Whether to ignore data processing errors during training.

normalizers: Mapping[str, Sequence[NormalizerConfig]]

Normalizers for the properties.

Any property can be associated with multiple normalizers. This is useful for cases where we want to normalize the same property in different ways. For example, we may want to normalize the energy by subtracting the atomic reference energies, as well as by mean and standard deviation normalization.

The normalizers are applied in the order they are defined in the list.

abstract classmethod ensure_dependencies()[source]

Ensure that all dependencies are installed.

This method should raise an exception if any dependencies are missing, with a message indicating which dependencies are missing and how to install them.

abstract create_model()[source]

Creates an instance of the finetune module for this configuration.

Return type:

FinetuneModuleBase

class mattertune.configs.backbones.jmp.model.JMPBackboneConfig(*, properties, optimizer, lr_scheduler=None, ignore_gpu_batch_transform_error=True, normalizers={}, name='jmp', ckpt_path, graph_computer)[source]
Parameters:

  • properties (Sequence[PropertyConfig])

  • optimizer (OptimizerConfig)

  • lr_scheduler (LRSchedulerConfig | None)

  • ignore_gpu_batch_transform_error (bool)

  • normalizers (Mapping[str, Sequence[NormalizerConfig]])

  • name (Literal['jmp'])

  • ckpt_path (Path | CachedPath)

  • graph_computer (JMPGraphComputerConfig)

name: Literal['jmp']

The type of the backbone.

ckpt_path: Path | CE.CachedPath

The path to the pre-trained model checkpoint.

graph_computer: JMPGraphComputerConfig

The configuration for the graph computer.

create_model()[source]

Creates an instance of the finetune module for this configuration.

classmethod ensure_dependencies()[source]

Ensure that all dependencies are installed.

This method should raise an exception if any dependencies are missing, with a message indicating which dependencies are missing and how to install them.

properties: Sequence[PropertyConfig]

Properties to predict.

optimizer: OptimizerConfig

Optimizer.

lr_scheduler: LRSchedulerConfig | None

Learning Rate Scheduler

ignore_gpu_batch_transform_error: bool

Whether to ignore data processing errors during training.

normalizers: Mapping[str, Sequence[NormalizerConfig]]

Normalizers for the properties.

Any property can be associated with multiple normalizers. This is useful for cases where we want to normalize the same property in different ways. For example, we may want to normalize the energy by subtracting the atomic reference energies, as well as by mean and standard deviation normalization.

The normalizers are applied in the order they are defined in the list.

class mattertune.configs.backbones.jmp.model.JMPGraphComputerConfig(*, pbc, cutoffs=CutoffsConfig(main=12.0, aeaint=12.0, qint=12.0, aint=12.0), max_neighbors=MaxNeighborsConfig(main=30, aeaint=20, qint=8, aint=1000), per_graph_radius_graph=False)[source]
Parameters:
pbc: bool

Whether to use periodic boundary conditions.

cutoffs: CutoffsConfig

The cutoff for the radius graph.

max_neighbors: MaxNeighborsConfig

The maximum number of neighbors for the radius graph.

per_graph_radius_graph: bool

Whether to compute the radius graph per graph.

class mattertune.configs.backbones.jmp.model.MaxNeighborsConfig(*, main, aeaint, qint, aint)[source]
Parameters:

  • main (int)

  • aeaint (int)

  • qint (int)

  • aint (int)

main: int
aeaint: int
qint: int
aint: int
classmethod from_goc_base_proportions(max_neighbors)[source]
GOC base proportions:

max_neighbors: 30 max_neighbors_qint: 8 max_neighbors_aeaint: 20 max_neighbors_aint: 1000

Parameters:

max_neighbors (int)