add view transformations

Author: tpapp

src/aggregation.jl

@@ -106,6 +106,53 @@ function _domain_label(transformation::ArrayTransformation, index::Int)
     _array_domain_label(inner_transformation, dims, index)
 end
 
+####
+#### array view
+####
+
+"""
+$(TYPEDEF)
+
+View of an array with `dims`.
+
+!!! note
+    This feature is experimental, and not part of the stable API; it may disappear or change without
+    relevant changes in SemVer or deprecations. Inner transformations are not supported.
+"""
+struct ViewTransformation{M} <: VectorTransform
+    dims::NTuple{M, Int}
+end
+
+function as(::typeof(view), dims::Tuple{Vararg{Int}})
+    @argcheck all(d -> d ≥ 0, dims) "All dimensions need to be non-negative."
+    ViewTransformation(dims)
+end
+
+as(::typeof(view), dims::Int...) = as(view, dims)
+
+dimension(transformation::ViewTransformation) = prod(transformation.dims)
+
+function transform_with(flag::LogJacFlag, t::ViewTransformation, x, index)
+    index′ = index + dimension(t)
+    y = reshape(@view(x[index:(index′-1)]), t.dims)
+    y, logjac_zero(flag, robust_eltype(x)), index′
+end
+
+function _domain_label(transformation::ViewTransformation, index::Int)
+    @unpack dims = transformation
+    _array_domain_label(asℝ, dims, index)
+end
+
+inverse_eltype(transformation::ViewTransformation, y) = eltype(y)
+
+function inverse_at!(x::AbstractVector, index, transformation::ViewTransformation,
+                     y::AbstractArray)
+    @argcheck size(y) == transformation.dims
+    index′ = index + dimension(transformation)
+    copy!(@view(x[index:(index′-1)]), vec(y))
+    index′
+end
+
 ####
 #### static array
 ####

test/runtests.jl

@@ -644,3 +644,17 @@ end
 @testset "static arrays inference" begin
     @test @inferred transform_with(NOLOGJAC, as(SVector{3, Float64}), zeros(3), 1) == (SVector(0.0, 0.0, 0.0), NOLOGJAC, 4)
 end
+
+@testset "view transformations" begin
+    x = randn(10)
+    t = as((a = asℝ, b = as(view, 2, 4), c = asℝ))
+    y, lj = transform_and_logjac(t, x)
+    @test typeof(y.b) <: AbstractMatrix
+    @test size(y.b) == (2, 4)
+    # test inverse
+    @test inverse(t, y) == x
+    # test that it is a view
+    z = y.b[3]
+    y.b[3] += 1
+    @test x[4] == z + 1
+end