Add option for source terms in flux_update_densities

examples/advection_v2.f90

@@ -221,7 +221,8 @@ contains
 
     call flux_generic_tree(tree, 1, [i_phi+s_deriv], [i_flux], wmax, &
          max_wavespeed, get_flux)
-    call flux_update_densities(tree, dt, 1, [i_phi], [i_flux], s_prev, s_out)
+    call flux_update_densities(tree, dt, 1, [i_phi], [i_flux], &
+         s_deriv, s_prev, s_out)
 
     ! Compute maximal time step
     dt_lim = 1.0_dp / sum(wmax/af_lvl_dr(tree, tree%highest_lvl))

examples/euler_gas_dynamics.f90

@@ -177,7 +177,7 @@ contains
     call flux_generic_tree(tree, n_vars, variables+s_deriv, fluxes, wmax, &
          max_wavespeed, get_fluxes, to_primitive, to_conservative)
     call flux_update_densities(tree, dt, n_vars, variables, fluxes, &
-         s_prev, s_out)
+         s_deriv, s_prev, s_out)
 
     ! Compute new time step
     dt_lim = 1.0_dp / sum(wmax/af_lvl_dr(tree, tree%highest_lvl))

src/m_af_flux_schemes.f90

@@ -37,6 +37,16 @@ module m_af_flux_schemes
        type(box_t), intent(in) :: box             !< Current box
        integer, intent(in)     :: line_ix(NDIM-1) !< Index of line for dim /= flux_dim
      end subroutine subr_flux_from_prim
+
+     subroutine subr_source(box, dt, n_vars, i_cc, s_deriv, s_out)
+       import
+       type(box_t), intent(inout) :: box
+       real(dp), intent(in)       :: dt
+       integer, intent(in)        :: n_vars
+       integer, intent(in)        :: i_cc(n_vars)
+       integer, intent(in)        :: s_deriv
+       integer, intent(in)        :: s_out
+     end subroutine subr_source
   end interface
 
   public :: flux_diff_1d, flux_diff_2d, flux_diff_3d
@@ -220,14 +230,17 @@ contains
   end subroutine flux_kurganovTadmor_1d
 
   subroutine flux_update_densities(tree, dt, n_vars, i_cc, i_flux, &
-       s_prev, s_out)
+       s_deriv, s_prev, s_out, add_source_box)
     type(af_t), intent(inout) :: tree
     real(dp), intent(in)      :: dt             !< Time step
     integer, intent(in)       :: n_vars         !< Number of variables
     integer, intent(in)       :: i_cc(n_vars)   !< Cell-centered variables
     integer, intent(in)       :: i_flux(n_vars) !< Flux variables
+    integer, intent(in)       :: s_deriv        !< State to compute derivatives from
     integer, intent(in)       :: s_prev         !< Previous time state
     integer, intent(in)       :: s_out          !< Output time state
+    !> Method to include source terms
+    procedure(subr_source), optional :: add_source_box
     integer                   :: lvl, n, id, IJK, nc
     real(dp)                  :: dt_dr(NDIM)
     real(dp)                  :: rfac(2, tree%n_cell)
@@ -242,12 +255,20 @@ contains
           id    = tree%lvls(lvl)%leaves(n)
           dt_dr = dt/tree%boxes(id)%dr
 
+          tree%boxes(id)%cc(DTIMES(1:nc), i_cc+s_out) = &
+               tree%boxes(id)%cc(DTIMES(1:nc), i_cc+s_prev)
+
+          if (present(add_source_box)) then
+             call add_source_box(tree%boxes(id), dt, &
+                  n_vars, i_cc, s_deriv, s_out)
+          end if
+
           associate(cc => tree%boxes(id)%cc, fc => tree%boxes(id)%fc)
 #if NDIM == 2
             if (tree%coord_t == af_cyl) then
                call af_cyl_flux_factors(tree%boxes(id), rfac)
                do KJI_DO(1, nc)
-                  cc(i, j, i_cc+s_out) = cc(i, j, i_cc+s_prev) + &
+                  cc(i, j, i_cc+s_out) = cc(i, j, i_cc+s_out) + &
                        dt_dr(1) * (&
                        rfac(1, i) * fc(i, j, 1, i_flux) - &
                        rfac(2, i) * fc(i+1, j, 1, i_flux)) &
@@ -256,7 +277,7 @@ contains
                end do; CLOSE_DO
             else
                do KJI_DO(1, nc)
-                  cc(i, j, i_cc+s_out) = cc(i, j, i_cc+s_prev) + &
+                  cc(i, j, i_cc+s_out) = cc(i, j, i_cc+s_out) + &
                        dt_dr(1) * &
                        (fc(i, j, 1, i_flux) - fc(i+1, j, 1, i_flux)) &
                        + dt_dr(2) * &
@@ -265,7 +286,7 @@ contains
             end if
 #elif NDIM == 3
             do KJI_DO(1, nc)
-               cc(i, j, k, i_cc+s_out) = cc(i, j, k, i_cc+s_prev) + &
+               cc(i, j, k, i_cc+s_out) = cc(i, j, k, i_cc+s_out) + &
                     dt_dr(1) * &
                     (fc(i, j, k, 1, i_flux) - fc(i+1, j, k, 1, i_flux)) + &
                     dt_dr(2) * &