warp-lang 1.3.3__py3-none-manylinux2014_aarch64.whl → 1.4.0__py3-none-manylinux2014_aarch64.whl

warp/dlpack.py

@@ -124,6 +124,8 @@ def device_to_dlpack(wp_device) -> DLDevice:
 
 
 def dtype_to_dlpack(wp_dtype) -> DLDataType:
+    if wp_dtype == warp.bool:
+        return (DLDataTypeCode.kDLBool, 8, 1)
     if wp_dtype == warp.int8:
         return (DLDataTypeCode.kDLInt, 8, 1)
     elif wp_dtype == warp.uint8:

warp/examples/benchmarks/benchmark_cloth.py

@@ -219,6 +219,16 @@ def run_benchmark(mode, dim, timers, render=False):
 
             integrator = benchmark_cloth_jax.JxIntegrator(cloth)
 
+        elif mode == "paddle_cpu":
+            import benchmark_cloth_paddle
+
+            integrator = benchmark_cloth_paddle.TrIntegrator(cloth, "cpu")
+
+        elif mode == "paddle_gpu":
+            import benchmark_cloth_paddle
+
+            integrator = benchmark_cloth_paddle.TrIntegrator(cloth, "gpu")
+
         else:
             raise RuntimeError("Unknown simulation backend")
 

warp/examples/core/example_render_opengl.py

@@ -21,6 +21,9 @@ import warp.render
 
 class Example:
     def __init__(self, num_tiles=4, custom_tile_arrangement=False):
+        if num_tiles < 1:
+            raise ValueError("num_tiles must be greater than or equal to 1.")
+
         self.renderer = wp.render.OpenGLRenderer(vsync=False)
         instance_ids = []
 
@@ -31,16 +34,15 @@ class Example:
             positions = None
             sizes = None
 
-        if num_tiles > 1:
-            # set up instances to hide one of the capsules in each tile
-            for i in range(num_tiles):
-                instances = [j for j in np.arange(13) if j != i + 2]
-                instance_ids.append(instances)
-                if custom_tile_arrangement:
-                    angle = np.pi * 2.0 / num_tiles * i
-                    positions.append((int(np.cos(angle) * 150 + 250), int(np.sin(angle) * 150 + 250)))
-                    sizes.append((150, 150))
-            self.renderer.setup_tiled_rendering(instance_ids, tile_positions=positions, tile_sizes=sizes)
+        # set up instances to hide one of the capsules in each tile
+        for i in range(num_tiles):
+            instances = [j for j in np.arange(13) if j != i + 2]
+            instance_ids.append(instances)
+            if custom_tile_arrangement:
+                angle = np.pi * 2.0 / num_tiles * i
+                positions.append((int(np.cos(angle) * 150 + 250), int(np.sin(angle) * 150 + 250)))
+                sizes.append((150, 150))
+        self.renderer.setup_tiled_rendering(instance_ids, tile_positions=positions, tile_sizes=sizes)
 
         self.renderer.render_ground()
 

warp/examples/fem/example_adaptive_grid.py

