Merge pull request #278 from GEOS-ESM/feature/update_to_pyFV3

Update to pyFV3/NDSL structure & better interface
GEOS-ESM · Apr 18, 2024 · 09582d1 · 09582d1
2 parents aa51c36 + 5e5a765
commit 09582d1
Show file tree

Hide file tree

Showing 26 changed files with 1,717 additions and 1,176 deletions.
diff --git a/.gitignore b/.gitignore
@@ -2,7 +2,5 @@
 /@fvdycore
 /fvdycore
 /fvdycore@
-/@gtFV3
-/gtFV3
-/gtFV3@
 __pycache__
+/python/@pyFV3
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -1,6 +1,6 @@
 esma_set_this ()
 
-option(BUILD_GEOS_GTFV3_INTERFACE "Build GEOS-gtFV3 interface" OFF)
+option(BUILD_PYFV3_INTERFACE "Build pyFV3 interface" OFF)
 
 set (srcs
   sw.f90 jw.f90 testcases_3_4_5_6_stand_alone.f90
@@ -22,44 +22,46 @@ set (srcs
   fv_regridding_utils.F90
   rs_scaleMod.F90
   )
-if (BUILD_GEOS_GTFV3_INTERFACE)
+if (BUILD_PYFV3_INTERFACE)
   list (APPEND srcs
-  geos-gtfv3/geos_gtfv3_interface.f90
-  geos-gtfv3/geos_gtfv3_interface.c)
+  python/interface/cffi_lib/interface.f90
+  python/interface/cffi_lib/interface.c)
 endif ()
 
-if (BUILD_GEOS_GTFV3_INTERFACE)
+if (BUILD_PYFV3_INTERFACE)
 
-  message(STATUS "Building GEOS-gtFV3 interface")
+  message(STATUS "Building pyFV3 interface")
 
-  add_definitions(-DRUN_GTFV3)
+  add_definitions(-DRUN_PYFV3)
 
   # The Python library creation requires mpiexec/mpirun to run on a
   # compute node. Probably a weird SLURM thing?
   find_package(MPI REQUIRED)
   find_package(Python3 COMPONENTS Interpreter REQUIRED)
 
   # Set up some variables in case names change
-  set(GEOS_GTFV3_INTERFACE_LIBRARY     ${CMAKE_CURRENT_BINARY_DIR}/libgeos_gtfv3_interface_py.so)
-  set(GEOS_GTFV3_INTERFACE_HEADER_FILE ${CMAKE_CURRENT_BINARY_DIR}/geos_gtfv3_interface_py.h)
-  set(GEOS_GTFV3_INTERFACE_SRCS        ${CMAKE_CURRENT_SOURCE_DIR}/geos-gtfv3/geos_gtfv3_interface.py)
+  set(PYFV3_INTERFACE_LIBRARY     ${CMAKE_CURRENT_BINARY_DIR}/libpyFV3_interface_py.so)
+  set(PYFV3_INTERFACE_HEADER_FILE ${CMAKE_CURRENT_BINARY_DIR}/pyfv3_interface_py.h)
+  set(PYFV3_INTERFACE_FLAG_HEADER_FILE ${CMAKE_CURRENT_SOURCE_DIR}/python/interface/cffi_lib/fv_flags.h)
+  set(PYFV3_INTERFACE_SRCS        ${CMAKE_CURRENT_SOURCE_DIR}/python/interface/cffi_lib/interface.py)
 
   # This command creates the shared object library from Python
   add_custom_command(
-    OUTPUT ${GEOS_GTFV3_INTERFACE_LIBRARY}
+    OUTPUT ${PYFV3_INTERFACE_LIBRARY}
     # Note below is essentially:
     #  mpirun -np 1 python file
     # but we use the CMake options as much as we can for flexibility
-    COMMAND ${MPIEXEC_EXECUTABLE} ${MPIEXEC_NUMPROC_FLAG} 1 ${Python3_EXECUTABLE} ${GEOS_GTFV3_INTERFACE_SRCS}
-    BYPRODUCTS ${GEOS_GTFV3_INTERFACE_HEADER_FILE}
+    COMMAND ${CMAKE_COMMAND} -E copy_if_different ${PYFV3_INTERFACE_FLAG_HEADER_FILE} ${CMAKE_CURRENT_BINARY_DIR}
+    COMMAND ${MPIEXEC_EXECUTABLE} ${MPIEXEC_NUMPROC_FLAG} 1 ${Python3_EXECUTABLE} ${PYFV3_INTERFACE_SRCS}
+    BYPRODUCTS ${PYFV3_INTERFACE_HEADER_FILE}
     WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}
