warp-lang 1.9.0__py3-none-win_amd64.whl → 1.10.0rc2__py3-none-win_amd64.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Potentially problematic release.
This version of warp-lang might be problematic. Click here for more details.
- warp/__init__.py +301 -287
- warp/__init__.pyi +2220 -313
- warp/_src/__init__.py +14 -0
- warp/_src/autograd.py +1075 -0
- warp/_src/build.py +618 -0
- warp/_src/build_dll.py +640 -0
- warp/{builtins.py → _src/builtins.py} +1497 -226
- warp/_src/codegen.py +4359 -0
- warp/{config.py → _src/config.py} +178 -169
- warp/_src/constants.py +57 -0
- warp/_src/context.py +8294 -0
- warp/_src/dlpack.py +462 -0
- warp/_src/fabric.py +355 -0
- warp/_src/fem/__init__.py +14 -0
- warp/_src/fem/adaptivity.py +508 -0
- warp/_src/fem/cache.py +687 -0
- warp/_src/fem/dirichlet.py +188 -0
- warp/{fem → _src/fem}/domain.py +40 -30
- warp/_src/fem/field/__init__.py +131 -0
- warp/_src/fem/field/field.py +701 -0
- warp/{fem → _src/fem}/field/nodal_field.py +30 -15
- warp/{fem → _src/fem}/field/restriction.py +1 -1
- warp/{fem → _src/fem}/field/virtual.py +53 -27
- warp/_src/fem/geometry/__init__.py +32 -0
- warp/{fem → _src/fem}/geometry/adaptive_nanogrid.py +77 -163
- warp/_src/fem/geometry/closest_point.py +97 -0
- warp/{fem → _src/fem}/geometry/deformed_geometry.py +14 -22
- warp/{fem → _src/fem}/geometry/element.py +32 -10
- warp/{fem → _src/fem}/geometry/geometry.py +48 -20
- warp/{fem → _src/fem}/geometry/grid_2d.py +12 -23
- warp/{fem → _src/fem}/geometry/grid_3d.py +12 -23
- warp/{fem → _src/fem}/geometry/hexmesh.py +40 -63
- warp/{fem → _src/fem}/geometry/nanogrid.py +255 -248
- warp/{fem → _src/fem}/geometry/partition.py +121 -63
- warp/{fem → _src/fem}/geometry/quadmesh.py +26 -45
- warp/{fem → _src/fem}/geometry/tetmesh.py +40 -63
- warp/{fem → _src/fem}/geometry/trimesh.py +26 -45
- warp/{fem → _src/fem}/integrate.py +164 -158
- warp/_src/fem/linalg.py +383 -0
- warp/_src/fem/operator.py +396 -0
- warp/_src/fem/polynomial.py +229 -0
- warp/{fem → _src/fem}/quadrature/pic_quadrature.py +15 -20
- warp/{fem → _src/fem}/quadrature/quadrature.py +95 -47
- warp/_src/fem/space/__init__.py +248 -0
- warp/{fem → _src/fem}/space/basis_function_space.py +20 -11
- warp/_src/fem/space/basis_space.py +679 -0
- warp/{fem → _src/fem}/space/dof_mapper.py +3 -3
- warp/{fem → _src/fem}/space/function_space.py +14 -13
- warp/{fem → _src/fem}/space/grid_2d_function_space.py +4 -7
- warp/{fem → _src/fem}/space/grid_3d_function_space.py +4 -4
- warp/{fem → _src/fem}/space/hexmesh_function_space.py +4 -10
- warp/{fem → _src/fem}/space/nanogrid_function_space.py +3 -9
- warp/{fem → _src/fem}/space/partition.py +117 -60
- warp/{fem → _src/fem}/space/quadmesh_function_space.py +4 -10
- warp/{fem → _src/fem}/space/restriction.py +66 -33
- warp/_src/fem/space/shape/__init__.py +152 -0
- warp/{fem → _src/fem}/space/shape/cube_shape_function.py +9 -9
- warp/{fem → _src/fem}/space/shape/shape_function.py +8 -9
- warp/{fem → _src/fem}/space/shape/square_shape_function.py +6 -6
- warp/{fem → _src/fem}/space/shape/tet_shape_function.py +3 -3
- warp/{fem → _src/fem}/space/shape/triangle_shape_function.py +3 -3
- warp/{fem → _src/fem}/space/tetmesh_function_space.py +3 -9
- warp/_src/fem/space/topology.py +459 -0
- warp/{fem → _src/fem}/space/trimesh_function_space.py +3 -9
- warp/_src/fem/types.py +112 -0
- warp/_src/fem/utils.py +486 -0
- warp/_src/jax.py +186 -0
- warp/_src/jax_experimental/__init__.py +14 -0
- warp/_src/jax_experimental/custom_call.py +387 -0
- warp/_src/jax_experimental/ffi.py +1284 -0
- warp/_src/jax_experimental/xla_ffi.py +656 -0
- warp/_src/marching_cubes.py +708 -0
- warp/_src/math.py +414 -0
- warp/_src/optim/__init__.py +14 -0
- warp/_src/optim/adam.py +163 -0
- warp/_src/optim/linear.py +1606 -0
- warp/_src/optim/sgd.py +112 -0
- warp/_src/paddle.py +406 -0
- warp/_src/render/__init__.py +14 -0
- warp/_src/render/imgui_manager.py +289 -0
- warp/_src/render/render_opengl.py +3636 -0
- warp/_src/render/render_usd.py +937 -0
- warp/_src/render/utils.py +160 -0
- warp/_src/sparse.py +2716 -0
- warp/_src/tape.py +1206 -0
- warp/{thirdparty → _src/thirdparty}/unittest_parallel.py +9 -2
- warp/_src/torch.py +391 -0
- warp/_src/types.py +5870 -0
- warp/_src/utils.py +1693 -0
- warp/autograd.py +12 -1054
- warp/bin/warp-clang.dll +0 -0
- warp/bin/warp.dll +0 -0
- warp/build.py +8 -588
- warp/build_dll.py +6 -471
- warp/codegen.py +6 -4246
- warp/constants.py +6 -39
- warp/context.py +12 -7851
- warp/dlpack.py +6 -444
- warp/examples/distributed/example_jacobi_mpi.py +4 -5
- warp/examples/fem/example_adaptive_grid.py +1 -1
- warp/examples/fem/example_apic_fluid.py +1 -1
- warp/examples/fem/example_burgers.py +8 -8
- warp/examples/fem/example_diffusion.py +1 -1
- warp/examples/fem/example_distortion_energy.py +1 -1
- warp/examples/fem/example_mixed_elasticity.py +2 -2
- warp/examples/fem/example_navier_stokes.py +1 -1
- warp/examples/fem/example_nonconforming_contact.py +7 -7
- warp/examples/fem/example_stokes.py +1 -1
- warp/examples/fem/example_stokes_transfer.py +1 -1
- warp/examples/fem/utils.py +2 -2
- warp/examples/interop/example_jax_callable.py +1 -1
- warp/examples/interop/example_jax_ffi_callback.py +1 -1
- warp/examples/interop/example_jax_kernel.py +3 -2
- warp/examples/tile/example_tile_mcgp.py +191 -0
- warp/fabric.py +6 -337
- warp/fem/__init__.py +159 -97
- warp/fem/adaptivity.py +7 -489
- warp/fem/cache.py +9 -648
- warp/fem/dirichlet.py +6 -184
- warp/fem/field/__init__.py +8 -109
- warp/fem/field/field.py +7 -652
- warp/fem/geometry/__init__.py +7 -18
- warp/fem/geometry/closest_point.py +11 -77
- warp/fem/linalg.py +18 -366
- warp/fem/operator.py +11 -369
- warp/fem/polynomial.py +9 -209
- warp/fem/space/__init__.py +5 -211
- warp/fem/space/basis_space.py +6 -662
- warp/fem/space/shape/__init__.py +41 -118
- warp/fem/space/topology.py +6 -437
- warp/fem/types.py +6 -81
- warp/fem/utils.py +11 -444
- warp/jax.py +8 -165
- warp/jax_experimental/__init__.py +14 -1
- warp/jax_experimental/custom_call.py +8 -342
- warp/jax_experimental/ffi.py +17 -853
- warp/jax_experimental/xla_ffi.py +5 -596
- warp/marching_cubes.py +5 -689
- warp/math.py +16 -393
- warp/native/array.h +385 -37
- warp/native/builtin.h +316 -39
- warp/native/bvh.cpp +43 -9
- warp/native/bvh.cu +62 -27
- warp/native/bvh.h +310 -309
- warp/native/clang/clang.cpp +102 -97
- warp/native/coloring.cpp +0 -1
- warp/native/crt.h +208 -0
- warp/native/exports.h +156 -0
- warp/native/hashgrid.cu +2 -0
- warp/native/intersect.h +24 -1
- warp/native/intersect_tri.h +44 -35
- warp/native/mat.h +1456 -276
- warp/native/mesh.cpp +4 -4
- warp/native/mesh.cu +4 -2
- warp/native/mesh.h +176 -61
- warp/native/quat.h +0 -52
- warp/native/scan.cu +2 -0
- warp/native/sort.cu +22 -13
- warp/native/sort.h +2 -0
- warp/native/sparse.cu +7 -3
- warp/native/spatial.h +12 -0
- warp/native/tile.h +837 -70
- warp/native/tile_radix_sort.h +1 -1
- warp/native/tile_reduce.h +394 -46
- warp/native/tile_scan.h +4 -4
- warp/native/vec.h +469 -53
- warp/native/version.h +23 -0
- warp/native/volume.cpp +1 -1
- warp/native/volume.cu +1 -0
- warp/native/volume.h +1 -1
- warp/native/volume_builder.cu +2 -0
- warp/native/warp.cpp +60 -32
- warp/native/warp.cu +313 -201
- warp/native/warp.h +14 -11
- warp/optim/__init__.py +6 -3
- warp/optim/adam.py +6 -145
- warp/optim/linear.py +14 -1585
- warp/optim/sgd.py +6 -94
- warp/paddle.py +6 -388
- warp/render/__init__.py +8 -4
- warp/render/imgui_manager.py +7 -267
- warp/render/render_opengl.py +6 -3616
- warp/render/render_usd.py +6 -918
- warp/render/utils.py +6 -142
- warp/sparse.py +37 -2563
- warp/tape.py +6 -1188
- warp/tests/__main__.py +1 -1
- warp/tests/cuda/test_async.py +4 -4
- warp/tests/cuda/test_conditional_captures.py +1 -1
- warp/tests/cuda/test_multigpu.py +1 -1
- warp/tests/cuda/test_streams.py +58 -1
- warp/tests/geometry/test_bvh.py +157 -22
- warp/tests/geometry/test_hash_grid.py +38 -0
- warp/tests/geometry/test_marching_cubes.py +0 -1
- warp/tests/geometry/test_mesh.py +5 -3
- warp/tests/geometry/test_mesh_query_aabb.py +5 -12
- warp/tests/geometry/test_mesh_query_point.py +5 -2
- warp/tests/geometry/test_mesh_query_ray.py +15 -3
- warp/tests/geometry/test_volume_write.py +5 -5
- warp/tests/interop/test_dlpack.py +14 -14
- warp/tests/interop/test_jax.py +1382 -79
- warp/tests/interop/test_paddle.py +1 -1
- warp/tests/test_adam.py +0 -1
- warp/tests/test_arithmetic.py +9 -9
- warp/tests/test_array.py +529 -100
- warp/tests/test_array_reduce.py +3 -3
- warp/tests/test_atomic.py +12 -8
- warp/tests/test_atomic_bitwise.py +209 -0
- warp/tests/test_atomic_cas.py +4 -4
- warp/tests/test_bool.py +2 -2
- warp/tests/test_builtins_resolution.py +5 -571
- warp/tests/test_codegen.py +34 -15
- warp/tests/test_conditional.py +1 -1
- warp/tests/test_context.py +6 -6
- warp/tests/test_copy.py +242 -161
- warp/tests/test_ctypes.py +3 -3
- warp/tests/test_devices.py +24 -2
- warp/tests/test_examples.py +16 -84
- warp/tests/test_fabricarray.py +35 -35
- warp/tests/test_fast_math.py +0 -2
- warp/tests/test_fem.py +60 -14
- warp/tests/test_fixedarray.py +3 -3
- warp/tests/test_func.py +8 -5
- warp/tests/test_generics.py +1 -1
- warp/tests/test_indexedarray.py +24 -24
- warp/tests/test_intersect.py +39 -9
- warp/tests/test_large.py +1 -1
- warp/tests/test_lerp.py +3 -1
- warp/tests/test_linear_solvers.py +1 -1
- warp/tests/test_map.py +49 -4
- warp/tests/test_mat.py +52 -62
- warp/tests/test_mat_constructors.py +4 -5
- warp/tests/test_mat_lite.py +1 -1
- warp/tests/test_mat_scalar_ops.py +121 -121
- warp/tests/test_math.py +34 -0
- warp/tests/test_module_aot.py +4 -4
- warp/tests/test_modules_lite.py +28 -2
- warp/tests/test_print.py +11 -11
- warp/tests/test_quat.py +93 -58
- warp/tests/test_runlength_encode.py +1 -1
- warp/tests/test_scalar_ops.py +38 -10
- warp/tests/test_smoothstep.py +1 -1
- warp/tests/test_sparse.py +126 -15
- warp/tests/test_spatial.py +105 -87
- warp/tests/test_special_values.py +6 -6
- warp/tests/test_static.py +7 -7
- warp/tests/test_struct.py +13 -2
- warp/tests/test_triangle_closest_point.py +48 -1
- warp/tests/test_tuple.py +96 -0
- warp/tests/test_types.py +82 -9
- warp/tests/test_utils.py +52 -52
- warp/tests/test_vec.py +29 -29
- warp/tests/test_vec_constructors.py +5 -5
- warp/tests/test_vec_scalar_ops.py +97 -97
- warp/tests/test_version.py +75 -0
- warp/tests/tile/test_tile.py +239 -0
- warp/tests/tile/test_tile_atomic_bitwise.py +403 -0
- warp/tests/tile/test_tile_cholesky.py +7 -4
- warp/tests/tile/test_tile_load.py +26 -2
- warp/tests/tile/test_tile_mathdx.py +3 -3
- warp/tests/tile/test_tile_matmul.py +1 -1
- warp/tests/tile/test_tile_mlp.py +2 -4
- warp/tests/tile/test_tile_reduce.py +214 -13
- warp/tests/unittest_suites.py +6 -14
- warp/tests/unittest_utils.py +10 -9
- warp/tests/walkthrough_debug.py +3 -1
- warp/torch.py +6 -373
- warp/types.py +29 -5750
- warp/utils.py +10 -1659
- {warp_lang-1.9.0.dist-info → warp_lang-1.10.0rc2.dist-info}/METADATA +47 -103
- warp_lang-1.10.0rc2.dist-info/RECORD +468 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/Gaia-LICENSE.txt +6 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/appdirs-LICENSE.txt +22 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/asset_pixel_jpg-LICENSE.txt +3 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/cuda-LICENSE.txt +1582 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/dlpack-LICENSE.txt +201 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/fp16-LICENSE.txt +28 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/libmathdx-LICENSE.txt +220 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/llvm-LICENSE.txt +279 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/moller-LICENSE.txt +16 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/nanovdb-LICENSE.txt +2 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/nvrtc-LICENSE.txt +1592 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/svd-LICENSE.txt +23 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/unittest_parallel-LICENSE.txt +21 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/usd-LICENSE.txt +213 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/windingnumber-LICENSE.txt +21 -0
- warp/examples/assets/cartpole.urdf +0 -110
- warp/examples/assets/crazyflie.usd +0 -0
- warp/examples/assets/nv_ant.xml +0 -92
- warp/examples/assets/nv_humanoid.xml +0 -183
- warp/examples/assets/quadruped.urdf +0 -268
- warp/examples/optim/example_bounce.py +0 -266
- warp/examples/optim/example_cloth_throw.py +0 -228
- warp/examples/optim/example_drone.py +0 -870
- warp/examples/optim/example_inverse_kinematics.py +0 -182
- warp/examples/optim/example_inverse_kinematics_torch.py +0 -191
- warp/examples/optim/example_softbody_properties.py +0 -400
- warp/examples/optim/example_spring_cage.py +0 -245
- warp/examples/optim/example_trajectory.py +0 -227
- warp/examples/sim/example_cartpole.py +0 -143
- warp/examples/sim/example_cloth.py +0 -225
- warp/examples/sim/example_cloth_self_contact.py +0 -316
- warp/examples/sim/example_granular.py +0 -130
- warp/examples/sim/example_granular_collision_sdf.py +0 -202
- warp/examples/sim/example_jacobian_ik.py +0 -244
- warp/examples/sim/example_particle_chain.py +0 -124
- warp/examples/sim/example_quadruped.py +0 -203
- warp/examples/sim/example_rigid_chain.py +0 -203
- warp/examples/sim/example_rigid_contact.py +0 -195
- warp/examples/sim/example_rigid_force.py +0 -133
- warp/examples/sim/example_rigid_gyroscopic.py +0 -115
- warp/examples/sim/example_rigid_soft_contact.py +0 -140
- warp/examples/sim/example_soft_body.py +0 -196
- warp/examples/tile/example_tile_walker.py +0 -327
- warp/sim/__init__.py +0 -74
- warp/sim/articulation.py +0 -793
- warp/sim/collide.py +0 -2570
- warp/sim/graph_coloring.py +0 -307
- warp/sim/import_mjcf.py +0 -791
- warp/sim/import_snu.py +0 -227
- warp/sim/import_urdf.py +0 -579
- warp/sim/import_usd.py +0 -898
- warp/sim/inertia.py +0 -357
- warp/sim/integrator.py +0 -245
- warp/sim/integrator_euler.py +0 -2000
- warp/sim/integrator_featherstone.py +0 -2101
- warp/sim/integrator_vbd.py +0 -2487
- warp/sim/integrator_xpbd.py +0 -3295
- warp/sim/model.py +0 -4821
- warp/sim/particles.py +0 -121
- warp/sim/render.py +0 -431
- warp/sim/utils.py +0 -431
- warp/tests/sim/disabled_kinematics.py +0 -244
- warp/tests/sim/test_cloth.py +0 -863
- warp/tests/sim/test_collision.py +0 -743
- warp/tests/sim/test_coloring.py +0 -347
- warp/tests/sim/test_inertia.py +0 -161
- warp/tests/sim/test_model.py +0 -226
- warp/tests/sim/test_sim_grad.py +0 -287
- warp/tests/sim/test_sim_grad_bounce_linear.py +0 -212
- warp/tests/sim/test_sim_kinematics.py +0 -98
- warp/thirdparty/__init__.py +0 -0
- warp_lang-1.9.0.dist-info/RECORD +0 -456
- /warp/{fem → _src/fem}/quadrature/__init__.py +0 -0
- /warp/{tests/sim → _src/thirdparty}/__init__.py +0 -0
- /warp/{thirdparty → _src/thirdparty}/appdirs.py +0 -0
- /warp/{thirdparty → _src/thirdparty}/dlpack.py +0 -0
- {warp_lang-1.9.0.dist-info → warp_lang-1.10.0rc2.dist-info}/WHEEL +0 -0
- {warp_lang-1.9.0.dist-info → warp_lang-1.10.0rc2.dist-info}/licenses/LICENSE.md +0 -0
- {warp_lang-1.9.0.dist-info → warp_lang-1.10.0rc2.dist-info}/top_level.txt +0 -0
|
@@ -0,0 +1,701 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
2
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License.
|
|
15
|
+
|
|
16
|
+
from functools import cached_property
|
|
17
|
+
from typing import Any, ClassVar, Dict, Optional, Set
|
|
18
|
+
|
|
19
|
+
import warp as wp
|
|
20
|
+
from warp._src.codegen import Struct
|
|
21
|
+
from warp._src.fem import cache
|
|
22
|
+
from warp._src.fem.domain import GeometryDomain, Sides
|
|
23
|
+
from warp._src.fem.geometry import DeformedGeometry, Geometry
|
|
24
|
+
from warp._src.fem.operator import Operator, integrand
|
|
25
|
+
from warp._src.fem.space import FunctionSpace, SpacePartition
|
|
26
|
+
from warp._src.fem.types import NULL_ELEMENT_INDEX, ElementKind, Sample
|
|
27
|
+
from warp._src.fem.utils import type_zero_element
|
|
28
|
+
from warp._src.types import (
|
|
29
|
+
is_value,
|
|
30
|
+
type_is_matrix,
|
|
31
|
+
type_is_quaternion,
|
|
32
|
+
type_is_vector,
|
|
33
|
+
type_repr,
|
|
34
|
+
type_scalar_type,
|
|
35
|
+
type_size,
|
|
36
|
+
type_to_warp,
|
|
37
|
+
types_equal,
|
|
38
|
+
)
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
class FieldLike:
|
|
42
|
+
"""Base class for integrable fields"""
|
|
43
|
+
|
|
44
|
+
EvalArg: Struct
|
|
45
|
+
"""Structure containing field-level arguments passed to device functions for field evaluation"""
|
|
46
|
+
|
|
47
|
+
ElementEvalArg: Struct
|
|
48
|
+
"""Structure combining geometry-level and field-level arguments passed to device functions for field evaluation"""
|
|
49
|
+
|
|
50
|
+
def eval_arg_value(self, device) -> "EvalArg": # noqa: F821
|
|
51
|
+
"""Value of the field-level arguments to be passed to device functions"""
|
|
52
|
+
args = self.EvalArg()
|
|
53
|
+
self.fill_eval_arg(args, device)
|
|
54
|
+
return args
|
|
55
|
+
|
|
56
|
+
def fill_eval_arg(self, arg: "FieldLike.EvalArg", device):
|
|
57
|
+
"""Fill the field-level arguments to be passed to device functions"""
|
|
58
|
+
if self.eval_arg_value is __class__.eval_arg_value:
|
|
59
|
+
raise NotImplementedError()
|
|
60
|
+
arg.assign(self.eval_arg_value(device))
|
|
61
|
+
|
|
62
|
+
@property
|
|
63
|
+
def degree(self) -> int:
|
|
64
|
+
"""Polynomial degree of the field, used to estimate necessary quadrature order"""
|
|
65
|
+
raise NotImplementedError()
|
|
66
|
+
|
|
67
|
+
@property
|
|
68
|
+
def name(self) -> str:
|
|
69
|
+
raise NotImplementedError()
|
|
70
|
+
|
|
71
|
+
@property
|
|
72
|
+
def __str__(self) -> str:
|
|
73
|
+
return self.name
|
|
74
|
+
|
|
75
|
+
def gradient_valid(self) -> bool:
|
|
76
|
+
"""Whether the gradient operator is implemented for this field."""
|
|
77
|
+
return False
|
|
78
|
+
|
|
79
|
+
def divergence_valid(self) -> bool:
|
|
80
|
+
"""Whether the divergence operator is implemented for this field."""
|
|
81
|
+
return False
|
|
82
|
+
|
|
83
|
+
@staticmethod
|
|
84
|
+
def eval_inner(args: "ElementEvalArg", s: Sample): # noqa: F821
|
|
85
|
+
"""Device function evaluating the inner field value at a sample point"""
|
|
86
|
+
raise NotImplementedError()
|
|
87
|
+
|
|
88
|
+
@staticmethod
|
|
89
|
+
def eval_grad_inner(args: "ElementEvalArg", s: Sample): # noqa: F821
|
|
90
|
+
"""Device function evaluating the inner field gradient at a sample point"""
|
|
91
|
+
raise NotImplementedError()
|
|
92
|
+
|
|
93
|
+
@staticmethod
|
|
94
|
+
def eval_div_inner(args: "ElementEvalArg", s: Sample): # noqa: F821
|
|
95
|
+
"""Device function evaluating the inner field divergence at a sample point"""
|
|
96
|
+
raise NotImplementedError()
|
|
97
|
+
|
|
98
|
+
@staticmethod
|
|
99
|
+
def eval_outer(args: "ElementEvalArg", s: Sample): # noqa: F821
|
|
100
|
+
"""Device function evaluating the outer field value at a sample point"""
|
|
101
|
+
raise NotImplementedError()
|
|
102
|
+
|
|
103
|
+
@staticmethod
|
|
104
|
+
def eval_grad_outer(args: "ElementEvalArg", s: Sample): # noqa: F821
|
|
105
|
+
"""Device function evaluating the outer field gradient at a sample point"""
|
|
106
|
+
raise NotImplementedError()
|
|
107
|
+
|
|
108
|
+
@staticmethod
|
|
109
|
+
def eval_div_outer(args: "ElementEvalArg", s: Sample): # noqa: F821
|
|
110
|
+
"""Device function evaluating the outer field divergence at a sample point"""
|
|
111
|
+
raise NotImplementedError()
|
|
112
|
+
|
|
113
|
+
@staticmethod
|
|
114
|
+
def eval_degree(args: "ElementEvalArg"): # noqa: F821
|
|
115
|
+
"""Polynomial degree of the field is applicable, or hint for determination of interpolation order"""
|
|
116
|
+
raise NotImplementedError()
|
|
117
|
+
|
|
118
|
+
def notify_operator_usage(self, ops: Set[Operator]):
|
|
119
|
+
"""Makes the Domain aware that the operators `ops` will be applied"""
|
|
120
|
+
pass
|
|
121
|
+
|
|
122
|
+
|
|
123
|
+
class GeometryField(FieldLike):
|
|
124
|
+
"""Base class for fields defined over a geometry"""
|
|
125
|
+
|
|
126
|
+
@property
|
|
127
|
+
def geometry(self) -> Geometry:
|
|
128
|
+
"""Geometry over which the field is expressed"""
|
|
129
|
+
raise NotImplementedError
|
|
130
|
+
|
|
131
|
+
@property
|
|
132
|
+
def element_kind(self) -> ElementKind:
|
|
133
|
+
"""Kind of element over which the field is expressed"""
|
|
134
|
+
raise NotImplementedError
|
|
135
|
+
|
|
136
|
+
@property
|
|
137
|
+
def dtype(self) -> type:
|
|
138
|
+
raise NotImplementedError
|
|
139
|
+
|
|
140
|
+
@staticmethod
|
|
141
|
+
def eval_reference_grad_inner(args: "ElementEvalArg", s: Sample): # noqa: F821
|
|
142
|
+
"""Device function evaluating the inner field gradient with respect to reference element coordinates at a sample point"""
|
|
143
|
+
raise NotImplementedError
|
|
144
|
+
|
|
145
|
+
@staticmethod
|
|
146
|
+
def eval_reference_grad_outer(args: "ElementEvalArg", s: Sample): # noqa: F821
|
|
147
|
+
"""Device function evaluating the outer field gradient with respect to reference element coordinates at a sample point"""
|
|
148
|
+
raise NotImplementedError
|
|
149
|
+
|
|
150
|
+
def trace(self) -> FieldLike:
|
|
151
|
+
"""Trace of this field over lower-dimensional elements"""
|
|
152
|
+
raise NotImplementedError
|
|
153
|
+
|
|
154
|
+
def make_deformed_geometry(self, relative=True) -> Geometry:
|
|
155
|
+
"""Returns a deformed version of the underlying geometry, with positions displaced according to this field's values.
|
|
156
|
+
|
|
157
|
+
Args:
|
|
158
|
+
relative: If ``True``, the field is interpreted as a relative displacement over the original geometry.
|
|
159
|
+
If ``False``, the field values are interpreted as absolute positions.
|
|
160
|
+
|
|
161
|
+
"""
|
|
162
|
+
return DeformedGeometry(self, relative=relative)
|
|
163
|
+
|
|
164
|
+
@cached_property
|
|
165
|
+
def gradient_dtype(self):
|
|
166
|
+
"""Return type of the (world space) gradient operator. Assumes self.gradient_valid()"""
|
|
167
|
+
if type_is_matrix(self.dtype):
|
|
168
|
+
return None
|
|
169
|
+
|
|
170
|
+
if type_is_vector(self.dtype):
|
|
171
|
+
return cache.cached_mat_type(
|
|
172
|
+
shape=(type_size(self.dtype), self.geometry.dimension),
|
|
173
|
+
dtype=type_scalar_type(self.dtype),
|
|
174
|
+
)
|
|
175
|
+
if type_is_quaternion(self.dtype):
|
|
176
|
+
return cache.cached_mat_type(
|
|
177
|
+
shape=(4, self.geometry.dimension),
|
|
178
|
+
dtype=type_scalar_type(self.dtype),
|
|
179
|
+
)
|
|
180
|
+
return cache.cached_vec_type(length=self.geometry.dimension, dtype=type_scalar_type(self.dtype))
|
|
181
|
+
|
|
182
|
+
@cached_property
|
|
183
|
+
def reference_gradient_dtype(self):
|
|
184
|
+
"""Return type of the reference space gradient operator. Assumes self.gradient_valid()"""
|
|
185
|
+
if type_is_matrix(self.dtype):
|
|
186
|
+
return None
|
|
187
|
+
|
|
188
|
+
if type_is_vector(self.dtype):
|
|
189
|
+
return cache.cached_mat_type(
|
|
190
|
+
shape=(type_size(self.dtype), self.geometry.cell_dimension),
|
|
191
|
+
dtype=type_scalar_type(self.dtype),
|
|
192
|
+
)
|
|
193
|
+
if type_is_quaternion(self.dtype):
|
|
194
|
+
return cache.cached_mat_type(
|
|
195
|
+
shape=(4, self.geometry.cell_dimension),
|
|
196
|
+
dtype=type_scalar_type(self.dtype),
|
|
197
|
+
)
|
|
198
|
+
return cache.cached_vec_type(length=self.geometry.cell_dimension, dtype=type_scalar_type(self.dtype))
|
|
199
|
+
|
|
200
|
+
@cached_property
|
|
201
|
+
def divergence_dtype(self):
|
|
202
|
+
"""Return type of the divergence operator. Assumes self.divergence_valid()"""
|
|
203
|
+
if type_is_vector(self.dtype):
|
|
204
|
+
return type_scalar_type(self.dtype)
|
|
205
|
+
if type_is_matrix(self.dtype):
|
|
206
|
+
return cache.cached_vec_type(length=self.dtype._shape_[1], dtype=type_scalar_type(self.dtype))
|
|
207
|
+
return None
|
|
208
|
+
|
|
209
|
+
|
|
210
|
+
class SpaceField(GeometryField):
|
|
211
|
+
"""Base class for fields defined over a function space"""
|
|
212
|
+
|
|
213
|
+
def __init__(self, space: FunctionSpace, space_partition: SpacePartition):
|
|
214
|
+
self._space = space
|
|
215
|
+
self._space_partition = space_partition
|
|
216
|
+
|
|
217
|
+
self.gradient_valid = self.space.gradient_valid
|
|
218
|
+
self.divergence_valid = self.space.divergence_valid
|
|
219
|
+
|
|
220
|
+
@property
|
|
221
|
+
def geometry(self) -> Geometry:
|
|
222
|
+
return self._space.geometry
|
|
223
|
+
|
|
224
|
+
@property
|
|
225
|
+
def element_kind(self) -> ElementKind:
|
|
226
|
+
return self._space.element_kind
|
|
227
|
+
|
|
228
|
+
@property
|
|
229
|
+
def space(self) -> FunctionSpace:
|
|
230
|
+
return self._space
|
|
231
|
+
|
|
232
|
+
@property
|
|
233
|
+
def space_partition(self) -> SpacePartition:
|
|
234
|
+
return self._space_partition
|
|
235
|
+
|
|
236
|
+
@property
|
|
237
|
+
def degree(self) -> int:
|
|
238
|
+
return self.space.degree
|
|
239
|
+
|
|
240
|
+
@property
|
|
241
|
+
def dtype(self) -> type:
|
|
242
|
+
return self.space.dtype
|
|
243
|
+
|
|
244
|
+
@property
|
|
245
|
+
def dof_dtype(self) -> type:
|
|
246
|
+
return self.space.dof_dtype
|
|
247
|
+
|
|
248
|
+
def _make_eval_degree(self):
|
|
249
|
+
ORDER = self.space.ORDER
|
|
250
|
+
|
|
251
|
+
@cache.dynamic_func(suffix=self.name)
|
|
252
|
+
def degree(args: self.ElementEvalArg):
|
|
253
|
+
return ORDER
|
|
254
|
+
|
|
255
|
+
return degree
|
|
256
|
+
|
|
257
|
+
|
|
258
|
+
class DiscreteField(SpaceField):
|
|
259
|
+
"""Explicitly-valued field defined over a partition of a discrete function space"""
|
|
260
|
+
|
|
261
|
+
@property
|
|
262
|
+
def dof_values(self) -> wp.array:
|
|
263
|
+
"""Array of degrees of freedom values"""
|
|
264
|
+
raise NotImplementedError
|
|
265
|
+
|
|
266
|
+
@dof_values.setter
|
|
267
|
+
def dof_values(self, values: wp.array):
|
|
268
|
+
"""Sets degrees of freedom values from an array"""
|
|
269
|
+
raise NotImplementedError
|
|
270
|
+
|
|
271
|
+
@staticmethod
|
|
272
|
+
def set_node_value(args: "FieldLike.EvalArg", node_index: int, value: Any):
|
|
273
|
+
"""Device function setting the value at given node"""
|
|
274
|
+
raise NotImplementedError
|
|
275
|
+
|
|
276
|
+
@cached_property
|
|
277
|
+
def name(self) -> str:
|
|
278
|
+
return f"{self.__class__.__qualname__}_{self.space.name}_{self.space_partition.name}"
|
|
279
|
+
|
|
280
|
+
|
|
281
|
+
class ImplicitField(GeometryField):
|
|
282
|
+
"""Field defined from an arbitrary function over a domain.
|
|
283
|
+
Does not support autodiff yet, so if gradient/divergence evaluation is required corresponding functions must be provided.
|
|
284
|
+
|
|
285
|
+
Args:
|
|
286
|
+
domain: Domain over which the field is defined
|
|
287
|
+
func: Warp function evaluating the field at a given position. Must accept at least one argument, with the first argument being the evaluation position (``wp.vec2`` or ``wp.vec3``).
|
|
288
|
+
values: Optional dictionary of additional argument values to be passed to the evaluation function.
|
|
289
|
+
grad_func: Optional gradient evaluation function; must take same arguments as `func`
|
|
290
|
+
div_func: Optional divergence evaluation function; must take same arguments as `func`
|
|
291
|
+
degree: Optional hint for automatic determination of quadrature orders when integrating this field
|
|
292
|
+
"""
|
|
293
|
+
|
|
294
|
+
_dynamic_attribute_constructors: ClassVar = {
|
|
295
|
+
"ElementEvalArg": lambda obj: obj._make_element_eval_arg(),
|
|
296
|
+
"eval_degree": lambda obj: obj._make_eval_degree(),
|
|
297
|
+
"eval_inner": lambda obj: obj._make_eval_func(obj._func),
|
|
298
|
+
"eval_grad_inner": lambda obj: obj._make_eval_func(obj._grad_func),
|
|
299
|
+
"eval_div_inner": lambda obj: obj._make_eval_func(obj._div_func),
|
|
300
|
+
"eval_reference_grad_inner": lambda obj: obj._make_eval_reference_grad(),
|
|
301
|
+
"eval_outer": lambda obj: obj.eval_inner,
|
|
302
|
+
"eval_grad_outer": lambda obj: obj.eval_grad_inner,
|
|
303
|
+
"eval_div_outer": lambda obj: obj.eval_div_inner,
|
|
304
|
+
"eval_reference_grad_outer": lambda obj: obj.eval_reference_grad_inner,
|
|
305
|
+
}
|
|
306
|
+
|
|
307
|
+
def __init__(
|
|
308
|
+
self,
|
|
309
|
+
domain: GeometryDomain,
|
|
310
|
+
func: wp.Function,
|
|
311
|
+
values: Optional[Dict[str, Any]] = None,
|
|
312
|
+
grad_func: Optional[wp.Function] = None,
|
|
313
|
+
div_func: Optional[wp.Function] = None,
|
|
314
|
+
degree=0,
|
|
315
|
+
):
|
|
316
|
+
self.domain = domain
|
|
317
|
+
self._degree = degree
|
|
318
|
+
|
|
319
|
+
if not isinstance(func, wp.Function):
|
|
320
|
+
raise ValueError("Implicit field function must be a warp Function (decorated with `wp.func`)")
|
|
321
|
+
|
|
322
|
+
self._func = func
|
|
323
|
+
self._grad_func = grad_func
|
|
324
|
+
self._div_func = div_func
|
|
325
|
+
|
|
326
|
+
argspec = integrand(func.func).argspec
|
|
327
|
+
arg_types = {**argspec.annotations} # make a mutable copy
|
|
328
|
+
|
|
329
|
+
try:
|
|
330
|
+
first_arg_type = arg_types.pop(argspec.args[0])
|
|
331
|
+
if types_equal(first_arg_type, wp.vec(length=domain.geometry.dimension, dtype=float), match_generic=True):
|
|
332
|
+
self._qp_based = False
|
|
333
|
+
elif type_to_warp(first_arg_type) == wp.int32:
|
|
334
|
+
self._qp_based = True
|
|
335
|
+
else:
|
|
336
|
+
raise TypeError(f"Unsupported argument type `{type_repr(first_arg_type)}`")
|
|
337
|
+
except Exception as err:
|
|
338
|
+
raise ValueError(
|
|
339
|
+
f"Implicit field function '{func.func.__name__}' must accept either a position or a quadrature point index as its first argument"
|
|
340
|
+
) from err
|
|
341
|
+
|
|
342
|
+
self.EvalArg = cache.get_argument_struct(arg_types)
|
|
343
|
+
self._func_arg = self.EvalArg()
|
|
344
|
+
self.values = values
|
|
345
|
+
|
|
346
|
+
cache.setup_dynamic_attributes(self)
|
|
347
|
+
|
|
348
|
+
@property
|
|
349
|
+
def values(self):
|
|
350
|
+
return self._func_arg
|
|
351
|
+
|
|
352
|
+
@values.setter
|
|
353
|
+
def values(self, v):
|
|
354
|
+
self._values = v
|
|
355
|
+
cache.populate_argument_struct(self._func_arg, v, self._func.func.__name__)
|
|
356
|
+
|
|
357
|
+
@property
|
|
358
|
+
def geometry(self) -> Geometry:
|
|
359
|
+
return self.domain.geometry
|
|
360
|
+
|
|
361
|
+
@property
|
|
362
|
+
def element_kind(self) -> ElementKind:
|
|
363
|
+
return self.domain.element_kind
|
|
364
|
+
|
|
365
|
+
@property
|
|
366
|
+
def dtype(self):
|
|
367
|
+
# Cannot determine dtype from function
|
|
368
|
+
return None
|
|
369
|
+
|
|
370
|
+
def eval_arg_value(self, device):
|
|
371
|
+
return self._func_arg
|
|
372
|
+
|
|
373
|
+
@property
|
|
374
|
+
def degree(self) -> int:
|
|
375
|
+
return self._degree
|
|
376
|
+
|
|
377
|
+
@property
|
|
378
|
+
def name(self) -> str:
|
|
379
|
+
return f"Implicit_{self.domain.name}_{self.degree}_{self.EvalArg.key}"
|
|
380
|
+
|
|
381
|
+
def gradient_valid(self) -> bool:
|
|
382
|
+
return self._grad_func is not None
|
|
383
|
+
|
|
384
|
+
def divergence_valid(self) -> bool:
|
|
385
|
+
return self._div_func is not None
|
|
386
|
+
|
|
387
|
+
def _make_eval_func(self, func):
|
|
388
|
+
if func is None:
|
|
389
|
+
return None
|
|
390
|
+
|
|
391
|
+
@cache.dynamic_func(
|
|
392
|
+
suffix=(self.name, func.key),
|
|
393
|
+
code_transformers=[cache.ExpandStarredArgumentStruct({"args.eval_arg": self.EvalArg})],
|
|
394
|
+
)
|
|
395
|
+
def eval_inner(args: self.ElementEvalArg, s: Sample):
|
|
396
|
+
if wp.static(self._qp_based):
|
|
397
|
+
qp_index = s.qp_index
|
|
398
|
+
return func(qp_index, *args.eval_arg)
|
|
399
|
+
else:
|
|
400
|
+
pos = self.domain.element_position(args.elt_arg, s)
|
|
401
|
+
return func(pos, *args.eval_arg)
|
|
402
|
+
|
|
403
|
+
return eval_inner
|
|
404
|
+
|
|
405
|
+
def _make_eval_reference_grad(self):
|
|
406
|
+
if self.eval_grad_inner is None:
|
|
407
|
+
return None
|
|
408
|
+
|
|
409
|
+
@cache.dynamic_func(suffix=f"{self.eval_grad_inner.key}")
|
|
410
|
+
def eval_reference_grad_inner(args: self.ElementEvalArg, s: Sample):
|
|
411
|
+
return self.eval_grad_inner(args, s) * self.domain.element_deformation_gradient(args.elt_arg, s)
|
|
412
|
+
|
|
413
|
+
return eval_reference_grad_inner
|
|
414
|
+
|
|
415
|
+
def _make_element_eval_arg(self):
|
|
416
|
+
@cache.dynamic_struct(suffix=self.name)
|
|
417
|
+
class ElementEvalArg:
|
|
418
|
+
elt_arg: self.domain.ElementArg
|
|
419
|
+
eval_arg: self.EvalArg
|
|
420
|
+
|
|
421
|
+
return ElementEvalArg
|
|
422
|
+
|
|
423
|
+
def _make_eval_degree(self):
|
|
424
|
+
ORDER = wp.constant(self._degree)
|
|
425
|
+
|
|
426
|
+
@cache.dynamic_func(suffix=self.name)
|
|
427
|
+
def degree(args: self.ElementEvalArg):
|
|
428
|
+
return ORDER
|
|
429
|
+
|
|
430
|
+
return degree
|
|
431
|
+
|
|
432
|
+
def trace(self):
|
|
433
|
+
if self.element_kind == ElementKind.SIDE:
|
|
434
|
+
raise RuntimeError("Trace only available for field defined on cell elements")
|
|
435
|
+
|
|
436
|
+
return ImplicitField(
|
|
437
|
+
domain=Sides(self.domain.geometry_partition),
|
|
438
|
+
func=self._func,
|
|
439
|
+
values={name: getattr(self.values, name) for name in self.EvalArg.vars},
|
|
440
|
+
grad_func=self._grad_func,
|
|
441
|
+
div_func=self._div_func,
|
|
442
|
+
degree=self._degree,
|
|
443
|
+
)
|
|
444
|
+
|
|
445
|
+
|
|
446
|
+
class UniformField(GeometryField):
|
|
447
|
+
"""Field defined as a constant value over a domain.
|
|
448
|
+
|
|
449
|
+
Args:
|
|
450
|
+
domain: Domain over which the field is defined
|
|
451
|
+
value: Uniform value over the domain
|
|
452
|
+
"""
|
|
453
|
+
|
|
454
|
+
_dynamic_attribute_constructors: ClassVar = {
|
|
455
|
+
"EvalArg": lambda obj: obj._make_eval_arg(),
|
|
456
|
+
"ElementEvalArg": lambda obj: obj._make_element_eval_arg(),
|
|
457
|
+
"eval_degree": lambda obj: obj._make_eval_degree(),
|
|
458
|
+
"eval_inner": lambda obj: obj._make_eval_inner(),
|
|
459
|
+
"eval_grad_inner": lambda obj: obj._make_eval_zero(obj.gradient_dtype),
|
|
460
|
+
"eval_div_inner": lambda obj: obj._make_eval_zero(obj.divergence_dtype),
|
|
461
|
+
"eval_reference_grad_inner": lambda obj: obj._make_eval_zero(obj.reference_gradient_dtype),
|
|
462
|
+
"eval_outer": lambda obj: obj.eval_inner,
|
|
463
|
+
"eval_grad_outer": lambda obj: obj.eval_grad_inner,
|
|
464
|
+
"eval_div_outer": lambda obj: obj.eval_div_inner,
|
|
465
|
+
"eval_reference_grad_outer": lambda obj: obj.eval_reference_grad_inner,
|
|
466
|
+
}
|
|
467
|
+
|
|
468
|
+
def __init__(self, domain: GeometryDomain, value: Any):
|
|
469
|
+
self.domain = domain
|
|
470
|
+
|
|
471
|
+
if not is_value(value):
|
|
472
|
+
raise ValueError("value must be a Warp scalar, vector or matrix")
|
|
473
|
+
|
|
474
|
+
value_type = type_to_warp(type(value))
|
|
475
|
+
self._value = value_type(value)
|
|
476
|
+
|
|
477
|
+
cache.setup_dynamic_attributes(self)
|
|
478
|
+
|
|
479
|
+
@property
|
|
480
|
+
def value(self):
|
|
481
|
+
return self._value
|
|
482
|
+
|
|
483
|
+
@value.setter
|
|
484
|
+
def value(self, v):
|
|
485
|
+
value_type = type_to_warp(type(v))
|
|
486
|
+
assert types_equal(value_type, self.dtype)
|
|
487
|
+
self._value = self.dtype(v)
|
|
488
|
+
|
|
489
|
+
@property
|
|
490
|
+
def geometry(self) -> Geometry:
|
|
491
|
+
return self.domain.geometry
|
|
492
|
+
|
|
493
|
+
@property
|
|
494
|
+
def element_kind(self) -> ElementKind:
|
|
495
|
+
return self.domain.element_kind
|
|
496
|
+
|
|
497
|
+
@property
|
|
498
|
+
def dtype(self) -> type:
|
|
499
|
+
return type(self.value)
|
|
500
|
+
|
|
501
|
+
def fill_eval_arg(self, arg, device):
|
|
502
|
+
arg.value = self.value
|
|
503
|
+
|
|
504
|
+
@property
|
|
505
|
+
def degree(self) -> int:
|
|
506
|
+
return 0
|
|
507
|
+
|
|
508
|
+
def gradient_valid(self) -> bool:
|
|
509
|
+
return self.gradient_dtype is not None
|
|
510
|
+
|
|
511
|
+
def divergence_valid(self) -> bool:
|
|
512
|
+
return self.divergence_dtype is not None
|
|
513
|
+
|
|
514
|
+
@cached_property
|
|
515
|
+
def name(self) -> str:
|
|
516
|
+
return f"Uniform{self.domain.name}_{cache.pod_type_key(self.dtype)}"
|
|
517
|
+
|
|
518
|
+
def _make_eval_inner(self):
|
|
519
|
+
@cache.dynamic_func(suffix=self.name)
|
|
520
|
+
def eval_inner(args: self.ElementEvalArg, s: Sample):
|
|
521
|
+
return args.eval_arg.value
|
|
522
|
+
|
|
523
|
+
return eval_inner
|
|
524
|
+
|
|
525
|
+
def _make_eval_zero(self, dtype):
|
|
526
|
+
if dtype is None:
|
|
527
|
+
return None
|
|
528
|
+
|
|
529
|
+
zero_element = type_zero_element(dtype)
|
|
530
|
+
|
|
531
|
+
@cache.dynamic_func(suffix=f"{self.name}_{cache.pod_type_key(dtype)}")
|
|
532
|
+
def eval_zero(args: self.ElementEvalArg, s: Sample):
|
|
533
|
+
return zero_element()
|
|
534
|
+
|
|
535
|
+
return eval_zero
|
|
536
|
+
|
|
537
|
+
def _make_eval_arg(self):
|
|
538
|
+
@cache.dynamic_struct(suffix=self.name)
|
|
539
|
+
class EvalArg:
|
|
540
|
+
value: self.dtype
|
|
541
|
+
|
|
542
|
+
return EvalArg
|
|
543
|
+
|
|
544
|
+
def _make_element_eval_arg(self):
|
|
545
|
+
@cache.dynamic_struct(suffix=self.name)
|
|
546
|
+
class ElementEvalArg:
|
|
547
|
+
elt_arg: self.domain.ElementArg
|
|
548
|
+
eval_arg: self.EvalArg
|
|
549
|
+
|
|
550
|
+
return ElementEvalArg
|
|
551
|
+
|
|
552
|
+
def _make_eval_degree(self):
|
|
553
|
+
@cache.dynamic_func(suffix=self.name)
|
|
554
|
+
def degree(args: self.ElementEvalArg):
|
|
555
|
+
return 0
|
|
556
|
+
|
|
557
|
+
return degree
|
|
558
|
+
|
|
559
|
+
def trace(self):
|
|
560
|
+
if self.element_kind == ElementKind.SIDE:
|
|
561
|
+
raise RuntimeError("Trace only available for field defined on cell elements")
|
|
562
|
+
|
|
563
|
+
return UniformField(domain=Sides(self.domain.geometry_partition), value=self.value)
|
|
564
|
+
|
|
565
|
+
|
|
566
|
+
class NonconformingField(GeometryField):
|
|
567
|
+
"""Field defined as the map of a DiscreteField over a non-conforming geometry.
|
|
568
|
+
|
|
569
|
+
Args:
|
|
570
|
+
domain: The new domain over which the nonconforming field will be evaluated
|
|
571
|
+
field: Nonconforming discrete field
|
|
572
|
+
background: Uniform value or domain-conforming field determining the value outside of the geometry of definition of `field`
|
|
573
|
+
"""
|
|
574
|
+
|
|
575
|
+
_LOOKUP_EPS = wp.constant(1.0e-6)
|
|
576
|
+
|
|
577
|
+
_dynamic_attribute_constructors: ClassVar = {
|
|
578
|
+
"EvalArg": lambda obj: obj._make_eval_arg(),
|
|
579
|
+
"ElementEvalArg": lambda obj: obj._make_element_eval_arg(),
|
|
580
|
+
"eval_degree": lambda obj: obj._make_eval_degree(),
|
|
581
|
+
"eval_inner": lambda obj: obj._make_nonconforming_eval("eval_inner"),
|
|
582
|
+
"eval_grad_inner": lambda obj: obj._make_nonconforming_eval("eval_grad_inner"),
|
|
583
|
+
"eval_div_inner": lambda obj: obj._make_nonconforming_eval("eval_div_inner"),
|
|
584
|
+
"eval_reference_grad_inner": lambda obj: obj._make_eval_reference_grad(),
|
|
585
|
+
"eval_outer": lambda obj: obj.eval_inner,
|
|
586
|
+
"eval_grad_outer": lambda obj: obj.eval_grad_inner,
|
|
587
|
+
"eval_div_outer": lambda obj: obj.eval_div_inner,
|
|
588
|
+
"eval_reference_grad_outer": lambda obj: obj.eval_reference_grad_inner,
|
|
589
|
+
}
|
|
590
|
+
|
|
591
|
+
def __init__(self, domain: GeometryDomain, field: DiscreteField, background: Any = 0.0):
|
|
592
|
+
self.domain = domain
|
|
593
|
+
|
|
594
|
+
self.field = field
|
|
595
|
+
|
|
596
|
+
if not isinstance(background, GeometryField):
|
|
597
|
+
background = UniformField(domain, self.dtype(background))
|
|
598
|
+
elif background.geometry != domain.geometry or background.element_kind != domain.element_kind:
|
|
599
|
+
raise ValueError("Background field must be conforming to the domain")
|
|
600
|
+
self.background = background
|
|
601
|
+
|
|
602
|
+
cache.setup_dynamic_attributes(self)
|
|
603
|
+
|
|
604
|
+
@property
|
|
605
|
+
def geometry(self) -> Geometry:
|
|
606
|
+
return self.domain.geometry
|
|
607
|
+
|
|
608
|
+
@property
|
|
609
|
+
def element_kind(self) -> ElementKind:
|
|
610
|
+
return self.domain.element_kind
|
|
611
|
+
|
|
612
|
+
@property
|
|
613
|
+
def dtype(self) -> type:
|
|
614
|
+
return self.field.dtype
|
|
615
|
+
|
|
616
|
+
def fill_eval_arg(self, arg, device):
|
|
617
|
+
self.field.fill_eval_arg(arg.field_cell_eval_arg.eval_arg, device)
|
|
618
|
+
self.field.geometry.fill_cell_arg(arg.field_cell_eval_arg.elt_arg, device)
|
|
619
|
+
self.background.fill_eval_arg(arg.background_arg, device)
|
|
620
|
+
|
|
621
|
+
@property
|
|
622
|
+
def degree(self) -> int:
|
|
623
|
+
return self.field.degree
|
|
624
|
+
|
|
625
|
+
def gradient_valid(self) -> bool:
|
|
626
|
+
return self.field.gradient_valid() and self.background.gradient_valid()
|
|
627
|
+
|
|
628
|
+
def divergence_valid(self) -> bool:
|
|
629
|
+
return self.field.divergence_valid() and self.background.divergence_valid()
|
|
630
|
+
|
|
631
|
+
@cached_property
|
|
632
|
+
def name(self) -> str:
|
|
633
|
+
return f"{self.domain.name}_{self.field.name}_{self.background.name}"
|
|
634
|
+
|
|
635
|
+
def _make_nonconforming_eval(self, eval_func_name):
|
|
636
|
+
field_eval = getattr(self.field, eval_func_name)
|
|
637
|
+
bg_eval = getattr(self.background, eval_func_name)
|
|
638
|
+
|
|
639
|
+
if field_eval is None or bg_eval is None:
|
|
640
|
+
return None
|
|
641
|
+
|
|
642
|
+
cell_lookup = self.field.geometry.cell_lookup
|
|
643
|
+
|
|
644
|
+
@cache.dynamic_func(suffix=f"{eval_func_name}_{self.name}")
|
|
645
|
+
def eval_nc(args: self.ElementEvalArg, s: Sample):
|
|
646
|
+
pos = self.domain.element_position(args.elt_arg, s)
|
|
647
|
+
cell_arg = args.eval_arg.field_cell_eval_arg.elt_arg
|
|
648
|
+
nonconforming_s = cell_lookup(cell_arg, pos, NonconformingField._LOOKUP_EPS)
|
|
649
|
+
if nonconforming_s.element_index != NULL_ELEMENT_INDEX:
|
|
650
|
+
if (
|
|
651
|
+
wp.length_sq(pos - self.field.geometry.cell_position(cell_arg, nonconforming_s))
|
|
652
|
+
<= NonconformingField._LOOKUP_EPS
|
|
653
|
+
):
|
|
654
|
+
return field_eval(
|
|
655
|
+
self.field.ElementEvalArg(cell_arg, args.eval_arg.field_cell_eval_arg.eval_arg), nonconforming_s
|
|
656
|
+
)
|
|
657
|
+
|
|
658
|
+
return bg_eval(self.background.ElementEvalArg(args.elt_arg, args.eval_arg.background_arg), s)
|
|
659
|
+
|
|
660
|
+
return eval_nc
|
|
661
|
+
|
|
662
|
+
def _make_eval_reference_grad(self):
|
|
663
|
+
if self.eval_grad_inner is None:
|
|
664
|
+
return None
|
|
665
|
+
|
|
666
|
+
@cache.dynamic_func(suffix=f"{self.eval_grad_inner.key}")
|
|
667
|
+
def eval_reference_grad_inner(args: self.ElementEvalArg, s: Sample):
|
|
668
|
+
return self.eval_grad_inner(args, s) * self.domain.element_deformation_gradient(args.elt_arg, s)
|
|
669
|
+
|
|
670
|
+
return eval_reference_grad_inner
|
|
671
|
+
|
|
672
|
+
def _make_eval_arg(self):
|
|
673
|
+
@cache.dynamic_struct(suffix=self.name)
|
|
674
|
+
class EvalArg:
|
|
675
|
+
field_cell_eval_arg: self.field.ElementEvalArg
|
|
676
|
+
background_arg: self.background.EvalArg
|
|
677
|
+
|
|
678
|
+
return EvalArg
|
|
679
|
+
|
|
680
|
+
def _make_element_eval_arg(self):
|
|
681
|
+
@cache.dynamic_struct(suffix=self.name)
|
|
682
|
+
class ElementEvalArg:
|
|
683
|
+
elt_arg: self.domain.ElementArg
|
|
684
|
+
eval_arg: self.EvalArg
|
|
685
|
+
|
|
686
|
+
return ElementEvalArg
|
|
687
|
+
|
|
688
|
+
def _make_eval_degree(self):
|
|
689
|
+
@cache.dynamic_func(suffix=self.name)
|
|
690
|
+
def degree(args: self.ElementEvalArg):
|
|
691
|
+
return self.field.eval_degree(args.eval_arg.field_cell_eval_arg)
|
|
692
|
+
|
|
693
|
+
return degree
|
|
694
|
+
|
|
695
|
+
def trace(self):
|
|
696
|
+
if self.element_kind == ElementKind.SIDE:
|
|
697
|
+
raise RuntimeError("Trace only available for field defined on cell elements")
|
|
698
|
+
|
|
699
|
+
return NonconformingField(
|
|
700
|
+
domain=Sides(self.domain.geometry_partition), field=self.field, background=self.background.trace()
|
|
701
|
+
)
|