def get_geom_vertices()

in robogym/envs/rearrange/common/utils.py [0:0]

• 30 lines of code
• 9 McCabe index (conditional complexity)

def get_geom_vertices(sim, geom_id):
    geom_type = sim.model.geom_type[geom_id]
    geom_size = sim.model.geom_size[geom_id]

    if geom_type == const.GEOM_BOX:
        dx, dy, dz = geom_size
        return np.array(list(itertools.product([dx, -dx], [dy, -dy], [dz, -dz])))
    elif geom_type in (const.GEOM_SPHERE, const.GEOM_ELLIPSOID):
        if geom_type == const.GEOM_SPHERE:
            r = [geom_size[0]] * 3
        else:
            r = geom_size[:3]

        # https://stats.stackexchange.com/a/30622
        vertices = []
        phi = np.linspace(0, np.pi * 2, 20)
        cos_theta = np.linspace(-1, 1, 20)
        for p, c in itertools.product(phi, cos_theta):
            x = np.sqrt(1 - c ** 2) * np.cos(p)
            y = np.sqrt(1 - c ** 2) * np.sin(p)
            z = c
            vertices.append(np.array([x, y, z]))

        return np.array(vertices) * r
    elif geom_type in (const.GEOM_CYLINDER, const.GEOM_CAPSULE):
        # We treat cylinder and capsule the same.
        r, h = geom_size[0], geom_size[2]
        points = np.array(
            [[r * np.cos(x), r * np.sin(x), 0.0] for x in np.linspace(0, np.pi * 2, 50)]
        )

        return np.concatenate([points + h, points - h])
    elif geom_type == const.GEOM_MESH:
        return sim.model.mesh_vert[mesh_vert_range_of_geom(sim, geom_id)]
    else:
        raise AssertionError(f"Unexpected geom type {geom_type}")