Merge pull request #99 from itsdfish/flexible_sigma

make sigma flexible

Author: itsdfish

src/LBA.jl

@@ -1,12 +1,12 @@
 """
-    LBA{T<:Real} <: AbstractLBA
+    LBA{T <: Real, T1 <: Union{<: T, Vector{<: T}}} <: AbstractLBA
 
 A model object for the linear ballistic accumulator.
 
 # Parameters
 
 - `ν::Vector{T}`: a vector of drift rates
-- `σ::Vector{T}`: a vector of drift rate standard deviation
+- `σ::T1`: a scalar or vector of drift rate standard deviation
 - `A::T`: max start point
 - `k::T`: A + k = b, where b is the decision threshold
 - `τ::T`: an encoding-response offset
@@ -19,7 +19,7 @@ Two constructors are defined below. The first constructor uses positional argume
 
 The second constructor uses keywords with default values, and is not order dependent: 
 
-    LBA(;τ=.3, A=.8, k=.5, ν=[2.0,1.75], σ=[1.0,1.0])
+    LBA(;τ = .3, A = .8, k = .5, ν = [2.0,1.75], σ = 1)
 
 # Example 
 
@@ -35,22 +35,22 @@ loglike = logpdf.(dist, choice, rt)
 
 Brown, S. D., & Heathcote, A. (2008). The simplest complete model of choice response time: Linear ballistic accumulation. Cognitive psychology, 57(3), 153-178.
 """
-mutable struct LBA{T <: Real} <: AbstractLBA
+mutable struct LBA{T <: Real, T1 <: Union{<:T, Vector{<:T}}} <: AbstractLBA{T, T1}
     ν::Vector{T}
-    σ::Vector{T}
+    σ::T1
     A::T
     k::T
     τ::T
 end
 
-function LBA(ν, σ, A, k, τ)
+function LBA(ν, σ, A, k, τ::T) where {T}
     _, _, A, k, τ = promote(ν[1], σ[1], A, k, τ)
-    ν = convert(Vector{typeof(k)}, ν)
-    σ = convert(Vector{typeof(k)}, σ)
+    ν = convert(Vector{T}, ν)
+    σ = isa(σ, Vector) ? convert(Vector{T}, σ) : convert(T, σ)
     return LBA(ν, σ, A, k, τ)
 end
 
-LBA(; τ = 0.3, A = 0.8, k = 0.5, ν = [2.0, 1.75], σ = fill(1.0, length(ν))) =
+LBA(; τ = 0.3, A = 0.8, k = 0.5, ν = [2.0, 1.75], σ = 1) =
     LBA(ν, σ, A, k, τ)
 
 function params(d::LBA)
@@ -79,7 +79,7 @@ function sample_drift_rates(rng::AbstractRNG, ν, σ)
     v = similar(ν)
     n_options = length(ν)
     while negative
-        v = [rand(rng, Normal(ν[i], σ[i])) for i ∈ 1:n_options]
+        v = @. rand(rng, Normal(ν, σ))
         negative = any(x -> x > 0, v) ? false : true
     end
     return v
@@ -97,7 +97,7 @@ function rand(rng::AbstractRNG, d::AbstractLBA)
     return (; choice, rt)
 end
 
-function logpdf(d::AbstractLBA, c, rt)
+function logpdf(d::AbstractLBA{T, T1}, c, rt) where {T, T1 <: Vector{<:Real}}
     (; τ, A, k, ν, σ) = d
     b = A + k
     LL = 0.0
@@ -114,7 +114,24 @@ function logpdf(d::AbstractLBA, c, rt)
     return max(LL, -1000.0)
 end
 
-function pdf(d::AbstractLBA, c, rt)
+function logpdf(d::AbstractLBA{T, T1}, c, rt) where {T, T1 <: Real}
+    (; τ, A, k, ν, σ) = d
+    b = A + k
+    LL = 0.0
+    rt < τ ? (return -Inf) : nothing
+    for i ∈ 1:length(ν)
+        if c == i
+            LL += log_dens(d, ν[i], σ, rt)
+        else
+            LL += log(max(0.0, 1 - cummulative(d, ν[i], σ, rt)))
+        end
+    end
+    pneg = pnegative(d)
+    LL = LL - log(1 - pneg)
+    return max(LL, -1000.0)
+end
+
+function pdf(d::AbstractLBA{T, T1}, c, rt) where {T, T1 <: Vector{<:Real}}
     (; τ, A, k, ν, σ) = d
     b = A + k
     den = 1.0
