arviz-devs / ArviZ.jl
1
"""
2
    stack(x::NamedTuple) -> NamedTuple
3
    stack(x::AbstractArray{NamedTuple}) -> NamedTuple
4
    stack(x::AbstractArray{AbstractArray{<:NamedTuple}}) -> NamedTuple
5

6
Given a container of `NamedTuple`s, concatenate them, using the container dimensions as the
7
dimensions of the resulting arrays.
8

9
# Examples
10

11
```@example
12
using ArviZ
13
nchains, ndraws = 4, 100
14
data = [(x=rand(), y=randn(2), z=randn(2, 3)) for _ in 1:nchains, _ in 1:ndraws];
15
stacked_data = ArviZ.stack(data);
16
```
17
"""
18 20
stack(x) = x
19 20
stack(x::AbstractArray{T}) where {T<:Number} = Array(x)
20 20
stack(x::AbstractArray) = stack(stack.(x))
21 0
stack(x::NamedTuple) = (; (k => stack(v) for (k, v) in pairs(x))...)
22
function stack(x::AbstractArray{S}) where {T<:Number,N,S<:AbstractArray{T,N}}
23 20
    ret = Array{T}(undef, (size(x)..., size(x[1])...))
24 20
    @simd for k in keys(x)
25 20
        @inbounds setindex!(ret, x[k], k, (Colon() for _ in 1:N)...)
26
    end
27 20
    return ret
28
end
29
function stack(x::AbstractArray{<:NamedTuple{K}}) where {K}
30 20
    length(x) == 0 && return nothing
31 20
    @inbounds x1 = x[1]
32 20
    ret = NamedTuple()
33 20
    for k in K
34 20
        v = replacemissing.(stack.(getproperty.(x, k)))
35 20
        ret = merge(ret, (k => stack(v),))
36
    end
37 20
    return ret
38
end
39

