warp-lang 1.0.0b2__py3-none-manylinux2014_x86_64.whl → 1.0.0b6__py3-none-manylinux2014_x86_64.whl

warp/fem/quadrature/quadrature.py

@@ -2,8 +2,9 @@ from typing import Any
 
 import warp as wp
 
-from warp.fem import domain
+from warp.fem import domain, cache
 from warp.fem.types import ElementIndex, Coords
+from warp.fem.space import FunctionSpace
 
 from ..polynomial import Polynomial
 
@@ -11,8 +12,11 @@ from ..polynomial import Polynomial
 class Quadrature:
     """Interface class for quadrature rules"""
 
-    Arg: wp.codegen.Struct
-    """Structure containing arguments to be passed to device functions"""
+    @wp.struct
+    class Arg:
+        """Structure containing arguments to be passed to device functions"""
+
+        pass
 
     def __init__(self, domain: domain.GeometryDomain):
         self._domain = domain
@@ -22,27 +26,46 @@ class Quadrature:
         """Domain over which this quadrature is defined"""
         return self._domain
 
-    def eval_arg_value(self, device) -> wp.codegen.StructInstance:
+    def arg_value(self, device) -> "Arg":
         """
         Value of the argument to be passed to device
         """
-        pass
+        arg = RegularQuadrature.Arg()
+        return arg
 
     def total_point_count(self):
         """Total number of quadrature points over the domain"""
-        pass
+        raise NotImplementedError()
+
+    def points_per_element(self):
+        """Number of points per element if constant, or ``None`` if varying"""
+        return None
 
-    def point_count(arg: Any, element_index: ElementIndex):
+    @staticmethod
+    def point_count(elt_arg: "domain.GeometryDomain.ElementArg", qp_arg: Arg, element_index: ElementIndex):
         """Number of quadrature points for a given element"""
-        pass
+        raise NotImplementedError()
 
-    def point_coords(arg: Any, element_index: ElementIndex, qp_index: int):
-        """Coordinates in element of the qp_index'th quadrature point"""
-        pass
+    @staticmethod
+    def point_coords(
+        elt_arg: "domain.GeometryDomain.ElementArg", qp_arg: Arg, element_index: ElementIndex, qp_index: int
+    ):
+        """Coordinates in element of the element's qp_index'th quadrature point"""
+        raise NotImplementedError()
 
-    def point_weight(arg: Any, element_index: ElementIndex, qp_index: int):
-        """Weight of the qp_index'th quadrature point"""
-        pass
+    @staticmethod
+    def point_weight(
+        elt_arg: "domain.GeometryDomain.ElementArg", qp_arg: Arg, element_index: ElementIndex, qp_index: int
+    ):
+        """Weight of the element's qp_index'th quadrature point"""
+        raise NotImplementedError()
+
+    @staticmethod
+    def point_index(
+        elt_arg: "domain.GeometryDomain.ElementArg", qp_arg: Arg, element_index: ElementIndex, qp_index: int
+    ):
+        """Global index of the element's qp_index'th  quadrature point"""
+        raise NotImplementedError()
 
     def __str__(self) -> str:
         return self.name
@@ -64,31 +87,29 @@ class RegularQuadrature(Quadrature):
 
         self._element_quadrature = domain.reference_element().instantiate_quadrature(order, family)
 
-        N = wp.constant(len(self.points))
+        self._N = wp.constant(len(self.points))
 
-        WeightVec = wp.vec(length=N, dtype=wp.float32)
-        CoordMat = wp.mat(shape=(N, 3), dtype=wp.float32)
+        WeightVec = wp.vec(length=self._N, dtype=wp.float32)
+        CoordMat = wp.mat(shape=(self._N, 3), dtype=wp.float32)
 
-        POINTS = wp.constant(CoordMat(self.points))
-        WEIGHTS = wp.constant(WeightVec(self.weights))
+        self._POINTS = wp.constant(CoordMat(self.points))
+        self._WEIGHTS = wp.constant(WeightVec(self.weights))
 
-        self.point_count = self._make_point_count(N)
-        self.point_index = self._make_point_index(N)
-        self.point_coords = self._make_point_coords(POINTS, self.name)
-        self.point_weight = self._make_point_weight(WEIGHTS, self.name)
+        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()
 
     @property
     def name(self):
