Merge pull request #437 from SciML/CJMselectionrework

Adapt selection

Author: AlCap23

docs/src/solvers/common.md

@@ -47,4 +47,16 @@ As we can see above, the use of a [`Basis`](@ref) is optional to invoke the esti
 
 The [`DataDrivenSolution`](@ref) `res` contains a `result` which is the inferred system and a [`Basis`](@ref).
 
+## Model Selection
 
+Most estimation and model inference algorithms require hyperparameters ,e.g., the sparsity controlling penalty, train-test splits. To account for this, the keyword `selector` can be passed to the [`DataDrivenCommonOptions`](@ref). This allows the user to control the selection criteria and returns the **minimum** selector. 
+
+Common choices for `selector` are `rss`, `bic`, `aic`, `aicc`, and `r2`. Given that each subresult of the algorithm extends the `StatsBase` api, we can also use different schemes like:
+
+```julia
+options = DataDrivenCommonOptions(
+    selector = (x)->rss(x) / nobs(x)
+    )
+```
+
+Which results in the mean squared error of the system.

lib/DataDrivenDMD/src/result.jl

@@ -1,4 +1,4 @@
-struct KoopmanResult{K, B, C, Q, P, T, TE} <: AbstractDataDrivenResult
+struct KoopmanResult{K, B, C, Q, P, T} <: AbstractDataDrivenResult
     """Matrix representation of the operator / generator"""
     k::K
     """Matrix representation of the inputs mapping"""
@@ -9,26 +9,54 @@ struct KoopmanResult{K, B, C, Q, P, T, TE} <: AbstractDataDrivenResult
     q::Q
     """Internal matrix used for updating"""
     p::P
-    """L2 norm error of the testing dataset"""
-    testerror::T
-    """L2 norm error of the training dataset"""
-    trainerror::TE
+    # StatsBase results
+    """Residual sum of squares"""
+    rss::T
+    """Loglikelihood"""
+    loglikelihood::T
+    """Nullloglikelihood"""
+    nullloglikelihood::T
+    """Degrees of freedom"""
+    dof::Int
+    """Number of observations"""
+    nobs::Int
+
     """Returncode"""
     retcode::DDReturnCode
-end
 
-is_success(k::KoopmanResult) = getfield(k, :retcode) == DDReturnCode(1)
-l2error(k::KoopmanResult) = is_success(k) ? getfield(k, :testerror) : Inf
+    function KoopmanResult(k_::K, b::B, c::C, q::Q, p::P, X::AbstractMatrix{T},
+                           Y::AbstractMatrix{T}, U::AbstractMatrix) where {K, B, C, Q, P, T}
+        k = Matrix(k_)
+        rss = isempty(b) ? sum(abs2, Y .- c * k * X) : sum(abs2, Y .- c * (k * X .+ b * U))
+        dof = sum(!iszero, k)
+        dof += isempty(b) ? 0 : sum(!iszero, b)
+        nobs = prod(size(Y))
+        ll = -nobs / 2 * log(rss / nobs)
+        nll = -nobs / 2 * log(mean(abs2, Y .- vec(mean(Y, dims = 2))))
 
-function l2error(k::KoopmanResult{<:Any, <:Any, <:Any, <:Any, <:Any, Nothing})
-    is_success(k) ? getfield(k, :traineerror) : Inf
+        new{K, B, C, Q, P, T}(k_, b, c, q, p, rss, ll, nll, dof, nobs, DDReturnCode(1))
+    end
 end
 
+is_success(k::KoopmanResult) = getfield(k, :retcode) == DDReturnCode(1)
+
 get_operator(k::KoopmanResult) = getfield(k, :k)
 get_generator(k::KoopmanResult) = getfield(k, :k)
 
 get_inputmap(k::KoopmanResult) = getfield(k, :b)
 get_outputmap(k::KoopmanResult) = getfield(k, :c)
 
-get_trainerror(k::KoopmanResult) = getfield(k, :trainerror)
-get_testerror(k::KoopmanResult) = getfield(k, :testerror)
+# StatsBase Overload
+StatsBase.coef(x::KoopmanResult) = getfield(x, :k)
+
+StatsBase.rss(x::KoopmanResult) = getfield(x, :rss)
+
+StatsBase.dof(x::KoopmanResult) = getfield(x, :dof)
+
+StatsBase.nobs(x::KoopmanResult) = getfield(x, :nobs)
+
+StatsBase.loglikelihood(x::KoopmanResult) = getfield(x, :loglikelihood)
+
+StatsBase.nullloglikelihood(x::KoopmanResult) = getfield(x, :nullloglikelihood)
+
+StatsBase.r2(x::KoopmanResult) = r2(x, :CoxSnell)

