Merge Pull Request #2950 from E3SM-Project/scream/mahf708/nb-pyscream

Automatically Merged using E3SM Pull Request AutoTester
PR Title: move pyscream to nb
PR Author: mahf708
PR LABELS: AT: AUTOMERGE, python

Author: E3SM-Bot

components/eamxx/CMakeLists.txt

@@ -272,8 +272,8 @@ set(SCREAM_BASE_DIR ${CMAKE_CURRENT_SOURCE_DIR})
 set(SCREAM_SRC_DIR  ${CMAKE_CURRENT_SOURCE_DIR}/src)
 set(SCREAM_BIN_DIR  ${CMAKE_CURRENT_BINARY_DIR})
 
-option (EAMXX_ENABLE_PYBIND "Whether to enable python interface to eamxx, via pybind11" OFF)
-if (EAMXX_ENABLE_PYBIND)
+option (EAMXX_ENABLE_PYSCREAM "Whether to enable python interface to eamxx" OFF)
+if (EAMXX_ENABLE_PYSCREAM)
   # Pybind11 requires shared libraries
   set (BUILD_SHARED_LIBS ON)
 endif()

components/eamxx/src/CMakeLists.txt

@@ -8,6 +8,6 @@ if (PROJECT_NAME STREQUAL "E3SM")
   add_subdirectory(mct_coupling)
 endif()
 
-if (EAMXX_ENABLE_PYBIND)
+if (EAMXX_ENABLE_PYSCREAM)
   add_subdirectory(python)
 endif()

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

@@ -1,7 +1,12 @@
-find_package(pybind11 REQUIRED)
+# Detect the installed nanobind package and import it into CMake
+find_package(Python COMPONENTS Interpreter Development REQUIRED)
+execute_process(
+  COMMAND "${PYTHON_EXECUTABLE}" -m nanobind --cmake_dir
+  OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE nanobind_ROOT)
+find_package(nanobind REQUIRED)
 find_package(mpi4py REQUIRED)
 