-        return (
-            f"{self.__class__.__name__}_{self.domain.reference_element().__class__.__name__}_{self.family}_{self.order}"
-        )
-
-    def __str__(self) -> str:
-        return self.name
+        return f"{self.__class__.__name__}_{self.domain.name}_{self.family}_{self.order}"
 
     def total_point_count(self):
         return len(self.points) * self.domain.geometry_element_count()
 
+    def points_per_element(self):
+        return self._N
+
     @property
     def points(self):
         return self._element_quadrature[0]
@@ -97,46 +118,177 @@ class RegularQuadrature(Quadrature):
     def weights(self):
         return self._element_quadrature[1]
 
-    @wp.struct
-    class Arg:
-        pass
+    def _make_point_count(self):
+        N = self._N
 
-    def arg_value(self, device) -> Arg:
-        arg = RegularQuadrature.Arg()
+        @cache.dynamic_func(suffix=self.name)
+        def point_count(elt_arg: self.domain.ElementArg, qp_arg: self.Arg, element_index: ElementIndex):
+            return N
+
+        return point_count
+
+    def _make_point_coords(self):
+        POINTS = self._POINTS
+
+        @cache.dynamic_func(suffix=self.name)
+        def point_coords(elt_arg: self.domain.ElementArg, qp_arg: self.Arg, element_index: ElementIndex, qp_index: int):
+            return Coords(POINTS[qp_index, 0], POINTS[qp_index, 1], POINTS[qp_index, 2])
+
+        return point_coords
+
+    def _make_point_weight(self):
+        WEIGHTS = self._WEIGHTS
+
+        @cache.dynamic_func(suffix=self.name)
+        def point_weight(elt_arg: self.domain.ElementArg, qp_arg: self.Arg, element_index: ElementIndex, qp_index: int):
+            return WEIGHTS[qp_index]
+
+        return point_weight
+
+    def _make_point_index(self):
+        N = self._N
+
+        @cache.dynamic_func(suffix=self.name)
+        def point_index(elt_arg: self.domain.ElementArg, qp_arg: self.Arg, element_index: ElementIndex, qp_index: int):
+            return N * element_index + qp_index
+
+        return point_index
+
+
+class NodalQuadrature(Quadrature):
+    """Quadrature using space node points as quadrature points
+
+    Note that in contrast to the `nodal=True` flag for :func:`integrate`, this quadrature odes not make any assumption
+    about orthogonality of shape functions, and is thus safe to use for arbitrary integrands.
+    """
+
+    def __init__(self, domain: domain.GeometryDomain, space: FunctionSpace):
+        super().__init__(domain)
+
+        self._space = space
+
+        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()
+
+    @property
+    def name(self):
+        return f"{self.__class__.__name__}_{self._space.name}"
+
+    def total_point_count(self):
+        return self._space.node_count()
+
+    def points_per_element(self):
+        return self._space.topology.NODES_PER_ELEMENT
+
+    def _make_arg(self):
+        @cache.dynamic_struct(suffix=self.name)
+        class Arg:
+            space_arg: self._space.SpaceArg
+            topo_arg: self._space.topology.TopologyArg
+
+        return Arg
+
+    @cache.cached_arg_value
+    def arg_value(self, device):
+        arg = self.Arg()
+        arg.space_arg = self._space.space_arg_value(device)
+        arg.topo_arg = self._space.topology.topo_arg_value(device)
         return arg
 
-    @staticmethod
-    def _make_point_count(N):
-        def point_count(arg: RegularQuadrature.Arg, element_index: ElementIndex):
+    def _make_point_count(self):
+        N = self._space.topology.NODES_PER_ELEMENT
+
+        @cache.dynamic_func(suffix=self.name)
+        def point_count(elt_arg: self.domain.ElementArg, qp_arg: self.Arg, element_index: ElementIndex):
             return N
 
-        from warp.fem.cache import get_func
+        return point_count
 
-        return get_func(point_count, str(N))
+    def _make_point_coords(self):
+        @cache.dynamic_func(suffix=self.name)
+        def point_coords(elt_arg: self.domain.ElementArg, qp_arg: self.Arg, element_index: ElementIndex, qp_index: int):
+            return self._space.node_coords_in_element(elt_arg, qp_arg.space_arg, element_index, qp_index)
 
