add some functionality for ProductDistribution

Author: itsdfish

Project.toml

@@ -1,7 +1,7 @@
 name = "SequentialSamplingModels"
 uuid = "0e71a2a6-2b30-4447-8742-d083a85e82d1"
 authors = ["itsdfish"]
-version = "0.12.6"
+version = "0.12.7"
 
 [deps]
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"

src/SequentialSamplingModels.jl

@@ -12,6 +12,7 @@ using PrettyTables
 using Random
 using SpecialFunctions
 
+using Distributions: ProductDistribution
 using HCubature: hcubature
 using StatsBase: Weights
 
@@ -28,6 +29,7 @@ import Distributions: mean
 import Distributions: minimum
 import Distributions: pdf
 import Distributions: rand
+import Distributions: rand!
 import Distributions: std
 import StatsAPI: params
 import StatsBase: cor2cov
@@ -109,4 +111,5 @@ include("MDFT.jl")
 include("ClassicMDFT.jl")
 include("MLBA.jl")
 include("ShiftedLogNormal.jl")
+include("product_distribution.jl")
 end

src/product_distribution.jl

@@ -0,0 +1,42 @@
+function rand(
+    rng::AbstractRNG,
+    s::Sampleable{T, R}
+) where {T <: Matrixvariate, R <: SequentialSamplingModels.Mixed}
+    n = size(s, 2)
+    data = (; choice = fill(0, n), rt = fill(0.0, n))
+    return rand!(rng, s, data)
+end
+
+function rand(
+    rng::AbstractRNG,
+    s::Sampleable{T, R},
+    dims::Dims
+) where {T <: Matrixvariate, R <: SequentialSamplingModels.Mixed}
+    n = size(s, 2)
+    ax = map(Base.OneTo, dims)
+    data = [(; choice = fill(0, n), rt = fill(0.0, n)) for _ in Iterators.product(ax...)]
+    return [rand!(rng, s, d) for d ∈ data]
+end
+
+function rand!(
+    rng::AbstractRNG,
+    s::Sampleable{T, R},
+    data::NamedTuple
+) where {T <: Matrixvariate, R <: SequentialSamplingModels.Mixed}
+    for i ∈ 1:size(s, 2)
+        data.choice[i], data.rt[i] = rand(rng, s.dists[i])
+    end
+    return data
+end
+
+function logpdf(d::ProductDistribution, data_array::Array{<:NamedTuple, N}) where {N}
+    return [logpdf(d, data) for data ∈ data_array]
+end
+
+function logpdf(d::ProductDistribution, data::NamedTuple)
+    LL = 0.0
+    for i ∈ 1:length(d.dists)
+        LL += logpdf(d.dists[i], data.choice[i], data.rt[i])
+    end
+    return LL
+end

src/type_system.jl

