python wrapper

clima/cython/AdiabatClimate.pyx

@@ -394,7 +394,7 @@ cdef class AdiabatClimate:
     return T_surf
 
   def RCE(self, ndarray[double, ndim=1] P_i_surf, double T_surf_guess, ndarray[double, ndim=1] T_guess,
-          convecting_with_below = None):
+          convecting_with_below = None, custom_dry_mix = None):
     """Compute full radiative-convective equilibrium.
 
     Parameters
@@ -416,6 +416,7 @@ cdef class AdiabatClimate:
     """
     cdef int ng = P_i_surf.shape[0]
     cdef int dim_T_guess = T_guess.shape[0]
+
     cdef ndarray[cbool, ndim=1] convecting_with_below_ = np.array([False],dtype=np.bool_)
     cdef cbool convecting_with_below_present
     if convecting_with_below is None:
@@ -424,12 +425,49 @@ cdef class AdiabatClimate:
       convecting_with_below_present = True
       convecting_with_below_ = convecting_with_below
     cdef int dim_convecting_with_below = convecting_with_below_.shape[0]
+
+    # Custom mixing ratio
+    cdef cbool custom_present = False
+    # initialize arrays
+    cdef int dim_sp_custom = 1
+    cdef ndarray sp_custom_c = np.zeros(dim_sp_custom*S_STR_LEN + 1, 'S1')
+    cdef int dim_P_custom = 1
+    cdef ndarray[double, ndim=1] P_custom = np.empty(dim_P_custom, np.double)
+    cdef int dim1_mix_custom = 1
+    cdef int dim2_mix_custom = 1
+    cdef ndarray[double, ndim=2, mode='fortran'] mix_custom = np.empty((dim1_mix_custom,dim2_mix_custom), np.double)
+    cdef list species
+
+    if custom_dry_mix is not None:
+      custom_present = True
+
+      species = list(custom_dry_mix.keys())
+      if 'pressure' not in species:
+        raise ClimaException('`pressure` must be a key in `custom_dry_mix`')
+      species.remove('pressure')
+      dim_sp_custom = len(species)
+      sp_custom_c = list2cstring(species, S_STR_LEN)
+
+      P_custom = custom_dry_mix['pressure']
+      dim_P_custom = len(P_custom)
+
+      dim1_mix_custom = dim_P_custom
+      dim2_mix_custom = dim_sp_custom
+      mix_custom = np.empty((dim_P_custom,dim_sp_custom), np.double, order='F')
+      for i,sp in enumerate(species):
+        if sp != 'pressure':
+          mix_custom[:,i] = custom_dry_mix[sp]
+
     cdef cbool converged
     cdef char err[ERR_LEN+1]
 
-    wa_pxd.adiabatclimate_rce_wrapper(self._ptr, &ng, <double *>P_i_surf.data, &T_surf_guess,
-                               &dim_T_guess, <double *>T_guess.data, &convecting_with_below_present,
-                               &dim_convecting_with_below, <cbool *>convecting_with_below_.data, &converged, err)
+    wa_pxd.adiabatclimate_rce_wrapper(
+      self._ptr, &ng, <double *>P_i_surf.data, &T_surf_guess, &dim_T_guess, <double *>T_guess.data, 
+      &convecting_with_below_present, &dim_convecting_with_below, <cbool *>convecting_with_below_.data,
+      &custom_present, &dim_sp_custom, <char*> sp_custom_c.data, &dim_P_custom, <double *> P_custom.data, 
+      &dim1_mix_custom, &dim2_mix_custom, <double *> mix_custom.data, 
+      &converged, err
+    )
     if len(err.strip()) > 0:
       raise ClimaException(err.decode("utf-8").strip())
     return converged

clima/cython/AdiabatClimate_pxd.pxd

@@ -65,9 +65,13 @@ cdef extern void adiabatclimate_surface_temperature_bg_gas_wrapper(AdiabatClimat
                                       double *P_i_surf, double *P_surf, char *bg_gas, 
                                       double *T_guess, double *T_surf, char *err)
 
-cdef extern void adiabatclimate_rce_wrapper(AdiabatClimate *ptr, int *ng, double *P_i_surf, double *T_surf_guess,
-                                      int *dim_T_guess, double *T_guess, bool *convecting_with_below_present, 
-                                      int *dim_convecting_with_below, bool *convecting_with_below, bool *converged, char *err)
+cdef extern void adiabatclimate_rce_wrapper(
+  AdiabatClimate *ptr, int *ng, double *P_i_surf, double *T_surf_guess, int *dim_T_guess, double *T_guess, 
+  bool *convecting_with_below_present, int *dim_convecting_with_below, bool *convecting_with_below,
+  bool *custom_present, int *dim_sp_custom, char* sp_custom, int *dim_P_custom, double *P_custom, 
+  int *dim1_mix_custom, int *dim2_mix_custom, double *mix_custom,
+  bool *converged, char *err
+)
 
 cdef extern void adiabatclimate_set_ocean_solubility_fcn_wrapper(AdiabatClimate *ptr, char *species_c, ocean_solubility_fcn fcn, char *err)
 

