How should I selectively sum multiple axes of an array?

Question

What is the preferred approach in J for selectively summing multiple axes of an array?

For instance, suppose that a is the following rank 3 array:

   ]a =: i. 2 3 4
 0  1  2  3
 4  5  6  7
 8  9 10 11

12 13 14 15
16 17 18 19
20 21 22 23

My goal is to define a dyad "sumAxes" to sum over multiple axes of my choosing:

   0 1 sumAxes a      NB. 0+4+8+12+16+20 ...
60 66 72 78

   0 2 sumAxes a      NB. 0+1+2+3+12+13+14+15 ...
60 92 124

   1 2 sumAxes a      NB. 0+1+2+3+4+5+6+7+8+9+10+11 ...
66 210

The way that I am currently trying to implement this verb is to use the dyad |: to first permute the axes of a, and then ravel the items of the necessary rank using ,"n (where n is the number axes I want to sum over) before summing the resulting items:

   sumAxes =: dyad : '(+/ @ ,"(#x)) x |: y'

This appears to work as I want, but as a beginner in J I am unsure if I am overlooking some aspect of rank or particular verbs that would enable a cleaner definition. More generally I wonder whether permuting axes, ravelling and summing is idiomatic or efficient in this language.

For context, most of my previous experience with array programming is with Python's NumPy library.

NumPy does not have J's concept of rank and instead expects the user to explicitly label the axes of an array to reduce over:

>>> import numpy
>>> a = numpy.arange(2*3*4).reshape(2, 3, 4) # a =: i. 2 3 4
>>> a.sum(axis=(0, 2))                       # sum over specified axes
array([ 60,  92, 124])

As a footnote, my current implementation of sumAxes has the disadvantage of working "incorrectly" compared to NumPy when just a single axis is specified (as rank is not interchangeable with "axis").

Answer 1

Motivation

J has incredible facilities for handling arbitrarily-ranked arrays. But there's one facet of the language which is simultaneously almost universally useful as well as justified, but also somewhat antithetical to this dimensionality-agnostic nature.

The major axis (in fact, leading axes in general) are implicitly privileged. This is the concept that underlies, e.g. # being the count of items (i.e. the dimension of the first axis), the understated elegance and generality of +/ without further modification, and a host of other beautiful parts of the language.

But it's also what accounts for the obstacles you're meeting in trying to solve this problem.

Standard approach

So the general approach to solving the problem is just as you have it: transpose or otherwise rearrange the data so the axes that interest you become leading axes. Your approach is classic and unimpeachable. You can use it in good conscience.

Alternative approaches

But, like you, it niggles me a bit that we are forced to jump through such hoops in similar circumstances. One clue that we're kind of working against the grain of the language is the dynamic argument to the conjunction "(#x); usually arguments to conjunctions are fixed, and calculating them at runtime often forces us to use either explicit code (as in your example) or dramatically more complicated code. When the language makes something hard to do, it's usually a sign you're cutting against the grain.

Another is that ravel (,). It's not just that we want to transpose some axes; it's that we want to focus on one specific axis, and then run all the elements trailing it into a flat vector. Though I actually think this reflects more a constraint imposed by how we're framing the problem, rather than one in the notation. More on in the final section of this post.

With that, we might feel justified in our desire to address a non-leading axis directly. And, here and there, J provides primitives that allow us to do exactly that, which might be a hint that the language's designers also felt the need to include certain exceptions to the primacy of leading axes.

Introductory examples

For example, dyadic |. (rotate) has ranks 1 _, i.e. it takes a vector on the left.

This is sometimes surprising to people who have been using it for years, never having passed more than a scalar on the left. That, along with the unbound right rank, is another subtle consequence of J's leading-axis bias: we think of the right argument as a vector of items, and the left argument as a simple, scalar rotation value of that vector.

Thus:

   3 |. 1 2 3 4 5 6
4 5 6 1 2 3

and

   1 |. 1 2 , 3 4 ,: 5 6
3 4
5 6
1 2

But in this latter case, what if we didn't want to treat the table as a vector of rows, but as a vector of columns?

Of course, the classic approach is to use rank, to explicitly denote the the axis we're interested in (because leaving it implicit always selects the leading axis):

   1 |."1 ] 1 2 , 3 4 ,: 5 6
2 1
4 3
6 5

Now, this is perfectly idiomatic, standard, and ubiquitous in J code: J encourages us to think in terms of rank. No one would blink an eye on reading this code.

But, as described at the outset, in another sense it can feel like a cop-out, or manual adjustment. Especially when we want to dynamically choose the rank at runtime. Notationally, we are now no longer addressing the array as a whole, but addressing each row.

And this is where the left rank of |. comes in: it's one of those few primitives which can address non-leading axes directly.

   0 1 |. 1 2 , 3 4 ,: 5 6
2 1
4 3
6 5

Look ma, no rank! Of course, we now have to specify a rotation value for each axis independently, but that's not only ok, it's useful, because now that left argument smells much more like something which can be calculated from the input, in true J spirit.

Summing non-leading axes directly

So, now that we know J lets us address non-leading axes in certain cases, we simply have to survey those cases and identify one which seems fit for our purpose here.

The primitive I've found most generally useful for non-leading-axis work is ;. with a boxed left-hand argument. So my instinct is to reach for that first.

