From 669f9028fb66f316f251b0991ae3b590b4c7a6ea Mon Sep 17 00:00:00 2001 From: Nick Wogan Date: Thu, 19 Dec 2024 12:09:31 -0800 Subject: [PATCH] python wrapper --- clima/cython/AdiabatClimate.pyx | 46 ++++++++++++++++++++++++++--- clima/cython/AdiabatClimate_pxd.pxd | 10 +++++-- clima/cython/_clima.pyx | 9 ++++++ clima/fortran/AdiabatClimate.f90 | 33 +++++++++++++++++---- 4 files changed, 86 insertions(+), 12 deletions(-) diff --git a/clima/cython/AdiabatClimate.pyx b/clima/cython/AdiabatClimate.pyx index 67b61a0..d43c90b 100644 --- a/clima/cython/AdiabatClimate.pyx +++ b/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, P_i_surf.data, &T_surf_guess, - &dim_T_guess, T_guess.data, &convecting_with_below_present, - &dim_convecting_with_below, convecting_with_below_.data, &converged, err) + wa_pxd.adiabatclimate_rce_wrapper( + self._ptr, &ng, P_i_surf.data, &T_surf_guess, &dim_T_guess, T_guess.data, + &convecting_with_below_present, &dim_convecting_with_below, convecting_with_below_.data, + &custom_present, &dim_sp_custom, sp_custom_c.data, &dim_P_custom, P_custom.data, + &dim1_mix_custom, &dim2_mix_custom, mix_custom.data, + &converged, err + ) if len(err.strip()) > 0: raise ClimaException(err.decode("utf-8").strip()) return converged diff --git a/clima/cython/AdiabatClimate_pxd.pxd b/clima/cython/AdiabatClimate_pxd.pxd index 73bfb91..c6d6e1b 100644 --- a/clima/cython/AdiabatClimate_pxd.pxd +++ b/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) diff --git a/clima/cython/_clima.pyx b/clima/cython/_clima.pyx index fda8541..71c97c4 100644 --- a/clima/cython/_clima.pyx +++ b/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 diff --git a/clima/fortran/AdiabatClimate.f90 b/clima/fortran/AdiabatClimate.f90 index 8e4c8cc..575d401 100644 --- a/clima/fortran/AdiabatClimate.f90 +++ b/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