From eef8c4a44b8f3a7b00e88a52a068c5124f0ea14a Mon Sep 17 00:00:00 2001 From: Se Eun Choi Date: Fri, 19 Jun 2026 04:03:36 +0200 Subject: [PATCH 1/7] Add Riemannian Levenberg-Marquardt solver via Manopt. Wire LMSolver into solve dispatch and cpd/approx/btd APIs, add orthonormal coordinate support on pullback metrics for Jacobian assembly, and cover residual/Jacobian consistency with tests. --- src/api/approx.jl | 3 +- src/api/btd.jl | 1 + src/api/cpd.jl | 9 +- src/backend.jl | 1 + src/core/types.jl | 2 +- src/cpd/core/cp_cost.jl | 2 +- src/manifolds/softplus_metric.jl | 45 +++- src/manifolds/squaring_metric.jl | 43 ++- src/solvers/abstract.jl | 1 + src/solvers/lm.jl | 446 +++++++++++++++++++++++++++++++ src/solvers/manopt_helpers.jl | 2 +- src/solvers/solve_dispatch.jl | 14 +- test/basic_tests.jl | 77 ++++++ 13 files changed, 631 insertions(+), 15 deletions(-) create mode 100644 src/solvers/lm.jl diff --git a/src/api/approx.jl b/src/api/approx.jl index 3aa4866..e380061 100644 --- a/src/api/approx.jl +++ b/src/api/approx.jl @@ -127,11 +127,12 @@ For the generic join path: - `rgd_fixed`: Riemannian gradient descent with fixed step size - `rcg`: Riemannian conjugate gradient - `lbfgs`: Limited-memory quasi-Newton + - `lm`: Levenberg-Marquardt on residual/Jacobian least squares ##Notes## * `:als` is not a solver option for `approx(...)`. However, if `approx(...)` auto-routes to `cpd(...)` or `btd(...)`, then those specialized pipelines may support ALS separately. * `warm_steps` and `warm_init` are not part of the generic `approx(...)` path. Generic joins start from random initial point and then use manifold solvers for refinement. -* For generic mixed joins, use manifold solvers such as `:rgd`, `:rcg`, or `:lbfgs`. +* For generic mixed joins, use manifold solvers such as `:rgd`, `:rcg`, `:lbfgs`, or `:lm`. """ function _approx_manifold_collection( dispatch::AutoApproxDispatch, diff --git a/src/api/btd.jl b/src/api/btd.jl index 97d6f02..1127881 100644 --- a/src/api/btd.jl +++ b/src/api/btd.jl @@ -159,6 +159,7 @@ refines it. Returns a [`BTDResult`](@ref). - `:als`: Alternating least squares. - `:rcg`: Riemannian conjugate gradient. - `:lbfgs`: Limited-memory quasi-Newton refinement. + - `:lm`: Levenberg-Marquardt refinement. ## Extended Options diff --git a/src/api/cpd.jl b/src/api/cpd.jl index e03ffab..1d07f56 100644 --- a/src/api/cpd.jl +++ b/src/api/cpd.jl @@ -134,7 +134,7 @@ function _component_energy_summary(energies::AbstractVector{<:Real}) top1 = ordered[1], top2 = sum(@view ordered[1:min(2, r)]), top3 = sum(@view ordered[1:min(3, r)]), - effective = hhi > 0 ? inv(hhi) : NaN, + effective = hhi > 0 ? one(hhi) / hhi : NaN, argmax_component = argmax(shares), ) end @@ -579,13 +579,13 @@ end function _validate_cpd_solver_supported(solver::AbstractSolver) throw( ArgumentError( - "Unsupported CPD solver $(typeof(solver)). Use :als, :rgd, :rgd_fixed, :rcg, or :lbfgs.", + "Unsupported CPD solver $(typeof(solver)). Use :als, :rgd, :rgd_fixed, :rcg, :lbfgs, or :lm.", ), ) end _validate_cpd_solver_supported( - ::Union{ALSSolver,RGDSolver,RGDFixedSolver,RCGSolver,LBFGSSolver}, + ::Union{ALSSolver,RGDSolver,RGDFixedSolver,RCGSolver,LBFGSSolver,LMSolver}, ) = nothing function _validate_cpd_solver_options( @@ -750,7 +750,7 @@ end function _cpd_manifold_grad_tol( model::JoinModel{<:AbstractFloat,<:CPDBackend}, - solver::Union{RGDSolver,RGDFixedSolver,RCGSolver,LBFGSSolver}, + solver::Union{RGDSolver,RGDFixedSolver,RCGSolver,LBFGSSolver,LMSolver}, tol::Real, ) return tol @@ -992,6 +992,7 @@ If `r` is omitted, uses the smallest tensor mode as a heuristic rank. - `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 diff --git a/src/backend.jl b/src/backend.jl index 5a96538..3f497fc 100644 --- a/src/backend.jl +++ b/src/backend.jl @@ -8,6 +8,7 @@ include("solvers/rgd.jl") include("solvers/btd_tsd.jl") include("solvers/rcg.jl") include("solvers/lbfgs.jl") +include("solvers/lm.jl") include("results/reconstruct.jl") include("results/rel_error.jl") include("solvers/solve_dispatch.jl") diff --git a/src/core/types.jl b/src/core/types.jl index d39ee6a..9549dcc 100644 --- a/src/core/types.jl +++ b/src/core/types.jl @@ -446,7 +446,7 @@ function LinearAlgebra.cond(R::CPDResult{T}) where {T<:AbstractFloat} U = hcat(Us...) s = svdvals(U) smin = minimum(s) - return iszero(smin) ? T(Inf) : inv(smin) + return iszero(smin) ? T(Inf) : one(T) / smin end """ diff --git a/src/cpd/core/cp_cost.jl b/src/cpd/core/cp_cost.jl index c99547b..a10ef67 100644 --- a/src/cpd/core/cp_cost.jl +++ b/src/cpd/core/cp_cost.jl @@ -69,7 +69,7 @@ end end @inline function _softplus_derivative(x::Real) - x >= 0 ? inv(one(x) + exp(-x)) : begin + x >= 0 ? one(x) / (one(x) + exp(-x)) : begin ex = exp(x) ex / (one(x) + ex) end diff --git a/src/manifolds/softplus_metric.jl b/src/manifolds/softplus_metric.jl index 6edfb39..ed85a5f 100644 --- a/src/manifolds/softplus_metric.jl +++ b/src/manifolds/softplus_metric.jl @@ -7,7 +7,7 @@ export SoftplusEuclidean, softplus_metric_inverse -@inline _sp_sigmoid(x::Real) = x >= 0 ? inv(one(x) + exp(-x)) : begin +@inline _sp_sigmoid(x::Real) = x >= 0 ? one(x) / (one(x) + exp(-x)) : begin ex = exp(x) ex / (one(x) + ex) end @@ -57,8 +57,7 @@ function softplus_metric_diag(M::SoftplusEuclidean, p::AbstractVector) end function softplus_metric_inverse(M::SoftplusEuclidean, p::AbstractVector, X::AbstractVector) - g_inv = inv.(softplus_metric_diag(M, p)) - return g_inv .* X + return X ./ softplus_metric_diag(M, p) end pullback_metric_inverse(M::SoftplusEuclidean, p::AbstractVector, X::AbstractVector) = @@ -68,3 +67,43 @@ function ManifoldsBase.inner(M::SoftplusEuclidean, p, X::AbstractVector, Y::Abst g = softplus_metric_diag(M, p) return dot(X, g .* Y) end + +function ManifoldsBase.get_coordinates_orthonormal( + M::SoftplusEuclidean, + p::AbstractVector, + X::AbstractVector, + ::ManifoldsBase.RealNumbers, +) + return sqrt.(softplus_metric_diag(M, p)) .* X +end + +function ManifoldsBase.get_coordinates_orthonormal!( + M::SoftplusEuclidean, + c, + p::AbstractVector, + X::AbstractVector, + ::ManifoldsBase.RealNumbers, +) + c .= sqrt.(softplus_metric_diag(M, p)) .* X + return c +end + +function ManifoldsBase.get_vector_orthonormal( + M::SoftplusEuclidean, + p::AbstractVector, + c::AbstractVector, + ::ManifoldsBase.RealNumbers, +) + return c ./ sqrt.(softplus_metric_diag(M, p)) +end + +function ManifoldsBase.get_vector_orthonormal!( + M::SoftplusEuclidean, + X, + p::AbstractVector, + c::AbstractVector, + ::ManifoldsBase.RealNumbers, +) + X .= c ./ sqrt.(softplus_metric_diag(M, p)) + return X +end diff --git a/src/manifolds/squaring_metric.jl b/src/manifolds/squaring_metric.jl index 696d1c3..b1d61a3 100644 --- a/src/manifolds/squaring_metric.jl +++ b/src/manifolds/squaring_metric.jl @@ -72,8 +72,7 @@ function pullback_metric_diag(M::SqEuclidean, p::AbstractVector) end function pullback_metric_inverse(M::SqEuclidean, p::AbstractVector, X::AbstractVector) - g_inv = 1.0 ./ pullback_metric_diag(M, p) - return g_inv .* X + return X ./ pullback_metric_diag(M, p) end function ManifoldsBase.inner(M::SqEuclidean, p, X::AbstractVector, Y::AbstractVector) @@ -81,6 +80,46 @@ function ManifoldsBase.inner(M::SqEuclidean, p, X::AbstractVector, Y::AbstractVe return dot(X, g .* Y) end +function ManifoldsBase.get_coordinates_orthonormal( + M::SqEuclidean, + p::AbstractVector, + X::AbstractVector, + ::ManifoldsBase.RealNumbers, +) + return sqrt.(pullback_metric_diag(M, p)) .* X +end + +function ManifoldsBase.get_coordinates_orthonormal!( + M::SqEuclidean, + c, + p::AbstractVector, + X::AbstractVector, + ::ManifoldsBase.RealNumbers, +) + c .= sqrt.(pullback_metric_diag(M, p)) .* X + return c +end + +function ManifoldsBase.get_vector_orthonormal( + M::SqEuclidean, + p::AbstractVector, + c::AbstractVector, + ::ManifoldsBase.RealNumbers, +) + return c ./ sqrt.(pullback_metric_diag(M, p)) +end + +function ManifoldsBase.get_vector_orthonormal!( + M::SqEuclidean, + X, + p::AbstractVector, + c::AbstractVector, + ::ManifoldsBase.RealNumbers, +) + X .= c ./ sqrt.(pullback_metric_diag(M, p)) + return X +end + # ----------------------------- # 2D benchmark objectives for squaring/nonnegative experiments # p = [x, y] diff --git a/src/solvers/abstract.jl b/src/solvers/abstract.jl index 5aa7608..bab4b9e 100644 --- a/src/solvers/abstract.jl +++ b/src/solvers/abstract.jl @@ -471,6 +471,7 @@ function _prepare_solver_problem( ) return ( + model = model, M = M, p0 = p0_local, model_cost = model_cost, diff --git a/src/solvers/lm.jl b/src/solvers/lm.jl new file mode 100644 index 0000000..6568729 --- /dev/null +++ b/src/solvers/lm.jl @@ -0,0 +1,446 @@ +# solvers/lm.jl — Riemannian Levenberg-Marquardt via Manopt nonlinear least squares +export LMSolver + +struct LMSolver <: AbstractSecondOrderROSolver + η::Float64 + damping_term_min::Float64 + β::Float64 + expect_zero_residual::Bool + linear_subsolver::Any +end + +function LMSolver(; + η::Real = 0.2, + damping_term_min::Real = 0.1, + β::Real = 5.0, + expect_zero_residual::Bool = false, + linear_subsolver = Manopt.default_lm_lin_solve!, +) + 0 < η < 1 || throw(ArgumentError("η must satisfy 0 < η < 1, got $η")) + damping_term_min > 0 || + throw(ArgumentError("damping_term_min must be > 0, got $damping_term_min")) + β > 1 || throw(ArgumentError("β must be > 1, got $β")) + return LMSolver( + Float64(η), + Float64(damping_term_min), + Float64(β), + expect_zero_residual, + linear_subsolver, + ) +end + +solver_symbol(::LMSolver) = :lm + +@inline function _lm_objective_scale( + ::Type{T}, + normA2, + normalized_objective::Bool, +) where {T} + return normalized_objective && !isnothing(normA2) && normA2 > 0 ? + one(T) / sqrt(T(normA2)) : one(T) +end + +function _ambient_tangent_vector!(out::AbstractVector, M, p, X) + emb = ManifoldsBase.embed(M, p, X) + length(emb) == length(out) || throw( + DimensionMismatch( + "Tangent embedding length $(length(emb)) does not match output length $(length(out)).", + ), + ) + copyto!(out, vec(emb)) + return out +end + +function _ambient_tangent_vector!(out::AbstractVector, M::Manifolds.Segre, p, X) + copyto!(out, _segre_tangent_tensorvec(p, X)) + return out +end + +function _join_tangent_ambient_vector!( + out::AbstractVector, + backend::Union{JoinBackend,BTDBackend}, + p, + X, +) + parts = point_parts(p) + xparts = point_parts(X) + _check_parts_len(parts, backend.r, "_join_tangent_ambient_vector!") + _check_parts_len(xparts, backend.r, "_join_tangent_ambient_vector!") + fill!(out, zero(eltype(out))) + @inbounds for k = 1:backend.r + _ambient_tangent_vector!( + backend.component_bufs[k], + backend.manifolds[k], + parts[k], + xparts[k], + ) + out .+= backend.component_bufs[k] + end + return out +end + +function _lm_raw_residual_vector( + model::JoinModel{<:AbstractFloat,<:Union{JoinBackend,BTDBackend}}, + p, +) + return copy(_join_residual!(model.backend, p)) +end + +function _lm_raw_jacobian_matrix( + model::JoinModel{<:AbstractFloat,<:Union{JoinBackend,BTDBackend}}, + M, + p; + basis = ManifoldsBase.DefaultOrthonormalBasis(), +) + p_work = _solver_point(M, p) + T = _scalar_eltype(p_work) + ambient_dim = length(tensor(model)) + d = manifold_dimension(M) + J = Matrix{T}(undef, ambient_dim, d) + coeff = zeros(T, d) + column = similar(model.backend.work_rec, T, ambient_dim) + @inbounds for j = 1:d + fill!(coeff, zero(T)) + coeff[j] = one(T) + Xj = ManifoldsBase.get_vector(M, p_work, coeff, basis) + _join_tangent_ambient_vector!(column, model.backend, p_work, Xj) + J[:, j] .= column + end + return J +end + +@inline function _cp_scaled_tangent_factors(λ, U, λ̇, U̇, ::Val{:identity}) + return λ, U, λ̇, U̇ +end + +@inline function _cp_scaled_tangent_factors(λ̃, Ũ, λ̇̃, U̇̃, ::Val{:square}) + λ = λ̃ .^ 2 + U = [Ũ[m] .^ 2 for m in eachindex(Ũ)] + λ̇ = 2 .* λ̃ .* λ̇̃ + U̇ = [2 .* Ũ[m] .* U̇̃[m] for m in eachindex(Ũ)] + return λ, U, λ̇, U̇ +end + +@inline function _cp_scaled_tangent_factors(λ̃, Ũ, λ̇̃, U̇̃, ::Val{:softplus}) + λ = _softplus_value.(λ̃) + U = [_softplus_value.(Ũ[m]) for m in eachindex(Ũ)] + λ̇ = _softplus_derivative.(λ̃) .* λ̇̃ + U̇ = [_softplus_derivative.(Ũ[m]) .* U̇̃[m] for m in eachindex(Ũ)] + return λ, U, λ̇, U̇ +end + +function _cp_rankr_tangent_tensorvec!( + out::AbstractVector{T}, + λ::AbstractVector{T}, + U::Vector{<:AbstractMatrix{T}}, + λ̇::AbstractVector{T}, + U̇::Vector{<:AbstractMatrix{T}}, +) where {T<:AbstractFloat} + fill!(out, zero(T)) + r = length(λ) + @inbounds for k = 1:r + comp = ([λ[k]], [Vector(@view U[m][:, k]) for m in eachindex(U)]...) + xcomp = ([λ̇[k]], [Vector(@view U̇[m][:, k]) for m in eachindex(U̇)]...) + out .+= _segre_tangent_tensorvec(comp, xcomp) + end + return out +end + +function _cp_rank1_tangent_tensorvec!( + out::AbstractVector{T}, + λ::T, + U::Vector{<:AbstractVector{T}}, + λ̇::T, + U̇::Vector{<:AbstractVector{T}}, +) where {T<:AbstractFloat} + comp = ([λ], U...) + xcomp = ([λ̇], U̇...) + copyto!(out, _segre_tangent_tensorvec(comp, xcomp)) + return out +end + +function _lm_raw_residual_vector(model::JoinModel{<:AbstractFloat,<:CPDBackend}, p) + return _lm_raw_residual_vector(cpd_model(model), p) +end + +function _lm_raw_jacobian_matrix( + model::JoinModel{<:AbstractFloat,<:CPDBackend}, + M, + p; + basis = ManifoldsBase.DefaultOrthonormalBasis(), +) + return _lm_raw_jacobian_matrix(cpd_model(model), M, p; basis) +end + +function _lm_raw_residual_vector(model::Rank1CPDModel{T}, p) where {T<:AbstractFloat} + return vec(embed_point(model, p)) .- vec(model.A) +end + +function _lm_raw_jacobian_matrix( + model::Rank1CPDModel{T}, + M, + p; + basis = ManifoldsBase.DefaultOrthonormalBasis(), +) where {T<:AbstractFloat} + p_work = _solver_point(M, p) + ambient_dim = length(model.A) + d = manifold_dimension(M) + J = Matrix{T}(undef, ambient_dim, d) + coeff = zeros(T, d) + λp, Up = unpack_point_rank1(p_work, model.dims) + column = Vector{T}(undef, ambient_dim) + @inbounds for j = 1:d + fill!(coeff, zero(T)) + coeff[j] = one(T) + Xj = ManifoldsBase.get_vector(M, p_work, coeff, basis) + if model.nonnegative + λ̇p, U̇p = unpack_point_rank1(Xj, model.dims) + kind = _rank1_uses_softplus_metric(model.M) ? Val(:softplus) : Val(:square) + λ, U, λ̇, U̇ = _cp_scaled_tangent_factors( + [λp], + [reshape(u, :, 1) for u in Up], + [λ̇p], + [reshape(u, :, 1) for u in U̇p], + kind, + ) + U_vec = [Vector(@view U[m][:, 1]) for m in eachindex(U)] + U̇_vec = [Vector(@view U̇[m][:, 1]) for m in eachindex(U̇)] + _cp_rank1_tangent_tensorvec!(column, λ[1], U_vec, λ̇[1], U̇_vec) + else + copyto!(column, vec(ManifoldsBase.embed(M, p_work, Xj))) + end + J[:, j] .= column + end + return J +end + +function _lm_raw_residual_vector(model::RankRCPDModel{T}, p) where {T<:AbstractFloat} + return vec(embed_point(model, p)) .- vec(model.A) +end + +function _lm_raw_jacobian_matrix( + model::RankRCPDModel{T}, + M, + p; + basis = ManifoldsBase.DefaultOrthonormalBasis(), +) where {T<:AbstractFloat} + p_work = _solver_point(M, p) + ambient_dim = length(model.A) + d = manifold_dimension(M) + J = Matrix{T}(undef, ambient_dim, d) + coeff = zeros(T, d) + column = Vector{T}(undef, ambient_dim) + if model.geometry == :native && !model.nonnegative + pparts = point_parts(p_work) + @inbounds for j = 1:d + fill!(coeff, zero(T)) + coeff[j] = one(T) + Xj = ManifoldsBase.get_vector(M, p_work, coeff, basis) + xparts = point_parts(Xj) + fill!(column, zero(T)) + for k = 1:model.r + column .+= _segre_tangent_tensorvec(pparts[k], xparts[k]) + end + J[:, j] .= column + end + return J + end + + λp, Up = + model.nonnegative ? unpack_point_rankr(p_work, model.dims, model.r) : + unpack_rankr_canonical(p_work, model.dims, model.r) + kind = + model.nonnegative ? + (model.geometry == :softplus_metric ? Val(:softplus) : Val(:square)) : + Val(:identity) + @inbounds for j = 1:d + fill!(coeff, zero(T)) + coeff[j] = one(T) + Xj = ManifoldsBase.get_vector(M, p_work, coeff, basis) + λ̇p, U̇p = + model.nonnegative ? unpack_point_rankr(Xj, model.dims, model.r) : + unpack_rankr_canonical(Xj, model.dims, model.r) + λ, U, λ̇, U̇ = _cp_scaled_tangent_factors(λp, Up, λ̇p, U̇p, kind) + _cp_rankr_tangent_tensorvec!(column, λ, U, λ̇, U̇) + J[:, j] .= column + end + return J +end + +function _lm_raw_residual_vector(model::AbstractDecompositionModel, p) + throw(ArgumentError("LMSolver residual is not implemented for model $(typeof(model)).")) +end + +function _lm_raw_jacobian_matrix(model::AbstractDecompositionModel, M, p; basis) + throw(ArgumentError("LMSolver Jacobian is not implemented for model $(typeof(model)).")) +end + +function _lm_residual_function( + model::AbstractDecompositionModel, + ::Type{T}, + normA2, + normalized_objective::Bool, +) where {T<:AbstractFloat} + scale = _lm_objective_scale(T, normA2, normalized_objective) + return (M, p) -> scale .* _lm_raw_residual_vector(model, p) +end + +function _lm_jacobian_function( + model::AbstractDecompositionModel, + ::Type{T}, + normA2, + normalized_objective::Bool; + basis = ManifoldsBase.DefaultOrthonormalBasis(), +) where {T<:AbstractFloat} + scale = _lm_objective_scale(T, normA2, normalized_objective) + return (M, p) -> scale .* _lm_raw_jacobian_matrix(model, M, p; basis) +end + +function solve_lm( + model, + model_cost, + model_egrad, + M, + p0; + maxiter::Int = 1000, + tol::Real = 1e-6, + verbose::Bool = true, + return_stats::Bool = false, + normA2 = nothing, + model_grad = nothing, + vector_transport_method::Union{ManifoldsBase.AbstractVectorTransportMethod,Nothing} = nothing, + post_step_callback = nothing, + diagnostics_recorder = nothing, + iteration_callbacks = (), + η::Real = 0.2, + damping_term_min::Real = 0.1, + β::Real = 5.0, + expect_zero_residual::Bool = false, + linear_subsolver = Manopt.default_lm_lin_solve!, + grad_tol = nothing, + normalized_objective::Bool = true, +) + setup = _prepare_manopt_solver_functions( + model_cost, + model_egrad, + M, + p0; + normA2, + model_grad, + tol, + grad_tol, + normalized_objective, + ) + p0_local = setup.p0 + T = setup.T + basis = ManifoldsBase.DefaultOrthonormalBasis() + residual = _lm_residual_function(model, T, normA2, setup.uses_relative_objective) + jacobian = _lm_jacobian_function(model, T, normA2, setup.uses_relative_objective; basis) + retraction_method = _solver_retraction_method(M, p0_local) + stopping = StopWhenAny( + StopAfterIteration(maxiter), + StopWhenGradientNormLess(setup.grad_stop_tol), + StopWhenStepsizeLess(T(tol)), + StopWhenCostRelChangeAndGradientLess(T(tol), setup.dual_grad_tol), + ) + callbacks = _manopt_callbacks( + n -> make_manopt_family_progress( + n; + enabled = verbose, + phase = :refinement, + method = "LM", + dt = 0.2, + ), + maxiter, + verbose, + setup.solver_cost, + setup.solver_grad, + M; + diagnostics_recorder, + post_step_callback, + iteration_callbacks, + ) + state = Manopt.LevenbergMarquardt( + M, + residual, + jacobian, + p0_local; + evaluation = Manopt.AllocatingEvaluation(), + function_type = Manopt.FunctionVectorialType(), + jacobian_type = Manopt.CoordinateVectorialType(basis), + retraction_method = retraction_method, + stopping_criterion = stopping, + η = η, + damping_term_min = damping_term_min, + β = β, + expect_zero_residual = expect_zero_residual, + linear_subsolver! = linear_subsolver, + debug = callbacks.debug_actions, + return_state = true, + ) + + return _manopt_finish_result( + _tk_get_solver_result(state), + state, + callbacks.progress, + diagnostics_recorder, + setup.solver_cost, + setup.solver_grad, + M, + normA2; + tol_T = T(tol), + maxiter, + solver = :lm, + tiny_grad_tol = setup.dual_grad_tol, + return_stats, + verbose, + normalized_objective = setup.uses_relative_objective, + solver_info_extra = ( + η = Float64(η), + damping_term_min = Float64(damping_term_min), + β = Float64(β), + expect_zero_residual = expect_zero_residual, + uses_vector_transport = !isnothing(vector_transport_method), + ), + ) +end + +function run_second_order_solver( + solver::LMSolver, + setup; + maxiter::Int, + tol::Real, + verbose::Bool, + return_stats::Bool, + vector_transport_method::Union{ManifoldsBase.AbstractVectorTransportMethod,Nothing} = nothing, + post_step_callback, + diagnostics_recorder, + iteration_callbacks, + grad_tol = nothing, + normalized_objective::Bool = true, +) + return solve_lm( + setup.model, + setup.model_cost, + setup.model_egrad, + setup.M, + setup.p0; + maxiter, + tol, + verbose, + return_stats, + normA2 = setup.normA2, + model_grad = setup.model_grad, + vector_transport_method, + post_step_callback, + diagnostics_recorder, + iteration_callbacks, + η = solver.η, + damping_term_min = solver.damping_term_min, + β = solver.β, + expect_zero_residual = solver.expect_zero_residual, + linear_subsolver = solver.linear_subsolver, + grad_tol, + normalized_objective, + ) +end diff --git a/src/solvers/manopt_helpers.jl b/src/solvers/manopt_helpers.jl index 67b8ac8..e6bc167 100644 --- a/src/solvers/manopt_helpers.jl +++ b/src/solvers/manopt_helpers.jl @@ -300,7 +300,7 @@ end function _relative_solver_functions(model_cost, model_grad, scale::Real) scale > 0 || return model_cost, model_grad, false scale == one(scale) && return model_cost, model_grad, false - inv_scale = inv(scale) + inv_scale = one(scale) / scale return ( (M, p) -> model_cost(M, p) * inv_scale, (M, p) -> _scale_solver_tangent(model_grad(M, p), inv_scale), diff --git a/src/solvers/solve_dispatch.jl b/src/solvers/solve_dispatch.jl index 9bdc801..70e1bf6 100644 --- a/src/solvers/solve_dispatch.jl +++ b/src/solvers/solve_dispatch.jl @@ -3,7 +3,7 @@ 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.", + "Unsupported solver specification $(typeof(solver)). Use a solver symbol such as :als, :rgd, :rgd_fixed, :rcg, :lbfgs, :lm, or :btd_tsd, or pass an AbstractSolver object.", ), ) end @@ -44,6 +44,16 @@ function _solver_object(::Val{:lbfgs}, ::Real; kwargs...) ) end +function _solver_object(::Val{:lm}, ::Real; kwargs...) + return LMSolver(; + η = get(kwargs, :η, 0.2), + damping_term_min = get(kwargs, :damping_term_min, 0.1), + β = get(kwargs, :β, 5.0), + expect_zero_residual = get(kwargs, :expect_zero_residual, false), + linear_subsolver = get(kwargs, :linear_subsolver, Manopt.default_lm_lin_solve!), + ) +end + function _solver_object(::Val{:btd_tsd}, stepsize::Real; kwargs...) return BTDTSDSolver(; stepsize, @@ -58,7 +68,7 @@ end function _solver_object(::Val{S}, ::Real; kwargs...) where {S} throw( ArgumentError( - "Unknown solver=$S. Use :als, :rgd, :rgd_fixed, :rcg, :lbfgs, or :btd_tsd.", + "Unknown solver=$S. Use :als, :rgd, :rgd_fixed, :rcg, :lbfgs, :lm, or :btd_tsd.", ), ) end diff --git a/test/basic_tests.jl b/test/basic_tests.jl index 6a818c6..fccb9c9 100644 --- a/test/basic_tests.jl +++ b/test/basic_tests.jl @@ -175,6 +175,83 @@ end ) end +@testset "LM residual/Jacobian smoke check matches gradient" begin + cases = ( + JoinModel((Manifolds.Segre((5, 4, 3)), Manifolds.Segre((5, 4, 3))), randn(5, 4, 3)), + JoinModel(abs.(randn(5, 4, 3)), 2; geometry = :softplus_metric, nonnegative = true), + ) + for model in cases + M = TensorKitchen.manifold(model) + p = TensorKitchen._solver_point( + M, + TensorKitchen.initial_point(model, :random; verbose = false), + ) + basis = ManifoldsBase.DefaultOrthonormalBasis() + r = TensorKitchen._lm_raw_residual_vector(model, p) + J = TensorKitchen._lm_raw_jacobian_matrix(model, M, p; basis) + g_coord = transpose(J) * r + g_from_J = ManifoldsBase.get_vector(M, p, g_coord, basis) + g_model = TensorKitchen.rgrad(model, p) + @test norm(M, p, g_from_J - g_model) ≤ 1e-7 * max(1.0, norm(M, p, g_model)) + end +end + +@testset "LM normalized and unnormalized objectives take the same step" begin + A = randn(6, 5, 4) + model = JoinModel((Manifolds.Segre((6, 5, 4)), Manifolds.Segre((6, 5, 4))), A) + p0 = TensorKitchen._solver_point( + TensorKitchen.manifold(model), + TensorKitchen.initial_point(model, :random; verbose = false), + ) + res_rel = solve( + LMSolver(), + model; + p0, + maxiter = 1, + tol = 0.0, + verbose = false, + return_stats = true, + normalized_objective = true, + ) + res_abs = solve( + LMSolver(), + model; + p0, + maxiter = 1, + tol = 0.0, + verbose = false, + return_stats = true, + normalized_objective = false, + ) + buf_rel = similar(model.backend.work_rec) + buf_abs = similar(model.backend.work_rec) + TensorKitchen._join_reconstruct!(buf_rel, model.backend, TensorKitchen.point(res_rel)) + TensorKitchen._join_reconstruct!(buf_abs, model.backend, TensorKitchen.point(res_abs)) + @test maximum(abs.(buf_rel .- buf_abs)) ≤ 1e-10 + @test isapprox(res_rel.rel_error, res_abs.rel_error; rtol = 1e-10, atol = 1e-10) +end + +@testset "cpd/approx accept LMSolver" begin + A = randn(5, 4, 3) + res_cpd_symbol = cpd(A, 2; solver = :lm, maxiter = 2, tol = 1e-6, verbose = false) + @test res_cpd_symbol.solver == :lm + + res_cpd_object = cpd( + A, + 2; + solver = LMSolver(damping_term_min = 1e-2), + maxiter = 2, + tol = 1e-6, + verbose = false, + ) + @test res_cpd_object.solver == :lm + + target = [1.2, -0.4, 0.8] + res_approx = + approx(Manifolds.Sphere(2), target; solver = :lm, maxiter = 2, verbose = false) + @test res_approx.solver == :lm +end + # ========================================================================= # cpd/cp_rank.jl (cost/egrad functions) # ========================================================================= From 6d47cff37ec22758b3a9ddefb24d596422dcd528 Mon Sep 17 00:00:00 2001 From: Se Eun Choi Date: Fri, 19 Jun 2026 08:41:56 +0200 Subject: [PATCH 2/7] LM solver accepts nested arrays for BTD, temporarily fix, eventually raise the issue to manopt --- src/solvers/lm.jl | 45 +++++++++++++++++---------------------------- test/basic_tests.jl | 42 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 59 insertions(+), 28 deletions(-) diff --git a/src/solvers/lm.jl b/src/solvers/lm.jl index 6568729..eb9b8a8 100644 --- a/src/solvers/lm.jl +++ b/src/solvers/lm.jl @@ -31,13 +31,9 @@ end solver_symbol(::LMSolver) = :lm -@inline function _lm_objective_scale( - ::Type{T}, - normA2, - normalized_objective::Bool, -) where {T} - return normalized_objective && !isnothing(normA2) && normA2 > 0 ? - one(T) / sqrt(T(normA2)) : one(T) +@inline function _lm_scaling_factor(::Type{T}, normA2, normalized_objective::Bool) where {T} + return normalized_objective && !isnothing(normA2) && normA2 > 0 ? inv(sqrt(T(normA2))) : + one(T) end function _ambient_tangent_vector!(out::AbstractVector, M, p, X) @@ -51,11 +47,6 @@ function _ambient_tangent_vector!(out::AbstractVector, M, p, X) return out end -function _ambient_tangent_vector!(out::AbstractVector, M::Manifolds.Segre, p, X) - copyto!(out, _segre_tangent_tensorvec(p, X)) - return out -end - function _join_tangent_ambient_vector!( out::AbstractVector, backend::Union{JoinBackend,BTDBackend}, @@ -92,8 +83,7 @@ function _lm_raw_jacobian_matrix( p; basis = ManifoldsBase.DefaultOrthonormalBasis(), ) - p_work = _solver_point(M, p) - T = _scalar_eltype(p_work) + T = _scalar_eltype(p) ambient_dim = length(tensor(model)) d = manifold_dimension(M) J = Matrix{T}(undef, ambient_dim, d) @@ -102,8 +92,8 @@ function _lm_raw_jacobian_matrix( @inbounds for j = 1:d fill!(coeff, zero(T)) coeff[j] = one(T) - Xj = ManifoldsBase.get_vector(M, p_work, coeff, basis) - _join_tangent_ambient_vector!(column, model.backend, p_work, Xj) + Xj = ManifoldsBase.get_vector(M, p, coeff, basis) + _join_tangent_ambient_vector!(column, model.backend, p, Xj) J[:, j] .= column end return J @@ -182,17 +172,16 @@ function _lm_raw_jacobian_matrix( p; basis = ManifoldsBase.DefaultOrthonormalBasis(), ) where {T<:AbstractFloat} - p_work = _solver_point(M, p) ambient_dim = length(model.A) d = manifold_dimension(M) J = Matrix{T}(undef, ambient_dim, d) coeff = zeros(T, d) - λp, Up = unpack_point_rank1(p_work, model.dims) + λp, Up = unpack_point_rank1(p, model.dims) column = Vector{T}(undef, ambient_dim) @inbounds for j = 1:d fill!(coeff, zero(T)) coeff[j] = one(T) - Xj = ManifoldsBase.get_vector(M, p_work, coeff, basis) + Xj = ManifoldsBase.get_vector(M, p, coeff, basis) if model.nonnegative λ̇p, U̇p = unpack_point_rank1(Xj, model.dims) kind = _rank1_uses_softplus_metric(model.M) ? Val(:softplus) : Val(:square) @@ -207,7 +196,7 @@ function _lm_raw_jacobian_matrix( U̇_vec = [Vector(@view U̇[m][:, 1]) for m in eachindex(U̇)] _cp_rank1_tangent_tensorvec!(column, λ[1], U_vec, λ̇[1], U̇_vec) else - copyto!(column, vec(ManifoldsBase.embed(M, p_work, Xj))) + copyto!(column, vec(ManifoldsBase.embed(M, p, Xj))) end J[:, j] .= column end @@ -224,18 +213,17 @@ function _lm_raw_jacobian_matrix( p; basis = ManifoldsBase.DefaultOrthonormalBasis(), ) where {T<:AbstractFloat} - p_work = _solver_point(M, p) ambient_dim = length(model.A) d = manifold_dimension(M) J = Matrix{T}(undef, ambient_dim, d) coeff = zeros(T, d) column = Vector{T}(undef, ambient_dim) if model.geometry == :native && !model.nonnegative - pparts = point_parts(p_work) + pparts = point_parts(p) @inbounds for j = 1:d fill!(coeff, zero(T)) coeff[j] = one(T) - Xj = ManifoldsBase.get_vector(M, p_work, coeff, basis) + Xj = ManifoldsBase.get_vector(M, p, coeff, basis) xparts = point_parts(Xj) fill!(column, zero(T)) for k = 1:model.r @@ -247,8 +235,8 @@ function _lm_raw_jacobian_matrix( end λp, Up = - model.nonnegative ? unpack_point_rankr(p_work, model.dims, model.r) : - unpack_rankr_canonical(p_work, model.dims, model.r) + model.nonnegative ? unpack_point_rankr(p, model.dims, model.r) : + unpack_rankr_canonical(p, model.dims, model.r) kind = model.nonnegative ? (model.geometry == :softplus_metric ? Val(:softplus) : Val(:square)) : @@ -256,7 +244,7 @@ function _lm_raw_jacobian_matrix( @inbounds for j = 1:d fill!(coeff, zero(T)) coeff[j] = one(T) - Xj = ManifoldsBase.get_vector(M, p_work, coeff, basis) + Xj = ManifoldsBase.get_vector(M, p, coeff, basis) λ̇p, U̇p = model.nonnegative ? unpack_point_rankr(Xj, model.dims, model.r) : unpack_rankr_canonical(Xj, model.dims, model.r) @@ -281,7 +269,7 @@ function _lm_residual_function( normA2, normalized_objective::Bool, ) where {T<:AbstractFloat} - scale = _lm_objective_scale(T, normA2, normalized_objective) + scale = _lm_scaling_factor(T, normA2, normalized_objective) return (M, p) -> scale .* _lm_raw_residual_vector(model, p) end @@ -292,7 +280,7 @@ function _lm_jacobian_function( normalized_objective::Bool; basis = ManifoldsBase.DefaultOrthonormalBasis(), ) where {T<:AbstractFloat} - scale = _lm_objective_scale(T, normA2, normalized_objective) + scale = _lm_scaling_factor(T, normA2, normalized_objective) return (M, p) -> scale .* _lm_raw_jacobian_matrix(model, M, p; basis) end @@ -376,6 +364,7 @@ function solve_lm( expect_zero_residual = expect_zero_residual, linear_subsolver! = linear_subsolver, debug = callbacks.debug_actions, + count = [:Cost, :Gradient], return_state = true, ) diff --git a/test/basic_tests.jl b/test/basic_tests.jl index fccb9c9..487de48 100644 --- a/test/basic_tests.jl +++ b/test/basic_tests.jl @@ -252,6 +252,48 @@ end @test res_approx.solver == :lm end +@testset "BTD accepts LMSolver on nested Tucker layouts" begin + A = randn(7, 6, 5) + ranks = (2, 2, 2) + manifolds = TensorKitchen._as_join_manifold_tuple(TuckerJoin(size(A), ranks, 2)) + backend = TensorKitchen._sum_backend_instance(TensorKitchen.BTDBackend, manifolds, A) + model = TensorKitchen.JoinModel{Float64,typeof(backend)}(backend) + p0 = TensorKitchen.initial_point(model, :random; verbose = false) + + @test p0 isa ArrayPartition + @test TensorKitchen.point_parts(p0)[1] isa Manifolds.TuckerPoint + + low = solve( + LMSolver(), + model; + p0, + maxiter = 2, + tol = 1e-6, + verbose = false, + return_stats = true, + ) + low_parts = TensorKitchen.point_parts(low.point) + @test low.solver == :lm + @test low.point isa ArrayPartition + @test length(low_parts) == 2 + @test low_parts[1] isa Manifolds.TuckerPoint + + res_btd = btd( + A, + 2, + ranks; + solver = :lm, + warm_rel_error_gate = nothing, + maxiter = 2, + tol = 1e-6, + verbose = false, + ) + @test res_btd isa BTDResult + @test res_btd.solver == :lm + @test length(res_btd.components) == 2 + @test !get(res_btd.solver_info, :btd_skipped_manifold_polish, false) +end + # ========================================================================= # cpd/cp_rank.jl (cost/egrad functions) # ========================================================================= From 51e1a2dcf5b71449649780b3c6c2e12d09df851d Mon Sep 17 00:00:00 2001 From: Se Eun Choi Date: Sat, 20 Jun 2026 08:54:48 +0200 Subject: [PATCH 3/7] refactor abstract joinModel for RLM --- src/join/join_backend.jl | 54 ++++++++++++++++++++++++++++++++++++---- src/solvers/lm.jl | 27 ++++++-------------- test/basic_tests.jl | 40 +++++++++++++++++++++++------ 3 files changed, 89 insertions(+), 32 deletions(-) diff --git a/src/join/join_backend.jl b/src/join/join_backend.jl index e8d88d7..d7abd0f 100644 --- a/src/join/join_backend.jl +++ b/src/join/join_backend.jl @@ -448,12 +448,19 @@ function rgrad(model::JoinModel{<:AbstractFloat,<:JoinBackend}, p) return wrap_like_point(p, vals) end -function _ambient_vector!(out::AbstractVector, M, p) +""" + _component_ambient_embedding!(out, M, p) + +Write the ambient embedding of one join component into `out`. +Component manifolds may specialize this hook when their native point structure +admits a more direct embedding than the generic `embed!` path. +""" +function _component_ambient_embedding!(out::AbstractVector, M, p) ManifoldsBase.embed!(M, out, p) return out end -function _ambient_vector!( +function _component_ambient_embedding!( out::AbstractVector, M::Manifolds.Tucker, p::Manifolds.TuckerPoint, @@ -464,7 +471,16 @@ function _ambient_vector!( return out end -function _ambient_vector!(out::AbstractVector, M::Manifolds.Tucker, p) +function _component_ambient_embedding!( + out::AbstractVector, + M::Manifolds.Segre, + p, +) + copyto!(out, _segre_component_tensorvec(p)) + return out +end + +function _component_ambient_embedding!(out::AbstractVector, M::Manifolds.Tucker, p) throw( ArgumentError( "Expected native TuckerPoint for Manifolds.Tucker, got $(typeof(p)).", @@ -472,6 +488,34 @@ function _ambient_vector!(out::AbstractVector, M::Manifolds.Tucker, p) ) end +""" + _component_ambient_pushforward!(out, M, p, X) + +Write the ambient pushforward `DΦ(p)[X]` of one join component into `out`. +This is the component-level differential used by LM Jacobian assembly on +generic `JoinModel`s. +""" +function _component_ambient_pushforward!(out::AbstractVector, M, p, X) + emb = ManifoldsBase.embed(M, p, X) + length(emb) == length(out) || throw( + DimensionMismatch( + "Tangent embedding length $(length(emb)) does not match output length $(length(out)).", + ), + ) + copyto!(out, vec(emb)) + return out +end + +function _component_ambient_pushforward!( + out::AbstractVector, + M::Manifolds.Segre, + p, + X, +) + copyto!(out, _segre_tangent_tensorvec(p, X)) + return out +end + function _subtract_ambient_tensor!( residual::AbstractArray{T,N}, M, @@ -483,7 +527,7 @@ function _subtract_ambient_tensor!( "_subtract_ambient_tensor!: work length $(length(work_vec)) != residual length $(length(residual))", ), ) - _ambient_vector!(work_vec, M, p) + _component_ambient_embedding!(work_vec, M, p) residual_vec = vec(residual) @inbounds for i in eachindex(residual_vec, work_vec) residual_vec[i] -= work_vec[i] @@ -531,7 +575,7 @@ function _join_reconstruct!(out::AbstractArray, backend::Union{JoinBackend,BTDBa @inbounds for k = 1:r # Reconstruct each component into its preallocated workspace. - _ambient_vector!(bufs[k], manifolds[k], parts[k]) + _component_ambient_embedding!(bufs[k], manifolds[k], parts[k]) # Accumulate into the output tensor without allocating a Khatri-Rao-sized object. out .+= bufs[k] diff --git a/src/solvers/lm.jl b/src/solvers/lm.jl index eb9b8a8..5625a4f 100644 --- a/src/solvers/lm.jl +++ b/src/solvers/lm.jl @@ -32,19 +32,8 @@ end solver_symbol(::LMSolver) = :lm @inline function _lm_scaling_factor(::Type{T}, normA2, normalized_objective::Bool) where {T} - return normalized_objective && !isnothing(normA2) && normA2 > 0 ? inv(sqrt(T(normA2))) : - one(T) -end - -function _ambient_tangent_vector!(out::AbstractVector, M, p, X) - emb = ManifoldsBase.embed(M, p, X) - length(emb) == length(out) || throw( - DimensionMismatch( - "Tangent embedding length $(length(emb)) does not match output length $(length(out)).", - ), - ) - copyto!(out, vec(emb)) - return out + return normalized_objective && !isnothing(normA2) && normA2 > 0 ? + one(T) / sqrt(T(normA2)) : one(T) end function _join_tangent_ambient_vector!( @@ -59,12 +48,7 @@ function _join_tangent_ambient_vector!( _check_parts_len(xparts, backend.r, "_join_tangent_ambient_vector!") fill!(out, zero(eltype(out))) @inbounds for k = 1:backend.r - _ambient_tangent_vector!( - backend.component_bufs[k], - backend.manifolds[k], - parts[k], - xparts[k], - ) + _component_ambient_pushforward!(backend.component_bufs[k], backend.manifolds[k], parts[k], xparts[k]) out .+= backend.component_bufs[k] end return out @@ -324,6 +308,8 @@ function solve_lm( basis = ManifoldsBase.DefaultOrthonormalBasis() residual = _lm_residual_function(model, T, normA2, setup.uses_relative_objective) jacobian = _lm_jacobian_function(model, T, normA2, setup.uses_relative_objective; basis) + initial_residual_values = residual(M, p0_local) + initial_jacobian_f = jacobian(M, p0_local) retraction_method = _solver_retraction_method(M, p0_local) stopping = StopWhenAny( StopAfterIteration(maxiter), @@ -358,13 +344,14 @@ function solve_lm( jacobian_type = Manopt.CoordinateVectorialType(basis), retraction_method = retraction_method, stopping_criterion = stopping, + initial_residual_values = initial_residual_values, + initial_jacobian_f = initial_jacobian_f, η = η, damping_term_min = damping_term_min, β = β, expect_zero_residual = expect_zero_residual, linear_subsolver! = linear_subsolver, debug = callbacks.debug_actions, - count = [:Cost, :Gradient], return_state = true, ) diff --git a/test/basic_tests.jl b/test/basic_tests.jl index 487de48..496705b 100644 --- a/test/basic_tests.jl +++ b/test/basic_tests.jl @@ -177,7 +177,7 @@ end @testset "LM residual/Jacobian smoke check matches gradient" begin cases = ( - JoinModel((Manifolds.Segre((5, 4, 3)), Manifolds.Segre((5, 4, 3))), randn(5, 4, 3)), + JoinModel(randn(5, 4, 3), 2; geometry = :canonical), JoinModel(abs.(randn(5, 4, 3)), 2; geometry = :softplus_metric, nonnegative = true), ) for model in cases @@ -196,9 +196,36 @@ end end end +@testset "LM generic join Jacobian uses component pushforwards" begin + A = randn(5, 4, 3) + model = JoinModel((Manifolds.Segre((5, 4, 3)), Manifolds.Segre((5, 4, 3))), A) + M = TensorKitchen.manifold(model) + p = TensorKitchen._solver_point(M, TensorKitchen.initial_point(model, :random; verbose = false)) + basis = ManifoldsBase.DefaultOrthonormalBasis() + J = TensorKitchen._lm_raw_jacobian_matrix(model, M, p; basis) + @test all(isfinite, J) + + retraction_method = TensorKitchen._solver_retraction_method(M, p) + ϵ = 1e-6 + buf_plus = similar(model.backend.work_rec) + buf_minus = similar(model.backend.work_rec) + d = manifold_dimension(M) + for j = 1:min(d, 3) + coeff = zeros(Float64, d) + coeff[j] = 1.0 + Xj = ManifoldsBase.get_vector(M, p, coeff, basis) + p_plus = ManifoldsBase.retract(M, p, ϵ * Xj, retraction_method) + p_minus = ManifoldsBase.retract(M, p, -ϵ * Xj, retraction_method) + TensorKitchen._join_reconstruct!(buf_plus, model.backend, p_plus) + TensorKitchen._join_reconstruct!(buf_minus, model.backend, p_minus) + fd = (buf_plus .- buf_minus) ./ (2 * ϵ) + @test maximum(abs.(fd .- J[:, j])) ≤ 1e-7 + end +end + @testset "LM normalized and unnormalized objectives take the same step" begin A = randn(6, 5, 4) - model = JoinModel((Manifolds.Segre((6, 5, 4)), Manifolds.Segre((6, 5, 4))), A) + model = JoinModel(A, 2; geometry = :canonical) p0 = TensorKitchen._solver_point( TensorKitchen.manifold(model), TensorKitchen.initial_point(model, :random; verbose = false), @@ -223,11 +250,10 @@ end return_stats = true, normalized_objective = false, ) - buf_rel = similar(model.backend.work_rec) - buf_abs = similar(model.backend.work_rec) - TensorKitchen._join_reconstruct!(buf_rel, model.backend, TensorKitchen.point(res_rel)) - TensorKitchen._join_reconstruct!(buf_abs, model.backend, TensorKitchen.point(res_abs)) - @test maximum(abs.(buf_rel .- buf_abs)) ≤ 1e-10 + rawmodel = TensorKitchen.cpd_model(model) + X_rel = TensorKitchen.embed_point(rawmodel, TensorKitchen.point(res_rel)) + X_abs = TensorKitchen.embed_point(rawmodel, TensorKitchen.point(res_abs)) + @test maximum(abs.(X_rel .- X_abs)) ≤ 1e-10 @test isapprox(res_rel.rel_error, res_abs.rel_error; rtol = 1e-10, atol = 1e-10) end From be39355aa6716db88ca30fb5d9c0c28249f99d00 Mon Sep 17 00:00:00 2001 From: Se Eun Choi Date: Sat, 20 Jun 2026 08:56:30 +0200 Subject: [PATCH 4/7] apply formatter --- src/join/join_backend.jl | 13 ++----------- src/solvers/lm.jl | 7 ++++++- test/basic_tests.jl | 5 ++++- 3 files changed, 12 insertions(+), 13 deletions(-) diff --git a/src/join/join_backend.jl b/src/join/join_backend.jl index d7abd0f..89fc48e 100644 --- a/src/join/join_backend.jl +++ b/src/join/join_backend.jl @@ -471,11 +471,7 @@ function _component_ambient_embedding!( return out end -function _component_ambient_embedding!( - out::AbstractVector, - M::Manifolds.Segre, - p, -) +function _component_ambient_embedding!(out::AbstractVector, M::Manifolds.Segre, p) copyto!(out, _segre_component_tensorvec(p)) return out end @@ -506,12 +502,7 @@ function _component_ambient_pushforward!(out::AbstractVector, M, p, X) return out end -function _component_ambient_pushforward!( - out::AbstractVector, - M::Manifolds.Segre, - p, - X, -) +function _component_ambient_pushforward!(out::AbstractVector, M::Manifolds.Segre, p, X) copyto!(out, _segre_tangent_tensorvec(p, X)) return out end diff --git a/src/solvers/lm.jl b/src/solvers/lm.jl index 5625a4f..77391f4 100644 --- a/src/solvers/lm.jl +++ b/src/solvers/lm.jl @@ -48,7 +48,12 @@ function _join_tangent_ambient_vector!( _check_parts_len(xparts, backend.r, "_join_tangent_ambient_vector!") fill!(out, zero(eltype(out))) @inbounds for k = 1:backend.r - _component_ambient_pushforward!(backend.component_bufs[k], backend.manifolds[k], parts[k], xparts[k]) + _component_ambient_pushforward!( + backend.component_bufs[k], + backend.manifolds[k], + parts[k], + xparts[k], + ) out .+= backend.component_bufs[k] end return out diff --git a/test/basic_tests.jl b/test/basic_tests.jl index 496705b..2d88071 100644 --- a/test/basic_tests.jl +++ b/test/basic_tests.jl @@ -200,7 +200,10 @@ end A = randn(5, 4, 3) model = JoinModel((Manifolds.Segre((5, 4, 3)), Manifolds.Segre((5, 4, 3))), A) M = TensorKitchen.manifold(model) - p = TensorKitchen._solver_point(M, TensorKitchen.initial_point(model, :random; verbose = false)) + p = TensorKitchen._solver_point( + M, + TensorKitchen.initial_point(model, :random; verbose = false), + ) basis = ManifoldsBase.DefaultOrthonormalBasis() J = TensorKitchen._lm_raw_jacobian_matrix(model, M, p; basis) @test all(isfinite, J) From aebd1a3eb8f91bcb5544302c81e77e1e2f9cd441 Mon Sep 17 00:00:00 2001 From: Se Eun Choi Date: Sat, 20 Jun 2026 09:14:03 +0200 Subject: [PATCH 5/7] Keep the RLM path for canonical CP (ALS) --- src/core/model.jl | 3 +- src/core/unpack_points.jl | 44 +++++++++++++++++++- src/join/join_backend.jl | 33 +++++++++++++++ test/basic_tests.jl | 85 +++++++++++++++++++++++++++++++++++++++ 4 files changed, 163 insertions(+), 2 deletions(-) diff --git a/src/core/model.jl b/src/core/model.jl index bf2be86..f598072 100644 --- a/src/core/model.jl +++ b/src/core/model.jl @@ -1,5 +1,6 @@ # core/model.jl — Top-level decomposition model interface -export AbstractDecompositionModel, rgrad, supports_rgrad, tensor, cost, post_step! +export AbstractDecompositionModel, + manifold, initial_point, egrad, rgrad, supports_rgrad, tensor, cost, post_step! """ AbstractDecompositionModel{T} diff --git a/src/core/unpack_points.jl b/src/core/unpack_points.jl index 8828e86..e8d3d02 100644 --- a/src/core/unpack_points.jl +++ b/src/core/unpack_points.jl @@ -1,6 +1,11 @@ # core/unpack_points.jl — Point unpacking and legacy/vector interop export pack_point_rank1, - unpack_point_rank1, pack_point_rankr, unpack_point_rankr, unpack_point_rankr_components + unpack_point_rank1, + pack_point_rankr, + unpack_point_rankr, + unpack_point_rankr_components, + canonical_to_joinpoint, + joinpoint_to_canonical function unpack_rankr_native(p, dims::NTuple{N,Int}, r::Int) where {N} parts = normalize_rankr_native_point(p, dims, r) @@ -72,6 +77,43 @@ function unpack_rankr_join(p, dims::NTuple{N,Int}, r::Int) where {N} return λ, U end +""" + canonical_to_joinpoint(λ, U) + canonical_to_joinpoint(p_canonical, dims, r) + +Convert a CPD point from canonical factor-matrix storage to the native +rank-`r` Segre join point layout used by generic `JoinModel((Segre, ...), A)`. + +The conversion preserves the represented tensor but may renormalize component +gauges the same way `pack_rankr_native` does. +""" +function canonical_to_joinpoint( + λ::AbstractVector{T}, + U::Vector{<:AbstractMatrix{T}}, +) where {T<:AbstractFloat} + r = length(λ) + return pack_rankr_native(λ, U, r) +end + +function canonical_to_joinpoint(p, dims::NTuple{N,Int}, r::Int) where {N} + λ, U = unpack_rankr_canonical(p, dims, r) + return canonical_to_joinpoint(λ, U) +end + +""" + joinpoint_to_canonical(p_join, dims, r) + +Convert a native Segre join point layout back to the canonical CPD point +layout `(λ, (u₁¹, …, uᵣ¹), …, (u₁ᴺ, …, uᵣᴺ))`. + +The conversion preserves the represented tensor but may renormalize component +gauges the same way `pack_rankr_canonical` does. +""" +function joinpoint_to_canonical(p, dims::NTuple{N,Int}, r::Int) where {N} + λ, U = unpack_rankr_native(p, dims, r) + return pack_rankr_canonical(λ, U, r) +end + function pack_point_rank1(λ::T, U::Vector{Vector{T}}) where {T<:AbstractFloat} parts = Vector{Vector{T}}(undef, length(U) + 1) parts[1] = T[λ] diff --git a/src/join/join_backend.jl b/src/join/join_backend.jl index 89fc48e..cf8417e 100644 --- a/src/join/join_backend.jl +++ b/src/join/join_backend.jl @@ -22,6 +22,16 @@ function _as_join_manifold_tuple(manifolds::AbstractVector) end _as_join_manifold_tuple(M::ProductManifold) = Tuple(M.manifolds) + +function _uniform_segre_dims(manifolds::Tuple) + isempty(manifolds) && return nothing + first_manifold = first(manifolds) + first_manifold isa Manifolds.Segre || return nothing + dims = factor_dims(first_manifold) + all(M -> M isa Manifolds.Segre && factor_dims(M) == dims, manifolds) || return nothing + return dims +end + @inline function _check_parts_len(parts, expected::Int, where_fn::AbstractString) length(parts) == expected || throw( DimensionMismatch( @@ -344,6 +354,29 @@ function initial_point( return ArrayPartition(parts...) end +function initial_point( + model::JoinModel{<:AbstractFloat,<:JoinBackend}, + init::ALSWarmStartInit; + verbose::Bool = false, + kwargs..., +) + backend = model.backend + dims = _uniform_segre_dims(backend.manifolds) + isnothing(dims) && throw( + ArgumentError( + "ALSWarmStartInit for a generic JoinModel requires all component manifolds to be Manifolds.Segre with identical factor_dims.", + ), + ) + dims == backend.target_shape || throw( + DimensionMismatch( + "Uniform Segre factor_dims $dims must match target size $(backend.target_shape) for ALS warm start.", + ), + ) + warm_model = JoinModel(backend.target, backend.r; geometry = :canonical) + p_canonical = initial_point(warm_model, init; verbose, kwargs...) + return canonical_to_joinpoint(p_canonical, backend.target_shape, backend.r) +end + # Gradient path: always recomputes the ambient reconstruction and marks the # WORO cache fresh so that the immediately following cost evaluation can reuse it. function _join_residual_grad!(backend::JoinBackend, p) diff --git a/test/basic_tests.jl b/test/basic_tests.jl index 2d88071..45d55a4 100644 --- a/test/basic_tests.jl +++ b/test/basic_tests.jl @@ -175,6 +175,59 @@ end ) end +@testset "JoinModel: generic (Segre, Segre, ...) backend" begin + dims = (5, 4, 3) + rng = MersenneTwister(4242) + A = randn(rng, dims...) + segres = (Manifolds.Segre(dims), Manifolds.Segre(dims), Manifolds.Segre(dims)) + + model = JoinModel(segres, A) + @test model isa JoinModel + @test model.backend isa TensorKitchen.JoinBackend + @test model.backend.r == 3 + @test tensor(model) == A + @test model.backend.manifolds == segres + + M = TensorKitchen.manifold(model) + @test M isa ProductManifold + @test length(M.manifolds) == 3 + @test all(Mk -> Mk isa Manifolds.Segre && factor_dims(Mk) == dims, M.manifolds) + + model_repeat = JoinModel(Manifolds.Segre(dims), 3, A) + @test model_repeat.backend.r == 3 + @test model_repeat.backend.manifolds == segres + + p = TensorKitchen.initial_point(model, :random; verbose = false) + @test length(TensorKitchen.point_parts(p)) == 3 + + f = cost(model, p) + g = TensorKitchen.egrad(model, p) + rg = rgrad(model, p) + @test isfinite(f) && f >= 0 + @test isfinite(norm(M, p, g)) + @test isfinite(norm(M, p, rg)) + + rec = similar(model.backend.work_rec) + TensorKitchen._join_reconstruct!(rec, model.backend, p) + @test f ≈ 0.5 * sum(abs2, rec .- vec(A)) + + comps = TensorKitchen.extract_components(model, p) + @test length(comps) == 3 + @test all(c -> c.manifold isa Manifolds.Segre, comps) + @test all(c -> size(c.tensor) == dims, comps) + + out = solve( + RGDSolver(1.0), + model; + p0 = p, + maxiter = 2, + tol = 1e-6, + verbose = false, + return_stats = true, + ) + @test isfinite(out.cost) && isfinite(out.rel_error) +end + @testset "LM residual/Jacobian smoke check matches gradient" begin cases = ( JoinModel(randn(5, 4, 3), 2; geometry = :canonical), @@ -672,6 +725,28 @@ end p_warm_sym = TensorKitchen.initial_point(model, :alswarm) @test TensorKitchen.cost(model, p_warm) <= TensorKitchen.cost(model, p_base) + 1e-10 @test isfinite(TensorKitchen.cost(model, p_warm_sym)) + generic_join = JoinModel((Manifolds.Segre(dims), Manifolds.Segre(dims)), A) + p_warm_canonical_match = TensorKitchen.initial_point( + model, + ALSWarmStartInit(2; base_init = TuckerInit()); + verbose = false, + ) + p_join_warm = TensorKitchen.initial_point( + generic_join, + ALSWarmStartInit(2; base_init = TuckerInit()); + verbose = false, + ) + p_join_from_canonical = + TensorKitchen.canonical_to_joinpoint(p_warm_canonical_match, dims, r) + A_join_warm = reconstruct_cpd_rankr( + components_from_factors(TensorKitchen.unpack_rankr_native(p_join_warm, dims, r)...), + ) + A_join_from_canonical = reconstruct_cpd_rankr( + components_from_factors( + TensorKitchen.unpack_rankr_native(p_join_from_canonical, dims, r)..., + ), + ) + @test A_join_warm ≈ A_join_from_canonical res_p0 = cpd(A, r; solver = :rgd, p0 = p0, maxiter = 5, tol = 1e-6, verbose = false) @test res_p0 isa CPDResult @@ -2618,6 +2693,16 @@ end A_native = reconstruct_cpd_rankr(components_from_factors(λn, Un)) @test all(λn .>= 0) @test A_native ≈ A_in + p_canonical = TensorKitchen.pack_rankr_canonical(λ, U, r) + p_join = TensorKitchen.canonical_to_joinpoint(p_canonical, dims, r) + p_canonical_roundtrip = TensorKitchen.joinpoint_to_canonical(p_join, dims, r) + λ_join, U_join = TensorKitchen.unpack_rankr_native(p_join, dims, r) + λ_canon_rt, U_canon_rt = + TensorKitchen.unpack_rankr_canonical(p_canonical_roundtrip, dims, r) + A_join = reconstruct_cpd_rankr(components_from_factors(λ_join, U_join)) + A_canon_rt = reconstruct_cpd_rankr(components_from_factors(λ_canon_rt, U_canon_rt)) + @test A_join ≈ A_in + @test A_canon_rt ≈ A_in A = randn(8, 6, 5) λ0, U0 = cp_init_tucker(A, 3) From 55ebe67ec5005c05c7d75106567383271ecdc5f5 Mon Sep 17 00:00:00 2001 From: Se Eun Choi Date: Sat, 20 Jun 2026 11:40:24 +0200 Subject: [PATCH 6/7] add generic JoinModel fitness checks for RLM in basic_tests --- test/basic_tests.jl | 33 +++++++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/test/basic_tests.jl b/test/basic_tests.jl index 45d55a4..0f5ab1f 100644 --- a/test/basic_tests.jl +++ b/test/basic_tests.jl @@ -328,6 +328,20 @@ end ) @test res_cpd_object.solver == :lm + res_cpd_alswarm_lm = cpd( + A, + 2; + solver = :lm, + init = :alswarm, + warm_steps = 2, + warm_init = :tucker, + maxiter = 2, + tol = 1e-6, + verbose = false, + ) + @test res_cpd_alswarm_lm.solver == :lm + @test isfinite(res_cpd_alswarm_lm.rel_error) + target = [1.2, -0.4, 0.8] res_approx = approx(Manifolds.Sphere(2), target; solver = :lm, maxiter = 2, verbose = false) @@ -374,6 +388,25 @@ end @test res_btd.solver == :lm @test length(res_btd.components) == 2 @test !get(res_btd.solver_info, :btd_skipped_manifold_polish, false) + + res_btd_alswarm_lm = btd( + A, + 2, + ranks; + solver = :lm, + init = :alswarm, + warm_init = BTDHOSVDMultistartInit(2; screening_steps = 0, block_maxiter = 1), + warm_steps = 1, + warm_block_maxiter = 1, + warm_rel_error_gate = nothing, + maxiter = 2, + tol = 1e-6, + verbose = false, + ) + @test res_btd_alswarm_lm isa BTDResult + @test res_btd_alswarm_lm.solver == :lm + @test hasproperty(res_btd_alswarm_lm.solver_info, :btd_als_warm_start_iters) + @test res_btd_alswarm_lm.solver_info.btd_als_warm_start_requested_solver == :lm end # ========================================================================= From caa638a579b94a4d5d2daca7d15f5d1a07cd813e Mon Sep 17 00:00:00 2001 From: Se Eun Choi Date: Sat, 20 Jun 2026 12:50:11 +0200 Subject: [PATCH 7/7] refactor JoinComponent Abstraction --- src/join/join_backend.jl | 220 +++++++++++++++++++++++++++++++-------- src/join/join_model.jl | 18 +++- src/solvers/lm.jl | 3 +- test/basic_tests.jl | 10 +- 4 files changed, 200 insertions(+), 51 deletions(-) diff --git a/src/join/join_backend.jl b/src/join/join_backend.jl index cf8417e..5e828b4 100644 --- a/src/join/join_backend.jl +++ b/src/join/join_backend.jl @@ -2,6 +2,16 @@ _is_manifold_like(::AbstractManifold) = true _is_manifold_like(_) = false +_is_join_component_like(::JoinComponent) = true +_is_join_component_like(x) = _is_manifold_like(x) + +_wrap_join_component(component::JoinComponent) = component +_wrap_join_component(manifold::AbstractManifold) = JoinComponent(manifold) + +_component_manifold(component::JoinComponent) = component.manifold +_component_manifold(manifold::AbstractManifold) = manifold +_backend_components(backend::JoinBackend) = backend.components +_backend_components(backend::BTDBackend) = backend.manifolds function _as_join_manifold_tuple(manifolds::Tuple) all(_is_manifold_like, manifolds) || throw( @@ -23,12 +33,35 @@ end _as_join_manifold_tuple(M::ProductManifold) = Tuple(M.manifolds) -function _uniform_segre_dims(manifolds::Tuple) - isempty(manifolds) && return nothing - first_manifold = first(manifolds) +function _as_join_component_tuple(components::Tuple) + all(_is_join_component_like, components) || throw( + ArgumentError( + "All join components must be AbstractManifold or JoinComponent. Got types: $(map(typeof, components)).", + ), + ) + return ntuple(k -> _wrap_join_component(components[k]), length(components)) +end + +function _as_join_component_tuple(components::AbstractVector) + all(_is_join_component_like, components) || throw( + ArgumentError( + "All join components must be AbstractManifold or JoinComponent. Got types: $(map(typeof, components)).", + ), + ) + return Tuple(_wrap_join_component(c) for c in components) +end + +_as_join_component_tuple(M::ProductManifold) = _as_join_component_tuple(Tuple(M.manifolds)) + +function _uniform_segre_dims(components::Tuple) + isempty(components) && return nothing + first_manifold = _component_manifold(first(components)) first_manifold isa Manifolds.Segre || return nothing dims = factor_dims(first_manifold) - all(M -> M isa Manifolds.Segre && factor_dims(M) == dims, manifolds) || return nothing + all(c -> begin + M = _component_manifold(c) + M isa Manifolds.Segre && factor_dims(M) == dims + end, components) || return nothing return dims end @@ -48,6 +81,9 @@ _manifold_init(M::Manifolds.Sphere, target, init::Symbol) = _sphere_init(M, targ _manifold_init(M::Manifolds.Segre, target, init::Symbol) = _segre_init(M, target, init) _manifold_init(M::Manifolds.Tucker, target, init::Symbol) = _tucker_init(M, target, init) +_component_init(component, target, init) = + _manifold_init(_component_manifold(component), target, init) + function _manifold_init(M, target, init_sym::Symbol) init_sym == :random && return rand(M) throw( @@ -58,15 +94,16 @@ function _manifold_init(M, target, init_sym::Symbol) end """ - _manifold_egrad(M, p, residual) + _component_egrad(component, p, residual) Compute the component Euclidean gradient induced by a join residual. Tucker components use their native tensor gradient; other components copy the residual. """ -_manifold_egrad(M, p, residual) = copy(residual) -_manifold_egrad(M::Manifolds.Tucker, p, residual) = _tucker_egrad(M, p, residual) +_component_egrad(::DefaultJoinEmbedding, M, p, residual) = copy(residual) +_component_egrad(::DefaultJoinEmbedding, M::Manifolds.Tucker, p, residual) = + _tucker_egrad(M, p, residual) -function _manifold_egrad(M::Manifolds.Segre, p, residual) +function _component_egrad(::DefaultJoinEmbedding, M::Manifolds.Segre, p, residual) dims = factor_dims(M) R = reshape(residual, dims) parts = point_parts(p) @@ -81,6 +118,12 @@ function _manifold_egrad(M::Manifolds.Segre, p, residual) return pack_tangent_rank1_segre(grad_λ, grad_U) end +_component_egrad(component::JoinComponent, p, residual) = + _component_egrad(component.embedding, component.manifold, p, residual) +_component_egrad(M, p, residual) = _component_egrad(DefaultJoinEmbedding(), M, p, residual) + +_manifold_egrad(M, p, residual) = _component_egrad(M, p, residual) + """ _ambient_vector(M, p, target_len) returns AbstractVector @@ -146,6 +189,7 @@ end ambient_length(M::Manifolds.Segre) = prod(factor_dims(M)) ambient_length(M::Manifolds.Tucker) = prod(factor_dims(M)) +ambient_length(component::JoinComponent) = ambient_length(component.manifold) """ _join_vector_workspace_like(target, n) returns AbstractVector @@ -167,15 +211,15 @@ end Ensure every join component embeds into the same flattened ambient space as the target tensor. """ -function _validate_join_ambient_compatibility(manifolds::Tuple, target::AbstractArray) +function _validate_join_ambient_compatibility(components::Tuple, target::AbstractArray) target_len = length(target) failures = String[] - @inbounds for k in eachindex(manifolds) - mk = ambient_length(manifolds[k]) + @inbounds for k in eachindex(components) + mk = ambient_length(components[k]) if mk != target_len push!( failures, - "[$k] $(typeof(manifolds[k])) has ambient length $mk but target has length $target_len", + "[$k] $(typeof(components[k])) has ambient length $mk but target has length $target_len", ) end end @@ -203,14 +247,14 @@ function _sum_backend_instance( end function _sum_backend_parts( - manifolds, + components, target::AbstractArray{T,N}; init_point = nothing, ) where {T<:AbstractFloat,N} - r = length(manifolds) + r = length(components) # Keep the original target representation instead of eagerly materializing Array. tgt = target - _validate_join_ambient_compatibility(manifolds, tgt) + _validate_join_ambient_compatibility(components, tgt) tflat = vec(tgt) tgt_len = length(tgt) @@ -218,8 +262,10 @@ function _sum_backend_parts( component_bufs = [_join_vector_workspace_like(tgt, tgt_len) for _ = 1:r] work_rec = _join_vector_workspace_like(tgt, tgt_len) work_residual = _join_vector_workspace_like(tgt, tgt_len) + manifolds = ntuple(k -> _component_manifold(components[k]), r) return (; + components, manifolds, r, target = tgt, @@ -236,13 +282,13 @@ end function _sum_backend_instance( ::Type{JoinBackend}, - manifolds, + components, target::AbstractArray{T,N}; init_point = nothing, ) where {T<:AbstractFloat,N} - parts = _sum_backend_parts(manifolds, target; init_point) + parts = _sum_backend_parts(components, target; init_point) return JoinBackend( - parts.manifolds, + parts.components, parts.r, parts.target, parts.target_size, @@ -258,11 +304,11 @@ end function _sum_backend_instance( ::Type{BTDBackend}, - manifolds, + components, target::AbstractArray{T,N}; init_point = nothing, ) where {T<:AbstractFloat,N} - parts = _sum_backend_parts(manifolds, target; init_point) + parts = _sum_backend_parts(components, target; init_point) return BTDBackend( parts.manifolds, parts.r, @@ -286,13 +332,14 @@ Construct a generic sum-of-manifolds approximation model from explicit component manifolds and a target tensor. """ function JoinModel( - manifolds::Tuple{Vararg{AbstractManifold}}, + components::Tuple, target::AbstractArray{T,N}; init_point = nothing, ) where {T<:AbstractFloat,N} - r = length(manifolds) + components_tuple = _as_join_component_tuple(components) + r = length(components_tuple) r >= 1 || throw(ArgumentError("JoinModel(manifolds, target) needs r >= 1, got r=$r")) - b = _sum_backend_instance(JoinBackend, manifolds, target; init_point) + b = _sum_backend_instance(JoinBackend, components_tuple, target; init_point) return JoinModel{T,typeof(b)}(b) end @@ -302,11 +349,11 @@ end Construct a join model by repeating a base manifold `r` times. """ function JoinModel( - manifolds::AbstractVector, + components::AbstractVector, target::AbstractArray{T,N}; init_point = nothing, ) where {T<:AbstractFloat,N} - return JoinModel(_as_join_manifold_tuple(manifolds), target; init_point) + return JoinModel(_as_join_component_tuple(components), target; init_point) end function JoinModel( @@ -314,7 +361,7 @@ function JoinModel( target::AbstractArray{T,N}; init_point = nothing, ) where {T<:AbstractFloat,N} - return JoinModel(_as_join_manifold_tuple(M), target; init_point) + return JoinModel(_as_join_component_tuple(M), target; init_point) end function JoinModel( @@ -326,6 +373,15 @@ function JoinModel( return JoinModel(ntuple(_ -> base, r), target; init_point) end +function JoinModel( + base::JoinComponent, + r::Int, + target::AbstractArray{T,N}; + init_point = nothing, +) where {T<:AbstractFloat,N} + return JoinModel(ntuple(_ -> base, r), target; init_point) +end + function JoinModel( base::AbstractManifold, target::AbstractArray{T,N}; @@ -334,6 +390,14 @@ function JoinModel( return JoinModel((base,), target; init_point) end +function JoinModel( + base::JoinComponent, + target::AbstractArray{T,N}; + init_point = nothing, +) where {T<:AbstractFloat,N} + return JoinModel((base,), target; init_point) +end + manifold(model::JoinModel{<:AbstractFloat,<:Union{JoinBackend,BTDBackend}}) = model.backend.M_product tensor(model::JoinModel{<:AbstractFloat,<:Union{JoinBackend,BTDBackend}}) = @@ -350,7 +414,7 @@ function initial_point( return backend.init_point(M, init) end parts = - ntuple(k -> _manifold_init(backend.manifolds[k], backend.target, init), backend.r) + ntuple(k -> _component_init(backend.components[k], backend.target, init), backend.r) return ArrayPartition(parts...) end @@ -361,7 +425,7 @@ function initial_point( kwargs..., ) backend = model.backend - dims = _uniform_segre_dims(backend.manifolds) + dims = _uniform_segre_dims(backend.components) isnothing(dims) && throw( ArgumentError( "ALSWarmStartInit for a generic JoinModel requires all component manifolds to be Manifolds.Segre with identical factor_dims.", @@ -407,19 +471,21 @@ function egrad(model::JoinModel{<:AbstractFloat,<:JoinBackend}, p) backend = model.backend residual = _join_residual_grad!(backend, p) parts = point_parts(p) - vals = ntuple(k -> _manifold_egrad(backend.manifolds[k], parts[k], residual), backend.r) + vals = + ntuple(k -> _component_egrad(backend.components[k], parts[k], residual), backend.r) return wrap_like_point(p, vals) end -function _join_basis_project(manifolds::Tuple, p, residual) +function _join_basis_project(components::Tuple, p, residual) parts = point_parts(p) - _check_parts_len(parts, length(manifolds), "_join_basis_project") + _check_parts_len(parts, length(components), "_join_basis_project") vals = ntuple(k -> begin - Mk = manifolds[k] + ck = components[k] + Mk = _component_manifold(ck) pk = parts[k] - eg = _manifold_egrad(Mk, pk, residual) + eg = _component_egrad(ck, pk, residual) egrad_to_rgrad(Mk, pk, eg) - end, length(manifolds)) + end, length(components)) return wrap_like_point(p, vals) end @@ -437,7 +503,7 @@ model_exact_join_basis_function(model::JoinModel{<:AbstractFloat,<:JoinBackend}) (M, p) -> begin backend = model.backend residual = _join_residual!(backend, p) - _join_basis_project(backend.manifolds, p, residual) + _join_basis_project(backend.components, p, residual) end function extract_components( @@ -445,6 +511,7 @@ function extract_components( p, ) backend = model.backend + components = _backend_components(backend) parts = point_parts(p) _check_parts_len(parts, backend.r, "extract_components") T = eltype(backend.target) @@ -453,14 +520,14 @@ function extract_components( manifold_type = Union{} @inbounds for k = 1:backend.r point_type = typejoin(point_type, typeof(parts[k])) - manifold_type = typejoin(manifold_type, typeof(backend.manifolds[k])) + manifold_type = typejoin(manifold_type, typeof(_component_manifold(components[k]))) end comps = Vector{DecompositionComponent{T,N,point_type,manifold_type}}(undef, backend.r) @inbounds for k = 1:backend.r # Components keep only point/manifold metadata and reconstruct derived tensors on demand. comps[k] = DecompositionComponent{T,N,point_type,manifold_type}( parts[k], - backend.manifolds[k], + _component_manifold(components[k]), backend.target_shape, ) end @@ -473,28 +540,34 @@ function rgrad(model::JoinModel{<:AbstractFloat,<:JoinBackend}, p) _check_parts_len(parts, backend.r, "rgrad") residual = _join_residual_grad!(backend, p) vals = ntuple(k -> begin - Mk = backend.manifolds[k] + ck = backend.components[k] + Mk = _component_manifold(ck) pk = parts[k] - eg = _manifold_egrad(Mk, pk, residual) + eg = _component_egrad(ck, pk, residual) egrad_to_rgrad(Mk, pk, eg) end, backend.r) return wrap_like_point(p, vals) end """ - _component_ambient_embedding!(out, M, p) + _component_ambient_embedding!(out, component, p) Write the ambient embedding of one join component into `out`. Component manifolds may specialize this hook when their native point structure admits a more direct embedding than the generic `embed!` path. """ function _component_ambient_embedding!(out::AbstractVector, M, p) + return _component_ambient_embedding!(out, DefaultJoinEmbedding(), M, p) +end + +function _component_ambient_embedding!(out::AbstractVector, ::DefaultJoinEmbedding, M, p) ManifoldsBase.embed!(M, out, p) return out end function _component_ambient_embedding!( out::AbstractVector, + ::DefaultJoinEmbedding, M::Manifolds.Tucker, p::Manifolds.TuckerPoint, ) @@ -505,11 +578,29 @@ function _component_ambient_embedding!( end function _component_ambient_embedding!(out::AbstractVector, M::Manifolds.Segre, p) + return _component_ambient_embedding!(out, DefaultJoinEmbedding(), M, p) +end + +function _component_ambient_embedding!( + out::AbstractVector, + ::DefaultJoinEmbedding, + M::Manifolds.Segre, + p, +) copyto!(out, _segre_component_tensorvec(p)) return out end function _component_ambient_embedding!(out::AbstractVector, M::Manifolds.Tucker, p) + return _component_ambient_embedding!(out, DefaultJoinEmbedding(), M, p) +end + +function _component_ambient_embedding!( + out::AbstractVector, + ::DefaultJoinEmbedding, + M::Manifolds.Tucker, + p, +) throw( ArgumentError( "Expected native TuckerPoint for Manifolds.Tucker, got $(typeof(p)).", @@ -517,14 +608,28 @@ function _component_ambient_embedding!(out::AbstractVector, M::Manifolds.Tucker, ) end +function _component_ambient_embedding!(out::AbstractVector, component::JoinComponent, p) + return _component_ambient_embedding!(out, component.embedding, component.manifold, p) +end + """ - _component_ambient_pushforward!(out, M, p, X) + _component_ambient_pushforward!(out, component, p, X) Write the ambient pushforward `DΦ(p)[X]` of one join component into `out`. This is the component-level differential used by LM Jacobian assembly on generic `JoinModel`s. """ function _component_ambient_pushforward!(out::AbstractVector, M, p, X) + return _component_ambient_pushforward!(out, DefaultJoinEmbedding(), M, p, X) +end + +function _component_ambient_pushforward!( + out::AbstractVector, + ::DefaultJoinEmbedding, + M, + p, + X, +) emb = ManifoldsBase.embed(M, p, X) length(emb) == length(out) || throw( DimensionMismatch( @@ -536,13 +641,38 @@ function _component_ambient_pushforward!(out::AbstractVector, M, p, X) end function _component_ambient_pushforward!(out::AbstractVector, M::Manifolds.Segre, p, X) + return _component_ambient_pushforward!(out, DefaultJoinEmbedding(), M, p, X) +end + +function _component_ambient_pushforward!( + out::AbstractVector, + ::DefaultJoinEmbedding, + M::Manifolds.Segre, + p, + X, +) copyto!(out, _segre_tangent_tensorvec(p, X)) return out end +function _component_ambient_pushforward!( + out::AbstractVector, + component::JoinComponent, + p, + X, +) + return _component_ambient_pushforward!( + out, + component.embedding, + component.manifold, + p, + X, + ) +end + function _subtract_ambient_tensor!( residual::AbstractArray{T,N}, - M, + component, p, work_vec::AbstractVector{T}, ) where {T<:AbstractFloat,N} @@ -551,7 +681,7 @@ function _subtract_ambient_tensor!( "_subtract_ambient_tensor!: work length $(length(work_vec)) != residual length $(length(residual))", ), ) - _component_ambient_embedding!(work_vec, M, p) + _component_ambient_embedding!(work_vec, component, p) residual_vec = vec(residual) @inbounds for i in eachindex(residual_vec, work_vec) residual_vec[i] -= work_vec[i] @@ -588,7 +718,7 @@ then accumulates those buffers into `out`. This avoids allocating one dense ambient tensor per component during solver iterations. """ function _join_reconstruct!(out::AbstractArray, backend::Union{JoinBackend,BTDBackend}, p) - manifolds = backend.manifolds + components = _backend_components(backend) r = backend.r bufs = backend.component_bufs @@ -599,7 +729,7 @@ function _join_reconstruct!(out::AbstractArray, backend::Union{JoinBackend,BTDBa @inbounds for k = 1:r # Reconstruct each component into its preallocated workspace. - _component_ambient_embedding!(bufs[k], manifolds[k], parts[k]) + _component_ambient_embedding!(bufs[k], components[k], parts[k]) # Accumulate into the output tensor without allocating a Khatri-Rao-sized object. out .+= bufs[k] diff --git a/src/join/join_model.jl b/src/join/join_model.jl index c6218a1..1b92e2c 100644 --- a/src/join/join_model.jl +++ b/src/join/join_model.jl @@ -1,10 +1,22 @@ # join/join_model.jl — Join-front-end model and backend type definitions -export AbstractJoinBackend, JoinModel, CPDBackend, JoinBackend, BTDBackend +export AbstractJoinBackend, JoinComponent, JoinModel, CPDBackend, JoinBackend, BTDBackend # BTDBackend is defined in `btd/model.jl` (includes contraction workspace). abstract type AbstractJoinBackend end +struct JoinComponent{M,E} + manifold::M + embedding::E +end + +struct DefaultJoinEmbedding end + +JoinComponent(manifold::M) where {M} = + JoinComponent{M,DefaultJoinEmbedding}(manifold, DefaultJoinEmbedding()) + +manifold(component::JoinComponent) = component.manifold + struct JoinModel{T<:AbstractFloat,B<:AbstractJoinBackend} <: AbstractDecompositionModel{T} backend::B end @@ -24,7 +36,7 @@ _JoinResidualWORO(residual::V) where {V} = struct JoinBackend{ T, N, - MT<:Tuple, + CT<:Tuple, A<:AbstractArray{T,N}, V, MP<:ProductManifold, @@ -32,7 +44,7 @@ struct JoinBackend{ W<:_JoinResidualWORO{T}, C, } <: AbstractJoinBackend - manifolds::MT + components::CT r::Int target::A target_shape::NTuple{N,Int} diff --git a/src/solvers/lm.jl b/src/solvers/lm.jl index 77391f4..d0f6d34 100644 --- a/src/solvers/lm.jl +++ b/src/solvers/lm.jl @@ -47,10 +47,11 @@ function _join_tangent_ambient_vector!( _check_parts_len(parts, backend.r, "_join_tangent_ambient_vector!") _check_parts_len(xparts, backend.r, "_join_tangent_ambient_vector!") fill!(out, zero(eltype(out))) + components = _backend_components(backend) @inbounds for k = 1:backend.r _component_ambient_pushforward!( backend.component_bufs[k], - backend.manifolds[k], + components[k], parts[k], xparts[k], ) diff --git a/test/basic_tests.jl b/test/basic_tests.jl index 0f5ab1f..ab7913b 100644 --- a/test/basic_tests.jl +++ b/test/basic_tests.jl @@ -186,7 +186,8 @@ end @test model.backend isa TensorKitchen.JoinBackend @test model.backend.r == 3 @test tensor(model) == A - @test model.backend.manifolds == segres + @test length(model.backend.components) == 3 + @test map(TensorKitchen.manifold, model.backend.components) == segres M = TensorKitchen.manifold(model) @test M isa ProductManifold @@ -195,7 +196,12 @@ end model_repeat = JoinModel(Manifolds.Segre(dims), 3, A) @test model_repeat.backend.r == 3 - @test model_repeat.backend.manifolds == segres + @test map(TensorKitchen.manifold, model_repeat.backend.components) == segres + + model_component = + JoinModel((TensorKitchen.JoinComponent(segres[1]), segres[2], segres[3]), A) + @test model_component.backend.components[1] isa TensorKitchen.JoinComponent + @test map(TensorKitchen.manifold, model_component.backend.components) == segres p = TensorKitchen.initial_point(model, :random; verbose = false) @test length(TensorKitchen.point_parts(p)) == 3