Merge pull request #38 from mcabbott/index

Allow indices to appear on RHS

Author: mcabbott

README.md

@@ -6,8 +6,9 @@
 [![Build Status](https://github.com/mcabbott/TensorCast.jl/workflows/CI/badge.svg)](https://github.com/mcabbott/TensorCast.jl/actions?query=workflow%3ACI)
 
 This package lets you work with multi-dimensional arrays in index notation, 
-by defining a few macros. The first is `@cast`, which deals both with "casting" into 
-new shapes (including going to and from an array-of-arrays) and with broadcasting:
+by defining a few macros. 
+
+The first is `@cast`, which deals both with "casting" into new shapes (including going to and from an array-of-arrays) and with broadcasting:
 
 ```julia
 @cast A[row][col] := B[row, col]       # slice a matrix B into rows, also @cast A[r] := B[r,:]
@@ -16,6 +17,8 @@ new shapes (including going to and from an array-of-arrays) and with broadcastin
 
 @cast E[φ,γ] = F[φ]^2 * exp(G[γ])      # broadcast E .= F.^2 .* exp.(G') into existing E
 
+@cast _[i] := isodd(i) ? log(i) : V[i] # broadcast a function of the index values
+
 @cast T[x,y,n] := outer(M[:,n])[x,y]   # generalised mapslices, vector -> matrix function
 ```
 

docs/src/basics.md

@@ -236,6 +236,42 @@ julia> colwise == Int.(reduce(hcat, vec.(list)))
 true
 ```
 
+## Index values
+
+Mostly the indices appearing in `@cast` expressions are just notation, to indicate what permutation / reshape is required. 
+But if an index appears outside of square brackets, this is understood as a value, implemented by broadcasting over a range (appropriately permuted):
+
+```jldoctest
+julia> @cast rat[i,j] := 0 * M[i,j] + i // j
+3×4 Array{Rational{Int64},2}:
+ 1//1  1//2  1//3  1//4
+ 2//1  1//1  2//3  1//2
+ 3//1  3//2  1//1  3//4
+
+julia> rat == @cast _[i,j] := axes(M,1)[i] // axes(M,2)[j]  # more verbose @cast
+true
+
+julia> rat == axes(M,1) .// transpose(axes(M,2))  # what it aims to generate
+true
+
+julia> @cast _[r,c] := r + 10c  (r in 1:2, c in 1:7)  # no array for range inference
+2×7 Array{Int64,2}:
+ 11  21  31  41  51  61  71
+ 12  22  32  42  52  62  72
+```
+
+Writing `$i` will interpolate the variable `i`, distinct from the index `i`:
+
+```jldoctest
+julia> i, k = 10, 100;
+
+julia> @cast ones(3)[i] = i + $i + k
+3-element Vector{Float64}:
+ 111.0
+ 112.0
+ 113.0
+```
+
 ## Reverse & shuffle
 
 A minus in front of an index will reverse that direction, and a tilde will shuffle it. 
@@ -248,6 +284,9 @@ julia> @cast M2[i,j] := M[i,-j]
  11  8  5  2
  12  9  6  3
 
+julia> all(M2[i,j] == M[i, end+begin-j] for i in 1:3, j in 1:4)
+true
+
 julia> using Random; Random.seed!(42); 
 
 julia> @cast M3[i,j] |= M[i,~j]
@@ -257,6 +296,9 @@ julia> @cast M3[i,j] |= M[i,~j]
  9  6  3  12
 ```
 
+Note that the minus is a slight deviation from the rule that left equals right for all indices,
+it should really be `M[i, end+1-j]`.
+
 ## Primes `'`
 
 Acting on indices, `A[i']` is normalised to `A[i′]` unicode \prime (which looks identical in some fonts).

docs/src/index.md

@@ -19,6 +19,10 @@ Version 0.4 has significant changes:
 - It uses [LazyStack.jl](https://github.com/mcabbott/LazyStack.jl) to combine handles slices, simplifying earlier code. This is lazier by default, write `@cast A[i,k] := log(B[k][i]) lazy=false` (with a new keyword option) to glue into an `Array` before broadcasting.
 - It uses [TransmuteDims.jl](https://github.com/mcabbott/TransmuteDims.jl) to handle all permutations & many reshapes. This is lazier by default -- the earlier code sometimes copied to avoid reshaping a `PermutedDimsArray`. This isn't always faster, though, and can be disabled by `lazy=false`.
 
+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. 
+
 ## Pages
 
 1. Use of `@cast` for broadcasting, dealing with arrays of arrays, and generalising `mapslices`

docs/src/reduce.md

@@ -17,7 +17,8 @@ julia> @reduce S[i] := sum(j) M[i,j] + 1000
 
 julia> @pretty @reduce S[i] := sum(j) M[i,j] + 1000
 begin
-    S = transmute(sum(@__dot__(M + 1000), dims = 2), (1,))
+    ndims(M) == 2 || throw(ArgumentError("expected a 2-tensor M[i, j]"))
+    S = dropdims(sum(@__dot__(M + 1000), dims = 2), dims=2)
 end
 ```
 
@@ -26,9 +27,6 @@ Note that:
 * The sum applies to the whole right side (including the 1000 here). 
 * And the summed dimensions are always dropped (unless you explicitly say `S[i,_] := ...`).
 
-Here `transmute(..., (1,))` is equivalent to `dropdims(..., dims=2)`, 
-keeping just the first dimension by reshaping.
-
 ## Not just `sum`
 
 You may use any reduction funciton which understands keyword `dims=...`, like `sum` does. 
@@ -112,10 +110,10 @@ There is no need to name the intermediate array, here `termite[x]`, but you must
 ```julia-repl
 julia> @pretty @reduce sum(x,θ) L[x,θ] * p[θ] * log(L[x,θ] / @reduce _[x] := sum(θ′) L[x,θ′] * p[θ′])
 begin
-    local fish = transmute(p, (nothing, 1))
-    termite = transmute(sum(@__dot__(L * fish), dims = 2), (1,))
-    local wallaby = transmute(p, (nothing, 1))
-    rat = sum(@__dot__(L * wallaby * log(L / termite)))
+    ndims(L) == 2 || error()  # etc, some checks
+    local goshawk = transmute(p, (nothing, 1))
+    sandpiper = dropdims(sum(@__dot__(L * goshawk), dims = 2), dims = 2)  # inner sum
+    bison = sum(@__dot__(L * goshawk * log(L / sandpiper)))
 end
 ```
 
@@ -131,8 +129,9 @@ before summing over one index:
 ```julia-repl
 julia> @pretty @reduce R[i,k] := sum(j) M[i,j] * N[j,k]
 begin
-    local fish = transmute(N, (nothing, 1, 2))                # fish = reshape(N, 1, size(N)...)
-    R = transmute(sum(@__dot__(M * fish), dims = 2), (1, 3))  # R = dropdims(sum(...), dims=2)
+    size(M, 2) == size(N, 1) || error()  # etc, some checks
+    local fish = transmute(N, (nothing, 1, 2))   # fish = reshape(N, 1, size(N)...)
+    R = dropdims(sum(@__dot__(M * fish), dims = 2), dims = 2)
 end
 ```