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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of warp-lang might be problematic. Click here for more details.

Files changed (132) hide show
  1. warp/__init__.py +5 -0
  2. warp/autograd.py +414 -191
  3. warp/bin/warp-clang.dll +0 -0
  4. warp/bin/warp.dll +0 -0
  5. warp/build.py +40 -12
  6. warp/build_dll.py +13 -6
  7. warp/builtins.py +1124 -497
  8. warp/codegen.py +261 -136
  9. warp/config.py +1 -1
  10. warp/context.py +357 -119
  11. warp/examples/assets/square_cloth.usd +0 -0
  12. warp/examples/benchmarks/benchmark_gemm.py +27 -18
  13. warp/examples/benchmarks/benchmark_interop_paddle.py +3 -3
  14. warp/examples/benchmarks/benchmark_interop_torch.py +3 -3
  15. warp/examples/core/example_torch.py +18 -34
  16. warp/examples/fem/example_apic_fluid.py +1 -0
  17. warp/examples/fem/example_mixed_elasticity.py +1 -1
  18. warp/examples/optim/example_bounce.py +1 -1
  19. warp/examples/optim/example_cloth_throw.py +1 -1
  20. warp/examples/optim/example_diffray.py +4 -15
  21. warp/examples/optim/example_drone.py +1 -1
  22. warp/examples/optim/example_softbody_properties.py +392 -0
  23. warp/examples/optim/example_trajectory.py +1 -3
  24. warp/examples/optim/example_walker.py +5 -0
  25. warp/examples/sim/example_cartpole.py +0 -2
  26. warp/examples/sim/example_cloth.py +3 -1
  27. warp/examples/sim/example_cloth_self_contact.py +260 -0
  28. warp/examples/sim/example_granular_collision_sdf.py +4 -5
  29. warp/examples/sim/example_jacobian_ik.py +0 -2
  30. warp/examples/sim/example_quadruped.py +5 -2
  31. warp/examples/tile/example_tile_cholesky.py +79 -0
  32. warp/examples/tile/example_tile_convolution.py +2 -2
  33. warp/examples/tile/example_tile_fft.py +2 -2
  34. warp/examples/tile/example_tile_filtering.py +3 -3
  35. warp/examples/tile/example_tile_matmul.py +4 -4
  36. warp/examples/tile/example_tile_mlp.py +12 -12
  37. warp/examples/tile/example_tile_nbody.py +180 -0
  38. warp/examples/tile/example_tile_walker.py +319 -0
  39. warp/fem/geometry/geometry.py +0 -2
  40. warp/math.py +147 -0
  41. warp/native/array.h +12 -0
  42. warp/native/builtin.h +0 -1
  43. warp/native/bvh.cpp +149 -70
  44. warp/native/bvh.cu +287 -68
  45. warp/native/bvh.h +195 -85
  46. warp/native/clang/clang.cpp +5 -1
  47. warp/native/coloring.cpp +5 -1
  48. warp/native/cuda_util.cpp +91 -53
  49. warp/native/cuda_util.h +5 -0
  50. warp/native/exports.h +40 -40
  51. warp/native/intersect.h +17 -0
  52. warp/native/mat.h +41 -0
  53. warp/native/mathdx.cpp +19 -0
  54. warp/native/mesh.cpp +25 -8
  55. warp/native/mesh.cu +153 -101
  56. warp/native/mesh.h +482 -403
  57. warp/native/quat.h +40 -0
  58. warp/native/solid_angle.h +7 -0
  59. warp/native/sort.cpp +85 -0
  60. warp/native/sort.cu +34 -0
  61. warp/native/sort.h +3 -1
  62. warp/native/spatial.h +11 -0
  63. warp/native/tile.h +1187 -669
  64. warp/native/tile_reduce.h +8 -6
  65. warp/native/vec.h +41 -0
  66. warp/native/warp.cpp +8 -1
  67. warp/native/warp.cu +263 -40
  68. warp/native/warp.h +19 -5
  69. warp/optim/linear.py +22 -4
  70. warp/render/render_opengl.py +130 -64
  71. warp/sim/__init__.py +6 -1
  72. warp/sim/collide.py +270 -26
  73. warp/sim/import_urdf.py +8 -8
  74. warp/sim/integrator_euler.py +25 -7
  75. warp/sim/integrator_featherstone.py +154 -35
  76. warp/sim/integrator_vbd.py +842 -40
  77. warp/sim/model.py +134 -72
  78. warp/sparse.py +1 -1
  79. warp/stubs.py +265 -132
  80. warp/tape.py +28 -30
  81. warp/tests/aux_test_module_unload.py +15 -0
  82. warp/tests/{test_sim_grad.py → flaky_test_sim_grad.py} +104 -63
  83. warp/tests/test_array.py +74 -0
  84. warp/tests/test_assert.py +242 -0
  85. warp/tests/test_codegen.py +14 -61
  86. warp/tests/test_collision.py +2 -2
  87. warp/tests/test_coloring.py +12 -2
  88. warp/tests/test_examples.py +12 -1
  89. warp/tests/test_func.py +21 -4
  90. warp/tests/test_grad_debug.py +87 -2
  91. warp/tests/test_hash_grid.py +1 -1
  92. warp/tests/test_ipc.py +116 -0
  93. warp/tests/test_lerp.py +13 -87
  94. warp/tests/test_mat.py +138 -167
  95. warp/tests/test_math.py +47 -1
  96. warp/tests/test_matmul.py +17 -16
  97. warp/tests/test_matmul_lite.py +10 -15
  98. warp/tests/test_mesh.py +84 -60
  99. warp/tests/test_mesh_query_aabb.py +165 -0
  100. warp/tests/test_mesh_query_point.py +328 -286
  101. warp/tests/test_mesh_query_ray.py +134 -121
  102. warp/tests/test_mlp.py +2 -2
  103. warp/tests/test_operators.py +43 -0
  104. warp/tests/test_overwrite.py +47 -2
  105. warp/tests/test_quat.py +77 -0
  106. warp/tests/test_reload.py +29 -0
  107. warp/tests/test_sim_grad_bounce_linear.py +204 -0
  108. warp/tests/test_smoothstep.py +17 -83
  109. warp/tests/test_static.py +19 -3
  110. warp/tests/test_tape.py +25 -0
  111. warp/tests/test_tile.py +178 -191
  112. warp/tests/test_tile_load.py +356 -0
  113. warp/tests/test_tile_mathdx.py +61 -8
  114. warp/tests/test_tile_mlp.py +17 -17
  115. warp/tests/test_tile_reduce.py +24 -18
  116. warp/tests/test_tile_shared_memory.py +66 -17
  117. warp/tests/test_tile_view.py +165 -0
  118. warp/tests/test_torch.py +35 -0
  119. warp/tests/test_utils.py +36 -24
  120. warp/tests/test_vec.py +110 -0
  121. warp/tests/unittest_suites.py +29 -4
  122. warp/tests/unittest_utils.py +30 -13
  123. warp/thirdparty/unittest_parallel.py +2 -2
  124. warp/types.py +411 -101
  125. warp/utils.py +10 -7
  126. {warp_lang-1.5.0.dist-info → warp_lang-1.6.0.dist-info}/METADATA +92 -69
  127. {warp_lang-1.5.0.dist-info → warp_lang-1.6.0.dist-info}/RECORD +130 -119
  128. {warp_lang-1.5.0.dist-info → warp_lang-1.6.0.dist-info}/WHEEL +1 -1
  129. warp/examples/benchmarks/benchmark_tile.py +0 -179
  130. warp/native/tile_gemm.h +0 -341
  131. {warp_lang-1.5.0.dist-info → warp_lang-1.6.0.dist-info}/LICENSE.md +0 -0
  132. {warp_lang-1.5.0.dist-info → warp_lang-1.6.0.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,392 @@
1
+ # Copyright (c) 2025 NVIDIA CORPORATION. All rights reserved.
2
+ # NVIDIA CORPORATION and its licensors retain all intellectual property
3
+ # and proprietary rights in and to this software, related documentation
4
+ # and any modifications thereto. Any use, reproduction, disclosure or
5
+ # distribution of this software and related documentation without an express
6
+ # license agreement from NVIDIA CORPORATION is strictly prohibited.
7
+
8
+ ###########################################################################
9
+ # Example FEM Bounce
10
+ #
11
+ # Shows how to use Warp to optimize for the material parameters of a soft body,
12
+ # such that it bounces off the wall and floor in order to hit a target.
13
+ #
14
+ # This example uses the built-in wp.Tape() object to compute gradients of
15
+ # the distance to target (loss) w.r.t the material parameters, followed by
16
+ # a simple gradient-descent optimization step.
17
+ #
18
+ ###########################################################################
19
+
20
+ import numpy as np
21
+
22
+ import warp as wp
23
+ import warp.optim
24
+ import warp.sim
25
+ import warp.sim.render
26
+
27
+
28
+ @wp.kernel
29
+ def assign_param(params: wp.array(dtype=wp.float32), tet_materials: wp.array2d(dtype=wp.float32)):
30
+ tid = wp.tid()
31
+ params_idx = 2 * wp.tid() % params.shape[0]
32
+ tet_materials[tid, 0] = params[params_idx]
33
+ tet_materials[tid, 1] = params[params_idx + 1]
34
+
35
+
36
+ @wp.kernel
37
+ def com_kernel(particle_q: wp.array(dtype=wp.vec3), com: wp.array(dtype=wp.vec3)):
38
+ tid = wp.tid()
39
+ point = particle_q[tid]
40
+ a = point / wp.float32(particle_q.shape[0])
41
+
42
+ # Atomically add the point coordinates to the accumulator
43
+ wp.atomic_add(com, 0, a)
44
+
45
+
46
+ @wp.kernel
47
+ def loss_kernel(
48
+ target: wp.vec3,
49
+ com: wp.array(dtype=wp.vec3),
50
+ pos_error: wp.array(dtype=float),
51
+ loss: wp.array(dtype=float),
52
+ ):
53
+ diff = com[0] - target
54
+ pos_error[0] = wp.dot(diff, diff)
55
+ norm = pos_error[0]
56
+ loss[0] = norm
57
+
58
+
59
+ @wp.kernel
60
+ def enforce_constraint_kernel(lower_bound: wp.float32, upper_bound: wp.float32, x: wp.array(dtype=wp.float32)):
61
+ tid = wp.tid()
62
+ if x[tid] < lower_bound:
63
+ x[tid] = lower_bound
64
+ elif x[tid] > upper_bound:
65
+ x[tid] = upper_bound
66
+
67
+
68
+ class Example:
69
+ def __init__(
70
+ self,
71
+ stage_path="example_softbody_properties.usd",
72
+ material_behavior="anisotropic",
73
+ verbose=False,
74
+ ):
75
+ self.verbose = verbose
76
+ self.material_behavior = material_behavior
77
+
78
+ # seconds
79
+ sim_duration = 1.0
80
+
81
+ # control frequency
82
+ fps = 60
83
+ self.frame_dt = 1.0 / fps
84
+ frame_steps = int(sim_duration / self.frame_dt)
85
+
86
+ # sim frequency
87
+ self.sim_substeps = 16
88
+ self.sim_steps = frame_steps * self.sim_substeps
89
+ self.sim_dt = self.frame_dt / self.sim_substeps
90
+
91
+ self.iter = 0
92
+ self.render_time = 0.0
93
+
94
+ self.train_rate = 1e7
95
+
96
+ self.losses = []
97
+
98
+ self.hard_lower_bound = wp.float32(500.0)
99
+ self.hard_upper_bound = wp.float32(4e6)
100
+
101
+ # Create FEM model.
102
+ self.cell_dim = 2
103
+ self.cell_size = 0.1
104
+ center = self.cell_size * self.cell_dim * 0.5
105
+ self.grid_origin = wp.vec3(-0.5, 1.0, -center)
106
+ self.create_model()
107
+
108
+ self.integrator = wp.sim.SemiImplicitIntegrator()
109
+
110
+ self.target = wp.vec3(-1.0, 1.5, 0.0)
111
+ # Initialize material parameters
112
+ if self.material_behavior == "anisotropic":
113
+ # Different Lame parameters for each tet
114
+ self.material_params = wp.array(
115
+ self.model.tet_materials.numpy()[:, :2].flatten(),
116
+ dtype=wp.float32,
117
+ requires_grad=True,
118
+ )
119
+ else:
120
+ # Same Lame parameters for all tets
121
+ self.material_params = wp.array(
122
+ self.model.tet_materials.numpy()[0, :2].flatten(),
123
+ dtype=wp.float32,
124
+ requires_grad=True,
125
+ )
126
+
127
+ self.optimizer = wp.optim.SGD(
128
+ [self.material_params],
129
+ lr=self.train_rate,
130
+ nesterov=False,
131
+ )
132
+
133
+ self.com = wp.array([wp.vec3(0.0, 0.0, 0.0)], dtype=wp.vec3, requires_grad=True)
134
+ self.pos_error = wp.zeros(1, dtype=wp.float32, requires_grad=True)
135
+ self.loss = wp.zeros(1, dtype=wp.float32, requires_grad=True)
136
+
137
+ # allocate sim states for trajectory
138
+ self.states = []
139
+ for _i in range(self.sim_steps + 1):
140
+ self.states.append(self.model.state())
141
+
142
+ if stage_path:
143
+ self.renderer = wp.sim.render.SimRenderer(self.model, stage_path, scaling=1.0)
144
+ else:
145
+ self.renderer = None
146
+
147
+ # capture forward/backward passes
148
+ self.use_cuda_graph = wp.get_device().is_cuda
149
+ if self.use_cuda_graph:
150
+ with wp.ScopedCapture() as capture:
151
+ self.tape = wp.Tape()
152
+ with self.tape:
153
+ self.forward()
154
+ self.tape.backward(self.loss)
155
+ self.graph = capture.graph
156
+
157
+ def create_model(self):
158
+ builder = wp.sim.ModelBuilder()
159
+ builder.default_particle_radius = 0.0005
160
+
161
+ total_mass = 0.2
162
+ num_particles = (self.cell_dim + 1) ** 3
163
+ particle_mass = total_mass / num_particles
164
+ particle_density = particle_mass / (self.cell_size**3)
165
+ if self.verbose:
166
+ print(f"Particle density: {particle_density}")
167
+
168
+ young_mod = 1.5 * 1e4
169
+ poisson_ratio = 0.3
170
+ k_mu = 0.5 * young_mod / (1.0 + poisson_ratio)
171
+ k_lambda = young_mod * poisson_ratio / ((1 + poisson_ratio) * (1 - 2 * poisson_ratio))
172
+
173
+ builder.add_soft_grid(
174
+ pos=self.grid_origin,
175
+ rot=wp.quat_identity(),
176
+ vel=wp.vec3(5.0, -5.0, 0.0),
177
+ dim_x=self.cell_dim,
178
+ dim_y=self.cell_dim,
179
+ dim_z=self.cell_dim,
180
+ cell_x=self.cell_size,
181
+ cell_y=self.cell_size,
182
+ cell_z=self.cell_size,
183
+ density=particle_density,
184
+ k_mu=k_mu,
185
+ k_lambda=k_lambda,
186
+ k_damp=0.0,
187
+ tri_ke=1e-4,
188
+ tri_ka=1e-4,
189
+ tri_kd=1e-4,
190
+ tri_drag=0.0,
191
+ tri_lift=0.0,
192
+ fix_bottom=False,
193
+ )
194
+
195
+ ke = 1.0e3
196
+ kf = 0.0
197
+ kd = 1.0e0
198
+ mu = 0.2
199
+ builder.add_shape_box(
200
+ body=-1,
201
+ pos=wp.vec3(2.0, 1.0, 0.0),
202
+ hx=0.25,
203
+ hy=1.0,
204
+ hz=1.0,
205
+ ke=ke,
206
+ kf=kf,
207
+ kd=kd,
208
+ mu=mu,
209
+ )
210
+
211
+ # use `requires_grad=True` to create a model for differentiable simulation
212
+ self.model = builder.finalize(requires_grad=True)
213
+ self.model.ground = True
214
+
215
+ self.model.soft_contact_ke = ke
216
+ self.model.soft_contact_kf = kf
217
+ self.model.soft_contact_kd = kd
218
+ self.model.soft_contact_mu = mu
219
+ self.model.soft_contact_margin = 0.001
220
+ self.model.soft_contact_restitution = 1.0
221
+
222
+ def forward(self):
223
+ wp.launch(
224
+ kernel=assign_param,
225
+ dim=self.model.tet_count,
226
+ inputs=(self.material_params,),
227
+ outputs=(self.model.tet_materials,),
228
+ )
229
+ # run control loop
230
+ for i in range(self.sim_steps):
231
+ wp.sim.collide(self.model, self.states[i])
232
+ self.states[i].clear_forces()
233
+
234
+ self.integrator.simulate(self.model, self.states[i], self.states[i + 1], self.sim_dt)
235
+
236
+ # Update loss
237
+ # Compute the center of mass for the last time step.
238
+ wp.launch(
239
+ kernel=com_kernel,
240
+ dim=self.model.particle_count,
241
+ inputs=(self.states[-1].particle_q,),
242
+ outputs=(self.com,),
243
+ )
244
+
245
+ # calculate loss
246
+ wp.launch(
247
+ kernel=loss_kernel,
248
+ dim=1,
249
+ inputs=(
250
+ self.target,
251
+ self.com,
252
+ ),
253
+ outputs=(self.pos_error, self.loss),
254
+ )
255
+
256
+ return self.loss
257
+
258
+ def step(self):
259
+ with wp.ScopedTimer("step"):
260
+ if self.use_cuda_graph:
261
+ wp.capture_launch(self.graph)
262
+ else:
263
+ self.tape = wp.Tape()
264
+ with self.tape:
265
+ self.forward()
266
+ self.tape.backward(loss=self.loss)
267
+
268
+ if self.verbose:
269
+ self.log_step()
270
+
271
+ self.optimizer.step([self.material_params.grad])
272
+
273
+ wp.launch(
274
+ kernel=enforce_constraint_kernel,
275
+ dim=self.material_params.shape[0],
276
+ inputs=(
277
+ self.hard_lower_bound,
278
+ self.hard_upper_bound,
279
+ ),
280
+ outputs=(self.material_params,),
281
+ )
282
+
283
+ self.losses.append(self.loss.numpy()[0])
284
+
285
+ # clear grads for next iteration
286
+ self.tape.zero()
287
+ self.loss.zero_()
288
+ self.com.zero_()
289
+ self.pos_error.zero_()
290
+
291
+ self.iter = self.iter + 1
292
+
293
+ def log_step(self):
294
+ x = self.material_params.numpy().reshape(-1, 2)
295
+ x_grad = self.material_params.grad.numpy().reshape(-1, 2)
296
+
297
+ print(f"Iter: {self.iter} Loss: {self.loss.numpy()[0]}")
298
+
299
+ print(f"Pos error: {np.sqrt(self.pos_error.numpy()[0])}")
300
+
301
+ print(
302
+ f"Max Mu: {np.max(x[:, 0])}, Min Mu: {np.min(x[:, 0])}, "
303
+ f"Max Lambda: {np.max(x[:, 1])}, Min Lambda: {np.min(x[:, 1])}"
304
+ )
305
+
306
+ print(
307
+ f"Max Mu Grad: {np.max(x_grad[:, 0])}, Min Mu Grad: {np.min(x_grad[:, 0])}, "
308
+ f"Max Lambda Grad: {np.max(x_grad[:, 1])}, Min Lambda Grad: {np.min(x_grad[:, 1])}"
309
+ )
310
+
311
+ def render(self):
312
+ if self.renderer is None:
313
+ return
314
+
315
+ with wp.ScopedTimer("render"):
316
+ # draw trajectory
317
+ traj_verts = [np.mean(self.states[0].particle_q.numpy(), axis=0).tolist()]
318
+ for i in range(0, self.sim_steps, self.sim_substeps):
319
+ traj_verts.append(np.mean(self.states[i].particle_q.numpy(), axis=0).tolist())
320
+
321
+ self.renderer.begin_frame(self.render_time)
322
+ self.renderer.render(self.states[i])
323
+ self.renderer.render_box(
324
+ pos=self.target,
325
+ rot=wp.quat_identity(),
326
+ extents=(0.1, 0.1, 0.1),
327
+ name="target",
328
+ color=(0.0, 0.0, 0.0),
329
+ )
330
+ self.renderer.render_line_strip(
331
+ vertices=traj_verts,
332
+ color=wp.render.bourke_color_map(0.0, self.losses[0], self.losses[-1]),
333
+ radius=0.02,
334
+ name=f"traj_{self.iter - 1}",
335
+ )
336
+ self.renderer.end_frame()
337
+
338
+ from pxr import Gf, UsdGeom
339
+
340
+ particles_prim = self.renderer.stage.GetPrimAtPath("/root/particles")
341
+ particles = UsdGeom.Points.Get(self.renderer.stage, particles_prim.GetPath())
342
+ particles.CreateDisplayColorAttr().Set([Gf.Vec3f(1.0, 1.0, 1.0)], time=self.renderer.time)
343
+
344
+ self.render_time += self.frame_dt
345
+
346
+
347
+ if __name__ == "__main__":
348
+ import argparse
349
+
350
+ parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
351
+ parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")
352
+ parser.add_argument(
353
+ "--stage_path",
354
+ type=lambda x: None if x == "None" else str(x),
355
+ default="example_softbody_properties.usd",
356
+ help="Path to the output USD file.",
357
+ )
358
+ parser.add_argument(
359
+ "--train_iters",
360
+ type=int,
361
+ default=300,
362
+ help="Total number of training iterations.",
363
+ )
364
+ parser.add_argument(
365
+ "--material_behavior",
366
+ default="anisotropic",
367
+ choices=["anisotropic", "isotropic"],
368
+ help="Set material behavior to be Anisotropic or Isotropic.",
369
+ )
370
+ parser.add_argument(
371
+ "--verbose",
372
+ action="store_true",
373
+ help="Print out additional status messages during execution.",
374
+ )
375
+
376
+ args = parser.parse_known_args()[0]
377
+
378
+ with wp.ScopedDevice(args.device):
379
+ example = Example(
380
+ stage_path=args.stage_path,
381
+ material_behavior=args.material_behavior,
382
+ verbose=args.verbose,
383
+ )
384
+
385
+ # replay and optimize
386
+ for i in range(args.train_iters):
387
+ example.step()
388
+ if i == 0 or i % 50 == 0 or i == args.train_iters - 1:
389
+ example.render()
390
+
391
+ if example.renderer:
392
+ example.renderer.save()
@@ -60,8 +60,6 @@ class Example:
60
60
 
61
61
  self.iter = 0
62
62
 
63
- builder = wp.sim.ModelBuilder()
64
-
65
63
  # add planar joints
66
64
  builder = wp.sim.ModelBuilder(gravity=0.0)
67
65
  builder.add_articulation()
@@ -153,7 +151,7 @@ class Example:
153
151
  tape.backward(loss=self.loss)
154
152
 
155
153
  if self.verbose and (self.iter + 1) % 10 == 0:
156
- print(f"Iter {self.iter+1} Loss: {self.loss.numpy()[0]:.3f}")
154
+ print(f"Iter {self.iter + 1} Loss: {self.loss.numpy()[0]:.3f}")
157
155
 
158
156
  assert not np.isnan(self.actions.grad.numpy()).any(), "NaN in gradient"
159
157
 
@@ -15,6 +15,11 @@
15
15
  # simulated forward in time and then evaluated based on the center of mass
16
16
  # momentum of the mesh.
17
17
  #
18
+ # This example uses the deprecated wp.matmul() for matrix multiplication,
19
+ # which will be removed in a future version. See the updated version of
20
+ # this example, example_tile_walker.py, in examples/tile for the new
21
+ # approach to GEMMs using Warp's tile API.
22
+ #
18
23
  ###########################################################################
19
24
 
20
25
  import math
@@ -27,8 +27,6 @@ import warp.sim.render
27
27
 
28
28
  class Example:
29
29
  def __init__(self, stage_path="example_cartpole.usd", num_envs=8):
30
- builder = wp.sim.ModelBuilder()
31
-
32
30
  self.num_envs = num_envs
33
31
 
34
32
  articulation_builder = wp.sim.ModelBuilder()
@@ -100,7 +100,6 @@ class Example:
100
100
  tri_ka=1e4,
101
101
  tri_kd=1e-5,
102
102
  edge_ke=100,
103
- color_particles=True,
104
103
  )
