Lens and derivatives

[jw3126]

@tkf's amazing gist shows how we can use lenses to compute partial derivatives.

Often I am in the situation, where I want (partial) gradients, hessians etc. instead. Optimizing a struct is a typical example.

Basic case

We have

struct Widget{T <: Number}
    parameter1::T
    parameter2::T
end

cost(w::Widget) = ...

and want to optimize cost. However many algorithms won't accept the cost function as is,
but need cost(::AbstractVector). Also they often need gradients+hessians in the same form as well.

So to use such algorithms, a way is needed to convert our struct to a vector and back.

Variations:

• Instead of a flat struct, one often has nested structs.
• Some parameters of the struct might be fixed, while others can be optimized.

Currently I write a lot of ad hoc functions to pack unpack the nested fields I need into vector form and back.

I think there should be a better way. I would like to just provide for each of my structs an
isomorphism (e.g. a special type of lens) with some type of vectors and the lens machinery should
be able to do the following for me:

• It should be easy to define the isomorphism for a big new struct, in terms of the isomorphisms of its fields.
• If should be easy to construct lenses that focus on a selection of nested fields.

[tkf]

This is interesting direction! How about providing AbstractVector (or even AbstractArray) interface/wrapper given an original object (e.g., struct) and a vector/array of lenses? Or is equivalent to what you mean by isomorphism? I'm thinking something like this:

using Setfield

struct VectorInterface{T, O, LS <: Tuple} <: AbstractVector{T}
    obj::O
    lenses::LS
end

function VectorInterface(obj::O, lenses) where O
    orig = promote((get(l, obj) for l in lenses)...)
    lenses = tuple(lenses...)
    return VectorInterface{typeof(orig[1]),
                           O,
                           typeof(lenses)}(obj, lenses)
end

Base.length(vi::VectorInterface) = length(vi.lenses)
Base.getindex(vi::VectorInterface, i) = get(vi.lenses[i], vi.obj)
Base.setindex(vi::VectorInterface, x, i) =
    set(Setfield.compose(vi.lenses[i], (@lens _.obj)), vi, x)
Setfield.hassetindex!(::VectorInterface) = false


struct Widget
    x1
    x2
    x3
    y1
    y2
end

# Accessing only x* as a vector:
vi = VectorInterface(Widget(1:5...),
                     ((@lens _.x1),
                      (@lens _.x2),
                      (@lens _.x3),))
@assert vi[1] == 1
@assert vi[2] == 2
@assert vi[3] == 3
let vi2 = @set vi[1] = 111
    @assert vi2[1] == 111
end

[jw3126]

Yeah I am not exactly sure, what I mean by isomorphism :) An isomorphism between two types T,S would be a pair of functions f::T->S and g::S->T which are inverse e.g. g(f(t)) == t for all t::T and f(g(s)) == s for all s::S.
Such an isomorphism f::T->S uniquely defines a lens by the rules

get(lens, t) = f(t)
set(lens, t, s) = g(s)

A lens of this form is an isomorphism. But these definitions might be too narrow.
For instance

struct Point2d{T};x::T,y::T end

Then

f(pt::Point2d) = @SVector [pt.x, pt.y]

is an isomorphism, while

f(pt::Point2d) = [pt.x, pt.y]

is not with the above definition. Anyway enough distracting abstract nonsense.

[jw3126]

VectorInterfaceis very interesting. For flat structs, it does the job!
However I think it is not easy to obtain a VectorInterface for a nested
struct from VectorInterface of its components.
Here is something, that I think works better for deep structs.
Unfortunately this is hacky and has horrible performance. I think it will be some work to clean it up.

using Setfield
using Setfield: compose

struct VectorLens{LS <: Tuple} <: Lens
    scalar_lenses::LS
end

function Setfield.get(lens::VectorLens, obj)
    [get(l, obj) for l in lens.scalar_lenses]
end
function Setfield.set(l::VectorLens, obj, v)
    for (li, vi) in zip(l.scalar_lenses,v)
        obj = set(li, obj, vi)
    end
    obj
end

