warp-lang 1.6.1__py3-none-macosx_10_13_universal2.whl → 1.7.0__py3-none-macosx_10_13_universal2.whl

warp/fem/adaptivity.py

@@ -1,3 +1,18 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from typing import List, Optional, Tuple
 
 import numpy as np
@@ -353,7 +368,7 @@ def _fill_graded_cells(
     level = int(coarse_level[cell])
     refinement = wp.min(1, coarse_refinement[cell])
 
-    count = wp.select(refinement > 0, 1, 8)
+    count = wp.where(refinement > 0, 8, 1)
     offset = wp.atomic_sub(fine_count, 0, count) - count
 
     f_level = level - refinement
@@ -388,7 +403,7 @@ def _sample_refinement(
             if sampled_level >= 0:
                 min_level = wp.min(sampled_level, min_level)
 
-    return wp.select(min_level < level_count, -1, cur_level - wp.clamp(min_level, 0, cur_level))
+    return wp.where(min_level < level_count, cur_level - wp.clamp(min_level, 0, cur_level), -1)
 
 
 @integrand
@@ -416,7 +431,7 @@ def _count_refined_voxels(
 
     coarse_refinement[cell] = wp.int8(refinement)
     if refinement >= 0:
-        wp.atomic_add(fine_count, 0, wp.select(refinement > 0, 1, 8))
+        wp.atomic_add(fine_count, 0, wp.where(refinement > 0, 8, 1))
 
 
 @wp.kernel
@@ -434,7 +449,7 @@ def _fill_refined_voxels(
     refinement = wp.min(1, int(coarse_refinement[cell]))
 
     if refinement >= 0:
-        count = wp.select(refinement > 0, 1, 8)
+        count = wp.where(refinement > 0, 8, 1)
         offset = wp.atomic_sub(fine_count, 0, count) - count
 
         f_level = level - refinement

warp/fem/cache.py

@@ -1,3 +1,18 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import ast
 import bisect
 import re

warp/fem/dirichlet.py

@@ -1,3 +1,18 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from typing import Any, Optional
 
 import warp as wp

warp/fem/domain.py

@@ -1,3 +1,18 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from typing import Any, Optional, Set, Union
 
 import warp as wp

warp/fem/field/__init__.py

@@ -1,3 +1,18 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from typing import Optional, Union
 
 from warp.fem.domain import Cells, GeometryDomain

warp/fem/field/field.py

@@ -1,3 +1,18 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from typing import Any, Dict, Optional, Set
 
 import warp as wp

warp/fem/field/nodal_field.py

@@ -1,3 +1,18 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from typing import Any
 
 import warp as wp
@@ -68,22 +83,15 @@ class NodalFieldBase(DiscreteField):
         @cache.dynamic_func(suffix=self.name)
         def eval_inner(args: self.ElementEvalArg, s: Sample):
             local_value_map = self.space.local_value_map_inner(args.elt_arg, s.element_index, s.element_coords)
-            res = self.space.space_value(
-                self._read_node_value(args, s.element_index, 0),
-                self.space.element_inner_weight(
-                    args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
-                ),
-                local_value_map,
-            )
-
             node_count = self.space.topology.element_node_count(
                 args.elt_arg, args.eval_arg.topology_arg, s.element_index
             )
-            for k in range(1, node_count):
+            res = self.space.dtype(0.0)
+            for k in range(node_count):
                 res += self.space.space_value(
                     self._read_node_value(args, s.element_index, k),
                     self.space.element_inner_weight(
-                        args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k
+                        args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k, s.qp_index
                     ),
                     local_value_map,
                 )
@@ -98,24 +106,16 @@ class NodalFieldBase(DiscreteField):
         @cache.dynamic_func(suffix=self.name)
         def eval_grad_inner(args: self.ElementEvalArg, s: Sample, grad_transform: Any):
             local_value_map = self.space.local_value_map_inner(args.elt_arg, s.element_index, s.element_coords)
-
-            res = self.space.space_gradient(
-                self._read_node_value(args, s.element_index, 0),
-                self.space.element_inner_weight_gradient(
-                    args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
-                ),
-                local_value_map,
-                grad_transform,
-            )
-
             node_count = self.space.topology.element_node_count(
                 args.elt_arg, args.eval_arg.topology_arg, s.element_index
             )
-            for k in range(1, node_count):
+
+            res = self.gradient_dtype(0.0)
+            for k in range(node_count):
                 res += self.space.space_gradient(
                     self._read_node_value(args, s.element_index, k),
                     self.space.element_inner_weight_gradient(
-                        args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k
+                        args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k, s.qp_index
                     ),
                     local_value_map,
                     grad_transform,
@@ -142,24 +142,16 @@ class NodalFieldBase(DiscreteField):
         def eval_div_inner(args: self.ElementEvalArg, s: Sample):
             grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
             local_value_map = self.space.local_value_map_inner(args.elt_arg, s.element_index, s.element_coords)
-
-            res = self.space.space_divergence(
-                self._read_node_value(args, s.element_index, 0),
-                self.space.element_inner_weight_gradient(
-                    args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
-                ),
-                local_value_map,
-                grad_transform,
-            )
-
             node_count = self.space.topology.element_node_count(
                 args.elt_arg, args.eval_arg.topology_arg, s.element_index
             )
-            for k in range(1, node_count):
+
+            res = self.divergence_dtype(0.0)
+            for k in range(node_count):
                 res += self.space.space_divergence(
                     self._read_node_value(args, s.element_index, k),
                     self.space.element_inner_weight_gradient(
-                        args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k
+                        args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k, s.qp_index
                     ),
                     local_value_map,
                     grad_transform,
@@ -172,23 +164,16 @@ class NodalFieldBase(DiscreteField):
         @cache.dynamic_func(suffix=self.name)
         def eval_outer(args: self.ElementEvalArg, s: Sample):
             local_value_map = self.space.local_value_map_outer(args.elt_arg, s.element_index, s.element_coords)
-            res = self.space.space_value(
-                self._read_node_value(args, s.element_index, 0),
-                self.space.element_outer_weight(
-                    args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
-                ),
-                local_value_map,
-            )
-
             node_count = self.space.topology.element_node_count(
                 args.elt_arg, args.eval_arg.topology_arg, s.element_index
             )
 
-            for k in range(1, node_count):
+            res = self.dtype(0.0)
+            for k in range(node_count):
                 res += self.space.space_value(
                     self._read_node_value(args, s.element_index, k),
                     self.space.element_outer_weight(
-                        args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k
+                        args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k, s.qp_index
                     ),
                     local_value_map,
                 )
@@ -203,24 +188,16 @@ class NodalFieldBase(DiscreteField):
         @cache.dynamic_func(suffix=self.name)
         def eval_grad_outer(args: self.ElementEvalArg, s: Sample, grad_transform: Any):
             local_value_map = self.space.local_value_map_outer(args.elt_arg, s.element_index, s.element_coords)
-
-            res = self.space.space_gradient(
-                self._read_node_value(args, s.element_index, 0),
-                self.space.element_outer_weight_gradient(
-                    args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
-                ),
-                local_value_map,
-                grad_transform,
-            )
-
             node_count = self.space.topology.element_node_count(
                 args.elt_arg, args.eval_arg.topology_arg, s.element_index
             )
-            for k in range(1, node_count):
+
+            res = self.gradient_dtype(0.0)
+            for k in range(node_count):
                 res += self.space.space_gradient(
                     self._read_node_value(args, s.element_index, k),
                     self.space.element_outer_weight_gradient(
-                        args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k
+                        args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k, s.qp_index
                     ),
                     local_value_map,
                     grad_transform,
@@ -247,24 +224,16 @@ class NodalFieldBase(DiscreteField):
         def eval_div_outer(args: self.ElementEvalArg, s: Sample):
             grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
             local_value_map = self.space.local_value_map_outer(args.elt_arg, s.element_index, s.element_coords)
-
-            res = self.space.space_divergence(
-                self._read_node_value(args, s.element_index, 0),
-                self.space.element_outer_weight_gradient(
-                    args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
-                ),
-                local_value_map,
-                grad_transform,
-            )
-
             node_count = self.space.topology.element_node_count(
                 args.elt_arg, args.eval_arg.topology_arg, s.element_index
             )
-            for k in range(1, node_count):
+
+            res = self.divergence_dtype(0.0)
+            for k in range(node_count):
                 res += self.space.space_divergence(
                     self._read_node_value(args, s.element_index, k),
                     self.space.element_outer_weight_gradient(
-                        args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k
+                        args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k, s.qp_index
                     ),
                     local_value_map,
                     grad_transform,

warp/fem/field/restriction.py

@@ -1,3 +1,18 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from warp.fem.space import SpaceRestriction
 
 from .field import DiscreteField

warp/fem/field/virtual.py

@@ -1,3 +1,18 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from typing import Any, Set
 
 import warp as wp
@@ -54,7 +69,12 @@ class AdjointField(SpaceField):
         def eval_test_inner(args: self.ElementEvalArg, s: Sample):
             dof = self._get_dof(s)
             node_weight = self.space.element_inner_weight(
-                args.elt_arg, args.eval_arg, s.element_index, s.element_coords, get_node_index_in_element(dof)
+                args.elt_arg,
+                args.eval_arg,
+                s.element_index,
+                s.element_coords,
+                get_node_index_in_element(dof),
+                s.qp_index,
             )
             local_value_map = self.space.local_value_map_inner(args.elt_arg, s.element_index, s.element_coords)
             dof_value = self.space.node_basis_element(get_node_coord(dof))
@@ -75,6 +95,7 @@ class AdjointField(SpaceField):
                 s.element_index,
                 s.element_coords,
                 get_node_index_in_element(dof),
+                s.qp_index,
             )
             grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
             local_value_map = self.space.local_value_map_inner(args.elt_arg, s.element_index, s.element_coords)
@@ -96,6 +117,7 @@ class AdjointField(SpaceField):
                 s.element_index,
                 s.element_coords,
                 get_node_index_in_element(dof),
+                s.qp_index,
             )
             grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
             local_value_map = self.space.local_value_map_inner(args.elt_arg, s.element_index, s.element_coords)
@@ -109,7 +131,12 @@ class AdjointField(SpaceField):
         def eval_test_outer(args: self.ElementEvalArg, s: Sample):
             dof = self._get_dof(s)
             node_weight = self.space.element_outer_weight(
-                args.elt_arg, args.eval_arg, s.element_index, s.element_coords, get_node_index_in_element(dof)
+                args.elt_arg,
+                args.eval_arg,
+                s.element_index,
+                s.element_coords,
+                get_node_index_in_element(dof),
+                s.qp_index,
             )
             local_value_map = self.space.local_value_map_outer(args.elt_arg, s.element_index, s.element_coords)
             dof_value = self.space.node_basis_element(get_node_coord(dof))
@@ -130,6 +157,7 @@ class AdjointField(SpaceField):
                 s.element_index,
                 s.element_coords,
                 get_node_index_in_element(dof),
+                s.qp_index,
             )
             grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
             local_value_map = self.space.local_value_map_outer(args.elt_arg, s.element_index, s.element_coords)
@@ -151,6 +179,7 @@ class AdjointField(SpaceField):
                 s.element_index,
                 s.element_coords,
                 get_node_index_in_element(dof),
+                s.qp_index,
             )
             grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
             local_value_map = self.space.local_value_map_outer(args.elt_arg, s.element_index, s.element_coords)
@@ -355,9 +384,8 @@ class LocalAdjointField(SpaceField):
 
         @cache.dynamic_func(suffix=str(TAYLOR_DOF_COUNT))
         def split_dof(dof_index: DofIndex, dof_begin: int):
-            dof = get_node_coord(dof_index)
-            value_dof = dof // TAYLOR_DOF_COUNT
-            taylor_dof = dof - value_dof * TAYLOR_DOF_COUNT - dof_begin
+            taylor_dof = get_node_index_in_element(dof_index) - dof_begin
+            value_dof = get_node_coord(dof_index)
             return value_dof, taylor_dof
 
         return split_dof
@@ -371,7 +399,7 @@ class LocalAdjointField(SpaceField):
 
             local_value_map = self.space.local_value_map_inner(args.elt_arg, s.element_index, s.element_coords)
             dof_value = self.space.value_basis_element(value_dof, local_value_map)
-            return wp.select(taylor_dof == 0, self.dtype(0.0), dof_value)
+            return wp.where(taylor_dof == 0, dof_value, self.dtype(0.0))
 
         return eval_test_inner
 
@@ -426,7 +454,7 @@ class LocalAdjointField(SpaceField):
 
             local_value_map = self.space.local_value_map_outer(args.elt_arg, s.element_index, s.element_coords)
             dof_value = self.space.value_basis_element(value_dof, local_value_map)
-            return wp.select(taylor_dof == 0, self.dtype(0.0), dof_value)
+            return wp.where(taylor_dof == 0, dof_value, self.dtype(0.0))
 
         return eval_test_outer
 
@@ -536,17 +564,25 @@ def make_linear_dispatch_kernel(test: LocalTestField, quadrature: Quadrature, ac
                 qp_index = quadrature.point_index(
                     domain_arg, qp_arg, test_element_index.domain_element_index, element_index, k
                 )
+                qp_eval_index = quadrature.point_evaluation_index(
+                    domain_arg, qp_arg, test_element_index.domain_element_index, element_index, k
+                )
                 coords = quadrature.point_coords(
                     domain_arg, qp_arg, test_element_index.domain_element_index, element_index, k
                 )
 
-                qp_result = local_result[qp_index]
+                qp_result = local_result[qp_eval_index]
 
                 qp_sum = float(0.0)
 
                 if wp.static(0 != TEST_INNER_COUNT):
                     w = test.space.element_inner_weight(
-                        domain_arg, test_space_arg, element_index, coords, test_element_index.node_index_in_element
+                        domain_arg,
+                        test_space_arg,
+                        element_index,
+                        coords,
+                        test_element_index.node_index_in_element,
+                        qp_index,
                     )
                     for val_dof in range(TEST_NODE_DOF_DIM):
                         test_dof = test_node_dof * TEST_NODE_DOF_DIM + val_dof
@@ -554,7 +590,12 @@ def make_linear_dispatch_kernel(test: LocalTestField, quadrature: Quadrature, ac
 
                 if wp.static(0 != TEST_OUTER_COUNT):
                     w = test.space.element_outer_weight(
-                        domain_arg, test_space_arg, element_index, coords, test_element_index.node_index_in_element
+                        domain_arg,
+                        test_space_arg,
+                        element_index,
+                        coords,
+                        test_element_index.node_index_in_element,
+                        qp_index,
                     )
                     for val_dof in range(TEST_NODE_DOF_DIM):
                         test_dof = test_node_dof * TEST_NODE_DOF_DIM + val_dof
@@ -562,7 +603,12 @@ def make_linear_dispatch_kernel(test: LocalTestField, quadrature: Quadrature, ac
 
                 if wp.static(0 != TEST_INNER_GRAD_COUNT):
                     w_grad = test.space.element_inner_weight_gradient(
-                        domain_arg, test_space_arg, element_index, coords, test_element_index.node_index_in_element
+                        domain_arg,
+                        test_space_arg,
+                        element_index,
+                        coords,
+                        test_element_index.node_index_in_element,
+                        qp_index,
                     )
                     for val_dof in range(TEST_NODE_DOF_DIM):
                         test_dof = test_node_dof * TEST_NODE_DOF_DIM + val_dof
@@ -574,7 +620,12 @@ def make_linear_dispatch_kernel(test: LocalTestField, quadrature: Quadrature, ac
 
                 if wp.static(0 != TEST_OUTER_GRAD_COUNT):
                     w_grad = test.space.element_outer_weight_gradient(
-                        domain_arg, test_space_arg, element_index, coords, test_element_index.node_index_in_element
+                        domain_arg,
+                        test_space_arg,
+                        element_index,
+                        coords,
+                        test_element_index.node_index_in_element,
+                        qp_index,
                     )
                     for val_dof in range(TEST_NODE_DOF_DIM):
                         test_dof = test_node_dof * TEST_NODE_DOF_DIM + val_dof
@@ -654,10 +705,10 @@ def make_bilinear_dispatch_kernel(
                 domain_arg, trial_topology_arg, element_index
             )
 
-            qp_point_count = wp.select(
+            qp_point_count = wp.where(
                 trial_node < element_trial_node_count,
-                0,
                 quadrature.point_count(domain_arg, qp_arg, test_element_index.domain_element_index, element_index),
+                0,
             )
 
             val_sum = accumulate_dtype(0.0)
@@ -666,51 +717,54 @@ def make_bilinear_dispatch_kernel(
                 qp_index = quadrature.point_index(
                     domain_arg, qp_arg, test_element_index.domain_element_index, element_index, k
                 )
+                qp_eval_index = quadrature.point_evaluation_index(
+                    domain_arg, qp_arg, test_element_index.domain_element_index, element_index, k
+                )
                 coords = quadrature.point_coords(
                     domain_arg, qp_arg, test_element_index.domain_element_index, element_index, k
                 )
 
-                qp_result = local_result[qp_index]
+                qp_result = local_result[qp_eval_index]
                 trial_result = float(0.0)
 
                 if wp.static(0 != TEST_INNER_COUNT):
                     w_test_inner = test.space.element_inner_weight(
-                        domain_arg, test_space_arg, element_index, coords, test_node
+                        domain_arg, test_space_arg, element_index, coords, test_node, qp_index
                     )
 
                 if wp.static(0 != TEST_OUTER_COUNT):
                     w_test_outer = test.space.element_outer_weight(
-                        domain_arg, test_space_arg, element_index, coords, test_node
+                        domain_arg, test_space_arg, element_index, coords, test_node, qp_index
                     )
 
                 if wp.static(0 != TEST_INNER_GRAD_COUNT):
                     w_test_grad_inner = test.space.element_inner_weight_gradient(
-                        domain_arg, test_space_arg, element_index, coords, test_node
+                        domain_arg, test_space_arg, element_index, coords, test_node, qp_index
                     )
 
                 if wp.static(0 != TEST_OUTER_GRAD_COUNT):
                     w_test_grad_outer = test.space.element_outer_weight_gradient(
-                        domain_arg, test_space_arg, element_index, coords, test_node
+                        domain_arg, test_space_arg, element_index, coords, test_node, qp_index
                     )
 
                 if wp.static(0 != TRIAL_INNER_COUNT):
                     w_trial_inner = trial.space.element_inner_weight(
-                        domain_arg, trial_space_arg, element_index, coords, trial_node
+                        domain_arg, trial_space_arg, element_index, coords, trial_node, qp_index
                     )
 
                 if wp.static(0 != TRIAL_OUTER_COUNT):
                     w_trial_outer = trial.space.element_outer_weight(
-                        domain_arg, trial_space_arg, element_index, coords, trial_node
+                        domain_arg, trial_space_arg, element_index, coords, trial_node, qp_index
                     )
 
                 if wp.static(0 != TRIAL_INNER_GRAD_COUNT):
                     w_trial_grad_inner = trial.space.element_inner_weight_gradient(
-                        domain_arg, trial_space_arg, element_index, coords, trial_node
+                        domain_arg, trial_space_arg, element_index, coords, trial_node, qp_index
                     )
 
                 if wp.static(0 != TRIAL_OUTER_GRAD_COUNT):
                     w_trial_grad_outer = trial.space.element_outer_weight_gradient(
-                        domain_arg, trial_space_arg, element_index, coords, trial_node
+                        domain_arg, trial_space_arg, element_index, coords, trial_node, qp_index
                     )
 
                 for trial_val_dof in range(TRIAL_NODE_DOF_DIM):

warp/fem/geometry/__init__.py

@@ -1,3 +1,18 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from .adaptive_nanogrid import AdaptiveNanogrid
 from .deformed_geometry import DeformedGeometry
 from .element import Element