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

warp/fem/space/basis_function_space.py

@@ -1,10 +1,25 @@
+# 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, Optional
 
 import warp as wp
 from warp.fem import cache
 from warp.fem.geometry import Geometry
 from warp.fem.linalg import basis_element, generalized_inner, generalized_outer
-from warp.fem.types import Coords, ElementIndex, make_free_sample
+from warp.fem.types import NULL_QP_INDEX, Coords, ElementIndex, make_free_sample
 
 from .basis_space import BasisSpace
 from .dof_mapper import DofMapper, IdentityMapper
@@ -290,7 +305,9 @@ class VectorValuedFunctionSpace(FunctionSpace):
             space_value: self.dtype,
         ):
             coords = self.node_coords_in_element(elt_arg, space_arg, element_index, node_index_in_elt)
-            weight = self.element_inner_weight(elt_arg, space_arg, element_index, coords, node_index_in_elt)
+            weight = self.element_inner_weight(
+                elt_arg, space_arg, element_index, coords, node_index_in_elt, NULL_QP_INDEX
+            )
             local_value_map = self.local_value_map_inner(elt_arg, element_index, coords)
 
             unit_value = local_value_map * weight

warp/fem/space/basis_space.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
 
 import numpy as np
@@ -6,7 +21,14 @@ import warp as wp
 from warp.fem import cache
 from warp.fem.geometry import Geometry
 from warp.fem.quadrature import Quadrature
-from warp.fem.types import NULL_ELEMENT_INDEX, Coords, ElementIndex, make_free_sample
+from warp.fem.types import (
+    NULL_ELEMENT_INDEX,
+    NULL_QP_INDEX,
+    Coords,
+    ElementIndex,
+    QuadraturePointIndex,
+    make_free_sample,
+)
 
 from .shape import ShapeFunction
 from .topology import RegularDiscontinuousSpaceTopology, SpaceTopology
@@ -220,6 +242,7 @@ class ShapeBasisSpace(BasisSpace):
             element_index: ElementIndex,
             coords: Coords,
             node_index_in_elt: int,
+            qp_index: QuadraturePointIndex,
         ):
             if wp.static(self.value == ShapeFunction.Value.Scalar):
                 return shape_element_inner_weight(coords, node_index_in_elt)
@@ -239,6 +262,7 @@ class ShapeBasisSpace(BasisSpace):
             element_index: ElementIndex,
             coords: Coords,
             node_index_in_elt: int,
+            qp_index: QuadraturePointIndex,
         ):
             if wp.static(self.value == ShapeFunction.Value.Scalar):
                 return shape_element_inner_weight_gradient(coords, node_index_in_elt)
@@ -358,6 +382,7 @@ class TraceBasisSpace(BasisSpace):
             element_index: ElementIndex,
             coords: Coords,
             node_index_in_elt: int,
+            qp_index: QuadraturePointIndex,
         ):
             cell_index, index_in_cell = self.topology.inner_cell_index(geo_side_arg, element_index, node_index_in_elt)
             if cell_index == NULL_ELEMENT_INDEX:
@@ -366,13 +391,7 @@ class TraceBasisSpace(BasisSpace):
             cell_coords = self.geometry.side_inner_cell_coords(geo_side_arg, element_index, coords)
 
             geo_cell_arg = self.geometry.side_to_cell_arg(geo_side_arg)
-            return cell_inner_weight(
-                geo_cell_arg,
-                basis_arg,
-                cell_index,
-                cell_coords,
-                index_in_cell,
-            )
+            return cell_inner_weight(geo_cell_arg, basis_arg, cell_index, cell_coords, index_in_cell, NULL_QP_INDEX)
 
         return trace_element_inner_weight
 
@@ -386,6 +405,7 @@ class TraceBasisSpace(BasisSpace):
             element_index: ElementIndex,
             coords: Coords,
             node_index_in_elt: int,
