warp-lang 1.1.0__py3-none-win_amd64.whl → 1.2.1__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 (218) hide show
  1. warp/bin/warp-clang.dll +0 -0
  2. warp/bin/warp.dll +0 -0
  3. warp/build.py +10 -37
  4. warp/build_dll.py +2 -2
  5. warp/builtins.py +274 -6
  6. warp/codegen.py +51 -4
  7. warp/config.py +2 -2
  8. warp/constants.py +4 -0
  9. warp/context.py +422 -203
  10. warp/examples/benchmarks/benchmark_api.py +0 -2
  11. warp/examples/benchmarks/benchmark_cloth_warp.py +0 -1
  12. warp/examples/benchmarks/benchmark_launches.py +0 -2
  13. warp/examples/core/example_dem.py +0 -2
  14. warp/examples/core/example_fluid.py +0 -2
  15. warp/examples/core/example_graph_capture.py +0 -2
  16. warp/examples/core/example_marching_cubes.py +0 -2
  17. warp/examples/core/example_mesh.py +0 -2
  18. warp/examples/core/example_mesh_intersect.py +0 -2
  19. warp/examples/core/example_nvdb.py +0 -2
  20. warp/examples/core/example_raycast.py +0 -2
  21. warp/examples/core/example_raymarch.py +0 -2
  22. warp/examples/core/example_render_opengl.py +0 -2
  23. warp/examples/core/example_sph.py +0 -2
  24. warp/examples/core/example_torch.py +0 -3
  25. warp/examples/core/example_wave.py +0 -2
  26. warp/examples/fem/example_apic_fluid.py +140 -115
  27. warp/examples/fem/example_burgers.py +262 -0
  28. warp/examples/fem/example_convection_diffusion.py +0 -2
  29. warp/examples/fem/example_convection_diffusion_dg.py +0 -2
  30. warp/examples/fem/example_deformed_geometry.py +0 -2
  31. warp/examples/fem/example_diffusion.py +0 -2
  32. warp/examples/fem/example_diffusion_3d.py +5 -4
  33. warp/examples/fem/example_diffusion_mgpu.py +0 -2
  34. warp/examples/fem/example_mixed_elasticity.py +0 -2
  35. warp/examples/fem/example_navier_stokes.py +0 -2
  36. warp/examples/fem/example_stokes.py +0 -2
  37. warp/examples/fem/example_stokes_transfer.py +0 -2
  38. warp/examples/optim/example_bounce.py +0 -2
  39. warp/examples/optim/example_cloth_throw.py +0 -2
  40. warp/examples/optim/example_diffray.py +0 -2
  41. warp/examples/optim/example_drone.py +0 -2
  42. warp/examples/optim/example_inverse_kinematics.py +0 -2
  43. warp/examples/optim/example_inverse_kinematics_torch.py +0 -2
  44. warp/examples/optim/example_spring_cage.py +0 -2
  45. warp/examples/optim/example_trajectory.py +0 -2
  46. warp/examples/optim/example_walker.py +0 -2
  47. warp/examples/sim/example_cartpole.py +0 -2
  48. warp/examples/sim/example_cloth.py +0 -2
  49. warp/examples/sim/example_granular.py +0 -2
  50. warp/examples/sim/example_granular_collision_sdf.py +0 -2
  51. warp/examples/sim/example_jacobian_ik.py +0 -2
  52. warp/examples/sim/example_particle_chain.py +0 -2
  53. warp/examples/sim/example_quadruped.py +0 -2
  54. warp/examples/sim/example_rigid_chain.py +0 -2
  55. warp/examples/sim/example_rigid_contact.py +0 -2
  56. warp/examples/sim/example_rigid_force.py +0 -2
  57. warp/examples/sim/example_rigid_gyroscopic.py +0 -2
  58. warp/examples/sim/example_rigid_soft_contact.py +0 -2
  59. warp/examples/sim/example_soft_body.py +0 -2
  60. warp/fem/__init__.py +1 -0
  61. warp/fem/cache.py +3 -1
  62. warp/fem/geometry/__init__.py +1 -0
  63. warp/fem/geometry/element.py +4 -0
  64. warp/fem/geometry/grid_3d.py +0 -4
  65. warp/fem/geometry/nanogrid.py +455 -0
  66. warp/fem/integrate.py +63 -9
  67. warp/fem/space/__init__.py +43 -158
  68. warp/fem/space/basis_space.py +34 -0
  69. warp/fem/space/collocated_function_space.py +1 -1
  70. warp/fem/space/grid_2d_function_space.py +13 -132
  71. warp/fem/space/grid_3d_function_space.py +16 -154
  72. warp/fem/space/hexmesh_function_space.py +37 -134
  73. warp/fem/space/nanogrid_function_space.py +202 -0
  74. warp/fem/space/quadmesh_2d_function_space.py +12 -119
  75. warp/fem/space/restriction.py +4 -1
  76. warp/fem/space/shape/__init__.py +77 -0
  77. warp/fem/space/shape/cube_shape_function.py +5 -15
  78. warp/fem/space/tetmesh_function_space.py +6 -76
  79. warp/fem/space/trimesh_2d_function_space.py +6 -76
  80. warp/native/array.h +12 -3
  81. warp/native/builtin.h +48 -5
  82. warp/native/bvh.cpp +14 -10
  83. warp/native/bvh.cu +23 -15
  84. warp/native/bvh.h +1 -0
  85. warp/native/clang/clang.cpp +2 -1
  86. warp/native/crt.cpp +11 -1
  87. warp/native/crt.h +18 -1
  88. warp/native/exports.h +187 -0
  89. warp/native/mat.h +47 -0
  90. warp/native/mesh.cpp +1 -1
  91. warp/native/mesh.cu +1 -2
  92. warp/native/nanovdb/GridHandle.h +366 -0
  93. warp/native/nanovdb/HostBuffer.h +590 -0
  94. warp/native/nanovdb/NanoVDB.h +3999 -2157
  95. warp/native/nanovdb/PNanoVDB.h +936 -99
  96. warp/native/quat.h +28 -1
  97. warp/native/rand.h +5 -1
  98. warp/native/vec.h +45 -1
  99. warp/native/volume.cpp +335 -103
  100. warp/native/volume.cu +39 -13
  101. warp/native/volume.h +725 -303
  102. warp/native/volume_builder.cu +381 -360
  103. warp/native/volume_builder.h +16 -1
  104. warp/native/volume_impl.h +61 -0
  105. warp/native/warp.cu +8 -2
  106. warp/native/warp.h +15 -7
  107. warp/render/render_opengl.py +191 -52
  108. warp/sim/integrator_featherstone.py +10 -3
  109. warp/sim/integrator_xpbd.py +16 -22
  110. warp/sparse.py +89 -27
  111. warp/stubs.py +83 -0
  112. warp/tests/assets/test_index_grid.nvdb +0 -0
  113. warp/tests/aux_test_dependent.py +0 -2
  114. warp/tests/aux_test_grad_customs.py +0 -2
  115. warp/tests/aux_test_reference.py +0 -2
  116. warp/tests/aux_test_reference_reference.py +0 -2
  117. warp/tests/aux_test_square.py +0 -2
  118. warp/tests/disabled_kinematics.py +0 -2
  119. warp/tests/test_adam.py +0 -2
  120. warp/tests/test_arithmetic.py +0 -36
  121. warp/tests/test_array.py +9 -11
  122. warp/tests/test_array_reduce.py +0 -2
  123. warp/tests/test_async.py +0 -2
  124. warp/tests/test_atomic.py +0 -2
  125. warp/tests/test_bool.py +58 -50
  126. warp/tests/test_builtins_resolution.py +0 -2
  127. warp/tests/test_bvh.py +0 -2
  128. warp/tests/test_closest_point_edge_edge.py +0 -1
  129. warp/tests/test_codegen.py +0 -4
  130. warp/tests/test_compile_consts.py +130 -10
  131. warp/tests/test_conditional.py +0 -2
  132. warp/tests/test_copy.py +0 -2
  133. warp/tests/test_ctypes.py +6 -8
  134. warp/tests/test_dense.py +0 -2
  135. warp/tests/test_devices.py +0 -2
  136. warp/tests/test_dlpack.py +9 -11
  137. warp/tests/test_examples.py +42 -39
  138. warp/tests/test_fabricarray.py +0 -3
  139. warp/tests/test_fast_math.py +0 -2
  140. warp/tests/test_fem.py +75 -54
  141. warp/tests/test_fp16.py +0 -2
  142. warp/tests/test_func.py +0 -2
  143. warp/tests/test_generics.py +27 -2
  144. warp/tests/test_grad.py +147 -8
  145. warp/tests/test_grad_customs.py +0 -2
  146. warp/tests/test_hash_grid.py +1 -3
  147. warp/tests/test_import.py +0 -2
  148. warp/tests/test_indexedarray.py +0 -2
  149. warp/tests/test_intersect.py +0 -2
  150. warp/tests/test_jax.py +0 -2
  151. warp/tests/test_large.py +11 -9
  152. warp/tests/test_launch.py +0 -2
  153. warp/tests/test_lerp.py +10 -54
  154. warp/tests/test_linear_solvers.py +3 -5
  155. warp/tests/test_lvalue.py +0 -2
  156. warp/tests/test_marching_cubes.py +0 -2
  157. warp/tests/test_mat.py +0 -2
  158. warp/tests/test_mat_lite.py +0 -2
  159. warp/tests/test_mat_scalar_ops.py +0 -2
  160. warp/tests/test_math.py +0 -2
  161. warp/tests/test_matmul.py +35 -37
  162. warp/tests/test_matmul_lite.py +29 -31
  163. warp/tests/test_mempool.py +0 -2
  164. warp/tests/test_mesh.py +0 -3
  165. warp/tests/test_mesh_query_aabb.py +0 -2
  166. warp/tests/test_mesh_query_point.py +0 -2
  167. warp/tests/test_mesh_query_ray.py +0 -2
  168. warp/tests/test_mlp.py +0 -2
  169. warp/tests/test_model.py +0 -2
  170. warp/tests/test_module_hashing.py +111 -0
  171. warp/tests/test_modules_lite.py +0 -3
  172. warp/tests/test_multigpu.py +0 -2
  173. warp/tests/test_noise.py +0 -4
  174. warp/tests/test_operators.py +0 -2
  175. warp/tests/test_options.py +0 -2
  176. warp/tests/test_peer.py +0 -2
  177. warp/tests/test_pinned.py +0 -2
  178. warp/tests/test_print.py +0 -2
  179. warp/tests/test_quat.py +0 -2
  180. warp/tests/test_rand.py +41 -5
  181. warp/tests/test_reload.py +0 -10
  182. warp/tests/test_rounding.py +0 -2
  183. warp/tests/test_runlength_encode.py +0 -2
  184. warp/tests/test_sim_grad.py +0 -2
  185. warp/tests/test_sim_kinematics.py +0 -2
  186. warp/tests/test_smoothstep.py +0 -2
  187. warp/tests/test_snippet.py +0 -2
  188. warp/tests/test_sparse.py +0 -2
  189. warp/tests/test_spatial.py +0 -2
  190. warp/tests/test_special_values.py +362 -0
  191. warp/tests/test_streams.py +0 -2
  192. warp/tests/test_struct.py +0 -2
  193. warp/tests/test_tape.py +0 -2
  194. warp/tests/test_torch.py +0 -2
  195. warp/tests/test_transient_module.py +0 -2
  196. warp/tests/test_types.py +0 -2
  197. warp/tests/test_utils.py +0 -2
  198. warp/tests/test_vec.py +0 -2
  199. warp/tests/test_vec_lite.py +0 -2
  200. warp/tests/test_vec_scalar_ops.py +0 -2
  201. warp/tests/test_verify_fp.py +0 -2
  202. warp/tests/test_volume.py +237 -13
  203. warp/tests/test_volume_write.py +86 -3
  204. warp/tests/unittest_serial.py +10 -9
  205. warp/tests/unittest_suites.py +6 -2
  206. warp/tests/unittest_utils.py +2 -171
  207. warp/tests/unused_test_misc.py +0 -2
  208. warp/tests/walkthrough_debug.py +1 -1
  209. warp/thirdparty/unittest_parallel.py +37 -40
  210. warp/types.py +526 -85
  211. {warp_lang-1.1.0.dist-info → warp_lang-1.2.1.dist-info}/METADATA +61 -31
  212. warp_lang-1.2.1.dist-info/RECORD +359 -0
  213. warp/examples/fem/example_convection_diffusion_dg0.py +0 -204
  214. warp/native/nanovdb/PNanoVDBWrite.h +0 -295
  215. warp_lang-1.1.0.dist-info/RECORD +0 -352
  216. {warp_lang-1.1.0.dist-info → warp_lang-1.2.1.dist-info}/LICENSE.md +0 -0
  217. {warp_lang-1.1.0.dist-info → warp_lang-1.2.1.dist-info}/WHEEL +0 -0
  218. {warp_lang-1.1.0.dist-info → warp_lang-1.2.1.dist-info}/top_level.txt +0 -0