-    @staticmethod
-    def _make_point_coords(POINTS, name):
-        def point_coords(arg: RegularQuadrature.Arg, element_index: ElementIndex, index: int):
-            return Coords(POINTS[index, 0], POINTS[index, 1], POINTS[index, 2])
+        return point_coords
 
-        from warp.fem.cache import get_func
+    def _make_point_weight(self):
+        @cache.dynamic_func(suffix=self.name)
+        def point_weight(elt_arg: self.domain.ElementArg, qp_arg: self.Arg, element_index: ElementIndex, qp_index: int):
+            return self._space.node_quadrature_weight(elt_arg, qp_arg.space_arg, element_index, qp_index)
 
-        return get_func(point_coords, name)
+        return point_weight
 
-    @staticmethod
-    def _make_point_weight(WEIGHTS, name):
-        def point_weight(arg: RegularQuadrature.Arg, element_index: ElementIndex, index: int):
-            return WEIGHTS[index]
+    def _make_point_index(self):
+        @cache.dynamic_func(suffix=self.name)
+        def point_index(elt_arg: self.domain.ElementArg, qp_arg: self.Arg, element_index: ElementIndex, qp_index: int):
+            return self._space.topology.element_node_index(elt_arg, qp_arg.topo_arg, element_index, qp_index)
 
-        from warp.fem.cache import get_func
+        return point_index
 
-        return get_func(point_weight, name)
 
-    @staticmethod
-    def _make_point_index(N):
-        def point_index(arg: RegularQuadrature.Arg, element_index: ElementIndex, index: int):
-            return N * element_index + index
+class ExplicitQuadrature(Quadrature):
+    """Quadrature using explicit per-cell points and weights. The number of quadrature points per cell is assumed
+    to be constant and deduced from the shape of the points and weights arrays.
+
+    Args:
+        domain: Domain of definition of the quadrature formula
+        points: 2d array of shape ``(domain.geometry_element-count(), points_per_cell)`` containing the coordinates of each quadrature point.
+        weights: 2d array of shape ``(domain.geometry_element-count(), points_per_cell)`` containing the weight for each quadrature point.
+
+    See also: :class:`PicQuadrature`
+    """
+
+    @wp.struct
+    class Arg:
+        points_per_cell: int
+        points: wp.array2d(dtype=Coords)
+        weights: wp.array2d(dtype=float)
+
+    def __init__(self, domain: domain.GeometryDomain, points: "wp.array2d(dtype=Coords)", weights: "wp.array2d(dtype=float)"):
+        super().__init__(domain)
+
+        if points.shape != weights.shape:
+            raise ValueError("Points and weights arrays must have the same shape")
+
+        self._points_per_cell = points.shape[1]
+        self._points = points
+        self._weights = weights
+
+    @property
+    def name(self):
+        return f"{self.__class__.__name__}"
+
+    def total_point_count(self):
+        return self._weights.size
+
+    def points_per_element(self):
+        return self._points_per_cell
+
+    @cache.cached_arg_value
+    def arg_value(self, device):
+        arg = self.Arg()
+        arg.points_per_cell = self._points_per_cell
+        arg.points = self._points.to(device)
+        arg.weights = self._weights.to(device)
+
+        return arg
+
+    @wp.func
+    def point_count(elt_arg: Any, qp_arg: Arg, element_index: ElementIndex):
+        return qp_arg.points_per_cell
+
+    @wp.func
+    def point_coords(elt_arg: Any, qp_arg: Arg, element_index: ElementIndex, qp_index: int):
+        return qp_arg.points[element_index, qp_index]
 
-        from warp.fem.cache import get_func
+    @wp.func
+    def point_weight(elt_arg: Any, qp_arg: Arg, element_index: ElementIndex, qp_index: int):
+        return qp_arg.weights[element_index, qp_index]
 
-        return get_func(point_index, str(N))
+    @wp.func
+    def point_index(elt_arg: Any, qp_arg: Arg, element_index: ElementIndex, qp_index: int):
+        return qp_arg.points_per_cell * element_index + qp_index

warp/fem/space/__init__.py

@@ -1,115 +1,292 @@
 from typing import Optional
+from enum import Enum
 
