Fortran - Logical Indexing

Suppose I have a matrix A which is (m x n) and a vector B which is (m x 1). This vector B is a vector of zeros and ones.

Also let the scalar s be the sum of the elements in B .

I want to create a matrix C which is s x n corresponding to the rows of B which equal 1, and a vector D which is s x 1, with the position of those elements in 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]

As of now my fortran code looks like:

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

In matlab I could create C simply by making: C = A(logical(B),:), and D in a similar way.

How can I do it in fortran, avoiding loops? I have looked into the where and forall statements, but they are not exactly what I am looking for.

Solution

As @francescalus suggested, the PACK intrinsic function is the Fortran equivalent of Matlab logical slicing, but only in part. There are two types of Matlab logical indexing:

Whole array logical indexing: the mask must have the same shape as the array and the returned value is of rank 1 (a vector in Matlab). This is so because the position of the trues is arbitrary, and thus you cannot guarantee that the result of, basically, poking holes into a matrix will be rectangular. This is what PACK does in Fortran.
```
c = a(a < 1); % Matlab: a(m,n) -> c(s)
C = PACK(A, A < 1) ! Fortran: A(m,n) -> C(s)
```
Single-dimension logical indexing: the mask must be 1-D, and is used to select inside a single dimension of the array. Other dimensions can be selected whole or with indexing of their own. This is what you want.
```
b = a(:,1) > 0; % a(m,n) gives logical b(m)
c = a(b,:); % a(m,n) selected with b(m) -> c(s,n)
```

However, PACK does not admit this usage directly. @high-performance-mark's solution shows you how to perform this feat: SPREAD is basically repmat, so his solution would look like this in 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)

Where the final reshape is required because a(bMat) is not a matrix, but a vector of size s*n due to the usage of the whole-matrix selection form. The other answer's Fortran code is doing exactly that:

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)

The code that assigns to D is very similar, but instead of using A we are selecting from an array that is generated on-the-fly to contains numbers from 1 to the length of the first dimension of A (which is the same as the size of B). This time there does not need to be any SPREAD of the mask or RESHAPE to the result because the source array is 1-D.

Incidentally, Fortran does support integer-vector indexing directly, so you could use the code to generate the D vector (with the indices of the true elements of B) first and then do:

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

saving yourself the spreading and reshaping.