diff --git a/src/api/cpd.jl b/src/api/cpd.jl index 69c7701..43a555c 100644 --- a/src/api/cpd.jl +++ b/src/api/cpd.jl @@ -1,39 +1,6 @@ # api/cpd.jl — user-facing CP decomposition entry points export cpd -Base.@kwdef struct CPDOpts{T<:Real} - solver::Symbol = :rgd - geometry::Symbol = :canonical - init::Any = :auto - maxiter::Int = (solver in (:rgd, :rcg)) ? 2000 : 500 - tol::T = 1e-6 - gradient_mode::Symbol = :riemannian - nonnegative::Bool = false - stepsize::T = 1.0 - normalization::Symbol = :separate - pullback_eps::T = 1e-8 - verbose::Bool = true -end - -function _validate_opts(opts::CPDOpts) - if opts.solver == :als && opts.geometry != :canonical - throw( - ArgumentError("ALS solver requires geometry=:canonical. Got: $(opts.geometry)"), - ) - end - if opts.nonnegative && - opts.solver in (:rgd, :rcg) && - opts.geometry ∉ (:squaring_metric, :softplus_metric) - @warn "Nonnegative RGD/RCG usually requires geometry=:softplus_metric or :squaring_metric." - end - return opts -end - -@inline function _resolve_cpd_init(init, solver::Symbol) - init == :auto || return init - return solver == :als ? :tucker : :alswarm -end - function _pullback_eps_value(::Type{T}, pullback_eps) where {T<:AbstractFloat} ε = T(pullback_eps) isfinite(ε) && ε > zero(T) || @@ -74,7 +41,7 @@ function _cpd_als_warm_then_pack( A = tensor(target) r = inner isa RankRCPDModel ? inner.r : 1 T = eltype(A) - warm_init = _resolve_cpd_init(init.base_init, :als) + warm_init = init.base_init == :auto ? :tucker : init.base_init if r == 1 warm_out = fit_cp_als( A, @@ -104,7 +71,7 @@ function _cpd_als_warm_then_pack( warm_result = _solve_model( warm_model; init = warm_init, - solver = :als, + solver = ALSSolver(), maxiter = init.nsteps, stepsize = one(T), tol = tol, @@ -178,6 +145,53 @@ function initial_point( ) end +function _run_cpd_solver( + model; + init_eff, + p0, + solver::AbstractSolver, + maxiter::Int, + stepsize, + tol, + gradient_mode, + normalization, + warm_normalization, + verbose::Bool, + vector_transport_method, + pullback_eps, + kwargs..., +) + p_solve = if init_eff isa ALSWarmStartInit && isnothing(p0) && !(solver isa ALSSolver) + _cpd_als_warm_then_pack( + model, + init_eff; + tol, + normalization = warm_normalization, + verbose, + pullback_eps, + kwargs..., + ) + else + _pack_cpd_explicit_p0(model, p0) + end + + return _solve_model( + model; + init = init_eff, + p0 = p_solve, + solver = solver, + maxiter, + stepsize, + tol, + gradient_mode, + normalization, + verbose, + refinement_verbose = verbose, + vector_transport_method, + kwargs..., + ) +end + function _cpd_impl( A::AbstractArray{T,N}, r::Int; @@ -205,58 +219,49 @@ function _cpd_impl( "softplus_beta has been removed. Use pullback_eps to tune softplus pullback regularization.", ), ) - dims = size(A) - init_resolved = _resolve_cpd_init(init, solver) + + solver_obj = _solver_object(solver, stepsize; kwargs...) + init_resolved = init == :auto ? (solver_obj isa ALSSolver ? :tucker : :alswarm) : init init_eff = init_resolved == :alswarm ? ALSWarmStartInit(warm_steps; base_init = warm_init) : init_resolved geometry_eff = _is_native_rankr_geometry(geometry) ? :native : geometry - manifold_solvers = (:rgd, :rgd_fixed, :rcg) - als_family = (:als,) - normalization_eff = - normalization == :auto ? - ( - solver ∈ als_family ? + pullback_eps_eff = _pullback_eps_value(T, pullback_eps) + if normalization == :auto + normalization_eff = + solver_obj isa ALSSolver ? (nonnegative ? NoNormalization() : SeparateLambdaNormalization()) : NoNormalization() - ) : _normalization_policy(normalization) - als_normalization_eff = - normalization == :auto ? - (nonnegative ? NoNormalization() : SeparateLambdaNormalization()) : - _normalization_policy(normalization) - pullback_eps_eff = _pullback_eps_value(T, pullback_eps) + warm_normalization_eff = + nonnegative ? NoNormalization() : SeparateLambdaNormalization() + else + normalization_eff = _normalization_policy(normalization) + warm_normalization_eff = normalization_eff + end r >= 1 || throw(ArgumentError("rank r must be >= 1, got r=$r")) - nonnegative && - solver ∉ (:als, :rgd, :rgd_fixed, :rcg) && - throw( - ArgumentError( - "nonnegative=true requires solver=:als, :rgd, :rgd_fixed, or :rcg. Got solver=$solver.", - ), - ) - geometry_eff ∈ (:native, :canonical, :squaring_metric, :softplus_metric) || throw( + solver_obj isa Union{ALSSolver,RGDSolver,RGDFixedSolver,RCGSolver} || throw( + ArgumentError( + "Unsupported CPD solver $(typeof(solver_obj)). Use :als, :rgd, :rgd_fixed, or :rcg.", + ), + ) + geometry_eff in (:native, :canonical, :squaring_metric, :softplus_metric) || throw( ArgumentError( - "Unknown geometry=$geometry. Use :native, :canonical, :squaring_metric, or :softplus_metric (regularized pullback-style geometries require nonnegative=true).", + "Unknown geometry=$geometry. Use :native, :canonical, :squaring_metric, or :softplus_metric.", ), ) - (geometry_eff ∉ (:squaring_metric, :softplus_metric) || nonnegative) || + if geometry_eff in (:squaring_metric, :softplus_metric) && !nonnegative throw(ArgumentError("geometry=$geometry_eff requires nonnegative=true.")) - solver ∈ manifold_solvers || - solver ∈ als_family || - throw( - ArgumentError( - "Unknown solver=$solver. Use a manifold solver (:rgd, :rgd_fixed, :rcg) or CP-ALS (:als).", - ), - ) - if solver ∈ als_family + end + if solver_obj isa ALSSolver geometry_eff == :canonical || throw( ArgumentError( - "solver=$solver does not use manifold geometry. Pass geometry=:canonical (default) or switch to a manifold solver (:rgd, :rgd_fixed, :rcg) for geometry=:native/:squaring_metric/:softplus_metric.", + "solver=:als does not use manifold geometry. Use geometry=:canonical.", ), ) gradient_mode == :riemannian || throw( ArgumentError( - "solver=$solver does not use gradient_mode. Pass gradient_mode=:riemannian (default) or switch to a manifold solver.", + "solver=:als does not use gradient_mode. Use gradient_mode=:riemannian.", ), ) end @@ -271,89 +276,35 @@ function _cpd_impl( use_pullback_metric = (geometry_eff == :squaring_metric), pullback_eps = pullback_eps_eff, ) - result_rgd = with_phase_progress() do - p_solve = - if init_eff isa ALSWarmStartInit && isnothing(p0) && solver ∈ manifold_solvers - _cpd_als_warm_then_pack( - model, - init_eff; - tol, - normalization = als_normalization_eff, - verbose, - pullback_eps = pullback_eps_eff, - kwargs..., - ) - else - _pack_cpd_explicit_p0(model, p0) - end - _solve_model( + + raw_result = with_phase_progress() do + _run_cpd_solver( model; - init = init_eff, - p0 = p_solve, - solver = solver, - maxiter = maxiter, - stepsize = stepsize, - tol = tol, - gradient_mode = gradient_mode, + init_eff, + p0, + solver = solver_obj, + maxiter, + stepsize, + tol, + gradient_mode, normalization = normalization_eff, - verbose = verbose, - refinement_verbose = verbose, - vector_transport_method = vector_transport_method, + warm_normalization = warm_normalization_eff, + verbose, + vector_transport_method, + pullback_eps = pullback_eps_eff, nonnegative, kwargs..., ) end - result = result_rgd - if nonnegative && geometry_eff ∈ (:squaring_metric, :softplus_metric) - result = _merge_res_solver_info(result, (nncp_pullback_eps = pullback_eps_eff,)) - end - return _to_cpd_result(model, result, dims, r) -end - -#### MAIN CPD #### - -function cpd(A::AbstractArray{T,N}, r::Int, opts::CPDOpts) where {T,N} - _validate_opts(opts) - - if opts.nonnegative - return nncpd( - A, - r; - solver = opts.solver, - init = opts.init, - geometry = opts.geometry, - maxiter = opts.maxiter, - stepsize = opts.stepsize, - tol = opts.tol, - gradient_mode = opts.gradient_mode, - normalization = opts.normalization, - pullback_eps = opts.pullback_eps, - verbose = opts.verbose, - ) - end - return cpd( - A, - r; - init = opts.init, - solver = opts.solver, - geometry = opts.geometry, - maxiter = opts.maxiter, - tol = opts.tol, - gradient_mode = opts.gradient_mode, - pullback_eps = opts.pullback_eps, - verbose = opts.verbose, - ) + result = + nonnegative && geometry_eff in (:squaring_metric, :softplus_metric) ? + _merge_res_solver_info(raw_result, (nncp_pullback_eps = pullback_eps_eff,)) : + raw_result + return _to_cpd_result(model, result, size(A), r) end -function cpd(A::AbstractArray{T,N}, r::Int; kwargs...) where {T,N} - valid_keys = fieldnames(CPDOpts) - for k in keys(kwargs) - k in valid_keys || throw(ArgumentError("Unknown keyword argument: $k")) - end - opts = CPDOpts{eltype(A)}(; kwargs...) - return cpd(A, r, opts) -end +#### MAIN CPD #### function cpd( A::AbstractArray{T,N}; @@ -455,11 +406,12 @@ function cpd( kwargs..., ) where {T<:AbstractFloat,N} if nonnegative + solver_obj = _solver_object(solver, stepsize; kwargs...) # Align effective defaults with nncpd() on the nonnegative route. # Explicitly passed non-default values are preserved. init_nn = init == :tucker ? :alswarm : init warm_steps_nn = warm_steps - geometry_nn = if solver ∈ (:als,) + geometry_nn = if solver_obj isa ALSSolver :canonical elseif geometry == :canonical :softplus_metric @@ -474,7 +426,7 @@ function cpd( p0 = p0, warm_steps = warm_steps_nn, warm_init = warm_init, - solver = solver, + solver = solver_obj, geometry = geometry_nn, maxiter = maxiter, stepsize = stepsize_nn, diff --git a/src/api/nncpd.jl b/src/api/nncpd.jl index 409ef76..947163b 100644 --- a/src/api/nncpd.jl +++ b/src/api/nncpd.jl @@ -56,7 +56,7 @@ function nncpd( warm_steps = 500, warm_init = TuckerInit(), solver = :rgd, - geometry = (solver == :als ? :canonical : :softplus_metric), + geometry = nothing, maxiter = 500, stepsize = 0.01, tol = 1e-6, @@ -69,6 +69,10 @@ function nncpd( pullback_eps = 1e-8, kwargs..., ) where {T<:AbstractFloat,N} + solver_obj = _solver_object(solver, stepsize; kwargs...) + geometry_eff = + isnothing(geometry) ? (solver_obj isa ALSSolver ? :canonical : :softplus_metric) : + geometry return _cpd_impl( A, r; @@ -76,8 +80,8 @@ function nncpd( p0 = p0, warm_steps = warm_steps, warm_init = warm_init, - solver = solver, - geometry = geometry, + solver = solver_obj, + geometry = geometry_eff, maxiter = maxiter, stepsize = stepsize, tol = tol, diff --git a/src/solvers/solve_dispatch.jl b/src/solvers/solve_dispatch.jl index fc8c537..0ee6b71 100644 --- a/src/solvers/solve_dispatch.jl +++ b/src/solvers/solve_dispatch.jl @@ -1,4 +1,12 @@ -# solvers/solve_dispatch.jl — shared symbol-based solver dispatch +# solvers/solve_dispatch.jl — shared solver normalization and dispatch + +function _solver_object(solver, ::Real; kwargs...) + throw( + ArgumentError( + "Unsupported solver specification $(typeof(solver)). Use a solver symbol such as :als, :rgd, :rgd_fixed, :rcg, :lbfgs, or :btd_tsd, or pass an AbstractSolver object.", + ), + ) +end function _solver_object(solver::Symbol, stepsize::Real; kwargs...) solvers = ( @@ -26,14 +34,6 @@ function _solver_object(solver::Symbol, stepsize::Real; kwargs...) armijo_sufficient_decrease = get(kwargs, :armijo_sufficient_decrease, 1e-4), armijo_alpha_min = get(kwargs, :armijo_alpha_min, 1e-12), ), - tsd = () -> BTDTSDSolver(; - stepsize, - schedule = get(kwargs, :schedule, :cyclic), - block_repeats = get(kwargs, :block_repeats, 1), - armijo_contraction = get(kwargs, :armijo_contraction, 0.5), - armijo_sufficient_decrease = get(kwargs, :armijo_sufficient_decrease, 1e-4), - armijo_alpha_min = get(kwargs, :armijo_alpha_min, 1e-12), - ), ) f = get(solvers, solver) do throw( @@ -45,6 +45,8 @@ function _solver_object(solver::Symbol, stepsize::Real; kwargs...) return f() end +_solver_object(solver::AbstractSolver, ::Real; kwargs...) = solver + function _solve_with_solver( solver_obj::AbstractROSolver, model; @@ -80,10 +82,17 @@ function _solve_with_solver( p0 = nothing, maxiter::Int, tol::Real, + gradient_mode::Symbol = :riemannian, normalization = SeparateLambdaNormalization(), verbose::Bool, kwargs..., ) + gradient_mode == :riemannian || throw( + ArgumentError( + "ALS solvers do not use gradient_mode. Use gradient_mode=:riemannian.", + ), + ) + return solve( solver_obj, model; @@ -102,14 +111,14 @@ end _solve_model(model; solver, init, maxiter, stepsize, tol, kwargs...) Top-level internal solver dispatcher shared by CPD, BTD, NNCPD, and generic -`approx`. Converts the public solver symbol into a concrete solver and returns -a result-like `NamedTuple`. +`approx`. Accepts either a public solver symbol or a concrete `AbstractSolver` +object, normalizes it to a solver object, and returns a result-like `NamedTuple`. """ function _solve_model( model::AbstractDecompositionModel; init, p0 = nothing, - solver::Symbol, + solver, maxiter::Int, stepsize::Real, tol::Real, diff --git a/test/basic_tests.jl b/test/basic_tests.jl index c828ba4..e49344a 100644 --- a/test/basic_tests.jl +++ b/test/basic_tests.jl @@ -235,6 +235,43 @@ end @test nncpd_res isa CPDResult @test nncpd_res.solver == :rgd + cpd_rgd_object = cpd( + A, + 2; + solver = RGDSolver(), + init = :alswarm, + warm_steps = 2, + maxiter = 1, + verbose = false, + ) + @test cpd_rgd_object isa CPDResult + @test cpd_rgd_object.solver == :rgd + + cpd_als_object = + cpd(A, 2; solver = ALSSolver(), init = :tucker, maxiter = 1, verbose = false) + @test cpd_als_object isa CPDResult + @test cpd_als_object.solver == :cp_als + @test_throws ArgumentError cpd( + A, + 2; + solver = ALSSolver(), + gradient_mode = :egrad_project, + maxiter = 1, + verbose = false, + ) + @test_throws ArgumentError cpd(A, 2; solver = "rgd", maxiter = 1, verbose = false) + + nncpd_als_object = nncpd( + abs.(A), + 2; + solver = ALSSolver(), + init = :tucker, + maxiter = 1, + verbose = false, + ) + @test nncpd_als_object isa CPDResult + @test nncpd_als_object.solver == :cp_als + btd_res = btd(A, 2, (3, 2, 2); maxiter = 1, verbose = false) @test btd_res isa BTDResult @test btd_res.solver == :rgd