-    MAIN_DEPENDENCY ${GEOS_GTFV3_INTERFACE_SRCS}
-    COMMENT "Building gtfv3 interface library with Python"
+    MAIN_DEPENDENCY ${PYFV3_INTERFACE_SRCS}
+    COMMENT "Building pyFV3 interface library with Python"
     VERBATIM
     )
 
   # This creates a target we can use for dependencies and post build
-  add_custom_target(generate_python_interface_library DEPENDS ${GEOS_GTFV3_INTERFACE_LIBRARY})
+  add_custom_target(generate_python_interface_library DEPENDS ${PYFV3_INTERFACE_LIBRARY})
 
   # Because of the weird hacking of INTERFACE libraries below, we cannot
   # use the "usual" CMake calls to install() the .so. I think it's because
@@ -73,31 +75,31 @@ if (BUILD_GEOS_GTFV3_INTERFACE)
     # of a race condition (if install/lib/ isn't mkdir'd first)
     COMMAND ${CMAKE_COMMAND} -E make_directory ${CMAKE_INSTALL_PREFIX}/lib
     # Now we copy the file (if different...though not sure if this is useful)
-    COMMAND ${CMAKE_COMMAND} -E copy_if_different "${GEOS_GTFV3_INTERFACE_LIBRARY}" ${CMAKE_INSTALL_PREFIX}/lib
+    COMMAND ${CMAKE_COMMAND} -E copy_if_different "${PYFV3_INTERFACE_LIBRARY}" ${CMAKE_INSTALL_PREFIX}/lib
     )
 
   # We use INTERFACE libraries to create a sort of "fake" target library we can use
   # to make libFVdycoreCubed_GridComp.a depend on. It seems to work!
-  add_library(geos_gtfv3_interface_py INTERFACE)
+  add_library(pyfv3_interface_py INTERFACE)
 
   # The target_include_directories bits were essentially stolen from the esma_add_library
   # code...
-  target_include_directories(geos_gtfv3_interface_py INTERFACE
+  target_include_directories(pyfv3_interface_py INTERFACE
     $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}>
     $<BUILD_INTERFACE:${CMAKE_CURRENT_BINARY_DIR}> # stubs
     # modules and copied *.h, *.inc
     $<BUILD_INTERFACE:${esma_include}/${this}>
     $<INSTALL_INTERFACE:include/${this}>
     )
-  target_link_libraries(geos_gtfv3_interface_py INTERFACE ${GEOS_GTFV3_INTERFACE_LIBRARY})
+  target_link_libraries(pyfv3_interface_py INTERFACE ${PYFV3_INTERFACE_LIBRARY})
 
   # This makes sure the library is built first
-  add_dependencies(geos_gtfv3_interface_py generate_python_interface_library)
+  add_dependencies(pyfv3_interface_py generate_python_interface_library)
 
   # This bit is to resolve an issue and Google told me to do this. I'm not
   # sure that the LIBRARY DESTINATION bit actually does anything since
   # this is using INTERFACE
