From 6554b9adc5b23dedf0f3691e5afcda5beb5eb3b5 Mon Sep 17 00:00:00 2001
From: jschepers <judith.schepers@vis.uni-stuttgart.de>
Date: Thu, 7 Mar 2024 16:57:22 +0000
Subject: [PATCH 01/13] tested traits for sim and fit functions

---
 benchmark/sim_and_fit.jl | 297 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 297 insertions(+)
 create mode 100644 benchmark/sim_and_fit.jl

diff --git a/benchmark/sim_and_fit.jl b/benchmark/sim_and_fit.jl
new file mode 100644
index 0000000..ce406ef
--- /dev/null
+++ b/benchmark/sim_and_fit.jl
@@ -0,0 +1,297 @@
+using UnfoldSim
+using Unfold
+using StableRNGs
+using DataFrames
+
+function define_simulation(sim_type, β, σs; n_subjects = 30, n_items = 100, noiselevel = 2)
+
+    # Create design
+    conditions = get_conditions(sim_type)
+
+    design = MultiSubjectDesign(;
+        n_subjects = n_subjects,
+        n_items = n_items,
+        items_between = conditions,
+    )
+
+    # Specify component
+    basis = p100()
+    formula = get_formula(sim_type)
+    signal = create_component(sim_type, basis, formula, β, σs)
+
+    # Specify inter-onset distribution
+    onset = UniformOnset(; width = 50, offset = 1)
+
+    # Specify noise
+    noise = PinkNoise(; noiselevel = noiselevel)
+
+    return Simulation(design, signal, onset, noise)
+end
+
+function sim_and_fit(
+    sim_type::SimulationType,
+    simulation::Simulation,
+    model_type::Type{<:UnfoldModel};
+    seed::Integer = 1,
+)
+
+    # At the moment, the function is not implemented for mixed models
+    if model_type in [UnfoldLinearMixedModel, UnfoldLinearMixedModelContinuousTime]
+        throw("Not implemented.")
+    end
+
+    # Set parameter(s) for data simulation
+    if model_type == UnfoldLinearModel
+        return_epoched = true
+    else # UnfoldLinearModelContinuousTime
+        return_epoched = false
+    end
+
+    # Simulate data
+    data, events = simulate_data(sim_type, simulation, return_epoched, seed)
+
+    # Create event dict containing basis function(s)/times and formula(s) for all events
+    event_dict = create_event_dict(sim_type, model_type, simulation)
+
+    # Fit an Unfold model for each subject
+    subject_list = unique(events.subject)
+    model_list = model_type[]
+
+    # Slice the data by its last dimension (i.e. the subject dimension)
+    data_slices = eachslice(data, dims = ndims(data))
+
+    for s = 1:size(data, ndims(data))
+        m = fit(
+            UnfoldModel,
+            event_dict,
+            subset(events, :subject => ByRow(==(subject_list[s]))),
+            data_slices[s],
+        )
+        push!(model_list, m)
+    end
+
+    models = DataFrame(subject = subject_list, unfoldmodel = model_list)
+
+    return models
+end
+
+simulate_data(sim_type::T, simulation, return_epoched, seed) where {T} =
+    simulate_data(EventStyle(T), simulation, return_epoched, seed)
+
+function simulate_data(::SingleEventType, simulation, return_epoched, seed)
+
+    # Simulate data
+    data, events = simulate(StableRNG(seed), simulation; return_epoched = return_epoched)
+
+    return data, events
+end
+
+function simulate_data(::MultipleEventTypes, simulation, return_epoched, seed)
+
+    # Simulate data
+    data, events = simulate(StableRNG(seed), simulation; return_epoched = return_epoched)
+
+    # Add an event column to the events df and assign each event to an event type
+    events[!, :event] = repeat(["stim", "fix"], size(events, 1) ÷ 2)
+
+    return data, events
+end
+
+create_event_dict(sim_type::T, model_type, simulation) where {T} = create_event_dict(
+    EventStyle(T),
+    PredictorStyle(T),
+    model_type::Type{<:UnfoldModel},
+    simulation,
+)
+
+function create_event_dict(
+    ::MultipleEventTypes,
+    ::ManyPredictors,
+    model_type::Type{<:UnfoldModel},
+    simulation,
+)
+    # Create times vector/basis function(s) (for model fitting)
+    if model_type == UnfoldLinearModel
+        #times = axes(data, 1)
+        times = 1:UnfoldSim.maxlength(simulation.components)
+        t_stim = times
+        t_fix = times
+    else # UnfoldLinearModelContinuousTime
+        t_stim = firbasis(τ = (-0.1, 1.5), sfreq = 100, name = "stimulus")
+        t_fix = firbasis(τ = (-0.1, 1), sfreq = 100, name = "fixation")
+    end
+
+    # Define formula(s)
+    f_stim = @formula 0 ~ 1 + continuous
+    f_fix = @formula 0 ~ 1 + spl(continuous, 4) + continuous + condition * pet
+
+    # Combine basis function(s)/times and formula(s) with the corresponding event
+    event_dict = Dict("stim" => (f_stim, t_stim), "fix" => (f_fix, t_fix))
+
+    return event_dict
+end
+
+function create_event_dict(
+    ::MultipleEventTypes,
+    ::OnlySplines,
+    model_type::Type{<:UnfoldModel},
+    simulation,
+)
+    # Create times vector/basis function(s) (for model fitting)
+    if model_type == UnfoldLinearModel
+        #times = axes(data, 1)
+        times = 1:UnfoldSim.maxlength(simulation.components)
+        t_stim = times
+        t_fix = times
+    else # UnfoldLinearModelContinuousTime
+        t_stim = firbasis(τ = (-0.1, 1.5), sfreq = 100, name = "stimulus")
+        t_fix = firbasis(τ = (-0.1, 1), sfreq = 100, name = "fixation")
+    end
+
+    # Define formula(s)
+    f_stim = @formula 0 ~ 1
+    f_fix = @formula 0 ~ 1 + spl(continuous, 4)
+
+    # Combine basis function(s)/times and formula(s) with the corresponding event
+    event_dict = Dict("stim" => (f_stim, t_stim), "fix" => (f_fix, t_fix))
+
+    return event_dict
+end
+
+function create_event_dict(
+    ::SingleEventType,
+    ::ManyPredictors,
+    model_type::Type{<:UnfoldModel},
+    simulation,
+)
+    # Create times vector/basis function(s) (for model fitting)
+    if model_type == UnfoldLinearModel
+        #times = axes(data, 1)
+        t = 1:UnfoldSim.maxlength(simulation.components)
+    else # UnfoldLinearModelContinuousTime
+        t = firbasis((-0.1, 1.0), 100)
+    end
+
+    # Define formula(s)
+    f = @formula 0 ~ 1 + spl(continuous, 4) + continuous + condition * pet
+
+    # Combine basis function(s)/times and formula(s) with the corresponding event
+    event_dict = Dict(Any => (f, t))
+
+    return event_dict
+end
+
+function create_event_dict(
+    ::SingleEventType,
+    ::OnlySplines,
+    model_type::Type{<:UnfoldModel},
+    simulation,
+)
+    # Create times vector/basis function(s) (for model fitting)
+    if model_type == UnfoldLinearModel
+        #times = axes(data, 1)
+        t = 1:UnfoldSim.maxlength(simulation.components)
+    else # UnfoldLinearModelContinuousTime
+        t = firbasis((-0.1, 1.0), 100)
+    end
+
+    # Define formula(s)
+    f = @formula 0 ~ 1 + spl(continuous, 4)
+
+    # Combine basis function(s)/times and formula(s) with the corresponding event
+    event_dict = Dict(Any => (f, t))
+
+    return event_dict
+end
+
+get_conditions(sim_type::T) where {T} = get_conditions(PredictorStyle(T))
+
+function get_conditions(::OnlySplines)
+    conditions = Dict(:continuous => range(-5, 5, length = 50))
+    return conditions
+end
+
+function get_conditions(::ManyPredictors)
+    conditions = Dict(
+        :continuous => range(-5, 5, length = 25),
+        :condition => ["face", "car"],
+        :pet => ["cat", "dog"],
+    )
+    return conditions
+end
+
+get_formula(sim_type::T) where {T} = get_formula(PredictorStyle(T))
+
+function get_formula(::OnlySplines)
+    formula = @formula(
+        0 ~
+            1 +
+            continuous +
+            continuous^2 +
+            continuous^3 +
+            (1 + continuous + continuous^2 + continuous^3 | subject)
+    )
+    return formula
+end
+
+function get_formula(::ManyPredictors)
+    formula = @formula(
+        0 ~
+            1 +
+            continuous +
+            continuous^2 +
+            continuous^3 +
+            condition +
+            pet +
+            (1 + continuous + continuous^2 + continuous^3 | subject)
+    )
+    return formula
+end
+
+create_component(sim_type::T, basis, formula, β, σs) where {T} =
+    create_component(ChannelStyle(T), basis, formula, β, σs)
+
+function create_component(::SingleChannel, basis, formula, β, σs)
+    signal = MixedModelComponent(; basis = basis, formula = formula, β = β, σs = σs)
+    return signal
+end
+
+function create_component(::MultiChannel, basis, formula, β, σs)
+    signal = MixedModelComponent(; basis = basis, formula = formula, β = β, σs = σs)
+    # Wrap the component in a multichannel component
+    # Load headmodel
+    hart = headmodel(type = "hartmut")
+    source_idx = UnfoldSim.closest_src(hart, "Left Postcentral Gyrus")
+    projection = UnfoldSim.magnitude(hart)
+    # Only use the first channels/electrodes
+    channels = 1:3
+
+    multichannel_signal =
+        MultichannelComponent(signal, projection[channels, source_idx], NoNoise())
+    return multichannel_signal
+end
+
+abstract type SimulationType end
+struct UnitTestSimulation <: SimulationType end
+struct BenchmarkSimulation <: SimulationType end
+
+abstract type PredictorStyle end
+struct OnlySplines <: PredictorStyle end
+struct ManyPredictors <: PredictorStyle end
+
+abstract type ChannelStyle end
+struct SingleChannel <: ChannelStyle end
+struct MultiChannel <: ChannelStyle end
+
+abstract type EventStyle end
+struct SingleEventType <: EventStyle end
+struct MultipleEventTypes <: EventStyle end
+
+PredictorStyle(::Type) = OnlySplines()
+PredictorStyle(::Type{UnitTestSimulation}) = ManyPredictors()
+
+ChannelStyle(::Type) = SingleChannel()
+ChannelStyle(::Type{UnitTestSimulation}) = MultiChannel()
+
+EventStyle(::Type) = SingleEventType()
+EventStyle(::Type{UnitTestSimulation}) = MultipleEventTypes()
\ No newline at end of file

