ai2cm · rheacangeo · Jul 28, 2021 · Jul 28, 2021 · Aug 2, 2021 · Aug 2, 2021
diff --git a/.jenkins/actions/run_parallel_regression_tests.sh b/.jenkins/actions/run_parallel_regression_tests.sh
@@ -12,3 +12,9 @@ if [ ${python_env} == "virtualenv" ]; then
 else
     make savepoint_tests_mpi
 fi
+export TEST_ARGS="${TEST_ARGS} --compute_grid"
+if [ ${python_env} == "virtualenv" ]; then
+    CONTAINER_CMD="" make savepoint_tests_mpi
+else
+    make savepoint_tests_mpi
+fi
diff --git a/examples/standalone/runfile/acoustics.py b/examples/standalone/runfile/acoustics.py
@@ -39,18 +39,6 @@ def initialize_serializer(data_directory: str, rank: int = 0) -> serialbox.Seria
     )
 
 
-def read_grid(serializer: serialbox.Serializer, rank: int = 0) -> Grid:
-    """Uses the serializer to generate a Grid object from serialized data"""
-    grid_savepoint = serializer.get_savepoint("Grid-Info")[0]
-    grid_data = {}
-    grid_fields = serializer.fields_at_savepoint(grid_savepoint)
-    for field in grid_fields:
-        grid_data[field] = serializer.read(field, grid_savepoint)
-        if len(grid_data[field].flatten()) == 1:
-            grid_data[field] = grid_data[field][0]
-    return fv3core.testing.TranslateGrid(grid_data, rank).python_grid()
-
-
 def initialize_fv3core(backend: str, disable_halo_exchange: bool) -> None:
     """
     Initializes globalfv3core config to the arguments for single runs
@@ -155,7 +143,9 @@ def driver(
         serializer = initialize_serializer(data_directory)
         initialize_fv3core(backend, disable_halo_exchange)
         mpi_comm, communicator = set_up_communicator(disable_halo_exchange)
-        grid = read_grid(serializer)
+        grid = spec.make_grid_with_data_from_namelist(
+            spec.namelist, communicator, backend
+        )
         spec.set_grid(grid)
 
         input_data = read_input_data(grid, serializer)
@@ -169,8 +159,6 @@ def driver(
             grid.nested,
             grid.stretched_grid,
             spec.namelist.dynamical_core.acoustic_dynamics,
-            input_data["ak"],
-            input_data["bk"],
             input_data["pfull"],
             input_data["phis"],
         )

diff --git a/examples/standalone/runfile/dynamics.py b/examples/standalone/runfile/dynamics.py
@@ -205,21 +205,15 @@ def collect_data_and_write_to_file(
         else:
             mpi_comm = MPI.COMM_WORLD
 
-        # get grid from serialized data
-        grid_savepoint = serializer.get_savepoint("Grid-Info")[0]
-        grid_data = {}
-        grid_fields = serializer.fields_at_savepoint(grid_savepoint)
-        for field in grid_fields:
-            grid_data[field] = serializer.read(field, grid_savepoint)
-            if len(grid_data[field].flatten()) == 1:
-                grid_data[field] = grid_data[field][0]
-        grid = fv3core.testing.TranslateGrid(grid_data, rank).python_grid()
-        spec.set_grid(grid)
-
+        namelist = spec.namelist
         # set up grid-dependent helper structures
-        layout = spec.namelist.layout
-        partitioner = util.CubedSpherePartitioner(util.TilePartitioner(layout))
+        partitioner = util.CubedSpherePartitioner(util.TilePartitioner(namelist.layout))
         communicator = util.CubedSphereCommunicator(mpi_comm, partitioner)
+        # generate the grid
+        grid = spec.make_grid_with_data_from_namelist(
+            namelist, communicator, args.backend
+        )
+        spec.set_grid(grid)
 
         # create a state from serialized data
         savepoint_in = serializer.get_savepoint("FVDynamics-In")[0]
@@ -229,12 +223,10 @@ def collect_data_and_write_to_file(
         state = driver_object.state_from_inputs(input_data)
         dycore = fv3core.DynamicalCore(
             comm=communicator,
-            grid_data=spec.grid.grid_data,
-            stencil_factory=spec.grid.stencil_factory,
-            damping_coefficients=spec.grid.damping_coefficients,
+            grid_data=grid.grid_data,
+            stencil_factory=grid.stencil_factory,
+            damping_coefficients=grid.damping_coefficients,
             config=spec.namelist.dynamical_core,
-            ak=state["atmosphere_hybrid_a_coordinate"],
-            bk=state["atmosphere_hybrid_b_coordinate"],
             phis=state["surface_geopotential"],
         )
 
@@ -245,12 +237,10 @@ def collect_data_and_write_to_file(
             print("timestep 1")
         dycore.step_dynamics(
             state,
-            input_data["consv_te"],
+            namelist.consv_te,
             input_data["do_adiabatic_init"],
             input_data["bdt"],
-            input_data["ptop"],
-            input_data["n_split"],
-            input_data["ks"],
+            namelist.n_split,
         )
 
     if profiler is not None:
@@ -267,12 +257,10 @@ def collect_data_and_write_to_file(
                 print(f"timestep {i+2}")
             dycore.step_dynamics(
                 state,
-                input_data["consv_te"],
+                namelist.consv_te,
                 input_data["do_adiabatic_init"],
                 input_data["bdt"],
-                input_data["ptop"],
-                input_data["n_split"],
-                input_data["ks"],
+                namelist.n_split,
                 timestep_timer,
             )
         times_per_step.append(timestep_timer.times)

diff --git a/examples/wrapped/runfiles/baroclinic_test.py b/examples/wrapped/runfiles/baroclinic_test.py
@@ -275,8 +275,6 @@ def convert_3d_to_1d(state, field_names):
         stencil_factory=spec.grid.stencil_factory,
         damping_coefficients=spec.grid.damping_coefficients,
         config=spec.namelist.dynamical_core,
-        ak=state["atmosphere_hybrid_a_coordinate"],
-        bk=state["atmosphere_hybrid_b_coordinate"],
         phis=state["surface_geopotential"],
     )
     fvsubgridz = fv3core.DryConvectiveAdjustment(
@@ -308,9 +306,7 @@ def convert_3d_to_1d(state, field_names):
             wrapper.flags.consv_te,
             wrapper.flags.do_adiabatic_init,
             dt_atmos,
-            wrapper.flags.ptop,
-            wrapper.flags.n_split,
-            wrapper.flags.ks,
+            wrapper.flags.n_,
         )
         if spec.namelist.fv_sg_adj > 0:
             state["eastward_wind_tendency_due_to_physics"] = u_tendency

diff --git a/examples/wrapped/runfiles/fv3core_test.py b/examples/wrapped/runfiles/fv3core_test.py
@@ -265,8 +265,6 @@ def convert_3d_to_1d(state, field_names):
         stencil_factory=spec.grid.stencil_factory,
         damping_coefficients=spec.grid.damping_coefficients,
         config=spec.namelist.dynamical_core,
-        ak=state["atmosphere_hybrid_a_coordinate"],
-        bk=state["atmosphere_hybrid_b_coordinate"],
         phis=state["surface_geopotential"],
     )
 
@@ -298,9 +296,7 @@ def convert_3d_to_1d(state, field_names):
             wrapper.flags.consv_te,
             wrapper.flags.do_adiabatic_init,
             dt_atmos,
-            wrapper.flags.ptop,
             wrapper.flags.n_split,
-            wrapper.flags.ks,
         )
         if spec.namelist.fv_sg_adj > 0:
             state["eastward_wind_tendency_due_to_physics"] = u_tendency

diff --git a/fv3core/_config.py b/fv3core/_config.py
@@ -866,8 +866,6 @@ def namelist_to_flatish_dict(nml_input):
     return flatter_namelist
 
 
-# TODO: Before this can be used, we need to write a module to make the grid data
-# from files on disk and call it
 def make_grid_from_namelist(namelist, rank):
     shape_params = {}
     for narg in ["npx", "npy", "npz"]:
@@ -885,6 +883,18 @@ def make_grid_from_namelist(namelist, rank):
     return Grid(indices, shape_params, rank, namelist.layout)
 
 
+def make_grid_with_data_from_namelist(namelist, communicator, backend):
+    grid = make_grid_from_namelist(namelist, communicator.rank)
+    grid.make_grid_data(
+        npx=namelist.npx,
+        npy=namelist.npy,
+        npz=namelist.npz,
+        communicator=communicator,
+        backend=backend,
+    )
+    return grid
+
+
 def set_grid(in_grid):
     """Updates the global grid given another.