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