Curvature Estimation¶

def delta_hyperbolicity(
    distance_matrix: Float[torch.Tensor, "n_points n_points"],
    samples: int | None = None,
    reference_idx: int = 0,
    relative: bool = True,
) -> Float[torch.Tensor, "n_points n_points n_points"] | Float[torch.Tensor, "samples"]:
    r"""Computes δ-hyperbolicity from a distance matrix.

    For each triplet of points (x,y,z) and reference point w, computes:
    δ(x,y,z) = min((x,y)_w, (y,z)_w) - (x,z)_w

    where (a,b)_w = ½(d(w,a) + d(w,b) - d(a,b)) is the Gromov product.

    Args:
        distance_matrix: Pairwise distance matrix.
        samples: Number of triplets to sample. If None, computes full δ tensor over all triplets.
        reference_idx: Index of the reference point w.
        relative: Whether to normalize by maximum distance.

    Returns:
        δ-hyperbolicity estimates:
        - When samples is not None: torch.Tensor of shape (samples,)
        - When samples is None: torch.Tensor of shape (n_points, n_points, n_points)

    Note:
        For global statistics, call .max() or other aggregation functions on the result.
    """
    if not isinstance(distance_matrix, torch.Tensor):
        raise TypeError(f"distance_matrix must be a torch.Tensor, got {type(distance_matrix)}")

    D = distance_matrix.float()

    if samples is not None:
        return _sample_delta_values(D, samples, reference_idx, relative)
    else:
        return _compute_full_delta_tensor(D, reference_idx, relative)

def greedy_signature_selection(
    candidate_components: Iterable[tuple[float, int]] = ((-1.0, 2), (0.0, 2), (1.0, 2)),
    max_components: int = 3,
    pipeline: Callable[..., float] = distortion_pipeline,
    verbose: bool = False,
    **kwargs: dict[str, Any],
) -> tuple[ProductManifold, list[float]]:
    r"""Greedily estimates an optimal product manifold signature.

    This implements the greedy signature selection algorithm that incrementally builds a product manifold
    by selecting components that best preserve distances. At each step, it chooses the manifold component
    that maximizes distortion reduction.

    Args:
        candidate_components: Candidate (curvature, dimension) pairs to consider.
        max_components: Maximum number of components to include.
        pipeline: Function that takes a ProductManifold, plus additional arguments, and returns a loss value.
        verbose: If True, prints progress information.
        **kwargs: Additional keyword arguments to pass to the pipeline function.

    Returns:
        optimal_pm: Optimized product manifold with the selected signature.
    """
    # Initialize variables
    signature: list[tuple[float, int]] = []
    loss_history: list[float] = []
    current_loss = float("inf")
    candidate_components_list = list(candidate_components)  # For type safe iteration

    # Greedy loop
    for i in range(max_components):
        if verbose:
            print(f"Iteration {i + 1}/{max_components}")
        best_loss, best_idx = current_loss, -1

        # Try each candidate
        for idx, comp in enumerate(candidate_components_list):
            if verbose:
                print(f"  Trying component {comp} (index {idx})")
            pm = ProductManifold(signature=signature.copy() + [comp])
            loss = pipeline(pm, **kwargs)
            if loss < best_loss:
                best_loss, best_idx = loss, idx

        # If no improvement, stop
        if best_idx < 0:
            if verbose:
                print("No improvement found, stopping.")
            break

        # Otherwise accept that component
        signature.append(candidate_components_list[best_idx])
        current_loss = best_loss
        loss_history.append(current_loss)
        if verbose:
            print(f"  Accepted component {candidate_components_list[best_idx]} with loss {current_loss:.4f}")
            print(f"  Current signature: {signature}")
            print()

    # Return final manifold
    optimal_pm = ProductManifold(signature=signature.copy())
    return optimal_pm, loss_history

def sectional_curvature(
    adjacency_matrix: Float[torch.Tensor, "n_points n_points"],
    distance_matrix: Float[torch.Tensor, "n_points n_points"],
    samples: int | None = None,
    relative: bool = True,
) -> Float[torch.Tensor, "n_points"] | Float[torch.Tensor, "samples"]:
    r"""Estimates sectional curvature of a graph.

    Uses discrete triangle comparison theorem to estimate local curvature.
    Positive values indicate spherical-like regions, negative values indicate
    hyperbolic-like regions, zero indicates flat regions.

    Args:
        adjacency_matrix: Binary adjacency matrix indicating graph connections.
        distance_matrix: Pairwise shortest path distance matrix.
        samples: Number of triangle configurations to sample. If None, computes per-node curvatures.
        relative: Whether to normalize by maximum distance.

    Returns:
        Sectional curvature estimates:
        - When samples is not None: torch.Tensor of shape (samples,)
        - When samples is None: torch.Tensor of shape (n_points,) with per-node curvatures

    Note:
        For global statistics, call .mean() or other aggregation functions on the result.
    """
    if not isinstance(adjacency_matrix, torch.Tensor) or not isinstance(distance_matrix, torch.Tensor):
        raise TypeError("Both adjacency_matrix and distance_matrix must be torch.Tensors")

    if adjacency_matrix.shape != distance_matrix.shape:
        raise ValueError("Adjacency matrix and distance matrix must have the same shape")

    A = adjacency_matrix.float()
    D = distance_matrix.float()

    if samples is not None:
        return _sample_curvatures(D, samples, relative)
    else:
        return _compute_node_curvatures(A, D, relative)