function veclens(lenses::Lens...)
    scalar_lenses = []
    for l in lenses
        if l isa VectorLens
            append!(scalar_lenses, l.scalar_lenses)
        else
            push!(scalar_lenses, l)
        end
    end
    VectorLens(tuple(scalar_lenses...))
end

function Setfield.compose(vl::VectorLens, lens::Lens)
    VectorLens(map(vl.scalar_lenses) do l
        compose(l, lens)
    end)
end
Base.:(∘)(l1::Lens, l2::Lens) = Setfield.compose(l1,l2)

struct Point{T}
    x::T
    y::T
end
struct Line{T}
    p1::Point{T}
    p2::Point{T}
end

point_lens = veclens((@lens _.x), (@lens _.y))

line_lens = veclens(
    point_lens ∘ (@lens _.p1),
    point_lens ∘ (@lens _.p2))

line = Line(Point(1,2), Point(3,4))


@show get(line_lens, line)
set(line_lens, line, [10, 20, 30, 40])

[tkf]

Cool. I like the compositionality of VectorLens. What is the source of the slowness? Is it because you go back and forth from vector to struct? Maybe emitting a special code for the case when the vector is a static array help (though you can't do the same trick for get [1])? Or another way would be to make VectorInterface on top of VectorLens, so that you don't need to "materialize" a vector.

What you trying to achieve with Line and line_lens looks like related to MultiScaleArrays.jl: https://github.com/JuliaDiffEq/MultiScaleArrays.jl
Did you look at it? Though I guess what you want to achieve here (beyond the Line example) is more general purpose.

BTW, I usually like some small amount of abstract nonsense :)

[1] Maybe we can add interface Setfield.get(::Type, ::Lens, obj) for vectorized lens to decide the container type? Or maybe the vectorized lens has to have the "destination" type?

[jw3126]

There are multiple issues:

1. The get operation is allocating. I think it shouldn't be hard to return an SVector instead. OTOH Vector is better in the case of large structs. Also I am not so keen on depending on StaticArrays? Decisions, decisions :)
2. The veclens operation is not typestable. Again not that hard to make typestable. Its more of an annoyance then a performance problem, since this function is probably not called in a hot loop.
3. The set operation updates each field individually. This is the part, were the real problem lies. We need a way to apply multiple lenses in a single transaction. I have a vague idea how to do this.
   We need a lens that can be used something like this with a single constructor call:

struct T;a;b;c end
l = @lens _.(a,b) # Creates MultiFieldLens{(:a,:b)}()
obj = T(1,2,3)
set(l, obj, (a=10, b=20)) # T(10, 20, 3)
get(l, obj)    # (a=1, b=2)

I think such a lens that allows translating between objects and named tuples would be useful in lots of other situations aswell. For example it could replace the reconstruct function with:

@set obj.(a=1, b=2)

That snipped is a good starting point, but I am not yet sure how to nest this. For instance it would be nice to be able to do something like this:

obj = T(1,2,T(3,4,5))
l = @lens _.(a, c=(a,b))
set(obj, l, (a=10, c=(a=30, b=40)))   # T(10, 2, T(30,40,5))

Syntax and implementation both need more thought in the nested case.

[jw3126]

I was not aware of MultiScaleArrays.jl, it looks interesting. thanks. I need to study it.
The idea to allow to decide on the container type of a vector lens sounds good, I think we need to do that. VectorInterface can indeed a possible output type here.

[tkf]

Thanks for describing the issues. MultiFieldLens would be awesome. I guess supporting destination container types like (named) tuple and static arrays would be equivalent to having VectorLens.

[rafaqz]

It would be great to get something solid that did this. There is also https://github.com/rdeits/Flatten.jl and my awful hack https://github.com/rafaqz/CompositeFieldVectors.jl.

I've been playing with flattening nested structs in various ways but mostly without great performance unless I do bad things with eval, and definitely in less theoretically sound ways than you are describing! But I'm keen to get something working using @generated functions similar to https://github.com/rdeits/Flatten.jl but more flexible.

My main use-case is rebuilding nested structs with Dual numbers for AutoDiff and updating parameters for parameter estimation routines, while using composed nested parameter structures. Often this means extracting floats, adding a Dual number wrapper type, and then rewrapping with the correct units for each field.

