Эффективная обработка списка малых векторов с помощью Compile

Часто нам нужно обработать данные, состоящие из списка координат: data = {{x1, y1}, { x2, y2}, ..., {xn, yn}} . Это могут быть двухмерные или трехмерные координаты или любой другой список небольших векторов фиксированной длины произвольной длины.

Позвольте мне проиллюстрировать, как использовать Compile для таких задач, используя простой пример суммирования списка 2D-векторов:

data = RandomReal[1, {1000000, 2}];

Во-первых, очевидная версия:

fun1 = Compile[{{vec, _Real, 2}},
  Module[{sum = vec[[1]]},
   Do[sum += vec[[i]], {i, 2, Length[vec]}];
   sum
   ]
  ]

Насколько это быстро?

In[13]:= Do[fun1[data], {10}] // Timing
Out[13]= {4.812, Null}

Вторая, менее очевидная версия:

fun2 = Compile[{{vec, _Real, 1}},
  Module[{sum = vec[[1]]},
   Do[sum += vec[[i]], {i, 2, Length[vec]}];
   sum
   ]
  ]

In[18]:= Do[
  fun2 /@ Transpose[data],
  {10}
  ] // Timing

Out[18]= {1.078, Null}

Как видите, вторая версия намного быстрее. Почему? Поскольку важнейшая операция, sum + = ... - это сложение чисел в fun2 , тогда как это добавление векторов произвольной длины в fun1 .

Вы можете увидеть практическое применение той же «оптимизации» в этом моем ответе , но там, где это уместно, можно привести множество других примеров.

В этом простом примере код, использующий fun2 , не длиннее и не намного сложнее, чем fun1 , но в общем случае вполне может быть.

Как я могу сказать Compile , что один из его аргументов не является произвольной матрицей n * m , а специальной n * 2 или n * 3 один, поэтому он может выполнять эту оптимизацию автоматически, а не использовать общую функцию сложения векторов для добавления крошечных векторов длины 2 или длины 3?

Приложение

Чтобы прояснить, что происходит, мы можем использовать CompilePrint :

CompilePrint [fun1] дает

        1 argument
        5 Integer registers
        5 Tensor registers
        Underflow checking off
        Overflow checking off
        Integer overflow checking on
        RuntimeAttributes -> {}

        T(R2)0 = A1
        I1 = 2
        I0 = 1
        Result = T(R1)3

1   T(R1)3 = Part[ T(R2)0, I0]
2   I3 = Length[ T(R2)0]
3   I4 = I0
4   goto 8
5   T(R1)2 = Part[ T(R2)0, I4]
6   T(R1)4 = T(R1)3 + T(R1)2
7   T(R1)3 = CopyTensor[ T(R1)4]]
8   if[ ++ I4 < I3] goto 5
9   Return

CompilePrint [fun2] дает

        1 argument
        5 Integer registers
        4 Real registers
        1 Tensor register
        Underflow checking off
        Overflow checking off
        Integer overflow checking on
        RuntimeAttributes -> {}

        T(R1)0 = A1
        I1 = 2
        I0 = 1
        Result = R2

1   R2 = Part[ T(R1)0, I0]
2   I3 = Length[ T(R1)0]
3   I4 = I0
4   goto 8
5   R1 = Part[ T(R1)0, I4]
6   R3 = R2 + R1
7   R2 = R3
8   if[ ++ I4 < I3] goto 5
9   Return

Я решил включить это, а не значительно более длинную версию C, где разница во времени еще более заметна.