Array improvements (nanobind)

python/CMakeLists.txt

@@ -6,7 +6,7 @@ project(dolfinx_nanobind)
 set(CMAKE_CXX_STANDARD 20)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
 set(CMAKE_CXX_EXTENSIONS OFF)
-set(CMAKE_BUILD_TYPE Debug)
+# set(CMAKE_BUILD_TYPE Debug)
 
 # Development component which includes both Development.Module and
 # Development.Embed is not required for building a Python module.  Correct
@@ -52,7 +52,7 @@ find_package(DOLFINX REQUIRED CONFIG)
 nanobind_add_module(
   cpp
   MODULE
-  # NB_SHARED
+  NB_SHARED
   dolfinx/wrappers/dolfinx.cpp
   dolfinx/wrappers/assemble.cpp
   dolfinx/wrappers/common.cpp
@@ -68,11 +68,11 @@ nanobind_add_module(
 )
 
 # Add strict compiler flags
-include(CheckCXXCompilerFlag)
-check_cxx_compiler_flag("-Wall -Werror -pedantic" HAVE_PEDANTIC)
-if(HAVE_PEDANTIC)
-#  target_compile_options(cpp PRIVATE -Wall;-Werror;-pedantic)
-endif()
+# include(CheckCXXCompilerFlag)
+# check_cxx_compiler_flag("-Wall -Werror -pedantic" HAVE_PEDANTIC)
+# if(HAVE_PEDANTIC)
+#   target_compile_options(cpp PRIVATE -Wall;-Werror;-pedantic)
+# endif()
 
 # Add DOLFINx libraries
 target_link_libraries(cpp PRIVATE dolfinx)

python/dolfinx/fem/assemble.py

@@ -326,4 +326,7 @@ def set_bc(b: np.ndarray, bcs: typing.List[DirichletBC],
 
     """
     _bcs = [bc._cpp_object for bc in bcs]
-    _cpp.fem.set_bc(b, _bcs, x0, scale)
+    if x0 is None:
+        _cpp.fem.set_bc(b, _bcs, scale)
+    else:
+        _cpp.fem.set_bc(b, _bcs, x0, scale)

python/dolfinx/fem/petsc.py

@@ -24,7 +24,7 @@
 import ufl
 from dolfinx.cpp.fem import pack_coefficients as _pack_coefficients
 from dolfinx.cpp.fem import pack_constants as _pack_constants
-from dolfinx.fem import assemble
+from dolfinx.fem import assemble as _assemble
 from dolfinx.fem.bcs import DirichletBC
 from dolfinx.fem.bcs import bcs_by_block as _bcs_by_block
 from dolfinx.fem.forms import Form
@@ -191,7 +191,7 @@ def assemble_vector(L: typing.Any, constants=None, coeffs=None) -> PETSc.Vec:
     b = create_petsc_vector(L.function_spaces[0].dofmap.index_map,
                             L.function_spaces[0].dofmap.index_map_bs)
     with b.localForm() as b_local:
-        assemble._assemble_vector_array(b_local.array_w, L, constants, coeffs)
+        _assemble._assemble_vector_array(b_local.array_w, L, constants, coeffs)
     return b
 
 
@@ -213,7 +213,7 @@ def _assemble_vector_vec(b: PETSc.Vec, L: Form, constants=None, coeffs=None) ->
 
     """
     with b.localForm() as b_local:
-        assemble._assemble_vector_array(b_local.array_w, L, constants, coeffs)
+        _assemble._assemble_vector_array(b_local.array_w, L, constants, coeffs)
     return b
 
 
@@ -244,7 +244,7 @@ def _assemble_vector_nest_vec(b: PETSc.Vec, L: typing.List[Form], constants=None
     coeffs = [None] * len(L) if coeffs is None else coeffs
     for b_sub, L_sub, const, coeff in zip(b.getNestSubVecs(), L, constants, coeffs):
         with b_sub.localForm() as b_local:
-            assemble._assemble_vector_array(b_local.array_w, L_sub, const, coeff)
+            _assemble._assemble_vector_array(b_local.array_w, L_sub, const, coeff)
     return b
 
 
@@ -321,7 +321,10 @@ def _assemble_vector_block_vec(b: PETSc.Vec,
     b_array = b.getArray(readonly=False)
     for submap, bc, _x0 in zip(maps, bcs0, x0_sub):
         size = submap[0].size_local * submap[1]
-        _cpp.fem.set_bc(b_array[offset: offset + size], bc, _x0, scale)
+        if _x0 is None:
+            _cpp.fem.set_bc(b_array[offset: offset + size], bc, scale)
+        else:
+            _cpp.fem.set_bc(b_array[offset: offset + size], bc, _x0, scale)
         offset += size
 
     return b
@@ -511,7 +514,7 @@ def apply_lifting(b: PETSc.Vec, a: typing.List[Form],
         x0 = [stack.enter_context(x.localForm()) for x in x0]
         x0_r = [x.array_r for x in x0]
         b_local = stack.enter_context(b.localForm())
-        assemble.apply_lifting(b_local.array_w, a, bcs, x0_r, scale, constants, coeffs)
+        _assemble.apply_lifting(b_local.array_w, a, bcs, x0_r, scale, constants, coeffs)
 
 
 def apply_lifting_nest(b: PETSc.Vec, a: typing.List[typing.List[Form]],
@@ -536,7 +539,7 @@ def set_bc(b: PETSc.Vec, bcs: typing.List[DirichletBC],
     """Apply the function :func:`dolfinx.fem.set_bc` to a PETSc Vector."""
     if x0 is not None:
         x0 = x0.array_r
-    assemble.set_bc(b.array_w, bcs, x0, scale)
+    _assemble.set_bc(b.array_w, bcs, x0, scale)
 
 
 def set_bc_nest(b: PETSc.Vec, bcs: typing.List[typing.List[DirichletBC]],

python/dolfinx/graph.py

@@ -8,9 +8,9 @@
 from __future__ import annotations
 
 import numpy as np
+from dolfinx.cpp.graph import partitioner
 
 from dolfinx import cpp as _cpp
-from dolfinx.cpp.graph import partitioner
 
 # Import graph partitioners, which may or may not be available
 # (dependent on build configuration)
@@ -46,16 +46,12 @@ def adjacencylist(data: np.ndarray, offsets=None):
     """
 
     if offsets is None:
-        if data.dtype == np.int32:
+        try:
             return _cpp.graph.AdjacencyList_int32(data)
-        elif data.dtype == np.int64:
+        except TypeError:
             return _cpp.graph.AdjacencyList_int64(data)
-        else:
-            raise TypeError("Wrong data type")
     else:
-        if data.dtype == np.int32:
+        try:
             return _cpp.graph.AdjacencyList_int32(data, offsets)
-        elif data.dtype == np.int64:
+        except TypeError:
             return _cpp.graph.AdjacencyList_int64(data, offsets)
-        else:
-            raise TypeError("Wrong data type")

python/dolfinx/wrappers/array.h

@@ -6,6 +6,7 @@
 
 #pragma once
 
+#include <array>
 #include <memory>
 #include <nanobind/nanobind.h>
 #include <nanobind/ndarray.h>
@@ -22,44 +23,51 @@ namespace dolfinx_wrappers
 /// From https://github.com/pybind/pybind11/issues/1042
 
 template <typename V, typename U>
-auto as_nbndarray_new(V&& array, U&& shape)
+auto as_nbndarray_new(V&& x, U&& shape)
 {
   using _V = std::decay_t<V>;
   std::size_t dim = shape.size();
-  typename _V::value_type* data = array.data();
-  std::unique_ptr<_V> x_ptr = std::make_unique<_V>(std::move(array));
+  typename _V::value_type* data = x.data();
+  std::unique_ptr<_V> x_ptr = std::make_unique<_V>(std::move(x));
   auto capsule = nb::capsule(x_ptr.get(), [](void* p) noexcept
                              { delete reinterpret_cast<_V*>(p); });
   x_ptr.release();
   return nb::ndarray<nb::numpy, typename _V::value_type>(
       static_cast<typename _V::value_type*>(data), dim, shape.data(), capsule);
 }
 
-template <typename Sequence, typename U>
-auto as_nbarray(Sequence&& seq, U&& shape)
+/// Create a nb::ndarray that shares data with a std::vector. The
+/// std::vector owns the data, and the nb::ndarray object keeps the
+/// std::vector alive.
+// From https://github.com/pybind/pybind11/issues/1042
+template <typename V>
+auto as_nbndarray_new(V&& x)
+{
+  return as_nbndarray_new(std::move(x), std::array{x.size()});
+}
+
+template <typename V, typename U>
+auto as_nbarray(V&& x, U&& shape)
 {
   std::size_t dim = shape.size();
-  auto data = seq.data();
-  std::unique_ptr<Sequence> seq_ptr
-      = std::make_unique<Sequence>(std::move(seq));
-  auto capsule = nb::capsule(
-      seq_ptr.get(), [](void* p) noexcept
-      { std::unique_ptr<Sequence>(reinterpret_cast<Sequence*>(p)); });
-  seq_ptr.release();
+  auto data = x.data();
+  std::unique_ptr<V> x_ptr = std::make_unique<V>(std::move(x));
+  auto capsule = nb::capsule(x_ptr.get(), [](void* p) noexcept
+                             { std::unique_ptr<V>(reinterpret_cast<V*>(p)); });
+  x_ptr.release();
 
-  return nb::ndarray<typename Sequence::value_type, nb::numpy>(
-      static_cast<typename Sequence::value_type*>(data), dim, shape.data(),
-      capsule);
+  return nb::ndarray<typename V::value_type, nb::numpy>(
+      static_cast<typename V::value_type*>(data), dim, shape.data(), capsule);
 }
 
 /// Create a nb::ndarray that shares data with a std::vector. The
 /// std::vector owns the data, and the nb::ndarray object keeps the std::vector
 /// alive.
 // From https://github.com/pybind/pybind11/issues/1042
-template <typename Sequence>
-auto as_nbarray(Sequence&& seq)
+template <typename V>
+auto as_nbarray(V&& x)
 {
-  return as_nbarray(std::move(seq), std::array{seq.size()});
+  return as_nbarray(std::move(x), std::array{x.size()});
 }
 
 } // namespace dolfinx_wrappers

python/dolfinx/wrappers/assemble.cpp

@@ -239,17 +239,18 @@ void declare_assembly_functions(nb::module_& m)
       nb::arg("A"), nb::arg("rows"), nb::arg("diagonal"), "Experimental.");
   m.def(
       "assemble_matrix",
-      [](const std::function<int(
-             const nb::ndarray<const std::int32_t, nb::numpy, nb::c_contig>&,
-             const nb::ndarray<const std::int32_t, nb::numpy, nb::c_contig>&,
-             const nb::ndarray<const T, nb::numpy, nb::c_contig>&)>& fin,
+      [](std::function<int(
+             nb::ndarray<const std::int32_t, nb::numpy, nb::c_contig>,
+             nb::ndarray<const std::int32_t, nb::numpy, nb::c_contig>,
+             nb::ndarray<const T, nb::numpy, nb::c_contig>)>
+             fin,
          const dolfinx::fem::Form<T, U>& form,
          const std::vector<
              std::shared_ptr<const dolfinx::fem::DirichletBC<T, U>>>& bcs)
       {
-        auto f = [&fin](const std::span<const std::int32_t>& rows,
-                        const std::span<const std::int32_t>& cols,
-                        const std::span<const T>& data)
+        auto f = [&fin](std::span<const std::int32_t> rows,
+                        std::span<const std::int32_t> cols,
+                        std::span<const T> data)
         {
           std::size_t rsize = rows.size();
           std::size_t csize = cols.size();
@@ -301,36 +302,29 @@ void declare_assembly_functions(nb::module_& m)
         dolfinx::fem::apply_lifting<T>(std::span<T>(b.data(), b.size()), a,
                                        _constants, _coeffs, bcs1, _x0, scale);
       },
-      nb::arg("b"), nb::arg("a"), nb::arg("constants"), nb::arg("coeffs"),
-      nb::arg("bcs1"), nb::arg("x0"), nb::arg("scale"),
+      nb::arg("b").noconvert(), nb::arg("a"), nb::arg("constants"),
+      nb::arg("coeffs"), nb::arg("bcs1"), nb::arg("x0"), nb::arg("scale"),
       "Modify vector for lifted boundary conditions");
   m.def(
       "set_bc",
       [](nb::ndarray<T, nb::ndim<1>, nb::c_contig> b,
          const std::vector<
              std::shared_ptr<const dolfinx::fem::DirichletBC<T, U>>>& bcs,
-         nb::ndarray<const T, nb::c_contig> x0, U scale)
+         nb::ndarray<const T, nb::ndim<1>, nb::c_contig> x0, T scale)
       {
-        if (x0.ndim() == 0)
-          dolfinx::fem::set_bc<T>(std::span(b.data(), b.size()), bcs, scale);
-        else if (x0.ndim() == 1)
-        {
-          dolfinx::fem::set_bc<T>(std::span(b.data(), b.size()), bcs,
-                                  std::span(x0.data(), x0.shape(0)), scale);
-        }
-        else
-          throw std::runtime_error("Wrong array dimension.");
+        dolfinx::fem::set_bc<T>(std::span(b.data(), b.size()), bcs,
+                                std::span(x0.data(), x0.shape(0)), scale);
       },
-      nb::arg("b"), nb::arg("bcs"), nb::arg("x0"), nb::arg("scale") = T(1));
+      nb::arg("b").noconvert(), nb::arg("bcs"), nb::arg("x0").noconvert(),
+      nb::arg("scale"));
   m.def(
       "set_bc",
       [](nb::ndarray<T, nb::ndim<1>, nb::c_contig> b,
          const std::vector<
              std::shared_ptr<const dolfinx::fem::DirichletBC<T, U>>>& bcs,
-         nb::object x0, T scale)
+         T scale)
       { dolfinx::fem::set_bc<T>(std::span(b.data(), b.size()), bcs, scale); },
-      nb::arg("b"), nb::arg("bcs"), nb::arg("x0") = nb::none(),
-      nb::arg("scale") = T(1));
+      nb::arg("b").noconvert(), nb::arg("bcs"), nb::arg("scale"));
 }
 
 } // namespace

python/dolfinx/wrappers/common.cpp

@@ -131,8 +131,6 @@ void common(nb::module_& m)
           [](const dolfinx::common::IndexMap& self,
              nb::ndarray<const std::int32_t, nb::ndim<1>, nb::c_contig> local)
           {
-            if (local.ndim() != 1)
-              throw std::runtime_error("Array of local indices must be 1D.");
             std::vector<std::int64_t> global(local.size());
             self.local_to_global(
                 std::span(local.data(), local.size()),