WIP: Prototypes for Calculators

Project.toml

@@ -9,6 +9,7 @@ PeriodicTable = "7b2266bf-644c-5ea3-82d8-af4bbd25a884"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 UnitfulAtomic = "a7773ee8-282e-5fa2-be4e-bd808c38a91a"
 

docs/src/calculations-interface.md

@@ -0,0 +1,172 @@
+# Calculation Interface
+
+AtomsBase includes a calculation interface meant ease implement 
+
+
+## Definitions
+
+There are three main targets for calculations `potential_energy`, `forces` and [virial](https://en.wikipedia.org/wiki/Virial_stress). 
+
+Individual calls are implemented by overloading `AtomsBase` functions
+
+- `AtomsBase.potential_energy` for potential energy calculation
+- `AtomsBase.forces` for allocating force calculation and/or...
+- `AtomsBase.forces!` for non-allocating force calculation
+- `AtomsBase.virial` for [virial](https://en.wikipedia.org/wiki/Virial_stress) calculation
+
+You do not implement all of these. To implement force calculation you need to implement either allocating or non-allocating force call - see below for more details on how to implement force calculation.
+
+Each call has will have two inputs `AtomsBase.AbstractSystem` compatible structure and a `calculator` that incudes details of the calculation method. Additionally keywords can be give. These can be ignored, but they need to be present in the function definition.
+
+- 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 first input, to which evaluated force values are stored (look for more details below)
+
+Outputs for the functions need to have following properties
+
+- Energy output is a subtype of `Number` that has a unit with dimensions of energy (mass * length^2 / time^2)
+- Force 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 you can reinterpret it as a matrix
+- Virial is a square matrix (3x3 in 3D) that has units of energy times length
+
+
+### Example implementations
+
+Example potential_energy implementation
+
+```julia
+function AtomsBase.potential_energy(system, calculator::MyType; kwargs...)
+    # we can ignore kwargs... or use them to tune the calculation
+    # or give extra information like pairlist
+
+    # add your own definition here
+    return 0.0u"eV"
+end
+```
+
+Example virial implementation
+
+```julia
+function AtomsBase.virial(system, calculator::MyType; kwargs...)
+    # we can ignore kwargs... or use them to tune the calculation
+    # or give extra information like pairlist
+
+    # add your own definition here
+    return zeros(3,3) * u"eV*Å"
+end
+```
+
+### Implementing forces call
+
+There are two optional implementations for force call allocating and non-allocating. The reason for this is that for very fast potentials allocation has noticeable effect on total evaluation time. So in order to reduce the evaluation time there is non-allocating option. On the other hand some expensive methods, like those in quantum chemistry, always allocate output data.
+
+To make implementation easy for everyone we made it so that you only need to define either allocating or non-allocating force call. The other call is generated is then generated with macro `AtomsBase.@generate_complement`.
+
+Example
+
+```julia
+struct MyType
+end
+
+AtomsBase.@generate_complement function AtomsBase.forces(system, calculator::Main.MyType; kwargs...)
+    # we can ignore kwargs... or use them to tune the calculation
+    # or give extra information like pairlist
+
+    # add your own definition
+    return zeros(AtomsBase.default_force_eltype, length(system))
+end
+```
+
+This creates both `forces` and `forces!`. `AtomsBase.default_force_eltype` is a type that can be used to allocate force data. You can also allocate for some other type.
+
+!!! note "For type definition under @generate_complement macro"
+    You need to use explicit definition of type when using
+    `@generate_complement` macro. `Main.MyType` is fine `MyType` is not!
+
+    You also need to define the type before macro call.
+
+Alternatively the definition could have been done with
+
+```julia
+struct MyOtherType
+end
+
+AtomsBase.@generate_complement function AtomsBase.forces!(f::AbstractVector, system, calculator::Main.MyOtherType; kwargs...)
+    @assert length(f) == length(system)
+    # we can ignore kwargs... or use them to tune the calculation
+    # or give extra information like pairlist
+
+    # add your own definition
+    for i in eachindex(f)
+        f[i] = zero(AtomsBase.default_force_eltype)
+    end
+
+    return f
+end
+```
+
+### Other Automatically Generated Calls
+
+Many methods have optimized calls when energy and forces (and/or virial) are calculated. To allow access to these calls there are also calls
+
+- `energy_forces` for potential energy and allocating forces
+- `energy_forces!` for potential energy and non-allocating forces
+- `energy_forces_virial` for potential energy, allocating forces and virial
+- `energy_forces_virial!` for potential energy, non-allocating forces and virial
+
+These all are generated automatically, if you have defined the corresponding individual methods. The main idea here is that you can implement more efficient methods by your self.
+
+Example implementation
+
+```julia
+function AtomsBase.potential_energy_forces(system, calculator::MyType; kwargs...)
+    # we can ignore kwargs... or use them to tune the calculation
+    # or give extra information like pairlist
+
+    # add your own definition here
+    E = 0.0u"eV"
+    f = zeros(AtomsBase.default_force_eltype, length(system))
+    return (;
+        :energy => E,
+        :forces => f
+    )
+end
+```
+
+Defining this does not overload the corresponding non-allocating call - you need to do that separately.
+
+Output for combination methods is defined to have keys `:energy`, `:forces` and `:virial`. You can access them with
+
+- `output[:energy]` for energy
+- `output[:forces]` for forces
+- `output[:virial]` for viral
+
+The type of the output can be [NamedTuple](https://docs.julialang.org/en/v1/base/base/#Core.NamedTuple), as was in the example above, [Dictionary](https://docs.julialang.org/en/v1/base/collections/#Dictionaries) or any structure that has the keys implemented. The reason for this is that this allows everyone to implement the performance functions without being restricted to certain output type.
+
+## Testing Function Calls
+
+We have implemented function calls to help testing the API. There is one call for each type of call 
+
+- `AtomsBase.test_potential_energy` to test potential_energy call
+- `AtomsBase.test_forces` to test both allocating and non-allocating force call
+- `AtomsBase.test_virial` to test virial call
+
+These functions take the same (non-allocating) input than the API calls.
+
+To test our example potential `MyType` we can do
+
+```julia
+hydrogen = isolated_system([
+    :H => [0, 0, 1.]u"bohr",
+    :H => [0, 0, 3.]u"bohr"
+])
+
+AtomsBase.test_potential_energy(hydrogen, MyType())
+AtomsBase.test_forces(hydrogen, MyType())
+AtomsBase.test_virial(hydrogen, MyType())
+
+AtomsBase.test_forces(hydrogen, MyOtherType()) # this works
+AtomsBase.test_virial(hydrogen, MyOtherType()) # this will fail
+```
+
+*It is recommended that you use the test functions to test that your implementation supports the API fully!*

src/AtomsBase.jl

@@ -12,6 +12,7 @@ include("flexible_system.jl")
 include("atomview.jl")
 include("atom.jl")
 include("fast_system.jl")
+include("calculators.jl")
 
 function __init__()
     @require AtomsView="ee286e10-dd2d-4ff2-afcb-0a3cd50c8041" begin

src/calculators.jl

@@ -0,0 +1,155 @@
+using Test
+
+
+function potential_energy end 
+
+function forces end 
+
+function forces! end 
+
+function virial end 
+
+function energy_forces(system, calculator; kwargs...)
+    e = potential_energy(system, calculator; kwargs...)
+    f = forces(system, calculator; kwargs...)
+    return (;
+        :energy => e,
+        :forces => f
+    )
+end 
+
+function energy_forces!(f::AbstractVector, system, calculator; kwargs...)
+    e = potential_energy(system, calculator; kwargs...)
+    forces!(f, system, calculator; kwargs...)
+    return (;
+        :energy => e,
+        :forces => f
+    )
+end 
+
+function energy_forces_virial(system, calculator; kwargs...)
+    ef = energy_forces(system, calculator; kwargs...)
+    v = virial(system, calculator; kwargs...)
+    return (;
+        :energy => ef[:energy],
+        :forces => ef[:forces],
+        :virial => v
+    )
+end 
+
+function energy_forces_virial!(f::AbstractVector, system, calculator; kwargs...)
+    ef = energy_forces!(f, system, calculator; kwargs)
+    v = virial(system, calculator; kwargs...)
+    return (;
+        :energy => ef[:energy],
+        :forces => ef[:forces],
+        :virial => v
+    )
+end 
+
+
+const default_force_eltype = SVector(1., 1., 1.) * u"eV/Å" |> typeof
+
+"""
+    @generate_complement
+
+Gnereate complementary function for given function expression.
+This is intended to generate non-allocating force call from
+allocating force call and viseversa.
+"""
+macro generate_complement(expr)
+    oldname = nothing
+    try
+        oldname = "$(expr.args[1].args[1])"
+    catch _
+        error("Not valid input")
+    end
+
+    try
+        has_kwargs = any( [ Symbol("...") == x.head  for x in expr.args[1].args[2].args ] )
+        !has_kwargs && error()
+    catch _
+        error("Call does not catch kwargs...")
+    end
+
+    calc_type = nothing
+    try
+        calc_type = expr.args[1].args[end].args[2]
+    catch _
+        throw(error("Calculator does not have defined type"))
+    end
+
+    if oldname[end] == '!'
+        length(expr.args[1].args) != 5 && error("Number of inputs does not match the call")
+        name = oldname[begin:end-1] 
+        q = Meta.parse(
+            "function $name(system, calculator::$calc_type; kwargs...)
+                final_data = zeros( default_force_eltype, length(system) )
+                $oldname(final_data, system, calculator; kwargs...)
+                return final_data
+            end"
+        )  
+    else
+        length(expr.args[1].args) != 4 && error("Number of inputs does not match the call")
+        name = oldname * "!"
+        q = Meta.parse(
+            "function $name(final_data::AbstractVector, system, calculator::$calc_type; kwargs...)
+                @assert length(final_data) == length(system)
+                final_data .= $oldname(system, calculator; kwargs...)
+                return final_data
+            end"
+        )
+    end
+    return quote
+        $expr
+        $q
+    end
+end
+
+
+function test_forces(sys, calculator; force_eltype=default_force_eltype, kwargs...)
+    @testset "Test forces for $(typeof(calculator))" begin
+        f = AtomsBase.forces(sys, calculator; kwargs...)
+        @test typeof(f) <: AbstractVector
+        @test eltype(f) <: AbstractVector
+        @test length(f) == length(sys)
+        T = (eltype ∘ eltype)( f )
+        f_matrix = reinterpret(reshape, T, f)
+        @test typeof(f_matrix) <: AbstractMatrix
+        @test eltype(f_matrix) <: Number
+        @test size(f_matrix) == (3, length(f))
+        @test all( AtomsBase.forces(sys, calculator; dummy_kword659234=1, kwargs...) .≈ f )
+        @test dimension(f[1][1]) == dimension(u"N")
+        @test length(f[1]) == (length ∘ position)(sys,1)
+        f_nonallocating = zeros(force_eltype, length(f))
+        AtomsBase.forces!(f_nonallocating, sys, calculator; kwargs...)
+        @test all( f_nonallocating .≈ f  )
+        AtomsBase.forces!(f_nonallocating, sys, calculator; dummy_kword659254=1, kwargs...)
+        @test all( f_nonallocating .≈ f  )
+    end
+end
+
+
+function test_potential_energy(sys, calculator; kwargs...)
+    @testset "Test potential_energy for $(typeof(calculator))" begin
+        e = AtomsBase.potential_energy(sys, calculator; kwargs...)
+        @test typeof(e) <: Number
+        @test dimension(e) == dimension(u"J")
+        e2 = AtomsBase.potential_energy(sys, calculator; dummy_kword6594254=1, kwargs...)
+        @test e ≈ e2
+    end
+end
+
+
+function test_virial(sys, calculator; kwargs...)
+    @testset "Test virial for $(typeof(calculator))" begin
+        v = AtomsBase.virial(sys, calculator; kwargs...)
+        @test typeof(v) <: AbstractMatrix
+        @test eltype(v) <: Number
+        @test dimension(v[begin]) == dimension(u"J*m")
+        l = (length ∘ position)(sys,1) 
+        @test size(v) == (l,l) # allow different dimensions than 3
+        v2 = AtomsBase.virial(sys, calculator; dummy_kword6594254=1, kwargs...)
+        @test all( v .≈ v2 )
+    end
+end