Expose finer-grained control over physics subroutines

ai2cm · Nov 1, 2023 · 53f19d0 · 53f19d0
1 parent 812a7c1
commit 53f19d0
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 4 deletions.
diff --git a/wrapper/lib/coupler_lib.F90 b/wrapper/lib/coupler_lib.F90
@@ -31,7 +31,12 @@ module coupler_lib
                                update_atmos_model_dynamics,        &
                                update_atmos_radiation_physics,     &
                                update_atmos_model_state,           &
-                               atmos_data_type, atmos_model_restart
+                               atmos_data_type,                    &
+                               atmos_model_restart,                &
+                               update_atmos_pre_radiation,         &
+                               update_atmos_radiation,             &
+                               update_atmos_physics
+
     !--- FMS old io
 #ifdef use_deprecated_io
     use fms_io_mod, only: fms_io_exit!< This can't be removed until fms_io is not used at all
@@ -194,12 +199,20 @@ subroutine do_dynamics() bind(c)
         Time_atmos = Time_atmos + Time_step_atmos
         call update_atmos_model_dynamics (Atm)
     end subroutine do_dynamics
+
+  ! substepped physics
+    subroutine do_pre_radiation() bind(c)
+        call update_atmos_pre_radiation (Atm)
+    end subroutine do_pre_radiation
+
+    subroutine do_radiation() bind(c)
+        call update_atmos_radiation (Atm)
+    end subroutine do_radiation
 
     subroutine do_physics() bind(c)
-        call update_atmos_radiation_physics (Atm)
-        call update_atmos_model_state (Atm)
+        call update_atmos_physics (Atm)
     end subroutine do_physics
-    
+
     subroutine compute_physics_subroutine() bind(c)
         call update_atmos_radiation_physics (Atm)
     end subroutine compute_physics_subroutine

diff --git a/wrapper/shield/wrapper/__init__.py b/wrapper/shield/wrapper/__init__.py
@@ -5,6 +5,9 @@
     initialize,
     step,
     step_dynamics,
+    step_radiation,
+    step_pre_radiation,
+    step_post_radiation_physics,
     step_physics,
     save_intermediate_restart_if_enabled,
     save_fortran_restart,

diff --git a/wrapper/templates/_wrapper.pyx b/wrapper/templates/_wrapper.pyx
@@ -539,3 +539,25 @@ def _get_diagnostic_data(int idx):
 
 
     return pace.util.Quantity(array, dims, units=units)
+
+
+cdef extern:
+    void do_pre_radiation()
+    void do_radiation()
+    void do_physics()
+
+
+def step_pre_radiation():
+    """Do pre-radiation computations (e.g. time varying logic)"""
+    do_pre_radiation()
+
+
+def step_radiation():
+    """Compute Radiative transfer scheme"""
+    do_radiation()
+
+
+def step_post_radiation_physics():
+    """Compute Post-radiation physics (e.g. moist physics turbulence)"""
+    # TODO ensure that IPD_control.first_step is set in this routine
+    do_physics()