warp-lang 1.7.0__py3-none-manylinux_2_34_aarch64.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 (429) hide show
  1. warp/__init__.py +139 -0
  2. warp/__init__.pyi +1 -0
  3. warp/autograd.py +1142 -0
  4. warp/bin/warp-clang.so +0 -0
  5. warp/bin/warp.so +0 -0
  6. warp/build.py +557 -0
  7. warp/build_dll.py +405 -0
  8. warp/builtins.py +6855 -0
  9. warp/codegen.py +3969 -0
  10. warp/config.py +158 -0
  11. warp/constants.py +57 -0
  12. warp/context.py +6812 -0
  13. warp/dlpack.py +462 -0
  14. warp/examples/__init__.py +24 -0
  15. warp/examples/assets/bear.usd +0 -0
  16. warp/examples/assets/bunny.usd +0 -0
  17. warp/examples/assets/cartpole.urdf +110 -0
  18. warp/examples/assets/crazyflie.usd +0 -0
  19. warp/examples/assets/cube.usd +0 -0
  20. warp/examples/assets/nonuniform.usd +0 -0
  21. warp/examples/assets/nv_ant.xml +92 -0
  22. warp/examples/assets/nv_humanoid.xml +183 -0
  23. warp/examples/assets/nvidia_logo.png +0 -0
  24. warp/examples/assets/pixel.jpg +0 -0
  25. warp/examples/assets/quadruped.urdf +268 -0
  26. warp/examples/assets/rocks.nvdb +0 -0
  27. warp/examples/assets/rocks.usd +0 -0
  28. warp/examples/assets/sphere.usd +0 -0
  29. warp/examples/assets/square_cloth.usd +0 -0
  30. warp/examples/benchmarks/benchmark_api.py +389 -0
  31. warp/examples/benchmarks/benchmark_cloth.py +296 -0
  32. warp/examples/benchmarks/benchmark_cloth_cupy.py +96 -0
  33. warp/examples/benchmarks/benchmark_cloth_jax.py +105 -0
  34. warp/examples/benchmarks/benchmark_cloth_numba.py +161 -0
  35. warp/examples/benchmarks/benchmark_cloth_numpy.py +85 -0
  36. warp/examples/benchmarks/benchmark_cloth_paddle.py +94 -0
  37. warp/examples/benchmarks/benchmark_cloth_pytorch.py +94 -0
  38. warp/examples/benchmarks/benchmark_cloth_taichi.py +120 -0
  39. warp/examples/benchmarks/benchmark_cloth_warp.py +153 -0
  40. warp/examples/benchmarks/benchmark_gemm.py +164 -0
  41. warp/examples/benchmarks/benchmark_interop_paddle.py +166 -0
  42. warp/examples/benchmarks/benchmark_interop_torch.py +166 -0
  43. warp/examples/benchmarks/benchmark_launches.py +301 -0
  44. warp/examples/benchmarks/benchmark_tile_load_store.py +103 -0
  45. warp/examples/browse.py +37 -0
  46. warp/examples/core/example_cupy.py +86 -0
  47. warp/examples/core/example_dem.py +241 -0
  48. warp/examples/core/example_fluid.py +299 -0
  49. warp/examples/core/example_graph_capture.py +150 -0
  50. warp/examples/core/example_marching_cubes.py +194 -0
  51. warp/examples/core/example_mesh.py +180 -0
  52. warp/examples/core/example_mesh_intersect.py +211 -0
  53. warp/examples/core/example_nvdb.py +182 -0
  54. warp/examples/core/example_raycast.py +111 -0
  55. warp/examples/core/example_raymarch.py +205 -0
  56. warp/examples/core/example_render_opengl.py +193 -0
  57. warp/examples/core/example_sample_mesh.py +300 -0
  58. warp/examples/core/example_sph.py +411 -0
  59. warp/examples/core/example_torch.py +211 -0
  60. warp/examples/core/example_wave.py +269 -0
  61. warp/examples/fem/example_adaptive_grid.py +286 -0
  62. warp/examples/fem/example_apic_fluid.py +423 -0
  63. warp/examples/fem/example_burgers.py +261 -0
  64. warp/examples/fem/example_convection_diffusion.py +178 -0
  65. warp/examples/fem/example_convection_diffusion_dg.py +204 -0
  66. warp/examples/fem/example_deformed_geometry.py +172 -0
  67. warp/examples/fem/example_diffusion.py +196 -0
  68. warp/examples/fem/example_diffusion_3d.py +225 -0
  69. warp/examples/fem/example_diffusion_mgpu.py +220 -0
  70. warp/examples/fem/example_distortion_energy.py +228 -0
  71. warp/examples/fem/example_magnetostatics.py +240 -0
  72. warp/examples/fem/example_mixed_elasticity.py +291 -0
  73. warp/examples/fem/example_navier_stokes.py +261 -0
  74. warp/examples/fem/example_nonconforming_contact.py +298 -0
  75. warp/examples/fem/example_stokes.py +213 -0
  76. warp/examples/fem/example_stokes_transfer.py +262 -0
  77. warp/examples/fem/example_streamlines.py +352 -0
  78. warp/examples/fem/utils.py +1000 -0
  79. warp/examples/interop/example_jax_callable.py +116 -0
  80. warp/examples/interop/example_jax_ffi_callback.py +132 -0
  81. warp/examples/interop/example_jax_kernel.py +205 -0
  82. warp/examples/optim/example_bounce.py +266 -0
  83. warp/examples/optim/example_cloth_throw.py +228 -0
  84. warp/examples/optim/example_diffray.py +561 -0
  85. warp/examples/optim/example_drone.py +870 -0
  86. warp/examples/optim/example_fluid_checkpoint.py +497 -0
  87. warp/examples/optim/example_inverse_kinematics.py +182 -0
  88. warp/examples/optim/example_inverse_kinematics_torch.py +191 -0
  89. warp/examples/optim/example_softbody_properties.py +400 -0
  90. warp/examples/optim/example_spring_cage.py +245 -0
  91. warp/examples/optim/example_trajectory.py +227 -0
  92. warp/examples/sim/example_cartpole.py +143 -0
  93. warp/examples/sim/example_cloth.py +225 -0
  94. warp/examples/sim/example_cloth_self_contact.py +322 -0
  95. warp/examples/sim/example_granular.py +130 -0
  96. warp/examples/sim/example_granular_collision_sdf.py +202 -0
  97. warp/examples/sim/example_jacobian_ik.py +244 -0
  98. warp/examples/sim/example_particle_chain.py +124 -0
  99. warp/examples/sim/example_quadruped.py +203 -0
  100. warp/examples/sim/example_rigid_chain.py +203 -0
  101. warp/examples/sim/example_rigid_contact.py +195 -0
  102. warp/examples/sim/example_rigid_force.py +133 -0
  103. warp/examples/sim/example_rigid_gyroscopic.py +115 -0
  104. warp/examples/sim/example_rigid_soft_contact.py +140 -0
  105. warp/examples/sim/example_soft_body.py +196 -0
  106. warp/examples/tile/example_tile_cholesky.py +87 -0
  107. warp/examples/tile/example_tile_convolution.py +66 -0
  108. warp/examples/tile/example_tile_fft.py +55 -0
  109. warp/examples/tile/example_tile_filtering.py +113 -0
  110. warp/examples/tile/example_tile_matmul.py +85 -0
  111. warp/examples/tile/example_tile_mlp.py +383 -0
  112. warp/examples/tile/example_tile_nbody.py +199 -0
  113. warp/examples/tile/example_tile_walker.py +327 -0
  114. warp/fabric.py +355 -0
  115. warp/fem/__init__.py +106 -0
  116. warp/fem/adaptivity.py +508 -0
  117. warp/fem/cache.py +572 -0
  118. warp/fem/dirichlet.py +202 -0
  119. warp/fem/domain.py +411 -0
  120. warp/fem/field/__init__.py +125 -0
  121. warp/fem/field/field.py +619 -0
  122. warp/fem/field/nodal_field.py +326 -0
  123. warp/fem/field/restriction.py +37 -0
  124. warp/fem/field/virtual.py +848 -0
  125. warp/fem/geometry/__init__.py +32 -0
  126. warp/fem/geometry/adaptive_nanogrid.py +857 -0
  127. warp/fem/geometry/closest_point.py +84 -0
  128. warp/fem/geometry/deformed_geometry.py +221 -0
  129. warp/fem/geometry/element.py +776 -0
  130. warp/fem/geometry/geometry.py +362 -0
  131. warp/fem/geometry/grid_2d.py +392 -0
  132. warp/fem/geometry/grid_3d.py +452 -0
  133. warp/fem/geometry/hexmesh.py +911 -0
  134. warp/fem/geometry/nanogrid.py +571 -0
  135. warp/fem/geometry/partition.py +389 -0
  136. warp/fem/geometry/quadmesh.py +663 -0
  137. warp/fem/geometry/tetmesh.py +855 -0
  138. warp/fem/geometry/trimesh.py +806 -0
  139. warp/fem/integrate.py +2335 -0
  140. warp/fem/linalg.py +419 -0
  141. warp/fem/operator.py +293 -0
  142. warp/fem/polynomial.py +229 -0
  143. warp/fem/quadrature/__init__.py +17 -0
  144. warp/fem/quadrature/pic_quadrature.py +299 -0
  145. warp/fem/quadrature/quadrature.py +591 -0
  146. warp/fem/space/__init__.py +228 -0
  147. warp/fem/space/basis_function_space.py +468 -0
  148. warp/fem/space/basis_space.py +667 -0
  149. warp/fem/space/dof_mapper.py +251 -0
  150. warp/fem/space/function_space.py +309 -0
  151. warp/fem/space/grid_2d_function_space.py +177 -0
  152. warp/fem/space/grid_3d_function_space.py +227 -0
  153. warp/fem/space/hexmesh_function_space.py +257 -0
  154. warp/fem/space/nanogrid_function_space.py +201 -0
  155. warp/fem/space/partition.py +367 -0
  156. warp/fem/space/quadmesh_function_space.py +223 -0
  157. warp/fem/space/restriction.py +179 -0
  158. warp/fem/space/shape/__init__.py +143 -0
  159. warp/fem/space/shape/cube_shape_function.py +1105 -0
  160. warp/fem/space/shape/shape_function.py +133 -0
  161. warp/fem/space/shape/square_shape_function.py +926 -0
  162. warp/fem/space/shape/tet_shape_function.py +834 -0
  163. warp/fem/space/shape/triangle_shape_function.py +672 -0
  164. warp/fem/space/tetmesh_function_space.py +271 -0
  165. warp/fem/space/topology.py +424 -0
  166. warp/fem/space/trimesh_function_space.py +194 -0
  167. warp/fem/types.py +99 -0
  168. warp/fem/utils.py +420 -0
  169. warp/jax.py +187 -0
  170. warp/jax_experimental/__init__.py +16 -0
  171. warp/jax_experimental/custom_call.py +351 -0
  172. warp/jax_experimental/ffi.py +698 -0
  173. warp/jax_experimental/xla_ffi.py +602 -0
  174. warp/math.py +244 -0
  175. warp/native/array.h +1145 -0
  176. warp/native/builtin.h +1800 -0
  177. warp/native/bvh.cpp +492 -0
  178. warp/native/bvh.cu +791 -0
  179. warp/native/bvh.h +554 -0
  180. warp/native/clang/clang.cpp +536 -0
  181. warp/native/coloring.cpp +613 -0
  182. warp/native/crt.cpp +51 -0
  183. warp/native/crt.h +362 -0
  184. warp/native/cuda_crt.h +1058 -0
  185. warp/native/cuda_util.cpp +646 -0
  186. warp/native/cuda_util.h +307 -0
  187. warp/native/error.cpp +77 -0
  188. warp/native/error.h +36 -0
  189. warp/native/exports.h +1878 -0
  190. warp/native/fabric.h +245 -0
  191. warp/native/hashgrid.cpp +311 -0
  192. warp/native/hashgrid.cu +87 -0
  193. warp/native/hashgrid.h +240 -0
  194. warp/native/initializer_array.h +41 -0
  195. warp/native/intersect.h +1230 -0
  196. warp/native/intersect_adj.h +375 -0
  197. warp/native/intersect_tri.h +339 -0
  198. warp/native/marching.cpp +19 -0
  199. warp/native/marching.cu +514 -0
  200. warp/native/marching.h +19 -0
  201. warp/native/mat.h +2220 -0
  202. warp/native/mathdx.cpp +87 -0
  203. warp/native/matnn.h +343 -0
  204. warp/native/mesh.cpp +266 -0
  205. warp/native/mesh.cu +404 -0
  206. warp/native/mesh.h +1980 -0
  207. warp/native/nanovdb/GridHandle.h +366 -0
  208. warp/native/nanovdb/HostBuffer.h +590 -0
  209. warp/native/nanovdb/NanoVDB.h +6624 -0
  210. warp/native/nanovdb/PNanoVDB.h +3390 -0
  211. warp/native/noise.h +859 -0
  212. warp/native/quat.h +1371 -0
  213. warp/native/rand.h +342 -0
  214. warp/native/range.h +139 -0
  215. warp/native/reduce.cpp +174 -0
  216. warp/native/reduce.cu +364 -0
  217. warp/native/runlength_encode.cpp +79 -0
  218. warp/native/runlength_encode.cu +61 -0
  219. warp/native/scan.cpp +47 -0
  220. warp/native/scan.cu +53 -0
  221. warp/native/scan.h +23 -0
  222. warp/native/solid_angle.h +466 -0
  223. warp/native/sort.cpp +251 -0
  224. warp/native/sort.cu +277 -0
  225. warp/native/sort.h +33 -0
  226. warp/native/sparse.cpp +378 -0
  227. warp/native/sparse.cu +524 -0
  228. warp/native/spatial.h +657 -0
  229. warp/native/svd.h +702 -0
  230. warp/native/temp_buffer.h +46 -0
  231. warp/native/tile.h +2584 -0
  232. warp/native/tile_reduce.h +264 -0
  233. warp/native/vec.h +1426 -0
  234. warp/native/volume.cpp +501 -0
  235. warp/native/volume.cu +67 -0
  236. warp/native/volume.h +969 -0
  237. warp/native/volume_builder.cu +477 -0
  238. warp/native/volume_builder.h +52 -0
  239. warp/native/volume_impl.h +70 -0
  240. warp/native/warp.cpp +1082 -0
  241. warp/native/warp.cu +3636 -0
  242. warp/native/warp.h +381 -0
  243. warp/optim/__init__.py +17 -0
  244. warp/optim/adam.py +163 -0
  245. warp/optim/linear.py +1137 -0
  246. warp/optim/sgd.py +112 -0
  247. warp/paddle.py +407 -0
  248. warp/render/__init__.py +18 -0
  249. warp/render/render_opengl.py +3518 -0
  250. warp/render/render_usd.py +784 -0
  251. warp/render/utils.py +160 -0
  252. warp/sim/__init__.py +65 -0
  253. warp/sim/articulation.py +793 -0
  254. warp/sim/collide.py +2395 -0
  255. warp/sim/graph_coloring.py +300 -0
  256. warp/sim/import_mjcf.py +790 -0
  257. warp/sim/import_snu.py +227 -0
  258. warp/sim/import_urdf.py +579 -0
  259. warp/sim/import_usd.py +894 -0
  260. warp/sim/inertia.py +324 -0
  261. warp/sim/integrator.py +242 -0
  262. warp/sim/integrator_euler.py +1997 -0
  263. warp/sim/integrator_featherstone.py +2101 -0
  264. warp/sim/integrator_vbd.py +2048 -0
  265. warp/sim/integrator_xpbd.py +3292 -0
  266. warp/sim/model.py +4791 -0
  267. warp/sim/particles.py +121 -0
  268. warp/sim/render.py +427 -0
  269. warp/sim/utils.py +428 -0
  270. warp/sparse.py +2057 -0
  271. warp/stubs.py +3333 -0
  272. warp/tape.py +1203 -0
  273. warp/tests/__init__.py +1 -0
  274. warp/tests/__main__.py +4 -0
  275. warp/tests/assets/curlnoise_golden.npy +0 -0
  276. warp/tests/assets/mlp_golden.npy +0 -0
  277. warp/tests/assets/pixel.npy +0 -0
  278. warp/tests/assets/pnoise_golden.npy +0 -0
  279. warp/tests/assets/spiky.usd +0 -0
  280. warp/tests/assets/test_grid.nvdb +0 -0
  281. warp/tests/assets/test_index_grid.nvdb +0 -0
  282. warp/tests/assets/test_int32_grid.nvdb +0 -0
  283. warp/tests/assets/test_vec_grid.nvdb +0 -0
  284. warp/tests/assets/torus.nvdb +0 -0
  285. warp/tests/assets/torus.usda +105 -0
  286. warp/tests/aux_test_class_kernel.py +34 -0
  287. warp/tests/aux_test_compile_consts_dummy.py +18 -0
  288. warp/tests/aux_test_conditional_unequal_types_kernels.py +29 -0
  289. warp/tests/aux_test_dependent.py +29 -0
  290. warp/tests/aux_test_grad_customs.py +29 -0
  291. warp/tests/aux_test_instancing_gc.py +26 -0
  292. warp/tests/aux_test_module_unload.py +23 -0
  293. warp/tests/aux_test_name_clash1.py +40 -0
  294. warp/tests/aux_test_name_clash2.py +40 -0
  295. warp/tests/aux_test_reference.py +9 -0
  296. warp/tests/aux_test_reference_reference.py +8 -0
  297. warp/tests/aux_test_square.py +16 -0
  298. warp/tests/aux_test_unresolved_func.py +22 -0
  299. warp/tests/aux_test_unresolved_symbol.py +22 -0
  300. warp/tests/cuda/__init__.py +0 -0
  301. warp/tests/cuda/test_async.py +676 -0
  302. warp/tests/cuda/test_ipc.py +124 -0
  303. warp/tests/cuda/test_mempool.py +233 -0
  304. warp/tests/cuda/test_multigpu.py +169 -0
  305. warp/tests/cuda/test_peer.py +139 -0
  306. warp/tests/cuda/test_pinned.py +84 -0
  307. warp/tests/cuda/test_streams.py +634 -0
  308. warp/tests/geometry/__init__.py +0 -0
  309. warp/tests/geometry/test_bvh.py +200 -0
  310. warp/tests/geometry/test_hash_grid.py +221 -0
  311. warp/tests/geometry/test_marching_cubes.py +74 -0
  312. warp/tests/geometry/test_mesh.py +316 -0
  313. warp/tests/geometry/test_mesh_query_aabb.py +399 -0
  314. warp/tests/geometry/test_mesh_query_point.py +932 -0
  315. warp/tests/geometry/test_mesh_query_ray.py +311 -0
  316. warp/tests/geometry/test_volume.py +1103 -0
  317. warp/tests/geometry/test_volume_write.py +346 -0
  318. warp/tests/interop/__init__.py +0 -0
  319. warp/tests/interop/test_dlpack.py +729 -0
  320. warp/tests/interop/test_jax.py +371 -0
  321. warp/tests/interop/test_paddle.py +800 -0
  322. warp/tests/interop/test_torch.py +1001 -0
  323. warp/tests/run_coverage_serial.py +39 -0
  324. warp/tests/sim/__init__.py +0 -0
  325. warp/tests/sim/disabled_kinematics.py +244 -0
  326. warp/tests/sim/flaky_test_sim_grad.py +290 -0
  327. warp/tests/sim/test_collision.py +604 -0
  328. warp/tests/sim/test_coloring.py +258 -0
  329. warp/tests/sim/test_model.py +224 -0
  330. warp/tests/sim/test_sim_grad_bounce_linear.py +212 -0
  331. warp/tests/sim/test_sim_kinematics.py +98 -0
  332. warp/tests/sim/test_vbd.py +597 -0
  333. warp/tests/test_adam.py +163 -0
  334. warp/tests/test_arithmetic.py +1096 -0
  335. warp/tests/test_array.py +2972 -0
  336. warp/tests/test_array_reduce.py +156 -0
  337. warp/tests/test_assert.py +250 -0
  338. warp/tests/test_atomic.py +153 -0
  339. warp/tests/test_bool.py +220 -0
  340. warp/tests/test_builtins_resolution.py +1298 -0
  341. warp/tests/test_closest_point_edge_edge.py +327 -0
  342. warp/tests/test_codegen.py +810 -0
  343. warp/tests/test_codegen_instancing.py +1495 -0
  344. warp/tests/test_compile_consts.py +215 -0
  345. warp/tests/test_conditional.py +252 -0
  346. warp/tests/test_context.py +42 -0
  347. warp/tests/test_copy.py +238 -0
  348. warp/tests/test_ctypes.py +638 -0
  349. warp/tests/test_dense.py +73 -0
  350. warp/tests/test_devices.py +97 -0
  351. warp/tests/test_examples.py +482 -0
  352. warp/tests/test_fabricarray.py +996 -0
  353. warp/tests/test_fast_math.py +74 -0
  354. warp/tests/test_fem.py +2003 -0
  355. warp/tests/test_fp16.py +136 -0
  356. warp/tests/test_func.py +454 -0
  357. warp/tests/test_future_annotations.py +98 -0
  358. warp/tests/test_generics.py +656 -0
  359. warp/tests/test_grad.py +893 -0
  360. warp/tests/test_grad_customs.py +339 -0
  361. warp/tests/test_grad_debug.py +341 -0
  362. warp/tests/test_implicit_init.py +411 -0
  363. warp/tests/test_import.py +45 -0
  364. warp/tests/test_indexedarray.py +1140 -0
  365. warp/tests/test_intersect.py +73 -0
  366. warp/tests/test_iter.py +76 -0
  367. warp/tests/test_large.py +177 -0
  368. warp/tests/test_launch.py +411 -0
  369. warp/tests/test_lerp.py +151 -0
  370. warp/tests/test_linear_solvers.py +193 -0
  371. warp/tests/test_lvalue.py +427 -0
  372. warp/tests/test_mat.py +2089 -0
  373. warp/tests/test_mat_lite.py +122 -0
  374. warp/tests/test_mat_scalar_ops.py +2913 -0
  375. warp/tests/test_math.py +178 -0
  376. warp/tests/test_mlp.py +282 -0
  377. warp/tests/test_module_hashing.py +258 -0
  378. warp/tests/test_modules_lite.py +44 -0
  379. warp/tests/test_noise.py +252 -0
  380. warp/tests/test_operators.py +299 -0
  381. warp/tests/test_options.py +129 -0
  382. warp/tests/test_overwrite.py +551 -0
  383. warp/tests/test_print.py +339 -0
  384. warp/tests/test_quat.py +2315 -0
  385. warp/tests/test_rand.py +339 -0
  386. warp/tests/test_reload.py +302 -0
  387. warp/tests/test_rounding.py +185 -0
  388. warp/tests/test_runlength_encode.py +196 -0
  389. warp/tests/test_scalar_ops.py +105 -0
  390. warp/tests/test_smoothstep.py +108 -0
  391. warp/tests/test_snippet.py +318 -0
  392. warp/tests/test_sparse.py +582 -0
  393. warp/tests/test_spatial.py +2229 -0
  394. warp/tests/test_special_values.py +361 -0
  395. warp/tests/test_static.py +592 -0
  396. warp/tests/test_struct.py +734 -0
  397. warp/tests/test_tape.py +204 -0
  398. warp/tests/test_transient_module.py +93 -0
  399. warp/tests/test_triangle_closest_point.py +145 -0
  400. warp/tests/test_types.py +562 -0
  401. warp/tests/test_utils.py +588 -0
  402. warp/tests/test_vec.py +1487 -0
  403. warp/tests/test_vec_lite.py +80 -0
  404. warp/tests/test_vec_scalar_ops.py +2327 -0
  405. warp/tests/test_verify_fp.py +100 -0
  406. warp/tests/tile/__init__.py +0 -0
  407. warp/tests/tile/test_tile.py +780 -0
  408. warp/tests/tile/test_tile_load.py +407 -0
  409. warp/tests/tile/test_tile_mathdx.py +208 -0
  410. warp/tests/tile/test_tile_mlp.py +402 -0
  411. warp/tests/tile/test_tile_reduce.py +447 -0
  412. warp/tests/tile/test_tile_shared_memory.py +247 -0
  413. warp/tests/tile/test_tile_view.py +173 -0
  414. warp/tests/unittest_serial.py +47 -0
  415. warp/tests/unittest_suites.py +427 -0
  416. warp/tests/unittest_utils.py +468 -0
  417. warp/tests/walkthrough_debug.py +93 -0
  418. warp/thirdparty/__init__.py +0 -0
  419. warp/thirdparty/appdirs.py +598 -0
  420. warp/thirdparty/dlpack.py +145 -0
  421. warp/thirdparty/unittest_parallel.py +570 -0
  422. warp/torch.py +391 -0
  423. warp/types.py +5230 -0
  424. warp/utils.py +1137 -0
  425. warp_lang-1.7.0.dist-info/METADATA +516 -0
  426. warp_lang-1.7.0.dist-info/RECORD +429 -0
  427. warp_lang-1.7.0.dist-info/WHEEL +5 -0
  428. warp_lang-1.7.0.dist-info/licenses/LICENSE.md +202 -0
  429. warp_lang-1.7.0.dist-info/top_level.txt +1 -0
