warp-lang 1.4.1__py3-none-win_amd64.whl → 1.5.0__py3-none-win_amd64.whl

warp/examples/fem/example_magnetostatics.py

@@ -8,17 +8,18 @@
 ###########################################################################
 # Example Magnetostatics
 #
-# This example demonstrates solving an in-plane magnetostatics problem
-# using a curl-curl formulation
+# This example demonstrates solving a 3d magnetostatics problem
+# (a copper coil with radial current around a cylindrical iron core)
+# using a curl-curl formulation and H(curl)-conforming function space
 #
 # 1/mu Curl B + j = 0
 # Div. B = 0
 #
-# solved over field A such that B = Curl A, A = (0, 0, a_z),
-# and a_z = 0 on the domain boundary
+# solved over field A such that B = Curl A,
+# and Direchlet homogeneous essential boundary conditions
 #
-# This example also illustrates using an ImplictField to warp a square mesh
-# to a circular domain
+# This example also illustrates using an ImplictField to warp a grid mesh
+# to a cylindrical domain
 ###########################################################################
 
 import numpy as np
@@ -27,138 +28,179 @@ import warp as wp
 import warp.examples.fem.utils as fem_example_utils
 import warp.fem as fem
 
-# Vacuum and copper magnetic permeabilities
-MU_0 = wp.constant(np.pi * 4.0e-7)
-MU_C = wp.constant(1.25e-6)
+# Physics constants
+MU_0 = wp.constant(np.pi * 4.0e-7)  # Vacuum magnetic permeability
+MU_c = wp.constant(1.25e-6)  # Copper magnetic permeability
+MU_i = wp.constant(6.0e-3)  # Iron magnetic permeability
 
 
 @wp.func
-def square_to_disk(x: wp.vec2):
+def cube_to_cylinder(x: wp.vec3):
     # mapping from unit square to unit disk
-    return wp.normalize(x) * wp.max(wp.abs(x))
+    pos_xz = wp.vec3(x[0], 0.0, x[2])
+    return wp.max(wp.abs(pos_xz)) * wp.normalize(pos_xz) + wp.vec3(0.0, x[1], 0.0)
 
 
 @wp.func
-def square_to_disk_grad(x: wp.vec2):
+def cube_to_cylinder_grad(x: wp.vec3):
     # gradient of mapping from unit square to unit disk
-    if x == wp.vec2(0.0):
-        return wp.mat22(0.0)
+    pos_xz = wp.vec3(x[0], 0.0, x[2])
+    if pos_xz == wp.vec3(0.0):
+        grad = wp.mat33(0.0)
+    else:
+        dir_xz = wp.normalize(pos_xz)
+        dir_grad = (wp.identity(n=3, dtype=float) - wp.outer(dir_xz, dir_xz)) / wp.length(pos_xz)
+
+        abs_xz = wp.abs(pos_xz)
+        xinf_grad = wp.select(
+            abs_xz[0] > abs_xz[2], wp.vec3(0.0, 0.0, wp.sign(pos_xz[2])), wp.vec(wp.sign(pos_xz[0]), 0.0, 0.0)
+        )
+        grad = dir_grad * wp.max(abs_xz) + wp.outer(dir_xz, xinf_grad)
 
-    d = wp.normalize(x)
-    d_grad = (wp.identity(n=2, dtype=float) - wp.outer(d, d)) / wp.length(x)
+    grad[1, 1] = 1.0
+    return grad
 
-    ax = wp.abs(x)
-    xinf = wp.max(ax)
 
-    xinf_grad = wp.select(ax[0] > ax[1], wp.vec2(0.0, wp.sign(x[1])), wp.vec(wp.sign(x[0]), 0.0))
+@wp.func
+def permeability_field(
+    pos: wp.vec3,
+    core_radius: float,
+    core_height: float,
+    coil_internal_radius: float,
+    coil_external_radius: float,
+    coil_height: float,
+):
+    x = wp.abs(pos[0])
+    y = wp.abs(pos[1])
+    z = wp.abs(pos[2])
 
-    return d_grad * xinf + wp.outer(d, xinf_grad)
+    r = wp.sqrt(x * x + z * z)
 
+    if r <= core_radius:
+        return wp.select(y < core_height, MU_0, MU_i)
 
