Fix dtype handling.

Author: duburcqa

genesis/engine/entities/emitter.py

@@ -101,14 +101,15 @@ def emit(
         else:
             gs.raise_exception(f"Unsupported nozzle shape: {droplet_shape}.")
 
+        direction = np.asarray(direction, dtype=gs.np_float)
         if np.linalg.norm(direction) < gs.EPS:
             gs.raise_exception("Zero-length direction.")
         else:
             direction = gu.normalize(direction)
 
         p_size = self._solver.particle_size if p_size is None else p_size
 
-        pos = np.array(pos)
+        pos = np.asarray(pos, dtype=gs.np_float)
         if droplet_length is None:
             # Use the speed to determine the length of the droplet in the emitting direction
             droplet_length = speed * self._solver.substep_dt * self._sim.substeps + self._acc_droplet_len
@@ -148,10 +149,10 @@ def emit(
                 gs.raise_exception()
 
             positions = gu.transform_by_trans_R(
-                positions,
+                positions.astype(gs.np_float, copy=False),
                 pos,
                 gu.z_up_to_R(direction) @ gu.axis_angle_to_R(np.array([0.0, 0.0, 1.0], dtype=gs.np_float), theta),
-            ).astype(gs.np_float)
+            )
 
             positions = np.tile(positions[np.newaxis], (self._sim._B, 1, 1))
 
@@ -161,8 +162,7 @@ def emit(
             n_particles = positions.shape[1]
 
             # Expand vels with batch dimension
-            vels = np.tile(direction * speed, (n_particles, 1)).astype(gs.np_float)
-            vels = np.tile(vels[np.newaxis], (self._sim._B, 1, 1))
+            vels = np.tile((speed * direction).reshape((1, 1, -1)), (self._sim._B, n_particles, 1))
 
             if n_particles > self._entity.n_particles:
                 gs.logger.warning(

genesis/engine/entities/fem_entity.py

@@ -62,28 +62,26 @@ def __init__(self, scene, solver, material, morph, surface, idx, v_start=0, el_s
         el2tri = np.array(
             [  # follow the order with correct normal
                 [[v[0], v[2], v[1]], [v[1], v[2], v[3]], [v[0], v[1], v[3]], [v[0], v[3], v[2]]] for v in self.elems
-            ]
+            ],
+            dtype=gs.np_int,
         )
-        all_tri = el2tri.reshape(-1, 3)
+        all_tri = el2tri.reshape((-1, 3))
         all_tri_sorted = np.sort(all_tri, axis=1)
         _, unique_idcs, cnt = np.unique(all_tri_sorted, axis=0, return_counts=True, return_index=True)
         unique_tri = all_tri[unique_idcs]
         surface_tri = unique_tri[cnt == 1]
 
-        self._surface_tri_np = surface_tri.astype(gs.np_int)
+        self._surface_tri_np = surface_tri
         self._n_surfaces = len(self._surface_tri_np)
 
         if self._n_surfaces > 0:
             self._n_surface_vertices = len(np.unique(self._surface_tri_np))
         else:
             self._n_surface_vertices = 0
 
-        tri2el = np.repeat(np.arange(self.elems.shape[0])[:, None], 4, axis=-1)
-        all_el = tri2el.reshape(
-            -1,
-        )
-        unique_el = all_el[unique_idcs]
-        self._surface_el_np = unique_el[cnt == 1].astype(gs.np_int)
+        tri2el = np.repeat(np.arange(self.elems.shape[0], dtype=gs.np_int)[:, np.newaxis], 4, axis=1)
+        unique_el = tri2el.flat[unique_idcs]
+        self._surface_el_np = unique_el[cnt == 1]
 
         if isinstance(self.sim.coupler, SAPCoupler):
             self.compute_pressure_field()
@@ -259,7 +257,7 @@ def set_muscle(self, muscle_group=None, muscle_direction=None):
         if muscle_direction is not None:
             muscle_direction = to_gs_tensor(muscle_direction)
             assert muscle_direction.shape == (self.n_elements, 3)
-            assert torch.allclose(muscle_direction.norm(dim=-1), torch.Tensor([1.0]).to(muscle_direction))
+            assert ((1.0 - muscle_direction.norm(dim=-1)).abs() < gs.EPS).all()
 
             self.set_muscle_direction(muscle_direction)
 
@@ -308,8 +306,8 @@ def instantiate(self, verts, elems):
         Exception
             If no vertices are provided.
         """
-        verts = verts.astype(gs.np_float)
-        elems = elems.astype(gs.np_int)
+        verts = verts.astype(gs.np_float, copy=False)
+        elems = elems.astype(gs.np_int, copy=False)
 
         # rotate
         R = gu.quat_to_R(np.array(self.morph.quat, dtype=gs.np_float))
@@ -368,8 +366,8 @@ def _add_to_solver(self, in_backward=False):
             )
 
         # Convert to appropriate numpy array types
-        elems_np = self.elems.astype(gs.np_int)
-        verts_numpy = self.init_positions.cpu().numpy().astype(gs.np_float)
+        elems_np = self.elems.astype(gs.np_int, copy=False)
+        verts_numpy = tensor_to_array(self.init_positions, dtype=gs.np_float)
 
         self._solver._kernel_add_elements(
             f=self._sim.cur_substep_local,

genesis/engine/entities/hybrid_entity.py

@@ -144,7 +144,7 @@ def __init__(
 
             # set muscle group
             mat_soft._n_groups = len(link_idcs)
-            self._muscle_group_cache = muscle_group
+            self._muscle_group_cache = muscle_group.astype(gs.np_int)
 
             # set up info in taichi field
             if isinstance(mat_soft, gs.materials.MPM.Base):
@@ -154,10 +154,10 @@ def __init__(
                     trans_local_to_global=gs.ti_vec3,
                     quat_local_to_global=gs.ti_vec4,
                 ).field(shape=(mat_soft.n_groups,), needs_grad=False, layout=ti.Layout.SOA)
-                part_soft_info.link_idx.from_numpy(np.array(link_idcs).astype(gs.np_int))
-                part_soft_info.geom_idx.from_numpy(np.array(geom_idcs).astype(gs.np_int))
-                part_soft_info.trans_local_to_global.from_numpy(np.array(trans_local_to_global).astype(gs.np_float))
-                part_soft_info.quat_local_to_global.from_numpy(np.array(quat_local_to_global).astype(gs.np_float))
+                part_soft_info.link_idx.from_numpy(np.asarray(link_idcs, dtype=gs.np_int))
+                part_soft_info.geom_idx.from_numpy(np.asarray(geom_idcs, dtype=gs.np_int))
+                part_soft_info.trans_local_to_global.from_numpy(np.asarray(trans_local_to_global, dtype=gs.np_float))
+                part_soft_info.quat_local_to_global.from_numpy(np.asarray(quat_local_to_global, dtype=gs.np_float))
 
                 part_soft_init_positions = ti.field(dtype=gs.ti_vec3, shape=(part_soft.init_particles.shape[0],))
                 part_soft_init_positions.from_torch(gs.Tensor(part_soft.init_particles))
@@ -352,7 +352,7 @@ def build(self):
         # can only be called here (at sim build)
         if not self.material.use_default_coupling and self._muscle_group_cache is not None:
             self._part_soft.set_muscle(
-                muscle_group=gs.tensor(self._muscle_group_cache.astype(gs.np_int)),
+                muscle_group=gs.tensor(self._muscle_group_cache),
                 # no muscle direction as the soft body is actuated by rigid parts
             )
 
@@ -618,7 +618,7 @@ def default_func_instantiate_rigid_from_soft(
     # add rigid entity
     mesh_center = (mesh.vertices.max(0) + mesh.vertices.min(0)) / 2.0
     offset = (graph_pos[src_node] - mesh_center) * morph.scale
-    pos_rigid = morph.pos + offset.astype(gs.np_float)
+    pos_rigid = morph.pos + offset.astype(gs.np_float, copy=False)
     quat_rigid = morph.quat
     scale_rigid = morph.scale
     morph_rigid = gs.morphs.URDF(
@@ -700,11 +700,10 @@ def default_func_instantiate_rigid_soft_association_from_soft(
         dist_to_p1 = np.linalg.norm(positions - p1, axis=-1)
         dist_to_line = np.sqrt(dist_to_p0**2 - ((positions_proj_on_line_t[:, None] * line_vec) ** 2).sum(-1))
 
-        dist_to_link = (
-            dist_to_p0 * (positions_proj_on_line_t < 0).astype(float)
-            + dist_to_p1 * (positions_proj_on_line_t > 1).astype(float)
-            + dist_to_line * np.logical_and(positions_proj_on_line_t >= 0, positions_proj_on_line_t <= 1).astype(float)
-        )
+        is_clipped_low = positions_proj_on_line_t < 0.0
+        is_clipped_high = positions_proj_on_line_t > 1.0
+        is_valid = ~is_clipped_low & ~is_clipped_high
+        dist_to_link = dist_to_p0 * is_clipped_low + dist_to_p1 * is_clipped_high + dist_to_line * is_valid
 
         trans, quat = gu.transform_pos_quat_by_trans_quat(
             geom.init_pos, geom.init_quat, link.init_x_pos, link.init_x_quat

genesis/engine/entities/mpm_entity.py

@@ -368,7 +368,7 @@ def get_free(self):
         """
         self._assert_active()
 
-        free = gs.zeros((self._n_particles,), dtype=int, requires_grad=False, scene=self._scene)
+        free = gs.zeros((self._n_particles,), dtype=torch.bool, requires_grad=False, scene=self._scene)
         self.solver._kernel_get_free(
             self._particle_start,
             self._n_particles,

genesis/engine/entities/particle_entity.py

@@ -176,18 +176,18 @@ def _kernel_add_vverts_to_solver(
                     self.solver.vverts_info.support_idxs[i_vv][j] = support_idxs_local[i_vv_, j] + self._particle_start
 
         kdtree = KDTree(self._particles)
-        _, support_idxs = kdtree.query(self._vverts, k=self.solver._n_vvert_supports)
+        _, support_idxs = kdtree.query(self._vverts, k=self.solver._n_vvert_supports).astype(gs.np_int, copy=False)
         support_idxs = np.clip(support_idxs, 0, len(self._particles) - 1)
         all_ps = self._particles[support_idxs]
-        Ps = np.stack([all_ps[:, i, :] for i in range(self.solver._n_vvert_supports - 1)], axis=2) - np.expand_dims(
-            all_ps[:, -1, :], axis=2
-        )
+        Ps = np.stack(
+            [all_ps[:, i, :] for i in range(self.solver._n_vvert_supports - 1)], axis=2, dtype=gs.np_float
+        ) - np.expand_dims(all_ps[:, -1, :], axis=2)
 
         _kernel_add_vverts_to_solver(
-            vverts=self._vverts.astype(gs.np_float),
-            particles=self._particles.astype(gs.np_float),
-            P_invs=np.linalg.pinv(Ps).astype(gs.np_float),
-            support_idxs_local=support_idxs.astype(gs.np_int),
+            vverts=self._vverts,
+            particles=self._particles,
+            P_invs=np.linalg.pinv(Ps),
+            support_idxs_local=support_idxs,
         )
 
     def sample(self):
@@ -232,7 +232,9 @@ def sample(self):
                 particles.append(particles_i)
 
         elif isinstance(self._morph, (gs.options.morphs.Primitive, gs.options.morphs.Mesh)):
-            particles = self._vmesh.particlize(self._particle_size, self.sampler)
+            particles = self._vmesh.particlize(self._particle_size, self.sampler).astype(
+                gs.np_float, order="C", copy=False
+            )
 
         elif isinstance(self._morph, gs.options.morphs.Nowhere):
             particles = pu.nowhere_particles(self._morph.n_particles)
@@ -244,14 +246,14 @@ def sample(self):
             gs.raise_exception("Entity has zero particles.")
 
         if isinstance(self._morph, gs.options.morphs.Nowhere):
-            origin = gu.nowhere().astype(gs.np_float)
-            self._vverts = np.array([])
-            self._vfaces = np.array([])
+            origin = gu.nowhere()
+            self._vverts = np.array([], dtype=gs.np_float)
+            self._vfaces = np.array([], dtype=gs.np_float)
 
         elif isinstance(self._morph, gs.options.morphs.MeshSet):
             for i in range(len(self._morph.files)):
-                pos_i = np.array(self._morph.poss[i])
-                quat_i = np.array(gu.euler_to_quat(self._morph.eulers[i]))
+                pos_i, euler_i = map(np.asarray, (self._morph.poss[i], self._morph.eulers[i]))
+                quat_i = gs.utils.geom.xyz_to_quat(euler_i, rpy=True, degrees=True)
                 self._vmesh[i].apply_transform(gu.trans_quat_to_T(pos_i, quat_i))
 
                 # NOTE: particles are transformed already
@@ -260,7 +262,7 @@ def sample(self):
             self.mesh_set_group_ids = np.concatenate(
                 [np.ones((v.shape[0],), dtype=gs.np_int) * i for i, v in enumerate(particles)]
             )
-            particles = np.concatenate(particles)
+            particles = np.concatenate(particles, dtype=gs.np_float)
             if not self._solver.boundary.is_inside(particles):  # HACK no check
                 gs.raise_exception(
                     f"Entity has particles outside solver boundary. Note that for MPMSolver, boundary is slightly tighter than the specified domain due to safety padding.\n\nCurrent boundary:\n{self._solver.boundary}\n\nEntity to be added:\nmin: {particles.min(0)}\nmax: {particles.max(0)}\n"
@@ -277,35 +279,43 @@ def sample(self):
                 vertex_normals=combined_vert_normals,
             )
             self._vmesh = mu.trimesh_to_mesh(combined_tmesh, 1, self._surface)
+
             if self._need_skinning:
-                self._vverts = np.array(self._vmesh.verts)
-                self._vfaces = np.array(self._vmesh.faces)
+                self._vverts = np.asarray(self._vmesh.verts, dtype=gs.np_float)
+                self._vfaces = np.asarray(self._vmesh.faces, dtype=gs.np_float)
             else:
-                self._vverts = np.array([])
-                self._vfaces = np.array([])
+                self._vverts = np.asarray([], dtype=gs.np_float)
+                self._vfaces = np.asarray([], dtype=gs.np_float)
             origin = np.mean(self._morph.poss, dtype=gs.np_float)
 
         else:
             # transform vmesh
-            self._vmesh.apply_transform(gu.trans_quat_to_T(np.array(self._morph.pos), np.array(self._morph.quat)))
+            pos, quat = map(np.asarray, (self._morph.pos, self._morph.quat))
+            self._vmesh.apply_transform(gu.trans_quat_to_T(pos, quat))
             # transform particles
-            origin = np.array(self._morph.pos, dtype=gs.np_float)
-            particles = gu.transform_by_trans_quat(particles, np.array(self._morph.pos), np.array(self._morph.quat))
-            # rotate
+            particles = gu.transform_by_trans_quat(
+                particles,
+                np.asarray(self._morph.pos, dtype=gs.np_float),
+                np.asarray(self._morph.quat, dtype=gs.np_float),
+            )
 
             if not self._solver.boundary.is_inside(particles):
                 gs.raise_exception(
-                    f"Entity has particles outside solver boundary. Note that for MPMSolver, boundary is slightly tighter than the specified domain due to safety padding.\n\nCurrent boundary:\n{self._solver.boundary}\n\nEntity to be added:\nmin: {particles.min(0)}\nmax: {particles.max(0)}\n"
+                    "Entity has particles outside solver boundary. Note that for MPMSolver, boundary is slightly "
+                    "tighter than the specified domain due to safety padding.\n\n"
+                    "Current boundary:\n{self._solver.boundary}\n\nEntity to be added:\nmin: {particles.min(0)}\n"
+                    "max: {particles.max(0)}\n"
                 )
 
             if self._need_skinning:
-                self._vverts = np.array(self._vmesh.verts)
-                self._vfaces = np.array(self._vmesh.faces)
+                self._vverts = np.asarray(self._vmesh.verts, dtype=gs.np_float)
+                self._vfaces = np.asarray(self._vmesh.faces, dtype=gs.np_float)
             else:
-                self._vverts = np.array([])
-                self._vfaces = np.array([])
+                self._vverts = np.asarray([], dtype=gs.np_float)
+                self._vfaces = np.asarray([], dtype=gs.np_float)
+            origin = np.array(self._morph.pos, dtype=gs.np_float)
 
-        self._particles = particles.astype(gs.np_float, order="C", copy=False)
+        self._particles = particles
         self._init_particles_offset = gs.tensor(self._particles) - gs.tensor(origin)
         self._n_particles = len(self._particles)
 

genesis/engine/entities/pbd_entity.py

@@ -80,9 +80,9 @@ def sample(self):
     def _add_to_solver_(self):
         self._kernel_add_particles_edges_to_solver(
             f=self._scene.sim.cur_substep_local,
-            particles=self._particles.astype(gs.np_float),
-            edges=self._edges.astype(gs.np_int),
-            edges_len_rest=self._edges_len_rest.astype(gs.np_float),
+            particles=self._particles,
+            edges=self._edges,
+            edges_len_rest=self._edges_len_rest,
             mat_type=self._mat_type,
             active=True,
         )
@@ -331,16 +331,16 @@ def sample(self):
             gs.raise_exception("Input mesh has zero surface area.")
         self._mass = self._vmesh.area * self.material.rho
 
-        self._particles = np.array(self._mesh.verts)
-        self._edges = np.array(self._mesh.get_unique_edges())
+        self._particles = np.asarray(self._mesh.verts, dtype=gs.np_float)
+        self._edges = np.asarray(self._mesh.get_unique_edges(), dtype=gs.np_int)
 
         self._particle_mass = self._mass / len(self._particles)
 
         # Inner edges are two diagonal edges of each quadrilateral formed by adjacent face pairs
         adjacency, inner_edges = trimesh.graph.face_adjacency(mesh=self._mesh.trimesh, return_edges=True)
         v3 = np.sum(self._mesh.faces[adjacency[:, 0]], axis=1) - inner_edges[:, 0] - inner_edges[:, 1]
         v4 = np.sum(self._mesh.faces[adjacency[:, 1]], axis=1) - inner_edges[:, 0] - inner_edges[:, 1]
-        self._inner_edges = np.stack([inner_edges[:, 0], inner_edges[:, 1], v3, v4], axis=1)
+        self._inner_edges = np.stack([inner_edges[:, 0], inner_edges[:, 1], v3, v4], axis=1, dtype=gs.np_int)
 
         self._edges_len_rest = np.linalg.norm(
             self._particles[self._edges[:, 0]] - self._particles[self._edges[:, 1]], axis=1
@@ -359,8 +359,8 @@ def _add_to_solver_(self):
 
         self._kernel_add_inner_edges_to_solver(
             f=self._scene.sim.cur_substep_local,
-            inner_edges=self._inner_edges.astype(gs.np_int),
-            inner_edges_len_rest=self._inner_edges_len_rest.astype(gs.np_float),
+            inner_edges=self._inner_edges,
+            inner_edges_len_rest=self._inner_edges_len_rest,
         )
 
     @ti.kernel
@@ -469,34 +469,33 @@ def sample(self):
         self._mass = self._vmesh.volume * self.material.rho
 
         tet_cfg = mu.generate_tetgen_config_from_morph(self.morph)
-        self._particles, self._elems = self._mesh.tetrahedralize(tet_cfg)
+        particles, elems = self._mesh.tetrahedralize(tet_cfg)
+        self._particles = particles.astype(gs.np_float, copy=False)
+        self._elems = elems.astype(gs.np_int, copy=False)
         self._edges = np.array(
             list(
                 set(
-                    tuple(sorted([self._elems[i, j], self._elems[i, k]]))
-                    for i in range(self._elems.shape[0])
+                    tuple(sorted((self._elems[i, j], self._elems[i, k])))
+                    for i in range(len(self._elems))
                     for j in range(4)
                     for k in range(j + 1, 4)
                 )
-            )
+            ),
+            dtype=gs.np_int,
         )
         self._particle_mass = self._mass / len(self._particles)
 
         self._edges_len_rest = np.linalg.norm(
             self._particles[self._edges[:, 0]] - self._particles[self._edges[:, 1]], axis=1
         )
         self._elems_vol_rest = (
-            np.linalg.det(self._particles[self._elems[:, 1:]] - self._particles[self._elems[:, 0:1]]) / 6.0
+            np.linalg.det(self._particles[self._elems[:, 1:]] - self._particles[self._elems[:, :1]]) / 6.0
         )
         self._n_particles = len(self._particles)
 
     def _add_to_solver_(self):
         super()._add_to_solver_()
-
-        self._kernel_add_elems_to_solver(
-            elems=self._elems.astype(gs.np_int),
-            elems_vol_rest=self._elems_vol_rest.astype(gs.np_float),
-        )
+        self._kernel_add_elems_to_solver(elems=self._elems, elems_vol_rest=self._elems_vol_rest)
 
     @ti.kernel
     def _kernel_add_elems_to_solver(

genesis/engine/entities/rigid_entity/rigid_entity.py

@@ -272,7 +272,7 @@ def _load_mesh(self, morph, surface):
 
     def _load_terrain(self, morph, surface):
         vmesh, mesh, self.terrain_hf = tu.parse_terrain(morph, surface)
-        self.terrain_scale = np.array([morph.horizontal_scale, morph.vertical_scale])
+        self.terrain_scale = np.array((morph.horizontal_scale, morph.vertical_scale), dtype=gs.np_float)
 
         g_infos = []
         if morph.visualization: