mlx3d.ops

mlx3d.ops

decimate_mesh(meshes, voxel_size)

Simplify a mesh by vertex clustering on a regular voxel grid.

Vertices in the same voxel_size cell collapse to their centroid; faces are remapped and degenerate ones (with a repeated vertex) dropped. Fast and robust — coarser than quadric-error decimation, but topology-agnostic.

Returns a new single-mesh :class:~mlx3d.structures.Meshes.

Source code in src/mlx3d/ops/geometry.py

def decimate_mesh(meshes: Meshes, voxel_size: float) -> Meshes:
    """Simplify a mesh by vertex clustering on a regular voxel grid.

    Vertices in the same ``voxel_size`` cell collapse to their centroid; faces
    are remapped and degenerate ones (with a repeated vertex) dropped. Fast and
    robust — coarser than quadric-error decimation, but topology-agnostic.

    Returns a new single-mesh :class:`~mlx3d.structures.Meshes`.
    """
    v = np.asarray(meshes.verts_packed(), dtype=np.float64)
    f = np.asarray(meshes.faces_packed())
    cells = np.floor(v / voxel_size).astype(np.int64)
    _, inv = np.unique(cells, axis=0, return_inverse=True)
    inv = inv.reshape(-1)

    n_new = int(inv.max()) + 1
    new_v = np.zeros((n_new, 3))
    counts = np.zeros(n_new)
    np.add.at(new_v, inv, v)
    np.add.at(counts, inv, 1)
    new_v /= counts[:, None]

    new_f = inv[f]  # (F, 3) remapped
    nondeg = (
        (new_f[:, 0] != new_f[:, 1]) & (new_f[:, 1] != new_f[:, 2]) & (new_f[:, 0] != new_f[:, 2])
    )
    new_f = new_f[nondeg]
    new_f = np.unique(np.sort(new_f, axis=1), axis=0) if new_f.shape[0] else new_f
    return Meshes([mx.array(new_v.astype(np.float32))], [mx.array(new_f.astype(np.int32))])

estimate_point_normals(points, k=16, orient_towards=None)

Estimate per-point normals by PCA over each point's k nearest neighbors.

The normal is the eigenvector of the local covariance with the smallest eigenvalue. Normal orientation is inherently ambiguous; pass orient_towards (e.g. the camera position) to flip normals consistently toward that point.

Parameters:

Name	Type	Description	Default
points	array	(P, 3) point cloud.	required
k	int	neighborhood size.	16
orient_towards	tuple[float, float, float] | array | None	optional (3,) location to orient normals toward.	None

Returns:

Type	Description
array	(P, 3) unit normals.

Source code in src/mlx3d/ops/geometry.py

def estimate_point_normals(
    points: mx.array,
    k: int = 16,
    orient_towards: tuple[float, float, float] | mx.array | None = None,
) -> mx.array:
    """Estimate per-point normals by PCA over each point's ``k`` nearest neighbors.

    The normal is the eigenvector of the local covariance with the smallest
    eigenvalue. Normal orientation is inherently ambiguous; pass
    ``orient_towards`` (e.g. the camera position) to flip normals consistently
    toward that point.

    Args:
        points: ``(P, 3)`` point cloud.
        k: neighborhood size.
        orient_towards: optional ``(3,)`` location to orient normals toward.

    Returns:
        ``(P, 3)`` unit normals.
    """
    pts = np.asarray(points, dtype=np.float64)
    _, idx = knn_points(mx.array(pts.astype(np.float32)), mx.array(pts.astype(np.float32)), K=k)
    nbrs = pts[np.asarray(idx)]  # (P, k, 3)
    centered = nbrs - nbrs.mean(axis=1, keepdims=True)
    cov = np.einsum("pki,pkj->pij", centered, centered) / k  # (P, 3, 3)
    _, vecs = np.linalg.eigh(cov)  # ascending eigenvalues
    normals = vecs[:, :, 0]  # smallest-eigenvalue direction
    if orient_towards is not None:
        target = np.asarray(orient_towards, dtype=np.float64).reshape(3)
        flip = ((target - pts) * normals).sum(-1) < 0
        normals[flip] *= -1.0
    normals /= np.maximum(np.linalg.norm(normals, axis=-1, keepdims=True), 1e-12)
    return mx.array(normals.astype(np.float32))

