Adding optional INDEX clause to WHERE construct

Often, when working with Fortran sources with WHERE constructs, I have felt handicapped by not being able to evaluate one expression for the subset that satisfy the condition clause and a different expression for the subset that do not, with both expressions having access to the indices of the elements that satisfy or do not satisfy the condition. Here is an example:

program whereIdx
   implicit none
   integer, parameter :: N = 5
   integer i, idx(N),ival(N)
   idx  = [ (i, i=1,N) ]
   ival = [-3, -1, 1, 3, 5]
   where (ival > 2)
      ival = ival+idx*2
   elsewhere
      ival = ival-idx*3
   end where
   print '(5I3)',ival
end program

It would be possible to make the code shorter and more readable if we are able to write, instead, something similar to

program whereIdx
   implicit none
   integer, parameter :: N = 5
   integer i, ival(N)
   integer, allocatable :: idx(:)
   ival = [-3, -1, 1, 3, 5]
   where (ival > 2, INDEX = idx)
      ival = ival+idx*2             ! expression uses indices of .TRUE. elements of ival > 2
   elsewhere
      ival = ival-idx*3             ! expression uses indices of .FALSE. elements of ival > 2
   end where
   print '(5I3)',ival
end program

Your opinions, please. Thanks.

2 Likes

To my taste, both the original and proposed are no better than a loop.

  implicit none
  integer, parameter :: n = 5
  integer :: i, ival(n)
  integer, allocatable :: idx(:)
  ival = [-3, -1, 1, 3, 5]
  do i = 1, n
    if (ival(i) > 2) then
      ival(i) = ival(i) + 2*i
    else
      ival(i) = ival(i) - 3*i
    end if
  end do
  print '(5i3)', ival
end

I’m generally not a fan of the WHERE construct, though, so take my opinion with a grain of salt.

3 Likes

I couldn’t figure out what your code was doing, neither the original, nor the proposed version, until I saw the loop in @nshaffer’s answer, which is clear to me. I personally use where when things are relatively simple, but once it becomes complicated like here, I prefer a loop. A good compiler should be able to make the performance of both equivalent.