Draft: Julia-based MACE interface for more efficient ring-polymer simulations

Project.toml

@@ -2,7 +2,7 @@ name = "NQCDynamics"
 uuid = "36248dfb-79eb-4f4d-ab9c-e29ea5f33e14"
 authors = ["James <[email protected]>"]
 
-version = "0.13.7"
+version = "0.14.0"
 
 
 [deps]
@@ -48,6 +48,9 @@ UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 UnitfulAtomic = "a7773ee8-282e-5fa2-be4e-bd808c38a91a"
 
+[weakdeps]
+MACEModel = "29bafc4c-4f38-420f-a153-8a5a5e0bd5f6"
+
 [compat]
 AdvancedHMC = "0.5, 0.6"
 AdvancedMH = "0.8"
@@ -86,7 +89,7 @@ UnPack = "1"
 UnicodePlots = "2, 3"
 Unitful = "1"
 UnitfulAtomic = "1"
-julia = "1.7"
+julia = "≥1.9"
 
 [extras]
 CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
@@ -99,11 +102,11 @@ JuLIP = "945c410c-986d-556a-acb1-167a618e0462"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 MKL = "33e6dc65-8f57-5167-99aa-e5a354878fb2"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
-PyCall = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0"
+PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "SafeTestsets", "CSV", "DataFrames", "DiffEqNoiseProcess", "FiniteDiff", "DiffEqDevTools", "Logging", "MKL", "PyCall", "JuLIP", "Symbolics", "Statistics", "Plots", "JLD2"]
+test = ["Test", "SafeTestsets", "CSV", "DataFrames", "DiffEqNoiseProcess", "FiniteDiff", "DiffEqDevTools", "Logging", "MKL", "PythonCall", "JuLIP", "Symbolics", "Statistics", "Plots", "JLD2"]

docs/Project.toml

@@ -11,13 +11,13 @@ DocumenterCitations = "daee34ce-89f3-4625-b898-19384cb65244"
 FileIO = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+MACEModels = "29bafc4c-4f38-420f-a153-8a5a5e0bd5f6"
 NNInterfaces = "07253886-9aba-4a5d-b3d5-29b8b100a664"
 NQCBase = "78c76ebc-5665-4934-b512-82d81b5cbfb7"
 NQCDistributions = "657014a1-bd25-4aa2-92cb-cbcede0310ad"
 NQCDynamics = "36248dfb-79eb-4f4d-ab9c-e29ea5f33e14"
 NQCModels = "c814dc9f-a51f-4eaf-877f-82eda4edad48"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
-PyCall = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0"
 PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
 RingPolymerArrays = "81cd853e-d334-476c-b156-f95acc8a32dc"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"

docs/make.jl

@@ -1,6 +1,6 @@
 using Documenter
 using DocumenterCitations
-using NQCBase, NQCModels, NQCDistributions, NQCDynamics
+using NQCBase, NQCModels, NQCDistributions, NQCDynamics, MACEModels
 using CubeLDFAModel, NNInterfaces
 
 bib = CitationBibliography(joinpath(@__DIR__, "references.bib"))
