list manipulation - How can I make threading more flexible?

Threading automatically with Listable functions requires the argument expressions to have the same length (or for one of them to be atomic). For nested lists the threading will continue down through the levels, provided the lengths are the same at each level. So, for example, these all work because the Dimensions of the two lists are the same at the first $n$ levels:

Array[#&, {10, 5, 3}] + Array[#&, {10}];
Array[#&, {10, 5, 3}] + Array[#&, {10, 5}];
Array[#&, {10, 5, 3}] + Array[#&, {10, 5, 3}];

whereas this doesn't work because the outer Dimensions don't match ($10\neq 5$):

Array[#&, {10, 5, 3}] + Array[#&, {5, 3}];

(*  Thread::tdlen: Objects of unequal length ... cannot be combined. *)

But there is an obvious interpretation of the above code, which is to map the addition over the outer level of the first argument, i.e. to add the second 5x3 array to each of the ten 5x3 arrays in the first argument.

A more easily visualised example is adding an offset to a list of coordinates:

coords = {{1, 2}, {3, 4}, {5, 6}, {7, 8}};
offset = {0, 10};

One way is to explicity Map the addition over the coordinate list:

result = # + offset & /@ coords
(* {{1, 12}, {3, 14}, {5, 16}, {7, 18}} *)

If the coordinate list was very long, a more efficient approach using Transpose might be preferred:

result = Transpose[Transpose[coords] + offset]
(* {{1, 12}, {3, 14}, {5, 16}, {7, 18}} *)

Neither of these is particularly readable though. It would be nice to have a "smart" threading function that would identify that the second dimension of coords (length 2) matches the first dimensions of offset (also length 2), allowing the code to be written very readably:

result = smartThread[ coords + offset ]
(* {{1, 12}, {3, 14}, {5, 16}, {7, 18}} *)

How can I write such a smartThread function, which will take an expression of the form func_[a_, b_] and match up the various dimensions in a and b to do this kind of flexible threading?

Answer

My own solution looks like this:

SetAttributes[smartThread, HoldAll];

smartThread[f_[a_?ArrayQ, b_?ArrayQ], dir : -1 | 1 : -1] := Module[{da, db, or, o, g, t},
  or = Function[{x, y, d}, Select[Permutations @ Range @ Length[x], 
   x[[#[[;; Length[y]]]]] == y &][[d]] ~Check~ Return[$Failed, Module]];
  t = #1 ~Transpose~ Ordering[#2] &;
  g = If[MemberQ[Attributes[f], Listable], f, Function[Null, f[##], Listable]];
  {da, db} = Dimensions /@ {a, b};

  If[Length[da] >= Length[db],
    t[a, o = or[da, db, dir]] ~g~ b, 
    a ~g~ t[b, o = or[db, da, dir]]]
   ~Transpose~ o]

The function works by examining the dimensions of both input lists to determine which dimensions are "shared" (ie. have the same length) by both inputs. One of the input lists (the one with the greater ArrayDepth) is transposed such that the shared dimensions are outermost. This allows the function to thread over those dimensions, after which the Transpose is reversed.

Straightforward examples

By "straightforward", I mean cases where there is no ambiguity about how to match up the dimensions of the two input lists. For Listable functions like Plus and Times the smartThread function works as you would expect:

smartThread[{{1, 2}, {3, 4}, {5, 6}} + {10, 0}]
(* {{11, 2}, {13, 4}, {15, 6}} *)


smartThread[{{1, 2}, {3, 4}, {5, 6}} * {10, 0}]
(* {{10, 0}, {30, 0}, {50, 0}} *)

Functions which are not normally Listable are replaced internally by versions which are, so you get Thread-like behaviour:

smartThread[{{1, 2}, {3, 4}, {5, 6}} ~ Max ~ {10, 0}]
(* {{10, 2}, {10, 4}, {10, 6}} *)

smartThread[{{1, 2}, {3, 4}, {5, 6}} ~ f ~ {10, 0}]
(* {{f[1, 10], f[2, 0]}, {f[3, 10], f[4, 0]}, {f[5, 10], f[6, 0]}} *)

It should work with nested lists of any depth, provided the dimensions can be matched up. Note that the output has the same shape as whichever input has the greater ArrayDepth:

smartThread[Array[#&, {5, 7, 2, 10, 3, 4}] + Array[#&, {10, 7, 3}]] // Dimensions
(* {5, 7, 2, 10, 3, 4} *)

smartThread[Array[#&, {10, 7, 3}] + Array[#&, {5, 7, 2, 10, 3, 4}]] // Dimensions
(* {5, 7, 2, 10, 3, 4} *)

Equal ArrayDepth case

If both input lists have the same ArrayDepth, the output will be the same shape as the first one:

smartThread[Array[# &, {10, 2}] + Array[# &, {2, 10}]] // Dimensions
(* {10, 2} *)

smartThread[Array[# &, {2, 10}] + Array[# &, {10, 2}]] // Dimensions
(* {2, 10} *)

Dimension matching ambiguites

There is not always a single unique way to match up the dimensions of the input lists. Consider for example smartThread[a + b] where Dimensions[a] = {2, 3, 2} and Dimensions[b] = {3, 2}. The default behaviour of smartThread is to match the dimensions innermost first, so the result would be equivalent to {a[[1]] + b, a[[2]] + b}:

smartThread[Array[0 &, {2, 3, 2}] + {{1, 2}, {3, 4}, {5, 6}}]
(* {{{1, 2}, {3, 4}, {5, 6}}, {{1, 2}, {3, 4}, {5, 6}}} *)

If it is required to match the dimensions outermost first, you can supply +1 as a second argument to smartThread:

smartThread[Array[0 &, {2, 3, 2}] + {{1, 2}, {3, 4}, {5, 6}}, 1]
(* {{{1, 1}, {3, 3}, {5, 5}}, {{2, 2}, {4, 4}, {6, 6}}} *)

More generally, setting the second argument to n causes the code to select the n'th valid permutation of the dimensions.

Note: I would personally recommend thinking twice before using smartThread in cases where there is ambiguity over how to match dimensions in the input lists. The motivation for writing it was to allow simple constructs like smartThread[coordinateList + offset], increasing readability for code where the intention is intuitively obvious. In situations where that isn't the case, it potentially makes the code less clear than using something like Map or Transpose explicitly.