-import warp.fem.domain
-import warp.fem.geometry
-import warp.fem.polynomial
+import warp.fem.domain as _domain
+import warp.fem.geometry as _geometry
+import warp.fem.polynomial as _polynomial
 
 from .function_space import FunctionSpace
-from .nodal_function_space import NodalFunctionSpace
+from .topology import SpaceTopology
+from .basis_space import BasisSpace, PointBasisSpace
+from .collocated_function_space import CollocatedFunctionSpace
 
 from .grid_2d_function_space import (
-    GridPiecewiseConstantSpace,
-    GridBipolynomialSpace,
-    GridDGBipolynomialSpace,
+    GridPiecewiseConstantBasis,
+    GridBipolynomialBasisSpace,
+    GridDGBipolynomialBasisSpace,
+    GridSerendipityBasisSpace,
+    GridDGSerendipityBasisSpace,
+    GridDGPolynomialBasisSpace,
 )
 from .grid_3d_function_space import (
-    GridTripolynomialSpace,
-    GridDGTripolynomialSpace,
-    Grid3DPiecewiseConstantSpace,
+    GridTripolynomialBasisSpace,
+    GridDGTripolynomialBasisSpace,
+    Grid3DPiecewiseConstantBasis,
+    Grid3DSerendipityBasisSpace,
+    Grid3DDGSerendipityBasisSpace,
+    Grid3DDGPolynomialBasisSpace,
 )
 from .trimesh_2d_function_space import (
-    Trimesh2DPiecewiseConstantSpace,
-    Trimesh2DPolynomialSpace,
-    Trimesh2DDGPolynomialSpace,
+    Trimesh2DPiecewiseConstantBasis,
+    Trimesh2DPolynomialBasisSpace,
+    Trimesh2DDGPolynomialBasisSpace,
+    Trimesh2DNonConformingPolynomialBasisSpace,
+)
+from .tetmesh_function_space import (
+    TetmeshPiecewiseConstantBasis,
+    TetmeshPolynomialBasisSpace,
+    TetmeshDGPolynomialBasisSpace,
+    TetmeshNonConformingPolynomialBasisSpace,
+)
+from .quadmesh_2d_function_space import (
+    Quadmesh2DPiecewiseConstantBasis,
+    Quadmesh2DBipolynomialBasisSpace,
+    Quadmesh2DDGBipolynomialBasisSpace,
+    Quadmesh2DSerendipityBasisSpace,
+    Quadmesh2DDGSerendipityBasisSpace,
+    Quadmesh2DPolynomialBasisSpace,
+)
+from .hexmesh_function_space import (
+    HexmeshPiecewiseConstantBasis,
+    HexmeshTripolynomialBasisSpace,
+    HexmeshDGTripolynomialBasisSpace,
+    HexmeshSerendipityBasisSpace,
+    HexmeshDGSerendipityBasisSpace,
+    HexmeshPolynomialBasisSpace,
 )
-from .tetmesh_function_space import TetmeshPiecewiseConstantSpace, TetmeshPolynomialSpace, TetmeshDGPolynomialSpace
 
 from .partition import SpacePartition, make_space_partition
 from .restriction import SpaceRestriction
 
 