@@ -132,6 +149,24 @@ function pdf(d::AbstractLBA, c, rt)
     isnan(den) ? (return 0.0) : (return den)
 end
 
+function pdf(d::AbstractLBA{T, T1}, c, rt) where {T, T1 <: Real}
+    (; τ, A, k, ν, σ) = d
+    b = A + k
+    den = 1.0
+    rt < τ ? (return 1e-10) : nothing
+    for i ∈ 1:length(ν)
+        if c == i
+            den *= dens(d, ν[i], σ, rt)
+        else
+            den *= (1 - cummulative(d, ν[i], σ, rt))
+        end
+    end
+    pneg = pnegative(d)
+    den = den / (1 - pneg)
+    den = max(den, 1e-10)
+    isnan(den) ? (return 0.0) : (return den)
+end
+
 function dens(d::AbstractLBA, v, σ, rt)
     (; τ, A, k) = d
     dt = rt - τ
@@ -164,7 +199,7 @@ function cummulative(d::AbstractLBA, v, σ, rt)
     return cm
 end
 
-function pnegative(d::AbstractLBA)
+function pnegative(d::AbstractLBA{T, T1}) where {T, T1 <: Vector{<:Real}}
     (; ν, σ) = d
     p = 1.0
     for i ∈ 1:length(ν)
@@ -173,6 +208,15 @@ function pnegative(d::AbstractLBA)
     return p
 end
 
+function pnegative(d::AbstractLBA{T, T1}) where {T, T1 <: Real}
+    (; ν, σ) = d
+    p = 1.0
+    for i ∈ 1:length(ν)
+        p *= Φ(-ν[i] / σ)
+    end
+    return p
+end
+
 """
     simulate(model::AbstractLBA; n_steps=100, _...)
 

src/LCA.jl

@@ -49,9 +49,9 @@ mutable struct LCA{T <: Real} <: AbstractLCA
     τ::T
 end
 
-function LCA(ν, σ, β, λ, α, τ)
+function LCA(ν, σ, β, λ, α, τ::T) where {T}
     _, σ, β, λ, α, τ = promote(ν[1], σ, β, λ, α, τ)
-    ν = convert(Vector{typeof(τ)}, ν)
+    ν = convert(Vector{T}, ν)
     return LCA(ν, σ, β, λ, α, τ)
 end
 

src/LNR.jl

@@ -4,7 +4,7 @@
 # Parameters 
 
 - `ν`: a vector of means in log-space
-- `σ`: a vector of standard deviation parameter in log-space
+- `σ`: a scalar or vector of standard deviation parameter in log-space
 - `τ`: a encoding-response offset
 
 # Constructors
@@ -15,13 +15,13 @@ Two constructors are defined below. The first constructor uses positional argume
     
 The second constructor uses keywords with default values, and is not order dependent: 
 
-    LNR(; ν = [-1, -2], σ = fill(1.0, length(ν)), τ = 0.20)
+    LNR(; ν = [-1, -2], σ = 1, τ = 0.20)
 
 # Example
 
 ```julia
 using SequentialSamplingModels
-dist = LNR(ν=[-2,-3], σ=[1.0,1.0], τ=.3)
+dist = LNR(ν = [-2,-3], σ = 1, τ = .3)
 choice,rt = rand(dist, 10)
 like = pdf.(dist, choice, rt)
 loglike = logpdf.(dist, choice, rt)
@@ -32,24 +32,24 @@ Rouder, J. N., Province, J. M., Morey, R. D., Gomez, P., & Heathcote, A. (2015).
 The lognormal race: A cognitive-process model of choice and latency with desirable 
 psychometric properties. Psychometrika, 80(2), 491-513.
 """
-struct LNR{T <: Real} <: AbstractLNR
+struct LNR{T <: Real, T1 <: Union{<:T, Vector{<:T}}} <: AbstractLNR{T, T1}
     ν::Vector{T}
