Optional output of tensor product computations (#919)

* Change to optional return type of tensor product functionality

* Remove not needed wrappers

* Fix

* Include nanobind/optional.

Author: schnellerhase

cpp/basix/finite-element.cpp

@@ -23,6 +23,8 @@
 #include <concepts>
 #include <limits>
 #include <numeric>
+#include <optional>
+#include <stdexcept>
 
 #define str_macro(X) #X
 #define str(X) str_macro(X)
@@ -328,10 +330,15 @@ basix::create_tp_element(element::family family, cell::type cell, int degree,
                          element::lagrange_variant lvariant,
                          element::dpc_variant dvariant, bool discontinuous)
 {
-  std::vector<int> dof_ordering = tp_dof_ordering(
+  std::optional<std::vector<int>> dof_ordering = tp_dof_ordering(
       family, cell, degree, lvariant, dvariant, discontinuous);
+
+  if (!dof_ordering.has_value())
+    throw std::runtime_error(
+        "Element does not have tensor product factorisation.");
+
   return create_element<T>(family, cell, degree, lvariant, dvariant,
-                           discontinuous, dof_ordering);
+                           discontinuous, *dof_ordering);
 }
 //-----------------------------------------------------------------------------
 template basix::FiniteElement<float>
@@ -342,60 +349,54 @@ basix::create_tp_element(element::family, cell::type, int,
                          element::lagrange_variant, element::dpc_variant, bool);
 //-----------------------------------------------------------------------------
 template <std::floating_point T>
-std::vector<std::vector<FiniteElement<T>>>
+std::optional<std::vector<std::vector<FiniteElement<T>>>>
 basix::tp_factors(element::family family, cell::type cell, int degree,
                   element::lagrange_variant lvariant,
                   element::dpc_variant dvariant, bool discontinuous,
                   const std::vector<int>& dof_ordering)
 {
-  std::vector<int> tp_dofs = tp_dof_ordering(family, cell, degree, lvariant,
-                                             dvariant, discontinuous);
-  if (!tp_dofs.empty() && tp_dofs == dof_ordering)
+  std::optional<std::vector<int>> tp_dofs = tp_dof_ordering(
+      family, cell, degree, lvariant, dvariant, discontinuous);
+  if (!tp_dofs.has_value() || tp_dofs->empty()
+      || tp_dofs.value() != dof_ordering)
+    return std::nullopt;
+
+  switch (family)
   {
-    switch (family)
-    {
-    case element::family::P:
+  case element::family::P:
+  {
+    FiniteElement<T> sub_element
+        = create_element<T>(element::family::P, cell::type::interval, degree,
+                            lvariant, dvariant, true);
+    switch (cell)
     {
-      FiniteElement<T> sub_element
-          = create_element<T>(element::family::P, cell::type::interval, degree,
-                              lvariant, dvariant, true);
-      switch (cell)
-      {
-      case cell::type::quadrilateral:
-      {
-        return {{sub_element, sub_element}};
-      }
-      case cell::type::hexahedron:
-      {
-        return {{sub_element, sub_element, sub_element}};
-      }
-      default:
-      {
-        throw std::runtime_error("Invalid celltype.");
-      }
-      }
-      break;
-    }
+    case cell::type::quadrilateral:
+      return {{{sub_element, sub_element}}};
+    case cell::type::hexahedron:
+      return {{{sub_element, sub_element, sub_element}}};
     default:
-    {
-      throw std::runtime_error("Invalid family.");
-    }
+      return std::nullopt;
     }
+    break;
   }
-  throw std::runtime_error(
-      "Element does not have tensor product factorisation.");
+  default:
+    return std::nullopt;
+  }
+  // C++ 23:
+  // std::unreachable()
+  return std::nullopt;
 }
 //-----------------------------------------------------------------------------
-template std::vector<std::vector<basix::FiniteElement<float>>>
+template std::optional<std::vector<std::vector<basix::FiniteElement<float>>>>
 basix::tp_factors(element::family, cell::type, int, element::lagrange_variant,
                   element::dpc_variant, bool, const std::vector<int>&);
-template std::vector<std::vector<basix::FiniteElement<double>>>
+template std::optional<std::vector<std::vector<basix::FiniteElement<double>>>>
 basix::tp_factors(element::family, cell::type, int, element::lagrange_variant,
                   element::dpc_variant, bool, const std::vector<int>&);
 //-----------------------------------------------------------------------------
-std::vector<int> basix::tp_dof_ordering(element::family family, cell::type cell,
-                                        int degree, element::lagrange_variant,
-                                        element::dpc_variant, bool)
+std::optional<std::vector<int>>
+basix::tp_dof_ordering(element::family family, cell::type cell, int degree,
+                       element::lagrange_variant, element::dpc_variant, bool)
 {
   std::vector<int> dof_ordering;
   std::vector<int> perm;
@@ -488,21 +489,17 @@ std::vector<int> basix::tp_dof_ordering(element::family family, cell::type cell,
       break;
     }
     default:
-    {
-    }
+      return std::nullopt;
     }
     break;
   }
   default:
