Merge pull request #290 from MineralsCloud:ConvergenceTestWorkflow

Rewrite `ConvergenceTestWorkflow`

src/ConvergenceTestWorkflow/Config.jl

@@ -1,20 +1,19 @@
 module Config
 
 using Configurations: from_dict, @option
-using ExpressBase: Action, SoftwareConfig
+using EasyConfig: Config as Conf
+using ExpressBase.Config: AbstractConfig, SoftwareConfig, SamplingPoints, IO, list_io
 using Unitful: FreeUnits
 
-using ...Config: SamplingPoints, DirStructure, iofiles
-
-@option "ecutwfc" struct CutoffEnergies <: SamplingPoints
+@option "ecut" struct CutoffEnergies <: SamplingPoints
     numbers::Vector{Float64}
     unit::FreeUnits
     CutoffEnergies(numbers, unit="Ry") = new(numbers, unit)
 end
 
-@option "k_mesh" struct MonkhorstPackGrids
+@option "kmesh" struct MonkhorstPackGrids <: AbstractConfig
     meshes::AbstractVector{<:AbstractVector{<:Integer}}
-    shifts::AbstractVector{<:AbstractVector{<:Integer}} = fill([0, 0, 0], length(meshes))
+    shifts::AbstractVector{<:AbstractVector{Bool}} = fill(falses(3), length(meshes))
     function MonkhorstPackGrids(meshes, shifts)
         if length(meshes) != length(shifts)
             throw(DimensionMismatch("`meshes` and `shifts` should have the same length!"))
@@ -27,58 +26,57 @@ end
     end
 end
 
-@option struct Save
-    raw::String = "raw.json"
-    status::String = ""
+@option struct Data <: AbstractConfig
+    raw::String = "energies.json"
 end
 
-@option struct RuntimeConfig
+@option struct StaticConfig <: AbstractConfig
     recipe::String
     template::String
-    parameters::Union{CutoffEnergies,MonkhorstPackGrids}
-    dirstructure::DirStructure = DirStructure()
-    save::Save = Save()
+    with::Union{CutoffEnergies,MonkhorstPackGrids}
+    io::IO = IO()
+    data::Data = Data()
     cli::SoftwareConfig
-    function RuntimeConfig(recipe, template, parameters, dirstructure, save, cli)
-        @assert recipe in ("ecut", "k_mesh")
+    function StaticConfig(recipe, template, with, io, data, cli)
+        @assert recipe in ("ecut", "kmesh")
         if !isfile(template)
             @warn "I cannot find template file `$template`!"
         end
-        return new(recipe, template, parameters, dirstructure, save, cli)
+        return new(recipe, template, with, io, data, cli)
     end
 end
 
-struct ExpandConfig{T} <: Action{T} end
-(::ExpandConfig)(energies::CutoffEnergies) = energies.numbers .* energies.unit
-function (::ExpandConfig)(x::MonkhorstPackGrids)
-    return map(x.meshes, x.shifts) do mesh, shift
-        (mesh, shift)
-    end
-end
-function (::ExpandConfig{T})(ds::DirStructure, energies::CutoffEnergies) where {T}
-    return iofiles(ds, energies.numbers, string(nameof(T)))
+function _update!(conf::Conf, io::IO, energies::CutoffEnergies)
+    conf.io = collect(
+        list_io(io, number, string(nameof(typeof(conf.calculation)))) for
+        number in energies.numbers
+    )
+    return conf
 end
-function (::ExpandConfig{T})(ds::DirStructure, grids::MonkhorstPackGrids) where {T}
-    return iofiles(ds, zip(grids.meshes, grids.shifts), string(nameof(T)))
+function _update!(conf::Conf, io::IO, grids::MonkhorstPackGrids)
+    conf.io = collect(
+        list_io(
+            io,
+            join(mesh, '×') * '_' * join(shift, '+'),
+            string(nameof(typeof(conf.calculation))),
+        ) for (mesh, shift) in zip(grids.meshes, grids.shifts)
+    )
+    return conf
 end
-function (::ExpandConfig)(save::Save)
-    return map((:raw, :status)) do f
-        v = getfield(save, f)
-        isempty(v) ? abspath(mktemp(; cleanup=false)[1]) : abspath(expanduser(v))
-    end
+function _update!(conf::Conf, data::Data)
+    conf.data.raw = abspath(expanduser(data.raw))
+    return conf
 end
