Utils¶

def clone_optional_tensor(tensor: Tensor | None) -> Tensor | None:
    """
    Clone a tensor when it is provided.

    Args:
        tensor: Optional tensor to clone.

    Returns:
        tensor: A cloned tensor, or ``None`` when no tensor is provided.
    """
    return tensor.clone() if tensor is not None else None

def empty_edgeattr(num_edges: int) -> Tensor:
    """
    Create an empty edge attribute tensor for a fixed number of edges.

    Args:
        num_edges: Number of edge rows to allocate.

    Returns:
        edge_attr: Empty floating-point tensor of shape ``(num_edges, 0)``.
    """
    return torch.empty((num_edges, 0), dtype=torch.float)

def empty_hyperedgeindex() -> Tensor:
    """
    Create an empty hyperedge index tensor.

    Returns:
        hyperedge_index: Empty long tensor of shape ``(2, 0)``.
    """
    return torch.empty((2, 0), dtype=torch.long)

def empty_nodefeatures() -> Tensor:
    """
    Create an empty node feature tensor.

    Returns:
        features: Empty floating-point tensor of shape ``(0, 0)``.
    """
    return torch.empty((0, 0), dtype=torch.float)

def escape(text: str, escaped_characters_table: dict[str, str]) -> str:
    """
    Escape characters in text according to the provided escape table.

    Args:
        text: The input string to escape.
        escaped_characters_table: A dictionary mapping characters to their escaped versions.

    Returns:
        escaped_text: The escaped string based on the escape table.
    """
    return text.translate(str.maketrans(escaped_characters_table))

def to_non_empty_edgeattr(edge_attr: Tensor | None) -> Tensor:
    """
    Convert optional edge attributes to a tensor with an edge dimension.

    Args:
        edge_attr: Optional edge attribute tensor.

    Returns:
        edge_attr: The provided tensor, or an empty tensor with the inferred edge count.
    """
    num_edges = edge_attr.size(0) if edge_attr is not None else 0
    return empty_edgeattr(num_edges) if edge_attr is None else edge_attr

def to_0based_ids(original_ids: Tensor, ids_to_rebase: Tensor | None = None) -> Tensor:
    """
    Remap IDs to contiguous 0-based indices.

    If ``ids_to_rebase`` is provided, only IDs present in it are kept and remapped.
    If ``ids_to_rebase`` is not provided, all unique IDs in ``original_ids`` are remapped.

    Examples:
        >>> to_0based_ids(torch.tensor([1, 3, 3, 7]), torch.tensor([3, 7]))
        ... -> tensor([0, 0, 1])  # 1 is excluded, 3 -> 0, 7 -> 1

        >>> to_0based_ids(torch.tensor([5, 3, 5, 8]))
        ... -> tensor([1, 0, 1, 2])  # 3 -> 0, 5 -> 1, 8 -> 2

    Args:
        original_ids: Tensor of original IDs.
        ids_to_rebase: Optional tensor of IDs to keep and remap.
            If ``None``, all unique IDs are used. Defaults to ``None``.

    Returns:
        ids: Tensor of 0-based IDs.
    """
    if ids_to_rebase is None:
        sorted_unique_original_ids = original_ids.unique(sorted=True)
        return torch.searchsorted(sorted_unique_original_ids, original_ids)

    keep_mask = torch.isin(original_ids, ids_to_rebase)
    ids_to_keep = original_ids[keep_mask]
    sorted_unique_ids_to_rebase = ids_to_rebase.unique(sorted=True)
    return torch.searchsorted(sorted_unique_ids_to_rebase, ids_to_keep)

def validate_is_between(
    name: str,
    value: int | float,
    min_value: int | float,
    max_value: int | float,
) -> None:
    """
    Validate that a numeric value is finite and lies within inclusive bounds.

    Args:
        name: Name of the validated value.
        value: Numeric value to validate.
        min_value: Inclusive lower bound.
        max_value: Inclusive upper bound.

    Raises:
        ValueError: If the bounds are invalid or the value is outside them.
    """
    if min_value > max_value:
        raise ValueError(
            f"Invalid bounds for {name!r}: 'min_value' ({min_value}) cannot "
            f"be greater than 'max_value' ({max_value})."
        )
    if not math.isfinite(value) or value < min_value or value > max_value:
        raise ValueError(
            f"{name!r} must be between {min_value} and {max_value} inclusive, got {value}."
        )

def validate_is_finite(name: str, value: int | float) -> None:
    """
    Validate that a numeric value is finite.

    Args:
        name: Name of the validated value.
        value: Numeric value to validate.

    Raises:
        ValueError: If the value is not finite.
    """
    if not math.isfinite(value):
        raise ValueError(f"{name!r} must be finite, got {value}.")

