define composite

Author: oxinabox

src/differential_arithmetic.jl

@@ -87,3 +87,75 @@ for T in (:Any,)
     @eval Base.:*(a::AbstractThunk, b::$T) = unthunk(a) * b
     @eval Base.:*(a::$T, b::AbstractThunk) = a * unthunk(b)
 end
+
+################## Composite ##############################################################
+
+# We intentionally do not define, `Base.*(::Composite, ::Composite)` as that is not meaningful
+# In general one doesn't have to represent multiplications of 2 differentials
+# Only of a differential and a scaling factor (generally `Real`)
+Base.*(s::Any, comp::Composite) = map(x->s*x, comp)
+Base.*(comp::Composite, s::Any) = s*comp
+
+function Base.:+(a::Composite{Primal, NamedTuple{an}}, b::Composite{Primal, NamedTuple{bn}}) where Primal
+    # Base on the `merge(:;NamedTuple, ::NamedTuple)` code from Base.
+    # https://github.com/JuliaLang/julia/blob/592748adb25301a45bd6edef3ac0a93eed069852/base/namedtuple.jl#L220-L231
+    if @generated
+        names = Base.merge_names(an, bn)
+        types = Base.merge_types(names, a, b)
+
+        vals = map(names) do field
+            a_field = :(getproperty(:a, $(QuoteNode(field))))
+            b_field = :(getproperty(:b, $(QuoteNode(field))))
+            val_expr = if Base.sym_in(field, an)
+                if Base.sym_in(field, bn)
+                    # in both
+                    :($a_field + $b_field)
+                else
+                    # only in `an`
+                    a_field
+                end
+            else # must be in `b` only
+                b_field
+            end
+        end
+        return :(NamedTuple{$names, $types}(($(vals...),)))
+    else
+        names = Base.merge_names(an, bn)
+        types = Base.merge_types(names, typeof(a), typeof(b))
+        vals = map(names) do field
+            val_expr = if Base.sym_in(field, an)
+                a_field = getproperty(a, field)
+                if Base.sym_in(field, bn)
+                    # in both
+                    b_field = getproperty(a, field)
+                    :($a_field + $b_field)
+                else
+                    # only in `an`
+                    a_field
+                end
+            else # must be in `b` only
+                b_field = getproperty(a, field)
+                b_field
+            end
+        end
+        NamedTuple{names,types}(map(n->getfield(sym_in(n, bn) ? b : a, n), names))
+    end
+    end
+end
+
+# this should not need to be generated, # TODO test that
+function Base.:+(a::Composite{Primal, <:Tuple}, b::Composite{Primal, <:Tuple}) where Primal
+    # TODO: should we even allow it on different lengths?
+    short, long =  length(a) < length(b) ? (a.backing, b.backing) : (b.backing, a.backing)
+    backing = ntuple(length(long)) do ii
+        long_val = getfield(long, ii)
+        if ii <= length(short)
+            short_val = getfield(short, ii)
+            return short_val + long_val
+        else
+            return long_val
+        end
+    end
+
+    return Composite{Primal, typeof(backing)}(backing)
+end

src/differentials.jl

@@ -284,6 +284,66 @@ function itself, when that function is not a closure.
 """
 const NO_FIELDS = DoesNotExist()
 
+
+"""
+    Composite{Primal, T} <: AbstractDifferential
+
+This type represents the differential for a `struct`/`NamedTuple`, or `Tuple`.
+`Primal` is the the corresponding primal type that this is a differential for.
+
+`Composite{Primal}` should have fields (technically properties), that match to a subset of the
+fields of the primal type; and each should be a differential type matching to the primal
+type of that field.
+Fields of the Primal that are not present in the Composite are treated as `Zero`.
+
+`T` is an implementation detail representing the backing datastructure.
+For Tuple it will be a Tuple, and for everything else it will be a `NamedTuple`.
+It should not be passed in by user.
+"""
+struct Composite{Primal, T} <: AbstractDifferential
+    backing::T
+end
+
+
+function Composite{Primal}(;kwargs...) where Primal
+    backing = (; kwargs...)
+    return Composite{Primal, typeof(backing)}(backing)
+end
+
+function Composite{Primal}(args...) where Primal
+    return Composite{Primal, typeof(args)}(args)
+end
+
+function Base.show(io::IO, comp::Composite{Primal})
+    print(io, "Composite{")
+    show(io, Primal)
+    print(io, "}")
+    # allow Tuple or NamedTuple `show` to do the rendering of brackets etc
+    show(io, comp.backing)
+end
+
+#TODO think about this, for if we are missing fields
+#Base.convert(::Type{Primal}, comp::Composite{Primal})
+Base.convert(::Type{<:NamedTuple}, comp::Composite{<:Any, <:NamedTuple}) = comp.backing
+Base.convert(::Type{<:Tuple}, comp::Composite{<:Any, <:Tuple}) = comp.backing
+
+Base.getindex(comp::Composite, idx) = getindex(comp.backing)
+Base.getproperty(comp::Composite, idx) = getproperty(comp.backing, idx)
+Base.propertynames(comp::Composite) = propertynames(comp.backing)
+Base.iterate(comp::Compositem, args...) = iterate(comp.backing, args...)
+Base.length(comp::Composite) = length(comp.backing)
+
+map(f, comp::Composite{Primal, <:Tuple}) where Primal = Composite{Primal}(map(f, comp.backing))
+function map(f, comp::Composite{Primal, <:NamedTuple{L}}) where{Primal, L}
+    vals = map(f, Tuple(comp.backing))
+    named_vals = NamedTuple{L, typeof(vals)}(vals)
+    return Composite{Primal}(named_vals)
+end
+
+Base.conj(comp::Composite{Primal}) = map(conj, comp)
+
+#==============================================================================#
+
 """
     refine_differential(𝒟::Type, der)