warp-lang 1.9.1__py3-none-win_amd64.whl → 1.10.0rc2__py3-none-win_amd64.whl

warp/{fem → _src/fem}/quadrature/pic_quadrature.py

@@ -16,10 +16,11 @@
 from typing import Any, Optional, Tuple, Union
 
 import warp as wp
-from warp.fem.cache import TemporaryStore, borrow_temporary, cached_arg_value, dynamic_kernel
-from warp.fem.domain import GeometryDomain
-from warp.fem.types import NULL_ELEMENT_INDEX, OUTSIDE, Coords, ElementIndex, make_free_sample
-from warp.fem.utils import compress_node_indices
+from warp._src.fem.cache import TemporaryStore, borrow_temporary, dynamic_kernel
+from warp._src.fem.domain import GeometryDomain
+from warp._src.fem.types import NULL_ELEMENT_INDEX, OUTSIDE, Coords, ElementIndex, make_free_sample
+from warp._src.fem.utils import compress_node_indices
+from warp._src.types import is_array
 
 from .quadrature import Quadrature
 
@@ -63,7 +64,7 @@ class PicQuadrature(Quadrature):
 
     @property
     def name(self):
-        return f"{self.__class__.__name__}"
+        return self.__class__.__name__
 
     @Quadrature.domain.setter
     def domain(self, domain: GeometryDomain):
@@ -84,15 +85,9 @@ class PicQuadrature(Quadrature):
         particle_fraction: wp.array(dtype=float)
         particle_coords: wp.array(dtype=Coords)
 
-    @cached_arg_value
-    def arg_value(self, device) -> Arg:
-        arg = PicQuadrature.Arg()
-        self.fill_arg(arg, device)
-        return arg
-
     def fill_arg(self, args: Arg, device):
-        args.cell_particle_offsets = self._cell_particle_offsets.array.to(device)
-        args.cell_particle_indices = self._cell_particle_indices.array.to(device)
+        args.cell_particle_offsets = self._cell_particle_offsets.to(device)
+        args.cell_particle_indices = self._cell_particle_indices.to(device)
         args.particle_fraction = self._particle_fraction.to(device)
         args.particle_coords = self.particle_coords.to(device)
 
@@ -101,16 +96,16 @@ class PicQuadrature(Quadrature):
 
     def active_cell_count(self):
         """Number of cells containing at least one particle"""
-        return self._cell_count
+        return self._cell_count.numpy()[0]
 
     def max_points_per_element(self):
         if self._max_particles_per_cell is None:
-            max_ppc = wp.zeros(shape=(1,), dtype=int, device=self._cell_particle_offsets.array.device)
+            max_ppc = wp.zeros(shape=(1,), dtype=int, device=self._cell_particle_offsets.device)
             wp.launch(
                 PicQuadrature._max_particles_per_cell_kernel,
-                self._cell_particle_offsets.array.shape[0] - 1,
+                self._cell_particle_offsets.shape[0] - 1,
                 device=max_ppc.device,
-                inputs=[self._cell_particle_offsets.array, max_ppc],
+                inputs=[self._cell_particle_offsets, max_ppc],
             )
             self._max_particles_per_cell = int(max_ppc.numpy()[0])
         return self._max_particles_per_cell
@@ -157,7 +152,7 @@ class PicQuadrature(Quadrature):
             kernel=PicQuadrature._fill_mask_kernel,
             dim=self.domain.geometry_element_count(),
             device=mask.device,
-            inputs=[self._cell_particle_offsets.array, mask],
+            inputs=[self._cell_particle_offsets, mask],
         )
 
     @wp.kernel
@@ -184,7 +179,7 @@ class PicQuadrature(Quadrature):
             cell_fraction[p] = 1.0 / float(cell_particle_count)
 
     def _bin_particles(self, positions, measures, max_dist: float, temporary_store: TemporaryStore):
