arrays - 对 m 维上的 n 维数组求和的最紧凑方法

Question

假设我有一个 n 维数组，并想对其维度的 m 求和，所以最后我得到一个 n-m 维数组 (n>m)。如果我只想对单个维度的数组求和，我可以传递 scalar integer到 SUM - 内在的，例如SUM(数组，DIM = 2)。但是，如果我想对不止一个维度(但不是全部)求和，我没有找到任何信息如何以简单的方式做到这一点，所以我唯一想到的就是链接几个 SUM 命令。例如，假设我想对 3d 数组的两个维度求和:

PROGRAM sum_test

  IMPLICIT NONE

  INTEGER :: a(2,2,2) = reshape( (/1, 2, 3, 4, 5, 6, 7, 8/), (/2,2,2/), order = (/ 1, 2, 3 /) )
  INTEGER :: b(2,2)
  INTEGER :: c(2)
  INTEGER :: d(2)

  d = SUM(SUM(a, DIM=3), DIM=2)

  WRITE (*,*) d

END PROGRAM sum_test

这给出了预期的结果。但是，如果我有更多维度，这可能会变得很难阅读。那么有没有一种“更好”的方法来实现这一目标呢？例如，在 python 的 numpy 中，可以将整数序列传递给 sum 命令。

编辑:

为了测试@francescalus 建议的不同解决方案的性能，我编写了一个小程序并在我的笔记本电脑上使用两个不同的编译器进行了编译(Mac 和 High Sierra，编译器是 g95 和 gfortran-mp-7) 和我可用的两台 linux 机器上。对于每个编译器，我都使用 -O2 作为优化选项。

因为 g95 不支持 do concurrent，所以我忽略了那个版本。我还对求和的等级组合如何影响性能感兴趣，我对一个 4d 数组求和一次超过第二和第四位，一次求和一次超过第三和第四位。这是代码:

PROGRAM sum_performance
  IMPLICIT NONE

  INTEGER, PARAMETER :: size = 100
  REAL :: array(size, size, size, size)
  REAL, DIMENSION(size, size) :: res11, res12, res13, res14
  REAL, DIMENSION(size, size) :: res21, res22, res23, res24

  INTEGER :: i,j,k,l
  INTEGER :: clock_start, clock_stop, clock_rate
  INTEGER :: count
  INTEGER, PARAMETER :: repeat = 100

  !getting clock rate:
  CALL SYSTEM_CLOCK(count_rate = clock_rate)

  !setting up array:
  CALL RANDOM_NUMBER(array)

  !summing second and fourth index
  WRITE (*,*) 'second and fourth rank'

  CALL SYSTEM_CLOCK(count = clock_start)
  DO count = 1,repeat
     res11 = SUM(SUM(array, DIM = 4), DIM = 2)
  END DO
  CALL SYSTEM_CLOCK(count = clock_stop)
  WRITE(*,*) 'chained sum:  ', (REAL(clock_stop - clock_start)/&
       REAL(clock_rate))/REAL(repeat)

  CALL SYSTEM_CLOCK(count = clock_start)
  DO count = 1,repeat
     DO l = 1,size
        DO j = 1,size
           res12(j,l) = SUM(array(:,j,:,l))
        END DO
     END DO
  END DO
  CALL SYSTEM_CLOCK(count = clock_stop)
  WRITE(*,*) 'nested do:    ', (REAL(clock_stop - clock_start)/&
       REAL(clock_rate))/REAL(repeat)

  CALL SYSTEM_CLOCK(count = clock_start)
  DO count = 1,repeat
     FORALL (j = 1:size, l = 1:size)
           res13(j,l) = SUM(array(:,j,:,l))
     END FORALL
  END DO
  CALL SYSTEM_CLOCK(count = clock_stop)
  WRITE(*,*) 'forall:       ', (REAL(clock_stop - clock_start)/&
       REAL(clock_rate))/REAL(repeat)


 !summing second and fourth index
  WRITE (*,*)
  WRITE (*,*) 'third and fourth rank'

  CALL SYSTEM_CLOCK(count = clock_start)
  DO count = 1,repeat
     res21 = SUM(SUM(array, DIM = 4), DIM = 3)
  END DO
  CALL SYSTEM_CLOCK(count = clock_stop)
  WRITE(*,*) 'chained sum:  ', (REAL(clock_stop - clock_start)/&
       REAL(clock_rate))/REAL(repeat)

  CALL SYSTEM_CLOCK(count = clock_start)
  DO count = 1,repeat
     DO l = 1,size
        DO k = 1,size
           res22(k,l) = SUM(array(:,:,k,l))
        END DO
     END DO
  END DO
  CALL SYSTEM_CLOCK(count = clock_stop)
  WRITE(*,*) 'nested do:    ', (REAL(clock_stop - clock_start)/&
       REAL(clock_rate))/REAL(repeat)

  CALL SYSTEM_CLOCK(count = clock_start)
  DO count = 1,repeat
     FORALL (k = 1:size, l = 1:size)
           res23(k,l) = SUM(array(:,:,k,l))
     END FORALL
  END DO
  CALL SYSTEM_CLOCK(count = clock_stop)
  WRITE(*,*) 'forall:       ', (REAL(clock_stop - clock_start)/&
       REAL(clock_rate))/REAL(repeat)

END PROGRAM

结果在很多方面都令人惊讶。对于 mac g95 我得到:

second and fourth rank
chained sum:   0.193214
nested do:     0.140472
forall:        0.18884899

third and fourth rank
chained sum:   0.196938
nested do:     0.114286005
forall:        0.115414

对于 mac gfortran-mp-7 我明白了

second and fourth rank
chained sum:    0.279830009    
nested do:      0.131999999    
forall:         0.130150005    

third and fourth rank
chained sum:     3.01672006    
nested do:      0.111460000    
forall:         0.110610001    

对于两台 linux 机器，相对性能与 mac g95 结果相似，尽管绝对性能不同。如您所见，链式 SUM 解决方案显示出最差的性能(可能因为临时数组创建而导致时间损失？)。此外，在对排名 3 和排名 4 求和时会发生一些有趣的事情(这是一个强大的功能，我检查了几次)。我猜这与临时数组的创建方式有关(如果创建的话)。有趣的是，它只发生在两个编译器之一上。嵌套的 DO 循环和 FORALL 看起来非常相似，所以我想这取决于个人喜好。

Answer 1

我不会低估在这种情况下循环结构的可读性。

很难明智地给出一个通用的食谱，所以我只举几个例子。

问题举例:

do i=1,SIZE(d)
  d(i) = SUM(a(i,:,:))
end do

对于更高排名的结果，可以嵌套这样的 do 结构，或者使用 forall 结构或 do concurrent

forall (i=1:SIZE(d,1), j=1:SIZE(d,2))
  d(i,j) = SUM(a(i,:,:,j)
end

或

do concurrent (i=1:SIZE(d,1), j=1:SIZE(d,2))
  d(i,j) = SUM(a(i,:,:,j)
end

如果确实需要，可以使用 reshape 避免循环/foralls:

d = SUM(RESHAPE(a,[2,4]),dim=2)

但是很可能需要重新排序整形才能减少所需的尺寸。这很容易变得不可读或不清楚。