Expressing finite difference as a function in the tsunami simulator

You now understand how to define a function and how to call it from the main program. Finally, we get to the fun part–applying our new knowledge about functions to refactor our tsunami simulator. Let’s look back at the main time loop in our program, as reprised in the following listing.

Listing 3.9 The time integration loop from the minimal working tsunami simulator

time_loop: do n = 1, num_time_steps     ❶
 
  dh(1) = h(1) - h(grid_size)           ❷
 
  do concurrent (i = 2:grid_size)       ❸
    dh(i) = h(i) - h(i-1)               ❸
  end do                                ❸
 
  do concurrent (i = 1:grid_size)       ❹
    h(i) = h(i) - c * dh(i) / dx * dt   ❹
  end do                                ❹
 
end do time_loop

❶ Iterates for num_time_steps time steps

❷ Calculates the difference on the left boundary

❸ Calculates the difference in the rest of the domain

❹ Computes and stores the value of h at the next time step

At the beginning of this chapter, I mentioned that we’ll use the finite difference calculation quite a bit as we move toward a more realistic wave simulator. A good first step, then, may be to replace the following

dh(1) = h(1) - h(grid_size)
 
do concurrent (i = 2:grid_size)
  dh(i) = h(i) - h(i-1)
end do

with a function call like this:

dh = diff(h)

In a nutshell, we’ll pack both the finite difference calculation (dh(i) = h(i) - h(i-1)) and the handling of the boundary condition (dh(1) = h(1) - h(grid_size)) into a single function diff that we can reuse whenever needed. This will be quite useful down the road as we add more physics terms to our solver. If coded correctly, the new program will output exactly the same results as the original version. Our time_loop in the main program should now look like the following listing.

Listing 3.10 Delegating the finite differencing to a function

time_loop: do n = 1, num_time_steps
 
  dh = diff(h)                           ❶
  do concurrent (i = 1:grid_size)        ❷
    h(i) = h(i) - c * dh(i) / dx * dt    ❷
  end do                                 ❷
 
  print *, n, h                          ❸
 
end do time_loop

❶ Calculates the difference in a function

❷ Computes and stores the new value of h

❸ Write the values to screen

And the following listing shows the definition of the diff function.

Listing 3.11 Finite difference calculation expressed as a function

function diff(x) result(dx)
  real, intent(in) :: x(:)            ❶
  real :: dx(size(x))                 ❷
  integer :: im
  im = size(x)
  dx(1) = x(1) - x(im)                ❸
  dx(2:im) = x(2:im) - x(1:im-1)      ❹
end function diff

❶ Assumed-shape real array as input argument

❷ The result will be a real array of the same size as x.

❸ Calculates the boundary value

❹ Calculates the finite difference for all other elements of x

We’re now calculating the difference in space in the function and are down to only one do loop inside the main time loop. Before we move on to subroutines, I’ll give you a sneak peek into one of Fortran’s most powerful features–its array-oriented syntax. While I only mentioned this briefly in chapter 1 when discussing the strengths and weaknesses of Fortran, we haven’t had the opportunity yet to cast arithmetic operations on whole arrays. We’ll go into more depth with everything about arrays in chapter 5, but for now, let’s rewrite the main time loop to greatly simplify it, as the following listing demonstrates.

Listing 3.12 Solving the advection equation with a single expression

time_loop: do n = 1, num_time_steps
  h = h - c * diff(h) / dx * dt       ❶
  print *, n, h                       ❷
end do time_loop

❶ Invokes diff(h) directly to update the new value of h

❷ Writes the values to screen

Now this is pretty sweet! We have a solver that not only fits in a single line of code, but also appears almost exactly the same as our original math equation. The internal details of the finite difference calculation are now hidden in the implementation of the function diff, and here we simply call it to calculate the difference when we need it. Substituting a whole loop with an array operation is possible because h and diff(h) are of the same shape (one-dimensional) and size. Notice also that c, dx, and dt are all scalar variables, and they’re compatible with array operations. Stay tuned for more in chapter 5.