-        if wp.types.is_array(positions):
+        if is_array(positions):
             device = positions.device
             if not self.domain.supports_lookup(device):
                 raise RuntimeError(
@@ -272,7 +267,7 @@ class PicQuadrature(Quadrature):
                 kernel=PicQuadrature._compute_uniform_fraction,
                 inputs=[
                     cell_index,
-                    self._cell_particle_offsets.array,
+                    self._cell_particle_offsets,
                     self._particle_fraction,
                 ],
                 device=device,

warp/{fem → _src/fem}/quadrature/quadrature.py

@@ -13,14 +13,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from functools import cached_property
 from typing import Any, ClassVar, Optional
 
 import warp as wp
-from warp.fem import cache
-from warp.fem.domain import GeometryDomain
-from warp.fem.geometry import Element
-from warp.fem.space import FunctionSpace
-from warp.fem.types import NULL_ELEMENT_INDEX, Coords, ElementIndex, QuadraturePointIndex
+from warp._src.fem import cache
+from warp._src.fem.domain import GeometryDomain
+from warp._src.fem.geometry import Element
+from warp._src.fem.space.function_space import FunctionSpace
+from warp._src.fem.types import NULL_ELEMENT_INDEX, Coords, ElementIndex, QuadraturePointIndex
 
 from ..polynomial import Polynomial
 
@@ -48,18 +49,22 @@ class Quadrature:
         """Domain over which this quadrature is defined"""
         return self._domain
 
+    @cache.cached_arg_value
     def arg_value(self, device) -> "Arg":
         """
         Value of the argument to be passed to device
         """
-        arg = Quadrature.Arg()
+        arg = self.Arg()
+        self.fill_arg(arg, device)
         return arg
 
     def fill_arg(self, arg: Arg, device):
         """
         Fill the argument with the value of the argument to be passed to device
         """
-        pass
+        if self.arg_value is __class__.arg_value:
+            raise NotImplementedError()
+        arg.assign(self.arg_value(device))
 
     def total_point_count(self):
         """Number of unique quadrature points that can be indexed by this rule.
@@ -160,6 +165,8 @@ class Quadrature:
         ):
             domain_element_index = wp.tid()
             element_index = self.domain.element_index(domain_index_arg, domain_element_index)
+            if element_index == NULL_ELEMENT_INDEX:
+                return
 
             qp_point_count = self.point_count(domain_arg, qp_arg, domain_element_index, element_index)
             for k in range(qp_point_count):
@@ -217,7 +224,9 @@ class _QuadratureWithRegularEvaluationPoints(Quadrature):
         cache.setup_dynamic_attributes(self, cls=__class__)
 
     ElementIndexArg = Quadrature.Arg
-    element_index_arg_value = Quadrature.arg_value
+
+    def element_index_arg_value(self, device):
+        return Quadrature.Arg()
 
     def evaluation_point_count(self):
         return self.domain.element_count() * self._EVALUATION_POINTS_PER_ELEMENT
@@ -267,22 +276,33 @@ class RegularQuadrature(_QuadratureWithRegularEvaluationPoints):
         _cache: ClassVar = {}
 
         def __init__(self, element: Element, order: int, family: Polynomial):
-            self.points, self.weights = element.instantiate_quadrature(order, family)
+            self.points, self.weights = element.prototype.instantiate_quadrature(order, family)
             self.count = wp.constant(len(self.points))
 
         @cache.cached_arg_value
         def arg_value(self, device):
             arg = RegularQuadrature.Arg()
-            self.fill_arg(arg, device)
-            return arg
 
-        def fill_arg(self, arg: "RegularQuadrature.Arg", device):
+            # pause graph capture while we copy from host
+            # we want the cached result to be available outside of the graph
+            if device.is_capturing:
+                graph = wp.context.capture_pause()
+            else:
+                graph = None
+
             arg.points = wp.array(self.points, device=device, dtype=Coords)
             arg.weights = wp.array(self.weights, device=device, dtype=float)
 
+            if graph is not None:
+                wp.context.capture_resume(graph)
+            return arg
+
+        def fill_arg(self, arg: "RegularQuadrature.Arg", device):
+            arg.assign(self.arg_value(device))
+
         @staticmethod
         def get(element: Element, order: int, family: Polynomial):
-            key = (element.__class__.__name__, order, family)
+            key = (element.value, order, family)
             try:
                 return RegularQuadrature.CachedFormula._cache[key]
             except KeyError:
@@ -311,7 +331,7 @@ class RegularQuadrature(_QuadratureWithRegularEvaluationPoints):
 
         cache.setup_dynamic_attributes(self)
 
-    @property
+    @cached_property
     def name(self):
         return f"{self.__class__.__name__}_{self.domain.name}_{self.family}_{self.order}"
 
@@ -329,9 +349,6 @@ class RegularQuadrature(_QuadratureWithRegularEvaluationPoints):
     def weights(self):
         return self._formula.weights
 
-    def arg_value(self, device):
-        return self._formula.arg_value(device)
-
     def fill_arg(self, arg: "RegularQuadrature.Arg", device):
         self._formula.fill_arg(arg, device)
 
@@ -398,20 +415,27 @@ class NodalQuadrature(Quadrature):
     any assumption about orthogonality of shape functions, and is thus safe to use for arbitrary integrands.
     """
 
-    def __init__(self, domain: Optional[GeometryDomain], space: FunctionSpace):
+    _dynamic_attribute_constructors: ClassVar = {
+        "Arg": lambda obj: obj._make_arg(),
+        "point_count": lambda obj: obj._make_point_count(),
+        "point_index": lambda obj: obj._make_point_index(),
+        "point_coords": lambda obj: obj._make_point_coords(),
+        "point_weight": lambda obj: obj._make_point_weight(),
+        "point_evaluation_index": lambda obj: obj._make_point_evaluation_index(),
+    }
+
+    def __init__(
+        self,
+        domain: Optional[GeometryDomain],
+        space: Optional[FunctionSpace],
+    ):
         self._space = space
 
         super().__init__(domain)
 
-        self.Arg = self._make_arg()
-
-        self.point_count = self._make_point_count()
-        self.point_index = self._make_point_index()
-        self.point_coords = self._make_point_coords()
-        self.point_weight = self._make_point_weight()
-        self.point_evaluation_index = self._make_point_evaluation_index()
+        cache.setup_dynamic_attributes(self)
 
-    @property
+    @cached_property
     def name(self):
         return f"{self.__class__.__name__}_{self._space.name}"
 
@@ -429,12 +453,6 @@ class NodalQuadrature(Quadrature):
 
         return Arg
 
-    @cache.cached_arg_value
-    def arg_value(self, device):
-        arg = self.Arg()
-        self.fill_arg(arg, device)
-        return arg
-
     def fill_arg(self, arg: "NodalQuadrature.Arg", device):
         self._space.fill_space_arg(arg.space_arg, device)
         self._space.topology.fill_topo_arg(arg.topo_arg, device)
@@ -486,7 +504,8 @@ class NodalQuadrature(Quadrature):
             element_index: ElementIndex,
             qp_index: int,
         ):