-  install(TARGETS geos_gtfv3_interface_py
+  install(TARGETS pyfv3_interface_py
     EXPORT ${PROJECT_NAME}-targets
     LIBRARY DESTINATION ${CMAKE_INSTALL_PREFIX}/lib
     )
@@ -106,13 +108,13 @@ endif ()
 
 set(dependencies MAPL GFTL_SHARED::gftl-shared GMAO_hermes GEOS_Shared esmf OpenMP::OpenMP_Fortran)
 
-if (BUILD_GEOS_GTFV3_INTERFACE)
+if (BUILD_PYFV3_INTERFACE)
   esma_add_library (${this}
     SRCS ${srcs}
     SUBCOMPONENTS fvdycore
     DEPENDENCIES ${dependencies}
     DEPENDENCIES ${GFDL}
-    DEPENDENCIES geos_gtfv3_interface_py) # Make the main library depend on the Python library
+    DEPENDENCIES pyfv3_interface_py) # Make the main library depend on the Python library
 else ()
   esma_add_library (${this}
     SRCS ${srcs}
@@ -192,17 +194,5 @@ ecbuild_add_executable (
    SOURCES interp_restarts_bin.F90
    LIBS ${this} OpenMP::OpenMP_Fortran)
 
-if (BUILD_GEOS_GTFV3_INTERFACE)
-  ecbuild_add_executable (
-    TARGET fv3_driver.x
-    SOURCES
-    geos-gtfv3/driver/FV_State_Utilities.F90
-    geos-gtfv3/driver/input/domain_dim.f90
-    geos-gtfv3/driver/input/grid_bounds.f90
-    geos-gtfv3/driver/input/input_scalars.f90
-    geos-gtfv3/driver/input/input_arrays.f90
-    geos-gtfv3/driver/fv3_driver.F90
-    LIBS ${this})
-endif ()
 
 add_subdirectory(scripts)
diff --git a/FV_StateMod.F90 b/FV_StateMod.F90
@@ -38,10 +38,11 @@ module FV_StateMod
 
    use fv_diagnostics_mod, only: prt_maxmin, prt_minmax, range_check, &
                                  get_vorticity, updraft_helicity, bunkers_vector, helicity_relative_CAPS
-#ifdef RUN_GTFV3
+#ifdef RUN_PYFV3
    use ieee_exceptions, only: ieee_get_halting_mode, ieee_set_halting_mode, ieee_all
-   use geos_gtfv3_interface_mod, only: geos_gtfv3_interface_f
-   use geos_gtfv3_interface_mod, only: geos_gtfv3_interface_f_init, geos_gtfv3_interface_f_finalize
+   use pyfv3_interface_mod, only: pyfv3_interface_f_run
+   use pyfv3_interface_mod, only: pyfv3_interface_f_init, pyfv3_interface_f_finalize, &
+                                  fv_flags_interface_type, make_fv_flags_C_interop
 #endif
 
 implicit none
@@ -264,8 +265,8 @@ module FV_StateMod
    real(REAL8), parameter ::  D180_0                  = 180.0
    real(REAL8), parameter ::  ratmax                  =  0.81
 
-#ifdef RUN_GTFV3
-   integer :: run_gtfv3 = 0
+#ifdef RUN_PYFV3
+   integer :: run_pyfv3 = 0
 #endif
 
 contains
@@ -749,8 +750,8 @@ subroutine FV_Setup(GC,LAYOUT_FILE, RC)
   call MAPL_MemUtilsWrite(VM, trim(Iam), RC=STATUS )
   VERIFY_(STATUS)
 
-#ifdef RUN_GTFV3
-  call MAPL_GetResource(MAPL, run_gtfv3, 'RUN_GTFV3:', default=0, RC=STATUS)
+#ifdef RUN_PYFV3
+  call MAPL_GetResource(MAPL, run_pyfv3, 'RUN_PYFV3:', default=0, RC=STATUS)
   VERIFY_(STATUS)
 #endif
 
@@ -819,11 +820,6 @@ subroutine FV_InitState (STATE, CLOCK, INTERNAL, IMPORT, GC, RC)
   integer    :: tile_in
   integer    :: gid, masterproc
 
-#ifdef RUN_GTFV3
-  logical :: halting_mode(5)
-  integer :: comm
-#endif
-
 ! BEGIN
 
 ! Retrieve the pointer to the state
@@ -1148,22 +1144,6 @@ subroutine FV_InitState (STATE, CLOCK, INTERNAL, IMPORT, GC, RC)
   call MAPL_MemUtilsWrite(VM, 'FV_StateMod: FV Initialize', RC=STATUS )
   VERIFY_(STATUS)
 
-#ifdef RUN_GTFV3
-  if (run_gtfv3 /= 0) then
-     ! call ESMF_VMGetCurrent(VM, _RC)
-     call ESMF_VMGet(VM, mpiCommunicator=comm, _RC)
-     ! A workaround to the issue of SIGFPE abort during importing of numpy, is to
-     ! disable trapping of FPEs temporarily, call the Python interface and resume trapping
-     call ieee_get_halting_mode(ieee_all, halting_mode)
-     call ieee_set_halting_mode(ieee_all, .false.)
-     call geos_gtfv3_interface_f_init( &
-          comm, &
-          FV_Atm(1)%npx, FV_Atm(1)%npy, FV_Atm(1)%npz, FV_Atm(1)%flagstruct%ntiles, &
-          IS, IE, JS, JE, ISD, IED, JSD, JED, real(STATE%DT), 7)
-     call ieee_set_halting_mode(ieee_all, halting_mode)
-  end if
-#endif
-
   RETURN_(ESMF_SUCCESS)
 
 end subroutine FV_InitState
@@ -1259,15 +1239,17 @@ subroutine FV_Run (STATE, EXPORT, CLOCK, GC, RC)
 
   logical :: NWAT_TEST
 
-#ifdef RUN_GTFV3
+#ifdef RUN_PYFV3
   type(ESMF_VM) :: vm
   integer :: comm, rank, mpierr
   real :: start, finish
+  logical :: halting_mode(5)
+  type(fv_flags_interface_type) :: c_fv_flags
 #endif
 
 ! Begin
 
-#ifdef RUN_GTFV3
+#ifdef RUN_PYFV3
   call ESMF_VMGetCurrent(vm, rc=status) ! pchakrab: replace with ESMF_GridCompGet(gc, VM=VM, _RC)
   call ESMF_VMGet(vm, mpiCommunicator=comm)
   call MPI_Comm_rank(comm, rank, mpierr)
@@ -1391,9 +1373,30 @@ subroutine FV_Run (STATE, EXPORT, CLOCK, GC, RC)
      _ASSERT( NWAT_TEST , 'NWAT must be either 0, 1, 3 or 6')
      FV_Atm(1)%ncnst = STATE%GRID%NQ
      deallocate( FV_Atm(1)%q )
-     allocate  ( FV_Atm(1)%q(isd:ied  ,jsd:jed  ,npz, FV_Atm(1)%ncnst) )
-    ! Echo FV3 setup
+     allocate  ( FV_Atm(1)%q(isd:ied  ,jsd:jed  ,npz, FV_Atm(1)%ncnst) )  
+     ! Echo FV3 setup
      call echo_fv3_setup()
+     ! Setup pyFV3 here since we need to know the exact nwat
+     ! We can do this because this is trigger _only once_.
+#ifdef RUN_PYFV3
+     if (run_pyfv3 /= 0) then
+      ! A workaround to the issue of SIGFPE abort during importing of numpy, is to
+      ! disable trapping of FPEs temporarily, call the Python interface and resume trapping
+      call ieee_get_halting_mode(ieee_all, halting_mode)
+      call ieee_set_halting_mode(ieee_all, .false.)
+      call make_fv_flags_C_interop(FV_Atm(1)%flagstruct, FV_Atm(1)%layout, c_fv_flags)
+      call pyfv3_interface_f_init( &
+            c_fv_flags, &
+            comm, &
+            FV_Atm(1)%npx, FV_Atm(1)%npy, FV_Atm(1)%npz, FV_Atm(1)%flagstruct%ntiles, &
+            FV_Atm(1)%bd%isc, FV_Atm(1)%bd%iec, FV_Atm(1)%bd%jsc, FV_Atm(1)%bd%jec, &
+            FV_Atm(1)%bd%isd, FV_Atm(1)%bd%ied, FV_Atm(1)%bd%jsd, FV_Atm(1)%bd%jed, &
+            real(STATE%DT), 7, &
+            FV_Atm(1)%ak, FV_Atm(1)%bk)
+      call ieee_set_halting_mode(ieee_all, halting_mode)
+      end if
+#endif
+
    endif
 
    select case ( FV_Atm(1)%flagstruct%nwat )
@@ -2003,8 +2006,8 @@ subroutine FV_Run (STATE, EXPORT, CLOCK, GC, RC)
     t_dt(:,:,:) = 0.0
     w_dt(:,:,:) = 0.0
 
-#ifdef RUN_GTFV3
-    if (run_gtfv3 == 0) then
+#ifdef RUN_PYFV3
+    if (run_pyfv3 == 0) then
        call cpu_time(start)
 #endif
 
@@ -2026,12 +2029,12 @@ subroutine FV_Run (STATE, EXPORT, CLOCK, GC, RC)
             FV_Atm(1)%neststruct, FV_Atm(1)%idiag, FV_Atm(1)%bd, FV_Atm(1)%parent_grid, FV_Atm(1)%domain, &
             FV_Atm(1)%diss_est, u_dt, v_dt, w_dt, t_dt, &
             time_total)
