Merge branch 'main' into tutorial_08

Author: Christoph Ortner

Project.toml

@@ -13,6 +13,7 @@ Bumper = "8ce10254-0962-460f-a3d8-1f77fea1446e"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ChunkSplitters = "ae650224-84b6-46f8-82ea-d812ca08434e"
 DynamicPolynomials = "7c1d4256-1411-5781-91ec-d7bc3513ac07"
+EmpiricalPotentials = "38527215-9240-4c91-a638-d4250620c9e2"
 EquivariantModels = "73ee3e68-46fd-466f-9c56-451dc0291ebc"
 ExtXYZ = "352459e4-ddd7-4360-8937-99dcb397b478"
 Folds = "41a02a25-b8f0-4f67-bc48-60067656b558"

docs/src/tutorials/basic_julia_workflow.jl

@@ -85,3 +85,20 @@ end
 
 # Finally, we delete the model to clean up.
 rm("TiAl_model.json")
+
+# ### Fast Evaluator 
+#
+# `ACEpotentials.jl` provides an experimental "fast evaluator". This tries to 
+# merge some of the operations in the full model resulting in a "slimmer" and 
+# usually faster evaluator. In some cases the performance gain can be multiple 
+# factors up to an order of magnitude. This is particularly important when 
+# using a parameter estimation solver that sparsifies. In that case, the 
+# performance gain can be significant. 
+#
+# To construct the fast evaluator, simply use 
+# ```julia 
+# fpot = fast_evaluator(model)
+# ```
+# An optional keyword argument `aa_static = true` can be used to optimize the 
+# n-correlation layer for very small models (at most a few hundred parameters). 
+# For larger models this leads to a stack overflow.

examples/zuobench/zuo_asp.jl

@@ -2,6 +2,7 @@
 using Distributed, Random, SparseArrays 
 addprocs(10, exeflags="--project=$(Base.active_project())")
 @everywhere using ACEpotentials, PrettyTables
+using ACEpotentials.Models: fast_evaluator
 
 ##
 
@@ -33,32 +34,30 @@ At, yt, Wt = ACEpotentials.assemble(train_data, model)
 Av, yv, Wv = ACEpotentials.assemble(val_data, model)
 
 @info("Compute ASP Path")
-solver = ACEfit.ASP(; P = P, select = :final, tsvd = true, 
-                     actMax = 1000, traceFlag=true )
+solver = ACEfit.ASP(; P = P, select = :final, tsvd = true, actMax = 1000 )
 asp_result = ACEfit.solve(solver, Wt .* At, Wt .* yt, Wv .* Av, Wv .* yv)
 
 ## 
 
 @info("Pick solutions for 100, 300, 1000 parameters, compute errors")
 
-@show length(asp_result["path"])
-path = asp_result["path"]
-nnzs = [ nnz(p.solution) for p in path ]
-I1000 = length(nnzs) 
-I300 = findfirst(nnzs .>= 300)
-I100 = findfirst(nnzs .>= 100)
-
-model_1000 = deepcopy(model)
-set_parameters!(model_1000, path[I1000].solution)
-model_300 = deepcopy(model)
-set_parameters!(model_300, path[I300].solution)
-model_100 = deepcopy(model)
-set_parameters!(model_100, path[I100].solution)
-
-err_100 = ACEpotentials.linear_errors(test_data, model_100)
-err_300 = ACEpotentials.linear_errors(test_data, model_300)
-err_1000 = ACEpotentials.linear_errors(test_data, model_1000)
-
+# select models from the model path 
+model_1000 = set_parameters!( deepcopy(model), 
+                  ACEfit.asp_select(asp_result, :final)[1])
+model_300  = set_parameters!( deepcopy(model), 
+                  ACEfit.asp_select(asp_result, (:bysize, 300))[1])
+model_100  = set_parameters!( deepcopy(model), 
+                  ACEfit.asp_select(asp_result, (:bysize, 100))[1])
+
+# generate sparsified, faster evaluators 
+pot_1000 = fast_evaluator(model_1000; aa_static = false) # static can cause stack overflow
+pot_300 = fast_evaluator(model_300; aa_static = true)
+pot_100 = fast_evaluator(model_100; aa_static = true)
+
+@info("Evaluate errors on the test set") 
+err_100 = ACEpotentials.linear_errors(test_data,  pot_100)
+err_300 = ACEpotentials.linear_errors(test_data,  pot_300)
+err_1000 = ACEpotentials.linear_errors(test_data, pot_1000)
 
 ##
 

src/ACEpotentials.jl

@@ -15,19 +15,18 @@ include("ace1_compat.jl")
 # Fitting
 include("atoms_data.jl")
 include("fit_model.jl")