-function (x::ExpandConfig)(config::AbstractDict)
-    config = from_dict(RuntimeConfig, config)
-    save_raw, save_status = x(config.save)
-    return (
-        template=x(config.template),
-        parameters=x(config.parameters),
-        root=config.files.dirs.root,
-        files=x(config.files, config.parameters),
-        save_raw=save_raw,
-        save_status=save_status,
-        cli=config.cli,
-    )
+
+function expand(config::StaticConfig, calculation::Calculation)
+    conf = Conf()
+    conf.cli = config.cli
+    conf.calculation = calculation
+    _update!(conf, config.template)
+    _update!(conf, config.with)
+    _update!(conf, config.io, config.with)
+    _update!(conf, config.data)
+    return conf
 end
 
 end

src/ConvergenceTestWorkflow/ConvergenceTestWorkflow.jl

@@ -2,6 +2,7 @@ module ConvergenceTestWorkflow
 
 include("Config.jl")
 include("actions.jl")
+include("think.jl")
 include("Recipes.jl")
 
 end

src/ConvergenceTestWorkflow/Recipes.jl

@@ -1,45 +1,79 @@
 module Recipes
 
-using ExpressBase: SCF
+using Configurations: from_dict
+using EasyJobsBase: Job, ConditionalJob, ArgDependentJob, →
+using ExpressBase: SelfConsistentField, think
 using ExpressBase.Files: load
 using ExpressBase.Recipes: Recipe
-using SimpleWorkflows.Workflows: Workflow, run!, →, ⇉, ⇶, ⭃
+using SimpleWorkflows: Workflow, eachjob, run!
 
-using ...Express: DownloadPotentials, RunCmd, jobify
-using ..ConvergenceTestWorkflow: MakeInput, TestConvergence
+using ..Config: StaticConfig, expand
+using ..ConvergenceTestWorkflow:
+    DownloadPotentials,
+    CreateInput,
+    WriteInput,
+    RunCmd,
+    ExtractData,
+    SaveData,
+    TestConvergence
 
 export build, run!
 
 struct TestCutoffEnergyRecipe <: Recipe
     config
 end
-struct TestKPointsRecipe <: Recipe
+struct TestKMeshRecipe <: Recipe
     config
 end
 
-function build(::Type{Workflow}, r::TestCutoffEnergyRecipe)
-    a = jobify(DownloadPotentials{SCF}(), r.config)
-    c = jobify(MakeInput{SCF}(), r.config)
-    d = jobify(RunCmd{SCF}(), r.config)
-    e = jobify(TestConvergence{SCF}(), r.config)
-    a → b ⇉ c ⇶ d ⭃ e
-    return Workflow(a)
+function stage(::SelfConsistentField, r::TestCutoffEnergyRecipe)
+    conf = expand(r.config, SelfConsistentField())
+    steps = map((
+        DownloadPotentials(SelfConsistentField()),
+        CreateInput(SelfConsistentField()),
+        WriteInput(SelfConsistentField()),
+        RunCmd(SelfConsistentField()),
+        ExtractData(SelfConsistentField()),
+        SaveData(SelfConsistentField()),
+        TestConvergence(SelfConsistentField()),
+    )) do action
+        think(action, conf)
+    end
+    download = Job(steps[1]; name="download potentials")
+    makeinputs = map(thunk -> Job(thunk; name="update input in SCF"), steps[2])
+    writeinputs = map(thunk -> ArgDependentJob(thunk; name="write input in SCF"), steps[3])
+    runcmds = map(
+        thunk -> ConditionalJob(thunk; name="run ab initio software in SCF"), steps[4]
+    )
+    extractdata = map(
+        thunk -> ConditionalJob(thunk; name="extract energies in SCF"), steps[5]
+    )
+    savedata = ArgDependentJob(steps[6]; name="save energies in SCF")
+    testconv = ArgDependentJob(steps[9]; name="test convergence in SCF")
+    download .→ makeinputs .→ writeinputs .→ runcmds .→ extractdata .→ testconv
+    extractdata .→ savedata
+    return steps = (;
+        download=download,
+        makeinputs=makeinputs,
+        writeinputs=writeinputs,
+        runcmds=runcmds,
+        extractdata=extractdata,
+        savedata=savedata,
+        testconv=testconv,
+    )
 end