40
@doc doc"""
41
    from_namedtuple(posterior::NamedTuple; kwargs...) -> InferenceData
42
    from_namedtuple(posterior::Vector{<:NamedTuple}; kwargs...) -> InferenceData
43
    from_namedtuple(posterior::Matrix{<:NamedTuple}; kwargs...) -> InferenceData
44
    from_namedtuple(posterior::Vector{Vector{<:NamedTuple}}; kwargs...) -> InferenceData
45
    from_mcmcchains(
46
        posterior::NamedTuple,
47
        sample_stats::Any,
48
        posterior_predictive::Any,
49
        predictions::Any,
50
        log_likelihood::Any;
51
        kwargs...
52
    ) -> InferenceData
53

54
Convert a `NamedTuple` or container of `NamedTuple`s to an `InferenceData`.
55

56
If containers are passed, they are flattened into a single `NamedTuple` with array elements
57
whose first dimensions correspond to the dimensions of the containers.
58

59
# Arguments
60

61
- `posterior`: The data to be converted. It may be of the following types:
62
    + `::NamedTuple`: The keys are the variable names and the values are arrays with
63
        dimensions `(nchains, ndraws, sizes...)`.
64
    + `::Vector{<:NamedTuple}`: Each element is a `NamedTuple` from a chain with
65
        `Array`/`MonteCarloMeasurements.Particle` values with dimensions
66
        `(ndraws, sizes...)`.
67
    + `::Matrix{<:NamedTuple}`: Each element is a single draw from a single chain, with
68
        array/scalar values with dimensions `sizes`. The dimensions of the matrix container
69
        are `(nchains, ndraws)`
70
    + `::Vector{Vector{<:NamedTuple}}`: The same as the above case.
71

72
# Keywords
73

74
- `posterior_predictive::Any=nothing`: Draws from the posterior predictive distribution
75
- `sample_stats::Any=nothing`: Statistics of the posterior sampling process
76
- `predictions::Any=nothing`: Out-of-sample predictions for the posterior.
77
- `prior::Any=nothing`: Draws from the prior
78
- `prior_predictive::Any=nothing`: Draws from the prior predictive distribution
79
- `sample_stats_prior::Any=nothing`: Statistics of the prior sampling process
80
- `observed_data::Dict{String,Array}=nothing`: Observed data on which the `posterior` is
81
     conditional. It should only contain data which is modeled as a random variable. Keys
82
     are parameter names and values.
83
- `constant_data::Dict{String,Array}=nothing`: Model constants, data included in the model
84
     which is not modeled as a random variable. Keys are parameter names and values.
85
- `predictions_constant_data::Dict{String,Array}=nothing`: Constants relevant to the model
86
     predictions (i.e. new `x` values in a linear regression).
87
- `log_likelihood::Any=nothing`: Pointwise log-likelihood for the data. It is recommended
88
     to use this argument as a dictionary whose keys are observed variable names and whose
89
     values are log likelihood arrays.
90
- `library=nothing`: Name of library that generated the draws
91
- `coords::Dict{String,Vector}=nothing`: Map from named dimension to named indices
92
- `dims::Dict{String,Vector{String}}=nothing`: Map from variable name to names of its
93
     dimensions
94

95
# Returns
96

97
- `InferenceData`: The data with groups corresponding to the provided data
98

99
# Examples
100

101
```@example
102
using ArviZ
103
nchains, ndraws = 2, 10
104

105
data1 = (
106
    x = rand(nchains, ndraws),
107
    y = randn(nchains, ndraws, 2),
108
    z = randn(nchains, ndraws, 3, 2),
109
)
110
idata1 = from_namedtuple(data1)
111

112
data2 = [(x = rand(ndraws), y = randn(ndraws, 2), z = randn(ndraws, 3, 2)) for _ = 1:nchains];
113
idata2 = from_namedtuple(data2)
114

115
data3 = [(x = rand(), y = randn(2), z = randn(3, 2)) for _ = 1:nchains, _ = 1:ndraws];
116
idata3 = from_namedtuple(data3)
117

118
data4 = [[(x = rand(), y = randn(2), z = randn(3, 2)) for _ = 1:ndraws] for _ = 1:nchains];
119
idata4 = from_namedtuple(data4)
120
```
121
"""
122
from_namedtuple
123

124
function from_namedtuple(
125
    posterior,
126
    posterior_predictive,
127
    sample_stats,
128
    predictions,
129
    log_likelihood;
130
    library=nothing,
131
    kwargs...,
132
)
133 20
    all_idata = InferenceData()
134 20
    post_dict = posterior === nothing ? nothing : convert(Dict, posterior)
135 20
    for (group, group_data) in [
136
        :posterior_predictive => posterior_predictive,
137
        :sample_stats => sample_stats,
138
        :predictions => predictions,
139
        :log_likelihood => log_likelihood,
140
    ]
141 20
        group_data === nothing && continue
142 20
        if post_dict !== nothing
143 20
            if group_data isa Union{Symbol,String}
144 0
                group_data = [Symbol(group_data)]
145
            end
146 20
            if group_data isa Union{AbstractVector{Symbol},NTuple{N,Symbol} where {N}}
147 20
                group_data = popsubdict!(post_dict, group_data)
148
            end
149 20
            isempty(group_data) && continue
150
        end
151 20
        group_dataset = convert_to_dataset(group_data; kwargs...)
152 20
        if library !== nothing
153 20
            setattribute!(group_dataset, "inference_library", string(library))
154
        end
155 20
        concat!(all_idata, InferenceData(; group => group_dataset))
156
    end
157

158 20
    (post_dict === nothing || isempty(post_dict)) && return all_idata
159

160 20
    group_dataset = convert_to_dataset(post_dict; kwargs...)
161 20
    if library !== nothing
162 20
        setattribute!(group_dataset, "inference_library", string(library))
163
    end
164 20
    concat!(all_idata, InferenceData(; posterior=group_dataset))