icp(source, target, iters=20, tol=1e-06)

Rigidly align source to target with Iterative Closest Point.

Each iteration matches every source point to its nearest target point and solves for the best rigid transform (Kabsch/SVD), accumulating the total R/t such that aligned = source @ R.T + t.

Parameters:

Name	Type	Description	Default
source	array	(N, 3) points to move.	required
target	array	(M, 3) reference points.	required
iters	int	maximum iterations.	20
tol	float	stop when the mean matched-distance improves by less than this.	1e-06

Returns:

Type	Description
dict[str, array]	dict with R (3, 3), t (3,), aligned (N, 3), and
dict[str, array]	rmse (final root-mean-square matched distance).

Source code in src/mlx3d/ops/geometry.py

def icp(
    source: mx.array,
    target: mx.array,
    iters: int = 20,
    tol: float = 1e-6,
) -> dict[str, mx.array]:
    """Rigidly align ``source`` to ``target`` with Iterative Closest Point.

    Each iteration matches every source point to its nearest target point and
    solves for the best rigid transform (Kabsch/SVD), accumulating the total
    ``R``/``t`` such that ``aligned = source @ R.T + t``.

    Args:
        source: ``(N, 3)`` points to move.
        target: ``(M, 3)`` reference points.
        iters: maximum iterations.
        tol: stop when the mean matched-distance improves by less than this.

    Returns:
        dict with ``R`` ``(3, 3)``, ``t`` ``(3,)``, ``aligned`` ``(N, 3)``, and
        ``rmse`` (final root-mean-square matched distance).
    """
    tgt = np.asarray(target, dtype=np.float64)
    cur = np.asarray(source, dtype=np.float64)
    R_tot = np.eye(3)
    t_tot = np.zeros(3)
    prev = np.inf
    rmse = np.inf
    tgt_mx = mx.array(tgt.astype(np.float32))
    for _ in range(iters):
        d2, idx = knn_points(mx.array(cur.astype(np.float32)), tgt_mx, K=1)
        matched = tgt[np.asarray(idx)[:, 0]]
        rmse = float(np.sqrt(np.asarray(d2)[:, 0].mean()))
        mu_s, mu_t = cur.mean(0), matched.mean(0)
        h = (cur - mu_s).T @ (matched - mu_t)
        u, _, vt = np.linalg.svd(h)
        d = np.sign(np.linalg.det(vt.T @ u.T))
        r_step = vt.T @ np.diag([1.0, 1.0, d]) @ u.T
        t_step = mu_t - r_step @ mu_s
        cur = cur @ r_step.T + t_step
        R_tot = r_step @ R_tot
        t_tot = r_step @ t_tot + t_step
        if abs(prev - rmse) < tol:
            break
        prev = rmse
    return {
        "R": mx.array(R_tot.astype(np.float32)),
        "t": mx.array(t_tot.astype(np.float32)),
        "aligned": mx.array(cur.astype(np.float32)),
        "rmse": mx.array(float(rmse)),
    }

poisson_reconstruction(points, normals, resolution=64, padding=0.1)

Reconstruct a watertight surface mesh from oriented points (Poisson).

Solves the Poisson equation on a regular grid: the oriented points define a vector field whose divergence drives an indicator function chi (via an FFT Poisson solve), then the surface is extracted with marching cubes at the iso-level passing through the input points. This is the regular-grid form of Screened Poisson Surface Reconstruction (Kazhdan et al.) -- the same physics without an octree.

Parameters:

Name	Type	Description	Default
points	array	(P, 3) surface samples.	required
normals	array	(P, 3) oriented normals (need not be unit; consistent orientation matters, so estimate/orient them first).	required
resolution	int	grid cells per axis.	64
padding	float	fractional padding added around the point bounding box.	0.1

Returns:

Type	Description
Meshes	A single-mesh :class:~mlx3d.structures.Meshes.

Source code in src/mlx3d/ops/geometry.py

