From dfb887fd9325ec7e231706007b53960387d33309 Mon Sep 17 00:00:00 2001
From: "Viral B. Shah" <viral@mayin.org>
Date: Sat, 20 Aug 2011 19:29:08 +0530
Subject: [PATCH] Use map() in @vectorize operators to generalize to N-d. Close
 #172.

---
 j/abstractarray.j | 64 +++++++++++++++++++++++++++++++++++------------
 j/operators.j     | 35 +++++++-------------------
 2 files changed, 57 insertions(+), 42 deletions(-)

diff --git a/j/abstractarray.j b/j/abstractarray.j
index b6854cb0bbb88..28561d7918fdf 100644
--- a/j/abstractarray.j
+++ b/j/abstractarray.j
@@ -843,6 +843,8 @@ isempty(a::AbstractArray) = (numel(a) == 0)
 #map(f, M::AbstractMatrix) = [ f(M[i,j]) | i=1:size(M,1), j=1:size(M,2) ]
 
 ## TODO: How to generalize map to n input args?
+
+## 1-d
 function map_to(dest::AbstractArray, f, A::AbstractArray)
     for i=1:numel(A)
         dest[i] = f(A[i])
@@ -850,6 +852,14 @@ function map_to(dest::AbstractArray, f, A::AbstractArray)
     return dest
 end
 
+function map(f, A::AbstractArray)
+    if isempty(A); return A; end
+    first = f(A[1])
+    dest = similar(A, typeof(first))
+    return map_to(dest, f, A)
+end
+
+## 2-d
 function map_to(dest::AbstractArray, f, A::AbstractArray, B::AbstractArray)
     for i=1:numel(A)
         dest[i] = f(A[i], B[i])
@@ -857,32 +867,54 @@ function map_to(dest::AbstractArray, f, A::AbstractArray, B::AbstractArray)
     return dest
 end
 
-function map(f, A::AbstractArray)
+function map(f, A::AbstractArray, B::AbstractArray)
+    if size(A) != size(B); error("Input size and shape should be same"); end
     if isempty(A); return A; end
-    first = f(A[1])
+    first = f(A[1], B[1])
     dest = similar(A, typeof(first))
-    return map_to(dest, f, A)
+    return map_to(dest, f, A, B)
 end
 
-function map(f, A::AbstractArray, B::AbstractArray)
-    if size(A) != size(B); error("Input size and shape should be same"); end
+function map_to(dest::AbstractArray, f, A::AbstractArray, B::Number)
+    for i=1:numel(A)
+        dest[i] = f(A[i], B)
+    end
+    return dest
+end
+
+function map(f, A::AbstractArray, B::Number)
     if isempty(A); return A; end
-    first = f(A[1], B[1])
+    first = f(A[1], B)
     dest = similar(A, typeof(first))
     return map_to(dest, f, A, B)
 end
 
+function map_to(dest::AbstractArray, f, A::Number, B::AbstractArray)
+    for i=1:numel(B)
+        dest[i] = f(A, B[i])
+    end
+    return dest
+end
+
+function map(f, A::Number, B::AbstractArray)
+    if isempty(A); return A; end
+    first = f(A, B[1])
+    dest = similar(B, typeof(first))
+    return map_to(dest, f, A, B)
+end
+
 ## Obsolete - Mainly here for reference purposes, use gen_cartesian_map
-# function cartesian_map(body, t::Tuple, it...)
-#     idx = length(t)-length(it)
-#     if idx == 0
-#         body(it)
-#     else
-#         for i = t[idx]
-#             cartesian_map(body, t, i, it...)
-#         end
-#     end
-# end
+## Still used in show()
+function cartesian_map(body, t::Tuple, it...)
+    idx = length(t)-length(it)
+    if idx == 0
+        body(it)
+    else
+        for i = t[idx]
+            cartesian_map(body, t, i, it...)
+        end
+    end
+end
 
 ## Transpose, Permute ##
 
diff --git a/j/operators.j b/j/operators.j
index b7e24ce6b7714..a674cbbcc01bb 100644
--- a/j/operators.j
+++ b/j/operators.j
@@ -101,40 +101,23 @@ one(::Type{Function}) = identity
 
 macro vectorize_1arg(S,f)
     quote
-        function ($f){T<:$S}(x::AbstractArray{T,1})
-            [ ($f)(x[i]) | i=1:length(x) ]
-        end
-        function ($f){T<:$S}(x::AbstractArray{T,2})
-            [ ($f)(x[i,j]) | i=1:size(x,1), j=1:size(x,2) ]
+        function ($f){T<:$S}(x::AbstractArray{T})
+            return map(($f), x)
         end
     end
 end
 
 macro vectorize_2arg(S,f)
     quote
-        function ($f){T<:$S}(x::T, y::AbstractArray{T,1})
-            [ ($f)(x,y[i]) | i=1:length(y) ]
-        end
-        function ($f){T<:$S}(x::AbstractArray{T,1}, y::T)
-            [ ($f)(x[i],y) | i=1:length(x) ]
-        end
-        function ($f){T<:$S}(x::T, y::AbstractArray{T,2})
-            [ ($f)(x,y[i,j]) | i=1:size(y,1), j=1:size(y,2) ]
+        function ($f){T1<:$S, T2<:$S}(x::T1, y::AbstractArray{T2})
+            return map(($f), x, y)
         end
-        function ($f){T<:$S}(x::AbstractArray{T,2}, y::T)
-            [ ($f)(x[i,j],y) | i=1:size(x,1), j=1:size(x,2) ]
+        function ($f){T1<:$S, T2<:$S}(x::AbstractArray{T1}, y::T2)
+            return map(($f), x, y)
         end
-        function ($f){T<:$S}(x::AbstractArray{T,1}, y::AbstractArray{T,1})
-            if size(x) != size(y)
-                error("vector length mismatch")
-            end
-            [ ($f)(x[i],y[i]) | i=1:length(x) ]
-        end
-        function ($f){T<:$S}(x::AbstractArray{T,2}, y::AbstractArray{T,2})
-            if size(x) != size(y)
-                error("matrix dimension mismatch")
-            end
-            [ ($f)(x[i,j],y[i,j]) | i=1:size(x,1), j=1:size(x,2) ]
+        function ($f){T1<:$S, T2<:$S}(x::AbstractArray{T1}, y::AbstractArray{T2})
+            if size(x) != size(y); error("Input sizes and shapes should match"); end
+            return map(($f), x, y)
         end
     end
 end