Search code examples
arraysfunctionfortranfortran90

Defining a function returning an array


I have the following code:

    Program function_as_an_array
    implicit none
    integer:: i
    integer, parameter:: N=10
    real*8:: x(N),y(N),f(N)

    do i=1,N
      x(i)=float(i)
    end do

    call func(f,N,x)

    open(unit=20, file='test.dat')
    do i=1,N
      y(i)=f(i)
      write(20,*) x(i),y(i) 
    end do
    close(20)
    Stop 
    End Program function_as_an_array


    Subroutine func(f,N,x)
    implicit none
    integer i,N
    real*8:: x(N),f(N) 

    do i=1,N
       f(i)=x(i)**2
    end do

    end Subroutine func

I want to make the program indeed be meant for "function as an arrray", i.e. I would like to replace the Subroutine func by a function f and get the same result (In the main program, I wish to keep a statement like y=f(x,N)). How can I do that?


Solution

  • There's no problem having a function return an array, as with this question and answer: the main issue is that you need the function to be in a module (or contained within the program) so that there's an automatic explicit interface: (Edit to add: or explicitly defining the interface as with Alexander Vogt's answer)

    module functions
    contains
    
        function func(N,x)
        implicit none
        integer, intent(in) :: N
        double precision, intent(in) :: x(N)
        double precision, dimension(N) :: func
    
        integer :: i
    
        do i=1,N
           func(i)=x(i)**2
        end do
    
    end function func
    
    end module functions
    
    Program function_as_an_array
    use functions
    implicit none
    integer:: i
    integer, parameter:: N=10
    double precision:: x(N),y(N)
    
    do i=1,N
      x(i)=float(i)
    end do
    
    y = func(N,x)
    
    open(unit=20, file='test.dat')
    do i=1,N
      write(20,*) x(i),y(i)
    end do
    close(20)
    Stop
    End Program function_as_an_array
    

    But note that this sort of function - applying the same operation to every element in array - is somewhat more nicely done with a Fortran elemental function, defined to work simply on a scalar and Fortran will automatically map it over all elements of an array for you:

    module functions
    contains
    
        elemental double precision function f(x)
        implicit none
        double precision, intent(in) :: x
    
        f = x**2
    
        end function f
    
    end module functions
    
    Program function_as_an_array
        use functions
        implicit none
        integer:: i
        integer, parameter:: N=10
        double precision:: x(N),y(N)
    
        do i=1,N
          x(i)=float(i)
        end do
    
        y = f(x)
    
        open(unit=20, file='test.dat')
        do i=1,N
          write(20,*) x(i),y(i)
        end do
        close(20)
        Stop
    End Program function_as_an_array
    

    The nice thing about this is that it will now work on scalars, and arrays of any rank automatically. Wherever possible, it's good to have the compiler do your work for you.