diff --git a/Project.toml b/Project.toml index 6e1c292..de5d629 100644 --- a/Project.toml +++ b/Project.toml @@ -8,6 +8,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Manifolds = "1cead3c2-87b3-11e9-0ccd-23c62b72b94e" ManifoldsBase = "3362f125-f0bb-47a3-aa74-596ffd7ef2fb" Manopt = "0fc0a36d-df90-57f3-8f93-d78a9fc72bb5" +PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" @@ -18,6 +19,7 @@ JuliaFormatter = "=2.10.1" Manifolds = "0.11.28" ManifoldsBase = "2.3.5" Manopt = "0.6" +PrecompileTools = "1" ProgressMeter = "1.11.0" RecursiveArrayTools = "4.3" TensorOperations = "5.6" diff --git a/docs/src/PIPELINE.md b/docs/src/PIPELINE.md index b1ef8bd..18868f8 100644 --- a/docs/src/PIPELINE.md +++ b/docs/src/PIPELINE.md @@ -42,7 +42,8 @@ Most optimization APIs share this core pattern: 3. Convert to a public result struct `_solve_model` lives in `src/solvers/solve_dispatch.jl` and is the common -symbol-to-solver dispatch layer (`:rgd`, `:rcg`, `:lbfgs`, `:als`, `:btd_tsd`). +symbol-to-solver dispatch layer (`:rgd`, `:rgd_fixed`, `:rcg`, `:lbfgs`, `:lm`, +`:als`, `:btd_tsd`). ## API flows @@ -58,7 +59,7 @@ symbol-to-solver dispatch layer (`:rgd`, `:rcg`, `:lbfgs`, `:als`, `:btd_tsd`). Notes: - `:als` means CP-ALS. -- Manifold solvers (`:rgd`, `:rgd_fixed`, `:rcg`, `:lbfgs`) share dispatch with other pipelines. +- Manifold solvers (`:rgd`, `:rgd_fixed`, `:rcg`, `:lbfgs`, `:lm`) share dispatch with other pipelines. - For `solver != :als`, `init = :auto` resolves to `:alswarm`, so CPD and NNCPD start from an ALS warm point before manifold refinement. - Generic `approx(...)` does not use CPD's ALS warm-start path unless it auto-routes to `cpd(...)`. diff --git a/src/TensorKitchen.jl b/src/TensorKitchen.jl index 58cd1d4..0ec35ef 100644 --- a/src/TensorKitchen.jl +++ b/src/TensorKitchen.jl @@ -6,6 +6,7 @@ using Random using ManifoldsBase using Manifolds using Manopt +using PrecompileTools import ProgressMeter using RecursiveArrayTools using TensorOperations @@ -19,6 +20,7 @@ include("manifolds.jl") include("decompositions.jl") include("backend.jl") include("api.jl") +include("precompile.jl") function __init__() return nothing diff --git a/src/api/cpd.jl b/src/api/cpd.jl index 2339fa1..363be15 100644 --- a/src/api/cpd.jl +++ b/src/api/cpd.jl @@ -988,20 +988,20 @@ If `r` is omitted, uses the smallest tensor mode as a heuristic rank. - `:random`: Uses a random initial point. - `:hosvd`: Uses a HOSVD initial point. * `solver = :rgd`: Sets the algorithm for refinement. Possible options are: - - `rgd` (default): Riemannian gradient descent - - `rgd_fixed`: Riemannian gradient descent with fixed step size - - `rcg`: Riemannian conjugate gradient - - `lbfgs`: Limited-memory Riemannian quasi-Newton - - `lm`: Levenberg-Marquardt using residual/Jacobian least squares - - `als`: Alternating Least Squares + - `rgd` (default): Riemannian gradient descent + - `rgd_fixed`: Riemannian gradient descent with fixed step size + - `rcg`: Riemannian conjugate gradient + - `lbfgs`: Limited-memory Riemannian quasi-Newton + - `lm`: Levenberg-Marquardt using residual/Jacobian least squares + - `als`: Alternating Least Squares ## Extended Options * `p0 = nothing`: Explicit initial point. If provided, it overrides the default initial point. * `:alswarm`: ALS warm start option. - `warm_init = TuckerInit()`: Before finding the warm start initial point, this sets the good starting point for ALS. - `warm_steps = 500`: Once finding the best initial point from warm_init, it runs this many ALS iterations to refine the initial point. -* `maxiter = 500`: Maximum number of Riemannian gradient descent iterations. -* `stepsize`: Initial step size for line search in Riemannian gradient descent. Defaults to `1.0` for ordinary CPD and `0.01` for the nonnegative route. +* `maxiter = 500`: Maximum number of refinement iterations. +* `stepsize`: Initial step size for first-order line-search solvers. Defaults to `1.0` for ordinary CPD and `0.01` for the nonnegative route. * `armijo_alpha_min = 1e-8`: Minimum Armijo line-search step size for `solver = :rgd`. * `tol = 1e-6`: Convergence tolerance. * `gradient_mode = :riemannian`: Gradient rule for manifold solvers. diff --git a/src/api/nncpd.jl b/src/api/nncpd.jl index 20ac168..84def07 100644 --- a/src/api/nncpd.jl +++ b/src/api/nncpd.jl @@ -14,6 +14,8 @@ end Computes a nonnegative rank-`r` CP approximation of `A` in two steps: (1) the first step finds an initial point; (2) the second step refines the initial point. Returns a [`CPDResult`](@ref). If `r` is omitted, uses the smallest tensor mode as a heuristic rank. `cpd(A, r; nonnegative=true, ...)` routes here and adopts the same effective defaults. +With the default manifold refinement path, this means an ALS warm start followed +by `solver = :rgd`. ## Options The options are the same as for [`cpd`](@ref). diff --git a/src/precompile.jl b/src/precompile.jl new file mode 100644 index 0000000..88fd48e --- /dev/null +++ b/src/precompile.jl @@ -0,0 +1,113 @@ +if Base.VERSION >= v"1.9" + using PrecompileTools: @compile_workload, @setup_workload + + @setup_workload begin + rng = MersenneTwister(7) + A = randn(rng, 4, 3, 2) + B = abs.(A) + A_med = randn(rng, 6, 5, 4) + + model = JoinModel(A, 2; geometry = :canonical) + p0 = _solver_point(manifold(model), initial_point(model, :random; verbose = false)) + model_med = JoinModel(A_med, 2; geometry = :canonical) + p0_med = _solver_point( + manifold(model_med), + initial_point(model_med, :random; verbose = false), + ) + + nn_model = JoinModel(B, 2; nonnegative = true, geometry = :softplus_metric) + p0_nn = _solver_point( + manifold(nn_model), + initial_point(nn_model, :random; verbose = false), + ) + + # The first LM call is dominated by JIT compilation in the operator-action + # path. Precompiling a minimal canonical and softplus NNCP workload moves + # that cost to package precompile time instead of first user execution. + @compile_workload begin + # Direct LM solve path on an already-prepared canonical CP model. + solve( + LMSolver(), + model; + init = :random, + p0 = p0, + maxiter = 1, + tol = 1e-6, + verbose = false, + return_stats = true, + ) + solve( + LMSolver(), + model_med; + init = :random, + p0 = p0_med, + maxiter = 1, + tol = 1e-6, + verbose = false, + return_stats = true, + ) + cpd( + A, + 2; + solver = :lm, + init = :random, + p0 = p0, + maxiter = 1, + tol = 1e-6, + verbose = false, + ) + cpd( + A_med, + 2; + solver = :lm, + init = :random, + p0 = p0_med, + maxiter = 1, + tol = 1e-6, + verbose = false, + ) + # Frontend LM route with ALS warm start, which is the path users hit + # most often when they request cpd(...; solver=:lm). + cpd( + A, + 2; + solver = :lm, + init = :alswarm, + warm_steps = 2, + maxiter = 1, + tol = 1e-6, + verbose = false, + ) + cpd( + A_med, + 2; + solver = :lm, + init = :alswarm, + warm_steps = 2, + maxiter = 1, + tol = 1e-6, + verbose = false, + ) + nncpd( + B, + 2; + solver = :lm, + init = :random, + p0 = p0_nn, + maxiter = 1, + tol = 1e-6, + verbose = false, + ) + nncpd( + B, + 2; + solver = :lm, + init = :alswarm, + warm_steps = 2, + maxiter = 1, + tol = 1e-6, + verbose = false, + ) + end + end +end