dense jacobian

src/CMakeLists.txt

@@ -4,3 +4,9 @@ add_library(forwarddiff
   forwarddiff_derivative.f90
   forwarddiff.f90
 )
+if ("${CMAKE_Fortran_COMPILER_ID}" MATCHES "GNU")
+  target_compile_options(forwarddiff PRIVATE -Wunused -Wimplicit-interface -fimplicit-none)
+  if (CMAKE_BUILD_TYPE STREQUAL "Debug")
+    target_compile_options(forwarddiff PRIVATE -fcheck=all,no-array-temps)
+  endif()
+endif()

src/forwarddiff_derivative.f90

@@ -6,6 +6,7 @@ module forwarddiff_derivative
 
   public :: derivative, derivative_sig
   public :: grad, grad_sig
+  public :: jacobian, jacobian_sig
 
   abstract interface
     function derivative_sig(x) result(res)
@@ -19,41 +20,102 @@ function grad_sig(x) result(res)
       type(dual), intent(in) :: x(:)
       type(dual) :: res
     end function
+
+    subroutine jacobian_sig(x, f)
+      import :: dual
+      type(dual), intent(in) :: x(:)
+      type(dual), intent(out) :: f(:)
+    end subroutine
   end interface
 
 contains
 
-  subroutine derivative(fcn, x, fcn_x, dfcn_x)
+  subroutine derivative(fcn, x, f, dfdx)
     procedure(derivative_sig) :: fcn
     real(wp), intent(in) :: x
-    real(wp), intent(out) :: fcn_x
-    real(wp), intent(out) :: dfcn_x
-    type(dual) :: f
-    f = fcn(dual(x, [1.0_wp]))
-    fcn_x = f%val
-    dfcn_x = f%der(1)
+    real(wp), intent(out) :: f
+    real(wp), intent(out) :: dfdx
+    type(dual) :: ff
+    ff = fcn(dual(x, [1.0_wp]))
+    f = ff%val
+    dfdx = ff%der(1)
   end subroutine
 
-  subroutine grad(fcn, x, fcn_x, dfcn_x)
+  subroutine grad(fcn, x, f, dfdx, err)
     procedure(grad_sig) :: fcn
     real(wp), intent(in) :: x(:)
-    real(wp), intent(out) :: fcn_x
-    real(wp), intent(out) :: dfcn_x(:)
+    real(wp), intent(out) :: f
+    real(wp), intent(out) :: dfdx(:)
+    character(:), allocatable :: err
 
     type(dual) :: xx(size(x))
-    type(dual) :: f
+    type(dual) :: ff
     integer :: i
 
+    if (size(x) /= size(dfdx)) then
+      err = 'Output dfdx array is not the right size.'
+      return
+    endif
+
+    ! Set x
     xx%val = x
+
+    ! Seed the dual number
     do i = 1,size(x)
       allocate(xx(i)%der(size(x)))
       xx(i)%der = 0.0_wp
       xx(i)%der(i) = 1.0_wp
     enddo
 
-    f = fcn(xx)
-    fcn_x = f%val
-    dfcn_x = f%der(:)
+    ! Do differentiation
+    ff = fcn(xx)
+
+    ! Unpack f(x)
+    f = ff%val
+
+    ! Unpack gradient
+    dfdx(:) = ff%der(:)
+
+  end subroutine
+
+  subroutine jacobian(fcn, x, f, dfdx, err)
+    procedure(jacobian_sig) :: fcn
+    real(wp), intent(in) :: x(:)
+    real(wp), intent(out) :: f(:)
+    real(wp), intent(out) :: dfdx(:,:)
+    character(:), allocatable :: err
+
+    type(dual) :: xx(size(x))
+    type(dual) :: ff(size(x))
+    integer :: i
+
+    if (size(x) /= size(f)) then
+      err = 'Output f array is not the right size.'
+      return
+    endif
+    if (size(x) /= size(dfdx,1) .or. size(x) /= size(dfdx,2)) then
+      err = 'Output dfdx array is not the right size.'
+      return
+    endif
+
+    xx%val = x
+    do i = 1,size(x)
+      allocate(xx(i)%der(size(x)))
+      xx(i)%der = 0.0_wp
+      xx(i)%der(i) = 1.0_wp
+    enddo
+
+    ! Do differentiation
+    call fcn(xx, ff)
+
+    ! Unpack f(x)
+    f = ff%val
+
+    ! Unpack jacobian
+    do i = 1,size(x)
+      dfdx(i,:) = ff(i)%der(:)
+    enddo
+
   end subroutine
 
 

