Consolidating whole atmosphere calculations into single module

.github/workflows/create_cache.yaml

@@ -35,14 +35,14 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v5
+        uses: actions/checkout@v6
 
       - name: Setup Python 3.11
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6
         with:
           python-version: '3.11'
 
-      - uses: actions/cache@v4
+      - uses: actions/cache@v5
         id: cache
         with:
           path: ~/.cache/pre-commit
@@ -64,14 +64,14 @@ jobs:
 
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v5
+        uses: actions/checkout@v6
 
       - name: Setup Python 3.11
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6
         with:
           python-version: '3.11'
 
-      - uses: actions/cache@v4
+      - uses: actions/cache@v5
         id: cache
         with:
           path: pyFV3/test_data

.github/workflows/gh_pages_doc.yaml

@@ -16,15 +16,15 @@ jobs:
 
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v5
+        uses: actions/checkout@v6
 
       - name: Configure Git Credentials
         run: |
           git config user.name github-actions[bot]
           git config user.email github-actions[bot]@users.noreply.github.com
 
       - name: Setup python
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6
         with:
           python-version: 3.11
 

.github/workflows/lint.yaml

@@ -16,16 +16,16 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v5
+        uses: actions/checkout@v6
 
       - name: Setup Python 3.11
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6
         with:
           python-version: '3.11'
 
       # Only restore (don't save) caches on PRs. New caches created from PRs won't be
       # accessible from other PRs, see workflows/create_cache.yaml.
-      - uses: actions/cache/restore@v4
+      - uses: actions/cache/restore@v5
         with:
           path: ~/.cache/pre-commit
           key: pre-commit_${{ env.pythonLocation }}_${{ hashFiles('.pre-commit-config.yaml') }}

.github/workflows/translate.yaml

@@ -35,7 +35,7 @@ jobs:
 
     steps:
       - name: Setup Python 3.11
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6
         with:
           python-version: '3.11'
 
@@ -44,14 +44,14 @@ jobs:
 
       - name: External trigger Checkout pyFV3
         if: ${{inputs.component_trigger}}
-        uses: actions/checkout@v5
+        uses: actions/checkout@v6
         with:
           submodules: 'recursive'
           repository: noaa-gfdl/pyFV3
           path: pyFV3
 
       - name: Checkout hash that triggered CI
-        uses: actions/checkout@v5
+        uses: actions/checkout@v6
         with:
           submodules: 'recursive'
           path: pyFV3/${{inputs.component_name}}
@@ -73,7 +73,7 @@ jobs:
       # accessible from other PRs, see workflows/create_cache.yaml.
       - name: Restore test_data (if cached)
         id: cache-restore
-        uses: actions/cache/restore@v4
+        uses: actions/cache/restore@v5
         with:
           key: ${{ env.DATA_PATH }}
           path: pyFV3/test_data

.github/workflows/unit_tests.yaml

@@ -22,12 +22,12 @@ jobs:
     steps:
 
       - name: Checkout repository
-        uses: actions/checkout@v5
+        uses: actions/checkout@v6
         with:
           submodules: 'recursive'
 
       - name: Setup Python 3.11
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6
         with:
           python-version: '3.11'
 

examples/notebook/test_functionality.ipynb

@@ -131,10 +131,11 @@
     "    layout=layout,\n",
     "    tile_partitioner=partitioner.tile,\n",
     "    tile_rank=cs_communicator.tile.rank,\n",
+    "    backend=backend,\n",
     ")\n",
     "\n",
     "# useful for easily allocating distributed data storages (fields)\n",
-    "quantity_factory = QuantityFactory.from_backend(sizer=sizer, backend=backend)\n",
+    "quantity_factory = QuantityFactory(sizer=sizer, backend=backend)\n",
     "\n",
     "compilation_config = CompilationConfig(backend=backend, communicator=cs_communicator)\n",
     "\n",

examples/standalone/runfile/acoustics.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 # type: ignore
 from types import SimpleNamespace
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Tuple
 
 import click
 import f90nml
@@ -14,24 +14,20 @@
     CubedSphereCommunicator,
     CubedSpherePartitioner,
     DaceConfig,
-    NullComm,
+    LocalComm,
+    MPIComm,
     StencilConfig,
     StencilFactory,
     TilePartitioner,
 )
+from ndsl.comm import Comm
 from ndsl.performance import Timer
 from ndsl.stencils.testing import Grid
 from pyfv3 import DynamicalCoreConfig
 from pyfv3.stencils import AcousticDynamics
 from pyfv3.testing import TranslateDynCore
 
 
