Precompiling dynamics without making run_dynamics type unstable

Project.toml

@@ -90,6 +90,11 @@ Unitful = "1"
 UnitfulAtomic = "1"
 julia = "≥1.9"
 
+[sources]
+NQCDInterfASE = {url = "https://github.com/NQCD/NQCDInterfASE.jl"}
+NQCModels = {url="https://github.com/NQCD/NQCModels.jl", rev="main"}
+NQCBase = {url="https://github.com/NQCD/NQCBase.jl", rev="main"}
+
 [extras]
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
@@ -101,6 +106,7 @@ JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
 JuLIP = "945c410c-986d-556a-acb1-167a618e0462"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
+NQCDInterfASE = "4606675f-2246-4a0b-99d5-75b910de637b"
 MKL = "33e6dc65-8f57-5167-99aa-e5a354878fb2"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
@@ -110,4 +116,4 @@ Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "SafeTestsets", "CSV", "DataFrames", "DiffEqNoiseProcess", "FiniteDiff", "DiffEqDevTools", "Logging", "MKL", "PythonCall", "JuLIP", "Symbolics", "Statistics", "CairoMakie", "JLD2", "JSON"]
+test = ["Test", "SafeTestsets", "CSV", "DataFrames", "DiffEqNoiseProcess", "FiniteDiff", "DiffEqDevTools", "Logging", "MKL", "PythonCall", "JuLIP", "Symbolics", "Statistics", "CairoMakie", "JLD2", "JSON", "NQCDInterfASE"]

src/Ensembles/Ensembles.jl

@@ -88,49 +88,23 @@ function run_dynamics(
     reduction = AppendReduction(),
     ensemble_algorithm = SciMLBase.EnsembleSerial(),
     algorithm = DynamicsMethods.select_algorithm(sim),
-    trajectories = 1,
+    trajectories::Int = 1,
     savetime = true,
     precompile_dynamics = true,
     kwargs...,
 )
 
-    """
-    Due to how Julia compiles code, running the same dynamics for a single time step forces precompilation, which will make subsequent simulations faster. 
-    This effect is negligible for smaller models, but begins to matter for more computationally intensive dynamics. 
+    #@debug "Setting up dynamics with" tspan = tspan
+    if !(output isa Tuple)
+        output = (output,)
+    end
 
-    See this GitHub issue for discussion: https://github.com/NQCD/NQCDynamics.jl/issues/365
-    """
-    function dynamics(tspan)
-        @debug "Setting up dynamics with" tspan = tspan
-        if !(output isa Tuple)
-            output = (output,)
-        end
-
-        kwargs = NQCBase.austrip_kwargs(; kwargs...)
-        trajectories = convert(Int, trajectories)
-        tspan = austrip.(tspan)
-        prob_func = Selection(distribution, selection, trajectories)
-
-        u0 = sample_distribution(sim, distribution)
-        problem = DynamicsMethods.create_problem(u0, tspan, sim)
-
-        output_func = EnsembleSaver(output, savetime)
+    kwargs = NQCBase.austrip_kwargs(; kwargs...)
+    tspan = austrip.(tspan)
+    prob_func = Selection(distribution, selection, trajectories)
+
+    u0 = sample_distribution(sim, distribution)
 
-        ensemble_problem = SciMLBase.EnsembleProblem(
-            problem;
-            prob_func = prob_func,
-            output_func = output_func,
-            reduction = deepcopy(reduction), # deepcopy reduction function because it's used twice
-        )
-        return @timed SciMLBase.solve(
-            ensemble_problem,
-            algorithm,
-            ensemble_algorithm;
-            trajectories,
-            kwargs...,
-        )
-    end
-
     if precompile_dynamics
         @debug "Beginning to precompile dynamics"
         # Get a short time limit that runs through saving at least one time step of data
@@ -140,32 +114,49 @@ function run_dynamics(
             get(
                 kwargs,
                 :dt,
-                1.0 # dt should be a number / unitful qty
-            )
-        ) |> first # Get the number or the first list entry. 
+                1.0 # dt should be a number
+            )::Real
+        )::Union{Vector{Real}, Real} |> first # Get the number or the first list entry. 
         tspan_short = (0.0, short_time)
-        precompile_time = dynamics(tspan_short)
-        @debug "Reduction is" reduction=reduction
-        if isa(reduction, FileReduction)
-            @debug "FileReduction wrote files, which will now be deleted. "
+        temp_prob = DynamicsMethods.create_problem(u0, tspan_short, sim)
+        precomp = SciMLBase.solve(temp_prob, algorithm; kwargs...)
+#=         if isa(reduction, FileReduction)
             rm(reduction.filename)
-        end
-        @info "Pre-compiled dynamics in $(precompile_time.time) seconds."
+        end =#
+        #@info "Pre-compiled dynamics in $(precompile_time.time) seconds."
     end
 
+    problem = DynamicsMethods.create_problem(u0, tspan, sim)
+    output_func = EnsembleSaver(output, savetime)
+
+    ensemble_problem = SciMLBase.EnsembleProblem(
+        problem;
+        prob_func = prob_func,
+        output_func = output_func,
+        reduction = deepcopy(reduction), # deepcopy reduction function because it's used twice
+    )
+
     if trajectories == 1
         @info "Performing 1 trajectory."
     else
         @info "Performing $trajectories trajectories."
     end
 
-    stats = dynamics(tspan)
-    log_simulation_duration(stats.time)
+    stats = @timed SciMLBase.solve(
+            ensemble_problem,
+            algorithm,
+            ensemble_algorithm;
+            trajectories,
+            kwargs...
+        )
+    output = stats.value
+    time = stats.time
+    log_simulation_duration(time)
 
     if trajectories == 1
-        return stats.value.u[1]
+        return output.u[1]
     else
-        return stats.value.u
+        return output.u
     end
 end