Repository: Genesis-Embodied-AI/Genesis
File: genesis/utils/terrain.py (Python)

Code before:
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Parameters
    ----------
    path : str
        .glb / .obj / … file containing the terrain mesh.
    spacing : float  |  (float, float)
        Desired grid spacing  Δx  (and  Δy).  Units must match the mesh.
    oversample : int, default 3
        Number of extra rays per coarse‑grid axis.  k = 1 reproduces the
        simple single‑ray method; k ≥ 2 captures peaks inside a cell.
    up_axis : {"z", "y"}, default "z"
        Mesh’s up direction.  If "y" (common in glTF), the function
        rotates the mesh so Z is up before sampling.

    Returns
    -------
    heights : (nx, ny)  float32 ndarray
        Highest elevation per coarse cell (NaN if no hit).
    xs      : (nx,)     float32 ndarray
    ys      : (ny,)     float32 ndarray
        Coordinates of grid lines (cell centres).

    Notes
    -----
    • Requires `pyembree` or `rtree` for fast ray‑tracing.
    • Memory cost grows as  oversample².  For very large terrains keep
      `oversample` ≤ 4 unless you have ample RAM.
    """
    mesh = trimesh.load(path, force="mesh")

    # -------------------------------- axis handling ---------------------------
    if up_axis.lower() == "y":
        # rotate so Z becomes up (‑90° around X)
        T = trimesh.transformations.rotation_matrix(np.deg2rad(-90), [1, 0, 0])
        mesh.apply_transform(T)

    (minx, miny, _), (maxx, maxy, maxz) = mesh.bounds

    # -------------------------------- coarse grid ----------------------------
    dx, dy = (spacing, spacing) if isinstance(spacing, (int, float)) else spacing
    nx = math.ceil((maxx - minx) / dx) + 1
    ny = math.ceil((maxy - miny) / dy) + 1

    xs = np.linspace(minx, maxx, nx, dtype=np.float32)
    ys = np.linspace(miny, maxy, ny, dtype=np.float32)

    # -------------------------------- fine grid ------------------------------
    fx = nx * oversample
    fy = ny * oversample
    fx_lin = np.linspace(minx, maxx, fx, dtype=np.float32)

Review comment (refactor):
I would remove this comment entirely. `rtree` is already installed automatically as part of the requirements.

Code after:
    up_axis: str = "z",
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Parameters
    ----------
    path : str
        .glb / .obj / … file containing the terrain mesh.
    spacing : float  |  (float, float)
        Desired grid spacing  Δx  (and  Δy).  Units must match the mesh.
    oversample : int, default 3
        Number of extra rays per coarse‑grid axis.  k = 1 reproduces the
        simple single‑ray method; k ≥ 2 captures peaks inside a cell.
    up_axis : {"z", "y"}, default "z"
        Mesh’s up direction.  If "y" (common in glTF), the function
        rotates the mesh so Z is up before sampling.

    Returns
    -------
    heights : (nx, ny)  float32 ndarray
        Highest elevation per coarse cell (NaN if no hit).
    xs      : (nx,)     float32 ndarray
    ys      : (ny,)     float32 ndarray
        Coordinates of grid lines (cell centres).

    Notes
    -----
    • Memory cost grows as  oversample².
    """
    if np.isscalar(spacing):
        spacing = (spacing, spacing)

    mesh = trimesh.load(path, force="mesh")

    # -------------------------------- axis handling ---------------------------
    if up_axis.lower() == "y":
        # rotate so Z becomes up (‑90° around X)
        T = trimesh.transformations.rotation_matrix(np.deg2rad(-90), [1, 0, 0])
        mesh.apply_transform(T)

    (minx, miny, _), (maxx, maxy, maxz) = mesh.bounds

    # -------------------------------- coarse grid ----------------------------
    dx, dy = spacing
    nx = math.ceil((maxx - minx) / dx) + 1
    ny = math.ceil((maxy - miny) / dy) + 1

    xs = np.linspace(minx, maxx, nx, dtype=np.float32)
    ys = np.linspace(miny, maxy, ny, dtype=np.float32)

    # -------------------------------- fine grid ------------------------------
    fx = nx * oversample