Port map special handling

use fieldcount

taylor bundle not taylor tangent

Author: oxinabox

src/stage1/forward.jl

@@ -194,12 +194,9 @@ struct FwdMap{N, T<:AbstractTangentBundle{N}}
 end
 (f::FwdMap{N})(args::AbstractTangentBundle{N}...) where {N} = ∂☆{N}()(f.f, args...)
 
-#==
-# TODO port this to TaylorBundle over composite structure
-function (::∂☆{N})(::ZeroBundle{N, typeof(map)}, f::ATB{N}, tup::CompositeBundle{N, <:Tuple}) where {N}
-    ∂vararg{N}()(map(FwdMap(f), tup.tup)...)
+function (::∂☆{N})(::ZeroBundle{N, typeof(map)}, f::ATB{N}, tup::TaylorBundle{N, <:Tuple}) where {N}
+    ∂vararg{N}()(map(FwdMap(f), destructure(tup))...)
 end
-==#
 
 function (::∂☆{N})(::ZeroBundle{N, typeof(map)}, f::ATB{N}, args::ATB{N, <:AbstractArray}...) where {N}
     # TODO: This could do an inplace map! to avoid the extra rebundling

src/tangent.jl

@@ -228,6 +228,18 @@ function Base.getindex(tb::TaylorBundle, tti::CanonicalTangentIndex)
     tb.tangent.coeffs[count_ones(tti.i)]
 end
 
+"for a TaylorTangent{N, <:Tuple} this breaks it up unto 1 TaylorTangent{N} for each element of the primal tuple"
+function destructure(r::TaylorBundle{N, B}) where {N, B<:Tuple}
+    return ntuple(fieldcount(B)) do field_ii
+        the_primal = primal(r)[field_ii]
+        the_partials = ntuple(N) do  order_ii
+            partial(r, order_ii)[field_ii]
+        end
+        return TaylorBundle{N}(the_primal, the_partials)
+    end
+end
+
+
 function truncate(tt::TaylorTangent, order::Val{N}) where {N}
     TaylorTangent(tt.coeffs[1:N])
 end