def validate_is_finite_when_provided(name: str, value: int | float | None) -> None:
    """
    Validate that an optional numeric value is finite when provided.

    Args:
        name: Name of the validated value.
        value: Optional numeric value to validate.

    Raises:
        ValueError: If the provided value is not finite.
    """
    if value is not None and not math.isfinite(value):
        raise ValueError(f"{name!r} must be finite when provided, got {value}.")

def validate_is_non_empty(name: str, value: Sequence) -> None:
    """
    Validate that a sequence is not empty.

    Args:
        name: Name of the validated sequence.
        value: Sequence to validate.

    Raises:
        ValueError: If the sequence is empty.
    """
    if len(value) < 1:
        raise ValueError(f"{name!r} cannot be empty.")

def validate_is_non_negative(name: str, value: int | float) -> None:
    """
    Validate that a numeric value is non-negative.

    Args:
        name: Name of the validated value.
        value: Numeric value to validate.

    Raises:
        ValueError: If the value is negative.
    """
    if value < 0:
        raise ValueError(f"{name!r} must be non-negative, got {value}.")

def validate_is_positive(name: str, value: int | float) -> None:
    """
    Validate that a numeric value is positive.

    Args:
        name: Name of the validated value.
        value: Numeric value to validate.

    Raises:
        ValueError: If the value is not positive.
    """
    if value <= 0:
        raise ValueError(f"{name!r} must be positive, got {value}.")

def validate_ratios(ratios: list[int | float]) -> None:
    """
    Validate split ratios.

    Args:
        ratios: Ratios that must be positive, finite, non-empty, and sum to one.

    Raises:
        ValueError: If any ratio is invalid or the ratios do not sum to one.
    """
    validate_is_non_empty("ratios", ratios)

    for ratio in ratios:
        validate_is_finite("ratios", ratio)
        validate_is_positive("ratios", ratio)

    # Allow small imprecision in sum of ratios, but raise error if it's significant
    # Example: ratios = [0.8, 0.1, 0.1] -> sum = 1.0 (valid)
    #          ratios = [0.8, 0.1, 0.05] -> sum = 0.95 (invalid, raises ValueError)
    #          (valid, allows small imprecision)
    #          ratios = [0.8, 0.1, 0.1, 0.0000001] -> sum = 1.0000001
    ratio_sum = float(sum(ratios))
    if abs(ratio_sum - 1.0) > 1e-6:
        raise ValueError(f"'ratios' must sum to 1.0, got {ratio_sum}.")

def get_hf_datasets_shas(
    dataset_names: list[str], namespace: str = "HypernetworkRG"
) -> dict[str, str | None]:
    """
    Retrieve Hugging Face dataset commit SHAs for multiple datasets.

    Args:
        dataset_names: Dataset names to query.
        namespace: Hugging Face namespace containing the datasets.
            Defaults to ``"HypernetworkRG"``.

    Returns:
        shas: Mapping from dataset name to commit SHA, or ``None`` when unavailable.
    """
    shas: dict[str, str | None] = {}

    for dataset_name in dataset_names:
        shas[dataset_name] = get_hf_dataset_sha(dataset_name, namespace)
    return shas

def get_hf_dataset_sha(dataset_name: str, namespace: str = "HypernetworkRG") -> str | None:
    """
    Retrieve the latest Hugging Face commit SHA for a dataset.

    Args:
        dataset_name: Dataset name to query.
        namespace: Hugging Face namespace containing the dataset.
            Defaults to ``"HypernetworkRG"``.

    Returns:
        sha: Dataset repository commit SHA, or ``None`` when unavailable.
    """
    api = HfApi()
    repo_id = f"{namespace}/{dataset_name}"
    try:
        info = api.dataset_info(repo_id=repo_id)
        return info.sha
    except Exception as e:
        warnings.warn(
            f"{dataset_name}: failed to retrieve SHA ({e})",
            category=UserWarning,
            stacklevel=2,
        )
        return None

def get_gh_datasets_shas(
    dataset_names: list[str],
    owner: str = "hypernetwork-research-group",
    repository: str = "datasets",
) -> dict[str, str | None]:
    """
    Retrieve GitHub commit SHAs for multiple compressed dataset files.

    Args:
        dataset_names: Dataset file stems to query.
        owner: GitHub repository owner.
        repository: GitHub repository name. Defaults to ``"datasets"``.

    Returns:
        shas: Mapping from dataset name to commit SHA, or ``None`` when unavailable.
    """
    shas: dict[str, str | None] = {}

    for dataset_name in dataset_names:
        shas[dataset_name] = get_gh_dataset_sha(dataset_name, owner, repository)
    return shas

