From fb6775e8a6a1f8deb3b2648f53b7662663c7b714 Mon Sep 17 00:00:00 2001
From: Nick Wogan <nicholaswogan@gmail.com>
Date: Fri, 29 Mar 2024 14:15:01 -0700
Subject: [PATCH] block jacobian

---
 src/forwarddiff_const.f90      |  10 ++
 src/forwarddiff_derivative.f90 | 288 +++++++++++++++++++++++----------
 test/test_sparse.f90           | 112 ++++++++++++-
 3 files changed, 324 insertions(+), 86 deletions(-)

diff --git a/src/forwarddiff_const.f90 b/src/forwarddiff_const.f90
index 7083d18..42dbec3 100644
--- a/src/forwarddiff_const.f90
+++ b/src/forwarddiff_const.f90
@@ -1,4 +1,14 @@
 module forwarddiff_const
   use iso_fortran_env, only: wp => real64
   implicit none
+  public
+
+  enum, bind(c)
+    ! Jacobian Type
+    enumerator :: &
+      DenseJacobian = 1, &
+      BandedJacobian = 2, &
+      BlockDiagonalJacobian = 3
+  end enum
+
 end module
\ No newline at end of file
diff --git a/src/forwarddiff_derivative.f90 b/src/forwarddiff_derivative.f90
index 8ed64e2..1290f5b 100644
--- a/src/forwarddiff_derivative.f90
+++ b/src/forwarddiff_derivative.f90
@@ -78,71 +78,238 @@ subroutine grad(fcn, x, f, dfdx, err)
 
   end subroutine
 
-  subroutine jacobian(fcn, x, f, dfdx, bandwidth, err)
+  subroutine jacobian(fcn, x, f, dfdx, jt, bandwidth, blocksize, err)
+    use forwarddiff_const, only: DenseJacobian, BandedJacobian, BlockDiagonalJacobian
     procedure(jacobian_sig) :: fcn
     real(wp), intent(in) :: x(:)
     real(wp), intent(out) :: f(:)
     real(wp), intent(out) :: dfdx(:,:)
+    integer, optional, intent(in) :: jt
     integer, optional, intent(in) :: bandwidth
+    integer, optional, intent(in) :: blocksize
     character(:), allocatable :: err
 
-    ! real(wp), allocatable :: JacT(:,:) 
-    integer :: bandwidth_, hbw
-    type(dual) :: xx(size(x))
-    type(dual) :: ff(size(x))
-    integer :: i, j, ii, jj, kk
+    integer :: jt_
 
-    bandwidth_ = size(x)
-    if (present(bandwidth)) then
-      bandwidth_ = bandwidth
-      if (bandwidth_ > size(x)) then
-        err = '`bandwidth` can not be > size(x).'
-        return
-      endif
-      if (bandwidth_ < 1) then
-        err = '`bandwidth` can not be < 1.'
+    ! Check dimensions work out
+    if (size(x) /= size(f)) then
+      err = 'Output `f` array is not the right size.'
+      return
+    endif
+
+    ! Determine type of jacobian
+    if (present(jt)) then
+      jt_ = jt
+    else
+      jt_ = DenseJacobian
+    endif
+
+    if (jt_ == DenseJacobian) then
+      call jacobian_dense(fcn, x, f, dfdx, err)
+      if (allocated(err)) return
+    elseif (jt_ == BandedJacobian) then
+      if (.not.present(bandwidth)) then
+        err = '`bandwidth` must be an argument when computing a banded jacobian.'
         return
       endif
-      if (mod(bandwidth_,2) == 0) then
-        err = '`bandwidth` must be odd.'
+      call jacobian_banded(fcn, x, f, dfdx, bandwidth, err)
+      if (allocated(err)) return
+    elseif (jt_ == BlockDiagonalJacobian) then
+      if (.not.present(blocksize)) then
+        err = '`blocksize` must be an argument when computing a block diagonal jacobian.'
         return
       endif
+      call jacobian_blockdiagonal(fcn, x, f, dfdx, blocksize, err)
+      if (allocated(err)) return
     endif
 
