[NOAA:Update] Bring back #15 & doubly periodic domain

README.md

@@ -11,6 +11,7 @@ Pace is an implementation of the FV3GFS / SHiELD atmospheric model developed by
 Full Sphinx documentation can be found at [https://ai2cm.github.io/pace/](https://ai2cm.github.io/pace/).
 
 **WARNING** This repo is under active development - supported features and procedures can change rapidly and without notice.
+
 ## Quickstart - bare metal
 
 ### Build
@@ -27,10 +28,13 @@ export BOOST_ROOT=BOOST/ROOT/boost_1_79_0
 ```
 
 When cloning Pace you will need to update the repository's submodules as well:
+
 ```shell
 git clone --recursive https://github.com/ai2cm/pace.git
 ```
+
 or if you have already cloned the repository:
+
 ```
 git submodule update --init --recursive
 ```
@@ -43,6 +47,7 @@ source venv_name/bin/activate
 ```
 
 Inside of your pace `venv` or conda environment pip install the Python requirements, GT4Py, and Pace:
+
 ```shell
 pip3 install -r requirements_dev.txt -c constraints.txt
 ```
@@ -52,6 +57,7 @@ Shell scripts to install Pace on specific machines such as Gaea can be found in
 ### Run
 
 With the environment activated, you can run an example baroclinic test case with the following command:
+
 ```shell
 mpirun -n 6 python3 -m pace.driver.run driver/examples/configs/baroclinic_c12.yaml
 
@@ -64,23 +70,30 @@ After the run completes, you will see an output direcotry `output.zarr`. An exam
 ### Environment variable configuration
 
 - `PACE_CONSTANTS`: Pace is bundled with various constants (see _util/pace/util/constants.py_).
+  - `FV3DYCORE` NOAA's FV3 dynamical core constants (original port)
+  - `GFS` Constant as defined in NOAA GFS
+  - `GEOS`  Constant as defined in GEOS v13
 - `PACE_FLOAT_PRECISION`: default precision of the field & scalars in the numerics. Default to 64.
 - `PACE_LOGLEVEL`: logging level to display (DEBUG, INFO, WARNING, ERROR, CRITICAL). Default to INFO.
 
 ## Quickstart - Docker
+
 ### Build
 
 While it is possible to install and build pace bare-metal, we can ensure all system libraries are installed with the correct versions by using a Docker container to test and develop pace.
 
 First, you will need to update the git submodules so that any dependencies are cloned and at the correct version:
+
 ```shell
 git submodule update --init --recursive
 ```
 
 Then build the `pace` docker image at the top level.
+
 ```shell
 make build
 ```
+
 ### Run
 
 ```shell
@@ -100,7 +113,6 @@ This git repository is laid out as a mono-repo, containing multiple independent
 
 ![Graph of interdependencies of Pace modules, generated from dependences.dot](./dependencies.svg)
 
-
 ## ML emulation
 
 An example of integration of an ML model replacing the microphysics parametrization is available on the `feature/microphysics-emulator` branch.

driver/examples/configs/baroclinic_c12_dp.yaml

@@ -0,0 +1,102 @@
+stencil_config:
+  compilation_config:
+    backend: numpy
+    rebuild: false
+    validate_args: true
+    format_source: false
+    device_sync: false
+grid_config:
+  type: generated
+  config:
+    grid_type: 4
+    dx_const: 3000.0
+    dy_const: 3000.0
+    deglat: 10.0
+initialization:
+  type: baroclinic
+performance_config:
+  collect_performance: true
+  experiment_name: c12_baroclinic
+nx_tile: 12
+nz: 79
+dt_atmos: 225
+minutes: 15
+layout:
+  - 1
+  - 1
+diagnostics_config:
+  path: output
+  output_format: netcdf
+  names:
+    - u
+    - v
+    - ua
+    - va
+    - pt
+    - delp
+    - qvapor
+    - qliquid
+    - qice
+    - qrain
+    - qsnow
+    - qgraupel
+  z_select:
+    - level: 65
+      names:
+        - pt
+dycore_config:
+  a_imp: 1.0
+  beta: 0.
+  consv_te: 0.
+  d2_bg: 0.
+  d2_bg_k1: 0.2
+  d2_bg_k2: 0.1
+  d4_bg: 0.15
+  d_con: 1.0
+  d_ext: 0.0
+  dddmp: 0.5
+  delt_max: 0.002
+  do_sat_adj: true
+  do_vort_damp: true
+  fill: true
+  hord_dp: 6
+  hord_mt: 6
+  hord_tm: 6
+  hord_tr: 8
+  hord_vt: 6
+  hydrostatic: false
+  k_split: 1
+  ke_bg: 0.
+  kord_mt: 9
+  kord_tm: -9
+  kord_tr: 9
+  kord_wz: 9
+  n_split: 1
+  nord: 3
+  nwat: 6
+  p_fac: 0.05
+  rf_cutoff: 3000.
+  rf_fast: true
+  tau: 10.
+  vtdm4: 0.06
+  z_tracer: true
+  do_qa: true
+  tau_i2s: 1000.
+  tau_g2v: 1200.
+  ql_gen: 0.001
+  ql_mlt: 0.002
+  qs_mlt: 0.000001
+  qi_lim: 1.0
+  dw_ocean: 0.1
+  dw_land: 0.15
+  icloud_f: 0
+  tau_l2v: 300.
+  tau_v2l: 90.
+  fv_sg_adj: 0
+  n_sponge: 48
+  u_max: 355.0
+
+physics_config:
+  hydrostatic: false
+  nwat: 6
+  do_qa: true

driver/pace/driver/driver.py

@@ -273,6 +273,10 @@ def from_dict(cls, kwargs: Dict[str, Any]) -> "DriverConfig":
             kwargs["grid_config"] = GridInitializerSelector.from_dict(
                 kwargs["grid_config"]
             )
+            grid_type = kwargs["grid_config"].config.grid_type
+            # Copy grid_type to the DycoreConfig if it's not the default value
+            if grid_type != 0:
+                kwargs["dycore_config"].grid_type = grid_type
 
         if (
             isinstance(kwargs["stencil_config"], dict)

driver/pace/driver/grid.py

@@ -85,12 +85,20 @@ class GeneratedGridConfig(GridInitializer):
         lon_target: desired center longitude for refined tile (deg)
         lat_target: desired center latitude for refined tile (deg)
         restart_path: if given, load vertical grid from restart file
+        grid_type: type of grid, 0 is a gnomonic cubed-sphere, 4 is doubly-periodic
+        dx_const: constant x-width of grid cells on a dp-grid
+        dy_const: constant y-width of grid cells on a dp-grid
+        deglat: latitude to use for coriolis calculations on a dp-grid
     """
 
     stretch_factor: Optional[float] = 1.0
     lon_target: Optional[float] = 350.0
     lat_target: Optional[float] = -90.0
     restart_path: Optional[str] = None
+    grid_type: Optional[int] = 0
+    dx_const: Optional[float] = 1000.0
+    dy_const: Optional[float] = 1000.0
+    deglat: Optional[float] = 15.0
 
     def get_grid(
         self,
@@ -99,7 +107,12 @@ def get_grid(
     ) -> Tuple[DampingCoefficients, DriverGridData, GridData]:
 
         metric_terms = MetricTerms(
-            quantity_factory=quantity_factory, communicator=communicator
+            quantity_factory=quantity_factory,
+            communicator=communicator,
+            grid_type=self.grid_type,
+            dx_const=self.dx_const,
+            dy_const=self.dy_const,
+            deglat=self.deglat,
         )
         if self.stretch_factor != 1:  # do horizontal grid transformation
             _transform_horizontal_grid(

driver/pace/driver/run.py

@@ -1,59 +1,15 @@
 import dataclasses
 import gc
-import logging
 from typing import Optional
 
 import click
 import yaml
 
-from pace.util import pace_log
-from pace.util.mpi import MPI
+from pace.util import AVAILABLE_LOG_LEVELS, pace_log
 
 from .driver import Driver, DriverConfig
 
 
-logger = logging.getLogger(__name__)
-
-
-log_levels = {
-    "info": logging.INFO,
-    "debug": logging.DEBUG,
-    "warning": logging.WARNING,
-    "error": logging.ERROR,
-    "critical": logging.CRITICAL,
-}
-
-
-def configure_logging(log_rank: Optional[int], log_level: str):
-    """
-    Configure logging for the driver.
-
-    Args:
-        log_rank: rank to log from, or 'all' to log to all ranks,
-            forced to 'all' if running without MPI
-        log_level: log level to use
-    """
-    level = log_levels[log_level.lower()]
-    if MPI is None:
-        logging.basicConfig(
-            level=level,
-            format="%(asctime)s [%(levelname)s] %(name)s:%(message)s",
-            handlers=[logging.StreamHandler()],
-            datefmt="%Y-%m-%d %H:%M:%S",
-        )
-    else:
-        if log_rank is None or int(log_rank) == MPI.COMM_WORLD.Get_rank():
-            logging.basicConfig(
-                level=level,
-                format=(
-                    f"%(asctime)s [%(levelname)s] (rank {MPI.COMM_WORLD.Get_rank()}) "
-                    "%(name)s:%(message)s"
-                ),
-                handlers=[logging.StreamHandler()],
-                datefmt="%Y-%m-%d %H:%M:%S",
-            )
-
-
 @click.command()
 @click.argument(
     "CONFIG_PATH",
@@ -76,7 +32,8 @@ def command_line(config_path: str, log_rank: Optional[int], log_level: str):
 
     CONFIG_PATH is the path to a DriverConfig yaml file.
     """
-    configure_logging(log_rank=log_rank, log_level=log_level)
+    level = AVAILABLE_LOG_LEVELS[log_level.lower()]
+    pace_log.setLevel(level)
     pace_log.info("loading DriverConfig from yaml")
     with open(config_path, "r") as f:
         config = yaml.safe_load(f)

fv3core/pace/fv3core/_config.py

@@ -195,6 +195,7 @@ class DynamicalCoreConfig:
     do_qa: bool = DEFAULT_BOOL
     layout: Tuple[int, int] = NamelistDefaults.layout
     grid_type: int = NamelistDefaults.grid_type
+    u_max: float = NamelistDefaults.u_max  # max windspeed for dp config
     do_f3d: bool = NamelistDefaults.do_f3d
     inline_q: bool = NamelistDefaults.inline_q
     do_skeb: bool = NamelistDefaults.do_skeb  # save dissipation estimate
@@ -334,6 +335,7 @@ def from_namelist(cls, namelist: Namelist) -> "DynamicalCoreConfig":
             do_qa=namelist.do_qa,
             layout=namelist.layout,
             grid_type=namelist.grid_type,
+            u_max=namelist.u_max,
             do_f3d=namelist.do_f3d,
             inline_q=namelist.inline_q,
             do_skeb=namelist.do_skeb,

tests/main/driver/test_example_configs.py

@@ -13,6 +13,7 @@
 
 TESTED_CONFIGS: List[str] = [
     "baroclinic_c12.yaml",
+    "baroclinic_c12_dp.yaml",
     "baroclinic_c12_comm_read.yaml",
     "baroclinic_c12_comm_write.yaml",
     "baroclinic_c12_null_comm.yaml",

util/HISTORY.md

@@ -4,7 +4,10 @@ History
 latest
 ------
 
+- Added `dx_const`, `dy_const`, `deglat`, and `u_max` namelist settings for doubly-periodic grids
+- Added `dx_const`, `dy_const`, and `deglat` to grid generation code for doubly-periodic grids
 - Added f32 support to halo exchange data transformation
+- Use one single logger, from logging.py
 
 v0.10.0
 -------

util/pace/util/__init__.py

@@ -54,7 +54,7 @@
 from .initialization import GridSizer, QuantityFactory, SubtileGridSizer
 from .io import read_state, write_state
 from .local_comm import LocalComm
-from .logging import pace_log
+from .logging import AVAILABLE_LOG_LEVELS, pace_log
 from .monitor import Monitor, NetCDFMonitor, ZarrMonitor
 from .mpi import MPIComm
 from .namelist import Namelist, NamelistDefaults

util/pace/util/grid/generation.py

@@ -222,6 +222,9 @@ def __init__(
         quantity_factory: util.QuantityFactory,
         communicator: util.CubedSphereCommunicator,
         grid_type: int = 0,
+        dx_const: float = 1000.0,
+        dy_const: float = 1000.0,
+        deglat: float = 15.0,
     ):
         assert grid_type < 3
         self._grid_type = grid_type
@@ -375,6 +378,9 @@ def from_tile_sizing(
         communicator: util.CubedSphereCommunicator,
         backend: str,
         grid_type: int = 0,
+        dx_const: float = 1000.0,
+        dy_const: float = 1000.0,
+        deglat: float = 15.0,
     ) -> "MetricTerms":
         sizer = util.SubtileGridSizer.from_tile_params(
             nx_tile=npx - 1,
@@ -393,6 +399,9 @@ def from_tile_sizing(
             quantity_factory=quantity_factory,
             communicator=communicator,
             grid_type=grid_type,
+            dx_const=dx_const,
+            dy_const=dy_const,
+            deglat=deglat,
         )
 
     @property

util/pace/util/logging.py

@@ -7,6 +7,16 @@
 
 LOGLEVEL = os.environ.get("PACE_LOGLEVEL", "INFO").upper()
 
+# Python log levels are hierarchical, therefore setting INFO
+# means DEBUG and everything lower will be logged.
+AVAILABLE_LOG_LEVELS = {
+    "info": logging.INFO,
+    "debug": logging.DEBUG,
+    "warning": logging.WARNING,
+    "error": logging.ERROR,
+    "critical": logging.CRITICAL,
+}
+
 
 def _pace_logger():
     name_log = logging.getLogger(__name__)