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

warp/fem/geometry/adaptive_nanogrid.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 Optional
 
 import numpy as np
@@ -15,7 +30,6 @@ from .nanogrid import (
     _extract_axis_flag,
     _get_boundary_mask,
     _make_face_flags,
-    _mat32,
 )
 
 _FACE_LEVEL_BIT = wp.constant(wp.uint8(4))  # follows nanogrid.FACE_OUTER_OFFSET_BIT
@@ -199,9 +213,9 @@ class AdaptiveNanogrid(Geometry):
             coords = uvw - wp.vec3(ijk)
 
             if wp.min(coords) == 0.0 or wp.max(coords) == 1.0:
-                il = wp.select(coords[0] > 0.5, -1, 0)
-                jl = wp.select(coords[1] > 0.5, -1, 0)
-                kl = wp.select(coords[2] > 0.5, -1, 0)
+                il = wp.where(coords[0] > 0.5, 0, -1)
+                jl = wp.where(coords[1] > 0.5, 0, -1)
+                kl = wp.where(coords[2] > 0.5, 0, -1)
 
                 for n in range(8):
                     ni = n >> 2
@@ -316,7 +330,7 @@ class AdaptiveNanogrid(Geometry):
         flip = Nanogrid._get_face_inner_offset(flags)
         scale = AdaptiveNanogrid._get_face_scale(flags)
         v1, v2 = Nanogrid._face_tangent_vecs(args.cell_arg.cell_grid, axis, flip)
-        return _mat32(v1, v2) * scale
+        return wp.matrix_from_cols(v1, v2) * scale
 
     @wp.func
     def side_inner_inverse_deformation_gradient(args: SideArg, s: Sample):
@@ -432,10 +446,10 @@ class AdaptiveNanogrid(Geometry):
             and wp.max(same_level_cell_coords) <= 1.0
         )
 
