Problem with Compile and Flatten

Question

I want to get two lists within a compiled function, where r is an integer:

Insert[
  DeleteCases[
   Flatten[
    Table[{i, j}, {i, 0, r - 1}, {j, i, r - 1}], 1],{0, 0}], {0, r}, r];

Flatten[Table[{i, j}, {i, r - 1}, {j, 0, i - 1}], 1];

But Flatten and Compile don't seem to work very well together:

Compile[{{r, _Integer}},
    Module[{list, i = 0, j},
        list = Table[{i, j}, {i, 0, r - 1}, {j, i, r - 1}];
        Flatten[list, 1];
    ]
]

This yields Compile::cplist: list should be a tensor of type Integer, Real, or Complex; evaluation will use the uncompiled function.. It works fine without the Flatten though.

Answer 1

It looks like the problem you are running in to is because list is a ragged array, which is not allowed. Replace {j,i,r-1} with {j,r-1} and you get a function that compiles and runs (despite the fact that it doesn't give you the desired output).

One way to get your desired output is to create a regular array with some criterion that can be selected for removal afterwards. Note I specifically avoid the use of pattern since pattern matching is not allowed in Compiled functions.

This seems to work:

f = Compile[{{r, _Integer}}, 
  Module[{list, i = 0, j}, 
   list = Table[{i, If[j < i, -1, j]}, {i, 0, r - 1}, {j, 0, r - 1}];
   Select[Flatten[list, 1], #[[2]] =!= -1 &]]]

Answer 2

One obvious solution that comes to my mind, especially when I see your DeleteCases- and Insert-calls, is to use Internal`Bag. This makes is somewhat easy to collect all elements of your result. Let's start with the compiled equivalent of

Flatten[Table[{i, j}, {i, r - 1}, {j, 0, i - 1}], 1];

Inside a compiled function, you would start by creating a new bag for your result. Important is that I'm storing your two-dimensional tensor into a one-dimensional list. This is not a problem, because we always add a pair e.g. {0,2} so we know that your final result can easily be rebuilt by using Partition[..,2] at the end:

f1 = Compile[{{r, _Integer}},
  Module[{res = Internal`Bag[Most[{0}]]},
   Do[Internal`StuffBag[res, {i, j}, 1], {i, r - 1}, {j, 0, i - 1}];
   Partition[Internal`BagPart[res, All], 2]
  ]
]

Knowing this, makes the creation of your other list easy. Just replace the DeleteCases and Insert calls by appropriate conditional events:

f2 = Compile[{{r, _Integer}},
  Module[{res = Internal`Bag[Most[{0}]], c = 0},
   Do[
    If[c++ == r, Internal`StuffBag[res, {0, r}, 1]];
    If[i == 0 && j == 0, Continue[]];
    Internal`StuffBag[res, {i, j}, 1], {i, 0, r - 1}, {j, i, r - 1}
    ];
   Partition[Internal`BagPart[res, All], 2]
  ]
]

You can use CompilePrint from the <<CompiledFunctionTools` package to verify that everything is properly compiled down.

Answer 3

All what you have to do is to tell Compile that list and result are in fact tensors. You can do that by adding the optional arguments at the end of Compile: {list,_Integer,3} meaning that list is an integer tensor of rank 3, and {result, _Integer,2} meaning that result is a tensor of rank 2. list should go out of the module in this case.

Compile should look like:

f1 = Compile[{{r, _Integer}}, 
   result = Module[{i = 0, j}, 
      list = Table[{i, j}, {i, 0, r - 1}, {j, i, r - 1}];
      Flatten[list, 1]
      ]; (* End Module *)
   result (* Return result *)
   , {{list, _Integer, 3}, {result, _Integer, 2}}
]; (* End Compile *)

which yields:

{{0,0},{0,1},{0,2},{0,3},{0,4},{1,1},{1,2},{1,3},{1,4},{2,2},{2,3},{2,4},{3,3},{3,4},{4,4}}