-function build(::Type{Workflow}, r::TestKPointsRecipe)
-    a = jobify(DownloadPotentials{SCF}(), r.config)
-    c = jobify(MakeInput{SCF}(), r.config)
-    d = jobify(RunCmd{SCF}(), r.config)
-    e = jobify(TestConvergence{SCF}(), r.config)
-    a → b ⇉ c ⇶ d ⭃ e
-    return Workflow(a)
+
+function build(::Type{Workflow}, r::Union{TestCutoffEnergyRecipe,TestKMeshRecipe})
+    stage = stage(SelfConsistentField(), r)
+    return Workflow(stage.download)
 end
 function build(::Type{Workflow}, file)
     dict = load(file)
     recipe = dict["recipe"]
     if recipe == "ecut"
         return build(Workflow, TestCutoffEnergyRecipe(dict))
     elseif recipe == "k_mesh"
-        return build(Workflow, TestKPointsRecipe(dict))
+        return build(Workflow, TestKMeshRecipe(dict))
     else
         error("unsupported recipe $recipe.")
     end

src/ConvergenceTestWorkflow/actions.jl

@@ -1,66 +1,36 @@
-using AbInitioSoftwareBase: Input, writetxt
-using ExpressBase: Action, SCF
+using AbInitioSoftwareBase: Input
+using ExpressBase:
+    Calculation, Action, SelfConsistentField, DownloadPotentials, RunCmd, WriteInput
 using ExpressBase.Files: save, load
 using EasyJobsBase: Job
 using Unitful: ustrip
 
-using ..Express: distribute_procs
-using .Config: ExpandConfig
-
-import ..Express: jobify
-
 struct CreateInput{T} <: Action{T}
     calculation::T
 end
-
-function jobify(x::CreateInput{SCF}, cfgfile)
-    dict = load(cfgfile)
-    config = ExpandConfig{SCF}()(dict)
-    inputs = first.(config.files)
-    if config.parameters isa AbstractVector{<:Tuple}
-        return map(inputs, config.parameters) do input, (mesh, shift)
-            Job(() -> x(input, config.template, mesh, shift, "Y-m-d_H:M:S"))
-        end
-    else
-        return map(inputs, config.parameters) do input, energy
-            Job(() -> x(input, config.template, energy, "Y-m-d_H:M:S"))
-        end
-    end
-end
-
-function parseoutput end
+(action::CreateInput)(template::Input) = Base.Fix1(action, template)
 
 struct ExtractData{T} <: Action{T}
     calculation::T
 end
 
-function jobify(x::ExtractData{T}, cfgfile) where {T}
-    dict = load(cfgfile)
-    config = ExpandConfig{T}()(dict)
-    return Job(function ()
-        data = x(last.(config.files))
-        saved = Dict("results" => (ustrip ∘ last).(data))
-        save(config.save_raw, saved)
-        return data
-    end)
+struct SaveData{T} <: Action{T}
+    calculation::T
+end
+function (action::SaveData)(path, raw_data)
+    raw_data = sort(collect(raw_data))  # In case the data is not sorted
+    data = Dict(
+        "ecut_kmesh" => (string ∘ first).(raw_data), "energy" => (string ∘ last).(raw_data)
+    )
+    return save(path, data)
 end
+(action::SaveData)(path) = Base.Fix1(action, path)
 
 struct TestConvergence{T} <: Action{T}
     calculation::T
 end
 (x::TestConvergence)(data) = isconvergent(data)
 
-function jobify(x::TestConvergence{T}, cfgfile) where {T}
-    dict = load(cfgfile)
-    config = ExpandConfig{T}()(dict)
-    return Job(function ()
-        data = ExtractData{T}()(last.(config.files))
-        saved = Dict("results" => (ustrip ∘ last).(data))
-        save(config.save_raw, saved)
-        return x(data)
-    end)
-end
-
 function isconvergent(a::AbstractVector)
     terms = abs.(diff(a))
     x, y, z = last(terms, 3)

src/ConvergenceTestWorkflow/think.jl

@@ -0,0 +1,18 @@
+using Thinkers: Thunk
+using EasyConfig: Config as Conf
+
+using .Config: StaticConfig, expand
+
+import ExpressBase: think
+
+think(action::CreateInput, conf::Conf) =
+    collect(Thunk(action(conf.template)) for _ in Base.OneTo(length(conf.at)))
+think(action::WriteInput, conf::Conf) =
+    collect(Thunk(action(file)) for file in first.(conf.io))
+think(action::ExtractData, conf::Conf) =
+    collect(Thunk(action, file) for file in last.(conf.io))
+think(action::SaveData, conf::Conf) = Thunk(action(conf.data.raw))
+function think(action::Action{T}, config::StaticConfig) where {T}
+    config = expand(config, T())
+    return think(action, config::Conf)
+end