@@ -0,0 +1,251 @@
+# Copyright (c) 2024 NVIDIA CORPORATION.  All rights reserved.
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+###########################################################################
+# Example Adaptive Grid
+#
+# Demonstrates using an adaptive grid to increase the simulation resolition
+# near a collider boundary.
+#
+###########################################################################
+import os.path
+
+import numpy as np
+
+import warp as wp
+import warp.examples
+import warp.examples.fem.utils as fem_example_utils
+import warp.fem as fem
+from warp.examples.fem.example_apic_fluid import divergence_form, solve_incompressibility
+
+
+@fem.integrand
+def inflow_velocity(
+    s: fem.Sample,
+    domain: fem.Domain,
+    bounds_lo: wp.vec3,
+    bounds_hi: wp.vec3,
+):
+    x = fem.position(domain, s)
+
+    if x[1] <= bounds_lo[1] or x[2] <= bounds_lo[2] or x[1] >= bounds_hi[1] or x[2] >= bounds_hi[2]:
+        return wp.vec3(0.0)
+
+    if x[0] <= bounds_lo[0] or x[0] >= bounds_hi[0]:
+        return wp.vec3(1.0, 0.0, 0.0)
+
+    return wp.vec3(0.0)
+
+
+@fem.integrand
+def noslip_projector_form(
+    s: fem.Sample,
+    u: fem.Field,
+    v: fem.Field,
+):
+    return wp.dot(u(s), v(s))
+
+
+@fem.integrand
+def mass_form(
+    s: fem.Sample,
+    u: fem.Field,
+    v: fem.Field,
+):
+    return u(s) * v(s)
+
+
+@fem.integrand
+def side_divergence_form(s: fem.Sample, domain: fem.Domain, u: fem.Field, psi: fem.Field):
+    # normal velocity jump (non-zero at resolution boundaries)
+    return -wp.dot(fem.jump(u, s), fem.normal(domain, s)) * psi(s)
+
+
+@wp.func
+def refinement_field(xyz: wp.vec3, volume: wp.uint64):
+    # use distance to collider as refinement function
+    uvw = wp.volume_world_to_index(volume, xyz)
+    sdf = wp.volume_sample_f(volume, uvw, wp.Volume.LINEAR)
+
+    if sdf < 0.0:
+        return sdf
+
+    # combine with  heuristical distance to keep coarsening past nvdb narrowband
+    return 0.5 * wp.max(wp.length(xyz) - 20.0, sdf)
+
+
+@fem.integrand
+def pressure_anomaly_field(s: fem.Sample, domain: fem.Domain, pressure: fem.Field):
+    # for visualization, deduce affine part such that grad P = u_x
+    x = domain(s)
+    return pressure(s) + x[0]
+
+
+class Example:
+    def __init__(self, quiet=False, degree=2, base_resolution=8, level_count=4, headless: bool = False):
+        self._quiet = quiet
+        self._degree = degree
+
+        # Start from a coarse, dense grid
+        res = wp.vec3i(2 * base_resolution, base_resolution // 2, base_resolution)
+        bounds_lo = wp.vec3(-50.0, 0.0, -17.5)
+        bounds_hi = wp.vec3(50.0, 12.5, 17.5)
+        sim_vol = fem_example_utils.gen_volume(res=res, bounds_lo=bounds_lo, bounds_hi=bounds_hi)
+
+        # load collision volume
+        collider_path = os.path.join(warp.examples.get_asset_directory(), "rocks.nvdb")
+        with open(collider_path, "rb") as file:
+            # create Volume object
+            collider = wp.Volume.load_from_nvdb(file)
+
+        # Make adaptive grid from coarse base and refinement field
+        refinement = fem.ImplicitField(
+            domain=fem.Cells(fem.Nanogrid(sim_vol)), func=refinement_field, values={"volume": collider.id}
+        )
+        self._geo = fem.adaptivity.adaptive_nanogrid_from_field(
+            sim_vol, level_count, refinement_field=refinement, grading="face"
+        )
+
+        # Function spaces for velocity, scalars and pressure (Pk / Pk / Pk-1)
+        self._u_basis = fem.make_polynomial_basis_space(geo=self._geo, degree=self._degree)
+        u_space = fem.make_collocated_function_space(self._u_basis, dtype=wp.vec3)
+        p_space = fem.make_polynomial_space(geo=self._geo, degree=self._degree - 1, dtype=float)
+
+        self.pressure_field = p_space.make_field()
+        self.pressure_anomaly_field = p_space.make_field()
+        self.velocity_field = u_space.make_field()
+
+        # Initialize velocity field with BC
+        bounds_scale = 0.9999  # account for difference between bounds and actual grid extents
+        bounds_center = 0.5 * (bounds_hi + bounds_lo)
+        bounds_extent = 0.5 * (bounds_hi - bounds_lo)
+        fem.interpolate(
+            inflow_velocity,
+            dest=fem.make_restriction(self.velocity_field, domain=fem.BoundarySides(self._geo)),
+            values={
+                "bounds_lo": bounds_center - bounds_scale * bounds_extent,
+                "bounds_hi": bounds_center + bounds_scale * bounds_extent,
+            },
+        )
+
+        self.plot = fem_example_utils.Plot()
+
+    def render(self):
+        # self.renderer.add_field("solution", self.pressure_field)
+        self.plot.add_field("pressure_anomaly", self.pressure_anomaly_field)
+        self.plot.add_field("velocity", self.velocity_field)
+
+    def step(self):
+        u_space = self.velocity_field.space
+        p_space = self.pressure_field.space
+
+        # Boundary condition projector and matrices
+        boundary = fem.BoundarySides(self._geo)
+        bd_test = fem.make_test(u_space, domain=boundary)
+        bd_trial = fem.make_trial(u_space, domain=boundary)
+        dirichlet_projector = fem.integrate(
+            noslip_projector_form, fields={"u": bd_test, "v": bd_trial}, nodal=True, output_dtype=float
+        )
+        fem.normalize_dirichlet_projector(dirichlet_projector)
+
+        # (Diagonal) mass matrix
+        s_space = fem.make_collocated_function_space(self._u_basis, dtype=float)
+        rho_test = fem.make_test(s_space)
+        rho_trial = fem.make_trial(s_space)
+        inv_mass_matrix = fem.integrate(
+            mass_form, fields={"u": rho_trial, "v": rho_test}, nodal=True, output_dtype=float
+        )
+        fem_example_utils.invert_diagonal_bsr_matrix(inv_mass_matrix)
+
+        # Assemble divergence operator matrix
+        p_test = fem.make_test(p_space)
+        u_trial = fem.make_trial(u_space)
+        divergence_matrix = fem.integrate(
+            divergence_form,
+            fields={"u": u_trial, "psi": p_test},
+            output_dtype=float,
+        )
+
+        # need to account for discontinuities at resolution boundaries (t-junctions)
+        p_side_test = fem.make_test(p_space, domain=fem.Sides(self._geo))
+        u_side_trial = fem.make_trial(u_space, domain=fem.Sides(self._geo))
+        divergence_matrix += fem.integrate(
+            side_divergence_form,
+            fields={"u": u_side_trial, "psi": p_side_test},
+            output_dtype=float,
+        )
+
+        # Solve incompressibility
+        solve_incompressibility(
+            divergence_matrix,
+            dirichlet_projector,
+            inv_mass_matrix.values,
+            self.pressure_field.dof_values,
+            self.velocity_field.dof_values,
+            quiet=self._quiet,
+        )
+
+        fem.interpolate(
+            pressure_anomaly_field,
+            dest=self.pressure_anomaly_field,
+            fields={"pressure": self.pressure_field},
+        )
+
+
+if __name__ == "__main__":
+    import argparse
+
+    wp.set_module_options({"enable_backward": False})
+
+    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=8, help="Grid resolution.")
+    parser.add_argument("--degree", type=int, default=2, help="Polynomial degree of shape functions.")
+    parser.add_argument("--level_count", type=int, default=4, help="Number of refinement levels.")
+    parser.add_argument(
+        "--headless",
+        action="store_true",
+        help="Run in headless mode, suppressing the opening of any graphical windows.",
+    )
+    parser.add_argument("--quiet", action="store_true", help="Suppresses the printing out of iteration residuals.")
+
+    args = parser.parse_known_args()[0]
+
+    with wp.ScopedDevice(args.device):
+        example = Example(
+            quiet=args.quiet,
+            degree=args.degree,
+            base_resolution=args.resolution,
+            level_count=args.level_count,
+            headless=args.headless,
+        )
+
+        example.step()
+        example.render()
+
+        if not args.headless:
+            ref_geom = None
+            try:
+                from pxr import Usd, UsdGeom
+
+                stage = Usd.Stage.Open(os.path.join(warp.examples.get_asset_directory(), "rocks.usd"))
+                mesh = UsdGeom.Mesh(stage.GetPrimAtPath("/root/rocks"))
+                points = np.array((mesh.GetPointsAttr().Get()))
+                counts = np.array((mesh.GetFaceVertexCountsAttr().Get()))
+                indices = np.array(mesh.GetFaceVertexIndicesAttr().Get())
+                ref_geom = (points, counts, indices)
+            except Exception:
+                pass
+
+            example.plot.plot(
+                {
+                    "rows": 2,
+                    "ref_geom": ref_geom,
+                    "velocity": {"streamlines": {"density": 25, "glyph_scale": 0.01}},
+                    "pressure_anomaly": {},
+                }
+            )

