Merge Pull Request #2905 from E3SM-Project/scream/mahf708/eamxx/pyscream

Automatically Merged using E3SM Pull Request AutoTester
PR Title: pyscream 0.0.2 (misc fixes, reorg)
PR Author: mahf708
PR LABELS: AT: AUTOMERGE, bugfix, python

components/eamxx/src/python/CMakeLists.txt

@@ -1,12 +1 @@
-find_package(pybind11 REQUIRED)
-find_package(mpi4py REQUIRED)
-
-pybind11_add_module(pyeamxx pyeamxx.cpp)
-target_link_libraries(pyeamxx PUBLIC
-  mpi4py
-  scream_share
-  scream_io
-  diagnostics
-  eamxx_physics
-  scream_test_support
-)
+add_subdirectory(libpyscream)

components/eamxx/src/python/libpyscream/CMakeLists.txt

@@ -0,0 +1,12 @@
+find_package(pybind11 REQUIRED)
+find_package(mpi4py REQUIRED)
+
+pybind11_add_module(pyscream_ext pyscream_ext.cpp)
+target_link_libraries(pyscream_ext PUBLIC
+  mpi4py
+  scream_share
+  scream_io
+  diagnostics
+  eamxx_physics
+  scream_test_support
+)

components/eamxx/src/python/pyatmproc.hpp → components/eamxx/src/python/libpyscream/pyatmproc.hpp

@@ -8,7 +8,7 @@
 #include "pygrid.hpp"
 #include "pyfield.hpp"
 #include "pyparamlist.hpp"
-#include "pyeamxx.hpp"
+#include "pyscream_ext.hpp"
 
 #include <ekat/io/ekat_yaml.hpp>
 
@@ -131,6 +131,14 @@ struct PyAtmProc {
     return pybind11::cast(missing);
   }
 
+  pybind11::list dump_fields() {
+    std::vector<std::string> all_fields;
+    for (auto it : fields) {
+      all_fields.push_back(it.first);
+    }
+    return pybind11::cast(all_fields);
+  }
+
   void setup_output (const std::string& yaml_file) {
     auto comm = PySession::get().comm;
 
@@ -175,7 +183,8 @@ inline void pybind_pyatmproc(pybind11::module& m)
     .def("get_params",&PyAtmProc::get_params)
     .def("setup_output",&PyAtmProc::setup_output)
     .def("run",&PyAtmProc::run)
-    .def("read_ic",&PyAtmProc::read_ic);
+    .def("read_ic",&PyAtmProc::read_ic)
+    .def("dump_fields",&PyAtmProc::dump_fields);
 }
 } // namespace scream
 

components/eamxx/src/python/pygrid.hpp → components/eamxx/src/python/libpyscream/pygrid.hpp

@@ -3,7 +3,7 @@
 
 #include "share/grid/mesh_free_grids_manager.hpp"
 
-#include "pyeamxx.hpp"
+#include "pyscream_ext.hpp"
 
 #include <pybind11/pybind11.h>
 
