feat: add non-autonomous stiff3 overload with dFdx callback and test

Agent-Logs-Url: https://github.com/ivan-pi/stiff3/sessions/67f692d9-43f0-4337-94d3-193b5edcdf71

Co-authored-by: ivan-pi <21085643+ivan-pi@users.noreply.github.com>

Author: Copilot

README.md

@@ -1,6 +1,6 @@
 # stiff3
 
-`stiff3` is a Fortran subprogram for solving stiff autonomous systems of ordinary differential equations (ODE's) using a semi-implicit Runge-Kutta method with three steps (SIRK3). The `stiff3` source code was originally published in the work:
+`stiff3` is a Fortran subprogram for solving stiff systems of ordinary differential equations (ODE's) using a semi-implicit Runge-Kutta method with three steps (SIRK3). Both autonomous and non-autonomous systems are supported through overloaded `stiff3` interfaces. The `stiff3` source code was originally published in the work:
 
 > Villadsen, J., & Michelsen, M. L. (1978). *Solution of differential equation models by polynomial approximation*. Prentice-Hall, Inc.
 
@@ -111,6 +111,5 @@ For students interested in CSE, here are some contribution ideas:
 - Support for banded or sparse Jacobian matrices
 - Use BLAS kernels for vector operations
 - Continuous (dense) output of variables
-- Extend `stiff3` to non-autonomous systems of ODE's
 - Advanced stepsize control settings
-- Write a tutorial on how to use `stiff3`
+- Write a tutorial on how to use `stiff3`

src/stiff3_solver.f90

@@ -18,13 +18,19 @@ module stiff3_solver
 
   public :: stiff3, stiff3_wp
   public :: rhs_sub, jacobian_sub, output_sub
+  public :: rhs_nonaut_sub, jacobian_nonaut_sub, rhs_time_derivative_sub
 
   !> Kind parameter for working precision of stiff3 reals
   integer, parameter :: stiff3_wp = kind(1.0d0)
 
   !> Working precision with short name for internal use
   integer, parameter :: wp = stiff3_wp
 
+  interface stiff3
+    module procedure stiff3_autonomous
+    module procedure stiff3_nonautonomous
+  end interface
+
   abstract interface
     !> Function to evaluate the right-hand side of a system of ODEs.
     !> It is assumed the system of ODEs is autonomous, meaning that
@@ -58,6 +64,33 @@ subroutine output_sub(x,y,iha,qa)
         !! Step-length acceleration factor
     end subroutine
 
+    !> User supplied subprogram for non-autonomous rhs evaluation.
+    subroutine rhs_nonaut_sub(n,x,y,f)
+      import wp
+      integer, intent(in) :: n
+      real(wp), intent(in) :: x
+      real(wp), intent(in) :: y(n)
+      real(wp), intent(inout) :: f(n)
+    end subroutine
+
+    !> User supplied subprogram for non-autonomous Jacobian evaluation.
+    subroutine jacobian_nonaut_sub(n,x,y,df)
+      import wp
+      integer, intent(in) :: n
+      real(wp), intent(in) :: x
+      real(wp), intent(in) :: y(n)
+      real(wp), intent(inout) :: df(n,n)
+    end subroutine
+
+    !> User supplied subprogram for evaluation of dF/dx.
+    subroutine rhs_time_derivative_sub(n,x,y,fx)
+      import wp
+      integer, intent(in) :: n
+      real(wp), intent(in) :: x
+      real(wp), intent(in) :: y(n)
+      real(wp), intent(inout) :: fx(n)
+    end subroutine
+
   end interface
 
 
@@ -68,7 +101,7 @@ subroutine output_sub(x,y,iha,qa)
 
   !> Semi-implicit Runge-Kutta integrator routine
   !
-  subroutine stiff3(n,fun,dfun,out,nprint,x0,x1,h0,eps,w,y)
+  subroutine stiff3_autonomous(n,fun,dfun,out,nprint,x0,x1,h0,eps,w,y)
     integer, intent(in) :: n
       !! Number of equations to be integrated.
     procedure(rhs_sub) :: fun
@@ -228,6 +261,74 @@ subroutine stiff3(n,fun,dfun,out,nprint,x0,x1,h0,eps,w,y)
 
   end subroutine
 
+  !> Semi-implicit Runge-Kutta integrator routine for non-autonomous systems
+  !
+  subroutine stiff3_nonautonomous(n,fun,dfun,dfdx,out,nprint,x0,x1,h0,eps,w,y)
+    integer, intent(in) :: n
+      !! Number of equations to be integrated.
+    procedure(rhs_nonaut_sub) :: fun
+      !! User supplied subprogram for function evaluation.
+    procedure(jacobian_nonaut_sub) :: dfun
+      !! User supplied subprogram for evaluation of the Jacobian wrt y.
+    procedure(rhs_time_derivative_sub) :: dfdx
+      !! User supplied subprogram for evaluation of dF/dx.
+    procedure(output_sub) :: out
+      !! User supplied subprogram for output.
+    integer, intent(in) :: nprint
+      !! Printing interval. For `nprint = k` the solution is only printed.
+      !! at every kth step.
+    real(wp), intent(in) :: x0, x1
+      !! Limits of the independent variable between which the differential
+      !! equation is solved.
+    real(wp), intent(inout) :: h0
+      !! Suggested initial half-step length. On exit `h0` contains suggested
+      !! value of half-step length for continued integration beyond `x1`.
+    real(wp), intent(in) :: eps, w(n)
+      !! Tolerance parameters.
+    real(wp), intent(inout) :: y(n)
+      !! Vector of dependent variables at `x0`. On exit `y` is the vector of
+      !! dependent variables at `x1`.
+
+    real(wp) :: yaug(n+1), waug(n+1)
+
+    yaug(1) = x0
+    yaug(2:) = y
+
+    waug(1) = 0.0_wp
+    waug(2:) = w
+
+    call stiff3_autonomous(n+1,fun_aug,jac_aug,out_aug,nprint,x0,x1,h0,eps,waug,yaug)
+    y = yaug(2:)
+
+  contains
+
+    subroutine fun_aug(naug,yi,fi)
+      integer, intent(in) :: naug
+      real(wp), intent(in) :: yi(naug)
+      real(wp), intent(inout) :: fi(naug)
+      call fun(n,yi(1),yi(2:),fi(2:))
+      fi(1) = 1.0_wp
+    end subroutine
+
+    subroutine jac_aug(naug,yi,dfi)
+      integer, intent(in) :: naug
+      real(wp), intent(in) :: yi(naug)
+      real(wp), intent(inout) :: dfi(naug,naug)
+      dfi = 0.0_wp
+      call dfdx(n,yi(1),yi(2:),dfi(2:,1))
+      call dfun(n,yi(1),yi(2:),dfi(2:,2:))
+    end subroutine
+
+    subroutine out_aug(x,yi,iha,qa)
+      real(wp), intent(in) :: x
+      real(wp), intent(in) :: yi(:)
+      integer, intent(in) :: iha
+      real(wp), intent(in) :: qa
+      call out(x,yi(2:),iha,qa)
+    end subroutine
+
+  end subroutine
+
 
   !> Single-step semi-implicit integration
   !

test/nonautonomous.f90

@@ -0,0 +1,58 @@
+program nonautonomous
+
+  use stiff3_solver, only: stiff3, wp => stiff3_wp
+
+  implicit none
+
+  integer, parameter :: n = 1, nout = 1000000
+  real(wp), parameter :: atol = 5.0e-6_wp
+  real(wp) :: y(n), w(n), x0, x1, h0, eps, y_exact
+
+  y = [0.0_wp]
+  w = [1.0_wp]
+  x0 = 0.0_wp
+  x1 = 1.0_wp
+  h0 = 1.0e-3_wp
+  eps = 1.0e-8_wp
+
+  call stiff3(n,fun,jac,dfdx,out,nout,x0,x1,h0,eps,w,y)
+
+  y_exact = 0.5_wp*(sin(x1) - cos(x1) + exp(-x1))
+  if (abs(y(1)-y_exact) > atol) then
+    error stop 1
+  end if
+
+contains
+
+  subroutine fun(n,t,y,f)
+    integer, intent(in) :: n
+    real(wp), intent(in) :: t
+    real(wp), intent(in) :: y(n)
+    real(wp), intent(inout) :: f(n)
+    f(1) = -y(1) + sin(t)
+  end subroutine
+
+  subroutine jac(n,t,y,df)
+    integer, intent(in) :: n
+    real(wp), intent(in) :: t
+    real(wp), intent(in) :: y(n)
+    real(wp), intent(inout) :: df(n,n)
+    df(1,1) = -1.0_wp
+  end subroutine
+
+  subroutine dfdx(n,t,y,fx)
+    integer, intent(in) :: n
+    real(wp), intent(in) :: t
+    real(wp), intent(in) :: y(n)
+    real(wp), intent(inout) :: fx(n)
+    fx(1) = cos(t)
+  end subroutine
+
+  subroutine out(t,y,ih,qa)
+    real(wp), intent(in) :: t
+    real(wp), intent(in) :: y(:)
+    integer, intent(in) :: ih
+    real(wp), intent(in) :: qa
+  end subroutine
+
+end program