-            return self._space.topology.element_node_index(elt_arg, qp_arg.topo_arg, element_index, qp_index)
+            node_index = self._space.topology.element_node_index(elt_arg, qp_arg.topo_arg, element_index, qp_index)
+            return node_index
 
         return point_index
 
@@ -529,7 +548,12 @@ class ExplicitQuadrature(_QuadratureWithRegularEvaluationPoints):
         points: wp.array2d(dtype=Coords)
         weights: wp.array2d(dtype=float)
 
-    def __init__(self, domain: GeometryDomain, points: "wp.array2d(dtype=Coords)", weights: "wp.array2d(dtype=float)"):
+    def __init__(
+        self,
+        domain: GeometryDomain,
+        points: "wp.array2d(dtype=Coords)",
+        weights: "wp.array2d(dtype=float)",
+    ):
         if points.shape != weights.shape:
             raise ValueError("Points and weights arrays must have the same shape")
 
@@ -554,7 +578,7 @@ class ExplicitQuadrature(_QuadratureWithRegularEvaluationPoints):
         self._points = points
         self._weights = weights
 
-    @property
+    @cached_property
     def name(self):
         return f"{self.__class__.__name__}_{self._whole_geo}_{self._points_per_cell}"
 
@@ -564,52 +588,76 @@ class ExplicitQuadrature(_QuadratureWithRegularEvaluationPoints):
     def max_points_per_element(self):
         return self._points_per_cell
 
-    def arg_value(self, device):
-        arg = self.Arg()
-        self.fill_arg(arg, device)
-        return arg
-
     def fill_arg(self, arg: "ExplicitQuadrature.Arg", device):
         arg.points_per_cell = self._points_per_cell
         arg.points = self._points.to(device)
         arg.weights = self._weights.to(device)
 
     @wp.func
