13
$\begingroup$

I have a tensor of dimensions $(2, 100, 100, 2, 100, 100)$ and I want to reshape it to a form of $(2*100*100,2*100*100)$, e.g. Flatten[A,{{1,5,6},{4,2,3}}]. If I use any standard method like Transpose plus ArrayReshape or Flatten, I end up being $~6-10$ times slower than Matlab. I would appreciate if there are some ideas for efficient transpose.

$\endgroup$
  • 4
    $\begingroup$ Since Flatten with the second argument is tailor-made for this, I doubt that you can do much better, short of writing custom C code via LibraryLink. $\endgroup$ – Leonid Shifrin Oct 13 '15 at 22:41
  • 2
    $\begingroup$ Only alternative I can think of is Join@@@Join@@@Join@@@(Join@@Join@@Join@@list). I am guessing that is not faster. $\endgroup$ – Verbeia Oct 13 '15 at 22:43
  • $\begingroup$ @LeonidShifrin If you have some C functions for this or know where to find some plz feel free to share :). $\endgroup$ – Alexis Michailidis Oct 13 '15 at 22:46
  • $\begingroup$ @AlexisMichailidis Alas, I don't. But that shouldn't be particularly hard to write, as long as you have fixed number of dimensions. $\endgroup$ – Leonid Shifrin Oct 13 '15 at 22:49
  • 9
    $\begingroup$ Make sure that your tensor is a packed array. $\endgroup$ – user21 Oct 14 '15 at 9:18
5
$\begingroup$

Mathematica changes over time, and perhaps it's more efficient than four years ago. I doubt it, though, in this area. I don't know how to do in MATLAB what the OP claims is done much faster in MATLAB, and there was no code provided to test. But we can compare the computation Flatten[A,{{1,5,6},{4,2,3}}] with the speed with which the array can be copied, which will give us some sort of gauge of how efficient Mathematica is.

Let's make an array. It's a big one, 3GB, and I had to make sure I had enough RAM available. (I have 16GB, but I also have a lot of things running. This example will take about 9GB, plus whatever overhead we incur.)

Clear[array, flatA, transA]

array = ConstantArray[1., {2, 100, 100, 2, 100, 100}]; // AbsoluteTiming
ByteCount@array
(*  {1.46419, Null}  <- time to write a 3GB array *)
(*  3200000184  *)

To copy an array, you have to store in another variable and then change one of the arrays. The first assignment just copies a pointer, which is very fast. When you change one of the elements of one of the arrays, Mathematica then copies the array and makes the change.

brray = array; // AbsoluteTiming
(*  {2.*10^-6, Null}  <- time to copy a pointer *)

brray[[1, 1, 1, 1, 1, 1]] = 2.; // AbsoluteTiming
(*  {1.88523, Null}   <- time to copy (read & write) 3GB array *)

brray[[1, 1, 1, 1, 1, 2]] = 2.; // AbsoluteTiming
(*  {7.*10^-6, Null}  <- changing a second element is fast *)

brray =.  (* clear B to free up memory *)

Comparing Flatten with ArrayReshape + Transpose, we see that Flatten is rather slow, while the other method takes the same time as copying the array twice. Copying it twice should be expected (because of the two function calls), so the second method is as efficient as it could be. Theoretically, one expects that the computation could be done with one copy, but, alas, Flatten does not do that. (@Leonid Shifrin in a comment suggests it might be done in a C function via LibraryLink.)

flatA = Flatten[array, {{1, 5, 6}, {4, 2, 3}}]; // AbsoluteTiming
(*  {9.91492, Null}  *)

transA = ArrayReshape[
    Transpose[array, {1, 5, 6, 4, 2, 3}],
    {2*100*100, 2*100*100}]; // AbsoluteTiming
(*  {3.84417, Null}  *)

transA === flatA
(*  True  *)
$\endgroup$
  • 1
    $\begingroup$ Flatten is so slow here, even for a packed array, that I'd suggest filing this as "unexpected behavior" with Wolfram Support. $\endgroup$ – Roman Jul 13 at 7:36
  • $\begingroup$ It may be worthwhile to check if the neural network framework offers something that works for this problem, since neural networks are the closes thing WL has to a dedicated numerics framework. $\endgroup$ – Sjoerd Smit Jul 13 at 13:12
3
$\begingroup$

I've followed Leonid's suggestion and written a C function that does this flattening job. It is faster than Flatten, but slower than Michael's ArrayReshape@*Transpose solution, even though I've tried to optimize it as much as I could. WR's code is impressive!

Here's the C code:

code = "
  #include \"WolframLibrary.h\"

  DLLEXPORT mint WolframLibrary_getVersion( ) {
     return WolframLibraryVersion;
  }

  DLLEXPORT int WolframLibrary_initialize(WolframLibraryData libData) {
     return LIBRARY_NO_ERROR;
  }

  DLLEXPORT void WolframLibrary_uninitialize( ) {
     return;
  }

  DLLEXPORT int myflatten(WolframLibraryData libData, mint Argc,
                          MArgument *Args, MArgument Res) {
    // read the input tensor and check it
    MTensor T0 = MArgument_getMTensor(Args[0]);
    if (libData->MTensor_getType(T0) != MType_Real)
      return LIBRARY_TYPE_ERROR;
    if (libData->MTensor_getRank(T0) != 6)
      return LIBRARY_RANK_ERROR;
    mint const* dims;
    dims = libData->MTensor_getDimensions(T0);
    // allocate a new tensor
    int err = LIBRARY_NO_ERROR;
    MTensor Tnew;
    mint newdims[2];
    newdims[0] = dims[0]*dims[4]*dims[5];
    newdims[1] = dims[3]*dims[1]*dims[2];
    err = libData->MTensor_new(MType_Real, 2, newdims, &Tnew);
    if (err) return err;
    // fill new tensor
    mint pos[6];
    double* t0 = libData->MTensor_getRealData(T0);
    double* tnew = libData->MTensor_getRealData(Tnew);
    mint A = dims[3]*dims[4]*dims[5];
    mint B = (dims[1]*dims[2]*dims[3]-1)*dims[4]*dims[5];
    for (pos[0]=0; pos[0]<dims[0]; pos[0]++) {
      for (pos[4]=0; pos[4]<dims[4]; pos[4]++) {
        for (pos[5]=0; pos[5]<dims[5]; pos[5]++) {
          for (pos[3]=0; pos[3]<dims[3]; pos[3]++) {
            for (pos[1]=0; pos[1]<dims[1]; pos[1]++) {
              for (pos[2]=0; pos[2]<dims[2]; pos[2]++) {
                *(tnew++) = *t0;
                t0 += A;
              }
            }
            t0 -= B;
          }
          t0 -= A-1;
        }
      }
      t0 += B;
    }
    // return new tensor
    MArgument_setMTensor(Res, Tnew);
    return err;
  }
";

Compile it and make it into a Mathematica function with

Needs["CCompilerDriver`"]
lib = CreateLibrary[code, "testDLL", "ShellOutputFunction" -> Print, 
                                     "ShellCommandFunction" -> Print]
myflatten = LibraryFunctionLoad[lib, "myflatten",
                                {{Real, 6, "Constant"}}, {Real, 2}];

Check timing of various flattening strategies:

S = 50;
X = RandomReal[{0, 1}, {2, S, S, 2, S, S}];

(* original *)
Y1 = Flatten[X, {{1, 5, 6}, {4, 2, 3}}]; // RepeatedTiming // First
(*    0.569    *)

(* my C code *)
Y2 = myflatten[X]; // RepeatedTiming // First
(*    0.274    *)

(* Michael's ArrayReshape@*Transpose solution *)
Y3 = ArrayReshape[Transpose[X, {1, 5, 6, 4, 2, 3}], {2 S^2, 2 S^2}]; // RepeatedTiming // First
(*    0.16    *)

Make sure all procedures give the same result:

Y1 == Y2 == Y3
(*    True    *)
$\endgroup$
  • $\begingroup$ Thanks for trying this. I was thinking about doing it later, but now I don't need to. :) +1 $\endgroup$ – Michael E2 Jul 13 at 22:22
  • $\begingroup$ @MichaelE2 Any ideas on how to improve the C code? It's strangely slow. The entire time is spent in the six-fold for loop, and the MTensor allocation takes no time at all. $\endgroup$ – Roman Jul 13 at 22:24
  • $\begingroup$ I've done no active C development in decades, so take with a grain of salt: AFAICT, The new array is written in memory order and the old array is read out of order: reading is faster than writing, so reading a block multiple times is better than writing. That seems good. -- Are all the pos[k]++ being done in registers or some other less efficient memory device? I don't know how C compilers handle that nowadays, but in old ones, indexed arrays would not be in registers. Maybe the same consideration for dim[k], but they are only read. -- Can memcpy() be leveraged? (I don't see how.) $\endgroup$ – Michael E2 Jul 13 at 22:40
  • 1
    $\begingroup$ Well, I tried changing pos[6] to register int p0,..p6 but no difference in speed. Apparently I haven't got a clue... $\endgroup$ – Michael E2 Jul 13 at 22:44
  • $\begingroup$ @MichaelE2 I couldn't figure out how to use memcpy either. As for using indexed arrays vs. symbols, I leave this to the compiler: at -O2 this kind of nitty-gritty should be taken care of automagically these days (you can check with the Compiler Explorer if unsure). $\endgroup$ – Roman Jul 14 at 0:31

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.