Merge pull request #21 from TensorKitchen/refactor

Refactor dispatch routing and dual stopping criteria

Author: sChoiKr

src/api/approx.jl

@@ -85,7 +85,7 @@ approx(M, target; verbose = false)
 # returns an BTDResult
 ```
 
-* `approx(base, r, target; kwargs...)` : builds a rank-r Segre join and routes to CPD when base isa Manifolds.Segre and BTD when base isa Manifolds.Tucker unless generic
+* `approx(base, r, target; kwargs...)` : builds a rank-r Segre join and routes by the type of `base`: `Manifolds.Segre` uses CPD and `Manifolds.Tucker` uses BTD unless generic
 dispatch is explicitly requested. 
     - base means one manifold template, not yet a full join.
     - `approx(base, r, target; ...)` repeats that same manifold r times to build a join.
@@ -133,31 +133,81 @@ For the generic join path:
 * `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`.
 """
+function _approx_manifold_collection(
+    dispatch::AutoApproxDispatch,
+    manifolds,
+    target::AbstractArray;
+    kwargs...,
+)
+    _all_segre_uniform(manifolds) && return _approx_manifold_collection(
+        CPDApproxDispatch(),
+        manifolds,
+        target;
+        kwargs...,
+    )
+    _all_tucker_uniform(manifolds, size(target)) && return _approx_manifold_collection(
+        BTDApproxDispatch(),
+        manifolds,
+        target;
+        kwargs...,
+    )
+    return _approx_manifold_collection(
+        GenericApproxDispatch(),
+        manifolds,
+        target;
+        kwargs...,
+    )
+end
+
+function _approx_manifold_collection(
+    ::CPDApproxDispatch,
+    manifolds,
+    target::AbstractArray;
+    kwargs...,
+)
+    _all_segre_uniform(manifolds) || throw(
+        ArgumentError(
+            "approx(...; dispatch=:cpd) requires all manifolds to be Manifolds.Segre with identical factor_dims.",
+        ),
+    )
+    return cpd(target, length(manifolds); kwargs...)
+end
+
+function _approx_manifold_collection(
+    ::BTDApproxDispatch,
+    manifolds,
+    target::AbstractArray;
+    kwargs...,
+)
+    _all_tucker_uniform(manifolds, size(target)) || throw(
+        ArgumentError(
+            "approx(...; dispatch=:btd) requires all manifolds to be Manifolds.Tucker with identical factor_dims/multilinear_rank matching the target.",
+        ),
+    )
+    return btd(target, length(manifolds), multilinear_rank(first(manifolds)); kwargs...)
+end
+
+function _approx_manifold_collection(
+    ::GenericApproxDispatch,
+    manifolds,
+    target::AbstractArray;
+    kwargs...,
+)
+    return approx(JoinModel(manifolds, target); kwargs...)
+end
+
 function approx(
     manifolds::Tuple{Vararg{AbstractManifold}},
     target::AbstractArray{T,N};
     dispatch::Symbol = :auto,
     kwargs...,
 ) where {T<:AbstractFloat,N}
-    dispatch = _normalize_approx_dispatch(dispatch)
-    if dispatch == :cpd || (dispatch == :auto && _all_segre_uniform(manifolds))
-        _all_segre_uniform(manifolds) || throw(
-            ArgumentError(
-                "approx(...; dispatch=:cpd) requires all manifolds to be Manifolds.Segre with identical factor_dims.",
-            ),
-        )
-        return cpd(target, length(manifolds); kwargs...)
-    end
-    if dispatch == :btd ||
-       (dispatch == :auto && _all_tucker_uniform(manifolds, size(target)))
-        _all_tucker_uniform(manifolds, size(target)) || throw(
-            ArgumentError(
-                "approx(...; dispatch=:btd) requires all manifolds to be Manifolds.Tucker with identical factor_dims/multilinear_rank matching the target.",
-            ),
-        )
-        return btd(target, length(manifolds), multilinear_rank(first(manifolds)); kwargs...)
-    end
-    return approx(JoinModel(manifolds, target); kwargs...)
+    return _approx_manifold_collection(
+        approx_dispatch(dispatch),
+        manifolds,
+        target;
+        kwargs...,
+    )
 end
 
 function approx(
@@ -166,25 +216,12 @@ function approx(
     dispatch::Symbol = :auto,
     kwargs...,
 ) where {T<:AbstractFloat,N}
-    dispatch = _normalize_approx_dispatch(dispatch)
-    if dispatch == :cpd || (dispatch == :auto && _all_segre_uniform(manifolds))
-        _all_segre_uniform(manifolds) || throw(
-            ArgumentError(
-                "approx(...; dispatch=:cpd) requires all manifolds to be Manifolds.Segre with identical factor_dims.",
-            ),
-        )
-        return cpd(target, length(manifolds); kwargs...)
-    end
-    if dispatch == :btd ||
-       (dispatch == :auto && _all_tucker_uniform(manifolds, size(target)))
-        _all_tucker_uniform(manifolds, size(target)) || throw(
-            ArgumentError(
-                "approx(...; dispatch=:btd) requires all manifolds to be Manifolds.Tucker with identical factor_dims/multilinear_rank matching the target.",
-            ),
-        )
-        return btd(target, length(manifolds), multilinear_rank(first(manifolds)); kwargs...)
-    end
-    return approx(JoinModel(manifolds, target); kwargs...)
+    return _approx_manifold_collection(
+        approx_dispatch(dispatch),
+        manifolds,
+        target;
+        kwargs...,
+    )
 end
 
 """
@@ -199,27 +236,46 @@ function approx(
     dispatch::Symbol = :auto,
     kwargs...,
 ) where {T<:AbstractFloat,N}
-    dispatch = _normalize_approx_dispatch(dispatch)
     mfs = Tuple(M.manifolds)
-    if dispatch == :cpd || (dispatch == :auto && _all_segre_uniform(mfs))
-        _all_segre_uniform(mfs) || throw(
-            ArgumentError(
-                "approx(...; dispatch=:cpd) requires all manifolds to be Manifolds.Segre with identical factor_dims.",
-            ),
-        )
-        return cpd(target, length(mfs); kwargs...)
-    end
-    if dispatch == :btd || (dispatch == :auto && _all_tucker_uniform(mfs, size(target)))
-        _all_tucker_uniform(mfs, size(target)) || throw(
-            ArgumentError(
-                "approx(...; dispatch=:btd) requires all manifolds to be Manifolds.Tucker with identical factor_dims/multilinear_rank matching the target.",
-            ),
-        )
-        return btd(target, length(mfs), multilinear_rank(first(mfs)); kwargs...)
-    end
+    return _approx_product_manifold(approx_dispatch(dispatch), M, mfs, target; kwargs...)
+end
+
+function _approx_product_manifold(
+    dispatch::Union{AutoApproxDispatch,CPDApproxDispatch,BTDApproxDispatch},
+    M::ProductManifold,
+    manifolds,
+    target::AbstractArray;
+    kwargs...,
+)
+    return _approx_manifold_collection(dispatch, manifolds, target; kwargs...)
+end
+
+function _approx_product_manifold(
+    ::GenericApproxDispatch,
+    M::ProductManifold,
+    manifolds,
+    target::AbstractArray;
+    kwargs...,
+)
     return approx(JoinModel(M, target); kwargs...)
 end
 
+_approx_segre_rank(
+    ::Union{AutoApproxDispatch,CPDApproxDispatch},
+    base::Manifolds.Segre,
+    r::Int,
+    target::AbstractArray;
+    kwargs...,
+) = cpd(target, r; kwargs...)
+
+_approx_segre_rank(
+    ::AbstractApproxDispatch,
+    base::Manifolds.Segre,
+    r::Int,
+    target::AbstractArray;
+    kwargs...,
+) = approx(JoinModel(base, r, target); kwargs...)
+
 """
     approx(base::Manifolds.Segre, r, target; dispatch=:auto, kwargs...)
 
@@ -233,11 +289,25 @@ function approx(
     dispatch::Symbol = :auto,
     kwargs...,
 ) where {T<:AbstractFloat,N}
-    dispatch = _normalize_approx_dispatch(dispatch)
-    return dispatch in (:auto, :cpd) ? cpd(target, r; kwargs...) :
-           approx(JoinModel(base, r, target); kwargs...)
+    return _approx_segre_rank(approx_dispatch(dispatch), base, r, target; kwargs...)
 end
 
+_approx_tucker_rank(
+    ::Union{AutoApproxDispatch,BTDApproxDispatch},
+    base::Manifolds.Tucker,
+    r::Int,
+    target::AbstractArray;
+    kwargs...,
+) = btd(target, r, multilinear_rank(base); kwargs...)
+
+_approx_tucker_rank(
+    ::AbstractApproxDispatch,
+    base::Manifolds.Tucker,
+    r::Int,
+    target::AbstractArray;
+    kwargs...,
+) = approx(JoinModel(base, r, target); kwargs...)
+
 """
     approx(base::Manifolds.Tucker, r, target; dispatch=:auto, kwargs...)
 
@@ -251,9 +321,18 @@ function approx(
     dispatch::Symbol = :auto,
     kwargs...,
 ) where {T<:AbstractFloat,N}
-    dispatch = _normalize_approx_dispatch(dispatch)
-    return dispatch in (:auto, :btd) ? btd(target, r, multilinear_rank(base); kwargs...) :
-           approx(JoinModel(base, r, target); kwargs...)
+    return _approx_tucker_rank(approx_dispatch(dispatch), base, r, target; kwargs...)
+end
+
+_reject_generic_rank_dispatch(::AutoApproxDispatch) = nothing
+_reject_generic_rank_dispatch(::GenericApproxDispatch) = nothing
+
+function _reject_generic_rank_dispatch(::CPDApproxDispatch)
+    throw(ArgumentError("approx(...; dispatch=:cpd) requires Manifolds.Segre inputs."))
+end
+
+function _reject_generic_rank_dispatch(::BTDApproxDispatch)
+    throw(ArgumentError("approx(...; dispatch=:btd) requires Manifolds.Tucker inputs."))
 end
 
 """
@@ -269,11 +348,7 @@ function approx(
     dispatch::Symbol = :auto,
     kwargs...,
 ) where {T<:AbstractFloat,N}
-    dispatch = _normalize_approx_dispatch(dispatch)
-    dispatch == :cpd &&
-        throw(ArgumentError("approx(...; dispatch=:cpd) requires Manifolds.Segre inputs."))
-    dispatch == :btd &&
-        throw(ArgumentError("approx(...; dispatch=:btd) requires Manifolds.Tucker inputs."))
+    _reject_generic_rank_dispatch(approx_dispatch(dispatch))
     return approx(JoinModel(base, r, target); kwargs...)
 end
 
@@ -289,11 +364,7 @@ function approx(
     dispatch::Symbol = :auto,
     kwargs...,
 ) where {T<:AbstractFloat,N}
-    dispatch = _normalize_approx_dispatch(dispatch)
-    dispatch == :cpd &&
-        throw(ArgumentError("approx(...; dispatch=:cpd) requires Manifolds.Segre inputs."))
-    dispatch == :btd &&
-        throw(ArgumentError("approx(...; dispatch=:btd) requires Manifolds.Tucker inputs."))
+    _reject_generic_rank_dispatch(approx_dispatch(dispatch))
     return approx(JoinModel(base, target); kwargs...)
 end
 

src/api/btd.jl

@@ -45,40 +45,23 @@ function _polish_btd_with_als(
     )
 end
 
-@inline function _resolve_btd_init(init, solver::Symbol)
+_default_btd_init(::ALSSolver) = BTDHOSVDMultistartInit()
+_default_btd_init(::AbstractSolver) = :alswarm
+
+@inline function _resolve_btd_init(init, solver::AbstractSolver)
     init == :auto || return init
-    return solver == :als ? BTDHOSVDMultistartInit() : :alswarm
+    return _default_btd_init(solver)
 end
 
 @inline _btd_solver_symbol(::ALSSolver) = :als
 @inline _btd_solver_symbol(solver::AbstractSolver) = solver_symbol(solver)
 
-function _btd_effective_init(
-    solver::Symbol,
-    init,
-    warm_steps::Int,
-    warm_init,
-    warm_block_method::Symbol,
-    warm_block_maxiter::Int,
-)
-    if init == :alswarm
-        return BTDALSWarmStartInit(
-            warm_steps;
-            base_init = warm_init,
-            block_method = warm_block_method,
-            block_maxiter = warm_block_maxiter,
-        )
-    elseif solver ∉ (:als,)
-        return BTDALSWarmStartInit(
-            warm_steps;
-            base_init = init,
-            block_method = warm_block_method,
-            block_maxiter = warm_block_maxiter,
-        )
-    end
-    return init
-end
+_btd_uses_warm_start(::AbstractSolver, _) = false
+_btd_uses_warm_start(::ALSSolver, ::BTDALSWarmStartInit) = false
+_btd_uses_warm_start(::AbstractSolver, ::BTDALSWarmStartInit) = true
 
+_btd_should_polish(::ALSSolver, ::Integer) = false
+_btd_should_polish(::AbstractSolver, polish_n::Integer) = polish_n > 0
 
 function _btd_warm_start_result(
     model::JoinModel{T,<:BTDBackend},
@@ -251,7 +234,7 @@ function btd(
 ) where {T<:AbstractFloat,N}
     solver_obj = _solver_object(solver, stepsize; kwargs...)
     solver_sym = _btd_solver_symbol(solver_obj)
-    init_resolved = _resolve_btd_init(init, solver_sym)
+    init_resolved = _resolve_btd_init(init, solver_obj)
     init_eff =
         init_resolved == :alswarm ?
         BTDALSWarmStartInit(
@@ -269,7 +252,7 @@ function btd(
     warm_info = (;)
     short_circuited = Ref(false)
     result = with_phase_progress() do
-        if solver_sym ∉ (:als,) && init_eff isa BTDALSWarmStartInit
+        if _btd_uses_warm_start(solver_obj, init_eff)
             warm = _btd_warm_start_result(
                 model,
                 b,
@@ -320,7 +303,7 @@ function btd(
     else
         btd_als_polish_maxiter
     end
-    if polish_n > 0 && solver_sym ∉ (:als,)
+    if _btd_should_polish(solver_obj, polish_n)
         result = _polish_btd_with_als(
             b,
             result,