LMM cluster permutation

.gitignore

@@ -6,3 +6,4 @@
 /docs/Manifest.toml
 /docs/build/
 /docs/src/generated
+/docs/dev

Project.toml

@@ -5,11 +5,32 @@ version = "1.0.0-DEV"
 
 [deps]
 BSplineKit = "093aae92-e908-43d7-9660-e50ee39d5a0a"
+ClusterDepth = "c8d8bbfa-f476-4995-adff-2987f04015d1"
+Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
+MixedModels = "ff71e718-51f3-5ec2-a782-8ffcbfa3c316"
+MixedModelsPermutations = "647c4018-d7ef-4d03-a0cc-8889a722319e"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+StatsAPI = "82ae8749-77ed-4fe6-ae5f-f523153014b0"
 StatsModels = "3eaba693-59b7-5ba5-a881-562e759f1c8d"
 Unfold = "181c99d8-e21b-4ff3-b70b-c233eddec679"
 
+[weakdeps]
+ClusterDepth = "c8d8bbfa-f476-4995-adff-2987f04015d1"
+MixedModels = "ff71e718-51f3-5ec2-a782-8ffcbfa3c316"
+MixedModelsPermutations = "647c4018-d7ef-4d03-a0cc-8889a722319e"
+
+[extensions]
+UnfoldStatsMixedModelsExt = ["MixedModels"]
+UnfoldStatsMixedModelsPermutationsExt = ["ClusterDepth", "MixedModelsPermutations"]
+
 [compat]
-julia = "1"
+BSplineKit = "0.16,0.17"
+ClusterDepth = "0.2"
+MixedModels = "4"
+MixedModelsPermutations = "0.2"
+StatsModels = "0.7"
+Unfold = "0.6,0.7"
+julia = "1.9,1.10"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

docs/Project.toml

@@ -1,12 +1,15 @@
 [deps]
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 Chain = "8be319e6-bccf-4806-a6f7-6fae938471bc"
+ClusterDepth = "c8d8bbfa-f476-4995-adff-2987f04015d1"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 Glob = "c27321d9-0574-5035-807b-f59d2c89b15c"
 HypothesisTests = "09f84164-cd44-5f33-b23f-e6b0d136a0d5"
 Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 LiveServer = "16fef848-5104-11e9-1b77-fb7a48bbb589"
+MixedModels = "ff71e718-51f3-5ec2-a782-8ffcbfa3c316"
+MixedModelsPermutations = "647c4018-d7ef-4d03-a0cc-8889a722319e"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Unfold = "181c99d8-e21b-4ff3-b70b-c233eddec679"

docs/literate/tutorial/lmm_clusterdepth.jl

@@ -0,0 +1,80 @@
+
+# get some data
+using UnfoldSim
+using Unfold
+using MixedModelsPermutations, ClusterDepth # both necessary to activate correct extension!
+using UnfoldStats
+using StatsModels
+using Random
+
+srate = 25
+design = MultiSubjectDesign(;
+    n_subjects = 30,
+    n_items = 40,
+    items_between = Dict(:stimtype => ["car", "face"]),
+)
+#both_within = Dict(:condition=>["scrambled","intact"]))
+contrasts = Dict(:stimtype => DummyCoding())
+p1 = MixedModelComponent(;
+    basis = UnfoldSim.p100(; sfreq = srate),
+    formula = @formula(dv ~ 1 + (1 | subject) + (1 | item)),
+    β = [5.0],
+    σs = Dict(:subject => [0.0], :item => [0.0]),
+    contrasts = contrasts,
+);
+
+n1 = MixedModelComponent(;
+    basis = UnfoldSim.n170(; sfreq = srate),
+    formula = @formula(dv ~ 1 + stimtype + (1 + stimtype | subject) + (1 | item)),
+    β = [1.0, 4], # n170-basis is negative
+    σs = Dict(:subject => [2.0, 0.25], :item => [0.25]),
+    contrasts = contrasts,
+);
+
+p3 = MixedModelComponent(;
+    basis = UnfoldSim.p300(; sfreq = srate),
+    formula = @formula(dv ~ 1 + (1 | subject) + (1 + stimtype | item)),
+    β = [4.0],
+    σs = Dict(:subject => [1.0], :item => [0.5, 2]),
+    contrasts = contrasts,
+);
+
+
+
+data_e, events = UnfoldSim.simulate(
+    design,
+    [p1, n1, p3],
+    UniformOnset(srate * 2, 10),
+    PinkNoise(; noiselevel = 1);
+    return_epoched = true,
+) # 18
+times = range(-0.1, 0.5, length = size(data_e, 1))
+data_e = reshape(data_e, size(data_e, 1), :)
+#events.latency .+= repeat(range(0,length=size(data,2),step=size(data,1)),inner=size(events[events.subject.=="S01",:],1))
+
+
+
+# # Fit LMM 
+m = fit(
+    UnfoldModel,
+    [
+        Any => (
+            @formula(0 ~ 1 + stimtype + (1 + stimtype | item) + (1 + stimtype | subject)),
+            times,
+        ),
+    ],
+    events,
+    data_e,
+);
+
+
+# # Cluster Permute :)
+coefficient = 2
+pvalues(
+    MersenneTwister(1),
+    m,
+    data_e,
+    coefficient;
+    n_permutations = 10,
+    clusterforming_threshold = 1.8,
+)