-#ifdef RUN_GTFV3
+#ifdef RUN_PYFV3
        call cpu_time(finish)
        if (rank == 0) print *, '0: fv_dynamics: time taken = ', finish - start, 's'
     else
        call cpu_time(start)
-       call geos_gtfv3_interface_f( &
+       call pyfv3_interface_f_run( &
             comm, &
             FV_Atm(1)%npx, FV_Atm(1)%npy, FV_Atm(1)%npz, FV_Atm(1)%flagstruct%ntiles, &
             FV_Atm(1)%bd%is, FV_Atm(1)%bd%ie, FV_Atm(1)%bd%js, FV_Atm(1)%bd%je, &
@@ -2044,12 +2047,10 @@ subroutine FV_Run (STATE, EXPORT, CLOCK, GC, RC)
             FV_Atm(1)%ps, FV_Atm(1)%pe, FV_Atm(1)%pk, FV_Atm(1)%peln, FV_Atm(1)%pkz, &
             FV_Atm(1)%phis, FV_Atm(1)%q_con, FV_Atm(1)%omga, &
             FV_Atm(1)%ua, FV_Atm(1)%va, FV_Atm(1)%uc, FV_Atm(1)%vc, &
-            ! input
-            FV_Atm(1)%ak, FV_Atm(1)%bk, &
             ! input/output
             FV_Atm(1)%mfx, FV_Atm(1)%mfy, FV_Atm(1)%cx, FV_Atm(1)%cy, FV_Atm(1)%diss_est)
        call cpu_time(finish)
-       print *, rank, ', geos_gtfv3_interface_f: time taken = ', finish - start, 's'
+       print *, rank, ', pyfv3_interface_f_run: time taken = ', finish - start, 's'
     end if
 #endif
 
@@ -2447,8 +2448,8 @@ subroutine FV_Finalize (STATE)
    !    call ESMF_GridDestroy  (STATE%GRID%GRID)
 #endif
 
-#ifdef RUN_GTFV3
-   if (run_gtfv3 /= 0) call geos_gtfv3_interface_f_finalize()
+#ifdef RUN_PYFV3
+   if (run_pyfv3 /= 0) call pyfv3_interface_f_finalize()
 #endif
 
  end subroutine FV_Finalize

diff --git a/geos-gtfv3/README.md b/geos-gtfv3/README.md
diff --git a/geos-gtfv3/driver/CMakeLists.txt b/geos-gtfv3/driver/CMakeLists.txt