From 3e04bde6feb7774bd99df4b12c63d21fc4d788fa Mon Sep 17 00:00:00 2001 From: Sam Date: Thu, 13 Jun 2024 22:31:47 +0100 Subject: [PATCH 01/17] add prefixing to ConcatLatentModels --- .../manipulators/CombineLatentModels.jl | 35 +++++++++++++------ .../manipulators/CombineLatentModels.jl | 17 ++++++--- 2 files changed, 37 insertions(+), 15 deletions(-) diff --git a/EpiAware/src/EpiLatentModels/manipulators/CombineLatentModels.jl b/EpiAware/src/EpiLatentModels/manipulators/CombineLatentModels.jl index 4fc59eaef..88b6812aa 100644 --- a/EpiAware/src/EpiLatentModels/manipulators/CombineLatentModels.jl +++ b/EpiAware/src/EpiLatentModels/manipulators/CombineLatentModels.jl @@ -4,9 +4,9 @@ The `CombineLatentModels` struct. This struct is used to combine multiple latent models into a single latent model. # Constructors - -- `CombineLatentModels(models::M) where {M <: AbstractVector{<:AbstractTuringLatentModel}}`: Constructs a `CombineLatentModels` instance with specified models, ensuring that there are at least two models. -- `CombineLatentModels(; models::M) where {M <: AbstractVector{<:AbstractTuringLatentModel}}`: Constructs a `CombineLatentModels` instance with specified models, ensuring that there are at least two models. +- `CombineLatentModels(models::M, prefixes::P) where {M <: AbstractVector{<:AbstractTuringLatentModel}, P <: AbstractVector{<:String}}`: Constructs a `CombineLatentModels` instance with specified models and prefixes, ensuring that there are at least two models and the number of models and prefixes are equal. +- `CombineLatentModels(models::M) where {M <: AbstractVector{<:AbstractTuringLatentModel}}`: Constructs a `CombineLatentModels` instance with specified models, automatically generating prefixes for each model. The +automatic prefixes are of the form `Combine.1`, `Combine.2`, etc. # Examples @@ -17,15 +17,27 @@ latent_model = generate_latent(combined_model, 10) latent_model() ``` " -@kwdef struct CombineLatentModels{M <: AbstractVector{<:AbstractTuringLatentModel}} <: +@kwdef struct CombineLatentModels{ + M <: AbstractVector{<:AbstractTuringLatentModel}, P <: AbstractVector{<:String}} <: AbstractTuringLatentModel "A vector of latent models" models::M + prefixes::P - function CombineLatentModels(models::M) where {M <: - AbstractVector{<:AbstractTuringLatentModel}} + function CombineLatentModels(models::M, + prefixes::P) where { + M <: AbstractVector{<:AbstractTuringLatentModel}, + P <: AbstractVector{<:String}} @assert length(models)>1 "At least two models are required" - return new{AbstractVector{<:AbstractTuringLatentModel}}(models) + @assert length(models)==length(prefixes) "The number of models and prefixes must be equal" + return new{AbstractVector{<:AbstractTuringLatentModel}, AbstractVector{<:String}}( + models, prefixes) + end + + function CombineLatentModels(models::M) where { + M <: AbstractVector{<:AbstractTuringLatentModel}} + prefixes = "Combine." .* string.(1:length(models)) + return CombineLatentModels(models, prefixes) end end @@ -44,19 +56,20 @@ Generate latent variables using a combination of multiple latent models. " @model function EpiAwareBase.generate_latent(latent_models::CombineLatentModels, n) @submodel final_latent, latent_aux = _accumulate_latents( - latent_models.models, 1, fill(0.0, n), [], n, length(latent_models.models)) + latent_models.models, 1, fill(0.0, n), [], n, length(latent_models.models), latent_models.prefixes) return final_latent, (; latent_aux...) end -@model function _accumulate_latents(models, index, acc_latent, acc_aux, n, n_models) +@model function _accumulate_latents( + models, index, acc_latent, acc_aux, n, n_models, prefixes) if index > n_models return acc_latent, (; acc_aux...) else - @submodel latent, new_aux = generate_latent(models[index], n) + @submodel prefix=prefixes[index] latent, new_aux=generate_latent(models[index], n) @submodel updated_latent, updated_aux = _accumulate_latents( models, index + 1, acc_latent .+ latent, - (; acc_aux..., new_aux...), n, n_models) + (; acc_aux..., new_aux...), n, n_models, prefixes) return updated_latent, (; updated_aux...) end end diff --git a/EpiAware/test/EpiLatentModels/manipulators/CombineLatentModels.jl b/EpiAware/test/EpiLatentModels/manipulators/CombineLatentModels.jl index 36ee54341..6f12bfcba 100644 --- a/EpiAware/test/EpiLatentModels/manipulators/CombineLatentModels.jl +++ b/EpiAware/test/EpiLatentModels/manipulators/CombineLatentModels.jl @@ -5,6 +5,11 @@ comb = CombineLatentModels([int, ar]) @test typeof(comb) <: AbstractTuringLatentModel @test comb.models == [int, ar] + @test comb.prefixes == ["Combine.1", "Combine.2"] + + comb = CombineLatentModels([int, ar], ["Int", "AR"]) + @test comb.models == [int, ar] + @test comb.prefixes == ["Int", "AR"] end @testitem "CombineLatentModels generate_latent method works as expected: FixedIntecept + custom" begin @@ -43,7 +48,8 @@ end comb_model = generate_latent(comb, n) # Test constant if conditioning on zero residuals - no_residual_mdl = comb_model | (ϵ_t = zeros(n - 1), ar_init = [0.0]) + no_residual_mdl = comb_model | + (var"Combine.2.ϵ_t" = zeros(n - 1), var"Combine.2.ar_init" = [0.0]) y_const, θ_const = no_residual_mdl() @test all(y_const .== fill(θ_const.intercept, n)) @@ -51,7 +57,9 @@ end # Check against linear regression by conditioning on normal residuals # Generate data fix_intercept = 0.5 - normal_res_mdl = comb_model | (damp_AR = [0.0], σ_AR = 1.0, intercept = fix_intercept) + normal_res_mdl = comb_model | + (var"Combine.2.damp_AR" = [0.0], var"Combine.2.σ_AR" = 1.0, + var"Combine.1.intercept" = fix_intercept) y, θ = normal_res_mdl() # Fit no-slope linear regression as a model test @@ -61,7 +69,8 @@ end end ns_regression_mdl = no_slope_linear_regression(y) | - (damp_AR = [0.0], σ_AR = 1.0, ϵ_t = zeros(n - 1), ar_init = [0.0]) + (var"Combine.2.damp_AR" = [0.0], var"Combine.2.σ_AR" = 1.0, + var"Combine.2.ϵ_t" = zeros(n - 1), var"Combine.2.ar_init" = [0.0]) chain = sample(ns_regression_mdl, NUTS(), 5000, progress = false) # Theoretical posterior distribution for intercept @@ -76,7 +85,7 @@ end post_var = var(int.intercept_prior) / (n * var(int.intercept_prior) + 1) post_dist = Normal(post_mean, sqrt(post_var)) - samples = get(chain, :intercept).intercept |> vec + samples = get(chain, :var"Combine.1.intercept").var"Combine.1.intercept" |> vec ks_test_pval = ExactOneSampleKSTest(samples, post_dist) |> pvalue @test ks_test_pval > 1e-6 end From ed783002a368208abeb0d0ffa8fb8c49e2c053de Mon Sep 17 00:00:00 2001 From: Sam Date: Thu, 13 Jun 2024 23:16:46 +0100 Subject: [PATCH 02/17] start working on concatlatentmodels --- .../manipulators/CombineLatentModels.jl | 1 + .../manipulators/ConcatLatentModels.jl | 53 ++++++++++++------- 2 files changed, 36 insertions(+), 18 deletions(-) diff --git a/EpiAware/src/EpiLatentModels/manipulators/CombineLatentModels.jl b/EpiAware/src/EpiLatentModels/manipulators/CombineLatentModels.jl index 88b6812aa..0f4561749 100644 --- a/EpiAware/src/EpiLatentModels/manipulators/CombineLatentModels.jl +++ b/EpiAware/src/EpiLatentModels/manipulators/CombineLatentModels.jl @@ -22,6 +22,7 @@ latent_model() AbstractTuringLatentModel "A vector of latent models" models::M + "A vector of prefixes for the latent models" prefixes::P function CombineLatentModels(models::M, diff --git a/EpiAware/src/EpiLatentModels/manipulators/ConcatLatentModels.jl b/EpiAware/src/EpiLatentModels/manipulators/ConcatLatentModels.jl index e80b3209d..8c8463b04 100644 --- a/EpiAware/src/EpiLatentModels/manipulators/ConcatLatentModels.jl +++ b/EpiAware/src/EpiLatentModels/manipulators/ConcatLatentModels.jl @@ -5,9 +5,10 @@ This struct is used to concatenate multiple latent models into a single latent m # Constructors -- `ConcatLatentModels(models::M, no_models::Int, dimension_adaptor::Function) where {M <: AbstractVector{<:AbstractTuringLatentModel}}`: Constructs a `ConcatLatentModels` instance with specified models, number of models, and dimension adaptor. -- `ConcatLatentModels(models::M, dimension_adaptor::Function) where {M <: AbstractVector{<:AbstractTuringLatentModel}}`: Constructs a `ConcatLatentModels` instance with specified models and dimension adaptor, ensuring that there are at least two models. The default dimension adaptor is `equal_dimensions`. -- `ConcatLatentModels(; models::M, dimension_adaptor::Function) where {M <: AbstractVector{<:AbstractTuringLatentModel}}`: Constructs a `ConcatLatentModels` instance with specified models and dimension adaptor, ensuring that there are at least two models. The default dimension adaptor is `equal_dimensions`. +- `ConcatLatentModels(models::M, no_models::I, dimension_adaptor::F, prefixes::P) where {M <: AbstractVector{<:AbstractTuringLatentModel}, I <: Int, F <: Function, P <: AbstractVector{String}}`: Constructs a `ConcatLatentModels` instance with specified models, number of models, dimension adaptor, and prefixes. +- `ConcatLatentModels(models::M, dimension_adaptor::F) where {M <: AbstractVector{<:AbstractTuringLatentModel}, F <: Function}`: Constructs a `ConcatLatentModels` instance with specified models and dimension adaptor. The number of models is automatically determined as are the prefixes (of the form `Concat.1`, `Concat.2`, etc.). +- `ConcatLatentModels(models::M; dimension_adaptor::Function, prefixes::P) where {M <: AbstractVector{<:AbstractTuringLatentModel}, P <: AbstractVector{String}}`: Constructs a `ConcatLatentModels` instance with specified models, dimension adaptor, prefixes, and automatically determines the number of models.The default dimension adaptor is `equal_dimensions`. The default prefixes are of the form `Concat.1`, `Concat.2`, etc. +- `ConcatLatentModels(; models::M, dimension_adaptor::Function, prefixes::P) where {M <: AbstractVector{<:AbstractTuringLatentModel}, P <: AbstractVector{String}}`: Constructs a `ConcatLatentModels` instance with specified models, dimension adaptor, prefixes, and automatically determines the number of models. The default dimension adaptor is `equal_dimensions`. The default prefixes are of the form `Concat.1`, `Concat.2`, etc. # Examples @@ -19,7 +20,8 @@ latent_model() ``` " struct ConcatLatentModels{ - M <: AbstractVector{<:AbstractTuringLatentModel}, N <: Int, F <: Function} <: + M <: AbstractVector{<:AbstractTuringLatentModel}, N <: Int, F <: Function, P <: + AbstractVector{String}} <: AbstractTuringLatentModel "A vector of latent models" models::M @@ -27,38 +29,51 @@ struct ConcatLatentModels{ no_models::N "The dimension function for the latent variables. By default this divides the number of latent variables by the number of models and returns a vector of dimensions rounding up the first element and rounding down the rest." dimension_adaptor::F + "A vector of prefixes for the latent models" + prefixes::P function ConcatLatentModels(models::M, no_models::I, - dimension_adaptor::F) where { + dimension_adaptor::F, prefixes::P) where { M <: AbstractVector{<:AbstractTuringLatentModel}, I <: Int, - F <: Function} + F <: Function, P <: AbstractVector{String}} @assert length(models)>1 "At least two models are required" @assert length(models)==no_models "no_models must be equal to the number of models" # check all dimension functions take a single n and return an integer check_dim = dimension_adaptor(no_models, no_models) @assert typeof(check_dim)<:AbstractVector{Int} "Output of dimension_adaptor must be a vector of integers" @assert length(check_dim)==no_models "The vector of dimensions must have the same length as the number of models" - return new{AbstractVector{<:AbstractTuringLatentModel}, Int, Function}( - models, no_models, dimension_adaptor) + @assert length(prefixes)==no_models "The number of models and prefixes must be equal" + return new{ + AbstractVector{<:AbstractTuringLatentModel}, Int, Function, + P <: AbstractVector{String}}( + models, no_models, dimension_adaptor, prefixes) end function ConcatLatentModels(models::M, - dimension_adaptor::Function) where { + dimension_adaptor::Function, prefixes) where { M <: AbstractVector{<:AbstractTuringLatentModel}} - return ConcatLatentModels(models, length(models), dimension_adaptor) + no_models = length(models) + if isempty(prefixes) + new_prefixes = "Concat." .* string.(1:no_models) + else + new_prefixes = prefixes + end + return ConcatLatentModels(models, no_models, dimension_adaptor, new_prefixes) end function ConcatLatentModels(models::M; - dimension_adaptor::Function = equal_dimensions) where { + dimension_adaptor::Function = equal_dimensions, prefixes = "Concat." .* + string.(1:length(models))) where { M <: AbstractVector{<:AbstractTuringLatentModel}} - return ConcatLatentModels(models, dimension_adaptor) + return ConcatLatentModels(models, dimension_adaptor, prefixes) end function ConcatLatentModels(; models::M, - dimension_adaptor::Function = equal_dimensions) where { + dimension_adaptor::Function = equal_dimensions, prefixes = "Concat" .* + string(1:length(models))) where { M <: AbstractVector{<:AbstractTuringLatentModel}} - return ConcatLatentModels(models, dimension_adaptor) + return ConcatLatentModels(models, dimension_adaptor, prefixes) end end @@ -96,20 +111,22 @@ Generate latent variables by concatenating multiple latent models. @assert sum(dims)==n "Sum of dimensions must be equal to the dimension of the latent variables" @submodel final_latent, latent_aux = _concat_latents( - latent_models.models, 1, [], [], dims, latent_models.no_models) + latent_models.models, 1, [], [], dims, latent_models.no_models, latent_models.prefixes) return final_latent, (; latent_aux...) end @model function _concat_latents( - models, index::Int, acc_latent, acc_aux, dims::AbstractVector{<:Int}, n_models::Int) + models, index::Int, acc_latent, acc_aux, + dims::AbstractVector{<:Int}, n_models::Int, prefixes) if index > n_models return acc_latent, (; acc_aux...) else - @submodel latent, new_aux = generate_latent(models[index], dims[index]) + @submodel prefix=prefixes[index] latent, new_aux=generate_latent( + models[index], dims[index]) @submodel updated_latent, updated_aux = _concat_latents( models, index + 1, vcat(acc_latent, latent), - (; acc_aux..., new_aux...), dims, n_models) + (; acc_aux..., new_aux...), dims, n_models, prefixes) return updated_latent, (; updated_aux...) end end From 6c77af552ef2f296735e13fa92f67418e8b4d727 Mon Sep 17 00:00:00 2001 From: Sam Date: Fri, 14 Jun 2024 11:52:20 +0100 Subject: [PATCH 03/17] fix concat constructor --- .../manipulators/ConcatLatentModels.jl | 31 +++++++++---------- 1 file changed, 14 insertions(+), 17 deletions(-) diff --git a/EpiAware/src/EpiLatentModels/manipulators/ConcatLatentModels.jl b/EpiAware/src/EpiLatentModels/manipulators/ConcatLatentModels.jl index 8c8463b04..2b1cc80bb 100644 --- a/EpiAware/src/EpiLatentModels/manipulators/ConcatLatentModels.jl +++ b/EpiAware/src/EpiLatentModels/manipulators/ConcatLatentModels.jl @@ -6,7 +6,7 @@ This struct is used to concatenate multiple latent models into a single latent m # Constructors - `ConcatLatentModels(models::M, no_models::I, dimension_adaptor::F, prefixes::P) where {M <: AbstractVector{<:AbstractTuringLatentModel}, I <: Int, F <: Function, P <: AbstractVector{String}}`: Constructs a `ConcatLatentModels` instance with specified models, number of models, dimension adaptor, and prefixes. -- `ConcatLatentModels(models::M, dimension_adaptor::F) where {M <: AbstractVector{<:AbstractTuringLatentModel}, F <: Function}`: Constructs a `ConcatLatentModels` instance with specified models and dimension adaptor. The number of models is automatically determined as are the prefixes (of the form `Concat.1`, `Concat.2`, etc.). +- `ConcatLatentModels(models::M, dimension_adaptor::F; prefixes::P = \"Concat.\" * string.(1:length(models))) where {M <: AbstractVector{<:AbstractTuringLatentModel}, F <: Function}`: Constructs a `ConcatLatentModels` instance with specified models and dimension adaptor. The number of models is automatically determined as are the prefixes (of the form `Concat.1`, `Concat.2`, etc.) by default. - `ConcatLatentModels(models::M; dimension_adaptor::Function, prefixes::P) where {M <: AbstractVector{<:AbstractTuringLatentModel}, P <: AbstractVector{String}}`: Constructs a `ConcatLatentModels` instance with specified models, dimension adaptor, prefixes, and automatically determines the number of models.The default dimension adaptor is `equal_dimensions`. The default prefixes are of the form `Concat.1`, `Concat.2`, etc. - `ConcatLatentModels(; models::M, dimension_adaptor::Function, prefixes::P) where {M <: AbstractVector{<:AbstractTuringLatentModel}, P <: AbstractVector{String}}`: Constructs a `ConcatLatentModels` instance with specified models, dimension adaptor, prefixes, and automatically determines the number of models. The default dimension adaptor is `equal_dimensions`. The default prefixes are of the form `Concat.1`, `Concat.2`, etc. @@ -21,7 +21,7 @@ latent_model() " struct ConcatLatentModels{ M <: AbstractVector{<:AbstractTuringLatentModel}, N <: Int, F <: Function, P <: - AbstractVector{String}} <: + AbstractVector{<:String}} <: AbstractTuringLatentModel "A vector of latent models" models::M @@ -36,7 +36,7 @@ struct ConcatLatentModels{ no_models::I, dimension_adaptor::F, prefixes::P) where { M <: AbstractVector{<:AbstractTuringLatentModel}, I <: Int, - F <: Function, P <: AbstractVector{String}} + F <: Function, P <: AbstractVector{<:String}} @assert length(models)>1 "At least two models are required" @assert length(models)==no_models "no_models must be equal to the number of models" # check all dimension functions take a single n and return an integer @@ -46,34 +46,31 @@ struct ConcatLatentModels{ @assert length(prefixes)==no_models "The number of models and prefixes must be equal" return new{ AbstractVector{<:AbstractTuringLatentModel}, Int, Function, - P <: AbstractVector{String}}( + AbstractVector{<:String}}( models, no_models, dimension_adaptor, prefixes) end - function ConcatLatentModels(models::M, - dimension_adaptor::Function, prefixes) where { + function ConcatLatentModels(models::M, dimension_adaptor::Function; + prefixes = nothing) where { M <: AbstractVector{<:AbstractTuringLatentModel}} no_models = length(models) - if isempty(prefixes) - new_prefixes = "Concat." .* string.(1:no_models) - else - new_prefixes = prefixes + if isnothing(prefixes) + prefixes = "Concat." .* string.(1:no_models) end - return ConcatLatentModels(models, no_models, dimension_adaptor, new_prefixes) + return ConcatLatentModels(models, no_models, dimension_adaptor, prefixes) end function ConcatLatentModels(models::M; - dimension_adaptor::Function = equal_dimensions, prefixes = "Concat." .* - string.(1:length(models))) where { + dimension_adaptor::Function = equal_dimensions, + prefixes = nothing) where { M <: AbstractVector{<:AbstractTuringLatentModel}} - return ConcatLatentModels(models, dimension_adaptor, prefixes) + return ConcatLatentModels(models, dimension_adaptor; prefixes = prefixes) end function ConcatLatentModels(; models::M, - dimension_adaptor::Function = equal_dimensions, prefixes = "Concat" .* - string(1:length(models))) where { + dimension_adaptor::Function = equal_dimensions, prefixes = nothing) where { M <: AbstractVector{<:AbstractTuringLatentModel}} - return ConcatLatentModels(models, dimension_adaptor, prefixes) + return ConcatLatentModels(models, dimension_adaptor; prefixes = prefixes) end end From 68a8a39b9f3b03388800024e52f60b28033738ac Mon Sep 17 00:00:00 2001 From: Sam Date: Fri, 14 Jun 2024 12:00:11 +0100 Subject: [PATCH 04/17] add tests for ConcatLatentModels --- .../EpiLatentModels/manipulators/ConcatLatentModels.jl | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/EpiAware/test/EpiLatentModels/manipulators/ConcatLatentModels.jl b/EpiAware/test/EpiLatentModels/manipulators/ConcatLatentModels.jl index 3d20a7d75..04adbd975 100644 --- a/EpiAware/test/EpiLatentModels/manipulators/ConcatLatentModels.jl +++ b/EpiAware/test/EpiLatentModels/manipulators/ConcatLatentModels.jl @@ -7,6 +7,7 @@ @test concat.models == [int, ar] @test concat.no_models == 2 @test concat.dimension_adaptor == equal_dimensions + @test concat.prefixes == ["Concat.1", "Concat.2"] function custom_dim(n::Int, no_models::Int)::Vector{Int} return vcat(4, equal_dimensions(n - 4, no_models - 1)) @@ -18,6 +19,14 @@ @test concat_custom.no_models == 2 @test concat_custom.dimension_adaptor == custom_dim @test concat_custom.dimension_adaptor(10, 4) == [4, 2, 2, 2] + @test concat_custom.prefixes == ["Concat.1", "Concat.2"] + + concat_prefix = ConcatLatentModels([int, ar]; prefixes = ["Int", "AR"]) + + @test concat_prefix.models == [int, ar] + @test concat_prefix.no_models == 2 + @test concat_prefix.dimension_adaptor == equal_dimensions + @test concat_prefix.prefixes == ["Int", "AR"] end @testitem "ConcatLatentModels generate_latent method works as expected: FixedIntecept + custom" begin From e4ec573bca1a9d0ba1e471eb2f807ed6d184e82b Mon Sep 17 00:00:00 2001 From: Sam Date: Fri, 14 Jun 2024 12:17:08 +0100 Subject: [PATCH 05/17] add draft ascertainment method --- .../ascertainment/Ascertainment.jl | 26 ++++++++++++++----- EpiAware/test/EpiAwareUtils/turing-methods.jl | 2 +- 2 files changed, 20 insertions(+), 8 deletions(-) diff --git a/EpiAware/src/EpiObsModels/ascertainment/Ascertainment.jl b/EpiAware/src/EpiObsModels/ascertainment/Ascertainment.jl index 0b76e0101..b1d496999 100644 --- a/EpiAware/src/EpiObsModels/ascertainment/Ascertainment.jl +++ b/EpiAware/src/EpiObsModels/ascertainment/Ascertainment.jl @@ -1,27 +1,39 @@ @doc raw" -The `Ascertainment` struct represents an observation model that incorporates ascertainment bias. It is parametrized by two types: `M` which represents the underlying observation model, and `T` which represents the latent model. +The `Ascertainment` struct represents an observation model that incorporates a ascertainment model. # Constructors -- `Ascertainment(model::M, latentmodel::T; link::F = exp) where {M <: AbstractTuringObservationModel, T <: AbstractTuringLatentModel, F <: Function}`: Constructs a new `Ascertainment` object with the specified `model`, `latentmodel`, and `link` function. `link` is a named keyword and defaults to `exp`. -- `Ascertainment(model::M, latentmodel::T, link::F = exp) where {M <: AbstractTuringObservationModel, T <: AbstractTuringLatentModel, F <: Function}`: Constructs a new `Ascertainment` object with the specified `model`, `latentmodel`, and `link` function. +- `Ascertainment(model::M, latentmodel::T, link::F, latent_prefix::P) where {M <: AbstractTuringObservationModel, T <: AbstractTuringLatentModel, F <: Function, P <: String}`: Constructs an `Ascertainment` instance with the specified observation model, latent model, link function, and latent prefix. +- `Ascertainment(; model::M, latentmodel::T, link::F, latent_prefix::P) where {M <: AbstractTuringObservationModel, T <: AbstractTuringLatentModel, F <: Function, P <: String}`: Constructs an `Ascertainment` instance with the specified observation model, latent model, link function, and latent prefix. # Examples ```julia using EpiAware, Turing -obs = Ascertainment(NegativeBinomialError(), FixedIntercept(0.1), x -> x) +obs = Ascertainment(NegativeBinomialError(), FixedIntercept(0.1); link = x -> x) gen_obs = generate_observations(obs, missing, fill(100, 10)) rand(gen_obs) ``` " @kwdef struct Ascertainment{ - M <: AbstractTuringObservationModel, T <: AbstractTuringLatentModel, F <: Function} <: + M <: AbstractTuringObservationModel, T <: AbstractTuringLatentModel, F <: Function, P <: + String} <: AbstractTuringObservationModel "The underlying observation model." model::M "The latent model." latentmodel::T "The link function used to transform the latent model to the observed data." - link::F = x -> x .|> exp + link::F = x -> exp.(x) + latent_prefix::P = "Ascertainment" + + function Ascertainment(model::M, + latentmodel::T; + link::F = x -> exp.(x), + latent_prefix::P = "Ascertainment") where { + M <: AbstractTuringObservationModel, T <: AbstractTuringLatentModel, F <: + Function, P <: + String} + return new{M, T, F, P}(model, latentmodel, link, latent_prefix) + end end @doc raw" @@ -38,7 +50,7 @@ Generates observations based on the `LatentDelay` observation model. - `obs_aux`: Additional observation-related variables. " @model function EpiAwareBase.generate_observations(obs_model::Ascertainment, y_t, Y_t) - @submodel expected_obs_mod, expected_aux = generate_latent( + @submodel prefix=obs_model.latent_prefix expected_obs_mod, expected_aux=generate_latent( obs_model.latentmodel, length(Y_t)) expected_obs = Y_t .* obs_model.link(expected_obs_mod) diff --git a/EpiAware/test/EpiAwareUtils/turing-methods.jl b/EpiAware/test/EpiAwareUtils/turing-methods.jl index bff970cd2..027ce8a96 100644 --- a/EpiAware/test/EpiAwareUtils/turing-methods.jl +++ b/EpiAware/test/EpiAwareUtils/turing-methods.jl @@ -75,7 +75,7 @@ end # Used again in obs model - obs_ascert = Ascertainment(PoissonError(), ar_process, x -> exp.(x)) + obs_ascert = Ascertainment(PoissonError(), ar_process; link = x -> exp.(x)) #Epi model gen_int = [0.2, 0.3, 0.5] From ff62d655bcea8420b15d5d5486bd396448c31f35 Mon Sep 17 00:00:00 2001 From: Sam Date: Fri, 14 Jun 2024 14:05:51 +0100 Subject: [PATCH 06/17] work on tests --- .../test/EpiObsModels/ascertainment/Ascertainment.jl | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/EpiAware/test/EpiObsModels/ascertainment/Ascertainment.jl b/EpiAware/test/EpiObsModels/ascertainment/Ascertainment.jl index 6d583c2a2..f64f26533 100644 --- a/EpiAware/test/EpiObsModels/ascertainment/Ascertainment.jl +++ b/EpiAware/test/EpiObsModels/ascertainment/Ascertainment.jl @@ -4,10 +4,19 @@ return x end - asc = Ascertainment(NegativeBinomialError(), FixedIntercept(0.1), natural) + asc = Ascertainment(NegativeBinomialError(), FixedIntercept(0.1); link = natural) @test asc.model == NegativeBinomialError() @test asc.latentmodel == FixedIntercept(0.1) @test asc.link == natural + @test asc.latent_prefix == "Ascertainment" + + asc_prefix = Ascertainment( + model = NegativeBinomialError(), latentmodel = FixedIntercept(0.1), + link = natural, latent_prefix = "A") + @test asc_prefix.model == NegativeBinomialError() + @test asc_prefix.latentmodel == FixedIntercept(0.1) + @test asc_prefix.link == natural + @test asc_prefix.latent_prefix == "A" end # make a test based on above example From 32209f3a5af58d5dda1d77f7a50249cacff04bb1 Mon Sep 17 00:00:00 2001 From: Sam Date: Fri, 14 Jun 2024 16:30:46 +0100 Subject: [PATCH 07/17] use latent_model not latentmodel --- .../src/EpiLatentModels/modifiers/DiffLatentModel.jl | 6 +++--- .../src/EpiObsModels/ascertainment/Ascertainment.jl | 12 ++++++------ .../test/EpiObsModels/ascertainment/Ascertainment.jl | 6 +++--- 3 files changed, 12 insertions(+), 12 deletions(-) diff --git a/EpiAware/src/EpiLatentModels/modifiers/DiffLatentModel.jl b/EpiAware/src/EpiLatentModels/modifiers/DiffLatentModel.jl index d076fcd9e..addde126a 100644 --- a/EpiAware/src/EpiLatentModels/modifiers/DiffLatentModel.jl +++ b/EpiAware/src/EpiLatentModels/modifiers/DiffLatentModel.jl @@ -18,12 +18,12 @@ terms ``Z_1, \ldots, Z_d`` are inferred. ## Constructors -- `DiffLatentModel(latentmodel, init_prior_distribution::Distribution; d::Int)` - Constructs a `DiffLatentModel` for `d`-fold differencing with `latentmodel` as the +- `DiffLatentModel(latent_model, init_prior_distribution::Distribution; d::Int)` + Constructs a `DiffLatentModel` for `d`-fold differencing with `latent_model` as the undifferenced latent process. All initial terms have common prior `init_prior_distribution`. - `DiffLatentModel(;model, init_priors::Vector{D} where {D <: Distribution})` - Constructs a `DiffLatentModel` for `d`-fold differencing with `latentmodel` as the + Constructs a `DiffLatentModel` for `d`-fold differencing with `latent_model` as the undifferenced latent process. The `d` initial terms have priors given by the vector `init_priors`, therefore `length(init_priors)` sets `d`. diff --git a/EpiAware/src/EpiObsModels/ascertainment/Ascertainment.jl b/EpiAware/src/EpiObsModels/ascertainment/Ascertainment.jl index b1d496999..52137bc60 100644 --- a/EpiAware/src/EpiObsModels/ascertainment/Ascertainment.jl +++ b/EpiAware/src/EpiObsModels/ascertainment/Ascertainment.jl @@ -2,8 +2,8 @@ The `Ascertainment` struct represents an observation model that incorporates a ascertainment model. # Constructors -- `Ascertainment(model::M, latentmodel::T, link::F, latent_prefix::P) where {M <: AbstractTuringObservationModel, T <: AbstractTuringLatentModel, F <: Function, P <: String}`: Constructs an `Ascertainment` instance with the specified observation model, latent model, link function, and latent prefix. -- `Ascertainment(; model::M, latentmodel::T, link::F, latent_prefix::P) where {M <: AbstractTuringObservationModel, T <: AbstractTuringLatentModel, F <: Function, P <: String}`: Constructs an `Ascertainment` instance with the specified observation model, latent model, link function, and latent prefix. +- `Ascertainment(model::M, latent_model::T, link::F, latent_prefix::P) where {M <: AbstractTuringObservationModel, T <: AbstractTuringLatentModel, F <: Function, P <: String}`: Constructs an `Ascertainment` instance with the specified observation model, latent model, link function, and latent prefix. +- `Ascertainment(; model::M, latent_model::T, link::F, latent_prefix::P) where {M <: AbstractTuringObservationModel, T <: AbstractTuringLatentModel, F <: Function, P <: String}`: Constructs an `Ascertainment` instance with the specified observation model, latent model, link function, and latent prefix. # Examples ```julia @@ -20,19 +20,19 @@ rand(gen_obs) "The underlying observation model." model::M "The latent model." - latentmodel::T + latent_model::T "The link function used to transform the latent model to the observed data." link::F = x -> exp.(x) latent_prefix::P = "Ascertainment" function Ascertainment(model::M, - latentmodel::T; + latent_model::T; link::F = x -> exp.(x), latent_prefix::P = "Ascertainment") where { M <: AbstractTuringObservationModel, T <: AbstractTuringLatentModel, F <: Function, P <: String} - return new{M, T, F, P}(model, latentmodel, link, latent_prefix) + return new{M, T, F, P}(model, latent_model, link, latent_prefix) end end @@ -51,7 +51,7 @@ Generates observations based on the `LatentDelay` observation model. " @model function EpiAwareBase.generate_observations(obs_model::Ascertainment, y_t, Y_t) @submodel prefix=obs_model.latent_prefix expected_obs_mod, expected_aux=generate_latent( - obs_model.latentmodel, length(Y_t)) + obs_model.latent_model, length(Y_t)) expected_obs = Y_t .* obs_model.link(expected_obs_mod) diff --git a/EpiAware/test/EpiObsModels/ascertainment/Ascertainment.jl b/EpiAware/test/EpiObsModels/ascertainment/Ascertainment.jl index f64f26533..47056d800 100644 --- a/EpiAware/test/EpiObsModels/ascertainment/Ascertainment.jl +++ b/EpiAware/test/EpiObsModels/ascertainment/Ascertainment.jl @@ -6,15 +6,15 @@ asc = Ascertainment(NegativeBinomialError(), FixedIntercept(0.1); link = natural) @test asc.model == NegativeBinomialError() - @test asc.latentmodel == FixedIntercept(0.1) + @test asc.latent_model == FixedIntercept(0.1) @test asc.link == natural @test asc.latent_prefix == "Ascertainment" asc_prefix = Ascertainment( - model = NegativeBinomialError(), latentmodel = FixedIntercept(0.1), + model = NegativeBinomialError(), latent_model = FixedIntercept(0.1), link = natural, latent_prefix = "A") @test asc_prefix.model == NegativeBinomialError() - @test asc_prefix.latentmodel == FixedIntercept(0.1) + @test asc_prefix.latent_model == FixedIntercept(0.1) @test asc_prefix.link == natural @test asc_prefix.latent_prefix == "A" end From 10177847a4276d1a4c386e8ac8266f7f8358ab37 Mon Sep 17 00:00:00 2001 From: Sam Date: Fri, 14 Jun 2024 22:08:50 +0100 Subject: [PATCH 08/17] use latent_model not latentmodel --- .../src/EpiObsModels/ascertainment/Ascertainment.jl | 10 ++++------ EpiAware/test/EpiAwareUtils/turing-methods.jl | 10 ++++++---- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/EpiAware/src/EpiObsModels/ascertainment/Ascertainment.jl b/EpiAware/src/EpiObsModels/ascertainment/Ascertainment.jl index 52137bc60..b541bd24c 100644 --- a/EpiAware/src/EpiObsModels/ascertainment/Ascertainment.jl +++ b/EpiAware/src/EpiObsModels/ascertainment/Ascertainment.jl @@ -14,9 +14,8 @@ rand(gen_obs) ``` " @kwdef struct Ascertainment{ - M <: AbstractTuringObservationModel, T <: AbstractTuringLatentModel, F <: Function, P <: - String} <: - AbstractTuringObservationModel + M <: AbstractTuringObservationModel, T <: AbstractTuringLatentModel, + F <: Function, P <: String} <: AbstractTuringObservationModel "The underlying observation model." model::M "The latent model." @@ -29,9 +28,8 @@ rand(gen_obs) latent_model::T; link::F = x -> exp.(x), latent_prefix::P = "Ascertainment") where { - M <: AbstractTuringObservationModel, T <: AbstractTuringLatentModel, F <: - Function, P <: - String} + M <: AbstractTuringObservationModel, T <: AbstractTuringLatentModel, + F <: Function, P <: String} return new{M, T, F, P}(model, latent_model, link, latent_prefix) end end diff --git a/EpiAware/test/EpiAwareUtils/turing-methods.jl b/EpiAware/test/EpiAwareUtils/turing-methods.jl index 027ce8a96..30dd60dfc 100644 --- a/EpiAware/test/EpiAwareUtils/turing-methods.jl +++ b/EpiAware/test/EpiAwareUtils/turing-methods.jl @@ -87,19 +87,21 @@ end direct_inf_model = DirectInfections(data, log_init_incidence_prior) #use generate_epiaware - mdl = generate_epiaware(missing, 10, direct_inf_model; latent_model = ar_process, - observation_model = obs_ascert) + mdl = generate_epiaware( + missing, 10, direct_inf_model; + latent_model = ar_process, observation_model = obs_ascert + ) #Check that can sample from model and has appropriate keys/variables θ = rand(mdl) #Both latent and obs processes should be present - @test haskey(θ, Symbol("obs.ϵ_t")) + @test haskey(θ, Symbol("obs.Ascertainment.ϵ_t")) @test haskey(θ, Symbol("latent.ϵ_t")) #Check can sample from model prior chn = sample(mdl, Prior(), 1000; progress = false) @test Symbol("latent.ϵ_t[1]") ∈ keys(chn) - @test Symbol("obs.ϵ_t[1]") ∈ keys(chn) + @test Symbol("obs.Ascertainment.ϵ_t[1]") ∈ keys(chn) #Check that can generate quantities gens = generated_quantities(mdl, chn) @test gens isa Matrix From 88e3e8c9412db2a1ab354df37c04f71728dbb1c8 Mon Sep 17 00:00:00 2001 From: Sam Date: Fri, 14 Jun 2024 23:16:26 +0100 Subject: [PATCH 09/17] fix Ascertainment constructors --- .../ascertainment/Ascertainment.jl | 28 +++++++++++++++---- .../src/EpiObsModels/ascertainment/helpers.jl | 11 ++++---- .../ascertainment/Ascertainment.jl | 5 ++-- .../EpiObsModels/ascertainment/helpers.jl | 3 +- 4 files changed, 33 insertions(+), 14 deletions(-) diff --git a/EpiAware/src/EpiObsModels/ascertainment/Ascertainment.jl b/EpiAware/src/EpiObsModels/ascertainment/Ascertainment.jl index b541bd24c..4fad2f272 100644 --- a/EpiAware/src/EpiObsModels/ascertainment/Ascertainment.jl +++ b/EpiAware/src/EpiObsModels/ascertainment/Ascertainment.jl @@ -8,12 +8,12 @@ The `Ascertainment` struct represents an observation model that incorporates a a # Examples ```julia using EpiAware, Turing -obs = Ascertainment(NegativeBinomialError(), FixedIntercept(0.1); link = x -> x) +obs = Ascertainment(model = NegativeBinomialError(), latent_model = FixedIntercept(0.1)) gen_obs = generate_observations(obs, missing, fill(100, 10)) rand(gen_obs) ``` " -@kwdef struct Ascertainment{ +struct Ascertainment{ M <: AbstractTuringObservationModel, T <: AbstractTuringLatentModel, F <: Function, P <: String} <: AbstractTuringObservationModel "The underlying observation model." @@ -21,8 +21,17 @@ rand(gen_obs) "The latent model." latent_model::T "The link function used to transform the latent model to the observed data." - link::F = x -> exp.(x) - latent_prefix::P = "Ascertainment" + link::F + latent_prefix::P + + function Ascertainment(model::M, + latent_model::T, + link::F, + latent_prefix::P) where { + M <: AbstractTuringObservationModel, T <: AbstractTuringLatentModel, + F <: Function, P <: String} + return new{M, T, F, P}(model, latent_model, link, latent_prefix) + end function Ascertainment(model::M, latent_model::T; @@ -30,7 +39,16 @@ rand(gen_obs) latent_prefix::P = "Ascertainment") where { M <: AbstractTuringObservationModel, T <: AbstractTuringLatentModel, F <: Function, P <: String} - return new{M, T, F, P}(model, latent_model, link, latent_prefix) + return Ascertainment(model, latent_model, link, latent_prefix) + end + + function Ascertainment(; model::M, + latent_model::T, + link::F = x -> exp.(x), + latent_prefix::P = "Ascertainment") where { + M <: AbstractTuringObservationModel, T <: AbstractTuringLatentModel, + F <: Function, P <: String} + return Ascertainment(model, latent_model, link, latent_prefix) end end diff --git a/EpiAware/src/EpiObsModels/ascertainment/helpers.jl b/EpiAware/src/EpiObsModels/ascertainment/helpers.jl index 9dd3101df..8c0de61f2 100644 --- a/EpiAware/src/EpiObsModels/ascertainment/helpers.jl +++ b/EpiAware/src/EpiObsModels/ascertainment/helpers.jl @@ -5,9 +5,10 @@ Create an `Ascertainment` object that models the ascertainment process based on # Arguments - `model::AbstractTuringObservationModel`: The observation model to be used. - `latent_model::AbstractTuringLatentModel`: The latent model to be used. Default is `HierarchicalNormal()` which is a hierarchical normal distribution. -- `link`: The link function to be used. Default is the identity map `x -> x`. This -function is used to transform the latent model _after_ broadcasting to periodic -weekly has been applied. +- `link`: The link function to be used. Default is the identity map `x -> x`. +This function is used to transform the latent model _after_ broadcasting to +periodic weekly has been applied. +- `latent_prefix`: The prefix to be used for the latent model. Default is `\"DayofWeek\"`. # Returns - `Ascertainment`: The `Ascertainment` object that models the ascertainment process based on the day of the week. @@ -24,6 +25,6 @@ rand(gen_obs) " function ascertainment_dayofweek(model::AbstractTuringObservationModel; latent_model::AbstractTuringLatentModel = HierarchicalNormal(), - link = x -> x) - return Ascertainment(model, broadcast_dayofweek(latent_model), link) + link = x -> x, latent_prefix = "DayofWeek") + return Ascertainment(model, broadcast_dayofweek(latent_model), link, latent_prefix) end diff --git a/EpiAware/test/EpiObsModels/ascertainment/Ascertainment.jl b/EpiAware/test/EpiObsModels/ascertainment/Ascertainment.jl index 47056d800..4a9af34e1 100644 --- a/EpiAware/test/EpiObsModels/ascertainment/Ascertainment.jl +++ b/EpiAware/test/EpiObsModels/ascertainment/Ascertainment.jl @@ -10,8 +10,7 @@ @test asc.link == natural @test asc.latent_prefix == "Ascertainment" - asc_prefix = Ascertainment( - model = NegativeBinomialError(), latent_model = FixedIntercept(0.1), + asc_prefix = Ascertainment(NegativeBinomialError(), FixedIntercept(0.1); link = natural, latent_prefix = "A") @test asc_prefix.model == NegativeBinomialError() @test asc_prefix.latent_model == FixedIntercept(0.1) @@ -22,7 +21,7 @@ end # make a test based on above example @testitem "Test Ascertainment generate_observations" begin using Turing, DynamicPPL - obs = Ascertainment(NegativeBinomialError(), FixedIntercept(0.1), x -> x) + obs = Ascertainment(NegativeBinomialError(), FixedIntercept(0.1); link = x -> x) gen_obs = generate_observations(obs, missing, fill(100, 10)) samples = sample(gen_obs, Prior(), 100; progress = false) gen = mapreduce(vcat, generated_quantities(gen_obs, samples)) do gen diff --git a/EpiAware/test/EpiObsModels/ascertainment/helpers.jl b/EpiAware/test/EpiObsModels/ascertainment/helpers.jl index 818e9fe00..0fdae8669 100644 --- a/EpiAware/test/EpiObsModels/ascertainment/helpers.jl +++ b/EpiAware/test/EpiObsModels/ascertainment/helpers.jl @@ -8,7 +8,8 @@ obs_model = generate_observations(obs, missing, fill(incidence_each_ts, nweeks * 7)) dayofweek_effect = [-0.1, 0.1, 0.2, 0.2, -0.4, 0.1, 0] expected_obs = repeat(7 * softmax(dayofweek_effect) .* incidence_each_ts, nweeks) - fix_obs_model = fix(obs_model, (ϵ_t = dayofweek_effect, std = 1)) + fix_obs_model = fix( + obs_model, (var"DayofWeek.ϵ_t" = dayofweek_effect, var"DayofWeek.std" = 1)) gq_expected_obs = fix_obs_model()[2].expected_obs @test expected_obs ≈ gq_expected_obs end From b2dddee528d75aaaa0da6fe591d442c6cd1290da Mon Sep 17 00:00:00 2001 From: Sam Date: Sat, 15 Jun 2024 14:48:38 +0100 Subject: [PATCH 10/17] add tests for prefix_submodel --- EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl | 5 ++-- EpiAware/src/EpiAwareUtils/prefix_submodel.jl | 30 +++++++++++++++++++ .../test/EpiAwareUtils/prefix_submodel.jl | 13 ++++++++ 3 files changed, 46 insertions(+), 2 deletions(-) create mode 100644 EpiAware/src/EpiAwareUtils/prefix_submodel.jl create mode 100644 EpiAware/test/EpiAwareUtils/prefix_submodel.jl diff --git a/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl b/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl index ee8e0d65a..4a17780b4 100644 --- a/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl +++ b/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl @@ -6,7 +6,7 @@ module EpiAwareUtils using ..EpiAwareBase using DataFramesMeta: DataFrame, @rename! -using DynamicPPL: Model, fix, condition +using DynamicPPL: Model, fix, condition, @submodel, @model using MCMCChains: Chains using Random: AbstractRNG using Tables: rowtable @@ -17,12 +17,13 @@ using Distributions, DocStringExtensions, QuadGK, Statistics, Turing export HalfNormal, DirectSample #Export functions -export scan, spread_draws, censored_pmf, get_param_array +export scan, spread_draws, censored_pmf, get_param_array, prefix_submodel include("docstrings.jl") include("censored_pmf.jl") include("HalfNormal.jl") include("scan.jl") +include("prefix_submodel.jl") include("turing-methods.jl") include("DirectSample.jl") include("post-inference.jl") diff --git a/EpiAware/src/EpiAwareUtils/prefix_submodel.jl b/EpiAware/src/EpiAwareUtils/prefix_submodel.jl new file mode 100644 index 000000000..e6441d8e2 --- /dev/null +++ b/EpiAware/src/EpiAwareUtils/prefix_submodel.jl @@ -0,0 +1,30 @@ +@doc raw" +Generate a submodel with an optional prefix. A lightweight wrapper around the `@submodel` macro from DynamicPPL.jl. + +# Arguments + +- `model::AbstractModel`: The model to be used. +- `fn::Function`: The Turing @model function to be applied to the model. +- `prefix::String`: The prefix to be used. If the prefix is an empty string, the submodel is created without a prefix. + +# Returns + +- `submodel`: The returns from the submodel are passed through. + +# Examples + +```julia +using EpiAware +submodel = prefix_submodel(CombineLatentModels([FixedIntercept(0.1), AR()]), generate_latent, \"Test\", 10) +rand(submodel) +``` +" +@model function prefix_submodel( + model::AbstractModel, fn::Function, prefix::String, kwargs...) + if prefix == "" + @submodel submodel = fn(model, kwargs...) + else + @submodel prefix=eval(prefix) submodel=fn(model, kwargs...) + end + return submodel +end diff --git a/EpiAware/test/EpiAwareUtils/prefix_submodel.jl b/EpiAware/test/EpiAwareUtils/prefix_submodel.jl new file mode 100644 index 000000000..2b3b5ef4b --- /dev/null +++ b/EpiAware/test/EpiAwareUtils/prefix_submodel.jl @@ -0,0 +1,13 @@ +@testitem "Test prefix_submodel can handle empty prefix" begin + submodel = prefix_submodel( + CombineLatentModels([FixedIntercept(0.1), AR()]), generate_latent, "", 10) + draw = rand(submodel) + @test typeof(draw[:var"Combine.2.ϵ_t"]) <: AbstractVector +end + +@testitem "Test prefix_submodel can handle non-empty prefix" begin + submodel = prefix_submodel( + CombineLatentModels([FixedIntercept(0.1), AR()]), generate_latent, "Test", 10) + draw = rand(submodel) + @test typeof(draw[:var"Test.Combine.2.ϵ_t"]) <: AbstractVector +end From c0b9ddeec4d47758b9010afee6d274a5a460bd8e Mon Sep 17 00:00:00 2001 From: Sam Date: Sat, 15 Jun 2024 15:05:02 +0100 Subject: [PATCH 11/17] start to add helper strucs --- .../src/EpiLatentModels/EpiLatentModels.jl | 6 ++-- .../modifiers/PrefixLatentModel.jl | 28 ++++++++++++++++++ EpiAware/src/EpiObsModels/EpiObsModels.jl | 12 +++++--- .../{ => modifiers}/LatentDelay.jl | 0 .../modifiers/PrefixObservationModel.jl | 29 +++++++++++++++++++ .../ascertainment/Ascertainment.jl | 0 .../{ => modifiers}/ascertainment/helpers.jl | 0 7 files changed, 68 insertions(+), 7 deletions(-) create mode 100644 EpiAware/src/EpiLatentModels/modifiers/PrefixLatentModel.jl rename EpiAware/src/EpiObsModels/{ => modifiers}/LatentDelay.jl (100%) create mode 100644 EpiAware/src/EpiObsModels/modifiers/PrefixObservationModel.jl rename EpiAware/src/EpiObsModels/{ => modifiers}/ascertainment/Ascertainment.jl (100%) rename EpiAware/src/EpiObsModels/{ => modifiers}/ascertainment/helpers.jl (100%) diff --git a/EpiAware/src/EpiLatentModels/EpiLatentModels.jl b/EpiAware/src/EpiLatentModels/EpiLatentModels.jl index a3e71764b..78252e1ec 100644 --- a/EpiAware/src/EpiLatentModels/EpiLatentModels.jl +++ b/EpiAware/src/EpiLatentModels/EpiLatentModels.jl @@ -5,7 +5,7 @@ module EpiLatentModels using ..EpiAwareBase -using ..EpiAwareUtils: HalfNormal +using ..EpiAwareUtils: HalfNormal, prefix_submodel using LogExpFunctions: softmax @@ -26,7 +26,7 @@ export RepeatEach, RepeatBlock export broadcast_dayofweek, broadcast_weekly, equal_dimensions # Export tools for modifying latent models -export DiffLatentModel, TransformLatentModel +export DiffLatentModel, TransformLatentModel, PrefixLatentModel include("docstrings.jl") include("models/Intercept.jl") @@ -35,9 +35,9 @@ include("models/AR.jl") include("models/HierarchicalNormal.jl") include("modifiers/DiffLatentModel.jl") include("modifiers/TransformLatentModel.jl") +include("modifiers/PrefixLatentModel.jl") include("manipulators/CombineLatentModels.jl") include("manipulators/ConcatLatentModels.jl") - include("manipulators/broadcast/LatentModel.jl") include("manipulators/broadcast/rules.jl") include("manipulators/broadcast/helpers.jl") diff --git a/EpiAware/src/EpiLatentModels/modifiers/PrefixLatentModel.jl b/EpiAware/src/EpiLatentModels/modifiers/PrefixLatentModel.jl new file mode 100644 index 000000000..2af97dc6c --- /dev/null +++ b/EpiAware/src/EpiLatentModels/modifiers/PrefixLatentModel.jl @@ -0,0 +1,28 @@ +@doc raw" + Generate a latent model with a prefix. A lightweight wrapper around `EpiAwareUtils.prefix_submodel`. + + # Constructors + - `PrefixLatentModel(model::M, prefix::P)`: Create a `PrefixLatentModel` with the latent model `model` and the prefix `prefix`. + - `PrefixLatentModel(; model::M, prefix::P)`: Create a `PrefixLatentModel` with the latent model `model` and the prefix `prefix`. + + # Examples + ```julia + using EpiAware + latent_model = PrefixLatentModel(model = HierarchicalNormal(), prefix = \"Test\") + mdl = generate_latent(latent_model, 10) + rand(mdl) + ``` +" +@kwdef struct PrefixLatentModel{M <: AbstractTuringLatentModel, P <: String} <: + AbstractTuringLatentModel + "The latent model" + model::M + "The prefix for the latent model" + prefix::P +end + +@model function EpiAwareBase.generate_latent(latent_model::PrefixLatentModel, n) + @submodel submodel = prefix_submodel( + latent_model.model, generate_latent, latent_model.prefix, n) + return submodel +end diff --git a/EpiAware/src/EpiObsModels/EpiObsModels.jl b/EpiAware/src/EpiObsModels/EpiObsModels.jl index a18d116aa..1370aaab4 100644 --- a/EpiAware/src/EpiObsModels/EpiObsModels.jl +++ b/EpiAware/src/EpiObsModels/EpiObsModels.jl @@ -18,15 +18,19 @@ export PoissonError, NegativeBinomialError export generate_observation_error_priors, observation_error # Observation model modifiers -export LatentDelay, Ascertainment, StackObservationModels +export LatentDelay, Ascertainment, PrefixObservationModel + +# Observation model manipulators +export StackObservationModels # helper functions export ascertainment_dayofweek include("docstrings.jl") -include("LatentDelay.jl") -include("ascertainment/Ascertainment.jl") -include("ascertainment/helpers.jl") +include("modifiers/LatentDelay.jl") +include("modifiers/ascertainment/Ascertainment.jl") +include("modifiers/ascertainment/helpers.jl") +include("modifiers/PrefixObservationModel.jl") include("StackObservationModels.jl") include("ObservationErrorModels/methods.jl") include("ObservationErrorModels/NegativeBinomialError.jl") diff --git a/EpiAware/src/EpiObsModels/LatentDelay.jl b/EpiAware/src/EpiObsModels/modifiers/LatentDelay.jl similarity index 100% rename from EpiAware/src/EpiObsModels/LatentDelay.jl rename to EpiAware/src/EpiObsModels/modifiers/LatentDelay.jl diff --git a/EpiAware/src/EpiObsModels/modifiers/PrefixObservationModel.jl b/EpiAware/src/EpiObsModels/modifiers/PrefixObservationModel.jl new file mode 100644 index 000000000..5be17d99c --- /dev/null +++ b/EpiAware/src/EpiObsModels/modifiers/PrefixObservationModel.jl @@ -0,0 +1,29 @@ +@doc raw" + Generate an observation model with a prefix. A lightweight wrapper around `EpiAwareUtils.prefix_submodel`. + + # Constructors + - `PrefixObservationModel(model::M, prefix::P)`: Create a `PrefixObservationModel` with the observation model `model` and the prefix `prefix`. + - `PrefixObservationModel(; model::M, prefix::P)`: Create a `PrefixObservationModel` with the observation model `model` and the prefix `prefix`. + + # Examples + ```julia + using EpiAware + observation_model = PrefixObservationModel(Poisson(), \"Test\") + obs = generate_observations(observation_model, 10) + rand(obs) + ``` +" +@kwdef struct PrefixObservationModel{M <: AbstractTuringObservationModel, P <: String} <: + AbstractTuringObservationModel + "The observation model" + model::M + "The prefix for the observation model" + prefix::P +end + +@model function EpiAwareBase.generate_observations( + observation_model::PrefixObservationModel, y_t, Y_t) + @submodel submodel = prefix_submodel( + observation_model.model, generate_observations, observation_model.prefix, y_t, Y_t) + return submodel +end diff --git a/EpiAware/src/EpiObsModels/ascertainment/Ascertainment.jl b/EpiAware/src/EpiObsModels/modifiers/ascertainment/Ascertainment.jl similarity index 100% rename from EpiAware/src/EpiObsModels/ascertainment/Ascertainment.jl rename to EpiAware/src/EpiObsModels/modifiers/ascertainment/Ascertainment.jl diff --git a/EpiAware/src/EpiObsModels/ascertainment/helpers.jl b/EpiAware/src/EpiObsModels/modifiers/ascertainment/helpers.jl similarity index 100% rename from EpiAware/src/EpiObsModels/ascertainment/helpers.jl rename to EpiAware/src/EpiObsModels/modifiers/ascertainment/helpers.jl From ed5eb2d0f4cb92aeaf02219442df894a1f4e3f17 Mon Sep 17 00:00:00 2001 From: Sam Date: Tue, 18 Jun 2024 11:50:14 +0100 Subject: [PATCH 12/17] add tests and constructors for Prefix wrappers around prefix_submodel --- EpiAware/src/EpiObsModels/EpiObsModels.jl | 2 +- .../{ => manipulators}/broadcast/LatentModel.jl | 0 .../{ => manipulators}/broadcast/helpers.jl | 0 .../{ => manipulators}/broadcast/rules.jl | 0 .../EpiLatentModels/modifiers/PrefixLatentModel.jl | 13 +++++++++++++ .../EpiObsModels/{ => modifiers}/LatentDelay.jl | 0 .../modifiers/PrefixObservationModel.jl | 14 ++++++++++++++ .../{ => modifiers}/ascertainment/Ascertainment.jl | 0 .../{ => modifiers}/ascertainment/helpers.jl | 0 9 files changed, 28 insertions(+), 1 deletion(-) rename EpiAware/test/EpiLatentModels/{ => manipulators}/broadcast/LatentModel.jl (100%) rename EpiAware/test/EpiLatentModels/{ => manipulators}/broadcast/helpers.jl (100%) rename EpiAware/test/EpiLatentModels/{ => manipulators}/broadcast/rules.jl (100%) create mode 100644 EpiAware/test/EpiLatentModels/modifiers/PrefixLatentModel.jl rename EpiAware/test/EpiObsModels/{ => modifiers}/LatentDelay.jl (100%) create mode 100644 EpiAware/test/EpiObsModels/modifiers/PrefixObservationModel.jl rename EpiAware/test/EpiObsModels/{ => modifiers}/ascertainment/Ascertainment.jl (100%) rename EpiAware/test/EpiObsModels/{ => modifiers}/ascertainment/helpers.jl (100%) diff --git a/EpiAware/src/EpiObsModels/EpiObsModels.jl b/EpiAware/src/EpiObsModels/EpiObsModels.jl index 1370aaab4..95958481a 100644 --- a/EpiAware/src/EpiObsModels/EpiObsModels.jl +++ b/EpiAware/src/EpiObsModels/EpiObsModels.jl @@ -5,7 +5,7 @@ module EpiObsModels using ..EpiAwareBase -using ..EpiAwareUtils: censored_pmf, HalfNormal +using ..EpiAwareUtils: censored_pmf, HalfNormal, prefix_submodel using ..EpiLatentModels: HierarchicalNormal, broadcast_dayofweek diff --git a/EpiAware/test/EpiLatentModels/broadcast/LatentModel.jl b/EpiAware/test/EpiLatentModels/manipulators/broadcast/LatentModel.jl similarity index 100% rename from EpiAware/test/EpiLatentModels/broadcast/LatentModel.jl rename to EpiAware/test/EpiLatentModels/manipulators/broadcast/LatentModel.jl diff --git a/EpiAware/test/EpiLatentModels/broadcast/helpers.jl b/EpiAware/test/EpiLatentModels/manipulators/broadcast/helpers.jl similarity index 100% rename from EpiAware/test/EpiLatentModels/broadcast/helpers.jl rename to EpiAware/test/EpiLatentModels/manipulators/broadcast/helpers.jl diff --git a/EpiAware/test/EpiLatentModels/broadcast/rules.jl b/EpiAware/test/EpiLatentModels/manipulators/broadcast/rules.jl similarity index 100% rename from EpiAware/test/EpiLatentModels/broadcast/rules.jl rename to EpiAware/test/EpiLatentModels/manipulators/broadcast/rules.jl diff --git a/EpiAware/test/EpiLatentModels/modifiers/PrefixLatentModel.jl b/EpiAware/test/EpiLatentModels/modifiers/PrefixLatentModel.jl new file mode 100644 index 000000000..a1f5ebbe4 --- /dev/null +++ b/EpiAware/test/EpiLatentModels/modifiers/PrefixLatentModel.jl @@ -0,0 +1,13 @@ +@testitem "Test PrefixLatentModel constructor" begin + model = PrefixLatentModel(model = HierarchicalNormal(), prefix = "Test") + + @test typeof(model.model) <: HierarchicalNormal + @test model.prefix == "Test" +end + +@testitem "Test generate_latent dispatches to prefix_submodel as expected" begin + model = PrefixLatentModel(model = HierarchicalNormal(), prefix = "Test") + mdl = generate_latent(model, 10) + draw = rand(mdl) + @test typeof(draw[:var"Test.ϵ_t"]) <: AbstractVector +end diff --git a/EpiAware/test/EpiObsModels/LatentDelay.jl b/EpiAware/test/EpiObsModels/modifiers/LatentDelay.jl similarity index 100% rename from EpiAware/test/EpiObsModels/LatentDelay.jl rename to EpiAware/test/EpiObsModels/modifiers/LatentDelay.jl diff --git a/EpiAware/test/EpiObsModels/modifiers/PrefixObservationModel.jl b/EpiAware/test/EpiObsModels/modifiers/PrefixObservationModel.jl new file mode 100644 index 000000000..b46ba54c4 --- /dev/null +++ b/EpiAware/test/EpiObsModels/modifiers/PrefixObservationModel.jl @@ -0,0 +1,14 @@ +@testitem "Test PrefixObservationModel constructor" begin + model = PrefixObservationModel(model = PoissonError(), prefix = "Test") + + @test typeof(model.model) <: PoissonError + @test model.prefix == "Test" +end + +@testitem "Test generate_observations dispatches to prefix_submodel as expected" begin + model = PrefixObservationModel(model = PoissonError(), prefix = "Test") + + mdl = generate_observations(model, missing, 10) + draw = rand(mdl) + @test typeof(draw[:var"Test.y_t[1]"]) <: Int +end diff --git a/EpiAware/test/EpiObsModels/ascertainment/Ascertainment.jl b/EpiAware/test/EpiObsModels/modifiers/ascertainment/Ascertainment.jl similarity index 100% rename from EpiAware/test/EpiObsModels/ascertainment/Ascertainment.jl rename to EpiAware/test/EpiObsModels/modifiers/ascertainment/Ascertainment.jl diff --git a/EpiAware/test/EpiObsModels/ascertainment/helpers.jl b/EpiAware/test/EpiObsModels/modifiers/ascertainment/helpers.jl similarity index 100% rename from EpiAware/test/EpiObsModels/ascertainment/helpers.jl rename to EpiAware/test/EpiObsModels/modifiers/ascertainment/helpers.jl From 326e71b5bc6d369d4c54d3dc06120c961f851614 Mon Sep 17 00:00:00 2001 From: Sam Date: Tue, 18 Jun 2024 12:03:45 +0100 Subject: [PATCH 13/17] reduce custom code by using Prefix constructors --- .../manipulators/CombineLatentModels.jl | 15 ++++++++++----- .../manipulators/ConcatLatentModels.jl | 14 +++++++++----- EpiAware/src/EpiObsModels/EpiObsModels.jl | 2 +- .../modifiers/ascertainment/Ascertainment.jl | 7 +++++-- 4 files changed, 25 insertions(+), 13 deletions(-) diff --git a/EpiAware/src/EpiLatentModels/manipulators/CombineLatentModels.jl b/EpiAware/src/EpiLatentModels/manipulators/CombineLatentModels.jl index 0f4561749..e230f1947 100644 --- a/EpiAware/src/EpiLatentModels/manipulators/CombineLatentModels.jl +++ b/EpiAware/src/EpiLatentModels/manipulators/CombineLatentModels.jl @@ -1,7 +1,7 @@ @doc raw" The `CombineLatentModels` struct. -This struct is used to combine multiple latent models into a single latent model. +This struct is used to combine multiple latent models into a single latent model. If a prefix is supplied wraps each model with `PrefixLatentModel`. # Constructors - `CombineLatentModels(models::M, prefixes::P) where {M <: AbstractVector{<:AbstractTuringLatentModel}, P <: AbstractVector{<:String}}`: Constructs a `CombineLatentModels` instance with specified models and prefixes, ensuring that there are at least two models and the number of models and prefixes are equal. @@ -31,6 +31,11 @@ latent_model() P <: AbstractVector{<:String}} @assert length(models)>1 "At least two models are required" @assert length(models)==length(prefixes) "The number of models and prefixes must be equal" + for i in eachindex(models) + if (prefixes[i] != "") + models[i] = PrefixLatentModel(models[i], prefixes[i]) + end + end return new{AbstractVector{<:AbstractTuringLatentModel}, AbstractVector{<:String}}( models, prefixes) end @@ -57,20 +62,20 @@ Generate latent variables using a combination of multiple latent models. " @model function EpiAwareBase.generate_latent(latent_models::CombineLatentModels, n) @submodel final_latent, latent_aux = _accumulate_latents( - latent_models.models, 1, fill(0.0, n), [], n, length(latent_models.models), latent_models.prefixes) + latent_models.models, 1, fill(0.0, n), [], n, length(latent_models.models)) return final_latent, (; latent_aux...) end @model function _accumulate_latents( - models, index, acc_latent, acc_aux, n, n_models, prefixes) + models, index, acc_latent, acc_aux, n, n_models) if index > n_models return acc_latent, (; acc_aux...) else - @submodel prefix=prefixes[index] latent, new_aux=generate_latent(models[index], n) + @submodel latent, new_aux = generate_latent(models[index], n) @submodel updated_latent, updated_aux = _accumulate_latents( models, index + 1, acc_latent .+ latent, - (; acc_aux..., new_aux...), n, n_models, prefixes) + (; acc_aux..., new_aux...), n, n_models) return updated_latent, (; updated_aux...) end end diff --git a/EpiAware/src/EpiLatentModels/manipulators/ConcatLatentModels.jl b/EpiAware/src/EpiLatentModels/manipulators/ConcatLatentModels.jl index 2b1cc80bb..6d31e58e0 100644 --- a/EpiAware/src/EpiLatentModels/manipulators/ConcatLatentModels.jl +++ b/EpiAware/src/EpiLatentModels/manipulators/ConcatLatentModels.jl @@ -44,6 +44,11 @@ struct ConcatLatentModels{ @assert typeof(check_dim)<:AbstractVector{Int} "Output of dimension_adaptor must be a vector of integers" @assert length(check_dim)==no_models "The vector of dimensions must have the same length as the number of models" @assert length(prefixes)==no_models "The number of models and prefixes must be equal" + for i in eachindex(models) + if (prefixes[i] != "") + models[i] = PrefixLatentModel(models[i], prefixes[i]) + end + end return new{ AbstractVector{<:AbstractTuringLatentModel}, Int, Function, AbstractVector{<:String}}( @@ -108,22 +113,21 @@ Generate latent variables by concatenating multiple latent models. @assert sum(dims)==n "Sum of dimensions must be equal to the dimension of the latent variables" @submodel final_latent, latent_aux = _concat_latents( - latent_models.models, 1, [], [], dims, latent_models.no_models, latent_models.prefixes) + latent_models.models, 1, [], [], dims, latent_models.no_models) return final_latent, (; latent_aux...) end @model function _concat_latents( models, index::Int, acc_latent, acc_aux, - dims::AbstractVector{<:Int}, n_models::Int, prefixes) + dims::AbstractVector{<:Int}, n_models::Int) if index > n_models return acc_latent, (; acc_aux...) else - @submodel prefix=prefixes[index] latent, new_aux=generate_latent( - models[index], dims[index]) + @submodel latent, new_aux = generate_latent(models[index], dims[index]) @submodel updated_latent, updated_aux = _concat_latents( models, index + 1, vcat(acc_latent, latent), - (; acc_aux..., new_aux...), dims, n_models, prefixes) + (; acc_aux..., new_aux...), dims, n_models) return updated_latent, (; updated_aux...) end end diff --git a/EpiAware/src/EpiObsModels/EpiObsModels.jl b/EpiAware/src/EpiObsModels/EpiObsModels.jl index 95958481a..4e35eb802 100644 --- a/EpiAware/src/EpiObsModels/EpiObsModels.jl +++ b/EpiAware/src/EpiObsModels/EpiObsModels.jl @@ -7,7 +7,7 @@ using ..EpiAwareBase using ..EpiAwareUtils: censored_pmf, HalfNormal, prefix_submodel -using ..EpiLatentModels: HierarchicalNormal, broadcast_dayofweek +using ..EpiLatentModels: HierarchicalNormal, broadcast_dayofweek, PrefixLatentModel using Turing, Distributions, DocStringExtensions, SparseArrays diff --git a/EpiAware/src/EpiObsModels/modifiers/ascertainment/Ascertainment.jl b/EpiAware/src/EpiObsModels/modifiers/ascertainment/Ascertainment.jl index 4fad2f272..5e454789c 100644 --- a/EpiAware/src/EpiObsModels/modifiers/ascertainment/Ascertainment.jl +++ b/EpiAware/src/EpiObsModels/modifiers/ascertainment/Ascertainment.jl @@ -1,5 +1,5 @@ @doc raw" -The `Ascertainment` struct represents an observation model that incorporates a ascertainment model. +The `Ascertainment` struct represents an observation model that incorporates a ascertainment model. If a `latent_prefix`is supplied the `latent_model` is wrapped in a call to `PrefixLatentModel`. # Constructors - `Ascertainment(model::M, latent_model::T, link::F, latent_prefix::P) where {M <: AbstractTuringObservationModel, T <: AbstractTuringLatentModel, F <: Function, P <: String}`: Constructs an `Ascertainment` instance with the specified observation model, latent model, link function, and latent prefix. @@ -30,6 +30,9 @@ struct Ascertainment{ latent_prefix::P) where { M <: AbstractTuringObservationModel, T <: AbstractTuringLatentModel, F <: Function, P <: String} + if (latent_prefix != "") + latent_model = PrefixLatentModel(model, prefix) + end return new{M, T, F, P}(model, latent_model, link, latent_prefix) end @@ -66,7 +69,7 @@ Generates observations based on the `LatentDelay` observation model. - `obs_aux`: Additional observation-related variables. " @model function EpiAwareBase.generate_observations(obs_model::Ascertainment, y_t, Y_t) - @submodel prefix=obs_model.latent_prefix expected_obs_mod, expected_aux=generate_latent( + @submodel expected_obs_mod, expected_aux = generate_latent( obs_model.latent_model, length(Y_t)) expected_obs = Y_t .* obs_model.link(expected_obs_mod) From 7278fde4f2a4d9d15e66b26df9dec78638513ae5 Mon Sep 17 00:00:00 2001 From: Sam Date: Tue, 18 Jun 2024 12:09:04 +0100 Subject: [PATCH 14/17] switch stackobservation models to using new prefixwrapper --- EpiAware/src/EpiObsModels/StackObservationModels.jl | 9 +++++++-- .../modifiers/ascertainment/Ascertainment.jl | 2 +- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/EpiAware/src/EpiObsModels/StackObservationModels.jl b/EpiAware/src/EpiObsModels/StackObservationModels.jl index 3712c2f7f..fd84f4e5f 100644 --- a/EpiAware/src/EpiObsModels/StackObservationModels.jl +++ b/EpiAware/src/EpiObsModels/StackObservationModels.jl @@ -1,7 +1,7 @@ @doc raw" A stack of observation models that are looped over to generate observations for -each model in the stack. Note that the model names are used to prefix the parameters in each model (so if I have a model named `cases` and a parameter `y_t`, the parameter in the model will be `cases.y_t`). +each model in the stack. Note that the model names are used to prefix the parameters in each model (so if I have a model named `cases` and a parameter `y_t`, the parameter in the model will be `cases.y_t`). Inside the constructor `PrefixObservationModel` is wrapped around each observation model. ## Constructors @@ -48,6 +48,11 @@ deaths_y_t N <: AbstractString } @assert length(models)==length(model_names) "The number of models and model names must be equal." + for i in eachindex(models) + if (model_names[i] != "") + models[i] = PrefixObservationModel(models[i], model_names[i]) + end + end new{typeof(models), typeof(model_names)}(models, model_names) end @@ -77,7 +82,7 @@ Generate observations from a stack of observation models. Assumes a 1 to 1 mappi obs = () for (model, model_name) in zip(obs_model.models, obs_model.model_names) - @submodel prefix=eval(model_name) obs_tmp=generate_observations( + @submodel obs_tmp = generate_observations( model, y_t[Symbol(model_name)], Y_t[Symbol(model_name)]) obs = obs..., obs_tmp... end diff --git a/EpiAware/src/EpiObsModels/modifiers/ascertainment/Ascertainment.jl b/EpiAware/src/EpiObsModels/modifiers/ascertainment/Ascertainment.jl index 5e454789c..b5fe3041a 100644 --- a/EpiAware/src/EpiObsModels/modifiers/ascertainment/Ascertainment.jl +++ b/EpiAware/src/EpiObsModels/modifiers/ascertainment/Ascertainment.jl @@ -31,7 +31,7 @@ struct Ascertainment{ M <: AbstractTuringObservationModel, T <: AbstractTuringLatentModel, F <: Function, P <: String} if (latent_prefix != "") - latent_model = PrefixLatentModel(model, prefix) + latent_model = PrefixLatentModel(model, latent_prefix) end return new{M, T, F, P}(model, latent_model, link, latent_prefix) end From 9eba553bba664c81ccc7f6fdd71edbdb618cd13c Mon Sep 17 00:00:00 2001 From: Sam Date: Tue, 18 Jun 2024 12:10:25 +0100 Subject: [PATCH 15/17] fix ascertaiment to use latent models --- .../src/EpiObsModels/modifiers/ascertainment/Ascertainment.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/EpiAware/src/EpiObsModels/modifiers/ascertainment/Ascertainment.jl b/EpiAware/src/EpiObsModels/modifiers/ascertainment/Ascertainment.jl index b5fe3041a..88778a021 100644 --- a/EpiAware/src/EpiObsModels/modifiers/ascertainment/Ascertainment.jl +++ b/EpiAware/src/EpiObsModels/modifiers/ascertainment/Ascertainment.jl @@ -31,7 +31,7 @@ struct Ascertainment{ M <: AbstractTuringObservationModel, T <: AbstractTuringLatentModel, F <: Function, P <: String} if (latent_prefix != "") - latent_model = PrefixLatentModel(model, latent_prefix) + latent_model = PrefixLatentModel(latent_model, latent_prefix) end return new{M, T, F, P}(model, latent_model, link, latent_prefix) end From c880a7492ca55a2e1794cf36c8615759c14d50f2 Mon Sep 17 00:00:00 2001 From: Sam Date: Tue, 18 Jun 2024 12:31:39 +0100 Subject: [PATCH 16/17] fix testing issues related to prefix_submodel --- EpiAware/src/EpiObsModels/StackObservationModels.jl | 10 ++++------ .../modifiers/ascertainment/Ascertainment.jl | 3 ++- .../manipulators/CombineLatentModels.jl | 8 ++++++-- .../manipulators/ConcatLatentModels.jl | 11 +++++++---- .../modifiers/ascertainment/Ascertainment.jl | 4 ++-- 5 files changed, 21 insertions(+), 15 deletions(-) diff --git a/EpiAware/src/EpiObsModels/StackObservationModels.jl b/EpiAware/src/EpiObsModels/StackObservationModels.jl index fd84f4e5f..513b68e16 100644 --- a/EpiAware/src/EpiObsModels/StackObservationModels.jl +++ b/EpiAware/src/EpiObsModels/StackObservationModels.jl @@ -48,12 +48,10 @@ deaths_y_t N <: AbstractString } @assert length(models)==length(model_names) "The number of models and model names must be equal." - for i in eachindex(models) - if (model_names[i] != "") - models[i] = PrefixObservationModel(models[i], model_names[i]) - end - end - new{typeof(models), typeof(model_names)}(models, model_names) + wrapped_models = [PrefixObservationModel(models[i], model_names[i]) + for i in eachindex(models)] + new{AbstractVector{<:AbstractTuringObservationModel}, typeof(model_names)}( + wrapped_models, model_names) end function StackObservationModels(models::NamedTuple{ diff --git a/EpiAware/src/EpiObsModels/modifiers/ascertainment/Ascertainment.jl b/EpiAware/src/EpiObsModels/modifiers/ascertainment/Ascertainment.jl index 88778a021..80efa74f2 100644 --- a/EpiAware/src/EpiObsModels/modifiers/ascertainment/Ascertainment.jl +++ b/EpiAware/src/EpiObsModels/modifiers/ascertainment/Ascertainment.jl @@ -33,7 +33,8 @@ struct Ascertainment{ if (latent_prefix != "") latent_model = PrefixLatentModel(latent_model, latent_prefix) end - return new{M, T, F, P}(model, latent_model, link, latent_prefix) + return new{M, AbstractTuringLatentModel, F, P}( + model, latent_model, link, latent_prefix) end function Ascertainment(model::M, diff --git a/EpiAware/test/EpiLatentModels/manipulators/CombineLatentModels.jl b/EpiAware/test/EpiLatentModels/manipulators/CombineLatentModels.jl index 6f12bfcba..8ef32a591 100644 --- a/EpiAware/test/EpiLatentModels/manipulators/CombineLatentModels.jl +++ b/EpiAware/test/EpiLatentModels/manipulators/CombineLatentModels.jl @@ -2,13 +2,17 @@ using Distributions: Normal int = Intercept(Normal(0, 1)) ar = AR() + prefix_int = PrefixLatentModel(int, "Combine.1") + prefix_ar = PrefixLatentModel(ar, "Combine.2") comb = CombineLatentModels([int, ar]) @test typeof(comb) <: AbstractTuringLatentModel - @test comb.models == [int, ar] + @test comb.models == [prefix_int, prefix_ar] @test comb.prefixes == ["Combine.1", "Combine.2"] comb = CombineLatentModels([int, ar], ["Int", "AR"]) - @test comb.models == [int, ar] + prefix_int = PrefixLatentModel(int, "Int") + prefix_ar = PrefixLatentModel(ar, "AR") + @test comb.models == [prefix_int, prefix_ar] @test comb.prefixes == ["Int", "AR"] end diff --git a/EpiAware/test/EpiLatentModels/manipulators/ConcatLatentModels.jl b/EpiAware/test/EpiLatentModels/manipulators/ConcatLatentModels.jl index 04adbd975..100aa710b 100644 --- a/EpiAware/test/EpiLatentModels/manipulators/ConcatLatentModels.jl +++ b/EpiAware/test/EpiLatentModels/manipulators/ConcatLatentModels.jl @@ -2,9 +2,11 @@ using Distributions: Normal int = Intercept(Normal(0, 1)) ar = AR() + prefix_int = PrefixLatentModel(int, "Concat.1") + prefix_ar = PrefixLatentModel(ar, "Concat.2") concat = ConcatLatentModels([int, ar]) @test typeof(concat) <: AbstractTuringLatentModel - @test concat.models == [int, ar] + @test concat.models == [prefix_int, prefix_ar] @test concat.no_models == 2 @test concat.dimension_adaptor == equal_dimensions @test concat.prefixes == ["Concat.1", "Concat.2"] @@ -15,15 +17,16 @@ concat_custom = ConcatLatentModels([int, ar]; dimension_adaptor = custom_dim) - @test concat_custom.models == [int, ar] + @test concat_custom.models == [prefix_int, prefix_ar] @test concat_custom.no_models == 2 @test concat_custom.dimension_adaptor == custom_dim @test concat_custom.dimension_adaptor(10, 4) == [4, 2, 2, 2] @test concat_custom.prefixes == ["Concat.1", "Concat.2"] concat_prefix = ConcatLatentModels([int, ar]; prefixes = ["Int", "AR"]) - - @test concat_prefix.models == [int, ar] + prefix_ar = PrefixLatentModel(ar, "AR") + prefix_int = PrefixLatentModel(int, "Int") + @test concat_prefix.models == [prefix_int, prefix_ar] @test concat_prefix.no_models == 2 @test concat_prefix.dimension_adaptor == equal_dimensions @test concat_prefix.prefixes == ["Int", "AR"] diff --git a/EpiAware/test/EpiObsModels/modifiers/ascertainment/Ascertainment.jl b/EpiAware/test/EpiObsModels/modifiers/ascertainment/Ascertainment.jl index 4a9af34e1..8afdc4b39 100644 --- a/EpiAware/test/EpiObsModels/modifiers/ascertainment/Ascertainment.jl +++ b/EpiAware/test/EpiObsModels/modifiers/ascertainment/Ascertainment.jl @@ -6,14 +6,14 @@ asc = Ascertainment(NegativeBinomialError(), FixedIntercept(0.1); link = natural) @test asc.model == NegativeBinomialError() - @test asc.latent_model == FixedIntercept(0.1) + @test asc.latent_model == PrefixLatentModel(FixedIntercept(0.1), "Ascertainment") @test asc.link == natural @test asc.latent_prefix == "Ascertainment" asc_prefix = Ascertainment(NegativeBinomialError(), FixedIntercept(0.1); link = natural, latent_prefix = "A") @test asc_prefix.model == NegativeBinomialError() - @test asc_prefix.latent_model == FixedIntercept(0.1) + @test asc_prefix.latent_model == PrefixLatentModel(FixedIntercept(0.1), "A") @test asc_prefix.link == natural @test asc_prefix.latent_prefix == "A" end From 1770c26a6e1af156e9ee3214c7f9c4762ec8dc4a Mon Sep 17 00:00:00 2001 From: Sam Date: Tue, 18 Jun 2024 12:36:13 +0100 Subject: [PATCH 17/17] fix final tests --- EpiAware/test/EpiObsModels/StackObservationModels.jl | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/EpiAware/test/EpiObsModels/StackObservationModels.jl b/EpiAware/test/EpiObsModels/StackObservationModels.jl index f0c00fa8c..418892853 100644 --- a/EpiAware/test/EpiObsModels/StackObservationModels.jl +++ b/EpiAware/test/EpiObsModels/StackObservationModels.jl @@ -2,13 +2,17 @@ obs = StackObservationModels( [PoissonError(), NegativeBinomialError()], ["Poisson", "NegativeBinomial"]) - @test obs.models == [PoissonError(), NegativeBinomialError()] + prefix_p = PrefixObservationModel(PoissonError(), "Poisson") + prefix_n = PrefixObservationModel(NegativeBinomialError(), "NegativeBinomial") + @test obs.models == [prefix_p, prefix_n] @test obs.model_names == ["Poisson", "NegativeBinomial"] obs_named = StackObservationModels(( Cases = PoissonError(), Deaths = NegativeBinomialError())) - @test obs_named.models == [PoissonError(), NegativeBinomialError()] + prefix_p = PrefixObservationModel(PoissonError(), "Cases") + prefix_n = PrefixObservationModel(NegativeBinomialError(), "Deaths") + @test obs_named.models == [prefix_p, prefix_n] @test obs_named.model_names == ["Cases", "Deaths"] end