Merge pull request #64 from jjuyeonkim/20250618_baroclinic12_ss_2_clean

Baroclinic Steady Test Case Initialization

Author: jjuyeonkim

.github/workflows/pace_tests.yaml

@@ -1,4 +1,8 @@
 # TODO: Re-activate after PR 60
+# TODO: Temporarily removing tests due to dependent changes in
+#       Pace#129, PyFV3#64 PRs; tests should be reverted back after
+#       these PRs are merged.
+
 # name: "pace main tests"
 # on:
 #   pull_request:

pyfv3/initialization/analytic_init.py

@@ -9,40 +9,43 @@
 from pyfv3.dycore_state import DycoreState
 
 
-class Cases(Enum, metaclass=MetaEnumStr):
-    baroclinic = "baroclinic"
+class AnalyticCase(Enum, metaclass=MetaEnumStr):
+    baroclinic_instability = "baroclinic_instability"
+    baroclinic_steady = "baroclinic_steady"
     rossby = "rossby"
     tropicalcyclone = "tropicalcyclone"
 
 
 def init_analytic_state(
-    analytic_init_case: str,
+    analytic_init_case: AnalyticCase,
     grid_data: GridData,
     quantity_factory: QuantityFactory,
     adiabatic: bool,
     hydrostatic: bool,
     moist_phys: bool,
+    sw_dynamics: bool,
     comm: Communicator,
 ) -> DycoreState:
     """
     This method initializes the chosen analytic test case type
     Args:
-        analytic_init_str:      test case specifier
+        analytic_init_case:     test case specifier
         grid_data:              current selected grid data values
         quantity_factory:       inclusion of QuantityFactory class
         adiabatic:              flag for adiabatic methods
         hydrostatic:            flag for hydrostatic methods
         moist_phys:             flag for including moisture physics methods
+        sw_dynamics:            flag for shallow water conditions (e.g., rossby)
         comm:                   inclusion of CubedSphereCommunicator class
 
     Returns:
         an instance of DycoreState class
     """
-    # Cases that expect Cubed Sphere Communicator
     spherical_cases = [
-        Cases.baroclinic.value,
-        Cases.tropicalcyclone.value,
-        Cases.rossby.value,
+        AnalyticCase.baroclinic_instability,
+        AnalyticCase.baroclinic_steady,
+        AnalyticCase.tropicalcyclone,
+        AnalyticCase.rossby,
     ]
 
     if analytic_init_case in spherical_cases:  # type: ignore
@@ -53,23 +56,40 @@ def init_analytic_state(
                 f"got {type(comm).__name__} instead."
             )
 
-        if analytic_init_case == Cases.baroclinic.value:  # type: ignore
+        if analytic_init_case == AnalyticCase.baroclinic_instability:  # type: ignore
             return bc.init_baroclinic_state(
                 grid_data=grid_data,
                 quantity_factory=quantity_factory,
                 adiabatic=adiabatic,
                 hydrostatic=hydrostatic,
                 moist_phys=moist_phys,
+                is_steady=False,
                 comm=comm,
             )
-        elif analytic_init_case == Cases.tropicalcyclone.value:  # type: ignore
+        elif analytic_init_case == AnalyticCase.baroclinic_steady:  # type: ignore
+            return bc.init_baroclinic_state(
+                grid_data=grid_data,
+                quantity_factory=quantity_factory,
+                adiabatic=adiabatic,
+                hydrostatic=hydrostatic,
+                moist_phys=moist_phys,
+                is_steady=True,
+                comm=comm,
+            )
+        elif analytic_init_case == AnalyticCase.tropicalcyclone:  # type: ignore
             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
