在 Fortran 的 where 块内循环

Question

Even though I do not exactly know why, it seems that Fortran (90) does not allow do loops inside where blocks. A code structured as follows does not compile with gfortran ( Unexpected DO statement in WHERE block at (1) ):

real :: a(30), b(30,10,5)
integer :: i, j

do i=1,10
  where(a(:) > 0.)
    ! do lots of calculations
    ! modify a
    do j=1,5
      b(:,i,j)=...
    enddo
  endwhere
enddo

The only workaround that I can think of would be

real :: a2(30)
do i=1,10
  a2(:)=a(:)
  where(a(:) > 0.)
    ! do lots of calculations
    ! modify a
  endwhere
  
  do j=1,5
    where(a2(:) > 0.)
      b(:,i,j)=...
    endwhere
  enddo
enddo

I suppose that there are more elegant solutions? Especially if the where condition is less straightforward, this will look messy pretty soon... Thanks!

Answer 1

If your arrays are all 1-indexed, you can replace where constructs by explicitly storing and using array masks, for example

program test
  implicit none
  
  real :: a(30), b(30,10,5)
  integer :: i, j
  
  integer, allocatable :: mask(:)
  
  do i=1,10
    mask = filter(a(:)>0.)
    ! do lots of calculations
    ! modify a
    do j=1,5
      b(mask,i,j)=...
    enddo
  enddo
contains

! Returns the array of indices `i` for which `input(i)` is `true`.
function filter(input) result(output)
  logical, intent(in)  :: input(:)
  integer, allocatable :: output(:)
  
  integer :: i
  
  output = pack([(i,i=1,size(input))], mask=input)
end function
end program