Interface re-design: state and Lux model (#22)

Co-authored-by: Christoph Ortner <[email protected]> Co-authored-by: Teemu Järvinen <[email protected]>
JuliaMolSim · Jul 5, 2024 · f15dd7f · f15dd7f
1 parent 758dd4a
commit f15dd7f
Show file tree

Hide file tree

Showing 14 changed files with 868 additions and 151 deletions.
diff --git a/Project.toml b/Project.toml
@@ -8,6 +8,7 @@ AtomsBase = "a963bdd2-2df7-4f54-a1ee-49d51e6be12a"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
+UnitfulAtomic = "a7773ee8-282e-5fa2-be4e-bd808c38a91a"
 
 [compat]
 AtomsBase = "0.3"

diff --git a/docs/make.jl b/docs/make.jl
@@ -16,7 +16,10 @@ makedocs(;
     ),
     pages=[
         "Home" => "index.md",
-        "Interface Definition" => "interface-definition.md",
+        "Interface" => "interface.md",
+        "Utilities" => "utilities.md",
+        "Example" => "example.md",
+        "Index" => "api.md",
     ],
 )
 

diff --git a/docs/src/api.md b/docs/src/api.md
@@ -0,0 +1,8 @@
+ # Index
+
+```@index
+```
+
+```@autodocs
+Modules = [AtomsCalculators, AtomsCalculators.AtomsCalculatorsTesting, AtomsCalculators.Testing]
+```
diff --git a/docs/src/interface-definition.md → docs/src/example.md b/docs/src/interface-definition.md → docs/src/example.md
@@ -1,54 +1,11 @@
-# Interface Definition
-
-There are two alternative ways to call the interface: using functions `potential_energy`, `forces` and [virial](https://en.wikipedia.org/wiki/Virial_stress), or using `calculate`
-function together with `Energy`, `Forces` and `Virial`.
-
-Individual calls are implemented by dispatching `AtomsCalculators` functions
-
-- `AtomsCalculators.potential_energy` for potential energy calculation
-- `AtomsCalculators.forces` for allocating force calculation and/or...
-- `AtomsCalculators.forces!` for non-allocating force calculation
-- `AtomsCalculators.virial` for [virial](https://en.wikipedia.org/wiki/Virial_stress) calculation
-
-The `calculate` interface is implemented by dispatching to
-
-- `AtomsCalculators.calculate` using `AtomsCalculators.Energy()` as the first argument for energy calculation
-- `AtomsCalculators.calculate` using `AtomsCalculators.Forces()` as the first argument for forces calculaton
-- `AtomsCalculators.calculate` using `AtomsCalculators.Virial()` as the first argument for virial calculation
-
-You do not need to implement all of these by yourself. There is macro that will help implement the other calls. 
-
-Each of the individual calls have two common inputs: `AtomsBase.AbstractSystem` compatible structure and a `calculator` that incudes details of the calculation method. Calculate interface has additionally the type of calculation as the first input. You can tune calculation by passing keyword arguments, which can be ignored, but they need to be present in the function definition.
-
-`potential_energy`, `forces`, `forces!` and `virial`:
-
-- First input is `AtomsBase.AbstractSystem` compatible structure
-- Second input is `calculator` structure
-- Method has to accept keyword arguments (they can be ignored)
-- Non-allocating force call `force!` has an AbstractVector as the first input, to which the evaluated force values are stored (look for more details below)
-
-`calculate`:
-
-- First input is either `Energy()`, `Forces()` or `Virial()`
-- Second is `AtomsBase.AbstractSystem` compatible structure
-- Third is `calculator` structure
-- Method has to accept keyword arguments (they can be ignored)
-
-## Output
-
-Outputs for the functions need to have following properties
-
-- Energy is a subtype of `Number` that has a unit with dimensions of energy (mass * length^2 / time^2)
-- Force output is a subtype of `AbstractVector` with element type also a subtype of AbstractVector (length 3 in 3D) and unit with dimensions of force (mass * length / time^2). With additional property that it can be reinterpret as a matrix
-- Virial is a square matrix (3x3 in 3D) that has units of force times length or energy
-- Calculate methods return a [NamedTuple](https://docs.julialang.org/en/v1/base/base/#Core.NamedTuple) that uses keys `:energy`, `:forces` and `:virial` to identify the results, which have the types defined above
-
 ## Implementing the interface
 
+**Note, this section is partly outdated!**
+
 You can either implement both of the calls e.g. for energy
 
 `AtomsCalculators.potential_energy(system, calculator; kwargs...)` and
-`AtomsCalculators(AtomsCalculators.Energy(), system, calculator; kwargs...)`
+`AtomsCalculators.calculate(AtomsCalculators.Energy(), system, calculator, ps=nothing, st=nothing; kwargs...)`
 
 ### Example implementations
 
@@ -72,7 +29,14 @@ end
 Completely equivalent implementation is
 
 ```julia
-AtomsCalculators.@generate_interface function AtomsCalculators.calculate(::AtomsCalculators.Energy, system, calculator::MyType; kwargs...)
+AtomsCalculators.@generate_interface function AtomsCalculators.calculate(
+    ::AtomsCalculators.Energy, 
+    system, 
+    calculator::MyType,
+    ps=nothing,
+    st=nothing; 
+    kwargs...
+)
     # we can ignore kwargs... or use them to tune the calculation
     # or give extra information like pairlist
 
@@ -96,12 +60,19 @@ end
 Equivalent implementation is
 
 ```julia
-AtomsCalculators.@generate_interface function AtomsCalculators.calculate(::AtomsCalculators.Virial, system, calculator::MyType; kwargs...)
+AtomsCalculators.@generate_interface function AtomsCalculators.calculate(
+    ::AtomsCalculators.Virial, 
+    system, 
+    calculator::MyType,
+    ps=nothing,
+    st=nothing;
+    kwargs...
+)
     # we can ignore kwargs... or use them to tune the calculation
     # or give extra information like pairlist
 
     # add your own definition here
-    return ( virial = zeros(3,3) * u"eV", )
+    return ( virial = zeros(3,3) * u"eV", state=nothing)
 end
 ```
 
@@ -128,12 +99,19 @@ Same way `AtomsCalculators.promote_force_type(system, calculator)` creates a for
 Alternatively the definition could have been done with
 
 ```julia
-AtomsCalculators.@generate_interface function AtomsCalculators.calculate(::AtomsCalculators.Forces, system, calculator::MyType; kwargs...)
+AtomsCalculators.@generate_interface function AtomsCalculators.calculate(
+    ::AtomsCalculators.Forces, 
+    system, 
+    calculator::MyType,
+    ps=nothing,
+    st=nothing;
+    kwargs...
+)
     # we can ignore kwargs... or use them to tune the calculation
     # or give extra information like pairlist
 
     # add your own definition
-    return ( forces = zeros(AtomsCalculators.promote_force_type(system, calculator), length(system)), )
+    return ( forces = zeros(AtomsCalculators.promote_force_type(system, calculator), length(system)), state=nothing )
 end
 ```
 

diff --git a/docs/src/index.md b/docs/src/index.md
@@ -6,16 +6,11 @@ CurrentModule = AtomsCalculators
 
 Documentation for [AtomsCalculators](https://github.com/JuliaMolSim/AtomsCalculators.jl).
 
-This is a calculation interface for [AtomsBase](https://github.com/JuliaMolSim/AtomsBase.jl) meant to extend the interface for calculators.
-At this moment this is a preliminary version with molecular dynamics and various geometry optimizations
-in mid. In the future there can be more functionality too.
+This package specifies the calculation interface for the [AtomsBase](https://github.com/JuliaMolSim/AtomsBase.jl) ecosystem.
+At this moment it provides a first version of a molecular mechanics interface (energy, forces, virial). Suggestions for extensions are welcome. 
+
+
 
 
-## Index
 
-```@index
-```
 
-```@autodocs
-Modules = [AtomsCalculators, AtomsCalculators.AtomsCalculatorsTesting]
-```
diff --git a/docs/src/intdocback.md b/docs/src/intdocback.md
@@ -0,0 +1,73 @@
+# Previous Version of Interface Docs
+
+There are two alternative ways to call the interface: using functions `potential_energy`, `forces` and [virial](https://en.wikipedia.org/wiki/Virial_stress), or using `calculate`
+function together with `Energy`, `Forces` and `Virial`.
+
+Individual calls are implemented by dispatching `AtomsCalculators` functions
+
+- `AtomsCalculators.potential_energy` for potential energy calculation
+- `AtomsCalculators.forces` for allocating force calculation and/or...
+- `AtomsCalculators.forces!` for non-allocating force calculation
+- `AtomsCalculators.virial` for [virial](https://en.wikipedia.org/wiki/Virial_stress) calculation
+
+The `calculate` interface is implemented by dispatching to
+
+- `AtomsCalculators.calculate` using `AtomsCalculators.Energy()` as the first argument for energy calculation
+- `AtomsCalculators.calculate` using `AtomsCalculators.Forces()` as the first argument for forces calculaton
+- `AtomsCalculators.calculate` using `AtomsCalculators.Virial()` as the first argument for virial calculation
+
+You do not need to implement all of these by yourself. There is macro that will help implement the other calls. 
+
+Each of the individual calls have two common inputs: `AtomsBase.AbstractSystem` compatible structure and a `calculator` that incudes details of the calculation method. Calculate interface has additionally the type of calculation as the first input. You can tune calculation by passing keyword arguments, which can be ignored, but they need to be present in the function definition.
+
+## High-level vs Low-level Interface
+
+The philosophy behing AtomsCalculators.jl is to provide high-level calls with default 
+parameters through `potential_energy`, `forces`, `forces!` and `virial` and low-level calls 
+with user-specifiable parameters through `calculate`. 
+
+The low-level `calculate` interface follows the [Lux](https://lux.csail.mit.edu/stable/) model for parameters and state. This means that calculators are **immutable** structs that are passed 
+to the `calculate` function together with `parameters` and `state`. All calculations then return 
+an output and a state.
+
+When implementing a new calculator, two paths can be chosen:
+
+1. **Providing a full low level implementation**
+    In this setup, one implements `calculate` for the various type of 
+    calculations to be supported, with full handling of parameters and state. **On top of that**, 
+    one should provide sensible **defaults** for calls with parameters and state set to `nothing`. 
+    These will be used to generate the high-level calls.
+2. **Only implementing high-level calls**
+    In this case, one only implements the high-level `potential_energy`, `forces` and `virial` 
+    (or some of them). Appropriate low-level implementation for `calculate` is then automatically 
+    generated by using and returning dummy state and parameters set to `nothing`.
+
+Note that when using the high-level approach, fixed parameters are allowed to be bundled inside the calculator.
+
+### Method Signatures
+
+`potential_energy`, `forces`, `forces!` and `virial`:
+
+- First input is `AtomsBase.AbstractSystem` compatible structure
+- Second input is `calculator` structure
+- Method has to accept keyword arguments (they can be ignored)
+- Non-allocating force call `force!` has an AbstractVector as the first input, to which the evaluated force values are stored (look for more details below)
+
+`calculate`:
+
+- First input is either `Energy()`, `Forces()` or `Virial()`
+- Second is `AtomsBase.AbstractSystem` compatible structure
+- Third is `calculator` structure
+- Fourth is `parameters`
+- Fifth is `state`
+- Method has to accept keyword arguments (they can be ignored)
+
+## Output
+
+Outputs for the functions need to have following properties
+
+- Energy is a subtype of `Number` that has a unit with dimensions of energy (mass * length^2 / time^2)
+- Force output is a subtype of `AbstractVector` with element type also a subtype of AbstractVector (length 3 in 3D) and unit with dimensions of force (mass * length / time^2). With additional property that it can be reinterpret as a matrix
+- Virial is a square matrix (3x3 in 3D) that has units of force times length or energy
+- Calculate methods return a [NamedTuple](https://docs.julialang.org/en/v1/base/base/#Core.NamedTuple) that uses keys `:energy`, `:forces` and `:virial` to identify the results, which have the types defined above, and additionally returns the updated `:state`.
+