diff --git a/.github/workflows/format_check.yml b/.github/workflows/format_check.yml index 2b61af9..aaefa39 100644 --- a/.github/workflows/format_check.yml +++ b/.github/workflows/format_check.yml @@ -18,8 +18,7 @@ jobs: - uses: actions/checkout@v6 - name: Install JuliaFormatter and format run: | - julia --project=. -e 'using Pkg; Pkg.instantiate()' - julia --project=. -e 'using JuliaFormatter; format(["./src", "./test"], verbose=true)' + julia -e 'using Pkg; Pkg.activate(temp=true); Pkg.add(PackageSpec(name="JuliaFormatter", version="2.10.1")); using JuliaFormatter; format(["./src", "./test"], verbose=true)' - name: Format check run: | julia -e ' diff --git a/Project.toml b/Project.toml index a3bb5e9..6e1c292 100644 --- a/Project.toml +++ b/Project.toml @@ -14,10 +14,10 @@ RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2" [compat] -JuliaFormatter = "2.8.5" +JuliaFormatter = "=2.10.1" Manifolds = "0.11.28" ManifoldsBase = "2.3.5" -Manopt = "0.5.37" +Manopt = "0.6" ProgressMeter = "1.11.0" RecursiveArrayTools = "4.3" TensorOperations = "5.6" diff --git a/src/api/btd.jl b/src/api/btd.jl index 1127881..f2307af 100644 --- a/src/api/btd.jl +++ b/src/api/btd.jl @@ -63,6 +63,15 @@ _btd_uses_warm_start(::AbstractSolver, ::BTDALSWarmStartInit) = true _btd_should_polish(::ALSSolver, ::Integer) = false _btd_should_polish(::AbstractSolver, polish_n::Integer) = polish_n > 0 +function _reject_unsupported_btd_solver(solver_obj) + solver_obj isa LMSolver || return nothing + throw( + ArgumentError( + "BTD currently does not support LM refinement because the required Manopt operator path is not yet available for nested Tucker layouts. Use :rgd, :rcg, :lbfgs, :als, or :btd_tsd instead.", + ), + ) +end + function _btd_warm_start_result( model::JoinModel{T,<:BTDBackend}, backend::BTDBackend, @@ -159,7 +168,10 @@ refines it. Returns a [`BTDResult`](@ref). - `:als`: Alternating least squares. - `:rcg`: Riemannian conjugate gradient. - `:lbfgs`: Limited-memory quasi-Newton refinement. - - `:lm`: Levenberg-Marquardt refinement. + - `:btd_tsd`: Blockwise tangent-subspace descent for BTD. + +`solver = :lm` is currently not supported for BTD because the required Manopt +LM operator path is not yet available for nested Tucker layouts. ## Extended Options @@ -234,6 +246,7 @@ function btd( kwargs..., ) where {T<:AbstractFloat,N} solver_obj = _solver_object(solver, stepsize; kwargs...) + _reject_unsupported_btd_solver(solver_obj) solver_sym = _btd_solver_symbol(solver_obj) init_resolved = _resolve_btd_init(init, solver_obj) init_eff = diff --git a/src/btd/model.jl b/src/btd/model.jl index 509905d..400ef95 100644 --- a/src/btd/model.jl +++ b/src/btd/model.jl @@ -31,8 +31,18 @@ Backend state for block-term decomposition as a sum of Tucker blocks. Stores the target tensor, product manifold, reusable work buffers, and per-block ambient reconstruction buffers used by cost, gradient, and ALS routines. """ -struct BTDBackend{T,N,MT<:Tuple,A<:AbstractArray{T,N},V,MP<:ProductManifold,I,C} <: - AbstractJoinBackend +struct BTDBackend{ + T, + N, + CT<:Tuple, + MT<:Tuple, + A<:AbstractArray{T,N}, + V, + MP<:ProductManifold, + I, + C, +} <: AbstractJoinBackend + components::CT manifolds::MT r::Int # Preserve the target array/backend so BTD shares the generic join storage behavior. @@ -60,7 +70,7 @@ model_exact_join_basis_function(model::JoinModel{<:AbstractFloat,<:BTDBackend}) (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 @@ -71,14 +81,11 @@ function _btd_sequential_tucker_init(model::JoinModel{<:AbstractFloat,<:BTDBacke residual = copy(backend.target) parts = Vector{Manifolds.TuckerPoint{eltype(backend.target)}}(undef, backend.r) for k = 1:backend.r - pk = _manifold_init(backend.manifolds[k], residual, init) + component = _backend_component(backend, k) + Mk = _backend_manifold(backend, k) + pk = _component_init(component, residual, init) parts[k] = pk - _subtract_ambient_tensor!( - residual, - backend.manifolds[k], - pk, - backend.component_bufs[k], - ) + _subtract_ambient_tensor!(residual, Mk, pk, backend.component_bufs[k]) end return ArrayPartition(parts...) end @@ -86,7 +93,7 @@ end function _btd_block_ranks_by_mode(backend::BTDBackend{T,N}) where {T,N} ranks = Vector{NTuple{N,Int}}(undef, backend.r) for b = 1:backend.r - M = backend.manifolds[b] + M = _backend_manifold(backend, b) M isa Manifolds.Tucker || throw( ArgumentError( "BTD HOSVD multistart expects Tucker manifolds, got $(typeof(M)) at block $b.", @@ -160,7 +167,7 @@ function _btd_hosvd_split_candidate( parts[b] = pk _subtract_ambient_tensor!( residual, - backend.manifolds[b], + _backend_manifold(backend, b), pk, backend.component_bufs[b], ) @@ -190,7 +197,7 @@ function initial_point( init_sym = _builtin_initializer_symbol(init) if init_sym == :random parts = ntuple( - k -> _manifold_init(backend.manifolds[k], backend.target, init), + k -> _component_init(_backend_component(backend, k), backend.target, init), backend.r, ) return ArrayPartition(parts...) @@ -282,6 +289,9 @@ function rgrad(model::JoinModel{<:AbstractFloat,<:BTDBackend}, p) parts = point_parts(p) _check_parts_len(parts, backend.r, "BTD rgrad") eg = point_parts(_btd_egrad(backend, p)) - vals = ntuple(k -> egrad_to_rgrad(backend.manifolds[k], parts[k], eg[k]), backend.r) + vals = ntuple( + k -> egrad_to_rgrad(_backend_manifold(backend, k), parts[k], eg[k]), + backend.r, + ) return wrap_like_point(p, vals) end diff --git a/src/core/model.jl b/src/core/model.jl index f598072..ac18139 100644 --- a/src/core/model.jl +++ b/src/core/model.jl @@ -1,6 +1,17 @@ # core/model.jl — Top-level decomposition model interface export AbstractDecompositionModel, - manifold, initial_point, egrad, rgrad, supports_rgrad, tensor, cost, post_step! + manifold, + initial_point, + egrad, + rgrad, + supports_rgrad, + tensor, + cost, + post_step!, + residual, + differential_action, + differential_action!, + adjoint_action """ AbstractDecompositionModel{T} @@ -113,3 +124,49 @@ end function tensor(model::AbstractDecompositionModel) error("tensor not implemented for $(typeof(model))") end + +function residual(model::AbstractDecompositionModel, p) + error("residual not implemented for $(typeof(model))") +end + +function differential_action!(out::AbstractVector, model::AbstractDecompositionModel, p, X) + error("differential_action! not implemented for $(typeof(model))") +end + +function differential_action(model::AbstractDecompositionModel, p, X) + T = eltype(tensor(model)) + out = Vector{T}(undef, length(tensor(model))) + differential_action!(out, model, p, X) + return out +end + +function adjoint_action( + model::AbstractDecompositionModel, + p, + a::AbstractVector; + basis = ManifoldsBase.DefaultOrthonormalBasis(), +) + M = manifold(model) + d = manifold_dimension(M) + length(a) == length(tensor(model)) || throw( + DimensionMismatch( + "adjoint_action expected ambient vector of length $(length(tensor(model))) for $(typeof(model)), got $(length(a)).", + ), + ) + T = _scalar_eltype(p) + coeff = zeros(T, d) + e_j = zeros(T, d) + col = Vector{T}(undef, length(a)) + @inbounds for j = 1:d + fill!(e_j, zero(T)) + e_j[j] = one(T) + Xj = ManifoldsBase.get_vector(M, p, e_j, basis) + differential_action!(col, model, p, Xj) + coeff[j] = dot(col, a) + end + return ManifoldsBase.get_vector(M, p, coeff, basis) +end + +function adjoint_action(model::AbstractDecompositionModel, p, a::AbstractArray; kwargs...) + return adjoint_action(model, p, vec(a); kwargs...) +end diff --git a/src/cpd/model/parameterizations.jl b/src/cpd/model/parameterizations.jl new file mode 100644 index 0000000..e4a6070 --- /dev/null +++ b/src/cpd/model/parameterizations.jl @@ -0,0 +1,352 @@ +# cpd/model/parameterizations.jl — CP point parameterization helpers + +abstract type AbstractCPParameterization end + +struct NativeCPParam <: AbstractCPParameterization end +struct CanonicalCPParam <: AbstractCPParameterization end +struct SquaredNNCPParam <: AbstractCPParameterization end +struct SoftplusNNCPParam <: AbstractCPParameterization end + +@inline _cp_softplus_encode_value(x::T) where {T<:AbstractFloat} = + _invsoftplus(max(x, eps(T))) + +function _cp_rank1_decode_factors(::NativeCPParam, dims, p) + return unpack_point_rank1(p, dims) +end + +function _cp_rank1_decode_factors(::SquaredNNCPParam, dims, p) + λ̃, Ũ = unpack_point_rank1(p, dims) + return λ̃^2, [Ũ[m] .^ 2 for m in eachindex(Ũ)] +end + +function _cp_rank1_decode_factors(::SoftplusNNCPParam, dims, p) + λ̃, Ũ = unpack_point_rank1(p, dims) + return _softplus_value(λ̃), [_softplus_value.(Ũ[m]) for m in eachindex(Ũ)] +end + +function _cp_rank1_encode_point(::NativeCPParam, λ, U) + return pack_point_rank1_segre(λ, U) +end + +function _cp_rank1_encode_point(::SquaredNNCPParam, λ, U) + return pack_point_rank1( + sqrt(max(λ, zero(λ))), + [sqrt.(max.(u, zero(eltype(u)))) for u in U], + ) +end + +function _cp_rank1_encode_point(::SoftplusNNCPParam, λ, U) + T = typeof(λ) + return pack_point_rank1( + _cp_softplus_encode_value(max(λ, zero(T))), + [_cp_softplus_encode_value.(max.(u, zero(eltype(u)))) for u in U], + ) +end + +function _cp_rank1_seed_point(::NativeCPParam, λ, U) + return pack_point_rank1_segre(λ, U) +end + +function _cp_rank1_seed_point(::SquaredNNCPParam, λ, U) + T = typeof(λ) + return pack_point_rank1( + sqrt(max(abs(λ), eps(T))), + [sqrt.(max.(abs.(u), eps(eltype(u)))) for u in U], + ) +end + +function _cp_rank1_seed_point(::SoftplusNNCPParam, λ, U) + T = typeof(λ) + return pack_point_rank1( + _invsoftplus(max(abs(λ), eps(T))), + [_invsoftplus.(max.(abs.(u), eps(eltype(u)))) for u in U], + ) +end + +function _cp_rank1_embed_tensor(embedding::AbstractCPParameterization, dims, p) + λ, U = _cp_rank1_decode_factors(embedding, dims, p) + return reconstruct_cp_rank1(λ, U) +end + +function _cp_rank1_tangent_tensorvec!( + out::AbstractVector{T}, + λ::T, + U::AbstractVector{<:AbstractVector{T}}, + λ̇::T, + U̇::AbstractVector{<:AbstractVector{T}}, +) where {T<:AbstractFloat} + comp = ([λ], U...) + xcomp = ([λ̇], U̇...) + copyto!(out, _segre_tangent_tensorvec(comp, xcomp)) + return out +end + +function _cp_rank1_decode_tangent_factors(::NativeCPParam, dims, p, X) + λ, U = unpack_point_rank1(p, dims) + λ̇, U̇ = unpack_point_rank1(X, dims) + return λ, U, λ̇, U̇ +end + +function _cp_rank1_decode_tangent_factors(::SquaredNNCPParam, dims, p, X) + λ̃, Ũ = unpack_point_rank1(p, dims) + λ̇̃, U̇̃ = unpack_point_rank1(X, dims) + λ = λ̃^2 + U = [Ũ[m] .^ 2 for m in eachindex(Ũ)] + λ̇ = 2 * λ̃ * λ̇̃ + U̇ = [2 .* Ũ[m] .* U̇̃[m] for m in eachindex(Ũ)] + return λ, U, λ̇, U̇ +end + +function _cp_rank1_decode_tangent_factors(::SoftplusNNCPParam, dims, p, X) + λ̃, Ũ = unpack_point_rank1(p, dims) + λ̇̃, U̇̃ = unpack_point_rank1(X, dims) + λ = _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_rank1_linear_egrad(::NativeCPParam, dims, p, A) + λ, U = unpack_point_rank1(p, dims) + T = typeof(λ) + grad_λ = rank1_inner(A, U) + grad_U = Vector{Vector{T}}(undef, length(dims)) + @inbounds for m in eachindex(dims) + g = Vector{T}(undef, dims[m]) + rank1_mode_contract!(g, A, U, m) + rmul!(g, λ) + grad_U[m] = g + end + return pack_tangent_rank1_segre(grad_λ, grad_U) +end + +function _cp_rank1_linear_egrad(::SquaredNNCPParam, dims, p, A) + λ̃, Ũ = unpack_point_rank1(p, dims) + λ = λ̃^2 + U = [Ũ[m] .^ 2 for m in eachindex(Ũ)] + grad_λ = rank1_inner(A, U) + grad_U = Vector{Vector{eltype(λ̃)}}(undef, length(dims)) + @inbounds for m in eachindex(dims) + g = Vector{eltype(λ̃)}(undef, dims[m]) + rank1_mode_contract!(g, A, U, m) + g .*= λ .* 2 .* Ũ[m] + grad_U[m] = g + end + grad_λ̃ = grad_λ * 2 * λ̃ + return p isa Vector ? pack_point_rank1_to_vector(grad_λ̃, grad_U) : + pack_point_rank1(grad_λ̃, grad_U) +end + +function _cp_rank1_linear_egrad(::SoftplusNNCPParam, dims, p, A) + λ̃, Ũ = unpack_point_rank1(p, dims) + λ = _softplus_value(λ̃) + U = [_softplus_value.(Ũ[m]) for m in eachindex(Ũ)] + grad_λ = rank1_inner(A, U) + grad_U = Vector{Vector{eltype(λ̃)}}(undef, length(dims)) + @inbounds for m in eachindex(dims) + g = Vector{eltype(λ̃)}(undef, dims[m]) + rank1_mode_contract!(g, A, U, m) + g .*= λ .* _softplus_derivative.(Ũ[m]) + grad_U[m] = g + end + grad_λ̃ = grad_λ * _softplus_derivative(λ̃) + return p isa Vector ? pack_point_rank1_to_vector(grad_λ̃, grad_U) : + pack_point_rank1(grad_λ̃, grad_U) +end + +function _cp_rankr_decode_factors(::NativeCPParam, dims, r, p) + return unpack_rankr_native(p, dims, r) +end + +function _cp_rankr_decode_factors(::CanonicalCPParam, dims, r, p) + return unpack_rankr_canonical(p, dims, r) +end + +function _cp_rankr_decode_factors(::SquaredNNCPParam, dims, r, p) + λ̃, Ũ = unpack_point_rankr(p, dims, r) + return λ̃ .^ 2, [Ũ[m] .^ 2 for m in eachindex(Ũ)] +end + +function _cp_rankr_decode_factors(::SoftplusNNCPParam, dims, r, p) + λ̃, Ũ = unpack_point_rankr(p, dims, r) + return _softplus_value.(λ̃), [_softplus_value.(Ũ[m]) for m in eachindex(Ũ)] +end + +function _cp_rankr_encode_point(::NativeCPParam, λ, U, r) + return pack_rankr_native(λ, U, r) +end + +function _cp_rankr_encode_point(::CanonicalCPParam, λ, U, r) + return pack_rankr_canonical(λ, U, r) +end + +function _cp_rankr_encode_point(::SquaredNNCPParam, λ, U, r) + T = eltype(λ) + return pack_point_rankr( + sqrt.(max.(λ, zero(T))), + [sqrt.(max.(F, zero(eltype(F)))) for F in U], + r, + ) +end + +function _cp_rankr_encode_point(::SoftplusNNCPParam, λ, U, r) + T = eltype(λ) + return pack_point_rankr( + _cp_softplus_encode_value.(max.(λ, zero(T))), + [_cp_softplus_encode_value.(max.(F, zero(eltype(F)))) for F in U], + r, + ) +end + +function _cp_rankr_seed_point(::NativeCPParam, λ, U, r) + return pack_rankr_native(λ, U, r) +end + +function _cp_rankr_seed_point(::CanonicalCPParam, λ, U, r) + return pack_rankr_canonical(λ, U, r) +end + +function _cp_rankr_seed_point(::SquaredNNCPParam, λ, U, r) + T = eltype(λ) + return pack_point_rankr( + sqrt.(max.(abs.(λ), eps(T))), + [sqrt.(max.(abs.(F), eps(eltype(F)))) for F in U], + r, + ) +end + +function _cp_rankr_seed_point(::SoftplusNNCPParam, λ, U, r) + T = eltype(λ) + return pack_point_rankr( + _invsoftplus.(max.(abs.(λ), eps(T))), + [_invsoftplus.(max.(abs.(F), eps(eltype(F)))) for F in U], + r, + ) +end + +function _cp_rankr_embed_tensor(embedding::AbstractCPParameterization, dims, r, p) + λ, U = _cp_rankr_decode_factors(embedding, dims, r, p) + return reconstruct_cpd_rankr(λ, 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(λ) + d = length(U) + Ucols = Vector{AbstractVector{T}}(undef, d) + U̇cols = Vector{AbstractVector{T}}(undef, d) + @inbounds for k = 1:r + for m = 1:d + Ucols[m] = @view U[m][:, k] + U̇cols[m] = @view U̇[m][:, k] + end + out .+= _segre_tangent_tensorvec(([λ[k]], Ucols...), ([λ̇[k]], U̇cols...)) + end + return out +end + +function _cp_rankr_decode_tangent_factors(::NativeCPParam, dims, r, p, X) + λ, U = unpack_rankr_native(p, dims, r) + xparts = parts_tuple(X) + length(xparts) == r || throw( + DimensionMismatch("expected $r native tangent components, got $(length(xparts))"), + ) + T = eltype(λ) + d = length(dims) + λ̇ = similar(λ) + U̇ = [zeros(T, dims[m], r) for m = 1:d] + @inbounds for k = 1:r + xk_λ, xk_U = unpack_point_rank1(xparts[k], dims) + λ̇[k] = xk_λ + for m = 1:d + U̇[m][:, k] .= xk_U[m] + end + end + return λ, U, λ̇, U̇ +end + +function _cp_rankr_decode_tangent_factors(::CanonicalCPParam, dims, r, p, X) + λ, U = unpack_rankr_canonical(p, dims, r) + λ̇, U̇ = unpack_rankr_canonical(X, dims, r) + return λ, U, λ̇, U̇ +end + +function _cp_rankr_decode_tangent_factors(::SquaredNNCPParam, dims, r, p, X) + λ̃, Ũ = unpack_point_rankr(p, dims, r) + λ̇̃, U̇̃ = unpack_point_rankr(X, dims, r) + λ = λ̃ .^ 2 + U = [Ũ[m] .^ 2 for m in eachindex(Ũ)] + λ̇ = 2 .* λ̃ .* λ̇̃ + U̇ = [2 .* Ũ[m] .* U̇̃[m] for m in eachindex(Ũ)] + return λ, U, λ̇, U̇ +end + +function _cp_rankr_decode_tangent_factors(::SoftplusNNCPParam, dims, r, p, X) + λ̃, Ũ = unpack_point_rankr(p, dims, r) + λ̇̃, U̇̃ = unpack_point_rankr(X, dims, r) + λ = _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_linear_terms(A, U, dims, r) + Nmodes = length(dims) + contracts = [mttkrp(A, U, m; method = :auto) for m = 1:Nmodes] + grad_λ = _inner_from_mttkrp_first_mode(U, contracts[1]) + return grad_λ, contracts +end + +function _cp_rankr_linear_egrad(::NativeCPParam, dims, r, p, A) + λ, U = unpack_rankr_native(p, dims, r) + grad_λ, contracts = _cp_rankr_linear_terms(A, U, dims, r) + grad_parts = Vector{Vector{Vector{eltype(λ)}}}(undef, r) + @inbounds for k = 1:r + grad_Uk = [λ[k] .* Vector(@view contracts[m][:, k]) for m = 1:length(dims)] + grad_parts[k] = pack_tangent_rank1_segre(grad_λ[k], grad_Uk) + end + return hasproperty(p, :x) ? ArrayPartition(grad_parts...) : (grad_parts...,) +end + +function _cp_rankr_linear_egrad(::CanonicalCPParam, dims, r, p, A) + λ, U = unpack_rankr_canonical(p, dims, r) + grad_λ, contracts = _cp_rankr_linear_terms(A, U, dims, r) + gradU = [contracts[m] .* transpose(λ) for m = 1:length(dims)] + return wrap_rankr_canonical_tangent_like(p, grad_λ, gradU, r) +end + +function _cp_rankr_linear_egrad(::SquaredNNCPParam, dims, r, p, A) + λ̃, Ũ = unpack_point_rankr(p, dims, r) + λ = λ̃ .^ 2 + U = [Ũ[m] .^ 2 for m in eachindex(Ũ)] + grad_λ, contracts = _cp_rankr_linear_terms(A, U, dims, r) + grad_λ̃ = grad_λ .* 2 .* λ̃ + gradU = Vector{Matrix{eltype(λ̃)}}(undef, length(dims)) + @inbounds for m in eachindex(dims) + gradU[m] = (contracts[m] .* transpose(λ)) .* (2 .* Ũ[m]) + end + return p isa Vector ? pack_point_rankr_to_vector(grad_λ̃, gradU, r) : + pack_point_rankr(grad_λ̃, gradU, r) +end + +function _cp_rankr_linear_egrad(::SoftplusNNCPParam, dims, r, p, A) + λ̃, Ũ = unpack_point_rankr(p, dims, r) + λ = _softplus_value.(λ̃) + U = [_softplus_value.(Ũ[m]) for m in eachindex(Ũ)] + grad_λ, contracts = _cp_rankr_linear_terms(A, U, dims, r) + grad_λ̃ = grad_λ .* _softplus_derivative.(λ̃) + gradU = Vector{Matrix{eltype(λ̃)}}(undef, length(dims)) + @inbounds for m in eachindex(dims) + gradU[m] = (contracts[m] .* transpose(λ)) .* _softplus_derivative.(Ũ[m]) + end + return p isa Vector ? pack_point_rankr_to_vector(grad_λ̃, gradU, r) : + pack_point_rankr(grad_λ̃, gradU, r) +end diff --git a/src/cpd/model/rank1.jl b/src/cpd/model/rank1.jl index 8903818..ff642d5 100644 --- a/src/cpd/model/rank1.jl +++ b/src/cpd/model/rank1.jl @@ -58,22 +58,49 @@ end tensor(model::Rank1CPDModel) = model.A manifold(model::Rank1CPDModel) = model.M +_cp_parameterization(model::Rank1CPDModel) = + model.nonnegative ? + (_rank1_uses_softplus_metric(model.M) ? SoftplusNNCPParam() : SquaredNNCPParam()) : + NativeCPParam() + function embed_point(model::Rank1CPDModel{T,N}, p) where {T,N} - if model.nonnegative - if _rank1_uses_softplus_metric(model.M) - λ̃, Ũ = unpack_point_rank1(p, model.dims) - λ = _softplus_value(λ̃) - U = [_softplus_value.(Ũ[m]) for m in eachindex(Ũ)] - return reconstruct_cp_rank1(λ, U) - end - return embed_point_rank1_nn(p, model.dims) - end - M = manifold(model) - M isa Manifolds.Segre || return embed_point_rank1(p, model.dims) - return reshape( - ManifoldsBase.embed!(M, Vector{eltype(p[1])}(undef, prod(model.dims)), p), - model.dims, + return _cp_rank1_embed_tensor(_cp_parameterization(model), model.dims, p) +end + +function residual(model::Rank1CPDModel{T,N}, p) where {T<:AbstractFloat,N} + return vec(embed_point(model, p)) .- vec(model.A) +end + +function differential_action!( + out::AbstractVector{T}, + model::Rank1CPDModel{T,N}, + p, + X, +) where {T<:AbstractFloat,N} + length(out) == length(model.A) || throw( + DimensionMismatch( + "differential_action! output length $(length(out)) != ambient length $(length(model.A)).", + ), + ) + λ, U, λ̇, U̇ = + _cp_rank1_decode_tangent_factors(_cp_parameterization(model), model.dims, p, X) + return _cp_rank1_tangent_tensorvec!(out, λ, U, λ̇, U̇) +end + +function adjoint_action( + model::Rank1CPDModel{T,N}, + p, + a::AbstractVector; + kwargs..., +) where {T<:AbstractFloat,N} + length(a) == length(model.A) || throw( + DimensionMismatch( + "adjoint_action expected ambient vector of length $(length(model.A)) for $(typeof(model)), got $(length(a)).", + ), ) + Aadj = reshape(a, model.dims) + eg = _cp_rank1_linear_egrad(_cp_parameterization(model), model.dims, p, Aadj) + return egrad_to_rgrad(model.M, p, eg) end function cost(model::Rank1CPDModel{T,N}, p) where {T,N} @@ -166,16 +193,7 @@ function initial_point( init == :alswarm && return initial_point(model, ALSWarmStartInit(); verbose) init_sym = _builtin_initializer_symbol(init) U0 = init_cp_rank1(model.A; init = init_sym) - if model.nonnegative - λ̃0 = _rank1_uses_softplus_metric(model.M) ? _invsoftplus(one(T)) : one(T) - Ũ = if _rank1_uses_softplus_metric(model.M) - [_invsoftplus.(max.(abs.(u), eps(T))) for u in U0] - else - [sqrt.(max.(abs.(u), eps(T))) for u in U0] - end - return pack_point_rank1(λ̃0, Ũ) - end - return pack_point_rank1_segre(one(T), U0) + return _cp_rank1_seed_point(_cp_parameterization(model), one(T), U0) end initial_point(model::Rank1CPDModel, init::PointInit; kwargs...) = init.point @@ -196,17 +214,7 @@ nonnegative squared parameterization so backend postprocessing can work with a uniform `lambda + factors` representation. """ function cpd_point(model::Rank1CPDModel{T,N}, p) where {T<:AbstractFloat,N} - if model.nonnegative - λ̃, Ũ = unpack_point_rank1(p, model.dims) - if _rank1_uses_softplus_metric(model.M) - return CPDPoint( - T[_softplus_value(λ̃)], - [reshape(_softplus_value.(Ũ[m]), :, 1) for m in eachindex(Ũ)], - ) - end - return CPDPoint(T[λ̃^2], [reshape(Ũ[m] .^ 2, :, 1) for m in eachindex(Ũ)]) - end - λ, U = unpack_point_rank1(p, model.dims) + λ, U = _cp_rank1_decode_factors(_cp_parameterization(model), model.dims, p) return CPDPoint(T[λ], [reshape(U[m], :, 1) for m in eachindex(U)]) end @@ -221,16 +229,7 @@ function pack_cpd_point( ) λ = lambda(point)[1] U = [Vector(@view F[:, 1]) for F in factors(point)] - if model.nonnegative - if _rank1_uses_softplus_metric(model.M) - return pack_point_rank1( - _invsoftplus(max(λ, zero(T))), - [_invsoftplus.(max.(u, zero(T))) for u in U], - ) - end - return pack_point_rank1(sqrt(max(λ, zero(T))), [sqrt.(max.(u, zero(T))) for u in U]) - end - return pack_point_rank1_segre(λ, U) + return _cp_rank1_encode_point(_cp_parameterization(model), λ, U) end function post_step!( diff --git a/src/cpd/model/rankr.jl b/src/cpd/model/rankr.jl index 8df2904..363f1e0 100644 --- a/src/cpd/model/rankr.jl +++ b/src/cpd/model/rankr.jl @@ -100,20 +100,54 @@ end tensor(model::RankRCPDModel) = model.A manifold(model::RankRCPDModel) = model.M +_cp_parameterization(model::RankRCPDModel) = + model.nonnegative ? + (model.geometry == :softplus_metric ? SoftplusNNCPParam() : SquaredNNCPParam()) : + (model.geometry == :native ? NativeCPParam() : CanonicalCPParam()) + function embed_point(model::RankRCPDModel{T,N}, p) where {T,N} - if model.nonnegative - λ̃, Ũ = unpack_point_rankr(p, model.dims, model.r) - if model.geometry == :softplus_metric - λ = _softplus_value.(λ̃) - U = [_softplus_value.(Ũ[m]) for m in eachindex(Ũ)] - return reconstruct_cpd_rankr(λ, U) - end - return embed_point_rankr_nn(p, model.dims, model.r) - end - λ, U = - model.geometry == :native ? unpack_rankr_native(p, model.dims, model.r) : - unpack_rankr_canonical(p, model.dims, model.r) - return reconstruct_cpd_rankr(λ, U) + return _cp_rankr_embed_tensor(_cp_parameterization(model), model.dims, model.r, p) +end + +function residual(model::RankRCPDModel{T,N}, p) where {T<:AbstractFloat,N} + return vec(embed_point(model, p)) .- vec(model.A) +end + +function differential_action!( + out::AbstractVector{T}, + model::RankRCPDModel{T,N}, + p, + X, +) where {T<:AbstractFloat,N} + length(out) == length(model.A) || throw( + DimensionMismatch( + "differential_action! output length $(length(out)) != ambient length $(length(model.A)).", + ), + ) + λ, U, λ̇, U̇ = _cp_rankr_decode_tangent_factors( + _cp_parameterization(model), + model.dims, + model.r, + p, + X, + ) + return _cp_rankr_tangent_tensorvec!(out, λ, U, λ̇, U̇) +end + +function adjoint_action( + model::RankRCPDModel{T,N}, + p, + a::AbstractVector; + kwargs..., +) where {T<:AbstractFloat,N} + length(a) == length(model.A) || throw( + DimensionMismatch( + "adjoint_action expected ambient vector of length $(length(model.A)) for $(typeof(model)), got $(length(a)).", + ), + ) + Aadj = reshape(a, model.dims) + eg = _cp_rankr_linear_egrad(_cp_parameterization(model), model.dims, model.r, p, Aadj) + return egrad_to_rgrad(model.M, p, eg) end function cost(model::RankRCPDModel{T,N}, p) where {T,N} @@ -422,18 +456,7 @@ function initial_point( init == :alswarm && return initial_point(model, ALSWarmStartInit(); verbose) init_sym = _builtin_initializer_symbol(init) λ0, U0 = init_cpd_factors(model.A, model.r; init = init_sym) # initialize the factors - if model.nonnegative - if model.geometry == :softplus_metric - λ̃0 = _invsoftplus.(max.(abs.(λ0), eps(T))) - Ũ0 = [_invsoftplus.(max.(abs.(U0[m]), eps(T))) for m in eachindex(U0)] - else - λ̃0 = sqrt.(max.(abs.(λ0), eps(T))) - Ũ0 = [sqrt.(max.(abs.(U0[m]), eps(T))) for m in eachindex(U0)] - end - return pack_point_rankr(λ̃0, Ũ0, model.r) # structured join layout - end - return model.geometry == :native ? pack_rankr_native(λ0, U0, model.r) : - pack_rankr_canonical(λ0, U0, model.r) # native: ArrayPartition, canonical: tuple + return _cp_rankr_seed_point(_cp_parameterization(model), λ0, U0, model.r) end initial_point(model::RankRCPDModel, init::PointInit; kwargs...) = init.point @@ -445,19 +468,7 @@ supports_normalization_policy(model::RankRCPDModel, policy::AbstractNormalizatio policy isa Union{NoNormalization,SeparateLambdaNormalization} function cpd_point(model::RankRCPDModel{T,N}, p) where {T<:AbstractFloat,N} - if model.nonnegative - λ̃, Ũ = unpack_point_rankr(p, model.dims, model.r) - if model.geometry == :softplus_metric - return CPDPoint( - _softplus_value.(λ̃), - [_softplus_value.(Ũ[m]) for m in eachindex(Ũ)], - ) - end - return CPDPoint(λ̃ .^ 2, [Ũ[m] .^ 2 for m in eachindex(Ũ)]) - end - λ, U = - model.geometry == :native ? unpack_rankr_native(p, model.dims, model.r) : - unpack_rankr_canonical(p, model.dims, model.r) + λ, U = _cp_rankr_decode_factors(_cp_parameterization(model), model.dims, model.r, p) return CPDPoint(λ, U) end @@ -470,19 +481,12 @@ function pack_cpd_point( "CPDPoint has $(length(lambda(point))) weights, expected rank $(model.r)", ), ) - if model.nonnegative - if model.geometry == :softplus_metric - λ̃ = _invsoftplus.(max.(lambda(point), zero(T))) - Ũ = [_invsoftplus.(max.(F, zero(T))) for F in factors(point)] - else - λ̃ = sqrt.(max.(lambda(point), zero(T))) - Ũ = [sqrt.(max.(F, zero(T))) for F in factors(point)] - end - return pack_point_rankr(λ̃, Ũ, model.r) - end - return model.geometry == :native ? - pack_rankr_native(lambda(point), factors(point), model.r) : - pack_rankr_canonical(lambda(point), factors(point), model.r) + return _cp_rankr_encode_point( + _cp_parameterization(model), + lambda(point), + factors(point), + model.r, + ) end function post_step!( diff --git a/src/decompositions.jl b/src/decompositions.jl index 7e8c143..1fd2713 100644 --- a/src/decompositions.jl +++ b/src/decompositions.jl @@ -8,6 +8,7 @@ include("cpd/core/mttkrp.jl") include("cpd/core/cpd_init.jl") include("cpd/core/cp_cost.jl") include("cpd/core/reconstruct.jl") +include("cpd/model/parameterizations.jl") include("cpd/model/rank1.jl") include("cpd/model/rankr.jl") diff --git a/src/join/cpd_backend.jl b/src/join/cpd_backend.jl index 123a6b5..ab3d7e3 100644 --- a/src/join/cpd_backend.jl +++ b/src/join/cpd_backend.jl @@ -59,6 +59,19 @@ initial_point( ) = initial_point(cpd_model(model), init; kwargs...) cost(model::JoinModel{<:AbstractFloat,<:CPDBackend}, p) = cost(cpd_model(model), p) egrad(model::JoinModel{<:AbstractFloat,<:CPDBackend}, p) = egrad(cpd_model(model), p) +residual(model::JoinModel{<:AbstractFloat,<:CPDBackend}, p) = residual(cpd_model(model), p) +differential_action!( + out::AbstractVector, + model::JoinModel{<:AbstractFloat,<:CPDBackend}, + p, + X, +) = differential_action!(out, cpd_model(model), p, X) +adjoint_action( + model::JoinModel{<:AbstractFloat,<:CPDBackend}, + p, + a::AbstractVector; + kwargs..., +) = adjoint_action(cpd_model(model), p, a; kwargs...) supports_rgrad(model::JoinModel{<:AbstractFloat,<:CPDBackend}) = supports_rgrad(cpd_model(model)) rgrad(model::JoinModel{<:AbstractFloat,<:CPDBackend}, p) = rgrad(cpd_model(model), p) @@ -96,60 +109,34 @@ function _public_cpd_factors(m, λ, U) return λ, U end -@inline function _decode_nonnegative_cpd(m, λ̃, Ũ) - use_softplus = - hasproperty(m, :geometry) ? (m.geometry == :softplus_metric) : - (hasproperty(m, :M) && _rank1_uses_softplus_metric(m.M)) - if use_softplus - return _softplus_value.(λ̃), [_softplus_value.(Ũ[j]) for j in eachindex(Ũ)] +@inline _cpd_solver_point(m, p, solver_sym::Symbol) = cpd_point(m, p) + +function _cpd_solver_point( + m::RankRCPDModel{T}, + p, + solver_sym::Symbol, +) where {T<:AbstractFloat} + if m.nonnegative && (solver_sym in _CP_ALS_FAMILY_SOLVERS) + λ, U = unpack_rankr_canonical(p, m.dims, m.r) + return CPDPoint(λ, U) end - return λ̃ .^ 2, [Ũ[j] .^ 2 for j in eachindex(Ũ)] + return cpd_point(m, p) end function _cpd_result(model::JoinModel{<:AbstractFloat,<:CPDBackend}, result, dims, r) m = cpd_model(model) solver_sym = _result_solver_symbol(solver(result)) si = _result_solver_info(result) - als_family = solver_sym in _CP_ALS_FAMILY_SOLVERS - - if r == 1 - λ̃, U_vec = unpack_point_rank1(point(result), dims) - if m.nonnegative && !als_family - λ_vec, U_sq = _decode_nonnegative_cpd(m, [λ̃], U_vec) - λ = λ_vec[1] - else - λ = λ̃ - U_sq = U_vec - end - λ_pub, U_pub = _public_cpd_factors(m, [λ], [reshape(u, :, 1) for u in U_sq]) - return CPDResult( - λ_pub, - U_pub, - cost(result), - rel_error(result), - grad_norm(result), - iterations(result), - converged(result), - solver_sym, - si, - ) - end - - comps = if m.nonnegative && !als_family - λ̃, Ũ = unpack_point_rankr(point(result), dims, r) - λ, U = _decode_nonnegative_cpd(m, λ̃, Ũ) - components_from_factors(λ, U) - else - unpack_point_rankr_components(point(result), dims, r) - end + q = _cpd_solver_point(m, point(result), solver_sym) + λ_raw = lambda(q) + U_raw = factors(q) rel_err = rel_error(result) if !isfinite(rel_err) - Xhat = reconstruct_cpd_rankr([c.λ for c in comps], factors_from_components(comps)) + Xhat = reconstruct_cpd_rankr(λ_raw, U_raw) rel_err = rel_error(m.A, Xhat) end - λ_pub, U_pub = - _public_cpd_factors(m, [c.λ for c in comps], factors_from_components(comps)) + λ_pub, U_pub = _public_cpd_factors(m, λ_raw, U_raw) return CPDResult( λ_pub, U_pub, @@ -167,31 +154,12 @@ function extract_components(model::JoinModel{<:AbstractFloat,<:CPDBackend}, p) return _extract_cpd_components(cpd_model(model), p) end -function _extract_cpd_components(m::RankRCPDModel, p) - comps = - m.nonnegative ? begin - λ̃, Ũ = unpack_point_rankr(p, m.dims, m.r) - λ, U = _decode_nonnegative_cpd(m, λ̃, Ũ) - components_from_factors(λ, U) - end : unpack_point_rankr_components(p, m.dims, m.r) +function _extract_cpd_components(m::Union{Rank1CPDModel,RankRCPDModel}, p) + q = cpd_point(m, p) + comps = components_from_factors(lambda(q), factors(q)) return [CPDComponent(pack_point_rank1(c.λ, c.vectors), c) for c in comps] end -function _extract_cpd_components(m::Rank1CPDModel, p) - λ, U = unpack_point_rank1(p, m.dims) - if m.nonnegative - if _rank1_uses_softplus_metric(m.M) - λ = _softplus_value(λ) - U = [_softplus_value.(u) for u in U] - else - λ = λ^2 - U = [u .^ 2 for u in U] - end - end - c = RankOneTensor(λ, U) - return [CPDComponent(pack_point_rank1(λ, U), c)] -end - function _extract_cpd_components(m, p) throw( ArgumentError( diff --git a/src/join/join_backend.jl b/src/join/join_backend.jl index 5e828b4..3be7835 100644 --- a/src/join/join_backend.jl +++ b/src/join/join_backend.jl @@ -8,10 +8,14 @@ _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_embedding(component::JoinComponent) = component_embedding(component) +_component_embedding(::AbstractManifold) = DefaultJoinEmbedding() +_component_manifold(component::JoinComponent) = component_manifold(component) _component_manifold(manifold::AbstractManifold) = manifold _backend_components(backend::JoinBackend) = backend.components -_backend_components(backend::BTDBackend) = backend.manifolds +_backend_components(backend::BTDBackend) = backend.components +_backend_component(backend, k::Int) = _backend_components(backend)[k] +_backend_manifold(backend, k::Int) = _component_manifold(_backend_component(backend, k)) function _as_join_manifold_tuple(manifolds::Tuple) all(_is_manifold_like, manifolds) || throw( @@ -118,8 +122,12 @@ function _component_egrad(::DefaultJoinEmbedding, M::Manifolds.Segre, p, residua 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(component::JoinComponent, p, residual) = _component_egrad( + _component_embedding(component), + _component_manifold(component), + p, + residual, +) _component_egrad(M, p, residual) = _component_egrad(DefaultJoinEmbedding(), M, p, residual) _manifold_egrad(M, p, residual) = _component_egrad(M, p, residual) @@ -189,7 +197,32 @@ 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) +ambient_length(component::JoinComponent) = ambient_length(component_manifold(component)) + +function component_basis_vector( + component, + p, + coeffs; + basis = ManifoldsBase.DefaultOrthonormalBasis(), +) + return ManifoldsBase.get_vector(component_manifold(component), p, coeffs, basis) +end + +function component_basis_vector( + component, + p, + j::Integer; + basis = ManifoldsBase.DefaultOrthonormalBasis(), +) + T = _scalar_eltype(p) + d = component_tangent_dimension(component, p) + 1 <= j <= d || throw( + BoundsError("Component basis index $j is out of bounds for tangent dimension $d."), + ) + coeffs = zeros(T, d) + coeffs[j] = one(T) + return component_basis_vector(component, p, coeffs; basis) +end """ _join_vector_workspace_like(target, n) returns AbstractVector @@ -251,10 +284,11 @@ function _sum_backend_parts( target::AbstractArray{T,N}; init_point = nothing, ) where {T<:AbstractFloat,N} - r = length(components) + components_tuple = _as_join_component_tuple(components) + r = length(components_tuple) # Keep the original target representation instead of eagerly materializing Array. tgt = target - _validate_join_ambient_compatibility(components, tgt) + _validate_join_ambient_compatibility(components_tuple, tgt) tflat = vec(tgt) tgt_len = length(tgt) @@ -262,10 +296,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) + manifolds = ntuple(k -> _component_manifold(components_tuple[k]), r) return (; - components, + components = components_tuple, manifolds, r, target = tgt, @@ -310,6 +344,7 @@ function _sum_backend_instance( ) where {T<:AbstractFloat,N} parts = _sum_backend_parts(components, target; init_point) return BTDBackend( + parts.components, parts.manifolds, parts.r, parts.target, @@ -609,15 +644,24 @@ function _component_ambient_embedding!( end function _component_ambient_embedding!(out::AbstractVector, component::JoinComponent, p) - return _component_ambient_embedding!(out, component.embedding, component.manifold, p) + return _component_ambient_embedding!( + out, + _component_embedding(component), + _component_manifold(component), + p, + ) end +component_ambient_embedding!(out::AbstractVector, component, p) = + _component_ambient_embedding!(out, component, p) """ _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. +generic `JoinModel`s. Implementations must overwrite `out` completely rather +than accumulating into it, since LM Jacobian assembly reuses the same work +buffer across columns. """ function _component_ambient_pushforward!(out::AbstractVector, M, p, X) return _component_ambient_pushforward!(out, DefaultJoinEmbedding(), M, p, X) @@ -663,13 +707,15 @@ function _component_ambient_pushforward!( ) return _component_ambient_pushforward!( out, - component.embedding, - component.manifold, + _component_embedding(component), + _component_manifold(component), p, X, ) end +component_ambient_pushforward!(out::AbstractVector, component, p, X) = + _component_ambient_pushforward!(out, component, p, X) function _subtract_ambient_tensor!( residual::AbstractArray{T,N}, component, @@ -744,3 +790,51 @@ function _join_residual!(backend::Union{JoinBackend,BTDBackend}, p) backend.work_residual .-= backend.target_flat return backend.work_residual end + +function residual(model::JoinModel{<:AbstractFloat,<:Union{JoinBackend,BTDBackend}}, p) + return copy(_join_residual!(model.backend, p)) +end + +function differential_action!( + out::AbstractVector{T}, + model::JoinModel{<:AbstractFloat,<:Union{JoinBackend,BTDBackend}}, + p, + X, +) where {T<:AbstractFloat} + backend = model.backend + parts = point_parts(p) + xparts = point_parts(X) + _check_parts_len(parts, backend.r, "differential_action!") + _check_parts_len(xparts, backend.r, "differential_action!") + length(out) == length(backend.target_flat) || throw( + DimensionMismatch( + "differential_action! output length $(length(out)) != ambient length $(length(backend.target_flat)).", + ), + ) + fill!(out, zero(T)) + @inbounds for k = 1:backend.r + component_ambient_pushforward!( + backend.component_bufs[k], + _backend_component(backend, k), + parts[k], + xparts[k], + ) + out .+= backend.component_bufs[k] + end + return out +end + +function adjoint_action( + model::JoinModel{<:AbstractFloat,<:Union{JoinBackend,BTDBackend}}, + p, + a::AbstractVector; + kwargs..., +) + backend = model.backend + length(a) == length(backend.target_flat) || throw( + DimensionMismatch( + "adjoint_action expected ambient vector of length $(length(backend.target_flat)), got $(length(a)).", + ), + ) + return _join_basis_project(_backend_components(backend), p, a) +end diff --git a/src/join/join_model.jl b/src/join/join_model.jl index 1b92e2c..1a42c87 100644 --- a/src/join/join_model.jl +++ b/src/join/join_model.jl @@ -1,6 +1,17 @@ # join/join_model.jl — Join-front-end model and backend type definitions -export AbstractJoinBackend, JoinComponent, JoinModel, CPDBackend, JoinBackend, BTDBackend +export AbstractJoinBackend, + JoinComponent, + JoinModel, + CPDBackend, + JoinBackend, + BTDBackend, + component_manifold, + component_embedding, + component_tangent_dimension, + component_basis_vector, + component_ambient_embedding!, + component_ambient_pushforward! # BTDBackend is defined in `btd/model.jl` (includes contraction workspace). abstract type AbstractJoinBackend end @@ -15,7 +26,18 @@ struct DefaultJoinEmbedding end JoinComponent(manifold::M) where {M} = JoinComponent{M,DefaultJoinEmbedding}(manifold, DefaultJoinEmbedding()) -manifold(component::JoinComponent) = component.manifold +component_manifold(component::JoinComponent) = component.manifold +component_embedding(component::JoinComponent) = component.embedding +component_tangent_dimension(component::JoinComponent) = + manifold_dimension(component_manifold(component)) +component_tangent_dimension(component::JoinComponent, p) = + component_tangent_dimension(component) +component_manifold(M::AbstractManifold) = M +component_embedding(::AbstractManifold) = DefaultJoinEmbedding() +component_tangent_dimension(M::AbstractManifold) = manifold_dimension(M) +component_tangent_dimension(M::AbstractManifold, p) = component_tangent_dimension(M) + +manifold(component::JoinComponent) = component_manifold(component) struct JoinModel{T<:AbstractFloat,B<:AbstractJoinBackend} <: AbstractDecompositionModel{T} backend::B diff --git a/src/solvers/abstract.jl b/src/solvers/abstract.jl index bab4b9e..e4c50b2 100644 --- a/src/solvers/abstract.jl +++ b/src/solvers/abstract.jl @@ -31,21 +31,6 @@ end StandardSolverTrace{T}() where {T<:AbstractFloat} = StandardSolverTrace{T}(IterationRecord{T}[], 0, 0, 0.0) -""" - _dual_stop_grad_tol(T, tol; grad_tol=nothing) - -Gradient tolerance paired with `StopWhenCostRelChangeAndGradientLess`. -Defaults to `sqrt(tol)`; callers may pass an explicit `grad_tol` (for example -`grad_tol = tol` on the nonnegative CPD manifold route). -""" -@inline function _dual_stop_grad_tol( - ::Type{T}, - tol::Real, - grad_tol = nothing, -) where {T<:Real} - return isnothing(grad_tol) ? sqrt(T(tol)) : T(grad_tol) -end - function record!( trace::StandardSolverTrace{T}, iter::Int, @@ -411,35 +396,6 @@ function _model_gradient_closure( return model_exact_join_basis_function(model) end -@inline _unwrap_solver_manifold(M) = hasproperty(M, :M) ? getproperty(M, :M) : M - -# The actual methods depend on the registered defaults, e.g. custom manifolds such -# as Segre or SoftplusEuclidean may choose ExponentialRetraction, while sphere-like -# factors may choose their ManifoldsBase default. -@inline function _default_component_retraction_method(Mi, pi) - return ManifoldsBase.default_retraction_method(Mi, typeof(pi)) -end - -function _solver_retraction_method(M, p) - return _solver_retraction_method_unwrapped(_unwrap_solver_manifold(M), p) -end - -function _solver_retraction_method_unwrapped(M::ProductManifold, p) - pparts0 = point_parts(p) - pparts = pparts0 isa Tuple ? pparts0 : Tuple(pparts0) - n = length(M.manifolds) - length(pparts) == n || throw( - ArgumentError( - "Cannot derive solver retraction method: ProductManifold has $n factors but point has $(length(pparts)) parts.", - ), - ) - methods = - ntuple(i -> _default_component_retraction_method(M.manifolds[i], pparts[i]), n) - return ManifoldsBase.ProductRetraction(methods) -end - -_solver_retraction_method_unwrapped(M, p) = _default_component_retraction_method(M, p) - """ _prepare_solver_problem(model; init, gradient_mode) diff --git a/src/solvers/btd_als.jl b/src/solvers/btd_als.jl index 2852167..1c68a28 100644 --- a/src/solvers/btd_als.jl +++ b/src/solvers/btd_als.jl @@ -2,7 +2,7 @@ export fit_btd_als @inline function _btd_block_ranks(backend::BTDBackend, b::Int) - M = backend.manifolds[b] + M = _backend_manifold(backend, b) M isa Manifolds.Tucker || throw( ArgumentError( "BTD ALSSolver expects Tucker manifolds, got $(typeof(M)) at block $b.", diff --git a/src/solvers/btd_tsd.jl b/src/solvers/btd_tsd.jl index 3f12788..7fef7d2 100644 --- a/src/solvers/btd_tsd.jl +++ b/src/solvers/btd_tsd.jl @@ -116,7 +116,7 @@ end _check_parts_len(parts, backend.r, "BTD block descent direction") pk = parts[b] eg_b = _btd_block_egrad(backend, parts, b) - rg_b = egrad_to_rgrad(backend.manifolds[b], pk, eg_b) + rg_b = egrad_to_rgrad(_backend_manifold(backend, b), pk, eg_b) decrease = _btd_tangent_dot(eg_b, rg_b) return pk, rg_b, decrease end @@ -135,7 +135,7 @@ function _btd_tsd_block_step( return p, c0, zero(T), 0, false end - Mk = backend.manifolds[b] + Mk = _backend_manifold(backend, b) α = T(solver.stepsize) α_min = T(solver.armijo_alpha_min) contraction = T(solver.armijo_contraction) diff --git a/src/solvers/lbfgs.jl b/src/solvers/lbfgs.jl index 297b8db..ca17557 100644 --- a/src/solvers/lbfgs.jl +++ b/src/solvers/lbfgs.jl @@ -3,7 +3,7 @@ export LBFGSSolver """ LBFGSSolver(; memory_size=1, cautious_update=true, initial_scale=1.0, - linesearch=:wolfe, preconditioner=nothing) + nonpositive_curvature_behavior=:ignore, linesearch=:wolfe, preconditioner=nothing) Limited-memory Riemannian BFGS wrapper built on `Manopt.quasi_Newton`. """ @@ -11,6 +11,7 @@ struct LBFGSSolver <: AbstractSecondOrderROSolver memory_size::Int cautious_update::Bool initial_scale::Float64 + nonpositive_curvature_behavior::Symbol linesearch::Symbol preconditioner::Any end @@ -19,6 +20,7 @@ function LBFGSSolver(; memory_size::Int = 1, cautious_update::Bool = true, initial_scale::Real = 1.0, + nonpositive_curvature_behavior::Symbol = :ignore, linesearch::Symbol = :wolfe, preconditioner = nothing, ) @@ -35,6 +37,7 @@ function LBFGSSolver(; memory_size, cautious_update, Float64(initial_scale), + nonpositive_curvature_behavior, linesearch, preconditioner, ) @@ -80,6 +83,7 @@ function solve_lbfgs( memory_size::Int = 1, cautious_update::Bool = true, initial_scale::Real = 1.0, + nonpositive_curvature_behavior::Symbol = :ignore, linesearch::Symbol = :wolfe, preconditioner = nothing, grad_tol = nothing, @@ -163,8 +167,10 @@ function solve_lbfgs( memory_size = memory_size, cautious_update = cautious_update, initial_scale = initial_scale, + nonpositive_curvature_behavior = nonpositive_curvature_behavior, linesearch = linesearch, has_preconditioner = !isnothing(preconditioner), + uses_nonpositive_curvature_behavior = false, ), ) end @@ -202,6 +208,7 @@ function run_second_order_solver( memory_size = solver.memory_size, cautious_update = solver.cautious_update, initial_scale = solver.initial_scale, + nonpositive_curvature_behavior = solver.nonpositive_curvature_behavior, linesearch = solver.linesearch, preconditioner = solver.preconditioner, normalized_objective, diff --git a/src/solvers/lm.jl b/src/solvers/lm.jl index d0f6d34..6f6df41 100644 --- a/src/solvers/lm.jl +++ b/src/solvers/lm.jl @@ -36,39 +36,10 @@ solver_symbol(::LMSolver) = :lm one(T) / sqrt(T(normA2)) : one(T) 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))) - components = _backend_components(backend) - @inbounds for k = 1:backend.r - _component_ambient_pushforward!( - backend.component_bufs[k], - components[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 +_lm_raw_residual_vector(model::AbstractDecompositionModel, p) = residual(model, p) function _lm_raw_jacobian_matrix( - model::JoinModel{<:AbstractFloat,<:Union{JoinBackend,BTDBackend}}, + model::AbstractDecompositionModel, M, p; basis = ManifoldsBase.DefaultOrthonormalBasis(), @@ -78,200 +49,68 @@ function _lm_raw_jacobian_matrix( 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, coeff, basis) - _join_tangent_ambient_vector!(column, model.backend, p, 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} - 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, 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, 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, Xj))) - end + differential_action!(column, model, p, Xj) 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(), +function _lm_residual_function( + model::AbstractDecompositionModel, + ::Type{T}, + normA2, + normalized_objective::Bool, ) where {T<:AbstractFloat} - 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) - @inbounds for j = 1:d - fill!(coeff, zero(T)) - coeff[j] = one(T) - Xj = ManifoldsBase.get_vector(M, p, 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, model.dims, model.r) : - unpack_rankr_canonical(p, 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, 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)).")) + scale = _lm_scaling_factor(T, normA2, normalized_objective) + return (M, p) -> scale .* _lm_raw_residual_vector(model, p) end -function _lm_raw_jacobian_matrix(model::AbstractDecompositionModel, M, p; basis) - throw(ArgumentError("LMSolver Jacobian is not implemented for model $(typeof(model)).")) +function _lm_differential_action_function( + model::AbstractDecompositionModel, + ::Type{T}, + normA2, + normalized_objective::Bool, +) where {T<:AbstractFloat} + scale = _lm_scaling_factor(T, normA2, normalized_objective) + return (M, p, X) -> scale .* differential_action(model, p, X) end -function _lm_residual_function( +function _lm_adjoint_action_function( model::AbstractDecompositionModel, ::Type{T}, normA2, normalized_objective::Bool, ) where {T<:AbstractFloat} scale = _lm_scaling_factor(T, normA2, normalized_objective) - return (M, p) -> scale .* _lm_raw_residual_vector(model, p) + return (M, p, a) -> adjoint_action(model, p, scale .* a) end -function _lm_jacobian_function( +function _lm_vector_differential_function( model::AbstractDecompositionModel, ::Type{T}, normA2, - normalized_objective::Bool; - basis = ManifoldsBase.DefaultOrthonormalBasis(), + normalized_objective::Bool, ) where {T<:AbstractFloat} - scale = _lm_scaling_factor(T, normA2, normalized_objective) - return (M, p) -> scale .* _lm_raw_jacobian_matrix(model, M, p; basis) + ambient_dim = length(tensor(model)) + residual_f = _lm_residual_function(model, T, normA2, normalized_objective) + differential_f = + _lm_differential_action_function(model, T, normA2, normalized_objective) + adjoint_f = _lm_adjoint_action_function(model, T, normA2, normalized_objective) + return Manopt.VectorDifferentialFunction( + residual_f, + differential_f, + adjoint_f, + ambient_dim; + evaluation = Manopt.AllocatingEvaluation(), + function_type = Manopt.FunctionVectorialType(), + jacobian_type = Manopt.FunctionVectorialType(), + adjoint_jacobian_type = Manopt.FunctionVectorialType(), + ) end function solve_lm( @@ -311,11 +150,32 @@ function solve_lm( ) 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) - initial_residual_values = residual(M, p0_local) - initial_jacobian_f = jacobian(M, p0_local) + vdf = _lm_vector_differential_function(model, T, normA2, setup.uses_relative_objective) + initial_residual_values = residual(model, p0_local) + scale = _lm_scaling_factor(T, normA2, setup.uses_relative_objective) + if scale != one(T) + initial_residual_values .*= scale + end + nlso = Manopt.ManifoldNonlinearLeastSquaresObjective( + vdf, + Manopt.ComponentwiseRobustifierFunction(Manopt.IdentityRobustifier()), + ) + initial_jacobian_matrices = fill(nothing, 1) + sub_objective = Manopt.construct_lm_subobjective( + false, + nlso, + damping_term_min, + 1.0e-6, + :Strict, + initial_residual_values, + initial_jacobian_matrices, + ) + sub_state = Manopt.ConjugateResidualState( + TangentSpace(M, p0_local), + sub_objective; + stopping_criterion = StopAfterIteration(max(20 * manifold_dimension(M), 200)) | + StopWhenGradientNormLess(T(1e-16)), + ) retraction_method = _solver_retraction_method(M, p0_local) stopping = StopWhenAny( StopAfterIteration(maxiter), @@ -342,21 +202,21 @@ function solve_lm( ) state = Manopt.LevenbergMarquardt( M, - residual, - jacobian, + nlso, p0_local; - evaluation = Manopt.AllocatingEvaluation(), - function_type = Manopt.FunctionVectorialType(), - jacobian_type = Manopt.CoordinateVectorialType(basis), retraction_method = retraction_method, stopping_criterion = stopping, initial_residual_values = initial_residual_values, - initial_jacobian_f = initial_jacobian_f, - η = η, + candidate_acceptance_threshold = η, + damping_increase_factor = β, + damping_increase_threshold = η, + damping_reduction_threshold = expect_zero_residual ? η : Inf, + damping_reduction_factor = inv(T(β)), damping_term_min = damping_term_min, - β = β, - expect_zero_residual = expect_zero_residual, - linear_subsolver! = linear_subsolver, + initial_damping_term = damping_term_min, + use_unified_basis = false, + sub_objective = sub_objective, + sub_state = sub_state, debug = callbacks.debug_actions, return_state = true, ) @@ -382,6 +242,10 @@ function solve_lm( damping_term_min = Float64(damping_term_min), β = Float64(β), expect_zero_residual = expect_zero_residual, + uses_operator_jacobian = true, + uses_direct_adjoint_action = true, + uses_coordinate_linear_solver = false, + uses_user_linear_subsolver = linear_subsolver !== Manopt.default_lm_lin_solve!, uses_vector_transport = !isnothing(vector_transport_method), ), ) diff --git a/src/solvers/manopt_helpers.jl b/src/solvers/manopt_helpers.jl index e6bc167..c6a849a 100644 --- a/src/solvers/manopt_helpers.jl +++ b/src/solvers/manopt_helpers.jl @@ -13,6 +13,21 @@ Base.unsafe_write(::_SolverDebugSink, ::Ptr{UInt8}, n::UInt) = Int(n) # Shared no-op IO used by Manopt debug groups. const _SOLVER_DEBUG_SINK = _SolverDebugSink() +# Gradient tolerance paired with StopWhenCostRelChangeAndGradientLess. +""" + _dual_stop_grad_tol(T, tol; grad_tol=nothing) + +Gradient tolerance paired with `StopWhenCostRelChangeAndGradientLess`. +Defaults to `sqrt(tol)`; callers may pass an explicit `grad_tol` (for example +`grad_tol = tol` on the nonnegative CPD manifold route). +""" +@inline function _dual_stop_grad_tol( + ::Type{T}, + tol::Real, + grad_tol = nothing, +) where {T<:Real} + return isnothing(grad_tol) ? sqrt(T(tol)) : T(grad_tol) +end # Stop when both the relative cost change and Riemannian gradient norm are small. mutable struct StopWhenCostRelChangeAndGradientLess{T<:Real} <: Manopt.StoppingCriterion tol_cost::T @@ -136,6 +151,62 @@ function _solver_point(M, p0) end +# Unwrap solver manifold wrappers down to the underlying manifold object. +@inline _unwrap_solver_manifold(M) = hasproperty(M, :M) ? getproperty(M, :M) : M + + +# The actual methods depend on the registered defaults, e.g. custom manifolds such +# as Segre or SoftplusEuclidean may choose ExponentialRetraction, while sphere-like +# factors may choose their ManifoldsBase default. +@inline function _default_component_retraction_method(Mi, pi) + return ManifoldsBase.default_retraction_method(Mi, typeof(pi)) +end + + +# Choose a solver retraction method, including per-factor product retractions. +function _solver_retraction_method(M, p) + return _solver_retraction_method_unwrapped(_unwrap_solver_manifold(M), p) +end + +function _solver_retraction_method_unwrapped(M::ProductManifold, p) + pparts0 = point_parts(p) + pparts = pparts0 isa Tuple ? pparts0 : Tuple(pparts0) + n = length(M.manifolds) + length(pparts) == n || throw( + ArgumentError( + "Cannot derive solver retraction method: ProductManifold has $n factors but point has $(length(pparts)) parts.", + ), + ) + methods = + ntuple(i -> _default_component_retraction_method(M.manifolds[i], pparts[i]), n) + return ManifoldsBase.ProductRetraction(methods) +end + +_solver_retraction_method_unwrapped(M, p) = _default_component_retraction_method(M, p) + + +# Conservative compatibility probe for vector transports used by Manopt solvers. +function _supports_vector_transport_to(M, p, vt, retraction_method) + try + X = zero_vector(M, p) + q = retract(M, p, X, retraction_method) + Y = vector_transport_to(M, p, X, q, vt) + return isnothing(check_vector(M, q, Y)) + catch + return false + end +end + + +# Choose a vector transport method that works with the current manifold/point layout. +function _default_vector_transport_method(M, p, retraction_method) + vt = ManifoldsBase.ProjectionTransport() + if _supports_vector_transport_to(M, p, vt, retraction_method) + return vt + end + + return ManifoldsBase.default_vector_transport_method(M, typeof(p)) +end # Detect pullback nonnegative geometries that need conservative line search. function _contains_sqeuclidean_manifold(M) M2 = _unwrap_solver_manifold(M) diff --git a/src/solvers/rcg.jl b/src/solvers/rcg.jl index f58f039..0bac9d7 100644 --- a/src/solvers/rcg.jl +++ b/src/solvers/rcg.jl @@ -1,42 +1,6 @@ # solvers/rcg.jl — Riemannian Conjugate Gradient export RCGSolver -# Vector transport selection -""" - _supports_vector_transport_to(M, p, vt, retraction_method) - -Return `true` if `vt` can transport a zero tangent vector from `p` to the -corresponding retracted point and the result is accepted as a tangent vector. -This is a conservative compatibility probe for Manifolds.jl / ManifoldsBase -vector transports. -""" -function _supports_vector_transport_to(M, p, vt, retraction_method) - try - X = zero_vector(M, p) - q = retract(M, p, X, retraction_method) - Y = vector_transport_to(M, p, X, q, vt) - return isnothing(check_vector(M, q, Y)) - catch - return false - end -end - -""" - _default_vector_transport_method(M, p, retraction_method) - -Return the default vector transport method for the given manifold and point. -If the manifold and point layout support it, use `ProjectionTransport()`. -Otherwise, use the manifold's default vector transport method. -""" -function _default_vector_transport_method(M, p, retraction_method) - vt = ManifoldsBase.ProjectionTransport() - if _supports_vector_transport_to(M, p, vt, retraction_method) - return vt - end - - return ManifoldsBase.default_vector_transport_method(M, typeof(p)) -end - # RCG coefficient and restart rule selection function _rcg_coefficient_rule( M, @@ -122,12 +86,10 @@ function solve_rcg( grad_tol, normalized_objective, ) - # Get the initial point and the tangent space type p0_local = setup.p0 T = setup.T retraction_method = _solver_retraction_method(M, p0_local) - transport = isnothing(vector_transport_method) ? _default_vector_transport_method(M, p0_local, retraction_method) : @@ -174,7 +136,7 @@ function solve_rcg( ) return _manopt_finish_result( - get_solver_result(state), + _tk_get_solver_result(state), state, callbacks.progress, diagnostics_recorder, @@ -203,30 +165,10 @@ function solve_rcg( ) end -# RCGSolver object """ RCGSolver(; coefficient=:hager_zhang, restart=:non_descent, ...) - Riemannian conjugate gradient solver. - -Useful options: - -- `coefficient = :hager_zhang` -- `coefficient = :polak_ribiere` -- `coefficient = :fletcher_reeves` -- `coefficient = :dai_yuan` -- `coefficient = :hestenes_stiefel` -- `coefficient = :conjugate_descent` -- `coefficient = :steepest` - -Restart options: - -- `restart = :non_descent` -- `restart = :non_sufficient_descent` -- `restart = :never` - -The default is chosen for CPD swamp experiments: -RCGSolver(; coefficient=:hager_zhang, restart=:non_descent) +Riemannian conjugate gradient solver. """ Base.@kwdef struct RCGSolver <: AbstractFirstOrderROSolver coefficient::Symbol = :hager_zhang diff --git a/src/solvers/rgd.jl b/src/solvers/rgd.jl index cd45cbd..732db5a 100644 --- a/src/solvers/rgd.jl +++ b/src/solvers/rgd.jl @@ -38,7 +38,8 @@ function solve_rgd( p0_local = setup.p0 T = setup.T retraction_method = _solver_retraction_method(M, p0_local) - armijo_alpha_min_T = T(armijo_alpha_min) + stepsize_eff_base = T(stepsize) * setup.objective_scale + armijo_alpha_min_T = T(armijo_alpha_min) * setup.objective_scale tol_g = setup.dual_grad_tol dual_stop = StopWhenCostRelChangeAndGradientLess(T(tol), tol_g) stopping = _manopt_stopping( @@ -57,9 +58,9 @@ function solve_rgd( p0_local, setup.solver_grad, retraction_method, - T(stepsize); + stepsize_eff_base; alpha_min = armijo_alpha_min_T, - ) : T(stepsize) + ) : stepsize_eff_base armijo_contraction = use_squaring_armijo ? T(0.5) : T(0.85) armijo_sufficient_decrease = use_squaring_armijo ? T(1e-4) : T(1e-3) armijo_stop_decreasing = @@ -109,7 +110,7 @@ function solve_rgd( ) return _manopt_finish_result( - get_solver_result(state), + _tk_get_solver_result(state), state, callbacks.progress, diagnostics_recorder, @@ -185,7 +186,7 @@ function solve_rgd_fixed( setup.solver_grad, p0_local; retraction_method = retraction_method, - stepsize = Manopt.ConstantStepsize(M, T(stepsize)), + stepsize = Manopt.ConstantStepsize(M, T(stepsize) * setup.objective_scale), stopping_criterion = stopping, debug = callbacks.debug_actions, count = [:Cost, :Gradient], @@ -193,7 +194,7 @@ function solve_rgd_fixed( ) return _manopt_finish_result( - get_solver_result(state), + _tk_get_solver_result(state), state, callbacks.progress, diagnostics_recorder, @@ -211,13 +212,10 @@ function solve_rgd_fixed( ) end -# ========== RGDSolver (AbstractFirstOrderSolver) ========== - """ RGDSolver(stepsize=1.0; armijo_alpha_min=1e-8) -Riemannian gradient descent with Armijo backtracking line search. Call via -`solve(RGDSolver(...), model; init=:random, gradient_mode=:riemannian)`. +Riemannian gradient descent with Armijo backtracking line search. """ struct RGDSolver <: AbstractFirstOrderSolver stepsize::Float64 @@ -270,14 +268,10 @@ function run_first_order_solver( ) end -# ========== RGDFixedSolver (AbstractFirstOrderSolver) ========== - """ RGDFixedSolver(stepsize=1.0) Riemannian gradient descent with a constant stepsize. -Used as the stable fallback for Tucker-product BTD on dependency stacks where -Armijo's rand/allocate_result path is still unreliable. """ struct RGDFixedSolver <: AbstractFirstOrderSolver stepsize::Float64 diff --git a/test/basic_tests.jl b/test/basic_tests.jl index ab7913b..00056e6 100644 --- a/test/basic_tests.jl +++ b/test/basic_tests.jl @@ -202,6 +202,10 @@ end 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 + @test TensorKitchen.component_manifold(model_component.backend.components[1]) == + segres[1] + @test TensorKitchen.component_embedding(model_component.backend.components[1]) isa + TensorKitchen.DefaultJoinEmbedding p = TensorKitchen.initial_point(model, :random; verbose = false) @test length(TensorKitchen.point_parts(p)) == 3 @@ -255,6 +259,73 @@ end end end +@testset "Operator interface matches Jacobian and gradient" begin + A = randn(5, 4, 3) + cases = ( + ( + "generic_join", + JoinModel((Manifolds.Segre((5, 4, 3)), Manifolds.Segre((5, 4, 3))), A), + ), + ("cp_canonical", JoinModel(A, 2; geometry = :canonical)), + ( + "cp_softplus", + JoinModel(abs.(A), 2; geometry = :softplus_metric, nonnegative = true), + ), + ) + for (label, model) in cases + M = TensorKitchen.manifold(model) + p = TensorKitchen._solver_point( + M, + TensorKitchen.initial_point(model, :random; verbose = false), + ) + basis = ManifoldsBase.DefaultOrthonormalBasis() + r = TensorKitchen.residual(model, p) + J = TensorKitchen._lm_raw_jacobian_matrix(model, M, p; basis) + @testset "$label" begin + @test r ≈ TensorKitchen._lm_raw_residual_vector(model, p) + 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) + @test TensorKitchen.differential_action(model, p, Xj) ≈ J[:, j] + end + g_adj = TensorKitchen.adjoint_action(model, p, r; basis) + g_model = TensorKitchen.rgrad(model, p) + @test norm(M, p, g_adj - g_model) ≤ 1e-7 * max(1.0, norm(M, p, g_model)) + end + end +end + +function _reference_join_jacobian_from_product_basis(model, M, p; basis) + backend = model.backend + parts = TensorKitchen.point_parts(p) + T = TensorKitchen._scalar_eltype(p) + ambient_dim = length(TensorKitchen.tensor(model)) + d = manifold_dimension(M) + J = Matrix{T}(undef, ambient_dim, d) + coeff = zeros(T, d) + col = Vector{T}(undef, ambient_dim) + buf = Vector{T}(undef, ambient_dim) + for j = 1:d + fill!(coeff, zero(T)) + coeff[j] = one(T) + Xj = ManifoldsBase.get_vector(M, p, coeff, basis) + xparts = TensorKitchen.point_parts(Xj) + fill!(col, zero(T)) + for k = 1:backend.r + TensorKitchen.component_ambient_pushforward!( + buf, + TensorKitchen._backend_component(backend, k), + parts[k], + xparts[k], + ) + col .+= buf + end + J[:, j] .= col + end + return J +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) @@ -266,6 +337,33 @@ end basis = ManifoldsBase.DefaultOrthonormalBasis() J = TensorKitchen._lm_raw_jacobian_matrix(model, M, p; basis) @test all(isfinite, J) + @test sum( + TensorKitchen.component_tangent_dimension( + TensorKitchen._backend_component(model.backend, k), + TensorKitchen.point_parts(p)[k], + ) for k = 1:model.backend.r + ) == manifold_dimension(M) + J_ref = _reference_join_jacobian_from_product_basis(model, M, p; basis) + @test maximum(abs.(J .- J_ref)) ≤ 1e-12 + + parts = TensorKitchen.point_parts(p) + c1 = TensorKitchen._backend_component(model.backend, 1) + ξ1 = TensorKitchen.component_basis_vector(c1, parts[1], 1; basis) + buf = similar(model.backend.work_rec) + TensorKitchen.component_ambient_pushforward!(buf, c1, parts[1], ξ1) + @test maximum(abs.(buf .- J[:, 1])) ≤ 1e-10 + + c2 = TensorKitchen._backend_component(model.backend, 2) + ξ2 = TensorKitchen.component_basis_vector(c2, parts[2], 1; basis) + TensorKitchen.component_ambient_pushforward!(buf, c2, parts[2], ξ2) + offset2 = TensorKitchen.component_tangent_dimension(c1, parts[1]) + 1 + @test maximum(abs.(buf .- J[:, offset2])) ≤ 1e-10 + + fill!(buf, NaN) + TensorKitchen.component_ambient_pushforward!(buf, c1, parts[1], ξ1) + fresh = similar(buf) + TensorKitchen.component_ambient_pushforward!(fresh, c1, parts[1], ξ1) + @test buf == fresh retraction_method = TensorKitchen._solver_retraction_method(M, p) ϵ = 1e-6 @@ -285,6 +383,109 @@ end end end +@testset "LM CPD rank-r Jacobian matches finite differences across geometries" begin + A = randn(5, 4, 3) + r = 2 + cases = ( + (JoinModel(A, r; geometry = :canonical), 1e-7), + (JoinModel(A, r; geometry = :native), 1e-7), + (JoinModel(abs.(A), r; nonnegative = true, geometry = :squaring_metric), 5e-6), + (JoinModel(abs.(A), r; nonnegative = true, geometry = :softplus_metric), 5e-6), + ) + + for (model, tol_fd) in cases + 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 + 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) + r_plus = TensorKitchen._lm_raw_residual_vector(model, p_plus) + r_minus = TensorKitchen._lm_raw_residual_vector(model, p_minus) + fd = (r_plus .- r_minus) ./ (2 * ϵ) + @test maximum(abs.(fd .- J[:, j])) ≤ tol_fd + end + end +end + +@testset "LM CPD Jacobian finite differences across direct parameterizations" begin + dims = (5, 4, 3) + A = randn(dims...) + r = 2 + cases = ( + ("rank1 native", TensorKitchen.Rank1CPDModel(A), 5e-7), + ("rank1 squared", TensorKitchen.Rank1CPDModel(abs.(A); nonnegative = true), 5e-6), + ( + "rank1 softplus", + TensorKitchen.Rank1CPDModel( + abs.(A); + nonnegative = true, + use_softplus_metric = true, + ), + 5e-6, + ), + ("rankr native", TensorKitchen.RankRCPDModel(A, r; geometry = :native), 5e-7), + ("rankr canonical", TensorKitchen.RankRCPDModel(A, r; geometry = :canonical), 5e-7), + ( + "rankr squared", + TensorKitchen.RankRCPDModel( + abs.(A), + r; + nonnegative = true, + geometry = :squaring_metric, + ), + 5e-6, + ), + ( + "rankr softplus", + TensorKitchen.RankRCPDModel( + abs.(A), + r; + nonnegative = true, + geometry = :softplus_metric, + ), + 5e-6, + ), + ) + + for (label, model, tol_fd) in cases + 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) + @testset "$label" begin + @test all(isfinite, J) + retraction_method = TensorKitchen._solver_retraction_method(M, p) + ϵ = 1e-6 + d = manifold_dimension(M) + r0 = TensorKitchen._lm_raw_residual_vector(model, p) + 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) + r_plus = TensorKitchen._lm_raw_residual_vector(model, p_plus) + fd = (r_plus .- r0) ./ ϵ + @test maximum(abs.(fd .- J[:, j])) ≤ tol_fd + end + end + end +end @testset "LM normalized and unnormalized objectives take the same step" begin A = randn(6, 5, 4) model = JoinModel(A, 2; geometry = :canonical) @@ -319,6 +520,55 @@ end @test isapprox(res_rel.rel_error, res_abs.rel_error; rtol = 1e-10, atol = 1e-10) end +@testset "CP parameterization tangent decode helpers" begin + dims = (3, 2, 2) + r = 2 + λ̃ = [1.5, -0.4] + Ũ = [randn(dims[m], r) for m = 1:length(dims)] + λ̇̃ = randn(r) + U̇̃ = [randn(dims[m], r) for m = 1:length(dims)] + p = TensorKitchen.pack_point_rankr(λ̃, Ũ, r) + X = TensorKitchen.pack_point_rankr(λ̇̃, U̇̃, r) + + λ_sq, U_sq, λ̇_sq, U̇_sq = TensorKitchen._cp_rankr_decode_tangent_factors( + TensorKitchen.SquaredNNCPParam(), + dims, + r, + p, + X, + ) + @test λ_sq ≈ λ̃ .^ 2 + @test all(U_sq[m] ≈ Ũ[m] .^ 2 for m in eachindex(U_sq)) + @test λ̇_sq ≈ 2 .* λ̃ .* λ̇̃ + @test all(U̇_sq[m] ≈ 2 .* Ũ[m] .* U̇̃[m] for m in eachindex(U̇_sq)) + + λ_sp, U_sp, λ̇_sp, U̇_sp = TensorKitchen._cp_rankr_decode_tangent_factors( + TensorKitchen.SoftplusNNCPParam(), + dims, + r, + p, + X, + ) + @test λ_sp ≈ TensorKitchen._softplus_value.(λ̃) + @test all(U_sp[m] ≈ TensorKitchen._softplus_value.(Ũ[m]) for m in eachindex(U_sp)) + @test λ̇_sp ≈ TensorKitchen._softplus_derivative.(λ̃) .* λ̇̃ + @test all( + U̇_sp[m] ≈ TensorKitchen._softplus_derivative.(Ũ[m]) .* U̇̃[m] for m in eachindex(U̇_sp) + ) + + model_sp = TensorKitchen.RankRCPDModel( + randn(dims...), + r; + nonnegative = true, + geometry = :softplus_metric, + ) + q_zero = + CPDPoint(zeros(Float64, r), [zeros(Float64, dims[m], r) for m = 1:length(dims)]) + p_zero = TensorKitchen.pack_cpd_point(model_sp, q_zero) + λ_lat, U_lat = TensorKitchen.unpack_point_rankr(p_zero, dims, r) + @test all(isfinite, λ_lat) + @test all(F -> all(isfinite, F), U_lat) +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) @@ -354,65 +604,67 @@ end @test res_approx.solver == :lm end -@testset "BTD accepts LMSolver on nested Tucker layouts" begin +@testset "BTD exposes LM residual/Jacobian hooks 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) + M = TensorKitchen.manifold(model) + p0_solver = TensorKitchen._solver_point(M, p0) + basis = ManifoldsBase.DefaultOrthonormalBasis() @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, + @test p0_solver isa ArrayPartition + @test TensorKitchen.point_parts(p0_solver)[1] isa Manifolds.TuckerPoint + + residual0 = TensorKitchen._lm_raw_residual_vector(model, p0_solver) + J0 = TensorKitchen._lm_raw_jacobian_matrix(model, M, p0_solver; basis = basis) + @test length(residual0) == length(A) + @test size(J0) == (length(A), manifold_dimension(M)) + @test all(isfinite, residual0) + @test all(isfinite, J0) + + coeff = zeros(Float64, manifold_dimension(M)) + coeff[1] = 1.0 + X = ManifoldsBase.get_vector(M, p0_solver, coeff, basis) + JX = TensorKitchen.differential_action(model, p0_solver, X) + ambient = randn(size(A)) + lhs = dot(JX, vec(ambient)) + rhs = ManifoldsBase.inner( + M, + p0_solver, + X, + TensorKitchen.adjoint_action(model, p0_solver, vec(ambient)), ) - 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 + @test isapprox(lhs, rhs; atol = 1e-8, rtol = 1e-8) +end - res_btd = btd( +@testset "BTD rejects LMSolver until nested Tucker LM support lands" begin + A = randn(7, 6, 5) + ranks = (2, 2, 2) + + @test_throws ArgumentError 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) - res_btd_alswarm_lm = btd( + @test_throws ArgumentError 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, + solver = LMSolver(), 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 # ========================================================================= @@ -2215,10 +2467,44 @@ end manifolds = TensorKitchen._as_join_manifold_tuple(TuckerJoin(size(A), (2, 2, 2), 2)) backend = TensorKitchen._sum_backend_instance(TensorKitchen.BTDBackend, manifolds, A) + @test length(backend.components) == 2 + @test all(c -> c isa TensorKitchen.JoinComponent, backend.components) + @test map(TensorKitchen._component_manifold, backend.components) == manifolds model_btd = JoinModel{Float64,typeof(backend)}(backend) M_btd = TensorKitchen.manifold(model_btd) p_btd = TensorKitchen._solver_point(M_btd, TensorKitchen.initial_point(model_btd, :random)) + basis_btd = ManifoldsBase.DefaultOrthonormalBasis() + J_btd = + TensorKitchen._lm_raw_jacobian_matrix(model_btd, M_btd, p_btd; basis = basis_btd) + @test all(isfinite, J_btd) + @test sum( + TensorKitchen.component_tangent_dimension( + TensorKitchen._backend_component(backend, k), + TensorKitchen.point_parts(p_btd)[k], + ) for k = 1:backend.r + ) == manifold_dimension(M_btd) + J_btd_ref = _reference_join_jacobian_from_product_basis( + model_btd, + M_btd, + p_btd; + basis = basis_btd, + ) + @test maximum(abs.(J_btd .- J_btd_ref)) ≤ 1e-12 + + retraction_method_btd = TensorKitchen._solver_retraction_method(M_btd, p_btd) + residual0_btd = TensorKitchen._lm_raw_residual_vector(model_btd, p_btd) + ϵ_btd = 1e-6 + for j = 1:min(manifold_dimension(M_btd), 2) + coeff = zeros(Float64, manifold_dimension(M_btd)) + coeff[j] = 1.0 + Xj = ManifoldsBase.get_vector(M_btd, p_btd, coeff, basis_btd) + p_plus = ManifoldsBase.retract(M_btd, p_btd, ϵ_btd * Xj, retraction_method_btd) + r_plus = TensorKitchen._lm_raw_residual_vector(model_btd, p_plus) + fd = (r_plus .- residual0_btd) ./ ϵ_btd + @test maximum(abs.(fd .- J_btd[:, j])) ≤ 5e-6 + end + parts_btd = TensorKitchen.point_parts(p_btd) residual_btd = TensorKitchen._join_residual!(backend, p_btd) tangent_dot_btd(a, b) = begin @@ -2230,8 +2516,11 @@ end end for b = 1:backend.r fast_eg = TensorKitchen._btd_block_egrad(backend, parts_btd, b) - residual_eg = - TensorKitchen._tucker_egrad(backend.manifolds[b], parts_btd[b], residual_btd) + residual_eg = TensorKitchen._tucker_egrad( + TensorKitchen._backend_manifold(backend, b), + parts_btd[b], + residual_btd, + ) @test norm(getproperty(fast_eg, :Ċ) - getproperty(residual_eg, :Ċ)) < 1e-10 @test all( norm(F - R) < 1e-10 for @@ -2243,7 +2532,11 @@ end q_btd = TensorKitchen._replace_block_part( p_btd, b, - retract(backend.manifolds[b], parts_btd[b], (-h) * block_grad), + retract( + TensorKitchen._backend_manifold(backend, b), + parts_btd[b], + (-h) * block_grad, + ), ) fd = (TensorKitchen.cost(model_btd, q_btd) - TensorKitchen.cost(model_btd, p_btd)) /