-    if (size(x) /= size(f)) then
-      err = 'Output `f` array is not the right size.'
+  end subroutine
+
+  subroutine jacobian_dense(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, j
+
+    if (size(x) /= size(dfdx,1) .or. size(x) /= size(dfdx,2)) then
+      err = 'Output `dfdx` array is not the right size.'
       return
     endif
-    if (bandwidth_ /= size(dfdx,1) .or. size(x) /= size(dfdx,2)) then
-      err = 'Output `dfdx` array is not the right size.'
+
+    ! Set x
+    xx%val = x
+
+    ! Allocate and seed dual
+    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.
+    ! | df(1,1) df(1,2) df(1,3) |
+    ! | df(2,1) df(2,2) df(2,3) |
+    ! | df(3,1) df(3,2) df(3,3) |
+
+    do j = 1,size(x)
+      do i = 1,size(x)
+        dfdx(i,j) = ff(i)%der(j)
+      enddo
+    enddo
+  
+  end subroutine
+
+  subroutine jacobian_banded(fcn, x, f, dfdx, bandwidth, err)
+    procedure(jacobian_sig) :: fcn
+    real(wp), intent(in) :: x(:)
+    real(wp), intent(out) :: f(:)
+    real(wp), intent(out) :: dfdx(:,:)
+    integer, intent(in) :: bandwidth
+    character(:), allocatable :: err
+
+    type(dual) :: xx(size(x))
+    type(dual) :: ff(size(x))
+    integer :: hbw
+    integer :: i, j, ii, jj, kk
+
+    ! Check bandwidth
+    if (bandwidth > size(x)) then
+      err = '`bandwidth` can not be > size(x).'
+      return
+    endif
+    if (bandwidth < 1) then
+      err = '`bandwidth` can not be < 1.'
+      return
+    endif
+    if (mod(bandwidth,2) == 0) then
+      err = '`bandwidth` must be odd.'
       return
     endif
 
-    ! allocate(JacT(size(x),bandwidth_))
+    if (bandwidth /= size(dfdx,1) .or. size(x) /= size(dfdx,2)) then
+      err = 'Output `dfdx` array is not the right size.'
+      return
+    endif
 
     ! Set x
     xx%val = x
 
     ! Allocate dual number
     do i = 1,size(x)
-      allocate(xx(i)%der(bandwidth_))
+      allocate(xx(i)%der(bandwidth))
       xx(i)%der = 0.0_wp
     enddo
 
     ! Seed the dual number
     j = 1
     outer : do
-      do i = 1,bandwidth_
+      do i = 1,bandwidth
         if (j > size(x)) exit outer 
         xx(j)%der(i) = 1.0_wp
         j = j + 1
       enddo
     enddo outer
 
+    ! Do differentiation
+    call fcn(xx, ff)
+
+    ! Unpack f(x)
+    f = ff%val
+
+    ! Unpack banded jacobian
+    ! In this case, we load diagonals of jacobian into each row of dfdx.
+    ! So, dfdx(1,:) is the "highest" diagonal. Illustration:
+    !
+    ! | df(1,1) df(1,2) 0       0       0       |      
+    ! | df(2,1) df(2,2) df(2,3) 0       0       |  ->  | 0       df(1,2) df(2,3) df(3,4) df(4,5) |
+    ! | 0       df(3,2) df(3,3) df(3,4) 0       |  ->  | df(1,1) df(2,2) df(3,3) df(4,4) df(5,5) |
+    ! | 0       0       df(4,3) df(4,4) df(4,5) |  ->  | df(2,1) df(3,2) df(4,3) df(5,4) 0       |
+    ! | 0       0       0       df(5,4) df(5,5) |
+    !
+
+    hbw = (bandwidth - 1)/2 ! halfbandwidth 
+    j = 1
+    outer1 : do
+      do i = 1,bandwidth
+        if (j > size(x)) exit outer1
+        do jj = -hbw,hbw
+          kk = jj + hbw + 1
+          ii = j + jj
+          if (ii < 1) then
+            dfdx(kk,j) = 0.0_wp
+          elseif (ii > size(x)) then
+            dfdx(kk,j) = 0.0_wp
+          else
+            dfdx(kk,j) = ff(ii)%der(i)
+          endif
+        enddo
+        j = j + 1
+      enddo
+    enddo outer1
+
+  end subroutine
+
+  subroutine jacobian_blockdiagonal(fcn, x, f, dfdx, blocksize, err)
+    procedure(jacobian_sig) :: fcn
+    real(wp), intent(in) :: x(:)
+    real(wp), intent(out) :: f(:)
+    real(wp), intent(out) :: dfdx(:,:)
+    integer, intent(in) :: blocksize
+    character(:), allocatable :: err
+
+    type(dual) :: xx(size(x))
+    type(dual) :: ff(size(x))
+    integer :: i, j, ii, jj
+
+    ! Check blocksize
+    if (blocksize > size(x)) then
+      err = '`blocksize` can not be > size(x).'
+      return
+    endif
+    if (blocksize < 1) then
+      err = '`blocksize` can not be < 1.'
+      return
+    endif
+
+    if (mod(size(x),blocksize) /= 0) then
+      err = 'size(x) must be an integer multiple of `blocksize`.'
+      return
+    endif
+
+    if (blocksize /= size(dfdx,1) .or. size(x) /= size(dfdx,2)) then
+      err = 'Output `dfdx` array is not the right size.'
+      return
+    endif
+
+    ! Set x
+    xx%val = x
+
+    ! Allocate dual number
     do i = 1,size(x)
-      print*,xx(i)%der(:)
+      allocate(xx(i)%der(blocksize))
+      xx(i)%der = 0.0_wp
     enddo
-    print*,"hi"
+
+    ! Seed the dual number
+    j = 1
+    outer : do
+      do i = 1,blocksize
+        if (j > size(x)) exit outer 
+        xx(j)%der(i) = 1.0_wp
+        j = j + 1
+      enddo
+    enddo outer
 
     ! Do differentiation
     call fcn(xx, ff)
@@ -150,65 +317,20 @@ subroutine jacobian(fcn, x, f, dfdx, bandwidth, err)
     ! Unpack f(x)
     f = ff%val
 
-    ! Unpack jacobian
-    ! do j = 1,bandwidth_
-    !   do i = 1,size(x)
-    !     JacT(i,j) = ff(i)%der(j)
-    !     ! dfdx(j,i) = ff(i)%der(j)
-    !     ! dfdx(1,:) is the first column of the jacobian
-    !   enddo
-    ! enddo
-    ! JacT(:,1) is the first column of the jacobian
-
-    ! We now want to disentangle the jacobian
-    ! dfdx(1,:) is a row. It should contain diagonals. Row 1 is the highest diagonal
-
-    ! First row of dfdx
-    ! dfdx(1,1:hbw) = 0.0_dp, unused
-    ! JacT(1,hbw+1) should be at dfdx(1,hbw+1)
-    ! JacT(2,hbw+2) should be at dfdx(1,hbw+2)
-    ! JacT(3,hbw) should be at dfdx(1,hbw+3)
-    ! JacT(4,hbw+1) should be at dfdx(1,hbw+4)
-    ! JacT(5,hbw+2) should be at dfdx(1,hbw+5)
-    ! JacT(6,hbw-1) should be at dfdx(1,hbw+6)
-
-    ! Second row
-
-    ! JacT(1,bw) is first diagonal entry
-
-    if (present(bandwidth)) then
-
-      hbw = (bandwidth_ - 1)/2 ! halfbandwidth 
-      j = 1
-      outer1 : do
-        do i = 1,bandwidth_
-          if (j > size(x)) exit outer1
-          do jj = -hbw,hbw
-            kk = jj + hbw + 1
-            ii = j + jj
-            if (ii < 1) then
-              dfdx(kk,j) = 0.0_wp
-            elseif (ii > size(x)) then
-              dfdx(kk,j) = 0.0_wp
-            else
-              ! dfdx(kk,j) = JacT(ii,i)
-              dfdx(kk,j) = ff(ii)%der(i)
-            endif
-          enddo
-          j = j + 1
-        enddo
-      enddo outer1
-
-    else
-      
-      ! Dense matrix
-      do j = 1,bandwidth_
-        do i = 1,size(x)
-          dfdx(i,j) = ff(i)%der(j)
+    ! Unpack block jacobian
+    !
+    ! | df(1,1) df(1,2) 0       0       |      
+    ! | df(2,1) df(2,2) 0       0       |  ->  | df(1,1) df(1,2) df(3,3) df(3,4) |
+    ! | 0       0       df(3,3) df(3,4) |  ->  | df(2,1) df(2,2) df(4,3) df(4,4) |
+    ! | 0       0       df(4,3) df(4,4) |
+    do ii = 1,size(x)/blocksize
+      jj = (ii - 1)*blocksize
+      do i = 1,blocksize
+        do j = 1,blocksize
+          dfdx(i,j+jj) = ff(i+jj)%der(j)
         enddo
       enddo
-    
-    endif
+    enddo
 
   end subroutine
 
diff --git a/test/test_sparse.f90 b/test/test_sparse.f90
index f5a2190..a3d1ce5 100644
--- a/test/test_sparse.f90
+++ b/test/test_sparse.f90
@@ -2,13 +2,16 @@
 program main
   use forwarddiff, only: wp, jacobian
   implicit none
-  integer, parameter :: nz = 7
+  integer, parameter :: nz = 6
 
-  call test_banded()
+  ! call test_banded()
+  call test_blockdiagonal1()
+  call test_blockdiagonal2()
 
 contains
 
   subroutine test_banded()
+    use forwarddiff, only: BandedJacobian
     integer, parameter :: bandwidth = 3
     real(wp) :: u(nz), f(nz), dfdu(bandwidth,nz)
     real(wp) :: f1(nz), dfdu1(nz,nz)
@@ -18,7 +21,7 @@ subroutine test_banded()
     do i = 1,nz
       u(i) = i
     enddo
-    call jacobian(rhs_banded_dual, u, f, dfdu, bandwidth=bandwidth, err=err)
+    call jacobian(rhs_banded_dual, u, f, dfdu, jt=BandedJacobian, bandwidth=bandwidth, err=err)
     if (allocated(err)) then
       print*,err
       stop
@@ -83,4 +86,107 @@ subroutine jac_banded(u, pd)
     
   end subroutine
 
+  subroutine test_blockdiagonal1()
+    use forwarddiff, only: BlockDiagonalJacobian
+    integer, parameter :: blocksize = 2
+    real(wp) :: u(nz), f(nz), dfdu(blocksize,nz)
+    real(wp) :: dfdu1(nz,nz)
+    character(:), allocatable :: err
+    integer :: i
+
+    do i = 1,nz
+      u(i) = i
+    enddo
+    call jacobian(rhs_blocked1_dual, u, f, dfdu, jt=BlockDiagonalJacobian, blocksize=blocksize, err=err)
+    if (allocated(err)) then
+      print*,err
+      stop 1
+    endif
+
+    print*,''
+    do i = 1,blocksize
+      print*,dfdu(i,:) 
+    enddo
+
+    call jacobian(rhs_blocked1_dual, u, f, dfdu1, err=err)
+    if (allocated(err)) then
+      print*,err
+      stop 1
+    endif
+
+    print*,''
+    do i = 1,nz
+      print*,dfdu1(i,:) 
+    enddo
+
+  end subroutine
+
+  subroutine test_blockdiagonal2()
+    use forwarddiff, only: BlockDiagonalJacobian
+    integer, parameter :: blocksize = 3
+    real(wp) :: u(nz), f(nz), dfdu(blocksize,nz)
+    real(wp) :: dfdu1(nz,nz)
+    character(:), allocatable :: err
+    integer :: i
+
+    do i = 1,nz
+      u(i) = i
+    enddo
+    call jacobian(rhs_blocked2_dual, u, f, dfdu, jt=BlockDiagonalJacobian, blocksize=blocksize, err=err)
+    if (allocated(err)) then
+      print*,err
+      stop 1
+    endif
+
+    print*,''
+    do i = 1,blocksize
+      print*,dfdu(i,:) 
+    enddo
+
+    call jacobian(rhs_blocked2_dual, u, f, dfdu1, err=err)
+    if (allocated(err)) then
+      print*,err
+      stop 1
+    endif
+
+    print*,''
+    do i = 1,nz
+      print*,dfdu1(i,:) 
+    enddo
+
+  end subroutine
+
+  subroutine rhs_blocked1_dual(u, du)
+    use forwarddiff
+    type(dual), intent(in) :: u(:)
+    type(dual), intent(out) :: du(:)
+    integer :: i
+
+    du(1) = u(1) + u(2)*u(1)
+    du(2) = u(2) + u(2)*u(1)
+
+    du(3) = u(3)*u(4) + u(3)
+    du(4) = u(4) + u(3)*u(3)
+
+    du(5) = 2.0_wp*u(5) + 0.5_wp*u(6)
+    du(6) = u(5) + u(6)*u(6)
+
+  end subroutine
+
+  subroutine rhs_blocked2_dual(u, du)
+    use forwarddiff
+    type(dual), intent(in) :: u(:)
+    type(dual), intent(out) :: du(:)
+    integer :: i
+
+    du(1) = u(1) + u(2)*u(1)
+    du(2) = u(2) + u(2)*u(1)
+    du(3) = u(3) + u(3)*u(1)
+
+    du(4) = u(4) + u(5)*u(6)
+    du(5) = 2.0_wp*u(5) + 0.5_wp*u(6)
+    du(6) = u(5) + u(6)*u(6)
+
+  end subroutine
+
 end program
\ No newline at end of file