Feature/fv3core fortran api #395

oelbert · 2022-11-29T20:23:33Z

Purpose

This PR adds an API to call the Pace dycore from GEOS.

Code changes:

fv3core/pace/fv3core/initialization/geos_wrapper.py: Added a wrapper api to access the dynamical core from GEOS. Initialize it with a namelist, communicator, and backend, call it with numpy arrays and it returns a dictionary of numpy arrays.
tests/main/fv3core/test_init_from_geos.py: A simple test that checks that the GEOS wrapper is initialized correctly, populates the dycore state correctly, steps through the dycore, and outputs arrays in the output dictionary.
fv3core/pace/fv3core/init.py: Added the geos wrapper to init.py to make it importable at the base level.
util/pace/util/grid/eta.py: Added code to setup atmosphere with 91 vertical levels for GEOS runs.
util/pace/util/initialization/sizer.py: Revised the from_namelist method to work with our updated namelist configuration that doesn't use fv_core_nml as a key.

Checklist

Before submitting this PR, please make sure:

You have followed the coding standards guidelines established at Code Review Checklist.
Docstrings and type hints are added to new and updated routines, as appropriate
All relevant documentation has been updated or added (e.g. README, CONTRIBUTING docs)
For each public change and fix in pace-util, HISTORY has been updated
Unit tests are added or updated for non-stencil code changes

elynnwu · 2022-12-01T17:26:39Z

fv3core/pace/fv3core/initialization/geos_wrapper.py

+from pace import fv3core
+
+
+class GeosDycoreWrapper:


[nit] I believe NASA always refers it as GEOS

elynnwu · 2022-12-01T17:30:23Z

fv3core/pace/fv3core/initialization/geos_wrapper.py

+
+        return self.output_dictionary
+
+    def _put_fortran_data_in_dycore(


Just wanted to check - is the idea here to not use translate fvdynamics anymore?

Yup, that's the idea

elynnwu · 2022-12-01T17:36:10Z

tests/main/fv3core/test_init_from_geos.py

+        },
+        "initialization": {"type": "baroclinic"},
+        "nx_tile": 12,
+        "nz": 79,


I think you want to test the additional vertical case you added, either 91 or 72

elynnwu · 2022-12-01T18:06:12Z

util/pace/util/grid/eta.py

@@ -204,14 +204,372 @@ def set_hybrid_pressure_coefficients(km: int) -> HybridPressureCoefficients:
                1,
            ]
        )
+
+    elif km == 91:


I think this is fine since this is how we set up our npz=79 case, but we should probably merge this so it is possible to replace ak/bk with custom values.

elynnwu · 2022-12-01T18:09:38Z

fv3core/pace/fv3core/initialization/geos_wrapper.py

