Add lazy Summed (#17)

Author: mtfishman

Project.toml

@@ -1,7 +1,7 @@
 name = "MapBroadcast"
 uuid = "ebd9b9da-f48d-417c-9660-449667d60261"
 authors = ["ITensor developers <[email protected]> and contributors"]
-version = "0.1.9"
+version = "0.1.10"
 
 [deps]
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"

src/MapBroadcast.jl

@@ -1,187 +1,6 @@
 module MapBroadcast
-# Convert broadcast call to map call by capturing array arguments
-# with `map_args` and creating a map function with `map_function`.
-# Logic from https://github.com/Jutho/Strided.jl/blob/v2.0.4/src/broadcast.jl.
 
-using Base.Broadcast:
-  Broadcast, BroadcastStyle, Broadcasted, broadcasted, combine_eltypes, instantiate
-using BlockArrays: mortar
-using Compat: allequal
-using FillArrays: Fill
-
-const WrappedScalarArgs = Union{AbstractArray{<:Any,0},Ref{<:Any}}
-
-# Get the arguments of the map expression that
-# is equivalent to the broadcast expression.
-function map_args(bc::Broadcasted)
-  return map_args_flatten(bc)
-end
-
-function map_args_flatten(bc::Broadcasted, args_rest...)
-  return (map_args_flatten(bc.args...)..., map_args_flatten(args_rest...)...)
-end
-function map_args_flatten(arg1::AbstractArray, args_rest...)
-  return (arg1, map_args_flatten(args_rest...)...)
-end
-map_args_flatten(arg1, args_rest...) = map_args_flatten(args_rest...)
-map_args_flatten() = ()
-
-struct MapFunction{F,Args<:Tuple} <: Function
-  f::F
-  args::Args
-end
-struct Arg end
-
-# Get the function of the map expression that
-# is equivalent to the broadcast expression.
-# Returns a `MapFunction`.
-function map_function(bc::Broadcasted)
-  return map_function_arg(bc)
-end
-map_function_args(args::Tuple{}) = args
-function map_function_args(args::Tuple)
-  return (map_function_arg(args[1]), map_function_args(Base.tail(args))...)
-end
-function map_function_arg(bc::Broadcasted)
-  return MapFunction(bc.f, map_function_args(bc.args))
-end
-map_function_arg(a::WrappedScalarArgs) = a[]
-map_function_arg(a::AbstractArray) = Arg()
-map_function_arg(a) = a
-
-# Evaluate MapFunction
-(f::MapFunction)(args...) = apply_arg(f, args)[1]
-function apply_arg(f::MapFunction, args)
-  mapfunction_args, args′ = apply_args(f.args, args)
-  return f.f(mapfunction_args...), args′
-end
-apply_arg(mapfunction_arg::Arg, args) = args[1], Base.tail(args)
-apply_arg(mapfunction_arg, args) = mapfunction_arg, args
-function apply_args(mapfunction_args::Tuple, args)
-  mapfunction_args1, args′ = apply_arg(mapfunction_args[1], args)
-  mapfunction_args_rest, args′′ = apply_args(Base.tail(mapfunction_args), args′)
-  return (mapfunction_args1, mapfunction_args_rest...), args′′
-end
-apply_args(mapfunction_args::Tuple{}, args) = mapfunction_args, args
-
-is_map_expr_or_arg(arg::AbstractArray) = true
-is_map_expr_or_arg(arg::Any) = false
-function is_map_expr_or_arg(bc::Broadcasted)
-  return all(is_map_expr_or_arg, bc.args)
-end
-function is_map_expr(bc::Broadcasted)
-  return is_map_expr_or_arg(bc)
-end
-
-abstract type ExprStyle end
-struct MapExpr <: ExprStyle end
-struct NotMapExpr <: ExprStyle end
-
-ExprStyle(bc::Broadcasted) = is_map_expr(bc) ? MapExpr() : NotMapExpr()
-
-abstract type AbstractMapped <: Base.AbstractBroadcasted end
-
-function check_shape(::Type{Bool}, args...)
-  return allequal(axes, args)
-end
-function check_shape(args...)
-  if !check_shape(Bool, args...)
-    throw(DimensionMismatch("Mismatched shapes $(axes.(args))."))
-  end
-  return nothing
-end
-
-# Promote the shape of the arguments to support broadcasting
-# over dimensions by expanding singleton dimensions.
-function promote_shape(ax, f, args::AbstractArray...)
-  if allequal((ax, axes.(args)...))
-    return f, args
-  end
-  return f, promote_shape_tile(ax, args...)
-end
-function promote_shape_tile(common_axes, args::AbstractArray...)
-  return map(arg -> tile(arg, common_axes), args)
-end
-
-# Catch the case of zero arguments, like `a .= 2`.
-function promote_shape(ax, f)
-  return identity, (Fill(f(), ax),)
-end
-
-# Extend by repeating value up to length.
-function extend(t::Tuple, value, length)
-  return ntuple(i -> get(t, i, value), length)
-end
-
-# Handles logic of expanding singleton dimensions
-# to match an array shape in broadcasting
-# by tiling or repeating the input array.
-function tile(a::AbstractArray, ax)
-  axes(a) == ax && return a
-  # Must be one-based for now.
-  @assert all(isone, first.(ax))
-  @assert all(isone, first.(axes(a)))
-  ndim = length(ax)
-  size′ = extend(size(a), 1, ndim)
-  a′ = reshape(a, size′)
-  target_size = length.(ax)
-  fillsize = ntuple(ndim) do dim
-    size′[dim] == target_size[dim] && return 1
-    isone(size′[dim]) && return target_size[dim]
-    return throw(DimensionMismatch("Dimensions $(axes(a)) and $ax don't match."))
-  end
-  return mortar(Fill(a′, fillsize))
-end
-
-struct Mapped{Style<:Union{Nothing,BroadcastStyle},Axes,F,Args<:Tuple} <: AbstractMapped
-  style::Style
-  f::F
-  args::Args
-  axes::Axes
-  function Mapped(style, f, args, axes)
-    check_shape(args...)
-    return new{typeof(style),typeof(axes),typeof(f),typeof(args)}(style, f, args, axes)
-  end
-end
-
-function Mapped(bc::Broadcasted)
-  return Mapped(ExprStyle(bc), bc)
-end
-function Mapped(::NotMapExpr, bc::Broadcasted)
-  f = map_function(bc)
-  ax = axes(bc)
-  f, args = promote_shape(ax, f, map_args(bc)...)
-  return Mapped(bc.style, f, args, ax)
-end
-function Mapped(::MapExpr, bc::Broadcasted)
-  f = bc.f
-  ax = axes(bc)
-  f, args = promote_shape(ax, f, bc.args...)
-  return Mapped(bc.style, f, args, ax)
-end
-
-function Broadcast.Broadcasted(m::Mapped)
-  return Broadcasted(m.style, m.f, m.args, m.axes)
-end
-
-function mapped(f, args...)
-  check_shape(args...)
-  return Mapped(broadcasted(f, args...))
-end
-
-Base.similar(m::Mapped, elt::Type) = similar(Broadcasted(m), elt)
-Base.similar(m::Mapped, elt::Type, ax) = similar(Broadcasted(m), elt, ax)
-Base.axes(m::Mapped) = axes(Broadcasted(m))
-# Equivalent to:
-# map(m.f, m.args...)
-# copy(Broadcasted(m))
-function Base.copy(m::Mapped)
-  elt = combine_eltypes(m.f, m.args)
-  # TODO: Handle case of non-concrete eltype.
-  @assert Base.isconcretetype(elt)
-  return copyto!(similar(m, elt), m)
-end
-Base.copyto!(dest::AbstractArray, m::Mapped) = map!(m.f, dest, m.args...)
-Broadcast.instantiate(m::Mapped) = Mapped(instantiate(Broadcasted(m)))
+include("mapped.jl")
+include("linearcombination.jl")
 
 end

src/linearcombination.jl

@@ -0,0 +1,72 @@
+using Base.Broadcast: Broadcasted
+struct LinearCombination{C} <: Function
+  coefficients::C
+end
+coefficients(a::LinearCombination) = a.coefficients
+function (f::LinearCombination)(args...)
+  return mapreduce(*,+,coefficients(f),args)
+end
+
+struct Summed{Style,N,C<:NTuple{N},A<:NTuple{N}}
+  style::Style
+  coefficients::C
+  arguments::A
+end
+Summed(a::Summed) = a
+coefficients(a::Summed) = a.coefficients
+arguments(a::Summed) = a.arguments
+style(a::Summed) = a.style
+LinearCombination(a::Summed) = LinearCombination(coefficients(a))
+using Base.Broadcast: combine_axes
+Base.axes(a::Summed) = combine_axes(a.arguments...)
+function Base.eltype(a::Summed)
+  cts = typeof.(coefficients(a))
+  elts = eltype.(arguments(a))
+  ts = map((ct, elt) -> Base.promote_op(*, ct, elt), cts, elts)
+  return Base.promote_op(+, ts...)
+end
+function Base.getindex(a::Summed, I...)
+  return mapreduce(+, coefficients(a), arguments(a)) do c, a
+    return c * a[I...]
+  end
+end
+using Base.Broadcast: combine_styles
+function Summed(coefficients::Tuple, arguments::Tuple)
+  return Summed(combine_styles(arguments...), coefficients, arguments)
+end
+Summed(a) = Summed((one(eltype(a)),), (a,))
+function Base.:+(a::Summed, b::Summed)
+  return Summed(
+    (coefficients(a)..., coefficients(b)...), (arguments(a)..., arguments(b)...)
+  )
+end
+Base.:-(a::Summed, b::Summed) = a + (-b)
+Base.:+(a::Summed, b::AbstractArray) = a + Summed(b)
+Base.:-(a::Summed, b::AbstractArray) = a - Summed(b)
+Base.:+(a::AbstractArray, b::Summed) = Summed(a) + b
+Base.:-(a::AbstractArray, b::Summed) = Summed(a) - b
+Base.:*(c::Number, a::Summed) = Summed(c .* coefficients(a), arguments(a))
+Base.:*(a::Summed, c::Number) = c * a
+Base.:/(a::Summed, c::Number) = Summed(coefficients(a) ./ c, arguments(a))
+Base.:-(a::Summed) = -one(eltype(a)) * a
+
+Base.similar(a::Summed) = similar(a, eltype(a))
+Base.similar(a::Summed, elt::Type) = similar(a, elt, axes(a))
+Base.similar(a::Summed, ax::Tuple) = similar(a, eltype(a), ax)
+function Base.similar(a::Summed, elt::Type, ax::Tuple)
+  return similar(Broadcasted(a), elt, ax)
+end
+Base.copy(a::Summed) = copyto!(similar(a), a)
+function Base.copyto!(dest::AbstractArray, a::Summed)
+  return copyto!(dest, Broadcasted(a))
+end
+function Broadcast.Broadcasted(a::Summed)
+  f = LinearCombination(a)
+  return Broadcasted(style(a), f, arguments(a), axes(a))
+end
+
+using Base.Broadcast: Broadcast
+Broadcast.BroadcastStyle(a::Type{<:Summed{<:Style}}) where {Style} = Style()
+Broadcast.broadcastable(a::Summed) = a
+Broadcast.materialize(a::Summed) = copy(a)
+Broadcast.materialize!(dest, a::Summed) = copyto!(dest, a)