+            qp_index: QuadraturePointIndex,
         ):
             cell_index, index_in_cell = self.topology.outer_cell_index(geo_side_arg, element_index, node_index_in_elt)
             if cell_index == NULL_ELEMENT_INDEX:
@@ -394,13 +414,7 @@ class TraceBasisSpace(BasisSpace):
             cell_coords = self.geometry.side_outer_cell_coords(geo_side_arg, element_index, coords)
 
             geo_cell_arg = self.geometry.side_to_cell_arg(geo_side_arg)
-            return cell_outer_weight(
-                geo_cell_arg,
-                basis_arg,
-                cell_index,
-                cell_coords,
-                index_in_cell,
-            )
+            return cell_outer_weight(geo_cell_arg, basis_arg, cell_index, cell_coords, index_in_cell, NULL_QP_INDEX)
 
         return trace_element_outer_weight
 
@@ -414,6 +428,7 @@ class TraceBasisSpace(BasisSpace):
             element_index: ElementIndex,
             coords: Coords,
             node_index_in_elt: int,
+            qp_index: QuadraturePointIndex,
         ):
             cell_index, index_in_cell = self.topology.inner_cell_index(geo_side_arg, element_index, node_index_in_elt)
             if cell_index == NULL_ELEMENT_INDEX:
@@ -421,7 +436,9 @@ class TraceBasisSpace(BasisSpace):
 
             cell_coords = self.geometry.side_inner_cell_coords(geo_side_arg, element_index, coords)
             geo_cell_arg = self.geometry.side_to_cell_arg(geo_side_arg)
-            return cell_inner_weight_gradient(geo_cell_arg, basis_arg, cell_index, cell_coords, index_in_cell)
+            return cell_inner_weight_gradient(
+                geo_cell_arg, basis_arg, cell_index, cell_coords, index_in_cell, NULL_QP_INDEX
+            )
 
         return trace_element_inner_weight_gradient
 
@@ -435,6 +452,7 @@ class TraceBasisSpace(BasisSpace):
             element_index: ElementIndex,
             coords: Coords,
             node_index_in_elt: int,
+            qp_index: QuadraturePointIndex,
         ):
             cell_index, index_in_cell = self.topology.outer_cell_index(geo_side_arg, element_index, node_index_in_elt)
             if cell_index == NULL_ELEMENT_INDEX:
@@ -442,7 +460,9 @@ class TraceBasisSpace(BasisSpace):
 
             cell_coords = self.geometry.side_outer_cell_coords(geo_side_arg, element_index, coords)
             geo_cell_arg = self.geometry.side_to_cell_arg(geo_side_arg)
-            return cell_outer_weight_gradient(geo_cell_arg, basis_arg, cell_index, cell_coords, index_in_cell)
+            return cell_outer_weight_gradient(
+                geo_cell_arg, basis_arg, cell_index, cell_coords, index_in_cell, NULL_QP_INDEX
+            )
 
         return trace_element_outer_weight_gradient
 
@@ -609,11 +629,12 @@ class PointBasisSpace(BasisSpace):
             element_index: ElementIndex,
             coords: Coords,
             node_index_in_elt: int,
+            qp_index: QuadraturePointIndex,
         ):
             qp_coord = self._quadrature.point_coords(
                 elt_arg, basis_arg, element_index, element_index, node_index_in_elt
             )
-            return wp.select(wp.length_sq(coords - qp_coord) < _DIRAC_INTEGRATION_RADIUS, 0.0, 1.0)
+            return wp.where(wp.length_sq(coords - qp_coord) < _DIRAC_INTEGRATION_RADIUS, 1.0, 0.0)
 
         return element_inner_weight
 
@@ -627,6 +648,7 @@ class PointBasisSpace(BasisSpace):
             element_index: ElementIndex,
             coords: Coords,
             node_index_in_elt: int,
+            qp_index: QuadraturePointIndex,
         ):
             return gradient_vec(0.0)
 

warp/fem/space/dof_mapper.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 math
 from enum import Enum
 from typing import Any

warp/fem/space/function_space.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

warp/fem/space/grid_2d_function_space.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 numpy as np
 
 import warp as wp

warp/fem/space/grid_3d_function_space.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 numpy as np
 
 import warp as wp

warp/fem/space/hexmesh_function_space.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 warp as wp
 from warp.fem import cache
 from warp.fem.geometry import Hexmesh
@@ -222,13 +237,13 @@ class HexmeshSpaceTopology(SpaceTopology):
                     hex_edge = _CUBE_TO_HEX_EDGE[type_instance]
                     v0 = geo_arg.hex_vertex_indices[element_index, EDGE_VERTEX_INDICES[hex_edge, 0]]
                     v1 = geo_arg.hex_vertex_indices[element_index, EDGE_VERTEX_INDICES[hex_edge, 1]]
-                    return wp.select(v0 > v1, 1.0, -1.0)
+                    return wp.where(v0 > v1, -1.0, 1.0)
 
             if wp.static(FACE_NODE_COUNT > 0):
                 if node_type == CubeShapeFunction.FACE:
                     face_index_and_ori = topo_arg.hex_face_indices[element_index, type_instance]
                     flip = face_index_and_ori[1] & 1
-                    return wp.select(flip == 0, -1.0, 1.0)
+                    return wp.where(flip == 0, 1.0, -1.0)
 
             return 1.0
 

warp/fem/space/nanogrid_function_space.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 Union
 
 import warp as wp

warp/fem/space/partition.py

@@ -1,10 +1,25 @@
+# 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
 import warp.fem.cache as cache
 from warp.fem.geometry import GeometryPartition, WholeGeometryPartition
 from warp.fem.types import NULL_NODE_INDEX
-from warp.fem.utils import _iota_kernel, compress_node_indices
+from warp.fem.utils import compress_node_indices
 
 from .function_space import FunctionSpace
 from .topology import SpaceTopology
@@ -72,7 +87,7 @@ class WholeSpacePartition(SpacePartition):
         """Return the global function space indices for nodes in this partition"""
         if self._node_indices is None:
             self._node_indices = cache.borrow_temporary(temporary_store=None, shape=(self.node_count(),), dtype=int)
-            wp.launch(kernel=_iota_kernel, dim=self.node_count(), inputs=[self._node_indices.array, 1])
+            wp.launch(kernel=self._iota_kernel, dim=self.node_count(), inputs=[self._node_indices.array])
         return self._node_indices.array
 
     def partition_arg_value(self, device):
@@ -89,6 +104,10 @@ class WholeSpacePartition(SpacePartition):
     def name(self) -> str:
         return "Whole"
 
+    @wp.kernel
+    def _iota_kernel(indices: wp.array(dtype=int)):
+        indices[wp.tid()] = wp.tid()
+
 
 class NodeCategory:
     OWNED_INTERIOR = wp.constant(0)

warp/fem/space/quadmesh_function_space.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.
+
 import warp as wp
 from warp.fem import cache
 from warp.fem.geometry import Quadmesh2D
@@ -151,10 +166,10 @@ class QuadmeshSpaceTopology(SpaceTopology):
 
             if wp.static(EDGE_NODE_COUNT > 0):
                 # EDGE_X, EDGE_Y