def poisson_reconstruction(
    points: mx.array,
    normals: mx.array,
    resolution: int = 64,
    padding: float = 0.1,
) -> Meshes:
    """Reconstruct a watertight surface mesh from oriented points (Poisson).

    Solves the Poisson equation on a regular grid: the oriented points define a
    vector field whose divergence drives an indicator function ``chi`` (via an
    FFT Poisson solve), then the surface is extracted with marching cubes at the
    iso-level passing through the input points. This is the regular-grid form of
    Screened Poisson Surface Reconstruction (Kazhdan et al.) -- the same physics
    without an octree.

    Args:
        points: ``(P, 3)`` surface samples.
        normals: ``(P, 3)`` oriented normals (need not be unit; consistent
            orientation matters, so estimate/orient them first).
        resolution: grid cells per axis.
        padding: fractional padding added around the point bounding box.

    Returns:
        A single-mesh :class:`~mlx3d.structures.Meshes`.
    """
    pts = np.asarray(points, dtype=np.float64)
    nrm = np.asarray(normals, dtype=np.float64)
    nrm = nrm / np.maximum(np.linalg.norm(nrm, axis=-1, keepdims=True), 1e-12)

    lo = pts.min(0)
    hi = pts.max(0)
    span = (hi - lo).max()
    pad = span * padding
    lo = lo - pad
    hi = lo + (span + 2 * pad)  # cube bounds
    size = hi - lo
    res = int(resolution)

    # Grid coordinates of each point in [0, res).
    g = (pts - lo) / size * res
    i0 = np.clip(np.floor(g).astype(np.int64), 0, res - 1)
    frac = g - i0

    # Trilinearly splat normals onto a (res, res, res, 3) vector field.
    vfield = np.zeros((res, res, res, 3))
    for dx in (0, 1):
        for dy in (0, 1):
            for dz in (0, 1):
                ix = np.clip(i0[:, 0] + dx, 0, res - 1)
                iy = np.clip(i0[:, 1] + dy, 0, res - 1)
                iz = np.clip(i0[:, 2] + dz, 0, res - 1)
                w = (
                    (frac[:, 0] if dx else 1 - frac[:, 0])
                    * (frac[:, 1] if dy else 1 - frac[:, 1])
                    * (frac[:, 2] if dz else 1 - frac[:, 2])
                )
                np.add.at(vfield, (ix, iy, iz), w[:, None] * nrm)

    # Divergence of the field, then solve Laplacian(chi) = div via FFT.
    div = (
        np.gradient(vfield[..., 0], axis=0)
        + np.gradient(vfield[..., 1], axis=1)
        + np.gradient(vfield[..., 2], axis=2)
    )
    k = 2.0 * np.pi * np.fft.fftfreq(res)
    kx, ky, kz = np.meshgrid(k, k, k, indexing="ij")
    denom = -(kx**2 + ky**2 + kz**2)
    denom[0, 0, 0] = 1.0  # avoid divide-by-zero for the DC term
    chi_hat = np.fft.fftn(div) / denom
    chi_hat[0, 0, 0] = 0.0
    chi = np.real(np.fft.ifftn(chi_hat)).astype(np.float32)

    # Iso-level: the indicator value the surface (the input points) sits at.
    iso = float(chi[i0[:, 0], i0[:, 1], i0[:, 2]].mean())
    spacing = tuple((size / res).astype(float))
    return marching_cubes(mx.array(chi), level=iso, spacing=spacing, origin=tuple(lo.astype(float)))

knn_gather(x, idx)

Gather neighbor features: x (N, P2, C), idx (N, P1, K) -> (N, P1, K, C).

Source code in src/mlx3d/ops/knn.py

def knn_gather(x: mx.array, idx: mx.array) -> mx.array:
    """Gather neighbor features: ``x`` (N, P2, C), ``idx`` (N, P1, K) -> (N, P1, K, C)."""
    N, P1, K = idx.shape
    C = x.shape[-1]
    flat = idx.reshape(N, P1 * K)
    gathered = mx.take_along_axis(x, flat[..., None].astype(mx.int32), axis=1)
    return gathered.reshape(N, P1, K, C)

knn_points(p1, p2, K=1, chunk_size=4096)

For each point in p1 find its K nearest neighbors in p2.

Parameters:

Name	Type	Description	Default
p1	array	(N, P1, D) or (P1, D) query points.	required
p2	array	(N, P2, D) or (P2, D) reference points.	required
K	int	number of neighbors.	1
chunk_size	int	queries processed per tile to bound the (P1, P2) distance matrix memory.	4096

Returns:

Type	Description
array	(dists, idx) of shapes (..., P1, K): squared distances and indices
array	into p2, sorted by increasing distance.

Source code in src/mlx3d/ops/knn.py

def knn_points(
    p1: mx.array,
    p2: mx.array,
    K: int = 1,
    chunk_size: int = 4096,
) -> tuple[mx.array, mx.array]:
    """For each point in ``p1`` find its ``K`` nearest neighbors in ``p2``.

    Args:
        p1: (N, P1, D) or (P1, D) query points.
        p2: (N, P2, D) or (P2, D) reference points.
        K: number of neighbors.
        chunk_size: queries processed per tile to bound the (P1, P2) distance
            matrix memory.

    Returns:
        ``(dists, idx)`` of shapes (..., P1, K): squared distances and indices
        into ``p2``, sorted by increasing distance.
    """
    if K < 1:
        raise ValueError("K must be at least 1.")
    if p2.shape[-2] == 0:
        raise ValueError("p2 must contain at least one point.")
    K = min(K, p2.shape[-2])
    if _can_use_metal_knn(p1, p2, K):
        return _knn_points_3d_smallk(p1, p2, K)

    squeeze = p1.ndim == 2
    if squeeze:
        p1, p2 = p1[None], p2[None]
    P1, P2 = p1.shape[1], p2.shape[1]

    dists_out, idx_out = [], []
    for start in range(0, P1, chunk_size):
        d = _pairwise_sqdist(p1[:, start : start + chunk_size], p2)  # (N, c, P2)
        if K == 1:
            idx = mx.argmin(d, axis=-1, keepdims=True)
            dist = mx.take_along_axis(d, idx, axis=-1)
        elif K < P2:
            # argpartition + small sort instead of a full argsort of the
            # (chunk, P2) matrix, which would allocate K-independent memory.
            part = mx.argpartition(d, kth=K - 1, axis=-1)[..., :K]
            dist_part = mx.take_along_axis(d, part, axis=-1)
            order = mx.argsort(dist_part, axis=-1)
            idx = mx.take_along_axis(part, order, axis=-1)
            dist = mx.take_along_axis(dist_part, order, axis=-1)
        else:
            idx = mx.argsort(d, axis=-1)
            dist = mx.take_along_axis(d, idx, axis=-1)
        dists_out.append(dist)
        idx_out.append(idx)

    dists = mx.concatenate(dists_out, axis=1)
    idx = mx.concatenate(idx_out, axis=1)
    if squeeze:
        dists, idx = dists[0], idx[0]
    return dists, idx

ray_mesh_intersect(meshes, origins, directions, face_chunk_size=4096, eps=1e-08, aabb_cull=True, spatial_sort=False, return_stats=False)

Nearest ray-triangle hit per ray (Möller-Trumbore).

Parameters:

Name	Type	Description	Default
meshes	Meshes	a single-mesh :class:~mlx3d.structures.Meshes.	required
origins	array	(R, 3) ray origins.	required
directions	array	(R, 3) ray directions (need not be normalized; t is measured in units of directions).	required
face_chunk_size	int	faces per chunk, bounding the (R, F) memory.	4096
eps	float	parallel / degenerate-triangle tolerance.	1e-08
aabb_cull	bool	if True, test each face chunk's axis-aligned bounding box before running the expensive ray-triangle math. This preserves exact results and skips whole chunks for coherent rays or spatially sorted meshes.	True
spatial_sort	bool	sort faces by Morton code before chunking. This one-time CPU preprocessing step is useful for arbitrary face orderings and maps returned face_idx values back to the original packed faces.	False
return_stats	bool	include simple broad-phase counters useful for tests and profiling.	False

Returns:

Type	Description
dict[str, object]	dict with hit (R,) bool, t (R,) hit distance (inf on
dict[str, object]	miss), face_idx (R,) int32 (-1 on miss), bary (R, 3)
dict[str, object]	barycentric coords, and points (R, 3) world hit positions. The
dict[str, object]	hit position is differentiable w.r.t. the mesh vertices.

Source code in src/mlx3d/ops/ray_mesh.py