-        return wp.select(
+        return wp.where(
             on_side,
-            Coords(OUTSIDE),
             Coords(same_level_cell_coords[(axis + 1) % 3], same_level_cell_coords[(axis + 2) % 3], 0.0),
+            Coords(OUTSIDE),
         )
 
     def _build_face_grid(self, temporary_store: Optional[cache.TemporaryStore] = None):
@@ -512,7 +526,7 @@ class AdaptiveNanogrid(Geometry):
         for ax in range(3):
             coord = ijk[ax]
             level_flag = ((level >> ax) & 1) << _GRID_LEVEL_BIT
-            ijk[ax] = wp.select(coord < 0, coord | level_flag, coord & ~level_flag)
+            ijk[ax] = wp.where(coord < 0, coord & ~level_flag, coord | level_flag)
 
         return _add_axis_flag(ijk, axis)
 
@@ -830,8 +844,8 @@ def _build_face_indices_and_flags(
 
     plus_cell_index, minus_cell_index = _find_face_neighbours(cell_grid, ijk, axis, level_count, cell_level)
 
-    inner_cell = wp.select(minus_cell_index == -1, minus_cell_index, plus_cell_index)
-    outer_cell = wp.select(plus_cell_index == -1, plus_cell_index, minus_cell_index)
+    inner_cell = wp.where(minus_cell_index == -1, plus_cell_index, minus_cell_index)
+    outer_cell = wp.where(plus_cell_index == -1, minus_cell_index, plus_cell_index)
 
     face_level = wp.min(cell_level[inner_cell], cell_level[outer_cell])
 

warp/fem/geometry/closest_point.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
@@ -44,7 +59,7 @@ def project_on_tri_at_origin(q: Any, e1: Any, e2: Any):
 
 @wp.func
 def project_on_tet_at_origin(q: wp.vec3, e1: wp.vec3, e2: wp.vec3, e3: wp.vec3):
-    mat = wp.inverse(wp.mat33(e1, e2, e3))
+    mat = wp.inverse(wp.matrix_from_cols(e1, e2, e3))
     coords = mat * q
 
     if wp.min(coords) >= 0.0 and coords[0] + coords[1] + coords[2] <= 1.0:

warp/fem/geometry/deformed_geometry.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.types import Coords, ElementIndex, Sample
@@ -31,7 +46,6 @@ class DeformedGeometry(Geometry):
         self._relative = relative
 
         self.field: GeometryField = field
-        self.base = self.field.geometry
         self.dimension = self.base.dimension
 
         self.CellArg = self.field.ElementEvalArg
@@ -66,9 +80,13 @@ class DeformedGeometry(Geometry):
         self._make_default_dependent_implementations()
 
     @property
-    def name(self):
+    def name(self) -> str:
         return f"DefGeo_{self.field.name}_{'rel' if self._relative else 'abs'}"
 
+    @property
+    def base(self) -> Geometry:
+        return self.field.geometry.base
+
     # Geometry device interface
 
     @cache.cached_arg_value

warp/fem/geometry/element.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 List, Tuple
 
 from warp.fem.polynomial import Polynomial, quadrature_1d

warp/fem/geometry/geometry.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
@@ -43,6 +58,11 @@ class Geometry:
         """Manifold dimension of the geometry cells"""
         return self.reference_cell().dimension
 
+    @property
+    def base(self) -> "Geometry":
+        """Returns the base geometry from which this geometry derives its topology. Usually `self`"""
+        return self
+
     @property
     def name(self) -> str:
         return self.__class__.__name__

warp/fem/geometry/grid_2d.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 warp as wp
@@ -158,7 +173,7 @@ class Grid2D(Geometry):
             return Grid2D.Side(axis, origin)
 
         axis_side_index = side_index - 2 * arg.cell_count
-        axis = wp.select(axis_side_index < arg.axis_offsets[1], 1, 0)
+        axis = wp.where(axis_side_index < arg.axis_offsets[1], 0, 1)
 
         altitude = arg.cell_arg.res[Grid2D.ROTATION[axis, 0]]
         longitude = axis_side_index - arg.axis_offsets[axis]
@@ -258,7 +273,7 @@ class Grid2D(Geometry):
     def side_position(args: SideArg, s: Sample):
         side = Grid2D.get_side(args, s.element_index)
 
-        coord = wp.select((side.origin[0] == 0) == (side.axis == 0), 1.0 - s.element_coords[0], s.element_coords[0])
+        coord = wp.where((side.origin[0] == 0) == (side.axis == 0), s.element_coords[0], 1.0 - s.element_coords[0])
 
         local_pos = wp.vec2(
             float(side.origin[0]),
@@ -273,7 +288,7 @@ class Grid2D(Geometry):
     def side_deformation_gradient(args: SideArg, s: Sample):
         side = Grid2D.get_side(args, s.element_index)
 
-        sign = wp.select((side.origin[0] == 0) == (side.axis == 0), -1.0, 1.0)
+        sign = wp.where((side.origin[0] == 0) == (side.axis == 0), 1.0, -1.0)
 
         return wp.cw_mul(Grid2D._rotate(side.axis, wp.vec2(0.0, sign)), args.cell_arg.cell_size)
 
@@ -301,7 +316,7 @@ class Grid2D(Geometry):
     def side_normal(args: SideArg, s: Sample):
         side = Grid2D.get_side(args, s.element_index)
 
-        sign = wp.select(side.origin[0] == 0, 1.0, -1.0)
+        sign = wp.where(side.origin[0] == 0, -1.0, 1.0)
 
         local_n = wp.vec2(sign, 0.0)
         return Grid2D._rotate(side.axis, local_n)
@@ -310,7 +325,7 @@ class Grid2D(Geometry):
     def side_inner_cell_index(arg: SideArg, side_index: ElementIndex):
         side = Grid2D.get_side(arg, side_index)
 
-        inner_alt = wp.select(side.origin[0] == 0, side.origin[0] - 1, 0)
+        inner_alt = wp.where(side.origin[0] == 0, 0, side.origin[0] - 1)
 
         inner_origin = wp.vec2i(inner_alt, side.origin[1])
 
@@ -322,8 +337,8 @@ class Grid2D(Geometry):
         side = Grid2D.get_side(arg, side_index)
 
         alt_axis = Grid2D.ROTATION[side.axis, 0]
-        outer_alt = wp.select(
-            side.origin[0] == arg.cell_arg.res[alt_axis], side.origin[0], arg.cell_arg.res[alt_axis] - 1
+        outer_alt = wp.where(
+            side.origin[0] == arg.cell_arg.res[alt_axis], arg.cell_arg.res[alt_axis] - 1, side.origin[0]
         )
 
         outer_origin = wp.vec2i(outer_alt, side.origin[1])
@@ -335,9 +350,9 @@ class Grid2D(Geometry):
     def side_inner_cell_coords(args: SideArg, side_index: ElementIndex, side_coords: Coords):
         side = Grid2D.get_side(args, side_index)
 
-        inner_alt = wp.select(side.origin[0] == 0, 1.0, 0.0)
+        inner_alt = wp.where(side.origin[0] == 0, 0.0, 1.0)
 
-        side_coord = wp.select((side.origin[0] == 0) == (side.axis == 0), 1.0 - side_coords[0], side_coords[0])
+        side_coord = wp.where((side.origin[0] == 0) == (side.axis == 0), side_coords[0], 1.0 - side_coords[0])
 
         coords = Grid2D._rotate(side.axis, wp.vec2(inner_alt, side_coord))
         return Coords(coords[0], coords[1], 0.0)
@@ -347,9 +362,9 @@ class Grid2D(Geometry):
         side = Grid2D.get_side(args, side_index)
 
         alt_axis = Grid2D.ROTATION[side.axis, 0]
-        outer_alt = wp.select(side.origin[0] == args.cell_arg.res[alt_axis], 0.0, 1.0)
+        outer_alt = wp.where(side.origin[0] == args.cell_arg.res[alt_axis], 1.0, 0.0)
 
-        side_coord = wp.select((side.origin[0] == 0) == (side.axis == 0), 1.0 - side_coords[0], side_coords[0])
+        side_coord = wp.where((side.origin[0] == 0) == (side.axis == 0), side_coords[0], 1.0 - side_coords[0])
 
         coords = Grid2D._rotate(side.axis, wp.vec2(outer_alt, side_coord))
         return Coords(coords[0], coords[1], 0.0)
@@ -367,7 +382,7 @@ class Grid2D(Geometry):
         if float(side.origin[0] - cell[side.axis]) == element_coords[side.axis]:
             long_axis = Grid2D.ROTATION[side.axis, 1]
             axis_coord = element_coords[long_axis]
-            side_coord = wp.select((side.origin[0] == 0) == (side.axis == 0), 1.0 - axis_coord, axis_coord)
+            side_coord = wp.where((side.origin[0] == 0) == (side.axis == 0), axis_coord, 1.0 - axis_coord)
             return Coords(side_coord, 0.0, 0.0)
 
         return Coords(OUTSIDE)

warp/fem/geometry/grid_3d.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
@@ -15,9 +30,6 @@ class Grid3DCellArg:
     origin: wp.vec3
 
 
-_mat32 = wp.mat(shape=(3, 2), dtype=float)
-
-
 class Grid3D(Geometry):
     """Three-dimensional regular grid geometry"""
 
@@ -316,7 +328,7 @@ class Grid3D(Geometry):
     def side_position(args: SideArg, s: Sample):
         side = Grid3D.get_side(args, s.element_index)
 
-        coord0 = wp.select(side.origin[0] == 0, s.element_coords[0], 1.0 - s.element_coords[0])
+        coord0 = wp.where(side.origin[0] == 0, 1.0 - s.element_coords[0], s.element_coords[0])
 
         local_pos = wp.vec3(
             float(side.origin[0]),
@@ -332,9 +344,9 @@ class Grid3D(Geometry):
     def side_deformation_gradient(args: SideArg, s: Sample):
         side = Grid3D.get_side(args, s.element_index)
 
-        sign = wp.select(side.origin[0] == 0, 1.0, -1.0)
+        sign = wp.where(side.origin[0] == 0, -1.0, 1.0)
 
-        return _mat32(
+        return wp.matrix_from_cols(
             wp.cw_mul(Grid3D._local_to_world(side.axis, wp.vec3(0.0, sign, 0.0)), args.cell_arg.cell_size),
             wp.cw_mul(Grid3D._local_to_world(side.axis, wp.vec3(0.0, 0.0, 1.0)), args.cell_arg.cell_size),
         )
@@ -364,7 +376,7 @@ class Grid3D(Geometry):
     def side_normal(args: SideArg, s: Sample):
         side = Grid3D.get_side(args, s.element_index)
 
-        sign = wp.select(side.origin[0] == 0, 1.0, -1.0)
+        sign = wp.where(side.origin[0] == 0, -1.0, 1.0)
 
         local_n = wp.vec3(sign, 0.0, 0.0)
         return Grid3D._local_to_world(side.axis, local_n)
@@ -373,7 +385,7 @@ class Grid3D(Geometry):
     def side_inner_cell_index(arg: SideArg, side_index: ElementIndex):
         side = Grid3D.get_side(arg, side_index)
 
-        inner_alt = wp.select(side.origin[0] == 0, side.origin[0] - 1, 0)
+        inner_alt = wp.where(side.origin[0] == 0, 0, side.origin[0] - 1)
 
         inner_origin = wp.vec3i(inner_alt, side.origin[1], side.origin[2])
 
@@ -386,8 +398,8 @@ class Grid3D(Geometry):
 
         alt_axis = Grid3D._local_to_world_axis(side.axis, 0)
 
-        outer_alt = wp.select(
-            side.origin[0] == arg.cell_arg.res[alt_axis], side.origin[0], arg.cell_arg.res[alt_axis] - 1
+        outer_alt = wp.where(
+            side.origin[0] == arg.cell_arg.res[alt_axis], arg.cell_arg.res[alt_axis] - 1, side.origin[0]
         )
 
         outer_origin = wp.vec3i(outer_alt, side.origin[1], side.origin[2])
@@ -399,9 +411,9 @@ class Grid3D(Geometry):
     def side_inner_cell_coords(args: SideArg, side_index: ElementIndex, side_coords: Coords):
         side = Grid3D.get_side(args, side_index)
 
-        inner_alt = wp.select(side.origin[0] == 0, 1.0, 0.0)
+        inner_alt = wp.where(side.origin[0] == 0, 0.0, 1.0)
 
-        side_coord0 = wp.select(side.origin[0] == 0, side_coords[0], 1.0 - side_coords[0])
+        side_coord0 = wp.where(side.origin[0] == 0, 1.0 - side_coords[0], side_coords[0])
 
         return Grid3D._local_to_world(side.axis, wp.vec3(inner_alt, side_coord0, side_coords[1]))
 
@@ -410,9 +422,9 @@ class Grid3D(Geometry):
         side = Grid3D.get_side(args, side_index)
 
         alt_axis = Grid3D._local_to_world_axis(side.axis, 0)
-        outer_alt = wp.select(side.origin[0] == args.cell_arg.res[alt_axis], 0.0, 1.0)
+        outer_alt = wp.where(side.origin[0] == args.cell_arg.res[alt_axis], 1.0, 0.0)
 
-        side_coord0 = wp.select(side.origin[0] == 0, side_coords[0], 1.0 - side_coords[0])
+        side_coord0 = wp.where(side.origin[0] == 0, 1.0 - side_coords[0], side_coords[0])
 
         return Grid3D._local_to_world(side.axis, wp.vec3(outer_alt, side_coord0, side_coords[1]))
 
@@ -430,7 +442,7 @@ class Grid3D(Geometry):
             long_axis = Grid3D._local_to_world_axis(side.axis, 1)
             lat_axis = Grid3D._local_to_world_axis(side.axis, 2)
             long_coord = element_coords[long_axis]
-            long_coord = wp.select(side.origin[0] == 0, long_coord, 1.0 - long_coord)
+            long_coord = wp.where(side.origin[0] == 0, 1.0 - long_coord, long_coord)
             return Coords(long_coord, element_coords[lat_axis], 0.0)
 
         return Coords(OUTSIDE)

warp/fem/geometry/hexmesh.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 warp as wp
@@ -31,8 +46,6 @@ class HexmeshSideArg:
     face_hex_face_orientation: wp.array(dtype=wp.vec4i)
 
 
-_mat32 = wp.mat(shape=(3, 2), dtype=float)
-
 FACE_VERTEX_INDICES = wp.constant(
     wp.mat(shape=(6, 4), dtype=int)(
         [
@@ -307,7 +320,7 @@ class Hexmesh(Geometry):
     def side_deformation_gradient(args: SideArg, s: Sample):
         """Transposed side deformation gradient at `coords`"""
         v1, v2 = Hexmesh._side_deformation_vecs(args, s.element_index, s.element_coords)
-        return _mat32(v1, v2)
+        return wp.matrix_from_cols(v1, v2)
 
     @wp.func
     def side_inner_cell_index(arg: SideArg, side_index: ElementIndex):
@@ -327,7 +340,7 @@ class Hexmesh(Geometry):
         )
 
         normal_coord = hex_coords[_FACE_COORD_INDICES[face_index, 2]]
-        normal_coord = wp.select(_FACE_COORD_INDICES[face_index, 3] == 0, normal_coord - 1.0, -normal_coord)
+        normal_coord = wp.where(_FACE_COORD_INDICES[face_index, 3] == 0, -normal_coord, normal_coord - 1.0)
 
         return face_coords, normal_coord
 
@@ -338,7 +351,7 @@ class Hexmesh(Geometry):
         hex_coords = Coords()
         hex_coords[_FACE_COORD_INDICES[face_index, 0]] = face_coords[0]
         hex_coords[_FACE_COORD_INDICES[face_index, 1]] = face_coords[1]
-        hex_coords[_FACE_COORD_INDICES[face_index, 2]] = wp.select(_FACE_COORD_INDICES[face_index, 3] == 0, 1.0, 0.0)
+        hex_coords[_FACE_COORD_INDICES[face_index, 2]] = wp.where(_FACE_COORD_INDICES[face_index, 3] == 0, 0.0, 1.0)
 
         return hex_coords
 
@@ -387,8 +400,8 @@ class Hexmesh(Geometry):
             face_orientation = args.face_hex_face_orientation[side_index][3]
 
         face_coords, normal_coord = Hexmesh._hex_local_face_coords(hex_coords, local_face_index)
-        return wp.select(
-            normal_coord == 0.0, Coords(OUTSIDE), Hexmesh._local_to_oriented_face_coords(face_orientation, face_coords)
+        return wp.where(
+            normal_coord == 0.0, Hexmesh._local_to_oriented_face_coords(face_orientation, face_coords), Coords(OUTSIDE)
         )
 
     @wp.func

warp/fem/geometry/nanogrid.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 Optional
 
 import numpy as np
@@ -18,13 +33,11 @@ FACE_AXIS_MASK = wp.constant(wp.uint8((1 << 2) - 1))
 FACE_INNER_OFFSET_BIT = wp.constant(wp.uint8(2))
 FACE_OUTER_OFFSET_BIT = wp.constant(wp.uint8(3))
 
-_mat32 = wp.mat(shape=(3, 2), dtype=float)
-
 
 @wp.func
 def _add_axis_flag(ijk: wp.vec3i, axis: int):
     coord = ijk[axis]
-    ijk[axis] = wp.select(coord < 0, coord | GRID_AXIS_FLAG, coord & (~GRID_AXIS_FLAG))
+    ijk[axis] = wp.where(coord < 0, coord & (~GRID_AXIS_FLAG), coord | GRID_AXIS_FLAG)
     return ijk
 
 
@@ -176,9 +189,9 @@ class Nanogrid(Geometry):
         coords = uvw - wp.vec3(ijk)
         if cell_index == -1:
             if wp.min(coords) == 0.0 or wp.max(coords) == 1.0:
-                il = wp.select(coords[0] > 0.5, -1, 0)
-                jl = wp.select(coords[1] > 0.5, -1, 0)
-                kl = wp.select(coords[2] > 0.5, -1, 0)
+                il = wp.where(coords[0] > 0.5, 0, -1)
+                jl = wp.where(coords[1] > 0.5, 0, -1)
+                kl = wp.where(coords[2] > 0.5, 0, -1)
 
                 for n in range(8):
                     ni = n >> 2
@@ -312,7 +325,7 @@ class Nanogrid(Geometry):
         axis = Nanogrid._get_face_axis(flags)
         flip = Nanogrid._get_face_inner_offset(flags)
         v1, v2 = Nanogrid._face_tangent_vecs(args.cell_arg.cell_grid, axis, flip)
-        return _mat32(v1, v2)
+        return wp.matrix_from_cols(v1, v2)
 
     @wp.func
     def side_inner_inverse_deformation_gradient(args: SideArg, s: Sample):
@@ -394,8 +407,8 @@ class Nanogrid(Geometry):
 
         on_side = float(side_ijk[axis] - cell_ijk[axis]) == element_coords[axis]
 
-        return wp.select(
-            on_side, Coords(OUTSIDE), Coords(element_coords[(axis + 1) % 3], element_coords[(axis + 2) % 3], 0.0)
+        return wp.where(
+            on_side, Coords(element_coords[(axis + 1) % 3], element_coords[(axis + 2) % 3], 0.0), Coords(OUTSIDE)
         )
 
     @wp.func
@@ -523,8 +536,8 @@ def _build_edge_grid(cell_ijk, grid: wp.Volume, temporary_store: cache.Temporary
 
 @wp.func
 def _make_face_flags(axis: int, plus_cell_index: int, minus_cell_index: int):
-    plus_boundary = wp.uint8(wp.select(plus_cell_index == -1, 0, 1)) << FACE_OUTER_OFFSET_BIT
-    minus_boundary = wp.uint8(wp.select(minus_cell_index == -1, 0, 1)) << FACE_INNER_OFFSET_BIT
+    plus_boundary = wp.uint8(wp.where(plus_cell_index == -1, 1, 0)) << FACE_OUTER_OFFSET_BIT
+    minus_boundary = wp.uint8(wp.where(minus_cell_index == -1, 1, 0)) << FACE_INNER_OFFSET_BIT
 
     return wp.uint8(axis) | plus_boundary | minus_boundary
 

warp/fem/geometry/partition.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/geometry/quadmesh.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 Optional
 
 import warp as wp
@@ -150,13 +165,13 @@ class Quadmesh(Geometry):
         s = side_coords[0]
 
         if vs == quad_vidx[0]:
-            return wp.select(ve == quad_vidx[1], Coords(0.0, s, 0.0), Coords(s, 0.0, 0.0))
+            return wp.where(ve == quad_vidx[1], Coords(s, 0.0, 0.0), Coords(0.0, s, 0.0))
         elif vs == quad_vidx[1]:
-            return wp.select(ve == quad_vidx[2], Coords(1.0 - s, 0.0, 0.0), Coords(1.0, s, 0.0))
+            return wp.where(ve == quad_vidx[2], Coords(1.0, s, 0.0), Coords(1.0 - s, 0.0, 0.0))
         elif vs == quad_vidx[2]:
-            return wp.select(ve == quad_vidx[3], Coords(1.0, 1.0 - s, 0.0), Coords(1.0 - s, 1.0, 0.0))
+            return wp.where(ve == quad_vidx[3], Coords(1.0 - s, 1.0, 0.0), Coords(1.0, 1.0 - s, 0.0))
 
-        return wp.select(ve == quad_vidx[0], Coords(s, 1.0, 0.0), Coords(0.0, 1.0 - s, 0.0))
+        return wp.where(ve == quad_vidx[0], Coords(0.0, 1.0 - s, 0.0), Coords(s, 1.0, 0.0))
 
     @wp.func
     def _quad_to_edge_coords(
@@ -175,18 +190,18 @@ class Quadmesh(Geometry):
         cy = quad_coords[1]
 
         if vs == quad_vidx[0]:
-            oc = wp.select(ve == quad_vidx[1], cx, cy)
-            ec = wp.select(ve == quad_vidx[1], cy, cx)
+            oc = wp.where(ve == quad_vidx[1], cy, cx)
+            ec = wp.where(ve == quad_vidx[1], cx, cy)
         elif vs == quad_vidx[1]:
-            oc = wp.select(ve == quad_vidx[2], cy, 1.0 - cx)
-            ec = wp.select(ve == quad_vidx[2], 1.0 - cx, cy)
+            oc = wp.where(ve == quad_vidx[2], 1.0 - cx, cy)
+            ec = wp.where(ve == quad_vidx[2], cy, 1.0 - cx)
         elif vs == quad_vidx[2]:
-            oc = wp.select(ve == quad_vidx[3], 1.0 - cx, 1.0 - cy)
-            ec = wp.select(ve == quad_vidx[3], 1.0 - cy, 1.0 - cx)
+            oc = wp.where(ve == quad_vidx[3], 1.0 - cy, 1.0 - cx)
+            ec = wp.where(ve == quad_vidx[3], 1.0 - cx, 1.0 - cy)
         else:
-            oc = wp.select(ve == quad_vidx[0], 1.0 - cy, cx)
-            ec = wp.select(ve == quad_vidx[0], cx, 1.0 - cy)
-        return wp.select(oc == 0.0, Coords(OUTSIDE), Coords(ec, 0.0, 0.0))
+            oc = wp.where(ve == quad_vidx[0], cx, 1.0 - cy)
+            ec = wp.where(ve == quad_vidx[0], 1.0 - cy, cx)
+        return wp.where(oc == 0.0, Coords(ec, 0.0, 0.0), Coords(OUTSIDE))
 
     @wp.func
     def boundary_side_index(args: SideIndexArg, boundary_side_index: int):