-  {
-  }
+    return std::nullopt;
   }
 
   if (perm.size() == 0)
-  {
-    throw std::runtime_error(
-        "Element does not have tensor product factorisation.");
-  }
+    return std::nullopt;
+
   dof_ordering.resize(perm.size());
   for (std::size_t i = 0; i < perm.size(); ++i)
     dof_ordering[perm[i]] = i;
@@ -1037,7 +1034,7 @@ FiniteElement<F>::FiniteElement(
       _embedded_superdegree(embedded_superdegree),
       _embedded_subdegree(embedded_subdegree), _value_shape(value_shape),
       _map_type(map_type), _sobolev_space(sobolev_space),
-      _discontinuous(discontinuous), _dof_ordering(dof_ordering)
+      _discontinuous(discontinuous), _dof_ordering(std::move(dof_ordering))
 {
   // Check that discontinuous elements only have DOFs on interior
   if (discontinuous)
@@ -1055,14 +1052,10 @@ FiniteElement<F>::FiniteElement(
     }
   }
 
-  try
-  {
-    _tensor_factors = tp_factors<F>(family, cell_type, degree, lvariant,
-                                    dvariant, discontinuous, dof_ordering);
-  }
-  catch (...)
-  {
-  }
+  auto factors = tp_factors<F>(family, cell_type, degree, lvariant, dvariant,
+                               discontinuous, _dof_ordering);
+  if (factors.has_value())
+    _tensor_factors = factors.value();
 
   std::vector<F> wcoeffs_b(wcoeffs.extent(0) * wcoeffs.extent(1));
   std::copy(wcoeffs.data_handle(), wcoeffs.data_handle() + wcoeffs.size(),

cpp/basix/finite-element.h

@@ -17,6 +17,7 @@
 #include <functional>
 #include <map>
 #include <numeric>
+#include <optional>
 #include <span>
 #include <string>
 #include <tuple>
@@ -1623,11 +1624,12 @@ FiniteElement<T> create_element(element::family family, cell::type cell,
 /// @param[in] discontinuous Indicates whether the element is discontinuous
 /// between cells points of the element. The discontinuous element will have the
 /// same DOFs, but they will all be associated with the interior of the cell.
-/// @return A vector containing the dof ordering
-std::vector<int> tp_dof_ordering(element::family family, cell::type cell,
-                                 int degree, element::lagrange_variant lvariant,
-                                 element::dpc_variant dvariant,
-                                 bool discontinuous);
+/// @return An optional vector containing the dof ordering, if has tensor
+/// structure
+std::optional<std::vector<int>>
+tp_dof_ordering(element::family family, cell::type cell, int degree,
+                element::lagrange_variant lvariant,
+                element::dpc_variant dvariant, bool discontinuous);
 
 /// Get the lexicographic DOF ordering for an element
 /// @param[in] family The element family
@@ -1655,9 +1657,10 @@ std::vector<int> lex_dof_ordering(element::family family, cell::type cell,
 /// between cells points of the element. The discontinuous element will have the
 /// same DOFs, but they will all be associated with the interior of the cell.
 /// @param[in] dof_ordering The ordering of the DOFs
-/// @return A list of lists of finite element factors
+/// @return An optioanl list of lists of finite element factors if family has
+/// tensor structure
 template <std::floating_point T>
-std::vector<std::vector<FiniteElement<T>>>
+std::optional<std::vector<std::vector<FiniteElement<T>>>>
 tp_factors(element::family family, cell::type cell, int degree,
            element::lagrange_variant lvariant, element::dpc_variant dvariant,
            bool discontinuous, const std::vector<int>& dof_ordering);

python/basix/finite_element.py

@@ -778,7 +778,7 @@ def tp_factors(
     discontinuous: bool = False,
     dof_ordering: typing.Optional[list[int]] = None,
     dtype: npt.DTypeLike = np.float64,
-) -> list[list[FiniteElement]]:
+) -> typing.Optional[list[list[FiniteElement]]]:
     """Elements in the tensor product factorisation of an element.
 
     If the element has no factorisation, raises a RuntimeError.
@@ -799,19 +799,21 @@ def tp_factors(
     Returns:
         A list of finite elements.
     """
-    return [
-        [FiniteElement(e) for e in elements]
-        for elements in _tp_factors(
-            family,
-            celltype,
-            degree,
-            lagrange_variant,
-            dpc_variant,
-            discontinuous,
-            dof_ordering if dof_ordering is not None else [],
-            np.dtype(dtype).char,
-        )
-    ]
+    factors = _tp_factors(
+        family,
+        celltype,
+        degree,
+        lagrange_variant,
+        dpc_variant,
+        discontinuous,
+        dof_ordering if dof_ordering is not None else [],
+        np.dtype(dtype).char,
+    )
+
+    if factors is None:
+        return None
+
+    return [[FiniteElement(e) for e in elements] for elements in factors]
 
 
 def tp_dof_ordering(
@@ -821,7 +823,7 @@ def tp_dof_ordering(
     lagrange_variant: LagrangeVariant = LagrangeVariant.unset,
     dpc_variant: DPCVariant = DPCVariant.unset,
     discontinuous: bool = False,
-) -> list[int]:
+) -> typing.Optional[list[int]]:
     """Tensor product DOF ordering for an element.
 
     This DOF ordering can be passed into create_element to create the

python/wrapper.cpp

@@ -17,6 +17,7 @@
 #include <memory>
 #include <nanobind/nanobind.h>
 #include <nanobind/ndarray.h>
+#include <nanobind/stl/optional.h>
 #include <nanobind/stl/pair.h>
 #include <nanobind/stl/string.h>
 #include <nanobind/stl/tuple.h>
@@ -522,8 +523,7 @@ NB_MODULE(_basixcpp, m)
       .value("prism", cell::type::prism)
       .value("pyramid", cell::type::pyramid);
 
-  m.def("cell_volume", [](cell::type cell_type) -> double
-        { return cell::volume<double>(cell_type); });
+  m.def("cell_volume", &cell::volume<double>);
   m.def("cell_facet_normals", [](cell::type cell_type)
         { return as_nbarrayp(cell::facet_normals<double>(cell_type)); });
   m.def(
@@ -641,8 +641,9 @@ NB_MODULE(_basixcpp, m)
            element::lagrange_variant lagrange_variant,
            element::dpc_variant dpc_variant, bool discontinuous,
            const std::vector<int>& dof_ordering, char dtype)
-            -> std::variant<std::vector<std::vector<FiniteElement<float>>>,
-                            std::vector<std::vector<FiniteElement<double>>>>
+            -> std::optional<
+                std::variant<std::vector<std::vector<FiniteElement<float>>>,
+                             std::vector<std::vector<FiniteElement<double>>>>>
         {
           if (dtype == 'd')
           {
@@ -660,40 +661,16 @@ NB_MODULE(_basixcpp, m)
             throw std::runtime_error("Unsupported finite element dtype.");
         });
 
-  m.def("tp_dof_ordering",
-        [](element::family family_name, cell::type cell, int degree,
-           element::lagrange_variant lagrange_variant,
-           element::dpc_variant dpc_variant,
-           bool discontinuous) -> std::vector<int>
-        {
-          return basix::tp_dof_ordering(family_name, cell, degree,
-                                        lagrange_variant, dpc_variant,
-                                        discontinuous);
-        });
-
-  m.def("lex_dof_ordering",
-        [](element::family family_name, cell::type cell, int degree,
-           element::lagrange_variant lagrange_variant,
-           element::dpc_variant dpc_variant,
-           bool discontinuous) -> std::vector<int>
-        {
-          return basix::lex_dof_ordering(family_name, cell, degree,
-                                         lagrange_variant, dpc_variant,
-                                         discontinuous);
-        });
+  m.def("tp_dof_ordering", &basix::tp_dof_ordering);
+  m.def("lex_dof_ordering", &basix::lex_dof_ordering);
 
   nb::enum_<polyset::type>(m, "PolysetType", nb::is_arithmetic(),
                            "Polyset type.")
       .value("standard", polyset::type::standard)
       .value("macroedge", polyset::type::macroedge);
 
-  m.def("superset",
-        [](cell::type cell, polyset::type type1, polyset::type type2)
-        { return polyset::superset(cell, type1, type2); });
-
-  m.def("restriction",
-        [](polyset::type ptype, cell::type cell, cell::type restriction_cell)
-        { return polyset::restriction(ptype, cell, restriction_cell); });
+  m.def("superset", &polyset::superset);
+  m.def("restriction", &polyset::restriction);
 
   m.def(
       "make_quadrature",
@@ -707,15 +684,13 @@ NB_MODULE(_basixcpp, m)
                          as_nbarray(std::move(w)));
       });
 
-  m.def(
-      "gauss_jacobi_rule",
-      [](double a, int m)
-      {
-        auto [pts, w]
-            = quadrature::gauss_jacobi_rule<double>(a, m);
-        return std::pair(as_nbarray(std::move(pts)),
-                         as_nbarray(std::move(w)));
-      });
+  m.def("gauss_jacobi_rule",
+        [](double a, int m)
+        {
+          auto [pts, w] = quadrature::gauss_jacobi_rule<double>(a, m);
+          return std::pair(as_nbarray(std::move(pts)),
+                           as_nbarray(std::move(w)));
+        });
 
   m.def("index", nb::overload_cast<int>(&basix::indexing::idx));
   m.def("index", nb::overload_cast<int, int>(&basix::indexing::idx));