[MISC] Various code cleanup and performance improvements. (#1314)

* Refactor math helper methods to be more generic and improve performance.
* Speed up primitive mesh generation for viewer.

Author: duburcqa

genesis/engine/entities/emitter.py

@@ -147,8 +147,10 @@ def emit(
             else:
                 gs.raise_exception()
 
-            positions = gu.transform_by_T(
-                positions, gu.trans_R_to_T(pos, gu.z_to_R(direction) @ gu.axis_angle_to_R(np.array([0, 0, 1]), theta))
+            positions = gu.transform_by_trans_R(
+                positions,
+                pos,
+                gu.z_up_to_R(direction) @ gu.axis_angle_to_R(np.array([0, 0, 1]), theta),
             ).astype(gs.np_float)
 
             positions = np.tile(positions[np.newaxis], (self._sim._B, 1, 1))

genesis/engine/entities/rigid_entity/rigid_link.py

@@ -144,7 +144,8 @@ def _build(self):
 
                 if inertia_mesh.is_watertight and self._init_mesh.mass > 0:
                     # TODO: check if this is correct. This is correct if the inertia frame is w.r.t to link frame
-                    T_inertia = gu.trans_quat_to_T(self._inertial_pos, self._inertial_quat)
+                    dtype = np.result_type(self._inertial_pos, self._inertial_quat)
+                    T_inertia = gu.trans_quat_to_T(self._inertial_pos.astype(dtype), self._inertial_quat.astype(dtype))
                     self._inertial_i = (
                         self._init_mesh.moment_inertia_frame(T_inertia) / self._init_mesh.mass * self._inertial_mass
                     )

genesis/engine/solvers/rigid/constraint_solver_decomp.py

@@ -122,7 +122,7 @@ def add_collision_constraints(self):
                 link_a_maybe_batch = [link_a, i_b] if ti.static(self._solver._options.batch_links_info) else link_a
                 link_b_maybe_batch = [link_b, i_b] if ti.static(self._solver._options.batch_links_info) else link_b
 
-                d1, d2 = gu.orthogonals(contact_data.normal)
+                d1, d2 = gu.ti_orthogonals(contact_data.normal)
 
                 invweight = self._solver.links_info[link_a_maybe_batch].invweight[0]
                 if link_b > -1:
@@ -775,7 +775,7 @@ def _func_update_contact_force(self):
                 contact_data = self._collider.contact_data[i_c, i_b]
 
                 force = ti.Vector.zero(gs.ti_float, 3)
-                d1, d2 = gu.orthogonals(contact_data.normal)
+                d1, d2 = gu.ti_orthogonals(contact_data.normal)
                 for i_dir in range(4):
                     d = (2 * (i_dir % 2) - 1) * (d1 if i_dir < 2 else d2)
                     n = d * contact_data.friction - contact_data.normal

genesis/engine/solvers/rigid/constraint_solver_decomp_island.py

@@ -139,7 +139,7 @@ def add_collision_constraints(self, island, i_b):
             link_a_maybe_batch = [link_a, i_b] if ti.static(self._solver._options.batch_links_info) else link_a
             link_b_maybe_batch = [link_b, i_b] if ti.static(self._solver._options.batch_links_info) else link_b
 
-            d1, d2 = gu.orthogonals(contact_data.normal)
+            d1, d2 = gu.ti_orthogonals(contact_data.normal)
 
             invweight = self._solver.links_info[link_a_maybe_batch].invweight[0] + self._solver.links_info[
                 link_b_maybe_batch
@@ -526,7 +526,7 @@ def _func_update_contact_force(self, island, i_b):
             contact_data = self._collider.contact_data[i_col, i_b]
 
             force = ti.Vector.zero(gs.ti_float, 3)
-            d1, d2 = gu.orthogonals(contact_data.normal)
+            d1, d2 = gu.ti_orthogonals(contact_data.normal)
             for i in range(4):
                 d = (2 * (i % 2) - 1) * (d1 if i < 2 else d2)
                 n = d * contact_data.friction - contact_data.normal

genesis/engine/solvers/rigid/rigid_solver_decomp.py

@@ -1006,7 +1006,7 @@ def _init_geom_fields(self):
         self.geoms_state = struct_geom_state.field(
             shape=self._batch_shape(self.n_geoms_), needs_grad=False, layout=ti.Layout.SOA
         )
-        self._geoms_render_T = np.empty((self.n_geoms_, self._B, 4, 4), order="F", dtype=np.float32)
+        self._geoms_render_T = np.empty((self.n_geoms_, self._B, 4, 4), dtype=np.float32)
 
         if self.n_geoms > 0:
             # Make sure that the constraints parameters are valid
@@ -1189,7 +1189,7 @@ def _init_vgeom_fields(self):
         self.vgeoms_state = struct_vgeom_state.field(
             shape=self._batch_shape(self.n_vgeoms_), needs_grad=False, layout=ti.Layout.SOA
         )
-        self._vgeoms_render_T = np.empty((self.n_vgeoms_, self._B, 4, 4), order="F", dtype=np.float32)
+        self._vgeoms_render_T = np.empty((self.n_vgeoms_, self._B, 4, 4), dtype=np.float32)
 
         if self.n_vgeoms > 0:
             vgeoms = self.vgeoms