avoid duplicating code by re-structuring if statement/do loop

Question

Hi I am trying to impose a specific condition on my function at many different spatial points in my grid. However I am duplicating lots of code and it's becoming increasingly inefficient.

How can I do what I need to by simply using a do loop? The specific condition I am trying to impose on my function is the same at all the different spatial points so I figure theres a way to do all of this in a single loop. Or how can I combine all these If/else if statements into a single statement? There must be a more efficient way than what I am doing.

I provided a sample code below.

FUNCTION grad(psi)
IMPLICIT NONE
INTEGER :: i,j,kk,ll
INTEGER, PARAMETER :: nx = 24, ny = 24
COMPLEX,DIMENSION(3,3,-nx:nx, -ny:ny) :: psi, grad
REAL :: pi
REAL :: f0
INTEGER :: nxx, nyy

nxx = nx/2
nyy = ny/2

pi = 4*atan(1.0)
f0 = pi**2*1.3

DO i=-nx+1,nx-1 !spatial points
DO j=-ny+1,ny-1 !spatial points

   IF ( i == 0 .AND. j == 0 .AND. i == j) THEN ! I have lots of statements like this

    DO kk=1,3

         grad(kk,1,i,j)  = psi(kk,1,i+1,j) - f0*psi(kk,1,i,j)

         grad(kk,2,i,j)  = psi(kk,2,i+1,j) - f0*psi(kk,2,i,j)

         grad(kk,3,i,j)  = psi(kk,3,i+1,j) - f0*psi(kk,3,i,j)

    END DO

    ELSE IF ( i == nxx .AND. j == nyy .AND. i == j) THEN ! I have lots of statements like this

    DO kk=1,3

         grad(kk,1,i,j)  = psi(kk,1,i+1,j) - f0*psi(kk,1,i,j)

         grad(kk,2,i,j)  = psi(kk,2,i+1,j) - f0*psi(kk,2,i,j)

         grad(kk,3,i,j)  = psi(kk,3,i+1,j) - f0*psi(kk,3,i,j)

    END DO

  ELSE IF ( i == -nxx .AND. j == -nyy .AND. i == j) THEN ! I have lots of statements like this

    DO kk=1,3

         grad(kk,1,i,j)  = psi(kk,1,i+1,j) - f0*psi(kk,1,i,j)

         grad(kk,2,i,j)  = psi(kk,2,i+1,j) - f0*psi(kk,2,i,j)

         grad(kk,3,i,j)  = psi(kk,3,i+1,j) - f0*psi(kk,3,i,j)

    END DO

   ELSE IF ( i == nxx .AND. j == -nyy) THEN ! I have lots of statements like this

    DO kk=1,3

         grad(kk,1,i,j)  = psi(kk,1,i+1,j) - f0*psi(kk,1,i,j)

         grad(kk,2,i,j)  = psi(kk,2,i+1,j) - f0*psi(kk,2,i,j)

         grad(kk,3,i,j)  = psi(kk,3,i+1,j) - f0*psi(kk,3,i,j)

    END DO

   ELSE IF ( i == -nxx .AND. j == nyy) THEN

    DO kk=1,3

         grad(kk,1,i,j)  = psi(kk,1,i+1,j) - f0*psi(kk,1,i,j)

         grad(kk,2,i,j)  = psi(kk,2,i+1,j) - f0*psi(kk,2,i,j)

         grad(kk,3,i,j)  = psi(kk,3,i+1,j) - f0*psi(kk,3,i,j)

    END DO

  ELSE IF ( i == nxx .AND. j == ny) THEN

    DO kk=1,3

         grad(kk,1,i,j)  = psi(kk,1,i+1,j) - f0*psi(kk,1,i,j)

         grad(kk,2,i,j)  = psi(kk,2,i+1,j) - f0*psi(kk,2,i,j)

         grad(kk,3,i,j)  = psi(kk,3,i+1,j) - f0*psi(kk,3,i,j)

    END DO

  ELSE IF ( i == -nxx .AND. j == ny) THEN

    DO kk=1,3

         grad(kk,1,i,j)  = psi(kk,1,i+1,j) - f0*psi(kk,1,i,j)

         grad(kk,2,i,j)  = psi(kk,2,i+1,j) - f0*psi(kk,2,i,j)

         grad(kk,3,i,j)  = psi(kk,3,i+1,j) - f0*psi(kk,3,i,j)

    END DO

   ELSE IF ( i == nx .AND. j == -nyy) THEN

    DO kk=1,3

         grad(kk,1,i,j)  = psi(kk,1,i+1,j) - f0*psi(kk,1,i,j)

         grad(kk,2,i,j)  = psi(kk,2,i+1,j) - f0*psi(kk,2,i,j)

         grad(kk,3,i,j)  = psi(kk,3,i+1,j) - f0*psi(kk,3,i,j)

    END DO

   ELSE IF ( i == nx .AND. j == nyy) THEN

    DO kk=1,3

         grad(kk,1,i,j)  = psi(kk,1,i+1,j) - f0*psi(kk,1,i,j)

         grad(kk,2,i,j)  = psi(kk,2,i+1,j) - f0*psi(kk,2,i,j)

         grad(kk,3,i,j)  = psi(kk,3,i+1,j) - f0*psi(kk,3,i,j)

    END DO

    ELSE 

        DO kk=1,3

        grad(kk,1,i,j)  = psi(kk,1,i+1,j) 

        grad(kk,2,i,j)  = psi(kk,2,i+1,j)

        grad(kk,3,i,j)  = psi(kk,3,i+1,j)

    END DO

   END IF