-    def point_count(elt_arg: Any, qp_arg: Arg, domain_element_index: ElementIndex, element_index: ElementIndex):
+    def point_count(
+        elt_arg: Any,
+        qp_arg: Arg,
+        domain_element_index: ElementIndex,
+        element_index: ElementIndex,
+    ):
         return qp_arg.points.shape[1]
 
     @wp.func
     def _point_coords_domain(
-        elt_arg: Any, qp_arg: Arg, domain_element_index: ElementIndex, element_index: ElementIndex, qp_index: int
+        elt_arg: Any,
+        qp_arg: Arg,
+        domain_element_index: ElementIndex,
+        element_index: ElementIndex,
+        qp_index: int,
     ):
         return qp_arg.points[domain_element_index, qp_index]
 
     @wp.func
     def _point_weight_domain(
-        elt_arg: Any, qp_arg: Arg, domain_element_index: ElementIndex, element_index: ElementIndex, qp_index: int
+        elt_arg: Any,
+        qp_arg: Arg,
+        domain_element_index: ElementIndex,
+        element_index: ElementIndex,
+        qp_index: int,
     ):
         return qp_arg.weights[domain_element_index, qp_index]
 
     @wp.func
     def _point_index_domain(
-        elt_arg: Any, qp_arg: Arg, domain_element_index: ElementIndex, element_index: ElementIndex, qp_index: int
+        elt_arg: Any,
+        qp_arg: Arg,
+        domain_element_index: ElementIndex,
+        element_index: ElementIndex,
+        qp_index: int,
     ):
         return qp_arg.points_per_cell * domain_element_index + qp_index
 
     @wp.func
     def _point_coords_geo(
-        elt_arg: Any, qp_arg: Arg, domain_element_index: ElementIndex, element_index: ElementIndex, qp_index: int
+        elt_arg: Any,
+        qp_arg: Arg,
+        domain_element_index: ElementIndex,
+        element_index: ElementIndex,
+        qp_index: int,
     ):
         return qp_arg.points[element_index, qp_index]
 
     @wp.func
     def _point_weight_geo(
-        elt_arg: Any, qp_arg: Arg, domain_element_index: ElementIndex, element_index: ElementIndex, qp_index: int
+        elt_arg: Any,
+        qp_arg: Arg,
+        domain_element_index: ElementIndex,
+        element_index: ElementIndex,
+        qp_index: int,
     ):
         return qp_arg.weights[element_index, qp_index]
 
     @wp.func
     def _point_index_geo(
-        elt_arg: Any, qp_arg: Arg, domain_element_index: ElementIndex, element_index: ElementIndex, qp_index: int
+        elt_arg: Any,
+        qp_arg: Arg,
+        domain_element_index: ElementIndex,
+        element_index: ElementIndex,
+        qp_index: int,
     ):
         return qp_arg.points_per_cell * element_index + qp_index

warp/_src/fem/space/__init__.py

