add checkpoints

Author: DexuanZhou

Project.toml

@@ -12,6 +12,7 @@ EquivariantTensors = "5e107534-7145-4f8f-b06f-47a52840c895"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 HyperDualNumbers = "50ceba7f-c3ee-5a84-a6e8-3ad40456ec97"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LowRankApprox = "898213cb-b102-5a47-900c-97e73b919f73"
@@ -29,9 +30,10 @@ UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
-julia = "1.10"
 EquivariantTensors = "0.1.2"
 Polynomials4ML = "0.5.0"
+julia = "1.10"
+
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 

example.jl

@@ -1,40 +1,34 @@
 # add https://github.com/ACEsuit/ChemBasisSets.jl.git
 # add ACEpsi #lux
-using Distributed
-N_procs = 8
-
-import Pkg
-Pkg.activate(Base.current_project())
-
-if nprocs() == 1
-    addprocs(N_procs - 1, exeflags="--project=$(Base.current_project())")
-end
- 
 using ACEpsi
 
+method = 1
 mol = ACEpsi.molecules.Be
+mol_name = "Be"
+setup(mol, mol_name, method)
 
-basis_set = "cc-pvtz"
-totdeg = [["1f", "1f"]]
-ν = 2
+x0, optimizer, model_list, ps_list, st_list, spec_list, spec1p_list, totdeg_list, ν_list = load_setup(mol_name);
+solver = (SPRINGSolver(), SketchSolver(800, 50, 50, 1.4), SVDSolver(800, 50, 50, 1.4))
+optimizer.sr_method = solver[method]
+string = method == 1 ? "SPRING" :
+             method == 2 ? "SKETCH" :
+             method == 3 ? "WSSR"   : error("Invalid method")
+optimizer.res_path = "$mol_name/$string/"
 
-model_list, ps_list, st_list, spec_list, spec1p_list, totdeg_list, ν_list = model_generator(mol, basis_set, totdeg, ν; ratio = 0.5);
-test_wavefunction(model_list, ps_list, st_list, spec_list, spec1p_list, mol)
-
-solver = (SPRINGSolver(), SketchSolver(800, 50, 50, 1.4), SVDSolver(800, 20, 20, 1.4))
+N_procs = 8
+using Pkg
+using Distributed
+Pkg.activate(Base.current_project())
 
-iterations = 1000 * ones(Int64, length(spec_list))
-iterations[end] = 30000
-optimizer = OPTSETTING(solver[1], iterations = iterations, burnin = 1000, lag = 10, nchains = 2^8, 
-                                 Δt = 0.08, acc_step = 10, acc_range = [0.45, 0.80],  
-                                 clip = 5.0, lr = 0.02, lr_dc = 10000,  m = 0.99, 
-                                 damping = 0.001, damping_decay = 100, damping_min = 0.001, 
-                                 norm_constrain = 0.001, η = 0.95, res_path = "Be/SPRING/")
+if nprocs() == 1
+    addprocs(N_procs - 1, exeflags="--project=$(Base.current_project())")
+end
 
 @everywhere begin 
     using ACEpsi
     using Statistics
     using LinearAlgebra
     using Optimisers: destructure
 end
-model_list, ps_list, st_list, val_list, var_list, rank_list = train(mol, model_list, ps_list, st_list, spec_list, spec1p_list, totdeg_list, ν_list, optimizer);
+
+model_list, ps_list, st_list, val_list, var_list, rank_list = train(x0, mol, model_list, ps_list, st_list, spec_list, spec1p_list, totdeg_list, ν_list, optimizer);

src/model/spec.jl

@@ -146,12 +146,12 @@ end
 
 
 function build_totdeglevels(mol, basis_set, totdeg, ν; ratio = 0.5, max_level::Union{Nothing, Int} = nothing)
-    _, orbital = ACEpsi.auto_load_basis(mol, basis_set; return_spec = true)
+    _, orbital = auto_load_basis(mol, basis_set; return_spec = true)
     n_atom = length(orbital)
 