-from .dof_mapper import DofMapper, IdentityMapper, SymmetricTensorMapper
+from .dof_mapper import DofMapper, IdentityMapper, SymmetricTensorMapper, SkewSymmetricTensorMapper
 
 
 def make_space_restriction(
-    space: FunctionSpace,
+    space: Optional[FunctionSpace] = None,
     space_partition: Optional[SpacePartition] = None,
-    domain: Optional[warp.fem.domain.GeometryDomain] = None,
+    domain: Optional[_domain.GeometryDomain] = None,
+    space_topology: Optional[SpaceTopology] = None,
     device=None,
+    temporary_store: "Optional[warp.fem.cache.TemporaryStore]" = None,
 ) -> SpaceRestriction:
     """
-    Restricts a function space to a Domain, i.e. a subset of its elements. 
+    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: the space to be restricted
-        space_partition: if provided, the subset of nodes from ``space`` to consider
+        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 domain is None:
-        if space_partition is None:
-            domain = warp.fem.domain.Cells(geometry=space.geometry)
-        else:
-            domain = warp.fem.domain.Cells(geometry=space_partition.geo_partition)
-    return SpaceRestriction(space=space, space_partition=space_partition, domain=domain, device=device)
 
+    if space_partition is None:
+        if space_topology is None:
+            assert space is not None
+            space_topology = space.topology
 
-def make_polynomial_space(
-    geo: warp.fem.geometry.Geometry,
-    dtype: type = float,
-    dof_mapper: Optional[DofMapper] = None,
+        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
+    )
+
+
+class ElementBasis(Enum):
+    """Choice of basis function to equip individual elements"""
+
+    LAGRANGE = 0
+    """Lagrange basis functions :math:`P_k` for simplices, tensor products :math:`Q_k` for squares and cubes"""
+    SERENDIPITY = 1
+    """Serendipity elements :math:`S_k`, corresponding to Lagrange nodes with interior points removed (for degree <= 3)"""
+    NONCONFORMING_POLYNOMIAL = 2
+    """Simplex Lagrange basis functions :math:`P_{kd}` embedded into non conforming reference elements (e.g. squares or cubes). Discontinuous only."""
+
+
+def make_polynomial_basis_space(
+    geo: _geometry.Geometry,
     degree: int = 1,
+    element_basis: Optional[ElementBasis] = None,
     discontinuous: bool = False,
-    family: Optional[warp.fem.polynomial.Polynomial] = None,
-) -> FunctionSpace:
+    family: Optional[_polynomial.Polynomial] = None,
+) -> BasisSpace:
     """
-    Equip elements of a geometry with a Lagrange polynomial function space
+    Equips a geometry with a polynomial basis.
 
     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