I also wrote this package which turned out to be very useful for this kind of thing. For ignoring fields where they need to be on the struct but not in the vector, and separating out units and other wrapper metadata to be applied during rebuild:
https://github.com/rafaqz/MetaFields.jl

[jw3126]

Derivatives of nested structs is also one of my use cases. I think in order to specify what pieces of your struct should be derived named tuples are a must. Currently I am still on julia 0.6 however, where named tuple support is limited. Once I am on 0.7 I will probably toy around with the ideas in this thread again.

Do you currently have a way to derive pieces of nested structs that is nice but slow? Can you provide an example?

[rafaqz]

I currently use this branch:
https://github.com/rafaqz/CompositeFieldVectors.jl/tree/manual

I stopped using the getfield/setfield methods I was originally using with structs of type Any and now I just do a complete rebuild of nested structs/tuples to aim for type stability. Importantly for me it takes the value of units without the unit wrapper because ForwardDiff doesn't handle units.

I prefer not to use named tuples etc, as I really want it to be automatic so I can change my models without specifying anything and they just work, and there are hundreds of possible parameter combinations to write tuples for.

To do that the switch needs to be attached to fields directly, and https://github.com/rafaqz/MetaFields.jl that I built for other things happens to be great at that. It generates the @flattenable macro and methods. There are a few fields where I need to do this:

@flattenable mutable struct Test
   a::Int     | false 
   b::Float64 | true
end

Edit: to clarify, flattenable(Test) would return (false, true)

They are ignored and the current value used in the reconstruct() function. Basically I'm storing the metadata in function methods. I think it's actually type-stable too, as is the flatten() function. But I'm not sure how to make the reconstruct() type stable, it probably needs to be@generated.

The whole thing has a hacky feel to it and works in limited conditions: it wont handle arrays or other data types, just Tuples and Structs with fields in Number.

[rafaqz]

The interesting things are here:
https://github.com/rafaqz/CompositeFieldVectors.jl/blob/manual/src/constructors.jl

I'm currently playing with using @pure for type stability but it's fickle and I don't think it's a good strategy

[jw3126]

Cool I like the @flattenable! I also think its not too ugly to get performance using @generated functions! I usually prefer @generated to @pure, since @pure is so fragile.

In my usecase I need quite some flexibility. For instance one optimization routine may tune some of the fileds and the next may tune some other fields etc. So I don't want to be forced to specify the diff fields at type definition.

If you change annotations with MetaFields.jl, do you need to restart the julia session?

[rafaqz]

That makes sense, I don't really have any fields like that. But you could just define another layer of metafield:

@metafield flattenable2 true

@flattenable2 @flattenable mutable struct Test
   a::Int     | false | true
   b::Float64 | true  | false
end

You can also redefine the methods after type definition, with @reflattenable. Just write the struct again, you can ignore types its just needs the field names and the metafield, but thats really only meant to happen once to overwrite a packages defaults.

I think @generated is the way to go too, its my current code thats ugly! I just need to skill up on writing @generated...

Edit: missed that question: no metaparam macros can be updated at any time, although I might have some @pure function in there right now that might mess with that - but my tests pass with @re- macros updating things

[rafaqz]

Updating: my fork of Flatten.jl now incorporates the @flattenable metafield and is type stable and probably as fast as this can get for flatten/reconstruct of nested structs.

But there are world age issues with @flattenable and @generated, so using Flatten has to be last, which breaks your need for switching fields, you actually do have to restart your session if @generated is used.

So I'm tending to agree that passing in tuples is the safest most generic way of doing this. But it might be slower and harder to get type stable than Flatten.jl as is, although you could use tuples of Val to make it all happen at compile time. For me type stability is more important currently as sensitivity analysis needs thousands of reconstruct and it ends up taking far longer than the actual algorithm.

Edit: Revise actually solves this.

revise(Flattenable)
revise(Flatten)

And updates work. Its admittedly still kind of ugly.

[jw3126]

Nice! Why are there world age issues? Is this because generated function body is not pure? Do you know which part of your generated functions are problematic? I would guess that in this could be solved in principle. Since you have to annotate type definitions anyway, you could try to do the problematic work in the macro, instead generated functions. It might make the code more ugly though.