From 3cada57d044e46d3a1464e73410a7a91e443bc1a Mon Sep 17 00:00:00 2001
From: Nick Wogan <nicholaswogan@gmail.com>
Date: Fri, 12 Apr 2024 14:33:12 -0700
Subject: [PATCH] python wrapper for prep_atmosphere

---
 photochem/cython/EvoAtmosphere.pyx            | 22 +++++++++++++++++++
 photochem/cython/EvoAtmosphere_pxd.pxd        |  2 ++
 photochem/fortran/EvoAtmosphere_wrapper.f90   | 18 +++++++++++++++
 src/evoatmosphere/photochem_evoatmosphere.f90 |  2 ++
 .../photochem_evoatmosphere_utils.f90         |  2 ++
 5 files changed, 46 insertions(+)

diff --git a/photochem/cython/EvoAtmosphere.pyx b/photochem/cython/EvoAtmosphere.pyx
index b609f54..53bbdc2 100644
--- a/photochem/cython/EvoAtmosphere.pyx
+++ b/photochem/cython/EvoAtmosphere.pyx
@@ -75,6 +75,28 @@ cdef class EvoAtmosphere:
       wrk._ptr = self._wrk_ptr
       return wrk
 
+  def prep_atmosphere(self, ndarray[double, ndim=2] usol):
+    """Given `usol`, the densities of each species in the atmosphere,
+    this subroutine calculates reaction rates, photolysis rates, etc.
+    and puts this information into self.wrk. self.wrk contains all the
+    information needed for `dochem` to compute chemistry.
+
+    Parameters
+    ----------
+    usol : ndarray[double,ndim=2]
+        densities
+    """
+    cdef char err[ERR_LEN+1]
+    cdef int nq = self.dat.nq
+    cdef int nz = self.var.nz
+    cdef ndarray usol_ = np.asfortranarray(usol)
+    if usol_.shape[0] != nq or usol_.shape[1] != nz:
+      raise PhotoException("Input usol is the wrong size.")
+      
+    ea_pxd.evoatmosphere_prep_atmosphere_wrapper(&self._ptr, &nq, &nz, <double *>usol_.data, err)
+    if len(err.strip()) > 0:
+      raise PhotoException(err.decode("utf-8").strip())
+
   def out2atmosphere_txt(self,filename = None, bool overwrite = False, bool clip = True):
     """Saves state of the atmosphere using the mixing ratios in self.wrk.usol.
 
diff --git a/photochem/cython/EvoAtmosphere_pxd.pxd b/photochem/cython/EvoAtmosphere_pxd.pxd
index 84daa04..5906b0d 100644
--- a/photochem/cython/EvoAtmosphere_pxd.pxd
+++ b/photochem/cython/EvoAtmosphere_pxd.pxd
@@ -16,6 +16,8 @@ cdef extern void evoatmosphere_create_wrapper(void *ptr, char *mechanism_file,
                                             char *atmosphere_txt, char *data_dir, void *dat_ptr, void *var_ptr,
                                             void *wrk_ptr, char *err);
 
+cdef extern void evoatmosphere_prep_atmosphere_wrapper(void *ptr, int *nq, int *nz, double *usol, char *err)
+
 cdef extern void evoatmosphere_out2atmosphere_txt_wrapper(void *ptr, char *filename, bool *overwrite, bool *clip, char *err)   
 cdef extern void evoatmosphere_gas_fluxes_wrapper(void *ptr, int *nq, double *surf_fluxes, double *top_fluxes, char *err)
 cdef extern void evoatmosphere_set_lower_bc_wrapper(void *ptr, char *species, char *bc_type, 
diff --git a/photochem/fortran/EvoAtmosphere_wrapper.f90 b/photochem/fortran/EvoAtmosphere_wrapper.f90
index 176a9a0..25f2e35 100644
--- a/photochem/fortran/EvoAtmosphere_wrapper.f90
+++ b/photochem/fortran/EvoAtmosphere_wrapper.f90
@@ -74,6 +74,24 @@ subroutine evoatmosphere_create_wrapper(ptr, mechanism_file, &
     wrk_ptr = c_loc(pc%wrk) 
   end subroutine
 
+  subroutine evoatmosphere_prep_atmosphere_wrapper(ptr, nq, nz, usol, err) bind(c)
+    type(c_ptr), intent(in) :: ptr
+    integer(c_int), intent(in) :: nq, nz
+    real(c_double), intent(in) :: usol(nq, nz)
+    character(kind=c_char), intent(out) :: err(err_len+1)
+    
+    character(:), allocatable :: err_f
+    type(EvoAtmosphere), pointer :: pc
+    
+    call c_f_pointer(ptr, pc)
+    
+    call pc%prep_atmosphere(usol, err_f)
+    err(1) = c_null_char
+    if (allocated(err_f)) then
+      call copy_string_ftoc(err_f, err)
+    endif
+  end subroutine
+
   subroutine evoatmosphere_out2atmosphere_txt_wrapper(ptr, filename, overwrite, clip, err) bind(c)
     type(c_ptr), intent(in) :: ptr
     character(kind=c_char), intent(in) :: filename(*)
diff --git a/src/evoatmosphere/photochem_evoatmosphere.f90 b/src/evoatmosphere/photochem_evoatmosphere.f90
index 1ac0afa..efaec68 100644
--- a/src/evoatmosphere/photochem_evoatmosphere.f90
+++ b/src/evoatmosphere/photochem_evoatmosphere.f90
@@ -314,6 +314,8 @@ module subroutine update_vertical_grid(self, TOA_alt, TOA_pressure, err)
       character(:), allocatable, intent(out) :: err
     end subroutine
 
+    ! Both below are needed only for `evolve` routine and probably should not be used most of the time.
+
     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_utils.f90 b/src/evoatmosphere/photochem_evoatmosphere_utils.f90
index 4e668b6..0ddb4aa 100644
--- a/src/evoatmosphere/photochem_evoatmosphere_utils.f90
+++ b/src/evoatmosphere/photochem_evoatmosphere_utils.f90
@@ -793,6 +793,8 @@ module subroutine update_vertical_grid(self, TOA_alt, TOA_pressure, err)
 
   end subroutine
 
+  ! Below is mostly stuff needed in `evolve` routine
+
   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(:,:)