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

sample_sdf_grid(sdf_fn, resolution=64, bounds=1.5)

Evaluate an SDF on a regular grid, ready for :func:marching_cubes.

Parameters:

Name	Type	Description	Default
sdf_fn	Callable[[array], array]	callable mapping (N, 3) points to (N,) signed distances.	required
resolution	int	number of grid samples per axis.	64
bounds	float | tuple[float, float] | tuple[tuple[float, float], ...]	sampling extent. A scalar b gives [-b, b]^3; a (min, max) pair applies to all axes; three (min, max) pairs set per-axis extents.	1.5

Returns:

Type	Description
array	(volume, spacing, origin) where volume is a (R, R, R) grid
tuple[float, float, float]	of signed distances indexed [x, y, z], and spacing / origin
tuple[float, float, float]	are the voxel size and grid origin to pass straight to
tuple[array, tuple[float, float, float], tuple[float, float, float]]	func:marching_cubes.

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

def sample_sdf_grid(
    sdf_fn: Callable[[mx.array], mx.array],
    resolution: int = 64,
    bounds: float | tuple[float, float] | tuple[tuple[float, float], ...] = 1.5,
) -> tuple[mx.array, tuple[float, float, float], tuple[float, float, float]]:
    """Evaluate an SDF on a regular grid, ready for :func:`marching_cubes`.

    Args:
        sdf_fn: callable mapping ``(N, 3)`` points to ``(N,)`` signed distances.
        resolution: number of grid samples per axis.
        bounds: sampling extent. A scalar ``b`` gives ``[-b, b]^3``; a
            ``(min, max)`` pair applies to all axes; three ``(min, max)`` pairs
            set per-axis extents.

    Returns:
        ``(volume, spacing, origin)`` where ``volume`` is a ``(R, R, R)`` grid
        of signed distances indexed ``[x, y, z]``, and ``spacing`` / ``origin``
        are the voxel size and grid origin to pass straight to
        :func:`marching_cubes`.
    """
    if resolution < 2:
        raise ValueError("resolution must be at least 2.")
    lo, hi = _resolve_bounds(bounds)
    xs = mx.linspace(float(lo[0]), float(hi[0]), resolution)
    ys = mx.linspace(float(lo[1]), float(hi[1]), resolution)
    zs = mx.linspace(float(lo[2]), float(hi[2]), resolution)
    gx, gy, gz = mx.meshgrid(xs, ys, zs, indexing="ij")
    pts = mx.stack([gx, gy, gz], axis=-1).reshape(-1, 3)
    volume = sdf_fn(pts).reshape(resolution, resolution, resolution)
    spacing = tuple(float((hi[i] - lo[i]) / (resolution - 1)) for i in range(3))
    origin = (float(lo[0]), float(lo[1]), float(lo[2]))
    return volume, spacing, origin

sdf_box(points, half_extents=(1.0, 1.0, 1.0), center=(0.0, 0.0, 0.0))

Signed distance to an axis-aligned box.

Parameters:

Name	Type	Description	Default
points	array	(..., 3) query positions.	required
half_extents	tuple[float, float, float] | array	(3,) half-sizes of the box along each axis.	(1.0, 1.0, 1.0)
center	tuple[float, float, float] | array	(3,) box center.	(0.0, 0.0, 0.0)

Returns:

Type	Description
array	(...) signed distances (negative inside).

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

def sdf_box(
    points: mx.array,
    half_extents: tuple[float, float, float] | mx.array = (1.0, 1.0, 1.0),
    center: tuple[float, float, float] | mx.array = (0.0, 0.0, 0.0),
) -> mx.array:
    """Signed distance to an axis-aligned box.

    Args:
        points: ``(..., 3)`` query positions.
        half_extents: ``(3,)`` half-sizes of the box along each axis.
        center: ``(3,)`` box center.

    Returns:
        ``(...)`` signed distances (negative inside).
    """
    q = mx.abs(points - _as_vec3(center)) - _as_vec3(half_extents)
    outside = mx.linalg.norm(mx.maximum(q, 0.0), axis=-1)
    inside = mx.minimum(mx.max(q, axis=-1), 0.0)
    return outside + inside

sdf_difference(distance_a, distance_b)

Subtract shape b from shape a (a minus b).

Parameters:

Name	Type	Description	Default
distance_a	array	distances of the shape to keep.	required
distance_b	array	distances of the shape to carve away.	required

Returns:

Type	Description
array	Combined signed distances.

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

def sdf_difference(distance_a: mx.array, distance_b: mx.array) -> mx.array:
    """Subtract shape ``b`` from shape ``a`` (``a`` minus ``b``).

    Args:
        distance_a: distances of the shape to keep.
        distance_b: distances of the shape to carve away.

    Returns:
        Combined signed distances.
    """
    return mx.maximum(distance_a, -distance_b)

