opt: Use PrecomileTools

Use PrecompileTools to precompile functions required for a very basic
simulation.  This should help with TTFX.

Could be extended to `PencilGrids`, functions inside `simulate!` that
only run intermittently, and summary statistics.

Project.toml

@@ -14,6 +14,7 @@ MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 NPZ = "15e1cf62-19b3-5cfa-8e77-841668bca605"
 PencilArrays = "0e08944d-e94e-41b1-9406-dcf66b6a9d2e"
 PencilFFTs = "4a48f351-57a6-4416-9ec4-c37015456aae"
+PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Strided = "5e0ebb24-38b0-5f93-81fe-25c709ecae67"
 
@@ -28,6 +29,7 @@ MPI = "0.20"
 NPZ = "0.4"
 PencilArrays = "0.11, 0.12, 0.13, 0.14, 0.15, 0.16, 0.17, 0.18"
 PencilFFTs = "0.12, 0.13, 0.14, 0.15"
+PrecompileTools = "1.2"
 Statistics = "1.8, 1.9"
 Strided = "2.0.4"
 julia = "1.8, 1.9"

src/UltraDark.jl

@@ -5,6 +5,7 @@ using AbstractFFTs: fftfreq, rfftfreq
 import Folds
 using LinearAlgebra
 using PencilFFTs, MPI
+using PrecompileTools
 
 using FFTW
 
@@ -273,4 +274,31 @@ function simulate!(
 
 end
 
+@setup_workload begin
+
+    resol = 64
+    output_dir = mktempdir()
+
+    @compile_workload begin
+        output_times = [0, 1e-9]
+
+        output_config = OutputConfig(
+            output_dir,
+            output_times;
+            box = true,
+            slice = true,
+            npy = true,
+            h5 = true,
+        )
+
+        options = Config.SimulationConfig()
+
+        grids = Grids(10.0, resol)
+        # Set ψx to something non-zero
+        grids.ψx .= (grids.x .^ 2 .+ grids.y .^ 2 .+ grids.z .^ 2) .^ 0.5 ./ 1e9
+
+        simulate!(grids, options, output_config)
+    end
+end
+
 end # module