diff --git a/_downloads/5fdddbed2260616231dbf7b0d94bb665/train.txt b/_downloads/5fdddbed2260616231dbf7b0d94bb665/train.txt index ac8753ffee..ae1027207b 100644 --- a/_downloads/5fdddbed2260616231dbf7b0d94bb665/train.txt +++ b/_downloads/5fdddbed2260616231dbf7b0d94bb665/train.txt @@ -1,6 +1,6 @@ -2024-11-19 06:16:16 (INFO): Running in local mode without elastic launch (single gpu only) -2024-11-19 06:16:16 (INFO): Setting env PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True -2024-11-19 06:16:17 (INFO): Project root: /home/runner/work/fairchem/fairchem/src/fairchem +2024-12-04 00:11:01 (INFO): Running in local mode without elastic launch (single gpu only) +2024-12-04 00:11:01 (INFO): Setting env PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True +2024-12-04 00:11:01 (INFO): Project root: /home/runner/work/fairchem/fairchem/src/fairchem /home/runner/work/fairchem/fairchem/src/fairchem/core/models/escn/so3.py:23: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature. _Jd = torch.load(os.path.join(os.path.dirname(__file__), "Jd.pt")) /home/runner/work/fairchem/fairchem/src/fairchem/core/models/scn/spherical_harmonics.py:23: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature. @@ -15,17 +15,17 @@ @torch.cuda.amp.autocast(enabled=False) /home/runner/work/fairchem/fairchem/src/fairchem/core/models/equiformer_v2/layer_norm.py:357: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead. @torch.cuda.amp.autocast(enabled=False) -2024-11-19 06:16:17 (INFO): amp: false +2024-12-04 00:11:02 (INFO): amp: false cmd: - checkpoint_dir: fine-tuning/checkpoints/2024-11-19-06-15-28-ft-oxides - commit: aa298ac + checkpoint_dir: fine-tuning/checkpoints/2024-12-04-00-10-40-ft-oxides + commit: 3f695ef identifier: ft-oxides - logs_dir: fine-tuning/logs/tensorboard/2024-11-19-06-15-28-ft-oxides + logs_dir: fine-tuning/logs/tensorboard/2024-12-04-00-10-40-ft-oxides print_every: 10 - results_dir: fine-tuning/results/2024-11-19-06-15-28-ft-oxides + results_dir: fine-tuning/results/2024-12-04-00-10-40-ft-oxides seed: 0 - timestamp_id: 2024-11-19-06-15-28-ft-oxides - version: 0.1.dev1+gaa298ac + timestamp_id: 2024-12-04-00-10-40-ft-oxides + version: 1.3.0 dataset: a2g_args: r_energy: true @@ -155,70 +155,70 @@ val_dataset: format: ase_db src: val.db -2024-11-19 06:16:17 (INFO): Loading model: gemnet_oc -2024-11-19 06:16:17 (WARNING): Unrecognized arguments: ['symmetric_edge_symmetrization'] -2024-11-19 06:16:20 (INFO): Loaded GemNetOC with 38864438 parameters. -2024-11-19 06:16:20 (WARNING): log_summary for Tensorboard not supported -2024-11-19 06:16:20 (INFO): Loading dataset: ase_db -2024-11-19 06:16:20 (WARNING): Could not find dataset metadata.npz files in '[PosixPath('train.db')]' -2024-11-19 06:16:20 (WARNING): Disabled BalancedBatchSampler because num_replicas=1. -2024-11-19 06:16:20 (WARNING): Failed to get data sizes, falling back to uniform partitioning. BalancedBatchSampler requires a dataset that has a metadata attributed with number of atoms. -2024-11-19 06:16:20 (INFO): rank: 0: Sampler created... -2024-11-19 06:16:20 (INFO): Created BalancedBatchSampler with sampler=, batch_size=4, drop_last=False -2024-11-19 06:16:20 (WARNING): Could not find dataset metadata.npz files in '[PosixPath('val.db')]' -2024-11-19 06:16:20 (WARNING): Disabled BalancedBatchSampler because num_replicas=1. -2024-11-19 06:16:20 (WARNING): Failed to get data sizes, falling back to uniform partitioning. BalancedBatchSampler requires a dataset that has a metadata attributed with number of atoms. -2024-11-19 06:16:20 (INFO): rank: 0: Sampler created... -2024-11-19 06:16:20 (INFO): Created BalancedBatchSampler with sampler=, batch_size=16, drop_last=False -2024-11-19 06:16:20 (WARNING): Could not find dataset metadata.npz files in '[PosixPath('test.db')]' -2024-11-19 06:16:20 (WARNING): Disabled BalancedBatchSampler because num_replicas=1. -2024-11-19 06:16:20 (WARNING): Failed to get data sizes, falling back to uniform partitioning. BalancedBatchSampler requires a dataset that has a metadata attributed with number of atoms. -2024-11-19 06:16:20 (INFO): rank: 0: Sampler created... -2024-11-19 06:16:20 (INFO): Created BalancedBatchSampler with sampler=, batch_size=16, drop_last=False -2024-11-19 06:16:20 (WARNING): Using `weight_decay` from `optim` instead of `optim.optimizer_params`.Please update your config to use `optim.optimizer_params.weight_decay`.`optim.weight_decay` will soon be deprecated. -2024-11-19 06:16:20 (INFO): Attemping to load user specified checkpoint at /tmp/fairchem_checkpoints/gnoc_oc22_oc20_all_s2ef.pt -2024-11-19 06:16:20 (INFO): Loading checkpoint from: /tmp/fairchem_checkpoints/gnoc_oc22_oc20_all_s2ef.pt -/home/runner/work/fairchem/fairchem/src/fairchem/core/trainers/base_trainer.py:602: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature. +2024-12-04 00:11:02 (INFO): Loading model: gemnet_oc +2024-12-04 00:11:02 (WARNING): Unrecognized arguments: ['symmetric_edge_symmetrization'] +2024-12-04 00:11:05 (INFO): Loaded GemNetOC with 38864438 parameters. +2024-12-04 00:11:05 (WARNING): log_summary for Tensorboard not supported +2024-12-04 00:11:05 (INFO): Loading dataset: ase_db +2024-12-04 00:11:05 (WARNING): Could not find dataset metadata.npz files in '[PosixPath('train.db')]' +2024-12-04 00:11:05 (WARNING): Disabled BalancedBatchSampler because num_replicas=1. +2024-12-04 00:11:05 (WARNING): Failed to get data sizes, falling back to uniform partitioning. BalancedBatchSampler requires a dataset that has a metadata attributed with number of atoms. +2024-12-04 00:11:05 (INFO): rank: 0: Sampler created... +2024-12-04 00:11:05 (INFO): Created BalancedBatchSampler with sampler=, batch_size=4, drop_last=False +2024-12-04 00:11:05 (WARNING): Could not find dataset metadata.npz files in '[PosixPath('val.db')]' +2024-12-04 00:11:05 (WARNING): Disabled BalancedBatchSampler because num_replicas=1. +2024-12-04 00:11:05 (WARNING): Failed to get data sizes, falling back to uniform partitioning. BalancedBatchSampler requires a dataset that has a metadata attributed with number of atoms. +2024-12-04 00:11:05 (INFO): rank: 0: Sampler created... +2024-12-04 00:11:05 (INFO): Created BalancedBatchSampler with sampler=, batch_size=16, drop_last=False +2024-12-04 00:11:05 (WARNING): Could not find dataset metadata.npz files in '[PosixPath('test.db')]' +2024-12-04 00:11:05 (WARNING): Disabled BalancedBatchSampler because num_replicas=1. +2024-12-04 00:11:05 (WARNING): Failed to get data sizes, falling back to uniform partitioning. BalancedBatchSampler requires a dataset that has a metadata attributed with number of atoms. +2024-12-04 00:11:05 (INFO): rank: 0: Sampler created... +2024-12-04 00:11:05 (INFO): Created BalancedBatchSampler with sampler=, batch_size=16, drop_last=False +2024-12-04 00:11:05 (WARNING): Using `weight_decay` from `optim` instead of `optim.optimizer_params`.Please update your config to use `optim.optimizer_params.weight_decay`.`optim.weight_decay` will soon be deprecated. +2024-12-04 00:11:05 (INFO): Attemping to load user specified checkpoint at /tmp/fairchem_checkpoints/gnoc_oc22_oc20_all_s2ef.pt +2024-12-04 00:11:05 (INFO): Loading checkpoint from: /tmp/fairchem_checkpoints/gnoc_oc22_oc20_all_s2ef.pt +/home/runner/work/fairchem/fairchem/src/fairchem/core/trainers/base_trainer.py:604: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature. checkpoint = torch.load(checkpoint_path, map_location=map_location) -2024-11-19 06:16:20 (INFO): Overwriting scaling factors with those loaded from checkpoint. If you're generating predictions with a pretrained checkpoint, this is the correct behavior. To disable this, delete `scale_dict` from the checkpoint. -/home/runner/work/fairchem/fairchem/src/fairchem/core/trainers/ocp_trainer.py:155: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead. +2024-12-04 00:11:05 (INFO): Overwriting scaling factors with those loaded from checkpoint. If you're generating predictions with a pretrained checkpoint, this is the correct behavior. To disable this, delete `scale_dict` from the checkpoint. +/home/runner/work/fairchem/fairchem/src/fairchem/core/trainers/ocp_trainer.py:164: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead. with torch.cuda.amp.autocast(enabled=self.scaler is not None): /home/runner/work/fairchem/fairchem/src/fairchem/core/models/gemnet_oc/gemnet_oc.py:1270: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead. with torch.cuda.amp.autocast(False): -2024-11-19 06:16:44 (INFO): energy_mae: 9.48e+00, forcesx_mae: 7.26e-02, forcesy_mae: 3.95e-02, forcesz_mae: 5.74e-02, forces_mae: 5.65e-02, forces_cosine_similarity: 1.12e-01, forces_magnitude_error: 1.11e-01, energy_forces_within_threshold: 0.00e+00, loss: 9.61e+00, lr: 5.00e-04, epoch: 1.69e-01, step: 1.00e+01 -2024-11-19 06:16:45 (INFO): Evaluating on val. - device 0: 0%| | 0/2 [00:00, batch_size=16, drop_last=False -2024-11-19 06:22:13 (INFO): Attemping to load user specified checkpoint at /tmp/fairchem_checkpoints/gndt_oc22_all_s2ef.pt -2024-11-19 06:22:13 (INFO): Loading checkpoint from: /tmp/fairchem_checkpoints/gndt_oc22_all_s2ef.pt -/home/runner/work/fairchem/fairchem/src/fairchem/core/trainers/base_trainer.py:602: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature. +2024-12-04 00:17:01 (INFO): Loading model: gemnet_t +2024-12-04 00:17:02 (INFO): Loaded GemNetT with 31671825 parameters. +2024-12-04 00:17:02 (WARNING): log_summary for Tensorboard not supported +2024-12-04 00:17:02 (WARNING): Could not find dataset metadata.npz files in '[PosixPath('data.db')]' +2024-12-04 00:17:02 (WARNING): Disabled BalancedBatchSampler because num_replicas=1. +2024-12-04 00:17:02 (WARNING): Failed to get data sizes, falling back to uniform partitioning. BalancedBatchSampler requires a dataset that has a metadata attributed with number of atoms. +2024-12-04 00:17:02 (INFO): rank: 0: Sampler created... +2024-12-04 00:17:02 (INFO): Created BalancedBatchSampler with sampler=, batch_size=16, drop_last=False +2024-12-04 00:17:03 (INFO): Attemping to load user specified checkpoint at /tmp/fairchem_checkpoints/gndt_oc22_all_s2ef.pt +2024-12-04 00:17:03 (INFO): Loading checkpoint from: /tmp/fairchem_checkpoints/gndt_oc22_all_s2ef.pt +/home/runner/work/fairchem/fairchem/src/fairchem/core/trainers/base_trainer.py:604: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature. checkpoint = torch.load(checkpoint_path, map_location=map_location) -2024-11-19 06:22:14 (INFO): Overwriting scaling factors with those loaded from checkpoint. If you're generating predictions with a pretrained checkpoint, this is the correct behavior. To disable this, delete `scale_dict` from the checkpoint. -2024-11-19 06:22:14 (WARNING): Scale factor comment not found in model -2024-11-19 06:22:14 (INFO): Predicting on test. - device 0: 0%| | 0/3 [00:00 None + .. py:method:: __call__(dict_config: omegaconf.DictConfig) -> None - .. py:method:: checkpoint(*args, **kwargs) + .. py:method:: _init_logger() -> None + + + .. py:method:: checkpoint(*args, **kwargs) -> submitit.helpers.DelayedSubmission Resubmits the same callable with the same arguments -.. py:function:: map_cli_args_to_dist_config(cli_args: argparse.Namespace) -> dict +.. py:function:: map_cli_args_to_dist_config(cli_args: omegaconf.DictConfig) -> dict .. py:function:: get_hydra_config_from_yaml(config_yml: str, overrides_args: list[str]) -> omegaconf.DictConfig -.. py:function:: runner_wrapper(config: omegaconf.DictConfig, cli_args: argparse.Namespace) +.. py:function:: runner_wrapper(config: omegaconf.DictConfig) .. py:function:: main(args: argparse.Namespace | None = None, override_args: list[str] | None = None) diff --git a/_sources/autoapi/core/common/relaxation/ase_utils/index.rst b/_sources/autoapi/core/common/relaxation/ase_utils/index.rst index 66c7451869..9ce117929d 100644 --- a/_sources/autoapi/core/common/relaxation/ase_utils/index.rst +++ b/_sources/autoapi/core/common/relaxation/ase_utils/index.rst @@ -17,6 +17,14 @@ core.common.relaxation.ase_utils +Attributes +---------- + +.. autoapisummary:: + + core.common.relaxation.ase_utils.ASE_PROP_RESHAPE + + Classes ------- @@ -36,31 +44,30 @@ Functions Module Contents --------------- -.. py:function:: batch_to_atoms(batch) +.. py:data:: ASE_PROP_RESHAPE -.. py:class:: OCPCalculator(config_yml: str | None = None, checkpoint_path: str | None = None, model_name: str | None = None, local_cache: str | None = None, trainer: str | None = None, cutoff: int = 6, max_neighbors: int = 50, cpu: bool = True, seed: int | None = None) +.. py:function:: batch_to_atoms(batch: torch_geometric.data.Batch, results: dict[str, torch.Tensor] | None = None, wrap_pos: bool = True, eps: float = 1e-07) -> list[ase.Atoms] - Bases: :py:obj:`ase.calculators.calculator.Calculator` + Convert a data batch to ase Atoms + :param batch: data batch + :param results: dictionary with predicted result tensors that will be added to a SinglePointCalculator. If no results + are given no calculator will be added to the atoms objects. + :param wrap_pos: wrap positions back into the cell. + :param eps: Small number to prevent slightly negative coordinates from being wrapped. - Base-class for all ASE calculators. + :returns: list of Atoms - A calculator must raise PropertyNotImplementedError if asked for a - property that it can't calculate. So, if calculation of the - stress tensor has not been implemented, get_stress(atoms) should - raise PropertyNotImplementedError. This can be achieved simply by not - including the string 'stress' in the list implemented_properties - which is a class member. These are the names of the standard - properties: 'energy', 'forces', 'stress', 'dipole', 'charges', - 'magmom' and 'magmoms'. +.. py:class:: OCPCalculator(config_yml: str | None = None, checkpoint_path: str | pathlib.Path | None = None, model_name: str | None = None, local_cache: str | None = None, trainer: str | None = None, cpu: bool = True, seed: int | None = None) - .. py:attribute:: implemented_properties - :type: ClassVar[list[str]] - :value: ['energy', 'forces'] + Bases: :py:obj:`ase.calculators.calculator.Calculator` - Properties calculator can handle (energy, forces, ...) + ASE based calculator using an OCP model + + + .. py:attribute:: _reshaped_props .. py:attribute:: config @@ -72,44 +79,25 @@ Module Contents .. py:attribute:: a2g + .. py:attribute:: implemented_properties + + Properties calculator can handle (energy, forces, ...) + + .. py:method:: load_checkpoint(checkpoint_path: str, checkpoint: dict | None = None) -> None Load existing trained model :param checkpoint_path: string Path to trained model + :param checkpoint: dict + A pretrained checkpoint dict - .. py:method:: calculate(atoms: ase.Atoms, properties, system_changes) -> None - - Do the calculation. - - properties: list of str - List of what needs to be calculated. Can be any combination - of 'energy', 'forces', 'stress', 'dipole', 'charges', 'magmom' - and 'magmoms'. - system_changes: list of str - List of what has changed since last calculation. Can be - any combination of these six: 'positions', 'numbers', 'cell', - 'pbc', 'initial_charges' and 'initial_magmoms'. - - Subclasses need to implement this, but can ignore properties - and system_changes if they want. Calculated properties should - be inserted into results dictionary like shown in this dummy - example:: - - self.results = {'energy': 0.0, - 'forces': np.zeros((len(atoms), 3)), - 'stress': np.zeros(6), - 'dipole': np.zeros(3), - 'charges': np.zeros(len(atoms)), - 'magmom': 0.0, - 'magmoms': np.zeros(len(atoms))} + .. py:method:: calculate(atoms: ase.Atoms | torch_geometric.data.Batch, properties, system_changes) -> None - The subclass implementation should first call this - implementation to set the atoms attribute and create any missing - directories. + Calculate implemented properties for a single Atoms object or a Batch of them. diff --git a/_sources/autoapi/core/common/relaxation/index.rst b/_sources/autoapi/core/common/relaxation/index.rst index 062af74e4f..af38070aad 100644 --- a/_sources/autoapi/core/common/relaxation/index.rst +++ b/_sources/autoapi/core/common/relaxation/index.rst @@ -3,6 +3,14 @@ core.common.relaxation .. py:module:: core.common.relaxation +.. autoapi-nested-parse:: + + Copyright (c) Meta, Inc. and its affiliates. + + This source code is licensed under the MIT license found in the + LICENSE file in the root directory of this source tree. + + Submodules ---------- @@ -12,6 +20,290 @@ Submodules /autoapi/core/common/relaxation/ase_utils/index /autoapi/core/common/relaxation/ml_relaxation/index + /autoapi/core/common/relaxation/optimizable/index /autoapi/core/common/relaxation/optimizers/index +Classes +------- + +.. autoapisummary:: + + core.common.relaxation.OptimizableBatch + core.common.relaxation.OptimizableUnitCellBatch + + +Functions +--------- + +.. autoapisummary:: + + core.common.relaxation.ml_relax + + +Package Contents +---------------- + +.. py:function:: ml_relax(batch: torch_geometric.data.Batch, model: fairchem.core.trainers.BaseTrainer, steps: int, fmax: float, relax_opt: dict[str] | None = None, relax_cell: bool = False, relax_volume: bool = False, save_full_traj: bool = True, transform: torch.nn.Module | None = None, mask_converged: bool = True) + + Runs ML-based relaxations. + + :param batch: a data batch object. + :param model: a trainer object with model. + :param steps: Max number of steps in the structure relaxation. + :param fmax: Structure relaxation terminates when the max force of the system is no bigger than fmax. + :param relax_opt: Optimizer parameters to be used for structure relaxations. + :param relax_cell: if true will use stress predictions to relax crystallographic cell. + The model given must predict stress + :param relax_volume: if true will relax the cell isotropically. the given model must predict stress. + :param save_full_traj: Whether to save out the full ASE trajectory. If False, only save out initial and final frames. + :param mask_converged: whether to mask batches where all atoms are below convergence threshold + :param cumulative_mask: if true, once system is masked then it remains masked even if new predictions give forces + above threshold, ie. once masked always masked. Note if this is used make sure to check convergence with + the same fmax always + + +.. py:class:: OptimizableBatch(batch: torch_geometric.data.Batch, trainer: fairchem.core.trainers.BaseTrainer, transform: torch.nn.Module | None = None, mask_converged: bool = True, numpy: bool = False, masked_eps: float = 1e-08) + + Bases: :py:obj:`ase.optimize.optimize.Optimizable` + + + A Batch version of ase Optimizable Atoms + + This class can be used with ML relaxations in fairchem.core.relaxations.ml_relaxation + or in ase relaxations classes, i.e. ase.optimize.lbfgs + + + .. py:attribute:: ignored_changes + :type: ClassVar[set[str]] + + + .. py:attribute:: batch + + + .. py:attribute:: trainer + + + .. py:attribute:: transform + + + .. py:attribute:: numpy + + + .. py:attribute:: mask_converged + + + .. py:attribute:: _cached_batch + :value: None + + + + .. py:attribute:: _update_mask + :value: None + + + + .. py:attribute:: torch_results + + + .. py:attribute:: results + + + .. py:attribute:: _eps + + + .. py:attribute:: otf_graph + :value: True + + + + .. py:property:: device + + + .. py:property:: batch_indices + + Get the batch indices specifying which position/force corresponds to which batch. + + + .. py:property:: converged_mask + + + .. py:property:: update_mask + + + .. py:method:: check_state(batch: torch_geometric.data.Batch, tol: float = 1e-12) -> bool + + Check for any system changes since last calculation. + + + + .. py:method:: _predict() -> None + + Run prediction if batch has any changes. + + + + .. py:method:: get_property(name, no_numpy: bool = False) -> torch.Tensor | numpy.typing.NDArray + + Get a predicted property by name. + + + + .. py:method:: get_positions() -> torch.Tensor | numpy.typing.NDArray + + Get the batch positions + + + + .. py:method:: set_positions(positions: torch.Tensor | numpy.typing.NDArray) -> None + + Set the atom positions in the batch. + + + + .. py:method:: get_forces(apply_constraint: bool = False, no_numpy: bool = False) -> torch.Tensor | numpy.typing.NDArray + + Get predicted batch forces. + + + + .. py:method:: get_potential_energy(**kwargs) -> torch.Tensor | numpy.typing.NDArray + + Get predicted energy as the sum of all batch energies. + + + + .. py:method:: get_potential_energies() -> torch.Tensor | numpy.typing.NDArray + + Get the predicted energy for each system in batch. + + + + .. py:method:: get_cells() -> torch.Tensor + + Get batch crystallographic cells. + + + + .. py:method:: set_cells(cells: torch.Tensor | numpy.typing.NDArray) -> None + + Set batch cells. + + + + .. py:method:: get_volumes() -> torch.Tensor + + Get a tensor of volumes for each cell in batch + + + + .. py:method:: iterimages() -> torch_geometric.data.Batch + + + .. py:method:: get_max_forces(forces: torch.Tensor | None = None, apply_constraint: bool = False) -> torch.Tensor + + Get the maximum forces per structure in batch + + + + .. py:method:: converged(forces: torch.Tensor | numpy.typing.NDArray | None, fmax: float, max_forces: torch.Tensor | None = None) -> bool + + Check if norm of all predicted forces are below fmax + + + + .. py:method:: get_atoms_list() -> list[ase.Atoms] + + Get ase Atoms objects corresponding to the batch + + + + .. py:method:: update_graph() + + Update the graph if model does not use otf_graph. + + + + .. py:method:: __len__() -> int + + +.. py:class:: OptimizableUnitCellBatch(batch: torch_geometric.data.Batch, trainer: fairchem.core.trainers.BaseTrainer, transform: torch.nn.Module | None = None, numpy: bool = False, mask_converged: bool = True, mask: collections.abc.Sequence[bool] | None = None, cell_factor: float | torch.Tensor | None = None, hydrostatic_strain: bool = False, constant_volume: bool = False, scalar_pressure: float = 0.0, masked_eps: float = 1e-08) + + Bases: :py:obj:`OptimizableBatch` + + + Modify the supercell and the atom positions in relaxations. + + Based on ase UnitCellFilter to work on data batches + + + .. py:attribute:: orig_cells + + + .. py:attribute:: stress + :value: None + + + + .. py:attribute:: hydrostatic_strain + + + .. py:attribute:: constant_volume + + + .. py:attribute:: pressure + + + .. py:attribute:: cell_factor + + + .. py:attribute:: _batch_trace + + + .. py:attribute:: _batch_diag + + + .. py:property:: batch_indices + + Get the batch indices specifying which position/force corresponds to which batch. + + We augment this to specify the batch indices for augmented positions and forces. + + + .. py:method:: deform_grad() + + Get the cell deformation matrix + + + + .. py:method:: get_positions() + + Get positions and cell deformation gradient. + + + + .. py:method:: set_positions(positions: torch.Tensor | numpy.typing.NDArray) + + Set positions and cell. + + positions has shape (natoms + ncells * 3, 3). + the first natoms rows are the positions of the atoms, the last nsystems * three rows are the deformation tensor + for each cell. + + + + .. py:method:: get_potential_energy(**kwargs) + + returns potential energy including enthalpy PV term. + + + + .. py:method:: get_forces(apply_constraint: bool = False, no_numpy: bool = False) -> torch.Tensor | numpy.typing.NDArray + + Get forces and unit cell stress. + + + + .. py:method:: __len__() + + diff --git a/_sources/autoapi/core/common/relaxation/ml_relaxation/index.rst b/_sources/autoapi/core/common/relaxation/ml_relaxation/index.rst index 6596764e47..a402d55df7 100644 --- a/_sources/autoapi/core/common/relaxation/ml_relaxation/index.rst +++ b/_sources/autoapi/core/common/relaxation/ml_relaxation/index.rst @@ -23,19 +23,22 @@ Functions Module Contents --------------- -.. py:function:: ml_relax(batch, model, steps: int, fmax: float, relax_opt, save_full_traj, device: str = 'cuda:0', transform=None, early_stop_batch: bool = False) +.. py:function:: ml_relax(batch: torch_geometric.data.Batch, model: fairchem.core.trainers.BaseTrainer, steps: int, fmax: float, relax_opt: dict[str] | None = None, relax_cell: bool = False, relax_volume: bool = False, save_full_traj: bool = True, transform: torch.nn.Module | None = None, mask_converged: bool = True) Runs ML-based relaxations. - :param batch: object - :param model: object - :param steps: int - Max number of steps in the structure relaxation. - :param fmax: float - Structure relaxation terminates when the max force - of the system is no bigger than fmax. - :param relax_opt: str - Optimizer and corresponding parameters to be used for structure relaxations. - :param save_full_traj: bool - Whether to save out the full ASE trajectory. If False, only save out initial and final frames. + + :param batch: a data batch object. + :param model: a trainer object with model. + :param steps: Max number of steps in the structure relaxation. + :param fmax: Structure relaxation terminates when the max force of the system is no bigger than fmax. + :param relax_opt: Optimizer parameters to be used for structure relaxations. + :param relax_cell: if true will use stress predictions to relax crystallographic cell. + The model given must predict stress + :param relax_volume: if true will relax the cell isotropically. the given model must predict stress. + :param save_full_traj: Whether to save out the full ASE trajectory. If False, only save out initial and final frames. + :param mask_converged: whether to mask batches where all atoms are below convergence threshold + :param cumulative_mask: if true, once system is masked then it remains masked even if new predictions give forces + above threshold, ie. once masked always masked. Note if this is used make sure to check convergence with + the same fmax always diff --git a/_sources/autoapi/core/common/relaxation/optimizable/index.rst b/_sources/autoapi/core/common/relaxation/optimizable/index.rst new file mode 100644 index 0000000000..c4765ded5e --- /dev/null +++ b/_sources/autoapi/core/common/relaxation/optimizable/index.rst @@ -0,0 +1,306 @@ +core.common.relaxation.optimizable +================================== + +.. py:module:: core.common.relaxation.optimizable + +.. autoapi-nested-parse:: + + Copyright (c) Meta, Inc. and its affiliates. + + This source code is licensed under the MIT license found in the + LICENSE file in the root directory of this source tree. + + Code based on ase.optimize + + + +Attributes +---------- + +.. autoapisummary:: + + core.common.relaxation.optimizable.ALL_CHANGES + + +Classes +------- + +.. autoapisummary:: + + core.common.relaxation.optimizable.Optimizable + core.common.relaxation.optimizable.OptimizableBatch + core.common.relaxation.optimizable.OptimizableUnitCellBatch + + +Functions +--------- + +.. autoapisummary:: + + core.common.relaxation.optimizable.compare_batches + + +Module Contents +--------------- + +.. py:class:: Optimizable + +.. py:data:: ALL_CHANGES + :type: set[str] + +.. py:function:: compare_batches(batch1: torch_geometric.data.Batch | None, batch2: torch_geometric.data.Batch, tol: float = 1e-06, excluded_properties: set[str] | None = None) -> list[str] + + Compare properties between two batches + + :param batch1: atoms batch + :param batch2: atoms batch + :param tol: tolerance used to compare equility of floating point properties + :param excluded_properties: list of properties to exclude from comparison + + :returns: list of system changes, property names that are differente between batch1 and batch2 + + +.. py:class:: OptimizableBatch(batch: torch_geometric.data.Batch, trainer: fairchem.core.trainers.BaseTrainer, transform: torch.nn.Module | None = None, mask_converged: bool = True, numpy: bool = False, masked_eps: float = 1e-08) + + Bases: :py:obj:`ase.optimize.optimize.Optimizable` + + + A Batch version of ase Optimizable Atoms + + This class can be used with ML relaxations in fairchem.core.relaxations.ml_relaxation + or in ase relaxations classes, i.e. ase.optimize.lbfgs + + + .. py:attribute:: ignored_changes + :type: ClassVar[set[str]] + + + .. py:attribute:: batch + + + .. py:attribute:: trainer + + + .. py:attribute:: transform + + + .. py:attribute:: numpy + + + .. py:attribute:: mask_converged + + + .. py:attribute:: _cached_batch + :value: None + + + + .. py:attribute:: _update_mask + :value: None + + + + .. py:attribute:: torch_results + + + .. py:attribute:: results + + + .. py:attribute:: _eps + + + .. py:attribute:: otf_graph + :value: True + + + + .. py:property:: device + + + .. py:property:: batch_indices + + Get the batch indices specifying which position/force corresponds to which batch. + + + .. py:property:: converged_mask + + + .. py:property:: update_mask + + + .. py:method:: check_state(batch: torch_geometric.data.Batch, tol: float = 1e-12) -> bool + + Check for any system changes since last calculation. + + + + .. py:method:: _predict() -> None + + Run prediction if batch has any changes. + + + + .. py:method:: get_property(name, no_numpy: bool = False) -> torch.Tensor | numpy.typing.NDArray + + Get a predicted property by name. + + + + .. py:method:: get_positions() -> torch.Tensor | numpy.typing.NDArray + + Get the batch positions + + + + .. py:method:: set_positions(positions: torch.Tensor | numpy.typing.NDArray) -> None + + Set the atom positions in the batch. + + + + .. py:method:: get_forces(apply_constraint: bool = False, no_numpy: bool = False) -> torch.Tensor | numpy.typing.NDArray + + Get predicted batch forces. + + + + .. py:method:: get_potential_energy(**kwargs) -> torch.Tensor | numpy.typing.NDArray + + Get predicted energy as the sum of all batch energies. + + + + .. py:method:: get_potential_energies() -> torch.Tensor | numpy.typing.NDArray + + Get the predicted energy for each system in batch. + + + + .. py:method:: get_cells() -> torch.Tensor + + Get batch crystallographic cells. + + + + .. py:method:: set_cells(cells: torch.Tensor | numpy.typing.NDArray) -> None + + Set batch cells. + + + + .. py:method:: get_volumes() -> torch.Tensor + + Get a tensor of volumes for each cell in batch + + + + .. py:method:: iterimages() -> torch_geometric.data.Batch + + + .. py:method:: get_max_forces(forces: torch.Tensor | None = None, apply_constraint: bool = False) -> torch.Tensor + + Get the maximum forces per structure in batch + + + + .. py:method:: converged(forces: torch.Tensor | numpy.typing.NDArray | None, fmax: float, max_forces: torch.Tensor | None = None) -> bool + + Check if norm of all predicted forces are below fmax + + + + .. py:method:: get_atoms_list() -> list[ase.Atoms] + + Get ase Atoms objects corresponding to the batch + + + + .. py:method:: update_graph() + + Update the graph if model does not use otf_graph. + + + + .. py:method:: __len__() -> int + + +.. py:class:: OptimizableUnitCellBatch(batch: torch_geometric.data.Batch, trainer: fairchem.core.trainers.BaseTrainer, transform: torch.nn.Module | None = None, numpy: bool = False, mask_converged: bool = True, mask: collections.abc.Sequence[bool] | None = None, cell_factor: float | torch.Tensor | None = None, hydrostatic_strain: bool = False, constant_volume: bool = False, scalar_pressure: float = 0.0, masked_eps: float = 1e-08) + + Bases: :py:obj:`OptimizableBatch` + + + Modify the supercell and the atom positions in relaxations. + + Based on ase UnitCellFilter to work on data batches + + + .. py:attribute:: orig_cells + + + .. py:attribute:: stress + :value: None + + + + .. py:attribute:: hydrostatic_strain + + + .. py:attribute:: constant_volume + + + .. py:attribute:: pressure + + + .. py:attribute:: cell_factor + + + .. py:attribute:: _batch_trace + + + .. py:attribute:: _batch_diag + + + .. py:property:: batch_indices + + Get the batch indices specifying which position/force corresponds to which batch. + + We augment this to specify the batch indices for augmented positions and forces. + + + .. py:method:: deform_grad() + + Get the cell deformation matrix + + + + .. py:method:: get_positions() + + Get positions and cell deformation gradient. + + + + .. py:method:: set_positions(positions: torch.Tensor | numpy.typing.NDArray) + + Set positions and cell. + + positions has shape (natoms + ncells * 3, 3). + the first natoms rows are the positions of the atoms, the last nsystems * three rows are the deformation tensor + for each cell. + + + + .. py:method:: get_potential_energy(**kwargs) + + returns potential energy including enthalpy PV term. + + + + .. py:method:: get_forces(apply_constraint: bool = False, no_numpy: bool = False) -> torch.Tensor | numpy.typing.NDArray + + Get forces and unit cell stress. + + + + .. py:method:: __len__() + + diff --git a/_sources/autoapi/core/common/relaxation/optimizers/index.rst b/_sources/autoapi/core/common/relaxation/optimizers/index.rst index 7d2b58ba15..24ae214571 100644 --- a/_sources/autoapi/core/common/relaxation/optimizers/index.rst +++ b/_sources/autoapi/core/common/relaxation/optimizers/index.rst @@ -3,6 +3,14 @@ core.common.relaxation.optimizers .. py:module:: core.common.relaxation.optimizers +.. autoapi-nested-parse:: + + Copyright (c) Meta, Inc. and its affiliates. + + This source code is licensed under the MIT license found in the + LICENSE file in the root directory of this source tree. + + Submodules ---------- @@ -13,3 +21,95 @@ Submodules /autoapi/core/common/relaxation/optimizers/lbfgs_torch/index +Classes +------- + +.. autoapisummary:: + + core.common.relaxation.optimizers.LBFGS + + +Package Contents +---------------- + +.. py:class:: LBFGS(optimizable_batch: core.common.relaxation.optimizers.optimizable.OptimizableBatch, maxstep: float = 0.02, memory: int = 100, damping: float = 1.2, alpha: float = 100.0, save_full_traj: bool = True, traj_dir: pathlib.Path | None = None, traj_names: list[str] | None = None) + + Limited memory BFGS optimizer for batch ML relaxations. + + + .. py:attribute:: optimizable + + + .. py:attribute:: maxstep + + + .. py:attribute:: memory + + + .. py:attribute:: damping + + + .. py:attribute:: alpha + + + .. py:attribute:: H0 + + + .. py:attribute:: save_full + + + .. py:attribute:: traj_dir + + + .. py:attribute:: traj_names + + + .. py:attribute:: trajectories + :value: None + + + + .. py:attribute:: fmax + :value: None + + + + .. py:attribute:: steps + :value: None + + + + .. py:attribute:: s + + + .. py:attribute:: y + + + .. py:attribute:: rho + + + .. py:attribute:: r0 + :value: None + + + + .. py:attribute:: f0 + :value: None + + + + .. py:method:: run(fmax, steps) + + + .. py:method:: determine_step(dr) + + + .. py:method:: _batched_dot(x: torch.Tensor, y: torch.Tensor) + + + .. py:method:: step(iteration: int) -> None + + + .. py:method:: write() -> None + + diff --git a/_sources/autoapi/core/common/relaxation/optimizers/lbfgs_torch/index.rst b/_sources/autoapi/core/common/relaxation/optimizers/lbfgs_torch/index.rst index 7cf51017c3..3f3414a018 100644 --- a/_sources/autoapi/core/common/relaxation/optimizers/lbfgs_torch/index.rst +++ b/_sources/autoapi/core/common/relaxation/optimizers/lbfgs_torch/index.rst @@ -18,18 +18,17 @@ Classes .. autoapisummary:: core.common.relaxation.optimizers.lbfgs_torch.LBFGS - core.common.relaxation.optimizers.lbfgs_torch.TorchCalc Module Contents --------------- -.. py:class:: LBFGS(batch: torch_geometric.data.Batch, model: TorchCalc, maxstep: float = 0.01, memory: int = 100, damping: float = 0.25, alpha: float = 100.0, force_consistent=None, device: str = 'cuda:0', save_full_traj: bool = True, traj_dir: pathlib.Path | None = None, traj_names=None, early_stop_batch: bool = False) +.. py:class:: LBFGS(optimizable_batch: core.common.relaxation.optimizers.optimizable.OptimizableBatch, maxstep: float = 0.02, memory: int = 100, damping: float = 1.2, alpha: float = 100.0, save_full_traj: bool = True, traj_dir: pathlib.Path | None = None, traj_names: list[str] | None = None) - .. py:attribute:: batch + Limited memory BFGS optimizer for batch ML relaxations. - .. py:attribute:: model + .. py:attribute:: optimizable .. py:attribute:: maxstep @@ -47,58 +46,61 @@ Module Contents .. py:attribute:: H0 - .. py:attribute:: force_consistent + .. py:attribute:: save_full - .. py:attribute:: device + .. py:attribute:: traj_dir - .. py:attribute:: save_full + .. py:attribute:: traj_names - .. py:attribute:: traj_dir + .. py:attribute:: trajectories + :value: None - .. py:attribute:: traj_names + .. py:attribute:: fmax + :value: None - .. py:attribute:: early_stop_batch - .. py:attribute:: otf_graph - :value: True + .. py:attribute:: steps + :value: None - .. py:method:: get_energy_and_forces(apply_constraint: bool = True) + .. py:attribute:: s - .. py:method:: set_positions(update, update_mask) -> None + .. py:attribute:: y - .. py:method:: check_convergence(iteration, forces=None, energy=None) + .. py:attribute:: rho - .. py:method:: run(fmax, steps) + .. py:attribute:: r0 + :value: None - .. py:method:: step(iteration: int, forces: torch.Tensor | None, update_mask: torch.Tensor) -> None + .. py:attribute:: f0 + :value: None - .. py:method:: write(energy, forces, update_mask) -> None -.. py:class:: TorchCalc(model, transform=None) + .. py:method:: run(fmax, steps) + - .. py:attribute:: model + .. py:method:: determine_step(dr) - .. py:attribute:: transform + .. py:method:: _batched_dot(x: torch.Tensor, y: torch.Tensor) - .. py:method:: get_energy_and_forces(atoms, apply_constraint: bool = True) + .. py:method:: step(iteration: int) -> None - .. py:method:: update_graph(atoms) + .. py:method:: write() -> None diff --git a/_sources/autoapi/core/datasets/ase_datasets/index.rst b/_sources/autoapi/core/datasets/ase_datasets/index.rst index d269cc4d00..4f83fee43a 100644 --- a/_sources/autoapi/core/datasets/ase_datasets/index.rst +++ b/_sources/autoapi/core/datasets/ase_datasets/index.rst @@ -331,8 +331,3 @@ Module Contents .. py:method:: sample_property_metadata(num_samples: int = 100) -> dict - .. py:method:: get_relaxed_energy(identifier) - :abstractmethod: - - - diff --git a/_sources/autoapi/core/datasets/index.rst b/_sources/autoapi/core/datasets/index.rst index b8f2ee82de..6a7a010dd6 100644 --- a/_sources/autoapi/core/datasets/index.rst +++ b/_sources/autoapi/core/datasets/index.rst @@ -129,11 +129,6 @@ Package Contents .. py:method:: sample_property_metadata(num_samples: int = 100) -> dict - .. py:method:: get_relaxed_energy(identifier) - :abstractmethod: - - - .. py:class:: AseReadDataset(config: dict, atoms_transform: Callable[[ase.Atoms, Any, Ellipsis], ase.Atoms] = apply_one_tags) Bases: :py:obj:`AseAtomsDataset` diff --git a/_sources/autoapi/core/models/equiformer_v2/equiformer_v2_dens/index.rst b/_sources/autoapi/core/models/equiformer_v2/equiformer_v2_dens/index.rst new file mode 100644 index 0000000000..8c5cbcea13 --- /dev/null +++ b/_sources/autoapi/core/models/equiformer_v2/equiformer_v2_dens/index.rst @@ -0,0 +1,219 @@ +core.models.equiformer_v2.equiformer_v2_dens +============================================ + +.. py:module:: core.models.equiformer_v2.equiformer_v2_dens + +.. autoapi-nested-parse:: + + Copyright (c) Meta, Inc. and its affiliates. + + This source code is licensed under the MIT license found in the + LICENSE file in the root directory of this source tree. + + + +Classes +------- + +.. autoapisummary:: + + core.models.equiformer_v2.equiformer_v2_dens.EqV2DeNSBackbone + core.models.equiformer_v2.equiformer_v2_dens.DeNSScalarHead + core.models.equiformer_v2.equiformer_v2_dens.DeNSVectorHead + core.models.equiformer_v2.equiformer_v2_dens.DeNSRank2Head + + +Module Contents +--------------- + +.. py:class:: EqV2DeNSBackbone(use_pbc: bool = True, use_pbc_single: bool = False, regress_forces: bool = True, otf_graph: bool = True, max_neighbors: int = 500, max_radius: float = 5.0, max_num_elements: int = 90, num_layers: int = 12, sphere_channels: int = 128, attn_hidden_channels: int = 128, num_heads: int = 8, attn_alpha_channels: int = 32, attn_value_channels: int = 16, ffn_hidden_channels: int = 512, norm_type: str = 'rms_norm_sh', lmax_list: list[int] | None = None, mmax_list: list[int] | None = None, grid_resolution: int | None = None, num_sphere_samples: int = 128, edge_channels: int = 128, use_atom_edge_embedding: bool = True, share_atom_edge_embedding: bool = False, use_m_share_rad: bool = False, distance_function: str = 'gaussian', num_distance_basis: int = 512, attn_activation: str = 'scaled_silu', use_s2_act_attn: bool = False, use_attn_renorm: bool = True, ffn_activation: str = 'scaled_silu', use_gate_act: bool = False, use_grid_mlp: bool = False, use_sep_s2_act: bool = True, alpha_drop: float = 0.1, drop_path_rate: float = 0.05, proj_drop: float = 0.0, weight_init: str = 'normal', enforce_max_neighbors_strictly: bool = True, avg_num_nodes: float | None = None, avg_degree: float | None = None, use_energy_lin_ref: bool | None = False, load_energy_lin_ref: bool | None = False, activation_checkpoint: bool | None = False, use_force_encoding=True, use_noise_schedule_sigma_encoding: bool = False) + + Bases: :py:obj:`fairchem.core.models.equiformer_v2.equiformer_v2.EquiformerV2Backbone` + + + DeNS extra Args: + use_force_encoding (bool): For ablation study, whether to encode forces during denoising positions. Default: True. + use_noise_schedule_sigma_encoding (bool): For ablation study, whether to encode the sigma (sampled std of Gaussian noises) during + denoising positions when `fixed_noise_std` = False in config files. Default: False. + + + .. py:attribute:: use_force_encoding + + + .. py:attribute:: use_noise_schedule_sigma_encoding + + + .. py:attribute:: irreps_sh + + + .. py:attribute:: force_embedding + + + .. py:method:: forward(data) -> dict[str, torch.Tensor] + + +.. py:class:: DeNSScalarHead(backbone: fairchem.core.models.base.BackboneInterface, output_name: str = 'energy', reduce: Literal['sum', 'mean'] = 'sum', use_denoising: bool = True) + + Bases: :py:obj:`torch.nn.Module`, :py:obj:`fairchem.core.models.base.HeadInterface` + + + Base class for all neural network modules. + + Your models should also subclass this class. + + Modules can also contain other Modules, allowing to nest them in + a tree structure. You can assign the submodules as regular attributes:: + + import torch.nn as nn + import torch.nn.functional as F + + class Model(nn.Module): + def __init__(self): + super().__init__() + self.conv1 = nn.Conv2d(1, 20, 5) + self.conv2 = nn.Conv2d(20, 20, 5) + + def forward(self, x): + x = F.relu(self.conv1(x)) + return F.relu(self.conv2(x)) + + Submodules assigned in this way will be registered, and will have their + parameters converted too when you call :meth:`to`, etc. + + .. note:: + As per the example above, an ``__init__()`` call to the parent class + must be made before assignment on the child. + + :ivar training: Boolean represents whether this module is in training or + evaluation mode. + :vartype training: bool + + + .. py:attribute:: reduce + + + .. py:attribute:: avg_num_nodes + + + .. py:attribute:: scalar_block + + + .. py:attribute:: output_name + + + .. py:attribute:: use_denoising + + + .. py:method:: forward(data: torch_geometric.data.batch.Batch, emb: dict[str, torch.Tensor | fairchem.core.models.base.GraphData]) -> dict[str, torch.Tensor] + + Head forward. + + :param data: Atomic systems as input + :type data: DataBatch + :param emb: Embeddings of the input as generated by the backbone + :type emb: dict[str->torch.Tensor] + + :returns: **outputs** -- Return one or more targets generated by this head + :rtype: dict[str->torch.Tensor] + + + +.. py:class:: DeNSVectorHead(backbone: fairchem.core.models.base.BackboneInterface, output_name: str = 'forces') + + Bases: :py:obj:`torch.nn.Module`, :py:obj:`fairchem.core.models.base.HeadInterface` + + + Base class for all neural network modules. + + Your models should also subclass this class. + + Modules can also contain other Modules, allowing to nest them in + a tree structure. You can assign the submodules as regular attributes:: + + import torch.nn as nn + import torch.nn.functional as F + + class Model(nn.Module): + def __init__(self): + super().__init__() + self.conv1 = nn.Conv2d(1, 20, 5) + self.conv2 = nn.Conv2d(20, 20, 5) + + def forward(self, x): + x = F.relu(self.conv1(x)) + return F.relu(self.conv2(x)) + + Submodules assigned in this way will be registered, and will have their + parameters converted too when you call :meth:`to`, etc. + + .. note:: + As per the example above, an ``__init__()`` call to the parent class + must be made before assignment on the child. + + :ivar training: Boolean represents whether this module is in training or + evaluation mode. + :vartype training: bool + + + .. py:attribute:: output_name + + + .. py:attribute:: activation_checkpoint + + + .. py:attribute:: vector_block + + + .. py:attribute:: denoising_pos_block + + + .. py:method:: forward(data: torch_geometric.data.batch.Batch, emb: dict[str, torch.Tensor]) -> dict[str, torch.Tensor] + + Head forward. + + :param data: Atomic systems as input + :type data: DataBatch + :param emb: Embeddings of the input as generated by the backbone + :type emb: dict[str->torch.Tensor] + + :returns: **outputs** -- Return one or more targets generated by this head + :rtype: dict[str->torch.Tensor] + + + +.. py:class:: DeNSRank2Head(backbone: fairchem.core.models.base.BackboneInterface, *args, use_denoising: bool = True, **kwargs) + + Bases: :py:obj:`fairchem.core.models.equiformer_v2.heads.rank2.Rank2SymmetricTensorHead` + + + A rank 2 symmetric tensor prediction head. + + .. attribute:: ouput_name + + name of output prediction property (ie, stress) + + .. attribute:: sphharm_norm + + layer normalization for spherical harmonic edge weights + + .. attribute:: xedge_layer_norm + + embedding layer norm + + .. attribute:: block + + rank 2 equivariant symmetric tensor block + + + .. py:attribute:: use_denoising + + + .. py:method:: forward(data: torch_geometric.data.batch.Batch, emb: dict[str, torch.Tensor]) -> dict[str, torch.Tensor] + + :param data: data batch + :param emb: dictionary with embedding object and graph data + + Returns: dict of {output property name: predicted value} + + + diff --git a/_sources/autoapi/core/models/equiformer_v2/index.rst b/_sources/autoapi/core/models/equiformer_v2/index.rst index 4dec7216cd..c08899cfee 100644 --- a/_sources/autoapi/core/models/equiformer_v2/index.rst +++ b/_sources/autoapi/core/models/equiformer_v2/index.rst @@ -14,6 +14,7 @@ Submodules /autoapi/core/models/equiformer_v2/drop/index /autoapi/core/models/equiformer_v2/edge_rot_mat/index /autoapi/core/models/equiformer_v2/equiformer_v2/index + /autoapi/core/models/equiformer_v2/equiformer_v2_dens/index /autoapi/core/models/equiformer_v2/equiformer_v2_deprecated/index /autoapi/core/models/equiformer_v2/eqv2_to_eqv2_hydra/index /autoapi/core/models/equiformer_v2/gaussian_rbf/index diff --git a/_sources/autoapi/core/models/equiformer_v2/trainers/dens_trainer/index.rst b/_sources/autoapi/core/models/equiformer_v2/trainers/dens_trainer/index.rst new file mode 100644 index 0000000000..3a490873a4 --- /dev/null +++ b/_sources/autoapi/core/models/equiformer_v2/trainers/dens_trainer/index.rst @@ -0,0 +1,213 @@ +core.models.equiformer_v2.trainers.dens_trainer +=============================================== + +.. py:module:: core.models.equiformer_v2.trainers.dens_trainer + +.. autoapi-nested-parse:: + + Copyright (c) Meta, Inc. and its affiliates. + + This source code is licensed under the MIT license found in the + LICENSE file in the root directory of this source tree. + + + +Classes +------- + +.. autoapisummary:: + + core.models.equiformer_v2.trainers.dens_trainer.DenoisingPosParams + core.models.equiformer_v2.trainers.dens_trainer.DenoisingForcesTrainer + + +Functions +--------- + +.. autoapisummary:: + + core.models.equiformer_v2.trainers.dens_trainer.add_gaussian_noise_to_position + core.models.equiformer_v2.trainers.dens_trainer.add_gaussian_noise_schedule_to_position + core.models.equiformer_v2.trainers.dens_trainer.denoising_pos_eval + core.models.equiformer_v2.trainers.dens_trainer.compute_atomwise_denoising_pos_and_force_hybrid_loss + + +Module Contents +--------------- + +.. py:class:: DenoisingPosParams + + .. py:attribute:: prob + :type: float + :value: 0.0 + + + + .. py:attribute:: fixed_noise_std + :type: bool + :value: False + + + + .. py:attribute:: std + :type: float + :value: None + + + + .. py:attribute:: num_steps + :type: int + :value: None + + + + .. py:attribute:: std_low + :type: float + :value: None + + + + .. py:attribute:: std_high + :type: float + :value: None + + + + .. py:attribute:: corrupt_ratio + :type: float + :value: None + + + + .. py:attribute:: all_atoms + :type: bool + :value: False + + + + .. py:attribute:: denoising_pos_coefficient + :type: float + :value: None + + + +.. py:function:: add_gaussian_noise_to_position(batch, std, corrupt_ratio=None, all_atoms=False) + + 1. Update `pos` in `batch`. + 2. Add `noise_vec` to `batch`, which will serve as the target for denoising positions. + 3. Add `denoising_pos_forward` to switch to denoising mode during training. + 4. Add `noise_mask` for partially corrupted structures when `corrupt_ratio` is not None. + 5. If `all_atoms` == True, we add noise to all atoms including fixed ones. + 6. Check whether `batch` has `md`. We do not add noise to structures from MD split. + + +.. py:function:: add_gaussian_noise_schedule_to_position(batch, std_low, std_high, num_steps, corrupt_ratio=None, all_atoms=False) + + 1. Similar to above, update positions in batch with gaussian noise, but + additionally, also save the sigmas the noise vectors are sampled from. + 2. Add `noise_mask` for partially corrupted structures when `corrupt_ratio` + is not None. + 3. If `all_atoms` == True, we add noise to all atoms including fixed ones. + 4. Check whether `batch` has `md`. We do not add noise to structures from MD split. + + +.. py:function:: denoising_pos_eval(evaluator: fairchem.core.modules.evaluator.Evaluator, prediction: dict[str, torch.Tensor], target: dict[str, torch.Tensor], denoising_targets: tuple[str], prev_metrics: dict[str, torch.Tensor] | None = None, denoising_pos_forward: bool = False) + + 1. Overwrite the original Evaluator.eval() here: https://github.com/Open-Catalyst-Project/ocp/blob/5a7738f9aa80b1a9a7e0ca15e33938b4d2557edd/ocpmodels/modules/evaluator.py#L69-L81 + 2. This is to make sure we separate forces MAE and denoising positions MAE. + + +.. py:function:: compute_atomwise_denoising_pos_and_force_hybrid_loss(pred, target, noise_mask, force_mult, denoising_pos_mult, mask=None) + +.. py:class:: DenoisingForcesTrainer(task: dict[str, str | Any], model: dict[str, Any], outputs: dict[str, str | int], dataset: dict[str, str | float], optimizer: dict[str, str | float], loss_functions: dict[str, str | float], evaluation_metrics: dict[str, str], identifier: str, local_rank: int, timestamp_id: str | None = None, run_dir: str | None = None, is_debug: bool = False, print_every: int = 100, seed: int | None = None, logger: str = 'wandb', amp: bool = False, cpu: bool = False, name: str = 'ocp', slurm: dict | None = None, gp_gpus: int | None = None, inference_only: bool = False) + + Bases: :py:obj:`core.models.equiformer_v2.trainers.forces_trainer.EquiformerV2ForcesTrainer` + + + 1. We add a denoising objective to the original S2EF task. + 2. The denoising objective is that we take as input + atom types, node-wise forces and 3D coordinates perturbed with Gaussian noises and then + output energy of the original structure (3D coordinates without any perturbation) and + the node-wise noises added to the original structure. + 3. This should make models leverage more from training data and enable data augmentation for + the S2EF task. + 4. We should only modify the training part. + 5. For normalizing the outputs of noise prediction, if we use `fixed_noise_std = True`, we use + `std` for the normalization factor. Otherwise, we use `std_high` when `fixed_noise_std = False`. + + :param task: Task configuration. + :type task: dict + :param model: Model configuration. + :type model: dict + :param outputs: Dictionary of model output configuration. + :type outputs: dict + :param dataset: Dataset configuration. The dataset needs to be a SinglePointLMDB dataset. + :type dataset: dict + :param optimizer: Optimizer configuration. + :type optimizer: dict + :param loss_functions: Loss function configuration. + :type loss_functions: dict + :param evaluation_metrics: Evaluation metrics configuration. + :type evaluation_metrics: dict + :param identifier: Experiment identifier that is appended to log directory. + :type identifier: str + :param run_dir: Path to the run directory where logs are to be saved. + (default: :obj:`None`) + :type run_dir: str, optional + :param timestamp_id: timestamp identifier. + :type timestamp_id: str, optional + :param run_dir: Run directory used to save checkpoints and results. + :type run_dir: str, optional + :param is_debug: Run in debug mode. + (default: :obj:`False`) + :type is_debug: bool, optional + :param print_every: Frequency of printing logs. + (default: :obj:`100`) + :type print_every: int, optional + :param seed: Random number seed. + (default: :obj:`None`) + :type seed: int, optional + :param logger: Type of logger to be used. + (default: :obj:`wandb`) + :type logger: str, optional + :param local_rank: Local rank of the process, only applicable for distributed training. + (default: :obj:`0`) + :type local_rank: int, optional + :param amp: Run using automatic mixed precision. + (default: :obj:`False`) + :type amp: bool, optional + :param cpu: If True will run on CPU. Default is False, will attempt to use cuda. + :type cpu: bool + :param name: Trainer name. + :type name: str + :param slurm: Slurm configuration. Currently just for keeping track. + (default: :obj:`{}`) + :type slurm: dict + :param gp_gpus: Number of graph parallel GPUs. + :type gp_gpus: int, optional + :param inference_only: If true trainer will be loaded for inference only. + (ie datasets, optimizer, schedular, etc, will not be instantiated) + :type inference_only: bool + + + .. py:attribute:: use_denoising_pos + + + .. py:attribute:: denoising_pos_params + + + .. py:property:: denoising_targets + + + .. py:method:: train(disable_eval_tqdm=False) + + + .. py:method:: _compute_loss(out, batch) + + + .. py:method:: _compute_metrics(out, batch, evaluator, metrics=None) + + + .. py:method:: predict(data_loader, per_image: bool = True, results_file: str | None = None, disable_tqdm: bool = False) + + diff --git a/_sources/autoapi/core/models/equiformer_v2/trainers/energy_trainer/index.rst b/_sources/autoapi/core/models/equiformer_v2/trainers/energy_trainer/index.rst index 196a4b0b02..b6e40fbabb 100644 --- a/_sources/autoapi/core/models/equiformer_v2/trainers/energy_trainer/index.rst +++ b/_sources/autoapi/core/models/equiformer_v2/trainers/energy_trainer/index.rst @@ -23,7 +23,7 @@ Classes Module Contents --------------- -.. py:class:: EquiformerV2EnergyTrainer(task: dict[str, str | Any], model: dict[str, Any], outputs: dict[str, str | int], dataset: dict[str, str | float], optimizer: dict[str, str | float], loss_functions: dict[str, str | float], evaluation_metrics: dict[str, str], identifier: str, local_rank: int, timestamp_id: str | None = None, run_dir: str | None = None, is_debug: bool = False, print_every: int = 100, seed: int | None = None, logger: str = 'wandb', amp: bool = False, cpu: bool = False, name: str = 'ocp', slurm=None, gp_gpus: int | None = None, inference_only: bool = False) +.. py:class:: EquiformerV2EnergyTrainer(task: dict[str, str | Any], model: dict[str, Any], outputs: dict[str, str | int], dataset: dict[str, str | float], optimizer: dict[str, str | float], loss_functions: dict[str, str | float], evaluation_metrics: dict[str, str], identifier: str, local_rank: int, timestamp_id: str | None = None, run_dir: str | None = None, is_debug: bool = False, print_every: int = 100, seed: int | None = None, logger: str = 'wandb', amp: bool = False, cpu: bool = False, name: str = 'ocp', slurm: dict | None = None, gp_gpus: int | None = None, inference_only: bool = False) Bases: :py:obj:`fairchem.core.trainers.OCPTrainer` @@ -41,7 +41,7 @@ Module Contents :type task: dict :param model: Model configuration. :type model: dict - :param outputs: Output property configuration. + :param outputs: Dictionary of model output configuration. :type outputs: dict :param dataset: Dataset configuration. The dataset needs to be a SinglePointLMDB dataset. :type dataset: dict @@ -56,6 +56,10 @@ Module Contents :param run_dir: Path to the run directory where logs are to be saved. (default: :obj:`None`) :type run_dir: str, optional + :param timestamp_id: timestamp identifier. + :type timestamp_id: str, optional + :param run_dir: Run directory used to save checkpoints and results. + :type run_dir: str, optional :param is_debug: Run in debug mode. (default: :obj:`False`) :type is_debug: bool, optional @@ -68,12 +72,24 @@ Module Contents :param logger: Type of logger to be used. (default: :obj:`wandb`) :type logger: str, optional + :param local_rank: Local rank of the process, only applicable for distributed training. + (default: :obj:`0`) + :type local_rank: int, optional :param amp: Run using automatic mixed precision. (default: :obj:`False`) :type amp: bool, optional + :param cpu: If True will run on CPU. Default is False, will attempt to use cuda. + :type cpu: bool + :param name: Trainer name. + :type name: str :param slurm: Slurm configuration. Currently just for keeping track. (default: :obj:`{}`) :type slurm: dict + :param gp_gpus: Number of graph parallel GPUs. + :type gp_gpus: int, optional + :param inference_only: If true trainer will be loaded for inference only. + (ie datasets, optimizer, schedular, etc, will not be instantiated) + :type inference_only: bool .. py:method:: load_extras() diff --git a/_sources/autoapi/core/models/equiformer_v2/trainers/forces_trainer/index.rst b/_sources/autoapi/core/models/equiformer_v2/trainers/forces_trainer/index.rst index fcde30d1e2..dca458b1ba 100644 --- a/_sources/autoapi/core/models/equiformer_v2/trainers/forces_trainer/index.rst +++ b/_sources/autoapi/core/models/equiformer_v2/trainers/forces_trainer/index.rst @@ -23,7 +23,7 @@ Classes Module Contents --------------- -.. py:class:: EquiformerV2ForcesTrainer(task: dict[str, str | Any], model: dict[str, Any], outputs: dict[str, str | int], dataset: dict[str, str | float], optimizer: dict[str, str | float], loss_functions: dict[str, str | float], evaluation_metrics: dict[str, str], identifier: str, local_rank: int, timestamp_id: str | None = None, run_dir: str | None = None, is_debug: bool = False, print_every: int = 100, seed: int | None = None, logger: str = 'wandb', amp: bool = False, cpu: bool = False, name: str = 'ocp', slurm=None, gp_gpus: int | None = None, inference_only: bool = False) +.. py:class:: EquiformerV2ForcesTrainer(task: dict[str, str | Any], model: dict[str, Any], outputs: dict[str, str | int], dataset: dict[str, str | float], optimizer: dict[str, str | float], loss_functions: dict[str, str | float], evaluation_metrics: dict[str, str], identifier: str, local_rank: int, timestamp_id: str | None = None, run_dir: str | None = None, is_debug: bool = False, print_every: int = 100, seed: int | None = None, logger: str = 'wandb', amp: bool = False, cpu: bool = False, name: str = 'ocp', slurm: dict | None = None, gp_gpus: int | None = None, inference_only: bool = False) Bases: :py:obj:`fairchem.core.trainers.OCPTrainer` @@ -41,7 +41,7 @@ Module Contents :type task: dict :param model: Model configuration. :type model: dict - :param outputs: Output property configuration. + :param outputs: Dictionary of model output configuration. :type outputs: dict :param dataset: Dataset configuration. The dataset needs to be a SinglePointLMDB dataset. :type dataset: dict @@ -56,6 +56,10 @@ Module Contents :param run_dir: Path to the run directory where logs are to be saved. (default: :obj:`None`) :type run_dir: str, optional + :param timestamp_id: timestamp identifier. + :type timestamp_id: str, optional + :param run_dir: Run directory used to save checkpoints and results. + :type run_dir: str, optional :param is_debug: Run in debug mode. (default: :obj:`False`) :type is_debug: bool, optional @@ -68,12 +72,24 @@ Module Contents :param logger: Type of logger to be used. (default: :obj:`wandb`) :type logger: str, optional + :param local_rank: Local rank of the process, only applicable for distributed training. + (default: :obj:`0`) + :type local_rank: int, optional :param amp: Run using automatic mixed precision. (default: :obj:`False`) :type amp: bool, optional + :param cpu: If True will run on CPU. Default is False, will attempt to use cuda. + :type cpu: bool + :param name: Trainer name. + :type name: str :param slurm: Slurm configuration. Currently just for keeping track. (default: :obj:`{}`) :type slurm: dict + :param gp_gpus: Number of graph parallel GPUs. + :type gp_gpus: int, optional + :param inference_only: If true trainer will be loaded for inference only. + (ie datasets, optimizer, schedular, etc, will not be instantiated) + :type inference_only: bool .. py:method:: load_extras() -> None diff --git a/_sources/autoapi/core/models/equiformer_v2/trainers/index.rst b/_sources/autoapi/core/models/equiformer_v2/trainers/index.rst index d6393ec61e..39df39ac70 100644 --- a/_sources/autoapi/core/models/equiformer_v2/trainers/index.rst +++ b/_sources/autoapi/core/models/equiformer_v2/trainers/index.rst @@ -10,6 +10,7 @@ Submodules .. toctree:: :maxdepth: 1 + /autoapi/core/models/equiformer_v2/trainers/dens_trainer/index /autoapi/core/models/equiformer_v2/trainers/energy_trainer/index /autoapi/core/models/equiformer_v2/trainers/forces_trainer/index /autoapi/core/models/equiformer_v2/trainers/lr_scheduler/index diff --git a/_sources/autoapi/core/trainers/base_trainer/index.rst b/_sources/autoapi/core/trainers/base_trainer/index.rst index 07ef372881..bf5c7780c5 100644 --- a/_sources/autoapi/core/trainers/base_trainer/index.rst +++ b/_sources/autoapi/core/trainers/base_trainer/index.rst @@ -51,6 +51,11 @@ Module Contents + .. py:attribute:: ema + :value: None + + + .. py:attribute:: timestamp_id :type: str diff --git a/_sources/autoapi/core/trainers/index.rst b/_sources/autoapi/core/trainers/index.rst index 3c83608998..abac26fa14 100644 --- a/_sources/autoapi/core/trainers/index.rst +++ b/_sources/autoapi/core/trainers/index.rst @@ -54,6 +54,11 @@ Package Contents + .. py:attribute:: ema + :value: None + + + .. py:attribute:: timestamp_id :type: str @@ -176,7 +181,7 @@ Package Contents .. py:method:: save_results(predictions: dict[str, numpy.typing.NDArray], results_file: str | None, keys: collections.abc.Sequence[str] | None = None) -> None -.. py:class:: OCPTrainer(task: dict[str, str | Any], model: dict[str, Any], outputs: dict[str, str | int], dataset: dict[str, str | float], optimizer: dict[str, str | float], loss_functions: dict[str, str | float], evaluation_metrics: dict[str, str], identifier: str, local_rank: int, timestamp_id: str | None = None, run_dir: str | None = None, is_debug: bool = False, print_every: int = 100, seed: int | None = None, logger: str = 'wandb', amp: bool = False, cpu: bool = False, name: str = 'ocp', slurm=None, gp_gpus: int | None = None, inference_only: bool = False) +.. py:class:: OCPTrainer(task: dict[str, str | Any], model: dict[str, Any], outputs: dict[str, str | int], dataset: dict[str, str | float], optimizer: dict[str, str | float], loss_functions: dict[str, str | float], evaluation_metrics: dict[str, str], identifier: str, local_rank: int, timestamp_id: str | None = None, run_dir: str | None = None, is_debug: bool = False, print_every: int = 100, seed: int | None = None, logger: str = 'wandb', amp: bool = False, cpu: bool = False, name: str = 'ocp', slurm: dict | None = None, gp_gpus: int | None = None, inference_only: bool = False) Bases: :py:obj:`fairchem.core.trainers.base_trainer.BaseTrainer` @@ -194,7 +199,7 @@ Package Contents :type task: dict :param model: Model configuration. :type model: dict - :param outputs: Output property configuration. + :param outputs: Dictionary of model output configuration. :type outputs: dict :param dataset: Dataset configuration. The dataset needs to be a SinglePointLMDB dataset. :type dataset: dict @@ -209,6 +214,10 @@ Package Contents :param run_dir: Path to the run directory where logs are to be saved. (default: :obj:`None`) :type run_dir: str, optional + :param timestamp_id: timestamp identifier. + :type timestamp_id: str, optional + :param run_dir: Run directory used to save checkpoints and results. + :type run_dir: str, optional :param is_debug: Run in debug mode. (default: :obj:`False`) :type is_debug: bool, optional @@ -221,12 +230,24 @@ Package Contents :param logger: Type of logger to be used. (default: :obj:`wandb`) :type logger: str, optional + :param local_rank: Local rank of the process, only applicable for distributed training. + (default: :obj:`0`) + :type local_rank: int, optional :param amp: Run using automatic mixed precision. (default: :obj:`False`) :type amp: bool, optional + :param cpu: If True will run on CPU. Default is False, will attempt to use cuda. + :type cpu: bool + :param name: Trainer name. + :type name: str :param slurm: Slurm configuration. Currently just for keeping track. (default: :obj:`{}`) :type slurm: dict + :param gp_gpus: Number of graph parallel GPUs. + :type gp_gpus: int, optional + :param inference_only: If true trainer will be loaded for inference only. + (ie datasets, optimizer, schedular, etc, will not be instantiated) + :type inference_only: bool .. py:method:: train(disable_eval_tqdm: bool = False) -> None @@ -253,6 +274,6 @@ Package Contents .. py:method:: predict(data_loader, per_image: bool = True, results_file: str | None = None, disable_tqdm: bool = False) - .. py:method:: run_relaxations(split='val') + .. py:method:: run_relaxations() diff --git a/_sources/autoapi/core/trainers/ocp_trainer/index.rst b/_sources/autoapi/core/trainers/ocp_trainer/index.rst index fb56f92e63..50b5df027d 100644 --- a/_sources/autoapi/core/trainers/ocp_trainer/index.rst +++ b/_sources/autoapi/core/trainers/ocp_trainer/index.rst @@ -23,7 +23,7 @@ Classes Module Contents --------------- -.. py:class:: OCPTrainer(task: dict[str, str | Any], model: dict[str, Any], outputs: dict[str, str | int], dataset: dict[str, str | float], optimizer: dict[str, str | float], loss_functions: dict[str, str | float], evaluation_metrics: dict[str, str], identifier: str, local_rank: int, timestamp_id: str | None = None, run_dir: str | None = None, is_debug: bool = False, print_every: int = 100, seed: int | None = None, logger: str = 'wandb', amp: bool = False, cpu: bool = False, name: str = 'ocp', slurm=None, gp_gpus: int | None = None, inference_only: bool = False) +.. py:class:: OCPTrainer(task: dict[str, str | Any], model: dict[str, Any], outputs: dict[str, str | int], dataset: dict[str, str | float], optimizer: dict[str, str | float], loss_functions: dict[str, str | float], evaluation_metrics: dict[str, str], identifier: str, local_rank: int, timestamp_id: str | None = None, run_dir: str | None = None, is_debug: bool = False, print_every: int = 100, seed: int | None = None, logger: str = 'wandb', amp: bool = False, cpu: bool = False, name: str = 'ocp', slurm: dict | None = None, gp_gpus: int | None = None, inference_only: bool = False) Bases: :py:obj:`fairchem.core.trainers.base_trainer.BaseTrainer` @@ -41,7 +41,7 @@ Module Contents :type task: dict :param model: Model configuration. :type model: dict - :param outputs: Output property configuration. + :param outputs: Dictionary of model output configuration. :type outputs: dict :param dataset: Dataset configuration. The dataset needs to be a SinglePointLMDB dataset. :type dataset: dict @@ -56,6 +56,10 @@ Module Contents :param run_dir: Path to the run directory where logs are to be saved. (default: :obj:`None`) :type run_dir: str, optional + :param timestamp_id: timestamp identifier. + :type timestamp_id: str, optional + :param run_dir: Run directory used to save checkpoints and results. + :type run_dir: str, optional :param is_debug: Run in debug mode. (default: :obj:`False`) :type is_debug: bool, optional @@ -68,12 +72,24 @@ Module Contents :param logger: Type of logger to be used. (default: :obj:`wandb`) :type logger: str, optional + :param local_rank: Local rank of the process, only applicable for distributed training. + (default: :obj:`0`) + :type local_rank: int, optional :param amp: Run using automatic mixed precision. (default: :obj:`False`) :type amp: bool, optional + :param cpu: If True will run on CPU. Default is False, will attempt to use cuda. + :type cpu: bool + :param name: Trainer name. + :type name: str :param slurm: Slurm configuration. Currently just for keeping track. (default: :obj:`{}`) :type slurm: dict + :param gp_gpus: Number of graph parallel GPUs. + :type gp_gpus: int, optional + :param inference_only: If true trainer will be loaded for inference only. + (ie datasets, optimizer, schedular, etc, will not be instantiated) + :type inference_only: bool .. py:method:: train(disable_eval_tqdm: bool = False) -> None @@ -100,6 +116,6 @@ Module Contents .. py:method:: predict(data_loader, per_image: bool = True, results_file: str | None = None, disable_tqdm: bool = False) - .. py:method:: run_relaxations(split='val') + .. py:method:: run_relaxations() diff --git a/_sources/core/model_checkpoints.md b/_sources/core/model_checkpoints.md index 6cfba70555..342498e14e 100644 --- a/_sources/core/model_checkpoints.md +++ b/_sources/core/model_checkpoints.md @@ -149,7 +149,7 @@ OC22 dataset or pretrained models, as well as the original paper for each model: | GemNet-OC-S2EF-ODAC | GemNet-OC | [checkpoint](https://dl.fbaipublicfiles.com/dac/checkpoints_20231018/Gemnet-OC.pt) | [config](https://github.com/FAIR-Chem/fairchem/tree/main/configs/odac/s2ef/gemnet-oc.yml) | | eSCN-S2EF-ODAC | eSCN | [checkpoint](https://dl.fbaipublicfiles.com/dac/checkpoints_20231018/eSCN.pt) | [config](https://github.com/FAIR-Chem/fairchem/tree/main/configs/odac/s2ef/eSCN.yml) | | EquiformerV2-S2EF-ODAC | EquiformerV2 | [checkpoint](https://dl.fbaipublicfiles.com/dac/checkpoints_20231116/eqv2_31M.pt) | [config](https://github.com/FAIR-Chem/fairchem/tree/main/configs/odac/s2ef/eqv2_31M.yml) | -| EquiformerV2-Large-S2EF-ODAC | EquiformerV2 (Large) | [checkpoint](https://dl.fbaipublicfiles.com/dac/checkpoints_20231018/Equiformer_V2_Large.pt) | [config](https://github.com/FAIR-Chem/fairchem/tree/main/configs/odac/s2ef/eqv2_153M.yml) | +| EquiformerV2-Large-S2EF-ODAC | EquiformerV2 (Large) | [checkpoint](https://dl.fbaipublicfiles.com/dac/checkpoints_20231116/Equiformer_V2_Large.pt) | [config](https://github.com/FAIR-Chem/fairchem/tree/main/configs/odac/s2ef/eqv2_153M.yml) | ## IS2RE Direct models @@ -157,7 +157,7 @@ OC22 dataset or pretrained models, as well as the original paper for each model: |-------------------------|--------------|--- | --- | | Gemnet-OC-IS2RE-ODAC | Gemnet-OC | [checkpoint](https://dl.fbaipublicfiles.com/dac/checkpoints_20231018/Gemnet-OC_Direct.pt) | [config](https://github.com/FAIR-Chem/fairchem/tree/main/configs/odac/is2re/gemnet-oc.yml) | | eSCN-IS2RE-ODAC | eSCN | [checkpoint](https://dl.fbaipublicfiles.com/dac/checkpoints_20231018/eSCN_Direct.pt) | [config](https://github.com/FAIR-Chem/fairchem/tree/main/configs/odac/is2re/eSCN.yml) | -| EquiformerV2-IS2RE-ODAC | EquiformerV2 | [checkpoint](https://dl.fbaipublicfiles.com/dac/checkpoints_20231018/Equiformer_V2_Direct.pt) | [config](https://github.com/FAIR-Chem/fairchem/tree/main/configs/odac/is2re/eqv2_31M.yml) | +| EquiformerV2-IS2RE-ODAC | EquiformerV2 | [checkpoint](https://dl.fbaipublicfiles.com/dac/checkpoints_20231116/Equiformer_V2_Direct.pt) | [config](https://github.com/FAIR-Chem/fairchem/tree/main/configs/odac/is2re/eqv2_31M.yml) | The models in the table above were trained to predict relaxed energy directly. Relaxed energies can also be predicted by running structural relaxations using the S2EF models from the previous section. diff --git a/autoapi/core/_cli_hydra/index.html b/autoapi/core/_cli_hydra/index.html index d6113a591a..170ee8fb48 100644 --- a/autoapi/core/_cli_hydra/index.html +++ b/autoapi/core/_cli_hydra/index.html @@ -617,6 +617,7 @@

Contents

  • logger
  • Submitit
    • @@ -674,7 +675,7 @@

    Functions

    get_hydra_config_from_yaml(→ omegaconf.DictConfig)

    -

    runner_wrapper(config, cli_args)

    +

    runner_wrapper(config)

    main([args, override_args])

    @@ -707,12 +708,17 @@

    Module Contents
    -__call__(dict_config: omegaconf.DictConfig, cli_args: argparse.Namespace) → None#
    +__call__(dict_config: omegaconf.DictConfig) → None# +
    + +
    +
    +_init_logger() → None#
    -checkpoint(*args, **kwargs)#
    +checkpoint(*args, **kwargs) → submitit.helpers.DelayedSubmission#

    Resubmits the same callable with the same arguments

    @@ -720,7 +726,7 @@

    Module Contents
    -core._cli_hydra.map_cli_args_to_dist_config(cli_args: argparse.Namespace) → dict#
    +core._cli_hydra.map_cli_args_to_dist_config(cli_args: omegaconf.DictConfig) → dict#
    @@ -730,7 +736,7 @@

    Module Contents
    -core._cli_hydra.runner_wrapper(config: omegaconf.DictConfig, cli_args: argparse.Namespace)#
    +core._cli_hydra.runner_wrapper(config: omegaconf.DictConfig)#

    @@ -813,6 +819,7 @@

    Module Contentslogger
  • Submitit
    • diff --git a/autoapi/core/common/relaxation/ase_utils/index.html b/autoapi/core/common/relaxation/ase_utils/index.html index 00daa76607..ce8cd39baf 100644 --- a/autoapi/core/common/relaxation/ase_utils/index.html +++ b/autoapi/core/common/relaxation/ase_utils/index.html @@ -610,15 +610,18 @@

    Contents

    -class core.common.relaxation.ase_utils.OCPCalculator(config_yml: str | None = None, checkpoint_path: str | None = None, model_name: str | None = None, local_cache: str | None = None, trainer: str | None = None, cutoff: int = 6, max_neighbors: int = 50, cpu: bool = True, seed: int | None = None)#
    +class core.common.relaxation.ase_utils.OCPCalculator(config_yml: str | None = None, checkpoint_path: str | pathlib.Path | None = None, model_name: str | None = None, local_cache: str | None = None, trainer: str | None = None, cpu: bool = True, seed: int | None = None)#

    Bases: ase.calculators.calculator.Calculator

    -

    Base-class for all ASE calculators.

    -

    A calculator must raise PropertyNotImplementedError if asked for a -property that it can’t calculate. So, if calculation of the -stress tensor has not been implemented, get_stress(atoms) should -raise PropertyNotImplementedError. This can be achieved simply by not -including the string ‘stress’ in the list implemented_properties -which is a class member. These are the names of the standard -properties: ‘energy’, ‘forces’, ‘stress’, ‘dipole’, ‘charges’, -‘magmom’ and ‘magmoms’.

    +

    ASE based calculator using an OCP model

    -
    -implemented_properties: ClassVar[list[str]] = ['energy', 'forces']#
    -

    Properties calculator can handle (energy, forces, …)

    -
    +
    +_reshaped_props#
    +
    @@ -706,48 +731,32 @@

    Module Contentsa2g#

    +
    +
    +implemented_properties#
    +

    Properties calculator can handle (energy, forces, …)

    +
    +
    load_checkpoint(checkpoint_path: str, checkpoint: dict | None = None) → None#

    Load existing trained model

    Parameters:
    -

    checkpoint_path – string -Path to trained model

    +
      +

    • checkpoint_path – string +Path to trained model

      • +

    • checkpoint – dict +A pretrained checkpoint dict

      • +
    -calculate(atoms: ase.Atoms, properties, system_changes) → None#
    -

    Do the calculation.

    -
    -
    properties: list of str

    List of what needs to be calculated. Can be any combination -of ‘energy’, ‘forces’, ‘stress’, ‘dipole’, ‘charges’, ‘magmom’ -and ‘magmoms’.

    -
    -
    system_changes: list of str

    List of what has changed since last calculation. Can be -any combination of these six: ‘positions’, ‘numbers’, ‘cell’, -‘pbc’, ‘initial_charges’ and ‘initial_magmoms’.

    -
    -
    -

    Subclasses need to implement this, but can ignore properties -and system_changes if they want. Calculated properties should -be inserted into results dictionary like shown in this dummy -example:

    -
    self.results = {'energy': 0.0,
-                'forces': np.zeros((len(atoms), 3)),
-                'stress': np.zeros(6),
-                'dipole': np.zeros(3),
-                'charges': np.zeros(len(atoms)),
-                'magmom': 0.0,
-                'magmoms': np.zeros(len(atoms))}
-
    -
    -

    The subclass implementation should first call this -implementation to set the atoms attribute and create any missing -directories.

    +calculate(atoms: ase.Atoms | torch_geometric.data.Batch, properties, system_changes) → None# +

    Calculate implemented properties for a single Atoms object or a Batch of them.

    @@ -820,15 +829,18 @@

    Module Contents @@ -665,10 +660,7 @@

    Classes#<
    - - - - +

    LBFGS

    TorchCalc

    Limited memory BFGS optimizer for batch ML relaxations.
    @@ -678,15 +670,11 @@

    Classes#<

    Module Contents#

    -class core.common.relaxation.optimizers.lbfgs_torch.LBFGS(batch: torch_geometric.data.Batch, model: TorchCalc, maxstep: float = 0.01, memory: int = 100, damping: float = 0.25, alpha: float = 100.0, force_consistent=None, device: str = 'cuda:0', save_full_traj: bool = True, traj_dir: pathlib.Path | None = None, traj_names=None, early_stop_batch: bool = False)#
    -
    -
    -batch#
    -
    - +class core.common.relaxation.optimizers.lbfgs_torch.LBFGS(optimizable_batch: core.common.relaxation.optimizers.optimizable.OptimizableBatch, maxstep: float = 0.02, memory: int = 100, damping: float = 1.2, alpha: float = 100.0, save_full_traj: bool = True, traj_dir: pathlib.Path | None = None, traj_names: list[str] | None = None)# +

    Limited memory BFGS optimizer for batch ML relaxations.

    -
    -model#
    +
    +optimizable#
    @@ -714,16 +702,6 @@

    Module ContentsH0#

    -
    -
    -force_consistent#
    -
    - -
    -
    -device#
    -
    -
    save_full#
    @@ -740,68 +718,68 @@

    Module Contents -
    -early_stop_batch#
    +
    +trajectories = None#

    -
    -otf_graph = True#
    +
    +fmax = None#
    -
    -
    -get_energy_and_forces(apply_constraint: bool = True)#
    +
    +
    +steps = None#
    -
    -
    -set_positions(update, update_mask) → None#
    +
    +
    +s#
    -
    -
    -check_convergence(iteration, forces=None, energy=None)#
    +
    +
    +y#
    -
    -
    -run(fmax, steps)#
    +
    +
    +rho#
    -
    -
    -step(iteration: int, forces: torch.Tensor | None, update_mask: torch.Tensor) → None#
    +
    +
    +r0 = None#
    -
    -
    -write(energy, forces, update_mask) → None#
    +
    +
    +f0 = None#
    -
    +
    +
    +run(fmax, steps)#
    +
    -
    -
    -class core.common.relaxation.optimizers.lbfgs_torch.TorchCalc(model, transform=None)#
    -
    -
    -model#
    +
    +
    +determine_step(dr)#
    -
    -
    -transform#
    +
    +
    +_batched_dot(x: torch.Tensor, y: torch.Tensor)#
    -
    -get_energy_and_forces(atoms, apply_constraint: bool = True)#
    +
    +step(iteration: int) → None#
    -
    -update_graph(atoms)#
    +
    +write() → None#
    @@ -877,35 +855,30 @@

    Module ContentsClasses
  • Module Contents
    • diff --git a/autoapi/core/datasets/ase_datasets/index.html b/autoapi/core/datasets/ase_datasets/index.html index e457dde4ae..76b786788e 100644 --- a/autoapi/core/datasets/ase_datasets/index.html +++ b/autoapi/core/datasets/ase_datasets/index.html @@ -648,7 +648,6 @@

    Contents

  • AseDBDataset.connect_db()
  • AseDBDataset.__del__()
  • AseDBDataset.sample_property_metadata()
    • -
  • AseDBDataset.get_relaxed_energy()
    • @@ -1034,11 +1033,6 @@

    Module Contentssample_property_metadata(num_samples: int = 100) → dict#
    -
    -
    -abstract get_relaxed_energy(identifier)#
    -
    - @@ -1147,7 +1141,6 @@

    Module ContentsAseDBDataset.connect_db()
  • AseDBDataset.__del__()
  • AseDBDataset.sample_property_metadata()
    • -
  • AseDBDataset.get_relaxed_energy()
    • diff --git a/autoapi/core/datasets/index.html b/autoapi/core/datasets/index.html index d7f8508d0a..76f6d346dc 100644 --- a/autoapi/core/datasets/index.html +++ b/autoapi/core/datasets/index.html @@ -620,7 +620,6 @@

    Contents

  • AseDBDataset.connect_db()
  • AseDBDataset.__del__()
  • AseDBDataset.sample_property_metadata()
    • -
  • AseDBDataset.get_relaxed_energy()
  • AseReadDataset
  • AseReadDataset + -
  • cell2planes() (data.oc.utils.geometry.Geometry static method)
    • +
  • cell_factor (core.common.relaxation.optimizable.OptimizableUnitCellBatch attribute) + +
  • cell_offsets (core.models.base.GraphData attribute)
  • centering (core.models.equiformer_v2.layer_norm.EquivariantRMSNormArraySphericalHarmonicsV2 attribute) @@ -2247,8 +2343,12 @@

    C

  • change_mat (core.models.equiformer_v2.heads.rank2.Rank2DecompositionEdgeBlock attribute)
    • -
  • check_convergence() (core.common.relaxation.optimizers.lbfgs_torch.LBFGS method) +
  • check_state() (core.common.relaxation.optimizable.OptimizableBatch method) + +
  • check_traj_files() (in module core.common.utils)
  • checkpoint() (core._cli.Runner method) @@ -2330,6 +2430,8 @@

    C

  • collected_params (core.modules.exponential_moving_average.ExponentialMovingAverage attribute)
  • CombineYRotations() (core.models.scn.spherical_harmonics.SphericalHarmonicsHelper method) +
    • +
  • compare_batches() (in module core.common.relaxation.optimizable)
  • compile (core.models.escn.escn_exportable.eSCN attribute)
    • @@ -2365,6 +2467,8 @@

    C

  • (core.models.escn.so3_exportable.CoefficientMapping method)
    • +
  • compute_atomwise_denoising_pos_and_force_hybrid_loss() (in module core.models.equiformer_v2.trainers.dens_trainer) +
  • compute_hybrid_success() (in module adsorbml.scripts.dense_eval)
  • compute_neighbors() (in module core.common.utils) @@ -2453,6 +2557,12 @@

    C

  • (core.datasets.LmdbDataset method)
  • (core.datasets.oc22_lmdb_dataset.OC22LmdbDataset method) +
    • + +
  • constant_volume (core.common.relaxation.optimizable.OptimizableUnitCellBatch attribute) + +
  • construct_mapping() (core.common.utils.UniqueKeyLoader method) @@ -2461,6 +2571,18 @@

    C

    • +
  • converged() (core.common.relaxation.optimizable.OptimizableBatch method) + +
    • +
  • converged_mask (core.common.relaxation.optimizable.OptimizableBatch property) + +
  • converged_oszicar() (in module adsorbml.scripts.utils) @@ -2590,6 +2712,13 @@

    C

    • +
  • + core.common.relaxation.optimizable + +
  • @@ -2800,6 +2929,13 @@

    C

    • +
  • + core.models.equiformer_v2.equiformer_v2_dens + +
  • @@ -2858,6 +2994,8 @@

    C

  • module
    • + +
    • core.models.equiformer_v2.layer_norm @@ -2879,8 +3017,6 @@

      C

    • module
    - -
    • core.models.equiformer_v2.so2_ops @@ -2900,6 +3036,13 @@

      C

      • +
    • + core.models.equiformer_v2.trainers.dens_trainer + +
    • @@ -3630,6 +3773,8 @@

      C

    • module
    +
  • corrupt_ratio (core.models.equiformer_v2.trainers.dens_trainer.DenoisingPosParams attribute) +
  • cosine_lr_lambda() (in module core.models.equiformer_v2.trainers.lr_scheduler)
  • cosine_similarity() (in module core.modules.evaluator) @@ -3740,8 +3885,12 @@

    C

    D

    +
  • denoising_pos_block (core.models.equiformer_v2.equiformer_v2_dens.DeNSVectorHead attribute) +
    • +
  • denoising_pos_coefficient (core.models.equiformer_v2.trainers.dens_trainer.DenoisingPosParams attribute) +
    • +
  • denoising_pos_eval() (in module core.models.equiformer_v2.trainers.dens_trainer) +
    • +
  • denoising_pos_params (core.models.equiformer_v2.trainers.dens_trainer.DenoisingForcesTrainer attribute) +
    • +
  • denoising_targets (core.models.equiformer_v2.trainers.dens_trainer.DenoisingForcesTrainer property) +
    • +
  • DenoisingForcesTrainer (class in core.models.equiformer_v2.trainers.dens_trainer) +
    • +
  • DenoisingPosParams (class in core.models.equiformer_v2.trainers.dens_trainer) +
  • denorm() (core.modules.normalization.normalizer.Normalizer method)
  • Dense (class in core.models.gemnet.layers.base_layers) @@ -4113,6 +4282,12 @@

    D

  • dense_rbf_F (core.models.gemnet_gp.layers.atom_update_block.OutputBlock attribute) +
    • +
  • DeNSRank2Head (class in core.models.equiformer_v2.equiformer_v2_dens) +
    • +
  • DeNSScalarHead (class in core.models.equiformer_v2.equiformer_v2_dens) +
    • +
  • DeNSVectorHead (class in core.models.equiformer_v2.equiformer_v2_dens)
  • dereference() (core.modules.normalization.element_references.LinearReferences method)
    • @@ -4132,12 +4307,20 @@

    D

    +
  • determine_step() (core.common.relaxation.optimizers.LBFGS method) + +
  • device (core.common.data_parallel.BalancedBatchSampler attribute)