From 8e32a9a151124f20d9c569430342c3dc8a593ea8 Mon Sep 17 00:00:00 2001
From: jschepers <judith.schepers@vis.uni-stuttgart.de>
Date: Thu, 7 Mar 2024 17:23:22 +0000
Subject: [PATCH 02/13] add extract_coefs functions and helper functions

---
 src/extract_coefs.jl | 123 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 123 insertions(+)
 create mode 100644 src/extract_coefs.jl

diff --git a/src/extract_coefs.jl b/src/extract_coefs.jl
new file mode 100644
index 0000000..8007e2b
--- /dev/null
+++ b/src/extract_coefs.jl
@@ -0,0 +1,123 @@
+# Helper functions
+const BSplineTerm = Base.get_extension(Unfold, :UnfoldBSplineKitExt).BSplineTerm
+
+extract_symbol(t::AbstractTerm) = t.sym
+
+extract_symbol(t::InterceptTerm) = "(Intercept)"
+extract_symbol(t::BSplineTerm) = repeat([t.term.sym], t.df - 1)
+extract_symbol(t::InteractionTerm) = extract_symbol.(t.terms)
+extract_symbol(t::FunctionTerm) = extract_symbol.(t.args)
+
+extract_symbol(t::MatrixTerm) = extract_symbol(t.terms)
+extract_symbol(t::Unfold.TimeExpandedTerm) =
+    repeat(extract_symbol(t.term), inner = length(Unfold.colnames(t.basisfunction)))
+extract_symbol(f::FormulaTerm) = vcat(extract_symbol.(f.rhs)...)
+extract_symbol(t::Vector) = vcat(extract_symbol.(t)...)
+extract_symbol(t::Tuple) = vcat(extract_symbol.(t)...)
+
+
+contained_or_equal(p, e) = (p == e)
+contained_or_equal(p, e::Tuple) = (p in e)
+
+get_predictor_string(p::Symbol) = ":$p"
+get_predictor_string(p::String) = "\"$p\""
+get_predictor_string(p::Tuple) = "$p"
+
+function get_basisnames(model::UnfoldLinearModel)
+    # Extract the event names from the design
+    design_keys = keys(Unfold.design(model))
+
+    # Create a list of the basis names corresponding to each model term
+    basisnames = String[]
+    for (ix, event) in enumerate(design_keys)
+        push!(basisnames, repeat(["event: $(event)"], size(modelmatrix(model)[ix], 2))...)
+    end
+    return basisnames
+end
+
+get_basisnames(model::UnfoldLinearModelContinuousTime) =
+    Unfold.extract_coef_info(Unfold.get_coefnames(model), 1)
+
+
+function extract_coefs(model::UnfoldModel, predictor, basisname)
+
+    # Get vector with underlying predictor variable (symbol) for all coefficients
+    symbols = extract_symbol(Unfold.formula(model))
+
+    # Check whether `predictor` is a predictor in the model
+    if predictor ∉ symbols
+        # TODO: Interactions will be listed separately at the moment, maybe don't list them
+        allowed_predictors = join(get_predictor_string.(unique(symbols)), ", ")
+
+        throw(
+            ArgumentError(
+                "The given predictor $(get_predictor_string(predictor)) was not found in the model. Possible predictors are: $allowed_predictors.",
+            ),
+        )
+    end
+
+    basisname_list = get_basisnames(model)
+
+    # Check whether given `basisname` is in the basisname list of the model
+    if basisname ∉ basisname_list
+        allowed_basisnames = join(["\"$b\"" for b in unique(basisname_list)], ", ")
+
+        throw(
+            ArgumentError(
+                "The given basisname \"$basisname\" was not found in the model. Possible basisnames are: $allowed_basisnames.",
+            ),
+        )
+    end
+
+    # Create a boolean mask which is true for the model coefficients that belong to the given predictor
+    mask_predictor = contained_or_equal.(predictor, symbols)
+
+    # Create a boolean mask which is true for the model coefficients that belong to the given basis name
+    mask_basisfunction = basisname_list .== basisname
+
+    mask_combined = mask_predictor .* mask_basisfunction
+
+    # Check whether the given combination between predictor variable and basisname exist in the model
+    if sum(mask_combined) == 0
+        # TODO: Is `ArgumentError` the right exception to use?
+        throw(
+            ArgumentError(
+                "The given predictor $(get_predictor_string(predictor)) does not exist for the given basisname \"$basisname\".",
+            ),
+        )
+    end
+
+    # Extract the requested coefficients from the coefficient array
+    if typeof(model) == UnfoldLinearModel
+        coef_subset = coef(model)[:, :, mask_combined]
+
+    elseif typeof(model) == UnfoldLinearModelContinuousTime
+
+        n_coefs = length(
+            unique(Unfold.extract_coef_info(Unfold.get_coefnames(model), 2)[mask_combined]),
+        )
+
+        coef_subset_temp = @view(coef(model)[:, mask_combined])
+        # Reshape the coefficient array such that time points and coefficient types get separate dimensions
+        coef_subset = reshape(coef_subset_temp, size(coef_subset_temp, 1), :, n_coefs)
+    else
+        throw("Not implemented.")
+    end
+
+    return coef_subset#::Array{<:Union{<:Missing,<:Float64},3}
+end
+
+function extract_coefs(models::Vector{<:UnfoldModel}, predictor, basisname)
+
+    # Extract the coefficients for all subjects
+    coefs_vector = extract_coefs.(models, predictor, basisname)
+
+    # Check that all coefficient arrays have the same size
+    @assert length(unique(size.(coefs_vector))) == 1
+
+    # Concatenate the coefficients to one array
+    # Dimensions: (channels x times x coefficients x subjects)
+    coefs_all_subjects = cat(coefs_vector..., dims = ndims(coefs_vector[1]) + 1)#::Array{<:Union{<:Missing,<:Float64},4}
+
+    return coefs_all_subjects
+end
\ No newline at end of file