def ray_mesh_intersect(
    meshes: Meshes,
    origins: mx.array,
    directions: mx.array,
    face_chunk_size: int = 4096,
    eps: float = 1e-8,
    aabb_cull: bool = True,
    spatial_sort: bool = False,
    return_stats: bool = False,
) -> dict[str, object]:
    """Nearest ray-triangle hit per ray (Möller-Trumbore).

    Args:
        meshes: a single-mesh :class:`~mlx3d.structures.Meshes`.
        origins: ``(R, 3)`` ray origins.
        directions: ``(R, 3)`` ray directions (need not be normalized; ``t`` is
            measured in units of ``directions``).
        face_chunk_size: faces per chunk, bounding the ``(R, F)`` memory.
        eps: parallel / degenerate-triangle tolerance.
        aabb_cull: if ``True``, test each face chunk's axis-aligned bounding
            box before running the expensive ray-triangle math. This preserves
            exact results and skips whole chunks for coherent rays or spatially
            sorted meshes.
        spatial_sort: sort faces by Morton code before chunking. This one-time
            CPU preprocessing step is useful for arbitrary face orderings and
            maps returned ``face_idx`` values back to the original packed faces.
        return_stats: include simple broad-phase counters useful for tests and
            profiling.

    Returns:
        dict with ``hit`` ``(R,)`` bool, ``t`` ``(R,)`` hit distance (``inf`` on
        miss), ``face_idx`` ``(R,)`` int32 (``-1`` on miss), ``bary`` ``(R, 3)``
        barycentric coords, and ``points`` ``(R, 3)`` world hit positions. The
        hit position is differentiable w.r.t. the mesh vertices.
    """
    face_remap = None
    if spatial_sort:
        meshes, face_remap = spatially_sort_faces(meshes)
    verts = meshes.verts_packed()
    faces = meshes.faces_packed().astype(mx.int32)
    o = origins[:, None, :]  # (R, 1, 3)
    d = directions[:, None, :]  # (R, 1, 3)

    best_t = mx.full((origins.shape[0],), _INF)
    best_f = mx.full((origins.shape[0],), -1, dtype=mx.int32)
    best_uv = mx.zeros((origins.shape[0], 2))
    chunks_total = 0
    chunks_skipped = 0
    face_tests = 0

    for start in range(0, faces.shape[0], face_chunk_size):
        fchunk = faces[start : start + face_chunk_size]  # (C, 3)
        tri = verts[fchunk]  # (C, 3, 3)
        chunks_total += 1
        if aabb_cull:
            bmin = mx.min(tri.reshape(-1, 3), axis=0)
            bmax = mx.max(tri.reshape(-1, 3), axis=0)
            chunk_hit = _rays_intersect_aabb(origins, directions, bmin, bmax, eps)
            if not bool(mx.any(chunk_hit)):
                chunks_skipped += 1
                continue
        face_tests += int(origins.shape[0]) * int(fchunk.shape[0])
        v0, v1, v2 = tri[:, 0, :], tri[:, 1, :], tri[:, 2, :]
        e1 = (v1 - v0)[None]  # (1, C, 3)
        e2 = (v2 - v0)[None]
        pvec = mx.linalg.cross(d, e2)  # (R, C, 3)
        det = mx.sum(e1 * pvec, axis=-1)  # (R, C)
        valid = mx.abs(det) > eps
        inv_det = 1.0 / mx.where(valid, det, mx.ones_like(det))

        tvec = o - v0[None]  # (R, C, 3)
        u = mx.sum(tvec * pvec, axis=-1) * inv_det
        qvec = mx.linalg.cross(tvec, e1)
        v = mx.sum(d * qvec, axis=-1) * inv_det
        t = mx.sum(e2 * qvec, axis=-1) * inv_det

        hit = valid & (u >= 0.0) & (v >= 0.0) & (u + v <= 1.0) & (t > eps)
        t_hit = mx.where(hit, t, mx.full(t.shape, _INF))  # (R, C)

        local = mx.argmin(t_hit, axis=-1)  # (R,)
        local_t = mx.take_along_axis(t_hit, local[:, None], axis=-1)[:, 0]
        closer = local_t < best_t
        best_t = mx.where(closer, local_t, best_t)
        best_f = mx.where(closer, (start + local).astype(mx.int32), best_f)
        cu = mx.take_along_axis(u, local[:, None], axis=-1)[:, 0]
        cv = mx.take_along_axis(v, local[:, None], axis=-1)[:, 0]
        best_uv = mx.where(closer[:, None], mx.stack([cu, cv], axis=-1), best_uv)

    hit_mask = best_f >= 0
    if face_remap is not None and face_remap.shape[0] > 0:
        safe_f = mx.where(hit_mask, best_f, 0)
        best_f = mx.where(hit_mask, face_remap[safe_f], best_f)
    u, v = best_uv[:, 0], best_uv[:, 1]
    bary = mx.stack([1.0 - u - v, u, v], axis=-1) * hit_mask[:, None]
    t_out = mx.where(hit_mask, best_t, mx.full(best_t.shape, mx.inf))
    points = origins + mx.where(hit_mask, best_t, 0.0)[:, None] * directions
    out = {
        "hit": hit_mask,
        "t": t_out,
        "face_idx": best_f,
        "bary": bary,
        "points": points,
    }
    if return_stats:
        out["stats"] = {
            "chunks_total": chunks_total,
            "chunks_skipped": chunks_skipped,
            "face_tests": face_tests,
            "spatial_sort": bool(spatial_sort),
        }
    return out