docs/make.jl

@@ -35,6 +35,7 @@ makedocs(;
         "Home" => "index.md",
         "Tutorials" =>
             ["Two-stage EEG analysis" => "generated/tutorial/two_stage_analysis.md"],
+        ["LMM EEG + ClusterDepth" => "generated/tutorial/lmm_clusterdepth.md"],
     ],
 )
 

ext/UnfoldStatsMixedModelsExt.jl

@@ -0,0 +1,23 @@
+module UnfoldStatsMixedModelsExt
+using Unfold
+using UnfoldStats
+using MixedModels
+import StatsAPI: pvalue
+
+lmm_ext = Base.get_extension(Unfold, :UnfoldMixedModelsExt)
+
+if isnothing(lmm_ext)
+    error("Something went wrong with getting the Unfold UnfoldMixedModelsExt extension")
+end
+# Currently, `extract_coefs` is not implemented for mixed-effects models
+UnfoldStats.extract_coefs(
+    model::Union{
+        lmm_ext.UnfoldLinearMixedModel,
+        lmm_ext.UnfoldLinearMixedModelContinuousTime,
+    },
+    predictor,
+    basisname,
+) = throw(
+    "The `extract_coefs` function is currently not implemented for mixed-effects models.",
+)
+end

ext/UnfoldStatsMixedModelsPermutationsExt.jl

