diff --git a/photochem/cython/EvoAtmosphere.pyx b/photochem/cython/EvoAtmosphere.pyx index da098a8..3899459 100644 --- a/photochem/cython/EvoAtmosphere.pyx +++ b/photochem/cython/EvoAtmosphere.pyx @@ -227,6 +227,42 @@ cdef class EvoAtmosphere: &veff_c, &flux_c, &missing, err); if len(err.strip()) > 0: raise PhotoException(err.decode("utf-8").strip()) + + def set_rate_fcn(self, str species, object fcn): + """Sets a function describing a custom rate for a species. + This could be useful for modeling external processes not in the + model. + + Parameters + ---------- + species : str + Species name + fcn : function + A Numba cfunc that describes the time-dependent rate + """ + cdef bytes species_b = pystring2cstring(species) + cdef char *species_c = species_b + cdef char err[ERR_LEN+1] + cdef uintptr_t fcn_l + cdef ea_pxd.time_dependent_rate_fcn fcn_c + + if fcn is None: + fcn_l = 0 + fcn_c = NULL + else: + argtypes = (ct.c_double, ct.c_int32, ct.POINTER(ct.c_double)) + restype = None + if not fcn.ctypes.argtypes == argtypes: + raise PhotoException("The callback function has the wrong argument types.") + if not fcn.ctypes.restype == restype: + raise PhotoException("The callback function has the wrong return type.") + + fcn_l = fcn.address + fcn_c = fcn_l + + ea_pxd.evoatmosphere_set_rate_fcn_wrapper(&self._ptr, species_c, fcn_c, err) + if len(err.strip()) > 0: + raise PhotoException(err.decode("utf-8").strip()) def regrid_prep_atmosphere(self, ndarray[double, ndim=2] usol, double top_atmos): """This subroutine calculates re-grids the model so that the top of the model domain diff --git a/photochem/cython/EvoAtmosphere_pxd.pxd b/photochem/cython/EvoAtmosphere_pxd.pxd index b9e377e..6fb1313 100644 --- a/photochem/cython/EvoAtmosphere_pxd.pxd +++ b/photochem/cython/EvoAtmosphere_pxd.pxd @@ -1,3 +1,4 @@ +from Atmosphere_pxd cimport time_dependent_rate_fcn from libcpp cimport bool cdef extern from "": pass @@ -21,6 +22,7 @@ cdef extern void evoatmosphere_set_lower_bc_wrapper(void *ptr, char *species, ch double *vdep, double *den, double *press, double *flux, double *height, bool *missing, char *err) cdef extern void evoatmosphere_set_upper_bc_wrapper(void *ptr, char *species, char *bc_type, double *veff, double *flux, bool *missing, char *err) +cdef extern void evoatmosphere_set_rate_fcn_wrapper(void *ptr, char *species_c, time_dependent_rate_fcn fcn, char *err) cdef extern void evoatmosphere_regrid_prep_atmosphere_wrapper(void *ptr, int *nq, int *nz, double *usol, double *top_atmos, char *err) cdef extern void evoatmosphere_evolve_wrapper(void *ptr, char *filename, diff --git a/photochem/fortran/EvoAtmosphere_wrapper.f90 b/photochem/fortran/EvoAtmosphere_wrapper.f90 index fcdc32a..ecdb004 100644 --- a/photochem/fortran/EvoAtmosphere_wrapper.f90 +++ b/photochem/fortran/EvoAtmosphere_wrapper.f90 @@ -187,6 +187,38 @@ subroutine evoatmosphere_set_upper_bc_wrapper(ptr, species, bc_type, veff, flux, endif end subroutine + subroutine evoatmosphere_set_rate_fcn_wrapper(ptr, species_c, fcn_c, err) bind(c) + use photochem_types, only: time_dependent_rate_fcn + type(c_ptr), intent(in) :: ptr + character(kind=c_char), intent(in) :: species_c(*) + type(c_funptr), value, intent(in) :: fcn_c + character(kind=c_char), intent(out) :: err(err_len+1) + + type(EvoAtmosphere), pointer :: pc + procedure(time_dependent_rate_fcn), pointer :: fcn_f + character(:), allocatable :: species_f + character(:), allocatable :: err_f + + ! Cast the void pointer to wrapped object. + call c_f_pointer(ptr, pc) + + ! Copy c string to f string. + allocate(character(len=len_cstring(species_c))::species_f) + call copy_string_ctof(species_c, species_f) + + ! Convert c function pointer to a fortran function pointer (unsafe). + call c_f_procpointer(fcn_c, fcn_f) + + ! Call the function. + call pc%set_rate_fcn(species_f, fcn_f, err_f) + + ! Set error, if there is one. + err(1) = c_null_char + if (allocated(err_f)) then + call copy_string_ftoc(err_f, err) + endif + end subroutine + subroutine evoatmosphere_regrid_prep_atmosphere_wrapper(ptr, nq, nz, usol, top_atmos, err) bind(c) type(c_ptr), intent(in) :: ptr integer(c_int), intent(in) :: nq, nz diff --git a/src/evoatmosphere/photochem_evoatmosphere.f90 b/src/evoatmosphere/photochem_evoatmosphere.f90 index 01fc47e..1b5d21a 100644 --- a/src/evoatmosphere/photochem_evoatmosphere.f90 +++ b/src/evoatmosphere/photochem_evoatmosphere.f90 @@ -59,6 +59,7 @@ function temp_dependent_albedo_fcn(T_surf) result(albedo) procedure :: gas_fluxes procedure :: set_lower_bc procedure :: set_upper_bc + procedure :: set_rate_fcn procedure :: rebin_update_vertical_grid procedure :: regrid_prep_atmosphere @@ -116,9 +117,10 @@ module subroutine right_hand_side_chem(self, usol, rhs, err) character(:), allocatable, intent(out) :: err end subroutine - module subroutine rhs_evo_gas(self, neqs, usol_flat, rhs, err) + module subroutine rhs_evo_gas(self, neqs, tn, usol_flat, rhs, err) class(EvoAtmosphere), target, intent(inout) :: self integer, intent(in) :: neqs + real(dp), intent(in) :: tn real(dp), target, intent(in) :: usol_flat(neqs) real(dp), intent(out) :: rhs(neqs) character(:), allocatable, intent(out) :: err @@ -257,6 +259,14 @@ module subroutine set_upper_bc(self, species, bc_type, veff, flux, err) character(:), allocatable, intent(out) :: err end subroutine + module subroutine set_rate_fcn(self, species, fcn, err) + use photochem_types, only: time_dependent_rate_fcn + class(EvoAtmosphere), target, intent(inout) :: self + character(*), intent(in) :: species + procedure(time_dependent_rate_fcn), pointer :: fcn + character(:), allocatable, intent(inout) :: err + end subroutine + module subroutine rebin_update_vertical_grid(self, usol_old, top_atmos, usol_new, err) class(EvoAtmosphere), target, intent(inout) :: self real(dp), intent(in) :: usol_old(:,:) diff --git a/src/evoatmosphere/photochem_evoatmosphere_integrate.f90 b/src/evoatmosphere/photochem_evoatmosphere_integrate.f90 index f9cafdf..19f738a 100644 --- a/src/evoatmosphere/photochem_evoatmosphere_integrate.f90 +++ b/src/evoatmosphere/photochem_evoatmosphere_integrate.f90 @@ -39,7 +39,7 @@ module function RhsFn_evo(tn, sunvec_y, sunvec_f, user_data) & fvec(1:self%var%neqs) => FN_VGetArrayPointer(sunvec_f) ! fill RHS vector - call self%right_hand_side(self%var%neqs, yvec, fvec, err) + call self%right_hand_side(self%var%neqs, tn, yvec, fvec, err) loc_ierr = FCVodeGetNumSteps(self%wrk%sun%cvode_mem, nsteps) if (nsteps(1) /= self%wrk%nsteps_previous .and. self%var%verbose > 0) then diff --git a/src/evoatmosphere/photochem_evoatmosphere_rhs.f90 b/src/evoatmosphere/photochem_evoatmosphere_rhs.f90 index 1ea3bef..0299ad0 100644 --- a/src/evoatmosphere/photochem_evoatmosphere_rhs.f90 +++ b/src/evoatmosphere/photochem_evoatmosphere_rhs.f90 @@ -595,6 +595,11 @@ module subroutine prep_all_evo_gas(self, usol_in, err) !!! reaction rates call reaction_rates(self%dat, self%var, wrk%density, wrk%densities, wrk%rx_rates) + ! Update the photon_flux if the function is associated. + ! we use time wrk%tn, which MUST be updated. + if (associated(var%photon_flux_fcn)) then + call var%photon_flux_fcn(wrk%tn, dat%nw, var%photon_flux) + endif call photorates(dat, var, wrk%densities, & wrk%prates, wrk%surf_radiance, wrk%amean_grd, wrk%optical_depth, err) if (allocated(err)) return @@ -612,7 +617,7 @@ module subroutine prep_all_evo_gas(self, usol_in, err) end subroutine - module subroutine rhs_evo_gas(self, neqs, usol_flat, rhs, err) + module subroutine rhs_evo_gas(self, neqs, tn, usol_flat, rhs, err) use photochem_enum, only: MosesBC, VelocityBC, DensityBC, PressureBC, FluxBC, VelocityDistributedFluxBC use photochem_enum, only: ZahnleHydrogenEscape use iso_c_binding, only: c_ptr, c_f_pointer @@ -620,6 +625,7 @@ module subroutine rhs_evo_gas(self, neqs, usol_flat, rhs, err) class(EvoAtmosphere), target, intent(inout) :: self integer, intent(in) :: neqs + real(dp), intent(in) :: tn real(dp), target, intent(in) :: usol_flat(neqs) real(dp), intent(out) :: rhs(neqs) character(:), allocatable, intent(out) :: err @@ -643,6 +649,9 @@ module subroutine rhs_evo_gas(self, neqs, usol_flat, rhs, err) 'related to some mixing ratios getting too negative.' return endif + + ! time + wrk%tn = tn ! fills self%wrk with data call prep_all_evo_gas(self, usol_in, err) @@ -653,6 +662,17 @@ module subroutine rhs_evo_gas(self, neqs, usol_flat, rhs, err) wrk%H2O_sat_mix, wrk%H2O_rh, wrk%rainout_rates, & wrk%density, wrk%mix, wrk%densities, wrk%xp, wrk%xl, rhs) + ! Extra functions specifying production or destruction + do i = 1,dat%nq + if (associated(var%rate_fcns(i)%fcn)) then + call var%rate_fcns(i)%fcn(tn, var%nz, wrk%xp) ! using wrk%xp space. + do j = 1,var%nz + k = i + (j-1)*dat%nq + rhs(k) = rhs(k) + wrk%xp(j) ! (molecules/cm^3/s) + enddo + endif + enddo + ! diffusion (interior grid points) do j = 2,var%nz-1 do i = 1,dat%nq diff --git a/src/evoatmosphere/photochem_evoatmosphere_utils.f90 b/src/evoatmosphere/photochem_evoatmosphere_utils.f90 index 559347a..ad497f0 100644 --- a/src/evoatmosphere/photochem_evoatmosphere_utils.f90 +++ b/src/evoatmosphere/photochem_evoatmosphere_utils.f90 @@ -236,6 +236,26 @@ module subroutine set_upper_bc(self, species, bc_type, veff, flux, err) end subroutine + module subroutine set_rate_fcn(self, species, fcn, err) + use photochem_types, only: time_dependent_rate_fcn + class(EvoAtmosphere), target, intent(inout) :: self + character(*), intent(in) :: species + procedure(time_dependent_rate_fcn), pointer :: fcn + character(:), allocatable, intent(inout) :: err + + integer :: ind + + ind = findloc(self%dat%species_names(1:self%dat%nq), species, 1) + if (ind == 0) then + err = 'Species "'//species//'" is not in the list of species, '// & + 'or is a background or short-lived species.' + return + endif + + self%var%rate_fcns(ind)%fcn => fcn + + end subroutine + module subroutine rebin_update_vertical_grid(self, usol_old, top_atmos, usol_new, err) class(EvoAtmosphere), target, intent(inout) :: self real(dp), intent(in) :: usol_old(:,:)