-    σ::Vector{T}
+    σ::T1
     τ::T
 end
 
-function LNR(ν, σ, τ)
+function LNR(ν, σ, τ::T) where {T}
     _, _, τ = promote(ν[1], σ[1], τ)
-    ν = convert(Vector{typeof(τ)}, ν)
-    σ = convert(Vector{typeof(τ)}, σ)
+    ν = convert(Vector{T}, ν)
+    σ = isa(σ, Vector) ? convert(Vector{T}, σ) : convert(T, σ)
     return LNR(ν, σ, τ)
 end
 
 function params(d::AbstractLNR)
     return (d.ν, d.σ, d.τ)
 end
 
-LNR(; ν = [-1, -2], σ = fill(1.0, length(ν)), τ = 0.20) = LNR(ν, σ, τ)
+LNR(; ν = [-1, -2], σ = 1, τ = 0.20) = LNR(ν, σ, τ)
 
 function rand(rng::AbstractRNG, dist::AbstractLNR)
     (; ν, σ, τ) = dist
@@ -58,7 +58,7 @@ function rand(rng::AbstractRNG, dist::AbstractLNR)
     return (; choice, rt)
 end
 
-function logpdf(d::AbstractLNR, r::Int, t::Float64)
+function logpdf(d::AbstractLNR{T, T1}, r::Int, t::Float64) where {T, T1 <: Vector{<:Real}}
     (; ν, σ, τ) = d
     LL = 0.0
     for i ∈ 1:length(ν)
@@ -71,7 +71,20 @@ function logpdf(d::AbstractLNR, r::Int, t::Float64)
     return LL
 end
 
-function pdf(d::AbstractLNR, r::Int, t::Float64)
+function logpdf(d::AbstractLNR{T, T1}, r::Int, t::Float64) where {T, T1 <: Real}
+    (; ν, σ, τ) = d
+    LL = 0.0
+    for i ∈ 1:length(ν)
+        if i == r
+            LL += logpdf(LogNormal(ν[i], σ), t - τ)
+        else
+            LL += logccdf(LogNormal(ν[i], σ), t - τ)
+        end
+    end
+    return LL
+end
+
+function pdf(d::AbstractLNR{T, T1}, r::Int, t::Float64) where {T, T1 <: Vector{<:Real}}
     (; ν, σ, τ) = d
     density = 1.0
     for i ∈ 1:length(ν)
@@ -83,3 +96,16 @@ function pdf(d::AbstractLNR, r::Int, t::Float64)
     end
     return density
 end
+
+function pdf(d::AbstractLNR{T, T1}, r::Int, t::Float64) where {T, T1 <: Real}
+    (; ν, σ, τ) = d
+    density = 1.0
+    for i ∈ 1:length(ν)
+        if i == r
+            density *= pdf(LogNormal(ν[i], σ), t - τ)
+        else
+            density *= (1 - cdf(LogNormal(ν[i], σ), t - τ))
+        end
+    end
+    return density
+end

src/MDFT.jl

@@ -90,10 +90,10 @@ mutable struct MDFT{T <: Real} <: AbstractMDFT
     _att_idx::Int
 end
 
-function MDFT(σ, α, τ, γ, κ, ϕ1, ϕ2, β, C)
+function MDFT(σ, α, τ::T, γ, κ, ϕ1, ϕ2, β, C) where {T}
     σ, α, τ, γ, _, ϕ1, ϕ2, β, = promote(σ, α, τ, γ, κ[1], ϕ1, ϕ2, β)
-    κ = convert(Vector{typeof(τ)}, κ)
-    C = convert(Array{typeof(τ), 2}, C)
+    κ = convert(Vector{T}, κ)
+    C = convert(Array{T, 2}, C)
     _CM = zeros(size(C, 1), length(κ))
     S = similar(C)
     return MDFT(σ, α, τ, γ, κ, ϕ1, ϕ2, β, S, C, _CM, 0)

src/MLBA.jl