+include("repulsion_restraint.jl")
 
 # Data 
 include("example_data.jl")
 
 # Misc 
-# TODO: all of this just needs to be moved from JuLIP to AtomsBase
 include("analysis/dataset_analysis.jl")
 include("analysis/potential_analysis.jl")
 include("descriptor.jl")
 
 
 # TODO: to be completely rewritten
-# include("io.jl")
 # include("export.jl")
 
 # Experimental 
@@ -41,19 +40,23 @@ import ACEpotentials.ACE1compat: ace1_model
 import ACEpotentials.Models: algebraic_smoothness_prior, 
                              exp_smoothness_prior, 
                              gaussian_smoothness_prior, 
-                             set_parameters!
+                             set_parameters!, 
+                             fast_evaluator, 
+                             @committee, 
+                             set_committee!
 import JSON 
 
 export ace1_model,
        length_basis, 
        algebraic_smoothness_prior, 
        exp_smoothness_prior, 
        gaussian_smoothness_prior, 
-         set_parameters!
-
+       set_parameters!, 
+       fast_evaluator, 
+         @committee,
+         set_committee!
 
 include("json_interface.jl")
 
 
-
 end

src/ace1_compat.jl

@@ -37,6 +37,7 @@ const _kw_defaults = Dict(:elements => nothing,
                           :pair_envelope => (:r, 2),
                           #
                           :Eref => missing,
+                          :ZBL => false, 
                           #
                           :variable_cutoffs => false,
                           )
@@ -300,8 +301,12 @@ function _pair_basis(kwargs)
       # here we use a similar convention, just need to convert to ace1-style 
       envelope = kwargs[:pair_envelope]
       if envelope isa Tuple && envelope[1] == :r 
+         if kwargs[:ZBL] 
+            @warn("""It is not recommended to combine the ZBL reference potential 
+                     with a repulsive pair basis. Use `pair_envelope = (:x, 0, q)` instead.""")
+         end
          envelope = (:r_ace1, envelope[2])