def get_gh_dataset_sha(dataset_name: str, owner: str, repository: str) -> str | None:
    """
    Retrieve the latest GitHub commit SHA for a dataset file.

    Args:
        dataset_name: Dataset file stem without the ``.json.zst`` suffix.
        owner: GitHub repository owner.
        repository: GitHub repository name.

    Returns:
        sha: Latest commit SHA for the dataset file, or ``None`` when unavailable.
    """
    url = f"https://api.github.com/repos/{owner}/{repository}/commits"
    file_path = f"{dataset_name}.json.zst"

    params = {
        "path": file_path,
        "per_page": 1,  # Latest commit only
    }

    try:
        response = requests.get(url, params=params, timeout=10)
        response.raise_for_status()
    except requests.RequestException as e:
        warnings.warn(
            f"{dataset_name}: failed to retrieve SHA ({e})",
            category=UserWarning,
            stacklevel=2,
        )
        return None

    data = response.json()

    if data:
        commit_sha = data[0]["sha"]
    else:
        warnings.warn(
            f"{dataset_name}: no commits found for {file_path}",
            category=UserWarning,
            stacklevel=2,
        )
        return None
    return commit_sha

def validate_hif_data(hif_data: dict[str, Any]) -> bool:
    """
    Validate a Python object against the HIF schema.

    Args:
        hif_data: Parsed HIF data to validate.

    Returns:
        valid: ``True`` when the data conforms to the schema, otherwise ``False``.
    """
    schema = __load_hif_schema()
    validator = fastjsonschema.compile(schema)
    try:
        validator(hif_data)
        return True
    except Exception:
        return False

def validate_hif_json(filename: str) -> bool:
    """
    Validate a JSON file against the HIF (Hypergraph Interchange Format) schema.

    Args:
        filename: Path to the JSON file to validate.

    Returns:
        valid: ``True`` if the file is valid HIF, ``False`` otherwise.

    Raises:
        ValueError: If the JSON file cannot be read.
    """
    try:
        with open(filename, encoding="utf-8") as f:
            hiftext = json.load(f)
            return validate_hif_data(hiftext)
    except Exception as e:
        raise ValueError(f"Failed to read JSON file {filename!r}: {e!s}.") from e

def is_input_layer(layer_idx: int) -> bool:
    """
    Check whether a layer index points to the input layer.

    Args:
        layer_idx: Layer index to inspect.

    Returns:
        result: ``True`` when ``layer_idx`` is the input layer index.
    """
    return is_layer(layer_idx, INPUT_LAYER)

def is_layer(layer_idx: int, desired_layer: int) -> bool:
    """
    Check whether a layer index matches a desired index.

    Args:
        layer_idx: Layer index to inspect.
        desired_layer: Target layer index.

    Returns:
        result: ``True`` when the indices match.
    """
    return layer_idx == desired_layer

def node_labels_from_node_degrees(
    node_incidences: Tensor,
    num_nodes: int,
    num_classes: int = 3,
) -> Tensor:
    """
    Create node labels by binning nodes according to hypergraph degree.

    The input is the node row of a hyperedge index, where each occurrence
    represents one node-hyperedge incidence. The function counts those
    occurrences to get one degree per node, then splits the degree distribution
    into ``num_classes`` quantile bins.

    Examples:
        >>> node_incidences = torch.tensor([1, 2, 2, 3, 3, 3, 4, 4, 4, 4])
        >>> labels = node_labels_from_node_degrees(node_incidences, num_nodes=5, num_classes=3)
        >>> labels
        ... tensor([0, 0, 1, 2, 2])

        Here, ``node_degrees = [0, 1, 2, 3, 4]``:
        node ``0`` is isolated, node ``1`` appears once, node ``2`` appears twice,
        node ``3`` appears three times, and node ``4`` appears four times. With
        three classes, the labels represent low, medium, and high degree bins.

    Args:
        node_incidences: A 1D tensor containing one node ID per hypergraph incidence,
            typically ``hdata.hyperedge_index[0]``.
        num_nodes: Total number of nodes, including isolated nodes that may not
            appear in ``node_incidences``.
        num_classes: Number of degree classes to produce.

    Returns:
        labels: A long tensor of shape ``[num_nodes]`` containing one class label
            per node.
    """
    # Count one occurrence per node-hyperedge incidence.
    # Example: node_incidences = [1, 2, 2, 3, 3, 3, 4, 4, 4, 4], num_nodes = 5
    #                       nodes 0  1  2  3  4  # node 0 is isolated
    #          -> node_degrees = [0, 1, 2, 3, 4]
    node_degrees = torch.bincount(input=node_incidences, minlength=num_nodes).float()

    # Build quantile cut points between classes.
    # Example: num_classes = 3
    #          -> quantiles = [1/3, 2/3]
    #          -> thresholds ~= [1.333, 2.667]
    #          -> thresholds split nodes into low, medium, and high degree bins
    thresholds = torch.quantile(
        input=node_degrees,
        q=torch.tensor(
            [i / num_classes for i in range(1, num_classes)],
            device=node_degrees.device,
        ),
    )

    # Convert each degree into the index of its quantile bin.
    # Example: node_degrees = [0, 1, 2, 3, 4], thresholds ~= [1.333, 2.667]
    #          -> labels = [0, 0, 1, 2, 2]
    labels = torch.bucketize(input=node_degrees, boundaries=thresholds).long()
    return labels