@@ -0,0 +1,248 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# isort: skip_file
+
+from enum import Enum
+from typing import Optional
+
+import warp._src.fem.domain as _domain
+import warp._src.fem.geometry as _geometry
+import warp._src.fem.polynomial as _polynomial
+
+from .function_space import FunctionSpace
+from .basis_function_space import CollocatedFunctionSpace, ContravariantFunctionSpace, CovariantFunctionSpace
+from .topology import SpaceTopology, RegularDiscontinuousSpaceTopology
+from .basis_space import BasisSpace, ShapeBasisSpace
+from .shape import ElementBasis, make_element_shape_function, ShapeFunction
+
+from .grid_2d_function_space import make_grid_2d_space_topology
+
+from .grid_3d_function_space import make_grid_3d_space_topology
+
+from .trimesh_function_space import make_trimesh_space_topology
+
+from .tetmesh_function_space import make_tetmesh_space_topology
+
+from .quadmesh_function_space import make_quadmesh_space_topology
+
+from .hexmesh_function_space import make_hexmesh_space_topology
+
+from .nanogrid_function_space import make_nanogrid_space_topology
+
+
+from .partition import SpacePartition, make_space_partition
+from .restriction import SpaceRestriction
+
+
+from .dof_mapper import DofMapper, IdentityMapper, SymmetricTensorMapper, SkewSymmetricTensorMapper
+
+
+def make_space_restriction(
+    space: Optional[FunctionSpace] = None,
+    space_partition: Optional[SpacePartition] = None,
+    domain: Optional[_domain.GeometryDomain] = None,
+    space_topology: Optional[SpaceTopology] = None,
+    device=None,
+    temporary_store: "Optional[warp.fem.TemporaryStore]" = None,  # noqa: F821
+) -> SpaceRestriction:
+    """
+    Restricts a function space partition to a Domain, i.e. a subset of its elements.
+
+    One of `space_partition`, `space_topology`, or `space` must be provided (and will be considered in that order).
+
+    Args:
+        space: (deprecated) if neither `space_partition` nor `space_topology` are provided, the space defining the topology to restrict
+        space_partition: the subset of nodes from the space topology to consider
+        domain: the domain to restrict the space to, defaults to all cells of the space geometry or partition.
+        space_topology: the space topology to be restricted, if `space_partition` is ``None``.
+        device: device on which to perform and store computations
+        temporary_store: shared pool from which to allocate temporary arrays
+    """
+
+    if space_partition is None:
+        if space_topology is None:
+            assert space is not None
+            space_topology = space.topology
+
+        if domain is None:
+            domain = _domain.Cells(geometry=space_topology.geometry)
+
+        space_partition = make_space_partition(
+            space_topology=space_topology, geometry_partition=domain.geometry_partition
+        )
+    elif domain is None:
+        domain = _domain.Cells(geometry=space_partition.geo_partition)
+
+    return SpaceRestriction(
+        space_partition=space_partition, domain=domain, device=device, temporary_store=temporary_store
+    )
+
+
+def make_polynomial_basis_space(
+    geo: _geometry.Geometry,
+    degree: int = 1,
+    element_basis: Optional[ElementBasis] = None,
+    discontinuous: bool = False,
+    family: Optional[_polynomial.Polynomial] = None,
+) -> BasisSpace:
+    """
+    Equips a geometry with a polynomial basis.
+
+    Args:
+        geo: the Geometry on which to build the space
+        degree: polynomial degree of the per-element shape functions
+        discontinuous: if True, use Discontinuous Galerkin shape functions. Discontinuous is implied if degree is 0, i.e, piecewise-constant shape functions.
+        element_basis: type of basis function for the individual elements
+        family: Polynomial family used to generate the shape function basis. If not provided, a reasonable basis is chosen.
+
+    Returns:
+        the constructed basis space
+    """
+
+    shape = make_element_shape_function(geo.reference_cell(), degree=degree, element_basis=element_basis, family=family)
+
+    discontinuous = discontinuous or element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL
+    topology = make_element_based_space_topology(geo, shape, discontinuous)
+
+    return ShapeBasisSpace(topology, shape)
+
+
+def make_element_based_space_topology(
+    geo: _geometry.Geometry,
+    shape: ShapeFunction,
+    discontinuous: bool = False,
+) -> SpaceTopology:
+    """
+    Makes a space topology from a geometry and an element-based shape function.
+
+    Args:
+        geo: The geometry to make the topology for
+        shape: The shape function to make the topology for
+        discontinuous: Whether to make a discontinuous topology
+
+    Returns:
+        The element-based space topology
+
+    Raises:
+        NotImplementedError: If the geometry type is not supported
+        ValueError: If the shape function is not supported for the given geometry`
+    """
+
+    topology = None
+    base_geo = geo.base
+
+    if discontinuous or shape.ORDER == 0:
+        topology = RegularDiscontinuousSpaceTopology(geo, shape.NODES_PER_ELEMENT)
+    elif isinstance(base_geo, _geometry.Grid2D):
+        topology = make_grid_2d_space_topology(geo, shape)
+    elif isinstance(base_geo, _geometry.Grid3D):
+        topology = make_grid_3d_space_topology(geo, shape)
+    elif isinstance(base_geo, _geometry.Trimesh):
+        topology = make_trimesh_space_topology(geo, shape)
+    elif isinstance(base_geo, _geometry.Tetmesh):
+        topology = make_tetmesh_space_topology(geo, shape)
+    elif isinstance(base_geo, _geometry.Quadmesh):
+        topology = make_quadmesh_space_topology(geo, shape)
+    elif isinstance(base_geo, _geometry.Hexmesh):
+        topology = make_hexmesh_space_topology(geo, shape)
+    elif isinstance(base_geo, _geometry.Nanogrid) or isinstance(base_geo, _geometry.AdaptiveNanogrid):
+        topology = make_nanogrid_space_topology(geo, shape)
+
+    if topology is None:
+        raise NotImplementedError(f"Unsupported geometry type {geo.name}")
+
+    return topology
+
+
+def make_collocated_function_space(
+    basis_space: BasisSpace, dtype: type = float, dof_mapper: Optional[DofMapper] = None
+) -> CollocatedFunctionSpace:
+    """
+    Constructs a function space from a scalar-valued basis space and a value type, such that all degrees of freedom of the value type are stored at each of the basis nodes.
+
+    Args:
+        geo: the Geometry on which to build the space
+        dtype: value type the function space. If ``dof_mapper`` is provided, the value type from the DofMapper will be used instead.
+        dof_mapper: mapping from node degrees of freedom to function values, defaults to Identity. Useful for reduced coordinates, e.g. :py:class:`SymmetricTensorMapper` maps 2x2 (resp 3x3) symmetric tensors to 3 (resp 6) degrees of freedom.
+
+    Returns:
+        the constructed function space
+    """
+
+    if basis_space.value != ShapeFunction.Value.Scalar:
+        raise ValueError("Collocated function spaces may only be constructed from scalar-valued basis")
+
+    return CollocatedFunctionSpace(basis_space, dtype=dtype, dof_mapper=dof_mapper)
+
+
+def make_covariant_function_space(
+    basis_space: BasisSpace,
+) -> CovariantFunctionSpace:
+    """
+    Constructs a covariant function space from a vector-valued basis space
+    """
+
+    if basis_space.value != ShapeFunction.Value.CovariantVector:
+        raise ValueError("Covariant function spaces may only be constructed from covariant vector-valued basis")
+    return CovariantFunctionSpace(basis_space)
+
+
+def make_contravariant_function_space(
+    basis_space: BasisSpace,
+) -> ContravariantFunctionSpace:
+    """
+    Constructs a contravariant function space from a vector-valued basis space
+    """
+
+    if basis_space.value != ShapeFunction.Value.ContravariantVector:
+        raise ValueError("Contravariant function spaces may only be constructed from contravariant vector-valued basis")
+    return ContravariantFunctionSpace(basis_space)
+
+
+def make_polynomial_space(
+    geo: _geometry.Geometry,
+    dtype: type = float,
+    dof_mapper: Optional[DofMapper] = None,
+    degree: int = 1,
+    element_basis: Optional[ElementBasis] = None,
+    discontinuous: bool = False,
+    family: Optional[_polynomial.Polynomial] = None,
+) -> CollocatedFunctionSpace:
+    """
+    Equips a geometry with a collocated, polynomial function space.
+    Equivalent to successive calls to :func:`make_polynomial_basis_space` then `make_collocated_function_space`, `make_covariant_function_space` or `make_contravariant_function_space`.
+
+    Args:
+        geo: the Geometry on which to build the space
+        dtype: value type the function space. If ``dof_mapper`` is provided, the value type from the DofMapper will be used instead.
+        dof_mapper: mapping from node degrees of freedom to function values, defaults to Identity. Useful for reduced coordinates, e.g. :py:class:`SymmetricTensorMapper` maps 2x2 (resp 3x3) symmetric tensors to 3 (resp 6) degrees of freedom.
+        degree: polynomial degree of the per-element shape functions
+        discontinuous: if True, use Discontinuous Galerkin shape functions. Discontinuous is implied if degree is 0, i.e, piecewise-constant shape functions.
+        element_basis: type of basis function for the individual elements
+        family: Polynomial family used to generate the shape function basis. If not provided, a reasonable basis is chosen.
+
+    Returns:
+        the constructed function space
+    """
+
+    basis_space = make_polynomial_basis_space(geo, degree, element_basis, discontinuous, family)
+
+    if basis_space.value == ShapeFunction.Value.CovariantVector:
+        return make_covariant_function_space(basis_space)
+    if basis_space.value == ShapeFunction.Value.ContravariantVector:
+        return make_contravariant_function_space(basis_space)
+
+    return make_collocated_function_space(basis_space, dtype=dtype, dof_mapper=dof_mapper)