test: more tests are now working

Author: avik-pal

lib/DataDrivenLux/src/DataDrivenLux.jl

@@ -13,6 +13,7 @@ using CommonSolve: CommonSolve, solve!
 using ConcreteStructs: @concrete
 using Setfield: Setfield, @set!
 
+# TODO: Get rid of Optim and Optimisers in favor of Optimization.jl
 using Optim: Optim, LBFGS
 using Optimisers: Optimisers, ADAM
 
@@ -64,17 +65,20 @@ export AdditiveError, MultiplicativeError
 export ObservedModel
 
 # Simplex
-include("./lux/simplex.jl")
+include("lux/simplex.jl")
 export Softmax, GumbelSoftmax, DirectSimplex
 
 # Nodes and Layers
-include("./lux/path_state.jl")
+include("lux/path_state.jl")
 export PathState
-include("./lux/node.jl")
+
+include("lux/node.jl")
 export FunctionNode
-include("./lux/layer.jl")
+
+include("lux/layer.jl")
 export FunctionLayer
-include("./lux/graph.jl")
+
+include("lux/graph.jl")
 export LayeredDAG
 
 include("caches/dataset.jl")

lib/DataDrivenLux/src/caches/cache.jl

@@ -9,10 +9,7 @@ struct SearchCache{ALG, PTYPE, O} <: AbstractAlgorithmCache
     optimiser_state::O
 end
 
-function Base.show(io::IO, cache::SearchCache)
-    print(io, "SearchCache : $(cache.alg)")
-    return
-end
+Base.show(io::IO, cache::SearchCache) = print(io, "SearchCache : $(cache.alg)")
 
 function init_model(x::AbstractDAGSRAlgorithm, basis::Basis, dataset::Dataset, intervals)
     (; simplex, n_layers, arities, functions, use_protected, skip) = x
@@ -116,7 +113,7 @@ function update_cache!(cache::SearchCache)
         sortperm!(cache.sorting, cache.candidates, by = loss)
         permute!(cache.candidates, cache.sorting)
         loss_quantile = quantile(losses, keep, sorted = true)
-        cache.keeps .= (losses .<= loss_quantile)
+        @. cache.keeps = losses ≤ loss_quantile
     end
 
     return