spatially_sort_faces(meshes)

Return a copy with faces sorted by Morton code of their centroids.

Spatial ordering makes adjacent face chunks geometrically coherent, which improves the effectiveness of chunk-level AABB culling in :func:ray_mesh_intersect. The second return value maps sorted face rows back to original packed face indices.

Source code in src/mlx3d/ops/ray_mesh.py

def spatially_sort_faces(meshes: Meshes) -> tuple[Meshes, mx.array]:
    """Return a copy with faces sorted by Morton code of their centroids.

    Spatial ordering makes adjacent face chunks geometrically coherent, which
    improves the effectiveness of chunk-level AABB culling in
    :func:`ray_mesh_intersect`. The second return value maps sorted face rows
    back to original packed face indices.
    """
    verts_list = meshes.verts_list()
    sorted_faces = []
    inverse_parts = []
    face_offset = 0
    for verts, faces in zip(verts_list, meshes.faces_list(), strict=True):
        if faces.shape[0] == 0:
            sorted_faces.append(faces)
            continue
        v_np = np.asarray(verts)
        f_np = np.asarray(faces, dtype=np.int32)
        centroids = v_np[f_np].mean(axis=1)
        order = np.argsort(_morton_codes(centroids), kind="stable").astype(np.int32)
        sorted_faces.append(mx.array(f_np[order]))
        inverse_parts.append(mx.array(order + face_offset, dtype=mx.int32))
        face_offset += faces.shape[0]
    inverse = (
        mx.concatenate(inverse_parts, axis=0) if inverse_parts else mx.zeros((0,), dtype=mx.int32)
    )
    return Meshes(verts_list, sorted_faces), inverse

sample_points_from_meshes(meshes, num_samples=10000, return_normals=False)

Uniformly sample points from mesh surfaces (area-weighted).

Sampling locations are differentiable with respect to vertex positions (the face choice and barycentric coordinates are treated as constants, as in PyTorch3D).

Parameters:

Name	Type	Description	Default
meshes	Meshes	a :class:Meshes batch of size N.	required
num_samples	int	points per mesh.	10000
return_normals	bool	also return the face normal at each sample.	False

Returns:

Type	Description
array | tuple[array, array]	An (N, num_samples, 3) array of samples, plus an (N, num_samples, 3)
array | tuple[array, array]	array of normals if return_normals is set.

Source code in src/mlx3d/ops/sample_points.py