From 9575c1bbd2c2e20b0e4c727835a1b4044a21a19a Mon Sep 17 00:00:00 2001
From: jschepers <judith.schepers@vis.uni-stuttgart.de>
Date: Mon, 11 Mar 2024 12:02:11 +0000
Subject: [PATCH 03/13] added not implemented exception for mixed-effects
 models

---
 src/extract_coefs.jl | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/src/extract_coefs.jl b/src/extract_coefs.jl
index 8007e2b..9f3ec9c 100644
--- a/src/extract_coefs.jl
+++ b/src/extract_coefs.jl
@@ -107,6 +107,15 @@ function extract_coefs(model::UnfoldModel, predictor, basisname)
     return coef_subset#::Array{<:Union{<:Missing,<:Float64},3}
 end
 
+# Currently, `extract_coefs` is not implemented for mixed-effects models
+extract_coefs(
+    model::Union{UnfoldLinearMixedModel,UnfoldLinearMixedModelContinuousTime},
+    predictor,
+    basisname,
+) = throw(
+    "The `extract_coefs` function is currently not implemented for mixed-effects models.",
+)
+
 function extract_coefs(models::Vector{<:UnfoldModel}, predictor, basisname)
 
     # Extract the coefficients for all subjects

From b827b2f4d7a5b7c20c744b349c50a0722fad8ae6 Mon Sep 17 00:00:00 2001
From: jschepers <judith.schepers@vis.uni-stuttgart.de>
Date: Mon, 11 Mar 2024 13:56:04 +0000
Subject: [PATCH 04/13] setup module and test infrastructure