def maxmin_scatter(
    src: Tensor,
    index: Tensor,
    dim: int,
    dim_size: int | None = None,
) -> Tensor:
    """
    Performs a scatter reduction that computes the channel-wise range (max - min) for each
    index group.

    Args:
        src: The source tensor containing the values to scatter.
        index: The indices of elements to scatter.
        dim: The axis along which to index.
        dim_size: The size of the output tensor along the scatter dimension.
            If not provided, it will be inferred from the maximum index value.
            Defaults to ``None``.

    Returns:
        values: A tensor containing the max-min values for each index group.
    """
    max_embeddings = scatter(src=src, index=index, dim=dim, dim_size=dim_size, reduce="max")
    min_embeddings = scatter(src=src, index=index, dim=dim, dim_size=dim_size, reduce="min")
    return max_embeddings - min_embeddings

def validate_floating_tensor_dtype(name: str, tensor: Tensor) -> None:
    """
    Validate that a tensor has a floating-point dtype.

    Args:
        name: Name of the validated tensor.
        tensor: Tensor to validate.

    Raises:
        ValueError: If the tensor does not have a floating-point dtype.
    """
    if not tensor.is_floating_point():
        raise ValueError(f"{name!r} must have a floating-point dtype, got {tensor.dtype}.")

def validate_long_tensor_dtype(name: str, tensor: Tensor) -> None:
    """
    Validate that a tensor has dtype ``torch.long``.

    Args:
        name: Name of the validated tensor.
        tensor: Tensor to validate.

    Raises:
        ValueError: If the tensor does not have dtype ``torch.long``.
    """
    if tensor.dtype != torch.long:
        raise ValueError(f"{name!r} must have dtype torch.long, got {tensor.dtype}.")

def assign_hyperedge_label_to_nodes(
    hyperedge_index: Tensor,
    y: Tensor,
    num_hyperedges: int,
) -> dict[frozenset[int], float]:
    """
    Build a mapping from each hyperedge node set to its label.

    Args:
        hyperedge_index: Hyperedge incidence tensor in COO format.
        y: Hyperedge labels indexed by hyperedge ID.
        num_hyperedges: Number of hyperedges to inspect.

    Returns:
        labels_by_nodes: Mapping from frozen node sets to labels.
    """
    labels_by_nodes: dict[frozenset[int], float] = {}
    for hyperedge_id in range(num_hyperedges):
        mask = hyperedge_index[1] == hyperedge_id
        nodes = frozenset(hyperedge_index[0][mask].tolist())
        labels_by_nodes[nodes] = y[hyperedge_id].item()
    return labels_by_nodes

def is_inductive_setting(node_space_setting: NodeSpaceSetting | None) -> bool:
    """
    Check whether a node space setting is inductive.

    Args:
        node_space_setting: Node space setting to inspect.

    Returns:
        result: ``True`` when the setting is ``"inductive"``.
    """
    return node_space_setting == "inductive"

def is_transductive_setting(node_space_setting: NodeSpaceSetting | None) -> bool:
    """
    Check whether a node space setting is transductive.

    Args:
        node_space_setting: Node space setting to inspect.

    Returns:
        result: ``True`` when the setting is ``"transductive"``.
    """
    return node_space_setting == "transductive"

def validate_node_space_setting(node_space_setting: NodeSpaceSetting) -> None:
    """
    Validate that the node space setting is one of the supported values.

    Args:
        node_space_setting: The node space setting to validate, which should be either "inductive"
            or "transductive".

    Raises:
        ValueError: If the node space setting is not one of the supported values.
    """
    if is_transductive_setting(node_space_setting) or is_inductive_setting(node_space_setting):
        return

    raise ValueError(
        f"'node_space_setting' must be one of 'transductive' or 'inductive', "
        f"got {node_space_setting!r}."
    )

