diff --git a/FV3/atmos_cubed_sphere/tools/external_ic.F90 b/FV3/atmos_cubed_sphere/tools/external_ic.F90 index 66157aa5e..f81ffcd81 100644 --- a/FV3/atmos_cubed_sphere/tools/external_ic.F90 +++ b/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. diff --git a/FV3/wrapper/fv3gfs/wrapper/__init__.py b/FV3/wrapper/fv3gfs/wrapper/__init__.py index 57d00b105..bee770cfa 100644 --- a/FV3/wrapper/fv3gfs/wrapper/__init__.py +++ b/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 diff --git a/FV3/wrapper/fv3gfs/wrapper/physics_properties.json b/FV3/wrapper/fv3gfs/wrapper/physics_properties.json index 75ec9355e..3bb6bf85e 100644 --- a/FV3/wrapper/fv3gfs/wrapper/physics_properties.json +++ b/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", diff --git a/FV3/wrapper/templates/_wrapper.pyx b/FV3/wrapper/templates/_wrapper.pyx index 62ea91cd4..ab4845b28 100644 --- a/FV3/wrapper/templates/_wrapper.pyx +++ b/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 diff --git a/FV3/wrapper/templates/dynamics_data.F90 b/FV3/wrapper/templates/dynamics_data.F90 index 98780d3e1..c5d37ea17 100644 --- a/FV3/wrapper/templates/dynamics_data.F90 +++ b/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 diff --git a/FV3/wrapper/tests/test_all_mpi_requiring.py b/FV3/wrapper/tests/test_all_mpi_requiring.py index 06b2cf7bc..9f81165ca 100644 --- a/FV3/wrapper/tests/test_all_mpi_requiring.py +++ b/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() diff --git a/FV3/wrapper/tests/test_wind_transformations.py b/FV3/wrapper/tests/test_wind_transformations.py new file mode 100644 index 000000000..3cf927fc6 --- /dev/null +++ b/FV3/wrapper/tests/test_wind_transformations.py @@ -0,0 +1,153 @@ +import unittest +import os +import numpy as np +import pytest +import fv3gfs.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): + fv3gfs.wrapper.step() + state = fv3gfs.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 = fv3gfs.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. + + fv3gfs.wrapper.step_dynamics() + fv3gfs.wrapper.compute_physics() + pre_applied_physics_state = fv3gfs.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, + ) = fv3gfs.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" + + fv3gfs.wrapper.apply_physics() + updated_dynamical_core_state = fv3gfs.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 = fv3gfs.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"): + fv3gfs.wrapper.transform_dgrid_winds_to_agrid_winds(x_wind, y_wind) + + def test_transform_agrid_winds_to_dgrid_winds_invalid_units(self): + state = fv3gfs.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"): + fv3gfs.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)