Add Python-wrapped versions of the wind transformation subroutines

FV3/atmos_cubed_sphere/tools/external_ic.F90

@@ -199,7 +199,7 @@ module external_ic_mod
    real(kind=R_GRID), parameter :: cnst_0p20=0.20d0
    real :: deg2rad
    character (len = 80) :: source   ! This tells what the input source was for the data
-   public get_external_ic, get_cubed_sphere_terrain
+   public get_external_ic, get_cubed_sphere_terrain, cubed_a2d
 
 ! version number of this module
 ! Include variable "version" to be written to log file.

FV3/wrapper/fv3gfs/wrapper/__init__.py

@@ -20,7 +20,9 @@
     DiagnosticInfo,
     step_pre_radiation,
     step_radiation,
-    step_post_radiation_physics
+    step_post_radiation_physics,
+    transform_agrid_winds_to_dgrid_winds,
+    transform_dgrid_winds_to_agrid_winds
 )
 from ._restart import get_restart_names, open_restart
 from . import examples

FV3/wrapper/fv3gfs/wrapper/physics_properties.json

@@ -1,4 +1,18 @@
 [
+    {
+        "name": "northward_wind_before_physics",
+        "fortran_name": "vgrs",
+        "units": "m/s",
+        "container": "Statein",
+        "dims": ["z", "y", "x"]
+    },
+    {
+        "name": "eastward_wind_before_physics",
+        "fortran_name": "ugrs",
+        "units": "m/s",
+        "container": "Statein",
+        "dims": ["z", "y", "x"]
+    },
     {
         "name": "air_temperature_after_physics",
         "fortran_name": "gt0",

FV3/wrapper/templates/_wrapper.pyx

@@ -64,6 +64,8 @@ cdef extern:
 {% for item in flagstruct_properties %}
     void get_{{ item.fortran_name }}({{item.type_cython}} *{{ item.fortran_name }}_out)
 {% endfor %}
+    void transform_agrid_winds_to_dgrid_winds_subroutine(REAL_t *ua, REAL_t *va, REAL_t *u, REAL_t *v)
+    void transform_dgrid_winds_to_agrid_winds_subroutine(REAL_t *u, REAL_t *v, REAL_t *ua, REAL_t *va)
 
 cdef get_quantity_factory():
     cdef int nx, ny, nz, nz_soil, n_topo_variables
@@ -581,3 +583,91 @@ 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()
+
+
+def transform_agrid_winds_to_dgrid_winds(ua, va):
+    """Transform A-grid lat-lon winds to D-grid cubed-sphere winds.
+
+    This function wraps the cubed_a2d subroutine from the fortran model, making
+    it possible to call from Python without having to provide information about
+    the grid or domain decomposition.
+
+    Args:
+        ua : Quantity
+            The eastward wind defined on the A-grid
+        va : Quantity
+            The northward wind defined on the A-grid
+
+    Returns:
+        u, v : Quantity, Quantity
+            The cubed-sphere components of the wind defined on the D-grid
+    """
+    cdef REAL_t[:, :, :] buffer_ua
+    cdef REAL_t[:, :, :] buffer_va
+    cdef REAL_t[:, :, :] buffer_u
+    cdef REAL_t[:, :, :] buffer_v
+
+    if ua.units != va.units:
+        raise ValueError(
+            f"Input wind components have differing units from each other. "
+            f"ua has units {ua.units!r} and va has units {va.units!r}."
+        )
+
+    units = ua.units
+    ua = ua.transpose([pace.util.Z_DIM, pace.util.Y_DIM, pace.util.X_DIM])
+    va = va.transpose([pace.util.Z_DIM, pace.util.Y_DIM, pace.util.X_DIM])
+    buffer_ua = np.ascontiguousarray(ua.view[:])
+    buffer_va = np.ascontiguousarray(va.view[:])
+
+    allocator = get_quantity_factory()
+    u = allocator.empty([pace.util.Z_DIM, pace.util.Y_INTERFACE_DIM, pace.util.X_DIM], units, dtype=real_type)
+    v = allocator.empty([pace.util.Z_DIM, pace.util.Y_DIM, pace.util.X_INTERFACE_DIM], units, dtype=real_type)
+
+    with pace.util.recv_buffer(u.np.empty, u.view[:]) as buffer_u, pace.util.recv_buffer(v.np.empty, v.view[:]) as buffer_v:
+        transform_agrid_winds_to_dgrid_winds_subroutine(&buffer_ua[0, 0, 0], &buffer_va[0, 0, 0], &buffer_u[0, 0, 0], &buffer_v[0, 0, 0])
+
+    return u, v
+
+
+def transform_dgrid_winds_to_agrid_winds(u, v):
+    """Transform D-grid cubed-sphere winds to A-grid lat-lon winds.
+
+    This function wraps the cubed_to_latlon subroutine from the fortran model,
+    making it possible to call from Python without having to provide information
+    about the grid or domain decomposition.
+
+    Args:
+        u : Quantity
+            The x-wind defined on the D-grid
+        v : Quantity
+            The y-wind defined on the D-grid
+
+    Returns:
+        ua, va : Quantity, Quantity
+            The lat-lon components of the wind defined on the A-grid
+    """
+    cdef REAL_t[:, :, :] buffer_ua
+    cdef REAL_t[:, :, :] buffer_va
+    cdef REAL_t[:, :, :] buffer_u
+    cdef REAL_t[:, :, :] buffer_v
+
+    if u.units != v.units:
+        raise ValueError(
+            f"Input wind components have differing units from each other. "
+            f"u has units {u.units!r} and v has units {v.units!r}."
+        )
+
+    units = u.units
+    u = u.transpose([pace.util.Z_DIM, pace.util.Y_INTERFACE_DIM, pace.util.X_DIM])
+    v = v.transpose([pace.util.Z_DIM, pace.util.Y_DIM, pace.util.X_INTERFACE_DIM])
+    buffer_u = np.ascontiguousarray(u.view[:])
+    buffer_v = np.ascontiguousarray(v.view[:])
+
+    allocator = get_quantity_factory()
+    ua = allocator.empty([pace.util.Z_DIM, pace.util.Y_DIM, pace.util.X_DIM], units, dtype=real_type)
+    va = allocator.empty([pace.util.Z_DIM, pace.util.Y_DIM, pace.util.X_DIM], units, dtype=real_type)
+
+    with pace.util.recv_buffer(ua.np.empty, ua.view[:]) as buffer_ua, pace.util.recv_buffer(va.np.empty, va.view[:]) as buffer_va:
+        transform_dgrid_winds_to_agrid_winds_subroutine(&buffer_u[0, 0, 0], &buffer_v[0, 0, 0], &buffer_ua[0, 0, 0], &buffer_va[0, 0, 0])
+
+    return ua, va

FV3/wrapper/templates/dynamics_data.F90

@@ -1,6 +1,10 @@
 module dynamics_data_mod
 
 use atmosphere_mod, only: Atm, mytile
+use external_ic_mod, only: cubed_a2d
+use fv_grid_utils_mod, only: cubed_to_latlon
+use mpp_domains_mod, only: DGRID_NE, mpp_update_domains
+use fv_mp_mod,       only: start_group_halo_update, complete_group_halo_update, group_halo_update_type
 use tracer_manager_mod, only: get_tracer_names, get_number_tracers, get_tracer_index
 use field_manager_mod,  only: MODEL_ATMOS
 use iso_c_binding
@@ -112,4 +116,111 @@ subroutine get_tracer_name(tracer_index, tracer_name_out, tracer_long_name_out,
         enddo
     end subroutine get_tracer_name
 
+    subroutine transform_agrid_winds_to_dgrid_winds_subroutine(ua, va, u, v) bind(c)
+        ! Wraps the internal cubed_a2d subroutine in a more convenient way, handling
+        ! halo updates and the additional grid arguments within fortran for convenience.
+        real, intent(in), dimension(i_start():i_end(),j_start():j_end(),1:nz()) :: ua, va
+        real, intent(out), dimension(i_start():i_end(),j_start():j_end()+1,1:nz()) :: u
+        real, intent(out), dimension(i_start():i_end()+1,j_start():j_end(),1:nz()) :: v
+
+        real, allocatable, dimension(:,:,:) :: ua_halo, va_halo, u_halo, v_halo
+
+        integer :: is, ie, js, je, isd, ied, jsd, jed, npx, npy, npz
+
+        is = i_start()
+        ie = i_end()
+        js = j_start()
+        je = j_end()
+        isd = Atm(mytile)%bd%isd
+        ied = Atm(mytile)%bd%ied
+        jsd = Atm(mytile)%bd%jsd
+        jed = Atm(mytile)%bd%jed
+        npx = Atm(mytile)%npx
+        npy = Atm(mytile)%npy
+        npz = nz()
+
+        allocate(ua_halo(isd:ied,jsd:jed,1:npz))
+        allocate(va_halo(isd:ied,jsd:jed,1:npz))
+        allocate(u_halo(isd:ied,jsd:jed+1,1:npz))
+        allocate(v_halo(isd:ied+1,jsd:jed,1:npz))
+
+        ! Note we do not need to do a halo update here, since cubed_a2d takes
+        ! of this internally.
+        ua_halo(is:ie,js:je,1:npz) = ua
+        va_halo(is:ie,js:je,1:npz) = va
+
+        call cubed_a2d(&
+            npx, &
+            npy, &
+            npz, &
+            ua_halo, &
+            va_halo, &
+            u_halo, &
+            v_halo, &
+            Atm(mytile)%gridstruct, &
+            Atm(mytile)%domain, &
+            Atm(mytile)%bd &
+        )
+
+        u = u_halo(is:ie,js:je+1,1:npz)
+        v = v_halo(is:ie+1,js:je,1:npz)
+    end subroutine transform_agrid_winds_to_dgrid_winds_subroutine
+
+    subroutine transform_dgrid_winds_to_agrid_winds_subroutine(u, v, ua, va) bind(c)
+        ! Wraps the internal cubed_to_latlon subroutine in a more convenient way, 
+        ! handling halo updates and the additional grid arguments within fortran
+        ! for convenience.
+        real, intent(in), dimension(i_start():i_end(),j_start():j_end()+1,1:nz()) :: u
+        real, intent(in), dimension(i_start():i_end()+1,j_start():j_end(),1:nz()) :: v
+        real, intent(out), dimension(i_start():i_end(),j_start():j_end(),1:nz()) :: ua, va
+
+        real, allocatable, dimension(:,:,:) :: ua_halo, va_halo, u_halo, v_halo
+
+        type(group_halo_update_type), save :: i_pack
+        integer :: is, ie, js, je, isd, ied, jsd, jed, npx, npy, npz, mode
+
+        is = i_start()
+        ie = i_end()
+        js = j_start()
+        je = j_end()
+        isd = Atm(mytile)%bd%isd
+        ied = Atm(mytile)%bd%ied
+        jsd = Atm(mytile)%bd%jsd
+        jed = Atm(mytile)%bd%jed
+        npx = Atm(mytile)%npx
+        npy = Atm(mytile)%npy
+        npz = nz()
+        mode = 1
+
+        allocate(ua_halo(isd:ied,jsd:jed,1:npz))
+        allocate(va_halo(isd:ied,jsd:jed,1:npz))
+        allocate(u_halo(isd:ied,jsd:jed+1,1:npz))
+        allocate(v_halo(isd:ied+1,jsd:jed,1:npz))
+
+        u_halo(is:ie,js:je+1,1:npz) = u
+        v_halo(is:ie+1,js:je,1:npz) = v
+        call start_group_halo_update(i_pack, u_halo, v_halo, Atm(mytile)%domain, gridtype=DGRID_NE)
+        call complete_group_halo_update(i_pack, Atm(mytile)%domain)
+
+        call cubed_to_latlon(&
+            u_halo, &
+            v_halo, &
+            ua_halo, &
+            va_halo, &
+            Atm(mytile)%gridstruct, &
+            npx, &
+            npy, &
+            npz, &
+            mode, &
+            Atm(mytile)%gridstruct%grid_type, &
+            Atm(mytile)%domain, &
+            Atm(mytile)%gridstruct%nested, &
+            Atm(mytile)%flagstruct%c2l_ord, &
+            Atm(mytile)%bd &
+        )
+
+        ua = ua_halo(is:ie,js:je,1:npz)
+        va = va_halo(is:ie,js:je,1:npz)
+    end subroutine transform_dgrid_winds_to_agrid_winds_subroutine
+
 end module dynamics_data_mod

FV3/wrapper/tests/test_all_mpi_requiring.py

@@ -46,6 +46,9 @@ def test_flags(self):
     def test_set_ocean_surface_temperature(self):
         run_unittest_script("test_set_ocean_surface_temperature.py")
 
+    def test_wind_transformations(self):
+        run_unittest_script("test_wind_transformations.py")
+
 
 if __name__ == "__main__":
     unittest.main()