Let's start with your examples, slightly modified to see what we're summing.

    ]a =: i. 2 3 4
    sumAxes =: dyad : '(< @ ,"(#x)) x |: y'

     0 1 sumAxes a
+--------------+--------------+---------------+---------------+
|0 4 8 12 16 20|1 5 9 13 17 21|2 6 10 14 18 22|3 7 11 15 19 23|
+--------------+--------------+---------------+---------------+ 
     0 2 sumAxes a
+-------------------+-------------------+---------------------+
|0 1 2 3 12 13 14 15|4 5 6 7 16 17 18 19|8 9 10 11 20 21 22 23|
+-------------------+-------------------+---------------------+
    1 2 sumAxes a
+-------------------------+-----------------------------------+
|0 1 2 3 4 5 6 7 8 9 10 11|12 13 14 15 16 17 18 19 20 21 22 23|
+-------------------------+-----------------------------------+

The relevant part of the definition of for dyads derived from ;.1 and friends is:

The frets in the dyadic cases 1, _1, 2 , and _2 are determined by the 1s in boolean vector x; an empty vector x and non-zero #y indicates the entire of y. If x is the atom 0 or 1 it is treated as (#y)#x. In general, boolean vector >j{x specifies how axis j is to be cut, with an atom treated as (j{$y)#>j{x.

What this means is: if we're just trying to slice an array along its dimensions with no internal segmentation, we can simply use dyad cut with a left argument consisting solely of 1s and a:s. The number of 1s in the vector (ie. the sum) determines the rank of the resulting array.

Thus, to reproduce the examples above:

     ('';'';1) <@:,;.1 a
+--------------+--------------+---------------+---------------+
|0 4 8 12 16 20|1 5 9 13 17 21|2 6 10 14 18 22|3 7 11 15 19 23|
+--------------+--------------+---------------+---------------+
     ('';1;'') <@:,;.1 a
+-------------------+-------------------+---------------------+
|0 1 2 3 12 13 14 15|4 5 6 7 16 17 18 19|8 9 10 11 20 21 22 23|
+-------------------+-------------------+---------------------+
     (1;'';'') <@:,;.1 a
+-------------------------+-----------------------------------+
|0 1 2 3 4 5 6 7 8 9 10 11|12 13 14 15 16 17 18 19 20 21 22 23|
+-------------------------+-----------------------------------+

Et voila. Also, notice the pattern in the left hand argument? The two aces are exactly at the indices of your original calls to sumAxe. See what I mean by the fact that providing a value for each dimension smelling like a good thing, in the J spirit?

So, to use this approach to provide an analog to sumAxe with the same interface:

   sax =: dyad : 'y +/@:,;.1~ (1;a:#~r-1) |.~ - {. x -.~ i. r=.#$y'     NB. Explicit
   sax =: ]  +/@:,;.1~  ( (] (-@{.@] |. 1 ; a: #~ <:@[) (-.~ i.) ) #@$) NB. Tacit

Results elided for brevity, but they're identical to your sumAxe.

Final considerations

There's one more thing I'd like to point out. The interface to your sumAxe call, calqued from Python, names the two axes you'd like "run together". That's definitely one way of looking at it.

Another way of looking at it, which draws upon the J philosophies I've touched on here, is to name the axis you want to sum along. The fact that this is our actual focus is confirmed by the fact that we ravel each "slice", because we do not care about its shape, only its values.

This change in perspective to talk about the thing you're interested in, has the advantage that it is always a single thing, and this singularity permits certain simplifications in our code (again, especially in J, where we usually talk about the [new, i.e. post-transpose] leading axis)¹.

Let's look again at our ones-and-aces vector arguments to ;., to illustrate what I mean:

     ('';'';1) <@:,;.1 a
     ('';1;'') <@:,;.1 a
     (1;'';'') <@:,;.1 a

Now consider the three parenthesized arguments as a single matrix of three rows. What stands out to you? To me, it's the ones along the anti-diagonal. They are less numerous, and have values; by contrast the aces form the "background" of the matrix (the zeros). The ones are the true content.

Which is in contrast to how our sumAxe interface stands now: it asks us to specify the aces (zeros). How about instead we specify the 1, i.e. the axis that actually interests us?

If we do that, we can rewrite our functions thus:

  xas  =: dyad : 'y +/@:,;.1~ (-x) |. 1 ; a: #~ _1 + #$y'  NB. Explicit
  xas  =: ]  +/@:,;.1~  -@[ |. 1 ; a: #~ <:@#@$@]          NB. Tacit

And instead of calling 0 1 sax a, you'd call 2 xas a, instead of 0 2 sax a, you'd call 1 xas a, etc.

The relative simplicity of these two verbs suggests J agrees with this inversion of focus.

---

¹ ^{In this code I'm assuming you always want to collapse all axes except 1. This assumption is encoded in the approach I use to generate the ones-and-aces vector, using |..}

^{However, your footnote sumAxes has the disadvantage of working "incorrectly" compared to NumPy when just a single axis is specified suggests sometimes you want to only collapse one axis.}

^{That's perfectly possible and the ;. approach can take arbitrary (orthotopic) slices; we'd only need to alter the method by which we instruct it (generate the 1s-and-aces vector). If you provide a couple examples of generalizations you'd like, I'll update the post here. Probably just a matter of using (<1) x} a: #~ #$y or ((1;'') {~ (e.~ i.@#@$)) instead of (-x) |. 1 ; a:#~<:#$y.}