Online Zhao-Carr microphysics emulation (#218)

This PR adds the online emulation hooks for the ZC microphysics scheme and a barebones module to test it. Downstream users should override the emulation package with two top-level functions, store and microphysics. During runtime, two new namelist parameters under gfs_physics_nml are added which can be used to turn on/off the emulation components: emulate_zc_microphysics: calls emulation.microphysics save_zc_microphysics: calls emulation.store * Add gcsfs to requirements * Add prognostic hooks and module * Add emulation namelist parameters * Add conditionals to physics driver * Add specific configuration for training creation vs emulation * remove unused global variable * Add namelist flags for emulation * emulation.emulate debugging * Ignore egg-info * Add dummy model for testing * Fixed tensorflow import for emulation hooks * Change data saving namelist param name * Emulation integration test finished * Fix dummy model to correct vertical levels * Remove storage option from emulator, use monitor instead * Undo accidental import shift * Move emulation back into fv3net, remove from here * Reduce to basic call_py_fort interface for tests * Change to basic interface for testing call_py_fort usage * Update test config to run both call_py_fort functions * Update integration test * Fix integration test * Cleanup * Update emulation/README.md Co-authored-by: Noah D. Brenowitz <[email protected]> * Remove subprocess directory removal from callpyfort integration * Move emulation package to tests subdirectory * Update readme with more information about emulation * Add gotcha notes * Update language for emulation outputs Co-authored-by: Noah D. Brenowitz <[email protected]>
ai2cm · Sep 14, 2021 · 13d1638 · 13d1638
1 parent 201f4cc
commit 13d1638
Show file tree

Hide file tree

Showing 17 changed files with 128 additions and 686 deletions.
diff --git a/.gitignore b/.gitignore
@@ -15,3 +15,4 @@ depend
 **/.vscode
 fv3.exe
 *.tmp.f90
+**/*.egg-info
diff --git a/FV3/gfsphysics/GFS_layer/GFS_physics_driver.F90 b/FV3/gfsphysics/GFS_layer/GFS_physics_driver.F90
@@ -4566,8 +4566,27 @@ subroutine GFS_physics_driver                                 &
             call set_state("total_precipitation", rain1)
             call set_state("ratio_of_snowfall_to_rainfall", Diag%sr)
             call set_state("tendency_of_rain_water_mixing_ratio_due_to_microphysics", rainp)
-            call call_function("emulation.monitor", "store_zarr")
-            call call_function("emulation.monitor", "store_netcdf")
+
+            if (Model%emulate_zc_microphysics) then
+              ! apply microphysics emulator
+              call call_function("emulation", "microphysics")
+            endif
+
+            if (Model%save_zc_microphysics) then
+              call call_function("emulation", "store")
+            endif
+
+            call get_state("air_temperature_output", Stateout%gt0)
+            call get_state("specific_humidity_output", qv_post_precpd)
+            call get_state("cloud_water_mixing_ratio_output", qc_post_precpd)
+            call get_state("total_precipitation", rain1)
+
+            do k=1,levs
+              do i=1,im
+                Stateout%gq0(i,k,1) = qv_post_precpd(i,k)
+                Stateout%gq0(i,k,ntcw) = qc_post_precpd(i,k)
+              enddo
+            enddo
 #endif
 
           endif

diff --git a/FV3/gfsphysics/GFS_layer/GFS_typedefs.F90 b/FV3/gfsphysics/GFS_layer/GFS_typedefs.F90
@@ -1085,6 +1085,8 @@ module GFS_typedefs
     real(kind=kind_phys) :: sst_perturbation  ! Sea surface temperature perturbation to climatology or nudging SST (default 0.0 K)
     logical :: override_surface_radiative_fluxes  ! Whether to use Statein to override the surface radiative fluxes
     logical :: use_climatological_sst  ! Whether to allow the Python wrapper to override the sea surface temperature
+    logical :: emulate_zc_microphysics ! Use an emulator in place of ZC microphysics
+    logical :: save_zc_microphysics ! Save ZC microphysics state
 #ifdef CCPP
     ! From physcons.F90, updated/set in control_initialize
     real(kind=kind_phys) :: dxinv           ! inverse scaling factor for critical relative humidity, replaces dxinv in physcons.F90
@@ -3130,6 +3132,8 @@ subroutine control_initialize (Model, nlunit, fn_nml, me, master, &
     real(kind=kind_phys) :: sst_perturbation = 0.0  ! Sea surface temperature perturbation [K]
     logical :: override_surface_radiative_fluxes = .false.
     logical :: use_climatological_sst = .true.
+    logical :: emulate_zc_microphysics = .false.
+    logical :: save_zc_microphysics = .false.
 !--- END NAMELIST VARIABLES
 
     NAMELIST /gfs_physics_nml/                                                              &
@@ -3221,7 +3225,8 @@ subroutine control_initialize (Model, nlunit, fn_nml, me, master, &
                           !--- aerosol scavenging factors ('name:value' string array)
                                fscav_aero, &
                                sst_perturbation,                                            & 
-                               override_surface_radiative_fluxes, use_climatological_sst
+                               override_surface_radiative_fluxes, use_climatological_sst,   &
+                               emulate_zc_microphysics, save_zc_microphysics
 
 !--- other parameters 
     integer :: nctp    =  0                !< number of cloud types in CS scheme
@@ -3691,6 +3696,11 @@ subroutine control_initialize (Model, nlunit, fn_nml, me, master, &
     Model%sst_perturbation = sst_perturbation
     Model%override_surface_radiative_fluxes = override_surface_radiative_fluxes
     Model%use_climatological_sst = use_climatological_sst
+
+    !--- emulation parameters
+    Model%emulate_zc_microphysics = emulate_zc_microphysics
+    Model%save_zc_microphysics = save_zc_microphysics
+
 !--- tracer handling
     Model%ntrac            = size(tracer_names)
 #ifdef CCPP
@@ -4462,6 +4472,10 @@ subroutine control_print(Model)
         print *, ' evpco             : ', Model%evpco
         print *, ' wminco            : ', Model%wminco
         print *, ' '
+        print *, ' Z-C Emulation parameters'
+        print *, ' use ZC emulator   : ', Model%emulate_zc_microphysics
+        print *, ' save ZC state     : ', Model%save_zc_microphysics
+        print *, ' '
       endif
       if (Model%imp_physics == Model%imp_physics_wsm6 .or. Model%imp_physics == Model%imp_physics_thompson) then
         print *, ' Thompson microphysical parameters'

diff --git a/docker/Dockerfile b/docker/Dockerfile
@@ -136,7 +136,7 @@ RUN cd ${CALLPY} && make && make install && ldconfig
 RUN python3 -m pip install git+https://github.com/VulcanClimateModeling/fv3gfs-util.git@1d7c302b836befe905d776b0a972f464bfd3a255
 
 # copy emulation target
-COPY emulation/ /emulation
+COPY tests/emulation/ /emulation
 RUN cd /emulation && python3 -m pip install .
 
 ENV VAR_META_PATH=/emulation/microphysics_parameter_metadata.yaml \

diff --git a/emulation/README.md b/emulation/README.md
diff --git a/emulation/emulation/__init__.py b/emulation/emulation/__init__.py
diff --git a/emulation/emulation/_filesystem.py b/emulation/emulation/_filesystem.py
diff --git a/emulation/emulation/debug.py b/emulation/emulation/debug.py