-try:
-    from mpi4py import MPI
-except ImportError:
-    MPI = None
-
-
 def dycore_config_from_namelist(data_directory: str) -> DynamicalCoreConfig:
     """
     Reads the namelist at the given directory and sets
@@ -89,17 +85,15 @@ def get_state_from_input(
 def set_up_communicator(
     disable_halo_exchange: bool,
     layout: Tuple[int, int],
-) -> Tuple[Optional[MPI.Comm], Optional[CubedSphereCommunicator]]:
-    partitioner = CubedSpherePartitioner(TilePartitioner(layout))
-    if MPI is not None:
-        comm = MPI.COMM_WORLD
-    else:
-        comm = None
-    if not disable_halo_exchange:
-        assert comm is not None
-        cube_comm = CubedSphereCommunicator(comm, partitioner)
-    else:
-        cube_comm = CubedSphereCommunicator(NullComm(0, 0), partitioner)
+) -> Tuple[Comm, CubedSphereCommunicator]:
+    comm = (
+        LocalComm(rank=0, total_ranks=1, buffer={})
+        if disable_halo_exchange
+        else MPIComm()
+    )
+    cube_comm = CubedSphereCommunicator(
+        comm, CubedSpherePartitioner(TilePartitioner(layout))
+    )
     return comm, cube_comm
 
 

examples/standalone/runfile/compile.py

@@ -7,16 +7,12 @@
 
 import f90nml
 import gt4py.cartesian.config
+from mpi4py import MPI
 
-from ndsl import NullComm
+from ndsl import LocalComm
 from pyfv3 import DynamicalCoreConfig
 
 
-try:
-    from mpi4py import MPI
-except ImportError:
-    MPI = None
-
 local = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 sys.path.insert(0, local)
 from runfile.dynamics import get_experiment_info, setup_dycore  # noqa: E402
@@ -68,10 +64,8 @@ def parse_args() -> Namespace:
     for iteration in range(iterations):
         top_tile_rank = global_rank + size * iteration
         if top_tile_rank < sub_tiles:
-            mpi_comm = NullComm(
-                rank=top_tile_rank,
-                total_ranks=6 * sub_tiles,
-                fill_value=0.0,
+            mpi_comm = LocalComm(
+                rank=top_tile_rank, total_ranks=6 * sub_tiles, buffer_dict={}
             )
             gt4py.cartesian.config.cache_settings["dir_name"] = os.environ.get(
                 "GT_CACHE_ROOT", f".gt_cache_{mpi_comm.Get_rank():06}"

examples/standalone/runfile/dynamics.py

@@ -22,7 +22,8 @@
     CubedSphereCommunicator,
     CubedSpherePartitioner,
     DaceConfig,
-    NullComm,
+    LocalComm,
+    MPIComm,
     StencilConfig,
     StencilFactory,
     TilePartitioner,
@@ -286,10 +287,15 @@ def setup_dycore(
         namelist = f90nml.read(args.data_dir + "/input.nml")
         dycore_config = DynamicalCoreConfig.from_f90nml(namelist)
         experiment_name, is_baroclinic_test_case = get_experiment_info(args.data_dir)
-        if args.disable_halo_exchange:
-            mpi_comm = NullComm(MPI.COMM_WORLD.Get_rank(), MPI.COMM_WORLD.Get_size())
-        else:
-            mpi_comm = MPI.COMM_WORLD
+        mpi_comm = (
+            LocalComm(
+                rank=MPI.COMM_WORLD.Get_rank(),
+                total_ranks=MPI.COMM_WORLD.Get_size(),
+                buffer_dict={},
+            )
+            if args.disable_halo_exchange
+            else MPIComm()
+        )
         dycore, state, stencil_factory = setup_dycore(
             dycore_config,
             mpi_comm,

pyfv3/dycore_state.py

@@ -312,7 +312,7 @@ class DycoreState:
     rdg_var: Quantity = field(
         metadata={
             "name": "gas constant for dry air over variable gravity (RDGAS / grav_var) for Whole Atmosphere calculations",
-            "units": "(J/kg/deg) / (m s^-2)", # JK TODO: What are the units?
+            "units": "(J/kg/deg) / (m s^-2)",  # JK TODO: What are the units?
             "dims": [X_DIM, Y_DIM, Z_DIM],
             "intent": "inout",
         }
@@ -345,7 +345,9 @@ def init_zeros(cls, quantity_factory: QuantityFactory) -> Self:
                     dtype=Float,
                 ).data
         return cls.init_from_storages(
-            storages=initial_storages, sizer=quantity_factory.sizer
+            storages=initial_storages,
+            sizer=quantity_factory.sizer,
+            backend=quantity_factory.backend,
         )
 
     @classmethod
@@ -368,10 +370,9 @@ def init_from_numpy_arrays(
                     _field.metadata["units"],
                     origin=sizer.get_origin(dims),
                     extent=sizer.get_extent(dims),
-                    gt4py_backend=backend,
+                    backend=backend,
                 )
-        state = cls(**dict_state)
-        return state
+        return cls(**dict_state)  # type: ignore[arg-type,unused-ignore]
 
     @classmethod
     def init_from_storages(
@@ -380,6 +381,7 @@ def init_from_storages(
         sizer: GridSizer,
         bdt: float = 0.0,
         mdt: float = 0.0,
+        backend: str | None = None,
     ) -> Self:
         inputs = {}
         for _field in fields(cls):
@@ -391,6 +393,7 @@ def init_from_storages(
                     _field.metadata["units"],
                     origin=sizer.get_origin(dims),
                     extent=sizer.get_extent(dims),
+                    backend=backend,
                 )
                 inputs[_field.name] = quantity
         return cls(**inputs, bdt=bdt, mdt=mdt)