warp-lang 0.10.1__py3-none-win_amd64.whl → 0.11.0__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 +10 -4
- warp/__init__.pyi +1 -0
- warp/bin/warp-clang.dll +0 -0
- warp/bin/warp.dll +0 -0
- warp/build.py +5 -3
- warp/build_dll.py +29 -9
- warp/builtins.py +868 -507
- warp/codegen.py +1074 -638
- warp/config.py +3 -3
- warp/constants.py +6 -0
- warp/context.py +715 -222
- warp/fabric.py +326 -0
- warp/fem/__init__.py +27 -0
- warp/fem/cache.py +389 -0
- warp/fem/dirichlet.py +181 -0
- warp/fem/domain.py +263 -0
- warp/fem/field/__init__.py +101 -0
- warp/fem/field/field.py +149 -0
- warp/fem/field/nodal_field.py +299 -0
- warp/fem/field/restriction.py +21 -0
- warp/fem/field/test.py +181 -0
- warp/fem/field/trial.py +183 -0
- warp/fem/geometry/__init__.py +19 -0
- warp/fem/geometry/closest_point.py +70 -0
- warp/fem/geometry/deformed_geometry.py +271 -0
- warp/fem/geometry/element.py +744 -0
- warp/fem/geometry/geometry.py +186 -0
- warp/fem/geometry/grid_2d.py +373 -0
- warp/fem/geometry/grid_3d.py +435 -0
- warp/fem/geometry/hexmesh.py +953 -0
- warp/fem/geometry/partition.py +376 -0
- warp/fem/geometry/quadmesh_2d.py +532 -0
- warp/fem/geometry/tetmesh.py +840 -0
- warp/fem/geometry/trimesh_2d.py +577 -0
- warp/fem/integrate.py +1616 -0
- warp/fem/operator.py +191 -0
- warp/fem/polynomial.py +213 -0
- warp/fem/quadrature/__init__.py +2 -0
- warp/fem/quadrature/pic_quadrature.py +245 -0
- warp/fem/quadrature/quadrature.py +294 -0
- warp/fem/space/__init__.py +292 -0
- warp/fem/space/basis_space.py +489 -0
- warp/fem/space/collocated_function_space.py +105 -0
- warp/fem/space/dof_mapper.py +236 -0
- warp/fem/space/function_space.py +145 -0
- warp/fem/space/grid_2d_function_space.py +267 -0
- warp/fem/space/grid_3d_function_space.py +306 -0
- warp/fem/space/hexmesh_function_space.py +352 -0
- warp/fem/space/partition.py +350 -0
- warp/fem/space/quadmesh_2d_function_space.py +369 -0
- warp/fem/space/restriction.py +160 -0
- warp/fem/space/shape/__init__.py +15 -0
- warp/fem/space/shape/cube_shape_function.py +738 -0
- warp/fem/space/shape/shape_function.py +103 -0
- warp/fem/space/shape/square_shape_function.py +611 -0
- warp/fem/space/shape/tet_shape_function.py +567 -0
- warp/fem/space/shape/triangle_shape_function.py +429 -0
- warp/fem/space/tetmesh_function_space.py +292 -0
- warp/fem/space/topology.py +295 -0
- warp/fem/space/trimesh_2d_function_space.py +221 -0
- warp/fem/types.py +77 -0
- warp/fem/utils.py +495 -0
- warp/native/array.h +147 -44
- warp/native/builtin.h +122 -149
- warp/native/bvh.cpp +73 -325
- warp/native/bvh.cu +406 -23
- warp/native/bvh.h +34 -43
- warp/native/clang/clang.cpp +13 -8
- warp/native/crt.h +2 -0
- warp/native/cuda_crt.h +5 -0
- warp/native/cuda_util.cpp +15 -3
- warp/native/cuda_util.h +3 -1
- warp/native/cutlass/tools/library/scripts/conv2d_operation.py +463 -0
- warp/native/cutlass/tools/library/scripts/conv3d_operation.py +321 -0
- warp/native/cutlass/tools/library/scripts/gemm_operation.py +988 -0
- warp/native/cutlass/tools/library/scripts/generator.py +4625 -0
- warp/native/cutlass/tools/library/scripts/library.py +799 -0
- warp/native/cutlass/tools/library/scripts/manifest.py +402 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/docs/source/conf.py +96 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/profile/conv/conv2d_f16_sm80.py +106 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/profile/gemm/gemm_f32_sm80.py +91 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/setup.py +80 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/__init__.py +48 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/arguments.py +118 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/c_types.py +241 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/compiler.py +432 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/conv2d_operation.py +631 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/epilogue.py +1026 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/frontend.py +104 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/gemm_operation.py +1276 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/library.py +744 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/memory_manager.py +74 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/operation.py +110 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/parser.py +619 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/reduction_operation.py +398 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/tensor_ref.py +70 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/test/__init__.py +4 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/test/conv2d_testbed.py +646 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/test/gemm_grouped_testbed.py +235 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/test/gemm_testbed.py +557 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/test/profiler.py +70 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/type_hint.py +39 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/utils/__init__.py +1 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/utils/device.py +76 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/utils/reference_model.py +255 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/__init__.py +0 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_dgrad_implicit_gemm_f16nhwc_f16nhwc_f16nhwc_tensor_op_f16_sm80.py +201 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_dgrad_implicit_gemm_f16nhwc_f16nhwc_f32nhwc_tensor_op_f32_sm80.py +177 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_dgrad_implicit_gemm_f32nhwc_f32nhwc_f32nhwc_simt_f32_sm80.py +98 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_dgrad_implicit_gemm_tf32nhwc_tf32nhwc_f32nhwc_tensor_op_f32_sm80.py +95 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_fprop_few_channels_f16nhwc_f16nhwc_f16nhwc_tensor_op_f32_sm80.py +163 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_fprop_fixed_channels_f16nhwc_f16nhwc_f16nhwc_tensor_op_f32_sm80.py +187 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_fprop_implicit_gemm_f16nhwc_f16nhwc_f16nhwc_tensor_op_f16_sm80.py +309 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_fprop_implicit_gemm_f16nhwc_f16nhwc_f32nhwc_tensor_op_f32_sm80.py +54 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_fprop_implicit_gemm_f32nhwc_f32nhwc_f32nhwc_simt_f32_sm80.py +96 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_fprop_implicit_gemm_tf32nhwc_tf32nhwc_f32nhwc_tensor_op_f32_sm80.py +107 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_strided_dgrad_implicit_gemm_f16nhwc_f16nhwc_f32nhwc_tensor_op_f32_sm80.py +253 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_wgrad_implicit_gemm_f16nhwc_f16nhwc_f16nhwc_tensor_op_f16_sm80.py +97 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_wgrad_implicit_gemm_f16nhwc_f16nhwc_f32nhwc_tensor_op_f32_sm80.py +242 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_wgrad_implicit_gemm_f32nhwc_f32nhwc_f32nhwc_simt_f32_sm80.py +96 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_wgrad_implicit_gemm_tf32nhwc_tf32nhwc_f32nhwc_tensor_op_f32_sm80.py +107 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/run_all_tests.py +10 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/test/frontend/test_frontend.py +146 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/test/gemm/__init__.py +0 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/test/gemm/gemm_bf16_sm80.py +96 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/test/gemm/gemm_f16_sm80.py +447 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/test/gemm/gemm_f32_sm80.py +146 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/test/gemm/gemm_f64_sm80.py +102 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/test/gemm/gemm_grouped_sm80.py +203 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/test/gemm/gemm_s8_sm80.py +229 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/test/gemm/run_all_tests.py +9 -0
- warp/native/cutlass/tools/library/scripts/pycutlass/test/unit/test_sm80.py +453 -0
- warp/native/cutlass/tools/library/scripts/rank_2k_operation.py +398 -0
- warp/native/cutlass/tools/library/scripts/rank_k_operation.py +387 -0
- warp/native/cutlass/tools/library/scripts/rt.py +796 -0
- warp/native/cutlass/tools/library/scripts/symm_operation.py +400 -0
- warp/native/cutlass/tools/library/scripts/trmm_operation.py +407 -0
- warp/native/cutlass_gemm.cu +5 -3
- warp/native/exports.h +1240 -952
- warp/native/fabric.h +228 -0
- warp/native/hashgrid.cpp +4 -4
- warp/native/hashgrid.h +22 -2
- warp/native/intersect.h +22 -7
- warp/native/intersect_adj.h +8 -8
- warp/native/intersect_tri.h +1 -1
- warp/native/marching.cu +157 -161
- warp/native/mat.h +80 -19
- warp/native/matnn.h +2 -2
- warp/native/mesh.cpp +33 -108
- warp/native/mesh.cu +114 -23
- warp/native/mesh.h +446 -46
- warp/native/noise.h +272 -329
- warp/native/quat.h +51 -8
- warp/native/rand.h +45 -35
- warp/native/range.h +6 -2
- warp/native/reduce.cpp +1 -1
- warp/native/reduce.cu +10 -12
- warp/native/runlength_encode.cu +6 -10
- warp/native/scan.cu +8 -11
- warp/native/sparse.cpp +4 -4
- warp/native/sparse.cu +164 -154
- warp/native/spatial.h +2 -2
- warp/native/temp_buffer.h +14 -30
- warp/native/vec.h +107 -23
- warp/native/volume.h +120 -0
- warp/native/warp.cpp +560 -30
- warp/native/warp.cu +431 -44
- warp/native/warp.h +13 -4
- warp/optim/__init__.py +1 -0
- warp/optim/linear.py +922 -0
- warp/optim/sgd.py +92 -0
- warp/render/render_opengl.py +335 -119
- warp/render/render_usd.py +11 -11
- warp/sim/__init__.py +2 -2
- warp/sim/articulation.py +385 -185
- warp/sim/collide.py +8 -0
- warp/sim/import_mjcf.py +297 -106
- warp/sim/import_urdf.py +389 -210
- warp/sim/import_usd.py +198 -97
- warp/sim/inertia.py +17 -18
- warp/sim/integrator_euler.py +14 -8
- warp/sim/integrator_xpbd.py +158 -16
- warp/sim/model.py +795 -291
- warp/sim/render.py +3 -3
- warp/sim/utils.py +3 -0
- warp/sparse.py +640 -150
- warp/stubs.py +606 -267
- warp/tape.py +61 -10
- warp/tests/__main__.py +3 -6
- warp/tests/assets/curlnoise_golden.npy +0 -0
- warp/tests/assets/pnoise_golden.npy +0 -0
- warp/tests/{test_class_kernel.py → aux_test_class_kernel.py} +9 -1
- warp/tests/aux_test_conditional_unequal_types_kernels.py +21 -0
- warp/tests/{test_dependent.py → aux_test_dependent.py} +2 -2
- warp/tests/{test_reference.py → aux_test_reference.py} +1 -1
- warp/tests/aux_test_unresolved_func.py +14 -0
- warp/tests/aux_test_unresolved_symbol.py +14 -0
- warp/tests/disabled_kinematics.py +239 -0
- warp/tests/run_coverage_serial.py +31 -0
- warp/tests/test_adam.py +103 -106
- warp/tests/test_arithmetic.py +128 -74
- warp/tests/test_array.py +212 -97
- warp/tests/test_array_reduce.py +57 -23
- warp/tests/test_atomic.py +64 -28
- warp/tests/test_bool.py +99 -0
- warp/tests/test_builtins_resolution.py +1292 -0
- warp/tests/test_bvh.py +42 -18
- warp/tests/test_closest_point_edge_edge.py +54 -57
- warp/tests/test_codegen.py +208 -130
- warp/tests/test_compile_consts.py +28 -20
- warp/tests/test_conditional.py +108 -24
- warp/tests/test_copy.py +10 -12
- warp/tests/test_ctypes.py +112 -88
- warp/tests/test_dense.py +21 -14
- warp/tests/test_devices.py +98 -0
- warp/tests/test_dlpack.py +75 -75
- warp/tests/test_examples.py +277 -0
- warp/tests/test_fabricarray.py +955 -0
- warp/tests/test_fast_math.py +15 -11
- warp/tests/test_fem.py +1271 -0
- warp/tests/test_fp16.py +53 -19
- warp/tests/test_func.py +187 -86
- warp/tests/test_generics.py +194 -49
- warp/tests/test_grad.py +178 -109
- warp/tests/test_grad_customs.py +176 -0
- warp/tests/test_hash_grid.py +52 -37
- warp/tests/test_import.py +10 -23
- warp/tests/test_indexedarray.py +32 -31
- warp/tests/test_intersect.py +18 -9
- warp/tests/test_large.py +141 -0
- warp/tests/test_launch.py +14 -41
- warp/tests/test_lerp.py +64 -65
- warp/tests/test_linear_solvers.py +154 -0
- warp/tests/test_lvalue.py +493 -0
- warp/tests/test_marching_cubes.py +12 -13
- warp/tests/test_mat.py +517 -2898
- warp/tests/test_mat_lite.py +115 -0
- warp/tests/test_mat_scalar_ops.py +2889 -0
- warp/tests/test_math.py +103 -9
- warp/tests/test_matmul.py +305 -69
- warp/tests/test_matmul_lite.py +410 -0
- warp/tests/test_mesh.py +71 -14
- warp/tests/test_mesh_query_aabb.py +41 -25
- warp/tests/test_mesh_query_point.py +140 -22
- warp/tests/test_mesh_query_ray.py +39 -22
- warp/tests/test_mlp.py +30 -22
- warp/tests/test_model.py +92 -89
- warp/tests/test_modules_lite.py +39 -0
- warp/tests/test_multigpu.py +88 -114
- warp/tests/test_noise.py +12 -11
- warp/tests/test_operators.py +16 -20
- warp/tests/test_options.py +11 -11
- warp/tests/test_pinned.py +17 -18
- warp/tests/test_print.py +32 -11
- warp/tests/test_quat.py +275 -129
- warp/tests/test_rand.py +18 -16
- warp/tests/test_reload.py +38 -34
- warp/tests/test_rounding.py +50 -43
- warp/tests/test_runlength_encode.py +168 -20
- warp/tests/test_smoothstep.py +9 -11
- warp/tests/test_snippet.py +143 -0
- warp/tests/test_sparse.py +261 -63
- warp/tests/test_spatial.py +276 -243
- warp/tests/test_streams.py +110 -85
- warp/tests/test_struct.py +268 -63
- warp/tests/test_tape.py +39 -21
- warp/tests/test_torch.py +118 -89
- warp/tests/test_transient_module.py +12 -13
- warp/tests/test_types.py +614 -0
- warp/tests/test_utils.py +494 -0
- warp/tests/test_vec.py +354 -2050
- warp/tests/test_vec_lite.py +73 -0
- warp/tests/test_vec_scalar_ops.py +2099 -0
- warp/tests/test_volume.py +457 -293
- warp/tests/test_volume_write.py +124 -134
- warp/tests/unittest_serial.py +35 -0
- warp/tests/unittest_suites.py +341 -0
- warp/tests/unittest_utils.py +568 -0
- warp/tests/unused_test_misc.py +71 -0
- warp/tests/{test_debug.py → walkthough_debug.py} +3 -17
- warp/thirdparty/appdirs.py +36 -45
- warp/thirdparty/unittest_parallel.py +549 -0
- warp/torch.py +9 -6
- warp/types.py +1089 -366
- warp/utils.py +93 -387
- warp_lang-0.11.0.dist-info/METADATA +238 -0
- warp_lang-0.11.0.dist-info/RECORD +332 -0
- {warp_lang-0.10.1.dist-info → warp_lang-0.11.0.dist-info}/WHEEL +1 -1
- warp/tests/test_all.py +0 -219
- warp/tests/test_array_scan.py +0 -60
- warp/tests/test_base.py +0 -208
- warp/tests/test_unresolved_func.py +0 -7
- warp/tests/test_unresolved_symbol.py +0 -7
- warp_lang-0.10.1.dist-info/METADATA +0 -21
- warp_lang-0.10.1.dist-info/RECORD +0 -188
- /warp/tests/{test_compile_consts_dummy.py → aux_test_compile_consts_dummy.py} +0 -0
- /warp/tests/{test_reference_reference.py → aux_test_reference_reference.py} +0 -0
- /warp/tests/{test_square.py → aux_test_square.py} +0 -0
- {warp_lang-0.10.1.dist-info → warp_lang-0.11.0.dist-info}/LICENSE.md +0 -0
- {warp_lang-0.10.1.dist-info → warp_lang-0.11.0.dist-info}/top_level.txt +0 -0
|
@@ -0,0 +1,299 @@
|
|
|
1
|
+
import warp as wp
|
|
2
|
+
|
|
3
|
+
from warp.fem.space import CollocatedFunctionSpace, SpacePartition
|
|
4
|
+
from warp.fem import cache, utils
|
|
5
|
+
from warp.fem.types import Sample, ElementIndex, NULL_NODE_INDEX
|
|
6
|
+
|
|
7
|
+
from .field import DiscreteField
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
class NodalFieldBase(DiscreteField):
|
|
11
|
+
"""Base class for nodal field and nodal field traces. Does not hold values"""
|
|
12
|
+
|
|
13
|
+
def __init__(self, space: CollocatedFunctionSpace, space_partition: SpacePartition):
|
|
14
|
+
super().__init__(space, space_partition)
|
|
15
|
+
|
|
16
|
+
self.EvalArg = self._make_eval_arg()
|
|
17
|
+
self.ElementEvalArg = self._make_element_eval_arg()
|
|
18
|
+
self.eval_degree = DiscreteField._make_eval_degree(self)
|
|
19
|
+
|
|
20
|
+
self._read_node_value = self._make_read_node_value()
|
|
21
|
+
|
|
22
|
+
self.eval_inner = self._make_eval_inner()
|
|
23
|
+
self.eval_outer = self._make_eval_outer()
|
|
24
|
+
self.eval_grad_inner = self._make_eval_grad_inner(world_space=True)
|
|
25
|
+
self.eval_grad_outer = self._make_eval_grad_outer(world_space=True)
|
|
26
|
+
self.eval_reference_grad_inner = self._make_eval_grad_inner(world_space=False)
|
|
27
|
+
self.eval_reference_grad_outer = self._make_eval_grad_outer(world_space=False)
|
|
28
|
+
self.eval_div_inner = self._make_eval_div_inner()
|
|
29
|
+
self.eval_div_outer = self._make_eval_div_outer()
|
|
30
|
+
|
|
31
|
+
self.set_node_value = self._make_set_node_value()
|
|
32
|
+
|
|
33
|
+
def _make_eval_arg(self):
|
|
34
|
+
@cache.dynamic_struct(suffix=self.name)
|
|
35
|
+
class EvalArg:
|
|
36
|
+
dof_values: wp.array(dtype=self.space.dof_dtype)
|
|
37
|
+
space_arg: self.space.SpaceArg
|
|
38
|
+
topology_arg: self.space.topology.TopologyArg
|
|
39
|
+
partition_arg: self.space_partition.PartitionArg
|
|
40
|
+
|
|
41
|
+
return EvalArg
|
|
42
|
+
|
|
43
|
+
def _make_element_eval_arg(self):
|
|
44
|
+
@cache.dynamic_struct(suffix=self.name)
|
|
45
|
+
class ElementEvalArg:
|
|
46
|
+
elt_arg: self.space.topology.ElementArg
|
|
47
|
+
eval_arg: self.EvalArg
|
|
48
|
+
|
|
49
|
+
return ElementEvalArg
|
|
50
|
+
|
|
51
|
+
def _make_read_node_value(self):
|
|
52
|
+
@cache.dynamic_func(suffix=self.name)
|
|
53
|
+
def read_node_value(args: self.ElementEvalArg, geo_element_index: ElementIndex, node_index_in_elt: int):
|
|
54
|
+
nidx = self.space.topology.element_node_index(
|
|
55
|
+
args.elt_arg, args.eval_arg.topology_arg, geo_element_index, node_index_in_elt
|
|
56
|
+
)
|
|
57
|
+
pidx = self.space_partition.partition_node_index(args.eval_arg.partition_arg, nidx)
|
|
58
|
+
if pidx == NULL_NODE_INDEX:
|
|
59
|
+
return self.space.dtype(0.0)
|
|
60
|
+
|
|
61
|
+
return self.space.dof_mapper.dof_to_value(args.eval_arg.dof_values[pidx])
|
|
62
|
+
|
|
63
|
+
return read_node_value
|
|
64
|
+
|
|
65
|
+
def _make_eval_inner(self):
|
|
66
|
+
NODES_PER_ELEMENT = self.space.topology.NODES_PER_ELEMENT
|
|
67
|
+
|
|
68
|
+
@cache.dynamic_func(suffix=self.name)
|
|
69
|
+
def eval_inner(args: self.ElementEvalArg, s: Sample):
|
|
70
|
+
res = self.space.element_inner_weight(
|
|
71
|
+
args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
|
|
72
|
+
) * self._read_node_value(args, s.element_index, 0)
|
|
73
|
+
for k in range(1, NODES_PER_ELEMENT):
|
|
74
|
+
res += self.space.element_inner_weight(
|
|
75
|
+
args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k
|
|
76
|
+
) * self._read_node_value(args, s.element_index, k)
|
|
77
|
+
return res
|
|
78
|
+
|
|
79
|
+
return eval_inner
|
|
80
|
+
|
|
81
|
+
def _make_eval_grad_inner(self, world_space: bool):
|
|
82
|
+
NODES_PER_ELEMENT = self.space.topology.NODES_PER_ELEMENT
|
|
83
|
+
|
|
84
|
+
if not self.gradient_valid():
|
|
85
|
+
return None
|
|
86
|
+
|
|
87
|
+
@cache.dynamic_func(suffix=self.name)
|
|
88
|
+
def eval_grad_inner_ref_space(args: self.ElementEvalArg, s: Sample):
|
|
89
|
+
res = utils.generalized_outer(
|
|
90
|
+
self._read_node_value(args, s.element_index, 0),
|
|
91
|
+
self.space.element_inner_weight_gradient(
|
|
92
|
+
args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
|
|
93
|
+
),
|
|
94
|
+
)
|
|
95
|
+
for k in range(1, NODES_PER_ELEMENT):
|
|
96
|
+
res += utils.generalized_outer(
|
|
97
|
+
self._read_node_value(args, s.element_index, k),
|
|
98
|
+
self.space.element_inner_weight_gradient(
|
|
99
|
+
args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k
|
|
100
|
+
),
|
|
101
|
+
)
|
|
102
|
+
return res
|
|
103
|
+
|
|
104
|
+
@cache.dynamic_func(suffix=self.name)
|
|
105
|
+
def eval_grad_inner_world_space(args: self.ElementEvalArg, s: Sample):
|
|
106
|
+
grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
|
|
107
|
+
res = eval_grad_inner_ref_space(args, s)
|
|
108
|
+
return utils.apply_right(res, grad_transform)
|
|
109
|
+
|
|
110
|
+
return eval_grad_inner_world_space if world_space else eval_grad_inner_ref_space
|
|
111
|
+
|
|
112
|
+
def _make_eval_div_inner(self):
|
|
113
|
+
NODES_PER_ELEMENT = self.space.topology.NODES_PER_ELEMENT
|
|
114
|
+
|
|
115
|
+
if not self.divergence_valid():
|
|
116
|
+
return None
|
|
117
|
+
|
|
118
|
+
@cache.dynamic_func(suffix=self.name)
|
|
119
|
+
def eval_div_inner(args: self.ElementEvalArg, s: Sample):
|
|
120
|
+
grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
|
|
121
|
+
|
|
122
|
+
res = utils.generalized_inner(
|
|
123
|
+
self._read_node_value(args, s.element_index, 0),
|
|
124
|
+
utils.apply_right(
|
|
125
|
+
self.space.element_inner_weight_gradient(
|
|
126
|
+
args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
|
|
127
|
+
),
|
|
128
|
+
grad_transform,
|
|
129
|
+
),
|
|
130
|
+
)
|
|
131
|
+
|
|
132
|
+
for k in range(1, NODES_PER_ELEMENT):
|
|
133
|
+
res += utils.generalized_inner(
|
|
134
|
+
self._read_node_value(args, s.element_index, k),
|
|
135
|
+
utils.apply_right(
|
|
136
|
+
self.space.element_inner_weight_gradient(
|
|
137
|
+
args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k
|
|
138
|
+
),
|
|
139
|
+
grad_transform,
|
|
140
|
+
),
|
|
141
|
+
)
|
|
142
|
+
return res
|
|
143
|
+
|
|
144
|
+
return eval_div_inner
|
|
145
|
+
|
|
146
|
+
def _make_eval_outer(self):
|
|
147
|
+
NODES_PER_ELEMENT = self.space.topology.NODES_PER_ELEMENT
|
|
148
|
+
|
|
149
|
+
@cache.dynamic_func(suffix=self.name)
|
|
150
|
+
def eval_outer(args: self.ElementEvalArg, s: Sample):
|
|
151
|
+
res = self.space.element_outer_weight(
|
|
152
|
+
args.elt_arg,
|
|
153
|
+
args.eval_arg.space_arg,
|
|
154
|
+
s.element_index,
|
|
155
|
+
s.element_coords,
|
|
156
|
+
0,
|
|
157
|
+
) * self._read_node_value(args, s.element_index, 0)
|
|
158
|
+
for k in range(1, NODES_PER_ELEMENT):
|
|
159
|
+
res += self.space.element_outer_weight(
|
|
160
|
+
args.elt_arg,
|
|
161
|
+
args.eval_arg.space_arg,
|
|
162
|
+
s.element_index,
|
|
163
|
+
s.element_coords,
|
|
164
|
+
k,
|
|
165
|
+
) * self._read_node_value(args, s.element_index, k)
|
|
166
|
+
return res
|
|
167
|
+
|
|
168
|
+
return eval_outer
|
|
169
|
+
|
|
170
|
+
def _make_eval_grad_outer(self, world_space: bool):
|
|
171
|
+
NODES_PER_ELEMENT = self.space.topology.NODES_PER_ELEMENT
|
|
172
|
+
|
|
173
|
+
if not self.gradient_valid():
|
|
174
|
+
return None
|
|
175
|
+
|
|
176
|
+
@cache.dynamic_func(suffix=self.name)
|
|
177
|
+
def eval_grad_outer_ref_space(args: self.ElementEvalArg, s: Sample):
|
|
178
|
+
res = utils.generalized_outer(
|
|
179
|
+
self._read_node_value(args, s.element_index, 0),
|
|
180
|
+
self.space.element_outer_weight_gradient(
|
|
181
|
+
args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
|
|
182
|
+
),
|
|
183
|
+
)
|
|
184
|
+
for k in range(1, NODES_PER_ELEMENT):
|
|
185
|
+
res += utils.generalized_outer(
|
|
186
|
+
self._read_node_value(args, s.element_index, k),
|
|
187
|
+
self.space.element_outer_weight_gradient(
|
|
188
|
+
args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k
|
|
189
|
+
),
|
|
190
|
+
)
|
|
191
|
+
return res
|
|
192
|
+
|
|
193
|
+
@cache.dynamic_func(suffix=self.name)
|
|
194
|
+
def eval_grad_outer_world_space(args: self.ElementEvalArg, s: Sample):
|
|
195
|
+
grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
|
|
196
|
+
res = eval_grad_outer_ref_space(args, s)
|
|
197
|
+
return utils.apply_right(res, grad_transform)
|
|
198
|
+
|
|
199
|
+
return eval_grad_outer_world_space if world_space else eval_grad_outer_ref_space
|
|
200
|
+
|
|
201
|
+
def _make_eval_div_outer(self):
|
|
202
|
+
NODES_PER_ELEMENT = self.space.topology.NODES_PER_ELEMENT
|
|
203
|
+
|
|
204
|
+
if not self.divergence_valid():
|
|
205
|
+
return None
|
|
206
|
+
|
|
207
|
+
@cache.dynamic_func(suffix=self.name)
|
|
208
|
+
def eval_div_outer(args: self.ElementEvalArg, s: Sample):
|
|
209
|
+
grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
|
|
210
|
+
|
|
211
|
+
res = utils.generalized_inner(
|
|
212
|
+
self._read_node_value(args, s.element_index, 0),
|
|
213
|
+
utils.apply_right(
|
|
214
|
+
self.space.element_outer_weight_gradient(
|
|
215
|
+
args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
|
|
216
|
+
),
|
|
217
|
+
grad_transform,
|
|
218
|
+
),
|
|
219
|
+
)
|
|
220
|
+
for k in range(1, NODES_PER_ELEMENT):
|
|
221
|
+
res += utils.generalized_inner(
|
|
222
|
+
self._read_node_value(args, s.element_index, k),
|
|
223
|
+
utils.apply_right(
|
|
224
|
+
self.space.element_outer_weight_gradient(
|
|
225
|
+
args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k
|
|
226
|
+
),
|
|
227
|
+
grad_transform,
|
|
228
|
+
),
|
|
229
|
+
)
|
|
230
|
+
return res
|
|
231
|
+
|
|
232
|
+
return eval_div_outer
|
|
233
|
+
|
|
234
|
+
def _make_set_node_value(self):
|
|
235
|
+
@cache.dynamic_func(suffix=self.name)
|
|
236
|
+
def set_node_value(args: self.EvalArg, partition_node_index: int, value: self.space.dtype):
|
|
237
|
+
args.dof_values[partition_node_index] = self.space.dof_mapper.value_to_dof(value)
|
|
238
|
+
|
|
239
|
+
return set_node_value
|
|
240
|
+
|
|
241
|
+
|
|
242
|
+
class NodalField(NodalFieldBase):
|
|
243
|
+
"""A field holding values for all degrees of freedom at each node of the underlying function space partition
|
|
244
|
+
|
|
245
|
+
See also: warp.fem.space.CollocatedFunctionSpace.make_field
|
|
246
|
+
"""
|
|
247
|
+
|
|
248
|
+
def __init__(self, space: CollocatedFunctionSpace, space_partition: SpacePartition):
|
|
249
|
+
if space.topology != space_partition.space_topology:
|
|
250
|
+
raise ValueError("Incompatible space and space partition topologies")
|
|
251
|
+
|
|
252
|
+
super().__init__(space, space_partition)
|
|
253
|
+
|
|
254
|
+
self._dof_values = wp.zeros(n=self.space_partition.node_count(), dtype=self.dof_dtype)
|
|
255
|
+
|
|
256
|
+
def eval_arg_value(self, device):
|
|
257
|
+
arg = self.EvalArg()
|
|
258
|
+
arg.dof_values = self._dof_values.to(device)
|
|
259
|
+
arg.space_arg = self.space.space_arg_value(device)
|
|
260
|
+
arg.partition_arg = self.space_partition.partition_arg_value(device)
|
|
261
|
+
arg.topology_arg = self.space.topology.topo_arg_value(device)
|
|
262
|
+
|
|
263
|
+
return arg
|
|
264
|
+
|
|
265
|
+
@property
|
|
266
|
+
def dof_values(self) -> wp.array:
|
|
267
|
+
"""Returns a warp array containing the values at all degrees of freedom of the underlying space partition"""
|
|
268
|
+
return self._dof_values
|
|
269
|
+
|
|
270
|
+
@dof_values.setter
|
|
271
|
+
def dof_values(self, values):
|
|
272
|
+
"""Sets the degrees-of-freedom values
|
|
273
|
+
|
|
274
|
+
Args:
|
|
275
|
+
values: Array that is convertible to a warp array of length ``self.space_partition.node_count()`` and data type ``self.space.dof_dtype``
|
|
276
|
+
"""
|
|
277
|
+
|
|
278
|
+
if isinstance(values, wp.array):
|
|
279
|
+
self._dof_values = values
|
|
280
|
+
else:
|
|
281
|
+
self._dof_values = wp.array(values, dtype=self.dof_dtype)
|
|
282
|
+
|
|
283
|
+
class Trace(NodalFieldBase):
|
|
284
|
+
def __init__(self, field):
|
|
285
|
+
self._field = field
|
|
286
|
+
super().__init__(field.space.trace(), field.space_partition)
|
|
287
|
+
|
|
288
|
+
def eval_arg_value(self, device):
|
|
289
|
+
arg = self.EvalArg()
|
|
290
|
+
arg.dof_values = self._field.dof_values.to(device)
|
|
291
|
+
arg.space_arg = self.space.space_arg_value(device)
|
|
292
|
+
arg.partition_arg = self.space_partition.partition_arg_value(device)
|
|
293
|
+
arg.topology_arg = self.space.topology.topo_arg_value(device)
|
|
294
|
+
|
|
295
|
+
return arg
|
|
296
|
+
|
|
297
|
+
def trace(self) -> Trace:
|
|
298
|
+
trace_field = NodalField.Trace(self)
|
|
299
|
+
return trace_field
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
from warp.fem.space import SpaceRestriction
|
|
2
|
+
from .field import DiscreteField
|
|
3
|
+
|
|
4
|
+
|
|
5
|
+
class FieldRestriction:
|
|
6
|
+
"""Restriction of a discrete field to a given GeometryDomain"""
|
|
7
|
+
|
|
8
|
+
def __init__(self, space_restriction: SpaceRestriction, field: DiscreteField):
|
|
9
|
+
if field.space.dimension - 1 == space_restriction.space_topology.dimension:
|
|
10
|
+
field = field.trace()
|
|
11
|
+
|
|
12
|
+
if field.space.dimension != space_restriction.space_topology.dimension:
|
|
13
|
+
raise ValueError("Incompatible space and field dimensions")
|
|
14
|
+
|
|
15
|
+
if field.space.topology != space_restriction.space_topology:
|
|
16
|
+
raise ValueError("Incompatible field and space restriction topologies")
|
|
17
|
+
|
|
18
|
+
self.space_restriction = space_restriction
|
|
19
|
+
self.domain = self.space_restriction.domain
|
|
20
|
+
self.field = field
|
|
21
|
+
self.space = self.field.space
|
warp/fem/field/test.py
ADDED
|
@@ -0,0 +1,181 @@
|
|
|
1
|
+
import warp as wp
|
|
2
|
+
|
|
3
|
+
from warp.fem.space import SpaceRestriction, FunctionSpace
|
|
4
|
+
from warp.fem.types import Sample, get_node_index_in_element
|
|
5
|
+
from warp.fem import utils, cache
|
|
6
|
+
|
|
7
|
+
from .field import SpaceField
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
class TestField(SpaceField):
|
|
11
|
+
"""Field defined over a space restriction that can be used as a test function.
|
|
12
|
+
|
|
13
|
+
In order to reuse computations, it is possible to define the test field using a SpaceRestriction
|
|
14
|
+
defined for a different value type than the test function value type, as long as the node topology is similar.
|
|
15
|
+
"""
|
|
16
|
+
|
|
17
|
+
def __init__(self, space_restriction: SpaceRestriction, space: FunctionSpace):
|
|
18
|
+
if space_restriction.domain.dimension == space.dimension - 1:
|
|
19
|
+
space = space.trace()
|
|
20
|
+
|
|
21
|
+
if space_restriction.domain.dimension != space.dimension:
|
|
22
|
+
raise ValueError("Incompatible space and domain dimensions")
|
|
23
|
+
|
|
24
|
+
if space.topology != space_restriction.space_topology:
|
|
25
|
+
raise ValueError("Incompatible space and space partition topologies")
|
|
26
|
+
|
|
27
|
+
super().__init__(space, space_restriction.space_partition)
|
|
28
|
+
|
|
29
|
+
self.space_restriction = space_restriction
|
|
30
|
+
self.domain = self.space_restriction.domain
|
|
31
|
+
|
|
32
|
+
self.EvalArg = self.space.SpaceArg
|
|
33
|
+
self.ElementEvalArg = self._make_element_eval_arg()
|
|
34
|
+
|
|
35
|
+
self.eval_degree = self._make_eval_degree()
|
|
36
|
+
self.eval_inner = self._make_eval_inner()
|
|
37
|
+
self.eval_grad_inner = self._make_eval_grad_inner()
|
|
38
|
+
self.eval_div_inner = self._make_eval_div_inner()
|
|
39
|
+
self.eval_outer = self._make_eval_outer()
|
|
40
|
+
self.eval_grad_outer = self._make_eval_grad_outer()
|
|
41
|
+
self.eval_div_outer = self._make_eval_div_outer()
|
|
42
|
+
self.at_node = self._make_at_node()
|
|
43
|
+
|
|
44
|
+
@property
|
|
45
|
+
def name(self) -> str:
|
|
46
|
+
return self.space.name + "Test"
|
|
47
|
+
|
|
48
|
+
def eval_arg_value(self, device) -> wp.codegen.StructInstance:
|
|
49
|
+
return self.space.space_arg_value(device)
|
|
50
|
+
|
|
51
|
+
def _make_element_eval_arg(self):
|
|
52
|
+
from warp.fem import cache
|
|
53
|
+
|
|
54
|
+
@cache.dynamic_struct(suffix=self.name)
|
|
55
|
+
class ElementEvalArg:
|
|
56
|
+
elt_arg: self.domain.ElementArg
|
|
57
|
+
eval_arg: self.EvalArg
|
|
58
|
+
|
|
59
|
+
return ElementEvalArg
|
|
60
|
+
|
|
61
|
+
def _make_eval_inner(self):
|
|
62
|
+
@cache.dynamic_func(suffix=self.name)
|
|
63
|
+
def eval_test_inner(args: self.ElementEvalArg, s: Sample):
|
|
64
|
+
weight = self.space.element_inner_weight(
|
|
65
|
+
args.elt_arg,
|
|
66
|
+
args.eval_arg,
|
|
67
|
+
s.element_index,
|
|
68
|
+
s.element_coords,
|
|
69
|
+
get_node_index_in_element(s.test_dof),
|
|
70
|
+
)
|
|
71
|
+
return weight * self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.test_dof)
|
|
72
|
+
|
|
73
|
+
return eval_test_inner
|
|
74
|
+
|
|
75
|
+
def _make_eval_grad_inner(self):
|
|
76
|
+
if not self.gradient_valid():
|
|
77
|
+
return None
|
|
78
|
+
|
|
79
|
+
@cache.dynamic_func(suffix=self.name)
|
|
80
|
+
def eval_nabla_test_inner(args: self.ElementEvalArg, s: Sample):
|
|
81
|
+
nabla_weight = self.space.element_inner_weight_gradient(
|
|
82
|
+
args.elt_arg,
|
|
83
|
+
args.eval_arg,
|
|
84
|
+
s.element_index,
|
|
85
|
+
s.element_coords,
|
|
86
|
+
get_node_index_in_element(s.test_dof),
|
|
87
|
+
)
|
|
88
|
+
grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
|
|
89
|
+
return utils.generalized_outer(
|
|
90
|
+
self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.test_dof),
|
|
91
|
+
utils.apply_right(nabla_weight, grad_transform),
|
|
92
|
+
)
|
|
93
|
+
|
|
94
|
+
return eval_nabla_test_inner
|
|
95
|
+
|
|
96
|
+
def _make_eval_div_inner(self):
|
|
97
|
+
if not self.divergence_valid():
|
|
98
|
+
return None
|
|
99
|
+
|
|
100
|
+
@cache.dynamic_func(suffix=self.name)
|
|
101
|
+
def eval_div_test_inner(args: self.ElementEvalArg, s: Sample):
|
|
102
|
+
nabla_weight = self.space.element_inner_weight_gradient(
|
|
103
|
+
args.elt_arg,
|
|
104
|
+
args.eval_arg,
|
|
105
|
+
s.element_index,
|
|
106
|
+
s.element_coords,
|
|
107
|
+
get_node_index_in_element(s.test_dof),
|
|
108
|
+
)
|
|
109
|
+
grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
|
|
110
|
+
return utils.generalized_inner(
|
|
111
|
+
self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.test_dof),
|
|
112
|
+
utils.apply_right(nabla_weight, grad_transform),
|
|
113
|
+
)
|
|
114
|
+
|
|
115
|
+
return eval_div_test_inner
|
|
116
|
+
|
|
117
|
+
def _make_eval_outer(self):
|
|
118
|
+
@cache.dynamic_func(suffix=self.name)
|
|
119
|
+
def eval_test_outer(args: self.ElementEvalArg, s: Sample):
|
|
120
|
+
weight = self.space.element_outer_weight(
|
|
121
|
+
args.elt_arg,
|
|
122
|
+
args.eval_arg,
|
|
123
|
+
s.element_index,
|
|
124
|
+
s.element_coords,
|
|
125
|
+
get_node_index_in_element(s.test_dof),
|
|
126
|
+
)
|
|
127
|
+
return weight * self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.test_dof)
|
|
128
|
+
|
|
129
|
+
return eval_test_outer
|
|
130
|
+
|
|
131
|
+
def _make_eval_grad_outer(self):
|
|
132
|
+
if not self.gradient_valid():
|
|
133
|
+
return None
|
|
134
|
+
|
|
135
|
+
@cache.dynamic_func(suffix=self.name)
|
|
136
|
+
def eval_nabla_test_outer(args: self.ElementEvalArg, s: Sample):
|
|
137
|
+
nabla_weight = self.space.element_outer_weight_gradient(
|
|
138
|
+
args.elt_arg,
|
|
139
|
+
args.eval_arg,
|
|
140
|
+
s.element_index,
|
|
141
|
+
s.element_coords,
|
|
142
|
+
get_node_index_in_element(s.test_dof),
|
|
143
|
+
)
|
|
144
|
+
grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
|
|
145
|
+
return utils.generalized_outer(
|
|
146
|
+
self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.test_dof),
|
|
147
|
+
utils.apply_right(nabla_weight, grad_transform),
|
|
148
|
+
)
|
|
149
|
+
|
|
150
|
+
return eval_nabla_test_outer
|
|
151
|
+
|
|
152
|
+
def _make_eval_div_outer(self):
|
|
153
|
+
if not self.divergence_valid():
|
|
154
|
+
return None
|
|
155
|
+
|
|
156
|
+
@cache.dynamic_func(suffix=self.name)
|
|
157
|
+
def eval_div_test_outer(args: self.ElementEvalArg, s: Sample):
|
|
158
|
+
nabla_weight = self.space.element_outer_weight_gradient(
|
|
159
|
+
args.elt_arg,
|
|
160
|
+
args.eval_arg,
|
|
161
|
+
s.element_index,
|
|
162
|
+
s.element_coords,
|
|
163
|
+
get_node_index_in_element(s.test_dof),
|
|
164
|
+
)
|
|
165
|
+
grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
|
|
166
|
+
return utils.generalized_inner(
|
|
167
|
+
self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.test_dof),
|
|
168
|
+
utils.apply_right(nabla_weight, grad_transform),
|
|
169
|
+
)
|
|
170
|
+
|
|
171
|
+
return eval_div_test_outer
|
|
172
|
+
|
|
173
|
+
def _make_at_node(self):
|
|
174
|
+
@cache.dynamic_func(suffix=self.name)
|
|
175
|
+
def at_node(args: self.ElementEvalArg, s: Sample):
|
|
176
|
+
node_coords = self.space.node_coords_in_element(
|
|
177
|
+
args.elt_arg, args.eval_arg, s.element_index, get_node_index_in_element(s.test_dof)
|
|
178
|
+
)
|
|
179
|
+
return Sample(s.element_index, node_coords, s.qp_index, s.qp_weight, s.test_dof, s.trial_dof)
|
|
180
|
+
|
|
181
|
+
return at_node
|
warp/fem/field/trial.py
ADDED
|
@@ -0,0 +1,183 @@
|
|
|
1
|
+
import warp as wp
|
|
2
|
+
from warp.fem.domain import GeometryDomain
|
|
3
|
+
from warp.fem.space import FunctionSpace, SpacePartition
|
|
4
|
+
from warp.fem.types import Sample, get_node_index_in_element
|
|
5
|
+
from warp.fem import utils, cache
|
|
6
|
+
|
|
7
|
+
from .field import SpaceField
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
class TrialField(SpaceField):
|
|
11
|
+
"""Field defined over a domain that can be used as a trial function"""
|
|
12
|
+
|
|
13
|
+
def __init__(
|
|
14
|
+
self,
|
|
15
|
+
space: FunctionSpace,
|
|
16
|
+
space_partition: SpacePartition,
|
|
17
|
+
domain: GeometryDomain,
|
|
18
|
+
):
|
|
19
|
+
if domain.dimension == space.dimension - 1:
|
|
20
|
+
space = space.trace()
|
|
21
|
+
|
|
22
|
+
if domain.dimension != space.dimension:
|
|
23
|
+
raise ValueError("Incompatible space and domain dimensions")
|
|
24
|
+
|
|
25
|
+
if not space.topology.is_derived_from(space_partition.space_topology):
|
|
26
|
+
raise ValueError("Incompatible space and space partition topologies")
|
|
27
|
+
|
|
28
|
+
super().__init__(space, space_partition)
|
|
29
|
+
|
|
30
|
+
self.domain = domain
|
|
31
|
+
|
|
32
|
+
self.EvalArg = self.space.SpaceArg
|
|
33
|
+
self.ElementEvalArg = self._make_element_eval_arg()
|
|
34
|
+
|
|
35
|
+
self.eval_degree = self._make_eval_degree()
|
|
36
|
+
self.eval_inner = self._make_eval_inner()
|
|
37
|
+
self.eval_grad_inner = self._make_eval_grad_inner()
|
|
38
|
+
self.eval_div_inner = self._make_eval_div_inner()
|
|
39
|
+
self.eval_outer = self._make_eval_outer()
|
|
40
|
+
self.eval_grad_outer = self._make_eval_grad_outer()
|
|
41
|
+
self.eval_div_outer = self._make_eval_div_outer()
|
|
42
|
+
self.at_node = self._make_at_node()
|
|
43
|
+
|
|
44
|
+
def partition_node_count(self) -> int:
|
|
45
|
+
"""Returns the number of nodes in the associated space topology partition"""
|
|
46
|
+
return self.space_partition.node_count()
|
|
47
|
+
|
|
48
|
+
@property
|
|
49
|
+
def name(self) -> str:
|
|
50
|
+
return self.space.name + "Trial"
|
|
51
|
+
|
|
52
|
+
def eval_arg_value(self, device) -> wp.codegen.StructInstance:
|
|
53
|
+
return self.space.space_arg_value(device)
|
|
54
|
+
|
|
55
|
+
def _make_element_eval_arg(self):
|
|
56
|
+
@cache.dynamic_struct(suffix=self.name)
|
|
57
|
+
class ElementEvalArg:
|
|
58
|
+
elt_arg: self.domain.ElementArg
|
|
59
|
+
eval_arg: self.EvalArg
|
|
60
|
+
|
|
61
|
+
return ElementEvalArg
|
|
62
|
+
|
|
63
|
+
def _make_eval_inner(self):
|
|
64
|
+
@cache.dynamic_func(suffix=self.name)
|
|
65
|
+
def eval_trial_inner(args: self.ElementEvalArg, s: Sample):
|
|
66
|
+
weight = self.space.element_inner_weight(
|
|
67
|
+
args.elt_arg,
|
|
68
|
+
args.eval_arg,
|
|
69
|
+
s.element_index,
|
|
70
|
+
s.element_coords,
|
|
71
|
+
get_node_index_in_element(s.trial_dof),
|
|
72
|
+
)
|
|
73
|
+
return weight * self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof)
|
|
74
|
+
|
|
75
|
+
return eval_trial_inner
|
|
76
|
+
|
|
77
|
+
def _make_eval_grad_inner(self):
|
|
78
|
+
if not self.gradient_valid():
|
|
79
|
+
return None
|
|
80
|
+
|
|
81
|
+
@cache.dynamic_func(suffix=self.name)
|
|
82
|
+
def eval_nabla_trial_inner(args: self.ElementEvalArg, s: Sample):
|
|
83
|
+
nabla_weight = self.space.element_inner_weight_gradient(
|
|
84
|
+
args.elt_arg,
|
|
85
|
+
args.eval_arg,
|
|
86
|
+
s.element_index,
|
|
87
|
+
s.element_coords,
|
|
88
|
+
get_node_index_in_element(s.trial_dof),
|
|
89
|
+
)
|
|
90
|
+
grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
|
|
91
|
+
return utils.generalized_outer(
|
|
92
|
+
self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof),
|
|
93
|
+
utils.apply_right(nabla_weight, grad_transform),
|
|
94
|
+
)
|
|
95
|
+
|
|
96
|
+
return eval_nabla_trial_inner
|
|
97
|
+
|
|
98
|
+
def _make_eval_div_inner(self):
|
|
99
|
+
if not self.divergence_valid():
|
|
100
|
+
return None
|
|
101
|
+
|
|
102
|
+
@cache.dynamic_func(suffix=self.name)
|
|
103
|
+
def eval_div_trial_inner(args: self.ElementEvalArg, s: Sample):
|
|
104
|
+
nabla_weight = self.space.element_inner_weight_gradient(
|
|
105
|
+
args.elt_arg,
|
|
106
|
+
args.eval_arg,
|
|
107
|
+
s.element_index,
|
|
108
|
+
s.element_coords,
|
|
109
|
+
get_node_index_in_element(s.trial_dof),
|
|
110
|
+
)
|
|
111
|
+
grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
|
|
112
|
+
return utils.generalized_inner(
|
|
113
|
+
self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof),
|
|
114
|
+
utils.apply_right(nabla_weight, grad_transform),
|
|
115
|
+
)
|
|
116
|
+
|
|
117
|
+
return eval_div_trial_inner
|
|
118
|
+
|
|
119
|
+
def _make_eval_outer(self):
|
|
120
|
+
@cache.dynamic_func(suffix=self.name)
|
|
121
|
+
def eval_trial_outer(args: self.ElementEvalArg, s: Sample):
|
|
122
|
+
weight = self.space.element_outer_weight(
|
|
123
|
+
args.elt_arg,
|
|
124
|
+
args.eval_arg,
|
|
125
|
+
s.element_index,
|
|
126
|
+
s.element_coords,
|
|
127
|
+
get_node_index_in_element(s.trial_dof),
|
|
128
|
+
)
|
|
129
|
+
return weight * self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof)
|
|
130
|
+
|
|
131
|
+
return eval_trial_outer
|
|
132
|
+
|
|
133
|
+
def _make_eval_grad_outer(self):
|
|
134
|
+
if not self.gradient_valid():
|
|
135
|
+
return None
|
|
136
|
+
|
|
137
|
+
@cache.dynamic_func(suffix=self.name)
|
|
138
|
+
def eval_nabla_trial_outer(args: self.ElementEvalArg, s: Sample):
|
|
139
|
+
nabla_weight = self.space.element_outer_weight_gradient(
|
|
140
|
+
args.elt_arg,
|
|
141
|
+
args.eval_arg,
|
|
142
|
+
s.element_index,
|
|
143
|
+
s.element_coords,
|
|
144
|
+
get_node_index_in_element(s.trial_dof),
|
|
145
|
+
)
|
|
146
|
+
grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
|
|
147
|
+
return utils.generalized_outer(
|
|
148
|
+
self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof),
|
|
149
|
+
utils.apply_right(nabla_weight, grad_transform),
|
|
150
|
+
)
|
|
151
|
+
|
|
152
|
+
return eval_nabla_trial_outer
|
|
153
|
+
|
|
154
|
+
def _make_eval_div_outer(self):
|
|
155
|
+
if not self.divergence_valid():
|
|
156
|
+
return None
|
|
157
|
+
|
|
158
|
+
@cache.dynamic_func(suffix=self.name)
|
|
159
|
+
def eval_div_trial_outer(args: self.ElementEvalArg, s: Sample):
|
|
160
|
+
nabla_weight = self.space.element_outer_weight_gradient(
|
|
161
|
+
args.elt_arg,
|
|
162
|
+
args.eval_arg,
|
|
163
|
+
s.element_index,
|
|
164
|
+
s.element_coords,
|
|
165
|
+
get_node_index_in_element(s.trial_dof),
|
|
166
|
+
)
|
|
167
|
+
grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
|
|
168
|
+
return utils.generalized_inner(
|
|
169
|
+
self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof),
|
|
170
|
+
utils.apply_right(nabla_weight, grad_transform),
|
|
171
|
+
)
|
|
172
|
+
|
|
173
|
+
return eval_div_trial_outer
|
|
174
|
+
|
|
175
|
+
def _make_at_node(self):
|
|
176
|
+
@cache.dynamic_func(suffix=self.name)
|
|
177
|
+
def at_node(args: self.ElementEvalArg, s: Sample):
|
|
178
|
+
node_coords = self.space.node_coords_in_element(
|
|
179
|
+
args.elt_arg, args.eval_arg, s.element_index, get_node_index_in_element(s.trial_dof)
|
|
180
|
+
)
|
|
181
|
+
return Sample(s.element_index, node_coords, s.qp_index, s.qp_weight, s.test_dof, s.trial_dof)
|
|
182
|
+
|
|
183
|
+
return at_node
|