@@ -15,7 +15,7 @@ end
 @time makedocs(;
     plugins=[bib],
     sitename="NQCDynamics.jl",
-    modules=[NQCDynamics, NQCDistributions, NQCModels, NQCBase, CubeLDFAModel],
+    modules=[NQCDynamics, NQCDistributions, NQCModels, NQCBase, CubeLDFAModel, MACEModels],
     doctest=false,
     format=Documenter.HTML(
         prettyurls=get(ENV, "CI", nothing) == "true",
@@ -32,9 +32,9 @@ end
         "NQCModels.jl" => Any[
             "NQCModels/overview.md"
             "NQCModels/analyticmodels.md"
-            "NQCModels/ase.md"
-            "NQCModels/neuralnetworkmodels.md"
+            "NQCModels/fullsizemodels.md"
             "NQCModels/frictionmodels.md"
+            "NQCModels/neuralnetworkmodels.md"
         ]
         "NQCDistributions.jl" => Any[
             "NQCDistributions/overview.md"

docs/src/NQCModels/ase.md → docs/src/NQCModels/fullsizemodels.md

@@ -1,15 +1,19 @@
 ```@setup logging
-@info "Expanding src/NQCModels/ase.md..."
+@info "Expanding src/NQCModels/fullsizemodels.md..."
 start_time = time()
 ```
-# ASE interface
+
+# Full dimensional model library
+
+## ASE interface
 
 The easiest way to obtain potentials and forces from established codes is to
 use the interfaces implemented in [ASE](https://wiki.fysik.dtu.dk/ase/).
 
 We provide the [`AdiabaticASEModel`](@ref) which wraps an ASE atoms object and its
 associated calculator to implement the required [`potential`](@ref) and
 [`derivative`](@ref) functions.
+Several examples for connecting common machine-learning interatomic potentials to NQCModels.jl through the ASE interface are shown in the [MLIP examples](@ref) section.
 
 !!! note
 
@@ -21,11 +25,12 @@ associated calculator to implement the required [`potential`](@ref) and
     pre-installed Python or install its own.
     Refer to the [PyCall README](https://github.com/JuliaPy/PyCall.jl) for installation
     and configuration instructions.
-    
-## Example
+
+### Example
 
 First, it is necessary to import `ase` and create the `ase.Atoms` object and attach
 the desired calculator. This works exactly as in Python:
+
 ```@example ase
 using PythonCall: pyimport, pylist
 using ASEconvert
@@ -39,24 +44,63 @@ nothing # hide
 
 Next, the [`AdiabaticASEModel`](@ref) is created by passing the `ase.Atoms` object directly
 to the model:
+
 ```@repl ase
 using NQCModels
 model = AdiabaticASEModel(h2)
 ```
+
 Now the model can be used in the same way as any of the previously introduced
 analytic models.
+
 ```@repl ase
 potential(model, rand(3, 2))
 derivative(model, rand(3, 2))
 ```
 
-!!! tip 
+!!! tip
 
     In theory, this should work with any of the ASE calculators that correctly implement
     the `get_potential_energy` and `get_forces` functions. For instance, you can use
     [SchNetPack (SPK)](https://github.com/atomistic-machine-learning/schnetpack) by
     passing their ASE calculator to the `AdiabaticASEModel`.
-    Take a look at [Neural network models](@ref) to learn more.
+    Take a look at [MLIP examples](@ref) to learn more.
+
+## MACE interface
+
+[MACEModels.jl](https://github.com/NQCD/MACEModels.jl) is an interface to the [MACE](https://github.com/ACEsuit/mace) code. The package attempts to improve calculation speeds by directly interfacing with `MACE`, instead of going through `ase`.
+
+More information on how to use MACE models is available in the [MACEModels.jl](@ref) API documentation. 
+
+!!! tip
+
+      To make the installation of a python environment to execute MACE in optional, the `MACEModel` is provided in a separate package, which needs to be installed with `]add MACEModels`. 
+
+### Example
+
+```julia
+using NQCModels, PyCall, NQCDynamics.jl, MACEModels
+
+ensemble_paths = [
+    "model01.model",
+    "model02.model",
+    "model03.model",
+]
+
+structure = Atoms([:N, :H, :H, :H]) # Define the structure
+cell = InfiniteCell() # Set periodicity (none in this case)
+
+# Create the MACE model
+model = MACEModel(
+    atoms, # Atoms (used for neighbour lists)
+    cell, # Cell (used for neighbour lists)
+    ensemble_paths; # Vector containing paths to one or more models
+    device = "cpu", # Device (or vector of devices) to use
+    default_dtype = Float32, # Data type of the models (Float32 or Float64)
+    batch_size = 1, # Number of structures to evaluate at once (improves overall throughput)
+)
+```
+
 ```@setup logging
 runtime = round(time() - start_time; digits=2)
 @info "...done after $runtime s."

docs/src/NQCModels/neuralnetworkmodels.md

@@ -2,7 +2,10 @@
 @info "Expanding src/NQCModels/neuralnetworkmodels.md..."
 start_time = time()
 ```
-# Neural network models
+
+# MLIP Examples
+
+## SchNet (SchNetPack) models
 
 Using the [ASE interface](@ref) we can directly use models trained using
 [SchNetPack](https://github.com/atomistic-machine-learning/schnetpack).
@@ -23,8 +26,9 @@ provide the relative path.
 
 First we load the model into an `ase` calculator and attach it to our diatomic
 hydrogen molecule.
+
 ```julia
-using PyCall
+using PythonCall
 
 ase = pyimport("ase")
 spkutils = pyimport("schnetpack.utils")
@@ -40,10 +44,11 @@ h2.set_calculator(calc)
 
 We can obtain the energies and forces from `ase` directly in the usual way, converting
 them to atomic units using [UnitfulAtomic](https://github.com/sostock/UnitfulAtomic.jl).
+
 ```julia-repl
 using Unitful, UnitfulAtomic;
-austrip(h2.get_total_energy() * u"eV")
-austrip.(h2.get_forces() .* u"eV/Å")
+austrip(pyconvert(Float64,h2.get_total_energy()) * u"eV")
+austrip.(pyconvert(Matrix{Float64}, h2.get_forces()) .* u"eV/Å")
 ```
 
 !!! warning
@@ -54,6 +59,7 @@ austrip.(h2.get_forces() .* u"eV/Å")
 Then, we can convert the ASE output into the format used in NQCModels,
 which makes it possible to use the SchNet model e.g. for molecular dynamics calculations
 within NQCDynamics.jl:
+
 ```julia-repl
 using NQCModels;
 model = AdiabaticASEModel(h2);
@@ -63,6 +69,34 @@ r = [0 0; 0 0; 0 ustrip(auconvert(0.74u"Å"))]
 potential(model, r)
 derivative(model, r)
 ```
+
+## MACE models
+
+!!! warning
+
+      A more direct interface for the use of MACE models is now available. @Alex link to example page and API docs here.
+
+      The following example is still valid but the new interface is recommended for new projects.
+
+The following example shows how to connect a trained [MACE](https://github.com/ACEsuit/mace) model to NQCDynamics.jl using the [ASE interface](@ref).
+
+```julia
+using PythonCall, NQCModels
+
+ase_io = pyimport("ase.io") # make sure ase is installed in your environment first
+mace_module = pyimport("mace.calculators") # make sure mace is installed in your environment first
+
+structure = ase_io.read("starting_structure.xyz")
+mace_calculator = mace_module.MACECalculator(
+    "path/to/model.model",
+    device = "cpu", # or cuda to run on GPU
+    default_dtype = "float32" # if this was set during training as well
+)
+
+structure.set_calculator(mace_calculator)
+model = AdiabaticASEModel(structure)
+```
+
 ```@setup logging
 runtime = round(time() - start_time; digits=2)
 @info "...done after $runtime s."

docs/src/NQCModels/overview.md

@@ -2,13 +2,20 @@
 @info "Expanding src/NQCModels/overview.md..."
 start_time = time()
 ```
+
 # NQCModels.jl
 
+!!! details "Overview of all model types currently implemented."
+
+    ```@raw html
+    <img src="../assets/figures/Model_types.svg"  width=100% caption="Overview of all Model types currently implemented.">
+    ```
+
 To perform nonadiabatic molecular dynamics simulations, it is necessary to define
 the system Hamiltonian.
 For simple models, this often comes in the form of small matrix in the diabatic
 representation but equally the electronic Hamiltonian could be obtained directly
-from *ab initio* electronic structure theory.
+from _ab initio_ electronic structure theory.
 
 `NQCModels.jl` is a package that aims to provide a common interface
 for defining these models that is flexible enough to allow for a wide range
@@ -30,9 +37,11 @@ and [`DiabaticModel`](@ref NQCModels.DiabaticModels.DiabaticModel).
 The [`AdiabaticModel`](@ref NQCModels.AdiabaticModels.AdiabaticModel)
 is used for adiabatic dynamics, providing only the potential
 and force used in classical mechanics.
-The [`DiabaticModel`](@ref NQCModels.DiabaticModels.DiabaticModel) is used for nonadiabatic dynamics, 
+The [`DiabaticModel`](@ref NQCModels.DiabaticModels.DiabaticModel) is used for nonadiabatic dynamics,
 where the potential is instead a `Hermitian` matrix.
 
+## Using `Model`s
+
 In the [Getting started](@ref) section we briefly touched on how the
 [`AdiabaticModel`](@ref NQCModels.AdiabaticModels.AdiabaticModel)
 works and introduced one of the included models.
@@ -87,25 +96,18 @@ with respect to each degree of freedom.
 In this case, the `Matrix` has `size = (1, 2)`, but it should be clear how this can extend
 to arbitrary numbers of atoms and degrees of freedom for complex models.
 
-The models currently available can be seen in type tree of the
-[`Model`](@ref NQCModels.Model) below.
-The leaves of the tree are the concrete models, whereas each branch is one of the abstract
-types.
-Each of these models can be seen in the [Analytic model library](@ref) and
-many shall return later when we investigate the dynamics methods.
+## Included models
 
-```@example
-import AbstractTrees # hide
-import InteractiveUtils: subtypes # hide
-import NQCModels: Model # hide
-AbstractTrees.children(x::Type) = subtypes(x) # hide
-AbstractTrees.print_tree(Model) # hide
-```
+`NQCModels.jl` includes both analytical models as well as interfaces to connect to full-dimensional potential energy surfaces.
+An overview of the implemented analytical models can be seen in the [Analytic model library](@ref) and
+many shall return later when we investigate the dynamics methods.
+An overview of interfaces to full-dimensional potential energy surfaces can be seen in the [Full dimensional model library](@ref).
 
 !!! note "Contributing new models"
 
     To learn more about NQCModels.jl and learn how to implement new models,
     visit the [developer documentation](@ref devdocs-model).
+
 ```@setup logging
 runtime = round(time() - start_time; digits=2)
 @info "...done after $runtime s."

docs/src/api/NQCDynamics/structure.md

@@ -2,20 +2,21 @@
 @info "Expanding src/api/NQCDynamics/structure.md..."
 start_time = time()
 ```
+
 # Structure
-This submodule contains utility functions to analyse and modify atomic structure data, such as interatomic distances, centres of mass, both with and without support for periodic boundary conditions. 
 
-These functions can be used to build more sophisticated output functions, or for basic analysis of simulation results in post. 
+This submodule contains utility functions to analyse and modify atomic structure data, such as interatomic distances, centres of mass, both with and without support for periodic boundary conditions.
+
+These functions can be used to build more sophisticated output functions, or for basic analysis of simulation results in post.
 
 This module **doesn't contain**:
+
 - Basic definitions of atomic structures (e.g. `Atoms`, `PeriodicCell`, ...). These are defined in [`NQCBase`](@ref).
 - Functions to generate atomic structures. These should be added to [`NQCDynamics.InitialConditions`](@ref).
-- Analysis functions for specific systems (e.g. molecules on surfaces). These should be added to [`NQCDynamics.Analysis`](@ref). 
-
+- Analysis functions for specific systems (e.g. molecules on surfaces). These should be added to [`NQCDynamics.Analysis`](@ref).
 
 ## Method reference
 
-
 ```@autodocs
 Modules=[NQCDynamics.Structure]
 ```

docs/src/api/NQCModels/mace.md

@@ -0,0 +1,16 @@
+```@setup logging
+@info "Expanding src/api/NQCModels/mace.md..."
+start_time = time()
+```
+
+# MACEModels.jl
+
+```@autodocs
+Modules = [MACEModels]
+Private = true
+```
+
+```@setup logging
+runtime = round(time() - start_time; digits=2)
+@info "...done after $runtime s."
+```