Array-like blocks

In a number of VNT use cases, it is necessary to associate multiple indices in a VarNamedTuple with an object that is not necessarily the same number of elements.

Consider, for example, a case where x[1:3] has a Dirichlet prior. With a Dict{VarName,Distribution}, we can do this:

using VarNames

struct Dirichlet{T}
    α::Vector{T}
end

d = Dict{VarName,Dirichlet}(@vn(x[1:3]) => Dirichlet(ones(3)))

Dict{VarName, Main.Dirichlet} with 1 entry:
  x[1:3] => Dirichlet{Float64}([1.0, 1.0, 1.0])

but we incur all the costs associated with the use of a Dict, as described before.

With a VarNamedTuple, we cannot store this directly:

vnt.data.x = some_array
vnt.data.x[1:3] = Dirichlet(ones(3))  # will error

because Dirichlet is not an array, and setindex! will fail. Nor can we write

vnt.data.x = some_array
vnt.data.x[1:3] .= Dirichlet(ones(3))

because although this will not error, it is semantically different: this means that every element x[1], x[2], and x[3] will be assigned the same Dirichlet(ones(3)) object, which is not what we want.

The current solution to this is to use ArrayLikeBlocks, which are thin wrappers around the actual value, but additionally also store the indices used to set the value. The second and third arguments here are the indices (positional and keyword) used to set the value, and the fourth argument is the size of the block.

using VarNames: ArrayLikeBlock

alb = ArrayLikeBlock(Dirichlet(ones(3)), 1:3, (;), (3,))

VarNames.ArrayLikeBlock{Main.Dirichlet{Float64}, UnitRange{Int64}, @NamedTuple{}, Tuple{Int64}}(Main.Dirichlet{Float64}([1.0, 1.0, 1.0]), 1:3, NamedTuple(), (3,))

We then set this ArrayLikeBlock in all the relevant indices. The extra information in the ArrayLikeBlock allows us to forbid partial indexing into it later on. In particular, we want to ensure that users can only retrieve the entire block at once, and not e.g. just x[1] or x[2:3].

Getting and setting: in practice

As a user you should not have to deal with ArrayLikeBlocks directly. Under the hood, templated_setindex!! will automatically wrap values in ArrayLikeBlocks when necessary:

x = zeros(5)
vnt = templated_setindex!!(VarNamedTuple(), Dirichlet(ones(3)), @vn(x[1:3]), x)

VarNamedTuple
└─ x => PartialArray size=(5,) data::Vector{VarNames.ArrayLikeBlock{Main.Dirichlet{Float64}, Tuple{UnitRange{Int64}}, @NamedTuple{}, Tuple{Int64}}}
        ├─ (1,) => VarNames.ArrayLikeBlock{Main.Dirichlet{Float64}, Tuple{UnitRange{Int64}}, @NamedTuple{}, Tuple{Int64}}(Main.Dirichlet{Float64}([1.0, 1.0, 1.0]), (1:3,), NamedTuple(), (3,))
        ├─ (2,) => VarNames.ArrayLikeBlock{Main.Dirichlet{Float64}, Tuple{UnitRange{Int64}}, @NamedTuple{}, Tuple{Int64}}(Main.Dirichlet{Float64}([1.0, 1.0, 1.0]), (1:3,), NamedTuple(), (3,))
        └─ (3,) => VarNames.ArrayLikeBlock{Main.Dirichlet{Float64}, Tuple{UnitRange{Int64}}, @NamedTuple{}, Tuple{Int64}}(Main.Dirichlet{Float64}([1.0, 1.0, 1.0]), (1:3,), NamedTuple(), (3,))

You can access the value again as long as you refer to the full range:

vnt[@vn(x[1:3])]

Main.Dirichlet{Float64}([1.0, 1.0, 1.0])

Because we provided template information, you can access this via any other combination of indexing, as long as it refers to all three indices:

vnt[@vn(x[begin:(end-2)])]

Main.Dirichlet{Float64}([1.0, 1.0, 1.0])

However, if you try to access only part of the block, you will get an error:

julia> vnt[@vn(x[1])]ERROR: ArgumentError: A non-Array value set with a range of indices must be retrieved with the same
range of indices.

Furthermore, if you set a value into any of the indices covered by the block, the entire block is invalidated and thus removed:

struct Normal end # Again a fake struct to avoid importing Distributions.

vnt = templated_setindex!!(vnt, Normal(), @vn(x[2]), x)

VarNamedTuple
└─ x => PartialArray size=(5,) data::Vector{Main.Normal}
        └─ (2,) => Main.Normal()