---
 src/UnfoldStats.jl | 4 +++-
 test/runtests.jl   | 4 ++--
 test/setup.jl      | 6 ++++++
 3 files changed, 11 insertions(+), 3 deletions(-)
 create mode 100644 test/setup.jl

diff --git a/src/UnfoldStats.jl b/src/UnfoldStats.jl
index 0043820..f29f41b 100644
--- a/src/UnfoldStats.jl
+++ b/src/UnfoldStats.jl
@@ -1,5 +1,7 @@
 module UnfoldStats
 
-# Write your package code here.
+using Unfold
+
+include("extract_coefs.jl")
 
 end
diff --git a/test/runtests.jl b/test/runtests.jl
index 71e110f..5bf04da 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,6 +1,6 @@
 using UnfoldStats
-using Test
+include("setup.jl")
 
 @testset "UnfoldStats.jl" begin
-    # Write your tests here.
+    include("extract_coefs.jl")
 end
diff --git a/test/setup.jl b/test/setup.jl
new file mode 100644
index 0000000..23f5e72
--- /dev/null
+++ b/test/setup.jl
@@ -0,0 +1,6 @@
+using Test
+using UnfoldSim
+using Unfold
+using DataFrames
+
+include("../benchmark/sim_and_fit.jl")
\ No newline at end of file

From 84a3e93ffd84e73871ec69ffa3f0472269c4e8a7 Mon Sep 17 00:00:00 2001
From: jschepers <judith.schepers@vis.uni-stuttgart.de>
Date: Mon, 11 Mar 2024 15:19:59 +0000
Subject: [PATCH 05/13] adapt test infrastructure and add first version of
 tests for extract_coefs

---
 Project.toml             |   5 ++
 benchmark/Project.toml   |   5 ++
 benchmark/sim_and_fit.jl |  27 +---------
 benchmark/types.jl       |  24 +++++++++
 src/UnfoldStats.jl       |   4 ++
 test/Project.toml        |   7 +++
 test/extract_coefs.jl    | 109 +++++++++++++++++++++++++++++++++++++++
 test/setup.jl            |   2 +
 8 files changed, 157 insertions(+), 26 deletions(-)
 create mode 100644 benchmark/Project.toml
 create mode 100644 benchmark/types.jl
 create mode 100644 test/Project.toml
 create mode 100644 test/extract_coefs.jl

diff --git a/Project.toml b/Project.toml
index 62ce7c9..f8841fc 100644
--- a/Project.toml
+++ b/Project.toml
@@ -3,6 +3,11 @@ uuid = "96fd419a-8306-4ce8-ba5b-cd907cb7647c"
 authors = ["Judith Schepers", "Benedikt V. Ehinger"]
 version = "1.0.0-DEV"
 
+[deps]
+BSplineKit = "093aae92-e908-43d7-9660-e50ee39d5a0a"
+StatsModels = "3eaba693-59b7-5ba5-a881-562e759f1c8d"
+Unfold = "181c99d8-e21b-4ff3-b70b-c233eddec679"
+
 [compat]
 julia = "1"
 
diff --git a/benchmark/Project.toml b/benchmark/Project.toml
new file mode 100644
index 0000000..c702867
--- /dev/null
+++ b/benchmark/Project.toml
@@ -0,0 +1,5 @@
+[deps]
+DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
+StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
+Unfold = "181c99d8-e21b-4ff3-b70b-c233eddec679"
+UnfoldSim = "ed8ae6d2-84d3-44c6-ab46-0baf21700804"
diff --git a/benchmark/sim_and_fit.jl b/benchmark/sim_and_fit.jl
index ce406ef..b8f6881 100644
--- a/benchmark/sim_and_fit.jl
+++ b/benchmark/sim_and_fit.jl
@@ -269,29 +269,4 @@ function create_component(::MultiChannel, basis, formula, β, σs)
     multichannel_signal =
         MultichannelComponent(signal, projection[channels, source_idx], NoNoise())
     return multichannel_signal
