Rossby Test Case Initialization

pyFV3/_config.py

@@ -274,6 +274,7 @@ class DynamicalCoreConfig:
     nf_omega: int = NamelistDefaults.nf_omega
     fv_sg_adj: int = NamelistDefaults.fv_sg_adj
     n_sponge: int = NamelistDefaults.n_sponge
+    sw_dynamics: bool = False  # TODO: Change to NamelistDefaults.sw_dynamics
     namelist_override: Optional[str] = None
 
     def __post_init__(self):

pyFV3/initialization/analytic_init.py

@@ -1,5 +1,8 @@
 from enum import Enum
 
+import pyFV3.initialization.test_cases.initialize_baroclinic as bc
+import pyFV3.initialization.test_cases.initialize_rossby as rossby
+import pyFV3.initialization.test_cases.initialize_tc as tc
 from ndsl import CubedSphereCommunicator, MetaEnumStr, QuantityFactory
 from ndsl.grid import GridData
 from ndsl.typing import Communicator
@@ -8,6 +11,7 @@
 
 class Cases(Enum, metaclass=MetaEnumStr):
     baroclinic = "baroclinic"
+    rossby = "rossby"
     tropicalcyclone = "tropicalcyclone"
 
 
@@ -34,12 +38,22 @@ def init_analytic_state(
     Returns:
         an instance of DycoreState class
     """
-    if analytic_init_case in Cases:  # type: ignore
-        if analytic_init_case == Cases.baroclinic.value:  # type: ignore
-            import pyFV3.initialization.test_cases.initialize_baroclinic as bc
+    # Cases that expect Cubed Sphere Communicator
+    spherical_cases = [
+        Cases.baroclinic.value,
+        Cases.tropicalcyclone.value,
+        Cases.rossby.value,
+    ]
 
-            assert isinstance(comm, CubedSphereCommunicator)
+    if analytic_init_case in spherical_cases:  # type: ignore
+        # TODO: Consider CubedSphereCommunicator check within individual init_*() calls
+        if not isinstance(comm, CubedSphereCommunicator):
+            raise TypeError(
+                f"Expected CubedSphereCommunicator instance for 'comm', "
+                f"got {type(comm).__name__} instead."
+            )
 
+        if analytic_init_case == Cases.baroclinic.value:  # type: ignore
             return bc.init_baroclinic_state(
                 grid_data=grid_data,
                 quantity_factory=quantity_factory,
@@ -48,18 +62,19 @@ def init_analytic_state(
                 moist_phys=moist_phys,
                 comm=comm,
             )
-
         elif analytic_init_case == Cases.tropicalcyclone.value:  # type: ignore
-            import pyFV3.initialization.test_cases.initialize_tc as tc
-
-            assert isinstance(comm, CubedSphereCommunicator)
-
             return tc.init_tc_state(
                 grid_data=grid_data,
                 quantity_factory=quantity_factory,
                 hydrostatic=hydrostatic,
                 comm=comm,
             )
+        elif analytic_init_case == Cases.rossby.value:  # type: ignore
+            return rossby.init_rossby_state(
+                grid_data=grid_data,
+                quantity_factory=quantity_factory,
+                comm=comm,
+            )
         else:
             raise ValueError(f"Case {analytic_init_case} not implemented")
     else:

pyFV3/initialization/test_cases/initialize_rossby.py

@@ -0,0 +1,210 @@
+""" Test case initialization for Rossby-Haurwitz wave 4
+
+Corresponds to Fortran shallow-water test #6 found in tools/test_cases.F90 of
+https://github.com/NOAA-GFDL/GFDL_atmos_cubed_sphere.git
+"""
+
+import numpy as np
+
+from ndsl import CubedSphereCommunicator, QuantityFactory, constants
+from ndsl.dsl.typing import Float
+from ndsl.grid import GridData
+from pyFV3.dycore_state import DycoreState
+from pyFV3.initialization import init_utils
+
+
+NHALO = constants.N_HALO_DEFAULT
+OMG = Float(7.848e-6)
+RK = Float(7.848e-6)
+R = Float(4.0)  # Wave Number (likely)
+GH0 = Float(8.0e3) * constants.GRAV
+
+
+def _preinit_for_all_sw(numpy_state: DycoreState, shape):
+    """Pre-initialization for all shallow water tests
+
+    Args:
+        numpy_state: DycoreState modified to update pe, pt, delp
+        shape: tuple
+    """
+    numpy_state.pe[:] = 0.0
+    numpy_state.pt[:] = 1.0
+
+    # Initialize Halo Corners
+    nx, ny, _ = init_utils.local_compute_size(shape)
+    numpy_state.delp[:NHALO, :NHALO] = 0.0
+    numpy_state.delp[:NHALO, NHALO + ny :] = 0.0
+    numpy_state.delp[NHALO + nx :, :NHALO] = 0.0
+    numpy_state.delp[NHALO + nx :, NHALO + ny :] = 0.0
+
+
+def _calc_rossby_winds(p1, p2):
+    """Calculates initial u or v winds specific to Rossby-Haurwitz wave test
+
+    Args
+        p1: np.ndarray
+        p2: np.ndarray
+
+    Returns
+        np.ndarray: representing u or v (D-winds)
+    """
+    muv = init_utils._find_midpoint_unit_vectors(
+        p1, p2
+    )  # TODO: Refactor to non-protected call
+    p3 = muv["midpoint"]
+    e2 = muv["unit_dir"]
+    ex = muv["exv"]
+    ey = muv["eyv"]
+    utmp = constants.RADIUS * OMG * np.cos(p3[:, :, 1]) + constants.RADIUS * RK * (
+        np.cos(p3[:, :, 1]) ** (R - 1)
+    ) * (R * np.sin(p3[:, :, 1]) ** 2 - np.cos(p3[:, :, 1]) ** 2) * np.cos(
+        R * p3[:, :, 0]
+    )
+    vtmp = (
+        -1
+        * constants.RADIUS
+        * RK
+        * R
+        * np.sin(p3[:, :, 1])
+        * np.sin(R * p3[:, :, 0])
+        * np.cos(p3[:, :, 1]) ** (R - 1)
+    )
+    return utmp * np.sum(e2 * ex, 2) + vtmp * np.sum(e2 * ey, 2)
+
+
+def _calc_rossby_delp(grid_data: GridData):
+    """Calculates initial delp, specific to Rossby-Haurwitz wave test
+
+    Args
+        grid_Data GridData
+
+    Returns
+        np.ndarray representing delp values
+    """
+    agd0 = grid_data.lon_agrid.data[:]
+    agd1 = grid_data.lat_agrid.data[:]
+
+    a = Float(0.5) * OMG * (2 * constants.OMEGA + OMG) * (np.cos(agd1) ** 2) + Float(
+        0.25
+    ) * RK * RK * (np.cos(agd1) ** (R + R)) * (
+        (R + 1) * (np.cos(agd1) ** 2)
+        + (2 * R * R - R - 2)
+        - 2 * (R * R) * np.cos(agd1) ** (-2)
+    )
+    b = (
+        (2 * (constants.OMEGA + OMG) * RK / ((R + 1) * (R + 2)))
+        * (np.cos(agd1) ** R)
+        * ((R * R + 2 * R + 2) - ((R + 1) * np.cos(agd1)) ** 2)
+    )
+    c = (
+        Float(0.25)
+        * RK
+        * RK
+        * (np.cos(agd1) ** (2 * R))
+        * ((R + 1) * (np.cos(agd1) ** 2) - (R + 2))
+    )
+    return GH0 + constants.RADIUS * constants.RADIUS * (
+        a + b * np.cos(R * agd0) + c * np.cos(2 * R * agd0)
+    )
+
+
+def _init_for_rossby(numpy_state: DycoreState, grid_data: GridData, shape):
+    """Initialization specific to Rossby-Haurwitz wave test
+
+    Args
+        numpy_state: DycoreState, modified to update the phis, delp, u, v
+        grid_Data: GridData
+    """
+    numpy_state.phis[:] = 0.0
+
+    # Calculate helper slices for delp, u+v winds
+    # similar to init_utils.compute_slices(nx, ny)
+    nx, ny, _ = init_utils.local_compute_size(shape)
+    islice = slice(NHALO, NHALO + nx)
+    islice_xtra = slice(NHALO, NHALO + nx + 1)
+    jslice = slice(NHALO, NHALO + ny)
+    jslice_xtra = slice(NHALO, NHALO + ny + 1)
+
+    # Initialize delp
+    delp_2d_buffer = (islice_xtra, jslice_xtra)
+    delp_buffer_0 = (islice_xtra, jslice_xtra, 0)
+    numpy_state.delp[delp_buffer_0] = _calc_rossby_delp(grid_data)[delp_2d_buffer]
+
+    grid = np.transpose(
+        np.stack(  # TODO: Refactor to non-protected _horizontal_data
+            [grid_data._horizontal_data.lon.data, grid_data._horizontal_data.lat.data]
+        ),
+        [1, 2, 0],
+    )
+
+    # Initialize u winds
+    p1 = grid[:-1, :, :]
+    p2 = grid[1:, :, :]
+    u_2d_buffer = (islice, jslice_xtra)
+    u_buffer_0 = (islice, jslice_xtra, 0)
+    numpy_state.u[u_buffer_0] = _calc_rossby_winds(p1, p2)[u_2d_buffer]
+
+    # Initialize v winds
+    p1 = grid[:, :-1, :]
+    p2 = grid[:, 1:, :]
+    v_2d_buffer = (islice_xtra, jslice)
+    v_buffer_0 = (islice_xtra, jslice, 0)
+    numpy_state.v[v_buffer_0] = _calc_rossby_winds(p1, p2)[v_2d_buffer]
+
+    # NOTE: test_cases.F90 has dtoa and atoc calls, but not implemented here.
+
+
+def _postinit_for_all_sw(numpy_state: DycoreState):
+    """Post-initialization from test_cases.F90 that applies to all shallow water tests
+
+    Args
+        numpy_state: DycoreState - modified
+    """
+
+    # NOTE: The cl/cl2 tracers from the original test_cases.F90 aren't brought over.
+    # NOTE: A-grid and C-grid winds are not initialized here.
+
+    numpy_state.delp[:, :, 1:] = numpy_state.delp[:, :, 0][:, :, np.newaxis]
+    numpy_state.u[:, :, 1:] = numpy_state.u[:, :, 0][:, :, np.newaxis]
+    numpy_state.v[:, :, 1:] = numpy_state.v[:, :, 0][:, :, np.newaxis]
+    numpy_state.ps[:] = numpy_state.delp[:, :, 0]
+
+
+def init_rossby_state(
+    grid_data: GridData,
+    quantity_factory: QuantityFactory,
+    comm: CubedSphereCommunicator,
+) -> DycoreState:
+    """
+    Create a DycoreState TODO: explain more
+
+    Args:
+        grid_data:              current selected grid data values
+        quantity_factory:       QuantityFactory
+        comm:                   CubedSphereCommunicator
+
+    Returns:
+        DycoreState
+    """
+
+    # TODO: Check sw_dunamics is True (https://github.com/NOAA-GFDL/PyFV3/pull/50)
+    #       May require a change to pass a config here in order to check.
+
+    sample_quantity = grid_data.lat
+    shape = (*sample_quantity.data.shape[0:2], grid_data.ak.data.shape[0])
+    numpy_state = init_utils.empty_numpy_dycore_state(shape)
+
+    _preinit_for_all_sw(numpy_state, shape)
+    _init_for_rossby(numpy_state, grid_data, shape)
+    _postinit_for_all_sw(numpy_state)
+
+    state = DycoreState.init_from_numpy_arrays(
+        numpy_state.__dict__,
+        sizer=quantity_factory.sizer,
+        backend=sample_quantity.metadata.gt4py_backend,
+    )
+
+    comm.halo_update(state.phis, n_points=NHALO)
+    comm.vector_halo_update(state.u, state.v, n_points=NHALO)
+
+    return state