END DO
END DO

END FUNCTION grad

Answer 1

If you are looking for conciseness, I'd say you can be much, much more concise than you are. The whole function you provided can be rewritten just like this:

function grad(psi)
  implicit none
  integer, parameter :: nx = 24, ny = 24, nxx = nx / 2, nyy = ny / 2
  real, parameter :: pi = 4 * atan(1.0), f0 = pi ** 2 * 1.3
  complex, dimension(3,3,-nx:nx,-ny:ny) :: psi, grad

  grad(:,:,-nx+1:nx-1,-ny+1:ny-1) = psi(:,:,-nx+2:nx,-ny+1:ny-1)
  grad(:,:,0,0) = psi(:,:,1,0)
  grad(:,:,[-nxx,nxx],[-nyy,nyy,ny]) = psi(:,:,[-nxx+1,nxx+1],[-nyy,nyy,ny]) - f0 * psi(:,:,[-nxx,nxx],[-nyy,nyy,ny])
  !grad(:,:,nx,[-nyy,nyy]) = psi(:,:,nx+1,[-nyy,nyy]) - f0 * psi(:,:,nx,[-nyy,nyy])
end

As said by @IanBush, assigning the default values then modifying the special cases seems a good aproach. Also, array sections notation is one of the distinctive and powerful features of Fortran language, and can be used to increase expressiveness without compromising clarity.

Pure colons mean all values in this dimension, and a colon between values means only values within this range in this dimension.

So, when I write grad(:,:,-nx+1:nx-1,-ny+1:ny-1) = psi(:,:,-nx+2:nx,-ny+1:ny-1) I mean: i'm assigning values from the array psi to grad; I include all values from the two first dimensions, but only a subset of the two last dimensions (I'm excluding the fisrt and last in each); also, they are mapped directly except for the third dimension, that maps to the next of the equivalent position in psi.

When I write grad(:,:,[-nxx,nxx],[-nyy,nyy,ny]), I am specifying a list of indices instead of a range for the third and fourth dimensions. This will include the total combinations of the two lists: -nxx,-nyy, -nxx,nyy, -nxx,ny, nxx,-nyy...

One advantage of this notation is that, as it is more obvious and closer to the mathematical notation, it is easier to catch inconsistencies. That is why the last line is commented out: an index nx+1, as you would have in the 8th and 9th conditions in the code you wrote, would be out of bounds. I don't know if the sample code you presented is official; if it is, you should correct your algorithm (well, as you are looping only from the second to the second-last indices, you'd actually never touch those conditions...).

As an additional advice, you could put your custom functions in a module, so you could pass all those parameter declarations to the module scope. Moreover, you could then consider assumed-shape array arguments.