test/test_forwarddiff.f90

@@ -1,13 +1,15 @@
 program test_forwarddiff
-  use forwarddiff, only: wp, derivative, grad
+  use forwarddiff, only: wp, derivative, grad, jacobian
   implicit none
-  call test()
+  call test_dual()
+  call test_jacobian1()
 
 contains
 
-  subroutine test()
+  subroutine test_dual()
     real(wp) :: x, f, dfdx
     real(wp) :: xx(2), dfdx1(2)
+    character(:), allocatable :: err
 
     open(unit=2,file='test.dat',status='replace',form='unformatted')
 
@@ -27,12 +29,12 @@ subroutine test()
     write(2) f, dfdx
 
     xx = [1.0_wp, 2.0_wp]
-    call grad(func_grad1, xx, f, dfdx1)
+    call grad(func_grad1, xx, f, dfdx1, err)
     print*,f, dfdx1
     write(2) f, dfdx1
 
     xx = [3.0_wp, 4.0_wp]
-    call grad(func_grad2, xx, f, dfdx1)
+    call grad(func_grad2, xx, f, dfdx1, err)
     print*,f, dfdx1
     write(2) f, dfdx1
 
@@ -100,4 +102,75 @@ function func_grad2(x) result(res)
     res = sum(x*3.14_wp)
   end function
 
+  subroutine test_jacobian1()
+    real(wp) :: u(3), f(3), dfdu(3,3)
+    real(wp) :: f1(3), dfdu1(3,3)
+    character(:), allocatable :: err
+
+    u = [1.0_wp, 2.0_wp, 3.0_wp]
+    call jacobian(rhs_rober_dual, u, f, dfdu, err)
+
+    print*,f
+    print*,''
+    print*,dfdu(:,1)
+    print*,dfdu(:,2)
+    print*,dfdu(:,3)
+
+    ! Check against analytical solution
+    call rhs_rober(u, f1)
+    call jac_rober(u, dfdu1)
+
+    if (.not. all(f == f1)) then
+      error stop "test_jacobian1 failed"
+    endif
+    if (.not. all(dfdu == dfdu1)) then
+      error stop "test_jacobian1 failed"
+    endif
+
+  end subroutine
+
+  subroutine rhs_rober_dual(u, du)
+    use forwarddiff_dual
+    type(dual), intent(in) :: u(:)
+    type(dual), intent(out) :: du(:)
+
+    real(wp), parameter :: k1 = 0.04_wp, &
+                           k2 = 3.0e7_wp, &
+                           k3 = 1.0e4_wp
+
+    du(1) = -k1*u(1) + k3*u(2)*u(3)
+    du(2) =  k1*u(1) - k2*u(2)**2.0_wp - k3*u(2)*u(3)
+    du(3) =  k2*u(2)**2.0_wp
+
+  end subroutine
+
+  subroutine rhs_rober(u, du)
+    use forwarddiff_dual
+    real(wp), intent(in) :: u(:)
+    real(wp), intent(out) :: du(:)
+
+    real(wp), parameter :: k1 = 0.04_wp, &
+                           k2 = 3.0e7_wp, &
+                           k3 = 1.0e4_wp
+
+    du(1) = -k1*u(1) + k3*u(2)*u(3)
+    du(2) =  k1*u(1) - k2*u(2)**2.0_wp - k3*u(2)*u(3)
+    du(3) =  k2*u(2)**2.0_wp
+
+  end subroutine
+
+  subroutine jac_rober(u, pd)
+    real(wp), intent(in) :: u(:)
+    real(wp), intent(out) :: pd(:,:)
+
+    real(wp), parameter :: k1 = 0.04_wp, &
+                           k2 = 3.0e7_wp, &
+                           k3 = 1.0e4_wp
+
+    pd(:,1) = [-k1, k1, 0.0_wp]
+    pd(:,2) = [k3*u(3), -2.0_wp*k2*u(2) - k3*u(3), 2.0_wp*k2*u(2)]
+    pd(:,3) = [k3*u(2), -k3*u(2), 0.0_wp]
+
+  end subroutine
+
 end program