-      end 
+      end
 
       pair_basis = ace_learnable_Rnlrzz(; spec = pair_spec, 
                                     maxq = maxq,  
@@ -322,6 +327,11 @@ end
 
 function ace1_model(; kwargs...)
 
+   # change the default for the envelope if ZBL is used 
+   if haskey(kwargs, :ZBL) && kwargs[:ZBL] && !haskey(kwargs, :envelope)
+      kwargs = (; pair_envelope = (:x, 0, 2), kwargs...)
+   end
+
    model_spec = Dict{Symbol, Any}(:model_name => "ACE1", kwargs...)
 
    kwargs = _clean_args(kwargs)
@@ -372,9 +382,11 @@ function ace1_model(; kwargs...)
       E0s = Dict([ key => val * u"eV" for (key, val) in Eref]...) 
    end
 
+
    model = Models.ace_model(; elements=elements, 
                        order = cor_order, 
                        Ytype = :spherical, 
+                       ZBL = kwargs[:ZBL],
                        E0s = E0s, 
                        rbasis = rbasis, 
                        pair_basis = pairbasis, 

src/atoms_data.jl

@@ -33,6 +33,9 @@ function _getfuzzy(coll, key)
 end
 
 _issimilarkey(k1, k2) = lowercase(String(k1)) == lowercase(String(k2))
+_issimilarkey(k1::Nothing, k2) = false 
+_issimilarkey(k1, k2::Nothing) = false
+_issimilarkey(k1::Nothing, k2::Nothing) = false
 
 function _find_similar_key(coll, key) 
     for k in keys(coll) 
@@ -43,38 +46,38 @@ function _find_similar_key(coll, key)
     return nothing 
 end
 
-function _find_similar_key(sys::ExtXYZ.Atoms, key) 
-    for k in keys(sys.system_data)
+function _find_similar_key(sys::AbstractSystem, key) 
+    for k in keys(sys)
         if _issimilarkey(k, key)
             return k 
         end
     end
-    for k in keys(sys.atom_data)
+    for k in atomkeys(sys)
         if _issimilarkey(k, key)
             return k 
         end
     end
     return nothing 
 end
 
-function _get_data_fuzzy(sys::ExtXYZ.Atoms, key)
+function _get_data_fuzzy(sys::AbstractSystem, key)
     k = _find_similar_key(sys, key) 
     if k == nothing 
         error("Couldn't find $key or similar in collection with keys $(keys(sys))")
     end 
-    if haskey(sys.system_data, k)
-        return sys.system_data[k]
+    if haskey(sys, k)
+        return sys[k]
     end
-    return sys.atom_data[k]    
+    return sys[:, k]
 end
 
 _has_similar_key(coll, key) = (_find_similar_key(coll, key) != nothing)
 
-function _get_data(sys::ExtXYZ.Atoms, key)
-    if haskey(sys.system_data, key)
-        return sys.system_data[key]
-    elseif haskey(sys.atom_data, key)
-        return sys.atom_data[key]
+function _get_data(sys::AbstractSystem, key)
+    if haskey(sys, key)
+        return sys[key]
+    elseif hasatomkey(sys, key)
+        return sys[:, key]
     else
         error("Couldn't find $key in System")
     end

src/fit_model.jl

@@ -110,7 +110,15 @@ function acefit!(raw_data::AbstractArray{<: AbstractSystem}, model;
    end 
 
    if repulsion_restraint 
-      error("Repulsion restraint is currently not implemented")
+      restraint_data = _rep_dimer_data_atomsbase(
+               model; 
+               weight = restraint_weight,
+               energy_key = Symbol(energy_key),
+               kwargs...
+            )
+      append!(data, restraint_data)
+      # return nothing 
+      # error("Repulsion restraint is currently not implemented")
       # if eltype(data) == AtomsData
       #    append!(data, _rep_dimer_data(model; weight = restraint_weight))
       # else
@@ -190,11 +198,11 @@ function linear_errors(raw_data::AbstractArray{<: AbstractSystem}, model;
                        virial_key = "virial", 
                        weights = default_weights(), 
                        verbose = true,
-                       return_efv = false
+                       return_efv = false, 
                        )
    data = [ AtomsData(at; energy_key = energy_key, force_key=force_key, 
                           virial_key = virial_key, weights = weights, 
-                          v_ref = _get_Vref(model)) 
+                          v_ref = nothing)
             for at in raw_data ]
    return linear_errors(data, model; verbose=verbose, return_efv = return_efv)
 end

src/models/ace.jl

@@ -92,35 +92,16 @@ function _make_idx_AA_spec(AA_spec, A_spec)
    return AA_spec_idx
 end 
 
-function _make_Vref_E0s(rbasis, E0s::Nothing)
-   NZ = _get_nz(rbasis)
-   return _make_Vref_E0s(rbasis, [ _i2z(rbasis, i) => 0.0 for i = 1:NZ ] )
-end
-
-# E0s can be anything with (key, value) pairs 
-function _make_Vref_E0s(rbasis, E0s)
-   _convert_E0s(E0s::Union{Dict, NamedTuple}) = E0s 
-   _convert_E0s(E0s::Union{AbstractVector, Tuple}) = Dict(E0s...)
-   _convert_E0s(E0s) = error("E0s must be nothing, a NamedTuple, Dict or list of pairs")
-   
-   NZ = _get_nz(rbasis)
-   V0 = OneBody(_convert_E0s(E0s))
-   if length(V0.E0) != NZ 
-      error("E0s must have the right number of elements")
-   end
-
-   return V0 
-end
-
 
 function _generate_ace_model(rbasis, Ytype::Symbol, AA_spec::AbstractVector, 
+                             Vref, 
                              level = TotalDegree(), 
                              pair_basis = nothing, 
-                             E0s = nothing, 
-                             Vref = _make_Vref_E0s(rbasis, E0s), )
+                             ) 
 
-   # storing E0s with unit
-   model_meta = Dict{String, Any}("E0s" => deepcopy(E0s))
+   # # storing E0s with unit
+   # model_meta = Dict{String, Any}("E0s" => deepcopy(E0s))
+   model_meta = Dict{String, Any}()
 
    # generate the coupling coefficients 
    cgen = EquivariantModels.Rot3DCoeffs_real(0)
@@ -170,8 +151,8 @@ end
 #       since it is implicitly already encoded in AA_spec. We need a 
 #       function `auto_level` that generates level automagically from AA_spec.
 function ace_model(rbasis, Ytype, AA_spec::AbstractVector, level, 
-                   pair_basis, E0s = nothing)
-   return _generate_ace_model(rbasis, Ytype, AA_spec, level, pair_basis, E0s)
+                   pair_basis, Vref)
+   return _generate_ace_model(rbasis, Ytype, AA_spec, Vref, level, pair_basis)
 end 
 
 # NOTE : a nicer convenience constructor is also provided in `ace_heuristics.jl`
@@ -320,9 +301,8 @@ function evaluate(model::ACEModel,
       val += dot(Apair, (@view ps.Wpair[:, i_z0]))
    end
    # ------------------- 
-   #  TODO - Vref : assume it is a OneBody 
-   @assert model.Vref isa OneBody
-   val += model.Vref.E0[Z0]
+   #  Vref
+   val += eval_site(model.Vref, Rs, Zs, Z0)
    # ------------------- 
 
    end # @no_escape