@@ -393,3 +393,8 @@ Increments the evidence states `x` on each time step.
 """
 increment!(model::SSM1D, x, μΔ; Δt = 0.001) =
     increment!(Random.default_rng(), model, x, μΔ; Δt)
+
+Base.eltype(::Type{<:Sampleable{F, Mixed}}) where {F} =
+    @NamedTuple{choice::Vector{Int64}, rt::Vector{Float64}}
+
+Base.length(s::ContinuousMultivariateSSM) = 2

test/product_distribution_tests.jl

@@ -0,0 +1,120 @@
+@safetestset "ProductDistribution Tests" begin
+    @safetestset "rand SSM1D 1" begin
+        using Distributions
+        using SequentialSamplingModels
+        using Test
+
+        walds = [Wald(; ν = 2.5, α = 0.1, τ = 0.2), Wald(; ν = 1.5, α = 1, τ = 10)]
+        dist = product_distribution(walds)
+        data = rand(dist)
+        @test length(data) == 2
+        @test data[1] < 10
+        @test data[2] > 10
+    end
+
+    @safetestset "rand SSM1D 2" begin
+        using Distributions
+        using SequentialSamplingModels
+        using Test
+
+        walds = [Wald(; ν = 2.5, α = 0.1, τ = 0.2), Wald(; ν = 1.5, α = 1, τ = 10)]
+        dist = product_distribution(walds)
+        data = rand(dist, 3)
+        @test size(data) == (2, 3)
+        @test all(data[1, :] .< 10)
+        @test all(data[2, :] .> 10)
+    end
+
+    @safetestset "rand logpdf 1" begin
+        using Distributions
+        using SequentialSamplingModels
+        using Test
+
+        walds = [Wald(; ν = 2.5, α = 0.1, τ = 0.2), Wald(; ν = 1.5, α = 1, τ = 10)]
+        dist = product_distribution(walds)
+        data = rand(dist)
+        LL1 = logpdf(dist, data)
+        LL2 = sum(i -> logpdf(walds[i], data[i]), 1:2)
+        @test LL1 ≈ LL2
+    end
+
+    @safetestset "logpdf SSM1D 2" begin
+        using Distributions
+        using SequentialSamplingModels
+        using Test
+
+        walds = [Wald(; ν = 2.5, α = 0.1, τ = 0.2), Wald(; ν = 1.5, α = 1, τ = 10)]
+        dist = product_distribution(walds)
+        data = rand(dist, 3)
+        LL1 = logpdf(dist, data)
+        LL2 = sum(i -> logpdf(walds[i], data[i, :]), 1:2)
+        @test LL1 ≈ LL2
+    end
+
+    @safetestset "rand SSM2D 1" begin
+        using Distributions
+        using SequentialSamplingModels
+        using Test
+
+        lbas = [
+            LBA(; ν = [3, 2], A = 0.8, k = 0.2, τ = 0.1),
+            LBA(ν = [1, 2], A = 0.5, k = 0.3, τ = 10)
+        ]
+        dist = product_distribution(lbas)
+        data = rand(dist)
+        @test length(data.rt) == 2
+        @test data.rt[1] < 10
+        @test data.rt[2] > 10
+    end
+
+    @safetestset "rand SSM2D 2" begin
+        using Distributions
+        using SequentialSamplingModels
+        using Test
+
+        lbas = [
+            LBA(; ν = [3, 2], A = 0.8, k = 0.2, τ = 0.1),
+            LBA(ν = [1, 2], A = 0.5, k = 0.3, τ = 10)
+        ]
+        dist = product_distribution(lbas)
+        data = rand(dist, 3)
+        @test length(data) == 3
+        @test all(map(i -> data[i].rt[1], 1:3) .< 10)
+        @test all(map(i -> data[i].rt[2], 1:3) .> 10)
+    end
+
+    @safetestset "logpdf SSM2D 1" begin
+        using Distributions
+        using SequentialSamplingModels
+        using Test
+
+        lbas = [
+            LBA(; ν = [3, 2], A = 0.8, k = 0.2, τ = 0.1),
+            LBA(ν = [1, 2], A = 0.5, k = 0.3, τ = 10)
+        ]
+        dist = product_distribution(lbas)
+        data = rand(dist)
+
+        LL1 = logpdf(dist, data)
+        LL2 = sum(i -> logpdf(lbas[i], data.choice[i], data.rt[i]), 1:2)
+        @test LL1 ≈ LL2
+    end
+
+    @safetestset "logpdf SSM2D 2" begin
+        using Distributions
+        using SequentialSamplingModels
+        using Test
+
+        lbas = [
+            LBA(; ν = [3, 2], A = 0.8, k = 0.2, τ = 0.1),
+            LBA(ν = [1, 2], A = 0.5, k = 0.3, τ = 10)
+        ]
+        dist = product_distribution(lbas)
+        data = rand(dist, 3)
+
+        LL1 = logpdf(dist, data)
+        LL2 =
+            map(j -> sum(i -> logpdf(lbas[i], data[j].choice[i], data[j].rt[i]), 1:2), 1:3)
+        @test LL1 ≈ LL2
+    end
+end