-@wp.func
-def permeability_field(pos: wp.vec2, coil_height: float, coil_internal_radius: float, coil_external_radius: float):
-    # space-varying permeability
+    if r >= coil_internal_radius and r <= coil_external_radius:
+        return wp.select(y < coil_height, MU_0, MU_c)
 
-    x = wp.abs(pos[0])
-    y = wp.abs(pos[1])
-    return wp.select(
-        y < coil_height and x > coil_internal_radius and x < coil_external_radius,
-        MU_0,
-        MU_C,
-    )
+    return MU_0
 
 
 @wp.func
 def current_field(
-    pos: wp.vec2, coil_height: float, coil_internal_radius: float, coil_external_radius: float, current: float
+    pos: wp.vec3,
+    current: float,
+    coil_internal_radius: float,
+    coil_external_radius: float,
+    coil_height: float,
 ):
-    # space-varying current direction along z axis (0, +1 or -1)
-    x = wp.abs(pos[0])
+    x = pos[0]
     y = wp.abs(pos[1])
-    return (
-        wp.select(
-            y < coil_height and x > coil_internal_radius and x < coil_external_radius,
-            0.0,
-            wp.sign(pos[0]),
-        )
-        * current
-    )
+    z = pos[2]
 
+    r = wp.sqrt(x * x + z * z)
 
-@fem.integrand
-def curl_z(u: fem.Field, s: fem.Sample):
-    # projection of curl((0, 0, u)) over z axis
-    du = fem.grad(u, s)
-    return wp.vec2(du[1], -du[0])
+    return wp.select(
+        y < coil_height and r >= coil_internal_radius and r <= coil_external_radius,
+        wp.vec3(0.0),
+        wp.vec3(z, 0.0, -x) * current / r,
+    )
 
 
 @fem.integrand
 def curl_curl_form(s: fem.Sample, domain: fem.Domain, u: fem.Field, v: fem.Field, mu: fem.Field):
-    return wp.dot(curl_z(u, s), curl_z(v, s)) / mu(s)
+    return wp.dot(fem.curl(u, s), fem.curl(v, s)) / mu(s)
 
 
 @fem.integrand
 def mass_form(s: fem.Sample, domain: fem.Domain, v: fem.Field, u: fem.Field):
-    return u(s) * v(s)
+    return wp.dot(u(s), v(s))
+
+
+@fem.integrand
+def curl_expr(s: fem.Sample, u: fem.Field):
+    return fem.curl(u, s)
 
 
 class Example:
-    def __init__(self, quiet=False, degree=2, resolution=32, domain_radius=2.0, current=1.0e6):
+    def __init__(self, quiet=False, mesh: str = "grid", resolution=32, domain_radius=2.0, current=1.0e6):
         # We mesh the unit disk by first meshing the unit square, then building a deformed geometry
         # from an implicit mapping field
-        square_geo = fem.Grid2D(
-            bounds_lo=wp.vec2(-domain_radius, -domain_radius),
-            bounds_hi=wp.vec2(domain_radius, domain_radius),
-            res=wp.vec2i(resolution, resolution),
-        )
 
-        def_field = fem.ImplicitField(domain=fem.Cells(square_geo), func=square_to_disk, grad_func=square_to_disk_grad)
-        disk_geo = def_field.make_deformed_geometry(relative=False)
+        if mesh == "hex":
+            positions, hex_vidx = fem_example_utils.gen_hexmesh(
+                bounds_lo=wp.vec3(-domain_radius, -domain_radius, -domain_radius),
+                bounds_hi=wp.vec3(domain_radius, domain_radius, domain_radius),
+                res=wp.vec3i(resolution, resolution, resolution),
+            )
+            cube_geo = fem.Hexmesh(hex_vertex_indices=hex_vidx, positions=positions)
+        elif mesh == "tet":
+            positions, tet_vidx = fem_example_utils.gen_tetmesh(
+                bounds_lo=wp.vec3(-domain_radius, -domain_radius, -domain_radius),
+                bounds_hi=wp.vec3(domain_radius, domain_radius, domain_radius),
+                res=wp.vec3i(resolution, resolution, resolution),
+            )
+            cube_geo = fem.Tetmesh(tet_vertex_indices=tet_vidx, positions=positions)
+        elif mesh == "nano":
+            vol = fem_example_utils.gen_volume(
+                bounds_lo=wp.vec3(-domain_radius, -domain_radius, -domain_radius),
+                bounds_hi=wp.vec3(domain_radius, domain_radius, domain_radius),
+                res=wp.vec3i(resolution, resolution, resolution),
+            )
+            cube_geo = fem.Nanogrid(grid=vol)
+        else:
+            cube_geo = fem.Grid3D(
+                bounds_lo=wp.vec3(-domain_radius, -domain_radius, -domain_radius),
+                bounds_hi=wp.vec3(domain_radius, domain_radius, domain_radius),
+                res=wp.vec3i(resolution, resolution, resolution),
+            )
+
+        def_field = fem.ImplicitField(
+            domain=fem.Cells(cube_geo), func=cube_to_cylinder, grad_func=cube_to_cylinder_grad
+        )
+        sim_geo = def_field.make_deformed_geometry(relative=False)
 
