Merge pull request #577 from ChrisRackauckas-Claude/precompile-improvements-20260101-033127

Add PrecompileTools workload to improve TTFX

Author: ChrisRackauckas

Project.toml

@@ -12,6 +12,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a"
+PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -33,6 +34,7 @@ MLUtils = "0.3, 0.4"
 ModelingToolkit = "11"
 OrdinaryDiffEqTsit5 = "1"
 Parameters = "0.12"
+PrecompileTools = "1"
 ProgressMeter = "1.6"
 QuadGK = "2.4"
 RecipesBase = "1"

src/DataDrivenDiffEq.jl

@@ -145,4 +145,7 @@ export get_algorithm, get_results, get_basis, is_converged, get_problem
 include("./utils/plot_recipes.jl")
 include("./utils/build_basis.jl")
 
+# Precompilation workload to improve startup time and TTFX
+include("./precompilation.jl")
+
 end # module

src/precompilation.jl

@@ -0,0 +1,46 @@
+# Precompilation workload for DataDrivenDiffEq
+# This file contains code that will be executed during precompilation to
+# improve startup time and time-to-first-execution (TTFX)
+
+using PrecompileTools
+
+@setup_workload begin
+    # Set up minimal test data and variables
+    # Note: We avoid heavy computations here, just trigger code paths
+    @compile_workload begin
+        # Create symbolic variables for Basis construction
+        # This is the main TTFX bottleneck (33+ seconds without precompilation)
+        @variables x y t
+
+        # Precompile Basis construction - the most important operation
+        # Use simple expressions to minimize precompilation time
+        basis = Basis([x, y, x * y, x^2, y^2], [x, y])
+
+        # Precompile basis evaluation (both scalar and vector forms)
+        u_test = [1.0, 2.0]
+        p_test = Float64[]
+        t_test = 0.0
+        _ = basis(u_test, p_test, t_test)
+
+        # Precompile problem constructors
+        X = randn(2, 10)
+        Y = randn(2, 10)
+        t_span = collect(range(0.0, 1.0, length = 10))
+
+        # DirectDataDrivenProblem
+        prob_direct = DirectDataDrivenProblem(X, Y)
+
+        # DiscreteDataDrivenProblem
+        prob_discrete = DiscreteDataDrivenProblem(X)
+
+        # ContinuousDataDrivenProblem (uses collocation which takes time)
+        prob_cont = ContinuousDataDrivenProblem(X, t_span)
+
+        # Precompile basis generators
+        poly_basis = polynomial_basis([x, y], 2)
+        mono_basis = monomial_basis([x, y], 2)
+
+        # Precompile basis evaluation with problem
+        _ = basis(prob_direct)
+    end
+end