Port wind transformations (#3)

This PR ports the wind transformations added in ai2cm/fv3gfs-fortran#345 to SHiELD-wrapper. Conveniently this did not require any updates to SHiELD that were not already present in #1.

wrapper/shield/wrapper/__init__.py

@@ -17,7 +17,9 @@
     compute_physics,
     apply_physics,
     flags,
-    DiagnosticInfo
+    DiagnosticInfo,
+    transform_agrid_winds_to_dgrid_winds,
+    transform_dgrid_winds_to_agrid_winds
 )
 from ._restart import get_restart_names, open_restart
 from . import examples

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
@@ -539,3 +541,85 @@ def _get_diagnostic_data(int idx):
 
 
     return pace.util.Quantity(array, dims, units=units)
+
+
+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

wrapper/templates/dynamics_data.F90

@@ -1,6 +1,7 @@
 module dynamics_data_mod
 
 use atmosphere_mod, only: Atm, mygrid
+use fv_grid_utils_mod, only: cubed_a2d, 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
@@ -114,4 +115,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(mygrid)%bd%isd
+        ied = Atm(mygrid)%bd%ied
+        jsd = Atm(mygrid)%bd%jsd
+        jed = Atm(mygrid)%bd%jed
+        npx = Atm(mygrid)%npx
+        npy = Atm(mygrid)%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(mygrid)%gridstruct, &
+            Atm(mygrid)%domain, &
+            Atm(mygrid)%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(mygrid)%bd%isd
+        ied = Atm(mygrid)%bd%ied
+        jsd = Atm(mygrid)%bd%jsd
+        jed = Atm(mygrid)%bd%jed
+        npx = Atm(mygrid)%npx
+        npy = Atm(mygrid)%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(mygrid)%domain, gridtype=DGRID_NE)
+        call complete_group_halo_update(i_pack, Atm(mygrid)%domain)
+
+        call cubed_to_latlon(&
+            u_halo, &
+            v_halo, &
+            ua_halo, &
+            va_halo, &
+            Atm(mygrid)%gridstruct, &
+            npx, &
+            npy, &
+            npz, &
+            mode, &
+            Atm(mygrid)%gridstruct%grid_type, &
+            Atm(mygrid)%domain, &
+            Atm(mygrid)%gridstruct%nested, &
+            Atm(mygrid)%flagstruct%c2l_ord, &
+            Atm(mygrid)%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

wrapper/tests/test_all_mpi_requiring.py

@@ -37,5 +37,8 @@ def test_get_initialization_time(self):
     def test_flags(self):
         run_unittest_script("test_flags.py")
 
+    def test_wind_transformations(self):
+        run_unittest_script("test_wind_transformations.py")
+
 if __name__ == "__main__":
     unittest.main()

wrapper/tests/test_wind_transformations.py

@@ -0,0 +1,153 @@
+import unittest
+import os
+import numpy as np
+import pytest
+import shield.wrapper
+from copy import deepcopy
+from mpi4py import MPI
+from util import (
+    get_default_config,
+    main,
+)
+
+
+test_dir = os.path.dirname(os.path.abspath(__file__))
+
+
+class WindTransformationTests(unittest.TestCase):
+    def __init__(self, *args, **kwargs):
+        super(WindTransformationTests, self).__init__(*args, **kwargs)
+
+    def setUp(self):
+        pass
+
+    def tearDown(self):
+        MPI.COMM_WORLD.barrier()
+
+    def test_transform_dgrid_winds_to_agrid_winds(self):
+        shield.wrapper.step()
+        state = shield.wrapper.get_state(
+            ["x_wind", "y_wind", "eastward_wind", "northward_wind"]
+        )
+        x_wind_fortran = state["x_wind"]
+        y_wind_fortran = state["y_wind"]
+        u_fortran = state["eastward_wind"]
+        v_fortran = state["northward_wind"]
+
+        # We expect exact reproduction of the fortran here, since the wrapper
+        # wraps the same routine that is used to convert the D-grid winds to the
+        # A-grid winds within the fortran model.
+        u_wrapper, v_wrapper = shield.wrapper.transform_dgrid_winds_to_agrid_winds(
+            x_wind_fortran, y_wind_fortran
+        )
+        np.testing.assert_equal(u_wrapper.view[:], u_fortran.view[:])
+        np.testing.assert_equal(v_wrapper.view[:], v_fortran.view[:])
+
+        assert u_wrapper.units == "m/s"
+        assert v_wrapper.units == "m/s"
+
+    def test_transform_agrid_winds_to_dgrid_winds(self):
+        # This test mimics the wind updating procedure from the
+        # atmosphere_state_update subroutine in the atmosphere.F90 module, but
+        # uses the Python wrapped cubed_a2d subroutine to convert the physics
+        # wind increments from the A to the D grid.
+
+        shield.wrapper.step_dynamics()
+        shield.wrapper.compute_physics()
+        pre_applied_physics_state = shield.wrapper.get_state(
+            [
+                "x_wind",
+                "y_wind",
+                "eastward_wind_before_physics",
+                "northward_wind_before_physics",
+                "eastward_wind_after_physics",
+                "northward_wind_after_physics",
+            ]
+        )
+        x_wind_before_physics = pre_applied_physics_state["x_wind"]
+        y_wind_before_physics = pre_applied_physics_state["y_wind"]
+
+        u_before_physics = pre_applied_physics_state["eastward_wind_before_physics"]
+        u_after_physics = pre_applied_physics_state["eastward_wind_after_physics"]
+
+        v_before_physics = pre_applied_physics_state["northward_wind_before_physics"]
+        v_after_physics = pre_applied_physics_state["northward_wind_after_physics"]
+
+        # Note that 3D variables from the physics component of the model have
+        # their vertical dimension flipped with respect to 3D variables from the
+        # dynamical core.  Therefore we need to flip the vertical dimension of
+        # these A-grid wind increments so that when we convert them to D-grid
+        # wind increments, they align with the D-grid winds in the dynamical
+        # core.
+        #
+        # deepcopy calls here are used out of convenience to construct Quantity
+        # objects of the same shape and metadata as others.  Their data is
+        # overwritten immediately.
+        u_increment = deepcopy(u_after_physics)
+        u_increment.view[:] = u_after_physics.view[:] - u_before_physics.view[:]
+        u_increment.view[:] = u_increment.view[:][::-1, :, :]
+
+        v_increment = deepcopy(v_after_physics)
+        v_increment.view[:] = v_after_physics.view[:] - v_before_physics.view[:]
+        v_increment.view[:] = v_increment.view[:][::-1, :, :]
+
+        (
+            x_wind_physics_increment,
+            y_wind_physics_increment,
+        ) = shield.wrapper.transform_agrid_winds_to_dgrid_winds(
+            u_increment, v_increment
+        )
+
+        assert x_wind_physics_increment.units == "m/s"
+        assert y_wind_physics_increment.units == "m/s"
+
+        shield.wrapper.apply_physics()
+        updated_dynamical_core_state = shield.wrapper.get_state(["x_wind", "y_wind"])
+        x_wind_after_physics = updated_dynamical_core_state["x_wind"]
+        y_wind_after_physics = updated_dynamical_core_state["y_wind"]
+
+        np.testing.assert_allclose(
+            x_wind_before_physics.view[:] + x_wind_physics_increment.view[:],
+            x_wind_after_physics.view[:],
+        )
+        np.testing.assert_allclose(
+            y_wind_before_physics.view[:] + y_wind_physics_increment.view[:],
+            y_wind_after_physics.view[:],
+        )
+
+    def test_transform_dgrid_winds_to_agrid_winds_invalid_units(self):
+        state = shield.wrapper.get_state(["x_wind", "y_wind"])
+        x_wind = state["x_wind"]
+        y_wind = state["y_wind"]
+        x_wind.metadata.units = "m/s/s"
+
+        with pytest.raises(ValueError, match="Input wind components"):
+            shield.wrapper.transform_dgrid_winds_to_agrid_winds(x_wind, y_wind)
+
+    def test_transform_agrid_winds_to_dgrid_winds_invalid_units(self):
+        state = shield.wrapper.get_state(["eastward_wind", "northward_wind"])
+        eastward_wind = state["eastward_wind"]
+        northward_wind = state["northward_wind"]
+        eastward_wind.metadata.units = "m/s/s"
+
+        with pytest.raises(ValueError, match="Input wind components"):
+            shield.wrapper.transform_agrid_winds_to_dgrid_winds(
+                eastward_wind, northward_wind
+            )
+
+
+if __name__ == "__main__":
+    config = get_default_config()
+    # Deactivate fv_subgrid_z for these tests.  Due to how the
+    # atmosphere_state_update subroutine is implemented in the atmosphere.F90
+    # module, the A-grid wind tendencies from the fv_subgrid_z subroutine are
+    # lumped into the A-grid wind tendencies from the physics.  We do not
+    # currently have easy access to the fv_subgrid_z tendency from the wrapper,
+    # and rather than implement that for the sake of the A-grid to D-grid
+    # transformation test, it is easier to simply turn it off. This way the
+    # A-grid wind tendency that is applied in the atmosphere_state_update
+    # subroutine comes only from the difference in the A-grid winds at the start
+    # of the physics and at the end of the physics, which is something that we
+    # can easily compute using existing wrapper infrastructure.
+    config["namelist"]["fv_core_nml"]["fv_sg_adj"] = -1
+    main(test_dir, config)