-end
-
-abstract type SimulationType end
-struct UnitTestSimulation <: SimulationType end
-struct BenchmarkSimulation <: SimulationType end
-
-abstract type PredictorStyle end
-struct OnlySplines <: PredictorStyle end
-struct ManyPredictors <: PredictorStyle end
-
-abstract type ChannelStyle end
-struct SingleChannel <: ChannelStyle end
-struct MultiChannel <: ChannelStyle end
-
-abstract type EventStyle end
-struct SingleEventType <: EventStyle end
-struct MultipleEventTypes <: EventStyle end
-
-PredictorStyle(::Type) = OnlySplines()
-PredictorStyle(::Type{UnitTestSimulation}) = ManyPredictors()
-
-ChannelStyle(::Type) = SingleChannel()
-ChannelStyle(::Type{UnitTestSimulation}) = MultiChannel()
-
-EventStyle(::Type) = SingleEventType()
-EventStyle(::Type{UnitTestSimulation}) = MultipleEventTypes()
\ No newline at end of file
+end
\ No newline at end of file
diff --git a/benchmark/types.jl b/benchmark/types.jl
new file mode 100644
index 0000000..758d243
--- /dev/null
+++ b/benchmark/types.jl
@@ -0,0 +1,24 @@
+abstract type SimulationType end
+struct UnitTestSimulation <: SimulationType end
+struct BenchmarkSimulation <: SimulationType end
+
+abstract type PredictorStyle end
+struct OnlySplines <: PredictorStyle end
+struct ManyPredictors <: PredictorStyle end
+
+abstract type ChannelStyle end
+struct SingleChannel <: ChannelStyle end
+struct MultiChannel <: ChannelStyle end
+
+abstract type EventStyle end
+struct SingleEventType <: EventStyle end
+struct MultipleEventTypes <: EventStyle end
+
+PredictorStyle(::Type) = OnlySplines()
+PredictorStyle(::Type{UnitTestSimulation}) = ManyPredictors()
+
+ChannelStyle(::Type) = SingleChannel()
+ChannelStyle(::Type{UnitTestSimulation}) = MultiChannel()
+
+EventStyle(::Type) = SingleEventType()
+EventStyle(::Type{UnitTestSimulation}) = MultipleEventTypes()
\ No newline at end of file
diff --git a/src/UnfoldStats.jl b/src/UnfoldStats.jl
index f29f41b..ae1a2f8 100644
--- a/src/UnfoldStats.jl
+++ b/src/UnfoldStats.jl
@@ -1,7 +1,11 @@
 module UnfoldStats
 
 using Unfold
+using BSplineKit
+using StatsModels
 
 include("extract_coefs.jl")
 
+# export functions to extract model coefficients 
+export extract_coefs
 end
diff --git a/test/Project.toml b/test/Project.toml
new file mode 100644
index 0000000..78d6ee2
--- /dev/null
+++ b/test/Project.toml
@@ -0,0 +1,7 @@
+[deps]
+DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
+StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
+StatsModels = "3eaba693-59b7-5ba5-a881-562e759f1c8d"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Unfold = "181c99d8-e21b-4ff3-b70b-c233eddec679"
+UnfoldSim = "ed8ae6d2-84d3-44c6-ab46-0baf21700804"
diff --git a/test/extract_coefs.jl b/test/extract_coefs.jl
new file mode 100644
index 0000000..34c3078
--- /dev/null
+++ b/test/extract_coefs.jl
@@ -0,0 +1,109 @@
+@testset "extract_coefs" begin
+    sim_type = UnitTestSimulation()
+
+    # Print termnames to check how many β and σs values need to be defined
+    # print(StatsModels.termnames(get_formula(sim_type).rhs))
+
+    β = [1, 0, 0, 0, 0, 0]
+    σs = Dict(:subject => [1, 0.1, 0.1, 0.1])
+
+    simulation = define_simulation(sim_type, β, σs, n_subjects = 5)
+
+    for model_type in [UnfoldLinearModel, UnfoldLinearModelContinuousTime]
+
+        models = sim_and_fit(sim_type, simulation, model_type, seed = 1)
+        model_1 = models[1, :unfoldmodel]
+
+        symbols = UnfoldStats.extract_symbol(Unfold.formula(model_1))
+
+        if model_type == UnfoldLinearModel
+
+            # Test helper functions
+            @test symbols == [
+                "(Intercept)",
+                :continuous,
+                "(Intercept)",
+                :continuous,
+                :continuous,
+                :continuous,
+                :continuous,
+                :condition,
+                :pet,
+                (:condition, :pet),
+            ]
+            @test UnfoldStats.contained_or_equal.(:continuous, symbols) ==
+                  [0, 1, 0, 1, 1, 1, 1, 0, 0, 0]
+
+            @test UnfoldStats.get_predictor_string(:continuous) == ":continuous"
+            @test UnfoldStats.get_predictor_string("(Intercept)") == "\"(Intercept)\""
+            @test UnfoldStats.get_predictor_string((:condition, :pet)) ==
+                  "(:condition, :pet)"
+
+            @test UnfoldStats.get_basisnames(model_1) ==
+                  reduce(vcat, fill.(["event: stim", "event: fix"], [2, 8]))
+
+            # Test exceptions
+            @test_throws ArgumentError extract_coefs(model_1, "continuous", "event: fix")
+            @test_throws ArgumentError extract_coefs(model_1, :continuous, "fix")
+            @test_throws ArgumentError extract_coefs(model_1, :pet, "event: stim")
+
+            # Test `extract_coefs` method for single Unfold model
+            coefs_1 = extract_coefs(model_1, :continuous, "event: fix")
+            @test size(coefs_1) == (3, 15, 4)
+            @test coefs_1[:] ==
+                  subset(
+                coeftable(model_1),
+                :basisname => ByRow(==("event: fix")),
+                :coefname => ByRow(
+                    x -> x in [
+                        "continuous",
+                        "spl(continuous,1)",
+                        "spl(continuous,2)",
+                        "spl(continuous,3)",
+                    ],
+                ),
+            ).estimate
+
+            # Test `extract_coefs` method for a vector of Unfold models
+            coefs = extract_coefs(models.unfoldmodel, :continuous, "event: stim")
+            @test coefs[:, :, 1, 4][:] ==
+                  subset(
+                coeftable(models[4, :unfoldmodel]),
+                :basisname => ByRow(==("event: stim")),
+                :coefname => ByRow(==("continuous")),
+            ).estimate
+
+        else
+            # Test helper functions
+            @test length(symbols) == 1210
+            @test UnfoldStats.get_basisnames(model_1) ==
+                  reduce(vcat, fill.(["stimulus", "fixation"], [161 * 2, 111 * 8]))
+
+            # Test `extract_coefs` method for single Unfold model
+            coefs_1 = extract_coefs(model_1, :continuous, "fixation")
+            @test coefs_1[:, 1, :][:] ==
+                  subset(
+                coeftable(model_1),
+                :basisname => ByRow(==("fixation")),
+                :coefname => ByRow(
+                    x -> x in [
+                        "continuous",
+                        "spl(continuous,1)",
+                        "spl(continuous,2)",
+                        "spl(continuous,3)",
+                    ],
+                ),
+                :time => (ByRow(==(-0.1))),
+            ).estimate
+
+            # Test `extract_coefs` method for a vector of Unfold models
+            coefs = extract_coefs(models.unfoldmodel, :continuous, "stimulus")
+            @test coefs[:, :, 1, 4][:] ==
+                  subset(
+                coeftable(models[4, :unfoldmodel]),
+                :basisname => ByRow(==("stimulus")),
+                :coefname => ByRow(==("continuous")),
+            ).estimate
+        end
+    end
+end
\ No newline at end of file
diff --git a/test/setup.jl b/test/setup.jl
index 23f5e72..4f0248b 100644
--- a/test/setup.jl
+++ b/test/setup.jl
@@ -2,5 +2,7 @@ using Test
 using UnfoldSim
 using Unfold
 using DataFrames