clima/cython/_clima.pyx

@@ -36,5 +36,14 @@ cdef c2stringarr(ndarray c_str_arr, int str_len, int arr_len):
   bs = c_str_arr[:-1].tobytes()
   return [bs[i:i+str_len].decode().strip() for i in range(0, str_len*arr_len, str_len)]
 
+cdef list2cstring(list arr, int str_len):
+  arr_c = np.zeros(len(arr)*str_len + 1, 'S1')
+  for i in range(len(arr)):
+    if len(arr[i]) > str_len:
+          raise Exception('Failed to convert Python list to a C string')
+    arr_c[i*str_len:(i+1)*str_len] = b' '
+    arr_c[i*str_len:i*str_len+len(arr[i])] = np.array([elem.encode('utf-8') for elem in arr[i]])
+  return arr_c
+
 class ClimaException(Exception):
     pass

clima/fortran/AdiabatClimate.f90

@@ -342,9 +342,11 @@ subroutine adiabatclimate_surface_temperature_bg_gas_wrapper(ptr, ng, P_i_surf,
 end subroutine
 
 subroutine adiabatclimate_rce_wrapper(ptr, ng, P_i_surf, T_surf_guess, dim_T_guess, T_guess, &
-                                      convecting_with_below_present, dim_convecting_with_below, &
-                                      convecting_with_below, converged, err) bind(c)
-  use clima, only: AdiabatClimate
+                                      convecting_with_below_present, dim_convecting_with_below, convecting_with_below, &
+                                      custom_present, dim_sp_custom, sp_custom, dim_P_custom, P_custom, &
+                                      dim1_mix_custom, dim2_mix_custom, mix_custom, &
+                                      converged, err) bind(c)
+  use clima, only: AdiabatClimate, s_str_len
   type(c_ptr), value, intent(in) :: ptr
   integer(c_int), intent(in) :: ng
   real(c_double), intent(in) :: P_i_surf(ng)
@@ -354,9 +356,18 @@ subroutine adiabatclimate_rce_wrapper(ptr, ng, P_i_surf, T_surf_guess, dim_T_gue
   logical(c_bool), intent(in) :: convecting_with_below_present
   integer(c_int), intent(in) :: dim_convecting_with_below
   logical(c_bool), intent(in) :: convecting_with_below(dim_convecting_with_below)
+  logical(c_bool), intent(in) :: custom_present
+  integer(c_int), intent(in) :: dim_sp_custom
+  character(c_char), intent(out) :: sp_custom(dim_sp_custom*s_str_len+1)
+  integer(c_int), intent(in) :: dim_P_custom
+  real(c_double), intent(in) :: P_custom(dim_P_custom)
+  integer(c_int), intent(in) :: dim1_mix_custom, dim2_mix_custom
+  real(c_double), intent(in) :: mix_custom(dim1_mix_custom,dim2_mix_custom)
   logical(c_bool), intent(out) :: converged
   character(c_char), intent(out) :: err(err_len+1)
 
+  character(s_str_len), allocatable :: sp_custom_f(:)
+  integer :: i, j, k
   logical, allocatable :: convecting_with_below_f(:)
   logical :: converged_f
   character(:), allocatable :: err_f
@@ -366,10 +377,22 @@ subroutine adiabatclimate_rce_wrapper(ptr, ng, P_i_surf, T_surf_guess, dim_T_gue
 
   allocate(convecting_with_below_f(dim_convecting_with_below))
   convecting_with_below_f = convecting_with_below
+
+  allocate(sp_custom_f(dim_sp_custom))
+  do i = 1,dim_sp_custom
+    do j = 1,s_str_len
+      k = j + (i - 1) * s_str_len
+      sp_custom_f(i)(j:j) = sp_custom(k)
+    enddo
+  enddo
 
   converged_f = .false.
-  if (convecting_with_below_present) then
-    converged_f = c%RCE(P_i_surf, T_surf_guess, T_guess, convecting_with_below_f, err_f)
+  if (convecting_with_below_present .and. custom_present) then
+    converged_f = c%RCE(P_i_surf, T_surf_guess, T_guess, convecting_with_below_f, sp_custom_f, P_custom, mix_custom, err_f)
+  elseif (convecting_with_below_present .and. .not.custom_present) then
+    converged_f = c%RCE(P_i_surf, T_surf_guess, T_guess, convecting_with_below_f, err=err_f)
+  elseif (.not.convecting_with_below_present .and. custom_present) then
+    converged_f = c%RCE(P_i_surf, T_surf_guess, T_guess, sp_custom=sp_custom_f, P_custom=P_custom, mix_custom=mix_custom, err=err_f)
   else
     converged_f = c%RCE(P_i_surf, T_surf_guess, T_guess, err=err_f)
   endif