def create_seeded_torch_generator(
    device: torch.device,
    seed: int | None,
) -> Generator | None:
    """
    Create a seeded torch generator when a seed is provided.

    Args:
        device: Device where the generator should be created.
        seed: Optional seed for deterministic random operations.

    Returns:
        generator: A seeded torch.Generator when ``seed`` is provided, otherwise ``None``.
    """
    if seed is None:
        return None
    generator = Generator(device=device)
    generator.manual_seed(seed)
    return generator

def sparse_dropout(
    sparse_tensor: Tensor,
    dropout_prob: float,
    fill_value: float = 0.0,
) -> Tensor:
    """
    Dropout function for sparse matrix.

    Returns a new sparse matrix with the same shape as the input sparse matrix,
    but with some elements dropped out.

    Args:
        sparse_tensor: The sparse matrix with format ``torch.sparse_coo_tensor``.
        dropout_prob: Probability of an element to be dropped.
        fill_value: The fill value for dropped elements. Defaults to ``0.0``.

    Returns:
        matrix: A new sparse matrix with the same shape as the input sparse matrix,
            but with some elements dropped out.

    Raises:
        ValueError: If ``dropout_prob`` is outside the range ``[0, 1]``.
    """
    device = sparse_tensor.device

    # Sparse tensors may be unsorted indices or have duplicate entries
    # 'coalesce()' will sum duplicates and sort indices to have a consistent format for dropout
    sparse_tensor = sparse_tensor.coalesce()

    if dropout_prob > 1 or dropout_prob < 0:
        raise ValueError("Dropout probability must be in the range [0, 1]")

    # Nothing to drop, return the original sparse tensor
    if dropout_prob == 0:
        return sparse_tensor

    values = sparse_tensor.values()
    indices = sparse_tensor.indices()

    keep_prob = 1 - dropout_prob

    # Generate a binary mask matching the shape of values for elements to keep
    # 'torch.bernoulli()' samples 1 with probability keep_prob and 0 with probability dropout_prob
    # Example:
    #   values = [0.5, 1.2, 3.4], keep_prob = 0.8
    #   -> keep_mask might be [1, 0, 1], meaning we keep the 1st and 3rd elements, drop the 2nd
    keep_mask = torch.bernoulli(torch.full_like(values, keep_prob)).to(device)

    if fill_value == 0.0:
        # If fill_value is 0, just zero out the dropped elements,
        # as keep_mask will be 0 for dropped elements and 1 for kept elements
        # Example: values = [0.5, 1.2, 3.4], keep_mask = [1, 0, 1], fill_value = 0.0
        #          -> new_values = [0.5*1, 1.2*0, 3.4*1] = [0.5, 0.0, 3.4]
        new_values = values * keep_mask
    else:
        # If fill_value is non-zero, we must fill the dropped elements with the
        # specified fill_value instead of zero
        # 'torch.logical_not(keep_mask)' identifies dropped elements where mask is 0 and
        # Example: values = [0.5, 1.2, 3.4], keep_mask = [1, 0, 1], fill_value = 9.9
        #          -> values_to_fill_mask = [0, 1, 0]
        #          -> fill_values = [0*9.9, 1*9.9, 0*9.9] = [0.0, 9.9, 0.0]
        #          -> new_values = [0.5*1 + 0.0, 1.2*0 + 9.9, 3.4*1 + 0.0] = [0.5, 9.9, 3.4]
        values_to_fill_mask = torch.logical_not(keep_mask)
        fill_values = values_to_fill_mask * fill_value
        new_values = values * keep_mask + fill_values

    # Reuse the original indices and shape to preserve spasity but change values
    dropout_sparse_tensor = torch.sparse_coo_tensor(
        indices=indices,
        values=new_values,
        size=sparse_tensor.size(),
        dtype=sparse_tensor.dtype,
        device=device,
    )

    return dropout_sparse_tensor

def validate_http_url(value: str) -> str:
    """
    Validate that a URL uses HTTP or HTTPS.

    Args:
        value: URL string to validate.

    Returns:
        value: The validated URL.

    Raises:
        ValueError: If the URL is not an absolute HTTP or HTTPS URL.
    """
    parsed = urlparse(value)
    if parsed.scheme not in {"http", "https"} or not parsed.netloc:
        raise ValueError(f"Invalid URL: {value}")
    return value