lib/DataDrivenDMD/src/solve.jl

@@ -64,14 +64,14 @@ function CommonSolve.solve!(prob::InternalDataDrivenProblem{A}) where {
                                                                        AbstractKoopmanAlgorithm
                                                                        }
     @unpack alg, basis, testdata, traindata, control_idx, options, problem, kwargs = prob
-
+    @unpack selector = options
     # Check for 
     results = alg(prob; kwargs...)
 
-    # Get the best result based on test error, if applicable else use testerror
-    sort!(results, by = l2error)
+    # Get the best result based on selector
+    idx = argmin(map(selector, results))
+    best_res = results[idx]
     # Convert to basis
-    best_res = first(results)
     new_basis = convert_to_basis(best_res, basis, problem, options, control_idx)
     # Build DataDrivenResult
     DataDrivenSolution(new_basis, problem, alg, results, prob, best_res.retcode)
@@ -93,13 +93,6 @@ function convert_to_basis(res::KoopmanResult, basis::Basis, prob, options, contr
     DataDrivenDiffEq.__construct_basis(Θ, basis, prob, options)
 end
 
-function __compute_rss(Z, C, K, B, X, U)
-    begin
-        (isempty(U) || isempty(B)) && return sum(abs2, Z .- C * (K * X))
-        return sum(abs2, Z .- C * (K * X + B * U))
-    end
-end
-
 function (algorithm::AbstractKoopmanAlgorithm)(prob::InternalDataDrivenProblem;
                                                control_input = nothing, kwargs...)
     @unpack traindata, testdata, control_idx, options = prob
@@ -127,14 +120,6 @@ function (algorithm::AbstractKoopmanAlgorithm)(prob::InternalDataDrivenProblem;
         Q = Y_ * X'
         P = X * X'
         C = Z / Y_
-        trainerror = __compute_rss(Z, C, Matrix(K), B, X_, U_)
-        if !isempty(X̃)
-            testerror = __compute_rss(Z̃, C, Matrix(K), B, X̃, Ũ)
-            retcode = testerror <= abstol ? DDReturnCode(1) : DDReturnCode(5)
-        else
-            testerror = nothing
-            retcode = trainerror <= abstol ? DDReturnCode(1) : DDReturnCode(5)
-        end
-        KoopmanResult(K, B, C, Q, P, testerror, trainerror, retcode)
+        KoopmanResult(K, B, C, Q, P, X_, Z, U_)
     end
 end

lib/DataDrivenDMD/test/linear_autonomous.jl

@@ -31,8 +31,7 @@ rng = StableRNG(42)
                 @test Matrix(get_operator(operator_res)) ≈ A
                 @test isempty(get_inputmap(operator_res))
                 @test get_outputmap(operator_res) ≈ I(2)
-                @test get_trainerror(operator_res) <= 1e-10
-                @test isnothing(get_testerror(operator_res))
+                @test rss(operator_res) <= 1e-10
             end
         end
     end
@@ -52,8 +51,7 @@ rng = StableRNG(42)
                 @test Matrix(get_operator(operator_res))≈A atol=1e-2
                 @test isempty(get_inputmap(operator_res))
                 @test get_outputmap(operator_res) ≈ I(2)
-                @test get_trainerror(operator_res) <= 1e-2
-                @test isnothing(get_testerror(operator_res))
+                @test rss(operator_res) <= 1e-2
             end
         end
     end
@@ -80,8 +78,7 @@ end
                 @test Matrix(get_operator(operator_res)) ≈ A
                 @test isempty(get_inputmap(operator_res))
                 @test get_outputmap(operator_res) ≈ I(2)
-                @test get_trainerror(operator_res) <= 1e-10
-                @test isnothing(get_testerror(operator_res))
+                @test rss(operator_res) <= 1e-10
             end
         end
     end

lib/DataDrivenSR/src/DataDrivenSR.jl

@@ -44,79 +44,107 @@ $(FIELDS)
     eq_options::SymbolicRegression.Options = SymbolicRegression.Options()
 end
 
-struct SRResult{H, P, T, TE} <: AbstractDataDrivenResult
+struct SRResult{H, P, T} <: AbstractDataDrivenResult
+    "The resulting basis"
+    basis::Basis
+    "The Hall of Fame"
     halloffame::H
+    """The Paretofrontier"""
     paretofrontier::P
-    testerror::T
-    trainerror::TE
+    # StatsBase results
+    """Residual sum of squares"""
+    rss::T
+    """Loglikelihood"""
+    loglikelihood::T
+    """Nullloglikelihood"""
+    nullloglikelihood::T
+    """Degrees of freedom"""
+    dof::Int
+    """Number of observations"""
+    nobs::Int
+    """Returncode"""
     retcode::DDReturnCode
 end
 
-is_success(k::SRResult) = getfield(k, :retcode) == DDReturnCode(1)
-l2error(k::SRResult) = is_success(k) ? getfield(k, :testerror) : Inf
-function l2error(k::SRResult{<:Any, <:Any, <:Any, Nothing})
-    is_success(k) ? getfield(k, :traineerror) : Inf
+function SRResult(prob, hof, paretos)
+    @unpack basis, problem = prob
+    bs = convert_to_basis(paretos, prob)
+    ps = get_parameter_values(bs)
+    problem = DataDrivenDiffEq.remake_problem(problem, p = ps)
+    y = DataDrivenDiffEq.get_implicit_data(problem)
+    rss = sum(abs2, y .- bs(problem))
+    dof = length(ps)
+    nobs = prod(size(y))
+    ll = iszero(rss) ? convert(eltype(rss), Inf) : -nobs / 2 * log(rss / nobs)
+    ll0 = -nobs / 2 * log.(sum(abs2, y .- mean(y, dims = 2)[:, 1]) / nobs)
+    return SRResult(bs, hof, paretos,
+                    rss, ll, ll0, dof, nobs,
+                    DDReturnCode(1))
 end
 
-# apply the algorithm on each dataset
-function (x::EQSearch)(ps::InternalDataDrivenProblem{EQSearch}, X, Y)
-    @unpack problem, testdata, options = ps
-    @unpack maxiters, abstol = options
-    @unpack weights, eq_options, numprocs, procs, parallelism, runtests = x
+is_success(k::SRResult) = getfield(k, :retcode) == DDReturnCode(1)
 
-    hofs = SymbolicRegression.EquationSearch(X, Y,
-                                             niterations = maxiters,
-                                             weights = weights,
-                                             options = eq_options,
-                                             numprocs = numprocs,
-                                             procs = procs, parallelism = parallelism,
-                                             runtests = runtests)
+# StatsBase Overload
+StatsBase.coef(x::SRResult) = getfield(x, :k)
 
-    # We always want something which is a vector or tuple
-    hofs = !isa(hofs, AbstractVector) ? [hofs] : hofs
+StatsBase.rss(x::SRResult) = getfield(x, :rss)
 
-    # Evaluate over the full training data
-    paretos = map(enumerate(hofs)) do (i, hof)
-        SymbolicRegression.calculate_pareto_frontier(X, Y[i, :], hof, eq_options)
+StatsBase.dof(x::SRResult) = getfield(x, :dof)
+
+StatsBase.nobs(x::SRResult) = getfield(x, :nobs)
+
+StatsBase.loglikelihood(x::SRResult) = getfield(x, :loglikelihood)
+
+StatsBase.nullloglikelihood(x::SRResult) = getfield(x, :nullloglikelihood)
+
+StatsBase.r2(x::SRResult) = r2(x, :CoxSnell)
+
+function collect_numerical_parameters(eq, options = DataDrivenCommonOptions())
+    ps = Any[]
+    eqs = map(eq) do eqi
+        _collect_numerical_parameters!(ps, eqi, options)
     end
+    return eqs, ps
+end
 
-    # Trainingerror
-    trainerror = mean(x -> x[end].loss, paretos)
-    # Testerror 
-    X̃, Ỹ = testdata
-    if !isempty(X̃)
-        testerror = mean(map(enumerate(hofs)) do (i, hof)
-                             doms = SymbolicRegression.calculate_pareto_frontier(X̃,
-                                                                                 Ỹ[i, :],
-                                                                                 hof,
-                                                                                 eq_options)
-                             doms[end].loss
-                         end)
-        retcode = testerror <= abstol ? DDReturnCode(1) : DDReturnCode(5)
+function _collect_numerical_parameters!(ps::AbstractVector, eq, options)
+    if Symbolics.istree(eq)
+        args_ = map(Symbolics.arguments(eq)) do (eqi)
+            _collect_numerical_parameters!(ps, eqi, options)
+        end
+        return Symbolics.operation(eq)(args_...)
+    elseif isa(eq, Number)
+        pval = round(eq, options.roundingmode, digits = options.digits)
+        # We do not collect zeros or ones
+        iszero(pval) && return zero(eltype(pval))
+        (abs(pval) ≈ 1) & return sign(pval) * one(eltype(pval))
+        p_ = Symbolics.variable(:p, length(ps) + 1)
+        p_ = Symbolics.setdefaultval(p_, pval)
+        p_ = ModelingToolkit.toparam(p_)
+        push!(ps, p_)
+        return p_
     else
-        testerror = nothing
-        retcode = trainerror <= abstol ? DDReturnCode(1) : DDReturnCode(5)
+        return eq
     end
-
-    return SRResult(hofs, paretos, testerror, trainerror, retcode)
 end
 
-function convert_to_basis(res::SRResult, prob)
-    @unpack paretofrontier = res
+function convert_to_basis(paretofrontier, prob)
     @unpack alg, basis, problem, options = prob
     @unpack eq_options = alg
     @unpack maxiters, eval_expresssion, generate_symbolic_parameters, digits, roundingmode = options
 
-    eqs_ = Num.(map(paretofrontier) do dom
-                    node_to_symbolic(dom[end].tree, eq_options)
-                end)
+    eqs_ = map(paretofrontier) do dom
+        node_to_symbolic(dom[end].tree, eq_options)
+    end
 
     # Substitute with the basis elements
     atoms = map(xi -> xi.rhs, equations(basis))
 
     subs = Dict([SymbolicUtils.Sym{LiteralReal}(Symbol("x$(i)")) => x
                  for (i, x) in enumerate(atoms)]...)
-    eqs = map(Base.Fix2(substitute, subs), eqs_)
+
+    eqs, ps = collect_numerical_parameters(eqs_)
+    eqs = map(Base.Fix2(substitute, subs), eqs)
 
     # Get the lhs
     causality, dt = DataDrivenDiffEq.assert_lhs(problem)
@@ -135,40 +163,61 @@ function convert_to_basis(res::SRResult, prob)
         eqs = [phi[i] ~ eq for (i, eq) in enumerate(eqs)]
     end
 
-    ps = parameters(basis)
+    ps_ = parameters(basis)
     @unpack p = problem
 
     p_new = map(eachindex(p)) do i
-        DataDrivenDiffEq._set_default_val(Num(ps[i]), p[i])
+        DataDrivenDiffEq._set_default_val(Num(ps_[i]), p[i])
     end
 
     Basis(eqs, states(basis),
-          parameters = p_new, iv = get_iv(basis),
+          parameters = [p_new; ps], iv = get_iv(basis),
           controls = controls(basis), observed = observed(basis),
           implicits = implicit_variables(basis),
           name = gensym(:Basis),
           eval_expression = eval_expresssion)
 end
 
+# apply the algorithm on each dataset
+function (x::EQSearch)(ps::InternalDataDrivenProblem{EQSearch}, X, Y)
+    @unpack problem, testdata, options = ps
+    @unpack maxiters, abstol = options
+    @unpack weights, eq_options, numprocs, procs, parallelism, runtests = x
+
+    hofs = SymbolicRegression.EquationSearch(X, Y,
+                                             niterations = maxiters,
+                                             weights = weights,
+                                             options = eq_options,
+                                             numprocs = numprocs,
+                                             procs = procs, parallelism = parallelism,
+                                             runtests = runtests)
+
+    # We always want something which is a vector or tuple
+    hofs = !isa(hofs, AbstractVector) ? [hofs] : hofs
+
+    # Evaluate over the full training data
+    paretos = map(enumerate(hofs)) do (i, hof)
+        SymbolicRegression.calculate_pareto_frontier(X, Y[i, :], hof, eq_options)
+    end
+
+    return SRResult(ps, hofs, paretos)
+end
+
 function CommonSolve.solve!(ps::InternalDataDrivenProblem{EQSearch})
     @unpack alg, basis, testdata, traindata, kwargs = ps
     @unpack weights, numprocs, procs, addprocs_function, parallelism, runtests, eq_options = alg
     @unpack traindata, testdata, basis, options = ps
-    @unpack maxiters, eval_expresssion, generate_symbolic_parameters, digits, roundingmode = options
+    @unpack maxiters, eval_expresssion, generate_symbolic_parameters, digits, roundingmode, selector = options
     @unpack problem = ps
 
     results = map(traindata) do (X, Y)
         alg(ps, X, Y)
     end
 
-    # Get the best result based on test error, if applicable else use testerror
-    sort!(results, by = l2error)
-    # Convert to basis
-    best_res = first(results)
-
-    new_basis = convert_to_basis(best_res, ps)
+    idx = argmin(map(selector, results))
+    best_res = results[idx]
 
-    DataDrivenSolution(new_basis, problem, alg, results, ps, best_res.retcode)
+    DataDrivenSolution(best_res.basis, problem, alg, results, ps, best_res.retcode)
 end
 
 export EQSearch