165

166 20
    return all_idata
167
end
168
function from_namedtuple(
169
    posterior::Union{NamedTuple,Nothing}=nothing;
170
    posterior_predictive=nothing,
171
    sample_stats=nothing,
172
    predictions=nothing,
173
    prior=nothing,
174
    prior_predictive=nothing,
175
    sample_stats_prior=nothing,
176
    observed_data=nothing,
177
    constant_data=nothing,
178
    predictions_constant_data=nothing,
179
    log_likelihood=nothing,
180
    library=nothing,
181
    kwargs...,
182
)
183 20
    all_idata = from_namedtuple(
184
        posterior,
185
        posterior_predictive,
186
        sample_stats,
187
        predictions,
188
        log_likelihood;
189
        library=library,
190
        kwargs...,
191
    )
192

193 20
    if any(x -> x !== nothing, [prior, prior_predictive, sample_stats_prior])
194 20
        pre_prior_idata = convert_to_inference_data(
195
            prior;
196
            posterior_predictive=prior_predictive,
197
            sample_stats=sample_stats_prior,
198
            library=library,
199
            kwargs...,
200
        )
201 20
        prior_idata = rekey(
202
            pre_prior_idata,
203
            Dict(
204
                :posterior => :prior,
205
                :posterior_predictive => :prior_predictive,
206
                :sample_stats => :sample_stats_prior,
207
            ),
208
        )
209 20
        concat!(all_idata, prior_idata)
210
    end
211

212 20
    for (group, group_data) in [
213
        :observed_data => observed_data,
214
        :constant_data => constant_data,
215
        :predictions_constant_data => predictions_constant_data,
216
    ]
217 20
        group_data === nothing && continue
218 20
        group_dict = convert(Dict, group_data)
219 20
        group_dataset = convert_to_constant_dataset(group_dict; library=library, kwargs...)
220 20
        concat!(all_idata, InferenceData(; group => group_dataset))
221
    end
222

223 20
    return all_idata
224
end
225
function from_namedtuple(data::AbstractVector{<:NamedTuple}; kwargs...)
226 20
    return from_namedtuple(stack(data); kwargs...)
227
end
228
function from_namedtuple(data::AbstractMatrix{<:NamedTuple}; kwargs...)
229 20
    return from_namedtuple(stack(data); kwargs...)
230
end
231
function from_namedtuple(data::AbstractVector{<:AbstractVector{<:NamedTuple}}; kwargs...)
232 20
    return from_namedtuple(stack(data); kwargs...)
233
end
234

235
"""
236
    convert_to_inference_data(obj::NamedTuple; kwargs...) -> InferenceData
237
    convert_to_inference_data(obj::Vector{<:NamedTuple}; kwargs...) -> InferenceData
238
    convert_to_inference_data(obj::Matrix{<:NamedTuple}; kwargs...) -> InferenceData
239
    convert_to_inference_data(obj::Vector{Vector{<:NamedTuple}}; kwargs...) -> InferenceData
240

241
Convert `obj` to an [`InferenceData`](@ref). See [`from_namedtuple`](@ref) for a description
242
of `obj` possibilities and `kwargs`.
243
"""
244 20
convert_to_inference_data(data::NamedTuple; kwargs...) = from_namedtuple(data; kwargs...)
245
function convert_to_inference_data(data::AbstractVector{<:NamedTuple}; kwargs...)
246 20
    return from_namedtuple(data; kwargs...)
247
end
248
function convert_to_inference_data(data::AbstractMatrix{<:NamedTuple}; kwargs...)
249 20
    return from_namedtuple(data; kwargs...)
250
end
251
function convert_to_inference_data(
252
    data::AbstractVector{<:AbstractVector{<:NamedTuple}}; kwargs...
253
)
254 20
    return from_namedtuple(data; kwargs...)
255
end

Read our documentation on viewing source code .

Loading