-                side_start = wp.select(
+                side_start = wp.where(
                     node_type == SquareShapeFunction.EDGE_X,
-                    wp.select(type_instance == 0, 1, 3),
-                    wp.select(type_instance == 0, 2, 0),
+                    wp.where(type_instance == 0, 0, 2),
+                    wp.where(type_instance == 0, 3, 1),
                 )
 
                 side_index = topo_arg.quad_edge_indices[element_index, side_start]
@@ -163,7 +178,7 @@ class QuadmeshSpaceTopology(SpaceTopology):
 
                 # Flip indexing direction
                 flipped = int(side_start >= 2) ^ int(local_vs != global_vs)
-                index_in_side = wp.select(flipped, type_index, EDGE_NODE_COUNT - 1 - type_index)
+                index_in_side = wp.where(flipped, EDGE_NODE_COUNT - 1 - type_index, type_index)
 
                 return global_offset + EDGE_NODE_COUNT * side_index + index_in_side
 
@@ -182,10 +197,10 @@ class QuadmeshSpaceTopology(SpaceTopology):
             node_type, type_instance, type_index = self._shape.node_type_and_type_index(node_index_in_elt)
 
             if node_type == SquareShapeFunction.EDGE_X or node_type == SquareShapeFunction.EDGE_Y:
-                side_start = wp.select(
+                side_start = wp.where(
                     node_type == SquareShapeFunction.EDGE_X,
-                    wp.select(type_instance == 0, 1, 3),
-                    wp.select(type_instance == 0, 2, 0),
+                    wp.where(type_instance == 0, 0, 2),
+                    wp.where(type_instance == 0, 3, 1),
                 )
 
                 side_index = topo_arg.quad_edge_indices[element_index, side_start]
@@ -194,7 +209,7 @@ class QuadmeshSpaceTopology(SpaceTopology):
 
                 # Flip indexing direction
                 flipped = int(side_start >= 2) ^ int(local_vs != global_vs)
-                return wp.select(flipped, 1.0, -1.0)
+                return wp.where(flipped, -1.0, 1.0)
 
             return 1.0
 

warp/fem/space/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.
+
 import warp as wp
 from warp.fem.cache import TemporaryStore, borrow_temporary, borrow_temporary_like, cached_arg_value
 from warp.fem.domain import GeometryDomain

warp/fem/space/shape/__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 enum import Enum
 from typing import Optional
 

warp/fem/space/shape/cube_shape_function.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 math
 
 import numpy as np
@@ -48,7 +63,7 @@ class CubeTripolynomialShapeFunctions(CubeShapeFunction):
 
         self.ORDER = wp.constant(degree)
         self.NODES_PER_ELEMENT = wp.constant((degree + 1) ** 3)
-        self.NODES_PER_EDGE = wp.constant(degree + 1)
+        self.NODES_PER_SIDE = wp.constant((degree + 1) ** 2)
 
         if is_closed(self.family):
             self.VERTEX_NODE_COUNT = wp.constant(1)
@@ -137,13 +152,13 @@ class CubeTripolynomialShapeFunctions(CubeShapeFunction):
         ):
             i, j, k = self._node_ijk(node_index_in_elt)
 
-            zi = wp.select(i == 0, 0, 1)
-            zj = wp.select(j == 0, 0, 1)
-            zk = wp.select(k == 0, 0, 1)
+            zi = wp.where(i == 0, 1, 0)
+            zj = wp.where(j == 0, 1, 0)
+            zk = wp.where(k == 0, 1, 0)
 
-            mi = wp.select(i == ORDER, 0, 1)
-            mj = wp.select(j == ORDER, 0, 1)
-            mk = wp.select(k == ORDER, 0, 1)
+            mi = wp.where(i == ORDER, 1, 0)
+            mj = wp.where(j == ORDER, 1, 0)
+            mk = wp.where(k == ORDER, 1, 0)
 
             if zi + mi == 1:
                 if zj + mj == 1:
@@ -489,7 +504,7 @@ class CubeSerendipityShapeFunctions(CubeShapeFunction):
 
         self.ORDER = wp.constant(degree)
         self.NODES_PER_ELEMENT = wp.constant(8 + 12 * (degree - 1))
-        self.NODES_PER_EDGE = wp.constant(degree + 1)
+        self.NODES_PER_SIDE = wp.constant(4 * degree)
 
         self.VERTEX_NODE_COUNT = wp.constant(1)
         self.EDGE_NODE_COUNT = wp.constant(degree - 1)
@@ -619,9 +634,9 @@ class CubeSerendipityShapeFunctions(CubeShapeFunction):
             if node_type == CubeSerendipityShapeFunctions.VERTEX:
                 node_ijk = CubeSerendipityShapeFunctions._vertex_coords(type_instance)
 
-                cx = wp.select(node_ijk[0] == 0, coords[0], 1.0 - coords[0])
-                cy = wp.select(node_ijk[1] == 0, coords[1], 1.0 - coords[1])
-                cz = wp.select(node_ijk[2] == 0, coords[2], 1.0 - coords[2])
+                cx = wp.where(node_ijk[0] == 0, 1.0 - coords[0], coords[0])
+                cy = wp.where(node_ijk[1] == 0, 1.0 - coords[1], coords[1])
+                cz = wp.where(node_ijk[2] == 0, 1.0 - coords[2], coords[2])
 
                 w = cx * cy * cz
 
@@ -644,8 +659,8 @@ class CubeSerendipityShapeFunctions(CubeShapeFunction):
             local_coords = Grid3D._world_to_local(axis, coords)
 
             w = float(1.0)
-            w *= wp.select(node_all[1] == 0, local_coords[1], 1.0 - local_coords[1])
-            w *= wp.select(node_all[2] == 0, local_coords[2], 1.0 - local_coords[2])
+            w *= wp.where(node_all[1] == 0, 1.0 - local_coords[1], local_coords[1])
+            w *= wp.where(node_all[2] == 0, 1.0 - local_coords[2], local_coords[2])
 
             for k in range(ORDER_PLUS_ONE):
                 if k != node_all[0]:
@@ -675,13 +690,13 @@ class CubeSerendipityShapeFunctions(CubeShapeFunction):
             if node_type == CubeSerendipityShapeFunctions.VERTEX:
                 node_ijk = CubeSerendipityShapeFunctions._vertex_coords(type_instance)
 
-                cx = wp.select(node_ijk[0] == 0, coords[0], 1.0 - coords[0])
-                cy = wp.select(node_ijk[1] == 0, coords[1], 1.0 - coords[1])
-                cz = wp.select(node_ijk[2] == 0, coords[2], 1.0 - coords[2])
+                cx = wp.where(node_ijk[0] == 0, 1.0 - coords[0], coords[0])
+                cy = wp.where(node_ijk[1] == 0, 1.0 - coords[1], coords[1])
+                cz = wp.where(node_ijk[2] == 0, 1.0 - coords[2], coords[2])
 
-                gx = wp.select(node_ijk[0] == 0, 1.0, -1.0)
-                gy = wp.select(node_ijk[1] == 0, 1.0, -1.0)
-                gz = wp.select(node_ijk[2] == 0, 1.0, -1.0)
+                gx = wp.where(node_ijk[0] == 0, -1.0, 1.0)
+                gy = wp.where(node_ijk[1] == 0, -1.0, 1.0)
+                gz = wp.where(node_ijk[2] == 0, -1.0, 1.0)
 
                 if wp.static(ORDER == 2):
                     w = cx + cy + cz - 3.0 + LOBATTO_COORDS[1]
@@ -713,11 +728,11 @@ class CubeSerendipityShapeFunctions(CubeShapeFunction):
 
             local_coords = Grid3D._world_to_local(axis, coords)
 
-            w_long = wp.select(node_all[1] == 0, local_coords[1], 1.0 - local_coords[1])
-            w_lat = wp.select(node_all[2] == 0, local_coords[2], 1.0 - local_coords[2])
+            w_long = wp.where(node_all[1] == 0, 1.0 - local_coords[1], local_coords[1])
+            w_lat = wp.where(node_all[2] == 0, 1.0 - local_coords[2], local_coords[2])
 
-            g_long = wp.select(node_all[1] == 0, 1.0, -1.0)
-            g_lat = wp.select(node_all[2] == 0, 1.0, -1.0)
+            g_long = wp.where(node_all[1] == 0, -1.0, 1.0)
+            g_lat = wp.where(node_all[2] == 0, -1.0, 1.0)
 
             w_alt = LAGRANGE_SCALE[node_all[0]]
             g_alt = float(0.0)

warp/fem/space/shape/shape_function.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 enum import Enum
 
 import numpy as np