def sample_points_from_meshes(
    meshes: Meshes,
    num_samples: int = 10000,
    return_normals: bool = False,
) -> mx.array | tuple[mx.array, mx.array]:
    """Uniformly sample points from mesh surfaces (area-weighted).

    Sampling locations are differentiable with respect to vertex positions
    (the face choice and barycentric coordinates are treated as constants,
    as in PyTorch3D).

    Args:
        meshes: a :class:`Meshes` batch of size N.
        num_samples: points per mesh.
        return_normals: also return the face normal at each sample.

    Returns:
        An (N, num_samples, 3) array of samples, plus an (N, num_samples, 3)
        array of normals if ``return_normals`` is set.
    """
    verts = meshes.verts_packed()
    faces = meshes.faces_packed()
    areas = meshes.faces_areas_packed()
    first_idx = meshes.mesh_to_faces_packed_first_idx().tolist()
    num_faces = meshes.num_faces_per_mesh.tolist()

    samples = []
    normals = []
    for i in range(len(meshes)):
        f0, nf = first_idx[i], num_faces[i]
        if nf == 0:
            samples.append(mx.zeros((num_samples, 3)))
            if return_normals:
                normals.append(mx.zeros((num_samples, 3)))
            continue
        a = areas[f0 : f0 + nf]
        logits = mx.log(mx.maximum(a, 1e-12))
        face_idx = mx.random.categorical(logits[None], num_samples=num_samples)[0]
        tri = faces[f0 + face_idx]  # (S, 3)
        v0, v1, v2 = verts[tri[:, 0]], verts[tri[:, 1]], verts[tri[:, 2]]

        # Uniform barycentric sampling via the square-root trick.
        u = mx.random.uniform(shape=(num_samples, 1))
        v = mx.random.uniform(shape=(num_samples, 1))
        su = mx.sqrt(u)
        w0 = 1.0 - su
        w1 = su * (1.0 - v)
        w2 = su * v
        samples.append(w0 * v0 + w1 * v1 + w2 * v2)

        if return_normals:
            n = mx.linalg.cross(v1 - v0, v2 - v0)
            n = n / mx.maximum(mx.linalg.norm(n, axis=-1, keepdims=True), 1e-12)
            normals.append(n)

    out = mx.stack(samples, axis=0)
    if return_normals:
        return out, mx.stack(normals, axis=0)
    return out

subdivide_meshes(meshes)

Subdivide each triangle into four by inserting edge midpoints.

Each face (v0, v1, v2) becomes four faces using the midpoints of its three edges; midpoints are shared between adjacent faces (so the result is watertight where the input was). New vertex positions are the differentiable averages of their two endpoints, so gradients flow back to the originals.

Operates on a single mesh and returns a new :class:~mlx3d.structures.Meshes.

Source code in src/mlx3d/ops/subdivide.py

def subdivide_meshes(meshes: Meshes) -> Meshes:
    """Subdivide each triangle into four by inserting edge midpoints.

    Each face ``(v0, v1, v2)`` becomes four faces using the midpoints of its
    three edges; midpoints are shared between adjacent faces (so the result is
    watertight where the input was). New vertex positions are the differentiable
    averages of their two endpoints, so gradients flow back to the originals.

    Operates on a single mesh and returns a new :class:`~mlx3d.structures.Meshes`.
    """
    if len(meshes) != 1:
        raise ValueError("subdivide_meshes operates on one mesh at a time.")
    verts = meshes.verts_packed()
    faces = np.asarray(meshes.faces_packed())
    if faces.shape[0] == 0:
        return Meshes([verts], [meshes.faces_packed()])

    v = faces.shape[0]
    # Per-corner edges (canonicalized) -> unique edge ids, in face-corner order.
    corner_edges = np.concatenate(
        [faces[:, [0, 1]], faces[:, [1, 2]], faces[:, [2, 0]]], axis=0
    )  # (3F, 2)
    keyed = np.sort(corner_edges, axis=1)
    unique_edges, inverse = np.unique(keyed, axis=0, return_inverse=True)
    inverse = inverse.reshape(3, v).T  # (F, 3): edge id for (v0v1, v1v2, v2v0)

    num_verts = int(verts.shape[0])
    # New vertices = original verts followed by edge midpoints (differentiable).
    edges = mx.array(unique_edges.astype(np.int32))
    midpoints = 0.5 * (verts[edges[:, 0]] + verts[edges[:, 1]])
    new_verts = mx.concatenate([verts, midpoints], axis=0)

    # Midpoint vertex index for each face corner-edge.
    m = inverse + num_verts  # (F, 3)
    v0, v1, v2 = faces[:, 0], faces[:, 1], faces[:, 2]
    m01, m12, m20 = m[:, 0], m[:, 1], m[:, 2]
    new_faces = np.concatenate(
        [
            np.stack([v0, m01, m20], axis=1),
            np.stack([v1, m12, m01], axis=1),
            np.stack([v2, m20, m12], axis=1),
            np.stack([m01, m12, m20], axis=1),
        ],
        axis=0,
    ).astype(np.int32)

    return Meshes([new_verts], [mx.array(new_faces)])