precompile

Project.toml

@@ -8,6 +8,7 @@ AstroSimBase = "c6a34d98-f626-4d8d-b450-a3f124eaada6"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 PhysicalMeshes = "97d9904f-034f-4fb7-aeaa-03a173434233"
+PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
@@ -17,4 +18,4 @@ CUDA = "5"
 FFTW = "1"
 PhysicalMeshes = "0.1"
 Unitful = "1"
-julia = "1.6"
+julia = "1.6"

src/PhysicalFFT.jl

@@ -2,12 +2,14 @@ module PhysicalFFT
 
 using FFTW
 using CUDA
+using PrecompileTools
 
 using Unitful
 
 using AstroSimBase
 using PhysicalMeshes
 
+
 export fft_poisson!, fft_poisson
 
 function fft_grid_kk(N, eps = 1e-6)
@@ -265,4 +267,6 @@ function fft_poisson(m::AbstractMesh, G::Number)
     fft_poisson!(m, ustrip.(m.rho .* g), m.config.boundary)
 end
 
+include("precompile.jl")
+
 end # module PhysicalFFT

src/precompile.jl

@@ -0,0 +1,78 @@
+@setup_workload begin
+    @compile_workload begin
+        sol(x::Real) = sin(2*pi*x) + sin(32*pi*x) / 256
+        init_rho(x::Real) = -4*pi*pi*sin(2*pi*x) - 4*pi*pi*sin(32*pi*x)
+
+        sol(p::PVector2D) =  sin(2*pi*p.x) * sin(2*pi*p.y) + sin(32*pi*p.x) * sin(32*pi*p.y) / 256
+        init_rho(p::PVector2D) = -8 * pi * pi * sin(2*pi*p.x) * sin(2*pi*p.y) - 8 * pi * pi * sin(32*pi*p.x) * sin(32*pi*p.y)
+
+        sol(p::PVector) =  sin(2*pi*p.x) * sin(2*pi*p.y) * sin(2*pi*p.z) + sin(32*pi*p.x) * sin(32*pi*p.y) * sin(2*pi*p.z) / 256
+        init_rho(p::PVector) = -12 * pi * pi * sin(2*pi*p.x) * sin(2*pi*p.y) * sin(2*pi*p.z) - 12 * pi * pi * sin(32*pi*p.x) * sin(32*pi*p.y) * sin(32*pi*p.z)
+
+        function test_fft1D(Nx, boundary=Periodic())
+            m = MeshCartesianStatic(;
+                xMin = 0.0,
+                xMax = 1.0,
+                Nx = Nx - 1,
+                NG = 0,
+                dim = 1,
+                boundary,
+            )
+            m.rho .= init_rho.(m.pos)
+            fft_poisson!(m, m.rho, m.config.boundary)
+            s = sol.(m.pos)
+            r = m.phi .- s
+
+            return L2norm(r)
+        end
+        test_fft1D(8, Periodic())
+        test_fft1D(8, Dirichlet())
+
+        function test_fft2D(Nx, boundary=Periodic())
+            m = MeshCartesianStatic(;
+                xMin = 0.0,
+                yMin = 0.0,
+                xMax = 1.0,
+                yMax = 1.0,
+                Nx = Nx - 1,
+                Ny = Nx - 1,
+                NG = 0,
+                dim = 2,
+                boundary,
+            )
+            m.rho .= init_rho.(m.pos)
+            fft_poisson!(m, m.rho, m.config.boundary)
+            s = sol.(m.pos)
+            r = m.phi .- s
+
+            return L2norm(r)
+        end
+        test_fft2D(8, Periodic())
+        test_fft2D(8, Dirichlet())
+
+        function test_fft3D(Nx, boundary=Periodic())
+            m = MeshCartesianStatic(;
+                xMin = 0.0,
+                yMin = 0.0,
+                zMin = 0.0,
+                xMax = 1.0,
+                yMax = 1.0,
+                zMax = 1.0,
+                Nx = Nx - 1,
+                Ny = Nx - 1,
+                Nz = Nx - 1,
+                NG = 0,
+                dim = 3,
+                boundary,
+            )
+            m.rho .= init_rho.(m.pos)
+            fft_poisson!(m, m.rho, m.config.boundary)
+            s = sol.(m.pos)
+            r = m.phi .- s
+
+            return L2norm(r)
+        end
+        test_fft3D(8, Periodic())
+        test_fft3D(8, Dirichlet())
+    end
+end