sdf_intersection(*distances)

Boolean intersection of shapes: the pointwise maximum of their distances.

Parameters:

Name	Type	Description	Default
*distances	array	two or more broadcastable distance arrays.	()

Returns:

Type	Description
array	Combined signed distances.

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

def sdf_intersection(*distances: mx.array) -> mx.array:
    """Boolean intersection of shapes: the pointwise maximum of their distances.

    Args:
        *distances: two or more broadcastable distance arrays.

    Returns:
        Combined signed distances.
    """
    if len(distances) < 2:
        raise ValueError("sdf_intersection needs at least two distance arrays.")
    out = distances[0]
    for d in distances[1:]:
        out = mx.maximum(out, d)
    return out

sdf_plane(points, normal=(0.0, 1.0, 0.0), offset=0.0)

Signed distance to a plane dot(p, n) = offset.

Parameters:

Name	Type	Description	Default
points	array	(..., 3) query positions.	required
normal	tuple[float, float, float] | array	plane normal (need not be unit length).	(0.0, 1.0, 0.0)
offset	float	signed distance of the plane from the origin along normal.	0.0

Returns:

Type	Description
array	(...) signed distances (negative on the side opposite normal).

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

def sdf_plane(
    points: mx.array,
    normal: tuple[float, float, float] | mx.array = (0.0, 1.0, 0.0),
    offset: float = 0.0,
) -> mx.array:
    """Signed distance to a plane ``dot(p, n) = offset``.

    Args:
        points: ``(..., 3)`` query positions.
        normal: plane normal (need not be unit length).
        offset: signed distance of the plane from the origin along ``normal``.

    Returns:
        ``(...)`` signed distances (negative on the side opposite ``normal``).
    """
    n = _as_vec3(normal)
    n = n / mx.linalg.norm(n)
    return mx.sum(points * n, axis=-1) - offset

sdf_smooth_difference(distance_a, distance_b, k=0.1)

Smooth (blended) subtraction of shape b from shape a.

Parameters:

Name	Type	Description	Default
distance_a	array	distances of the shape to keep.	required
distance_b	array	distances of the shape to carve away.	required
k	float	blend radius; larger values round the join more.	0.1

Returns:

Type	Description
array	Combined signed distances.

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

def sdf_smooth_difference(distance_a: mx.array, distance_b: mx.array, k: float = 0.1) -> mx.array:
    """Smooth (blended) subtraction of shape ``b`` from shape ``a``.

    Args:
        distance_a: distances of the shape to keep.
        distance_b: distances of the shape to carve away.
        k: blend radius; larger values round the join more.

    Returns:
        Combined signed distances.
    """
    return sdf_smooth_intersection(distance_a, -distance_b, k)

sdf_smooth_intersection(distance_a, distance_b, k=0.1)

Smooth (blended) intersection.

Parameters:

Name	Type	Description	Default
distance_a	array	first shape's distances.	required
distance_b	array	second shape's distances.	required
k	float	blend radius; larger values round the join more.	0.1

Returns:

Type	Description
array	Combined signed distances.

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

def sdf_smooth_intersection(distance_a: mx.array, distance_b: mx.array, k: float = 0.1) -> mx.array:
    """Smooth (blended) intersection.

    Args:
        distance_a: first shape's distances.
        distance_b: second shape's distances.
        k: blend radius; larger values round the join more.

    Returns:
        Combined signed distances.
    """
    return -sdf_smooth_union(-distance_a, -distance_b, k)

sdf_smooth_union(distance_a, distance_b, k=0.1)

Smooth (blended) union, after Inigo Quilez's polynomial smin.

Parameters:

Name	Type	Description	Default
distance_a	array	first shape's distances.	required
distance_b	array	second shape's distances.	required
k	float	blend radius; larger values round the join more. As k -> 0 this approaches :func:sdf_union.	0.1

Returns:

Type	Description
array	Combined signed distances.

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

def sdf_smooth_union(distance_a: mx.array, distance_b: mx.array, k: float = 0.1) -> mx.array:
    """Smooth (blended) union, after Inigo Quilez's polynomial smin.

    Args:
        distance_a: first shape's distances.
        distance_b: second shape's distances.
        k: blend radius; larger values round the join more. As ``k -> 0`` this
            approaches :func:`sdf_union`.

    Returns:
        Combined signed distances.
    """
    h = mx.clip(0.5 + 0.5 * (distance_b - distance_a) / k, 0.0, 1.0)
    return distance_b * (1.0 - h) + distance_a * h - k * h * (1.0 - h)

sdf_sphere(points, radius=1.0, center=(0.0, 0.0, 0.0))