-        coil_config = {"coil_height": 1.0, "coil_internal_radius": 0.1, "coil_external_radius": 0.3}
+        coil_config = {"coil_height": 0.25, "coil_internal_radius": 0.3, "coil_external_radius": 0.4}
+        core_config = {"core_height": 1.0, "core_radius": 0.2}
 
-        domain = fem.Cells(disk_geo)
-        self._permeability_field = fem.ImplicitField(domain, func=permeability_field, values=coil_config)
+        domain = fem.Cells(sim_geo)
+        self._permeability_field = fem.ImplicitField(
+            domain, func=permeability_field, values=dict(**coil_config, **core_config)
+        )
         self._current_field = fem.ImplicitField(domain, func=current_field, values=dict(current=current, **coil_config))
 
-        z_space = fem.make_polynomial_space(disk_geo, degree=degree, element_basis=fem.ElementBasis.LAGRANGE)
-        xy_space = fem.make_polynomial_space(
-            disk_geo, degree=degree, element_basis=fem.ElementBasis.LAGRANGE, dtype=wp.vec2
+        A_space = fem.make_polynomial_space(
+            sim_geo, degree=1, element_basis=fem.ElementBasis.NEDELEC_FIRST_KIND, dtype=wp.vec3
         )
+        self.A_field = A_space.make_field()
 
-        self.A_field = z_space.make_field()
-        self.B_field = xy_space.make_field()
+        B_space = fem.make_polynomial_space(sim_geo, degree=1, element_basis=fem.ElementBasis.LAGRANGE, dtype=wp.vec3)
+        self.B_field = B_space.make_field()
 
         self.renderer = fem_example_utils.Plot()
 
     def step(self):
-        z_space = self.A_field.space
-        disk_geo = z_space.geometry
+        A_space = self.A_field.space
+        sim_geo = A_space.geometry
 
-        u = fem.make_trial(space=z_space)
-        v = fem.make_test(space=z_space)
-        lhs = fem.integrate(curl_curl_form, fields={"u": u, "v": v, "mu": self._permeability_field})
-        rhs = fem.integrate(mass_form, fields={"v": v, "u": self._current_field})
+        u = fem.make_trial(space=A_space)
+        v = fem.make_test(space=A_space)
+        lhs = fem.integrate(curl_curl_form, fields={"u": u, "v": v, "mu": self._permeability_field}, output_dtype=float)
+        rhs = fem.integrate(mass_form, fields={"v": v, "u": self._current_field}, output_dtype=float)
 
         # Dirichlet BC
-        boundary = fem.BoundarySides(disk_geo)
-        u_bd = fem.make_trial(space=z_space, domain=boundary)
-        v_bd = fem.make_test(space=z_space, domain=boundary)
-        dirichlet_bd_proj = fem.integrate(mass_form, fields={"u": u_bd, "v": v_bd}, nodal=True)
+        boundary = fem.BoundarySides(sim_geo)
+        u_bd = fem.make_trial(space=A_space, domain=boundary)
+        v_bd = fem.make_test(space=A_space, domain=boundary)
+        dirichlet_bd_proj = fem.integrate(mass_form, fields={"u": u_bd, "v": v_bd}, nodal=True, output_dtype=float)
         fem.project_linear_system(lhs, rhs, dirichlet_bd_proj)
 
-        x = wp.zeros_like(rhs)
-        fem_example_utils.bsr_cg(lhs, b=rhs, x=x, tol=1.0e-8, quiet=False)
-
-        # make sure result is exactly zero outside of circle
-        wp.sparse.bsr_mv(dirichlet_bd_proj, x=x, y=x, alpha=-1.0, beta=1.0)
-        wp.utils.array_cast(in_array=x, out_array=self.A_field.dof_values)
+        # solve using Conjugate Residual (numerically rhs may not be in image of lhs)
+        fem_example_utils.bsr_cg(lhs, b=rhs, x=self.A_field.dof_values, method="cr", max_iters=250, quiet=False)
 
         # compute B as curl(A)
-        fem.interpolate(curl_z, dest=self.B_field, fields={"u": self.A_field})
+        fem.interpolate(curl_expr, dest=self.B_field, fields={"u": self.A_field})
 
     def render(self):
-        self.renderer.add_field("A", self.A_field)
         self.renderer.add_field("B", self.B_field)
 
 
@@ -170,7 +212,7 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
     parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")
     parser.add_argument("--resolution", type=int, default=32, help="Grid resolution.")
-    parser.add_argument("--degree", type=int, default=2, help="Polynomial degree of shape functions.")
+    parser.add_argument("--mesh", type=str, default="grid", choices=["tet", "hex", "grid", "nano"], help="Mesh type.")
     parser.add_argument("--radius", type=float, default=2.0, help="Radius of simulation domain.")
     parser.add_argument(
         "--headless",
@@ -182,9 +224,9 @@ if __name__ == "__main__":
     args = parser.parse_known_args()[0]
 
     with wp.ScopedDevice(args.device):
-        example = Example(quiet=args.quiet, degree=args.degree, resolution=args.resolution, domain_radius=args.radius)
+        example = Example(quiet=args.quiet, mesh=args.mesh, resolution=args.resolution, domain_radius=args.radius)
         example.step()
         example.render()
 
         if not args.headless:
-            example.renderer.plot({"A": {"contours": {"levels": 30}}, "B": {"streamlines": {"density": 1.0}}})
+            example.renderer.plot({"B": {"streamlines": {"density": 1.0}}})

warp/examples/fem/example_nonconforming_contact.py

@@ -210,11 +210,11 @@ class Example:
 
         # Add boundary stress from other solid field
         other_stress_field = fem.field.field.NonconformingField(boundary, other_stress_field)
-        fem.utils.array_axpy(
-            y=u_rhs,
-            x=fem.integrate(
-                boundary_stress_form, fields={"u": u_bd_test, "tau": other_stress_field}, output_dtype=wp.vec2d
-            ),
+        fem.integrate(
+            boundary_stress_form,
+            fields={"u": u_bd_test, "tau": other_stress_field},
+            output=u_rhs,
+            add=True,
         )
 
         # Enforce boundary conditions

warp/examples/fem/example_stokes.py

@@ -200,4 +200,6 @@ if __name__ == "__main__":
         example.render()
 
         if not args.headless:
-            example.renderer.plot(options={"velocity": {"streamlines": {}}, "pressure": {"contours": {}}})
+            example.renderer.plot(
+                options={"velocity": {"streamlines": {}}, "pressure": {"contours": {}}}, backend="matplotlib"
+            )

warp/examples/fem/example_streamlines.py

@@ -82,7 +82,7 @@ def mass_form(
     u: fem.Field,
     v: fem.Field,
 ):
-    return u(s) * v(s)
+    return wp.dot(u(s), v(s))
 
 
 @fem.integrand
@@ -199,7 +199,7 @@ class Example:
             domain=self._inflow, order=self._degree, family=fem.Polynomial.GAUSS_LEGENDRE
         )
         n_streamlines = streamline_spawn.total_point_count()
-        spawn_points = wp.array(dtype=wp.vec3, shape=n_streamlines)
+        spawn_points = wp.empty(dtype=wp.vec3, shape=n_streamlines)
 
         jitter_amount = self._streamline_dx / self._degree
         fem.interpolate(
@@ -212,8 +212,8 @@ class Example:
         # to populate the per-point data
 
         point_count = self._streamline_point_count
-        points = wp.array(dtype=wp.vec3, shape=(n_streamlines, point_count))
-        speed = wp.array(dtype=float, shape=(n_streamlines, point_count))
+        points = wp.empty(dtype=wp.vec3, shape=(n_streamlines, point_count))
+        speed = wp.empty(dtype=float, shape=(n_streamlines, point_count))
 
         fem.interpolate(
             gen_streamlines,
@@ -235,7 +235,7 @@ class Example:
     def render(self):
         # self.renderer.add_field("solution", self.pressure_field)
         self.plot.add_field("pressure", self.pressure_field)
-        self.plot.add_field("velocity", self.velocity_field)
+        # self.plot.add_field("velocity", self.velocity_field)
 
         if self.renderer is not None:
             streamline_count = self._points.shape[0]
@@ -259,10 +259,11 @@ class Example:
             self.renderer.end_frame()
 
     def _generate_incompressible_flow(self):
-        # Function spaces for velocity, scalars and pressure (Pk / Pk / Pk-1)
-        u_space = fem.make_polynomial_space(geo=self._geo, degree=self._degree, dtype=wp.vec3)
-        s_space = fem.make_polynomial_space(geo=self._geo, degree=self._degree, dtype=float)
-        p_space = fem.make_polynomial_space(geo=self._geo, degree=self._degree - 1, dtype=float)
+        # Function spaces for velocity and pressure (RT1 / P0)
+        u_space = fem.make_polynomial_space(
+            geo=self._geo, element_basis=fem.ElementBasis.RAVIART_THOMAS, degree=1, dtype=wp.vec3
+        )
+        p_space = fem.make_polynomial_space(geo=self._geo, degree=0, dtype=float)
 
         self.pressure_field = p_space.make_field()
         self.velocity_field = u_space.make_field()
@@ -288,8 +289,8 @@ class Example:
         fem.interpolate(inflow_velocity, dest=fem.make_restriction(self.velocity_field, domain=self._inflow))
 
         # (Diagonal) mass matrix
-        rho_test = fem.make_test(s_space)
-        rho_trial = fem.make_trial(s_space)
+        rho_test = fem.make_test(u_space)
+        rho_trial = fem.make_trial(u_space)
         inv_mass_matrix = fem.integrate(
             mass_form, fields={"u": rho_trial, "v": rho_test}, nodal=True, output_dtype=float
         )
@@ -341,11 +342,3 @@ if __name__ == "__main__":
 
         example.step()
         example.render()
-
-        if not args.headless:
-            example.plot.plot(
-                {
-                    "velocity": {"streamlines": {"density": 2}},
-                    "pressure": {"contours": {}},
-                }
-            )

warp/examples/fem/utils.py

@@ -31,7 +31,7 @@ def gen_trimesh(res, bounds_lo: Optional[wp.vec2] = None, bounds_hi: Optional[wp
     Args:
         res: Resolution of the grid along each dimension
         bounds_lo: Position of the lower bound of the axis-aligned grid
-        bounds_up: Position of the upper bound of the axis-aligned grid
+        bounds_hi: Position of the upper bound of the axis-aligned grid
 
     Returns:
         Tuple of ndarrays: (Vertex positions, Triangle vertex indices)
@@ -62,7 +62,7 @@ def gen_tetmesh(res, bounds_lo: Optional[wp.vec3] = None, bounds_hi: Optional[wp
     Args:
         res: Resolution of the grid along each dimension
         bounds_lo: Position of the lower bound of the axis-aligned grid
-        bounds_up: Position of the upper bound of the axis-aligned grid
+        bounds_hi: Position of the upper bound of the axis-aligned grid
 
     Returns:
         Tuple of ndarrays: (Vertex positions, Tetrahedron vertex indices)
@@ -95,7 +95,7 @@ def gen_quadmesh(res, bounds_lo: Optional[wp.vec2] = None, bounds_hi: Optional[w
     Args:
         res: Resolution of the grid along each dimension
         bounds_lo: Position of the lower bound of the axis-aligned grid
-        bounds_up: Position of the upper bound of the axis-aligned grid
+        bounds_hi: Position of the upper bound of the axis-aligned grid
 
     Returns:
         Tuple of ndarrays: (Vertex positions, Triangle vertex indices)
@@ -125,7 +125,7 @@ def gen_hexmesh(res, bounds_lo: Optional[wp.vec3] = None, bounds_hi: Optional[wp
     Args:
         res: Resolution of the grid along each dimension
         bounds_lo: Position of the lower bound of the axis-aligned grid
-        bounds_up: Position of the upper bound of the axis-aligned grid
+        bounds_hi: Position of the upper bound of the axis-aligned grid
 
     Returns:
         Tuple of ndarrays: (Vertex positions, Triangle vertex indices)
@@ -158,7 +158,7 @@ def gen_volume(res, bounds_lo: Optional[wp.vec3] = None, bounds_hi: Optional[wp.
     Args:
         res: Resolution of the grid along each dimension
         bounds_lo: Position of the lower bound of the axis-aligned grid
-        bounds_up: Position of the upper bound of the axis-aligned grid
+        bounds_hi: Position of the upper bound of the axis-aligned grid
         device: Cuda device on which to allocate the grid
     """
 
@@ -575,6 +575,7 @@ class Plot:
 
     def _plot_pyvista(self, options: Dict[str, Any]):
         import pyvista
+        import pyvista.themes
 
         grids = {}
         scales = {}
@@ -702,7 +703,7 @@ class Plot:
         subplot_rows = options.get("rows", 1)
         subplot_shape = (subplot_rows, (len(grids) + subplot_rows - 1) // subplot_rows)
 
-        plotter = pyvista.Plotter(shape=subplot_shape)
+        plotter = pyvista.Plotter(shape=subplot_shape, theme=pyvista.themes.DocumentProTheme())
         plotter.link_views()
         plotter.add_camera_orientation_widget()
         for index, (name, grid) in enumerate(grids.items()):
@@ -717,7 +718,7 @@ class Plot:
                     plotter.view_xy()
                 else:
                     plotter.add_mesh(marker)
-            elif field.space.dimension == 3:
+            elif field.space.geometry.cell_dimension == 3:
                 plotter.add_mesh_clip_plane(grid, show_edges=True, clim=value_range, assign_to_axis="z")
             else:
                 plotter.add_mesh(grid, show_edges=True, clim=value_range)
@@ -809,6 +810,8 @@ class Plot:
                 if "arrows" in args or "streamlines" in args:
                     plot_opts["glyph_scale"] = args.get("arrows", {}).get("glyph_scale", 1.0)
                     plot_fn = _plot_quivers_3d
+                elif field.space.geometry.cell_dimension == 2:
+                    plot_fn = _plot_surface
                 else:
                     plot_fn = _plot_3d_scatter
                 plot_3d = True
@@ -856,23 +859,43 @@ def _field_triangulation(field):
 
 
 def _plot_surface(field, axes, **kwargs):
-    Z = _value_or_magnitude(field.dof_values.numpy())
+    from matplotlib.cm import get_cmap
+    from matplotlib.colors import Normalize
 
-    if "clim" in kwargs:
-        axes.set_zlim(*kwargs["clim"])
+    C = _value_or_magnitude(field.dof_values.numpy())
+
+    positions = field.space.node_positions().numpy().T
+    if field.space.dimension == 3:
+        X, Y, Z = positions
+    else:
+        X, Y = positions
+        Z = C
+        axes.set_zlim(kwargs["clim"])
 
     if hasattr(field.space, "node_grid"):
         X, Y = field.space.node_grid()
-        Z = Z.reshape(X.shape)
-        return axes.plot_surface(X, Y, Z, linewidth=0.1, antialiased=False, **kwargs)
+        C = C.reshape(X.shape)
+        return axes.plot_surface(X, Y, C, linewidth=0.1, antialiased=False, **kwargs)
 
     if hasattr(field.space, "node_triangulation"):
         triangulation = _field_triangulation(field)
-        return axes.plot_trisurf(triangulation, Z, linewidth=0.1, antialiased=False, **kwargs)
+
+        if field.space.dimension == 3:
+            plot = axes.plot_trisurf(triangulation, Z, linewidth=0.1, antialiased=False)
+            # change colors -- recompute color map manually
+            vmin, vmax = kwargs["clim"]
+            norm = Normalize(vmin=vmin, vmax=vmax)
+            values = np.mean(C[triangulation.triangles], axis=1)
+            colors = get_cmap(kwargs["cmap"])(norm(values))
+            plot.set_norm(norm)
+            plot.set_fc(colors)
+        else:
+            plot = axes.plot_trisurf(triangulation, C, linewidth=0.1, antialiased=False, **kwargs)
+
+        return plot
 
     # scatter
-    X, Y = field.space.node_positions().numpy().T
-    return axes.scatter(X, Y, Z, c=Z, **kwargs)
+    return axes.scatter(X, Y, Z, c=C, **kwargs)
 
 
 def _plot_displaced_tri_mesh(field, axes, **kwargs):

warp/examples/optim/example_walker.py

@@ -21,7 +21,7 @@ import math
 import os
 
 import numpy as np
-from pxr import Usd, UsdGeom
+from pxr import Gf, Usd, UsdGeom
 
 import warp as wp
 import warp.examples
@@ -93,7 +93,7 @@ class Example:
         geom = UsdGeom.Mesh(asset_stage.GetPrimAtPath("/root/bear"))
         points = geom.GetPointsAttr().Get()
 
-        xform = geom.ComputeLocalToWorldTransform(0.0)
+        xform = Gf.Matrix4f(geom.ComputeLocalToWorldTransform(0.0))
         for i in range(len(points)):
             points[i] = xform.Transform(points[i])
 

warp/examples/sim/example_cloth.py

@@ -26,29 +26,6 @@ import warp.sim
 import warp.sim.render
 
 
-def color_lattice_grid(num_x, num_y):
-    colors = []
-    for _i in range(4):
-        colors.append([])
-
-    for xi in range(num_x + 1):
-        for yi in range(num_y + 1):
-            vId = xi * (num_y + 1) + yi
-
-            a = 1 if xi % 2 else 0
-            b = 1 if yi % 2 else 0
-
-            c = a * 2 + b
-
-            colors[c].append(vId)
-
-    colors_wp = []
-    for i_color in range(len(colors)):
-        colors_wp.append(wp.array(colors[i_color], dtype=wp.int32))
-
-    return colors_wp
-
-
 class IntegratorType(Enum):
     EULER = "euler"
     XPBD = "xpbd"
@@ -122,6 +99,8 @@ class Example:
                 tri_ke=1e4,
                 tri_ka=1e4,
                 tri_kd=1e-5,
+                edge_ke=100,
+                color_particles=True,
             )
 
         usd_stage = Usd.Stage.Open(os.path.join(warp.examples.get_asset_directory(), "bunny.usd"))
@@ -154,8 +133,6 @@ class Example:
             self.integrator = wp.sim.XPBDIntegrator(iterations=1)
         else:
             self.integrator = wp.sim.VBDIntegrator(self.model, iterations=1)
-            # we need to give VBD coloring information
-            self.model.particle_coloring = color_lattice_grid(width, height)
 
         self.state_0 = self.model.state()
         self.state_1 = self.model.state()

warp/examples/sim/example_jacobian_ik.py

@@ -41,6 +41,8 @@ def compute_endeffector_position(
 
 class Example:
     def __init__(self, stage_path="example_jacobian_ik.usd", num_envs=10):
+        rng = np.random.default_rng(42)
+
         builder = wp.sim.ModelBuilder()
 
         self.num_envs = num_envs
@@ -81,7 +83,7 @@ class Example:
             )
             self.target_origin.append((i * 2.0, 4.0, 0.0))
             # joint initial positions
-            builder.joint_q[-3:] = np.random.uniform(-0.5, 0.5, size=3)
+            builder.joint_q[-3:] = rng.uniform(-0.5, 0.5, size=3)
         self.target_origin = np.array(self.target_origin)
 
         # finalize model
@@ -207,6 +209,8 @@ if __name__ == "__main__":
 
     args = parser.parse_known_args()[0]
 
+    rng = np.random.default_rng(42)
+
     with wp.ScopedDevice(args.device):
         example = Example(stage_path=args.stage_path, num_envs=args.num_envs)
 
@@ -218,7 +222,7 @@ if __name__ == "__main__":
         for _ in range(args.num_rollouts):
             # select new random target points for all envs
             example.targets = example.target_origin.copy()
-            example.targets[:, 1:] += np.random.uniform(-0.5, 0.5, size=(example.num_envs, 2))
+            example.targets[:, 1:] += rng.uniform(-0.5, 0.5, size=(example.num_envs, 2))
 
             for iter in range(args.train_iters):
                 example.step()

warp/examples/sim/example_quadruped.py

@@ -115,10 +115,11 @@ class Example:
 
         self.model.joint_attach_ke = 16000.0
         self.model.joint_attach_kd = 200.0
+        self.use_tile_gemm = False
 
         # self.integrator = wp.sim.XPBDIntegrator()
         # self.integrator = wp.sim.SemiImplicitIntegrator()
-        self.integrator = wp.sim.FeatherstoneIntegrator(self.model)
+        self.integrator = wp.sim.FeatherstoneIntegrator(self.model, use_tile_gemm=self.use_tile_gemm)
 
         if stage_path:
             self.renderer = wp.sim.render.SimRenderer(self.model, stage_path)