@@ -0,0 +1,174 @@
+module UnfoldStatsMixedModelsPermutationsExt
+using Unfold
+import Unfold: pvalues
+using UnfoldStats
+#using MixedModels
+using MixedModelsPermutations
+using ClusterDepth
+using Logging
+using Random
+const MixedModels = MixedModelsPermutations.MixedModels
+LMMext = Base.get_extension(Unfold, :UnfoldMixedModelsExt)
+
+
+function Unfold.pvalues(
+    rng,
+    model::LMMext.UnfoldLinearMixedModel,
+    data,
+    coefficient;
+    type = "clusterdepth",
+    clusterforming_threshold,
+    kwargs...,
+)
+    if type != "clusterdepth"
+        error("other types (e.g. FDR currently not implemented")
+    elseif type == "clusterdepth"
+        return lmm_clusterdepth(
+            rng,
+            model,
+            data,
+            coefficient;
+            clusterforming_threshold,
+            kwargs...,
+        )
+    end
+end
+
+
+
+
+function lmm_clusterdepth(
+    rng,
+    model,
+    data,
+    coefficient;
+    lmm_statistic = :z,
+    clusterforming_threshold,
+    kwargs...,
+)
+    permuted = lmm_permutations(rng, model, data, coefficient; kwargs...)
+    observed = get_lmm_statistic(model, coefficient, lmm_statistic)
+    return lmm_clusterdepth_pvalues(
+        rng,
+        observed,
+        permuted;
+        clusterforming_threshold,
+        kwargs...,
+    )
+end
+function lmm_clusterdepth_pvalues(
+    rng,
+    observed,
+    permuted;
+    clusterforming_threshold,
+    kwargs...,
+)
+
+    # we need global variables here (yes, sorry...), because instead of actually
+    # letting ClusterDepth do the permutation, we just have to index the already
+    # permuted data given in the function (`permuted`)
+    global n_permutation_count
+    n_permutation_count = 0
+    function _fake_permutation_fun(r, data)
+        global n_permutation_count
+        n_permutation_count = n_permutation_count + 1
+        return permuted[:, n_permutation_count]
+    end
+    J_tuple = ClusterDepth.perm_clusterdepths_both(
+        rng,
+        abs.(permuted),
+        _fake_permutation_fun,
+        clusterforming_threshold;
+        statfun = x -> abs.(x),
+        nₚ = size(permuted, 2),
+    )
+
+    pvals = ClusterDepth.pvals(observed, J_tuple, clusterforming_threshold)
+
+end
+
+function lmm_permutations(
+    rng::AbstractRNG,
+    model,
+    data::AbstractArray{<:Real,3},
+    coefficient::Int;
+    n_permutations = 500,
+    lmm_statistic = :z,
+    time_selection = 1:size(data, 2),
+)
+    permdata = Matrix{Float64}(undef, length(time_selection), n_permutations)
+
+
+    Xs = LMMext.prepare_modelmatrix(model)
+
+    mm_outer = LMMext.LinearMixedModel_wrapper(Unfold.formulas(model), data[1, 1, :], Xs)
+    mm_outer.optsum.maxtime = 0.1 # 
+
+    chIx = 1 # for now we only support 1 channel anyway
+    #
+    #p = Progress(length(time_selection))
+    #Threads.@threads for tIx =1:length(time_selection)
+    #@showprogress "Processing Timepoints" for 
+    Threads.@threads for tIx = 1:length(time_selection)
+
+        # splice in the correct dataa for residual calculation
+        mm = deepcopy(mm_outer)
+        mm.y .= data[chIx, time_selection[tIx], :]
+
+        # set the previous calculated model-fit
+        θ = Vector(modelfit(model).θ[time_selection[tIx]])
+        @debug size(θ)
+        MixedModels.updateL!(MixedModels.setθ!(mm, θ))
+
+        # get the coefficient 
+        H0 = coef(mm)
+        # set the one of interest to 0
+        H0[coefficient] = 0
+        # run the permutation
+
+        permutations = undef
+        Logging.with_logger(NullLogger()) do   # remove NLopt warnings  
+            permutations = permutation(
+                deepcopy(rng), # important here is to set the same seed to keep flip all time-points the same
+                n_permutations,
+                mm;
+                β = H0,
+                hide_progress = true,
+                #blup_method = MixedModelsPermutations.olsranef,
+            ) # constant rng to keep autocorr & olsranef for singular models
+        end
+
+        # extract the test-statistic
+
+        permdata[tIx, :] =
+            get_lmm_statistic(model, permutations, coefficient, lmm_statistic)
+
+        #next!(p)
+    end # end for
+    return permdata
+end
+
+function get_lmm_statistic(
+    model,
+    permutations::MixedModelsPermutations.MixedModels.MixedModelFitCollection,
+    coefficient::Int,
+    lmm_statistic,
+)
+    [
+        getproperty(m, lmm_statistic) for m in permutations.coefpvalues if
+        String(m.coefname) == Unfold.coefnames(Unfold.formulas(model))[1][coefficient]
+    ]
+
+end
+function get_lmm_statistic(
+    model::LMMext.UnfoldLinearMixedModel,
+    coefficient::Int,
+    lmm_statistic,
+)
+    return get_lmm_statistic(model, modelfit(model), coefficient, lmm_statistic)
+    #    r = coeftable(m)
+    #    r = subset(r, :group => (x -> isnothing.(x)), :coefname => (x -> x .!== "(Intercept)"))
+    #    tvals = abs.(r.estimate ./ r.stderror)
+    #    return tvals
+end
+end

src/UnfoldStats.jl

@@ -1,5 +1,6 @@
 module UnfoldStats
 
+using BSplineKit: Reexport
 using Unfold
 using BSplineKit
 using StatsModels
@@ -8,4 +9,5 @@ include("extract_coefs.jl")
 
 # export functions to extract model coefficients 
 export extract_coefs
+
 end

src/extract_coefs.jl

@@ -170,14 +170,6 @@ 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.",
-)
 
 """
     extract_coefs(models::Vector{<:UnfoldModel}, predictor, basisname)