@@ -158,7 +155,6 @@ function optimize_cache!(cache::SearchCache{<:Any, __PROCESSUSE(2)}, p = cache.p
 end
 
 # Distributed
-
 function optimize_cache!(cache::SearchCache{<:Any, __PROCESSUSE(3)}, p = cache.p)
     (; optimizer, optim_options) = cache.alg
 
@@ -176,4 +172,4 @@ function optimize_cache!(cache::SearchCache{<:Any, __PROCESSUSE(3)}, p = cache.p
     return
 end
 
-function convert_to_basis(cache::SearchCache) end
+function convert_to_basis(::SearchCache) end

lib/DataDrivenLux/src/caches/candidate.jl

@@ -22,17 +22,16 @@ StatsBase.nullloglikelihood(stats::PathStatistics) = getfield(stats, :nullloglik
 StatsBase.dof(stats::PathStatistics) = getfield(stats, :dof)
 StatsBase.r2(c::PathStatistics) = r2(c, :CoxSnell)
 
-struct ComponentModel{B, M}
-    basis::B
-    model::M
+@concrete struct ComponentModel
+    basis
+    model
 end
 
-function (c::ComponentModel)(dataset::Dataset{T}, ps, st::NamedTuple{fieldnames},
-        p::AbstractVector{T}) where {T, fieldnames}
+function (c::ComponentModel)(dataset::Dataset{T}, ps, st::NamedTuple,
+        p::AbstractVector{T}) where {T}
     return first(c.model(c.basis(dataset, p), ps, st))
 end
-function (c::ComponentModel)(ps, st::NamedTuple{fieldnames},
-        paths::Vector{<:AbstractPathState}) where {fieldnames}
+function (c::ComponentModel)(ps, st::NamedTuple, paths::Vector{<:AbstractPathState})
     return get_loglikelihood(c.model, ps, st, paths)
 end
 
@@ -45,29 +44,29 @@ to the symbolic regression problem.
 # Fields
 $(FIELDS)
 """
-struct Candidate{S <: NamedTuple} <: StatsBase.StatisticalModel
+@concrete struct Candidate <: StatsBase.StatisticalModel
     "Random seed"
-    rng::AbstractRNG
+    rng <: AbstractRNG
     "The current state"
-    st::S
+    st <: NamedTuple
     "The current parameters"
-    ps::AbstractVector
+    ps <: AbstractVector
     "Incoming paths"
-    incoming_path::Vector{AbstractPathState}
+    incoming_path <: Vector{<:AbstractPathState}
     "Outgoing path"
-    outgoing_path::Vector{AbstractPathState}
+    outgoing_path <: Vector{<:AbstractPathState}
     "Statistics"
-    statistics::PathStatistics
+    statistics <: PathStatistics
     "The observed model"
-    observed::ObservedModel
+    observed <: ObservedModel
     "The parameter distribution"
-    parameterdist::ParameterDistributions
+    parameterdist <: ParameterDistributions
     "The optimal scales"
-    scales::AbstractVector
+    scales <: AbstractVector
     "The optimal parameters"
-    parameters::AbstractVector
+    parameters <: AbstractVector
     "The component model"
-    model::ComponentModel
+    model <: ComponentModel
 end
 
 function (c::Candidate)(dataset::Dataset{T}, ps = c.ps, p = c.parameters) where {T}
@@ -89,12 +88,9 @@ StatsBase.r2(c::Candidate) = r2(c, :CoxSnell)
 get_parameters(c::Candidate) = transform_parameter(c.parameterdist, c.parameters)
 get_scales(c::Candidate) = transform_scales(c.observed, c.scales)
 
-function Candidate(rng, model, basis, dataset; observed = ObservedModel(dataset.y),
-        parameterdist = ParameterDistributions(basis), ptype = Float32)
-    (; y, x) = dataset
-
-    T = eltype(dataset)
-
+function Candidate(
+        rng, model, basis, dataset::Dataset{T}; observed = ObservedModel(dataset.y),
+        parameterdist = ParameterDistributions(basis), ptype = Float32) where {T}
     # Create the initial state and path
     dataset_intervals = interval_eval(basis, dataset, get_interval(parameterdist))
 
@@ -110,21 +106,21 @@ function Candidate(rng, model, basis, dataset; observed = ObservedModel(dataset.
 
     ŷ, _ = model(basis(dataset, transform_parameter(parameterdist, parameters)), ps, st)
 
-    lls = logpdf(observed, y, ŷ, scales)
+    lls = logpdf(observed, dataset.y, ŷ, scales)
     lls += logpdf(parameterdist, parameters)
 
-    rss = sum(abs2, y .- ŷ)
+    rss = sum(abs2, dataset.y .- ŷ)
     dof_ = get_dof(outgoing_path)
 
-    ȳ = vec(mean(y, dims = 2))
+    ȳ = vec(mean(dataset.y; dims = 2))
 
-    null_ll = logpdf(observed, y, ȳ, scales) + logpdf(parameterdist, parameters)
+    null_ll = logpdf(observed, dataset.y, ȳ, scales) + logpdf(parameterdist, parameters)
 
-    stats = PathStatistics(rss, lls, null_ll, dof_, prod(size(y)))
+    stats = PathStatistics(rss, lls, null_ll, dof_, prod(size(dataset.y)))
 
-    return Candidate{typeof(st)}(
-        Lux.replicate(rng), st, ComponentVector(ps), incoming_path, outgoing_path, stats,
-        observed, parameterdist, scales, parameters, ComponentModel(basis, model))
+    return Candidate(Lux.replicate(rng), st, ComponentVector(ps), incoming_path,
+        outgoing_path, stats, observed, parameterdist, scales, parameters,
+        ComponentModel(basis, model))
 end
 
 function update_values!(c::Candidate, ps, dataset)
@@ -136,34 +132,24 @@ function update_values!(c::Candidate, ps, dataset)
     dataloglikelihood = logpdf(observed, y, ŷ, scales) + logpdf(parameterdist, parameters)
     rss = sum(abs2, y .- ŷ)
     dof = get_dof(outgoing_path)
-    ȳ = vec(mean(y, dims = 2))
+    ȳ = vec(mean(y; dims = 2))
     nullloglikelihood = logpdf(observed, y, ȳ, scales) + logpdf(parameterdist, parameters)
     update_stats!(statistics, rss, dataloglikelihood, nullloglikelihood, dof)
     return
 end
 
 @views function Distributions.logpdf(
         c::Candidate, p::ComponentVector, dataset::Dataset{T}, ps = c.ps) where {T}
-    (; observed, parameterdist) = c
-    (; scales, parameters) = p
-    (; y) = dataset
-
-    ŷ = c(dataset, ps, parameters)
-    return logpdf(c, p, y, ŷ)
+    ŷ = c(dataset, ps, p.parameters)
+    return logpdf(c, p, dataset.y, ŷ)
 end
 
 function Distributions.logpdf(c::Candidate, p::AbstractVector, y::AbstractMatrix{T},
         ŷ::AbstractMatrix{T}) where {T}
-    (; scales, parameters) = p
-    (; observed, parameterdist) = c
-
-    return logpdf(observed, y, ŷ, scales) + logpdf(parameterdist, parameters)
+    return logpdf(c.observed, y, ŷ, p.scales) + logpdf(c.parameterdist, p.parameters)
 end
 
-function initial_values(c::Candidate)
-    (; scales, parameters) = c
-    return ComponentVector((; scales = scales, parameters = parameters))
-end
+initial_values(c::Candidate) = ComponentVector(; c.scales, c.parameters)
 
 function optimize_candidate!(
         c::Candidate, dataset::Dataset{T}, ps = c.ps; optimizer = Optim.LBFGS(),
@@ -195,16 +181,10 @@ function optimize_candidate!(
     return
 end
 
-function check_intervals(paths::AbstractArray{<:AbstractPathState})::Bool
-    @inbounds for path in paths
-        check_intervals(path) || return false
-    end
-    return true
-end
+check_intervals(paths::AbstractArray{<:AbstractPathState}) = all(check_intervals, paths)
 
 function sample(c::Candidate, ps, i = 0, max_sample = 10)
-    (; incoming_path, st) = c
-    return sample(c.model.model, incoming_path, ps, st, i, max_sample)
+    return sample(c.model.model, c.incoming_path, ps, c.st, i, max_sample)
 end
 
 function sample(model, incoming, ps, st, i = 0, max_sample = 10)

lib/DataDrivenLux/src/caches/dataset.jl

@@ -1,12 +1,12 @@
-struct Dataset{T}
-    x::AbstractMatrix{T}
-    y::AbstractMatrix{T}
-    u::AbstractMatrix{T}
-    t::AbstractVector{T}
-    x_intervals::AbstractVector{Interval{T}}
-    y_intervals::AbstractVector{Interval{T}}
-    u_intervals::AbstractVector{Interval{T}}
-    t_interval::Interval{T}
+@concrete struct Dataset{T}
+    x <: AbstractMatrix{T}
+    y <: AbstractMatrix{T}
+    u <: AbstractMatrix{T}
+    t <: AbstractVector{T}
+    x_intervals <: AbstractVector{Interval{T}}
+    y_intervals <: AbstractVector{Interval{T}}
+    u_intervals <: AbstractVector{Interval{T}}
+    t_interval <: Interval{T}
 end
 
 Base.eltype(::Dataset{T}) where {T} = T
@@ -20,10 +20,10 @@ function Dataset(X::AbstractMatrix, Y::AbstractMatrix,
     U = convert.(T, U)
     t = convert.(T, t)
     t = isempty(t) ? convert.(T, LinRange(0, size(Y, 2) - 1, size(Y, 2))) : convert.(T, t)
-    x_intervals = Interval.(map(extrema, eachrow(X)))
-    y_intervals = Interval.(map(extrema, eachrow(Y)))
-    u_intervals = Interval.(map(extrema, eachrow(U)))
-    t_intervals = isempty(t) ? Interval{T}(zero(T), zero(T)) : Interval(extrema(t))
+    x_intervals = interval.(map(extrema, eachrow(X)))
+    y_intervals = interval.(map(extrema, eachrow(Y)))
+    u_intervals = interval.(map(extrema, eachrow(U)))
+    t_intervals = isempty(t) ? Interval{T}(zero(T), zero(T)) : interval(extrema(t))
     return Dataset{T}(X, Y, U, t, x_intervals, y_intervals, u_intervals, t_intervals)
 end
 
@@ -35,50 +35,40 @@ end
 
 function (b::Basis{false, false})(d::Dataset{T}, p::P) where {T, P}
     f = DataDrivenDiffEq.get_f(b)
-    (; x, t) = d
-    return f(x, p, t)
+    return f(d.x, p, d.t)
 end
 
 function (b::Basis{false, true})(d::Dataset{T}, p::P) where {T, P}
     f = DataDrivenDiffEq.get_f(b)
-    (; x, t, u) = d
-    return f(x, p, t, u)
+    return f(d.x, p, d.t, d.u)
 end
 
 function (b::Basis{true, false})(d::Dataset{T}, p::P) where {T, P}
     f = DataDrivenDiffEq.get_f(b)
-    (; y, x, t) = d
-    return f(y, x, p, t)
+    return f(d.y, d.x, p, d.t)
 end
 
 function (b::Basis{true, true})(d::Dataset{T}, p::P) where {T, P}
     f = DataDrivenDiffEq.get_f(b)
-    (; y, x, t, u) = d
-    return f(y, x, p, t, u)
+    return f(d.y, d.x, p, d.t, d.u)
 end
 
-##
-
 function interval_eval(b::Basis{false, false}, d::Dataset{T}, p::P) where {T, P}
     f = DataDrivenDiffEq.get_f(b)
-    (; x_intervals, t_interval) = d
-    return f(x_intervals, p, t_interval)
+    return f(d.x_intervals, p, d.t_interval)
 end
 
 function interval_eval(b::Basis{false, true}, d::Dataset{T}, p::P) where {T, P}
     f = DataDrivenDiffEq.get_f(b)
-    (; x_intervals, t_interval, u_intervals) = d
-    return f(x_intervals, p, t_interval, u_intervals)
+    return f(d.x_intervals, p, d.t_interval, d.u_intervals)
 end
 
 function interval_eval(b::Basis{true, false}, d::Dataset{T}, p::P) where {T, P}
     f = DataDrivenDiffEq.get_f(b)
-    (; y_intervals, x_intervals, t_interval) = d
-    return f(y_intervals, x_intervals, p, t_interval)
+    return f(d.y_intervals, d.x_intervals, p, d.t_interval)
 end
 
 function interval_eval(b::Basis{true, true}, d::Dataset{T}, p::P) where {T, P}
     f = DataDrivenDiffEq.get_f(b)
-    (; y_intervals, x_intervals, t_interval, u_intervals) = d
-    return f(y_intervals, x_intervals, p, t_interval, u_intervals)
+    return f(d.y_intervals, d.x_intervals, p, d.t_interval, d.u_intervals)
 end

lib/DataDrivenLux/src/custom_priors.jl

@@ -147,7 +147,7 @@ function ParameterDistribution(
     upper_t = isinf(upper) ? TransformVariables.∞ : upper
     transform = as(Real, lower_t, upper_t)
     init = convert.(T, TransformVariables.inverse(transform, init))
-    return ParameterDistribution(d, Interval(lower, upper), transform, init)
+    return ParameterDistribution(d, interval(lower, upper), transform, init)
 end
 
 function Base.summary(io::IO, p::ParameterDistribution)

lib/DataDrivenLux/src/lux/path_state.jl

@@ -1,12 +1,21 @@
 abstract type AbstractPathState end
 
-@concrete struct PathState{T} <: AbstractPathState
+struct PathState{T, PO <: Tuple, PI <: Tuple} <: AbstractPathState
     "Accumulated loglikelihood of the state"
     path_interval::Interval{T}
     "All the operators of the path"
-    path_operators <: Tuple
+    path_operators::PO
     "The unique identifier of nodes in the path"
-    path_ids <: Tuple
+    path_ids::PI
+
+    function PathState{T}(
+            interval::Interval{T}, path_operators::PO, path_ids::PI) where {T, PO, PI}
+        return new{T, PO, PI}(interval, path_operators, path_ids)
+    end
+    function PathState{T}(
+            interval::Interval, path_operators::PO, path_ids::PI) where {T, PO, PI}
+        return new{T, PO, PI}(Interval{T}(interval), path_operators, path_ids)
+    end
 end
 
 function PathState(interval::Interval{T}, id::Tuple{Int, Int} = (1, 1)) where {T}

lib/DataDrivenLux/test/candidate.jl

@@ -26,8 +26,7 @@ using StableRNGs
 
     @test DataDrivenLux.get_scales(candidate) ≈ ones(Float64, 1)
     @test isempty(DataDrivenLux.get_parameters(candidate))
-    @test_nowarn DataDrivenLux.optimize_candidate!(
-        candidate, dataset; options = Optim.Options())
+    @test_nowarn DataDrivenLux.optimize_candidate!(candidate, dataset)
 end
 
 @testset "Candidate with parametes" begin