105
104
 
106
105
  usd_stage = Usd.Stage.Open(os.path.join(warp.examples.get_asset_directory(), "bunny.usd"))
@@ -122,6 +121,9 @@ class Example:
122
121
  kf=1.0e1,
123
122
  )
124
123
 
124
+ if self.integrator_type == IntegratorType.VBD:
125
+ builder.color()
126
+
125
127
  self.model = builder.finalize()
126
128
  self.model.ground = True
127
129
  self.model.soft_contact_ke = 1.0e4
@@ -0,0 +1,260 @@
1
+ # Copyright (c) 2024 NVIDIA CORPORATION. All rights reserved.
2
+ # NVIDIA CORPORATION and its licensors retain all intellectual property
3
+ # and proprietary rights in and to this software, related documentation
4
+ # and any modifications thereto. Any use, reproduction, disclosure or
5
+ # distribution of this software and related documentation without an express
6
+ # license agreement from NVIDIA CORPORATION is strictly prohibited.
7
+
8
+ ###########################################################################
9
+ # Example Sim Cloth Self Contact
10
+ #
11
+ # This simulation demonstrates twisting an FEM cloth model using the VBD
12
+ # integrator, showcasing its ability to handle complex self-contacts while
13
+ # ensuring it remains intersection-free.
14
+ #
15
+ ###########################################################################
16
+
17
+ import math
18
+ import os
19
+
20
+ import numpy as np
21
+ from pxr import Usd, UsdGeom
22
+
23
+ import warp as wp
24
+ import warp.examples
25
+ import warp.sim
26
+ import warp.sim.render
27
+ from warp.sim.model import PARTICLE_FLAG_ACTIVE
28
+
29
+
30
+ @wp.kernel
31
+ def initialize_rotation(
32
+ particle_indices_to_rot: wp.array(dtype=wp.int32),
33
+ pos: wp.array(dtype=wp.vec3),
34
+ rot_centers: wp.array(dtype=wp.vec3),
35
+ rot_axes: wp.array(dtype=wp.vec3),
36
+ roots: wp.array(dtype=wp.vec3),
37
+ roots_to_ps: wp.array(dtype=wp.vec3),
38
+ ):
39
+ tid = wp.tid()
40
+ particle_index = particle_indices_to_rot[wp.tid()]
41
+
42
+ p = pos[particle_index]
43
+ rot_center = rot_centers[tid]
44
+ rot_axis = rot_axes[tid]
45
+ op = p - rot_center
46
+
47
+ root = wp.dot(op, rot_axis) * rot_axis
48
+
49
+ root_to_p = p - root
50
+
51
+ roots[tid] = root
52
+ roots_to_ps[tid] = root_to_p
53
+
54
+
55
+ @wp.kernel
56
+ def apply_rotation(
57
+ time: float,
58
+ angular_velocity: float,
59
+ particle_indices_to_rot: wp.array(dtype=wp.int32),
60
+ rot_centers: wp.array(dtype=wp.vec3),
61
+ rot_axes: wp.array(dtype=wp.vec3),
62
+ roots: wp.array(dtype=wp.vec3),
63
+ roots_to_ps: wp.array(dtype=wp.vec3),
64
+ pos_0: wp.array(dtype=wp.vec3),
65
+ pos_1: wp.array(dtype=wp.vec3),
66
+ ):
67
+ tid = wp.tid()
68
+ particle_index = particle_indices_to_rot[wp.tid()]
69
+
70
+ rot_axis = rot_axes[tid]
71
+
72
+ ux = rot_axis[0]
73
+ uy = rot_axis[1]
74
+ uz = rot_axis[2]
75
+
76
+ theta = time * angular_velocity
77
+
78
+ R = wp.mat33(
79
+ wp.cos(theta) + ux * ux * (1.0 - wp.cos(theta)),
80
+ ux * uy * (1.0 - wp.cos(theta)) - uz * wp.sin(theta),
81
+ ux * uz * (1.0 - wp.cos(theta)) + uy * wp.sin(theta),
82
+ uy * ux * (1.0 - wp.cos(theta)) + uz * wp.sin(theta),
83
+ wp.cos(theta) + uy * uy * (1.0 - wp.cos(theta)),
84
+ uy * uz * (1.0 - wp.cos(theta)) - ux * wp.sin(theta),
85
+ uz * ux * (1.0 - wp.cos(theta)) - uy * wp.sin(theta),
86
+ uz * uy * (1.0 - wp.cos(theta)) + ux * wp.sin(theta),
87
+ wp.cos(theta) + uz * uz * (1.0 - wp.cos(theta)),
88
+ )
89
+
90
+ root = roots[tid]
91
+ root_to_p = roots_to_ps[tid]
92
+ root_to_p_rot = R * root_to_p
93
+ p_rot = root + root_to_p_rot
94
+
95
+ pos_0[particle_index] = p_rot
96
+ pos_1[particle_index] = p_rot
97
+
98
+
99
+ class Example:
100
+ def __init__(self, stage_path="example_cloth_self_contact.usd", num_frames=1500):
101
+ fps = 60
102
+ self.frame_dt = 1.0 / fps
103
+ self.num_substeps = 10
104
+ self.iterations = 10
105
+ self.dt = self.frame_dt / self.num_substeps
106
+
107
+ self.num_frames = num_frames
108
+ self.sim_time = 0.0
109
+ self.profiler = {}
110
+
111
+ self.rot_angular_velocity = math.pi / 6
112
+ self.rot_end_time = 21
113
+
114
+ usd_stage = Usd.Stage.Open(os.path.join(warp.examples.get_asset_directory(), "square_cloth.usd"))
115
+ usd_geom = UsdGeom.Mesh(usd_stage.GetPrimAtPath("/root/cloth/cloth"))
116
+
117
+ mesh_points = np.array(usd_geom.GetPointsAttr().Get())
118
+ mesh_indices = np.array(usd_geom.GetFaceVertexIndicesAttr().Get())
119
+
120
+ self.input_scale_factor = 1.0
121
+ self.renderer_scale_factor = 0.01
122
+
123
+ vertices = [wp.vec3(v) * self.input_scale_factor for v in mesh_points]
124
+ self.faces = mesh_indices.reshape(-1, 3)
125
+
126
+ builder = wp.sim.ModelBuilder()
127
+ builder.add_cloth_mesh(
128
+ pos=wp.vec3(0.0, 0.0, 0.0),
129
+ rot=wp.quat_identity(),
130
+ scale=1.0,
131
+ vertices=vertices,
132
+ indices=mesh_indices,
133
+ vel=wp.vec3(0.0, 0.0, 0.0),
134
+ density=0.02,
135
+ tri_ke=1.0e5,
136
+ tri_ka=1.0e5,
137
+ tri_kd=3.0e-5,
138
+ )
139
+ builder.color()
140
+ self.model = builder.finalize()
141
+ self.model.ground = False
142
+ self.model.soft_contact_ke = 1.0e5
143
+ self.model.soft_contact_kd = 1.0e-6
144
+ self.model.soft_contact_mu = 0.2
145
+
146
+ # set up contact query and contact detection distances
147
+ self.model.soft_contact_radius = 0.2
148
+ self.model.soft_contact_margin = 0.35
149
+
150
+ cloth_size = 50
151
+ left_side = [cloth_size - 1 + i * cloth_size for i in range(cloth_size)]
152
+ right_side = [i * cloth_size for i in range(cloth_size)]
153
+ rot_point_indices = left_side + right_side
154
+
155
+ if len(rot_point_indices):
156
+ flags = self.model.particle_flags.numpy()
157
+ for fixed_vertex_id in rot_point_indices:
158
+ flags[fixed_vertex_id] = wp.uint32(int(flags[fixed_vertex_id]) & ~int(PARTICLE_FLAG_ACTIVE))
159
+
160
+ self.model.particle_flags = wp.array(flags)
161
+
162
+ self.integrator = wp.sim.VBDIntegrator(
163
+ self.model,
164
+ self.iterations,
165
+ handle_self_contact=True,
166
+ )
167
+ self.state0 = self.model.state()
168
+ self.state1 = self.model.state()
169
+
170
+ rot_axes = [[1, 0, 0]] * len(right_side) + [[-1, 0, 0]] * len(left_side)
171
+
172
+ self.rot_point_indices = wp.array(rot_point_indices, dtype=int)
173
+ self.rot_centers = wp.zeros(len(rot_point_indices), dtype=wp.vec3)
174
+ self.rot_axes = wp.array(rot_axes, dtype=wp.vec3)
175
+
176
+ self.roots = wp.zeros_like(self.rot_centers)
177
+ self.roots_to_ps = wp.zeros_like(self.rot_centers)
178
+
179
+ wp.launch(
180
+ kernel=initialize_rotation,
181
+ dim=self.rot_point_indices.shape[0],
182
+ inputs=[
183
+ self.rot_point_indices,
184
+ self.state0.particle_q,
185
+ self.rot_centers,
186
+ self.rot_axes,
187
+ self.roots,
188
+ self.roots_to_ps,
189
+ ],
190
+ )
191
+
192
+ if stage_path:
193
+ self.renderer = wp.sim.render.SimRenderer(self.model, stage_path, scaling=40.0)
194
+ else:
195
+ self.renderer = None
196
+
197
+ def step(self):
198
+ with wp.ScopedTimer("step", print=False, dict=self.profiler):
199
+ for _ in range(self.num_substeps):
200
+ if self.sim_time < self.rot_end_time:
201
+ wp.launch(
202
+ kernel=apply_rotation,
203
+ dim=self.rot_point_indices.shape[0],
204
+ inputs=[
205
+ self.sim_time,
206
+ self.rot_angular_velocity,
207
+ self.rot_point_indices,
208
+ self.rot_centers,
209
+ self.rot_axes,
210
+ self.roots,
211
+ self.roots_to_ps,
212
+ self.state0.particle_q,
213
+ self.state1.particle_q,
214
+ ],
215
+ )
216
+
217
+ self.integrator.simulate(self.model, self.state0, self.state1, self.dt)
218
+
219
+ (self.state0, self.state1) = (self.state1, self.state0)
220
+
221
+ self.sim_time += self.dt
222
+
223
+ def render(self):
224
+ if self.renderer is None:
225
+ return
226
+
227
+ with wp.ScopedTimer("render", print=False):
228
+ self.renderer.begin_frame(self.sim_time)
229
+ self.renderer.render(self.state0)
230
+ self.renderer.end_frame()
231
+
232
+
233
+ if __name__ == "__main__":
234
+ import argparse
235
+
236
+ parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
237
+ parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")
238
+ parser.add_argument(
239
+ "--stage_path",
240
+ type=lambda x: None if x == "None" else str(x),
241
+ default="example_cloth_self_contact.usd",
242
+ help="Path to the output USD file.",
243
+ )
244
+ parser.add_argument("--num_frames", type=int, default=1500, help="Total number of frames.")
245
+
246
+ args = parser.parse_known_args()[0]
247
+
248
+ with wp.ScopedDevice(args.device):
249
+ example = Example(stage_path=args.stage_path, num_frames=args.num_frames)
250
+
251
+ for i in range(example.num_frames):
252
+ example.step()
253
+ example.render()
254
+ print(f"[{i:4d}/{example.num_frames}]")
255
+
256
+ frame_times = example.profiler["step"]
257
+ print("\nAverage frame sim time: {:.2f} ms".format(sum(frame_times) / len(frame_times)))
258
+
259
+ if example.renderer:
260
+ example.renderer.save()