Fortran - 逻辑索引

假设我有一个矩阵A这是(m x n)和一个向量B这是(m x 1)。这个向量B是一个由 0 和 1 组成的向量。

还让标量s是其中元素的总和B .

我想创建一个矩阵C这是s x n对应于行B等于 1，并且是一个向量D即 s x 1，这些元素的位置在A.

Take as an example: 

      A = [1, 2, 3; 
           4, 5, 6; 
           7, 8, 9; 
           10, 11, 12; 
           13, 14, 15 ]

        B = [0; 1; 0; 1; 1]

    Then:
C = [ 4, 5, 6; 
     10, 11, 12; 
     13, 14, 15 ]
and
D = [2; 
     4
     5]

到目前为止，我的 Fortran 代码如下所示：

PROGRAM test1
IMPLICIT NONE

 REAL, DIMENSION(5,3) :: A
INTEGER, DIMENSION(5,1) :: B = 0
INTEGER :: i, j, k

k = 1

!Create A matrix
do i=1,5
    do j=1,3
        A(i,j) = k
        k = k+1
    end do
end do

!Create B matrix
B(2,1) = 1
B(4,1) = 1
B(5,1) = 1


end program

在 matlab 中，我可以简单地通过创建 C = A(逻辑(B),:) 和 D 以类似的方式创建 C。

我怎样才能在 Fortran 中做到这一点，避免循环？我已经调查过where and forall陈述，但它们并不完全是我正在寻找的。

正如 @francescalus 所建议的，PACK内在函数在 Fortran 中相当于 Matlab 逻辑切片，但只是部分功能。 Matlab 逻辑索引有两种类型：

• 整个数组逻辑索引：掩码必须与数组具有相同的形状，并且返回值的等级为 1（Matlab 中的向量）。之所以如此，是因为真值的位置是任意的，因此你不能保证在矩阵中打孔的结果基本上是矩形的。这就是 PACK 在 Fortran 中所做的事情.

  c = a(a < 1); % Matlab: a(m,n) -> c(s)
  C = PACK(A, A < 1) ! Fortran: A(m,n) -> C(s)
  

• 单维逻辑索引：掩码必须是一维的，用于在数组的单个维度内进行选择。其他维度可以选择整体或带有自己的索引。这就是你想要的.

  b = a(:,1) > 0; % a(m,n) gives logical b(m)
  c = a(b,:); % a(m,n) selected with b(m) -> c(s,n)
  

然而，PACK并不直接承认这种用法。 @high-performance-mark 的解决方案向您展示了如何实现这一壮举：SPREAD基本上是repmat，所以他的解决方案在 Matlab 中如下所示：

b = a(:,1) > 0; % a(m,n) gives logical b(m)
bMat = repmat(b, 1, size(a,2)); % n copies of b(m) -> bMat(m,n)
c = reshape(a(bMat), [sum(b), size(a,2)]); % a(m,n) -> c(s,n)

需要最终重塑的地方，因为a(bMat)由于使用了整体矩阵选择形式，所以不是矩阵，而是大小为 s*n 的向量。另一个答案的 Fortran 代码正是这样做的：

c = RESHAPE( &
        PACK(a, & ! The matrix we are selecting from
            SPREAD(b==1, ! The == 1 is because you were using an INTEGER array, not LOGICAL
                   dim=2, ncopies=SIZE(a,2)) & ! This is the repmat-like code
        ), & ! The result of this PACK is cVec(s*n)
        [COUNT(b==1),SIZE(a,2)]) ! Reshape into (s,n)

分配给 D 的代码非常相似，但我们不是使用 A，而是从动态生成的数组中进行选择，该数组包含从 1 到 A 第一个维度的长度的数字（与 A 的第一维长度相同） B) 的大小。这次不需要对结果进行任何掩码的 SPREAD 或 RESHAPE，因为源数组是一维的。

顺便说一下，Fortrandoes直接支持整数向量索引，因此您可以使用代码首先生成 D 向量（带有 B 的真实元素的索引），然后then do:

C = A(D,:) ! Works the same in Fortran!

避免传播和重塑。