Merge pull request #294 from MineralsCloud/GatherData

Use `GatherData` in `stage` in src/EquationOfStateWorkflow/Recipes.jl
MineralsCloud · Oct 15, 2023 · dbcfb3a · dbcfb3a
2 parents 497184a + 5dd85d1
commit dbcfb3a
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 10 deletions.
diff --git a/src/EquationOfStateWorkflow/Recipes.jl b/src/EquationOfStateWorkflow/Recipes.jl
@@ -15,6 +15,7 @@ using ..EquationOfStateWorkflow:
     WriteInput,
     RunCmd,
     ExtractData,
+    GatherData,
     SaveData,
     FitEquationOfState,
     SaveParameters
@@ -34,6 +35,7 @@ function stage(::SelfConsistentField, r::ParallelEosFittingRecipe)
         WriteInput(SelfConsistentField()),
         RunCmd(SelfConsistentField()),
         ExtractData(SelfConsistentField()),
+        GatherData(SelfConsistentField()),
         SaveData(SelfConsistentField()),
         FitEquationOfState(SelfConsistentField()),
         SaveParameters(SelfConsistentField()),
@@ -50,19 +52,22 @@ function stage(::SelfConsistentField, r::ParallelEosFittingRecipe)
     extractdata = map(
         thunk -> ConditionalJob(thunk; name="extract E(V) data in SCF"), steps[6]
     )
-    savedata = ArgDependentJob(steps[7]; name="save E(V) data in SCF")
-    fiteos = ArgDependentJob(steps[8]; name="fit E(V) data in SCF")
-    saveparams = ArgDependentJob(steps[9]; name="save EOS parameters in SCF")
+    gatherdata = ArgDependentJob(steps[7]; name="gather E(V) data in SCF")
+    savedata = ArgDependentJob(steps[8]; name="save E(V) data in SCF")
+    fiteos = ArgDependentJob(steps[9]; name="fit E(V) data in SCF")
+    saveparams = ArgDependentJob(steps[10]; name="save EOS parameters in SCF")
     download .→
-    compute .→ makeinputs .→ writeinputs .→ runcmds .→ extractdata .→ fiteos → saveparams
-    extractdata .→ savedata
+    compute .→
+    makeinputs .→ writeinputs .→ runcmds .→ extractdata .→ gatherdata → fiteos → saveparams
+    gatherdata → savedata
     return steps = (;
         download=download,
         compute=compute,
         makeinputs=makeinputs,
         writeinputs=writeinputs,
         runcmds=runcmds,
         extractdata=extractdata,
+        gatherdata=gatherdata,
         savedata=savedata,
         fiteos=fiteos,
         saveparams=saveparams,
@@ -76,6 +81,7 @@ function stage(::VariableCellOptimization, r::ParallelEosFittingRecipe)
         WriteInput(VariableCellOptimization()),
         RunCmd(VariableCellOptimization()),
         ExtractData(VariableCellOptimization()),
+        GatherData(VariableCellOptimization()),
         SaveData(VariableCellOptimization()),
         FitEquationOfState(VariableCellOptimization()),
         SaveParameters(VariableCellOptimization()),
@@ -95,17 +101,20 @@ function stage(::VariableCellOptimization, r::ParallelEosFittingRecipe)
     extractdata = map(
         thunk -> ConditionalJob(thunk; name="extract E(V) data in vc-relax"), steps[5]
     )
-    savedata = ArgDependentJob(steps[6]; name="save E(V) data in vc-relax")
-    fiteos = ArgDependentJob(steps[7]; name="fit E(V) data in vc-relax")
-    saveparams = ArgDependentJob(steps[8]; name="save EOS parameters in vc-relax")
-    compute .→ makeinputs .→ writeinputs .→ runcmds .→ extractdata .→ fiteos → saveparams
-    extractdata .→ savedata
+    gatherdata = ArgDependentJob(steps[6]; name="gather E(V) data in vc-relax")
+    savedata = ArgDependentJob(steps[7]; name="save E(V) data in vc-relax")
+    fiteos = ArgDependentJob(steps[8]; name="fit E(V) data in vc-relax")
+    saveparams = ArgDependentJob(steps[9]; name="save EOS parameters in vc-relax")
+    compute .→
+    makeinputs .→ writeinputs .→ runcmds .→ extractdata .→ gatherdata → fiteos → saveparams
+    gatherdata → savedata
     return steps = (;
         compute=compute,
         makeinputs=makeinputs,
         writeinputs=writeinputs,
         runcmds=runcmds,
         extractdata=extractdata,
+        gatherdata=gatherdata,
         savedata=savedata,
         fiteos=fiteos,
         saveparams=saveparams,

diff --git a/src/EquationOfStateWorkflow/actions.jl b/src/EquationOfStateWorkflow/actions.jl
@@ -54,6 +54,11 @@ struct ExtractData{T} <: Action{T}
     calculation::T
 end
 
+struct GatherData{T} <: Action{T}
+    calculation::T
+end
+(::GatherData)(iter) = collect(iter)
+
 struct SaveData{T} <: Action{T}
     calculation::T
 end

diff --git a/src/EquationOfStateWorkflow/think.jl b/src/EquationOfStateWorkflow/think.jl
@@ -23,6 +23,7 @@ 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::GatherData, ::Conf) = Thunk(action)
 think(action::SaveData, conf::Conf) = Thunk(action(conf.data.raw))
 think(action::SaveParameters, conf::Conf) = Thunk(action(conf.data.eos_params))
 function think(action::FitEquationOfState, conf::Conf)