+class GeosDycoreWrapper:
+    """
+    Provides an interface for the Geos model to access the Pace dycore.
+    Takes numpy arrays as inputs, returns a dictionary of numpy arrays as outputs


So this only works for numpy and cpu backend? We should extend this to take cupy array as well?

Ideally we can also make this work on GPU. The reason for the current input/output setup is to comport with Purnendu's code, which puts the Fortran data in numpy arrays and reads a dictionary back out.

Should be gpu-functional now, I hope?

This should be safe. Probably you want to had some comments in there to explain the upload mechanism.

Fortran comes as CPU
Dycore state is init two a GPU storage thanks to the backend in QuantityFactory
safe_assign_arrayis capable of recognizing that source & destination are not in the same memory space and deals with the upload / download.

In case Fortran memory is on GPU (unlikely) and comes as a cupy array this will take a copy

This upload mechanism is slow and requires it's own optimization via staggered cudaMemCpyAsync. This is not present at the moment

fv3core/pace/fv3core/initialization/geos_wrapper.py

jdahm · 2022-12-01T18:53:14Z

util/pace/util/initialization/sizer.py

+        if "fv_core_nml" in namelist.keys():
+            layout = namelist["fv_core_nml"]["layout"]
+            # npx and npy in the namelist are cell centers, but npz is mid levels
+            nx_tile = namelist["fv_core_nml"]["npx"] - 1
+            ny_tile = namelist["fv_core_nml"]["npy"] - 1
+            nz = namelist["fv_core_nml"]["npz"]
+        elif "nx_tile" in namelist.keys():
+            layout = namelist["layout"]
+            # everything is cell centered in this format
+            nx_tile = namelist["nx_tile"]
+            ny_tile = namelist["nx_tile"]
+            nz = namelist["nz"]
        return cls.from_tile_params(


Warning: these seems to be no else here.

should be ok

util/pace/util/grid/eta.py

…y allocation to __init__

Co-authored-by: Johann Dahm <[email protected]>

…1 vertical levels

…m/pace into feature/fv3core_fortran_api dont edit in browser

FlorianDeconinck · 2022-12-02T14:42:03Z

fv3core/pace/fv3core/initialization/geos_wrapper.py

+        diss_estd: np.ndarray,
+    ) -> dict:
+
+        self._put_fortran_data_in_dycore(


Grab the timer from pace.util and slap one around each of those function call. See fv_dynamics for usage. They are GPU robust already

FlorianDeconinck · 2022-12-02T17:18:11Z

launch jenkins

…place

jdahm · 2022-12-02T19:59:01Z

fv3core/pace/fv3core/initialization/geos_wrapper.py

@@ -114,38 +115,41 @@ def __call__(
        diss_estd: np.ndarray,
    ) -> dict:


Is there a reason for not making this more strongly typed? Could it be Dict[str, np.ndarray]?

small oversight, thought I'd done that

oelbert · 2022-12-05T14:34:13Z

launch jenkins

oelbert added 11 commits November 17, 2022 20:39

revised api, tests. added etas

317cbe9

updated tests, name, wrapper code

a59e12e

linting

2f05dc0

linting again

1ef4fb8

updated sizer to init from current configs, tests pass

87f7f48

Merge branch 'main' into feature/fv3core_fortran_api

30780d7

added docstring to geos wrapper class

ca3a1ec

updated docstring

72dfd12

made output dict persist to prevent reallocation

cb0b84b

updated wrapper call

0529e2f

added typehint for output_dictionary

12444ad

oelbert marked this pull request as ready for review December 1, 2022 16:40

oelbert requested review from elynnwu and jdahm December 1, 2022 16:42

elynnwu reviewed Dec 1, 2022

View reviewed changes

jdahm reviewed Dec 1, 2022

View reviewed changes

fv3core/pace/fv3core/initialization/geos_wrapper.py Show resolved Hide resolved

jdahm reviewed Dec 1, 2022

View reviewed changes

fv3core/pace/fv3core/initialization/geos_wrapper.py Outdated Show resolved Hide resolved

jdahm reviewed Dec 1, 2022

View reviewed changes

util/pace/util/grid/eta.py Outdated Show resolved Hide resolved

oelbert and others added 4 commits December 1, 2022 20:27

changed to save_assign_array for data movement, added output_director…

a9dcab0

…y allocation to __init__

condense ks assignment in eta.py

f5e177b

Co-authored-by: Johann Dahm <[email protected]>

added further handling for sizer.from_namelist and now testing with 9…

f052559

…1 vertical levels

Merge branch 'feature/fv3core_fortran_api' of https://github.com/ai2c…

45ac2f7

…m/pace into feature/fv3core_fortran_api dont edit in browser

FlorianDeconinck reviewed Dec 2, 2022

View reviewed changes

private methods return state/dict instead of modifying attributes in-…

a427243

…place

jdahm reviewed Dec 2, 2022

View reviewed changes

oelbert added 2 commits December 2, 2022 20:25

added api to fv3core readme, fixed assertion in test

154e1b9

stronger typing, changed dictionary names

c28c5b5

jdahm approved these changes Dec 2, 2022

View reviewed changes

oelbert enabled auto-merge (squash) December 2, 2022 21:52

Merge branch 'main' into feature/fv3core_fortran_api

a7b66cf

oelbert disabled auto-merge December 2, 2022 22:18

oelbert enabled auto-merge (squash) December 2, 2022 22:23

elynnwu approved these changes Dec 2, 2022

View reviewed changes

oelbert and others added 3 commits December 6, 2022 21:27

reshaping arrays to match GEOS

dbc3e30

fixing shapes for pressure variables

df94e78

Timer in critical parts

cfb63aa

oelbert disabled auto-merge December 9, 2022 18:48

oelbert merged commit 5dddd21 into main Dec 9, 2022

oelbert deleted the feature/fv3core_fortran_api branch December 9, 2022 18:55

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Feature/fv3core fortran api #395

Feature/fv3core fortran api #395

oelbert commented Nov 29, 2022 •

edited

Loading

elynnwu Dec 1, 2022

elynnwu Dec 1, 2022

oelbert Dec 1, 2022

elynnwu Dec 1, 2022

oelbert Dec 1, 2022

elynnwu Dec 1, 2022

elynnwu Dec 1, 2022

oelbert Dec 1, 2022

oelbert Dec 1, 2022

FlorianDeconinck Dec 2, 2022

jdahm Dec 1, 2022

oelbert Dec 1, 2022

FlorianDeconinck Dec 2, 2022

FlorianDeconinck commented Dec 2, 2022

jdahm Dec 2, 2022

oelbert Dec 2, 2022

oelbert commented Dec 5, 2022


		return self.output_dictionary

		def _put_fortran_data_in_dycore(

		@@ -114,38 +115,41 @@ def __call__(
		diss_estd: np.ndarray,
		) -> dict:

Feature/fv3core fortran api #395

Feature/fv3core fortran api #395

Conversation

oelbert commented Nov 29, 2022 • edited Loading

Purpose

Code changes:

Checklist

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

FlorianDeconinck commented Dec 2, 2022

Choose a reason for hiding this comment

Choose a reason for hiding this comment

oelbert commented Dec 5, 2022

oelbert commented Nov 29, 2022 •

edited

Loading