Spaces:
Sleeping
Sleeping
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)
|