@@ -14,7 +14,7 @@ namespace scream {
 inline void create_grids_manager (int ncols, int nlevs, const std::string& latlon_nc_file)
 {
   EKAT_REQUIRE_MSG (PySession::get().inited,
-      "Error! You did not initialize pyeamxx, or you already finalized it!\n");
+      "Error! You did not initialize pyscream, or you already finalized it!\n");
   auto& comm = PySession::get().comm;
   ekat::ParameterList gm_params;
   std::vector<std::string> grids_names = {"Physics"};
@@ -23,8 +23,10 @@ inline void create_grids_manager (int ncols, int nlevs, const std::string& latlo
   pl.set("number_of_global_columns",ncols);
   pl.set("number_of_vertical_levels",nlevs);
   gm_params.set("grids_names",grids_names);
+  gm_params.set("geo_data_source",std::string("CREATE_EMPTY_DATA"));
 
   if (latlon_nc_file!="") {
+    gm_params.set("geo_data_source",std::string("IC_FILE"));
     gm_params.set("ic_filename",latlon_nc_file);
   }
 

components/eamxx/src/python/pyeamxx.cpp → components/eamxx/src/python/libpyscream/pyscream_ext.cpp

@@ -45,7 +45,7 @@ void finalize () {
   finalize_scream_session();
 }
 
-PYBIND11_MODULE (pyeamxx,m) {
+PYBIND11_MODULE (pyscream_ext,m) {
 
   m.doc() = "Python interfaces to certain EAMxx infrastructure code";
 

components/eamxx/src/python/pyeamxx.hpp → components/eamxx/src/python/libpyscream/pyscream_ext.hpp

@@ -1,5 +1,5 @@
-#ifndef PYEAMXX_HPP
-#define PYEAMXX_HPP
+#ifndef PYSCREAM_HPP
+#define PYSCREAM_HPP
 
 #include "physics/register_physics.hpp"
 #include "diagnostics/register_diagnostics.hpp"
@@ -27,4 +27,4 @@ struct PySession {
 
 } // namespace scream
 
-#endif // PYEAMXX_HPP
+#endif // PYSCREAM_HPP

components/eamxx/src/python/pyproject.toml

@@ -3,15 +3,15 @@ requires = ["setuptools"]
 build-backend = "setuptools.build_meta"
 
 [project]
-name = "pyeamxx"
+name = "pyscream"
 version = "0.0.2"
 dependencies = ["numpy", "mpi4py"]
 
 [tool.setuptools.packages.find]
 where = ["."]
-include = ["pyeamxx", "libpyeamxx"]
+include = ["pyscream", "libpyscream"]
 exclude = ["build_src", "tests"]
 namespaces = true
 
 [tool.setuptools.package-data]
-"libpyeamxx" = ["*.so*"]
+"libpyscream" = ["*.so*"]

components/eamxx/src/python/pyscream/__init__.py

@@ -0,0 +1,20 @@
+"""
+    This file will serve as a way to organize and expose
+    libpyscream internals to the rest of pyscream
+"""
+
+from libpyscream.pyscream_ext import init
+from libpyscream.pyscream_ext import finalize
+from libpyscream.pyscream_ext import Field
+from libpyscream.pyscream_ext import AtmProc
+from libpyscream.pyscream_ext import ParameterList
+from libpyscream.pyscream_ext import create_grids_manager
+
+__all__ = [
+    "init",
+    "finalize",
+    "Field",
+    "AtmProc",
+    "ParameterList",
+    "create_grids_manager",
+]

components/eamxx/src/python/readme

@@ -1,7 +1,7 @@
 INFO:
 - EAMxx python bindings
-- pyeamxx is where we will house the python code
-- libpyeamxx is where we will house the extensions
+- pyscream is where we will house the python code
+- libpyscream is where we will house the extensions
 - packaging moved to https://github.com/mahf708/experimental-scream-feedstock
 
 TODO:
@@ -11,5 +11,5 @@ TODO:
 - decide archs/pythons/mpis to tgt
 
 USER:
-- conda install pyeamxx -c mahf708/label/$mac (mac is chrysalis or pm-cpu)
+- conda install pyscream -c mahf708/label/$mac (mac is chrysalis or pm-cpu)
 - see example in components/eamxx/tests/python/pyp3

components/eamxx/tests/python/pyp3/CMakeLists.txt

@@ -53,6 +53,6 @@ foreach (rank IN LISTS MpiRanks)
     SRC_FILE p3_standalone_cxx.INSTANT.nsteps_x1.np${rank}.2021-10-12-45000.nc
     TGT_FILE p3_standalone_py.INSTANT.nsteps_x1.np${rank}.2021-10-12-45000.nc
     FIXTURES_REQUIRED ${FIXTURES_BASE_NAME}_cxx${suffix} ${FIXTURES_BASE_NAME}_py${suffix}
-    LABELS pyeamxx
+    LABELS pyscream
   )
 endforeach()

components/eamxx/tests/python/pyp3/p3_standalone_py

@@ -5,7 +5,7 @@ import sys
 # Add path to scream libs
 sys.path.append('@SCREAM_BASE_DIR@/scripts')
 
-# Add path to pyeamxx libs
+# Add path to pyscream libs
 sys.path.append('@CMAKE_BINARY_DIR@/src/python')
 
 # Without these, and manual init/finalize, on my laptop I get
@@ -15,7 +15,7 @@ mpi4py.rc.initialize = False  # do not initialize MPI automatically
 mpi4py.rc.finalize = False    # do not finalize MPI automatically
 
 from mpi4py import MPI
-import pyeamxx
+import pyscream
 from pathlib import Path
 
 from utils import ensure_yaml
@@ -38,9 +38,9 @@ def main ():
     # Create the grid
     ncols = 218
     nlevs = 72
-    pyeamxx.create_grids_manager(ncols,nlevs,str(ic_file))
+    pyscream.create_grids_manager(ncols,nlevs,str(ic_file))
 
-    p3 = pyeamxx.AtmProc(yaml_input['atmosphere_processes']['p3'],'p3')
+    p3 = pyscream.AtmProc(yaml_input['atmosphere_processes']['p3'],'p3')
     params = p3.get_params()
     old = params.get_dbl('max_total_ni')
     print (f"max_total_ni: {params.get_dbl('max_total_ni')}")
@@ -64,7 +64,7 @@ if  __name__  == "__main__":
     # This level of indirection ensures all pybind structs are destroyed
     # before we finalize eamxx (and hence kokkos)
     MPI.Init()
-    pyeamxx.init()
+    pyscream.init()
     main ()
-    pyeamxx.finalize()
+    pyscream.finalize()
     MPI.Finalize()