+using StableRNGs
 
+include("../benchmark/types.jl")
 include("../benchmark/sim_and_fit.jl")
\ No newline at end of file

From d693021f7081a6833cf18b756f8dc59daf4c1152 Mon Sep 17 00:00:00 2001
From: jschepers <judith.schepers@vis.uni-stuttgart.de>
Date: Mon, 11 Mar 2024 15:21:29 +0000
Subject: [PATCH 06/13] updated .gitignore

---
 .gitignore | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.gitignore b/.gitignore
index 95731a5..71fac49 100644
--- a/.gitignore
+++ b/.gitignore
@@ -2,5 +2,6 @@
 *.jl.cov
 *.jl.mem
 /Manifest.toml
+/test/Manifest.toml
 /docs/Manifest.toml
 /docs/build/

From 12f41cca80482e87a21d5c68332f6555f214e891 Mon Sep 17 00:00:00 2001
From: jschepers <judith.schepers@vis.uni-stuttgart.de>
Date: Tue, 12 Mar 2024 13:30:38 +0000
Subject: [PATCH 07/13] Minor edits

---
 docs/Project.toml     | 1 +
 test/extract_coefs.jl | 1 +
 2 files changed, 2 insertions(+)

diff --git a/docs/Project.toml b/docs/Project.toml
index fdc7a14..695771c 100644
--- a/docs/Project.toml
+++ b/docs/Project.toml
@@ -1,3 +1,4 @@
 [deps]
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 UnfoldStats = "96fd419a-8306-4ce8-ba5b-cd907cb7647c"
diff --git a/test/extract_coefs.jl b/test/extract_coefs.jl
index 34c3078..985ef50 100644
--- a/test/extract_coefs.jl
+++ b/test/extract_coefs.jl
@@ -81,6 +81,7 @@
 
             # Test `extract_coefs` method for single Unfold model
             coefs_1 = extract_coefs(model_1, :continuous, "fixation")