-pybind11_add_module(pyscream_ext pyscream_ext.cpp)
+nanobind_add_module(pyscream_ext pyscream_ext.cpp)
 target_link_libraries(pyscream_ext PUBLIC
   mpi4py
   scream_share

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

@@ -12,8 +12,11 @@
 
 #include <ekat/io/ekat_yaml.hpp>
 
-#include <pybind11/pybind11.h>
-#include <pybind11/stl.h>
+#include <nanobind/nanobind.h>
+#include <nanobind/stl/list.h>
+#include <nanobind/stl/vector.h>
+
+namespace nb = nanobind;
 
 namespace scream {
 
@@ -26,7 +29,7 @@ struct PyAtmProc {
 
   std::shared_ptr<OutputManager> output_mgr;
 
-  PyAtmProc (const pybind11::dict& d, const std::string& name)
+  PyAtmProc (const nb::dict& d, const std::string& name)
   {
     PyParamList params(d,name);
 
@@ -102,7 +105,7 @@ struct PyAtmProc {
     ap->initialize(t0,RunType::Initial);
   }
 
-  pybind11::list read_ic (const std::string& ic_filename) {
+  std::vector<std::string> read_ic (const std::string& ic_filename) {
     // Get input fields, and read them from file (if present).
     // If field is not in the IC, user is responsible for setting
     // it to an initial value
@@ -128,10 +131,10 @@ struct PyAtmProc {
     }
     scorpio::release_file(ic_filename);
 
-    return pybind11::cast(missing);
+    return missing;
   }
 
-  pybind11::list list_fields(std::string ftype) {
+  std::vector<std::string> list_fields(std::string ftype) {
     std::vector<std::string> fields_list;
     for (const auto& field_pair : fields) {
       const auto& field_identifier = field_pair.second.f.get_header().get_identifier();
@@ -144,18 +147,18 @@ struct PyAtmProc {
         fields_list.push_back(field_pair.first);
       }
     }
-  return pybind11::cast(fields_list);
+  return fields_list;
   }
 
-  pybind11::list list_all_fields() {
+  std::vector<std::string> list_all_fields() {
     return list_fields("all");
   }
 
-  pybind11::list list_required_fields() {
+  std::vector<std::string> list_required_fields() {
     return list_fields("required");
   }
 
-  pybind11::list list_computed_fields() {
+  std::vector<std::string> list_computed_fields() {
     return list_fields("computed");
   }
 
@@ -194,10 +197,10 @@ struct PyAtmProc {
 };
 
 // Register type in the py module
-inline void pybind_pyatmproc(pybind11::module& m)
+inline void nb_pyatmproc(nb::module_& m)
 {
-  pybind11::class_<PyAtmProc>(m,"AtmProc")
-    .def(pybind11::init<const pybind11::dict&,const std::string&>())
+  nb::class_<PyAtmProc>(m,"AtmProc")
+    .def(nb::init<const nb::dict&,const std::string&>())
     .def("get_field",&PyAtmProc::get_field)
     .def("initialize",&PyAtmProc::initialize)
     .def("get_params",&PyAtmProc::get_params)

components/eamxx/src/python/libpyscream/pyfield.hpp

@@ -4,9 +4,11 @@
 #include "share/field/field.hpp"
 #include "share/field/field_utils.hpp"
 
-#include <pybind11/pybind11.h>
-#include <pybind11/numpy.h>
-#include <pybind11/stl.h>
+#include <nanobind/nanobind.h>
+#include <nanobind/ndarray.h>
+#include <nanobind/stl/list.h>
+
+namespace nb = nanobind;
 
 namespace scream {
 
@@ -24,7 +26,8 @@ struct PyField {
     f.allocate_view();
   }
 
-  pybind11::array get () const {
+  template <typename FRAMEWORK>
+  nb::ndarray<FRAMEWORK> get () const {
     const auto& fh  = f.get_header();
     const auto& fid = fh.get_identifier();
 
@@ -39,10 +42,14 @@ struct PyField {
     // NOTE: since the field may be padded, the strides do not necessarily
     //       match the dims. Also, the strides must be grabbed from the
     //       actual view, since the layout doesn't know them.
-    pybind11::array::ShapeContainer shape (fid.get_layout().dims());
+    int shape_t = f.rank();
+    size_t shape[shape_t] = {0};
+    for (int i=0; i<shape_t; ++i) {
+      shape[i] = fid.get_layout().dims()[i];
+    }
     std::vector<ssize_t> strides;
 
-    pybind11::dtype dt;
+    nb::dlpack::dtype dt;
     switch (fid.data_type()) {
       case DataType::IntType:
         dt = get_dt_and_set_strides<int>(strides);
@@ -59,8 +66,8 @@ struct PyField {
 
     // NOTE: you MUST set the parent handle, or else you won't have view semantic
     auto data = f.get_internal_view_data_unsafe<void,Host>();
-    auto this_obj = pybind11::cast(this);
-    return pybind11::array(dt,shape,strides,data,pybind11::handle(this_obj));
+    auto this_obj = nb::cast(this);
+    return nb::ndarray<FRAMEWORK>(data, shape_t, shape, nb::handle(this_obj), strides.data(), dt);
   }
 
   void sync_to_host () {
@@ -75,7 +82,7 @@ struct PyField {
 private:
 
   template<typename T>
-  pybind11::dtype get_dt_and_set_strides (std::vector<ssize_t>& strides) const
+  nb::dlpack::dtype get_dt_and_set_strides (std::vector<ssize_t>& strides) const
   {
     strides.resize(f.rank());
     switch (f.rank()) {
@@ -126,15 +133,15 @@ struct PyField {
             " - field rnak: " + std::to_string(f.rank()) + "\n");
     }
 
-    return pybind11::dtype::of<T>();
+    return nb::dtype<T>();
   }
 };
 
-inline void pybind_pyfield (pybind11::module& m) {
+inline void nb_pyfield (nb::module_& m) {
   // Field class
-  pybind11::class_<PyField>(m,"Field")
-    .def(pybind11::init<>())
-    .def("get",&PyField::get)
+  nb::class_<PyField>(m,"Field")
+    .def(nb::init<>())
+    .def("get",&PyField::get<nb::numpy>)
     .def("sync_to_host",&PyField::sync_to_host)
     .def("sync_to_dev",&PyField::sync_to_dev)
     .def("print",&PyField::print);

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

@@ -5,10 +5,13 @@
 
 #include "pyscream_ext.hpp"
 
-#include <pybind11/pybind11.h>
+#include <nanobind/nanobind.h>
+#include <nanobind/stl/string.h>
 
 #include <mpi.h>
 
+namespace nb = nanobind;
+
 namespace scream {
 
 inline void create_grids_manager (int ncols, int nlevs, const std::string& latlon_nc_file)
@@ -38,9 +41,9 @@ inline void create_grids_manager (int ncols, int nlevs)
   create_grids_manager(ncols,nlevs,"");
 }
 
-inline void pybind_pygrid (pybind11::module& m) {
-  m.def("create_grids_manager",pybind11::overload_cast<int,int>(&create_grids_manager));
-  m.def("create_grids_manager",pybind11::overload_cast<int,int,const std::string&>(&create_grids_manager));
+inline void nb_pygrid (nb::module_& m) {
+  m.def("create_grids_manager",nb::overload_cast<int,int>(&create_grids_manager));
+  m.def("create_grids_manager",nb::overload_cast<int,int,const std::string&>(&create_grids_manager));
 }
 
 } // namespace scream

components/eamxx/src/python/libpyscream/pyparamlist.hpp

@@ -3,12 +3,15 @@
 
 #include <ekat/ekat_parameter_list.hpp>
 
-#include <pybind11/pybind11.h>
-#include <pybind11/numpy.h>
-#include <pybind11/stl.h>
+#include <nanobind/nanobind.h>
+#include <nanobind/stl/list.h>
+#include <nanobind/stl/string.h>
+#include <nanobind/stl/vector.h>
 
 #include <functional>
 
+namespace nb = nanobind;
+
 namespace scream {
 
 struct PyParamList {
@@ -19,11 +22,11 @@ struct PyParamList {
    : pl_ref(src)
   {}
 
-  PyParamList(const pybind11::dict& d)
+  PyParamList(const nb::dict& d)
    : PyParamList(d,"")
   {}
 
-  PyParamList(const pybind11::dict& d, const std::string& name)
+  PyParamList(const nb::dict& d, const std::string& name)
    : pl(name)
    , pl_ref(pl)
   {
@@ -69,70 +72,70 @@ struct PyParamList {
 
 private:
 
-  void parse_dict(const pybind11::dict& d, ekat::ParameterList& p) {
+  void parse_dict(const nb::dict& d, ekat::ParameterList& p) {
     for (auto item : d) {
-      const std::string key = pybind11::str(item.first);
-      if (pybind11::isinstance<pybind11::str>(item.second)) {
-        auto pystr = pybind11::str(item.second);
-        p.set<std::string>(key,pystr.cast<std::string>());
-      } else if (pybind11::isinstance<pybind11::bool_>(item.second)) {
-        auto pyint = pybind11::cast<pybind11::bool_>(item.second);
-        p.set(key,pyint.cast<bool>());
-      } else if (pybind11::isinstance<pybind11::int_>(item.second)) {
-        auto pyint = pybind11::cast<pybind11::int_>(item.second);
-        p.set(key,pyint.cast<int>());
-      } else if (pybind11::isinstance<pybind11::float_>(item.second)) {
-        auto pydouble = pybind11::cast<pybind11::float_>(item.second);
-        p.set(key,pydouble.cast<double>());
-      } else if (pybind11::isinstance<pybind11::list>(item.second)) {
-        auto pylist = pybind11::cast<pybind11::list>(item.second);
+      auto key = nb::cast<std::string>(item.first);
+      if (nb::isinstance<nb::str>(item.second)) {
+        auto pystr = nb::str(item.second);
+        p.set<std::string>(key,nb::cast<std::string>(pystr));
+      } else if (nb::isinstance<nb::bool_>(item.second)) {
+        auto pyint = nb::cast<nb::bool_>(item.second);
+        p.set(key,nb::cast<bool>(pyint));
+      } else if (nb::isinstance<nb::int_>(item.second)) {
+        auto pyint = nb::cast<nb::int_>(item.second);
+        p.set(key,nb::cast<int>(pyint));
+      } else if (nb::isinstance<nb::float_>(item.second)) {
+        auto pydouble = nb::cast<nb::float_>(item.second);
+        p.set(key,nb::cast<double>(pydouble));
+      } else if (nb::isinstance<nb::list>(item.second)) {
+        auto pylist = nb::cast<nb::list>(item.second);
         parse_list(pylist,p,key);
-      } else if (pybind11::isinstance<pybind11::dict>(item.second)) {
-        auto pydict = pybind11::cast<pybind11::dict>(item.second);
+      } else if (nb::isinstance<nb::dict>(item.second)) {
+        auto pydict = nb::cast<nb::dict>(item.second);
         parse_dict(pydict,p.sublist(key));
       } else {
         EKAT_ERROR_MSG ("Unsupported/unrecognized dict entry type.\n");
       }
     }
   }
 
-  void parse_list (const pybind11::list& l, ekat::ParameterList&p, const std::string& key) {
-    EKAT_REQUIRE_MSG (pybind11::len(l)>0,
+  void parse_list (const nb::list& l, ekat::ParameterList&p, const std::string& key) {
+    EKAT_REQUIRE_MSG (nb::len(l)>0,
         "Error! Cannot deduce type for dictionary list entry '" + key + "'\n");
     auto first = l[0];
-    bool are_ints = pybind11::isinstance<pybind11::int_>(first);
-    bool are_floats = pybind11::isinstance<pybind11::float_>(first);
-    bool are_strings = pybind11::isinstance<pybind11::str>(first);
+    bool are_ints = nb::isinstance<nb::int_>(first);
+    bool are_floats = nb::isinstance<nb::float_>(first);
+    bool are_strings = nb::isinstance<nb::str>(first);
     if (are_ints) {
-      parse_list_impl<int,pybind11::int_>(l,p,key);
+      parse_list_impl<int,nb::int_>(l,p,key);
     } else if (are_floats) {
-      parse_list_impl<double,pybind11::float_>(l,p,key);
+      parse_list_impl<double,nb::float_>(l,p,key);
     } else if (are_strings) {
-      parse_list_impl<std::string,pybind11::str>(l,p,key);
+      parse_list_impl<std::string,nb::str>(l,p,key);
     } else {
       EKAT_ERROR_MSG ("Unrecognized/unsupported list entry type.\n");
     }
   }
 
   template<typename Txx, typename Tpy>
-  void parse_list_impl(const pybind11::list& l, ekat::ParameterList& p, const std::string& key) {
+  void parse_list_impl(const nb::list& l, ekat::ParameterList& p, const std::string& key) {
     std::vector<Txx> vals;
     for (auto item : l) {
-      EKAT_REQUIRE_MSG (pybind11::isinstance<Tpy>(item),
+      EKAT_REQUIRE_MSG (nb::isinstance<Tpy>(item),
           "Error! Inconsistent types in list entries.\n");
-      auto item_py = pybind11::cast<Tpy>(item);
-      vals.push_back(item_py.template cast<Txx>());
+      auto item_py = nb::cast<Tpy>(item);
+      vals.push_back(nb::cast<Txx>(item_py));
     }
     p.set(key,vals);
   }
 };
 
-inline void pybind_pyparamlist (pybind11::module& m)
+inline void nb_pyparamlist (nb::module_& m)
 {
   // Param list
-  pybind11::class_<PyParamList>(m,"ParameterList")
-    .def(pybind11::init<const pybind11::dict&>())
-    .def(pybind11::init<const pybind11::dict&,const std::string&>())
+  nb::class_<PyParamList>(m,"ParameterList")
+    .def(nb::init<const nb::dict&>())
+    .def(nb::init<const nb::dict&,const std::string&>())
     .def("sublist",&PyParamList::sublist)
     .def("print",&PyParamList::print)
     .def("set",&PyParamList::set<bool>)