Added stepper to EvoAtmosphere

photochem/cython/EvoAtmosphere.pyx

@@ -157,6 +157,46 @@ cdef class EvoAtmosphere:
       raise PhotoException(err.decode("utf-8").strip())
     return success
 
+  def initialize_stepper(self, ndarray[double, ndim=2] usol_start):
+    """Initializes an integration starting at `usol_start`.
+
+    Parameters
+    ----------
+    usol_start : ndarray[double,ndim=2]
+        Initial mixing ratios
+    """   
+    cdef char err[ERR_LEN+1]
+    cdef int nq = self.dat.nq
+    cdef int nz = self.var.nz
+    cdef ndarray usol_start_ = np.asfortranarray(usol_start)
+    if usol_start_.shape[0] != nq or usol_start_.shape[1] != nz:
+      raise PhotoException("Input usol_start is the wrong size.")
+    ea_pxd.evoatmosphere_initialize_stepper_wrapper(&self._ptr, &nq, &nz,  <double *>usol_start_.data, err)
+    if len(err.strip()) > 0:
+      raise PhotoException(err.decode("utf-8").strip())
+
+  def step(self):
+    """Takes one internal integration step. Function `initialize_stepper`.
+    must have been called before this
+
+    Returns
+    -------
+    float
+        Current time in the integration.
+    """
+    cdef char err[ERR_LEN+1]
+    cdef double tn = ea_pxd.evoatmosphere_step_wrapper(&self._ptr, err)
+    if len(err.strip()) > 0:
+      raise PhotoException(err.decode("utf-8").strip())
+    return tn
+
+  def destroy_stepper(self):
+    "Deallocates memory created during `initialize_stepper`"
+    cdef char err[ERR_LEN+1]
+    ea_pxd.evoatmosphere_destroy_stepper_wrapper(&self._ptr, err)
+    if len(err.strip()) > 0:
+      raise PhotoException(err.decode("utf-8").strip())
+
   def production_and_loss(self, str species, ndarray[double, ndim=2] usol, double top_atmos):
     """Computes the production and loss of input `species`.
     See ProductionLoss object in photochem_types.f90.

photochem/cython/EvoAtmosphere_pxd.pxd

@@ -23,6 +23,12 @@ cdef extern void evoatmosphere_evolve_wrapper(void *ptr, char *filename,
                 double *tstart, int *nq, int *nz, double *usol, 
                 int *nt, double *t_eval, bool *overwrite, bool *restart_from_file, bool *success, char *err)
 
+cdef extern void evoatmosphere_initialize_stepper_wrapper(void *ptr, int *nq, int *nz, double *usol_start, char *err)
+
+cdef extern double evoatmosphere_step_wrapper(void *ptr, char *err)
+
+cdef extern void evoatmosphere_destroy_stepper_wrapper(void *ptr, char *err)
+
 cdef extern void evoatmosphere_production_and_loss_wrapper(void *ptr, char *species, int *nq, 
                                                         int *nz, double *usol, double *top_atmos, void *pl_ptr, char *err)
 

photochem/fortran/EvoAtmosphere_wrapper.f90

@@ -154,6 +154,56 @@ subroutine evoatmosphere_evolve_wrapper(ptr, filename, tstart, nq, nz, usol, nt,
     endif
   end subroutine
 
+  subroutine evoatmosphere_initialize_stepper_wrapper(ptr, nq, nz, usol_start, err) bind(c)
+    type(c_ptr), intent(in) :: ptr
+    integer(c_int), intent(in) :: nq, nz
+    real(c_double), intent(in) :: usol_start(nq, nz)
+    character(len=c_char), intent(out) :: err(err_len+1)
+
+    character(:), allocatable :: err_f
+
+    type(EvoAtmosphere), pointer :: pc
+    call c_f_pointer(ptr, pc)
+
+    call pc%initialize_stepper(usol_start, err_f)
+    err(1) = c_null_char
+    if (allocated(err_f)) then
+      call copy_string_ftoc(err_f, err)
+    endif
+  end subroutine
+
+  function evoatmosphere_step_wrapper(ptr, err) result(tn) bind(c)
+    type(c_ptr), intent(in) :: ptr
+    character(len=c_char), intent(out) :: err(err_len+1)
+    real(c_double) :: tn
+
+    character(:), allocatable :: err_f
+    type(EvoAtmosphere), pointer :: pc
+    call c_f_pointer(ptr, pc)
+
+    tn = pc%step(err_f)
+    err(1) = c_null_char
+    if (allocated(err_f)) then
+      call copy_string_ftoc(err_f, err)
+    endif
+  end function
+
+  subroutine evoatmosphere_destroy_stepper_wrapper(ptr, err) bind(c)
+    type(c_ptr), intent(in) :: ptr
+    character(len=c_char), intent(out) :: err(err_len+1)
+
+    character(:), allocatable :: err_f
+
+    type(EvoAtmosphere), pointer :: pc
+    call c_f_pointer(ptr, pc)
+
+    call pc%destroy_stepper(err_f)
+    err(1) = c_null_char
+    if (allocated(err_f)) then
+      call copy_string_ftoc(err_f, err)
+    endif
+  end subroutine
+
   subroutine evoatmosphere_production_and_loss_wrapper(ptr, species, nq, nz, usol, top_atmos, pl_ptr, err) bind(c)
     use photochem, only: ProductionLoss
     type(c_ptr), intent(in) :: ptr

src/evoatmosphere/photochem_evoatmosphere.f90

@@ -48,6 +48,9 @@ function temp_dependent_albedo_fcn(T_surf) result(albedo)
 
     !!! photochem_evoatmosphere_integrate.f90 !!!
     procedure :: evolve
+    procedure :: initialize_stepper
+    procedure :: step
+    procedure :: destroy_stepper
 
     !!! photochem_evoatmosphere_utils.f90 !!!
     procedure :: out2atmosphere_txt
@@ -153,6 +156,27 @@ module function evolve(self, filename, tstart, usol_start, t_eval, overwrite, re
       character(:), allocatable, intent(out) :: err
       ! Evolve atmosphere through time, and saves output in a binary Fortran file.
     end function
+
+    !> Initializes an integration starting at `usol_start`
+    module subroutine initialize_stepper(self, usol_start, err)      
+      class(EvoAtmosphere), target, intent(inout) :: self
+      real(dp), intent(in) :: usol_start(:,:) !! Initial mixing ratios
+      character(:), allocatable, intent(out) :: err
+    end subroutine
+
+    !> Takes one internal integration step. Function `initialize_stepper`
+    !> must have been called befe this
+    module function step(self, err) result(tn)
+      class(EvoAtmosphere), target, intent(inout) :: self
+      character(:), allocatable, intent(out) :: err
+      real(dp) :: tn
+    end function
+
+    !> Deallocates memory created during `initialize_stepper`
+    module subroutine destroy_stepper(self, err)
+      class(EvoAtmosphere), target, intent(inout) :: self
+      character(:), allocatable, intent(out) :: err
+    end subroutine
 
     module function RhsFn_evo(tn, sunvec_y, sunvec_f, user_data) &
                           result(ierr) bind(c, name='RhsFn_evo')

src/evoatmosphere/photochem_evoatmosphere_integrate.f90

@@ -682,5 +682,219 @@ subroutine read_end_of_evo_file(self, filename, t_eval, restart_index, tcur, top
     endif
 
   end subroutine
+
+  module subroutine initialize_stepper(self, usol_start, err)
+    use, intrinsic :: iso_c_binding
+    use fcvode_mod, only: CV_BDF, CV_NORMAL, CV_ONE_STEP, FCVodeInit, FCVodeSVtolerances, &
+                          FCVodeSetLinearSolver, FCVode, FCVodeCreate, FCVodeFree, &
+                          FCVodeSetMaxNumSteps, FCVodeSetJacFn, FCVodeSetInitStep, &
+                          FCVodeGetCurrentStep, FCVodeSetMaxErrTestFails, FCVodeSetMaxOrd, &
+                          FCVodeSetUserData, FCVodeSetErrHandlerFn
+    use fsundials_nvector_mod, only: N_Vector, FN_VDestroy
+    use fnvector_serial_mod, only: FN_VMake_Serial   
+    use fsunmatrix_band_mod, only: FSUNBandMatrix
+    use fsundials_matrix_mod, only: SUNMatrix, FSUNMatDestroy
+    use fsundials_linearsolver_mod, only: SUNLinearSolver, FSUNLinSolFree
+    use fsunlinsol_band_mod, only: FSUNLinSol_Band
+
+    use photochem_enum, only: DensityBC
+
+    class(EvoAtmosphere), target, intent(inout) :: self
+    real(dp), intent(in) :: usol_start(:,:)
+    character(:), allocatable, intent(out) :: err
+
+    real(c_double), pointer :: yvec_usol(:,:)
+    real(c_double) :: tstart
+    integer(c_int) :: ierr       ! error flag from C functions
+    integer(c_int64_t) :: neqs_long
+    integer(c_int64_t) :: mu, ml
+    integer(c_long) :: mxsteps_
+    type(c_ptr)    :: user_data
+
+    type(PhotochemData), pointer :: dat
+    type(PhotochemVars), pointer :: var
+    type(PhotochemWrkEvo), pointer :: wrk
+    type(EvoAtmosphere), pointer :: self_ptr
+
+    integer :: i, j, k
+
+    dat => self%dat
+    var => self%var
+    wrk => self%wrk
+
+    if (size(usol_start,1) /= dat%nq .or. size(usol_start,2) /= var%nz) then
+      err = "Input 'usol_start' to 'initialize_stepper' is the wrong dimension"
+      return
+    endif
+
+    if (self%evolve_climate) then
+      err = 'You can not integrate with this stepper when climate evolution is turned on. '// &
+            'Use the `evolve` routine instead for integrations with changing climate.'
+      return
+    endif
+
+    ! settings
+    mxsteps_ = var%mxsteps
+    neqs_long = var%neqs
+    tstart = 0
+    mu = dat%nq
+    ml = dat%nq
+    self_ptr => self
+    user_data = c_loc(self_ptr)
+
+    call self%destroy_stepper(err)
+    if (allocated(err)) return
+
+    ! initialize solution vector
+    allocate(wrk%sun%yvec(var%neqs))
+    yvec_usol(1:dat%nq,1:var%nz) => wrk%sun%yvec
+    do j=1,var%nz
+      do i=1,dat%nq
+        k = i + (j-1)*dat%nq
+        wrk%sun%yvec(k) = usol_start(i,j)
+      enddo
+    enddo
+    do i = 1,dat%nq
+      if (var%lowerboundcond(i) == DensityBC) then
+        wrk%sun%yvec(i) = var%lower_fix_den(i)
+      endif
+    enddo
+    ! set abstol
+    allocate(wrk%sun%abstol(var%neqs))
+    call self%set_trop_ind(yvec_usol, err)
+    if (allocated(err)) return
+    do j=1,var%nz
+      do i=1,dat%nq
+        k = i + (j-1)*dat%nq
+        wrk%sun%abstol(k) = self%wrk%density_hydro(j)*var%atol
+      enddo
+    enddo
+
+    wrk%sun%abstol_nvec => FN_VMake_Serial(neqs_long, wrk%sun%abstol)
+    if (.not. associated(wrk%sun%abstol_nvec)) then
+      err = "CVODE setup error."
+      return
+    end if
+
+    ! create SUNDIALS N_Vector
+    wrk%sun%sunvec_y => FN_VMake_Serial(neqs_long, wrk%sun%yvec)
+    if (.not. associated(wrk%sun%sunvec_y)) then
+      err = "CVODE setup error."
+      return
+    end if
+
+    ! create CVode memory
+    wrk%sun%cvode_mem = FCVodeCreate(CV_BDF)
+    if (.not. c_associated(wrk%sun%cvode_mem)) then
+      err = "CVODE setup error."
+      return
+    end if
+
+    ! set user data
+    ierr = FCVodeSetUserData(wrk%sun%cvode_mem, user_data)
+    if (ierr /= 0) then
+      err = "CVODE setup error."
+      return
+    end if
+
+    ierr = FCVodeInit(wrk%sun%cvode_mem, c_funloc(RhsFn_evo), tstart, wrk%sun%sunvec_y)
+    if (ierr /= 0) then
+      err = "CVODE setup error."
+      return
+    end if
+
+    ierr = FCVodeSVtolerances(wrk%sun%cvode_mem, var%rtol, wrk%sun%abstol_nvec)
+    if (ierr /= 0) then
+      err = "CVODE setup error."
+      return
+    end if
+
+    wrk%sun%sunmat => FSUNBandMatrix(neqs_long, mu, ml)
+    wrk%sun%sunlin => FSUNLinSol_Band(wrk%sun%sunvec_y, wrk%sun%sunmat)
+
+    ierr = FCVodeSetLinearSolver(wrk%sun%cvode_mem, wrk%sun%sunlin, wrk%sun%sunmat)
+    if (ierr /= 0) then
+      err = "CVODE setup error."
+      return
+    end if
+
+    ierr = FCVodeSetJacFn(wrk%sun%cvode_mem, c_funloc(JacFn_evo))
+    if (ierr /= 0) then
+      err = "CVODE setup error."
+      return
+    end if
+
+    ierr = FCVodeSetMaxNumSteps(wrk%sun%cvode_mem, mxsteps_)
+    if (ierr /= 0) then
+      err = "CVODE setup error."
+      return
+    end if
+
+    ierr = FCVodeSetInitStep(wrk%sun%cvode_mem, var%initial_dt)
+    if (ierr /= 0) then
+      err = "CVODE setup error."
+      return
+    end if
+
+    ierr = FCVodeSetMaxErrTestFails(wrk%sun%cvode_mem, var%max_err_test_failures)
+    if (ierr /= 0) then
+      err = "CVODE setup error."
+      return
+    end if
+
+    ierr = FCVodeSetMaxOrd(wrk%sun%cvode_mem, var%max_order)
+    if (ierr /= 0) then
+      err = "CVODE setup error."
+      return
+    end if
+
+    if (var%verbose == 0) then
+      ierr = FCVodeSetErrHandlerFn(wrk%sun%cvode_mem, c_funloc(ErrHandlerFn_evo), c_null_ptr)
+      if (ierr /= 0) then
+        err = "CVODE setup error."
+        return
+      end if
+    endif
+
+  end subroutine
+
+  module function step(self, err) result(tn)
+    use iso_c_binding, only: c_null_ptr, c_int, c_double, c_associated
+    use fcvode_mod, only: CV_ONE_STEP, FCVode
+    class(EvoAtmosphere), target, intent(inout) :: self
+    character(:), allocatable, intent(out) :: err
+    real(dp) :: tn
+
+    integer(c_int) :: ierr
+    real(c_double), parameter :: dum = 0.0_dp
+    real(c_double) :: tcur(1)
+
+    if (.not.c_associated(self%wrk%sun%cvode_mem)) then
+      err = "You must first initialize the stepper with 'initialize_stepper'"
+      return 
+    endif
+
+    ierr = FCVode(self%wrk%sun%cvode_mem, dum, self%wrk%sun%sunvec_y, tcur, CV_ONE_STEP)
+    if (ierr /= 0) then
+      err = "CVODE step failed"
+      return
+    endif
+    tn = tcur(1)
+  end function
+
+  module subroutine destroy_stepper(self, err)
+    use iso_c_binding, only: c_int, c_associated, c_null_ptr
+    use fcvode_mod, only: FCVodeFree
+    use fsundials_nvector_mod, only: FN_VDestroy
+    use fsundials_matrix_mod, only: FSUNMatDestroy
+    use fsundials_linearsolver_mod, only: FSUNLinSolFree
+
+    class(EvoAtmosphere), target, intent(inout) :: self
+    character(:), allocatable, intent(out) :: err
+
+    call self%wrk%sun%finalize(err)
+    if (allocated(err)) return
+
+  end subroutine
 
 end submodule