How to loop over an array sorted by one dimension in Chapel?

Question

Suppose our darling A looks like

60.0    3.0
675.0   3.0
1050.0  4.0
0.0     0.0

I want to loop through the rows in decreasing order of the first column. Something like...

for r in sorted(A, by=A[1,], reverse=True) {
  writeln(r);
}

I'd like

1050.0  4.0
675.0   3.0
60.0    3.0
0.0     0.0

The Sort documentation does not address 2D arrays.

Answer 1

As you have observed, the Sort module currently only supports 1D arrays.

As a work-around to this, you could write a sort() wrapper that restructures your data into an array of arrays and sort on that, if you're willing to accept a penalty in performance, due to creating a temp-array and iterating over it twice.

An example of this:

use Sort;

var A: [1..4, 1..2] real = ((60.0, 3.0), (675.0, 3.0), (1050.0, 4.0), (0.0, 0.0));

writeln('A:');
writeln(A);

sort2D(A, axis=2, reversed=true);

writeln('A sorted:');
writeln(A);

/* 2D sort wrapper using temp array */
proc sort2D(A: [?D] ?t, axis=1, reversed=false) {
  const (rows, cols) = A.shape;

  // array of arrays
  var tmp: [D.dim(1)] [D.dim(2)] t;
  for row in D.dim(1) {
    tmp[row] = A[row, ..];
  }

  var cmp = new Comparator2D(axis, reversed);
  sort(tmp, comparator=cmp);

  for row in D.dim(1) {
    A[row, ..] = tmp[row];
  }
}

/* Comparator for arrays of arrays */
record Comparator2D {
  const axis = 1,
        reversed = false;

  proc key(a) {
    if reversed then
      return -a[axis];
    else
      return a[axis];
  }
}