@@ -0,0 +1,3518 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
+ # SPDX-License-Identifier: Apache-2.0
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+
16
+ import ctypes
17
+ import sys
18
+ import time
19
+ from collections import defaultdict
20
+ from typing import List, Optional, Tuple, Union
21
+
22
+ import numpy as np
23
+
24
+ import warp as wp
25
+
26
+ from .utils import tab10_color_map
27
+
28
+ Mat44 = Union[List[float], List[List[float]], np.ndarray]
29
+
30
+
31
+ wp.set_module_options({"enable_backward": False})
32
+
33
+ shape_vertex_shader = """
34
+ #version 330 core
35
+ layout (location = 0) in vec3 aPos;
36
+ layout (location = 1) in vec3 aNormal;
37
+ layout (location = 2) in vec2 aTexCoord;
38
+
39
+ // column vectors of the instance transform matrix
40
+ layout (location = 3) in vec4 aInstanceTransform0;
41
+ layout (location = 4) in vec4 aInstanceTransform1;
42
+ layout (location = 5) in vec4 aInstanceTransform2;
43
+ layout (location = 6) in vec4 aInstanceTransform3;
44
+
45
+ // colors to use for the checkerboard pattern
46
+ layout (location = 7) in vec3 aObjectColor1;
47
+ layout (location = 8) in vec3 aObjectColor2;
48
+
49
+ uniform mat4 view;
50
+ uniform mat4 model;
51
+ uniform mat4 projection;
52
+
53
+ out vec3 Normal;
54
+ out vec3 FragPos;
55
+ out vec2 TexCoord;
56
+ out vec3 ObjectColor1;
57
+ out vec3 ObjectColor2;
58
+
59
+ void main()
60
+ {
61
+ mat4 transform = model * mat4(aInstanceTransform0, aInstanceTransform1, aInstanceTransform2, aInstanceTransform3);
62
+ vec4 worldPos = transform * vec4(aPos, 1.0);
63
+ gl_Position = projection * view * worldPos;
64
+ FragPos = vec3(worldPos);
65
+ Normal = mat3(transpose(inverse(transform))) * aNormal;
66
+ TexCoord = aTexCoord;
67
+ ObjectColor1 = aObjectColor1;
68
+ ObjectColor2 = aObjectColor2;
69
+ }
70
+ """
71
+
72
+ shape_fragment_shader = """
73
+ #version 330 core
74
+ out vec4 FragColor;
75
+
76
+ in vec3 Normal;
77
+ in vec3 FragPos;
78
+ in vec2 TexCoord;
79
+ in vec3 ObjectColor1;
80
+ in vec3 ObjectColor2;
81
+
82
+ uniform vec3 viewPos;
83
+ uniform vec3 lightColor;
84
+ uniform vec3 sunDirection;
85
+
86
+ void main()
87
+ {
88
+ float ambientStrength = 0.3;
89
+ vec3 ambient = ambientStrength * lightColor;
90
+ vec3 norm = normalize(Normal);
91
+
92
+ float diff = max(dot(norm, sunDirection), 0.0);
93
+ vec3 diffuse = diff * lightColor;
94
+
95
+ vec3 lightDir2 = normalize(vec3(1.0, 0.3, -0.3));
96
+ diff = max(dot(norm, lightDir2), 0.0);
97
+ diffuse += diff * lightColor * 0.3;
98
+
99
+ float specularStrength = 0.5;
100
+ vec3 viewDir = normalize(viewPos - FragPos);
101
+
102
+ vec3 reflectDir = reflect(-sunDirection, norm);
103
+ float spec = pow(max(dot(viewDir, reflectDir), 0.0), 32);
104
+ vec3 specular = specularStrength * spec * lightColor;
105
+
106
+ reflectDir = reflect(-lightDir2, norm);
107
+ spec = pow(max(dot(viewDir, reflectDir), 0.0), 64);
108
+ specular += specularStrength * spec * lightColor * 0.3;
109
+
110
+ // checkerboard pattern
111
+ float u = TexCoord.x;
112
+ float v = TexCoord.y;
113
+ // blend the checkerboard pattern dependent on the gradient of the texture coordinates
114
+ // to void Moire patterns
115
+ vec2 grad = abs(dFdx(TexCoord)) + abs(dFdy(TexCoord));
116
+ float blendRange = 1.5;
117
+ float blendFactor = max(grad.x, grad.y) * blendRange;
118
+ float scale = 2.0;
119
+ float checker = mod(floor(u * scale) + floor(v * scale), 2.0);
120
+ checker = mix(checker, 0.5, smoothstep(0.0, 1.0, blendFactor));
121
+ vec3 checkerColor = mix(ObjectColor1, ObjectColor2, checker);
122
+
123
+ vec3 result = (ambient + diffuse + specular) * checkerColor;
124
+ FragColor = vec4(result, 1.0);
125
+ }
126
+ """
127
+
128
+ grid_vertex_shader = """
129
+ #version 330 core
130
+
131
+ uniform mat4 view;
132
+ uniform mat4 model;
133
+ uniform mat4 projection;
134
+
135
+ in vec3 position;
136
+
137
+ void main() {
138
+ gl_Position = projection * view * model * vec4(position, 1.0);
139
+ }
140
+ """
141
+
142
+ # Fragment shader source code
143
+ grid_fragment_shader = """
144
+ #version 330 core
145
+
146
+ out vec4 outColor;
147
+
148
+ void main() {
149
+ outColor = vec4(0.5, 0.5, 0.5, 1.0);
150
+ }
151
+ """
152
+
153
+ sky_vertex_shader = """
154
+ #version 330 core
155
+
156
+ layout (location = 0) in vec3 aPos;
157
+ layout (location = 1) in vec3 aNormal;
158
+ layout (location = 2) in vec2 aTexCoord;
159
+
160
+ uniform mat4 view;
161
+ uniform mat4 inv_model;
162
+ uniform mat4 projection;
163
+ uniform vec3 viewPos;
164
+
165
+ out vec3 FragPos;
166
+ out vec2 TexCoord;
167
+
168
+ void main()
169
+ {
170
+ vec4 worldPos = vec4(aPos + viewPos, 1.0);
171
+ gl_Position = projection * view * inv_model * worldPos;
172
+
173
+ FragPos = vec3(worldPos);
174
+ TexCoord = aTexCoord;
175
+ }
176
+ """
177
+
178
+ sky_fragment_shader = """
179
+ #version 330 core
180
+
181
+ out vec4 FragColor;
182
+
183
+ in vec3 FragPos;
184
+ in vec2 TexCoord;
185
+
186
+ uniform vec3 color1;
187
+ uniform vec3 color2;
188
+ uniform float farPlane;
189
+
190
+ uniform vec3 sunDirection;
191
+
192
+ void main()
193
+ {
194
+ float y = tanh(FragPos.y/farPlane*10.0)*0.5+0.5;
195
+ float height = sqrt(1.0-y);
196
+
197
+ float s = pow(0.5, 1.0 / 10.0);
198
+ s = 1.0 - clamp(s, 0.75, 1.0);
199
+
200
+ vec3 haze = mix(vec3(1.0), color2 * 1.3, s);
201
+ vec3 sky = mix(color1, haze, height / 1.3);
202
+
203
+ float diff = max(dot(sunDirection, normalize(FragPos)), 0.0);
204
+ vec3 sun = pow(diff, 32) * vec3(1.0, 0.8, 0.6) * 0.5;
205
+
206
+ FragColor = vec4(sky + sun, 1.0);
207
+ }
208
+ """
209
+
210
+ frame_vertex_shader = """
211
+ #version 330 core
212
+ layout (location = 0) in vec3 aPos;
213
+ layout (location = 1) in vec2 aTexCoord;
214
+
215
+ out vec2 TexCoord;
216
+
217
+ void main() {
218
+ gl_Position = vec4(aPos, 1.0);
219
+ TexCoord = aTexCoord;
220
+ }
221
+ """
222
+
223
+ frame_fragment_shader = """
224
+ #version 330 core
225
+ in vec2 TexCoord;
226
+
227
+ out vec4 FragColor;
228
+
229
+ uniform sampler2D textureSampler;
230
+
231
+ void main() {
232
+ FragColor = texture(textureSampler, TexCoord);
233
+ }
234
+ """
235
+
236
+ frame_depth_fragment_shader = """
237
+ #version 330 core
238
+ in vec2 TexCoord;
239
+
240
+ out vec4 FragColor;
241
+
242
+ uniform sampler2D textureSampler;
243
+
244
+ vec3 bourkeColorMap(float v) {
245
+ vec3 c = vec3(1.0, 1.0, 1.0);
246
+
247
+ v = clamp(v, 0.0, 1.0); // Ensures v is between 0 and 1
248
+
249
+ if (v < 0.25) {
250
+ c.r = 0.0;
251
+ c.g = 4.0 * v;
252
+ } else if (v < 0.5) {
253
+ c.r = 0.0;
254
+ c.b = 1.0 + 4.0 * (0.25 - v);
255
+ } else if (v < 0.75) {
256
+ c.r = 4.0 * (v - 0.5);
257
+ c.b = 0.0;
258
+ } else {
259
+ c.g = 1.0 + 4.0 * (0.75 - v);
260
+ c.b = 0.0;
261
+ }
262
+
263
+ return c;
264
+ }
265
+
266
+ void main() {
267
+ float depth = texture(textureSampler, TexCoord).r;
268
+ FragColor = vec4(bourkeColorMap(sqrt(1.0 - depth)), 1.0);
269
+ }
270
+ """
271
+
272
+
273
+ @wp.kernel
274
+ def update_vbo_transforms(
275
+ instance_id: wp.array(dtype=int),
276
+ instance_body: wp.array(dtype=int),
277
+ instance_transforms: wp.array(dtype=wp.transform),
278
+ instance_scalings: wp.array(dtype=wp.vec3),
279
+ body_q: wp.array(dtype=wp.transform),
280
+ # outputs
281
+ vbo_transforms: wp.array(dtype=wp.mat44),
282
+ ):
283
+ tid = wp.tid()
284
+ i = instance_id[tid]
285
+ X_ws = instance_transforms[i]
286
+ if instance_body:
287
+ body = instance_body[i]
288
+ if body >= 0:
289
+ if body_q:
290
+ X_ws = body_q[body] * X_ws
291
+ else:
292
+ return
293
+ p = wp.transform_get_translation(X_ws)
294
+ q = wp.transform_get_rotation(X_ws)
295
+ s = instance_scalings[i]
296
+ rot = wp.quat_to_matrix(q)
297
+ # transposed definition
298
+ vbo_transforms[tid] = wp.mat44(
299
+ rot[0, 0] * s[0],
300
+ rot[1, 0] * s[0],
301
+ rot[2, 0] * s[0],
302
+ 0.0,
303
+ rot[0, 1] * s[1],
304
+ rot[1, 1] * s[1],
305
+ rot[2, 1] * s[1],
306
+ 0.0,
307
+ rot[0, 2] * s[2],
308
+ rot[1, 2] * s[2],
309
+ rot[2, 2] * s[2],
310
+ 0.0,
311
+ p[0],
312
+ p[1],
313
+ p[2],
314
+ 1.0,
315
+ )
316
+
317
+
318
+ @wp.kernel
319
+ def update_vbo_vertices(
320
+ points: wp.array(dtype=wp.vec3),
321
+ scale: wp.vec3,
322
+ # outputs
323
+ vbo_vertices: wp.array(dtype=float, ndim=2),
324
+ ):
325
+ tid = wp.tid()
326
+ p = points[tid]
327
+ vbo_vertices[tid, 0] = p[0] * scale[0]
328
+ vbo_vertices[tid, 1] = p[1] * scale[1]
329
+ vbo_vertices[tid, 2] = p[2] * scale[2]
330
+
331
+
332
+ @wp.kernel
333
+ def update_points_positions(
334
+ instance_positions: wp.array(dtype=wp.vec3),
335
+ instance_scalings: wp.array(dtype=wp.vec3),
336
+ # outputs
337
+ vbo_transforms: wp.array(dtype=wp.mat44),
338
+ ):
339
+ tid = wp.tid()
340
+ p = instance_positions[tid]
341
+ s = wp.vec3(1.0)
342
+ if instance_scalings:
343
+ s = instance_scalings[tid]
344
+ # transposed definition
345
+ # fmt: off
346
+ vbo_transforms[tid] = wp.mat44(
347
+ s[0], 0.0, 0.0, 0.0,
348
+ 0.0, s[1], 0.0, 0.0,
349
+ 0.0, 0.0, s[2], 0.0,
350
+ p[0], p[1], p[2], 1.0)
351
+ # fmt: on
352
+
353
+
354
+ @wp.kernel
355
+ def update_line_transforms(
356
+ lines: wp.array(dtype=wp.vec3, ndim=2),
357
+ # outputs
358
+ vbo_transforms: wp.array(dtype=wp.mat44),
359
+ ):
360
+ tid = wp.tid()
361
+ p0 = lines[tid, 0]
362
+ p1 = lines[tid, 1]
363
+ p = 0.5 * (p0 + p1)
364
+ d = p1 - p0
365
+ s = wp.length(d)
366
+ axis = wp.normalize(d)
367
+ y_up = wp.vec3(0.0, 1.0, 0.0)
368
+ angle = wp.acos(wp.dot(axis, y_up))
369
+ axis = wp.normalize(wp.cross(axis, y_up))
370
+ q = wp.quat_from_axis_angle(axis, -angle)
371
+ rot = wp.quat_to_matrix(q)
372
+ # transposed definition
373
+ # fmt: off
374
+ vbo_transforms[tid] = wp.mat44(
375
+ rot[0, 0], rot[1, 0], rot[2, 0], 0.0,
376
+ s * rot[0, 1], s * rot[1, 1], s * rot[2, 1], 0.0,
377
+ rot[0, 2], rot[1, 2], rot[2, 2], 0.0,
378
+ p[0], p[1], p[2], 1.0,
379
+ )
380
+ # fmt: on
381
+
382
+
383
+ @wp.kernel
384
+ def compute_gfx_vertices(
385
+ indices: wp.array(dtype=int, ndim=2),
386
+ vertices: wp.array(dtype=wp.vec3, ndim=1),
387
+ scale: wp.vec3,
388
+ # outputs
389
+ gfx_vertices: wp.array(dtype=float, ndim=2),
390
+ ):
391
+ tid = wp.tid()
392
+ v0 = vertices[indices[tid, 0]] * scale[0]
393
+ v1 = vertices[indices[tid, 1]] * scale[1]
394
+ v2 = vertices[indices[tid, 2]] * scale[2]
395
+ i = tid * 3
396
+ j = i + 1
397
+ k = i + 2
398
+ gfx_vertices[i, 0] = v0[0]
399
+ gfx_vertices[i, 1] = v0[1]
400
+ gfx_vertices[i, 2] = v0[2]
401
+ gfx_vertices[j, 0] = v1[0]
402
+ gfx_vertices[j, 1] = v1[1]
403
+ gfx_vertices[j, 2] = v1[2]
404
+ gfx_vertices[k, 0] = v2[0]
405
+ gfx_vertices[k, 1] = v2[1]
406
+ gfx_vertices[k, 2] = v2[2]
407
+ n = wp.normalize(wp.cross(v1 - v0, v2 - v0))
408
+ gfx_vertices[i, 3] = n[0]
409
+ gfx_vertices[i, 4] = n[1]
410
+ gfx_vertices[i, 5] = n[2]
411
+ gfx_vertices[j, 3] = n[0]
412
+ gfx_vertices[j, 4] = n[1]
413
+ gfx_vertices[j, 5] = n[2]
414
+ gfx_vertices[k, 3] = n[0]
415
+ gfx_vertices[k, 4] = n[1]
416
+ gfx_vertices[k, 5] = n[2]
417
+
418
+
419
+ @wp.kernel
420
+ def compute_average_normals(
421
+ indices: wp.array(dtype=int, ndim=2),
422
+ vertices: wp.array(dtype=wp.vec3),
423
+ scale: wp.vec3,
424
+ # outputs
425
+ normals: wp.array(dtype=wp.vec3),
426
+ faces_per_vertex: wp.array(dtype=int),
427
+ ):
428
+ tid = wp.tid()
429
+ i = indices[tid, 0]
430
+ j = indices[tid, 1]
431
+ k = indices[tid, 2]
432
+ v0 = vertices[i] * scale[0]
433
+ v1 = vertices[j] * scale[1]
434
+ v2 = vertices[k] * scale[2]
435
+ n = wp.normalize(wp.cross(v1 - v0, v2 - v0))
436
+ wp.atomic_add(normals, i, n)
437
+ wp.atomic_add(faces_per_vertex, i, 1)
438
+ wp.atomic_add(normals, j, n)
439
+ wp.atomic_add(faces_per_vertex, j, 1)
440
+ wp.atomic_add(normals, k, n)
441
+ wp.atomic_add(faces_per_vertex, k, 1)
442
+
443
+
444
+ @wp.kernel
445
+ def assemble_gfx_vertices(
446
+ vertices: wp.array(dtype=wp.vec3, ndim=1),
447
+ normals: wp.array(dtype=wp.vec3),
448
+ faces_per_vertex: wp.array(dtype=int),
449
+ scale: wp.vec3,
450
+ # outputs
451
+ gfx_vertices: wp.array(dtype=float, ndim=2),
452
+ ):
453
+ tid = wp.tid()
454
+ v = vertices[tid]
455
+ n = normals[tid] / float(faces_per_vertex[tid])
456
+ gfx_vertices[tid, 0] = v[0] * scale[0]
457
+ gfx_vertices[tid, 1] = v[1] * scale[1]
458
+ gfx_vertices[tid, 2] = v[2] * scale[2]
459
+ gfx_vertices[tid, 3] = n[0]
460
+ gfx_vertices[tid, 4] = n[1]
461
+ gfx_vertices[tid, 5] = n[2]
462
+
463
+
464
+ @wp.kernel
465
+ def copy_rgb_frame(
466
+ input_img: wp.array(dtype=wp.uint8),
467
+ width: int,
468
+ height: int,
469
+ # outputs
470
+ output_img: wp.array(dtype=float, ndim=3),
471
+ ):
472
+ w, v = wp.tid()
473
+ pixel = v * width + w
474
+ pixel *= 3
475
+ r = float(input_img[pixel + 0])
476
+ g = float(input_img[pixel + 1])
477
+ b = float(input_img[pixel + 2])
478
+ # flip vertically (OpenGL coordinates start at bottom)
479
+ v = height - v - 1
480
+ output_img[v, w, 0] = r / 255.0
481
+ output_img[v, w, 1] = g / 255.0
482
+ output_img[v, w, 2] = b / 255.0
483
+
484
+
485
+ @wp.kernel
486
+ def copy_rgb_frame_uint8(
487
+ input_img: wp.array(dtype=wp.uint8),
488
+ width: int,
489
+ height: int,
490
+ # outputs
491
+ output_img: wp.array(dtype=wp.uint8, ndim=3),
492
+ ):
493
+ w, v = wp.tid()
494
+ pixel = v * width + w
495
+ pixel *= 3
496
+ # flip vertically (OpenGL coordinates start at bottom)
497
+ v = height - v - 1
498
+ output_img[v, w, 0] = input_img[pixel + 0]
499
+ output_img[v, w, 1] = input_img[pixel + 1]
500
+ output_img[v, w, 2] = input_img[pixel + 2]
501
+
502
+
503
+ @wp.kernel
504
+ def copy_depth_frame(
505
+ input_img: wp.array(dtype=wp.float32),
506
+ width: int,
507
+ height: int,
508
+ near: float,
509
+ far: float,
510
+ # outputs
511
+ output_img: wp.array(dtype=wp.float32, ndim=3),
512
+ ):
513
+ w, v = wp.tid()
514
+ pixel = v * width + w
515
+ # flip vertically (OpenGL coordinates start at bottom)
516
+ v = height - v - 1
517
+ d = 2.0 * input_img[pixel] - 1.0
518
+ d = 2.0 * near * far / ((far - near) * d - near - far)
519
+ output_img[v, w, 0] = -d
520
+
521
+
522
+ @wp.kernel
523
+ def copy_rgb_frame_tiles(
524
+ input_img: wp.array(dtype=wp.uint8),
525
+ positions: wp.array(dtype=int, ndim=2),
526
+ screen_width: int,
527
+ screen_height: int,
528
+ tile_height: int,
529
+ # outputs
530
+ output_img: wp.array(dtype=float, ndim=4),
531
+ ):
532
+ tile, x, y = wp.tid()
533
+ p = positions[tile]
534
+ qx = x + p[0]
535
+ qy = y + p[1]
536
+ pixel = qy * screen_width + qx
537
+ # flip vertically (OpenGL coordinates start at bottom)
538
+ y = tile_height - y - 1
539
+ if qx >= screen_width or qy >= screen_height:
540
+ output_img[tile, y, x, 0] = 0.0
541
+ output_img[tile, y, x, 1] = 0.0
542
+ output_img[tile, y, x, 2] = 0.0
543
+ return # prevent out-of-bounds access
544
+ pixel *= 3
545
+ r = float(input_img[pixel + 0])
546
+ g = float(input_img[pixel + 1])
547
+ b = float(input_img[pixel + 2])
548
+ output_img[tile, y, x, 0] = r / 255.0
549
+ output_img[tile, y, x, 1] = g / 255.0
550
+ output_img[tile, y, x, 2] = b / 255.0
551
+
552
+
553
+ @wp.kernel
554
+ def copy_rgb_frame_tiles_uint8(
555
+ input_img: wp.array(dtype=wp.uint8),
556
+ positions: wp.array(dtype=int, ndim=2),
557
+ screen_width: int,
558
+ screen_height: int,
559
+ tile_height: int,
560
+ # outputs
561
+ output_img: wp.array(dtype=wp.uint8, ndim=4),
562
+ ):
563
+ tile, x, y = wp.tid()
564
+ p = positions[tile]
565
+ qx = x + p[0]
566
+ qy = y + p[1]
567
+ pixel = qy * screen_width + qx
568
+ # flip vertically (OpenGL coordinates start at bottom)
569
+ y = tile_height - y - 1
570
+ if qx >= screen_width or qy >= screen_height:
571
+ output_img[tile, y, x, 0] = wp.uint8(0)
572
+ output_img[tile, y, x, 1] = wp.uint8(0)
573
+ output_img[tile, y, x, 2] = wp.uint8(0)
574
+ return # prevent out-of-bounds access
575
+ pixel *= 3
576
+ output_img[tile, y, x, 0] = input_img[pixel + 0]
577
+ output_img[tile, y, x, 1] = input_img[pixel + 1]
578
+ output_img[tile, y, x, 2] = input_img[pixel + 2]
579
+
580
+
581
+ @wp.kernel
582
+ def copy_depth_frame_tiles(
583
+ input_img: wp.array(dtype=wp.float32),
584
+ positions: wp.array(dtype=int, ndim=2),
585
+ screen_width: int,
586
+ screen_height: int,
587
+ tile_height: int,
588
+ near: float,
589
+ far: float,
590
+ # outputs
591
+ output_img: wp.array(dtype=wp.float32, ndim=4),
592
+ ):
593
+ tile, x, y = wp.tid()
594
+ p = positions[tile]
595
+ qx = x + p[0]
596
+ qy = y + p[1]
597
+ pixel = qy * screen_width + qx
598
+ # flip vertically (OpenGL coordinates start at bottom)
599
+ y = tile_height - y - 1
600
+ if qx >= screen_width or qy >= screen_height:
601
+ output_img[tile, y, x, 0] = far
602
+ return # prevent out-of-bounds access
603
+ d = 2.0 * input_img[pixel] - 1.0
604
+ d = 2.0 * near * far / ((far - near) * d - near - far)
605
+ output_img[tile, y, x, 0] = -d
606
+
607
+
608
+ @wp.kernel
609
+ def copy_rgb_frame_tile(
610
+ input_img: wp.array(dtype=wp.uint8),
611
+ offset_x: int,
612
+ offset_y: int,
613
+ screen_width: int,
614
+ screen_height: int,
615
+ tile_height: int,
616
+ # outputs
617
+ output_img: wp.array(dtype=float, ndim=4),
618
+ ):
619
+ tile, x, y = wp.tid()
620
+ qx = x + offset_x
621
+ qy = y + offset_y
622
+ pixel = qy * screen_width + qx
623
+ # flip vertically (OpenGL coordinates start at bottom)
624
+ y = tile_height - y - 1
625
+ if qx >= screen_width or qy >= screen_height:
626
+ output_img[tile, y, x, 0] = 0.0
627
+ output_img[tile, y, x, 1] = 0.0
628
+ output_img[tile, y, x, 2] = 0.0
629
+ return # prevent out-of-bounds access
630
+ pixel *= 3
631
+ r = float(input_img[pixel + 0])
632
+ g = float(input_img[pixel + 1])
633
+ b = float(input_img[pixel + 2])
634
+ output_img[tile, y, x, 0] = r / 255.0
635
+ output_img[tile, y, x, 1] = g / 255.0
636
+ output_img[tile, y, x, 2] = b / 255.0
637
+
638
+
639
+ @wp.kernel
640
+ def copy_rgb_frame_tile_uint8(
641
+ input_img: wp.array(dtype=wp.uint8),
642
+ offset_x: int,
643
+ offset_y: int,
644
+ screen_width: int,
645
+ screen_height: int,
646
+ tile_height: int,
647
+ # outputs
648
+ output_img: wp.array(dtype=wp.uint8, ndim=4),
649
+ ):
650
+ tile, x, y = wp.tid()
651
+ qx = x + offset_x
652
+ qy = y + offset_y
653
+ pixel = qy * screen_width + qx
654
+ # flip vertically (OpenGL coordinates start at bottom)
655
+ y = tile_height - y - 1
656
+ if qx >= screen_width or qy >= screen_height:
657
+ output_img[tile, y, x, 0] = wp.uint8(0)
658
+ output_img[tile, y, x, 1] = wp.uint8(0)
659
+ output_img[tile, y, x, 2] = wp.uint8(0)
660
+ return # prevent out-of-bounds access
661
+ pixel *= 3
662
+ output_img[tile, y, x, 0] = input_img[pixel + 0]
663
+ output_img[tile, y, x, 1] = input_img[pixel + 1]
664
+ output_img[tile, y, x, 2] = input_img[pixel + 2]
665
+
666
+
667
+ def check_gl_error():
668
+ from pyglet import gl
669
+
670
+ error = gl.glGetError()
671
+ if error != gl.GL_NO_ERROR:
672
+ print(f"OpenGL error: {error}")
673
+
674
+
675
+ class ShapeInstancer:
676
+ """
677
+ Handles instanced rendering for a mesh.
678
+ Note the vertices must be in the 8-dimensional format:
679
+ [3D point, 3D normal, UV texture coordinates]
680
+ """
681
+
682
+ gl = None # Class-level variable to hold the imported module
683
+
684
+ @classmethod
685
+ def initialize_gl(cls):
686
+ if cls.gl is None: # Only import if not already imported
687
+ from pyglet import gl
688
+
689
+ cls.gl = gl
690
+
691
+ def __new__(cls, *args, **kwargs):
692
+ instance = super(ShapeInstancer, cls).__new__(cls)
693
+ instance.instance_transform_gl_buffer = None
694
+ instance.vao = None
695
+ return instance
696
+
697
+ def __init__(self, shape_shader, device):
698
+ self.shape_shader = shape_shader
699
+ self.device = device
700
+ self.face_count = 0
701
+ self.instance_color1_buffer = None
702
+ self.instance_color2_buffer = None
703
+ self.color1 = (1.0, 1.0, 1.0)
704
+ self.color2 = (0.0, 0.0, 0.0)
705
+ self.num_instances = 0
706
+ self.transforms = None
707
+ self.scalings = None
708
+ self._instance_transform_cuda_buffer = None
709
+
710
+ ShapeInstancer.initialize_gl()
711
+
712
+ def __del__(self):
713
+ gl = ShapeInstancer.gl
714
+
715
+ if self.instance_transform_gl_buffer is not None:
716
+ try:
717
+ gl.glDeleteBuffers(1, self.instance_transform_gl_buffer)
718
+ gl.glDeleteBuffers(1, self.instance_color1_buffer)
719
+ gl.glDeleteBuffers(1, self.instance_color2_buffer)
720
+ except gl.GLException:
721
+ pass
722
+ if self.vao is not None:
723
+ try:
724
+ gl.glDeleteVertexArrays(1, self.vao)
725
+ gl.glDeleteBuffers(1, self.vbo)
726
+ gl.glDeleteBuffers(1, self.ebo)
727
+ except gl.GLException:
728
+ pass
729
+
730
+ def register_shape(self, vertices, indices, color1=(1.0, 1.0, 1.0), color2=(0.0, 0.0, 0.0)):
731
+ gl = ShapeInstancer.gl
732
+
733
+ if color1 is not None and color2 is None:
734
+ color2 = np.clip(np.array(color1) + 0.25, 0.0, 1.0)
735
+ self.color1 = color1
736
+ self.color2 = color2
737
+
738
+ gl.glUseProgram(self.shape_shader.id)
739
+
740
+ # Create VAO, VBO, and EBO
741
+ self.vao = gl.GLuint()
742
+ gl.glGenVertexArrays(1, self.vao)
743
+ gl.glBindVertexArray(self.vao)
744
+
745
+ self.vbo = gl.GLuint()
746
+ gl.glGenBuffers(1, self.vbo)
747
+ gl.glBindBuffer(gl.GL_ARRAY_BUFFER, self.vbo)
748
+ gl.glBufferData(gl.GL_ARRAY_BUFFER, vertices.nbytes, vertices.ctypes.data, gl.GL_STATIC_DRAW)
749
+
750
+ self.ebo = gl.GLuint()
751
+ gl.glGenBuffers(1, self.ebo)
752
+ gl.glBindBuffer(gl.GL_ELEMENT_ARRAY_BUFFER, self.ebo)
753
+ gl.glBufferData(gl.GL_ELEMENT_ARRAY_BUFFER, indices.nbytes, indices.ctypes.data, gl.GL_STATIC_DRAW)
754
+
755
+ # Set up vertex attributes
756
+ vertex_stride = vertices.shape[1] * vertices.itemsize
757
+ # positions
758
+ gl.glVertexAttribPointer(0, 3, gl.GL_FLOAT, gl.GL_FALSE, vertex_stride, ctypes.c_void_p(0))
759
+ gl.glEnableVertexAttribArray(0)
760
+ # normals
761
+ gl.glVertexAttribPointer(1, 3, gl.GL_FLOAT, gl.GL_FALSE, vertex_stride, ctypes.c_void_p(3 * vertices.itemsize))
762
+ gl.glEnableVertexAttribArray(1)
763
+ # uv coordinates
764
+ gl.glVertexAttribPointer(2, 2, gl.GL_FLOAT, gl.GL_FALSE, vertex_stride, ctypes.c_void_p(6 * vertices.itemsize))
765
+ gl.glEnableVertexAttribArray(2)
766
+
767
+ gl.glBindVertexArray(0)
768
+
769
+ self.face_count = len(indices)
770
+
771
+ def update_colors(self, colors1, colors2):
772
+ gl = ShapeInstancer.gl
773
+
774
+ if colors1 is None:
775
+ colors1 = np.tile(self.color1, (self.num_instances, 1))
776
+ if colors2 is None:
777
+ colors2 = np.tile(self.color2, (self.num_instances, 1))
778
+ if np.shape(colors1) != (self.num_instances, 3):
779
+ colors1 = np.tile(colors1, (self.num_instances, 1))
780
+ if np.shape(colors2) != (self.num_instances, 3):
781
+ colors2 = np.tile(colors2, (self.num_instances, 1))
782
+ colors1 = np.array(colors1, dtype=np.float32)
783
+ colors2 = np.array(colors2, dtype=np.float32)
784
+
785
+ gl.glBindVertexArray(self.vao)
786
+
787
+ # create buffer for checkerboard colors
788
+ if self.instance_color1_buffer is None:
789
+ self.instance_color1_buffer = gl.GLuint()
790
+ gl.glGenBuffers(1, self.instance_color1_buffer)
791
+ gl.glBindBuffer(gl.GL_ARRAY_BUFFER, self.instance_color1_buffer)
792
+ gl.glBufferData(gl.GL_ARRAY_BUFFER, colors1.nbytes, colors1.ctypes.data, gl.GL_STATIC_DRAW)
793
+
794
+ if self.instance_color2_buffer is None:
795
+ self.instance_color2_buffer = gl.GLuint()
796
+ gl.glGenBuffers(1, self.instance_color2_buffer)
797
+ gl.glBindBuffer(gl.GL_ARRAY_BUFFER, self.instance_color2_buffer)
798
+ gl.glBufferData(gl.GL_ARRAY_BUFFER, colors2.nbytes, colors2.ctypes.data, gl.GL_STATIC_DRAW)
799
+
800
+ gl.glBindBuffer(gl.GL_ARRAY_BUFFER, self.instance_color1_buffer)
801
+ gl.glVertexAttribPointer(7, 3, gl.GL_FLOAT, gl.GL_FALSE, colors1[0].nbytes, ctypes.c_void_p(0))
802
+ gl.glEnableVertexAttribArray(7)
803
+ gl.glVertexAttribDivisor(7, 1)
804
+
805
+ gl.glBindBuffer(gl.GL_ARRAY_BUFFER, self.instance_color2_buffer)
806
+ gl.glVertexAttribPointer(8, 3, gl.GL_FLOAT, gl.GL_FALSE, colors2[0].nbytes, ctypes.c_void_p(0))
807
+ gl.glEnableVertexAttribArray(8)
808
+ gl.glVertexAttribDivisor(8, 1)
809
+
810
+ def allocate_instances(self, positions, rotations=None, colors1=None, colors2=None, scalings=None):
811
+ gl = ShapeInstancer.gl
812
+
813
+ gl.glBindVertexArray(self.vao)
814
+
815
+ self.num_instances = len(positions)
816
+
817
+ # Create instance buffer and bind it as an instanced array
818
+ if self.instance_transform_gl_buffer is None:
819
+ self.instance_transform_gl_buffer = gl.GLuint()
820
+ gl.glGenBuffers(1, self.instance_transform_gl_buffer)
821
+ gl.glBindBuffer(gl.GL_ARRAY_BUFFER, self.instance_transform_gl_buffer)
822
+
823
+ self.instance_ids = wp.array(np.arange(self.num_instances), dtype=wp.int32, device=self.device)
824
+ if rotations is None:
825
+ self.instance_transforms = wp.array(
826
+ [(*pos, 0.0, 0.0, 0.0, 1.0) for pos in positions], dtype=wp.transform, device=self.device
827
+ )
828
+ else:
829
+ self.instance_transforms = wp.array(
830
+ [(*pos, *rot) for pos, rot in zip(positions, rotations)],
831
+ dtype=wp.transform,
832
+ device=self.device,
833
+ )
834
+
835
+ if scalings is None:
836
+ self.instance_scalings = wp.array(
837
+ np.tile((1.0, 1.0, 1.0), (self.num_instances, 1)), dtype=wp.vec3, device=self.device
838
+ )
839
+ else:
840
+ self.instance_scalings = wp.array(scalings, dtype=wp.vec3, device=self.device)
841
+
842
+ vbo_transforms = wp.zeros(dtype=wp.mat44, shape=(self.num_instances,), device=self.device)
843
+
844
+ wp.launch(
845
+ update_vbo_transforms,
846
+ dim=self.num_instances,
847
+ inputs=[
848
+ self.instance_ids,
849
+ None,
850
+ self.instance_transforms,
851
+ self.instance_scalings,
852
+ None,
853
+ ],
854
+ outputs=[
855
+ vbo_transforms,
856
+ ],
857
+ device=self.device,
858
+ record_tape=False,
859
+ )
860
+
861
+ vbo_transforms = vbo_transforms.numpy()
862
+ gl.glBufferData(gl.GL_ARRAY_BUFFER, vbo_transforms.nbytes, vbo_transforms.ctypes.data, gl.GL_DYNAMIC_DRAW)
863
+
864
+ # Create CUDA buffer for instance transforms
865
+ self._instance_transform_cuda_buffer = wp.RegisteredGLBuffer(
866
+ int(self.instance_transform_gl_buffer.value), self.device
867
+ )
868
+
869
+ self.update_colors(colors1, colors2)
870
+
871
+ # Set up instance attribute pointers
872
+ matrix_size = vbo_transforms[0].nbytes
873
+
874
+ gl.glBindBuffer(gl.GL_ARRAY_BUFFER, self.instance_transform_gl_buffer)
875
+
876
+ # we can only send vec4s to the shader, so we need to split the instance transforms matrix into its column vectors
877
+ for i in range(4):
878
+ gl.glVertexAttribPointer(
879
+ 3 + i, 4, gl.GL_FLOAT, gl.GL_FALSE, matrix_size, ctypes.c_void_p(i * matrix_size // 4)
880
+ )
881
+ gl.glEnableVertexAttribArray(3 + i)
882
+ gl.glVertexAttribDivisor(3 + i, 1)
883
+
884
+ gl.glBindVertexArray(0)
885
+
886
+ def update_instances(self, transforms: wp.array = None, scalings: wp.array = None, colors1=None, colors2=None):
887
+ gl = ShapeInstancer.gl
888
+
889
+ if transforms is not None:
890
+ if transforms.device.is_cuda:
891
+ wp_transforms = transforms
892
+ else:
893
+ wp_transforms = transforms.to(self.device)
894
+ self.transforms = wp_transforms
895
+ if scalings is not None:
896
+ if transforms.device.is_cuda:
897
+ wp_scalings = scalings
898
+ else:
899
+ wp_scalings = scalings.to(self.device)
900
+ self.scalings = wp_scalings
901
+
902
+ if transforms is not None or scalings is not None:
903
+ gl.glBindVertexArray(self.vao)
904
+ vbo_transforms = self._instance_transform_cuda_buffer.map(dtype=wp.mat44, shape=(self.num_instances,))
905
+
906
+ wp.launch(
907
+ update_vbo_transforms,
908
+ dim=self.num_instances,
909
+ inputs=[
910
+ self.instance_ids,
911
+ None,
912
+ self.instance_transforms,
913
+ self.instance_scalings,
914
+ None,
915
+ ],
916
+ outputs=[
917
+ vbo_transforms,
918
+ ],
919
+ device=self.device,
920
+ record_tape=False,
921
+ )
922
+
923
+ self._instance_transform_cuda_buffer.unmap()
924
+
925
+ if colors1 is not None or colors2 is not None:
926
+ self.update_colors(colors1, colors2)
927
+
928
+ def render(self):
929
+ gl = ShapeInstancer.gl
930
+
931
+ gl.glUseProgram(self.shape_shader.id)
932
+
933
+ gl.glBindVertexArray(self.vao)
934
+ gl.glDrawElementsInstanced(gl.GL_TRIANGLES, self.face_count, gl.GL_UNSIGNED_INT, None, self.num_instances)
935
+ gl.glBindVertexArray(0)
936
+
937
+ # scope exposes VBO transforms to be set directly by a warp kernel
938
+ def __enter__(self):
939
+ gl = ShapeInstancer.gl
940
+
941
+ gl.glBindVertexArray(self.vao)
942
+ self.vbo_transforms = self._instance_transform_cuda_buffer.map(dtype=wp.mat44, shape=(self.num_instances,))
943
+ return self
944
+
945
+ def __exit__(self, exc_type, exc_value, traceback):
946
+ self._instance_transform_cuda_buffer.unmap()
947
+
948
+
949
+ def str_buffer(string: str):
950
+ return ctypes.c_char_p(string.encode("utf-8"))
951
+
952
+
953
+ def arr_pointer(arr: np.ndarray):
954
+ return arr.astype(np.float32).ctypes.data_as(ctypes.POINTER(ctypes.c_float))
955
+
956
+
957
+ class OpenGLRenderer:
958
+ """
959
+ OpenGLRenderer is a simple OpenGL renderer for rendering 3D shapes and meshes.
960
+ """
961
+
962
+ # number of segments to use for rendering spheres, capsules, cones and cylinders
963
+ default_num_segments = 32
964
+
965
+ gl = None # Class-level variable to hold the imported module
966
+
967
+ @classmethod
968
+ def initialize_gl(cls):
969
+ if cls.gl is None: # Only import if not already imported
970
+ from pyglet import gl
971
+
972
+ cls.gl = gl
973
+
974
+ def __init__(
975
+ self,
976
+ title="Warp",
977
+ scaling=1.0,
978
+ fps=60,
979
+ up_axis="Y",
980
+ screen_width=1024,
981
+ screen_height=768,
982
+ near_plane=1.0,
983
+ far_plane=100.0,
984
+ camera_fov=45.0,
985
+ camera_pos=(0.0, 2.0, 10.0),
986
+ camera_front=(0.0, 0.0, -1.0),
987
+ camera_up=(0.0, 1.0, 0.0),
988
+ background_color=(0.53, 0.8, 0.92),
989
+ draw_grid=True,
990
+ draw_sky=True,
991
+ draw_axis=True,
992
+ show_info=True,
993
+ render_wireframe=False,
994
+ render_depth=False,
995
+ axis_scale=1.0,
996
+ vsync=False,
997
+ headless=None,
998
+ enable_backface_culling=True,
999
+ enable_mouse_interaction=True,
1000
+ enable_keyboard_interaction=True,
1001
+ device=None,
1002
+ ):
1003
+ """
1004
+ Args:
1005
+
1006
+ title (str): The window title.
1007
+ scaling (float): The scaling factor for the scene.
1008
+ fps (int): The target frames per second.
1009
+ up_axis (str): The up axis of the scene. Can be "X", "Y", or "Z".
1010
+ screen_width (int): The width of the window.
1011
+ screen_height (int): The height of the window.
1012
+ near_plane (float): The near clipping plane.
1013
+ far_plane (float): The far clipping plane.
1014
+ camera_fov (float): The camera field of view in degrees.
1015
+ camera_pos (tuple): The initial camera position.
1016
+ camera_front (tuple): The initial camera front direction.
1017
+ camera_up (tuple): The initial camera up direction.
1018
+ background_color (tuple): The background color of the scene.
1019
+ draw_grid (bool): Whether to draw a grid indicating the ground plane.
1020
+ draw_sky (bool): Whether to draw a sky sphere.
1021
+ draw_axis (bool): Whether to draw the coordinate system axes.
1022
+ show_info (bool): Whether to overlay rendering information.
1023
+ render_wireframe (bool): Whether to render scene shapes as wireframes.
1024
+ render_depth (bool): Whether to show the depth buffer instead of the RGB image.
1025
+ axis_scale (float): The scale of the coordinate system axes being rendered (only if ``draw_axis`` is True).
1026
+ vsync (bool): Whether to enable vertical synchronization.
1027
+ headless (bool): Whether to run in headless mode (no window is created). If None, the value is determined by the Pyglet configuration defined in ``pyglet.options["headless"]``.
1028
+ enable_backface_culling (bool): Whether to enable backface culling.
1029
+ enable_mouse_interaction (bool): Whether to enable mouse interaction.
1030
+ enable_keyboard_interaction (bool): Whether to enable keyboard interaction.
1031
+ device (Devicelike): Where to store the internal data.
1032
+
1033
+ Note:
1034
+
1035
+ :class:`OpenGLRenderer` requires Pyglet (version >= 2.0, known to work on 2.0.7) to be installed.
1036
+
1037
+ Headless rendering is supported via EGL on UNIX operating systems. To enable headless rendering, set the following pyglet options before importing ``warp.render``:
1038
+
1039
+ .. code-block:: python
1040
+
1041
+ import pyglet
1042
+
1043
+ pyglet.options["headless"] = True
1044
+
1045
+ import warp.render
1046
+
1047
+ # OpenGLRenderer is instantiated with headless=True by default
1048
+ renderer = warp.render.OpenGLRenderer()
1049
+ """
1050
+ try:
1051
+ import pyglet
1052
+
1053
+ # disable error checking for performance
1054
+ pyglet.options["debug_gl"] = False
1055
+
1056
+ from pyglet.graphics.shader import Shader, ShaderProgram
1057
+ from pyglet.math import Vec3 as PyVec3
1058
+
1059
+ OpenGLRenderer.initialize_gl()
1060
+ gl = OpenGLRenderer.gl
1061
+ except ImportError as e:
1062
+ raise Exception("OpenGLRenderer requires pyglet (version >= 2.0) to be installed.") from e
1063
+
1064
+ self.camera_near_plane = near_plane
1065
+ self.camera_far_plane = far_plane
1066
+ self.camera_fov = camera_fov
1067
+
1068
+ self.background_color = background_color
1069
+ self.draw_grid = draw_grid
1070
+ self.draw_sky = draw_sky
1071
+ self.draw_axis = draw_axis
1072
+ self.show_info = show_info
1073
+ self.render_wireframe = render_wireframe
1074
+ self.render_depth = render_depth
1075
+ self.enable_backface_culling = enable_backface_culling
1076
+
1077
+ if device is None:
1078
+ self._device = wp.get_preferred_device()
1079
+ else:
1080
+ self._device = wp.get_device(device)
1081
+
1082
+ self._title = title
1083
+
1084
+ self.window = pyglet.window.Window(
1085
+ width=screen_width, height=screen_height, caption=title, resizable=True, vsync=vsync, visible=not headless
1086
+ )
1087
+ if headless is None:
1088
+ self.headless = pyglet.options.get("headless", False)
1089
+ else:
1090
+ self.headless = headless
1091
+ self.app = pyglet.app
1092
+
1093
+ # making window current opengl rendering context
1094
+ self._switch_context()
1095
+
1096
+ self.screen_width, self.screen_height = self.window.get_framebuffer_size()
1097
+
1098
+ self.enable_mouse_interaction = enable_mouse_interaction
1099
+ self.enable_keyboard_interaction = enable_keyboard_interaction
1100
+
1101
+ self._camera_speed = 0.04
1102
+ if isinstance(up_axis, int):
1103
+ self._camera_axis = up_axis
1104
+ else:
1105
+ self._camera_axis = "XYZ".index(up_axis.upper())
1106
+ self._last_x, self._last_y = self.screen_width // 2, self.screen_height // 2
1107
+ self._first_mouse = True
1108
+ self._left_mouse_pressed = False
1109
+ self._keys_pressed = defaultdict(bool)
1110
+ self._input_processors = []
1111
+ self._key_callbacks = []
1112
+
1113
+ self.render_2d_callbacks = []
1114
+ self.render_3d_callbacks = []
1115
+
1116
+ self._camera_pos = PyVec3(0.0, 0.0, 0.0)
1117
+ self._camera_front = PyVec3(0.0, 0.0, -1.0)
1118
+ self._camera_up = PyVec3(0.0, 1.0, 0.0)
1119
+ self._scaling = scaling
1120
+
1121
+ self._model_matrix = self.compute_model_matrix(self._camera_axis, scaling)
1122
+ self._inv_model_matrix = np.linalg.inv(self._model_matrix.reshape(4, 4)).flatten()
1123
+ self.update_view_matrix(cam_pos=camera_pos, cam_front=camera_front, cam_up=camera_up)
1124
+ self.update_projection_matrix()
1125
+
1126
+ self._camera_front = self._camera_front.normalize()
1127
+ self._pitch = np.rad2deg(np.arcsin(self._camera_front.y))
1128
+ self._yaw = -np.rad2deg(np.arccos(self._camera_front.x / np.cos(np.deg2rad(self._pitch))))
1129
+
1130
+ self._frame_dt = 1.0 / fps
1131
+ self.time = 0.0
1132
+ self._start_time = time.time()
1133
+ self.clock_time = 0.0
1134
+ self._paused = False
1135
+ self._frame_speed = 0.0
1136
+ self.skip_rendering = False
1137
+ self._skip_frame_counter = 0
1138
+ self._fps_update = 0.0
1139
+ self._fps_render = 0.0
1140
+ self._fps_alpha = 0.1 # low pass filter rate to update FPS stats
1141
+
1142
+ self._body_name = {}
1143
+ self._shapes = []
1144
+ self._shape_geo_hash = {}
1145
+ self._shape_gl_buffers = {}
1146
+ self._shape_instances = defaultdict(list)
1147
+ self._instances = {}
1148
+ self._instance_custom_ids = {}
1149
+ self._instance_shape = {}
1150
+ self._instance_gl_buffers = {}
1151
+ self._instance_transform_gl_buffer = None
1152
+ self._instance_transform_cuda_buffer = None
1153
+ self._instance_color1_buffer = None
1154
+ self._instance_color2_buffer = None
1155
+ self._instance_count = 0
1156
+ self._wp_instance_ids = None
1157
+ self._wp_instance_custom_ids = None
1158
+ self._np_instance_visible = None
1159
+ self._instance_ids = None
1160
+ self._inverse_instance_ids = None
1161
+ self._wp_instance_transforms = None
1162
+ self._wp_instance_scalings = None
1163
+ self._wp_instance_bodies = None
1164
+ self._update_shape_instances = False
1165
+ self._add_shape_instances = False
1166
+
1167
+ # additional shape instancer used for points and line rendering
1168
+ self._shape_instancers = {}
1169
+
1170
+ # instancer for the arrow shapes sof the coordinate system axes
1171
+ self._axis_instancer = None
1172
+
1173
+ # toggle tiled rendering
1174
+ self._tiled_rendering = False
1175
+ self._tile_instances = None
1176
+ self._tile_ncols = 0
1177
+ self._tile_nrows = 0
1178
+ self._tile_width = 0
1179
+ self._tile_height = 0
1180
+ self._tile_viewports = None
1181
+ self._tile_view_matrices = None
1182
+ self._tile_projection_matrices = None
1183
+
1184
+ self._frame_texture = None
1185
+ self._frame_depth_texture = None
1186
+ self._frame_fbo = None
1187
+ self._frame_pbo = None
1188
+
1189
+ if not headless:
1190
+ self.window.push_handlers(on_draw=self._draw)
1191
+ self.window.push_handlers(on_resize=self._window_resize_callback)
1192
+ self.window.push_handlers(on_key_press=self._key_press_callback)
1193
+ self.window.push_handlers(on_close=self._close_callback)
1194
+
1195
+ self._key_handler = pyglet.window.key.KeyStateHandler()
1196
+ self.window.push_handlers(self._key_handler)
1197
+
1198
+ self.window.on_mouse_scroll = self._scroll_callback
1199
+ self.window.on_mouse_drag = self._mouse_drag_callback
1200
+
1201
+ gl.glClearColor(*self.background_color, 1)
1202
+ gl.glEnable(gl.GL_DEPTH_TEST)
1203
+ gl.glDepthMask(True)
1204
+ gl.glDepthRange(0.0, 1.0)
1205
+
1206
+ self._shape_shader = ShaderProgram(
1207
+ Shader(shape_vertex_shader, "vertex"), Shader(shape_fragment_shader, "fragment")
1208
+ )
1209
+ self._grid_shader = ShaderProgram(
1210
+ Shader(grid_vertex_shader, "vertex"), Shader(grid_fragment_shader, "fragment")
1211
+ )
1212
+
1213
+ self._sun_direction = np.array((-0.2, 0.8, 0.3))
1214
+ self._sun_direction /= np.linalg.norm(self._sun_direction)
1215
+ with self._shape_shader:
1216
+ gl.glUniform3f(
1217
+ gl.glGetUniformLocation(self._shape_shader.id, str_buffer("sunDirection")), *self._sun_direction
1218
+ )
1219
+ gl.glUniform3f(gl.glGetUniformLocation(self._shape_shader.id, str_buffer("lightColor")), 1, 1, 1)
1220
+ self._loc_shape_model = gl.glGetUniformLocation(self._shape_shader.id, str_buffer("model"))
1221
+ self._loc_shape_view = gl.glGetUniformLocation(self._shape_shader.id, str_buffer("view"))
1222
+ self._loc_shape_projection = gl.glGetUniformLocation(self._shape_shader.id, str_buffer("projection"))
1223
+ self._loc_shape_view_pos = gl.glGetUniformLocation(self._shape_shader.id, str_buffer("viewPos"))
1224
+ gl.glUniform3f(self._loc_shape_view_pos, 0, 0, 10)
1225
+
1226
+ # create grid data
1227
+ limit = 10.0
1228
+ ticks = np.linspace(-limit, limit, 21)
1229
+ grid_vertices = []
1230
+ for i in ticks:
1231
+ if self._camera_axis == 0:
1232
+ grid_vertices.extend([0, -limit, i, 0, limit, i])
1233
+ grid_vertices.extend([0, i, -limit, 0, i, limit])
1234
+ elif self._camera_axis == 1:
1235
+ grid_vertices.extend([-limit, 0, i, limit, 0, i])
1236
+ grid_vertices.extend([i, 0, -limit, i, 0, limit])
1237
+ elif self._camera_axis == 2:
1238
+ grid_vertices.extend([-limit, i, 0, limit, i, 0])
1239
+ grid_vertices.extend([i, -limit, 0, i, limit, 0])
1240
+ grid_vertices = np.array(grid_vertices, dtype=np.float32)
1241
+ self._grid_vertex_count = len(grid_vertices) // 3
1242
+
1243
+ with self._grid_shader:
1244
+ self._grid_vao = gl.GLuint()
1245
+ gl.glGenVertexArrays(1, self._grid_vao)
1246
+ gl.glBindVertexArray(self._grid_vao)
1247
+
1248
+ self._grid_vbo = gl.GLuint()
1249
+ gl.glGenBuffers(1, self._grid_vbo)
1250
+ gl.glBindBuffer(gl.GL_ARRAY_BUFFER, self._grid_vbo)
1251
+ gl.glBufferData(gl.GL_ARRAY_BUFFER, grid_vertices.nbytes, grid_vertices.ctypes.data, gl.GL_STATIC_DRAW)
1252
+
1253
+ self._loc_grid_view = gl.glGetUniformLocation(self._grid_shader.id, str_buffer("view"))
1254
+ self._loc_grid_model = gl.glGetUniformLocation(self._grid_shader.id, str_buffer("model"))
1255
+ self._loc_grid_projection = gl.glGetUniformLocation(self._grid_shader.id, str_buffer("projection"))
1256
+
1257
+ self._loc_grid_pos_attribute = gl.glGetAttribLocation(self._grid_shader.id, str_buffer("position"))
1258
+ gl.glVertexAttribPointer(self._loc_grid_pos_attribute, 3, gl.GL_FLOAT, gl.GL_FALSE, 0, None)
1259
+ gl.glEnableVertexAttribArray(self._loc_grid_pos_attribute)
1260
+
1261
+ # create sky data
1262
+ self._sky_shader = ShaderProgram(Shader(sky_vertex_shader, "vertex"), Shader(sky_fragment_shader, "fragment"))
1263
+
1264
+ with self._sky_shader:
1265
+ self._loc_sky_view = gl.glGetUniformLocation(self._sky_shader.id, str_buffer("view"))
1266
+ self._loc_sky_inv_model = gl.glGetUniformLocation(self._sky_shader.id, str_buffer("inv_model"))
1267
+ self._loc_sky_projection = gl.glGetUniformLocation(self._sky_shader.id, str_buffer("projection"))
1268
+
1269
+ self._loc_sky_color1 = gl.glGetUniformLocation(self._sky_shader.id, str_buffer("color1"))
1270
+ self._loc_sky_color2 = gl.glGetUniformLocation(self._sky_shader.id, str_buffer("color2"))
1271
+ self._loc_sky_far_plane = gl.glGetUniformLocation(self._sky_shader.id, str_buffer("farPlane"))
1272
+ gl.glUniform3f(self._loc_sky_color1, *background_color)
1273
+ # glUniform3f(self._loc_sky_color2, *np.clip(np.array(background_color)+0.5, 0.0, 1.0))
1274
+ gl.glUniform3f(self._loc_sky_color2, 0.8, 0.4, 0.05)
1275
+ gl.glUniform1f(self._loc_sky_far_plane, self.camera_far_plane)
1276
+ self._loc_sky_view_pos = gl.glGetUniformLocation(self._sky_shader.id, str_buffer("viewPos"))
1277
+ gl.glUniform3f(
1278
+ gl.glGetUniformLocation(self._sky_shader.id, str_buffer("sunDirection")), *self._sun_direction
1279
+ )
1280
+
1281
+ # create VAO, VBO, and EBO
1282
+ self._sky_vao = gl.GLuint()
1283
+ gl.glGenVertexArrays(1, self._sky_vao)
1284
+ gl.glBindVertexArray(self._sky_vao)
1285
+
1286
+ vertices, indices = self._create_sphere_mesh(self.camera_far_plane * 0.9, 32, 32, reverse_winding=True)
1287
+ self._sky_tri_count = len(indices)
1288
+
1289
+ self._sky_vbo = gl.GLuint()
1290
+ gl.glGenBuffers(1, self._sky_vbo)
1291
+ gl.glBindBuffer(gl.GL_ARRAY_BUFFER, self._sky_vbo)
1292
+ gl.glBufferData(gl.GL_ARRAY_BUFFER, vertices.nbytes, vertices.ctypes.data, gl.GL_STATIC_DRAW)
1293
+
1294
+ self._sky_ebo = gl.GLuint()
1295
+ gl.glGenBuffers(1, self._sky_ebo)
1296
+ gl.glBindBuffer(gl.GL_ELEMENT_ARRAY_BUFFER, self._sky_ebo)
1297
+ gl.glBufferData(gl.GL_ELEMENT_ARRAY_BUFFER, indices.nbytes, indices.ctypes.data, gl.GL_STATIC_DRAW)
1298
+
1299
+ # set up vertex attributes
1300
+ vertex_stride = vertices.shape[1] * vertices.itemsize
1301
+ # positions
1302
+ gl.glVertexAttribPointer(0, 3, gl.GL_FLOAT, gl.GL_FALSE, vertex_stride, ctypes.c_void_p(0))
1303
+ gl.glEnableVertexAttribArray(0)
1304
+ # normals
1305
+ gl.glVertexAttribPointer(1, 3, gl.GL_FLOAT, gl.GL_FALSE, vertex_stride, ctypes.c_void_p(3 * vertices.itemsize))
1306
+ gl.glEnableVertexAttribArray(1)
1307
+ # uv coordinates
1308
+ gl.glVertexAttribPointer(2, 2, gl.GL_FLOAT, gl.GL_FALSE, vertex_stride, ctypes.c_void_p(6 * vertices.itemsize))
1309
+ gl.glEnableVertexAttribArray(2)
1310
+
1311
+ gl.glBindVertexArray(0)
1312
+
1313
+ self._last_time = time.time()
1314
+ self._last_begin_frame_time = self._last_time
1315
+ self._last_end_frame_time = self._last_time
1316
+
1317
+ # create arrow shapes for the coordinate system axes
1318
+ vertices, indices = self._create_arrow_mesh(
1319
+ base_radius=0.02 * axis_scale, base_height=0.85 * axis_scale, cap_height=0.15 * axis_scale
1320
+ )
1321
+ self._axis_instancer = ShapeInstancer(self._shape_shader, self._device)
1322
+ self._axis_instancer.register_shape(vertices, indices)
1323
+ sqh = np.sqrt(0.5)
1324
+ self._axis_instancer.allocate_instances(
1325
+ positions=[(0.0, 0.0, 0.0), (0.0, 0.0, 0.0), (0.0, 0.0, 0.0)],
1326
+ rotations=[(0.0, 0.0, 0.0, 1.0), (0.0, 0.0, -sqh, sqh), (sqh, 0.0, 0.0, sqh)],
1327
+ colors1=[(0.0, 1.0, 0.0), (1.0, 0.0, 0.0), (0.0, 0.0, 1.0)],
1328
+ colors2=[(0.0, 1.0, 0.0), (1.0, 0.0, 0.0), (0.0, 0.0, 1.0)],
1329
+ )
1330
+
1331
+ # create frame buffer for rendering to a texture
1332
+ self._frame_texture = None
1333
+ self._frame_depth_texture = None
1334
+ self._frame_fbo = None
1335
+ self._setup_framebuffer()
1336
+
1337
+ # fmt: off
1338
+ # set up VBO for the quad that is rendered to the user window with the texture
1339
+ self._frame_vertices = np.array([
1340
+ # Positions TexCoords
1341
+ -1.0, -1.0, 0.0, 0.0,
1342
+ 1.0, -1.0, 1.0, 0.0,
1343
+ 1.0, 1.0, 1.0, 1.0,
1344
+ -1.0, 1.0, 0.0, 1.0
1345
+ ], dtype=np.float32)
1346
+ # fmt: on
1347
+
1348
+ self._frame_indices = np.array([0, 1, 2, 2, 3, 0], dtype=np.uint32)
1349
+
1350
+ self._frame_vao = gl.GLuint()
1351
+ gl.glGenVertexArrays(1, self._frame_vao)
1352
+ gl.glBindVertexArray(self._frame_vao)
1353
+
1354
+ self._frame_vbo = gl.GLuint()
1355
+ gl.glGenBuffers(1, self._frame_vbo)
1356
+ gl.glBindBuffer(gl.GL_ARRAY_BUFFER, self._frame_vbo)
1357
+ gl.glBufferData(
1358
+ gl.GL_ARRAY_BUFFER, self._frame_vertices.nbytes, self._frame_vertices.ctypes.data, gl.GL_STATIC_DRAW
1359
+ )
1360
+
1361
+ self._frame_ebo = gl.GLuint()
1362
+ gl.glGenBuffers(1, self._frame_ebo)
1363
+ gl.glBindBuffer(gl.GL_ELEMENT_ARRAY_BUFFER, self._frame_ebo)
1364
+ gl.glBufferData(
1365
+ gl.GL_ELEMENT_ARRAY_BUFFER, self._frame_indices.nbytes, self._frame_indices.ctypes.data, gl.GL_STATIC_DRAW
1366
+ )
1367
+
1368
+ gl.glVertexAttribPointer(0, 2, gl.GL_FLOAT, gl.GL_FALSE, 4 * self._frame_vertices.itemsize, ctypes.c_void_p(0))
1369
+ gl.glEnableVertexAttribArray(0)
1370
+ gl.glVertexAttribPointer(
1371
+ 1, 2, gl.GL_FLOAT, gl.GL_FALSE, 4 * self._frame_vertices.itemsize, ctypes.c_void_p(2 * vertices.itemsize)
1372
+ )
1373
+ gl.glEnableVertexAttribArray(1)
1374
+
1375
+ self._frame_shader = ShaderProgram(
1376
+ Shader(frame_vertex_shader, "vertex"), Shader(frame_fragment_shader, "fragment")
1377
+ )
1378
+ gl.glUseProgram(self._frame_shader.id)
1379
+ self._frame_loc_texture = gl.glGetUniformLocation(self._frame_shader.id, str_buffer("textureSampler"))
1380
+
1381
+ self._frame_depth_shader = ShaderProgram(
1382
+ Shader(frame_vertex_shader, "vertex"), Shader(frame_depth_fragment_shader, "fragment")
1383
+ )
1384
+ gl.glUseProgram(self._frame_depth_shader.id)
1385
+ self._frame_loc_depth_texture = gl.glGetUniformLocation(
1386
+ self._frame_depth_shader.id, str_buffer("textureSampler")
1387
+ )
1388
+
1389
+ # unbind the VBO and VAO
1390
+ gl.glBindBuffer(gl.GL_ARRAY_BUFFER, 0)
1391
+ gl.glBindVertexArray(0)
1392
+
1393
+ # update model matrix
1394
+ self.scaling = scaling
1395
+
1396
+ check_gl_error()
1397
+
1398
+ # create text to render stats on the screen
1399
+ self._info_label = pyglet.text.Label(
1400
+ "",
1401
+ font_name="Arial",
1402
+ font_size=12,
1403
+ color=(255, 255, 255, 255),
1404
+ x=10,
1405
+ y=10,
1406
+ anchor_x="left",
1407
+ anchor_y="top",
1408
+ multiline=True,
1409
+ width=400,
1410
+ )
1411
+
1412
+ if not headless:
1413
+ # set up our own event handling so we can synchronously render frames
1414
+ # by calling update() in a loop
1415
+ from pyglet.window import Window
1416
+
1417
+ Window._enable_event_queue = False
1418
+
1419
+ self.window.dispatch_pending_events()
1420
+
1421
+ platform_event_loop = self.app.platform_event_loop
1422
+ platform_event_loop.start()
1423
+
1424
+ # start event loop
1425
+ self.app.event_loop.dispatch_event("on_enter")
1426
+
1427
+ @property
1428
+ def paused(self):
1429
+ return self._paused
1430
+
1431
+ @paused.setter
1432
+ def paused(self, value):
1433
+ self._paused = value
1434
+ if value:
1435
+ self.window.set_caption(f"{self._title} (paused)")
1436
+ else:
1437
+ self.window.set_caption(self._title)
1438
+
1439
+ @property
1440
+ def has_exit(self):
1441
+ return self.app.event_loop.has_exit
1442
+
1443
+ def clear(self):
1444
+ gl = OpenGLRenderer.gl
1445
+
1446
+ self._switch_context()
1447
+
1448
+ if not self.headless:
1449
+ self.app.event_loop.dispatch_event("on_exit")
1450
+ self.app.platform_event_loop.stop()
1451
+
1452
+ if self._instance_transform_gl_buffer is not None:
1453
+ try:
1454
+ gl.glDeleteBuffers(1, self._instance_transform_gl_buffer)
1455
+ gl.glDeleteBuffers(1, self._instance_color1_buffer)
1456
+ gl.glDeleteBuffers(1, self._instance_color2_buffer)
1457
+ except gl.GLException:
1458
+ pass
1459
+ for vao, vbo, ebo, _, _vertex_cuda_buffer in self._shape_gl_buffers.values():
1460
+ try:
1461
+ gl.glDeleteVertexArrays(1, vao)
1462
+ gl.glDeleteBuffers(1, vbo)
1463
+ gl.glDeleteBuffers(1, ebo)
1464
+ except gl.GLException:
1465
+ pass
1466
+
1467
+ self._body_name.clear()
1468
+ self._shapes.clear()
1469
+ self._shape_geo_hash.clear()
1470
+ self._shape_gl_buffers.clear()
1471
+ self._shape_instances.clear()
1472
+ self._instances.clear()
1473
+ self._instance_shape.clear()
1474
+ self._instance_gl_buffers.clear()
1475
+ self._instance_transform_gl_buffer = None
1476
+ self._instance_transform_cuda_buffer = None
1477
+ self._instance_color1_buffer = None
1478
+ self._instance_color2_buffer = None
1479
+ self._wp_instance_ids = None
1480
+ self._wp_instance_custom_ids = None
1481
+ self._wp_instance_transforms = None
1482
+ self._wp_instance_scalings = None
1483
+ self._wp_instance_bodies = None
1484
+ self._np_instance_visible = None
1485
+ self._update_shape_instances = False
1486
+
1487
+ def close(self):
1488
+ self.clear()
1489
+ self.window.close()
1490
+
1491
+ @property
1492
+ def tiled_rendering(self):
1493
+ return self._tiled_rendering
1494
+
1495
+ @tiled_rendering.setter
1496
+ def tiled_rendering(self, value):
1497
+ if value:
1498
+ assert self._tile_instances is not None, "Tiled rendering is not set up. Call setup_tiled_rendering first."
1499
+ self._tiled_rendering = value
1500
+
1501
+ def setup_tiled_rendering(
1502
+ self,
1503
+ instances: List[List[int]],
1504
+ rescale_window: bool = False,
1505
+ tile_width: Optional[int] = None,
1506
+ tile_height: Optional[int] = None,
1507
+ tile_ncols: Optional[int] = None,
1508
+ tile_nrows: Optional[int] = None,
1509
+ tile_positions: Optional[List[Tuple[int]]] = None,
1510
+ tile_sizes: Optional[List[Tuple[int]]] = None,
1511
+ projection_matrices: Optional[List[Mat44]] = None,
1512
+ view_matrices: Optional[List[Mat44]] = None,
1513
+ ):
1514
+ """
1515
+ Set up tiled rendering where the render buffer is split into multiple tiles that can visualize
1516
+ different shape instances of the scene with different view and projection matrices.
1517
+ See :meth:`get_pixels` which allows to retrieve the pixels of for each tile.
1518
+ See :meth:`update_tile` which allows to update the shape instances, projection matrix, view matrix, tile size, or tile position for a given tile.
1519
+
1520
+ :param instances: A list of lists of shape instance ids. Each list of shape instance ids
1521
+ will be rendered into a separate tile.
1522
+ :param rescale_window: If True, the window will be resized to fit the tiles.
1523
+ :param tile_width: The width of each tile in pixels (optional).
1524
+ :param tile_height: The height of each tile in pixels (optional).
1525
+ :param tile_ncols: The number of tiles rendered horizontally (optional). Will be considered
1526
+ if `tile_width` is set to compute the tile positions, unless `tile_positions` is defined.
1527
+ :param tile_positions: A list of (x, y) tuples specifying the position of each tile in pixels.
1528
+ If None, the tiles will be arranged in a square grid, or, if `tile_ncols` and `tile_nrows`
1529
+ is set, in a grid with the specified number of columns and rows.
1530
+ :param tile_sizes: A list of (width, height) tuples specifying the size of each tile in pixels.
1531
+ If None, the tiles will have the same size as specified by `tile_width` and `tile_height`.
1532
+ :param projection_matrices: A list of projection matrices for each tile (each view matrix is
1533
+ either a flattened 16-dimensional array or a 4x4 matrix).
1534
+ If the entire array is None, or only a view instances, the projection matrices for all, or these
1535
+ instances, respectively, will be derived from the current render settings.
1536
+ :param view_matrices: A list of view matrices for each tile (each view matrix is either a flattened
1537
+ 16-dimensional array or a 4x4 matrix).
1538
+ If the entire array is None, or only a view instances, the view matrices for all, or these
1539
+ instances, respectively, will be derived from the current camera settings and be
1540
+ updated when the camera is moved.
1541
+ """
1542
+
1543
+ assert len(instances) > 0 and all(isinstance(i, list) for i in instances), "Invalid tile instances."
1544
+
1545
+ self._tile_instances = instances
1546
+ n = len(self._tile_instances)
1547
+
1548
+ if tile_positions is None or tile_sizes is None:
1549
+ if tile_ncols is None or tile_nrows is None:
1550
+ # try to fit the tiles into a square
1551
+ self._tile_ncols = int(np.ceil(np.sqrt(n)))
1552
+ self._tile_nrows = int(np.ceil(n / float(self._tile_ncols)))
1553
+ else:
1554
+ self._tile_ncols = tile_ncols
1555
+ self._tile_nrows = tile_nrows
1556
+ self._tile_width = tile_width or max(32, self.screen_width // self._tile_ncols)
1557
+ self._tile_height = tile_height or max(32, self.screen_height // self._tile_nrows)
1558
+ self._tile_viewports = [
1559
+ (i * self._tile_width, j * self._tile_height, self._tile_width, self._tile_height)
1560
+ for i in range(self._tile_ncols)
1561
+ for j in range(self._tile_nrows)
1562
+ ]
1563
+ if rescale_window:
1564
+ self.window.set_size(self._tile_width * self._tile_ncols, self._tile_height * self._tile_nrows)
1565
+ else:
1566
+ assert len(tile_positions) == n and len(tile_sizes) == n, (
1567
+ "Number of tiles does not match number of instances."
1568
+ )
1569
+ self._tile_ncols = None
1570
+ self._tile_nrows = None
1571
+ self._tile_width = None
1572
+ self._tile_height = None
1573
+ if all(tile_sizes[i][0] == tile_sizes[0][0] for i in range(n)):
1574
+ # tiles all have the same width
1575
+ self._tile_width = tile_sizes[0][0]
1576
+ if all(tile_sizes[i][1] == tile_sizes[0][1] for i in range(n)):
1577
+ # tiles all have the same height
1578
+ self._tile_height = tile_sizes[0][1]
1579
+ self._tile_viewports = [(x, y, w, h) for (x, y), (w, h) in zip(tile_positions, tile_sizes)]
1580
+
1581
+ if projection_matrices is None:
1582
+ projection_matrices = [None] * n
1583
+ self._tile_projection_matrices = []
1584
+ for i, p in enumerate(projection_matrices):
1585
+ if p is None:
1586
+ w, h = self._tile_viewports[i][2:]
1587
+ self._tile_projection_matrices.append(
1588
+ self.compute_projection_matrix(
1589
+ self.camera_fov, w / h, self.camera_near_plane, self.camera_far_plane
1590
+ )
1591
+ )
1592
+ else:
1593
+ self._tile_projection_matrices.append(np.array(p).flatten())
1594
+
1595
+ if view_matrices is None:
1596
+ self._tile_view_matrices = [None] * n
1597
+ else:
1598
+ self._tile_view_matrices = [np.array(m).flatten() for m in view_matrices]
1599
+
1600
+ self._tiled_rendering = True
1601
+
1602
+ def update_tile(
1603
+ self,
1604
+ tile_id,
1605
+ instances: Optional[List[int]] = None,
1606
+ projection_matrix: Optional[Mat44] = None,
1607
+ view_matrix: Optional[Mat44] = None,
1608
+ tile_size: Optional[Tuple[int]] = None,
1609
+ tile_position: Optional[Tuple[int]] = None,
1610
+ ):
1611
+ """
1612
+ Update the shape instances, projection matrix, view matrix, tile size, or tile position
1613
+ for a given tile given its index.
1614
+
1615
+ :param tile_id: The index of the tile to update.
1616
+ :param instances: A list of shape instance ids (optional).
1617
+ :param projection_matrix: A projection matrix (optional).
1618
+ :param view_matrix: A view matrix (optional).
1619
+ :param tile_size: A (width, height) tuple specifying the size of the tile in pixels (optional).
1620
+ :param tile_position: A (x, y) tuple specifying the position of the tile in pixels (optional).
1621
+ """
1622
+
1623
+ assert self._tile_instances is not None, "Tiled rendering is not set up. Call setup_tiled_rendering first."
1624
+ assert tile_id < len(self._tile_instances), "Invalid tile id."
1625
+
1626
+ if instances is not None:
1627
+ self._tile_instances[tile_id] = instances
1628
+ if projection_matrix is not None:
1629
+ self._tile_projection_matrices[tile_id] = np.array(projection_matrix).flatten()
1630
+ if view_matrix is not None:
1631
+ self._tile_view_matrices[tile_id] = np.array(view_matrix).flatten()
1632
+ (x, y, w, h) = self._tile_viewports[tile_id]
1633
+ if tile_size is not None:
1634
+ w, h = tile_size
1635
+ if tile_position is not None:
1636
+ x, y = tile_position
1637
+ self._tile_viewports[tile_id] = (x, y, w, h)
1638
+
1639
+ def _setup_framebuffer(self):
1640
+ gl = OpenGLRenderer.gl
1641
+
1642
+ self._switch_context()
1643
+
1644
+ if self._frame_texture is None:
1645
+ self._frame_texture = gl.GLuint()
1646
+ gl.glGenTextures(1, self._frame_texture)
1647
+ if self._frame_depth_texture is None:
1648
+ self._frame_depth_texture = gl.GLuint()
1649
+ gl.glGenTextures(1, self._frame_depth_texture)
1650
+
1651
+ # set up RGB texture
1652
+ gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, 0)
1653
+ gl.glBindBuffer(gl.GL_PIXEL_UNPACK_BUFFER, 0)
1654
+ gl.glBindTexture(gl.GL_TEXTURE_2D, self._frame_texture)
1655
+ gl.glTexImage2D(
1656
+ gl.GL_TEXTURE_2D,
1657
+ 0,
1658
+ gl.GL_RGB,
1659
+ self.screen_width,
1660
+ self.screen_height,
1661
+ 0,
1662
+ gl.GL_RGB,
1663
+ gl.GL_UNSIGNED_BYTE,
1664
+ None,
1665
+ )
1666
+ gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MIN_FILTER, gl.GL_LINEAR)
1667
+ gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MAG_FILTER, gl.GL_LINEAR)
1668
+
1669
+ # set up depth texture
1670
+ gl.glBindTexture(gl.GL_TEXTURE_2D, self._frame_depth_texture)
1671
+ gl.glTexImage2D(
1672
+ gl.GL_TEXTURE_2D,
1673
+ 0,
1674
+ gl.GL_DEPTH_COMPONENT32,
1675
+ self.screen_width,
1676
+ self.screen_height,
1677
+ 0,
1678
+ gl.GL_DEPTH_COMPONENT,
1679
+ gl.GL_FLOAT,
1680
+ None,
1681
+ )
1682
+ gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MIN_FILTER, gl.GL_LINEAR)
1683
+ gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MAG_FILTER, gl.GL_LINEAR)
1684
+ gl.glBindTexture(gl.GL_TEXTURE_2D, 0)
1685
+
1686
+ # create a framebuffer object (FBO)
1687
+ if self._frame_fbo is None:
1688
+ self._frame_fbo = gl.GLuint()
1689
+ gl.glGenFramebuffers(1, self._frame_fbo)
1690
+ gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, self._frame_fbo)
1691
+
1692
+ # attach the texture to the FBO as its color attachment
1693
+ gl.glFramebufferTexture2D(
1694
+ gl.GL_FRAMEBUFFER, gl.GL_COLOR_ATTACHMENT0, gl.GL_TEXTURE_2D, self._frame_texture, 0
1695
+ )
1696
+ # attach the depth texture to the FBO as its depth attachment
1697
+ gl.glFramebufferTexture2D(
1698
+ gl.GL_FRAMEBUFFER, gl.GL_DEPTH_ATTACHMENT, gl.GL_TEXTURE_2D, self._frame_depth_texture, 0
1699
+ )
1700
+
1701
+ if gl.glCheckFramebufferStatus(gl.GL_FRAMEBUFFER) != gl.GL_FRAMEBUFFER_COMPLETE:
1702
+ print("Framebuffer is not complete!", flush=True)
1703
+ gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, 0)
1704
+ sys.exit(1)
1705
+
1706
+ # unbind the FBO (switch back to the default framebuffer)
1707
+ gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, 0)
1708
+
1709
+ if self._frame_pbo is None:
1710
+ self._frame_pbo = gl.GLuint()
1711
+ gl.glGenBuffers(1, self._frame_pbo) # generate 1 buffer reference
1712
+ gl.glBindBuffer(gl.GL_PIXEL_PACK_BUFFER, self._frame_pbo) # binding to this buffer
1713
+
1714
+ # allocate memory for PBO
1715
+ rgb_bytes_per_pixel = 3
1716
+ depth_bytes_per_pixel = 4
1717
+ pixels = np.zeros(
1718
+ (self.screen_height, self.screen_width, rgb_bytes_per_pixel + depth_bytes_per_pixel), dtype=np.uint8
1719
+ )
1720
+ gl.glBufferData(gl.GL_PIXEL_PACK_BUFFER, pixels.nbytes, pixels.ctypes.data, gl.GL_DYNAMIC_DRAW)
1721
+ gl.glBindBuffer(gl.GL_PIXEL_PACK_BUFFER, 0)
1722
+
1723
+ @staticmethod
1724
+ def compute_projection_matrix(
1725
+ fov: float,
1726
+ aspect_ratio: float,
1727
+ near_plane: float,
1728
+ far_plane: float,
1729
+ ) -> Mat44:
1730
+ """
1731
+ Compute a projection matrix given the field of view, aspect ratio, near plane, and far plane.
1732
+
1733
+ :param fov: The field of view in degrees.
1734
+ :param aspect_ratio: The aspect ratio (width / height).
1735
+ :param near_plane: The near plane.
1736
+ :param far_plane: The far plane.
1737
+ :return: A projection matrix.
1738
+ """
1739
+
1740
+ from pyglet.math import Mat4 as PyMat4
1741
+
1742
+ return np.array(PyMat4.perspective_projection(aspect_ratio, near_plane, far_plane, fov))
1743
+
1744
+ def update_projection_matrix(self):
1745
+ if self.screen_height == 0:
1746
+ return
1747
+ aspect_ratio = self.screen_width / self.screen_height
1748
+ self._projection_matrix = self.compute_projection_matrix(
1749
+ self.camera_fov, aspect_ratio, self.camera_near_plane, self.camera_far_plane
1750
+ )
1751
+
1752
+ @property
1753
+ def camera_pos(self):
1754
+ return self._camera_pos
1755
+
1756
+ @camera_pos.setter
1757
+ def camera_pos(self, value):
1758
+ self.update_view_matrix(cam_pos=value)
1759
+
1760
+ @property
1761
+ def camera_front(self):
1762
+ return self._camera_front
1763
+
1764
+ @camera_front.setter
1765
+ def camera_front(self, value):
1766
+ self.update_view_matrix(cam_front=value)
1767
+
1768
+ @property
1769
+ def camera_up(self):
1770
+ return self._camera_up
1771
+
1772
+ @camera_up.setter
1773
+ def camera_up(self, value):
1774
+ self.update_view_matrix(cam_up=value)
1775
+
1776
+ def compute_view_matrix(self, cam_pos, cam_front, cam_up):
1777
+ from pyglet.math import Mat4, Vec3
1778
+
1779
+ model = np.array(self._model_matrix).reshape((4, 4))
1780
+ cp = model @ np.array([*cam_pos / self._scaling, 1.0])
1781
+ cf = model @ np.array([*cam_front / self._scaling, 1.0])
1782
+ up = model @ np.array([*cam_up / self._scaling, 0.0])
1783
+ cp = Vec3(*cp[:3])
1784
+ cf = Vec3(*cf[:3])
1785
+ up = Vec3(*up[:3])
1786
+ return np.array(Mat4.look_at(cp, cp + cf, up), dtype=np.float32)
1787
+
1788
+ def update_view_matrix(self, cam_pos=None, cam_front=None, cam_up=None, stiffness=1.0):
1789
+ from pyglet.math import Vec3
1790
+
1791
+ if cam_pos is not None:
1792
+ self._camera_pos = self._camera_pos * (1.0 - stiffness) + Vec3(*cam_pos) * stiffness
1793
+ if cam_front is not None:
1794
+ self._camera_front = self._camera_front * (1.0 - stiffness) + Vec3(*cam_front) * stiffness
1795
+ if cam_up is not None:
1796
+ self._camera_up = self._camera_up * (1.0 - stiffness) + Vec3(*cam_up) * stiffness
1797
+
1798
+ self._view_matrix = self.compute_view_matrix(self._camera_pos, self._camera_front, self._camera_up)
1799
+
1800
+ @staticmethod
1801
+ def compute_model_matrix(camera_axis: int, scaling: float):
1802
+ if camera_axis == 0:
1803
+ return np.array((0, 0, scaling, 0, scaling, 0, 0, 0, 0, scaling, 0, 0, 0, 0, 0, 1), dtype=np.float32)
1804
+ elif camera_axis == 2:
1805
+ return np.array((0, scaling, 0, 0, 0, 0, scaling, 0, scaling, 0, 0, 0, 0, 0, 0, 1), dtype=np.float32)
1806
+
1807
+ return np.array((scaling, 0, 0, 0, 0, scaling, 0, 0, 0, 0, scaling, 0, 0, 0, 0, 1), dtype=np.float32)
1808
+
1809
+ def update_model_matrix(self, model_matrix: Optional[Mat44] = None):
1810
+ gl = OpenGLRenderer.gl
1811
+
1812
+ self._switch_context()
1813
+
1814
+ if model_matrix is None:
1815
+ self._model_matrix = self.compute_model_matrix(self._camera_axis, self._scaling)
1816
+ else:
1817
+ self._model_matrix = np.array(model_matrix).flatten()
1818
+ self._inv_model_matrix = np.linalg.inv(self._model_matrix.reshape((4, 4))).flatten()
1819
+ # update model view matrix in shaders
1820
+ ptr = arr_pointer(self._model_matrix)
1821
+ gl.glUseProgram(self._shape_shader.id)
1822
+ gl.glUniformMatrix4fv(self._loc_shape_model, 1, gl.GL_FALSE, ptr)
1823
+ gl.glUseProgram(self._grid_shader.id)
1824
+ gl.glUniformMatrix4fv(self._loc_grid_model, 1, gl.GL_FALSE, ptr)
1825
+ # sky shader needs inverted model view matrix
1826
+ gl.glUseProgram(self._sky_shader.id)
1827
+ inv_ptr = arr_pointer(self._inv_model_matrix)
1828
+ gl.glUniformMatrix4fv(self._loc_sky_inv_model, 1, gl.GL_FALSE, inv_ptr)
1829
+
1830
+ @property
1831
+ def num_tiles(self):
1832
+ return len(self._tile_instances)
1833
+
1834
+ @property
1835
+ def tile_width(self):
1836
+ return self._tile_width
1837
+
1838
+ @property
1839
+ def tile_height(self):
1840
+ return self._tile_height
1841
+
1842
+ @property
1843
+ def num_shapes(self):
1844
+ return len(self._shapes)
1845
+
1846
+ @property
1847
+ def num_instances(self):
1848
+ return self._instance_count
1849
+
1850
+ @property
1851
+ def scaling(self):
1852
+ return self._scaling
1853
+
1854
+ @scaling.setter
1855
+ def scaling(self, scaling):
1856
+ self._scaling = scaling
1857
+ self.update_model_matrix()
1858
+
1859
+ def begin_frame(self, t: float = None):
1860
+ self._last_begin_frame_time = time.time()
1861
+ self.time = t or self.clock_time
1862
+
1863
+ def end_frame(self):
1864
+ self._last_end_frame_time = time.time()
1865
+ if self._add_shape_instances:
1866
+ self.allocate_shape_instances()
1867
+ if self._update_shape_instances:
1868
+ self.update_shape_instances()
1869
+ self.update()
1870
+ while self.paused and self.is_running():
1871
+ self.update()
1872
+
1873
+ def update(self):
1874
+ self.clock_time = time.time() - self._start_time
1875
+ update_duration = self.clock_time - self._last_time
1876
+ frame_duration = self._last_end_frame_time - self._last_begin_frame_time
1877
+ self._last_time = self.clock_time
1878
+ self._frame_speed = update_duration * 100.0
1879
+
1880
+ if not self.headless:
1881
+ self.app.platform_event_loop.step(self._frame_dt * 1e-3)
1882
+
1883
+ if not self.skip_rendering:
1884
+ self._skip_frame_counter += 1
1885
+ if self._skip_frame_counter > 100:
1886
+ self._skip_frame_counter = 0
1887
+
1888
+ if frame_duration > 0.0:
1889
+ if self._fps_update is None:
1890
+ self._fps_update = 1.0 / frame_duration
1891
+ else:
1892
+ update = 1.0 / frame_duration
1893
+ self._fps_update = (1.0 - self._fps_alpha) * self._fps_update + self._fps_alpha * update
1894
+ if update_duration > 0.0:
1895
+ if self._fps_render is None:
1896
+ self._fps_render = 1.0 / update_duration
1897
+ else:
1898
+ update = 1.0 / update_duration
1899
+ self._fps_render = (1.0 - self._fps_alpha) * self._fps_render + self._fps_alpha * update
1900
+
1901
+ if not self.headless:
1902
+ self.app.event_loop._redraw_windows(self._frame_dt * 1e-3)
1903
+ else:
1904
+ self._draw()
1905
+
1906
+ def _draw(self):
1907
+ gl = OpenGLRenderer.gl
1908
+
1909
+ self._switch_context()
1910
+
1911
+ if not self.headless:
1912
+ # catch key hold events
1913
+ self._process_inputs()
1914
+
1915
+ if self.enable_backface_culling:
1916
+ gl.glEnable(gl.GL_CULL_FACE)
1917
+ else:
1918
+ gl.glDisable(gl.GL_CULL_FACE)
1919
+
1920
+ if self._frame_fbo is not None:
1921
+ gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, self._frame_fbo)
1922
+
1923
+ gl.glClearColor(*self.background_color, 1)
1924
+ gl.glClear(gl.GL_COLOR_BUFFER_BIT | gl.GL_DEPTH_BUFFER_BIT)
1925
+ gl.glBindVertexArray(0)
1926
+
1927
+ if not self._tiled_rendering:
1928
+ if self.draw_grid:
1929
+ self._draw_grid()
1930
+
1931
+ if self.draw_sky:
1932
+ self._draw_sky()
1933
+
1934
+ view_mat_ptr = arr_pointer(self._view_matrix)
1935
+ projection_mat_ptr = arr_pointer(self._projection_matrix)
1936
+ gl.glUseProgram(self._shape_shader.id)
1937
+ gl.glUniformMatrix4fv(self._loc_shape_view, 1, gl.GL_FALSE, view_mat_ptr)
1938
+ gl.glUniform3f(self._loc_shape_view_pos, *self._camera_pos)
1939
+ gl.glUniformMatrix4fv(self._loc_shape_view, 1, gl.GL_FALSE, view_mat_ptr)
1940
+ gl.glUniformMatrix4fv(self._loc_shape_projection, 1, gl.GL_FALSE, projection_mat_ptr)
1941
+
1942
+ if self.render_wireframe:
1943
+ gl.glPolygonMode(gl.GL_FRONT_AND_BACK, gl.GL_LINE)
1944
+
1945
+ if self._tiled_rendering:
1946
+ self._render_scene_tiled()
1947
+ else:
1948
+ self._render_scene()
1949
+
1950
+ for cb in self.render_3d_callbacks:
1951
+ cb()
1952
+
1953
+ gl.glPolygonMode(gl.GL_FRONT_AND_BACK, gl.GL_FILL)
1954
+
1955
+ gl.glBindBuffer(gl.GL_PIXEL_UNPACK_BUFFER, 0)
1956
+ gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, 0)
1957
+ gl.glClear(gl.GL_COLOR_BUFFER_BIT | gl.GL_DEPTH_BUFFER_BIT)
1958
+ gl.glViewport(0, 0, self.screen_width, self.screen_height)
1959
+
1960
+ # render frame buffer texture to screen
1961
+ if self._frame_fbo is not None:
1962
+ if self.render_depth:
1963
+ with self._frame_depth_shader:
1964
+ gl.glActiveTexture(gl.GL_TEXTURE0)
1965
+ gl.glBindTexture(gl.GL_TEXTURE_2D, self._frame_depth_texture)
1966
+ gl.glUniform1i(self._frame_loc_depth_texture, 0)
1967
+
1968
+ gl.glBindVertexArray(self._frame_vao)
1969
+ gl.glDrawElements(gl.GL_TRIANGLES, len(self._frame_indices), gl.GL_UNSIGNED_INT, None)
1970
+ gl.glBindVertexArray(0)
1971
+ gl.glBindTexture(gl.GL_TEXTURE_2D, 0)
1972
+ else:
1973
+ with self._frame_shader:
1974
+ gl.glActiveTexture(gl.GL_TEXTURE0)
1975
+ gl.glBindTexture(gl.GL_TEXTURE_2D, self._frame_texture)
1976
+ gl.glUniform1i(self._frame_loc_texture, 0)
1977
+
1978
+ gl.glBindVertexArray(self._frame_vao)
1979
+ gl.glDrawElements(gl.GL_TRIANGLES, len(self._frame_indices), gl.GL_UNSIGNED_INT, None)
1980
+ gl.glBindVertexArray(0)
1981
+ gl.glBindTexture(gl.GL_TEXTURE_2D, 0)
1982
+
1983
+ # check for OpenGL errors
1984
+ # check_gl_error()
1985
+
1986
+ if self.show_info:
1987
+ gl.glClear(gl.GL_DEPTH_BUFFER_BIT)
1988
+ gl.glBlendFunc(gl.GL_SRC_ALPHA, gl.GL_ONE_MINUS_SRC_ALPHA)
1989
+ gl.glEnable(gl.GL_BLEND)
1990
+
1991
+ text = f"""Sim Time: {self.time:.1f}
1992
+ Update FPS: {self._fps_update:.1f}
1993
+ Render FPS: {self._fps_render:.1f}
1994
+
1995
+ Shapes: {len(self._shapes)}
1996
+ Instances: {len(self._instances)}"""
1997
+ if self.paused:
1998
+ text += "\nPaused (press space to resume)"
1999
+
2000
+ self._info_label.text = text
2001
+ self._info_label.y = self.screen_height - 5
2002
+ self._info_label.draw()
2003
+
2004
+ for cb in self.render_2d_callbacks:
2005
+ cb()
2006
+
2007
+ def _draw_grid(self, is_tiled=False):
2008
+ gl = OpenGLRenderer.gl
2009
+
2010
+ self._switch_context()
2011
+
2012
+ if not is_tiled:
2013
+ gl.glUseProgram(self._grid_shader.id)
2014
+
2015
+ gl.glUniformMatrix4fv(self._loc_grid_view, 1, gl.GL_FALSE, arr_pointer(self._view_matrix))
2016
+ gl.glUniformMatrix4fv(self._loc_grid_projection, 1, gl.GL_FALSE, arr_pointer(self._projection_matrix))
2017
+
2018
+ gl.glBindVertexArray(self._grid_vao)
2019
+ gl.glDrawArrays(gl.GL_LINES, 0, self._grid_vertex_count)
2020
+ gl.glBindVertexArray(0)
2021
+
2022
+ def _draw_sky(self, is_tiled=False):
2023
+ gl = OpenGLRenderer.gl
2024
+
2025
+ self._switch_context()
2026
+
2027
+ if not is_tiled:
2028
+ gl.glUseProgram(self._sky_shader.id)
2029
+
2030
+ gl.glUniformMatrix4fv(self._loc_sky_view, 1, gl.GL_FALSE, arr_pointer(self._view_matrix))
2031
+ gl.glUniformMatrix4fv(self._loc_sky_projection, 1, gl.GL_FALSE, arr_pointer(self._projection_matrix))
2032
+ gl.glUniform3f(self._loc_sky_view_pos, *self._camera_pos)
2033
+
2034
+ gl.glBindVertexArray(self._sky_vao)
2035
+ gl.glDrawElements(gl.GL_TRIANGLES, self._sky_tri_count, gl.GL_UNSIGNED_INT, None)
2036
+ gl.glBindVertexArray(0)
2037
+
2038
+ def _render_scene(self):
2039
+ gl = OpenGLRenderer.gl
2040
+
2041
+ self._switch_context()
2042
+
2043
+ start_instance_idx = 0
2044
+
2045
+ for shape, (vao, _, _, tri_count, _) in self._shape_gl_buffers.items():
2046
+ num_instances = len(self._shape_instances[shape])
2047
+
2048
+ gl.glBindVertexArray(vao)
2049
+ gl.glDrawElementsInstancedBaseInstance(
2050
+ gl.GL_TRIANGLES, tri_count, gl.GL_UNSIGNED_INT, None, num_instances, start_instance_idx
2051
+ )
2052
+
2053
+ start_instance_idx += num_instances
2054
+
2055
+ if self.draw_axis:
2056
+ self._axis_instancer.render()
2057
+
2058
+ for instancer in self._shape_instancers.values():
2059
+ instancer.render()
2060
+
2061
+ gl.glBindVertexArray(0)
2062
+
2063
+ def _render_scene_tiled(self):
2064
+ gl = OpenGLRenderer.gl
2065
+
2066
+ self._switch_context()
2067
+
2068
+ for i, viewport in enumerate(self._tile_viewports):
2069
+ projection_matrix_ptr = arr_pointer(self._tile_projection_matrices[i])
2070
+ view_matrix_ptr = arr_pointer(
2071
+ self._tile_view_matrices[i] if self._tile_view_matrices[i] is not None else self._view_matrix
2072
+ )
2073
+
2074
+ gl.glViewport(*viewport)
2075
+ if self.draw_grid:
2076
+ gl.glUseProgram(self._grid_shader.id)
2077
+ gl.glUniformMatrix4fv(self._loc_grid_projection, 1, gl.GL_FALSE, projection_matrix_ptr)
2078
+ gl.glUniformMatrix4fv(self._loc_grid_view, 1, gl.GL_FALSE, view_matrix_ptr)
2079
+ self._draw_grid(is_tiled=True)
2080
+
2081
+ if self.draw_sky:
2082
+ gl.glUseProgram(self._sky_shader.id)
2083
+ gl.glUniformMatrix4fv(self._loc_sky_projection, 1, gl.GL_FALSE, projection_matrix_ptr)
2084
+ gl.glUniformMatrix4fv(self._loc_sky_view, 1, gl.GL_FALSE, view_matrix_ptr)
2085
+ self._draw_sky(is_tiled=True)
2086
+
2087
+ gl.glUseProgram(self._shape_shader.id)
2088
+ gl.glUniformMatrix4fv(self._loc_shape_projection, 1, gl.GL_FALSE, projection_matrix_ptr)
2089
+ gl.glUniformMatrix4fv(self._loc_shape_view, 1, gl.GL_FALSE, view_matrix_ptr)
2090
+
2091
+ instances = self._tile_instances[i]
2092
+
2093
+ for instance in instances:
2094
+ shape = self._instance_shape[instance]
2095
+
2096
+ vao, _, _, tri_count, _ = self._shape_gl_buffers[shape]
2097
+
2098
+ start_instance_idx = self._inverse_instance_ids[instance]
2099
+
2100
+ gl.glBindVertexArray(vao)
2101
+ gl.glDrawElementsInstancedBaseInstance(
2102
+ gl.GL_TRIANGLES, tri_count, gl.GL_UNSIGNED_INT, None, 1, start_instance_idx
2103
+ )
2104
+
2105
+ if self.draw_axis:
2106
+ self._axis_instancer.render()
2107
+
2108
+ for instancer in self._shape_instancers.values():
2109
+ instancer.render()
2110
+
2111
+ gl.glBindVertexArray(0)
2112
+
2113
+ def _close_callback(self):
2114
+ self.close()
2115
+
2116
+ def _mouse_drag_callback(self, x, y, dx, dy, buttons, modifiers):
2117
+ if not self.enable_mouse_interaction:
2118
+ return
2119
+
2120
+ import pyglet
2121
+ from pyglet.math import Vec3 as PyVec3
2122
+
2123
+ if buttons & pyglet.window.mouse.LEFT:
2124
+ sensitivity = 0.1
2125
+ dx *= sensitivity
2126
+ dy *= sensitivity
2127
+
2128
+ self._yaw += dx
2129
+ self._pitch += dy
2130
+
2131
+ self._pitch = max(min(self._pitch, 89.0), -89.0)
2132
+
2133
+ self._camera_front = PyVec3(
2134
+ np.cos(np.deg2rad(self._yaw)) * np.cos(np.deg2rad(self._pitch)),
2135
+ np.sin(np.deg2rad(self._pitch)),
2136
+ np.sin(np.deg2rad(self._yaw)) * np.cos(np.deg2rad(self._pitch)),
2137
+ ).normalize()
2138
+
2139
+ self.update_view_matrix()
2140
+
2141
+ def _scroll_callback(self, x, y, scroll_x, scroll_y):
2142
+ if not self.enable_mouse_interaction:
2143
+ return
2144
+
2145
+ self.camera_fov -= scroll_y
2146
+ self.camera_fov = max(min(self.camera_fov, 90.0), 15.0)
2147
+ self.update_projection_matrix()
2148
+
2149
+ def _process_inputs(self):
2150
+ import pyglet
2151
+ from pyglet.math import Vec3 as PyVec3
2152
+
2153
+ for cb in self._input_processors:
2154
+ if cb(self._key_handler) == pyglet.event.EVENT_HANDLED:
2155
+ return
2156
+
2157
+ if self._key_handler[pyglet.window.key.W] or self._key_handler[pyglet.window.key.UP]:
2158
+ self._camera_pos += self._camera_front * (self._camera_speed * self._frame_speed)
2159
+ self.update_view_matrix()
2160
+ if self._key_handler[pyglet.window.key.S] or self._key_handler[pyglet.window.key.DOWN]:
2161
+ self._camera_pos -= self._camera_front * (self._camera_speed * self._frame_speed)
2162
+ self.update_view_matrix()
2163
+ if self._key_handler[pyglet.window.key.A] or self._key_handler[pyglet.window.key.LEFT]:
2164
+ camera_side = PyVec3.cross(self._camera_front, self._camera_up).normalize()
2165
+ self._camera_pos -= camera_side * (self._camera_speed * self._frame_speed)
2166
+ self.update_view_matrix()
2167
+ if self._key_handler[pyglet.window.key.D] or self._key_handler[pyglet.window.key.RIGHT]:
2168
+ camera_side = PyVec3.cross(self._camera_front, self._camera_up).normalize()
2169
+ self._camera_pos += camera_side * (self._camera_speed * self._frame_speed)
2170
+ self.update_view_matrix()
2171
+
2172
+ def register_input_processor(self, callback):
2173
+ self._input_processors.append(callback)
2174
+
2175
+ def _key_press_callback(self, symbol, modifiers):
2176
+ import pyglet
2177
+
2178
+ if not self.enable_keyboard_interaction:
2179
+ return
2180
+
2181
+ for cb in self._key_callbacks:
2182
+ if cb(symbol, modifiers) == pyglet.event.EVENT_HANDLED:
2183
+ return pyglet.event.EVENT_HANDLED
2184
+
2185
+ if symbol == pyglet.window.key.ESCAPE:
2186
+ self.close()
2187
+ if symbol == pyglet.window.key.SPACE:
2188
+ self.paused = not self.paused
2189
+ if symbol == pyglet.window.key.TAB:
2190
+ self.skip_rendering = not self.skip_rendering
2191
+ if symbol == pyglet.window.key.C:
2192
+ self.draw_axis = not self.draw_axis
2193
+ if symbol == pyglet.window.key.G:
2194
+ self.draw_grid = not self.draw_grid
2195
+ if symbol == pyglet.window.key.I:
2196
+ self.show_info = not self.show_info
2197
+ if symbol == pyglet.window.key.X:
2198
+ self.render_wireframe = not self.render_wireframe
2199
+ if symbol == pyglet.window.key.T:
2200
+ self.render_depth = not self.render_depth
2201
+ if symbol == pyglet.window.key.B:
2202
+ self.enable_backface_culling = not self.enable_backface_culling
2203
+
2204
+ def register_key_press_callback(self, callback):
2205
+ self._key_callbacks.append(callback)
2206
+
2207
+ def _window_resize_callback(self, width, height):
2208
+ self._first_mouse = True
2209
+ self.screen_width, self.screen_height = self.window.get_framebuffer_size()
2210
+ self.update_projection_matrix()
2211
+ self._setup_framebuffer()
2212
+
2213
+ def register_shape(self, geo_hash, vertices, indices, color1=None, color2=None):
2214
+ gl = OpenGLRenderer.gl
2215
+
2216
+ self._switch_context()
2217
+
2218
+ shape = len(self._shapes)
2219
+ if color1 is None:
2220
+ color1 = tab10_color_map(len(self._shape_geo_hash))
2221
+ if color2 is None:
2222
+ color2 = np.clip(np.array(color1) + 0.25, 0.0, 1.0)
2223
+ # TODO check if we actually need to store the shape data
2224
+ self._shapes.append((vertices, indices, color1, color2, geo_hash))
2225
+ self._shape_geo_hash[geo_hash] = shape
2226
+
2227
+ gl.glUseProgram(self._shape_shader.id)
2228
+
2229
+ # Create VAO, VBO, and EBO
2230
+ vao = gl.GLuint()
2231
+ gl.glGenVertexArrays(1, vao)
2232
+ gl.glBindVertexArray(vao)
2233
+
2234
+ vbo = gl.GLuint()
2235
+ gl.glGenBuffers(1, vbo)
2236
+ gl.glBindBuffer(gl.GL_ARRAY_BUFFER, vbo)
2237
+ gl.glBufferData(gl.GL_ARRAY_BUFFER, vertices.nbytes, vertices.ctypes.data, gl.GL_STATIC_DRAW)
2238
+
2239
+ vertex_cuda_buffer = wp.RegisteredGLBuffer(int(vbo.value), self._device)
2240
+
2241
+ ebo = gl.GLuint()
2242
+ gl.glGenBuffers(1, ebo)
2243
+ gl.glBindBuffer(gl.GL_ELEMENT_ARRAY_BUFFER, ebo)
2244
+ gl.glBufferData(gl.GL_ELEMENT_ARRAY_BUFFER, indices.nbytes, indices.ctypes.data, gl.GL_STATIC_DRAW)
2245
+
2246
+ # Set up vertex attributes
2247
+ vertex_stride = vertices.shape[1] * vertices.itemsize
2248
+ # positions
2249
+ gl.glVertexAttribPointer(0, 3, gl.GL_FLOAT, gl.GL_FALSE, vertex_stride, ctypes.c_void_p(0))
2250
+ gl.glEnableVertexAttribArray(0)
2251
+ # normals
2252
+ gl.glVertexAttribPointer(1, 3, gl.GL_FLOAT, gl.GL_FALSE, vertex_stride, ctypes.c_void_p(3 * vertices.itemsize))
2253
+ gl.glEnableVertexAttribArray(1)
2254
+ # uv coordinates
2255
+ gl.glVertexAttribPointer(2, 2, gl.GL_FLOAT, gl.GL_FALSE, vertex_stride, ctypes.c_void_p(6 * vertices.itemsize))
2256
+ gl.glEnableVertexAttribArray(2)
2257
+
2258
+ gl.glBindVertexArray(0)
2259
+
2260
+ self._shape_gl_buffers[shape] = (vao, vbo, ebo, len(indices), vertex_cuda_buffer)
2261
+
2262
+ return shape
2263
+
2264
+ def deregister_shape(self, shape):
2265
+ gl = OpenGLRenderer.gl
2266
+
2267
+ self._switch_context()
2268
+
2269
+ if shape not in self._shape_gl_buffers:
2270
+ return
2271
+
2272
+ vao, vbo, ebo, _, vertex_cuda_buffer = self._shape_gl_buffers[shape]
2273
+ try:
2274
+ gl.glDeleteVertexArrays(1, vao)
2275
+ gl.glDeleteBuffers(1, vbo)
2276
+ gl.glDeleteBuffers(1, ebo)
2277
+ except gl.GLException:
2278
+ pass
2279
+
2280
+ _, _, _, _, geo_hash = self._shapes[shape]
2281
+ del self._shape_geo_hash[geo_hash]
2282
+ del self._shape_gl_buffers[shape]
2283
+ self._shapes.pop(shape)
2284
+
2285
+ def add_shape_instance(
2286
+ self,
2287
+ name: str,
2288
+ shape: int,
2289
+ body,
2290
+ pos,
2291
+ rot,
2292
+ scale=(1.0, 1.0, 1.0),
2293
+ color1=None,
2294
+ color2=None,
2295
+ custom_index: int = -1,
2296
+ visible: bool = True,
2297
+ ):
2298
+ if color1 is None:
2299
+ color1 = self._shapes[shape][2]
2300
+ if color2 is None:
2301
+ color2 = self._shapes[shape][3]
2302
+ instance = len(self._instances)
2303
+ self._shape_instances[shape].append(instance)
2304
+ body = self._resolve_body_id(body)
2305
+ self._instances[name] = (instance, body, shape, [*pos, *rot], scale, color1, color2, visible)
2306
+ self._instance_shape[instance] = shape
2307
+ self._instance_custom_ids[instance] = custom_index
2308
+ self._add_shape_instances = True
2309
+ self._instance_count = len(self._instances)
2310
+ return instance
2311
+
2312
+ def remove_shape_instance(self, name: str):
2313
+ if name not in self._instances:
2314
+ return
2315
+
2316
+ instance, _, shape, _, _, _, _, _ = self._instances[name]
2317
+
2318
+ self._shape_instances[shape].remove(instance)
2319
+ self._instance_count = len(self._instances)
2320
+ self._add_shape_instances = self._instance_count > 0
2321
+ del self._instance_shape[instance]
2322
+ del self._instance_custom_ids[instance]
2323
+ del self._instances[name]
2324
+
2325
+ def update_instance_colors(self):
2326
+ gl = OpenGLRenderer.gl
2327
+
2328
+ self._switch_context()
2329
+
2330
+ colors1, colors2 = [], []
2331
+ all_instances = list(self._instances.values())
2332
+ for _shape, instances in self._shape_instances.items():
2333
+ for i in instances:
2334
+ if i >= len(all_instances):
2335
+ continue
2336
+ instance = all_instances[i]
2337
+ colors1.append(instance[5])
2338
+ colors2.append(instance[6])
2339
+ colors1 = np.array(colors1, dtype=np.float32)
2340
+ colors2 = np.array(colors2, dtype=np.float32)
2341
+
2342
+ gl.glBindBuffer(gl.GL_ARRAY_BUFFER, self._instance_color1_buffer)
2343
+ gl.glBufferData(gl.GL_ARRAY_BUFFER, colors1.nbytes, colors1.ctypes.data, gl.GL_STATIC_DRAW)
2344
+
2345
+ gl.glBindBuffer(gl.GL_ARRAY_BUFFER, self._instance_color2_buffer)
2346
+ gl.glBufferData(gl.GL_ARRAY_BUFFER, colors2.nbytes, colors2.ctypes.data, gl.GL_STATIC_DRAW)
2347
+
2348
+ def allocate_shape_instances(self):
2349
+ gl = OpenGLRenderer.gl
2350
+
2351
+ self._switch_context()
2352
+
2353
+ self._add_shape_instances = False
2354
+ self._wp_instance_transforms = wp.array(
2355
+ [instance[3] for instance in self._instances.values()], dtype=wp.transform, device=self._device
2356
+ )
2357
+ self._wp_instance_scalings = wp.array(
2358
+ [instance[4] for instance in self._instances.values()], dtype=wp.vec3, device=self._device
2359
+ )
2360
+ self._wp_instance_bodies = wp.array(
2361
+ [instance[1] for instance in self._instances.values()], dtype=wp.int32, device=self._device
2362
+ )
2363
+
2364
+ gl.glUseProgram(self._shape_shader.id)
2365
+ if self._instance_transform_gl_buffer is not None:
2366
+ gl.glDeleteBuffers(1, self._instance_transform_gl_buffer)
2367
+ gl.glDeleteBuffers(1, self._instance_color1_buffer)
2368
+ gl.glDeleteBuffers(1, self._instance_color2_buffer)
2369
+
2370
+ # create instance buffer and bind it as an instanced array
2371
+ self._instance_transform_gl_buffer = gl.GLuint()
2372
+ gl.glGenBuffers(1, self._instance_transform_gl_buffer)
2373
+ gl.glBindBuffer(gl.GL_ARRAY_BUFFER, self._instance_transform_gl_buffer)
2374
+
2375
+ transforms = np.tile(np.diag(np.ones(4, dtype=np.float32)), (len(self._instances), 1, 1))
2376
+ gl.glBufferData(gl.GL_ARRAY_BUFFER, transforms.nbytes, transforms.ctypes.data, gl.GL_DYNAMIC_DRAW)
2377
+
2378
+ # create CUDA buffer for instance transforms
2379
+ self._instance_transform_cuda_buffer = wp.RegisteredGLBuffer(
2380
+ int(self._instance_transform_gl_buffer.value), self._device
2381
+ )
2382
+
2383
+ # create color buffers
2384
+ self._instance_color1_buffer = gl.GLuint()
2385
+ gl.glGenBuffers(1, self._instance_color1_buffer)
2386
+ self._instance_color2_buffer = gl.GLuint()
2387
+ gl.glGenBuffers(1, self._instance_color2_buffer)
2388
+
2389
+ self.update_instance_colors()
2390
+
2391
+ # set up instance attribute pointers
2392
+ matrix_size = transforms[0].nbytes
2393
+
2394
+ instance_ids = []
2395
+ instance_custom_ids = []
2396
+ instance_visible = []
2397
+ instances = list(self._instances.values())
2398
+ inverse_instance_ids = {}
2399
+ instance_count = 0
2400
+ colors_size = np.zeros(3, dtype=np.float32).nbytes
2401
+ for shape, (vao, _vbo, _ebo, _tri_count, _vertex_cuda_buffer) in self._shape_gl_buffers.items():
2402
+ gl.glBindVertexArray(vao)
2403
+
2404
+ gl.glBindBuffer(gl.GL_ARRAY_BUFFER, self._instance_transform_gl_buffer)
2405
+
2406
+ # we can only send vec4s to the shader, so we need to split the instance transforms matrix into its column vectors
2407
+ for i in range(4):
2408
+ gl.glVertexAttribPointer(
2409
+ 3 + i, 4, gl.GL_FLOAT, gl.GL_FALSE, matrix_size, ctypes.c_void_p(i * matrix_size // 4)
2410
+ )
2411
+ gl.glEnableVertexAttribArray(3 + i)
2412
+ gl.glVertexAttribDivisor(3 + i, 1)
2413
+
2414
+ gl.glBindBuffer(gl.GL_ARRAY_BUFFER, self._instance_color1_buffer)
2415
+ gl.glVertexAttribPointer(7, 3, gl.GL_FLOAT, gl.GL_FALSE, colors_size, ctypes.c_void_p(0))
2416
+ gl.glEnableVertexAttribArray(7)
2417
+ gl.glVertexAttribDivisor(7, 1)
2418
+
2419
+ gl.glBindBuffer(gl.GL_ARRAY_BUFFER, self._instance_color2_buffer)
2420
+ gl.glVertexAttribPointer(8, 3, gl.GL_FLOAT, gl.GL_FALSE, colors_size, ctypes.c_void_p(0))
2421
+ gl.glEnableVertexAttribArray(8)
2422
+ gl.glVertexAttribDivisor(8, 1)
2423
+
2424
+ instance_ids.extend(self._shape_instances[shape])
2425
+ for i in self._shape_instances[shape]:
2426
+ inverse_instance_ids[i] = instance_count
2427
+ instance_count += 1
2428
+ instance_custom_ids.append(self._instance_custom_ids[i])
2429
+ instance_visible.append(instances[i][7])
2430
+
2431
+ # trigger update to the instance transforms
2432
+ self._update_shape_instances = True
2433
+
2434
+ self._wp_instance_ids = wp.array(instance_ids, dtype=wp.int32, device=self._device)
2435
+ self._wp_instance_custom_ids = wp.array(instance_custom_ids, dtype=wp.int32, device=self._device)
2436
+ self._np_instance_visible = np.array(instance_visible)
2437
+ self._instance_ids = instance_ids
2438
+ self._inverse_instance_ids = inverse_instance_ids
2439
+
2440
+ gl.glBindVertexArray(0)
2441
+
2442
+ def update_shape_instance(self, name, pos=None, rot=None, color1=None, color2=None, visible=None):
2443
+ """Update the instance transform of the shape
2444
+
2445
+ Args:
2446
+ name: The name of the shape
2447
+ pos: The position of the shape
2448
+ rot: The rotation of the shape
2449
+ color1: The first color of the checker pattern
2450
+ color2: The second color of the checker pattern
2451
+ visible: Whether the shape is visible
2452
+ """
2453
+ gl = OpenGLRenderer.gl
2454
+
2455
+ self._switch_context()
2456
+
2457
+ if name in self._instances:
2458
+ i, body, shape, tf, scale, old_color1, old_color2, v = self._instances[name]
2459
+ if visible is None:
2460
+ visible = v
2461
+ new_tf = np.copy(tf)
2462
+ if pos is not None:
2463
+ new_tf[:3] = pos
2464
+ if rot is not None:
2465
+ new_tf[3:] = rot
2466
+ self._instances[name] = (
2467
+ i,
2468
+ body,
2469
+ shape,
2470
+ new_tf,
2471
+ scale,
2472
+ old_color1 if color1 is None else color1,
2473
+ old_color2 if color2 is None else color2,
2474
+ visible,
2475
+ )
2476
+ self._update_shape_instances = True
2477
+ if color1 is not None or color2 is not None:
2478
+ vao, vbo, ebo, tri_count, vertex_cuda_buffer = self._shape_gl_buffers[shape]
2479
+ gl.glBindVertexArray(vao)
2480
+ self.update_instance_colors()
2481
+ gl.glBindVertexArray(0)
2482
+ return True
2483
+ return False
2484
+
2485
+ def update_shape_instances(self):
2486
+ with self._shape_shader:
2487
+ self._update_shape_instances = False
2488
+ self._wp_instance_transforms = wp.array(
2489
+ [instance[3] for instance in self._instances.values()], dtype=wp.transform, device=self._device
2490
+ )
2491
+ self.update_body_transforms(None)
2492
+
2493
+ def update_body_transforms(self, body_tf: wp.array):
2494
+ if self._instance_transform_cuda_buffer is None:
2495
+ return
2496
+
2497
+ body_q = None
2498
+ if body_tf is not None:
2499
+ if body_tf.device.is_cuda:
2500
+ body_q = body_tf
2501
+ else:
2502
+ body_q = body_tf.to(self._device)
2503
+
2504
+ vbo_transforms = self._instance_transform_cuda_buffer.map(dtype=wp.mat44, shape=(self._instance_count,))
2505
+
2506
+ wp.launch(
2507
+ update_vbo_transforms,
2508
+ dim=self._instance_count,
2509
+ inputs=[
2510
+ self._wp_instance_ids,
2511
+ self._wp_instance_bodies,
2512
+ self._wp_instance_transforms,
2513
+ self._wp_instance_scalings,
2514
+ body_q,
2515
+ ],
2516
+ outputs=[
2517
+ vbo_transforms,
2518
+ ],
2519
+ device=self._device,
2520
+ record_tape=False,
2521
+ )
2522
+
2523
+ self._instance_transform_cuda_buffer.unmap()
2524
+
2525
+ def register_body(self, name):
2526
+ # register body name and return its ID
2527
+ if name not in self._body_name:
2528
+ self._body_name[name] = len(self._body_name)
2529
+ return self._body_name[name]
2530
+
2531
+ def _resolve_body_id(self, body):
2532
+ if body is None:
2533
+ return -1
2534
+ if isinstance(body, int):
2535
+ return body
2536
+ return self._body_name[body]
2537
+
2538
+ def is_running(self):
2539
+ return not self.app.event_loop.has_exit
2540
+
2541
+ def save(self):
2542
+ # save just keeps the window open to allow the user to interact with the scene
2543
+ while not self.app.event_loop.has_exit:
2544
+ self.update()
2545
+ if self.app.event_loop.has_exit:
2546
+ self.clear()
2547
+ self.app.event_loop.exit()
2548
+
2549
+ def get_pixels(self, target_image: wp.array, split_up_tiles=True, mode="rgb", use_uint8=False):
2550
+ """
2551
+ Read the pixels from the frame buffer (RGB or depth are supported) into the given array.
2552
+
2553
+ If `split_up_tiles` is False, array must be of shape (screen_height, screen_width, 3) for RGB mode or
2554
+ (screen_height, screen_width, 1) for depth mode.
2555
+ If `split_up_tiles` is True, the pixels will be split up into tiles (see :attr:`tile_width` and :attr:`tile_height` for dimensions):
2556
+ array must be of shape (num_tiles, tile_height, tile_width, 3) for RGB mode or (num_tiles, tile_height, tile_width, 1) for depth mode.
2557
+
2558
+ Args:
2559
+ target_image (array): The array to read the pixels into. Must have float32 as dtype and be on a CUDA device.
2560
+ split_up_tiles (bool): Whether to split up the viewport into tiles, see :meth:`setup_tiled_rendering`.
2561
+ mode (str): can be either "rgb" or "depth"
2562
+ use_uint8 (bool): Whether to use uint8 as dtype in RGB mode for the target_image array and return values in the range [0, 255]. Otherwise, float32 is assumed as dtype with values in the range [0, 1].
2563
+
2564
+ Returns:
2565
+ bool: Whether the pixels were successfully read.
2566
+ """
2567
+ gl = OpenGLRenderer.gl
2568
+
2569
+ self._switch_context()
2570
+
2571
+ channels = 3 if mode == "rgb" else 1
2572
+
2573
+ if split_up_tiles:
2574
+ assert self._tile_width is not None and self._tile_height is not None, (
2575
+ "Tile width and height are not set, tiles must all have the same size"
2576
+ )
2577
+ assert all(vp[2] == self._tile_width for vp in self._tile_viewports), (
2578
+ "Tile widths do not all equal global tile_width, use `get_tile_pixels` instead to retrieve pixels for a single tile"
2579
+ )
2580
+ assert all(vp[3] == self._tile_height for vp in self._tile_viewports), (
2581
+ "Tile heights do not all equal global tile_height, use `get_tile_pixels` instead to retrieve pixels for a single tile"
2582
+ )
2583
+ assert target_image.shape == (
2584
+ self.num_tiles,
2585
+ self._tile_height,
2586
+ self._tile_width,
2587
+ channels,
2588
+ ), (
2589
+ f"Shape of `target_image` array does not match {self.num_tiles} x {self._tile_height} x {self._tile_width} x {channels}"
2590
+ )
2591
+ else:
2592
+ assert target_image.shape == (
2593
+ self.screen_height,
2594
+ self.screen_width,
2595
+ channels,
2596
+ ), f"Shape of `target_image` array does not match {self.screen_height} x {self.screen_width} x {channels}"
2597
+
2598
+ gl.glBindBuffer(gl.GL_PIXEL_PACK_BUFFER, self._frame_pbo)
2599
+ if mode == "rgb":
2600
+ gl.glBindTexture(gl.GL_TEXTURE_2D, self._frame_texture)
2601
+ if mode == "depth":
2602
+ gl.glBindTexture(gl.GL_TEXTURE_2D, self._frame_depth_texture)
2603
+ try:
2604
+ # read screen texture into PBO
2605
+ if mode == "rgb":
2606
+ gl.glGetTexImage(gl.GL_TEXTURE_2D, 0, gl.GL_RGB, gl.GL_UNSIGNED_BYTE, ctypes.c_void_p(0))
2607
+ elif mode == "depth":
2608
+ gl.glGetTexImage(gl.GL_TEXTURE_2D, 0, gl.GL_DEPTH_COMPONENT, gl.GL_FLOAT, ctypes.c_void_p(0))
2609
+ except gl.GLException:
2610
+ # this can happen if the window is closed/being moved to a different display
2611
+ gl.glBindTexture(gl.GL_TEXTURE_2D, 0)
2612
+ gl.glBindBuffer(gl.GL_PIXEL_PACK_BUFFER, 0)
2613
+ return False
2614
+ gl.glBindTexture(gl.GL_TEXTURE_2D, 0)
2615
+ gl.glBindBuffer(gl.GL_PIXEL_PACK_BUFFER, 0)
2616
+
2617
+ pbo_buffer = wp.RegisteredGLBuffer(int(self._frame_pbo.value), self._device, wp.RegisteredGLBuffer.READ_ONLY)
2618
+ screen_size = self.screen_height * self.screen_width
2619
+ if mode == "rgb":
2620
+ img = pbo_buffer.map(dtype=wp.uint8, shape=(screen_size * channels))
2621
+ elif mode == "depth":
2622
+ img = pbo_buffer.map(dtype=wp.float32, shape=(screen_size * channels))
2623
+ img = img.to(target_image.device)
2624
+ if split_up_tiles:
2625
+ positions = wp.array(self._tile_viewports, ndim=2, dtype=wp.int32, device=target_image.device)
2626
+ if mode == "rgb":
2627
+ wp.launch(
2628
+ copy_rgb_frame_tiles_uint8 if use_uint8 else copy_rgb_frame_tiles,
2629
+ dim=(self.num_tiles, self._tile_width, self._tile_height),
2630
+ inputs=[img, positions, self.screen_width, self.screen_height, self._tile_height],
2631
+ outputs=[target_image],
2632
+ device=target_image.device,
2633
+ )
2634
+ elif mode == "depth":
2635
+ wp.launch(
2636
+ copy_depth_frame_tiles,
2637
+ dim=(self.num_tiles, self._tile_width, self._tile_height),
2638
+ inputs=[
2639
+ img,
2640
+ positions,
2641
+ self.screen_width,
2642
+ self.screen_height,
2643
+ self._tile_height,
2644
+ self.camera_near_plane,
2645
+ self.camera_far_plane,
2646
+ ],
2647
+ outputs=[target_image],
2648
+ device=target_image.device,
2649
+ )
2650
+ else:
2651
+ if mode == "rgb":
2652
+ wp.launch(
2653
+ copy_rgb_frame_uint8 if use_uint8 else copy_rgb_frame,
2654
+ dim=(self.screen_width, self.screen_height),
2655
+ inputs=[img, self.screen_width, self.screen_height],
2656
+ outputs=[target_image],
2657
+ device=target_image.device,
2658
+ )
2659
+ elif mode == "depth":
2660
+ wp.launch(
2661
+ copy_depth_frame,
2662
+ dim=(self.screen_width, self.screen_height),
2663
+ inputs=[img, self.screen_width, self.screen_height, self.camera_near_plane, self.camera_far_plane],
2664
+ outputs=[target_image],
2665
+ device=target_image.device,
2666
+ )
2667
+ pbo_buffer.unmap()
2668
+ return True
2669
+
2670
+ # def create_image_texture(self, file_path):
2671
+ # from PIL import Image
2672
+ # img = Image.open(file_path)
2673
+ # img_data = np.array(list(img.getdata()), np.uint8)
2674
+ # texture = glGenTextures(1)
2675
+ # glBindTexture(GL_TEXTURE_2D, texture)
2676
+ # glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, img.width, img.height, 0, GL_RGB, GL_UNSIGNED_BYTE, img_data)
2677
+ # glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_REPEAT)
2678
+ # glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_REPEAT)
2679
+ # glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)
2680
+ # glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
2681
+ # return texture
2682
+
2683
+ # def create_check_texture(self, color1=(0, 0.5, 1.0), color2=None, width=default_texture_size, height=default_texture_size):
2684
+ # if width == 1 and height == 1:
2685
+ # pixels = np.array([np.array(color1)*255], dtype=np.uint8)
2686
+ # else:
2687
+ # pixels = np.zeros((width, height, 3), dtype=np.uint8)
2688
+ # half_w = width // 2
2689
+ # half_h = height // 2
2690
+ # color1 = np.array(np.array(color1)*255, dtype=np.uint8)
2691
+ # pixels[0:half_w, 0:half_h] = color1
2692
+ # pixels[half_w:width, half_h:height] = color1
2693
+ # if color2 is None:
2694
+ # color2 = np.array(np.clip(np.array(color1, dtype=np.float32) + 50, 0, 255), dtype=np.uint8)
2695
+ # else:
2696
+ # color2 = np.array(np.array(color2)*255, dtype=np.uint8)
2697
+ # pixels[half_w:width, 0:half_h] = color2
2698
+ # pixels[0:half_w, half_h:height] = color2
2699
+ # texture = glGenTextures(1)
2700
+ # glBindTexture(GL_TEXTURE_2D, texture)
2701
+ # glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, width, height, 0, GL_RGB, GL_UNSIGNED_BYTE, pixels.flatten())
2702
+ # glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_REPEAT)
2703
+ # glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_REPEAT)
2704
+ # glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)
2705
+ # glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
2706
+ # return texture
2707
+
2708
+ def render_plane(
2709
+ self,
2710
+ name: str,
2711
+ pos: tuple,
2712
+ rot: tuple,
2713
+ width: float,
2714
+ length: float,
2715
+ color: tuple = (1.0, 1.0, 1.0),
2716
+ color2=None,
2717
+ parent_body: str = None,
2718
+ is_template: bool = False,
2719
+ u_scaling=1.0,
2720
+ v_scaling=1.0,
2721
+ ):
2722
+ """Add a plane for visualization
2723
+
2724
+ Args:
2725
+ name: The name of the plane
2726
+ pos: The position of the plane
2727
+ rot: The rotation of the plane
2728
+ width: The width of the plane
2729
+ length: The length of the plane
2730
+ color: The color of the plane
2731
+ texture: The texture of the plane (optional)
2732
+ """
2733
+ geo_hash = hash(("plane", width, length))
2734
+ if geo_hash in self._shape_geo_hash:
2735
+ shape = self._shape_geo_hash[geo_hash]
2736
+ if self.update_shape_instance(name, pos, rot):
2737
+ return shape
2738
+ else:
2739
+ faces = np.array([0, 1, 2, 2, 3, 0], dtype=np.uint32)
2740
+ normal = (0.0, 1.0, 0.0)
2741
+ width = width if width > 0.0 else 100.0
2742
+ length = length if length > 0.0 else 100.0
2743
+ aspect = width / length
2744
+ u = width * aspect * u_scaling
2745
+ v = length * v_scaling
2746
+ gfx_vertices = np.array(
2747
+ [
2748
+ [-width, 0.0, -length, *normal, 0.0, 0.0],
2749
+ [-width, 0.0, length, *normal, 0.0, v],
2750
+ [width, 0.0, length, *normal, u, v],
2751
+ [width, 0.0, -length, *normal, u, 0.0],
2752
+ ],
2753
+ dtype=np.float32,
2754
+ )
2755
+ shape = self.register_shape(geo_hash, gfx_vertices, faces, color1=color, color2=color2)
2756
+ if not is_template:
2757
+ body = self._resolve_body_id(parent_body)
2758
+ self.add_shape_instance(name, shape, body, pos, rot)
2759
+ return shape
2760
+
2761
+ def render_ground(self, size: float = 1000.0, plane=None):
2762
+ """Add a ground plane for visualization
2763
+
2764
+ Args:
2765
+ size: The size of the ground plane
2766
+ """
2767
+ color1 = (200 / 255, 200 / 255, 200 / 255)
2768
+ color2 = (150 / 255, 150 / 255, 150 / 255)
2769
+ sqh = np.sqrt(0.5)
2770
+ if self._camera_axis == 0:
2771
+ q = (0.0, 0.0, -sqh, sqh)
2772
+ elif self._camera_axis == 1:
2773
+ q = (0.0, 0.0, 0.0, 1.0)
2774
+ elif self._camera_axis == 2:
2775
+ q = (sqh, 0.0, 0.0, sqh)
2776
+ pos = (0.0, 0.0, 0.0)
2777
+ if plane is not None:
2778
+ normal = np.array(plane[:3])
2779
+ normal /= np.linalg.norm(normal)
2780
+ pos = plane[3] * normal
2781
+ if np.allclose(normal, (0.0, 1.0, 0.0)):
2782
+ # no rotation necessary
2783
+ q = (0.0, 0.0, 0.0, 1.0)
2784
+ else:
2785
+ c = np.cross(normal, (0.0, 1.0, 0.0))
2786
+ angle = np.arcsin(np.linalg.norm(c))
2787
+ axis = np.abs(c) / np.linalg.norm(c)
2788
+ q = wp.quat_from_axis_angle(axis, angle)
2789
+ return self.render_plane(
2790
+ "ground",
2791
+ pos,
2792
+ q,
2793
+ size,
2794
+ size,
2795
+ color1,
2796
+ color2=color2,
2797
+ u_scaling=1.0,
2798
+ v_scaling=1.0,
2799
+ )
2800
+
2801
+ def render_sphere(
2802
+ self,
2803
+ name: str,
2804
+ pos: tuple,
2805
+ rot: tuple,
2806
+ radius: float,
2807
+ parent_body: str = None,
2808
+ is_template: bool = False,
2809
+ color=None,
2810
+ ):
2811
+ """Add a sphere for visualization
2812
+
2813
+ Args:
2814
+ pos: The position of the sphere
2815
+ radius: The radius of the sphere
2816
+ name: A name for the USD prim on the stage
2817
+ color: The color of the sphere
2818
+ """
2819
+ geo_hash = hash(("sphere", radius))
2820
+ if geo_hash in self._shape_geo_hash:
2821
+ shape = self._shape_geo_hash[geo_hash]
2822
+ if self.update_shape_instance(name, pos, rot, color1=color, color2=color):
2823
+ return shape
2824
+ else:
2825
+ vertices, indices = self._create_sphere_mesh(radius)
2826
+ shape = self.register_shape(geo_hash, vertices, indices, color1=color, color2=color)
2827
+ if not is_template:
2828
+ body = self._resolve_body_id(parent_body)
2829
+ self.add_shape_instance(name, shape, body, pos, rot, color1=color, color2=color)
2830
+ return shape
2831
+
2832
+ def render_capsule(
2833
+ self,
2834
+ name: str,
2835
+ pos: tuple,
2836
+ rot: tuple,
2837
+ radius: float,
2838
+ half_height: float,
2839
+ parent_body: str = None,
2840
+ is_template: bool = False,
2841
+ up_axis: int = 1,
2842
+ color: tuple = None,
2843
+ ):
2844
+ """Add a capsule for visualization
2845
+
2846
+ Args:
2847
+ pos: The position of the capsule
2848
+ radius: The radius of the capsule
2849
+ half_height: The half height of the capsule
2850
+ name: A name for the USD prim on the stage
2851
+ up_axis: The axis of the capsule that points up (0: x, 1: y, 2: z)
2852
+ color: The color of the capsule
2853
+ """
2854
+ geo_hash = hash(("capsule", radius, half_height, up_axis))
2855
+ if geo_hash in self._shape_geo_hash:
2856
+ shape = self._shape_geo_hash[geo_hash]
2857
+ if self.update_shape_instance(name, pos, rot):
2858
+ return shape
2859
+ else:
2860
+ vertices, indices = self._create_capsule_mesh(radius, half_height, up_axis=up_axis)
2861
+ shape = self.register_shape(geo_hash, vertices, indices, color1=color, color2=color)
2862
+ if not is_template:
2863
+ body = self._resolve_body_id(parent_body)
2864
+ self.add_shape_instance(name, shape, body, pos, rot)
2865
+ return shape
2866
+
2867
+ def render_cylinder(
2868
+ self,
2869
+ name: str,
2870
+ pos: tuple,
2871
+ rot: tuple,
2872
+ radius: float,
2873
+ half_height: float,
2874
+ parent_body: str = None,
2875
+ is_template: bool = False,
2876
+ up_axis: int = 1,
2877
+ color: tuple = None,
2878
+ ):
2879
+ """Add a cylinder for visualization
2880
+
2881
+ Args:
2882
+ pos: The position of the cylinder
2883
+ radius: The radius of the cylinder
2884
+ half_height: The half height of the cylinder
2885
+ name: A name for the USD prim on the stage
2886
+ up_axis: The axis of the cylinder that points up (0: x, 1: y, 2: z)
2887
+ color: The color of the capsule
2888
+ """
2889
+ geo_hash = hash(("cylinder", radius, half_height, up_axis))
2890
+ if geo_hash in self._shape_geo_hash:
2891
+ shape = self._shape_geo_hash[geo_hash]
2892
+ if self.update_shape_instance(name, pos, rot):
2893
+ return shape
2894
+ else:
2895
+ vertices, indices = self._create_cylinder_mesh(radius, half_height, up_axis=up_axis)
2896
+ shape = self.register_shape(geo_hash, vertices, indices, color1=color, color2=color)
2897
+ if not is_template:
2898
+ body = self._resolve_body_id(parent_body)
2899
+ self.add_shape_instance(name, shape, body, pos, rot)
2900
+ return shape
2901
+
2902
+ def render_cone(
2903
+ self,
2904
+ name: str,
2905
+ pos: tuple,
2906
+ rot: tuple,
2907
+ radius: float,
2908
+ half_height: float,
2909
+ parent_body: str = None,
2910
+ is_template: bool = False,
2911
+ up_axis: int = 1,
2912
+ color: tuple = None,
2913
+ ):
2914
+ """Add a cone for visualization
2915
+
2916
+ Args:
2917
+ pos: The position of the cone
2918
+ radius: The radius of the cone
2919
+ half_height: The half height of the cone
2920
+ name: A name for the USD prim on the stage
2921
+ up_axis: The axis of the cone that points up (0: x, 1: y, 2: z)
2922
+ color: The color of the cone
2923
+ """
2924
+ geo_hash = hash(("cone", radius, half_height, up_axis))
2925
+ if geo_hash in self._shape_geo_hash:
2926
+ shape = self._shape_geo_hash[geo_hash]
2927
+ if self.update_shape_instance(name, pos, rot):
2928
+ return shape
2929
+ else:
2930
+ vertices, indices = self._create_cone_mesh(radius, half_height, up_axis=up_axis)
2931
+ shape = self.register_shape(geo_hash, vertices, indices, color1=color, color2=color)
2932
+ if not is_template:
2933
+ body = self._resolve_body_id(parent_body)
2934
+ self.add_shape_instance(name, shape, body, pos, rot)
2935
+ return shape
2936
+
2937
+ def render_box(
2938
+ self,
2939
+ name: str,
2940
+ pos: tuple,
2941
+ rot: tuple,
2942
+ extents: tuple,
2943
+ parent_body: str = None,
2944
+ is_template: bool = False,
2945
+ color: tuple = None,
2946
+ ):
2947
+ """Add a box for visualization
2948
+
2949
+ Args:
2950
+ pos: The position of the box
2951
+ extents: The extents of the box
2952
+ name: A name for the USD prim on the stage
2953
+ color: The color of the box
2954
+ """
2955
+ geo_hash = hash(("box", tuple(extents)))
2956
+ if geo_hash in self._shape_geo_hash:
2957
+ shape = self._shape_geo_hash[geo_hash]
2958
+ if self.update_shape_instance(name, pos, rot):
2959
+ return shape
2960
+ else:
2961
+ vertices, indices = self._create_box_mesh(extents)
2962
+ shape = self.register_shape(geo_hash, vertices, indices, color1=color, color2=color)
2963
+ if not is_template:
2964
+ body = self._resolve_body_id(parent_body)
2965
+ self.add_shape_instance(name, shape, body, pos, rot)
2966
+ return shape
2967
+
2968
+ def render_mesh(
2969
+ self,
2970
+ name: str,
2971
+ points,
2972
+ indices,
2973
+ colors=None,
2974
+ pos=(0.0, 0.0, 0.0),
2975
+ rot=(0.0, 0.0, 0.0, 1.0),
2976
+ scale=(1.0, 1.0, 1.0),
2977
+ update_topology=False,
2978
+ parent_body: str = None,
2979
+ is_template: bool = False,
2980
+ smooth_shading: bool = True,
2981
+ ):
2982
+ """Add a mesh for visualization
2983
+
2984
+ Args:
2985
+ points: The points of the mesh
2986
+ indices: The indices of the mesh
2987
+ colors: The colors of the mesh
2988
+ pos: The position of the mesh
2989
+ rot: The rotation of the mesh
2990
+ scale: The scale of the mesh
2991
+ name: A name for the USD prim on the stage
2992
+ smooth_shading: Whether to average face normals at each vertex or introduce additional vertices for each face
2993
+ """
2994
+ if colors is not None:
2995
+ colors = np.array(colors, dtype=np.float32)
2996
+
2997
+ points = np.array(points, dtype=np.float32)
2998
+ point_count = len(points)
2999
+
3000
+ indices = np.array(indices, dtype=np.int32).reshape((-1, 3))
3001
+ idx_count = len(indices)
3002
+
3003
+ geo_hash = hash((points.tobytes(), indices.tobytes()))
3004
+
3005
+ if name in self._instances:
3006
+ # We've already registered this mesh instance and its associated shape.
3007
+ shape = self._instances[name][2]
3008
+ else:
3009
+ if geo_hash in self._shape_geo_hash:
3010
+ # We've only registered the shape, which can happen when `is_template` is `True`.
3011
+ shape = self._shape_geo_hash[geo_hash]
3012
+ else:
3013
+ shape = None
3014
+
3015
+ # Check if we already have that shape registered and can perform
3016
+ # minimal updates since the topology is not changing, before exiting.
3017
+ if not update_topology:
3018
+ if name in self._instances:
3019
+ # Update the instance's transform.
3020
+ self.update_shape_instance(name, pos, rot, color1=colors)
3021
+
3022
+ if shape is not None:
3023
+ # Update the shape's point positions.
3024
+ self.update_shape_vertices(shape, points, scale)
3025
+
3026
+ if not is_template and name not in self._instances:
3027
+ # Create a new instance.
3028
+ body = self._resolve_body_id(parent_body)
3029
+ self.add_shape_instance(name, shape, body, pos, rot, color1=colors)
3030
+
3031
+ return shape
3032
+
3033
+ # No existing shape for the given mesh was found, or its topology may have changed,
3034
+ # so we need to define a new one either way.
3035
+ if smooth_shading:
3036
+ normals = wp.zeros(point_count, dtype=wp.vec3)
3037
+ vertices = wp.array(points, dtype=wp.vec3)
3038
+ faces_per_vertex = wp.zeros(point_count, dtype=int)
3039
+ wp.launch(
3040
+ compute_average_normals,
3041
+ dim=idx_count,
3042
+ inputs=[wp.array(indices, dtype=int), vertices, scale],
3043
+ outputs=[normals, faces_per_vertex],
3044
+ )
3045
+ gfx_vertices = wp.zeros((point_count, 8), dtype=float)
3046
+ wp.launch(
3047
+ assemble_gfx_vertices,
3048
+ dim=point_count,
3049
+ inputs=[vertices, normals, faces_per_vertex, scale],
3050
+ outputs=[gfx_vertices],
3051
+ )
3052
+ gfx_vertices = gfx_vertices.numpy()
3053
+ gfx_indices = indices.flatten()
3054
+ else:
3055
+ gfx_vertices = wp.zeros((idx_count * 3, 8), dtype=float)
3056
+ wp.launch(
3057
+ compute_gfx_vertices,
3058
+ dim=idx_count,
3059
+ inputs=[wp.array(indices, dtype=int), wp.array(points, dtype=wp.vec3), scale],
3060
+ outputs=[gfx_vertices],
3061
+ )
3062
+ gfx_vertices = gfx_vertices.numpy()
3063
+ gfx_indices = np.arange(idx_count * 3)
3064
+
3065
+ # If there was a shape for the given mesh, clean it up.
3066
+ if shape is not None:
3067
+ self.deregister_shape(shape)
3068
+
3069
+ # If there was an instance for the given mesh, clean it up.
3070
+ if name in self._instances:
3071
+ self.remove_shape_instance(name)
3072
+
3073
+ # Register the new shape.
3074
+ shape = self.register_shape(geo_hash, gfx_vertices, gfx_indices)
3075
+
3076
+ if not is_template:
3077
+ # Create a new instance if necessary.
3078
+ body = self._resolve_body_id(parent_body)
3079
+ self.add_shape_instance(name, shape, body, pos, rot, color1=colors)
3080
+
3081
+ return shape
3082
+
3083
+ def render_arrow(
3084
+ self,
3085
+ name: str,
3086
+ pos: tuple,
3087
+ rot: tuple,
3088
+ base_radius: float,
3089
+ base_height: float,
3090
+ cap_radius: float = None,
3091
+ cap_height: float = None,
3092
+ parent_body: str = None,
3093
+ is_template: bool = False,
3094
+ up_axis: int = 1,
3095
+ color: Tuple[float, float, float] = None,
3096
+ ):
3097
+ """Add a arrow for visualization
3098
+
3099
+ Args:
3100
+ pos: The position of the arrow
3101
+ base_radius: The radius of the cylindrical base of the arrow
3102
+ base_height: The height of the cylindrical base of the arrow
3103
+ cap_radius: The radius of the conical cap of the arrow
3104
+ cap_height: The height of the conical cap of the arrow
3105
+ name: A name for the USD prim on the stage
3106
+ up_axis: The axis of the arrow that points up (0: x, 1: y, 2: z)
3107
+ """
3108
+ geo_hash = hash(("arrow", base_radius, base_height, cap_radius, cap_height, up_axis))
3109
+ if geo_hash in self._shape_geo_hash:
3110
+ shape = self._shape_geo_hash[geo_hash]
3111
+ if self.update_shape_instance(name, pos, rot, color1=color, color2=color):
3112
+ return shape
3113
+ else:
3114
+ vertices, indices = self._create_arrow_mesh(
3115
+ base_radius, base_height, cap_radius, cap_height, up_axis=up_axis
3116
+ )
3117
+ shape = self.register_shape(geo_hash, vertices, indices, color1=color, color2=color)
3118
+ if not is_template:
3119
+ body = self._resolve_body_id(parent_body)
3120
+ self.add_shape_instance(name, shape, body, pos, rot, color1=color, color2=color)
3121
+ return shape
3122
+
3123
+ def render_ref(self, name: str, path: str, pos: tuple, rot: tuple, scale: tuple, color: tuple = None):
3124
+ """
3125
+ Create a reference (instance) with the given name to the given path.
3126
+ """
3127
+
3128
+ if path in self._instances:
3129
+ _, body, shape, _, original_scale, color1, color2 = self._instances[path]
3130
+ if color is not None:
3131
+ color1 = color2 = color
3132
+ self.add_shape_instance(name, shape, body, pos, rot, scale or original_scale, color1, color2)
3133
+ return
3134
+
3135
+ raise Exception("Cannot create reference to path: " + path)
3136
+
3137
+ def render_points(self, name: str, points, radius, colors=None):
3138
+ """Add a set of points
3139
+
3140
+ Args:
3141
+ points: The points to render
3142
+ radius: The radius of the points (scalar or list)
3143
+ colors: The colors of the points
3144
+ name: A name for the USD prim on the stage
3145
+ """
3146
+
3147
+ if len(points) == 0:
3148
+ return
3149
+
3150
+ if isinstance(points, wp.array):
3151
+ wp_points = points
3152
+ else:
3153
+ wp_points = wp.array(points, dtype=wp.vec3, device=self._device)
3154
+
3155
+ if name not in self._shape_instancers:
3156
+ np_points = points.numpy() if isinstance(points, wp.array) else points
3157
+ instancer = ShapeInstancer(self._shape_shader, self._device)
3158
+ radius_is_scalar = np.isscalar(radius)
3159
+ if radius_is_scalar:
3160
+ vertices, indices = self._create_sphere_mesh(radius)
3161
+ else:
3162
+ vertices, indices = self._create_sphere_mesh(1.0)
3163
+ if colors is None:
3164
+ color = tab10_color_map(len(self._shape_geo_hash))
3165
+ elif len(colors) == 3:
3166
+ color = colors
3167
+ else:
3168
+ color = colors[0]
3169
+ instancer.register_shape(vertices, indices, color, color)
3170
+ scalings = None if radius_is_scalar else np.tile(radius, (3, 1)).T
3171
+ instancer.allocate_instances(np_points, colors1=colors, colors2=colors, scalings=scalings)
3172
+ self._shape_instancers[name] = instancer
3173
+ else:
3174
+ instancer = self._shape_instancers[name]
3175
+ if len(points) != instancer.num_instances:
3176
+ np_points = points.numpy() if isinstance(points, wp.array) else points
3177
+ instancer.allocate_instances(np_points)
3178
+
3179
+ with instancer:
3180
+ wp.launch(
3181
+ update_points_positions,
3182
+ dim=len(points),
3183
+ inputs=[wp_points, instancer.instance_scalings],
3184
+ outputs=[instancer.vbo_transforms],
3185
+ device=self._device,
3186
+ )
3187
+
3188
+ def _render_lines(self, name: str, lines, color: tuple, radius: float = 0.01):
3189
+ if len(lines) == 0:
3190
+ return
3191
+
3192
+ if name not in self._shape_instancers:
3193
+ instancer = ShapeInstancer(self._shape_shader, self._device)
3194
+ vertices, indices = self._create_capsule_mesh(radius, 0.5)
3195
+ if color is None or isinstance(color, list) and len(color) > 0 and isinstance(color[0], list):
3196
+ color = tab10_color_map(len(self._shape_geo_hash))
3197
+ instancer.register_shape(vertices, indices, color, color)
3198
+ instancer.allocate_instances(np.zeros((len(lines), 3)))
3199
+ self._shape_instancers[name] = instancer
3200
+ else:
3201
+ instancer = self._shape_instancers[name]
3202
+ if len(lines) != instancer.num_instances:
3203
+ instancer.allocate_instances(np.zeros((len(lines), 3)))
3204
+ instancer.update_colors(color, color)
3205
+
3206
+ lines_wp = wp.array(lines, dtype=wp.vec3, ndim=2, device=self._device)
3207
+ with instancer:
3208
+ wp.launch(
3209
+ update_line_transforms,
3210
+ dim=len(lines),
3211
+ inputs=[lines_wp],
3212
+ outputs=[instancer.vbo_transforms],
3213
+ device=self._device,
3214
+ )
3215
+
3216
+ def render_line_list(self, name: str, vertices, indices, color: tuple = None, radius: float = 0.01):
3217
+ """Add a line list as a set of capsules
3218
+
3219
+ Args:
3220
+ vertices: The vertices of the line-list
3221
+ indices: The indices of the line-list
3222
+ color: The color of the line
3223
+ radius: The radius of the line
3224
+ """
3225
+ lines = []
3226
+ for i in range(len(indices) // 2):
3227
+ lines.append((vertices[indices[2 * i]], vertices[indices[2 * i + 1]]))
3228
+ lines = np.array(lines)
3229
+ self._render_lines(name, lines, color, radius)
3230
+
3231
+ def render_line_strip(self, name: str, vertices, color: tuple = None, radius: float = 0.01):
3232
+ """Add a line strip as a set of capsules
3233
+
3234
+ Args:
3235
+ vertices: The vertices of the line-strip
3236
+ color: The color of the line
3237
+ radius: The radius of the line
3238
+ """
3239
+ lines = []
3240
+ for i in range(len(vertices) - 1):
3241
+ lines.append((vertices[i], vertices[i + 1]))
3242
+ lines = np.array(lines)
3243
+ self._render_lines(name, lines, color, radius)
3244
+
3245
+ def update_shape_vertices(self, shape, points, scale):
3246
+ if isinstance(points, wp.array):
3247
+ wp_points = points.to(self._device)
3248
+ else:
3249
+ wp_points = wp.array(points, dtype=wp.vec3, device=self._device)
3250
+
3251
+ cuda_buffer = self._shape_gl_buffers[shape][4]
3252
+ vertices_shape = self._shapes[shape][0].shape
3253
+ vbo_vertices = cuda_buffer.map(dtype=wp.float32, shape=vertices_shape)
3254
+
3255
+ wp.launch(
3256
+ update_vbo_vertices,
3257
+ dim=vertices_shape[0],
3258
+ inputs=[wp_points, scale],
3259
+ outputs=[vbo_vertices],
3260
+ device=self._device,
3261
+ )
3262
+
3263
+ cuda_buffer.unmap()
3264
+
3265
+ @staticmethod
3266
+ def _create_sphere_mesh(
3267
+ radius=1.0,
3268
+ num_latitudes=default_num_segments,
3269
+ num_longitudes=default_num_segments,
3270
+ reverse_winding=False,
3271
+ ):
3272
+ vertices = []
3273
+ indices = []
3274
+
3275
+ for i in range(num_latitudes + 1):
3276
+ theta = i * np.pi / num_latitudes
3277
+ sin_theta = np.sin(theta)
3278
+ cos_theta = np.cos(theta)
3279
+
3280
+ for j in range(num_longitudes + 1):
3281
+ phi = j * 2 * np.pi / num_longitudes
3282
+ sin_phi = np.sin(phi)
3283
+ cos_phi = np.cos(phi)
3284
+
3285
+ x = cos_phi * sin_theta
3286
+ y = cos_theta
3287
+ z = sin_phi * sin_theta
3288
+
3289
+ u = float(j) / num_longitudes
3290
+ v = float(i) / num_latitudes
3291
+
3292
+ vertices.append([x * radius, y * radius, z * radius, x, y, z, u, v])
3293
+
3294
+ for i in range(num_latitudes):
3295
+ for j in range(num_longitudes):
3296
+ first = i * (num_longitudes + 1) + j
3297
+ second = first + num_longitudes + 1
3298
+
3299
+ if reverse_winding:
3300
+ indices.extend([first, second, first + 1, second, second + 1, first + 1])
3301
+ else:
3302
+ indices.extend([first, first + 1, second, second, first + 1, second + 1])
3303
+
3304
+ return np.array(vertices, dtype=np.float32), np.array(indices, dtype=np.uint32)
3305
+
3306
+ @staticmethod
3307
+ def _create_capsule_mesh(radius, half_height, up_axis=1, segments=default_num_segments):
3308
+ vertices = []
3309
+ indices = []
3310
+
3311
+ x_dir, y_dir, z_dir = ((1, 2, 0), (2, 0, 1), (0, 1, 2))[up_axis]
3312
+ up_vector = np.zeros(3)
3313
+ up_vector[up_axis] = half_height
3314
+
3315
+ for i in range(segments + 1):
3316
+ theta = i * np.pi / segments
3317
+ sin_theta = np.sin(theta)
3318
+ cos_theta = np.cos(theta)
3319
+
3320
+ for j in range(segments + 1):
3321
+ phi = j * 2 * np.pi / segments
3322
+ sin_phi = np.sin(phi)
3323
+ cos_phi = np.cos(phi)
3324
+
3325
+ z = cos_phi * sin_theta
3326
+ y = cos_theta
3327
+ x = sin_phi * sin_theta
3328
+
3329
+ u = cos_theta * 0.5 + 0.5
3330
+ v = cos_phi * sin_theta * 0.5 + 0.5
3331
+
3332
+ xyz = x, y, z
3333
+ x, y, z = xyz[x_dir], xyz[y_dir], xyz[z_dir]
3334
+ xyz = np.array((x, y, z), dtype=np.float32) * radius
3335
+ if j < segments // 2:
3336
+ xyz += up_vector
3337
+ else:
3338
+ xyz -= up_vector
3339
+
3340
+ vertices.append([*xyz, x, y, z, u, v])
3341
+
3342
+ nv = len(vertices)
3343
+ for i in range(segments + 1):
3344
+ for j in range(segments + 1):
3345
+ first = (i * (segments + 1) + j) % nv
3346
+ second = (first + segments + 1) % nv
3347
+ indices.extend([first, second, (first + 1) % nv, second, (second + 1) % nv, (first + 1) % nv])
3348
+
3349
+ vertex_data = np.array(vertices, dtype=np.float32)
3350
+ index_data = np.array(indices, dtype=np.uint32)
3351
+
3352
+ return vertex_data, index_data
3353
+
3354
+ @staticmethod
3355
+ def _create_cone_mesh(radius, half_height, up_axis=1, segments=default_num_segments):
3356
+ # render it as a cylinder with zero top radius so we get correct normals on the sides
3357
+ return OpenGLRenderer._create_cylinder_mesh(radius, half_height, up_axis, segments, 0.0)
3358
+
3359
+ @staticmethod
3360
+ def _create_cylinder_mesh(radius, half_height, up_axis=1, segments=default_num_segments, top_radius=None):
3361
+ if up_axis not in (0, 1, 2):
3362
+ raise ValueError("up_axis must be between 0 and 2")
3363
+
3364
+ x_dir, y_dir, z_dir = (
3365
+ (1, 2, 0),
3366
+ (0, 1, 2),
3367
+ (2, 0, 1),
3368
+ )[up_axis]
3369
+
3370
+ indices = []
3371
+
3372
+ cap_vertices = []
3373
+ side_vertices = []
3374
+
3375
+ # create center cap vertices
3376
+ position = np.array([0, -half_height, 0])[[x_dir, y_dir, z_dir]]
3377
+ normal = np.array([0, -1, 0])[[x_dir, y_dir, z_dir]]
3378
+ cap_vertices.append([*position, *normal, 0.5, 0.5])
3379
+ cap_vertices.append([*-position, *-normal, 0.5, 0.5])
3380
+
3381
+ if top_radius is None:
3382
+ top_radius = radius
3383
+ side_slope = -np.arctan2(top_radius - radius, 2 * half_height)
3384
+
3385
+ # create the cylinder base and top vertices
3386
+ for j in (-1, 1):
3387
+ center_index = max(j, 0)
3388
+ if j == 1:
3389
+ radius = top_radius
3390
+ for i in range(segments):
3391
+ theta = 2 * np.pi * i / segments
3392
+
3393
+ cos_theta = np.cos(theta)
3394
+ sin_theta = np.sin(theta)
3395
+
3396
+ x = cos_theta
3397
+ y = j * half_height
3398
+ z = sin_theta
3399
+
3400
+ position = np.array([radius * x, y, radius * z])
3401
+
3402
+ normal = np.array([x, side_slope, z])
3403
+ normal = normal / np.linalg.norm(normal)
3404
+ uv = (i / (segments - 1), (j + 1) / 2)
3405
+ vertex = np.hstack([position[[x_dir, y_dir, z_dir]], normal[[x_dir, y_dir, z_dir]], uv])
3406
+ side_vertices.append(vertex)
3407
+
3408
+ normal = np.array([0, j, 0])
3409
+ uv = (cos_theta * 0.5 + 0.5, sin_theta * 0.5 + 0.5)
3410
+ vertex = np.hstack([position[[x_dir, y_dir, z_dir]], normal[[x_dir, y_dir, z_dir]], uv])
3411
+ cap_vertices.append(vertex)
3412
+
3413
+ cs = center_index * segments
3414
+ indices.extend([center_index, i + cs + 2, (i + 1) % segments + cs + 2][::-j])
3415
+
3416
+ # create the cylinder side indices
3417
+ for i in range(segments):
3418
+ index1 = len(cap_vertices) + i + segments
3419
+ index2 = len(cap_vertices) + ((i + 1) % segments) + segments
3420
+ index3 = len(cap_vertices) + i
3421
+ index4 = len(cap_vertices) + ((i + 1) % segments)
3422
+
3423
+ indices.extend([index1, index2, index3, index2, index4, index3])
3424
+
3425
+ vertex_data = np.array(np.vstack((cap_vertices, side_vertices)), dtype=np.float32)
3426
+ index_data = np.array(indices, dtype=np.uint32)
3427
+
3428
+ return vertex_data, index_data
3429
+
3430
+ @staticmethod
3431
+ def _create_arrow_mesh(
3432
+ base_radius, base_height, cap_radius=None, cap_height=None, up_axis=1, segments=default_num_segments
3433
+ ):
3434
+ if up_axis not in (0, 1, 2):
3435
+ raise ValueError("up_axis must be between 0 and 2")
3436
+ if cap_radius is None:
3437
+ cap_radius = base_radius * 1.8
3438
+ if cap_height is None:
3439
+ cap_height = base_height * 0.18
3440
+
3441
+ up_vector = np.array([0, 0, 0])
3442
+ up_vector[up_axis] = 1
3443
+
3444
+ base_vertices, base_indices = OpenGLRenderer._create_cylinder_mesh(
3445
+ base_radius, base_height / 2, up_axis, segments
3446
+ )
3447
+ cap_vertices, cap_indices = OpenGLRenderer._create_cone_mesh(cap_radius, cap_height / 2, up_axis, segments)
3448
+
3449
+ base_vertices[:, :3] += base_height / 2 * up_vector
3450
+ # move cap slightly lower to avoid z-fighting
3451
+ cap_vertices[:, :3] += (base_height + cap_height / 2 - 1e-3 * base_height) * up_vector
3452
+
3453
+ vertex_data = np.vstack((base_vertices, cap_vertices))
3454
+ index_data = np.hstack((base_indices, cap_indices + len(base_vertices)))
3455
+
3456
+ return vertex_data, index_data
3457
+
3458
+ @staticmethod
3459
+ def _create_box_mesh(extents):
3460
+ x_extent, y_extent, z_extent = extents
3461
+
3462
+ vertices = [
3463
+ # Position Normal UV
3464
+ [-x_extent, -y_extent, -z_extent, -1, 0, 0, 0, 0],
3465
+ [-x_extent, -y_extent, z_extent, -1, 0, 0, 1, 0],
3466
+ [-x_extent, y_extent, z_extent, -1, 0, 0, 1, 1],
3467
+ [-x_extent, y_extent, -z_extent, -1, 0, 0, 0, 1],
3468
+ [x_extent, -y_extent, -z_extent, 1, 0, 0, 0, 0],
3469
+ [x_extent, -y_extent, z_extent, 1, 0, 0, 1, 0],
3470
+ [x_extent, y_extent, z_extent, 1, 0, 0, 1, 1],
3471
+ [x_extent, y_extent, -z_extent, 1, 0, 0, 0, 1],
3472
+ [-x_extent, -y_extent, -z_extent, 0, -1, 0, 0, 0],
3473
+ [-x_extent, -y_extent, z_extent, 0, -1, 0, 1, 0],
3474
+ [x_extent, -y_extent, z_extent, 0, -1, 0, 1, 1],
3475
+ [x_extent, -y_extent, -z_extent, 0, -1, 0, 0, 1],
3476
+ [-x_extent, y_extent, -z_extent, 0, 1, 0, 0, 0],
3477
+ [-x_extent, y_extent, z_extent, 0, 1, 0, 1, 0],
3478
+ [x_extent, y_extent, z_extent, 0, 1, 0, 1, 1],
3479
+ [x_extent, y_extent, -z_extent, 0, 1, 0, 0, 1],
3480
+ [-x_extent, -y_extent, -z_extent, 0, 0, -1, 0, 0],
3481
+ [-x_extent, y_extent, -z_extent, 0, 0, -1, 1, 0],
3482
+ [x_extent, y_extent, -z_extent, 0, 0, -1, 1, 1],
3483
+ [x_extent, -y_extent, -z_extent, 0, 0, -1, 0, 1],
3484
+ [-x_extent, -y_extent, z_extent, 0, 0, 1, 0, 0],
3485
+ [-x_extent, y_extent, z_extent, 0, 0, 1, 1, 0],
3486
+ [x_extent, y_extent, z_extent, 0, 0, 1, 1, 1],
3487
+ [x_extent, -y_extent, z_extent, 0, 0, 1, 0, 1],
3488
+ ]
3489
+
3490
+ # fmt: off
3491
+ indices = [
3492
+ 0, 1, 2,
3493
+ 0, 2, 3,
3494
+ 4, 6, 5,
3495
+ 4, 7, 6,
3496
+ 8, 10, 9,
3497
+ 8, 11, 10,
3498
+ 12, 13, 14,
3499
+ 12, 14, 15,
3500
+ 16, 17, 18,
3501
+ 16, 18, 19,
3502
+ 20, 22, 21,
3503
+ 20, 23, 22,
3504
+ ]
3505
+ # fmt: on
3506
+ return np.array(vertices, dtype=np.float32), np.array(indices, dtype=np.uint32)
3507
+
3508
+ def _switch_context(self):
3509
+ try:
3510
+ self.window.switch_to()
3511
+ except AttributeError:
3512
+ # The window could be in the process of being closed, in which case
3513
+ # its corresponding context might have been destroyed and set to `None`.
3514
+ pass
3515
+
3516
+
3517
+ if __name__ == "__main__":
3518
+ renderer = OpenGLRenderer()