+        the constructed basis space
     """
 
-    if isinstance(geo, warp.fem.geometry.Grid2D):
+    base_geo = geo.base if isinstance(geo, _geometry.DeformedGeometry) else geo
+
+    if element_basis is None:
+        element_basis = ElementBasis.LAGRANGE
+
+    if isinstance(base_geo, _geometry.Grid2D):
         if degree == 0:
-            return GridPiecewiseConstantSpace(geo, dtype=dtype, dof_mapper=dof_mapper)
+            return GridPiecewiseConstantBasis(geo)
+
+        if element_basis == ElementBasis.SERENDIPITY and degree > 1:
+            if discontinuous:
+                return GridDGSerendipityBasisSpace(geo, degree=degree, family=family)
+            else:
+                return GridSerendipityBasisSpace(geo, degree=degree, family=family)
+
+        if element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL:
+            return GridDGPolynomialBasisSpace(geo, degree=degree)
 
         if discontinuous:
-            return GridDGBipolynomialSpace(geo, dtype=dtype, dof_mapper=dof_mapper, degree=degree, family=family)
+            return GridDGBipolynomialBasisSpace(geo, degree=degree, family=family)
         else:
-            return GridBipolynomialSpace(geo, dtype=dtype, dof_mapper=dof_mapper, degree=degree, family=family)
+            return GridBipolynomialBasisSpace(geo, degree=degree, family=family)
 
-    if isinstance(geo, warp.fem.geometry.Grid3D):
+    if isinstance(base_geo, _geometry.Grid3D):
         if degree == 0:
-            return Grid3DPiecewiseConstantSpace(geo, dtype=dtype, dof_mapper=dof_mapper)
+            return Grid3DPiecewiseConstantBasis(geo)
+
+        if element_basis == ElementBasis.SERENDIPITY and degree > 1:
+            if discontinuous:
+                return Grid3DDGSerendipityBasisSpace(geo, degree=degree, family=family)
+            else:
+                return Grid3DSerendipityBasisSpace(geo, degree=degree, family=family)
+
+        if element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL:
+            return Grid3DDGPolynomialBasisSpace(geo, degree=degree)
 
         if discontinuous:
-            return GridDGTripolynomialSpace(geo, dtype=dtype, dof_mapper=dof_mapper, degree=degree, family=family)
+            return GridDGTripolynomialBasisSpace(geo, degree=degree, family=family)
         else:
-            return GridTripolynomialSpace(geo, dtype=dtype, dof_mapper=dof_mapper, degree=degree, family=family)
+            return GridTripolynomialBasisSpace(geo, degree=degree, family=family)
 
-    if isinstance(geo, warp.fem.geometry.Trimesh2D):
+    if isinstance(base_geo, _geometry.Trimesh2D):
         if degree == 0:
-            return Trimesh2DPiecewiseConstantSpace(geo, dtype=dtype, dof_mapper=dof_mapper)
+            return Trimesh2DPiecewiseConstantBasis(geo)
+
+        if element_basis == ElementBasis.SERENDIPITY and degree > 2:
+            raise NotImplementedError("Serendipity variant not implemented yet")
+
+        if element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL:
+            return Trimesh2DNonConformingPolynomialBasisSpace(geo, degree=degree)
 
         if discontinuous:
-            return Trimesh2DDGPolynomialSpace(geo, dtype=dtype, dof_mapper=dof_mapper, degree=degree)
+            return Trimesh2DDGPolynomialBasisSpace(geo, degree=degree)
         else:
-            return Trimesh2DPolynomialSpace(geo, dtype=dtype, dof_mapper=dof_mapper, degree=degree)
+            return Trimesh2DPolynomialBasisSpace(geo, degree=degree)
 
-    if isinstance(geo, warp.fem.geometry.Tetmesh):
+    if isinstance(base_geo, _geometry.Tetmesh):
         if degree == 0:
-            return TetmeshPiecewiseConstantSpace(geo, dtype=dtype, dof_mapper=dof_mapper)
+            return TetmeshPiecewiseConstantBasis(geo)
+
+        if element_basis == ElementBasis.SERENDIPITY and degree > 2:
+            raise NotImplementedError("Serendipity variant not implemented yet")
+
+        if element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL:
+            return TetmeshNonConformingPolynomialBasisSpace(geo, degree=degree)
 
         if discontinuous:
-            return TetmeshDGPolynomialSpace(geo, dtype=dtype, dof_mapper=dof_mapper, degree=degree)
+            return TetmeshDGPolynomialBasisSpace(geo, degree=degree)
         else:
-            return TetmeshPolynomialSpace(geo, dtype=dtype, dof_mapper=dof_mapper, degree=degree)
+            return TetmeshPolynomialBasisSpace(geo, degree=degree)
+
+    if isinstance(base_geo, _geometry.Quadmesh2D):
+        if degree == 0:
+            return Quadmesh2DPiecewiseConstantBasis(geo)
+
+        if element_basis == ElementBasis.SERENDIPITY and degree > 1:
+            if discontinuous:
+                return Quadmesh2DDGSerendipityBasisSpace(geo, degree=degree, family=family)
+            else:
+                return Quadmesh2DSerendipityBasisSpace(geo, degree=degree, family=family)
+
+        if element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL:
+            return Quadmesh2DPolynomialBasisSpace(geo, degree=degree)
+
+        if discontinuous:
+            return Quadmesh2DDGBipolynomialBasisSpace(geo, degree=degree, family=family)
+        else:
+            return Quadmesh2DBipolynomialBasisSpace(geo, degree=degree, family=family)
+
+    if isinstance(base_geo, _geometry.Hexmesh):
+        if degree == 0:
+            return HexmeshPiecewiseConstantBasis(geo)
+
+        if element_basis == ElementBasis.SERENDIPITY and degree > 1:
+            if discontinuous:
+                return HexmeshDGSerendipityBasisSpace(geo, degree=degree, family=family)
+            else:
+                return HexmeshSerendipityBasisSpace(geo, degree=degree, family=family)
+
+        if element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL:
+            return HexmeshPolynomialBasisSpace(geo, degree=degree)
+
+        if discontinuous:
+            return HexmeshDGTripolynomialBasisSpace(geo, degree=degree, family=family)
+        else:
+            return HexmeshTripolynomialBasisSpace(geo, degree=degree, family=family)
+
+    raise NotImplementedError()
+
+
+def make_collocated_function_space(
+    basis_space: BasisSpace, dtype: type = float, dof_mapper: Optional[DofMapper] = None
+) -> CollocatedFunctionSpace:
+    """
+    Constructs a function space from a 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
+    """
+    return CollocatedFunctionSpace(basis_space, dtype=dtype, dof_mapper=dof_mapper)
+
+
+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` and `make_collocated_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
+    """
 
-    raise NotImplementedError
+    basis_space = make_polynomial_basis_space(geo, degree, element_basis, discontinuous, family)
+    return CollocatedFunctionSpace(basis_space, dtype=dtype, dof_mapper=dof_mapper)