@@ -972,11 +972,9 @@ def apply_body_delta_velocities(
972
972
 
973
973
 
974
974
  @wp.kernel
975
- def apply_joint_torques(
975
+ def apply_joint_actions(
976
976
  body_q: wp.array(dtype=wp.transform),
977
977
  body_com: wp.array(dtype=wp.vec3),
978
- joint_q_start: wp.array(dtype=int),
979
- joint_qd_start: wp.array(dtype=int),
980
978
  joint_type: wp.array(dtype=int),
981
979
  joint_parent: wp.array(dtype=int),
982
980
  joint_child: wp.array(dtype=int),
@@ -1028,8 +1026,6 @@ def apply_joint_torques(
1028
1026
  # q_c = wp.transform_get_rotation(X_wc)
1029
1027
 
1030
1028
  # joint properties (for 1D joints)
1031
- # q_start = joint_q_start[tid]
1032
- qd_start = joint_qd_start[tid]
1033
1029
  axis_start = joint_axis_start[tid]
1034
1030
  lin_axis_count = joint_axis_dim[tid, 0]
1035
1031
  ang_axis_count = joint_axis_dim[tid, 1]
@@ -1043,14 +1039,14 @@ def apply_joint_torques(
1043
1039
  mode = joint_axis_mode[axis_start]
1044
1040
  if mode == wp.sim.JOINT_MODE_FORCE:
1045
1041
  axis = joint_axis[axis_start]
1046
- act = joint_act[qd_start]
1042
+ act = joint_act[axis_start]
1047
1043
  a_p = wp.transform_vector(X_wp, axis)
1048
1044
  t_total += act * a_p
1049
1045
  elif type == wp.sim.JOINT_PRISMATIC:
1050
1046
  mode = joint_axis_mode[axis_start]
1051
1047
  if mode == wp.sim.JOINT_MODE_FORCE:
1052
1048
  axis = joint_axis[axis_start]
1053
- act = joint_act[qd_start]
1049
+ act = joint_act[axis_start]
1054
1050
  a_p = wp.transform_vector(X_wp, axis)
1055
1051
  f_total += act * a_p
1056
1052
  elif type == wp.sim.JOINT_COMPOUND:
@@ -1075,13 +1071,13 @@ def apply_joint_torques(
1075
1071
 
1076
1072
  if joint_axis_mode[axis_start + 0] == wp.sim.JOINT_MODE_FORCE:
1077
1073
  axis_0 = joint_axis[axis_start + 0]
1078
- t_total += joint_act[qd_start + 0] * wp.transform_vector(X_wp, axis_0)
1074
+ t_total += joint_act[axis_start + 0] * wp.transform_vector(X_wp, axis_0)
1079
1075
  if joint_axis_mode[axis_start + 1] == wp.sim.JOINT_MODE_FORCE:
1080
1076
  axis_1 = joint_axis[axis_start + 1]
1081
- t_total += joint_act[qd_start + 1] * wp.transform_vector(X_wp, axis_1)
1077
+ t_total += joint_act[axis_start + 1] * wp.transform_vector(X_wp, axis_1)
1082
1078
  if joint_axis_mode[axis_start + 2] == wp.sim.JOINT_MODE_FORCE:
1083
1079
  axis_2 = joint_axis[axis_start + 2]
1084
- t_total += joint_act[qd_start + 2] * wp.transform_vector(X_wp, axis_2)
1080
+ t_total += joint_act[axis_start + 2] * wp.transform_vector(X_wp, axis_2)
1085
1081
 
1086
1082
  elif type == wp.sim.JOINT_UNIVERSAL:
1087
1083
  # q_off = wp.transform_get_rotation(X_cj)
@@ -1107,10 +1103,10 @@ def apply_joint_torques(
1107
1103
 
1108
1104
  if joint_axis_mode[axis_start + 0] == wp.sim.JOINT_MODE_FORCE:
1109
1105
  axis_0 = joint_axis[axis_start + 0]
1110
- t_total += joint_act[qd_start + 0] * wp.transform_vector(X_wp, axis_0)
1106
+ t_total += joint_act[axis_start + 0] * wp.transform_vector(X_wp, axis_0)
1111
1107
  if joint_axis_mode[axis_start + 1] == wp.sim.JOINT_MODE_FORCE:
1112
1108
  axis_1 = joint_axis[axis_start + 1]
1113
- t_total += joint_act[qd_start + 1] * wp.transform_vector(X_wp, axis_1)
1109
+ t_total += joint_act[axis_start + 1] * wp.transform_vector(X_wp, axis_1)
1114
1110
 
1115
1111
  elif type == wp.sim.JOINT_D6:
1116
1112
  # unroll for loop to ensure joint actions remain differentiable
@@ -1119,43 +1115,43 @@ def apply_joint_torques(
1119
1115
  if lin_axis_count > 0:
1120
1116
  if joint_axis_mode[axis_start + 0] == wp.sim.JOINT_MODE_FORCE:
1121
1117
  axis = joint_axis[axis_start + 0]
1122
- act = joint_act[qd_start + 0]
1118
+ act = joint_act[axis_start + 0]
1123
1119
  a_p = wp.transform_vector(X_wp, axis)
1124
1120
  f_total += act * a_p
1125
1121
  if lin_axis_count > 1:
1126
1122
  if joint_axis_mode[axis_start + 1] == wp.sim.JOINT_MODE_FORCE:
1127
1123
  axis = joint_axis[axis_start + 1]
1128
- act = joint_act[qd_start + 1]
1124
+ act = joint_act[axis_start + 1]
1129
1125
  a_p = wp.transform_vector(X_wp, axis)
1130
1126
  f_total += act * a_p
1131
1127
  if lin_axis_count > 2:
1132
1128
  if joint_axis_mode[axis_start + 2] == wp.sim.JOINT_MODE_FORCE:
1133
1129
  axis = joint_axis[axis_start + 2]
1134
- act = joint_act[qd_start + 2]
1130
+ act = joint_act[axis_start + 2]
1135
1131
  a_p = wp.transform_vector(X_wp, axis)
1136
1132
  f_total += act * a_p
1137
1133
 
1138
1134
  if ang_axis_count > 0:
1139
1135
  if joint_axis_mode[axis_start + lin_axis_count + 0] == wp.sim.JOINT_MODE_FORCE:
1140
1136
  axis = joint_axis[axis_start + lin_axis_count + 0]
1141
- act = joint_act[qd_start + lin_axis_count + 0]
1137
+ act = joint_act[axis_start + lin_axis_count + 0]
1142
1138
  a_p = wp.transform_vector(X_wp, axis)
1143
1139
  t_total += act * a_p
1144
1140
  if ang_axis_count > 1:
1145
1141
  if joint_axis_mode[axis_start + lin_axis_count + 1] == wp.sim.JOINT_MODE_FORCE:
1146
1142
  axis = joint_axis[axis_start + lin_axis_count + 1]
1147
- act = joint_act[qd_start + lin_axis_count + 1]
1143
+ act = joint_act[axis_start + lin_axis_count + 1]
1148
1144
  a_p = wp.transform_vector(X_wp, axis)
1149
1145
  t_total += act * a_p
1150
1146
  if ang_axis_count > 2:
1151
1147
  if joint_axis_mode[axis_start + lin_axis_count + 2] == wp.sim.JOINT_MODE_FORCE:
1152
1148
  axis = joint_axis[axis_start + lin_axis_count + 2]
1153
- act = joint_act[qd_start + lin_axis_count + 2]
1149
+ act = joint_act[axis_start + lin_axis_count + 2]
1154
1150
  a_p = wp.transform_vector(X_wp, axis)
1155
1151
  t_total += act * a_p
1156
1152
 
1157
1153
  else:
1158
- print("joint type not handled in apply_joint_torques")
1154
+ print("joint type not handled in apply_joint_actions")
1159
1155
 
1160
1156
  # write forces
1161
1157
  if id_p >= 0:
@@ -2838,13 +2834,11 @@ class XPBDIntegrator(Integrator):
2838
2834
 
2839
2835
  if model.joint_count:
2840
2836
  wp.launch(
2841
- kernel=apply_joint_torques,
2837
+ kernel=apply_joint_actions,
2842
2838
  dim=model.joint_count,
2843
2839
  inputs=[
2844
2840
  state_in.body_q,
2845
2841
  model.body_com,
2846
- model.joint_q_start,
2847
- model.joint_qd_start,
2848
2842
  model.joint_type,
2849
2843
  model.joint_parent,
2850
2844
  model.joint_child,
warp/sparse.py CHANGED
@@ -128,8 +128,8 @@ def bsr_zeros(
128
128
 
129
129
  bsr = bsr_matrix_t(block_type)()
130
130
 
131
- bsr.nrow = rows_of_blocks
132
- bsr.ncol = cols_of_blocks
131
+ bsr.nrow = int(rows_of_blocks)
132
+ bsr.ncol = int(cols_of_blocks)
133
133
  bsr.nnz = 0
134
134
  bsr.columns = wp.empty(shape=(0,), dtype=int, device=device)
135
135
  bsr.values = wp.empty(shape=(0,), dtype=block_type, device=device)
@@ -152,7 +152,11 @@ def _bsr_ensure_fits(bsr: BsrMatrix, nrow: int = None, nnz: int = None):
152
152
  bsr.values = wp.empty(shape=(nnz,), dtype=bsr.values.dtype, device=bsr.values.device)
153
153
 
154
154
 
155
- def bsr_set_zero(bsr: BsrMatrix, rows_of_blocks: Optional[int] = None, cols_of_blocks: Optional[int] = None):
155
+ def bsr_set_zero(
156
+ bsr: BsrMatrix,
157
+ rows_of_blocks: Optional[int] = None,
158
+ cols_of_blocks: Optional[int] = None,
159
+ ):
156
160
  """
157
161
  Sets a BSR matrix to zero, possibly changing its size
158
162
 
@@ -163,9 +167,9 @@ def bsr_set_zero(bsr: BsrMatrix, rows_of_blocks: Optional[int] = None, cols_of_b
163
167
  """
164
168
 
165
169
  if rows_of_blocks is not None:
166
- bsr.nrow = rows_of_blocks
170
+ bsr.nrow = int(rows_of_blocks)
167
171
  if cols_of_blocks is not None:
168
- bsr.ncol = cols_of_blocks
172
+ bsr.ncol = int(cols_of_blocks)
169
173
  bsr.nnz = 0
170
174
  _bsr_ensure_fits(bsr)
171
175
  bsr.offsets.zero_()
@@ -254,7 +258,10 @@ def bsr_set_from_triplets(
254
258
  )
255
259
 
256
260
 
257
- def bsr_assign(dest: BsrMatrix[BlockType[Rows, Cols, Scalar]], src: BsrMatrix[BlockType[Rows, Cols, Any]]):
261
+ def bsr_assign(
262
+ dest: BsrMatrix[BlockType[Rows, Cols, Scalar]],
263
+ src: BsrMatrix[BlockType[Rows, Cols, Any]],
264
+ ):
258
265
  """Copies the content of the `src` matrix to `dest`, casting the block values if the two matrices use distinct scalar types."""
259
266
 
260
267
  if dest.values.device != src.values.device:
@@ -288,12 +295,20 @@ def bsr_copy(A: BsrMatrix, scalar_type: Optional[Scalar] = None):
288
295
  else:
289
296
  block_type = wp.types.matrix(shape=A.block_shape, dtype=scalar_type)
290
297
 
291
- copy = bsr_zeros(rows_of_blocks=A.nrow, cols_of_blocks=A.ncol, block_type=block_type, device=A.values.device)
298
+ copy = bsr_zeros(
299
+ rows_of_blocks=A.nrow,
300
+ cols_of_blocks=A.ncol,
301
+ block_type=block_type,
302
+ device=A.values.device,
303
+ )
292
304
  bsr_assign(dest=copy, src=A)
293
305
  return copy
294
306
 
295
307
 
296
- def bsr_set_transpose(dest: BsrMatrix[BlockType[Cols, Rows, Scalar]], src: BsrMatrix[BlockType[Rows, Cols, Scalar]]):
308
+ def bsr_set_transpose(
309
+ dest: BsrMatrix[BlockType[Cols, Rows, Scalar]],
310
+ src: BsrMatrix[BlockType[Rows, Cols, Scalar]],
311
+ ):
297
312
  """Assigns the transposed matrix `src` to matrix `dest`"""
298
313
 
299
314
  if dest.values.device != src.values.device:
@@ -356,7 +371,12 @@ def bsr_transposed(A: BsrMatrix):
356
371
  else:
357
372
  block_type = wp.types.matrix(shape=A.block_shape[::-1], dtype=A.scalar_type)
358
373
 
359
- transposed = bsr_zeros(rows_of_blocks=A.ncol, cols_of_blocks=A.nrow, block_type=block_type, device=A.values.device)
374
+ transposed = bsr_zeros(
375
+ rows_of_blocks=A.ncol,
376
+ cols_of_blocks=A.nrow,
377
+ block_type=block_type,
378
+ device=A.values.device,
379
+ )
360
380
  bsr_set_transpose(dest=transposed, src=A)
361
381
  return transposed
362
382
 
@@ -399,7 +419,10 @@ def bsr_get_diag(A: BsrMatrix[_BlockType], out: "Optional[Array[BlockType]]" = N
399
419
  raise ValueError(f"Output array must be of length at least {dim}")
400
420
 
401
421
  wp.launch(
402
- kernel=_bsr_get_diag_kernel, dim=dim, device=A.values.device, inputs=[A.offsets, A.columns, A.values, out]
422
+ kernel=_bsr_get_diag_kernel,
423
+ dim=dim,
424
+ device=A.values.device,
425
+ inputs=[A.offsets, A.columns, A.values, out],
403
426
  )
404
427
 
405
428
  return out
@@ -567,7 +590,9 @@ def bsr_set_identity(A: BsrMatrix, rows_of_blocks: Optional[int] = None):
567
590
 
568
591
 
569
592
  def bsr_identity(
570
- rows_of_blocks: int, block_type: BlockType[Rows, Rows, Scalar], device: wp.context.Devicelike = None
593
+ rows_of_blocks: int,
594
+ block_type: BlockType[Rows, Rows, Scalar],
595
+ device: wp.context.Devicelike = None,
571
596
  ) -> BsrMatrix[BlockType[Rows, Rows, Scalar]]:
572
597
  """Creates and returns a square identity matrix.
573
598
 
@@ -576,7 +601,12 @@ def bsr_identity(
576
601
  block_type: Block type for the newly created matrix -- must be square
577
602
  device: Device onto which to allocate the data arrays
578
603
  """
579
- A = bsr_zeros(rows_of_blocks=rows_of_blocks, cols_of_blocks=rows_of_blocks, block_type=block_type, device=device)
604
+ A = bsr_zeros(
605
+ rows_of_blocks=rows_of_blocks,
606
+ cols_of_blocks=rows_of_blocks,
607
+ block_type=block_type,
608
+ device=device,
609
+ )
580
610
  bsr_set_identity(A)
581
611
  return A
582
612
 
@@ -601,7 +631,12 @@ def bsr_scale(x: BsrMatrix, alpha: Scalar) -> BsrMatrix:
601
631
  if not isinstance(alpha, x.scalar_type):
602
632
  alpha = x.scalar_type(alpha)
603
633
 
604
- wp.launch(kernel=_bsr_scale_kernel, dim=x.nnz, device=x.values.device, inputs=[alpha, x.values])
634
+ wp.launch(
635
+ kernel=_bsr_scale_kernel,
636
+ dim=x.nnz,
637
+ device=x.values.device,
638
+ inputs=[alpha, x.values],
639
+ )
605
640
 
606
641
  return x
607
642
 
@@ -723,10 +758,20 @@ def bsr_axpy(
723
758
  work_arrays._allocate(device, y, sum_nnz)
724
759
 
725
760
  wp.copy(work_arrays._sum_cols, y.columns, 0, 0, y.nnz)
726
- wp.launch(kernel=_bsr_get_block_row, device=device, dim=y.nnz, inputs=[0, y.offsets, work_arrays._sum_rows])
761
+ wp.launch(
762
+ kernel=_bsr_get_block_row,
763
+ device=device,
764
+ dim=y.nnz,
765
+ inputs=[0, y.offsets, work_arrays._sum_rows],
766
+ )
727
767
 
728
768
  wp.copy(work_arrays._sum_cols, x.columns, y.nnz, 0, x.nnz)
729
- wp.launch(kernel=_bsr_get_block_row, device=device, dim=x.nnz, inputs=[y.nnz, x.offsets, work_arrays._sum_rows])
769
+ wp.launch(
770
+ kernel=_bsr_get_block_row,
771
+ device=device,
772
+ dim=x.nnz,
773
+ inputs=[y.nnz, x.offsets, work_arrays._sum_rows],
774
+ )
730
775
 
731
776
  # Save old y values before overwriting matrix
732
777
  wp.copy(dest=work_arrays._old_y_values, src=y.values, count=y.nnz)
@@ -858,22 +903,25 @@ def _bsr_mm_compute_values(
858
903
  mm_cols: wp.array(dtype=int),
859
904
  mm_values: wp.array(dtype=Any),
860
905
  ):
861
- row = wp.tid()
862
- mm_beg = mm_offsets[row]
863
- mm_end = mm_offsets[row + 1]
906
+ mm_block = wp.tid()
907
+
908
+ row = wp.lower_bound(mm_offsets, mm_block + 1) - 1
909
+ col = mm_cols[mm_block]
910
+
911
+ mm_val = mm_values.dtype(type(alpha)(0.0))
864
912
 
865
913
  x_beg = x_offsets[row]
866
914
  x_end = x_offsets[row + 1]
867
915
  for x_block in range(x_beg, x_end):
868
916
  x_col = x_columns[x_block]
869
- ax_val = alpha * x_values[x_block]
870
-
871
917
  y_beg = y_offsets[x_col]
872
918
  y_end = y_offsets[x_col + 1]
873
919
 
874
- for y_block in range(y_beg, y_end):
875
- mm_block = wp.lower_bound(mm_cols, mm_beg, mm_end, y_columns[y_block])
876
- mm_values[mm_block] = mm_values[mm_block] + ax_val * y_values[y_block]
920
+ y_block = wp.lower_bound(y_columns, y_beg, y_end, col)
921
+ if y_block < y_end and y_columns[y_block] == col:
922
+ mm_val += x_values[x_block] * y_values[y_block]
923
+
924
+ mm_values[mm_block] += alpha * mm_val
877
925
 
878
926
 
879
927
  class bsr_mm_work_arrays:
@@ -995,13 +1043,24 @@ def bsr_mm(
995
1043
  kernel=_bsr_mm_count_coeffs,
996
1044
  device=device,
997
1045
  dim=z.nrow,
998
- inputs=[copied_z_nnz, x.offsets, x.columns, y.offsets, work_arrays._mm_row_counts],
1046
+ inputs=[
1047
+ copied_z_nnz,
1048
+ x.offsets,
1049
+ x.columns,
1050
+ y.offsets,
1051
+ work_arrays._mm_row_counts,
1052
+ ],
999
1053
  )
1000
1054
  warp.utils.array_scan(work_arrays._mm_row_counts, work_arrays._mm_row_counts)
1001
1055
 
1002
1056
  # Get back total counts on host
1003
1057
  if device.is_cuda:
1004
- wp.copy(dest=work_arrays._pinned_count_buffer, src=work_arrays._mm_row_counts, src_offset=z.nrow, count=1)
1058
+ wp.copy(
1059
+ dest=work_arrays._pinned_count_buffer,
1060
+ src=work_arrays._mm_row_counts,
1061
+ src_offset=z.nrow,
1062
+ count=1,
1063
+ )
1005
1064
  wp.synchronize_stream(wp.get_stream(device))
1006
1065
  mm_nnz = int(work_arrays._pinned_count_buffer.numpy()[0])
1007
1066
  else:
@@ -1014,7 +1073,10 @@ def bsr_mm(
1014
1073
  # Copy z row and column indices
1015
1074
  wp.copy(dest=work_arrays._mm_cols, src=z.columns, count=copied_z_nnz)
1016
1075
  wp.launch(
1017
- kernel=_bsr_get_block_row, device=device, dim=copied_z_nnz, inputs=[0, z.offsets, work_arrays._mm_rows]
1076
+ kernel=_bsr_get_block_row,
1077
+ device=device,
1078
+ dim=copied_z_nnz,
1079
+ inputs=[0, z.offsets, work_arrays._mm_rows],
1018
1080
  )
1019
1081
  # Save current z values in temporary buffer
1020
1082
  wp.copy(src=z.values, dest=work_arrays._old_z_values, count=copied_z_nnz)
@@ -1097,7 +1159,7 @@ def bsr_mm(
1097
1159
  wp.launch(
1098
1160
  kernel=_bsr_mm_compute_values,
1099
1161
  device=device,
1100
- dim=z.nrow,
1162
+ dim=z.nnz,
1101
1163
  inputs=[
1102
1164
  alpha,
1103
1165
  work_arrays._old_z_offsets if x == z else x.offsets,
warp/stubs.py CHANGED
@@ -348,6 +348,78 @@ def frac(x: Float) -> Float:
348
348
  ...
349
349
 
350
350
 
351
+ @over
352
+ def isfinite(x: Scalar) -> bool:
353
+ """Return ``True`` if x is a finite number, otherwise return ``False``."""
354
+ ...
355
+
356
+
357
+ @over
358
+ def isfinite(x: Vector[Any, Scalar]) -> bool:
359
+ """Return ``True`` if all elements of the vector ``x`` are finite, otherwise return ``False``."""
360
+ ...
361
+
362
+
363
+ @over
364
+ def isfinite(x: Quaternion[Scalar]) -> bool:
365
+ """Return ``True`` if all elements of the quaternion ``x`` are finite, otherwise return ``False``."""
366
+ ...
367
+
368
+
369
+ @over
370
+ def isfinite(m: Matrix[Any, Any, Scalar]) -> bool:
371
+ """Return ``True`` if all elements of the matrix ``m`` are finite, otherwise return ``False``."""
372
+ ...
373
+
374
+
375
+ @over
376
+ def isnan(x: Scalar) -> bool:
377
+ """Return ``True`` if ``x`` is NaN, otherwise return ``False``."""
378
+ ...
379
+
380
+
381
+ @over
382
+ def isnan(x: Vector[Any, Scalar]) -> bool:
383
+ """Return ``True`` if any element of the vector ``x`` is NaN, otherwise return ``False``."""
384
+ ...
385
+
386
+
387
+ @over
388
+ def isnan(x: Quaternion[Scalar]) -> bool:
389
+ """Return ``True`` if any element of the quaternion ``x`` is NaN, otherwise return ``False``."""
390
+ ...
391
+
392
+
393
+ @over
394
+ def isnan(m: Matrix[Any, Any, Scalar]) -> bool:
395
+ """Return ``True`` if any element of the matrix ``m`` is NaN, otherwise return ``False``."""
396
+ ...
397
+
398
+
399
+ @over
400
+ def isinf(x: Scalar) -> bool:
401
+ """Return ``True`` if x is positive or negative infinity, otherwise return ``False``."""
402
+ ...
403
+
404
+
405
+ @over
406
+ def isinf(x: Vector[Any, Scalar]) -> bool:
407
+ """Return ``True`` if any element of the vector ``x`` is positive or negative infinity, otherwise return ``False``."""
408
+ ...
409
+
410
+
411
+ @over
412
+ def isinf(x: Quaternion[Scalar]) -> bool:
413
+ """Return ``True`` if any element of the quaternion ``x`` is positive or negative infinity, otherwise return ``False``."""
414
+ ...
415
+
416
+
417
+ @over
418
+ def isinf(m: Matrix[Any, Any, Scalar]) -> bool:
419
+ """Return ``True`` if any element of the matrix ``m`` is positive or negative infinity, otherwise return ``False``."""
420
+ ...
421
+
422
+
351
423
  @over
352
424
  def dot(x: Vector[Any, Scalar], y: Vector[Any, Scalar]) -> Scalar:
353
425
  """Compute the dot product between two vectors."""
@@ -1040,6 +1112,17 @@ def volume_store_i(id: uint64, i: int32, j: int32, k: int32, value: int32):
1040
1112
  ...
1041
1113
 
1042
1114
 
1115
+ @over
1116
+ def volume_lookup_index(id: uint64, i: int32, j: int32, k: int32) -> int32:
1117
+ """Returns the index associated to the voxel with coordinates ``i``, ``j``, ``k``.
1118
+
1119
+ If the voxel at this index does not exist, this function returns -1.
1120
+ This function is available for both index grids and classical volumes.
1121
+
1122
+ """
1123
+ ...
1124
+
1125
+
1043
1126
  @over
1044
1127
  def volume_index_to_world(id: uint64, uvw: vec3f) -> vec3f:
1045
1128
  """Transform a point ``uvw`` defined in volume index space to world space given the volume's intrinsic affine transformation."""
Binary file
@@ -10,8 +10,6 @@
10
10
  import warp as wp
11
11
  import warp.tests.aux_test_reference as ref
12
12
 
13
- wp.init()
14
-
15
13
 
16
14
  @wp.kernel
17
15
  def kern(expect: float):
@@ -9,8 +9,6 @@
9
9
 
10
10
  import warp as wp
11
11
 
12
- wp.init()
13
-
14
12
 
15
13
  @wp.func
16
14
  def aux_custom_fn(x: float, y: float):
@@ -3,8 +3,6 @@
3
3
  import warp as wp
4
4
  import warp.tests.aux_test_reference_reference as refref
5
5
 
6
- wp.init()
7
-
8
6
 
9
7
  @wp.func
10
8
  def magic():
@@ -2,8 +2,6 @@
2
2
 
3
3
  import warp as wp
4
4
 
5
- wp.init()
6
-
7
5
 
8
6
  @wp.func
9
7
  def more_magic():
@@ -1,7 +1,5 @@
1
1
  import warp as wp
2
2
 
3
- wp.init()
4
-
5
3
 
6
4
  @wp.func
7
5
  def multiply(x: float):
@@ -12,8 +12,6 @@ import warp as wp
12
12
  import warp.sim
13
13
  from warp.tests.unittest_utils import *
14
14
 
15
- wp.init()
16
-
17
15
 
18
16
  def build_ant(num_envs):
19
17
  builder = wp.sim.ModelBuilder()
warp/tests/test_adam.py CHANGED
@@ -14,8 +14,6 @@ import warp.optim
14
14
  import warp.sim
15
15
  from warp.tests.unittest_utils import *
16
16
 
17
- wp.init()
18
-
19
17
 
20
18
  @wp.kernel
21
19
  def objective(params: wp.array(dtype=float), score: wp.array(dtype=float)):
@@ -5,7 +5,6 @@
5
5
  # distribution of this software and related documentation without an express
6
6
  # license agreement from NVIDIA CORPORATION is strictly prohibited.
7
7
 
8
- import math
9
8
  import unittest
10
9
 
11
10
  import numpy as np
@@ -13,8 +12,6 @@ import numpy as np
13
12
  import warp as wp
14
13
  from warp.tests.unittest_utils import *
15
14
 
16
- wp.init()
17
-
18
15
  np_signed_int_types = [
19
16
  np.int8,
20
17
  np.int16,
@@ -838,36 +835,6 @@ def test_float_to_int(test, device, dtype, register_kernels=False):
838
835
  tape.zero()
839
836
 
840
837
 
841
- def test_infinity(test, device, dtype, register_kernels=False):
842
- wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
843
-
844
- def check_infinity(
845
- outputs: wp.array(dtype=wptype),
846
- ):
847
- outputs[0] = wptype(wp.inf)
848
- outputs[1] = wptype(-wp.inf)
849
- outputs[2] = wptype(2.0 * wp.inf)
850
- outputs[3] = wptype(-2.0 * wp.inf)
851
- outputs[4] = wptype(2.0 / 0.0)
852
- outputs[5] = wptype(-2.0 / 0.0)
853
-
854
- kernel = getkernel(check_infinity, suffix=dtype.__name__)
855
-
856
- if register_kernels:
857
- return
858
-
859
- outputs = wp.zeros(6, dtype=wptype, device=device)
860
-
861
- wp.launch(kernel, dim=1, inputs=[], outputs=[outputs], device=device)
862
-
863
- test.assertEqual(outputs.numpy()[0], math.inf)
864
- test.assertEqual(outputs.numpy()[1], -math.inf)
865
- test.assertEqual(outputs.numpy()[2], math.inf)
866
- test.assertEqual(outputs.numpy()[3], -math.inf)
867
- test.assertEqual(outputs.numpy()[4], math.inf)
868
- test.assertEqual(outputs.numpy()[5], -math.inf)
869
-
870
-
871
838
  def test_interp(test, device, dtype, register_kernels=False):
872
839
  rng = np.random.default_rng(123)
873
840
 
@@ -1102,9 +1069,6 @@ for dtype in np_float_types:
1102
1069
  add_function_test_register_kernel(
1103
1070
  TestArithmetic, f"test_float_to_int_{dtype.__name__}", test_float_to_int, devices=devices, dtype=dtype
1104
1071
  )
1105
- add_function_test_register_kernel(
1106
- TestArithmetic, f"test_infinity_{dtype.__name__}", test_infinity, devices=devices, dtype=dtype
1107
- )
1108
1072
 
1109
1073
  for dtype in np_scalar_types:
1110
1074
  add_function_test_register_kernel(
warp/tests/test_array.py CHANGED
@@ -12,8 +12,6 @@ import numpy as np
12
12
  import warp as wp
13
13
  from warp.tests.unittest_utils import *
14
14
 
15
- wp.init()
16
-
17
15
 
18
16
  @wp.kernel
19
17
  def kernel_1d(a: wp.array(dtype=int, ndim=1)):
@@ -422,10 +420,10 @@ def test_view(test, device):
422
420
  wp_arr_e = wp.array(np_arr_e, dtype=wp.vec4, device=device)
423
421
  wp_arr_f = wp.array(np_arr_e, dtype=wp.quat, device=device)
424
422
 
425
- assert np.array_equal(np_arr_a.view(dtype=np.float32), wp_arr_a.view(dtype=wp.float32).numpy())
426
- assert np.array_equal(np_arr_b.view(dtype=np.uint32), wp_arr_b.view(dtype=wp.uint32).numpy())
427
- assert np.array_equal(np_arr_c.view(dtype=np.float16), wp_arr_c.view(dtype=wp.float16).numpy())
428
- assert np.array_equal(np_arr_d.view(dtype=np.uint16), wp_arr_d.view(dtype=wp.uint16).numpy())
423
+ assert_np_equal(wp_arr_a.view(dtype=wp.float32).numpy(), np_arr_a.view(dtype=np.float32))
424
+ assert_np_equal(wp_arr_b.view(dtype=wp.uint32).numpy(), np_arr_b.view(dtype=np.uint32))
425
+ assert_np_equal(wp_arr_c.view(dtype=wp.float16).numpy(), np_arr_c.view(dtype=np.float16))
426
+ assert_np_equal(wp_arr_d.view(dtype=wp.uint16).numpy(), np_arr_d.view(dtype=np.uint16))
429
427
  assert_array_equal(wp_arr_e.view(dtype=wp.quat), wp_arr_f)
430
428
 
431
429
 
@@ -486,7 +484,7 @@ def test_transpose(test, device):
486
484
  check = wp.zeros(shape=(2, 3), dtype=int, device=device)
487
485
 
488
486
  wp.launch(compare_2darrays, dim=(2, 3), inputs=[arr_transpose, arr_compare, check], device=device)
489
- assert np.array_equal(check.numpy(), np.ones((2, 3), dtype=int))
487
+ assert_np_equal(check.numpy(), np.ones((2, 3), dtype=int))
490
488
 
491
489
  # test transpose in square 3d case
492
490
  # wp does not support copying from/to non-contiguous arrays so check in kernel
@@ -497,7 +495,7 @@ def test_transpose(test, device):
497
495
  check = wp.zeros(shape=(3, 2, 2), dtype=int, device=device)
498
496
 
499
497
  wp.launch(compare_3darrays, dim=(3, 2, 2), inputs=[arr_transpose, arr_compare, check], device=device)
500
- assert np.array_equal(check.numpy(), np.ones((3, 2, 2), dtype=int))
498
+ assert_np_equal(check.numpy(), np.ones((3, 2, 2), dtype=int))
501
499
 
502
500
  # test transpose in square 3d case without axes supplied
503
501
  arr_transpose = arr.transpose()
@@ -505,13 +503,13 @@ def test_transpose(test, device):
505
503
  check = wp.zeros(shape=(2, 2, 3), dtype=int, device=device)
506
504
 
507
505
  wp.launch(compare_3darrays, dim=(2, 2, 3), inputs=[arr_transpose, arr_compare, check], device=device)
508
- assert np.array_equal(check.numpy(), np.ones((2, 2, 3), dtype=int))
506
+ assert_np_equal(check.numpy(), np.ones((2, 2, 3), dtype=int))
509
507
 
510
508
  # test transpose in 1d case (should be noop)
511
509
  np_arr = np.array([1, 2, 3], dtype=float)
512
510
  arr = wp.array(np_arr, dtype=float, device=device)
513
511
 
514
- assert np.array_equal(np_arr.transpose(), arr.transpose().numpy())
512
+ assert_np_equal(arr.transpose().numpy(), np_arr.transpose())
515
513
 
516
514
 
517
515
  def test_fill_scalar(test, device):
@@ -2335,7 +2333,7 @@ devices = get_test_devices()
2335
2333
 
2336
2334
  class TestArray(unittest.TestCase):
2337
2335
  def test_array_new_del(self):
2338
- # test the scenario in which an array instance is created but not initalized before gc
2336
+ # test the scenario in which an array instance is created but not initialized before gc
2339
2337
  instance = wp.array.__new__(wp.array)
2340
2338
  instance.__del__()
2341
2339
 
@@ -13,8 +13,6 @@ import warp as wp
13
13
  from warp.tests.unittest_utils import *
14
14
  from warp.utils import array_inner, array_sum
15
15
 
16
- wp.init()
17
-
18
16
 
19
17
  def make_test_array_sum(dtype):
20
18
  N = 1000