Function fit_hyperplane

plib/utils.py:1822–1892

Fit a hyperplane for the n points such that the plane minimizes the point-to-plane distances. .. math:: \min_{n,c} \sum_i w_i (n^T * (p_i - c))^2, where :math:`n` is the normal of the hyperplane, :math:`c` is the anchor of the plane. Args: points: (*

(
        points: torch.Tensor,  # (*, n, d)
        centers: torch.Tensor = None,  # (*, d)
        weights: torch.Tensor = None,  # (*, n)
        th_eig_val: float = 1.e-3,
)

def fit_hyperplane(
        points: torch.Tensor,  # (*, n, d)
        centers: torch.Tensor = None,  # (*, d)
        weights: torch.Tensor = None,  # (*, n)
        th_eig_val: float = 1.e-3,
) -> T.Dict[str, T.Any]:
    """
    Fit a hyperplane for the n points such that the plane minimizes the point-to-plane distances.

    .. math::

        \min_{n,c}  \sum_i w_i (n^T * (p_i - c))^2,

    where :math:`n` is the normal of the hyperplane, :math:`c` is the anchor of the plane.

    Args:
        points:
            (*, n, d) the points in d-dimensinoal space
        centers:
            (*, d) the given anchors of the hyperplanes
        weights:
            (*, n) the weight to each points.  If None, all points have unit weights.
        # ray_origins:
        #     (*, d) the ray origins.  If given, we will also compute the intersection of the ray on the plane.
        # ray_directions:
        #     (*, d) the ray directions. If given, we will also compute the intersection of the ray on the plane.

    Returns:
        plane_normals:
            (*, d)  Note that the plane normal can points to one of the two directions
        centers:
            (*, d)
        ts:
            (*,) or None.  ts on the ray to the plane.

    """
    *p_shape, d = points.shape
    if weights is None:
        weights = torch.ones(*p_shape, 1, dtype=points.dtype, device=points.device)  # (*, n, 1)
    else:
        weights = weights.unsqueeze(-1)  # (*, n, 1)

    if centers is None:
        centers = (points * weights).sum(dim=-2) / weights.sum(dim=-2)  # (*, d)

    points_centered = points - centers.unsqueeze(-2)  # (*, n, d)
    PTP = (weights * points_centered).transpose(-1, -2) @ points_centered  # (*, d, d)

    # we will make sure PTP is at least rank 2 (at least two points not on the same line)
    with torch.no_grad():
        eig_vals, eig_vecs = torch.linalg.eigh(PTP)  # eig_vecs: (*, d, d),  eig_vals: (*, d) small to large
        valid_mask = (eig_vals[..., -2] > th_eig_val)  # (*,)
        valid_mask = torch.logical_and(
            valid_mask,
            (eig_vals[..., 1] - eig_vals[..., 0]) > th_eig_val,
        )  # (*,)

    # fill PTP with full rank matrix if invalid