+            @test size(coefs_1) == (3, 111, 4)
             @test coefs_1[:, 1, :][:] ==
                   subset(
                 coeftable(model_1),

From 98780757c9a7e38b303462d02509718ffc1bd063 Mon Sep 17 00:00:00 2001
From: jschepers <judith.schepers@vis.uni-stuttgart.de>
Date: Tue, 12 Mar 2024 14:05:26 +0000
Subject: [PATCH 08/13] added docstrings for extract_coefs function and its
 methods

---
 src/extract_coefs.jl | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/src/extract_coefs.jl b/src/extract_coefs.jl
index 9f3ec9c..8ec788e 100644
--- a/src/extract_coefs.jl
+++ b/src/extract_coefs.jl
@@ -38,7 +38,21 @@ end
 get_basisnames(model::UnfoldLinearModelContinuousTime) =
     Unfold.extract_coef_info(Unfold.get_coefnames(model), 1)
 
+"""
+    extract_coefs(model::UnfoldModel, predictor, basisname)
 
+Return the coefficients of an Unfold model for a certain `predictor` and `basisname`.
+
+For extracting the terms of a predictor variable `predictor` must be a symbol e.g. :continuous.
+For extracting the intercept `predictor` should be a String, i.e. "(Intercept)".
+
+`basisname` must match one of the basis names which can be found in `coeftable(model)`.
+
+Note: If a predictor variable has more than one term in the formula (e.g. a spline set, a categorical variable with several levels or an interaction),
+the coefficients for all terms are returned.
+
+The dimensions of the returned coefficients are channel x times x coefs.
+"""
 function extract_coefs(model::UnfoldModel, predictor, basisname)
 
     # Get vector with underlying predictor variable (symbol) for all coefficients
@@ -116,6 +130,13 @@ extract_coefs(
     "The `extract_coefs` function is currently not implemented for mixed-effects models.",
 )
 
+"""
+    extract_coefs(models::Vector{<:UnfoldModel}, predictor, basisname)
+
+When applied to a vector of Unfold models, extracts the coefficients (matching the predictor and basisname) for all models (usually subjects) and concatenates them.
+
+The dimensions of the returned coefficients are channel x times x coefs x subjects.
+"""
 function extract_coefs(models::Vector{<:UnfoldModel}, predictor, basisname)
 
     # Extract the coefficients for all subjects

From 170221c21113b63f3f1e8f54e8128b7f04ba599a Mon Sep 17 00:00:00 2001
From: jschepers <judith.schepers@vis.uni-stuttgart.de>
Date: Tue, 12 Mar 2024 16:11:43 +0000
Subject: [PATCH 09/13] Add docstrings for helper functions

---
 src/extract_coefs.jl | 58 ++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 56 insertions(+), 2 deletions(-)

diff --git a/src/extract_coefs.jl b/src/extract_coefs.jl
index 8ec788e..27f48ee 100644
--- a/src/extract_coefs.jl
+++ b/src/extract_coefs.jl
@@ -1,6 +1,30 @@
 # Helper functions
 const BSplineTerm = Base.get_extension(Unfold, :UnfoldBSplineKitExt).BSplineTerm
 
+"""
+    extract_symbol(t::AbstractTerm)
+
+Return the symbol(s) underlying a term from a model formula, repeated by their actual coefficient number (after StatsModels.apply_schema).
+
+# Examples
+```julia
+julia> f = @formula 0 ~ 1 + spl(continuous, 4) + continuous + condition + pet + condition & pet
+julia> ... # apply schema using an event dataframe, according to StatsModels
+julia> extract_symbol(f)
+8-element Vector{Any}:
+ "(Intercept)"
+ :continuous
+ :continuous
+ :continuous
+ :continuous
+ :condition
+ :pet
+ (:condition, :pet)
+```
+We get the actual symbols of each predictor - this is different to a function that would return the symbol for each term, which would be `["(Intercept)", :continuous,:continuous,:condition,:pet,(:condition,:pet) ]`
+
+The difference between those two cases would get even more stark, if a basisfunction is in play as it timeexpand terms into many more predictors.
+"""
 extract_symbol(t::AbstractTerm) = t.sym
 
 extract_symbol(t::InterceptTerm) = "(Intercept)"
@@ -15,14 +39,44 @@ extract_symbol(f::FormulaTerm) = vcat(extract_symbol.(f.rhs)...)
 extract_symbol(t::Vector) = vcat(extract_symbol.(t)...)
 extract_symbol(t::Tuple) = vcat(extract_symbol.(t)...)
 
+"""
+    contained_or_equal(p, e)
 
+Test if `p` equals `e` or whether `e` contains `p` if `e` is a tuple.
+"""
 contained_or_equal(p, e) = (p == e)
 contained_or_equal(p, e::Tuple) = (p in e)
 
+"""
+    get_predictor_string(p)
+
+Return string representation based on the type of `p`.
+
+This function is used for a useful display of variables e.g. in an error message.
+
+# Examples
+```jldoctest
+julia> get_predictor_string(:condition)
+":condition"
+```
+
+```jldoctest
+julia> get_predictor_string("(Intercept)")
+"\"(Intercept)\""
+```
+"""
 get_predictor_string(p::Symbol) = ":$p"
 get_predictor_string(p::String) = "\"$p\""
 get_predictor_string(p::Tuple) = "$p"
 
+"""
+    get_basisnames(model::UnfoldModel)
+
+Return the basisnames for all predictor terms as a vector.
+
+The returned vector contains the name of the event type/basis, repeated by their actual coefficient number (after StatsModels.apply_schema).
+If a model has more than one event type (e.g. stimulus and fixation), the vectors are concatenated.
+"""
 function get_basisnames(model::UnfoldLinearModel)
     # Extract the event names from the design
     design_keys = keys(Unfold.design(model))
@@ -51,7 +105,7 @@ For extracting the intercept `predictor` should be a String, i.e. "(Intercept)".
 Note: If a predictor variable has more than one term in the formula (e.g. a spline set, a categorical variable with several levels or an interaction),
 the coefficients for all terms are returned.
 
-The dimensions of the returned coefficients are channel x times x coefs.
+The dimensions of the returned coefficients are channel x times x coefficients.
 """
 function extract_coefs(model::UnfoldModel, predictor, basisname)
 
@@ -135,7 +189,7 @@ extract_coefs(
 
 When applied to a vector of Unfold models, extracts the coefficients (matching the predictor and basisname) for all models (usually subjects) and concatenates them.
 
-The dimensions of the returned coefficients are channel x times x coefs x subjects.
+The dimensions of the returned coefficients are channel x times x coefficients x subjects.
 """
 function extract_coefs(models::Vector{<:UnfoldModel}, predictor, basisname)
 

From 14da6088fd308775b2d7852adf52af31f1a13ccc Mon Sep 17 00:00:00 2001
From: jschepers <judith.schepers@vis.uni-stuttgart.de>
Date: Tue, 12 Mar 2024 16:52:04 +0000
Subject: [PATCH 10/13] fix doctests

---
 src/extract_coefs.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/extract_coefs.jl b/src/extract_coefs.jl
index 27f48ee..e0d988e 100644
--- a/src/extract_coefs.jl
+++ b/src/extract_coefs.jl
@@ -56,12 +56,12 @@ This function is used for a useful display of variables e.g. in an error message
 
 # Examples
 ```jldoctest
-julia> get_predictor_string(:condition)
+julia> UnfoldStats.get_predictor_string(:condition)
 ":condition"
 ```
 
 ```jldoctest
-julia> get_predictor_string("(Intercept)")
+julia> UnfoldStats.get_predictor_string("(Intercept)")
 "\"(Intercept)\""
 ```
 """

From 0d4cc7eb73aef89811937717974ba13acd13c293 Mon Sep 17 00:00:00 2001
From: jschepers <judith.schepers@vis.uni-stuttgart.de>
Date: Tue, 12 Mar 2024 17:01:49 +0000
Subject: [PATCH 11/13] adapted julia version in CI workflow

---
 .github/workflows/CI.yml | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
index 605229b..ef00fe1 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/CI.yml
@@ -19,8 +19,7 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.0'
-          - 'nightly'
+          - '1'
         os:
           - ubuntu-latest
         arch:

From 98bed8522fdcea5aa706a8403bb5fbe698e1a7e2 Mon Sep 17 00:00:00 2001
From: jschepers <judith.schepers@vis.uni-stuttgart.de>
Date: Wed, 13 Mar 2024 09:34:14 +0000
Subject: [PATCH 12/13] Formatting

---
 benchmark/sim_and_fit.jl | 2 +-
 benchmark/types.jl       | 2 +-
 src/extract_coefs.jl     | 2 +-
 test/extract_coefs.jl    | 2 +-
 test/setup.jl            | 2 +-
 5 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/benchmark/sim_and_fit.jl b/benchmark/sim_and_fit.jl
index b8f6881..aec830b 100644
--- a/benchmark/sim_and_fit.jl
+++ b/benchmark/sim_and_fit.jl
@@ -269,4 +269,4 @@ function create_component(::MultiChannel, basis, formula, β, σs)
     multichannel_signal =
         MultichannelComponent(signal, projection[channels, source_idx], NoNoise())
     return multichannel_signal
-end
\ No newline at end of file
+end
diff --git a/benchmark/types.jl b/benchmark/types.jl
index 758d243..e763bf4 100644
--- a/benchmark/types.jl
+++ b/benchmark/types.jl
@@ -21,4 +21,4 @@ ChannelStyle(::Type) = SingleChannel()
 ChannelStyle(::Type{UnitTestSimulation}) = MultiChannel()
 
 EventStyle(::Type) = SingleEventType()
-EventStyle(::Type{UnitTestSimulation}) = MultipleEventTypes()
\ No newline at end of file
+EventStyle(::Type{UnitTestSimulation}) = MultipleEventTypes()
diff --git a/src/extract_coefs.jl b/src/extract_coefs.jl
index e0d988e..4e06344 100644
--- a/src/extract_coefs.jl
+++ b/src/extract_coefs.jl
@@ -204,4 +204,4 @@ function extract_coefs(models::Vector{<:UnfoldModel}, predictor, basisname)
     coefs_all_subjects = cat(coefs_vector..., dims = ndims(coefs_vector[1]) + 1)#::Array{<:Union{<:Missing,<:Float64},4}
 
     return coefs_all_subjects
-end
\ No newline at end of file
+end
diff --git a/test/extract_coefs.jl b/test/extract_coefs.jl
index 985ef50..6bc20b6 100644
--- a/test/extract_coefs.jl
+++ b/test/extract_coefs.jl
@@ -107,4 +107,4 @@
             ).estimate
         end
     end
-end
\ No newline at end of file
+end
diff --git a/test/setup.jl b/test/setup.jl
index 4f0248b..8a72add 100644
--- a/test/setup.jl
+++ b/test/setup.jl
@@ -5,4 +5,4 @@ using DataFrames
 using StableRNGs
 
 include("../benchmark/types.jl")
-include("../benchmark/sim_and_fit.jl")
\ No newline at end of file
+include("../benchmark/sim_and_fit.jl")

From 638486181b150031ef77e76970b18fa29034194c Mon Sep 17 00:00:00 2001
From: jschepers <judith.schepers@vis.uni-stuttgart.de>
Date: Wed, 13 Mar 2024 10:00:50 +0000
Subject: [PATCH 13/13] remove one of the doc tests

---
 src/extract_coefs.jl | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/src/extract_coefs.jl b/src/extract_coefs.jl
index 4e06344..0ac0d2c 100644
--- a/src/extract_coefs.jl
+++ b/src/extract_coefs.jl
@@ -59,11 +59,6 @@ This function is used for a useful display of variables e.g. in an error message
 julia> UnfoldStats.get_predictor_string(:condition)
 ":condition"
 ```
-
-```jldoctest
-julia> UnfoldStats.get_predictor_string("(Intercept)")
-"\"(Intercept)\""
-```
 """
 get_predictor_string(p::Symbol) = ":$p"
 get_predictor_string(p::String) = "\"$p\""