@@ -1,5 +1,5 @@
 """
-    MLBA{T <: Real} <: AbstractMLBA
+    MLBA{T <: Real, T1 <: Union{<: T, Vector{<: T}}} <: AbstractMLBA{T,T1}
 
 # Fields 
 
@@ -8,7 +8,7 @@
 - `λₚ::T`: decay constant for attention weights of positive differences
 - `λₙ::T`: decay constant for attention weights of negative differences  
 - `γ::T`: risk aversion exponent for subjective values
-- `σ::Vector{T}`: a vector of drift rate standard deviation
+- `σ::T1`: a scalar or vector of drift rate standard deviation
 - `A::T`: max start point
 - `k::T`: A + k = b, where b is the decision threshold
 - `τ::T`: an encoding-response offset
@@ -27,7 +27,7 @@
         τ = 0.3,
         A = 1.0,
         k = 1.0,
-        σ = fill(1.0, n_alternatives)
+        σ = 1.0
     )
 
 # Example 
@@ -59,22 +59,22 @@ loglike = logpdf.(dist, choice, rt, (M,))
 
 Trueblood, J. S., Brown, S. D., & Heathcote, A. (2014). The multiattribute linear ballistic accumulator model of context effects in multialternative choice. Psychological Review, 121(2), 179.
 """
-mutable struct MLBA{T <: Real} <: AbstractMLBA
+mutable struct MLBA{T <: Real, T1 <: Union{<:T, Vector{<:T}}} <: AbstractMLBA{T, T1}
     ν::Vector{T}
     β₀::T
     λₚ::T
     λₙ::T
     γ::T
-    σ::Vector{T}
+    σ::T1
     A::T
     k::T
     τ::T
 end
 
-function MLBA(ν, β₀, λₚ, λₙ, γ, σ, A, k, τ)
+function MLBA(ν, β₀, λₚ, λₙ, γ, σ, A, k::T, τ) where {T}
     _, β₀, λₚ, λₙ, γ, _, A, k, τ = promote(ν[1], β₀, λₚ, λₙ, γ, σ[1], A, k, τ)
-    ν = convert(Vector{typeof(k)}, ν)
-    σ = convert(Vector{typeof(k)}, σ)
+    ν = convert(Vector{T}, ν)
+    σ = isa(σ, Vector) ? convert(Vector{T}, σ) : convert(T, σ)
     return MLBA(ν, β₀, λₚ, λₙ, γ, σ, A, k, τ)
 end
 
@@ -88,7 +88,7 @@ MLBA(;
     τ = 0.3,
     A = 1.0,
     k = 1.0,
-    σ = fill(1.0, n_alternatives)
+    σ = 1.0
 ) =
     MLBA(ν, β₀, λₚ, λₙ, γ, σ, A, k, τ)
 

src/RDM.jl

@@ -121,9 +121,9 @@ struct RDM{T <: Real} <: AbstractRDM
     τ::T
 end
 
-function RDM(ν, k, A, τ)
+function RDM(ν, k::T, A, τ) where {T}
     _, A, k, τ = promote(ν[1], k, A, τ)
-    ν = convert(Vector{typeof(k)}, ν)
+    ν = convert(Vector{T}, ν)
     return RDM(ν, A, k, τ)
 end
 

src/poisson_race.jl

@@ -36,9 +36,9 @@ struct PoissonRace{T <: Real} <: AbstractPoissonRace
     τ::T
 end
 
-function PoissonRace(ν, α, τ)
+function PoissonRace(ν, α, τ::T) where {T}
     _, τ = promote(ν[1], τ)
-    ν = convert(Vector{typeof(τ)}, ν)
+    ν = convert(Vector{T}, ν)
     return PoissonRace(ν, α, τ)
 end
 

src/stDDM.jl

@@ -64,10 +64,10 @@ mutable struct stDDM{T <: Real} <: AbstractstDDM
     τ::T
 end
 
-function stDDM(ν, σ, s, z, η, ρ, α, τ)
+function stDDM(ν, σ, s, z, η, ρ, α, τ::T) where {T}
     _, σ, s, z, _, ρ, α, τ = promote(ν[1], σ, s, z, η[1], ρ, α, τ)
-    ν = convert(Vector{typeof(τ)}, ν)
-    η = convert(Vector{typeof(τ)}, η)
+    ν = convert(Vector{T}, ν)
+    η = convert(Vector{T}, η)
     return stDDM(ν, σ, s, z, η, ρ, α, τ)
 end