-    lj = maximum(length.(ACEpsi.iteratespec1p.(ACEpsi.orbital_for_iatom_jord.(Ref(orbital), Ref(totdeg), 1:n_atom, 1), Ref(1))))
+    lj = maximum(length.(iteratespec1p.(orbital_for_iatom_jord.(Ref(orbital), Ref(totdeg), 1:n_atom, 1), Ref(1))))
     _lj = max(Int(ceil(lj * ratio)), 1)
-    spec1pl = [ACEpsi.iteratespec1p(ACEpsi.orbital_for_iatom_jord(orbital, totdeg, i, 1), lj) for i = 1:n_atom]
+    spec1pl = [iteratespec1p(orbital_for_iatom_jord(orbital, totdeg, i, 1), lj) for i = 1:n_atom]
 
     totdeglevels = Vector{Vector{Vector{String}}}()
     νlevels = Int[]

src/molecule/molecules.jl

@@ -24,15 +24,15 @@ CH4 = create_molecule(SVector(
     Nuc("H", [-1.18886,  1.18886, -1.18886])
 ))
 
-# === Diatomic and small molecules with spacing parameter ===
-Li2(spacing) = create_molecule(SVector(
-    Nuc("Li", [-spacing/2, 0.0, 0.0]),
-    Nuc("Li", [ spacing/2, 0.0, 0.0])
+Li2 = create_molecule(SVector(
+    Nuc("Li", [-5.051/2, 0.0, 0.0]),
+    Nuc("Li", [ 5.051/2, 0.0, 0.0])
 ))
-LiH(spacing) = create_molecule(SVector(
-    Nuc("Li", [-spacing/2, 0.0, 0.0]),
-    Nuc("H",  [ spacing/2, 0.0, 0.0])
+LiH = create_molecule(SVector(
+    Nuc("Li", [-3.015/2, 0.0, 0.0]),
+    Nuc("H",  [ 3.015/2, 0.0, 0.0])
 ))
+# === Diatomic and small molecules with spacing parameter ===
 N2(spacing) = create_molecule(SVector(
     Nuc("N", [-spacing/2, 0.0, 0.0]),
     Nuc("N", [ spacing/2, 0.0, 0.0])

src/train/opt/sr.jl

@@ -1,7 +1,7 @@
 export opts!
 
 function opts!(i, OptParams::OPTPARAMS, optimizer::DirectSolver, Eloc, o::Matrix{T}, nchains, damping::Float64, dim_ps::Int64, η::Float64, norm_constrain, γ, m) where {T}
-    s = o * o'
+    s = o * o' # O * O'
     @inbounds @simd for i = 1:dim_ps
         s[i, i] += damping
     end
@@ -19,45 +19,41 @@ end
 
 function opts!(i, OptParams::OPTPARAMSPRING, optimizer::SPRINGSolver, Eloc, o::Matrix{T}, nchains, damping::Float64, dim_ps::Int64, η::Float64, norm_constrain, γ, m) where {T}
     res = norm(OptParams.f)
-    lmul!(-γ, Eloc)
+    lmul!(-γ, Eloc) # -delta tau * (E - E_mean)
     s = o' * o
 
     s .+= 1/(nchains)
-    Tvecs, Tvals = svd(Symmetric(s))
+    Tvals, Tvecs = eigen(Symmetric(s))
     Tvals = max.(Tvals, 0.0) .+ damping
 
-    if OptParams.dw_tot[1] !== 0.0
-        mul!(OptParams.dow, transpose(o), OptParams.dw_tot) 
-        epsilon_tilde = Eloc .- η * OptParams.dow
-    else
-        epsilon_tilde = Eloc
-    end
+    mul!(OptParams.dow, transpose(o), OptParams.dw_tot) 
+    epsilon_tilde = Eloc .- η * OptParams.dow
 
     mul!(OptParams.dow, Tvecs', epsilon_tilde)
-    OptParams.dow = Diagonal(1 ./ Tvals) * OptParams.dow
+    ldiv!(Diagonal(Tvals), OptParams.dow)
     OptParams.dow = Tvecs * OptParams.dow
-    mul!(OptParams.f, o, OptParams.dow) 
-    OptParams.dw_tot .*= η
-    OptParams.dw_tot .+= OptParams.f
+    OptParams.dow .-= mean(OptParams.dow)
+
+    mul!(OptParams.f, o, OptParams.dow)
+    OptParams.dw_tot .= η * OptParams.dw_tot .+ OptParams.f / sqrt(nchains) 
     OptParams.dw_tot .*= min(1, sqrt(norm_constrain)/norm(OptParams.dw_tot))
     return OptParams.dw_tot, length(OptParams.f), res
 end
 
 function opts!(i, OptParams::OPTPARAMMINSR, optimizer::MINSRSolver, Eloc, o::Matrix{T}, nchains, damping::Float64, dim_ps::Int64, η::Float64, norm_constrain, γ, m) where {T}
-    lmul!(sqrt(nchains), Eloc)
-    ldiv!(sqrt(nchains), o)
-    res = norm(o * Eloc)
-    lmul!(-γ, Eloc)
+    res = norm(OptParams.f)
+    lmul!(-γ, Eloc) # -delta tau * (E - E_mean)
     s = o' * o
 
-    Tvecs, Tvals = svd(Symmetric(s))
+    Tvals, Tvecs = eigen(Symmetric(s))
     Tvals = max.(Tvals, 0.0) .+ damping
+
     mul!(OptParams.dow, Tvecs', Eloc)
-    OptParams.dow = Diagonal(1 ./ Tvals) * OptParams.dow
+    ldiv!(Diagonal(Tvals), OptParams.dow)
     OptParams.dow = Tvecs * OptParams.dow
-    mul!(OptParams.f, o, OptParams.dow) 
-    OptParams.dw_tot .*= η
-    OptParams.dw_tot .+= (1-η) * OptParams.f
+
+    mul!(OptParams.f, o, OptParams.dow)
+    OptParams.dw_tot .= η * OptParams.dw_tot .+ (1-η) * OptParams.f / sqrt(nchains) 
     OptParams.dw_tot .*= min(1, sqrt(norm_constrain)/norm(OptParams.dw_tot))
     return OptParams.dw_tot, length(OptParams.f), res
 end

src/train/opt/struct.jl

@@ -4,7 +4,7 @@ export MINSRSolver, SPRINGSolver, DirectSolver, SketchSolver, SVDSolver
 
 abstract type SR_method end
 
-struct OPTSETTING
+mutable struct OPTSETTING
     sr_method::SR_method
     iterations::Vector{Int64}
     burnin::Int
@@ -13,6 +13,7 @@ struct OPTSETTING
     Δt::Float64
     acc_step::Int
     acc_range::Vector{Float64}
+    acc_opt
     clip::Float64
     lr::Float64
     lr_dc::Int
@@ -23,16 +24,17 @@ struct OPTSETTING
     norm_constrain::Float64
     η::Float64
     res_path
+    checkpoints
 end
 
 function OPTSETTING(sr_method::SR_method; iterations::Vector{Int}, burnin::Int, lag::Int, nchains::Int,
-                    Δt::Float64, acc_step::Int, acc_range::Vector{Float64},
+                    Δt::Float64, acc_step::Int, acc_range::Vector{Float64}, acc_opt, 
                     clip::Float64, lr::Float64, lr_dc::Int, m::Float64,
                     damping::Float64, damping_decay::Int, damping_min::Float64,
-                    norm_constrain::Float64, η::Float64, res_path)
-    return OPTSETTING(sr_method, iterations, burnin, lag, nchains, Δt, acc_step, acc_range,
+                    norm_constrain::Float64, η::Float64, res_path, checkpoints)
+    return OPTSETTING(sr_method, iterations, burnin, lag, nchains, Δt, acc_step, acc_range, acc_opt, 
                       clip, lr, lr_dc, m, damping, damping_decay, damping_min,
-                      norm_constrain, η, res_path)
+                      norm_constrain, η, res_path, checkpoints)
 end
 
 struct MINSRSolver <: SR_method