[FR] Transform which maps a tuple onto a user-provided struct (#163)

* Sketch out the idea

* Rename to Transform, not Transformation

* Add ConstructionBase

* Try implementing proper methods

* minor fixes

* Add methods for tuple, namedtuple, and scalar; add tests for each, with inverse

* Add printing for tuple and scalar cases

* Make a separate ScalarWrapperTransform

* Eliminate method ambiguity by repeating a StaticArrays method for `as`

* Drop full scalar transform interface for typewrapper, in favor of convenience constructor wrapping in tuple

* Make precompile work

* match namedtuple to struct fields

* Switch semantics: NamedTuples unpack as kwargs

* Implement inverses which do their best; add inference tests by using utility function

* Help 1.10 constprop out a little bit to reach full inference, add a couple more tests

* Formatting improvements

Co-authored-by: Tamas K. Papp <tkpapp@gmail.com>

* Dead code

* Error for different number of fields than tuple fields, rather than returning first n fields

* Test pretty printing on aggregations, make array of scalars go inline

* Add an inline array transform test for longer scalar transform

* Chase coverage on array pretty printing test

* Docs suggestions from review

Co-authored-by: David Müller-Widmann <devmotion@users.noreply.github.com>
Co-authored-by: Tamas K. Papp <tkpapp@gmail.com>

* Remove scalar case from struct transforms

---------

Co-authored-by: Tamás K. Papp <tkpapp@gmail.com>
Co-authored-by: David Müller-Widmann <devmotion@users.noreply.github.com>

Author: 3 people

Project.toml

@@ -7,6 +7,7 @@ version = "0.8.20"
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 CompositionsBase = "a33af91c-f02d-484b-be07-31d278c5ca2b"
+ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -32,6 +33,7 @@ ArgCheck = "1, 2"
 ChangesOfVariables = "0.1"
 Compat = "4.10.0"
 CompositionsBase = "0.1.2"
+ConstructionBase = "1.6.0"
 DocStringExtensions = "0.8, 0.9"
 ForwardDiff = "0.10, 1"
 InverseFunctions = "0.1"

docs/src/index.md

@@ -98,6 +98,39 @@ t_abc = merge(t_a, t_b, t_c)
 t_collision = merge(t_a, as((;a = asℝ₋))) # Will have a = asℝ₋, from rightmost
 ```
 
+In some of these cases, it may be helpful for some of the transformed variables to be wrapped in a user-provided type.
+For example, given
+```julia
+struct Foo
+a
+b
+end
+struct Bar
+c
+d
+end
+```
+it may be useful to transform a flat vector `[1, 2, 3, 4]` into a named tuple like `(e = 1, foo = Foo(2, Bar(3, 4)))`. This can be achieved with a transform like the following:
+```julia
+tb = as(Bar, (Identity(), Identity()))
+tf = as(Foo, (Identity(), tb))
+t = as((; e = Identity(), foo = tf))
+```
+where each instance of `Identity()` here could be replaced with an arbitrary scalar transform.
+
+This relies on [`constructorof` from ConstructionBase](https://juliaobjects.github.io/ConstructionBase.jl/stable/#ConstructionBase.constructorof). If a `Tuple` transform is wrapped in a type this way, transform results will be unpacked and passed to the constructor, hence in the same order as the transform. 
+If a `NamedTuple` transform is used, it will be unpacked as keyword arguments to the constructor of the type. 
+`Tuple` or `NamedTuple` transforms that do not provide proper arguments to the constructor of a given type will simply result in `MethodError`s.
+
+For structs that are constructed from a single argument, wrap a single transform in a `Tuple`, like so:
+```julia
+struct Baz; f; end
+tbaz = as(Baz, (Identity(),))
+```
+
+Inverting these transforms would, in the most general case, require inverting the constructor of the given type, which may itself have several valid dispatches. 
+Rather than conduct a close inspection of user types and their available constructors, inverting a `Tuple`-based struct transform will check that the struct has exactly `n` fields, where `n` is the length of the `Tuple` transform, and use those fields in their struct order; if there are not `n` fields the inversion will error. This will work for any structs with only the default constructor. Inverting a `NamedTuple`-based struct transform will attempt to use struct fields with names matching the names in the `NamedTuple` transform, which fails if the struct does not have matching fields. 
+
 ## Scalar transforms
 
 The symbol `∞` is a placeholder for infinity. It does not correspond to `Inf`, but acts as a placeholder for the correct dispatch. `-∞` is valid.
@@ -131,7 +164,7 @@ TVShift
 TVNeg
 ```
 
-Consistent with common notation, transforms are applied right-to-left; for example, `as(Real, ∞, 3)` is equivalent to `TVShift(3) ∘ TVNeg() ∘ TVExp()`.
+Consistent with common notation, transforms are applied right-to-left; for example, `as(Real, -∞, 3)` is equivalent to `TVShift(3) ∘ TVNeg() ∘ TVExp()`.
 If you are working in an editor where typing Unicode is difficult, `TransformVariables.compose` is also available, as in `TransformVariables.compose(TVScale(5.0), TVNeg(), TVExp())`.
 
 This composition works with any scalar transform in any order, so `TVScale(4) ∘ as(Real, 2, ∞) ∘ TVShift(1e3)` is a valid transform.

src/TransformVariables.jl

@@ -9,6 +9,7 @@ using LinearAlgebra: UpperTriangular, logabsdet, norm, rmul!
 using Random: AbstractRNG, GLOBAL_RNG
 using StaticArrays: MMatrix, SMatrix, SArray, SVector, pushfirst
 using CompositionsBase
+using ConstructionBase
 
 include("utilities.jl")
 include("generic.jl")

src/aggregation.jl

@@ -17,6 +17,9 @@ end
 function _summary_rows(transformation::ArrayTransformation, mime)
     (; inner_transformation, dims) = transformation
     _dims = foldr((a,b) -> "$(string(a))×$(string(b))", dims, init = "")
+    if inner_transformation isa ScalarTransform
+        return _summary_row(transformation, _dims*string(inner_transformation))
+    end
     rows = _summary_row(transformation, _dims)
     for row in _summary_rows(inner_transformation, mime)
         push!(rows, (level = row.level + 1, indices = nothing, repr = row.repr))
@@ -177,7 +180,14 @@ as(SArray{2,3}, asℝ₊, 2, 3)     # transform to a 2x3 SMatrix of positive num
 as(SVector{3})                   # ℝ³ → ℝ³, identity, but an SVector
 ```
 """
-function as(::Type{<:SArray{S}}, inner_transformation = Identity()) where S
+function as(::Type{<:SArray{S}}, inner_transformation::AbstractTransform) where S
+    dim = fieldtypes(S)
+    @argcheck all(x -> x ≥ 1, dim)
+    StaticArrayTransformation{prod(dim),S,typeof(inner_transformation)}(inner_transformation)
+end
+# Repeated with more specific typing to eliminate method ambiguity with 
+# the ScalarWrapperTransform method for `as`
+function as(::Type{<:SArray{S}}, inner_transformation::ScalarTransform = Identity()) where S
     dim = fieldtypes(S)
     @argcheck all(x -> x ≥ 1, dim)
     StaticArrayTransformation{prod(dim),S,typeof(inner_transformation)}(inner_transformation)
@@ -264,7 +274,6 @@ struct TransformTuple{T} <: VectorTransform
     end
 end
 
-
 """
 $(SIGNATURES)
 
@@ -282,20 +291,25 @@ Base.getindex(t::TransformTuple, i::Int) = getindex(_inner(t), i)
 Base.propertynames(t::TransformTuple) = propertynames(_inner(t))
 Base.getproperty(t::TransformTuple, i::Int) = getproperty(_inner(t), i)
 Base.getproperty(t::TransformTuple{<:NamedTuple}, i::Symbol) = getproperty(_inner(t), i)
+
 """
 $(SIGNATURES)
 
 Merge multiple `TransformTuple{<:NamedTuple}` by merging the underlying `NamedTuple`s.
 """
-function Base.merge(t1::TransformTuple{<:NamedTuple}, 
+function Base.merge(t1::TransformTuple{<:NamedTuple},
                     ts::Vararg{TransformTuple{<:NamedTuple}})
     TransformTuple(merge(_inner(t1), map(_inner, ts)...))
 end
 
 function _summary_rows(transformation::TransformTuple, mime)
     inner = _inner(transformation)
     repr1 = string(nameof(typeof(inner)), " of transformations")
-    rows = _summary_row(transformation, repr1)
+    _tuple_summary_rows(repr1, transformation, mime)
+end
+function _tuple_summary_rows(named, transformation::TransformTuple, mime)
+    inner = _inner(transformation)
+    rows = _summary_row(transformation, named)
     _index = 0
     for (key, t) in pairs(inner)
         for row in _summary_rows(t, mime)
@@ -482,3 +496,74 @@ function _domain_label(t::TransformTuple, index::Int)
     end
     error("internal error")
 end
+
+####
+#### type wrapper transformation
+####
+
+"""
+$(TYPEDEF)
+"""
+struct TypeWrapperTransform{T,S} <: VectorTransform
+    inner_transformation::S
+end
+
+function as(::Type{T}, inner_transformation::S) where {T,S<:TransformTuple}
+    @argcheck isstructtype(T)
+    TypeWrapperTransform{T,S}(inner_transformation)
+end
+
+as(::Type{T}, inner_transformation::NTransforms) where T  = as(T, as(inner_transformation))
+
+dimension(t::TypeWrapperTransform) = dimension(t.inner_transformation)
+
+function _summary_rows(transformation::TypeWrapperTransform{T, S}, mime) where {T, S<:TransformTuple}
+    (; inner_transformation) = transformation
+    innerinner = _inner(inner_transformation)
+    name = string("$T wrapper on ", nameof(typeof(innerinner)), " of transformations") 
+    _tuple_summary_rows(name, inner_transformation, mime)
+end
+
+function transform_with(flag::LogJacFlag, t::TypeWrapperTransform{T}, x, index) where T
+    (; inner_transformation) = t
+    y, ℓ, index′ = transform_with(flag, inner_transformation, x, index)
+    ctor = constructorof(T)
+    ctor(y...), ℓ, index′
+end
+function transform_with(flag::LogJacFlag, t::TypeWrapperTransform{C, T}, x, index) where {C, N, T<:TransformTuple{<:NamedTuple{N}}}
+    (; inner_transformation) = t
+    y, ℓ, index′ = transform_with(flag, inner_transformation, x, index)
+    ctor = constructorof(C)
+    ctor(;y...), ℓ, index′
+end
+
+# NamedTuple inner transformations
+function inverse_eltype(t::TypeWrapperTransform{C, S}, ::Type{T}) where {C, N, T<:C, S<:TransformTuple{<:NamedTuple{N}}}
+    used_names = filter(n->n ∈ fieldnames(T), N)
+    types = map(n -> fieldtype(T, n), used_names)
+    inverse_eltype(t.inner_transformation, NamedTuple{used_names,Tuple{types...}})
+end
+function inverse_at!(x, index, t::TypeWrapperTransform{T, S}, y::T) where {T, N, S<:TransformTuple{<:NamedTuple{N}}}
+    yvals = NamedTuple{N}(map(n->(getfield(y,n)), N))
+    inverse_at!(x, index, t.inner_transformation, yvals)
+end
+
+# Regular Tuple inner transformation
+function inverse_eltype(t::TypeWrapperTransform{C, S}, ::Type{T}) where {C, T<:C, S<:TransformTuple}
+    num_args = length(t.inner_transformation)
+    inverse_eltype(t.inner_transformation, Tuple{fieldtypes(T)[begin:num_args]...})
+end
+function inverse_at!(x, index, t::TypeWrapperTransform{T, S}, y::T) where {T, S<:TransformTuple}
+    inner = t.inner_transformation
+    num_args = length(inner)
+    if length(fieldnames(typeof(y))) != num_args
+        throw(ArgumentError("The provided type $T has a different number of fields than the inner transformation, so it cannot be inverted."))
+    end
+    yvals = Tuple(getfield(y, i) for i in 1:num_args)
+    inverse_at!(x, index, inner, yvals)
+end
+
+# Informative error for trying to invert an incompatible type
+function inverse_eltype(t::TypeWrapperTransform{C, S}, ::Type{T}) where {C, T, S<:TransformTuple}
+    throw(ArgumentError("Cannot invert a $T as if it were a $C"))
+end

test/runtests.jl

@@ -412,6 +412,7 @@ end
         @test tt[1] == inner[1]
         @test tt[2] == inner[2]
         @test tt[3] == inner[3]
+        @test tt[1] == inner[1]
         @test_throws BoundsError tt[4]
         @test propertynames(tt) == propertynames(inner)
         @test (@set tt[3] = t2) == as((t1, t2, t2))
@@ -542,6 +543,107 @@ end
 
 end
 
+@testset "transform to custom type" begin
+    
+    struct CustomType{A, B}
+        a::A
+        b::B
+    end
+    @kwdef struct KwCustomType{A, B}
+        a::A
+        b::B
+    end
+    struct MyType{C}
+        c::C
+    end
+    struct YourType
+        d::Float64
+    end
+
+    # From named tuple to type
+    t1 = as((a = asℝ, b = asℝ))
+    t2 = as(CustomType, t1)
+    @test_throws MethodError transform(t2, [1.0, 2.0])
+
+    t1 = as((a = asℝ, b = asℝ))
+    t2 = as(KwCustomType, t1)
+    x = [1.0, 2.0]
+    test_transformation(t2, y -> y isa KwCustomType; N=1, jac=false)
+    @test_throws ArgumentError inverse(t2, [1.0, 2.0])
+    @test_throws ArgumentError inverse(t2, MyType(3.0))
+
+    # Named tuple with different ordering
+    t1 = as((b = asℝ, a = asℝ))    
+    t2 = as(KwCustomType, t1)
+    y = @inferred transform(t2, [1.0, 2.0])
+    @test y == KwCustomType(a = 2.0, b = 1.0) 
+    test_transformation(t2, y -> y isa KwCustomType; N=1, jac=false)
+
+    # Named tuple with wrong number or names of fields
+    t1 = as((;b = asℝ))    
+    t2 = as(KwCustomType, t1)
+    @test_throws UndefKeywordError transform(t2, [1.0])
+    @test inverse(t2, KwCustomType(1.0, 3.0)) == [3.0]
+    t1 = as((a = asℝ, c = asℝ))    
+    t2 = as(KwCustomType, t1)
+    @test_throws UndefKeywordError transform(t2, [1.0, 3.0])
+    @test_throws ArgumentError inverse(t2, KwCustomType(1.0, 3.0))
+    t1 = as((b = asℝ, a = asℝ, c = asℝ))    
+    t2 = as(KwCustomType, t1)
+    @test_throws MethodError transform(t2, [1.0, 2.0, 3.0])
+    @test_throws ArgumentError inverse(t2, KwCustomType(1.0, 3.0))
+
+    # Type with shortened constructor
+    struct MaskedType{A, B}
+        a::A
+        b::B
+    end
+    MaskedType(x) = MaskedType(x, nothing)
+    MaskedType(;b=0.0) = MaskedType(nothing, b)
+    t = as(MaskedType, as((;b=asℝ)))
+    test_transformation(t, y -> y isa MaskedType; N=1, jac=false)
+
+    # No kwarg constructor accepts `a` arg, so errors
+    t = as(MaskedType, as((;a=asℝ)))
+    @test_throws MethodError transform(t, [1.0])
+
+    # When constructor accepts less args than struct has fields, 
+    # inverse errors
+    t = as(MaskedType, (asℝ,))
+    x = [1.0]
+    y = transform(t, x)
+    @test y == MaskedType(1.0, nothing)
+    @test_throws ArgumentError inverse(t, MaskedType(1.0, nothing))
+    # test_transformation(t, y -> y isa MaskedType; N=1, jac=false)
+
+    # Insufficient arguments in provided tuple for constructor
+    # Not specially caught, but good to check
+    t1 = as(CustomType, (asℝ₊,))
+    @test_throws MethodError transform(t1, [1.0])
+
+    # From tuple to type
+    t1 = as(ntuple(i->asℝ₊, Val(2)))
+    t2 = as(CustomType, t1)
+    test_transformation(t2, y -> y isa CustomType; N=1, jac=false)
+    # Trying to invert from another type should error
+    @test_throws ArgumentError inverse(t2, [1.0, 2.0])
+    @test_throws ArgumentError inverse(t2, MyType(3.0))
+
+    # Nested custom types
+    t1 = as(MyType, (asℝ₋,))
+    t2 = as(YourType, (asℝ₋,))
+    t3 = as(KwCustomType, as((a = t1, b = t2)))
+    x = [0.0, -1]
+    y = transform(t3, x)
+    test_transformation(t3, y -> y isa KwCustomType; N=1, jac=false)
+    # Switched order should still work
+    @test y == KwCustomType(;b = YourType(-exp(-1)), a = MyType(-1.0))
+    # Inverting with wrong type
+    @test_throws ArgumentError inverse(t3, KwCustomType(-1.0, YourType(-exp(-1))))
+    @test_throws ArgumentError inverse(t3, CustomType(MyType(-1.0), YourType(-exp(-1))))
+
+end
+
 ####
 #### log density correctness checks
 ####
@@ -842,21 +944,51 @@ end
 end
 
 @testset "pretty printing" begin
+    struct SmallType
+        f1
+    end
+    struct LargerType
+        f2
+        f3
+    end
+    LargerType(;f2, f3) = LargerType(f2, f3)
     t = as((a = asℝ₊,
             b = as(Array, asℝ₋, 3, 3),
             c = corr_cholesky_factor(13),
-            d = as((asℝ, corr_cholesky_factor(SMatrix{3,3}), UnitSimplex(3), unit_vector_norm(4)))))
+            d = as((asℝ, corr_cholesky_factor(SMatrix{3,3}), UnitSimplex(3), unit_vector_norm(4))),
+            e = as(LargerType, as((f3 = as(SmallType, (asℝ₊,)), f2 = as𝕀))),
+            ))
     repr_t = """
-[1:97] NamedTuple of transformations
+[1:100] NamedTuple of transformations
   [1:1] :a → asℝ₊
   [2:10] :b → 3×3×asℝ₋
   [11:88] :c → 13×13 correlation cholesky factor
-  [89:97] :d → Tuple of transformations
-    [98:98] 1 → asℝ
-    [108:110] 2 → SMatrix{3,3} correlation cholesky factor
-    [120:121] 3 → 3 element unit simplex transformation
-    [131:133] 4 → 4 element (unit vector, norm) transformation"""
-    repr(MIME("text/plain"), t) == repr_t
+  [89:98] :d → Tuple of transformations
+    [89:89] 1 → asℝ
+    [90:92] 2 → SMatrix{3,3} correlation cholesky factor
+    [93:94] 3 → 3 element unit simplex transformation
+    [95:98] 4 → 4 element (unit vector, norm) transformation
+  [99:100] :e → LargerType wrapper on NamedTuple of transformations
+    [99:99] :f3 → SmallType wrapper on Tuple of transformations
+      [99:99] 1 → asℝ₊
+    [100:100] :f2 → as𝕀"""
+    @test repr(MIME("text/plain"), t) == repr_t
+
+    t = as((as(Array, asℝ₊, 3),
+            as(Array, asℝ₋, 3, 3),
+            as(Array, TVScale(5.0) ∘ asℝ₋, 3, 3, 3),
+            as(Array, as((a=asℝ₊, b = as𝕀)), 3)
+            ))
+    repr_t = """
+[1:45] Tuple of transformations
+  [1:3] 1 → 3×asℝ₊
+  [4:12] 2 → 3×3×asℝ₋
+  [13:39] 3 → 3×3×3×TVScale(5.0) ∘ TVNeg() ∘ asℝ₊
+  [40:45] 4 → 3×
+    NamedTuple of transformations
+      :a → asℝ₊
+      :b → as𝕀"""
+    @test repr(MIME("text/plain"), t) == repr_t
 end
 
 @testset "print ∞" begin