diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 9cc48234..c7cfb20b 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -18,6 +18,19 @@ concurrency: # yamllint disable rule:line-length jobs: + lint: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Install pre-commit + run: pip install pre-commit + - name: Run pre-commit + run: pre-commit run --all-files --show-diff-on-failure + run-tests: runs-on: ${{ github.event.repository.private && 'self-hosted' || 'ubuntu-latest' }} strategy: @@ -32,12 +45,6 @@ jobs: run: | apt update && apt install -y python3 python3-pip git python3 -m pip install --upgrade pip - - name: Linting - run: | - pip install ruff yamllint - ruff check . - ruff format --check . - yamllint . - name: Install PyTorch env: TORCH: "${{ matrix.torch-version }}" diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 0e3784c6..405c4cbc 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,15 +1,30 @@ exclude: '.git|.tox' -default_stages: [commit] +default_stages: [pre-commit] fail_fast: true +ci: + autoupdate_schedule: quarterly + repos: - - repo: https://github.com/psf/black - rev: 24.4.2 + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.12.2 + hooks: + - id: ruff + args: [--fix] + - id: ruff-format + + - repo: https://github.com/adrienverge/yamllint + rev: v1.37.1 hooks: - - id: black + - id: yamllint - - repo: https://github.com/pycqa/flake8 - rev: 7.1.0 + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v5.0.0 hooks: - - id: flake8 - additional_dependencies: [Flake8-pyproject] + - id: check-case-conflict + - id: check-symlinks + - id: check-yaml + - id: destroyed-symlinks + - id: end-of-file-fixer + - id: mixed-line-ending + - id: trailing-whitespace diff --git a/AUTHORS.md b/AUTHORS.md index d6411c42..fd9d0775 100644 --- a/AUTHORS.md +++ b/AUTHORS.md @@ -25,4 +25,4 @@ Present and past contributors include, in alphabetical order by last name: - Menghang (David) Wang - William C. Witt - Albert Zhu - - Laura Zichi \ No newline at end of file + - Laura Zichi diff --git a/docs/_static/custom.css b/docs/_static/custom.css index 0bada211..64bb0803 100644 --- a/docs/_static/custom.css +++ b/docs/_static/custom.css @@ -17,4 +17,4 @@ /* Justify text in list items that contain multiple sentences */ li p { text-align: justify; -} \ No newline at end of file +} diff --git a/docs/api/data_stats.rst b/docs/api/data_stats.rst index f37e82bc..176ac8d9 100644 --- a/docs/api/data_stats.rst +++ b/docs/api/data_stats.rst @@ -23,7 +23,7 @@ As an example, we show how one can configure the full :class:`~nequip.data.DataS - name: num_neighbors_mean field: _target_: nequip.data.NumNeighbors - metric: + metric: _target_: nequip.data.Mean - name: per_atom_energy_mean field: @@ -55,12 +55,12 @@ As an example, we show how one can configure the full :class:`~nequip.data.DataS .. autoclass:: nequip.data.StandardDeviation :members: - + .. autoclass:: nequip.data.Min :members: - + .. autoclass:: nequip.data.Max :members: .. autoclass:: nequip.data.Count - :members: \ No newline at end of file + :members: diff --git a/docs/api/data_transforms.rst b/docs/api/data_transforms.rst index fcf1c40c..cdc3178b 100644 --- a/docs/api/data_transforms.rst +++ b/docs/api/data_transforms.rst @@ -13,4 +13,4 @@ Data transforms convert the raw data from the :class:`~nequip.data.dataset.Atomi :members: .. autoclass:: nequip.data.transforms.StressSignFlipTransform - :members: \ No newline at end of file + :members: diff --git a/docs/api/datamodule.rst b/docs/api/datamodule.rst index cf31e6aa..6ceb6b8b 100644 --- a/docs/api/datamodule.rst +++ b/docs/api/datamodule.rst @@ -32,4 +32,4 @@ For usage examples and configuration guidance, see the :doc:`../guide/configurat :members: .. autoclass:: nequip.data.datamodule.WaterDataModule - :members: \ No newline at end of file + :members: diff --git a/docs/api/lightning_module.rst b/docs/api/lightning_module.rst index 75e22257..cc446059 100644 --- a/docs/api/lightning_module.rst +++ b/docs/api/lightning_module.rst @@ -14,4 +14,4 @@ NequIPLightningModule :members: .. autoclass:: nequip.train.ScheduleFreeLightningModule - :members: \ No newline at end of file + :members: diff --git a/docs/api/model.rst b/docs/api/model.rst index a22ca451..f9d6d994 100644 --- a/docs/api/model.rst +++ b/docs/api/model.rst @@ -5,5 +5,3 @@ nequip.model nequip_model save_model - - diff --git a/docs/api/nequip_model.rst b/docs/api/nequip_model.rst index 0bfb9f56..5562eb34 100644 --- a/docs/api/nequip_model.rst +++ b/docs/api/nequip_model.rst @@ -3,4 +3,4 @@ NequIP Message Passing GNN Model .. autofunction:: nequip.model.NequIPGNNModel -.. autofunction:: nequip.model.NequIPGNNEnergyModel \ No newline at end of file +.. autofunction:: nequip.model.NequIPGNNEnergyModel diff --git a/docs/api/nn.rst b/docs/api/nn.rst index 5f66c930..551d9a89 100644 --- a/docs/api/nn.rst +++ b/docs/api/nn.rst @@ -6,4 +6,4 @@ nequip.nn :members: .. autoclass:: nequip.nn.PerTypeScaleShift - :members: \ No newline at end of file + :members: diff --git a/docs/api/save_model.rst b/docs/api/save_model.rst index 973b54b4..ffeaf2c2 100644 --- a/docs/api/save_model.rst +++ b/docs/api/save_model.rst @@ -18,8 +18,8 @@ These files can be loaded using the following :class:`~nequip.model.ModelFromChe - Use absolute paths instead of relative paths. - Do not change the directory structure or move your files when using the model loaders. - Ideally, store the original checkpoint/package files somewhere that makes their association with the new training run clear to you. - - Be aware that iterated nested use of :class:`~nequip.model.ModelFromCheckpoint` will result in a checkpoint chaining phenomenon where loading the checkpoint at the end of the chain requires successfully loading every intermediate checkpoint file in the chain. One can break this chain if necessary by using ``nequip-package`` to convert the checkpoint file into a packaged model, and then using :class:`~nequip.model.ModelFromPackage`. + + Be aware that iterated nested use of :class:`~nequip.model.ModelFromCheckpoint` will result in a checkpoint chaining phenomenon where loading the checkpoint at the end of the chain requires successfully loading every intermediate checkpoint file in the chain. One can break this chain if necessary by using ``nequip-package`` to convert the checkpoint file into a packaged model, and then using :class:`~nequip.model.ModelFromPackage`. .. autofunction:: nequip.model.ModelFromCheckpoint @@ -36,4 +36,4 @@ This is typically used by the ASE calculator or other inference applications. Modifying Saved Models ###################### -.. autofunction:: nequip.model.modify \ No newline at end of file +.. autofunction:: nequip.model.modify diff --git a/docs/api/train.rst b/docs/api/train.rst index 94bb6b13..1258b996 100644 --- a/docs/api/train.rst +++ b/docs/api/train.rst @@ -7,4 +7,4 @@ nequip.train lightning_module metrics callbacks - ddp \ No newline at end of file + ddp diff --git a/docs/dev/contributing.md b/docs/dev/contributing.md index 1c92f8f6..97d33ba6 100644 --- a/docs/dev/contributing.md +++ b/docs/dev/contributing.md @@ -1,7 +1,7 @@ # Contributing to NequIP ```{note} -If you want to make a major change, add a new feature, or have any uncertainty about where the change should go or how it should be designed, please reach out to discuss with us first. Details on contacting the development team can be found in `README.md`. +If you want to make a major change, add a new feature, or have any uncertainty about where the change should go or how it should be designed, please reach out to discuss with us first. Details on contacting the development team can be found in `README.md`. ``` ## Design Philosophy @@ -60,7 +60,7 @@ ruff format --check . ``` The formatter can be easily set up to run automatically on file save in various editors. - + You can also use ``pre-commit install`` to install a [pre-commit](https://pre-commit.com/) hook. ```{tip} @@ -73,4 +73,4 @@ PRs should generally come from feature branches. All PRs should be onto the `develop` branch or other feature branches, and **NOT** `main` (unless they fix a significant bug with immediate impact). -On feature branches, it is preferable to [rebase](https://docs.github.com/en/get-started/using-git/about-git-rebase) wherever possible. We strive towards a clean and easily-readable commit history. Note that this does not apply to `develop` or `main` -- the commit history on these core branches are sacred. \ No newline at end of file +On feature branches, it is preferable to [rebase](https://docs.github.com/en/get-started/using-git/about-git-rebase) wherever possible. We strive towards a clean and easily-readable commit history. Note that this does not apply to `develop` or `main` -- the commit history on these core branches are sacred. diff --git a/docs/dev/dev.rst b/docs/dev/dev.rst index d02ea4d7..c9d712ed 100644 --- a/docs/dev/dev.rst +++ b/docs/dev/dev.rst @@ -6,4 +6,4 @@ Developer Guide contributing understanding_nequip extension_packages - ../README \ No newline at end of file + ../README diff --git a/docs/dev/extension_packages.rst b/docs/dev/extension_packages.rst index 9dbb4180..a873da43 100644 --- a/docs/dev/extension_packages.rst +++ b/docs/dev/extension_packages.rst @@ -22,4 +22,4 @@ Example Extension Packages Feel free to use the following extension packages for inspiration and reference: - **Allegro** (`GitHub `__, `Docs `__, `Paper `__): Strictly local equivariant models with excellent scalability for multirank molecular dynamics simulations. -- **NequIP-LES** (`Github `__, `Paper `__): An extension of NequIP and Allegro that adds long-range electrostatics via the Latent Ewald Summation (LES) algorithm. \ No newline at end of file +- **NequIP-LES** (`Github `__, `Paper `__): An extension of NequIP and Allegro that adds long-range electrostatics via the Latent Ewald Summation (LES) algorithm. diff --git a/docs/dev/extension_packages/data.rst b/docs/dev/extension_packages/data.rst index 236d0158..626ec35e 100644 --- a/docs/dev/extension_packages/data.rst +++ b/docs/dev/extension_packages/data.rst @@ -12,7 +12,7 @@ Field Types ~~~~~~~~~~~ - **Graph fields**: per-frame (e.g., ``total_energy``, ``stress``) -- **Node fields**: per-atom (e.g., ``forces``, ``positions``) +- **Node fields**: per-atom (e.g., ``forces``, ``positions``) - **Edge fields**: per-edge (e.g., ``edge_vectors``) - **Long fields**: integer dtype (e.g., ``atomic_numbers``) - **Cartesian tensors**: physical tensors (e.g., ``stress``) @@ -26,7 +26,7 @@ Registration Example register_fields( graph_fields=["custom_energy"], - node_fields=["custom_forces"], + node_fields=["custom_forces"], edge_fields=["custom_edge_attr"], long_fields=["custom_indices"], cartesian_tensor_fields={"custom_tensor": "ij=ji"} @@ -66,4 +66,4 @@ Data Transforms Extension packages can implement custom data transforms to preprocess data during loading. Transforms are classes that implement a ``__call__`` method to modify ``AtomicDataDict`` objects. -See the :doc:`transforms API documentation <../../api/data_transforms>` for available transform classes and their patterns. \ No newline at end of file +See the :doc:`transforms API documentation <../../api/data_transforms>` for available transform classes and their patterns. diff --git a/docs/dev/extension_packages/getting_started.rst b/docs/dev/extension_packages/getting_started.rst index 1e7951ba..a0dec6d6 100644 --- a/docs/dev/extension_packages/getting_started.rst +++ b/docs/dev/extension_packages/getting_started.rst @@ -19,4 +19,4 @@ When the NequIP framework is imported, it automatically discovers and loads all - Register custom OmegaConf resolvers - Perform any other initialization needed for your package -Your package's ``__init__.py`` will be automatically imported whenever the NequIP framework is used, ensuring your extensions are always available. \ No newline at end of file +Your package's ``__init__.py`` will be automatically imported whenever the NequIP framework is used, ensuring your extensions are always available. diff --git a/docs/dev/extension_packages/models.rst b/docs/dev/extension_packages/models.rst index 128f4e56..26ee3e9e 100644 --- a/docs/dev/extension_packages/models.rst +++ b/docs/dev/extension_packages/models.rst @@ -53,7 +53,7 @@ Train-time via config: H: 0.12 model: _target_: nequip.model.ModelFromPackage - # ... + # ... Compile-time for accelerations: @@ -64,4 +64,4 @@ Compile-time for accelerations: Implementation Tips ------------------- -When implementing custom :class:`torch.nn.Module` subclasses, use :meth:`torch.nn.Module.extra_repr` to provide crucial model information for debugging (visible when ``_NEQUIP_LOG_LEVEL=DEBUG`` is set). \ No newline at end of file +When implementing custom :class:`torch.nn.Module` subclasses, use :meth:`torch.nn.Module.extra_repr` to provide crucial model information for debugging (visible when ``_NEQUIP_LOG_LEVEL=DEBUG`` is set). diff --git a/docs/dev/extension_packages/training.rst b/docs/dev/extension_packages/training.rst index 2a19bb8a..5f0ddf15 100644 --- a/docs/dev/extension_packages/training.rst +++ b/docs/dev/extension_packages/training.rst @@ -22,5 +22,5 @@ Examples include :class:`~nequip.train.EMALightningModule` (exponential moving a When composing multiple LightningModule subclasses, follow these rules to avoid the "deadly diamond of death" multiple inheritance pattern: - Use unique, distinguishable names for class attributes to prevent overwrites during inheritance -- Be very careful of inheritance order when creating "diamond" subclasses -- Use assertions to verify the composed module behaves as intended \ No newline at end of file +- Be very careful of inheritance order when creating "diamond" subclasses +- Use assertions to verify the composed module behaves as intended diff --git a/docs/dev/understanding_nequip.rst b/docs/dev/understanding_nequip.rst index f50e3898..2ded89a5 100644 --- a/docs/dev/understanding_nequip.rst +++ b/docs/dev/understanding_nequip.rst @@ -6,4 +6,4 @@ This section provides comprehensive documentation for understanding NequIP's int .. toctree:: :maxdepth: 2 - understanding_nequip/workflows/index \ No newline at end of file + understanding_nequip/workflows/index diff --git a/docs/dev/understanding_nequip/workflows/index.rst b/docs/dev/understanding_nequip/workflows/index.rst index 67d33b72..4018dcb0 100644 --- a/docs/dev/understanding_nequip/workflows/index.rst +++ b/docs/dev/understanding_nequip/workflows/index.rst @@ -8,4 +8,4 @@ This section covers key workflows and processes within the NequIP framework. .. toctree:: :maxdepth: 2 - packaging \ No newline at end of file + packaging diff --git a/docs/guide/accelerations/cuequivariance.md b/docs/guide/accelerations/cuequivariance.md index ed34238c..94ad75c5 100644 --- a/docs/guide/accelerations/cuequivariance.md +++ b/docs/guide/accelerations/cuequivariance.md @@ -83,4 +83,4 @@ nequip-prepare-lmp-mliap \ Where `model_file` can be either a [checkpoint file](../getting-started/files.md#checkpoint-files) (`.ckpt`) or [package file](../getting-started/files.md#package-files). The resulting `.nequip.lmp.pt` file can be used directly in LAMMPS scripts with the `pair_style mliap` command. -See the [ML-IAP documentation](../../integrations/lammps/mliap.md) for complete usage instructions and examples. \ No newline at end of file +See the [ML-IAP documentation](../../integrations/lammps/mliap.md) for complete usage instructions and examples. diff --git a/docs/guide/accelerations/ddp_training.md b/docs/guide/accelerations/ddp_training.md index a4f03824..1c951e56 100644 --- a/docs/guide/accelerations/ddp_training.md +++ b/docs/guide/accelerations/ddp_training.md @@ -6,7 +6,7 @@ There are two ways to set up multi-rank distributed data parallel (DDP) training If train-time compilation is not used, you can use PyTorch Lightning's {class}`~lightning.pytorch.strategies.DDPStrategy`. -See [Lightning's docs](https://lightning.ai/docs/pytorch/stable/accelerators/gpu_intermediate.html#distributed-data-parallel) for how to set it up through Lightning's {class}`~lightning.pytorch.trainer.trainer.Trainer`. +See [Lightning's docs](https://lightning.ai/docs/pytorch/stable/accelerators/gpu_intermediate.html#distributed-data-parallel) for how to set it up through Lightning's {class}`~lightning.pytorch.trainer.trainer.Trainer`. **NOTE** that it is usually not necessary to explicitly set the {class}`~lightning.pytorch.strategies.DDPStrategy` as an input to the Lightning {class}`~lightning.pytorch.trainer.trainer.Trainer` if the cluster environment is set up to facilitate Lightning's automatic detection of cluster variables. The main aspect that deserves user attention is configuring the job submission script and the relevant `Trainer` arguments (`num_nodes` and sometimes `devices`) correctly. @@ -28,7 +28,7 @@ A minimal SLURM example for doing DDP training with 2 nodes with 4 GPUs per node # ... set up (e.g. module load, activate Python env, etc) -# ... cluster specific set up such as network interface +# ... cluster specific set up such as network interface # (e.g. MASTER_PORT, MASTER_ADDR, NCCL_SOCKET_IFNAME) srun nequip-train -cn config.yaml ++trainer.num_nodes=${SLURM_NNODES} @@ -62,7 +62,7 @@ The `batch_size` configured under the dataloaders in the `data` [section of the As increasing the number of ranks (while holding the per-rank batch size constant) increases the effective batch size, you should consider adjusting other hyperparameters that you would typically adjust when raising the batch size, such as the learning rate (see [Lightning's docs](https://lightning.ai/docs/pytorch/stable/accelerators/gpu_faq.html#how-should-i-adjust-the-learning-rate-when-using-multiple-devices) for similar advice). -It may be helpful to use a combination of {mod}`omegaconf`'s [variable interpolation](https://omegaconf.readthedocs.io/en/latest/usage.html#variable-interpolation), [environment variable resolver](https://omegaconf.readthedocs.io/en/latest/custom_resolvers.html#oc-env) and NequIP's custom arithmetic resolver `int_div` to dynamically configure these parameters based on the runtime environment. +It may be helpful to use a combination of {mod}`omegaconf`'s [variable interpolation](https://omegaconf.readthedocs.io/en/latest/usage.html#variable-interpolation), [environment variable resolver](https://omegaconf.readthedocs.io/en/latest/custom_resolvers.html#oc-env) and NequIP's custom arithmetic resolver `int_div` to dynamically configure these parameters based on the runtime environment. For example, to get the world size as a SLURM environment variable and set the per-rank batch size as the desired effective global batch size divided by the world size, you can use: @@ -78,4 +78,4 @@ When using DDP, the {class}`torch.utils.data.distributed.DistributedSampler` may Be very careful when reporting validation or test metrics in DDP settings, as data duplication can lead to incorrect metrics. Either ensure data samples can be evenly distributed to all ranks, or perform validation/testing on a single rank. -For more details on handling validation in distributed settings, see [Lightning's documentation](https://lightning.ai/docs/pytorch/stable/common/lightning_module.html#test-loop). \ No newline at end of file +For more details on handling validation in distributed settings, see [Lightning's documentation](https://lightning.ai/docs/pytorch/stable/common/lightning_module.html#test-loop). diff --git a/docs/guide/accelerations/gpu_kernel_modifiers.md b/docs/guide/accelerations/gpu_kernel_modifiers.md index d1bfd31c..5d454799 100644 --- a/docs/guide/accelerations/gpu_kernel_modifiers.md +++ b/docs/guide/accelerations/gpu_kernel_modifiers.md @@ -18,4 +18,4 @@ NequIP GNN models support several GPU kernel modifiers that can significantly sp openequivariance cuequivariance -``` \ No newline at end of file +``` diff --git a/docs/guide/accelerations/precision.md b/docs/guide/accelerations/precision.md index da6edfef..513313fd 100644 --- a/docs/guide/accelerations/precision.md +++ b/docs/guide/accelerations/precision.md @@ -38,7 +38,7 @@ callbacks: callbacks: - _target_: nequip.train.callbacks.TF32Scheduler schedule: - 0: true # Enable TF32 for faster early training + 0: true # Enable TF32 for faster early training 100: false # Disable TF32 at epoch 100 for precise convergence ``` @@ -58,4 +58,4 @@ nequip-compile model.ckpt compiled_model.pt --tf32 ... nequip-compile model.ckpt compiled_model.pt --no-tf32 ... ``` -The default behavior is to compile without TF32, regardless of training settings. \ No newline at end of file +The default behavior is to compile without TF32, regardless of training settings. diff --git a/docs/guide/accelerations/pt2_compilation.md b/docs/guide/accelerations/pt2_compilation.md index 79c88b75..ed41fe9f 100644 --- a/docs/guide/accelerations/pt2_compilation.md +++ b/docs/guide/accelerations/pt2_compilation.md @@ -13,7 +13,7 @@ or model: _target_: allegro.model.AllegroModel compile_mode: compile - # other model hyperparameters + # other model hyperparameters ``` Note that `compile_mode` can only be `compile` (use [`torch.compile`](https://pytorch.org/docs/stable/generated/torch.compile.html)), or `eager` (no compilation used). If `compile_mode` is unspecified, it defaults to `eager`. It will take a bit of time (around a minute or more) for the model to be compiled with [`torch.compile`](https://pytorch.org/docs/stable/generated/torch.compile.html) before training proceeds, but the speed-ups are worth it. @@ -26,4 +26,4 @@ Train-time compilation will not work if any of the batch dimensions are all neve ```{warning} At present we advise against using train-time compilation on CPUs. As of PyTorch 2.6.0, there are known cases of **CPU**-specific train-time compilation issues for certain configurations of NequIP and Allegro models. Be cautious when trying to use train-time compilation with CPUs. If you encounter such issues, please open a GitHub issue. -``` \ No newline at end of file +``` diff --git a/docs/guide/configuration/config.md b/docs/guide/configuration/config.md index 9289427d..12ed4103 100644 --- a/docs/guide/configuration/config.md +++ b/docs/guide/configuration/config.md @@ -70,7 +70,7 @@ The full set of options are found in the documentation of the [underlying object `training_module` defines the {class}`~nequip.train.NequIPLightningModule` (or its subclasses). Users are directed to the [`nequip.train.NequIPLightningModule` API documentation](../../api/lightning_module.rst) to learn how to configure it. Usually the {class}`~nequip.train.EMALightningModule` is the right choice. The following important objects are configured as part of the `training_module`: - + ### `model` This section configures the model itself, including hyperparameters and the choice of architecture (for example, the NequIP message-passing E(3)-equivariant GNN, or the Allegro architecture). Refer to the [model documentation page](../../api/model.rst) to learn how to configure this section. @@ -79,7 +79,7 @@ The following important objects are configured as part of the `training_module`: ### `optimizer` and `lr_scheduler` - The `optimizer` can be any PyTorch-compatible optimizer. Options from PyTorch can be found in {mod}`torch.optim`. The {class}`~torch.optim.Adam` optimizer, for example, can be configured as follows: + The `optimizer` can be any PyTorch-compatible optimizer. Options from PyTorch can be found in {mod}`torch.optim`. The {class}`~torch.optim.Adam` optimizer, for example, can be configured as follows: ```yaml optimizer: _target_: torch.optim.Adam @@ -99,4 +99,3 @@ lr_scheduler: frequency: 1 ``` The `scheduler` is a PyTorch-compatible learning rate scheduler. Options from PyTorch can be found in {mod}`torch.optim.lr_scheduler`. - diff --git a/docs/guide/configuration/data.md b/docs/guide/configuration/data.md index fd9d7de8..100bd8d1 100644 --- a/docs/guide/configuration/data.md +++ b/docs/guide/configuration/data.md @@ -10,7 +10,7 @@ The data processing in NequIP follows this pipeline from raw files to model-read ``` This entire pipeline is coordinated and managed by a {class}`~nequip.data.datamodule.NequIPDataModule` object, which also: -- Manages train/val/test dataset splits +- Manages train/val/test dataset splits - Computes dataset statistics **Key Components:** @@ -151,4 +151,4 @@ Statistics are computed once during data setup, not during training. For advanced use cases, you should use the base {class}`~nequip.data.DataStatisticsManager` directly for more flexible configuration. See the [dataset statistics API documentation](../../api/data_stats.rst) for configuration options. -For guidance on using computed statistics to initialize model parameters, see [Training data statistics as hyperparameters](model.md/#training-data-statistics-as-hyperparameters). \ No newline at end of file +For guidance on using computed statistics to initialize model parameters, see [Training data statistics as hyperparameters](model.md/#training-data-statistics-as-hyperparameters). diff --git a/docs/guide/configuration/index.rst b/docs/guide/configuration/index.rst index 9a0c1ba7..bfe90d43 100644 --- a/docs/guide/configuration/index.rst +++ b/docs/guide/configuration/index.rst @@ -9,4 +9,4 @@ Learn how to configure NequIP models, data processing, and training. config data model - metrics \ No newline at end of file + metrics diff --git a/docs/guide/configuration/ml-data-flow-chart.svg b/docs/guide/configuration/ml-data-flow-chart.svg index 234d6757..e9cef0ed 100644 --- a/docs/guide/configuration/ml-data-flow-chart.svg +++ b/docs/guide/configuration/ml-data-flow-chart.svg @@ -1,67 +1,67 @@ - - + - - Raw Data File - .xyz, .npz, etc. - + - - + - - Dataset Object(s) - load raw data and parses into AtomicDataDicts - + - - + - - Data Transforms - applied to each data point - neighbor lists, type mapping, etc. - + - - + - - DataLoader - combine transformed data into batches - + - - + - - Model Training - nequip-train - \ No newline at end of file + diff --git a/docs/guide/configuration/model.md b/docs/guide/configuration/model.md index 4a916e8f..ade79bcd 100644 --- a/docs/guide/configuration/model.md +++ b/docs/guide/configuration/model.md @@ -27,7 +27,7 @@ Once statistics are configured, use the resolver syntax `${training_data_stats:s ```yaml model: _target_: nequip.model.NequIPGNNModel - + # Use computed statistics for model initialization avg_num_neighbors: ${training_data_stats:num_neighbors_mean} per_type_energy_shifts: ${training_data_stats:per_atom_energy_mean} @@ -37,7 +37,7 @@ model: The {class}`~nequip.data.CommonDataStatisticsManager` computes these essential statistics: - `num_neighbors_mean` - for edge normalization (`avg_num_neighbors`) -- `per_atom_energy_mean` - for energy shifts (`per_type_energy_shifts`) +- `per_atom_energy_mean` - for energy shifts (`per_type_energy_shifts`) - `forces_rms` - for energy scales (`per_type_energy_scales`) - `per_type_forces_rms` - for per-type energy scales (`per_type_energy_scales`) @@ -56,7 +56,7 @@ The statistic name must match exactly between the `stats_manager` configuration ## Energy Shifts & Scales -Both NequIP and Allegro models can be thought of as predicting an **unscaled** per-atom energy contribution $\tilde{E}_i$ where $i$ is an atom index. The per-atom energies are then computed as +Both NequIP and Allegro models can be thought of as predicting an **unscaled** per-atom energy contribution $\tilde{E}_i$ where $i$ is an atom index. The per-atom energies are then computed as $E_i = \alpha^{t_i} \tilde{E}_i + E_{0}^{t_i}$ @@ -69,7 +69,7 @@ It is recommended to use isolated atom energies (computed by the same reference model: _target_: nequip.model.NequIPGNNModel # other args - per_type_energy_shifts: + per_type_energy_shifts: C: 1.234 H: 2.345 O: 3.456 @@ -102,7 +102,7 @@ For practical molecular dynamics simulations, it may be favorable to train model ```yaml training_module: _target_: nequip.train.EMALightningModule - + # other `EMALightningModule` arguments model: @@ -112,12 +112,12 @@ training_module: pair_potential: _target_: nequip.nn.pair_potential.ZBL - units: metal - chemical_species: ${chemical_symbols} + units: metal + chemical_species: ${chemical_symbols} ``` `units` refer to LAMMPS unit system names, and can be `metal` (eV and Angstroms) or `real` (kcal/mol and Angstroms). Note that one must also specify `chemical_species`, which are a list of elements in the same order as the `type_names`. Here, the config snippet assumes that the list is defined elsewhere as `chemical_symbols` and is interpolating that variable (since the variable is also needed in configuring the `data` part of the config file). ```{warning} If you use ZBL, your data must agree with the system of units specified in the `units` parameter. -``` \ No newline at end of file +``` diff --git a/docs/guide/getting-started/index.rst b/docs/guide/getting-started/index.rst index e6702c78..4b10e9d5 100644 --- a/docs/guide/getting-started/index.rst +++ b/docs/guide/getting-started/index.rst @@ -8,4 +8,4 @@ Framework-level information for using NequIP. install workflow - files \ No newline at end of file + files diff --git a/docs/guide/getting-started/install.md b/docs/guide/getting-started/install.md index 20ab3299..62de5339 100644 --- a/docs/guide/getting-started/install.md +++ b/docs/guide/getting-started/install.md @@ -21,7 +21,7 @@ After installing `torch`, NequIP can be installed in the following ways. ```bash git clone --depth 1 https://github.com/mir-group/nequip.git cd nequip - pip install . + pip install . ``` 3. from source, using the latest `develop` branch @@ -29,7 +29,7 @@ After installing `torch`, NequIP can be installed in the following ways. git clone https://github.com/mir-group/nequip.git cd nequip git checkout develop - pip install . + pip install . ``` If you want to track your training runs with third-party services, like Weights and Biases, that are supported by [PyTorch Lightning's loggers](https://lightning.ai/docs/pytorch/stable/extensions/logging.html), you may need to install extra packages. For example, one needs to `pip install wandb` to use Lightning's {class}`~lightning.pytorch.loggers.WandbLogger`. @@ -55,4 +55,4 @@ To run the full tests, including a set of longer/more intensive integration test pytest tests/ ``` -If a GPU is present, the unit tests will use it. \ No newline at end of file +If a GPU is present, the unit tests will use it. diff --git a/docs/guide/guide.rst b/docs/guide/guide.rst index 3fb20073..be709d81 100644 --- a/docs/guide/guide.rst +++ b/docs/guide/guide.rst @@ -20,4 +20,4 @@ This guide covers using NequIP, from initial setup to advanced optimization tech configuration/index training-techniques/index accelerations/index - reference/index \ No newline at end of file + reference/index diff --git a/docs/guide/reference/faq.md b/docs/guide/reference/faq.md index 3c155e8b..b65fc03f 100644 --- a/docs/guide/reference/faq.md +++ b/docs/guide/reference/faq.md @@ -9,7 +9,7 @@ ## Units **Q**: What units do `nequip` framework models use? - **A**: `nequip` has no preferred system of units and uses the units of the data provided. Model inputs, outputs, error metrics, and all other quantities follow the units of the dataset. Users **must** use consistent input and output units. For example, if the length unit is Å and the energy labels are in eV, the force predictions from the model will be in eV/Å. The provided force labels should hence also be in eV/Å. + **A**: `nequip` has no preferred system of units and uses the units of the data provided. Model inputs, outputs, error metrics, and all other quantities follow the units of the dataset. Users **must** use consistent input and output units. For example, if the length unit is Å and the energy labels are in eV, the force predictions from the model will be in eV/Å. The provided force labels should hence also be in eV/Å. ```{warning} `nequip` cannot and does not check the consistency of units in inputs you provide, and it is your responsibility to ensure consistent treatment of input and output units @@ -71,4 +71,4 @@ Importing, restarting, or migrating models or training runs from pre-0.7.0 versi **Q**: Are losses still sometimes in normalized internal units? - **A**: No, in `nequip >= 0.7.0`, the loss components are all in physical units. \ No newline at end of file + **A**: No, in `nequip >= 0.7.0`, the loss components are all in physical units. diff --git a/docs/guide/reference/index.rst b/docs/guide/reference/index.rst index 2136440c..b43364f6 100644 --- a/docs/guide/reference/index.rst +++ b/docs/guide/reference/index.rst @@ -8,4 +8,4 @@ Reference materials, conventions, troubleshooting, and frequently asked question conventions faq - troubleshoot \ No newline at end of file + troubleshoot diff --git a/docs/guide/training-techniques/index.rst b/docs/guide/training-techniques/index.rst index df980c32..defc23cb 100644 --- a/docs/guide/training-techniques/index.rst +++ b/docs/guide/training-techniques/index.rst @@ -8,4 +8,4 @@ Advanced training methods and techniques. fine_tuning multitask_training - wandb_example \ No newline at end of file + wandb_example diff --git a/docs/integrations/ase.md b/docs/integrations/ase.md index 1efd9ca3..ae73b317 100644 --- a/docs/integrations/ase.md +++ b/docs/integrations/ase.md @@ -1,7 +1,7 @@ # ASE ## Introduction -The [Atomic Simulation Environment (ASE)](https://wiki.fysik.dtu.dk/ase/) is a popular Python package providing a framework for working with atomic data, reading and writing common formats, and running various simulations and calculations. +The [Atomic Simulation Environment (ASE)](https://wiki.fysik.dtu.dk/ase/) is a popular Python package providing a framework for working with atomic data, reading and writing common formats, and running various simulations and calculations. INequIP provides the {class}`~nequip.ase.NequIPCalculator` for integration with ASE. See the [ASE calculator API](../api/ase.rst) for detailed documentation. ## Creating an ASE Calculator @@ -67,9 +67,9 @@ import torch # Initialize the nequip calculator calculator = NequIPCalculator.from_compiled_model( - compile_path="path/to/compiled_model.nequip.pt2", - chemical_symbols=["Si"], - device="cuda" if torch.cuda.is_available() else "cpu", + compile_path="path/to/compiled_model.nequip.pt2", + chemical_symbols=["Si"], + device="cuda" if torch.cuda.is_available() else "cpu", ) # use GPUs if available # Range of scaling factors for lattice constant @@ -132,7 +132,7 @@ force_max = 0.05 # run until the forces are smaller than this in eV/A # Initialize Nequip ASE Calculator from checkpoint calculator = NequIPCalculator.from_compiled_model( - compile_path=compile_path, + compile_path=compile_path, device="cuda" if torch.cuda.is_available() else "cpu", ) # use GPUs if available @@ -182,4 +182,4 @@ for struct in tqdm(list_of_structs_to_relax, desc="Relaxing"): relaxed_df = pd.DataFrame(relax_results).T relaxed_df.to_csv("relaxed_structures.csv") -``` \ No newline at end of file +``` diff --git a/docs/integrations/lammps/mliap.md b/docs/integrations/lammps/mliap.md index 5436b925..17053515 100644 --- a/docs/integrations/lammps/mliap.md +++ b/docs/integrations/lammps/mliap.md @@ -104,4 +104,4 @@ Replace `output.nequip.lmp.pt` with the path to your ML-IAP interface file, `sys Example command to run the LAMMPS simulation with Kokkos GPU support: ```bash srun -n 1 /path/to/lammps/build/lmp -in in.lammps -k on g 1 -sf kk -pk kokkos newton on neigh half -``` \ No newline at end of file +``` diff --git a/docs/integrations/lammps/pair_styles.md b/docs/integrations/lammps/pair_styles.md index 1cdf220e..747006df 100644 --- a/docs/integrations/lammps/pair_styles.md +++ b/docs/integrations/lammps/pair_styles.md @@ -34,4 +34,4 @@ nequip-compile \ ## Installation and Usage -For detailed installation instructions, compilation of LAMMPS with the pair styles, and usage examples, please refer to the README in the [`pair_nequip_allegro`](https://github.com/mir-group/pair_nequip_allegro) repository. \ No newline at end of file +For detailed installation instructions, compilation of LAMMPS with the pair styles, and usage examples, please refer to the README in the [`pair_nequip_allegro`](https://github.com/mir-group/pair_nequip_allegro) repository. diff --git a/docs/requirements.txt b/docs/requirements.txt index 5275634b..c04b06f3 100644 --- a/docs/requirements.txt +++ b/docs/requirements.txt @@ -2,4 +2,4 @@ myst-parser sphinx_copybutton furo torch -pygments-lammps \ No newline at end of file +pygments-lammps diff --git a/tests/integration/test_compile.py b/tests/integration/test_compile.py new file mode 100644 index 00000000..86d312da --- /dev/null +++ b/tests/integration/test_compile.py @@ -0,0 +1,33 @@ +import pytest +import torch +from pathlib import Path +from nequip.scripts.compile import main + + +@pytest.mark.parametrize( + "device", ["cpu"] + (["cuda"] if torch.cuda.is_available() else []) +) +@pytest.mark.parametrize( + "mode,extension", [("aotinductor", ".nequip.pt2"), ("torchscript", ".nequip.pth")] +) +def test_compile_modes(tmp_path_factory, device, mode, extension): + tmp_path = tmp_path_factory.mktemp("nequip_compiled") + output_model_name = f"mir-group__NequIP-OAM-L__0.1{extension}" + output_path = Path(tmp_path) / output_model_name + + args = [ + "nequip.net:mir-group/NequIP-OAM-L:0.1", + str(output_path), + "--mode", + mode, + "--device", + device, + ] + + if mode == "aotinductor": + args.extend(["--target", "ase"]) + + main(args=args) + + assert output_path.exists() + assert output_path.is_file() diff --git a/tests/unit/model/test_pair/zbl_data.lmps b/tests/unit/model/test_pair/zbl_data.lmps index 6afa6e0f..e7f79bd1 100644 --- a/tests/unit/model/test_pair/zbl_data.lmps +++ b/tests/unit/model/test_pair/zbl_data.lmps @@ -43,4 +43,4 @@ label Ziloop jump SELF Zjloop next Zi -jump SELF Ziloop \ No newline at end of file +jump SELF Ziloop