Signed distance to a sphere.

Parameters:

Name	Type	Description	Default
points	array	(..., 3) query positions.	required
radius	float	sphere radius.	1.0
center	tuple[float, float, float] | array	(3,) sphere center.	(0.0, 0.0, 0.0)

Returns:

Type	Description
array	(...) signed distances (negative inside).

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

def sdf_sphere(
    points: mx.array,
    radius: float = 1.0,
    center: tuple[float, float, float] | mx.array = (0.0, 0.0, 0.0),
) -> mx.array:
    """Signed distance to a sphere.

    Args:
        points: ``(..., 3)`` query positions.
        radius: sphere radius.
        center: ``(3,)`` sphere center.

    Returns:
        ``(...)`` signed distances (negative inside).
    """
    return mx.linalg.norm(points - _as_vec3(center), axis=-1) - radius

sdf_to_mesh(sdf_fn, resolution=64, bounds=1.5, level=0.0)

Extract a triangle mesh from an SDF callable.

Convenience wrapper that samples sdf_fn on a grid with :func:sample_sdf_grid and runs :func:marching_cubes at level.

Parameters:

Name	Type	Description	Default
sdf_fn	Callable[[array], array]	callable mapping (N, 3) points to (N,) signed distances.	required
resolution	int	grid samples per axis.	64
bounds	float | tuple[float, float] | tuple[tuple[float, float], ...]	sampling extent (see :func:sample_sdf_grid).	1.5
level	float	isovalue to extract (0 is the surface).	0.0

Returns:

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

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

def sdf_to_mesh(
    sdf_fn: Callable[[mx.array], mx.array],
    resolution: int = 64,
    bounds: float | tuple[float, float] | tuple[tuple[float, float], ...] = 1.5,
    level: float = 0.0,
) -> Meshes:
    """Extract a triangle mesh from an SDF callable.

    Convenience wrapper that samples ``sdf_fn`` on a grid with
    :func:`sample_sdf_grid` and runs :func:`marching_cubes` at ``level``.

    Args:
        sdf_fn: callable mapping ``(N, 3)`` points to ``(N,)`` signed distances.
        resolution: grid samples per axis.
        bounds: sampling extent (see :func:`sample_sdf_grid`).
        level: isovalue to extract (``0`` is the surface).

    Returns:
        A single-mesh :class:`~mlx3d.structures.Meshes`.
    """
    volume, spacing, origin = sample_sdf_grid(sdf_fn, resolution, bounds)
    return marching_cubes(volume, level=level, spacing=spacing, origin=origin)

sdf_torus(points, major_radius=1.0, minor_radius=0.25, center=(0.0, 0.0, 0.0))

Signed distance to a torus lying in the xz-plane (y is the axis).

Parameters:

Name	Type	Description	Default
points	array	(..., 3) query positions.	required
major_radius	float	distance from the center to the tube center.	1.0
minor_radius	float	tube radius.	0.25
center	tuple[float, float, float] | array	(3,) torus center.	(0.0, 0.0, 0.0)

Returns:

Type	Description
array	(...) signed distances (negative inside the tube).

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

def sdf_torus(
    points: mx.array,
    major_radius: float = 1.0,
    minor_radius: float = 0.25,
    center: tuple[float, float, float] | mx.array = (0.0, 0.0, 0.0),
) -> mx.array:
    """Signed distance to a torus lying in the xz-plane (y is the axis).

    Args:
        points: ``(..., 3)`` query positions.
        major_radius: distance from the center to the tube center.
        minor_radius: tube radius.
        center: ``(3,)`` torus center.

    Returns:
        ``(...)`` signed distances (negative inside the tube).
    """
    p = points - _as_vec3(center)
    xz = mx.stack([p[..., 0], p[..., 2]], axis=-1)
    radial = mx.linalg.norm(xz, axis=-1) - major_radius  # sqrt(x^2 + z^2) - R
    q = mx.stack([radial, p[..., 1]], axis=-1)
    return mx.linalg.norm(q, axis=-1) - minor_radius

sdf_union(*distances)

Boolean union of shapes: the pointwise minimum of their distances.

Parameters:

Name	Type	Description	Default
*distances	array	two or more broadcastable distance arrays.	()

Returns:

Type	Description
array	Combined signed distances.

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

def sdf_union(*distances: mx.array) -> mx.array:
    """Boolean union of shapes: the pointwise minimum of their distances.

    Args:
        *distances: two or more broadcastable distance arrays.

    Returns:
        Combined signed distances.
    """
    if len(distances) < 2:
        raise ValueError("sdf_union needs at least two distance arrays.")
    out = distances[0]
    for d in distances[1:]:
        out = mx.minimum(out, d)
    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)])