在 m 维上求和 n 维数组的最紧凑方法

Question

假设我有一个n 维数组并且想要对其维度的m 求和，所以最后我得到一个n-m 维数组（n>m）。如果我只想在一个维度上对数组求和，我可以将scalar integer 传递给SUM-intrinsic，例如SUM(array, 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)

但可能需要在重塑时重新排序以减少所需尺寸。这很容易变得不可读或不清楚。