+        elif analytic_init_case == AnalyticCase.rossby:  # type: ignore
+            # TODO sw_dynamics check is awkward here, and should be moved.
+            if sw_dynamics is False:
+                raise ValueError(
+                    "Rossby initialization requires dynamical core config "
+                    "sw_dynamics flag to be True."
+                )
             return rossby.init_rossby_state(
                 grid_data=grid_data,
                 quantity_factory=quantity_factory,

pyfv3/initialization/test_cases/initialize_baroclinic.py

@@ -1,10 +1,9 @@
-import math
-
 import numpy as np
 
 import ndsl.constants as constants
 import ndsl.dsl.gt4py_utils as utils
 from ndsl import CubedSphereCommunicator, QuantityFactory
+from ndsl.dsl.typing import Float
 from ndsl.grid import GridData
 from ndsl.grid.gnomonic import great_circle_distance_lon_lat, lon_lat_midpoint
 from pyfv3.dycore_state import DycoreState
@@ -13,22 +12,27 @@
 
 # maximum windspeed amplitude - close to windspeed of zonal-mean time-mean
 # jet stream in troposphere
-U0 = 35.0  # From Table VI of DCMIP2016
+U0 = Float(35.0)  # From Table VI of DCMIP2016
 # [lon, lat] of zonal wind perturbation centerpoint at 20E, 40N
-PCEN = [math.pi / 9.0, 2.0 * math.pi / 9.0]  # From Table VI of DCMIP2016
-U1 = 1.0
-SURFACE_PRESSURE = 1.0e5  # units of (Pa), from Table VI of DCMIP2016
-# NOTE RADIUS = 6.3712e6 in FV3 vs Jabowski paper 6.371229e6
-R = constants.RADIUS / 10.0  # Perturbation radiusfor test case 13
+PCEN = [
+    constants.PI / Float(9.0),
+    Float(2.0) * constants.PI / Float(9.0),
+]  # From Table VI of DCMIP2016
+SURFACE_PRESSURE = Float(1.0e5)  # units of (Pa), from Table VI of DCMIP2016
 NHALO = constants.N_HALO_DEFAULT
 
 
-def apply_perturbation(u_component, up, lon, lat):
+def apply_perturbation(u_component, up, lon, lat, is_steady: bool = False):
     """
     Apply a Gaussian perturbation to intiate a baroclinic wave in JRMS2006
     up is the maximum amplitude of the perturbation
     modifies u_component to include the perturbation of radius R
     """
+    if is_steady:
+        R = Float(1.0)  # Steady State radius for test case 12
+    else:
+        # NOTE RADIUS = 6.3712e6 in FV3 vs Jabowski paper 6.371229e6
+        R = constants.RADIUS / Float(10.0)  # Perturbation radius for test case 13
     r = np.zeros((u_component.shape[0], u_component.shape[1], 1))
     # Equation (11), distance from perturbation at 20E, 40N in JRMS2006
     r = great_circle_distance_lon_lat(PCEN[0], lon, PCEN[1], lat, constants.RADIUS, np)[
@@ -43,9 +47,21 @@ def apply_perturbation(u_component, up, lon, lat):
     )
 
 
-def baroclinic_perturbed_zonal_wind(eta_v, lon, lat):
+def baroclinic_perturbed_zonal_wind(
+    eta_v,
+    lon,
+    lat,
+    is_steady: bool = False,
+):
     u = zonal_wind(eta_v, lat)
-    apply_perturbation(u, U1, lon, lat)
+
+    if is_steady:
+        u1 = Float(0.0)  # Steady State for test case 12
+        # TODO: Check if Fortran side is actually 0 (not initialized)
+    else:
+        u1 = Float(1.0)  # Perturbation case for test case 13
+
+    apply_perturbation(u, u1, lon, lat, is_steady=is_steady)
     return u
 
 
@@ -59,12 +75,16 @@ def wind_component_calc(
     islice_grid,
     jslice,
     jslice_grid,
+    is_steady: bool = False,
 ):
     slice_grid = (islice_grid, jslice_grid)
     slice_3d = (islice, jslice, slice(None))
     u_component = np.zeros(shape)
     u_component[slice_3d] = baroclinic_perturbed_zonal_wind(
-        eta_v, lon[slice_grid], lat[slice_grid]
+        eta_v,
+        lon[slice_grid],
+        lat[slice_grid],
+        is_steady=is_steady,
     )
     u_component[slice_3d] = init_utils.local_coordinate_transformation(
         u_component[slice_3d],
@@ -95,6 +115,7 @@ def initialize_zonal_wind(
     jslice,
     jslice_grid,
     axis,
+    is_steady: bool = False,
 ):
     shape = u.shape
     uu1 = wind_component_calc(
@@ -107,6 +128,7 @@ def initialize_zonal_wind(
         islice,
         jslice,
         jslice_grid,
+        is_steady=is_steady,
     )
     uu3 = wind_component_calc(
         shape,
@@ -118,6 +140,7 @@ def initialize_zonal_wind(
         islice_grid,
         jslice,
         jslice,
+        is_steady=is_steady,
     )
     upper = (slice(None),) * axis + (slice(0, -1),)
     lower = (slice(None),) * axis + (slice(1, None),)
@@ -132,6 +155,7 @@ def initialize_zonal_wind(
         islice,
         jslice,
         jslice,
+        is_steady=is_steady,
     )
     u[islice, jslice, :] = 0.25 * (uu1 + 2.0 * uu2 + uu3)[islice, jslice, :]
 
@@ -161,6 +185,7 @@ def baroclinic_initialization(
     hydrostatic,
     nx,
     ny,
+    is_steady: bool = False,
 ):
     """
     Calls methods that compute initial state via the Jablonowski perturbation test case
@@ -191,6 +216,7 @@ def baroclinic_initialization(
         jslice=slice(0, ny),
         jslice_grid=slice(1, ny + 1),
         axis=1,
+        is_steady=is_steady,
     )
 
     initialize_zonal_wind(
@@ -206,6 +232,7 @@ def baroclinic_initialization(
         jslice=slice(0, ny + 1),
         jslice_grid=slice(0, ny + 1),
         axis=0,
+        is_steady=is_steady,
     )
 
     slice_3d = (slice(0, nx), slice(0, ny), slice(None))
@@ -251,6 +278,7 @@ def init_baroclinic_state(
     hydrostatic: bool,
     moist_phys: bool,
     comm: CubedSphereCommunicator,
+    is_steady: bool = False,
 ) -> DycoreState:
     """
     Create a DycoreState object with quantities initialized to the Jablonowski &
@@ -322,6 +350,7 @@ def init_baroclinic_state(
         hydrostatic=hydrostatic,
         nx=nx,
         ny=ny,
+        is_steady=is_steady,
     )
 
     init_utils.p_var(

pyfv3/initialization/test_cases/initialize_rossby.py

@@ -1,9 +1,11 @@
-""" Test case initialization for Rossby-Haurwitz wave 4
+"""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
 """
 
+from types import SimpleNamespace
+
 import numpy as np
 
 from ndsl import CubedSphereCommunicator, QuantityFactory, constants
@@ -20,11 +22,11 @@
 GH0 = Float(8.0e3) * constants.GRAV
 
 
-def _preinit_for_all_sw(numpy_state: DycoreState, shape):
+def _preinit_for_all_sw(numpy_state: SimpleNamespace, shape):
     """Pre-initialization for all shallow water tests
 
     Args:
-        numpy_state: DycoreState modified to update pe, pt, delp
+        numpy_state: SimpleNamespace modified to update pe, pt, delp
         shape: tuple
     """
     numpy_state.pe[:] = 0.0
@@ -108,11 +110,11 @@ def _calc_rossby_delp(grid_data: GridData):
     )
 
 
-def _init_for_rossby(numpy_state: DycoreState, grid_data: GridData, shape):
+def _init_for_rossby(numpy_state: SimpleNamespace, grid_data: GridData, shape):
     """Initialization specific to Rossby-Haurwitz wave test
 
     Args
-        numpy_state: DycoreState, modified to update the phis, delp, u, v
+        numpy_state: SimpleNamespace, modified to update the phis, delp, u, v
         grid_Data: GridData
     """
     numpy_state.phis[:] = 0.0
@@ -154,11 +156,11 @@ def _init_for_rossby(numpy_state: DycoreState, grid_data: GridData, shape):
     # NOTE: test_cases.F90 has dtoa and atoc calls, but not implemented here.
 
 
-def _postinit_for_all_sw(numpy_state: DycoreState):
+def _postinit_for_all_sw(numpy_state: SimpleNamespace):
     """Post-initialization from test_cases.F90 that applies to all shallow water tests
 
     Args
-        numpy_state: DycoreState - modified
+        numpy_state: SimpleNamespace - modified
     """
 
     # NOTE: The cl/cl2 tracers from the original test_cases.F90 aren't brought over.
@@ -176,7 +178,7 @@ def init_rossby_state(
     comm: CubedSphereCommunicator,
 ) -> DycoreState:
     """
-    Create a DycoreState TODO: explain more
+    Create an initial DycoreState for Rossby
 
     Args:
         grid_data:              current selected grid data values