Allow repeat-like patterns (#51)

* allow broadcast over indices not present on RHS

* partial fix for function on left, e.g. rand(3)[i] = i

* a test

* fix tests, docs

Author: mcabbott

docs/Project.toml

@@ -2,6 +2,7 @@
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 Einsum = "b7d42ee7-0b51-5a75-98ca-779d3107e4c0"
+OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Strided = "5e0ebb24-38b0-5f93-81fe-25c709ecae67"
 

docs/make.jl

@@ -1,10 +1,10 @@
 
 using Documenter
-using TensorCast
+using TensorCast, OffsetArrays
 
 makedocs(
     sitename = "TensorCast",
-    modules = [TensorCast],
+    modules = [TensorCast, OffsetArrays],
     pages = [
         "Home" => "index.md",
         "Basics" => "basics.md",

docs/src/basics.md

@@ -83,7 +83,7 @@ Combining this with slicing from `M[:,j]` is a convenient way to perform `mapsli
 
 ```jldoctest mylabel
 julia> @cast C[i,j] := cumsum(M[:,j])[i]
-3×4 stack(::Vector{Vector{Int64}}) with eltype Int64:
+3×4 lazystack(::Vector{Vector{Int64}}) with eltype Int64:
  1   4   7  10
  3   9  15  21
  6  15  24  33
@@ -205,28 +205,23 @@ as `size(A)` is known:
 julia> @cast A[i,j] = 10 * collect(1:12)[i⊗j];
 ```
 
+## Repeating
+
 If the right hand side is independent of an index, then the same result is repeated. 
 The range of the index must still be known:
 
 ```jldoctest mylabel
 julia> @cast R[r,(n,c)] := M[r,c]^2  (n in 1:3)
-ERROR: LoadError: index n appears only on the left
-    @cast R[r, (n, c)] := M[r, c] ^ 2  n in 1:3
-    @ Main none:1
-Stacktrace:
- [1] checkallseen(canon::Vector{Any}, store::NamedTuple{(:dict, :assert, :mustassert, :seen, :need, :top, :main), Tuple{Dict{Any, Any}, Vector{Any}, Vector{Any}, Vector{Any}, Vector{Any}, Vector{Any}, Vector{Any}}}, call::TensorCast.CallInfo)
-   @ TensorCast ~/.julia/dev/TensorCast/src/macro.jl:1466
- [2] _macro(exone::Expr, extwo::Expr, exthree::Nothing; call::TensorCast.CallInfo, dict::Dict{Any, Any})
-   @ TensorCast ~/.julia/dev/TensorCast/src/macro.jl:199
- [3] var"@cast"(__source__::LineNumberNode, __module__::Module, exs::Vararg{Any})
-   @ TensorCast ~/.julia/dev/TensorCast/src/macro.jl:74
-in expression starting at none:1
+3×12 Matrix{Int64}:
+ 1  1  1  16  16  16  49  49  49  100  100  100
+ 4  4  4  25  25  25  64  64  64  121  121  121
+ 9  9  9  36  36  36  81  81  81  144  144  144
 
 julia> R == repeat(M .^ 2, inner=(1,3))
 true
 
-julia> @cast R[r,(c,n)] = M[r,c]  # repeat(M, outer=(1,3)), uses size(R)
-3×12 Array{Int64,2}:
+julia> @cast similar(R)[r,(c,n)] = M[r,c]  # repeat(M, outer=(1,3)), uses size(R)
+3×12 Matrix{Int64}:
  1  4  7  10  1  4  7  10  1  4  7  10
  2  5  8  11  2  5  8  11  2  5  8  11
  3  6  9  12  3  6  9  12  3  6  9  12

docs/src/index.md

@@ -23,6 +23,7 @@ New features in 0.4:
 - Indices can appear ouside of indexing: `@cast A[i,j] = i+j` translates to `A .= axes(A,1) .+ axes(A,2)'`
 - The ternary operator `? :` can appear on the right, and will be broadcast correctly.
 - All operations should now support [OffsetArrays.jl](https://github.com/JuliaArrays/OffsetArrays.jl).
+- You can `repeat` by broadcasting over indices not appearing on the right, such as `@cast r[i,(k,j)] = m[i,j]`
 
 ## Pages
 

docs/src/options.md

@@ -22,12 +22,12 @@ M = rand(1:99, 3,4)
 @cast S[k][i] := M[i,k] lazy=false  # the same
 ```
 
-The default way of un-slicing is `reduce(hcat, ...)`, which creates a new array. 
-But there are other options, controlled by keywords after the expression:
+The default way of un-slicing uses [LazyStack.jl](https://github.com/mcabbott/LazyStack.jl) to create a view.
+The keyword `lazy=false` after the expression will turn this off, to make a solid array using Base's code:
 
 ```julia
-@cast A[i,k] := S[k][i] lazy=false  # A = reduce(hcat, B)
 @cast A[i,k] := S[k][i]             # A = LazyStack.stack(B)
+@cast A[i,k] := S[k][i] lazy=false  # A = reduce(hcat, B)
 
 size(A) == (3, 4) # true
 ```

src/TensorCast.jl

@@ -32,7 +32,7 @@ module Fast # shield non-macro code from @optlevel 1
     export sliceview, slicecopy, copy_glue, glue!, iscodesorted, countcolons
 
     include("view.jl")
-    export diagview, mul!, rview, star
+    export diagview, mul!, rview, star, onlyfirst
 
     include("static.jl")
     export static_slice, static_glue

src/macro.jl

@@ -187,6 +187,7 @@ function _macro(exone, extwo=nothing, exthree=nothing; call::CallInfo=CallInfo()
 
     # Third pass to standardise & then glue, postwalk sees A[i] before A[i][j]
     right3 = MacroTools.postwalk(x -> standardglue(x, canon, store, call), right2)
+    right3 = checkallseen(right3, canon, store, call)
 
     if !(:matmul in call.flags)
         # Then finally broadcasting if necc (or just permutedims etc. if not):
@@ -196,8 +197,6 @@ function _macro(exone, extwo=nothing, exthree=nothing; call::CallInfo=CallInfo()
         right4 = matmultarget(right3, canon, parsed, store, call)
     end
 
-    checkallseen(canon, store, call) # this must be run before inplaceoutput()
-
     # Return to LHS, build up what it requested:
     if :inplace in call.flags
         rightlist = inplaceoutput(right4, canon, parsed, store, call)
@@ -1344,6 +1343,25 @@ end
 #     return sortperm(xi) == sortperm(yi)
 # end
 
+"""
+    checkallseen(rhs, ...)
+
+Now not just a check, but also inserts trivial broadcasting, if needed, 
+over indices omitted from the rhs.
+"""
+function checkallseen(ex, canon, store, call)
+    right = setdiff(store.seen, canon)
+    length(right) > 0 && throw(MacroError("index $(right[1]) appears only on the right", call))
+    left = setdiff(canon, unique!(store.seen))
+    # length(left) > 0 && throw(MacroError("index $(left[1]) appears only on the left", call))
+    if isempty(left)
+        ex
+    else
+        fake = map(i -> recursemacro(i, canon, store, call), left)
+        :( TensorCast.onlyfirst($ex, $(fake...)) )
+    end
+end
+
 """
     needview!([:, 3, A])   # true, need view(A, :,3,:)
     needview!([:, :_, :])  # false, can use rview(A, :,1,:)
@@ -1467,14 +1485,6 @@ function checknorepeats(flat, call=nothing, msg=nothing)
     end
 end
 
-function checkallseen(canon, store, call)
-    left = setdiff(canon, unique!(store.seen))
-    length(left) > 0 && throw(MacroError("index $(left[1]) appears only on the left", call))
-    right = setdiff(store.seen, canon)
-    length(right) > 0 && throw(MacroError("index $(right[1]) appears only on the right", call))
-end
-
-# this may never be necessary with checkallseen?
 function findcheck(i::Symbol, flat::Vector, call=nothing, msg=nothing)
     msg == nothing && (msg = " in " * string(:( [$(flat...)] )))
     res = findfirst(isequal(i), flat)
@@ -1590,11 +1600,8 @@ function inplaceoutput(ex, canon, parsed, store::NamedTuple, call::CallInfo)
         newleft = standardise(parsed.left, store, call)
         @capture(newleft, zed_[ijk__]) || throw(MacroError("failed to parse LHS correctly, $(parsed.left) -> $newleft"))
 
-        if !(zed isa Symbol) # then standardise did something!
+        if newleft != parsed.left  # then standardise did something!
             push!(call.flags, :showfinal)
-            Zsym = gensym(:reverse)
-            push!(out, :( local $Zsym = $zed ) )
-            zed = Zsym
         end
     end
 
@@ -1619,7 +1626,12 @@ function inplaceoutput(ex, canon, parsed, store::NamedTuple, call::CallInfo)
     end
 
     if :showfinal in call.flags
-        push!(out, parsed.name)
+        A = parsed.name
+        if A isa Symbol || @capture(A, AA_.ff_)
+        else
+            A = Symbol(A,"_val") # exact same symbol is used by standardise()
+        end
+        push!(out, A)
     end
 
     return out

src/view.jl

@@ -28,7 +28,15 @@ mul!(Z::AbstractArray{T,0}, A,B) where {T} = copyto!(Z, A * B)
 """
     star(x,y,...)
 
-Used for multiplying axes now, not sizes.
+Used for multiplying axes now, always producing a `OneTo` whose length
+is the product of the given ranges.
 """
 star(x, y) = Base.OneTo(length(x) * length(y))
 star(x,y,zs...) = star(star(x,y), zs...)
+
+"""
+    onlyfirst(x, ys...) = x
+
+Used with arguments which are there just to set the shape of broadcasting.
+"""
+onlyfirst(x, ys...) = x

test/four.jl

@@ -81,6 +81,26 @@ end
     @test E ≈ sum((1:2) ./ (1:4)')
 end
 
+@testset "repeats" begin
+    @test [1,1,1] == @cast _[i] := 1  (i in 1:3)
+    @test [1,1,1] == @cast rand(3)[i] = 1
+
+    M = reshape(1:12, 3,4)
+    repeat(M, inner=(1,3)) == @cast _[i, (r,j)] := M[i,j]  r in 1:3
+    repeat(M, outer=(1,3)) == @cast _[i, (j,r)] := M[i,j]  r in 1:3
+
+    @cast T[i,r,j] := M[i,j]  r in 1:3
+    @test T[:,3,:] == M
+
+    @test 100 .+ sum(M, dims=2) == @cast _[i,k] := sum(M[i,:])+100  k in 1:1
+    @test 4 .+ sum(M, dims=2) == @reduce _[i,k] := sum(j) M[i,j]+1  k in 1:1
+    @test repeat(sum(M, dims=2), inner=(1,2)) == @reduce _[i,k] := sum(j) M[i,j]  k in 1:2
+    @test repeat(M, inner=(1,2)) == @cast rand(3,8)[i,(k,j)] = M[i,j]
+
+    @test_throws Exception @macroexpand @cast _[i] := A[i,j]  # j doesn't appear on left
+    # @test_throws Exception @macroexpand @cast _[i] := i+j  (i in 1:2, j in 1:3)
+end
+
 @testset "offset handling" begin
     using OffsetArrays
     @cast A[i,j] := i+10j  (i in 0:1, j in 7:15)
@@ -91,6 +111,13 @@ end
     @cast B[(i,k),j] := A[i,(j,k)]  k in 1:3
     @test axes(B) === (Base.OneTo(6), Base.OneTo(3))
 
+    @cast E[(j,i),k] := A[i,j]^2  (k in 10:11)  # RHS indep k
+    @test E[:,11] == vec(A') .^ 2
+
+    # dropdims
+    @cast F[i,k,j] := A[i,j]  k in 5:5  # a trivial dimension not indexed from 1
+    @test A == @cast _[i,j] := F[i,_,j]
+
     # reduction
     @reduce C[_,j] := sum(i) A[i,j]
     @test axes(C) == (1:1, 7:15)
@@ -156,8 +183,9 @@ end
     @test y == @cast _[i,j,k] := tuple(y[i,:,k]...)[j]
 end
 
+using TensorCast: MacroError, _macro, CallInfo
 @testset "bugs" begin
     # scatter not handled, previously ignored -- https://github.com/mcabbott/TensorCast.jl/issues/49
-    @test_throws TensorCast.MacroError @pretty @reduce v[ixs[j]] := mean(i) vs[i][j]
-    @test_throws TensorCast.MacroError @pretty @cast v[i][ixs[j]] := vs[j][i]
+    @test_throws MacroError _macro(:(  v[ixs[j]] := mean(i) ),:(  vs[i][j] ), call=CallInfo(:reduce))
+    @test_throws MacroError _macro(:(  v[i][ixs[j]] := vs[j][i] ), call=CallInfo())
 end