File size: 12,106 Bytes
66c9c8a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
from typing import Optional
from enum import Enum

import warp.fem.domain as _domain
import warp.fem.geometry as _geometry
import warp.fem.polynomial as _polynomial

from .function_space import FunctionSpace
from .topology import SpaceTopology
from .basis_space import BasisSpace, PointBasisSpace
from .collocated_function_space import CollocatedFunctionSpace

from .grid_2d_function_space import (
    GridPiecewiseConstantBasis,
    GridBipolynomialBasisSpace,
    GridDGBipolynomialBasisSpace,
    GridSerendipityBasisSpace,
    GridDGSerendipityBasisSpace,
    GridDGPolynomialBasisSpace,
)
from .grid_3d_function_space import (
    GridTripolynomialBasisSpace,
    GridDGTripolynomialBasisSpace,
    Grid3DPiecewiseConstantBasis,
    Grid3DSerendipityBasisSpace,
    Grid3DDGSerendipityBasisSpace,
    Grid3DDGPolynomialBasisSpace,
)
from .trimesh_2d_function_space import (
    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 .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.cache.TemporaryStore]" = None,
) -> 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
    )


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[_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
    """

    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 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 GridDGBipolynomialBasisSpace(geo, degree=degree, family=family)
        else:
            return GridBipolynomialBasisSpace(geo, degree=degree, family=family)

    if isinstance(base_geo, _geometry.Grid3D):
        if degree == 0:
            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 GridDGTripolynomialBasisSpace(geo, degree=degree, family=family)
        else:
            return GridTripolynomialBasisSpace(geo, degree=degree, family=family)

    if isinstance(base_geo, _geometry.Trimesh2D):
        if degree == 0:
            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 Trimesh2DDGPolynomialBasisSpace(geo, degree=degree)
        else:
            return Trimesh2DPolynomialBasisSpace(geo, degree=degree)

    if isinstance(base_geo, _geometry.Tetmesh):
        if degree == 0:
            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 TetmeshDGPolynomialBasisSpace(geo, degree=degree)
        else:
            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
    """

    basis_space = make_polynomial_basis_space(geo, degree, element_basis, discontinuous, family)
    return CollocatedFunctionSpace(basis_space, dtype=dtype, dof_mapper=dof_mapper)