warp/examples/fem/example_apic_fluid.py

@@ -169,7 +169,7 @@ def solve_incompressibility(
     # For simplicity, assemble Schur complement and solve with CG
     schur = bsr_mm(divergence_mat, transposed_divergence_mat)
 
-    fem_example_utils.bsr_cg(schur, b=rhs, x=pressure, quiet=quiet, tol=1.0e-6)
+    fem_example_utils.bsr_cg(schur, b=rhs, x=pressure, quiet=quiet, tol=1.0e-6, method="cr", max_iters=1000)
 
     # Apply pressure to velocity
     bsr_mv(A=transposed_divergence_mat, x=pressure, y=velocity, alpha=1.0, beta=1.0)

warp/examples/fem/example_diffusion_3d.py

@@ -146,7 +146,7 @@ if __name__ == "__main__":
     parser.add_argument(
         "--boundary_compliance", type=float, default=0.0, help="Dirichlet boundary condition compliance."
     )
-    parser.add_argument("--mesh", choices=("grid", "tet", "hex", "nano"), default="grid", help="Mesh type.")
+    parser.add_argument("--mesh", choices=("grid", "tet", "hex", "nano", "anano"), default="grid", help="Mesh type.")
     parser.add_argument(
         "--headless",
         action="store_true",
@@ -171,4 +171,4 @@ if __name__ == "__main__":
         example.render()
 
         if not args.headless:
-            example.renderer.plot()
+            example.renderer.plot(backend="matplotlib")

warp/examples/fem/example_magnetostatics.py

@@ -150,7 +150,7 @@ class Example:
         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 outisde of circle
+        # 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)
 

warp/examples/fem/example_streamlines.py

@@ -255,6 +255,7 @@ class Example:
             self.renderer.render_line_list("streamlines", vertices, indices)
             self.renderer.render_line_list("streamlines", vertices, indices, colors)
 
+            self.renderer.paused = True
             self.renderer.end_frame()
 
     def _generate_incompressible_flow(self):

warp/examples/fem/utils.py

@@ -177,7 +177,9 @@ def gen_volume(res, bounds_lo: Optional[wp.vec3] = None, bounds_hi: Optional[wp.
 
     ijk = np.transpose(np.meshgrid(x, y, z), axes=(1, 2, 3, 0)).reshape(-1, 3)
     ijk = wp.array(ijk, dtype=wp.vec3i, device=device)
-    return wp.Volume.allocate_by_voxels(ijk, voxel_size=voxel_size, translation=0.5 * voxel_size, device=device)
+    return wp.Volume.allocate_by_voxels(
+        ijk, voxel_size=voxel_size, translation=bounds_lo + 0.5 * voxel_size, device=device
+    )
 
 
 #
@@ -579,6 +581,19 @@ class Plot:
 
         animate = False
 
+        ref_geom = options.get("ref_geom", None)
+        if ref_geom is not None:
+            if isinstance(ref_geom, tuple):
+                vertices, counts, indices = ref_geom
+                offsets = np.cumsum(counts)
+                ranges = np.array([offsets - counts, offsets]).T
+                faces = np.concatenate(
+                    [[count] + list(indices[beg:end]) for (count, (beg, end)) in zip(counts, ranges)]
+                )
+                ref_geom = pyvista.PolyData(vertices, faces)
+            else:
+                ref_geom = pyvista.PolyData(ref_geom)
+
         for name, (field, values) in self._fields.items():
             cells, types = field.space.vtk_cells()
             node_pos = field.space.node_positions().numpy()
@@ -667,7 +682,7 @@ class Plot:
                         density = field_args["streamlines"].get("density", 1.0)
                         cell_size = 1.0 / np.sqrt(field.space.geometry.cell_count())
                         lines = grid.streamlines(vectors=name, n_points=int(100 * density))
-                        marker = lines.tube(radius=0.0025 * grid_scale / density)
+                        marker = lines.tube(radius=0.0025 * grid_scale / np.sqrt(density))
                     elif "displacement" in field_args:
                         grid.points = field.space.node_positions().numpy() + v
 
@@ -700,11 +715,15 @@ class Plot:
                     plotter.add_mesh(grid, opacity=0.25, clim=value_range)
                     plotter.view_xy()
                 else:
-                    plotter.add_mesh(marker, opacity=0.25)
+                    plotter.add_mesh(marker)
             elif field.space.dimension == 3:
-                plotter.add_mesh_clip_plane(grid, show_edges=True, clim=value_range)
+                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)
+
+            if ref_geom:
+                plotter.add_mesh(ref_geom)
+
         plotter.show(interactive_update=animate)
 
         frame = 0

warp/examples/sim/example_cloth.py

@@ -26,9 +26,33 @@ 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"
+    VBD = "vbd"
 
     def __str__(self):
         return self.value
@@ -67,7 +91,7 @@ class Example:
                 tri_ka=1.0e3,
                 tri_kd=1.0e1,
             )
-        else:
+        elif self.integrator_type == IntegratorType.XPBD:
             builder.add_cloth_grid(
                 pos=wp.vec3(0.0, 4.0, 0.0),
                 rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), math.pi * 0.5),
@@ -83,6 +107,22 @@ class Example:
                 spring_ke=1.0e3,
                 spring_kd=0.0,
             )
+        else:
+            # VBD
+            builder.add_cloth_grid(
+                pos=wp.vec3(0.0, 4.0, 0.0),
+                rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), math.pi * 0.5),
+                vel=wp.vec3(0.0, 0.0, 0.0),
+                dim_x=self.sim_width,
+                dim_y=self.sim_height,
+                cell_x=0.1,
+                cell_y=0.1,
+                mass=0.1,
+                fix_left=True,
+                tri_ke=1e4,
+                tri_ka=1e4,
+                tri_kd=1e-5,
+            )
 
         usd_stage = Usd.Stage.Open(os.path.join(warp.examples.get_asset_directory(), "bunny.usd"))
         usd_geom = UsdGeom.Mesh(usd_stage.GetPrimAtPath("/root/bunny"))
@@ -103,16 +143,20 @@ class Example:
             kf=1.0e1,
         )
 
-        if self.integrator_type == IntegratorType.EULER:
-            self.integrator = wp.sim.SemiImplicitIntegrator()
-        else:
-            self.integrator = wp.sim.XPBDIntegrator(iterations=1)
-
         self.model = builder.finalize()
         self.model.ground = True
         self.model.soft_contact_ke = 1.0e4
         self.model.soft_contact_kd = 1.0e2
 
+        if self.integrator_type == IntegratorType.EULER:
+            self.integrator = wp.sim.SemiImplicitIntegrator()
+        elif self.integrator_type == IntegratorType.XPBD:
+            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/fem/__init__.py

@@ -1,3 +1,4 @@
+from .adaptivity import adaptive_nanogrid_from_field, adaptive_nanogrid_from_hierarchy
 from .cache import TemporaryStore, borrow_temporary, borrow_temporary_like, set_default_temporary_store
 from .dirichlet import normalize_dirichlet_projector, project_linear_system
 from .domain import BoundarySides, Cells, FrontierSides, GeometryDomain, Sides, Subdomain
@@ -13,6 +14,7 @@ from .field import (
     make_trial,
 )
 from .geometry import (
+    AdaptiveNanogrid,
     ExplicitGeometryPartition,
     Geometry,
     GeometryPartition,